From 4c4dc41e3b921c78d6590c66b3ff8e1bf57e60c5 Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Tue, 3 Nov 2026 20:41:54 -0800
Subject: [PATCH 01/15] [None][feat] modeling_v2 qwen4_exp: family operations
The Qwen3.8-Flash-Next operations of the qwen4_exp family, in `models/qwen4_exp/ops/`
(activation, attention/QSA, embedding, gemm, hc, moe, norm, ple, quantization, ssm, vision),
and the CuteDSL MoE finalize-fusion kernel in `models/qwen4_exp/kernels/moe/`, each entry with
its cells and tests under `tests/unittest/_torch/modeling_v2/models/qwen4_exp//`. They are
the TP1 target's operations that the family's target calls; the next commit switches the target
to them and removes the TP1 target with its own copies.
Beyond the move:
- `qsa_canonical_selection` runs as one Triton kernel.
- `qsa_paged_index_scores` shares one key gather across a verify step's rows of one request
(its request-group kernel), and its launch metadata does not specialize on pointer alignment.
- Each entry's cells include the target's calls at every layout it serves: tp1, and one rank
of tp4, tep4, dep4 and dep2.
An entry's tests are its cells against its reference, plus the behaviors the target relies on:
CUDA-graph replays that read after a PDL wait or on new inputs, bitwise equivalences between
variants, padding and batch invariance.
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
.../models/qwen4_exp/kernels/__init__.py | 3 +
.../models/qwen4_exp/kernels/moe/__init__.py | 2 +
...contiguous_grouped_gemm_finalize_fusion.py | 2738 +++++++++++++++++
.../models/qwen4_exp/ops/__init__.py | 3 +
.../qwen4_exp/ops/activation/__init__.py | 2 +
.../activation/fused_sigmoid_mul_inplace.py | 227 ++
.../qwen4_exp/ops/activation/gelu_erf.py | 106 +
.../qwen4_exp/ops/activation/gelu_tanh.py | 116 +
.../ops/activation/silu_and_mul_rows.py | 106 +
.../qwen4_exp/ops/attention/__init__.py | 2 +
.../ops/attention/qsa_canonical_selection.py | 608 ++++
.../ops/attention/qsa_expand_block_indices.py | 221 ++
.../ops/attention/qsa_mtp_reuse_topk.py | 322 ++
.../ops/attention/qsa_nvfp4_pages_to_bf16.py | 391 +++
.../ops/attention/qsa_nvfp4_verify_ptx.py | 251 ++
.../ops/attention/qsa_paged_index_scores.py | 1158 +++++++
.../ops/attention/qsa_paged_index_store.py | 473 +++
.../ops/attention/qsa_paged_kv_store.py | 314 ++
.../ops/attention/qsa_paged_kv_store_nvfp4.py | 572 ++++
.../ops/attention/qsa_paged_sparse_gqa.py | 484 +++
.../attention/qsa_paged_sparse_gqa_decode.py | 820 +++++
.../attention/qsa_paged_sparse_gqa_nvfp4.py | 1159 +++++++
.../qsa_paged_sparse_gqa_nvfp4_decode.py | 1606 ++++++++++
.../qsa_paged_sparse_gqa_nvfp4_verify.py | 2147 +++++++++++++
.../ops/attention/qsa_prefill_compress.py | 375 +++
.../ops/attention/qsa_rollback_planes.py | 241 ++
.../attention/qsa_rowtile_attention_cute.py | 1383 +++++++++
.../ops/attention/qsa_sparse_gqa_rowtile.py | 938 ++++++
.../ops/attention/qsa_topk_expand_decode.py | 1556 ++++++++++
.../ops/attention/qsa_topk_select.py | 299 ++
.../qwen4_exp_attn_decode_prologue.py | 1023 ++++++
.../ops/attention/qwen4_exp_qk_norm_rope.py | 1193 +++++++
.../ops/attention/spec_rollback_commit.py | 487 +++
.../ops/attention/verify_step_setup.py | 216 ++
.../qwen4_exp/ops/embedding/__init__.py | 2 +
.../qwen4_exp/ops/embedding/ngram_hash.py | 286 ++
.../models/qwen4_exp/ops/gemm/__init__.py | 2 +
.../qwen4_exp/ops/gemm/bf16_gemv_pdl.py | 388 +++
.../models/qwen4_exp/ops/gemm/fp8_gemv.py | 2016 ++++++++++++
.../qwen4_exp/ops/gemm/fp8_per_tensor_mm.py | 192 ++
.../qwen4_exp/ops/gemm/low_m_bf16_mm.py | 256 ++
.../models/qwen4_exp/ops/gemm/mid_rows_mm.py | 1380 +++++++++
.../models/qwen4_exp/ops/gemm/nvfp4_gemv.py | 992 ++++++
.../models/qwen4_exp/ops/hc/__init__.py | 2 +
.../ops/hc/qwen4_exp_gated_residual.py | 196 ++
.../ops/hc/qwen4_exp_hc_decode_first_cute.py | 605 ++++
.../hc/qwen4_exp_hc_decode_mix_fp4_cute.py | 534 ++++
.../ops/hc/qwen4_exp_hc_decode_rows_tc.py | 1353 ++++++++
.../ops/hc/qwen4_exp_hc_first_rescale_tc.py | 976 ++++++
.../ops/hc/qwen4_exp_hc_inject_norm.py | 1123 +++++++
.../ops/hc/qwen4_exp_hc_inject_rstd.py | 258 ++
.../ops/hc/qwen4_exp_hc_lowrank_act.py | 184 ++
.../ops/hc/qwen4_exp_hc_up_gate_mean.py | 922 ++++++
.../ops/hc/qwen4_exp_hc_up_gate_mean_rstd.py | 1126 +++++++
.../qwen4_exp/ops/hc/qwen4_exp_post_inject.py | 153 +
.../models/qwen4_exp/ops/moe/__init__.py | 2 +
...sl_nvfp4_gather_grouped_gemm_act_fusion.py | 484 +++
.../ops/moe/cute_dsl_nvfp4_grouped_gemm.py | 367 +++
.../fp4_block_scale_moe_renorm_permuted.py | 139 +
.../moe/fp4_block_scale_moe_renorm_routed.py | 484 +++
.../ops/moe/fp4_moe_renorm_cute_dsl.py | 314 ++
.../ops/moe/fp8_block_scale_moe_runner.py | 434 +++
.../ops/moe/fused_sigmoid_gate_mul_add.py | 171 +
.../models/qwen4_exp/ops/moe/moe_sort.py | 272 ++
.../ops/moe/moe_unpermute_inplace.py | 246 ++
.../ops/moe/renorm_moe_routing_op.py | 198 ++
.../models/qwen4_exp/ops/norm/__init__.py | 2 +
.../ops/norm/flashinfer_gemma_rmsnorm.py | 139 +
.../models/qwen4_exp/ops/norm/layer_norm.py | 145 +
.../ops/norm/rms_norm_gated_token_major.py | 249 ++
.../models/qwen4_exp/ops/ple/__init__.py | 2 +
.../qwen4_exp/ops/ple/ngram_token_history.py | 357 +++
.../ops/ple/ple_gate_conv_inject_norm.py | 761 +++++
.../qwen4_exp/ops/ple/ple_ngram_embed.py | 624 ++++
.../qwen4_exp/ops/quantization/__init__.py | 2 +
.../ops/quantization/bf16_to_fp32.py | 80 +
.../ops/quantization/fp8_quantize_1x128.py | 161 +
.../models/qwen4_exp/ops/ssm/__init__.py | 2 +
.../ops/ssm/causal_conv1d_triton_update.py | 388 +++
.../qwen4_exp/ops/ssm/causal_conv1d_update.py | 588 ++++
.../ops/ssm/cu_seqlens_from_seq_lens.py | 180 ++
.../ssm/flashinfer_chunk_gated_delta_rule.py | 626 ++++
...t_gated_delta_rule_cached_replay_update.py | 616 ++++
.../ops/ssm/gdn_decode_gated_norm_pdl.py | 559 ++++
.../ops/ssm/gdn_decode_rows_gated_norm.py | 1155 +++++++
.../ops/ssm/gdn_prefill_conv_post.py | 1232 ++++++++
.../models/qwen4_exp/ops/vision/__init__.py | 2 +
.../ops/vision/patch_embed_conv3d.py | 179 ++
...t_modeling_v2_fused_sigmoid_mul_inplace.py | 74 +
.../activation/test_modeling_v2_gelu_erf.py | 93 +
.../activation/test_modeling_v2_gelu_tanh.py | 84 +
.../test_modeling_v2_silu_and_mul_rows.py | 44 +
...est_modeling_v2_qsa_canonical_selection.py | 282 ++
...st_modeling_v2_qsa_expand_block_indices.py | 134 +
.../test_modeling_v2_qsa_mtp_reuse_topk.py | 255 ++
...est_modeling_v2_qsa_nvfp4_pages_to_bf16.py | 151 +
...test_modeling_v2_qsa_paged_index_scores.py | 460 +++
.../test_modeling_v2_qsa_paged_kv_store.py | 346 +++
...st_modeling_v2_qsa_paged_kv_store_nvfp4.py | 325 ++
.../test_modeling_v2_qsa_paged_sparse_gqa.py | 353 +++
...modeling_v2_qsa_paged_sparse_gqa_decode.py | 110 +
..._modeling_v2_qsa_paged_sparse_gqa_nvfp4.py | 332 ++
...ng_v2_qsa_paged_sparse_gqa_nvfp4_decode.py | 268 ++
...ng_v2_qsa_paged_sparse_gqa_nvfp4_verify.py | 268 ++
.../test_modeling_v2_qsa_prefill_compress.py | 237 ++
.../test_modeling_v2_qsa_rollback_planes.py | 219 ++
...test_modeling_v2_qsa_sparse_gqa_rowtile.py | 238 ++
...test_modeling_v2_qsa_topk_expand_decode.py | 179 ++
.../test_modeling_v2_qsa_topk_select.py | 105 +
...eling_v2_qwen4_exp_attn_decode_prologue.py | 219 ++
...test_modeling_v2_qwen4_exp_qk_norm_rope.py | 154 +
.../test_modeling_v2_spec_rollback_commit.py | 319 ++
.../test_modeling_v2_verify_step_setup.py | 139 +
.../embedding/test_modeling_v2_ngram_hash.py | 99 +
.../gemm/test_modeling_v2_bf16_gemv_pdl.py | 102 +
.../gemm/test_modeling_v2_fp8_gemv.py | 178 ++
.../test_modeling_v2_fp8_per_tensor_mm.py | 70 +
.../gemm/test_modeling_v2_low_m_bf16_mm.py | 113 +
.../gemm/test_modeling_v2_mid_rows_mm.py | 102 +
.../gemm/test_modeling_v2_nvfp4_gemv.py | 288 ++
...st_modeling_v2_qwen4_exp_gated_residual.py | 126 +
...eling_v2_qwen4_exp_hc_decode_first_cute.py | 188 ++
...ing_v2_qwen4_exp_hc_decode_mix_fp4_cute.py | 152 +
...modeling_v2_qwen4_exp_hc_decode_rows_tc.py | 172 ++
...deling_v2_qwen4_exp_hc_first_rescale_tc.py | 92 +
...st_modeling_v2_qwen4_exp_hc_inject_norm.py | 206 ++
...st_modeling_v2_qwen4_exp_hc_inject_rstd.py | 139 +
...st_modeling_v2_qwen4_exp_hc_lowrank_act.py | 57 +
...t_modeling_v2_qwen4_exp_hc_up_gate_mean.py | 67 +
...eling_v2_qwen4_exp_hc_up_gate_mean_rstd.py | 137 +
.../test_modeling_v2_qwen4_exp_post_inject.py | 76 +
...sl_nvfp4_gather_grouped_gemm_act_fusion.py | 366 +++
...modeling_v2_cute_dsl_nvfp4_grouped_gemm.py | 348 +++
..._v2_fp4_block_scale_moe_renorm_permuted.py | 86 +
...ng_v2_fp4_block_scale_moe_renorm_routed.py | 199 ++
...4_block_scale_moe_renorm_routed_windows.py | 207 ++
..._modeling_v2_fp8_block_scale_moe_runner.py | 257 ++
..._modeling_v2_fused_sigmoid_gate_mul_add.py | 101 +
.../moe/test_modeling_v2_moe_sort.py | 203 ++
.../test_modeling_v2_moe_unpermute_inplace.py | 242 ++
.../test_modeling_v2_renorm_moe_routing_op.py | 115 +
...st_modeling_v2_flashinfer_gemma_rmsnorm.py | 60 +
.../norm/test_modeling_v2_layer_norm.py | 63 +
...v2_qwen4_exp_rms_norm_gated_token_major.py | 78 +
.../test_modeling_v2_ngram_token_history.py | 134 +
...t_modeling_v2_ple_gate_conv_inject_norm.py | 199 ++
.../ple/test_modeling_v2_ple_ngram_embed.py | 181 ++
.../test_modeling_v2_bf16_to_fp32.py | 44 +
.../test_modeling_v2_fp8_quantize_1x128.py | 130 +
...modeling_v2_causal_conv1d_triton_update.py | 130 +
.../test_modeling_v2_causal_conv1d_update.py | 123 +
...st_modeling_v2_cu_seqlens_from_seq_lens.py | 72 +
...ng_v2_flashinfer_chunk_gated_delta_rule.py | 377 +++
...t_gated_delta_rule_cached_replay_update.py | 759 +++++
...t_modeling_v2_gdn_decode_gated_norm_pdl.py | 311 ++
..._modeling_v2_gdn_decode_rows_gated_norm.py | 471 +++
.../test_modeling_v2_gdn_prefill_conv_post.py | 197 ++
.../test_modeling_v2_patch_embed_conv3d.py | 113 +
158 files changed, 58966 insertions(+)
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/kernels/__init__.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/kernels/moe/__init__.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/kernels/moe/blockscaled_contiguous_grouped_gemm_finalize_fusion.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/__init__.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/__init__.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/fused_sigmoid_mul_inplace.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/gelu_erf.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/gelu_tanh.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/silu_and_mul_rows.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/__init__.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_canonical_selection.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_expand_block_indices.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_mtp_reuse_topk.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_nvfp4_pages_to_bf16.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_nvfp4_verify_ptx.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_index_scores.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_index_store.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_kv_store.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_kv_store_nvfp4.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_decode.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4_decode.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4_verify.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_prefill_compress.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_rollback_planes.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_rowtile_attention_cute.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_sparse_gqa_rowtile.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_topk_expand_decode.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_topk_select.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qwen4_exp_attn_decode_prologue.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qwen4_exp_qk_norm_rope.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/spec_rollback_commit.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/verify_step_setup.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/embedding/__init__.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/embedding/ngram_hash.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/__init__.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/bf16_gemv_pdl.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/fp8_gemv.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/fp8_per_tensor_mm.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/low_m_bf16_mm.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/mid_rows_mm.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/nvfp4_gemv.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/__init__.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_gated_residual.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_first_cute.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_rows_tc.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_first_rescale_tc.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_inject_norm.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_inject_rstd.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_lowrank_act.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_up_gate_mean.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_up_gate_mean_rstd.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_post_inject.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/__init__.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/cute_dsl_nvfp4_grouped_gemm.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_block_scale_moe_renorm_permuted.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_block_scale_moe_renorm_routed.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_moe_renorm_cute_dsl.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp8_block_scale_moe_runner.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fused_sigmoid_gate_mul_add.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/moe_sort.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/moe_unpermute_inplace.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/renorm_moe_routing_op.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/__init__.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/flashinfer_gemma_rmsnorm.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/layer_norm.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/rms_norm_gated_token_major.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/__init__.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ngram_token_history.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_gate_conv_inject_norm.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_ngram_embed.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/quantization/__init__.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/quantization/bf16_to_fp32.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/quantization/fp8_quantize_1x128.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/__init__.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_triton_update.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_update.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/cu_seqlens_from_seq_lens.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/flashinfer_chunk_gated_delta_rule.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_decode_gated_norm_pdl.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_decode_rows_gated_norm.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_prefill_conv_post.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/vision/__init__.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/vision/patch_embed_conv3d.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_erf.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_tanh.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_silu_and_mul_rows.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_canonical_selection.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_expand_block_indices.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_mtp_reuse_topk.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_nvfp4_pages_to_bf16.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_index_scores.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_kv_store.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_decode.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_prefill_compress.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_rollback_planes.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_expand_decode.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_select.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_spec_rollback_commit.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_verify_step_setup.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/embedding/test_modeling_v2_ngram_hash.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_bf16_gemv_pdl.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_gemv.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_per_tensor_mm.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_low_m_bf16_mm.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_mid_rows_mm.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_nvfp4_gemv.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_gated_residual.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_rows_tc.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_first_rescale_tc.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_rstd.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_rstd.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_post_inject.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed_windows.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp8_block_scale_moe_runner.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_sort.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_unpermute_inplace.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_renorm_moe_routing_op.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_layer_norm.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ngram_token_history.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_gate_conv_inject_norm.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_ngram_embed.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/quantization/test_modeling_v2_bf16_to_fp32.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/quantization/test_modeling_v2_fp8_quantize_1x128.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_triton_update.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_update.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_prefill_conv_post.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/vision/test_modeling_v2_patch_embed_conv3d.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/kernels/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/kernels/__init__.py
new file mode 100644
index 000000000000..a5adae93fbad
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/kernels/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Kernels used only by this family, called through its ops."""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/kernels/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/kernels/moe/__init__.py
new file mode 100644
index 000000000000..52a7a9daf028
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/kernels/moe/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/kernels/moe/blockscaled_contiguous_grouped_gemm_finalize_fusion.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/kernels/moe/blockscaled_contiguous_grouped_gemm_finalize_fusion.py
new file mode 100644
index 000000000000..078eced39888
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/kernels/moe/blockscaled_contiguous_grouped_gemm_finalize_fusion.py
@@ -0,0 +1,2738 @@
+# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+"""The target's copy of TRT-LLM's CuTe DSL finalize-fused grouped GEMM (the CuteDSL MoE
+FC2), with three options the target's prefill MoE sets (`fp4_moe_renorm_cute_dsl`): the
+MMA instructions per K stage (`mma_inst_tile_k`), an epilogue in which each warp owns
+whole output rows (`row_owned_epilogue`), and L2 cache hints on the operand loads and the
+scatter-add (`l2_hints`). Their defaults keep the shared kernel's behavior."""
+
+from typing import Tuple, Type, Union
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+import cutlass.utils as utils
+import cutlass.utils.blackwell_helpers as sm100_utils
+import cutlass.utils.blockscaled_layout as blockscaled_utils
+from cutlass._mlir.dialects import llvm
+from cutlass.cute.nvgpu import cpasync, tcgen05
+from cutlass.cutlass_dsl import dsl_user_op
+
+from tensorrt_llm._torch.kernels.blackwell.utils import (
+ TRTLLM_ENABLE_PDL,
+ blk_reduce_bf16,
+ blk_reduce_fp16,
+ blk_reduce_fp32,
+ griddepcontrol_launch_dependents,
+ griddepcontrol_wait,
+ is_power_of_2,
+)
+
+#: L2 cache-hint policies (CUTLASS CacheHintSm90 encodings).
+_EVICT_FIRST = 0x12F0000000000000
+_EVICT_LAST = 0x14F0000000000000
+
+
+@dsl_user_op
+def _blk_reduce_bf16_evict_last(dst_gemm, src_smem, size, loc=None, ip=None):
+ llvm.inline_asm(
+ None,
+ [
+ dst_gemm.iterator.llvm_ptr,
+ src_smem.iterator.llvm_ptr,
+ size.ir_value(),
+ cutlass.Int64(_EVICT_LAST).ir_value(),
+ ],
+ "cp.reduce.async.bulk.global.shared::cta.bulk_group.L2::cache_hint.add.noftz.bf16 "
+ "[$0], [$1], $2, $3;",
+ "l,l,r,l",
+ has_side_effects=True,
+ loc=loc,
+ ip=ip,
+ )
+
+
+"""
+High-performance persistent blockscaled contiguous grouped dense GEMM (C = alpha * (SFA * A) * (SFB * B)) example for
+the NVIDIA Blackwell architecture using CUTE DSL.
+- Matrix A is MxKx1, A can be row-major("K"), ValidM is composed of valid m in different groups
+- Matrix B is NxKxL, B can be column-major("K"), L is grouped dimension
+- Matrix C is SxNX1, C can be row-major("N"), ValidM is composed of valid m in different groups
+- Matrix SFA layout is filled internally according to A shape and BlockScaledBasicChunk, which has
+ M x ceil_div(K, sf_vec_size) x L elements respectively
+- Matrix SFB layout is filled internally according to B shape and BlockScaledBasicChunk, which has
+ N x ceil_div(K, sf_vec_size) x L elements respectively
+
+Matrix A/C Memory Layout Diagrams:
+
+ ```
+ Group 0 Group 1 Group 2
+ -+---------+---------+---------+
+ | | | |
+ K| ValidM0 | ValidM1 | ValidM2 |
+ | | | |
+ -+---------+---------+---------+
+ |<- ValidM ->|
+ ```
+ Note: the Group(L) dimension will be flatted into M dimension, and the rest Group(L) size is 1.
+ each ValidM will be aligned to 256 or 128. The alignment is determined by the mma_tiler_mn parameter.
+ For NVFP4, 2CTA, the alignment is 256. For NVFP4, 1CTA, the alignment is 128.
+
+This GEMM kernel supports the following features:
+ - Utilizes Tensor Memory Access (TMA) for efficient memory operations
+ - Utilizes Blackwell's tcgen05.mma for matrix multiply-accumulate (MMA) operations
+ - Implements TMA multicast with cluster to reduce L2 memory traffic
+ - Support persistent tile scheduling to better overlap memory load/store with mma between tiles
+ - Support warp specialization to avoid explicit pipelining between mainloop load and mma
+
+This GEMM works as follows:
+1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations.
+2. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using non-TMA
+ operations.
+2. MMA warp:
+ - Load scale factor A/B from shared memory (SMEM) to tensor memory (TMEM) using tcgen05.cp instruction.
+ - Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction.
+3. EPILOGUE warp (with Fused Finalize for MoE):
+ - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld.
+ - Apply alpha scaling: acc_scaled = alpha * acc
+ - **Fused Finalize Logic** (following TensorRT-LLM's sm90_visitor_scatter.hpp pattern):
+ a) Use permuted_idx_to_expanded_idx to map from permuted row to token/topk indices
+ b) Load router_scale directly from global memory to register (no shared memory)
+ c) Apply router_scale: Final = router_scale * acc_scaled
+ - Type convert Final matrix to output type.
+ - Store C matrix from registers (RMEM) to shared memory (SMEM) to global memory (GMEM) with TMA operations.
+
+SM100 tcgen05.mma.kind.block_scale instructions operate as follows:
+- Read matrix A from SMEM
+- Read matrix B from SMEM
+- Read scalefactor A from TMEM
+- Read scalefactor B from TMEM
+- Write accumulator to TMEM
+The accumulator in TMEM must then be loaded to registers before writing back to GMEM.
+
+.. code-block:: bash
+
+ python blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py \
+ --ab_dtype Float4E2M1FN --out_dtype BFloat16 \
+ --sf_dtype Float8E4M3FN --sf_vec_size 16 \
+ --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \
+ --benchmark 1024x7168x2048x64
+
+To collect performance with NCU profiler:
+
+.. code-block:: bash
+
+ ncu python blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py \\ \
+ --ab_dtype Float4E2M1FN --out_dtype BFloat16 \
+ --sf_dtype Float8E4M3FN --sf_vec_size 16 \
+ --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \
+ --benchmark [80,120,160]x7168x2048x64
+
+Constraints:
+* Supported input data types: mxf8, mxf4, nvf4
+ see detailed valid dtype combinations in below Sm100BlockScaledPersistentDenseGemmKernel class documentation
+* A/B tensor must have the same data type, mixed data type is not supported (e.g., mxf8 x mxf4)
+* Mma tiler M must be 128 or 256(use_2cta_instrs)
+* Mma tiler N must be 64/128/192/256
+* Cluster shape M/N must be positive and power of 2, total cluster size <= 16
+* Cluster shape M must be multiple of 2 if Mma tiler M is 256(use_2cta_instrs)
+* The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned,
+ i.e, number of elements is a multiple of 16 and 32 for Float8 and Float4, respectively.
+
+CUDA Graph Support:
+* For CUDA graph support, the tile_idx_to_expert_idx, A/C matrices, and scale factor A can be padded to a larger size
+ (e.g., permuted_m = m*topK + num_local_experts*(256-1), example: 4096*8 + (256/32)*255 = 34808)
+* Use create_tensors() with permuted_m parameter to automatically pad:
+ - tile_idx_to_expert_idx: padded for invalid tiles
+ - A matrix: padded to permuted_m rows (padding rows contain dummy data)
+ - C matrix: padded to permuted_m rows (output buffer for cuda_graph)
+ - Scale factor A: padded to match A matrix dimensions
+* Kernel handling of padding (similar to masked_grouped_gemm.py):
+ - Scheduler warp checks if tile_idx >= num_non_exiting_tiles to exit
+ - Only valid tiles (tile_idx < num_non_exiting_tiles) are written to tile_info pipeline
+ - When no more valid tiles exist, outer loop exits and calls producer_tail()
+ - Consumer warps process only valid tiles from pipeline
+ - No deadlock or synchronization issues
+* Consumer warps check initial tile against num_non_exiting_tiles and set is_valid_tile=False if
+ tile_idx >= num_non_exiting_tiles
+* Only rows within (aligned_groupm[0]+aligned_groupm[1]+...) contain valid data
+* Padding rows in C matrix will not be written by the kernel
+"""
+
+
+class Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel:
+ """This class implements batched matrix multiplication (C = A x SFA x B x SFB) with support for various data types
+ and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization.
+
+ :param sf_vec_size: Scalefactor vector size.
+ :type sf_vec_size: int
+ :param mma_tiler_mn: Shape of the Matrix Multiply-Accumulate (MMA) tile (M,N)
+ :type mma_tiler_mn: Tuple[int, int]
+ :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing
+ :type cluster_shape_mn: Tuple[int, int]
+
+ :note: In current version, A and B tensor must have the same data type
+ - i.e., Float8E4M3FN for A and Float8E5M2 for B is not supported
+
+ :note: Supported combinations of A/B data types, SF data typs and SF vector size:
+ - MXF8: A/B: Float8E5M2/Float8E4M3FN + SF: Float8E8M0FNU + sf_vec_size: 32
+ - MXF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU + sf_vec_size: 32
+ - NVF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU/Float8E4M3FN + sf_vec_size: 16
+
+ :note: Supported accumulator data types:
+ - Float32
+
+ :note: Supported C data types:
+ - Float32
+ - Float16/BFloat16
+ - Float8E4M3FN/Float8E5M2
+
+ :note: Constraints:
+ - MMA tiler M must be 128 or 256 (use_2cta_instrs)
+ - MMA tiler N must be 64/128/192/256
+ - Cluster shape M must be multiple of 2 if Mma tiler M is 256
+ - Cluster shape M/N must be positive and power of 2, total cluster size <= 16
+ - Also, Cluster shape M/N must be <= 4 for scale factor multicasts due to limited size of scale factors
+
+ Example:
+ >>> gemm = Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel(
+ ... sf_vec_size=16, mma_tiler_mn=(256, 128), cluster_shape_mn=(2, 1)
+ ... )
+ >>> gemm(
+ ... a_tensor, b_tensor, sfa_tensor, sfb_tensor, out_tensor, max_active_clusters, stream
+ ... )
+ """
+
+ def __init__(
+ self,
+ sf_vec_size: int,
+ mma_tiler_mn: Tuple[int, int],
+ cluster_shape_mn: Tuple[int, int],
+ raster_along_m: bool = False,
+ use_expert_counts: bool = False,
+ num_local_experts: int = 0,
+ mma_inst_tile_k: int = 4,
+ l2_hints: bool = False,
+ row_owned_epilogue: bool = False,
+ ):
+ """Initializes the configuration for a Blackwell blockscaled dense GEMM kernel.
+
+ This configuration includes several key aspects:
+
+ 1. MMA Instruction Settings (tcgen05):
+ - mma_tiler_mn: The (M, N) shape of the MMA instruction tiler.
+
+ 2. Cluster Shape:
+ - cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster.
+
+ :param mma_tiler_mn: Tuple (M, N) shape of the MMA instruction.
+ :type mma_tiler_mn: Tuple[int, int]
+ :param cluster_shape_mn: Tuple (ClusterM, ClusterN) shape of the cluster.
+ :type cluster_shape_mn: Tuple[int, int]
+ :param raster_along_m: Boolean, True to use raster along M.
+ :type raster_along_m: bool
+ """
+
+ self.sf_vec_size = sf_vec_size
+ self.acc_dtype = cutlass.Float32
+ self.use_2cta_instrs = mma_tiler_mn[0] == 256
+ self.cluster_shape_mn = cluster_shape_mn
+ self.raster_along_m = raster_along_m
+ # MMA instructions (K 64 each for NVFP4) per A/B pipeline stage. At N 256 the 66 KB C
+ # staging buffer leaves room for only 2 stages of the default K 256, so with a short K
+ # (an FC2 of K 640 is 2.5 such tiles) the TMA producer cannot run ahead of the MMA;
+ # 2 (K 128) fits 5 stages.
+ self.mma_inst_tile_k = mma_inst_tile_k
+ # Operand TMA loads EVICT_FIRST and the bulk scatter-add EVICT_LAST, so the [T, H]
+ # output the reductions hit is kept in L2 over the streamed operands.
+ self.l2_hints = l2_hints
+ # Each epilogue thread writes only its own row of the C staging buffer and bulk-reduces
+ # only that row, so it can wait for its own previous bulk read right before overwriting
+ # the row instead of a per-tile bulk wait + CTA barrier.
+ self.row_owned_epilogue = row_owned_epilogue
+ self.use_expert_counts = use_expert_counts
+ self.num_local_experts = num_local_experts
+ self.num_tile_info_fields = 6 if self.use_expert_counts else 5
+ if self.use_expert_counts:
+ assert self.num_local_experts > 0
+ # K dimension is deferred in _setup_attributes
+ self.mma_tiler = (*mma_tiler_mn, 1)
+
+ self.cta_group = tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE
+
+ self.occupancy = 1
+ self.epilog_warp_id = (0, 1, 2, 3)
+ self.mma_warp_id = 4
+ self.tma_warp_id = 5
+ self.sched_warp_id = 6
+ self.meta_load_warp_id = 7
+ self.threads_per_warp = 32
+ self.threads_per_cta = self.threads_per_warp * len(
+ (
+ *self.epilog_warp_id,
+ self.mma_warp_id,
+ self.tma_warp_id,
+ self.sched_warp_id,
+ self.meta_load_warp_id,
+ )
+ )
+ self.threads_wo_sched = self.threads_per_warp * len(
+ (
+ *self.epilog_warp_id,
+ self.mma_warp_id,
+ self.tma_warp_id,
+ self.meta_load_warp_id,
+ )
+ )
+
+ # Set barrier for cta sync, epilogue sync and tmem ptr sync
+ self.cta_sync_barrier = pipeline.NamedBarrier(
+ barrier_id=1,
+ num_threads=self.threads_per_cta,
+ )
+ self.epilog_sync_barrier = pipeline.NamedBarrier(
+ barrier_id=2,
+ num_threads=32 * len(self.epilog_warp_id),
+ )
+ self.tmem_alloc_barrier = pipeline.NamedBarrier(
+ barrier_id=3,
+ num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)),
+ )
+ self.sched_sync_barrier = pipeline.NamedBarrier(
+ barrier_id=4,
+ num_threads=self.threads_per_warp,
+ )
+ self.num_smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
+ # TMEM offset for final accumulator
+ self.tmem_final_offset = 384
+
+ def _setup_attributes(self):
+ """Set up configurations that are dependent on GEMM inputs
+
+ This method configures various attributes based on the input tensor properties
+ (data types, leading dimensions) and kernel settings:
+ - Configuring tiled MMA
+ - Computing MMA/cluster/tile shapes
+ - Computing cluster layout
+ - Computing multicast CTAs for A/B
+ - Computing epilogue subtile
+ - Setting up A/B/C stage counts in shared memory
+ - Computing A/B/C shared memory layout
+ - Computing tensor memory allocation columns
+ """
+ self.mma_inst_shape_mn = (
+ self.mma_tiler[0],
+ self.mma_tiler[1],
+ )
+ # (CTA_Tile_Shape_M, Round_Up(MMA_Tile_Shape_N, 128), MMA_Inst_Shape_K)
+ self.mma_inst_shape_mn_sfb = (
+ self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1),
+ cute.round_up(self.mma_inst_shape_mn[1], 128),
+ )
+
+ # Configure tiled mma
+ tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
+ self.a_dtype,
+ self.a_major_mode,
+ self.b_major_mode,
+ self.sf_dtype,
+ self.sf_vec_size,
+ self.cta_group,
+ self.mma_inst_shape_mn,
+ )
+
+ tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
+ self.a_dtype,
+ self.a_major_mode,
+ self.b_major_mode,
+ self.sf_dtype,
+ self.sf_vec_size,
+ cute.nvgpu.tcgen05.CtaGroup.ONE,
+ self.mma_inst_shape_mn_sfb,
+ )
+
+ # Compute mma/cluster/tile shapes
+ mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
+ mma_inst_tile_k = self.mma_inst_tile_k
+ self.mma_tiler = (
+ self.mma_tiler[0],
+ self.mma_tiler[1],
+ mma_inst_shape_k * mma_inst_tile_k,
+ )
+
+ self.mma_tiler_sfb = (
+ self.mma_inst_shape_mn_sfb[0],
+ self.mma_inst_shape_mn_sfb[1],
+ mma_inst_shape_k * mma_inst_tile_k,
+ )
+
+ self.cta_tile_shape_mnk = (
+ self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
+ self.mma_tiler[1],
+ self.mma_tiler[2],
+ )
+
+ self.cta_tile_shape_mnk_sfb = (
+ self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape),
+ self.mma_tiler_sfb[1],
+ self.mma_tiler_sfb[2],
+ )
+
+ # Compute cluster layout
+ self.cluster_layout_vmnk = cute.tiled_divide(
+ cute.make_layout((*self.cluster_shape_mn, 1)),
+ (tiled_mma.thr_id.shape,),
+ )
+
+ self.cluster_layout_sfb_vmnk = cute.tiled_divide(
+ cute.make_layout((*self.cluster_shape_mn, 1)),
+ (tiled_mma_sfb.thr_id.shape,),
+ )
+
+ # Compute number of multicast CTAs for A/B
+ self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2])
+ self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1])
+ self.is_a_mcast = self.num_mcast_ctas_a > 1
+ self.is_b_mcast = self.num_mcast_ctas_b > 1
+
+ # Compute epilogue subtile
+ self.epi_tile = sm100_utils.compute_epilogue_tile_shape(
+ self.cta_tile_shape_mnk,
+ self.use_2cta_instrs,
+ self.gemm_output_layout,
+ self.out_dtype,
+ )
+
+ self.epi_tile_n = cute.size(self.epi_tile[1])
+
+ # Setup A/B/C/Scale stage count in shared memory and ACC stage count in tensor memory
+ (
+ self.num_acc_stage,
+ self.num_ab_stage,
+ self.num_c_stage,
+ self.num_tile_stage,
+ self.num_meta_stage,
+ ) = self._compute_stages(
+ tiled_mma,
+ self.mma_tiler,
+ self.a_dtype,
+ self.b_dtype,
+ self.out_dtype,
+ self.cta_tile_shape_mnk,
+ self.sf_dtype,
+ self.sf_vec_size,
+ self.final_scale_dtype,
+ self.num_smem_capacity,
+ self.occupancy,
+ )
+
+ # Compute A/B/C/Scale shared memory layout
+ self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
+ tiled_mma,
+ self.mma_tiler,
+ self.a_dtype,
+ self.num_ab_stage,
+ )
+ self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
+ tiled_mma,
+ self.mma_tiler,
+ self.b_dtype,
+ self.num_ab_stage,
+ )
+ self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
+ tiled_mma,
+ self.mma_tiler,
+ self.sf_vec_size,
+ self.num_ab_stage,
+ )
+ self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(
+ tiled_mma,
+ self.mma_tiler,
+ self.sf_vec_size,
+ self.num_ab_stage,
+ )
+
+ swizzled_pad = 16 // (self.out_dtype.width // 8)
+ self.c_smem_layout_staged = cute.make_layout(
+ (self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[1], self.num_c_stage),
+ stride=(
+ self.cta_tile_shape_mnk[1] + swizzled_pad,
+ 1,
+ self.cta_tile_shape_mnk[0] * (self.cta_tile_shape_mnk[1] + 8),
+ ),
+ )
+
+ # Overlap and double buffer accumulator when num_acc_stage == 1 for cta_tile_n = 256 case
+ self.overlapping_accum = self.num_acc_stage == 1
+
+ sf_atom_mn = 32
+
+ self.num_sfa_tmem_cols = (self.cta_tile_shape_mnk[0] // sf_atom_mn) * mma_inst_tile_k
+ self.num_sfb_tmem_cols = (self.cta_tile_shape_mnk_sfb[1] // sf_atom_mn) * mma_inst_tile_k
+ self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols
+ self.num_accumulator_tmem_cols = (
+ self.cta_tile_shape_mnk[1] * self.num_acc_stage
+ if not self.overlapping_accum
+ else self.cta_tile_shape_mnk[1] * 2 - self.num_sf_tmem_cols
+ )
+
+ self.iter_acc_early_release_in_epilogue = self.num_sf_tmem_cols // self.epi_tile_n
+
+ # Compute the number of tensor memory allocation columns
+ self.num_tmem_alloc_cols = 512
+
+ @cute.jit
+ def _expert_count_tile_info(
+ self,
+ expert_counts: cute.Tensor,
+ tile_idx: cutlass.Int32,
+ expert_capacity: cutlass.Int32,
+ ):
+ """Resolve a compact tile directly from expert-major receive counts."""
+ tile_size = cutlass.Int32(self.mma_tiler[0])
+ num_valid_tiles = cutlass.Int32(0)
+ expert_idx = cutlass.Int32(-1)
+ expert_tile_start = cutlass.Int32(0)
+ expert_count = cutlass.Int32(0)
+ for candidate_expert in cutlass.range_constexpr(self.num_local_experts):
+ count = cutlass.min(
+ cutlass.max(expert_counts[candidate_expert], cutlass.Int32(0)),
+ expert_capacity,
+ )
+ num_expert_tiles = (count + tile_size - 1) // tile_size
+ is_target = (tile_idx >= num_valid_tiles) and (
+ tile_idx < num_valid_tiles + num_expert_tiles
+ )
+ if is_target:
+ expert_idx = cutlass.Int32(candidate_expert)
+ expert_tile_start = num_valid_tiles
+ expert_count = count
+ num_valid_tiles += num_expert_tiles
+
+ row_in_expert = (tile_idx - expert_tile_start) * tile_size
+ rows_in_tile = cutlass.min(
+ tile_size,
+ cutlass.max(expert_count - row_in_expert, cutlass.Int32(0)),
+ )
+ mn_limit = tile_idx * tile_size + rows_in_tile
+ expanded_row_start = expert_idx * expert_capacity + row_in_expert
+ return num_valid_tiles, expert_idx, mn_limit, expanded_row_start
+
+ @cute.jit
+ def __call__(
+ self,
+ a: cute.Tensor,
+ b: Union[cute.Tensor, Tuple[cute.Tensor, ...]],
+ out: cute.Tensor,
+ sfa: cute.Tensor,
+ sfb: Union[cute.Tensor, Tuple[cute.Tensor, ...]],
+ tile_idx_to_expert_idx: cute.Tensor,
+ num_non_exiting_tiles: cute.Tensor,
+ tile_idx_to_mn_limit: cute.Tensor,
+ alpha: Union[cute.Tensor, Tuple[cute.Tensor, ...]],
+ expert_capacity: cutlass.Int32,
+ max_active_clusters: cutlass.Constexpr,
+ stream: cuda.CUstream,
+ permuted_idx_to_expanded_idx: cute.Tensor,
+ token_final_scales: cute.Tensor,
+ epilogue_op: cutlass.Constexpr = lambda x: x,
+ ):
+ """Execute the GEMM operation in steps:
+ - Setup static attributes before smem/grid/tma computation
+ - Setup TMA load/store atoms and tensors
+ - Compute grid size with regard to hardware constraints
+ - Define shared storage for kernel
+ - Launch the kernel synchronously
+
+ :param a: Input tensor A
+ :type a: cute.Tensor
+ :param b: Input tensor B
+ :type b: cute.Tensor
+ :param out: Finalized output tensor (shape [seq_len, n])
+ :type out: cute.Tensor
+ :param sfa: Scale factor tensor A
+ :type sfa: cute.Tensor
+ :param sfb: Scale factor tensor B
+ :type sfb: cute.Tensor
+ :param tile_idx_to_expert_idx: Mapping from tile index to expert ID,
+ shape (permuted_m/cta_tile_m,) where cta_tile_m is the CTA tile M size
+ :type tile_idx_to_expert_idx: cute.Tensor
+ :param num_non_exiting_tiles: Number of valid tiles (valid_m/cta_tile_m), shape (1,)
+ :type num_non_exiting_tiles: cute.Tensor
+ :param alpha: Alpha tensor for each group
+ :type alpha: Union[cute.Tensor, Tuple[cute.Tensor, ...]]
+ :param max_active_clusters: Maximum number of active clusters
+ :type max_active_clusters: cutlass.Constexpr
+ :param stream: CUDA stream for asynchronous execution
+ :type stream: cuda.CUstream
+ :param permuted_idx_to_expanded_idx: Mapping from permuted index to expanded index, shape (permuted_m,)
+ :type permuted_idx_to_expanded_idx: cute.Tensor
+ :param token_final_scales: Token-wise scaling factors, shape (m, topK)
+ :type token_final_scales: cute.Tensor
+ :param epilogue_op: Optional elementwise lambda function to apply to the output tensor
+ :type epilogue_op: cutlass.Constexpr
+ :raises TypeError: If input data types are incompatible with the MMA instruction.
+ """
+ # Setup static attributes before smem/grid/tma computation
+ self.a_dtype: Type[cutlass.Numeric] = a.element_type
+ # Handle tuple of B tensors
+ b_tuple = b if isinstance(b, tuple) else (b,)
+ sfb_tuple = sfb if isinstance(sfb, tuple) else (sfb,)
+ alpha_tuple = alpha if isinstance(alpha, tuple) else (alpha,)
+ self.b_dtype: Type[cutlass.Numeric] = b_tuple[0].element_type
+ self.out_dtype: Type[cutlass.Numeric] = out.element_type
+ self.sf_dtype: Type[cutlass.Numeric] = sfa.element_type
+ self.final_scale_dtype: Type[cutlass.Numeric] = token_final_scales.element_type
+ self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode()
+ self.b_major_mode = utils.LayoutEnum.from_tensor(b_tuple[0]).mma_major_mode()
+ self.gemm_output_layout = utils.LayoutEnum.ROW_MAJOR
+
+ self.topK = token_final_scales.shape[1]
+ # Check if input data types are compatible with MMA instruction
+ if cutlass.const_expr(self.a_dtype != self.b_dtype):
+ raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}")
+
+ # Setup attributes that dependent on gemm inputs
+ self._setup_attributes()
+ # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout
+ # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL)
+ sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(a.shape, self.sf_vec_size)
+ sfa = cute.make_tensor(sfa.iterator, sfa_layout)
+
+ # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL)
+ sfb_layout_0 = blockscaled_utils.tile_atom_to_shape_SF(b_tuple[0].shape, self.sf_vec_size)
+ sfb_tensor_0 = cute.make_tensor(sfb_tuple[0].iterator, sfb_layout_0)
+ sfb_tensors = [sfb_tensor_0]
+ sfb_tuple = tuple(sfb_tensors)
+ # Backward compat alias
+ sfb = sfb_tuple[0]
+
+ tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
+ self.a_dtype,
+ self.a_major_mode,
+ self.b_major_mode,
+ self.sf_dtype,
+ self.sf_vec_size,
+ self.cta_group,
+ self.mma_inst_shape_mn,
+ )
+
+ tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
+ self.a_dtype,
+ self.a_major_mode,
+ self.b_major_mode,
+ self.sf_dtype,
+ self.sf_vec_size,
+ cute.nvgpu.tcgen05.CtaGroup.ONE,
+ self.mma_inst_shape_mn_sfb,
+ )
+ atom_thr_size = cute.size(tiled_mma.thr_id.shape)
+
+ # Setup TMA load for A
+ a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id)
+ a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
+ tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
+ a_op,
+ a,
+ a_smem_layout,
+ self.mma_tiler,
+ tiled_mma,
+ self.cluster_layout_vmnk.shape,
+ )
+
+ # Setup TMA load for B
+ b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id)
+ b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
+
+ # Setup TMA load for SFA
+ sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id)
+ sfa_smem_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0))
+ tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A(
+ sfa_op,
+ sfa,
+ sfa_smem_layout,
+ self.mma_tiler,
+ tiled_mma,
+ self.cluster_layout_vmnk.shape,
+ internal_type=cutlass.Int16,
+ )
+
+ # Setup TMA load for SFB
+ sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, tiled_mma.thr_id)
+ sfb_smem_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0))
+
+ # Helper to create TMA for one B tensor
+ def _make_tma_b(b_tensor, sfb_tensor):
+ atom_b, tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
+ b_op,
+ b_tensor,
+ b_smem_layout,
+ self.mma_tiler,
+ tiled_mma,
+ self.cluster_layout_vmnk.shape,
+ )
+ atom_sfb, tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B(
+ sfb_op,
+ sfb_tensor,
+ sfb_smem_layout,
+ self.mma_tiler_sfb,
+ tiled_mma_sfb,
+ self.cluster_layout_sfb_vmnk.shape,
+ internal_type=cutlass.Int16,
+ )
+ if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192):
+ x = tensor_sfb.stride[0][1]
+ y = cute.ceil_div(tensor_sfb.shape[0][1], 4)
+
+ new_shape = (
+ (tensor_sfb.shape[0][0], ((2, 2), y)),
+ tensor_sfb.shape[1],
+ tensor_sfb.shape[2],
+ )
+ # Use right multiplication for ScaledBasis (3 * x instead of x * 3)
+ x_times_3 = 3 * x
+ new_stride = (
+ (tensor_sfb.stride[0][0], ((x, x), x_times_3)),
+ tensor_sfb.stride[1],
+ tensor_sfb.stride[2],
+ )
+ tensor_sfb_new_layout = cute.make_layout(new_shape, stride=new_stride)
+ tensor_sfb = cute.make_tensor(tensor_sfb.iterator, tensor_sfb_new_layout)
+ return atom_b, tensor_b, atom_sfb, tensor_sfb
+
+ # Create TMA for the B tensor
+ atom_b_0, tensor_b_0, atom_sfb_0, tensor_sfb_0 = _make_tma_b(b_tuple[0], sfb_tuple[0])
+ tma_atoms_b = [atom_b_0]
+ tma_tensors_b = [tensor_b_0]
+ tma_atoms_sfb = [atom_sfb_0]
+ tma_tensors_sfb = [tensor_sfb_0]
+ tma_atoms_b = tuple(tma_atoms_b)
+ tma_tensors_b = tuple(tma_tensors_b)
+ tma_atoms_sfb = tuple(tma_atoms_sfb)
+ tma_tensors_sfb = tuple(tma_tensors_sfb)
+
+ a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout)
+ b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout)
+ sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout)
+ sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout)
+ self.num_tma_load_bytes = (
+ a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size
+ ) * atom_thr_size
+
+ self.tile_sched_params, grid = self._compute_grid(
+ (a.shape[0], b_tuple[0].shape[0], a.shape[2]),
+ self.cta_tile_shape_mnk,
+ self.cluster_shape_mn,
+ max_active_clusters,
+ self.raster_along_m,
+ )
+
+ self.buffer_align_bytes = 1024
+
+ #### finalized epi layout ####
+ epi_tile_m = cute.size(self.epi_tile[0])
+ epi_tile_n = cute.size(self.epi_tile[1])
+ epi_tile_size = epi_tile_m * epi_tile_n
+ num_epilogue_threads = 32 * len(self.epilog_warp_id)
+ self.ttr_racc_size = epi_tile_size // num_epilogue_threads
+ self.copy_size = self.cta_tile_shape_mnk[1] * (self.out_dtype.width // 8)
+
+ if cutlass.const_expr(self.out_dtype == cutlass.BFloat16):
+ # 8-element vectorization for BF16
+ self.epi_layout = cute.make_layout(
+ shape=(self.ttr_racc_size // 8, 4, 2), stride=(8, 2, 1)
+ )
+ self.epi_loop_size = self.ttr_racc_size // 8
+ self.element_offset = 8
+
+ elif cutlass.const_expr(self.out_dtype == cutlass.Float32):
+ # 2-element vectorization for FP32
+ self.epi_layout = cute.make_layout(shape=(self.ttr_racc_size // 2, 2), stride=(2, 1))
+ self.epi_loop_size = self.ttr_racc_size // 2
+ self.element_offset = 2
+ else:
+ # Scalar fallback
+ self.epi_layout = cute.make_layout(shape=(self.ttr_racc_size,), stride=(1,))
+ self.epi_loop_size = self.ttr_racc_size
+ self.element_offset = 1
+
+ # Define shared storage for kernel
+ @cute.struct
+ class SharedStorage:
+ # (bidx, bidy, expert, valid, mn_limit, expanded_row_start)
+ sInfo: cute.struct.Align[
+ cute.struct.MemRange[
+ cutlass.Int32, self.num_tile_info_fields * self.num_tile_stage
+ ],
+ # 1 byte alignment
+ 1,
+ ]
+ ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]
+ acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]
+ tile_info_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_tile_stage * 2]
+ meta_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_meta_stage * 2]
+ tmem_dealloc_mbar_ptr: cutlass.Int64
+ tmem_holding_buf: cutlass.Int32
+ # (MMA, MMA_M, MMA_K, STAGE)
+ sA: cute.struct.Align[
+ cute.struct.MemRange[self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)],
+ self.buffer_align_bytes,
+ ]
+ # (MMA, MMA_N, MMA_K, STAGE)
+ sB: cute.struct.Align[
+ cute.struct.MemRange[self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)],
+ self.buffer_align_bytes,
+ ]
+ # (granularity_m, repeat_m), (granularity_k, repeat_k), num_scale_stage)
+ sSFA: cute.struct.Align[
+ cute.struct.MemRange[self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged)],
+ self.buffer_align_bytes,
+ ]
+ # (granularity_n, repeat_n), (granularity_k, repeat_k), num_scale_stage)
+ sSFB: cute.struct.Align[
+ cute.struct.MemRange[self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)],
+ self.buffer_align_bytes,
+ ]
+
+ sC: cute.struct.Align[
+ cute.struct.MemRange[self.out_dtype, cute.cosize(self.c_smem_layout_staged)],
+ self.buffer_align_bytes,
+ ]
+ meta_token_idx: cute.struct.Align[
+ cute.struct.MemRange[
+ cutlass.Int32, self.cta_tile_shape_mnk[0] * self.num_meta_stage
+ ],
+ 1,
+ ]
+ meta_scale: cute.struct.Align[
+ cute.struct.MemRange[
+ self.final_scale_dtype, self.cta_tile_shape_mnk[0] * self.num_meta_stage
+ ],
+ 1,
+ ]
+
+ self.shared_storage = SharedStorage
+
+ # Launch the kernel synchronously
+ self.kernel(
+ tiled_mma,
+ tiled_mma_sfb,
+ tma_atom_a,
+ tma_tensor_a,
+ tma_atoms_b,
+ tma_tensors_b,
+ tma_atom_sfa,
+ tma_tensor_sfa,
+ tma_atoms_sfb,
+ tma_tensors_sfb,
+ out,
+ tile_idx_to_expert_idx,
+ num_non_exiting_tiles,
+ tile_idx_to_mn_limit,
+ alpha_tuple,
+ permuted_idx_to_expanded_idx,
+ token_final_scales,
+ expert_capacity,
+ self.cluster_layout_vmnk,
+ self.cluster_layout_sfb_vmnk,
+ self.a_smem_layout_staged,
+ self.b_smem_layout_staged,
+ self.sfa_smem_layout_staged,
+ self.sfb_smem_layout_staged,
+ self.c_smem_layout_staged,
+ self.epi_tile,
+ self.epi_layout,
+ self.topK,
+ self.tile_sched_params,
+ epilogue_op,
+ ).launch(
+ grid=grid,
+ block=[self.threads_per_cta, 1, 1],
+ cluster=(*self.cluster_shape_mn, 1),
+ smem=self.shared_storage.size_in_bytes(),
+ stream=stream,
+ min_blocks_per_mp=1,
+ use_pdl=TRTLLM_ENABLE_PDL,
+ )
+ return
+
+ def mainloop_s2t_copy_and_partition(
+ self,
+ sSF: cute.Tensor,
+ tSF: cute.Tensor,
+ ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
+ """
+ Make tiledCopy for smem to tmem load for scale factor tensor, then use it to partition smem memory (source) and
+ tensor memory (destination).
+
+ :param sSF: The scale factor tensor in smem
+ :type sSF: cute.Tensor
+ :param tSF: The scale factor tensor in tmem
+ :type tSF: cute.Tensor
+
+ :return: A tuple containing (tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t) where:
+ - tiled_copy_s2t: The tiled copy operation for smem to tmem load for scale factor tensor(s2t)
+ - tCsSF_compact_s2t: The partitioned scale factor tensor in smem
+ - tSF_compact_s2t: The partitioned scale factor tensor in tmem
+ :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]
+ """
+ # (MMA, MMA_MN, MMA_K, STAGE)
+ tCsSF_compact = cute.filter_zeros(sSF)
+ # (MMA, MMA_MN, MMA_K)
+ tCtSF_compact = cute.filter_zeros(tSF)
+
+ # Make S2T CopyAtom and tiledCopy
+ copy_atom_s2t = cute.make_copy_atom(
+ tcgen05.Cp4x32x128bOp(self.cta_group),
+ self.sf_dtype,
+ )
+ tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact)
+ thr_copy_s2t = tiled_copy_s2t.get_slice(0)
+
+ # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
+ tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact)
+ # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
+ tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy_s2t, tCsSF_compact_s2t_)
+ # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
+ tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact)
+
+ return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t
+
+ @cute.kernel
+ def kernel(
+ self,
+ tiled_mma: cute.TiledMma,
+ tiled_mma_sfb: cute.TiledMma,
+ tma_atom_a: cute.CopyAtom,
+ mA_mkl: cute.Tensor,
+ tma_atoms_b: Tuple[cute.CopyAtom, ...],
+ mB_nkl_tuple: Tuple[cute.Tensor, ...],
+ tma_atom_sfa: cute.CopyAtom,
+ mSFA_mkl: cute.Tensor,
+ tma_atoms_sfb: Tuple[cute.CopyAtom, ...],
+ mSFB_nkl_tuple: Tuple[cute.Tensor, ...],
+ out: cute.Tensor,
+ tile_idx_to_expert_idx: cute.Tensor,
+ num_non_exiting_tiles: cute.Tensor,
+ tile_idx_to_mn_limit: cute.Tensor,
+ alpha_tuple: Tuple[cute.Tensor, ...],
+ permuted_idx_to_expanded_idx: cute.Tensor,
+ token_final_scales: cute.Tensor,
+ expert_capacity: cutlass.Int32,
+ cluster_layout_vmnk: cute.Layout,
+ cluster_layout_sfb_vmnk: cute.Layout,
+ a_smem_layout_staged: cute.ComposedLayout,
+ b_smem_layout_staged: cute.ComposedLayout,
+ sfa_smem_layout_staged: cute.Layout,
+ sfb_smem_layout_staged: cute.Layout,
+ c_smem_layout_staged: cute.Layout,
+ epi_tile: cute.Tile,
+ epi_layout: cute.Layout,
+ topK: cutlass.Int32,
+ tile_sched_params: utils.PersistentTileSchedulerParams,
+ epilogue_op: cutlass.Constexpr,
+ ):
+ """
+ GPU device kernel performing the Persistent batched GEMM computation.
+ """
+ warp_idx = cute.arch.warp_idx()
+ warp_idx = cute.arch.make_warp_uniform(warp_idx)
+
+ #
+ # Prefetch tma desc
+ #
+ if warp_idx == self.tma_warp_id:
+ cpasync.prefetch_descriptor(tma_atom_a)
+ cpasync.prefetch_descriptor(tma_atom_sfa)
+ cpasync.prefetch_descriptor(tma_atoms_b[0])
+ cpasync.prefetch_descriptor(tma_atoms_sfb[0])
+
+ use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2
+
+ #
+ # Setup cta/thread coordinates
+ #
+ # Coords inside cluster
+ bidx, bidy, bidz = cute.arch.block_idx()
+ mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape)
+ is_leader_cta = mma_tile_coord_v == 0
+ cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+ block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
+
+ block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord(
+ cta_rank_in_cluster
+ )
+
+ # Coord inside cta
+ tidx, _, _ = cute.arch.thread_idx()
+
+ #
+ # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier
+ #
+ smem = utils.SmemAllocator()
+ storage = smem.allocate(self.shared_storage)
+
+ # Initialize mainloop ab_pipeline (barrier) and states
+ ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+ num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
+ ab_pipeline_consumer_group = pipeline.CooperativeGroup(
+ pipeline.Agent.Thread, num_tma_producer
+ )
+ ab_pipeline = pipeline.PipelineTmaUmma.create(
+ barrier_storage=storage.ab_mbar_ptr.data_ptr(),
+ num_stages=self.num_ab_stage,
+ producer_group=ab_pipeline_producer_group,
+ consumer_group=ab_pipeline_consumer_group,
+ tx_count=self.num_tma_load_bytes,
+ cta_layout_vmnk=cluster_layout_vmnk,
+ )
+
+ # Initialize acc_pipeline (barrier) and states
+ acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+ num_acc_consumer_threads = (
+ len(self.epilog_warp_id) * self.threads_per_warp * (2 if use_2cta_instrs else 1)
+ )
+ acc_pipeline_consumer_group = pipeline.CooperativeGroup(
+ pipeline.Agent.Thread, num_acc_consumer_threads
+ )
+ acc_pipeline = pipeline.PipelineUmmaAsync.create(
+ barrier_storage=storage.acc_mbar_ptr.data_ptr(),
+ num_stages=self.num_acc_stage,
+ producer_group=acc_pipeline_producer_group,
+ consumer_group=acc_pipeline_consumer_group,
+ cta_layout_vmnk=cluster_layout_vmnk,
+ )
+
+ # Initialize tile info pipeline (barrier) and states
+ tile_info_pipeline_producer_group = pipeline.CooperativeGroup(
+ pipeline.Agent.Thread,
+ self.threads_per_warp * 1,
+ )
+ tile_info_pipeline_consumer_group = pipeline.CooperativeGroup(
+ pipeline.Agent.Thread,
+ self.threads_wo_sched,
+ )
+ tile_info_pipeline = pipeline.PipelineAsync.create(
+ barrier_storage=storage.tile_info_mbar_ptr.data_ptr(),
+ num_stages=self.num_tile_stage,
+ producer_group=tile_info_pipeline_producer_group,
+ consumer_group=tile_info_pipeline_consumer_group,
+ )
+
+ # Initialize metadata pipeline (meta loader warp -> epilogue warps)
+ meta_pipeline_producer_group = pipeline.CooperativeGroup(
+ pipeline.Agent.Thread,
+ self.threads_per_warp * 1,
+ )
+ meta_pipeline_consumer_group = pipeline.CooperativeGroup(
+ pipeline.Agent.Thread,
+ self.threads_per_warp * len(self.epilog_warp_id),
+ )
+ meta_pipeline = pipeline.PipelineAsync.create(
+ barrier_storage=storage.meta_mbar_ptr.data_ptr(),
+ num_stages=self.num_meta_stage,
+ producer_group=meta_pipeline_producer_group,
+ consumer_group=meta_pipeline_consumer_group,
+ )
+
+ # Tensor memory dealloc barrier init
+ tmem = utils.TmemAllocator(
+ storage.tmem_holding_buf,
+ barrier_for_retrieve=self.tmem_alloc_barrier,
+ allocator_warp_id=self.epilog_warp_id[0],
+ is_two_cta=use_2cta_instrs,
+ two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
+ )
+
+ # Cluster arrive after barrier init
+ if cute.size(self.cluster_shape_mn) > 1:
+ cute.arch.cluster_arrive_relaxed()
+
+ #
+ # Setup smem tensor A/B/C/Scale/ExpandedIdx
+ #
+ # (MMA, MMA_M, MMA_K, STAGE)
+ sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner)
+ # (MMA, MMA_N, MMA_K, STAGE)
+ sB = storage.sB.get_tensor(b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner)
+ # (granularity_m, repeat_m), (granularity_k, repeat_k), num_scale_stage)
+ sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged)
+ # (granularity_n, repeat_n), (granularity_k, repeat_k), num_scale_stage)
+ sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged)
+
+ sC = storage.sC.get_tensor(c_smem_layout_staged)
+
+ # (bidx, bidy, expert, valid, mn_limit, expanded_row_start)
+ info_layout = cute.make_layout(
+ (self.num_tile_info_fields, self.num_tile_stage),
+ stride=(1, self.num_tile_info_fields),
+ )
+ sInfo = storage.sInfo.get_tensor(info_layout)
+
+ # Per-row finalize metadata staged by the meta loader warp: (row, stage)
+ meta_layout = cute.make_layout(
+ (self.cta_tile_shape_mnk[0], self.num_meta_stage),
+ stride=(1, self.cta_tile_shape_mnk[0]),
+ )
+ sMetaTokenIdx = storage.meta_token_idx.get_tensor(meta_layout)
+ sMetaScale = storage.meta_scale.get_tensor(meta_layout)
+
+ #
+ # Compute multicast mask for A/B buffer full
+ #
+ a_full_mcast_mask = None
+ b_full_mcast_mask = None
+ sfa_full_mcast_mask = None
+ sfb_full_mcast_mask = None
+ if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs):
+ a_full_mcast_mask = cpasync.create_tma_multicast_mask(
+ cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
+ )
+ b_full_mcast_mask = cpasync.create_tma_multicast_mask(
+ cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1
+ )
+ sfa_full_mcast_mask = cpasync.create_tma_multicast_mask(
+ cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
+ )
+ sfb_full_mcast_mask = cpasync.create_tma_multicast_mask(
+ cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_vmnk, mcast_mode=1
+ )
+
+ #
+ # Local_tile partition global tensors
+ #
+ # (bM, bK, loopM, loopK, loopL)
+ gA_mkl = cute.local_tile(
+ mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
+ )
+ # (bN, bK, loopN, loopK, loopL)
+ gB_nkl_0 = cute.local_tile(
+ mB_nkl_tuple[0],
+ cute.slice_(self.mma_tiler, (0, None, None)),
+ (None, None, None),
+ )
+
+ # (bM, bK, RestM, RestK, RestL)
+ gSFA_mkl = cute.local_tile(
+ mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
+ )
+
+ # (bN, bK, RestN, RestK, RestL)
+ gSFB_nkl_0 = cute.local_tile(
+ mSFB_nkl_tuple[0],
+ cute.slice_(self.mma_tiler_sfb, (0, None, None)),
+ (None, None, None),
+ )
+
+ k_tile_cnt = cutlass.Int32(cute.size(gA_mkl, mode=[3]))
+
+ #
+ # Partition global tensor for TiledMMA_A/B
+ #
+ thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
+ thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v)
+ # (MMA, MMA_M, MMA_K, loopM, loopK, loopL)
+ tCgA = thr_mma.partition_A(gA_mkl)
+ # (MMA, MMA_N, MMA_K, loopN, loopK, loopL)
+ tCgB_0 = thr_mma.partition_B(gB_nkl_0)
+ # (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
+ tCgSFA = thr_mma.partition_A(gSFA_mkl)
+ # (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
+ tCgSFB_0 = thr_mma_sfb.partition_B(gSFB_nkl_0)
+
+ #
+ # Partition global/shared tensor for TMA load A/B
+ #
+ # TMA load A partition_S/D
+ a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape)
+ # ((atom_v, rest_v), STAGE)
+ # ((atom_v, rest_v), loopM, loopK, loopL)
+ tAsA, tAgA = cpasync.tma_partition(
+ tma_atom_a,
+ block_in_cluster_coord_vmnk[2],
+ a_cta_layout,
+ cute.group_modes(sA, 0, 3),
+ cute.group_modes(tCgA, 0, 3),
+ )
+ # TMA load B partition_S/D
+ b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape)
+ # ((atom_v, rest_v), STAGE)
+ # ((atom_v, rest_v), loopM, loopK, loopL)
+ tBsB_0, tBgB_0 = cpasync.tma_partition(
+ tma_atoms_b[0],
+ block_in_cluster_coord_vmnk[1],
+ b_cta_layout,
+ cute.group_modes(sB, 0, 3),
+ cute.group_modes(tCgB_0, 0, 3),
+ )
+
+ # TMA load SFA partition_S/D
+ sfa_cta_layout = a_cta_layout
+ # ((atom_v, rest_v), STAGE)
+ # ((atom_v, rest_v), RestM, RestK, RestL)
+
+ tAsSFA, tAgSFA = cute.nvgpu.cpasync.tma_partition(
+ tma_atom_sfa,
+ block_in_cluster_coord_vmnk[2],
+ sfa_cta_layout,
+ cute.group_modes(sSFA, 0, 3),
+ cute.group_modes(tCgSFA, 0, 3),
+ )
+
+ tAsSFA = cute.filter_zeros(tAsSFA)
+ tAgSFA = cute.filter_zeros(tAgSFA)
+
+ # TMA load SFB partition_S/D
+ sfb_cta_layout = cute.make_layout(
+ cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape
+ )
+ # ((atom_v, rest_v), STAGE)
+ # ((atom_v, rest_v), RestN, RestK, RestL)
+ tBsSFB_0, tBgSFB_0 = cute.nvgpu.cpasync.tma_partition(
+ tma_atoms_sfb[0],
+ block_in_cluster_coord_sfb_vmnk[1],
+ sfb_cta_layout,
+ cute.group_modes(sSFB, 0, 3),
+ cute.group_modes(tCgSFB_0, 0, 3),
+ )
+ tBsSFB_0 = cute.filter_zeros(tBsSFB_0)
+ tBgSFB_0 = cute.filter_zeros(tBgSFB_0)
+
+ #
+ # Partition shared/tensor memory tensor for TiledMMA_A/B/C
+ #
+ # (MMA, MMA_M, MMA_K, STAGE)
+ tCrA = tiled_mma.make_fragment_A(sA)
+ # (MMA, MMA_N, MMA_K, STAGE)
+ tCrB = tiled_mma.make_fragment_B(sB)
+ # (MMA, MMA_M, MMA_N)
+ acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
+
+ if cutlass.const_expr(self.overlapping_accum):
+ num_acc_stage_overlapped = 2
+ tCtAcc_fake = tiled_mma.make_fragment_C(
+ cute.append(acc_shape, num_acc_stage_overlapped)
+ )
+ # (MMA, MMA_M, MMA_N, STAGE)
+ tCtAcc_fake = cute.make_tensor(
+ tCtAcc_fake.iterator,
+ cute.make_layout(
+ tCtAcc_fake.shape,
+ stride=(
+ tCtAcc_fake.stride[0],
+ tCtAcc_fake.stride[1],
+ tCtAcc_fake.stride[2],
+ (256 - self.num_sf_tmem_cols) * tCtAcc_fake.stride[0][1],
+ ),
+ ),
+ )
+ else:
+ # (MMA, MMA_M, MMA_N, STAGE)
+ tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage))
+
+ gC_mnl = cute.local_tile(
+ out, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
+ )
+
+ # (MMA, MMA_M, MMA_N, loopM, loopN, loopL)
+ tCgC = thr_mma.partition_C(gC_mnl)
+
+ #
+ # Cluster wait before tensor memory alloc
+ #
+ if cute.size(self.cluster_shape_mn) > 1:
+ cute.arch.cluster_wait()
+ else:
+ self.cta_sync_barrier.arrive_and_wait()
+
+ griddepcontrol_wait()
+
+ #
+ # Specialized Schedule warp
+ #
+ if warp_idx == self.sched_warp_id:
+ #
+ # Persistent tile scheduling loop
+ #
+ tile_sched = utils.StaticPersistentTileScheduler.create(
+ tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
+ )
+ # First tile
+ work_tile = tile_sched.initial_work_tile_info()
+
+ tile_info_producer_state = pipeline.make_pipeline_state(
+ pipeline.PipelineUserType.Producer, self.num_tile_stage
+ )
+
+ if cutlass.const_expr(self.use_expert_counts):
+ num_valid_tiles, _, _, _ = self._expert_count_tile_info(
+ tile_idx_to_expert_idx, cutlass.Int32(0), expert_capacity
+ )
+ else:
+ num_valid_tiles = num_non_exiting_tiles[0]
+
+ if cutlass.const_expr(self.raster_along_m):
+ while work_tile.is_valid_tile:
+ cur_tile_coord = work_tile.tile_idx
+ mma_tile_coord_m = cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape)
+ tile_idx = mma_tile_coord_m
+ if tile_idx < num_valid_tiles:
+ tile_info_pipeline.producer_acquire(tile_info_producer_state)
+ expanded_row_start = cutlass.Int32(-1)
+ if cutlass.const_expr(self.use_expert_counts):
+ _, expert_idx, mn_limit, expanded_mma_row_start = (
+ self._expert_count_tile_info(
+ tile_idx_to_expert_idx, mma_tile_coord_m, expert_capacity
+ )
+ )
+ expanded_row_start = (
+ expanded_mma_row_start
+ + mma_tile_coord_v * self.cta_tile_shape_mnk[0]
+ )
+ else:
+ expert_idx = tile_idx_to_expert_idx[mma_tile_coord_m]
+ mn_limit = tile_idx_to_mn_limit[tile_idx]
+ with cute.arch.elect_one():
+ sInfo[(0, tile_info_producer_state.index)] = cur_tile_coord[0]
+ sInfo[(1, tile_info_producer_state.index)] = cur_tile_coord[1]
+ sInfo[(2, tile_info_producer_state.index)] = expert_idx
+ sInfo[(3, tile_info_producer_state.index)] = cutlass.Int32(
+ work_tile.is_valid_tile
+ )
+ sInfo[(4, tile_info_producer_state.index)] = mn_limit
+ if cutlass.const_expr(self.use_expert_counts):
+ sInfo[(5, tile_info_producer_state.index)] = expanded_row_start
+ # fence view async shared
+ cute.arch.fence_proxy(
+ "async.shared",
+ space="cta",
+ )
+
+ self.sched_sync_barrier.arrive_and_wait()
+ tile_info_pipeline.producer_commit(tile_info_producer_state)
+ tile_info_producer_state.advance()
+
+ tile_sched.advance_to_next_work()
+ work_tile = tile_sched.get_current_work()
+ else:
+ is_continue = cutlass.Boolean(1)
+ while work_tile.is_valid_tile and is_continue:
+ cur_tile_coord = work_tile.tile_idx
+ mma_tile_coord_m = cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape)
+ tile_idx = mma_tile_coord_m
+ if tile_idx < num_valid_tiles:
+ tile_info_pipeline.producer_acquire(tile_info_producer_state)
+ expanded_row_start = cutlass.Int32(-1)
+ if cutlass.const_expr(self.use_expert_counts):
+ _, expert_idx, mn_limit, expanded_mma_row_start = (
+ self._expert_count_tile_info(
+ tile_idx_to_expert_idx, mma_tile_coord_m, expert_capacity
+ )
+ )
+ expanded_row_start = (
+ expanded_mma_row_start
+ + mma_tile_coord_v * self.cta_tile_shape_mnk[0]
+ )
+ else:
+ expert_idx = tile_idx_to_expert_idx[mma_tile_coord_m]
+ mn_limit = tile_idx_to_mn_limit[tile_idx]
+ with cute.arch.elect_one():
+ sInfo[(0, tile_info_producer_state.index)] = cur_tile_coord[0]
+ sInfo[(1, tile_info_producer_state.index)] = cur_tile_coord[1]
+ sInfo[(2, tile_info_producer_state.index)] = expert_idx
+ sInfo[(3, tile_info_producer_state.index)] = cutlass.Int32(
+ work_tile.is_valid_tile
+ )
+ sInfo[(4, tile_info_producer_state.index)] = mn_limit
+ if cutlass.const_expr(self.use_expert_counts):
+ sInfo[(5, tile_info_producer_state.index)] = expanded_row_start
+ # fence view async shared
+ cute.arch.fence_proxy(
+ "async.shared",
+ space="cta",
+ )
+
+ self.sched_sync_barrier.arrive_and_wait()
+ tile_info_pipeline.producer_commit(tile_info_producer_state)
+ tile_info_producer_state.advance()
+
+ else:
+ is_continue = cutlass.Boolean(0)
+
+ tile_sched.advance_to_next_work()
+ work_tile = tile_sched.get_current_work()
+
+ tile_info_pipeline.producer_acquire(tile_info_producer_state)
+ with cute.arch.elect_one():
+ sInfo[(0, tile_info_producer_state.index)] = work_tile.tile_idx[0]
+ sInfo[(1, tile_info_producer_state.index)] = work_tile.tile_idx[1]
+ sInfo[(2, tile_info_producer_state.index)] = -1
+ sInfo[(3, tile_info_producer_state.index)] = cutlass.Int32(0)
+ sInfo[(4, tile_info_producer_state.index)] = cutlass.Int32(0)
+ if cutlass.const_expr(self.use_expert_counts):
+ sInfo[(5, tile_info_producer_state.index)] = cutlass.Int32(-1)
+ cute.arch.fence_proxy(
+ "async.shared",
+ space="cta",
+ )
+ self.sched_sync_barrier.arrive_and_wait()
+ tile_info_pipeline.producer_commit(tile_info_producer_state)
+ tile_info_producer_state.advance()
+ tile_info_pipeline.producer_tail(tile_info_producer_state)
+
+ #
+ # Specialized TMA load warp
+ #
+ if warp_idx == self.tma_warp_id:
+ ab_producer_state = pipeline.make_pipeline_state(
+ pipeline.PipelineUserType.Producer, self.num_ab_stage
+ )
+
+ tile_info_consumer_state = pipeline.make_pipeline_state(
+ pipeline.PipelineUserType.Consumer, self.num_tile_stage
+ )
+
+ # Get the first tile info from pipeline (scheduler has filtered out tiles >= num_non_exiting_tiles)
+ tile_info = cute.make_rmem_tensor((self.num_tile_info_fields,), cutlass.Int32)
+ tile_info_pipeline.consumer_wait(tile_info_consumer_state)
+ tile_info[0] = sInfo[(0, tile_info_consumer_state.index)]
+ tile_info[1] = sInfo[(1, tile_info_consumer_state.index)]
+ tile_info[2] = sInfo[(2, tile_info_consumer_state.index)]
+ tile_info[3] = sInfo[(3, tile_info_consumer_state.index)]
+ is_valid_tile = tile_info[3] == 1
+ cute.arch.fence_proxy(
+ "async.shared",
+ space="cta",
+ )
+ tile_info_pipeline.consumer_release(tile_info_consumer_state)
+ tile_info_consumer_state.advance()
+
+ while is_valid_tile:
+ mma_tile_coord_mnl = (
+ tile_info[0] // cute.size(tiled_mma.thr_id.shape),
+ tile_info[1],
+ tile_info[2],
+ )
+ #
+ # Slice to per mma tile index
+ #
+ # ((atom_v, rest_v), loopK)
+ tAgA_slice = tAgA[(None, mma_tile_coord_mnl[0], None, 0)]
+
+ # ((atom_v, rest_v), RestK)
+ tAgSFA_slice = tAgSFA[(None, mma_tile_coord_mnl[0], None, 0)]
+
+ expert_idx = mma_tile_coord_mnl[2]
+ slice_n = mma_tile_coord_mnl[1]
+ if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
+ slice_n = mma_tile_coord_mnl[1] // 2
+
+ # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt
+ ab_producer_state.reset_count()
+ peek_ab_empty_status = cutlass.Boolean(1)
+ if ab_producer_state.count < k_tile_cnt:
+ peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state)
+ #
+ # Tma load loop
+ #
+ for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):
+ tAgA_k = tAgA_slice[(None, ab_producer_state.count)]
+ tAgSFA_k = tAgSFA_slice[(None, ab_producer_state.count)]
+ tAsA_pipe = tAsA[(None, ab_producer_state.index)]
+ tBsB_pipe = tBsB_0[(None, ab_producer_state.index)]
+ tAsSFA_pipe = tAsSFA[(None, ab_producer_state.index)]
+ tBsSFB_pipe = tBsSFB_0[(None, ab_producer_state.index)]
+
+ tma_bar = ab_pipeline.producer_get_barrier(ab_producer_state)
+
+ # Conditionally wait for AB buffer empty
+ ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status)
+
+ # TMA load A/B
+ tBgB_slice = tBgB_0[(None, mma_tile_coord_mnl[1], None, expert_idx)]
+ tBgSFB_slice = tBgSFB_0[(None, slice_n, None, expert_idx)]
+ if cutlass.const_expr(self.l2_hints):
+ cute.copy(
+ tma_atom_a,
+ tAgA_k,
+ tAsA_pipe,
+ tma_bar_ptr=tma_bar,
+ mcast_mask=a_full_mcast_mask,
+ cache_policy=cutlass.Int64(_EVICT_FIRST),
+ )
+ cute.copy(
+ tma_atom_sfa,
+ tAgSFA_k,
+ tAsSFA_pipe,
+ tma_bar_ptr=tma_bar,
+ mcast_mask=sfa_full_mcast_mask,
+ cache_policy=cutlass.Int64(_EVICT_FIRST),
+ )
+ cute.copy(
+ tma_atoms_b[0],
+ tBgB_slice[(None, ab_producer_state.count)],
+ tBsB_pipe,
+ tma_bar_ptr=tma_bar,
+ mcast_mask=b_full_mcast_mask,
+ cache_policy=cutlass.Int64(_EVICT_FIRST),
+ )
+ cute.copy(
+ tma_atoms_sfb[0],
+ tBgSFB_slice[(None, ab_producer_state.count)],
+ tBsSFB_pipe,
+ tma_bar_ptr=tma_bar,
+ mcast_mask=sfb_full_mcast_mask,
+ cache_policy=cutlass.Int64(_EVICT_FIRST),
+ )
+ else:
+ cute.copy(
+ tma_atom_a,
+ tAgA_k,
+ tAsA_pipe,
+ tma_bar_ptr=tma_bar,
+ mcast_mask=a_full_mcast_mask,
+ )
+ cute.copy(
+ tma_atom_sfa,
+ tAgSFA_k,
+ tAsSFA_pipe,
+ tma_bar_ptr=tma_bar,
+ mcast_mask=sfa_full_mcast_mask,
+ )
+ cute.copy(
+ tma_atoms_b[0],
+ tBgB_slice[(None, ab_producer_state.count)],
+ tBsB_pipe,
+ tma_bar_ptr=tma_bar,
+ mcast_mask=b_full_mcast_mask,
+ )
+ cute.copy(
+ tma_atoms_sfb[0],
+ tBgSFB_slice[(None, ab_producer_state.count)],
+ tBsSFB_pipe,
+ tma_bar_ptr=tma_bar,
+ mcast_mask=sfb_full_mcast_mask,
+ )
+
+ # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1
+ ab_producer_state.advance()
+ peek_ab_empty_status = cutlass.Boolean(1)
+ if ab_producer_state.count < k_tile_cnt:
+ peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state)
+
+ #
+ # Advance to next tile
+ #
+ tile_info_pipeline.consumer_wait(tile_info_consumer_state)
+ tile_info[0] = sInfo[(0, tile_info_consumer_state.index)]
+ tile_info[1] = sInfo[(1, tile_info_consumer_state.index)]
+ tile_info[2] = sInfo[(2, tile_info_consumer_state.index)]
+ tile_info[3] = sInfo[(3, tile_info_consumer_state.index)]
+ is_valid_tile = tile_info[3] == 1
+ cute.arch.fence_proxy(
+ "async.shared",
+ space="cta",
+ )
+ tile_info_pipeline.consumer_release(tile_info_consumer_state)
+ tile_info_consumer_state.advance()
+ #
+ # Wait A/B buffer empty
+ #
+ ab_pipeline.producer_tail(ab_producer_state)
+
+ #
+ # Specialized MMA warp
+ #
+ if warp_idx == self.mma_warp_id:
+ #
+ # Bar sync for retrieve tensor memory ptr from shared mem
+ #
+ tmem.wait_for_alloc()
+
+ #
+ # Retrieving tensor memory ptr and make accumulator tensor
+ #
+ acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+ # (MMA, MMA_M, MMA_N, STAGE)
+ tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
+
+ # Make SFA tmem tensor
+ sfa_tmem_ptr = cute.recast_ptr(
+ acc_tmem_ptr + self.num_accumulator_tmem_cols,
+ dtype=self.sf_dtype,
+ )
+ # (MMA, MMA_M, MMA_K)
+ tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
+ tiled_mma,
+ self.mma_tiler,
+ self.sf_vec_size,
+ cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)),
+ )
+ tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout)
+
+ # Make SFB tmem tensor
+ # Make SFB tmem tensor
+ sfb_tmem_ptr = cute.recast_ptr(
+ acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols,
+ dtype=self.sf_dtype,
+ )
+ # (MMA, MMA_N, MMA_K)
+ tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
+ tiled_mma,
+ self.mma_tiler,
+ self.sf_vec_size,
+ cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)),
+ )
+ tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout)
+
+ # Partition for S2T copy of SFA/SFB
+ #
+ (
+ tiled_copy_s2t_sfa,
+ tCsSFA_compact_s2t,
+ tCtSFA_compact_s2t,
+ ) = self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA)
+ (
+ tiled_copy_s2t_sfb,
+ tCsSFB_compact_s2t,
+ tCtSFB_compact_s2t,
+ ) = self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB)
+
+ ab_consumer_state = pipeline.make_pipeline_state(
+ pipeline.PipelineUserType.Consumer, self.num_ab_stage
+ )
+ acc_producer_state = pipeline.make_pipeline_state(
+ pipeline.PipelineUserType.Producer, self.num_acc_stage
+ )
+
+ tile_info_consumer_state = pipeline.make_pipeline_state(
+ pipeline.PipelineUserType.Consumer, self.num_tile_stage
+ )
+
+ # Get the first tile info from pipeline (scheduler has filtered out tiles >= num_non_exiting_tiles)
+ tile_info = cute.make_rmem_tensor((5,), cutlass.Int32)
+ tile_info_pipeline.consumer_wait(tile_info_consumer_state)
+ tile_info[0] = sInfo[(0, tile_info_consumer_state.index)]
+ tile_info[1] = sInfo[(1, tile_info_consumer_state.index)]
+ tile_info[2] = sInfo[(2, tile_info_consumer_state.index)]
+ tile_info[3] = sInfo[(3, tile_info_consumer_state.index)]
+ is_valid_tile = tile_info[3] == 1
+ cute.arch.fence_proxy(
+ "async.shared",
+ space="cta",
+ )
+ tile_info_pipeline.consumer_release(tile_info_consumer_state)
+ tile_info_consumer_state.advance()
+
+ while is_valid_tile:
+ # Peek (try_wait) AB buffer full for k_tile = 0
+ ab_consumer_state.reset_count()
+ peek_ab_full_status = cutlass.Boolean(1)
+ if ab_consumer_state.count < k_tile_cnt and is_leader_cta:
+ peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state)
+
+ mma_tile_coord_mnl = (
+ tile_info[0] // cute.size(tiled_mma.thr_id.shape),
+ tile_info[1],
+ tile_info[2],
+ )
+
+ # Get accumulator stage index
+ if cutlass.const_expr(self.overlapping_accum):
+ acc_stage_index = acc_producer_state.phase ^ 1
+ else:
+ acc_stage_index = acc_producer_state.index
+
+ tCtAcc = tCtAcc_base[(None, None, None, acc_stage_index)]
+
+ tCtSFB_mma = tCtSFB
+ if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192):
+ # If this is an ODD tile, shift the TMEM start address for cta_tile_shape_n=192 case by two words
+ # (ignores first 64 columns of SFB)
+ offset = (
+ cutlass.Int32(2) if mma_tile_coord_mnl[1] % 2 == 1 else cutlass.Int32(0)
+ )
+ shifted_ptr = cute.recast_ptr(
+ acc_tmem_ptr
+ + self.num_accumulator_tmem_cols
+ + self.num_sfa_tmem_cols
+ + offset,
+ dtype=self.sf_dtype,
+ )
+ tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout)
+ elif cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
+ # Move in increments of 64 columns of SFB
+ offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2)
+ shifted_ptr = cute.recast_ptr(
+ acc_tmem_ptr
+ + self.num_accumulator_tmem_cols
+ + self.num_sfa_tmem_cols
+ + offset,
+ dtype=self.sf_dtype,
+ )
+ tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout)
+ #
+ # Wait for accumulator buffer empty
+ #
+ if is_leader_cta:
+ acc_pipeline.producer_acquire(acc_producer_state)
+ #
+ # Reset the ACCUMULATE field for each tile
+ #
+ tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+
+ #
+ # Mma mainloop
+ #
+ for k_tile in cutlass.range(k_tile_cnt):
+ if is_leader_cta:
+ # Conditionally wait for AB buffer full
+ ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status)
+
+ # Copy SFA/SFB from smem to tmem
+ s2t_stage_coord = (
+ None,
+ None,
+ None,
+ None,
+ ab_consumer_state.index,
+ )
+ tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord]
+ tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord]
+ cute.copy(
+ tiled_copy_s2t_sfa,
+ tCsSFA_compact_s2t_staged,
+ tCtSFA_compact_s2t,
+ )
+ cute.copy(
+ tiled_copy_s2t_sfb,
+ tCsSFB_compact_s2t_staged,
+ tCtSFB_compact_s2t,
+ )
+
+ # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB
+ num_kblocks = cute.size(tCrA, mode=[2])
+
+ for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):
+ kblock_coord = (
+ None,
+ None,
+ kblock_idx,
+ ab_consumer_state.index,
+ )
+
+ # Set SFA/SFB tensor to tiled_mma
+ sf_kblock_coord = (None, None, kblock_idx)
+ tiled_mma.set(
+ tcgen05.Field.SFA,
+ tCtSFA[sf_kblock_coord].iterator,
+ )
+ tiled_mma.set(
+ tcgen05.Field.SFB,
+ tCtSFB_mma[sf_kblock_coord].iterator,
+ )
+
+ cute.gemm(
+ tiled_mma,
+ tCtAcc,
+ tCrA[kblock_coord],
+ tCrB[kblock_coord],
+ tCtAcc,
+ )
+
+ # Enable accumulate on tCtAcc after first kblock
+ tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+
+ # Async arrive AB buffer empty
+ ab_pipeline.consumer_release(ab_consumer_state)
+
+ # Peek (try_wait) AB buffer full for k_tile = k_tile + 1
+ ab_consumer_state.advance()
+ peek_ab_full_status = cutlass.Boolean(1)
+ if ab_consumer_state.count < k_tile_cnt:
+ if is_leader_cta:
+ peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state)
+
+ #
+ # Async arrive accumulator buffer full(each kblock)
+ #
+ if is_leader_cta:
+ acc_pipeline.producer_commit(acc_producer_state)
+
+ # Peek (try_wait) Acc buffer empty for k_tile = k_tile + 1
+ acc_producer_state.advance()
+ #
+ # Advance to next tile
+ #
+ tile_info_pipeline.consumer_wait(tile_info_consumer_state)
+ tile_info[0] = sInfo[(0, tile_info_consumer_state.index)]
+ tile_info[1] = sInfo[(1, tile_info_consumer_state.index)]
+ tile_info[2] = sInfo[(2, tile_info_consumer_state.index)]
+ tile_info[3] = sInfo[(3, tile_info_consumer_state.index)]
+ is_valid_tile = tile_info[3] == 1
+ cute.arch.fence_proxy(
+ "async.shared",
+ space="cta",
+ )
+ tile_info_pipeline.consumer_release(tile_info_consumer_state)
+ tile_info_consumer_state.advance()
+ #
+ # Wait for accumulator buffer empty
+ #
+ acc_pipeline.producer_tail(acc_producer_state)
+
+ #
+ # Specialized metadata loader warp
+ #
+ if warp_idx == self.meta_load_warp_id:
+ meta_lane = tidx % self.threads_per_warp
+
+ tile_info_consumer_state = pipeline.make_pipeline_state(
+ pipeline.PipelineUserType.Consumer, self.num_tile_stage
+ )
+ meta_producer_state = pipeline.make_pipeline_state(
+ pipeline.PipelineUserType.Producer, self.num_meta_stage
+ )
+ tile_info = cute.make_rmem_tensor((self.num_tile_info_fields,), cutlass.Int32)
+ tile_info_pipeline.consumer_wait(tile_info_consumer_state)
+ tile_info[0] = sInfo[(0, tile_info_consumer_state.index)]
+ tile_info[2] = sInfo[(2, tile_info_consumer_state.index)]
+ tile_info[3] = sInfo[(3, tile_info_consumer_state.index)]
+ tile_info[4] = sInfo[(4, tile_info_consumer_state.index)]
+ if cutlass.const_expr(self.use_expert_counts):
+ tile_info[5] = sInfo[(5, tile_info_consumer_state.index)]
+ is_valid_tile = tile_info[3] == 1
+ cute.arch.fence_proxy("async.shared", space="cta")
+ tile_info_pipeline.consumer_release(tile_info_consumer_state)
+ tile_info_consumer_state.advance()
+
+ while is_valid_tile:
+ tile_m_start = tile_info[0] * self.cta_tile_shape_mnk[0]
+ expert_idx = tile_info[2]
+ alpha_val = alpha_tuple[0][expert_idx]
+
+ meta_pipeline.producer_acquire(meta_producer_state)
+ meta_stage = meta_producer_state.index
+ # Strided row assignment keeps the permuted_idx loads and smem
+ # stores coalesced (each fixed j: 32 lanes touch 32 contiguous
+ # rows). Rows beyond mn_limit are padding: their expanded_idx is
+ # not guaranteed to be in range, so gate the token_final_scales
+ # gather to token 0 for them (branchless) to keep it in-bounds.
+ # The epilogue ignores padding rows via its own is_valid_row, so
+ # the value staged for them is irrelevant.
+ for j in cutlass.range(
+ self.cta_tile_shape_mnk[0] // self.threads_per_warp,
+ unroll_full=True,
+ ):
+ r = meta_lane + j * self.threads_per_warp
+ permuted_row = tile_m_start + r
+ if cutlass.const_expr(self.use_expert_counts):
+ expanded_idx = tile_info[5] + r
+ else:
+ expanded_idx = permuted_idx_to_expanded_idx[permuted_row]
+ # max(., 0) keeps topk_idx in-bounds for padding rows (-1);
+ # the is_valid gate keeps token_idx in-bounds (padding rows
+ # may hold out-of-range garbage, not just -1).
+ safe_idx = cutlass.max(expanded_idx, cutlass.Int32(0))
+ token_idx = safe_idx // topK
+ topk_idx = safe_idx % topK
+ is_valid_row = cutlass.Int32(permuted_row < tile_info[4])
+ gather_tok = token_idx * is_valid_row
+ token_scale = token_final_scales[(gather_tok, topk_idx)]
+ sMetaTokenIdx[(r, meta_stage)] = token_idx
+ sMetaScale[(r, meta_stage)] = alpha_val * token_scale
+ cute.arch.fence_proxy("async.shared", space="cta")
+ meta_pipeline.producer_commit(meta_producer_state)
+ meta_producer_state.advance()
+
+ tile_info_pipeline.consumer_wait(tile_info_consumer_state)
+ tile_info[0] = sInfo[(0, tile_info_consumer_state.index)]
+ tile_info[2] = sInfo[(2, tile_info_consumer_state.index)]
+ tile_info[3] = sInfo[(3, tile_info_consumer_state.index)]
+ tile_info[4] = sInfo[(4, tile_info_consumer_state.index)]
+ if cutlass.const_expr(self.use_expert_counts):
+ tile_info[5] = sInfo[(5, tile_info_consumer_state.index)]
+ is_valid_tile = tile_info[3] == 1
+ cute.arch.fence_proxy("async.shared", space="cta")
+ tile_info_pipeline.consumer_release(tile_info_consumer_state)
+ tile_info_consumer_state.advance()
+ meta_pipeline.producer_tail(meta_producer_state)
+
+ #
+ # Specialized epilogue warps
+ #
+ if warp_idx < self.mma_warp_id:
+ #
+ # Alloc tensor memory buffer
+ #
+ tmem.allocate(self.num_tmem_alloc_cols)
+
+ #
+ # Bar sync for retrieve tensor memory ptr from shared memory
+ #
+ tmem.wait_for_alloc()
+
+ #
+ # Retrieving tensor memory ptr and make accumulator tensor
+ #
+ tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+ # (MMA, MMA_M, MMA_N, STAGE)
+ tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
+
+ #
+ # Partition for epilogue
+ #
+ epi_tidx = tidx % 128
+ (
+ tiled_copy_t2r,
+ tTR_tAcc_base,
+ tTR_rAcc,
+ ) = self.epilog_tmem_copy_and_partition(
+ epi_tidx, tCtAcc_base, tCgC, epi_tile, use_2cta_instrs
+ )
+
+ tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.out_dtype)
+ tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition(
+ epi_tidx, tTR_rC, sC, tiled_copy_t2r
+ )
+
+ acc_consumer_state = pipeline.make_pipeline_state(
+ pipeline.PipelineUserType.Consumer, self.num_acc_stage
+ )
+
+ tile_info_consumer_state = pipeline.make_pipeline_state(
+ pipeline.PipelineUserType.Consumer, self.num_tile_stage
+ )
+ meta_consumer_state = pipeline.make_pipeline_state(
+ pipeline.PipelineUserType.Consumer, self.num_meta_stage
+ )
+
+ # Get the first tile info
+ tile_info = cute.make_rmem_tensor((5,), cutlass.Int32)
+
+ tile_info_pipeline.consumer_wait(tile_info_consumer_state)
+ tile_info[0] = sInfo[(0, tile_info_consumer_state.index)]
+ tile_info[1] = sInfo[(1, tile_info_consumer_state.index)]
+ tile_info[3] = sInfo[(3, tile_info_consumer_state.index)]
+ tile_info[4] = sInfo[(4, tile_info_consumer_state.index)]
+ is_valid_tile = tile_info[3] == 1
+ cute.arch.fence_proxy(
+ "async.shared",
+ space="cta",
+ )
+ tile_info_pipeline.consumer_release(tile_info_consumer_state)
+ tile_info_consumer_state.advance()
+
+ while is_valid_tile:
+ tile_m_start = tile_info[0] * self.cta_tile_shape_mnk[0]
+ permuted_row = tile_m_start + epi_tidx
+ is_valid_row = permuted_row < tile_info[4]
+ if cutlass.const_expr(self.row_owned_epilogue):
+ cute.arch.cp_async_bulk_wait_group(0, read=True)
+
+ # Read per-row finalize metadata prefetched by the meta loader
+ # warp (token_idx for scatter, combined_scale = alpha * token_scale).
+ meta_pipeline.consumer_wait(meta_consumer_state)
+ token_idx = sMetaTokenIdx[(epi_tidx, meta_consumer_state.index)]
+ meta_scale = sMetaScale[(epi_tidx, meta_consumer_state.index)]
+
+ # Get accumulator stage index
+ if cutlass.const_expr(self.overlapping_accum):
+ acc_stage_index = acc_consumer_state.phase
+ reverse_subtile = (
+ cutlass.Boolean(True) if acc_stage_index == 0 else cutlass.Boolean(False)
+ )
+ else:
+ acc_stage_index = acc_consumer_state.index
+
+ # Set tensor memory buffer for current tile
+ # (T2R, T2R_M, T2R_N, EPI_M, EPI_M)
+ tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_stage_index)]
+
+ #
+ # Wait for accumulator buffer full
+ #
+ acc_pipeline.consumer_wait(acc_consumer_state)
+
+ tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
+ #
+ # Process sub-tiles with vectorized scatter-add
+ # Following TensorRT-LLM's direct G2R (global to register) approach
+ #
+ subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
+
+ for subtile_idx in cutlass.range(subtile_cnt):
+ real_subtile_idx = subtile_idx
+ if cutlass.const_expr(self.overlapping_accum):
+ if reverse_subtile:
+ real_subtile_idx = subtile_cnt - 1 - subtile_idx
+ #
+ # Load accumulator from tensor memory buffer to register
+ #
+ tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)]
+
+ cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc)
+
+ #
+ # Async arrive accumulator buffer empty earlier when overlapping_accum is enabled
+ #
+ if cutlass.const_expr(self.overlapping_accum):
+ if subtile_idx == self.iter_acc_early_release_in_epilogue:
+ # Fence for TMEM load
+ cute.arch.fence_view_async_tmem_load()
+ acc_pipeline.consumer_release(acc_consumer_state)
+ acc_consumer_state.advance()
+
+ # Get vectorized accumulator and apply the combined finalize scale
+ acc_vec = tTR_rAcc.load()
+ acc_vec_final = meta_scale * acc_vec
+
+ tRS_rC.store(acc_vec_final.to(self.out_dtype))
+ if is_valid_row:
+ cute.copy(
+ tiled_copy_r2s,
+ tRS_rC,
+ tRS_sC[(None, None, real_subtile_idx, None)],
+ )
+
+ # Make all R2S smem writes visible to the async bulk-reduce proxy.
+ cute.arch.fence_proxy(
+ "async.shared",
+ space="cta",
+ )
+ #
+ # Async arrive accumulator buffer empty
+ #
+ if cutlass.const_expr(not self.overlapping_accum):
+ cute.arch.fence_view_async_tmem_load()
+ acc_pipeline.consumer_release(acc_consumer_state)
+ acc_consumer_state.advance()
+
+ # Whole-row async bulk reduce (smem -> global scatter-add).
+ if is_valid_row:
+ coord_n = tile_info[1] * self.cta_tile_shape_mnk[1]
+ scatter_out_offset = cute.domain_offset((token_idx, coord_n, 0), out)
+ if cutlass.const_expr(self.out_dtype == cutlass.BFloat16 and self.l2_hints):
+ _blk_reduce_bf16_evict_last(
+ scatter_out_offset,
+ sC[epi_tidx, None, 0],
+ cutlass.Int32(self.copy_size),
+ )
+ elif cutlass.const_expr(self.out_dtype == cutlass.BFloat16):
+ blk_reduce_bf16(
+ scatter_out_offset,
+ sC[epi_tidx, None, 0],
+ cutlass.Int32(self.copy_size),
+ )
+ elif cutlass.const_expr(self.out_dtype == cutlass.Float32):
+ blk_reduce_fp32(
+ scatter_out_offset,
+ sC[epi_tidx, None, 0],
+ cutlass.Int32(self.copy_size),
+ )
+ elif cutlass.const_expr(self.out_dtype == cutlass.Float16):
+ blk_reduce_fp16(
+ scatter_out_offset,
+ sC[epi_tidx, None, 0],
+ cutlass.Int32(self.copy_size),
+ )
+
+ cute.arch.cp_async_bulk_commit_group()
+ if cutlass.const_expr(not self.row_owned_epilogue):
+ cute.arch.cp_async_bulk_wait_group(0, read=True)
+ self.epilog_sync_barrier.arrive_and_wait()
+
+ # Release the prefetched metadata slot for this tile.
+ meta_pipeline.consumer_release(meta_consumer_state)
+ meta_consumer_state.advance()
+
+ # Advance to next tile
+ #
+ tile_info_pipeline.consumer_wait(tile_info_consumer_state)
+ tile_info[0] = sInfo[(0, tile_info_consumer_state.index)]
+ tile_info[1] = sInfo[(1, tile_info_consumer_state.index)]
+ tile_info[3] = sInfo[(3, tile_info_consumer_state.index)]
+ tile_info[4] = sInfo[(4, tile_info_consumer_state.index)]
+ is_valid_tile = tile_info[3] == 1
+ cute.arch.fence_proxy(
+ "async.shared",
+ space="cta",
+ )
+ tile_info_pipeline.consumer_release(tile_info_consumer_state)
+ tile_info_consumer_state.advance()
+ if cutlass.const_expr(self.row_owned_epilogue):
+ cute.arch.cp_async_bulk_wait_group(0, read=True)
+ #
+ # Dealloc the tensor memory buffer
+ #
+ tmem.relinquish_alloc_permit()
+ self.epilog_sync_barrier.arrive_and_wait()
+ tmem.free(tmem_ptr)
+
+ griddepcontrol_launch_dependents()
+
+ def epilog_tmem_copy_and_partition(
+ self,
+ tidx: cutlass.Int32,
+ tAcc: cute.Tensor,
+ gC_mnl: cute.Tensor,
+ epi_tile: cute.Tile,
+ use_2cta_instrs: Union[cutlass.Boolean, bool],
+ ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
+ """
+ Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array
+ (destination).
+
+ :param tidx: The thread index in epilogue warp groups
+ :type tidx: cutlass.Int32
+ :param tAcc: The accumulator tensor to be copied and partitioned
+ :type tAcc: cute.Tensor
+ :param gC_mnl: The global tensor C
+ :type gC_mnl: cute.Tensor
+ :param epi_tile: The epilogue tiler
+ :type epi_tile: cute.Tile
+ :param use_2cta_instrs: Whether use_2cta_instrs is enabled
+ :type use_2cta_instrs: bool
+
+ :return: A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where:
+ - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r)
+ - tTR_tAcc: The partitioned accumulator tensor
+ - tTR_rAcc: The accumulated tensor in register used to hold t2r results
+ :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]
+ """
+ # Make tiledCopy for tensor memory load
+ copy_atom_t2r = sm100_utils.get_tmem_load_op(
+ self.cta_tile_shape_mnk,
+ self.gemm_output_layout,
+ self.out_dtype,
+ self.acc_dtype,
+ epi_tile,
+ use_2cta_instrs,
+ )
+
+ # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE)
+ tAcc_epi = cute.flat_divide(
+ tAcc[((None, None), 0, 0, None)],
+ epi_tile,
+ )
+ # (EPI_TILE_M, EPI_TILE_N)
+ 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)
+ # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE)
+ tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi)
+
+ # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, loopM, loopN, loopL)
+ gC_mnl_epi = cute.flat_divide(gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile)
+
+ # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, loopM, loopN, loopL)
+ tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi)
+ # (T2R, T2R_M, T2R_N)
+ tTR_rAcc = cute.make_rmem_tensor(
+ tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype
+ )
+
+ return tiled_copy_t2r, tTR_tAcc, tTR_rAcc
+
+ def epilog_smem_copy_and_partition(
+ self,
+ tidx: cutlass.Int32,
+ tTR_rC: cute.Tensor,
+ sC: cute.Tensor,
+ tiled_copy_t2r: cute.TiledCopy,
+ ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
+ """
+ Create tiled copy for register to shared memory (R2S).
+ """
+ atom = cute.make_copy_atom(
+ cute.nvgpu.CopyUniversalOp(),
+ self.out_dtype,
+ )
+
+ tiled_copy_r2s = cute.make_tiled_copy_D(atom, tiled_copy_t2r)
+ # (R2S, R2S_M, R2S_N, PIPE_D)
+ thr_copy_r2s = tiled_copy_r2s.get_slice(tidx)
+ tRS_sC = thr_copy_r2s.partition_D(sC)
+ tRS_rC = tiled_copy_r2s.retile(tTR_rC)
+ return tiled_copy_r2s, tRS_rC, tRS_sC
+
+ @staticmethod
+ def _compute_stages(
+ tiled_mma: cute.TiledMma,
+ mma_tiler_mnk: Tuple[int, int, int],
+ a_dtype: Type[cutlass.Numeric],
+ b_dtype: Type[cutlass.Numeric],
+ out_dtype: Type[cutlass.Numeric],
+ cta_tile: cute.Tile,
+ sf_dtype: Type[cutlass.Numeric],
+ sf_vec_size: int,
+ final_scale_dtype: Type[cutlass.Numeric],
+ num_smem_capacity: int,
+ occupancy: int,
+ ) -> Tuple[int, int, int, int, int]:
+ """Computes the number of stages for A/B/C operands based on heuristics.
+
+ :param tiled_mma: The tiled MMA object defining the core computation.
+ :type tiled_mma: cute.TiledMma
+ :param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler.
+ :type mma_tiler_mnk: tuple[int, int, int]
+ :param a_dtype: Data type of operand A.
+ :type a_dtype: type[cutlass.Numeric]
+ :param b_dtype: Data type of operand B.
+ :type b_dtype: type[cutlass.Numeric]
+ :param out_dtype: Data type of operand C (output).
+ :type out_dtype: type[cutlass.Numeric]
+ :param cta_tile: The CTA tile shape.
+ :type cta_tile: cute.Tile
+ :param sf_dtype: Data type of scale factor.
+ :type sf_dtype: type[cutlass.Numeric]
+ :param sf_vec_size: Vector size of scale factor.
+ :type sf_vec_size: int
+ :param num_smem_capacity: Total available shared memory capacity in bytes.
+ :type num_smem_capacity: int
+ :param occupancy: Target number of CTAs per SM (occupancy).
+ :type occupancy: int
+
+ :return: A tuple containing the computed number of stages for:
+ (ACC stages, A/B operand stages, C stages)
+ :rtype: tuple[int, int, int]
+ """
+ # Default ACC stages
+ num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2
+
+ # Default C stages
+ num_c_stage = 1
+
+ # Default Tile info stages
+ num_tile_stage = 2
+
+ # Metadata loader pipeline depth (meta warp -> epilogue). Lets the loader
+ # prefetch the per-row {token_idx, combined_scale} ahead of the epilogue.
+ num_meta_stage = 2
+ # Per-row metadata smem held by the loader: token_idx (int32) +
+ # combined_scale (final_scale_dtype), cta_M rows per stage.
+ token_idx_bytes = cutlass.Int32.width // 8
+ scale_bytes = final_scale_dtype.width // 8
+ meta_smem_bytes = cta_tile[0] * (token_idx_bytes + scale_bytes) * num_meta_stage
+
+ # Calculate smem layout and size for one stage of A, B, and C
+ a_smem_layout_stage_one = sm100_utils.make_smem_layout_a(
+ tiled_mma,
+ mma_tiler_mnk,
+ a_dtype,
+ 1, # a tmp 1 stage is provided
+ )
+ b_smem_layout_staged_one = sm100_utils.make_smem_layout_b(
+ tiled_mma,
+ mma_tiler_mnk,
+ b_dtype,
+ 1, # a tmp 1 stage is provided
+ )
+
+ sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa(
+ tiled_mma,
+ mma_tiler_mnk,
+ sf_vec_size,
+ 1, # a tmp 1 stage is provided
+ )
+
+ sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb(
+ tiled_mma,
+ mma_tiler_mnk,
+ sf_vec_size,
+ 1, # a tmp 1 stage is provided
+ )
+
+ # satisfy 16B alignment for the output tensor
+ swizzled_pad = 16 // (out_dtype.width // 8)
+ c_smem_layout_staged_one = cute.make_layout(
+ (cta_tile[0], cta_tile[1]), stride=(cta_tile[1] + swizzled_pad, 1)
+ )
+
+ ab_bytes_per_stage = (
+ cute.size_in_bytes(a_dtype, a_smem_layout_stage_one)
+ + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one)
+ + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one)
+ + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one)
+ )
+ # 1024B alignment for mbar
+ mbar_helpers_bytes = 1024
+
+ c_bytes_per_stage = cute.size_in_bytes(out_dtype, c_smem_layout_staged_one)
+ c_bytes = c_bytes_per_stage * num_c_stage
+
+ # Calculate A/B stages:
+ # Start with total smem per CTA (capacity / occupancy)
+ # Subtract reserved bytes, initial C stages bytes, and the per-row
+ # metadata smem held by the loader warp
+ # Divide remaining by bytes needed per A/B stage
+ num_ab_stage = (
+ num_smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes + meta_smem_bytes)
+ ) // ab_bytes_per_stage
+
+ return num_acc_stage, num_ab_stage, num_c_stage, num_tile_stage, num_meta_stage
+
+ @staticmethod
+ def _compute_grid(
+ gemm_shape: Tuple[int, int, int],
+ cta_tile_shape_mnk: Tuple[int, int, int],
+ cluster_shape_mn: Tuple[int, int],
+ max_active_clusters: cutlass.Constexpr,
+ raster_along_m: bool,
+ ) -> Tuple[utils.PersistentTileSchedulerParams, Tuple[int, int, int]]:
+ """Use persistent tile scheduler to compute the grid size based on GEMM shape.
+
+ :param gemm_shape: The GEMM computation shape (M, N, L)
+ :type gemm_shape: tuple[int, int, int]
+ :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile.
+ :type cta_tile_shape_mnk: tuple[int, int, int]
+ :param cluster_shape_mn: Shape of each cluster in M, N dimensions.
+ :type cluster_shape_mn: tuple[int, int]
+ :param max_active_clusters: Maximum number of active clusters.
+ :type max_active_clusters: cutlass.Constexpr
+ :param raster_along_m: Boolean, True to use raster along M.
+ :type raster_along_m: bool
+
+ :return: A tuple containing:
+ - tile_sched_params: Parameters for the persistent tile scheduler.
+ - grid: Grid shape for kernel launch.
+ :rtype: Tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]]
+ """
+ (m, n, l) = gemm_shape # noqa: E741
+
+ num_ctas_m = cute.ceil_div(m, cta_tile_shape_mnk[0])
+ num_ctas_n = cute.ceil_div(n, cta_tile_shape_mnk[1])
+ num_ctas_l = l
+
+ num_ctas_mnl = (num_ctas_m, num_ctas_n, num_ctas_l)
+ cluster_shape_mnl = (*cluster_shape_mn, 1)
+
+ tile_sched_params = utils.PersistentTileSchedulerParams(
+ num_ctas_mnl, cluster_shape_mnl, raster_along_m=raster_along_m
+ )
+ grid = utils.StaticPersistentTileScheduler.get_grid_shape(
+ tile_sched_params, max_active_clusters
+ )
+
+ return tile_sched_params, grid
+
+ @staticmethod
+ def _get_tma_atom_kind(
+ atom_sm_cnt: cutlass.Int32, mcast: cutlass.Boolean
+ ) -> Union[cpasync.CopyBulkTensorTileG2SMulticastOp, cpasync.CopyBulkTensorTileG2SOp]:
+ """
+ Select the appropriate TMA copy atom based on the number of SMs and the multicast flag.
+
+ :param atom_sm_cnt: The number of SMs
+ :type atom_sm_cnt: cutlass.Int32
+ :param mcast: The multicast flag
+ :type mcast: cutlass.Boolean
+
+ :return: The appropriate TMA copy atom kind
+ :rtype: cpasync.CopyBulkTensorTileG2SMulticastOp or cpasync.CopyBulkTensorTileG2SOp
+
+ :raise ValueError: If the atom_sm_cnt is invalid
+ """
+ if atom_sm_cnt == 2 and mcast:
+ return cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.TWO)
+ elif atom_sm_cnt == 2 and not mcast:
+ return cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.TWO)
+ elif atom_sm_cnt == 1 and mcast:
+ return cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.ONE)
+ elif atom_sm_cnt == 1 and not mcast:
+ return cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE)
+
+ raise ValueError(f"Invalid atom_sm_cnt: {atom_sm_cnt} and {mcast}")
+
+ @staticmethod
+ def is_valid_dtypes_and_scale_factor_vec_size(
+ ab_dtype: Type[cutlass.Numeric],
+ sf_dtype: Type[cutlass.Numeric],
+ sf_vec_size: int,
+ out_dtype: Type[cutlass.Numeric],
+ ) -> bool:
+ """
+ Check if the dtypes are valid
+
+ :param ab_dtype: The data type of the A and B operands
+ :type ab_dtype: Type[cutlass.Numeric]
+ :param sf_dtype: The data type of the scale factor
+ :type sf_dtype: Type[cutlass.Numeric]
+ :param sf_vec_size: The vector size of the scale factor
+ :type sf_vec_size: int
+ :param out_dtype: The data type of the output tensor
+ :type out_dtype: Type[cutlass.Numeric]
+
+ :return: True if the dtypes are valid, False otherwise
+ :rtype: bool
+ """
+ is_valid = True
+ if ab_dtype not in {
+ cutlass.Float4E2M1FN,
+ cutlass.Float8E5M2,
+ cutlass.Float8E4M3FN,
+ }:
+ is_valid = False
+
+ # Check valid sf_vec_size
+ if sf_vec_size not in {16, 32}:
+ is_valid = False
+
+ # Check valid sf_dtype
+ if sf_dtype not in {cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN}:
+ is_valid = False
+
+ # Check valid sf_dtype and sf_vec_size combinations
+ if sf_dtype == cutlass.Float8E4M3FN and sf_vec_size == 32:
+ is_valid = False
+ if ab_dtype in {cutlass.Float8E5M2, cutlass.Float8E4M3FN} and sf_vec_size == 16:
+ is_valid = False
+
+ if out_dtype not in {cutlass.Float32, cutlass.Float16, cutlass.BFloat16}:
+ is_valid = False
+
+ return is_valid
+
+ @staticmethod
+ def is_valid_layouts(
+ ab_dtype: Type[cutlass.Numeric],
+ out_dtype: Type[cutlass.Numeric],
+ a_major: str,
+ b_major: str,
+ out_major: str,
+ ) -> bool:
+ """
+ Check if layouts and dtypes are valid combinations
+
+ :param ab_dtype: The data type of the A and B operands
+ :type ab_dtype: Type[cutlass.Numeric]
+ :param out_dtype: The data type of the output tensor
+ :type out_dtype: Type[cutlass.Numeric]
+ :param a_major: The major dimension of the A tensor
+ :type a_major: str
+ :param b_major: The major dimension of the B tensor
+ :type b_major: str
+ :param out_major: The major dimension of the C tensor
+ :type out_major: str
+
+ :return: True if the layouts are valid, False otherwise
+ :rtype: bool
+ """
+ is_valid = True
+
+ if ab_dtype is cutlass.Float4E2M1FN and not (a_major == "k" and b_major == "k"):
+ is_valid = False
+ if out_dtype is cutlass.Float4E2M1FN and out_major == "m":
+ is_valid = False
+ return is_valid
+
+ @staticmethod
+ def is_valid_mma_tiler_and_cluster_shape(
+ mma_tiler_mn: Tuple[int, int],
+ cluster_shape_mn: Tuple[int, int],
+ ) -> bool:
+ """
+ Check if the mma tiler and cluster shape are valid
+
+ :param use_2cta_instrs: Whether to use 2 CTA groups
+ :type use_2cta_instrs: bool
+ :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler
+ :type mma_tiler_mn: Tuple[int, int]
+ :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster
+ :type cluster_shape_mn: Tuple[int, int]
+
+ :return: True if the mma tiler and cluster shape are valid, False otherwise
+ :rtype: bool
+ """
+ is_valid = True
+
+ # Skip invalid mma tile shape
+ if mma_tiler_mn[0] not in (128, 256):
+ is_valid = False
+ # Skip invalid mma tile n
+ if mma_tiler_mn[1] not in (64, 128, 192, 256):
+ is_valid = False
+
+ # Skip illegal cluster shape
+ if (mma_tiler_mn[0] // cluster_shape_mn[0]) != 128:
+ is_valid = False
+
+ if (
+ cluster_shape_mn[0] * cluster_shape_mn[1] > 16
+ or cluster_shape_mn[0] <= 0
+ or cluster_shape_mn[1] <= 0
+ # Special cluster shape check for scale factor multicasts.
+ # Due to limited size of scale factors, we can't multicast among more than 4 CTAs.
+ or cluster_shape_mn[0] > 4
+ or cluster_shape_mn[1] > 4
+ or not is_power_of_2(cluster_shape_mn[0])
+ or not is_power_of_2(cluster_shape_mn[1])
+ ):
+ is_valid = False
+
+ return is_valid
+
+ @staticmethod
+ def is_valid_tensor_alignment(
+ m: cutlass.Int64,
+ n: cutlass.Int64,
+ k: cutlass.Int64,
+ l: cutlass.Int64, # noqa: E741
+ ab_dtype: Type[cutlass.Numeric],
+ out_dtype: Type[cutlass.Numeric],
+ a_major: str,
+ b_major: str,
+ out_major: str,
+ ) -> bool:
+ """
+ Check if the tensor alignment is valid
+
+ :param m: The number of rows in the A tensor
+ :type m: cutlass.Int64
+ :param n: The number of columns in the B tensor
+ :type n: cutlass.Int64
+ :param k: The number of columns in the A tensor
+ :type k: cutlass.Int64
+ :param l: The number of columns in the C tensor
+ :type l: cutlass.Int64
+ :param ab_dtype: The data type of the A and B operands
+ :type ab_dtype: Type[cutlass.Numeric]
+ :param out_dtype: The data type of the output tensor
+ :type out_dtype: Type[cutlass.Numeric]
+ :param a_major: The major axis of the A tensor
+ :type a_major: str
+ :param b_major: The major axis of the B tensor
+ :type b_major: str
+ :param out_major: The major axis of the C tensor
+ :type out_major: str
+
+ :return: True if the problem shape is valid, False otherwise
+ :rtype: bool
+ """
+ is_valid = True
+
+ def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape):
+ major_mode_idx = 0 if is_mode0_major else 1
+ num_major_elements = tensor_shape[major_mode_idx]
+ num_contiguous_elements = 16 * 8 // dtype.width
+ return num_major_elements % num_contiguous_elements == 0
+
+ if (
+ not check_contigous_16B_alignment(ab_dtype, a_major == "m", (m, k, l))
+ or not check_contigous_16B_alignment(ab_dtype, b_major == "n", (n, k, l))
+ or not check_contigous_16B_alignment(out_dtype, out_major == "m", (m, n, l))
+ ):
+ is_valid = False
+ return is_valid
+
+ @classmethod
+ def can_implement(
+ cls,
+ ab_dtype: Type[cutlass.Numeric],
+ sf_dtype: Type[cutlass.Numeric],
+ sf_vec_size: int,
+ out_dtype: Type[cutlass.Numeric],
+ mma_tiler_mn: Tuple[int, int],
+ cluster_shape_mn: Tuple[int, int],
+ m: cutlass.Int64,
+ n: cutlass.Int64,
+ k: cutlass.Int64,
+ l: cutlass.Int64, # noqa: E741
+ a_major: str,
+ b_major: str,
+ out_major: str,
+ ) -> bool:
+ """
+ Check if the gemm can be implemented
+
+ :param ab_dtype: The data type of the A and B operands
+ :type ab_dtype: Type[cutlass.Numeric]
+ :param sf_dtype: The data type of the scale factor
+ :type sf_dtype: Type[cutlass.Numeric]
+ :param sf_vec_size: The vector size of the scale factor
+ :type sf_vec_size: int
+ :param out_dtype: The data type of the output tensor
+ :type out_dtype: Type[cutlass.Numeric]
+ :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler
+ :type mma_tiler_mn: Tuple[int, int]
+ :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster
+ :type cluster_shape_mn: Tuple[int, int]
+ :param m: The number of rows in the A tensor
+ :type m: cutlass.Int64
+ :param n: The number of columns in the B tensor
+ :type n: cutlass.Int64
+ :param k: The number of columns in the A tensor
+ :type k: cutlass.Int64
+ :param l: The number of columns in the C tensor
+ :type l: cutlass.Int64
+ :param a_major: The major axis of the A tensor
+ :type a_major: str
+ :param b_major: The major axis of the B tensor
+ :type b_major: str
+ :param out_major: The major axis of the C tensor
+ :type out_major: str
+
+ :return: True if the gemm can be implemented, False otherwise
+ :rtype: bool
+ """
+ can_implement = True
+ # Skip unsupported types
+ if not cls.is_valid_dtypes_and_scale_factor_vec_size(
+ ab_dtype, sf_dtype, sf_vec_size, out_dtype
+ ):
+ can_implement = False
+
+ # Skip unsupported layouts
+ if not cls.is_valid_layouts(ab_dtype, out_dtype, a_major, b_major, out_major):
+ can_implement = False
+
+ # Skip invalid mma tile shape and cluster shape
+ if not cls.is_valid_mma_tiler_and_cluster_shape(mma_tiler_mn, cluster_shape_mn):
+ can_implement = False
+ # Skip illegal problem shape for load/store alignment
+ if not cls.is_valid_tensor_alignment(
+ m, n, k, l, ab_dtype, out_dtype, a_major, b_major, out_major
+ ):
+ can_implement = False
+ # Skip unsupported A/B layout
+ if not (a_major == "k" and b_major == "k"):
+ can_implement = False
+ return can_implement
+
+ @cute.jit
+ def wrapper(
+ self,
+ a_ptr: cute.Pointer,
+ b_ptr: cute.Pointer,
+ a_sf_ptr: cute.Pointer,
+ b_sf_ptr: cute.Pointer,
+ c_ptr: cute.Pointer,
+ alpha_ptr: cute.Pointer,
+ tile_idx_to_group_idx_ptr: cute.Pointer,
+ tile_idx_to_mn_limit_ptr: cute.Pointer,
+ permuted_idx_to_expanded_idx_ptr: cute.Pointer,
+ num_non_exiting_tiles_ptr: cute.Pointer,
+ token_final_scales_ptr: cute.Pointer,
+ m: cutlass.Int64,
+ n: cutlass.Int64,
+ k: cutlass.Int64,
+ l: cutlass.Int64, # noqa: E741
+ num_tokens: cutlass.Int64,
+ top_k: cutlass.Int64,
+ expert_capacity: cutlass.Int32,
+ tile_size: cutlass.Constexpr,
+ scaling_vector_size: cutlass.Constexpr,
+ max_active_clusters: cutlass.Constexpr,
+ stream: cuda.CUstream,
+ epilogue_op: cutlass.Constexpr = lambda x: x,
+ c_stride_row: cutlass.Int64 = cutlass.Int64(0),
+ ):
+ """Single-B wrapper.
+
+ ``l`` is the number of experts in the (sole) B tensor.
+ """
+ scale_k = k // scaling_vector_size
+ num_tiles = m // tile_size
+
+ a = cute.make_tensor(a_ptr, layout=cute.make_ordered_layout((m, k, 1), order=(1, 0, 2)))
+ a_sf = cute.make_tensor(
+ a_sf_ptr,
+ layout=cute.make_ordered_layout(
+ (32, 4, m // 128, 4, scale_k // 4, 1), order=(2, 1, 4, 0, 3, 5)
+ ),
+ )
+ # c supports strided output for locality domain shared buffers.
+ # c_stride_row: row stride in output elements (0 = default = n)
+ actual_c_stride_row = n if c_stride_row == 0 else c_stride_row
+ c = cute.make_tensor(
+ c_ptr,
+ layout=cute.make_layout(
+ (num_tokens, n, 1),
+ stride=(actual_c_stride_row, 1, num_tokens * actual_c_stride_row),
+ ),
+ )
+
+ alpha = cute.make_tensor(alpha_ptr, layout=cute.make_layout((l,)))
+ b = cute.make_tensor(b_ptr, layout=cute.make_ordered_layout((n, k, l), order=(1, 0, 2)))
+ b_sf = cute.make_tensor(
+ b_sf_ptr,
+ layout=cute.make_ordered_layout(
+ (32, 4, n // 128, 4, scale_k // 4, l), order=(2, 1, 4, 0, 3, 5)
+ ),
+ )
+
+ tile_idx_to_group_idx = cute.make_tensor(
+ tile_idx_to_group_idx_ptr, layout=cute.make_layout((num_tiles,))
+ )
+ tile_idx_to_mn_limit = cute.make_tensor(
+ tile_idx_to_mn_limit_ptr, layout=cute.make_layout((num_tiles,))
+ )
+ permuted_idx_to_expanded_idx = cute.make_tensor(
+ permuted_idx_to_expanded_idx_ptr, layout=cute.make_layout((m,))
+ )
+ num_non_exiting_tiles = cute.make_tensor(
+ num_non_exiting_tiles_ptr, layout=cute.make_layout((1,))
+ )
+ token_final_scales = cute.make_tensor(
+ token_final_scales_ptr,
+ layout=cute.make_ordered_layout((num_tokens, top_k), order=(1, 0)),
+ )
+
+ return self(
+ a,
+ b,
+ c,
+ a_sf,
+ b_sf,
+ tile_idx_to_group_idx,
+ num_non_exiting_tiles,
+ tile_idx_to_mn_limit,
+ alpha,
+ expert_capacity,
+ max_active_clusters=max_active_clusters,
+ stream=stream,
+ permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx,
+ token_final_scales=token_final_scales,
+ epilogue_op=epilogue_op,
+ )
+
+
+@cute.jit
+def cvt_sf_MKL_to_M32x4xrm_K4xrk_L(
+ sf_ref_tensor: cute.Tensor,
+ sf_mma_tensor: cute.Tensor,
+):
+ """Convert scale factor tensor from MKL layout to mma specification M(32x4xrest_m)xK(4xrest_k)xL layout"""
+ # sf_mma_tensor has flatten shape (32, 4, rest_m, 4, rest_k, l)
+ # group to ((32, 4, rest_m), (4, rest_k), l)
+ sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3)
+ sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3)
+ for i in cutlass.range(cute.size(sf_ref_tensor)):
+ mkl_coord = sf_ref_tensor.layout.get_hier_coord(i)
+ sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/__init__.py
new file mode 100644
index 000000000000..a8aa08cbff1d
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Operations of the qwen4_exp family's target."""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/__init__.py
new file mode 100644
index 000000000000..52a7a9daf028
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/__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/ops/activation/fused_sigmoid_mul_inplace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/fused_sigmoid_mul_inplace.py
new file mode 100644
index 000000000000..98189e91514c
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/fused_sigmoid_mul_inplace.py
@@ -0,0 +1,227 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""In-place sigmoid output gating: ``attention_output *= sigmoid(gate)``.
+
+The gated attention's output gate, between the attention core and o_proj, in
+one Triton kernel (`torch.ops.trtllm.fused_sigmoid_mul_inplace`). Called by the
+qwen4_exp target (`qwen3_8_flash_next_nvfp4__sm_103__default`) at every layout -- its
+trunk, gate off the fused input-projection row, and its MTP layer, gate as the strided
+half of the q projection.
+"""
+
+from __future__ import annotations
+
+import torch
+
+# Registered by the module that defines it: after `import
+# tensorrt_llm._torch.custom_ops` alone, torch.ops.trtllm has no
+# `fused_sigmoid_mul_inplace` (measured).
+import tensorrt_llm._torch.attention.kernels.fused_qk_norm_rope_gate # noqa: F401
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+#: The checkpoint's gated attention: 24 q heads of 256, q projection [T, 24, 512].
+_CKPT = dict(heads=24, head_dim=256, dtype=torch.bfloat16, gate="interleaved")
+
+#: The default target at tp4 / tep4: a rank's fused input projection row is
+#: [q 6 x 256 | k 256 | v 256 | gate 6 x 256 | index scratch 768] = 4352 columns,
+#: so the gate is the 1536-wide window at column qkv_width = 1536 + 2 * 256 = 2048.
+_TP4_ROW = dict(
+ heads=6,
+ head_dim=256,
+ dtype=torch.bfloat16,
+ gate="fused_row",
+ row_stride=4352,
+ gate_col=2048,
+)
+
+#: The default target at dep4: tp1's widths, 14080 columns, the gate at 7168.
+_DEP4_ROW = dict(
+ heads=24,
+ head_dim=256,
+ dtype=torch.bfloat16,
+ gate="fused_row",
+ row_stride=14080,
+ gate_col=7168,
+)
+
+#: The default target's MTP layer head-sharded at tep4: q_proj 12288 -> 3072 columns a
+#: rank, read as [T, 6, 512], the gate the second 256 of each head.
+_MTP_TEP4 = dict(heads=6, head_dim=256, dtype=torch.bfloat16, gate="interleaved")
+
+
+class FusedSigmoidMulInplace(OpWrapper):
+ """`attention_output[t, i] *= sigmoid(gate[t, i // D, i % D])` in place, `D = gate.shape[-1]`.
+
+ Returns `attention_output`, the tensor it wrote. Both operands are widened to
+ fp32 and the product is rounded once to `attention_output.dtype`. `gate` is
+ read through its own token and head strides, so the gate half of an
+ interleaved `[T, heads, 2 * D]` q projection is passed as a strided view with
+ no repack (bitwise equal to its compacted copy, measured). A 2-D `gate` of
+ `attention_output`'s shape is one head spanning the row.
+
+ Fusion boundary: sigmoid and multiply only. The caller owns producing both
+ operands and everything downstream (the o projection). `gate` is read-only.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ # The target at tp1: hidden 6144, the trunk's gate off the 14080-wide fused
+ # row (the test's default row stride and gate column), the MTP layer's
+ # off its [T, 24, 512] q_proj.
+ Cell(
+ why="decode 128 tokens: 2-D gate off the fused projection row (stride 14080)",
+ spec=dict(tokens=128, heads=24, head_dim=256, dtype=torch.bfloat16, gate="fused_row"),
+ ),
+ Cell(
+ why="prefill 8192 tokens: 2-D gate off the fused projection row, 2048-wide blocks",
+ spec=dict(tokens=8192, heads=24, head_dim=256, dtype=torch.bfloat16, gate="fused_row"),
+ ),
+ Cell(
+ why="MTP decode: 1 token, 24x256 gate as the strided half of q_proj",
+ spec=dict(tokens=1, **_CKPT),
+ ),
+ # The default target's trunk. Hidden 1536 is not a multiple of the kernel's
+ # block, so every tp4 / tep4 row ends in 512 masked lanes: two 1024-wide
+ # blocks below 1024 rows, one 2048-wide block from 1024 rows.
+ Cell(
+ why="default tp4/tep4 decode, one token: 6 heads off the 4352-wide row",
+ spec=dict(tokens=1, **_TP4_ROW),
+ ),
+ Cell(
+ why="default tp4/tep4 largest captured decode batch: 128 rows, 6 heads",
+ spec=dict(tokens=128, **_TP4_ROW),
+ ),
+ Cell(
+ why="default tep4 MTP-7 verify 128 x 8 = 1024 rows: the first 2048-wide block count",
+ spec=dict(tokens=1024, **_TP4_ROW),
+ ),
+ Cell(
+ why="default tp4/tep4 prefill chunk at max_num_tokens: 8192 rows, 6 heads",
+ spec=dict(tokens=8192, **_TP4_ROW),
+ ),
+ Cell(
+ why="default dep4 rank's largest captured decode batch: 32 rows, 24 heads",
+ spec=dict(tokens=32, **_DEP4_ROW),
+ ),
+ # The default target's MTP layer at tep4, on the same 1536-wide masked rows.
+ Cell(
+ why="default MTP drafter, 1 row: 6-head gate half of the [T, 6, 512] q_proj",
+ spec=dict(tokens=1, **_MTP_TEP4),
+ ),
+ Cell(
+ why="default MTP drafter at 128 requests: 128 rows, 6 heads, 1024-wide blocks",
+ spec=dict(tokens=128, **_MTP_TEP4),
+ ),
+ Cell(
+ why="default MTP-3 verify 128 x 4 = 512 rows, 6 heads",
+ spec=dict(tokens=512, **_MTP_TEP4),
+ ),
+ Cell(
+ why="default MTP-7 verify 128 x 8 = 1024 rows: the first 2048-wide block count",
+ spec=dict(tokens=1024, **_MTP_TEP4),
+ ),
+ Cell(
+ why="default MTP prefill chunk at max_num_tokens: 8192 rows, 6 heads",
+ spec=dict(tokens=8192, **_MTP_TEP4),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none. Source: `torch.ops.trtllm.fused_sigmoid_mul_inplace`,
+ a Python `torch.library.custom_op` registered by
+ `tensorrt_llm/_torch/attention/kernels/fused_qk_norm_rope_gate.py` around the
+ Triton kernel `_fused_sigmoid_mul_kernel` there. One kernel, JIT-compiled for
+ the running device at the first call of each (hidden, head_dim, block)
+ specialization, no autotuner: grid `(T, cdiv(hidden, block))`,
+ `block = min(next_pow2(hidden), 1024 if T < 1024 else 2048)`, 4 warps. The
+ cells ran on sm_103.
+
+ Call sites, each the attention output gate between the attention core and
+ o_proj:
+ - The target at tp1 (`qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default`): the
+ trunk's gate as the 2-D window of the fused input projection row
+ (14080 wide, the gate at column 7168, hidden 6144) at decode and
+ prefill row counts; its MTP layer's gate as the strided second half of
+ each head of a `[T, 24, 512]` q_proj.
+ - The default target's trunk
+ (`qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default`):
+ `DecodeTarget._attention_block` (decode, 1..128 rows at tp4 / tep4,
+ 1..32 at a dep4 rank), its MTP verify (up to 128 x 8 = 1024 rows, tep4,
+ draft 7) and `PrefillTarget._attention_block` (one chunk of up to 8192
+ rows, `max_num_tokens`). `gate` is the 2-D window of the fused input
+ projection row at column qkv_width: at tp4 / tep4 the row is 4352 wide
+ (q 1536 | k 256 | v 256 | gate 1536 | index scratch 768), the gate at
+ column 2048 and hidden 1536; at dep4 tp1's 14080-wide row, gate at
+ 7168, hidden 6144.
+ - The default target's MTP layer, head-sharded at tep4 (and tp4): q_proj
+ 12288 -> 3072 columns a rank, read as `[T, 6, 512]`, `gate` the second
+ 256 of each head (token stride 3072, head stride 512), hidden 1536.
+ Rows: the drafter 1..128, the depth-3 verify 128 x 4 = 512, the depth-7
+ verify 128 x 8 = 1024, a prefill chunk of up to 8192
+ (`max_num_tokens`).
+ Every row of the default target's 6-head forms ends in masked lanes (see
+ the cells); the test places `attention_output` at the head of a larger
+ buffer and holds the elements past its last row unchanged.
+
+ Preconditions, every one enforced loudly by the op's own bare `assert`
+ (an AssertionError with no message): `attention_output` 2-D with
+ `stride(-1) == 1` (row stride free); `gate` 2-D of `attention_output`'s
+ shape, or 3-D with `shape[0] == T` and `shape[1] * shape[2] == hidden`,
+ `stride(-1) == 1`. A CPU tensor raises ValueError from Triton (measured).
+ Hence no `is_valid`. Cross-device operands were not probed.
+
+ The pre-call contents of `attention_output` are destroyed; the gate is
+ left bitwise unchanged.
+
+ Numerics, measured on GB300 at every cell: bf16 bitwise equal to the fp32
+ reference except at most ~2 in 10^6 elements one ulp off (`tl.sigmoid`'s
+ fast exp). Against a float64 product every cell sits at bf16 output
+ rounding, the op and the reference alike (relative L2 1.64e-3 to 1.69e-3).
+ Torch's default bf16 tolerance holds, nothing loosened. A wrong read lands
+ far outside it (1.6 to 3.8 max abs): the q half instead of the gate half,
+ the fused row's q window instead of its gate window, or at the 6-head 3-D
+ gate the q half of each head or the flat row with the head stride ignored;
+ the test holds the default band refusing the last two. So the column offset
+ and the per-head indexing are under test.
+
+ Certified: bf16 operands, a compact `attention_output`; the 2-D window of
+ the fused projection row at hidden 6144 (row stride 14080, gate column
+ 7168) at 32, 128 and 8192 rows, and at hidden 1536 (row stride 4352, gate
+ column 2048) at 1, 128, 1024 and 8192 rows; the strided q_proj half (3-D
+ gate) at 24 heads (`[T, 24, 512]`, hidden 6144) at 1 token, and at 6 heads
+ (a compact `[T, 6, 512]`, hidden 1536) at 1, 128, 512, 1024 and 8192 rows;
+ CUDA-graph capture and replay of the 1536-wide fused-row form at 8 rows and
+ of the 6-head 3-D gate at 1, 128 and 1024 rows.
+
+ Hard limits
+ T * row stride token offset computed in int64 boundary 2^63 not probed
+ in-row offset int32; < hidden (6144) at the 2-D form,
+ head * head stride + dim < 2 * hidden at the 3-D form
+ (largest run 8192 x 14080 = 1.2e8 elements)
+
+ Not certified: 0 tokens; rows past 8192; the 3-D gate at a token stride
+ other than its own q_proj's (heads x 2 x head_dim, e.g. a gate inside a
+ fused q|k|v row); the 24-head strided q_proj half above 1 token (a dep4
+ rank's MTP layer: drafter 1..32 rows and its verify); a strided
+ `attention_output`; a gate dtype other than `attention_output.dtype`;
+ fp16/fp32 operands; hidden sizes other than 1536 and 6144; CUDA-graph
+ capture of the 6144-wide forms; torch.compile.
+ """
+
+ def __call__(self, attention_output: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
+ torch.ops.trtllm.fused_sigmoid_mul_inplace(attention_output, gate)
+ return attention_output
+
+ def reference(self, attention_output: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
+ """What `attention_output` holds after the call; neither input is mutated.
+
+ `gate.reshape(T, hidden)` is the per-head `(t, i // D, i % D)` read, for
+ the 2-D and 3-D forms alike.
+ """
+ num_tokens, hidden = attention_output.shape
+ g = gate.reshape(num_tokens, hidden).float()
+ return (attention_output.float() * torch.sigmoid(g)).to(attention_output.dtype)
+
+
+fused_sigmoid_mul_inplace = FusedSigmoidMulInplace()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/gelu_erf.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/gelu_erf.py
new file mode 100644
index 000000000000..a31ca88bb4a3
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/gelu_erf.py
@@ -0,0 +1,106 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Non-gated exact (erf) GELU -- what `torch.nn.GELU()` is -- the vision patch
+merger's activation, which no config key names.
+
+Native torch (rung 4 of the sourcing ladder): one `torch.nn.functional.gelu(x)`
+call, `approximate="none"` -- one elementwise CUDA kernel, the same call the
+stock tower makes (`modeling_qwen3vl.py`'s merger `nn.GELU()`).
+`torch.ops.trtllm.moe_gelu` computes this function bitwise but is not a
+substitute: its `tile_idx_to_mn_limit` / `num_non_exiting_tiles` operands must
+be int32 *device* tensors (on the host it raises nothing and returns garbage,
+nan, or an illegal memory access), and it rejects fp32 and non-2-D inputs.
+FlashInfer's `activation` module has only gated forms (`gelu_and_mul`,
+`gelu_tanh_and_mul`, `silu_and_mul`), as has
+`torch.ops.trtllm.flashinfer_gelu_tanh_and_mul`.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn.functional as F
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+#: `merger.linear_fc1` is `[4608, 4608]`: spatial merge 2 over a 1152-wide tower.
+MERGER_WIDTH = 4608
+
+_SQRT_2 = math.sqrt(2.0)
+
+
+class GeluErf(OpWrapper):
+ """`0.5 * x * (1 + erf(x / sqrt(2)))`, elementwise, any shape, new tensor
+ out in `x`'s dtype.
+
+ Fusion boundary: the activation only. `merger.linear_fc1` before it and
+ `merger.linear_fc2` after it are the caller's; there is no gate half and no
+ multiply. This is the erf GELU and not the tanh one (`GeluTanh`, the block
+ MLPs' `hidden_act`); both appear in this one tower.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: shape, dtype; `dist` "uniform" over [lo, hi] (default) or "merger"
+ # (a tower hidden state scaled by the merger norm's measured gamma, grouped
+ # 4-to-1, through an fc1 at the checkpoint's measured weight/bias scale).
+ 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="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),
+ ),
+ Cell(
+ why="the merger-built activand at the tower call ceiling: 8192 merged rows x 4608",
+ spec=dict(shape=(8192, MERGER_WIDTH), dtype=torch.bfloat16, dist="merger"),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none. Out-of-place; `x` is not mutated.
+
+ Certified on contiguous bf16 `[rows, 4608]` activands (what the merger's fc1
+ returns) at 1, 512 and 8192 rows (8192, from 32768 patch rows, the most one
+ tower call carries); 512 and 8192 rows also on the merger-built
+ distribution. The default target's call site is the vision tower's patch
+ merger, eager and rank-invariant: every rank runs the whole tower at tp4,
+ tep4 and dep4, so the shapes are tp1's.
+
+ Measured on GB300 against `reference` (fp64, rounded once to bf16): the
+ merger-built activand is bitwise equal or nearly so; the uniform [-16, 16]
+ spread differs in ~12% of elements. Every difference is at most 2.4e-7
+ absolute, in the negative tail where the output is near zero. One
+ compiled-in torch elementwise kernel per call: no JIT path.
+
+ Identity is pinned by fp32 only (the identity test): max |erf GELU - tanh
+ GELU| is 4.73e-4 (x = -2.699), 31x outside the default fp32 gate. In bf16
+ one ulp at the output scale exceeds that gap, so a bf16 match does not say
+ which GELU was computed.
+
+ Measured, not certified: `-inf -> nan` (the closed form does the same);
+ `+inf -> +inf`; large negatives give `-0.0`; `finfo.max` comes back
+ unchanged; CUDA-graph capturable; the output is a fresh allocation.
+
+ Not certified: fp32 beyond the identity test, non-contiguous inputs, fp16,
+ other widths, float8 operands. No index-typed argument;
+ numel past 2^31 not probed.
+ """
+
+ def __call__(self, x: torch.Tensor) -> torch.Tensor:
+ return F.gelu(x)
+
+ def reference(self, x: torch.Tensor) -> torch.Tensor:
+ """The closed form in fp64 from primitive ops, rounded once to `x.dtype`."""
+ x64 = x.double()
+ return (0.5 * x64 * (1.0 + torch.erf(x64 / _SQRT_2))).to(x.dtype)
+
+
+gelu_erf = GeluErf()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/gelu_tanh.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/gelu_tanh.py
new file mode 100644
index 000000000000..e28a9c0b4d6b
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/gelu_tanh.py
@@ -0,0 +1,116 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Non-gated tanh GELU -- HuggingFace's `gelu_pytorch_tanh` -- the vision
+tower's block-MLP activation.
+
+Tier 4 (native torch), a macro entry of one call: `torch.nn.functional.gelu(x,
+approximate="tanh")`, the aten elementwise kernel. No registered op computes a
+plain non-gated GELU over an unquantized tensor: `flashinfer_gelu_tanh_and_mul` (and FlashInfer's
+`gelu_tanh_and_mul` / `gelu_and_mul`) is the gated form, and silently returns a
+half-width `gelu(a) * b` when handed this tower's non-gated activand.
+`moe_gelu` is the erf GELU behind MoE tile metadata, and the
+`cute_dsl_nvfp4_dense_gemm_gelu*` ops are NVFP4 GEMM epilogues (this tower's
+fc1 is bf16).
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn.functional as F
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+#: The block MLP's `linear_fc1` is `[4304, 1152]`: one width, nothing split.
+INTERMEDIATE_SIZE = 4304
+
+_SQRT_2_OVER_PI = math.sqrt(2.0 / math.pi)
+
+
+class GeluTanh(OpWrapper):
+ """`0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x**3)))`, elementwise,
+ any shape, new tensor out in `x`'s dtype.
+
+ Fusion boundary: the activation only. `linear_fc1` before it and
+ `linear_fc2` after it are the caller's; there is no gate half and no
+ multiply. This is the tanh GELU and not the erf one (`GeluErf`, the patch
+ merger's); which one a module wants is a property of the checkpoint.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: shape, dtype; `dist` "uniform" over [lo, hi] (default) or "fc1" (a
+ # unit-normal hidden state through an fc1 at the checkpoint's measured
+ # weight std 0.0197 / bias std 1.676).
+ CELLS: tuple[Cell, ...] = (
+ 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="fc1-scaled activand at the tower call ceiling, 32768 patch rows",
+ spec=dict(shape=(32768, INTERMEDIATE_SIZE), dtype=torch.bfloat16, dist="fc1"),
+ ),
+ Cell(
+ why="block MLP activand at the tower call ceiling, 32768 patch rows x 4304 bf16",
+ spec=dict(shape=(32768, INTERMEDIATE_SIZE), dtype=torch.bfloat16, lo=-16.0, hi=16.0),
+ ),
+ Cell(
+ why="block MLP activand, 1 patch row x 4304 bf16",
+ spec=dict(shape=(1, INTERMEDIATE_SIZE), dtype=torch.bfloat16, lo=-16.0, hi=16.0),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none. Out-of-place; `x` is not mutated. Call site: the
+ vision tower's block MLP, `gelu_tanh(cublas_mm(layer_norm(h), fc1, fc1_b))`,
+ eager only, at most 32768 patch rows a call (`_VISION_CALL_MAX_PATCHES`).
+ The tower is rank-invariant, so every member (tp4, tep4, dep4) calls it at
+ the unsplit width 4304.
+
+ Source: the aten CUDA kernel compiled into torch (no JIT path), one launch a
+ call, computing in fp32 and rounding once to bf16.
+
+ Certified on contiguous bf16 `[rows, 4304]` activands (what the block MLP's
+ fc1 returns) at 1, 512 and 32768 rows (32768, the tower call ceiling):
+ fc1-scaled activands (x std ~1.8, |x| <= 8.8) at 512 and 32768 rows,
+ uniform [-16, 16] at 1 and 32768. Gate: torch's bf16 default against the
+ fp64 reference.
+
+ Error, measured on GB300 at every cell: within one bf16 ulp of the reference
+ everywhere except x in [-7.19, -4.56], where fp32 cancellation in `1 + tanh`
+ leaves the output up to 254 ulp of its own size off (abs error <= 7.6e-8 at
+ |ref| <= 3.6e-6). HF's `gelu_pytorch_tanh` is this same call. x >= 0 is
+ bitwise the reference. Against the unrounded fp64 form the fc1 cells sit at
+ bf16 output rounding (relative L2 1.4e-3, max abs 7.755e-3: half an ulp at
+ |y| in [2, 4)).
+
+ Identity is pinned by fp32 only (the identity test): max |tanh GELU - erf
+ GELU| is 4.73e-4 (|x| ~ 2.7), 31x outside the default fp32 gate. In bf16 one
+ ulp at the output scale exceeds that gap, so a bf16 match does not say which
+ GELU was computed.
+
+ Measured, not certified: `-inf -> nan` (the closed form does the same: -inf
+ times a factor saturated to 0); `+inf -> +inf`; large negatives give
+ `-0.0`; `finfo.max` comes back unchanged in bf16/fp16/fp32 (the cube is not
+ taken in the input dtype); CUDA-graph capturable, replay bitwise eager at 1,
+ 512 and 32768 rows (the call site is eager); the output is a fresh
+ allocation.
+
+ Not certified: fp32 beyond the identity test, non-contiguous inputs, fp16,
+ other widths, float8 operands, row counts between the cells. No index-typed
+ argument; numel past 2^31 (498951 rows of 4304) not probed.
+ """
+
+ def __call__(self, x: torch.Tensor) -> torch.Tensor:
+ return F.gelu(x, approximate="tanh")
+
+ def reference(self, x: torch.Tensor) -> torch.Tensor:
+ """The closed form in fp64 from primitive ops, rounded once to `x.dtype`."""
+ x64 = x.double()
+ inner = _SQRT_2_OVER_PI * (x64 + 0.044715 * x64 * x64 * x64)
+ return (0.5 * x64 * (1.0 + torch.tanh(inner))).to(x.dtype)
+
+
+gelu_tanh = GeluTanh()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/silu_and_mul_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/silu_and_mul_rows.py
new file mode 100644
index 000000000000..74088ea716fd
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/silu_and_mul_rows.py
@@ -0,0 +1,106 @@
+# 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: a Triton kernel. `flashinfer_silu_and_mul` addresses row `t` at `t * 2 * inter`, so
+its input must be exactly the gate|up columns, contiguous; here they are the head of a wider
+row (the same GEMM also forms the per-token gate logit) and are read in place.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+#: This checkpoint's shared expert: intermediate 640, the fused GEMM output row 2 * 640 + 8.
+_INTER = 640
+_WIDTH = 2 * _INTER + 8
+_BLOCK = 1024
+
+
+@triton.jit
+def _silu_and_mul_rows_kernel(x_ptr, out_ptr, stride_x, INTER: tl.constexpr, BLOCK: tl.constexpr):
+ row = tl.program_id(0).to(tl.int64)
+ cols = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+ mask = cols < INTER
+ g = tl.load(x_ptr + row * stride_x + cols, mask=mask, other=0.0).to(tl.float32)
+ u = tl.load(x_ptr + row * stride_x + INTER + cols, mask=mask, other=0.0).to(tl.float32)
+ y = g / (1.0 + tl.exp(-g)) * u
+ tl.store(out_ptr + row * INTER + cols, y.to(out_ptr.dtype.element_ty), mask=mask)
+
+
+class SiluAndMulRows(OpWrapper):
+ """`out[t, i] = silu(x[t, i]) * x[t, inter + i]` for `i < inter`: the gate half first, as
+ `flashinfer_silu_and_mul` reads it, in fp32 and rounded once to `x`'s dtype. Columns of
+ `x` past `2 * inter` are not read. Returns a fresh contiguous `[T, inter]` tensor.
+
+ Fusion boundary: the activation only; the GEMM that forms `x` and the one that reads
+ `out` are the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="prefill chunk 16384 rows of the fused 1288-wide output", spec=dict(rows=16384)),
+ Cell(why="short prefill step, 128 rows of the fused output", spec=dict(rows=128)),
+ Cell(why="zero rows: no launch", spec=dict(rows=0)),
+ Cell(why="decode, one request: the shared expert's single row", spec=dict(rows=1)),
+ Cell(
+ why="decode, 32 rows: a dep4 rank's local share of a 128-request batch",
+ spec=dict(rows=32),
+ ),
+ Cell(why="MTP-7 verify: 128 requests x 8 rows", spec=dict(rows=1024)),
+ Cell(why="prefill chunk at max_num_tokens 8192", spec=dict(rows=8192)),
+ )
+
+ note = """
+ Metadata consumed: none; every row independent. Call sites: the decode and prefill
+ targets' MoE block, on the replicated shared expert's fused
+ `[gate | up | gate logit | pad]` GEMM output, so inter 640 and the 1288-wide row at
+ every member (tp4, tep4, dep4). Decode rows: 1..128 per rank at tp4/tep4, 1..32 local
+ at dep4, 1024 at an MTP-7 verify of 128 requests; prefill chunks of up to
+ max_num_tokens.
+
+ Preconditions: 2-D `x` with unit column stride, `x.shape[1] >= 2 * inter`; bf16 or fp16.
+
+ Source: Triton, JIT-compiled at first call for the running device; the cells ran that
+ path on sm_103 (GB300). Launch: `(rows, cdiv(inter, 1024))` programs of 4 warps, the
+ same kernel at every row count (no path switch). Row offsets are int64, so no index
+ product overflows inside the envelope (2^31 elements of the 1288-wide input is 1.67M
+ rows).
+
+ Numerics: `g / (1 + exp(-g)) * u` in fp32 (Triton's `exp` and division), rounded once;
+ bitwise `flashinfer_silu_and_mul` on the same values at every non-empty cell, including
+ 0, -0, 88 and -100.
+
+ Certified: bf16 at inter 640 on the 1288-wide fused output, 0, 1, 32, 128, 1024, 8192
+ and 16384 rows; CUDA-graph replay at 1, 32 and 128 rows. Not certified: fp16, other
+ widths; row counts between the cells run the same kernel but were not driven.
+ """
+
+ def __call__(self, x: torch.Tensor, inter: int) -> torch.Tensor:
+ 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/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/__init__.py
new file mode 100644
index 000000000000..52a7a9daf028
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/__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/ops/attention/qsa_canonical_selection.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_canonical_selection.py
new file mode 100644
index 000000000000..7ec7256b8eff
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_canonical_selection.py
@@ -0,0 +1,608 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""QSA block selection, made canonical: one answer per row, whichever tied columns the
+selection engine happened to keep.
+
+Tier 5 (Triton). No registered op rewrites a selection, and none
+of the top-k engines (`indexer_topk_prefill`, the CuTe DSL kernels, FlashInfer's top-k)
+takes a tie-break rule or fixes the order of its output. A row's k-th value is read off
+the engine's picks; the window is counted against it once and compacted into the row in
+a second, early-exiting pass. Only the window's columns are read, and nothing
+`[rows, cols]`-sized is allocated.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+#: This checkpoint's block budget: `indexer_budget` 2048 tokens over
+#: `indexer_compress_ratio` 4.
+_BLOCK_TOPK = 512
+#: Score-row width the target hands over: a 260-page block table at 32 tokens
+#: per block, compressed 4:1.
+_SCORE_COLUMNS = 260 * 32 // 4
+#: The score-row width at the 256k-token context: 65536 compressed blocks.
+_LONG_SCORE_COLUMNS = 262144 // 4
+#: From this many rows on, one program per row keeps the GPU busy and rereads its window
+#: from L2 (`_canonical_rows_kernel`). Below it a row's window is cut into segments of
+#: `_SEGMENT` columns, each its own program, so a decode-sized launch is not one program
+#: streaming a 256 KiB row.
+_SPLIT_ROWS = 1024
+_SEGMENT = 4096
+#: Columns per step of the counting pass, of the writing pass (whose in-block prefix sum
+#: costs more per column), and warps per program, tuned on GB300 for both launch shapes.
+_ROWS_COUNT_BLOCK = 2048
+_SPLIT_COUNT_BLOCK = 4096
+_WRITE_BLOCK = 1024
+_NUM_WARPS = 4
+
+
+def _cell(why: str, rows, windows, logits: str, tail: str = "nan", **extra) -> Cell:
+ return Cell(
+ why=why,
+ spec=dict(
+ rows=rows,
+ cols=extra.pop("cols", _SCORE_COLUMNS),
+ index_topk=_BLOCK_TOPK,
+ windows=windows,
+ logits=logits,
+ tail=tail,
+ selection=extra.pop("selection", "random"),
+ **extra,
+ ),
+ )
+
+
+@triton.jit
+def _kth_value(
+ row_logits,
+ row_indices,
+ start,
+ num_cols,
+ stride_col,
+ stride_index,
+ K: tl.constexpr,
+ K_P2: tl.constexpr,
+):
+ """The smallest score the engine kept, read at the clamped column `start + pick`: an
+ empty slot (`pick < 0`) reads +inf, and a NaN pick makes it NaN, which no score
+ compares to, so the row comes out empty."""
+ slots = tl.arange(0, K_P2)
+ in_k = slots < K
+ picks = tl.load(row_indices + slots * stride_index, mask=in_k, other=-1)
+ columns = tl.minimum(tl.maximum(picks + start, 0), num_cols - 1)
+ picked = tl.load(
+ row_logits + columns.to(tl.int64) * stride_col,
+ mask=in_k & (picks >= 0),
+ other=float("inf"),
+ ).to(tl.float32)
+ kth = tl.min(picked, axis=0)
+ picked_nan = tl.max((picked != picked).to(tl.int32), axis=0)
+ return tl.where(picked_nan > 0, float("nan"), kth)
+
+
+@triton.jit
+def _load_window(row_logits, cols, lo, hi, stride_col):
+ # NaN past the window: it compares false, so it is never counted or kept.
+ return tl.load(
+ row_logits + cols.to(tl.int64) * stride_col,
+ mask=(cols >= lo) & (cols < hi),
+ other=float("nan"),
+ ).to(tl.float32)
+
+
+@triton.jit
+def _count(x, kth):
+ """Columns of `x` above the k-th value and tied with it (float equality: -0.0 == +0.0,
+ NaN equals nothing)."""
+ return tl.sum((x > kth).to(tl.int32), axis=0), tl.sum((x == kth).to(tl.int32), axis=0)
+
+
+@triton.jit
+def _write_block(
+ row_indices,
+ stride_index,
+ cols,
+ start,
+ x,
+ kth,
+ drop,
+ kept,
+ ties_seen,
+ K: tl.constexpr,
+ WRITE_BLOCK: tl.constexpr,
+):
+ """Compact one block of the window into the row; returns the updated running counts.
+
+ Kept: every column above the k-th value, and every tie whose 0-based rank among the
+ window's ties is at least `drop` -- the lowest-index `drop` ties are the surplus. The
+ j-th kept column, ascending, goes to slot j as a window-local offset; slots from k on
+ are not written. One prefix sum carries both counts, above in the high half-word."""
+ tl.static_assert(WRITE_BLOCK < 65536)
+ is_above = (x > kth).to(tl.int32)
+ is_tied = (x == kth).to(tl.int32)
+ packed = is_above * 65536 + is_tied
+ before = tl.cumsum(packed, axis=0) - packed
+ ties_before = ties_seen + (before & 65535)
+ kept_ties = tl.maximum(ties_seen - drop, 0)
+ slot = kept + (before >> 16) + tl.maximum(ties_before - drop, 0) - kept_ties
+ keep = (is_above != 0) | ((is_tied != 0) & (ties_before >= drop))
+ tl.store(row_indices + slot * stride_index, cols - start, mask=keep & (slot < K))
+ block = tl.sum(packed, axis=0)
+ ties_seen += block & 65535
+ kept += (block >> 16) + tl.maximum(ties_seen - drop, 0) - kept_ties
+ return kept, ties_seen
+
+
+@triton.jit
+def _canonical_rows_kernel(
+ logits_ptr,
+ starts_ptr,
+ ends_ptr,
+ indices_ptr,
+ num_cols,
+ stride_starts,
+ stride_ends,
+ stride_logits_row,
+ stride_logits_col,
+ stride_indices_row,
+ stride_indices_col,
+ K: tl.constexpr,
+ K_P2: tl.constexpr,
+ COUNT_BLOCK: tl.constexpr,
+ WRITE_BLOCK: tl.constexpr,
+):
+ # One program per row: both passes, the second reading the window back from L2. Each
+ # loop loads its next block before reducing the current one. Every pick is read
+ # before the k-th value's reduction completes and every store depends on it, so the
+ # in-place rewrite never overwrites a pick another thread has yet to read.
+ row = tl.program_id(0).to(tl.int64)
+ row_logits = logits_ptr + row * stride_logits_row
+ row_indices = indices_ptr + row * stride_indices_row
+ start = tl.load(starts_ptr + row * stride_starts).to(tl.int32)
+ end = tl.load(ends_ptr + row * stride_ends).to(tl.int32)
+ kth = _kth_value(
+ row_logits, row_indices, start, num_cols, stride_logits_col, stride_indices_col, K, K_P2
+ )
+ # The window's columns inside the matrix, in blocks aligned to the block size so the
+ # loads vectorize.
+ lo = tl.maximum(start, 0)
+ hi = tl.minimum(end, num_cols)
+
+ above = tl.zeros((), dtype=tl.int32)
+ tied = tl.zeros((), dtype=tl.int32)
+ first = (lo // COUNT_BLOCK) * COUNT_BLOCK
+ x = _load_window(row_logits, first + tl.arange(0, COUNT_BLOCK), lo, hi, stride_logits_col)
+ for c0 in range(first, hi, COUNT_BLOCK):
+ c0 = tl.multiple_of(c0, COUNT_BLOCK)
+ cols = c0 + COUNT_BLOCK + tl.arange(0, COUNT_BLOCK)
+ x_next = _load_window(row_logits, cols, lo, hi, stride_logits_col)
+ block_above, block_tied = _count(x, kth)
+ above += block_above
+ tied += block_tied
+ x = x_next
+ # A window of at most k has no surplus.
+ drop = tl.maximum(above + tied - K, 0)
+
+ kept = tl.zeros((), dtype=tl.int32)
+ ties_seen = tl.zeros((), dtype=tl.int32)
+ c0 = (lo // WRITE_BLOCK) * WRITE_BLOCK
+ x = _load_window(row_logits, c0 + tl.arange(0, WRITE_BLOCK), lo, hi, stride_logits_col)
+ while (c0 < hi) & (kept < K):
+ c0 = tl.multiple_of(c0, WRITE_BLOCK)
+ cols = c0 + tl.arange(0, WRITE_BLOCK)
+ x_next = _load_window(row_logits, cols + WRITE_BLOCK, lo, hi, stride_logits_col)
+ kept, ties_seen = _write_block(
+ row_indices,
+ stride_indices_col,
+ cols,
+ start,
+ x,
+ kth,
+ drop,
+ kept,
+ ties_seen,
+ K,
+ WRITE_BLOCK,
+ )
+ x = x_next
+ c0 += WRITE_BLOCK
+ slots = tl.arange(0, K_P2)
+ tl.store(row_indices + slots * stride_indices_col, -1, mask=(slots < K) & (slots >= kept))
+
+
+@triton.jit
+def _split_kth_kernel(
+ logits_ptr,
+ starts_ptr,
+ indices_ptr,
+ kth_ptr,
+ num_cols,
+ stride_starts,
+ stride_logits_row,
+ stride_logits_col,
+ stride_indices_row,
+ stride_indices_col,
+ K: tl.constexpr,
+ K_P2: tl.constexpr,
+):
+ # Its own launch: the segments' programs that write the row come after every pick of
+ # it has been read.
+ row = tl.program_id(0).to(tl.int64)
+ start = tl.load(starts_ptr + row * stride_starts).to(tl.int32)
+ kth = _kth_value(
+ logits_ptr + row * stride_logits_row,
+ indices_ptr + row * stride_indices_row,
+ start,
+ num_cols,
+ stride_logits_col,
+ stride_indices_col,
+ K,
+ K_P2,
+ )
+ tl.store(kth_ptr + row, kth)
+
+
+@triton.jit
+def _split_count_kernel(
+ logits_ptr,
+ starts_ptr,
+ ends_ptr,
+ kth_ptr,
+ counts_ptr,
+ num_rows,
+ num_cols,
+ num_segments,
+ stride_starts,
+ stride_ends,
+ stride_logits_row,
+ stride_logits_col,
+ SEGMENT: tl.constexpr,
+ COUNT_BLOCK: tl.constexpr,
+):
+ # One program per (row, segment): the segment's window columns above and tied with
+ # the k-th value, into `counts[0 | 1, row, segment]`.
+ row = tl.program_id(0).to(tl.int64)
+ segment = tl.program_id(1)
+ row_logits = logits_ptr + row * stride_logits_row
+ seg_lo = segment * SEGMENT
+ start = tl.load(starts_ptr + row * stride_starts).to(tl.int32)
+ end = tl.load(ends_ptr + row * stride_ends).to(tl.int32)
+ lo = tl.maximum(tl.maximum(start, 0), seg_lo)
+ hi = tl.minimum(tl.minimum(end, num_cols), seg_lo + SEGMENT)
+ kth = tl.load(kth_ptr + row)
+ above = tl.zeros((), dtype=tl.int32)
+ tied = tl.zeros((), dtype=tl.int32)
+ for c0 in range(seg_lo, hi, COUNT_BLOCK):
+ c0 = tl.multiple_of(c0, COUNT_BLOCK)
+ x = _load_window(row_logits, c0 + tl.arange(0, COUNT_BLOCK), lo, hi, stride_logits_col)
+ block_above, block_tied = _count(x, kth)
+ above += block_above
+ tied += block_tied
+ tl.store(counts_ptr + row * num_segments + segment, above)
+ tl.store(counts_ptr + (num_rows + row) * num_segments + segment, tied)
+
+
+@triton.jit
+def _split_write_kernel(
+ logits_ptr,
+ starts_ptr,
+ ends_ptr,
+ kth_ptr,
+ counts_ptr,
+ indices_ptr,
+ num_rows,
+ num_cols,
+ num_segments,
+ stride_starts,
+ stride_ends,
+ stride_logits_row,
+ stride_logits_col,
+ stride_indices_row,
+ stride_indices_col,
+ K: tl.constexpr,
+ K_P2: tl.constexpr,
+ SEGMENT: tl.constexpr,
+ WRITE_BLOCK: tl.constexpr,
+ COUNTS_BLOCK: tl.constexpr,
+):
+ # One program per (row, segment): the row's totals and the counts of the segments
+ # before this one place the segment's kept columns; segment 0 also writes the `-1`s.
+ row = tl.program_id(0).to(tl.int64)
+ segment = tl.program_id(1)
+ row_logits = logits_ptr + row * stride_logits_row
+ row_indices = indices_ptr + row * stride_indices_row
+ start = tl.load(starts_ptr + row * stride_starts).to(tl.int32)
+ end = tl.load(ends_ptr + row * stride_ends).to(tl.int32)
+ seg_lo = segment * SEGMENT
+ lo = tl.maximum(tl.maximum(start, 0), seg_lo)
+ hi = tl.minimum(tl.minimum(end, num_cols), seg_lo + SEGMENT)
+
+ above = tl.zeros((), dtype=tl.int32)
+ tied = tl.zeros((), dtype=tl.int32)
+ above_before = tl.zeros((), dtype=tl.int32)
+ ties_before = tl.zeros((), dtype=tl.int32)
+ for j0 in range(0, num_segments, COUNTS_BLOCK):
+ j = j0 + tl.arange(0, COUNTS_BLOCK)
+ seg_above = tl.load(counts_ptr + row * num_segments + j, mask=j < num_segments, other=0)
+ seg_tied = tl.load(
+ counts_ptr + (num_rows + row) * num_segments + j, mask=j < num_segments, other=0
+ )
+ above += tl.sum(seg_above, axis=0)
+ tied += tl.sum(seg_tied, axis=0)
+ above_before += tl.sum(tl.where(j < segment, seg_above, 0), axis=0)
+ ties_before += tl.sum(tl.where(j < segment, seg_tied, 0), axis=0)
+ drop = tl.maximum(above + tied - K, 0)
+ if segment == 0:
+ total = above + tl.maximum(tied - drop, 0)
+ slots = tl.arange(0, K_P2)
+ tl.store(row_indices + slots * stride_indices_col, -1, mask=(slots < K) & (slots >= total))
+
+ kth = tl.load(kth_ptr + row)
+ kept = above_before + tl.maximum(ties_before - drop, 0)
+ ties_seen = ties_before
+ c0 = seg_lo
+ while (c0 < hi) & (kept < K):
+ c0 = tl.multiple_of(c0, WRITE_BLOCK)
+ cols = c0 + tl.arange(0, WRITE_BLOCK)
+ x = _load_window(row_logits, cols, lo, hi, stride_logits_col)
+ kept, ties_seen = _write_block(
+ row_indices,
+ stride_indices_col,
+ cols,
+ start,
+ x,
+ kth,
+ drop,
+ kept,
+ ties_seen,
+ K,
+ WRITE_BLOCK,
+ )
+ c0 += WRITE_BLOCK
+
+
+class QsaCanonicalSelection(OpWrapper):
+ """Rewrite an exact top-k selection, in place, into the window's canonical selection.
+
+ For row `r` with window `[s, e) = [row_starts[r], row_ends[r])` of `logits` and
+ `k = indices.shape[1]`, `indices[r]` holds on entry what a top-k engine wrote:
+ window-local offsets of an exact top-k of the window, in any order, `-1` past
+ `min(k, e - s)`. Let `v` be the smallest selected score. On return the row is
+
+ * every window column scoring strictly above `v`, then the highest-index window
+ columns scoring exactly `v` until there are `k` -- the k largest
+ `(score, column)` pairs; or the whole window when `e - s <= k`;
+ * as window-local offsets in ascending order, `-1` in every remaining slot.
+
+ So the row is a function of `(logits, row_starts, row_ends, k)` alone: two engines,
+ or two launches of one engine, that broke a tie at the k-th value differently produce
+ the same row. On a row with no tie at the k-th value this is the engine's own set,
+ sorted. Ties are float equality: `-0.0` and `+0.0` are one tie class. Returns
+ `indices`.
+
+ Fusion boundary: the rewrite only. Selecting is the engine's call before it; the
+ engine must be exact -- `v` is read off its picks, and a pick below the true k-th
+ value would admit more than `k` columns above it.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # The target calls this right after `qsa_topk_select`, at `k = 512` over the
+ # `[rows, cols]` score matrix whose columns past each `row_end` are allocator
+ # garbage (the cells poison them). `selection` names what produced the input row:
+ # `random` is an exact top-k that keeps a random subset of the tie class at the k-th
+ # value, in a random order (the adversary); `radix` is the engine the target calls.
+ CELLS: tuple[Cell, ...] = (
+ _cell(
+ "decode: 1, 128 and 256 rows (MTP3 and MTP7 at batch 32)",
+ (1, 128, 256),
+ ("positions_random", 2051, 8191),
+ "relu",
+ ),
+ _cell(
+ "prefill: 8192 rows from position 0",
+ 8192,
+ ("chunk", 0),
+ "relu",
+ ),
+ _cell(
+ "input from the radix engine (qsa_topk_select), zero ties",
+ 64,
+ ("positions_random", 2051, 8191),
+ "relu_sparse",
+ selection="radix",
+ ),
+ _cell(
+ "decode at up to 256k context: 65536-wide table, relu-zero ties",
+ (1, 16, 64),
+ ("positions_random", 2051, 262143),
+ "relu_sparse",
+ cols=_LONG_SCORE_COLUMNS,
+ ),
+ _cell(
+ "prefill: a 2048-row chunk ending at the 256k context",
+ 2048,
+ ("chunk", 262144 - 2048),
+ "relu",
+ cols=_LONG_SCORE_COLUMNS,
+ ),
+ )
+
+ note = """
+ Metadata consumed: none. Stateless; every input is an argument, and `indices` is
+ both the engine's selection read and the canonical row written.
+
+ Preconditions (none checked):
+ - `indices` holds an EXACT top-k of each window: the smallest selected score is
+ the window's k-th largest. `qsa_topk_select` is exact; an approximate selection
+ would be rewritten into a row with more than `k` columns at or above its k-th
+ pick, truncated to the lowest `k` column indices -- wrong, silently.
+ - `0 <= row_starts[r] <= row_ends[r] <= logits.shape[1]`. Anything outside the
+ window is never read for a decision, so garbage, NaN and +inf there are fine
+ (the cells poison it).
+ - Scores inside the window are not NaN: a NaN column is never kept, and a NaN pick
+ empties the row. (The engine returns the identity for a window of at most k
+ without reading it; this rewrite reads it.)
+
+ Cost: no host read, no data-dependent shape, so it can be captured into a CUDA graph
+ (the entry's test replays one). A row reads its k picks, its window once to count
+ and again, up to its k-th kept column, to write; columns outside the window are not
+ read, so narrowing `logits` to a bound on the windows (`logits[:, :bound]`) gives
+ the same rows and saves nothing. From 1024 rows: one launch, one program per row.
+ Below: three launches -- the k-th values, then counts and writes over 4096-column
+ segments, each its own program -- through `[rows]` fp32 and `[2, rows, cols / 4096]`
+ int32 scratch.
+ """
+
+ def __call__(
+ self,
+ logits: torch.Tensor,
+ row_starts: torch.Tensor,
+ row_ends: torch.Tensor,
+ indices: torch.Tensor,
+ ) -> torch.Tensor:
+ num_rows, num_cols = logits.shape
+ index_topk = indices.shape[1]
+ if num_rows == 0 or index_topk == 0:
+ return indices
+ # A one-element window broadcasts over every row -- the hazard `is_valid`
+ # refuses -- as a stride-0 view rather than a read past the tensor's end.
+ row_starts = row_starts.expand(num_rows)
+ row_ends = row_ends.expand(num_rows)
+ k_p2 = triton.next_power_of_2(index_topk)
+ if num_rows >= _SPLIT_ROWS:
+ _canonical_rows_kernel[(num_rows,)](
+ logits,
+ row_starts,
+ row_ends,
+ indices,
+ num_cols,
+ row_starts.stride(0),
+ row_ends.stride(0),
+ logits.stride(0),
+ logits.stride(1),
+ indices.stride(0),
+ indices.stride(1),
+ K=index_topk,
+ K_P2=k_p2,
+ COUNT_BLOCK=_ROWS_COUNT_BLOCK,
+ WRITE_BLOCK=_WRITE_BLOCK,
+ num_warps=_NUM_WARPS,
+ )
+ return indices
+ num_segments = triton.cdiv(num_cols, _SEGMENT)
+ kth = torch.empty(num_rows, dtype=torch.float32, device=logits.device)
+ counts = torch.empty(2, num_rows, num_segments, dtype=torch.int32, device=logits.device)
+ _split_kth_kernel[(num_rows,)](
+ logits,
+ row_starts,
+ indices,
+ kth,
+ num_cols,
+ row_starts.stride(0),
+ logits.stride(0),
+ logits.stride(1),
+ indices.stride(0),
+ indices.stride(1),
+ K=index_topk,
+ K_P2=k_p2,
+ num_warps=_NUM_WARPS,
+ )
+ _split_count_kernel[(num_rows, num_segments)](
+ logits,
+ row_starts,
+ row_ends,
+ kth,
+ counts,
+ num_rows,
+ num_cols,
+ num_segments,
+ row_starts.stride(0),
+ row_ends.stride(0),
+ logits.stride(0),
+ logits.stride(1),
+ SEGMENT=_SEGMENT,
+ COUNT_BLOCK=_SPLIT_COUNT_BLOCK,
+ num_warps=_NUM_WARPS,
+ )
+ _split_write_kernel[(num_rows, num_segments)](
+ logits,
+ row_starts,
+ row_ends,
+ kth,
+ counts,
+ indices,
+ num_rows,
+ num_cols,
+ num_segments,
+ row_starts.stride(0),
+ row_ends.stride(0),
+ logits.stride(0),
+ logits.stride(1),
+ indices.stride(0),
+ indices.stride(1),
+ K=index_topk,
+ K_P2=k_p2,
+ SEGMENT=_SEGMENT,
+ WRITE_BLOCK=_WRITE_BLOCK,
+ COUNTS_BLOCK=64,
+ num_warps=_NUM_WARPS,
+ )
+ return indices
+
+ def reference(
+ self,
+ logits: torch.Tensor,
+ row_starts: torch.Tensor,
+ row_ends: torch.Tensor,
+ indices: torch.Tensor,
+ ) -> torch.Tensor:
+ """The k largest `(score, column)` pairs of each window, ascending by column,
+ `-1`-padded: a top-k over one int64 key per column whose high word orders the
+ scores as fp32 does and whose low word is the column. `indices` supplies only the
+ width `k`; its contents are ignored."""
+ num_rows, num_cols = logits.shape
+ index_topk = indices.shape[1]
+ columns = torch.arange(num_cols, device=logits.device)[None, :]
+ starts, ends = row_starts.long()[:, None], row_ends.long()[:, None]
+ in_window = (columns >= starts) & (columns < ends)
+ # `+ 0.0` folds -0.0 into +0.0, so the two are one tie class.
+ bits = (logits.float() + 0.0).view(torch.int32).long()
+ order = torch.where(bits < 0, bits ^ 0x7FFFFFFF, bits)
+ key = torch.where(in_window, (order << 32) | columns, torch.iinfo(torch.int64).min)
+ width = min(index_topk, num_cols)
+ top = torch.topk(key, width, dim=1).indices
+ taken = torch.gather(in_window, 1, top)
+ pad = torch.iinfo(torch.int64).max
+ ordered = torch.sort(torch.where(taken, top, pad), dim=1).values
+ out = torch.where(ordered == pad, -1, ordered - starts).to(torch.int32)
+ if width < index_topk:
+ out = torch.nn.functional.pad(out, (0, index_topk - width), value=-1)
+ return out
+
+ def is_valid(
+ self,
+ logits: torch.Tensor,
+ row_starts: torch.Tensor,
+ row_ends: torch.Tensor,
+ indices: torch.Tensor,
+ ) -> None:
+ # A one-element `row_starts` / `row_ends`, or a one-row `indices`, broadcasts over
+ # every row of `logits` without an error: a [1] `row_starts` against 4 rows
+ # returns every row framed by row 0's start.
+ num_rows = logits.shape[0]
+ assert row_starts.shape == (num_rows,) and row_ends.shape == (num_rows,), (
+ f"row_starts and row_ends need one entry per logits row ({num_rows}); got "
+ f"{tuple(row_starts.shape)} and {tuple(row_ends.shape)}"
+ )
+ assert indices.dim() == 2 and indices.shape[0] == num_rows, (
+ f"indices must be [{num_rows}, k]; got {tuple(indices.shape)}"
+ )
+
+
+qsa_canonical_selection = QsaCanonicalSelection()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_expand_block_indices.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_expand_block_indices.py
new file mode 100644
index 000000000000..2133bf12c615
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_expand_block_indices.py
@@ -0,0 +1,221 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""QSA Top-K group ids expanded into raw KV token indices, plus the causal tail."""
+
+from __future__ import annotations
+
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+# Tier-2 sourcing: a vendored Triton callable reached as ordinary Python. QSA's
+# kernel plane has no torch.ops surface.
+from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels as _qsa_kernels
+
+#: This checkpoint's selection geometry: `indexer_compress_ratio` and
+#: `indexer_budget`, so `block_topk` 512 and an output 2051 wide.
+_CKPT = dict(compress_ratio=4, token_topk=2048)
+
+
+#: First positions of 32 MTP verify blocks in the sparse branch (KV past 2048).
+_VERIFY_STARTS = tuple(2100 + (499 * g) % 6000 for g in range(32))
+
+#: The longest sequence the default target serves: `max_seq_len` 262144, with
+#: `max_num_tokens` 8192 and chunked prefill.
+_ENVELOPE_SEQ_LEN = 262144
+
+#: First positions of 128 MTP7 verify blocks, spread over the sparse branch up to
+#: the envelope (the last block ends at 262107).
+_VERIFY_STARTS_128 = tuple(2100 + (2039 * g) % 260000 for g in range(128))
+
+
+def _spec(batch, **overrides) -> dict:
+ """A cell spec at this checkpoint's geometry and the engine's metadata dtypes.
+
+ `batch` is one of:
+ ("sweep", n) n rows at 4096 + 13 i, length = position + 1
+ ("decode", n, top) n one-row requests spread over [1, top),
+ every tail width, length = position + 1
+ ("serving", requests, context_lens) one packed step: context chunks, then
+ two MTP rows per remaining request
+ ("verify", starts, p) one MTP verify block of p rows per start,
+ every row's length the block's end
+ ("chunk", start, n) one context chunk of n rows from `start`,
+ every row's length the chunk's end
+ """
+ base = dict(
+ batch=batch,
+ position_dtype=torch.int64, # qsa_logical_positions
+ length_dtype=torch.int32, # qsa_sequence_lengths
+ **_CKPT,
+ )
+ return {**base, **overrides}
+
+
+class QsaExpandBlockIndices(OpWrapper):
+ """Expand each row's selected compressed groups into token ids, then append
+ the incomplete group at the causal frontier.
+
+ Wraps `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels.
+ triton_expand_qsa_block_indices`: one Triton launch,
+ `_expand_qsa_block_indices_kernel`. Per row (one per query token), with
+ `final_topk = token_topk + compress_ratio - 1`:
+
+ 1. column `c < token_topk` holds member `c % cr` of group
+ `block_indices[row, c // cr]`; `-1` if that group is `-1` or the
+ token is `>= sequence_lengths[row]`;
+ 2. the last `(query_positions[row] + 1) % cr` visible tokens -- the
+ group the indexer never scored -- go in the columns right after the
+ selected groups, i.e. from column `(# non-negative groups) * cr`;
+ 3. everything else is `-1`.
+
+ Selection order is kept (nothing sorted, compacted or de-duplicated).
+ `query_positions` feeds the tail only; the expansion is masked by
+ `sequence_lengths` alone, so the causal bound on the *selection* is the
+ producer's. Returns a fresh contiguous int32 `[rows, final_topk]`; inputs
+ are read-only.
+
+ Fusion boundary: index arithmetic and the tail only. Scoring, Top-K, the
+ per-row causal group bound and the geometry are the caller's.
+
+ Called by the qwen4_exp target at every layout; its row counts and sequence
+ lengths are among the cells.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ 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))),
+ # ---- 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; 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))),
+ ),
+ # ---- the default target's decode, verify and prefill extents (tp4 / tep4 / dep4
+ # alike: one selection per query row, no head or rank axis).
+ Cell(
+ why="decode batch 128: the first whole-row launch, KV up to the 262144 envelope",
+ spec=_spec(("decode", 128, _ENVELOPE_SEQ_LEN)),
+ ),
+ Cell(
+ why="verify MTP7 at 128 requests: 1024 rows, KV up to the 262144 envelope",
+ spec=_spec(("verify", _VERIFY_STARTS_128, 8)),
+ ),
+ Cell(
+ why="context chunk of max_num_tokens 8192 rows ending a 262144-token prompt",
+ spec=_spec(("chunk", _ENVELOPE_SEQ_LEN - 8192, 8192)),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none -- every input is an argument. The target passes
+ `qsa_topk_indices[:rows]` (int32, written by `qsa_topk_select`),
+ `qsa_logical_positions` (int64) and `qsa_sequence_lengths` (int32: the
+ request's whole KV length broadcast over its rows, so inside a context
+ chunk it exceeds `position + 1` on every row but the last).
+
+ Preconditions that are values, so not `is_valid` guards (both 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
+ 12..15 lost. `qsa_topk_select` emits packed rows.
+ - selected groups are complete and visible (`g < (pos + 1) // cr`). An
+ incomplete group comes back with `-1` holes in place and its real
+ tokens repeated in the tail; a group past the query's horizon but
+ inside `sequence_lengths` is expanded verbatim.
+ Duplicate group ids expand to duplicate tokens.
+
+ 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.
+
+ Source: a Triton kernel JIT-compiled on first call for the device it runs
+ on; the cells were driven through that JIT path on sm_103 (GB300).
+
+ Launch: one program per (row, 256 columns) below 128 rows, one per row
+ from 128 up. No autotune. CUDA-graph capturable (captured at 1, 128, 256
+ and 1024 rows: both launch shapes); a replay recomputes from the live
+ buffers, and the captured row count fixes the launch shape.
+
+ Rank-invariant: one selection per query row, no head or rank axis, so tp4,
+ tep4 and dep4 call it with the same geometry; only the row count and the
+ positions move.
+
+ Hard limits (read from the kernel, not probed): both row strides are
+ `tl.constexpr` and the row offset is `program_id * stride` in int32. The
+ output's stride is `final_topk` (2051), so the store wraps first, at
+ `rows > 2^31 / 2051` (about 1.05 M rows); the input read wraps at
+ `rows * stride(0) >= 2^31` (4 Mi rows at width 512). Token ids are int32,
+ exact far past the envelope. The envelope reaches neither: `max_num_tokens`
+ 8192 rows, positions below `max_seq_len` 262144.
+ Certified extent: 8192 rows x 2051 columns, positions up to 262143.
+
+ Not certified: multiple devices, `torch.compile`, rows above 8192, and the
+ column-split launch at 2..127 rows (only its one-row shape is a cell;
+ decode batches 2..120 are captured there).
+ """
+
+ def __call__(
+ self,
+ block_indices: torch.Tensor,
+ query_positions: torch.Tensor,
+ sequence_lengths: torch.Tensor,
+ *,
+ compress_ratio: int,
+ token_topk: int,
+ ) -> torch.Tensor:
+ return _qsa_kernels.triton_expand_qsa_block_indices(
+ block_indices,
+ query_positions,
+ sequence_lengths,
+ compress_ratio=compress_ratio,
+ token_topk=token_topk,
+ )
+
+ def reference(
+ self,
+ block_indices: torch.Tensor,
+ query_positions: torch.Tensor,
+ sequence_lengths: torch.Tensor,
+ *,
+ compress_ratio: int,
+ token_topk: int,
+ ) -> torch.Tensor:
+ """The expansion as the contract states it, in int64 torch."""
+ rows, block_topk = block_indices.shape
+ final_topk = token_topk + compress_ratio - 1
+ device = block_indices.device
+ blocks = block_indices.long()
+ lengths = sequence_lengths.long()
+ offsets = torch.arange(compress_ratio, device=device)
+
+ expanded = (blocks[:, :, None] * compress_ratio + offsets).reshape(rows, token_topk)
+ chosen = (blocks >= 0).repeat_interleave(compress_ratio, dim=1)
+ keep = chosen & (expanded < lengths[:, None])
+ out = torch.full((rows, final_topk), -1, dtype=torch.int64, device=device)
+ out[:, :token_topk] = torch.where(keep, expanded, -1)
+
+ first_free = (blocks >= 0).sum(dim=1) * compress_ratio
+ visible = query_positions.long() + 1
+ tail_start = visible // compress_ratio * compress_ratio
+ row = torch.arange(rows, device=device)
+ for offset in range(compress_ratio - 1):
+ column = first_free + offset
+ token = tail_start + offset
+ live = (offset < visible - tail_start) & (token < lengths) & (column < final_topk)
+ out[row[live], column[live]] = token[live]
+ return out.to(torch.int32)
+
+
+qsa_expand_block_indices = QsaExpandBlockIndices()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_mtp_reuse_topk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_mtp_reuse_topk.py
new file mode 100644
index 000000000000..3ed07e74c1c1
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_mtp_reuse_topk.py
@@ -0,0 +1,322 @@
+# 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.
+#: This is dep4's layout (attention DP keeps the draft layer in the trunk's manager).
+_COMBINED_SLOTS = (13, 12)
+#: tp4 / tep4's layout: attention TP puts the draft layer in a separate draft KV
+#: cache, so the metadata sizes the buffers by the larger manager (the trunk's 12
+#: sparse layers) and the draft loop reads the draft manager's only slot, 0.
+_SEPARATE_DRAFT_SLOTS = (12, 0)
+
+#: The engines' `max_batch_size` at every member and variant, and so the requests
+#: one draft step carries (one row each) and the shared buffers' height
+#: (`max_num_sequences` = `max_batch_size * max_beam_width`, beam width 1).
+_MAX_BATCH = 128
+#: The longest sequence the default target serves: `max_seq_len` 262144 (16384
+#: on the MTP member).
+_ENVELOPE_SEQ_LEN = 262144
+
+
+def _spec(contexts, **overrides) -> dict:
+ """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), plus one at the
+ last position `max_seq_len` allows, the rest seeded
+ ("full", B) captured positions >= 2047: all 512 slots valid
+ ("explicit", P) one request per captured position in P
+ A request captured at position `p` (its last accepted row) holds
+ `(p + 1) // 4` complete groups; its draft row `d` positions later sees
+ `(p + d + 1) // 4`. `drafted` cycles over the requests. Positions stay below
+ `max_seq_len - d`.
+ """
+ base = dict(
+ contexts=contexts,
+ drafted=(1, 2, 3),
+ rows_per_request=1,
+ shared_rows=64,
+ slots=_COMBINED_SLOTS,
+ block_topk=_BLOCK_TOPK,
+ metadata_dtype=torch.int32,
+ graph_batches=(),
+ max_seq_len=8192,
+ )
+ return {**base, **overrides}
+
+
+def _default(drafted, shared_rows) -> dict:
+ """The default target's draft step: 128 requests, tp4 / tep4's separate-draft
+ slots, positions over the 262144 envelope, captured at 1 and 128 requests."""
+ return _spec(
+ ("mixed", _MAX_BATCH),
+ drafted=drafted,
+ shared_rows=shared_rows,
+ slots=_SEPARATE_DRAFT_SLOTS,
+ graph_batches=(1, _MAX_BATCH),
+ max_seq_len=_ENVELOPE_SEQ_LEN,
+ )
+
+
+class QsaMtpReuseTopk(OpWrapper):
+ """Rebuild each query row's group selection from its request's captured row,
+ with the groups that completed since the capture written into it.
+
+ 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.
+
+ Called by the qwen4_exp target at every layout.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ # ---- decode batch: one row per request; the first 16 requests sit at the
+ # boundaries (empty, short, 511-id, 512-id and full rows), fresh 0 and 1.
+ Cell(
+ why="decode MTP: 32 requests (envelope), 256-row shared buffer",
+ spec=_spec(("mixed", 32), shared_rows=256),
+ ),
+ # ---- the riskiest edge: two fresh groups displace a full row's two oldest
+ Cell(
+ why="8 requests at fresh 2: full, 511-id and short rows gain two groups",
+ spec=_spec(
+ ("explicit", (2050, 402, 2046, 8186, 6, 2047, 1022, 2049)),
+ drafted=(5, 6, 5, 6, 6, 4, 5, 6),
+ shared_rows=8,
+ ),
+ ),
+ # ---- CUDA graph -----------------------------------------------------
+ Cell(
+ why="CUDA graph, drafted 1..6: captured at 1, 8 and 32 requests, new-data replay",
+ spec=_spec(("mixed", 32), drafted=(1, 2, 3, 4, 5, 6), graph_batches=(1, 8, 32)),
+ ),
+ # ---- the default target (tp4 / tep4 / dep4 alike: one selection row per
+ # request, no head or rank axis): its 128-request draft steps.
+ Cell(
+ why="128 requests 3 positions past the capture, 512-row shared buffer (tep4 slots)",
+ spec=_default(drafted=(3,), shared_rows=512),
+ ),
+ Cell(
+ why="128 requests 1 position past the capture, 256-row shared buffer (tep4 slots)",
+ spec=_default(drafted=(1,), shared_rows=256),
+ ),
+ Cell(
+ why="128 requests 7 positions past the capture, 1024-row shared buffer (tep4 slots)",
+ spec=_default(drafted=(7,), shared_rows=1024),
+ ),
+ Cell(
+ why="served mtp7 step: 128 requests filling the engine's 128-row buffer, drafted 1..6",
+ spec=_default(drafted=(1, 2, 3, 4, 5, 6), shared_rows=_MAX_BATCH),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none; every input is an argument, all int32 at the call
+ site:
+ shared_indices `qsa_shared_topk_indices[slot]`, [max_num_sequences, block_topk]
+ shared_visible_blocks `qsa_shared_topk_visible_blocks[slot]`, [max_num_sequences]
+ request_indices `qsa_req_idx_per_token[:rows]`
+ visible_blocks `qsa_visible_blocks[:rows]`
+ output `qsa_topk_indices[:rows]`, the [max_num_tokens, block_topk] buffer
+ `slot` is the layer's row in the cache manager's `qsa_shared_topk_slots`.
+
+ It runs in the MTP draft loop's decode steps after the first: no context
+ rows, one row per request, so rows = requests, and the request map is the
+ row index. Draft step `i` sits `i` positions past the captured row, so
+ `max_draft_len` 1 never reaches this op, 3 reaches 1 and 2 positions, and 7
+ (the target's ceiling) reaches 1..6. With compress_ratio 4, `fresh =
+ ((p + 1) % 4 + i) // 4` for a capture at position `p`: 0 or 1 up to i = 3
+ and 0, 1 or 2 from i = 4. In a full row
+ the fresh ids displace the row's oldest ids (one at fresh 1, two at fresh
+ 2), and a 511-id row at fresh 2 loses its oldest id.
+
+ Preconditions that are values, so not `is_valid` guards:
+ - each captured row is prefix-packed (every valid id before every -1), as
+ `qsa_topk_select` emits it. The append slot comes from the COUNT of
+ non-negative ids, so a -1 hole puts a fresh id over a valid one and
+ loses a group silently;
+ - captured ids are below that row's `shared_visible_blocks`, so a fresh id
+ never duplicates a captured one;
+ - every `request_indices[r]` is a row of both shared tensors;
+ - `shared_visible_blocks[q]` is the visible count of the very row that was
+ captured, and `visible_blocks[r]` is `(position + 1) // cr` of row `r`.
+
+ Loud rejections (not pinned by a test): `ValueError` when the shared width differs from the output
+ width, when `request_indices` or `visible_blocks` does not hold one element
+ per output row, or on CPU tensors; a 1-D `output` or `shared_indices`; a
+ width that is not a power of two (Triton compilation).
+
+ Accepted and silently wrong, hence `is_valid`: any non-contiguous operand
+ (the kernel addresses raw pointers densely and ignores strides); a 3-D
+ `shared_indices` whose middle extent equals `block_topk` (read as slot 0);
+ a 2-D `shared_visible_blocks` (the whole `[slots, seqs]` buffer is read as
+ slot 0's counts). The test pins the last.
+
+ Source: tier 2, an in-tree Triton kernel in the QSA attention backend,
+ JIT-compiled on first call for the device it runs on; the cells were driven
+ through that JIT path on sm_103 (GB300). One program per output row, no
+ autotune, no PDL.
+
+ The default target (tp4 / tep4 / dep4) calls it with the same geometry on every
+ rank: one selection row per request, no head or rank axis (the indexer is
+ replicated). The slots are `_SEPARATE_DRAFT_SLOTS` under attention TP (tp4,
+ tep4) and `_COMBINED_SLOTS` under attention DP. The buffers are `max_num_sequences` rows
+ tall, `max_batch_size * max_beam_width` (`build_attention_metadata`): 128 at
+ every member. The 256-, 512- and 1024-row cells are taller than any engine
+ allocates and are certified as stated. A draft step carries at most 128
+ requests, one row each.
+
+ The `drafted` axis is the offset past the capture (above), not
+ `max_draft_len`: mtp7's offsets 1..6 are the "served mtp7 step" cell, and the
+ offset-7 cell is beyond any served draft step. Every cell holds a full row
+ gaining the widest fresh count its offsets allow.
+
+ Hard limits (read from the kernel, not probed): both row offsets are int64
+ (`request.to(tl.int64) * BLOCK_TOPK`, `row.to(tl.int64) * BLOCK_TOPK`), so no
+ buffer size the envelope reaches wraps them. Ids and visible counts are int32;
+ positions below `max_seq_len` 262144 give at most 65536 groups. `block_topk`
+ must be a power of two (`tl.arange`).
+
+ Certified: int32 ids and metadata, `block_topk` 512, the identity request map,
+ up to 128 requests in 8- to 1024-row shared buffers (128 requests filling a
+ 128-row buffer included), slots (13, 12) and (12, 0), captured positions up to
+ the 262144 envelope, CUDA graph capture at 1, 8, 32 and 128 requests with a
+ new-data replay.
+ Not certified: a permuted request map, more than one row per request, int64
+ metadata, 128 requests at dep4's (13, 12) slots (MTP is served at tep4).
+ """
+
+ def __call__(
+ 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/ops/attention/qsa_nvfp4_pages_to_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_nvfp4_pages_to_bf16.py
new file mode 100644
index 000000000000..32f1b145a305
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_nvfp4_pages_to_bf16.py
@@ -0,0 +1,391 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""A batch's pages of an NVFP4 paged K/V pool, dequantized into a bf16 scratch pool, and the
+block table that addresses them there.
+
+Tier 5 (Triton written for this target). The target's context launches over an NVFP4 pool
+read their requests' pages through this scratch with the bf16 row-tile kernel.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import triton
+import triton.language as tl
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4 import (
+ _e2m1_to_f32,
+)
+
+#: Values per e4m3 block scale, e2m1 values per byte, tokens per V scale group.
+_SF_VEC = 16
+_PER_BYTE = 2
+_V_SF_QUAD = 4
+_E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
+
+#: This checkpoint's K/V geometry and the engine's QSA table (`width` pages a slot);
+#: `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(
+ kv_heads=2,
+ head_dim=256,
+ tokens_per_block=32,
+ width=320,
+ slots=40,
+ global_k=1.0,
+ global_v=1.0,
+ slot_pages=2,
+ layer=0,
+ holes=False,
+)
+
+
+def _spec(kv_lens, **overrides) -> dict:
+ """`kv_lens` is one KV length per listed request, in order."""
+ return {**_CHECKPOINT, "kv_lens": tuple(kv_lens), **overrides}
+
+
+@triton.jit
+def _dequantize_head_page(
+ data,
+ scales,
+ page,
+ head,
+ global_scale,
+ valid,
+ d_stride_page,
+ d_stride_head,
+ d_stride_token,
+ s_stride_page,
+ s_stride_head,
+ GROUPED: tl.constexpr,
+ TOKENS_PER_BLOCK: tl.constexpr,
+ HEAD_DIM: tl.constexpr,
+):
+ """[tpb, D] fp32 values of one head page: e2m1 x e4m3 block scale x second-level
+ scale, in the sparse GQA kernels' order."""
+ tokens = tl.arange(0, TOKENS_PER_BLOCK)[:, None]
+ dims = tl.arange(0, HEAD_DIM)[None, :]
+ packed = tl.load(
+ data + page * d_stride_page + head * d_stride_head + tokens * d_stride_token + dims // 2,
+ mask=valid,
+ other=0,
+ ).to(tl.int32)
+ codes = (packed >> ((dims % 2) * 4)) & 0xF
+ if GROUPED:
+ # V: `[tpb / 4, D / 16, 4]` per head page.
+ offsets = (tokens // 4) * (HEAD_DIM // 4) + (dims // 16) * 4 + tokens % 4
+ else:
+ # K: token-major `[tpb, D / 16]` per head page.
+ offsets = tokens * (HEAD_DIM // 16) + dims // 16
+ block_scales = (
+ tl.load(scales + page * s_stride_page + head * s_stride_head + offsets, mask=valid, other=0)
+ .to(tl.float8e4nv, bitcast=True)
+ .to(tl.float32)
+ )
+ return _e2m1_to_f32(codes) * block_scales * tl.load(global_scale)
+
+
+@triton.jit
+def _nvfp4_pages_to_bf16_kernel(
+ k_cache,
+ v_cache,
+ k_scales,
+ v_scales,
+ kv_scale_quant_orig,
+ block_table,
+ slots,
+ counts,
+ offsets,
+ table,
+ k_out,
+ v_out,
+ num_pool_pages,
+ listed,
+ d_stride_page,
+ d_stride_head,
+ d_stride_token,
+ s_stride_page,
+ s_stride_head,
+ bt_stride,
+ table_stride,
+ o_stride_page,
+ o_stride_head,
+ o_stride_token,
+ NUM_KV_HEADS: tl.constexpr,
+ TOKENS_PER_BLOCK: tl.constexpr,
+ HEAD_DIM: tl.constexpr,
+):
+ request = tl.program_id(0)
+ page = tl.program_id(1)
+ slot = tl.load(slots + request).to(tl.int64)
+ dst = tl.load(offsets + request).to(tl.int64) + page
+ listed_page = (page < tl.load(counts + request)) & (dst < listed)
+ src = tl.load(block_table + slot * bt_stride + page, mask=listed_page, other=-1).to(tl.int64)
+ resident = listed_page & (src >= 0) & (src < num_pool_pages)
+ safe = tl.where(resident, src, 0)
+ tl.store(
+ table + slot * table_stride + page,
+ tl.where(resident, dst, -1).to(tl.int32),
+ mask=listed_page,
+ )
+ tokens = tl.arange(0, TOKENS_PER_BLOCK)[:, None]
+ dims = tl.arange(0, HEAD_DIM)[None, :]
+ out_offsets = dst * o_stride_page + tokens * o_stride_token + dims
+ for head in tl.static_range(NUM_KV_HEADS):
+ k = _dequantize_head_page(
+ k_cache,
+ k_scales,
+ safe,
+ head,
+ kv_scale_quant_orig + 1,
+ resident,
+ d_stride_page,
+ d_stride_head,
+ d_stride_token,
+ s_stride_page,
+ s_stride_head,
+ False,
+ TOKENS_PER_BLOCK,
+ HEAD_DIM,
+ )
+ v = _dequantize_head_page(
+ v_cache,
+ v_scales,
+ safe,
+ head,
+ kv_scale_quant_orig + 2,
+ resident,
+ d_stride_page,
+ d_stride_head,
+ d_stride_token,
+ s_stride_page,
+ s_stride_head,
+ True,
+ TOKENS_PER_BLOCK,
+ HEAD_DIM,
+ )
+ out = out_offsets + head * o_stride_head
+ tl.store(k_out + out, k.to(k_out.dtype.element_ty), mask=listed_page)
+ tl.store(v_out + out, v.to(v_out.dtype.element_ty), mask=listed_page)
+
+
+def _dequantize_pages(data, scales, pages, global_scale, grouped: bool) -> torch.Tensor:
+ """[n, Hkv, tpb, D] fp32 values of `pages` (in the pool) of one role plane."""
+ heads, tokens_per_block, packed = data.shape[1], data.shape[2], data.shape[3]
+ head_dim, groups = packed * _PER_BYTE, packed * _PER_BYTE // _SF_VEC
+ raw = data[pages].long() # [n, H, tpb, 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)
+ flat = scales[pages].reshape(pages.shape[0], heads, tokens_per_block * groups)
+ tokens = torch.arange(tokens_per_block, device=data.device)[:, None]
+ dims = torch.arange(head_dim, device=data.device)[None, :]
+ if grouped:
+ columns = (tokens // _V_SF_QUAD) * (_V_SF_QUAD * groups) + (dims // _SF_VEC) * _V_SF_QUAD
+ columns = columns + tokens % _V_SF_QUAD
+ else:
+ columns = tokens * groups + dims // _SF_VEC
+ sf = flat[:, :, columns.flatten()].reshape(pages.shape[0], heads, tokens_per_block, head_dim)
+ return values * sf.view(torch.float8_e4m3fn).float() * global_scale
+
+
+class QsaNvfp4PagesToBf16(OpWrapper):
+ """K and V of listed requests' pages of an NVFP4 HND paged pool, as bf16 scratch pools,
+ and the block table that addresses them there.
+
+ Request `r` lists its first `counts[r]` table pages, `block_table[slots[r], j]` for `j
+ < counts[r]`, at scratch pages `offsets[r] + j` (those below `listed`, the scratch's
+ page count)::
+
+ k_out[offsets[r] + j, h, t, d] = bf16(e2m1(K code at (s, h, t, d)) x
+ e4m3(K scale of (t, d // 16)) x g_k)
+
+ with `s = block_table[slots[r], j]` and `g_k = kv_scale_quant_orig[1]`, `v_out` likewise
+ from the V planes and `g_v = kv_scale_quant_orig[2]`. K's block scales are token-major
+ `[tpb, D/16]` per head page, V's in four-token groups `[tpb/4, D/16, 4]`. The product is
+ formed in fp32 in the sparse GQA kernels' order and rounded once to bf16, so a bf16
+ kernel over the scratch reads the values those kernels dequantize. `table[slots[r], j]`
+ becomes `offsets[r] + j`, or -1 where the pool's entry is unmapped or out of the pool
+ (whose scratch page is then zeros); the call writes no other entry of `table`, which the
+ caller fills with -1. Returns `(k_out, v_out, table)`.
+
+ Fusion boundary: the listing's dequantize and table; choosing the requests and their
+ page counts (and `offsets`, their exclusive prefix sum) is the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="one request, one page", spec=_spec((17,))),
+ Cell(why="con 1 prefill: one 8192-token request", spec=_spec((8192,))),
+ Cell(why="max_num_tokens: two 8192-token requests", spec=_spec((8192, 8192))),
+ Cell(
+ why="three requests, a chunked continuation and one filling the table",
+ spec=_spec((5000 + 3192, 3000, 10240)),
+ ),
+ Cell(why="many short requests", spec=_spec(tuple(33 + 61 * i for i in range(30)))),
+ Cell(why="holes: unmapped and out-of-pool entries", spec=_spec((4000, 2000), holes=True)),
+ Cell(
+ why="second-level scales g_k 0.37, g_v 2.5: the product rounds to bf16",
+ spec=_spec((1500, 700), global_k=0.37, global_v=2.5),
+ ),
+ Cell(
+ why="the engine's trunk layout: layer 11's planes among the 24 role pages",
+ spec=_spec((3000, 900), slot_pages=24, layer=11),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none -- every input is an argument. Call site: the target's
+ context/mixed forward over an NVFP4 pool, once per attention layer and group of the
+ step's context requests (`_nvfp4_context_plan`: slots from the context rows' request
+ indices, counts `ceil(kv_len / tpb)` from the device KV lengths, the groups and the
+ scratch's pages from the host ones); the row-tile kernel then reads the scratch
+ through `table`.
+
+ One Triton kernel, a program per (request, table page), both K/V heads and both
+ roles: a head page is 32 tokens x 256 dims, 8 warps. Programs past a request's
+ count store nothing.
+
+ Preconditions: uint8 planes `[pages, Hkv, tpb, D/2]` (data) and `[pages, Hkv, tpb,
+ D/16]` (scales), K's and V's at equal strides, unit last-dim stride and contiguous
+ head pages (the QSA manager's role planes); `kv_scale_quant_orig` an fp32 [3];
+ `table` an int32 tensor of `block_table`'s shape. Bitwise deterministic.
+ """
+
+ def __call__(
+ self,
+ 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,
+ slots: torch.Tensor,
+ counts: torch.Tensor,
+ offsets: torch.Tensor,
+ table: torch.Tensor,
+ k_out: torch.Tensor,
+ v_out: torch.Tensor,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ requests, listed = slots.shape[0], k_out.shape[0]
+ if requests and listed:
+ _nvfp4_pages_to_bf16_kernel[(requests, block_table.shape[1])](
+ k_cache,
+ v_cache,
+ k_scale_cache,
+ v_scale_cache,
+ kv_scale_quant_orig,
+ block_table,
+ slots,
+ counts,
+ offsets,
+ table,
+ k_out,
+ v_out,
+ k_cache.shape[0],
+ listed,
+ k_cache.stride(0),
+ k_cache.stride(1),
+ k_cache.stride(2),
+ k_scale_cache.stride(0),
+ k_scale_cache.stride(1),
+ block_table.stride(0),
+ table.stride(0),
+ k_out.stride(0),
+ k_out.stride(1),
+ k_out.stride(2),
+ NUM_KV_HEADS=k_cache.shape[1],
+ TOKENS_PER_BLOCK=k_cache.shape[2],
+ HEAD_DIM=k_cache.shape[3] * _PER_BYTE,
+ num_warps=8,
+ )
+ return k_out, v_out, table
+
+ def reference(
+ self,
+ 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,
+ slots: torch.Tensor,
+ counts: torch.Tensor,
+ offsets: torch.Tensor,
+ table: torch.Tensor,
+ k_out: torch.Tensor,
+ v_out: torch.Tensor,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ k_ref, v_ref, table_ref = k_out.clone(), v_out.clone(), table.clone()
+ listed = k_out.shape[0]
+ for r in range(slots.shape[0]):
+ slot, base = int(slots[r]), int(offsets[r])
+ for j in range(min(int(counts[r]), block_table.shape[1])):
+ dst = base + j
+ if dst >= listed:
+ break
+ src = int(block_table[slot, j])
+ resident = 0 <= src < k_cache.shape[0]
+ table_ref[slot, j] = dst if resident else -1
+ page = torch.tensor([src if resident else 0], device=k_cache.device)
+ for out, data, sc, g, grouped in (
+ (k_ref, k_cache, k_scale_cache, kv_scale_quant_orig[1], False),
+ (v_ref, v_cache, v_scale_cache, kv_scale_quant_orig[2], True),
+ ):
+ vals = _dequantize_pages(data, sc, page, g, grouped)[0]
+ out[dst] = (vals if resident else torch.zeros_like(vals)).to(out.dtype)
+ return k_ref, v_ref, table_ref
+
+ def is_valid(
+ self,
+ 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,
+ slots: torch.Tensor,
+ counts: torch.Tensor,
+ offsets: torch.Tensor,
+ table: torch.Tensor,
+ k_out: torch.Tensor,
+ v_out: torch.Tensor,
+ ) -> None:
+ pages, heads, tokens_per_block, packed = k_cache.shape
+ head_dim = packed * _PER_BYTE
+ assert tokens_per_block % _V_SF_QUAD == 0, "tokens_per_block: a multiple of 4"
+ assert head_dim % _SF_VEC == 0 and math.log2(head_dim).is_integer(), "head_dim"
+ for name, plane, width in (
+ ("k_cache", k_cache, packed),
+ ("v_cache", v_cache, packed),
+ ("k_scale_cache", k_scale_cache, head_dim // _SF_VEC),
+ ("v_scale_cache", v_scale_cache, head_dim // _SF_VEC),
+ ):
+ assert plane.dtype == torch.uint8, f"{name}: uint8"
+ assert tuple(plane.shape) == (pages, heads, tokens_per_block, width), f"{name}: shape"
+ assert plane.stride(3) == 1, f"{name}: unit last-dim stride"
+ assert v_cache.stride() == k_cache.stride(), "K and V data planes at equal strides"
+ assert v_scale_cache.stride() == k_scale_cache.stride(), "K and V scales at equal strides"
+ assert k_scale_cache.stride(2) == head_dim // _SF_VEC, "scales: contiguous head pages"
+ assert kv_scale_quant_orig.dtype == torch.float32 and kv_scale_quant_orig.shape == (3,)
+ assert block_table.dim() == 2 and block_table.dtype == torch.int32
+ assert block_table.stride(1) == 1, "block_table: row-contiguous"
+ assert table.shape == block_table.shape and table.dtype == torch.int32
+ assert table.stride(1) == 1, "table: row-contiguous"
+ for name, t in (("slots", slots), ("counts", counts), ("offsets", offsets)):
+ assert t.dim() == 1 and t.shape == slots.shape, f"{name}: one per request"
+ assert t.dtype in (torch.int32, torch.int64), f"{name}: integer"
+ for out in (k_out, v_out):
+ assert out.dtype == torch.bfloat16, "out: bf16"
+ assert out.shape[1:] == (heads, tokens_per_block, head_dim), (
+ "out: [listed, Hkv, tpb, D]"
+ )
+ assert out.is_contiguous(), "out: contiguous"
+
+
+qsa_nvfp4_pages_to_bf16 = QsaNvfp4PagesToBf16()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_nvfp4_verify_ptx.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_nvfp4_verify_ptx.py
new file mode 100644
index 000000000000..d184882c004c
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_nvfp4_verify_ptx.py
@@ -0,0 +1,251 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Inline-PTX pieces of `qsa_paged_sparse_gqa_nvfp4_verify`'s kernel beyond the ones the
+per-row kernel (`qsa_paged_sparse_gqa_nvfp4_decode`) already has: the union scan's bit
+operations, vector global loads and stores, the tcgen05 thread-sync fences, and tcgen05.mma issue with
+precomputed descriptors.
+
+`gemm_ts` issues every K step of a GEMM whose A operand sits in tensor memory from one
+inline-PTX block, the shared-memory descriptor of B advanced by immediates; the encodings are
+those of CUTLASS's `cute/arch/mma_sm100_desc.hpp`. `cute.gemm` rebuilds the descriptors per K
+step on the uniform datapath (about 18 instructions a UMMA), which the verify kernel's single
+MMA warp cannot hide.
+"""
+
+from __future__ import annotations
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32
+from cutlass._mlir.dialects import llvm
+from cutlass._mlir.dialects import nvvm as _nvvm
+from cutlass.cute.nvgpu import tcgen05
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4_decode import ( # noqa: E501
+ _asm,
+ _ir,
+)
+
+# ---------------------------------------------------------------------------------------
+# scalar and memory helpers
+# ---------------------------------------------------------------------------------------
+
+
+@dsl_user_op
+def popc(x, *, loc=None, ip=None):
+ return cutlass.Int32(
+ _asm(T.i32(), [_ir(x, cutlass.Int32, loc, ip)], "popc.b32 $0, $1;", "=r,r", loc, ip)
+ )
+
+
+@dsl_user_op
+def ctz(x, *, loc=None, ip=None):
+ """Trailing zeros of a nonzero word."""
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [_ir(x, cutlass.Int32, loc, ip)],
+ "{ .reg .b32 t; brev.b32 t, $1; clz.b32 $0, t; }",
+ "=r,r",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def red_or_shared(saddr, v, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [_ir(saddr, cutlass.Int32, loc, ip), _ir(v, cutlass.Int32, loc, ip)],
+ "{ red.shared.or.b32 [$1], $2; mov.u32 $0, 0; }",
+ "=r,r,r",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def stg4(gaddr, a, b, c, d, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [_ir(gaddr, cutlass.Int64, loc, ip)]
+ + [_ir(v, cutlass.Int32, loc, ip) 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=True,
+ )
+
+
+@dsl_user_op
+def stg4f(gaddr, a, b, c, d, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [_ir(gaddr, cutlass.Int64, loc, ip)]
+ + [_ir(v, cutlass.Float32, loc, ip) for v in (a, b, c, d)],
+ "{ st.global.v4.f32 [$1], {$2, $3, $4, $5}; mov.u32 $0, 0; }",
+ "=r,l,f,f,f,f",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def ldg2(gaddr, *, loc=None, ip=None):
+ res = _asm(
+ llvm.StructType.get_literal([T.i32()] * 2),
+ [_ir(gaddr, cutlass.Int64, loc, ip)],
+ "ld.global.v2.b32 {$0, $1}, [$2];",
+ "=r,=r,l",
+ loc,
+ ip,
+ side=True,
+ )
+ return [cutlass.Int32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(2)]
+
+
+@dsl_user_op
+def stg2(gaddr, a, b, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [
+ _ir(gaddr, cutlass.Int64, loc, ip),
+ _ir(a, cutlass.Int32, loc, ip),
+ _ir(b, cutlass.Int32, loc, ip),
+ ],
+ "{ st.global.v2.b32 [$1], {$2, $3}; mov.u32 $0, 0; }",
+ "=r,l,r,r",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+def groups8(x):
+ """Bit i of the result: any of bits 4i .. 4i + 3 of `x` (the eight four-token groups of a
+ 32-token bitmap word)."""
+ y = x | (x >> 1)
+ y = y | (y >> 2)
+ z = y & 1
+ for i in range(1, 8):
+ z = z | (((y >> (4 * i)) & 1) << i)
+ return z
+
+
+def spread4(nib):
+ """Bits 0..3 of `nib` to bits 0, 8, 16, 24."""
+ return (nib & 1) | ((nib & 2) << 7) | ((nib & 4) << 14) | ((nib & 8) << 21)
+
+
+@dsl_user_op
+def tcgen05_fence_before(*, loc=None, ip=None):
+ _nvvm.tcgen05_fence(_nvvm.Tcgen05FenceKind.BEFORE_THREAD_SYNC, loc=loc, ip=ip)
+
+
+@dsl_user_op
+def tcgen05_fence_after(*, loc=None, ip=None):
+ _nvvm.tcgen05_fence(_nvvm.Tcgen05FenceKind.AFTER_THREAD_SYNC, loc=loc, ip=ip)
+
+
+# ---------------------------------------------------------------------------------------
+# tcgen05.mma issue with precomputed descriptors
+# ---------------------------------------------------------------------------------------
+
+#: Swizzle: B -> the descriptor's layout type (none, 32B, 64B, 128B).
+_SW_TYPES = {0: 0, 1: 6, 2: 4, 3: 2}
+#: Layout type -> swizzle atom size along MN (16-byte units).
+_SW_ATOM_MN = {0: 1, 6: 2, 4: 4, 2: 8}
+
+
+def instr_desc(op) -> int:
+ """The 32-bit instruction descriptor of an f16/bf16 x f16/bf16 -> f32 tcgen05 MMA op."""
+ fmt = {cutlass.Float16: 0, cutlass.BFloat16: 1}
+ a_major = 0 if op.a_major_mode == tcgen05.OperandMajorMode.K else 1
+ b_major = 0 if op.b_major_mode == tcgen05.OperandMajorMode.K else 1
+ m, n = op.shape_mnk[0], op.shape_mnk[1]
+ assert op.acc_dtype is cutlass.Float32 and m in (64, 128) and n % 8 == 0 and 8 <= n <= 256
+ d = 0
+ d |= 1 << 4 # c_format f32
+ d |= fmt[op.a_dtype] << 7
+ d |= fmt[op.b_dtype] << 10
+ d |= a_major << 15
+ d |= b_major << 16
+ d |= (n >> 3) << 17
+ d |= (m >> 4) << 24
+ return d
+
+
+def smem_desc_base(layout_u128, swizzle, k_major: bool) -> int:
+ """The 64-bit shared-memory descriptor without its start address, for an operand whose
+ (MN, K) layout `layout_u128` is in 16-byte units."""
+ assert swizzle.num_base == 4 and swizzle.num_shift == 3
+ ltype = _SW_TYPES[swizzle.num_bits]
+ atom_mn = _SW_ATOM_MN[ltype]
+ if not k_major:
+ canon = cute.logical_divide(layout_u128, (atom_mn, 8))
+ s00 = canon.stride[0][0]
+ s10 = canon.stride[1][0]
+ assert (ltype == 0 or s00 == 1) and s10 == atom_mn, "not a canonical UMMA MN layout"
+ s01, s11 = canon.stride[0][1], canon.stride[1][1]
+ sbo, lbo = (s01, s11) if ltype == 0 else (s11, s01)
+ else:
+ canon = cute.logical_divide(layout_u128, (8, 2))
+ s00 = canon.stride[0][0]
+ s10 = canon.stride[1][0]
+ assert s00 == atom_mn and (ltype == 0 or s10 == 1), "not a canonical UMMA K layout"
+ sbo, lbo = canon.stride[0][1], s10
+ d = 0
+ d |= (lbo & 0x3FFF) << 16
+ d |= (sbo & 0x3FFF) << 32
+ d |= 1 << 46 # descriptor version
+ d |= (ltype & 0x7) << 61
+ return d
+
+
+def gemm_ts(op, acc_tmem, a_tmem, tCrA, tCrB, sB, zero_init: bool):
+ """acc (+)= A B over every K step: A in tensor memory at `a_tmem` (fragment `tCrA`), B the
+ shared-memory operand `sB` (fragment `tCrB`, one stage). Issued by one elected thread of
+ the calling warp. Plain Python run while the kernel is traced: the K loop unrolls into one
+ PTX block."""
+ idesc = instr_desc(op)
+ k_major_b = op.b_major_mode == tcgen05.OperandMajorMode.K
+ base_b = smem_desc_base(
+ cute.recast_layout(128, op.b_dtype.width, sB.layout[0]),
+ sB.iterator.type.swizzle_type,
+ k_major_b,
+ )
+ b_lo = base_b & 0xFFFFFFFF
+ b_hi = (base_b >> 32) & 0xFFFFFFFF
+ a_layout = cute.recast_layout(32, tCrA.element_type.width, tCrA.layout)
+ n_k = cute.size(tCrA.shape[2])
+ off_a = [cute.crd2idx((0, 0, k), a_layout) for k in range(n_k)]
+ off_b = [cute.crd2idx((0, 0, k), tCrB.layout) for k in range(n_k)]
+ start_b = Int32(b_lo) | Int32((sB[None, None, 0].iterator.toint() & 0x3FFFF) >> 4)
+ body = "{\n\t.reg .pred leader;\n\t.reg .b32 db_lo;\n\t.reg .b64 db;\n\telect.sync _|leader, -1;\n\t"
+ for k in range(n_k):
+ acc = "0" if (zero_init and k == 0) else "1"
+ body += (
+ f"add.u32 db_lo, $2, {hex(off_b[k])};\n\t"
+ f"mov.b64 db, {{db_lo, {hex(b_hi)}}};\n\t"
+ f"@leader tcgen05.mma.cta_group::1.kind::f16 [$0], [$1 + {hex(off_a[k])}], db, "
+ f"{hex(idesc)}, {acc};\n\t"
+ )
+ body += "}\n"
+ llvm.inline_asm(
+ None,
+ [
+ Int32(cute.arch.make_warp_uniform(acc_tmem)).ir_value(),
+ Int32(cute.arch.make_warp_uniform(a_tmem)).ir_value(),
+ Int32(cute.arch.make_warp_uniform(start_b)).ir_value(),
+ ],
+ body,
+ "r,r,r",
+ has_side_effects=True,
+ is_align_stack=False,
+ asm_dialect=llvm.AsmDialect.AD_ATT,
+ )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_index_scores.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_index_scores.py
new file mode 100644
index 000000000000..6c74ae3a15c7
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_index_scores.py
@@ -0,0 +1,1158 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""QSA compressed-block index scoring, read straight out of the paged index-K side cache."""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+# Approved exception: QSA's selection path is Triton called as Python, with no
+# torch.ops boundary. The engine's own sparse hook (`select_qsa_paged_tokens`
+# in `_torch/attention/backends/sparse/qsa/module.py`) calls this same callable.
+from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import (
+ triton_qsa_paged_index_scores as _triton_qsa_paged_index_scores,
+)
+
+_INTEGER_DTYPES = (torch.int32, torch.int64)
+
+#: At or below this many rows a call is decode-sized: the engine's own threshold
+#: between its one-tile-a-program and its deeper-tile launches.
+_DECODE_ROW_LIMIT = 256
+#: Dot columns of one visible-split program: whole query rows of `heads` columns each
+#: (`tl.dot` needs at least 16), so the rows of one request share each key gather.
+_SPLIT_DOT_WIDTH = 16
+#: Score columns a tile gathers keys for.
+_SPLIT_BLOCK_N = 64
+#: Programs a launch holds per SM, empty request slots included. One row tile (one
+#: request's verify rows: a chain of dependent loads) wants a tile a program, up to the
+#: 1024 tiles of a 262143-position window; wider batches stream several requests at once.
+_SPLIT_PROGRAMS_PER_SM = 16
+_SPLIT_PROGRAMS_PER_SM_ONE_TILE = 32
+_SPLIT_NUM_WARPS = 2
+_SM_COUNT: dict = {}
+
+
+# Metadata views can begin at any row of a mixed context/decode step. Their
+# alignment does not change these scalar loads or the score computation.
+@triton.jit(
+ do_not_specialize=["num_rows", "split"],
+ do_not_specialize_on_alignment=["query_positions", "request_indices"],
+)
+def _qsa_index_scores_visible_split_kernel(
+ q,
+ index_cache,
+ block_table,
+ query_positions,
+ request_indices,
+ output,
+ score_scale,
+ num_rows,
+ split,
+ q_stride_row: tl.constexpr,
+ q_stride_head: tl.constexpr,
+ q_stride_dim: tl.constexpr,
+ cache_stride_page: tl.constexpr,
+ cache_stride_token: tl.constexpr,
+ cache_stride_dim: tl.constexpr,
+ block_table_stride_request: tl.constexpr,
+ block_table_stride_page: tl.constexpr,
+ output_stride_row: tl.constexpr,
+ output_stride_block: tl.constexpr,
+ NUM_INDEX_HEADS: tl.constexpr,
+ INDEX_HEAD_DIM: tl.constexpr,
+ BLOCK_D: tl.constexpr,
+ TOKENS_PER_BLOCK: tl.constexpr,
+ COMPRESS_RATIO: tl.constexpr,
+ MAX_COMPRESSED_BLOCKS: tl.constexpr,
+ ROWS_PER_TILE: tl.constexpr,
+ QUERY_TILE: tl.constexpr,
+ BLOCK_H: tl.constexpr,
+ BLOCK_N: tl.constexpr,
+ NUM_STAGES: tl.constexpr,
+ USE_PDL: tl.constexpr,
+):
+ """Visible-only scores of one request's rows among ROWS_PER_TILE consecutive rows (in
+ QUERY_TILE dot lanes). Program (t, g, s) takes the g-th distinct request of row tile t,
+ in row order, and walks its column tiles s, s + split, ... up to its visible bound: the
+ grid is sized by the device, not by the table, a tile's keys are gathered once for every
+ row of their request in the row tile, and different requests run side by side.
+ """
+ if USE_PDL:
+ tl.extra.cuda.gdc_wait()
+ tl.extra.cuda.gdc_launch_dependents()
+ row_base = tl.program_id(0) * ROWS_PER_TILE
+ group_slot = tl.program_id(1)
+ first_tile = tl.program_id(2)
+
+ # One dot column per (query row, index head) pair; q is read before the metadata
+ # decides whether this program has work, so the two loads overlap.
+ dot_columns = tl.arange(0, BLOCK_H)
+ column_row = dot_columns // NUM_INDEX_HEADS
+ column_head = dot_columns % NUM_INDEX_HEADS
+ rows = row_base + column_row
+ present = (column_row < ROWS_PER_TILE) & (rows < num_rows)
+ dim_offsets = tl.arange(0, BLOCK_D)
+ query_values = tl.load(
+ q
+ + rows[None, :] * q_stride_row
+ + dim_offsets[:, None] * q_stride_dim
+ + column_head[None, :] * q_stride_head,
+ mask=(dim_offsets[:, None] < INDEX_HEAD_DIM) & present[None, :],
+ other=0.0,
+ )
+
+ # Each tile row (lane): its request and visible bound.
+ lanes = tl.arange(0, QUERY_TILE)
+ lane_rows = row_base + lanes
+ lane_present = (lanes < ROWS_PER_TILE) & (lane_rows < num_rows)
+ lane_requests = tl.load(request_indices + lane_rows, mask=lane_present, other=-1).to(tl.int64)
+ lane_positions = tl.load(query_positions + lane_rows, mask=lane_present, other=0)
+ lane_visible = tl.where(
+ lane_present, tl.minimum((lane_positions + 1) // COMPRESS_RATIO, MAX_COMPRESSED_BLOCKS), 0
+ )
+ # A lane leads its request's group when no earlier lane holds the same request; groups
+ # are numbered by their leader's lane.
+ earlier_same = (
+ (lane_requests[:, None] == lane_requests[None, :])
+ & (lanes[:, None] < lanes[None, :])
+ & lane_present[:, None]
+ )
+ leads = lane_present & (tl.max(earlier_same.to(tl.int32), axis=0) == 0)
+ rank = tl.cumsum(leads.to(tl.int32), axis=0) - 1
+ request = tl.max(tl.where(leads & (rank == group_slot), lane_requests, -1))
+ lane_group = lane_present & (lane_requests == request)
+ group_visible = tl.max(tl.where(lane_group, lane_visible, 0))
+ # No such group (request -1) leaves group_visible at 0.
+ if first_tile * BLOCK_N >= group_visible:
+ return
+
+ column_offsets = tl.arange(0, BLOCK_N)
+ group_tiles = tl.cdiv(group_visible, BLOCK_N)
+ for tile in tl.range(first_tile, group_tiles, split, num_stages=NUM_STAGES):
+ block_columns = tile * BLOCK_N + column_offsets
+ valid_blocks = block_columns < group_visible
+ anchor_positions = block_columns * COMPRESS_RATIO + COMPRESS_RATIO - 1
+ logical_pages = anchor_positions // TOKENS_PER_BLOCK
+ token_in_page = anchor_positions % TOKENS_PER_BLOCK
+ physical_pages = tl.load(
+ block_table
+ + request * block_table_stride_request
+ + logical_pages * block_table_stride_page,
+ mask=valid_blocks,
+ other=0,
+ ).to(tl.int64)
+ keys = tl.load(
+ index_cache
+ + physical_pages[:, None] * cache_stride_page
+ + token_in_page[:, None] * cache_stride_token
+ + dim_offsets[None, :] * cache_stride_dim,
+ mask=valid_blocks[:, None] & (dim_offsets[None, :] < INDEX_HEAD_DIM),
+ other=0.0,
+ eviction_policy="evict_first",
+ )
+ positive_scores = tl.maximum(tl.dot(keys, query_values, out_dtype=tl.float32), 0.0)
+ # Each row's heads are adjacent dot columns: sum them per row.
+ scores = tl.sum(tl.reshape(positive_scores, [BLOCK_N, QUERY_TILE, NUM_INDEX_HEADS]), axis=2)
+ tl.store(
+ output
+ + lane_rows[None, :] * output_stride_row
+ + block_columns[:, None] * output_stride_block,
+ scores * score_scale,
+ mask=lane_group[None, :] & (block_columns[:, None] < lane_visible[None, :]),
+ )
+
+
+def _sm_count(device: torch.device) -> int:
+ index = device.index if device.index is not None else torch.cuda.current_device()
+ if index not in _SM_COUNT:
+ _SM_COUNT[index] = torch.cuda.get_device_properties(index).multi_processor_count
+ return _SM_COUNT[index]
+
+
+def _rows_per_tile(rows: int, query_tile: int) -> int:
+ """Rows a program's dot lanes hold: the whole batch when it fits one dot tile, else a
+ full tile.
+ """
+ return min(rows, query_tile)
+
+
+def _takes_visible_split(q: torch.Tensor, only_visible_blocks: bool, prefetch: bool) -> bool:
+ """The visible-split kernel serves decode-sized, visible-only calls whose head count
+ divides its dot width; `prefetch` keeps the engine's PDL prefetch kernel.
+ """
+ heads = q.shape[1] if q.ndim == 3 else 0
+ return (
+ only_visible_blocks
+ and not prefetch
+ and q.ndim == 3
+ and q.shape[0] <= _DECODE_ROW_LIMIT
+ and 0 < heads <= _SPLIT_DOT_WIDTH
+ and _SPLIT_DOT_WIDTH % heads == 0
+ )
+
+
+def _qsa_index_scores_visible_split(
+ q: torch.Tensor,
+ index_cache: torch.Tensor,
+ block_table: torch.Tensor,
+ query_positions: torch.Tensor,
+ request_indices: torch.Tensor,
+ tokens_per_block: int,
+ compress_ratio: int,
+ use_pdl: bool,
+ *,
+ dot_width: int = _SPLIT_DOT_WIDTH,
+ rows_per_tile: int | None = None,
+ block_n: int = _SPLIT_BLOCK_N,
+ programs_per_sm: int | None = None,
+ num_warps: int = _SPLIT_NUM_WARPS,
+ num_stages: int = 2,
+) -> torch.Tensor:
+ """One launch of the visible-split kernel after the callable's own checks."""
+ if index_cache.ndim != 4 or index_cache.shape[2] != 1:
+ raise ValueError("QSA index cache must have shape [pages, tokens, 1, dim]")
+ if block_table.ndim != 2 or block_table.dtype != torch.int32:
+ raise ValueError("QSA index scoring requires a two-dimensional int32 block table")
+ rows, num_index_heads, index_head_dim = q.shape
+ if index_head_dim <= 0 or index_cache.shape[3] != index_head_dim:
+ raise ValueError("QSA index query and cache dimensions do not match")
+ if q.dtype != index_cache.dtype:
+ raise ValueError("QSA index query and cache dtypes must match")
+ if query_positions.shape != (rows,) or request_indices.shape != (rows,):
+ raise ValueError("QSA index-scoring metadata must have one entry per row")
+ if query_positions.dtype not in _INTEGER_DTYPES or request_indices.dtype not in _INTEGER_DTYPES:
+ raise ValueError("QSA index-scoring metadata must use integer storage")
+ if tokens_per_block <= 0 or compress_ratio < 2:
+ raise ValueError("QSA index scoring requires positive pages and compression ratio >= 2")
+ for t in (index_cache, block_table, query_positions, request_indices):
+ if t.device != q.device:
+ raise ValueError("QSA paged index scoring tensors must share one CUDA device")
+ max_compressed_blocks = block_table.shape[1] * tokens_per_block // compress_ratio
+ if max_compressed_blocks <= 0:
+ raise ValueError("QSA index scoring requires a non-empty block table")
+ output = torch.empty((rows, max_compressed_blocks), dtype=torch.float32, device=q.device)
+ if rows == 0:
+ return output
+ query_tile = dot_width // num_index_heads
+ if rows_per_tile is None:
+ rows_per_tile = _rows_per_tile(rows, query_tile)
+ row_tiles = triton.cdiv(rows, rows_per_tile)
+ if programs_per_sm is None:
+ programs_per_sm = (
+ _SPLIT_PROGRAMS_PER_SM_ONE_TILE if row_tiles == 1 else _SPLIT_PROGRAMS_PER_SM
+ )
+ # Every (row tile, group slot) pair gets `split` programs; slots past a tile's request
+ # count exit at once.
+ split = min(
+ triton.cdiv(max_compressed_blocks, block_n),
+ max(1, triton.cdiv(programs_per_sm * _sm_count(q.device), row_tiles * rows_per_tile)),
+ )
+ _qsa_index_scores_visible_split_kernel[(row_tiles, rows_per_tile, split)](
+ q,
+ index_cache,
+ block_table,
+ query_positions,
+ request_indices,
+ output,
+ index_head_dim**-0.5,
+ rows,
+ split,
+ q.stride(0),
+ q.stride(1),
+ q.stride(2),
+ index_cache.stride(0),
+ index_cache.stride(1),
+ index_cache.stride(3),
+ block_table.stride(0),
+ block_table.stride(1),
+ output.stride(0),
+ output.stride(1),
+ NUM_INDEX_HEADS=num_index_heads,
+ INDEX_HEAD_DIM=index_head_dim,
+ BLOCK_D=max(16, triton.next_power_of_2(index_head_dim)),
+ TOKENS_PER_BLOCK=tokens_per_block,
+ COMPRESS_RATIO=compress_ratio,
+ MAX_COMPRESSED_BLOCKS=max_compressed_blocks,
+ ROWS_PER_TILE=rows_per_tile,
+ QUERY_TILE=query_tile,
+ BLOCK_H=dot_width,
+ BLOCK_N=block_n,
+ NUM_STAGES=num_stages,
+ USE_PDL=use_pdl,
+ num_warps=num_warps,
+ launch_pdl=use_pdl,
+ )
+ return output
+
+
+@triton.jit
+def _qsa_index_scores_group_kernel(
+ q,
+ index_cache,
+ block_table,
+ query_positions,
+ request_indices,
+ output,
+ score_scale,
+ rows,
+ q_stride_row: tl.constexpr,
+ q_stride_head: tl.constexpr,
+ q_stride_dim: tl.constexpr,
+ cache_stride_page: tl.constexpr,
+ cache_stride_token: tl.constexpr,
+ cache_stride_dim: tl.constexpr,
+ block_table_stride_request: tl.constexpr,
+ block_table_stride_page: tl.constexpr,
+ output_stride_row: tl.constexpr,
+ output_stride_block: tl.constexpr,
+ NUM_INDEX_HEADS: tl.constexpr,
+ INDEX_HEAD_DIM: tl.constexpr,
+ BLOCK_D: tl.constexpr,
+ TOKENS_PER_BLOCK: tl.constexpr,
+ COMPRESS_RATIO: tl.constexpr,
+ MAX_COMPRESSED_BLOCKS: tl.constexpr,
+ GROUP_ROWS: tl.constexpr,
+ BLOCK_G: tl.constexpr,
+ BLOCK_N: tl.constexpr,
+ TILES_PER_PROGRAM: tl.constexpr,
+ USE_PDL: tl.constexpr,
+):
+ """The engine's row-wise kernel (visible blocks only) over a run of ``GROUP_ROWS``
+ consecutive rows a program: when the run's rows name one request (a verify step's
+ rows), one key gather feeds every row's heads as the dot's columns; otherwise each
+ row gathers its own request's keys. Each score is the same dot, ReLU and head sum as
+ the row-wise kernel's, bit for bit."""
+ BLOCK_H: tl.constexpr = BLOCK_G * NUM_INDEX_HEADS
+ if USE_PDL:
+ tl.extra.cuda.gdc_wait()
+ tl.extra.cuda.gdc_launch_dependents()
+ group = tl.program_id(0)
+ slots = tl.arange(0, BLOCK_G)
+ group_rows = group * GROUP_ROWS + slots
+ row_ok = (slots < GROUP_ROWS) & (group_rows < rows)
+ lead = tl.load(request_indices + group * GROUP_ROWS).to(tl.int64)
+ requests = tl.load(request_indices + group_rows, mask=row_ok, other=0).to(tl.int64)
+ requests = tl.where(row_ok, requests, lead)
+ positions = tl.load(query_positions + group_rows, mask=row_ok, other=0)
+ visible = tl.where(row_ok, (positions + 1) // COMPRESS_RATIO, 0)
+ max_visible = tl.max(visible, axis=0)
+ first_tile = tl.program_id(1) * TILES_PER_PROGRAM
+ if first_tile * BLOCK_N >= max_visible:
+ return
+ last_tile = tl.minimum(first_tile + TILES_PER_PROGRAM, tl.cdiv(max_visible, BLOCK_N))
+ shared = tl.max((requests != lead).to(tl.int32), axis=0) == 0
+
+ dim_offsets = tl.arange(0, BLOCK_D)
+ # Dot column c is head c % heads of the run's row c // heads.
+ columns = tl.arange(0, BLOCK_H)
+ column_slot = columns // NUM_INDEX_HEADS
+ column_rows = group * GROUP_ROWS + column_slot
+ column_ok = (column_slot < GROUP_ROWS) & (column_rows < rows)
+ query_values = tl.load(
+ q
+ + column_rows[None, :] * q_stride_row
+ + dim_offsets[:, None] * q_stride_dim
+ + (columns % NUM_INDEX_HEADS)[None, :] * q_stride_head,
+ mask=(dim_offsets[:, None] < INDEX_HEAD_DIM) & column_ok[None, :],
+ other=0.0,
+ )
+ dim_ok = dim_offsets[None, :] < INDEX_HEAD_DIM
+ block_offsets = tl.arange(0, BLOCK_N)
+ for tile in tl.range(first_tile, last_tile):
+ block_columns = tile * BLOCK_N + block_offsets
+ valid_blocks = (block_columns < MAX_COMPRESSED_BLOCKS) & (block_columns < max_visible)
+ anchor_positions = block_columns * COMPRESS_RATIO + COMPRESS_RATIO - 1
+ logical_pages = anchor_positions // TOKENS_PER_BLOCK
+ token_in_page = anchor_positions % TOKENS_PER_BLOCK
+ if shared:
+ physical_pages = tl.load(
+ block_table
+ + lead * block_table_stride_request
+ + logical_pages * block_table_stride_page,
+ mask=valid_blocks,
+ other=0,
+ ).to(tl.int64)
+ keys = tl.load(
+ index_cache
+ + physical_pages[:, None] * cache_stride_page
+ + token_in_page[:, None] * cache_stride_token
+ + dim_offsets[None, :] * cache_stride_dim,
+ mask=valid_blocks[:, None] & dim_ok,
+ other=0.0,
+ eviction_policy="evict_first",
+ )
+ per_head_scores = tl.dot(keys, query_values, out_dtype=tl.float32)
+ else:
+ per_head_scores = tl.zeros((BLOCK_N, BLOCK_H), dtype=tl.float32)
+ for slot in tl.static_range(GROUP_ROWS):
+ request = tl.sum(tl.where(slots == slot, requests, 0), axis=0)
+ physical_pages = tl.load(
+ block_table
+ + request * block_table_stride_request
+ + logical_pages * block_table_stride_page,
+ mask=valid_blocks,
+ other=0,
+ ).to(tl.int64)
+ keys = tl.load(
+ index_cache
+ + physical_pages[:, None] * cache_stride_page
+ + token_in_page[:, None] * cache_stride_token
+ + dim_offsets[None, :] * cache_stride_dim,
+ mask=valid_blocks[:, None] & dim_ok,
+ other=0.0,
+ )
+ per_head_scores = tl.where(
+ (column_slot == slot)[None, :],
+ tl.dot(keys, query_values, out_dtype=tl.float32),
+ per_head_scores,
+ )
+ # The row-wise kernel's head sum is (h0 + h1) + (h2 + h3): pairs, then the pairs.
+ pairs = tl.sum(
+ tl.reshape(
+ tl.maximum(per_head_scores, 0.0), (BLOCK_N, BLOCK_G, NUM_INDEX_HEADS // 2, 2)
+ ),
+ axis=3,
+ )
+ scores = tl.sum(pairs, axis=2)
+ tl.store(
+ output
+ + group_rows[None, :] * output_stride_row
+ + block_columns[:, None] * output_stride_block,
+ scores * score_scale,
+ mask=(block_columns[:, None] < visible[None, :]) & row_ok[None, :],
+ )
+
+
+#: Index-head counts the request-group kernel serves: its head sum is the row-wise kernel's
+#: at up to 8 heads, two a thread then across the lanes; at 16 the row-wise kernel holds
+#: four a thread.
+_GROUP_HEADS = (2, 4, 8)
+_GROUP_BLOCK_N = 128
+
+
+def _takes_request_groups(
+ q: torch.Tensor,
+ only_visible_blocks: bool,
+ context_rows: bool,
+ prefetch: bool,
+ rows_per_request: int,
+) -> bool:
+ """The request-group kernel serves verify-shaped, visible-only calls: runs of
+ `rows_per_request` rows of one request at 2, 4 or 8 index heads. A prefill-sized
+ `context_rows` call keeps the engine's query-tiled kernel."""
+ return (
+ rows_per_request > 1
+ and only_visible_blocks
+ and not prefetch
+ and q.ndim == 3
+ and q.shape[1] in _GROUP_HEADS
+ and not (context_rows and q.shape[0] > _DECODE_ROW_LIMIT)
+ )
+
+
+def _qsa_index_scores_request_groups(
+ q: torch.Tensor,
+ index_cache: torch.Tensor,
+ block_table: torch.Tensor,
+ query_positions: torch.Tensor,
+ request_indices: torch.Tensor,
+ tokens_per_block: int,
+ compress_ratio: int,
+ use_pdl: bool,
+ rows_per_request: int,
+) -> torch.Tensor:
+ """One launch of the request-group kernel after the callable's own checks."""
+ if index_cache.ndim != 4 or index_cache.shape[2] != 1:
+ raise ValueError("QSA index cache must have shape [pages, tokens, 1, dim]")
+ if block_table.ndim != 2 or block_table.dtype != torch.int32:
+ raise ValueError("QSA index scoring requires a two-dimensional int32 block table")
+ rows, num_index_heads, index_head_dim = q.shape
+ if index_head_dim <= 0 or index_cache.shape[3] != index_head_dim:
+ raise ValueError("QSA index query and cache dimensions do not match")
+ if q.dtype != index_cache.dtype:
+ raise ValueError("QSA index query and cache dtypes must match")
+ if query_positions.shape != (rows,) or request_indices.shape != (rows,):
+ raise ValueError("QSA index-scoring metadata must have one entry per row")
+ if query_positions.dtype not in _INTEGER_DTYPES or request_indices.dtype not in _INTEGER_DTYPES:
+ raise ValueError("QSA index-scoring metadata must use integer storage")
+ if tokens_per_block <= 0 or compress_ratio < 2:
+ raise ValueError("QSA index scoring requires positive pages and compression ratio >= 2")
+ for t in (index_cache, block_table, query_positions, request_indices):
+ if t.device != q.device:
+ raise ValueError("QSA paged index scoring tensors must share one CUDA device")
+ max_compressed_blocks = block_table.shape[1] * tokens_per_block // compress_ratio
+ if max_compressed_blocks <= 0:
+ raise ValueError("QSA index scoring requires a non-empty block table")
+ output = torch.empty((rows, max_compressed_blocks), dtype=torch.float32, device=q.device)
+ if rows == 0:
+ return output
+ # A verify step's rows of one request read the same keys: one gather a run of rows,
+ # wider tiles, and a deeper column range once there are runs enough to fill the GPU.
+ groups = triton.cdiv(rows, rows_per_request)
+ tiles = 1 if groups <= 1 else 2 if groups <= 4 else 4
+ grid = (groups, triton.cdiv(max_compressed_blocks, _GROUP_BLOCK_N * tiles))
+ _qsa_index_scores_group_kernel[grid](
+ q,
+ index_cache,
+ block_table,
+ query_positions,
+ request_indices,
+ output,
+ index_head_dim**-0.5,
+ rows,
+ q.stride(0),
+ q.stride(1),
+ q.stride(2),
+ index_cache.stride(0),
+ index_cache.stride(1),
+ index_cache.stride(3),
+ block_table.stride(0),
+ block_table.stride(1),
+ output.stride(0),
+ output.stride(1),
+ NUM_INDEX_HEADS=num_index_heads,
+ INDEX_HEAD_DIM=index_head_dim,
+ BLOCK_D=max(16, triton.next_power_of_2(index_head_dim)),
+ TOKENS_PER_BLOCK=tokens_per_block,
+ COMPRESS_RATIO=compress_ratio,
+ MAX_COMPRESSED_BLOCKS=max_compressed_blocks,
+ GROUP_ROWS=rows_per_request,
+ # The dot's columns: every row's heads, at least the 16 the row-wise kernel pads to.
+ BLOCK_G=max(triton.next_power_of_2(rows_per_request), 16 // num_index_heads, 1),
+ BLOCK_N=_GROUP_BLOCK_N,
+ TILES_PER_PROGRAM=tiles,
+ USE_PDL=use_pdl,
+ num_warps=4,
+ launch_pdl=use_pdl,
+ )
+ return output
+
+
+_CHECKPOINT = dict(
+ heads=4,
+ head_dim=128,
+ tokens_per_block=32,
+ compress_ratio=4,
+ dtype=torch.bfloat16,
+ # `qsa_logical_positions` is int64, `qsa_req_idx_per_token` int32.
+ positions_dtype=torch.int64,
+ request_dtype=torch.int32,
+ # The target's q is a head slice of its packed [T, 4 + 1 + 1, 128] index
+ # scratch buffer, so its row stride is 768 rather than 512.
+ q_layout="scratch_slice",
+)
+
+
+def _spec(**overrides) -> dict:
+ """A cell spec: this checkpoint's geometry unless the cell says otherwise."""
+ return {**_CHECKPOINT, **overrides}
+
+
+#: The serving tables: [8, 260] on the warm-up capture path and [32, 260] in mixed
+#: steps.
+_WIDTH = 260
+
+
+#: The same plane under an NVFP4 KV pool: the QSA position page joins each physical
+#: page after the index buffers (`page_pad` elements), so the trunk's view is
+#: `[pages, 32, 1, 128]` at page stride 12 x 4096 + 192 = 49344.
+_TRUNK_POOL_NVFP4 = dict(pool_buffers=12, page_pad=192)
+
+#: Core's plane (bf16 K/V) at every four-rank layout, as the stock manager allocates it:
+#: the 12 attention layers' `[32, 1, 128]` INDEX_KEY buffers coalesce into one pool, layer k at
+#: byte 8192 k of a 98304-byte page, so one layer's view is `[pages, 32, 1, 128]` at page stride
+#: 12 x 4096 = 49152. The cells take the last layer's plane.
+_TRUNK_POOL_BF16 = dict(pool_buffers=12, page_pad=0)
+
+#: Pages in one rank's pool: the stock engine's `maxNumBlocks` at core's configuration (bf16 K/V
+#: and recurrent state, KV fraction 0.6, max_batch_size 128, max_seq_len 9216, the n-gram table
+#: host-resident). The largest per member; the index plane shares the attention pages.
+_TEP4_POOL_PAGES = 258983
+_DEP4_POOL_PAGES = 120530
+
+#: `QSAAttentionMetadata.qsa_block_table` is `[max_num_requests, max_blocks_per_seq]`: 128 rows
+#: on a tep4 or dep4 rank (the engine's max_num_requests), and ceil((max_seq_len + generation
+#: headroom) / 32) padded to a multiple of 4 -- 292 at core's 9216, 516 at MTP's 16384.
+_TABLE_SLOTS = 128
+_CORE_TABLE_WIDTH = 292
+_MTP_TABLE_WIDTH = 516
+
+#: The longest histories each envelope serves: the last position below max_seq_len.
+_CORE_LAST_POSITION = 9216 - 1
+_MTP_LAST_POSITION = 16384 - 1
+
+#: The target's AgentX server at tp1 (max_seq_len 262144): an 8200-page table (65600 score
+#: columns, 1025 tiles of 64), histories near 100k tokens and to the end.
+_AGENTX_TABLE_WIDTH = 8200
+_AGENTX_CONTEXT_LAST = 100000 - 1
+_AGENTX_LAST_POSITION = 262144 - 1
+
+
+#: The tune's table: max_seq_len 262144 is 8192 pages, and MTP's headroom makes it 8196,
+#: 65568 columns.
+_TUNE_TABLE_WIDTH = 8196
+_TUNE_LAST_POSITION = 262144 - 1
+#: Windows of 0 (positions 0..2), 1, 8, 32 and 33 blocks, ~20k-94k decode contexts, and the
+#: envelope's last position.
+_TUNE_DEP4_POSITIONS = (
+ 94000, 0, 2, 3, 31, 127, 131, 20000, 34800, 49600, 64400, 65535, 79200, 88888, 131071,
+ _TUNE_LAST_POSITION,
+) # fmt: skip
+
+
+def _decode_positions(requests: int, last: int, rows_per_request: int = 1) -> tuple[int, ...]:
+ """Each request's first row position; the first request's last row sits at `last`."""
+ first = last - rows_per_request + 1
+ return (first,) + tuple(1 + (7919 * g) % first for g in range(1, requests))
+
+
+class QsaPagedIndexScores(OpWrapper):
+ """Score every causally visible compressed block of each query row.
+
+ One Triton launch after allocating the output: the entry's own visible-split
+ kernel (`_qsa_index_scores_visible_split_kernel`) for ungrouped decode-sized
+ visible-only calls, its request-group kernel (`_qsa_index_scores_group_kernel`) for
+ verify-shaped ones, otherwise `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels.
+ triton_qsa_paged_index_scores`. For row `r` and compressed block `b` of request
+ `request_indices[r]`::
+
+ anchor = b * compress_ratio + compress_ratio - 1
+ key_b = index_cache[block_table[req, anchor // tpb], anchor % tpb, 0]
+ out[r, b] = sum_h relu(q[r, h] . key_b) * index_head_dim ** -0.5
+
+ The block's key is whatever the compression step left at the group's LAST
+ position; nothing is pooled here. Row `r` sees `(query_positions[r] + 1) //
+ compress_ratio` blocks (complete groups only). Columns at or past that are
+ `-inf` when `only_visible_blocks` is False and **not written** -- the fresh
+ allocation's garbage -- when it is True. The column count is the block
+ table's width, `block_table.shape[1] * tpb // compress_ratio`.
+
+ Kernel choice. With `rows_per_request=1`, `only_visible_blocks`, no `prefetch`,
+ at most 256 rows and an index-head count dividing 16, the visible-split kernel
+ runs: a `(row tiles, request slots, split)` grid sized by the device (16 programs
+ an SM, 32 when the batch is one tile) whose programs stride over each request's
+ visible column tiles, one key gather serving every row of that request among
+ `16 // heads` consecutive rows -- so the grid does not scale with the table's
+ columns. Otherwise the engine chooses a kernel: `context_rows` is a layout
+ declaration ("consecutive rows mostly share a request"), not a phase one. With
+ more than 256 rows and an index-head count that divides 128 into more than one
+ row (1, 2, 4, ..., 64), it selects the query-tiled kernel sharing one key gather
+ across `128 // heads` rows; otherwise the row-wise kernel runs. Before the engine's
+ choice, `rows_per_request` above 1 declares runs of that many rows of one request
+ from row 0 (a verify step's `predicted_tokens_per_seq`): with `only_visible_blocks`,
+ no `prefetch` and 2, 4 or 8 index heads, the entry's request-group kernel shares one
+ key gather across each run (a run whose rows name different requests is gathered
+ row by row). Each cell names the kernel it was
+ measured on. The request-group kernel is bitwise the row-wise one; the
+ visible-split and query-tiled kernels are not.
+
+ Fusion boundary: the table walk, the gather, the dots, relu, head sum,
+ scale and causal bound are inside. Producing `q` (norm, RoPE), writing and
+ compressing the index cache, and everything after the scores are not.
+ Returns a fresh fp32 `[rows, columns]`; nothing is written in place.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="NVFP4 trunk plane, page stride 49344: decode MTP 4 requests x 2 rows",
+ spec=_spec(
+ layout="decode",
+ gen_positions=(7000, 5000, 4200, 2100),
+ rows_per_request=2,
+ slots=32,
+ pages=_WIDTH,
+ only_visible_blocks=True,
+ context_rows=False,
+ call_rows_per_request=2,
+ branch="request_group",
+ **_TRUNK_POOL_NVFP4,
+ ),
+ ),
+ Cell(
+ why="decode MTP7: 32 slots x 8 rows = 256 rows, the last decode-sized count",
+ spec=_spec(
+ layout="decode",
+ gen_positions=tuple(100 + (499 * g) % 8000 for g in range(32)),
+ rows_per_request=8,
+ slots=32,
+ pages=_WIDTH,
+ only_visible_blocks=True,
+ context_rows=False,
+ call_rows_per_request=8,
+ branch="request_group",
+ ),
+ ),
+ # The target's AgentX verify steps at tp1 (MTP3: 4 rows a request) over the NVFP4 trunk
+ # plane and its 8200-page table, NaN on every page no request maps.
+ Cell(
+ why="AgentX MTP3 verify: 16 requests x 4 rows to 99999, 65600 columns",
+ spec=_spec(
+ layout="decode",
+ gen_positions=_decode_positions(16, _AGENTX_CONTEXT_LAST, rows_per_request=4),
+ rows_per_request=4,
+ call_rows_per_request=4,
+ slots=32,
+ pages=_AGENTX_TABLE_WIDTH,
+ pool_pages=60000,
+ only_visible_blocks=True,
+ context_rows=False,
+ use_pdl=True,
+ branch="request_group",
+ **_TRUNK_POOL_NVFP4,
+ ),
+ ),
+ Cell(
+ why="AgentX MTP3 verify: 16 requests x 4 rows to 262143, 65600 columns",
+ spec=_spec(
+ layout="decode",
+ gen_positions=_decode_positions(16, _AGENTX_LAST_POSITION, rows_per_request=4),
+ rows_per_request=4,
+ call_rows_per_request=4,
+ slots=32,
+ pages=_AGENTX_TABLE_WIDTH,
+ pool_pages=140000,
+ only_visible_blocks=True,
+ context_rows=False,
+ use_pdl=True,
+ branch="request_group",
+ **_TRUNK_POOL_NVFP4,
+ ),
+ ),
+ Cell(
+ why="AgentX MTP3 verify: 1 request x 4 rows to 262143, 65600 columns",
+ spec=_spec(
+ layout="decode",
+ gen_positions=_decode_positions(1, _AGENTX_LAST_POSITION, rows_per_request=4),
+ rows_per_request=4,
+ call_rows_per_request=4,
+ slots=32,
+ pages=_AGENTX_TABLE_WIDTH,
+ pool_pages=60000,
+ only_visible_blocks=True,
+ context_rows=False,
+ use_pdl=True,
+ branch="request_group",
+ **_TRUNK_POOL_NVFP4,
+ ),
+ ),
+ Cell(
+ why="capture: 8191 rows over 8 ragged requests, visible-only",
+ spec=_spec(
+ layout="capture",
+ lengths=(1024, 512, 2048, 256, 1500, 700, 1151, 1000),
+ slots=8,
+ pages=_WIDTH,
+ only_visible_blocks=True,
+ context_rows=True,
+ branch="query_tile",
+ ),
+ ),
+ # The four-rank layouts' sites over one rank's measured pool: the call-site flags each
+ # site passes (the decode block's PDL launch, with prefetch at up to 4 rows), the
+ # engine's [128, width] table with requests on seeded rows, and pages drawn over the
+ # whole pool, both ends included.
+ Cell(
+ why="tep4 decode: one request at position 9215, prefetch, 258983-page pool",
+ spec=_spec(
+ layout="decode",
+ gen_positions=_decode_positions(1, _CORE_LAST_POSITION),
+ rows_per_request=1,
+ slots=_TABLE_SLOTS,
+ pages=_CORE_TABLE_WIDTH,
+ pool_pages=_TEP4_POOL_PAGES,
+ only_visible_blocks=True,
+ context_rows=False,
+ use_pdl=True,
+ prefetch=True,
+ branch="row_wise",
+ **_TRUNK_POOL_BF16,
+ ),
+ ),
+ Cell(
+ why="tep4 decode: 128 requests x 1 row, the largest captured batch",
+ spec=_spec(
+ layout="decode",
+ gen_positions=_decode_positions(128, _CORE_LAST_POSITION),
+ rows_per_request=1,
+ slots=_TABLE_SLOTS,
+ pages=_CORE_TABLE_WIDTH,
+ pool_pages=_TEP4_POOL_PAGES,
+ only_visible_blocks=True,
+ context_rows=False,
+ use_pdl=True,
+ branch="visible_split",
+ **_TRUNK_POOL_BF16,
+ ),
+ ),
+ Cell(
+ why="tep4 MTP7 verify: 128 requests x 8 rows = 1024, table 516, to 16383",
+ spec=_spec(
+ layout="decode",
+ gen_positions=_decode_positions(128, _MTP_LAST_POSITION, rows_per_request=8),
+ rows_per_request=8,
+ slots=_TABLE_SLOTS,
+ pages=_MTP_TABLE_WIDTH,
+ pool_pages=_TEP4_POOL_PAGES,
+ only_visible_blocks=True,
+ context_rows=False,
+ use_pdl=True,
+ branch="row_wise",
+ **_TRUNK_POOL_BF16,
+ ),
+ ),
+ Cell(
+ why="tep4 context: one fresh 8192-token prompt (max_num_tokens)",
+ spec=_spec(
+ layout="capture",
+ lengths=(8192,),
+ slots=_TABLE_SLOTS,
+ pages=_CORE_TABLE_WIDTH,
+ pool_pages=_TEP4_POOL_PAGES,
+ only_visible_blocks=True,
+ context_rows=True,
+ branch="query_tile",
+ **_TRUNK_POOL_BF16,
+ ),
+ ),
+ Cell(
+ why="dep4 decode: 32 requests x 1 row over a 120530-page pool",
+ spec=_spec(
+ layout="decode",
+ gen_positions=_decode_positions(32, _CORE_LAST_POSITION),
+ rows_per_request=1,
+ slots=_TABLE_SLOTS,
+ pages=_CORE_TABLE_WIDTH,
+ pool_pages=_DEP4_POOL_PAGES,
+ only_visible_blocks=True,
+ context_rows=False,
+ use_pdl=True,
+ branch="visible_split",
+ **_TRUNK_POOL_BF16,
+ ),
+ ),
+ Cell(
+ why="envelope: an 8192-row chunk at positions 1024..9215, table column 287",
+ spec=_spec(
+ layout="decode",
+ gen_positions=(_CORE_LAST_POSITION + 1 - 8192,),
+ rows_per_request=8192,
+ slots=_TABLE_SLOTS,
+ pages=_CORE_TABLE_WIDTH,
+ pool_pages=_TEP4_POOL_PAGES,
+ only_visible_blocks=True,
+ context_rows=True,
+ branch="query_tile",
+ **_TRUNK_POOL_BF16,
+ ),
+ ),
+ # The tune's envelope (max_seq_len 262144): an 8196-page table, 65568 columns.
+ Cell(
+ why="tune ll MTP3 verify: 1 request x 4 rows at 84000..84003, 8196-page table",
+ spec=_spec(
+ layout="decode",
+ gen_positions=(84000,),
+ rows_per_request=4,
+ slots=_TABLE_SLOTS,
+ pages=_TUNE_TABLE_WIDTH,
+ pool_pages=_TEP4_POOL_PAGES,
+ only_visible_blocks=True,
+ context_rows=False,
+ branch="visible_split",
+ **_TRUNK_POOL_BF16,
+ ),
+ ),
+ Cell(
+ why="tune draft: 1 row at 262143, the envelope's last position",
+ spec=_spec(
+ layout="decode",
+ gen_positions=(_TUNE_LAST_POSITION,),
+ rows_per_request=1,
+ slots=_TABLE_SLOTS,
+ pages=_TUNE_TABLE_WIDTH,
+ pool_pages=_TEP4_POOL_PAGES,
+ only_visible_blocks=True,
+ context_rows=False,
+ branch="visible_split",
+ **_TRUNK_POOL_BF16,
+ ),
+ ),
+ Cell(
+ why="tune MTP7 verify: 4 requests x 8 rows, the first to 262143",
+ spec=_spec(
+ layout="decode",
+ gen_positions=_decode_positions(4, _TUNE_LAST_POSITION, rows_per_request=8),
+ rows_per_request=8,
+ slots=_TABLE_SLOTS,
+ pages=_TUNE_TABLE_WIDTH,
+ pool_pages=_TEP4_POOL_PAGES,
+ only_visible_blocks=True,
+ context_rows=False,
+ use_pdl=True,
+ branch="visible_split",
+ **_TRUNK_POOL_BF16,
+ ),
+ ),
+ Cell(
+ why="tune dep4 decode: 16 requests x 1 row, empty to 262143-position windows",
+ spec=_spec(
+ layout="decode",
+ gen_positions=_TUNE_DEP4_POSITIONS,
+ rows_per_request=1,
+ slots=_TABLE_SLOTS,
+ pages=_TUNE_TABLE_WIDTH,
+ pool_pages=_DEP4_POOL_PAGES,
+ only_visible_blocks=True,
+ context_rows=False,
+ branch="visible_split",
+ **_TRUNK_POOL_BF16,
+ ),
+ ),
+ Cell(
+ why="tune context step: 2 rows of one request at 84000 (the snapshot split)",
+ spec=_spec(
+ layout="decode",
+ gen_positions=(84000,),
+ rows_per_request=2,
+ slots=_TABLE_SLOTS,
+ pages=_TUNE_TABLE_WIDTH,
+ pool_pages=_TEP4_POOL_PAGES,
+ only_visible_blocks=True,
+ context_rows=True,
+ branch="visible_split",
+ **_TRUNK_POOL_BF16,
+ ),
+ ),
+ Cell(
+ why="small context: fresh 200- and 56-token prompts, 256 rows from position 0",
+ spec=_spec(
+ layout="capture",
+ lengths=(200, 56),
+ slots=_TABLE_SLOTS,
+ pages=_TUNE_TABLE_WIDTH,
+ pool_pages=_DEP4_POOL_PAGES,
+ only_visible_blocks=True,
+ context_rows=True,
+ branch="visible_split",
+ **_TRUNK_POOL_BF16,
+ ),
+ ),
+ )
+
+ note = """
+ Source: the entry's own Triton kernel (`_qsa_index_scores_visible_split_kernel`,
+ this file) and the engine's beside the wrapped callable
+ (`_qsa_paged_index_scores_kernel`, `_qsa_paged_index_scores_qtile_kernel`,
+ `_qsa_paged_index_scores_group_kernel`),
+ JIT-compiled for the running device on first call; the cells ran that JIT
+ path on sm_103.
+
+ Visible-split kernel (decode-sized visible-only calls). Its grid is
+ `(row tiles, request slots, split)`: a row tile is `min(rows, 16 // heads)`
+ consecutive rows, program `(t, g, s)` takes the g-th distinct request of
+ tile t (a slot past the tile's request count exits at once) and its column
+ tiles s, s + split, ... of 64 blocks up to that request's bound, `split`
+ fixed by the SM count (16 programs an SM, 32 for a one-tile batch) and the
+ table, never by the positions -- so a captured graph follows any window
+ (pinned by `test_captured_visible_split_follows_the_positions`). The dot is
+ `[64, 128] x [128, rows x heads]` with fp32 accumulation, relu, a per-row
+ head sum and the scale, as in the engine's kernels; the columns written are
+ exactly each row's visible ones. `use_pdl` waits on its predecessor before the first read and
+ releases dependents at once; `prefetch` keeps the engine's kernel (the
+ visible-split one does not read before the wait).
+
+ Metadata consumed: none -- every input is an argument. What the target
+ passes: `block_table` is `QSAAttentionMetadata.qsa_block_table` (raw
+ pool-page ids, int32, -1 past each sequence); `query_positions` /
+ `request_indices` are `qsa_logical_positions` / `qsa_req_idx_per_token`;
+ `index_cache` is the layer's bf16 index-K plane (`get_index_k_buffer`).
+
+ Preconditions (not checked): every block-table entry the causal bound
+ reaches is a valid page, and `request_indices[r]` is a row of the table --
+ the loads are masked by the bound, not by the pool size, so a violation is
+ an out-of-bounds read (not probed). The compression step must already have
+ run: an uncompressed anchor scores as a raw index key. `rows == 0` returns
+ an empty `[0, columns]` tensor without launching.
+
+ The op validates its own metadata loudly (`ValueError`s; the test pins the
+ one-entry-per-row metadata refusal), so `is_valid` adds nothing.
+
+ `only_visible_blocks=True` really leaves garbage in the whole unscored tail
+ (measured, not pinned by a test). A consumer that scans `[row_start,
+ row_end)` -- `qsa_topk_select` -- is unaffected.
+
+ fp32 operands are accepted but `tl.dot` rounds them to a tensor-core
+ format: 8.1e-3 max abs against an fp32 reference on scores to 7.8 (measured),
+ outside the test's band. No cell covers fp32; the side cache is bf16.
+
+ CUDA graphs: the test captures the decode launch (`use_pdl`, `prefetch`) at
+ 4 rows and replays it after `q` and the anchor keys were rewritten in
+ place; every replay matches an eager call bitwise. The same holds captured
+ at each measured-pool decode and verify cell (1, 128 and 1024 rows at tep4,
+ 32 at dep4, the AgentX MTP3 verify cells' 4 and 64 rows, with their sites' flags)
+ after q and every mapped key changed.
+
+ The request-group kernel (`rows_per_request` > 1): each program takes one run of
+ rows and 128 blocks a tile (1, 2 or 4 tiles a program at 1, up to 4, and more runs);
+ the dot's columns are the run's rows' heads, padded to 16 as the row-wise kernel
+ pads its heads, and the head sum is the row-wise kernel's pairwise order, so every
+ score is bitwise the row-wise kernel's -- the test compares the bit patterns at 2,
+ 3, 4 and 8 rows a run, 2, 4 and 8 heads, partial last runs and runs that mix
+ requests; 16 heads stay row-wise.
+ The run reads up to its longest row's visible blocks and stores each row's own.
+
+ At the measured-pool cells: at most 2.9e-6 against the reference (scores to
+ 11.3), relative L2 below 6.4e-8.
+
+ `prefetch` (with `use_pdl` and `only_visible_blocks`, at generation-sized
+ batches) reads the pages and compressed keys before the PDL wait and q after
+ it. Sound only behind the decode step's own prologue
+ (`qwen4_exp_attn_decode_prologue`), which writes no compressed key but the last
+ visible block's, at a row whose position ends a group; that key is read again
+ after the wait.
+
+ `prefetch` takes effect only where the call is one tile a program -- 256 rows or
+ fewer -- and is otherwise dropped silently: a 1024-row verify step passing it
+ runs the plain PDL kernel.
+
+ Every cell runs `only_visible_blocks=True` (the target never passes False),
+ 4 index heads x 128, compress ratio 4, 32-token pages, bf16 operands, int64
+ positions and int32 request indices. The NVFP4-pool page stride (49344) is
+ certified on the row-wise kernel at the tp1 cell's small pool and on the
+ request-group kernel at the AgentX cells' 60000- and 140000-page pools.
+
+ The four-rank layouts' cells run over one rank's measured pool, core's coalesced
+ plane at page stride 49152 (`page_stride` is a compile-time constant, so this
+ is its own kernel variant): 258983 pages at tep4, 120530 at dep4, the plane
+ poisoned with NaN outside the mapped pages, which include page 0 and the
+ pool's last page. The last page's elements sit near 1.27e10 elements from the
+ plane's base: both kernels widen the page and request index to int64 before
+ the stride product, and the cells prove it.
+
+ Index arithmetic left in int32: `row * q_stride_row` and the output's
+ `row * output_stride_row` (program ids times compile-time strides). The output
+ term passes 2^31 at rows x columns >= 2^31 -- 919k rows at 2336 columns; the
+ envelope's 8192 rows x 65568 columns (max_seq_len 262144) is 5.4e8. Positions
+ are read at their own width (int64 here), so the fp32 2^24 bound does not apply.
+
+ The visible-split kernel's index arithmetic: pages and request rows are widened
+ to int64 before the stride products (the measured-pool cells read the pool's last
+ page); `lane_rows * output_stride_row` stays int32, at most 256 x 65568 here.
+
+ Certified: up to 8192 rows (one request's 8192-row context, fresh or ending at
+ position 9215, on the query-tile kernel; 1024 verify rows on the row-wise
+ kernel with 8 tiles a program), tables of 128 x 292 and 128 x 516 pages,
+ positions to 16383; on the visible-split kernel also the tune's 128 x 8196
+ table (65568 columns) with positions to 262143 at 1, 2, 4, 16 and 32 rows
+ (empty windows, 1 and 2 rows a request, 8-row requests across two tiles, small
+ context steps from position 0); on the request-group kernel 2, 4 and 8 rows a
+ request, up to 256 rows, a 32 x 8200 table and positions to 262143. Not certified:
+ more than 8192 rows, tables wider than 516 pages on the row-wise and query-tile
+ kernels (the tune's context-sized calls), any other head geometry, compress ratio,
+ page size or operand dtype, a context step's mixed request rows on the pool
+ layout, the fp8 side-plane stride at a measured pool, and the `-inf` padded mode.
+ """
+
+ def __call__(
+ self,
+ q: torch.Tensor,
+ index_cache: torch.Tensor,
+ block_table: torch.Tensor,
+ query_positions: torch.Tensor,
+ request_indices: torch.Tensor,
+ tokens_per_block: int,
+ compress_ratio: int,
+ only_visible_blocks: bool,
+ context_rows: bool,
+ use_pdl: bool = False,
+ prefetch: bool = False,
+ rows_per_request: int = 1,
+ ) -> torch.Tensor:
+ if rows_per_request < 1:
+ raise ValueError("QSA index scoring requires rows_per_request >= 1")
+ if rows_per_request == 1 and _takes_visible_split(q, only_visible_blocks, prefetch):
+ return _qsa_index_scores_visible_split(
+ q,
+ index_cache,
+ block_table,
+ query_positions,
+ request_indices,
+ tokens_per_block,
+ compress_ratio,
+ use_pdl,
+ )
+ if _takes_request_groups(q, only_visible_blocks, context_rows, prefetch, rows_per_request):
+ return _qsa_index_scores_request_groups(
+ q,
+ index_cache,
+ block_table,
+ query_positions,
+ request_indices,
+ tokens_per_block,
+ compress_ratio,
+ use_pdl,
+ rows_per_request,
+ )
+ return _triton_qsa_paged_index_scores(
+ q=q,
+ index_cache=index_cache,
+ block_table=block_table,
+ query_positions=query_positions,
+ request_indices=request_indices,
+ tokens_per_block=tokens_per_block,
+ compress_ratio=compress_ratio,
+ only_visible_blocks=only_visible_blocks,
+ context_rows=context_rows,
+ use_pdl=use_pdl,
+ prefetch=prefetch,
+ )
+
+ def reference(
+ self,
+ q: torch.Tensor,
+ index_cache: torch.Tensor,
+ block_table: torch.Tensor,
+ query_positions: torch.Tensor,
+ request_indices: torch.Tensor,
+ tokens_per_block: int,
+ compress_ratio: int,
+ only_visible_blocks: bool,
+ context_rows: bool,
+ use_pdl: bool = False,
+ prefetch: bool = False,
+ rows_per_request: int = 1,
+ ) -> torch.Tensor:
+ """fp32 scores, `-inf` past each row's bound -- or NaN there, meaning
+ "unspecified", when `only_visible_blocks` is True. `context_rows` and
+ `rows_per_request` are ignored: they pick a kernel, not a result.
+
+ Walked in row chunks only to bound the gathered keys, which reach
+ 8.7 GiB in fp32 at 8191 rows x 2080 columns.
+ """
+ rows, _, head_dim = q.shape
+ num_columns = block_table.shape[1] * tokens_per_block // compress_ratio
+ blocks = torch.arange(num_columns, device=q.device)
+ anchors = blocks * compress_ratio + compress_ratio - 1
+ logical_pages, in_page = anchors // tokens_per_block, anchors % tokens_per_block
+ table = block_table.long()
+ requests = request_indices.long()
+ scores = torch.empty(rows, num_columns, dtype=torch.float32, device=q.device)
+ chunk = max(1, 2**26 // max(num_columns * head_dim, 1))
+ for lo in range(0, rows, chunk):
+ hi = min(lo + chunk, rows)
+ # Pages past a sequence are -1; their columns are masked below.
+ physical = table[requests[lo:hi]][:, logical_pages].clamp_min(0)
+ keys = index_cache[physical, in_page.expand(hi - lo, num_columns)][:, :, 0, :]
+ per_head = torch.einsum("rhd,rbd->rbh", q[lo:hi].float(), keys.float())
+ scores[lo:hi] = torch.relu(per_head).sum(-1) * (head_dim**-0.5)
+ visible = blocks[None, :] < ((query_positions.long() + 1) // compress_ratio)[:, None]
+ tail = float("nan") if only_visible_blocks else float("-inf")
+ return scores.masked_fill(~visible, tail)
+
+
+qsa_paged_index_scores = QsaPagedIndexScores()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_index_store.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_index_store.py
new file mode 100644
index 000000000000..a014e6a74314
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_index_store.py
@@ -0,0 +1,473 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""QSA index side-cache store: raw index keys and position coordinates into their pages."""
+
+from __future__ import annotations
+
+from typing import Tuple
+
+import torch
+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."""
+ 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
+#: buffer; `slots` x `width` is `qsa_block_table`. `pages` is the pool.
+_CKPT = dict(
+ hkv=1,
+ head_dim=128,
+ index_heads=4,
+ tokens_per_block=32,
+ axes=3,
+ slots=32,
+ width=260,
+ key_dtype=torch.bfloat16,
+ request_dtype=torch.int32,
+ position_dtype=torch.int64,
+ key_layout="packed",
+ cache_layout="plain",
+ hole=None,
+)
+
+#: The target's serving envelope: ISL 8192 + OSL 1024 puts decode rows up to position
+#: 9215 (page column 287), 128 concurrent requests, and `max_num_tokens` 16384
+#: rows in one prefill or mixed step.
+_WIDE = dict(slots=160, width=300)
+_PREFILL_16K = ((0, 8192), (0, 8192))
+
+#: 32 requests of one MTP7 verify block each (8 rows), at starts that put some blocks across a
+#: page edge and some inside one.
+_MTP7_32 = tuple((29 + 61 * i, 8) for i in range(32))
+
+
+def _spec(segments, pages, **overrides) -> dict:
+ """`segments` is `(first position, rows)` per request, in row order."""
+ return {**_CKPT, "segments": segments, "pages": pages, **overrides}
+
+
+#: The default target (tp4 / tep4 / dep4) at one rank. Its fused input projection row
+#: `[q | k | v | gate | index q, k | pad of one index head]` is 4352 wide under attention TP (6 q
+#: and 1 KV heads of 256: 34 heads of 128, the raw index key head 32) and 14080 under attention DP
+#: (110 heads, key head 108). Its bf16 index pool is the stock manager's at every member: the 12
+#: sparse layers' INDEX_KEY buffers coalesce into one pool, so a layer's plane is `[pages, 32, 1,
+#: 128]` at page stride 12 x 4096 = 49152 elements, layer `j` at `j` x 8192 bytes into the slot,
+#: and the position plane is its own contiguous `[pages, 32, 3]` int32 pool (`cache_layout`
+#: "coalesced", the last layer's plane). `pages` is the tep4 pool: the stock tep4 engine's
+#: (bf16 KV) `kv_max_tokens` 7866720 a rank, 245835 pages of 32. The table
+#: is `qsa_block_table`, `[max_batch_size, max_blocks_per_seq]`: 128 x 292 (`max_seq_len` 9216) or
+#: 128 x 516 under MTP (16384). `straddle` puts the batch's last row on the pool's last page, past
+#: 2^31 elements of the index plane.
+_DEFAULT = dict(
+ _CKPT,
+ key_layout="fused_row",
+ cache_layout="coalesced",
+ layers=12,
+ layer=11,
+ page_stride=49152,
+ slots=128,
+ width=292,
+ pages=245835,
+ straddle=True,
+)
+
+#: 128 decode requests, one row each, at positions on both sides of page edges; the last at
+#: 9215 (`max_seq_len` - 1, table column 287).
+_DECODE_128 = tuple((31 + 72 * i, 1) for i in range(127)) + ((9215, 1),)
+#: 32 decode requests (one dep4 rank's share of 128), the last at 9215.
+_DECODE_32 = tuple((31 + 287 * i, 1) for i in range(31)) + ((9215, 1),)
+#: 128 MTP7 verify blocks of 8 rows, some across a page edge, up to position 16165.
+_MTP7_128 = tuple((29 + 127 * i, 8) for i in range(128))
+
+
+class QsaPagedIndexStore(OpWrapper):
+ """Write each row's raw index key and position coordinates into its paged slot.
+
+ For row t, with `page = block_table[request_indices[t], logical_positions[t]
+ // tokens_per_block]` and `slot = logical_positions[t] % tokens_per_block`::
+
+ index_cache[page, slot] = token_k[t] # cast to the cache dtype
+ position_cache[page, slot, :A] = position_coordinates[t] # A = coordinates' width
+
+ 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: one Triton launch (`_qsa_paged_index_store_kernel`). Nothing in
+ the engine performs this store for a packed, ragged batch: the registered
+ index-cache ops write fp8 planes or norm and rotate on the way in,
+ `triton_qsa_decode_pre_indexer` is one row per request, and
+ `QSAIndexer.update_cache_and_compress` bundles the compression.
+
+ Fusion boundary: page lookup, slot arithmetic, both scatters. The index
+ projection, the query's norm and rotation, group compression, scoring,
+ selection, sparse attention and the main K/V pool are the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ 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 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 "
+ "per-slot buffer holding 12 index-K pages and the position page; one row "
+ "on an unallocated (-1) table entry is skipped",
+ spec=_spec(
+ ((0, 1), (31, 1), (32, 1), (33, 40)),
+ pages=8,
+ cache_layout="shared_page",
+ layers=12,
+ page_stride=49344,
+ position_strides=(24672, 3, 1),
+ hole="unallocated_page",
+ ),
+ ),
+ Cell(
+ why="default tep4 decode: one row at position 9215 on the pool's last page, token_k "
+ "off the 4352-wide fused row, the coalesced bf16 index plane",
+ spec=dict(_DEFAULT, member="tep4", segments=((9215, 1),), row_width=4352),
+ ),
+ Cell(
+ why="default tep4 decode: 128 requests one row each (the largest captured batch), "
+ "4352-wide row",
+ spec=dict(_DEFAULT, member="tep4", segments=_DECODE_128, row_width=4352),
+ ),
+ Cell(
+ why="default tep4 MTP7 verify: 128 requests x 8 rows (1024), 16 rows per program, "
+ "a 516-column table",
+ spec=dict(_DEFAULT, member="tep4", segments=_MTP7_128, row_width=4352, width=516),
+ ),
+ Cell(
+ why="default tp4 step of 8192 rows (max_num_tokens): a context from 0 and a decode "
+ "row at 9215, 4352-wide row",
+ spec=dict(_DEFAULT, member="tp4", segments=((0, 8191), (9215, 1)), row_width=4352),
+ ),
+ Cell(
+ why="default dep4 decode: 32 requests one row each on a 128-row table, token_k off "
+ "the 14080-wide row",
+ spec=dict(_DEFAULT, member="dep4", segments=_DECODE_32, row_width=14080),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none -- every input is an argument. The target passes
+ `qsa_req_idx_per_token` (int32) and `qsa_logical_positions` (int64)
+ sliced to the step's rows, `qsa_block_table` unsliced, the layer's
+ `get_index_k_buffer(i)` (`[pages, tpb, 1, 128]` bf16) and the shared
+ `get_qsa_position_buffer()` (`[pages, tpb, 3]` int32, rewritten
+ idempotently by every sparse layer). Every operand is read through its
+ own strides, all of them runtime arguments; the cells drive the target's:
+ `token_k` off the packed index projection (row stride 640) and off the fused
+ input projection (14080), the `[T, 3]` transpose of the position buffer as
+ the coordinates, contiguous metadata and table. Under an
+ NVFP4 KV pool the QSA cache manager packs each page's index-K roles and
+ then its position page into one per-slot buffer, so both planes are strided
+ views of one allocation: the trunk's index plane is `[pages, 32, 1, 128]`
+ at page stride 49344 (12 x 4096 + 192) with the position plane
+ `[pages, 32, 3]` at (24672, 3, 1); the cell checks that the other index-K
+ roles of each page survive the call.
+
+ The default target (`_DEFAULT` cells) under bf16 KV, read off the stock
+ manager at tp4, tep4 and dep4 (the test re-reads it): each sparse layer's
+ `get_index_k_buffer(i)` is `[pages, 32, 1, 128]` at strides (49152, 128,
+ 128, 1), the 12 layers' planes 8192 bytes apart in one coalesced pool, and
+ `get_qsa_position_buffer()` a contiguous `[pages, 32, 3]` int32 pool. Under
+ fp8 and NVFP4 KV both are the shared-page layout above (49344 / 24672) at
+ every member. The cells drive the last layer's plane of a 245835-page pool
+ (the stock tep4 engine's), its other 11 planes checked untouched, with the
+ batch's last row on the pool's last page: index-plane offsets past 2^31
+ elements (1.2e10). `token_k` is the index-K head off the fused input
+ projection row, 4352 wide under attention TP (head 32) and 14080 under
+ attention DP (head 108); the table is `[128, 292]` (`max_seq_len` 9216) or
+ `[128, 516]` (MTP's `max_seq_len` 16384).
+
+ Preconditions:
+ - destination slots are unique across the batch: two rows naming one
+ (page, slot) both store, unordered. The engine's mapping cannot
+ collide.
+ Loud rejections (`ValueError`, before launching; the test pins the
+ unequal-length metadata one): metadata that is not
+ one-dimensional, of unequal length, or not int32/int64; `token_k` or the
+ coordinates with a row count other than the metadata's; an index plane
+ whose page size, head count or head width disagrees with
+ `tokens_per_block` and `token_k`; a position plane whose page size
+ disagrees or that is narrower than the coordinates; a table that is not a
+ 2-D integer tensor. `rows == 0` returns before launching.
+
+ Values the caller owns -- data, not metadata (measured, not pinned by a
+ test): coordinates wider than the plane's dtype wrap (int64 `2^31 + 5`
+ stores as -2147483643); floats truncate toward zero.
+
+ One launch per call; HEAD_DIM, the head count, the page size and the
+ coordinate width are constexprs, everything else (strides, pool and table
+ extents, rows) is a runtime argument and the row count and extents are
+ not specialized on, so a new batch or pool size does not recompile
+ (strides still specialize on 1 and 16-divisibility). Up to 256 rows run
+ one row per program (1 warp), above that 16 rows per program (4 warps).
+ CUDA-graph capturable (captured at 256 rows, 32 MTP7 verify blocks); a
+ replay follows positions rewritten in place.
+
+ Hard limits: every index is widened to int64 inside the kernel before
+ any arithmetic; an int64 position at `2^32 + 5` lands past the table and
+ is skipped rather than truncated (read from the kernel, not pinned by a
+ test). Rows are a grid axis: 2^31 programs.
+
+ Skip certified by a cell: an unallocated (`-1`) table entry, one row of the
+ NVFP4-plane cell. The other skip conditions (a negative request index among
+ them) share that mask and have no cell.
+
+ Not certified: an fp8 index plane, keys other than bf16, any index geometry
+ other than 1 head x 128 at 32-token pages, decode batches other than 1, 32
+ and 128 requests and verify batches other than 32 and 128 MTP7 blocks,
+ strided metadata or tables, stride-0 coordinates, multiple devices,
+ concurrent streams onto one pool, duplicate destinations. The tp4 and dep4
+ pools are driven at tep4's page count, not their own: both hold at most
+ that many pages a rank (the same or larger weights a rank, the same or
+ larger KV page).
+ """
+
+ def __call__(
+ self,
+ token_k: torch.Tensor,
+ position_coordinates: torch.Tensor,
+ index_cache: torch.Tensor,
+ position_cache: torch.Tensor,
+ request_indices: torch.Tensor,
+ logical_positions: torch.Tensor,
+ block_table: torch.Tensor,
+ tokens_per_block: int,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ 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
+
+ def reference(
+ self,
+ token_k: torch.Tensor,
+ position_coordinates: torch.Tensor,
+ index_cache: torch.Tensor,
+ position_cache: torch.Tensor,
+ request_indices: torch.Tensor,
+ logical_positions: torch.Tensor,
+ block_table: torch.Tensor,
+ tokens_per_block: int,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Both planes after the store, as fresh tensors; the inputs are not touched.
+
+ 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()
+ pool = index_cache.shape[0]
+ kept, pages, slots = [], [], []
+ rows = zip(request_indices.tolist(), logical_positions.tolist())
+ for row, (request, position) in enumerate(rows):
+ 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)
+ )
+ 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/ops/attention/qsa_paged_kv_store.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_kv_store.py
new file mode 100644
index 000000000000..5d5af38707ad
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_kv_store.py
@@ -0,0 +1,314 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""QSA sparse path's paged K/V append: this step's rows into the HND main pool."""
+
+from __future__ import annotations
+
+from typing import Tuple
+
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+# Tier-2 sourcing: a vendored Triton callable reached as ordinary Python. QSA's
+# kernel plane has no torch.ops surface.
+from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels as _qsa_kernels
+
+#: This checkpoint's main K/V geometry and the engine's QSA metadata. `heads_q`
+#: places k and v as head slices of the packed `[T, 24 + 2 + 2, 256]` QKV
+#: buffer; `slots` x `width` is `qsa_block_table` (max_num_sequences x
+#: max_blocks_per_seq). `pages` is the pool, independent of the batch.
+_CKPT = dict(
+ hkv=2,
+ head_dim=256,
+ heads_q=24,
+ tokens_per_block=32,
+ slots=32,
+ width=260,
+ row_dtype=torch.bfloat16,
+ cache_dtype=torch.bfloat16,
+ request_dtype=torch.int32,
+ position_dtype=torch.int64,
+ hole=None,
+ # Heads of further columns after v in each row: 27 is the target's fused input
+ # projection row [q | k | v | gate | index | pad] (55 heads of 256).
+ row_pad_heads=0,
+)
+
+_F8 = torch.float8_e4m3fn
+
+_DECODE = ((0, 1), (31, 1), (32, 1), (33, 1), (100, 1), (255, 1), (7, 1), (64, 1))
+
+#: 32 requests of one MTP7 verify block each (8 rows), at starts that put some blocks across a
+#: page edge and some inside one.
+_MTP7_32 = tuple((29 + 61 * i, 8) for i in range(32))
+
+#: One rank's QSA K/V pool as the stock `Qwen4ExpHybridCacheManager` lays it out at the default target's
+#: members: a page holds the K and V role planes of all 12 attention layers (24 role pages of
+#: `[hkv, 32, 256]`), so a layer's `get_buffers(i, kv_layout="HND")` view steps 24 role pages a
+#: page. `pool="engine"` cells drive layer 47's view (`layer_slot` 11, the page's last planes) at
+#: an engine's own page count. `slots` x `width` is the engine's `qsa_block_table`:
+#: `max_num_sequences` 128 by `max_blocks_per_seq` 292 at `max_seq_len` 9216 (516 at 16384).
+_ENGINE = dict(_CKPT, pool="engine", role_pages=24, layer_slot=11, slots=128, width=292)
+
+#: A tp4 / tep4 rank: one KV head, k and v at heads 6 and 7 of the fused projection row
+#: `[q 6 | k 1 | v 1 | gate 6 | index 3]` (17 heads of 256, row stride 4352).
+_TP4 = dict(_ENGINE, hkv=1, heads_q=6, row_pad_heads=9, row_stride=4352)
+
+#: A dep4 rank (attention DP, every head): the tp1 row `[q 24 | k 2 | v 2 | gate 24 | index 3]`
+#: (55 heads, row stride 14080) over the kv_fp8 variant's e4m3 pool.
+_DEP4_FP8 = dict(_ENGINE, row_pad_heads=27, row_stride=14080, cache_dtype=torch.float8_e4m3fn)
+
+#: Attention pool pages of an engine at KV fraction 0.6, bf16 recurrent state, `max_batch_size`
+#: 128 and `max_seq_len` 9216: its `maxNumBlocks` less the recurrent pool's slots, which that
+#: total also counts. tep4, bf16 KV: 258983 - 1874. dep4, kv_fp8: 172579 - 1217.
+PAGES_TEP4 = 257109
+PAGES_DEP4_FP8 = 171362
+
+#: 128 decode rows, one a request: page-aligned, mid-page, page-crossing and the last position
+#: `max_seq_len` 9216 admits, then a stride-577 walk over the rest.
+_DECODE_128 = ((0, 1), (31, 1), (32, 1), (9215, 1)) + tuple(
+ ((i * 577 + 45) % 9216, 1) for i in range(124)
+)
+
+#: 32 decode rows, one a request, the last at position 9215.
+_DECODE_32 = tuple(((i * 293 + 5) % 9216, 1) for i in range(31)) + ((9215, 1),)
+
+#: 128 requests of one MTP7 verify block each (8 rows, 1024), the last ending at position 16383
+#: (`max_seq_len` 16384 under MTP); starts at every offset in a page, so blocks straddle pages.
+_MTP7_128 = tuple(((29 + 127 * i) % 16376, 8) for i in range(127)) + ((16376, 8),)
+
+#: An 8192-row (`max_num_tokens`) context step: a chunk continuing from mid-page to position
+#: 9215, and a request's first 1023 tokens.
+_PREFILL_8192 = ((2047, 7169), (0, 1023))
+
+
+def _spec(segments, pages, **overrides) -> dict:
+ """`segments` is `(first position, rows)` per request, in row order."""
+ return {**_CKPT, "segments": segments, "pages": pages, **overrides}
+
+
+def _engine(base: dict, member: str, segments, pages: int, **overrides) -> dict:
+ """An engine-pool cell: `rows` is the step's row count, `member` whose pool `pages` is."""
+ rows = sum(n for _, n in segments)
+ return {
+ **base,
+ "member": member,
+ "segments": segments,
+ "pages": pages,
+ "rows": rows,
+ **overrides,
+ }
+
+
+def _bits(tensor: torch.Tensor) -> torch.Tensor:
+ """The same storage as integers: fp8 has no comparison ops, and -0.0 != 0.0."""
+ as_int = {1: torch.uint8, 2: torch.int16, 4: torch.int32}[tensor.element_size()]
+ return tensor.view(as_int)
+
+
+class QsaPagedKvStore(OpWrapper):
+ """Scatter every row's K and V into its request's page of the HND pool.
+
+ Wraps `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels.
+ triton_qsa_paged_kv_store`: one Triton launch over `(rows, Hkv)`. For row t::
+
+ page = block_table[request_indices[t], logical_positions[t] // tokens_per_block]
+ slot = logical_positions[t] % tokens_per_block
+ k_cache[page, :, slot, :] = k[t] v_cache[page, :, slot, :] = v[t]
+
+ `logical_positions` is absolute, so prefill, decode and MTP rows are the
+ same call. The store is a move, converted once to the pool's dtype
+ (bitwise `Tensor.to`; an fp8-e4m3 pool saturates at 448 and takes no
+ scale, 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
+ index plane is never touched.
+ """
+
+ 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 off the fused projection row (row stride 14080), "
+ "at aligned, mid and page-crossing positions, plus one row on an "
+ "unallocated (-1) table entry that is skipped",
+ spec=_spec(_DECODE, pages=64, row_pad_heads=27, hole="unallocated_page"),
+ ),
+ Cell(
+ why="prefill 1024 rows off the fused projection row (row stride 14080)",
+ spec=_spec(((0, 256),) * 4, pages=160, row_pad_heads=27),
+ ),
+ # MTP7 verify blocks: 8 rows per request, so one block spans up to two pages.
+ # The non-default pool the forward reaches: fp8-e4m3.
+ Cell(
+ why="fp8 decode MTP7: 32 requests x 8 rows (256), blocks straddling pages",
+ spec=_spec(_MTP7_32, pages=1200, slots=40, cache_dtype=_F8),
+ ),
+ # The default target's members over the engine's own pool: one rank's page count, the
+ # 24-role-page stride, layer 47's view, a 128-row table.
+ Cell(
+ why="tep4 decode, 1 row: hkv 1 off the 4352-wide row, position 9215, "
+ "257109-page bf16 pool",
+ spec=_engine(_TP4, "tep4", ((9215, 1),), PAGES_TEP4),
+ ),
+ Cell(
+ why="tep4 decode, 128 rows (max_batch_size), aligned, mid and page-crossing",
+ spec=_engine(_TP4, "tep4", _DECODE_128, PAGES_TEP4),
+ ),
+ Cell(
+ why="tep4 MTP7 verify, 128 requests x 8 rows (1024) to position 16383, table width 516",
+ spec=_engine(_TP4, "tep4", _MTP7_128, PAGES_TEP4, width=516),
+ ),
+ Cell(
+ why="tp4 context, 8192 rows (max_num_tokens): a mid-page chunk to 9215 and a "
+ "request's first 1023",
+ spec=_engine(_TP4, "tp4", _PREFILL_8192, PAGES_TEP4),
+ ),
+ Cell(
+ why="dep4 kv_fp8 decode, 32 rows: hkv 2 off the 14080-wide row into a "
+ "171362-page e4m3 pool",
+ spec=_engine(_DEP4_FP8, "dep4", _DECODE_32, PAGES_DEP4_FP8),
+ ),
+ )
+
+ note = """
+ Source: the in-tree Triton kernel `_qsa_paged_kv_store_kernel`, JIT-compiled
+ on its first call; `ARCHS` is the architecture it compiled for and ran on.
+
+ Metadata consumed: none -- every input is an argument. The target passes
+ `qsa_req_idx_per_token` (int32, the row's batch index) and
+ `qsa_logical_positions` (int64) sliced to the step's rows,
+ `qsa_block_table` unsliced (raw pool page ids, `-1` unallocated), and the
+ layer pool's K and V role views of `get_buffers(i, kv_layout="HND")`.
+
+ Preconditions:
+ - destination slots are unique across the batch: two rows naming one
+ (page, slot) both store, unordered. The engine's mapping cannot
+ collide.
+ - `request_indices` and `logical_positions` are contiguous: the kernel
+ gets no stride for them. A strided view is read through its base
+ pointer and rows land at other rows' slots, silently (`is_valid`).
+ Every other operand is read through its strides; the cells drive k/v as
+ head slices of the packed QKV row and of the fused projection row, the
+ pool's K and V role views, and a contiguous table. The callable rejects
+ loudly: mismatched or non-3-D/4-D K/V shapes, head geometry,
+ `k_cache.shape[2] != tokens_per_block`, row dtypes outside bf16/fp16,
+ pool dtypes outside bf16/fp16/fp8-e4m3, a non-int32 or non-2-D table,
+ non-integer or wrong-length metadata, CPU operands, an empty pool or
+ table (not pinned by a test). `rows == 0` returns before launching.
+
+ Strides and the pool/table extents are Triton constexprs: each distinct
+ layout or pool size compiles its own kernel (once per engine, at its
+ pool's page count). CUDA-graph capturable (captured at 256 rows, 32 MTP7
+ verify blocks, and at 128 tep4 decode rows over the 257109-page pool); a
+ replay follows positions rewritten in place.
+
+ The engine's pool, as the stock manager lays it out at tp4 / tep4 / dep4:
+ one page holds the K and V planes of all 12 attention layers, so a
+ layer's K view steps 24 role pages a page -- 196608 elements at one KV
+ head, 393216 at two, bf16 and fp8 alike -- and its page count is the
+ engine's (`PAGES_TEP4`, `PAGES_DEP4_FP8`). The `pool="engine"` cells drive
+ layer 47's view with the pool's last page among the touched ones, past
+ 2^32 elements of the view's base (byte 101099339776 at tep4, 67382247424
+ at dep4 kv_fp8). The kernel loads the page id as int64 before its stride
+ product, and every such row lands bitwise.
+
+ Hard limits: both metadata values are cast to int32 inside the kernel
+ (`.to(tl.int32)`); an int64 position of `2^32 + 5` lands on slot 5 of page
+ column 0 -- written, not skipped (not pinned by a test). The engine's
+ positions stay below 2^14 at `max_seq_len` 16384 with MTP.
+
+ Not certified: duplicate destinations, multiple devices, a scaled fp8
+ pool (there is no scale argument), fp8 conversion of values past 448, the
+ packed-QKV-row (MTP draft layer) layout over a bf16 pool or at one KV head,
+ fp16 rows, head geometry other than D 256 with Hkv 1 or 2 and 32-token
+ pages, int32 positions or int64 request indices, head-major rows, sliced
+ or transposed tables, `rows == 0`, and of the rows that resolve to no slot
+ any but an unallocated (`-1`) table entry (the reference skips each).
+ Pools at a page count not driven here: tp4's own engine (tep4's count
+ stands in: a tp4 rank's pool has the same per-page shape), the MTP
+ engine's at tep4 (its draft layer has a separate cache, so its trunk
+ count differs; the non-MTP count stands in), and dep4's bf16 pool (119680
+ pages). Positions past 16383 and tables wider than 516 columns.
+ """
+
+ def __call__(
+ self,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ k_cache: torch.Tensor,
+ v_cache: torch.Tensor,
+ request_indices: torch.Tensor,
+ logical_positions: torch.Tensor,
+ block_table: torch.Tensor,
+ tokens_per_block: int,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ _qsa_kernels.triton_qsa_paged_kv_store(
+ k=k,
+ v=v,
+ k_cache=k_cache,
+ v_cache=v_cache,
+ request_indices=request_indices,
+ logical_positions=logical_positions,
+ block_table=block_table,
+ tokens_per_block=tokens_per_block,
+ )
+ return k_cache, v_cache
+
+ def reference(
+ self,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ k_cache: torch.Tensor,
+ v_cache: torch.Tensor,
+ request_indices: torch.Tensor,
+ logical_positions: torch.Tensor,
+ block_table: torch.Tensor,
+ tokens_per_block: int,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Both caches after the scatter, as fresh tensors; the inputs are not touched."""
+ expected_k, expected_v = k_cache.clone(), v_cache.clone()
+ requests = request_indices.long()
+ positions = logical_positions.long()
+ columns = positions.div(tokens_per_block, rounding_mode="floor")
+ table_rows, table_width = block_table.shape
+ addressable = (
+ (requests >= 0) & (requests < table_rows) & (positions >= 0) & (columns < table_width)
+ )
+ pages = block_table[
+ requests.clamp(0, table_rows - 1), columns.clamp(0, table_width - 1)
+ ].long()
+ addressable &= (pages >= 0) & (pages < k_cache.shape[0])
+ page, slot = pages[addressable], (positions % tokens_per_block)[addressable]
+ expected_k[page, :, slot, :] = k[addressable].to(expected_k.dtype)
+ expected_v[page, :, slot, :] = v[addressable].to(expected_v.dtype)
+ return expected_k, expected_v
+
+ def is_valid(
+ self,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ k_cache: torch.Tensor,
+ v_cache: torch.Tensor,
+ request_indices: torch.Tensor,
+ logical_positions: torch.Tensor,
+ block_table: torch.Tensor,
+ tokens_per_block: int,
+ ) -> None:
+ # The kernel reads these two as flat unit-stride buffers; a strided view
+ # lands rows at other rows' slots, with nothing raised.
+ assert request_indices.is_contiguous(), "request_indices must be a contiguous 1-D tensor"
+ assert logical_positions.is_contiguous(), (
+ "logical_positions must be a contiguous 1-D tensor"
+ )
+
+
+qsa_paged_kv_store = QsaPagedKvStore()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_kv_store_nvfp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_kv_store_nvfp4.py
new file mode 100644
index 000000000000..719c548b5fd2
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_kv_store_nvfp4.py
@@ -0,0 +1,572 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""QSA sparse path's paged K/V append into an NVFP4 pool: quantize and scatter in one launch."""
+
+from __future__ import annotations
+
+from typing import Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+#: Values per e4m3 block scale, and e2m1 values per stored byte.
+_SF_VEC = 16
+_PER_BYTE = 2
+#: Tokens of a V scale group: V's scales are stored `[tpb / 4, D / 16, 4]` per head page.
+_V_SF_QUAD = 4
+#: e2m1's largest magnitude: a block's scale maps its amax onto it.
+_E2M1_MAX = 6.0
+#: What `cvt.rn.satfinite.e2m1x2.f32` makes of a NaN of either sign: +6. A NaN reaches
+#: it as 0 x inf, from a zero in a block whose nonzero amax rounds its scale to 0.
+_E2M1_NAN_CODE = 0x7
+#: The round-to-nearest-even decision points between neighbouring e2m1 magnitudes:
+#: `(threshold, inclusive)` steps the code up when `|v| > t`, or `|v| >= t` where the
+#: upper neighbour is the even code.
+_E2M1_STEPS = (
+ (0.25, False),
+ (0.75, True),
+ (1.25, False),
+ (1.75, True),
+ (2.5, False),
+ (3.5, True),
+ (5.0, False),
+)
+
+#: This checkpoint's K/V geometry and the engine's QSA table. `heads_q` places k and v as
+#: head slices of the packed `[T, 24 + 2 + 2, 256]` QKV row; `global_k` / `global_v` are
+#: `kv_scale_orig_quant[1]` / `[2]`; `slot_pages` role pages of one lifecycle slot hold
+#: the call's planes, layer `layer`'s K and V among them (the trunk pool holds 24).
+_CKPT = dict(
+ hkv=2,
+ head_dim=256,
+ heads_q=24,
+ tokens_per_block=32,
+ slots=32,
+ width=260,
+ row_dtype=torch.bfloat16,
+ request_dtype=torch.int32,
+ position_dtype=torch.int64,
+ global_k=1.0,
+ global_v=1.0,
+ values="randn",
+ hole=None,
+ thop_append=False,
+ slot_pages=2,
+ layer=0,
+)
+
+_DECODE = ((0, 1), (31, 1), (32, 1), (33, 1), (100, 1), (255, 1), (7, 1), (64, 1))
+
+#: One dep4 rank's kv_nvfp4 pool as the stock engine sizes it (NVFP4 K/V, chunked prefill):
+#: 213295 slots, its `maxNumBlocks`; the trunk layout of 24 role pages a slot, layer 11's
+#: planes (the last attention layer's); and the `qsa_block_table` of
+#: `max_num_requests` 128 rows by `max_blocks_per_seq` 292 columns (max_seq_len 9216).
+#: `last_page` puts that request's last page on the pool's last slot, so the stores land past
+#: 2^31 bytes of both the data and the scale plane. The pool alone is 47 GB, so `large`.
+_DEP4_PAGES = 213295
+_DEP4 = dict(slots=128, width=292, slot_pages=24, layer=11, last_page=0, large=True)
+#: 32 decode rows of 32 requests: max_seq_len - 1 first, page edges, then a stride over the table.
+_DEP4_DECODE = ((9215, 1), (0, 1), (31, 1), (32, 1), (33, 1)) + tuple(
+ (293 * i + 17, 1) for i in range(1, 28)
+)
+
+
+def _spec(segments, pages, **overrides) -> dict:
+ """`segments` is `(first position, rows)` per request, in row order."""
+ return {**_CKPT, "segments": segments, "pages": pages, **overrides}
+
+
+@triton.jit
+def _rcp_approx(x):
+ """`rcp.approx.ftz.f32`, the reciprocal the engine's FP4 quantizer uses."""
+ return tl.inline_asm_elementwise(
+ "rcp.approx.ftz.f32 $0, $1;", "=f,f", [x], dtype=tl.float32, is_pure=True, pack=1
+ )
+
+
+@triton.jit
+def _e2m1x2(lo, hi):
+ """One byte of two e2m1 codes, `lo` in the low nibble (`cvt.rn.satfinite`)."""
+ return tl.inline_asm_elementwise(
+ "{ .reg .b8 t; .reg .b16 s; cvt.rn.satfinite.e2m1x2.f32 t, $2, $1; "
+ "mov.b16 s, {t, t}; cvt.u32.u16 $0, s; }",
+ "=r,f,f",
+ [lo, hi],
+ dtype=tl.int32,
+ is_pure=True,
+ pack=1,
+ )
+
+
+@triton.jit
+def _quantize_row(
+ src,
+ stride_dim,
+ global_scale,
+ HEAD_DIM: tl.constexpr,
+):
+ """One head row: (packed e2m1 bytes [D/2], e4m3 scale bytes [D/16])."""
+ pairs = tl.arange(0, HEAD_DIM // 2)
+ lo = tl.load(src + (2 * pairs) * stride_dim).to(tl.float32)
+ hi = tl.load(src + (2 * pairs + 1) * stride_dim).to(tl.float32)
+ pair_max = tl.maximum(tl.abs(lo), tl.abs(hi))
+ block_max = tl.max(tl.reshape(pair_max, [HEAD_DIM // 16, 8]), axis=1)
+ sf_value = global_scale * (block_max * _rcp_approx(tl.full([HEAD_DIM // 16], 6.0, tl.float32)))
+ sf_code = sf_value.to(tl.float8e4nv)
+ sf_rounded = sf_code.to(tl.float32)
+ out_scale = tl.where(
+ block_max != 0.0,
+ _rcp_approx(sf_rounded * _rcp_approx(tl.zeros_like(sf_rounded) + global_scale)),
+ 0.0,
+ )
+ lo_scaled = tl.reshape(
+ tl.reshape(lo, [HEAD_DIM // 16, 8]) * out_scale[:, None], [HEAD_DIM // 2]
+ )
+ hi_scaled = tl.reshape(
+ tl.reshape(hi, [HEAD_DIM // 16, 8]) * out_scale[:, None], [HEAD_DIM // 2]
+ )
+ packed = (_e2m1x2(lo_scaled, hi_scaled) & 0xFF).to(tl.uint8)
+ return packed, sf_code.to(tl.uint8, bitcast=True)
+
+
+@triton.jit
+def _qsa_paged_kv_store_nvfp4_kernel(
+ k,
+ v,
+ k_cache,
+ v_cache,
+ k_scale_cache,
+ v_scale_cache,
+ request_indices,
+ logical_positions,
+ block_table,
+ kv_scale_orig_quant,
+ num_pages,
+ num_requests,
+ table_width,
+ k_stride_row,
+ k_stride_head,
+ k_stride_dim,
+ v_stride_row,
+ v_stride_head,
+ v_stride_dim,
+ kc_stride_page,
+ kc_stride_head,
+ kc_stride_token,
+ vc_stride_page,
+ vc_stride_head,
+ vc_stride_token,
+ ks_stride_page,
+ ks_stride_head,
+ vs_stride_page,
+ vs_stride_head,
+ bt_stride_request,
+ bt_stride_page,
+ HEAD_DIM: tl.constexpr,
+ TOKENS_PER_BLOCK: tl.constexpr,
+):
+ row = tl.program_id(0).to(tl.int64)
+ head = tl.program_id(1).to(tl.int64)
+ request = tl.load(request_indices + row).to(tl.int64)
+ logical = tl.load(logical_positions + row).to(tl.int64)
+ column = logical // TOKENS_PER_BLOCK
+ token = logical % TOKENS_PER_BLOCK
+ lookup = (request >= 0) & (request < num_requests) & (logical >= 0) & (column < table_width)
+ page = tl.load(
+ block_table + request * bt_stride_request + column * bt_stride_page, mask=lookup, other=-1
+ ).to(tl.int64)
+ valid = (page >= 0) & (page < num_pages)
+ page = tl.maximum(page, 0)
+
+ k_data, k_sf = _quantize_row(
+ k + row * k_stride_row + head * k_stride_head,
+ k_stride_dim,
+ tl.load(kv_scale_orig_quant + 1),
+ HEAD_DIM,
+ )
+ v_data, v_sf = _quantize_row(
+ v + row * v_stride_row + head * v_stride_head,
+ v_stride_dim,
+ tl.load(kv_scale_orig_quant + 2),
+ HEAD_DIM,
+ )
+ pairs = tl.arange(0, HEAD_DIM // 2)
+ blocks = tl.arange(0, HEAD_DIM // 16)
+ tl.store(
+ k_cache + page * kc_stride_page + head * kc_stride_head + token * kc_stride_token + pairs,
+ k_data,
+ mask=valid,
+ )
+ tl.store(
+ v_cache + page * vc_stride_page + head * vc_stride_head + token * vc_stride_token + pairs,
+ v_data,
+ mask=valid,
+ )
+ # K scales token-major within the head page; V scales in four-token groups.
+ tl.store(
+ k_scale_cache
+ + page * ks_stride_page
+ + head * ks_stride_head
+ + token * (HEAD_DIM // 16)
+ + blocks,
+ k_sf,
+ mask=valid,
+ )
+ tl.store(
+ v_scale_cache
+ + page * vs_stride_page
+ + head * vs_stride_head
+ + (token // 4) * (4 * (HEAD_DIM // 16))
+ + blocks * 4
+ + token % 4,
+ v_sf,
+ mask=valid,
+ )
+
+
+def _store_nvfp4(
+ *,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ k_cache: torch.Tensor,
+ v_cache: torch.Tensor,
+ k_scale_cache: torch.Tensor,
+ v_scale_cache: torch.Tensor,
+ request_indices: torch.Tensor,
+ logical_positions: torch.Tensor,
+ block_table: torch.Tensor,
+ tokens_per_block: int,
+ kv_scale_orig_quant: torch.Tensor,
+) -> None:
+ """Quantize rows of K and V to NVFP4 and store them into their pages (see the entry)."""
+ if k.ndim != 3 or v.shape != k.shape:
+ raise ValueError("NVFP4 K/V store: k and v must be matching [rows, heads, dim]")
+ rows, heads, head_dim = k.shape
+ if k.dtype not in (torch.bfloat16, torch.float16) or v.dtype != k.dtype:
+ raise ValueError("NVFP4 K/V store: k and v must both be BF16 or both FP16")
+ if head_dim < 32 or head_dim & (head_dim - 1):
+ raise ValueError("NVFP4 K/V store: head_dim must be a power of two of at least 32")
+ if tokens_per_block <= 0 or tokens_per_block % _V_SF_QUAD:
+ raise ValueError("NVFP4 K/V store: tokens_per_block must be a positive multiple of 4")
+ data_shape = (k_cache.shape[0], heads, tokens_per_block, head_dim // _PER_BYTE)
+ scale_shape = (k_cache.shape[0], heads, tokens_per_block, head_dim // _SF_VEC)
+ for name, cache, shape in (
+ ("k_cache", k_cache, data_shape),
+ ("v_cache", v_cache, data_shape),
+ ("k_scale_cache", k_scale_cache, scale_shape),
+ ("v_scale_cache", v_scale_cache, scale_shape),
+ ):
+ if cache.dtype != torch.uint8 or tuple(cache.shape) != shape:
+ raise ValueError(
+ f"NVFP4 K/V store: {name} must be uint8 {shape}, got {cache.dtype} "
+ f"{tuple(cache.shape)}"
+ )
+ if cache.stride(3) != 1:
+ raise ValueError(f"NVFP4 K/V store: {name} must be contiguous along its last dim")
+ for name, cache in (("k_scale_cache", k_scale_cache), ("v_scale_cache", v_scale_cache)):
+ if cache.stride(2) != head_dim // _SF_VEC:
+ raise ValueError(f"NVFP4 K/V store: each head page of {name} must be contiguous")
+ if block_table.ndim != 2 or block_table.dtype != torch.int32:
+ raise ValueError("NVFP4 K/V store: block_table must be a 2-D int32 tensor")
+ for name, meta in (
+ ("request_indices", request_indices),
+ ("logical_positions", logical_positions),
+ ):
+ if meta.shape != (rows,) or meta.dtype not in (torch.int32, torch.int64):
+ raise ValueError(f"NVFP4 K/V store: {name} must be one int32/int64 entry per row")
+ if not meta.is_contiguous():
+ raise ValueError(f"NVFP4 K/V store: {name} must be contiguous")
+ if (
+ kv_scale_orig_quant.dtype != torch.float32
+ or kv_scale_orig_quant.shape != (3,)
+ or not kv_scale_orig_quant.is_contiguous()
+ ):
+ raise ValueError("NVFP4 K/V store: kv_scale_orig_quant must be a contiguous fp32 [3]")
+ operands = (k, v, k_cache, v_cache, k_scale_cache, v_scale_cache, request_indices)
+ operands += (logical_positions, block_table, kv_scale_orig_quant)
+ if not all(t.is_cuda and t.device == k.device for t in operands):
+ raise ValueError("NVFP4 K/V store: every operand must be on k's CUDA device")
+ if rows == 0:
+ return
+ _qsa_paged_kv_store_nvfp4_kernel[(rows, heads)](
+ k,
+ v,
+ k_cache,
+ v_cache,
+ k_scale_cache,
+ v_scale_cache,
+ request_indices,
+ logical_positions,
+ block_table,
+ kv_scale_orig_quant,
+ k_cache.shape[0],
+ block_table.shape[0],
+ block_table.shape[1],
+ k.stride(0),
+ k.stride(1),
+ k.stride(2),
+ v.stride(0),
+ v.stride(1),
+ v.stride(2),
+ k_cache.stride(0),
+ k_cache.stride(1),
+ k_cache.stride(2),
+ v_cache.stride(0),
+ v_cache.stride(1),
+ v_cache.stride(2),
+ k_scale_cache.stride(0),
+ k_scale_cache.stride(1),
+ v_scale_cache.stride(0),
+ v_scale_cache.stride(1),
+ block_table.stride(0),
+ block_table.stride(1),
+ HEAD_DIM=head_dim,
+ TOKENS_PER_BLOCK=tokens_per_block,
+ num_warps=4,
+ )
+
+
+def _e2m1_codes(values: torch.Tensor) -> torch.Tensor:
+ """fp32 -> e2m1 codes 0..15, round to nearest even, saturating at 6; a NaN
+ (0 x inf, in a block whose scale underflows to zero) converts to +6."""
+ magnitude = values.abs()
+ code = torch.zeros(values.shape, dtype=torch.int32, device=values.device)
+ for threshold, inclusive in _E2M1_STEPS:
+ code += (magnitude >= threshold if inclusive else magnitude > threshold).to(torch.int32)
+ code |= torch.signbit(values).to(torch.int32) << 3
+ return torch.where(torch.isnan(values), torch.full_like(code, _E2M1_NAN_CODE), code)
+
+
+def _quantize_reference(rows: torch.Tensor, global_scale: torch.Tensor):
+ """[n, D] rows -> (packed bytes [n, D/2], e4m3 scale bytes [n, D/16]), in plain torch."""
+ n, head_dim = rows.shape
+ x = rows.float().reshape(n, head_dim // _SF_VEC, _SF_VEC)
+ block_max = x.abs().amax(dim=-1)
+ sixth = torch.tensor(_E2M1_MAX, dtype=torch.float32, device=rows.device).reciprocal()
+ scale_code = (global_scale * (block_max * sixth)).to(torch.float8_e4m3fn)
+ rounded = scale_code.float()
+ out_scale = torch.where(
+ block_max != 0,
+ (rounded * global_scale.reciprocal()).reciprocal(),
+ torch.zeros_like(rounded),
+ )
+ codes = _e2m1_codes(x * out_scale[..., None]).reshape(n, head_dim // _PER_BYTE, _PER_BYTE)
+ packed = (codes[..., 0] | (codes[..., 1] << 4)).to(torch.uint8)
+ return packed, scale_code.view(torch.uint8)
+
+
+class QsaPagedKvStoreNvfp4(OpWrapper):
+ """Quantize every row's K and V to NVFP4 and store them into the request's page.
+
+ A Triton kernel defined in this file, one launch over `(rows, Hkv)`. For row t
+ and KV head h, with `page = block_table[request_indices[t], logical_positions[t]
+ // tpb]` and `s = logical_positions[t] % tpb`, each 16-value block `b` of
+ `k[t, h]` (and alike for v, with `kv_scale_orig_quant[2]`) is stored as::
+
+ g = kv_scale_orig_quant[1]
+ sf = e4m3(g * (amax_b * rcp(6))) # one byte
+ code_i = e2m1(k[t, h, i] * rcp(float(sf) * rcp(g))) # 0 when amax_b == 0
+ k_cache[page, h, s, i // 2] nibble i % 2 (low nibble = even i)
+ k_scale_cache[page, h, s, b] = sf # token-major
+ v_scale_cache[page, h, (s // 4) * 4 * (D / 16) + 4 * b + s % 4] # four-token groups
+
+ `rcp` is `rcp.approx.ftz.f32`, `e4m3` and `e2m1` round to nearest even and
+ saturate. This is the attention op's own NVFP4 append (`cvt_warp_fp16_to_fp4`
+ behind `quantizeAndWriteFP4KVCache`), the same bits for the same row: a
+ token's stored value does not depend on which path wrote it. A row that
+ resolves to no slot -- negative request, request past the table, negative
+ position, page column past the table, a table entry that is negative or past
+ the pool -- is skipped. Returns the four caches, written in place.
+
+ Fusion boundary: quantization and the scatter. The projection, norm and RoPE
+ of K are the caller's; the QSA index plane is not touched.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="decode: 8 rows at aligned, mid-page and page-crossing positions, with zeros, "
+ "-0.0 and a 900.0 outlier inside blocks, and a row on an unallocated (-1) page: "
+ "zero scales, saturation and the skip",
+ spec=_spec(_DECODE, pages=69, slots=8, values="edge", hole="unallocated_page"),
+ ),
+ Cell(
+ why="prefill: 1024 rows, 4 requests x 256 walking 8 pages each, in the engine's "
+ "trunk layout: layer 11's planes among the 24 role pages of a slot",
+ spec=_spec(((0, 256),) * 4, pages=160, slot_pages=24, layer=11),
+ ),
+ Cell(
+ why="the attention op's own append over the same rows (context, chunk, decode "
+ "P=1 and P=3) at unit second-level scales: data and both scale layouts bitwise equal",
+ spec=_spec(((0, 115), (0, 101), (0, 36)), pages=96, thop_append=True),
+ ),
+ Cell(
+ why="dep4 kv_nvfp4 decode: one row at position 9215 (max_seq_len - 1) on the last "
+ "of the rank's 213295 pool slots, 128 x 292 table",
+ spec=_spec(((9215, 1),), pages=_DEP4_PAGES, **_DEP4),
+ ),
+ Cell(
+ why="dep4 kv_nvfp4 decode: 32 requests of the 128 table rows at positions 0..9215 "
+ "across page edges, on the rank's 213295-slot pool",
+ spec=_spec(_DEP4_DECODE, pages=_DEP4_PAGES, **_DEP4),
+ ),
+ Cell(
+ why="dep4 kv_nvfp4 context chunk: 8192 rows (max_num_tokens), one prompt's chunk at "
+ "2048..9215 and a new prompt's first 1024, on the rank's 213295-slot pool",
+ spec=_spec(((2048, 7168), (0, 1024)), pages=_DEP4_PAGES, **_DEP4),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none -- every input is an argument. The target passes the
+ step's `qsa_req_idx_per_token` (int32) and `qsa_logical_positions` (int64),
+ `qsa_block_table` unsliced (raw pool page ids, `-1` unallocated), the layer's
+ K and V planes of the QSA manager's `get_buffers(i, "HND")` (uint8 `[slots,
+ Hkv, tpb, D/2]`), the two planes of `get_block_scale_buffers(i, "HND")`
+ (uint8 `[slots, Hkv, tpb, D/16]`), and its own `kv_scale_orig_quant` -- the
+ same fp32 `[3]` tensor its dense attention call binds, so both writers apply
+ one second-level scale.
+
+ Preconditions:
+ - destination slots are unique across the batch: two rows naming one
+ (page, slot) both store, unordered.
+ - the dense path's reader assumes this layout: e2m1 nibbles HND per page,
+ K scales linear, V scales in four-token groups. The attention op's
+ append, run over the same rows, leaves bitwise the same pool (a cell),
+ at the unit second-level scales the target binds.
+ The callable rejects loudly (ValueError; the test pins the non-contiguous
+ metadata refusal): non-matching or non-3-D k/v, a row
+ dtype outside bf16/fp16, a head_dim that is not a power of two >= 32, a
+ tokens_per_block that is not a multiple of 4, caches that are not uint8 of
+ the stated shapes or whose last dim or scale head page is strided, a non-2-D
+ or non-int32 table, metadata that is not one contiguous int32/int64 entry per
+ row, a `kv_scale_orig_quant` that is not a contiguous fp32 `[3]`, operands on
+ another device. `rows == 0` returns before launching (not certified).
+
+ Every offset is formed in 64 bits: rows, heads, requests, positions and pages
+ are widened before any product. CUDA-graph capturable; a replay follows rows,
+ pages and positions rewritten in place (the test replays 8 MTP verify blocks of
+ 8 rows bitwise equal to eager).
+
+ Compiled variants: the page count (`num_pages`, the pool's slot count), the
+ table's height (`num_requests`) and width and every stride are plain-int
+ arguments, and Triton compiles the kernel once per class of each -- equal to
+ 1, a multiple of 16, anything else -- besides once per head_dim and
+ tokens_per_block. A served call puts the width and each stride in one fixed
+ class (the k/v head slices of the packed QKV row, the planes' slot, head and
+ token strides, the contiguous table's) and varies the page count and the
+ table height. The cells run (other, other), (multiple of 16, multiple of 16)
+ and, at dep4, (other, multiple of 16); (multiple of 16, other) and a table
+ height of 1 are not certified.
+
+ The default target at dep4 with NVFP4 K/V: a rank holds both KV heads in
+ the trunk layout, data `[slots, 24, 2, 32, 128]` and scales
+ `[slots, 24, 2, 32, 16]`, attention layer j's K and V at role pages 2j and
+ 2j + 1 (the stock manager's `get_buffers` / `get_block_scale_buffers`
+ views). The stock engine at that configuration (chunked prefill on,
+ max_num_tokens 8192) has `maxNumBlocks` 213295 a rank, and the dep4 cells
+ drive that many slots under its 128 x 292 `qsa_block_table`, one
+ request's last page on the pool's last slot: stores land 4.19e10 bytes
+ into the data plane and 5.24e9 into the scale plane. A cell's pool is
+ 47 GB, so the dep4 cells are `large`.
+
+ Not finite values: a NaN or +-inf in a block is not certified (the e2m1 and
+ e4m3 conversions saturate, the block maximum is the kernel's `max`).
+
+ Not certified either: second-level scales other than 1.0, int32 positions
+ or int64 request indices, and of the rows that resolve to no slot any but
+ an unallocated (`-1`) table entry (the reference skips each).
+ """
+
+ def __call__(
+ self,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ k_cache: torch.Tensor,
+ v_cache: torch.Tensor,
+ k_scale_cache: torch.Tensor,
+ v_scale_cache: torch.Tensor,
+ request_indices: torch.Tensor,
+ logical_positions: torch.Tensor,
+ block_table: torch.Tensor,
+ tokens_per_block: int,
+ kv_scale_orig_quant: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ _store_nvfp4(
+ k=k,
+ v=v,
+ k_cache=k_cache,
+ v_cache=v_cache,
+ k_scale_cache=k_scale_cache,
+ v_scale_cache=v_scale_cache,
+ request_indices=request_indices,
+ logical_positions=logical_positions,
+ block_table=block_table,
+ tokens_per_block=tokens_per_block,
+ kv_scale_orig_quant=kv_scale_orig_quant,
+ )
+ return k_cache, v_cache, k_scale_cache, v_scale_cache
+
+ def reference(
+ self,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ k_cache: torch.Tensor,
+ v_cache: torch.Tensor,
+ k_scale_cache: torch.Tensor,
+ v_scale_cache: torch.Tensor,
+ request_indices: torch.Tensor,
+ logical_positions: torch.Tensor,
+ block_table: torch.Tensor,
+ tokens_per_block: int,
+ kv_scale_orig_quant: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ """The four caches after the store, as fresh tensors; the inputs are untouched.
+
+ Quantization in fp32 torch with IEEE reciprocals where the kernel uses
+ `rcp.approx.ftz`; the two agree on every operand the cells drive (6.0, the
+ unit second-level scale, and the e4m3 scale values their rows produce), which
+ the cells' bitwise gate is what pins.
+ """
+ out = [
+ t.clone(memory_format=torch.contiguous_format)
+ for t in (k_cache, v_cache, k_scale_cache, v_scale_cache)
+ ]
+ rows, heads, head_dim = k.shape
+ requests = request_indices.long()
+ positions = logical_positions.long()
+ columns = positions.div(tokens_per_block, rounding_mode="floor")
+ table_rows, table_width = block_table.shape
+ ok = (requests >= 0) & (requests < table_rows) & (positions >= 0) & (columns < table_width)
+ pages = block_table[
+ requests.clamp(0, table_rows - 1), columns.clamp(0, table_width - 1)
+ ].long()
+ ok &= (pages >= 0) & (pages < k_cache.shape[0])
+ page, slot = pages[ok], (positions % tokens_per_block)[ok]
+ n = int(page.numel())
+ if n == 0:
+ return tuple(out)
+ groups = head_dim // _SF_VEC
+ quad_base = (slot // _V_SF_QUAD) * (_V_SF_QUAD * groups) + slot % _V_SF_QUAD
+ for src, data, scales, global_scale, linear in (
+ (k, out[0], out[2], kv_scale_orig_quant[1], True),
+ (v, out[1], out[3], kv_scale_orig_quant[2], False),
+ ):
+ packed, sf = _quantize_reference(src[ok].reshape(n * heads, head_dim), global_scale)
+ data[page, :, slot, :] = packed.reshape(n, heads, head_dim // _PER_BYTE)
+ sf = sf.reshape(n, heads, groups)
+ if linear:
+ scales[page, :, slot, :] = sf
+ else:
+ flat = scales.view(scales.shape[0], heads, tokens_per_block * groups)
+ offsets = quad_base[:, None] + _V_SF_QUAD * torch.arange(groups, device=sf.device)
+ for h in range(heads):
+ flat[page[:, None], h, offsets] = sf[:, h]
+ return tuple(out)
+
+
+qsa_paged_kv_store_nvfp4 = QsaPagedKvStoreNvfp4()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa.py
new file mode 100644
index 000000000000..b65b55eac0cf
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa.py
@@ -0,0 +1,484 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""QSA sparse GQA: attention over an explicitly selected set of paged K/V tokens."""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+# Approved exception: QSA's sparse path never reaches the thop attention op's
+# sparse argument group, so there is no torch.ops surface for this computation.
+from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import (
+ can_fuse_qsa_splitk_output_gate as _can_fuse_qsa_splitk_output_gate,
+)
+from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import (
+ triton_qsa_paged_sparse_gqa as _triton_qsa_paged_sparse_gqa,
+)
+
+_CHECKPOINT = dict(
+ q_heads=24,
+ kv_heads=2,
+ head_dim=256,
+ tokens_per_block=32,
+ compress_ratio=4,
+ token_topk=2048,
+ cache_dtype=torch.bfloat16,
+ q_dtype=torch.bfloat16,
+ # q is the first 24 heads of the packed [T, 24 + 2 + 2, 256] QKV buffer;
+ # "fused_row" is the target's fused input projection row, 55 heads of 256.
+ q_layout="qkv_slice",
+ request_dtype=torch.int32,
+ # The engine's table: 260 pages per slot, -1 past each sequence.
+ pages=260,
+)
+
+_SPLIT_K = ("splitk", "merge")
+_PREFILL = ("prefill",)
+_FP8 = dict(cache_dtype=torch.float8_e4m3fn)
+
+
+def _spec(**overrides) -> dict:
+ """A cell spec: this checkpoint's geometry unless the cell says otherwise.
+
+ `segments` is `(first position, rows)` per request, in row order;
+ `prefill_pair` says whether `query_positions` / `compress_ratio` are
+ passed, which the target does exactly when the step has a context request.
+ """
+ return {**_CHECKPOINT, **overrides}
+
+
+def _decode(segments, launches=_SPLIT_K, **overrides) -> dict:
+ base = dict(segments=segments, slots=32, prefill_pair=False, launches=launches)
+ return _spec(**{**base, **overrides})
+
+
+#: The engine's trunk layout (`Qwen4ExpHybridCacheManager.get_buffers`): a lifecycle slot
+#: holds the 24 role pages of the 12 attention layers, K then V each, and a layer's
+#: `k_cache` / `v_cache` are its two planes, so one pool page advances 24 role pages.
+#: Layer 11's planes are the slot's last.
+_TRUNK = dict(slot_pages=24, layer=11)
+
+#: The element offset the default target's cells straddle: the requests' pages lie on both
+#: sides of it in the layer's K and V planes. An engine's pool reaches far past it (tep4's
+#: bf16 engine: 245835 pages, 4.8e10 elements of trunk).
+_OFFSET_EDGE = 2**31
+
+#: One tp4 / tep4 rank (attention TP 4): 6 q heads and one of the 2 KV heads, q the first 6
+#: heads of the packed [T, 6 + 1 + 1, 256] QKV buffer; a `max_batch_size` 128-row table,
+#: `ceil((max_seq_len + 1) / 32)` padded to a multiple of 4 columns: 292 at max_seq_len 9216,
+#: 516 at MTP's 16384.
+_TP4 = dict(q_heads=6, kv_heads=1, slots=128, pages=292, **_TRUNK)
+_TP4_MTP = dict(_TP4, pages=516)
+#: One dep4 rank (attention DP): tp1's 24 / 2 heads and a 32-row table (128 requests over 4
+#: ranks); a bf16 pool, or the kv_fp8 variant's e4m3 pool.
+_DEP4 = dict(q_heads=24, kv_heads=2, slots=32, pages=292, **_TRUNK)
+_DEP4_FP8 = dict(_DEP4, **_FP8)
+
+#: q and the output gate off one rank's fused input projection row, `[q | k | v | gate |
+#: index q, k | pad]` in heads of 256 (tp1's `weights._attn_in_w`): 6 + 1 + 1 + 6 + 3 = 17
+#: heads (row stride 4352) at a tp4 / tep4 rank, 24 + 2 + 2 + 24 + 3 = 55 (14080) at dep4.
+_TP4_ROW = dict(q_layout="fused_row", gate="fused_row", row_heads=17)
+_DEP4_ROW = dict(q_layout="fused_row", gate="fused_row", row_heads=55)
+
+
+def _straddling(segments, kv_heads: int) -> int:
+ """The pool page count that puts a cell's requests -- the pool's last pages, as the
+ test lays them out -- on both sides of `_OFFSET_EDGE` elements of the K plane (half
+ of them below it, or the whole pool when they need more than twice the pages below)."""
+ tpb, d, slot_pages, layer = 32, 256, _TRUNK["slot_pages"], _TRUNK["layer"]
+ needed = sum(math.ceil((s + n) / tpb) for s, n in segments) + 7
+ role_page = kv_heads * tpb * d
+ edge_page = math.ceil((_OFFSET_EDGE - 2 * layer * role_page) / (slot_pages * role_page))
+ return max(edge_page + needed // 2, needed)
+
+
+def _default(segments, member: dict, launches, **overrides) -> dict:
+ """A cell of the default target's call: `member`'s heads, table and trunk layout, its pool
+ straddling `_OFFSET_EDGE`."""
+ return _decode(
+ segments,
+ launches=launches,
+ rows=sum(n for _, n in segments),
+ pool_pages=_straddling(segments, member["kv_heads"]),
+ offset_edge=_OFFSET_EDGE,
+ **{**member, **overrides},
+ )
+
+
+def _spread(count: int, rows: int, last_start: int) -> tuple:
+ """`count` requests of `rows` rows: one starting at `last_start` (its last row at the
+ envelope's last position), the others spread from position 0 up."""
+ span = last_start + 1
+ return tuple(((5 + 1129 * g) % span, rows) for g in range(count - 1)) + ((last_start, rows),)
+
+
+class QsaPagedSparseGqa(OpWrapper):
+ """Sparse GQA over the selected tokens of an HND paged K/V pool.
+
+ Wraps `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels.
+ triton_qsa_paged_sparse_gqa`. For row `r`, query head `h`, KV head
+ `h // (Hq // Hkv)` (contiguous groups, not interleaved), over the columns
+ `j` of `selected_tokens[r]` that name a resident token::
+
+ t = selected_tokens[r, j] # logical position
+ page = block_table[request_indices[r], t // tokens_per_block]
+ s_j = (q[r, h] . k_cache[page, kv, t % tpb]) * softmax_scale
+ out[r, h] = sum_j softmax(s)_j * v_cache[page, kv, t % tpb]
+
+ A negative index, a page past the table's width, a negative or
+ out-of-pool table entry, and a request index outside the table's rows all
+ contribute nothing; a row left with nothing returns zeros. Indices are not
+ deduplicated: a token named twice is attended twice. Returns a fresh,
+ contiguous `[rows, Hq, D]` tensor of `q.dtype`, even for a strided `q`.
+
+ Two kernels, chosen inside the call: with `query_positions` and
+ `compress_ratio` given and `rows * Hkv > 512`, a causal-prefix prefill
+ kernel; otherwise split-K, plus a merge kernel when it used more than one
+ split. Each cell names the launches it was measured on.
+
+ **The prefill pair is a bound, not a mask.** The prefill kernel reads only
+ `min(visible // cr, block_topk) * cr + visible % cr` leading columns of a
+ row (`visible = query_positions[r] + 1`, rounded up to a 64-wide tile) and
+ nothing after them. That is exactly what QSA's own top-k + expansion emit,
+ packed at the front; a selection with valid entries past the bound is
+ silently truncated. The split-K kernel ignores the pair.
+
+ **The output gate rides the merge when `output_gate` is passed.** Each
+ merged head row is rounded to `q.dtype`, multiplied in fp32 by
+ `sigmoid(output_gate)` and rounded again: the arithmetic of the ungated
+ call followed by `fused_sigmoid_mul_inplace`. `output_gate` is `[rows,
+ Hq * D]` or `[rows, Hq, D]` in `q.dtype`, unit stride in its last dim, any
+ row stride (the gate columns of the fused projection row). Only a split-K
+ launch with more than one split has a merge to carry it; `fuses_output_gate`
+ says whether a call does, and the op raises `ValueError` if handed a gate
+ it cannot carry.
+
+ Fusion boundary: the block-table gather, scale, online softmax, P.V, the
+ split-K merge and, when passed, the output gate are inside. Deciding what
+ to attend to, appending K/V to the pool, and q's projection, norm and RoPE
+ are not.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ # ---- generation: no prefill pair, split-K ---------------------------
+ Cell(
+ why="decode MTP7: 32 slots x 8 rows = 256 rows, 512 programs, 4-way split-K, "
+ "q and the gate off the fused row, the gate fused into the merge",
+ spec=_decode(
+ tuple((100 + (499 * g) % 8000, 8) for g in range(32)),
+ q_layout="fused_row",
+ gate="fused_row",
+ ),
+ ),
+ # ---- context: the prefill pair is passed ----------------------------
+ Cell(
+ why="fp8 context 272 rows from 2032: causal-prefix prefill kernel",
+ spec=_decode(((2032, 272),), launches=_PREFILL, prefill_pair=True, **_FP8),
+ ),
+ Cell(
+ why="q off the fused projection row (55 x 256), 8191-row prefill",
+ spec=_decode(((0, 8191),), launches=_PREFILL, prefill_pair=True, q_layout="fused_row"),
+ ),
+ # ---- the default target: one rank's heads, table and trunk-layout pool ------
+ Cell(
+ why="tp4/tep4 decode, 1 row at position 9215 (max_seq_len, table column 287): "
+ "6 q / 1 kv heads, 64-way split-K + merge with PDL, [128, 292] table",
+ spec=_default(_spread(1, 1, 9215), _TP4, _SPLIT_K, pdl=True),
+ ),
+ Cell(
+ why="tp4/tep4 decode, 128 requests x 1 row (the largest captured batch): "
+ "6 q / 1 kv heads, 8-way split-K + merge, [128, 292] table",
+ spec=_default(_spread(128, 1, 9215), _TP4, _SPLIT_K),
+ ),
+ Cell(
+ why="tep4 MTP7 verify, 128 requests x 8 rows = 1024: 6 q / 1 kv heads, one split "
+ "and no merge, [128, 516] table (max_seq_len 16384), rows to position 16383",
+ spec=_default(_spread(128, 8, 16376), _TP4_MTP, ("splitk",), rows_per_request=8),
+ ),
+ Cell(
+ why="tp4/tep4 8192-row step (max_num_tokens) with the prefill pair: a context "
+ "request 0..8064 and 127 generation rows to position 9215, all on the prefill "
+ "kernel, 6 q / 1 kv heads",
+ spec=_default(((0, 8065),) + _spread(127, 1, 9215), _TP4, _PREFILL, prefill_pair=True),
+ ),
+ Cell(
+ why="dep4 kv_fp8 decode, 32 requests x 1 row (a rank's batch): 24 q / 2 kv heads "
+ "over an e4m3 pool, 8-way split-K + merge, [32, 292] table",
+ spec=_default(_spread(32, 1, 9215), _DEP4_FP8, _SPLIT_K),
+ ),
+ # ---- the gated decode / verify site: q and the gate off the fused row ------
+ Cell(
+ why="tep4 gated decode 1 row in a CUDA graph: q and the gate off the 17-head "
+ "(4352-wide) fused row, 64-way split-K + gated merge with PDL, captured, refilled "
+ "and replayed",
+ spec=_default(
+ _spread(1, 1, 9215), _TP4, _SPLIT_K, splits=64, pdl=True, capture=True, **_TP4_ROW
+ ),
+ ),
+ Cell(
+ why="tep4 gated decode 3 rows (PDL's 2..4-row branch): 64-way split-K + gated merge "
+ "with PDL, 17-head fused row",
+ spec=_default(_spread(3, 1, 9215), _TP4, _SPLIT_K, splits=64, pdl=True, **_TP4_ROW),
+ ),
+ Cell(
+ why="tep4 gated decode 16 rows (the 5..31-row branch): 32-way split-K + gated merge, "
+ "17-head fused row",
+ spec=_default(_spread(16, 1, 9215), _TP4, _SPLIT_K, splits=32, **_TP4_ROW),
+ ),
+ Cell(
+ why="tep4 gated decode 128 x 1 row in a CUDA graph (the largest captured batch): "
+ "8-way split-K + gated merge, 17-head fused row, captured, refilled and replayed",
+ spec=_default(
+ _spread(128, 1, 9215), _TP4, _SPLIT_K, splits=8, capture=True, **_TP4_ROW
+ ),
+ ),
+ Cell(
+ why="tep4 MTP3 gated verify 128 x 4 rows = 512 (the 257..512-row branch): 4-way "
+ "split-K + gated merge, [128, 516] table, rows to position 16383, 17-head fused row",
+ spec=_default(
+ _spread(128, 4, 16380),
+ _TP4_MTP,
+ _SPLIT_K,
+ splits=4,
+ rows_per_request=4,
+ **_TP4_ROW,
+ ),
+ ),
+ Cell(
+ why="dep4 core gated decode 32 x 1 row: 24 q / 2 kv heads over a bf16 pool, q and "
+ "the gate off the 55-head (14080-wide) fused row, 8-way split-K + gated merge",
+ spec=_default(_spread(32, 1, 9215), _DEP4, _SPLIT_K, splits=8, **_DEP4_ROW),
+ ),
+ Cell(
+ why="dep4 kv_fp8 context 8192 rows from position 0 with the prefill pair: 24 q / 2 "
+ "kv heads over an e4m3 pool on the prefill kernel, [32, 292] table",
+ spec=_default(((0, 8192),), _DEP4_FP8, _PREFILL, prefill_pair=True),
+ ),
+ Cell(
+ why="tep4 8192-row context step (max_num_tokens) with the prefill pair, q off the "
+ "17-head fused row: a context request 0..8064 and 127 generation rows to position "
+ "9215, all on the prefill kernel",
+ spec=_default(
+ ((0, 8065),) + _spread(127, 1, 9215),
+ _TP4,
+ _PREFILL,
+ prefill_pair=True,
+ q_layout="fused_row",
+ row_heads=17,
+ ),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none -- every input is an argument. What the target
+ passes: `block_table` is `qsa_block_table`, physical page ids into dim 0
+ of the caches (the *unscaled* table, not logical == physical);
+ `selected_tokens` is `qsa_expand_block_indices`' output, logical
+ positions, `-1`-padded; `k_cache` / `v_cache` are the layer pool's role
+ planes.
+
+ Preconditions:
+ - Whenever the prefill kernel can run, each selection row is packed at
+ the front with exactly the bound's count of valid entries (see the
+ class docstring). QSA's producer -- `qsa_topk_select` with zero
+ `row_starts` and `row_ends = (pos + 1) // cr`, then the expansion --
+ satisfies it (measured, including with NaN/inf in the scores).
+ Checking it needs the contents, so it is not an `is_valid` guard.
+ - The prefill pair is a property of the launch: passed because some
+ rows are context, it puts the generation rows on the prefill kernel
+ too, and the bound applies to them (the default 8192-row cell drives
+ 127 generation rows through it).
+ - There is no KV scale argument: an fp8 pool is read at face value (each code
+ widens exactly to bf16 under a bf16 q), so it must hold K/V written at unit scale.
+ - `D` must be a power of two (192 is refused); `Hq % Hkv == 0`.
+
+ The op validates its own inputs loudly (`ValueError`s; the test pins the
+ refusal of positions passed without a compress ratio; a selection width
+ the compress ratio cannot produce is refused on the prefill kernel only),
+ so `is_valid` adds nothing. `rows == 0` returns an empty `[0, Hq, D]`
+ tensor (not certified).
+
+ PDL: `TRTLLM_ENABLE_PDL` (default on) enables programmatic dependent
+ launch for multi-split launches of 4 rows or fewer; the default 1-row
+ cell runs one (split-K and merge both launched with PDL, asserted). At tp1
+ the target sends gated decode batches that small to
+ `qsa_paged_sparse_gqa_decode`. Calls are bitwise deterministic (pinned
+ on a 32-split launch and its merge).
+
+ The decode site carries up to 8 rows per request (max_draft_len 7): the
+ decode cell runs 32 requests of 8 rows, 512 `rows * Hkv` programs and a
+ 4-way split-K launch with the output gate in its merge.
+
+ The default target (tp4 / tep4 / dep4). A tp4 or tep4 rank holds 6 q heads
+ and one KV head (`Hkv` 1, a group of 6 padded to 16 MMA rows), so its
+ launch follows `rows`: 64-way split-K up to 4 rows, 32-way below 32,
+ 8-way to 256, 4-way to 512, then one split and no merge (the 1024-row
+ MTP7 verify step, where the output gate cannot ride a merge and
+ `fuses_output_gate` is False); with the prefill pair, the prefill kernel
+ past 512 rows. A dep4 rank holds tp1's 24 / 2 heads, 32 requests, and
+ under kv_fp8 an e4m3 pool, read on the split-K kernel. Its cells drive
+ the engine's trunk layout (a page advances 24 role pages; layer 11's
+ planes) over `[128, 292]`, `[128, 516]` (MTP's max_seq_len 16384) and
+ `[32, 292]` tables, with pools whose mapped pages straddle 2^31 elements of
+ the K and V planes: both kernels widen the page id to int64 before its
+ stride product. Against `reference` over the five ungated cells: at most
+ 2.00 bf16 ulp of the row scale, relative RMS at most 0.572 ulp.
+
+ The gated decode and verify site. q and the output gate are views of one
+ rank's fused input projection row (`row_heads`): 17 heads of 256 at tp4 /
+ tep4 (q heads 0..5, the gate 8..13, row stride 4352), 55 at dep4 (q 0..23,
+ the gate 28..51, 14080). `q_stride_row` is a Triton constexpr, so each row
+ width compiles its own kernels. Cells assert their split count, so every
+ 6/1 launch branch with a merge carries the gate: 64-way split with PDL at
+ 1 and 3 rows, 32-way at 16, 8-way at 128 (decode), 4-way at 512 (MTP3
+ verify, 128 x 4 rows to position 16383). dep4's decode runs it at 24/2
+ over a bf16 pool (8-way, 32 rows).
+ The 1-row and 128-row cells are captured in a CUDA graph (the split-K
+ partials come from the graph's pool); a replay over q, the gate, the
+ selection and K/V refilled in place takes the eager call's launches (PDL
+ included) and matches an eager call on the new inputs bitwise.
+ A gated output is rounded twice (the attention, then the gated
+ product), each at most one bf16 ulp of the value rounded, so its element
+ error is bounded near the band's 4 ulp rather than 2: at most 3.32 ulp of
+ the row scale and 0.667 ulp relative RMS over the seven gated cells
+ (tp1's included); the same calls ungated at most 2.00.
+ Context rows: dep4's kv_fp8 context step runs
+ 8192 rows from position 0 on the prefill kernel over an e4m3 pool, and
+ tep4's 8192-row step (max_num_tokens) with q off the 17-head row. Both at
+ most 2.00 ulp and 0.573 ulp RMS.
+
+ Pool page count. `NUM_CACHE_PAGES` is a constexpr that only bounds the
+ page id (`physical_pages < NUM_CACHE_PAGES`). The cells' pools straddle
+ 2^31 elements at sizes a normal cell holds (5592 to 33167 pages). The
+ engines' K/V views hold more: 257109 pages (tep4 bf16), 119680 (dep4 bf16)
+ and 171362 (dep4 kv_fp8), 94.2, 87.7 and 62.8 GiB as trunk pools. Probed
+ at those exact counts, each matched its cell's numbers; no cell carries
+ them.
+
+ Not certified: head geometries other than 24/2 and 6/1 x 256;
+ `tokens_per_block` other than 32; more than 8192 rows; capture of the
+ one-split launch (the 1024-row MTP7 verify, which allocates only its
+ output), of the 512-row MTP3 verify and of the prefill kernel; dep4's
+ gated decode below 16 rows (its 64- and 32-way splits at 24/2); an fp8
+ pool at 6/1 (no tp4 / tep4 member stores one) and on the prefill kernel
+ with q off the 55-head row (the fp8 context cell reads the packed
+ `[q | k | v]` slice); positions past 16383; rows whose selection
+ resolves no token (all padding, a foreign
+ request, unmapped pages); an fp16 pool or q; standalone contiguous K/V
+ buffers; int64 request indices. Hard limits: not probed; the largest
+ certified extent is 8192 rows x 2051 selected tokens, a [128, 516] table
+ and a 33167-page trunk pool.
+ """
+
+ def __call__(
+ self,
+ *,
+ q: torch.Tensor,
+ k_cache: torch.Tensor,
+ v_cache: torch.Tensor,
+ block_table: torch.Tensor,
+ selected_tokens: torch.Tensor,
+ request_indices: torch.Tensor,
+ tokens_per_block: int,
+ softmax_scale: float,
+ query_positions: Optional[torch.Tensor] = None,
+ compress_ratio: Optional[int] = None,
+ output_gate: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ return _triton_qsa_paged_sparse_gqa(
+ q=q,
+ k_cache=k_cache,
+ v_cache=v_cache,
+ block_table=block_table,
+ selected_tokens=selected_tokens,
+ request_indices=request_indices,
+ tokens_per_block=tokens_per_block,
+ softmax_scale=softmax_scale,
+ query_positions=query_positions,
+ compress_ratio=compress_ratio,
+ output_gate=output_gate,
+ )
+
+ @staticmethod
+ def fuses_output_gate(
+ q: torch.Tensor,
+ k_cache: torch.Tensor,
+ selected_tokens: torch.Tensor,
+ output_gate: torch.Tensor,
+ ) -> bool:
+ """Whether a call with these operands and no prefill pair carries `output_gate`."""
+ return _can_fuse_qsa_splitk_output_gate(q, k_cache, selected_tokens, output_gate)
+
+ def reference(
+ self,
+ *,
+ q: torch.Tensor,
+ k_cache: torch.Tensor,
+ v_cache: torch.Tensor,
+ block_table: torch.Tensor,
+ selected_tokens: torch.Tensor,
+ request_indices: torch.Tensor,
+ tokens_per_block: int,
+ softmax_scale: float,
+ query_positions: Optional[torch.Tensor] = None,
+ compress_ratio: Optional[int] = None,
+ output_gate: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ """fp32 sparse GQA over every resident selected token, rounded once;
+ with `output_gate`, then gated in fp32 and rounded again.
+
+ `query_positions` and `compress_ratio` are ignored: on a prefix-packed
+ selection the prefill kernel's bound reads exactly the valid entries.
+ Walked 32 rows at a time to bound the gathered K/V.
+ """
+ rows, num_q_heads, head_dim = q.shape
+ num_pages, num_kv_heads = k_cache.shape[0], k_cache.shape[1]
+ groups = num_q_heads // num_kv_heads
+ num_requests, table_width = block_table.shape
+ requests = request_indices.long()
+ out = torch.empty(rows, num_q_heads, head_dim, dtype=torch.float32, device=q.device)
+ for start in range(0, rows, 32):
+ stop = min(start + 32, rows)
+ chunk = selected_tokens[start:stop].long()
+ logical_page = chunk.clamp_min(0) // tokens_per_block
+ in_page = chunk.clamp_min(0) % tokens_per_block
+ request = requests[start:stop]
+ owned = (request >= 0) & (request < num_requests)
+ physical = block_table[
+ request.clamp(0, num_requests - 1)[:, None].expand_as(logical_page),
+ logical_page.clamp(max=table_width - 1),
+ ].long()
+ resident = (
+ (chunk >= 0)
+ & owned[:, None]
+ & (logical_page < table_width)
+ & (physical >= 0)
+ & (physical < num_pages)
+ )
+ physical = physical.clamp(0, num_pages - 1)
+ # What a non-resident column gathers (a clamped page) is not part of the answer:
+ # zeroed, so a pool holding NaN outside the mapped pages cannot reach the sum.
+ absent = ~resident[:, :, None, None]
+ keys = k_cache[physical, :, in_page, :].float().masked_fill(absent, 0.0)
+ values = v_cache[physical, :, in_page, :].float().masked_fill(absent, 0.0)
+ grouped = q[start:stop].reshape(-1, num_kv_heads, groups, head_dim).float()
+ scores = torch.einsum("bhgd,bkhd->bhgk", grouped, keys) * softmax_scale
+ scores = scores.masked_fill(~resident[:, None, None, :], float("-inf"))
+ # A row with nothing resident softmaxes over all -inf: NaN, i.e. zeros.
+ weights = torch.nan_to_num(torch.softmax(scores, dim=-1), nan=0.0)
+ pooled = torch.einsum("bhgk,bkhd->bhgd", weights, values)
+ out[start:stop] = pooled.reshape(stop - start, num_q_heads, head_dim)
+ if output_gate is None:
+ return out.to(q.dtype)
+ gate = output_gate.reshape(rows, num_q_heads, head_dim).float()
+ return (out.to(q.dtype).float() * torch.sigmoid(gate)).to(q.dtype)
+
+
+qsa_paged_sparse_gqa = QsaPagedSparseGqa()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_decode.py
new file mode 100644
index 000000000000..7047ce0ec03b
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_decode.py
@@ -0,0 +1,820 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""QSA sparse GQA for decode rows: a split-K pass that loads its tokens in one round trip,
+then a merge spread over the SMs.
+
+Tier 5 (CuTe DSL written for this target), both launches.
+
+Split-K pass. One CTA (eight warps) covers `CAP` consecutive selected tokens of one KV
+head of one row, so a decode row's 2051 selected tokens take 33 CTAs a KV head:
+
+1. before its PDL wait the CTA stages the KV head's query heads in shared memory and from
+ there in every warp's MMA A fragments, stages the request's block-table row (up to
+ `BT_STAGE` pages) in shared memory, and releases its dependents;
+2. after it, its tokens' pages come from shared memory and every token's K and V row
+ (16 bytes a lane, a warp eight tokens) is requested at once with cp.async, so the CTA
+ pays one DRAM round trip -- QSA's split-K kernel walks two or three 16-token tiles,
+ each a chain of dependent loads (selection, page, K/V);
+3. scores = q K^T on warp MMAs (bf16 in, fp32 out; a warp eight tokens), scaled into the
+ log2 domain and masked; one softmax over the CTA's tokens (a warp two heads); P V on
+ warp MMAs with P in bf16 (a warp 32 head dims); each head's partial times its
+ reciprocal sum is staged in shared memory and stored in 16-byte pieces with its
+ log2-sum-exp, as QSA's split-K pass stores them.
+
+Merge. A CTA covers 128 elements of a head (two CTAs a head) with up to eight warps: warp
+`w` reads its splits' partials of four adjacent elements a lane, one 16-byte load a split,
+all in flight; every warp forms the split weights itself from two LSEs a lane; the warps'
+shares are added in shared memory.
+
+Both launches release their dependents at entry, before their own waits: the merge and,
+through it, the output projection that follows launch while the top-k ahead of the
+split-K pass still runs, and the projection streams its weight through that latency-bound
+stretch. Before its wait the split-K pass reads q and the merge the output gate: the
+caller's contract is that both are complete when the grid ahead of the split-K pass
+releases it (the target's top-k does so only after its own wait; a launch without PDL
+releases at its end).
+"""
+
+from __future__ import annotations
+
+import functools
+import math
+from typing import Optional
+
+import cuda.bindings.driver as cuda_drv
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass import const_expr
+from cutlass.cute import experimental as cute_ext
+from cutlass.cute.nvgpu import cpasync, warp
+from cutlass.cute.runtime import from_dlpack, make_ptr
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa import (
+ qsa_paged_sparse_gqa as _composition,
+)
+
+# Approved exception (as for `qsa_paged_sparse_gqa`): QSA's sparse path has no torch.ops
+# surface; its row-count PDL policy is reached as Python.
+from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import _qsa_pdl_enabled
+
+#: Output elements a merge CTA covers (four adjacent ones a lane of each warp).
+CHUNK = 128
+#: Splits the merge reads: two LSEs a lane.
+MAX_SPLITS = 64
+#: Rows the target calls this entry with: the launches that run with PDL (QSA enables it up
+#: to four rows); larger batches keep the composition.
+MAX_ROWS = 4
+#: Selected tokens a split-K CTA covers.
+CAP = 64
+#: MMA rows: the query heads of a KV head, padded.
+M = 16
+#: bf16 elements of padding after each shared K, V, q and P row (keeps the MMA fragment
+#: reads free of bank conflicts).
+PAD = 8
+WARPS = 8
+THREADS = WARPS * 32
+#: Block-table entries staged in shared memory before the wait; a wider table is read
+#: from global memory after it.
+BT_STAGE = 4096
+HEAD_DIM = 256
+LOG2E = 1.4426950408889634
+
+
+class _SplitKKernel:
+ def __init__(
+ self,
+ *,
+ group,
+ width,
+ splits,
+ tpb,
+ scale,
+ bt_width,
+ num_pages,
+ k_strides,
+ v_strides,
+ num_requests,
+ pdl,
+ ):
+ assert group <= M and CAP == 2 * 32 and CAP % WARPS == 0
+ self.group, self.width, self.splits, self.tpb = group, width, splits, tpb
+ self.scale = scale * LOG2E
+ self.bt_width, self.num_pages, self.num_requests = bt_width, num_pages, num_requests
+ self.stage_bt = bt_width <= BT_STAGE
+ self.k_strides, self.v_strides = k_strides, v_strides
+ self.ld = HEAD_DIM + PAD
+ self.pdl = pdl
+
+ def smem_bytes(self) -> int:
+ def r(n):
+ return -(-n // 128) * 128
+
+ b = (
+ r(M * self.ld * 2)
+ + 2 * r(CAP * self.ld * 2)
+ + r(M * (CAP + PAD) * 2)
+ + r(M * (CAP + 4) * 4)
+ )
+ b += r(M * 2 * 4) + 2 * r(CAP * 8)
+ if self.stage_bt:
+ b += r(self.bt_width * 4)
+ return b + 256
+
+ @cute.jit
+ def __call__(self, mQ, pK, pV, mBT, mSel, mReq, pPart, mLse, stream: cuda_drv.CUstream):
+ rows = cute.size(mSel, mode=[0])
+ kv_heads = cute.size(mQ, mode=[1]) // self.group
+ mma = cute.make_tiled_mma(
+ warp.MmaF16BF16Op(cutlass.BFloat16, cutlass.Float32, (16, 8, 16)),
+ cute.make_layout((1, WARPS, 1)),
+ )
+ g2s = cute.make_copy_atom(cpasync.CopyG2SOp(), cutlass.BFloat16, num_bits_per_copy=128)
+ st = cute.make_copy_atom(
+ cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=128
+ )
+ self.kernel(mQ, pK, pV, mBT, mSel, mReq, pPart, mLse, mma, g2s, st).launch(
+ grid=[rows * kv_heads * self.splits, 1, 1],
+ block=[THREADS, 1, 1],
+ smem=self.smem_bytes(),
+ stream=stream,
+ use_pdl=self.pdl,
+ )
+
+ @cute.kernel
+ def kernel(
+ self,
+ mQ,
+ pK,
+ pV,
+ mBT,
+ mSel,
+ mReq,
+ pPart,
+ mLse,
+ mma: cute.TiledMma,
+ g2s: cute.CopyAtom,
+ st: cute.CopyAtom,
+ ):
+ tidx, _, _ = cute.arch.thread_idx()
+ bidx, _, _ = cute.arch.block_idx()
+ warp_i = cute.arch.warp_idx()
+ lane = cute.arch.lane_idx()
+ group: cutlass.Constexpr = self.group
+ splits: cutlass.Constexpr = self.splits
+ width: cutlass.Constexpr = self.width
+ tpb: cutlass.Constexpr = self.tpb
+ bt_width: cutlass.Constexpr = self.bt_width
+ ld: cutlass.Constexpr = self.ld
+ dim: cutlass.Constexpr = HEAD_DIM
+ heads = cute.size(mQ, mode=[1])
+ kv_heads = heads // group
+ rows = cute.size(mSel, mode=[0])
+ row = bidx // (kv_heads * splits)
+ kvh = (bidx // splits) % kv_heads
+ split = bidx % splits
+ c0 = split * CAP
+
+ smem = cutlass.utils.SmemAllocator()
+ sQ = smem.allocate_tensor(
+ cutlass.BFloat16, cute.make_layout((M, dim), stride=(ld, 1)), byte_alignment=128
+ )
+ sK = smem.allocate_tensor(
+ cutlass.BFloat16, cute.make_layout((CAP, dim), stride=(ld, 1)), byte_alignment=128
+ )
+ sV = smem.allocate_tensor(
+ cutlass.BFloat16, cute.make_layout((CAP, dim), stride=(ld, 1)), byte_alignment=128
+ )
+ sP = smem.allocate_tensor(
+ cutlass.BFloat16, cute.make_layout((M, CAP), stride=(CAP + PAD, 1)), byte_alignment=128
+ )
+ sS = smem.allocate_tensor(
+ cutlass.Float32, cute.make_layout((M, CAP), stride=(CAP + 4, 1)), byte_alignment=128
+ )
+ sStat = smem.allocate_tensor(
+ cutlass.Float32, cute.make_layout((M, 2), stride=(2, 1)), byte_alignment=128
+ )
+ sKo = smem.allocate_tensor(
+ cutlass.Int64, cute.make_layout((CAP,), stride=(1,)), byte_alignment=128
+ )
+ sVo = smem.allocate_tensor(
+ cutlass.Int64, cute.make_layout((CAP,), stride=(1,)), byte_alignment=128
+ )
+
+ # 1. Before the wait: the KV head's query heads (complete before the grid ahead of
+ # this one released it) into shared memory and the A fragments, the block-table row
+ # (host-filled), the release.
+ cpr: cutlass.Constexpr = dim // 8 # 16-byte pieces a row: one a lane
+ q_vec = cute.logical_divide(mQ, (None, None, 8))
+ sq_vec = cute.logical_divide(sQ, (None, 8))
+ for i in cutlass.range_constexpr(-(-(group * cpr) // THREADS)):
+ c = tidx + i * THREADS
+ cc = cutlass.min(c, group * cpr - 1)
+ h = cc // cpr
+ ch = cc % cpr
+ if c < group * cpr:
+ cute.copy(g2s, q_vec[row, kvh * group + h, (None, ch)], sq_vec[h, (None, ch)])
+ cute.arch.cp_async_commit_group()
+ for i in cutlass.range_constexpr((M - group) * dim // THREADS):
+ e = tidx + i * THREADS
+ sQ[group + e // dim, e % dim] = cutlass.BFloat16(0.0)
+ req = mReq[row]
+ req_ok = (req >= 0) & (req < self.num_requests)
+ safe_req = cutlass.min(cutlass.max(req, 0), self.num_requests - 1)
+ if const_expr(self.stage_bt):
+ sBT = smem.allocate_tensor(
+ cutlass.Int32, cute.make_layout((bt_width,), stride=(1,)), byte_alignment=128
+ )
+ for i in cutlass.range_constexpr(-(-bt_width // THREADS)):
+ e = tidx + i * THREADS
+ ec = cutlass.min(e, bt_width - 1)
+ v = mBT[safe_req, ec]
+ if e < bt_width:
+ sBT[e] = v
+ cute.arch.cp_async_wait_group(0)
+ cute.arch.sync_threads()
+ thr = mma.get_slice(tidx)
+ tCsQ = thr.partition_A(sQ)
+ tCrQ = thr.make_fragment_A(tCsQ)
+ cute.autovec_copy(tCsQ, tCrQ)
+ if const_expr(self.pdl):
+ # The merge waits before reading the partials and reads nothing before its wait
+ # that the grid ahead of this one writes.
+ cute.arch.griddepcontrol_launch_dependents()
+ cute.arch.griddepcontrol_wait()
+
+ # 2. The CTA's tokens: each one's K and V row offsets in elements, or -1.
+ ks0, ks1, ks2, _ = self.k_strides
+ vs0, vs1, vs2, _ = self.v_strides
+ if tidx < CAP:
+ col = c0 + tidx
+ colc = cutlass.min(col, width - 1)
+ tok = mSel[row, colc]
+ ok = req_ok & (col < width) & (tok >= 0)
+ safe_tok = cutlass.max(tok, 0)
+ lp = safe_tok // tpb
+ ok = ok & (lp < bt_width)
+ lpc = cutlass.min(lp, bt_width - 1)
+ page = cutlass.Int32(0)
+ if const_expr(self.stage_bt):
+ page = sBT[lpc]
+ else:
+ page = mBT[safe_req, lpc]
+ ok = ok & (page >= 0) & (page < self.num_pages)
+ pg = cutlass.Int64(cutlass.max(page, 0))
+ tip = cutlass.Int64(safe_tok % tpb)
+ hk = cutlass.Int64(kvh)
+ sKo[tidx] = cutlass.Int64(
+ cutlass.select_(ok, pg * ks0 + hk * ks1 + tip * ks2, cutlass.Int64(-1))
+ )
+ sVo[tidx] = cutlass.Int64(
+ cutlass.select_(ok, pg * vs0 + hk * vs1 + tip * vs2, cutlass.Int64(-1))
+ )
+ cute.arch.sync_threads()
+
+ # Every token's K and V row: warp w the tokens 8w .. 8w + 7, a lane 16 bytes of each;
+ # an invalid token's rows are zeros (its score is masked; V must stay finite).
+ sk_vec = cute.logical_divide(sK, (None, 8))
+ sv_vec = cute.logical_divide(sV, (None, 8))
+ k_base = pK.toint() + cutlass.Int64(lane * 16)
+ v_base = pV.toint() + cutlass.Int64(lane * 16)
+ tpw: cutlass.Constexpr = CAP // WARPS
+ for i in cutlass.range_constexpr(tpw):
+ t = warp_i * tpw + i
+ ko = sKo[t]
+ vo = sVo[t]
+ ksrc = cute.make_tensor(
+ cute.make_ptr(
+ cutlass.BFloat16,
+ k_base + cutlass.max(ko, 0) * 2,
+ cute.AddressSpace.gmem,
+ assumed_align=16,
+ ),
+ cute.make_layout((8,), stride=(1,)),
+ )
+ vsrc = cute.make_tensor(
+ cute.make_ptr(
+ cutlass.BFloat16,
+ v_base + cutlass.max(vo, 0) * 2,
+ cute.AddressSpace.gmem,
+ assumed_align=16,
+ ),
+ cute.make_layout((8,), stride=(1,)),
+ )
+ if ko >= 0:
+ cute.copy(g2s, ksrc, sk_vec[t, (None, lane)])
+ cute.copy(g2s, vsrc, sv_vec[t, (None, lane)])
+ else:
+ for e in cutlass.range_constexpr(8):
+ sK[t, lane * 8 + e] = cutlass.BFloat16(0.0)
+ sV[t, lane * 8 + e] = cutlass.BFloat16(0.0)
+ cute.arch.cp_async_commit_group()
+ cute.arch.cp_async_wait_group(0)
+ cute.arch.sync_threads()
+
+ # 3. Scores: a warp eight tokens over the whole head dim, masked into shared memory.
+ tCsK = thr.partition_B(sK)
+ tCrK = thr.make_fragment_B(tCsK)
+ cute.autovec_copy(tCsK, tCrK)
+ scores = cute.make_rmem_tensor(thr.partition_shape_C((M, CAP)), cutlass.Float32)
+ scores.fill(0.0)
+ cute.gemm(mma, scores, tCrQ, tCrK, scores)
+ cs = thr.partition_C(cute.make_identity_tensor((M, CAP)))
+ neg_inf = cutlass.Float32(-math.inf)
+ for j in cutlass.range_constexpr(cute.size(scores)):
+ h, t = cs[j]
+ v = neg_inf
+ if sKo[t] >= 0:
+ v = scores[j] * self.scale
+ sS[h, t] = v
+ cute.arch.sync_threads()
+ # The softmax: warp w the heads w and w + WARPS, a lane two tokens; padding heads
+ # get P = 0.
+ for j in cutlass.range_constexpr(M // WARPS):
+ h = warp_i + WARPS * j
+ a = sS[h, lane]
+ b = sS[h, lane + 32]
+ mx = cute.arch.warp_reduction_max(cute.arch.fmax(a, b))
+ safe_mx = cutlass.Float32(cutlass.select_(mx == neg_inf, cutlass.Float32(0.0), mx))
+ keep = h < group
+ pa = cutlass.Float32(
+ cutlass.select_(keep, cute.math.exp2(a - safe_mx), cutlass.Float32(0.0))
+ )
+ pb = cutlass.Float32(
+ cutlass.select_(keep, cute.math.exp2(b - safe_mx), cutlass.Float32(0.0))
+ )
+ den = cute.arch.warp_reduction_sum(pa + pb)
+ sP[h, lane] = pa.to(cutlass.BFloat16)
+ sP[h, lane + 32] = pb.to(cutlass.BFloat16)
+ if lane == 0:
+ sStat[h, 0] = mx
+ sStat[h, 1] = den
+ cute.arch.sync_threads()
+
+ # P V: a warp 32 of the head dims (four 8-wide tiles).
+ tCsP = thr.partition_A(sP)
+ tCrP = thr.make_fragment_A(tCsP)
+ cute.autovec_copy(tCsP, tCrP)
+ sVt = cute.make_tensor(sV.iterator, cute.make_layout((dim, CAP), stride=(1, ld)))
+ tCsV = thr.partition_B(sVt)
+ tCrV = thr.make_fragment_B(tCsV)
+ cute.autovec_copy(tCsV, tCrV)
+ acc = cute.make_rmem_tensor(thr.partition_shape_C((M, dim)), cutlass.Float32)
+ acc.fill(0.0)
+ cute.gemm(mma, acc, tCrP, tCrV, acc)
+
+ # Each head's partial times its reciprocal sum, staged over sK (free since the
+ # softmax barrier) as [group, dim] fp32, then stored in 16-byte pieces: the CTA's
+ # heads are contiguous in the partial tensor.
+ sOut = cute.make_tensor(
+ cute.recast_ptr(sK.iterator, dtype=cutlass.Float32),
+ cute.make_layout((group, dim), stride=(dim, 1)),
+ )
+ co = thr.partition_C(cute.make_identity_tensor((M, dim)))
+ h0 = lane // 4 # an m16n8 C fragment's rows: h0 and h0 + 8
+ d0 = sStat[h0, 1]
+ d1 = sStat[h0 + 8, 1]
+ r0 = cutlass.Float32(0.0)
+ r1 = cutlass.Float32(0.0)
+ if d0 > 0.0:
+ r0 = 1.0 / d0
+ if d1 > 0.0:
+ r1 = 1.0 / d1
+ for j in cutlass.range_constexpr(cute.size(acc)):
+ h, d = co[j]
+ r = cutlass.Float32(cutlass.select_(h < 8, r0, r1))
+ if h < group:
+ sOut[h, d] = acc[j] * r
+ cute.arch.sync_threads()
+ base = ((cutlass.Int64(split) * rows + row) * heads + kvh * group) * dim
+ out_ptr = pPart.toint() + base * 4
+ so_vec = cute.make_tensor(
+ sOut.iterator, cute.make_layout((group * dim // 4, 4), stride=(4, 1))
+ )
+ reg = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), cutlass.Float32)
+ for i in cutlass.range_constexpr(-(-(group * dim // 4) // THREADS)):
+ c = tidx + i * THREADS
+ if c < group * dim // 4:
+ cute.autovec_copy(so_vec[c, None], reg)
+ dst = cute.make_tensor(
+ cute.make_ptr(
+ cutlass.Float32,
+ out_ptr + cutlass.Int64(c * 16),
+ cute.AddressSpace.gmem,
+ assumed_align=16,
+ ),
+ cute.make_layout((4,), stride=(1,)),
+ )
+ cute.copy(st, reg, dst)
+ if tidx < group:
+ mx = sStat[tidx, 0]
+ den = sStat[tidx, 1]
+ lse = neg_inf
+ if den > 0.0:
+ lse = mx + cute.math.log2(den)
+ mLse[split, row, kvh * group + tidx] = lse
+
+
+class _MergeKernel:
+ def __init__(self, *, heads, dim, splits, gated, pdl):
+ assert dim % CHUNK == 0 and 1 < splits <= MAX_SPLITS
+ self.heads, self.dim, self.splits, self.gated, self.pdl = heads, dim, splits, gated, pdl
+ self.chunks = dim // CHUNK
+ self.warps = min(8, splits)
+ self.spw = -(-splits // self.warps) # splits a warp, rounded up
+
+ @cute.jit
+ def __call__(self, mPart, mLse, mGate, mOut, stream: cuda_drv.CUstream):
+ rows = cute.size(mOut, mode=[0])
+ ld = cute.make_copy_atom(cute.nvgpu.CopyG2ROp(), cutlass.Float32, num_bits_per_copy=128)
+ self.kernel(mPart, mLse, mGate, mOut, ld).launch(
+ grid=[rows * self.heads * self.chunks, 1, 1],
+ block=[self.warps * 32, 1, 1],
+ smem=self.warps * CHUNK * 4 + 64,
+ stream=stream,
+ use_pdl=self.pdl,
+ )
+
+ @cute.kernel
+ def kernel(self, mPart, mLse, mGate, mOut, ld: cute.CopyAtom):
+ tidx, _, _ = cute.arch.thread_idx()
+ bidx, _, _ = cute.arch.block_idx()
+ warp_i = cute.arch.warp_idx()
+ lane = cute.arch.lane_idx()
+ per_row: cutlass.Constexpr = self.heads * self.chunks
+ spw: cutlass.Constexpr = self.spw
+ row = bidx // per_row
+ head = (bidx % per_row) // self.chunks
+ c0 = (bidx % self.chunks) * CHUNK
+ red = cutlass.utils.SmemAllocator().allocate_tensor(
+ cutlass.Float32,
+ cute.make_layout((self.warps, CHUNK), stride=(CHUNK, 1)),
+ byte_alignment=16,
+ )
+ # The gate, before the wait: see the module notes.
+ g = cutlass.Float32(0.0)
+ if const_expr(self.gated):
+ if tidx < CHUNK:
+ g = mGate[row, head * self.dim + c0 + tidx].to(cutlass.Float32)
+ if const_expr(self.pdl):
+ # The output projection after this reads only its weight before its own wait.
+ cute.arch.griddepcontrol_launch_dependents()
+ cute.arch.griddepcontrol_wait()
+
+ # 1. The split weights, in every warp: exp2(lse - max), and their sum.
+ neg_inf = cutlass.Float32(-math.inf)
+ l0 = neg_inf
+ l1 = neg_inf
+ if lane < self.splits:
+ l0 = mLse[lane, row, head]
+ if lane + 32 < self.splits:
+ l1 = mLse[lane + 32, row, head]
+ # 2. Warp w's splits w * spw .. w * spw + spw - 1, four adjacent elements a lane:
+ # one 16-byte load a split, all in flight. A slot past the split count reads the
+ # last split again and weighs it 0.
+ part_vec = cute.logical_divide(mPart, (None, None, None, 4))
+ p = cute.make_rmem_tensor(cute.make_layout((spw, 4), stride=(4, 1)), cutlass.Float32)
+ for j in cutlass.range_constexpr(spw):
+ sidx = cutlass.min(warp_i * spw + j, self.splits - 1)
+ cute.copy(ld, part_vec[sidx, row, head, (None, (c0 + 4 * lane) // 4)], p[j, None])
+ mx = cute.arch.warp_reduction_max(cute.arch.fmax(l0, l1))
+ w0 = cutlass.Float32(0.0)
+ w1 = cutlass.Float32(0.0)
+ if mx > neg_inf:
+ w0 = cute.math.exp2(l0 - mx)
+ w1 = cute.math.exp2(l1 - mx)
+ den = cute.arch.warp_reduction_sum(w0 + w1)
+ acc = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), cutlass.Float32)
+ acc.fill(0.0)
+ for j in cutlass.range_constexpr(spw):
+ sidx = warp_i * spw + j
+ wa = cute.arch.shuffle_sync(w0, sidx % 32)
+ wb = cute.arch.shuffle_sync(w1, sidx % 32)
+ w = cutlass.Float32(cutlass.select_(sidx < 32, wa, wb))
+ w = cutlass.Float32(cutlass.select_(sidx < self.splits, w, cutlass.Float32(0.0)))
+ for e in cutlass.range_constexpr(4):
+ acc[e] = w * p[j, e] + acc[e]
+ for e in cutlass.range_constexpr(4):
+ red[warp_i, 4 * lane + e] = acc[e]
+ cute.arch.sync_threads()
+
+ # 3. The warps' shares summed in warp order; the output gate on the rounded
+ # attention value, as the composition applies it.
+ if tidx < CHUNK:
+ tot = cutlass.Float32(0.0)
+ for w_ in cutlass.range_constexpr(self.warps):
+ tot = tot + red[w_, tidx]
+ merged = cutlass.Float32(0.0)
+ if den > 0.0:
+ merged = tot / den
+ if const_expr(self.gated):
+ att = merged.to(cutlass.BFloat16).to(cutlass.Float32)
+ merged = att * (1.0 / (1.0 + cute.math.exp(-g)))
+ mOut[row, head, c0 + tidx] = merged.to(mOut.element_type)
+
+
+def _dl(t: torch.Tensor, align: int = 4):
+ return from_dlpack(t, assumed_align=align)
+
+
+def _ptr(t: torch.Tensor, dtype):
+ return make_ptr(dtype, t.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
+
+
+def splits(q, k_cache, selected_tokens) -> int:
+ """The split-K CTAs a row and KV head: one per `CAP` selected tokens."""
+ return -(-selected_tokens.shape[1] // CAP)
+
+
+@functools.lru_cache(maxsize=None)
+def _compiled_splitk(
+ rows,
+ q_shape,
+ q_stride,
+ group,
+ width,
+ tpb,
+ scale,
+ bt_shape,
+ num_pages,
+ k_strides,
+ v_strides,
+ pdl,
+):
+ heads = q_shape[1]
+ n_splits = -(-width // CAP)
+ kern = _SplitKKernel(
+ group=group,
+ width=width,
+ splits=n_splits,
+ tpb=tpb,
+ scale=scale,
+ bt_width=bt_shape[1],
+ num_pages=num_pages,
+ k_strides=k_strides,
+ v_strides=v_strides,
+ num_requests=bt_shape[0],
+ pdl=pdl,
+ )
+ dev = "cuda"
+ q = torch.empty_strided(q_shape, q_stride, dtype=torch.bfloat16, device=dev)
+ kv = torch.empty(64, dtype=torch.bfloat16, device=dev)
+ bt = torch.empty(bt_shape, dtype=torch.int32, device=dev)
+ sel = torch.empty(rows, width, dtype=torch.int32, device=dev)
+ req = torch.empty(rows, dtype=torch.int32, device=dev)
+ part = torch.empty(64, dtype=torch.float32, device=dev)
+ lse = torch.empty(n_splits, rows, heads, dtype=torch.float32, device=dev)
+ stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream)
+ return cute_ext.compile(
+ kern,
+ _dl(q, 16),
+ _ptr(kv, cutlass.BFloat16),
+ _ptr(kv, cutlass.BFloat16),
+ _dl(bt),
+ _dl(sel),
+ _dl(req),
+ _ptr(part, cutlass.Float32),
+ _dl(lse),
+ stream,
+ )
+
+
+@functools.lru_cache(maxsize=None)
+def _compiled_merge(rows, heads, dim, n_splits, gate_stride, dtype, pdl):
+ kern = _MergeKernel(
+ heads=heads, dim=dim, splits=n_splits, gated=gate_stride is not None, pdl=pdl
+ )
+ dev = "cuda"
+ part = torch.empty(n_splits, rows, heads, dim, dtype=torch.float32, device=dev)
+ lse = torch.empty(n_splits, rows, heads, dtype=torch.float32, device=dev)
+ gate = torch.empty(rows, gate_stride or heads * dim, dtype=dtype, device=dev)[:, : heads * dim]
+ out = torch.empty(rows, heads, dim, dtype=dtype, device=dev)
+ stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream)
+ return cute_ext.compile(kern, _dl(part, 16), _dl(lse), _dl(gate, 2), _dl(out, 2), stream)
+
+
+def launch(
+ q,
+ k_cache,
+ v_cache,
+ block_table,
+ selected_tokens,
+ request_indices,
+ tokens_per_block,
+ softmax_scale,
+ output_gate=None,
+ pdl=True,
+):
+ """The launch behind `QsaPagedSparseGqaDecode.__call__`; returns `[rows, H, D]` in q's dtype."""
+ rows, num_q_heads, head_dim = q.shape
+ num_kv_heads = k_cache.shape[1]
+ group_size = num_q_heads // num_kv_heads
+ width = selected_tokens.shape[1]
+ n_splits = -(-width // CAP)
+ output = torch.empty((rows, num_q_heads, head_dim), dtype=q.dtype, device=q.device)
+ if rows == 0:
+ return output
+ dev = q.device
+ use_pdl = bool(pdl and _qsa_pdl_enabled(rows))
+ partial_output = torch.empty(
+ (n_splits, rows, num_q_heads, head_dim), dtype=torch.float32, device=dev
+ )
+ partial_lse = torch.empty((n_splits, rows, num_q_heads), dtype=torch.float32, device=dev)
+ fn = _compiled_splitk(
+ rows,
+ tuple(q.shape),
+ tuple(q.stride()),
+ group_size,
+ width,
+ tokens_per_block,
+ float(softmax_scale),
+ tuple(block_table.shape),
+ k_cache.shape[0],
+ tuple(k_cache.stride()),
+ tuple(v_cache.stride()),
+ use_pdl,
+ )
+ stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream)
+ fn(
+ _dl(q, 16),
+ _ptr(k_cache, cutlass.BFloat16),
+ _ptr(v_cache, cutlass.BFloat16),
+ _dl(block_table),
+ _dl(selected_tokens),
+ _dl(request_indices),
+ _ptr(partial_output, cutlass.Float32),
+ _dl(partial_lse),
+ stream,
+ )
+ gate2d = None if output_gate is None else output_gate.reshape(rows, num_q_heads * head_dim)
+ mfn = _compiled_merge(
+ rows,
+ num_q_heads,
+ head_dim,
+ n_splits,
+ None if gate2d is None else gate2d.stride(0),
+ q.dtype,
+ use_pdl,
+ )
+ mfn(
+ _dl(partial_output, 16),
+ _dl(partial_lse),
+ _dl(output.view(rows, num_q_heads * head_dim) if gate2d is None else gate2d, 2),
+ _dl(output, 2),
+ stream,
+ )
+ return output
+
+
+def supports(q, k_cache, selected_tokens) -> bool:
+ """Whether these operands are ones this entry serves (see `is_valid`)."""
+ return (
+ q.dim() == 3
+ and 0 < q.shape[0] <= MAX_ROWS
+ and q.shape[2] == HEAD_DIM
+ and q.dtype == torch.bfloat16
+ and k_cache.dtype == torch.bfloat16
+ and q.shape[1] % k_cache.shape[1] == 0
+ and q.shape[1] // k_cache.shape[1] <= M
+ and 1 < splits(q, k_cache, selected_tokens) <= MAX_SPLITS
+ and all(t.data_ptr() % 16 == 0 for t in (q, k_cache))
+ and all(s % 8 == 0 for s in (*q.stride()[:2], *k_cache.stride()[:3]))
+ )
+
+
+class QsaPagedSparseGqaDecode(OpWrapper):
+ """`qsa_paged_sparse_gqa` without a prefill pair, for decode rows: a split-K pass whose
+ CTAs load their tokens' K/V in one round trip, then the merge -- and, with
+ `output_gate`, the gate -- as one launch. Returns a new `[rows, H, D]` tensor in q's
+ dtype.
+
+ Fusion boundary: the split-K attention, the merge, the optional sigmoid output gate.
+ The selection, the K/V store and the output projection are the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="decode, 1 row at position 8700, gated", spec=dict(pos=(8700,), gated=True)),
+ Cell(
+ why="decode, 4 rows, one a short sequence, gated",
+ spec=dict(pos=(8700, 5000, 77, 3000), gated=True),
+ ),
+ Cell(
+ why="decode, 1 row with half the selection -1",
+ spec=dict(pos=(8700,), gated=True, holes=True),
+ ),
+ )
+
+ note = """
+ Metadata consumed: as `qsa_paged_sparse_gqa` (block table, request indices). Preconditions:
+ as that entry's, plus: no prefill pair; at most `MAX_ROWS` rows; bf16 q and K/V with a head
+ dim of 256 and at most 16 query heads a KV head; more than one and at most 64 split-K CTAs
+ a row and KV head (`splits(...)`: 64 selected tokens a CTA); `output_gate` (when given)
+ `[rows, H * D]` or `[rows, H, D]` with unit inner stride in q's dtype. `supports(...)` is
+ the predicate a caller can check.
+
+ PDL: the split-K pass reads q and the block table before its wait and the merge reads the
+ output gate before its own; the selection and the K/V rows are read after the split-K
+ pass's wait. So q and the gate must be complete when the grid ahead of this entry releases
+ it: a launch without PDL releases at its end; a PDL producer of the selection (the
+ target's top-k) must release only after its own wait.
+
+ Numerics: one softmax over a CTA's 64 tokens (QSA's kernel runs an online softmax over
+ 16-token tiles) and the merge sums the splits in split order (QSA's merge reduces them as
+ a tree), so the fp32 values can differ in the last bits and the bf16 output by an ulp on
+ some elements.
+
+ Not certified: a call without `output_gate` (the target always passes one).
+ """
+
+ def __call__(
+ self,
+ q: torch.Tensor,
+ k_cache: torch.Tensor,
+ v_cache: torch.Tensor,
+ block_table: torch.Tensor,
+ selected_tokens: torch.Tensor,
+ request_indices: torch.Tensor,
+ tokens_per_block: int,
+ softmax_scale: float,
+ output_gate: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ return launch(
+ q,
+ k_cache,
+ v_cache,
+ block_table,
+ selected_tokens,
+ request_indices,
+ tokens_per_block,
+ softmax_scale,
+ output_gate,
+ )
+
+ def reference(
+ self,
+ q: torch.Tensor,
+ k_cache: torch.Tensor,
+ v_cache: torch.Tensor,
+ block_table: torch.Tensor,
+ selected_tokens: torch.Tensor,
+ request_indices: torch.Tensor,
+ tokens_per_block: int,
+ softmax_scale: float,
+ output_gate: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ """The composition entry's reference with the same operands."""
+ return _composition.reference(
+ q=q,
+ k_cache=k_cache,
+ v_cache=v_cache,
+ block_table=block_table,
+ selected_tokens=selected_tokens,
+ request_indices=request_indices,
+ tokens_per_block=tokens_per_block,
+ softmax_scale=softmax_scale,
+ query_positions=None,
+ compress_ratio=None,
+ output_gate=output_gate,
+ )
+
+ def is_valid(
+ self,
+ q: torch.Tensor,
+ k_cache: torch.Tensor,
+ v_cache: torch.Tensor,
+ block_table: torch.Tensor,
+ selected_tokens: torch.Tensor,
+ request_indices: torch.Tensor,
+ tokens_per_block: int,
+ softmax_scale: float,
+ output_gate: Optional[torch.Tensor] = None,
+ ) -> None:
+ assert q.ndim == 3 and q.shape[2] == HEAD_DIM, "q [rows, H, 256]"
+ assert 0 < q.shape[0] <= MAX_ROWS, f"1 .. {MAX_ROWS} rows"
+ assert q.dtype == torch.bfloat16 and k_cache.dtype == v_cache.dtype == torch.bfloat16, (
+ "bf16 q and K/V"
+ )
+ assert q.shape[1] % k_cache.shape[1] == 0 and q.shape[1] // k_cache.shape[1] <= M, (
+ "<= 16 q heads a KV head"
+ )
+ assert 1 < splits(q, k_cache, selected_tokens) <= MAX_SPLITS, (
+ "65 .. 4096 selected tokens a row"
+ )
+ assert q.stride(2) == 1 and k_cache.stride(3) == 1 and v_cache.stride(3) == 1, (
+ "unit inner strides"
+ )
+ assert all(t.data_ptr() % 16 == 0 for t in (q, k_cache, v_cache)), (
+ "16-byte aligned q and K/V"
+ )
+ assert all(
+ s % 8 == 0 for s in (*q.stride()[:2], *k_cache.stride()[:3], *v_cache.stride()[:3])
+ ), "q and K/V rows 16-byte aligned"
+ if output_gate is not None:
+ assert output_gate.dtype == q.dtype and output_gate.stride(-1) == 1, "gate in q's dtype"
+ assert output_gate.numel() == q.numel(), "gate covers q"
+
+
+qsa_paged_sparse_gqa_decode = QsaPagedSparseGqaDecode()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4.py
new file mode 100644
index 000000000000..734370a39e0b
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4.py
@@ -0,0 +1,1159 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""QSA sparse GQA over an NVFP4 paged pool: attention over explicitly selected tokens."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+# Approved exception (as for `qsa_paged_sparse_gqa`): QSA's sparse path has no torch.ops
+# surface; its row-count PDL policy is reached as Python. The split-K buckets are this
+# entry's own: an NVFP4 tile needs 16 tokens where the bf16 one takes 64.
+from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import _qsa_pdl_enabled
+
+#: 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,
+)
+
+#: The prefill kernel's (BLOCK_N, warps).
+_PREFILL_TILE = (32, 2)
+
+_SPLIT_K = ("splitk", "merge")
+_PREFILL = ("prefill",)
+
+
+def _spec(**overrides) -> dict:
+ """A cell spec: this checkpoint's geometry unless the cell says otherwise.
+
+ `segments` is `(first position, rows)` per request, in row order; `prefill_pair` says whether `query_positions`
+ / `compress_ratio` are passed, which the target does exactly when the step has a
+ context request; `pool_pages` fixes the pool's page count (the requests then sit
+ on its last pages); `capture` sweeps the decode site under a CUDA graph.
+ """
+ return {**_CHECKPOINT, **overrides}
+
+
+def _decode(segments, launches=_SPLIT_K, **overrides) -> dict:
+ base = dict(segments=segments, slots=32, prefill_pair=False, launches=launches)
+ return _spec(**{**base, **overrides})
+
+
+#: One dep4 rank's engine under `kv_nvfp4` (attention DP, so tp1's 24/2 x 256 geometry):
+#: `max_batch_size` 32 a rank and `max_seq_len` 9216 give a `[32, 292]` table
+#: (`ceil((9216 + 1) / 32)` rounded up to a multiple of 4); the trunk layout's 24 role pages
+#: a slot, layer 11's planes the last of them.
+_DEP4 = dict(pages=292, slots=32, slot_pages=24, layer=11)
+
+#: The byte offset the dep4 cells' pools straddle: one page of layer 11's scale plane
+#: advances 24 x 2 x 32 x 16 bytes, so from page 87382 on it lies past 2^31 bytes (the data
+#: plane, 8x wider, from page 10923). A dep4 pool holds far more pages than either.
+_OFFSET_EDGE = 2**31
+
+
+def _dep4_context(segments, pool_pages) -> dict:
+ """A context step of dep4's kv_nvfp4 engine: the prefill pair, the prefill kernel."""
+ return _decode(
+ segments,
+ launches=_PREFILL,
+ prefill_pair=True,
+ pool_pages=pool_pages,
+ offset_edge=_OFFSET_EDGE,
+ **_DEP4,
+ )
+
+
+@triton.jit
+def _e2m1_to_f32(code):
+ """e2m1 code (0..15, bit 3 the sign) -> its value, exactly."""
+ 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 _nvfp4_halves(packed, scales):
+ """`[N, D/2]` packed e2m1 bytes (low nibble the even element) and one e4m3 scale per byte
+ -> the even and the odd elements as two `[N, D/2]` f16 tensors, each `e2m1 * scale`:
+ `cvt.rn.f16x2.e2m1x2` and `cvt.rn.f16x2.e4m3x2`, then one f16x2 multiply, exact (two by
+ four significant bits, inside f16's range)."""
+ return tl.inline_asm_elementwise(
+ asm="""
+ {
+ .reg .b8 b0, b1, b2, b3, s0, s1, s2, s3;
+ .reg .b16 t0, t1, t2, t3;
+ .reg .b32 h0, h1, h2, h3, c0, c1, c2, c3;
+ mov.b32 {b0, b1, b2, b3}, $4;
+ mov.b32 {s0, s1, s2, s3}, $5;
+ mov.b16 t0, {s0, s0};
+ mov.b16 t1, {s1, s1};
+ mov.b16 t2, {s2, s2};
+ mov.b16 t3, {s3, s3};
+ cvt.rn.f16x2.e4m3x2 c0, t0;
+ cvt.rn.f16x2.e4m3x2 c1, t1;
+ cvt.rn.f16x2.e4m3x2 c2, t2;
+ cvt.rn.f16x2.e4m3x2 c3, t3;
+ cvt.rn.f16x2.e2m1x2 h0, b0;
+ cvt.rn.f16x2.e2m1x2 h1, b1;
+ cvt.rn.f16x2.e2m1x2 h2, b2;
+ cvt.rn.f16x2.e2m1x2 h3, b3;
+ mul.rn.f16x2 h0, h0, c0;
+ mul.rn.f16x2 h1, h1, c1;
+ mul.rn.f16x2 h2, h2, c2;
+ mul.rn.f16x2 h3, h3, c3;
+ prmt.b32 $0, h0, h1, 0x5410;
+ prmt.b32 $1, h2, h3, 0x5410;
+ prmt.b32 $2, h0, h1, 0x7632;
+ prmt.b32 $3, h2, h3, 0x7632;
+ }
+ """,
+ constraints="=r,=r,=r,=r,r,r",
+ args=[packed, scales],
+ dtype=(tl.float16, tl.float16),
+ is_pure=True,
+ pack=4,
+ )
+
+
+@triton.jit
+def _load_halves(
+ cache,
+ scale_rows,
+ pages,
+ kv_head,
+ tokens,
+ valid,
+ stride_page,
+ stride_head,
+ stride_token,
+ SCALE_STEP: tl.constexpr,
+ BLOCK_N: tl.constexpr,
+ HEAD_DIM: tl.constexpr,
+):
+ """A tile's even and odd elements of one role plane. `scale_rows` points at each token's
+ first scale; its D/16 scales lie SCALE_STEP apart (K's token-major rows: 1; V's
+ four-token groups `[tpb/4, D/16, 4]`: 4). Each byte and each scale is read once."""
+ HALF: tl.constexpr = HEAD_DIM // 2
+ GROUPS: tl.constexpr = HEAD_DIM // 16
+ packed = tl.load(
+ cache
+ + pages[:, None] * stride_page
+ + kv_head * stride_head
+ + tokens[:, None] * stride_token
+ + tl.arange(0, HALF)[None, :],
+ mask=valid[:, None],
+ other=0,
+ )
+ scales = tl.load(
+ scale_rows[:, None] + tl.arange(0, GROUPS)[None, :] * SCALE_STEP,
+ mask=valid[:, None],
+ other=0,
+ )
+ scales = tl.reshape(tl.broadcast_to(scales[:, :, None], [BLOCK_N, GROUPS, 8]), [BLOCK_N, HALF])
+ return _nvfp4_halves(packed, scales)
+
+
+@triton.jit
+def _attend_tile(
+ q_even,
+ q_odd,
+ running_max,
+ running_sum,
+ acc_even,
+ acc_odd,
+ columns,
+ row,
+ safe_request,
+ valid_request,
+ kv_head,
+ score_scale,
+ selected_tokens,
+ selected_stride_row,
+ block_table,
+ block_table_stride_request,
+ page_table_width,
+ num_cache_pages,
+ k_cache,
+ v_cache,
+ k_scale_cache,
+ v_scale_cache,
+ 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,
+ TOKENS_PER_BLOCK: tl.constexpr,
+ TOPK: tl.constexpr,
+ BLOCK_N: tl.constexpr,
+ HEAD_DIM: tl.constexpr,
+):
+ """One online-softmax step over the selected columns `columns` of `row`: the even and odd
+ halves of K against q's, scores in the log2 domain, P rounded to f16 against V's halves."""
+ GROUPS: tl.constexpr = HEAD_DIM // 16
+ logical_tokens = tl.load(
+ selected_tokens + row * selected_stride_row + columns, mask=columns < TOPK, other=-1
+ )
+ safe_tokens = tl.maximum(logical_tokens, 0)
+ logical_pages = safe_tokens // TOKENS_PER_BLOCK
+ tokens = safe_tokens % TOKENS_PER_BLOCK
+ valid = (
+ valid_request
+ & (columns < TOPK)
+ & (logical_tokens >= 0)
+ & (logical_pages < page_table_width)
+ )
+ 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 &= (pages >= 0) & (pages < num_cache_pages)
+ pages = tl.maximum(pages, 0)
+
+ k_even, k_odd = _load_halves(
+ k_cache,
+ k_scale_cache + pages * ks_stride_page + kv_head * ks_stride_head + tokens * GROUPS,
+ pages,
+ kv_head,
+ tokens,
+ valid,
+ k_stride_page,
+ k_stride_head,
+ k_stride_token,
+ 1,
+ BLOCK_N,
+ HEAD_DIM,
+ )
+ scores = tl.dot(q_even, tl.trans(k_even))
+ scores = tl.dot(q_odd, tl.trans(k_odd), scores)
+ scores = tl.where(valid[None, :], scores * score_scale, -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)
+
+ v_even, v_odd = _load_halves(
+ v_cache,
+ v_scale_cache
+ + pages * vs_stride_page
+ + kv_head * vs_stride_head
+ + (tokens // 4) * (4 * GROUPS)
+ + tokens % 4,
+ pages,
+ kv_head,
+ tokens,
+ valid,
+ v_stride_page,
+ v_stride_head,
+ v_stride_token,
+ 4,
+ BLOCK_N,
+ HEAD_DIM,
+ )
+ weights = probabilities.to(tl.float16)
+ acc_even = tl.dot(weights, v_even, acc_even * correction[:, None])
+ acc_odd = tl.dot(weights, v_odd, acc_odd * correction[:, None])
+ running_sum = running_sum * correction + tl.sum(probabilities, axis=1)
+ return next_max, running_sum, acc_even, acc_odd
+
+
+@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)
+ # The global scales scale the scores and the normalized output, not each element.
+ score_scale = softmax_scale * _LOG2_E * tl.load(kv_scale_quant_orig + 1)
+ v_global = tl.load(kv_scale_quant_orig + 2)
+
+ HALF: tl.constexpr = HEAD_DIM // 2
+ head_offsets = tl.arange(0, BLOCK_M)
+ half_offsets = tl.arange(0, HALF)
+ token_offsets = tl.arange(0, BLOCK_N)
+ query_heads = kv_head * GROUP_SIZE + head_offsets
+ head_mask = (head_offsets < GROUP_SIZE)[:, None]
+ # q's even and odd elements as f16, so both MMA operands are f16 (q is bf16 within
+ # f16's range after its norm; an NVFP4 value times its scale is exact in f16).
+ q_even_ptr = (
+ q + row * q_stride_row + query_heads[:, None] * q_stride_head + 2 * half_offsets[None, :]
+ )
+ q_even = tl.load(q_even_ptr, mask=head_mask, other=0.0).to(tl.float16)
+ q_odd = tl.load(q_even_ptr + 1, mask=head_mask, other=0.0).to(tl.float16)
+
+ running_max = tl.full([BLOCK_M], -float("inf"), tl.float32)
+ running_sum = tl.zeros([BLOCK_M], tl.float32)
+ acc_even = tl.zeros([BLOCK_M, HALF], tl.float32)
+ acc_odd = tl.zeros([BLOCK_M, HALF], tl.float32)
+ 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):
+ columns = tile * BLOCK_N + token_offsets
+ running_max, running_sum, acc_even, acc_odd = _attend_tile(
+ q_even,
+ q_odd,
+ running_max,
+ running_sum,
+ acc_even,
+ acc_odd,
+ columns,
+ row,
+ safe_request,
+ valid_request,
+ kv_head,
+ score_scale,
+ selected_tokens,
+ selected_stride_row,
+ block_table,
+ block_table_stride_request,
+ page_table_width,
+ num_cache_pages,
+ k_cache,
+ v_cache,
+ k_scale_cache,
+ v_scale_cache,
+ 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,
+ TOKENS_PER_BLOCK,
+ TOPK,
+ BLOCK_N,
+ HEAD_DIM,
+ )
+
+ denominator = tl.maximum(running_sum[:, None], 1.0e-20)
+ has_values = running_sum[:, None] > 0
+ out_ptr = (
+ output
+ + row * output_stride_row
+ + query_heads[:, None] * output_stride_head
+ + 2 * half_offsets[None, :]
+ )
+ tl.store(out_ptr, tl.where(has_values, acc_even / denominator, 0.0) * v_global, mask=head_mask)
+ tl.store(
+ out_ptr + 1, tl.where(has_values, acc_odd / denominator, 0.0) * v_global, mask=head_mask
+ )
+
+
+@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)
+ # The global scales scale the scores and the normalized output, not each element.
+ score_scale = softmax_scale * _LOG2_E * tl.load(kv_scale_quant_orig + 1)
+ v_global = tl.load(kv_scale_quant_orig + 2)
+
+ HALF: tl.constexpr = HEAD_DIM // 2
+ head_offsets = tl.arange(0, BLOCK_M)
+ half_offsets = tl.arange(0, HALF)
+ token_offsets = tl.arange(0, BLOCK_N)
+ query_heads = kv_head * GROUP_SIZE + head_offsets
+ head_mask = (head_offsets < GROUP_SIZE)[:, None]
+ # q's even and odd elements as f16, so both MMA operands are f16 (q is bf16 within
+ # f16's range after its norm; an NVFP4 value times its scale is exact in f16).
+ q_even_ptr = (
+ q + row * q_stride_row + query_heads[:, None] * q_stride_head + 2 * half_offsets[None, :]
+ )
+ q_even = tl.load(q_even_ptr, mask=head_mask, other=0.0).to(tl.float16)
+ q_odd = tl.load(q_even_ptr + 1, mask=head_mask, other=0.0).to(tl.float16)
+
+ running_max = tl.full([BLOCK_M], -float("inf"), tl.float32)
+ running_sum = tl.zeros([BLOCK_M], tl.float32)
+ acc_even = tl.zeros([BLOCK_M, HALF], tl.float32)
+ acc_odd = tl.zeros([BLOCK_M, HALF], 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):
+ columns = tile * BLOCK_N + token_offsets
+ running_max, running_sum, acc_even, acc_odd = _attend_tile(
+ q_even,
+ q_odd,
+ running_max,
+ running_sum,
+ acc_even,
+ acc_odd,
+ columns,
+ row,
+ safe_request,
+ valid_request,
+ kv_head,
+ score_scale,
+ selected_tokens,
+ selected_stride_row,
+ block_table,
+ block_table_stride_request,
+ page_table_width,
+ num_cache_pages,
+ k_cache,
+ v_cache,
+ k_scale_cache,
+ v_scale_cache,
+ 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,
+ TOKENS_PER_BLOCK,
+ TOPK,
+ BLOCK_N,
+ HEAD_DIM,
+ )
+
+ has_values = running_sum > 0
+ denominator = tl.maximum(running_sum[:, None], 1.0e-20)
+ even = tl.where(has_values[:, None], acc_even / denominator, 0.0) * v_global
+ odd = tl.where(has_values[:, None], acc_odd / denominator, 0.0) * v_global
+ if NUM_SPLITS == 1:
+ out_ptr = (
+ output
+ + row * output_stride_row
+ + query_heads[:, None] * output_stride_head
+ + 2 * half_offsets[None, :]
+ )
+ else:
+ out_ptr = (
+ partial_output
+ + ((split * num_rows + row) * NUM_QUERY_HEADS + query_heads[:, None]) * HEAD_DIM
+ + 2 * half_offsets[None, :]
+ )
+ partial_lse_values = tl.where(
+ has_values,
+ running_max + tl.math.log2(tl.maximum(running_sum, 1.0e-20)),
+ -float("inf"),
+ )
+ tl.store(
+ partial_lse + (split * num_rows + row) * NUM_QUERY_HEADS + query_heads,
+ partial_lse_values,
+ mask=head_offsets < GROUP_SIZE,
+ )
+ tl.store(out_ptr, even, mask=head_mask)
+ tl.store(out_ptr + 1, odd, mask=head_mask)
+ 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) by programs = rows x Hkv."""
+ base_programs = rows * num_kv_heads
+ if base_programs <= 4:
+ block_n, target_splits, num_warps = 32, 64, 4
+ elif base_programs <= 16:
+ block_n, target_splits, num_warps = 32, 16, 4
+ elif base_programs <= 32:
+ block_n, target_splits, num_warps = 32, 8, 4
+ elif base_programs <= 64:
+ block_n, target_splits, num_warps = 32, 4, 4
+ elif base_programs <= 128:
+ block_n, target_splits, num_warps = 64, 4, 2
+ elif base_programs <= 256:
+ block_n, target_splits, num_warps = 32, 8, 2
+ elif base_programs <= 512:
+ block_n, target_splits, num_warps = 32, 4, 2
+ elif base_programs <= 1024:
+ block_n, target_splits, num_warps = 32, 2, 2
+ else:
+ block_n, target_splits, num_warps = 32, 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 _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=_PREFILL_TILE[0],
+ HEAD_DIM=head_dim,
+ num_warps=_PREFILL_TILE[1],
+ 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 _qsa_pdl_enabled(rows)
+ if num_splits == 1:
+ partial_output = partial_lse = output
+ else:
+ partial_output = torch.empty((num_splits, *q.shape), dtype=torch.float32, device=q.device)
+ partial_lse = torch.empty(
+ (num_splits, rows, num_q_heads), dtype=torch.float32, device=q.device
+ )
+ _nvfp4_sparse_gqa_splitk_kernel[(rows, num_kv_heads, num_splits)](
+ q,
+ k_cache,
+ v_cache,
+ k_scale_cache,
+ v_scale_cache,
+ kv_scale_quant_orig,
+ block_table,
+ selected_tokens,
+ request_indices,
+ partial_output,
+ partial_lse,
+ output,
+ softmax_scale,
+ rows,
+ num_pages,
+ block_table.shape[0],
+ block_table.shape[1],
+ q.stride(0),
+ q.stride(1),
+ *plane_strides,
+ block_table.stride(0),
+ selected_tokens.stride(0),
+ output.stride(0),
+ output.stride(1),
+ GROUP_SIZE=group_size,
+ NUM_QUERY_HEADS=num_q_heads,
+ TOKENS_PER_BLOCK=tokens_per_block,
+ TOPK=selected_tokens.shape[1],
+ NUM_SPLITS=num_splits,
+ NUM_TILES=num_tiles,
+ BLOCK_M=block_m,
+ BLOCK_N=block_n,
+ HEAD_DIM=head_dim,
+ USE_PDL=use_pdl,
+ num_warps=num_warps,
+ num_stages=2,
+ launch_pdl=use_pdl,
+ )
+ if num_splits == 1:
+ return output
+ _nvfp4_merge_splitk_kernel[(rows, num_q_heads)](
+ partial_output,
+ partial_lse,
+ output,
+ rows,
+ output.stride(0),
+ output.stride(1),
+ NUM_QUERY_HEADS=num_q_heads,
+ NUM_SPLITS=num_splits,
+ BLOCK_SPLITS=triton.next_power_of_2(num_splits),
+ HEAD_DIM=head_dim,
+ USE_PDL=use_pdl,
+ num_warps=2,
+ num_stages=1,
+ launch_pdl=use_pdl,
+ )
+ return output
+
+
+def _dequantize_tokens(data, scales, pages, slots, global_scale, grouped: bool) -> torch.Tensor:
+ """[..., Hkv, D] fp32 values of the tokens at (pages, slots) of one NVFP4 role plane."""
+ heads, tokens_per_block, packed = data.shape[1], data.shape[2], data.shape[3]
+ head_dim = packed * _PER_BYTE
+ groups = head_dim // _SF_VEC
+ raw = data[pages, :, slots, :].long() # [..., H, D/2]
+ lut = torch.tensor(_E2M1_VALUES, dtype=torch.float32, device=data.device)
+ lut = torch.cat([lut, -lut])
+ values = torch.stack([lut[raw & 0xF], lut[raw >> 4]], dim=-1).flatten(-2)
+ if grouped:
+ flat = scales.reshape(scales.shape[0], heads, tokens_per_block * groups)
+ base = (slots // _V_SF_QUAD) * (_V_SF_QUAD * groups) + slots % _V_SF_QUAD
+ columns = base[..., None] + _V_SF_QUAD * torch.arange(groups, device=data.device)
+ sf = torch.stack([flat[pages[..., None], h, columns] for h in range(heads)], dim=-2)
+ else:
+ sf = scales[pages, :, slots, :]
+ sf = sf.view(torch.float8_e4m3fn).float().repeat_interleave(_SF_VEC, dim=-1)
+ return values * sf * global_scale
+
+
+class QsaPagedSparseGqaNvfp4(OpWrapper):
+ """Sparse GQA over the selected tokens of an NVFP4 HND paged pool.
+
+ Triton kernels defined in this file. For row `r`, query head `h`, KV head
+ `h // (Hq // Hkv)` (contiguous groups), over the columns `j` of
+ `selected_tokens[r]` that name a resident token::
+
+ t = selected_tokens[r, j] # logical position
+ page = block_table[request_indices[r], t // tokens_per_block]
+ K[t] = e2m1(k_cache) * e4m3(k_scale_cache) # exact in f16
+ V[t] = e2m1(v_cache) * e4m3(v_scale_cache)
+ s_j = (q[r, h] . K[t]) * softmax_scale * kv_scale_quant_orig[1]
+ out[r, h] = kv_scale_quant_orig[2] * 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]`.
+ The kernels decode each byte with `cvt.rn.f16x2.e2m1x2` and scale it in f16,
+ which is exact, and run the scores and P.V on f16 MMAs with fp32
+ accumulation (q converted to f16, P rounded to f16 before P.V), with the
+ online softmax and split-K merge of QSA's bf16 kernels. The second-level
+ scales scale the scores and the normalized output, not each element.
+
+ A negative index, a page past the table's width, a negative or out-of-pool
+ table entry, and a request index outside the table's rows contribute
+ nothing; a row left with nothing returns zeros. Returns a fresh, contiguous
+ `[rows, Hq, D]` tensor of `q.dtype`.
+
+ Two kernels, chosen inside the call: with `query_positions` and
+ `compress_ratio` given and `rows * Hkv > 512`, a causal-prefix prefill
+ kernel; otherwise split-K, plus a merge kernel when it used more than one
+ split. **The prefill pair is a bound, not a mask**, exactly as on the bf16
+ entry: the prefill kernel reads only `min(visible // cr, block_topk) * cr +
+ visible % cr` leading columns of a row.
+
+ Fusion boundary: the block-table gather, the NVFP4 dequantize, scale,
+ online softmax, P.V and the split-K merge are inside. Selection, appending
+ K/V to the pool, and q's projection, norm, RoPE and output gate are not.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="decode: 4 requests x 4 rows, 8-way split-K + merge, in the engine's trunk "
+ "layout: layer 11's planes among the 24 role pages of a slot",
+ spec=_decode(
+ tuple((4000 + 37 * i, 4) for i in range(4)), slot_pages=24, layer=11, pool_pages=600
+ ),
+ ),
+ Cell(
+ why="context 4096 rows from 0: rows see 1..4096 tokens",
+ spec=_decode(((0, 4096),), launches=_PREFILL, prefill_pair=True),
+ ),
+ Cell(
+ why="captured decode: G 1, 8 and 32 requests x 1 and 8 rows, replay == eager",
+ spec=_decode((), slots=32, capture=dict(groups=(1, 8, 32), rows_per_request=(1, 8))),
+ ),
+ Cell(
+ why="dep4 8192-row context chunk (max_num_tokens): a request at 1152..9215 "
+ "(max_seq_len, table column 287) and one from 0, [32, 292] table, trunk layout, "
+ "pages straddling 2^31 bytes of the scale plane; 87408 pages",
+ spec=_dep4_context(((1152, 8064), (0, 128)), pool_pages=87408),
+ ),
+ Cell(
+ why="dep4 512-row context step (the chunked variant's max_num_tokens): four "
+ "requests' chunks, one ending at max_seq_len, one crossing the 2048-token "
+ "selection budget; 87401 pages (not a multiple of 16)",
+ spec=_dep4_context(((8960, 256), (0, 160), (2000, 64), (5000, 32)), pool_pages=87401),
+ ),
+ )
+
+ note = """
+ 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). In the default
+ target only dep4's `kv_nvfp4` engine binds an NVFP4 pool, with tp1's 24/2 x 256
+ geometry (attention DP): its context steps, up to `max_num_tokens` 8192 rows (512
+ in the chunked variant), over a `[32, 292]` table (`max_batch_size` 32 a rank,
+ `max_seq_len` 9216) and the trunk layout's 24 role pages a slot.
+
+ Launch shapes at Hkv 2 (programs = rows x 2): split-K at 32-token tiles
+ 64-way up to 2 rows, 16-way up to 8, 8-way up to 16 and 4-way up to 32
+ (4 warps); 64-token tiles 4-way up to 64 rows; 32-token tiles 8-way up to
+ 128, 4-way up to 256 and 2-way up to 512 rows (2 warps); each followed by
+ the merge. Past 256 rows the prefill kernel (32-token tiles, 2 warps) when the pair
+ is passed; past 512 rows otherwise one split and no merge. The decode site
+ reaches 256 rows at max_draft_len 7 (32 requests x 8).
+
+ Each byte is decoded once, into the even and the odd halves of the head
+ dimension, which the kernels keep apart: q's halves against K's on two MMAs,
+ P against V's halves into two accumulators, the output stored back
+ interleaved.
+
+ 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: P is rounded to f16 here and to
+ bf16 there, and the second-level scales act on the scores and the output.
+
+ 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 are Triton, JIT-compiled on first call for the
+ device that runs them; every cell ran that path on sm_103 (GB300). They take
+ the pool's page count (`num_pages`), the
+ table's row count (`num_requests`) and the split-K kernel the step's row
+ count as plain ints, and Triton compiles each apart at 1, at multiples of 16
+ and otherwise. Every cell runs a 32-row table (a multiple of 16); the other
+ table-height classes are not certified. The page count is driven at both the
+ multiple-of-16 class and the other (the dep4 cells: 87408 and 87401). Every other
+ plain int -- the strides, the table widths 260 (tp1) and 292 (dep4) and the
+ selection's row stride 2051 -- is fixed at the served values the cells use.
+
+ Large pools: every page index is widened to 64 bits before its product with a
+ plane's page stride. In the trunk layout one page of a layer's scale plane
+ advances 24 x 2 x 32 x 16 bytes and of its data plane 8x that, so a pool past
+ 87381 pages addresses both planes past 2^31 bytes; the dep4 cells place their
+ requests' pages on both sides of that edge in the scale plane (the test asserts
+ it), with every scale outside them NaN. Against `reference`: at most 2.00 bf16
+ ulp of a row's magnitude and 0.25 ulp relative RMS at 8192 rows, 1.98 / 0.24
+ ulp at 512 rows, inside the test's band of 4 per element and 2 RMS.
+
+ Preconditions:
+ - Whenever the prefill kernel can run, each selection row is packed at the
+ 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`.
+ - `q` within f16's range: the kernels convert it (the target's q follows its
+ RMS norm).
+ - An SM 100-family GPU: the byte decode is `cvt.rn.f16x2.e2m1x2`.
+
+ The callable rejects loudly (ValueError) every operand whose dtype, shape or
+ layout it cannot read (the test pins the half-passed prefill pair); `rows ==
+ 0` returns an empty `[0, Hq, D]` tensor (not certified). Calls are bitwise
+ deterministic.
+
+ Every row, request, page and position index is widened to 64 bits before any
+ product. Not certified: head geometries other than 24/2 x 256;
+ `tokens_per_block` other than 32; more than 8192 rows; capture of the
+ context/mixed site; a pool of one slot; second-level scales other than 1.0
+ (the target binds ones); the prefill pair on the split-K kernel (256 context
+ rows or fewer); one split without a merge; rows whose selection resolves no
+ token (all padding, a foreign request, unmapped pages); an fp16 or unsliced q;
+ int64 request indices; dep4's own pool page count (the cells' pools
+ are sized to the 2^31-byte edge rather than to an engine's allocation, which
+ is larger); positions past `max_seq_len` 9216.
+ """
+
+ def __call__(
+ 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/ops/attention/qsa_paged_sparse_gqa_nvfp4_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4_decode.py
new file mode 100644
index 000000000000..7dee4f99f9c0
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4_decode.py
@@ -0,0 +1,1606 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""QSA sparse GQA over an NVFP4 paged K/V pool, decode rows: CuTe DSL split-K + merge.
+
+Both launches are CuTe DSL. One CTA (eight warps) serves one
+(row, KV head, split) and walks the split's tiles of 64 selected tokens with an online
+softmax; the split count is chosen by row count so the grid stays about two waves deep
+(33 splits of one tile at 1 row, one split of 33 tiles from 128 rows up).
+
+Per tile:
+
+1. raw NVFP4 bytes land in shared memory with cp.async, one tile ahead (two buffers):
+ each token's 128 K bytes, its 16 K scale bytes (token-major), its 128 V bytes and the
+ 64-byte four-token V-scale block it sits in; a token that does not resolve (padding,
+ foreign request, unmapped page) gets zeros instead;
+2. the V tile is dequantized into an f16 tile (`cvt.rn.f16x2.e2m1x2` x `cvt.rn.f16x2.e4m3x2`,
+ exact) for `ldmatrix.trans`;
+3. S = q K^T on `mma.sync.m16n8k16` (f16 in, f32 out), a warp eight tokens over the whole
+ head: K's B fragments are decoded straight from the raw bytes in registers. The head
+ dimension is permuted identically in q and K (lane quad c of k-step pair m reads the
+ bytes 16m + 4c .. 16m + 4c + 3, dims 32m + 8c .. 32m + 8c + 7, inside one scale block),
+ which leaves the dot product unchanged;
+4. online softmax in the log2 domain across the warps (two shared-memory exchanges), P
+ rounded to f16 as the Triton kernel does;
+5. O += P V on `mma.sync` (a warp 32 head dims), P from shared memory with `ldmatrix`,
+ V with `ldmatrix.trans`.
+
+The second-level scales act on the scores (K) and the normalized output (V), as in the
+Triton op. With one split the CTA stores the output itself (optionally times the
+sigmoid of the output gate, applied to the bf16-rounded attention value); otherwise it
+stores an fp32 partial and its log2-sum-exp for the merge kernel.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass import const_expr
+from cutlass._mlir.dialects import llvm
+from cutlass.cute.runtime import from_dlpack
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4 import (
+ qsa_paged_sparse_gqa_nvfp4 as _composition,
+)
+from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
+
+TILE = 64
+WARPS = 8
+THREADS = WARPS * 32
+HEAD_DIM = 256
+#: Bytes of a raw K row in shared memory (128 + 16 of padding: conflict-free fragment reads).
+RAW_LD = 144
+#: Bytes of a raw V row (unpadded: the dequant pass reads 32 bytes a thread).
+RAW_LDV = 128
+#: Raw tiles in flight (cp.async ring).
+STAGES = 3
+#: Bytes of an f16 V row in shared memory (256 + 8 elements: conflict-free ldmatrix).
+SV_LD = 528
+#: Bytes of an f16 P row (64 + 8 elements).
+SP_LD = 144
+LOG2E = 1.4426950408889634
+#: Selected tokens a merge CTA covers per head (four adjacent ones a lane of each warp).
+CHUNK = 128
+MAX_SPLITS = 64
+
+
+# ---------------------------------------------------------------------------------------
+# inline PTX helpers
+# ---------------------------------------------------------------------------------------
+
+
+def _ir(x, cls, loc, ip):
+ return cls(x).ir_value(loc=loc, ip=ip)
+
+
+def _asm(ret, args, text, cons, loc, ip, side=False):
+ return llvm.inline_asm(
+ ret,
+ args,
+ text,
+ cons,
+ has_side_effects=side,
+ is_align_stack=False,
+ asm_dialect=llvm.AsmDialect.AD_ATT,
+ loc=loc,
+ ip=ip,
+ )
+
+
+def _unpack(res, n, ty, loc, ip):
+ return [
+ ty(llvm.extractvalue(T.i32() if ty is cutlass.Int32 else T.f32(), res, [i], loc=loc, ip=ip))
+ for i in range(n)
+ ]
+
+
+@dsl_user_op
+def _cp16(saddr, gaddr, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [_ir(saddr, cutlass.Int32, loc, ip), _ir(gaddr, cutlass.Int64, loc, ip)],
+ "{ cp.async.cg.shared.global [$1], [$2], 16; mov.u32 $0, 0; }",
+ "=r,r,l",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _cp16z(saddr, gaddr, nbytes, *, loc=None, ip=None):
+ """cp.async of 16 bytes, zero-filled past `nbytes` (0 or 16): no branch for a dead token."""
+ _asm(
+ T.i32(),
+ [
+ _ir(saddr, cutlass.Int32, loc, ip),
+ _ir(gaddr, cutlass.Int64, loc, ip),
+ _ir(nbytes, cutlass.Int32, loc, ip),
+ ],
+ "{ cp.async.cg.shared.global [$1], [$2], 16, $3; mov.u32 $0, 0; }",
+ "=r,r,l,r",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _sts4(saddr, a, b, c, d, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [_ir(saddr, cutlass.Int32, loc, ip)]
+ + [_ir(v, cutlass.Int32, loc, ip) for v in (a, b, c, d)],
+ "{ st.shared.v4.b32 [$1], {$2, $3, $4, $5}; mov.u32 $0, 0; }",
+ "=r,r,r,r,r,r",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _sts1(saddr, a, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [_ir(saddr, cutlass.Int32, loc, ip), _ir(a, cutlass.Int32, loc, ip)],
+ "{ st.shared.b32 [$1], $2; mov.u32 $0, 0; }",
+ "=r,r,r",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _lds1(saddr, *, loc=None, ip=None):
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [_ir(saddr, cutlass.Int32, loc, ip)],
+ "ld.shared.b32 $0, [$1];",
+ "=r,r",
+ loc,
+ ip,
+ side=True,
+ )
+ )
+
+
+@dsl_user_op
+def _lds4(saddr, *, loc=None, ip=None):
+ res = _asm(
+ llvm.StructType.get_literal([T.i32()] * 4),
+ [_ir(saddr, cutlass.Int32, loc, ip)],
+ "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];",
+ "=r,=r,=r,=r,r",
+ loc,
+ ip,
+ side=True,
+ )
+ return _unpack(res, 4, cutlass.Int32, loc, ip)
+
+
+@dsl_user_op
+def _ldg4(gaddr, *, loc=None, ip=None):
+ res = _asm(
+ llvm.StructType.get_literal([T.i32()] * 4),
+ [_ir(gaddr, cutlass.Int64, loc, ip)],
+ "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];",
+ "=r,=r,=r,=r,l",
+ loc,
+ ip,
+ side=True,
+ )
+ return _unpack(res, 4, cutlass.Int32, loc, ip)
+
+
+@dsl_user_op
+def _ldsm4(saddr, *, loc=None, ip=None):
+ res = _asm(
+ llvm.StructType.get_literal([T.i32()] * 4),
+ [_ir(saddr, cutlass.Int32, loc, ip)],
+ "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {$0, $1, $2, $3}, [$4];",
+ "=r,=r,=r,=r,r",
+ loc,
+ ip,
+ side=True,
+ )
+ return _unpack(res, 4, cutlass.Int32, loc, ip)
+
+
+@dsl_user_op
+def _ldsm4t(saddr, *, loc=None, ip=None):
+ res = _asm(
+ llvm.StructType.get_literal([T.i32()] * 4),
+ [_ir(saddr, cutlass.Int32, loc, ip)],
+ "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {$0, $1, $2, $3}, [$4];",
+ "=r,=r,=r,=r,r",
+ loc,
+ ip,
+ side=True,
+ )
+ return _unpack(res, 4, cutlass.Int32, loc, ip)
+
+
+@dsl_user_op
+def _mma(a, b, c, *, loc=None, ip=None):
+ """D = A B + C on m16n8k16 (f16 A/B, f32 C/D); a: 4 regs, b: 2 regs, c: 4 f32."""
+ res = _asm(
+ llvm.StructType.get_literal([T.f32()] * 4),
+ [_ir(v, cutlass.Int32, loc, ip) for v in a]
+ + [_ir(v, cutlass.Int32, loc, ip) for v in b]
+ + [_ir(v, cutlass.Float32, loc, ip) for v in c],
+ "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+ "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};",
+ "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
+ loc,
+ ip,
+ )
+ return _unpack(res, 4, cutlass.Float32, loc, ip)
+
+
+@dsl_user_op
+def _dq_word(word, scale2, *, loc=None, ip=None):
+ """Four e2m1 bytes (8 values, low nibble the even one) times an f16x2 scale -> 4 f16x2."""
+ res = _asm(
+ llvm.StructType.get_literal([T.i32()] * 4),
+ [_ir(word, cutlass.Int32, loc, ip), _ir(scale2, cutlass.Int32, loc, ip)],
+ "{ .reg .b8 b0, b1, b2, b3; .reg .b32 h0, h1, h2, h3; "
+ "mov.b32 {b0, b1, b2, b3}, $4; "
+ "cvt.rn.f16x2.e2m1x2 h0, b0; cvt.rn.f16x2.e2m1x2 h1, b1; "
+ "cvt.rn.f16x2.e2m1x2 h2, b2; cvt.rn.f16x2.e2m1x2 h3, b3; "
+ "mul.rn.f16x2 $0, h0, $5; mul.rn.f16x2 $1, h1, $5; "
+ "mul.rn.f16x2 $2, h2, $5; mul.rn.f16x2 $3, h3, $5; }",
+ "=r,=r,=r,=r,r,r",
+ loc,
+ ip,
+ )
+ return _unpack(res, 4, cutlass.Int32, loc, ip)
+
+
+@dsl_user_op
+def _e4m3x2(half, *, loc=None, ip=None):
+ """Two e4m3 bytes (the low 16 bits) -> f16x2, the low byte in the low half."""
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [_ir(half, cutlass.Int32, loc, ip)],
+ "{ .reg .b16 h; cvt.u16.u32 h, $1; cvt.rn.f16x2.e4m3x2 $0, h; }",
+ "=r,r",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _lds_u16(saddr, *, loc=None, ip=None):
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [_ir(saddr, cutlass.Int32, loc, ip)],
+ "ld.shared.u16 $0, [$1];",
+ "=r,r",
+ loc,
+ ip,
+ side=True,
+ )
+ )
+
+
+@dsl_user_op
+def _stg4(gaddr, a, b, c, d, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [_ir(gaddr, cutlass.Int64, loc, ip)]
+ + [_ir(v, cutlass.Int32, loc, ip) 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=True,
+ )
+
+
+@dsl_user_op
+def _scale2s(word, sel, *, loc=None, ip=None):
+ """The e4m3 byte `word.byte[sel & 3]` (sel = k * 0x11) -> f16x2 {s, s}: one prmt + one cvt."""
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [_ir(word, cutlass.Int32, loc, ip), _ir(sel, cutlass.Int32, loc, ip)],
+ "{ .reg .b32 t; .reg .b16 h; prmt.b32 t, $1, 0, $2; cvt.u16.u32 h, t; "
+ "cvt.rn.f16x2.e4m3x2 $0, h; }",
+ "=r,r,r",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _ex2(x, *, loc=None, ip=None):
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [_ir(x, cutlass.Float32, loc, ip)],
+ "ex2.approx.ftz.f32 $0, $1;",
+ "=f,f",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _bf2_to_f2(w, *, loc=None, ip=None):
+ """bf16x2 -> f16x2 (through f32, round to nearest)."""
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [_ir(w, cutlass.Int32, loc, ip)],
+ "{ .reg .b16 lo, hi; .reg .f32 fl, fh; mov.b32 {lo, hi}, $1; "
+ "cvt.f32.bf16 fl, lo; cvt.f32.bf16 fh, hi; cvt.rn.f16x2.f32 $0, fh, fl; }",
+ "=r,r",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _f2_pack(lo, hi, *, loc=None, ip=None):
+ """Two f32 -> f16x2 (lo in the low half)."""
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [_ir(lo, cutlass.Float32, loc, ip), _ir(hi, cutlass.Float32, loc, ip)],
+ "cvt.rn.f16x2.f32 $0, $2, $1;",
+ "=r,f,f",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _bf2_pack(lo, hi, *, loc=None, ip=None):
+ """Two f32 -> bf16x2 (lo in the low half)."""
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [_ir(lo, cutlass.Float32, loc, ip), _ir(hi, cutlass.Float32, loc, ip)],
+ "cvt.rn.bf16x2.f32 $0, $2, $1;",
+ "=r,f,f",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _stg1(gaddr, a, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [_ir(gaddr, cutlass.Int64, loc, ip), _ir(a, cutlass.Int32, loc, ip)],
+ "{ st.global.b32 [$1], $2; mov.u32 $0, 0; }",
+ "=r,l,r",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _stg2f(gaddr, a, b, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [
+ _ir(gaddr, cutlass.Int64, loc, ip),
+ _ir(a, cutlass.Float32, loc, ip),
+ _ir(b, cutlass.Float32, loc, ip),
+ ],
+ "{ st.global.v2.f32 [$1], {$2, $3}; mov.u32 $0, 0; }",
+ "=r,l,f,f",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _ldg1(gaddr, *, loc=None, ip=None):
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [_ir(gaddr, cutlass.Int64, loc, ip)],
+ "ld.global.b32 $0, [$1];",
+ "=r,l",
+ loc,
+ ip,
+ side=True,
+ )
+ )
+
+
+@dsl_user_op
+def _bf2_lo(w, *, loc=None, ip=None):
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [_ir(w, cutlass.Int32, loc, ip)],
+ "{ .reg .b16 lo, hi; mov.b32 {lo, hi}, $1; cvt.f32.bf16 $0, lo; }",
+ "=f,r",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _bf2_hi(w, *, loc=None, ip=None):
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [_ir(w, cutlass.Int32, loc, ip)],
+ "{ .reg .b16 lo, hi; mov.b32 {lo, hi}, $1; cvt.f32.bf16 $0, hi; }",
+ "=f,r",
+ loc,
+ ip,
+ )
+ )
+
+
+# ---------------------------------------------------------------------------------------
+# split kernel
+# ---------------------------------------------------------------------------------------
+
+
+class _SplitKernel:
+ def __init__(
+ self,
+ *,
+ group,
+ kv_heads,
+ width,
+ tpb,
+ splits,
+ tps,
+ scale,
+ num_pages,
+ bt_width,
+ num_requests,
+ k_strides,
+ v_strides,
+ ks_strides,
+ vs_strides,
+ q_strides,
+ req64,
+ gated,
+ pdl,
+ ):
+ assert group <= 16 and HEAD_DIM == 256 and tpb % 4 == 0
+ self.group, self.kvh, self.width, self.tpb = group, kv_heads, width, tpb
+ self.splits, self.tps = splits, tps
+ self.n_tiles = -(-width // TILE)
+ self.scale = scale * LOG2E
+ self.num_pages, self.bt_width, self.num_requests = num_pages, bt_width, num_requests
+ self.ks, self.vs, self.kss, self.vss = k_strides, v_strides, ks_strides, vs_strides
+ self.qs = q_strides
+ self.req64, self.gated, self.pdl = req64, gated, pdl
+ self.heads = group * kv_heads
+ # shared memory carve-out (bytes)
+ self.o_tok = 0
+ self.tok_bytes = -(-tps * TILE * 4 // 128) * 128
+ self.raw_ldv = RAW_LDV
+ self.raw_vsld = 64
+ self.raw_k = TILE * RAW_LD
+ self.raw_v = TILE * self.raw_ldv
+ self.raw_ks = TILE * 16
+ self.raw_vs = TILE * self.raw_vsld
+ self.raw_buf = self.raw_k + self.raw_v + self.raw_ks + self.raw_vs
+ self.stages = min(STAGES, tps)
+ self.o_raw = self.o_tok + self.tok_bytes
+ self.o_sv = self.o_raw + self.stages * self.raw_buf
+ self.o_sp = self.o_sv + TILE * SV_LD
+ self.o_red = self.o_sp + 16 * SP_LD
+ self.smem = self.o_red + 2 * WARPS * 16 * 4 + 128
+ assert k_strides == v_strides and ks_strides == vs_strides, (
+ "K and V planes share one layout"
+ )
+ assert tpb & (tpb - 1) == 0
+
+ @cute.jit
+ def __call__(
+ self,
+ sel_t: cute.Tensor,
+ q_ptr: cute.Pointer,
+ k_ptr: cute.Pointer,
+ v_ptr: cute.Pointer,
+ ks_ptr: cute.Pointer,
+ vs_ptr: cute.Pointer,
+ gs_ptr: cute.Pointer,
+ bt_ptr: cute.Pointer,
+ req_ptr: cute.Pointer,
+ out_ptr: cute.Pointer,
+ lse_ptr: cute.Pointer,
+ gate_ptr: cute.Pointer,
+ rows: cutlass.Int32,
+ gate_stride: cutlass.Int32,
+ stream,
+ ):
+ sel = cute.make_tensor(
+ sel_t.iterator, cute.make_layout((rows, self.width), stride=(self.width, 1))
+ )
+ bt = cute.make_tensor(
+ bt_ptr, cute.make_layout((self.num_requests, self.bt_width), stride=(self.bt_width, 1))
+ )
+ req = cute.make_tensor(req_ptr, cute.make_layout((rows,), stride=(1,)))
+ gs = cute.make_tensor(gs_ptr, cute.make_layout((3,), stride=(1,)))
+ self.kernel(
+ sel,
+ q_ptr.toint(),
+ k_ptr.toint(),
+ v_ptr.toint(),
+ ks_ptr.toint(),
+ vs_ptr.toint(),
+ gs,
+ bt,
+ req,
+ out_ptr.toint(),
+ lse_ptr.toint(),
+ gate_ptr.toint(),
+ rows,
+ gate_stride,
+ ).launch(
+ grid=[rows * self.kvh * self.splits, 1, 1],
+ block=[THREADS, 1, 1],
+ smem=self.smem,
+ stream=stream,
+ use_pdl=self.pdl,
+ min_blocks_per_mp=2,
+ )
+
+ @cute.jit
+ def _issue(self, j, buf, tidx, kvh, s_tok, s_raw, k_base, v_base, ks_base, vs_base):
+ """cp.async of tile `j`'s raw K, K-scale, V and V-scale-block rows into ring slot `buf`,
+ committed as one group: thread t serves token t // 4 (a quarter of its rows); a token
+ that does not resolve gets zeros (src-size 0)."""
+ tpb: cutlass.Constexpr = self.tpb
+ lt: cutlass.Constexpr = tpb.bit_length() - 1
+ ks0, ks1, ks2 = self.ks
+ kss0, kss1 = self.kss
+ rb = s_raw + buf * self.raw_buf
+ rk = rb
+ rv = rb + self.raw_k
+ rks = rv + self.raw_v
+ rvs = rks + self.raw_ks
+ tt = tidx // 4
+ q4 = tidx % 4
+ ent = _lds1(s_tok + (j * TILE + tt) * 4)
+ ok = ent >= 0
+ e = cutlass.max(ent, 0)
+ pg = cutlass.Int64(e >> lt)
+ sl = e & (tpb - 1)
+ n = cutlass.Int32(cutlass.select_(ok, cutlass.Int32(16), cutlass.Int32(0)))
+ off = pg * ks0 + cutlass.Int64(kvh * ks1 + sl * ks2 + 32 * q4)
+ offs = pg * kss0 + cutlass.Int64(kvh * kss1)
+ _cp16z(rk + tt * RAW_LD + 32 * q4, k_base + off, n)
+ _cp16z(rk + tt * RAW_LD + 32 * q4 + 16, k_base + off + 16, n)
+ _cp16z(rv + tt * self.raw_ldv + 32 * q4, v_base + off, n)
+ _cp16z(rv + tt * self.raw_ldv + 32 * q4 + 16, v_base + off + 16, n)
+ _cp16z(
+ rvs + tt * self.raw_vsld + 16 * q4,
+ vs_base + offs + cutlass.Int64((sl >> 2) * 64 + 16 * q4),
+ n,
+ )
+ if q4 == 0:
+ _cp16z(rks + tt * 16, ks_base + offs + cutlass.Int64(sl * 16), n)
+ cute.arch.cp_async_commit_group()
+
+ @cute.kernel
+ def kernel(
+ self,
+ mSel: cute.Tensor,
+ q_base: cutlass.Int64,
+ k_base: cutlass.Int64,
+ v_base: cutlass.Int64,
+ ks_base: cutlass.Int64,
+ vs_base: cutlass.Int64,
+ mGs: cute.Tensor,
+ mBT: cute.Tensor,
+ mReq: cute.Tensor,
+ out_base: cutlass.Int64,
+ lse_base: cutlass.Int64,
+ gate_base: cutlass.Int64,
+ rows: cutlass.Int32,
+ gate_stride: cutlass.Int32,
+ ):
+ tidx, _, _ = cute.arch.thread_idx()
+ bidx, _, _ = cute.arch.block_idx()
+ warp = cute.arch.warp_idx()
+ lane = cute.arch.lane_idx()
+ g = lane // 4
+ c = lane % 4
+ group: cutlass.Constexpr = self.group
+ kvh_n: cutlass.Constexpr = self.kvh
+ splits: cutlass.Constexpr = self.splits
+ tps: cutlass.Constexpr = self.tps
+ n_tiles_all: cutlass.Constexpr = self.n_tiles
+ width: cutlass.Constexpr = self.width
+ tpb: cutlass.Constexpr = self.tpb
+ heads: cutlass.Constexpr = self.heads
+ row = bidx // (kvh_n * splits)
+ kvh = (bidx // splits) % kvh_n
+ split = bidx % splits
+ t0 = split * tps
+ n_t = cutlass.min(tps, n_tiles_all - t0)
+
+ smem = cutlass.utils.SmemAllocator()
+ base_t = smem.allocate_tensor(
+ cutlass.Int8, cute.make_layout((self.smem - 128,), stride=(1,)), byte_alignment=128
+ )
+ sbase = cutlass.Int32(base_t.iterator.toint())
+ s_tok = sbase + self.o_tok
+ s_raw = sbase + self.o_raw
+ s_v = sbase + self.o_sv
+ s_p = sbase + self.o_sp
+ s_red = sbase + self.o_red
+
+ if const_expr(self.pdl):
+ cute.arch.griddepcontrol_launch_dependents()
+ cute.arch.griddepcontrol_wait()
+
+ # ---- the token table of this split: page * tpb + slot, or -1 ---------------------
+ rq = cutlass.Int64(mReq[row])
+ req_ok = (rq >= 0) & (rq < self.num_requests)
+ safe_req = cutlass.Int32(
+ cutlass.min(cutlass.max(rq, cutlass.Int64(0)), cutlass.Int64(self.num_requests - 1))
+ )
+ for i in cutlass.range_constexpr(-(-(tps * TILE) // THREADS)):
+ cl = tidx + i * THREADS
+ col = t0 * TILE + cl
+ 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 < self.bt_width)
+ page = mBT[safe_req, cutlass.min(lp, self.bt_width - 1)]
+ ok = ok & (page >= 0) & (page < self.num_pages)
+ ent = cutlass.Int32(cutlass.select_(ok, page * tpb + safe_tok % tpb, cutlass.Int32(-1)))
+ if cl < tps * TILE:
+ _sts1(s_tok + cl * 4, ent)
+ cute.arch.sync_threads()
+
+ # ---- q of this KV head: rows g and g + 8 of the m16 tile, as permuted A fragments --
+ # thread (g, c) holds dims 32m + 8c .. 32m + 8c + 7 of heads g and g + 8, m = 0..7; a
+ # row past the group (g + 8 always below 16 heads, g itself below 8) holds zeros and
+ # stores nothing
+ qa = cute.make_rmem_tensor(cute.make_layout((8, 4), stride=(4, 1)), cutlass.Int32)
+ qb = cute.make_rmem_tensor(cute.make_layout((8, 4), stride=(4, 1)), cutlass.Int32)
+ qs0, qs1 = self.qs
+ q_row = q_base + (cutlass.Int64(row) * qs0 + cutlass.Int64(kvh * group) * qs1) * 2
+ for m in cutlass.range_constexpr(8):
+ off = cutlass.Int64((32 * m + 8 * c) * 2)
+ hac = cutlass.min(g, group - 1)
+ w4 = _ldg4(q_row + cutlass.Int64(hac) * (qs1 * 2) + off)
+ for e in cutlass.range_constexpr(4):
+ qa[m, e] = cutlass.Int32(
+ cutlass.select_(g < group, _bf2_to_f2(w4[e]), cutlass.Int32(0))
+ )
+ hb = g + 8
+ hbc = cutlass.min(hb, group - 1)
+ w4b = _ldg4(q_row + cutlass.Int64(hbc) * (qs1 * 2) + off)
+ for e in cutlass.range_constexpr(4):
+ qb[m, e] = cutlass.Int32(
+ cutlass.select_(hb < group, _bf2_to_f2(w4b[e]), cutlass.Int32(0))
+ )
+
+ stages: cutlass.Constexpr = self.stages
+ self._issue(0, 0, tidx, kvh, s_tok, s_raw, k_base, v_base, ks_base, vs_base)
+ if const_expr(stages > 2):
+ if n_t > 1:
+ self._issue(1, 1, tidx, kvh, s_tok, s_raw, k_base, v_base, ks_base, vs_base)
+
+ neg_inf = cutlass.Float32(-math.inf)
+ st = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), cutlass.Float32)
+ st[0] = neg_inf # running max, head g
+ st[1] = neg_inf # running max, head g + 8
+ st[2] = cutlass.Float32(0.0) # this thread's share of the running sum, head g
+ st[3] = cutlass.Float32(0.0) # ... head g + 8 (reduced over the CTA at the end)
+ acc = cute.make_rmem_tensor(cute.make_layout((4, 4), stride=(4, 1)), cutlass.Float32)
+ acc.fill(0.0)
+ score_scale = cutlass.Float32(self.scale) * mGs[1]
+
+ # loop-invariant per-thread offsets
+ vt = tidx // 4
+ vq = tidx % 4
+ v_src = vt * RAW_LDV + 32 * vq
+ vs_src = vt * 64 + 16 * vq
+ v_dst = s_v + vt * SV_LD + 128 * vq
+ sel_v = (vt % 4) * 0x11
+ ktok = warp * 8 + g
+ k_src = ktok * RAW_LD + 4 * c
+ ks_src = ktok * 16
+ sel_k0 = (c // 2) * 0x11
+ sel_k1 = (2 + c // 2) * 0x11
+ ta = warp * 8 + 2 * c
+ p_st_g = s_p + g * SP_LD + ta * 2
+ p_st_h = p_st_g + 8 * SP_LD
+ p_ld = s_p + ((lane % 8) + 8 * ((lane // 8) % 2)) * SP_LD + 16 * (lane // 16)
+ v_ld_row = s_v + ((lane % 8) + 8 * ((lane // 8) % 2)) * SV_LD
+ ci0 = 4 * warp + lane // 16
+ ci1 = ci0 + 2
+ v_ld0 = v_ld_row + 16 * ((ci0 & ~7) | ((ci0 & 7) ^ (2 * ((ci0 >> 3) & 3))))
+ v_ld1 = v_ld_row + 16 * ((ci1 & ~7) | ((ci1 & 7) ^ (2 * ((ci1 >> 3) & 3))))
+ red_g = s_red + g * WARPS * 4
+ red_h = s_red + (g + 8) * WARPS * 4
+
+ for j in cutlass.range(n_t):
+ buf = j % stages
+ if const_expr(stages == 1):
+ cute.arch.cp_async_wait_group(0)
+ else:
+ ahead: cutlass.Constexpr = stages - 1
+ if j + ahead < n_t:
+ self._issue(
+ j + ahead,
+ (j + ahead) % stages,
+ tidx,
+ kvh,
+ s_tok,
+ s_raw,
+ k_base,
+ v_base,
+ ks_base,
+ vs_base,
+ )
+ cute.arch.cp_async_wait_group(ahead)
+ else:
+ if const_expr(stages > 2):
+ if j + 1 < n_t:
+ cute.arch.cp_async_wait_group(1)
+ else:
+ cute.arch.cp_async_wait_group(0)
+ else:
+ cute.arch.cp_async_wait_group(0)
+ cute.arch.sync_threads() # [A] tile j landed; previous tile's PV done
+ rb = s_raw + buf * self.raw_buf
+ rk = rb
+ rv = rb + self.raw_k
+ rks = rv + self.raw_v
+ rvs = rks + self.raw_ks
+
+ # -- V tile -> f16 (thread: token tidx // 4, quarter tidx % 4 = dims 64q .. 64q+63)
+ vw0 = _lds4(rv + v_src)
+ vw1 = _lds4(rv + v_src + 16)
+ vsw = _lds4(rvs + vs_src)
+ for gg in cutlass.range_constexpr(4):
+ sc = _scale2s(vsw[gg], sel_v)
+ wa = vw0[2 * gg] if gg < 2 else vw1[2 * gg - 4]
+ wb = vw0[2 * gg + 1] if gg < 2 else vw1[2 * gg - 3]
+ da = _dq_word(wa, sc)
+ db = _dq_word(wb, sc)
+ # logical 16-byte chunk 8 vq + k of the row sits at 8 vq + (k ^ 2 vq): the four
+ # threads of a token store to different banks; ldmatrix applies the same map
+ _sts4(v_dst + 16 * ((2 * gg) ^ (2 * vq)), da[0], da[1], da[2], da[3])
+ _sts4(v_dst + 16 * ((2 * gg + 1) ^ (2 * vq)), db[0], db[1], db[2], db[3])
+
+ # -- S = q K^T for this warp's eight tokens (two accumulator chains)
+ ksw = _lds4(rks + ks_src)
+ kws = [_lds1(rk + k_src + 16 * m) for m in range(8)]
+ zero = cutlass.Float32(0.0)
+ sa = [zero, zero, zero, zero]
+ sb = [zero, zero, zero, zero]
+ for m in cutlass.range_constexpr(8):
+ # scale block 2m + c // 2: word m // 2, byte 2 (m % 2) + c // 2
+ sc = _scale2s(ksw[m // 2], sel_k0 if m % 2 == 0 else sel_k1)
+ b = _dq_word(kws[m], sc)
+ if const_expr(m % 2 == 0):
+ sa = _mma((qa[m, 0], qb[m, 0], qa[m, 1], qb[m, 1]), (b[0], b[1]), sa)
+ sa = _mma((qa[m, 2], qb[m, 2], qa[m, 3], qb[m, 3]), (b[2], b[3]), sa)
+ else:
+ sb = _mma((qa[m, 0], qb[m, 0], qa[m, 1], qb[m, 1]), (b[0], b[1]), sb)
+ sb = _mma((qa[m, 2], qb[m, 2], qa[m, 3], qb[m, 3]), (b[2], b[3]), sb)
+
+ # mask + scale; tokens 8w + 2c, 8w + 2c + 1
+ ea = _lds1(s_tok + (j * TILE + ta) * 4)
+ eb = _lds1(s_tok + (j * TILE + ta + 1) * 4)
+ s0 = cutlass.Float32(cutlass.select_(ea >= 0, (sa[0] + sb[0]) * score_scale, neg_inf))
+ s1 = cutlass.Float32(cutlass.select_(eb >= 0, (sa[1] + sb[1]) * score_scale, neg_inf))
+ s2 = cutlass.Float32(cutlass.select_(ea >= 0, (sa[2] + sb[2]) * score_scale, neg_inf))
+ s3 = cutlass.Float32(cutlass.select_(eb >= 0, (sa[3] + sb[3]) * score_scale, neg_inf))
+ mg = cute.arch.fmax(s0, s1)
+ mh = cute.arch.fmax(s2, s3)
+ mg = cute.arch.fmax(mg, cute.arch.shuffle_sync_bfly(mg, offset=1))
+ mh = cute.arch.fmax(mh, cute.arch.shuffle_sync_bfly(mh, offset=1))
+ mg = cute.arch.fmax(mg, cute.arch.shuffle_sync_bfly(mg, offset=2))
+ mh = cute.arch.fmax(mh, cute.arch.shuffle_sync_bfly(mh, offset=2))
+ if c == 0:
+ _sts1(red_g + warp * 4, _f32_bits(mg))
+ _sts1(red_h + warp * 4, _f32_bits(mh))
+ cute.arch.sync_threads() # [B] maxes; V tile in shared memory
+ rg0 = _lds4(red_g)
+ rg1 = _lds4(red_g + 16)
+ rh0 = _lds4(red_h)
+ rh1 = _lds4(red_h + 16)
+ tmg = st[0]
+ tmh = st[1]
+ for e in cutlass.range_constexpr(4):
+ tmg = cute.arch.fmax(tmg, cute.arch.fmax(_bits_f32(rg0[e]), _bits_f32(rg1[e])))
+ tmh = cute.arch.fmax(tmh, cute.arch.fmax(_bits_f32(rh0[e]), _bits_f32(rh1[e])))
+ safe_g = cutlass.Float32(cutlass.select_(tmg == neg_inf, cutlass.Float32(0.0), tmg))
+ safe_h = cutlass.Float32(cutlass.select_(tmh == neg_inf, cutlass.Float32(0.0), tmh))
+ corr_g = _ex2(st[0] - safe_g)
+ corr_h = _ex2(st[1] - safe_h)
+ p0 = _ex2(s0 - safe_g)
+ p1 = _ex2(s1 - safe_g)
+ p2 = _ex2(s2 - safe_h)
+ p3 = _ex2(s3 - safe_h)
+ _sts1(p_st_g, _f2_pack(p0, p1))
+ _sts1(p_st_h, _f2_pack(p2, p3))
+ st[0] = tmg
+ st[1] = tmh
+ st[2] = st[2] * corr_g + (p0 + p1)
+ st[3] = st[3] * corr_h + (p2 + p3)
+ cute.arch.sync_threads() # [C] P
+
+ # -- O += P V: this warp's dims 32w .. 32w + 31 (n-tiles of 8)
+ if cute.arch.vote_any_sync((corr_g != 1.0) | (corr_h != 1.0)):
+ for nt in cutlass.range_constexpr(4):
+ acc[nt, 0] = acc[nt, 0] * corr_g
+ acc[nt, 1] = acc[nt, 1] * corr_g
+ acc[nt, 2] = acc[nt, 2] * corr_h
+ acc[nt, 3] = acc[nt, 3] * corr_h
+ for kk in cutlass.range_constexpr(4):
+ a = _ldsm4(p_ld + 32 * kk)
+ b0 = _ldsm4t(v_ld0 + 16 * kk * SV_LD)
+ b1 = _ldsm4t(v_ld1 + 16 * kk * SV_LD)
+ r0 = _mma(a, (b0[0], b0[1]), (acc[0, 0], acc[0, 1], acc[0, 2], acc[0, 3]))
+ r1 = _mma(a, (b0[2], b0[3]), (acc[1, 0], acc[1, 1], acc[1, 2], acc[1, 3]))
+ r2 = _mma(a, (b1[0], b1[1]), (acc[2, 0], acc[2, 1], acc[2, 2], acc[2, 3]))
+ r3 = _mma(a, (b1[2], b1[3]), (acc[3, 0], acc[3, 1], acc[3, 2], acc[3, 3]))
+ for e in cutlass.range_constexpr(4):
+ acc[0, e] = r0[e]
+ acc[1, e] = r1[e]
+ acc[2, e] = r2[e]
+ acc[3, e] = r3[e]
+
+ # ---- the running sums over the quad and the warps -----------------------------------
+ lsg = st[2]
+ lsh = st[3]
+ lsg = lsg + cute.arch.shuffle_sync_bfly(lsg, offset=1)
+ lsh = lsh + cute.arch.shuffle_sync_bfly(lsh, offset=1)
+ lsg = lsg + cute.arch.shuffle_sync_bfly(lsg, offset=2)
+ lsh = lsh + cute.arch.shuffle_sync_bfly(lsh, offset=2)
+ cute.arch.sync_threads()
+ if c == 0:
+ _sts1(red_g + warp * 4, _f32_bits(lsg))
+ _sts1(red_h + warp * 4, _f32_bits(lsh))
+ cute.arch.sync_threads()
+ ug0 = _lds4(red_g)
+ ug1 = _lds4(red_g + 16)
+ uh0 = _lds4(red_h)
+ uh1 = _lds4(red_h + 16)
+ tsg = cutlass.Float32(0.0)
+ tsh = cutlass.Float32(0.0)
+ for e in cutlass.range_constexpr(4):
+ tsg = tsg + _bits_f32(ug0[e])
+ tsh = tsh + _bits_f32(uh0[e])
+ for e in cutlass.range_constexpr(4):
+ tsg = tsg + _bits_f32(ug1[e])
+ tsh = tsh + _bits_f32(uh1[e])
+ st[2] = tsg
+ st[3] = tsh
+
+ # ---- epilogue -----------------------------------------------------------------------
+ gv = mGs[2]
+ lg = st[2]
+ lh = st[3]
+ ig = cutlass.Float32(0.0)
+ ih = cutlass.Float32(0.0)
+ if lg > 0.0:
+ ig = gv / lg
+ if lh > 0.0:
+ ih = gv / lh
+ hg = kvh * group + g
+ hh = kvh * group + g + 8
+ okg = g < group
+ okh = (g + 8) < group
+ for nt in cutlass.range_constexpr(4):
+ d = 32 * warp + 8 * nt + 2 * c
+ og0 = acc[nt, 0] * ig
+ og1 = acc[nt, 1] * ig
+ oh0 = acc[nt, 2] * ih
+ oh1 = acc[nt, 3] * ih
+ if const_expr(splits == 1):
+ ag = out_base + (cutlass.Int64(row) * heads + hg) * (HEAD_DIM * 2) + d * 2
+ ah = out_base + (cutlass.Int64(row) * heads + hh) * (HEAD_DIM * 2) + d * 2
+ if okg:
+ if const_expr(self.gated):
+ gw = _ldg1(
+ gate_base + (cutlass.Int64(row) * gate_stride + hg * HEAD_DIM + d) * 2
+ )
+ og0 = _round_bf16(og0) * _sigmoid(_bf2_lo(gw))
+ og1 = _round_bf16(og1) * _sigmoid(_bf2_hi(gw))
+ _stg1(ag, _bf2_pack(og0, og1))
+ if okh:
+ if const_expr(self.gated):
+ gw2 = _ldg1(
+ gate_base + (cutlass.Int64(row) * gate_stride + hh * HEAD_DIM + d) * 2
+ )
+ oh0 = _round_bf16(oh0) * _sigmoid(_bf2_lo(gw2))
+ oh1 = _round_bf16(oh1) * _sigmoid(_bf2_hi(gw2))
+ _stg1(ah, _bf2_pack(oh0, oh1))
+ else:
+ pg_ = (
+ out_base
+ + ((cutlass.Int64(split) * rows + row) * heads + hg) * (HEAD_DIM * 4)
+ + d * 4
+ )
+ ph_ = (
+ out_base
+ + ((cutlass.Int64(split) * rows + row) * heads + hh) * (HEAD_DIM * 4)
+ + d * 4
+ )
+ if okg:
+ _stg2f(pg_, og0, og1)
+ if okh:
+ _stg2f(ph_, oh0, oh1)
+ if const_expr(splits > 1):
+ if warp == 0:
+ if c == 0:
+ lse_g = neg_inf
+ if lg > 0.0:
+ lse_g = st[0] + cute.math.log2(lg)
+ if okg:
+ _stg1(
+ lse_base + ((cutlass.Int64(split) * rows + row) * heads + hg) * 4,
+ _f32_bits(lse_g),
+ )
+ if okh:
+ lse_h = neg_inf
+ if lh > 0.0:
+ lse_h = st[1] + cute.math.log2(lh)
+ _stg1(
+ lse_base + ((cutlass.Int64(split) * rows + row) * heads + hh) * 4,
+ _f32_bits(lse_h),
+ )
+
+
+@dsl_user_op
+def _f32_bits(x, *, loc=None, ip=None):
+ return cutlass.Int32(
+ _asm(T.i32(), [_ir(x, cutlass.Float32, loc, ip)], "mov.b32 $0, $1;", "=r,f", loc, ip)
+ )
+
+
+@dsl_user_op
+def _bits_f32(x, *, loc=None, ip=None):
+ return cutlass.Float32(
+ _asm(T.f32(), [_ir(x, cutlass.Int32, loc, ip)], "mov.b32 $0, $1;", "=f,r", loc, ip)
+ )
+
+
+@dsl_user_op
+def _round_bf16(x, *, loc=None, ip=None):
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [_ir(x, cutlass.Float32, loc, ip)],
+ "{ .reg .b16 h; cvt.rn.bf16.f32 h, $1; cvt.f32.bf16 $0, h; }",
+ "=f,f",
+ loc,
+ ip,
+ )
+ )
+
+
+def _sigmoid(x):
+ """1 / (1 + exp(-x)) as the bf16 merge computes it."""
+ return 1.0 / (1.0 + cute.math.exp(-cutlass.Float32(x)))
+
+
+# ---------------------------------------------------------------------------------------
+# merge kernel (splits > 1)
+# ---------------------------------------------------------------------------------------
+
+
+class _MergeKernel:
+ def __init__(self, *, heads, splits, gated, pdl):
+ assert HEAD_DIM % CHUNK == 0 and 1 < splits <= MAX_SPLITS
+ self.heads, self.splits, self.gated, self.pdl = heads, splits, gated, pdl
+ self.chunks = HEAD_DIM // CHUNK
+ # at least four warps: the final pass needs a thread per element of the chunk
+ self.warps = max(4, min(8, splits))
+ self.spw = -(-splits // self.warps)
+
+ @cute.jit
+ def __call__(
+ self,
+ gate_t: cute.Tensor,
+ part_ptr: cute.Pointer,
+ lse_ptr: cute.Pointer,
+ out_ptr: cute.Pointer,
+ rows: cutlass.Int32,
+ gate_stride: cutlass.Int32,
+ stream,
+ ):
+ lse = cute.make_tensor(
+ lse_ptr,
+ cute.make_layout(
+ (self.splits, rows, self.heads), stride=(rows * self.heads, self.heads, 1)
+ ),
+ )
+ out = cute.make_tensor(
+ out_ptr,
+ cute.make_layout(
+ (rows, self.heads, HEAD_DIM), stride=(self.heads * HEAD_DIM, HEAD_DIM, 1)
+ ),
+ )
+ self.kernel(gate_t.iterator.toint(), part_ptr.toint(), lse, out, rows, gate_stride).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,
+ gate_base: cutlass.Int64,
+ part_base: cutlass.Int64,
+ mLse,
+ mOut,
+ rows: cutlass.Int32,
+ gate_stride: cutlass.Int32,
+ ):
+ 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,
+ )
+ gval = cutlass.Float32(0.0)
+ if const_expr(self.gated):
+ if tidx < CHUNK:
+ gw = _ldg1(
+ gate_base
+ + (cutlass.Int64(row) * gate_stride + head * HEAD_DIM + c0 + (tidx & ~1)) * 2
+ )
+ gval = cutlass.Float32(cutlass.select_((tidx & 1) == 0, _bf2_lo(gw), _bf2_hi(gw)))
+ if const_expr(self.pdl):
+ cute.arch.griddepcontrol_launch_dependents()
+ cute.arch.griddepcontrol_wait()
+ 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]
+ p = cute.make_rmem_tensor(cute.make_layout((spw, 4), stride=(4, 1)), cutlass.Float32)
+ for jj in cutlass.range_constexpr(spw):
+ sidx = cutlass.min(warp_i * spw + jj, self.splits - 1)
+ w4 = _ldg4(
+ part_base
+ + ((cutlass.Int64(sidx) * rows + row) * self.heads + head) * (HEAD_DIM * 4)
+ + (c0 + 4 * lane) * 4
+ )
+ for e in cutlass.range_constexpr(4):
+ p[jj, e] = _bits_f32(w4[e])
+ 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 jj in cutlass.range_constexpr(spw):
+ sidx = warp_i * spw + jj
+ 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[jj, e] + acc[e]
+ for e in cutlass.range_constexpr(4):
+ red[warp_i, 4 * lane + e] = acc[e]
+ cute.arch.sync_threads()
+ 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(-gval)))
+ mOut[row, head, c0 + tidx] = merged.to(cutlass.BFloat16)
+
+
+# ---------------------------------------------------------------------------------------
+# launch
+# ---------------------------------------------------------------------------------------
+
+_COMPILED: dict = {}
+
+
+#: Splits by (rows x KV heads): about 250-290 CTAs where the merge is cheap, a single split
+#: from 128 rows up.
+_SPLIT_TABLE = ((4, 33), (8, 11), (16, 17), (32, 9), (64, 4), (128, 2))
+
+
+def plan_splits(rows: int, kv_heads: int, width: int) -> tuple[int, int]:
+ """(splits, tiles per split) for this many (row, KV head) pairs."""
+ n_tiles = -(-width // TILE)
+ pairs = rows * kv_heads
+ want = 1
+ for bound, sp in _SPLIT_TABLE:
+ if pairs <= bound:
+ want = sp
+ break
+ want = max(1, min(n_tiles, want, MAX_SPLITS))
+ tps = -(-n_tiles // want)
+ return -(-n_tiles // tps), tps
+
+
+def _ptr(t: torch.Tensor, dtype, align=16):
+ return make_ptr(dtype, t.data_ptr(), cute.AddressSpace.gmem, assumed_align=align)
+
+
+def nvfp4_sparse_gqa_decode(
+ 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,
+ output_gate: Optional[torch.Tensor] = None,
+ pdl: bool = True,
+ splits: Optional[int] = None,
+) -> torch.Tensor:
+ """Decode-path NVFP4 sparse GQA. Returns `[rows, Hq, D]` bf16 (contiguous); with
+ `output_gate` (`[rows, Hq * D]` bf16, unit column stride) the result is
+ `bf16(attn) * sigmoid(gate)` rounded to bf16, as `qsa_paged_sparse_gqa_decode` fuses it."""
+ rows, num_q_heads, head_dim = q.shape
+ num_pages, num_kv_heads, tpb, packed = k_cache.shape
+ assert head_dim == HEAD_DIM and packed == HEAD_DIM // 2 and tpb == tokens_per_block
+ assert q.dtype == torch.bfloat16 and q.stride(2) == 1
+ assert k_cache.dtype == torch.uint8 and v_cache.shape == k_cache.shape
+ assert k_cache.stride(3) == 1 and v_cache.stride(3) == 1
+ assert k_scale_cache.stride(2) == HEAD_DIM // 16 and v_scale_cache.stride(2) == HEAD_DIM // 16
+ assert block_table.dtype == torch.int32 and block_table.stride(1) == 1
+ assert block_table.stride(0) == block_table.shape[1], "row-contiguous block table"
+ assert selected_tokens.dtype == torch.int32 and selected_tokens.is_contiguous()
+ assert request_indices.is_contiguous() and request_indices.dtype in (torch.int32, torch.int64)
+ if selected_tokens.data_ptr() % 16:
+ # The kernel reads the selection with 16-byte loads. A row slice of a 2051-wide
+ # selection (a mixed step's generation rows after its context rows) starts 16-byte
+ # aligned only at every fourth row, so an unaligned one is copied (rows x 8 KB).
+ selected_tokens = selected_tokens.clone()
+ group = num_q_heads // num_kv_heads
+ width = selected_tokens.shape[1]
+ out = torch.empty((rows, num_q_heads, head_dim), dtype=torch.bfloat16, device=q.device)
+ if rows == 0:
+ return out
+ if splits is None:
+ n_splits, tps = plan_splits(rows, num_kv_heads, width)
+ else:
+ n_tiles = -(-width // TILE)
+ tps = -(-n_tiles // splits)
+ n_splits = -(-n_tiles // tps)
+ use_pdl = bool(pdl and get_env_enable_pdl())
+ gated = output_gate is not None
+ if gated:
+ assert output_gate.shape == (rows, num_q_heads * head_dim) and output_gate.stride(1) == 1
+ assert output_gate.dtype == torch.bfloat16
+ req64 = request_indices.dtype == torch.int64
+ key = (
+ group,
+ num_kv_heads,
+ width,
+ tpb,
+ n_splits,
+ tps,
+ float(softmax_scale),
+ num_pages,
+ tuple(block_table.shape),
+ tuple(k_cache.stride()[:3]),
+ tuple(v_cache.stride()[:3]),
+ tuple(k_scale_cache.stride()[:2]),
+ tuple(v_scale_cache.stride()[:2]),
+ tuple(q.stride()[:2]),
+ req64,
+ gated and n_splits == 1,
+ use_pdl,
+ )
+ if n_splits > 1:
+ part = torch.empty(
+ (n_splits, rows, num_q_heads, head_dim), dtype=torch.float32, device=q.device
+ )
+ lse = torch.empty((n_splits, rows, num_q_heads), dtype=torch.float32, device=q.device)
+ o_t, l_t = part, lse
+ else:
+ o_t, l_t = out, out
+ gate_t = output_gate if gated else q
+ gate_stride = output_gate.stride(0) if gated else 0
+ fn = _COMPILED.get(key)
+ if fn is None:
+ kern = _SplitKernel(
+ group=group,
+ kv_heads=num_kv_heads,
+ width=width,
+ tpb=tpb,
+ splits=n_splits,
+ tps=tps,
+ scale=float(softmax_scale),
+ num_pages=num_pages,
+ bt_width=block_table.shape[1],
+ num_requests=block_table.shape[0],
+ k_strides=tuple(k_cache.stride()[:3]),
+ v_strides=tuple(v_cache.stride()[:3]),
+ ks_strides=tuple(k_scale_cache.stride()[:2]),
+ vs_strides=tuple(v_scale_cache.stride()[:2]),
+ q_strides=tuple(q.stride()[:2]),
+ req64=req64,
+ gated=gated and n_splits == 1,
+ pdl=use_pdl,
+ )
+ fn = cute.compile(
+ kern,
+ from_dlpack(selected_tokens, assumed_align=16).mark_layout_dynamic(leading_dim=1),
+ _ptr(q, cutlass.BFloat16),
+ _ptr(k_cache, cutlass.Uint8),
+ _ptr(v_cache, cutlass.Uint8),
+ _ptr(k_scale_cache, cutlass.Uint8),
+ _ptr(v_scale_cache, cutlass.Uint8),
+ _ptr(kv_scale_quant_orig, cutlass.Float32, 4),
+ _ptr(block_table, cutlass.Int32, 4),
+ _ptr(request_indices, cutlass.Int64 if req64 else cutlass.Int32, 8 if req64 else 4),
+ _ptr(o_t, cutlass.Float32),
+ _ptr(l_t, cutlass.Float32),
+ _ptr(gate_t, cutlass.BFloat16),
+ cutlass.Int32(rows),
+ cutlass.Int32(gate_stride),
+ cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+ options="--opt-level 2 --enable-tvm-ffi",
+ )
+ _COMPILED[key] = fn
+ fn(
+ selected_tokens,
+ q.data_ptr(),
+ k_cache.data_ptr(),
+ v_cache.data_ptr(),
+ k_scale_cache.data_ptr(),
+ v_scale_cache.data_ptr(),
+ kv_scale_quant_orig.data_ptr(),
+ block_table.data_ptr(),
+ request_indices.data_ptr(),
+ o_t.data_ptr(),
+ l_t.data_ptr(),
+ gate_t.data_ptr(),
+ rows,
+ gate_stride,
+ )
+ if n_splits == 1:
+ return out
+ mkey = ("merge", num_q_heads, n_splits, gated, use_pdl)
+ mfn = _COMPILED.get(mkey)
+ if mfn is None:
+ mfn = cute.compile(
+ _MergeKernel(heads=num_q_heads, splits=n_splits, gated=gated, pdl=use_pdl),
+ from_dlpack(gate_t, assumed_align=16).mark_layout_dynamic(leading_dim=gate_t.dim() - 1),
+ _ptr(part, cutlass.Float32),
+ _ptr(lse, cutlass.Float32, 4),
+ _ptr(out, cutlass.BFloat16),
+ cutlass.Int32(rows),
+ cutlass.Int32(gate_stride),
+ cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+ options="--opt-level 2 --enable-tvm-ffi",
+ )
+ _COMPILED[mkey] = mfn
+ mfn(gate_t, part.data_ptr(), lse.data_ptr(), out.data_ptr(), rows, gate_stride)
+ return out
+
+
+def _decode(segments, **overrides) -> dict:
+ """A cell spec: `(first position, rows)` per request in row order, this checkpoint's
+ geometry (24 q heads over 2 KV heads of 256, 32 tokens a page, 2051 selected tokens)."""
+ return dict(segments=segments, **overrides)
+
+
+#: One tp4 / tep4 rank's heads: 6 q heads over 1 KV head, fewer q heads a KV head than the
+#: m16 tile's eight first rows.
+_TP4_HEADS = dict(q_heads=6, kv_heads=1)
+
+
+#: One dep4 rank's pool under `kv_nvfp4` (attention DP, so tp1's 24/2 x 256 geometry), as the
+#: stock engine allocates it (`max_batch_size` 128, `max_seq_len` 9216,
+#: `free_gpu_memory_fraction` 0.6, chunked prefill): a K/V view of 211797 pages of 32 tokens a
+#: rank (the manager's 213295 blocks less the 1498 recurrent-state slots its second pool group
+#: holds), the engine's `qsa_block_table` of `max_num_sequences` 128 rows by 292 columns
+#: (`ceil((9216 + 1) / 32)` rounded up to a multiple of 4), the trunk layout's 24 role pages a
+#: slot with layer 11's planes the last of them. The requests' pages are scattered over the
+#: whole pool, its last page among them.
+_DEP4_NVFP4 = dict(slots=128, pages=292, slot_pages=24, layer=11, pool_pages=211797, spread=True)
+
+#: The byte offsets the dep4 cells' pages must straddle in both planes: one page of layer 11's
+#: scale plane advances 24 x 2 x 32 x 16 bytes and of its data plane 8x that, so a 211797-page
+#: pool addresses the scale plane to 5.2 GB and the data plane to 41.6 GB.
+_OFFSET_EDGES = (2**31, 2**32)
+
+
+def _dep4_decode(positions) -> dict:
+ """A dep4 decode step: one row a request at each position, with the output gate, over
+ the pool above; captured, as the engine runs it. About 44 GiB of pool (`large`)."""
+ return _decode(
+ tuple((p, 1) for p in positions),
+ gated=True,
+ capture=True,
+ large=True,
+ offset_edges=_OFFSET_EDGES,
+ **_DEP4_NVFP4,
+ )
+
+
+#: Decode positions of a dep4 step's requests: the last token `max_seq_len` 9216 allows (table
+#: column 287), a first token, the last position whose history fits inside the 2048-token
+#: budget, then a spread across the rest.
+def _dep4_positions(requests: int) -> tuple[int, ...]:
+ edge = (9215, 0, 2047)
+ spread = tuple((100 + 499 * r) % 9216 for r in range(requests - len(edge)))
+ return (edge + spread)[:requests]
+
+
+class QsaPagedSparseGqaNvfp4Decode(OpWrapper):
+ """`qsa_paged_sparse_gqa_nvfp4` for decode rows (no prefill pair): sparse GQA over the
+ explicitly selected tokens of an NVFP4 paged pool, `[rows, Hq, D]` bf16. With
+ `output_gate` (`[rows, Hq * D]` bf16, unit column stride) the result is
+ `bf16(attn) * sigmoid(gate)` rounded once more to bf16, the composition
+ `fused_sigmoid_mul_inplace` makes.
+
+ Fusion boundary: the block-table gather, the NVFP4 dequantize, scale, online softmax,
+ P.V, the split-K merge and (optionally) the output gate. Selection, the K/V append and
+ q's projection, norm and RoPE are not inside.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="decode: 1 row, 33-way split + merge", spec=_decode(((4000, 1),))),
+ Cell(
+ why="MTP3 verify: 4 requests x 4 rows, gate fused into the merge",
+ spec=_decode(tuple((4000 + 37 * i, 4) for i in range(4)), gated=True),
+ ),
+ Cell(
+ why="128 rows (16 x 8), one split: gate fused into the split kernel, trunk layout",
+ spec=_decode(
+ tuple((100 + (499 * g) % 8000, 8) for g in range(16)),
+ gated=True,
+ slot_pages=24,
+ layer=11,
+ pool_pages=2048,
+ ),
+ ),
+ Cell(
+ why="tp4 rank heads (6 q over 1 KV): 1 row, 33-way split + merge",
+ spec=_decode(((4000, 1),), **_TP4_HEADS),
+ ),
+ Cell(
+ why="tp4 rank heads (6 q over 1 KV): 128 rows one a request, 2-way split + gated "
+ "merge, captured",
+ spec=_decode(
+ tuple((100 + (499 * r) % 8000, 1) for r in range(128)),
+ gated=True,
+ capture=True,
+ **_TP4_HEADS,
+ ),
+ ),
+ Cell(
+ why="tp4 rank heads (6 q over 1 KV): 160 rows (20 x 8), one split, gate fused into "
+ "the split kernel",
+ spec=_decode(
+ tuple((100 + (499 * g) % 8000, 8) for g in range(20)), gated=True, **_TP4_HEADS
+ ),
+ ),
+ Cell(
+ why="dep4 kv_nvfp4 decode: 1 row at position 9215, 33-way split, gate fused into "
+ "the merge, [128, 292] table, 211797-page pool, captured",
+ spec=_dep4_decode(_dep4_positions(1)),
+ ),
+ Cell(
+ why="dep4 kv_nvfp4 decode: 32 requests (a fourth of the slots), 4-way split + "
+ "gated merge, 211797-page pool, captured",
+ spec=_dep4_decode(_dep4_positions(32)),
+ ),
+ Cell(
+ why="dep4 kv_nvfp4 decode: 128 requests (max_batch_size, every slot), one split "
+ "with the gate fused, 211797-page pool, captured",
+ spec=_dep4_decode(_dep4_positions(128)),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none -- every input is an argument (as `qsa_paged_sparse_gqa_nvfp4`).
+ Call sites (the target's `modeling.py`): `_sparse_gqa` over an NVFP4 pool for rows without
+ the prefill pair -- the decode step's attention layers (with the output gate), the MTP
+ draft layer and a mixed step's generation rows.
+
+ Launch shapes: splits by (rows x KV heads) from `_SPLIT_TABLE` (33 splits of one 64-token
+ tile at one row, a single split of 33 tiles from 128 rows up), each split followed by the
+ merge. A one-split launch (65 rows and up) runs the kernel without the f16 V tile, bitwise
+ the standard kernel.
+
+ Numerics: dequantized K/V values are exact in f16; scores on f16 MMAs with fp32
+ accumulation; P rounded to f16 (as the Triton op); the second-level scales act on the
+ scores and the normalized output. Within the Triton op's band (4 ulp element, 2 ulp RMS)
+ of the fp32 reference; not bitwise the Triton op. Deterministic; graph-capturable
+ (row count is a runtime value; one compile per split variant).
+
+ The CuTe kernel puts a KV head's q heads in the m16 tile's rows g and g + 8; a group of
+ fewer than eight q heads (tp4 and tep4: 6 over 1 KV head a rank) leaves rows past it
+ zero and unstored. The tp4-heads cells drive 1, 128 and 160 rows at that geometry (33
+ splits, 2 splits with the gated merge, captured, and one split with the gate fused).
+
+ dep4's `kv_nvfp4` engine binds tp1's 24/2 x 256 geometry (attention DP): its decode
+ steps, one row a request with the output gate, up to `max_batch_size` 128 rows a rank
+ (the captured batches 1..128), over the engine's `qsa_block_table` ([128, 292] at
+ `max_seq_len` 9216) and one rank's pool as the stock engine allocates it: a K/V view of
+ 211797 pages of 32 tokens (the manager's 213295 blocks less its 1498 recurrent-state
+ slots), the trunk layout's 24 role pages a slot. The dep4 cells drive 1, 32 and 128 rows
+ there (33 splits, 4 splits, one split with the gate fused into the split kernel), each
+ also captured and replayed at its own compile variant. Against `reference`, in bf16 ulp
+ (2^-7) of a (row, head)'s magnitude, max / relative RMS: 1 row 0.55 / 0.12, 32 rows 1.15
+ / 0.03, 128 rows 1.21 / 0.05; QSA's bf16 kernel over the same values lies within 1.00
+ per element of each. Each dep4 cell holds about 44 GiB of pool (41.6 GB of data planes,
+ 5.2 GB of scales), past a normal cell's 40 GiB: `large`.
+
+ Compiled variants: CuTe DSL, JIT-compiled on the first call of each key for the device
+ that runs it; every cell ran that path on sm_103 (GB300). The key holds the pool's page
+ count, the table's shape, the planes' strides, the split plan and whether the gate is
+ fused, so the dep4 cells compile the variants a dep4 engine at those settings runs; any
+ other page count or table shape compiles a variant of its own.
+
+ Large pools: a page index is widened to 64 bits before its product with a plane's page
+ stride. A 211797-page pool addresses layer 11's scale planes to 5.2 GB and its data
+ planes to 41.6 GB; the dep4 cells scatter their requests' pages over the whole pool, its
+ last page among them, so both planes are read on both sides of 2^31 and 2^32 bytes (the
+ test asserts it), with every scale outside them NaN. Hard limit, read from the source and
+ not probed: the token table's entry `page * tokens_per_block + slot` is 32-bit, so a pool
+ of 2^26 pages or more at 32 tokens a page wraps it (1.6 TB of scale planes at this
+ geometry, beyond any device).
+
+ Not certified: the target engine's own dep4 pool (the cells take the stock engine's page
+ count); dep4 with MTP (the draft layer's two more role pages a slot, and the verify
+ entry); positions past `max_seq_len` 9216 (a wider table is another compiled variant);
+ second-level scales other than 1.0 (the target binds ones); int64 request indices.
+
+ Preconditions: bf16 q with head_dim 256; K and V planes (and their scale planes) share
+ strides; power-of-two tokens per block; int32 row-contiguous block table; contiguous
+ int32 selections.
+ """
+
+ def __call__(
+ self,
+ q,
+ k_cache,
+ v_cache,
+ k_scale_cache,
+ v_scale_cache,
+ kv_scale_quant_orig,
+ block_table,
+ selected_tokens,
+ request_indices,
+ tokens_per_block: int,
+ softmax_scale: float,
+ output_gate=None,
+ ) -> torch.Tensor:
+ return nvfp4_sparse_gqa_decode(
+ q,
+ k_cache,
+ v_cache,
+ k_scale_cache,
+ v_scale_cache,
+ kv_scale_quant_orig,
+ block_table,
+ selected_tokens,
+ request_indices,
+ tokens_per_block,
+ softmax_scale,
+ output_gate=output_gate,
+ )
+
+ def reference(
+ self,
+ q,
+ k_cache,
+ v_cache,
+ k_scale_cache,
+ v_scale_cache,
+ kv_scale_quant_orig,
+ block_table,
+ selected_tokens,
+ request_indices,
+ tokens_per_block: int,
+ softmax_scale: float,
+ output_gate=None,
+ ) -> torch.Tensor:
+ """`qsa_paged_sparse_gqa_nvfp4`'s fp32 reference, then the gate in fp32 over its
+ bf16 rounding, rounded once."""
+ out = _composition.reference(
+ q,
+ k_cache,
+ v_cache,
+ k_scale_cache,
+ v_scale_cache,
+ kv_scale_quant_orig,
+ block_table,
+ selected_tokens,
+ request_indices,
+ tokens_per_block,
+ softmax_scale,
+ )
+ if output_gate is None:
+ return out
+ rows, heads, dim = out.shape
+ gate = torch.sigmoid(output_gate.float().view(rows, heads, dim))
+ return (out.to(torch.bfloat16).float() * gate).to(torch.bfloat16)
+
+ def is_valid(
+ self,
+ q,
+ k_cache,
+ v_cache,
+ k_scale_cache,
+ v_scale_cache,
+ kv_scale_quant_orig,
+ block_table,
+ selected_tokens,
+ request_indices,
+ tokens_per_block: int,
+ softmax_scale: float,
+ output_gate=None,
+ ) -> None:
+ assert q.dtype == torch.bfloat16 and q.dim() == 3 and q.shape[2] == HEAD_DIM
+ assert q.stride(2) == 1, "q unit last stride"
+ assert k_cache.dtype == torch.uint8 and v_cache.shape == k_cache.shape
+ assert k_cache.stride() == v_cache.stride(), "K and V planes share strides"
+ assert k_scale_cache.stride() == v_scale_cache.stride(), "scale planes share strides"
+ assert tokens_per_block & (tokens_per_block - 1) == 0, "power-of-two tokens per block"
+ assert block_table.dtype == torch.int32 and block_table.is_contiguous()
+ assert selected_tokens.dtype == torch.int32 and selected_tokens.is_contiguous()
+ assert output_gate is None or (
+ output_gate.dtype == torch.bfloat16
+ and output_gate.shape == (q.shape[0], q.shape[1] * q.shape[2])
+ and output_gate.stride(1) == 1
+ )
+
+
+qsa_paged_sparse_gqa_nvfp4_decode = QsaPagedSparseGqaNvfp4Decode()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4_verify.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4_verify.py
new file mode 100644
index 000000000000..4b4aec3eff8f
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4_verify.py
@@ -0,0 +1,2147 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""QSA sparse GQA over an NVFP4 paged K/V pool for MTP verify rows, request-grouped: CuTe DSL
+(tcgen05) split-K, then a row merge of the split partials.
+
+Tier 5 (CuTe DSL written for this target). A verify step's rows come `T`
+(`rows_per_request`) consecutive rows a request, and the rows of one request select mostly
+the same tokens. One CTA serves one (group of T rows, KV head, split): 128 MMA rows, row
+m = t * 12 + h (verify row t, q head h of the KV head; rows past T * 12 idle), over the
+group's *union* selection, so each selected K/V token is loaded and dequantized once for the
+group's rows instead of once a row.
+
+Union (every thread, one code copy per pass): per-row token bitmaps in shared memory (the
+request's block-table row fetched in the same round trip), the union of four-token groups in
+ascending order (an expanded selection keeps column % 4 == token % 4), each group's K/V row
+offset and scale offset resolved; per tile a 64-bit selection word for each verify row. A
+(row, position) pair the row did not select is masked to -inf before its max, so the softmax
+offset is the row's own max and causality is the selection's. Splits take interleaved tiles of
+the union. Rows of a group that belong to different requests (a partial or misaligned group)
+are served in separate passes; rows of an invalid request give zeros.
+
+Launch (`plan_launch`): within one wave every (group, KV head) unit takes the same number of
+splits; past one wave the units of the full waves take one CTA each (bf16 output, the gate
+fused) and the units of the last, partial wave -- whole groups -- take two to four CTAs each,
+whose rows alone go through the merge.
+
+Roles (11 warps):
+
+* warps 5-6 `cp.async` each tile's raw NVFP4 K / V rows (eight threads a 128-byte row), K scales
+ and V-scale groups into a ring of `RAW_SLOTS` raw slots, completing on the slot's mbarrier;
+* warps 7-10 (one thread a position, K then V) load a raw row and its scales into registers,
+ release the raw slot, and dequantize into f16 SW128 K and V slots (`cvt.rn.f16x2.e2m1x2` x
+ `cvt.rn.f16x2.e4m3x2`, exact, as the per-row kernel);
+* warp 4 issues the tcgen05 MMAs with precomputed descriptors (`qsa_nvfp4_verify_ptx`):
+ S[j % 2] = Q K(j)^T (Q from tensor memory), then O += P(j - 1) V(j - 1) (P from tensor
+ memory, over the S it came from);
+* warps 0-3 (one thread an MMA row) copy their q row into a raw slot at once and into tensor
+ memory (f16) after the first resolve; per tile load S, take the max over the row's own
+ entries, move the row's reference max only when it grew past `RESCALE_TH` (log2), rescaling
+ O in tensor memory then, and write P = 2^(s - max) as f16 (0 off the row's selection);
+ finally normalize O (times the V global scale) and store it as bf16 (or an fp32 partial and
+ its log2-sum-exp for the merge). With the output gate the four warps read the CTA's gate
+ rows (coalesced) before the last P V completes, the live warps stage bf16(O / l) in the free
+ K slots, and all four apply the gate and store coalesced rows.
+
+Tensor memory: O in columns 0-255, S/P buffers at 256 and 320, Q at 384. Shared memory: two f16
+K and two f16 V slots (32 KB each), three raw slots, the union tables (up to ~225 KB): one CTA
+an SM.
+"""
+
+from __future__ import annotations
+
+import functools
+from typing import NamedTuple, Optional
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.cute.nvgpu.tcgen05 as tcgen05
+import cutlass.utils.blackwell_helpers as sm100_utils
+import torch
+from cutlass import const_expr
+from cutlass.cute.nvgpu import OperandMajorMode
+from cutlass.cute.runtime import from_dlpack
+from cutlass.cute.typing import Float32, Int32, Int64
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4_decode import ( # noqa: E501
+ MAX_SPLITS,
+ _bf2_hi,
+ _bf2_lo,
+ _bf2_pack,
+ _bf2_to_f2,
+ _bits_f32,
+ _cp16z,
+ _dq_word,
+ _f32_bits,
+ _ldg1,
+ _ldg4,
+ _lds1,
+ _lds4,
+ _ptr,
+ _round_bf16,
+ _scale2s,
+ _stg1,
+ _sts1,
+ _sts4,
+ nvfp4_sparse_gqa_decode,
+ qsa_paged_sparse_gqa_nvfp4_decode,
+)
+from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+
+from .qsa_nvfp4_verify_ptx import (
+ ctz,
+ gemm_ts,
+ groups8,
+ ldg2,
+ popc,
+ red_or_shared,
+ spread4,
+ stg2,
+ stg4,
+ stg4f,
+ tcgen05_fence_after,
+ tcgen05_fence_before,
+)
+
+try:
+ from cutlass.memory import SmemAllocator
+except ImportError: # DSL releases before the memory namespace split
+ from cutlass.utils import SmemAllocator
+
+F16 = cutlass.Float16
+LOG2E = 1.4426950408889634
+HEAD_DIM = 256
+#: MMA rows a CTA (T verify rows x the KV head's q heads) and positions a tile.
+BM = 128
+BN = 64
+#: Four-token groups a tile.
+GPT = BN // 4
+#: Bytes of a raw K or V row in shared memory (128 + 16 of padding: conflict-free reads).
+RAW_LD = 144
+RAW_SLOTS = 3
+RAW_K = BN * RAW_LD
+RAW_KS = BN * 16
+RAW_VS = GPT * 64
+#: One raw slot: K rows, V rows, K-scale rows, V-scale groups.
+RAW_BUF = 2 * RAW_K + RAW_KS + RAW_VS
+#: Bytes of the four f16 K / V slots (two each, 64 positions x 256 dims).
+KV_SLOTS_BYTES = 4 * BN * HEAD_DIM * 2
+#: Shared memory a CTA may claim (227 KB), less barriers, scan scratch and alignment.
+SMEM_BUDGET = 232448 - 4096
+#: A row's reference max moves only past this many log2 units (P <= 2^8, exact in f16 range).
+RESCALE_TH = 8.0
+SOFT_WARPS = 4
+MMA_WARP = 4
+LOAD_WARP0 = 5
+LOAD_WARPS = 2
+DQ_WARP0 = 7
+DQ_WARPS = 4
+THREADS = 32 * (DQ_WARP0 + DQ_WARPS)
+TMEM_O = 0
+TMEM_S0 = 256
+TMEM_Q = 384
+TMEM_COLS = 512
+CHUNK_COLS = 32
+#: Bytes of a staged q row (bf16, 512 + 16 of padding).
+Q_LD = 528
+#: Most splits of one (group, KV head) union.
+MAX_PLAN_SPLITS = 32
+#: Most splits of a last-wave unit (tail mode).
+MAX_TAIL_SPLITS = 4
+#: Fewer request groups than this take the per-row kernel (GB300: it is as fast or faster
+#: below 32 rows at four rows a request, the grouped kernel from 32 rows up).
+MIN_GROUPS = 8
+
+
+class _Layout(NamedTuple):
+ """The shared-memory plan of one compile: bitmap words (`ngw` groups of four), the union's
+ group capacity, local tiles a CTA may own in one pass, union-scratch offsets in raw slot 0
+ and the table size in int32 words.
+ """
+
+ ngw: int
+ nw: int
+ ug_max: int
+ tmax: int
+ u_ug: int
+ u_bt: int
+ tab_words: int
+
+
+def _layout(rpr: int, group: int, width: int, bt_width: int, tpb: int, splits: int):
+ """The kernel's shared-memory plan for `rpr` rows a group over a `width`-column expanded
+ selection and a `bt_width`-page block table, or None where it does not fit: the union
+ scratch (each row's token bitmap, the union's group list and the request's block-table
+ row) shares raw slot 0, the staged q rows raw slots 1-2.
+ """
+ if tpb % 4 or tpb & (tpb - 1) or rpr < 1 or rpr * group > BM:
+ return None
+ ngw = -(-(bt_width * tpb) // 128)
+ nw = 4 * ngw
+ # an expanded selection names at most ceil(width / 4) four-token groups a row
+ ug_max = min(rpr * -(-width // 4), 32 * ngw)
+ tmax = -(-(-(-ug_max // GPT)) // splits)
+ u_ug = -(-(rpr * nw * 4) // 16) * 16
+ u_bt = -(-(u_ug + ug_max * 4) // 16) * 16
+ tab_words = 3 * tmax * GPT + 2 * tmax * rpr
+ fits = (
+ ngw <= THREADS
+ and u_bt + bt_width * 4 <= RAW_BUF
+ and RAW_BUF + rpr * group * Q_LD <= RAW_SLOTS * RAW_BUF
+ and 4 * tab_words + RAW_SLOTS * RAW_BUF + KV_SLOTS_BYTES <= SMEM_BUDGET
+ )
+ return _Layout(ngw, nw, ug_max, tmax, u_ug, u_bt, tab_words) if fits else None
+
+
+def _sigmoid_fast(x):
+ """1 / (1 + 2^(-x log2 e)) on ex2.approx and rcp.approx (relative error ~1e-6)."""
+ return cute.arch.rcp_approx(1.0 + cute.math.exp2(-x * LOG2E, fastmath=True))
+
+
+def _tree_fmax(vals):
+ """Max of a short list, pairwise (traced as Python)."""
+ while len(vals) > 1:
+ vals = [
+ cute.arch.fmax(vals[j], vals[j + 1]) if j + 1 < len(vals) else vals[j]
+ for j in range(0, len(vals), 2)
+ ]
+ return vals[0]
+
+
+def _tree_sum(vals):
+ """Sum of a short list, pairwise (traced as Python)."""
+ while len(vals) > 1:
+ vals = [
+ vals[j] + vals[j + 1] if j + 1 < len(vals) else vals[j] for j in range(0, len(vals), 2)
+ ]
+ return vals[0]
+
+
+def _phase(g, n):
+ """Parity of use `g` of a ring of `n` slots."""
+ return (g // n) % 2
+
+
+class _VerifyKernel:
+ def __init__(
+ self,
+ *,
+ rpr,
+ group,
+ kv_heads,
+ width,
+ tpb,
+ splits,
+ tail,
+ scale,
+ num_pages,
+ bt_width,
+ num_requests,
+ k_strides,
+ ks_strides,
+ q_strides,
+ req64,
+ gated,
+ pdl,
+ ):
+ lay = _layout(rpr, group, width, bt_width, tpb, 1 if tail else splits)
+ assert lay is not None, "no shared-memory plan for this geometry (see `_layout`)"
+ self.T = rpr
+ self.group, self.kvh, self.width, self.tpb = group, kv_heads, width, tpb
+ self.splits = splits
+ #: tail mode: CTAs below `n_whole` serve whole units, the rest `splits`-way splits
+ self.tail = tail
+ self.scale_log2 = scale * LOG2E
+ self.num_pages, self.bt_width, self.num_requests = num_pages, bt_width, num_requests
+ self.ks, self.kss, self.qs = k_strides, ks_strides, q_strides
+ self.req64, self.gated, self.pdl = req64, gated, pdl
+ self.heads = group * kv_heads
+ self.rows_m = rpr * group
+ #: softmax warps holding live MMA rows (the rest only keep the barriers' counts)
+ self.live_warps = -(-self.rows_m // 32)
+ self.ngw, self.nw, self.ug_max, self.tmax = lay.ngw, lay.nw, lay.ug_max, lay.tmax
+ self.raw_k = RAW_K
+ self.raw_v = RAW_K
+ self.raw_ks = RAW_KS
+ self.raw_vs = RAW_VS
+ self.raw_buf = RAW_BUF
+ self.raw_bytes = RAW_SLOTS * RAW_BUF
+ # union scratch (raw slot 0, free between passes): bitmaps, group list, block-table row
+ self.u_bm = 0
+ self.u_ug = lay.u_ug
+ self.u_bt = lay.u_bt
+ # q rows (bf16, 528-byte rows: conflict-free) in raw slots 1-2 until Q is in tensor memory
+ self.u_q = RAW_BUF
+ # tables (int32 words). Per group: K/V row offset and scale offset of its first token
+ # (16-byte units, -1 when the group does not resolve) and its (token, row) selection
+ # bytes; per (tile, row): the 64-bit selection word.
+ self.t_roff = 0
+ self.t_soff = self.tmax * GPT
+ self.t_gmask = 2 * self.tmax * GPT
+ self.t_bits = 3 * self.tmax * GPT
+ self.tab_words = lay.tab_words
+
+ @cute.jit
+ def __call__(
+ self,
+ sel_t: cute.Tensor,
+ q_ptr: cute.Pointer,
+ k_ptr: cute.Pointer,
+ v_ptr: cute.Pointer,
+ ks_ptr: cute.Pointer,
+ vs_ptr: cute.Pointer,
+ gs_ptr: cute.Pointer,
+ bt_ptr: cute.Pointer,
+ req_ptr: cute.Pointer,
+ out_ptr: cute.Pointer,
+ part_ptr: cute.Pointer,
+ lse_ptr: cute.Pointer,
+ gate_ptr: cute.Pointer,
+ rows: cutlass.Int32,
+ gate_stride: cutlass.Int32,
+ n_whole: cutlass.Int32,
+ tail_row0: cutlass.Int32,
+ n_ctas: cutlass.Int32,
+ stream,
+ ):
+ def qk():
+ return sm100_utils.make_trivial_tiled_mma(
+ F16,
+ F16,
+ OperandMajorMode.K,
+ OperandMajorMode.K,
+ Float32,
+ tcgen05.CtaGroup.ONE,
+ (BM, BN),
+ tcgen05.OperandSource.TMEM,
+ )
+
+ def pv():
+ return sm100_utils.make_trivial_tiled_mma(
+ F16,
+ F16,
+ OperandMajorMode.K,
+ OperandMajorMode.MN,
+ Float32,
+ tcgen05.CtaGroup.ONE,
+ (BM, HEAD_DIM),
+ tcgen05.OperandSource.TMEM,
+ )
+
+ qk_mma, qk_mma_acc, pv_mma_acc = qk(), qk(), pv()
+ q_layout = sm100_utils.make_smem_layout_a(qk_mma, (BM, BN, HEAD_DIM), F16, 1)
+ k_layout = sm100_utils.make_smem_layout_b(qk_mma, (BM, BN, HEAD_DIM), F16, 2)
+ v_layout = sm100_utils.make_smem_layout_b(pv_mma_acc, (BM, HEAD_DIM, BN), F16, 2)
+ p_layout = sm100_utils.make_smem_layout_a(pv_mma_acc, (BM, HEAD_DIM, BN), F16, 1)
+
+ @cute.struct
+ class SharedStorage:
+ q_full: cute.struct.MemRange[Int64, 1]
+ raw_full: cute.struct.MemRange[Int64, RAW_SLOTS]
+ raw_empty: cute.struct.MemRange[Int64, RAW_SLOTS]
+ k_full: cute.struct.MemRange[Int64, 2]
+ k_empty: cute.struct.MemRange[Int64, 2]
+ v_full: cute.struct.MemRange[Int64, 2]
+ v_empty: cute.struct.MemRange[Int64, 2]
+ 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
+ scan: cute.struct.MemRange[Int32, 32]
+ tab: cute.struct.Align[cute.struct.MemRange[Int32, self.tab_words], 16]
+ raw: cute.struct.Align[cute.struct.MemRange[cutlass.Int8, self.raw_bytes], 128]
+ sK: cute.struct.Align[cute.struct.MemRange[F16, cute.cosize(k_layout)], 1024]
+ sV: cute.struct.Align[cute.struct.MemRange[F16, cute.cosize(k_layout)], 1024]
+
+ self.shared_storage = SharedStorage
+ bt = cute.make_tensor(
+ bt_ptr, cute.make_layout((self.num_requests, self.bt_width), stride=(self.bt_width, 1))
+ )
+ req = cute.make_tensor(req_ptr, cute.make_layout((rows,), stride=(1,)))
+ gs = cute.make_tensor(gs_ptr, cute.make_layout((3,), stride=(1,)))
+ self.kernel(
+ qk_mma,
+ qk_mma_acc,
+ pv_mma_acc,
+ q_layout,
+ k_layout,
+ v_layout,
+ p_layout,
+ sel_t.iterator.toint(),
+ q_ptr.toint(),
+ k_ptr.toint(),
+ v_ptr.toint(),
+ ks_ptr.toint(),
+ vs_ptr.toint(),
+ gs,
+ bt,
+ req,
+ out_ptr.toint(),
+ part_ptr.toint(),
+ lse_ptr.toint(),
+ gate_ptr.toint(),
+ rows,
+ gate_stride,
+ n_whole,
+ tail_row0,
+ ).launch(
+ grid=[n_ctas, 1, 1],
+ block=[THREADS, 1, 1],
+ stream=stream,
+ min_blocks_per_mp=1,
+ use_pdl=self.pdl,
+ )
+
+ # ---- pieces every role runs (every thread of the CTA, barrier 0) -----------------------
+
+ @cute.jit
+ def _pass_rows(self, rq: cute.Tensor, handled):
+ lead = cutlass.Int32(-1)
+ for t in tuple(range(self.T - 1, -1, -1)):
+ lead = cutlass.Int32(cutlass.select_(((handled >> t) & 1) == 0, rq[t], lead))
+ act = cutlass.Int32(0)
+ for t in cutlass.range_constexpr(self.T):
+ act = act | cutlass.Int32(
+ cutlass.select_(
+ (((handled >> t) & 1) == 0) & (rq[t] == lead),
+ cutlass.Int32(1 << t),
+ cutlass.Int32(0),
+ )
+ )
+ return lead, act
+
+ @cute.jit
+ def _union(
+ self, act, lead, r0, rows, sel_base, tidx, warp, lane, s_bm, s_scan, s_ug, s_bt, mBT
+ ):
+ rpr: cutlass.Constexpr = self.T
+ nw: cutlass.Constexpr = self.nw
+ ngw: cutlass.Constexpr = self.ngw
+ width: cutlass.Constexpr = self.width
+ ntok: cutlass.Constexpr = 32 * self.nw
+ n_q: cutlass.Constexpr = -(-width // 4) # four-column quads a row
+ n_i: cutlass.Constexpr = -(-n_q // THREADS)
+ cute.arch.sync_threads()
+ for i in cutlass.range_constexpr(-(-(rpr * nw) // THREADS)):
+ idx = tidx + i * THREADS
+ if idx < rpr * nw:
+ _sts1(s_bm + idx * 4, cutlass.Int32(0))
+ # thread: columns 4q .. 4q + 3 of each row (one selected group in an expanded selection);
+ # every load first, then one bit set a distinct word
+ toks = []
+ for t in cutlass.range_constexpr(rpr):
+ row_sel = sel_base + cutlass.Int64(cutlass.min(r0 + t, rows - 1)) * (width * 4)
+ for i in cutlass.range_constexpr(n_i):
+ for u in cutlass.range_constexpr(4):
+ colc = cutlass.min(4 * (tidx + i * THREADS) + u, width - 1)
+ toks.append(_ldg1(row_sel + cutlass.Int64(colc * 4)))
+ # the request's block-table row, in the same round trip
+ n_bt: cutlass.Constexpr = -(-self.bt_width // THREADS)
+ btv = [mBT[lead, cutlass.min(tidx + i * THREADS, self.bt_width - 1)] for i in range(n_bt)]
+ cute.arch.sync_threads()
+ for t in cutlass.range_constexpr(rpr):
+ for i in cutlass.range_constexpr(n_i):
+ q4 = tidx + i * THREADS
+ live_t = ((act >> t) & 1) != 0
+ okv = []
+ wds = []
+ bv = []
+ for u in cutlass.range_constexpr(4):
+ tok = toks[(t * n_i + i) * 4 + u]
+ okv.append(live_t & (4 * q4 + u < width) & (tok >= 0) & (tok < ntok))
+ wds.append(tok >> 5)
+ bv.append(
+ cutlass.Int32(
+ cutlass.select_(
+ okv[u], cutlass.Int32(1) << (tok & 31), cutlass.Int32(0)
+ )
+ )
+ )
+ # bits sharing token 0's word go in one atomic; any other token on its own
+ m0 = bv[0]
+ for u in cutlass.range_constexpr(1, 4):
+ m0 = m0 | cutlass.Int32(
+ cutlass.select_(wds[u] == wds[0], bv[u], cutlass.Int32(0))
+ )
+ if okv[0] | (m0 != 0):
+ if m0 != 0:
+ red_or_shared(s_bm + (t * nw + cutlass.max(wds[0], 0)) * 4, m0)
+ for u in cutlass.range_constexpr(1, 4):
+ if okv[u] & (wds[u] != wds[0]):
+ red_or_shared(s_bm + (t * nw + wds[u]) * 4, bv[u])
+ for i in cutlass.range_constexpr(n_bt):
+ if tidx + i * THREADS < self.bt_width:
+ _sts1(s_bt + (tidx + i * THREADS) * 4, btv[i])
+ cute.arch.sync_threads()
+ gmv = cutlass.Int32(0)
+ if tidx < ngw:
+ for k in cutlass.range_constexpr(4):
+ xo = cutlass.Int32(0)
+ for t in cutlass.range_constexpr(rpr):
+ xo = xo | _lds1(s_bm + (t * nw + 4 * tidx + k) * 4)
+ gmv = gmv | (groups8(xo) << (8 * k))
+ cnt = popc(gmv)
+ incl = cnt
+ for o in (1, 2, 4, 8, 16):
+ up = cute.arch.shuffle_sync_up(incl, offset=o)
+ if lane >= o:
+ incl = incl + up
+ if lane == 31:
+ _sts1(s_scan + warp * 4, incl)
+ cute.arch.sync_threads()
+ wsum = cutlass.Int32(0)
+ tot = cutlass.Int32(0)
+ if lane < THREADS // 32:
+ wv = _lds1(s_scan + lane * 4)
+ tot = wv
+ if lane < warp:
+ wsum = wv
+ before = cute.arch.warp_reduction_sum(wsum)
+ n_ug = cute.arch.warp_reduction_sum(tot)
+ excl = before + incl - cnt
+ if tidx < ngw:
+ xm = gmv
+ for j in cutlass.range(cnt):
+ b = ctz(xm)
+ if excl + j < self.ug_max:
+ _sts1(s_ug + (excl + j) * 4, 32 * tidx + b)
+ xm = xm & (xm - 1)
+ cute.arch.sync_threads()
+ return cutlass.min(n_ug, self.ug_max)
+
+ @cute.jit
+ def _resolve(self, lead, kvh, split, nsplit, n_loc, n_ug, tidx, s_bm, s_ug, s_tab, s_bt):
+ """Every local tile's groups (entry of the group's first token, (token, row) selection
+ bytes), then each (tile, verify row)'s 64-bit selection word.
+ """
+ rpr: cutlass.Constexpr = self.T
+ nw: cutlass.Constexpr = self.nw
+ tmax: cutlass.Constexpr = self.tmax
+ n_rs: cutlass.Constexpr = -(-(tmax * GPT) // THREADS)
+ s_roff = s_tab + self.t_roff * 4
+ s_soff = s_tab + self.t_soff * 4
+ s_gmask = s_tab + self.t_gmask * 4
+ s_bits = s_tab + self.t_bits * 4
+ ks0, ks1, ks2 = self.ks
+ kss0, kss1 = self.kss
+ pages = []
+ gids = []
+ for i in cutlass.range_constexpr(n_rs):
+ slot = tidx + i * THREADS
+ kl = slot // GPT
+ ugi = (split + kl * nsplit) * GPT + slot % GPT
+ gid = _lds1(s_ug + cutlass.min(ugi, self.ug_max - 1) * 4)
+ gid = cutlass.min(cutlass.max(gid, 0), 8 * nw - 1)
+ gids.append(gid)
+ pages.append(_lds1(s_bt + cutlass.min((gid * 4) // self.tpb, self.bt_width - 1) * 4))
+ for i in cutlass.range_constexpr(n_rs):
+ slot = tidx + i * THREADS
+ kl = slot // GPT
+ ugi = (split + kl * nsplit) * GPT + slot % GPT
+ gid = gids[i]
+ page = pages[i]
+ wsel = gid >> 3
+ sh = (gid & 7) * 4
+ mw = cutlass.Int32(0)
+ for t in cutlass.range_constexpr(rpr):
+ nib = (_lds1(s_bm + (t * nw + wsel) * 4) >> sh) & 0xF
+ mw = mw | (spread4(nib) << t)
+ okp = (kl < n_loc) & (ugi < n_ug) & (page >= 0) & (page < self.num_pages)
+ pg64 = cutlass.Int64(cutlass.max(page, 0))
+ sl0 = (gid * 4) % self.tpb
+ roff = (pg64 * ks0 + cutlass.Int64(kvh * ks1 + sl0 * ks2)) >> 4
+ soff = (pg64 * kss0 + cutlass.Int64(kvh * kss1 + sl0 * 16)) >> 4
+ roff = cutlass.Int32(cutlass.select_(okp, cutlass.Int32(roff), cutlass.Int32(-1)))
+ soff = cutlass.Int32(cutlass.select_(okp, cutlass.Int32(soff), cutlass.Int32(-1)))
+ mw = cutlass.Int32(cutlass.select_(okp, mw, cutlass.Int32(0)))
+ if slot < tmax * GPT:
+ _sts1(s_roff + slot * 4, roff)
+ _sts1(s_soff + slot * 4, soff)
+ _sts1(s_gmask + slot * 4, mw)
+ cute.arch.sync_threads()
+ for i in cutlass.range_constexpr(-(-(tmax * rpr) // THREADS)):
+ w = tidx + i * THREADS
+ if w < tmax * rpr:
+ kl = w // rpr
+ t = w % rpr
+ lo = cutlass.Int32(0)
+ hi = cutlass.Int32(0)
+ for gq in cutlass.range_constexpr(GPT):
+ nb = (_lds1(s_gmask + (kl * GPT + gq) * 4) >> t) & 0x01010101
+ nb = (nb | (nb >> 7) | (nb >> 14) | (nb >> 21)) & 0xF
+ if const_expr(gq < 8):
+ lo = lo | (nb << (4 * gq))
+ else:
+ hi = hi | (nb << (4 * (gq - 8)))
+ _sts1(s_bits + w * 8, lo)
+ _sts1(s_bits + w * 8 + 4, hi)
+ cute.arch.sync_threads()
+
+ @cute.jit
+ def _pass_tiles(
+ self,
+ act,
+ lead,
+ kvh,
+ r0,
+ rows,
+ sel_base,
+ split,
+ nsplit,
+ tidx,
+ warp,
+ lane,
+ s_raw,
+ s_scan,
+ s_tab,
+ mBT,
+ ):
+ """Union + resolve of one pass; returns the CTA's local tile count."""
+ s_bm = s_raw + self.u_bm
+ s_ug = s_raw + self.u_ug
+ s_bt = s_raw + self.u_bt
+ n_ug = self._union(
+ act, lead, r0, rows, sel_base, tidx, warp, lane, s_bm, s_scan, s_ug, s_bt, mBT
+ )
+ n_tiles = (n_ug + GPT - 1) // GPT
+ n_loc = cutlass.Int32(
+ cutlass.select_(
+ split < n_tiles,
+ (n_tiles - split + nsplit - 1) // nsplit,
+ cutlass.Int32(0),
+ )
+ )
+ self._resolve(lead, kvh, split, nsplit, n_loc, n_ug, tidx, s_bm, s_ug, s_tab, s_bt)
+ return n_loc
+
+ @cute.kernel
+ def kernel(
+ self,
+ qk_mma: cute.TiledMma,
+ qk_mma_acc: cute.TiledMma,
+ pv_mma_acc: cute.TiledMma,
+ q_layout: cute.ComposedLayout,
+ k_layout: cute.ComposedLayout,
+ v_layout: cute.ComposedLayout,
+ p_layout: cute.ComposedLayout,
+ sel_base: cutlass.Int64,
+ q_base: cutlass.Int64,
+ k_base: cutlass.Int64,
+ v_base: cutlass.Int64,
+ ks_base: cutlass.Int64,
+ vs_base: cutlass.Int64,
+ mGs: cute.Tensor,
+ mBT: cute.Tensor,
+ mReq: cute.Tensor,
+ out_base: cutlass.Int64,
+ part_base: cutlass.Int64,
+ lse_base: cutlass.Int64,
+ gate_base: cutlass.Int64,
+ rows: cutlass.Int32,
+ gate_stride: cutlass.Int32,
+ n_whole: cutlass.Int32,
+ tail_row0: cutlass.Int32,
+ ):
+ warp = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+ tidx, _, _ = cute.arch.thread_idx()
+ bidx, _, _ = cute.arch.block_idx()
+ lane = cute.arch.lane_idx()
+ qk_mma.set(tcgen05.Field.ACCUMULATE, False)
+ qk_mma_acc.set(tcgen05.Field.ACCUMULATE, True)
+ pv_mma_acc.set(tcgen05.Field.ACCUMULATE, True)
+ rpr: cutlass.Constexpr = self.T
+ group: cutlass.Constexpr = self.group
+ kvh_n: cutlass.Constexpr = self.kvh
+ splits: cutlass.Constexpr = self.splits
+ if const_expr(self.tail):
+ # units (group, KV head) below n_whole: one CTA each; the rest: `splits` CTAs each
+ is_whole = bidx < n_whole
+ t_cta = bidx - n_whole
+ unit = cutlass.Int32(cutlass.select_(is_whole, bidx, n_whole + t_cta // splits))
+ split = cutlass.Int32(cutlass.select_(is_whole, cutlass.Int32(0), t_cta % splits))
+ nsplit = cutlass.Int32(
+ cutlass.select_(is_whole, cutlass.Int32(1), cutlass.Int32(splits))
+ )
+ else:
+ unit = bidx // splits
+ split = bidx % splits
+ nsplit = cutlass.Int32(splits)
+ gi = unit // kvh_n
+ kvh = unit % kvh_n
+ r0 = gi * rpr
+
+ smem = SmemAllocator()
+ storage = smem.allocate(self.shared_storage)
+ q_full = storage.q_full.data_ptr()
+ raw_full = storage.raw_full.data_ptr()
+ raw_empty = storage.raw_empty.data_ptr()
+ k_full = storage.k_full.data_ptr()
+ k_empty = storage.k_empty.data_ptr()
+ v_full = storage.v_full.data_ptr()
+ v_empty = storage.v_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()
+ s_scan = cutlass.Int32(storage.scan.data_ptr().toint())
+ s_tab = cutlass.Int32(storage.tab.data_ptr().toint())
+ s_raw = cutlass.Int32(storage.raw.data_ptr().toint())
+ if tidx == 0:
+ cute.arch.mbarrier_init(q_full, 32 * SOFT_WARPS)
+ for i in cutlass.range_constexpr(RAW_SLOTS):
+ cute.arch.mbarrier_init(raw_full + i, 32 * LOAD_WARPS)
+ cute.arch.mbarrier_init(raw_empty + i, 32 * DQ_WARPS)
+ for i in cutlass.range_constexpr(2):
+ cute.arch.mbarrier_init(k_full + i, BN)
+ cute.arch.mbarrier_init(k_empty + i, 1)
+ cute.arch.mbarrier_init(v_full + i, BN)
+ cute.arch.mbarrier_init(v_empty + i, 1)
+ cute.arch.mbarrier_init(s_full + i, 1)
+ cute.arch.mbarrier_init(p_full + i, 32 * SOFT_WARPS)
+ cute.arch.mbarrier_init(o_full, 1)
+ cute.arch.mbarrier_init(o_final, 1)
+ cute.arch.mbarrier_init_fence()
+ if warp == MMA_WARP:
+ cute.arch.alloc_tmem(TMEM_COLS, storage.tmem_holding.ptr)
+ cute.arch.relinquish_tmem_alloc_permit()
+ cute.arch.sync_threads()
+ tmem_ptr = cute.arch.retrieve_tmem_ptr(Float32, 16, storage.tmem_holding.ptr)
+
+ qk_thr = qk_mma.get_slice(0)
+ pv_thr = pv_mma_acc.get_slice(0)
+ sK = storage.sK.get_tensor(k_layout.outer, swizzle=k_layout.inner)
+ sVr = storage.sV.get_tensor(k_layout.outer, swizzle=k_layout.inner)
+ sV = cute.make_tensor(cute.recast_ptr(sVr.iterator, v_layout.inner), v_layout.outer)
+ s_k16 = cutlass.Int32(storage.sK.data_ptr().toint())
+ s_v16 = cutlass.Int32(storage.sV.data_ptr().toint())
+
+ # ---- softmax-side tensor-memory views (warps 0-3: one thread an MMA row) -------------
+ st = tidx
+ tStS = qk_thr.make_fragment_C(qk_thr.partition_shape_C((BM, BN)))
+ tStS0 = cute.make_tensor(tmem_ptr + TMEM_S0, tStS.layout)
+ tOtO_f = pv_thr.make_fragment_C(pv_thr.partition_shape_C((BM, HEAD_DIM)))
+ tOtO = cute.make_tensor(tmem_ptr + TMEM_O, tOtO_f.layout)
+ tScS = qk_thr.partition_C(cute.make_identity_tensor((BM, BN)))
+ load_s = tcgen05.make_tmem_copy(
+ cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), 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 = BN // 2
+ tStP0 = cute.make_tensor(
+ tStS0.iterator, cute.composition(tStS0.layout, cute.make_layout((BM, p_cols)))
+ )
+ tScP = cute.make_tensor(
+ tScS.iterator, cute.composition(tScS.layout, cute.make_layout((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)
+ tOcO = pv_thr.partition_C(cute.make_identity_tensor((BM, HEAD_DIM)))
+ tOtO_c = cute.make_tensor(
+ tOtO.iterator, cute.composition(tOtO.layout, cute.make_layout((BM, CHUNK_COLS)))
+ )
+ tOcO_c = cute.make_tensor(
+ tOcO.iterator, cute.composition(tOcO.layout, cute.make_layout((BM, CHUNK_COLS)))
+ )
+ load_o = tcgen05.make_tmem_copy(
+ cute.make_copy_atom(
+ tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(CHUNK_COLS)), Float32
+ ),
+ tOtO_c,
+ )
+ store_o = tcgen05.make_tmem_copy(
+ cute.make_copy_atom(
+ tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(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)
+ tview = (load_s, tLtS0, tLcS, store_p, tPtP0, tPcP, load_o, store_o, tLtO, tLcO, tStO)
+
+ m = tidx
+ t_row = m // group
+ h = m % group
+ live = (m < self.rows_m) & (r0 + t_row < rows)
+ safe_row = cutlass.min(r0 + t_row, rows - 1)
+ qs0, qs1 = self.qs
+ q_row = q_base + (cutlass.Int64(safe_row) * qs0 + cutlass.Int64(kvh * group + h) * qs1) * 2
+ q_s = s_raw + self.u_q + cutlass.min(m, self.rows_m - 1) * Q_LD
+
+ if const_expr(self.pdl):
+ cute.arch.griddepcontrol_launch_dependents()
+ cute.arch.griddepcontrol_wait()
+
+ if warp < SOFT_WARPS:
+ # O = 0; this row's q copied into raw slots 1-2 now, into tensor memory (f16) after the
+ # first pass's resolve. Dead rows stage zeros.
+ rZ = cute.make_rmem_tensor(tLcO.shape, Float32)
+ rZ.fill(0.0)
+ for c in cutlass.range_constexpr(HEAD_DIM // CHUNK_COLS):
+ cute.copy(
+ store_o, rZ, cute.make_tensor(tStO.iterator + c * CHUNK_COLS, tStO.layout)
+ )
+ cute.arch.fence_view_async_tmem_store()
+ if m < self.rows_m:
+ qn = cutlass.Int32(cutlass.select_(live, cutlass.Int32(16), cutlass.Int32(0)))
+ for i in cutlass.range_constexpr(HEAD_DIM // 8):
+ _cp16z(q_s + 16 * i, q_row + cutlass.Int64(16 * i), qn)
+ cute.arch.cp_async_commit_group()
+
+ rq = cute.make_rmem_tensor(cute.make_layout((rpr,), stride=(1,)), cutlass.Int32)
+ handled = cutlass.Int32(0)
+ for t in cutlass.range_constexpr(rpr):
+ row = r0 + t
+ xr = cutlass.Int64(mReq[cutlass.min(row, rows - 1)])
+ okr = (row < rows) & (xr >= 0) & (xr < self.num_requests)
+ rq[t] = cutlass.Int32(cutlass.select_(okr, cutlass.Int32(xr), cutlass.Int32(-1)))
+ handled = handled | cutlass.Int32(
+ cutlass.select_(okr, cutlass.Int32(0), cutlass.Int32(1 << t))
+ )
+ full = (1 << rpr) - 1
+ scale_log2 = Float32(self.scale_log2) * mGs[1]
+ neg_inf = -Float32.inf
+ row_max = neg_inf
+ row_sum = Float32(0.0)
+ q_pending = cutlass.Int32(1)
+ g = cutlass.Int32(0)
+ for _pass in cutlass.range(rpr):
+ if handled != full:
+ lead, act = self._pass_rows(rq, handled)
+ n_loc = self._pass_tiles(
+ act,
+ lead,
+ kvh,
+ r0,
+ rows,
+ sel_base,
+ split,
+ nsplit,
+ tidx,
+ warp,
+ lane,
+ s_raw,
+ s_scan,
+ s_tab,
+ mBT,
+ )
+ if warp >= DQ_WARP0:
+ self.dequant_tiles(
+ tidx - 32 * DQ_WARP0,
+ g,
+ n_loc,
+ s_raw,
+ s_k16,
+ s_v16,
+ raw_full,
+ raw_empty,
+ k_full,
+ k_empty,
+ v_full,
+ v_empty,
+ )
+ elif warp >= LOAD_WARP0:
+ self.load_tiles(
+ tidx - 32 * LOAD_WARP0,
+ g,
+ n_loc,
+ s_raw,
+ s_tab,
+ k_base,
+ v_base,
+ ks_base,
+ vs_base,
+ raw_full,
+ raw_empty,
+ q_full,
+ )
+ elif warp == MMA_WARP:
+ self.mma_tiles(
+ g,
+ n_loc,
+ qk_mma,
+ pv_mma_acc,
+ qk_thr,
+ pv_thr,
+ q_layout,
+ p_layout,
+ sK,
+ sV,
+ tmem_ptr,
+ q_full,
+ k_full,
+ k_empty,
+ v_full,
+ v_empty,
+ s_full,
+ p_full,
+ o_full,
+ )
+ else:
+ if q_pending != 0:
+ self.stage_q(q_s, live, tview, q_full)
+ q_pending = cutlass.Int32(0)
+ row_max, row_sum = self.softmax_tiles(
+ g,
+ n_loc,
+ row_max,
+ row_sum,
+ live,
+ t_row,
+ warp,
+ s_tab,
+ scale_log2,
+ tview,
+ s_full,
+ p_full,
+ o_full,
+ )
+ g = g + n_loc
+ cute.arch.sync_threads() # [P] end of the pass
+ handled = handled | act
+
+ if warp == MMA_WARP:
+ with cute.arch.elect_one():
+ tcgen05.commit(o_final)
+ if warp < SOFT_WARPS:
+ cute.arch.cp_async_wait_group(0)
+ self.epilogue(
+ row_max,
+ row_sum,
+ live,
+ warp,
+ m,
+ r0,
+ rows,
+ kvh,
+ split,
+ nsplit,
+ tail_row0,
+ mGs,
+ scale_log2,
+ out_base,
+ part_base,
+ lse_base,
+ gate_base,
+ gate_stride,
+ tview,
+ o_final,
+ tidx,
+ s_k16,
+ )
+ # every role done with tensor memory (the softmax warps read O last)
+ cute.arch.sync_threads()
+ if warp == MMA_WARP:
+ cute.arch.dealloc_tmem(tmem_ptr, TMEM_COLS)
+
+ # ------------------------------------------------------------------ loaders (warps 5-6)
+ @cute.jit
+ def load_tiles(
+ self,
+ lt,
+ g,
+ n_loc,
+ s_raw,
+ s_tab,
+ k_base,
+ v_base,
+ ks_base,
+ vs_base,
+ raw_full,
+ raw_empty,
+ q_full,
+ ):
+ """Raw tile copies: K and V rows eight threads a row (16-byte chunk lt % 8 of rows
+ lt // 8 + 8 i: whole 128-byte lines a warp instruction); the K-scale row lt and chunk
+ lt % 4 of the V-scale group lt // 4 (both at the group's scale offset + lt % 4). Raw slots
+ 1-2 hold the q rows until Q is in tensor memory.
+ """
+ s_roff = s_tab + self.t_roff * 4
+ s_soff = s_tab + self.t_soff * 4
+ s_gmask = s_tab + self.t_gmask * 4
+ tok16: cutlass.Constexpr = self.ks[2] // 16
+ q8 = lt % 8
+ for kl in cutlass.range(n_loc):
+ gg = g + kl
+ r = gg % RAW_SLOTS
+ cute.arch.mbarrier_wait(raw_empty + r, _phase(gg, RAW_SLOTS) ^ 1)
+ if gg >= 1:
+ cute.arch.mbarrier_wait(q_full, 0)
+ rb = s_raw + r * self.raw_buf
+ rk = rb
+ rv = rb + self.raw_k
+ rks = rv + self.raw_v
+ rvs = rks + self.raw_ks
+ for i in cutlass.range_constexpr(BN // 8):
+ n = lt // 8 + 8 * i
+ slot = kl * GPT + n // 4
+ roff = _lds1(s_roff + slot * 4)
+ mw = _lds1(s_gmask + slot * 4)
+ ok = (roff >= 0) & (((mw >> (8 * (n % 4))) & 0xFF) != 0)
+ nb = cutlass.Int32(cutlass.select_(ok, cutlass.Int32(16), cutlass.Int32(0)))
+ off = cutlass.Int64(cutlass.max(roff, 0) + (n % 4) * tok16 + q8) * 16
+ _cp16z(rk + n * RAW_LD + 16 * q8, k_base + off, nb)
+ _cp16z(rv + n * RAW_LD + 16 * q8, v_base + off, nb)
+ slot = kl * GPT + lt // 4
+ soff = _lds1(s_soff + slot * 4)
+ mw = _lds1(s_gmask + slot * 4)
+ ok = (soff >= 0) & (((mw >> (8 * (lt % 4))) & 0xFF) != 0)
+ gok = (soff >= 0) & (mw != 0)
+ nb = cutlass.Int32(cutlass.select_(ok, cutlass.Int32(16), cutlass.Int32(0)))
+ ng = cutlass.Int32(cutlass.select_(gok, cutlass.Int32(16), cutlass.Int32(0)))
+ soff64 = cutlass.Int64(cutlass.max(soff, 0) + lt % 4) * 16
+ _cp16z(rks + lt * 16, ks_base + soff64, nb)
+ _cp16z(rvs + (lt // 4) * 64 + 16 * (lt % 4), vs_base + soff64, ng)
+ cute.arch.cp_async_mbarrier_arrive_noinc(raw_full + r)
+
+ # ------------------------------------------------------------------ dequant (warps 7-10)
+ @cute.jit
+ def dequant_tiles(
+ self,
+ dt,
+ g,
+ n_loc,
+ s_raw,
+ s_k16,
+ s_v16,
+ raw_full,
+ raw_empty,
+ k_full,
+ k_empty,
+ v_full,
+ v_empty,
+ ):
+ """Thread `dt` < 64 dequantizes K row `dt` of each tile, the others V row `dt - 64`: the raw
+ row and its scales into registers (the raw slot is then free), then into the f16 SW128
+ slot once the MMA has released it.
+ """
+ is_v = dt >= BN
+ n = dt % BN
+ for kl in cutlass.range(n_loc):
+ gg = g + kl
+ r = gg % RAW_SLOTS
+ b = gg % 2
+ cute.arch.mbarrier_wait(raw_full + r, _phase(gg, RAW_SLOTS))
+ rb = s_raw + r * self.raw_buf
+ row_off = (
+ rb
+ + cutlass.Int32(cutlass.select_(is_v, cutlass.Int32(self.raw_k), cutlass.Int32(0)))
+ + n * RAW_LD
+ )
+ w = [_lds4(row_off + 16 * c) for c in range(8)]
+ sc_off = rb + self.raw_k + self.raw_v
+ sc_off = cutlass.Int32(
+ cutlass.select_(is_v, sc_off + self.raw_ks + (n // 4) * 64, sc_off + n * 16)
+ )
+ sw = [_lds4(sc_off + 16 * i) for i in range(4)]
+ cute.arch.mbarrier_arrive(raw_empty + r)
+ if is_v:
+ cute.arch.mbarrier_wait(v_empty + b, _phase(gg, 2) ^ 1)
+ self._decode_row(w, sw, s_v16 + b * (BN * HEAD_DIM * 2), n, True)
+ cute.arch.fence_view_async_shared()
+ cute.arch.mbarrier_arrive(v_full + b)
+ else:
+ cute.arch.mbarrier_wait(k_empty + b, _phase(gg, 2) ^ 1)
+ self._decode_row(w, sw, s_k16 + b * (BN * HEAD_DIM * 2), n, False)
+ cute.arch.fence_view_async_shared()
+ cute.arch.mbarrier_arrive(k_full + b)
+
+ @cute.jit
+ def _decode_row(self, w, sw, dst, n, is_v: cutlass.Constexpr):
+ """Raw row words `w` (8 x 4) and scale words `sw` -> f16 SW128 row `n` at `dst`: 16-byte
+ chunk ci of the row's 64-dim block ci // 8 at 8192 (ci // 8) + 128 n + 16 ((ci % 8) ^ (n % 8)).
+ """
+ sel_v = (n % 4) * 0x11
+ for c in cutlass.range_constexpr(8):
+ # blocks 2c, 2c + 1. K: bytes 2c, 2c + 1 of the token's 16 (word c // 2);
+ # V: byte n % 4 of the group's words 2c, 2c + 1
+ if const_expr(is_v):
+ sc0 = _scale2s(sw[(2 * c) // 4][(2 * c) % 4], sel_v)
+ sc1 = _scale2s(sw[(2 * c + 1) // 4][(2 * c + 1) % 4], sel_v)
+ else:
+ sc0 = _scale2s(sw[0][c // 2], (2 * (c % 2)) * 0x11)
+ sc1 = _scale2s(sw[0][c // 2], (2 * (c % 2) + 1) * 0x11)
+ for e in cutlass.range_constexpr(4):
+ d4 = _dq_word(w[c][e], sc0 if e < 2 else sc1)
+ ci = 4 * c + e
+ _sts4(
+ dst + 8192 * (ci // 8) + 128 * n + 16 * ((ci % 8) ^ (n % 8)),
+ d4[0],
+ d4[1],
+ d4[2],
+ d4[3],
+ )
+
+ # ------------------------------------------------------------------ MMA (warp 4)
+ @cute.jit
+ def mma_tiles(
+ self,
+ g,
+ n_loc,
+ qk_mma,
+ pv_mma_acc,
+ qk_thr,
+ pv_thr,
+ q_layout,
+ p_layout,
+ sK,
+ sV,
+ tmem_ptr,
+ q_full,
+ k_full,
+ k_empty,
+ v_full,
+ v_empty,
+ s_full,
+ p_full,
+ o_full,
+ ):
+ """QK(g), then QK(j) / PV(j - 1) pairs, then the last PV: S/P buffer j % 2 at TMEM column
+ TMEM_S0 + 64 (j % 2); QK(j) overwrites it only after PV(j - 2) was issued (tcgen05 MMAs of
+ one thread execute in order).
+ """
+ tQ = cute.make_tensor(tmem_ptr + 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 + TMEM_Q, dtype=F16), tSrQ.layout)
+ tSrK = qk_thr.make_fragment_B(sK)
+ tOrV = pv_thr.make_fragment_B(sV)
+ tP = cute.make_tensor(tmem_ptr + TMEM_S0, p_layout.outer)
+ tOrP = pv_thr.make_fragment_A(tP)[None, None, None, 0]
+ tOrP = cute.make_tensor(cute.recast_ptr(tmem_ptr + TMEM_S0, dtype=F16), tOrP.layout)
+ tbase = cutlass.Int32(tmem_ptr.toint())
+ qk_op = qk_mma.op
+ pv_op = pv_mma_acc.op
+ if n_loc > 0:
+ cute.arch.mbarrier_wait(q_full, 0)
+ cute.arch.mbarrier_wait(k_full + g % 2, _phase(g, 2))
+ tcgen05_fence_after()
+ self.issue_qk(qk_op, tbase, tSrQ, tSrK, sK, g % 2)
+ with cute.arch.elect_one():
+ tcgen05.commit(s_full + g % 2)
+ tcgen05.commit(k_empty + g % 2)
+ for j in cutlass.range(1, n_loc + 1, 1, unroll=1):
+ gj = g + j
+ if j < n_loc:
+ cute.arch.mbarrier_wait(k_full + gj % 2, _phase(gj, 2))
+ tcgen05_fence_after()
+ self.issue_qk(qk_op, tbase, tSrQ, tSrK, sK, gj % 2)
+ with cute.arch.elect_one():
+ tcgen05.commit(s_full + gj % 2)
+ tcgen05.commit(k_empty + gj % 2)
+ gp = gj - 1
+ cute.arch.mbarrier_wait(v_full + gp % 2, _phase(gp, 2))
+ cute.arch.mbarrier_wait(p_full + gp % 2, _phase(gp, 2))
+ tcgen05_fence_after()
+ gemm_ts(
+ pv_op,
+ tbase + TMEM_O,
+ tbase + TMEM_S0 + (gp % 2) * BN,
+ tOrP,
+ tOrV[None, None, None, gp % 2],
+ sV[None, None, None, gp % 2],
+ False,
+ )
+ with cute.arch.elect_one():
+ tcgen05.commit(v_empty + gp % 2)
+ tcgen05.commit(o_full)
+
+ @cute.jit
+ def issue_qk(self, qk_op, tbase, tSrQ, tSrK, sK, slot):
+ """S[slot] = Q K(slot)^T: Q from tensor memory, K from SW128 slot `slot`."""
+ gemm_ts(
+ qk_op,
+ tbase + TMEM_S0 + slot * BN,
+ tbase + TMEM_Q,
+ tSrQ,
+ tSrK[None, None, None, slot],
+ sK[None, None, None, slot],
+ True,
+ )
+
+ # ------------------------------------------------------------------ softmax + epilogue (warps 0-3)
+ @cute.jit
+ def stage_q(self, q_s, live, tview, q_full):
+ """This MMA row's q (bf16 in raw slots 1-2) -> f16 into tensor memory (the QK A operand)."""
+ load_s, tLtS0, tLcS, store_p, tPtP0, tPcP, load_o, store_o, tLtO, tLcO, tStO = tview
+ p_cols: cutlass.Constexpr = BN // 2
+ cute.arch.cp_async_wait_group(0)
+ for c in cutlass.range_constexpr(HEAD_DIM // (2 * p_cols)):
+ rQ = cute.make_rmem_tensor(tPcP.shape, Float32)
+ for i in cutlass.range_constexpr(p_cols // 4):
+ w4 = _lds4(q_s + 16 * (c * (p_cols // 4) + i))
+ for e in cutlass.range_constexpr(4):
+ rQ[4 * i + e] = _bits_f32(
+ cutlass.Int32(cutlass.select_(live, _bf2_to_f2(w4[e]), cutlass.Int32(0)))
+ )
+ cute.copy(
+ store_p,
+ rQ,
+ cute.make_tensor(tPtP0.iterator + (TMEM_Q - TMEM_S0 + c * p_cols), tPtP0.layout),
+ )
+ cute.arch.fence_view_async_tmem_store()
+ tcgen05_fence_before()
+ cute.arch.mbarrier_arrive(q_full)
+
+ @cute.jit
+ def softmax_tiles(
+ self,
+ g,
+ n_loc,
+ row_max,
+ row_sum,
+ live,
+ t_row,
+ warp,
+ s_tab,
+ scale_log2,
+ tview,
+ s_full,
+ p_full,
+ o_full,
+ ):
+ """Per tile: S from tensor memory, the max over the row's own entries, the reference max
+ moved (and O rescaled) only past RESCALE_TH, P = 2^(s - max) as f16 over the S.
+ """
+ load_s, tLtS0, tLcS, store_p, tPtP0, tPcP, load_o, store_o, tLtO, tLcO, tStO = tview
+ rpr: cutlass.Constexpr = self.T
+ s_bits = s_tab + self.t_bits * 4
+ tau = Float32(RESCALE_TH)
+ neg_inf = -Float32.inf
+ live_warp = warp < self.live_warps
+ for kl in cutlass.range(n_loc, unroll=1):
+ gg = g + kl
+ buf = gg % 2
+ cute.arch.mbarrier_wait(s_full + buf, _phase(gg, 2))
+ tcgen05_fence_after()
+ if live_warp:
+ lo = cutlass.Int32(0)
+ hi = cutlass.Int32(0)
+ if live:
+ lo = _lds1(s_bits + (kl * rpr + t_row) * 8)
+ hi = _lds1(s_bits + (kl * rpr + t_row) * 8 + 4)
+ rS = cute.make_rmem_tensor(tLcS.shape, Float32)
+ tLtS_b = cute.make_tensor(tLtS0.iterator + buf * BN, tLtS0.layout)
+ for c in cutlass.range_constexpr(cute.size(rS, mode=[1])):
+ cute.copy(load_s, tLtS_b[None, c, 0, 0], rS[None, c, 0, 0])
+ cute.arch.fence_view_async_tmem_load()
+ tile_max = neg_inf
+ for i in cutlass.range_constexpr(BN):
+ word = lo if i < 32 else hi
+ x = rS[i]
+ x = x if (word & (1 << (i % 32))) != 0 else neg_inf
+ rS[i] = x
+ tile_max = cute.arch.fmax(tile_max, x)
+ new_max = cute.arch.fmax(row_max, tile_max)
+ use_max = row_max
+ factor = Float32(1.0)
+ if row_max == neg_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(..gg - 1): wait for it, then rescale this warp's rows
+ cute.arch.mbarrier_wait(o_full, (gg - 1) % 2)
+ tcgen05_fence_after()
+ for c in cutlass.range_constexpr(HEAD_DIM // CHUNK_COLS):
+ rO = cute.make_rmem_tensor(tLcO.shape, Float32)
+ cute.copy(
+ load_o,
+ cute.make_tensor(tLtO.iterator + c * 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 * CHUNK_COLS, tStO.layout),
+ )
+ cute.arch.fence_view_async_tmem_store()
+ row_max = use_max
+ safe = use_max
+ if use_max == neg_inf:
+ safe = Float32(0.0)
+ neg = -safe * scale_log2
+ rP = cute.make_rmem_tensor(tPcP.shape, Float32)
+ rP16 = cute.make_tensor(cute.recast_ptr(rP.iterator, dtype=F16), rS.layout)
+ acc0 = (Float32(0.0), Float32(0.0))
+ acc1 = (Float32(0.0), Float32(0.0))
+ for i in cutlass.range_constexpr(0, BN, 4):
+ 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):
+ rS[i + u] = cute.math.exp2(e[u], fastmath=True)
+ 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(F16))
+ cute.copy(store_p, rP, cute.make_tensor(tPtP0.iterator + buf * BN, tPtP0.layout))
+ cute.arch.fence_view_async_tmem_store()
+ tcgen05_fence_before()
+ cute.arch.mbarrier_arrive(p_full + buf)
+ return row_max, row_sum
+
+ @cute.jit
+ def epilogue(
+ self,
+ row_max,
+ row_sum,
+ live,
+ warp,
+ m,
+ r0,
+ rows,
+ kvh,
+ split,
+ nsplit,
+ tail_row0,
+ mGs,
+ scale_log2,
+ out_base,
+ part_base,
+ lse_base,
+ gate_base,
+ gate_stride,
+ tview,
+ o_final,
+ st,
+ s_stage,
+ ):
+ """O / l (times the V global scale): a whole unit's rows as bf16, a split's as an fp32
+ partial and its log2-sum-exp for the merge (rows from `tail_row0` on).
+
+ Gated whole units: the four softmax warps read the gate (coalesced, 16 bytes a chunk of
+ eight values) before the last P V completes; the live warps stage bf16(O / l) in the
+ free K slots; then all four apply the gate and store coalesced rows."""
+ load_s, tLtS0, tLcS, store_p, tPtP0, tPcP, load_o, store_o, tLtO, tLcO, tStO = tview
+ group: cutlass.Constexpr = self.group
+ heads: cutlass.Constexpr = self.heads
+ splits: cutlass.Constexpr = self.splits
+ tailm: cutlass.Constexpr = self.tail
+ has_whole: cutlass.Constexpr = tailm or splits == 1
+ has_split: cutlass.Constexpr = tailm or splits > 1
+ gated_whole: cutlass.Constexpr = has_whole and self.gated
+ rows_m: cutlass.Constexpr = self.rows_m
+ cpr: cutlass.Constexpr = HEAD_DIM // 8 # 16-byte chunks a row
+ n_gc: cutlass.Constexpr = -(-(rows_m * cpr) // (32 * SOFT_WARPS))
+ n_chunks: cutlass.Constexpr = HEAD_DIM // CHUNK_COLS
+ neg_inf = -Float32.inf
+ row = r0 + m // group
+ h = m % group
+ head = kvh * group + h
+ prow = cutlass.Int64(row - tail_row0)
+ prows = cutlass.Int64(rows - tail_row0)
+ whole = nsplit == 1
+ gregs = []
+ if const_expr(gated_whole):
+ for j in cutlass.range_constexpr(n_gc):
+ k = st + 32 * SOFT_WARPS * j
+ gm = cutlass.min(k // cpr, rows_m - 1)
+ grow = cutlass.min(r0 + gm // group, rows - 1)
+ ghead = kvh * group + gm % group
+ gregs.append(
+ _ldg4(
+ gate_base
+ + (cutlass.Int64(grow) * gate_stride + ghead * HEAD_DIM + (k % cpr) * 8) * 2
+ )
+ )
+ cute.arch.mbarrier_wait(o_final, 0)
+ tcgen05_fence_after()
+ if warp < self.live_warps:
+ gv = mGs[2]
+ inv = Float32(0.0)
+ if row_sum > Float32(0.0):
+ inv = gv / row_sum
+ 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 * CHUNK_COLS, tLtO.layout), rO)
+ for v8 in cutlass.range_constexpr(CHUNK_COLS // 8):
+ vals = [rO[8 * v8 + i] * inv for i in range(8)]
+ d0 = c * CHUNK_COLS + 8 * v8
+ oaddr = out_base + ((cutlass.Int64(row) * heads + head) * HEAD_DIM + d0) * 2
+ paddr = (
+ part_base
+ + (((cutlass.Int64(split) * prows + prow) * heads + head) * HEAD_DIM + d0)
+ * 4
+ )
+ packed = [_bf2_pack(vals[2 * e], vals[2 * e + 1]) for e in range(4)]
+ saddr = s_stage + cutlass.min(m, rows_m - 1) * Q_LD + d0 * 2
+ if const_expr(has_whole):
+ if whole:
+ if const_expr(gated_whole):
+ if m < rows_m:
+ _sts4(saddr, packed[0], packed[1], packed[2], packed[3])
+ else:
+ if live:
+ stg4(oaddr, packed[0], packed[1], packed[2], packed[3])
+ if const_expr(has_split):
+ if live & (nsplit > 1):
+ stg4f(paddr, vals[0], vals[1], vals[2], vals[3])
+ stg4f(paddr + 16, vals[4], vals[5], vals[6], vals[7])
+ if const_expr(has_split):
+ if live & (nsplit > 1):
+ lse = neg_inf
+ if row_sum > Float32(0.0):
+ lse = row_max * scale_log2 + cute.math.log2(row_sum)
+ _stg1(
+ lse_base + ((cutlass.Int64(split) * prows + prow) * heads + head) * 4,
+ _f32_bits(lse),
+ )
+ if const_expr(gated_whole):
+ if whole:
+ cute.arch.barrier(barrier_id=1, number_of_threads=32 * SOFT_WARPS)
+ for j in cutlass.range_constexpr(n_gc):
+ k = st + 32 * SOFT_WARPS * j
+ gm = cutlass.min(k // cpr, rows_m - 1)
+ grow = r0 + gm // group
+ ghead = kvh * group + gm % group
+ a4 = _lds4(s_stage + gm * Q_LD + (k % cpr) * 16)
+ g4 = gregs[j]
+ res = [
+ _bf2_pack(
+ _bf2_lo(a4[e]) * _sigmoid_fast(_bf2_lo(g4[e])),
+ _bf2_hi(a4[e]) * _sigmoid_fast(_bf2_hi(g4[e])),
+ )
+ for e in range(4)
+ ]
+ if (k < rows_m * cpr) & (grow < rows):
+ stg4(
+ out_base
+ + ((cutlass.Int64(grow) * heads + ghead) * HEAD_DIM + (k % cpr) * 8)
+ * 2,
+ res[0],
+ res[1],
+ res[2],
+ res[3],
+ )
+
+
+#: Threads of a merge CTA (16 bytes, four values, a thread at a time).
+MERGE_THREADS = 256
+#: Merge CTAs a row, each a contiguous share of the row's heads x head_dim values (measured on
+#: GB300: 8 is faster than 1-4 at 64-128 rows, 16 no faster).
+MERGE_PARTS = 8
+
+
+class _RowMergeKernel:
+ """The split partials of a row range merged as the decode op's merge kernel does
+ (log2-sum-exp weights, fp32), the gate applied to the bf16-rounded result: `parts` CTAs a
+ row, every load of a thread issued before the first use."""
+
+ def __init__(self, *, heads, splits, gated, pdl, parts):
+ self.heads, self.splits, self.gated, self.pdl = heads, splits, gated, pdl
+ self.parts = parts
+
+ @cute.jit
+ def __call__(
+ self,
+ gate_t: cute.Tensor,
+ part_ptr: cute.Pointer,
+ lse_ptr: cute.Pointer,
+ out_ptr: cute.Pointer,
+ rows: cutlass.Int32,
+ gate_stride: cutlass.Int32,
+ stream,
+ ):
+ self.kernel(
+ part_ptr.toint(),
+ lse_ptr.toint(),
+ out_ptr.toint(),
+ gate_t.iterator.toint(),
+ rows,
+ gate_stride,
+ ).launch(
+ grid=[rows * self.parts, 1, 1],
+ block=[MERGE_THREADS, 1, 1],
+ stream=stream,
+ use_pdl=self.pdl,
+ )
+
+ @cute.kernel
+ def kernel(
+ self,
+ part_base: cutlass.Int64,
+ lse_base: cutlass.Int64,
+ out_base: cutlass.Int64,
+ gate_base: cutlass.Int64,
+ rows: cutlass.Int32,
+ gate_stride: cutlass.Int32,
+ ):
+ tidx, _, _ = cute.arch.thread_idx()
+ bidx, _, _ = cute.arch.block_idx()
+ heads: cutlass.Constexpr = self.heads
+ sp: cutlass.Constexpr = self.splits
+ parts: cutlass.Constexpr = self.parts
+ n_all: cutlass.Constexpr = heads * HEAD_DIM // 4
+ n_c: cutlass.Constexpr = -(-n_all // parts)
+ per: cutlass.Constexpr = -(-n_c // MERGE_THREADS)
+ neg_inf = -Float32.inf
+ row = bidx // parts
+ c0 = (bidx % parts) * n_c
+ r64 = cutlass.Int64(row)
+ cs = [
+ cutlass.min(c0 + cutlass.min(tidx + MERGE_THREADS * j, n_c - 1), n_all - 1)
+ for j in range(per)
+ ]
+ gws = []
+ if const_expr(self.gated):
+ # the gate is an input of the attention call, ready before the split kernel ran
+ gws = [ldg2(gate_base + (r64 * gate_stride + 4 * c) * 2) for c in cs]
+ if const_expr(self.pdl):
+ cute.arch.griddepcontrol_launch_dependents()
+ cute.arch.griddepcontrol_wait()
+ lss = []
+ pss = []
+ for j in cutlass.range_constexpr(per):
+ head = cs[j] // (HEAD_DIM // 4)
+ lss.append(
+ [
+ _bits_f32(
+ _ldg1(lse_base + ((cutlass.Int64(s_) * rows + r64) * heads + head) * 4)
+ )
+ for s_ in range(sp)
+ ]
+ )
+ pss.append(
+ [
+ _ldg4(
+ part_base
+ + ((cutlass.Int64(s_) * rows + r64) * (heads * HEAD_DIM) + 4 * cs[j]) * 4
+ )
+ for s_ in range(sp)
+ ]
+ )
+ for j in cutlass.range_constexpr(per):
+ ls = lss[j]
+ mx = _tree_fmax(ls)
+ ws = [
+ Float32(cutlass.select_(mx > neg_inf, cute.math.exp2(lv - mx), Float32(0.0)))
+ for lv in ls
+ ]
+ den = _tree_sum(ws)
+ res = []
+ for e in cutlass.range_constexpr(4):
+ acc = _tree_sum([ws[s_] * _bits_f32(pss[j][s_][e]) for s_ in range(sp)])
+ merged = Float32(0.0)
+ if den > Float32(0.0):
+ merged = acc / den
+ if const_expr(self.gated):
+ if const_expr(e % 2 == 0):
+ gx = _bf2_lo(gws[j][e // 2])
+ else:
+ gx = _bf2_hi(gws[j][e // 2])
+ merged = _round_bf16(merged) * _sigmoid_fast(gx)
+ res.append(merged)
+ if (tidx + MERGE_THREADS * j < n_c) & (c0 + tidx + MERGE_THREADS * j < n_all):
+ stg2(
+ out_base + (r64 * (heads * HEAD_DIM) + 4 * cs[j]) * 2,
+ _bf2_pack(res[0], res[1]),
+ _bf2_pack(res[2], res[3]),
+ )
+
+
+# ---------------------------------------------------------------------------------------
+# launch
+# ---------------------------------------------------------------------------------------
+
+_COMPILED: dict = {}
+#: Geometries whose every reachable launch plan is compiled (`_compile_plans`).
+_PLANNED: set = set()
+
+
+@functools.lru_cache(maxsize=None)
+def _num_sms(device_index: int) -> int:
+ return torch.cuda.get_device_properties(device_index).multi_processor_count
+
+
+def plan_splits(n_groups: int, kv_heads: int, width: int, num_sms: int) -> int:
+ """Splits so the grid fills the GPU once (one CTA an SM): measured on GB300, a second
+ partial wave costs more than the extra splits save.
+ """
+ est_tiles = -(-int(width * 1.25) // BN)
+ return max(1, min(MAX_PLAN_SPLITS, est_tiles, num_sms // (n_groups * kv_heads)))
+
+
+def plan_launch(n_groups: int, kv_heads: int, width: int, num_sms: int):
+ """(splits, n_whole) of a launch: n_whole None = every unit (group, KV head) takes `splits`
+ CTAs (one wave: `plan_splits`); else units below n_whole take one CTA and the units of the
+ last, partial wave (whole groups) `splits` CTAs each, merged afterwards."""
+ units = n_groups * kv_heads
+ if units <= num_sms:
+ return plan_splits(n_groups, kv_heads, width, num_sms), None
+ waves = -(-units // num_sms)
+ tail = units - (waves - 1) * num_sms
+ tail = -(-tail // kv_heads) * kv_heads
+ tail_splits = min(MAX_TAIL_SPLITS, num_sms // tail)
+ if tail_splits < 2:
+ return 1, None
+ return tail_splits, units - tail
+
+
+def reachable_plans(max_groups: int, kv_heads: int, width: int, num_sms: int) -> list:
+ """Every `(splits, tail)` pair `plan_launch` returns for `MIN_GROUPS..max_groups` groups,
+ `tail` whether the last wave is split: the kernel variants a geometry can need."""
+ plans = set()
+ for n_groups in range(MIN_GROUPS, max(MIN_GROUPS, max_groups) + 1):
+ splits, n_whole = plan_launch(n_groups, kv_heads, width, num_sms)
+ plans.add((max(1, min(MAX_SPLITS, int(splits))), n_whole is not None))
+ return sorted(plans)
+
+
+def _plan(q, k_cache, block_table, selected_tokens, tokens_per_block, rows_per_request):
+ """(splits, layout) of the grouped kernel for this call, or None where the per-row kernel
+ takes it: fewer than `MIN_GROUPS` groups, one row a request, not SM103, or no shared-memory
+ plan (more than six rows a request, or a block table wider than the union scratch holds:
+ 612 pages of 32 tokens at four rows a request).
+ """
+ rows, num_q_heads, head_dim = q.shape
+ rpr = int(rows_per_request)
+ if rpr < 2 or -(-rows // rpr) < MIN_GROUPS or head_dim != HEAD_DIM:
+ return None
+ if torch.cuda.get_device_capability(q.device) != (10, 3):
+ return None
+ kv_heads = k_cache.shape[1]
+ width = selected_tokens.shape[1]
+ splits = plan_splits(-(-rows // rpr), kv_heads, width, _num_sms(q.device.index or 0))
+ layout = _layout(
+ rpr, num_q_heads // kv_heads, width, block_table.shape[1], tokens_per_block, splits
+ )
+ return None if layout is None else (splits, layout)
+
+
+def takes_grouped(q, k_cache, block_table, selected_tokens, tokens_per_block, rows_per_request):
+ """Whether `nvfp4_sparse_gqa_verify` runs this call on the grouped kernel (else on the
+ per-row kernel). A host-side decision on shapes only: a per-capture constant.
+ """
+ plan = _plan(q, k_cache, block_table, selected_tokens, tokens_per_block, rows_per_request)
+ return plan is not None
+
+
+def _compile_verify(key: tuple, kern_kw: dict, compile_args: tuple):
+ """The grouped kernel of `key` (its `splits` and `tail` at positions 5 and 6), compiled
+ once."""
+ fn = _COMPILED.get(key)
+ if fn is None:
+ kern = _VerifyKernel(splits=key[5], tail=key[6], **kern_kw)
+ fn = cute.compile(kern, *compile_args, options="--opt-level 2 --enable-tvm-ffi")
+ _COMPILED[key] = fn
+ return fn
+
+
+def _compile_merge(heads: int, splits: int, gated: bool, pdl: bool, merge_args: tuple):
+ """The row merge of `splits` partials, compiled once."""
+ mkey = ("rowmerge", heads, splits, gated, pdl, MERGE_PARTS)
+ mfn = _COMPILED.get(mkey)
+ if mfn is None:
+ mfn = cute.compile(
+ _RowMergeKernel(heads=heads, splits=splits, gated=gated, pdl=pdl, parts=MERGE_PARTS),
+ *merge_args,
+ options="--opt-level 2 --enable-tvm-ffi",
+ )
+ _COMPILED[mkey] = mfn
+ return mfn
+
+
+def nvfp4_sparse_gqa_verify_grouped(
+ 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,
+ rows_per_request: int,
+ output_gate: Optional[torch.Tensor] = None,
+ pdl: bool = True,
+ splits: Optional[int] = None,
+) -> torch.Tensor:
+ """The grouped kernel (and the merge with more than one split), whatever the row count;
+ `splits` overrides the plan. `nvfp4_sparse_gqa_verify` is the routed entry.
+ """
+ rows, num_q_heads, head_dim = q.shape
+ num_pages, num_kv_heads, tpb, packed = k_cache.shape
+ rpr = int(rows_per_request)
+ assert head_dim == HEAD_DIM and packed == HEAD_DIM // 2 and tpb == tokens_per_block
+ assert q.dtype == torch.bfloat16 and q.stride(2) == 1 and q.stride(1) == HEAD_DIM
+ assert (q.stride(0) * 2) % 16 == 0 and q.data_ptr() % 16 == 0, "16-byte aligned q rows"
+ assert k_cache.dtype == torch.uint8 and k_cache.stride() == v_cache.stride()
+ assert k_cache.stride(3) == 1 and v_cache.shape == k_cache.shape
+ assert k_scale_cache.stride() == v_scale_cache.stride()
+ assert k_scale_cache.stride(2) == HEAD_DIM // 16
+ assert block_table.dtype == torch.int32 and block_table.stride(1) == 1
+ assert block_table.stride(0) == block_table.shape[1], "row-contiguous block table"
+ assert selected_tokens.dtype == torch.int32 and selected_tokens.is_contiguous()
+ assert request_indices.is_contiguous() and request_indices.dtype in (torch.int32, torch.int64)
+ group = num_q_heads // num_kv_heads
+ width = selected_tokens.shape[1]
+ out = torch.empty((rows, num_q_heads, head_dim), dtype=torch.bfloat16, device=q.device)
+ if rows == 0:
+ return out
+ n_groups = -(-rows // rpr)
+ units = n_groups * num_kv_heads
+ n_whole = None
+ splits_given = splits
+ if splits is None:
+ splits, n_whole = plan_launch(n_groups, num_kv_heads, width, _num_sms(q.device.index or 0))
+ n_splits = max(1, min(MAX_SPLITS, int(splits)))
+ tail = n_whole is not None
+ assert (
+ _layout(rpr, group, width, block_table.shape[1], tpb, 1 if tail else n_splits) is not None
+ ), "no shared-memory plan for this geometry: route through nvfp4_sparse_gqa_verify"
+ use_pdl = bool(pdl and get_env_enable_pdl())
+ gated = output_gate is not None
+ if gated:
+ assert output_gate.shape == (rows, num_q_heads * head_dim) and output_gate.stride(1) == 1
+ assert output_gate.dtype == torch.bfloat16
+ assert (output_gate.stride(0) * 2) % 16 == 0 and output_gate.data_ptr() % 16 == 0
+ req64 = request_indices.dtype == torch.int64
+ if tail:
+ tail_row0 = (n_whole // num_kv_heads) * rpr
+ n_ctas = n_whole + (units - n_whole) * n_splits
+ elif n_splits > 1:
+ n_whole, tail_row0, n_ctas = 0, 0, units * n_splits
+ else:
+ n_whole, tail_row0, n_ctas = units, rows, units
+ merged = rows - tail_row0
+ key = (
+ rpr,
+ group,
+ num_kv_heads,
+ width,
+ tpb,
+ n_splits,
+ tail,
+ float(softmax_scale),
+ num_pages,
+ tuple(block_table.shape),
+ tuple(k_cache.stride()[:3]),
+ tuple(k_scale_cache.stride()[:2]),
+ tuple(q.stride()[:2]),
+ req64,
+ gated,
+ use_pdl,
+ )
+ if merged > 0:
+ part = torch.empty(
+ (n_splits, merged, num_q_heads, head_dim), dtype=torch.float32, device=q.device
+ )
+ lse = torch.empty((n_splits, merged, num_q_heads), dtype=torch.float32, device=q.device)
+ else:
+ part, lse = out, out
+ gate_t = output_gate if gated else q
+ gate_stride = output_gate.stride(0) if gated else 0
+ compile_args = (
+ # read with scalar loads: a row slice of the selection needs no copy
+ from_dlpack(selected_tokens, assumed_align=4).mark_layout_dynamic(leading_dim=1),
+ _ptr(q, cutlass.BFloat16),
+ _ptr(k_cache, cutlass.Uint8),
+ _ptr(v_cache, cutlass.Uint8),
+ _ptr(k_scale_cache, cutlass.Uint8),
+ _ptr(v_scale_cache, cutlass.Uint8),
+ _ptr(kv_scale_quant_orig, cutlass.Float32, 4),
+ _ptr(block_table, cutlass.Int32, 4),
+ _ptr(request_indices, cutlass.Int64 if req64 else cutlass.Int32, 8 if req64 else 4),
+ _ptr(out, cutlass.BFloat16),
+ _ptr(part, cutlass.Float32),
+ _ptr(lse, cutlass.Float32),
+ _ptr(gate_t, cutlass.BFloat16),
+ cutlass.Int32(rows),
+ cutlass.Int32(gate_stride),
+ cutlass.Int32(n_whole),
+ cutlass.Int32(tail_row0),
+ cutlass.Int32(n_ctas),
+ cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+ )
+ merge_args = (
+ from_dlpack(gate_t, assumed_align=16).mark_layout_dynamic(leading_dim=gate_t.dim() - 1),
+ _ptr(part, cutlass.Float32),
+ _ptr(lse, cutlass.Float32, 4),
+ _ptr(out, cutlass.BFloat16, 8),
+ cutlass.Int32(merged),
+ cutlass.Int32(gate_stride),
+ cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+ )
+ kern_kw = dict(
+ rpr=rpr,
+ group=group,
+ kv_heads=num_kv_heads,
+ width=width,
+ tpb=tpb,
+ scale=float(softmax_scale),
+ num_pages=num_pages,
+ bt_width=block_table.shape[1],
+ num_requests=block_table.shape[0],
+ k_strides=tuple(k_cache.stride()[:3]),
+ ks_strides=tuple(k_scale_cache.stride()[:2]),
+ q_strides=tuple(q.stride()[:2]),
+ req64=req64,
+ gated=gated,
+ pdl=use_pdl,
+ )
+ # A launch plan depends on the call's row count: an eager step (a mixed step's
+ # generation rows) meeting a new one would compile mid-serving, for seconds. The first
+ # call of a geometry compiles every plan its row counts can reach instead.
+ if splits_given is None and key[:5] + key[7:] not in _PLANNED:
+ _PLANNED.add(key[:5] + key[7:])
+ # A group is one request's rows: at most a block-table row's worth of them.
+ max_groups = block_table.shape[0]
+ for plan_splits_, plan_tail in reachable_plans(
+ max_groups, num_kv_heads, width, _num_sms(q.device.index or 0)
+ ):
+ if (
+ _layout(
+ rpr, group, width, block_table.shape[1], tpb, 1 if plan_tail else plan_splits_
+ )
+ is None
+ ):
+ continue
+ _compile_verify(key[:5] + (plan_splits_, plan_tail) + key[7:], kern_kw, compile_args)
+ if plan_tail or plan_splits_ > 1:
+ _compile_merge(num_q_heads, plan_splits_, gated, use_pdl, merge_args)
+ fn = _compile_verify(key, kern_kw, compile_args)
+ fn(
+ selected_tokens,
+ q.data_ptr(),
+ k_cache.data_ptr(),
+ v_cache.data_ptr(),
+ k_scale_cache.data_ptr(),
+ v_scale_cache.data_ptr(),
+ kv_scale_quant_orig.data_ptr(),
+ block_table.data_ptr(),
+ request_indices.data_ptr(),
+ out.data_ptr(),
+ part.data_ptr(),
+ lse.data_ptr(),
+ gate_t.data_ptr(),
+ rows,
+ gate_stride,
+ n_whole,
+ tail_row0,
+ n_ctas,
+ )
+ if merged <= 0:
+ return out
+ mfn = _compile_merge(num_q_heads, n_splits, gated, use_pdl, merge_args)
+ mfn(
+ gate_t[tail_row0:],
+ part.data_ptr(),
+ lse.data_ptr(),
+ out[tail_row0:].data_ptr(),
+ merged,
+ gate_stride,
+ )
+ return out
+
+
+def nvfp4_sparse_gqa_verify(
+ 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,
+ rows_per_request: int,
+ output_gate: Optional[torch.Tensor] = None,
+ pdl: bool = True,
+) -> torch.Tensor:
+ """`nvfp4_sparse_gqa_decode` for rows that come `rows_per_request` consecutive rows a
+ request: the grouped kernel where `takes_grouped`, the per-row kernel otherwise. Returns
+ `[rows, Hq, D]` bf16 (contiguous); with `output_gate` (`[rows, Hq * D]` bf16, unit column
+ stride) `bf16(attn) * sigmoid(gate)` rounded to bf16.
+ """
+ plan = _plan(q, k_cache, block_table, selected_tokens, tokens_per_block, rows_per_request)
+ if plan is None:
+ return nvfp4_sparse_gqa_decode(
+ q,
+ k_cache,
+ v_cache,
+ k_scale_cache,
+ v_scale_cache,
+ kv_scale_quant_orig,
+ block_table,
+ selected_tokens,
+ request_indices,
+ tokens_per_block,
+ softmax_scale,
+ output_gate=output_gate,
+ pdl=pdl,
+ )
+ return nvfp4_sparse_gqa_verify_grouped(
+ q,
+ k_cache,
+ v_cache,
+ k_scale_cache,
+ v_scale_cache,
+ kv_scale_quant_orig,
+ block_table,
+ selected_tokens,
+ request_indices,
+ tokens_per_block,
+ softmax_scale,
+ rows_per_request,
+ output_gate=output_gate,
+ pdl=pdl,
+ splits=None,
+ )
+
+
+def _verify(segments, overlap=0.05, **overrides) -> dict:
+ """A cell spec: `(first position, rows)` per request in row order, four rows a request
+ unless `rows_per_request` overrides it, this checkpoint's geometry (24 q heads over 2 KV
+ heads of 256, 32 tokens a page, 2051 selected tokens). `overlap` is the spread of a
+ request's rows' block scores around the request's (each row then takes its top 512 blocks
+ in ascending order, as `qsa_topk_expand_decode` writes them); None selects every row
+ independently at random, as the decode op's cells do.
+ """
+ return dict(segments=segments, rows_per_request=4, overlap=overlap, **overrides)
+
+
+class QsaPagedSparseGqaNvfp4Verify(OpWrapper):
+ """`qsa_paged_sparse_gqa_nvfp4_decode` for rows that come `rows_per_request` consecutive
+ rows a request (an MTP verify step: T = draft length + 1 rows a generation request):
+ sparse GQA over the explicitly selected tokens of an NVFP4 paged pool, `[rows, Hq, D]`
+ bf16, each selected K/V token loaded and dequantized once for a group of T rows. With
+ `output_gate` (`[rows, Hq * D]` bf16, unit column stride) the result is
+ `bf16(attn) * sigmoid(gate)` rounded once more to bf16.
+
+ Fusion boundary: the block-table gather, the union of the group's selections, the NVFP4
+ dequantize, scale, online softmax, P.V, the split-K merge and (optionally) the output gate.
+ Selection, the K/V append and q's projection, norm and RoPE are not inside.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="MTP3 verify, 8 requests x 4 rows: the fewest the grouped kernel takes, "
+ "9 splits, gate fused into the merge",
+ spec=_verify(tuple((4000 + 37 * i, 4) for i in range(8)), gated=True),
+ ),
+ Cell(
+ why="32 rows whose selections share little (independent random rows, union "
+ "~2.7x a row's)",
+ spec=_verify(tuple((3000 + 211 * i, 4) for i in range(8)), overlap=None),
+ ),
+ Cell(
+ why="MTP3 verify at 128 rows (32 requests x 4, the mid-concurrency step): two "
+ "splits a unit, merged by the row merge with the gate",
+ spec=_verify(tuple((2100 + (53 * i) % 1900, 4) for i in range(32)), gated=True),
+ ),
+ Cell(
+ why="MTP3 verify at 1024 rows (256 requests x 4): one CTA a unit with the gate fused "
+ "into its store, the last wave's units split two ways and merged",
+ spec=_verify(tuple((2100 + (37 * i) % 1900, 4) for i in range(256)), gated=True),
+ ),
+ Cell(
+ why="groups spanning requests (3, 4, 4, 1, 4 x 5, 2 rows: misaligned and partial "
+ "groups), early positions, int64 request indices, non-unit global scales, trunk "
+ "layout",
+ spec=_verify(
+ (
+ (0, 3),
+ (5, 4),
+ (37, 4),
+ (4000, 1),
+ (4100, 4),
+ (4200, 4),
+ (4300, 4),
+ (4400, 4),
+ (4500, 4),
+ (4600, 2),
+ ),
+ request_dtype=torch.int64,
+ global_k=0.37,
+ global_v=2.5,
+ slot_pages=24,
+ layer=11,
+ pool_pages=2048,
+ ),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none -- every input is an argument (as `qsa_paged_sparse_gqa_nvfp4`).
+ Call sites (the target's `modeling.py`): `_sparse_gqa` over an NVFP4 pool for rows without
+ the prefill pair that come `rows_per_request` > 1 a request -- the decode target's
+ attention layers on a verify step (with the output gate) and a mixed step's generation
+ rows. Draft-layer and one-row-a-request calls stay on `qsa_paged_sparse_gqa_nvfp4_decode`.
+
+ Routing (`takes_grouped`, host-side, on shapes only): the grouped kernel from
+ `MIN_GROUPS` (8) groups of T >= 2 rows up, on SM103, where its shared-memory plan fits
+ (T <= 6; a block table of at most 612 pages of 32 tokens at T = 4, 892 at T = 3, 1364 at
+ T = 2); the per-row kernel otherwise. Launch (`plan_launch`): within one wave of
+ (group, KV head) units every unit takes the same splits, filling the GPU once (one CTA an
+ SM), followed by the row merge; past one wave the units of the full waves take one
+ CTA each (bf16 output, the gate fused) and the units of the partial last wave two to four,
+ only their rows merged.
+
+ GB300, one layer, cold L2, T = 4, the decode step's call (output gate on) on MTP3 verify
+ selections dumped from a served run (a request's rows' union 1.4-1.7x a row's selection,
+ about 8.4K tokens of context, the server's pool geometry): 64 rows 29.7 us, 128 rows 46.1,
+ 256 rows 75.3, 512 rows 137, 1024 rows 233 -- against 39.5 / 69.4 / 135 / 237 / 461 us for
+ the per-row kernel. Below 32 rows the call is the per-row kernel's.
+
+ Numerics: dequantized K/V values are exact in f16; scores on f16 MMAs with fp32
+ accumulation; P rounded to f16 (as the per-row kernel); a row's softmax offset moves only
+ when its max grew by more than 2^8, the sum stays fp32 unrounded; the second-level scales
+ act on the scores and the normalized output; the output gate's sigmoid on `ex2.approx` and
+ `rcp.approx` (relative error ~1e-6) where the kernel applies it, as the merge's elsewhere.
+ Within the decode op's band (4 ulp element,
+ 2 ulp RMS) of the float64 reference and of the per-row kernel; not bitwise the per-row
+ kernel (tiles hold the group's union). Deterministic; graph-capturable (row count is a
+ runtime value; one compile per (T, splits, launch mode) variant).
+
+ Preconditions: as `qsa_paged_sparse_gqa_nvfp4_decode` (bf16 q with head_dim 256 and
+ 16-byte aligned rows; K and V planes, and their scale planes, share strides; power-of-two
+ tokens per block; int32 row-contiguous block table; contiguous int32 selections), and
+ expanded selections: each four-column quad of a row names one four-token group at its own
+ offsets (column % 4 == token % 4), no group twice in a row -- what
+ `qsa_topk_expand_decode` and `qsa_expand_block_indices` write. The rows of a group share
+ the K/V tiles, so a NaN in a token one row selects reaches the group's other rows of that
+ request.
+ """
+
+ def __call__(
+ self,
+ q,
+ k_cache,
+ v_cache,
+ k_scale_cache,
+ v_scale_cache,
+ kv_scale_quant_orig,
+ block_table,
+ selected_tokens,
+ request_indices,
+ tokens_per_block: int,
+ softmax_scale: float,
+ rows_per_request: int,
+ output_gate=None,
+ ) -> torch.Tensor:
+ return nvfp4_sparse_gqa_verify(
+ q,
+ k_cache,
+ v_cache,
+ k_scale_cache,
+ v_scale_cache,
+ kv_scale_quant_orig,
+ block_table,
+ selected_tokens,
+ request_indices,
+ tokens_per_block,
+ softmax_scale,
+ rows_per_request,
+ output_gate=output_gate,
+ )
+
+ def reference(
+ self,
+ q,
+ k_cache,
+ v_cache,
+ k_scale_cache,
+ v_scale_cache,
+ kv_scale_quant_orig,
+ block_table,
+ selected_tokens,
+ request_indices,
+ tokens_per_block: int,
+ softmax_scale: float,
+ rows_per_request: int,
+ output_gate=None,
+ ) -> torch.Tensor:
+ """`qsa_paged_sparse_gqa_nvfp4_decode`'s reference (`rows_per_request` only groups the
+ work, it does not change the result).
+ """
+ return qsa_paged_sparse_gqa_nvfp4_decode.reference(
+ q,
+ k_cache,
+ v_cache,
+ k_scale_cache,
+ v_scale_cache,
+ kv_scale_quant_orig,
+ block_table,
+ selected_tokens,
+ request_indices,
+ tokens_per_block,
+ softmax_scale,
+ output_gate=output_gate,
+ )
+
+ def is_valid(
+ self,
+ q,
+ k_cache,
+ v_cache,
+ k_scale_cache,
+ v_scale_cache,
+ kv_scale_quant_orig,
+ block_table,
+ selected_tokens,
+ request_indices,
+ tokens_per_block: int,
+ softmax_scale: float,
+ rows_per_request: int,
+ output_gate=None,
+ ) -> None:
+ qsa_paged_sparse_gqa_nvfp4_decode.is_valid(
+ q,
+ k_cache,
+ v_cache,
+ k_scale_cache,
+ v_scale_cache,
+ kv_scale_quant_orig,
+ block_table,
+ selected_tokens,
+ request_indices,
+ tokens_per_block,
+ softmax_scale,
+ output_gate=output_gate,
+ )
+ assert int(rows_per_request) >= 1, "rows_per_request >= 1"
+ # The union is sized for expanded selections: a quad of columns names one four-token
+ # group at its own offsets, and no group twice in a row. Anything else can overflow
+ # the union, or count a token once where the reference counts it twice.
+ rows, width = selected_tokens.shape
+ quads = -(-width // 4)
+ padded = torch.full((rows, 4 * quads), -1, dtype=torch.int64, device=selected_tokens.device)
+ padded[:, :width] = selected_tokens.long()
+ sel = padded.view(rows, quads, 4)
+ valid = sel >= 0
+ column = torch.arange(4, device=sel.device)
+ assert bool(((sel % 4 == column) | ~valid).all()), "column % 4 == token % 4"
+ group = torch.where(valid, sel // 4, torch.full_like(sel, -1)).amax(dim=2)
+ assert bool(((sel // 4 == group[..., None]) | ~valid).all()), "one group a quad"
+ listed = torch.where(group >= 0, group, torch.full_like(group, -1)).sort(dim=1).values
+ repeat = (listed[:, 1:] == listed[:, :-1]) & (listed[:, 1:] >= 0)
+ assert not bool(repeat.any()), "no four-token group twice in a row"
+
+
+qsa_paged_sparse_gqa_nvfp4_verify = QsaPagedSparseGqaNvfp4Verify()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_prefill_compress.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_prefill_compress.py
new file mode 100644
index 000000000000..0a051c7f90c0
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_prefill_compress.py
@@ -0,0 +1,375 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""QSA block compression of the paged index-key side cache: completed groups to anchors.
+
+Called by the qwen4_exp target (`qwen3_8_flash_next_nvfp4__sm_103__default`) at every
+layout.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+# Tier-2 sourcing: a vendored Triton callable reached as ordinary Python. QSA's
+# kernel plane has no torch.ops surface.
+from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels as _qsa_kernels
+
+#: The mRoPE sections the checkpoint declares. The target passes `None`
+#: (`_QSA_MROPE_SECTION`); these are certified on the NVFP4 trunk-plane cell.
+MROPE_SECTION = (11, 11, 10)
+
+#: This checkpoint's index geometry and the engine's QSA metadata: index head
+#: 128, rotary width 64 (32 cos/sin pairs), compress ratio 4, rms eps 1e-6,
+#: theta 1e7; `slots` x `width` is `qsa_block_table`. `pages` is the pool.
+_CKPT = dict(
+ head_dim=128,
+ tokens_per_block=32,
+ slots=32,
+ width=260,
+ compress_ratio=4,
+ eps=1e-6,
+ mrope_section=None,
+ logical_dtype=torch.int64,
+ request_dtype=torch.int32,
+ layout="plain",
+)
+
+#: 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))
+
+
+#: The default target's index pool at bf16 KV, as the stock `Qwen4ExpHybridCacheManager` lays it out on
+#: every member and rank: KV cache manager V2 coalesces the 12 sparse layers' same-size index
+#: pages into one pool, so a layer's `get_index_k_buffer` view is `[pages, 32, 1, 128]` with page
+#: stride 12 x 4096 elements, starting `ordinal x 4096` into each slot. The position plane is its
+#: own contiguous `[pages, 32, 3]` int32 pool.
+INDEX_ROLES = 12
+
+#: A rank's attention pool, in pages: the stock engine's `maxNumBlocks` at core config (bf16 KV,
+#: fraction 0.6, host-resident n-gram table), measured per member. The manager's index view holds
+#: slightly fewer pages than its primary block count, so these bound it from above.
+TEP4_PAGES = 258983
+DEP4_PAGES = 120530
+
+#: `qsa_block_table` is `[max_num_sequences, max_blocks_per_seq]`: 128 seats (`max_batch_size`, per
+#: rank under attention DP), 292 pages a sequence at max_seq_len 9216 and 516 at the MTP member's 16384.
+SEATS = 128
+CORE_WIDTH = 292
+MTP_WIDTH = 516
+
+
+def _spec(segments, pages, **overrides) -> dict:
+ """`segments` is `(first position, rows)` per request, in row order."""
+ return {**_CKPT, "segments": segments, "pages": pages, **overrides}
+
+
+def _mrope(segments, pages, **overrides) -> dict:
+ """Prefill from 0 under mRoPE (11, 11, 10)."""
+ return _spec(segments, pages, mrope_section=MROPE_SECTION, **overrides)
+
+
+def _pool(segments, pages, width, ordinal, **overrides) -> dict:
+ """The default target's call over a rank's coalesced index pool: layer `ordinal`'s plane, 128 seats,
+ and request 0's last page the pool's last (`top_page`)."""
+ return _spec(
+ segments,
+ pages,
+ slots=SEATS,
+ width=width,
+ layout="coalesced",
+ roles=INDEX_ROLES,
+ ordinal=ordinal,
+ top_page=True,
+ **overrides,
+ )
+
+
+def _decode(requests: int, last: int) -> tuple:
+ """`requests` decode rows spread over positions 3..`last`; every other one an anchor
+ (`4k + 3`), the others at `4k + 0..2`, the last request an anchor at `last`."""
+ span = (last - 3) // 4
+ segments = []
+ for r in range(requests):
+ k = r * span // (requests - 1)
+ offset = 3 if r % 2 == 0 or r == requests - 1 else r % 3
+ segments.append((4 * k + offset, 1))
+ return tuple(segments)
+
+
+#: 128 MTP7 verify blocks (8 rows, two anchors each) at starts spread to the MTP member's last
+#: position, 16383: some blocks inside one page, some across a page edge.
+_MTP7_128 = tuple((r * (16384 - 8) // 127, 8) for r in range(128))
+
+
+class QsaPrefillCompress(OpWrapper):
+ """Replace each completed group's anchor slot with its pooled, normed, rotated key.
+
+ Wraps `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels.
+ triton_qsa_prefill_compress`: one Triton launch, grid `(rows,)`. A row
+ with `(L + 1) % cr == 0` (L = its logical position) is the anchor of the
+ group `L - cr + 1 .. L`; every other row does nothing. Each member is
+ resolved through `block_table` on its own, so groups may straddle pages::
+
+ pooled = bf16(sum_f32(raw keys of the group) / cr)
+ y = f32(pooled) * rsqrt(mean(f32(pooled)^2) + eps) * (1 + k_norm_weight)
+ y = cos_sin.dtype(y)
+ (t, h, w) = position_cache at the group's FIRST member
+ pair q takes h if q % 3 == 1 and q < 3 * s1, w if q % 3 == 2 and
+ q < 3 * s2, else t (t for every pair when mrope_section is None)
+ NeoX-rotate y[:D/2] by cos_sin[pos[q], q], cos_sin[pos[q], R + q];
+ y[D/2:] passes through (R = cos_sin.shape[1] // 2 = D / 4)
+ index_cache[anchor] = bf16(y)
+
+ Returns `index_cache`, written in place at anchor slots only; everything
+ else, including the non-anchor members' raw keys, is untouched.
+
+ Not idempotent: the anchor slot is the group's last input and its only
+ output, so a second call over the same rows re-pools a compressed anchor.
+
+ Fusion boundary: pool, Gemma norm, position gather, rotation, store. The
+ raw-key and position stores (`qsa_paged_index_store`, which must run
+ first), the table and metadata, and the main K/V pool are the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ 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 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, position strides (24672, 3, 1).
+ Cell(
+ why="NVFP4 trunk planes: 12 index roles and the position page share each page",
+ spec=_mrope(((0, 64),) * 2, pages=4, layout="shared_page", roles=12),
+ ),
+ # The default target's calls over a rank's measured pool, every page reachable: the last
+ # page's plane sits past 2^31 elements of the pool's base.
+ Cell(
+ why="tep4 decode: 128 requests x 1 row, the 258983-page coalesced pool, ordinal 11",
+ spec=_pool(_decode(128, 9215), TEP4_PAGES, CORE_WIDTH, 11),
+ ),
+ Cell(
+ why="tep4 MTP7 verify: 128 requests x 8 rows = 1024, positions to 16383",
+ spec=_pool(_MTP7_128, TEP4_PAGES, MTP_WIDTH, 5),
+ ),
+ Cell(
+ why="tep4 prefill at max_num_tokens: 8192 rows (6144 + 2047 + 1, all from 0)",
+ spec=_pool(((0, 6144), (0, 2047), (0, 1)), TEP4_PAGES, CORE_WIDTH, 0, rounding_flips=5),
+ ),
+ Cell(
+ why="dep4 decode: 32 requests x 1 row over a rank's 120530-page pool, 128 seats",
+ spec=_pool(_decode(32, 9215), DEP4_PAGES, CORE_WIDTH, 11),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none -- every input is an argument. The target passes
+ `qsa_logical_positions` (int64) and `qsa_req_idx_per_token` (int32)
+ sliced to the step's rows, `qsa_block_table` unsliced, the layer's
+ `get_index_k_buffer(i)` and the shared `get_qsa_position_buffer()`, the
+ indexer's `k_layernorm` weight (bf16 [128]) and eps, a float32
+ `[max_positions, 64]` cos|sin table built at load, and `mrope_section`
+ None. It runs on every step, decode and prefill rows alike.
+
+ Preconditions:
+ - every member's raw key and the first member's coordinates are in the
+ caches before the call; this call stores neither.
+ - every logical position is non-negative. `L = -1` satisfies the anchor
+ test and addresses positions -4..-1, an illegal memory access that
+ poisons the process. A value, so not guarded; pad with 0 (an exact
+ no-op).
+ - `logical_positions`, `request_indices`, `k_norm_weight` and the last
+ dim of `cos_sin` are read with an assumed unit stride while the
+ callable checks only shape, so a strided one is accepted and wrong.
+ `is_valid` refuses them.
+ The callable rejects loudly (not pinned by a test): metadata shape/dtype mismatches, a non-int32
+ or non-2-D table, more than one index head, a position plane with fewer
+ than 3 axes or not int32, a non-bf16 cache, page sizes that disagree with
+ `tokens_per_block`, a malformed cos|sin table, `4 R != D`, a non-2^k `D`
+ or `compress_ratio`, a norm weight of the wrong length, CPU operands, and
+ `mrope_section` that is not 3 entries summing to `R`. `rows == 0` returns
+ before launching.
+
+ `cos_sin`'s dtype is the rotation's precision: bf16 is accepted, puts some
+ elements off by up to 1.56e-2, and is not certified. Text rows reduce
+ mRoPE to plain RoPE bit for bit (measured, not pinned by a test). Plain
+ `@triton.jit`, no autotune, JIT-compiled at first call; the cells drove
+ that JIT path on sm_103 (GB300). Every stride is a
+ `tl.constexpr`, so each pool layout (contiguous, the NVFP4 shared page,
+ the default target's coalesced bf16 pool) is its own compiled specialization, and
+ each is driven by a cell. CUDA-graph capturable at a fixed row count
+ (captured at 256 rows, 32 MTP7 verify blocks, and at 128 decode rows over
+ the default target's tep4 pool).
+
+ Under an NVFP4 KV pool both caches are strided views of one per-slot
+ allocation (each physical page holds the index roles, then the position
+ page); the trunk-plane cell checks the other index roles and the position
+ page untouched.
+
+ The `coalesced` cells hold the default target's bf16-KV index pool (layout
+ at `INDEX_ROLES`, read off the stock manager at tep4 and dep4 rank 0) at a
+ rank's page count -- tep4 258983, dep4 120530, an upper bound on the index
+ view -- with request 0's last page the pool's last, written at 258982 x
+ 49152 elements from the base, past 2^31. The table is the engine's
+ `[128, 292]` (core) or `[128, 516]` (MTP, max_seq_len 16384), 128 seats at
+ 128 and at 32 requests. The other 11 layers' planes are checked untouched.
+
+ The kernel's `tl.rsqrt` is the hardware approximation, an fp32 ulp from the
+ reference's: every cell is bitwise the reference but two. In the 32-block
+ MTP7 one, 1 of 8192 anchor elements rounds one bf16 step away. In the 8192-row
+ tep4 prefill, 5 of 262016 do, every one an exact bf16 tie in the reference's
+ fp32 value (low 16 bits 0x8000). The nearest wrong variants (sequential bf16
+ pooling, coordinates from the group's last member, the norm weight without
+ its `+1`) fall far outside that budget.
+
+ Hard limits (read from the kernel; the page offset also driven): positions,
+ requests and page ids are widened to int64 before any address arithmetic,
+ so no int32 boundary inside the call -- the coalesced cells address pages
+ past 2^31 elements of the pool's base. Certified extent: positions to 16383
+ through a [128, 516] table over 258983 pages; 8192 rows in one launch.
+
+ Each rank's call is local to its device (the indexer is replicated, no
+ collective), so a one-device cell stands for any rank's call at its pool.
+
+ Not certified: `compress_ratio` other than 4, `D` other than 128, a bf16
+ or fp16 `cos_sin`, positions past 16383 or the table's rows (max_seq_len
+ 262144 included), mRoPE outside the NVFP4
+ trunk-plane cell, the NVFP4 draft planes (one index role a page),
+ `tokens_per_block` other than 32 (no group straddles a page), eps other
+ than 1e-6, int32 positions or int64 request indices, a strided view of the
+ index cache, position plane or table other than the NVFP4 pool's shared
+ page and the default target's coalesced bf16 pool, the kv_fp8 / kv_nvfp4 dep4
+ pools at their measured page count (their packed `QSA_SIDE_PLANES` layout
+ is the shared-page cell's, at 4 pages), the MTP member's own pool count and
+ its separate draft cache (the 1024-row cell runs on the core tep4 pool), a
+ tp4 pool (never booted; the indexer is replicated, so tep4's stands in), two
+ rows naming one anchor in a launch, and `rows == 0`.
+ """
+
+ def __call__(
+ self,
+ *,
+ logical_positions: torch.Tensor,
+ request_indices: torch.Tensor,
+ block_table: torch.Tensor,
+ index_cache: torch.Tensor,
+ position_cache: torch.Tensor,
+ k_norm_weight: torch.Tensor,
+ cos_sin: torch.Tensor,
+ eps: float,
+ tokens_per_block: int,
+ compress_ratio: int,
+ mrope_section: Optional[Tuple[int, int, int]],
+ ) -> torch.Tensor:
+ _qsa_kernels.triton_qsa_prefill_compress(
+ logical_positions=logical_positions,
+ request_indices=request_indices,
+ block_table=block_table,
+ index_cache=index_cache,
+ position_cache=position_cache,
+ k_norm_weight=k_norm_weight,
+ cos_sin=cos_sin,
+ eps=eps,
+ tokens_per_block=tokens_per_block,
+ compress_ratio=compress_ratio,
+ mrope_section=mrope_section,
+ )
+ return index_cache
+
+ def reference(
+ self,
+ *,
+ logical_positions: torch.Tensor,
+ request_indices: torch.Tensor,
+ block_table: torch.Tensor,
+ index_cache: torch.Tensor,
+ position_cache: torch.Tensor,
+ k_norm_weight: torch.Tensor,
+ cos_sin: torch.Tensor,
+ eps: float,
+ tokens_per_block: int,
+ compress_ratio: int,
+ mrope_section: Optional[Tuple[int, int, int]],
+ ) -> torch.Tensor:
+ """The index cache after the call, as a fresh tensor, from the pre-call state."""
+ out = index_cache.clone()
+ logical = logical_positions.long()
+ anchored = (logical + 1) % compress_ratio == 0
+ anchor = logical[anchored]
+ if anchor.numel() == 0:
+ return out
+ device = index_cache.device
+ request = request_indices.long()[anchored]
+ members = (
+ anchor[:, None] - (compress_ratio - 1) + torch.arange(compress_ratio, device=device)
+ )
+ pages = block_table.long()[request[:, None], members // tokens_per_block]
+ slots = members % tokens_per_block
+
+ head_dim = index_cache.shape[3]
+ raw = index_cache[pages, slots, 0].float() # [anchors, cr, D]
+ pooled = (raw.sum(dim=1) / compress_ratio).to(torch.bfloat16).float()
+ scale = torch.rsqrt(pooled.mul(pooled).sum(dim=-1, keepdim=True) / head_dim + eps)
+ normed = (pooled * scale * (k_norm_weight.float() + 1.0)).to(cos_sin.dtype)
+
+ t, h, w = position_cache[pages[:, 0], slots[:, 0], :3].long().unbind(-1)
+ pairs_n = cos_sin.shape[1] // 2
+ pairs = torch.arange(pairs_n, device=device)
+ if mrope_section is None:
+ pos = t[:, None].expand(-1, pairs_n)
+ else:
+ pos = torch.where(
+ (pairs % 3 == 1) & (pairs < 3 * mrope_section[1]),
+ h[:, None],
+ torch.where(
+ (pairs % 3 == 2) & (pairs < 3 * mrope_section[2]), w[:, None], t[:, None]
+ ),
+ )
+ cos, sin = cos_sin[pos, pairs], cos_sin[pos, pairs_n + pairs]
+ first, second = normed[:, :pairs_n], normed[:, pairs_n : 2 * pairs_n]
+ rotated = torch.cat(
+ [first * cos - second * sin, second * cos + first * sin, normed[:, 2 * pairs_n :]],
+ dim=-1,
+ )
+ out[pages[:, -1], slots[:, -1], 0] = rotated.to(torch.bfloat16)
+ return out
+
+ def is_valid(
+ self,
+ *,
+ logical_positions: torch.Tensor,
+ request_indices: torch.Tensor,
+ block_table: torch.Tensor,
+ index_cache: torch.Tensor,
+ position_cache: torch.Tensor,
+ k_norm_weight: torch.Tensor,
+ cos_sin: torch.Tensor,
+ eps: float,
+ tokens_per_block: int,
+ compress_ratio: int,
+ mrope_section: Optional[Tuple[int, int, int]],
+ ) -> None:
+ # Read by the kernel with an assumed unit stride; the callable checks
+ # only their shapes. Each is accepted and silently wrong.
+ assert logical_positions.is_contiguous(), "logical_positions must be contiguous"
+ assert request_indices.is_contiguous(), "request_indices must be contiguous"
+ assert k_norm_weight.is_contiguous(), "k_norm_weight must be contiguous"
+ # `ndim != 2` is the callable's own loud error; leave it to the callable.
+ assert cos_sin.ndim != 2 or cos_sin.stride(1) == 1, (
+ "cos_sin must have unit stride along its last dimension"
+ )
+
+
+qsa_prefill_compress = QsaPrefillCompress()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_rollback_planes.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_rollback_planes.py
new file mode 100644
index 000000000000..c7e0ff27e19d
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_rollback_planes.py
@@ -0,0 +1,241 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The speculative rollback of every attention layer's QSA index-K plane in one launch each way:
+gather the verify rows' slots of all planes before the trunk writes them, and after the
+draft loop write the gathered values back into the slots of the rows the sampler rejected.
+
+Tier 5 (Triton, called by the qwen4_exp target, `qwen3_8_flash_next_nvfp4__sm_103__default`,
+at every layout). The planes are separate buffers of the cache manager (one per
+attention layer, each `[pages, tokens_per_block, heads, dim]` with a free page stride and a
+contiguous tail); the kernels reach them through a table of their base addresses and page
+strides, one program per (plane, row).
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+#: This checkpoint's index plane row: one index KV head of 128 bf16.
+ROW = 128
+
+#: The index planes of core's bf16 K/V pool at every four-rank layout, as the stock
+#: `Qwen4ExpHybridCacheManager` (KV cache manager V2) allocates them: the 12 attention layers'
+#: `[32, 1, 128]` INDEX_KEY buffers coalesce into one pool, layer k's buffer at element 4096 k of
+#: a 49152-element page, so every plane is a `[pages, 32, 1, 128]` view at page stride 49152.
+#: 257109 pages is a tep4 rank's attention pool at core's configuration, which the index pool
+#: shares: the stock engine's `maxNumBlocks` 258983 less the 1874 recurrent-state slots that
+#: total also counts. The largest measured per member.
+_TEP4_POOL = dict(planes=12, layout="coalesced", pool_pages=257109)
+
+
+def _verify(requests: int, p: int) -> dict:
+ """A verify step's rows: `requests` generation requests of `p` rows (`max_draft_len + 1`)."""
+ return dict(_TEP4_POOL, rows=requests * p, requests=requests, p=p, rejected="bool")
+
+
+@triton.jit
+def _gather_kernel(ptrs, page_strides, pages, within, out, rows, ROW: tl.constexpr):
+ plane = tl.program_id(0)
+ r = tl.program_id(1)
+ base = tl.load(ptrs + plane).to(tl.pointer_type(tl.bfloat16))
+ stride = tl.load(page_strides + plane)
+ page = tl.load(pages + r)
+ tok = tl.load(within + r)
+ offs = tl.arange(0, ROW)
+ vals = tl.load(base + page * stride + tok * ROW + offs)
+ tl.store(out + (plane.to(tl.int64) * rows + r) * ROW + offs, vals)
+
+
+@triton.jit
+def _restore_kernel(ptrs, page_strides, pages, within, rejected, old, rows, ROW: tl.constexpr):
+ plane = tl.program_id(0)
+ r = tl.program_id(1)
+ if tl.load(rejected + r) != 0:
+ base = tl.load(ptrs + plane).to(tl.pointer_type(tl.bfloat16))
+ stride = tl.load(page_strides + plane)
+ page = tl.load(pages + r)
+ tok = tl.load(within + r)
+ offs = tl.arange(0, ROW)
+ vals = tl.load(old + (plane.to(tl.int64) * rows + r) * ROW + offs)
+ tl.store(base + page * stride + tok * ROW + offs, vals)
+
+
+def plane_table(planes: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+ """`(base addresses, page strides)` of the planes, int64 on their device (built once per
+ cache-manager bind: the planes do not move)."""
+ dev = planes[0].device
+ for t in planes:
+ assert t.dtype == torch.bfloat16 and t.dim() == 4, "index planes [pages, tpb, H, D] bf16"
+ assert t.shape[2] * t.shape[3] == ROW, "one 128-wide row a token"
+ assert t.stride(1) == ROW and t.stride(2) == t.shape[3] and t.stride(3) == 1, (
+ "contiguous [tpb, H, D] tail"
+ )
+ ptrs = torch.tensor([t.data_ptr() for t in planes], dtype=torch.int64, device=dev)
+ strides = torch.tensor([t.stride(0) for t in planes], dtype=torch.int64, device=dev)
+ return ptrs, strides
+
+
+class QsaRollbackGather(OpWrapper):
+ """`torch.stack([plane[pages, within] for plane in planes])` as `[L, R, 128]` bf16, with
+ the planes given as `plane_table(planes)`: every row's slot of every plane, in plane
+ order. Nothing is written but the output.
+
+ Fusion boundary: the gather of all planes; the slots (`pages`, `within`) are the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+ # Both halves are certified at the same cells: the target's at tp1, and its MTP verify
+ # rows at `max_batch_size` 128 over a tep4 rank's pool.
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="12 planes, 4 verify rows (MTP3, one request)", spec=dict(planes=12, rows=4)),
+ Cell(why="12 planes, 512 verify rows (MTP3, 128 requests)", spec=dict(planes=12, rows=512)),
+ Cell(
+ why="tep4 MTP1 verify: 128 requests x 2 rows (256) over the 257109-page coalesced pool",
+ spec=_verify(128, 2),
+ ),
+ Cell(
+ why="tep4 MTP7 verify: 128 requests x 8 rows (1024, the envelope's largest), same pool",
+ spec=_verify(128, 8),
+ ),
+ )
+ note = """
+ Metadata consumed: none. Call site: `SpeculativeRollback.begin`, before the verify step's
+ trunk writes any plane; `pages` / `within` are the generation rows' QSA page and offset
+ (`block_table[req, pos // 32]`, `pos % 32`), int64 `[R]`, R = requests x (max_draft_len + 1).
+ One program per (plane, row), one 256-byte row each. Preconditions: the planes as
+ `plane_table` accepts them (bf16 NHD views, any page stride); every (page, within) inside
+ the pool.
+
+ Triton, JIT-compiled at first call for the running device (sm_103 here).
+
+ Addressing is 64-bit: `pages` and the page strides are int64, so a slot past 2^31 elements
+ from a plane's base is reached. The tep4 cells gather from (and restore into) the pool's
+ last page, 257108 x 49152 = 1.26e10 elements past every plane's base, and its first.
+
+ Layouts certified: separate contiguous planes and the packed side-plane page of an FP8 /
+ NVFP4 pool (page stride 12 x 4096 + 192) at the tp1 cells' 4 and 512 rows over 600 pages;
+ core's coalesced bf16 pool (page stride 49152) at 256 and 1024 rows over 257109 pages.
+ `test_cuda_graph_replay` captures the 1024-row call and replays it at new slots.
+
+ Hard limits (not probed): rows is the launch grid's y extent, at most 65535; out-of-pool or
+ duplicated slots are the caller's. Not certified: a pool page count measured under the MTP
+ unit's own engine (`max_seq_len` 16384; the cells use core's tep4 count), and the
+ side-plane layout at 256 / 1024 rows.
+ """
+
+ def __call__(self, ptrs, page_strides, pages, within, num_planes: int) -> torch.Tensor:
+ rows = pages.shape[0]
+ out = torch.empty([num_planes, rows, ROW], dtype=torch.bfloat16, device=pages.device)
+ if rows:
+ _gather_kernel[(num_planes, rows)](
+ ptrs, page_strides, pages, within, out, rows, ROW=ROW, num_warps=1
+ )
+ return out
+
+ def reference(self, planes, pages, within, num_planes: int) -> torch.Tensor:
+ """Every row's slot of the first `num_planes` planes, by torch's advanced indexing.
+
+ Takes the planes themselves where `__call__` takes `plane_table(planes)`: torch
+ cannot dereference the raw base addresses the kernel is handed, so the reference reads
+ the tensors the table was built from. A copy, exact.
+ """
+ return torch.stack([t[pages, within].reshape(-1, ROW) for t in planes[:num_planes]])
+
+ def is_valid(self, ptrs, page_strides, pages, within, num_planes: int) -> None:
+ assert ptrs.dtype == torch.int64 and ptrs.numel() == num_planes, (
+ "ptrs must be plane_table's int64 table of exactly num_planes addresses; a "
+ "num_planes past it reads addresses beyond the table"
+ )
+ assert page_strides.dtype == torch.int64 and page_strides.numel() == num_planes, (
+ "page_strides must be plane_table's int64 table of num_planes strides"
+ )
+ assert pages.dtype == torch.int64 and within.dtype == torch.int64, (
+ "pages and within must be int64"
+ )
+ assert pages.shape == within.shape and pages.dim() == 1, (
+ "pages and within must be one [R] pair: rows are counted off pages, so a shorter "
+ "within is read past its end"
+ )
+
+
+class QsaRollbackRestore(OpWrapper):
+ """For every row `r` with `rejected[r]` and every plane `l`: `plane_l[pages[r], within[r]]
+ = old[l, r]` -- the inverse of `qsa_rollback_gather` over the rejected rows. Rows not
+ rejected are not touched.
+
+ Fusion boundary: the restore of all planes; `rejected` is the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+ # Both halves are certified at the same cells: the target's at tp1, and its MTP verify
+ # rows at `max_batch_size` 128 over a tep4 rank's pool.
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="12 planes, 4 verify rows (MTP3, one request)", spec=dict(planes=12, rows=4)),
+ Cell(why="12 planes, 512 verify rows (MTP3, 128 requests)", spec=dict(planes=12, rows=512)),
+ Cell(
+ why="tep4 MTP1 verify: 128 requests x 2 rows (256) over the 257109-page coalesced pool",
+ spec=_verify(128, 2),
+ ),
+ Cell(
+ why="tep4 MTP7 verify: 128 requests x 8 rows (1024, the envelope's largest), same pool",
+ spec=_verify(128, 8),
+ ),
+ )
+ note = """
+ Metadata consumed: none. Call site: `SpeculativeRollback.commit_speculative_states`,
+ after the draft loop: `old` is the begin step's `qsa_rollback_gather` output, `rejected`
+ the bool `[R]` mask `ordinal >= accepted` (int32 / int8 / bool, nonzero = restore; the tp1
+ cells drive int32, the tep4 cells the call site's bool).
+
+ Returns nothing: it writes the planes, which it holds only as `plane_table` addresses.
+ Triton, JIT-compiled at first call for the running device (sm_103 here).
+
+ Every cell checks that the rejected rows' slots hold `old`, the others the step's values,
+ and that nothing else in the pool was written. Addressing, layouts and the CUDA-graph
+ replay as `qsa_rollback_gather`'s note states; the same limits.
+
+ Silent, refused by `is_valid`: a non-contiguous `old` (read as dense `[L, R, 128]`;
+ measured wrong in `test_guard_refuses_a_strided_old`).
+ """
+
+ def __call__(self, ptrs, page_strides, pages, within, rejected, old) -> None:
+ num_planes, rows, _ = old.shape
+ if rows:
+ _restore_kernel[(num_planes, rows)](
+ ptrs, page_strides, pages, within, rejected, old, rows, ROW=ROW, num_warps=1
+ )
+
+ def reference(self, planes, pages, within, rejected, old) -> torch.Tensor:
+ """What each row's slot of each plane holds after the call, `[L, R, 128]`: `old` where
+ `rejected`, else the slot's value now. Read before the call, as the planes stand.
+
+ Takes the planes where `__call__` takes `plane_table(planes)`, for the reason
+ `QsaRollbackGather.reference` states. Every location outside these slots is unchanged,
+ which the test checks over the whole pool.
+ """
+ now = torch.stack([t[pages, within].reshape(-1, ROW) for t in planes[: old.shape[0]]])
+ return torch.where(rejected.bool()[None, :, None], old, now)
+
+ def is_valid(self, ptrs, page_strides, pages, within, rejected, old) -> None:
+ assert old.dtype == torch.bfloat16 and old.dim() == 3 and old.shape[2] == ROW, (
+ "old must be qsa_rollback_gather's [L, R, 128] bf16 output"
+ )
+ assert old.is_contiguous(), (
+ "old must be contiguous; the kernel reads it as dense [L, R, 128] and a strided "
+ "view restores the wrong values"
+ )
+ assert ptrs.numel() == old.shape[0] and page_strides.numel() == old.shape[0], (
+ "plane_table must hold one address and one stride per plane of old"
+ )
+ assert pages.shape == (old.shape[1],) and within.shape == (old.shape[1],), (
+ "pages and within must hold one slot per row of old"
+ )
+ assert rejected.shape == (old.shape[1],), "rejected must hold one flag per row of old"
+
+
+qsa_rollback_gather = QsaRollbackGather()
+qsa_rollback_restore = QsaRollbackRestore()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_rowtile_attention_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_rowtile_attention_cute.py
new file mode 100644
index 000000000000..1a9cac72e24e
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_rowtile_attention_cute.py
@@ -0,0 +1,1383 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The row-tile pass of `qsa_sparse_gqa_rowtile` on SM103 tensor cores (CuTe DSL, tcgen05).
+
+Inputs and output: per group of `rows_per_tile` rows of one request, the plan's packed
+slot list (`slots[group]`, `listed` entries), each row's bitmap re-indexed onto it
+(`packed_bits`, as int32 halves) and `group_plan[group] = (tiles, listed)`; groups with
+`tiles < 0` belong to the lone-row pass.
+
+A persistent kernel: one CTA per SM, each taking work items -- a (group, KV head) pair, 128
+MMA rows = `rows_per_tile` query rows x `group_size` heads -- from a global ticket counter,
+longest-first (the plan's last group first). 9 warps:
+
+* warp 5 lane 0 draws the tickets and publishes the items through a shared-memory ring that
+ every role reads; warps 5-8 gather each item's Q rows (into a 60 KB Q buffer) and its list's
+ K and V rows (64 a tile) with `cp.async` into one ring of `kv_slots` 32 KB SW128 slots, filled
+ in the MMA warp's order K0, K1, V0, K2, V1, ...; the next item's Q is fetched while the
+ current item's last tiles load;
+* warp 4 issues the tcgen05 MMAs: S[t % 2] = Q K(t)^T (Q from tensor memory), then
+ O += P(t - 1) V(t - 1) (P from tensor memory, overwriting the S it came from), `t` counting
+ the CTA's tiles across items;
+* warps 0-3, one thread per MMA row, per tile load S with its row max (`tcgen05.ld.red`),
+ rescale O when that max grew by more than `rescale_threshold` in log2 units (lazily;
+ otherwise exponentiate against the stale max), write P = 2^(s - max) on the row's own bits
+ and 0 elsewhere. After an item's last tile they copy the next item's Q from the Q buffer
+ into tensor memory (so its first QK MMAs overlap this epilogue), then normalize O, stage it
+ in the Q buffer as four 64-dim SW128 boxes of the item's G heads x R rows and store it with
+ TMA (thread 0 hands the buffer back to the loaders once the store has read it).
+
+The Q buffer is shared by Q and the staged O: a softmax thread's Q row and its O row occupy
+the same bytes, so the next item's Q is read out before this item's O is staged over it. An
+item followed by a lone (or empty) group stores O row by row from registers instead, as the
+loaders may then fetch the following item's Q into the buffer at any time.
+
+Tensor memory: O in columns 0-255, S/P buffers at 256 and 320, Q at 384 (512 columns).
+The tile's row max covers every listed entry, an upper bound of the row's own: a valid
+softmax offset. P is rounded to bf16 for the PV MMA and summed in fp32 unrounded.
+"""
+
+from __future__ import annotations
+
+import functools
+
+import cuda.bindings.driver as cuda_drv
+import cutlass
+import cutlass.cute as cute
+import cutlass.cute.nvgpu.tcgen05 as tcgen05
+import cutlass.pipeline as pipeline
+import cutlass.utils.blackwell_helpers as sm100_utils
+import torch
+from cutlass._mlir.dialects import nvvm as _nvvm
+from cutlass.cute.nvgpu import OperandMajorMode, cpasync
+from cutlass.cute.runtime import make_ptr
+from cutlass.cute.typing import Float32, Int32, Int64
+from cutlass.cutlass_dsl import dsl_user_op
+
+try:
+ from cutlass.memory import SmemAllocator
+except ImportError: # DSL releases before the memory namespace split
+ from cutlass.utils import SmemAllocator
+
+BF16 = cutlass.BFloat16
+_LOG2_E = 1.4426950408889634
+#: The head dimension the tensor-memory plan is laid out for.
+HEAD_DIM = 256
+#: MMA rows per CTA and tokens per tile (one 64-bit bitmap word).
+BLOCK_M = 128
+BLOCK_N = 64
+#: K/V ring slots (32 KB each); with the 60 KB Q buffer, about 220 KB of shared memory.
+KV_SLOTS = 5
+#: Work items in flight in the shared-memory schedule ring.
+SCHED_SLOTS = 8
+#: The next item's Q is fetched once this many tiles of the current item remain to load.
+Q_LEAD = 3
+
+
+@dsl_user_op
+def tcgen05_fence_before(*, loc=None, ip=None):
+ """tcgen05.fence::before_thread_sync: order this thread's prior tcgen05 ops before a thread sync."""
+ _nvvm.tcgen05_fence(_nvvm.Tcgen05FenceKind.BEFORE_THREAD_SYNC, loc=loc, ip=ip)
+
+
+@dsl_user_op
+def tcgen05_fence_after(*, loc=None, ip=None):
+ """tcgen05.fence::after_thread_sync: order later tcgen05 ops after a thread sync."""
+ _nvvm.tcgen05_fence(_nvvm.Tcgen05FenceKind.AFTER_THREAD_SYNC, loc=loc, ip=ip)
+
+
+class _RowtileAttentionKernel:
+ def __init__(
+ self,
+ group_size: int,
+ rows_per_tile: int,
+ head_dim: int = 256,
+ tokens_per_block: int = 32,
+ kv_slots: int = KV_SLOTS,
+ rescale_threshold: float = 8.0,
+ use_pdl: bool = True,
+ ):
+ assert head_dim == 256, (
+ "the TMEM plan (O 256 + 2 x S 64 + Q 128 columns) is laid out for D = 256"
+ )
+ assert group_size * rows_per_tile <= 128
+ self.group_size = group_size
+ self.rows_per_tile = rows_per_tile
+ self.d = head_dim
+ self.tpb = tokens_per_block
+ self.bm = 128
+ self.bn = 64
+ self.kv_slots = kv_slots
+ self.sched_slots = SCHED_SLOTS
+ self.q_lead = Q_LEAD
+ self.tau = rescale_threshold
+ self.use_pdl = use_pdl
+ self.mma_warp = 4
+ self.load_warp0 = 5
+ self.threads = 32 * 9
+ self.tmem_o = 0
+ self.tmem_s0 = 256
+ self.tmem_q = 384 # Q as the QK MMA's A operand: 256 bf16 a row, two a column
+ self.tmem_cols = 512
+ self.chunk_cols = 32 # O columns a TMEM load/store moves per thread
+
+ @cute.jit
+ def __call__(
+ self,
+ q: cute.Pointer,
+ k: cute.Pointer,
+ v: cute.Pointer,
+ slots: cute.Pointer,
+ bits: cute.Pointer,
+ plan: cute.Pointer,
+ out: cute.Pointer,
+ work: cute.Pointer,
+ rows: Int32,
+ num_groups: Int32,
+ num_kv_heads: Int32,
+ num_ctas: Int32,
+ q_stride_row: Int32,
+ k_stride_page: Int32,
+ k_stride_head: Int32,
+ k_stride_token: Int32,
+ v_stride_page: Int32,
+ v_stride_head: Int32,
+ v_stride_token: Int32,
+ slots_stride_group: Int32,
+ bits_stride_row: Int32,
+ out_stride_row: Int32,
+ scale_log2: Float32,
+ stream,
+ ):
+ def qk():
+ return sm100_utils.make_trivial_tiled_mma(
+ BF16,
+ BF16,
+ OperandMajorMode.K,
+ OperandMajorMode.K,
+ Float32,
+ tcgen05.CtaGroup.ONE,
+ (self.bm, self.bn),
+ tcgen05.OperandSource.TMEM,
+ )
+
+ def pv():
+ return sm100_utils.make_trivial_tiled_mma(
+ BF16,
+ BF16,
+ OperandMajorMode.K,
+ OperandMajorMode.MN,
+ Float32,
+ tcgen05.CtaGroup.ONE,
+ (self.bm, self.d),
+ tcgen05.OperandSource.TMEM,
+ )
+
+ # Overwrite / accumulate twins of each MMA: the ACCUMULATE field is set once at kernel
+ # scope, so no loop or branch rebinds an MMA value.
+ qk_mma, qk_mma_acc, pv_mma, pv_mma_acc = qk(), qk(), pv(), pv()
+ q_layout = sm100_utils.make_smem_layout_a(qk_mma, (self.bm, self.bn, self.d), BF16, 1)
+ k_layout = sm100_utils.make_smem_layout_b(
+ qk_mma, (self.bm, self.bn, self.d), BF16, self.kv_slots
+ )
+ v_layout = sm100_utils.make_smem_layout_b(
+ pv_mma, (self.bm, self.d, self.bn), BF16, self.kv_slots
+ )
+ p_layout = sm100_utils.make_smem_layout_a(pv_mma, (self.bm, self.d, self.bn), BF16, 1)
+
+ @cute.struct
+ class SharedStorage:
+ q_full: cute.struct.MemRange[Int64, 1]
+ kv_full: cute.struct.MemRange[Int64, self.kv_slots]
+ kv_empty: cute.struct.MemRange[Int64, self.kv_slots]
+ s_full: cute.struct.MemRange[Int64, 2]
+ p_full: cute.struct.MemRange[Int64, 2]
+ o_full: cute.struct.MemRange[Int64, 1]
+ o_final: cute.struct.MemRange[Int64, 1]
+ qs_full: cute.struct.MemRange[Int64, 1]
+ qs_empty: cute.struct.MemRange[Int64, 1]
+ sch_full: cute.struct.MemRange[Int64, self.sched_slots]
+ sch_empty: cute.struct.MemRange[Int64, self.sched_slots]
+ sch_val: cute.struct.MemRange[Int32, self.sched_slots]
+ tmem_holding: Int32
+ sKV: cute.struct.Align[cute.struct.MemRange[BF16, cute.cosize(k_layout)], 1024]
+ # The Q buffer / O staging: four 64-dim boxes of the live G x R MMA rows.
+ sQ: cute.struct.Align[
+ cute.struct.MemRange[BF16, self.group_size * self.rows_per_tile * self.d], 1024
+ ]
+
+ self.shared_storage = SharedStorage
+
+ # The epilogue's TMA store: out viewed as (D, heads, rows), one box = 64 dims of the
+ # group's G heads x R rows, staged in the Q buffer as 128-byte SW128 rows (MMA row m =
+ # local row x G + head), four boxes 64 G R elements apart.
+ G = self.group_size
+ R = self.rows_per_tile
+ o_box = cute.make_composed_layout(
+ cute.make_swizzle(3, 4, 3), 0, cute.make_layout((64, G, R), stride=(1, 64, 64 * G))
+ )
+ mO = cute.make_tensor(
+ out,
+ cute.make_layout((self.d, num_kv_heads * G, rows), stride=(1, self.d, out_stride_row)),
+ )
+ tma_atom_o, tma_tensor_o = cpasync.make_tiled_tma_atom(
+ cpasync.CopyBulkTensorTileS2GOp(), mO, o_box, (64, G, R)
+ )
+
+ self.kernel(
+ qk_mma,
+ qk_mma_acc,
+ pv_mma,
+ pv_mma_acc,
+ q_layout,
+ k_layout,
+ v_layout,
+ p_layout,
+ tma_atom_o,
+ tma_tensor_o,
+ q,
+ k,
+ v,
+ slots,
+ bits,
+ plan,
+ out,
+ work,
+ rows,
+ num_groups,
+ num_kv_heads,
+ q_stride_row,
+ k_stride_page,
+ k_stride_head,
+ k_stride_token,
+ v_stride_page,
+ v_stride_head,
+ v_stride_token,
+ slots_stride_group,
+ bits_stride_row,
+ out_stride_row,
+ scale_log2,
+ ).launch(
+ grid=[num_ctas, 1, 1],
+ block=[self.threads, 1, 1],
+ stream=stream,
+ min_blocks_per_mp=1,
+ use_pdl=self.use_pdl,
+ )
+
+ @cute.kernel
+ def kernel(
+ self,
+ qk_mma: cute.TiledMma,
+ qk_mma_acc: cute.TiledMma,
+ pv_mma: cute.TiledMma,
+ pv_mma_acc: cute.TiledMma,
+ q_layout: cute.ComposedLayout,
+ k_layout: cute.ComposedLayout,
+ v_layout: cute.ComposedLayout,
+ p_layout: cute.ComposedLayout,
+ tma_atom_o: cute.CopyAtom,
+ tma_tensor_o: cute.Tensor,
+ q: cute.Pointer,
+ k: cute.Pointer,
+ v: cute.Pointer,
+ slots: cute.Pointer,
+ bits: cute.Pointer,
+ plan: cute.Pointer,
+ out: cute.Pointer,
+ work: cute.Pointer,
+ rows: Int32,
+ num_groups: Int32,
+ num_kv_heads: Int32,
+ q_stride_row: Int32,
+ k_stride_page: Int32,
+ k_stride_head: Int32,
+ k_stride_token: Int32,
+ v_stride_page: Int32,
+ v_stride_head: Int32,
+ v_stride_token: Int32,
+ slots_stride_group: Int32,
+ bits_stride_row: Int32,
+ out_stride_row: Int32,
+ scale_log2: Float32,
+ ):
+ warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+ tidx, _, _ = cute.arch.thread_idx()
+ qk_mma.set(tcgen05.Field.ACCUMULATE, False)
+ qk_mma_acc.set(tcgen05.Field.ACCUMULATE, True)
+ pv_mma.set(tcgen05.Field.ACCUMULATE, False)
+ pv_mma_acc.set(tcgen05.Field.ACCUMULATE, True)
+
+ smem = SmemAllocator()
+ storage = smem.allocate(self.shared_storage)
+ q_full = storage.q_full.data_ptr()
+ kv_full = storage.kv_full.data_ptr()
+ kv_empty = storage.kv_empty.data_ptr()
+ s_full = storage.s_full.data_ptr()
+ p_full = storage.p_full.data_ptr()
+ o_full = storage.o_full.data_ptr()
+ o_final = storage.o_final.data_ptr()
+ qs_full = storage.qs_full.data_ptr()
+ qs_empty = storage.qs_empty.data_ptr()
+ sch_full = storage.sch_full.data_ptr()
+ sch_empty = storage.sch_empty.data_ptr()
+ sch_val = storage.sch_val.data_ptr()
+ if tidx == 0:
+ cute.arch.mbarrier_init(q_full, 128)
+ for i in cutlass.range_constexpr(self.kv_slots):
+ cute.arch.mbarrier_init(kv_full + i, 128)
+ cute.arch.mbarrier_init(kv_empty + i, 1)
+ for i in cutlass.range_constexpr(2):
+ cute.arch.mbarrier_init(s_full + i, 1)
+ cute.arch.mbarrier_init(p_full + i, 128)
+ cute.arch.mbarrier_init(o_full, 1)
+ cute.arch.mbarrier_init(o_final, 1)
+ cute.arch.mbarrier_init(qs_full, 128)
+ # One arrival a hand-back: softmax thread 0, after the softmax warps are done with it.
+ cute.arch.mbarrier_init(qs_empty, 1)
+ for i in cutlass.range_constexpr(self.sched_slots):
+ cute.arch.mbarrier_init(sch_full + i, 1)
+ # One elected lane of each of the 9 warps releases an item.
+ cute.arch.mbarrier_init(sch_empty + i, 9)
+ cute.arch.mbarrier_init_fence()
+
+ tmem_users = pipeline.NamedBarrier(barrier_id=1, num_threads=32 * 5)
+ if warp_idx == self.mma_warp:
+ cute.arch.alloc_tmem(self.tmem_cols, storage.tmem_holding.ptr)
+ cute.arch.relinquish_tmem_alloc_permit()
+ cute.arch.sync_threads()
+
+ # Slices taken before the role branches: the MMA warp's `set(ACCUMULATE)` rebinds the
+ # tiled MMAs inside its branch, which the other branches must not see.
+ qk_thr = qk_mma.get_slice(0)
+ pv_thr = pv_mma.get_slice(0)
+ sK = storage.sKV.get_tensor(k_layout.outer, swizzle=k_layout.inner)
+ sV = cute.make_tensor(cute.recast_ptr(sK.iterator, v_layout.inner), v_layout.outer)
+ sQ = storage.sQ.data_ptr()
+ sO = storage.sQ.get_tensor(
+ cute.make_layout(
+ (64, self.group_size, self.rows_per_tile, self.d // 64),
+ stride=(1, 64, 64 * self.group_size, 64 * self.group_size * self.rows_per_tile),
+ ),
+ swizzle=cute.make_swizzle(3, 4, 3),
+ )
+
+ if warp_idx >= self.load_warp0:
+ self.load_role(
+ tidx - 32 * self.load_warp0,
+ sK,
+ sQ,
+ q,
+ k,
+ v,
+ slots,
+ plan,
+ work,
+ rows,
+ num_groups,
+ num_kv_heads,
+ q_stride_row,
+ k_stride_page,
+ k_stride_head,
+ k_stride_token,
+ v_stride_page,
+ v_stride_head,
+ v_stride_token,
+ slots_stride_group,
+ kv_full,
+ kv_empty,
+ qs_full,
+ qs_empty,
+ sch_full,
+ sch_empty,
+ sch_val,
+ )
+ elif warp_idx == self.mma_warp:
+ tmem_ptr = cute.arch.retrieve_tmem_ptr(Float32, 16, storage.tmem_holding.ptr)
+ self.mma_role(
+ qk_mma,
+ qk_mma_acc,
+ pv_mma,
+ pv_mma_acc,
+ qk_thr,
+ pv_thr,
+ q_layout,
+ p_layout,
+ sK,
+ sV,
+ tmem_ptr,
+ plan,
+ num_kv_heads,
+ q_full,
+ kv_full,
+ kv_empty,
+ s_full,
+ p_full,
+ o_full,
+ o_final,
+ sch_full,
+ sch_empty,
+ sch_val,
+ )
+ tmem_users.arrive_and_wait()
+ cute.arch.dealloc_tmem(tmem_ptr, self.tmem_cols)
+ else:
+ tmem_ptr = cute.arch.retrieve_tmem_ptr(Float32, 16, storage.tmem_holding.ptr)
+ self.softmax_role(
+ qk_thr,
+ pv_thr,
+ tmem_ptr,
+ tidx,
+ plan,
+ num_kv_heads,
+ rows,
+ bits,
+ bits_stride_row,
+ out,
+ out_stride_row,
+ scale_log2,
+ q_full,
+ s_full,
+ p_full,
+ o_full,
+ o_final,
+ qs_full,
+ qs_empty,
+ sch_full,
+ sch_empty,
+ sch_val,
+ sQ,
+ sO,
+ tma_atom_o,
+ tma_tensor_o,
+ )
+ tmem_users.arrive_and_wait()
+
+ if cutlass.const_expr(self.use_pdl):
+ # The lone-row pass may still run: finish only after it has, so a later launch sees every row.
+ cute.arch.griddepcontrol_wait()
+
+ # ------------------------------------------------------------------ schedule
+ @cute.jit
+ def read_item(self, n: Int32, sch_full, sch_val):
+ """Item `n` of this CTA (group * kv_heads + head, or -1 past the last), once published."""
+ slot = n % self.sched_slots
+ cute.arch.mbarrier_wait(sch_full + slot, (n // self.sched_slots) % 2)
+ return cute.make_tensor(sch_val, cute.make_layout((self.sched_slots,)))[slot]
+
+ @cute.jit
+ def release_item(self, n: Int32, sch_empty):
+ with cute.arch.elect_one():
+ cute.arch.mbarrier_arrive(sch_empty + n % self.sched_slots)
+
+ @cute.jit
+ def item_info(self, item: Int32, plan: cute.Pointer, num_kv_heads: Int32):
+ """(group, KV head, tiles, listed) of a published item; tiles = -1 past the last."""
+ group = Int32(0)
+ head = Int32(0)
+ tiles = Int32(-1)
+ listed = Int32(0)
+ if item >= 0:
+ group = item // num_kv_heads
+ head = item - group * num_kv_heads
+ mPlan = cute.make_tensor(plan, cute.make_layout((1 << 30,)))
+ tiles = mPlan[group * 2]
+ listed = mPlan[group * 2 + 1]
+ return group, head, tiles, listed
+
+ # ------------------------------------------------------------------ loaders (warps 5-8)
+ @cute.jit
+ def publish(self, n: Int32, ticket: Int32, num_items: Int32, num_groups: Int32,
+ num_kv_heads: Int32, sch_full, sch_empty, sch_val): # fmt: skip
+ """Item `n` from a ticket: the plan's groups last to first, both heads of a group in turn."""
+ item = Int32(-1)
+ if ticket < num_items:
+ back = ticket // num_kv_heads
+ item = (num_groups - 1 - back) * num_kv_heads + (ticket - back * num_kv_heads)
+ slot = n % self.sched_slots
+ cute.arch.mbarrier_wait(sch_empty + slot, ((n // self.sched_slots) % 2) ^ 1)
+ cute.make_tensor(sch_val, cute.make_layout((self.sched_slots,)))[slot] = item
+ cute.arch.mbarrier_arrive(sch_full + slot)
+
+ @cute.jit
+ def load_role(
+ self,
+ lt: Int32,
+ sK,
+ sQ: cute.Pointer,
+ q: cute.Pointer,
+ k: cute.Pointer,
+ v: cute.Pointer,
+ slots: cute.Pointer,
+ plan: cute.Pointer,
+ work: cute.Pointer,
+ rows: Int32,
+ num_groups: Int32,
+ num_kv_heads: Int32,
+ q_stride_row: Int32,
+ k_stride_page: Int32,
+ k_stride_head: Int32,
+ k_stride_token: Int32,
+ v_stride_page: Int32,
+ v_stride_head: Int32,
+ v_stride_token: Int32,
+ slots_stride_group: Int32,
+ kv_full,
+ kv_empty,
+ qs_full,
+ qs_empty,
+ sch_full,
+ sch_empty,
+ sch_val,
+ ):
+ """Per item: its Q rows into the Q buffer (unless fetched during the previous item), then
+ the ring in the MMA warp's order K0, K1, V0, K2, V1, ..., V(T-1). Thread `lt` gathers
+ tokens 16 (lt // 32) + 4 t + (lt % 32) // 8 (t < 4), one 16-byte chunk (lt % 8) of each of
+ their four 128-byte dim blocks; the slots of tile j + 1 are fetched while tile j's copies
+ issue. Lane 0 of warp 5 publishes item n + 1 when the warps start item n.
+ """
+ D = self.d
+ atom = cute.make_copy_atom(
+ cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), BF16, num_bits_per_copy=128
+ )
+ chunk = lt % 8
+ # Element-coordinate view of the swizzled ring: (8 elements, (chunk, 64-dim block), token, slot).
+ sKV_c = cute.make_tensor(
+ sK.iterator,
+ cute.make_layout(
+ (8, (8, D // 64), self.bn, self.kv_slots),
+ stride=(1, (8, self.bn * 64), 64, self.bn * D),
+ ),
+ )
+ num_items = num_groups * num_kv_heads
+ first_token = 16 * (lt // 32) + (lt % 32) // 8
+ ticket = Int32(0)
+ if lt == 0:
+ ticket = cute.arch.atomic_add(work, Int32(1))
+ self.publish(
+ Int32(0), ticket, num_items, num_groups, num_kv_heads, sch_full, sch_empty, sch_val
+ )
+ ticket = cute.arch.atomic_add(work, Int32(1))
+ n = Int32(0)
+ ring = Int32(0)
+ ring_phase = Int32(1)
+ qs_phase = Int32(1)
+ q_ready = Int32(0)
+ item = self.read_item(n, sch_full, sch_val)
+ while item >= 0:
+ if lt == 0:
+ self.publish(
+ n + 1, ticket, num_items, num_groups, num_kv_heads, sch_full, sch_empty, sch_val
+ )
+ ticket = cute.arch.atomic_add(work, Int32(1))
+ group, kv_head, tiles, listed = self.item_info(item, plan, num_kv_heads)
+ nitem = self.read_item(n + 1, sch_full, sch_val)
+ ngroup, nhead, ntiles, _ = self.item_info(nitem, plan, num_kv_heads)
+ if tiles > 0:
+ if q_ready == 0:
+ cute.arch.mbarrier_wait(qs_empty, qs_phase)
+ qs_phase = qs_phase ^ 1
+ self.issue_q(lt, group, kv_head, rows, q, q_stride_row, sQ)
+ cute.arch.cp_async_mbarrier_arrive_noinc(qs_full)
+ q_ready = Int32(0)
+ q_at = cutlass.max(tiles - self.q_lead, Int32(0))
+ mSlots = cute.make_tensor(
+ slots + Int64(group) * slots_stride_group, cute.make_layout((1 << 30,))
+ )
+ slot_next = cute.make_rmem_tensor((4,), Int32)
+ k_cur = cute.make_rmem_tensor((4,), Int64)
+ v_cur = cute.make_rmem_tensor((4,), Int64)
+ v_prev = cute.make_rmem_tensor((4,), Int64)
+ for t in cutlass.range_constexpr(4):
+ slot_next[t] = mSlots[cutlass.min(first_token + 4 * t, listed - 1)]
+ for j in cutlass.range(0, tiles, 1, unroll=1):
+ for t in cutlass.range_constexpr(4):
+ slot = slot_next[t]
+ page = slot // self.tpb
+ tok = slot % self.tpb
+ k_cur[t] = (
+ Int64(page) * k_stride_page
+ + kv_head * k_stride_head
+ + tok * k_stride_token
+ )
+ v_cur[t] = (
+ Int64(page) * v_stride_page
+ + kv_head * v_stride_head
+ + tok * v_stride_token
+ )
+ if j + 1 < tiles:
+ for t in cutlass.range_constexpr(4):
+ slot_next[t] = mSlots[
+ cutlass.min((j + 1) * self.bn + first_token + 4 * t, listed - 1)
+ ]
+ # K(j)
+ cute.arch.mbarrier_wait(kv_empty + ring, ring_phase)
+ self.issue_tile(atom, k, k_cur, sKV_c, ring, chunk, first_token)
+ cute.arch.cp_async_mbarrier_arrive_noinc(kv_full + ring)
+ ring, ring_phase = self.advance(ring, ring_phase)
+ if j > 0:
+ # V(j - 1)
+ cute.arch.mbarrier_wait(kv_empty + ring, ring_phase)
+ self.issue_tile(atom, v, v_prev, sKV_c, ring, chunk, first_token)
+ cute.arch.cp_async_mbarrier_arrive_noinc(kv_full + ring)
+ ring, ring_phase = self.advance(ring, ring_phase)
+ for t in cutlass.range_constexpr(4):
+ v_prev[t] = v_cur[t]
+ if j == q_at and ntiles > 0:
+ # The next item's Q, once this item's epilogue no longer needs the buffer.
+ cute.arch.mbarrier_wait(qs_empty, qs_phase)
+ qs_phase = qs_phase ^ 1
+ self.issue_q(lt, ngroup, nhead, rows, q, q_stride_row, sQ)
+ cute.arch.cp_async_mbarrier_arrive_noinc(qs_full)
+ q_ready = Int32(1)
+ # V(tiles - 1)
+ cute.arch.mbarrier_wait(kv_empty + ring, ring_phase)
+ self.issue_tile(atom, v, v_prev, sKV_c, ring, chunk, first_token)
+ cute.arch.cp_async_mbarrier_arrive_noinc(kv_full + ring)
+ ring, ring_phase = self.advance(ring, ring_phase)
+ self.release_item(n, sch_empty)
+ n = n + 1
+ item = nitem
+ # Tail: every slot released, so no load is in flight when the CTA retires.
+ for i in cutlass.range_constexpr(self.kv_slots):
+ cute.arch.mbarrier_wait(kv_empty + ring, ring_phase)
+ ring, ring_phase = self.advance(ring, ring_phase)
+
+ @cute.jit
+ def issue_q(
+ self,
+ lt: Int32,
+ group: Int32,
+ kv_head: Int32,
+ rows: Int32,
+ q: cute.Pointer,
+ q_stride_row: Int32,
+ sQ: cute.Pointer,
+ ):
+ """An item's live Q rows into the Q buffer: MMA row m = local row x G + head at m x D,
+ its 16-byte chunk c at chunk c ^ (m % 8). Consecutive threads take consecutive chunks of
+ a row's G heads (contiguous in q).
+ """
+ G = self.group_size
+ R = self.rows_per_tile
+ D = self.d
+ atom = cute.make_copy_atom(
+ cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), BF16, num_bits_per_copy=128
+ )
+ per_row = G * (D // 8)
+ total = R * per_row
+ for it in cutlass.range_constexpr((total + 127) // 128):
+ x = lt + 128 * it
+ local_row = x // per_row
+ rem = x - local_row * per_row
+ head = rem // (D // 8)
+ c = rem - head * (D // 8)
+ row = group * R + local_row
+ if x < total and row < rows:
+ m = local_row * G + head
+ src = cute.make_tensor(
+ q
+ + cute.assume(
+ Int64(row) * q_stride_row + (kv_head * G + head) * D + c * 8, divby=8
+ ),
+ cute.make_layout((8,)),
+ )
+ dst = cute.make_tensor(
+ sQ + cute.assume(self.box_offset(m, c), divby=8), cute.make_layout((8,))
+ )
+ cute.copy(atom, src, dst)
+
+ @cute.jit
+ def box_offset(self, m, c):
+ """Element offset of row m's 16-byte chunk c (of D / 8) in the Q buffer: the layout of the
+ epilogue's TMA boxes -- box c // 8 (64 dims of G heads x R rows), its 128-byte row m,
+ chunk (c % 8) ^ (m % 8) (SW128) -- so a softmax thread's Q row and its O row occupy the
+ same bytes.
+ """
+ return (
+ (c // 8) * (64 * self.group_size * self.rows_per_tile)
+ + m * 64
+ + (((c % 8) ^ (m % 8)) * 8)
+ )
+
+ @cute.jit
+ def issue_tile(
+ self,
+ atom,
+ src_base: cute.Pointer,
+ bases,
+ sKV_c,
+ ring: Int32,
+ chunk: Int32,
+ first_token: Int32,
+ ):
+ for t in cutlass.range_constexpr(4):
+ n = first_token + 4 * t
+ for b in cutlass.range_constexpr(self.d // 64):
+ offset = cute.assume(bases[t] + b * 64 + chunk * 8, divby=8)
+ src = cute.make_tensor(src_base + offset, cute.make_layout((8,)))
+ dst = sKV_c[(None, (chunk, b), n, ring)]
+ cute.copy(atom, src, dst)
+
+ # ------------------------------------------------------------------ MMA (warp 4)
+ @cute.jit
+ def mma_role(
+ self,
+ qk_mma: cute.TiledMma,
+ qk_mma_acc: cute.TiledMma,
+ pv_mma: cute.TiledMma,
+ pv_mma_acc: cute.TiledMma,
+ qk_thr,
+ pv_thr,
+ q_layout,
+ p_layout,
+ sK,
+ sV,
+ tmem_ptr,
+ plan: cute.Pointer,
+ num_kv_heads: Int32,
+ q_full,
+ kv_full,
+ kv_empty,
+ s_full,
+ p_full,
+ o_full,
+ o_final,
+ sch_full,
+ sch_empty,
+ sch_val,
+ ):
+ """Per item: QK(0), then QK(j) / PV(j - 1) pairs, then the last PV. S/P buffer b sits at
+ TMEM columns tmem_s0 + 64 b, b = the CTA's tile count mod 2; QK(t) overwrites S[t % 2] only
+ after PV(t - 2) was issued (tcgen05 MMAs of one thread execute in order), and P(t - 2) was
+ its last reader. An item's PV(0) overwrites O: it waits on P(0), which the softmax warps
+ publish only after reading the previous item's O out.
+ """
+ tQ = cute.make_tensor(tmem_ptr + self.tmem_q, q_layout.outer)
+ tSrQ = qk_thr.make_fragment_A(tQ)[None, None, None, 0]
+ tSrQ = cute.make_tensor(cute.recast_ptr(tmem_ptr + self.tmem_q, dtype=BF16), tSrQ.layout)
+ tSrK = qk_thr.make_fragment_B(sK)
+ tOrV = pv_thr.make_fragment_B(sV)
+ tStS = qk_thr.make_fragment_C(qk_thr.partition_shape_C((self.bm, self.bn)))
+ tOtO_f = pv_thr.make_fragment_C(pv_thr.partition_shape_C((self.bm, self.d)))
+ tOtO = cute.make_tensor(tmem_ptr + self.tmem_o, tOtO_f.layout)
+ tP = cute.make_tensor(tmem_ptr + self.tmem_s0, p_layout.outer)
+ tOrP = pv_thr.make_fragment_A(tP)[None, None, None, 0]
+ n = Int32(0)
+ g = Int32(0)
+ nq = Int32(0)
+ ring = Int32(0)
+ ring_phase = Int32(0)
+ item = self.read_item(n, sch_full, sch_val)
+ while item >= 0:
+ group, kv_head, tiles, listed = self.item_info(item, plan, num_kv_heads)
+ if tiles > 0:
+ cute.arch.mbarrier_wait(q_full, nq % 2)
+ tcgen05_fence_after()
+ # QK(0)
+ buf0 = g % 2
+ cute.arch.mbarrier_wait(kv_full + ring, ring_phase)
+ self.gemm_qk(
+ qk_mma,
+ qk_mma_acc,
+ tSrQ,
+ tSrK[None, None, None, ring],
+ cute.make_tensor(tmem_ptr + (self.tmem_s0 + buf0 * self.bn), tStS.layout),
+ )
+ with cute.arch.elect_one():
+ tcgen05.commit(s_full + buf0)
+ tcgen05.commit(kv_empty + ring)
+ ring, ring_phase = self.advance(ring, ring_phase)
+ for j in cutlass.range(1, tiles + 1, 1, unroll=1):
+ t = g + j
+ if j < tiles:
+ # QK(j) into S[t % 2]
+ buf = t % 2
+ cute.arch.mbarrier_wait(kv_full + ring, ring_phase)
+ self.gemm_qk(
+ qk_mma,
+ qk_mma_acc,
+ tSrQ,
+ tSrK[None, None, None, ring],
+ cute.make_tensor(
+ tmem_ptr + (self.tmem_s0 + buf * self.bn), tStS.layout
+ ),
+ )
+ with cute.arch.elect_one():
+ tcgen05.commit(s_full + buf)
+ tcgen05.commit(kv_empty + ring)
+ ring, ring_phase = self.advance(ring, ring_phase)
+ # PV(j - 1)
+ prev = t - 1
+ pbuf = prev % 2
+ cute.arch.mbarrier_wait(kv_full + ring, ring_phase)
+ cute.arch.mbarrier_wait(p_full + pbuf, (prev // 2) % 2)
+ tcgen05_fence_after()
+ tOrP_b = cute.make_tensor(
+ cute.recast_ptr(
+ tmem_ptr + (self.tmem_s0 + pbuf * self.bn), dtype=tOrP.element_type
+ ),
+ tOrP.layout,
+ )
+ if j == 1:
+ self.gemm_pv(pv_mma, pv_mma_acc, tOrP_b, tOrV[None, None, None, ring], tOtO)
+ else:
+ self.gemm_pv(
+ pv_mma_acc, pv_mma_acc, tOrP_b, tOrV[None, None, None, ring], tOtO
+ )
+ with cute.arch.elect_one():
+ tcgen05.commit(kv_empty + ring)
+ tcgen05.commit(o_full)
+ ring, ring_phase = self.advance(ring, ring_phase)
+ with cute.arch.elect_one():
+ tcgen05.commit(o_final)
+ g = g + tiles
+ nq = nq + 1
+ self.release_item(n, sch_empty)
+ n = n + 1
+ item = self.read_item(n, sch_full, sch_val)
+
+ @cute.jit
+ def advance(self, ring: Int32, ring_phase: Int32):
+ ring = ring + 1
+ if ring == self.kv_slots:
+ ring = Int32(0)
+ ring_phase = ring_phase ^ 1
+ return ring, ring_phase
+
+ @cute.jit
+ def gemm_qk(self, first: cute.TiledMma, rest: cute.TiledMma, tSrQ, tSrK, tStS):
+ for kb in cutlass.range_constexpr(cute.size(tSrQ, mode=[2])):
+ if cutlass.const_expr(kb == 0):
+ cute.gemm(first, tStS, tSrQ[None, None, kb], tSrK[None, None, kb], tStS)
+ else:
+ cute.gemm(rest, tStS, tSrQ[None, None, kb], tSrK[None, None, kb], tStS)
+
+ @cute.jit
+ def gemm_pv(self, first: cute.TiledMma, rest: cute.TiledMma, tOrP, tOrV, tOtO):
+ for kb in cutlass.range_constexpr(cute.size(tOrP, mode=[2])):
+ if cutlass.const_expr(kb == 0):
+ cute.gemm(first, tOtO, tOrP[None, None, kb], tOrV[None, None, kb], tOtO)
+ else:
+ cute.gemm(rest, tOtO, tOrP[None, None, kb], tOrV[None, None, kb], tOtO)
+
+ # ------------------------------------------------------------------ softmax + epilogue (warps 0-3)
+ @cute.jit
+ def copy_q(self, m: Int32, live, sQ: cute.Pointer, store_p, tPtP0, tPcP, p_cols: int, q_full):
+ """Row m's Q from the Q buffer into TMEM columns tmem_q.. (the QK MMA's A operand), 64
+ elements a store; dead rows store zeros (their S rows are masked out).
+ """
+ D = self.d
+ for c in cutlass.range_constexpr(D // (2 * p_cols)):
+ rQ = cute.make_rmem_tensor(tPcP.shape, Float32)
+ rQ16 = cute.make_tensor(
+ cute.recast_ptr(rQ.iterator, dtype=BF16), cute.make_layout((2 * p_cols,))
+ )
+ if live:
+ for i in cutlass.range_constexpr(2 * p_cols // 8):
+ cidx = c * (2 * p_cols // 8) + i
+ off = cute.assume(self.box_offset(m, cidx), divby=8)
+ cute.autovec_copy(
+ cute.make_tensor(sQ + off, cute.make_layout((8,))),
+ cute.make_tensor(rQ16.iterator + i * 8, cute.make_layout((8,))),
+ )
+ else:
+ rQ.fill(0.0)
+ cute.copy(
+ store_p,
+ rQ,
+ cute.make_tensor(
+ tPtP0.iterator + (self.tmem_q - self.tmem_s0 + c * p_cols), tPtP0.layout
+ ),
+ )
+ cute.arch.fence_view_async_tmem_store()
+ tcgen05_fence_before()
+ cute.arch.mbarrier_arrive(q_full)
+
+ @cute.jit
+ def softmax_role(
+ self,
+ qk_thr,
+ pv_thr,
+ tmem_ptr,
+ st: Int32,
+ plan: cute.Pointer,
+ num_kv_heads: Int32,
+ rows: Int32,
+ bits: cute.Pointer,
+ bits_stride_row: Int32,
+ out: cute.Pointer,
+ out_stride_row: Int32,
+ scale_log2: Float32,
+ q_full,
+ s_full,
+ p_full,
+ o_full,
+ o_final,
+ qs_full,
+ qs_empty,
+ sch_full,
+ sch_empty,
+ sch_val,
+ sQ: cute.Pointer,
+ sO,
+ tma_atom_o,
+ tma_tensor_o,
+ ):
+ G = self.group_size
+ R = self.rows_per_tile
+ D = self.d
+ tStS = qk_thr.make_fragment_C(qk_thr.partition_shape_C((self.bm, self.bn)))
+ tStS0 = cute.make_tensor(tmem_ptr + self.tmem_s0, tStS.layout)
+ tOtO_f = pv_thr.make_fragment_C(pv_thr.partition_shape_C((self.bm, self.d)))
+ tOtO = cute.make_tensor(tmem_ptr + self.tmem_o, tOtO_f.layout)
+ # The softmax warps' own barrier (hand-backs of the Q buffer, the TMA store's staging).
+ softmax_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=128)
+
+ # S: one 64-column load a tile; P: 32 packed-bf16 columns stored over the S it came from
+ # (or two 16-column halves).
+ cS = cute.make_identity_tensor((self.bm, self.bn))
+ tScS = qk_thr.partition_C(cS)
+ # The load also reduces the row max over the tile's 64 columns (tcgen05.ld.red, SM103): an
+ # upper bound of the row's masked max, which serves as the softmax offset just as well.
+ load_s = tcgen05.make_tmem_copy(
+ cute.make_copy_atom(
+ tcgen05.copy.LdRed32x32bOp(
+ tcgen05.copy.Repetition(32), redOp=tcgen05.TmemLoadRedOp.MAX
+ ),
+ Float32,
+ ),
+ tStS0,
+ )
+ thr_load_s = load_s.get_slice(st)
+ tLtS0 = thr_load_s.partition_S(tStS0)
+ tLcS = thr_load_s.partition_D(tScS)
+ p_cols = self.bn // 2
+ tStP_layout = cute.composition(tStS0.layout, cute.make_layout((self.bm, p_cols)))
+ tStP0 = cute.make_tensor(tStS0.iterator, tStP_layout)
+ tScP = cute.make_tensor(
+ tScS.iterator, cute.composition(tScS.layout, cute.make_layout((self.bm, p_cols)))
+ )
+ store_p = tcgen05.make_tmem_copy(
+ cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(p_cols)), Float32),
+ tStP0,
+ )
+ thr_store_p = store_p.get_slice(st)
+ tPtP0 = thr_store_p.partition_D(tStP0)
+ tPcP = thr_store_p.partition_S(tScP)
+ # O in 32-column chunks (rescale and epilogue).
+ cO = cute.make_identity_tensor((self.bm, self.d))
+ tOcO = pv_thr.partition_C(cO)
+ tOtO_c = cute.make_tensor(
+ tOtO.iterator,
+ cute.composition(tOtO.layout, cute.make_layout((self.bm, self.chunk_cols))),
+ )
+ tOcO_c = cute.make_tensor(
+ tOcO.iterator,
+ cute.composition(tOcO.layout, cute.make_layout((self.bm, self.chunk_cols))),
+ )
+ load_o = tcgen05.make_tmem_copy(
+ cute.make_copy_atom(
+ tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(self.chunk_cols)), Float32
+ ),
+ tOtO_c,
+ )
+ store_o = tcgen05.make_tmem_copy(
+ cute.make_copy_atom(
+ tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(self.chunk_cols)), Float32
+ ),
+ tOtO_c,
+ )
+ thr_load_o = load_o.get_slice(st)
+ thr_store_o = store_o.get_slice(st)
+ tLtO = thr_load_o.partition_S(tOtO_c)
+ tLcO = thr_load_o.partition_D(tOcO_c)
+ tStO = thr_store_o.partition_D(tOtO_c)
+ n_chunks = self.d // self.chunk_cols
+
+ # The TMA store's view: per item, box b of (group, KV head).
+ gO = cute.flat_divide(tma_tensor_o, (64, G, R))
+ tOsO, tOgO = cpasync.tma_partition(
+ tma_atom_o,
+ 0,
+ cute.make_layout(1),
+ cute.group_modes(sO, 0, 3),
+ cute.group_modes(gO, 0, 3),
+ )
+
+ m = st
+ local_row = m // G
+ head = m % G
+ tau = Float32(self.tau)
+ n = Int32(0)
+ g = Int32(0)
+ nq = Int32(0)
+ qs_phase = Int32(0)
+ q_ready = Int32(0)
+ store_pending = Int32(0)
+ item = self.read_item(n, sch_full, sch_val)
+ while item >= 0:
+ group, kv_head, tiles, listed = self.item_info(item, plan, num_kv_heads)
+ row = group * R + local_row
+ live = (local_row < R) and (row < rows)
+ if tiles > 0:
+ if q_ready == 0:
+ cute.arch.mbarrier_wait(qs_full, qs_phase)
+ qs_phase = qs_phase ^ 1
+ self.copy_q(m, live, sQ, store_p, tPtP0, tPcP, p_cols, q_full)
+ # Every row read out: hand the buffer back.
+ if st < 32:
+ softmax_bar.arrive_and_wait()
+ if st == 0:
+ cute.arch.mbarrier_arrive(qs_empty)
+ else:
+ softmax_bar.arrive()
+ q_ready = Int32(0)
+ safe_row = group * R
+ if live:
+ safe_row = row
+ mBits = cute.make_tensor(
+ bits + Int64(safe_row) * bits_stride_row, cute.make_layout((1 << 30,))
+ )
+ row_max = -Float32.inf
+ row_sum = Float32(0.0)
+ lo_next = Int32(0)
+ hi_next = Int32(0)
+ if live:
+ lo_next = mBits[0]
+ hi_next = mBits[1]
+ for j in cutlass.range(0, tiles, 1, unroll=1):
+ t = g + j
+ lo = lo_next
+ hi = hi_next
+ if live and j + 1 < tiles:
+ lo_next = mBits[2 * j + 2]
+ hi_next = mBits[2 * j + 3]
+ buf = t % 2
+ cute.arch.mbarrier_wait(s_full + buf, (t // 2) % 2)
+ tcgen05_fence_after()
+ rS = cute.make_rmem_tensor(tLcS.shape, Float32)
+ n_ld = cute.size(rS, mode=[1])
+ rMax = cute.make_rmem_tensor(cute.make_layout((1, n_ld)), Float32)
+ tLtS_b = cute.make_tensor(tLtS0.iterator + buf * self.bn, tLtS0.layout)
+ for c in cutlass.range_constexpr(n_ld):
+ cute.copy_atom_call(
+ load_s, tLtS_b[None, c, 0, 0], (rS[None, c, 0, 0], rMax[None, c])
+ )
+ cute.arch.fence_view_async_tmem_load()
+ tile_max = rMax[0]
+ for c in cutlass.range_constexpr(1, n_ld):
+ tile_max = cute.arch.fmax(tile_max, rMax[c])
+ new_max = cute.arch.fmax(row_max, tile_max)
+ use_max = row_max
+ factor = Float32(1.0)
+ if row_max == -Float32.inf:
+ use_max = new_max
+ elif (new_max - row_max) * scale_log2 > tau:
+ use_max = new_max
+ factor = cute.math.exp2((row_max - new_max) * scale_log2, fastmath=True)
+ need = factor != Float32(1.0)
+ if cute.arch.vote_any_sync(need):
+ # O holds PV(0..j-1): wait for PV(j-1), then rescale this warp's rows.
+ cute.arch.mbarrier_wait(o_full, (t - 1) % 2)
+ tcgen05_fence_after()
+ for c in cutlass.range_constexpr(n_chunks):
+ rO = cute.make_rmem_tensor(tLcO.shape, Float32)
+ cute.copy(
+ load_o,
+ cute.make_tensor(tLtO.iterator + c * self.chunk_cols, tLtO.layout),
+ rO,
+ )
+ for i in cutlass.range_constexpr(cute.size(rO)):
+ rO[i] = rO[i] * factor
+ cute.copy(
+ store_o,
+ rO,
+ cute.make_tensor(tStO.iterator + c * self.chunk_cols, tStO.layout),
+ )
+ cute.arch.fence_view_async_tmem_store()
+ row_max = use_max
+ neg = -use_max * scale_log2
+ rP = cute.make_rmem_tensor(tPcP.shape, Float32)
+ rP16 = cute.make_tensor(cute.recast_ptr(rP.iterator, dtype=BF16), rS.layout)
+ # p = 2^(s * scale - max * scale) on the row's own entries, 0 elsewhere; four
+ # partial sums.
+ acc0 = (Float32(0.0), Float32(0.0))
+ acc1 = (Float32(0.0), Float32(0.0))
+ for i in cutlass.range_constexpr(0, self.bn, 4):
+ word = lo if i < 32 else hi
+ e = cute.make_rmem_tensor((4,), Float32)
+ e[0], e[1] = cute.arch.fma_packed_f32x2(
+ (rS[i], rS[i + 1]), (scale_log2, scale_log2), (neg, neg)
+ )
+ e[2], e[3] = cute.arch.fma_packed_f32x2(
+ (rS[i + 2], rS[i + 3]), (scale_log2, scale_log2), (neg, neg)
+ )
+ for u in cutlass.range_constexpr(4):
+ x = cute.math.exp2(e[u], fastmath=True)
+ rS[i + u] = x if (word & (1 << ((i + u) % 32))) != 0 else Float32(0.0)
+ acc0 = cute.arch.add_packed_f32x2(acc0, (rS[i], rS[i + 1]))
+ acc1 = cute.arch.add_packed_f32x2(acc1, (rS[i + 2], rS[i + 3]))
+ row_sum = row_sum * factor + ((acc0[0] + acc0[1]) + (acc1[0] + acc1[1]))
+ rP16.store(rS.load().to(BF16))
+ cute.copy(
+ store_p, rP, cute.make_tensor(tPtP0.iterator + buf * self.bn, tPtP0.layout)
+ )
+ cute.arch.fence_view_async_tmem_store()
+ tcgen05_fence_before()
+ cute.arch.mbarrier_arrive(p_full + buf)
+ if j == 0 and store_pending != 0:
+ # The previous item's TMA store has read the Q buffer: hand it back.
+ if st == 0:
+ cute.arch.cp_async_bulk_wait_group(0, read=True)
+ cute.arch.mbarrier_arrive(qs_empty)
+ store_pending = Int32(0)
+
+ # The last QK of this item has completed (its S arrived): the next item's Q goes
+ # into tensor memory now, so its first QK MMAs overlap this epilogue.
+ nitem = self.read_item(n + 1, sch_full, sch_val)
+ ngroup, nhead, ntiles, _ = self.item_info(nitem, plan, num_kv_heads)
+ next_q = ntiles > 0
+ if next_q:
+ nrow = ngroup * R + local_row
+ nlive = (local_row < R) and (nrow < rows)
+ cute.arch.mbarrier_wait(qs_full, qs_phase)
+ qs_phase = qs_phase ^ 1
+ self.copy_q(m, nlive, sQ, store_p, tPtP0, tPcP, p_cols, q_full)
+ q_ready = Int32(1)
+ # The Q buffer is free for staging O unless a later item still has to load its Q
+ # there (a following lone or empty group): then each thread stores its own row.
+ staged = next_q or nitem < 0
+
+ # Epilogue: O / l in bf16.
+ cute.arch.mbarrier_wait(o_final, nq % 2)
+ tcgen05_fence_after()
+ inv = Float32(0.0)
+ if row_sum > Float32(0.0):
+ inv = Float32(1.0) / row_sum
+ if staged:
+ if nitem < 0:
+ # No next Q orders this staging after the previous item's TMA store read
+ # the buffer (thread 0 waited for it at this item's first tile): sync.
+ softmax_bar.arrive_and_wait()
+ for c in cutlass.range_constexpr(n_chunks):
+ rO = cute.make_rmem_tensor(tLcO.shape, Float32)
+ cute.copy(
+ load_o,
+ cute.make_tensor(tLtO.iterator + c * self.chunk_cols, tLtO.layout),
+ rO,
+ )
+ for i in cutlass.range_constexpr(cute.size(rO)):
+ rO[i] = rO[i] * inv
+ rB = cute.make_rmem_tensor((self.chunk_cols,), BF16)
+ rB.store(rO.load().to(BF16))
+ if m < G * R:
+ for i in cutlass.range_constexpr(self.chunk_cols // 8):
+ offset = cute.assume(
+ self.box_offset(m, c * (self.chunk_cols // 8) + i), divby=8
+ )
+ cute.autovec_copy(
+ cute.make_tensor(rB.iterator + i * 8, cute.make_layout((8,))),
+ cute.make_tensor(sQ + offset, cute.make_layout((8,))),
+ )
+ cute.arch.fence_view_async_shared()
+ if st < 32:
+ softmax_bar.arrive_and_wait()
+ if st == 0:
+ for b in cutlass.range_constexpr(self.d // 64):
+ cute.copy(tma_atom_o, tOsO[None, b], tOgO[None, b, kv_head, group])
+ cute.arch.cp_async_bulk_commit_group()
+ else:
+ softmax_bar.arrive()
+ store_pending = Int32(1)
+ if nitem < 0:
+ if st == 0:
+ cute.arch.cp_async_bulk_wait_group(0, read=True)
+ store_pending = Int32(0)
+ else:
+ for c in cutlass.range_constexpr(n_chunks):
+ rO = cute.make_rmem_tensor(tLcO.shape, Float32)
+ cute.copy(
+ load_o,
+ cute.make_tensor(tLtO.iterator + c * self.chunk_cols, tLtO.layout),
+ rO,
+ )
+ for i in cutlass.range_constexpr(cute.size(rO)):
+ rO[i] = rO[i] * inv
+ rB = cute.make_rmem_tensor((self.chunk_cols,), BF16)
+ rB.store(rO.load().to(BF16))
+ if live:
+ for i in cutlass.range_constexpr(self.chunk_cols // 8):
+ dst_off = cute.assume(
+ Int64(row) * out_stride_row
+ + (kv_head * G + head) * D
+ + c * self.chunk_cols
+ + i * 8,
+ divby=8,
+ )
+ cute.autovec_copy(
+ cute.make_tensor(rB.iterator + i * 8, cute.make_layout((8,))),
+ cute.make_tensor(out + dst_off, cute.make_layout((8,))),
+ )
+ g = g + tiles
+ nq = nq + 1
+ elif tiles == 0:
+ # Nothing resident for this group: zeros.
+ if live:
+ rZ = cute.make_rmem_tensor((8,), BF16)
+ rZ.fill(0.0)
+ for i in cutlass.range_constexpr(D // 8):
+ dst_off = cute.assume(
+ Int64(row) * out_stride_row + (kv_head * G + head) * D + i * 8,
+ divby=8,
+ )
+ cute.autovec_copy(
+ rZ, cute.make_tensor(out + dst_off, cute.make_layout((8,)))
+ )
+ self.release_item(n, sch_empty)
+ n = n + 1
+ item = self.read_item(n, sch_full, sch_val)
+
+
+@functools.lru_cache(maxsize=None)
+def _compiled(group_size: int, rows_per_tile: int, tokens_per_block: int, use_pdl: bool):
+ kernel = _RowtileAttentionKernel(
+ group_size, rows_per_tile, HEAD_DIM, tokens_per_block, kv_slots=KV_SLOTS, use_pdl=use_pdl
+ )
+ gmem = cute.AddressSpace.gmem
+
+ def fake(dtype, align):
+ return make_ptr(dtype, 16 * 1024, gmem, assumed_align=align)
+
+ args = (
+ fake(BF16, 16),
+ fake(BF16, 16),
+ fake(BF16, 16),
+ fake(Int32, 4),
+ fake(Int32, 4),
+ fake(Int32, 4),
+ fake(BF16, 16),
+ fake(Int32, 4),
+ 1,
+ 1,
+ 1,
+ 1,
+ 8,
+ 8,
+ 8,
+ 8,
+ 8,
+ 8,
+ 8,
+ 1,
+ 1,
+ 8,
+ 1.0,
+ cuda_drv.CUstream(0),
+ )
+ return cute.compile(kernel, *args)
+
+
+@functools.lru_cache(maxsize=None)
+def _sm_count(device_index: int) -> int:
+ return torch.cuda.get_device_properties(device_index).multi_processor_count
+
+
+def num_ctas(num_items: int, device: torch.device) -> int:
+ """The persistent grid: one CTA per SM, never more than there are items."""
+ return max(1, min(num_items, _sm_count(device.index if device.index is not None else 0)))
+
+
+def _aligned(t: torch.Tensor, strides: int) -> bool:
+ return t.data_ptr() % 16 == 0 and all(t.stride(i) % 8 == 0 for i in range(strides))
+
+
+def eligible(q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor) -> bool:
+ """Whether a launch fits this pass: SM103, bf16 q and K/V, D = 256, unit-stride head
+ dimension and 16-byte-aligned rows.
+ """
+ if q.shape[2] != HEAD_DIM:
+ return False
+ if not (q.dtype == k_cache.dtype == v_cache.dtype == torch.bfloat16):
+ return False
+ if torch.cuda.get_device_capability(q.device) != (10, 3):
+ return False
+ if (
+ q.stride(2) != 1
+ or q.stride(1) != HEAD_DIM
+ or k_cache.stride(3) != 1
+ or v_cache.stride(3) != 1
+ ):
+ return False
+ return _aligned(q, 1) and _aligned(k_cache, 3) and _aligned(v_cache, 3)
+
+
+class _Launcher:
+ """`rowtile_attention_cute[grid](...)`: grid-indexed like the entry's Triton launches;
+ the grid is `(ctas,)`, the persistent CTAs (`num_ctas`).
+ """
+
+ def __getitem__(self, grid):
+ return functools.partial(self._launch, tuple(grid))
+
+ @staticmethod
+ def _launch(
+ grid,
+ q,
+ k_cache,
+ v_cache,
+ slots,
+ packed_bits,
+ group_plan,
+ output,
+ work,
+ softmax_scale,
+ rows,
+ *,
+ group_size,
+ rows_per_tile,
+ tokens_per_block,
+ use_pdl,
+ ):
+ assert len(grid) == 1, grid
+ assert (
+ output.is_contiguous()
+ and packed_bits.dtype == torch.int64
+ and slots.dtype == torch.int32
+ and work.dtype == torch.int32
+ )
+ fn = _compiled(group_size, rows_per_tile, tokens_per_block, bool(use_pdl))
+ gmem = cute.AddressSpace.gmem
+
+ def ptr(t, dtype, align):
+ return make_ptr(dtype, t.data_ptr(), gmem, assumed_align=align)
+
+ fn(
+ ptr(q, BF16, 16),
+ ptr(k_cache, BF16, 16),
+ ptr(v_cache, BF16, 16),
+ ptr(slots, Int32, 4),
+ ptr(packed_bits, Int32, 4),
+ ptr(group_plan, Int32, 4),
+ ptr(output, BF16, 16),
+ ptr(work, Int32, 4),
+ rows,
+ group_plan.shape[0],
+ k_cache.shape[1],
+ grid[0],
+ q.stride(0),
+ k_cache.stride(0),
+ k_cache.stride(1),
+ k_cache.stride(2),
+ v_cache.stride(0),
+ v_cache.stride(1),
+ v_cache.stride(2),
+ slots.stride(0),
+ 2 * packed_bits.stride(0),
+ output.stride(0),
+ softmax_scale * _LOG2_E,
+ cuda_drv.CUstream(torch.cuda.current_stream(q.device).cuda_stream),
+ )
+
+
+rowtile_attention_cute = _Launcher()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_sparse_gqa_rowtile.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_sparse_gqa_rowtile.py
new file mode 100644
index 000000000000..a7b5ad5e6a8e
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_sparse_gqa_rowtile.py
@@ -0,0 +1,938 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""QSA sparse GQA for context batches: several query rows per program, K/V read once per tile.
+
+The first launch turns each row's selection (the columns the prefill bound
+lets it read) into a bitmap over logical token positions. Three plan launches
+form, per group of consecutive rows of one request, the union of its rows'
+resident selected tokens as a packed list of physical cache slots, and each
+row's bitmap re-indexed onto that list. The row-tile launch is persistent: one
+program per SM takes (group, KV head) items longest first and walks each packed
+list in 64-slot tiles, each tile's K/V loaded once for every row and head of the
+group, each row masked to its own bits. A group whose rows span requests
+(generation rows of a mixed batch) walks each row's own selected columns
+instead, as the per-row kernel does.
+"""
+
+from __future__ import annotations
+
+import math
+
+import cutlass
+import cutlass.cute as cute
+import torch
+import triton
+import triton.language as tl
+from cutlass._mlir.dialects import llvm
+from cutlass.cute.runtime import from_dlpack
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa import (
+ qsa_paged_sparse_gqa as _per_row_entry,
+)
+from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
+
+from .qsa_rowtile_attention_cute import BLOCK_M, BLOCK_N
+from .qsa_rowtile_attention_cute import eligible as _cute_eligible
+from .qsa_rowtile_attention_cute import num_ctas as _rowtile_ctas
+from .qsa_rowtile_attention_cute import rowtile_attention_cute as _rowtile_attention_cute
+
+_INTEGER_DTYPES = (torch.int32, torch.int64)
+# The per-row kernel's scan tile: the prefill bound is rounded up to it, so the
+# bitmap reads exactly the columns that kernel reads.
+_BOUND_TILE = 64
+# Programs of the lone-row pass per KV head; each walks the lone list with this stride.
+_LONE_PROGRAMS = 256 # capacity: fixed: a launch width; the lone list may be any length
+# Bitmap words one program of the plan's union and slot passes covers (a grid dimension).
+_CHUNK_WORDS = 32 # capacity: fixed: a launch tile; the grid covers every word
+_PLAN_WARPS = 4
+_LONE_WARPS = 2
+_LONE_STAGES = 2
+
+
+# --------------------------------------------------------------------------------------
+# The bitmap pass (CuTe DSL, Tier 6). Row `r`: one bit per readable selected token (a
+# column below the per-row kernel's prefill bound naming a token >= 0 whose page is inside
+# the table) set in `bits[r]`, and the highest such token and the number of columns read
+# in `row_info[r]`. A CTA owns one row: the row's selection is staged in shared memory with
+# coalesced loads, each thread walks `_BITS_RUN` consecutive columns, ORs the bits of one
+# 64-token word in a register and flushes the word into a shared-memory bitmap with one
+# `red.shared.or.b64` when the run leaves it; then the CTA writes every word of the row and
+# `row_info`. Any column order gives the same words (OR is idempotent and commutative); the
+# order only decides how few flushes a run makes, and the selection arrives in token order.
+# No memset of the bitmap is needed.
+# --------------------------------------------------------------------------------------
+
+_BITS_THREADS = 128
+#: Consecutive selected columns a thread walks.
+_BITS_RUN = 16
+
+
+@dsl_user_op
+def _red_or_shared(addr, val, *, loc=None, ip=None):
+ llvm.inline_asm(
+ T.i32(),
+ [cutlass.Int32(addr).ir_value(loc=loc, ip=ip), cutlass.Int64(val).ir_value(loc=loc, ip=ip)],
+ "{ red.shared.or.b64 [$1], $2; mov.u32 $0, 0; }",
+ "=r,r,l",
+ has_side_effects=True,
+ is_align_stack=False,
+ asm_dialect=llvm.AsmDialect.AD_ATT,
+ loc=loc,
+ ip=ip,
+ )
+
+
+class _RowBitsKernel:
+ """The bitmap pass's `bits` and `row_info`, one CTA a row."""
+
+ def __init__(self, topk, block_topk, compress_ratio, tokens_per_block, table_width, num_words):
+ self.topk, self.block_topk, self.cr = topk, block_topk, compress_ratio
+ self.tpb, self.width, self.num_words = tokens_per_block, table_width, num_words
+ self.staged = topk + topk // _BITS_RUN + 1
+
+ @cute.jit
+ def __call__(
+ self,
+ sel_t: cute.Tensor,
+ qp_ptr: cute.Pointer,
+ bits_ptr: cute.Pointer,
+ info_ptr: cute.Pointer,
+ rows: cutlass.Int32,
+ stream,
+ ):
+ qp = cute.make_tensor(qp_ptr, cute.make_layout((rows,), stride=(1,)))
+ bits = cute.make_tensor(
+ bits_ptr, cute.make_layout((rows, self.num_words), stride=(self.num_words, 1))
+ )
+ info = cute.make_tensor(info_ptr, cute.make_layout((rows, 2), stride=(2, 1)))
+ # `sel_t` arrives as a tensor: the TVM-FFI entry takes its stream from a GPU tensor.
+ self.kernel(sel_t, qp, bits, info).launch(
+ grid=[rows, 1, 1],
+ block=[_BITS_THREADS, 1, 1],
+ smem=self.num_words * 8 + self.staged * 4 + 128,
+ stream=stream,
+ )
+
+ @cute.kernel
+ def kernel(self, mSel: cute.Tensor, mQp: cute.Tensor, mBits: cute.Tensor, mInfo: cute.Tensor):
+ tidx, _, _ = cute.arch.thread_idx()
+ row, _, _ = cute.arch.block_idx()
+ smem = cutlass.utils.SmemAllocator()
+ words = smem.allocate_tensor(
+ cutlass.Int64, cute.make_layout((self.num_words,), stride=(1,)), byte_alignment=16
+ )
+ red = smem.allocate_tensor(
+ cutlass.Int32, cute.make_layout((4,), stride=(1,)), byte_alignment=16
+ )
+ # The row's selection, one padding slot every `_BITS_RUN` columns (conflict-free runs).
+ sel_s = smem.allocate_tensor(
+ cutlass.Int32, cute.make_layout((self.staged,), stride=(1,)), byte_alignment=16
+ )
+ words_base = cutlass.Int32(words.iterator.toint())
+ for i in cutlass.range_constexpr((self.num_words + _BITS_THREADS - 1) // _BITS_THREADS):
+ w = tidx + i * _BITS_THREADS
+ if w < self.num_words:
+ words[w] = cutlass.Int64(0)
+
+ # The columns the per-row kernel reads: the prefill bound rounded up to its tile.
+ visible = cutlass.Int32(mQp[row]) + 1
+ selected_blocks = cutlass.min(visible // self.cr, cutlass.Int32(self.block_topk))
+ valid = selected_blocks * self.cr + visible % self.cr
+ bounded = cutlass.min(valid, cutlass.Int32(self.topk))
+ columns_read = cutlass.min(
+ (bounded + _BOUND_TILE - 1) // _BOUND_TILE * _BOUND_TILE, cutlass.Int32(self.topk)
+ )
+ columns_read = cutlass.max(columns_read, cutlass.Int32(0))
+ for i in cutlass.range_constexpr((self.topk + _BITS_THREADS - 1) // _BITS_THREADS):
+ c = tidx + i * _BITS_THREADS
+ if c < columns_read:
+ sel_s[c + c // _BITS_RUN] = cutlass.Int32(mSel[row, c])
+ cute.arch.sync_threads()
+
+ highest = cutlass.Int32(-1)
+ runs = (self.topk + _BITS_THREADS * _BITS_RUN - 1) // (_BITS_THREADS * _BITS_RUN)
+ for run in cutlass.range_constexpr(runs):
+ c0 = (run * _BITS_THREADS + tidx) * _BITS_RUN
+ if c0 < columns_read:
+ cur_w = cutlass.Int32(-1)
+ cur_m = cutlass.Int64(0)
+ for k in cutlass.range_constexpr(_BITS_RUN):
+ c = c0 + k
+ t = cutlass.Int32(-1)
+ if c < columns_read:
+ t = sel_s[c + c // _BITS_RUN]
+ if t >= 0:
+ if t // self.tpb < self.width:
+ w = t >> 6
+ if w != cur_w:
+ if cur_w >= 0:
+ _red_or_shared(words_base + cur_w * 8, cur_m)
+ cur_w = w
+ cur_m = cutlass.Int64(0)
+ cur_m = cur_m | (cutlass.Int64(1) << cutlass.Int64(t & 63))
+ highest = cutlass.max(highest, t)
+ if cur_w >= 0:
+ _red_or_shared(words_base + cur_w * 8, cur_m)
+ for off in cutlass.range_constexpr(5):
+ highest = cutlass.max(highest, cute.arch.shuffle_sync_bfly(highest, offset=1 << off))
+ if tidx % 32 == 0:
+ red[tidx // 32] = highest
+ cute.arch.sync_threads()
+ if tidx == 0:
+ mInfo[row, 0] = cutlass.max(cutlass.max(red[0], red[1]), cutlass.max(red[2], red[3]))
+ mInfo[row, 1] = columns_read
+ for i in cutlass.range_constexpr((self.num_words + _BITS_THREADS - 1) // _BITS_THREADS):
+ w = tidx + i * _BITS_THREADS
+ if w < self.num_words:
+ mBits[row, w] = words[w]
+
+
+_ROW_BITS: dict = {}
+
+
+def _row_bits(
+ selected_tokens, query_positions, compress_ratio, tokens_per_block, table_width, num_words
+):
+ """`(bits [rows, num_words] int64, row_info [rows, 2] int32)` of the bitmap pass over a
+ selection with unit column stride."""
+ rows, final_topk = selected_tokens.shape
+ block_topk = (final_topk - compress_ratio + 1) // compress_ratio
+ device = selected_tokens.device
+ row_info = torch.empty((rows, 2), dtype=torch.int32, device=device)
+ query_positions = query_positions.contiguous()
+ bits = torch.empty((rows, num_words), dtype=torch.int64, device=device)
+ key = (
+ final_topk, compress_ratio, tokens_per_block, table_width,
+ query_positions.dtype, selected_tokens.dtype,
+ ) # fmt: skip
+ args = (selected_tokens, query_positions.data_ptr(), bits.data_ptr(), row_info.data_ptr())
+ fn = _ROW_BITS.get(key)
+ if fn is None:
+ gm = cute.AddressSpace.gmem
+ qty = cutlass.Int64 if query_positions.dtype == torch.int64 else cutlass.Int32
+ fn = cute.compile(
+ _RowBitsKernel(
+ final_topk, block_topk, compress_ratio, tokens_per_block, table_width, num_words
+ ),
+ from_dlpack(selected_tokens, assumed_align=4).mark_layout_dynamic(leading_dim=1),
+ make_ptr(qty, args[1], gm, assumed_align=4),
+ make_ptr(cutlass.Int64, args[2], gm, assumed_align=16),
+ make_ptr(cutlass.Int32, args[3], gm, assumed_align=8),
+ cutlass.Int32(1),
+ cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+ options="--opt-level 2 --enable-tvm-ffi",
+ )
+ _ROW_BITS[key] = fn
+ fn(*args, rows)
+ return bits, row_info
+
+
+@triton.jit
+def _or_combine(a, b):
+ return a | b
+
+
+@triton.jit
+def _compress_bits(x, m):
+ """Per word, the bits of `x` at the set positions of `m`, packed to the low
+ end in order (a parallel bit extract; Hacker's Delight `compress`, 64-bit).
+ `x` and `m` are uint64 and broadcast together."""
+ x = x & m
+ mk = ~m << 1
+ for i in tl.static_range(6):
+ mp = mk ^ (mk << 1)
+ mp = mp ^ (mp << 2)
+ mp = mp ^ (mp << 4)
+ mp = mp ^ (mp << 8)
+ mp = mp ^ (mp << 16)
+ mp = mp ^ (mp << 32)
+ mv = mp & m
+ m = (m ^ mv) | (mv >> (1 << i))
+ t = x & mv
+ x = (x ^ t) | (t >> (1 << i))
+ mk = mk & ~mp
+ return x
+
+
+@triton.jit
+def _popc64(x):
+ """Set bits of each uint64 of `x`, as int32."""
+ return tl.inline_asm_elementwise(
+ "popc.b64 $0, $1;", "=r,l", [x], dtype=tl.int32, is_pure=True, pack=1
+ )
+
+
+@triton.jit
+def _bits_below(n):
+ """uint64 with the low `n` bits set, `0 <= n <= 64`."""
+ ones = (n * 0 - 1).to(tl.int64).to(tl.uint64, bitcast=True)
+ return tl.where(n >= 64, ones, ~(ones << tl.minimum(n, 63).to(tl.uint64)))
+
+
+@triton.jit
+def _group_rows(group, request_indices, row_info, rows, NUM_REQUESTS,
+ ROWS_PER_TILE: tl.constexpr, ROWS_POW2: tl.constexpr): # fmt: skip
+ """A group's members, which of them exist, whether they are of one request (and which),
+ and the bitmap words its union spans (0 for a group of several requests)."""
+ offsets = tl.arange(0, ROWS_POW2)
+ members = group * ROWS_PER_TILE + offsets
+ live = (offsets < ROWS_PER_TILE) & (members < rows)
+ requests = tl.load(request_indices + members, mask=live, other=-1).to(tl.int64)
+ lowest = tl.min(tl.where(live, requests, NUM_REQUESTS), axis=0)
+ highest_request = tl.max(tl.where(live, requests, -1), axis=0)
+ highest_token = tl.max(tl.load(row_info + members * 2, mask=live, other=-1), axis=0)
+ one_request = (lowest == highest_request) & (lowest >= 0) & (lowest < NUM_REQUESTS)
+ words_needed = tl.where(one_request, (highest_token + 64) // 64, 0)
+ return offsets, members, live, one_request, lowest, words_needed
+
+
+@triton.jit
+def _rowtile_union_kernel(
+ request_indices, row_info, bits, block_table, union_words, chunk_counts, rows,
+ bits_stride_row, block_table_stride_request, block_table_stride_page,
+ union_stride_group, counts_stride_group,
+ NUM_REQUESTS, NUM_CACHE_PAGES,
+ ROWS_PER_TILE: tl.constexpr, ROWS_POW2: tl.constexpr,
+ PAGE_TABLE_WIDTH: tl.constexpr, TOKENS_PER_BLOCK: tl.constexpr, CHUNK_WORDS: tl.constexpr,
+ WORD_PAGES: tl.constexpr,
+): # fmt: skip
+ """Per (group, chunk of `CHUNK_WORDS` bitmap words): the union of the group's rows' bits
+ restricted to resident tokens (a mapped page inside the pool), into `union_words`, and
+ its token count into `chunk_counts`. A chunk past the group's highest token, or a group
+ of several requests, counts zero and writes no word."""
+ group = tl.program_id(0).to(tl.int64)
+ chunk = tl.program_id(1)
+ _, members, live, _, lowest, words_needed = _group_rows(
+ group, request_indices, row_info, rows, NUM_REQUESTS, ROWS_PER_TILE, ROWS_POW2
+ )
+ start = chunk * CHUNK_WORDS
+ count = tl.full([], 0, tl.int32)
+ if start < words_needed:
+ words = start + tl.arange(0, CHUNK_WORDS)
+ in_range = words < words_needed
+ row_words = tl.load(
+ bits + members[:, None] * bits_stride_row + words[None, :],
+ mask=live[:, None] & in_range[None, :],
+ other=0,
+ ).to(tl.uint64, bitcast=True) # fmt: skip
+ selected = tl.reduce(row_words, 0, _or_combine)
+ # The pages a word's 64 tokens span (at most WORD_PAGES), each as the word's bits it
+ # covers; a page is looked up only where the union selects one of its tokens.
+ first_token = words.to(tl.int32) * 64
+ logical_pages = first_token[:, None] // TOKENS_PER_BLOCK + tl.arange(0, WORD_PAGES)[None, :]
+ low = tl.maximum(logical_pages * TOKENS_PER_BLOCK - first_token[:, None], 0)
+ high = tl.minimum((logical_pages + 1) * TOKENS_PER_BLOCK - first_token[:, None], 64)
+ page_bits = tl.where(low < high, _bits_below(high) & ~_bits_below(low), 0)
+ page_bits = page_bits & selected[:, None]
+ page_base = block_table + lowest * block_table_stride_request
+ pages = tl.load(
+ page_base + tl.minimum(logical_pages, PAGE_TABLE_WIDTH - 1) * block_table_stride_page,
+ mask=(page_bits != 0) & (logical_pages < PAGE_TABLE_WIDTH),
+ other=-1,
+ ) # fmt: skip
+ kept = (pages >= 0) & (pages < NUM_CACHE_PAGES) & (logical_pages < PAGE_TABLE_WIDTH)
+ resident = tl.reduce(tl.where(kept, page_bits, 0), 1, _or_combine)
+ tl.store(union_words + group * union_stride_group + words,
+ resident.to(tl.int64, bitcast=True), mask=in_range) # fmt: skip
+ count = tl.sum(_popc64(resident), axis=0)
+ tl.store(chunk_counts + group * counts_stride_group + chunk, count)
+
+
+@triton.jit
+def _rowtile_group_kernel(
+ request_indices, row_info, chunk_counts, chunk_bases, group_plan, lone_rows, lone_count,
+ rows, counts_stride_group, num_chunks, NUM_REQUESTS,
+ ROWS_PER_TILE: tl.constexpr, ROWS_POW2: tl.constexpr,
+ BLOCK_N: tl.constexpr, CHUNKS_POW2: tl.constexpr,
+): # fmt: skip
+ """Per group: each chunk's first list position (`chunk_bases`), and `group_plan[group]` =
+ (tiles of the list for the row-tile pass, or -1; list length). A group of several
+ requests appends its rows to the lone list instead."""
+ group = tl.program_id(0).to(tl.int64)
+ offsets, members, live, one_request, _, _ = _group_rows(
+ group, request_indices, row_info, rows, NUM_REQUESTS, ROWS_PER_TILE, ROWS_POW2
+ )
+ chunks = tl.arange(0, CHUNKS_POW2)
+ counts = tl.load(chunk_counts + group * counts_stride_group + chunks,
+ mask=chunks < num_chunks, other=0) # fmt: skip
+ tl.store(chunk_bases + group * counts_stride_group + chunks,
+ tl.cumsum(counts, axis=0) - counts, mask=chunks < num_chunks) # fmt: skip
+ listed = tl.sum(counts, axis=0)
+ tl.store(group_plan + group * 2, tl.where(one_request, tl.cdiv(listed, BLOCK_N), -1))
+ tl.store(group_plan + group * 2 + 1, listed)
+ if not one_request:
+ base = tl.atomic_add(lone_count, tl.sum(live.to(tl.int32), axis=0))
+ tl.store(lone_rows + base + offsets, members.to(tl.int32), mask=live)
+
+
+@triton.jit
+def _rowtile_slots_kernel(
+ request_indices, row_info, bits, block_table, union_words, chunk_counts, chunk_bases,
+ group_plan, slots, packed_bits, rows,
+ bits_stride_row, packed_stride_row, block_table_stride_request, block_table_stride_page,
+ union_stride_group, counts_stride_group, slots_stride_group,
+ NUM_REQUESTS,
+ ROWS_PER_TILE: tl.constexpr, ROWS_POW2: tl.constexpr,
+ TOKENS_PER_BLOCK: tl.constexpr, CHUNK_WORDS: tl.constexpr, WORD_PAGES: tl.constexpr,
+): # fmt: skip
+ """Per (group, chunk) of a row-tile group: the chunk's listed tokens as physical cache
+ slots (`page * tokens_per_block + offset`) at their list positions in `slots[group]`,
+ and each row's bits re-indexed onto the list ORed into `packed_bits[row]` (bit `j` set
+ when the row selects the list's `j`-th token): a per-word parallel bit extract."""
+ group = tl.program_id(0).to(tl.int64)
+ chunk = tl.program_id(1)
+ count = tl.load(chunk_counts + group * counts_stride_group + chunk)
+ if (tl.load(group_plan + group * 2) >= 0) & (count > 0):
+ _, members, live, _, lowest, words_needed = _group_rows(
+ group, request_indices, row_info, rows, NUM_REQUESTS, ROWS_PER_TILE, ROWS_POW2
+ )
+ words = chunk * CHUNK_WORDS + tl.arange(0, CHUNK_WORDS)
+ in_range = words < words_needed
+ resident = tl.load(union_words + group * union_stride_group + words, mask=in_range,
+ other=0).to(tl.uint64, bitcast=True) # fmt: skip
+ word_count = _popc64(resident)
+ first = (tl.load(chunk_bases + group * counts_stride_group + chunk)
+ + tl.cumsum(word_count, axis=0) - word_count) # fmt: skip
+ # The pages the chunk's words span, looked up once a word (as in the union pass).
+ first_token = words.to(tl.int32) * 64
+ first_page = first_token // TOKENS_PER_BLOCK
+ word_pages = first_page[:, None] + tl.arange(0, WORD_PAGES)[None, :]
+ low = tl.maximum(word_pages * TOKENS_PER_BLOCK - first_token[:, None], 0)
+ high = tl.minimum((word_pages + 1) * TOKENS_PER_BLOCK - first_token[:, None], 64)
+ page_bits = tl.where(low < high, _bits_below(high) & ~_bits_below(low), 0)
+ pages_of_word = tl.load(
+ block_table + lowest * block_table_stride_request
+ + word_pages * block_table_stride_page,
+ mask=(page_bits & resident[:, None]) != 0,
+ other=0,
+ ) # fmt: skip
+ lanes = tl.arange(0, 64)
+ lanes64 = lanes.to(tl.uint64)
+ kept = ((resident[:, None] >> lanes64[None, :]) & 1) != 0
+ rank = first[:, None] + _popc64(resident[:, None] & _bits_below(lanes)[None, :])
+ tokens = first_token[:, None] + lanes[None, :]
+ which = tokens // TOKENS_PER_BLOCK - first_page[:, None]
+ pages = tl.zeros([CHUNK_WORDS, 64], tl.int32)
+ for k in tl.static_range(WORD_PAGES):
+ page_k = tl.sum(tl.where(tl.arange(0, WORD_PAGES)[None, :] == k, pages_of_word, 0), 1)
+ pages = tl.where(which == k, page_k[:, None], pages)
+ tl.store(slots + group * slots_stride_group + rank,
+ pages * TOKENS_PER_BLOCK + tokens % TOKENS_PER_BLOCK, mask=kept) # fmt: skip
+ row_words = tl.load(
+ bits + members[:, None] * bits_stride_row + words[None, :],
+ mask=live[:, None] & in_range[None, :],
+ other=0,
+ ).to(tl.uint64, bitcast=True) # fmt: skip
+ packed = _compress_bits(row_words, resident[None, :])
+ shift = (first % 64).to(tl.uint64)
+ destination = packed_bits + members[:, None] * packed_stride_row + (first // 64)[None, :]
+ writes = live[:, None] & in_range[None, :] & (packed != 0)
+ tl.atomic_or(destination, (packed << shift[None, :]).to(tl.int64, bitcast=True),
+ mask=writes, sem="relaxed") # fmt: skip
+ # Bits past the destination word spill into the next one (never for shift 0).
+ spill = (packed >> 1) >> (63 - shift[None, :])
+ tl.atomic_or(destination + 1, spill.to(tl.int64, bitcast=True),
+ mask=writes & ((first % 64) + word_count > 64)[None, :], sem="relaxed") # fmt: skip
+
+
+@triton.jit
+def _rowtile_lone_kernel(
+ q, k_cache, v_cache, block_table, selected_tokens, request_indices, row_info,
+ lone_rows, lone_count, output, softmax_scale, rows,
+ q_stride_row: tl.constexpr, q_stride_head: tl.constexpr, q_stride_dim: tl.constexpr,
+ k_stride_page: tl.constexpr, k_stride_head: tl.constexpr,
+ k_stride_token: tl.constexpr, k_stride_dim: tl.constexpr,
+ v_stride_page: tl.constexpr, v_stride_head: tl.constexpr,
+ v_stride_token: tl.constexpr, v_stride_dim: tl.constexpr,
+ block_table_stride_request: tl.constexpr, block_table_stride_page: tl.constexpr,
+ selected_stride_row: tl.constexpr, selected_stride_token: tl.constexpr,
+ output_stride_row: tl.constexpr, output_stride_head: tl.constexpr,
+ output_stride_dim: tl.constexpr,
+ NUM_CACHE_PAGES, NUM_REQUESTS, LONE_PROGRAMS,
+ PAGE_TABLE_WIDTH: tl.constexpr,
+ GROUP_SIZE: tl.constexpr, TOKENS_PER_BLOCK: tl.constexpr, ROWS_PER_TILE: tl.constexpr,
+ ROWS_POW2: tl.constexpr, BLOCK_H: tl.constexpr, BLOCK_N: tl.constexpr,
+ HEAD_DIM: tl.constexpr, USE_PDL: tl.constexpr,
+ LONE_STAGES: tl.constexpr,
+): # fmt: skip
+ """Rows of groups the row-tile pass does not serve: each its own columns, gathered."""
+ if USE_PDL:
+ # Independent of the row-tile pass (disjoint rows): let it start now.
+ tl.extra.cuda.gdc_launch_dependents()
+ kv_head = tl.program_id(1)
+ for slot in range(tl.program_id(0), tl.load(lone_count), LONE_PROGRAMS):
+ _rowtile_lone_row(
+ q, k_cache, v_cache, block_table, selected_tokens, request_indices, row_info, output,
+ softmax_scale, tl.load(lone_rows + slot).to(tl.int64), kv_head,
+ q_stride_row, q_stride_head, q_stride_dim,
+ k_stride_page, k_stride_head, k_stride_token, k_stride_dim,
+ v_stride_page, v_stride_head, v_stride_token, v_stride_dim,
+ block_table_stride_request, block_table_stride_page,
+ selected_stride_row, selected_stride_token,
+ output_stride_row, output_stride_head, output_stride_dim,
+ NUM_CACHE_PAGES, NUM_REQUESTS, PAGE_TABLE_WIDTH, GROUP_SIZE, TOKENS_PER_BLOCK,
+ BLOCK_H, BLOCK_N, HEAD_DIM, LONE_STAGES,
+ ) # fmt: skip
+
+
+@triton.jit
+def _rowtile_lone_row(
+ q, k_cache, v_cache, block_table, selected_tokens, request_indices, row_info, output,
+ softmax_scale, row, kv_head,
+ q_stride_row, q_stride_head, q_stride_dim,
+ k_stride_page, k_stride_head, k_stride_token, k_stride_dim,
+ v_stride_page, v_stride_head, v_stride_token, v_stride_dim,
+ block_table_stride_request, block_table_stride_page,
+ selected_stride_row, selected_stride_token,
+ output_stride_row, output_stride_head, output_stride_dim,
+ NUM_CACHE_PAGES, NUM_REQUESTS, PAGE_TABLE_WIDTH: tl.constexpr,
+ GROUP_SIZE: tl.constexpr, TOKENS_PER_BLOCK: tl.constexpr, BLOCK_H: tl.constexpr,
+ BLOCK_N: tl.constexpr, HEAD_DIM: tl.constexpr, LONE_STAGES: tl.constexpr,
+): # fmt: skip
+ request = tl.load(request_indices + row).to(tl.int64)
+ valid_request = (request >= 0) & (request < NUM_REQUESTS)
+ safe_request = tl.minimum(tl.maximum(request, 0), NUM_REQUESTS - 1)
+ page_base = block_table + safe_request * block_table_stride_request
+ columns_read = tl.load(row_info + row * 2 + 1)
+ head_offsets = tl.arange(0, BLOCK_H)
+ dim_offsets = tl.arange(0, HEAD_DIM)
+ query_heads = kv_head * GROUP_SIZE + head_offsets
+ head_live = head_offsets < GROUP_SIZE
+ q_tile = tl.load(
+ q + row * q_stride_row + query_heads[:, None] * q_stride_head
+ + dim_offsets[None, :] * q_stride_dim,
+ mask=head_live[:, None],
+ other=0.0,
+ ) # fmt: skip
+ accumulator = tl.zeros([BLOCK_H, HEAD_DIM], tl.float32)
+ running_max = tl.full([BLOCK_H], -1.0e30, tl.float32)
+ running_sum = tl.zeros([BLOCK_H], tl.float32)
+ token_offsets = tl.arange(0, BLOCK_N)
+ for tile in tl.range(0, tl.cdiv(columns_read, BLOCK_N), num_stages=LONE_STAGES):
+ columns = tile * BLOCK_N + token_offsets
+ tokens = tl.load(
+ selected_tokens + row * selected_stride_row + columns * selected_stride_token,
+ mask=columns < columns_read,
+ other=-1,
+ ).to(tl.int32)
+ safe_tokens = tl.maximum(tokens, 0)
+ logical_pages = safe_tokens // TOKENS_PER_BLOCK
+ token_in_page = safe_tokens % TOKENS_PER_BLOCK
+ usable = valid_request & (tokens >= 0) & (logical_pages < PAGE_TABLE_WIDTH)
+ physical_pages = tl.load(
+ page_base + tl.minimum(logical_pages, PAGE_TABLE_WIDTH - 1) * block_table_stride_page,
+ mask=usable,
+ other=-1,
+ ).to(tl.int64)
+ usable = usable & (physical_pages >= 0) & (physical_pages < NUM_CACHE_PAGES)
+ safe_pages = tl.maximum(physical_pages, 0)
+ keys = tl.load(
+ k_cache + safe_pages[None, :] * k_stride_page + kv_head * k_stride_head
+ + token_in_page[None, :] * k_stride_token + dim_offsets[:, None] * k_stride_dim,
+ mask=usable[None, :],
+ other=0.0,
+ ).to(q_tile.dtype) # fmt: skip
+ values = tl.load(
+ v_cache + safe_pages[:, None] * v_stride_page + kv_head * v_stride_head
+ + token_in_page[:, None] * v_stride_token + dim_offsets[None, :] * v_stride_dim,
+ mask=usable[:, None],
+ other=0.0,
+ ).to(q_tile.dtype) # fmt: skip
+ scores = tl.dot(q_tile, keys) * (softmax_scale * 1.4426950408889634)
+ scores = tl.where(usable[None, :], scores, -float("inf"))
+ next_max = tl.maximum(running_max, tl.max(scores, axis=1))
+ correction = tl.math.exp2(running_max - next_max)
+ probabilities = tl.where(usable[None, :], tl.math.exp2(scores - next_max[:, None]), 0.0)
+ accumulator = tl.dot(
+ probabilities.to(q_tile.dtype), values, accumulator * correction[:, None]
+ )
+ running_sum = running_sum * correction + tl.sum(probabilities, axis=1)
+ running_max = next_max
+ normalized = tl.where(
+ running_sum[:, None] > 0, accumulator / tl.maximum(running_sum[:, None], 1.0e-20), 0.0
+ )
+ tl.store(
+ output + row * output_stride_row + query_heads[:, None] * output_stride_head
+ + dim_offsets[None, :] * output_stride_dim,
+ normalized.to(output.dtype.element_ty),
+ mask=head_live[:, None],
+ ) # fmt: skip
+
+
+def _check_inputs(q, k_cache, v_cache, block_table, selected_tokens, request_indices,
+ tokens_per_block, softmax_scale, query_positions, compress_ratio): # fmt: skip
+ if q.ndim != 3:
+ raise ValueError("QSA queries must have shape [rows, query_heads, head_dim]")
+ if k_cache.ndim != 4 or v_cache.shape != k_cache.shape:
+ raise ValueError("QSA K/V caches must have matching HND paged shapes")
+ if block_table.ndim != 2 or block_table.dtype != torch.int32:
+ raise ValueError("QSA sparse GQA requires a two-dimensional int32 block table")
+ if selected_tokens.ndim != 2 or selected_tokens.dtype != torch.int32:
+ raise ValueError("QSA selected tokens must be a two-dimensional int32 tensor")
+ rows, num_q_heads, head_dim = q.shape
+ if selected_tokens.shape[0] != rows or selected_tokens.shape[1] <= 0:
+ raise ValueError("QSA selected-token rows must match non-empty query rows")
+ if request_indices.shape != (rows,) or request_indices.dtype not in _INTEGER_DTYPES:
+ raise ValueError("QSA request indices must have one integer entry per row")
+ if not request_indices.is_contiguous():
+ raise ValueError("QSA request indices must be contiguous")
+ if tokens_per_block <= 0 or k_cache.shape[2] != tokens_per_block:
+ raise ValueError("QSA cache token dimension must match positive tokens_per_block")
+ if k_cache.shape[3] != head_dim or k_cache.dtype != v_cache.dtype:
+ raise ValueError("QSA query and K/V cache head geometry is incompatible")
+ if q.dtype != torch.bfloat16:
+ raise ValueError("QSA queries must use BF16")
+ if k_cache.dtype != torch.bfloat16:
+ raise ValueError("QSA K/V cache must use BF16")
+ if selected_tokens.stride(1) != 1:
+ raise ValueError("QSA selected tokens must have unit column stride")
+ num_kv_heads = k_cache.shape[1]
+ if num_kv_heads <= 0 or num_q_heads <= 0 or num_q_heads % num_kv_heads:
+ raise ValueError("QSA query heads must be divisible by non-zero KV heads")
+ if k_cache.shape[0] <= 0 or block_table.shape[0] <= 0 or block_table.shape[1] <= 0:
+ raise ValueError("QSA sparse GQA requires non-empty cache and block-table geometry")
+ if k_cache.shape[0] * tokens_per_block >= 2**31:
+ raise ValueError("QSA row-tile cache slots must fit int32")
+ if not math.isfinite(softmax_scale) or softmax_scale <= 0:
+ raise ValueError("QSA sparse GQA softmax scale must be finite and positive")
+ if query_positions is None or compress_ratio is None:
+ raise ValueError("the row-tile kernel needs the prefill pair (positions, compress ratio)")
+ if query_positions.shape != (rows,) or query_positions.dtype not in _INTEGER_DTYPES:
+ raise ValueError("QSA query positions must be one integer entry per row")
+ if compress_ratio <= 0:
+ raise ValueError("QSA compression ratio must be positive")
+ token_topk = selected_tokens.shape[1] - compress_ratio + 1
+ if token_topk <= 0 or token_topk % compress_ratio:
+ raise ValueError("QSA selected-token width is incompatible with its compression ratio")
+ devices = {t.device for t in (q, k_cache, v_cache, block_table, selected_tokens,
+ request_indices, query_positions)} # fmt: skip
+ if len(devices) != 1 or next(iter(devices)).type != "cuda":
+ raise ValueError("QSA row-tile sparse GQA needs every tensor on one CUDA device")
+ if num_q_heads // num_kv_heads > BLOCK_M:
+ raise ValueError("a row's query heads per KV head must fit the 128-row tile")
+ if not _cute_eligible(q, k_cache, v_cache):
+ raise ValueError(
+ "the row-tile pass needs SM103, head_dim 256, unit-stride heads (q's heads packed) "
+ "and 16-byte-aligned rows"
+ )
+
+
+def qsa_rowtile_sparse_gqa(
+ *,
+ q: torch.Tensor,
+ k_cache: torch.Tensor,
+ v_cache: torch.Tensor,
+ block_table: torch.Tensor,
+ selected_tokens: torch.Tensor,
+ request_indices: torch.Tensor,
+ tokens_per_block: int,
+ softmax_scale: float,
+ query_positions: torch.Tensor,
+ compress_ratio: int,
+) -> torch.Tensor:
+ """Bitmap pass, plan pass, lone-row pass, row-tile pass (`qsa_rowtile_attention_cute`);
+ returns `[rows, Hq, D]` contiguous."""
+ _check_inputs(q, k_cache, v_cache, block_table, selected_tokens, request_indices,
+ tokens_per_block, softmax_scale, query_positions, compress_ratio) # fmt: skip
+ rows, num_q_heads, head_dim = q.shape
+ num_kv_heads = k_cache.shape[1]
+ group_size = num_q_heads // num_kv_heads
+ output = torch.empty((rows, num_q_heads, head_dim), dtype=q.dtype, device=q.device)
+ if rows == 0:
+ return output
+ table_width = block_table.shape[1]
+ num_words = triton.cdiv(table_width * tokens_per_block, 64)
+ bits, row_info = _row_bits(
+ selected_tokens, query_positions, compress_ratio, tokens_per_block, table_width, num_words
+ )
+ rows_per_tile = BLOCK_M // group_size
+ rows_pow2 = triton.next_power_of_2(rows_per_tile)
+ kv_strides = (
+ k_cache.stride(0), k_cache.stride(1), k_cache.stride(2), k_cache.stride(3),
+ v_cache.stride(0), v_cache.stride(1), v_cache.stride(2), v_cache.stride(3),
+ ) # fmt: skip
+ q_strides = (q.stride(0), q.stride(1), q.stride(2))
+ common = (*q_strides, *kv_strides, block_table.stride(0), block_table.stride(1))
+ geometry = dict(
+ NUM_CACHE_PAGES=k_cache.shape[0], NUM_REQUESTS=block_table.shape[0],
+ PAGE_TABLE_WIDTH=table_width, GROUP_SIZE=group_size,
+ TOKENS_PER_BLOCK=tokens_per_block, ROWS_PER_TILE=rows_per_tile, ROWS_POW2=rows_pow2,
+ HEAD_DIM=head_dim,
+ ) # fmt: skip
+ out_strides = (output.stride(0), output.stride(1), output.stride(2))
+ num_groups = triton.cdiv(rows, rows_per_tile)
+ # A group's list holds at most its rows' columns: the slot lists and the re-indexed
+ # bitmaps are that wide (in whole 64-slot tiles), not the table's width.
+ list_words = triton.cdiv(min(num_words * 64, rows_per_tile * selected_tokens.shape[1]), 64)
+ num_chunks = triton.cdiv(num_words, _CHUNK_WORDS)
+ slots = torch.empty((num_groups, list_words * 64), dtype=torch.int32, device=q.device)
+ packed_bits = torch.zeros((rows, list_words), dtype=torch.int64, device=q.device)
+ union_words = torch.empty((num_groups, num_words), dtype=torch.int64, device=q.device)
+ # [0]: per chunk, its resident tokens; [1]: its first list position.
+ chunk_info = torch.empty((2, num_groups, num_chunks), dtype=torch.int32, device=q.device)
+ group_plan = torch.empty((num_groups, 2), dtype=torch.int32, device=q.device)
+ lone_rows = torch.empty((rows,), dtype=torch.int32, device=q.device)
+ # [0]: the lone list's length; [1]: the row-tile pass's work-ticket counter.
+ counters = torch.zeros((2,), dtype=torch.int32, device=q.device)
+ lone_count = counters[:1]
+ rows_meta = dict(ROWS_PER_TILE=rows_per_tile, ROWS_POW2=rows_pow2)
+ # The pages an aligned 64-token word spans.
+ if 64 % tokens_per_block == 0:
+ word_pages = 64 // tokens_per_block
+ elif tokens_per_block % 64 == 0:
+ word_pages = 1
+ else:
+ word_pages = triton.cdiv(64, tokens_per_block) + 1
+ word_meta = dict(TOKENS_PER_BLOCK=tokens_per_block, CHUNK_WORDS=_CHUNK_WORDS,
+ WORD_PAGES=triton.next_power_of_2(word_pages)) # fmt: skip
+ _rowtile_union_kernel[(num_groups, num_chunks)](
+ request_indices, row_info, bits, block_table, union_words, chunk_info[0], rows,
+ bits.stride(0), block_table.stride(0), block_table.stride(1),
+ union_words.stride(0), chunk_info.stride(1),
+ block_table.shape[0], k_cache.shape[0],
+ PAGE_TABLE_WIDTH=table_width, num_warps=_PLAN_WARPS, **rows_meta, **word_meta,
+ ) # fmt: skip
+ _rowtile_group_kernel[(num_groups,)](
+ request_indices, row_info, chunk_info[0], chunk_info[1], group_plan, lone_rows,
+ lone_count, rows, chunk_info.stride(1), num_chunks, block_table.shape[0],
+ BLOCK_N=BLOCK_N, CHUNKS_POW2=triton.next_power_of_2(num_chunks),
+ num_warps=_PLAN_WARPS, **rows_meta,
+ ) # fmt: skip
+ _rowtile_slots_kernel[(num_groups, num_chunks)](
+ request_indices, row_info, bits, block_table, union_words, chunk_info[0], chunk_info[1],
+ group_plan, slots, packed_bits, rows,
+ bits.stride(0), packed_bits.stride(0), block_table.stride(0), block_table.stride(1),
+ union_words.stride(0), chunk_info.stride(1), slots.stride(0),
+ block_table.shape[0],
+ num_warps=_PLAN_WARPS, **rows_meta, **word_meta,
+ ) # fmt: skip
+ use_pdl = get_env_enable_pdl()
+ _rowtile_lone_kernel[(min(rows, _LONE_PROGRAMS), num_kv_heads)](
+ q, k_cache, v_cache, block_table, selected_tokens, request_indices, row_info,
+ lone_rows, lone_count, output, softmax_scale, rows,
+ *common, selected_tokens.stride(0), selected_tokens.stride(1), *out_strides,
+ BLOCK_H=max(16, triton.next_power_of_2(group_size)), BLOCK_N=64,
+ LONE_PROGRAMS=min(rows, _LONE_PROGRAMS), USE_PDL=use_pdl, LONE_STAGES=_LONE_STAGES,
+ num_warps=_LONE_WARPS, **geometry,
+ ) # fmt: skip
+ _rowtile_attention_cute[(_rowtile_ctas(num_groups * num_kv_heads, q.device),)](
+ q, k_cache, v_cache, slots, packed_bits, group_plan, output, counters[1:], softmax_scale,
+ rows,
+ group_size=group_size, rows_per_tile=rows_per_tile, tokens_per_block=tokens_per_block,
+ use_pdl=use_pdl,
+ ) # fmt: skip
+ return output
+
+
+_CHECKPOINT = dict(
+ q_heads=24,
+ kv_heads=2,
+ head_dim=256,
+ tokens_per_block=32,
+ compress_ratio=4,
+ token_topk=2048,
+ cache_dtype=torch.bfloat16,
+ q_dtype=torch.bfloat16,
+ # K and V are the two role planes of one [pages, 2, kv, tpb, D] pool.
+ cache_layout="pool_planes",
+ # q is the first 24 heads of the packed [T, 24 + 2 + 2, 256] QKV buffer.
+ q_layout="qkv_slice",
+ request_dtype=torch.int32,
+ # The engine's table at max_seq_len 10240: 320 pages per slot, -1 past each sequence.
+ pages=320,
+)
+
+
+def _spec(segments, slots=32, **overrides) -> dict:
+ """`segments` is `(first position, rows)` per request, in row order."""
+ return {**_CHECKPOINT, "segments": tuple(segments), "slots": slots, **overrides}
+
+
+def _decodes(count, rows_each=1, first=8200):
+ return [(first + (37 * g) % 1000, rows_each) for g in range(count)]
+
+
+class QsaSparseGqaRowtile(OpWrapper):
+ """Sparse GQA over the selected tokens of an HND paged K/V pool, for a
+ launch that carries the prefill pair (`query_positions`, `compress_ratio`).
+
+ For row `r`, query head `h`, KV head `h // (Hq // Hkv)`, over the columns
+ `j < bound(r)` of `selected_tokens[r]` that name a resident token::
+
+ t = selected_tokens[r, j] # logical position
+ page = block_table[request_indices[r], t // tokens_per_block]
+ s_j = (q[r, h] . k_cache[page, kv, t % tpb]) * softmax_scale
+ out[r, h] = sum_j softmax(s)_j * v_cache[page, kv, t % tpb]
+
+ `bound(r)` is the per-row kernel's prefill bound -- `min(visible // cr,
+ block_topk) * cr + visible % cr` columns with `visible = query_positions[r]
+ + 1`, rounded up to 64 -- so a row reads exactly the columns
+ `qsa_paged_sparse_gqa` reads on a context launch. Negative indices, pages
+ past the table, unmapped or out-of-pool pages and foreign requests
+ contribute nothing; a row left with nothing returns zeros. Returns a
+ fresh, contiguous `[rows, Hq, D]` tensor of `q.dtype`.
+
+ Six launches. A bitmap pass: one bit per selected logical token per row.
+ A plan per group of 10 consecutive rows of one request, in three passes
+ over (group, 32 bitmap words) programs: the union of the rows' selected
+ tokens whose page is mapped (a page lookup per word's page, not per
+ token) and its count per chunk; a per-group scan of those counts (list
+ length and tiles); the union listed in logical order as physical cache
+ slots, and each row's bitmap re-indexed onto that list (a per-word
+ parallel bit extract ORed into place). Every group of one request goes to
+ the row-tile pass, however long its list: its list is never longer than
+ its rows' own columns combined, and a listed tile costs no more than a
+ per-row tile. A group whose rows span requests joins the lone
+ list. A lone-row pass (256 programs
+ per KV head striding the list) serves listed rows over their own columns,
+ as the per-row kernel does. The row-tile pass is persistent (one program
+ per SM drawing (group, KV head) items from a ticket counter, the plan's
+ last groups first); per item it loads each listed tile's K/V once for the
+ 120 (row, head) pairs and masks each row to its own re-indexed bits, and
+ one item's epilogue overlaps the next item's first MMAs. It is launched
+ programmatically dependent on the lone pass, whose rows are disjoint, so
+ the two overlap. That pass is `qsa_rowtile_attention_cute` (tcgen05 MMAs,
+ S / P / O in tensor memory, a cp.async gather of the K/V rows, a TMA store
+ of O), so the entry serves SM103, bf16 q and K/V and D = 256.
+
+ Fusion boundary: as `qsa_paged_sparse_gqa` -- the block-table gather,
+ scale, online softmax and P.V are inside; deciding what to attend to,
+ appending K/V and q's projection, norm, RoPE and gate are not.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="con 1 prefill: 8192 rows of one request from 0, across the budget",
+ spec=_spec([(0, 8192)]),
+ ),
+ Cell(
+ why="sparse union: 4096 rows from 4096 ranking one score vector",
+ spec=_spec([(4096, 4096)], shared_scores=True),
+ ),
+ Cell(
+ why="con 128 mixed: 3 chunks (16240 rows) + 52 decodes on the lone pass",
+ spec=_spec([(0, 8192), (2192, 6000), (6144, 2048)] + _decodes(52), slots=128),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none -- every input is an argument, the same arguments
+ `qsa_paged_sparse_gqa` takes, with the prefill pair required. What the
+ target passes: `qsa_block_table` (physical page ids), the expansion's
+ `-1`-padded logical selections, the layer pool's K/V role planes, and the
+ logical positions with the checkpoint's compress ratio.
+
+ Preconditions (neither is checkable without reading the contents):
+ - Within the bound, a row names each token at most once. The row-tile
+ pass attends to the *set* of tokens a row names; the per-row kernel
+ attends to every column, so a repeated index would weigh twice there
+ and once here. QSA's top-k + expansion never repeats one (distinct
+ blocks, and a tail outside every complete block): measured zero
+ repeats on four layers of the target's own 8192-token prefill.
+ Rows served by the lone-row pass keep per-column semantics.
+ - As for the per-row kernel: each row is prefix-packed with exactly the
+ bound's count of valid entries; entries past the 64-rounded bound
+ are ignored.
+
+ Work: a row-tile group walks its packed list, so its MMA work grows
+ with the union of its rows' selections rather than with the group's
+ highest position. On the target's own selections (8192-token prefill,
+ layers 0/3/6/11) that union is 13-32% shorter than every 64-token tile up
+ to the highest token, and the walk loads ~5x fewer K/V tiles than the
+ per-row kernel. Measured on GB300 (eager, whole call) at 8192 context
+ rows of real selections: 1544-1565 us on layers 0-6 and 1274 us on layer
+ 11, against 1711-1715 us for a walk of every tile up to the highest token
+ and 2953 us for `qsa_paged_sparse_gqa`. Uniformly random selections
+ leave almost no gap in the union, and the call costs what the full walk
+ does (16292-row mixed step: 3859 vs 3836 us). Below a few hundred rows
+ it is slower than the per-row kernel (generation-only launches 1.4-2x);
+ the target's served context launches are thousands of rows. The tensor-core
+ row-tile pass alone, on six dumped 16384-row context launches of the
+ target (GB300, layers 0-5 of one step): 581-775 us, against 694-886 us for
+ the one-program-per-item pass it replaced (the same output, bit for bit).
+
+ Compiled once per (group size, rows per tile, tokens per block, PDL):
+ nothing in the row-tile pass's compile key depends on the row count or
+ the plan.
+
+ Scratch per call: the bitmap, `rows x ceil(table_width * tpb / 64)` int64
+ words (21 MB at 16384 rows and 320 pages), and per group its union over
+ as many words (2 MB); the re-indexed bitmaps, `rows x ceil(10 x
+ selection width / 64)` int64 words (42 MB at 16384 rows), and the slot
+ lists, `ceil(rows / 10) x` 64 times that many int32 slots (27 MB); the
+ per-chunk counts, the plan and the lone list. Calls allocate,
+ so capture in a CUDA graph is not certified; this entry serves the
+ prefill path, which runs eager. Results are deterministic: the bitmap
+ OR, the re-indexed bitmap OR and the lone list order do not change any
+ row's arithmetic.
+
+ Not certified: int64 request indices or positions; rows whose selection
+ resolves no token (all padding, a foreign request, unmapped pages);
+ `rows == 0`.
+ """
+
+ def __call__(
+ self,
+ *,
+ q: torch.Tensor,
+ k_cache: torch.Tensor,
+ v_cache: torch.Tensor,
+ block_table: torch.Tensor,
+ selected_tokens: torch.Tensor,
+ request_indices: torch.Tensor,
+ tokens_per_block: int,
+ softmax_scale: float,
+ query_positions: torch.Tensor,
+ compress_ratio: int,
+ ) -> torch.Tensor:
+ return qsa_rowtile_sparse_gqa(
+ q=q,
+ k_cache=k_cache,
+ v_cache=v_cache,
+ block_table=block_table,
+ selected_tokens=selected_tokens,
+ request_indices=request_indices,
+ tokens_per_block=tokens_per_block,
+ softmax_scale=softmax_scale,
+ query_positions=query_positions,
+ compress_ratio=compress_ratio,
+ )
+
+ def reference(
+ self,
+ *,
+ q: torch.Tensor,
+ k_cache: torch.Tensor,
+ v_cache: torch.Tensor,
+ block_table: torch.Tensor,
+ selected_tokens: torch.Tensor,
+ request_indices: torch.Tensor,
+ tokens_per_block: int,
+ softmax_scale: float,
+ query_positions: torch.Tensor,
+ compress_ratio: int,
+ ) -> torch.Tensor:
+ """The reference `qsa_paged_sparse_gqa` certifies against: fp32 sparse
+ GQA over every resident selected token, rounded once. The prefill pair
+ is ignored, as there: on a prefix-packed selection the bound reads
+ exactly the valid entries."""
+ return _per_row_entry.reference(
+ q=q,
+ k_cache=k_cache,
+ v_cache=v_cache,
+ block_table=block_table,
+ selected_tokens=selected_tokens,
+ request_indices=request_indices,
+ tokens_per_block=tokens_per_block,
+ softmax_scale=softmax_scale,
+ query_positions=query_positions,
+ compress_ratio=compress_ratio,
+ )
+
+
+qsa_sparse_gqa_rowtile = QsaSparseGqaRowtile()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_topk_expand_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_topk_expand_decode.py
new file mode 100644
index 000000000000..cf74ab9cefec
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_topk_expand_decode.py
@@ -0,0 +1,1556 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""QSA's block top-k and its expansion to token indices, one launch per batch of rows.
+
+A CuTe DSL kernel, an in-tree callable. One CTA per row holds four consecutive scores a
+thread and narrows to the k-th largest by radix passes over the scores' bit patterns
+(shared-memory histograms of 11 bits), starting below the bits every visible score
+shares. It stops as soon as the bin holding the threshold has at most `CAP` keys and
+ranks those exactly by pairwise comparison (ties to the higher block, the canonical
+selection's rule). Every selected block is then placed in block order by one block-wide
+scan, and expanded to token indices with the causal tail, as the expansion entry does.
+
+Scores are a sum of ReLUs times a positive scale, so non-negative: their fp32 bit
+patterns, read as int32, order like their values (a negative zero is taken as zero, a
+NaN ranks above everything).
+
+A table wider than one CTA's four keys a thread (`MAX_WIDTH`, e.g. the 65536 columns of a
+262144-token `max_seq_len`) takes `_TopkExpandWideKernel`, still one CTA a row and one
+launch: the row's window is staged in shared memory, a bound every top-k key clears is
+read off the lanes' maxima, the keys above it are compacted in block order, and the same
+selection runs over that short list (or over the window, slab by slab, when the list
+cannot hold them).
+"""
+
+from __future__ import annotations
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass import const_expr
+from cutlass.cute.nvgpu import cpasync
+from cutlass.cute.runtime import from_dlpack
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_expand_block_indices import (
+ qsa_expand_block_indices as _expand,
+)
+from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
+
+#: Consecutive scores a thread holds (one 16-byte load).
+KPT = 4
+#: Score columns a row may have: `KPT` a thread, at most 1024 threads a CTA.
+MAX_WIDTH = KPT * 1024
+#: CTAs an SM the register allocation is held to. At the 2592-wide score table a row is 672
+#: threads, which the per-subpartition register file holds once an SM at the compiler's own
+#: 48 registers a thread; at three (32 a thread) rows overlap each other's barriers and
+#: shared-memory round trips.
+MIN_BLOCKS = 3
+#: Radix passes at most (11 bits each; 31 magnitude bits).
+_PASSES = 3
+_HIST = 2048
+_I32_MAX = 0x7FFFFFFF
+#: At most this many keys in the threshold's bin are ranked pairwise instead of
+#: narrowed by another pass.
+CAP = 128
+#: Threads that own the 2048 bins of a scan, four each.
+_SCAN_THREADS = _HIST // 4
+#: Consecutive threads that rank one candidate of the threshold's bin, each over every
+#: `_RANK_LANES`-th candidate (`_SCAN_THREADS` threads cover `CAP` candidates).
+_RANK_LANES = _SCAN_THREADS // CAP
+
+#: Tables wider than `MAX_WIDTH` take the wide kernel: one CTA a row (1024 threads by
+#: default) walks its visible window in slabs of `KPT` columns a thread.
+_WIDE_THREADS = 1024
+#: The widest table the wide kernel serves (the row count times the row stride stays an
+#: Int32 product for 8192 rows).
+MAX_WIDE_WIDTH = 1 << 17
+#: Shared memory of an SM on sm_100/sm_103 (228 KiB), the most one CTA may hold (227 KiB),
+#: and what the driver reserves a CTA (1 KiB).
+_SMEM_SM = 233472
+_SMEM_LIMIT = 232448
+_SMEM_RESERVED = 1024
+#: The wide kernel's fixed shared state in 4-byte words: the histogram, the candidates,
+#: three sets of warp totals, the scan's double-buffered warp totals and the state words;
+#: and the alignment padding of its allocations.
+_WIDE_FIXED_WORDS = _HIST + 2 * CAP + 3 * 32 + 64 + 16
+_WIDE_PAD = 10 * 128
+
+
+def _threads(width: int) -> int:
+ return max(_SCAN_THREADS, -(-width // (KPT * 32)) * 32)
+
+
+class _TopkExpandKernel:
+ def __init__(self, *, width, key_stride, block_topk, compress, token_topk, pdl):
+ self.width, self.key_stride, self.k, self.cr, self.token_topk, self.pdl = (
+ width,
+ key_stride,
+ block_topk,
+ compress,
+ token_topk,
+ pdl,
+ )
+ self.final = token_topk + compress - 1
+ self.threads = _threads(width)
+ self.warps = self.threads // 32
+ self.flag_words = -(-width // 32)
+
+ @cute.jit
+ def __call__(
+ self,
+ keys_t: cute.Tensor,
+ end_ptr: cute.Pointer,
+ pos_ptr: cute.Pointer,
+ len_ptr: cute.Pointer,
+ topk_ptr: cute.Pointer,
+ sel_ptr: cute.Pointer,
+ rows: cutlass.Int32,
+ stream,
+ ):
+ # The keys arrive as a tensor, as the TVM-FFI entry takes its stream from a GPU
+ # tensor; their layout is rebuilt with the static width and row stride the 16-byte
+ # loads need. The row count is the call's.
+ mKeys = cute.make_tensor(
+ keys_t.iterator, cute.make_layout((rows, self.width), stride=(self.key_stride, 1))
+ )
+ per_row = cute.make_layout((rows,), stride=(1,))
+ mEnd = cute.make_tensor(end_ptr, per_row)
+ mPos = cute.make_tensor(pos_ptr, per_row)
+ mLen = cute.make_tensor(len_ptr, per_row)
+ mTopk = cute.make_tensor(topk_ptr, cute.make_layout((rows, self.k), stride=(self.k, 1)))
+ mSel = cute.make_tensor(
+ sel_ptr, cute.make_layout((rows, self.final), stride=(self.final, 1))
+ )
+ ld_keys = cute.make_copy_atom(
+ cute.nvgpu.CopyG2ROp(), cutlass.Int32, num_bits_per_copy=32 * KPT
+ )
+ self.kernel(mKeys, mEnd, mPos, mLen, mTopk, mSel, ld_keys).launch(
+ grid=[rows, 1, 1],
+ block=[self.threads, 1, 1],
+ smem=(_HIST + self.flag_words + 2 * CAP + 64 + 16) * 4 + 128,
+ stream=stream,
+ use_pdl=self.pdl,
+ min_blocks_per_mp=MIN_BLOCKS,
+ )
+
+ @cute.kernel
+ def kernel(self, mKeys, mEnd, mPos, mLen, mTopk, mSel, ld_keys: cute.CopyAtom):
+ tidx, _, _ = cute.arch.thread_idx()
+ row, _, _ = cute.arch.block_idx()
+ warp = cute.arch.warp_idx()
+ lane = cute.arch.lane_idx()
+ nt: cutlass.Constexpr = self.threads
+ nw: cutlass.Constexpr = self.warps
+ k: cutlass.Constexpr = self.k
+ cr: cutlass.Constexpr = self.cr
+ smem = cutlass.utils.SmemAllocator()
+ hist = smem.allocate_tensor(
+ cutlass.Int32, cute.make_layout((_HIST,), stride=(1,)), byte_alignment=16
+ )
+ flags = smem.allocate_tensor(
+ cutlass.Int32, cute.make_layout((self.flag_words,), stride=(1,)), byte_alignment=16
+ )
+ cand_key = smem.allocate_tensor(
+ cutlass.Int32, cute.make_layout((CAP,), stride=(1,)), byte_alignment=16
+ )
+ cand_idx = smem.allocate_tensor(
+ cutlass.Int32, cute.make_layout((CAP,), stride=(1,)), byte_alignment=16
+ )
+ wa = smem.allocate_tensor(
+ cutlass.Int32, cute.make_layout((32,), stride=(1,)), byte_alignment=16
+ )
+ wb = smem.allocate_tensor(
+ cutlass.Int32, cute.make_layout((32,), stride=(1,)), byte_alignment=16
+ )
+ state = smem.allocate_tensor(
+ cutlass.Int32, cute.make_layout((16,), stride=(1,)), byte_alignment=16
+ )
+
+ # 0. Shared state, before the wait (nothing global is read here).
+ for b in cutlass.range_constexpr(-(-_HIST // nt)):
+ if tidx + b * nt < _HIST:
+ hist[tidx + b * nt] = cutlass.Int32(0)
+ for b in cutlass.range_constexpr(-(-self.flag_words // nt)):
+ if tidx + b * nt < self.flag_words:
+ flags[tidx + b * nt] = cutlass.Int32(0)
+ if tidx == 0:
+ state[5] = cutlass.Int32(0) # candidates compacted
+ # The row's metadata, host-filled: read before the wait too.
+ n = mEnd[row]
+ seqlen = mLen[row]
+ visible = cutlass.Int32(mPos[row]) + 1
+ if const_expr(self.pdl):
+ cute.arch.griddepcontrol_wait()
+ cute.arch.griddepcontrol_launch_dependents()
+
+ # 1. Four consecutive keys a thread (a column past the window is no candidate:
+ # -1); the largest and smallest key in the window.
+ e0 = tidx * KPT
+ keys = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32)
+ keys.fill(-1)
+ if e0 < self.width:
+ cute.copy(ld_keys, cute.logical_divide(mKeys, (None, KPT))[row, (None, tidx)], keys)
+ kmax = cutlass.Int32(-1)
+ kmin = cutlass.Int32(_I32_MAX)
+ for c in cutlass.range_constexpr(KPT):
+ b = keys[c]
+ if b < 0: # -0.0 (or a negative): rank as zero
+ b = cutlass.Int32(0)
+ if e0 + c >= n:
+ b = cutlass.Int32(-1)
+ keys[c] = b
+ kmax = cutlass.max(kmax, b)
+ if b >= 0:
+ kmin = cutlass.min(kmin, b)
+ kmax = cute.arch.warp_reduction_max(kmax)
+ kmin = cute.arch.warp_reduction(kmin, lambda x, y: cutlass.min(x, y))
+ if lane == 0:
+ wa[warp] = kmax
+ wb[warp] = kmin
+ cute.arch.sync_threads()
+ gmax = cutlass.Int32(-1)
+ gmin = cutlass.Int32(_I32_MAX)
+ if lane < nw:
+ gmax = wa[lane]
+ gmin = wb[lane]
+ gmax = cute.arch.warp_reduction_max(gmax)
+ gmin = cute.arch.warp_reduction(gmin, lambda x, y: cutlass.min(x, y))
+
+ # 2. Every key in the window shares the bits above the highest bit where the
+ # largest and the smallest differ (`hb`); radix passes walk down from there, 11 bits
+ # at a time, until the threshold's bin holds at most CAP keys or no bits are left.
+ diff = gmax ^ gmin
+ hb = cutlass.Int32(-1)
+ if diff != 0:
+ hb = 31 - cutlass.Int32(cute.arch.clz(diff))
+ low = (cutlass.Int64(1) << cutlass.Int64(hb + 1)) - 1
+ mask0 = cutlass.Int32(cutlass.Int64(_I32_MAX) & (cutlass.Int64(-1) ^ low))
+ if tidx == 0:
+ state[0] = gmax & mask0 # threshold prefix
+ state[1] = mask0 # its mask
+ state[2] = cutlass.Int32(k) # still needed at the threshold
+ state[3] = n # keys in the threshold's bin
+ state[4] = cutlass.Int32(1) # narrowing done
+ state[6] = cutlass.max(hb - 10, 0) # the next pass's shift
+ if n > k:
+ if hb >= 0:
+ state[4] = cutlass.Int32(0)
+ cute.arch.sync_threads()
+ for _p in cutlass.range_constexpr(_PASSES):
+ if state[4] == 0:
+ prefix = state[0]
+ mask = state[1]
+ need = state[2]
+ shift = state[6]
+ for c in cutlass.range_constexpr(KPT):
+ if keys[c] >= 0:
+ if (keys[c] & mask) == prefix:
+ digit = (keys[c] >> shift) & (_HIST - 1)
+ cute.arch.atomic_add(
+ hist.iterator + digit, cutlass.Int32(1), scope="cta"
+ )
+ cute.arch.sync_threads()
+ # Thread t < 512 owns bins 2047-4t .. 2044-4t, counted from the top; each is
+ # cleared for the next pass once read.
+ h = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), cutlass.Int32)
+ h.fill(0)
+ own = cutlass.Int32(0)
+ if tidx < _SCAN_THREADS:
+ for j in cutlass.range_constexpr(4):
+ h[j] = hist[_HIST - 1 - 4 * tidx - j]
+ hist[_HIST - 1 - 4 * tidx - j] = cutlass.Int32(0)
+ own = own + h[j]
+ incl = own
+ for off in (1, 2, 4, 8, 16):
+ up = cute.arch.shuffle_sync_up(incl, offset=off)
+ if lane >= off:
+ incl = incl + up
+ if lane == 31:
+ wa[warp] = incl
+ cute.arch.sync_threads()
+ above = cutlass.Int32(0)
+ if lane < warp:
+ above = wa[lane]
+ before = cute.arch.warp_reduction_sum(above) + incl - own # keys in the bins above
+ if tidx < _SCAN_THREADS:
+ if before < need:
+ if need <= before + own:
+ run = before
+ for j in cutlass.range_constexpr(4):
+ if run < need:
+ if need <= run + h[j]:
+ state[0] = prefix | ((_HIST - 1 - 4 * tidx - j) << shift)
+ state[1] = mask | ((_HIST - 1) << shift)
+ state[2] = need - run
+ state[3] = h[j]
+ state[4] = cutlass.Int32(
+ cutlass.select_((h[j] <= CAP) | (shift == 0), 1, 0)
+ )
+ state[6] = cutlass.max(shift - 11, 0)
+ run = run + h[j]
+ cute.arch.sync_threads()
+
+ prefix = state[0]
+ mask = state[1]
+ need = state[2]
+ at_thr = state[3]
+ # 3. At most CAP keys in the threshold's bin: rank them exactly, ties to the
+ # higher block, and flag the `need` best.
+ if n > k:
+ if at_thr <= CAP:
+ for c in cutlass.range_constexpr(KPT):
+ if keys[c] >= 0:
+ if (keys[c] & mask) == prefix:
+ slot = cute.arch.atomic_add(
+ state.iterator + 5, cutlass.Int32(1), scope="cta"
+ )
+ cand_key[slot] = keys[c]
+ cand_idx[slot] = e0 + c
+ cute.arch.sync_threads()
+ # Candidate `ci` is ranked by `_RANK_LANES` consecutive lanes, lane `sub`
+ # counting the better ones among candidates sub, sub + _RANK_LANES, ...;
+ # butterfly shuffles add the counts up.
+ ci = tidx // _RANK_LANES
+ sub = tidx % _RANK_LANES
+ rank = cutlass.Int32(0)
+ it = cutlass.Int32(0)
+ if ci < at_thr:
+ kt = cand_key[ci]
+ it = cand_idx[ci]
+ for j in cutlass.range(sub, at_thr, _RANK_LANES, unroll=4):
+ kj = cand_key[j]
+ ij = cand_idx[j]
+ better = (kj > kt) | ((kj == kt) & (ij > it))
+ rank = rank + cutlass.Int32(cutlass.select_(better, 1, 0))
+ for off in cutlass.range_constexpr(_RANK_LANES.bit_length() - 1):
+ rank = rank + cute.arch.shuffle_sync_bfly(rank, offset=1 << off)
+ if ci < at_thr:
+ if sub == 0:
+ if rank < need:
+ cute.arch.atomic_or(
+ flags.iterator + (it >> 5),
+ cutlass.Int32(1) << (it & 31),
+ scope="cta",
+ )
+ cute.arch.sync_threads()
+
+ # 4. Block order: one scan of (selected, tied at the threshold) counts. A tied key
+ # is left only when the passes ran out of bits with more than CAP keys tied; the
+ # last `need` of those in block order are taken.
+ take = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32)
+ tied = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32)
+ own = cutlass.Int32(0)
+ for c in cutlass.range_constexpr(KPT):
+ t_c = cutlass.Int32(0)
+ q_c = cutlass.Int32(0)
+ if keys[c] >= 0:
+ if n <= k:
+ t_c = cutlass.Int32(1)
+ else:
+ mk = keys[c] & mask
+ if mk > prefix:
+ t_c = cutlass.Int32(1)
+ if mk == prefix:
+ if at_thr > CAP:
+ q_c = cutlass.Int32(1)
+ else:
+ e = e0 + c
+ t_c = (flags[e >> 5] >> (e & 31)) & 1
+ take[c] = t_c
+ tied[c] = q_c
+ own = own + t_c + (q_c << 16)
+ incl = own
+ for off in (1, 2, 4, 8, 16):
+ up = cute.arch.shuffle_sync_up(incl, offset=off)
+ if lane >= off:
+ incl = incl + up
+ if lane == 31:
+ wb[warp] = incl
+ cute.arch.sync_threads()
+ above = cutlass.Int32(0)
+ if lane < warp:
+ above = wb[lane]
+ pre = cute.arch.warp_reduction_sum(above) + incl - own
+ taken_before = pre & 0xFFFF
+ tied_before = pre >> 16
+ # The tied keys left out, the lowest-index ones (none unless the threshold's bin
+ # holds more keys than are still needed).
+ skip = cutlass.max(at_thr - need, 0)
+ for c in cutlass.range_constexpr(KPT):
+ e = e0 + c
+ sel = cutlass.Int32(0)
+ pos = cutlass.Int32(0)
+ if tied[c] == 1:
+ if tied_before >= skip:
+ sel = cutlass.Int32(1)
+ pos = taken_before + tied_before - skip
+ tied_before = tied_before + 1
+ if take[c] == 1:
+ sel = cutlass.Int32(1)
+ pos = taken_before + cutlass.max(tied_before - skip, 0)
+ taken_before = taken_before + 1
+ if sel == 1:
+ mTopk[row, pos] = e
+ for j in cutlass.range_constexpr(cr):
+ tok = e * cr + j
+ val = cutlass.Int32(-1)
+ if tok < seqlen:
+ val = tok
+ mSel[row, pos * cr + j] = val
+
+ # 5. Padding of the block list; the causal tail after the selected tokens, then -1.
+ total = n
+ if n > k:
+ total = cutlass.Int32(k)
+ for c in cutlass.range_constexpr(-(-k // nt)):
+ s = tidx + c * nt
+ if s < k:
+ if s >= total:
+ mTopk[row, s] = cutlass.Int32(-1)
+ tail_start = (visible // cr) * cr
+ tail_count = visible - tail_start
+ body = total * cr
+ if body > self.token_topk:
+ body = cutlass.Int32(self.token_topk)
+ for c in cutlass.range_constexpr(-(-self.final // nt)):
+ col = tidx + c * nt
+ if col < self.final:
+ if col >= body:
+ off = col - body
+ val = cutlass.Int32(-1)
+ if off < cr - 1:
+ if off < tail_count:
+ if tail_start + off < seqlen:
+ val = tail_start + off
+ mSel[row, col] = val
+
+
+def _cache_slabs(width: int, threads: int, ctas: int) -> int:
+ """Slabs of the wide kernel's window cache: as many of the table's as the shared memory
+ left after the fixed state and the candidate list hold with `ctas` CTAs an SM (11 of
+ 4096 columns at one CTA of 1024 threads: a window of up to 45056 blocks, 180224 tokens,
+ is read from global once)."""
+ slab = KPT * threads
+ per_cta = min(_SMEM_LIMIT, _SMEM_SM // ctas - _SMEM_RESERVED)
+ room = per_cta - (_WIDE_FIXED_WORDS + 2 * slab) * 4 - _WIDE_PAD
+ return max(0, min(-(-width // slab), room // (slab * 4)))
+
+
+class _TopkExpandWideKernel:
+ """`_TopkExpandKernel`'s selection over a table wider than one CTA's registers. The row's
+ visible window (`row_ends`, not the table width) is walked in slabs, `KPT` consecutive
+ keys a thread a slab; its first `cache_slabs` slabs are read from global once (cp.async
+ into shared memory) and every walk reads them from there, a later slab from global (L2).
+
+ Prefilter: every lane's largest key is a distinct key of the window, so with `per_warp`
+ the `k / warps` rounded up, the smallest over the warps of the `per_warp`-th largest lane
+ maximum of each warp (`lo`) has at least `k` keys at or above it. The `k`-th largest key,
+ every key tied with it and everything above are among those, so the keys below `lo` are
+ no candidates and the selection is unchanged.
+
+ Two walks of the window: the lane maxima, then the candidates (keys at or above `lo`)
+ compacted in block order into a shared list of up to `KPT` a thread (a window of ~21k
+ relu-sum scores keeps about a thousand). The narrow kernel's radix passes, pairwise rank
+ of the threshold's bin and block-order scan then run once over that list, `KPT`
+ consecutive entries a thread. A window with more candidates than the list holds (a
+ flat or degenerate row) takes the same steps over the whole window instead, slab by
+ slab, the threshold bin's `need` best kept as a cut (the worst of them, key and block)
+ and the block-order scan's counts carried from slab to slab."""
+
+ def __init__(
+ self,
+ *,
+ width,
+ key_stride,
+ block_topk,
+ compress,
+ token_topk,
+ pdl,
+ threads=_WIDE_THREADS,
+ ctas=1,
+ prefilter=True,
+ ):
+ self.width, self.key_stride, self.k, self.cr, self.token_topk, self.pdl = (
+ width,
+ key_stride,
+ block_topk,
+ compress,
+ token_topk,
+ pdl,
+ )
+ assert threads % 32 == 0 and _SCAN_THREADS <= threads <= 1024, threads
+ self.final = token_topk + compress - 1
+ self.threads = threads
+ self.warps = threads // 32
+ self.slab = KPT * threads
+ self.ctas = ctas
+ self.cache_slabs = _cache_slabs(width, threads, ctas)
+ self.per_warp = -(-block_topk // self.warps)
+ self.prefilter = bool(prefilter) and self.per_warp <= 32
+
+ def smem_bytes(self) -> int:
+ return (self.cache_slabs * self.slab + 2 * self.slab + _WIDE_FIXED_WORDS) * 4 + _WIDE_PAD
+
+ @cute.jit
+ def __call__(
+ self,
+ keys_t: cute.Tensor,
+ end_ptr: cute.Pointer,
+ pos_ptr: cute.Pointer,
+ len_ptr: cute.Pointer,
+ topk_ptr: cute.Pointer,
+ sel_ptr: cute.Pointer,
+ rows: cutlass.Int32,
+ stream,
+ ):
+ mKeys = cute.make_tensor(
+ keys_t.iterator, cute.make_layout((rows, self.width), stride=(self.key_stride, 1))
+ )
+ per_row = cute.make_layout((rows,), stride=(1,))
+ mEnd = cute.make_tensor(end_ptr, per_row)
+ mPos = cute.make_tensor(pos_ptr, per_row)
+ mLen = cute.make_tensor(len_ptr, per_row)
+ mTopk = cute.make_tensor(topk_ptr, cute.make_layout((rows, self.k), stride=(self.k, 1)))
+ mSel = cute.make_tensor(
+ sel_ptr, cute.make_layout((rows, self.final), stride=(self.final, 1))
+ )
+ ld_keys = cute.make_copy_atom(
+ cute.nvgpu.CopyG2ROp(), cutlass.Int32, num_bits_per_copy=32 * KPT
+ )
+ g2s = cute.make_copy_atom(cpasync.CopyG2SOp(), cutlass.Int32, num_bits_per_copy=32 * KPT)
+ self.kernel(mKeys, mEnd, mPos, mLen, mTopk, mSel, ld_keys, g2s).launch(
+ grid=[rows, 1, 1],
+ block=[self.threads, 1, 1],
+ smem=self.smem_bytes(),
+ stream=stream,
+ use_pdl=self.pdl,
+ min_blocks_per_mp=self.ctas,
+ )
+
+ @cute.jit
+ def _slab(self, keys, gK, sK, row, s, tidx, n, lo, ld_keys: cute.CopyAtom):
+ """Slab `s`'s `KPT` keys of thread `tidx` into `keys`: a negative (-0.0) as zero, -1
+ past the window and below `lo`."""
+ col4 = s * self.threads + tidx
+ e0 = col4 * KPT
+ if const_expr(self.cache_slabs > 0):
+ if s < self.cache_slabs:
+ cute.autovec_copy(sK[(None, col4)], keys)
+ else:
+ keys.fill(-1)
+ if e0 < n:
+ cute.copy(ld_keys, gK[row, (None, col4)], keys)
+ else:
+ keys.fill(-1)
+ if e0 < n:
+ cute.copy(ld_keys, gK[row, (None, col4)], keys)
+ for c in cutlass.range_constexpr(KPT):
+ b = cutlass.max(keys[c], cutlass.Int32(0))
+ if (b < lo) | (e0 + c >= n):
+ b = cutlass.Int32(-1)
+ keys[c] = b
+
+ @cute.jit
+ def _passes(self, keys, tidx, warp, lane, hist, wsum, state):
+ """The narrow kernel's radix passes over `keys` (this thread's share; -1 is none):
+ from the state's prefix, mask and shift, 11 bits a pass, until the threshold's bin
+ holds at most CAP keys or no bits are left."""
+ for _p in cutlass.range_constexpr(_PASSES):
+ if state[4] == 0:
+ prefix = state[0]
+ mask = state[1]
+ need = state[2]
+ shift = state[6]
+ for c in cutlass.range_constexpr(KPT):
+ if keys[c] >= 0:
+ if (keys[c] & mask) == prefix:
+ digit = (keys[c] >> shift) & (_HIST - 1)
+ cute.arch.atomic_add(
+ hist.iterator + digit, cutlass.Int32(1), scope="cta"
+ )
+ cute.arch.sync_threads()
+ self._narrow(tidx, warp, lane, hist, wsum, state, prefix, mask, need, shift)
+
+ @cute.jit
+ def _narrow(self, tidx, warp, lane, hist, wsum, state, prefix, mask, need, shift):
+ """One pass's histogram read (and cleared) from the top bin down: the bin holding
+ the `need`-th key becomes the threshold's."""
+ h = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), cutlass.Int32)
+ h.fill(0)
+ own = cutlass.Int32(0)
+ if tidx < _SCAN_THREADS:
+ for j in cutlass.range_constexpr(4):
+ h[j] = hist[_HIST - 1 - 4 * tidx - j]
+ hist[_HIST - 1 - 4 * tidx - j] = cutlass.Int32(0)
+ own = own + h[j]
+ incl = own
+ for off in (1, 2, 4, 8, 16):
+ up = cute.arch.shuffle_sync_up(incl, offset=off)
+ if lane >= off:
+ incl = incl + up
+ if lane == 31:
+ wsum[warp] = incl
+ cute.arch.sync_threads()
+ above = cutlass.Int32(0)
+ if lane < warp:
+ above = wsum[lane]
+ before = cute.arch.warp_reduction_sum(above) + incl - own
+ if tidx < _SCAN_THREADS:
+ if before < need:
+ if need <= before + own:
+ run = before
+ for j in cutlass.range_constexpr(4):
+ if run < need:
+ if need <= run + h[j]:
+ state[0] = prefix | ((_HIST - 1 - 4 * tidx - j) << shift)
+ state[1] = mask | ((_HIST - 1) << shift)
+ state[2] = need - run
+ state[3] = h[j]
+ state[4] = cutlass.Int32(
+ cutlass.select_((h[j] <= CAP) | (shift == 0), 1, 0)
+ )
+ state[6] = cutlass.max(shift - 11, 0)
+ run = run + h[j]
+ cute.arch.sync_threads()
+
+ @cute.jit
+ def _cut(self, tidx, state, cand_key, cand_idx, at_thr, need):
+ """The threshold bin's `at_thr` <= CAP keys (gathered into `cand_key`/`cand_idx`)
+ ranked pairwise, ties to the higher block: the `need`-th best is the cut, kept in
+ `state[7]` (key) and `state[8]` (block)."""
+ ci = tidx // _RANK_LANES
+ sub = tidx % _RANK_LANES
+ rank = cutlass.Int32(0)
+ kt = cutlass.Int32(0)
+ it = cutlass.Int32(0)
+ if ci < at_thr:
+ kt = cand_key[ci]
+ it = cand_idx[ci]
+ for j in cutlass.range(sub, at_thr, _RANK_LANES, unroll=4):
+ kj = cand_key[j]
+ ij = cand_idx[j]
+ better = (kj > kt) | ((kj == kt) & (ij > it))
+ rank = rank + cutlass.Int32(cutlass.select_(better, 1, 0))
+ for off in cutlass.range_constexpr(_RANK_LANES.bit_length() - 1):
+ rank = rank + cute.arch.shuffle_sync_bfly(rank, offset=1 << off)
+ if ci < at_thr:
+ if sub == 0:
+ if rank == need - 1:
+ state[7] = kt
+ state[8] = it
+ cute.arch.sync_threads()
+
+ @cute.jit
+ def _mark(self, keys, idx, take, tied, everything, prefix, mask, at_thr, cut_key, cut_idx):
+ """Each of this thread's keys taken (above the threshold's bin, or in it at or above
+ the cut) or tied (in it, when the passes ran out of bits with more than CAP keys);
+ `everything` takes every key. Returns the packed counts, taken + (tied << 16)."""
+ own = cutlass.Int32(0)
+ for c in cutlass.range_constexpr(KPT):
+ t_c = cutlass.Int32(0)
+ q_c = cutlass.Int32(0)
+ if keys[c] >= 0:
+ if everything:
+ t_c = cutlass.Int32(1)
+ else:
+ mk = keys[c] & mask
+ if mk > prefix:
+ t_c = cutlass.Int32(1)
+ if mk == prefix:
+ if at_thr > CAP:
+ q_c = cutlass.Int32(1)
+ elif (keys[c] > cut_key) | ((keys[c] == cut_key) & (idx[c] >= cut_idx)):
+ t_c = cutlass.Int32(1)
+ take[c] = t_c
+ tied[c] = q_c
+ own = own + t_c + (q_c << 16)
+ return own
+
+ @cute.jit
+ def _emit(self, row, idx, take, tied, taken_before, tied_before, skip, seqlen, mTopk, mSel):
+ """This thread's selected blocks at their block-order positions, and their tokens
+ below the sequence length (-1 past it). `tied_before - skip` of the tied keys before
+ these are taken (the lowest-index `skip` of the tie class are left)."""
+ for c in cutlass.range_constexpr(KPT):
+ e = idx[c]
+ sel = cutlass.Int32(0)
+ pos = cutlass.Int32(0)
+ if tied[c] == 1:
+ if tied_before >= skip:
+ sel = cutlass.Int32(1)
+ pos = taken_before + tied_before - skip
+ tied_before = tied_before + 1
+ if take[c] == 1:
+ sel = cutlass.Int32(1)
+ pos = taken_before + cutlass.max(tied_before - skip, 0)
+ taken_before = taken_before + 1
+ if sel == 1:
+ mTopk[row, pos] = e
+ for j in cutlass.range_constexpr(self.cr):
+ tok = e * self.cr + j
+ val = cutlass.Int32(-1)
+ if tok < seqlen:
+ val = tok
+ mSel[row, pos * self.cr + j] = val
+
+ @cute.jit
+ def _select_listed(
+ self,
+ ckey,
+ cidx,
+ keys,
+ row,
+ tidx,
+ warp,
+ lane,
+ m,
+ seqlen,
+ hist,
+ cand_key,
+ cand_idx,
+ wsum,
+ state,
+ mTopk,
+ mSel,
+ ):
+ """The narrow kernel's selection over the candidate list: `KPT` consecutive entries
+ a thread, so one block-wide scan puts them in block order."""
+ idx = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32)
+ take = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32)
+ tied = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32)
+ cute.autovec_copy(ckey[(None, tidx)], keys)
+ cute.autovec_copy(cidx[(None, tidx)], idx)
+ for c in cutlass.range_constexpr(KPT):
+ if tidx * KPT + c >= m:
+ keys[c] = cutlass.Int32(-1)
+ self._passes(keys, tidx, warp, lane, hist, wsum, state)
+ prefix = state[0]
+ mask = state[1]
+ need = state[2]
+ at_thr = state[3]
+ if m > self.k:
+ if at_thr <= CAP:
+ for c in cutlass.range_constexpr(KPT):
+ if keys[c] >= 0:
+ if (keys[c] & mask) == prefix:
+ slot = cute.arch.atomic_add(
+ state.iterator + 5, cutlass.Int32(1), scope="cta"
+ )
+ cand_key[slot] = keys[c]
+ cand_idx[slot] = idx[c]
+ cute.arch.sync_threads()
+ self._cut(tidx, state, cand_key, cand_idx, at_thr, need)
+ skip = cutlass.max(at_thr - need, 0)
+ own = self._mark(
+ keys, idx, take, tied, m <= self.k, prefix, mask, at_thr, state[7], state[8]
+ )
+ incl = own
+ for off in (1, 2, 4, 8, 16):
+ up = cute.arch.shuffle_sync_up(incl, offset=off)
+ if lane >= off:
+ incl = incl + up
+ if lane == 31:
+ wsum[warp] = incl
+ cute.arch.sync_threads()
+ above = cutlass.Int32(0)
+ if lane < warp:
+ above = wsum[lane]
+ pre = cute.arch.warp_reduction_sum(above) + incl - own
+ self._emit(row, idx, take, tied, pre & 0xFFFF, pre >> 16, skip, seqlen, mTopk, mSel)
+
+ @cute.jit
+ def _select_window(
+ self,
+ keys,
+ gK,
+ sK,
+ row,
+ tidx,
+ warp,
+ lane,
+ n,
+ nslabs,
+ lo,
+ m,
+ seqlen,
+ hist,
+ cand_key,
+ cand_idx,
+ wsum,
+ wsc,
+ state,
+ mTopk,
+ mSel,
+ ld_keys: cute.CopyAtom,
+ ):
+ """The same selection when the list cannot hold every candidate: each step walks
+ the window slab by slab, the block-order scan's counts carried from slab to slab (a
+ slab's fit 16 bits each; the warp totals alternate between two halves of `wsc`)."""
+ nt = self.threads
+ nw = self.warps
+ for _p in cutlass.range_constexpr(_PASSES):
+ if state[4] == 0:
+ prefix = state[0]
+ mask = state[1]
+ need = state[2]
+ shift = state[6]
+ for s in cutlass.range(nslabs):
+ self._slab(keys, gK, sK, row, s, tidx, n, lo, ld_keys)
+ for c in cutlass.range_constexpr(KPT):
+ if keys[c] >= 0:
+ if (keys[c] & mask) == prefix:
+ digit = (keys[c] >> shift) & (_HIST - 1)
+ cute.arch.atomic_add(
+ hist.iterator + digit, cutlass.Int32(1), scope="cta"
+ )
+ cute.arch.sync_threads()
+ self._narrow(tidx, warp, lane, hist, wsum, state, prefix, mask, need, shift)
+ prefix = state[0]
+ mask = state[1]
+ need = state[2]
+ at_thr = state[3]
+ if m > self.k:
+ if at_thr <= CAP:
+ for s in cutlass.range(nslabs):
+ self._slab(keys, gK, sK, row, s, tidx, n, lo, ld_keys)
+ e0 = (s * nt + tidx) * KPT
+ for c in cutlass.range_constexpr(KPT):
+ if keys[c] >= 0:
+ if (keys[c] & mask) == prefix:
+ slot = cute.arch.atomic_add(
+ state.iterator + 5, cutlass.Int32(1), scope="cta"
+ )
+ cand_key[slot] = keys[c]
+ cand_idx[slot] = e0 + c
+ cute.arch.sync_threads()
+ self._cut(tidx, state, cand_key, cand_idx, at_thr, need)
+ skip = cutlass.max(at_thr - need, 0)
+ cut_key = state[7]
+ cut_idx = state[8]
+ idx = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32)
+ take = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32)
+ tied = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32)
+ carry_taken = cutlass.Int32(0)
+ carry_tied = cutlass.Int32(0)
+ for s in cutlass.range(nslabs):
+ self._slab(keys, gK, sK, row, s, tidx, n, lo, ld_keys)
+ e0 = (s * nt + tidx) * KPT
+ for c in cutlass.range_constexpr(KPT):
+ idx[c] = e0 + c
+ own = self._mark(
+ keys, idx, take, tied, m <= self.k, prefix, mask, at_thr, cut_key, cut_idx
+ )
+ incl = own
+ for off in (1, 2, 4, 8, 16):
+ up = cute.arch.shuffle_sync_up(incl, offset=off)
+ if lane >= off:
+ incl = incl + up
+ half = (s & 1) * 32
+ if lane == 31:
+ wsc[half + warp] = incl
+ cute.arch.sync_threads()
+ above = cutlass.Int32(0)
+ if lane < warp:
+ above = wsc[half + lane]
+ pre = cute.arch.warp_reduction_sum(above) + incl - own
+ slab_total = cutlass.Int32(0)
+ if lane < nw:
+ slab_total = wsc[half + lane]
+ slab_total = cute.arch.warp_reduction_sum(slab_total)
+ self._emit(
+ row,
+ idx,
+ take,
+ tied,
+ carry_taken + (pre & 0xFFFF),
+ carry_tied + (pre >> 16),
+ skip,
+ seqlen,
+ mTopk,
+ mSel,
+ )
+ carry_taken = carry_taken + (slab_total & 0xFFFF)
+ carry_tied = carry_tied + (slab_total >> 16)
+
+ @cute.kernel
+ def kernel(
+ self, mKeys, mEnd, mPos, mLen, mTopk, mSel, ld_keys: cute.CopyAtom, g2s: cute.CopyAtom
+ ):
+ tidx, _, _ = cute.arch.thread_idx()
+ row, _, _ = cute.arch.block_idx()
+ warp = cute.arch.warp_idx()
+ lane = cute.arch.lane_idx()
+ nt: cutlass.Constexpr = self.threads
+ nw: cutlass.Constexpr = self.warps
+ slab: cutlass.Constexpr = self.slab
+ k: cutlass.Constexpr = self.k
+ cr: cutlass.Constexpr = self.cr
+ smem = cutlass.utils.SmemAllocator()
+ sK = None
+ if const_expr(self.cache_slabs > 0):
+ # `KPT` consecutive keys (mode 0) of each thread-slab column (mode 1).
+ sK = smem.allocate_tensor(
+ cutlass.Int32,
+ cute.make_layout((KPT, self.cache_slabs * slab // KPT), stride=(1, KPT)),
+ byte_alignment=128,
+ )
+ hist = smem.allocate_tensor(
+ cutlass.Int32, cute.make_layout((_HIST,), stride=(1,)), byte_alignment=128
+ )
+ cand_key = smem.allocate_tensor(
+ cutlass.Int32, cute.make_layout((CAP,), stride=(1,)), byte_alignment=128
+ )
+ cand_idx = smem.allocate_tensor(
+ cutlass.Int32, cute.make_layout((CAP,), stride=(1,)), byte_alignment=128
+ )
+ wa = smem.allocate_tensor(
+ cutlass.Int32, cute.make_layout((32,), stride=(1,)), byte_alignment=128
+ )
+ wb = smem.allocate_tensor(
+ cutlass.Int32, cute.make_layout((32,), stride=(1,)), byte_alignment=128
+ )
+ wc = smem.allocate_tensor(
+ cutlass.Int32, cute.make_layout((32,), stride=(1,)), byte_alignment=128
+ )
+ wsc = smem.allocate_tensor(
+ cutlass.Int32, cute.make_layout((64,), stride=(1,)), byte_alignment=128
+ )
+ state = smem.allocate_tensor(
+ cutlass.Int32, cute.make_layout((16,), stride=(1,)), byte_alignment=128
+ )
+ # The candidates in block order: keys and blocks, `KPT` consecutive a thread.
+ ckey = smem.allocate_tensor(
+ cutlass.Int32, cute.make_layout((KPT, nt), stride=(1, KPT)), byte_alignment=128
+ )
+ cidx = smem.allocate_tensor(
+ cutlass.Int32, cute.make_layout((KPT, nt), stride=(1, KPT)), byte_alignment=128
+ )
+
+ # 0. Shared state and the row's host-filled metadata, before the wait.
+ for b in cutlass.range_constexpr(-(-_HIST // nt)):
+ if tidx + b * nt < _HIST:
+ hist[tidx + b * nt] = cutlass.Int32(0)
+ if tidx == 0:
+ state[5] = cutlass.Int32(0) # candidates compacted
+ n = mEnd[row]
+ seqlen = mLen[row]
+ visible = cutlass.Int32(mPos[row]) + 1
+ if const_expr(self.pdl):
+ cute.arch.griddepcontrol_wait()
+ cute.arch.griddepcontrol_launch_dependents()
+
+ # 1. The window's first slabs into shared memory, every load in flight at once.
+ gK = cute.logical_divide(mKeys, (None, KPT))
+ nslabs = (n + (slab - 1)) // slab
+ if const_expr(self.cache_slabs > 0):
+ ncached = cutlass.min(nslabs, cutlass.Int32(self.cache_slabs))
+ for s in cutlass.range(ncached):
+ col4 = s * nt + tidx
+ if col4 * KPT < n:
+ cute.copy(g2s, gK[row, (None, col4)], sK[(None, col4)])
+ cute.arch.cp_async_commit_group()
+ cute.arch.cp_async_wait_group(0)
+ cute.arch.sync_threads()
+
+ # 2. The window's largest key, and the prefilter's bound `lo`: each warp's
+ # `per_warp`-th largest lane maximum (ranked by every lane against all 32, ties to
+ # the lower lane, so exactly one lane holds each rank), then their minimum.
+ keys = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32)
+ tmax = cutlass.Int32(-1)
+ for s in cutlass.range(nslabs):
+ self._slab(keys, gK, sK, row, s, tidx, n, cutlass.Int32(0), ld_keys)
+ for c in cutlass.range_constexpr(KPT):
+ tmax = cutlass.max(tmax, keys[c])
+ kmax = cute.arch.warp_reduction_max(tmax)
+ if lane == 0:
+ wa[warp] = kmax
+ if const_expr(self.prefilter):
+ ahead = cutlass.Int32(0)
+ for i in cutlass.range_constexpr(32):
+ v = cute.arch.shuffle_sync(tmax, offset=i)
+ if (v > tmax) | ((v == tmax) & (lane > i)):
+ ahead = ahead + 1
+ if ahead == self.per_warp - 1:
+ wc[warp] = tmax
+ cute.arch.sync_threads()
+ gmax = cutlass.Int32(-1)
+ if lane < nw:
+ gmax = wa[lane]
+ gmax = cute.arch.warp_reduction_max(gmax)
+ # Keys below `lo` are no candidates; all are when the window holds at most k.
+ lo = cutlass.Int32(0)
+ if const_expr(self.prefilter):
+ if n > k:
+ bound = cutlass.Int32(_I32_MAX)
+ if lane < nw:
+ bound = wc[lane]
+ bound = cute.arch.warp_reduction(bound, lambda x, y: cutlass.min(x, y))
+ lo = cutlass.max(bound, cutlass.Int32(0))
+
+ # 3. The candidates, compacted in block order: a scan of each slab's counts (warp
+ # totals alternating between two halves of `wsc`, one barrier a slab), carried
+ # from slab to slab. `m` counts them all, kept or not.
+ cflat_key = cute.make_tensor(ckey.iterator, cute.make_layout((KPT * nt,), stride=(1,)))
+ cflat_idx = cute.make_tensor(cidx.iterator, cute.make_layout((KPT * nt,), stride=(1,)))
+ carry = cutlass.Int32(0)
+ for s in cutlass.range(nslabs):
+ self._slab(keys, gK, sK, row, s, tidx, n, lo, ld_keys)
+ e0 = (s * nt + tidx) * KPT
+ cnt = cutlass.Int32(0)
+ for c in cutlass.range_constexpr(KPT):
+ if keys[c] >= 0:
+ cnt = cnt + 1
+ incl = cnt
+ for off in (1, 2, 4, 8, 16):
+ up = cute.arch.shuffle_sync_up(incl, offset=off)
+ if lane >= off:
+ incl = incl + up
+ half = (s & 1) * 32
+ if lane == 31:
+ wsc[half + warp] = incl
+ cute.arch.sync_threads()
+ above = cutlass.Int32(0)
+ if lane < warp:
+ above = wsc[half + lane]
+ pos = carry + cute.arch.warp_reduction_sum(above) + incl - cnt
+ slab_total = cutlass.Int32(0)
+ if lane < nw:
+ slab_total = wsc[half + lane]
+ slab_total = cute.arch.warp_reduction_sum(slab_total)
+ for c in cutlass.range_constexpr(KPT):
+ if keys[c] >= 0:
+ if pos < KPT * nt:
+ cflat_key[pos] = keys[c]
+ cflat_idx[pos] = e0 + c
+ pos = pos + 1
+ carry = carry + slab_total
+ m = carry
+ cute.arch.sync_threads()
+
+ # 4. The narrow kernel's selection over the list, or over the window when the list
+ # cannot hold every candidate.
+ diff = gmax ^ lo
+ hb = cutlass.Int32(-1)
+ if diff != 0:
+ hb = 31 - cutlass.Int32(cute.arch.clz(diff))
+ low = (cutlass.Int64(1) << cutlass.Int64(hb + 1)) - 1
+ mask0 = cutlass.Int32(cutlass.Int64(_I32_MAX) & (cutlass.Int64(-1) ^ low))
+ if tidx == 0:
+ state[0] = gmax & mask0 # threshold prefix
+ state[1] = mask0 # its mask
+ state[2] = cutlass.Int32(k) # still needed at the threshold
+ state[3] = m # keys in the threshold's bin
+ state[4] = cutlass.Int32(1) # narrowing done
+ state[6] = cutlass.max(hb - 10, 0) # the next pass's shift
+ if m > k:
+ if hb >= 0:
+ state[4] = cutlass.Int32(0)
+ cute.arch.sync_threads()
+ if m <= KPT * nt:
+ self._select_listed(
+ ckey,
+ cidx,
+ keys,
+ row,
+ tidx,
+ warp,
+ lane,
+ m,
+ seqlen,
+ hist,
+ cand_key,
+ cand_idx,
+ wb,
+ state,
+ mTopk,
+ mSel,
+ )
+ else:
+ self._select_window(
+ keys,
+ gK,
+ sK,
+ row,
+ tidx,
+ warp,
+ lane,
+ n,
+ nslabs,
+ lo,
+ m,
+ seqlen,
+ hist,
+ cand_key,
+ cand_idx,
+ wa,
+ wsc,
+ state,
+ mTopk,
+ mSel,
+ ld_keys,
+ )
+
+ # 5. Padding of the block list; the causal tail after the selected tokens, then -1.
+ total = n
+ if n > k:
+ total = cutlass.Int32(k)
+ for c in cutlass.range_constexpr(-(-k // nt)):
+ slot = tidx + c * nt
+ if slot < k:
+ if slot >= total:
+ mTopk[row, slot] = cutlass.Int32(-1)
+ tail_start = (visible // cr) * cr
+ tail_count = visible - tail_start
+ body = total * cr
+ if body > self.token_topk:
+ body = cutlass.Int32(self.token_topk)
+ for c in cutlass.range_constexpr(-(-self.final // nt)):
+ col = tidx + c * nt
+ if col < self.final:
+ if col >= body:
+ off = col - body
+ val = cutlass.Int32(-1)
+ if off < cr - 1:
+ if off < tail_count:
+ if tail_start + off < seqlen:
+ val = tail_start + off
+ mSel[row, col] = val
+
+
+#: One compiled launch per score-table geometry, k, compress ratio, token budget, index
+#: dtypes, PDL and wide-kernel configuration; the row count is a runtime argument.
+_COMPILED: dict = {}
+
+
+#: SMs of each device, read once.
+_SMS: dict = {}
+
+
+def _wide_config(rows: int, device) -> tuple:
+ """The wide kernel's `(threads, CTAs an SM, prefilter)` for a call of `rows` rows: one
+ 1024-thread CTA an SM (the largest window cache) while the rows fit one wave, two once
+ they do not, so one CTA's walks overlap the other's loads."""
+ sms = _SMS.get(device.index)
+ if sms is None:
+ sms = _SMS[device.index] = torch.cuda.get_device_properties(device).multi_processor_count
+ return (_WIDE_THREADS, 1 if rows <= sms else 2, True)
+
+
+_INDEX_TYPES = {torch.int32: cutlass.Int32, torch.int64: cutlass.Int64}
+
+
+def launch(
+ scores,
+ row_ends,
+ query_positions,
+ sequence_lengths,
+ topk_indices,
+ compress_ratio,
+ token_topk,
+ pdl=True,
+ out=None,
+):
+ """The launch behind `QsaTopkExpandDecode.__call__`: writes `topk_indices` and returns the
+ token selection `[rows, token_topk + compress_ratio - 1]` int32 (`out` when given)."""
+ rows, width = scores.shape
+ block_topk = topk_indices.shape[1]
+ sel = out
+ if sel is None:
+ sel = torch.empty(
+ (rows, token_topk + compress_ratio - 1), dtype=torch.int32, device=scores.device
+ )
+ if rows == 0:
+ return sel
+ pdl = bool(pdl and get_env_enable_pdl())
+ keys = scores.view(torch.int32)
+ key = (
+ width,
+ scores.stride(0),
+ block_topk,
+ compress_ratio,
+ token_topk,
+ query_positions.dtype,
+ sequence_lengths.dtype,
+ pdl,
+ None if width <= MAX_WIDTH else _wide_config(rows, scores.device),
+ )
+ ptrs = (
+ row_ends.data_ptr(),
+ query_positions.data_ptr(),
+ sequence_lengths.data_ptr(),
+ topk_indices.data_ptr(),
+ sel.data_ptr(),
+ )
+ fn = _COMPILED.get(key)
+ if fn is None:
+ gm = cute.AddressSpace.gmem
+ geometry = dict(
+ width=width,
+ key_stride=scores.stride(0),
+ block_topk=block_topk,
+ compress=compress_ratio,
+ token_topk=token_topk,
+ pdl=pdl,
+ )
+ if width <= MAX_WIDTH:
+ kernel = _TopkExpandKernel(**geometry)
+ else:
+ threads, ctas, prefilter = key[-1]
+ kernel = _TopkExpandWideKernel(
+ **geometry, threads=threads, ctas=ctas, prefilter=prefilter
+ )
+ fn = _COMPILED[key] = cute.compile(
+ kernel,
+ from_dlpack(keys, assumed_align=16).mark_layout_dynamic(leading_dim=1),
+ make_ptr(cutlass.Int32, ptrs[0], gm, assumed_align=4),
+ make_ptr(
+ _INDEX_TYPES[query_positions.dtype],
+ ptrs[1],
+ gm,
+ assumed_align=query_positions.element_size(),
+ ),
+ make_ptr(
+ _INDEX_TYPES[sequence_lengths.dtype],
+ ptrs[2],
+ gm,
+ assumed_align=sequence_lengths.element_size(),
+ ),
+ make_ptr(cutlass.Int32, ptrs[3], gm, assumed_align=4),
+ make_ptr(cutlass.Int32, ptrs[4], gm, assumed_align=4),
+ cutlass.Int32(rows),
+ cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+ options="--enable-tvm-ffi",
+ )
+ fn(keys, *ptrs, rows)
+ return sel
+
+
+def supports(scores) -> bool:
+ """Whether a score table this wide is one this entry serves (see `is_valid`)."""
+ return scores.shape[1] % KPT == 0 and scores.shape[1] <= MAX_WIDE_WIDTH
+
+
+#: The default target's score-table widths, `block_table.shape[1] * 32 // 4`: core's `max_seq_len`
+#: 9216 (288 pages) and MTP's 16384 (512 pages, `MAX_WIDTH` itself), served by the narrow
+#: kernel; the tune's `max_seq_len` 262144 with MTP's extra pages (8196 pages, 65568 columns)
+#: and without (65536), served by the wide one.
+_CORE_WIDTH = 2304
+_MTP_WIDTH = 4096
+_TUNE_WIDTH = 65568
+_TUNE_WIDTH_NO_MTP = 65536
+
+
+def _spread(edges: tuple, count: int, first: int, last: int) -> tuple:
+ """`edges`, then `count - len(edges)` values spread evenly over `[first, last]`."""
+ fill = count - len(edges)
+ return tuple(edges) + tuple(first + (r * (last - first)) // (fill - 1) for r in range(fill))
+
+
+#: 128 one-row decode requests at core's width. The edges are windows of 0 and 1 blocks,
+#: k - 1, k and k + 1 blocks (512, the `n <= k` copy-through and the first narrowing),
+#: and the full 2304-block window.
+_CORE_DECODE_128 = _spread(
+ (1, 2, 3, 6, 7, 2046, 2047, 2048, 2051, 2052, 9212, 9214, 9215), 128, 8, 9215
+)
+#: MTP7 verify at MTP's width: 128 requests x 8 consecutive rows, one starting where its
+#: rows cross k blocks and one ending at the full 4096-block window.
+_MTP7_VERIFY_1024 = tuple((start, 8) for start in _spread((1, 2041, 2044, 16376), 128, 9, 16375))
+#: Decode rows at the tune's 65536-wide table, by window: 0 and 1 blocks, k - 1, k and k + 1,
+#: one slab (4096 columns) - 1, itself and + 1, the 1024-thread CTA's window cache (11 slabs,
+#: 45056 blocks) - 1, itself and + 1, the whole table, and four between.
+_WIDE_EDGES_16 = (
+ 2,
+ 3,
+ 2043,
+ 2047,
+ 2051,
+ 16380,
+ 16384,
+ 16388,
+ 180220,
+ 180224,
+ 180228,
+ 262143,
+ 50000,
+ 99999,
+ 150001,
+ 220000,
+)
+
+
+class QsaTopkExpandDecode(OpWrapper):
+ """`qsa_topk_select` (zero row starts, `row_ends` the visible block counts),
+ `qsa_canonical_selection` and `qsa_expand_block_indices`, in one launch: one CTA a row, for
+ decode rows and a context chunk's rows alike.
+
+ Per row: the `block_topk` largest of `scores[row, :row_ends[row]]` (non-negative fp32),
+ ties at the threshold to the higher block, written to `topk_indices[row]` in block order
+ and padded with -1 -- the composition's selection after `qsa_canonical_selection`; the
+ returned token selection is the expansion entry's: each selected group's members below
+ `sequence_lengths[row]`, packed at the front, then the incomplete group at the causal
+ frontier (`query_positions[row] + 1` visible tokens), then -1. `topk_indices` is written
+ in place, and the selection into `out` when given (a contiguous int32
+ `[rows, token_topk + compress_ratio - 1]`, e.g. a row range of the step's whole
+ selection); nothing else is written.
+
+ Fusion boundary: the selection and the expansion. Scoring is the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # `width` is the score table's column count (and row stride): a compile key, and the
+ # thread count (`_threads`: 640 at 2560, 576 at 2304, 1024 at 4096).
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="decode, 4 rows across the budget, one below it",
+ spec=dict(rows=4, pos=(8700, 4095, 1500, 9900), width=2560),
+ ),
+ # The target writes a context chunk's rows through `out` into its rows of the
+ # step's selection; relu-sparse scores put the zero class (negative zeros among it)
+ # across the budget, the ties the canonical rule decides.
+ Cell(
+ why="context chunk: 2048 rows of one request (positions 6144..8191), relu-sparse "
+ "scores, written through `out` into rows of a step-sized selection",
+ spec=dict(rows=2048, chunk=(6144, 2048), relu_sparse=True, out_margin=64, width=2560),
+ ),
+ Cell(
+ why="decode, 1 row of scores within 5% of each other",
+ spec=dict(rows=1, pos=(8700,), narrow=True, width=2560),
+ ),
+ # Serving's decode rows at 8K: positive, log-normally spread over ~2 octaves, a
+ # zero or two; the first pass leaves tens to ~CAP candidates at the threshold.
+ Cell(
+ why="decode, 16 rows of log-normal scores with a few zeros: threshold bins of "
+ "up to CAP candidates ranked pairwise, ties included",
+ spec=dict(
+ rows=16, pos=tuple(8192 + 37 * r for r in range(16)), lognormal=True, width=2560
+ ),
+ ),
+ # The default target's call sites, at its own widths (`_CORE_WIDTH`, `_MTP_WIDTH`).
+ Cell(
+ why="core decode, 1 row at the 2304-wide table, its full window, relu-sparse",
+ spec=dict(rows=1, pos=(9215,), relu_sparse=True, width=_CORE_WIDTH),
+ ),
+ Cell(
+ why="core decode batch, 128 rows at the 2304-wide table (max_batch_size): windows "
+ "of 0, 1, k-1, k, k+1 and 2304 blocks and 115 spread between, relu-sparse",
+ spec=dict(rows=128, pos=_CORE_DECODE_128, relu_sparse=True, width=_CORE_WIDTH),
+ ),
+ Cell(
+ why="MTP7 verify, 1024 rows (128 requests x 8) at the 4096-wide table, up to its "
+ "full window, relu-sparse",
+ spec=dict(rows=1024, requests=_MTP7_VERIFY_1024, relu_sparse=True, width=_MTP_WIDTH),
+ ),
+ Cell(
+ why="core context chunk, 2048 rows (positions 6144..8191) at the 2304-wide table, "
+ "relu-sparse, written through `out`",
+ spec=dict(
+ rows=2048, chunk=(6144, 2048), relu_sparse=True, out_margin=64, width=_CORE_WIDTH
+ ),
+ ),
+ Cell(
+ why="envelope's largest row: an 8192-row chunk (max_num_tokens, positions "
+ "8192..16383) at the 4096-wide table, relu-sparse, written through `out`",
+ spec=dict(
+ rows=8192, chunk=(8192, 8192), relu_sparse=True, out_margin=64, width=_MTP_WIDTH
+ ),
+ ),
+ # The tune's tables (`max_seq_len` 262144): 65568 columns with MTP's extra pages,
+ # 65536 without, past `MAX_WIDTH`, so the wide kernel's. Relu-sum scores are the
+ # indexer's form (four heads' ReLUs summed): the prefilter keeps ~5% of a window.
+ Cell(
+ why="wide: MTP3 verify, 4 rows (positions 84000..84003) at the 65568-wide table, "
+ "relu-sum scores",
+ spec=dict(rows=4, requests=((84000, 4),), relu_sum=True, width=_TUNE_WIDTH, graph=True),
+ ),
+ Cell(
+ why="wide: 16 decode rows at the 65536-wide table, windows at the edges of k, a "
+ "slab, the window cache and the table, relu-sum scores",
+ spec=dict(
+ rows=16, pos=_WIDE_EDGES_16, relu_sum=True, width=_TUNE_WIDTH_NO_MTP, graph=True
+ ),
+ ),
+ Cell(
+ why="wide: 4 rows of log-normal scores at 65568: threshold bins ranked pairwise, "
+ "ties included",
+ spec=dict(rows=4, pos=(100000, 30000, 4095, 2047), lognormal=True, width=_TUNE_WIDTH),
+ ),
+ Cell(
+ why="wide: 2 rows of relu-sparse scores at 65536 (~1.4% positive): the prefilter "
+ "keeps the zeros, more candidates than the list holds, the whole-window path; at "
+ "30000 blocks the zero class (negative zeros among it) straddles the budget",
+ spec=dict(rows=2, pos=(262143, 120000), relu_sparse=True, width=_TUNE_WIDTH_NO_MTP),
+ ),
+ Cell(
+ why="wide: 3 rows of equal scores at 65568 (windows 25000, 513 and 512 blocks): "
+ "the tie class is the window, the highest blocks taken",
+ spec=dict(rows=3, pos=(100000, 2051, 2047), equal=True, width=_TUNE_WIDTH),
+ ),
+ Cell(
+ why="wide: ht_dep4's context chunk, 8160 rows (positions 57000..65159) at the "
+ "65536-wide table, relu-sum, written through `out`",
+ spec=dict(
+ rows=8160,
+ chunk=(57000, 8160),
+ relu_sum=True,
+ out_margin=64,
+ width=_TUNE_WIDTH_NO_MTP,
+ ),
+ ),
+ )
+
+ note = """
+ Metadata consumed: `qsa_visible_blocks` (`row_ends`), the rows' logical positions
+ and sequence lengths. Preconditions: zero row starts; `scores` fp32 `[rows, W]`, unit
+ column stride, 16-byte aligned rows, `W % 4 == 0` and `W <= 131072` (up to 4096 the
+ narrow kernel, four keys a thread and at most 1024 threads; wider tables the wide
+ kernel, `_TopkExpandWideKernel`); `topk_indices` int32 `[rows, K]` contiguous; `row_ends` int32,
+ the positions and lengths int32 or int64, each a contiguous `[rows]`. Ties: the higher
+ block at the threshold, as `qsa_canonical_selection` rewrites the CUDA top-k's pick.
+
+ Compiled once per score-table width and row stride (the block table's, fixed for an
+ engine), k, compress ratio, token budget, index dtypes and PDL: the row count is a
+ runtime argument, so an eager step's context rows, any count, reuse the kernel (read
+ from the launch, not pinned by a test).
+
+ JIT-compiled: `cute.compile` (`--enable-tvm-ffi`) on a key's first call, for the device
+ that makes it. The cells ran that path on sm_103 (GB300) at widths 2560, 2304 and 4096.
+
+ Certified, every cell bitwise against `reference` (block list and token selection):
+ - tp1's cells at 2560 columns: decode 1 row of near-equal scores, 4 rows, 16 rows of
+ log-normal scores (threshold bins ranked pairwise), and a 2048-row context chunk.
+ - The default target's cells at its own widths (`_CORE_WIDTH` 2304, `_MTP_WIDTH` 4096), with
+ relu-sparse scores: the zero class, negative zeros among it, straddles the 512-block
+ budget and is resolved by the tie rule. They are:
+ - decode, 1 and 128 rows at 2304, with windows of 0, 1, 511, 512, 513 and 2304 blocks among them;
+ - MTP7 verify, 1024 rows (128 requests x 8) at 4096;
+ - a 2048-row context chunk at 2304, through `out`;
+ - an 8192-row chunk (`max_num_tokens`) at 4096, through `out`.
+ - The tune's tables (`max_seq_len` 262144: 65568 columns with MTP's pages, 65536
+ without), the wide kernel: a 4-row MTP3 verify at ~84k tokens and 16 decode rows at
+ the edges of k, a 4096-column slab, the 45056-block window cache and the whole table
+ (relu-sum scores); 4 rows of log-normal scores (threshold bins ranked pairwise); 2
+ relu-sparse rows (more candidates than the list holds: the whole-window path, the
+ zero class straddling the budget); 3 rows of equal scores; and ht_dep4's 8160-row
+ context chunk through `out` (two CTAs an SM past one wave of rows).
+ - A CUDA-graph replay at 1x2560, 1x2304, 128x2304, 1024x4096, 4x65568, 16x65536 and
+ 2x65536. The scores are refilled before each replay, and the selection, read after
+ the wait, must equal the reference.
+ The widths are derived from `max_seq_len` (288, 512 and 8192 (+4) pages of 32 tokens),
+ not read off an engine. 4096 is `MAX_WIDTH`: 1024 threads, the largest CTA; wider
+ tables take the wide kernel, compiled per row-count class (one CTA an SM up to the
+ SM count, two past it).
+
+ `is_valid` refuses `token_topk != k * compress_ratio`. It sizes the selection's rows
+ while `k * compress_ratio` places the tokens, so a smaller budget writes past each row.
+
+ Not certified:
+ - threshold bins ranked pairwise at 2304 and 4096 columns (the log-normal cell is at 2560);
+ - decode batches other than 1, 4, 16 and 128 rows;
+ - a NaN inside a window (it ranks above every finite score: the caller keeps the window finite);
+ - index dtypes other than int64 positions and int32 lengths.
+
+ Index limits, read from the source, not probed. The row count is an Int32 argument and
+ the row layouts are built from it, so these products must stay below 2^31:
+ - rows x row stride: 524288 rows at 4096 columns, 16383 at 131072 (`MAX_WIDE_WIDTH`);
+ - rows x (token_topk + compress_ratio - 1): about 1.05 M rows at 2051.
+ Positions are read as Int32. The envelope's largest call is 8192 rows.
+ """
+
+ def __call__(
+ self,
+ scores,
+ row_ends,
+ query_positions,
+ sequence_lengths,
+ topk_indices,
+ compress_ratio,
+ token_topk,
+ out=None,
+ ):
+ return launch(
+ scores,
+ row_ends,
+ query_positions,
+ sequence_lengths,
+ topk_indices,
+ compress_ratio,
+ token_topk,
+ out=out,
+ )
+
+ def reference(
+ self,
+ scores,
+ row_ends,
+ query_positions,
+ sequence_lengths,
+ topk_indices,
+ compress_ratio,
+ token_topk,
+ out=None,
+ ):
+ """`(topk_indices_after, selection)`: a stable descending sort per row (ties to the
+ higher block), the k first in block order, then the expansion entry's reference."""
+ rows, _ = scores.shape
+ k = topk_indices.shape[1]
+ out = torch.full_like(topk_indices, -1)
+ for r in range(rows):
+ n = int(row_ends[r])
+ s = scores[r, :n].float().clamp(min=0.0)
+ # A stable sort of the reversed row puts the higher block first among equals.
+ order = (n - 1) - torch.sort(torch.flip(s, [0]), descending=True, stable=True).indices[
+ :k
+ ]
+ chosen = torch.sort(order).values
+ out[r, : chosen.numel()] = chosen.to(torch.int32)
+ sel = _expand.reference(
+ out,
+ query_positions,
+ sequence_lengths,
+ compress_ratio=compress_ratio,
+ token_topk=token_topk,
+ )
+ return out, sel
+
+ def is_valid(
+ self,
+ scores,
+ row_ends,
+ query_positions,
+ sequence_lengths,
+ topk_indices,
+ compress_ratio,
+ token_topk,
+ out=None,
+ ):
+ assert scores.dtype == torch.float32 and scores.stride(1) == 1, (
+ "fp32 scores, unit column stride"
+ )
+ assert scores.shape[1] % KPT == 0 and scores.shape[1] <= MAX_WIDE_WIDTH, (
+ "W % 4 == 0, W <= 131072"
+ )
+ assert scores.stride(0) % KPT == 0 and scores.data_ptr() % 16 == 0, "16-byte aligned rows"
+ assert topk_indices.dtype == torch.int32 and topk_indices.is_contiguous(), (
+ "int32 topk_indices"
+ )
+ rows = scores.shape[0]
+ assert row_ends.dtype == torch.int32 and row_ends.shape == (rows,), "int32 row_ends [rows]"
+ for t in (query_positions, sequence_lengths):
+ assert t.dtype in _INDEX_TYPES and t.shape == (rows,), "int32/int64 [rows]"
+ assert all(t.is_contiguous() for t in (row_ends, query_positions, sequence_lengths)), (
+ "contiguous row_ends, positions and lengths"
+ )
+ assert token_topk == topk_indices.shape[1] * compress_ratio, (
+ "token_topk = k * compress_ratio"
+ )
+ assert out is None or (
+ out.dtype == torch.int32
+ and out.is_contiguous()
+ and out.shape == (scores.shape[0], token_topk + compress_ratio - 1)
+ ), "out: contiguous int32 [rows, token_topk + compress_ratio - 1]"
+
+
+qsa_topk_expand_decode = QsaTopkExpandDecode()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_topk_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_topk_select.py
new file mode 100644
index 000000000000..7abaa54f8115
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_topk_select.py
@@ -0,0 +1,299 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""QSA block selection: an exact per-row top-k over a caller-supplied column window."""
+
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import torch
+
+import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+#: This checkpoint's block budget: `indexer_budget` 2048 tokens over
+#: `indexer_compress_ratio` 4. The op's own default, 2048, is the *token* budget.
+_BLOCK_TOPK = 512
+#: Score-row width the target hands over: a 260-page block table at 32 tokens
+#: per block, compressed 4:1.
+_SCORE_COLUMNS = 260 * 32 // 4
+#: The default target's widest score row: the tune envelope's `max_seq_len`
+#: 262144 is an 8192-page block table at 32 tokens per block, compressed 4:1.
+#: Core's 9216 (288 pages, 2304 columns) and MTP's 16384 (512 pages, 4096
+#: columns) are narrower; the kernel picks its path by each row's window, so
+#: the windows below span theirs.
+_ENVELOPE_SEQ_LEN = 262144
+_ENVELOPE_COLUMNS = _ENVELOPE_SEQ_LEN // 4
+#: 128 decode requests: the window edges first -- empty (position 2), one
+#: block, 511 / 512 / 513 blocks around the budget (512 is the kernel's
+#: copy-through shortcut, 513 its first histogram row), core's and MTP's widest
+#: rows, the envelope's less one column -- then 120 spread to the envelope.
+_DECODE_POSITIONS = (2, 3, 2046, 2047, 2051, 9215, 16383, _ENVELOPE_SEQ_LEN - 2) + tuple(
+ 4096 + (_ENVELOPE_SEQ_LEN - 1 - 4096) * i // 119 for i in range(120)
+)
+#: 128 verify requests at max_draft_len 7, 8 consecutive positions each: one
+#: from position 0, two straddling the budget, core's and MTP's last rows, the
+#: envelope's last (a full 65536-column window), then 122 spread between.
+_VERIFY_STARTS = (0, 2040, 2044, 9208, 16376, _ENVELOPE_SEQ_LEN - 8) + tuple(
+ 4096 + (_ENVELOPE_SEQ_LEN - 16 - 4096) * i // 121 for i in range(122)
+)
+#: Rows from which a contiguous launch runs the CuTe DSL sibling: the CUDA op
+#: serves rows at or past 12288 with its slower radix variant; below it the CUDA
+#: op needs no JIT compilation.
+_CUTE_DSL_MIN_ROWS = 12288
+
+
+class TopkExpectation(NamedTuple):
+ """What a correct selection must satisfy.
+
+ The op's output order is unspecified, so `reference` cannot return the
+ indices themselves; it returns each row's expected top values plus the
+ frame the test needs to read the op's offsets back as values.
+ """
+
+ #: `[rows, index_topk]` fp32: the window's top values, descending, `-inf`-padded.
+ values: torch.Tensor
+ #: The logits as fp32 with every out-of-window column at `-inf`.
+ logits: torch.Tensor
+ #: int64 `row_starts`, to lift window-local offsets to columns.
+ row_starts: torch.Tensor
+ #: int64 window length per row, `max(row_ends - row_starts, 0)`.
+ window: torch.Tensor
+
+
+class QsaTopkSelect(OpWrapper):
+ """Each row's top-`index_topk` window-local column offsets, written into `indices`.
+
+ Wraps `torch.ops.trtllm.indexer_topk_prefill` (CUDA, compiled into the
+ wheel) 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
+ is written, so `indices` needs no pre-fill. Returns `indices`.
+
+ Exact, not approximate: the kernel resolves the threshold over the full 32
+ bits of the fp32 pattern. The **order of the offsets is unspecified** and
+ not stable between identical launches (candidates are appended through a
+ shared-memory atomic). Ties are broken arbitrarily within the tied value.
+
+ Fusion boundary: selection only. Scoring, the causal bound the caller
+ encodes in `row_ends`, and the block-to-token expansion are outside.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # The target calls this with `index_topk = 512` over the `[rows, 2080]`
+ # score matrix `qsa_paged_index_scores` returns with
+ # `only_visible_blocks=True` -- so every column at or past a row's
+ # `row_end` is allocator garbage. Those cells poison it with NaN: a NaN
+ # inside the window would be selected (see `note`), so they prove the
+ # window is honoured, not merely that the tail happened to be small.
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="decode 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), the CUDA op",
+ spec=dict(
+ rows=8191,
+ cols=_SCORE_COLUMNS,
+ index_topk=_BLOCK_TOPK,
+ windows=("chunk", 0),
+ logits="relu",
+ tail="nan",
+ ),
+ ),
+ Cell(
+ 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",
+ ),
+ ),
+ # The default target (tp4 / tep4 / dep4): the selection is per row and
+ # replicated across ranks, so the member changes only the row count.
+ Cell(
+ why="decode: one row at the 262144-token envelope, a 65535-block window, NaN tail",
+ spec=dict(
+ rows=1,
+ cols=_ENVELOPE_COLUMNS,
+ index_topk=_BLOCK_TOPK,
+ windows=("positions", (_ENVELOPE_SEQ_LEN - 2,)),
+ logits="relu",
+ tail="nan",
+ ),
+ ),
+ Cell(
+ why="decode batch 128: windows empty, 511 / 512 / 513 blocks and up to the envelope",
+ spec=dict(
+ rows=128,
+ cols=_ENVELOPE_COLUMNS,
+ index_topk=_BLOCK_TOPK,
+ windows=("positions", _DECODE_POSITIONS),
+ logits="relu",
+ tail="nan",
+ ),
+ ),
+ Cell(
+ why="verify MTP7 at 128 requests: 1024 rows, windows up to the envelope's 65536",
+ spec=dict(
+ rows=1024,
+ cols=_ENVELOPE_COLUMNS,
+ index_topk=_BLOCK_TOPK,
+ windows=("positions", tuple(s + j for s in _VERIFY_STARTS for j in range(8))),
+ logits="relu",
+ tail="nan",
+ ),
+ ),
+ Cell(
+ why="context chunk: max_num_tokens 8192 rows ending a 262144-token prompt",
+ spec=dict(
+ rows=8192,
+ cols=_ENVELOPE_COLUMNS,
+ index_topk=_BLOCK_TOPK,
+ windows=("chunk", _ENVELOPE_SEQ_LEN - 8192),
+ logits="relu",
+ tail="nan",
+ ),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none. Stateless; every input is an argument.
+
+ Preconditions (the op checks none of these):
+ - `0 <= row_starts[r] <= row_ends[r] <= logits.shape[1]`. An empty window
+ is legal and yields an all-`-1` row; a `row_end` past the last column
+ reads out of bounds (not probed).
+ - Logits inside the window are finite. `-inf` is never selected, but a
+ NaN ranks above every finite value and IS selected (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, not pinned by a test): bf16/fp16 `logits`
+ ("expected scalar type Float"), int64 `indices` ("expected scalar type
+ Int"), `row_starts` / `row_ends` shorter than the row count ("one entry per
+ row").
+
+ The CUDA op dispatches by row index, not column count: rows below 12288 run
+ an insertion-sort variant, rows at or above a radix variant, in one launch.
+ A launch of `_CUTE_DSL_MIN_ROWS` (12288) rows or more with contiguous
+ logits runs the CuTe DSL sibling instead, as QSA's stock indexer does; it
+ selects identical sets in a different order, ties broken arbitrarily.
+ `max_num_tokens` 16384 reaches it at prefill-heavy mixed steps; smaller
+ launches, every decode one among them, keep the compiled CUDA op, so a
+ first captured decode JIT-compiles nothing. `torch.ops.trtllm.indexer_topk_decode`
+ is the request-grouped sibling; QSA does not use it.
+
+ The CUDA op picks its path per row by the window's length, never by the
+ row width: a window of at most `index_topk` copies offsets 0..len-1
+ through (unsorted, no histogram); a longer one runs the histogram passes
+ over the fp32 bits into a 2048-item final buffer. The width enters only as
+ the row stride, and a row that does not start 16-byte aligned is read with
+ a scalar head and tail. So the default target's cells, at the 262144-token
+ envelope's 65536-column rows, take windows from empty through 1, 511, 512,
+ 513, core's 2304 (`max_seq_len` 9216) and MTP's 4096 (16384) to 65535 and
+ 65536 -- rows 1, 128 (decode), 1024 (MTP7 verify at 128 requests) and 8192
+ (a `max_num_tokens` context chunk). All four run the compiled-in CUDA op
+ (below 12288 rows); the JIT-compiled CuTe DSL sibling runs at the 16384-row
+ cell, on sm_103 in the same test file. Attention DP and TP do not change the
+ op's shape, only its row count.
+
+ Measured, not a cell: a tied class far wider than the final buffer (this
+ checkpoint's relu scores leave rows with fewer positive blocks than the
+ budget, the rest exactly 0) still returns 512 distinct in-window offsets
+ whose values are the window's top 512, at 0 to 3000 positive scores in a
+ 65536-column window.
+
+ CUDA graphs: captured at `[256, 2048]` (32 requests at max_draft_len 7) and
+ at 1, 128 and 1024 rows of 65536 columns, `k=512`, each over full windows,
+ and replayed against new logits and a new `row_ends` through the captured
+ addresses (the 128-row replay at 511 blocks, the copy-through path).
+ Bitwise comparison of two replays is not a valid check (order is
+ unspecified).
+
+ Hard limits: `indices` is addressed as `row * index_topk + slot` and the
+ logits row as `row * stride(0)`, both in 64-bit; `num_columns` and
+ `stride(0)` pass through int32, so a row past 2^31 - 1 elements is out of
+ reach (the envelope's is 65536). The largest certified extents are logits
+ 8192 x 65536 (2^29 elements) and `indices` 16384 x 512.
+ """
+
+ def __call__(
+ self,
+ logits: torch.Tensor,
+ row_starts: torch.Tensor,
+ row_ends: torch.Tensor,
+ indices: torch.Tensor,
+ index_topk: int = 2048,
+ ) -> torch.Tensor:
+ 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(
+ self,
+ logits: torch.Tensor,
+ row_starts: torch.Tensor,
+ row_ends: torch.Tensor,
+ indices: torch.Tensor,
+ index_topk: int = 2048,
+ ) -> TopkExpectation:
+ """Each window's top values, sorted. `indices` (the output buffer) is ignored."""
+ num_rows, num_cols = logits.shape
+ columns = torch.arange(num_cols, device=logits.device)[None, :]
+ starts, ends = row_starts.long()[:, None], row_ends.long()[:, None]
+ in_window = (columns >= starts) & (columns < ends)
+ masked = logits.float().masked_fill(~in_window, float("-inf"))
+ values = torch.sort(masked, dim=1, descending=True).values
+ if num_cols < index_topk:
+ values = torch.nn.functional.pad(
+ values, (0, index_topk - num_cols), value=float("-inf")
+ )
+ window = (row_ends.long() - row_starts.long()).clamp_min(0)
+ return TopkExpectation(values[:, :index_topk], masked, row_starts.long(), window)
+
+ def is_valid(
+ self,
+ logits: torch.Tensor,
+ row_starts: torch.Tensor,
+ row_ends: torch.Tensor,
+ indices: torch.Tensor,
+ index_topk: int = 2048,
+ ) -> None:
+ # The op checks only `indices.size(1) >= index_topk`, then writes a flat
+ # run of `index_topk`-wide rows and ignores the real row stride, so in a
+ # wider buffer rows land across their neighbours and the last is never
+ # written.
+ assert indices.dim() == 2 and indices.shape[1] == index_topk, (
+ f"indices must be [num_rows, index_topk]; got {tuple(indices.shape)} "
+ f"for index_topk={index_topk}"
+ )
+ assert indices.is_contiguous(), "indices must be contiguous"
+
+
+qsa_topk_select = QsaTopkSelect()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qwen4_exp_attn_decode_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qwen4_exp_attn_decode_prologue.py
new file mode 100644
index 000000000000..a3b705eb80d5
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qwen4_exp_attn_decode_prologue.py
@@ -0,0 +1,1023 @@
+# 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 program per row and v head stores v; one
+more program per row stores 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.
+
+Over an NVFP4 K/V pool the K/V append is `qsa_paged_kv_store_nvfp4`'s quantizing store:
+the programs holding k heads quantize their rotated, bf16-rounded rows from registers and
+the v programs quantize v from the projection row, with the store's arithmetic, so the
+pool's e2m1 bytes and e4m3 block scales are bitwise the composition's.
+
+A verify step (`rows_per_request` > 1, each request's rows consecutive at consecutive
+positions) writes several raw keys of one group in the launch: at a group's last token
+the earlier members held by the same request's preceding rows are taken from their
+projection rows, the others from the index plane.
+
+At up to `SCORE_ROWS` decode rows the index scores can ride the launch as well
+(`index_scores`): further programs run `qsa_paged_index_scores`' decode tiles, each
+rotating the row's index query itself and, at a group's last token, compressing the last
+visible block's key itself -- the same operations on the same values as the programs that
+store them, so the scores are bitwise the separate launch's.
+
+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, the K/V second-level scales, the
+compressed keys a score tile reads, and at a group end the group's earlier keys from the
+index plane with their rotation -- is read before `griddepcontrol.wait`; after it a
+program reads only projection rows (a score tile of a row without a page re-reads its
+last visible key, which the launch does not write).
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_index_scores import (
+ qsa_paged_index_scores as _index_scores,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_index_store import (
+ qsa_paged_index_store as _index_store,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_kv_store import (
+ qsa_paged_kv_store as _kv_store,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_kv_store_nvfp4 import (
+ _e2m1x2,
+ _quantize_row,
+ _rcp_approx,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_kv_store_nvfp4 import (
+ qsa_paged_kv_store_nvfp4 as _kv_store_nvfp4,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_prefill_compress import (
+ qsa_prefill_compress as _compress,
+)
+from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+
+from .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
+
+#: Rows up to which the index scores can ride the launch, as `qsa_paged_index_scores`' decode
+#: tiles (SCORE_BN blocks a program, two warps; the norm + rope plan is two warps there too).
+#: GB300, 12 chained layers at ~8K tokens, us a layer: 1 row 15.0 -> 13.7, 2 rows 15.8 -> 14.3;
+#: at 4 rows the 324 extra programs a row make it slower than the separate launch (17.1 -> 21.5).
+SCORE_ROWS = 2
+SCORE_BN = 8
+
+
+@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_tile(src, heads, head_mask, w2, wp2, c, s, eps, D: tl.constexpr, ROT: tl.constexpr,
+ out_dtype: tl.constexpr): # 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 to `out_dtype`."""
+ 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
+ return (y * c[None, :] + yp * s[None, :]).to(out_dtype)
+
+
+@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
+ """`_norm_rope_tile` stored at `dst`. Returns the stored (rounded) tile."""
+ out = _norm_rope_tile(src, heads, head_mask, w2, wp2, c, s, eps, D, ROT, dst.dtype.element_ty)
+ d = tl.arange(0, D)
+ tl.store(dst + heads[:, None] * dst_head_stride + d[None, :], out, mask=head_mask[:, None])
+ return out
+
+
+@triton.jit
+def _quantize_tile(x, global_scale, ROWS: tl.constexpr, D: tl.constexpr):
+ """`qsa_paged_kv_store_nvfp4`'s `_quantize_row` over a `[ROWS, D]` register tile:
+ (packed e2m1 bytes `[ROWS, D/2]`, e4m3 scale bytes `[ROWS, D/16]`), the same operations
+ on the same values."""
+ lo, hi = tl.split(tl.reshape(x.to(tl.float32), [ROWS, D // 2, 2]))
+ pair_max = tl.maximum(tl.abs(lo), tl.abs(hi))
+ block_max = tl.max(tl.reshape(pair_max, [ROWS, D // 16, 8]), axis=2)
+ sf_value = global_scale * (block_max * _rcp_approx(tl.full([ROWS, D // 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, [ROWS, D // 16, 8]) * out_scale[:, :, None], [ROWS, D // 2]
+ )
+ hi_scaled = tl.reshape(
+ tl.reshape(hi, [ROWS, D // 16, 8]) * out_scale[:, :, None], [ROWS, D // 2]
+ )
+ packed = (_e2m1x2(lo_scaled, hi_scaled) & 0xFF).to(tl.uint8)
+ return packed, sf_code.to(tl.uint8, bitcast=True)
+
+
+@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,
+ k_scale_cache,
+ v_scale_cache,
+ s_ks_page,
+ s_ks_head,
+ s_vs_page,
+ s_vs_head,
+ kv_scale_orig_quant,
+ scores,
+ s_sc_row,
+ score_scale,
+ 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,
+ NVFP4: tl.constexpr,
+ IN_LAUNCH_GROUPS: tl.constexpr,
+ SCORE_TILES: tl.constexpr,
+ SCORE_BLOCKS: tl.constexpr,
+ SBN: tl.constexpr,
+ SBH: tl.constexpr,
+ USE_PDL: tl.constexpr,
+):
+ token = tl.program_id(0).to(tl.int64)
+ program = tl.program_id(1)
+ if USE_PDL:
+ # The consumer (the index scores, or with the scores here the top-k) reads
+ # nothing of this launch before its 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 NVFP4:
+ k_global = tl.load(kv_scale_orig_quant + 1)
+ 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:
+ is_k = (main_heads >= HQ) & (main_heads < HQ + HK) & live
+ k_heads = (main_heads - HQ).to(tl.int64)[:, None]
+ if NVFP4:
+ # The rounded rows the composition's store reads back, quantized here.
+ k_data, k_sf = _quantize_tile(out, k_global, HPB, D)
+ pairs = tl.arange(0, D // 2)
+ blocks = tl.arange(0, D // 16)
+ tl.store(
+ k_cache + page * s_kc_page + k_heads * s_kc_head + slot * s_kc_tok
+ + pairs[None, :],
+ k_data,
+ mask=is_k[:, None],
+ ) # fmt: skip
+ # K scales token-major within the head page.
+ tl.store(
+ k_scale_cache + page * s_ks_page + k_heads * s_ks_head
+ + slot * (D // 16) + blocks[None, :],
+ k_sf,
+ mask=is_k[:, None],
+ ) # fmt: skip
+ else:
+ d = tl.arange(0, D)
+ tl.store(
+ k_cache + page * s_kc_page + k_heads * s_kc_head + slot * s_kc_tok
+ + d[None, :],
+ out.to(k_cache.dtype.element_ty),
+ mask=is_k[:, None],
+ ) # fmt: skip
+ 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
+ elif program >= MAIN_PROGRAMS + 2 + HK:
+ if SCORE_TILES > 0:
+ # `qsa_paged_index_scores`' decode tile (its prefetching kernel at up to four
+ # rows: SBN blocks, the heads padded to SBH, two warps) with the row's index
+ # query rotated here as the index program does, and at a group's last token the
+ # last visible block's key compressed here as the index-key program does.
+ tile = program - (MAIN_PROGRAMS + 2 + HK)
+ visible_blocks = (logical + 1) // CR
+ if tile * SBN >= visible_blocks:
+ if USE_PDL:
+ tl.extra.cuda.gdc_wait()
+ return
+ dims = tl.arange(0, ID)
+ columns = tile * SBN + tl.arange(0, SBN)
+ valid = (columns < SCORE_BLOCKS) & (columns < visible_blocks)
+ anchors = columns * CR + CR - 1
+ key_pages = tl.load(
+ block_table + request * s_bt_req + (anchors // TPB) * s_bt_col, mask=valid, other=0
+ ).to(tl.int64)
+ key_ptrs = (
+ index_cache
+ + key_pages[:, None] * s_ic_page
+ + (anchors % TPB)[:, None] * s_ic_slot
+ + dims[None, :]
+ )
+ key_mask = valid[:, None] & (dims[None, :] < ID)
+ keys = tl.load(key_ptrs, mask=key_mask, other=0.0, eviction_policy="evict_first")
+ 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)
+ # The block this launch compresses (the last visible one, at a group's last
+ # token): the group's earlier raw keys, the first token's coordinates, their
+ # rotation and the key norm weight, as the index-key program reads them.
+ fresh_tile = ((logical + 1) % CR == 0) & ((visible_blocks - 1) // SBN == tile)
+ fresh = (columns == visible_blocks - 1) & fresh_tile
+ compress_here = fresh_tile & live
+ 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=compress_here,
+ 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=(compress_here & (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=compress_here,
+ 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=compress_here, other=0)
+ pos_h = tl.load(first_base + s_pc_axis, mask=compress_here, other=0)
+ pos_w = tl.load(first_base + 2 * s_pc_axis, mask=compress_here, other=0)
+ cosine, sine = _compress_factors(
+ pos_t,
+ pos_h,
+ pos_w,
+ cos_sin,
+ stride_cs,
+ ROTARY_PAIRS,
+ IS_MROPE,
+ MROPE_H,
+ MROPE_W,
+ compress_here,
+ )
+ key_weight = tl.load(index_k_weight + dims).to(tl.float32) + 1.0
+ if USE_PDL:
+ tl.extra.cuda.gdc_wait()
+ query = _norm_rope_tile(
+ row + IQ_COL, index_heads, index_heads < IH, w2, wp2, c, s, eps, ID, ROT,
+ index_q.dtype.element_ty,
+ ) # fmt: skip
+ if fresh_tile:
+ if compress_here:
+ current_k = tl.load(row + IQ_COL + IH * ID + dims)
+ group = tl.where(
+ offsets[:, None] == CR - 1,
+ current_k.to(index_cache.dtype.element_ty)[None, :],
+ group,
+ )
+ pooled = (tl.sum(group.to(tl.float32), axis=0) / CR).to(tl.bfloat16)
+ fresh_key = tl.reshape(
+ _compress_norm_rope(
+ pooled[None, :], key_weight, cosine, sine, eps, ID, ROTARY_PAIRS
+ ),
+ (ID,),
+ )
+ keys = tl.where(fresh[:, None], fresh_key[None, :].to(keys.dtype), keys)
+ else:
+ # No page: nothing is written; the key is the plane's.
+ keys = tl.where(
+ fresh[:, None],
+ tl.load(key_ptrs, mask=key_mask & fresh[:, None], other=0.0),
+ keys,
+ )
+ # The query as `[ID, SBH]`, heads past IH zero (the scores kernel's masked load).
+ padded = tl.reshape(tl.broadcast_to(query[None, :, :], (SBH // IH, IH, ID)), (SBH, ID))
+ head_rows = tl.arange(0, SBH)
+ padded = tl.where((head_rows < IH)[:, None], padded, tl.zeros_like(padded))
+ per_head = tl.dot(keys, tl.trans(padded), out_dtype=tl.float32)
+ block_scores = tl.sum(
+ tl.where(head_rows[None, :] < IH, tl.maximum(per_head, 0.0), 0.0), axis=1
+ )
+ tl.store(scores + token * s_sc_row + columns, block_scores * score_scale, mask=valid)
+ elif program >= MAIN_PROGRAMS + 2:
+ # One v head a program, as the K/V append stores it.
+ h = program - MAIN_PROGRAMS - 2
+ v_row = row + (HQ + HK + h) * D
+ if NVFP4:
+ v_global = tl.load(kv_scale_orig_quant + 2)
+ if USE_PDL:
+ tl.extra.cuda.gdc_wait()
+ if live:
+ if NVFP4:
+ v_data, v_sf = _quantize_row(v_row, 1, v_global, D)
+ pairs = tl.arange(0, D // 2)
+ blocks = tl.arange(0, D // 16)
+ tl.store(
+ v_cache + page * s_vc_page + h * s_vc_head + slot * s_vc_tok + pairs, v_data
+ )
+ # V scales in four-token groups within the head page.
+ tl.store(
+ v_scale_cache + page * s_vs_page + h * s_vs_head
+ + (slot // 4) * (4 * (D // 16)) + blocks * 4 + slot % 4,
+ v_sf,
+ ) # fmt: skip
+ else:
+ d = tl.arange(0, D)
+ tl.store(
+ v_cache + page * s_vc_page + h * s_vc_head + slot * s_vc_tok + d,
+ tl.load(v_row + d).to(v_cache.dtype.element_ty),
+ )
+ else:
+ # Before the wait: the coordinates; at a group's last token, the group's earlier
+ # raw keys (written by earlier steps), the first token's 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)
+ earlier = group_end & (offsets < CR - 1)
+ if IN_LAUNCH_GROUPS:
+ # A request's rows come consecutive at consecutive positions: the earlier
+ # member at row `token - (CR - 1) + o`, when that row is this request's at
+ # that position, is written by this launch -- its raw key is taken from its
+ # projection row after the wait, not from the index plane, and as the first
+ # member its coordinates from its row of `coords`.
+ member_rows = token - (CR - 1) + offsets
+ maybe = earlier & (member_rows >= 0)
+ member_requests = tl.load(request_indices + member_rows, mask=maybe, other=-1)
+ member_logical = tl.load(logical_positions + member_rows, mask=maybe, other=-1)
+ in_launch = (
+ maybe
+ & (member_requests.to(tl.int64) == request)
+ & (member_logical.to(tl.int64) == members)
+ )
+ else:
+ in_launch = offsets < 0
+ group = tl.load(
+ index_cache
+ + member_pages[:, None] * s_ic_page
+ + (tl.maximum(members, 0) % TPB)[:, None] * s_ic_slot
+ + dims[None, :],
+ mask=(earlier & ~in_launch)[:, None],
+ other=0.0,
+ )
+ first = tl.maximum(logical - (CR - 1), 0)
+ first_page = tl.load(
+ block_table + request * s_bt_req + (first // TPB) * s_bt_col, mask=group_end, other=0
+ ).to(tl.int64)
+ first_base = position_cache + first_page * s_pc_page + (first % TPB) * s_pc_slot
+ if IN_LAUNCH_GROUPS:
+ # A first member written by this launch has its coordinates in its row.
+ first_in_launch = tl.sum(tl.where(offsets == 0, in_launch, False).to(tl.int32)) > 0
+ cached = group_end & ~first_in_launch
+ first_row = coords + (token - (CR - 1)) * s_c_row
+ pos_t = tl.load(first_row, mask=first_in_launch, other=0)
+ pos_h = tl.load(first_row + s_c_axis, mask=first_in_launch, other=0)
+ pos_w = tl.load(first_row + 2 * s_c_axis, mask=first_in_launch, other=0)
+ pos_t = tl.where(cached, tl.load(first_base, mask=cached, other=0), pos_t)
+ pos_h = tl.where(cached, tl.load(first_base + s_pc_axis, mask=cached, other=0), pos_h)
+ pos_w = tl.where(
+ cached, tl.load(first_base + 2 * s_pc_axis, mask=cached, other=0), pos_w
+ )
+ else:
+ 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:
+ # The coordinates, as the index store writes them.
+ tl.store(
+ position_cache + page * s_pc_page + slot * s_pc_slot + axes * s_pc_axis,
+ coordinate.to(position_cache.dtype.element_ty),
+ mask=axes < 3,
+ )
+ # The raw index key -- or, at a group's last token, the group compressed in
+ # its place (the compression kernel overwrites the same slot).
+ current_k = tl.load(row + IQ_COL + IH * ID + dims)
+ stored_k = current_k.to(index_cache.dtype.element_ty)
+ if (logical + 1) % CR == 0:
+ if IN_LAUNCH_GROUPS:
+ launch_k = tl.load(
+ proj
+ + member_rows[:, None] * stride_proj
+ + IQ_COL
+ + IH * ID
+ + dims[None, :],
+ mask=in_launch[:, None],
+ other=0.0,
+ )
+ group = tl.where(
+ in_launch[:, None], launch_k.to(index_cache.dtype.element_ty), group
+ )
+ group = tl.where(offsets[:, None] == CR - 1, stored_k[None, :], group)
+ pooled = (tl.sum(group.to(tl.float32), axis=0) / CR).to(tl.bfloat16)
+ stored_k = tl.reshape(
+ _compress_norm_rope(
+ pooled[None, :], key_weight, cosine, sine, eps, ID, ROTARY_PAIRS
+ ),
+ (ID,),
+ )
+ tl.store(index_cache + page * s_ic_page + slot * s_ic_slot + dims, stored_k)
+
+
+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)`.
+
+ Over an NVFP4 pool (`k_scale_cache`, `v_scale_cache` and `kv_scale_orig_quant` given,
+ `k_cache` / `v_cache` its uint8 e2m1 planes) the K/V append is
+ `qsa_paged_kv_store_nvfp4`'s: e2m1 bytes and e4m3 block scales, K's token-major and
+ V's in four-token groups, at the second-level scales `kv_scale_orig_quant[1]` / `[2]`.
+ `rows_per_request` > 1 declares a verify step's layout: each request's rows
+ consecutive, at consecutive positions.
+
+ `index_scores` (fp32 `[T, table width * tpb / compress ratio]`, up to `SCORE_ROWS`
+ decode rows) takes `qsa_paged_index_scores`' output over the index plane and query
+ this launch leaves (`only_visible_blocks`: each row's visible columns written, the
+ rest untouched): more programs score `SCORE_BN` blocks each, rotating the row's index
+ query themselves and, at a group's last token, compressing the last visible block's
+ key themselves, as the programs writing them do.
+
+ Fusion boundary: as the four entries (with `index_scores`, and the index scores); the
+ projection before and the selection after are the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="decode, 1 row at position 8699 (a group's last token)", spec=dict(pos=(8699,))),
+ Cell(
+ why="decode, 4 rows: three group ends (two on a page's last slot), one with no compression",
+ spec=dict(pos=(8699, 8700, 4095, 31)),
+ ),
+ Cell(
+ why="decode, 2 rows, one unmapped (request out of range)",
+ spec=dict(pos=(8699, 1201), unmapped=1),
+ ),
+ Cell(
+ why="NVFP4 pool, 1 row at position 8699 (a group's last token), K/V the role "
+ "planes of the engine's [pages, 2, Hkv, tpb, D/2] pool",
+ spec=dict(pos=(8699,), nvfp4=True),
+ ),
+ Cell(
+ why="NVFP4 pool, 4 rows: three group ends (two on a page's last slot), one with "
+ "no compression, with zero, -0.0 and outlier blocks in k and v",
+ spec=dict(pos=(8699, 8700, 4095, 31), nvfp4=True, edge=True),
+ ),
+ Cell(
+ why="NVFP4 pool, 2 rows, one unmapped (request out of range)",
+ spec=dict(pos=(8699, 1201), unmapped=1, nvfp4=True),
+ ),
+ Cell(
+ why="NVFP4 pool, 40 rows (two heads a program): positions across groups and pages",
+ spec=dict(pos=tuple(2000 + 151 * r for r in range(40)), nvfp4=True),
+ ),
+ Cell(
+ why="NVFP4 pool, 72 rows (eight heads a program, k heads beside q heads)",
+ spec=dict(pos=tuple(1000 + 113 * r for r in range(72)), nvfp4=True),
+ ),
+ Cell(
+ why="NVFP4 pool, 1 row at a group's last token, the index scores in the launch",
+ spec=dict(pos=(8699,), nvfp4=True, scores=True),
+ ),
+ Cell(
+ why="NVFP4 pool, 2 rows (a group end on a page's last slot, one row without), the "
+ "index scores in the launch",
+ spec=dict(pos=(4095, 8700), nvfp4=True, scores=True, edge=True),
+ ),
+ Cell(
+ why="bf16 pool, 2 rows (a group end, and one at 8 tokens), the index scores in the "
+ "launch",
+ spec=dict(pos=(6143, 7), scores=True),
+ ),
+ Cell(
+ why="verify, NVFP4 pool, 3 requests x 4 rows: group ends at a request's third "
+ "row (two earlier members in the launch), fourth (three) and first (none)",
+ spec=dict(pos=(8697, 8696, 8699), rows_per_request=4, nvfp4=True, edge=True),
+ ),
+ Cell(
+ why="verify, bf16 pool, 2 requests x 4 rows, one group end crossing a page",
+ spec=dict(pos=(4093, 8698), rows_per_request=4),
+ ),
+ Cell(
+ why="verify, NVFP4 pool, 32 requests x 4 rows (eight heads a program)",
+ spec=dict(pos=tuple(3000 + 211 * r for r in range(32)), rows_per_request=4, nvfp4=True),
+ ),
+ )
+
+ 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]`; generation
+ rows: one per request (`rows_per_request` 1: a group's earlier tokens are already in
+ the index plane), or a verify step's `rows_per_request` consecutive rows a request at
+ consecutive positions (a group's earlier tokens in the index plane or in the
+ request's preceding rows). A row whose page lookup fails writes no cache, as the
+ stores and compression skip it.
+ Over an NVFP4 pool: the preconditions of `qsa_paged_kv_store_nvfp4` (uint8
+ `[pages, Hkv, tpb, D/2]` data and `[pages, Hkv, tpb, D/16]` scale planes, each scale
+ head page contiguous, a contiguous fp32 `[3]` `kv_scale_orig_quant`); the page
+ bound is the data plane's page count, as there.
+
+ Numerics: the same arithmetic as the four entries; outputs bitwise theirs (the
+ NVFP4 cells compare every byte of the four K/V planes).
+
+ Compiled variants: one per row-count band of the norm + rope launch plan (up to 16,
+ up to 64, more rows), pool kind (bf16 / NVFP4) and step kind (one row a request /
+ verify); the row count, table height and width and the page count are not
+ specialized on.
+ """
+
+ 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,
+ k_scale_cache: Optional[torch.Tensor] = None,
+ v_scale_cache: Optional[torch.Tensor] = None,
+ kv_scale_orig_quant: Optional[torch.Tensor] = None,
+ rows_per_request: int = 1,
+ index_scores: Optional[torch.Tensor] = 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)
+ nvfp4 = k_scale_cache is not None
+ if not nvfp4:
+ # Placeholders: the bf16 kernel does not read them.
+ k_scale_cache = v_scale_cache = kv_scale_orig_quant = k_cache
+ score_tiles = score_blocks = 0
+ if index_scores is not None:
+ assert num_tokens <= SCORE_ROWS and rows_per_request == 1 and warps == 2, (
+ "the index scores ride decode launches of up to SCORE_ROWS rows"
+ )
+ score_blocks = block_table.shape[1] * tokens_per_block // compress_ratio
+ score_tiles = triton.cdiv(score_blocks, SCORE_BN)
+ pdl = get_env_enable_pdl()
+ _attn_decode_prologue_kernel[(num_tokens, main_programs + 2 + num_heads_k + score_tiles)](
+ 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),
+ k_scale_cache,
+ v_scale_cache,
+ k_scale_cache.stride(0),
+ k_scale_cache.stride(1),
+ v_scale_cache.stride(0),
+ v_scale_cache.stride(1),
+ kv_scale_orig_quant,
+ proj if index_scores is None else index_scores,
+ 0 if index_scores is None else index_scores.stride(0),
+ index_head_dim**-0.5,
+ 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],
+ NVFP4=nvfp4,
+ IN_LAUNCH_GROUPS=rows_per_request > 1,
+ SCORE_TILES=score_tiles,
+ SCORE_BLOCKS=score_blocks,
+ SBN=SCORE_BN,
+ SBH=max(16, triton.next_power_of_2(index_heads)),
+ 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,
+ k_scale_cache: Optional[torch.Tensor] = None,
+ v_scale_cache: Optional[torch.Tensor] = None,
+ kv_scale_orig_quant: Optional[torch.Tensor] = None,
+ rows_per_request: int = 1,
+ index_scores: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """The four entries in their composition order, on the given buffers; with
+ `index_scores`, then `qsa_paged_index_scores`' visible columns written into it."""
+ 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]
+ if k_scale_cache is None:
+ _kv_store(
+ k,
+ v,
+ k_cache,
+ v_cache,
+ request_indices,
+ logical_positions,
+ block_table,
+ tokens_per_block,
+ )
+ else:
+ _kv_store_nvfp4(
+ k,
+ v,
+ k_cache,
+ v_cache,
+ k_scale_cache,
+ v_scale_cache,
+ request_indices,
+ logical_positions,
+ block_table,
+ tokens_per_block,
+ kv_scale_orig_quant,
+ )
+ 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
+ if index_scores is not None:
+ scores = _index_scores(
+ q=index_q, index_cache=index_cache, block_table=block_table,
+ query_positions=logical_positions, request_indices=request_indices,
+ tokens_per_block=tokens_per_block, compress_ratio=compress_ratio,
+ only_visible_blocks=True, context_rows=False,
+ ) # fmt: skip
+ for r in range(rows):
+ visible = (int(logical_positions[r]) + 1) // compress_ratio
+ index_scores[r, :visible] = scores[r, :visible]
+ 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)"
+ )
+ k_scale_cache = kwargs.get("k_scale_cache")
+ v_scale_cache = kwargs.get("v_scale_cache")
+ global_scales = kwargs.get("kv_scale_orig_quant")
+ nvfp4 = (k_scale_cache, v_scale_cache, global_scales)
+ assert all(t is None for t in nvfp4) or all(t is not None for t in nvfp4), (
+ "an NVFP4 pool takes both scale planes and kv_scale_orig_quant"
+ )
+ scores = kwargs.get("index_scores")
+ if scores is not None:
+ assert scores.dtype == torch.float32 and scores.stride(1) == 1, "fp32 index scores"
+ assert proj.shape[0] <= SCORE_ROWS and kwargs.get("rows_per_request", 1) == 1, (
+ "the index scores ride decode launches of up to SCORE_ROWS rows"
+ )
+ if k_scale_cache is not None:
+ for scales in (k_scale_cache, v_scale_cache):
+ assert scales.dtype == torch.uint8 and scales.stride(3) == 1, "uint8 scale planes"
+ assert scales.stride(2) == scales.shape[3], "each scale head page contiguous"
+ assert (
+ global_scales.dtype == torch.float32
+ and global_scales.shape == (3,)
+ and global_scales.is_contiguous()
+ ), "a contiguous fp32 [3] kv_scale_orig_quant"
+
+
+qwen4_exp_attn_decode_prologue = Qwen4ExpAttnDecodePrologue()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qwen4_exp_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qwen4_exp_qk_norm_rope.py
new file mode 100644
index 000000000000..27a17ed46ace
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qwen4_exp_qk_norm_rope.py
@@ -0,0 +1,1193 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Gemma QK norm + partial NeoX rope for both of an attention layer's planes, one launch.
+
+The main plane (q and k heads of the fused input projection) is rewritten in
+place; the QSA index plane's query heads are read from the same projection row
+and written, normalized and rotated, into a separate buffer. The rotation reads
+the target's own fp32 `[positions, rotary_dim]` cos|sin table.
+
+Two kernels, one result: a Triton kernel, and from `_CUTE_MIN_TOKENS` tokens (a mixed
+step's prefill chunk) a persistent CuTe DSL kernel that is bitwise the Triton kernel's
+eight-head launch and streams the projection at about twice its rate.
+"""
+
+from __future__ import annotations
+
+import struct
+from typing import Dict, Tuple
+
+import cutlass
+import cutlass.cute as cute
+import torch
+import triton
+import triton.language as tl
+from cutlass._mlir.dialects import llvm
+from cutlass.cute.runtime import from_dlpack
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
+
+
+@triton.jit
+def _norm_rope_tile(
+ src,
+ dst,
+ heads,
+ head_mask,
+ weight,
+ positions,
+ num_tokens,
+ token,
+ cos_sin,
+ stride_cs,
+ eps,
+ D: tl.constexpr,
+ ROT: tl.constexpr,
+ S1: tl.constexpr,
+ S2: tl.constexpr,
+ WEIGHT_SPLIT: tl.constexpr,
+ weight_k,
+ dst_head_stride,
+):
+ """One `[len(heads), D]` tile: rms over each head, `* (1 + w)`, rotate the
+ first ROT dims in NeoX halves. fp32 throughout, one rounding at the store."""
+ HALF: tl.constexpr = ROT // 2
+ d = tl.arange(0, D)
+ rot = d < ROT
+ partner = tl.where(d < HALF, d + HALF, tl.where(rot, d - HALF, d))
+ hm = head_mask[:, None]
+ x = tl.load(src + heads[:, None] * D + d[None, :], mask=hm, other=0.0).to(tl.float32)
+ xp = tl.load(src + heads[:, None] * D + partner[None, :], mask=hm & rot[None, :], other=0.0).to(
+ tl.float32
+ )
+ w = tl.load(weight + d).to(tl.float32)
+ wp = tl.load(weight + partner).to(tl.float32)
+ if WEIGHT_SPLIT > 0:
+ # Heads at or past WEIGHT_SPLIT take the second weight (the k heads).
+ is_k = (heads >= WEIGHT_SPLIT)[:, None]
+ wk = tl.load(weight_k + d).to(tl.float32)
+ wkp = tl.load(weight_k + partner).to(tl.float32)
+ w2 = tl.where(is_k, wk[None, :], w[None, :]) + 1.0
+ wp2 = tl.where(is_k, wkp[None, :], wp[None, :]) + 1.0
+ else:
+ w2 = w[None, :] + 1.0
+ wp2 = wp[None, :] + 1.0
+ rstd = tl.rsqrt(tl.sum(x * x, axis=1) / D + eps)
+ y = x * rstd[:, None] * w2
+ yp = xp * rstd[:, None] * wp2
+ # Interleaved mRoPE: pair j reads coordinate 1 when j % 3 == 1 and
+ # j < 3 * S1, coordinate 2 when j % 3 == 2 and j < 3 * S2, else 0.
+ j = d % HALF
+ p0 = tl.load(positions + token)
+ p1 = tl.load(positions + num_tokens + token)
+ p2 = tl.load(positions + 2 * num_tokens + token)
+ pos = tl.where((j % 3 == 1) & (j < 3 * S1), p1, tl.where((j % 3 == 2) & (j < 3 * S2), p2, p0))
+ row = cos_sin + pos.to(tl.int64) * stride_cs
+ c = tl.load(row + j, mask=rot, other=1.0)
+ s = tl.load(row + HALF + j, mask=rot, other=0.0)
+ s = tl.where(d < HALF, -s, s)
+ out = y * c[None, :] + yp * s[None, :]
+ tl.store(
+ dst + heads[:, None] * dst_head_stride + d[None, :],
+ out.to(dst.dtype.element_ty),
+ mask=hm,
+ )
+
+
+@triton.jit(do_not_specialize=["num_tokens"])
+def _qwen4_exp_qk_norm_rope_kernel(
+ proj,
+ stride_proj,
+ index_q,
+ stride_iq_token,
+ stride_iq_head,
+ positions,
+ num_tokens,
+ cos_sin,
+ stride_cs,
+ q_weight,
+ k_weight,
+ index_q_weight,
+ eps,
+ HQ: tl.constexpr,
+ HK: tl.constexpr,
+ D: tl.constexpr,
+ IH: tl.constexpr,
+ ID: tl.constexpr,
+ IQ_COL: tl.constexpr,
+ ROT: tl.constexpr,
+ S1: tl.constexpr,
+ S2: tl.constexpr,
+ IS1: tl.constexpr,
+ IS2: tl.constexpr,
+ HPB: tl.constexpr,
+ MAIN_PROGRAMS: tl.constexpr,
+):
+ token = tl.program_id(0).to(tl.int64)
+ program = tl.program_id(1)
+ row = proj + token * stride_proj
+ if program < MAIN_PROGRAMS:
+ main_heads = program * HPB + tl.arange(0, HPB)
+ _norm_rope_tile(
+ row, row, main_heads, main_heads < HQ + HK, q_weight, positions, num_tokens, token,
+ cos_sin, stride_cs, eps, D, ROT, S1, S2, HQ, k_weight, D,
+ ) # fmt: skip
+ else:
+ index_heads = tl.arange(0, IH)
+ _norm_rope_tile(
+ row + IQ_COL, index_q + token * stride_iq_token, index_heads, index_heads < IH,
+ index_q_weight, positions, num_tokens, token, cos_sin, stride_cs, eps,
+ ID, ROT, IS1, IS2, 0, index_q_weight, stride_iq_head,
+ ) # fmt: skip
+
+
+def _launch(num_tokens: int) -> Tuple[int, int]:
+ """Main-plane heads per program and warps per program, by token count.
+
+ Measured on GB300 (CUDA-graph replay, 12 rotating layers, us per call;
+ HPB/warps):
+
+ tokens 1/2 2/2 2/4 4/4 8/8 32/8
+ 1 1.76 1.84 1.87 1.99 1.92 2.82
+ 8 1.83 1.97 1.97 2.02 2.05 2.91
+ 32 2.11 2.14 2.18 2.14 2.16 2.99
+ 64 2.48 2.39 2.47 2.44 2.48 3.03
+ 128 3.24 3.19 3.46 3.12 2.96 3.12
+ 2048 29.7 25.4 30.8 24.0 22.6 25.8
+ 8192 115 98.5 121 95.7 89.5 122
+ """
+ if num_tokens <= 16:
+ return 1, 2
+ if num_tokens <= 64:
+ return 2, 2
+ return 8, 8
+
+
+#: This checkpoint's geometry: 24 q + 2 k heads of 256 at the head of the
+#: 14080-wide fused projection row, 4 index q heads of 128 at column 13312, a
+#: 64-dim rotary prefix, `mrope_section` [11, 11, 10] on the main plane and
+#: coordinate 0 alone on the index plane.
+CKPT = dict(
+ num_heads_q=24,
+ num_heads_k=2,
+ head_dim=256,
+ index_heads=4,
+ index_head_dim=128,
+ index_q_col=13312,
+ proj_width=14080,
+ index_scratch_heads=6,
+ rotary_dim=64,
+ mrope_section1=11,
+ mrope_section2=10,
+ index_mrope_section1=0,
+ index_mrope_section2=0,
+)
+
+
+# ---- CuTe DSL path: prefill-size token counts -------------------------------------------
+
+#: Token count from which the op runs the CuTe kernel. GB300, graph replay, us per call (Triton
+#: -> CuTe): 512 8.4 -> 10.3, 768 12.3 -> 8.2, 1024 14.5 -> 8.3, 4096 45.2 -> 18.5, 16384
+#: 182 -> 82.7. The Triton kernel runs its eight-head launch from 65 tokens on, and the CuTe
+#: kernel is bitwise that launch.
+_CUTE_MIN_TOKENS = 1024
+#: A token's stage in a warp's ring: the lanes' four chunks (32 x 16 bytes each), the cos and
+#: sin of the 32 rotary pairs, and three positions.
+_STAGE = 4 * 32 * 16 + 2 * 32 * 4 + 16
+_STAGE_CS = 4 * 32 * 16
+_STAGE_POS = _STAGE_CS + 2 * 32 * 4
+#: Stages in a warp's ring (the next token's loads in flight while one is formed).
+_RING = 2
+_CTAS_PER_SM = 3
+
+
+def _iv(x, loc, ip):
+ return x.ir_value(loc=loc, ip=ip)
+
+
+def _asm(ret, args, text, constraints, loc, ip, side_effects=False):
+ return llvm.inline_asm(
+ ret,
+ [_iv(a, loc, ip) for a in args],
+ text,
+ constraints,
+ has_side_effects=side_effects,
+ is_align_stack=False,
+ asm_dialect=llvm.AsmDialect.AD_ATT,
+ loc=loc,
+ ip=ip,
+ )
+
+
+def _asm4(args, text, constraints, loc, ip, side_effects=False):
+ ty = T.i32()
+ res = _asm(
+ llvm.StructType.get_literal([ty] * 4), args, text, constraints, loc, ip, side_effects
+ )
+ return [cutlass.Int32(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(4)]
+
+
+@dsl_user_op
+def _ldg_v4(addr, *, loc=None, ip=None):
+ return _asm4(
+ [cutlass.Int64(addr)],
+ "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];",
+ "=r,=r,=r,=r,l",
+ loc,
+ ip,
+ side_effects=True,
+ )
+
+
+@dsl_user_op
+def _ldg_b32(addr, *, loc=None, ip=None):
+ return cutlass.Int32(
+ _asm(T.i32(), [cutlass.Int64(addr)], "ld.global.b32 $0, [$1];", "=r,l", loc, ip, True)
+ )
+
+
+@dsl_user_op
+def _stg_v4(addr, off: int, a, b, c, d, *, loc=None, ip=None):
+ """16 bytes at `addr + off`."""
+ _asm(
+ T.i32(),
+ [cutlass.Int64(addr)] + [cutlass.Int32(v) for v in (a, b, c, d)],
+ "{ st.global.v4.b32 [$1+%d], {$2, $3, $4, $5}; mov.u32 $0, 0; }" % off,
+ "=r,l,r,r,r,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+
+
+@dsl_user_op
+def _lds_v4(addr, *, loc=None, ip=None):
+ return _asm4(
+ [cutlass.Int32(addr)],
+ "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];",
+ "=r,=r,=r,=r,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+
+
+@dsl_user_op
+def _lds_b32(addr, *, loc=None, ip=None):
+ return cutlass.Int32(
+ _asm(T.i32(), [cutlass.Int32(addr)], "ld.shared.b32 $0, [$1];", "=r,r", loc, ip, True)
+ )
+
+
+@dsl_user_op
+def _cp16(dst, src, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [cutlass.Int32(dst), cutlass.Int64(src)],
+ "{ cp.async.cg.shared.global [$1], [$2], 16; mov.u32 $0, 0; }",
+ "=r,r,l",
+ loc,
+ ip,
+ side_effects=True,
+ )
+
+
+@dsl_user_op
+def _cp4(dst, src, pred, *, loc=None, ip=None):
+ """4 bytes where `pred` is not 0."""
+ _asm(
+ T.i32(),
+ [cutlass.Int32(dst), cutlass.Int64(src), cutlass.Int32(pred)],
+ "{ .reg .pred p; setp.ne.b32 p, $3, 0; @p cp.async.ca.shared.global [$1], [$2], 4; "
+ "mov.u32 $0, 0; }",
+ "=r,r,l,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+
+
+@dsl_user_op
+def _cp_commit(*, loc=None, ip=None):
+ _asm(T.i32(), [], "{ cp.async.commit_group; mov.u32 $0, 0; }", "=r", loc, ip, True)
+
+
+@dsl_user_op
+def _cp_wait(n: int, *, loc=None, ip=None):
+ _asm(T.i32(), [], "{ cp.async.wait_group %d; mov.u32 $0, 0; }" % n, "=r", loc, ip, True)
+
+
+@dsl_user_op
+def _bf16_lo(w, *, loc=None, ip=None):
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [cutlass.Int32(w)],
+ "{ .reg .b32 t; shl.b32 t, $1, 16; mov.b32 $0, t; }",
+ "=f,r",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _bf16_hi(w, *, loc=None, ip=None):
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [cutlass.Int32(w)],
+ "{ .reg .b32 t; and.b32 t, $1, 0xFFFF0000; mov.b32 $0, t; }",
+ "=f,r",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _unpack2(w, *, loc=None, ip=None):
+ """A bf16x2 word -> its two values as an fp32x2 pair, the low element first."""
+ return cutlass.Int64(
+ _asm(
+ T.i64(),
+ [cutlass.Int32(w)],
+ "{ .reg .b32 a, b; shl.b32 a, $1, 16; and.b32 b, $1, 0xFFFF0000; mov.b64 $0, {a, b}; }",
+ "=l,r",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _pair(lo, hi, *, loc=None, ip=None):
+ """Two fp32 bit patterns -> one fp32x2 pair."""
+ return cutlass.Int64(
+ _asm(
+ T.i64(),
+ [cutlass.Int32(lo), cutlass.Int32(hi)],
+ "mov.b64 $0, {$1, $2};",
+ "=l,r,r",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _splat2(x, *, loc=None, ip=None):
+ return cutlass.Int64(
+ _asm(T.i64(), [cutlass.Float32(x)], "mov.b64 $0, {$1, $1};", "=l,f", loc, ip)
+ )
+
+
+@dsl_user_op
+def _mul2(a, b, *, loc=None, ip=None):
+ return cutlass.Int64(
+ _asm(
+ T.i64(),
+ [cutlass.Int64(a), cutlass.Int64(b)],
+ "mul.rn.f32x2 $0, $1, $2;",
+ "=l,l,l",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _fma2(a, b, c, *, loc=None, ip=None):
+ return cutlass.Int64(
+ _asm(
+ T.i64(),
+ [cutlass.Int64(a), cutlass.Int64(b), cutlass.Int64(c)],
+ "fma.rn.f32x2 $0, $1, $2, $3;",
+ "=l,l,l,l",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _add2(a, b, *, loc=None, ip=None):
+ return cutlass.Int64(
+ _asm(
+ T.i64(),
+ [cutlass.Int64(a), cutlass.Int64(b)],
+ "add.rn.f32x2 $0, $1, $2;",
+ "=l,l,l",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _hsum2(a, *, loc=None, ip=None):
+ """The low plus the high element of a pair."""
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [cutlass.Int64(a)],
+ "{ .reg .f32 x, y; mov.b64 {x, y}, $1; add.rn.f32 $0, x, y; }",
+ "=f,l",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _bf16x2_of(a, *, loc=None, ip=None):
+ """A pair rounded (RN) into one bf16x2 word, the low element low."""
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [cutlass.Int64(a)],
+ "{ .reg .f32 x, y; mov.b64 {x, y}, $1; cvt.rn.bf16x2.f32 $0, y, x; }",
+ "=r,l",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _neg_bits(w, *, loc=None, ip=None):
+ """An fp32 bit pattern with its sign flipped (the Triton kernel's `xor`)."""
+ return cutlass.Int32(
+ _asm(T.i32(), [cutlass.Int32(w)], "xor.b32 $0, $1, 0x80000000;", "=r,r", loc, ip)
+ )
+
+
+@dsl_user_op
+def _fma(a, b, c, *, loc=None, ip=None):
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)],
+ "fma.rn.f32 $0, $1, $2, $3;",
+ "=f,f,f,f",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _mul(a, b, *, loc=None, ip=None):
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [cutlass.Float32(a), cutlass.Float32(b)],
+ "mul.rn.f32 $0, $1, $2;",
+ "=f,f,f",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _add(a, b, *, loc=None, ip=None):
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [cutlass.Float32(a), cutlass.Float32(b)],
+ "add.rn.f32 $0, $1, $2;",
+ "=f,f,f",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _rstd(ss, eps, width_hex: str, *, loc=None, ip=None):
+ """`rsqrt(ss / width + eps)` as the Triton kernel's `div.full`, `add`, `rsqrt.approx.ftz`:
+ the `add` has no rounding mode there, so ptxas fuses it with the division (a multiply by
+ the reciprocal) into one FFMA -- here the same."""
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [cutlass.Float32(ss), cutlass.Float32(eps)],
+ "{ .reg .f32 t; div.full.f32 t, $1, " + width_hex + "; add.f32 t, t, $2; "
+ "rsqrt.approx.ftz.f32 $0, t; }",
+ "=f,f,f",
+ loc,
+ ip,
+ )
+ )
+
+
+def _hex(x: float) -> str:
+ return "0f%08X" % struct.unpack(" bool:
+ """The geometry and layouts the CuTe kernel covers; anything else runs the Triton kernel."""
+ bf = torch.bfloat16
+ return (
+ num_heads_q % 4 == 0
+ and num_heads_k == 2
+ and head_dim == 256
+ and index_heads == 4
+ and index_head_dim == 128
+ and rotary_dim == 64
+ and index_q_col % 8 == 0
+ and proj.dtype == index_q.dtype == bf
+ and all(w.dtype == bf and w.is_contiguous() and w.data_ptr() % 16 == 0 for w in weights)
+ and cos_sin.dtype == torch.float32
+ and cos_sin.stride(1) == 1
+ and proj.data_ptr() % 16 == 0
+ and proj.stride(0) % 8 == 0
+ and index_q.data_ptr() % 16 == 0
+ and index_q.stride(0) % 8 == 0
+ and index_q.stride(1) == index_head_dim
+ )
+
+
+def _launch_cute(
+ proj: torch.Tensor,
+ index_q: torch.Tensor,
+ position_ids: torch.Tensor,
+ cos_sin: torch.Tensor,
+ q_weight: torch.Tensor,
+ k_weight: torch.Tensor,
+ index_q_weight: torch.Tensor,
+ num_heads_q: int,
+ index_q_col: int,
+ eps: float,
+ sections: Tuple[int, int, int, int],
+) -> None:
+ key = (num_heads_q, index_q_col) + tuple(sections)
+ fn = _RING_KERNELS.get(key)
+ if fn is None:
+ gm = cute.AddressSpace.gmem
+ bf = cutlass.BFloat16
+ fn = cute.compile(
+ _NormRopeRing(num_heads_q, index_q_col, *sections),
+ from_dlpack(proj, assumed_align=16).mark_layout_dynamic(leading_dim=1),
+ make_ptr(bf, index_q.data_ptr(), gm, assumed_align=16),
+ make_ptr(cutlass.Int32, position_ids.data_ptr(), gm, assumed_align=4),
+ make_ptr(cutlass.Float32, cos_sin.data_ptr(), gm, assumed_align=4),
+ make_ptr(bf, q_weight.data_ptr(), gm, assumed_align=16),
+ make_ptr(bf, k_weight.data_ptr(), gm, assumed_align=16),
+ make_ptr(bf, index_q_weight.data_ptr(), gm, assumed_align=16),
+ *(cutlass.Int32(1),) * 4,
+ cutlass.Float32(1e-6),
+ cutlass.Int32(1),
+ cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+ options="--opt-level 2 --enable-tvm-ffi",
+ )
+ _RING_KERNELS[key] = fn
+ dev = proj.device.index
+ if dev not in _SM_COUNT:
+ _SM_COUNT[dev] = torch.cuda.get_device_properties(proj.device).multi_processor_count
+ tokens = proj.shape[0]
+ fn(
+ proj,
+ index_q.data_ptr(),
+ position_ids.data_ptr(),
+ cos_sin.data_ptr(),
+ q_weight.data_ptr(),
+ k_weight.data_ptr(),
+ index_q_weight.data_ptr(),
+ tokens,
+ proj.stride(0),
+ index_q.stride(0),
+ cos_sin.stride(0),
+ eps,
+ min(tokens, _SM_COUNT[dev] * _CTAS_PER_SM),
+ )
+
+
+def _launch_triton(
+ proj: torch.Tensor,
+ index_q: torch.Tensor,
+ position_ids: torch.Tensor,
+ cos_sin: torch.Tensor,
+ q_weight: torch.Tensor,
+ k_weight: torch.Tensor,
+ index_q_weight: torch.Tensor,
+ num_heads_q: int,
+ num_heads_k: int,
+ head_dim: int,
+ index_heads: int,
+ index_head_dim: int,
+ index_q_col: int,
+ rotary_dim: int,
+ eps: float,
+ sections: Tuple[int, int, int, int],
+) -> None:
+ num_tokens = proj.shape[0]
+ hpb, warps = _launch(num_tokens)
+ main_programs = triton.cdiv(num_heads_q + num_heads_k, hpb)
+ _qwen4_exp_qk_norm_rope_kernel[(num_tokens, main_programs + 1)](
+ proj,
+ proj.stride(0),
+ index_q,
+ index_q.stride(0),
+ index_q.stride(1),
+ position_ids,
+ num_tokens,
+ cos_sin,
+ cos_sin.stride(0),
+ q_weight,
+ k_weight,
+ index_q_weight,
+ eps,
+ HQ=num_heads_q,
+ HK=num_heads_k,
+ D=head_dim,
+ IH=index_heads,
+ ID=index_head_dim,
+ IQ_COL=index_q_col,
+ ROT=rotary_dim,
+ S1=sections[0],
+ S2=sections[1],
+ IS1=sections[2],
+ IS2=sections[3],
+ HPB=hpb,
+ MAIN_PROGRAMS=main_programs,
+ num_warps=warps,
+ )
+
+
+class Qwen4ExpQkNormRope(OpWrapper):
+ """Both planes of one attention layer's QK norm + rope, in one launch.
+
+ Main plane, in place on `proj`: heads `[0, Hq)` (weight `q_weight`) and
+ `[Hq, Hq + Hk)` (weight `k_weight`), each `head_dim` wide at the start of
+ every row, become `rope(rmsnorm(x) * (1 + w))`. Index plane: the
+ `index_heads` heads of `index_head_dim` starting at column `index_q_col` of
+ `proj` are read, normalized with `index_q_weight` and rotated, and written
+ to `index_q` (`[T, index_heads, index_head_dim]`, any token/head stride,
+ unit dim stride); `proj`'s index columns are not written. Every other
+ column of `proj` is left alone. Returns `(proj, index_q)`.
+
+ The rotation is NeoX over the first `rotary_dim` dims of a head: pair `j`
+ is `(j, j + rotary_dim / 2)`, rotated by `cos_sin[p, j]` and
+ `cos_sin[p, rotary_dim / 2 + j]` (`cos_sin` is the fp32 `[P, rotary_dim]`
+ table, cos block then sin block). `position_ids` is the contiguous int32
+ `[3, T]`; pair `j` reads coordinate 1 when `j % 3 == 1 and j < 3 * s1`,
+ coordinate 2 when `j % 3 == 2 and j < 3 * s2`, coordinate 0 otherwise --
+ interleaved mRoPE with the plane's own `(s1, s2)`; `(0, 0)` is the plain
+ rope on coordinate 0. fp32 throughout, one rounding at the store.
+
+ Fusion boundary: norm and rotation of both planes. The projection GEMM
+ before it, and the KV/index-cache writes and attention after it, are the
+ caller's. The index plane's raw key is not touched, so the index-cache
+ store may run before or after this call.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # Four cells, one per compiled launch the forward reaches: the Triton kernel's 1-, 2-
+ # and 8-head launches (by decode batch size) and the CuTe kernel (prefill).
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="decode M=1: one token on the 1-head/2-warp launch, three distinct mRoPE rows",
+ spec=dict(tokens=1, **CKPT),
+ ),
+ Cell(
+ why="decode M=17: first on the 2-head launch (13 main programs)",
+ spec=dict(tokens=17, **CKPT),
+ ),
+ Cell(
+ why="decode M=128: max concurrency, 8-head launch (last main program half masked), "
+ "positions up to the table end (262144 rows)",
+ spec=dict(tokens=128, max_position=262144, **CKPT),
+ ),
+ Cell(why="prefill 16384 tokens: max_num_tokens", spec=dict(tokens=16384, **CKPT)),
+ )
+
+ note = """
+ Metadata consumed: none (positions are an operand). Below `_CUTE_MIN_TOKENS`
+ tokens one Triton kernel, no autotuner: grid `(T, cdiv(Hq + Hk, HPB) + 1)`;
+ the last program of a token row runs the whole index plane as one
+ `[index_heads, index_head_dim]` tile. HPB and warps by token count from
+ `_launch` (its docstring has the table). Replaces, per attention layer, the
+ CUDA `fused_qk_norm_rope` main call, the index-scratch copy and the second
+ `fused_qk_norm_rope` call: 6.67 -> 1.76 us at M=1, 9.15 -> 2.96 at M=128, 185
+ -> 89.5 at 8192 tokens (graph replay).
+
+ From `_CUTE_MIN_TOKENS` tokens, where the geometry is this checkpoint's (`_cute_ok`:
+ head dims 256 / 128, 2 k heads, 4 index heads, a 64-dim rotary prefix, 16-byte
+ aligned rows, contiguous index heads), the CuTe DSL `_NormRopeRing`: three
+ persistent CTAs an SM, each forming every G-th token, a warp's next token
+ prefetched by cp.async (its class docstring has the mapping). 16384 tokens: 182 -> 82.7 us in the
+ microbenchmark (the Triton kernel was L1-bound; the ring runs at ~65% of DRAM peak,
+ bounded by load latency at three CTAs an SM).
+
+ Numerics: fp32 throughout, one rounding at the store; the CuTe kernel is bitwise
+ the Triton kernel's eight-head launch (measured at 1024-16384 tokens, signed zeros
+ included; not pinned by a test). Against the fp64
+ reference over the same table every cell is within one output rounding
+ step (measured max 1.48 ulp of the row scale, relative RMS <= 0.011 ulp).
+ Against the CUDA pair it replaces (measured, not pinned by a test), <= 1.41
+ ulp of the row scale and 0.09 ulp RMS (that op derives cos/sin with `__sincosf`, whose phase error grows
+ with position). Three agreeing position rows select the same table row for
+ every pair, so a text batch is exactly the plain rope on row 0.
+
+ Preconditions in `is_valid`: the main q/k columns end at or before
+ `index_q_col` (the planes do not overlap, so reading the index plane while
+ the main plane is rewritten in place is race-free); `index_q` does not
+ alias `proj`; positions below the table length (an out-of-range position
+ reads past the table).
+
+ Hard limits: token offsets in int64; positions index the table in int64.
+
+ CUDA graphs: the Triton kernel captured at M=4 replays bitwise an eager call on
+ inputs rewritten in place.
+
+ Not certified: non-zero index-plane mRoPE sections (the target passes (0, 0)),
+ capture of the CuTe kernel, `T == 0`, fp16/fp32 activations, `rotary_dim` > either head width,
+ non-power-of-two head widths or head counts (tl.arange), YaRN, non-Gemma
+ norm, a k plane on the index side (its scratch slot is not written).
+ """
+
+ def __call__(
+ self,
+ proj: torch.Tensor,
+ index_q: torch.Tensor,
+ position_ids: torch.Tensor,
+ cos_sin: torch.Tensor,
+ q_weight: torch.Tensor,
+ k_weight: torch.Tensor,
+ index_q_weight: torch.Tensor,
+ num_heads_q: int,
+ num_heads_k: int,
+ head_dim: int,
+ index_heads: int,
+ index_head_dim: int,
+ index_q_col: int,
+ rotary_dim: int,
+ eps: float,
+ mrope_section1: int,
+ mrope_section2: int,
+ index_mrope_section1: int = 0,
+ index_mrope_section2: int = 0,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ num_tokens = proj.shape[0]
+ if num_tokens == 0:
+ return proj, index_q
+ sections = (mrope_section1, mrope_section2, index_mrope_section1, index_mrope_section2)
+ weights = (q_weight, k_weight, index_q_weight)
+ if num_tokens >= _CUTE_MIN_TOKENS and _cute_ok(
+ proj, index_q, cos_sin, weights, num_heads_q, num_heads_k, head_dim, index_heads,
+ index_head_dim, index_q_col, rotary_dim,
+ ): # fmt: skip
+ _launch_cute(
+ proj, index_q, position_ids, cos_sin, *weights, num_heads_q, index_q_col, eps,
+ sections,
+ ) # fmt: skip
+ else:
+ _launch_triton(
+ proj, index_q, position_ids, cos_sin, *weights, num_heads_q, num_heads_k,
+ head_dim, index_heads, index_head_dim, index_q_col, rotary_dim, eps, sections,
+ ) # fmt: skip
+ return proj, index_q
+
+ def reference(
+ self,
+ proj: torch.Tensor,
+ index_q: torch.Tensor,
+ position_ids: torch.Tensor,
+ cos_sin: torch.Tensor,
+ q_weight: torch.Tensor,
+ k_weight: torch.Tensor,
+ index_q_weight: torch.Tensor,
+ num_heads_q: int,
+ num_heads_k: int,
+ head_dim: int,
+ index_heads: int,
+ index_head_dim: int,
+ index_q_col: int,
+ rotary_dim: int,
+ eps: float,
+ mrope_section1: int,
+ mrope_section2: int,
+ index_mrope_section1: int = 0,
+ index_mrope_section2: int = 0,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """`(proj after, index_q's written values)` in fp64 over the same table,
+ rounded once; neither input is mutated."""
+ num_tokens = proj.shape[0]
+ half = rotary_dim // 2
+ table = cos_sin.double()
+ pos = position_ids.long()
+ j = torch.arange(half, device=proj.device)
+
+ def angles(s1: int, s2: int):
+ coord = torch.zeros(half, dtype=torch.long, device=proj.device)
+ coord[(j % 3 == 1) & (j < 3 * s1)] = 1
+ coord[(j % 3 == 2) & (j < 3 * s2)] = 2
+ p = pos[coord, :].t() # [T, half]
+ return table[p, j][:, None, :], table[p, half + j][:, None, :]
+
+ def plane(x: torch.Tensor, w: torch.Tensor, cos, sin) -> torch.Tensor:
+ x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * (1.0 + w)
+ x1, x2 = x[..., :half], x[..., half:rotary_dim]
+ rotated = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], -1)
+ return torch.cat([rotated, x[..., rotary_dim:]], -1)
+
+ out = proj.double().clone()
+ main_heads = num_heads_q + num_heads_k
+ main_cols = main_heads * head_dim
+ w = torch.cat(
+ [
+ q_weight.double().expand(num_heads_q, head_dim),
+ k_weight.double().expand(num_heads_k, head_dim),
+ ]
+ )
+ main = out[:, :main_cols].reshape(num_tokens, main_heads, head_dim)
+ out[:, :main_cols] = plane(main, w, *angles(mrope_section1, mrope_section2)).reshape(
+ num_tokens, main_cols
+ )
+ iq = proj[:, index_q_col : index_q_col + index_heads * index_head_dim].double()
+ iq = plane(
+ iq.reshape(num_tokens, index_heads, index_head_dim),
+ index_q_weight.double(),
+ *angles(index_mrope_section1, index_mrope_section2),
+ )
+ return out.to(proj.dtype), iq.to(index_q.dtype)
+
+ def is_valid(
+ self,
+ proj: torch.Tensor,
+ index_q: torch.Tensor,
+ position_ids: torch.Tensor,
+ cos_sin: torch.Tensor,
+ q_weight: torch.Tensor,
+ k_weight: torch.Tensor,
+ index_q_weight: torch.Tensor,
+ num_heads_q: int,
+ num_heads_k: int,
+ head_dim: int,
+ index_heads: int,
+ index_head_dim: int,
+ index_q_col: int,
+ rotary_dim: int,
+ *unused_args: object,
+ **unused_kwargs: object,
+ ) -> None:
+ num_tokens = proj.shape[0]
+ assert proj.stride(-1) == 1, "proj must have a unit column stride"
+ assert (num_heads_q + num_heads_k) * head_dim <= index_q_col, (
+ "the index q columns overlap the main q/k heads: the in-place main write races "
+ "the index read"
+ )
+ assert index_q_col + index_heads * index_head_dim <= proj.shape[1], (
+ "the index q columns run past the projection row"
+ )
+ assert index_q.shape == (num_tokens, index_heads, index_head_dim), index_q.shape
+ assert index_q.stride(-1) == 1, "index_q must have a unit dim stride"
+ lo, hi = index_q.data_ptr(), proj.data_ptr()
+ assert lo >= hi + proj.numel() * proj.element_size() or hi >= lo + (
+ index_q.shape[0] * index_q.stride(0) * index_q.element_size()
+ ), "index_q aliases proj"
+ assert position_ids.shape == (3, num_tokens) and position_ids.is_contiguous(), (
+ f"position_ids must be a contiguous [3, {num_tokens}]; got {tuple(position_ids.shape)}"
+ )
+ assert position_ids.dtype == torch.int32, position_ids.dtype
+ assert cos_sin.dtype == torch.float32 and cos_sin.shape[1] == rotary_dim, cos_sin.shape
+ assert rotary_dim <= min(head_dim, index_head_dim) and rotary_dim % 2 == 0, rotary_dim
+ if num_tokens:
+ assert int(position_ids.min()) >= 0, "negative position"
+ assert int(position_ids.max()) < cos_sin.shape[0], "position past the rope table"
+
+
+qwen4_exp_qk_norm_rope = Qwen4ExpQkNormRope()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/spec_rollback_commit.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/spec_rollback_commit.py
new file mode 100644
index 000000000000..e90cfc7c171a
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/spec_rollback_commit.py
@@ -0,0 +1,487 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""`SpeculativeRollback`'s staging and commit of a verify step in one launch each.
+
+Tier 5 (Triton, called by the qwen4_exp target, `qwen3_8_flash_next_nvfp4__sm_103__default`,
+at every layout). Each `__call__` is one launch of a Triton kernel of this file,
+JIT-compiled on first use for the device it runs on. A verify step runs `p` rows per generation
+request through the trunk; the PLE pools and the QSA planes are written for every row. The staging
+launch copies what the commit rewrites from: each request's pre-step n-gram id history
+followed by its rows' ids (the `stream` whose length-`context` windows are the histories
+after each row), and its pre-step PLE conv state. After the draft loop the commit launch
+rewrites, for each generation request that accepted `a` rows:
+
+* its n-gram pool row to window `a` of its stream (the history after row `a - 1`);
+* its PLE conv pool row to columns `[a, a + state_len)` of `[pre-step state | the conv input
+ of its p rows]`;
+* the QSA position plane and every attention layer's index-K plane at each of its rows of
+ ordinal `>= a` to the values gathered before the trunk wrote them.
+
+The commit runs one program per (request, block): block `b` writes the request's conv
+channels `[b * 256, (b + 1) * 256)` and restores its rejected rows in index-K plane `b`; block 0
+also writes its n-gram row and position-plane slots.
+"""
+
+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
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_rollback_planes import (
+ ROW,
+)
+
+#: Conv channels a program covers.
+_BLOCK_C = 256
+
+#: A tep4 rank (the MTP member) as the stock `Qwen4ExpHybridCacheManager` lays its pools out at
+#: core's runtime: the PLE conv state `[slots, hc_count * hidden, (ple_kernel - 1) * ngram_size]`
+#: = `[1874, 10240, 9]` bf16 and the n-gram history `[1874, 2]` int64 (replicated, 1874 the final
+#: engine's state slots), and every attention layer's index-K plane a `[pages, 32, 1, 128]` bf16
+#: view of one coalesced pool at page stride 49152 elements, 257109 attention pages, beside the
+#: `[pages, 32, 3]` int32 position plane.
+_TEP4_PLE = dict(state_slots=1874, context=2, channels=10240, state_len=9)
+_TEP4 = dict(planes=12, layout="coalesced", pool_pages=257109, **_TEP4_PLE)
+
+
+@triton.jit
+def _stage_kernel(
+ ngram_pool,
+ ngram_stride,
+ conv_pool,
+ conv_s0,
+ conv_s1,
+ conv_s2,
+ slots,
+ ids,
+ stream,
+ conv_prev,
+ channels,
+ P: tl.constexpr,
+ CTX: tl.constexpr,
+ W_BLOCK: tl.constexpr,
+ STATE: tl.constexpr,
+ STATE_BLOCK: tl.constexpr,
+ BLOCK_C: tl.constexpr,
+):
+ i = tl.program_id(0)
+ cb = tl.program_id(1)
+ slot = tl.load(slots + i).to(tl.int64)
+ c = cb * BLOCK_C + tl.arange(0, BLOCK_C)
+ s = tl.arange(0, STATE_BLOCK)
+ m = (c < channels)[:, None] & (s < STATE)[None, :]
+ vals = tl.load(conv_pool + slot * conv_s0 + c[:, None] * conv_s1 + s[None, :] * conv_s2, mask=m)
+ tl.store(
+ conv_prev + (i.to(tl.int64) * channels + c[:, None]) * STATE + s[None, :], vals, mask=m
+ )
+ if cb == 0:
+ w = tl.arange(0, W_BLOCK)
+ hist = tl.load(ngram_pool + slot * ngram_stride + w, mask=w < CTX)
+ row = tl.load(ids + i * P + (w - CTX), mask=(w >= CTX) & (w < CTX + P)).to(tl.int64)
+ tl.store(stream + i * (CTX + P) + w, tl.where(w < CTX, hist, row), mask=w < CTX + P)
+
+
+@triton.jit
+def _commit_kernel(
+ accepted,
+ slots,
+ ngram_pool,
+ ngram_stride,
+ stream,
+ conv_pool,
+ conv_s0,
+ conv_s1,
+ conv_s2,
+ conv_prev,
+ conv_x,
+ x_s0,
+ x_s1,
+ x_s2,
+ position,
+ pos_s0,
+ pos_s1,
+ pos_s2,
+ position_old,
+ pages,
+ within,
+ plane_ptrs,
+ plane_strides,
+ index_old,
+ num_planes,
+ rows,
+ channels,
+ P: tl.constexpr,
+ P_BLOCK: tl.constexpr,
+ CTX: tl.constexpr,
+ CTX_BLOCK: tl.constexpr,
+ STATE: tl.constexpr,
+ STATE_BLOCK: tl.constexpr,
+ BLOCK_C: tl.constexpr,
+ ROW: tl.constexpr,
+):
+ i = tl.program_id(0)
+ cb = tl.program_id(1)
+ a = tl.load(accepted + i).to(tl.int32)
+ slot = tl.load(slots + i).to(tl.int64)
+ i64 = i.to(tl.int64)
+
+ # PLE conv row: column `a + s` of [pre-step state | conv input of rows 0..P-1].
+ c = cb * BLOCK_C + tl.arange(0, BLOCK_C)
+ s = tl.arange(0, STATE_BLOCK)
+ j = a + s
+ m = (c < channels)[:, None] & (s < STATE)[None, :]
+ from_prev = (j < STATE)[None, :]
+ prev = tl.load(
+ conv_prev + (i64 * channels + c[:, None]) * STATE + j[None, :], mask=m & from_prev
+ )
+ xin = tl.load(
+ conv_x + i64 * x_s0 + (j - STATE).to(tl.int64)[None, :] * x_s1 + c[:, None] * x_s2,
+ mask=m & ~from_prev,
+ )
+ tl.store(
+ conv_pool + slot * conv_s0 + c[:, None] * conv_s1 + s[None, :] * conv_s2,
+ tl.where(from_prev, prev, xin),
+ mask=m,
+ )
+
+ # The rows of ordinal >= a get back the plane values gathered before the trunk wrote
+ # them: index-K plane `cb` here, the position plane in block 0.
+ o = tl.arange(0, P_BLOCK)
+ r = i64 * P + o
+ rejected = (o < P) & (o >= a)
+ if cb < num_planes:
+ page = tl.load(pages + r, mask=rejected, other=0)
+ tok = tl.load(within + r, mask=rejected, other=0)
+ e = tl.arange(0, ROW)
+ base = tl.load(plane_ptrs + cb).to(tl.pointer_type(tl.bfloat16))
+ stride = tl.load(plane_strides + cb)
+ vals = tl.load(
+ index_old + (cb * rows + r[:, None]) * ROW + e[None, :], mask=rejected[:, None]
+ )
+ tl.store(
+ base + page[:, None] * stride + tok[:, None] * ROW + e[None, :],
+ vals,
+ mask=rejected[:, None],
+ )
+
+ if cb == 0:
+ # n-gram row: the id window after row `a - 1` (a negative row wraps, as torch
+ # indexing does).
+ t = a - 1
+ t = tl.where(t < 0, t + P, t)
+ k = tl.arange(0, CTX_BLOCK)
+ hist = tl.load(stream + i64 * (CTX + P) + t + 1 + k, mask=k < CTX)
+ tl.store(ngram_pool + slot * ngram_stride + k, hist, mask=k < CTX)
+
+ page = tl.load(pages + r, mask=rejected, other=0)
+ tok = tl.load(within + r, mask=rejected, other=0)
+ ax = tl.arange(0, 4)
+ pm = rejected[:, None] & (ax < 3)[None, :]
+ old_pos = tl.load(position_old + r[:, None] * 3 + ax[None, :], mask=pm)
+ tl.store(
+ position + page[:, None] * pos_s0 + tok[:, None] * pos_s1 + ax[None, :] * pos_s2,
+ old_pos,
+ mask=pm,
+ )
+
+
+def _pow2(n: int) -> int:
+ return triton.next_power_of_2(max(int(n), 1))
+
+
+class SpecRollbackStage(OpWrapper):
+ """`(torch.cat([ngram_pool[slots], ids.reshape(n, p).to(int64)], 1),
+ conv_pool[slots].contiguous())`: per generation request, its n-gram id history followed
+ by its `p` rows' ids, and its PLE conv state, read before the trunk writes either pool.
+
+ Fusion boundary: the two reads of the staging; the slots and ids are the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="MTP3, one request", spec=dict(n=1, p=4)),
+ Cell(why="MTP3, 128 requests", spec=dict(n=128, p=4)),
+ Cell(
+ why="tep4 MTP1 verify: 128 requests x 2 rows, the 10240 x 9 PLE conv over 1874 slots",
+ spec=dict(n=128, p=2, **_TEP4_PLE),
+ ),
+ Cell(
+ why="tep4 MTP7 verify: 128 requests x 8 rows, the 10240 x 9 PLE conv over 1874 slots",
+ spec=dict(n=128, p=8, **_TEP4_PLE),
+ ),
+ )
+ note = """
+ Metadata consumed: none. Call site: `SpeculativeRollback.stage_ple_verify`, before
+ `ple_ngram_embed` publishes the id history. Grid (requests, conv channel blocks).
+
+ The tp1 cells run a 640-channel, 3-column conv state; the default target's (`_TEP4_PLE`) the
+ checkpoint's 10240 x 9 state at a tep4 rank's measured slot count, `ids` int32 and `slots`
+ int32 as the call site hands them, slots 0 and 1873 among those read.
+ """
+
+ def __call__(self, ngram_pool, conv_pool, slots, ids, p: int):
+ n = slots.shape[0]
+ ctx = ngram_pool.shape[1]
+ channels, state = conv_pool.shape[1], conv_pool.shape[2]
+ stream = torch.empty([n, ctx + p], dtype=ngram_pool.dtype, device=ngram_pool.device)
+ conv_prev = torch.empty(
+ [n, channels, state], dtype=conv_pool.dtype, device=conv_pool.device
+ )
+ if n:
+ _stage_kernel[(n, triton.cdiv(channels, _BLOCK_C))](
+ ngram_pool,
+ ngram_pool.stride(0),
+ conv_pool,
+ conv_pool.stride(0),
+ conv_pool.stride(1),
+ conv_pool.stride(2),
+ slots,
+ ids,
+ stream,
+ conv_prev,
+ channels,
+ P=p,
+ CTX=ctx,
+ W_BLOCK=_pow2(ctx + p),
+ STATE=state,
+ STATE_BLOCK=_pow2(state),
+ BLOCK_C=_BLOCK_C,
+ num_warps=4,
+ )
+ return stream, conv_prev
+
+ def reference(self, ngram_pool, conv_pool, slots, ids, p: int):
+ n = slots.shape[0]
+ slots = slots.to(torch.long)
+ rows = torch.reshape(ids[: n * p], [n, p]).to(ngram_pool.dtype)
+ return torch.cat([ngram_pool[slots], rows], 1), conv_pool[slots]
+
+ def is_valid(self, ngram_pool, conv_pool, slots, ids, p: int) -> None:
+ assert ngram_pool.dim() == 2 and ngram_pool.stride(1) == 1
+ assert ngram_pool.dtype == torch.int64
+ assert conv_pool.dim() == 3
+ assert slots.dim() == 1 and ids.dim() == 1 and ids.is_contiguous()
+ assert ids.shape[0] >= slots.shape[0] * p
+
+
+class SpecRollbackCommit(OpWrapper):
+ """For each generation request `i` with `a = accepted[i]` accepted rows (golden token
+ included, `1 <= a <= p`) and slot `slots[i]`:
+
+ * `ngram_pool[slot] = stream[i, a : a + context]`;
+ * `conv_pool[slot] = cat([conv_prev[i], conv_x[i].T], 1)[:, a : a + state_len]`;
+ * for each row `r = i * p + t` with `t >= a`: `position[pages[r], within[r]] =
+ position_old[r]` and `plane_l[pages[r], within[r]] = index_old[l, r]` for every
+ index-K plane `l` of `plane_table`.
+
+ Fusion boundary: the whole commit of `SpeculativeRollback`; the staged tensors are what
+ `SpecRollbackStage`, `qsa_rollback_gather` and the position gather produced.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="MTP3, one request, 12 planes", spec=dict(n=1, p=4, planes=12)),
+ Cell(why="MTP3, 128 requests, 12 planes", spec=dict(n=128, p=4, planes=12)),
+ Cell(
+ why="tep4 MTP1 verify: 128 requests x 2 rows (256), 12 planes of the coalesced pool",
+ spec=dict(n=128, p=2, **_TEP4),
+ ),
+ Cell(
+ why="tep4 MTP7 verify: 128 requests x 8 rows (1024, the largest verify step), 12 planes",
+ spec=dict(n=128, p=8, **_TEP4),
+ ),
+ )
+ note = """
+ Metadata consumed: none. Call site: `SpeculativeRollback.commit_speculative_states`,
+ after the draft loop. Grid (requests, max(conv channel blocks, planes)). Slots or plane
+ slots shared by two requests (CUDA-graph padding dummies) are written by both; which one
+ lands is unordered.
+
+ Two grid regimes. The tp1 cells' 640 conv channels are 3 blocks under 12 planes, so blocks
+ 3..11 restore planes only. The default target's 10240 channels (`_TEP4`) are 40 blocks over 12
+ planes, so blocks 12..39 write conv channels only. The tep4 cells run that layout:
+ - `accepted` and `slots` are int32, as the MTP worker hands them.
+ - The planes are the 12 trunk layers' views of the coalesced pool at the measured tep4
+ page count and stride, addressed through `plane_table`.
+ - One request's rows straddle the pool's first and last page, and its rejected row is
+ restored on the last page, 257108 x 49152 elements (1.26e10) past the plane's base.
+ - Slots 0 and 1873 are among those written.
+
+ The commit is pure data movement, held bitwise to `reference`. The test also holds that
+ nothing outside the named slots of any pool is written, and that a CUDA-graph replay at
+ 1024 rows re-reads its inputs.
+ """
+
+ def __call__(
+ self,
+ accepted,
+ slots,
+ ngram_pool,
+ stream,
+ conv_pool,
+ conv_prev,
+ conv_x,
+ position,
+ position_old,
+ pages,
+ within,
+ plane_ptrs,
+ plane_strides,
+ index_old,
+ ) -> None:
+ n, p, channels = conv_x.shape
+ state = conv_pool.shape[2]
+ ctx = ngram_pool.shape[1]
+ if not n:
+ return
+ blocks = max(triton.cdiv(channels, _BLOCK_C), index_old.shape[0])
+ _commit_kernel[(n, blocks)](
+ accepted,
+ slots,
+ ngram_pool,
+ ngram_pool.stride(0),
+ stream,
+ conv_pool,
+ conv_pool.stride(0),
+ conv_pool.stride(1),
+ conv_pool.stride(2),
+ conv_prev,
+ conv_x,
+ conv_x.stride(0),
+ conv_x.stride(1),
+ conv_x.stride(2),
+ position,
+ position.stride(0),
+ position.stride(1),
+ position.stride(2),
+ position_old,
+ pages,
+ within,
+ plane_ptrs,
+ plane_strides,
+ index_old,
+ index_old.shape[0],
+ index_old.shape[1],
+ channels,
+ P=p,
+ P_BLOCK=_pow2(p),
+ CTX=ctx,
+ CTX_BLOCK=_pow2(ctx),
+ STATE=state,
+ STATE_BLOCK=_pow2(state),
+ BLOCK_C=_BLOCK_C,
+ ROW=ROW,
+ num_warps=4,
+ )
+
+ def reference(
+ self,
+ accepted,
+ slots,
+ ngram_pool,
+ stream,
+ conv_pool,
+ conv_prev,
+ conv_x,
+ position,
+ position_old,
+ pages,
+ within,
+ planes,
+ index_old,
+ ):
+ """What every slot the commit names holds after it.
+
+ `(ngram [n, context], conv [n, C, state_len], position [n * p, 3], index [planes,
+ n * p, ROW])`: the n-gram and conv pool rows at `slots`, and the position plane and
+ each index-K plane at every row's `(pages, within)`. A row the request accepted keeps
+ what the plane holds now; a rejected one gets `position_old` / `index_old`.
+
+ Takes the planes themselves (`planes`, the views `plane_table` built `plane_ptrs` and
+ `plane_strides` from) in place of their address table. `slots`, `ngram_pool` and
+ `conv_pool` only say which rows those are and are not read. Data movement only: the
+ values are the inputs' own, so it is the exact answer.
+ """
+ n, p, _ = conv_x.shape
+ context, state_len = ngram_pool.shape[1], conv_pool.shape[2]
+ a = accepted.to(torch.long)
+ req = torch.arange(n, device=a.device)
+ # Window k of the unfold is the id history after k rows.
+ ngram = stream.unfold(1, context, 1)[req, a]
+ # [n, C, state_len + p]; window k of the unfold is the conv state after k tokens.
+ history = torch.cat([conv_prev, conv_x.transpose(1, 2)], 2)
+ conv = history.unfold(2, state_len, 1)[req, :, a]
+ rejected = (torch.arange(p, device=a.device)[None, :] >= a[:, None]).reshape(-1)
+ pg, tok = pages.to(torch.long), within.to(torch.long)
+ pos = torch.where(rejected[:, None], position_old, position[pg, tok])
+ index = torch.stack(
+ [
+ torch.where(rejected[:, None], index_old[k], plane[pg, tok].reshape(-1, ROW))
+ for k, plane in enumerate(planes)
+ ]
+ )
+ return ngram, conv, pos, index
+
+ def is_valid(
+ self,
+ accepted,
+ slots,
+ ngram_pool,
+ stream,
+ conv_pool,
+ conv_prev,
+ conv_x,
+ position,
+ position_old,
+ pages,
+ within,
+ plane_ptrs,
+ plane_strides,
+ index_old,
+ ) -> None:
+ # The kernel reads every staged tensor by raw offset in the layout its producer writes,
+ # and the pools by their strides. A wrong shape is read short or past its end, and a
+ # strided staged tensor at the wrong elements, without an error.
+ n, p, channels = conv_x.shape
+ assert accepted.shape == (n,) and slots.shape == (n,), (
+ f"accepted {tuple(accepted.shape)} and slots {tuple(slots.shape)} must hold one "
+ f"entry per generation request ({n})"
+ )
+ assert ngram_pool.dtype == stream.dtype == torch.int64 and ngram_pool.stride(1) == 1, (
+ "ngram_pool and stream must be int64, the n-gram history's last dim unit-strided"
+ )
+ assert stream.shape == (n, ngram_pool.shape[1] + p) and stream.is_contiguous(), (
+ f"stream must be the contiguous [n, context + p] = [{n}, {ngram_pool.shape[1] + p}] "
+ f"spec_rollback_stage wrote; got {tuple(stream.shape)}"
+ )
+ assert conv_pool.shape[1] == channels and conv_prev.is_contiguous(), (
+ "conv_prev must be contiguous (read as packed [n, C, state_len] rows), and conv_pool "
+ "as wide as conv_x"
+ )
+ assert conv_prev.shape == (n, channels, conv_pool.shape[2]), (
+ f"conv_prev is {tuple(conv_prev.shape)}, expected {(n, channels, conv_pool.shape[2])}"
+ )
+ assert position.dim() == 3 and position.shape[2] == 3, (
+ f"position must be the [pages, tokens_per_block, 3] plane; got {tuple(position.shape)}"
+ )
+ assert position_old.shape == (n * p, 3) and position_old.is_contiguous(), (
+ f"position_old must be the contiguous [n * p, 3] = [{n * p}, 3] gather"
+ )
+ assert pages.shape == within.shape == (n * p,), (
+ f"pages and within must hold one entry per verify row ({n * p})"
+ )
+ assert pages.dtype == within.dtype == torch.int64, "pages and within must be int64"
+ assert index_old.is_contiguous() and index_old.shape[1:] == (n * p, ROW), (
+ f"index_old must be the contiguous [planes, n * p, {ROW}] qsa_rollback_gather "
+ f"returns; got {tuple(index_old.shape)}, contiguous {index_old.is_contiguous()}"
+ )
+ assert plane_ptrs.numel() == plane_strides.numel() == index_old.shape[0], (
+ f"the plane table holds {plane_ptrs.numel()} planes and index_old "
+ f"{index_old.shape[0]}; the kernel restores index_old's count"
+ )
+
+
+spec_rollback_stage = SpecRollbackStage()
+spec_rollback_commit = SpecRollbackCommit()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/verify_step_setup.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/verify_step_setup.py
new file mode 100644
index 000000000000..edce2a631e91
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/verify_step_setup.py
@@ -0,0 +1,216 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The decode target's per-step setup of a verify step in one launch.
+
+Tier 5 (Triton written for this target). Before the trunk runs, a verify step needs:
+
+* `pos3`: a contiguous int32 `[3, T]` copy of the step's mRoPE positions;
+* the PLE segmentation vector `cu_seqlens` (`[batch + 1]` int64, leading zero) of the
+ per-request row counts;
+* the speculative rollback's staging: each row's QSA slot (`pages`, `within`, int64; page
+ `block_table[request, position // tokens_per_block]` and offset `position % tokens_per_block`,
+ floor semantics, a negative column wrapping as torch indexing does), the requests' recurrent
+ slots as int64, and every attention layer's index-K plane at those slots (`[L, T, 128]` bf16,
+ read before the trunk writes any plane).
+
+Program `(0, b)` writes block `b` of the rows' positions and slots (block 0 also the
+segmentation vector and the int64 slots); program `(1 + l, b)` copies plane `l` at the same
+rows, resolving each row's slot as program `(0, b)` does.
+"""
+
+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
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_rollback_planes import (
+ ROW,
+)
+
+#: Rows a program covers.
+_BLOCK_T = 16
+#: Most requests the segmentation vector covers (`batch + 1` entries).
+_BLOCK_S = 512
+
+
+@triton.jit
+def _slot(req_idx, logical, block_table, bt_stride, bt_cols, rows, tpb):
+ req = tl.load(req_idx + rows).to(tl.int64)
+ pos = tl.load(logical + rows).to(tl.int64)
+ q = pos // tpb
+ r = pos % tpb
+ neg = (r != 0) & (pos < 0)
+ q = tl.where(neg, q - 1, q)
+ r = tl.where(neg, r + tpb, r)
+ q = tl.where(q < 0, q + bt_cols, q)
+ page = tl.load(block_table + req * bt_stride + q).to(tl.int64)
+ return page, r
+
+
+@triton.jit
+def _setup_kernel(
+ positions,
+ pos_stride,
+ seq_lens,
+ slots,
+ req_idx,
+ logical,
+ block_table,
+ bt_stride,
+ bt_cols,
+ plane_ptrs,
+ plane_strides,
+ pos3,
+ cu_seqlens,
+ slots_out,
+ pages_out,
+ within_out,
+ index_out,
+ num_rows,
+ batch,
+ tpb,
+ BLOCK_T: tl.constexpr,
+ BLOCK_S: tl.constexpr,
+ ROW: tl.constexpr,
+):
+ who = tl.program_id(0)
+ tb = tl.program_id(1)
+ t = tb * BLOCK_T + tl.arange(0, BLOCK_T)
+ tm = t < num_rows
+ tt = tl.where(tm, t, 0)
+ page, within = _slot(req_idx, logical, block_table, bt_stride, bt_cols, tt, tpb)
+ if who == 0:
+ tl.store(pages_out + t, page, mask=tm)
+ tl.store(within_out + t, within, mask=tm)
+ for a in tl.static_range(3):
+ p = tl.load(positions + a * pos_stride + t, mask=tm)
+ tl.store(pos3 + a * num_rows + t, p.to(tl.int32), mask=tm)
+ if tb == 0:
+ s = tl.arange(0, BLOCK_S)
+ sm = s < batch
+ lengths = tl.load(seq_lens + s, mask=sm, other=0).to(tl.int64)
+ tl.store(cu_seqlens + 1 + s, tl.cumsum(lengths, axis=0), mask=sm)
+ tl.store(cu_seqlens + s, tl.zeros([BLOCK_S], dtype=tl.int64), mask=s == 0)
+ tl.store(slots_out + s, tl.load(slots + s, mask=sm).to(tl.int64), mask=sm)
+ else:
+ plane = who - 1
+ base = tl.load(plane_ptrs + plane).to(tl.pointer_type(tl.bfloat16))
+ stride = tl.load(plane_strides + plane)
+ e = tl.arange(0, ROW)
+ vals = tl.load(
+ base + page[:, None] * stride + within[:, None] * ROW + e[None, :], mask=tm[:, None]
+ )
+ tl.store(
+ index_out + (plane.to(tl.int64) * num_rows + t[:, None]) * ROW + e[None, :],
+ vals,
+ mask=tm[:, None],
+ )
+
+
+class VerifyStepSetup(OpWrapper):
+ """`(pos3, cu_seqlens, slots.long(), pages, within, index_old)`:
+
+ * `pos3 = position_ids.reshape(3, -1)` as a fresh contiguous int32 tensor;
+ * `cu_seqlens = [0, cumsum(seq_lens)]` int64;
+ * `pages = block_table[request_indices.long(), logical_positions // tokens_per_block]`
+ and `within = logical_positions % tokens_per_block`, int64;
+ * `index_old = qsa_rollback_gather(plane_ptrs, plane_strides, pages, within, L)`.
+
+ Fusion boundary: the decode target's step setup before the trunk; which rows, slots and
+ planes are the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="MTP3, one request, 12 planes", spec=dict(batch=1, p=4, planes=12)),
+ Cell(why="MTP3, 32 requests, 12 planes", spec=dict(batch=32, p=4, planes=12)),
+ )
+ note = """
+ Metadata consumed: none (the caller passes `seq_lens_cuda[:batch]`, the QSA row mapping
+ and block table). Call site: `DecodeTarget.forward` of a captured verify step, after the
+ captured preprocessing corrected positions and lengths and before the trunk.
+ Preconditions: `position_ids` reshapes to `[3, T]` with unit last stride; at most
+ `_BLOCK_S - 1` requests.
+ """
+
+ def __call__(
+ self,
+ position_ids,
+ seq_lens,
+ slots,
+ request_indices,
+ logical_positions,
+ block_table,
+ tokens_per_block: int,
+ plane_ptrs,
+ plane_strides,
+ num_planes: int,
+ ):
+ pos2 = torch.reshape(position_ids, [3, -1])
+ rows = pos2.shape[1]
+ batch = seq_lens.shape[0]
+ dev = pos2.device
+ pos3 = torch.empty([3, rows], dtype=torch.int32, device=dev)
+ cu_seqlens = torch.empty(batch + 1, dtype=torch.int64, device=dev)
+ slots_out = torch.empty(batch, dtype=torch.int64, device=dev)
+ pages = torch.empty(rows, dtype=torch.int64, device=dev)
+ within = torch.empty(rows, dtype=torch.int64, device=dev)
+ index_old = torch.empty([num_planes, rows, ROW], dtype=torch.bfloat16, device=dev)
+ if rows:
+ _setup_kernel[(1 + num_planes, triton.cdiv(rows, _BLOCK_T))](
+ pos2,
+ pos2.stride(0),
+ seq_lens,
+ slots,
+ request_indices,
+ logical_positions,
+ block_table,
+ block_table.stride(0),
+ block_table.shape[1],
+ plane_ptrs,
+ plane_strides,
+ pos3,
+ cu_seqlens,
+ slots_out,
+ pages,
+ within,
+ index_old,
+ rows,
+ batch,
+ tokens_per_block,
+ BLOCK_T=_BLOCK_T,
+ BLOCK_S=_BLOCK_S,
+ ROW=ROW,
+ num_warps=4,
+ )
+ return pos3, cu_seqlens, slots_out, pages, within, index_old
+
+ def reference(self, *args):
+ raise NotImplementedError("test_modeling_v2_verify_step_setup compares the composed ops")
+
+ def is_valid(
+ self,
+ position_ids,
+ seq_lens,
+ slots,
+ request_indices,
+ logical_positions,
+ block_table,
+ tokens_per_block: int,
+ plane_ptrs,
+ plane_strides,
+ num_planes: int,
+ ) -> None:
+ pos2 = torch.reshape(position_ids, [3, -1])
+ assert pos2.stride(1) == 1
+ assert 0 < seq_lens.shape[0] < _BLOCK_S and seq_lens.is_contiguous()
+ assert slots.shape == seq_lens.shape
+ assert request_indices.shape[0] >= pos2.shape[1]
+ assert logical_positions.shape[0] >= pos2.shape[1]
+ assert block_table.stride(1) == 1 and tokens_per_block > 0
+ assert plane_ptrs.numel() == plane_strides.numel() == num_planes
+
+
+verify_step_setup = VerifyStepSetup()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/embedding/__init__.py
new file mode 100644
index 000000000000..52a7a9daf028
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/embedding/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/embedding/ngram_hash.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/embedding/ngram_hash.py
new file mode 100644
index 000000000000..420b5f4a9499
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/embedding/ngram_hash.py
@@ -0,0 +1,286 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""PLE n-gram hash: a token-id history becomes one global embedding row per head.
+
+A macro entry: composed from native torch integer ops (~30 small launches);
+exact, so gated bitwise. It serves the general form -- any history length and
+`num_query_tokens`, a chunk's tail included. The Triton op
+``torch.ops.trtllm.qwen4_exp_ple_ngram_hash`` covers the ``[T, 3]`` window,
+last-position form only, and is not wrapped here; see `note`.
+``EngramHashProvider.compute_hashes`` in
+``tensorrt_llm/_torch/modules/engram/engram.py`` is a different function: it
+hashes *compressed* ids (a tokenizer table this checkpoint does not ship),
+derives its multipliers from a seed, and pads instead of resetting at every
+EOS.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+#: This checkpoint's `ple_embedding` buffers: 3-gram, 8 heads per order.
+CKPT_MULTIPLIERS = (23703573157769, 20109073645365, 8052911324071)
+CKPT_HEAD_VOCAB_SIZES = (
+ 20000003, 20000023, 20000033, 20000047, 20000059, 20000063, 20000069, 20000077,
+ 20000081, 20000093, 20000107, 20000147, 20000153, 20000159, 20000161, 20000171,
+) # fmt: skip
+CKPT_HEAD_OFFSETS = (
+ 0, 20000003, 40000026, 60000059, 80000106, 100000165, 120000228, 140000297,
+ 160000374, 180000455, 200000548, 220000655, 240000802, 260000955, 280001114, 300001275,
+) # fmt: skip
+CKPT_EOS_TOKEN_ID = 248044
+CKPT_VOCAB_SIZE = 248320
+#: `ngram_size - 1`: the cached ids ahead of the new tokens.
+CKPT_CONTEXT_LEN = len(CKPT_MULTIPLIERS) - 1
+
+
+def _rows(batch: int, length: int, query: int, **spec) -> dict:
+ return dict(batch=batch, length=length, num_query_tokens=query, **spec)
+
+
+class NgramHash(OpWrapper):
+ """For a query token at history position `t` and head `h` of the `n`-gram
+ block (`m = layer_multipliers`, `^` int64 XOR, `%` non-negative)::
+
+ mixed[t] = (c[t] * m[0]) ^ (c[t - 1] * m[1]) ^ ... ^ (c[t - n + 1] * m[n - 1])
+ row[t, h] = mixed[t] % head_vocab_sizes[h] + head_offsets[h]
+
+ `c[t-k]` is `token_history[t-k]` when `t - k` lies in `t`'s segment and
+ `eos_token_id` otherwise; a segment starts one past the most recent EOS
+ **strictly before** `t` (an EOS closes its segment, so the token after it
+ hashes as a sequence's first). Shifts before position 0 read EOS too.
+ Heads run `[2-gram block | ... | ngram_size-gram block]`, `heads_per_ngram
+ = head_vocab_sizes.numel() // (ngram_size - 1)` each, `ngram_size =
+ layer_multipliers.numel()`.
+
+ Returns a fresh contiguous int64 `[batch, num_query_tokens, ngram_heads]`
+ for the **last** `num_query_tokens` positions; the leading ones only give
+ the first query token its context. Indices are **global** -- the offset is
+ applied, one concatenated table is addressed. No input is modified.
+
+ Fusion boundary: segment scan, shifted streams, mix, modulo and offset.
+ The id history (here built by `ngram_token_history`) and the table gather
+ are the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: batch, length, num_query_tokens; history is uniform random 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.
+ #
+ # The call site's form is one `[T, ngram_size]` window per packed token
+ # from `ngram_token_history`, num_query_tokens 1: T is the step's packed
+ # row count (requests x rows per request). The chunk-tail cell is the same
+ # rows hashed per request instead: `ngram_size - 1` carried ids ahead of
+ # the request's query tokens.
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="decode row whose cached context is EOS",
+ spec=_rows(4, 3, 1, eos_positions=(0, 1)),
+ ),
+ Cell(
+ why="call site: 32 windows, decode / MTP verify",
+ spec=_rows(32, 3, 1, lead_eos=False, eos_density=0.2),
+ ),
+ Cell(
+ why="call site: 8192 windows, largest prefill",
+ spec=_rows(8192, 3, 1, lead_eos=False, eos_density=0.05),
+ ),
+ Cell(
+ why="call site: 128 windows, a 128-request decode batch",
+ spec=_rows(128, 3, 1, lead_eos=False, eos_density=0.05),
+ ),
+ Cell(
+ why="call site: 1024 windows, 128 requests x 8 rows of an MTP-7 verify",
+ spec=_rows(1024, 3, 1, lead_eos=False, eos_density=0.05),
+ ),
+ Cell(
+ why="chunk tail: 128 requests, 2 carried ids then 8 query tokens each",
+ spec=_rows(128, CKPT_CONTEXT_LEN + 8, 8, lead_eos=False, eos_density=0.05),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none; rows are independent, and within a row the only
+ structure seen is `eos_token_id` -- packing two sequences into one row is
+ safe only when the first ends in EOS. Call site: `ngram_hash(window,
+ multipliers, head_vocab, head_offsets, eos_token_id, 1)` on the
+ `[T, 3]` int64 window, reshaped to `[T, 16]` by the caller.
+
+ Preconditions: `head_vocab_sizes`, `head_offsets` on `token_history`'s
+ device (raises otherwise; a cpu `layer_multipliers` is accepted and
+ gives the identical result). Certified on an int64 `token_history`, the
+ window `ngram_token_history` returns. Negative ids are accepted and hashed
+ as ordinary ids, not padding. A chunk's tail (`num_query_tokens > 1`) is
+ certified at 8 query tokens behind the 2 carried ids, 128 requests;
+ longer histories are not certified. `num_query_tokens` above the history
+ length returns a silently shorter `[batch, history_len, heads]`, so a
+ `(length 3, num_query_tokens 8)` call is `is_valid`'s to refuse.
+ A single multiplier raises `ZeroDivisionError`. The four silent mistakes
+ are `is_valid`'s.
+
+ A wrong hash is a valid row: an extra shift, no EOS reset, an inclusive
+ segment start, reversed multipliers, no offsets, swapped blocks all stay
+ in range and inside the table. The tail slice is taken before the mix
+ (bitwise equal to slicing after), and chunked prefill with the
+ `ngram_size - 1` carried ids reproduces one-shot bitwise.
+
+ CUDA-graph capture certified at 128 and 1024 windows (the captured decode
+ and MTP-7 verify widths), bitwise equal to the reference after refilling
+ the history.
+
+ Not wrapped: `torch.ops.trtllm.qwen4_exp_ple_ngram_hash` (Triton,
+ `tensorrt_llm/_torch/modules/qwen4_exp/ple_kernels.py`) on
+ `(window [T, 3], multipliers, head_vocab_sizes, head_offsets, eos)` ->
+ `[T, 16]` is bitwise equal to `reference` with num_query_tokens 1 at 1,
+ 128, 1024 and 8192 windows. It is that form only: a longer history raises
+ `ValueError`, it takes no `num_query_tokens`, and a negative id or a
+ product past 2^63 gives a negative row (C remainder) where this
+ composition's stays inside the table. The target's PLE hashes
+ inside `ple_ngram_embed`'s own launch and reads only `reference` here.
+
+ Hard limits
+ token id max_id * max(layer_multipliers) < 2^63 boundary 389114 for
+ this checkpoint's multipliers probed above it: accepted, no error,
+ indices still inside the table (int64 wraps silently). The largest
+ legal id 248319 gives 5.886e18. Not guarded: data, would sync.
+ """
+
+ def __call__(
+ self,
+ token_history: torch.Tensor,
+ layer_multipliers: torch.Tensor,
+ head_vocab_sizes: torch.Tensor,
+ head_offsets: torch.Tensor,
+ eos_token_id: int,
+ num_query_tokens: int,
+ ) -> torch.Tensor:
+ ids = token_history.long()
+ batch_size, history_len = ids.shape
+ ngram_size = layer_multipliers.numel()
+ total_heads = head_vocab_sizes.numel()
+ heads_per_ngram = total_heads // (ngram_size - 1)
+
+ # Segment bookkeeping over the FULL history: the most recent EOS
+ # strictly before each position starts that position's segment.
+ positions = torch.arange(history_len, device=ids.device, dtype=torch.long)
+ minus_one = positions.new_full((), -1)
+ eos_positions = torch.where(ids == eos_token_id, positions, minus_one)
+ previous_eos_inclusive = torch.cummax(eos_positions, dim=1).values
+ previous_eos = torch.cat(
+ [eos_positions.new_full((batch_size, 1), -1), previous_eos_inclusive[:, :-1]],
+ dim=1,
+ )
+ position_in_segment = positions.unsqueeze(0) - (previous_eos + 1)
+
+ # Elementwise along the token axis from here on, so slicing to the
+ # query tail now is bitwise equal to slicing after the modulo.
+ keep_from = history_len - num_query_tokens
+ eos_fill = ids.new_full((), eos_token_id)
+ shifted = [ids[:, keep_from:]]
+ for shift in range(1, ngram_size):
+ source_positions = positions - shift
+ gather_positions = source_positions.clamp_min(0).unsqueeze(0).expand(batch_size, -1)
+ gathered = ids.gather(1, gather_positions)
+ valid = (position_in_segment >= shift) & (source_positions.unsqueeze(0) >= 0)
+ shifted.append(torch.where(valid, gathered, eos_fill)[:, keep_from:])
+
+ blocks = []
+ for ngram in range(2, ngram_size + 1):
+ low = (ngram - 2) * heads_per_ngram
+ high = low + heads_per_ngram
+ mixed = shifted[0] * layer_multipliers[0]
+ for position in range(1, ngram):
+ mixed = torch.bitwise_xor(mixed, shifted[position] * layer_multipliers[position])
+ blocks.append(
+ torch.remainder(mixed.unsqueeze(-1), head_vocab_sizes[low:high].view(1, 1, -1))
+ + head_offsets[low:high].view(1, 1, -1)
+ )
+ return torch.cat(blocks, dim=-1)
+
+ def reference(
+ self,
+ token_history: torch.Tensor,
+ layer_multipliers: torch.Tensor,
+ head_vocab_sizes: torch.Tensor,
+ head_offsets: torch.Tensor,
+ eos_token_id: int,
+ num_query_tokens: int,
+ ) -> torch.Tensor:
+ """Scalar Python-int arithmetic, one token and one head at a time, the
+ segment start tracked by walking the row; no torch op computes it."""
+ multipliers = layer_multipliers.tolist()
+ sizes = head_vocab_sizes.tolist()
+ offsets = head_offsets.tolist()
+ ngram_size = len(multipliers)
+ heads_per_ngram = len(sizes) // (ngram_size - 1)
+ out = []
+ for row in token_history.tolist():
+ segment_start, last_eos = [], -1
+ for position, token in enumerate(row):
+ segment_start.append(last_eos + 1)
+ if token == eos_token_id:
+ last_eos = position
+ row_out = []
+ for position in range(len(row) - num_query_tokens, len(row)):
+ context = [
+ row[position - shift]
+ if position - segment_start[position] >= shift and position - shift >= 0
+ else eos_token_id
+ for shift in range(ngram_size)
+ ]
+ heads = []
+ for ngram in range(2, ngram_size + 1):
+ mixed = context[0] * multipliers[0]
+ for index in range(1, ngram):
+ mixed ^= context[index] * multipliers[index]
+ low = (ngram - 2) * heads_per_ngram
+ heads.extend(
+ mixed % sizes[low + h] + offsets[low + h] for h in range(heads_per_ngram)
+ )
+ row_out.append(heads)
+ out.append(row_out)
+ return torch.tensor(out, dtype=torch.long, device=token_history.device).reshape(
+ token_history.shape[0], num_query_tokens, len(sizes)
+ )
+
+ def is_valid(
+ self,
+ token_history: torch.Tensor,
+ layer_multipliers: torch.Tensor,
+ head_vocab_sizes: torch.Tensor,
+ head_offsets: torch.Tensor,
+ eos_token_id: int,
+ num_query_tokens: int,
+ ) -> None:
+ # Each measured to be accepted and to return plausible in-range rows.
+ history_len = token_history.shape[-1]
+ ngram_size = layer_multipliers.numel()
+ total_heads = head_vocab_sizes.numel()
+ assert 0 <= num_query_tokens <= history_len, (
+ f"num_query_tokens must be in [0, history_len]; got {num_query_tokens} "
+ f"against a history of {history_len} -- a value outside that range is "
+ f"silently truncated to a shorter output rather than rejected"
+ )
+ assert layer_multipliers.dtype == torch.int64, (
+ f"layer_multipliers must be int64; got {layer_multipliers.dtype} -- a "
+ f"narrower dtype is accepted and silently hashes with truncated "
+ f"multipliers"
+ )
+ assert ngram_size < 2 or total_heads % (ngram_size - 1) == 0, (
+ f"head_vocab_sizes.numel()={total_heads} must be a multiple of "
+ f"ngram_size-1={ngram_size - 1} -- the blocks are equal width, and a "
+ f"remainder is silently dropped rather than rejected"
+ )
+ assert head_offsets.numel() == total_heads, (
+ f"head_offsets and head_vocab_sizes must have the same length; got "
+ f"{head_offsets.numel()} and {total_heads} -- an over-long head_offsets "
+ f"is silently read for its first {total_heads} entries only"
+ )
+
+
+ngram_hash = NgramHash()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/__init__.py
new file mode 100644
index 000000000000..52a7a9daf028
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/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/ops/gemm/bf16_gemv_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/bf16_gemv_pdl.py
new file mode 100644
index 000000000000..75652e898bae
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/bf16_gemv_pdl.py
@@ -0,0 +1,388 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Decode-row bf16 projection launched with programmatic dependent launch (PDL),
+its weight streaming in while the kernel before it finishes.
+
+Tier 5 (CuTe DSL written for this target). At decode row counts a projection
+reads its weight once and its activation is a few KB, so the GEMV's time is the
+weight stream plus its launch and ramp. A normal launch cannot read a byte of
+that stream until the producer of the activation has exited. Launched with PDL,
+every CTA loads its weight rows -- static, so safe to read early -- into
+registers before `griddepcontrol.wait` and triggers its own dependents right
+before it; the activation is read, and the output written, only after the wait.
+
+One CTA computes `OUTS` output columns over the whole K for every row; element
+`k` of a row is held by thread `(k // VEC) % THREADS` as its vector
+`(k // VEC) // THREADS`, so each k-tile of loads is contiguous across a warp.
+Products and sums in fp32: a warp reduction, then a cross-warp one in shared
+memory; one rounding at the store.
+
+A K with a shared-memory tier (`SOUTS`) gives each CTA `SOUTS` more outputs whose
+weight rows it stages in shared memory with cp.async before the wait, each thread
+exactly the vectors it later reads (so `cp.async.wait_group` alone orders them). The
+2560 x 6144 projection then runs as one wave of 640 CTAs whose whole weight is on
+chip by the wait, instead of 1280 CTAs of which the last third start loading theirs
+only as the first ones exit.
+"""
+
+from __future__ import annotations
+
+import functools
+
+import cuda.bindings.driver as cuda_drv
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass import const_expr
+from cutlass._mlir.dialects import llvm
+from cutlass.cute import experimental as cute_ext
+from cutlass.cute.nvgpu import cpasync
+from cutlass.cute.runtime import from_dlpack
+from cutlass.cutlass_dsl import dsl_user_op
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import (
+ Arch,
+ Cell,
+ OpWrapper,
+ true_fp32_matmul,
+)
+from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+
+#: Largest row count one launch serves; every row re-reads the weight from registers.
+MAX_ROWS = 4
+#: (threads, outputs per CTA, elements per vector load) per K. Swept on GB300 at one row,
+#: graph replay over per-layer weights: K=2560 64 threads x 2 outputs; K=640 (80 16-byte
+#: vectors, which no warp multiple divides) 32 threads x 4 outputs on 8-byte vectors;
+#: K=1536 (a tp4/tep4 rank's o / GDN out projection) as K=2560.
+TILING = {2560: (64, 2, 8), 6144: (128, 2, 8), 640: (32, 4, 4), 1536: (64, 2, 8)}
+#: Outputs per CTA whose weight rows sit in shared memory (beside the `TILING` ones in
+#: registers). K=6144 swept behind the GDN recurrence and the sparse-attention merge
+#: (one row, graph replay): 2 + 2 a CTA beat 2 + 0, 2 + 6, 4 + 4 and 1 + 3.
+SOUTS = {6144: 2}
+_COMPILE_OPTIONS = "--ptxas-options -maxrregcount=64"
+_COMPILE_OPTIONS_SMEM = "--ptxas-options -maxrregcount=80"
+#: CTAs a multiprocessor must hold for a grid with a shared-memory tier: four, the
+#: 640-CTA grid's share of 160 SMs. Behind the sparse-attention merge (one row, graph
+#: replay) 4 runs the attention tail 13.5 us a layer against 14.6 at 1 and 15.6 at 2-3;
+#: the GDN tail is unchanged.
+MIN_BLOCKS_SMEM = 4
+
+
+def tiling(n: int, k: int) -> tuple[int, int, int, int]:
+ """`(threads, outs, souts, vec)` for a `[*, K] x [K, N]` product: register and shared-memory
+ outputs a CTA. Without a shared-memory tier that divides N the outputs shrink to divide N."""
+ threads, outs, vec = TILING[k]
+ souts = SOUTS.get(k, 0)
+ if souts and n % (outs + souts) == 0:
+ return threads, outs, souts, vec
+ while n % outs:
+ outs //= 2
+ return threads, outs, 0, vec
+
+
+#: Before its PDL wait a CTA also has the weight rows of CTA `bidx + dist` brought into
+#: L2: behind a latency-bound kernel only part of a large grid's CTAs fit, and the later
+#: waves' weights would otherwise start streaming only after it. `dist` is about the CTAs
+#: resident before the wait: 1216 behind the hyper-connection kernels (measured), at
+#: least 912 (a 1280-CTA out projection's first wave) otherwise. Grids smaller than `dist`
+#: prefetch nothing.
+L2_DIST = 1216
+L2_DIST_MIN = 912
+
+
+def l2_dist(grid: int) -> int:
+ return min(L2_DIST, max(grid // 2, L2_DIST_MIN))
+
+
+@dsl_user_op
+def _bulk_prefetch_l2(src, nbytes, *, loc=None, ip=None):
+ llvm.inline_asm(
+ None,
+ [
+ cutlass.Int64(src.toint(loc=loc, ip=ip)).ir_value(loc=loc, ip=ip),
+ cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip),
+ ],
+ "cp.async.bulk.prefetch.L2.global [$0], $1;",
+ "l,r",
+ has_side_effects=True,
+ is_align_stack=False,
+ asm_dialect=llvm.AsmDialect.AD_ATT,
+ loc=loc,
+ ip=ip,
+ )
+
+
+class _GemvKernel:
+ def __init__(self, *, rows, k, threads, outs, vec, pdl, dist=L2_DIST, souts=0):
+ assert k % (threads * vec) == 0 and threads % 32 == 0
+ self.rows, self.threads, self.outs, self.souts, self.vec = rows, threads, outs, souts, vec
+ self.dist = dist
+ self.k_tiles = k // (threads * vec)
+ self.warps = threads // 32
+ self.pdl = pdl
+
+ def smem_bytes(self) -> int:
+ red = self.rows * (self.outs + self.souts) * self.warps * 4
+ if not self.souts:
+ return red
+ w = self.souts * self.k_tiles * self.threads * self.vec * 2
+ return -(-w // 128) * 128 + -(-red // 128) * 128 + 128
+
+ @cute.jit
+ def __call__(self, mA, mW, mOut, stream: cuda_drv.CUstream):
+ n = cute.size(mW, mode=[0])
+ load_a = cute.make_copy_atom(
+ cute.nvgpu.CopyG2ROp(),
+ cutlass.BFloat16,
+ num_bits_per_copy=16 * self.vec,
+ load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS,
+ )
+ load_w = cute.make_copy_atom(
+ cute.nvgpu.CopyG2ROp(),
+ cutlass.BFloat16,
+ num_bits_per_copy=16 * self.vec,
+ load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING,
+ )
+ g2s = cute.make_copy_atom(
+ cpasync.CopyG2SOp(), cutlass.BFloat16, num_bits_per_copy=16 * self.vec
+ )
+ self.kernel(mA, mW, mOut, load_a, load_w, g2s).launch(
+ grid=[cute.ceil_div(n, self.outs + self.souts), 1, 1],
+ block=[self.threads, 1, 1],
+ smem=self.smem_bytes(),
+ stream=stream,
+ use_pdl=self.pdl,
+ min_blocks_per_mp=MIN_BLOCKS_SMEM if self.souts else 1,
+ )
+
+ @cute.kernel
+ def kernel(
+ self, mA, mW, mOut, load_a: cute.CopyAtom, load_w: cute.CopyAtom, g2s: cute.CopyAtom
+ ):
+ tidx, _, _ = cute.arch.thread_idx()
+ bidx, _, _ = cute.arch.block_idx()
+ warp = cute.arch.warp_idx()
+ rows: cutlass.Constexpr = self.rows
+ outs: cutlass.Constexpr = self.outs
+ threads: cutlass.Constexpr = self.threads
+ vec: cutlass.Constexpr = self.vec
+ k_tiles: cutlass.Constexpr = self.k_tiles
+ warps: cutlass.Constexpr = self.warps
+ souts: cutlass.Constexpr = self.souts
+ tot: cutlass.Constexpr = outs + souts
+ n0 = bidx * tot
+
+ # (row, (vec, (threads, k_tiles))): thread t's vectors of a row, k-tile by k-tile.
+ w_vec = cute.logical_divide(cute.logical_divide(mW, (None, vec)), (None, (None, threads)))
+ a_vec = cute.logical_divide(cute.logical_divide(mA, (None, vec)), (None, (None, threads)))
+ smem = cutlass.utils.SmemAllocator()
+
+ # The weight is static: it streams before the wait, while the producer runs. The
+ # shared-memory tier first ([output, k-tile, thread, vec]: a warp's vectors are
+ # consecutive), so it lands while the register tier streams.
+ if const_expr(souts > 0):
+ s_w = smem.allocate_tensor(
+ cutlass.BFloat16,
+ cute.make_layout(
+ (souts, k_tiles, threads, vec),
+ stride=(k_tiles * threads * vec, threads * vec, vec, 1),
+ ),
+ byte_alignment=128,
+ )
+ for j in cutlass.range_constexpr(souts):
+ s_row = w_vec[n0 + outs + j, (None, (tidx, None))]
+ for kt in cutlass.range_constexpr(k_tiles):
+ cute.copy(g2s, s_row[None, kt], s_w[j, kt, tidx, None])
+ cute.arch.cp_async_commit_group()
+ w_reg = cute.make_rmem_tensor(
+ cute.make_layout((outs, k_tiles, vec), stride=(k_tiles * vec, vec, 1)), cutlass.BFloat16
+ )
+ for j in cutlass.range_constexpr(outs):
+ w_row = w_vec[n0 + j, (None, (tidx, None))]
+ for kt in cutlass.range_constexpr(k_tiles):
+ cute.copy(load_w, w_row[None, kt], w_reg[j, kt, None])
+ if const_expr(self.pdl):
+ n_rows = cute.size(mW, mode=[0])
+ k_len = cute.size(mW, mode=[1])
+ if tidx == 0:
+ if (bidx + self.dist + 1) * tot <= n_rows:
+ _bulk_prefetch_l2(
+ mW.iterator + (bidx + self.dist) * tot * k_len, tot * k_len * 2
+ )
+ cute.arch.griddepcontrol_launch_dependents()
+ cute.arch.griddepcontrol_wait()
+
+ acc = cute.make_rmem_tensor(cute.make_layout((rows, tot), stride=(tot, 1)), cutlass.Float32)
+ acc.fill(0.0)
+ a_reg = cute.make_rmem_tensor(
+ cute.make_layout((k_tiles, vec), stride=(vec, 1)), cutlass.BFloat16
+ )
+ sw_reg = cute.make_rmem_tensor(cute.make_layout((vec,), stride=(1,)), cutlass.BFloat16)
+ for r in cutlass.range_constexpr(rows):
+ a_row = a_vec[r, (None, (tidx, None))]
+ for kt in cutlass.range_constexpr(k_tiles):
+ cute.copy(load_a, a_row[None, kt], a_reg[kt, None])
+ for kt in cutlass.range_constexpr(k_tiles):
+ for v in cutlass.range_constexpr(vec):
+ av = a_reg[kt, v].to(cutlass.Float32)
+ for j in cutlass.range_constexpr(outs):
+ acc[r, j] = acc[r, j] + av * w_reg[j, kt, v].to(cutlass.Float32)
+ if const_expr(souts > 0):
+ if const_expr(r == 0):
+ cute.arch.cp_async_wait_group(0)
+ for j in cutlass.range_constexpr(souts):
+ for kt in cutlass.range_constexpr(k_tiles):
+ cute.autovec_copy(s_w[j, kt, tidx, None], sw_reg)
+ for v in cutlass.range_constexpr(vec):
+ acc[r, outs + j] = acc[r, outs + j] + a_reg[kt, v].to(
+ cutlass.Float32
+ ) * sw_reg[v].to(cutlass.Float32)
+ for r in cutlass.range_constexpr(rows):
+ for j in cutlass.range_constexpr(tot):
+ acc[r, j] = cute.arch.warp_reduction_sum(acc[r, j])
+ part = smem.allocate_tensor(
+ cutlass.Float32,
+ cute.make_layout((rows, tot, warps), stride=(tot * warps, warps, 1)),
+ byte_alignment=16,
+ )
+ with cute.arch.elect_one():
+ for r in cutlass.range_constexpr(rows):
+ for j in cutlass.range_constexpr(tot):
+ part[r, j, warp] = acc[r, j]
+ cute.arch.sync_threads()
+ if tidx < tot:
+ for r in cutlass.range_constexpr(rows):
+ total = cutlass.Float32(0.0)
+ for w in cutlass.range_constexpr(warps):
+ total = total + part[r, tidx, w]
+ mOut[r, n0 + tidx] = total.to(cutlass.BFloat16)
+
+
+@functools.lru_cache(maxsize=None)
+def _compiled(rows, n, k, pdl):
+ threads, outs, souts, vec = tiling(n, k)
+ kern = _GemvKernel(
+ rows=rows,
+ k=k,
+ threads=threads,
+ outs=outs,
+ vec=vec,
+ pdl=pdl,
+ dist=l2_dist(-(-n // (outs + souts))),
+ souts=souts,
+ )
+ a = torch.empty(rows, k, dtype=torch.bfloat16, device="cuda")
+ w = torch.empty(n, k, dtype=torch.bfloat16, device="cuda")
+ out = torch.empty(rows, n, dtype=torch.bfloat16, device="cuda")
+ stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream)
+ return cute_ext.compile(
+ kern,
+ from_dlpack(a, assumed_align=32),
+ from_dlpack(w, assumed_align=32),
+ from_dlpack(out, assumed_align=32),
+ stream,
+ options=_COMPILE_OPTIONS_SMEM if souts else _COMPILE_OPTIONS,
+ )
+
+
+def launch(mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool) -> torch.Tensor:
+ """The launch behind `Bf16GemvPdl.__call__`. `pdl` is honoured unless
+ `TRTLLM_ENABLE_PDL=0`."""
+ m, k = mat_a.shape
+ n = mat_b.shape[1]
+ dev = mat_a.device
+ pdl = pdl and get_env_enable_pdl()
+ out = torch.empty((m, n), dtype=mat_a.dtype, device=dev)
+ if m > 0:
+ fn = _compiled(m, n, k, bool(pdl))
+ stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream)
+ fn(
+ from_dlpack(mat_a, assumed_align=32),
+ from_dlpack(mat_b.t(), assumed_align=32),
+ from_dlpack(out, assumed_align=32),
+ stream,
+ )
+ return out
+
+
+def _cell(why: str, m: int, n: int, k: int, pdl: bool = True) -> Cell:
+ return Cell(why=why, spec=dict(m=m, n=n, k=k, pdl=pdl))
+
+
+class Bf16GemvPdl(OpWrapper):
+ """`mat_a @ mat_b` in bf16 with fp32 accumulation, one PDL launch per call.
+
+ A drop-in for `ops/gemm/low_m_bf16_mm` at decode row counts: `mat_a` is
+ dense row-major `[M, K]` with `M <= MAX_ROWS`, `mat_b` the dense column-major
+ `[K, N]` view of a stored `[N, K]` weight (its `.t()`), the result a fresh
+ row-major `[M, N]` bf16 tensor. With `pdl` (and PDL available) the kernel is
+ launched as a programmatic dependent: it may start before the kernel ahead of
+ it on the stream exits, reads only its weight until `griddepcontrol.wait`,
+ and lets the kernel after it launch early as well.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # Every shape is a qwen3.8-flash-next-nvfp4 decode projection the target calls it at
+ # (`_decode_projection`): one row at tp1, up to MAX_ROWS rows of a tp4/tep4 rank's o /
+ # GDN out projection (K=1536). Each `TILING` key is its own compiled kernel.
+ CELLS: tuple[Cell, ...] = (
+ _cell("qwen3.8 decode M=1 GDN out / attention o, K=6144 shared-memory tier", 1, 2560, 6144),
+ _cell("qwen3.8 decode M=1 GDN in-proj concat N=16480, L2 prefetch grid", 1, 16480, 2560),
+ _cell("qwen3.8 decode M=1 shared down K=640", 1, 2560, 640),
+ _cell("qwen3.8 decode M=1 shared gate N=1, outputs per CTA halved to 1", 1, 1, 2560),
+ _cell("tp4 o / GDN out M=4 (MTP-3 verify, 1 request), PDL on", 4, 2560, 1536),
+ _cell("tp4 o / GDN out M=1 (decode without MTP), PDL on", 1, 2560, 1536),
+ _cell("tp4 o / GDN out M=4, PDL off", 4, 2560, 1536, pdl=False),
+ )
+
+ note = """
+ Metadata consumed: none. Preconditions: bf16 operands; `mat_a` dense
+ row-major `[M, K]`, `1 <= M <= MAX_ROWS`; `mat_b` the dense column-major
+ `[K, N]` view of a `[N, K]` weight (stride `(1, K)`); K a `TILING` key.
+
+ PDL: with `pdl` the launch carries the programmatic-dependent attribute.
+ Every read of `mat_a` and the store of the output come after
+ `griddepcontrol.wait`, so a producer launched before this kernel is
+ complete when they happen, whatever that producer is. Under a CUDA graph a
+ cross-stream dependency is an ordinary edge and stays a full dependency.
+
+ Compilation: CuTe DSL, JIT on the first call per (M, N, K, pdl); that call
+ must happen before CUDA-graph capture -- the engine's warmup forward at each
+ captured batch size does it.
+
+ Numerics: fp32 products and sums, one rounding to bf16; only the summation
+ order differs from the fp32 reference.
+
+ Hard limits
+ M rows per launch boundary MAX_ROWS=4 driven at 1 (cells)
+ K a TILING key 2560, 6144, 640, 1536 all driven
+ N any; the outputs per CTA halve until they divide N N=1 driven
+
+ The test also replays a PDL chain of these GEMVs under a CUDA graph, each
+ reading the one before it.
+
+ Not certified: M = 2..4 at K other than 1536; PDL disabled in the environment (the
+ PDL-off cell is the same kernel without the attribute).
+ """
+
+ def __call__(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = True) -> torch.Tensor:
+ 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/ops/gemm/fp8_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/fp8_gemv.py
new file mode 100644
index 000000000000..ff6fa1208e36
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/fp8_gemv.py
@@ -0,0 +1,2016 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Decode-row W8A16 projection: `y[M, N] = x[M, K] @ dequant(W)^T` with FP8 e4m3 weights
+(128 x 128 block scales, fp32) and bf16 activations that are never quantized.
+
+CuTe DSL (inline PTX), one PDL launch, CUDA-graph capturable, M = 1 .. 32.
+
+Math: swap-AB warp MMA (`mma.sync.m16n8k16`, bf16 in, fp32 accumulate). The A operand is
+16 weight rows dequantized in registers with `cvt.rn.bf16x2.e4m3x2` (exact: every e4m3
+value is a bf16); B is up to 8 activation rows per n8 tile, read as raw bf16. Each 128-wide
+k-block is accumulated on its own (four independent MMA chains) and folded into the running
+fp32 sum times its block scale, so the only roundings are fp32 sums and the final bf16 store.
+
+Weight layout (`pack_weight`): 2 KB units of 16 rows x 128 k in exactly the order the MMA
+fragments read them (one conflict-free `ld.shared.v4` per lane per two k16 steps), stored
+k-block-major `[K/128][ceil(N/16)][2048]` or, for small N, tile-major
+`[ceil(N/16)][K/128][2048]`. K is permuted identically for both operands inside every
+128-block (lane quad t of k16 step j holds k = 32 (j // 2) + 8 t + 4 (j % 2) + 0..3), which
+leaves the dot product unchanged and makes every activation read a 16-byte vector.
+
+Work split (`Plan`): a CTA owns a contiguous range of 16-row tiles over its K range (all of K,
+or 1/ks of it in a `ks`-CTA cluster reduced over DSMEM). A weight-producer warp streams the
+CTA's units with `cp.async.bulk` into a ring of shared-memory slots (a slot = `kw` k-blocks of
+every tile of the CTA) completed on mbarriers; the static weight needs no PDL wait, so the ring
+fills while the kernel ahead finishes. Consumer warp (kw_i, tw_i) takes k-block kw_i of a slot
+for tiles tw_i, tw_i + tw, .... Activations are read only after `griddepcontrol.wait`, either
+straight into registers a few stages ahead (few stages per warp) or by an activation-producer
+warp with 16-byte `cp.async` into the same ring slots (long K). The `kw` k-block partials are
+summed in shared memory in a fixed order and stored as bf16: results are deterministic.
+"""
+
+from __future__ import annotations
+
+import functools
+from dataclasses import dataclass
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass import const_expr
+from cutlass._mlir.dialects import llvm
+from cutlass.cute.runtime import from_dlpack, make_ptr
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import (
+ Arch,
+ Cell,
+ OpWrapper,
+ true_fp32_matmul,
+)
+
+TILE_N = 16
+BLOCK_K = 128
+UNIT = TILE_N * BLOCK_K # bytes of one (16-row, 128-k) unit
+#: Bytes of one NVFP4 (16-row, 128-k) unit: 1024 of e2m1 codes, then 128 of e4m3 scales.
+UNIT_W4 = TILE_N * BLOCK_K // 2 + TILE_N * BLOCK_K // 16
+FP8_MAX = 448.0
+MAX_ROWS = 32
+
+
+# ------------------------------------------------------------------------------------------
+# inline PTX helpers
+# ------------------------------------------------------------------------------------------
+
+
+def _ir(x, cls, loc, ip):
+ return cls(x).ir_value(loc=loc, ip=ip)
+
+
+def _asm(ret, args, text, cons, loc, ip, side=False):
+ return llvm.inline_asm(
+ ret,
+ args,
+ text,
+ cons,
+ has_side_effects=side,
+ is_align_stack=False,
+ asm_dialect=llvm.AsmDialect.AD_ATT,
+ loc=loc,
+ ip=ip,
+ )
+
+
+def _unpack(res, n, ty, loc, ip):
+ return [
+ ty(llvm.extractvalue(T.i32() if ty is cutlass.Int32 else T.f32(), res, [i], loc=loc, ip=ip))
+ for i in range(n)
+ ]
+
+
+@dsl_user_op
+def _lds4(saddr, *, loc=None, ip=None):
+ res = _asm(
+ llvm.StructType.get_literal([T.i32()] * 4),
+ [_ir(saddr, cutlass.Int32, loc, ip)],
+ "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];",
+ "=r,=r,=r,=r,r",
+ loc,
+ ip,
+ side=True,
+ )
+ return _unpack(res, 4, cutlass.Int32, loc, ip)
+
+
+@dsl_user_op
+def _lds2(saddr, *, loc=None, ip=None):
+ res = _asm(
+ llvm.StructType.get_literal([T.i32()] * 2),
+ [_ir(saddr, cutlass.Int32, loc, ip)],
+ "ld.shared.v2.b32 {$0, $1}, [$2];",
+ "=r,=r,r",
+ loc,
+ ip,
+ side=True,
+ )
+ return _unpack(res, 2, cutlass.Int32, loc, ip)
+
+
+@dsl_user_op
+def _lds_f32(saddr, *, loc=None, ip=None):
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [_ir(saddr, cutlass.Int32, loc, ip)],
+ "ld.shared.f32 $0, [$1];",
+ "=f,r",
+ loc,
+ ip,
+ side=True,
+ )
+ )
+
+
+@dsl_user_op
+def _sts_f32(saddr, v, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [_ir(saddr, cutlass.Int32, loc, ip), _ir(v, cutlass.Float32, loc, ip)],
+ "{ st.shared.f32 [$1], $2; mov.u32 $0, 0; }",
+ "=r,r,f",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _ldg4p(gaddr, pred, *, loc=None, ip=None):
+ """16 bytes from global (L1-cached, read-only) if `pred` != 0, else zeros."""
+ res = _asm(
+ llvm.StructType.get_literal([T.i32()] * 4),
+ [_ir(gaddr, cutlass.Int64, loc, ip), _ir(pred, cutlass.Int32, loc, ip)],
+ "{ .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.nc.v4.b32 {$0, $1, $2, $3}, [$4]; }",
+ "=r,=r,=r,=r,l,r",
+ loc,
+ ip,
+ side=True,
+ )
+ return _unpack(res, 4, cutlass.Int32, loc, ip)
+
+
+@dsl_user_op
+def _ldg_f32(gaddr, *, loc=None, ip=None):
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [_ir(gaddr, cutlass.Int64, loc, ip)],
+ "ld.global.nc.f32 $0, [$1];",
+ "=f,l",
+ loc,
+ ip,
+ side=True,
+ )
+ )
+
+
+@dsl_user_op
+def _stg_bf16(gaddr, v, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [_ir(gaddr, cutlass.Int64, loc, ip), _ir(v, cutlass.Float32, loc, ip)],
+ "{ .reg .b16 h; cvt.rn.bf16.f32 h, $2; st.global.b16 [$1], h; mov.u32 $0, 0; }",
+ "=r,l,f",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _ldg_bf16p(gaddr, pred, *, loc=None, ip=None):
+ """One bf16 from global as fp32 if `pred` != 0, else 0."""
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [_ir(gaddr, cutlass.Int64, loc, ip), _ir(pred, cutlass.Int32, loc, ip)],
+ "{ .reg .pred p; .reg .b16 h; setp.ne.b32 p, $2, 0; mov.b16 h, 0; @p ld.global.b16 h, [$1]; "
+ "cvt.f32.bf16 $0, h; }",
+ "=f,l,r",
+ loc,
+ ip,
+ side=True,
+ )
+ )
+
+
+@dsl_user_op
+def _ldg_s32p(gaddr, pred, dflt, *, loc=None, ip=None):
+ """One int32 from global if `pred` != 0, else `dflt`."""
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [
+ _ir(gaddr, cutlass.Int64, loc, ip),
+ _ir(pred, cutlass.Int32, loc, ip),
+ _ir(dflt, cutlass.Int32, loc, ip),
+ ],
+ "{ .reg .pred p; setp.ne.b32 p, $2, 0; mov.b32 $0, $3; @p ld.global.b32 $0, [$1]; }",
+ "=r,l,r,r",
+ loc,
+ ip,
+ side=True,
+ )
+ )
+
+
+@dsl_user_op
+def _ldg_u16p(gaddr, pred, *, loc=None, ip=None):
+ """One 16-bit value from global, zero-extended, if `pred` != 0, else 0. Nothing in the asm
+ reads it, so the thread runs on until the value's first use."""
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [_ir(gaddr, cutlass.Int64, loc, ip), _ir(pred, cutlass.Int32, loc, ip)],
+ "{ .reg .pred p; setp.ne.b32 p, $2, 0; mov.b32 $0, 0; @p ld.global.u16 $0, [$1]; }",
+ "=r,l,r",
+ loc,
+ ip,
+ side=True,
+ )
+ )
+
+
+@dsl_user_op
+def _ldg_v2p(gaddr, pred, *, loc=None, ip=None):
+ """8 bytes from global (read-only) if `pred` != 0, else zeros, as two words."""
+ res = _asm(
+ llvm.StructType.get_literal([T.i32()] * 2),
+ [_ir(gaddr, cutlass.Int64, loc, ip), _ir(pred, cutlass.Int32, loc, ip)],
+ "{ .reg .pred p; setp.ne.b32 p, $3, 0; mov.b32 $0, 0; mov.b32 $1, 0; "
+ "@p ld.global.nc.v2.b32 {$0, $1}, [$2]; }",
+ "=r,=r,l,r",
+ loc,
+ ip,
+ side=True,
+ )
+ return _unpack(res, 2, cutlass.Int32, loc, ip)
+
+
+@dsl_user_op
+def _bf16_lo(v, *, loc=None, ip=None):
+ """The bf16 in the low half of `v` as fp32."""
+ return cutlass.Float32(
+ _asm(T.f32(), [_ir(v, cutlass.Int32, loc, ip)], "shl.b32 $0, $1, 16;", "=f,r", loc, ip)
+ )
+
+
+@dsl_user_op
+def _round_bf16(x, *, loc=None, ip=None):
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [_ir(x, cutlass.Float32, loc, ip)],
+ "{ .reg .b16 h; cvt.rn.bf16.f32 h, $1; cvt.f32.bf16 $0, h; }",
+ "=f,f",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _conv4_silu_bf16(c0, c1, c2, x, w01, w23, *, loc=None, ip=None):
+ """`causal_conv1d_triton_update`'s arithmetic for one output of a width-4 window: the
+ window `(c0, c1, c2, x)` (bf16 bits in the low half of each word) times the taps (bf16x2
+ words `(w0 | w1 << 16, w2 | w3 << 16)`), each product rounded to bf16 (`mul.rn.bf16x2`),
+ the four summed in fp32 in tap order from +0, then SiLU as
+ `acc / (1 + ex2.approx(acc * -log2 e))` with `div.full.f32` -- the instruction sequence
+ the Triton kernel compiles to. Returns the fp32 result (rounded to bf16 by the store)."""
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [_ir(v, cutlass.Int32, loc, ip) for v in (c0, c1, c2, x, w01, w23)],
+ "{ .reg .b32 u, v; .reg .b16 p0, p1, p2, p3; .reg .f32 s, q, e; "
+ "prmt.b32 u, $1, $2, 0x5410; prmt.b32 v, $3, $4, 0x5410; "
+ "mul.rn.bf16x2 u, u, $5; mul.rn.bf16x2 v, v, $6; "
+ "mov.b32 {p0, p1}, u; mov.b32 {p2, p3}, v; "
+ "cvt.f32.bf16 s, p0; add.rn.f32 s, s, 0f00000000; "
+ "cvt.f32.bf16 q, p1; add.rn.f32 s, q, s; "
+ "cvt.f32.bf16 q, p2; add.rn.f32 s, q, s; "
+ "cvt.f32.bf16 q, p3; add.rn.f32 s, q, s; "
+ "mul.rn.f32 e, s, 0fBFB8AA3B; ex2.approx.f32 e, e; add.rn.f32 e, e, 0f3F800000; "
+ "div.full.f32 $0, s, e; }",
+ "=f,r,r,r,r,r,r",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _bf16_bits(x, *, loc=None, ip=None):
+ """`x` rounded to bf16, as its bits in the low half of a word."""
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [_ir(x, cutlass.Float32, loc, ip)],
+ "{ .reg .b16 h; cvt.rn.bf16.f32 h, $1; cvt.u32.u16 $0, h; }",
+ "=r,f",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _stg_b16(gaddr, v, *, loc=None, ip=None):
+ """The low half of `v` to global."""
+ _asm(
+ T.i32(),
+ [_ir(gaddr, cutlass.Int64, loc, ip), _ir(v, cutlass.Int32, loc, ip)],
+ "{ st.global.b16 [$1], $2; mov.u32 $0, 0; }",
+ "=r,l,r",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _dq_chunk(w0, w1, w2, w3, *, loc=None, ip=None):
+ """Four words of e4m3 (rows g / g+8 of k16 step 2c, then of step 2c+1) -> the eight bf16x2
+ A registers of the two steps: (R0, R1, R2, R3) = (g lo, g+8 lo, g hi, g+8 hi)."""
+ res = _asm(
+ llvm.StructType.get_literal([T.i32()] * 8),
+ [_ir(v, cutlass.Int32, loc, ip) for v in (w0, w1, w2, w3)],
+ "{ .reg .b16 a0, a1, b0, b1, c0, c1, d0, d1; "
+ "mov.b32 {a0, a1}, $8; mov.b32 {b0, b1}, $9; mov.b32 {c0, c1}, $10; mov.b32 {d0, d1}, $11; "
+ "cvt.rn.bf16x2.e4m3x2 $0, a0; cvt.rn.bf16x2.e4m3x2 $1, b0; "
+ "cvt.rn.bf16x2.e4m3x2 $2, a1; cvt.rn.bf16x2.e4m3x2 $3, b1; "
+ "cvt.rn.bf16x2.e4m3x2 $4, c0; cvt.rn.bf16x2.e4m3x2 $5, d0; "
+ "cvt.rn.bf16x2.e4m3x2 $6, c1; cvt.rn.bf16x2.e4m3x2 $7, d1; }",
+ "=r,=r,=r,=r,=r,=r,=r,=r,r,r,r,r",
+ loc,
+ ip,
+ )
+ return _unpack(res, 8, cutlass.Int32, loc, ip)
+
+
+@dsl_user_op
+def _mma(a, b, c, *, loc=None, ip=None):
+ """D = A B + C on m16n8k16, bf16 A/B, fp32 C/D; a: 4 regs, b: 2 regs, c: 4 f32."""
+ res = _asm(
+ llvm.StructType.get_literal([T.f32()] * 4),
+ [_ir(v, cutlass.Int32, loc, ip) for v in a]
+ + [_ir(v, cutlass.Int32, loc, ip) for v in b]
+ + [_ir(v, cutlass.Float32, loc, ip) for v in c],
+ "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+ "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};",
+ "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
+ loc,
+ ip,
+ )
+ return _unpack(res, 4, cutlass.Float32, loc, ip)
+
+
+@dsl_user_op
+def _scales4(sw, *, loc=None, ip=None):
+ """A unit's scale word `(s_g(2t), s_g8(2t), s_g(2t+1), s_g8(2t+1))` (e4m3) -> four f16x2
+ registers with each scale in both halves: (row g, steps 0..3), (row g + 8, steps 0..3),
+ (row g, steps 4..7), (row g + 8, steps 4..7)."""
+ res = _asm(
+ llvm.StructType.get_literal([T.i32()] * 4),
+ [_ir(sw, cutlass.Int32, loc, ip)],
+ "{ .reg .b16 lo, hi, p0, p1, q0, q1; .reg .b32 p, q; "
+ "mov.b32 {lo, hi}, $4; "
+ "cvt.rn.f16x2.e4m3x2 p, lo; cvt.rn.f16x2.e4m3x2 q, hi; "
+ "mov.b32 {p0, p1}, p; mov.b32 {q0, q1}, q; "
+ "mov.b32 $0, {p0, p0}; mov.b32 $1, {p1, p1}; mov.b32 $2, {q0, q0}; "
+ "mov.b32 $3, {q1, q1}; }",
+ "=r,=r,=r,=r,r",
+ loc,
+ ip,
+ )
+ return _unpack(res, 4, cutlass.Int32, loc, ip)
+
+
+@dsl_user_op
+def _dq_step(w, sg, sg8, *, loc=None, ip=None):
+ """One step's code word (bytes: row g k0|k1, row g+8 k0|k1, row g k2|k3, row g+8 k2|k3)
+ -> the four f16x2 A registers, each times its row's block scale (exact)."""
+ res = _asm(
+ llvm.StructType.get_literal([T.i32()] * 4),
+ [_ir(v, cutlass.Int32, loc, ip) for v in (w, sg, sg8)],
+ "{ .reg .b8 b0, b1, b2, b3; .reg .b32 h0, h1, h2, h3; "
+ "mov.b32 {b0, b1, b2, b3}, $4; "
+ "cvt.rn.f16x2.e2m1x2 h0, b0; cvt.rn.f16x2.e2m1x2 h1, b1; "
+ "cvt.rn.f16x2.e2m1x2 h2, b2; cvt.rn.f16x2.e2m1x2 h3, b3; "
+ "mul.rn.f16x2 $0, h0, $5; mul.rn.f16x2 $1, h1, $6; "
+ "mul.rn.f16x2 $2, h2, $5; mul.rn.f16x2 $3, h3, $6; }",
+ "=r,=r,=r,=r,r,r,r",
+ loc,
+ ip,
+ )
+ return _unpack(res, 4, cutlass.Int32, loc, ip)
+
+
+@dsl_user_op
+def _bf16x2_to_f16x2(v, *, loc=None, ip=None):
+ """bf16x2 -> f16x2 (round to nearest, saturating at +-65504)."""
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [_ir(v, cutlass.Int32, loc, ip)],
+ "{ .reg .b32 lo, hi; .reg .f32 flo, fhi; "
+ "shl.b32 lo, $1, 16; and.b32 hi, $1, 0xFFFF0000; "
+ "mov.b32 flo, lo; mov.b32 fhi, hi; "
+ "cvt.rn.satfinite.f16x2.f32 $0, fhi, flo; }",
+ "=r,r",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _mma_f16(a, b, c, *, loc=None, ip=None):
+ """D = A B + C on m16n8k16, f16 A/B, fp32 C/D; a: 4 regs, b: 2 regs, c: 4 f32."""
+ res = _asm(
+ llvm.StructType.get_literal([T.f32()] * 4),
+ [_ir(v, cutlass.Int32, loc, ip) for v in a]
+ + [_ir(v, cutlass.Int32, loc, ip) for v in b]
+ + [_ir(v, cutlass.Float32, loc, ip) for v in c],
+ "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+ "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};",
+ "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
+ loc,
+ ip,
+ )
+ return _unpack(res, 4, cutlass.Float32, loc, ip)
+
+
+@dsl_user_op
+def _mbar_init(addr, count, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [_ir(addr, cutlass.Int32, loc, ip), _ir(count, cutlass.Int32, loc, ip)],
+ "{ mbarrier.init.shared::cta.b64 [$1], $2; mov.u32 $0, 0; }",
+ "=r,r,r",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _mbar_init_fence(*, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [],
+ "{ fence.mbarrier_init.release.cluster; mov.u32 $0, 0; }",
+ "=r",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _mbar_expect_tx(addr, nbytes, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [_ir(addr, cutlass.Int32, loc, ip), _ir(nbytes, cutlass.Int32, loc, ip)],
+ "{ mbarrier.arrive.expect_tx.shared::cta.b64 _, [$1], $2; mov.u32 $0, 0; }",
+ "=r,r,r",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _mbar_arrive(addr, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [_ir(addr, cutlass.Int32, loc, ip)],
+ "{ mbarrier.arrive.shared::cta.b64 _, [$1]; mov.u32 $0, 0; }",
+ "=r,r",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _mbar_wait(addr, parity, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [_ir(addr, cutlass.Int32, loc, ip), _ir(parity, cutlass.Int32, loc, ip)],
+ "{ .reg .pred p; WAIT: mbarrier.try_wait.parity.shared::cta.b64 p, [$1], $2; @!p bra WAIT; "
+ "mov.u32 $0, 0; }",
+ "=r,r,r",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _bulk_g2s(saddr, gaddr, nbytes, mbar, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [
+ _ir(saddr, cutlass.Int32, loc, ip),
+ _ir(gaddr, cutlass.Int64, loc, ip),
+ _ir(nbytes, cutlass.Int32, loc, ip),
+ _ir(mbar, cutlass.Int32, loc, ip),
+ ],
+ "{ cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [$1], [$2], $3, [$4]; "
+ "mov.u32 $0, 0; }",
+ "=r,r,l,r,r",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _cp16(saddr, gaddr, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [_ir(saddr, cutlass.Int32, loc, ip), _ir(gaddr, cutlass.Int64, loc, ip)],
+ "{ cp.async.cg.shared.global [$1], [$2], 16; mov.u32 $0, 0; }",
+ "=r,r,l",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _cp_async_mbar_arrive(mbar, *, loc=None, ip=None):
+ """Arrive on `mbar` once this thread's earlier cp.async copies have landed."""
+ _asm(
+ T.i32(),
+ [_ir(mbar, cutlass.Int32, loc, ip)],
+ "{ cp.async.mbarrier.arrive.noinc.shared::cta.b64 [$1]; mov.u32 $0, 0; }",
+ "=r,r",
+ loc,
+ ip,
+ side=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(),
+ [_ir(addr, cutlass.Int32, loc, ip), _ir(rank, cutlass.Int32, loc, ip)],
+ "mapa.shared::cluster.u32 $0, $1, $2;",
+ "=r,r,r",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _st_dsmem_f32(addr, v, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [_ir(addr, cutlass.Int32, loc, ip), _ir(v, cutlass.Float32, loc, ip)],
+ "{ st.shared::cluster.f32 [$1], $2; mov.u32 $0, 0; }",
+ "=r,r,f",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _cluster_arrive_release(*, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [],
+ "{ barrier.cluster.arrive.release.aligned; mov.u32 $0, 0; }",
+ "=r",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _cluster_arrive_relaxed(*, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [],
+ "{ barrier.cluster.arrive.relaxed.aligned; mov.u32 $0, 0; }",
+ "=r",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _cluster_wait(*, loc=None, ip=None):
+ _asm(T.i32(), [], "{ barrier.cluster.wait.aligned; mov.u32 $0, 0; }", "=r", loc, ip, side=True)
+
+
+@dsl_user_op
+def _cluster_wait_acquire(*, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [],
+ "{ barrier.cluster.wait.acquire.aligned; mov.u32 $0, 0; }",
+ "=r",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _bulk_prefetch_l2(gaddr, nbytes, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [_ir(gaddr, cutlass.Int64, loc, ip), _ir(nbytes, cutlass.Int32, loc, ip)],
+ "{ cp.async.bulk.prefetch.L2.global [$1], $2; mov.u32 $0, 0; }",
+ "=r,l,r",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _bar_sync(bar_id, nthreads, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [_ir(bar_id, cutlass.Int32, loc, ip), _ir(nthreads, cutlass.Int32, loc, ip)],
+ "{ bar.sync $1, $2; mov.u32 $0, 0; }",
+ "=r,r,r",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+# ------------------------------------------------------------------------------------------
+# kernel
+# ------------------------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class Plan:
+ """Launch shape. `grid` CTAs, each a contiguous range of <= ceil(tiles / grid) 16-row
+ tiles; `kw` x `tw` consumer warps (k-blocks per ring slot x tile groups); `slots` ring
+ slots; `l2_ahead` slots past the ring prefetched into L2 at the start."""
+
+ grid: int
+ kw: int
+ tw: int
+ slots: int
+ l2_ahead: int = 0 # rolling window of slots past the ring prefetched into L2
+ xreg: int = 0 # activations: 0 auto, 1 registers, 2 through the ring
+ xdepth: int = 0 # register mode: stages of activations loaded ahead (0: auto)
+ ks: int = 1 # cluster K split: `ks` CTAs share a tile range, each 1/ks of K, reduced over DSMEM
+
+
+#: Most activation-fragment registers a consumer may hold to take the register path.
+XREG_MAX = 64
+XRS = 320 # bytes of one activation row slice (128 bf16 + 64 B pad: conflict-free fragment reads)
+
+
+class _Fp8Gemv:
+ def __init__(
+ self, *, n, k, m, plan: Plan, pdl, tile_major, conv_dim=0, pad_slot=-1, conv_p=0, w4=False
+ ):
+ assert k % BLOCK_K == 0 and 1 <= m <= MAX_ROWS
+ self.n, self.k, self.m, self.pdl = n, k, m, pdl
+ # NVFP4 units (`w4`, see `pack_weight_w4`) in place of FP8 ones: f16 MMA over the
+ # activation rounded to f16, the global scale applied to the fp32 partials.
+ self.w4 = w4
+ self.unit = UNIT_W4 if w4 else UNIT
+ self.tile_major = tile_major
+ # conv1d epilogue (`conv_dim` > 0): columns below conv_dim of a non-pad row get the
+ # causal conv step (window of 3, 4 taps) + SiLU; see `fp8_gemv_conv1d`. With
+ # `conv_p` > 1 the rows are m / conv_p requests of conv_p consecutive tokens each, and
+ # the epilogue runs the multi-token step with per-token window snapshots; see
+ # `fp8_gemv_conv1d_verify`.
+ self.conv_dim, self.pad_slot, self.conv_p = conv_dim, pad_slot, conv_p
+ assert conv_dim == 0 or plan.ks == 1
+ assert conv_p == 0 or (conv_dim > 0 and conv_p >= 2 and m % conv_p == 0), (conv_p, m)
+ self.nreq = m // conv_p if conv_p else 0
+ self.nt8 = -(-m // 8)
+ self.n_tiles = -(-n // TILE_N)
+ self.nkb = k // BLOCK_K
+ self.grid = plan.grid
+ self.ks = plan.ks
+ assert plan.grid % self.ks == 0 and self.nkb % self.ks == 0, (plan.grid, self.ks, self.nkb)
+ self.groups = plan.grid // self.ks
+ self.nkb_c = self.nkb // self.ks # k-blocks per CTA
+ self.rt = -(-self.n_tiles // self.groups)
+ self.kw, self.tw, self.slots, self.l2_ahead = plan.kw, plan.tw, plan.slots, plan.l2_ahead
+ assert self.nkb_c % self.kw == 0, (self.nkb_c, self.kw)
+ self.n_stages = self.nkb_c // self.kw
+ self.slots = min(self.slots, self.n_stages)
+ self.l2_ahead = max(0, min(self.l2_ahead, self.n_stages - self.slots))
+ self.rtw = -(-self.rt // self.tw)
+ self.cw = self.kw * self.tw
+ # Activations: with few stages a warp loads all of its fragments into registers right
+ # after the PDL wait (lowest latency); otherwise they stream through the ring slots.
+ self.xreg = self.n_stages * self.nt8 * 16 <= XREG_MAX if plan.xreg == 0 else plan.xreg == 1
+ # register mode: stages of activation fragments in flight ahead of the one in use
+ self.xdepth = plan.xdepth or max(1, XREG_MAX // (self.nt8 * 16) - 1)
+ self.xdepth = min(self.xdepth, self.n_stages)
+ # + weight producer warp (+ activation producer warp)
+ self.threads = (self.cw + (1 if self.xreg else 2)) * 32
+ self.w_bytes = self.kw * self.rt * self.unit
+ self.x_bytes = 0 if self.xreg else self.kw * m * XRS
+ self.slot_bytes = -(-(self.w_bytes + self.x_bytes) // 128) * 128
+ self.np = self.nt8 * 8
+ self.o_red = self.slots * self.slot_bytes
+ red_bytes = self.kw * self.rt * self.np * TILE_N * 4
+ # block scales of the (at most nsb) 128-row blocks the CTA's tiles touch
+ self.nsb = 0 if w4 else self.rt // 8 + 2
+ self.o_sc = self.o_red + -(-red_bytes // 128) * 128
+ self.o_bar = self.o_sc + -(-(self.nsb * self.nkb * 4) // 128) * 128
+ self.o_cred = self.o_bar + 3 * self.slots * 8
+ cred_bytes = self.ks * self.rt * m * TILE_N * 4 if self.ks > 1 else 0
+ self.smem = self.o_cred + cred_bytes + 128
+
+ @cute.jit
+ def __call__(
+ self,
+ out_t: cute.Tensor,
+ x_ptr: cute.Pointer,
+ w_ptr: cute.Pointer,
+ s_ptr: cute.Pointer,
+ x_stride: cutlass.Int32,
+ cs_ptr: cute.Pointer,
+ cw_ptr: cute.Pointer,
+ idx_ptr: cute.Pointer,
+ cs_s0: cutlass.Int32,
+ cs_s1: cutlass.Int32,
+ cs_s2: cutlass.Int32,
+ iw_ptr: cute.Pointer,
+ iidx_ptr: cute.Pointer,
+ iw_s0: cutlass.Int32,
+ iw_s1: cutlass.Int32,
+ iw_s2: cutlass.Int32,
+ iw_s3: cutlass.Int32,
+ stream,
+ ):
+ args = (
+ out_t.iterator.toint(),
+ x_ptr.toint(),
+ w_ptr.toint(),
+ s_ptr.toint(),
+ x_stride,
+ cs_ptr.toint(),
+ cw_ptr.toint(),
+ idx_ptr.toint(),
+ cs_s0,
+ cs_s1,
+ cs_s2,
+ iw_ptr.toint(),
+ iidx_ptr.toint(),
+ iw_s0,
+ iw_s1,
+ iw_s2,
+ iw_s3,
+ )
+ if const_expr(self.ks > 1):
+ self.kernel(*args).launch(
+ grid=[self.grid, 1, 1],
+ block=[self.threads, 1, 1],
+ smem=self.smem,
+ stream=stream,
+ cluster=[self.ks, 1, 1],
+ use_pdl=self.pdl,
+ )
+ else:
+ self.kernel(*args).launch(
+ grid=[self.grid, 1, 1],
+ block=[self.threads, 1, 1],
+ smem=self.smem,
+ stream=stream,
+ use_pdl=self.pdl,
+ )
+
+ @cute.jit
+ def _issue_w(self, dst, w_base, kbs, t0, nt, mbar):
+ """Bulk copies of a stage's units (k-blocks kbs .. kbs + kw, tiles t0 .. t0 + nt)."""
+ kw: cutlass.Constexpr = self.kw
+ rt: cutlass.Constexpr = self.rt
+ unit: cutlass.Constexpr = self.unit
+ if const_expr(self.tile_major):
+ for tl in cutlass.range_constexpr(rt):
+ if tl < nt:
+ _bulk_g2s(
+ dst + tl * (kw * unit),
+ w_base + (cutlass.Int64(t0 + tl) * self.nkb + kbs) * unit,
+ kw * unit,
+ mbar,
+ )
+ else:
+ for j in cutlass.range_constexpr(kw):
+ _bulk_g2s(
+ dst + j * (rt * unit),
+ w_base + (cutlass.Int64(kbs + j) * self.n_tiles + t0) * unit,
+ nt * unit,
+ mbar,
+ )
+
+ @cute.jit
+ def _prefetch_w(self, w_base, kbs, t0, nt):
+ kw: cutlass.Constexpr = self.kw
+ rt: cutlass.Constexpr = self.rt
+ unit: cutlass.Constexpr = self.unit
+ if const_expr(self.tile_major):
+ for tl in cutlass.range_constexpr(rt):
+ if tl < nt:
+ _bulk_prefetch_l2(
+ w_base + (cutlass.Int64(t0 + tl) * self.nkb + kbs) * unit, kw * unit
+ )
+ else:
+ for j in cutlass.range_constexpr(kw):
+ _bulk_prefetch_l2(
+ w_base + (cutlass.Int64(kbs + j) * self.n_tiles + t0) * unit, nt * unit
+ )
+
+ @cute.jit
+ def _compute(self, acc, cur, slot, kb, kwi, twi, t0, nt, b0, scs, lane):
+ """This warp's tiles of k-block `kb` (in ring slot `slot`) times the activation fragments
+ `cur`, folded into `acc` with the block scale."""
+ kw: cutlass.Constexpr = self.kw
+ tw: cutlass.Constexpr = self.tw
+ rt: cutlass.Constexpr = self.rt
+ nt8: cutlass.Constexpr = self.nt8
+ nkb: cutlass.Constexpr = self.nkb
+ unit: cutlass.Constexpr = self.unit
+ for i in cutlass.range_constexpr(self.rtw):
+ tl = twi + i * tw
+ if const_expr(self.w4):
+ if tl < nt:
+ self._compute_w4(acc, cur, slot, i, tl, kwi, lane)
+ elif tl < nt:
+ sc = _lds_f32(scs + (((t0 + tl) // 8 - b0) * nkb + kb) * 4)
+ if const_expr(self.tile_major):
+ ua = slot + (tl * kw + kwi) * unit + lane * 16
+ else:
+ ua = slot + (kwi * rt + tl) * unit + lane * 16
+ wv = [_lds4(ua + c * 512) for c in range(4)]
+ # one accumulator chain per 16-byte chunk (two k16 steps): 4-way ILP
+ zero = cutlass.Float32(0.0)
+ ch = [[[zero] * 4 for _ in range(nt8)] for _ in range(4)]
+ for c in cutlass.range_constexpr(4):
+ a = _dq_chunk(wv[c][0], wv[c][1], wv[c][2], wv[c][3])
+ for st in cutlass.range_constexpr(2):
+ for j in cutlass.range_constexpr(nt8):
+ ch[c][j] = _mma(
+ a[4 * st : 4 * st + 4],
+ [cur[j][4 * c + 2 * st], cur[j][4 * c + 2 * st + 1]],
+ ch[c][j],
+ )
+ for j in cutlass.range_constexpr(nt8):
+ for e in cutlass.range_constexpr(4):
+ part = (ch[0][j][e] + ch[1][j][e]) + (ch[2][j][e] + ch[3][j][e])
+ acc[i, j, e] = acc[i, j, e] + part * sc
+
+ @cute.jit
+ def _compute_w4(self, acc, cur, slot, i, tl, kwi, lane):
+ """`_compute` of one tile over an NVFP4 unit (`pack_weight_w4`): each step's e2m1 codes
+ times their 16-k block scales in f16 (exact), four f16 MMA chains, `cur` f16."""
+ kw: cutlass.Constexpr = self.kw
+ rt: cutlass.Constexpr = self.rt
+ nt8: cutlass.Constexpr = self.nt8
+ unit: cutlass.Constexpr = self.unit
+ if const_expr(self.tile_major):
+ ua = slot + (tl * kw + kwi) * unit
+ else:
+ ua = slot + (kwi * rt + tl) * unit
+ wv = [_lds4(ua + h * 512 + lane * 16) for h in range(2)]
+ sw = _lds2(ua + 1024 + ((lane // 4) * 2 + (lane % 4) // 2) * 8)
+ # (row g, row g + 8) block scales of chunks 0..3, as f16x2 pairs
+ sc = _scales4(sw[0]) + _scales4(sw[1])
+ zero = cutlass.Float32(0.0)
+ ch = [[[zero] * 4 for _ in range(nt8)] for _ in range(4)]
+ for c in cutlass.range_constexpr(4):
+ for st in cutlass.range_constexpr(2):
+ a = _dq_step(wv[c // 2][2 * (c % 2) + st], sc[2 * c], sc[2 * c + 1])
+ for j in cutlass.range_constexpr(nt8):
+ ch[c][j] = _mma_f16(
+ a, [cur[j][4 * c + 2 * st], cur[j][4 * c + 2 * st + 1]], ch[c][j]
+ )
+ for j in cutlass.range_constexpr(nt8):
+ for e in cutlass.range_constexpr(4):
+ part = (ch[0][j][e] + ch[1][j][e]) + (ch[2][j][e] + ch[3][j][e])
+ acc[i, j, e] = acc[i, j, e] + part
+
+ @cute.jit
+ def _issue_x(self, xs, src, xh, xc, x_stride):
+ """cp.async of one stage's activations into the slot's activation area `xs`."""
+ m: cutlass.Constexpr = self.m
+ n_rc: cutlass.Constexpr = self.kw * self.m
+ for i in cutlass.range_constexpr(-(-n_rc // 2)):
+ rc = xh + 2 * i
+ xj = rc // m
+ xr = rc - xj * m
+ ok = rc < n_rc
+ if const_expr(n_rc % 2 == 0):
+ ok = rc >= 0
+ if ok:
+ _cp16(
+ xs + rc * XRS + xc * 16,
+ src + cutlass.Int64(xr) * cutlass.Int64(x_stride) * 2 + cutlass.Int64(xj * 256),
+ )
+
+ @cute.kernel
+ def kernel(
+ self,
+ out_base: cutlass.Int64,
+ x_base: cutlass.Int64,
+ w_base: cutlass.Int64,
+ s_base: cutlass.Int64,
+ x_stride: cutlass.Int32,
+ cs_base: cutlass.Int64,
+ cw_base: cutlass.Int64,
+ idx_base: cutlass.Int64,
+ cs_s0: cutlass.Int32,
+ cs_s1: cutlass.Int32,
+ cs_s2: cutlass.Int32,
+ iw_base: cutlass.Int64,
+ iidx_base: cutlass.Int64,
+ iw_s0: cutlass.Int32,
+ iw_s1: cutlass.Int32,
+ iw_s2: cutlass.Int32,
+ iw_s3: cutlass.Int32,
+ ):
+ tidx, _, _ = cute.arch.thread_idx()
+ bidx, _, _ = cute.arch.block_idx()
+ warp = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+ lane = cute.arch.lane_idx()
+ g = lane // 4
+ t = lane % 4
+ m: cutlass.Constexpr = self.m
+ n_tiles: cutlass.Constexpr = self.n_tiles
+ nkb: cutlass.Constexpr = self.nkb
+ rt: cutlass.Constexpr = self.rt
+ kw: cutlass.Constexpr = self.kw
+ tw: cutlass.Constexpr = self.tw
+ cw: cutlass.Constexpr = self.cw
+ slots: cutlass.Constexpr = self.slots
+ n_stages: cutlass.Constexpr = self.n_stages
+ slot_bytes: cutlass.Constexpr = self.slot_bytes
+ w_bytes: cutlass.Constexpr = self.w_bytes
+ nt8: cutlass.Constexpr = self.nt8
+ np_: cutlass.Constexpr = self.np
+ rtw: cutlass.Constexpr = self.rtw
+ n_out: cutlass.Constexpr = self.n
+ l2a: cutlass.Constexpr = self.l2_ahead
+
+ smem = cutlass.utils.SmemAllocator()
+ base_t = smem.allocate_tensor(
+ cutlass.Int8, cute.make_layout((self.smem - 128,), stride=(1,)), byte_alignment=128
+ )
+ sbase = cutlass.Int32(base_t.iterator.toint())
+ ring = sbase
+ red = sbase + self.o_red
+ scs = sbase + self.o_sc
+ bars = sbase + self.o_bar # wfull[slots], xfull[slots], empty[slots]
+
+ ks: cutlass.Constexpr = self.ks
+ groups: cutlass.Constexpr = self.groups
+ grp = bidx // ks
+ rank = bidx % ks
+ kb0 = rank * self.nkb_c
+ t0 = (grp * n_tiles) // groups
+ nt = ((grp + 1) * n_tiles) // groups - t0
+ b0 = t0 // 8
+ n_blocks: cutlass.Constexpr = -(-self.n // 128)
+ n_sc: cutlass.Constexpr = self.nsb * nkb
+ if const_expr(self.w4):
+ # the NVFP4 global scale: `s_base` points at it
+ gsc = _ldg_f32(s_base)
+ for q in cutlass.range_constexpr(-(-n_sc // self.threads)):
+ qi = tidx + q * self.threads
+ if qi < n_sc:
+ blk = b0 + qi // nkb
+ if blk < n_blocks:
+ _sts_f32(
+ scs + qi * 4, _ldg_f32(s_base + (cutlass.Int64(blk) * nkb + qi % nkb) * 4)
+ )
+
+ if warp == cw:
+ if lane == 0:
+ for sl in cutlass.range_constexpr(slots):
+ _mbar_init(bars + 8 * sl, 1)
+ _mbar_init(bars + 8 * (slots + sl), 32)
+ _mbar_init(bars + 8 * (2 * slots + sl), cw)
+ _mbar_init_fence()
+ cute.arch.sync_threads()
+ if const_expr(self.ks > 1):
+ _cluster_arrive_relaxed() # matched by the wait before the first DSMEM store
+ if const_expr(self.pdl):
+ cute.arch.griddepcontrol_launch_dependents()
+
+ if warp == cw:
+ # ---- weight producer (static data: no PDL wait). Slot layout: unit (j, tile) at
+ # (j * rt + tile) * unit for a k-block-major weight, (tile * kw + j) * unit for a
+ # tile-major one; one bulk copy per k-block row resp. per tile. ----
+ if lane == 0:
+ for ps in cutlass.range(n_stages):
+ psl = ps % slots
+ if ps >= slots:
+ _mbar_wait(bars + 8 * (2 * slots + psl), ((ps // slots) - 1) & 1)
+ if const_expr(l2a > 0):
+ # rolling window: stages slots .. slots + l2a - 1 at the start, then one
+ # stage l2a past each refill
+ if ps == 0:
+ for la in cutlass.range_constexpr(slots, slots + l2a):
+ self._prefetch_w(w_base, kb0 + la * kw, t0, nt)
+ if ps >= slots:
+ if ps + l2a < n_stages:
+ self._prefetch_w(w_base, kb0 + (ps + l2a) * kw, t0, nt)
+ _mbar_expect_tx(bars + 8 * psl, kw * nt * self.unit)
+ self._issue_w(
+ ring + psl * slot_bytes, w_base, kb0 + ps * kw, t0, nt, bars + 8 * psl
+ )
+ elif warp == cw + 1:
+ # ---- activation producer (ring mode): after the PDL wait, a stage's rows into its slot ----
+ if const_expr(self.pdl):
+ cute.arch.griddepcontrol_wait()
+ xh = lane // 16
+ xc = lane % 16
+ x_lane = x_base + cutlass.Int64(xc * 16)
+ for xs_ in cutlass.range(n_stages):
+ xsl = xs_ % slots
+ if xs_ >= slots:
+ _mbar_wait(bars + 8 * (2 * slots + xsl), ((xs_ // slots) - 1) & 1)
+ self._issue_x(
+ ring + xsl * slot_bytes + w_bytes,
+ x_lane + cutlass.Int64((kb0 + xs_ * kw) * 256),
+ xh,
+ xc,
+ x_stride,
+ )
+ _cp_async_mbar_arrive(bars + 8 * (slots + xsl))
+ else:
+ # ---- consumers ----
+ kwi = warp % kw
+ twi = warp // kw
+ n_ep: cutlass.Constexpr = -(-(rt * m * TILE_N) // (cw * 32))
+ # verify epilogue items: (tile, request, row of the tile), each all conv_p tokens
+ n_ev: cutlass.Constexpr = -(-(rt * max(self.nreq, 1) * TILE_N) // (cw * 32))
+ if const_expr(self.conv_p > 0):
+ # Before the PDL wait, as below: each item's slot, snapshot row, 4 taps (two
+ # words) and the 3-value window it starts from, kept as raw bits: only the
+ # window's address waits on a load (the slot) before the main loop.
+ nreq: cutlass.Constexpr = self.nreq
+ cvs = cute.make_rmem_tensor(
+ cute.make_layout((n_ev, 8), stride=(8, 1)), cutlass.Int32
+ )
+ cslots = cute.make_rmem_tensor(
+ cute.make_layout((n_ev,), stride=(1,)), cutlass.Int32
+ )
+ crows = cute.make_rmem_tensor(cute.make_layout((n_ev,), stride=(1,)), cutlass.Int32)
+ for ic in cutlass.range_constexpr(n_ev):
+ ec = tidx + ic * (cw * 32)
+ rc_ = ec % TILE_N
+ sc_ = (ec // TILE_N) % nreq
+ tc_ = ec // (TILE_N * nreq)
+ oc_ = (t0 + tc_) * TILE_N + rc_
+ okc = (tc_ < nt) & (oc_ < self.conv_dim)
+ sl_c = _ldg_s32p(
+ idx_base + cutlass.Int64(sc_) * 4, cutlass.Int32(okc), self.pad_slot
+ )
+ cslots[ic] = sl_c
+ okw = cutlass.Int32(okc & (sl_c != self.pad_slot))
+ crows[ic] = _ldg_s32p(iidx_base + cutlass.Int64(sc_) * 4, okw, 0)
+ taps = _ldg_v2p(cw_base + cutlass.Int64(oc_) * 8, cutlass.Int32(okc))
+ cvs[ic, 0] = taps[0]
+ cvs[ic, 1] = taps[1]
+ wbase = cs_base + (cutlass.Int64(sl_c) * cs_s0 + cutlass.Int64(oc_) * cs_s1) * 2
+ for q in cutlass.range_constexpr(3):
+ cvs[ic, 4 + q] = _ldg_u16p(wbase + cutlass.Int64(q) * cs_s2 * 2, okw)
+ elif const_expr(self.conv_dim > 0):
+ # The conv operands of this thread's epilogue outputs -- slot, 4 taps and the 3-value
+ # window -- before the PDL wait: the slots are host-filled, and the window is this
+ # layer's, last written by this kernel a step ago.
+ cvs = cute.make_rmem_tensor(
+ cute.make_layout((n_ep, 8), stride=(8, 1)), cutlass.Float32
+ )
+ cslots = cute.make_rmem_tensor(
+ cute.make_layout((n_ep,), stride=(1,)), cutlass.Int32
+ )
+ for ic in cutlass.range_constexpr(n_ep):
+ ec = tidx + ic * (cw * 32)
+ rc_ = ec % TILE_N
+ nc_ = (ec // TILE_N) % m
+ tc_ = ec // (TILE_N * m)
+ oc_ = (t0 + tc_) * TILE_N + rc_
+ okc = (tc_ < nt) & (oc_ < self.conv_dim)
+ sl_c = _ldg_s32p(
+ idx_base + cutlass.Int64(nc_) * 4, cutlass.Int32(okc), self.pad_slot
+ )
+ cslots[ic] = sl_c
+ okw = cutlass.Int32(okc & (sl_c != self.pad_slot))
+ for q in cutlass.range_constexpr(4):
+ cvs[ic, q] = _ldg_bf16p(cw_base + (cutlass.Int64(oc_) * 4 + q) * 2, okw)
+ wbase = cs_base + (cutlass.Int64(sl_c) * cs_s0 + cutlass.Int64(oc_) * cs_s1) * 2
+ for q in cutlass.range_constexpr(3):
+ cvs[ic, 4 + q] = _ldg_bf16p(wbase + cutlass.Int64(q) * cs_s2 * 2, okw)
+ if const_expr(self.pdl):
+ cute.arch.griddepcontrol_wait()
+ acc = cute.make_rmem_tensor(
+ cute.make_layout((rtw, nt8, 4), stride=(nt8 * 4, 4, 1)), cutlass.Float32
+ )
+ acc.fill(0.0)
+ if const_expr(self.xreg):
+ # fragments of this warp's k-block of a stage: row 8j + g, lane quad t reads
+ # elements 32c + 8t .. + 7; loaded `xdepth` stages ahead (the stage loop is
+ # unrolled, so every buffer is its own registers)
+ xd: cutlass.Constexpr = self.xdepth
+ xsrc = [
+ x_base
+ + cutlass.Int64(8 * j + g) * cutlass.Int64(x_stride) * 2
+ + cutlass.Int64((kb0 + kwi) * (BLOCK_K * 2) + t * 16)
+ for j in range(nt8)
+ ]
+ xok = [cutlass.Int32(8 * j + g < m) for j in range(nt8)]
+ xbuf = {}
+ for s0 in cutlass.range_constexpr(xd):
+ xbuf[s0] = [
+ [
+ v
+ for c in range(4)
+ for v in _ldg4p(xsrc[j] + (s0 * kw * (BLOCK_K * 2) + 64 * c), xok[j])
+ ]
+ for j in range(nt8)
+ ]
+ for s in cutlass.range_constexpr(n_stages):
+ if const_expr(s + xd < n_stages):
+ xbuf[s + xd] = [
+ [
+ v
+ for c in range(4)
+ for v in _ldg4p(
+ xsrc[j] + ((s + xd) * kw * (BLOCK_K * 2) + 64 * c), xok[j]
+ )
+ ]
+ for j in range(nt8)
+ ]
+ sl = s % slots
+ slot = ring + sl * slot_bytes
+ _mbar_wait(bars + 8 * sl, (s // slots) & 1)
+ cur = xbuf.pop(s)
+ if const_expr(self.w4):
+ cur = [[_bf16x2_to_f16x2(v) for v in r] for r in cur]
+ self._compute(
+ acc, cur, slot, kb0 + s * kw + kwi, kwi, twi, t0, nt, b0, scs, lane
+ )
+ cute.arch.sync_warp()
+ if lane == 0:
+ _mbar_arrive(bars + 8 * (2 * slots + sl))
+ else:
+ # activation row 8j + g (clamped to m - 1: columns past m are discarded), k-block of
+ # this warp: lane quad t reads elements 32c + 8t .. + 7
+ x_rows = [cutlass.min(8 * j + g, m - 1) * XRS + t * 16 for j in range(nt8)]
+ for s in cutlass.range(n_stages):
+ sl = s % slots
+ par = (s // slots) & 1
+ slot = ring + sl * slot_bytes
+ _mbar_wait(bars + 8 * (slots + sl), par)
+ xa = slot + w_bytes + kwi * (m * XRS)
+ cur = [
+ [v for c in range(4) for v in _lds4(xa + x_rows[j] + 64 * c)]
+ for j in range(nt8)
+ ]
+ if const_expr(self.w4):
+ cur = [[_bf16x2_to_f16x2(v) for v in r] for r in cur]
+ _mbar_wait(bars + 8 * sl, par)
+ self._compute(
+ acc, cur, slot, kb0 + s * kw + kwi, kwi, twi, t0, nt, b0, scs, lane
+ )
+ cute.arch.sync_warp()
+ if lane == 0:
+ _mbar_arrive(bars + 8 * (2 * slots + sl))
+
+ # ---- k-block partials -> shared memory [kw][rt][np][16] fp32 ----
+ for i in cutlass.range_constexpr(rtw):
+ wtl = twi + i * tw
+ if wtl < nt:
+ for j in cutlass.range_constexpr(nt8):
+ n0 = 8 * j + 2 * t
+ p = red + (((kwi * rt + wtl) * np_ + n0) * TILE_N + g) * 4
+ v = [acc[i, j, e] for e in range(4)]
+ if const_expr(self.w4):
+ v = [ve * gsc for ve in v]
+ _sts_f32(p, v[0])
+ _sts_f32(p + TILE_N * 4, v[1])
+ _sts_f32(p + 8 * 4, v[2])
+ _sts_f32(p + TILE_N * 4 + 8 * 4, v[3])
+ _bar_sync(1, cw * 32)
+ if const_expr(self.conv_p > 0):
+ # Item (tile, request s, row): column `orow` of the request's rows s p .. s p + p - 1.
+ # A conv column runs the p tokens through the window in order -- output, snapshot
+ # t = the window after token t, then the advanced window into the pool; the other
+ # columns are the plain projection. Out: [m, conv_dim] conv output, then
+ # [m, n - conv_dim] projection, both row-major in the one buffer.
+ cp: cutlass.Constexpr = self.conv_p
+ cdim: cutlass.Constexpr = self.conv_dim
+ n_rest: cutlass.Constexpr = n_out - self.conv_dim
+ rest_base = out_base + cutlass.Int64(m * cdim * 2)
+ for it in cutlass.range_constexpr(n_ev):
+ e_ = tidx + it * (cw * 32)
+ er = e_ % TILE_N
+ es = (e_ // TILE_N) % nreq
+ etl = e_ // (TILE_N * nreq)
+ orow = (t0 + etl) * TILE_N + er
+ if (etl < nt) & (orow < n_out):
+ if orow < cdim:
+ live = cslots[it] != self.pad_slot
+ # bf16 bits throughout; taps as two bf16x2 words
+ c0 = cvs[it, 4]
+ c1 = cvs[it, 5]
+ c2 = cvs[it, 6]
+ iwb = (
+ iw_base
+ + (cutlass.Int64(crows[it]) * iw_s0 + cutlass.Int64(orow) * iw_s2)
+ * 2
+ )
+ for tk in cutlass.range_constexpr(cp):
+ en = es * cp + tk
+ tot = cutlass.Float32(0.0)
+ for q in cutlass.range_constexpr(kw):
+ tot = tot + _lds_f32(
+ red + (((q * rt + etl) * np_ + en) * TILE_N + er) * 4
+ )
+ xv = _bf16_bits(tot)
+ res = _bf16_lo(xv)
+ if live:
+ res = _conv4_silu_bf16(c0, c1, c2, xv, cvs[it, 0], cvs[it, 1])
+ sb = iwb + cutlass.Int64(tk) * cutlass.Int64(iw_s1) * 2
+ _stg_b16(sb, c1)
+ _stg_b16(sb + cutlass.Int64(iw_s3) * 2, c2)
+ _stg_b16(sb + cutlass.Int64(iw_s3) * 4, xv)
+ _stg_bf16(out_base + (cutlass.Int64(en) * cdim + orow) * 2, res)
+ c0 = c1
+ c1 = c2
+ c2 = xv
+ if live:
+ wb = (
+ cs_base
+ + (
+ cutlass.Int64(cslots[it]) * cs_s0
+ + cutlass.Int64(orow) * cs_s1
+ )
+ * 2
+ )
+ _stg_b16(wb, c0)
+ _stg_b16(wb + cutlass.Int64(cs_s2) * 2, c1)
+ _stg_b16(wb + cutlass.Int64(cs_s2) * 4, c2)
+ else:
+ for tk in cutlass.range_constexpr(cp):
+ en = es * cp + tk
+ tot = cutlass.Float32(0.0)
+ for q in cutlass.range_constexpr(kw):
+ tot = tot + _lds_f32(
+ red + (((q * rt + etl) * np_ + en) * TILE_N + er) * 4
+ )
+ _stg_bf16(
+ rest_base + (cutlass.Int64(en) * n_rest + orow - cdim) * 2,
+ tot,
+ )
+ elif const_expr(ks == 1):
+ for it in cutlass.range_constexpr(n_ep):
+ e_ = tidx + it * (cw * 32)
+ er = e_ % TILE_N
+ en = (e_ // TILE_N) % m
+ etl = e_ // (TILE_N * m)
+ orow = (t0 + etl) * TILE_N + er
+ if (etl < nt) & (orow < n_out):
+ tot = cutlass.Float32(0.0)
+ for q in cutlass.range_constexpr(kw):
+ tot = tot + _lds_f32(
+ red + (((q * rt + etl) * np_ + en) * TILE_N + er) * 4
+ )
+ if const_expr(self.conv_dim > 0):
+ p_ = _round_bf16(tot)
+ res = p_
+ if (orow < self.conv_dim) & (cslots[it] != self.pad_slot):
+ cv = (
+ cvs[it, 0] * cvs[it, 4]
+ + cvs[it, 1] * cvs[it, 5]
+ + cvs[it, 2] * cvs[it, 6]
+ + cvs[it, 3] * p_
+ )
+ res = cv / (1.0 + cute.math.exp(-cv))
+ wb = (
+ cs_base
+ + (
+ cutlass.Int64(cslots[it]) * cs_s0
+ + cutlass.Int64(orow) * cs_s1
+ )
+ * 2
+ )
+ _stg_bf16(wb, cvs[it, 5])
+ _stg_bf16(wb + cutlass.Int64(cs_s2) * 2, cvs[it, 6])
+ _stg_bf16(wb + cutlass.Int64(cs_s2) * 4, p_)
+ _stg_bf16(out_base + (cutlass.Int64(en) * n_out + orow) * 2, res)
+ else:
+ _stg_bf16(out_base + (cutlass.Int64(en) * n_out + orow) * 2, tot)
+ if const_expr(ks > 1):
+ # The CTA's K-share of each output goes to cluster rank (tile % ks), slot
+ # [sender][tile][row of M][16]; after the cluster barrier the owner sums the shares in
+ # rank order and stores. Every thread of the cluster takes part in the barriers.
+ _cluster_wait()
+ if warp < cw:
+ items1: cutlass.Constexpr = rt * m * TILE_N
+ for it1 in cutlass.range_constexpr(-(-items1 // (cw * 32))):
+ e1 = tidx + it1 * (cw * 32)
+ er1 = e1 % TILE_N
+ en1 = (e1 // TILE_N) % m
+ etl1 = e1 // (TILE_N * m)
+ if etl1 < nt:
+ tot1 = cutlass.Float32(0.0)
+ for q1 in cutlass.range_constexpr(kw):
+ tot1 = tot1 + _lds_f32(
+ red + (((q1 * rt + etl1) * np_ + en1) * TILE_N + er1) * 4
+ )
+ cslot = (
+ sbase
+ + self.o_cred
+ + (((rank * rt + etl1) * m + en1) * TILE_N + er1) * 4
+ )
+ _st_dsmem_f32(_mapa(cslot, etl1 % ks), tot1)
+ _cluster_arrive_release()
+ _cluster_wait_acquire()
+ if warp < cw:
+ items2: cutlass.Constexpr = rt * m * TILE_N
+ for it2 in cutlass.range_constexpr(-(-items2 // (cw * 32))):
+ f_ = tidx + it2 * (cw * 32)
+ fr = f_ % TILE_N
+ fn_ = (f_ // TILE_N) % m
+ ftl = f_ // (TILE_N * m)
+ frow = (t0 + ftl) * TILE_N + fr
+ if (ftl < nt) & (frow < n_out) & (ftl % ks == rank):
+ tot2 = cutlass.Float32(0.0)
+ for q2 in cutlass.range_constexpr(ks):
+ tot2 = tot2 + _lds_f32(
+ sbase
+ + self.o_cred
+ + (((q2 * rt + ftl) * m + fn_) * TILE_N + fr) * 4
+ )
+ _stg_bf16(out_base + (cutlass.Int64(fn_) * n_out + frow) * 2, tot2)
+
+
+# ------------------------------------------------------------------------------------------
+# host side
+# ------------------------------------------------------------------------------------------
+
+
+@dataclass
+class PackedFp8Weight:
+ """`data`: uint8 `[K/128, ceil(N/16), 2048]` (see `pack_weight`); `scale`: fp32
+ `[ceil(N/128), K/128]` contiguous; logical shape `[n, k]`."""
+
+ data: torch.Tensor
+ scale: torch.Tensor
+ n: int
+ k: int
+ tile_major: bool = False
+
+
+def quantize_fp8_block(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ """bf16/fp32 `[N, K]` -> (e4m3 codes `[N, K]`, fp32 scales `[ceil(N/128), K/128]`),
+ scale = amax(128x128 block) / 448, code = round_to_e4m3(w / scale)."""
+ n, k = w.shape
+ assert k % BLOCK_K == 0
+ nb = -(-n // 128)
+ wp = torch.zeros(nb * 128, k, dtype=torch.float32, device=w.device)
+ wp[:n] = w.float()
+ blocks = wp.reshape(nb, 128, k // BLOCK_K, BLOCK_K)
+ amax = blocks.abs().amax((1, 3))
+ scale = (amax / FP8_MAX).clamp_min(1e-30)
+ codes = (blocks / scale[:, None, :, None]).clamp(-FP8_MAX, FP8_MAX).reshape(nb * 128, k)[:n]
+ return codes.to(torch.float8_e4m3fn), scale.contiguous()
+
+
+def dequantize_fp8_block(codes: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+ """fp32 `[N, K]` = code * scale of its 128x128 block."""
+ n, k = codes.shape
+ s = scale.repeat_interleave(128, 0)[:n].repeat_interleave(BLOCK_K, 1)
+ return codes.float() * s
+
+
+#: Weights of at most this many 16-row tiles are stored tile-major (a CTA then holds one or two
+#: tiles, and one bulk copy brings a tile's k-blocks of a stage); larger ones k-block-major.
+TILE_MAJOR_MAX_TILES = 320
+
+
+def pack_weight(
+ codes: torch.Tensor, scale: torch.Tensor, tile_major: bool | None = None
+) -> PackedFp8Weight:
+ """e4m3 `[N, K]` (+ its block scales `[ceil(N/128), K/128]` fp32) -> the kernel's layout.
+
+ A unit (k-block kb, 16-row tile T) is 2048 bytes: rows 16T .. 16T + 15, k = 128 kb .. + 127, as
+ `[c (4)][lane (32)][16 B]`, lane = 4 g + t, the 16 bytes being
+ `W[16T + g][32c + 8t + 0..3], W[16T + g + 8][same], W[16T + g][32c + 8t + 4..7],
+ W[16T + g + 8][same]` (rows past N are zero). Units are stored `[K/128][ceil(N/16)]`
+ (k-block-major) or, with `tile_major`, `[ceil(N/16)][K/128]` (default: tile-major up to
+ TILE_MAJOR_MAX_TILES tiles)."""
+ n, k = codes.shape
+ assert k % BLOCK_K == 0 and codes.dtype == torch.float8_e4m3fn
+ n_tiles = -(-n // TILE_N)
+ if tile_major is None:
+ tile_major = n_tiles <= TILE_MAJOR_MAX_TILES
+ b = torch.zeros(n_tiles * TILE_N, k, dtype=torch.uint8, device=codes.device)
+ b[:n] = codes.view(torch.uint8)
+ # source (tile, h, g, kb, c, t, jj, e); unit order (c, g, t, jj, h, e)
+ v = b.reshape(n_tiles, 2, 8, k // BLOCK_K, 4, 4, 2, 4)
+ if tile_major:
+ data = v.permute(0, 3, 4, 2, 5, 6, 1, 7).contiguous().reshape(n_tiles, k // BLOCK_K, UNIT)
+ else:
+ data = v.permute(3, 0, 4, 2, 5, 6, 1, 7).contiguous().reshape(k // BLOCK_K, n_tiles, UNIT)
+ nb = -(-n // 128)
+ assert scale.shape == (nb, k // BLOCK_K), (scale.shape, nb)
+ return PackedFp8Weight(data, scale.float().contiguous(), n, k, tile_major)
+
+
+def pack_from_bf16(w: torch.Tensor) -> PackedFp8Weight:
+ """Converter from a bf16 `[N, K]` weight."""
+ codes, scale = quantize_fp8_block(w)
+ return pack_weight(codes, scale)
+
+
+#: Shared memory one CTA may take so that a CTA of the next launch fits beside it (PDL overlap).
+SMEM_BUDGET = 100 << 10
+
+
+#: Measured plans (GB300, 152 SMs, CUDA-graph chain with PDL, weights rotated over > L2) for the
+#: qwen3.8-flash-next-nvfp4 tp1 decode projections: (N, K) -> ((largest M, Plan fields), ...).
+#: `grid` 0 is "auto" (one CTA per 16-row tile up to 9/8 x SMs, else one per SM); -t is t tiles a
+#: CTA (-2: half the tiles).
+_TABLE: dict[tuple[int, int], tuple[tuple[int, dict], ...]] = {
+ # GDN in_proj (q|k|v|z|a|b)
+ (16480, 2560): (
+ (8, dict(grid=0, kw=2, tw=4, slots=3, l2_ahead=1, xreg=2)),
+ (16, dict(grid=0, kw=2, tw=4, slots=3, xreg=2)),
+ (32, dict(grid=0, kw=2, tw=4, slots=2, xreg=2)),
+ ),
+ # attention in-proj (q|gate|k|v|index)
+ (14080, 2560): (
+ (8, dict(grid=0, kw=2, tw=6, slots=3, l2_ahead=1, xreg=2)),
+ (16, dict(grid=0, kw=2, tw=6, slots=3, xreg=2)),
+ (32, dict(grid=0, kw=2, tw=6, slots=2, xreg=2)),
+ ),
+ # GDN out / attention o
+ (2560, 6144): (
+ (1, dict(grid=0, kw=12, tw=1, slots=4, xreg=1, xdepth=2)),
+ (8, dict(grid=0, kw=12, tw=1, slots=4, xreg=1, xdepth=1)),
+ (16, dict(grid=-2, kw=6, tw=2, slots=3, xreg=2)),
+ (32, dict(grid=-2, kw=6, tw=2, slots=2, xreg=2)),
+ ),
+ # shared expert gate_up with its gate row, and without it
+ (1281, 2560): ((32, dict(grid=0, kw=20, tw=1, slots=1, xreg=1)),),
+ (1280, 2560): ((32, dict(grid=0, kw=20, tw=1, slots=1, xreg=1)),),
+ # shared expert down
+ (2560, 640): ((32, dict(grid=0, kw=5, tw=1, slots=1, xreg=1)),),
+ # vocab projection (MTP draft head / verify head): 15520 tiles, 7 a CTA -> 2218 CTAs (about
+ # 15 a multiprocessor over the launch, ~66 KB shared memory each)
+ (248320, 2560): ((32, dict(grid=-7, kw=2, tw=4, slots=2, l2_ahead=1, xreg=2)),),
+}
+
+#: Measured plans of the NVFP4 units (`w4`), as `_TABLE`. GDN in_proj: four k-blocks a stage
+#: (5 stages), no L2 look-ahead: 6.1 / 6.6 / 7.1 us at 1 / 4 / 8 rows alone, 6.4 / 6.1 behind a
+#: PDL predecessor (1 / 4 rows); from 9 rows the heuristic plan, behind fp8_gemv there.
+_TABLE_W4: dict[tuple[int, int], tuple[tuple[int, dict], ...]] = {
+ (16480, 2560): ((8, dict(grid=0, kw=4, tw=4, slots=3, xreg=2)),),
+}
+
+#: Most 16-row tiles a CTA takes in the heuristic plan: a ring slot holds `kw` k-blocks of each of
+#: them, so a larger N gets more CTAs (several per multiprocessor over the launch) instead.
+MAX_TILES_PER_CTA = 7
+
+
+def default_plan(n: int, k: int, m: int, sms: int, w4: bool = False) -> Plan:
+ """The plan `fp8_gemv` launches with: the measured table entry for (N, K, M), else a heuristic
+ (`w4`: over NVFP4 units)."""
+ n_tiles = -(-n // TILE_N)
+ nkb = k // BLOCK_K
+ auto_grid = n_tiles if n_tiles <= sms + sms // 8 else sms
+ for bound, fields in (_TABLE_W4 if w4 else _TABLE).get((n, k), ()):
+ if m <= bound:
+ f = dict(fields)
+ g = f.pop("grid")
+ grid = auto_grid if g == 0 else -(-n_tiles // -g)
+ return Plan(grid=grid, **f)
+ grid = auto_grid
+ rt = -(-n_tiles // grid)
+ if rt > MAX_TILES_PER_CTA:
+ grid = -(-n_tiles // MAX_TILES_PER_CTA)
+ rt = -(-n_tiles // grid)
+ divs = [d for d in range(1, nkb + 1) if nkb % d == 0]
+ if rt >= 4:
+ kw = max(d for d in divs if d <= 2)
+ tw = min(rt, 4)
+ else:
+ kw = max(d for d in divs if d <= 12)
+ tw = rt
+ nt8 = -(-m // 8)
+ xreg = (nkb // kw) * nt8 * 16 <= XREG_MAX
+ slot = kw * rt * (UNIT_W4 if w4 else UNIT) + (0 if xreg else kw * m * XRS)
+ slots = max(1, min(nkb // kw, SMEM_BUDGET // slot))
+ return Plan(grid=grid, kw=kw, tw=tw, slots=slots, xreg=1 if xreg else 2)
+
+
+@functools.lru_cache(maxsize=None)
+def _compiled(
+ n,
+ k,
+ m,
+ plan: Plan,
+ pdl: bool,
+ tile_major: bool,
+ conv_dim: int = 0,
+ pad_slot: int = -1,
+ conv_p: int = 0,
+ w4: bool = False,
+):
+ kern = _Fp8Gemv(
+ n=n,
+ k=k,
+ m=m,
+ plan=plan,
+ pdl=pdl,
+ tile_major=tile_major,
+ conv_dim=conv_dim,
+ pad_slot=pad_slot,
+ conv_p=conv_p,
+ w4=w4,
+ )
+ dev = "cuda"
+ out = torch.empty(m, n, dtype=torch.bfloat16, device=dev)
+ x = torch.empty(m, k, dtype=torch.bfloat16, device=dev)
+ w = torch.empty(16, dtype=torch.uint8, device=dev)
+ s = torch.empty(4, dtype=torch.float32, device=dev)
+ return cute.compile(
+ kern,
+ from_dlpack(out, assumed_align=16).mark_layout_dynamic(leading_dim=1),
+ make_ptr(cutlass.BFloat16, x.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
+ make_ptr(cutlass.Uint8, w.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
+ make_ptr(cutlass.Float32, s.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
+ cutlass.Int32(k),
+ make_ptr(cutlass.BFloat16, x.data_ptr(), cute.AddressSpace.gmem, assumed_align=2),
+ make_ptr(cutlass.BFloat16, x.data_ptr(), cute.AddressSpace.gmem, assumed_align=2),
+ make_ptr(cutlass.Int32, s.data_ptr(), cute.AddressSpace.gmem, assumed_align=4),
+ cutlass.Int32(0),
+ cutlass.Int32(0),
+ cutlass.Int32(0),
+ make_ptr(cutlass.BFloat16, x.data_ptr(), cute.AddressSpace.gmem, assumed_align=2),
+ make_ptr(cutlass.Int32, s.data_ptr(), cute.AddressSpace.gmem, assumed_align=4),
+ cutlass.Int32(0),
+ cutlass.Int32(0),
+ cutlass.Int32(0),
+ cutlass.Int32(0),
+ cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+ options="--opt-level 2 --enable-tvm-ffi",
+ )
+
+
+@functools.lru_cache(maxsize=None)
+def _sms() -> int:
+ return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count
+
+
+#: The launch arguments past `x_stride` of a call without a conv epilogue.
+_NO_CONV = (0,) * 12
+
+#: Most tokens per request of `fp8_gemv_conv1d_verify` (`max_draft_len + 1` of the widest
+#: verify step the target runs).
+VERIFY_MAX_P = 8
+
+
+def _env_pdl() -> bool:
+ try:
+ from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+ except ImportError:
+ return True
+ return bool(get_env_enable_pdl())
+
+
+def fp8_gemv(
+ x: torch.Tensor, w: PackedFp8Weight, pdl: bool = True, plan: Plan | None = None
+) -> torch.Tensor:
+ """`x[M, K]` bf16 (row stride >= K, 16-byte aligned rows) @ dequant(w)^T -> fresh
+ `[M, N]` bf16. M = 1 .. MAX_ROWS. `w` a `PackedFp8Weight`, or NVFP4 units in this kernel's
+ order (`nvfp4_gemv.SplitNvfp4Weight`, `w.w4`)."""
+ m, k = x.shape
+ assert k == w.k and x.dtype == torch.bfloat16 and x.stride(1) == 1
+ assert 1 <= m <= MAX_ROWS
+ assert x.stride(0) % 8 == 0 and x.data_ptr() % 16 == 0
+ w4 = getattr(w, "w4", False)
+ out = torch.empty((m, w.n), dtype=torch.bfloat16, device=x.device)
+ if plan is None:
+ plan = default_plan(w.n, k, m, _sms(), w4)
+ fn = _compiled(w.n, k, m, plan, bool(pdl and _env_pdl()), w.tile_major, w4=w4)
+ fn(
+ out,
+ x.data_ptr(),
+ w.data.data_ptr(),
+ w.scale.data_ptr(),
+ x.stride(0),
+ *_NO_CONV,
+ )
+ return out
+
+
+def fp8_gemv_conv1d(
+ x: torch.Tensor,
+ w: PackedFp8Weight,
+ conv_state: torch.Tensor,
+ conv_weight: torch.Tensor,
+ conv_state_indices: torch.Tensor,
+ conv_dim: int,
+ pad_slot_id: int,
+ pdl: bool = True,
+ plan: Plan | None = None,
+) -> torch.Tensor:
+ """FP8-weight `gdn_in_proj_conv1d_pdl`: `p = round_bf16(x @ dequant(w)^T)`; for a column
+ c < conv_dim of a row r whose slot s = conv_state_indices[r] is not pad_slot_id,
+ `out[r, c] = round_bf16(silu(sum_j conv_weight[c, j] * [conv_state[s, c, :], p[r, c]][j]))` (fp32)
+ and `conv_state[s, c, :] = [conv_state[s, c, 1:], p[r, c]]`; every other output is `p`.
+ `conv_state` `[slots, conv_dim, 3]` bf16 (any strides), `conv_weight` `[conv_dim, 4]` bf16
+ contiguous, `conv_state_indices` int32 `[M]`. Returns the fresh `[M, N]` output. `w` as
+ `fp8_gemv`'s."""
+ m, k = x.shape
+ assert k == w.k and x.dtype == torch.bfloat16 and x.stride(1) == 1 and 1 <= m <= MAX_ROWS
+ assert x.stride(0) % 8 == 0 and x.data_ptr() % 16 == 0
+ assert conv_state.dtype == torch.bfloat16 and tuple(conv_state.shape[1:]) == (conv_dim, 3)
+ assert conv_weight.dtype == torch.bfloat16 and tuple(conv_weight.shape) == (conv_dim, 4)
+ assert conv_weight.is_contiguous() and conv_dim <= w.n
+ assert conv_state_indices.dtype == torch.int32 and conv_state_indices.is_contiguous()
+ assert conv_state_indices.numel() == m
+ w4 = getattr(w, "w4", False)
+ out = torch.empty((m, w.n), dtype=torch.bfloat16, device=x.device)
+ if plan is None:
+ plan = default_plan(w.n, k, m, _sms(), w4)
+ fn = _compiled(
+ w.n, k, m, plan, bool(pdl and _env_pdl()), w.tile_major, conv_dim, int(pad_slot_id), w4=w4
+ )
+ fn(
+ out,
+ x.data_ptr(),
+ w.data.data_ptr(),
+ w.scale.data_ptr(),
+ x.stride(0),
+ conv_state.data_ptr(),
+ conv_weight.data_ptr(),
+ conv_state_indices.data_ptr(),
+ *conv_state.stride(),
+ *_NO_CONV[6:],
+ )
+ return out
+
+
+def fp8_gemv_conv1d_verify(
+ x: torch.Tensor,
+ w: PackedFp8Weight,
+ conv_state: torch.Tensor,
+ conv_weight: torch.Tensor,
+ conv_state_indices: torch.Tensor,
+ intermediate_conv_window: torch.Tensor,
+ intermediate_state_indices: torch.Tensor,
+ conv_dim: int,
+ rows_per_seq: int,
+ pad_slot_id: int,
+ pdl: bool = True,
+ plan: Plan | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """`fp8_gemv` + `causal_conv1d_triton_update` of a speculative verify step in one launch.
+
+ `x` is `n` requests of `p = rows_per_seq` consecutive token rows (row `s p + t` is token
+ `t` of request `s`); `proj = round_bf16(x @ dequant(w)^T)`. For request `s` with slot
+ `c = conv_state_indices[s]` != pad_slot_id and snapshot row `r =
+ intermediate_state_indices[s]`, over the window sequence `win = [conv_state[c, :, 0..2],
+ proj[s p .. s p + p - 1, :conv_dim]]` (per channel): `conv[s p + t] = silu(sum_j
+ bf16(conv_weight[:, j] * win[t + j]))` (bf16 products, fp32 sum, the Triton kernel's
+ SiLU), `intermediate_conv_window[r, t, :conv_dim, :] = win[t + 1 .. t + 3]` for t < p, and
+ `conv_state[c] = win[p .. p + 2]`; a pad request's rows get the plain projection and
+ write no state. Returns `(conv [M, conv_dim], proj[:, conv_dim:] [M, N - conv_dim])`,
+ both contiguous views of one fresh buffer.
+
+ `conv_state` `[slots, conv_dim, 3]` and `intermediate_conv_window` `[rows, >= p,
+ conv_dim, 3]` bf16 (any strides), `conv_weight` `[conv_dim, 4]` bf16 contiguous, both
+ index vectors int32 `[n]` contiguous. The slots, snapshot rows, taps and the windows are
+ read before the PDL wait (as `fp8_gemv_conv1d` does). `w` as `fp8_gemv`'s."""
+ m, k = x.shape
+ p = int(rows_per_seq)
+ assert k == w.k and x.dtype == torch.bfloat16 and x.stride(1) == 1 and 2 <= m <= MAX_ROWS
+ assert x.stride(0) % 8 == 0 and x.data_ptr() % 16 == 0
+ assert 2 <= p <= VERIFY_MAX_P and m % p == 0, (m, p)
+ nreq = m // p
+ assert conv_state.dtype == torch.bfloat16 and tuple(conv_state.shape[1:]) == (conv_dim, 3)
+ assert conv_weight.dtype == torch.bfloat16 and tuple(conv_weight.shape) == (conv_dim, 4)
+ assert conv_weight.is_contiguous() and conv_weight.data_ptr() % 8 == 0 and 0 < conv_dim < w.n
+ iw, iidx = intermediate_conv_window, intermediate_state_indices
+ assert iw.dtype == torch.bfloat16 and iw.dim() == 4 and iw.shape[1] >= p
+ assert tuple(iw.shape[2:]) == (conv_dim, 3)
+ for v in (conv_state_indices, iidx):
+ assert v.dtype == torch.int32 and v.is_contiguous() and v.numel() == nreq
+ w4 = getattr(w, "w4", False)
+ buf = torch.empty(m * w.n, dtype=torch.bfloat16, device=x.device)
+ if plan is None:
+ plan = default_plan(w.n, k, m, _sms(), w4)
+ fn = _compiled(
+ w.n,
+ k,
+ m,
+ plan,
+ bool(pdl and _env_pdl()),
+ w.tile_major,
+ conv_dim,
+ int(pad_slot_id),
+ p,
+ w4=w4,
+ )
+ fn(
+ buf.view(m, w.n),
+ x.data_ptr(),
+ w.data.data_ptr(),
+ w.scale.data_ptr(),
+ x.stride(0),
+ conv_state.data_ptr(),
+ conv_weight.data_ptr(),
+ conv_state_indices.data_ptr(),
+ *conv_state.stride(),
+ iw.data_ptr(),
+ iidx.data_ptr(),
+ *iw.stride(),
+ )
+ return buf[: m * conv_dim].view(m, conv_dim), buf[m * conv_dim :].view(m, w.n - conv_dim)
+
+
+def _gemv_cell(why: str, n: int, k: int, m: int) -> Cell:
+ return Cell(why=why, spec=dict(n=n, k=k, m=m))
+
+
+class Fp8Gemv(OpWrapper):
+ """`fp8_gemv(x, w)`: `y[M, N] = x[M, K] @ dequant(w)^T` in bf16 for 1..MAX_ROWS rows, the
+ weight a `PackedFp8Weight` (e4m3 codes with 128x128 fp32 block scales, packed by
+ `pack_weight` / `pack_from_bf16`), the bf16 activation never quantized. fp32 accumulation per
+ 128-wide k-block, block scale applied in fp32, one bf16 rounding at the store.
+
+ Fusion boundary: the projection; `fp8_gemv_conv1d` also fuses the GDN causal-conv step of
+ decode rows, `fp8_gemv_conv1d_verify` the multi-token step of a verify step's rows."""
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ _gemv_cell("GDN in_proj, one decode row", 16480, 2560, 1),
+ _gemv_cell("attention in-proj, MTP3 verify at con1 (4 rows)", 14080, 2560, 4),
+ _gemv_cell("GDN out / attention o, 4 rows", 2560, 6144, 4),
+ _gemv_cell("shared expert gate|up, 16 rows", 1280, 2560, 16),
+ _gemv_cell("shared expert down, 8 rows", 2560, 640, 8),
+ _gemv_cell("vocab projection (draft / verify head), 4 rows", 248320, 2560, 4),
+ )
+
+ note = """
+ Metadata consumed: none. Call sites: the decode target's dense projections under the dense
+ FP8 mode (`TRTLLM_MODELING_V2_DENSE_FP8`) at up to `_FP8_GEMV_MAX_ROWS` rows.
+ Preconditions: x bf16 `[M, K]`, unit column stride, row stride a multiple of 8 elements and a
+ 16-byte aligned base; K a multiple of 128. One PDL launch (weights streamed before the
+ wait); CUDA-graph capturable; compiled per (N, K, M, plan) on first use.
+ """
+
+ def __call__(self, x: torch.Tensor, w: PackedFp8Weight) -> torch.Tensor:
+ return fp8_gemv(x, w)
+
+ def reference(self, x: torch.Tensor, codes: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+ """fp32 `[M, N]` over the dequantized weight (a true fp32 matmul)."""
+ with true_fp32_matmul():
+ return x.float() @ dequantize_fp8_block(codes, scale).t()
+
+
+class Fp8GemvConv1d(OpWrapper):
+ """`fp8_gemv_conv1d`: the FP8-weight `gdn_in_proj_conv1d_pdl` -- the GDN input projection
+ of decode rows, then for the conv columns of a live row the causal conv step against its
+ paged state (SiLU; the state advanced in place). Returns the `[M, N]` projection row with
+ the conv columns replaced by the conv output."""
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="STP decode: one row, live slot", spec=dict(slots=(37,))),
+ Cell(why="two rows, one padding slot", spec=dict(slots=(17, -1))),
+ )
+
+ note = """
+ Call site: the decode target's one-row GDN in_proj + conv step under the dense FP8 mode.
+ `conv_state` `[slots, conv_dim, 3]` bf16 (any strides), `conv_weight` `[conv_dim, 4]` bf16
+ contiguous, int32 slots; a `pad_slot_id` row gets the plain projection.
+ """
+
+ def __call__(self, x, w, conv_state, conv_weight, conv_state_indices, conv_dim, pad_slot_id):
+ return fp8_gemv_conv1d(
+ x, w, conv_state, conv_weight, conv_state_indices, conv_dim, pad_slot_id
+ )
+
+ def reference(self, *args, **kwargs):
+ raise NotImplementedError("the test builds the fp32 reference from the dequantized weight")
+
+
+#: The checkpoint's GDN in_proj (q|k|v|z|a|b) and its q|k|v conv width.
+GDN_IN_N, GDN_IN_K, GDN_CONV_DIM = 16480, 2560, 10240
+
+
+def _verify_cell(why: str, p: int, slots: tuple, steps: int | None = None) -> Cell:
+ return Cell(
+ why=why,
+ spec=dict(
+ n=GDN_IN_N,
+ k=GDN_IN_K,
+ conv_dim=GDN_CONV_DIM,
+ p=p,
+ slots=slots,
+ snapshot_steps=steps or p,
+ pool_slots=461,
+ snapshot_rows=32,
+ ),
+ )
+
+
+def verify_reference(
+ x,
+ deq,
+ conv_state,
+ conv_weight,
+ conv_state_indices,
+ intermediate_conv_window,
+ intermediate_state_indices,
+ conv_dim,
+ rows_per_seq,
+ pad_slot_id,
+):
+ """`fp8_gemv_conv1d_verify`'s reference over a dequantized fp32 weight `deq` `[N, K]`; see
+ `Fp8GemvConv1dVerify.reference`."""
+ with true_fp32_matmul():
+ proj = (x.float() @ deq.t()).to(torch.bfloat16)
+ m, p = x.shape[0], rows_per_seq
+ conv = proj[:, :conv_dim].clone()
+ state, window = conv_state.clone(), intermediate_conv_window.clone()
+ slots = conv_state_indices.tolist()
+ rows = intermediate_state_indices.tolist()
+ for s in range(m // p):
+ if slots[s] == pad_slot_id:
+ continue
+ seq = torch.cat([conv_state[slots[s]], proj[s * p : (s + 1) * p, :conv_dim].t()], 1)
+ for t in range(p):
+ acc = torch.zeros(conv_dim, dtype=torch.float32, device=x.device)
+ for j in range(4):
+ acc = acc + (conv_weight[:, j] * seq[:, t + j]).float()
+ conv[s * p + t] = (acc / (1.0 + torch.exp(-acc))).to(torch.bfloat16)
+ window[rows[s], t] = seq[:, t + 1 : t + 4]
+ state[slots[s]] = seq[:, p : p + 3]
+ return conv, proj[:, conv_dim:].contiguous(), state, window
+
+
+class Fp8GemvConv1dVerify(OpWrapper):
+ """`fp8_gemv_conv1d_verify`: the GDN input projection of a speculative verify step (n
+ requests x p tokens) with `causal_conv1d_triton_update`'s multi-token conv step as its
+ epilogue -- per conv column and request, the p tokens through the cached window (bf16
+ products, fp32 sum, SiLU), the window after each token into the per-step snapshot buffer,
+ and the pool slot advanced by all p tokens. Returns `(conv [M, conv_dim], projection
+ columns conv_dim.. [M, N - conv_dim])`.
+
+ Fusion boundary: projection, conv, SiLU, snapshots and the pool advance. The q/k/v and
+ z/a/b splits, the recurrence and the post-acceptance promotion of a snapshot are the
+ caller's."""
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec keys: n, k (the in_proj), conv_dim, p (tokens per request), slots (one per
+ # request; -1 the pad sentinel), snapshot_steps (the buffer's step depth), pool_slots,
+ # snapshot_rows. Every cell: bf16, width 4, no bias, SiLU, pad_slot_id -1,
+ # intermediate_state_indices = arange(n), the engine's 461-slot pool and 32-row buffer.
+ CELLS: tuple[Cell, ...] = (
+ _verify_cell("MTP3 verify, 1 request (4 rows)", 4, (37,)),
+ _verify_cell("MTP3 verify, 2 requests (8 rows)", 4, (0, 460)),
+ _verify_cell("MTP3 verify, 4 requests (16 rows)", 4, (3, 17, 233, 459)),
+ _verify_cell("MTP1 verify, 8 requests x 2 rows (16 rows)", 2, tuple(range(5, 461, 57))),
+ _verify_cell("MTP7 verify, 2 requests x 8 rows", 8, (100, 7)),
+ _verify_cell("3 rows a request into an 8-step buffer: steps 3..7 untouched", 3, (9, 41), 8),
+ _verify_cell("MTP3, a pad request beside a live one", 4, (-1, 21)),
+ )
+
+ note = """
+ Call site: the decode target's GDN in_proj of an MTP verify step under the dense FP8 mode
+ in a captured graph, 2..16 rows (p = runtime_draft_len + 1 = 2..8 tokens per request,
+ a request's rows consecutive); the conv output feeds the cached replay as its contiguous
+ packed q|k|v, z/a/b are column views of the second output.
+ `conv_state` `[slots, conv_dim, 3]` and `intermediate_conv_window` `[rows, >= p, conv_dim,
+ 3]` bf16 (any strides), `conv_weight` `[conv_dim, 4]` bf16 contiguous, int32 slot and
+ snapshot-row vectors of n entries. Same arithmetic as `causal_conv1d_triton_update`
+ (bitwise the two-launch fp8_gemv + Triton path under the same plan); a pad request gets
+ the plain projection and writes neither pool nor snapshots (the Triton op leaves its rows
+ unwritten). Slots, snapshot rows, taps and windows are read before the PDL wait: nothing
+ in a captured step writes this layer's pool ahead of it.
+ Not guarded (content): duplicate slots or snapshot rows race; slot / row values out of
+ range are illegal accesses.
+ """
+
+ def __call__(
+ self,
+ x,
+ w,
+ conv_state,
+ conv_weight,
+ conv_state_indices,
+ intermediate_conv_window,
+ intermediate_state_indices,
+ conv_dim,
+ rows_per_seq,
+ pad_slot_id,
+ ):
+ return fp8_gemv_conv1d_verify(
+ x,
+ w,
+ conv_state,
+ conv_weight,
+ conv_state_indices,
+ intermediate_conv_window,
+ intermediate_state_indices,
+ conv_dim,
+ rows_per_seq,
+ pad_slot_id,
+ )
+
+ def reference(
+ self,
+ x,
+ codes,
+ scale,
+ conv_state,
+ conv_weight,
+ conv_state_indices,
+ intermediate_conv_window,
+ intermediate_state_indices,
+ conv_dim,
+ rows_per_seq,
+ pad_slot_id,
+ ):
+ """`(conv, rest, conv_state, intermediate_conv_window)` as the call leaves them, from
+ an fp32 matmul over the dequantized weight rounded once to bf16, then the Triton
+ op's reference conv (products in bf16, fp32 sum, `acc / (1 + exp(-acc))`); mutates
+ nothing. A pad request's conv rows are its projection."""
+ return verify_reference(
+ x,
+ dequantize_fp8_block(codes, scale),
+ conv_state,
+ conv_weight,
+ conv_state_indices,
+ intermediate_conv_window,
+ intermediate_state_indices,
+ conv_dim,
+ rows_per_seq,
+ pad_slot_id,
+ )
+
+ def is_valid(
+ self,
+ x,
+ w,
+ conv_state,
+ conv_weight,
+ conv_state_indices,
+ intermediate_conv_window,
+ intermediate_state_indices,
+ conv_dim,
+ rows_per_seq,
+ pad_slot_id,
+ **kwargs,
+ ) -> None:
+ m = x.shape[0]
+ assert 2 <= rows_per_seq <= VERIFY_MAX_P and m % rows_per_seq == 0, (
+ f"{m} rows are not whole requests of {rows_per_seq} tokens (2..{VERIFY_MAX_P})"
+ )
+ assert m <= MAX_ROWS, f"{m} rows; the kernel holds at most {MAX_ROWS}"
+ assert intermediate_conv_window.shape[1] >= rows_per_seq, (
+ f"a snapshot buffer of {intermediate_conv_window.shape[1]} steps holds fewer than "
+ f"the {rows_per_seq} tokens a request writes"
+ )
+ assert conv_state.shape[2] == 3 and conv_weight.shape[1] == 4, "width-4 conv only"
+ assert conv_state.shape[1] == conv_dim and intermediate_conv_window.shape[2] == conv_dim
+
+
+fp8_gemv_op = Fp8Gemv()
+fp8_gemv_conv1d_op = Fp8GemvConv1d()
+fp8_gemv_conv1d_verify_op = Fp8GemvConv1dVerify()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/fp8_per_tensor_mm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/fp8_per_tensor_mm.py
new file mode 100644
index 000000000000..69106bdd5cd6
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/fp8_per_tensor_mm.py
@@ -0,0 +1,192 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""A dense projection in FP8: the bf16 activation quantized to e4m3 with one dynamic
+per-tensor scale, the weight stored in e4m3 with one per-tensor scale, cuBLASLt's FP8 GEMM
+with both scales in its epilogue, bf16 out.
+
+Tier 5 (Triton written for this target) for the activation quantization: two launches --
+per-block partial maxima of `|x|`, then the quantization, each of whose programs reduces the
+partials itself (no memset, no host sync) -- and `torch._scaled_mm` (cuBLASLt, tensor-wise
+scales) for the GEMM. Tensor-wise scales keep cuBLASLt on its fast FP8 kernels: on GB300 the
+checkpoint's GDN input projection runs ~1.9x bf16 at 16K rows and ~1.7x at 1-64 rows, where
+row-wise activation scales fall back to kernels only ~1.4x bf16.
+"""
+
+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
+
+E4M3_MAX = 448.0
+#: Floor of the activation maximum: an all-zero activation quantizes to zeros at this scale.
+AMAX_FLOOR = 1e-12
+#: Most programs of either launch; rows are grid-strided above it, and the quantization
+#: reduces at most this many partial maxima.
+_MAX_PROGRAMS = 1024
+_ROWS = 4
+
+
+def fp8_per_tensor_weight(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ """`(codes, scale)` of a bf16 `[N, K]` weight: e4m3 codes `[N, K]` (contiguous) and the
+ fp32 `[1]` dequantization scale `amax / 448`."""
+ scale = (weight.abs().amax().float().clamp_min(AMAX_FLOOR) / E4M3_MAX).reshape(1)
+ codes = (weight.float() / scale).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn)
+ return codes.contiguous(), scale
+
+
+@triton.jit
+def _amax_kernel(
+ x_ptr, part_ptr, M, K, stride_m, PROGRAMS, ROWS: tl.constexpr, BLOCK_K: tl.constexpr
+):
+ pid = tl.program_id(0)
+ acc = tl.zeros([ROWS, BLOCK_K], dtype=tl.float32)
+ for r0 in range(pid * ROWS, M, PROGRAMS * ROWS):
+ rows = r0 + tl.arange(0, ROWS)
+ for k0 in range(0, K, BLOCK_K):
+ cols = k0 + tl.arange(0, BLOCK_K)
+ mask = (rows[:, None] < M) & (cols[None, :] < K)
+ v = tl.load(
+ x_ptr + rows[:, None].to(tl.int64) * stride_m + cols[None, :], mask=mask, other=0.0
+ )
+ acc = tl.maximum(acc, tl.abs(v.to(tl.float32)))
+ tl.store(part_ptr + pid, tl.max(tl.max(acc, axis=1), axis=0))
+
+
+@triton.jit
+def _quant_kernel(
+ x_ptr,
+ part_ptr,
+ q_ptr,
+ scale_ptr,
+ M,
+ K,
+ stride_m,
+ PROGRAMS,
+ NUM_PARTS,
+ amax_floor,
+ ROWS: tl.constexpr,
+ BLOCK_K: tl.constexpr,
+ BLOCK_P: tl.constexpr,
+):
+ pid = tl.program_id(0)
+ p = tl.arange(0, BLOCK_P)
+ amax = tl.max(tl.load(part_ptr + p, mask=p < NUM_PARTS, other=0.0), axis=0)
+ amax = tl.maximum(amax, amax_floor)
+ scale = amax / 448.0
+ inv = 448.0 / amax
+ if pid == 0:
+ tl.store(scale_ptr, scale)
+ for r0 in range(pid * ROWS, M, PROGRAMS * ROWS):
+ rows = r0 + tl.arange(0, ROWS)
+ for k0 in range(0, K, BLOCK_K):
+ cols = k0 + tl.arange(0, BLOCK_K)
+ mask = (rows[:, None] < M) & (cols[None, :] < K)
+ v = tl.load(
+ x_ptr + rows[:, None].to(tl.int64) * stride_m + cols[None, :], mask=mask, other=0.0
+ )
+ q = tl.minimum(tl.maximum(v.to(tl.float32) * inv, -448.0), 448.0)
+ tl.store(
+ q_ptr + rows[:, None].to(tl.int64) * K + cols[None, :],
+ q.to(tl.float8e4nv),
+ mask=mask,
+ )
+
+
+def fp8_quantize_per_tensor(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ """`(codes [M, K] e4m3 contiguous, scale [1] fp32)` of a bf16 `[M, K]` activation with
+ unit column stride: scale = max(|x|) / 448 (floored), codes = round(x / scale)."""
+ m, k = x.shape
+ q = torch.empty((m, k), dtype=torch.float8_e4m3fn, device=x.device)
+ scale = torch.empty(1, dtype=torch.float32, device=x.device)
+ programs = max(1, min(triton.cdiv(m, _ROWS), _MAX_PROGRAMS))
+ block_k = min(1024, triton.next_power_of_2(k))
+ part = torch.empty(programs, dtype=torch.float32, device=x.device)
+ _amax_kernel[(programs,)](
+ x, part, m, k, x.stride(0), programs, ROWS=_ROWS, BLOCK_K=block_k, num_warps=4
+ )
+ _quant_kernel[(programs,)](
+ x,
+ part,
+ q,
+ scale,
+ m,
+ k,
+ x.stride(0),
+ programs,
+ part.shape[0],
+ AMAX_FLOOR,
+ ROWS=_ROWS,
+ BLOCK_K=block_k,
+ BLOCK_P=triton.next_power_of_2(part.shape[0]),
+ num_warps=4,
+ )
+ return q, scale
+
+
+def _ckpt(m: int, k: int, n: int) -> dict:
+ return dict(m=m, k=k, n=n)
+
+
+class Fp8PerTensorMm(OpWrapper):
+ """`y = (q(x) * sx) @ (codes * sw)^T` in bf16, accumulated in fp32, where `sx =
+ max(|x|) / 448` over the whole activation and `q(x) = round_e4m3(x / sx)`; `weight_t` is
+ the `[K, N]` transpose view of the `[N, K]` e4m3 codes `fp8_per_tensor_weight` returns and
+ `weight_scale` its `[1]` scale.
+
+ Lossy against the bf16 projection: ~3.7% relative error per output on random operands
+ (e4m3 keeps 3 mantissa bits on both operands). The activation scale is shared by every
+ row of the call, so a row's rounding depends on the largest element of its batch.
+
+ Fusion boundary: activation quantization + GEMM. The weight conversion
+ (`fp8_per_tensor_weight`) is done once at load by the caller.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="GDN input projection, one decode row", spec=_ckpt(1, 2560, 16480)),
+ Cell(why="GDN / attention output projection, 128 rows", spec=_ckpt(128, 6144, 2560)),
+ Cell(
+ why="attention input projection, a 16384-row prefill step",
+ spec=_ckpt(16384, 2560, 14080),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none. Call sites: the target's GDN in/out and attention in/o
+ projections when the dense FP8 mode is on (`TRTLLM_MODELING_V2_DENSE_FP8`).
+ Preconditions: `x` bf16 `[M, K]` with unit column stride, K a multiple of 16; the weight
+ `[K, N]` column-major view of contiguous `[N, K]` e4m3 codes with N a multiple of 16
+ (cuBLASLt FP8). Three launches (two Triton, one cuBLASLt); CUDA-graph capturable: no host
+ sync, the activation scale stays on the device.
+ """
+
+ def __call__(
+ self, x: torch.Tensor, weight_t: torch.Tensor, weight_scale: torch.Tensor
+ ) -> torch.Tensor:
+ q, scale = fp8_quantize_per_tensor(x)
+ return torch._scaled_mm(
+ q, weight_t, scale_a=scale, scale_b=weight_scale, out_dtype=torch.bfloat16
+ )
+
+ def reference(
+ self, x: torch.Tensor, weight_t: torch.Tensor, weight_scale: torch.Tensor
+ ) -> torch.Tensor:
+ """fp32 `[M, N]`: the activation quantized as the op does, both operands dequantized,
+ an fp32 matmul."""
+ scale = x.abs().amax().float().clamp_min(AMAX_FLOOR) / E4M3_MAX
+ xq = (x.float() / scale).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn).float() * scale
+ return xq @ (weight_t.float() * weight_scale.float())
+
+ def is_valid(self, x: torch.Tensor, weight_t: torch.Tensor, weight_scale: torch.Tensor) -> None:
+ assert x.dtype == torch.bfloat16 and x.dim() == 2 and x.stride(1) == 1
+ assert weight_t.dtype == torch.float8_e4m3fn and weight_t.stride(0) == 1
+ assert weight_t.shape[0] == x.shape[1], "K mismatch"
+ assert x.shape[1] % 16 == 0 and weight_t.shape[1] % 16 == 0, "cuBLASLt FP8 needs K, N % 16"
+ assert weight_scale.numel() == 1 and weight_scale.dtype == torch.float32
+
+
+fp8_per_tensor_mm = Fp8PerTensorMm()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/low_m_bf16_mm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/low_m_bf16_mm.py
new file mode 100644
index 000000000000..6a6722a355a6
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/low_m_bf16_mm.py
@@ -0,0 +1,256 @@
+# 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
+are qwen3.8-flash-next-nvfp4's decode projections (hidden 2560) at tp1 and at a
+tp4/tep4 rank.
+
+FlashInfer API: ``flashinfer.gemm.kernels.dense_bf16_gemm_direct.run_direct_dense`` and
+``flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk.run_splitk_dense``.
+"""
+
+from __future__ import annotations
+
+import functools
+from typing import NamedTuple
+
+import torch
+from flashinfer.gemm.kernels import dense_bf16_gemm_direct as _fi_direct
+from flashinfer.gemm.kernels import dense_bf16_gemm_sm100_splitk as _fi_splitk
+
+import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import (
+ Arch,
+ Cell,
+ OpWrapper,
+ true_fp32_matmul,
+)
+
+
+class _Plan(NamedTuple):
+ """`kind` is "direct", "splitk" or "cublas"; `tactic` is the kernel's own tactic tuple."""
+
+ kind: str
+ tactic: tuple = ()
+
+
+def _d(block_size: int, outputs_per_block: int) -> tuple:
+ """A direct-kernel choice; rows_per_block is the kernel's default for the runtime M."""
+ return ("direct", block_size, outputs_per_block)
+
+
+def _s(mma_m: int, mma_n: int, split_k: int, ab_stages: int) -> tuple:
+ return ("splitk", mma_m, mma_n, split_k, ab_stages)
+
+
+#: (N, K) -> ((M upper bound, choice), ...) in ascending M. An M above the last
+#: bound, and any (N, K) not listed, runs on cuBLAS. Every row below is a
+#: qwen3.8-flash-next-nvfp4 decode projection at tp1 or at a tp4/tep4 rank; other
+#: shapes are added after measuring them. Each choice beat cublas_mm by >= 5% at the
+#: bound's M in a sweep of these shapes (GB300, cold L2, graph replay, pdl off).
+#: M=3 uses the M=4 choice, and so on.
+_TABLE: dict[tuple[int, int], tuple[tuple[int, tuple], ...]] = {
+ # GDN in_qkv: 16 q + 16 k heads x 128, 48 v heads x 128
+ (10240, 2560): ((1, _d(64, 1)), (2, _d(64, 2)), (8, _s(128, 8, 1, 6))),
+ # GDN in_proj at a tp4/tep4 rank: q|k (4 heads x 128 each), v (12 x 128), z (12 x 128),
+ # a, b (12 each) concatenated: 4120 columns.
+ (4120, 2560): ((1, _d(64, 1)), (2, _d(64, 2)), (8, _s(64, 8, 2, 10))),
+ # Attention input projection at a tp4/tep4 rank: q|gate (6 heads x 256 each), k, v (one
+ # KV head x 256 each), the replicated index q|k (5 x 128) and one pad head of 128: 4352.
+ (4352, 2560): ((2, _d(64, 2)), (8, _s(64, 8, 2, 10))),
+ # Attention o and GDN out at a tp4/tep4 rank (6 q heads x 256 = 12 v heads x 128 = 1536
+ # in); M <= 4 go to `bf16_gemv_pdl`, these rows are what this entry runs at the same shape.
+ (2560, 1536): ((4, _d(64, 2)), (8, _s(64, 8, 3, 10))),
+ # GDN in_z
+ (6144, 2560): ((1, _d(64, 1)), (2, _d(64, 2))),
+ # GDN in_a / in_b (48 v heads)
+ (48, 2560): ((32, _d(64, 1)),),
+ # GDN out and attention o: [6144 -> 2560]
+ (2560, 6144): (
+ (2, _d(256, 2)),
+ (8, _s(64, 8, 3, 10)),
+ (16, _s(64, 16, 3, 8)),
+ (32, _s(64, 32, 3, 8)),
+ ),
+ # attention q (q and its output gate interleaved per head)
+ (12288, 2560): ((2, _d(64, 2)), (4, _s(128, 8, 1, 6)), (8, _s(128, 16, 1, 6))),
+ # attention k / v (2 kv heads x 256) and the MoE router (512 experts)
+ (512, 2560): ((2, _d(64, 2)), (8, _d(64, 1)), (16, _d(32, 1)), (32, _s(64, 8, 4, 10))),
+ # QSA index projection: 4 index heads x 128 + 1 key x 128
+ (640, 2560): ((4, _d(64, 1)), (8, _d(32, 1)), (16, _s(64, 8, 4, 10)), (32, _s(64, 16, 4, 8))),
+ # shared expert gate_up, and the same with its gate row appended (padded to 1296)
+ (1280, 2560): ((2, _d(64, 1)), (4, _d(64, 2)), (8, _s(64, 8, 4, 10))),
+ (1296, 2560): ((2, _d(64, 1)), (4, _d(64, 2)), (8, _s(64, 8, 4, 10))),
+ # shared expert gate [2560 -> 1]
+ (1, 2560): ((32, _d(64, 1)),),
+ # the four GDN input projections concatenated (10240 + 6144 + 48 + 48)
+ (16480, 2560): ((1, _d(64, 2)),),
+ # the attention input projections concatenated: q|gate (12288), k, v (512 each),
+ # index (640), zero-padded to 55 heads of 256 (14080)
+ (14080, 2560): ((2, _d(64, 2)), (8, _s(128, 8, 1, 6))),
+ # PLE value projection (decoder layer 1 only); the PLE key projection is [10240, 2560]
+ (2560, 2560): ((1, _d(64, 1)), (2, _d(64, 2)), (4, _d(32, 2)), (8, _s(64, 8, 2, 10))),
+}
+# Measured and left on cuBLAS: (2560, 640) shared-expert down (the direct kernel
+# cannot tile K=640; split-K and TGV were within 4% of cuBLAS), the GDN concat
+# (16480, 2560) at M >= 2 (direct 3.6-5.3% ahead at M=2 across two sweeps, behind
+# from M=4), the attention concat (14080, 2560) at M >= 16, and M >= 64
+# everywhere (outside both kernels' M <= 32 policy). Not servable at all: the
+# hyper-connection mixer's first projection (fp32 output; both kernels write
+# bf16 only) and its up projection [10240, 320] (K=320 tiles neither kernel).
+
+
+@functools.lru_cache(maxsize=None)
+def plan(m: int, n: int, k: int) -> _Plan:
+ """The kernel and tactic `LowMBf16Mm` launches for an [M, K] x [K, N] product."""
+ for bound, choice in _TABLE.get((n, k), ()):
+ if m <= bound:
+ if choice[0] == "direct":
+ t = _fi_direct.DirectTactic(
+ choice[1], choice[2], _fi_direct.default_tactic(m, n, k).rows_per_block
+ )
+ _fi_direct.validate_tactic(t, m, n, k)
+ return _Plan("direct", tuple(choice[1:]) + (t.rows_per_block,))
+ _fi_splitk.validate_tactic(_fi_splitk.SplitKTactic(*choice[1:]), m, n, k)
+ return _Plan("splitk", tuple(choice[1:]))
+ return _Plan("cublas")
+
+
+def _cell(why: str, m: int, n: int, k: int, kind: str, pdl: bool = False) -> Cell:
+ """`kind` is the branch `plan` must pick for this shape; the test holds it to that."""
+ return Cell(why=why, spec=dict(m=m, n=n, k=k, kind=kind, pdl=pdl))
+
+
+class LowMBf16Mm(OpWrapper):
+ """`mat_a @ mat_b` in bf16 with fp32 accumulation: one kernel launch per call.
+
+ A drop-in for `catalog/gemm/cublas_mm` without bias or `out_dtype`: `mat_a`
+ is dense row-major `[M, K]`, `mat_b` dense column-major `[K, N]` (a stored
+ `[N, K]` weight's `.t()`), the result a fresh row-major `[M, N]` bf16 tensor.
+ `plan(M, N, K)` picks the kernel on the host from three sources:
+
+ - `flashinfer.gemm.kernels.dense_bf16_gemm_direct.run_direct_dense` -- SIMT,
+ the weight row's whole K prefetched into registers; wins at M <= 2 and on
+ narrow N;
+ - `flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk.run_splitk_dense` --
+ tcgen05, swap-AB, split-K reduced across a cluster in shared memory (no
+ extra launch); wins at M 4-32 on the wide shapes;
+ - `torch.ops.trtllm.cublas_mm` -- everything else, including every M > 32.
+
+ `pdl` launches the FlashInfer kernels with programmatic dependent launch:
+ they wait on the previous grid (`griddepcontrol.wait`) before reading
+ `mat_a`, so it is safe after any kernel; it changes nothing on the cuBLAS
+ branch.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # Every cell is a qwen3.8-flash-next-nvfp4 decode projection, at tp1 (from M=2; one row
+ # takes `bf16_gemv_pdl` there) or at a tp4/tep4 rank; `kind` pins the branch `plan`
+ # takes for it. The test's host-only plan pins hold the table's other shapes.
+ CELLS: tuple[Cell, ...] = (
+ _cell(
+ "qwen3.8 decode M=2 GDN in_z: direct, 2-row tile, pdl on",
+ 2,
+ 6144,
+ 2560,
+ "direct",
+ pdl=True,
+ ),
+ _cell(
+ "qwen3.8 decode M=32 o_proj: split-K, mma_n 32, pdl on (M boundary)",
+ 32,
+ 2560,
+ 6144,
+ "splitk",
+ pdl=True,
+ ),
+ _cell("qwen3.8 decode M=128: cuBLAS branch above M=32", 128, 10240, 2560, "cublas"),
+ _cell(
+ "tp4 GDN in_proj M=4 (MTP3 verify, 1 request): split-K, pdl on",
+ 4,
+ 4120,
+ 2560,
+ "splitk",
+ True,
+ ),
+ _cell("tp4 GDN in_proj M=1: direct", 1, 4120, 2560, "direct", pdl=True),
+ _cell("tp4 attention in M=1: direct, 2-output tile, pdl on", 1, 4352, 2560, "direct", True),
+ _cell(
+ "tp4 o / GDN out M=8 (2 requests): split-K 3-way, pdl on", 8, 2560, 1536, "splitk", True
+ ),
+ _cell("router M=1: one decode token, direct 2-row tile", 1, 512, 2560, "direct", True),
+ _cell("router M=64: cuBLAS branch above M=32", 64, 512, 2560, "cublas", True),
+ )
+
+ note = """
+ Metadata consumed: none. Preconditions: bf16 operands, `mat_a` dense
+ row-major, `mat_b` dense column-major, both 32-byte aligned (a freshly
+ allocated activation and a `.t()` weight view are); the FlashInfer kernels
+ raise ValueError on a layout or alignment they cannot read, the cuBLAS
+ branch does not (see `is_valid`).
+
+ Compilation: the FlashInfer kernels are CuTe-DSL, JIT-compiled in-process
+ on the first call per (direct: M, N, K, tactic, pdl; split-K: tactic, pdl)
+ and cached for the process. That first call must happen before CUDA-graph
+ capture -- the engine's warmup forward at each captured batch size does it.
+
+ Host cost: the FlashInfer branches spend tens of microseconds of Python per
+ call (argument validation, DLPack conversion), which a CUDA graph removes
+ and an eager step does not. Route only graph-captured decode calls here.
+
+ Numerics: every branch accumulates in fp32 and rounds once; the branches
+ differ from each other only in summation order.
+
+ Hard limits
+ M FlashInfer branches only at M <= 32 (both kernels' policy) boundary 32 driven at 32
+ (cells) and 128 (cuBLAS cell)
+ N*K largest FlashInfer-branch weight 16480*2560 = 4.2e7 elements boundary 2^31 not probed (48x below it)
+ """
+
+ def __call__(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = False) -> torch.Tensor:
+ m, k = mat_a.shape
+ n = mat_b.shape[1]
+ p = plan(m, n, k)
+ if p.kind == "direct":
+ out = torch.empty((m, n), dtype=mat_a.dtype, device=mat_a.device)
+ return _fi_direct.run_direct_dense(
+ mat_a, mat_b, out, pdl, _fi_direct.DirectTactic(*p.tactic)
+ )
+ if p.kind == "splitk":
+ out = torch.empty((m, n), dtype=mat_a.dtype, device=mat_a.device)
+ return _fi_splitk.run_splitk_dense(
+ mat_a, mat_b, None, out, pdl, _fi_splitk.SplitKTactic(*p.tactic)
+ )
+ return torch.ops.trtllm.cublas_mm(mat_a, mat_b, None, None, 0, None)
+
+ def reference(
+ self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = False
+ ) -> torch.Tensor:
+ """fp32-accumulated, TF32 off; `pdl` is a launch attribute and is ignored."""
+ with true_fp32_matmul():
+ return (mat_a.float() @ mat_b.float()).to(mat_a.dtype)
+
+ def is_valid(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = False) -> None:
+ # The cuBLAS branch ignores both leading dimensions: a row-padded
+ # mat_a or a column-padded mat_b (a slice of a wider tensor) is read as
+ # if dense -- silently a different product. The FlashInfer branches
+ # raise on those. A dense row-major mat_b is outside the contract too
+ # (cuBLAS and the direct kernel raise on it; split-K computes it).
+ assert mat_a.dtype == torch.bfloat16 and mat_b.dtype == torch.bfloat16, "bf16 operands only"
+ assert mat_a.is_contiguous(), "mat_a must be dense row-major [M, K]"
+ assert mat_b.stride(0) == 1 and mat_b.stride(1) == mat_b.shape[0], (
+ "mat_b must be dense column-major [K, N] (e.g. weight.t())"
+ )
+
+
+low_m_bf16_mm = LowMBf16Mm()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/mid_rows_mm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/mid_rows_mm.py
new file mode 100644
index 000000000000..cd5c8a6eb386
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/mid_rows_mm.py
@@ -0,0 +1,1380 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Dense projections of 17..256 decode rows on tcgen05: `y[M, N] = x[M, K] @ W^T`, bf16 out.
+
+Two operand modes, one kernel:
+
+- FP8: `W` is e4m3 codes `[N, K]` with one fp32 scale (or one per weight row); the bf16
+ activation is quantized first by `fp8_row_quant` -- one launch, one e4m3 scale per row
+ (`max |x[m, :]| / 448`) -- and the GEMM applies `scale_x[m] * scale_w` in its epilogue.
+- bf16: `W` is the bf16 weight and `x` is read as it is.
+
+CuTe DSL (tcgen05), swap-AB: the MMA's A operand is 128 weight rows (`BM`), its B the
+activation rows (N = the call's rows rounded up to 32 / 64 / 128 / 256, rows past M
+zero-filled by TMA), so the ring holds mostly weight bytes. One CTA a 128-row weight tile and
+`K / KS` of K; the `KS` CTAs of one tile are a thread-block cluster. A TMA warp streams the
+weight tile and the activation rows into a ring of shared-memory stages (a rolling window of
+weight k-tiles past the ring is prefetched to L2), one MMA warp accumulates in tensor memory,
+eight epilogue warps (two a TMEM lane quadrant, alternate column tiles) read the accumulator.
+With `KS = 1` they scale it, round it to bf16 into shared memory transposed (over the drained
+ring), and every warp stores the tile in 16-byte rows. With `KS > 1` the tile is cut into
+units of 32 weight rows x 4 activation rows owned round-robin by the cluster's ranks: each
+CTA pushes its fp32 partial of a unit into the owner's mailbox (`st.async`, counted on the
+owner's mbarrier), and the owner sums the `KS` partials in rank order -- deterministic -- and
+stores them. The mailbox is shared memory of its own, or the drained ring behind a cluster
+barrier when it is large (256 rows).
+
+Programmatic dependent launch: the weight stages that fit the ring are issued before
+`griddepcontrol.wait` (the weight is static), the activation after it; dependents are released
+at kernel start, and the CTA gives up its tensor-memory allocation permit as soon as it has
+allocated, so a CTA of the next launch is scheduled beside it (two CTAs an SM) and fills its
+own ring meanwhile. `fp8_row_quant` releases its dependents at its start, so the GEMM after it
+launches and streams its weight while the quantization runs.
+"""
+
+from __future__ import annotations
+
+import functools
+from dataclasses import dataclass
+
+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
+from cutlass import const_expr
+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,
+ true_fp32_matmul,
+)
+from tensorrt_llm._torch.kernels.blackwell.custom_pipeline import PipelineTmaUmma, PipelineUmmaAsync
+from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
+
+#: Rows one launch takes: the largest MMA N.
+MAX_ROWS = 256
+E4M3_MAX = 448.0
+#: Floor of a row's maximum: an all-zero row quantizes to zeros at this scale.
+AMAX_FLOOR = 1e-12
+_SMEM_CAP = 232448 - 1024
+
+
+# ------------------------------------------------------------------------------------------
+# inline PTX
+# ------------------------------------------------------------------------------------------
+
+
+def _asm(ret, args, text, cons, loc, ip, side=False):
+ return llvm.inline_asm(
+ ret,
+ [a.ir_value(loc=loc, ip=ip) for a in args],
+ text,
+ cons,
+ has_side_effects=side,
+ is_align_stack=False,
+ asm_dialect=llvm.AsmDialect.AD_ATT,
+ loc=loc,
+ ip=ip,
+ )
+
+
+def _unpack(res, n, ty, loc, ip):
+ mlir_ty = T.i32() if ty is cutlass.Int32 else T.f32()
+ return [ty(llvm.extractvalue(mlir_ty, res, [i], loc=loc, ip=ip)) for i in range(n)]
+
+
+@dsl_user_op
+def _ldg_v4(addr, *, loc=None, ip=None):
+ """16 bytes from global as four words."""
+ res = _asm(
+ llvm.StructType.get_literal([T.i32()] * 4),
+ [cutlass.Int64(addr)],
+ "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];",
+ "=r,=r,=r,=r,l",
+ loc,
+ ip,
+ side=True,
+ )
+ return _unpack(res, 4, cutlass.Int32, loc, ip)
+
+
+@dsl_user_op
+def _ldg_f32(addr, *, loc=None, ip=None):
+ return cutlass.Float32(
+ _asm(T.f32(), [cutlass.Int64(addr)], "ld.global.nc.f32 $0, [$1];", "=f,l", loc, ip, True)
+ )
+
+
+@dsl_user_op
+def _stg_f32(addr, v, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [cutlass.Int64(addr), cutlass.Float32(v)],
+ "{ st.global.f32 [$1], $2; mov.u32 $0, 0; }",
+ "=r,l,f",
+ loc,
+ ip,
+ True,
+ )
+
+
+@dsl_user_op
+def _stg_bf16(addr, v, *, loc=None, ip=None):
+ """`v` rounded to bf16, to global."""
+ _asm(
+ T.i32(),
+ [cutlass.Int64(addr), cutlass.Float32(v)],
+ "{ .reg .b16 h; cvt.rn.bf16.f32 h, $2; st.global.b16 [$1], h; mov.u32 $0, 0; }",
+ "=r,l,f",
+ loc,
+ ip,
+ True,
+ )
+
+
+@dsl_user_op
+def _stg_v2(addr, a, b, *, loc=None, ip=None):
+ _asm(
+ 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,
+ True,
+ )
+
+
+@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,
+ True,
+ )
+
+
+@dsl_user_op
+def _bf16x2(lo, hi, *, loc=None, ip=None):
+ """Two fp32 rounded to bf16, `lo` in the low half."""
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [cutlass.Float32(lo), cutlass.Float32(hi)],
+ "cvt.rn.bf16x2.f32 $0, $2, $1;",
+ "=r,f,f",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _sts_v4f(saddr, a, b, c, d, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [cutlass.Int32(saddr)] + [cutlass.Float32(v) for v in (a, b, c, d)],
+ "{ st.shared.v4.f32 [$1], {$2, $3, $4, $5}; mov.u32 $0, 0; }",
+ "=r,r,f,f,f,f",
+ loc,
+ ip,
+ True,
+ )
+
+
+@dsl_user_op
+def _sts_f32(saddr, v, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [cutlass.Int32(saddr), cutlass.Float32(v)],
+ "{ st.shared.f32 [$1], $2; mov.u32 $0, 0; }",
+ "=r,r,f",
+ loc,
+ ip,
+ True,
+ )
+
+
+@dsl_user_op
+def _sts_bf16(saddr, v, *, loc=None, ip=None):
+ """`v` rounded to bf16, to shared memory."""
+ _asm(
+ T.i32(),
+ [cutlass.Int32(saddr), cutlass.Float32(v)],
+ "{ .reg .b16 h; cvt.rn.bf16.f32 h, $2; st.shared.b16 [$1], h; mov.u32 $0, 0; }",
+ "=r,r,f",
+ loc,
+ ip,
+ True,
+ )
+
+
+@dsl_user_op
+def _lds_v4f(saddr, *, loc=None, ip=None):
+ res = _asm(
+ llvm.StructType.get_literal([T.f32()] * 4),
+ [cutlass.Int32(saddr)],
+ "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];",
+ "=f,=f,=f,=f,r",
+ loc,
+ ip,
+ side=True,
+ )
+ return _unpack(res, 4, cutlass.Float32, loc, ip)
+
+
+@dsl_user_op
+def _lds_v4(saddr, *, loc=None, ip=None):
+ res = _asm(
+ llvm.StructType.get_literal([T.i32()] * 4),
+ [cutlass.Int32(saddr)],
+ "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];",
+ "=r,=r,=r,=r,r",
+ loc,
+ ip,
+ side=True,
+ )
+ return _unpack(res, 4, cutlass.Int32, loc, ip)
+
+
+@dsl_user_op
+def _st_async_v4f(raddr, rbar, a, b, c, d, *, loc=None, ip=None):
+ """Four fp32 into a cluster peer's shared memory (`raddr`, from `mapa`), crediting 16
+ bytes to the peer's mbarrier `rbar`."""
+ _asm(
+ T.i32(),
+ [cutlass.Int32(raddr), cutlass.Int32(rbar)] + [cutlass.Float32(v) for v in (a, b, c, d)],
+ "{ .reg .b32 x, y, z, w; mov.b32 x, $3; mov.b32 y, $4; mov.b32 z, $5; mov.b32 w, $6; "
+ "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [$1], {x, y, z, w}, [$2]; "
+ "mov.u32 $0, 0; }",
+ "=r,r,r,f,f,f,f",
+ loc,
+ ip,
+ True,
+ )
+
+
+@dsl_user_op
+def _mbar_expect_tx(addr, nbytes, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [cutlass.Int32(addr), cutlass.Int32(nbytes)],
+ "{ mbarrier.arrive.expect_tx.shared::cta.b64 _, [$1], $2; mov.u32 $0, 0; }",
+ "=r,r,r",
+ loc,
+ ip,
+ True,
+ )
+
+
+@dsl_user_op
+def _mbar_wait(addr, parity, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [cutlass.Int32(addr), cutlass.Int32(parity)],
+ "{ .reg .pred p; WAIT: mbarrier.try_wait.parity.shared::cta.b64 p, [$1], $2; "
+ "@!p bra WAIT; mov.u32 $0, 0; }",
+ "=r,r,r",
+ loc,
+ ip,
+ 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 _absmax_bf16x2(acc, w, *, loc=None, ip=None):
+ """Per half, max(acc, |w|) over bf16x2 words (non-negative halves)."""
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [cutlass.Int32(acc), cutlass.Int32(w)],
+ "{ .reg .b32 a; and.b32 a, $2, 0x7fff7fff; max.bf16x2 $0, $1, a; }",
+ "=r,r,r",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _row_amax(v, *, loc=None, ip=None):
+ """The warp's largest |x| from per-lane bf16x2 maxima, as fp32: the two halves' larger
+ one (as unsigned bits: the halves are non-negative), then `redux.sync.max.u32`."""
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [cutlass.Int32(v)],
+ "{ .reg .b32 l, h; and.b32 l, $1, 0xffff; shr.b32 h, $1, 16; max.u32 l, l, h; "
+ "redux.sync.max.u32 l, l, 0xffffffff; shl.b32 l, l, 16; mov.b32 $0, l; }",
+ "=f,r",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _e4m3x4(w0, w1, inv, *, loc=None, ip=None):
+ """Four bf16 (two words) times `inv`, each rounded to e4m3 (saturating), packed in order."""
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [cutlass.Int32(w0), cutlass.Int32(w1), cutlass.Float32(inv)],
+ "{ .reg .b32 t0, t1, t2, t3; .reg .f32 f0, f1, f2, f3; .reg .b16 p, q; "
+ "shl.b32 t0, $1, 16; and.b32 t1, $1, 0xffff0000; "
+ "shl.b32 t2, $2, 16; and.b32 t3, $2, 0xffff0000; "
+ "mov.b32 f0, t0; mov.b32 f1, t1; mov.b32 f2, t2; mov.b32 f3, t3; "
+ "mul.rn.f32 f0, f0, $3; mul.rn.f32 f1, f1, $3; "
+ "mul.rn.f32 f2, f2, $3; mul.rn.f32 f3, f3, $3; "
+ "cvt.rn.satfinite.e4m3x2.f32 p, f1, f0; cvt.rn.satfinite.e4m3x2.f32 q, f3, f2; "
+ "mov.b32 $0, {p, q}; }",
+ "=r,r,r,f",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _div_rn(a, b, *, loc=None, ip=None):
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [cutlass.Float32(a), cutlass.Float32(b)],
+ "div.rn.f32 $0, $1, $2;",
+ "=f,f,f",
+ loc,
+ ip,
+ )
+ )
+
+
+# ------------------------------------------------------------------------------------------
+# per-row activation quantization
+# ------------------------------------------------------------------------------------------
+
+
+class _RowQuant:
+ """One warp a row: `scale[m] = max(max |x[m, :]|, AMAX_FLOOR) / 448`,
+ `q[m, k] = e4m3_satfinite(x[m, k] * (448 / max))`. K a multiple of 256."""
+
+ WARPS = 4
+
+ def __init__(self, k: int, pdl: bool):
+ assert k % 256 == 0
+ self.k, self.pdl = k, pdl
+ self.nv = k // 256 # 16-byte vectors a lane
+
+ @cute.jit
+ def __call__(
+ self,
+ m: cutlass.Int32,
+ x_ptr: cute.Pointer,
+ x_stride: cutlass.Int32,
+ q_t: cute.Tensor,
+ s_ptr: cute.Pointer,
+ stream: cuda.CUstream,
+ ):
+ self.kernel(x_ptr.toint(), x_stride, q_t.iterator.toint(), s_ptr.toint(), m).launch(
+ grid=[(m + self.WARPS - 1) // self.WARPS, 1, 1],
+ block=[32 * self.WARPS, 1, 1],
+ stream=stream,
+ use_pdl=self.pdl,
+ )
+
+ @cute.kernel
+ def kernel(
+ self,
+ x_base: cutlass.Int64,
+ x_stride: cutlass.Int32,
+ q_base: cutlass.Int64,
+ s_base: cutlass.Int64,
+ m: cutlass.Int32,
+ ):
+ bidx, _, _ = cute.arch.block_idx()
+ warp = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+ lane = cute.arch.lane_idx()
+ if const_expr(self.pdl):
+ cute.arch.griddepcontrol_launch_dependents()
+ cute.arch.griddepcontrol_wait()
+ row = bidx * self.WARPS + warp
+ if row < m:
+ xr = (
+ x_base + cutlass.Int64(row) * cutlass.Int64(x_stride) * 2 + cutlass.Int64(lane * 16)
+ )
+ vals = [_ldg_v4(xr + v * 512) for v in range(self.nv)]
+ mx = cutlass.Int32(0)
+ for v in cutlass.range_constexpr(self.nv):
+ for j in cutlass.range_constexpr(4):
+ mx = _absmax_bf16x2(mx, vals[v][j])
+ amax = cutlass.max(_row_amax(mx), cutlass.Float32(AMAX_FLOOR))
+ inv = _div_rn(cutlass.Float32(E4M3_MAX), amax)
+ qr = q_base + cutlass.Int64(row) * self.k + cutlass.Int64(lane * 8)
+ for v in cutlass.range_constexpr(self.nv):
+ _stg_v2(
+ qr + v * 256,
+ _e4m3x4(vals[v][0], vals[v][1], inv),
+ _e4m3x4(vals[v][2], vals[v][3], inv),
+ )
+ if lane == 0:
+ _stg_f32(s_base + cutlass.Int64(row) * 4, _div_rn(amax, cutlass.Float32(E4M3_MAX)))
+
+
+# ------------------------------------------------------------------------------------------
+# GEMM
+# ------------------------------------------------------------------------------------------
+
+#: Weight rows one CTA computes: the MMA's M.
+BM = 128
+
+
+@dataclass(frozen=True)
+class Plan:
+ """`ks` K split (cluster size); `kpb` MMA k-blocks a ring stage (4: 128 bytes of K a row);
+ `stages` ring depth; `l2` 1 to prefetch the CTA's weight slice past the ring into L2 at
+ the start. A CTA computes `BM` weight rows for every activation row."""
+
+ ks: int = 1
+ kpb: int = 4
+ stages: int = 4
+ l2: int = 1
+ epi: int = 16
+
+
+def mma_rows(m: int) -> int:
+ """The MMA N (activation rows, zero-padded by TMA) a call of `m` rows runs at."""
+ for b in (32, 64, 128, 256):
+ if m <= b:
+ return b
+ raise ValueError(f"{m} rows: at most {MAX_ROWS}")
+
+
+class _MidMm:
+ """See the module docstring. Swap-AB: the MMA's A operand is `BM` weight rows, its B the
+ activation rows (N = `mpad`). Compile-time: N, K, the plan, `mpad`, the operand mode and
+ the weight-scale mode (0 none / 1 one fp32 / 2 one a weight row). Runtime: the row count
+ (<= mpad) and the pointers."""
+
+ def __init__(self, *, n, k, plan: Plan, mpad, fp8, w_mode, pdl):
+ epi_cols = min(plan.epi, mpad)
+ assert mpad % 16 == 0 and 16 <= mpad <= 256 and mpad % epi_cols == 0
+ self.n, self.k, self.plan, self.mpad, self.fp8, self.w_mode, self.pdl = (
+ n,
+ k,
+ plan,
+ mpad,
+ fp8,
+ w_mode,
+ pdl,
+ )
+ self.ks = plan.ks
+ self.io_dtype = cutlass.Float8E4M3FN if fp8 else cutlass.BFloat16
+ self.acc_dtype = cutlass.Float32
+ self.epi_cols = epi_cols
+ self.cta_group = tcgen05.CtaGroup.ONE
+ self.cluster_shape_mn = (1, 1)
+ # Two warps a TMEM lane quadrant: warpgroup h reads the column tiles `sub % 2 == h`.
+ self.epilogue_warp_id = tuple(range(8))
+ self.mma_warp_id = 8
+ self.tma_warp_id = 9
+ self.threads = 32 * 10
+ self.epilog_sync_bar_id = 1
+ self.tmem_alloc_sync_bar_id = 2
+ self.n_tiles = -(-n // BM)
+ # The tile staged for the stores, `[mpad][BM]` bf16 (KS = 1, over the drained ring).
+ self.stage_bytes = mpad * BM * 2 if plan.ks == 1 else 0
+ # The split-K mailbox (KS > 1): the most units a rank owns, KS slots of 512 bytes each;
+ # a large one reuses the drained ring (after a cluster barrier: every peer's mainloop is
+ # done before any partial lands).
+ mbox = -(-mpad // plan.ks) * plan.ks * 512 if plan.ks > 1 else 0
+ self.alias_mbox = mbox > (64 << 10)
+ self.mbox_bytes = 0 if self.alias_mbox else mbox
+ self.mbox_need = mbox
+
+ def _setup(self):
+ k_major = utils.LayoutEnum.ROW_MAJOR.mma_major_mode()
+ tiled_mma = utils.sm100.make_trivial_tiled_mma(
+ self.io_dtype,
+ self.io_dtype,
+ k_major,
+ k_major,
+ self.acc_dtype,
+ self.cta_group,
+ (BM, self.mpad),
+ )
+ inst_k = cute.size(tiled_mma.shape_mnk, mode=[2])
+ self.mma_tiler = (BM, self.mpad, inst_k * self.plan.kpb)
+ 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 = (BM, self.epi_cols)
+ assert self.k % (self.mma_tiler[2] * self.ks) == 0, (self.k, self.mma_tiler, self.ks)
+ self.k_tiles = self.k // self.mma_tiler[2] // self.ks
+ self.stages = self._stages()
+ self.n_pre = self.stages
+ self.a_smem_layout_staged = utils.sm100.make_smem_layout_a(
+ tiled_mma, self.mma_tiler, self.io_dtype, self.stages
+ )
+ self.b_smem_layout_staged = utils.sm100.make_smem_layout_b(
+ tiled_mma, self.mma_tiler, self.io_dtype, self.stages
+ )
+ acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
+ acc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, 1))
+ self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(acc_fake, arch="sm_100")
+ return tiled_mma
+
+ def _stages(self) -> int:
+ """The plan's ring depth, at most the CTA's K tiles and what fits beside the mailbox."""
+ bpe = 1 if self.fp8 else 2
+ inst_k = 32 if self.fp8 else 16
+ kb = inst_k * self.plan.kpb * bpe
+ k_tiles = self.k // (inst_k * self.plan.kpb) // self.ks
+ fit = (_SMEM_CAP - self.mbox_bytes - 4096) // ((BM + self.mpad) * kb)
+ return max(1, min(self.plan.stages, k_tiles, fit))
+
+ def smem_bytes(self) -> int:
+ """Host-side estimate of the dynamic shared memory (ring or staged tile, mailbox)."""
+ bpe = 1 if self.fp8 else 2
+ inst_k = 32 if self.fp8 else 16
+ ring = self._stages() * (BM + self.mpad) * inst_k * self.plan.kpb * bpe
+ return max(ring, self.stage_bytes, self.mbox_need if self.alias_mbox else 0) + (
+ self.mbox_bytes + 1024 + 2048
+ )
+
+ @cute.jit
+ def __call__(
+ self,
+ m: cutlass.Int32,
+ x_ptr: cute.Pointer,
+ w_ptr: cute.Pointer,
+ out_t: cute.Tensor,
+ sx_ptr: cute.Pointer,
+ sw_ptr: cute.Pointer,
+ stream: cuda.CUstream,
+ ):
+ k = self.k
+ a = cute.make_tensor(w_ptr, cute.make_layout((self.n, k, 1), stride=(k, 1, 0)))
+ b = cute.make_tensor(x_ptr, cute.make_layout((m, k, 1), stride=(k, 1, 0)))
+ tiled_mma = self._setup()
+ 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
+ )
+ kern = self.kernel(
+ tiled_mma,
+ tma_atom_a,
+ tma_a,
+ tma_atom_b,
+ tma_b,
+ out_t.iterator.toint(),
+ w_ptr.toint(),
+ sx_ptr.toint(),
+ sw_ptr.toint(),
+ m,
+ self.cluster_layout_vmnk,
+ self.a_smem_layout_staged,
+ self.b_smem_layout_staged,
+ self.epi_tile,
+ )
+ # All of the SM's L1 as shared memory: a CTA of the next launch then fits beside this
+ # one (programmatic dependent launch) and streams its weight stages early.
+ if const_expr(self.ks > 1):
+ kern.launch(
+ grid=[self.ks, self.n_tiles, 1],
+ block=[self.threads, 1, 1],
+ cluster=[self.ks, 1, 1],
+ stream=stream,
+ use_pdl=self.pdl,
+ preferred_smem_carveout=100,
+ )
+ else:
+ kern.launch(
+ grid=[1, self.n_tiles, 1],
+ block=[self.threads, 1, 1],
+ stream=stream,
+ use_pdl=self.pdl,
+ preferred_smem_carveout=100,
+ )
+
+ @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,
+ out_base: cutlass.Int64,
+ w_base: cutlass.Int64,
+ sx_base: cutlass.Int64,
+ sw_base: cutlass.Int64,
+ 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())
+ if const_expr(self.pdl):
+ 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 # weight-row tile
+ ks: cutlass.Constexpr = self.ks
+ mpad: cutlass.Constexpr = self.mpad
+ n0 = nt * BM
+
+ @cute.struct
+ class SharedStorage:
+ ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.stages * 2]
+ acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2]
+ tmem_dealloc_mbar: cutlass.Int64
+ red_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.stages,
+ 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=1,
+ 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=1024,
+ swizzle=a_smem_layout_staged.inner,
+ )
+ sB = smem.allocate_tensor(
+ element_type=self.io_dtype,
+ layout=b_smem_layout_staged.outer,
+ byte_alignment=1024,
+ swizzle=b_smem_layout_staged.inner,
+ )
+ ring_bytes = self.stages * (BM + mpad) * self.mma_tiler[2] * (1 if self.fp8 else 2)
+ if const_expr(self.stage_bytes > ring_bytes):
+ smem.allocate_tensor(
+ element_type=cutlass.Float32,
+ layout=cute.make_layout(((self.stage_bytes - ring_bytes) // 4 + 4,)),
+ byte_alignment=16,
+ )
+ sXt = smem.allocate_tensor(
+ element_type=cutlass.Float32, layout=cute.make_layout((256,)), byte_alignment=16
+ )
+ sx_s = cutlass.Int32(sXt.iterator.toint())
+ red_bar = cutlass.Int32(storage.red_mbar.ptr.toint())
+ mbox = cutlass.Int32(0)
+ if const_expr(ks > 1):
+ # The split-K mailbox: units of 32 weight rows x 4 activation rows (fp32, 512
+ # bytes), unit u owned by rank u % KS; slot (u // KS, sender) of the owner.
+ if const_expr(self.alias_mbox):
+ ring_b = self.stages * (BM + mpad) * self.mma_tiler[2] * (1 if self.fp8 else 2)
+ if const_expr(self.mbox_need > ring_b):
+ smem.allocate_tensor(
+ element_type=cutlass.Float32,
+ layout=cute.make_layout(((self.mbox_need - ring_b) // 4 + 4,)),
+ byte_alignment=16,
+ )
+ mbox = cutlass.Int32(sA.iterator.toint())
+ else:
+ sMb = smem.allocate_tensor(
+ element_type=cutlass.Float32,
+ layout=cute.make_layout((self.mbox_bytes // 4,)),
+ byte_alignment=16,
+ )
+ mbox = cutlass.Int32(sMb.iterator.toint())
+ if tidx == 0:
+ cute.arch.mbarrier_init(storage.red_mbar.ptr, 1)
+ cute.arch.mbarrier_init_fence()
+ # Peers st.async into this mailbox only after this CTA's barrier is initialized.
+ cute.arch.cluster_arrive_relaxed()
+ # The output tile (bf16, KS = 1) or the fp32 partial (KS > 1) reuses the ring (drained
+ # once the accumulator is complete): it starts at the A stages, the B stages follow.
+ p_base = cutlass.Int32(sA.iterator.toint())
+
+ 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, 1))
+
+ pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
+ if const_expr(ks > 1):
+ cute.arch.cluster_wait()
+
+ if warp_idx == self.tma_warp_id:
+ tAgA_slice = tAgA[(None, nt, None, 0)]
+ tBgB_slice = tBgB[(None, 0, None, 0)]
+ ab_producer.reset()
+ # The weight does not depend on the previous kernel: the ring's stages of it
+ # stream before the grid-dependency wait, the activation rows after it.
+ pre = []
+ for kk in cutlass.range_constexpr(self.n_pre):
+ handle = ab_producer.acquire_and_advance()
+ cute.copy(
+ tma_atom_a,
+ tAgA_slice[(None, k0 + kk)],
+ tAsA[(None, handle.index)],
+ tma_bar_ptr=handle.barrier,
+ )
+ pre.append(handle)
+ l2a: cutlass.Constexpr = min(self.plan.l2, self.k_tiles - self.n_pre)
+ if const_expr(l2a > 0):
+ # A rolling window of `l2` weight k-tiles past the ring prefetched towards L2
+ # (TMA prefetch): the HBM stream does not wait for ring slots.
+ for kk in cutlass.range_constexpr(self.n_pre, self.n_pre + l2a):
+ cute.prefetch(tma_atom_a, tAgA_slice[(None, k0 + kk)])
+ if const_expr(self.pdl):
+ cute.arch.griddepcontrol_wait()
+ for kk in cutlass.range_constexpr(self.n_pre):
+ handle = pre[kk]
+ cute.copy(
+ tma_atom_b,
+ tBgB_slice[(None, k0 + kk)],
+ tBsB[(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()
+ if const_expr(l2a > 0):
+ if k_tile + l2a < self.k_tiles:
+ cute.prefetch(tma_atom_a, tAgA_slice[(None, k0 + k_tile + l2a)])
+ 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 const_expr(ks > 1 and self.alias_mbox):
+ cute.arch.cluster_arrive_relaxed()
+ cute.arch.cluster_wait()
+ 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)
+ tCtAcc = tCtAcc_base[(None, None, None, 0)]
+ acc_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1)
+ ab_consumer.reset()
+ peek = ab_consumer.try_wait()
+ acc_pipeline.producer_acquire(acc_producer_state)
+ num_kblocks = cute.size(tCrA, mode=[2])
+ for k_tile in cutlass.range(0, self.k_tiles, 1, unroll=1):
+ handle = ab_consumer.wait_and_advance(peek)
+ for kb in cutlass.range_constexpr(num_kblocks):
+ if const_expr(kb == 0):
+ tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0)
+ else:
+ tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+ cute.gemm(
+ tiled_mma,
+ tCtAcc,
+ tCrA[(None, None, kb, handle.index)],
+ tCrB[(None, None, kb, handle.index)],
+ tCtAcc,
+ )
+ 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()
+ if const_expr(ks > 1 and self.alias_mbox):
+ cute.arch.cluster_arrive_relaxed()
+ acc_pipeline.producer_tail(acc_producer_state)
+ if const_expr(ks > 1 and self.alias_mbox):
+ cute.arch.cluster_wait()
+
+ if warp_idx < self.mma_warp_id:
+ tmem.allocate(self.num_tmem_alloc_cols)
+ # This CTA allocates no more tensor memory: a CTA of the next launch may be
+ # scheduled on this SM now (programmatic dependent launch).
+ tmem.relinquish_alloc_permit()
+ if const_expr(self.fp8):
+ # The activation scales (the previous launch's output) into shared memory.
+ if const_expr(self.pdl):
+ cute.arch.griddepcontrol_wait()
+ if tidx < m:
+ _sts_f32(sx_s + tidx * 4, _ldg_f32(sx_base + cutlass.Int64(tidx) * 4))
+ 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, 1)
+ 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 % BM)
+ half = tidx // BM
+ tTR_tAcc_base = thr_copy_t2r.partition_S(tAcc_epi)
+ tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, 0)]
+ tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
+ # This thread's accumulator row (a weight row) and first column (an activation
+ # row): the 32dp32b load gives a thread one row's `epi_cols` consecutive columns.
+ tCcC = thr_mma.partition_C(cute.make_identity_tensor((BM, mpad)))
+ tTR_cC = thr_copy_t2r.partition_D(
+ cute.flat_divide(tCcC[((None, None), 0, 0)], epi_tile)
+ )
+ frag_c = tTR_cC[(None, None, None, 0, 0)]
+ crd0 = frag_c[0]
+ r_loc = crd0[0]
+ c_loc = crd0[1]
+ tTR_rAcc = cute.make_rmem_tensor(frag_c.shape, self.acc_dtype)
+ nfrag: cutlass.Constexpr = self.epi_cols
+ n = n0 + r_loc
+ sw = cutlass.Float32(1.0)
+ if const_expr(ks == 1):
+ if const_expr(self.w_mode == 1):
+ sw = _ldg_f32(sw_base)
+ if const_expr(self.w_mode == 2):
+ if n < self.n:
+ sw = _ldg_f32(sw_base + cutlass.Int64(n) * 4)
+ cute.arch.barrier(
+ barrier_id=self.epilog_sync_bar_id,
+ number_of_threads=32 * len(self.epilogue_warp_id),
+ )
+ vc = (m + 3) // 4 # 4-row chunks of activation rows that hold real rows
+ n_own = (4 * vc - kr + ks - 1) // ks # units this rank owns
+ if const_expr(ks > 1):
+ if tidx == 0:
+ _mbar_expect_tx(red_bar, n_own * (ks - 1) * 512)
+ acc_pipeline.consumer_wait(acc_consumer_state)
+ if const_expr(ks > 1 and self.alias_mbox):
+ cute.arch.cluster_arrive_relaxed()
+ cute.arch.cluster_wait()
+ # The tile goes to shared memory transposed, `[activation row][weight row]` (over
+ # the drained ring): bf16 and scaled (KS = 1) or the fp32 partial (KS > 1).
+ for sub in cutlass.range_constexpr(mpad // self.epi_cols):
+ c0 = sub * self.epi_cols + c_loc
+ if (c0 < m) & (half == sub % 2):
+ cute.copy(tiled_copy_t2r, tTR_tAcc[(None, None, None, sub)], tTR_rAcc)
+ for q in cutlass.range_constexpr(nfrag // 4):
+ if const_expr(ks == 1):
+ if const_expr(self.fp8):
+ s4 = _lds_v4f(sx_s + (c0 + 4 * q) * 4)
+ for e in cutlass.range_constexpr(4):
+ i = 4 * q + e
+ v = tTR_rAcc[i] * sw
+ if const_expr(self.fp8):
+ v = v * s4[e]
+ _sts_bf16(p_base + ((c0 + i) * BM + r_loc) * 2, v)
+ else:
+ cq = (c0 // 4) + q
+ if cq * 4 < m:
+ u = (r_loc // 32) * vc + cq
+ owner = u % ks
+ slot = mbox + (((u // ks) * ks + kr) * 32 + r_loc % 32) * 16
+ if owner == kr:
+ _sts_v4f(
+ slot,
+ tTR_rAcc[4 * q],
+ tTR_rAcc[4 * q + 1],
+ tTR_rAcc[4 * q + 2],
+ tTR_rAcc[4 * q + 3],
+ )
+ else:
+ _st_async_v4f(
+ _mapa(slot, owner),
+ _mapa(red_bar, owner),
+ tTR_rAcc[4 * q],
+ tTR_rAcc[4 * q + 1],
+ tTR_rAcc[4 * q + 2],
+ tTR_rAcc[4 * q + 3],
+ )
+ 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.free(tmem_ptr)
+ # The tile's stores (KS = 1): an item is 8 weight rows of one activation row (16 bytes
+ # out), consecutive threads take consecutive items (coalesced).
+ g8: cutlass.Constexpr = BM // 8
+ if const_expr(ks == 1):
+ cute.arch.barrier()
+ for f in cutlass.range(tidx, m * g8, self.threads, unroll=1):
+ c = f // g8
+ g = f - c * g8
+ n2 = n0 + 8 * g
+ if n2 < self.n:
+ w4 = _lds_v4(p_base + (c * BM + 8 * g) * 2)
+ _stg_v4(
+ out_base + (cutlass.Int64(c) * self.n + cutlass.Int64(n2)) * 2,
+ w4[0],
+ w4[1],
+ w4[2],
+ w4[3],
+ )
+ else:
+ # The owner sums its units' KS partials (its own and the peers', which arrive on
+ # the mailbox barrier) in rank order and stores them: an item is one unit's weight
+ # row (lane) and its 4 activation rows; consecutive threads take consecutive rows.
+ if warp_idx < self.mma_warp_id:
+ _mbar_wait(red_bar, 0)
+ sw2 = cutlass.Float32(1.0)
+ if const_expr(self.w_mode == 1):
+ sw2 = _ldg_f32(sw_base)
+ vc2 = (m + 3) // 4
+ n_own2 = (4 * vc2 - kr + ks - 1) // ks
+ for f in cutlass.range(
+ tidx, n_own2 * 32, 32 * len(self.epilogue_warp_id), unroll=1
+ ):
+ j = f // 32
+ ln = f - j * 32
+ u = kr + j * ks
+ w_ = u // vc2
+ cq = u - w_ * vc2
+ n2 = n0 + w_ * 32 + ln
+ if n2 < self.n:
+ base = mbox + ((j * ks) * 32 + ln) * 16
+ parts = [_lds_v4f(base + s_ * 512) for s_ in range(ks)]
+ acc = parts[0]
+ for s_ in cutlass.range_constexpr(1, ks):
+ acc = [acc[e] + parts[s_][e] for e in range(4)]
+ sc = sw2
+ if const_expr(self.w_mode == 2):
+ sc = _ldg_f32(sw_base + cutlass.Int64(n2) * 4)
+ if const_expr(self.fp8):
+ s4 = _lds_v4f(sx_s + cq * 16)
+ for e in cutlass.range_constexpr(4):
+ v = acc[e] * sc
+ if const_expr(self.fp8):
+ v = v * s4[e]
+ if cq * 4 + e < m:
+ _stg_bf16(
+ out_base
+ + (cutlass.Int64(cq * 4 + e) * self.n + cutlass.Int64(n2)) * 2,
+ v,
+ )
+
+
+# ------------------------------------------------------------------------------------------
+# host side
+# ------------------------------------------------------------------------------------------
+
+
+def _env_pdl() -> bool:
+ try:
+ from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+ except ImportError:
+ return True
+ return bool(get_env_enable_pdl())
+
+
+def _stream():
+ return cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+
+
+@functools.lru_cache(maxsize=None)
+def _compiled_quant(k: int, pdl: bool):
+ kern = _RowQuant(k, pdl)
+ dev = "cuda"
+ x = torch.empty(4, k, dtype=torch.bfloat16, device=dev)
+ q = torch.empty(4, k, dtype=torch.uint8, device=dev)
+ s = torch.empty(4, dtype=torch.float32, device=dev)
+ return cute.compile(
+ kern,
+ cutlass.Int32(4),
+ make_ptr(cutlass.BFloat16, x.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
+ cutlass.Int32(k),
+ from_dlpack(q, assumed_align=16).mark_layout_dynamic(leading_dim=1),
+ make_ptr(cutlass.Float32, s.data_ptr(), cute.AddressSpace.gmem, assumed_align=4),
+ _stream(),
+ options="--opt-level 2 --enable-tvm-ffi",
+ )
+
+
+@functools.lru_cache(maxsize=None)
+def _compiled_mm(n: int, k: int, plan: Plan, mpad: int, fp8: bool, w_mode: int, pdl: bool):
+ kern = _MidMm(n=n, k=k, plan=plan, mpad=mpad, fp8=fp8, w_mode=w_mode, pdl=pdl)
+ assert kern.smem_bytes() <= _SMEM_CAP, (plan, kern.smem_bytes())
+ dev = "cuda"
+ dt = cutlass.Float8E4M3FN if fp8 else cutlass.BFloat16
+ x = torch.empty(16, dtype=torch.uint8, device=dev)
+ out = torch.empty(4, n, dtype=torch.bfloat16, device=dev)
+ s = torch.empty(4, dtype=torch.float32, device=dev)
+ return cute.compile(
+ kern,
+ cutlass.Int32(4),
+ make_ptr(dt, x.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
+ make_ptr(dt, x.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
+ from_dlpack(out, assumed_align=16).mark_layout_dynamic(leading_dim=1),
+ make_ptr(cutlass.Float32, s.data_ptr(), cute.AddressSpace.gmem, assumed_align=4),
+ make_ptr(cutlass.Float32, s.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
+ _stream(),
+ options="--opt-level 2 --enable-tvm-ffi",
+ )
+
+
+def fp8_row_quant(x: torch.Tensor, pdl: bool = True) -> tuple[torch.Tensor, torch.Tensor]:
+ """`(q [M, K] e4m3 contiguous, scale [M] fp32)` of a bf16 `[M, K]` activation (unit column
+ stride, 16-byte aligned rows): `scale = max(max |x[m]|, AMAX_FLOOR) / 448`,
+ `q = e4m3_satfinite(x * (448 / max))`."""
+ m, k = x.shape
+ assert x.dtype == torch.bfloat16 and x.stride(1) == 1 and k % 256 == 0
+ assert x.stride(0) % 8 == 0 and x.data_ptr() % 16 == 0
+ q = torch.empty((m, k), dtype=torch.float8_e4m3fn, device=x.device)
+ # The GEMM reads the scales of its whole MMA tile of rows (16-byte vectors): padded.
+ s = torch.empty(mma_rows(m) if m <= MAX_ROWS else m, dtype=torch.float32, device=x.device)
+ if m:
+ _compiled_quant(k, bool(pdl and _env_pdl()))(
+ m, x.data_ptr(), x.stride(0), q.view(torch.uint8), s.data_ptr()
+ )
+ return q, s[:m]
+
+
+def mid_mm_raw(
+ x: torch.Tensor,
+ w: torch.Tensor,
+ plan: Plan,
+ sx: torch.Tensor | None = None,
+ sw: torch.Tensor | None = None,
+ pdl: bool = True,
+) -> torch.Tensor:
+ """The GEMM alone: `x [M, K]` (e4m3 with `sx [M]`, or bf16) contiguous, `w [N, K]` of the
+ same dtype contiguous, `sw` None / fp32 [1] / fp32 [N]. Returns a fresh bf16 `[M, N]`."""
+ m, k = x.shape
+ n = w.shape[0]
+ fp8 = x.dtype == torch.float8_e4m3fn
+ assert w.dtype == x.dtype and w.shape[1] == k and x.is_contiguous() and w.is_contiguous()
+ assert 1 <= m <= MAX_ROWS and n % 16 == 0
+ assert fp8 == (sx is not None)
+ w_mode = 0 if sw is None else (1 if sw.numel() == 1 else 2)
+ out = torch.empty((m, n), dtype=torch.bfloat16, device=x.device)
+ fn = _compiled_mm(n, k, plan, mma_rows(m), fp8, w_mode, bool(pdl and _env_pdl()))
+ dummy = out.data_ptr()
+ fn(
+ m,
+ x.data_ptr(),
+ w.data_ptr(),
+ out,
+ dummy if sx is None else sx.data_ptr(),
+ dummy if sw is None else sw.data_ptr(),
+ )
+ return out
+
+
+# ------------------------------------------------------------------------------------------
+# plans, weights, entry
+# ------------------------------------------------------------------------------------------
+
+
+@dataclass
+class MidWeight:
+ """A projection weight as `mid_rows_mm` reads it: `w` `[N, K]` contiguous, e4m3 codes with
+ `scale` (fp32 `[1]`, or `[N]` one a weight row) or bf16 with `scale` None."""
+
+ w: torch.Tensor
+ scale: torch.Tensor | None = None
+
+ @property
+ def n(self) -> int:
+ return self.w.shape[0]
+
+ @property
+ def k(self) -> int:
+ return self.w.shape[1]
+
+ @property
+ def fp8(self) -> bool:
+ return self.w.dtype == torch.float8_e4m3fn
+
+
+#: Measured plans (GB300, 152 SMs, CUDA-graph chain with PDL over HBM-cold weights; the
+#: campaign's bench/gemm_mid) of the qwen3.8-flash-next-nvfp4 tp1 projections, by MMA rows:
+#: (N, K, fp8) -> ((largest rows, Plan), ...). The ring depth is clamped to what fits.
+_TABLE: dict[tuple[int, int, bool], tuple[tuple[int, Plan], ...]] = {
+ # GDN in_proj (q|k|v|z|a|b)
+ (16480, 2560, True): ((256, Plan(1, 4, 8, 2)),),
+ # attention in-proj (q|gate|k|v|index)
+ (14080, 2560, True): ((256, Plan(1, 4, 5, 2)),),
+ # GDN out / attention o
+ (2560, 6144, True): ((256, Plan(6, 4, 6, 2)),),
+ # shared expert gate|up, down; the router (bf16)
+ (1280, 2560, False): ((256, Plan(10, 4, 4, 0)),),
+ (2560, 640, False): ((256, Plan(5, 4, 2, 0)),),
+ (512, 2560, False): ((256, Plan(10, 4, 4, 0)),),
+ # vocab head (per-row e4m3 weight)
+ (248320, 2560, True): ((64, Plan(1, 4, 4, 0)), (256, Plan(1, 4, 6, 2))),
+}
+
+
+def default_plan(n: int, k: int, m: int, fp8: bool) -> Plan:
+ """The table's plan for (N, K, operand mode) at `m` rows, else a heuristic: split over K
+ until about one CTA an SM, as deep a ring as fits."""
+ for bound, plan in _TABLE.get((n, k, fp8), ()):
+ if m <= bound:
+ return plan
+ kpb = 4
+ kt = k // ((32 if fp8 else 16) * kpb)
+ tiles = -(-n // BM)
+ ks = 1
+ for d in range(1, 17):
+ if kt % d == 0 and tiles * d <= 152:
+ ks = d
+ stage = (BM + mma_rows(m)) * 128
+ return Plan(ks, kpb, max(2, min(10, (200 << 10) // stage)))
+
+
+def mid_rows_mm(x: torch.Tensor, mw: MidWeight, plan: Plan | None = None, pdl: bool = True):
+ """`x [M, K]` bf16 (unit column stride, 16-byte aligned rows) @ `mw`^T -> fresh bf16
+ `[M, N]`, M = 1 .. MAX_ROWS. An e4m3 weight quantizes `x` per row first
+ (`fp8_row_quant`)."""
+ m = x.shape[0]
+ if plan is None:
+ plan = default_plan(mw.n, mw.k, m, mw.fp8)
+ if mw.fp8:
+ q, sx = fp8_row_quant(x, pdl)
+ return mid_mm_raw(q, mw.w, plan, sx, mw.scale, pdl)
+ if not x.is_contiguous():
+ x = x.contiguous()
+ return mid_mm_raw(x, mw.w, plan, None, None, pdl)
+
+
+def precompile(mw: MidWeight, lo: int = 1, hi: int = MAX_ROWS, pdl: bool = True) -> int:
+ """Compile every (plan, MMA rows) `mid_rows_mm` reaches over `lo`..`hi` rows of this
+ weight's geometry now rather than at a first serving call; returns how many."""
+ pdl = bool(pdl and _env_pdl())
+ w_mode = 0 if mw.scale is None else (1 if mw.scale.numel() == 1 else 2)
+ keys = {(default_plan(mw.n, mw.k, m, mw.fp8), mma_rows(m)) for m in range(lo, hi + 1)}
+ for plan, mpad in sorted(keys, key=str):
+ _compiled_mm(mw.n, mw.k, plan, mpad, mw.fp8, w_mode, pdl)
+ if mw.fp8:
+ _compiled_quant(mw.k, pdl)
+ return len(keys)
+
+
+# ------------------------------------------------------------------------------------------
+# op wrappers
+# ------------------------------------------------------------------------------------------
+
+
+def _mm_cell(why: str, n: int, k: int, m: int, mode: str) -> Cell:
+ """`mode`: "tensor" (e4m3, one weight scale), "row" (e4m3, one scale a weight row), or
+ "bf16"."""
+ return Cell(why=why, spec=dict(n=n, k=k, m=m, mode=mode))
+
+
+class MidRowsMm(OpWrapper):
+ """`mid_rows_mm(x, mw)`: `y[M, N] = x[M, K] @ W^T` in bf16 for 1..MAX_ROWS rows on tcgen05.
+ With an e4m3 weight (`MidWeight(codes [N, K], scale)`, one fp32 scale or one a weight row)
+ the activation is quantized to e4m3 with one scale a row (`fp8_row_quant`, its own launch)
+ and the product is `(q(x) * scale_x[m]) @ (codes * scale_w)^T`; with a bf16 weight it is
+ the plain product. fp32 accumulation (tensor memory; split-K partials summed in rank
+ order), the scales applied in fp32, one bf16 rounding.
+
+ Lossy (e4m3 operands) against the bf16 projection in the FP8 mode: ~3% relative error per
+ output on random operands; a row's activation rounding depends only on its own maximum.
+
+ Fusion boundary: activation quantization + GEMM (two launches) / the GEMM (one)."""
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ _mm_cell("GDN in_proj, STP con32 (32 rows), e4m3", 16480, 2560, 32, "tensor"),
+ _mm_cell(
+ "attention in-proj, MTP3 con32 verify (128 rows), e4m3", 14080, 2560, 128, "tensor"
+ ),
+ _mm_cell("GDN in_proj, 256 rows (MMA N 256), e4m3", 16480, 2560, 256, "tensor"),
+ _mm_cell("GDN out / attention o, 17 rows, split-K, e4m3", 2560, 6144, 17, "tensor"),
+ _mm_cell(
+ "GDN out / attention o, 200 rows, split-K, mailbox over the ring, e4m3",
+ 2560,
+ 6144,
+ 200,
+ "tensor",
+ ),
+ _mm_cell("router, 128 rows, bf16", 512, 2560, 128, "bf16"),
+ _mm_cell("shared expert gate|up, 64 rows, bf16", 1280, 2560, 64, "bf16"),
+ _mm_cell("shared expert down, 256 rows, bf16", 2560, 640, 256, "bf16"),
+ _mm_cell("vocab head, per-row e4m3 weight, 32 rows", 248320, 2560, 32, "row"),
+ )
+
+ note = """
+ Tier 6 (CuTe DSL written for this target): at 17-256 decode rows cuBLAS(Lt) reaches
+ ~0.4-0.6 of HBM bandwidth on these weights, and the per-tensor FP8 entry adds two
+ quantization launches.
+ Metadata consumed: none. Call sites: the decode target's dense projections in their
+ `_MID_MM_ROWS` row ranges (`_dense_mm`), the router and shared expert, the FP8 vocab
+ heads above `fp8_gemv`'s rows.
+ Preconditions: `x` bf16 `[M, K]`, 1 <= M <= 256, unit column stride; the FP8 mode reads
+ any row stride that is a multiple of 8 elements (16-byte aligned rows), the bf16 mode a
+ contiguous `x` (others are copied). `mw.w` `[N, K]` contiguous, N a multiple of 16, K a
+ multiple of 128 (e4m3) / 64 (bf16) times the plan's K split.
+ Launch: grid (KS, ceil(N / 128)), cluster (KS, 1, 1) when KS > 1, 10 warps (8 epilogue,
+ MMA, TMA), shared memory up to ~110 KB at 32 rows so two CTAs fit an SM.
+ Programmatic dependent launch (TRTLLM_ENABLE_PDL=0 turns it off): the weight stages that
+ fit the ring stream before `griddepcontrol.wait`, the activation after it; dependents are
+ released at kernel start (the quantization launch releases the GEMM at its start).
+ Graph-capturable (current stream, no host sync). Deterministic. Compiled per (N, K, plan,
+ MMA rows (32 / 64 / 128 / 256), mode); `precompile` compiles every variant a weight's
+ dispatch range reaches.
+ """
+
+ def __call__(self, x: torch.Tensor, mw: MidWeight) -> torch.Tensor:
+ return mid_rows_mm(x, mw)
+
+ def reference(self, x: torch.Tensor, mw: MidWeight) -> torch.Tensor:
+ """fp32 `[M, N]`: the activation quantized as the op does (FP8 mode), both operands
+ dequantized, a true fp32 matmul."""
+ if mw.fp8:
+ q, s = fp8_row_quant(x)
+ xd = q.float() * s[:, None]
+ wd = mw.w.float() * (mw.scale[:, None] if mw.scale.numel() > 1 else mw.scale)
+ else:
+ xd, wd = x.float(), mw.w.float()
+ with true_fp32_matmul():
+ return xd @ wd.t()
+
+ def is_valid(self, x: torch.Tensor, mw: MidWeight) -> None:
+ m, k = x.shape
+ assert x.dtype == torch.bfloat16 and x.stride(1) == 1, "bf16 x with unit column stride"
+ assert 1 <= m <= MAX_ROWS, f"{m} rows: 1..{MAX_ROWS}"
+ assert mw.w.is_contiguous() and mw.k == k and mw.n % 16 == 0
+ if mw.fp8:
+ assert mw.scale is not None and mw.scale.dtype == torch.float32
+ assert mw.scale.numel() in (1, mw.n), "one weight scale, or one a weight row"
+ assert x.stride(0) % 8 == 0 and x.data_ptr() % 16 == 0, "16-byte aligned rows"
+ assert k % 256 == 0, "K a multiple of 256 (the quantization's row vectors)"
+ else:
+ assert mw.w.dtype == torch.bfloat16 and mw.scale is None
+
+
+class Fp8RowQuant(OpWrapper):
+ """`fp8_row_quant(x)`: `(q [M, K] e4m3, scale [M] fp32)` with `scale = max(max |x[m]|,
+ AMAX_FLOOR) / 448` and `q = e4m3_satfinite(x * (448 / max))` -- bitwise torch's cast of
+ `x * (448 / max)`. One warp a row, one launch; releases its dependents at its start."""
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="in-projection input, 37 rows", spec=dict(m=37, k=2560)),
+ Cell(
+ why="output-projection input, 256 rows, row stride K + 64",
+ spec=dict(m=256, k=6144, pad=64),
+ ),
+ Cell(why="one row", spec=dict(m=1, k=2560)),
+ )
+
+ note = """
+ Call site: `mid_rows_mm`'s FP8 mode. Preconditions: `x` bf16 `[M, K]`, unit column
+ stride, row stride a multiple of 8 elements, 16-byte aligned, K a multiple of 256.
+ """
+
+ def __call__(self, x: torch.Tensor):
+ return fp8_row_quant(x)
+
+ def reference(self, x: torch.Tensor):
+ """Both quotients correctly rounded (float64, rounded once), as the kernel's
+ `div.rn.f32`."""
+ amax = x.float().abs().amax(1).clamp_min(AMAX_FLOOR).double()
+ inv = (E4M3_MAX / amax).float()
+ return (x.float() * inv[:, None]).to(torch.float8_e4m3fn), (amax / E4M3_MAX).float()
+
+
+mid_rows_mm_op = MidRowsMm()
+fp8_row_quant_op = Fp8RowQuant()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/nvfp4_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/nvfp4_gemv.py
new file mode 100644
index 000000000000..bbe39917e7a4
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/nvfp4_gemv.py
@@ -0,0 +1,992 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Decode-row W4A16 projection: `y[M, N] = x[M, K] @ dequant(W)^T` with NVFP4 weights (e2m1
+codes, one e4m3 scale per 16 consecutive k of a row, one fp32 global scale) and bf16
+activations, M = 1 .. 16.
+
+CuTe DSL (inline PTX), one PDL launch, CUDA-graph capturable; compiled once per (N, K, plan)
+and row octet (1..8 or 9..16 rows): the row count is a launch argument.
+
+Math: swap-AB warp MMA (`mma.sync.m16n8k16`, f16 in, fp32 accumulate). A is 16 weight rows,
+each e2m1 pair widened by `cvt.rn.f16x2.e2m1x2` and multiplied by its block scale in f16
+(`cvt.rn.f16x2.e4m3x2`, `mul.rn.f16x2`): an e2m1 value times an e4m3 scale has at most five
+significant bits and lies in [2^-10, 2688], so the dequantized weight is exact. B is up to 8
+activation rows per n8 tile, rounded once from bf16 to f16 (exact for magnitudes in
+[2^-14, 65504]; smaller ones lose low bits, larger ones saturate). The MMA sums in fp32 over
+all of K; the sum times the global scale is rounded once to bf16.
+
+Weight layout (`pack_weight`): 1152-byte units of 16 rows x 128 k -- 1024 bytes of codes then
+128 bytes of scales -- stored k-block-major `[K/128][ceil(N/16)][1152]`. K is permuted
+identically for both operands inside every 128-block: lane quad t of k16 step s holds
+k = 32 t + 4 s + 0..3, so a thread's activations of a k-block are 32 consecutive values and
+its four k of a step share one 16-element scale block (2 t for steps 0..3, 2 t + 1 for 4..7).
+Codes: `[half h (2)][lane (32)][4 words]`, word i of half h the step 4 h + i, its bytes
+`(row g, k0 | k1 << 4), (row g + 8, same k), (row g, k2 | k3 << 4), (row g + 8, same k)` --
+one byte per A register. Scales: `[lane (32)][4 bytes]`: `(row g, block 2t), (row g + 8, 2t),
+(row g, 2t + 1), (row g + 8, 2t + 1)`.
+
+Work split (`Plan`): persistent, one or two resident CTAs per multiprocessor, each a contiguous
+range of 16-row tiles taken `rt` tiles (a chunk) at a time over all of K. A producer warp streams the units
+with `cp.async.bulk` into a ring of shared-memory slots (one k-block of the chunk's tiles a
+slot) completed on mbarriers, continuously across chunks and before the PDL wait (the weight
+is static). After `griddepcontrol.wait` the consumer warps convert the activations once to
+f16 into shared memory in MMA-fragment order; consumer warp w then takes tiles w, w + warps,
+... of every slot, its accumulators in registers for the chunk, and stores its tiles through a
+per-warp staging buffer.
+
+Split entries (`nvfp4_split_gemv*`, `SplitNvfp4Weight`): the same NVFP4 weight packed in
+`fp8_gemv`'s fragment order (`pack_split`) and streamed by that kernel, which splits K across
+warps and carries the GDN conv epilogues; for projections with few 16-row tiles a
+multiprocessor (the GDN in_proj), where one warp per tile over all of K is latency-bound.
+"""
+
+from __future__ import annotations
+
+import functools
+from dataclasses import dataclass
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass import const_expr
+from cutlass.cute.runtime import from_dlpack, make_ptr
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import (
+ Arch,
+ Cell,
+ OpWrapper,
+ true_fp32_matmul,
+)
+
+from .fp8_gemv import (
+ TILE_MAJOR_MAX_TILES,
+ _asm,
+ _bar_sync,
+ _bf16x2_to_f16x2,
+ _bulk_g2s,
+ _bulk_prefetch_l2,
+ _dq_step,
+ _env_pdl,
+ _ir,
+ _ldg4p,
+ _lds4,
+ _mbar_arrive,
+ _mbar_expect_tx,
+ _mbar_init,
+ _mbar_init_fence,
+ _mbar_wait,
+ _mma_f16,
+ _scales4,
+ _sms,
+ fp8_gemv,
+ fp8_gemv_conv1d,
+ fp8_gemv_conv1d_verify,
+ fp8_gemv_conv1d_verify_op,
+ verify_reference,
+)
+from .fp8_gemv import _compiled as _split_compiled
+from .fp8_gemv import default_plan as _split_plan
+
+TILE_N = 16
+BLOCK_K = 128
+SF_VEC = 16 # k per e4m3 block scale
+DATA = TILE_N * BLOCK_K // 2 # bytes of e2m1 codes of one (16-row, 128-k) unit
+SCALES = TILE_N * BLOCK_K // SF_VEC # bytes of its e4m3 block scales
+UNIT = DATA + SCALES
+E2M1_MAX = 6.0
+E4M3_MAX = 448.0
+MAX_ROWS = 16
+#: The e2m1 magnitudes, by the three low bits of a code (bit 3 is the sign).
+E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
+#: Rounding thresholds between consecutive e2m1 magnitudes (a tie rounds toward zero).
+_E2M1_BOUNDS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0)
+#: Shared memory one CTA may take (the SM100 opt-in maximum less a margin) and a
+#: multiprocessor holds (less a margin per resident CTA).
+SMEM_MAX = 225 << 10
+SMEM_SM = 228 << 10
+#: Activation loads a thread keeps in flight in the prologue (16 bytes each).
+_X_BATCH = 10
+
+
+# ------------------------------------------------------------------------------------------
+# inline PTX helpers
+# ------------------------------------------------------------------------------------------
+
+
+@dsl_user_op
+def _lds32(saddr, *, loc=None, ip=None):
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [_ir(saddr, cutlass.Int32, loc, ip)],
+ "ld.shared.b32 $0, [$1];",
+ "=r,r",
+ loc,
+ ip,
+ side=True,
+ )
+ )
+
+
+@dsl_user_op
+def _sts4(saddr, v0, v1, v2, v3, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [_ir(saddr, cutlass.Int32, loc, ip)]
+ + [_ir(v, cutlass.Int32, loc, ip) for v in (v0, v1, v2, v3)],
+ "{ st.shared.v4.b32 [$1], {$2, $3, $4, $5}; mov.u32 $0, 0; }",
+ "=r,r,r,r,r,r",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _stg4(gaddr, v0, v1, v2, v3, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [_ir(gaddr, cutlass.Int64, loc, ip)]
+ + [_ir(v, cutlass.Int32, loc, ip) for v in (v0, v1, v2, v3)],
+ "{ st.global.v4.b32 [$1], {$2, $3, $4, $5}; mov.u32 $0, 0; }",
+ "=r,l,r,r,r,r",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _lds_u16(saddr, *, loc=None, ip=None):
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [_ir(saddr, cutlass.Int32, loc, ip)],
+ "ld.shared.u16 $0, [$1];",
+ "=r,r",
+ loc,
+ ip,
+ side=True,
+ )
+ )
+
+
+@dsl_user_op
+def _stg_b16(gaddr, v, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [_ir(gaddr, cutlass.Int64, loc, ip), _ir(v, cutlass.Int32, loc, ip)],
+ "{ .reg .b16 h; cvt.u16.u32 h, $2; st.global.b16 [$1], h; mov.u32 $0, 0; }",
+ "=r,l,r",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+@dsl_user_op
+def _sts_bf16_scaled(saddr, v, s, *, loc=None, ip=None):
+ """bf16(v * s) to shared memory."""
+ _asm(
+ T.i32(),
+ [
+ _ir(saddr, cutlass.Int32, loc, ip),
+ _ir(v, cutlass.Float32, loc, ip),
+ _ir(s, cutlass.Float32, loc, ip),
+ ],
+ "{ .reg .f32 f; .reg .b16 h; mul.rn.f32 f, $2, $3; cvt.rn.bf16.f32 h, f; "
+ "st.shared.b16 [$1], h; mov.u32 $0, 0; }",
+ "=r,r,f,f",
+ loc,
+ ip,
+ side=True,
+ )
+
+
+# ------------------------------------------------------------------------------------------
+# kernel
+# ------------------------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class Plan:
+ """Launch shape: `grid` persistent CTAs (all resident: one or two per multiprocessor), each
+ taking its tiles `rt` at a time; `warps` consumer warps; `slots` ring slots (one k-block of
+ a chunk each); with `l2_ahead` > 0, issuing stage s also prefetches stage
+ s + slots + l2_ahead - 1 into L2."""
+
+ grid: int
+ rt: int
+ warps: int
+ slots: int
+ l2_ahead: int = 0
+
+
+class _Nvfp4Gemv:
+ def __init__(self, *, n, k, nt8, plan: Plan, pdl):
+ assert k % BLOCK_K == 0 and nt8 in (1, 2)
+ self.n, self.k, self.nt8, self.pdl = n, k, nt8, pdl
+ self.n_tiles = -(-n // TILE_N)
+ self.nkb = k // BLOCK_K
+ self.grid = plan.grid
+ self.rt, self.cw, self.slots, self.l2_ahead = plan.rt, plan.warps, plan.slots, plan.l2_ahead
+ self.rtw = -(-self.rt // self.cw)
+ self.threads = (self.cw + 1) * 32 # + the weight producer warp
+ self.np = nt8 * 8
+ self.slot_bytes = -(-(self.rt * UNIT) // 128) * 128
+ self.o_ring = -(-(self.nkb * nt8 * 4 * 512) // 128) * 128 # activation fragments first
+ self.o_stg = self.o_ring + self.slots * self.slot_bytes
+ self.o_bar = self.o_stg + self.cw * self.np * TILE_N * 2
+ self.smem = self.o_bar + 2 * self.slots * 8 + 128
+ assert self.smem <= SMEM_MAX, (self.smem, plan)
+
+ @cute.jit
+ def __call__(
+ self,
+ out_t: cute.Tensor,
+ x_ptr: cute.Pointer,
+ w_ptr: cute.Pointer,
+ x_stride: cutlass.Int32,
+ m: cutlass.Int32,
+ gscale: cutlass.Float32,
+ stream,
+ ):
+ self.kernel(
+ out_t.iterator.toint(), x_ptr.toint(), w_ptr.toint(), x_stride, m, gscale
+ ).launch(
+ grid=[self.grid, 1, 1],
+ block=[self.threads, 1, 1],
+ smem=self.smem,
+ stream=stream,
+ use_pdl=self.pdl,
+ )
+
+ @cute.jit
+ def _compute(self, acc, xf, slot, wi, ntc, lane):
+ """This warp's tiles of the k-block in ring slot `slot` times the activation fragments
+ `xf` ([nt8][16] f16x2), accumulated into `acc`."""
+ cw: cutlass.Constexpr = self.cw
+ nt8: cutlass.Constexpr = self.nt8
+ for i in cutlass.range_constexpr(self.rtw):
+ tl = wi + i * cw
+ if tl < ntc:
+ ua = slot + tl * UNIT
+ sc = _scales4(_lds32(ua + DATA + lane * 4))
+ wv = [_lds4(ua + h * 512 + lane * 16) for h in range(2)]
+ for s in cutlass.range_constexpr(8):
+ h = s // 4
+ a = _dq_step(wv[h][s % 4], sc[2 * h], sc[2 * h + 1])
+ for j in cutlass.range_constexpr(nt8):
+ r = _mma_f16(
+ a,
+ [xf[j][2 * s], xf[j][2 * s + 1]],
+ [acc[i, j, 0], acc[i, j, 1], acc[i, j, 2], acc[i, j, 3]],
+ )
+ for e in cutlass.range_constexpr(4):
+ acc[i, j, e] = r[e]
+
+ @cute.kernel
+ def kernel(
+ self,
+ out_base: cutlass.Int64,
+ x_base: cutlass.Int64,
+ w_base: cutlass.Int64,
+ x_stride: cutlass.Int32,
+ m: cutlass.Int32,
+ gscale: cutlass.Float32,
+ ):
+ tidx, _, _ = cute.arch.thread_idx()
+ bidx, _, _ = cute.arch.block_idx()
+ warp = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+ lane = cute.arch.lane_idx()
+ g = lane // 4
+ t = lane % 4
+ n_tiles: cutlass.Constexpr = self.n_tiles
+ nkb: cutlass.Constexpr = self.nkb
+ rt: cutlass.Constexpr = self.rt
+ cw: cutlass.Constexpr = self.cw
+ slots: cutlass.Constexpr = self.slots
+ slot_bytes: cutlass.Constexpr = self.slot_bytes
+ nt8: cutlass.Constexpr = self.nt8
+ np_: cutlass.Constexpr = self.np
+ rtw: cutlass.Constexpr = self.rtw
+ n_out: cutlass.Constexpr = self.n
+ l2a: cutlass.Constexpr = self.l2_ahead
+ grid: cutlass.Constexpr = self.grid
+
+ smem = cutlass.utils.SmemAllocator()
+ base_t = smem.allocate_tensor(
+ cutlass.Int8, cute.make_layout((self.smem - 128,), stride=(1,)), byte_alignment=128
+ )
+ sbase = cutlass.Int32(base_t.iterator.toint())
+ xs = sbase # activation fragments [kb][j][c][lane][16 B], f16
+ ring = sbase + self.o_ring
+ bars = sbase + self.o_bar # full[slots], empty[slots]
+
+ tile0 = (bidx * n_tiles) // grid
+ ntot = ((bidx + 1) * n_tiles) // grid - tile0
+ nch = (ntot + rt - 1) // rt
+
+ if warp == cw:
+ if lane == 0:
+ for sl in cutlass.range_constexpr(slots):
+ _mbar_init(bars + 8 * sl, 1)
+ _mbar_init(bars + 8 * (slots + sl), cw)
+ _mbar_init_fence()
+ cute.arch.sync_threads()
+ if const_expr(self.pdl):
+ cute.arch.griddepcontrol_launch_dependents()
+
+ if warp == cw:
+ # ---- weight producer (static data: no PDL wait); stage = (chunk, k-block) ----
+ if lane == 0:
+ for st in cutlass.range(nch * nkb):
+ sl = st % slots
+ if st >= slots:
+ _mbar_wait(bars + 8 * (slots + sl), ((st // slots) - 1) & 1)
+ c = st // nkb
+ kb = st - c * nkb
+ tc0 = tile0 + c * rt
+ ntc = cutlass.min(rt, ntot - c * rt)
+ if const_expr(l2a > 0):
+ pst = st + slots + l2a - 1
+ if pst < nch * nkb:
+ pc = pst // nkb
+ _bulk_prefetch_l2(
+ w_base
+ + (cutlass.Int64(pst - pc * nkb) * n_tiles + tile0 + pc * rt)
+ * UNIT,
+ cutlass.min(rt, ntot - pc * rt) * UNIT,
+ )
+ _mbar_expect_tx(bars + 8 * sl, ntc * UNIT)
+ _bulk_g2s(
+ ring + sl * slot_bytes,
+ w_base + (cutlass.Int64(kb) * n_tiles + tc0) * UNIT,
+ ntc * UNIT,
+ bars + 8 * sl,
+ )
+ else:
+ # ---- consumers ----
+ if const_expr(self.pdl):
+ cute.arch.griddepcontrol_wait()
+ # activations -> f16 fragments: item q = ((kb nt8 + j) 4 + c) 32 + lane holds row
+ # 8 j + g, k = 128 kb + 32 t + 8 c .. + 7 (rows past m zero)
+ n_items: cutlass.Constexpr = nkb * nt8 * 4 * 32
+ n_it: cutlass.Constexpr = -(-n_items // (cw * 32))
+ # loads issued a batch at a time ahead of their conversions (each load asm is
+ # ordered, so a load placed after a conversion would wait out its predecessor)
+ for b0 in cutlass.range_constexpr(0, n_it, _X_BATCH):
+ vals = []
+ for it in cutlass.range_constexpr(b0, min(b0 + _X_BATCH, n_it)):
+ q = tidx + it * (cw * 32)
+ ql = q % 32
+ row = 8 * ((q // 128) % nt8) + ql // 4
+ kx = (q // (128 * nt8)) * BLOCK_K + 32 * (ql % 4) + 8 * ((q // 32) % 4)
+ vals.append(
+ _ldg4p(
+ x_base
+ + cutlass.Int64(row) * cutlass.Int64(x_stride) * 2
+ + cutlass.Int64(kx * 2),
+ cutlass.Int32((row < m) & (q < n_items)),
+ )
+ )
+ for it in cutlass.range_constexpr(b0, min(b0 + _X_BATCH, n_it)):
+ q = tidx + it * (cw * 32)
+ if q < n_items:
+ v = vals[it - b0]
+ _sts4(
+ xs + q * 16,
+ _bf16x2_to_f16x2(v[0]),
+ _bf16x2_to_f16x2(v[1]),
+ _bf16x2_to_f16x2(v[2]),
+ _bf16x2_to_f16x2(v[3]),
+ )
+ _bar_sync(1, cw * 32)
+
+ stg = sbase + self.o_stg + warp * (np_ * TILE_N * 2) # [np][16] bf16 per warp
+ acc = cute.make_rmem_tensor(
+ cute.make_layout((rtw, nt8, 4), stride=(nt8 * 4, 4, 1)), cutlass.Float32
+ )
+ for c in cutlass.range(nch):
+ ntc = cutlass.min(rt, ntot - c * rt)
+ acc.fill(0.0)
+ for kb in cutlass.range_constexpr(nkb):
+ st = c * nkb + kb
+ sl = st % slots
+ xf = [
+ [
+ v
+ for cc in range(4)
+ for v in _lds4(xs + (((kb * nt8 + j) * 4 + cc) * 32 + lane) * 16)
+ ]
+ for j in range(nt8)
+ ]
+ _mbar_wait(bars + 8 * sl, (st // slots) & 1)
+ self._compute(acc, xf, ring + sl * slot_bytes, warp, ntc, lane)
+ cute.arch.sync_warp()
+ if lane == 0:
+ _mbar_arrive(bars + 8 * (slots + sl))
+ # ---- store: through the warp's staging buffer, a row's 16 outputs per 2 lanes ----
+ for i in cutlass.range_constexpr(rtw):
+ tl = warp + i * cw
+ if tl < ntc:
+ for j in cutlass.range_constexpr(nt8):
+ n0 = 8 * j + 2 * t
+ p = stg + (n0 * TILE_N + g) * 2
+ _sts_bf16_scaled(p, acc[i, j, 0], gscale)
+ _sts_bf16_scaled(p + TILE_N * 2, acc[i, j, 1], gscale)
+ _sts_bf16_scaled(p + 8 * 2, acc[i, j, 2], gscale)
+ _sts_bf16_scaled(p + TILE_N * 2 + 8 * 2, acc[i, j, 3], gscale)
+ cute.arch.sync_warp()
+ orow0 = (tile0 + c * rt + tl) * TILE_N
+ for rr in cutlass.range_constexpr(-(-(np_ * 2) // 32)):
+ li = lane + rr * 32
+ en = li // 2
+ hh = li % 2
+ if (li < np_ * 2) & (en < m):
+ if const_expr(n_out % 8 == 0):
+ if orow0 + 8 * hh < n_out:
+ v = _lds4(stg + (en * TILE_N + 8 * hh) * 2)
+ _stg4(
+ out_base
+ + (cutlass.Int64(en) * n_out + orow0 + 8 * hh) * 2,
+ v[0],
+ v[1],
+ v[2],
+ v[3],
+ )
+ else:
+ for e in cutlass.range_constexpr(8):
+ if orow0 + 8 * hh + e < n_out:
+ _stg_b16(
+ out_base
+ + (cutlass.Int64(en) * n_out + orow0 + 8 * hh + e)
+ * 2,
+ _lds_u16(stg + (en * TILE_N + 8 * hh + e) * 2),
+ )
+ cute.arch.sync_warp()
+
+
+# ------------------------------------------------------------------------------------------
+# host side
+# ------------------------------------------------------------------------------------------
+
+
+@dataclass
+class PackedNvfp4Weight:
+ """`data`: uint8 `[K/128, ceil(N/16), 1152]` (see `pack_weight`); `global_scale`: the fp32
+ dequantization multiplier (`w = e2m1 code x e4m3 block scale x global_scale`); logical
+ shape `[n, k]`."""
+
+ data: torch.Tensor
+ global_scale: float
+ n: int
+ k: int
+
+
+def _global_scale(w: torch.Tensor) -> float:
+ """amax(|w|) / (448 x 6): the largest block's scale lands on e4m3's largest value."""
+ return max(w.abs().amax().float().item(), 1e-30) / (E4M3_MAX * E2M1_MAX)
+
+
+def _quantize_rows(w: torch.Tensor, gscale: float):
+ """(codes uint8 `[n, K]`, e4m3 block scales `[n, K/16]`) of `w` under `gscale`."""
+ n, k = w.shape
+ blk = w.float().reshape(n, k // SF_VEC, SF_VEC)
+ sf = (blk.abs().amax(2) / (E2M1_MAX * gscale)).clamp(max=E4M3_MAX).to(torch.float8_e4m3fn)
+ step = sf.float()[..., None] * gscale
+ v = torch.where(step > 0, blk / step.clamp_min(1e-38), 0.0)
+ bounds = torch.tensor(_E2M1_BOUNDS, device=w.device)
+ idx = torch.bucketize(v.abs(), bounds).to(torch.uint8)
+ codes = idx | (((v < 0) & (idx > 0)).to(torch.uint8) << 3)
+ return codes.reshape(n, k), sf
+
+
+def quantize_nvfp4(w: torch.Tensor, rows_per_chunk: int = 16384):
+ """bf16/fp32 `[N, K]` -> (codes uint8 `[N, K]` -- bit 3 the sign, bits 0..2 the index into
+ `E2M1_VALUES` --, e4m3 block scales `[N, K/16]`, global scale): global = amax(|w|) /
+ (448 x 6), a block's scale = e4m3(amax(block) / 6 / global), a code the nearest e2m1 value
+ of w / (scale x global) (ties toward zero), saturated at 6."""
+ n, k = w.shape
+ assert k % SF_VEC == 0
+ gscale = _global_scale(w)
+ codes = torch.empty(n, k, dtype=torch.uint8, device=w.device)
+ scales = torch.empty(n, k // SF_VEC, dtype=torch.float8_e4m3fn, device=w.device)
+ for lo in range(0, n, rows_per_chunk):
+ hi = min(lo + rows_per_chunk, n)
+ codes[lo:hi], scales[lo:hi] = _quantize_rows(w[lo:hi], gscale)
+ return codes, scales, gscale
+
+
+def dequantize_nvfp4(codes: torch.Tensor, scales: torch.Tensor, gscale: float) -> torch.Tensor:
+ """fp32 `[N, K]` = e2m1(code) x its block scale x the global scale."""
+ lut = torch.tensor(E2M1_VALUES, dtype=torch.float32, device=codes.device)
+ val = lut[(codes & 7).long()] * (1.0 - 2.0 * (codes >> 3).float())
+ return val * scales.float().repeat_interleave(SF_VEC, 1) * gscale
+
+
+def _pack_into(out: torch.Tensor, tile0: int, codes: torch.Tensor, scales: torch.Tensor) -> None:
+ """Rows `16 tile0 ..` (codes `[n, K]`, scales `[n, K/16]`) into the packed buffer `out`
+ `[K/128, tiles, 1152]`; rows past n up to the tile boundary are zero."""
+ n, k = codes.shape
+ nkb = k // BLOCK_K
+ tiles = -(-n // TILE_N)
+ c = torch.zeros(tiles * TILE_N, k, dtype=torch.uint8, device=codes.device)
+ c[:n] = codes
+ # (tile, h', g, kb, t, hh, s4, e2, e1): row 8 h' + g, k = 128 kb + 32 t + 16 hh + 4 s4 + 2 e2 + e1
+ c = c.reshape(tiles, 2, 8, nkb, 4, 2, 4, 2, 2)
+ byte = c[..., 0] | (c[..., 1] << 4)
+ # unit byte order [hh][g][t][s4][e2][h']
+ out[:, tile0 : tile0 + tiles, :DATA] = byte.permute(3, 0, 5, 2, 4, 6, 7, 1).reshape(
+ nkb, tiles, DATA
+ )
+ s = torch.zeros(tiles * TILE_N, k // SF_VEC, dtype=torch.uint8, device=codes.device)
+ s[:n] = scales.view(torch.uint8)
+ # (tile, h', g, kb, t, hb): block 8 kb + 2 t + hb; unit order [g][t][hb][h']
+ out[:, tile0 : tile0 + tiles, DATA:] = (
+ s.reshape(tiles, 2, 8, nkb, 4, 2).permute(3, 0, 2, 4, 5, 1).reshape(nkb, tiles, SCALES)
+ )
+
+
+def pack_weight(codes: torch.Tensor, scales: torch.Tensor, gscale: float) -> PackedNvfp4Weight:
+ """(codes `[N, K]`, e4m3 scales `[N, K/16]`, global scale) -> the kernel's layout (module
+ docstring)."""
+ n, k = codes.shape
+ assert k % BLOCK_K == 0 and scales.shape == (n, k // SF_VEC)
+ out = torch.empty(k // BLOCK_K, -(-n // TILE_N), UNIT, dtype=torch.uint8, device=codes.device)
+ _pack_into(out, 0, codes, scales)
+ return PackedNvfp4Weight(out, float(gscale), n, k)
+
+
+def pack_from_bf16(w: torch.Tensor, rows_per_chunk: int = 16384) -> PackedNvfp4Weight:
+ """Converter from a bf16 `[N, K]` weight, quantized and packed in row chunks (the transient
+ memory stays a few hundred MB for the 248320 x 2560 vocabulary)."""
+ n, k = w.shape
+ assert k % BLOCK_K == 0 and rows_per_chunk % TILE_N == 0
+ gscale = _global_scale(w)
+ out = torch.empty(k // BLOCK_K, -(-n // TILE_N), UNIT, dtype=torch.uint8, device=w.device)
+ for lo in range(0, n, rows_per_chunk):
+ hi = min(lo + rows_per_chunk, n)
+ _pack_into(out, lo // TILE_N, *_quantize_rows(w[lo:hi], gscale))
+ return PackedNvfp4Weight(out, float(gscale), n, k)
+
+
+@dataclass
+class SplitNvfp4Weight:
+ """NVFP4 units in `fp8_gemv`'s fragment order, which that kernel (k split across warps, the
+ GDN conv epilogues) streams in place of FP8 ones: `data` uint8 `[K/128, ceil(N/16), 1152]`
+ (k-block-major) or `[ceil(N/16), K/128, 1152]` (`tile_major`), see `pack_split`; `scale` the
+ fp32 global scale as a one-element device tensor (the kernel reads it), also
+ `global_scale`; logical shape `[n, k]`."""
+
+ data: torch.Tensor
+ scale: torch.Tensor
+ global_scale: float
+ n: int
+ k: int
+ tile_major: bool = False
+ w4: bool = True
+
+
+def pack_split(
+ codes: torch.Tensor, scales: torch.Tensor, gscale: float, tile_major: bool | None = None
+) -> SplitNvfp4Weight:
+ """(codes `[N, K]`, e4m3 scales `[N, K/16]`, global scale) -> `fp8_gemv`'s NVFP4 units.
+
+ A unit (k-block kb, 16-row tile T) is 1152 bytes. Codes, 1024 bytes `[h (2)][lane (32)][4
+ words]`, lane = 4 g + t: word 2 cc + st of half h is k16 step st of chunk c = 2 h + cc,
+ covering k = 128 kb + 32 c + 8 t + 4 st + 0..3 (`fp8_gemv`'s k order), its bytes `(row g,
+ k0 | k1 << 4), (row g + 8, same k), (row g, k2 | k3 << 4), (row g + 8, same k)`. Scales,
+ 128 bytes `[g (8)][t // 2 (2)][c (4)][row g, row g + 8]`: the block 2 c + t // 2 of the
+ k-block. Rows past N are zero. Default `tile_major` as `fp8_gemv.pack_weight`'s."""
+ n, k = codes.shape
+ assert k % BLOCK_K == 0 and scales.shape == (n, k // SF_VEC)
+ n_tiles = -(-n // TILE_N)
+ nkb = k // BLOCK_K
+ if tile_major is None:
+ tile_major = n_tiles <= TILE_MAJOR_MAX_TILES
+ c = torch.zeros(n_tiles * TILE_N, k, dtype=torch.uint8, device=codes.device)
+ c[:n] = codes
+ # (tile, rh, g, kb, h, cc, t, st, e2, e1): row 8 rh + g, k = 128 kb + 32 (2 h + cc) + 8 t
+ # + 4 st + 2 e2 + e1
+ c = c.reshape(n_tiles, 2, 8, nkb, 2, 2, 4, 2, 2, 2)
+ byte = c[..., 0] | (c[..., 1] << 4)
+ # unit byte order [h][g][t][cc][st][e2][rh]
+ perm = (4, 2, 6, 5, 7, 8, 1)
+ body = (0, 3) if tile_major else (3, 0)
+ data = torch.empty(*(byte.shape[i] for i in body), UNIT, dtype=torch.uint8, device=codes.device)
+ data[..., :DATA] = byte.permute(*body, *perm).reshape(*data.shape[:2], DATA)
+ s = torch.zeros(n_tiles * TILE_N, k // SF_VEC, dtype=torch.uint8, device=codes.device)
+ s[:n] = scales.view(torch.uint8)
+ # (tile, rh, g, kb, c, th): block 8 kb + 2 c + th; unit order [g][th][c][rh]
+ s = s.reshape(n_tiles, 2, 8, nkb, 4, 2)
+ data[..., DATA:] = s.permute(*body, 2, 5, 4, 1).reshape(*data.shape[:2], SCALES)
+ scale = torch.tensor([gscale], dtype=torch.float32, device=codes.device)
+ return SplitNvfp4Weight(data, scale, float(gscale), n, k, tile_major)
+
+
+def pack_split_from_bf16(w: torch.Tensor, rows_per_chunk: int = 16384) -> SplitNvfp4Weight:
+ """Converter from a bf16 `[N, K]` weight (`quantize_nvfp4`, then `pack_split`)."""
+ return pack_split(*quantize_nvfp4(w, rows_per_chunk))
+
+
+#: Measured plans (GB300, 152 SMs, CUDA graph of back-to-back calls over weights rotated past
+#: L2), keyed (N, K, row octets): `ctas` resident CTAs per multiprocessor (two small ones let
+#: a launch's ramp overlap its neighbour's tail), tiles a chunk, consumer warps, ring slots,
+#: L2 look-ahead. The 16-row variant's activations (80 KB in shared memory) leave room for
+#: one CTA.
+_TABLE: dict[tuple[int, int, int], dict] = {
+ (248320, 2560, 1): dict(ctas=2, rt=8, warps=8, slots=5, l2_ahead=1),
+ (248320, 2560, 2): dict(ctas=1, rt=32, warps=16, slots=3, l2_ahead=1),
+}
+_DEFAULT = {
+ 1: dict(ctas=2, rt=8, warps=8, slots=6, l2_ahead=1),
+ 2: dict(ctas=1, rt=16, warps=16, slots=6, l2_ahead=1),
+}
+
+
+def default_plan(n: int, k: int, nt8: int, sms: int) -> Plan:
+ """The plan `nvfp4_gemv` launches with: the measured table entry for (N, K, nt8), else the
+ octet's default, its ring trimmed to the shared memory its CTAs per multiprocessor share."""
+ n_tiles = -(-n // TILE_N)
+ f = dict(_TABLE.get((n, k, nt8), _DEFAULT[nt8]))
+ ctas = f.pop("ctas")
+ f["rt"] = min(f["rt"], n_tiles)
+ slot = -(-(f["rt"] * UNIT) // 128) * 128
+ fixed = -(-((k // BLOCK_K) * nt8 * 2048) // 128) * 128 + f["warps"] * nt8 * 256 + 128
+ budget = min(SMEM_MAX, SMEM_SM // ctas - (1 << 10)) - fixed - 2 * 8 * f["slots"]
+ f["slots"] = max(2, min(f["slots"], budget // slot))
+ return Plan(grid=min(ctas * sms, n_tiles), **f)
+
+
+@functools.lru_cache(maxsize=None)
+def _compiled(n, k, nt8, plan: Plan, pdl: bool):
+ kern = _Nvfp4Gemv(n=n, k=k, nt8=nt8, plan=plan, pdl=pdl)
+ buf = torch.empty(16, dtype=torch.int32, device="cuda")
+ out = torch.empty(nt8 * 8, n, dtype=torch.bfloat16, device="cuda")
+ return cute.compile(
+ kern,
+ from_dlpack(out, assumed_align=16).mark_layout_dynamic(leading_dim=1),
+ make_ptr(cutlass.BFloat16, buf.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
+ make_ptr(cutlass.Uint8, buf.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
+ cutlass.Int32(k),
+ cutlass.Int32(1),
+ cutlass.Float32(1.0),
+ cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+ options="--opt-level 2 --enable-tvm-ffi",
+ )
+
+
+def _octet(m: int) -> int:
+ return 1 if m <= 8 else 2
+
+
+def precompile(w: PackedNvfp4Weight, pdl: bool = True) -> None:
+ """Compile both row octets of `w`'s shape now: a first call mid-serving would otherwise
+ stall on the compile."""
+ for nt8 in (1, 2):
+ _compiled(w.n, w.k, nt8, default_plan(w.n, w.k, nt8, _sms()), bool(pdl and _env_pdl()))
+
+
+def nvfp4_gemv(
+ x: torch.Tensor, w: PackedNvfp4Weight, pdl: bool = True, plan: Plan | None = None
+) -> torch.Tensor:
+ """`x[M, K]` bf16 (row stride a multiple of 8, 16-byte aligned base) @ dequant(w)^T ->
+ fresh `[M, N]` bf16. M = 1 .. MAX_ROWS."""
+ m, k = x.shape
+ assert k == w.k and x.dtype == torch.bfloat16 and x.stride(1) == 1
+ assert 1 <= m <= MAX_ROWS
+ assert x.stride(0) % 8 == 0 and x.data_ptr() % 16 == 0
+ out = torch.empty((m, w.n), dtype=torch.bfloat16, device=x.device)
+ nt8 = _octet(m)
+ if plan is None:
+ plan = default_plan(w.n, k, nt8, _sms())
+ fn = _compiled(w.n, k, nt8, plan, bool(pdl and _env_pdl()))
+ fn(out, x.data_ptr(), w.data.data_ptr(), x.stride(0), m, w.global_scale)
+ return out
+
+
+class Nvfp4Gemv(OpWrapper):
+ """`nvfp4_gemv(x, w)`: `y[M, N] = x[M, K] @ dequant(w)^T` in bf16 for 1..MAX_ROWS rows, the
+ weight a `PackedNvfp4Weight` (e2m1 codes, one e4m3 scale per 16 k of a row, an fp32 global
+ scale; `pack_from_bf16` / `pack_weight`). The dequantized weight is exact in f16; the
+ activation is rounded once to f16 on chip (exact for |x| in [2^-14, 65504]); fp32
+ accumulation over all of K, the global scale applied in fp32, one bf16 rounding.
+
+ Fusion boundary: the projection."""
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec keys: n, k, m (rows). Every cell: x a column slice of a wider bf16 row (row stride
+ # past K); the weight `quantize_nvfp4` of a random bf16 matrix with an outlier column and
+ # a block of near-zero values.
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="MTP draft head, one request's draft step (1 row)", spec=dict(n=248320, k=2560, m=1)
+ ),
+ Cell(why="draft head, 3 rows: a partial row octet", spec=dict(n=248320, k=2560, m=3)),
+ Cell(why="draft head, 8 rows: a full row octet", spec=dict(n=248320, k=2560, m=8)),
+ Cell(
+ why="draft head, 9 rows: the two-octet variant with one row in the second",
+ spec=dict(n=248320, k=2560, m=9),
+ ),
+ Cell(why="draft head, 16 rows: the largest dispatch", spec=dict(n=248320, k=2560, m=16)),
+ Cell(
+ why="N not a multiple of 8 (scalar stores, a partial last tile), 3 k-blocks, "
+ "fewer tiles than multiprocessors",
+ spec=dict(n=1001, k=384, m=5),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none. Call site: the MTP draft head's vocabulary projection at up to
+ `_FP8_GEMV_HEAD_ROWS` rows under `TRTLLM_MODELING_V2_MTP_HEAD_NVFP4=1` (the verify head
+ is untouched). Preconditions: x bf16 `[M, K]`, unit column stride, row stride a multiple
+ of 8 elements and a 16-byte aligned base; K a multiple of 128. One PDL launch (the weight
+ streamed before the wait); CUDA-graph capturable; compiled per (N, K, plan) for rows
+ 1..8 and 9..16 (`precompile` builds both), the row count a launch argument.
+ """
+
+ def __call__(self, x: torch.Tensor, w: PackedNvfp4Weight) -> torch.Tensor:
+ return nvfp4_gemv(x, w)
+
+ def reference(
+ self, x: torch.Tensor, codes: torch.Tensor, scales: torch.Tensor, global_scale: float
+ ) -> torch.Tensor:
+ """fp32 `[M, N]` over the dequantized weight (a true fp32 matmul); takes the unpacked
+ NVFP4 weight (`quantize_nvfp4`'s outputs) in place of the packed `w`."""
+ with true_fp32_matmul():
+ return x.float() @ dequantize_nvfp4(codes, scales, global_scale).t()
+
+ def is_valid(self, x: torch.Tensor, w: PackedNvfp4Weight, **kwargs) -> None:
+ assert x.float().abs().amax().item() <= 65504.0, (
+ "an activation beyond f16's range saturates at 65504 in the kernel"
+ )
+
+
+def precompile_split(
+ w: SplitNvfp4Weight,
+ max_rows: int,
+ conv_dim: int,
+ pad_slot_id: int,
+ verify_p: tuple[int, ...] = (),
+ pdl: bool = True,
+) -> None:
+ """Compile now every `fp8_gemv` variant over `w`'s shape a decode step can reach: the
+ projection and the single-token conv entry at 1..max_rows rows, and the verify entry at
+ each p of `verify_p` for every whole number of requests within max_rows. The kernel is
+ compiled per row count; a first call mid-serving would stall on the compile."""
+ pdl = bool(pdl and _env_pdl())
+ sms = _sms()
+ for m in range(1, max_rows + 1):
+ plan = _split_plan(w.n, w.k, m, sms, True)
+ _split_compiled(w.n, w.k, m, plan, pdl, w.tile_major, w4=True)
+ _split_compiled(w.n, w.k, m, plan, pdl, w.tile_major, conv_dim, pad_slot_id, w4=True)
+ for p in verify_p:
+ if m % p == 0 and m >= 2:
+ _split_compiled(
+ w.n, w.k, m, plan, pdl, w.tile_major, conv_dim, pad_slot_id, p, w4=True
+ )
+
+
+def _split_cell(why: str, n: int, k: int, m: int) -> Cell:
+ return Cell(why=why, spec=dict(n=n, k=k, m=m))
+
+
+class Nvfp4SplitGemv(OpWrapper):
+ """`nvfp4_split_gemv(x, w)` (`fp8_gemv` over a `SplitNvfp4Weight`): `y[M, N] = x[M, K] @
+ dequant(w)^T` in bf16 for 1..32 rows, the weight NVFP4 (e2m1 codes, one e4m3 scale per 16 k
+ of a row, an fp32 global scale; `pack_split_from_bf16` / `pack_split`). The dequantized
+ weight is exact in f16; the activation is rounded once to f16 on chip (exact for |x| in
+ [2^-14, 65504]); fp32 accumulation, the global scale applied in fp32, one bf16 rounding.
+
+ Fusion boundary: the projection; `Nvfp4SplitGemvConv1d` / `Nvfp4SplitGemvConv1dVerify`
+ also fuse the GDN conv step."""
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec keys: n, k, m (rows). Every cell: x a column slice of a wider bf16 row; the weight
+ # `quantize_nvfp4` of a random bf16 matrix with an outlier column and near-zero blocks.
+ CELLS: tuple[Cell, ...] = (
+ _split_cell("GDN in_proj, one decode row", 16480, 2560, 1),
+ _split_cell("GDN in_proj, 3 rows", 16480, 2560, 3),
+ _split_cell("GDN in_proj, 9 rows (two n8 tiles)", 16480, 2560, 9),
+ _split_cell("GDN in_proj, 16 rows", 16480, 2560, 16),
+ _split_cell("small N, tile-major units, 3 k-blocks", 1001, 384, 5),
+ )
+
+ note = """
+ Metadata consumed: none. Call site: the decode target's GDN in_proj at up to
+ `_INPROJ_NVFP4_MAX_ROWS` rows of an eager step under `TRTLLM_MODELING_V2_GDN_INPROJ_NVFP4=1`.
+ Preconditions: x bf16 `[M, K]`, unit column stride, row stride a multiple of 8 elements and
+ a 16-byte aligned base; K a multiple of 128. One PDL launch (the weight streamed before the
+ wait); CUDA-graph capturable; compiled per (N, K, M, plan) (`precompile_split`).
+ """
+
+ def __call__(self, x: torch.Tensor, w: SplitNvfp4Weight) -> torch.Tensor:
+ return fp8_gemv(x, w)
+
+ def reference(
+ self, x: torch.Tensor, codes: torch.Tensor, scales: torch.Tensor, global_scale: float
+ ) -> torch.Tensor:
+ """fp32 `[M, N]` over the dequantized weight (a true fp32 matmul); takes the unpacked
+ NVFP4 weight (`quantize_nvfp4`'s outputs) in place of the packed `w`."""
+ with true_fp32_matmul():
+ return x.float() @ dequantize_nvfp4(codes, scales, global_scale).t()
+
+ def is_valid(self, x: torch.Tensor, w: SplitNvfp4Weight, **kwargs) -> None:
+ assert x.float().abs().amax().item() <= 65504.0, (
+ "an activation beyond f16's range saturates at 65504 in the kernel"
+ )
+
+
+class Nvfp4SplitGemvConv1d(OpWrapper):
+ """`nvfp4_split_gemv_conv1d` (`fp8_gemv_conv1d` over a `SplitNvfp4Weight`): the GDN input
+ projection of decode rows over the NVFP4 copy, then for the conv columns of a live row the
+ causal conv step against its paged state (SiLU; the state advanced in place). Returns the
+ `[M, N]` projection with the conv columns replaced by the conv output."""
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec keys: slots (one per row; -1 the pad sentinel). Every cell: the GDN in_proj (16480 x
+ # 2560, conv_dim 10240), a strided 64-slot pool, width 4, SiLU.
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="STP decode: one row, live slot", spec=dict(slots=(37,))),
+ Cell(why="two rows, one padding slot", spec=dict(slots=(17, -1))),
+ Cell(why="four rows", spec=dict(slots=(0, 5, 63, 9))),
+ Cell(why="sixteen rows (two n8 tiles)", spec=dict(slots=tuple(range(2, 50, 3)))),
+ )
+
+ note = """
+ Call site: the decode target's GDN in_proj + conv step of a captured single-token step at
+ up to `_INPROJ_NVFP4_MAX_ROWS` rows under `TRTLLM_MODELING_V2_GDN_INPROJ_NVFP4=1`.
+ `conv_state` `[slots, conv_dim, 3]` bf16 (any strides), `conv_weight` `[conv_dim, 4]` bf16
+ contiguous, int32 slots; a `pad_slot_id` row gets the plain projection. Slots, taps and the
+ window are read before the PDL wait (nothing in a captured step writes this layer's pool
+ ahead of it).
+ """
+
+ def __call__(self, x, w, conv_state, conv_weight, conv_state_indices, conv_dim, pad_slot_id):
+ return fp8_gemv_conv1d(
+ x, w, conv_state, conv_weight, conv_state_indices, conv_dim, pad_slot_id
+ )
+
+ def reference(self, *args, **kwargs):
+ raise NotImplementedError("the test builds the fp32 reference from the dequantized weight")
+
+ def is_valid(self, x, w, **kwargs) -> None:
+ assert x.float().abs().amax().item() <= 65504.0, (
+ "an activation beyond f16's range saturates at 65504 in the kernel"
+ )
+
+
+class Nvfp4SplitGemvConv1dVerify(OpWrapper):
+ """`nvfp4_split_gemv_conv1d_verify` (`fp8_gemv_conv1d_verify` over a `SplitNvfp4Weight`):
+ the GDN input projection of a speculative verify step over the NVFP4 copy with
+ `causal_conv1d_triton_update`'s multi-token conv step as its epilogue. Returns `(conv [M,
+ conv_dim], projection columns conv_dim.. [M, N - conv_dim])`.
+
+ Fusion boundary: projection, conv, SiLU, snapshots and the pool advance."""
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec keys as `Fp8GemvConv1dVerify`'s.
+ CELLS: tuple[Cell, ...] = fp8_gemv_conv1d_verify_op.CELLS
+
+ note = """
+ Call site: the decode target's GDN in_proj of an MTP verify step in a captured graph, 2..16
+ rows, under `TRTLLM_MODELING_V2_GDN_INPROJ_NVFP4=1`; otherwise as `Fp8GemvConv1dVerify`
+ (bitwise the projection + `causal_conv1d_triton_update` over the same NVFP4 projection).
+ """
+
+ def __call__(
+ self,
+ x,
+ w,
+ conv_state,
+ conv_weight,
+ conv_state_indices,
+ intermediate_conv_window,
+ intermediate_state_indices,
+ conv_dim,
+ rows_per_seq,
+ pad_slot_id,
+ ):
+ return fp8_gemv_conv1d_verify(
+ x,
+ w,
+ conv_state,
+ conv_weight,
+ conv_state_indices,
+ intermediate_conv_window,
+ intermediate_state_indices,
+ conv_dim,
+ rows_per_seq,
+ pad_slot_id,
+ )
+
+ def reference(
+ self,
+ x,
+ codes,
+ scales,
+ global_scale,
+ conv_state,
+ conv_weight,
+ conv_state_indices,
+ intermediate_conv_window,
+ intermediate_state_indices,
+ conv_dim,
+ rows_per_seq,
+ pad_slot_id,
+ ):
+ """As `Fp8GemvConv1dVerify.reference`, over the dequantized NVFP4 weight."""
+ return verify_reference(
+ x,
+ dequantize_nvfp4(codes, scales, global_scale),
+ conv_state,
+ conv_weight,
+ conv_state_indices,
+ intermediate_conv_window,
+ intermediate_state_indices,
+ conv_dim,
+ rows_per_seq,
+ pad_slot_id,
+ )
+
+ def is_valid(
+ self,
+ x,
+ w,
+ conv_state,
+ conv_weight,
+ conv_state_indices,
+ intermediate_conv_window,
+ intermediate_state_indices,
+ conv_dim,
+ rows_per_seq,
+ pad_slot_id,
+ **kwargs,
+ ) -> None:
+ assert x.float().abs().amax().item() <= 65504.0, (
+ "an activation beyond f16's range saturates at 65504 in the kernel"
+ )
+ fp8_gemv_conv1d_verify_op.is_valid(
+ x,
+ w,
+ conv_state,
+ conv_weight,
+ conv_state_indices,
+ intermediate_conv_window,
+ intermediate_state_indices,
+ conv_dim,
+ rows_per_seq,
+ pad_slot_id,
+ )
+
+
+nvfp4_gemv_op = Nvfp4Gemv()
+nvfp4_split_gemv_op = Nvfp4SplitGemv()
+nvfp4_split_gemv_conv1d_op = Nvfp4SplitGemvConv1d()
+nvfp4_split_gemv_conv1d_verify_op = Nvfp4SplitGemvConv1dVerify()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/__init__.py
new file mode 100644
index 000000000000..52a7a9daf028
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_gated_residual.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_gated_residual.py
new file mode 100644
index 000000000000..8fe4616fcea7
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_gated_residual.py
@@ -0,0 +1,196 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Qwen4-Exp hyper-connection gating: low-rank stream mixer plus block-injection
+weights.
+
+A macro entry composed from native torch (three ``F.linear`` GEMMs and four
+elementwise passes); it launches no TensorRT-LLM kernel. The ``mhc_*`` op family
+computes a different hyper-connection (a per-stream scalar gate, a weighted sum,
+the raw stream gated) and does not instantiate hidden 2560. The semantics follow
+the HuggingFace ``Qwen4ExpTextGatedResidual`` minus its norm.
+
+Called by the qwen4_exp target (`qwen3_8_flash_next_nvfp4__sm_103__default`) at every
+layout.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+#: This checkpoint's geometry: 4 streams of 2560, rank-320 mixer bottleneck.
+CKPT_HC_COUNT = 4
+CKPT_HIDDEN = 2560
+
+
+def gate_ratio(out: torch.Tensor, ref: torch.Tensor) -> float:
+ """Worst element's distance from `ref` as a fraction of the test's band:
+ `rtol = eps`, `atol = eps * max|ref|`, both of `ref`'s dtype. Below 1 is
+ inside."""
+ eps = torch.finfo(ref.dtype).eps
+ ref64 = ref.double()
+ atol = eps * ref64.abs().max().item()
+ return ((out.double() - ref64).abs() / (atol + eps * ref64.abs())).max().item()
+
+
+class Qwen4ExpGatedResidual(OpWrapper):
+ """Collapse the `hc_count`-stream residual to one stream, and weight the
+ block's injection. With `C = hc_count`, `H = width // C` and `normed` the
+ `[..., C*H]` stream **after** its grouped norm, viewed stream-major as
+ `[..., C, H]`::
+
+ low = silu(normed @ down^T / C) # [..., lowrank]
+ gate = sigmoid(low @ up^T) # [..., C*H]
+ mixed = mean over streams of (gate * normed) # [..., H]
+ inject = 2 * sigmoid(normed @ block_inject^T / C) # [..., C]
+
+ Returns fresh contiguous `(mixed, inject)` in the activation dtype;
+ `inject` is `None` exactly when `block_inject_weight` is `None` (the
+ stack-final mixer), and `mixed` is then bitwise what it is with one.
+
+ Both divisions by `C` sit on the pre-activation logit; `mixed` gates the
+ **normed** stream; `inject` carries a factor of 2 (a zero logit is an
+ identity injection). GEMMs run in the activation dtype (fp32 accumulate),
+ the pointwise chain in fp32, each result rounded once.
+
+ Fusion boundary: both projections, both activations, the gating, the
+ stream mean and the injection projection are inside. The grouped gemma
+ norm producing `normed`, and the accumulate `qwen4_exp_post_inject` (into
+ the **pre-norm** stream the caller still holds), are not. Weights are used
+ as stored (`nn.Linear` layout), nothing is cached.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: leading (token dims), inject (True). Every cell is this
+ # checkpoint's 4x2560 bf16 geometry with synthetic weights of its measured
+ # element statistics, the normed stream reaching ~+-30 as the real one does.
+ CELLS: tuple[Cell, ...] = (
+ 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),
+ ),
+ # The MTP layer's shared_head mixer (block_inject_weight None) at the
+ # rows DecodeTarget.mtp_forward runs: one request, the largest decode
+ # graph batch, and mtp7's first draft over 128 x (1 + 7) verify rows.
+ Cell(
+ why="1 row, no injection: one request's MTP shared_head",
+ spec=dict(leading=(1,), inject=False),
+ ),
+ Cell(
+ why="128 rows, no injection: MTP shared_head at the largest decode graph batch",
+ spec=dict(leading=(128,), inject=False),
+ ),
+ Cell(
+ why="1024 rows, no injection: mtp7's first draft over 128 x 8 verify rows",
+ spec=dict(leading=(1024,), inject=False),
+ ),
+ Cell(
+ why="8192 rows, no injection: the envelope's largest, a max_num_tokens prefill chunk",
+ spec=dict(leading=(8192,), inject=False),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none; every row independent. Call sites: twice per
+ decoder layer (attention and MoE hyper-connections, 96 calls) plus the
+ stack-final mixer with `block_inject_weight=None` -- whose `mixed` *is*
+ the checkpoint's final norm -- and three more in the MTP layer (two
+ blocks and `shared_head`'s own mixer).
+
+ Preconditions, all loud (`RuntimeError` before producing anything): the
+ width divisible by `hc_count` (unflatten), the weights' input width equal
+ to it and their dtype equal to the activation's (`F.linear`).
+
+ float32 activations are accepted and not certified: `F.linear` on fp32
+ follows torch's TF32 default here, ~10 mantissa bits.
+
+ Operand range: with the real layer-0 weights the normed stream reaches
+ |30.1|, the logits 47.8 / 67.5, `inject` spans 0.430 .. 1.578. At up to
+ 4096x that stream every output stays finite, `|mixed| <= max|normed|`,
+ `inject` in `[0, 2]`; agreement with float64 degrades there because the
+ sigmoid saturates.
+
+ Certified: bf16 2-D `[rows, 10240]` at 8 and 8192 rows with the
+ injection; without it (the stack-final mixer, the MTP layer's
+ shared_head) at 1, 128, 256, 1024 and 8192 rows; CUDA-graph replay at 8
+ rows, bitwise equal to eager. On sm_103 the no-injection cells measure
+ `mixed` within 0.40 of the test's band, cosine >= 0.9999982 and relative
+ L2 <= 1.88e-3 against float64.
+
+ Not certified: other row counts; `inject` at 128 and 1024 rows (the MTP
+ layer's unfused attention and MoE sites; `mixed` there is bitwise the
+ no-injection cells'); more than one leading dim, hc_count other than 4,
+ hidden other than 2560, lowrank other than 320, float16 and float32,
+ weight/activation dtype mixing; the operand-range bounds above are
+ measured. Hard limits: not probed; a 2^31-element operand would take
+ 209716 rows, past the 8192-row envelope.
+ """
+
+ def __call__(
+ self,
+ hyper_input_normed: torch.Tensor,
+ input_mix_weight_down: torch.Tensor,
+ input_mix_weight_up: torch.Tensor,
+ block_inject_weight: Optional[torch.Tensor],
+ hc_count: int,
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ activation_dtype = hyper_input_normed.dtype
+ hidden_size = hyper_input_normed.shape[-1] // hc_count
+
+ # GEMMs in the activation dtype: a bf16 product accumulates in fp32 with
+ # no ambient state, an fp32 one would follow torch's TF32 default.
+ low_rank = F.linear(hyper_input_normed, input_mix_weight_down)
+ low_rank = F.silu(low_rank.float() / hc_count).to(activation_dtype)
+
+ mix_weight = F.linear(low_rank, input_mix_weight_up)
+ mix_weight = torch.sigmoid(mix_weight.float()).unflatten(-1, (hc_count, hidden_size))
+ streams = hyper_input_normed.float().unflatten(-1, (hc_count, hidden_size))
+ mixed_input = (mix_weight * streams).mean(dim=-2).to(activation_dtype)
+
+ injection_weights = None
+ if block_inject_weight is not None:
+ inject_logits = F.linear(hyper_input_normed, block_inject_weight)
+ injection_weights = 2.0 * torch.sigmoid(inject_logits.float() / hc_count)
+ injection_weights = injection_weights.to(activation_dtype)
+
+ return mixed_input, injection_weights
+
+ def reference(
+ self,
+ hyper_input_normed: torch.Tensor,
+ input_mix_weight_down: torch.Tensor,
+ input_mix_weight_up: torch.Tensor,
+ block_inject_weight: Optional[torch.Tensor],
+ hc_count: int,
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ """float64 throughout -- an fp32 GEMM here would be TF32 and the less
+ accurate side -- rounded once to the activation dtype."""
+ dtype = hyper_input_normed.dtype
+ hidden = hyper_input_normed.shape[-1] // hc_count
+ x = hyper_input_normed.double()
+ low = F.silu((x @ input_mix_weight_down.double().t()) / hc_count)
+ gate = torch.sigmoid(low @ input_mix_weight_up.double().t()).unflatten(
+ -1, (hc_count, hidden)
+ )
+ mixed = (gate * x.unflatten(-1, (hc_count, hidden))).mean(dim=-2).to(dtype)
+ inject = None
+ if block_inject_weight is not None:
+ inject = 2.0 * torch.sigmoid((x @ block_inject_weight.double().t()) / hc_count)
+ inject = inject.to(dtype)
+ return mixed, inject
+
+
+qwen4_exp_gated_residual = Qwen4ExpGatedResidual()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_first_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_first_cute.py
new file mode 100644
index 000000000000..8e58e9df35cc
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_first_cute.py
@@ -0,0 +1,605 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Decode-row hyper-connection site, first half, for one-row batches: the previous
+block's injection (with the MoE block's combine and the routed experts' finalize
+folded in, when given) and the first projection per stream, one PDL launch.
+
+Tier 5 (CuTe DSL written for this target). A stream is exactly `THREADS x VEC`
+elements, one 16-byte vector per thread (a power-of-two tile would pad a 2560-wide
+stream to 8 x 512, three eighths of it masked lanes holding registers), and the grid
+is sized so that every CTA is resident at once (one wave), so all of them stream their
+weight before the producer finishes.
+
+Order of work inside a CTA (`OUTS` weight rows x one stream):
+
+1. its weight rows and its slice of the norm weight into registers -- static, so
+ before `griddepcontrol.wait`, while the kernel ahead of this one runs;
+2. `griddepcontrol.wait`; then the stream's vector of the residual and of the
+ block output -- with the fold, the slot map and the weights first, then the
+ slot rows -- and the injection weight;
+3. the injected residual rounded once (stored by the stream's first CTA), its sum
+ of squares and its `(1 + w)`-scaled products with the weight rows, in fp32;
+4. a warp then a cross-warp reduction; thread `j < OUTS` stores partial `j`, thread
+ `OUTS` of the stream's first CTA the sum of squares.
+"""
+
+from __future__ import annotations
+
+import functools
+from typing import Optional, Tuple
+
+import cuda.bindings.driver as cuda_drv
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass import const_expr
+from cutlass.cute import experimental as cute_ext
+from cutlass.cute.runtime import from_dlpack
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+
+#: One decode row per launch: the forward runs only 1-row decode batches through the
+#: site's GEMV pair (`_HC_DECODE_GEMV_MAX_ROWS`).
+MAX_ROWS = 1
+VEC = 8
+#: Threads per CTA: one 16-byte vector of a 2560-wide stream each.
+THREADS = 320
+#: Streaming multiprocessors the one-wave grid must fit (GB300).
+_SMS = 152
+_COMPILE_OPTIONS = "--ptxas-options -maxrregcount=128"
+
+
+def outs_for(n_cols: int, hc_count: int) -> int:
+ """Weight rows per CTA: the fewest that divide `n_cols` and keep the grid in one wave."""
+ for outs in range(1, n_cols + 1):
+ if n_cols % outs == 0 and (n_cols // outs) * hc_count <= _SMS:
+ return outs
+ raise AssertionError(f"no one-wave tiling for {n_cols} columns")
+
+
+class _FirstKernel:
+ def __init__(self, *, hidden, outs, inject, combine, topk, pdl):
+ assert hidden == THREADS * VEC, "one vector per thread"
+ self.hidden, self.outs = hidden, outs
+ self.inject, self.combine, self.topk, self.pdl = inject, combine, topk, pdl
+ self.warps = THREADS // 32
+
+ @cute.jit
+ def __call__(
+ self,
+ mHc,
+ mBlock,
+ mInj,
+ mShared,
+ mGate,
+ mTopW,
+ mTopRows,
+ mNorm,
+ mW,
+ mHcOut,
+ mPart,
+ mSs,
+ stream: cuda_drv.CUstream,
+ ):
+ n_tiles = cute.size(mPart, mode=[1]) // self.outs
+ streams = cute.size(mPart, mode=[0])
+ load_static = cute.make_copy_atom(
+ cute.nvgpu.CopyG2ROp(),
+ cutlass.BFloat16,
+ num_bits_per_copy=16 * VEC,
+ load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING,
+ )
+ load_act = cute.make_copy_atom(
+ cute.nvgpu.CopyG2ROp(),
+ cutlass.BFloat16,
+ num_bits_per_copy=16 * VEC,
+ load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS,
+ )
+ store = cute.make_copy_atom(
+ cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=16 * VEC
+ )
+ self.kernel(
+ mHc,
+ mBlock,
+ mInj,
+ mShared,
+ mGate,
+ mTopW,
+ mTopRows,
+ mNorm,
+ mW,
+ mHcOut,
+ mPart,
+ mSs,
+ load_static,
+ load_act,
+ store,
+ ).launch(
+ grid=[n_tiles, streams, 1],
+ block=[THREADS, 1, 1],
+ smem=(self.outs + 1) * self.warps * 4,
+ stream=stream,
+ use_pdl=self.pdl,
+ # At most 96 registers a thread: the mix kernel that follows (512 threads at
+ # 56) then fits on the SM beside this CTA and streams its weights before
+ # this kernel ends.
+ min_blocks_per_mp=2,
+ )
+
+ @cute.kernel
+ def kernel(
+ self,
+ mHc,
+ mBlock,
+ mInj,
+ mShared,
+ mGate,
+ mTopW,
+ mTopRows,
+ mNorm,
+ mW,
+ mHcOut,
+ mPart,
+ mSs,
+ load_static: cute.CopyAtom,
+ load_act: cute.CopyAtom,
+ store: cute.CopyAtom,
+ ):
+ tidx, _, _ = cute.arch.thread_idx()
+ bidx, s, _ = cute.arch.block_idx()
+ warp = cute.arch.warp_idx()
+ outs: cutlass.Constexpr = self.outs
+ warps: cutlass.Constexpr = self.warps
+ topk: cutlass.Constexpr = self.topk
+ n0 = bidx * outs
+ vi = s * THREADS + tidx # this thread's vector of the width; `tidx` of the stream
+
+ w_vec = cute.logical_divide(mW, (None, VEC))
+ hc_vec = cute.logical_divide(mHc, (None, VEC))
+ nw_vec = cute.logical_divide(mNorm, (None, VEC))
+
+ # 1. Static operands, before the wait.
+ w_reg = cute.make_rmem_tensor(
+ cute.make_layout((outs, VEC), stride=(VEC, 1)), cutlass.BFloat16
+ )
+ for j in cutlass.range_constexpr(outs):
+ cute.copy(load_static, w_vec[n0 + j, (None, vi)], w_reg[j, None])
+ nw_reg = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16)
+ cute.copy(load_static, nw_vec[0, (None, vi)], nw_reg)
+ if const_expr(self.pdl):
+ cute.arch.griddepcontrol_launch_dependents()
+ cute.arch.griddepcontrol_wait()
+
+ # 2. The producers' outputs.
+ x_reg = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16)
+ cute.copy(load_act, hc_vec[0, (None, vi)], x_reg)
+ xr = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.Float32)
+ for v in cutlass.range_constexpr(VEC):
+ xr[v] = x_reg[v].to(cutlass.Float32)
+ if const_expr(self.inject):
+ b_vec = cute.logical_divide(mBlock, (None, VEC))
+ g = mInj[0, s].to(cutlass.Float32)
+ b = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.Float32)
+ if const_expr(topk > 0):
+ # The runner's finalize: fp32 FMAs in slot order from zero, a slot mapped
+ # to -1 skipped, one rounding. Every slot row is requested before any FMA.
+ rows = cute.make_rmem_tensor(cute.make_layout((topk,), stride=(1,)), cutlass.Int32)
+ wts = cute.make_rmem_tensor(cute.make_layout((topk,), stride=(1,)), cutlass.Float32)
+ for k in cutlass.range_constexpr(topk):
+ rows[k] = mTopRows[0, k]
+ wts[k] = mTopW[0, k].to(cutlass.Float32)
+ fc2 = cute.make_rmem_tensor(
+ cute.make_layout((topk, VEC), stride=(VEC, 1)), cutlass.BFloat16
+ )
+ for k in cutlass.range_constexpr(topk):
+ cute.copy(
+ load_act,
+ b_vec[cutlass.max(rows[k], cutlass.Int32(0)), (None, tidx)],
+ fc2[k, None],
+ )
+ for v in cutlass.range_constexpr(VEC):
+ acc = cutlass.Float32(0.0)
+ for k in cutlass.range_constexpr(topk):
+ if rows[k] >= 0:
+ acc = cute.math.fma(fc2[k, v].to(cutlass.Float32), wts[k], acc)
+ b[v] = acc.to(cutlass.BFloat16).to(cutlass.Float32)
+ else:
+ blk = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16)
+ cute.copy(load_act, b_vec[0, (None, tidx)], blk)
+ for v in cutlass.range_constexpr(VEC):
+ b[v] = blk[v].to(cutlass.Float32)
+ if const_expr(self.combine):
+ # The MoE block output, routed + sigmoid(gate) * shared in fp32, rounded once.
+ sh_vec = cute.logical_divide(mShared, (None, VEC))
+ sh = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16)
+ cute.copy(load_act, sh_vec[0, (None, tidx)], sh)
+ z = mGate[0].to(cutlass.Float32)
+ sg = 1.0 / (1.0 + cute.math.exp(-z))
+ for v in cutlass.range_constexpr(VEC):
+ b[v] = (
+ cute.math.fma(sg, sh[v].to(cutlass.Float32), b[v])
+ .to(cutlass.BFloat16)
+ .to(cutlass.Float32)
+ )
+ out = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16)
+ for v in cutlass.range_constexpr(VEC):
+ # Rounded once, like `addcmul`: the norm and every later site read this value.
+ out[v] = cute.math.fma(b[v], g, xr[v]).to(cutlass.BFloat16)
+ xr[v] = out[v].to(cutlass.Float32)
+ if bidx == 0:
+ o_vec = cute.logical_divide(mHcOut, (None, VEC))
+ cute.copy(store, out, o_vec[0, (None, vi)])
+
+ # 3. Sum of squares and the scaled products, fp32.
+ acc = cute.make_rmem_tensor(cute.make_layout((outs + 1,), stride=(1,)), cutlass.Float32)
+ acc.fill(0.0)
+ for v in cutlass.range_constexpr(VEC):
+ y = xr[v] * (1.0 + nw_reg[v].to(cutlass.Float32))
+ for j in cutlass.range_constexpr(outs):
+ acc[j] = y * w_reg[j, v].to(cutlass.Float32) + acc[j]
+ acc[outs] = xr[v] * xr[v] + acc[outs]
+
+ # 4. Reductions; thread j < OUTS stores partial j, thread OUTS the sum of squares.
+ for j in cutlass.range_constexpr(outs + 1):
+ acc[j] = cute.arch.warp_reduction_sum(acc[j])
+ smem = cutlass.utils.SmemAllocator()
+ red = smem.allocate_tensor(
+ cutlass.Float32,
+ cute.make_layout((outs + 1, warps), stride=(warps, 1)),
+ byte_alignment=16,
+ )
+ with cute.arch.elect_one():
+ for j in cutlass.range_constexpr(outs + 1):
+ red[j, warp] = acc[j]
+ cute.arch.sync_threads()
+ if tidx < outs:
+ total = cutlass.Float32(0.0)
+ for w in cutlass.range_constexpr(warps):
+ total = total + red[tidx, w]
+ mPart[s, n0 + tidx] = total
+ if tidx == outs:
+ if bidx == 0:
+ total = cutlass.Float32(0.0)
+ for w in cutlass.range_constexpr(warps):
+ total = total + red[outs, w]
+ mSs[s] = total
+
+
+def _dl(t: torch.Tensor, align: int = 16):
+ return from_dlpack(t, assumed_align=align)
+
+
+@functools.lru_cache(maxsize=None)
+def _compiled(
+ hidden, hc_count, n_cols, n_rows, inject, combine, topk, n_perm, pdl, hc_stride, block_stride
+):
+ """One compiled launch per variant and geometry (strides are part of a CuTe layout)."""
+ outs = outs_for(n_cols, hc_count)
+ kern = _FirstKernel(
+ hidden=hidden, outs=outs, inject=inject, combine=combine, topk=topk, pdl=pdl
+ )
+ dev, bf = "cuda", torch.bfloat16
+ width = hidden * hc_count
+ hc = torch.empty(1, hc_stride, dtype=bf, device=dev)[:, :width]
+ block = torch.empty(max(n_perm, 1), block_stride, dtype=bf, device=dev)[:, :hidden]
+ inj = torch.empty(1, hc_count, dtype=bf, device=dev)
+ shared = torch.empty(1, hidden, dtype=bf, device=dev)
+ gate = torch.empty(1, dtype=bf, device=dev)
+ topw = torch.empty(1, max(topk, 1), dtype=bf, device=dev)
+ toprows = torch.empty(1, max(topk, 1), dtype=torch.int32, device=dev)
+ norm = torch.empty(1, width, dtype=bf, device=dev)
+ w = torch.empty(n_rows, width, dtype=bf, device=dev)
+ hc_out = torch.empty(1, width, dtype=bf, device=dev)
+ part = torch.empty(hc_count, n_cols, dtype=torch.float32, device=dev)
+ ss = torch.empty(hc_count, dtype=torch.float32, device=dev)
+ stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream)
+ return cute_ext.compile(
+ kern,
+ _dl(hc),
+ _dl(block),
+ _dl(inj, 2),
+ _dl(shared),
+ _dl(gate, 2),
+ _dl(topw, 2),
+ _dl(toprows, 4),
+ _dl(norm),
+ _dl(w),
+ _dl(hc_out),
+ _dl(part, 4),
+ _dl(ss, 4),
+ stream,
+ options=_COMPILE_OPTIONS,
+ )
+
+
+def launch(
+ hyper_input,
+ block_out,
+ inject,
+ norm_weight,
+ first,
+ hc_count,
+ n_cols,
+ shared=None,
+ gate_logits=None,
+ topk_weights=None,
+ topk_rows=None,
+):
+ """The launch behind `Qwen4ExpHcDecodeFirstCute.__call__`. Returns `(hc_out, part, ss)`,
+ `part` fp32 `[C, 1, n_cols]` and `ss` fp32 `[C, 1]`."""
+ rows, width = hyper_input.shape
+ hidden = width // hc_count
+ dev = hyper_input.device
+ has_inject = block_out is not None
+ has_shared = shared is not None
+ top_k = topk_rows.shape[1] if topk_rows is not None else 0
+ hc_out = (
+ torch.empty((rows, width), dtype=hyper_input.dtype, device=dev)
+ if has_inject
+ else hyper_input
+ )
+ part = torch.empty((hc_count, rows, n_cols), dtype=torch.float32, device=dev)
+ ss = torch.empty((hc_count, rows), dtype=torch.float32, device=dev)
+ if rows == 0:
+ return hc_out, part, ss
+ pdl = get_env_enable_pdl()
+ weight = first.t() # the row-major [n_rows, width] weight
+ gate = gate_logits.reshape(1) if has_shared else hyper_input.new_empty(1)
+ n_perm = block_out.shape[0] if (has_inject and top_k) else 1
+ fn = _compiled(
+ hidden,
+ hc_count,
+ n_cols,
+ weight.shape[0],
+ has_inject,
+ has_shared,
+ top_k,
+ n_perm,
+ bool(pdl),
+ hyper_input.stride(0),
+ block_out.stride(0) if (has_inject and top_k) else hidden,
+ )
+ blk = block_out if has_inject else hyper_input.new_empty(1, hidden)
+ inj = inject if has_inject else hyper_input.new_empty(1, hc_count)
+ sh = shared if has_shared else hyper_input.new_empty(1, hidden)
+ topw = topk_weights if top_k else hyper_input.new_empty(1, 1)
+ toprows = topk_rows if top_k else torch.empty(1, 1, dtype=torch.int32, device=dev)
+ stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream)
+ fn(
+ _dl(hyper_input),
+ _dl(blk),
+ _dl(inj, 2),
+ _dl(sh),
+ _dl(gate, 2),
+ _dl(topw, 2),
+ _dl(toprows, 4),
+ _dl(norm_weight.reshape(1, -1)),
+ _dl(weight),
+ _dl(hc_out),
+ _dl(part[:, 0, :], 4),
+ _dl(ss[:, 0], 4),
+ stream,
+ )
+ return hc_out, part, ss
+
+
+class Qwen4ExpHcDecodeFirstCute(OpWrapper):
+ """The first half of a one-row decode hyper-connection site: the previous block's
+ injection and the mixer's first projection (`[down; block_inject]`) per stream. With
+ `C = hc_count`, `H = width // C`, `N = n_cols`, `first` the `[width, >= N]` view of
+ the load-time `[down; block_inject; 0]` weight::
+
+ hc_out[0, s*H+i] = hyper_input[0, s*H+i] + block_out[0, i] * inject[0, s] (rounded once)
+ part[s, 0, n] = sum_i hc_out[0, s*H+i] * (1 + norm_weight[s*H+i]) * first[s*H+i, n]
+ ss[s, 0] = sum_i hc_out[0, s*H+i]^2
+
+ Returns `(hc_out, part, ss)`: `hc_out` fresh in the activation dtype (it *is*
+ `hyper_input` when `block_out` / `inject` are None -- the norm-only site), `part`
+ fp32 `[C, 1, N]`, `ss` fp32 `[C, 1]`. The grouped RMS `rstd[s] = (ss[s, 0] / H +
+ eps)^-1/2` scales stream `s`'s partial: `normed @ first[:, :N] = sum_s rstd[s] *
+ part[s]`. `qwen4_exp_hc_decode_mix_cute` makes that combination.
+
+ With `shared` and `gate_logits` (the site after a MoE block), `block_out` is the
+ routed experts' output and the block output injected above is `block_out[0, i] +
+ sigmoid(gate_logits[0]) * shared[0, i]` in fp32, rounded once to `block_out`'s
+ dtype -- bitwise what `fused_sigmoid_gate_mul_add` stores. `gate_logits` is one
+ pre-sigmoid scalar per row, `[T, 1]` or `[T]`.
+
+ With `topk_weights` / `topk_rows` as well (the MoE runner left unfinalized),
+ `block_out` is its permuted FC2 output `[P, H]` and the routed row is the runner's
+ finalize, formed here: `routed[0, i] = sum_k topk_weights[0, k] *
+ block_out[topk_rows[0, k], i]` as fp32 FMAs in slot order from zero (a slot whose
+ row is -1 is skipped), rounded once to `block_out`'s dtype -- bitwise what the
+ runner's finalize stores. `topk_rows` is the runner's expanded->permuted index map
+ `[T, k]` int32, `topk_weights` its `[T, k]` routing weights. The fold needs the
+ combine (`shared`, `gate_logits`).
+
+ Fusion boundary: the routed experts' finalize (optional), the shared-expert combine
+ (optional), the previous block's injection, the `(1 + w)` scale and the first
+ projection per stream, plus the per-stream sum of squares. The routed and shared
+ experts and the gate projection, the RMS itself, the activations and everything
+ after are not inside.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: rows (1), inject (True), moe (False: shared + gate_logits [T, 1] too), fold
+ # (0: top-k of the finalize folded in, block_out then the permuted rows), drop
+ # (False: some slots -1), n_cols (324). The weight is the target's padded
+ # [width, 384] (columns past n_cols never read).
+ # Four cells: the forward reaches four compiled variants (inject at 324 cols, the
+ # fold at 324, the fold at 320 into the stack-final mixer, norm-only), none of which
+ # another covers; the -1 slots ride the 324-col fold cell.
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="checkpoint 4x2560, 324 cols of 384, bf16, inject, 1 row", spec=dict(rows=1)),
+ Cell(
+ why="finalize fold, top-10, last 3 slots mapped to -1, 1 row",
+ spec=dict(rows=1, moe=True, fold=10, drop=True),
+ ),
+ Cell(
+ why="finalize fold into the stack-final mixer: 320 cols",
+ spec=dict(rows=1, moe=True, fold=10, n_cols=320),
+ ),
+ Cell(why="norm-only site (stack start / after PLE)", spec=dict(rows=1, inject=False)),
+ )
+
+ note = """
+ Metadata consumed: none. Call sites: `DecodeTarget`'s trunk, every
+ hyper-connection site of a one-row decode step (97 per forward), followed by
+ `qwen4_exp_hc_decode_mix_cute` (or its NVFP4 variant at an MoE site); batches of
+ more rows take `qwen4_exp_hc_inject_norm` and the mixer. The 47 sites after a MoE
+ block other than the one before the PLE add pass `shared` and `gate_logits`, and
+ with them the unfinalized runner's `topk_weights` / `topk_rows`.
+
+ Preconditions: 2-D operands with unit column stride (row strides free); `first`
+ `[width, >= n_cols]` with unit stride along `width` (the `.t()` view of a row-major
+ `[N, width]` weight; columns past `n_cols` are never read); one row; `hidden` 2560
+ (320 threads x 8); bf16 operands; `n_cols` with a divisor that keeps `n_cols / outs x C <= 152`
+ CTAs (324 -> 9 rows x 36 tiles x 4 streams = 144; 320 -> 10 rows, 128); the
+ permuted rows' count is part of the compiled layout (one compile per value).
+
+ Numerics: `hc_out` as `qwen4_exp_hc_inject_norm` (fp32 FMAs, one rounding; the
+ finalize and the combine likewise); `part` and `ss` are fp32 reductions in a fixed
+ order, inside the float64 reference's band (the combine and the finalize can each
+ land one ulp of the block output from it, scaled into `hc_out` by `inject`).
+ Deterministic (measured). Graph-capturable: compiled on the first call per
+ variant, which must precede capture (the engine's warmup forward does it).
+
+ PDL: with `TRTLLM_ENABLE_PDL` (default on, sm_90+) the weight rows and the norm
+ slice stream before `griddepcontrol.wait`, every activation read and every
+ write after it, and dependents are released before the wait.
+
+ Certified: one bf16 row, the injection at 324 columns, the fold (top-10, three
+ slots -1) at 324 columns, the fold into the stack-final mixer at 320 columns,
+ the norm-only site; graph replay under PDL with the fold.
+
+ Not certified: the combine without the fold, a flat `[T]` gate, a row-slice
+ `hyper_input`, a weight narrower than the padded 384 rows (accepted; no cell
+ drives them).
+ """
+
+ def __call__(
+ self,
+ hyper_input: torch.Tensor,
+ block_out: Optional[torch.Tensor],
+ inject: Optional[torch.Tensor],
+ norm_weight: torch.Tensor,
+ first: torch.Tensor,
+ hc_count: int,
+ n_cols: int,
+ shared: Optional[torch.Tensor] = None,
+ gate_logits: Optional[torch.Tensor] = None,
+ topk_weights: Optional[torch.Tensor] = None,
+ topk_rows: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ return launch(
+ hyper_input,
+ block_out,
+ inject,
+ norm_weight,
+ first,
+ hc_count,
+ n_cols,
+ shared,
+ gate_logits,
+ topk_weights,
+ topk_rows,
+ )
+
+ def reference(
+ self,
+ hyper_input: torch.Tensor,
+ block_out: Optional[torch.Tensor],
+ inject: Optional[torch.Tensor],
+ norm_weight: torch.Tensor,
+ first: torch.Tensor,
+ hc_count: int,
+ n_cols: int,
+ shared: Optional[torch.Tensor] = None,
+ gate_logits: Optional[torch.Tensor] = None,
+ topk_weights: Optional[torch.Tensor] = None,
+ topk_rows: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """float64; the finalize (with `topk_rows`) rounded once to
+ `block_out`'s dtype, the combine (with `shared`) rounded once to `block_out`'s
+ dtype, `hc_out` rounded once to the activation dtype, `part` and `ss`
+ rounded once to fp32."""
+ dtype = hyper_input.dtype
+ rows, width = hyper_input.shape
+ hidden = width // hc_count
+ if block_out is None:
+ hc_out = hyper_input
+ else:
+ if topk_rows is not None:
+ valid = (topk_rows >= 0).unsqueeze(-1)
+ rows_ = block_out.double()[topk_rows.clamp(min=0).long()]
+ w = topk_weights.double().unsqueeze(-1)
+ block_out = torch.where(valid, rows_ * w, 0.0).sum(1).to(block_out.dtype)
+ if shared is not None:
+ sg = torch.sigmoid(gate_logits.double().reshape(rows, 1))
+ block_out = (block_out.double() + sg * shared.double()).to(block_out.dtype)
+ streams = hyper_input.double().unflatten(-1, (hc_count, hidden))
+ gated = block_out.double().unsqueeze(-2) * inject.double().unsqueeze(-1)
+ hc_out = (streams + gated).flatten(-2).to(dtype)
+ x = hc_out.double().unflatten(-1, (hc_count, hidden))
+ y = x * (1.0 + norm_weight.double()).unflatten(-1, (hc_count, hidden))
+ w = first[:, :n_cols].double().unflatten(0, (hc_count, hidden))
+ part = torch.einsum("tsi,sin->stn", y, w).float()
+ ss = x.pow(2).sum(dim=-1).t().float()
+ return hc_out, part, ss
+
+ def is_valid(
+ self,
+ hyper_input: torch.Tensor,
+ block_out: Optional[torch.Tensor],
+ inject: Optional[torch.Tensor],
+ norm_weight: torch.Tensor,
+ first: torch.Tensor,
+ hc_count: int,
+ n_cols: int,
+ shared: Optional[torch.Tensor] = None,
+ gate_logits: Optional[torch.Tensor] = None,
+ topk_weights: Optional[torch.Tensor] = None,
+ topk_rows: Optional[torch.Tensor] = None,
+ ) -> None:
+ # Raw row strides with unit column stride everywhere; a contiguous
+ # [width, N] weight would be read with its rows and columns swapped.
+ assert hyper_input.dim() == 2 and hyper_input.stride(1) == 1, "hyper_input [T, width]"
+ rows, width = hyper_input.shape
+ assert width % hc_count == 0 and (width // hc_count) % 16 == 0, "hidden multiple of 16"
+ assert (block_out is None) == (inject is None), "block_out and inject come together"
+ assert (shared is None) == (gate_logits is None), "shared and gate_logits come together"
+ assert shared is None or block_out is not None, "the combine needs block_out"
+ assert (topk_weights is None) == (topk_rows is None), (
+ "topk_weights and topk_rows come together"
+ )
+ assert topk_rows is None or shared is not None, "the finalize fold needs the combine"
+ if block_out is not None:
+ if topk_rows is None:
+ assert block_out.shape == (rows, width // hc_count) and block_out.stride(1) == 1
+ else:
+ # The permuted FC2 rows: any row count, read through topk_rows.
+ assert block_out.dim() == 2 and block_out.shape[1] == width // hc_count
+ assert block_out.stride(1) == 1, "permuted rows [P, hidden]"
+ assert topk_rows.dtype == torch.int32, "topk_rows int32 (the runner's index map)"
+ assert topk_rows.dim() == 2 and topk_rows.shape[0] == rows, "topk_rows [T, k]"
+ assert topk_rows.stride(1) == 1 and topk_weights.stride(1) == 1
+ assert topk_weights.shape == topk_rows.shape, "topk_weights [T, k]"
+ assert inject.shape == (rows, hc_count) and inject.stride(1) == 1
+ if shared is not None:
+ assert shared.shape == (rows, width // hc_count), "shared [T, hidden]"
+ assert shared.stride(1) == 1, "shared [T, hidden]"
+ assert shared.dtype == block_out.dtype, "shared in block_out's dtype"
+ assert gate_logits.shape in ((rows,), (rows, 1)), "gate_logits [T, 1] or [T]"
+ assert rows <= 1 or gate_logits.stride(0) != 0, "one gate logit per row"
+ assert norm_weight.shape == (width,) and norm_weight.is_contiguous(), "norm_weight [width]"
+ assert first.shape[0] == width and first.shape[1] >= n_cols, "first [width, >= n_cols]"
+ assert first.stride(0) == 1, "first must be the .t() view of a row-major [N, width] weight"
+ assert hyper_input.shape[0] <= MAX_ROWS, "one row"
+ assert hyper_input.shape[1] // hc_count == THREADS * VEC, "hidden 2560"
+ assert hyper_input.dtype == torch.bfloat16, "bf16"
+ outs_for(n_cols, hc_count)
+
+
+qwen4_exp_hc_decode_first_cute = Qwen4ExpHcDecodeFirstCute()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py
new file mode 100644
index 000000000000..df11a7567828
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py
@@ -0,0 +1,534 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Decode-row hyper-connection site at an MoE site, second half, for one-row batches:
+the grouped RMS, the SiLU bottleneck and the injection weights from the first half's
+partials, the up projection, the sigmoid gate, the stream mean and the NVFP4 of the
+mean for the routed experts -- one PDL launch.
+
+Tier 5 (CuTe DSL written for this target). A CTA owns one 16-column NVFP4 block of the
+hidden width across the 4 streams (64 rows of the up weight), and 8 threads share a
+row, 40 weights each; 160 CTAs of 512 threads, all resident at once.
+
+Order of work inside a CTA:
+
+1. its 64 weight rows and its norm-weight elements into registers (static), before
+ `griddepcontrol.wait`;
+2. after the wait: the four sums of squares (every thread), one partial column per
+ thread for the first 324 threads -> the shared `low` vector and, in CTA 0, the
+ injection weights; each thread's residual element;
+3. each row's dot with `low` over its 8 threads, the sigmoid gate on the normed
+ element, the mean over streams (16 threads), the store; with NVFP4, warp 0 also
+ quantizes the block as `torch.ops.trtllm.fp4_quantize` does (linear scales,
+ `rcp.approx.ftz` reciprocals).
+"""
+
+from __future__ import annotations
+
+import functools
+
+import cuda.bindings.driver as cuda_drv
+import cutlass
+import cutlass.cute as cute
+import torch
+import torch.nn.functional as F
+from cutlass import const_expr
+from cutlass._mlir.dialects import llvm
+from cutlass.cute import experimental as cute_ext
+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.quantization.fp4_quantize import (
+ fp4_quantize,
+)
+from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+
+MAX_ROWS = 1
+VEC = 8
+#: Hidden columns per CTA: one NVFP4 scale block.
+BLOCK_H = 16
+#: Threads sharing one up-weight row.
+GROUP = 8
+_COMPILE_OPTIONS = "--ptxas-options -maxrregcount=64"
+
+
+@dsl_user_op
+def _lds_v4f(addr, *, loc=None, ip=None):
+ """16 bytes of this CTA's shared memory as 4 fp32 (one `LDS.128`)."""
+ ty = T.f32()
+ res = llvm.inline_asm(
+ llvm.StructType.get_literal([ty] * 4),
+ [cutlass.Int32(addr).ir_value(loc=loc, ip=ip)],
+ "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];",
+ "=f,=f,=f,=f,r",
+ has_side_effects=True,
+ is_align_stack=False,
+ asm_dialect=llvm.AsmDialect.AD_ATT,
+ loc=loc,
+ ip=ip,
+ )
+ return [cutlass.Float32(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(4)]
+
+
+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: the logit is
+ # formed on every thread (only its store is predicated), so the partials' loads have
+ # a use outside any branch and are not sunk behind the sums of squares' divisions.
+ 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)
+ # This row's stream's rstd by selects (a dynamic register index would go to local memory).
+ rstd_s = rstd[0]
+ for t in cutlass.range_constexpr(1, streams):
+ if s == t:
+ rstd_s = rstd[t]
+ smem = cutlass.utils.SmemAllocator()
+ # `low` flat with 16 bytes of pad every 32 values: the 8 lanes of a row read 16-byte
+ # vectors of 8 different 8-value rows, in 8 different bank groups.
+ low = smem.allocate_tensor(
+ cutlass.Float32, cute.make_layout((rank + rank // 8,), stride=(1,)), byte_alignment=16
+ )
+ low_a = cutlass.Int32(low.iterator.toint())
+ vals = smem.allocate_tensor(
+ cutlass.Float32,
+ cute.make_layout((streams, BLOCK_H), stride=(BLOCK_H, 1)),
+ byte_alignment=16,
+ )
+ 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 + (tidx // 32) * 4] = z * sig
+ else:
+ if const_expr(self.inject_out):
+ if tidx < cols:
+ 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):
+ lrow = kt * GROUP + q
+ lv = _lds_v4f(low_a + (lrow * VEC + (lrow // 4) * 4) * 4) + _lds_v4f(
+ low_a + (lrow * VEC + (lrow // 4) * 4) * 4 + 16
+ )
+ for e in cutlass.range_constexpr(VEC):
+ sums[e % 2] = u_reg[kt, e].to(cutlass.Float32) * lv[e] + sums[e % 2]
+ dot = sums[0] + sums[1]
+ dot = cute.arch.warp_reduction(dot, lambda a, b: a + b, threads_in_group=GROUP)
+ if q == 0:
+ gate = 1.0 / (1.0 + cute.math.exp(-dot))
+ vals[s, i] = gate * (x * rstd_s * gam)
+ cute.arch.sync_threads()
+ if warp == 0:
+ m = cutlass.Float32(0.0)
+ if lane < BLOCK_H:
+ tot = cutlass.Float32(0.0)
+ for t in cutlass.range_constexpr(streams):
+ tot = tot + vals[t, lane]
+ mb = (tot * (1.0 / streams)).to(cutlass.BFloat16)
+ mMixed[0, i0 + lane] = mb
+ m = mb.to(cutlass.Float32)
+ if const_expr(self.fp4):
+ # NVFP4 of the 16 rounded values in lanes 0..15 (lanes 16..31 a second,
+ # discarded group), as `fp4_quantize`: e4m3(gs * amax / 6), then
+ # x * rcp(sf * rcp(gs)) rounded to e2m1, nearest, ties to even.
+ amax = cute.arch.warp_reduction_max(cute.math.abs(m), threads_in_group=BLOCK_H)
+ sf8 = (gs * (amax * cute.arch.rcp_approx(6.0))).to(cutlass.Float8E4M3FN)
+ osc = cute.arch.rcp_approx(sf8.to(cutlass.Float32) * cute.arch.rcp_approx(gs))
+ if amax == 0.0:
+ osc = cutlass.Float32(0.0)
+ sv = m * osc
+ a = cute.math.abs(sv)
+ code = cutlass.Int32(0)
+ if a > 0.25:
+ code = code + 1
+ if a >= 0.75:
+ code = code + 1
+ if a > 1.25:
+ code = code + 1
+ if a >= 1.75:
+ code = code + 1
+ if a > 2.5:
+ code = code + 1
+ if a >= 3.5:
+ code = code + 1
+ if a > 5.0:
+ code = code + 1
+ if cute.math.copysign(1.0, sv) < 0.0: # the sign bit, -0.0 included
+ code = code + 8
+ hi = cute.arch.shuffle_sync_down(code, offset=1)
+ if lane < BLOCK_H:
+ if lane % 2 == 0:
+ mXq[0, (i0 + lane) // 2] = (code + hi * 16).to(cutlass.Uint8)
+ if lane == 0:
+ mXsf[bidx] = sf8
+
+
+def _dl(t: torch.Tensor, align: int = 16):
+ return from_dlpack(t, assumed_align=align)
+
+
+@functools.lru_cache(maxsize=None)
+def _compiled(
+ hidden, streams, rank, part_cols, part_stride, ss_stride, hc_stride, inject_out, fp4, pdl
+):
+ kern = _MixKernel(
+ hidden=hidden, streams=streams, rank=rank, inject_out=inject_out, fp4=fp4, pdl=pdl
+ )
+ dev, bf = "cuda", torch.bfloat16
+ width = hidden * streams
+ part = torch.empty(streams, part_stride, dtype=torch.float32, device=dev)[:, :part_cols]
+ ss = torch.empty(streams * ss_stride, dtype=torch.float32, device=dev)[::ss_stride]
+ up = torch.empty(width, rank, dtype=bf, device=dev)
+ hc = torch.empty(1, hc_stride, dtype=bf, device=dev)[:, :width]
+ norm = torch.empty(1, width, dtype=bf, device=dev)
+ gs = torch.empty(1, dtype=torch.float32, device=dev)
+ mixed = torch.empty(1, hidden, dtype=bf, device=dev)
+ inj = torch.empty(1, streams, dtype=bf, device=dev)
+ xq = torch.empty(1, hidden // 2, dtype=torch.uint8, device=dev)
+ xsf = torch.empty(hidden // 16, dtype=torch.float8_e4m3fn, device=dev)
+ stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream)
+ return cute_ext.compile(
+ kern,
+ _dl(part, 4),
+ _dl(ss, 4),
+ _dl(up),
+ _dl(hc, 2),
+ _dl(norm, 2),
+ _dl(gs, 4),
+ _dl(mixed, 2),
+ _dl(inj, 2),
+ _dl(xq, 1),
+ _dl(xsf, 1),
+ cutlass.Float32(1e-6),
+ stream,
+ options=_COMPILE_OPTIONS,
+ )
+
+
+def launch(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale=None):
+ """The launch behind `Qwen4ExpHcDecodeMixFp4Cute.__call__`. Returns `(mixed,
+ inject_out, xq, xsf)`; without `global_scale` (the kernel's plain build, kept for
+ measurement) `xq` / `xsf` are None."""
+ rows, width = hc_out.shape
+ hidden = width // hc_count
+ rank = up.shape[0]
+ dev, dt = hc_out.device, hc_out.dtype
+ fp4 = global_scale is not None
+ mixed = torch.empty((rows, hidden), dtype=dt, device=dev)
+ inject_out = torch.empty((rows, hc_count), dtype=dt, device=dev) if with_inject else None
+ xq = torch.empty((rows, hidden // 2), dtype=torch.uint8, device=dev) if fp4 else None
+ xsf = torch.empty((rows * (hidden // 16),), dtype=torch.uint8, device=dev) if fp4 else None
+ if rows == 0:
+ return mixed, inject_out, xq, xsf
+ pdl = get_env_enable_pdl()
+ part2 = part[:, 0, :]
+ fn = _compiled(
+ hidden,
+ hc_count,
+ rank,
+ part2.shape[1],
+ part2.stride(0),
+ ss.stride(0),
+ hc_out.stride(0),
+ bool(with_inject),
+ fp4,
+ bool(pdl),
+ )
+ stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream)
+ gs = global_scale.reshape(1) if fp4 else mixed.new_empty(1, dtype=torch.float32)
+ inj = inject_out if with_inject else mixed.new_empty(1, hc_count)
+ q = xq if fp4 else mixed.new_empty(1, hidden // 2, dtype=torch.uint8)
+ sf = (
+ xsf.view(torch.float8_e4m3fn)
+ if fp4
+ else mixed.new_empty(hidden // 16, dtype=torch.float8_e4m3fn)
+ )
+ fn(
+ _dl(part2, 4),
+ _dl(ss[:, 0], 4),
+ _dl(up.t()),
+ _dl(hc_out, 2),
+ _dl(norm_weight.reshape(1, -1), 2),
+ _dl(gs, 4),
+ _dl(mixed, 2),
+ _dl(inj, 2),
+ _dl(q, 1),
+ _dl(sf, 1),
+ cutlass.Float32(eps),
+ stream,
+ )
+ return mixed, inject_out, xq, xsf
+
+
+def _check(part, ss, up, hc_out, norm_weight, hc_count, with_inject) -> None:
+ rows, width = hc_out.shape
+ rank = up.shape[0]
+ assert hc_out.stride(1) == 1, "hc_out unit column stride"
+ assert part.dtype == torch.float32 and part.stride(2) == 1, "part fp32, unit last stride"
+ assert part.shape[:2] == (hc_count, rows), "part [C, T, N]"
+ assert part.shape[2] >= rank + (hc_count if with_inject else 0), "part too narrow"
+ assert ss.dtype == torch.float32 and ss.shape == (hc_count, rows), "ss fp32 [C, T]"
+ assert up.shape == (rank, width), "up [R, width]"
+ assert norm_weight.shape == (width,) and norm_weight.is_contiguous()
+ assert rows <= MAX_ROWS, "one row"
+ assert width // hc_count == 2560 and rank == 320 and hc_count == 4, (
+ "the checkpoint's 4 x 2560, rank 320"
+ )
+ assert hc_out.dtype == torch.bfloat16 and up.dtype == torch.bfloat16, "bf16"
+ assert up.t().is_contiguous(), "up the [R, width] view of a row-major [width, R] weight"
+
+
+def _reference(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject):
+ """`(mixed, inject_out)` in float64, one rounding per output."""
+ dtype = hc_out.dtype
+ rank = up.shape[0]
+ hidden = hc_out.shape[1] // hc_count
+ rstd = torch.rsqrt(ss.double().t() / hidden + eps)
+ logits = torch.einsum("ts,stn->tn", rstd, part.double())
+ low = F.silu(logits[:, :rank] / hc_count)
+ gate = torch.sigmoid(low @ up.double())
+ x = hc_out.double().unflatten(-1, (hc_count, hidden)) * rstd.unsqueeze(-1)
+ normed = x.flatten(-2) * (1.0 + norm_weight.double())
+ mixed = (gate * normed).unflatten(-1, (hc_count, hidden)).mean(dim=-2).to(dtype)
+ inject_out = None
+ if with_inject:
+ inject_out = (2.0 * torch.sigmoid(logits[:, rank : rank + hc_count] / hc_count)).to(dtype)
+ return mixed, inject_out
+
+
+_INJECT = Cell(why="checkpoint 4x2560 r320 bf16, inject, 1 row (con 1)", spec=dict(rows=1))
+_ROW_SLICE = Cell(why="hc_out a row slice buf[::2], 1 row", spec=dict(rows=1, layout="row_slice"))
+
+
+class Qwen4ExpHcDecodeMixFp4Cute(OpWrapper):
+ """`qwen4_exp_hc_decode_mix_cute` plus the NVFP4 of its `mixed` for the routed
+ experts. With `C = hc_count`, `H = width // C`, `R = up.shape[0]`, `part` / `ss` the
+ per-stream partial logits `[C, T, >= R (+ C)]` and sums of squares `[C, T]` from
+ `qwen4_exp_hc_decode_first_cute`, `up` the `[R, width]` view of the up weight::
+
+ rstd[t, s] = 1 / sqrt(ss[s, t] / H + eps)
+ logits[t, n] = sum_s rstd[t, s] * part[s, t, n]
+ low = silu(logits[:, :R] / C)
+ gate = sigmoid(low @ up) # [T, C*H]
+ normed[t, s*H+i] = hc_out[t, s*H+i] * rstd[t, s] * (1 + norm_weight[s*H+i])
+ mixed[t, i] = mean over s of gate[t, s*H+i] * normed[t, s*H+i]
+ inject_out = 2 * sigmoid(logits[:, R:R+C] / C) # with_inject only
+ xq, xsf = NVFP4(mixed) # e2m1 pairs, e4m3 scale per 16, linear layout
+
+ fp32 inside, one rounding per output; the quantization reads the rounded `mixed`.
+ Returns fresh `(mixed, inject_out, xq [T, H/2] uint8, xsf [T*H/16] uint8)`;
+ `inject_out` is None without `with_inject`. `global_scale` is the 1-element fp32
+ NVFP4 global scale.
+
+ Fusion boundary: everything `qwen4_exp_hc_decode_mix_cute` fuses, plus the NVFP4
+ quantization of `mixed`. The partials are the producer's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+ # Its call sites always inject, so no stack-final (no-inject) cell.
+ CELLS: tuple[Cell, ...] = (_INJECT, _ROW_SLICE)
+
+ note = """
+ Metadata consumed: none. Call sites: `DecodeTarget`'s MoE sites at a one-row
+ decode step, after `qwen4_exp_hc_decode_first_cute`; `xq`/`xsf` feed
+ `fp4_block_scale_moe_renorm_routed` (linear scales). Batches of more rows take
+ `qwen4_exp_hc_up_gate_mean_fp4`.
+
+ Preconditions: those of `qwen4_exp_hc_decode_mix_cute`, and `global_scale` fp32 with
+ one element.
+
+ Numerics: `mixed` and `inject_out` as `qwen4_exp_hc_decode_mix_cute`; the NVFP4 is
+ `fp4_quantize`'s arithmetic on the rounded `mixed` (`rcp.approx.ftz` reciprocals,
+ e4m3 scale rounded to nearest, e2m1 nearest-even with saturation at 6).
+ Deterministic (measured); graph-capturable once compiled (the first call per
+ variant). PDL: the weight rows and norm elements stream before
+ `griddepcontrol.wait`.
+
+ Certified: one bf16 row with the injection, `hc_out` compact and a row slice;
+ graph replay under PDL. Not certified: `with_inject=False`.
+ """
+
+ def __call__(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale):
+ return launch(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale)
+
+ def reference(
+ self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale
+ ):
+ """float64 with one rounding per output, then `fp4_quantize`'s native-torch NVFP4
+ of the rounded `mixed` (linear scales)."""
+ mixed, inject_out = _reference(
+ part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject
+ )
+ return (mixed, inject_out) + fp4_quantize.reference(
+ mixed, global_scale, BLOCK_H, False, False
+ )
+
+ def is_valid(
+ self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale
+ ) -> None:
+ _check(part, ss, up, hc_out, norm_weight, hc_count, with_inject)
+ # The kernel loads element 0 only; a per-row scale would be applied as row 0's.
+ assert global_scale.numel() == 1 and global_scale.dtype == torch.float32, (
+ "global_scale fp32 [1]"
+ )
+
+
+qwen4_exp_hc_decode_mix_fp4_cute = Qwen4ExpHcDecodeMixFp4Cute()
+
+
+class Qwen4ExpHcDecodeMixCute(OpWrapper):
+ """The second half of a one-row decode hyper-connection site: the grouped RMS, the
+ SiLU bottleneck and the injection weights from `qwen4_exp_hc_decode_first_cute`'s
+ per-stream partials, then the up projection, the sigmoid gate on the normed stream
+ and the mean over streams -- `qwen4_exp_hc_decode_mix_fp4_cute`'s formulas without
+ the NVFP4 output. Returns fresh `(mixed, inject_out)` in `hc_out`'s dtype;
+ `inject_out` is None without `with_inject`.
+
+ Fusion boundary: the grouped RMS, the bottleneck SiLU, the injection weights, the
+ up projection, the sigmoid gate, the normed stream and the stream mean. The
+ partials are the producer's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+ CELLS: tuple[Cell, ...] = (
+ _INJECT,
+ Cell(why="stack-final mixer (no inject), 1 row", spec=dict(rows=1, with_inject=False)),
+ _ROW_SLICE,
+ )
+
+ note = """
+ Metadata consumed: none. Call sites: `DecodeTarget`'s non-MoE sites at a one-row
+ decode step, after `qwen4_exp_hc_decode_first_cute`; batches of more rows take
+ `qwen4_exp_hc_up_gate_mean`. Beside a first kernel capped at 96 registers, its 160
+ CTAs (512 threads at 49 registers) are all resident while that kernel runs, so the
+ up weight is in registers when the wait returns.
+
+ Preconditions: `part` fp32 `[C, T, >= R (+ C)]` and `ss` fp32 `[C, T]`, `part` with
+ unit last stride; `hc_out` `[T, width]` with unit column stride; `norm_weight`
+ contiguous `[width]`; one row, the checkpoint's 4 x 2560 width, rank 320, bf16;
+ `up` the `.t()` view of a contiguous `[width, R]` weight.
+
+ Numerics: fp32 inside (`low`, the gate logits and `normed` are never rounded), one
+ rounding per output, inside the float64 reference's band. Deterministic
+ (measured); graph-capturable once compiled (the first call per variant).
+
+ Certified: one bf16 row with and without the injection, `hc_out` compact and a
+ row slice. Not certified: graph replay of this build (the NVFP4 build's
+ replay test covers the shared kernel, not this compiled variant).
+ """
+
+ def __call__(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject):
+ mixed, inject_out, _, _ = launch(
+ part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject
+ )
+ return mixed, inject_out
+
+ def reference(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject):
+ """float64, one rounding per output."""
+ return _reference(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject)
+
+ def is_valid(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject) -> None:
+ _check(part, ss, up, hc_out, norm_weight, hc_count, with_inject)
+
+
+qwen4_exp_hc_decode_mix_cute = Qwen4ExpHcDecodeMixCute()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_rows_tc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_rows_tc.py
new file mode 100644
index 000000000000..39c27a557c7b
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_rows_tc.py
@@ -0,0 +1,1353 @@
+"""Decode hyper-connection site for 2..16 rows on warp tensor cores: two PDL launches, the
+first emitting the bottleneck `low`, the second the up projection, gate, stream mean and NVFP4.
+
+Tier 5 (CuTe DSL written for this target). A cluster spans the 4 streams x 2 K halves, so the
+first launch finishes the site's first half itself:
+
+first (`_First2`): grid (2 K halves, 4 streams, 14 N chunks of 24 columns), cluster (2, 4, 1),
+8 warps. A CTA forms its rows' injection over its 1280-element K half (rounded once; stored by the
+chunk-0 CTAs), the slice's sums of squares and `y = x * (1 + w)` in bf16; each warp multiplies a
+160-wide K slice of `y` by the CTA's 24 weight rows (staged before the PDL wait) on mma.sync;
+the warps' partials are summed in shared memory. Every CTA pushes rows (2o, 2o+1) of its partial
+tile and of its sums of squares to cluster rank o; rank o then forms, per row,
+`rstd_s = (ss_s / H + eps)^-1/2`, the logits `z = sum_s rstd_s (P_s,0 + P_s,1) / C`, and writes
+`low = bf16(silu(z))` (columns < 320), the injection weights `2 sigmoid(z)` (columns 320..323)
+and the rows' rstd.
+
+mix (`_Mix2`): 160 CTAs of 8 warps (16 hidden columns x 4 streams). Before the wait: the 64
+up-weight rows. After: `low` (640 bytes a row), rstd and the CTA's residual elements; the
+[16 x 64] up projection; gate x normed; the stream mean and its NVFP4.
+
+Accuracy as `hc_tc.py` (bf16 operands into both projections; `hc_out` bitwise).
+"""
+
+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.nvgpu import cpasync, warp
+from cutlass.cute.runtime import from_dlpack
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_inject_norm import (
+ _asm,
+ _mapa,
+)
+from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+
+#: This checkpoint's hyper-connection bottleneck width.
+RANK = 320
+#: Rows a launch serves (one m16 MMA tile); the forward sends 2..MAX_ROWS-row sites here.
+MAX_ROWS = 16
+M = 16
+VEC = 8
+PAD = 8
+KH = 2 # K halves
+KHW = 1280
+NC = 24 # columns a CTA
+NCH = 14 # 14 x 24 = 336 >= 324
+WARPS = 8
+THREADS = WARPS * 32
+KW = KHW // WARPS # K a warp
+CL = KH * 4 # cluster size
+RPO = M // CL
+#: Item iterations a thread whose activation loads issue together (all of them at <= 16 rows).
+GROUP_IT = 10
+#: Rows from which the finalize fold is split over a K half's 4 CTAs (each forms a quarter of
+#: the half's rows and pushes it to all four); below it each CTA folds its whole K half itself
+#: in its items (4x the gathers, no exchange: 1.4-1.7 us a layer faster at 2 and 3 rows behind
+#: the runner, 1.7-2.3 us slower from 4, in the chained [MoE + site + MoE site] bench).
+FOLD_SPLIT_ROWS = 4
+#: Fold vectors a thread whose gathers issue together, split (all of them at <= 16 rows) and not
+#: (ten 16-byte gathers a vector: register budget).
+GROUP_FOLD = 3
+GROUP_FOLD_ITEMS = 2
+#: Rows of First2's A stage allocated at least: its shared-memory footprint (~90 KB at 2..8
+#: rows, ~118 KB at 16) leaves room beside it for the mix kernel's CTAs of the same site.
+A_MIN_ROWS = 8
+BLOCK_H = 16
+
+
+@dsl_user_op
+def _st_dsmem_f32(addr, val, *, loc=None, ip=None):
+ """4 bytes into a cluster CTA's shared memory."""
+ _asm(
+ T.i32(),
+ [cutlass.Int32(addr), cutlass.Float32(val)],
+ "{ st.shared::cluster.f32 [$1], $2; mov.u32 $0, 0; }",
+ "=r,r,f",
+ loc,
+ ip,
+ side_effects=True,
+ )
+
+
+@dsl_user_op
+def _st_async_f32(addr, val, mbar, *, loc=None, ip=None):
+ """4 bytes into a cluster CTA's shared memory, completing on that CTA's mbarrier."""
+ _asm(
+ T.i32(),
+ [cutlass.Int32(addr), cutlass.Float32(val), cutlass.Int32(mbar)],
+ "{ st.async.shared::cluster.mbarrier::complete_tx::bytes.f32 [$1], $2, [$3]; "
+ "mov.u32 $0, 0; }",
+ "=r,r,f,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+
+
+@dsl_user_op
+def _mbar_init_tx(mbar, nbytes, *, loc=None, ip=None):
+ """Init an mbarrier (one arrival), make it visible to the cluster, arrive expecting
+ `nbytes` of transactions."""
+ _asm(
+ T.i32(),
+ [cutlass.Int32(mbar), cutlass.Int32(nbytes)],
+ "{ mbarrier.init.shared::cta.b64 [$1], 1; fence.mbarrier_init.release.cluster; "
+ "mbarrier.arrive.expect_tx.shared::cta.b64 _, [$1], $2; mov.u32 $0, 0; }",
+ "=r,r,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+
+
+@dsl_user_op
+def _mbar_wait0(mbar, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [cutlass.Int32(mbar)],
+ "{ .reg .pred p; WAIT: mbarrier.try_wait.parity.shared::cta.b64 p, [$1], 0; "
+ "@!p bra WAIT; mov.u32 $0, 0; }",
+ "=r,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+
+
+@dsl_user_op
+def _st_async_v4(addr, a, b, c, d, mbar, *, loc=None, ip=None):
+ """16 bytes into a cluster CTA's shared memory, completing on that CTA's mbarrier."""
+ _asm(
+ T.i32(),
+ [cutlass.Int32(addr)] + [cutlass.Int32(v) for v in (a, b, c, d)] + [cutlass.Int32(mbar)],
+ "{ st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [$1], {$2, $3, $4, $5}, "
+ "[$6]; mov.u32 $0, 0; }",
+ "=r,r,r,r,r,r,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+
+
+@dsl_user_op
+def _bf16x2(lo, hi, *, loc=None, ip=None):
+ """Two fp32 values rounded (RN) into one bf16x2 word, `lo` low (exact for bf16 values)."""
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [cutlass.Float32(hi), cutlass.Float32(lo)],
+ "cvt.rn.bf16x2.f32 $0, $1, $2;",
+ "=r,f,f",
+ loc,
+ ip,
+ )
+ )
+
+
+class _First2:
+ def __init__(self, *, hidden, rank, n_cols, rows, inject, combine, pdl, topk=0):
+ assert hidden == KH * KHW and n_cols <= NC * NCH and rows <= M
+ assert topk == 0 or combine, "the finalize fold needs the combine"
+ self.hidden, self.rank, self.n_cols, self.rows = hidden, rank, n_cols, rows
+ self.inject, self.combine, self.pdl, self.topk = inject, combine, pdl, topk
+ self.ld = KHW + PAD
+ self.rld = NC + 2 # pushed slot: NC partials + the sum of squares (+pad)
+
+ def smem_bytes(self):
+ def r(n):
+ return -(-n // 128) * 128
+
+ assert WARPS * M * NC * 4 + M * NC * 4 <= NC * self.ld * 2
+ b = r(max(self.rows, A_MIN_ROWS) * self.ld * 2) + r(NC * self.ld * 2)
+ b += r(self.rows * (KHW // VEC) * 4)
+ b += r(CL * RPO * self.rld * 4) + r(M * 4)
+ b += r(KHW * 2) # the norm-weight slice
+ if self.topk:
+ b += 2 * r(self.rows * self.topk * 4) # the fold's slot rows and weights
+ if self.topk and self.rows >= FOLD_SPLIT_ROWS:
+ b += r(self.rows * KHW * 2) + 128 # the block rows of the K half, and an mbarrier
+ return b + 256 + 128
+
+ @cute.jit
+ def __call__(
+ self,
+ mHc,
+ mBlock,
+ mInj,
+ mShared,
+ mGate,
+ mTopW,
+ mTopRows,
+ mNorm,
+ mW,
+ mHcOut,
+ mLow,
+ mInjOut,
+ mRstd,
+ eps: cutlass.Float32,
+ stream: cuda_drv.CUstream,
+ ):
+ mma = cute.make_tiled_mma(
+ warp.MmaF16BF16Op(cutlass.BFloat16, cutlass.Float32, (16, 8, 16)),
+ cute.make_layout((1, 1, 1)),
+ )
+ g2s = cute.make_copy_atom(cpasync.CopyG2SOp(), cutlass.BFloat16, num_bits_per_copy=128)
+ self.kernel(
+ mHc,
+ mBlock,
+ mInj,
+ mShared,
+ mGate,
+ mTopW,
+ mTopRows,
+ mNorm,
+ mW,
+ mHcOut,
+ mLow,
+ mInjOut,
+ mRstd,
+ eps,
+ mma,
+ g2s,
+ ).launch(
+ grid=[KH, 4, NCH],
+ block=[THREADS, 1, 1],
+ cluster=[KH, 4, 1],
+ smem=self.smem_bytes(),
+ stream=stream,
+ use_pdl=self.pdl,
+ min_blocks_per_mp=1,
+ )
+
+ @cute.kernel
+ def kernel(
+ self,
+ mHc,
+ mBlock,
+ mInj,
+ mShared,
+ mGate,
+ mTopW,
+ mTopRows,
+ mNorm,
+ mW,
+ mHcOut,
+ mLow,
+ mInjOut,
+ mRstd,
+ eps: cutlass.Float32,
+ mma: cute.TiledMma,
+ g2s: cute.CopyAtom,
+ ):
+ tidx, _, _ = cute.arch.thread_idx()
+ kh, s, nc = cute.arch.block_idx()
+ warp_i = cute.arch.warp_idx()
+ lane = cute.arch.lane_idx()
+ rows: cutlass.Constexpr = self.rows
+ hidden: cutlass.Constexpr = self.hidden
+ rank: cutlass.Constexpr = self.rank
+ n_cols: cutlass.Constexpr = self.n_cols
+ ld: cutlass.Constexpr = self.ld
+ rld: cutlass.Constexpr = self.rld
+ nvec: cutlass.Constexpr = KHW // VEC
+ k0 = kh * KHW
+ n0 = nc * NC
+ me = kh + KH * s # cluster rank
+
+ smem = cutlass.utils.SmemAllocator()
+ # sA holds the R live rows only; the MMA's 16-row A view runs on into sW (rows past R
+ # are garbage, their outputs discarded: the MMA's rows are independent).
+ sA = smem.allocate_tensor(
+ cutlass.BFloat16,
+ cute.make_layout((max(rows, A_MIN_ROWS), KHW), stride=(ld, 1)),
+ byte_alignment=128,
+ )
+ sW = smem.allocate_tensor(
+ cutlass.BFloat16, cute.make_layout((NC, KHW), stride=(ld, 1)), byte_alignment=128
+ )
+ sSs = smem.allocate_tensor(
+ cutlass.Float32, cute.make_layout((rows, nvec), stride=(nvec, 1)), byte_alignment=128
+ )
+ # The warps' partials and their sum reuse sW once every warp holds its B fragment.
+ wp_ptr = cute.recast_ptr(sW.iterator, dtype=cutlass.Float32)
+ sWp = cute.make_tensor(wp_ptr, cute.make_layout((WARPS, M, NC), stride=(M * NC, NC, 1)))
+ sAcc = cute.make_tensor(wp_ptr + WARPS * M * NC, cute.make_layout((M, NC), stride=(NC, 1)))
+ sRed = smem.allocate_tensor(
+ cutlass.Float32,
+ cute.make_layout((CL, RPO, rld), stride=(RPO * rld, rld, 1)),
+ byte_alignment=128,
+ )
+ sSsr = smem.allocate_tensor(
+ cutlass.Float32, cute.make_layout((M,), stride=(1,)), byte_alignment=128
+ )
+ sNw = smem.allocate_tensor(
+ cutlass.BFloat16, cute.make_layout((1, KHW), stride=(KHW, 1)), byte_alignment=128
+ )
+ topk: cutlass.Constexpr = self.topk
+ split: cutlass.Constexpr = topk > 0 and rows >= FOLD_SPLIT_ROWS
+ if const_expr(topk > 0):
+ sTopR = smem.allocate_tensor(
+ cutlass.Int32, cute.make_layout((rows * topk,), stride=(1,)), byte_alignment=128
+ )
+ sTopW = smem.allocate_tensor(
+ cutlass.Float32, cute.make_layout((rows * topk,), stride=(1,)), byte_alignment=128
+ )
+ sBar = smem.allocate_tensor(
+ cutlass.Int64, cute.make_layout((2 if split else 1,)), byte_alignment=8
+ )
+ rbar = cutlass.Int32(sBar.iterator.toint())
+ bbar = rbar + 8
+ if const_expr(split):
+ # The K half's block rows (finalize + combine), a quarter formed by each of the 4
+ # stream CTAs of this K half and pushed (st.async) into all four.
+ sBlk = smem.allocate_tensor(
+ cutlass.BFloat16, cute.make_layout((rows, KHW), stride=(KHW, 1)), byte_alignment=128
+ )
+ blk_a = cutlass.Int32(sBlk.iterator.toint())
+ sblk_vec = cute.logical_divide(sBlk, (None, VEC))
+ # This CTA receives rows [me * RPO, me * RPO + RPO) of the cluster's partials from all CL
+ # CTAs, by st.async on `rbar` (no cluster barrier, no fence, after the pushes).
+ rows_me = cutlass.max(cutlass.min(cutlass.Int32(RPO), rows - me * RPO), cutlass.Int32(0))
+ if tidx == 0:
+ _mbar_init_tx(rbar, rows_me * (CL * (NC + 1) * 4))
+ if const_expr(split):
+ _mbar_init_tx(bbar, rows * KHW * 2)
+ cute.arch.cluster_arrive()
+
+ # 1. Before the wait: the norm weight's slice (its own cp.async group, waited on right
+ # after the PDL wait), then the CTA's 24 weight rows over its K half.
+ nw_vec = cute.logical_divide(mNorm, (None, VEC))
+ snw_vec = cute.logical_divide(sNw, (None, VEC))
+ if tidx < KHW // VEC:
+ cute.copy(
+ g2s,
+ nw_vec[0, (None, (s * hidden + k0) // VEC + tidx)],
+ snw_vec[0, (None, tidx)],
+ )
+ cute.arch.cp_async_commit_group()
+ w_vec = cute.logical_divide(mW, (None, VEC))
+ sw_vec = cute.logical_divide(sW, (None, VEC))
+ pieces: cutlass.Constexpr = NC * nvec
+ col0 = (s * hidden + k0) // VEC
+ for i in cutlass.range_constexpr(-(-pieces // THREADS)):
+ c = tidx + i * THREADS
+ if c < pieces:
+ n = c // nvec
+ p = c % nvec
+ cute.copy(g2s, w_vec[n0 + n, (None, col0 + p)], sw_vec[n, (None, p)])
+ cute.arch.cp_async_commit_group()
+ if const_expr(self.pdl):
+ cute.arch.griddepcontrol_launch_dependents()
+ cute.arch.griddepcontrol_wait()
+
+ # 2. The K half's injection, sums of squares, y (bf16) into sA. Every activation load of
+ # a group of up to GROUP_IT items a thread (the row's injection weight and gate logit
+ # included) issues before the first is used: one L2 round trip a group.
+ hc_vec = cute.logical_divide(mHc, (None, VEC))
+ b_vec = cute.logical_divide(mBlock, (None, VEC))
+ sh_vec = cute.logical_divide(mShared, (None, VEC))
+ o_vec = cute.logical_divide(mHcOut, (None, VEC))
+ sa_vec = cute.logical_divide(sA, (None, VEC))
+ items: cutlass.Constexpr = rows * nvec
+ n_it: cutlass.Constexpr = -(-items // THREADS)
+ if const_expr(split):
+ # With the split fold, every load that does not need the slot map issues with it: the
+ # items' hc rows and injection weights (one group: <= 10 items a thread at <= 16
+ # rows), and the fold quarter's shared-expert rows and gate logits. The gathers
+ # are then the fold's one added L2 round trip.
+ assert n_it <= GROUP_IT
+ fvq: cutlass.Constexpr = KHW // 4 // VEC
+ f_items: cutlass.Constexpr = rows * fvq
+ n_f: cutlass.Constexpr = -(-f_items // THREADS)
+ assert n_f <= GROUP_FOLD
+ fcol0 = (kh * KHW + s * (KHW // 4)) // VEC
+ xg0 = cute.make_rmem_tensor(
+ cute.make_layout((n_it, VEC), stride=(VEC, 1)), cutlass.BFloat16
+ )
+ ig0 = cute.make_rmem_tensor(cute.make_layout((n_it,), stride=(1,)), cutlass.Float32)
+ for k in cutlass.range_constexpr(n_it):
+ c = cutlass.min(tidx + k * THREADS, cutlass.Int32(items - 1))
+ r = c // nvec
+ p = c % nvec
+ cute.autovec_copy(hc_vec[r, (None, col0 + p)], xg0[k, None])
+ ig0[k] = mInj[r, s].to(cutlass.Float32)
+ shq = cute.make_rmem_tensor(
+ cute.make_layout((n_f, VEC), stride=(VEC, 1)), cutlass.BFloat16
+ )
+ zq = cute.make_rmem_tensor(cute.make_layout((n_f,), stride=(1,)), cutlass.Float32)
+ for k in cutlass.range_constexpr(n_f):
+ c = cutlass.min(tidx + k * THREADS, cutlass.Int32(f_items - 1))
+ cute.autovec_copy(sh_vec[c // fvq, (None, fcol0 + c % fvq)], shq[k, None])
+ zq[k] = mGate[c // fvq].to(cutlass.Float32)
+ if const_expr(topk > 0):
+ # The rows' expanded->permuted slot map and routing weights (the gathers below
+ # depend on it).
+ for i in cutlass.range_constexpr(-(-(rows * topk) // THREADS)):
+ cs = tidx + i * THREADS
+ if cs < rows * topk:
+ sTopR[cs] = mTopRows[cs // topk, cs % topk]
+ sTopW[cs] = mTopW[cs // topk, cs % topk].to(cutlass.Float32)
+ cute.arch.cp_async_wait_group(1) # the norm slice (the weights may still be landing)
+ cute.arch.sync_threads()
+ if const_expr(split):
+ # The finalize and the combine of this CTA's quarter of its K half (40 vectors a
+ # row): the runner's fp32 FMAs in slot order from zero (a -1 slot skipped, a -1 slot
+ # reads row 0, unused), rounded once; then fma(sigmoid(gate), shared, routed) rounded
+ # once; pushed to the K half's 4 CTAs.
+ cute.arch.cluster_wait() # every peer has started (its mbarriers initialized)
+ for f0 in cutlass.range_constexpr(0, n_f, GROUP_FOLD):
+ f_n: cutlass.Constexpr = min(GROUP_FOLD, n_f - f0)
+ fgq = cute.make_rmem_tensor(
+ cute.make_layout((f_n, topk, VEC), stride=(topk * VEC, VEC, 1)),
+ cutlass.BFloat16,
+ )
+ for k in cutlass.range_constexpr(f_n):
+ c = cutlass.min(tidx + (f0 + k) * THREADS, cutlass.Int32(f_items - 1))
+ r = c // fvq
+ q = c % fvq
+ for j in cutlass.range_constexpr(topk):
+ cute.autovec_copy(
+ b_vec[
+ cutlass.max(sTopR[r * topk + j], cutlass.Int32(0)),
+ (None, fcol0 + q),
+ ],
+ fgq[k, j, None],
+ )
+ for k in cutlass.range_constexpr(f_n):
+ c = tidx + (f0 + k) * THREADS
+ if c < f_items:
+ r = c // fvq
+ q = c % fvq
+ sg = 1.0 / (1.0 + cute.math.exp(-zq[f0 + k]))
+ bq = [cutlass.Float32(0.0)] * VEC
+ for v in cutlass.range_constexpr(VEC):
+ acc = cutlass.Float32(0.0)
+ for j in cutlass.range_constexpr(topk):
+ if sTopR[r * topk + j] >= 0:
+ acc = cute.math.fma(
+ fgq[k, j, v].to(cutlass.Float32), sTopW[r * topk + j], acc
+ )
+ bv = acc.to(cutlass.BFloat16).to(cutlass.Float32)
+ bq[v] = (
+ cute.math.fma(sg, shq[f0 + k, v].to(cutlass.Float32), bv)
+ .to(cutlass.BFloat16)
+ .to(cutlass.Float32)
+ )
+ words = [_bf16x2(bq[2 * w], bq[2 * w + 1]) for w in range(4)]
+ dst = blk_a + (r * KHW + s * (KHW // 4) + q * VEC) * 2
+ for sp in cutlass.range_constexpr(4):
+ peer = kh + KH * sp
+ _st_async_v4(_mapa(dst, peer), *words, _mapa(bbar, peer))
+ nwr = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16)
+ outr = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16)
+ yr = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16)
+ group: cutlass.Constexpr = GROUP_FOLD_ITEMS if topk > 0 and not split else GROUP_IT
+ for g0 in cutlass.range_constexpr(0, n_it, group):
+ g_n: cutlass.Constexpr = min(group, n_it - g0)
+ if const_expr(split):
+ xg = xg0
+ else:
+ xg = cute.make_rmem_tensor(
+ cute.make_layout((g_n, VEC), stride=(VEC, 1)), cutlass.BFloat16
+ )
+ bg = cute.make_rmem_tensor(
+ cute.make_layout((g_n, VEC), stride=(VEC, 1)), cutlass.BFloat16
+ )
+ shg = cute.make_rmem_tensor(
+ cute.make_layout((g_n, VEC), stride=(VEC, 1)), cutlass.BFloat16
+ )
+ if const_expr(split):
+ ig = ig0
+ else:
+ ig = cute.make_rmem_tensor(cute.make_layout((g_n,), stride=(1,)), cutlass.Float32)
+ zg = cute.make_rmem_tensor(cute.make_layout((g_n,), stride=(1,)), cutlass.Float32)
+ if const_expr(topk > 0 and not split):
+ fg = cute.make_rmem_tensor(
+ cute.make_layout((g_n, topk, VEC), stride=(topk * VEC, VEC, 1)),
+ cutlass.BFloat16,
+ )
+ for k in cutlass.range_constexpr(0 if split else g_n):
+ # Past the last item a thread reads the last item's operands (unused).
+ c = cutlass.min(tidx + (g0 + k) * THREADS, cutlass.Int32(items - 1))
+ r = c // nvec
+ p = c % nvec
+ cute.autovec_copy(hc_vec[r, (None, col0 + p)], xg[k, None])
+ if const_expr(self.inject):
+ if const_expr(topk > 0):
+ # The finalize's slot rows of the permuted FC2 output (a -1 slot reads
+ # row 0, unused).
+ for j in cutlass.range_constexpr(topk):
+ cute.autovec_copy(
+ b_vec[
+ cutlass.max(sTopR[r * topk + j], cutlass.Int32(0)),
+ (None, k0 // VEC + p),
+ ],
+ fg[k, j, None],
+ )
+ else:
+ cute.autovec_copy(b_vec[r, (None, k0 // VEC + p)], bg[k, None])
+ ig[k] = mInj[r, s].to(cutlass.Float32)
+ if const_expr(self.combine):
+ cute.autovec_copy(sh_vec[r, (None, k0 // VEC + p)], shg[k, None])
+ zg[k] = mGate[r].to(cutlass.Float32)
+ if const_expr(split and g0 == 0):
+ _mbar_wait0(bbar) # the K half's block rows have landed
+ for k in cutlass.range_constexpr(g_n):
+ c = tidx + (g0 + k) * THREADS
+ if c < items:
+ r = c // nvec
+ p = c % nvec
+ cute.autovec_copy(snw_vec[0, (None, p)], nwr)
+ ssum = cutlass.Float32(0.0)
+ if const_expr(self.inject):
+ g = ig[k]
+ if const_expr(split):
+ cute.autovec_copy(sblk_vec[r, (None, p)], bg[k, None])
+ if const_expr(self.combine and not split):
+ sg = 1.0 / (1.0 + cute.math.exp(-zg[k]))
+ for v in cutlass.range_constexpr(VEC):
+ if const_expr(topk > 0 and not split):
+ # The runner's finalize: fp32 FMAs in slot order from zero, a
+ # -1 slot skipped, one rounding.
+ acc = cutlass.Float32(0.0)
+ for j in cutlass.range_constexpr(topk):
+ if sTopR[r * topk + j] >= 0:
+ acc = cute.math.fma(
+ fg[k, j, v].to(cutlass.Float32),
+ sTopW[r * topk + j],
+ acc,
+ )
+ bv = acc.to(cutlass.BFloat16).to(cutlass.Float32)
+ else:
+ bv = bg[k, v].to(cutlass.Float32)
+ if const_expr(self.combine and not split):
+ bv = (
+ cute.math.fma(sg, shg[k, v].to(cutlass.Float32), bv)
+ .to(cutlass.BFloat16)
+ .to(cutlass.Float32)
+ )
+ outr[v] = cute.math.fma(bv, g, xg[k, v].to(cutlass.Float32)).to(
+ cutlass.BFloat16
+ )
+ if nc == 0:
+ cute.autovec_copy(outr, o_vec[r, (None, col0 + p)])
+ else:
+ for v in cutlass.range_constexpr(VEC):
+ outr[v] = xg[k, v]
+ for v in cutlass.range_constexpr(VEC):
+ xv = outr[v].to(cutlass.Float32)
+ ssum = xv * xv + ssum
+ yr[v] = (xv * (1.0 + nwr[v].to(cutlass.Float32))).to(cutlass.BFloat16)
+ cute.autovec_copy(yr, sa_vec[r, (None, p)])
+ sSs[r, p] = ssum
+ cute.arch.cp_async_wait_group(0)
+ cute.arch.sync_threads()
+
+ # 3. Warp w: rows x its 160-wide K slice x the 24 columns; the warps' partials summed in
+ # warp order; the K half's sums of squares per row.
+ # The warp's K slice: 320-byte offsets, so 16-byte aligned (ldmatrix rows).
+ a_w = cute.make_tensor(
+ cute.make_ptr(
+ cutlass.BFloat16,
+ sA.iterator.toint() + warp_i * (KW * 2),
+ cute.AddressSpace.smem,
+ assumed_align=16,
+ ),
+ cute.make_layout((M, KW), stride=(ld, 1)),
+ )
+ b_w = cute.make_tensor(
+ cute.make_ptr(
+ cutlass.BFloat16,
+ sW.iterator.toint() + warp_i * (KW * 2),
+ cute.AddressSpace.smem,
+ assumed_align=16,
+ ),
+ cute.make_layout((NC, KW), stride=(ld, 1)),
+ )
+ thr = mma.get_slice(lane)
+ tCsA = thr.partition_A(a_w)
+ tCrA = thr.make_fragment_A(tCsA)
+ tCsB = thr.partition_B(b_w)
+ tCrB = thr.make_fragment_B(tCsB)
+ # Fragments by ldmatrix (4 / 2 8x8 matrices an instruction).
+ ldsm_a = cute.make_tiled_copy_A(
+ cute.make_copy_atom(warp.LdMatrix8x8x16bOp(False, 4), cutlass.BFloat16), mma
+ )
+ ldsm_b = cute.make_tiled_copy_B(
+ cute.make_copy_atom(warp.LdMatrix8x8x16bOp(False, 2), cutlass.BFloat16), mma
+ )
+ cpa = ldsm_a.get_slice(lane)
+ cpb = ldsm_b.get_slice(lane)
+ cute.copy(ldsm_a, cpa.partition_S(a_w), cpa.retile(tCrA))
+ cute.copy(ldsm_b, cpb.partition_S(b_w), cpb.retile(tCrB))
+ acc = cute.make_rmem_tensor(thr.partition_shape_C((M, NC)), cutlass.Float32)
+ acc.fill(0.0)
+ cute.gemm(mma, acc, tCrA, tCrB, acc)
+ cute.arch.sync_threads() # every warp's fragments are loaded: sW may be reused
+ co = thr.partition_C(cute.make_identity_tensor((M, NC)))
+ for j in cutlass.range_constexpr(cute.size(acc)):
+ rr, nn = co[j]
+ sWp[warp_i, rr, nn] = acc[j]
+ # Sums of squares: 16 threads a row (10 vectors each), then a 16-lane butterfly.
+ segs: cutlass.Constexpr = THREADS // M
+ per_seg: cutlass.Constexpr = nvec // segs
+ rq = tidx // segs
+ sq = tidx % segs
+ t = cutlass.Float32(0.0)
+ if rq < rows:
+ for p in cutlass.range_constexpr(per_seg):
+ t = t + sSs[rq, sq * per_seg + p]
+ t = cute.arch.warp_reduction_sum(t, threads_in_group=segs)
+ if sq == 0:
+ if rq < rows:
+ sSsr[rq] = t
+ cute.arch.sync_threads()
+ for i in cutlass.range_constexpr(-(-(M * NC) // THREADS)):
+ e = tidx + i * THREADS
+ if e < M * NC:
+ rr = e // NC
+ nn = e % NC
+ t = cutlass.Float32(0.0)
+ for w in cutlass.range_constexpr(WARPS):
+ t = t + sWp[w, rr, nn]
+ sAcc[rr, nn] = t
+ cute.arch.sync_threads()
+
+ # 4. Push rows (2o, 2o+1) (and their sums of squares) to cluster rank o, slot `me`.
+ if const_expr(not split):
+ cute.arch.cluster_wait()
+ red_addr = cutlass.Int32(sRed.iterator.toint())
+ per: cutlass.Constexpr = RPO * (NC + 1)
+ for it in cutlass.range_constexpr(-(-(CL * per) // THREADS)):
+ c = tidx + it * THREADS
+ if c < CL * per:
+ o = c // per
+ rem = c % per
+ rl = rem // (NC + 1)
+ nn = rem % (NC + 1)
+ rr = o * RPO + rl
+ if rr < rows:
+ v = cutlass.Float32(0.0)
+ if nn < NC:
+ v = sAcc[rr, nn]
+ else:
+ v = sSsr[rr]
+ dst = red_addr + ((me * RPO + rl) * rld + nn) * 4
+ _st_async_f32(_mapa(dst, o), v, _mapa(rbar, o))
+ _mbar_wait0(rbar)
+
+ # 5. Rank `me` finishes rows (2 me, 2 me + 1): rstd per stream, the logits, low / injection.
+ for it in cutlass.range_constexpr(-(-(RPO * NC) // THREADS)):
+ c = tidx + it * THREADS
+ if c < RPO * NC:
+ rl = c // NC
+ nn = c % NC
+ rr = me * RPO + rl
+ col = n0 + nn
+ if rr < rows:
+ if col < n_cols:
+ z = cutlass.Float32(0.0)
+ for st in cutlass.range_constexpr(4):
+ ssv = sRed[KH * st, rl, NC] + sRed[KH * st + 1, rl, NC]
+ rs = cute.math.rsqrt(ssv / hidden + eps)
+ z = (sRed[KH * st, rl, nn] + sRed[KH * st + 1, rl, nn]) * rs + z
+ z = z * 0.25
+ sig = 1.0 / (1.0 + cute.math.exp(-z))
+ if col < rank:
+ mLow[rr, col] = (z * sig).to(cutlass.BFloat16)
+ else:
+ mInjOut[rr, col - rank] = (2.0 * sig).to(cutlass.BFloat16)
+ if nc == 0:
+ if tidx < RPO * 4:
+ rl = tidx // 4
+ st = tidx % 4
+ rr = me * RPO + rl
+ if rr < rows:
+ ssv = sRed[KH * st, rl, NC] + sRed[KH * st + 1, rl, NC]
+ mRstd[rr, st] = cute.math.rsqrt(ssv / hidden + eps)
+
+
+def _dl(t: torch.Tensor, align: int = 16):
+ return from_dlpack(t, assumed_align=align)
+
+
+@functools.lru_cache(maxsize=None)
+def _compiled_first(
+ hidden,
+ hc_count,
+ rank,
+ n_cols,
+ n_wrows,
+ rows,
+ inject,
+ combine,
+ pdl,
+ hc_stride,
+ block_stride,
+ topk=0,
+ n_perm=0,
+ gate_stride=1,
+):
+ kern = _First2(
+ hidden=hidden,
+ rank=rank,
+ n_cols=n_cols,
+ rows=rows,
+ inject=inject,
+ combine=combine,
+ pdl=pdl,
+ topk=topk,
+ )
+ dev, bf = "cuda", torch.bfloat16
+ width = hidden * hc_count
+ hc = torch.empty(rows, hc_stride, dtype=bf, device=dev)[:, :width]
+ block = torch.empty(n_perm if topk else rows, block_stride, dtype=bf, device=dev)[:, :hidden]
+ topw = torch.empty(rows, max(topk, 1), dtype=bf, device=dev)
+ toprows = torch.empty(rows, max(topk, 1), dtype=torch.int32, device=dev)
+ inj = torch.empty(rows, hc_count, dtype=bf, device=dev)
+ shared = torch.empty(rows, hidden, dtype=bf, device=dev)
+ # The kernel reads the gate logits through this layout: a compile per row stride.
+ gate = torch.empty(rows, gate_stride, dtype=bf, device=dev)[:, 0]
+ norm = torch.empty(1, width, dtype=bf, device=dev)
+ w = torch.empty(n_wrows, width, dtype=bf, device=dev)
+ hc_out = torch.empty(rows, width, dtype=bf, device=dev)
+ low = torch.empty(rows, rank, dtype=bf, device=dev)
+ inj_out = torch.empty(rows, hc_count, dtype=bf, device=dev)
+ rstd = torch.empty(rows, 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(low),
+ _dl(inj_out, 2),
+ _dl(rstd, 4),
+ cutlass.Float32(1e-6),
+ stream,
+ )
+
+
+def first_rows(
+ hyper_input,
+ block_out,
+ inject,
+ norm_weight,
+ first,
+ hc_count,
+ n_cols,
+ eps,
+ shared=None,
+ gate_logits=None,
+ topk_weights=None,
+ topk_rows=None,
+ rank=320,
+):
+ """Returns `(hc_out, low [R, rank] bf16, inject_out [R, C] bf16 or None, rstd [R, C] fp32)`;
+ `n_cols` = rank + C with the injection, rank without (the stack-final mixer). With
+ `topk_weights` / `topk_rows` (the runner left unfinalized), `block_out` is its permuted FC2
+ output `[P, H]` and the routed rows are its finalize, formed here."""
+ 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
+ hc_out = (
+ torch.empty((rows, width), dtype=hyper_input.dtype, device=dev)
+ if has_inject
+ else hyper_input
+ )
+ low = torch.empty((rows, rank), dtype=hyper_input.dtype, device=dev)
+ inj_out = torch.empty((rows, hc_count), dtype=hyper_input.dtype, device=dev)
+ rstd = torch.empty((rows, hc_count), dtype=torch.float32, device=dev)
+ pdl = get_env_enable_pdl()
+ weight = first.t()
+ assert weight.shape[0] >= NC * NCH
+ top_k = topk_rows.shape[1] if topk_rows is not None else 0
+ # One logit a row at any row stride: a column of the shared expert's fused gate|up GEMM
+ # output is read in place.
+ gate = gate_logits.reshape(rows) if has_shared else hyper_input.new_empty(rows)
+ fn = _compiled_first(
+ hidden,
+ hc_count,
+ rank,
+ n_cols,
+ weight.shape[0],
+ rows,
+ has_inject,
+ has_shared,
+ bool(pdl),
+ hyper_input.stride(0),
+ block_out.stride(0) if has_inject else hidden,
+ top_k,
+ block_out.shape[0] if top_k else 0,
+ gate.stride(0),
+ )
+ blk = block_out if has_inject else hyper_input.new_empty(rows, hidden)
+ inj = inject if has_inject else hyper_input.new_empty(rows, hc_count)
+ sh = shared if has_shared else hyper_input.new_empty(rows, hidden)
+ topw = topk_weights if top_k else hyper_input.new_empty(rows, 1)
+ toprows = topk_rows if top_k else torch.empty(rows, 1, dtype=torch.int32, device=dev)
+ stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream)
+ 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(low),
+ _dl(inj_out, 2),
+ _dl(rstd, 4),
+ cutlass.Float32(eps),
+ stream,
+ )
+ return hc_out, low, (inj_out if n_cols > rank else None), rstd
+
+
+class _Mix2:
+ def __init__(self, *, hidden, streams, rank, rows, fp4, pdl):
+ assert streams * BLOCK_H == WARPS * 8 and rows <= M and rank % 16 == 0
+ self.hidden, self.streams, self.rank, self.rows, self.fp4, self.pdl = (
+ hidden,
+ streams,
+ rank,
+ rows,
+ fp4,
+ pdl,
+ )
+ self.ld = rank + PAD
+
+ def smem_bytes(self):
+ def r(n):
+ return -(-n // 128) * 128
+
+ nrow = self.streams * BLOCK_H
+ b = r(nrow * self.ld * 2) + r(M * self.ld * 2)
+ b += r(M * nrow * 4) + r(M * nrow * 2) + r(M * self.streams * 4) + r(nrow * 4)
+ return b + 256
+
+ @cute.jit
+ def __call__(
+ self, mLow, mRstd, mUp, mHc, mNorm, mGs, mMixed, mXq, mXsf, stream: cuda_drv.CUstream
+ ):
+ 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)
+ self.kernel(mLow, mRstd, mUp, mHc, mNorm, mGs, mMixed, mXq, mXsf, mma, g2s).launch(
+ grid=[self.hidden // BLOCK_H, 1, 1],
+ block=[THREADS, 1, 1],
+ smem=self.smem_bytes(),
+ stream=stream,
+ use_pdl=self.pdl,
+ min_blocks_per_mp=2,
+ )
+
+ @cute.kernel
+ def kernel(
+ self,
+ mLow,
+ mRstd,
+ mUp,
+ mHc,
+ mNorm,
+ mGs,
+ mMixed,
+ mXq,
+ mXsf,
+ mma: cute.TiledMma,
+ g2s: cute.CopyAtom,
+ ):
+ tidx, _, _ = cute.arch.thread_idx()
+ bidx, _, _ = cute.arch.block_idx()
+ warp_i = cute.arch.warp_idx()
+ lane = cute.arch.lane_idx()
+ hidden: cutlass.Constexpr = self.hidden
+ streams: cutlass.Constexpr = self.streams
+ rank: cutlass.Constexpr = self.rank
+ rows: cutlass.Constexpr = self.rows
+ ld: cutlass.Constexpr = self.ld
+ nrow: cutlass.Constexpr = streams * BLOCK_H
+ i0 = bidx * BLOCK_H
+
+ smem = cutlass.utils.SmemAllocator()
+ sU = smem.allocate_tensor(
+ cutlass.BFloat16, cute.make_layout((nrow, rank), stride=(ld, 1)), byte_alignment=128
+ )
+ sL = smem.allocate_tensor(
+ cutlass.BFloat16, cute.make_layout((M, rank), stride=(ld, 1)), byte_alignment=128
+ )
+ sV = smem.allocate_tensor(
+ cutlass.Float32, cute.make_layout((M, nrow), stride=(nrow, 1)), byte_alignment=128
+ )
+ sX = smem.allocate_tensor(
+ cutlass.BFloat16, cute.make_layout((M, nrow), stride=(nrow, 1)), byte_alignment=128
+ )
+ sR = smem.allocate_tensor(
+ cutlass.Float32, cute.make_layout((M, streams), stride=(streams, 1)), byte_alignment=128
+ )
+ sG = smem.allocate_tensor(
+ cutlass.Float32, cute.make_layout((nrow,), stride=(1,)), byte_alignment=128
+ )
+
+ u_vec = cute.logical_divide(mUp, (None, VEC))
+ su_vec = cute.logical_divide(sU, (None, VEC))
+ pieces: cutlass.Constexpr = nrow * (rank // VEC)
+ for i in cutlass.range_constexpr(-(-pieces // THREADS)):
+ c = tidx + i * THREADS
+ if c < pieces:
+ n = c // (rank // VEC)
+ p = c % (rank // VEC)
+ jrow = (n // BLOCK_H) * hidden + i0 + n % BLOCK_H
+ cute.copy(g2s, u_vec[jrow, (None, p)], su_vec[n, (None, p)])
+ cute.arch.cp_async_commit_group()
+ if tidx < nrow:
+ sG[tidx] = 1.0 + mNorm[0, (tidx // BLOCK_H) * hidden + i0 + tidx % BLOCK_H].to(
+ cutlass.Float32
+ )
+ # Rows of sL past R are never written (MMA rows are independent; discarded).
+ gs = cutlass.Float32(0.0)
+ if const_expr(self.fp4):
+ gs = mGs[0]
+ # Before the wait as well: the up rows land, and every thread's B fragments go to
+ # registers (the post-wait MMA then reads only its A fragments from shared memory).
+ if const_expr(self.pdl):
+ cute.arch.griddepcontrol_launch_dependents()
+ cute.arch.cp_async_wait_group(0)
+ cute.arch.sync_threads()
+ thr = mma.get_slice(tidx)
+ tCsB = thr.partition_B(sU)
+ tCrB = thr.make_fragment_B(tCsB)
+ cute.autovec_copy(tCsB, tCrB)
+ if const_expr(self.pdl):
+ cute.arch.griddepcontrol_wait()
+
+ # low rows (16-byte pieces), rstd, this CTA's residual elements.
+ l_vec = cute.logical_divide(mLow, (None, VEC))
+ sl_vec = cute.logical_divide(sL, (None, VEC))
+ lp: cutlass.Constexpr = rows * (rank // VEC)
+ for i in cutlass.range_constexpr(-(-lp // THREADS)):
+ c = tidx + i * THREADS
+ if c < lp:
+ rr = c // (rank // VEC)
+ p = c % (rank // VEC)
+ cute.copy(g2s, l_vec[rr, (None, p)], sl_vec[rr, (None, p)])
+ cute.arch.cp_async_commit_group()
+ if tidx < rows * streams:
+ sR[tidx // streams, tidx % streams] = mRstd[tidx // streams, tidx % streams]
+ h_vec = cute.logical_divide(mHc, (None, VEC))
+ sx_vec = cute.logical_divide(sX, (None, VEC))
+ xp: cutlass.Constexpr = rows * (nrow // VEC)
+ if tidx < xp:
+ rr = tidx // (nrow // VEC)
+ q = tidx % (nrow // VEC)
+ cute.copy(
+ g2s,
+ h_vec[
+ rr,
+ (
+ None,
+ ((q * VEC) // BLOCK_H) * (hidden // VEC)
+ + (i0 + (q * VEC) % BLOCK_H) // VEC,
+ ),
+ ],
+ sx_vec[rr, (None, q)],
+ )
+ cute.arch.cp_async_commit_group()
+ cute.arch.cp_async_wait_group(0)
+ cute.arch.sync_threads()
+
+ tCsA = thr.partition_A(sL)
+ tCrA = thr.make_fragment_A(tCsA)
+ ldsm_a = cute.make_tiled_copy_A(
+ cute.make_copy_atom(warp.LdMatrix8x8x16bOp(False, 4), cutlass.BFloat16), mma
+ )
+ cpa = ldsm_a.get_slice(tidx)
+ cute.copy(ldsm_a, cpa.partition_S(sL), cpa.retile(tCrA))
+ acc = cute.make_rmem_tensor(thr.partition_shape_C((M, nrow)), cutlass.Float32)
+ acc.fill(0.0)
+ cute.gemm(mma, acc, tCrA, tCrB, acc)
+ co = thr.partition_C(cute.make_identity_tensor((M, nrow)))
+ for j in cutlass.range_constexpr(cute.size(acc)):
+ rr, n = co[j]
+ if rr < rows:
+ gate = 1.0 / (1.0 + cute.math.exp(-acc[j]))
+ x = sX[rr, n].to(cutlass.Float32)
+ sV[rr, n] = gate * (x * sR[rr, n // BLOCK_H] * sG[n])
+ cute.arch.sync_threads()
+
+ for rb in cutlass.range_constexpr(0, rows, WARPS):
+ r = rb + warp_i
+ if r < rows:
+ m = cutlass.Float32(0.0)
+ if lane < BLOCK_H:
+ tot = cutlass.Float32(0.0)
+ for t in cutlass.range_constexpr(streams):
+ tot = tot + sV[r, t * BLOCK_H + lane]
+ mb = (tot * (1.0 / streams)).to(cutlass.BFloat16)
+ mMixed[r, i0 + lane] = mb
+ m = mb.to(cutlass.Float32)
+ if const_expr(self.fp4):
+ 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:
+ code = code + 8
+ hi = cute.arch.shuffle_sync_down(code, offset=1)
+ if lane < BLOCK_H:
+ if lane % 2 == 0:
+ mXq[r, (i0 + lane) // 2] = (code + hi * 16).to(cutlass.Uint8)
+ if lane == 0:
+ mXsf[r * (hidden // BLOCK_H) + bidx] = sf8
+
+
+@functools.lru_cache(maxsize=None)
+def _compiled_mix(hidden, streams, rank, rows, hc_stride, fp4, pdl):
+ kern = _Mix2(hidden=hidden, streams=streams, rank=rank, rows=rows, fp4=fp4, pdl=pdl)
+ dev, bf = "cuda", torch.bfloat16
+ width = hidden * streams
+ low = torch.empty(rows, rank, dtype=bf, device=dev)
+ rstd = torch.empty(rows, streams, dtype=torch.float32, device=dev)
+ up = torch.empty(width, rank, dtype=bf, device=dev)
+ hc = torch.empty(rows, 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(rows, hidden, dtype=bf, device=dev)
+ xq = torch.empty(rows, hidden // 2, dtype=torch.uint8, device=dev)
+ xsf = torch.empty(rows * (hidden // 16), dtype=torch.float8_e4m3fn, device=dev)
+ stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream)
+ return cute_ext.compile(
+ kern,
+ _dl(low),
+ _dl(rstd, 4),
+ _dl(up),
+ _dl(hc, 16),
+ _dl(norm, 2),
+ _dl(gs, 4),
+ _dl(mixed, 2),
+ _dl(xq, 1),
+ _dl(xsf, 1),
+ stream,
+ )
+
+
+def mix_rows(low, rstd, up, hc_out, norm_weight, hc_count, global_scale=None):
+ """Returns `(mixed, xq, xsf)` (xq / xsf None without `global_scale`)."""
+ 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)
+ 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
+ pdl = get_env_enable_pdl()
+ fn = _compiled_mix(hidden, hc_count, rank, rows, hc_out.stride(0), 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)
+ q = xq if fp4 else mixed.new_empty(rows, hidden // 2, dtype=torch.uint8)
+ sf = (
+ xsf.view(torch.float8_e4m3fn)
+ if fp4
+ else mixed.new_empty(rows * (hidden // 16), dtype=torch.float8_e4m3fn)
+ )
+ fn(
+ _dl(low),
+ _dl(rstd, 4),
+ _dl(up.t()),
+ _dl(hc_out, 16),
+ _dl(norm_weight.reshape(1, -1), 2),
+ _dl(gs, 4),
+ _dl(mixed, 2),
+ _dl(q, 1),
+ _dl(sf, 1),
+ stream,
+ )
+ return mixed, xq, xsf
+
+
+def _reference_first(
+ hyper_input,
+ block_out,
+ inject,
+ norm_weight,
+ first,
+ hc_count,
+ n_cols,
+ eps,
+ shared=None,
+ gate_logits=None,
+ topk_weights=None,
+ topk_rows=None,
+):
+ """float64: the finalize (with `topk_rows`), the combine and the injection each rounded once
+ as the 1-row kernel does, the grouped RMS, the first projection over `normed`, `low =
+ silu(z)` and `2 sigmoid(z)` (z = logits / C) each rounded once, rstd in fp32."""
+ dt = hyper_input.dtype
+ rows, width = hyper_input.shape
+ hidden = width // hc_count
+ rank = RANK
+ if block_out is None:
+ hc_out = hyper_input
+ else:
+ if topk_rows is not None:
+ valid = (topk_rows >= 0).unsqueeze(-1)
+ gathered = block_out.double()[topk_rows.clamp(min=0).long()]
+ w = topk_weights.double().unsqueeze(-1)
+ block_out = torch.where(valid, gathered * w, 0.0).sum(1).to(block_out.dtype)
+ blk = block_out.double()
+ if shared is not None:
+ sg = torch.sigmoid(gate_logits.double().reshape(rows, 1))
+ blk = (blk + sg * shared.double()).to(dt).double()
+ streams = hyper_input.double().unflatten(-1, (hc_count, hidden))
+ hc_out = (streams + blk.unsqueeze(-2) * inject.double().unsqueeze(-1)).flatten(-2).to(dt)
+ x = hc_out.double().unflatten(-1, (hc_count, hidden))
+ rstd = torch.rsqrt(x.pow(2).mean(-1) + eps)
+ normed = (x * rstd.unsqueeze(-1)).flatten(-2) * (1.0 + norm_weight.double())
+ z = normed @ first[:, :n_cols].double() / hc_count
+ low = torch.nn.functional.silu(z[:, :rank]).to(dt)
+ inj = (2.0 * torch.sigmoid(z[:, rank:])).to(dt) if n_cols > rank else None
+ return hc_out, low, inj, rstd.float()
+
+
+class Qwen4ExpHcDecodeRowsFirst(OpWrapper):
+ """First launch of a 2..16-row decode hyper-connection site: the previous block's injection
+ (and, with `shared` / `gate_logits`, the MoE block's combine `routed + sigmoid(gate) * shared`
+ rounded once), the grouped RMS, the first projection over `[down; block_inject]`, the SiLU
+ bottleneck and the injection weights. Returns `(hc_out, low [R, 320] bf16, inject_out [R, C]
+ or None, rstd [R, C] fp32)`; `hc_out` *is* `hyper_input` without `block_out`.
+
+ `hc_out` is bitwise the 1-row kernels'; `low` and the injection weights come from bf16
+ operands on tensor cores (`y = bf16(x * (1 + w))` against the bf16 weight) with fp32
+ accumulation, inside 1-2 bf16 ulp of the float64 composition.
+
+ 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 -1 slot skipped), rounded once -- bitwise what the runner's finalize
+ stores (as `qwen4_exp_hc_decode_first_cute`'s fold). The fold needs the combine.
+
+ `gate_logits` is one logit a row, `[R, 1]` at any row stride: the target passes the column
+ of the shared expert's fused gate|up GEMM output, read in place (a compile per stride).
+
+ Fusion boundary: the routed experts' finalize (optional), the combine, injection, grouped
+ RMS statistics, first projection and activations.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="site after a MoE block: combine + injection, 2 rows",
+ spec=dict(rows=2, site="combine"),
+ ),
+ Cell(
+ why="site after a MoE block, runner finalize folded in, 4 rows (MTP3 verify)",
+ spec=dict(rows=4, site="fold"),
+ ),
+ Cell(why="MoE site: injection, 8 rows", spec=dict(rows=8, site="inject")),
+ Cell(
+ why="norm-only site (stack start, MTP head), 16 rows, 320 cols",
+ spec=dict(rows=16, site="norm"),
+ ),
+ )
+ note = """
+ Metadata consumed: none. Call sites: `DecodeTarget._hc_site` / `_hc_moe_site` at 2..16 rows
+ (decode batches, MTP verify and draft-extend rows), followed by `qwen4_exp_hc_decode_rows_mix`;
+ after a TRTLLM-backend MoE block (`DecodeTarget._moe_block` on the permuted runner) with the
+ fold. Grid (2 K halves, 4 streams, 14 N chunks of 24 columns), cluster (2, 4, 1), 8 warps;
+ weights staged before the PDL wait. With the fold from 4 rows each of a K half's 4 CTAs
+ forms a quarter of the half's routed + shared rows and pushes it to all four (DSMEM); at 2
+ and 3 rows each CTA folds its K half in its items. One compile per row count and variant.
+ """
+
+ def __call__(
+ self,
+ hyper_input,
+ block_out,
+ inject,
+ norm_weight,
+ first,
+ hc_count,
+ n_cols,
+ eps,
+ shared=None,
+ gate_logits=None,
+ topk_weights=None,
+ topk_rows=None,
+ ):
+ return first_rows(
+ hyper_input,
+ block_out,
+ inject,
+ norm_weight,
+ first,
+ hc_count,
+ n_cols,
+ eps,
+ shared,
+ gate_logits,
+ topk_weights,
+ topk_rows,
+ )
+
+ def reference(
+ self,
+ hyper_input,
+ block_out,
+ inject,
+ norm_weight,
+ first,
+ hc_count,
+ n_cols,
+ eps,
+ shared=None,
+ gate_logits=None,
+ topk_weights=None,
+ topk_rows=None,
+ ):
+ return _reference_first(
+ hyper_input,
+ block_out,
+ inject,
+ norm_weight,
+ first,
+ hc_count,
+ n_cols,
+ eps,
+ shared,
+ gate_logits,
+ topk_weights,
+ topk_rows,
+ )
+
+ def is_valid(
+ self,
+ hyper_input,
+ block_out,
+ inject,
+ norm_weight,
+ first,
+ hc_count,
+ n_cols,
+ eps,
+ shared=None,
+ gate_logits=None,
+ topk_weights=None,
+ topk_rows=None,
+ ) -> None:
+ rows, width = hyper_input.shape
+ assert 1 <= rows <= MAX_ROWS, f"1..{MAX_ROWS} rows"
+ assert hyper_input.dtype == torch.bfloat16 and hyper_input.stride(1) == 1
+ assert width // hc_count == KH * KHW, "hidden 2560"
+ assert first.stride(0) == 1 and first.shape[1] >= NC * NCH, (
+ "the padded [width, 384] .t() view"
+ )
+ assert (block_out is None) == (inject is None)
+ assert (shared is None) == (gate_logits is None)
+ if gate_logits is not None:
+ assert gate_logits.dtype == torch.bfloat16 and gate_logits.numel() == rows, (
+ "one bf16 gate logit a row"
+ )
+ assert shared.is_contiguous() and shared.shape == (rows, width // hc_count)
+ assert (topk_weights is None) == (topk_rows is None), "topk_weights and topk_rows together"
+ if topk_rows is not None:
+ assert shared is not None, "the finalize fold needs the combine"
+ assert topk_rows.dtype == torch.int32 and topk_rows.shape[0] == rows
+ assert topk_rows.is_contiguous() and topk_weights.is_contiguous()
+ assert topk_weights.shape == topk_rows.shape and topk_weights.dtype == torch.bfloat16
+ assert block_out.dim() == 2 and block_out.shape[1] == width // hc_count
+ assert block_out.stride(1) == 1, "permuted rows [P, hidden]"
+
+
+class Qwen4ExpHcDecodeRowsMix(OpWrapper):
+ """Second launch of a 2..16-row decode hyper-connection site: the up projection of `low`, the
+ sigmoid gate on the normed stream (`hc_out * rstd * (1 + w)`), the stream mean, and with
+ `global_scale` the NVFP4 of the mean for the routed experts. Returns `(mixed, xq, xsf)`
+ (`xq` / `xsf` None without `global_scale`); one rounding per output.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="MoE site with NVFP4 output, 8 rows", spec=dict(rows=8, site="inject")),
+ Cell(why="plain mix, 2 rows", spec=dict(rows=2, site="combine")),
+ Cell(why="plain mix, 16 rows", spec=dict(rows=16, site="norm")),
+ )
+ note = """
+ Metadata consumed: none. Call sites: after `qwen4_exp_hc_decode_rows_first`. 160 CTAs of 8
+ warps (16 hidden columns x 4 streams); the up weight's 64 rows staged before the PDL wait.
+ """
+
+ def __call__(self, low, rstd, up, hc_out, norm_weight, hc_count, global_scale=None):
+ return mix_rows(low, rstd, up, hc_out, norm_weight, hc_count, global_scale)
+
+ def reference(self, low, rstd, up, hc_out, norm_weight, hc_count, global_scale=None):
+ """float64 mixed, rounded once (the NVFP4 is checked against the kernel's own mixed)."""
+ rows, width = hc_out.shape
+ hidden = width // hc_count
+ gate = torch.sigmoid(low.double() @ up.double())
+ x = hc_out.double().unflatten(-1, (hc_count, hidden)) * rstd.double().unsqueeze(-1)
+ normed = x.flatten(-2) * (1.0 + norm_weight.double())
+ mixed = (gate * normed).unflatten(-1, (hc_count, hidden)).mean(-2).to(hc_out.dtype)
+ return mixed, None, None
+
+ def is_valid(self, low, rstd, up, hc_out, norm_weight, hc_count, global_scale=None) -> None:
+ rows = hc_out.shape[0]
+ assert 1 <= rows <= MAX_ROWS and low.shape[0] == rows and rstd.shape == (rows, hc_count)
+ assert up.t().is_contiguous(), "up the [R, width] view of a row-major [width, R] weight"
+ assert global_scale is None or global_scale.numel() == 1
+
+
+qwen4_exp_hc_decode_rows_first = Qwen4ExpHcDecodeRowsFirst()
+qwen4_exp_hc_decode_rows_mix = Qwen4ExpHcDecodeRowsMix()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_first_rescale_tc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_first_rescale_tc.py
new file mode 100644
index 000000000000..2e21b5c45ddc
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_first_rescale_tc.py
@@ -0,0 +1,976 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Hyper-connection first projection on the un-normalized residual, one tcgen05
+launch: the A operand is the site's normed stream formed on chip from `hc_out`
+and its per-(row, stream) rstd, so `[T, C*H]` normed never exists in memory.
+
+Tier 6 (CuTe DSL written for this target). A 2-CTA tcgen05 GEMM (M = 256 rows a
+CTA pair, K = 64 a stage, N = 352 as two MMAs of 192 + 160 columns). Per CTA: a
+TMA warp lands the raw 128 x 64 `hc_out` tile (128B-swizzled) in the operand
+ring; eight transform warps scale it in place by `rstd * (1 + w)`, round to bf16
+and commit the stage to the pair's leader; a TMA warp streams both halves of the
+first weight (2-CTA multicast); the leader's MMA warp accumulates in tensor
+memory. The same eight warps then run the epilogue (TMEM -> registers -> SiLU /
+2 sigmoid -> HBM): `qwen4_exp_hc_lowrank_act`'s activations without the fp32
+logits round trip.
+
+With few row tiles (`_pick_split`) each CTA pair covers `1 / split` of K and
+writes fp32 partial logits `[split, T, 352]`; a Triton combine sums them in split
+order and applies the activations.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+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
+import triton
+import triton.language as tl
+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._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_inject_norm import (
+ _bf16_hi,
+ _bf16_lo,
+ _bf16x2,
+ _ldg_v4,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_up_gate_mean import (
+ _asm,
+ _st_global_v4,
+)
+from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+from tensorrt_llm._torch.kernels.blackwell.custom_pipeline import PipelineTmaUmma, PipelineUmmaAsync
+from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
+
+#: Programmatic dependent launch (sm_90+), off with TRTLLM_ENABLE_PDL=0.
+_PDL = get_env_enable_pdl()
+
+#: K a stage (one 128B-swizzled bf16 row of the operand tile).
+BK = 64
+#: Columns of the two MMAs a K step: 192 + 160 = 352 >= rank + C, the padded first
+#: weight's columns the cuBLAS path multiplies (`hc_first_cols`). The 2-CTA atom needs
+#: N / 2 a multiple of 16 per CTA, so two 176-column halves do not build.
+NH0 = 192
+NH1 = 160
+NCOL = NH0 + NH1
+#: Operand (A) ring depth: A is the DRAM stream (`hc_out`), B (the weight) comes from L2.
+_A_STAGES = 6
+_B_STAGES = 5
+#: Transform warp groups: group g transforms k blocks j = g (mod groups), so that many
+#: k blocks are in flight in the transform at once (its per-k-block latency is what
+#: bounds a single group).
+_GROUPS = 2
+#: Largest K split (`_pick_split`).
+_MAX_SPLIT = 8
+
+
+@dsl_user_op
+def _lds_v4(addr, *, loc=None, ip=None):
+ """16 bytes of shared memory. Side-effecting, so never moved across the barrier
+ wait that makes them valid.
+ """
+ ty = T.i32()
+ res = _asm(
+ llvm.StructType.get_literal([ty] * 4),
+ [cutlass.Int32(addr)],
+ "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];",
+ "=r,=r,=r,=r,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+ return [cutlass.Int32(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(4)]
+
+
+@dsl_user_op
+def _sts_v4(addr, a, b, c, d, *, loc=None, ip=None):
+ """16 bytes to shared memory."""
+ _asm(
+ T.i32(),
+ [cutlass.Int32(addr)] + [cutlass.Int32(v) for v in (a, b, c, d)],
+ "{ st.shared.v4.b32 [$1], {$2, $3, $4, $5}; mov.u32 $0, 0; }",
+ "=r,r,r,r,r,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+
+
+@dsl_user_op
+def _ldg_f32(addr, *, loc=None, ip=None):
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [cutlass.Int64(addr)],
+ "ld.global.f32 $0, [$1];",
+ "=f,l",
+ loc,
+ ip,
+ side_effects=True,
+ )
+ )
+
+
+@dsl_user_op
+def _bits(x, *, loc=None, ip=None):
+ return cutlass.Int32(_asm(T.i32(), [cutlass.Float32(x)], "mov.b32 $0, $1;", "=r,f", loc, ip))
+
+
+@dsl_user_op
+def _mul(a, b, *, loc=None, ip=None):
+ """`mul.rn.f32`: never contracted into an FMA, so the rounding is the norm's."""
+ 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 _add1(a, *, loc=None, ip=None):
+ """`1 + a` in fp32 (`add.rn`): the grouped norm's `(1 + w)`."""
+ return cutlass.Float32(
+ _asm(T.f32(), [cutlass.Float32(a)], "add.rn.f32 $0, $1, 0f3F800000;", "=f,f", loc, ip)
+ )
+
+
+@dsl_user_op
+def _sigmoid_fast(x, *, loc=None, ip=None):
+ """1 / (1 + 2^(-x log2 e)) with ex2.approx and rcp.approx (two MUFU ops, no
+ division).
+ """
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [cutlass.Float32(x)],
+ "{ .reg .f32 t; mul.f32 t, $1, 0fBFB8AA3B; ex2.approx.ftz.f32 t, t; "
+ "add.f32 t, t, 0f3F800000; rcp.approx.ftz.f32 $0, t; }",
+ "=f,f",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _st_global_v2(addr, a, b, *, loc=None, ip=None):
+ """8 bytes to global memory (a row's four bf16 injection weights)."""
+ _asm(
+ 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,
+ )
+
+
+class HcFirstRescaleKernel:
+ """`low = silu(A @ first[:rank]^T / C)`, `inject = 2 sigmoid(A @ first[rank:rank+C]^T / C)`
+ with `A = bf16(hc_out * rstd * (1 + w))` formed in the operand ring (`split == 1`), or the
+ fp32 partial logits of K slice `sp` (`split > 1`).
+ """
+
+ def __init__(self, hc: int, hidden: int, rank: int, split: int):
+ assert (hc * hidden) % BK == 0 and hidden % BK == 0 and (hc * hidden // BK) % split == 0
+ assert rank + hc <= NCOL and rank % 16 == 0 and hc == 4
+ self.hc, self.hidden, self.rank, self.split = hc, hidden, rank, split
+ self.kt = hc * hidden // BK // split # k blocks a CTA pair
+ self.num_stages = _A_STAGES
+ self.num_b_stages = _B_STAGES
+ self.groups = _GROUPS
+ self.acc_dtype = cutlass.Float32
+ self.io_dtype = cutlass.BFloat16
+ self.cta_group = tcgen05.CtaGroup.TWO
+ self.cluster_shape_mn = (2, 1)
+ self.mma_tiler = (256, NH0, BK)
+ self.mma_tiler1 = (256, NH1, BK)
+ self.tw = 8 # transform / epilogue warps
+ self.mma_warp_id = 8
+ self.tma_warp_id = 9
+ self.tmb_warp_id = 10
+ self.threads_per_cta = 32 * 11
+ self.epi_bar_id = 1
+ self.tmem_bar_id = 2
+
+ def _setup(self):
+ k_major = utils.LayoutEnum.ROW_MAJOR.mma_major_mode()
+ tiled_mma = utils.sm100.make_trivial_tiled_mma(
+ self.io_dtype,
+ self.io_dtype,
+ k_major,
+ k_major,
+ self.acc_dtype,
+ self.cta_group,
+ self.mma_tiler[:2],
+ )
+ tiled_mma1 = utils.sm100.make_trivial_tiled_mma(
+ self.io_dtype,
+ self.io_dtype,
+ k_major,
+ k_major,
+ self.acc_dtype,
+ self.cta_group,
+ self.mma_tiler1[:2],
+ )
+ self.cluster_layout_vmnk = cute.tiled_divide(
+ cute.make_layout((*self.cluster_shape_mn, 1)), (tiled_mma.thr_id.shape,)
+ )
+ self.a_smem_staged = utils.sm100.make_smem_layout_a(
+ tiled_mma, self.mma_tiler, self.io_dtype, self.num_stages
+ )
+ self.b_smem_staged = utils.sm100.make_smem_layout_b(
+ tiled_mma, self.mma_tiler, self.io_dtype, self.num_b_stages
+ )
+ self.b1_smem_staged = utils.sm100.make_smem_layout_b(
+ tiled_mma1, self.mma_tiler1, self.io_dtype, self.num_b_stages
+ )
+ self.num_tmem_cols = 512
+ return tiled_mma, tiled_mma1
+
+ @cute.jit
+ def __call__(
+ self,
+ m: cutlass.Int32,
+ hc_ptr: cute.Pointer,
+ w_ptr: cute.Pointer,
+ rstd: cute.Tensor,
+ b_ptr: cute.Pointer,
+ out_ptr: cute.Pointer,
+ inj_ptr: cute.Pointer,
+ stream: cuda.CUstream,
+ ):
+ width = self.hc * self.hidden
+ tiled_mma, tiled_mma1 = self._setup()
+ a = cute.make_tensor(hc_ptr, cute.make_layout((m, width, 1), stride=(width, 1, 0)))
+ b = cute.make_tensor(b_ptr, cute.make_layout((NH0, width, 1), stride=(width, 1, 0)))
+ b1 = cute.make_tensor(
+ b_ptr + NH0 * width, cute.make_layout((NH1, width, 1), stride=(width, 1, 0))
+ )
+ a_one = cute.slice_(self.a_smem_staged, (None, None, None, 0))
+ # The raw A tile of a CTA: a plain 1-CTA TMA into the operand stage, transformed
+ # in place.
+ tma_atom_a, tma_a = cute.nvgpu.make_tiled_tma_atom_A(
+ cpasync.CopyBulkTensorTileG2SOp(),
+ a,
+ a_one,
+ 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_one = cute.slice_(self.b_smem_staged, (None, None, None, 0))
+ tma_atom_b, tma_b = cute.nvgpu.make_tiled_tma_atom_B(
+ b_op, b, b_one, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape
+ )
+ b1_one = cute.slice_(self.b1_smem_staged, (None, None, None, 0))
+ tma_atom_b1, tma_b1 = cute.nvgpu.make_tiled_tma_atom_B(
+ b_op, b1, b1_one, self.mma_tiler1, tiled_mma1, self.cluster_layout_vmnk.shape
+ )
+ self.a_bytes = cute.size_in_bytes(self.io_dtype, a_one)
+ # Both MMAs' B, both CTAs of the pair.
+ self.b_bytes = (
+ cute.size_in_bytes(self.io_dtype, b_one) + cute.size_in_bytes(self.io_dtype, b1_one)
+ ) * 2
+ pairs = (m + 255) // 256
+ self.kernel(
+ tiled_mma,
+ tiled_mma1,
+ tma_atom_a,
+ tma_a,
+ tma_atom_b,
+ tma_b,
+ tma_atom_b1,
+ tma_b1,
+ w_ptr.toint(),
+ rstd.iterator.toint(),
+ out_ptr.toint(),
+ inj_ptr.toint(),
+ m,
+ self.cluster_layout_vmnk,
+ self.a_smem_staged,
+ self.b_smem_staged,
+ self.b1_smem_staged,
+ ).launch(
+ grid=[2 * pairs * self.split, 1, 1],
+ 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,
+ tiled_mma1: cute.TiledMma,
+ tma_atom_a: cute.CopyAtom,
+ mA_mkl: cute.Tensor,
+ tma_atom_b: cute.CopyAtom,
+ mB_nkl: cute.Tensor,
+ tma_atom_b1: cute.CopyAtom,
+ mB1_nkl: cute.Tensor,
+ w_base: cutlass.Int64,
+ rstd_base: cutlass.Int64,
+ out_base: cutlass.Int64,
+ inj_base: cutlass.Int64,
+ m: cutlass.Int32,
+ cluster_layout_vmnk: cute.Layout,
+ a_smem_staged: cute.ComposedLayout,
+ b_smem_staged: cute.ComposedLayout,
+ b1_smem_staged: cute.ComposedLayout,
+ ):
+ hc, hid, rank, split = self.hc, self.hidden, self.rank, self.split
+ warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+ tidx, _, _ = cute.arch.thread_idx()
+ bid, _, _ = cute.arch.block_idx()
+ v = bid % 2
+ is_leader = v == 0
+ pair = bid // 2
+ m_pair = pair // split
+ sp = pair % split
+ kt0 = sp * self.kt
+ cta_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+ cluster_coord = cluster_layout_vmnk.get_flat_coord(cta_rank)
+ if warp_idx == self.tma_warp_id:
+ cpasync.prefetch_descriptor(tma_atom_a)
+ if warp_idx == self.tmb_warp_id:
+ cpasync.prefetch_descriptor(tma_atom_b)
+ cpasync.prefetch_descriptor(tma_atom_b1)
+
+ @cute.struct
+ class SharedStorage:
+ raw_full: cute.struct.MemRange[cutlass.Int64, self.num_stages]
+ b_mbar: cute.struct.MemRange[cutlass.Int64, self.num_b_stages * 2]
+ a_mbar: cute.struct.MemRange[cutlass.Int64, self.num_stages * 2]
+ acc_mbar: cute.struct.MemRange[cutlass.Int64, 2]
+ tmem_dealloc_mbar: cutlass.Int64
+ tmem_holding_buf: cutlass.Int32
+
+ smem = utils.SmemAllocator()
+ storage = smem.allocate(SharedStorage)
+ # Raw tile landed (TMA transaction count), per CTA: the transform warps wait on it.
+ raw_full = storage.raw_full.data_ptr()
+ if warp_idx == self.tma_warp_id:
+ with cute.arch.elect_one():
+ for i in cutlass.range_constexpr(self.num_stages):
+ cute.arch.mbarrier_init(raw_full + i, 1)
+ b_producer, b_consumer = PipelineTmaUmma.create(
+ barrier_storage=storage.b_mbar.data_ptr(),
+ num_stages=self.num_b_stages,
+ producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+ consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1),
+ tx_count=self.b_bytes,
+ cta_layout_vmnk=cluster_layout_vmnk,
+ ).make_participants()
+ # A operand stages: the TMA warp acquires a stage (the MMA has released it), the
+ # transform warps of both CTAs commit it to the leader, the leader's MMA consumes
+ # and releases it.
+ a_pipe = pipeline.PipelineAsyncUmma.create(
+ barrier_storage=storage.a_mbar.data_ptr(),
+ num_stages=self.num_stages,
+ producer_group=pipeline.CooperativeGroup(
+ pipeline.Agent.Thread, 2 * (self.tw // self.groups)
+ ),
+ consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1),
+ cta_layout_vmnk=cluster_layout_vmnk,
+ )
+ a_consumer = a_pipe.make_participants()[1]
+ acc_pipeline = PipelineUmmaAsync.create(
+ barrier_storage=storage.acc_mbar.data_ptr(),
+ num_stages=1,
+ producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+ consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 2 * self.tw),
+ cta_layout_vmnk=cluster_layout_vmnk,
+ )
+ tmem_alloc_barrier = pipeline.NamedBarrier(
+ barrier_id=self.tmem_bar_id, num_threads=32 * (1 + self.tw)
+ )
+ tmem = utils.TmemAllocator(
+ storage.tmem_holding_buf.ptr,
+ barrier_for_retrieve=tmem_alloc_barrier,
+ allocator_warp_id=0,
+ is_two_cta=True,
+ 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_staged.outer,
+ byte_alignment=1024,
+ swizzle=a_smem_staged.inner,
+ )
+ sB = smem.allocate_tensor(
+ element_type=self.io_dtype,
+ layout=b_smem_staged.outer,
+ byte_alignment=1024,
+ swizzle=b_smem_staged.inner,
+ )
+ sB1 = smem.allocate_tensor(
+ element_type=self.io_dtype,
+ layout=b1_smem_staged.outer,
+ byte_alignment=1024,
+ swizzle=b1_smem_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)
+ )
+ thr_mma = tiled_mma.get_slice(v)
+ 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),
+ )
+ b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape)
+ tBsB, tBgB = cpasync.tma_partition(
+ tma_atom_b,
+ cluster_coord[1],
+ b_cta_layout,
+ cute.group_modes(sB, 0, 3),
+ cute.group_modes(tCgB, 0, 3),
+ )
+ gB1_nkl = cute.local_tile(
+ mB1_nkl, cute.slice_(self.mma_tiler1, (0, None, None)), (None, None, None)
+ )
+ tCgB1 = tiled_mma1.get_slice(v).partition_B(gB1_nkl)
+ tB1sB1, tB1gB1 = cpasync.tma_partition(
+ tma_atom_b1,
+ cluster_coord[1],
+ b_cta_layout,
+ cute.group_modes(sB1, 0, 3),
+ cute.group_modes(tCgB1, 0, 3),
+ )
+ b_mcast = cpasync.create_tma_multicast_mask(
+ cluster_layout_vmnk, cluster_coord, mcast_mode=1
+ )
+ tCrA = tiled_mma.make_fragment_A(sA)
+ tCrB = tiled_mma.make_fragment_B(sB)
+ tCrA1 = tiled_mma1.make_fragment_A(sA)
+ tCrB1 = tiled_mma1.make_fragment_B(sB1)
+ acc0_fake = tiled_mma.make_fragment_C(tiled_mma.partition_shape_C(self.mma_tiler[:2]))
+ acc1_fake = tiled_mma1.make_fragment_C(tiled_mma1.partition_shape_C(self.mma_tiler1[:2]))
+
+ pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
+
+ # ------------------------------------------------------------ TMA warp (A, raw)
+ if warp_idx == self.tma_warp_id:
+ cute.arch.griddepcontrol_wait()
+ tAgA_t = tAgA[(None, m_pair, None, 0)]
+ a_state = pipeline.make_pipeline_state(
+ pipeline.PipelineUserType.Producer, self.num_stages
+ )
+ for j in cutlass.range(self.kt, unroll=1):
+ ai = a_state.index
+ a_pipe.producer_acquire(a_state)
+ with cute.arch.elect_one():
+ cute.arch.mbarrier_arrive_and_expect_tx(raw_full + ai, self.a_bytes)
+ cute.copy(
+ tma_atom_a, tAgA_t[(None, kt0 + j)], tAsA[(None, ai)], tma_bar_ptr=raw_full + ai
+ )
+ a_state.advance()
+
+ # ------------------------------------------------------------ TMA warp (B)
+ if warp_idx == self.tmb_warp_id:
+ for j in cutlass.range(self.kt, unroll=1):
+ handle = b_producer.acquire_and_advance()
+ cute.copy(
+ tma_atom_b,
+ tBgB[(None, 0, kt0 + j, 0)],
+ tBsB[(None, handle.index)],
+ tma_bar_ptr=handle.barrier,
+ mcast_mask=b_mcast,
+ )
+ cute.copy(
+ tma_atom_b1,
+ tB1gB1[(None, 0, kt0 + j, 0)],
+ tB1sB1[(None, handle.index)],
+ tma_bar_ptr=handle.barrier,
+ mcast_mask=b_mcast,
+ )
+ b_producer.tail()
+
+ # ------------------------------------------------------------ MMA warp
+ if warp_idx == self.mma_warp_id:
+ tmem.wait_for_alloc()
+ tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+ tCtAcc0 = cute.make_tensor(tmem_ptr, acc0_fake.layout)
+ tCtAcc1 = cute.make_tensor(tmem_ptr + NH0, acc1_fake.layout)
+ acc_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1)
+ if is_leader:
+ acc_pipeline.producer_acquire(acc_state)
+ tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+ tiled_mma1.set(tcgen05.Field.ACCUMULATE, False)
+ for j in cutlass.range(self.kt, unroll=1):
+ b_handle = b_consumer.wait_and_advance()
+ a_handle = a_consumer.wait_and_advance()
+ num_kblocks = cute.size(tCrA, mode=[2])
+ for kk in cutlass.range(num_kblocks, unroll_full=True):
+ cute.gemm(
+ tiled_mma,
+ tCtAcc0,
+ tCrA[(None, None, kk, a_handle.index)],
+ tCrB[(None, None, kk, b_handle.index)],
+ tCtAcc0,
+ )
+ cute.gemm(
+ tiled_mma1,
+ tCtAcc1,
+ tCrA1[(None, None, kk, a_handle.index)],
+ tCrB1[(None, None, kk, b_handle.index)],
+ tCtAcc1,
+ )
+ tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+ tiled_mma1.set(tcgen05.Field.ACCUMULATE, True)
+ b_handle.release()
+ a_handle.release()
+ acc_pipeline.producer_commit(acc_state)
+ acc_state.advance()
+ acc_pipeline.producer_tail(acc_state)
+
+ # ------------------------------------------------------------ transform + epilogue
+ if warp_idx < self.tw:
+ cute.arch.griddepcontrol_wait()
+ tmem.allocate(self.num_tmem_cols)
+ # Before the transform loop: the MMA warp waits for this allocation, and the
+ # MMA must run for the operand ring to drain.
+ tmem.wait_for_alloc()
+ tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+ # Transform mapping (128B-swizzled K-major 128 x 64 bf16 tile): within its group
+ # of `gt` threads, thread t owns the 16-byte chunks at byte offsets 16 (t + gt j),
+ # i.e. rows (t >> 3) + (gt / 8) j and the logical 8-column chunk
+ # q = (t & 7) ^ ((t >> 3) & 7) of every row it owns. The tile is transformed in
+ # place: same offsets in, same out.
+ G = self.groups
+ gt = 32 * self.tw // G
+ nch = 1024 // gt # chunks a thread
+ grp = tidx // gt
+ tg = tidx % gt
+ q = (tg & 7) ^ ((tg >> 3) & 7)
+ row0 = m_pair * 256 + v * 128 + (tg >> 3)
+ op_addr = cutlass.Int32(sA.iterator.toint()) + tg * 16
+ op_stage = cute.size_in_bytes(
+ self.io_dtype, cute.slice_(a_smem_staged, (None, None, None, 0))
+ )
+ rows64 = [cutlass.Int64(cutlass.min(row0 + (gt // 8) * jj, m - 1)) for jj in range(nch)]
+ kpb = hid // BK # k blocks a stream
+ S = self.num_stages
+ # The group's next k block's scales are loaded one iteration ahead: rstd of the
+ # thread's rows, 8 bf16 of the norm weight.
+ rbuf = cute.make_rmem_tensor((nch,), cutlass.Float32)
+ wbuf = cute.make_rmem_tensor((4,), cutlass.Int32)
+ kfirst = kt0 + grp
+ for jj in cutlass.range_constexpr(nch):
+ rbuf[jj] = _ldg_f32(
+ rstd_base + (rows64[jj] * hc + cutlass.Int64(kfirst // kpb)) * 4
+ )
+ w0 = _ldg_v4(w_base + cutlass.Int64(kfirst * BK + q * 8) * 2, cutlass.Int32(1))
+ for e in cutlass.range_constexpr(4):
+ wbuf[e] = w0[e]
+ for j in cutlass.range(grp, self.kt, G, unroll=1):
+ r = [rbuf[jj] for jj in range(nch)]
+ ww = [wbuf[e] for e in range(4)]
+ if j + G < self.kt:
+ kn = kt0 + j + G
+ s1 = kn // kpb
+ for jj in cutlass.range_constexpr(nch):
+ rbuf[jj] = _ldg_f32(rstd_base + (rows64[jj] * hc + cutlass.Int64(s1)) * 4)
+ wn = _ldg_v4(w_base + cutlass.Int64(kn * BK + q * 8) * 2, cutlass.Int32(1))
+ for e in cutlass.range_constexpr(4):
+ wbuf[e] = wn[e]
+ w1 = []
+ for word in ww:
+ w1 += [_add1(_bf16_lo(word)), _add1(_bf16_hi(word))]
+ ai = j % S
+ cute.arch.mbarrier_wait(raw_full + ai, (j // S) % 2)
+ base = op_addr + ai * op_stage
+ for jj in cutlass.range_constexpr(nch):
+ words = _lds_v4(base + jj * (gt * 16))
+ out = []
+ for e in cutlass.range_constexpr(4):
+ # (x * rstd) * (1 + w), each product rounded in fp32, then bf16:
+ # `qwen4_exp_hc_inject_norm`'s normed, bitwise.
+ lo = _mul(_mul(_bf16_lo(words[e]), r[jj]), w1[2 * e])
+ hi = _mul(_mul(_bf16_hi(words[e]), r[jj]), w1[2 * e + 1])
+ out.append(_bf16x2(lo, hi))
+ _sts_v4(base + jj * (gt * 16), *out)
+ cute.arch.fence_proxy("async.shared", space="cta")
+ cute.arch.sync_warp()
+ with cute.arch.elect_one():
+ a_pipe.sync_object_full.arrive(ai, a_pipe.producer_mask)
+
+ # Epilogue: warp w reads TMEM lanes 32 (w % 4).. of accumulator w // 4
+ # (columns [0, NH0) and [NH0, NH0 + NH1)).
+ lt = tidx % 128
+ half = tidx // 128
+ row = m_pair * 256 + v * 128 + lt
+ row64 = cutlass.Int64(row)
+ acc_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1)
+ acc_pipeline.consumer_wait(acc_state)
+ inv = cutlass.Float32(1.0 / hc)
+ epi_tile = (128, 16)
+ for hh in cutlass.range_constexpr(2):
+ if half == hh:
+ nhh = NH0 if hh == 0 else NH1
+ acc_t = cute.make_tensor(
+ tmem_ptr + hh * NH0, (acc0_fake if hh == 0 else acc1_fake).layout
+ )
+ copy_atom_t2r = utils.sm100.get_tmem_load_op(
+ (128, nhh, BK),
+ utils.LayoutEnum.ROW_MAJOR,
+ self.io_dtype,
+ self.acc_dtype,
+ epi_tile,
+ True,
+ )
+ tAcc_epi = cute.flat_divide(acc_t[((None, None), 0, 0)], epi_tile)
+ tiled_copy_t2r = tcgen05.make_tmem_copy(
+ copy_atom_t2r, tAcc_epi[(None, None, 0, 0)]
+ )
+ thr_copy_t2r = tiled_copy_t2r.get_slice(lt)
+ tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi)
+ tTR_cE = thr_copy_t2r.partition_D(cute.make_identity_tensor(epi_tile))
+ frag = cute.make_rmem_tensor(tTR_cE.shape, self.acc_dtype)
+ for c in cutlass.range_constexpr(nhh // 16):
+ col = hh * NH0 + c * 16
+ if cutlass.const_expr(col < rank + hc):
+ cute.copy(tiled_copy_t2r, tTR_tAcc[(None, None, None, 0, c)], frag)
+ if row < m:
+ if cutlass.const_expr(split > 1):
+ pbase = (
+ out_base
+ + ((cutlass.Int64(sp) * m + row64) * NCOL + col) * 4
+ )
+ for qq in cutlass.range_constexpr(4):
+ _st_global_v4(
+ pbase + qq * 16,
+ *[_bits(frag[4 * qq + e]) for e in range(4)],
+ )
+ elif cutlass.const_expr(col + 16 <= rank):
+ vals = []
+ for e in cutlass.range_constexpr(16):
+ z = frag[e] * inv
+ vals.append(z * _sigmoid_fast(z))
+ words = [
+ _bf16x2(vals[2 * e], vals[2 * e + 1]) for e in range(8)
+ ]
+ obase = out_base + (row64 * rank + col) * 2
+ _st_global_v4(obase, *words[0:4])
+ _st_global_v4(obase + 16, *words[4:8])
+ else:
+ # The chunk holding the injection logits (rank % 16 == 0).
+ vals = []
+ for e in cutlass.range_constexpr(hc):
+ z = frag[e] * inv
+ vals.append(cutlass.Float32(2.0) * _sigmoid_fast(z))
+ _st_global_v2(
+ inj_base + row64 * (hc * 2),
+ _bf16x2(vals[0], vals[1]),
+ _bf16x2(vals[2], vals[3]),
+ )
+ cute.arch.sync_warp()
+ with cute.arch.elect_one():
+ acc_pipeline.consumer_release(acc_state)
+ cute.arch.barrier(barrier_id=self.epi_bar_id, number_of_threads=32 * self.tw)
+ tmem.relinquish_alloc_permit()
+ tmem.free(tmem_ptr)
+
+
+@triton.jit
+def _combine_act_kernel(
+ part_ptr,
+ low_ptr,
+ inj_ptr,
+ T,
+ SPLIT: tl.constexpr,
+ RANK: tl.constexpr,
+ HC: tl.constexpr,
+ NCOL: tl.constexpr,
+ BLOCK: tl.constexpr,
+):
+ # The K-split partial logits summed in split order (fp32), then the activations of
+ # the split-1 epilogue (`tl.sigmoid` here, the ex2 / rcp form there).
+ row = tl.program_id(0).to(tl.int64)
+ k = tl.arange(0, BLOCK)
+ mask = k < RANK + HC
+ z = tl.load(part_ptr + row * NCOL + k, mask=mask, other=0.0)
+ for sidx in tl.static_range(1, SPLIT):
+ z += tl.load(part_ptr + (sidx * T + row) * NCOL + k, mask=mask, other=0.0)
+ z = z * (1.0 / HC)
+ tl.store(
+ low_ptr + row * RANK + k, (z * tl.sigmoid(z)).to(low_ptr.dtype.element_ty), mask=k < RANK
+ )
+ ki = k - RANK
+ tl.store(
+ inj_ptr + row * HC + ki,
+ (2.0 * tl.sigmoid(z)).to(inj_ptr.dtype.element_ty),
+ mask=(ki >= 0) & (ki < HC),
+ )
+
+
+_COMPILED: dict = {}
+_SMS: dict = {}
+
+
+def _launcher(hc: int, hidden: int, rank: int, split: int, ptrs):
+ """Compile once per static configuration; later calls reuse the TVM-FFI entry,
+ which launches on the caller's current stream (graph-capturable). `ptrs` is
+ `(hc_out, norm_weight, rstd tensor, first, out, inject)`.
+ """
+ key = (hc, hidden, rank, split)
+ fn = _COMPILED.get(key)
+ if fn is None:
+ kern = HcFirstRescaleKernel(hc, hidden, rank, split)
+ gm = cute.AddressSpace.gmem
+ cptrs = [
+ make_ptr(cutlass.BFloat16, ptrs[0], gm, assumed_align=16),
+ make_ptr(cutlass.BFloat16, ptrs[1], gm, assumed_align=16),
+ from_dlpack(ptrs[2], assumed_align=16).mark_layout_dynamic(leading_dim=1),
+ make_ptr(cutlass.BFloat16, ptrs[3], gm, assumed_align=16),
+ make_ptr(cutlass.BFloat16, ptrs[4], gm, assumed_align=16),
+ make_ptr(cutlass.BFloat16, ptrs[5], gm, assumed_align=8),
+ ]
+ fn = cute.compile(
+ kern,
+ 128,
+ *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 _pick_split(rows: int, k_blocks: int, device: torch.device) -> int:
+ """CTA pairs = ceil(rows / 256) * split: split K in two while the grid stays within
+ one wave of SMs (and K divides), up to `_MAX_SPLIT`.
+ """
+ sms = _SMS.get(device.index)
+ if sms is None:
+ sms = _SMS[device.index] = torch.cuda.get_device_properties(device).multi_processor_count
+ pairs = -(-rows // 256)
+ split = 1
+ while split < _MAX_SPLIT and 2 * pairs * split * 2 <= sms and k_blocks % (2 * split) == 0:
+ split *= 2
+ return split
+
+
+def hc_first_rescale(
+ hc_out: torch.Tensor,
+ rstd: torch.Tensor,
+ norm_weight: torch.Tensor,
+ first_t: torch.Tensor,
+ hc_count: int,
+ rank: int,
+ with_inject: bool,
+):
+ """`(low [T, rank], inject [T, C] or None)`; `first_t` the contiguous padded
+ `[N, C*H]` first weight (N >= 352).
+ """
+ rows, width = hc_out.shape
+ hidden = width // hc_count
+ dev = hc_out.device
+ low = torch.empty((rows, rank), dtype=hc_out.dtype, device=dev)
+ # Written in every case (the kernel stores the injection chunk); dropped without
+ # `with_inject`.
+ inject = torch.empty((rows, hc_count), dtype=hc_out.dtype, device=dev)
+ if rows > 0:
+ split = _pick_split(rows, width // BK, dev)
+ if split == 1:
+ out = low
+ else:
+ out = torch.empty((split, rows, NCOL), dtype=torch.float32, device=dev)
+ ptrs = (
+ hc_out.data_ptr(),
+ norm_weight.data_ptr(),
+ rstd,
+ first_t.data_ptr(),
+ out.data_ptr(),
+ inject.data_ptr(),
+ )
+ _launcher(hc_count, hidden, rank, split, ptrs)(rows, *ptrs)
+ if split > 1:
+ _combine_act_kernel[(rows,)](
+ out,
+ low,
+ inject,
+ rows,
+ SPLIT=split,
+ RANK=rank,
+ HC=hc_count,
+ NCOL=NCOL,
+ BLOCK=triton.next_power_of_2(rank + hc_count),
+ num_warps=4,
+ )
+ return low, (inject if with_inject else None)
+
+
+class Qwen4ExpHcFirstRescaleTc(OpWrapper):
+ """With `C = hc_count`, `H = width // C`, `first` the `[C*H, N]` transposed view
+ of the padded `[N, C*H]` `[down; block_inject; 0]` weight (N >= 352)::
+
+ normed = bf16(hc_out[t, s * H + i] * rstd[t, s] * (1 + norm_weight[s * H + i]))
+ logits = normed @ first[:, : rank + C] # fp32
+ low = silu(logits[:, :rank] / C) # [T, rank]
+ inject = 2 * sigmoid(logits[:, rank : rank + C] / C) # [T, C]
+
+ `normed` is bitwise `qwen4_exp_hc_inject_norm`'s (the same fp32 products, one
+ rounding), formed in shared memory and never written. Returns fresh contiguous
+ `(low, inject)` in `hc_out`'s dtype; `inject` is `None` exactly when
+ `with_inject` is false (the stack-final mixer; `low` is unchanged).
+
+ Fusion boundary: the grouped norm's scaling, the first projection and both
+ activations -- `cublas_mm(normed, first[:, :352])` plus
+ `qwen4_exp_hc_lowrank_act` fed by a site that stored `normed`.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: rows, inject (True). The checkpoint's geometry throughout: hc_count 4,
+ # hidden 2560, rank 320, the padded 384-row first weight, bf16.
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="prefill chunk 16384, one K slice a CTA pair",
+ spec=dict(rows=16384),
+ ),
+ Cell(why="prefill 8192 rows, K split in two", spec=dict(rows=8192)),
+ Cell(why="prefill 2048 rows, K split in eight", spec=dict(rows=2048)),
+ Cell(
+ why="stack-final mixer (no injection), prefill chunk 16384",
+ spec=dict(rows=16384, inject=False),
+ ),
+ Cell(why="ragged tail, 16320 + 70 rows (mixed step)", spec=dict(rows=16390)),
+ Cell(why="five rows, less than one tile", spec=dict(rows=5)),
+ )
+
+ note = """
+ Tier 6 (CuTe DSL written for this target): a tcgen05 GEMM whose A operand is
+ produced by warps of the same CTA from a second tensor; no lower tier has an
+ operand prologue (cuBLAS needs `normed` materialized). Metadata consumed:
+ none. Call sites: `PrefillTarget`'s trunk mixers at or above
+ `_HC_RSTD_MIN_ROWS` rows (every attention / GDN, MoE and stack-final site),
+ fed by `qwen4_exp_hc_inject_rstd`.
+
+ Preconditions: bf16; `hc_out` `[T, C*H]` contiguous; rstd `[T, C]` fp32
+ contiguous; `norm_weight` `[C*H]` contiguous; `first` the transposed view of
+ a contiguous `[N, C*H]` weight, N >= 352; 16-byte aligned pointers;
+ `hc_count == 4`, `H % 64 == 0`, `rank % 16 == 0`, `rank + C <= 352`.
+ Compiles once per `(C, H, rank, split)` on first call; the row count is a
+ runtime argument. Graph-capturable (current stream, no host sync; the split
+ is a host function of the row count).
+
+ Launch: one CTA pair a 256-row tile, the whole K (split 1), or, when that
+ leaves more than half the SMs idle, `split` pairs a tile each over `1/split`
+ of K (2048 rows: split 8, 8192: split 2, 16384: split 1 on 152 SMs) with a
+ Triton combine. Launched with programmatic dependent launch
+ (TRTLLM_ENABLE_PDL=0 turns it off): the weight stream starts before
+ `griddepcontrol.wait`. GB300, 16384 rows: 82 us, against cuBLAS 77.6 +
+ lowrank_act 9.3 on a stored `normed` -- the site's saving is the `normed`
+ write and re-read this removes upstream (`qwen4_exp_hc_inject_rstd`).
+
+ Numerics: `low` / `inject` are one fp32 activation of fp32 logits rounded
+ once (split 1: `ex2.approx` / `rcp.approx` sigmoid; split > 1: the partials
+ summed in split order, `tl.sigmoid`). Measured against the cuBLAS +
+ `qwen4_exp_hc_lowrank_act` path on the same `normed`: split 1 bitwise at
+ 2048-16390 rows (where cuBLAS does not split K itself); otherwise the last
+ bit differs on about 0.5% of `low`, inside the same band.
+
+ Certified: bf16 at the checkpoint's 4 x 2560, rank 320: 16384 (split 1),
+ 8192 (split 2), 2048 (split 8), 5 rows, a partial last tile (16390), the
+ stack-final form; graph replay. Against the float64 product of the
+ once-rounded `normed`, gate band (`rtol = atol / max|ref| = eps`).
+
+ Not certified: zero rows, hc_count other than 4, hidden other than 2560,
+ rank other than 320.
+ """
+
+ def __call__(
+ self,
+ hc_out: torch.Tensor,
+ rstd: torch.Tensor,
+ norm_weight: torch.Tensor,
+ first: torch.Tensor,
+ hc_count: int,
+ rank: int,
+ with_inject: bool,
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ return hc_first_rescale(hc_out, rstd, norm_weight, first.t(), hc_count, rank, with_inject)
+
+ def reference(
+ self,
+ hc_out: torch.Tensor,
+ rstd: torch.Tensor,
+ norm_weight: torch.Tensor,
+ first: torch.Tensor,
+ hc_count: int,
+ rank: int,
+ with_inject: bool,
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ """float64: `normed` rounded once to `hc_out`'s dtype (the op's operand),
+ then the projection and both activations, each output rounded once.
+ """
+ rows, width = hc_out.shape
+ hidden = width // hc_count
+ x = hc_out.double().unflatten(-1, (hc_count, hidden)) * rstd.double().unsqueeze(-1)
+ normed = (x.flatten(-2) * (1.0 + norm_weight.double())).to(hc_out.dtype)
+ logits = F.linear(normed.double(), first[:, : rank + hc_count].t().double())
+ low = F.silu(logits[:, :rank] / hc_count).to(hc_out.dtype)
+ inject = (2.0 * torch.sigmoid(logits[:, rank:] / hc_count)).to(hc_out.dtype)
+ return low, (inject if with_inject else None)
+
+ def is_valid(
+ self,
+ hc_out: torch.Tensor,
+ rstd: torch.Tensor,
+ norm_weight: torch.Tensor,
+ first: torch.Tensor,
+ hc_count: int,
+ rank: int,
+ with_inject: bool,
+ ) -> None:
+ rows, width = hc_out.shape
+ assert hc_out.dtype == norm_weight.dtype == first.dtype == torch.bfloat16, "bf16"
+ assert hc_out.stride() == (width, 1), "hc_out: contiguous rows"
+ assert rstd.dtype == torch.float32 and rstd.shape == (rows, hc_count), "rstd [T, C] fp32"
+ assert rstd.is_contiguous(), "rstd contiguous"
+ assert norm_weight.shape == (width,) and norm_weight.is_contiguous(), "norm_weight [C*H]"
+ assert first.dim() == 2 and first.shape[0] == width and first.shape[1] >= NCOL, (
+ f"first [C*H, N >= {NCOL}]"
+ )
+ assert first.stride() == (1, width), "first must be the transposed view of [N, C*H]"
+ # The epilogue stores the injection chunk as two bf16x2 words (C == 4) and
+ # reads its 16-column chunks wholly as `low` or as injection.
+ assert hc_count == 4, "hc_count 4"
+ assert (width // hc_count) % BK == 0, f"hidden must be a multiple of {BK}"
+ assert rank % 16 == 0 and rank + hc_count <= NCOL, "rank % 16 == 0, rank + C <= 352"
+ assert isinstance(with_inject, bool)
+ for t in (hc_out, norm_weight, first, rstd):
+ assert t.data_ptr() % 16 == 0, "16-byte aligned operands"
+
+
+qwen4_exp_hc_first_rescale_tc = Qwen4ExpHcFirstRescaleTc()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_inject_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_inject_norm.py
new file mode 100644
index 000000000000..ba8d55b24a46
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_inject_norm.py
@@ -0,0 +1,1123 @@
+# 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). No ``torch.ops.trtllm.*`` op or FlashInfer callable
+accumulates a block output into ``hc_count`` streams with per-stream gates and
+normalizes each stream with its own RMS in one launch. It computes what
+``qwen4_exp_post_inject`` followed by a per-stream ``flashinfer_gemma_rmsnorm``
+computes; with ``shared`` / ``gate_logits`` also ``fused_sigmoid_gate_mul_add``
+before it, and with ``topk_weights`` / ``topk_rows`` the MoE runner's finalize.
+
+The decode MoE sites (the finalize fold below `_FEW_WARPS_MIN_ROWS` rows) run a
+CuTe DSL kernel instead (Tier 6), bitwise the Triton kernel's outputs there.
+"""
+
+from __future__ import annotations
+
+import struct
+from typing import Optional, Tuple
+
+import cutlass
+import cutlass.cute as cute
+import torch
+import triton
+import triton.language as tl
+from cutlass._mlir.dialects import llvm
+from cutlass.cute.runtime import from_dlpack
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
+
+#: Trigger dependents at kernel start (`griddepcontrol.launch_dependents`); a no-op
+#: unless the next kernel is launched with programmatic dependent launch.
+_PDL_TRIGGER = get_env_enable_pdl()
+
+#: Rows from which a wide (`hidden >= 2048`) launch runs two warps a program
+#: instead of eight (bandwidth-bound prefill-sized grids).
+_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.
+ 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. 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
+ (`qwen4_exp_post_inject`'s) and the grouped gemma RMS norm of the next
+ site. The routed and shared experts, the gate projection and the mixer
+ that consumes `normed` are not inside.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: rows, inject (True), moe (False: shared + gate_logits [T, 1] too), fold
+ # (0: top-k of the finalize folded in), drop (False: some slots -1), site
+ # ("trunk": the norm weight a trunk site's; "stack_final": the stack-final mixer's).
+ # The checkpoint's geometry throughout: hc_count 4, hidden 2560, bf16, contiguous rows.
+ # The first four reach the target's four distinct launches at tp1 -- the Triton kernel's
+ # eight-warp form (decode), its two-warp form (prefill, here with the combine),
+ # the CuTe fold kernel (decode MoE sites; the -1 slots ride it) and the norm-only
+ # variant -- none of which another covers. The rest are the other layouts': one
+ # form at every row count, its MoE sites the combine over the routed rows already
+ # finalized and reduced across ranks (no fold), so the Triton kernel with the
+ # combine at decode sizes, and every form from 1 row up to max_num_tokens.
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="checkpoint 4x2560 bf16, inject, decode batch 128 (eight-warp form)",
+ spec=dict(rows=128),
+ ),
+ Cell(
+ why="MoE combine, [T, 1] gate, prefill chunk 16384 (the ladder's max_num_tokens)",
+ spec=dict(rows=16384, moe=True),
+ ),
+ Cell(
+ why="MoE combine + finalize fold (top-10, CuTe kernel), three slots -1, decode 128",
+ spec=dict(rows=128, moe=True, fold=10, drop=True),
+ ),
+ Cell(
+ why="norm only (no inject: stack start / after PLE), decode 128",
+ spec=dict(rows=128, inject=False),
+ ),
+ Cell(
+ why="1 row, attention site: inject of the reduced block output, eight-warp form",
+ spec=dict(rows=1, moe=False, inject=True),
+ ),
+ Cell(
+ why="1 row, MoE site: combine over the reduced routed rows, no fold",
+ spec=dict(rows=1, moe=True, inject=True),
+ ),
+ Cell(
+ why="32 rows, MoE site: combine over the reduced routed rows, no fold",
+ spec=dict(rows=32, moe=True, inject=True),
+ ),
+ Cell(
+ why="128 rows, MoE site: combine without the fold (the con-128 decode batch)",
+ spec=dict(rows=128, moe=True, inject=True),
+ ),
+ Cell(
+ why="1024 rows, attention site: two-warp form (an MTP-7 verify of 128 requests)",
+ spec=dict(rows=1024, moe=False, inject=True),
+ ),
+ Cell(
+ why="8192 rows, MoE site: combine, two-warp form (max_num_tokens)",
+ spec=dict(rows=8192, moe=True, inject=True),
+ ),
+ Cell(
+ why="1 row, stack-final site: last MoE combine, the final mixer's norm weight",
+ spec=dict(rows=1, moe=True, inject=True, site="stack_final"),
+ ),
+ Cell(
+ why="1 row, norm only: the stack's first site at decode",
+ spec=dict(rows=1, moe=False, inject=False),
+ ),
+ Cell(
+ why="8192 rows, norm only: the stack's first site at max_num_tokens (two-warp)",
+ spec=dict(rows=8192, moe=False, inject=False),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none; every (row, stream) independent. Call sites: in
+ the trunk forward of the target, once after each block (attention / GDN
+ and MoE) of each layer, normalizing for the next site -- the layer's MoE
+ mixer, the next layer's attention mixer, or the stack-final mixer -- plus
+ norm-only calls at the stack's first site and after the PLE add. 97 calls
+ 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.
+
+ In the default (tp4 / tep4 / dep4) target HC is replicated and runs on the
+ rank's own rows, one form at every row count (1 to max_num_tokens 8192, an
+ MTP-7 verify 1024): the attention site injects the block output after its
+ all-reduce (none at dep4); the MoE site passes the routed rows already
+ finalized and all-reduced / reduce-scattered with this rank's replicated
+ shared expert output and gate logit -- the combine, never the fold; the
+ last layer's MoE site normalizes with the stack-final mixer's weight
+ (`hyper_connection_mixer.hc_norm`); the stack's first site is norm only.
+ The MTP draft layer's sites are the trunk's too, at a draft step's rows (1
+ to max_num_tokens): norm only over the trunk's residual (`pre_h_norm`, a
+ column view of the trunk's packed output on the first draft step) and
+ before the shared head's mixer; the attention site with the front end's
+ `x_h + x_e` as its injection at unit weight (`inject` all ones: bitwise
+ the bf16 add); the MoE-side site after the output projection's all-reduce.
+
+ Preconditions: 2-D operands, last-dim stride 1 (row strides free); the
+ width divisible by `hc_count`; `inject` `[T, hc_count]`, `block_out` and
+ `shared` `[T, hidden]` of one dtype (with the fold `block_out` is `[P,
+ hidden]`, any `P`, read through int32 `topk_rows` `[T, k]`); `gate_logits`
+ `T` elements along its first axis (`[T, 1]` or `[T]`, any row stride);
+ `norm_weight` contiguous `[width]`. All on one device. `shared` and
+ `gate_logits` come together and only with `block_out`.
+
+ Numerics: the combine is fp32 with one rounding to `block_out`'s dtype
+ (bitwise `fused_sigmoid_gate_mul_add`: the same `tl.sigmoid` and fma); the
+ accumulate is fp32 with one rounding (bitwise what `qwen4_exp_post_inject`
+ stores for 16-bit operands up to fma contraction); the norm is fp32 with
+ `rsqrt`, one rounding. Against the float64 reference the combine can land
+ one ulp of the block output away (its rounding midpoint), which the
+ accumulate carries scaled by `inject`; the test's band states that.
+ Graph-capturable by construction (no host sync, outputs allocated per
+ call); not certified.
+
+ Launch: one program per (row, stream), stream-minor in a 1-D grid, so a
+ row's `block_out` / `shared` are read from HBM once and from L2 by the
+ other `hc_count - 1` programs; with the fold, one program per two streams
+ of a row (the gathered rows are read and summed half as often). Every
+ program releases programmatic dependents at its start
+ (`griddepcontrol.launch_dependents`; off with TRTLLM_ENABLE_PDL=0): a
+ consumer launched with PDL (the decode mixer's
+ `qwen4_exp_hc_decode_tc_first`) streams its weights meanwhile and waits for
+ this grid before reading `normed`. A no-op for any other consumer. At a
+ hidden of 2048 or more a program runs eight warps below
+ `_FEW_WARPS_MIN_ROWS` rows and two from it; the RMS sum's fp32 reduction
+ order follows the warp count.
+
+ The decode MoE sites -- the fold below `_FEW_WARPS_MIN_ROWS` rows, bf16
+ operands, `hc_count` 4, a hidden in (2048, 4096], 16-byte aligned rows with
+ row strides a multiple of 16 -- run `_FoldNormKernel` (CuTe DSL) instead:
+ a row's two CTAs one cluster, CTA j forming half the block output (the
+ finalize and the combine) into both CTAs' shared memory, then each
+ accumulating and normalizing two streams, in the Triton kernel's layout and
+ reduction tree -- bitwise its `hc_out` and `normed` (tested at 511 rows, the
+ last before the switch). It releases dependents at its start as the Triton
+ kernel does.
+
+ Certified on sm_103 (GB300), bf16, the checkpoint's 4 x 2560, PDL trigger
+ on (the default): the Triton kernel and the CuTe fold kernel are JIT-
+ compiled at first call for the running GPU, and the cells ran that path.
+ tp1: the injection at 128 rows (eight-warp form), the combine at 16384
+ rows (two-warp form), the combine with the fold (three slots -1) at 128
+ rows (CuTe kernel), norm-only at 128 rows; the CuTe fold bitwise the
+ Triton kernel at 511 rows.
+ default: the injection at 1 row (eight-warp) and 1024 rows (two-warp);
+ the combine without the fold at 1, 32 and 128 rows (eight-warp) and
+ 8192 rows (two-warp); the combine at 1 row under the stack-final
+ weight (RMS 3.1, -0.37 to 12.4, so `1 + w` up to 13.4); norm-only at 1
+ and 8192 rows (two-warp).
+ Measured against the float64 reference rounded once: the decode cells (1,
+ 32 rows) bitwise or one ulp on 4 of 327680 `normed` elements; at 128 rows
+ one `hc_out` element and 5 `normed` elements of 1310720 differ; at 1024
+ rows `hc_out` bitwise, 74 of 10485760 `normed` one ulp; norm-only at 8192
+ rows 602 of 83886080 one ulp. The combine at 8192 / 16384 rows moves 131 /
+ 192 `hc_out` elements, each inside one ulp of `hc_out` plus one ulp of the
+ combine times `|inject|` (worst 0.95 of that band; up to 31 ulp of a
+ near-zero residual), and `normed` within one ulp of the norm of the
+ stored `hc_out`. Relative L2 at most 1.8e-5, cosine 1.0 on both outputs at
+ every cell.
+
+ Not certified: hc_count other than 4, hidden other than 2560, float16
+ operands, a flat `[T]` gate, row strides other than the width, zero rows,
+ TRTLLM_ENABLE_PDL=0, CUDA-graph replay (the kernels take them; no cell
+ drives them).
+
+ Hard limits: row offsets are int64 (`program_id(0)` widened); the grid is
+ `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; `hc_out`
+ is the float64 accumulate rounded once; `normed` is the float64 norm of
+ that rounded residual, rounded once. No other intermediate rounding."""
+ dtype = hyper_input.dtype
+ rows, width = hyper_input.shape
+ hidden = width // hc_count
+ 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/ops/hc/qwen4_exp_hc_inject_rstd.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_inject_rstd.py
new file mode 100644
index 000000000000..597216df40fc
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_inject_rstd.py
@@ -0,0 +1,258 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Hyper-connection post-block injection plus the next site's per-(row, stream) RMS
+statistic, without the normalized copy: `qwen4_exp_hc_inject_norm` with its
+`normed` store replaced by one fp32 rstd per (row, stream).
+
+Tier 5 (Triton written for this target). At prefill sizes the grouped norm's output
+(`[T, C*H]`, 335 MB at 16384 rows) is a full extra write whose only readers are the
+site's two projections; those form `bf16(hc_out * rstd * (1 + w))` themselves
+(`qwen4_exp_hc_first_rescale_tc` on its A tiles, `qwen4_exp_hc_up_gate_mean_rstd` in
+its epilogue), so the residual pass moves the residual in and out once and nothing
+else.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+#: Rows from which a wide (`hidden >= 2048`) launch runs two warps a program instead
+#: of eight (`qwen4_exp_hc_inject_norm`'s split; the reduction tree follows it).
+_FEW_WARPS_MIN_ROWS = 512
+
+
+@triton.jit
+def _hc_inject_rstd_kernel(
+ hc_ptr,
+ block_ptr,
+ inject_ptr,
+ shared_ptr,
+ gate_ptr,
+ hc_out_ptr,
+ rstd_ptr,
+ stride_hc,
+ stride_block,
+ stride_inject,
+ stride_shared,
+ stride_gate,
+ stride_out,
+ eps,
+ HC_COUNT: tl.constexpr,
+ HIDDEN: tl.constexpr,
+ BLOCK: tl.constexpr,
+ HAS_INJECT: tl.constexpr,
+ HAS_SHARED: tl.constexpr,
+):
+ # One program per (row, stream), stream-minor: a row's block output (and the
+ # combine's operands) come from HBM once and from L2 for the other programs. The
+ # loads, the fp32 arithmetic and the reduction are `qwen4_exp_hc_inject_norm`'s
+ # (without the fold), so `hc_out` and rstd are bitwise what that kernel stores and
+ # multiplies by at the same warp count.
+ pid = tl.program_id(0).to(tl.int64)
+ row = pid // HC_COUNT
+ stream = (pid % HC_COUNT).to(tl.int32)
+ cols = tl.arange(0, BLOCK)
+ mask = cols < HIDDEN
+ offs = stream * HIDDEN + cols
+ x = tl.load(hc_ptr + row * stride_hc + offs, mask=mask, other=0.0).to(tl.float32)
+ if HAS_INJECT:
+ 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.
+ 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)
+ g = tl.load(inject_ptr + row * stride_inject + stream).to(tl.float32)
+ 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)
+ tl.store(rstd_ptr + row * HC_COUNT + stream, rstd)
+
+
+class Qwen4ExpHcInjectRstd(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]
+ 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 rstd is computed from that
+ rounded value in fp32. Returns `(hc_out, rstd)`: `hc_out` a fresh `[T, C*H]` in
+ `hyper_input`'s dtype, rstd a fresh contiguous `[T, C]` fp32. A site's normalized
+ stream is `bf16(hc_out * rstd * (1 + norm_weight))`, bitwise
+ `qwen4_exp_hc_inject_norm`'s `normed` (same fp32 operations and rounding).
+
+ `block_out=None, inject=None` is the statistic alone: `hc_out` *is* `hyper_input`
+ (not a copy) and only rstd 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 injected block output is
+ `block_out + sigmoid(gate_logits) * shared` in fp32, rounded once to `block_out`'s
+ dtype (`qwen4_exp_hc_inject_norm`'s combine). None of the inputs is written.
+
+ Fusion boundary: the shared-expert combine (optional), the stream-broadcast gated
+ accumulate and the per-(row, stream) RMS statistic. The norm weight is not an
+ argument: the consumers apply it.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: rows, inject (True), moe (False: shared + gate_logits [T, 1] too). The
+ # checkpoint's geometry throughout: hc_count 4, hidden 2560, bf16, contiguous rows.
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="MoE combine, prefill chunk 16384 (the mixed step's size)",
+ spec=dict(rows=16384, moe=True),
+ ),
+ Cell(why="injection after attention / GDN, prefill chunk 16384", spec=dict(rows=16384)),
+ Cell(why="MoE combine, prefill 2048", spec=dict(rows=2048, moe=True)),
+ Cell(why="injection after attention / GDN, prefill 2048", spec=dict(rows=2048)),
+ Cell(
+ why="statistic only (stack start / after PLE), ragged 16390 rows",
+ spec=dict(rows=16390, inject=False),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none; every (row, stream) independent. Call sites: the trunk
+ sites of `PrefillTarget` at or above `_HC_RSTD_MIN_ROWS` rows, in place of
+ `qwen4_exp_hc_inject_norm`: once after each block (attention / GDN and MoE) of
+ each layer, plus the statistic-only calls at the stack's first site and after
+ the PLE add. Their mixers are `qwen4_exp_hc_first_rescale_tc` and
+ `qwen4_exp_hc_up_gate_mean_rstd(_fp4)`, which take the site's norm weight.
+
+ 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; `gate_logits` `T` elements along its first axis.
+ `shared` and `gate_logits` come together and only with `block_out`.
+
+ Numerics: `hc_out` and rstd are bitwise `qwen4_exp_hc_inject_norm`'s at the same
+ row count (that kernel's two-warp form from 512 rows, eight below); rstd is the
+ fp32 `rsqrt` of the fp32 mean of squares. Graph-capturable (no host sync, outputs
+ allocated per call).
+
+ Launch: one program per (row, stream), stream-minor in a 1-D grid. GB300, 16384
+ rows: 0.75 GB (0.84 GB with the combine) at about 7 TB/s, against
+ `qwen4_exp_hc_inject_norm`'s 1.09 / 1.17 GB.
+
+ Certified: bf16, the checkpoint's 4 x 2560: the combine and the injection at 2048
+ and 16384 rows, the statistic alone at 16390 rows; `hc_out` and the rebuilt
+ normed stream bitwise `qwen4_exp_hc_inject_norm`'s; graph replay.
+
+ Not certified: hc_count other than 4, hidden other than 2560, float16, the
+ finalize fold (`topk_rows`; the prefill sites do not fold), zero rows.
+
+ Hard limits: row offsets are int64; the grid is `rows * hc_count` programs along
+ x (2^31 - 1).
+ """
+
+ def __call__(
+ self,
+ hyper_input: torch.Tensor,
+ block_out: Optional[torch.Tensor],
+ inject: Optional[torch.Tensor],
+ eps: float,
+ hc_count: int,
+ shared: Optional[torch.Tensor] = None,
+ gate_logits: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ rows, width = hyper_input.shape
+ hidden = width // hc_count
+ rstd = torch.empty((rows, hc_count), dtype=torch.float32, device=hyper_input.device)
+ has_inject = block_out is not None
+ has_shared = shared is not None
+ hc_out = torch.empty_like(hyper_input) if has_inject else hyper_input
+ if rows > 0:
+ block = triton.next_power_of_2(hidden)
+ _hc_inject_rstd_kernel[(rows * 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,
+ hc_out,
+ rstd,
+ 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,
+ hc_out.stride(0),
+ eps,
+ HC_COUNT=hc_count,
+ HIDDEN=hidden,
+ BLOCK=block,
+ HAS_INJECT=has_inject,
+ HAS_SHARED=has_shared,
+ num_warps=(2 if rows >= _FEW_WARPS_MIN_ROWS else 8) if block >= 2048 else 4,
+ )
+ return hc_out, rstd
+
+ def reference(
+ self,
+ hyper_input: torch.Tensor,
+ block_out: Optional[torch.Tensor],
+ inject: Optional[torch.Tensor],
+ eps: float,
+ hc_count: int,
+ shared: Optional[torch.Tensor] = None,
+ gate_logits: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """float64: the combine (with `shared`) rounded once to `block_out`'s dtype,
+ `hc_out` the accumulate rounded once, rstd the float64 statistic of that
+ rounded residual returned in fp32 -- the dataflow the op states.
+ """
+ dtype = hyper_input.dtype
+ rows, width = hyper_input.shape
+ hidden = width // hc_count
+ if block_out is None:
+ hc_out = hyper_input
+ else:
+ 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))
+ rstd = torch.rsqrt(x.pow(2).mean(dim=-1) + eps).float()
+ return hc_out, rstd
+
+ def is_valid(
+ self,
+ hyper_input: torch.Tensor,
+ block_out: Optional[torch.Tensor],
+ inject: Optional[torch.Tensor],
+ eps: float,
+ hc_count: int,
+ shared: Optional[torch.Tensor] = None,
+ gate_logits: 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"
+ rows, width = hyper_input.shape
+ assert width % hc_count == 0, f"width {width} not divisible by hc_count {hc_count}"
+ if block_out is not None:
+ assert block_out.shape == (rows, width // hc_count) and block_out.stride(1) == 1
+ assert inject.shape == (rows, hc_count) and inject.stride(1) == 1
+ if shared is not None:
+ assert shared.shape == (rows, width // hc_count) and shared.stride(1) == 1
+ 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_rstd = Qwen4ExpHcInjectRstd()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_lowrank_act.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_lowrank_act.py
new file mode 100644
index 000000000000..bb5368c55c33
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_lowrank_act.py
@@ -0,0 +1,184 @@
+# 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.
+
+Called by the qwen4_exp target (`qwen3_8_flash_next_nvfp4__sm_103__default`) at every
+layout.
+"""
+
+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).
+ # The default target's HC mixers run this form at every row count, on the fresh
+ # fp32 `cublas_mm` output over the first `hc_first_cols` = ceil((320 + 4) / 32) * 32
+ # = 352 columns of the first projection: hence the 352-wide cells.
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="checkpoint rank 320 + 4, padded to 384, 128 rows", spec=dict(rows=128)),
+ Cell(why="checkpoint rank 320 + 4, padded to 384, prefill 8192", spec=dict(rows=8192)),
+ Cell(why="stack-final mixer: no inject, 8192 rows", spec=dict(rows=8192, inject=False)),
+ Cell(
+ why="default decode HC mixer: one token, 352-wide logits",
+ spec=dict(rows=1, logits_width=352),
+ ),
+ Cell(
+ why="default dep4 rank's largest captured decode batch: 32 rows",
+ spec=dict(rows=32, logits_width=352),
+ ),
+ Cell(
+ why="default tep4 MTP-7 verify: 128 requests x 8 rows = 1024",
+ spec=dict(rows=1024, logits_width=352),
+ ),
+ Cell(
+ why="default prefill chunk at max_num_tokens: 8192 rows, 352-wide",
+ spec=dict(rows=8192, logits_width=352),
+ ),
+ Cell(
+ why="default stack-final mixer (no inject): one decode token",
+ spec=dict(rows=1, inject=False, logits_width=352),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none. Call sites: every hyper-connection mixer, between
+ the first projection (`cublas_mm`, fp32 out, 352 columns) and the up
+ projection. HC is replicated, so the rows are this rank's tokens at every
+ member: decode 1..128 (tp4 / tep4) or 1..32 (a dep4 rank), MTP verify up to
+ 128 x 8 = 1024 rows (tep4, draft 7), prefill one chunk of up to 8192 rows
+ (`max_num_tokens`). Preconditions: `logits` fp32, unit column stride, at
+ least `rank (+ hc_count)` columns; `hc_count` a power of two.
+ Graph-capturable.
+
+ Source: a Triton kernel JIT-compiled at first call for the running device;
+ the cells ran on sm_103. One program per row with an int64 row offset;
+ 2^31 elements is 6.1e6 rows of 352, far past the envelope.
+
+ Certified: bf16 outputs from fp32 logits 352 wide (the call site) at 1, 32,
+ 1024 and 8192 rows with the injection and at 1 row without it, and 384 wide
+ at 128 and 8192 rows (without the injection at 8192); graph replay.
+ Against the float64 reference: bitwise at 1, 32 and 128 rows; at 1024 and
+ 8192 rows at most 0.99 ulp on a few `low` elements, relative L2 <= 1.6e-5.
+
+ Not certified: zero rows; rows past 8192; hc_count other than 4; rank
+ other than 320; float16 outputs. The kernel has no row-count switch, so row
+ counts between the cells take the same path, and the logits width enters
+ only as the row stride (128 rows is driven at 384 wide only).
+ """
+
+ def __call__(
+ 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/ops/hc/qwen4_exp_hc_up_gate_mean.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_up_gate_mean.py
new file mode 100644
index 000000000000..e4a4561ecd0c
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_up_gate_mean.py
@@ -0,0 +1,922 @@
+# 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.
+
+A persistent warp-specialized CuTe DSL kernel: one TMA warp loads `low` and
+the up weight, one MMA warp issues `tcgen05.mma` into a double-buffered TMEM
+accumulator, four epilogue warps read the accumulator, load the matching
+normed elements straight from global memory and write the mean. An MMA tile
+is 128 rows by `C * BH` up-projection outputs: the `BH`-wide hidden slice
+`[h0, h0 + BH)` of every stream, read through a strided view of the unmodified
+`[C*H, R]` weight (no copy, no re-layout), so the stream reduction stays inside
+one thread.
+"""
+
+from __future__ import annotations
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+import cutlass.utils as utils
+import torch
+import torch.nn.functional as F
+from cutlass._mlir.dialects import llvm
+from cutlass.cute.nvgpu import cpasync, tcgen05
+from cutlass.cute.runtime import from_dlpack
+from cutlass.cutlass_dsl import T, dsl_user_op
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+from tensorrt_llm._torch.kernels.blackwell.custom_pipeline import PipelineTmaUmma, PipelineUmmaAsync
+from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
+
+#: Programmatic dependent launch (sm_90+), off with TRTLLM_ENABLE_PDL=0.
+_PDL = get_env_enable_pdl()
+
+
+@dsl_user_op
+def _sigmoid(x, *, loc=None, ip=None):
+ """sigmoid(x) = 0.5 * tanh(0.5 * x) + 0.5: one MUFU op (tanh.approx.f32,
+ max relative error 2^-10.99) instead of ex2 + rcp."""
+ return cutlass.Float32(
+ llvm.inline_asm(
+ T.f32(),
+ [cutlass.Float32(x).ir_value(loc=loc, ip=ip)],
+ "{ .reg .f32 t; mul.f32 t, $1, 0f3F000000; tanh.approx.f32 t, t; "
+ "fma.rn.f32 $0, t, 0f3F000000, 0f3F000000; }",
+ "=f,f",
+ has_side_effects=False,
+ is_align_stack=False,
+ asm_dialect=llvm.AsmDialect.AD_ATT,
+ loc=loc,
+ ip=ip,
+ )
+ )
+
+
+# The NVFP4 epilogue (`fp4=True`) is the `torch.ops.trtllm.fp4_quantize` recipe
+# (`cvt_warp_fp16_to_fp4`, e4m3 scales) instruction for instruction, every
+# multiply an explicit `mul.rn.f32` so that no compiler contraction or
+# reassociation can move a rounding.
+
+
+def _asm(ret, args, text, constraints, loc, ip, side_effects=False):
+ return llvm.inline_asm(
+ ret,
+ [a.ir_value(loc=loc, ip=ip) for a in args],
+ text,
+ constraints,
+ has_side_effects=side_effects,
+ is_align_stack=False,
+ asm_dialect=llvm.AsmDialect.AD_ATT,
+ loc=loc,
+ ip=ip,
+ )
+
+
+@dsl_user_op
+def _rcp_approx_ftz(x, *, loc=None, ip=None):
+ """`rcp.approx.ftz.f32`, the quantizer's `reciprocal_approximate_ftz`."""
+ return cutlass.Float32(
+ _asm(T.f32(), [cutlass.Float32(x)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", loc, ip)
+ )
+
+
+@dsl_user_op
+def _absmax16(v, *, loc=None, ip=None):
+ """max |v_i| over 16 fp32 values (exact; NaN-quiet like the bf16 `__hmax2`)."""
+ body = "abs.f32 m, $1; " + "".join(f"abs.f32 t, ${i}; max.f32 m, m, t; " for i in range(2, 17))
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [cutlass.Float32(x) for x in v],
+ "{ .reg .f32 m, t; " + body + "mov.f32 $0, m; }",
+ "=f" + ",f" * 16,
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _e4m3_scale(amax, gs, rcp6, *, loc=None, ip=None):
+ """The block's scale byte: e4m3(gs * (amax * rcp(6))), `cvt.rn.satfinite`."""
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [cutlass.Float32(amax), cutlass.Float32(gs), cutlass.Float32(rcp6)],
+ "{ .reg .f32 t; .reg .b16 h; mul.rn.f32 t, $1, $3; mul.rn.f32 t, $2, t; "
+ "cvt.rn.satfinite.e4m3x2.f32 h, 0f00000000, t; cvt.u32.u16 $0, h; }",
+ "=r,f,f,f",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _e2m1_out_scale(sf, rcp_gs, amax, *, loc=None, ip=None):
+ """amax != 0 ? rcp(float(e4m3 sf) * rcp(gs)) : 0, the data multiplier."""
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [cutlass.Int32(sf), cutlass.Float32(rcp_gs), cutlass.Float32(amax)],
+ "{ .reg .b16 h, lo, hi; .reg .b32 f2; .reg .f32 s, t; .reg .pred p; "
+ "cvt.u16.u32 h, $1; cvt.rn.f16x2.e4m3x2 f2, h; mov.b32 {lo, hi}, f2; "
+ "cvt.f32.f16 s, lo; mul.rn.f32 t, s, $2; rcp.approx.ftz.f32 t, t; "
+ "setp.neu.f32 p, $3, 0f00000000; selp.f32 $0, t, 0f00000000, p; }",
+ "=f,r,f,f",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _e2m1x8(v, scale, *, loc=None, ip=None):
+ """Eight values times `scale`, each to e2m1 (`cvt.rn.satfinite`), packed
+ into 4 bytes; element `2k` in byte k's low nibble."""
+ muls = "".join(f"mul.rn.f32 a{i}, ${i + 1}, $9; " for i in range(8))
+ cvts = "".join(f"cvt.rn.satfinite.e2m1x2.f32 b{k}, a{2 * k + 1}, a{2 * k}; " for k in range(4))
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [cutlass.Float32(x) for x in v] + [cutlass.Float32(scale)],
+ "{ .reg .f32 a<8>; .reg .b8 b<4>; " + muls + cvts + "mov.b32 $0, {b0, b1, b2, b3}; }",
+ "=r" + ",f" * 9,
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _st_global_v4(addr, a, b, c, d, *, loc=None, ip=None):
+ """One 16-byte `st.global.v4.b32` (a row's packed e2m1 for 32 columns): a
+ row-per-thread epilogue pays per store instruction, not per byte."""
+ _asm(
+ T.i32(),
+ [cutlass.Int64(addr)] + [cutlass.Int32(x) for x in (a, b, c, d)],
+ "{ st.global.v4.b32 [$1], {$2, $3, $4, $5}; mov.u32 $0, 0; }",
+ "=r,l,r,r,r,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+
+
+@dsl_user_op
+def _st_global_b16(addr, x, *, loc=None, ip=None):
+ """One 2-byte store of the low half of `x` (two e4m3 scale bytes)."""
+ _asm(
+ T.i32(),
+ [cutlass.Int64(addr), cutlass.Int32(x)],
+ "{ .reg .b16 t; cvt.u16.u32 t, $2; st.global.b16 [$1], t; mov.u32 $0, 0; }",
+ "=r,l,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+
+
+class HcUpGateMeanKernel:
+ """`mixed[t, h] = 1/C * sum_s sigmoid((low @ up^T)[t, s*H + h]) * normed[t, s*H + h]`.
+
+ Compile-time: `block_h` (hidden columns per tile, 32 or 64), `hc` (C),
+ `hidden` (H), `rank` (R), `fp4` (also write the NVFP4 of `mixed`: packed
+ e2m1 `[T, H/2]` and linear e4m3 scales `[T, H/16]` under the global scale
+ `gs`). Runtime: the row count and the pointers; without `fp4` the `xq`,
+ `sf` and `gs` pointers are never read.
+ """
+
+ def __init__(
+ self, block_h: int, hc: int, hidden: int, rank: int, num_n_stage: int = 2, fp4: bool = False
+ ):
+ assert block_h in (32, 64) and hidden % block_h == 0 and rank % 64 == 0
+ self.bh, self.hc, self.hidden, self.rank = block_h, hc, hidden, rank
+ self.fp4 = fp4
+ self.acc_dtype = cutlass.Float32
+ self.io_dtype = cutlass.BFloat16
+ self.mma_tiler = (128, hc * block_h, 1)
+ self.epi_cols = 32
+ self.cta_group = tcgen05.CtaGroup.ONE
+ self.cluster_shape_mn = (1, 1)
+ self.epilogue_warp_id = (0, 1, 2, 3)
+ self.mma_warp_id = 4
+ self.tma_warp_id = 5
+ self.nrm_warp_id = 6
+ self.threads_per_cta = 32 * 7
+ self.num_n_stage = num_n_stage
+ self.epilog_sync_bar_id = 1
+ self.tmem_alloc_sync_bar_id = 2
+ self.num_acc_stage = 2
+
+ def _setup(self, a_major, b_major):
+ tiled_mma = utils.sm100.make_trivial_tiled_mma(
+ self.io_dtype, a_major, b_major, self.acc_dtype, self.cta_group, self.mma_tiler[:2]
+ )
+ inst_k = cute.size(tiled_mma.shape_mnk, mode=[2])
+ self.mma_tiler = (self.mma_tiler[0], self.mma_tiler[1], inst_k * 4)
+ self.cta_tile_shape_mnk = self.mma_tiler
+ self.cluster_layout_vmnk = cute.tiled_divide(
+ cute.make_layout((*self.cluster_shape_mn, 1)), (tiled_mma.thr_id.shape,)
+ )
+ self.epi_tile = (self.mma_tiler[0], self.epi_cols)
+ a_one = utils.sm100.make_smem_layout_a(tiled_mma, self.mma_tiler, self.io_dtype, 1)
+ b_one = utils.sm100.make_smem_layout_b(tiled_mma, self.mma_tiler, self.io_dtype, 1)
+ ab_bytes = cute.size_in_bytes(self.io_dtype, a_one) + cute.size_in_bytes(
+ self.io_dtype, b_one
+ )
+ smem_capacity = utils.get_smem_capacity_in_bytes()
+ # One normed stage is the 32-column slice of every stream for one tile half.
+ self.n_box = (self.mma_tiler[0], self.epi_cols)
+ self.num_n_bytes = self.hc * self.n_box[0] * self.n_box[1] * 2
+ self.num_ab_stage = min(
+ 8, (smem_capacity - 2048 - self.num_n_stage * self.num_n_bytes) // ab_bytes
+ )
+ self.n_pre = min(self.num_ab_stage, self.rank // self.mma_tiler[2])
+ self.n_smem_layout_staged = utils.sm100.make_smem_layout_epi(
+ self.io_dtype, utils.LayoutEnum.ROW_MAJOR, self.n_box, self.num_n_stage * self.hc
+ )
+ self.a_smem_layout_staged = utils.sm100.make_smem_layout_a(
+ tiled_mma, self.mma_tiler, self.io_dtype, self.num_ab_stage
+ )
+ self.b_smem_layout_staged = utils.sm100.make_smem_layout_b(
+ tiled_mma, self.mma_tiler, self.io_dtype, self.num_ab_stage
+ )
+ acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
+ acc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage))
+ self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(acc_fake, arch="sm_100")
+ return tiled_mma
+
+ @cute.jit
+ def __call__(
+ self,
+ m: cutlass.Int32,
+ low_ptr: cute.Pointer,
+ up_ptr: cute.Pointer,
+ normed_ptr: cute.Pointer,
+ mixed: cute.Tensor,
+ xq_ptr: cute.Pointer,
+ sf_ptr: cute.Pointer,
+ gs_ptr: cute.Pointer,
+ max_active_clusters: cutlass.Constexpr,
+ stream: cuda.CUstream,
+ ):
+ bh, hc, hid, rank = self.bh, self.hc, self.hidden, self.rank
+ nt = hid // bh
+ width = hc * hid
+ # low [m, R] (K-major A); the up weight [C*H, R] seen as N = (BH, C, H/BH) so
+ # that one N tile holds hidden slice [h0, h0+BH) of every stream.
+ a = cute.make_tensor(low_ptr, cute.make_layout((m, rank, 1), stride=(rank, 1, 0)))
+ b = cute.make_tensor(
+ up_ptr,
+ cute.make_layout(((bh, hc, nt), rank, 1), stride=((rank, hid * rank, bh * rank), 1, 0)),
+ )
+ # normed [m, C*H] is staged through shared memory by TMA in [128, 32] boxes;
+ # mixed [m, H] is written in the accumulator's N order, its stream mode
+ # with stride 0 (every stream maps onto the one output column).
+ nrm = cute.make_tensor(normed_ptr, cute.make_layout((m, width), stride=(width, 1)))
+ mix = cute.make_tensor(
+ mixed.iterator, cute.make_layout((m, (bh, hc, nt), 1), stride=(hid, (1, 0, bh), 0))
+ )
+ # NVFP4 outputs (`fp4` only): 8 e2m1 codes per Int32, one e4m3 byte per 16.
+ xq = cute.make_tensor(xq_ptr, cute.make_layout((m, hid // 8), stride=(hid // 8, 1)))
+ xsf = cute.make_tensor(sf_ptr, cute.make_layout((m, hid // 16), stride=(hid // 16, 1)))
+ gsc = cute.make_tensor(gs_ptr, cute.make_layout(1))
+ k_major = utils.LayoutEnum.ROW_MAJOR.mma_major_mode()
+ tiled_mma = self._setup(k_major, k_major)
+
+ a_op = utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id)
+ a_smem = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
+ tma_atom_a, tma_a = cute.nvgpu.make_tiled_tma_atom_A(
+ a_op, a, a_smem, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape
+ )
+ b_op = utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id)
+ b_smem = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
+ tma_atom_b, tma_b = cute.nvgpu.make_tiled_tma_atom_B(
+ b_op, b, b_smem, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape
+ )
+ tma_atom_n, tma_n = cpasync.make_tiled_tma_atom(
+ cpasync.CopyBulkTensorTileG2SOp(),
+ nrm,
+ cute.select(self.n_smem_layout_staged, mode=[0, 1]),
+ self.n_box,
+ )
+ self.num_tma_load_bytes = cute.size_in_bytes(self.io_dtype, a_smem) + cute.size_in_bytes(
+ self.io_dtype, b_smem
+ )
+ # Rastered along N: consecutive CTAs take consecutive hidden slices of the
+ # same rows, so the normed reads in flight cover whole rows rather than
+ # one 64-byte column block of many rows.
+ 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: MMA N = 4 * 32 = 128, five A/B stages.
+_BLOCK_H = 32
+
+
+class Qwen4ExpHcUpGateMean(OpWrapper):
+ """With `C = hc_count`, `H = width // C`, `low` the `[T, R]` SiLU bottleneck
+ and `up` the `[R, C*H]` transposed view of the `[C*H, R]` up weight::
+
+ mixed[t, i] = mean over s of sigmoid((low @ up)[t, s*H + i]) * normed[t, s*H + i]
+
+ The GEMM accumulates in fp32 in tensor memory and is never rounded or
+ written; sigmoid, product and stream mean are fp32 with one rounding to
+ `normed`'s dtype. Returns a fresh `[T, H]` tensor. Fusion boundary: the up
+ projection, the sigmoid gate and the stream mean (two launches as `cublas_mm(low,
+ up)` followed by a separate gate-and-mean pass).
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: rows, hc_count (4), hidden (2560), rank (320); bf16 only.
+ CELLS: tuple[Cell, ...] = (
+ tuple(
+ Cell(why=f"checkpoint 4x2560, rank 320, {why}", spec=dict(rows=t))
+ for t, why in (
+ (128, "decode 128 rows, one full tile"),
+ (16384, "prefill chunk 16384 (max_num_tokens)"),
+ )
+ )
+ + (Cell(why="ragged tail, 16320 + 70 rows (mixed step)", spec=dict(rows=16390)),)
+ # The tp4 / tep4 / dep4 target: the mixer runs on the full 4 x 2560
+ # residual on every rank, so only the row count moves per member.
+ + tuple(
+ Cell(why=why, spec=dict(rows=t))
+ for t, why in (
+ (1, "1 row: one decode token, the smallest captured batch; 1 of a tile's 128 rows"),
+ (32, "32 rows: dep4's per-rank decode cap (128 / 4), a quarter of one tile"),
+ (1024, "1024 rows: mtp7 verify at tep4, 128 requests x 8 rows"),
+ (8192, "8192 rows: one prefill chunk at max_num_tokens, the envelope's largest"),
+ )
+ )
+ )
+
+ note = """
+ Source: the CuTe DSL kernel in this module (no trtllm op or FlashInfer call
+ fuses this multi-input epilogue). It is JIT-compiled by `cute.compile` on
+ the first call per `(block_h, C, H, R)` for the running device (sm_103 in
+ the cells, no `CUTE_DSL_ARCH` override); the row count is a runtime
+ argument. Metadata consumed: none. Call sites: the hyper-connection mixers
+ at decode, MTP verify (4 rows per request) and prefill. Under tp4, tep4 and
+ dep4 every rank mixes the full 4 x 2560 residual, so only the row count
+ moves per member. Preconditions (bf16): `low` `[T, R]` contiguous, `up` the
+ transposed view of a contiguous `[C*H, R]` weight, `normed` `[T, C*H]`
+ contiguous, 16-byte aligned pointers, `R % 64 == 0`, `H % 32 == 0`,
+ `C * 32 <= 256`. Graph-capturable (launches on the current stream, no host
+ sync). Launched with programmatic dependent launch (TRTLLM_ENABLE_PDL=0
+ turns it off): each CTA's first tile streams its up-weight stages before
+ `griddepcontrol.wait` and reads `low` / `normed` after it.
+
+ Certified: bf16 at the checkpoint's 4 x 2560, rank 320, at 1 and 32 rows
+ (part of one tile), 128 (one full tile), 1024, 8192, 16384 and 16390 rows
+ (a partial last tile); graph replay.
+
+ Measured on GB300 against the float64 reference, with the band
+ `rtol = eps`, `atol = eps * max|ref|` (`eps` the bf16 epsilon): worst
+ element at most 0.40 of the band, relative L2 <= 8.1e-5, cosine
+ >= 0.999999997. The band does not separate a bf16-rounded gate logit
+ (0.28x to 0.46x), which is a legitimate two-launch result.
+
+ Not certified:
+ - rows that no cell drives: zero, 2-31, 33-127 (the same one-tile path
+ as 1 and 32), between the full-tile cells, and above 16390;
+ - hc_count other than 4, hidden other than 2560, rank other than 320 (the
+ kernel compiles for them);
+ - `hc_up_gate_mean_fp4`, which this module also defines: no cell here
+ drives it.
+ """
+
+ def __call__(
+ 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/ops/hc/qwen4_exp_hc_up_gate_mean_rstd.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_up_gate_mean_rstd.py
new file mode 100644
index 000000000000..a42b236a8f13
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_up_gate_mean_rstd.py
@@ -0,0 +1,1126 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Hyper-connection mixer second half on the un-normalized residual: the up
+projection as a 2-CTA tcgen05 GEMM whose epilogue gates `hc_out * rstd * (1 + w)`
+-- the site's normed stream, formed in registers -- and averages the streams
+(optionally also writing the routed experts' NVFP4 input).
+
+Tier 6 (CuTe DSL written for this target). `qwen4_exp_hc_up_gate_mean`'s epilogue
+on a CTA pair: `tcgen05.mma.cta_group::2` with M = 256 (each CTA 128 rows of
+`low`, resident in shared memory for its row tile) and N = 128 (C streams x 32
+hidden columns), each CTA loading half of the up weight's N per k block. Tiles
+are walked row-tile-major over a contiguous range per pair (one `low` load per
+row tile); the MMA runs up to four tiles ahead of the epilogue (four accumulator
+stages).
+
+The epilogue is two sets of eight warps that take alternate tiles. A load warp
+brings each (tile, stream)'s `hc_out` box (TMA) and `(1 + w)` slice (bulk copy) to
+its own shared-memory stage; a set hands a stream's stage back as soon as it has
+consumed it, so the next tile's boxes load under the current tile's math. Within a
+set, warp w and w + 4 read the same 128 accumulator rows and split the tile's 32
+columns; the gate and the stream mean run on packed fp32 pairs (`f32x2`), and each
+thread stores its 16 outputs as one 32-byte sector.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+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._experimental.modeling_v2.catalog.quantization.fp4_quantize import (
+ fp4_quantize,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_inject_norm import (
+ _bf16x2,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_up_gate_mean import (
+ _asm,
+ _rcp_approx_ftz,
+ _st_global_b16,
+ _st_global_v4,
+)
+from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+from tensorrt_llm._torch.kernels.blackwell.custom_pipeline import PipelineTmaUmma, PipelineUmmaAsync
+from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
+
+from .qwen4_exp_hc_first_rescale_tc import _lds_v4, _sts_v4
+
+#: Programmatic dependent launch (sm_90+), off with TRTLLM_ENABLE_PDL=0.
+_PDL = get_env_enable_pdl()
+
+#: Hidden columns a tile (MMA N = C * 32 = 128) and a thread's epilogue columns.
+_BLOCK_H = 32
+_SUB = 16
+#: Epilogue warp sets (alternate tiles) and warps per set.
+_SETS = 2
+_SET_WARPS = 8
+#: Up-weight ring, accumulator stages, and the hc_out / (1 + w) ring (one stage per
+#: (tile, stream)).
+_B_STAGES = 8
+_ACC_STAGES = 4
+_N_STAGES = 8
+#: NVFP4 block: one e4m3 scale per this many contiguous elements.
+SF_VEC = 16
+
+
+@dsl_user_op
+def _ldg_v4_f32(addr, *, loc=None, ip=None):
+ """Four fp32 values at a 16-byte aligned global address."""
+ ty = T.f32()
+ res = _asm(
+ llvm.StructType.get_literal([ty] * 4),
+ [cutlass.Int64(addr)],
+ "ld.global.v4.f32 {$0, $1, $2, $3}, [$4];",
+ "=f,=f,=f,=f,l",
+ loc,
+ ip,
+ )
+ return [cutlass.Float32(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(4)]
+
+
+@dsl_user_op
+def _lds_v4_f32(addr, *, loc=None, ip=None):
+ """Four fp32 values at a 16-byte aligned shared address (side-effecting: ordered
+ after the stage wait).
+ """
+ ty = T.f32()
+ res = _asm(
+ llvm.StructType.get_literal([ty] * 4),
+ [cutlass.Int32(addr)],
+ "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];",
+ "=f,=f,=f,=f,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+ return [cutlass.Float32(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(4)]
+
+
+@dsl_user_op
+def _bulk_g2s(dst, src, nbytes: int, mbar, *, loc=None, ip=None):
+ """One non-tensor bulk copy global -> this CTA's shared memory, completing on
+ `mbar`.
+ """
+ _asm(
+ T.i32(),
+ [cutlass.Int32(dst), cutlass.Int64(src), cutlass.Int32(mbar)],
+ "{ cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [$1], [$2], "
+ + str(nbytes)
+ + ", [$3]; mov.u32 $0, 0; }",
+ "=r,r,l,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+
+
+@dsl_user_op
+def _st_global_v8(addr, w, *, loc=None, ip=None):
+ """One 32-byte `st.global.v8.b32` (sm_100+): a thread's 16 bf16 outputs, one
+ full sector.
+ """
+ _asm(
+ T.i32(),
+ [cutlass.Int64(addr)] + [cutlass.Int32(x) for x in w],
+ "{ st.global.v8.b32 [$1], {$2, $3, $4, $5, $6, $7, $8, $9}; mov.u32 $0, 0; }",
+ "=r,l" + ",r" * 8,
+ loc,
+ ip,
+ side_effects=True,
+ )
+
+
+@dsl_user_op
+def _gate2(z0, z1, xw, wa, wb, hr, a0, a1, first: bool, *, loc=None, ip=None):
+ """One column pair of one stream on the packed fp32 pipe (`f32x2`: each lane is
+ the scalar `.rn` operation): `t = tanh.approx(z * 0.5)`, `g = fma(t, hr, hr)`,
+ `xw = x * w1`, `acc = fma(g, xw, acc)` (`g * xw` for the first stream). `xw` is
+ the stream's bf16x2 word (element 2j in the low half), widened exactly by bit
+ shifts.
+ """
+ tail = "mul.rn.f32x2 a, t, x; " if first else "mov.b64 a, {$8, $9}; fma.rn.f32x2 a, t, x, a; "
+ text = (
+ "{ .reg .b64 z, h, t, x, w, a; .reg .f32 c, t0, t1; .reg .b32 xl, xh; "
+ "mov.f32 c, 0f3F000000; mov.b64 h, {c, c}; "
+ "mov.b64 z, {$2, $3}; mul.rn.f32x2 z, z, h; "
+ "mov.b64 {t0, t1}, z; tanh.approx.f32 t0, t0; tanh.approx.f32 t1, t1; "
+ "mov.b64 t, {t0, t1}; mov.b64 h, {$7, $7}; fma.rn.f32x2 t, t, h, h; "
+ "shl.b32 xl, $4, 16; and.b32 xh, $4, 0xFFFF0000; mov.b64 x, {xl, xh}; "
+ "mov.b64 w, {$5, $6}; mul.rn.f32x2 x, x, w; " + tail + "mov.b64 {$0, $1}, a; }"
+ )
+ ty = T.f32()
+ res = _asm(
+ llvm.StructType.get_literal([ty, ty]),
+ [
+ cutlass.Float32(z0),
+ cutlass.Float32(z1),
+ cutlass.Int32(xw),
+ cutlass.Float32(wa),
+ cutlass.Float32(wb),
+ cutlass.Float32(hr),
+ cutlass.Float32(a0),
+ cutlass.Float32(a1),
+ ],
+ text,
+ "=f,=f,f,f,r,f,f,f,f,f",
+ loc,
+ ip,
+ )
+ return [cutlass.Float32(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(2)]
+
+
+@dsl_user_op
+def _nvfp4_16(words, gs, rcp_gs, rcp6, *, loc=None, ip=None):
+ """NVFP4 of one 16-element block given as 8 bf16x2 words (element 2k in word k's
+ low half): `(e2m1 word for elements 0..7, e2m1 word for 8..15, e4m3 scale byte)`.
+ `qwen4_exp_hc_up_gate_mean`'s recipe (`_absmax16`, `_e4m3_scale`,
+ `_e2m1_out_scale`, `_e2m1x8`) on the same bf16-rounded values, with the max taken
+ on the packed bf16 pairs, the bf16 -> fp32 unpack done by bit shifts (both exact)
+ and the scaling multiplies on `f32x2` pairs (each lane `mul.rn.f32`): bitwise its
+ outputs.
+ """
+ amax = "abs.bf16x2 m, $3; " + "".join(
+ f"abs.bf16x2 t, ${4 + k}; max.bf16x2 m, m, t; " for k in range(7)
+ )
+ unpack = "".join(
+ f"shl.b32 a{2 * k}, ${3 + k}, 16; and.b32 a{2 * k + 1}, ${3 + k}, 0xFFFF0000; "
+ for k in range(8)
+ )
+ muls = "mov.b64 oo, {o, o}; " + "".join(
+ f"mov.b64 pp, {{a{2 * k}, a{2 * k + 1}}}; mul.rn.f32x2 pp, pp, oo; "
+ f"mov.b64 {{a{2 * k}, a{2 * k + 1}}}, pp; "
+ for k in range(8)
+ )
+ cvts = "".join(f"cvt.rn.satfinite.e2m1x2.f32 b{k}, a{2 * k + 1}, a{2 * k}; " for k in range(8))
+ text = (
+ "{ .reg .b32 m, t, f2; .reg .b32 a<16>; .reg .f32 x, y, mx, o; .reg .b16 h, l16, h16; "
+ ".reg .pred p; .reg .b8 b<8>; .reg .b64 oo, pp; "
+ + amax
+ + "mov.b32 {l16, h16}, m; cvt.f32.bf16 x, l16; cvt.f32.bf16 y, h16; max.f32 mx, x, y; "
+ "mul.rn.f32 x, mx, $13; mul.rn.f32 x, $11, x; "
+ "cvt.rn.satfinite.e4m3x2.f32 h, 0f00000000, x; cvt.u32.u16 $2, h; "
+ "cvt.rn.f16x2.e4m3x2 f2, h; mov.b32 {l16, h16}, f2; cvt.f32.f16 y, l16; "
+ "mul.rn.f32 y, y, $12; rcp.approx.ftz.f32 y, y; setp.neu.f32 p, mx, 0f00000000; "
+ "selp.f32 o, y, 0f00000000, p; "
+ + unpack
+ + muls
+ + cvts
+ + "mov.b32 $0, {b0, b1, b2, b3}; mov.b32 $1, {b4, b5, b6, b7}; }"
+ )
+ ty = T.i32()
+ res = _asm(
+ llvm.StructType.get_literal([ty, ty, ty]),
+ [cutlass.Int32(w) for w in words]
+ + [cutlass.Float32(gs), cutlass.Float32(rcp_gs), cutlass.Float32(rcp6)],
+ text,
+ "=r,=r,=r" + ",r" * 8 + ",f,f,f",
+ loc,
+ ip,
+ )
+ return [cutlass.Int32(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(3)]
+
+
+class HcUpGateMeanRstdKernel:
+ """`mixed[t, i] = 1/C sum_s sigmoid((low @ up)[t, s*H + i]) * hc_out[t, s*H + i]
+ * rstd[t, s] * w1[s*H + i]` (+ its NVFP4 with `fp4`).
+ """
+
+ def __init__(self, hc: int, hidden: int, rank: int, fp4: bool):
+ assert rank % 64 == 0 and hidden % _BLOCK_H == 0 and hc == 4
+ self.bh = _BLOCK_H
+ self.sub = _SUB
+ self.hc, self.hidden, self.rank, self.fp4 = hc, hidden, rank, fp4
+ self.kb = rank // 64
+ self.acc_dtype = cutlass.Float32
+ self.io_dtype = cutlass.BFloat16
+ self.mma_tiler = (256, hc * self.bh, 64)
+ self.cta_group = tcgen05.CtaGroup.TWO
+ self.cluster_shape_mn = (2, 1)
+ self.epi_sets = _SETS
+ self.set_warps = _SET_WARPS
+ self.epi_warps = _SETS * _SET_WARPS
+ self.mma_warp_id = self.epi_warps
+ self.tma_warp_id = self.epi_warps + 1
+ self.nrm_warp_id = self.epi_warps + 2
+ self.threads_per_cta = 32 * (self.epi_warps + 3)
+ self.num_b_stage = _B_STAGES
+ self.num_n_stage = _N_STAGES
+ self.num_acc_stage = _ACC_STAGES
+ # The pair's j-th tile uses accumulator stage j % ACC and, for stream s, the
+ # hc_out stage (C * j + s) % NS.
+ assert _ACC_STAGES & (_ACC_STAGES - 1) == 0 and _SETS & (_SETS - 1) == 0
+ assert _N_STAGES >= hc * _SETS
+ self.acc_shift = _ACC_STAGES.bit_length() - 1
+ self.set_shift = _SETS.bit_length() - 1
+ self.epilog_sync_bar_id = 1
+ self.tmem_alloc_sync_bar_id = 2
+ # Named barriers 3..3 + 4 * SETS - 1: a set's epilogue warps w and w + 4 (the
+ # NVFP4 hand-off).
+ self.pair_bar_id0 = 3
+
+ def _setup(self):
+ k_major = utils.LayoutEnum.ROW_MAJOR.mma_major_mode()
+ tiled_mma = utils.sm100.make_trivial_tiled_mma(
+ self.io_dtype,
+ self.io_dtype,
+ k_major,
+ k_major,
+ self.acc_dtype,
+ self.cta_group,
+ self.mma_tiler[:2],
+ )
+ self.cluster_layout_vmnk = cute.tiled_divide(
+ cute.make_layout((*self.cluster_shape_mn, 1)), (tiled_mma.thr_id.shape,)
+ )
+ self.n_box = (128, self.bh)
+ # A stage: one stream of this CTA's tile, its hc_out box ([128, 32] bf16) and
+ # (1 + w) slice ([32] fp32).
+ self.num_n_bytes = 128 * self.bh * 2 + self.bh * 4
+ self.a_smem_layout_staged = utils.sm100.make_smem_layout_a(
+ tiled_mma, self.mma_tiler, self.io_dtype, self.kb
+ )
+ self.b_smem_layout_staged = utils.sm100.make_smem_layout_b(
+ tiled_mma, self.mma_tiler, self.io_dtype, self.num_b_stage
+ )
+ 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
+ )
+ 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,
+ hc_ptr: cute.Pointer,
+ mixed: cute.Tensor,
+ xq_ptr: cute.Pointer,
+ sf_ptr: cute.Pointer,
+ gs_ptr: cute.Pointer,
+ rstd_ptr: cute.Pointer,
+ w1_ptr: cute.Pointer,
+ max_clusters: cutlass.Constexpr,
+ stream: cuda.CUstream,
+ ):
+ bh, hc, hid, rank = self.bh, self.hc, self.hidden, self.rank
+ nt = hid // bh
+ width = hc * hid
+ a = cute.make_tensor(low_ptr, cute.make_layout((m, rank, 1), stride=(rank, 1, 0)))
+ # The `[C*H, R]` up weight viewed as N = (32 hidden, C streams, H/32 slices): an
+ # MMA tile's N is one hidden slice of every stream.
+ b = cute.make_tensor(
+ up_ptr,
+ cute.make_layout(((bh, hc, nt), rank, 1), stride=((rank, hid * rank, bh * rank), 1, 0)),
+ )
+ nrm = cute.make_tensor(hc_ptr, cute.make_layout((m, width), stride=(width, 1)))
+ gsc = cute.make_tensor(gs_ptr, cute.make_layout(1))
+ tiled_mma = self._setup()
+ 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,
+ )
+ # Both CTAs' A (all k blocks) / B stage land on the leader's barrier.
+ self.a_bytes = cute.size_in_bytes(self.io_dtype, a_smem) * self.kb * 2
+ self.b_bytes = cute.size_in_bytes(self.io_dtype, b_smem) * 2
+ total = ((m + 255) // 256) * nt
+ ncl = cutlass.min(total, max_clusters)
+ self.kernel(
+ tiled_mma,
+ tma_atom_a,
+ tma_a,
+ tma_atom_b,
+ tma_b,
+ tma_atom_n,
+ tma_n,
+ gsc,
+ mixed.iterator.toint(),
+ xq_ptr.toint(),
+ sf_ptr.toint(),
+ rstd_ptr.toint(),
+ w1_ptr.toint(),
+ m,
+ total,
+ self.cluster_layout_vmnk,
+ self.a_smem_layout_staged,
+ self.b_smem_layout_staged,
+ self.n_smem_layout_staged,
+ ).launch(
+ grid=[2 * ncl, 1, 1],
+ 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,
+ mGS: cute.Tensor,
+ mixed_base: cutlass.Int64,
+ q_base: cutlass.Int64,
+ sf_base: cutlass.Int64,
+ rstd_base: cutlass.Int64,
+ w1_base: cutlass.Int64,
+ m: cutlass.Int32,
+ total: 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,
+ ):
+ hc, hid, bh, sub = self.hc, self.hidden, self.bh, self.sub
+ nt = hid // bh
+ 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()
+ bid, _, _ = cute.arch.block_idx()
+ nblk, _, _ = cute.arch.grid_dim()
+ v = bid % 2
+ is_leader = v == 0
+ cl = bid // 2
+ ncl = nblk // 2
+ cta_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+ cluster_coord = cluster_layout_vmnk.get_flat_coord(cta_rank)
+ # The pair's contiguous range [t0, t1) of (256-row tile, hidden slice) tiles,
+ # row-tile-major.
+ t0 = (cl * total) // ncl
+ t1 = ((cl + 1) * total) // ncl
+ m_first = t0 // nt
+ m_last = (t1 - 1) // nt
+
+ @cute.struct
+ class SharedStorage:
+ a_mbar: cute.struct.MemRange[cutlass.Int64, 2]
+ b_mbar: cute.struct.MemRange[cutlass.Int64, self.num_b_stage * 2]
+ acc_mbar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]
+ n_mbar: 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)
+ a_producer, a_consumer = PipelineTmaUmma.create(
+ barrier_storage=storage.a_mbar.data_ptr(),
+ num_stages=1,
+ producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+ consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1),
+ tx_count=self.a_bytes,
+ cta_layout_vmnk=cluster_layout_vmnk,
+ ).make_participants()
+ b_producer, b_consumer = PipelineTmaUmma.create(
+ barrier_storage=storage.b_mbar.data_ptr(),
+ num_stages=self.num_b_stage,
+ producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+ consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1),
+ tx_count=self.b_bytes,
+ cta_layout_vmnk=cluster_layout_vmnk,
+ ).make_participants()
+ # An accumulator stage is consumed by one set's warps in each CTA.
+ acc_pipeline = PipelineUmmaAsync.create(
+ barrier_storage=storage.acc_mbar.data_ptr(),
+ num_stages=self.num_acc_stage,
+ producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+ consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 2 * self.set_warps),
+ cta_layout_vmnk=cluster_layout_vmnk,
+ )
+ # hc_out / (1 + w) stages are per CTA (each CTA its own rows), one set's warps
+ # consuming each.
+ n_pipe = pipeline.PipelineTmaAsync.create(
+ barrier_storage=storage.n_mbar.data_ptr(),
+ num_stages=self.num_n_stage,
+ producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+ consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, self.set_warps),
+ tx_count=self.num_n_bytes,
+ cta_layout_vmnk=cute.make_layout((1, 1, 1, 1)),
+ )
+ n_producer = n_pipe.make_participants()[0]
+ tmem_alloc_barrier = pipeline.NamedBarrier(
+ barrier_id=self.tmem_alloc_sync_bar_id, num_threads=32 * (1 + self.epi_warps)
+ )
+ tmem = utils.TmemAllocator(
+ storage.tmem_holding_buf.ptr,
+ barrier_for_retrieve=tmem_alloc_barrier,
+ allocator_warp_id=0,
+ is_two_cta=True,
+ 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=1024,
+ swizzle=a_smem_layout_staged.inner,
+ )
+ sB = smem.allocate_tensor(
+ element_type=self.io_dtype,
+ layout=b_smem_layout_staged.outer,
+ byte_alignment=1024,
+ swizzle=b_smem_layout_staged.inner,
+ )
+ sN = smem.allocate_tensor(
+ element_type=self.io_dtype,
+ layout=n_smem_layout_staged.outer,
+ byte_alignment=1024,
+ swizzle=n_smem_layout_staged.inner,
+ )
+ # NVFP4: a set's upper warp group hands its half of a row's e2m1 / scale bytes to
+ # the lower group through shared memory (two tile slots per set), which stores
+ # the whole row piece (one 16-byte e2m1 store, one 2-byte scale store).
+ sX = smem.allocate_tensor(
+ element_type=cutlass.Int32,
+ layout=cute.make_layout((2 * self.epi_sets, 128, 4), stride=(512, 4, 1)),
+ byte_alignment=128,
+ )
+ sW = smem.allocate_tensor(
+ element_type=cutlass.Float32,
+ layout=cute.make_layout((bh, self.num_n_stage), stride=(1, bh)),
+ byte_alignment=128,
+ )
+
+ 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)
+ )
+ gN_box = cute.local_tile(mN_mn, self.n_box, (None, None))
+ thr_mma = tiled_mma.get_slice(v)
+ tCgA = thr_mma.partition_A(gA_mkl)
+ tCgB = thr_mma.partition_B(gB_nkl)
+ a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape)
+ tAsA, tAgA = cpasync.tma_partition(
+ tma_atom_a,
+ cluster_coord[2],
+ a_cta_layout,
+ cute.group_modes(sA, 0, 3),
+ cute.group_modes(tCgA, 0, 3),
+ )
+ b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape)
+ tBsB, tBgB = cpasync.tma_partition(
+ tma_atom_b,
+ cluster_coord[1],
+ b_cta_layout,
+ 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),
+ )
+ a_mcast = cpasync.create_tma_multicast_mask(
+ cluster_layout_vmnk, cluster_coord, mcast_mode=2
+ )
+ b_mcast = cpasync.create_tma_multicast_mask(
+ cluster_layout_vmnk, cluster_coord, mcast_mode=1
+ )
+ 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)
+
+ # ---------------------------------------------------------- TMA (A, B): both CTAs
+ if warp_idx == self.tma_warp_id:
+ cute.arch.griddepcontrol_wait()
+ for mi in cutlass.range(m_first, m_last + 1, 1, unroll=1):
+ a_handle = a_producer.acquire_and_advance()
+ for kb in cutlass.range_constexpr(self.kb):
+ cute.copy(
+ tma_atom_a,
+ tAgA[(None, mi, kb, 0)],
+ tAsA[(None, kb)],
+ tma_bar_ptr=a_handle.barrier,
+ mcast_mask=a_mcast,
+ )
+ n_lo = cutlass.max(t0 - mi * nt, 0)
+ n_hi = cutlass.min(t1 - mi * nt, nt)
+ for n in cutlass.range(n_lo, n_hi, 1, unroll=1):
+ for kb in cutlass.range_constexpr(self.kb):
+ handle = b_producer.acquire_and_advance()
+ cute.copy(
+ tma_atom_b,
+ tBgB[(None, n, kb, 0)],
+ tBsB[(None, handle.index)],
+ tma_bar_ptr=handle.barrier,
+ mcast_mask=b_mcast,
+ )
+ a_producer.tail()
+ b_producer.tail()
+
+ # ---------------------------------------------------------- hc_out boxes, (1 + w)
+ if warp_idx == self.nrm_warp_id:
+ cute.arch.griddepcontrol_wait()
+ sw_base = cutlass.Int32(sW.iterator.toint())
+ for mi in cutlass.range(m_first, m_last + 1, 1, unroll=1):
+ n_lo = cutlass.max(t0 - mi * nt, 0)
+ n_hi = cutlass.min(t1 - mi * nt, nt)
+ for n in cutlass.range(n_lo, n_hi, 1, unroll=1):
+ for s in cutlass.range_constexpr(hc):
+ handle = n_producer.acquire_and_advance()
+ with cute.arch.elect_one():
+ _bulk_g2s(
+ sw_base + handle.index * (bh * 4),
+ w1_base + cutlass.Int64(s * hid + n * bh) * 4,
+ bh * 4,
+ cutlass.Int32(handle.barrier.toint()),
+ )
+ cute.copy(
+ tma_atom_n,
+ tNgN[(None, 2 * mi + v, s * nt + n)],
+ tNsN[(None, handle.index)],
+ tma_bar_ptr=handle.barrier,
+ )
+ n_producer.tail()
+
+ # ---------------------------------------------------------- MMA (the leader issues)
+ 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_state = pipeline.make_pipeline_state(
+ pipeline.PipelineUserType.Producer, self.num_acc_stage
+ )
+ # Both CTAs walk the tiles (the accumulator state advances in both); the
+ # leader issues.
+ for mi in cutlass.range(m_first, m_last + 1, 1, unroll=1):
+ n_lo = cutlass.max(t0 - mi * nt, 0)
+ n_hi = cutlass.min(t1 - mi * nt, nt)
+ if is_leader:
+ a_handle = a_consumer.wait_and_advance()
+ for n in cutlass.range(n_lo, n_hi, 1, unroll=1):
+ tCtAcc = tCtAcc_base[(None, None, None, acc_state.index)]
+ acc_pipeline.producer_acquire(acc_state)
+ tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+ for kb in cutlass.range_constexpr(self.kb):
+ handle = b_consumer.wait_and_advance()
+ num_kblocks = cute.size(tCrA, mode=[2])
+ for kk in cutlass.range(num_kblocks, unroll_full=True):
+ cute.gemm(
+ tiled_mma,
+ tCtAcc,
+ tCrA[(None, None, kk, kb)],
+ tCrB[(None, None, kk, handle.index)],
+ tCtAcc,
+ )
+ tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+ handle.release()
+ acc_pipeline.producer_commit(acc_state)
+ acc_state.advance()
+ a_handle.release()
+ else:
+ for n in cutlass.range(n_lo, n_hi, 1, unroll=1):
+ acc_state.advance()
+ acc_pipeline.producer_tail(acc_state)
+
+ # ---------------------------------------------------------- epilogue (each CTA its rows)
+ if warp_idx < self.epi_warps:
+ cute.arch.griddepcontrol_wait()
+ 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)
+ # Set `eset` takes the pair's tiles t0 + eset, t0 + eset + SETS, ...; thread
+ # (grp, lt) of a set owns row lt of the CTA's 128 and columns
+ # [grp * 16, grp * 16 + 16) of the tile's 32.
+ eset = warp_idx // self.set_warps
+ lt = tidx % 128
+ grp = (tidx // 128) % 2
+ epi_tile = (128, sub)
+ copy_atom_t2r = utils.sm100.get_tmem_load_op(
+ (128, self.mma_tiler[1], self.mma_tiler[2]),
+ utils.LayoutEnum.ROW_MAJOR,
+ self.io_dtype,
+ self.acc_dtype,
+ epi_tile,
+ True,
+ )
+ 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(lt)
+ tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi)
+ tTR_sN = thr_copy_t2r.partition_D(cute.flat_divide(sN, epi_tile))
+ sw_addr = cutlass.Int32(sW.iterator.toint())
+ sx_addr = cutlass.Int32(sX.iterator.toint()) + eset * (2 * 128 * 16)
+ frag_shape = tTR_sN[(None, None, None, 0, 0, 0)].shape
+ rAcc = cute.make_rmem_tensor(frag_shape, self.acc_dtype)
+ rN = cute.make_rmem_tensor(frag_shape, self.io_dtype)
+ # sigmoid(z) * r / C = fma(tanh(z / 2), r / 2C, r / 2C): the stream mean and the
+ # norm's rstd folded into the gate.
+ scale = cutlass.Float32(0.5 / hc)
+ if cutlass.const_expr(self.fp4):
+ gs = mGS[0]
+ rcp_gs = _rcp_approx_ftz(gs)
+ rcp6 = _rcp_approx_ftz(cutlass.Float32(6.0))
+
+ # The rows' rstd, reloaded when the row tile changes.
+ mi_cur = cutlass.Int32(-1)
+ rs0 = cutlass.Float32(0.0)
+ rs1 = cutlass.Float32(0.0)
+ rs2 = cutlass.Float32(0.0)
+ rs3 = cutlass.Float32(0.0)
+ for t in cutlass.range(t0 + eset, t1, self.epi_sets, unroll=1):
+ j = t - t0
+ acc_state = pipeline.PipelineState(
+ self.num_acc_stage,
+ j,
+ j & (self.num_acc_stage - 1),
+ (j >> self.acc_shift) & 1,
+ )
+ mi = t // nt
+ n = t - mi * nt
+ row = mi * 256 + v * 128 + lt
+ row_ok = row < m
+ if mi != mi_cur:
+ rs0, rs1, rs2, rs3 = _ldg_v4_f32(
+ rstd_base + cutlass.Int64(cutlass.min(row, m - 1)) * (hc * 4)
+ )
+ mi_cur = mi
+ rs = [rs0, rs1, rs2, rs3]
+ col0 = n * bh + grp * sub
+ acc_pipeline.consumer_wait(acc_state)
+ acc = [cutlass.Float32(0.0)] * sub
+ for s in cutlass.range_constexpr(hc):
+ cute.copy(
+ tiled_copy_t2r,
+ tTR_tAcc[(None, None, None, 0, 2 * s + grp, acc_state.index)],
+ rAcc,
+ )
+ # tcgen05.wait::ld: the stream's accumulator columns are in registers;
+ # after the last stream the stage goes back before the math.
+ cute.arch.fence_view_async_tmem_load()
+ if cutlass.const_expr(s == hc - 1):
+ cute.arch.sync_warp()
+ with cute.arch.elect_one():
+ acc_pipeline.consumer_release(acc_state)
+ hrs = rs[s] * scale
+ q = j * hc + s
+ if cutlass.const_expr(self.num_n_stage == hc * self.epi_sets):
+ # The set's stream-s stage is fixed; its phase flips every own tile.
+ n_state = pipeline.PipelineState(
+ self.num_n_stage, q, eset * hc + s, (j >> self.set_shift) & 1
+ )
+ else:
+ n_state = pipeline.PipelineState(
+ self.num_n_stage, q, q % self.num_n_stage, (q // self.num_n_stage) & 1
+ )
+ n_pipe.consumer_wait(n_state)
+ cute.autovec_copy(tTR_sN[(None, None, None, 0, grp, n_state.index)], rN)
+ waddr = sw_addr + (n_state.index * bh + grp * sub) * 4
+ w1 = []
+ for c4 in cutlass.range_constexpr(sub // 4):
+ w1 += _lds_v4_f32(waddr + c4 * 16)
+ xw = cute.recast_tensor(rN, cutlass.Int32)
+ for k in cutlass.range_constexpr(sub // 2):
+ acc[2 * k], acc[2 * k + 1] = _gate2(
+ rAcc[2 * k],
+ rAcc[2 * k + 1],
+ xw[k],
+ w1[2 * k],
+ w1[2 * k + 1],
+ hrs,
+ acc[2 * k],
+ acc[2 * k + 1],
+ s == 0,
+ )
+ # Every read of the stage (hc_out, (1 + w)) is performed before it is
+ # handed back.
+ cute.arch.fence_acq_rel_cta()
+ cute.arch.sync_warp()
+ n_pipe.consumer_release(n_state) # lane 0 of each warp arrives
+ words = [_bf16x2(acc[2 * k], acc[2 * k + 1]) for k in range(sub // 2)]
+ if row_ok:
+ _st_global_v8(mixed_base + (cutlass.Int64(row) * hid + col0) * 2, words)
+ if cutlass.const_expr(self.fp4):
+ q0, q1, sf = _nvfp4_16(words, gs, rcp_gs, rcp6)
+ slot = sx_addr + (((t // self.epi_sets) % 2) * 128 + lt) * 16
+ if grp == 1:
+ _sts_v4(slot, q0, q1, sf, sf)
+ cute.arch.barrier(
+ barrier_id=self.pair_bar_id0 + warp_idx % 4 + 4 * eset,
+ number_of_threads=64,
+ )
+ if grp == 0:
+ pq = _lds_v4(slot)
+ if row_ok:
+ row64 = cutlass.Int64(row)
+ _st_global_v4(
+ q_base + row64 * (hid // 2) + cutlass.Int64(n * bh // 2),
+ q0,
+ q1,
+ pq[0],
+ pq[1],
+ )
+ _st_global_b16(
+ sf_base + row64 * (hid // 16) + cutlass.Int64(n * bh // 16),
+ sf | (pq[2] << 8),
+ )
+
+ cute.arch.barrier(
+ barrier_id=self.epilog_sync_bar_id, number_of_threads=32 * self.epi_warps
+ )
+ tmem.relinquish_alloc_permit()
+ tmem.free(tmem_ptr)
+
+
+_COMPILED: dict = {}
+
+
+def _launcher(hc: int, hidden: int, rank: int, fp4: bool, ptrs):
+ """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, hc_out, mixed tensor, xq, sf, gs, rstd, norm_w1)`.
+ """
+ key = (hc, hidden, rank, fp4)
+ fn = _COMPILED.get(key)
+ if fn is None:
+ kern = HcUpGateMeanRstdKernel(hc, hidden, rank, fp4)
+ gm = cute.AddressSpace.gmem
+ cptrs = [
+ make_ptr(cutlass.BFloat16, ptrs[0], gm, assumed_align=16),
+ make_ptr(cutlass.BFloat16, ptrs[1], gm, assumed_align=16),
+ make_ptr(cutlass.BFloat16, ptrs[2], gm, assumed_align=16),
+ from_dlpack(ptrs[3], assumed_align=16).mark_layout_dynamic(leading_dim=1),
+ make_ptr(cutlass.Int32, ptrs[4], gm, assumed_align=16),
+ make_ptr(cutlass.Uint8, ptrs[5], gm, assumed_align=2),
+ make_ptr(cutlass.Float32, ptrs[6], gm, assumed_align=4),
+ make_ptr(cutlass.Float32, ptrs[7], gm, assumed_align=16),
+ make_ptr(cutlass.Float32, ptrs[8], gm, assumed_align=16),
+ ]
+ fn = cute.compile(
+ kern,
+ 128,
+ *cptrs,
+ max_clusters=utils.HardwareInfo().get_max_active_clusters(2),
+ 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_rstd(
+ low: torch.Tensor,
+ up: torch.Tensor,
+ hc_out: torch.Tensor,
+ rstd: torch.Tensor,
+ norm_w1: torch.Tensor,
+ hc_count: int,
+ global_scale: Optional[torch.Tensor] = None,
+):
+ """`mixed [T, H]`, and with `global_scale` also `xq [T, H/2]` and `xsf [T*H/16]`
+ (linear scales).
+ """
+ rows, rank = low.shape
+ hidden = hc_out.shape[1] // hc_count
+ dev = hc_out.device
+ fp4 = global_scale is not None
+ mixed = torch.empty((rows, hidden), dtype=hc_out.dtype, device=dev)
+ xq = xsf = None
+ if fp4:
+ xq = torch.empty((rows, hidden // 2), dtype=torch.uint8, device=dev)
+ xsf = torch.empty((rows * (hidden // SF_VEC),), dtype=torch.uint8, device=dev)
+ if rows > 0:
+ # The NVFP4 pointers are never read without `fp4`.
+ qp = (
+ (xq.data_ptr(), xsf.data_ptr(), global_scale.data_ptr())
+ if fp4
+ else (mixed.data_ptr(),) * 3
+ )
+ ptrs = (
+ low.data_ptr(),
+ up.data_ptr(),
+ hc_out.data_ptr(),
+ mixed,
+ *qp,
+ rstd.data_ptr(),
+ norm_w1.data_ptr(),
+ )
+ _launcher(hc_count, hidden, rank, fp4, ptrs)(rows, *ptrs)
+ return (mixed, xq, xsf) if fp4 else mixed
+
+
+def _mixed_reference(low, up, hc_out, rstd, norm_w1, hc_count) -> torch.Tensor:
+ """float64: the gate times the unrounded normed stream, stream mean, one rounding."""
+ hidden = hc_out.shape[1] // hc_count
+ gate = torch.sigmoid(F.linear(low.double(), up.double().t())).unflatten(-1, (hc_count, hidden))
+ x = hc_out.double().unflatten(-1, (hc_count, hidden)) * rstd.double().unsqueeze(-1)
+ x = x * norm_w1.double().unflatten(-1, (hc_count, hidden))
+ return (gate * x).mean(dim=-2).to(hc_out.dtype)
+
+
+def _check(low, up, hc_out, rstd, norm_w1, hc_count) -> None:
+ rows, rank = low.shape
+ width = hc_out.shape[1]
+ assert low.dtype == up.dtype == hc_out.dtype == torch.bfloat16, "bf16 operands"
+ assert hc_out.shape[0] == rows and up.shape == (rank, width), "shapes"
+ assert low.stride() == (rank, 1) and hc_out.stride() == (width, 1), "contiguous rows"
+ assert up.stride() == (1, rank), "up must be the transposed view of a [C*H, R] weight"
+ assert rstd.dtype == torch.float32 and rstd.shape == (rows, hc_count), "rstd [T, C] fp32"
+ assert rstd.is_contiguous(), "rstd contiguous"
+ assert norm_w1.dtype == torch.float32 and norm_w1.shape == (width,), "norm_w1 [C*H] fp32"
+ assert norm_w1.is_contiguous(), "norm_w1 contiguous"
+ # A row's rstd is one 16-byte load, and the MMA's N is C * 32 = 128.
+ assert hc_count == 4, "hc_count 4"
+ assert width % hc_count == 0 and (width // hc_count) % _BLOCK_H == 0, "H % 32 == 0"
+ assert rank % 64 == 0, "rank must be a multiple of 64"
+ for t in (low, up, hc_out, rstd, norm_w1):
+ assert t.data_ptr() % 16 == 0, "16-byte aligned operands"
+
+
+class Qwen4ExpHcUpGateMeanRstd(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, rstd the site's
+ `[T, C]` fp32 statistic and `norm_w1 = 1 + norm_weight` in fp32 `[C*H]`::
+
+ normed[t, s*H + i] = hc_out[t, s*H + i] * rstd[t, s] * norm_w1[s*H + i]
+ 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;
+ `normed` is fp32 in registers (not rounded to bf16 as a stored normed stream
+ is); gate, product and stream mean are fp32 with one rounding to `hc_out`'s
+ dtype. Per element and stream, in stream order: `t = tanh.approx(z * 0.5)`,
+ `g = fma(t, r, r)` with `r = rstd / 2C`, `xw = hc_out * norm_w1`,
+ `acc = fma(g, xw, acc)` (`g * xw` for stream 0), each a round-to-nearest fp32
+ operation. Returns a fresh `[T, H]` tensor. Fusion boundary: the up projection,
+ the grouped norm's scaling, the sigmoid gate and the stream mean.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: rows. The checkpoint's geometry: hc_count 4, hidden 2560, rank 320, bf16.
+ CELLS: tuple[Cell, ...] = tuple(
+ Cell(why=f"checkpoint 4x2560, rank 320, {why}", spec=dict(rows=t))
+ for t, why in (
+ (16384, "prefill chunk 16384 (max_num_tokens)"),
+ (2048, "prefill 2048 rows"),
+ (16390, "ragged tail, 16320 + 70 rows (mixed step)"),
+ (5, "five rows, less than one tile"),
+ )
+ )
+
+ note = """
+ Tier 6 (CuTe DSL written for this target): a 2-CTA tcgen05 GEMM with a
+ multi-input epilogue (`qwen4_exp_hc_up_gate_mean`'s, plus the norm's
+ scaling); no lower tier has one. Metadata consumed: none. Call sites:
+ `PrefillTarget`'s attention / GDN and stack-final mixers at or above
+ `_HC_RSTD_MIN_ROWS` rows, after `qwen4_exp_hc_first_rescale_tc`, on
+ `qwen4_exp_hc_inject_rstd`'s `hc_out` and rstd; `norm_w1` precomputed once
+ per site at load.
+
+ Preconditions (bf16 activations): `low` `[T, R]` contiguous, `up` the
+ transposed view of a contiguous `[C*H, R]` weight, `hc_out` `[T, C*H]`
+ contiguous, rstd `[T, C]` and `norm_w1` `[C*H]` fp32 contiguous, 16-byte
+ aligned pointers, `hc_count == 4`, `R % 64 == 0`, `H % 32 == 0`. Compiles
+ once per `(C, H, R)` on first call; the row count is a runtime argument.
+ Graph-capturable (current stream, no host sync). Launched with
+ programmatic dependent launch (TRTLLM_ENABLE_PDL=0 turns it off).
+
+ Launch: persistent, one cluster (CTA pair) per two SMs, each walking a
+ contiguous row-tile-major range of (256-row tile, 32-column slice) tiles;
+ 608 threads a CTA (two sets of eight epilogue warps on alternate tiles, a
+ weight-load, an hc_out-load and an MMA warp), about 217 KB of shared memory
+ (low tile, 8-stage weight ring, 8 per-stream hc_out stages) and all 512
+ tensor-memory columns (4 accumulator stages). GB300, 16 calls chained over
+ distinct weights: 4096 / 8192 / 16384 rows about 24 / 45 / 89 us.
+
+ Numerics: the gate is `tanh.approx`-based (as `qwen4_exp_hc_up_gate_mean`'s,
+ max relative error 2^-10.99) with rstd / C folded into it; against the float64
+ value of the unrounded normed stream the result sits in the gate band at
+ about 0.1-0.25 of it (the stored-normed op: 0.3), and is not bitwise
+ `qwen4_exp_hc_up_gate_mean`'s (that op gates the bf16-rounded normed). The
+ fp32 sequence above is bitwise what the GPU test emulates (zero logits).
+
+ Certified: bf16 at the checkpoint's 4 x 2560, rank 320: 16384, 2048, 5
+ rows and a partial last tile (16390); gate band (`rtol = atol / max|ref| =
+ eps`) against the float64 reference; the epilogue's fp32 sequence bitwise
+ (5, 300, 4100 rows); graph replay.
+
+ Not certified: zero rows, hc_count other than 4, hidden other than 2560,
+ rank other than 320.
+ """
+
+ def __call__(
+ self,
+ low: torch.Tensor,
+ up: torch.Tensor,
+ hc_out: torch.Tensor,
+ rstd: torch.Tensor,
+ norm_w1: torch.Tensor,
+ hc_count: int,
+ ) -> torch.Tensor:
+ return hc_up_gate_mean_rstd(low, up, hc_out, rstd, norm_w1, hc_count)
+
+ def reference(
+ self,
+ low: torch.Tensor,
+ up: torch.Tensor,
+ hc_out: torch.Tensor,
+ rstd: torch.Tensor,
+ norm_w1: torch.Tensor,
+ hc_count: int,
+ ) -> torch.Tensor:
+ """float64, one rounding."""
+ return _mixed_reference(low, up, hc_out, rstd, norm_w1, hc_count)
+
+ def is_valid(
+ self,
+ low: torch.Tensor,
+ up: torch.Tensor,
+ hc_out: torch.Tensor,
+ rstd: torch.Tensor,
+ norm_w1: torch.Tensor,
+ hc_count: int,
+ ) -> None:
+ _check(low, up, hc_out, rstd, norm_w1, hc_count)
+
+
+class Qwen4ExpHcUpGateMeanRstdFp4(OpWrapper):
+ """`qwen4_exp_hc_up_gate_mean_rstd` at an MoE site, whose epilogue also writes
+ the routed experts' NVFP4 input::
+
+ mixed = qwen4_exp_hc_up_gate_mean_rstd(low, up, hc_out, rstd, norm_w1, C)
+ xq, xsf = NVFP4(mixed) # e2m1 pairs, e4m3 scale per 16, linear layout
+
+ `mixed` is bitwise the bf16 entry's (same kernel, same epilogue arithmetic);
+ the quantization reads that bf16-rounded value with
+ `qwen4_exp_hc_up_gate_mean_fp4`'s recipe. 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. Fusion boundary: the up projection, the norm's scaling,
+ the sigmoid gate, the stream mean and the NVFP4 quantization of the mean.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: rows. The checkpoint's geometry: hc_count 4, hidden 2560, rank 320, bf16.
+ CELLS: tuple[Cell, ...] = tuple(
+ Cell(why=f"checkpoint 4x2560, rank 320, NVFP4, {why}", spec=dict(rows=t))
+ for t, why in (
+ (16384, "prefill chunk 16384 (max_num_tokens)"),
+ (2048, "prefill 2048 rows"),
+ (16390, "ragged tail, 16320 + 70 rows (mixed step)"),
+ (5, "five rows, less than one tile"),
+ )
+ )
+
+ note = """
+ Tier 6 (CuTe DSL written for this target): `qwen4_exp_hc_up_gate_mean_rstd`'s
+ kernel compiled with its NVFP4 epilogue. Metadata consumed: none. Call
+ sites: `PrefillTarget`'s 48 MoE-site mixers at or above `_HC_RSTD_MIN_ROWS`
+ rows; `xq` / `xsf` feed the routed experts (linear scales), `mixed` the
+ router and the shared expert. Preconditions: those of
+ `qwen4_exp_hc_up_gate_mean_rstd` and `global_scale` one fp32 element on the
+ device. Compiles once per `(C, H, R)`; graph-capturable.
+
+ Each thread holds 16 of a row's 32 tile columns (one scale block); within an
+ epilogue set the upper warp group hands its e2m1 word pair and scale byte to
+ the lower one through shared memory, which stores the row's 16 e2m1 bytes and
+ 2 scale bytes in one store each. GB300, 16 calls chained over distinct
+ weights: 4096 / 8192 / 16384 rows about 26 / 53 / 100 us.
+
+ Numerics: `mixed` as the bf16 entry. The quantization is the trtllm NVFP4
+ recipe instruction for instruction (`rcp.approx.ftz.f32` reciprocals,
+ `cvt.rn.satfinite` to e4m3 and to e2m1, every multiply `mul.rn.f32`); the
+ GPU test checks it against `fp4_quantize`'s native-torch reference on the
+ returned `mixed` (scales bitwise, e2m1 codes bitwise outside the reciprocal
+ tie window).
+
+ Certified: bf16 at the checkpoint's 4 x 2560, rank 320, a calibrated
+ global scale: 16384, 2048, 5 rows and a partial last tile (16390); graph
+ replay.
+
+ Not certified: zero rows, hc_count other than 4, hidden other than 2560,
+ rank other than 320.
+ """
+
+ def __call__(
+ self,
+ low: torch.Tensor,
+ up: torch.Tensor,
+ hc_out: torch.Tensor,
+ rstd: torch.Tensor,
+ norm_w1: torch.Tensor,
+ hc_count: int,
+ global_scale: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ return hc_up_gate_mean_rstd(low, up, hc_out, rstd, norm_w1, hc_count, global_scale)
+
+ def reference(
+ self,
+ low: torch.Tensor,
+ up: torch.Tensor,
+ hc_out: torch.Tensor,
+ rstd: torch.Tensor,
+ norm_w1: 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 = _mixed_reference(low, up, hc_out, rstd, norm_w1, hc_count)
+ return (mixed,) + fp4_quantize.reference(mixed, global_scale, SF_VEC, False, False)
+
+ def is_valid(
+ self,
+ low: torch.Tensor,
+ up: torch.Tensor,
+ hc_out: torch.Tensor,
+ rstd: torch.Tensor,
+ norm_w1: torch.Tensor,
+ hc_count: int,
+ global_scale: torch.Tensor,
+ ) -> None:
+ _check(low, up, hc_out, rstd, norm_w1, hc_count)
+ assert (hc_out.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_rstd = Qwen4ExpHcUpGateMeanRstd()
+qwen4_exp_hc_up_gate_mean_rstd_fp4 = Qwen4ExpHcUpGateMeanRstdFp4()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_post_inject.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_post_inject.py
new file mode 100644
index 000000000000..eadadb7e44ba
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_post_inject.py
@@ -0,0 +1,153 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Hyper-connection post-block injection: a per-stream gated accumulate of one
+block output into the multi-stream residual.
+
+Native torch (sourcing tier 4): one kernel-launching call (``torch.addcmul``)
+plus view glue. No ``torch.ops.trtllm.*`` op computes this accumulate from the
+gates with an *identity* cross-stream mix: ``mhc_post_mapping`` applies a
+Sinkhorn mix across streams, and ``trtllm::qwen4_exp_hc_combine`` takes the
+injection *logits* rather than the gates ``qwen4_exp_gated_residual`` returns.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+#: This checkpoint's hyper-connection geometry: 4 streams of 2560, packed
+#: stream-major into a 10240-wide residual.
+CKPT_HC_COUNT = 4
+CKPT_HIDDEN = 2560
+
+
+class Qwen4ExpPostInject(OpWrapper):
+ """`out[..., s * H + i] = hyper_input[..., s * H + i] + block_out[..., i] *
+ inject[..., s]`, with `H = block_out.shape[-1]` and `hc_count =
+ inject.shape[-1]`: one block output accumulated into **every** stream, each
+ scaled by its own gate. Returns a fresh contiguous tensor of `hyper_input`'s
+ shape; no input is mutated.
+
+ `hyper_input` is the **pre-norm** residual, packed **stream-major** (last
+ dim is `(hc_count, H)`, not `(H, hc_count)` -- both have the same shape, so
+ a stream-minor caller is silently wrong). The cross-stream mix is the
+ identity: no stream's value enters another.
+
+ Fusion boundary: the accumulate only. The grouped norm, the input mixer
+ that produced the block's input and `inject` (`qwen4_exp_gated_residual`),
+ and the block itself are the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: leading (token dims), hc_count, hidden, dtype; layout
+ # "contiguous" (default) | "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 (
+ (64, "decode batch / MTP rows, 64 tokens"),
+ (8192, "prefill 8192 tokens, largest measured"),
+ )
+ ) + (
+ Cell(
+ why="hyper_input a column slice buf[:, :width]",
+ spec=dict(leading=(8,), hc_count=CKPT_HC_COUNT, hidden=CKPT_HIDDEN, layout="col_slice"),
+ ),
+ )
+ # The tp4/tep4/dep4 layouts' sites, each with `hyper_input` a column
+ # slice of a wider row buffer (row stride twice the packed width). The op
+ # sees a rank's full 4 x 2560 residual at every member; only the row count
+ # moves: one decode request, max_batch_size, a verify step's trunk and a
+ # max_num_tokens prefill chunk.
+ CELLS += tuple(
+ Cell(
+ why=f"col_slice {t} rows: {why}",
+ spec=dict(leading=(t,), hc_count=CKPT_HC_COUNT, hidden=CKPT_HIDDEN, layout="col_slice"),
+ )
+ for t, why in (
+ (1, "one decode request; the slice reads as contiguous at one row"),
+ (128, "a max_batch_size decode step"),
+ (1024, "a verify step's trunk, 128 requests x verify width 8"),
+ (8192, "a max_num_tokens prefill chunk, the envelope's largest row"),
+ )
+ )
+
+ note = """
+ Metadata consumed: none. Call sites: twice per decoder layer (after the
+ attention/GDN block and after the MoE) and twice in the MTP layer, with
+ `inject` the `[T, 4]` weights `qwen4_exp_gated_residual` returned for that
+ block -- in `[0, 2]`, because they are `2 * sigmoid(.)`.
+
+ Preconditions, all loud except the guarded one: the packed width must equal
+ `inject.shape[-1] * block_out.shape[-1]` (the unflatten raises); all three
+ on one CUDA device (raises). The leading dims must be **equal**, not merely
+ broadcastable -- a 1-D `inject` or `block_out` is accepted and broadcasts
+ one row over every token (`is_valid`).
+
+ Mixed dtypes are not rejected, they promote: an fp32 `inject` beside a
+ bf16 stream returns an fp32 result -- correct values, twice the bytes.
+ Recorded, not guarded.
+
+ Rounding: `addcmul` forms product and sum in fp32 opmath and rounds once
+ on the store; bit-exact against an fp32 reference rounded once at 1..8192
+ tokens.
+
+ CUDA-graph capture certified; the output is a fresh allocation, so under
+ capture it comes from the graph's private pool.
+
+ Certified: bf16 2-D operands at the checkpoint's 4 x 2560, 64 and 8192
+ rows contiguous; `hyper_input` a column slice (`buf[:, :width]` of a
+ 2 x width row buffer) at 1, 8, 128, 1024 and 8192 rows; graph replay. At
+ one row the slice is contiguous by torch's definition (a size-1 dim's
+ stride is ignored), so that cell is the call site's layout, not a strided
+ read.
+
+ Not certified: other row counts, a row-slice `hyper_input`, hc_count other
+ than 4, hidden other than 2560, float16, float32 and float8 operands,
+ leading dims beyond one. Hard limits: not probed; 2^31 elements of the
+ packed residual is 209,715 rows, past the target's max_num_tokens 8192.
+ """
+
+ def __call__(
+ self,
+ hyper_input: torch.Tensor,
+ block_out: torch.Tensor,
+ inject: torch.Tensor,
+ ) -> torch.Tensor:
+ streams = hyper_input.unflatten(-1, (inject.shape[-1], block_out.shape[-1]))
+ return torch.addcmul(streams, block_out.unsqueeze(-2), inject.unsqueeze(-1)).flatten(-2)
+
+ def reference(
+ self,
+ hyper_input: torch.Tensor,
+ block_out: torch.Tensor,
+ inject: torch.Tensor,
+ ) -> torch.Tensor:
+ """fp32 product and sum, rounded once to the promoted dtype."""
+ hc_count, hidden = inject.shape[-1], block_out.shape[-1]
+ streams = hyper_input.float().unflatten(-1, (hc_count, hidden))
+ gated = block_out.float().unsqueeze(-2) * inject.float().unsqueeze(-1)
+ out_dtype = torch.promote_types(
+ torch.promote_types(hyper_input.dtype, block_out.dtype), inject.dtype
+ )
+ return (streams + gated).flatten(-2).to(out_dtype)
+
+ def is_valid(
+ self,
+ hyper_input: torch.Tensor,
+ block_out: torch.Tensor,
+ inject: torch.Tensor,
+ ) -> None:
+ # Measured: a 1-D `inject` or `block_out` is accepted and returns a full
+ # `[8, 10240]` result computed from one row shared by every token.
+ assert hyper_input.shape[:-1] == block_out.shape[:-1] == inject.shape[:-1], (
+ f"leading dims must match exactly; got hyper_input {tuple(hyper_input.shape)}, "
+ f"block_out {tuple(block_out.shape)}, inject {tuple(inject.shape)}"
+ )
+
+
+qwen4_exp_post_inject = Qwen4ExpPostInject()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/__init__.py
new file mode 100644
index 000000000000..52a7a9daf028
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py
new file mode 100644
index 000000000000..9ba3b8122735
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py
@@ -0,0 +1,484 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""CuteDSL NVFP4 FC1 of the routed experts: gather, grouped GEMM, SwiGLU and NVFP4
+requantization of the intermediate, in one trtllm call."""
+
+from __future__ import annotations
+
+from typing import NamedTuple, Optional
+
+import torch
+
+import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import SWIGLU_LIMIT_SCALAR_DISABLED
+from tensorrt_llm._torch.utils import ActivationType
+
+_SWIGLU = int(ActivationType.Swiglu)
+
+#: e2m1 magnitudes by 3-bit code, and the midpoints between neighbours. A value
+#: exactly on a midpoint rounds to the even code (`_TIE_UP[i]`: the even side is up).
+E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
+E2M1_MAX = 6.0
+E4M3_MAX = 448.0
+_E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0)
+_TIE_UP = (False, True, False, True, False, True, False)
+
+#: The pairing the fused SwiGLU epilogue reads: the weight's rows `[128 g, 128 g + 64)`
+#: are up rows, the next 64 their gate rows, and output column `64 g + r` pairs row
+#: `128 g + r` with row `128 g + 64 + r`. Fixed by the epilogue's 64-column subtile,
+#: whatever MMA N tile the autotuner picks.
+GATE_UP_GROUP = 64
+
+
+def _ckpt(
+ tokens: int,
+ tile_size: int,
+ mma_n: int = 128,
+ dist: str = "router",
+ gsf: str = "calibrated",
+ zero_tokens: int = 0,
+) -> dict:
+ """The checkpoint's routed experts: 512 experts, top-10, hidden 2560, I 640, all
+ local (tp1), SwiGLU, no clamp, 16-element e4m3 block scales.
+
+ `mma_n` is the autotuner tactic's MMA N tile (128 is also the fallback when the
+ tuner has no entry). `dist` names how the test draws the expert ids: `router` =
+ each row's top-10 of uniform scores. `gsf` sets `global_sf` against the
+ activations' magnitude: `calibrated` = the activation maximum the test
+ calibrates to maps to the e4m3 maximum, as the checkpoint's per-layer `1 /
+ down_proj.input_scale` does; `subnormal` = 2^-12 of it, so block scales fall
+ into e4m3 subnormals and zero. `zero_tokens` makes that many leading tokens
+ all-zero rows (zero block scales).
+ """
+ return dict(
+ tokens=tokens,
+ tile_size=tile_size,
+ mma_n=mma_n,
+ dist=dist,
+ gsf=gsf,
+ zero_tokens=zero_tokens,
+ num_experts=512,
+ top_k=10,
+ hidden=2560,
+ intermediate=640,
+ )
+
+
+class Fc1Expected(NamedTuple):
+ """What `reference` returns: the op's two outputs, plus what the test's compare reads.
+
+ `c` / `c_sf` are laid out exactly as the op returns them; rows the op does not
+ write hold zeros. `rows` marks the permuted rows the op writes and the compare
+ gates. `scaled` is `y * global_sf` on those rows, before rounding -- the value
+ the block scale and the codes are derived from, which rounding loses.
+ """
+
+ c: torch.Tensor
+ c_sf: torch.Tensor
+ rows: torch.Tensor
+ scaled: torch.Tensor
+
+
+class CuteDslNvfp4GatherGroupedGemmActFusion(OpWrapper):
+ """`torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell`:
+ FC1 of the routed experts on the CuteDSL MoE path, one kernel launch.
+
+ For every permuted row `p` the sort filled -- tile `i = p // tile_size` below
+ `num_non_exiting_tiles` and `p < tile_idx_to_mn_limit[i]`:
+
+ t = permuted_idx_to_expanded_idx[p] // top_k the token
+ e = tile_idx_to_group_idx[i] the local expert
+ acc = alpha[e] * deq(input[t]) @ deq(weight[e]).T [2 I], fp32
+ y[p, j] = up * silu(gate), up = acc[128 g + r], gate = acc[128 g + 64 + r],
+ j = 64 g + r
+ out[p], out_sf[p] = NVFP4(y[p] * global_sf) 16-element blocks
+
+ where `deq` is e2m1 times its e4m3 block scale, and NVFP4 takes per block
+ `sf = e4m3(max|v| / 6)` (saturating at 448) and codes `e2m1(v / sf)` (saturating
+ at 6). So `out * out_sf / global_sf` approximates `y`.
+
+ Arguments, as the CuteDSL MoE passes them (`run_moe_nvfp4_impl`,
+ `_torch/moe/fused_moe/fused_moe_cute_dsl.py`):
+
+ - `input` `[T, H/2]` float4_e2m1fn_x2 and `input_scale` `[T, H/16]` uint8: the
+ un-permuted activation as `fp4_quantize(x, 1 / input_scale_1, 16, False,
+ False)` returns it -- LINEAR (row-major) e4m3 scales, viewed 2-D. Gathered
+ inside; the caller never permutes it.
+ - `weight` `[E, 2I, H/2]` float4_e2m1fn_x2: per expert `cat([up, gate])`
+ interleaved in `GATE_UP_GROUP`-row groups (up rows 0:64, gate rows 0:64, up
+ rows 64:128, ...). No 32-row shuffle (that is trtllm-gen's layout).
+ - `weight_scale` `[E, 2I, H/16]` uint8: the e4m3 block scales with the same row
+ interleave, then 128x4-swizzled per expert (`block_scale_interleave`).
+ - `alpha` `[E]` fp32: per expert, `input_scale_1 * weight_scale_2[e]` -- one
+ factor for both halves, applied to the accumulator before SwiGLU.
+ - the four sort tensors, exactly as `moe_sort` returns them for the same
+ `top_k` and `tile_tokens_dim == tile_size`.
+ - `global_sf` fp32, one element: the FC2 input's global scale, `1 /
+ max(down_proj.input_scale)` of the layer.
+ - `num_experts`, `num_local_experts`, `local_expert_offset`: read only by the
+ autotuner's shape buckets and cache key; the kernel takes the expert count
+ from `weight.size(0)` and the expert from the tile.
+
+ Returns `(c, c_sf)`: `c` `[M, I/2]` float4_e2m1fn_x2 and `c_sf` `[M * I/16]`
+ uint8, `M = permuted_idx_to_expanded_idx.numel()`, one row per PERMUTED row and
+ the block scales 128x4-swizzled over `[M, I/16]` -- exactly the `input` /
+ `input_scale` the plain NVFP4 grouped GEMM (FC2) takes, with the same four sort
+ tensors. Rows the sort did not fill are not part of the result (see `note`).
+
+ Fusion boundary: the gather of the activation rows, the grouped GEMM, the
+ per-expert alpha, SwiGLU and the requantization happen inside. The caller owns
+ the activation quantization, the sort, FC2 and the routing-weighted combine;
+ routing weights are not read here.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="tile 128: decode, 1 request without MTP",
+ spec=_ckpt(1, 128),
+ ),
+ Cell(
+ why="tile 128: 8192 rows, max_num_tokens prefill",
+ spec=_ckpt(8192, 128),
+ ),
+ Cell(
+ why="tile 128: tiny global_sf and 7 all-zero token rows; subnormal and zero "
+ "block scales",
+ spec=_ckpt(2048, 128, gsf="subnormal", zero_tokens=7),
+ ),
+ )
+
+ note = """
+ Path certified: the CuTe DSL kernel `BlockScaledContiguousGatherGroupedGemmKernel`
+ (`_torch/moe/kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py`),
+ built by `Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner`
+ (`_torch/custom_ops/cute_dsl_custom_ops.py`), which JIT-compiles it with
+ `cute.compile` on the first call per (tile, tactic, top_k, activation) and
+ keeps it for the process. The op is registered only when the CuTe DSL imports
+ (`IS_CUTLASS_DSL_AVAILABLE`), and the runner accepts SM 100 and 103 only.
+
+ Metadata consumed: none beyond the arguments. Stateless: allocates `c` and
+ `c_sf` on every call, on the current stream, with no host synchronization, so
+ it captures into a CUDA graph together with moe_sort.
+
+ Call sites: the engine's CuteDSL MoE (`run_moe_nvfp4_impl`) and the target's
+ routed-expert step, eager for prefill and mixed batches (T up to
+ max_num_tokens 8192) and captured for decode (T = requests x MTP verify width,
+ up to 32 x 4 = 128). Tile 128 is the target's and the cells' tile; 256 is
+ certified through the per-pair invariance below.
+
+ Tactics. The op asks the AutoTuner for an MMA tile of (tile_size, 128) or
+ (tile_size, 256) and with no tuning entry falls back to (tile_size, 128);
+ serving warmup tunes per row bucket and may pick either. Every pair's codes
+ and block scales are bitwise identical across tile 128 / 256 x MMA N 128 / 256
+ x the fallback (the test checks it at 8192 rows), so cells driven at one
+ tactic certify the others.
+
+ Determinism. A (token, expert) pair's output row depends on that pair alone:
+ it is bitwise unchanged by reordering an expert's rows (moe_sort orders them by
+ atomics from T = 5), by repeating the call, and by cutting a batch to a prefix.
+
+ What is written: the rows the sort filled hold the result; the padding rows of
+ a written tile -- from its tile_idx_to_mn_limit to its end -- get zero block
+ scales and zero codes, so a plain grouped GEMM over whole tiles (FC2) reads
+ zeros there; the tiles at and past num_non_exiting_tiles are not written at
+ all, and under a graph keep whatever the buffer last held. FC2 must be tiled by
+ the same tile arrays, and the combine reads only
+ expanded_idx_to_permuted_idx's rows.
+
+ Scale range: with a tiny global_sf the block scales fall to e4m3 subnormals
+ and to zero (a cell). An all-zero token gives scale 0 and codes of magnitude
+ 0 (the same cell). With global_sf 25x the calibrated value the block scales
+ clip at 448 and the codes at 6 (no NaN) -- measured, not certified. A block
+ whose scale is 0 dequantizes to zero whatever its codes hold, and the test's
+ compare does not read them.
+
+ Preconditions the caller owns (a wrong preparation of each operand -- the
+ weight's rows, its scales, the activation's scales, alpha -- runs without
+ error, with finite scales, and puts most block scales and codes outside the
+ test's gate):
+ - `weight` / `weight_scale` in the layout the class docstring states: not
+ trtllm-gen's row interleave, up rows first, scales swizzled after the same
+ row interleave as the weight, alpha not shifted by an expert.
+ - One `alpha` per expert scales BOTH halves, so up_proj and gate_proj must
+ share `weight_scale_2`, and `input_scale` must be one value per layer; this
+ checkpoint satisfies both. `global_sf` inverts the largest
+ `down_proj.input_scale`; the spread of the rest is FC2's `alpha`.
+ - `input_scale` LINEAR. Swizzled scales have the same size when T % 128 == 0
+ and are read as garbage without error.
+ - The sort from moe_sort with the same `top_k` and tile_tokens_dim ==
+ `tile_size`, over the tokens `input` holds. Nothing checks `top_k` against it.
+
+ Refused by `is_valid`, each silently wrong unguarded: a non-contiguous input,
+ input_scale, weight, weight_scale or alpha (data_ptr() read as dense); an
+ alpha shorter than weight.size(0); a weight_scale with fewer bytes than the
+ weight needs (read past its end: zeros of the same allocation come back as zero
+ experts, a fresh allocation faults).
+
+ Loud (AssertionError from the runner before the launch): tile_size 256 with a
+ tile-128 sort and 128 with a tile-256 sort; input_scale not H/16 wide;
+ global_sf of 2 elements; alpha not fp32; weight_scale typed float8_e4m3fn;
+ input typed uint8. tile_size 64: ValueError.
+
+ Measured at bring-up and not re-driven by the test: the CUDA-graph capture
+ with moe_sort, the determinism and write statements, the wrong preparations'
+ distance from the gate, the short-alpha and short-weight_scale refusals and
+ the loud rejections. The test re-drives the cells, the tile / tactic
+ invariance and the `input` contiguity refusal.
+
+ Uncertified: Relu2 and SiTU; a SwiGLU clamp (swiglu_limit_scalar >= 0);
+ expert_counts (DeepEP count-native metadata); expert parallelism (the kernel
+ reads neither num_experts nor local_expert_offset, so EP changes only the sort
+ and the weight slice); other hidden / intermediate sizes and top_k;
+ scaling_vector_size other than 16; other architectures.
+
+ Hard limits (sm_103)
+ written rows x I/16 B of c_sf boundary 2^31 B: an illegal memory access
+ once rows written reach 53687092 at I = 640 (T ~ 5.37M at
+ top-10; T = 8192 writes at most 212480 rows)
+ input, input_scale, weight past 2^31 elements / bytes: exact (probed
+ up to 1536 experts)
+ permuted index int32, within moe_sort's 2^31-pair limit; no index passes
+ through fp32
+ """
+
+ def __call__(
+ self,
+ input: torch.Tensor,
+ weight: torch.Tensor,
+ input_scale: torch.Tensor,
+ weight_scale: torch.Tensor,
+ alpha: torch.Tensor,
+ tile_idx_to_group_idx: torch.Tensor,
+ tile_idx_to_mn_limit: torch.Tensor,
+ permuted_idx_to_expanded_idx: torch.Tensor,
+ num_non_exiting_tiles: torch.Tensor,
+ global_sf: torch.Tensor,
+ num_experts: int,
+ top_k: int,
+ num_local_experts: int,
+ local_expert_offset: int,
+ tile_size: int,
+ scaling_vector_size: int = 16,
+ activation_type: int = _SWIGLU,
+ swiglu_limit_scalar: float = SWIGLU_LIMIT_SCALAR_DISABLED,
+ expert_counts: Optional[torch.Tensor] = None,
+ expert_capacity: int = 0,
+ situ_beta: Optional[float] = None,
+ situ_linear_beta: Optional[float] = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ return torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell(
+ input=input,
+ weight=weight,
+ input_scale=input_scale,
+ weight_scale=weight_scale,
+ alpha=alpha,
+ tile_idx_to_group_idx=tile_idx_to_group_idx,
+ tile_idx_to_mn_limit=tile_idx_to_mn_limit,
+ permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx,
+ num_non_exiting_tiles=num_non_exiting_tiles,
+ global_sf=global_sf,
+ num_experts=num_experts,
+ top_k=top_k,
+ num_local_experts=num_local_experts,
+ local_expert_offset=local_expert_offset,
+ tile_size=tile_size,
+ scaling_vector_size=scaling_vector_size,
+ activation_type=activation_type,
+ swiglu_limit_scalar=swiglu_limit_scalar,
+ expert_counts=expert_counts,
+ expert_capacity=expert_capacity,
+ situ_beta=situ_beta,
+ situ_linear_beta=situ_linear_beta,
+ )
+
+ def reference(
+ self,
+ input: torch.Tensor,
+ weight: torch.Tensor,
+ input_scale: torch.Tensor,
+ weight_scale: torch.Tensor,
+ alpha: torch.Tensor,
+ tile_idx_to_group_idx: torch.Tensor,
+ tile_idx_to_mn_limit: torch.Tensor,
+ permuted_idx_to_expanded_idx: torch.Tensor,
+ num_non_exiting_tiles: torch.Tensor,
+ global_sf: torch.Tensor,
+ num_experts: int,
+ top_k: int,
+ num_local_experts: int,
+ local_expert_offset: int,
+ tile_size: int,
+ scaling_vector_size: int = 16,
+ activation_type: int = _SWIGLU,
+ swiglu_limit_scalar: float = SWIGLU_LIMIT_SCALAR_DISABLED,
+ expert_counts: Optional[torch.Tensor] = None,
+ expert_capacity: int = 0,
+ situ_beta: Optional[float] = None,
+ situ_linear_beta: Optional[float] = None,
+ ) -> Fc1Expected:
+ """Gather, dequantize, GEMM in float64, SwiGLU, exact NVFP4 rounding.
+
+ Products of e2m1 and e4m3 values are exact, and float64 sums of 2560 of
+ them are exact to ~1e-16 relative, so this side carries no accumulation
+ error of its own; the block scale and codes are the exact round-to-nearest-
+ even of the exact `y * global_sf`. Where a block scale rounds to zero the
+ codes here are zero: they dequantize to zero whatever they hold, and the
+ test's compare does not read them.
+
+ Certified surface only: SwiGLU without a clamp, the sort-metadata path
+ (`expert_counts` None), no SiTU. Anything else raises NotImplementedError.
+ `num_experts`, `num_local_experts`, `local_expert_offset` and
+ `expert_capacity` do not enter the computation.
+ """
+ if (
+ int(activation_type) != _SWIGLU
+ or swiglu_limit_scalar >= 0
+ or expert_counts is not None
+ or situ_beta is not None
+ or situ_linear_beta is not None
+ ):
+ raise NotImplementedError(
+ "reference covers SwiGLU without a clamp on sort metadata only"
+ )
+ vec = scaling_vector_size
+ dev = input.device
+ m = permuted_idx_to_expanded_idx.numel()
+ n2 = weight.shape[1]
+ inter = n2 // 2
+ hidden = input.shape[1] * 2
+ n_tiles = int(num_non_exiting_tiles.reshape(-1)[0])
+ span = n_tiles * tile_size
+
+ rows = torch.zeros(m, dtype=torch.bool, device=dev)
+ limit = tile_idx_to_mn_limit[:n_tiles].long().repeat_interleave(tile_size)
+ rows[:span] = torch.arange(span, device=dev) < limit
+ real = torch.nonzero(rows).squeeze(1)
+ token = permuted_idx_to_expanded_idx.long()[real] // top_k
+ expert = tile_idx_to_group_idx[:n_tiles].long().repeat_interleave(tile_size)[real]
+
+ a = _dequant(input.view(torch.uint8)[token], input_scale.view(torch.uint8)[token], vec)
+ y = torch.empty(real.numel(), inter, dtype=torch.float64, device=dev)
+ w_sf = weight_scale.view(torch.uint8).reshape(weight.shape[0], -1)
+ group = GATE_UP_GROUP
+ for e in torch.unique(expert).tolist():
+ sel = torch.nonzero(expert == e).squeeze(1)
+ w = _dequant(weight[e].view(torch.uint8), unswizzle(w_sf[e], n2, hidden // vec), vec)
+ acc = (a[sel] @ w.T) * float(alpha[e])
+ acc = acc.view(-1, n2 // (2 * group), 2, group)
+ up = acc[:, :, 0].reshape(-1, inter)
+ gate = acc[:, :, 1].reshape(-1, inter)
+ y[sel] = up * gate * torch.sigmoid(gate)
+ scaled = y * float(global_sf.reshape(-1)[0])
+
+ blocks = scaled.view(-1, inter // vec, vec)
+ pre = blocks.abs().amax(-1) / E2M1_MAX
+ sf = pre.clamp(max=E4M3_MAX).float().to(torch.float8_e4m3fn)
+ sf_val = sf.double().unsqueeze(-1)
+ q = torch.where(sf_val > 0, blocks / sf_val, torch.zeros_like(blocks))
+ codes = e2m1_codes(q.reshape(-1, inter))
+ c = torch.zeros(m, inter // 2, dtype=torch.uint8, device=dev)
+ c[real] = codes[:, 0::2] | (codes[:, 1::2] << 4)
+ sf_rows = torch.zeros(m, inter // vec, dtype=torch.uint8, device=dev)
+ sf_rows[real] = sf.view(torch.uint8)
+ return Fc1Expected(
+ c=c.view(torch.float4_e2m1fn_x2), c_sf=swizzle(sf_rows), rows=rows, scaled=scaled
+ )
+
+ def is_valid(
+ self,
+ input: torch.Tensor,
+ weight: torch.Tensor,
+ input_scale: torch.Tensor,
+ weight_scale: torch.Tensor,
+ alpha: torch.Tensor,
+ tile_idx_to_group_idx: torch.Tensor,
+ tile_idx_to_mn_limit: torch.Tensor,
+ permuted_idx_to_expanded_idx: torch.Tensor,
+ num_non_exiting_tiles: torch.Tensor,
+ global_sf: torch.Tensor,
+ num_experts: int,
+ top_k: int,
+ num_local_experts: int,
+ local_expert_offset: int,
+ tile_size: int,
+ scaling_vector_size: int = 16,
+ activation_type: int = _SWIGLU,
+ swiglu_limit_scalar: float = SWIGLU_LIMIT_SCALAR_DISABLED,
+ expert_counts: Optional[torch.Tensor] = None,
+ expert_capacity: int = 0,
+ situ_beta: Optional[float] = None,
+ situ_linear_beta: Optional[float] = None,
+ ) -> None:
+ # The runner hands data_ptr() to the kernel with dense row-major layouts
+ # assumed; it checks dims and dtypes, never strides.
+ for name, t in (
+ ("input", input),
+ ("input_scale", input_scale),
+ ("weight", weight),
+ ("weight_scale", weight_scale),
+ ("alpha", alpha),
+ ):
+ assert t.is_contiguous(), (
+ f"{name} must be contiguous; the kernel reads it from data_ptr() as a "
+ "dense row-major buffer and silently reads the wrong elements"
+ )
+ # One alpha per expert of `weight`; the runner checks only alpha.dim() == 1.
+ assert alpha.numel() == weight.size(0), (
+ f"alpha has {alpha.numel()} elements for {weight.size(0)} experts; the "
+ "kernel reads alpha[expert] past the end without raising"
+ )
+ # The runner checks weight_scale.dim() == 3 only, never its size.
+ expected = weight.numel() * 2 // scaling_vector_size
+ assert weight_scale.numel() == expected, (
+ f"weight_scale holds {weight_scale.numel()} bytes, the weight needs "
+ f"{expected}; the kernel reads past its end -- silently within the "
+ "allocation, with an illegal memory access beyond it"
+ )
+
+
+# ── native-torch NVFP4 helpers for reference() and the test's compare ────────
+
+
+def unswizzle(buf: torch.Tensor, rows: int, cols: int) -> torch.Tensor:
+ """128x4-swizzled scale bytes -> `[rows, cols]` (rows % 128 == 0, cols % 4 == 0).
+
+ Byte `(r, c)` sits at `(r // 128) * 128 * cols + (c // 4) * 512 + (r % 32) * 16
+ + ((r % 128) // 32) * 4 + c % 4`.
+ """
+ return buf.reshape(rows // 128, cols // 4, 32, 4, 4).permute(0, 3, 2, 1, 4).reshape(rows, cols)
+
+
+def swizzle(sf: torch.Tensor) -> torch.Tensor:
+ """`[rows, cols]` scale bytes -> the flat 128x4-swizzled buffer; inverse of `unswizzle`."""
+ rows, cols = sf.shape
+ return sf.reshape(rows // 128, 4, 32, cols // 4, 4).permute(0, 3, 2, 1, 4).reshape(-1)
+
+
+def e4m3_value(byte: torch.Tensor) -> torch.Tensor:
+ return byte.to(torch.uint8).view(torch.float8_e4m3fn).double()
+
+
+def _dequant(packed: torch.Tensor, sf: torch.Tensor, vec: int) -> torch.Tensor:
+ """`[R, K/2]` packed e2m1 (element 2j in the low nibble) and `[R, K/vec]` e4m3
+ bytes -> float64 `[R, K]`. Exact."""
+ lut = torch.tensor(E2M1, dtype=torch.float64, device=packed.device)
+ lo, hi = (packed & 0xF).long(), (packed >> 4).long()
+ codes = torch.stack([lo, hi], dim=-1).reshape(packed.shape[0], -1)
+ vals = lut[codes & 7] * torch.where(codes >= 8, -1.0, 1.0).double()
+ return vals * e4m3_value(sf).repeat_interleave(vec, dim=1)
+
+
+def e2m1_codes(v: torch.Tensor) -> torch.Tensor:
+ """float -> 4-bit e2m1 codes (bit 3 = sign), round-to-nearest-even, saturating at 6."""
+ mid = torch.tensor(_E2M1_MIDPOINTS, device=v.device, dtype=v.dtype)
+ tie_up = torch.tensor(_TIE_UP, device=v.device)
+ a = v.abs().unsqueeze(-1)
+ code = ((a > mid).sum(-1) + ((a == mid) & tie_up).any(-1)).to(torch.uint8)
+ negative = (v < 0) | ((v == 0) & torch.signbit(v))
+ return code | (negative.to(torch.uint8) << 3)
+
+
+cute_dsl_nvfp4_gather_grouped_gemm_act_fusion = CuteDslNvfp4GatherGroupedGemmActFusion()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/cute_dsl_nvfp4_grouped_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/cute_dsl_nvfp4_grouped_gemm.py
new file mode 100644
index 000000000000..3d7512323e9c
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/cute_dsl_nvfp4_grouped_gemm.py
@@ -0,0 +1,367 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""CuteDSL NVFP4 FC2 of the routed experts: the plain grouped GEMM over the permuted
+FC1 output, one bf16 row per permuted row, in one trtllm call."""
+
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import torch
+
+import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+from .cute_dsl_nvfp4_gather_grouped_gemm_act_fusion import _dequant, unswizzle
+
+
+def _ckpt(
+ tokens: int,
+ tile_size: int,
+ mma_n: int = 128,
+ cluster_n: int = 1,
+ dist: str = "router",
+ gsf: str = "calibrated",
+ zero_tokens: int = 0,
+) -> dict:
+ """The checkpoint's routed experts: 512 experts, top-10, hidden 2560, I 640, all
+ local (tp1), bf16 out, 16-element e4m3 block scales.
+
+ `mma_n` / `cluster_n` select the autotuner tactic: MMA tile `(tile_size,
+ mma_n)`, cluster `(tile_size // 128, cluster_n)`; (128, 1) is also the
+ fallback when the tuner has no entry. `dist` names how the test draws the
+ expert ids: `router` = each row's top-10 of uniform scores. `gsf` is the FC1
+ `global_sf` that quantized this op's input, against the SwiGLU output's
+ magnitude: `calibrated` = the output maximum the test calibrates to maps to
+ the e4m3 maximum, as the checkpoint's per-layer `1 / down_proj.input_scale`
+ does; `subnormal` = 2^-12 of it, so input block scales fall into e4m3
+ subnormals and zero. `zero_tokens` makes that many leading tokens' rows all
+ zero, as FC1 returns them for an all-zero token.
+ """
+ return dict(
+ tokens=tokens,
+ tile_size=tile_size,
+ mma_n=mma_n,
+ cluster_n=cluster_n,
+ dist=dist,
+ gsf=gsf,
+ zero_tokens=zero_tokens,
+ num_experts=512,
+ top_k=10,
+ hidden=2560,
+ intermediate=640,
+ )
+
+
+class Fc2Expected(NamedTuple):
+ """What `reference` returns: the op's output, plus what the test's compare reads.
+
+ `out` is laid out exactly as the op returns it; rows the op does not write
+ hold zeros. `rows` marks the rows the op writes -- always the prefix of
+ `num_non_exiting_tiles * tile_size` rows. `exact` is `alpha[e] * sum_k a_k
+ w_k` on those rows, before any rounding; `magnitude` is `|alpha[e]| * sum_k
+ |a_k w_k|`, the scale every fp32 rounding of the accumulation is relative to.
+ """
+
+ out: torch.Tensor
+ rows: torch.Tensor
+ exact: torch.Tensor
+ magnitude: torch.Tensor
+
+
+class CuteDslNvfp4GroupedGemm(OpWrapper):
+ """`torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_blackwell`: FC2 of the routed
+ experts on the CuteDSL MoE path, one kernel launch.
+
+ For every row `p` of the first `num_non_exiting_tiles` tiles (tile `i = p //
+ tile_size`):
+
+ e = tile_idx_to_group_idx[i] the local expert
+ acc[p] = deq(input[p]) @ deq(weight[e]).T [H], fp32
+ out[p] = bf16(alpha[e] * acc[p]) fp32 product
+
+ where `deq` is e2m1 times its e4m3 block scale. Every row of a live tile is
+ computed, padding rows included; no row is gathered, weighted or summed.
+
+ Arguments, as the CuteDSL MoE passes them (`run_moe_nvfp4_impl`,
+ `_torch/moe/fused_moe/fused_moe_cute_dsl.py`, its non-fused-finalize branch):
+
+ - `input` `[M, I/2]` float4_e2m1fn_x2 and `input_scale` `[M * I/16]` uint8:
+ FC1's output exactly as the CuteDSL FC1
+ (`cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell`) returns it --
+ one row per PERMUTED row, the block scales 128x4-swizzled over `[M,
+ I/16]`. `M = permuted_idx_to_expanded_idx.numel()` of the sort.
+ - `weight` `[E, H, I/2]` float4_e2m1fn_x2: `down_proj` per expert in the
+ checkpoint's row order. No 32-row shuffle (that is trtllm-gen's layout).
+ - `weight_scale` `[E, H, I/16]` uint8: its e4m3 block scales, 128x4-swizzled
+ per expert (`block_scale_interleave`).
+ - `alpha` `[E]` fp32: per expert, `max(down_proj.input_scale) * weight_scale_2[e]`
+ -- the inverse of FC1's `global_sf` times the expert's weight scale.
+ - `tile_idx_to_group_idx` `[M / tile_size]` int32 and `num_non_exiting_tiles`
+ `[1]` int32: moe_sort's `tile_idx_to_expert_idx` and
+ `num_non_exiting_tiles` from the sort FC1 ran on, with tile_tokens_dim ==
+ `tile_size`.
+ - `num_experts`, `top_k`, `num_local_experts`, `local_expert_offset`: read only
+ by the autotuner's shape buckets and cache key; the kernel takes the
+ expert count from `weight.size(0)` and the expert from the tile.
+ - `output_dtype`: bf16, the only one the runner builds.
+
+ Returns `[M, H]` bfloat16, one row per permuted row. The pair `(t, k)` is row
+ `expanded_idx_to_permuted_idx[t, k]`; rows outside the live tiles are not
+ written (see `note`).
+
+ Fusion boundary: the grouped GEMM, the per-expert alpha and the bf16
+ conversion happen inside. The rows arrive already permuted (FC1 gathered
+ them); the routing-weighted top-k combine is the caller's
+ (`moe_unpermute_inplace`). Routing weights are not read here.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="tile 128: decode, 1 request without MTP; the fallback tactic",
+ spec=_ckpt(1, 128),
+ ),
+ Cell(
+ why="tile 128: MMA N 256 and cluster N 2 tactic, 8192-row max_num_tokens prefill",
+ spec=_ckpt(8192, 128, mma_n=256, cluster_n=2),
+ ),
+ Cell(
+ why="tile 128: tiny FC1 global_sf and 7 all-zero token rows; subnormal and zero "
+ "input scales",
+ spec=_ckpt(2048, 128, gsf="subnormal", zero_tokens=7),
+ ),
+ )
+
+ note = """
+ Path certified: the CuTe DSL kernel `Sm100BlockScaledContiguousGroupedGemmKernel`
+ (`_torch/moe/kernels/blackwell/blockscaled_contiguous_grouped_gemm.py`),
+ built by `Sm100BlockScaledContiguousGroupedGemmRunner`
+ (`_torch/custom_ops/cute_dsl_custom_ops.py`), which JIT-compiles it with
+ `cute.compile` on the first call per (tile, tactic) and keeps it for the
+ process. The op is registered only when the CuTe DSL imports
+ (`IS_CUTLASS_DSL_AVAILABLE`), and the runner accepts SM 100 and 103 only.
+
+ Metadata consumed: none beyond the arguments. Stateless: allocates its output
+ with torch.empty on every call, on the current stream, with no host
+ synchronization; `num_non_exiting_tiles` is read on the device, so it
+ captures into a CUDA graph together with moe_sort and FC1.
+
+ Call sites: the engine's CuteDSL MoE (`run_moe_nvfp4_impl`, when the fused
+ finalize is off) and the target's routed-expert step, both followed by
+ `moe_unpermute_inplace`; eager for prefill and mixed batches (T up to
+ max_num_tokens 8192) and captured for decode (T = requests x MTP verify
+ width, up to 32 x 4 = 128). Tile 128 is the target's and the cells' tile;
+ 256 is certified through the per-pair invariance below.
+
+ Numerics: each written element equals the reference's `out` bit for bit on
+ every cell (and, at bring-up, on the real FC1 output at 8192 tokens and on the
+ checkpoint's own down_proj experts). Operands whose block scales span a wider range can round a
+ partial sum; the test's compare admits that at bf16 ties only.
+
+ Tactics. The op asks the AutoTuner for an MMA tile of (tile_size, 128) or
+ (tile_size, 256) and a cluster of (tile_size // 128, 1) or (tile_size //
+ 128, 2), and with no tuning entry falls back to ((tile_size, 128),
+ (tile_size // 128, 1)); serving warmup tunes per row bucket and may pick any
+ of the four. Every pair's row is bitwise identical across tile 128 / 256 x
+ the four tactics x the fallback (the test checks it at 8192 rows), so cells
+ driven at one tactic certify the others.
+
+ Determinism. A (token, expert) pair's output row depends on that pair's input
+ row and its expert alone: it is bitwise unchanged by moving it to another row
+ of its expert's run (moe_sort orders them by atomics from T = 5), by repeating
+ the call, and by cutting a batch to a prefix. This op does not make the MoE
+ run-to-run nondeterministic; the fused-finalize variant
+ (`cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell`) reduces across a
+ token's experts in launch order and does.
+
+ What is read and written: every row of the first num_non_exiting_tiles tiles
+ is written; the pair rows hold their results and the padding rows of a live
+ tile are computed from whatever the input holds there -- zero for FC1's
+ output, which zeroes them. Rows at and past num_non_exiting_tiles * tile_size
+ are never written: eager they hold whatever torch.empty returned, under a
+ graph whatever the pool last held. Input rows past that point reach no written
+ row. The combine must read only expanded_idx_to_permuted_idx's rows.
+
+ Preconditions the caller owns:
+ - `weight` / `weight_scale` in the layout the class docstring states, and
+ `input_scale` swizzled. A wrong preparation of each runs without error,
+ with finite output, and puts most of the pair rows' elements outside the
+ test's gate: trtllm-gen's 32-row shuffle of the weight, input scales
+ linear, one alpha for the whole layer.
+ - `alpha` per expert, `max(down_proj.input_scale) * weight_scale_2[e]`, with
+ the maximum the value FC1's `global_sf` inverts.
+ - `tile_idx_to_group_idx` / `num_non_exiting_tiles` from the same sort as
+ FC1's input, tile_tokens_dim == `tile_size`, and every expert index below
+ `weight.size(0)`: nothing checks the values (a tile naming expert 512 of
+ 512 ran without error and came back all zeros).
+ - Every tensor on the device. Nothing checks it: a host-resident `alpha` or
+ tile table was read correctly on a coherent-memory Grace host, unmeasured
+ where the GPU cannot dereference a host pointer.
+
+ Refused by `is_valid`, each silently wrong unguarded: a non-contiguous input,
+ input_scale, weight, weight_scale, alpha or tile_idx_to_group_idx (data_ptr()
+ read as dense; finite, wrong rows).
+
+ Loud (AssertionError from the runner before the launch): a tile_size other
+ than the sort's, either way round; output_dtype float16; alpha bf16 or short
+ of weight.size(0); weight_scale with fewer experts or typed float8_e4m3fn;
+ input typed uint8; input_scale 2-D or short; scaling_vector_size 32;
+ num_non_exiting_tiles of 2 elements; tile_idx_to_group_idx int64 or one tile
+ short. tile_size 64: ValueError.
+
+ Measured at bring-up and not re-driven by the test: the CUDA-graph capture
+ with moe_sort and FC1, the determinism and read / write statements, the wrong
+ preparations' distance from the gate, and the loud rejections. The test
+ re-drives the cells, the tile / tactic invariance and the `input` contiguity
+ refusal.
+
+ Uncertified: output dtypes other than bf16 (the runner builds none); expert
+ parallelism (the kernel reads neither local_expert_offset nor num_experts, so
+ EP changes only the sort and the weight slice); other hidden / intermediate
+ sizes; scaling_vector_size other than 16; other architectures.
+
+ Hard limits (sm_103)
+ input, input_scale, output, weight, weight_scale past 2^31 elements /
+ bytes: exact (probed up to M = 53687296 rows and 20973
+ experts)
+ tile index int32 tables, within moe_sort's 2^31-pair limit; no index
+ passes through fp32
+ """
+
+ def __call__(
+ self,
+ input: torch.Tensor,
+ weight: torch.Tensor,
+ input_scale: torch.Tensor,
+ weight_scale: torch.Tensor,
+ alpha: torch.Tensor,
+ tile_idx_to_group_idx: torch.Tensor,
+ num_non_exiting_tiles: torch.Tensor,
+ num_experts: int,
+ top_k: int,
+ num_local_experts: int,
+ local_expert_offset: int,
+ tile_size: int,
+ output_dtype: torch.dtype,
+ scaling_vector_size: int = 16,
+ ) -> torch.Tensor:
+ return torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_blackwell(
+ input=input,
+ weight=weight,
+ input_scale=input_scale,
+ weight_scale=weight_scale,
+ alpha=alpha,
+ tile_idx_to_group_idx=tile_idx_to_group_idx,
+ num_non_exiting_tiles=num_non_exiting_tiles,
+ num_experts=num_experts,
+ top_k=top_k,
+ num_local_experts=num_local_experts,
+ local_expert_offset=local_expert_offset,
+ tile_size=tile_size,
+ output_dtype=output_dtype,
+ scaling_vector_size=scaling_vector_size,
+ )
+
+ def reference(
+ self,
+ input: torch.Tensor,
+ weight: torch.Tensor,
+ input_scale: torch.Tensor,
+ weight_scale: torch.Tensor,
+ alpha: torch.Tensor,
+ tile_idx_to_group_idx: torch.Tensor,
+ num_non_exiting_tiles: torch.Tensor,
+ num_experts: int,
+ top_k: int,
+ num_local_experts: int,
+ local_expert_offset: int,
+ tile_size: int,
+ output_dtype: torch.dtype,
+ scaling_vector_size: int = 16,
+ ) -> Fc2Expected:
+ """Dequantize, GEMM in float64, scale by alpha in fp32, round once to bf16.
+
+ Products of e2m1 and e4m3 values are exact, and float64 sums of 640 of
+ them are exact to ~1e-16 relative, so `exact` carries no accumulation
+ error of its own; `out` is that sum rounded to fp32, multiplied by the
+ fp32 alpha, rounded to bf16 -- what an fp32-accumulating GEMM with this
+ epilogue returns whenever its fp32 partial sums are exact. Rows whose
+ input dequantizes to all zeros are zero without a GEMM.
+
+ bf16 output only (anything else raises NotImplementedError).
+ `num_experts`, `top_k`, `num_local_experts` and `local_expert_offset` do
+ not enter the computation.
+ """
+ if output_dtype != torch.bfloat16:
+ raise NotImplementedError("reference covers bf16 output only")
+ vec = scaling_vector_size
+ dev = input.device
+ m = input.shape[0]
+ hidden = weight.shape[1]
+ k_blocks = input.shape[1] * 2 // vec
+ n_tiles = int(num_non_exiting_tiles.reshape(-1)[0])
+ span = n_tiles * tile_size
+
+ rows = torch.zeros(m, dtype=torch.bool, device=dev)
+ rows[:span] = True
+ a = _dequant(
+ input.view(torch.uint8)[:span],
+ unswizzle(input_scale.view(torch.uint8), m, k_blocks)[:span],
+ vec,
+ )
+ live = a.abs().amax(dim=1) > 0
+ expert = tile_idx_to_group_idx[:n_tiles].long().repeat_interleave(tile_size)
+ acc = torch.zeros(span, hidden, dtype=torch.float64, device=dev)
+ mag = torch.zeros(span, hidden, dtype=torch.float64, device=dev)
+ w_sf = weight_scale.view(torch.uint8).reshape(weight.shape[0], -1)
+ for e in torch.unique(expert[live]).tolist():
+ sel = torch.nonzero(live & (expert == e)).squeeze(1)
+ w = _dequant(weight[e].view(torch.uint8), unswizzle(w_sf[e], hidden, k_blocks), vec)
+ acc[sel] = a[sel] @ w.T
+ mag[sel] = a[sel].abs() @ w.abs().T
+ scale = alpha.float()[expert].unsqueeze(1)
+ out = torch.zeros(m, hidden, dtype=torch.bfloat16, device=dev)
+ out[:span] = (acc.float() * scale).to(torch.bfloat16)
+ return Fc2Expected(
+ out=out,
+ rows=rows,
+ exact=acc * scale.double(),
+ magnitude=mag * scale.double().abs(),
+ )
+
+ def is_valid(
+ self,
+ input: torch.Tensor,
+ weight: torch.Tensor,
+ input_scale: torch.Tensor,
+ weight_scale: torch.Tensor,
+ alpha: torch.Tensor,
+ tile_idx_to_group_idx: torch.Tensor,
+ num_non_exiting_tiles: torch.Tensor,
+ num_experts: int,
+ top_k: int,
+ num_local_experts: int,
+ local_expert_offset: int,
+ tile_size: int,
+ output_dtype: torch.dtype,
+ scaling_vector_size: int = 16,
+ ) -> None:
+ # The runner hands data_ptr() to the kernel with dense row-major layouts
+ # assumed; it checks dims, sizes and dtypes, never strides.
+ for name, t in (
+ ("input", input),
+ ("input_scale", input_scale),
+ ("weight", weight),
+ ("weight_scale", weight_scale),
+ ("alpha", alpha),
+ ("tile_idx_to_group_idx", tile_idx_to_group_idx),
+ ):
+ assert t.is_contiguous(), (
+ f"{name} must be contiguous; the kernel reads it from data_ptr() as a "
+ "dense row-major buffer and silently reads the wrong elements"
+ )
+
+
+cute_dsl_nvfp4_grouped_gemm = CuteDslNvfp4GroupedGemm()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_block_scale_moe_renorm_permuted.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_block_scale_moe_renorm_permuted.py
new file mode 100644
index 000000000000..90e2623cf65d
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_block_scale_moe_renorm_permuted.py
@@ -0,0 +1,139 @@
+# 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
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+from .fp4_block_scale_moe_renorm_routed import _ckpt, run_renorm_moe
+from .fp4_block_scale_moe_renorm_routed import fp4_block_scale_moe_renorm_routed as _finalized
+
+
+def finalize(
+ rows: torch.Tensor, topk_weights: torch.Tensor, topk_rows: torch.Tensor
+) -> torch.Tensor:
+ """The runner's finalize, emulated: per token `sum_k w[t, k] * rows[topk_rows[t, k]]`
+ as fp32 FMAs in slot order from zero (a -1 slot skipped), rounded once to the
+ rows' dtype. Each fp32 FMA is formed exactly in float64 (a product of two
+ fp32 values is exact there) and rounded to fp32. Returns `[T, H]`."""
+ acc = torch.zeros(topk_rows.shape[0], rows.shape[1], dtype=torch.float32, device=rows.device)
+ for k in range(topk_rows.shape[1]):
+ idx = topk_rows[:, k].long()
+ r = rows[idx.clamp(min=0)].float().double()
+ w = topk_weights[:, k].float().double().unsqueeze(-1)
+ step = (r * w + acc.double()).float()
+ acc = torch.where((idx >= 0).unsqueeze(-1), step, acc)
+ return acc.to(rows.dtype)
+
+
+class Fp4BlockScaleMoeRenormPermuted(OpWrapper):
+ """`fp4_block_scale_moe_renorm_routed` with `do_finalize=False`.
+
+ Same operands, same routing, same GEMMs; the runner skips its finalize
+ launch and returns three tensors:
+
+ - `rows` `[P, H]` bf16: the FC2 output in the runner's permuted, padded row
+ order (`P` is the runner's padded maximum; rows no slot maps to hold
+ garbage);
+ - `topk_weights` `[T, k]` bf16: the routing weights, slot order;
+ - `topk_rows` `[T, k]` int32: slot `(t, k)`'s row in `rows`, or -1.
+
+ `finalize(rows, topk_weights, topk_rows)` (this module) is bitwise the
+ finalized entry's output on the same operands.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="_moe_block decode, finalize folded into the next site: one token; the block "
+ "routing kernel",
+ spec=_ckpt(1),
+ ),
+ Cell(
+ why="4 rows (MTP3 verify; 2..16-row batches): finalize folded into "
+ "qwen4_exp_hc_decode_rows_first",
+ spec=_ckpt(4),
+ ),
+ Cell(
+ why="128 rows: the con=128 decode batch; the cluster routing kernel",
+ spec=_ckpt(128),
+ ),
+ Cell(why="coarse 3-level logits: ties decide most slots", spec=_ckpt(64, dist="coarse")),
+ )
+
+ note = """
+ Measured against the finalized entry on identical operands: the emulated
+ finalize of the three outputs is bitwise its output at every cell, so the
+ routing semantics, the GEMMs and the weight layout are the finalized
+ entry's (see its note). Which finalize kernel the runner would have
+ launched (scalar below ~1184 blocks, vectorized above) does not change the
+ arithmetic: both accumulate fp32 FMAs in slot order and round once.
+
+ `rows` is the runner's own `[P, H]` buffer, allocated per call at the
+ padded maximum the runner computes from `T`, `k` and the expert count;
+ only rows some slot maps to hold data. `topk_weights` carries the output
+ dtype (bf16), as the finalize reads it.
+
+ Silent, refused by `is_valid`: the finalized entry's refusals (a
+ non-contiguous operand, an uncertified (num_experts, top_k) pair, hidden
+ not a multiple of 256).
+
+ Hard limits
+ T certified at 1, 4, 64 and 128 (decode); the decode forward folds at
+ every site that takes the combine
+ num_experts 512, top_k 10
+ """
+
+ def __call__(
+ self,
+ routing_logits: torch.Tensor,
+ hidden_states: torch.Tensor,
+ hidden_states_scale: torch.Tensor,
+ gemm1_weights: torch.Tensor,
+ gemm1_weights_scale: torch.Tensor,
+ gemm2_weights: torch.Tensor,
+ gemm2_weights_scale: torch.Tensor,
+ output1_scale_scalar: torch.Tensor,
+ output1_scale_gate_scalar: torch.Tensor,
+ output2_scale_scalar: torch.Tensor,
+ num_experts: int,
+ top_k: int,
+ intermediate_size: int,
+ ) -> tuple:
+ rows, topk_weights, topk_rows = run_renorm_moe(
+ routing_logits,
+ hidden_states,
+ hidden_states_scale,
+ gemm1_weights,
+ gemm1_weights_scale,
+ gemm2_weights,
+ gemm2_weights_scale,
+ output1_scale_scalar,
+ output1_scale_gate_scalar,
+ output2_scale_scalar,
+ num_experts,
+ top_k,
+ intermediate_size,
+ False,
+ )
+ 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/ops/moe/fp4_block_scale_moe_renorm_routed.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_block_scale_moe_renorm_routed.py
new file mode 100644
index 000000000000..f6257976bc31
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_block_scale_moe_renorm_routed.py
@@ -0,0 +1,484 @@
+# 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.
+
+It runs over an expert window (`local_expert_offset`, `local_num_experts`; the whole routing
+space by default) at the intermediate width a rank holds: the window of a tep4 / dep4 / dep2
+rank, every expert at a tp4 rank's padded slice, every expert at tp1. A decode call over the
+whole 512-expert layer at I 640 runs its pinned tactic (`DECODE_TACTICS`), every other call the
+AutoTuner's."""
+
+from __future__ import annotations
+
+from functools import cache
+
+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 tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.renorm_moe_routing_op import (
+ renorm_moe_routing_op as _renorm,
+)
+from tensorrt_llm._torch.autotuner import AutoTuner
+from tensorrt_llm._torch.custom_ops.trtllm_gen_custom_ops import FP4BlockScaleMoERunner
+from tensorrt_llm._torch.utils import last_positive_power_of_2
+from tensorrt_llm.logger import logger
+
+#: `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)})
+
+#: The geometry `DECODE_TACTICS` was measured on: (num_experts, top_k, hidden, intermediate).
+DECODE_TACTICS_GEOMETRY = (512, 10, 2560, 640)
+
+#: trtllm-gen tactic `(tileN, FC1 x FC2 config index)` per AutoTuner bucket (the row count
+#: rounded down to a power of two) for the decode row counts, at `DECODE_TACTICS_GEOMETRY`.
+#: Each is the fastest valid tactic of its bucket on GB300 over a 48-layer chain of distinct
+#: expert stacks (HBM-cold weights, as a decode step streams them), with router logits whose
+#: distinct experts per layer match serving's routing skew, in one CUDA graph with PDL. The
+#: tactics compared per bucket (the fastest few, the runner's default, the AutoTuner's picks)
+#: all gave bitwise the same output, so the choice moves time only; the entry's test holds
+#: each pinned tactic bitwise to the AutoTuner path.
+#:
+#: A call outside AutoTuner tuning mode whose bucket is listed runs its tactic directly; any
+#: other call goes through `torch.ops.trtllm.fp4_block_scale_moe_runner` (the AutoTuner), so
+#: the larger buckets are still profiled at warmup. The config index addresses the runner's
+#: compiled cubin list: a tactic the runner does not list as valid is not run; its bucket goes
+#: through the AutoTuner instead, with one warning.
+DECODE_TACTICS: dict[int, tuple[int, int]] = {
+ 1: (8, 55),
+ 2: (8, 53),
+ 4: (8, 44),
+ 8: (8, 0),
+ 16: (8, 44),
+ 32: (8, 44),
+ 64: (16, 38),
+ 128: (16, 36),
+ 256: (32, 32),
+ 512: (32, 32),
+ 1024: (64, 5),
+}
+
+
+@cache
+def _runner():
+ """The SwiGlu trtllm-gen runner object the AutoTuner path also uses (one per act type;
+ the other constructor arguments only shape the AutoTuner's tuning config)."""
+ num_experts, top_k, _, intermediate = DECODE_TACTICS_GEOMETRY
+ return FP4BlockScaleMoERunner(
+ num_experts=num_experts,
+ top_k=top_k,
+ n_group=None,
+ topk_group=None,
+ intermediate_size=intermediate,
+ local_expert_offset=0,
+ local_num_experts=num_experts,
+ routed_scaling_factor=None,
+ routing_method_type=ROUTING_RENORMALIZE,
+ do_finalize=True,
+ act_type=ACT_SWIGLU,
+ ).get_runner()
+
+
+@cache
+def _checked_tactic(bucket: int) -> tuple[int, int] | None:
+ """`DECODE_TACTICS[bucket]` if it is a valid config for the bucket's row count, else None
+ (the AutoTuner path), warned once per bucket."""
+ tactic = DECODE_TACTICS[bucket]
+ num_experts, top_k, hidden, intermediate = DECODE_TACTICS_GEOMETRY
+ valid = _runner().get_valid_configs(top_k, hidden, intermediate, num_experts, bucket)
+ if list(tactic) not in [list(t) for t in valid]:
+ logger.warning(
+ f"pinned tactic {tactic} for {bucket} rows is not a valid config of this build's "
+ "trtllm-gen FP4 MoE runner; using the AutoTuner for this bucket. Re-measure "
+ "DECODE_TACTICS."
+ )
+ return None
+ return tactic
+
+
+def _pinned_tactic(
+ num_tokens: int, num_experts: int, top_k: int, hidden: int, intermediate: int
+) -> tuple[int, int] | None:
+ """The pinned tactic for this call, or None for the AutoTuner path."""
+ if (num_experts, top_k, hidden, intermediate) != DECODE_TACTICS_GEOMETRY:
+ return None
+ bucket = last_positive_power_of_2(num_tokens)
+ if bucket not in DECODE_TACTICS or AutoTuner.get().is_tuning_mode:
+ return None
+ return _checked_tactic(bucket)
+
+
+def run_renorm_moe(
+ 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,
+ do_finalize: bool,
+ local_expert_offset: int = 0,
+ local_num_experts: int | None = None,
+) -> list[torch.Tensor]:
+ """One logits-routed trtllm-gen FP4 MoE call (Renormalize, SwiGlu) over the expert window
+ `[local_expert_offset, local_expert_offset + local_num_experts)` (all experts when
+ `local_num_experts` is None): `[out]` with `do_finalize`, else `[rows, topk_weights,
+ topk_rows]`. A `DECODE_TACTICS` bucket of the whole layer runs its pinned tactic; any other
+ call is the AutoTuner's."""
+ if local_num_experts is None:
+ local_num_experts = num_experts
+ whole = local_expert_offset == 0 and local_num_experts == num_experts
+ tactic = (
+ _pinned_tactic(
+ hidden_states.shape[0],
+ num_experts,
+ top_k,
+ hidden_states.shape[1] * 2,
+ intermediate_size,
+ )
+ if whole
+ else None
+ )
+ if tactic is None:
+ return _pre_routed(
+ 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,
+ local_expert_offset,
+ local_num_experts,
+ None,
+ ROUTING_RENORMALIZE,
+ do_finalize,
+ ACT_SWIGLU,
+ )
+ return _runner().run_moe(
+ 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,
+ do_finalize,
+ list(tactic),
+ None,
+ None,
+ None,
+ )
+
+
+def _ckpt(tokens: int, dist: str = "randn") -> dict:
+ """qwen4_exp's routed geometry: 512 experts, top-10, H 2560 / I 640."""
+ return dict(
+ tokens=tokens,
+ hidden=2560,
+ intermediate=640,
+ num_experts=512,
+ top_k=10,
+ logits_dtype=torch.bfloat16,
+ dist=dist,
+ )
+
+
+def _layer(tokens: int, off: int, local: int, inter: int = 640) -> dict:
+ """qwen4_exp's routed geometry (512 experts, top-10, H 2560, bf16 logits) at one rank's
+ window `[off, off + local)` and intermediate width `inter`."""
+ return dict(
+ tokens=tokens,
+ hidden=2560,
+ intermediate=inter,
+ num_experts=512,
+ top_k=10,
+ local_expert_offset=off,
+ local_num_experts=local,
+ logits_dtype=torch.bfloat16,
+ dist="randn",
+ )
+
+
+class Fp4BlockScaleMoeRenormRouted(OpWrapper):
+ """One NVFP4 MoE layer whose routing happens inside the call, over an expert window.
+
+ `routing_logits` is the router GEMM's `[T, num_experts]` output over the whole routing
+ space. Per token the kernel picks the `top_k` largest logits (ties to the LOWER expert
+ index) and weights them by an fp32 softmax over just those `top_k`, rounded to bf16;
+ slots routed outside `[local_expert_offset, local_expert_offset + local_num_experts)` add
+ nothing. The same function as `ops/moe/renorm_moe_routing_op` followed by the pre-routed
+ `catalog/moe/fp4_block_scale_moe_runner` over the same window, and bitwise equal to that
+ composition at every cell. Returns the combined `[T, H]` bf16 rows (the window's partial).
+
+ Fixed, not arguments: no routing bias, Renormalize routing, SwiGlu, `do_finalize=True`.
+ The weight layout (row shuffle + 128x4 scale swizzle), the window's weight stacks and the
+ three per-expert scalars are exactly the pre-routed entry's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # One cell per routing kernel the call picks by row count: block, single-cluster
+ # and the multi-kernel histogram path.
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="_moe_block decode: one token; the block routing kernel", spec=_ckpt(1)),
+ Cell(why="128 rows: the con=128 decode batch; the cluster routing kernel", spec=_ckpt(128)),
+ Cell(
+ why="16384 rows: max_num_tokens of the tuned config; past the 8192 tuning bucket",
+ spec=_ckpt(16384),
+ ),
+ # The windows: tep4 / dep4 a 128-expert window at I 640, dep2 a 256-expert window, tp4
+ # all 512 experts at its 160-wide slice padded to 256. Rows: 1 (decode), 4 (an MTP-3
+ # verify step), 5 (past the block routing kernel), 128 (dep4's 4 x 32 gathered rows),
+ # 1024 (MTP-7 verify of 128 requests) and 8192 (a prefill chunk); offsets cover every
+ # rank.
+ Cell(why="tep4 decode: one token, rank 0's 128-expert window", spec=_layer(1, 0, 128)),
+ Cell(why="tep4 MTP-3 verify at con1: 4 rows, rank 1's window", spec=_layer(4, 128, 128)),
+ Cell(why="ep4 at 5 rows, rank 2's window: past the block kernel", spec=_layer(5, 256, 128)),
+ Cell(
+ why="dep4 gathered decode 4 x 32 = 128 rows, rank 3's window",
+ spec=_layer(128, 384, 128),
+ ),
+ Cell(why="tep4 MTP-7 verify of 128 requests: 1024 rows", spec=_layer(1024, 0, 128)),
+ Cell(why="ep4 prefill chunk at 8192 rows, rank 3's window", spec=_layer(8192, 384, 128)),
+ Cell(why="tp4 MTP-3 verify: 4 rows, 512 experts at I 256", spec=_layer(4, 0, 512, 256)),
+ Cell(why="tp4 at 128 rows, 512 experts at I 256", spec=_layer(128, 0, 512, 256)),
+ Cell(
+ why="tp4 prefill chunk 8192 rows, 512 experts at I 256", spec=_layer(8192, 0, 512, 256)
+ ),
+ Cell(
+ why="dep2 gathered decode: 2 rows, rank 1's 256-expert window", spec=_layer(2, 256, 256)
+ ),
+ Cell(
+ why="dep2 prefill chunk at 8192 rows, rank 0's 256-expert window",
+ spec=_layer(8192, 0, 256),
+ ),
+ )
+
+ note = """
+ Routing semantics, measured against the composition it replaces
+ (`renorm_moe_routing_op` -> pre-routed runner): bitwise-equal finalized
+ output at every cell (the test re-drives this), and at bring-up on randn /
+ coarse / large-magnitude / all-equal logits, bf16 and fp32 logits,
+ 1..16384 rows. Both kernels rank by a packed
+ (twiddled value, 65535 - index) key (ties to the lower expert index), take
+ an fp32 `expf` softmax over the k with the same 32-lane warp reduction, and
+ round the weight to bf16 once.
+
+ Kernel choice inside the call is by row count (block kernel up to a small
+ count, single-cluster up to the cluster cap, then the multi-kernel
+ histogram path); every one of them shares the top-k and softmax code.
+ `routing_bias` is None (the thop defaults the bias dtype to bf16 and never
+ reads it on this policy). PDL (`TRTLLM_ENABLE_PDL`) is read by the routing
+ launch.
+
+ Silent, refused by `is_valid`: a non-contiguous operand (raw pointers,
+ dense layout assumed) and a (num_experts, top_k) pair outside
+ `CERTIFIED_PAIRS` (a tier-list miss logs and launches nothing; the test
+ holds this refusal).
+
+ Not certified: fp32 logits (the runner's other accepted logits dtype).
+
+ Tactics: a call of 1..2047 rows outside AutoTuner tuning mode runs the
+ `DECODE_TACTICS` tactic of its power-of-two bucket (measured HBM-cold with
+ serving's routing skew); the AutoTuner, which profiles one layer's experts
+ warm in L2 under uniform routing, takes the rest and the warmup's tuning
+ pass. The test holds each pinned tactic bitwise to the AutoTuner path.
+
+ Windows: a window's weight stacks are its own experts only (the op compares the scales'
+ expert dimension with `local_num_experts`; the whole stack beside a window's scales would be
+ read as its first experts, refused by `is_valid`). Not certified: windows of the padded
+ I 256 stack (tp4 does not split experts); window rows past 8192 (the caller chunks).
+
+ Hard limits
+ window the whole 512 at I 640 (tp1), 128 of 512 at I 640 (every offset), 256 of
+ 512 at I 640, 512 at I 256
+ T certified at 1, 128 and 16384 (max_num_tokens of the tuned config).
+ Autotuner buckets stop at 8192, the op's
+ `tune_max_num_tokens` default, which this entry does not
+ 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,
+ local_expert_offset: int = 0,
+ local_num_experts: int | None = None,
+ ) -> torch.Tensor:
+ return run_renorm_moe(
+ routing_logits,
+ hidden_states,
+ hidden_states_scale,
+ gemm1_weights,
+ gemm1_weights_scale,
+ gemm2_weights,
+ gemm2_weights_scale,
+ output1_scale_scalar,
+ output1_scale_gate_scalar,
+ output2_scale_scalar,
+ num_experts,
+ top_k,
+ intermediate_size,
+ True,
+ local_expert_offset,
+ local_num_experts,
+ )[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,
+ local_expert_offset: int = 0,
+ local_num_experts: int | None = None,
+ ) -> torch.Tensor:
+ """The routing, then the window's layer, both by their own entries' references.
+
+ Routing: `renorm_moe_routing_op.reference` (stable descending sort in the
+ logits' dtype, fp32 softmax over the k, bf16 weights -- the dtype the
+ 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,
+ local_expert_offset=local_expert_offset,
+ local_num_experts=local_num_experts,
+ )
+
+ 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,
+ local_expert_offset: int = 0,
+ local_num_experts: int | None = None,
+ ) -> None:
+ # Raw data pointers, dense row-major assumed: a strided view is read as
+ # if dense. A router GEMM output written into a wider buffer is the
+ # 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"
+ )
+ # The op checks the scales' expert dimension, not the weights'.
+ local = num_experts if local_num_experts is None else local_num_experts
+ for name, tensor in (("gemm1_weights", gemm1_weights), ("gemm2_weights", gemm2_weights)):
+ assert tensor.shape[0] == local, (
+ f"{name} holds {tensor.shape[0]} experts but the window has {local}; pass the "
+ "window's own weights"
+ )
+
+
+fp4_block_scale_moe_renorm_routed = Fp4BlockScaleMoeRenormRouted()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_moe_renorm_cute_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_moe_renorm_cute_dsl.py
new file mode 100644
index 000000000000..1d7cf6041e65
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_moe_renorm_cute_dsl.py
@@ -0,0 +1,314 @@
+# 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.*
+import tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops as _cdo
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_block_scale_moe_renorm_routed import (
+ CERTIFIED_PAIRS,
+ _ckpt,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_block_scale_moe_renorm_routed import (
+ fp4_block_scale_moe_renorm_routed as _routed,
+)
+from tensorrt_llm._torch.utils import ActivationType
+
+from ...kernels.moe.blockscaled_contiguous_grouped_gemm_finalize_fusion import (
+ Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel,
+)
+
+#: Token tile of the grouped GEMMs (the kernels' M tile).
+TILE_TOKENS = 128
+#: The GEMM tactics, pinned: `((mma M, mma N), cluster (M, N), raster along M)`. FC1's is what
+#: the engine's autotune warmup picks at every prefill bucket; FC2's likewise, on the finalize
+#: kernel run with a K-tile of 128 (`_FC2_MMA_INST_TILE_K`).
+FC1_TACTIC = ((128, 256), (1, 1), False)
+FC2_TACTIC = ((128, 256), (1, 2), False)
+#: MMA instructions per FC2 A/B pipeline stage (K 128): five stages beside the C staging
+#: buffer instead of two at the kernel's default K 256, so FC2's short K (640) pipelines.
+_FC2_MMA_INST_TILE_K = 2
+#: Launches of at least this many rows run FC2 with L2 cache hints (operands EVICT_FIRST,
+#: scatter-add EVICT_LAST): a gain at ~16K rows, a small loss at 8K (GB300 micro-bench).
+L2_HINT_MIN_ROWS = 12288 # capacity: fixed: dispatch threshold
+_FP4 = torch.float4_e2m1fn_x2
+_SWIGLU_LIMIT = _cdo._canonicalize_swiglu_limit_scalar(_cdo.SWIGLU_LIMIT_SCALAR_DISABLED)
+
+
+class _Fc2Kernel(Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel):
+ def __init__(self, *args, **kwargs):
+ super().__init__(
+ *args, mma_inst_tile_k=_FC2_MMA_INST_TILE_K, row_owned_epilogue=True, **kwargs
+ )
+
+
+class _Fc2HintKernel(Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel):
+ def __init__(self, *args, **kwargs):
+ super().__init__(
+ *args,
+ mma_inst_tile_k=_FC2_MMA_INST_TILE_K,
+ row_owned_epilogue=True,
+ l2_hints=True,
+ **kwargs,
+ )
+
+
+class _Fc2Runner(_cdo.Sm100BlockScaledContiguousGroupedGemmFinalizeFusionRunner):
+ kernel_class = _Fc2Kernel
+ kernel_cache = dict()
+ tuning_config_cache = dict()
+
+
+class _Fc2HintRunner(_cdo.Sm100BlockScaledContiguousGroupedGemmFinalizeFusionRunner):
+ kernel_class = _Fc2HintKernel
+ kernel_cache = dict()
+ tuning_config_cache = dict()
+
+
+_RUNNERS: dict = {}
+
+
+def _runners(num_experts: int, top_k: int) -> tuple:
+ """`(fc1, fc2, fc2 with L2 hints)` for a geometry, built once (they hold compiled kernels)."""
+ key = (num_experts, top_k)
+ if key not in _RUNNERS:
+ _RUNNERS[key] = (
+ _cdo.Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner(
+ num_experts,
+ top_k,
+ num_experts,
+ 0,
+ TILE_TOKENS,
+ 16,
+ activation_type=ActivationType.Swiglu,
+ swiglu_limit_scalar=_SWIGLU_LIMIT,
+ ),
+ _Fc2Runner(num_experts, top_k, num_experts, 0, TILE_TOKENS, torch.bfloat16, 16),
+ _Fc2HintRunner(num_experts, top_k, num_experts, 0, TILE_TOKENS, torch.bfloat16, 16),
+ )
+ return _RUNNERS[key]
+
+
+class Fp4MoeRenormCuteDsl(OpWrapper):
+ """One NVFP4 MoE layer whose routing happens inside the call, on the CuTe
+ DSL MoE kernels. Same function as `fp4_block_scale_moe_renorm_routed`:
+ per token the `top_k` largest logits (ties to the LOWER expert index),
+ weighted by an fp32 softmax over just those `top_k` rounded to bf16; the
+ experts' SwiGLU FC1, the NVFP4 requantization of its output, FC2, and the
+ weighted sum. Returns the combined `[T, H]` bf16 rows.
+
+ Operands, in this kernel family's layout (see `weights.py`): `fc1_w`
+ `[E, 2I, H/2]` uint8, `[up ; gate]` rows interleaved in 64-row groups;
+ `fc1_s` `[E, 2I, H/16]` e4m3 bytes in that row order, 128x4-swizzled;
+ `fc2_w` `[E, H, I/2]` as stored; `fc2_s` `[E, H, I/16]` swizzled.
+ `hidden_states` / `hidden_states_scale` are the NVFP4 activation with
+ **linear** block scales, quantized with the layer's FC1 activation global
+ scale. Scalars: `fc1_alpha` `[E]` dequantizes the FC1 accumulator (the
+ routed entry's `output1_scale_gate_scalar`); `fc2_input_scale` `[1]` is the
+ FC2 activation's quantization multiplier (the reciprocal of its dequant
+ scale), one number for every expert; `fc2_alpha` `[E]` dequantizes the FC2
+ accumulator (the routed entry's `output2_scale_scalar`).
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="4096 rows: where the target switches to this entry; the sort's cluster kernel",
+ spec=_ckpt(4096),
+ ),
+ Cell(
+ why="16377 rows: a con=128 mixed step (16256 context + 121 generation)",
+ spec=_ckpt(16377),
+ ),
+ Cell(
+ why="16384 rows: max_num_tokens of the tuned config; the sort's cooperative kernel",
+ spec=_ckpt(16384),
+ ),
+ )
+
+ note = """
+ Semantics against `fp4_block_scale_moe_renorm_routed` (trtllm-gen) on the
+ same experts: the same routing (both are the `renorm_moe_routing_op`
+ function, weights rounded to bf16 once) and the same NVFP4 arithmetic per
+ expert (e2m1 x e4m3-block x fp32-alpha products accumulated in fp32, SwiGLU
+ on the FC1 accumulator, the FC2 input requantized to NVFP4 with 16-element
+ blocks), but not bitwise: the GEMMs' accumulation order differs, and the
+ weighted sum over a token's experts is a scatter-add with fp32 atomics, so
+ its order -- and the output's last bit -- varies between launches. Both sit
+ inside the routed entry's band against its fp32 reference.
+
+ One FC2 activation scale for all experts: the routed entry takes it per
+ expert (folded into `output1_scale_scalar`); a checkpoint whose experts
+ disagree on `down_proj.input_scale` cannot use this entry as is.
+
+ Kernels (all TRT-LLM's): `trtllm::moe_topk_sort` (routing + tile metadata),
+ `trtllm::moe_output_memset_inplace` (zeroes the rows the scatter-add
+ touches; on `memset_stream` beside FC1 when given), the gathered FC1 runner
+ and the finalize-fused FC2 runner, both with pinned tactics (no AutoTuner
+ lookup); FC2's kernel runs with a K-tile of 128, a row-owned epilogue and,
+ from `L2_HINT_MIN_ROWS` rows, L2 cache hints. Not CUDA-graph certified
+ (prefill steps run eagerly or under the breakable prefill graph).
+
+ Hard limits
+ T certified 4096..16384 (the target calls it from
+ `_CUTE_DSL_MOE_MIN_ROWS` = 4096 rows)
+ num_experts 512, top_k 10 (the routed entry's CERTIFIED_PAIRS)
+ """
+
+ def __call__(
+ self,
+ routing_logits: torch.Tensor,
+ hidden_states: torch.Tensor,
+ hidden_states_scale: torch.Tensor,
+ fc1_w: torch.Tensor,
+ fc1_s: torch.Tensor,
+ fc1_alpha: torch.Tensor,
+ fc2_w: torch.Tensor,
+ fc2_s: torch.Tensor,
+ fc2_alpha: torch.Tensor,
+ fc2_input_scale: torch.Tensor,
+ num_experts: int,
+ top_k: int,
+ memset_stream: torch.cuda.Stream | None = None,
+ ) -> torch.Tensor:
+ tokens = hidden_states.shape[0]
+ hidden = hidden_states.shape[1] * 2
+ fc1, fc2, fc2_hint = _runners(num_experts, top_k)
+ if tokens >= L2_HINT_MIN_ROWS:
+ fc2 = fc2_hint
+ # Renormalize routing and the tile sort in one launch: the same (expert, weight) pairs
+ # (weights rounded to bf16 once) and slot order as `renorm_moe_routing_op` + `moe_sort`.
+ t2e, t2lim, e2p, p2e, _, live, weights = torch.ops.trtllm.moe_topk_sort(
+ routing_logits, None, num_experts, top_k, None, None, 0, num_experts, None,
+ TILE_TOKENS, 1,
+ ) # fmt: skip
+ scales = weights.float()
+ out = torch.empty((tokens, hidden), dtype=torch.bfloat16, device=hidden_states.device)
+
+ def memset() -> None:
+ 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,
+ )
+
+ # The memset of the rows FC2 scatter-adds into runs beside FC1 on `memset_stream` (a
+ # stream nothing else is queued on: the join before FC2 waits for all of it). `out` is
+ # a main-stream allocation the side stream is done with before main's next use of it
+ # (the join), so its block needs no `record_stream` -- which would make the caching
+ # allocator poll an event on every later allocation (~2000 cudaEventQuery a mixed step).
+ main = torch.cuda.current_stream()
+ if memset_stream is not None:
+ memset_stream.wait_stream(main)
+ with torch.cuda.stream(memset_stream):
+ memset()
+ h, h_sf = fc1.forward(
+ [
+ hidden_states.view(_FP4),
+ fc1_w.view(_FP4),
+ hidden_states_scale.view(torch.uint8).view(tokens, hidden // 16),
+ fc1_s.view(torch.uint8),
+ fc1_alpha,
+ t2e,
+ t2lim,
+ p2e,
+ live,
+ fc2_input_scale,
+ ],
+ tactic=FC1_TACTIC,
+ )
+ if memset_stream is not None:
+ main.wait_stream(memset_stream)
+ else:
+ memset()
+ fc2.forward(
+ [
+ h.view(_FP4),
+ fc2_w.view(_FP4),
+ h_sf.view(torch.uint8),
+ fc2_s.view(torch.uint8),
+ fc2_alpha,
+ out,
+ t2e,
+ t2lim,
+ p2e,
+ live,
+ scales,
+ ],
+ tactic=FC2_TACTIC,
+ )
+ 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,
+ memset_stream: torch.cuda.Stream | None = None,
+ ) -> 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/ops/moe/fp8_block_scale_moe_runner.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp8_block_scale_moe_runner.py
new file mode 100644
index 000000000000..392908f43983
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp8_block_scale_moe_runner.py
@@ -0,0 +1,434 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""DeepSeek-style FP8 block-scale mixture-of-experts layer: given top-k +
+grouped FC1 GEMM + SwiGLU + FP8 requantization + grouped FC2 GEMM +
+routing-weighted combine, in one trtllm-gen call.
+
+Called by the qwen4_exp target's MTP layer at every layout."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import (
+ Arch,
+ Cell,
+ OpWrapper,
+ true_fp32_matmul,
+)
+
+#: Block edge on every scale axis: 1x128 activations, 128x128 weights.
+BLOCK = 128
+
+#: Largest finite e4m3 magnitude.
+E4M3_MAX = 448.0
+
+#: The activation kernel's floor on a block amax (`kDsActAmaxEpsilon` in
+#: trtllmGenKernels/blockScaleMoe/DevKernel.cu). Only an all-zero block reaches it.
+ACT_AMAX_FLOOR = 1.0e-10
+
+
+def dequantize_block_weight(codes: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
+ """One expert's `[M, K]` fp32 weight from e4m3 codes and `[M/128, K/128]` scales."""
+ m, k = codes.shape
+ blocks = codes.float().reshape(m // BLOCK, BLOCK, k // BLOCK, BLOCK)
+ return (blocks * scales.float().reshape(m // BLOCK, 1, k // BLOCK, 1)).reshape(m, k)
+
+
+def requantize_1x128(x: torch.Tensor) -> torch.Tensor:
+ """The e4m3 round trip the kernel applies to both intermediates.
+
+ Per row and per 128 consecutive columns: `scale = max(amax, 1e-10) / 448`,
+ then round-to-nearest-even into e4m3 saturating at +-448. The plain
+ `amax / 448` scale, not the round-up power of two the MX members use.
+ """
+ shape = x.shape
+ blocks = x.reshape(*shape[:-1], shape[-1] // BLOCK, BLOCK)
+ scale = blocks.abs().amax(dim=-1, keepdim=True).clamp(min=ACT_AMAX_FLOOR) / E4M3_MAX
+ codes = (blocks / scale).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn)
+ return (codes.float() * scale).reshape(shape)
+
+
+def _geom(hidden, inter, num_experts, top_k, offset=0, local=None):
+ """A cell's MoE geometry; the window defaults to every expert local."""
+ return dict(
+ hidden=hidden,
+ inter=inter,
+ num_experts=num_experts,
+ top_k=top_k,
+ local_expert_offset=offset,
+ local_num_experts=num_experts if local is None else local,
+ )
+
+
+#: The MTP draft layer's routed experts: H 2560, I 640, 512 experts, top-10,
+#: every expert local (tp1).
+_CKPT = _geom(2560, 640, 512, 10)
+
+
+def _window(offset: int) -> dict:
+ """One rank's 128-expert window of the 512 (the qwen4 default target at tp4,
+ tep4 and dep4): the router still ranks all 512, so `num_experts` stays 512."""
+ return _geom(2560, 640, 512, 10, offset=offset, local=128)
+
+
+class Fp8BlockScaleMoeRunner(OpWrapper):
+ """One FP8 block-scale MoE layer over e4m3 activations, entered pre-routed.
+
+ Per token `t` and slot `j`, with `A[t]` the fp32 activations
+ (`hidden_states[t, k] * hidden_states_scale[k // 128, t]`) and `W_*` the
+ dequantized 128x128-block weights of local expert `e = topk_ids[t, j] -
+ local_expert_offset` (slots outside `[0, local_num_experts)` contribute
+ nothing):
+
+ up, gate = requant(A[t] @ [W_up; W_gate].T).split(I)
+ gate, up = min(gate, lim), clamp(up, -lim, lim) # gemm1_clamp_limit
+ out[t] += topk_weights[t, j] * (requant(silu(gate) * up) @ W_down.T)
+
+ accumulated in fp32 and stored bf16, `requant` being `requantize_1x128`.
+ FC1 rows are `[up ; gate]` -- up first, the engine loader's `[w3; w1]`.
+
+ Fusion boundary: permutation, both GEMMs with on-the-fly dequantization,
+ both intermediate requantizations, the clamped SwiGLU and the weighted
+ combine are inside. The router GEMM and top-k, quantizing the bf16 hidden
+ states (`fp8_quantize_1x128(x, use_ue8m0=False)`), the shared expert, any
+ all-reduce and the residual are the caller's.
+
+ `hidden_states_scale` is `[H/128, T]` -- transposed, exactly what
+ `fp8_quantize_1x128` returns on this architecture. Pass it straight through:
+ at `T == H/128` (20 tokens at H 2560) its transpose has the same shape and
+ passes every check, and is 41-43 ulp RMS wrong.
+
+ Returns a fresh `[T, H]` bf16 tensor, or an empty `[0]` bf16 tensor when
+ `output` is given (the result is then in `output`).
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ # The token column at the checkpoint geometry, the 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="20 tokens == H/128: the activation scale buffer is square",
+ spec=dict(tokens=20, **_CKPT),
+ ),
+ Cell(
+ why="16384 rows: max_num_tokens of the tuned config; past the last tuning bucket",
+ spec=dict(tokens=16384, **_CKPT),
+ ),
+ # The qwen4 default target: every member runs the MTP experts as four
+ # 128-wide windows, one per rank -- at tp4 too, where 640 / 4 = 160 is not
+ # a multiple of 128 and the intermediate cannot be split. One cell per
+ # window offset, each at a row count the MTP layer reaches.
+ Cell(
+ why="default MTP drafter decode: one token, rank 0's 128-expert window",
+ spec=dict(tokens=1, **_window(0)),
+ ),
+ Cell(
+ why="default: 128 drafter rows, the largest captured batch; rank 1's window",
+ spec=dict(tokens=128, **_window(128)),
+ ),
+ Cell(
+ why="default: 1024 rows, the first draft over mtp7's 128 x 8 verify rows; window 2",
+ spec=dict(tokens=1024, **_window(256)),
+ ),
+ Cell(
+ why="default: 8192 rows, a max_num_tokens context chunk; the last window (384)",
+ spec=dict(tokens=8192, **_window(384)),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none. Two process-global caches sit behind the call --
+ the autotuner's profiling cache and the C++ runner cache -- and a warm
+ tuned tactic is bitwise equal to the cold fallback (the test holds this at
+ 256 tokens, for the whole stack and for a 128-expert window, which tunes
+ under its own `local_num_experts` key).
+
+ The op is a Python `torch.library.custom_op`
+ (`trtllm::fp8_block_scale_moe_runner` in
+ `_torch/custom_ops/trtllm_gen_custom_ops.py`) over the compiled-in C++ class
+ `torch.classes.trtllm.FP8BlockScaleMoERunner`
+ (`cpp/tensorrt_llm/thop/moe/trtllmGen/fp8BlockScaleMoe.cpp`), which launches
+ trtllm-gen's precompiled batched-GEMM cubins: nothing is JIT-compiled.
+
+ Preconditions. `hidden_states` is `[T, H]` e4m3 and `hidden_states_scale`
+ `[H/128, T]` fp32 (see the class docstring); weights are the checkpoint's
+ per-expert tensors in plain row-major layout, no shuffle and no swizzle:
+ `gemm1_weights` `[E, 2I, H]` e4m3 rows `[up ; gate]`, `gemm1_weights_scale`
+ `[E, 2I/128, H/128]` fp32, `gemm2_weights` `[E, H, I]`, `gemm2_weights_scale`
+ `[E, H/128, I/128]`; `topk_ids` `[T, top_k]` int32 and `topk_weights`
+ `[T, top_k]` bf16, used as given (nothing is renormalized). H and I are
+ multiples of 128, `intermediate_size == gemm1_weights.shape[1] // 2`,
+ `0 < top_k < num_experts`, `top_k <= 32`. The op raises on each of those.
+ Supply distinct ids per row: a repeated id counted once at 16 tokens and
+ twice at 256 (the token count picks the routing kernel).
+
+ The window: local slot i answers for global id `local_expert_offset + i`,
+ and `num_experts` stays the routing space (512) on every window call. The
+ weights and all four scale tensors hold the window's experts only: the op
+ compares each one's dim 0 with `local_num_experts` and raises on the whole
+ stack (the test holds all four). A token with no slot in the window comes
+ back bitwise zero (every window cell holds this).
+
+ The qwen4 default target's four 128-wide windows of 512 are certified at the
+ row counts its MTP layer reaches, one window offset per cell, on sm_103: 1
+ to 8192 rows, at most 7.95 bf16 ulp per element of the row max and 0.227
+ ulp RMS, inside the test's 16 / 0.75 gate. Each cell asserts some slot routes
+ into its window. The four windows' bf16 outputs added in fp32 -- what the
+ target's all-reduce or reduce-scatter adds -- sit at most 4.41 / 0.297 ulp
+ from the 512-expert reference at 1, 128 and 1024 rows, and the sum with any
+ one window left out fails the gate. Not bitwise to the 512-expert call (each window rounds its
+ partial to bf16): at most 1.97 / 0.187 ulp from it.
+
+ `num_experts` passed as the window's own count (128) puts windows 1..3 past
+ the routing space: at small row counts the op answers all zeros with no
+ error, and at 16 rows it reads out of bounds (an illegal memory access) --
+ ids `>= num_experts` are dropped only at small row counts. `is_valid`
+ refuses a window that starts outside `[0, num_experts)`; window 0 of that
+ mistake still passes it, and ids past `num_experts` cannot be seen without
+ reading the ids. A window that starts inside the routing space and runs
+ past its end ([448, 576) of 512) computes its in-range experts correctly.
+
+ Inert on the certified path, bitwise: `routing_method_type` 0/1/2/4 (the
+ pre-routed path forces the routing postprocess), `n_group`/`topk_group`
+ (1, 1), `routed_scaling_factor` None/1.0/2.5, `tune_max_num_tokens`,
+ `use_dp`. `gemm1_clamp_limit` is live and one scalar for every expert,
+ applied after the FC1 requantization; None is bitwise equal to omitting it.
+
+ Silent on the raw op, refused by `is_valid` (the test holds the strided
+ `hidden_states` and the outside-the-routing-space refusals): any
+ non-contiguous tensor (read as dense, 55-115 ulp RMS off); `act_type != 0`
+ (the Python runner forces SwiGlu, so 1 and 2 are bitwise identical to 0);
+ `routing_bias` alongside `topk_ids` (ignored; a +-1e3 bias leaves the output
+ bitwise unchanged); a `local_expert_offset` outside `[0, num_experts)`.
+ `num_fused_shared_experts > 0` is rejected loudly on the pre-routed path,
+ which is the only path the C++ runner's small-tile tactic exclusion guards.
+
+ The registered fake returns `[T, 2H]` (it assumes packed fp4 hidden
+ states): an upstream defect, eager is unaffected.
+
+ Hard limits
+ T whole stack certified at 1, 20 and 16384 (max_num_tokens of the
+ tuned config); 128-expert windows at 1, 128, 1024 and 8192
+ (the default target's max_num_tokens chunk). Chunk above
+ them, not probed. Autotuner buckets stop at
+ `tune_max_num_tokens` (default 8192, which the target
+ leaves): past it a call runs the 8192 bucket's tactic; a
+ larger value makes the op append a bucket at that size
+ top_k <= 32 (raises past it); certified 10, the checkpoint's
+ E, H, I certified at the checkpoint's 512, 2560, 640
+ window the whole stack, and 128 wide at offsets 0, 128, 256 and 384 of 512
+
+ Not certified (measured, but not re-driven by the test): windows of other
+ widths (a 256-wide dep2 window, say) or offsets, a window past the routing
+ space's end, repeated or
+ out-of-range ids, `gemm1_clamp_limit`, the inert arguments, `output=`, the
+ registered fake, the op's other loud rejections. Nor: CUDA-graph capture
+ of the call (the cells run eager), the in-kernel routing path
+ (`routing_logits` given), multi-rank execution (every window here runs on
+ one device; the cross-rank sum is the collective's), geometries other than
+ the checkpoint's, sm_100.
+ """
+
+ def __call__(
+ self,
+ routing_logits: Optional[torch.Tensor],
+ routing_bias: Optional[torch.Tensor],
+ hidden_states: torch.Tensor,
+ hidden_states_scale: torch.Tensor,
+ gemm1_weights: torch.Tensor,
+ gemm1_weights_scale: torch.Tensor,
+ gemm2_weights: torch.Tensor,
+ gemm2_weights_scale: torch.Tensor,
+ num_experts: int,
+ top_k: int,
+ num_fused_shared_experts: Optional[int],
+ n_group: Optional[int],
+ topk_group: Optional[int],
+ intermediate_size: int,
+ local_expert_offset: int,
+ local_num_experts: int,
+ routed_scaling_factor: Optional[float],
+ routing_method_type: int,
+ topk_weights: Optional[torch.Tensor] = None,
+ topk_ids: Optional[torch.Tensor] = None,
+ act_type: int = 0,
+ gemm1_clamp_limit: Optional[float] = None,
+ output: Optional[torch.Tensor] = None,
+ tune_max_num_tokens: int = 8192,
+ use_dp: bool = False,
+ ) -> torch.Tensor:
+ return torch.ops.trtllm.fp8_block_scale_moe_runner(
+ routing_logits,
+ routing_bias,
+ hidden_states,
+ hidden_states_scale,
+ gemm1_weights,
+ gemm1_weights_scale,
+ gemm2_weights,
+ gemm2_weights_scale,
+ num_experts,
+ top_k,
+ num_fused_shared_experts,
+ n_group,
+ topk_group,
+ intermediate_size,
+ local_expert_offset,
+ local_num_experts,
+ routed_scaling_factor,
+ routing_method_type,
+ topk_weights=topk_weights,
+ topk_ids=topk_ids,
+ act_type=act_type,
+ gemm1_clamp_limit=gemm1_clamp_limit,
+ output=output,
+ tune_max_num_tokens=tune_max_num_tokens,
+ use_dp=use_dp,
+ )
+
+ def reference(
+ self,
+ routing_logits: Optional[torch.Tensor],
+ routing_bias: Optional[torch.Tensor],
+ hidden_states: torch.Tensor,
+ hidden_states_scale: torch.Tensor,
+ gemm1_weights: torch.Tensor,
+ gemm1_weights_scale: torch.Tensor,
+ gemm2_weights: torch.Tensor,
+ gemm2_weights_scale: torch.Tensor,
+ num_experts: int,
+ top_k: int,
+ num_fused_shared_experts: Optional[int],
+ n_group: Optional[int],
+ topk_group: Optional[int],
+ intermediate_size: int,
+ local_expert_offset: int,
+ local_num_experts: int,
+ routed_scaling_factor: Optional[float],
+ routing_method_type: int,
+ topk_weights: Optional[torch.Tensor] = None,
+ topk_ids: Optional[torch.Tensor] = None,
+ act_type: int = 0,
+ gemm1_clamp_limit: Optional[float] = None,
+ output: Optional[torch.Tensor] = None,
+ tune_max_num_tokens: int = 8192,
+ use_dp: bool = False,
+ ) -> torch.Tensor:
+ """The layer over bit-identical operands, fp32 GEMMs, bf16 store.
+
+ Pre-routed only: `routing_logits` must be None. Ignored: `routing_bias`,
+ `num_experts`, `top_k`, `num_fused_shared_experts`, `n_group`,
+ `topk_group`, `routed_scaling_factor`, `routing_method_type`,
+ `act_type`, `output`, `tune_max_num_tokens`, `use_dp` -- each inert on
+ that path. Returns the `[T, H]` result even when `output` is given.
+ """
+ assert routing_logits is None and topk_ids is not None and topk_weights is not None, (
+ "the reference models the pre-routed path only"
+ )
+ num_tokens, hidden = hidden_states.shape
+ nb = hidden // BLOCK
+ act_in = (
+ hidden_states.float().reshape(num_tokens, nb, BLOCK)
+ * hidden_states_scale.t().reshape(num_tokens, nb, 1).float()
+ ).reshape(num_tokens, hidden)
+ out = torch.zeros(num_tokens, hidden, dtype=torch.float32, device=hidden_states.device)
+ local = topk_ids.long() - local_expert_offset
+ live = local[(local >= 0) & (local < local_num_experts)]
+ with true_fp32_matmul():
+ for e in torch.unique(live).tolist():
+ tok, slot = (local == e).nonzero(as_tuple=True)
+ w1 = dequantize_block_weight(gemm1_weights[e], gemm1_weights_scale[e])
+ y1 = requantize_1x128(act_in[tok] @ w1.t())
+ up, gate = y1[:, :intermediate_size], y1[:, intermediate_size:]
+ if gemm1_clamp_limit is not None:
+ gate = gate.clamp(max=gemm1_clamp_limit)
+ up = up.clamp(-gemm1_clamp_limit, gemm1_clamp_limit)
+ act = requantize_1x128(torch.nn.functional.silu(gate) * up)
+ w2 = dequantize_block_weight(gemm2_weights[e], gemm2_weights_scale[e])
+ y = act @ w2.t()
+ out.index_add_(0, tok, y * topk_weights[tok, slot].float().unsqueeze(1))
+ return out.to(torch.bfloat16)
+
+ def is_valid(
+ self,
+ routing_logits: Optional[torch.Tensor],
+ routing_bias: Optional[torch.Tensor],
+ hidden_states: torch.Tensor,
+ hidden_states_scale: torch.Tensor,
+ gemm1_weights: torch.Tensor,
+ gemm1_weights_scale: torch.Tensor,
+ gemm2_weights: torch.Tensor,
+ gemm2_weights_scale: torch.Tensor,
+ num_experts: int,
+ top_k: int,
+ num_fused_shared_experts: Optional[int],
+ n_group: Optional[int],
+ topk_group: Optional[int],
+ intermediate_size: int,
+ local_expert_offset: int,
+ local_num_experts: int,
+ routed_scaling_factor: Optional[float],
+ routing_method_type: int,
+ topk_weights: Optional[torch.Tensor] = None,
+ topk_ids: Optional[torch.Tensor] = None,
+ act_type: int = 0,
+ gemm1_clamp_limit: Optional[float] = None,
+ output: Optional[torch.Tensor] = None,
+ tune_max_num_tokens: int = 8192,
+ use_dp: bool = False,
+ ) -> None:
+ """Four inputs the op accepts and answers wrongly, each measured.
+
+ The transposed activation-scale buffer at `T == H/128` is the fifth,
+ and no metadata check can see it.
+ """
+ # The kernel takes raw data pointers: a strided view is read as dense.
+ for name, tensor in (
+ ("routing_logits", routing_logits),
+ ("routing_bias", routing_bias),
+ ("hidden_states", hidden_states),
+ ("hidden_states_scale", hidden_states_scale),
+ ("gemm1_weights", gemm1_weights),
+ ("gemm1_weights_scale", gemm1_weights_scale),
+ ("gemm2_weights", gemm2_weights),
+ ("gemm2_weights_scale", gemm2_weights_scale),
+ ("topk_weights", topk_weights),
+ ("topk_ids", topk_ids),
+ ("output", output),
+ ):
+ if tensor is not None:
+ assert tensor.is_contiguous(), (
+ f"{name} must be contiguous; a strided view is read as if dense "
+ "and silently produces wrong results"
+ )
+ # The Python runner overwrites the activation with SwiGlu.
+ assert act_type == 0, (
+ f"act_type={act_type} is silently ignored -- this op always computes "
+ "SwiGLU. Use a different runner for a non-gated activation"
+ )
+ # Once topk_ids carry the routing, nothing reads the bias.
+ if routing_bias is not None:
+ assert topk_ids is None, (
+ "routing_bias is ignored when topk_ids/topk_weights are given "
+ "(routing has already happened); fold it into the router logits"
+ )
+ # `num_experts` is the routing space, not the window: a window that
+ # starts at or past it holds no routable id. Passing the window's own
+ # count as `num_experts` puts ranks 1..3 there: the op answers all zeros
+ # with no error at small row counts and reads out of bounds at 16 rows.
+ assert 0 <= local_expert_offset < num_experts, (
+ f"local_expert_offset={local_expert_offset} is outside the routing space "
+ f"[0, num_experts={num_experts}); num_experts is every routed expert "
+ "(512 for a 128-expert window of 512), not the window's count"
+ )
+
+
+fp8_block_scale_moe_runner = Fp8BlockScaleMoeRunner()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fused_sigmoid_gate_mul_add.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fused_sigmoid_gate_mul_add.py
new file mode 100644
index 000000000000..c6ac057e9432
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fused_sigmoid_gate_mul_add.py
@@ -0,0 +1,171 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Shared-expert combine: final_hidden_states + sigmoid(gate) * shared_expert_output."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+# The module callable is the surface: outside torch.compile its body is one Triton
+# launch after argument checks. `torch.ops.trtllm.fused_sigmoid_gate_mul_add_` (same
+# module) is only its torch.compile boundary, which the callable takes itself when
+# compiling, with a different signature (mutated `output`, separate addend).
+from tensorrt_llm._torch.moe.fused_shared_expert import (
+ fused_sigmoid_gate_mul_add as _fused_sigmoid_gate_mul_add,
+)
+
+
+def _ckpt(tokens: int, **extra) -> dict:
+ """The checkpoint's MoE combine: hidden 2560 bf16, one gate logit per token, in place."""
+ spec = dict(
+ tokens=tokens,
+ hidden=2560,
+ dtype=torch.bfloat16,
+ gate_dtype=torch.bfloat16,
+ gate_layout="column",
+ dest="inplace",
+ )
+ spec.update(extra)
+ return spec
+
+
+class FusedSigmoidGateMulAdd(OpWrapper):
+ """`out[t, i] = final[t, i] + sigmoid(gate[t]) * shared[t, i]`, fp32 inside, rounded once.
+
+ Wraps `tensorrt_llm._torch.moe.fused_shared_expert.fused_sigmoid_gate_mul_add`
+ (one Triton kernel, `_sigmoid_gate_mul_add_kernel`).
+
+ `gate_logits` is one pre-sigmoid scalar per token read as a flat vector of
+ `num_tokens`: `[T, 1]`, `[T]`, `[1, T]` and a strided column are the same
+ call, bit for bit, and its dtype is independent of the activations'.
+
+ Destination: with `output=None` the result overwrites `final_hidden_states`
+ and that tensor is returned; otherwise `output` is written and returned and
+ both inputs are untouched. Never a fresh tensor.
+
+ Fusion boundary: sigmoid, broadcast, multiply and add. The routed MoE, the
+ shared expert and the `[hidden, 1]` gate projection are the caller's. The
+ single rounding matters: rounding the product to bf16 before the add moves
+ the output by up to 0.0312 at [17, 2560], outside the bf16 tolerance.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="_moe_block / MTPLayer._moe decode: one token, in place; 8 warps",
+ spec=_ckpt(1),
+ ),
+ 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 "
+ "GEMM leaves it",
+ spec=_ckpt(33, gate_layout="strided"),
+ ),
+ # The default target's call sites, the decode PLE-layer combine and the MTP layer's
+ # MoE, on the gate the shared expert's fused GEMM leaves: column 1280 of the
+ # replicated [T, 1288] `gug` output.
+ Cell(
+ why="default decode PLE-layer / MTP MoE combine: one token, strided gate",
+ spec=_ckpt(1, gate_layout="strided"),
+ ),
+ Cell(
+ why="default tp4/tep4 decode batch at the largest captured size: 128 rows",
+ spec=_ckpt(128, gate_layout="strided"),
+ ),
+ Cell(
+ why="default tep4 MTP-7 at batch 128: 128 x 8 = 1024 rows, strided gate",
+ spec=_ckpt(1024, gate_layout="strided"),
+ ),
+ Cell(
+ why="the envelope's largest row: max_num_tokens 8192, strided gate",
+ spec=_ckpt(8192, gate_layout="strided"),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none; stateless, no autotuner, no architecture branch.
+ Source: a Triton kernel (`_sigmoid_gate_mul_add_kernel`) JIT-compiled on first
+ call for the running device; the cells ran on sm_103.
+
+ Call sites in the default target: the decode PLE-layer combine
+ (`DecodeTarget.forward`) and the MTP layer's MoE (`DecodeTarget.mtp_forward`),
+ in place on the routed output. The shared expert is replicated at every member,
+ so the gate is column 1280 of this rank's `[T, 1288]` `gug` output (gate|up 2 x
+ 640 plus 8 gate-logit rows) and hidden is 2560 everywhere. Rows are this rank's
+ tokens: decode 1..128 at tp4 / tep4 (1..32 a dep4 rank's own), MTP-7 at batch
+ 128 up to 128 x 8 = 1024. Every row count takes one path: hidden 2560 gives
+ one 4096-wide block and 8 warps per row, grid `(T, 1)`.
+
+ Preconditions, all enforced loudly by the callable (AssertionError before
+ any launch; the test holds the wrong-length gate): `final_hidden_states`
+ 2-D; `shared_expert_output` and `output` match it in shape and dtype;
+ `gate_logits.numel() == num_tokens`; the destination has `stride(-1) == 1`
+ (its row stride is free). A column-strided `final_hidden_states` or
+ `shared_expert_output` is densified inside the call (an extra copy kernel),
+ as is a non-contiguous gate. `num_tokens == 0` returns the destination
+ without launching (certified).
+
+ Numerics: the default `assert_close` band holds for bf16, the dtype the target
+ calls it with; fp16 and fp32 run the same kernel and are not certified.
+ bf16 at hidden 2560 is bitwise equal at 1-2 tokens and >= 99.993% bitwise
+ at 17-8192 tokens, one ulp otherwise -- `tl.sigmoid`'s fast exponential.
+ The bf16 band refuses the wrong variants: no sigmoid, multiply instead of
+ add, gate on the routed half, gate dropped, one gate for the whole batch.
+ At the default cells, wherever |out| >= 2^-6 each difference from
+ `reference` is exactly one bf16 ulp of the output; the rest sit where the
+ sum cancels toward zero, the largest 1.6e-2 absolute at 8192 rows (one ulp
+ of a [2, 4) value). Against a float64 reference the relative L2 is 1.66e-3
+ at every cell, equal to `reference`'s own (bf16 output rounding). A strided
+ `gug` column and the same logits as a dense `[T, 1]` are bitwise one call
+ (held at 1024 rows). Graph-capturable: one capture at the 128-row site,
+ replayed on refilled operands, matches `reference` (held).
+
+ Hard limits
+ num_tokens largest run 8192 (max_num_tokens). The row offset is int64 in
+ the kernel and column offsets are int32 within one row, so no
+ element index overflows int32.
+
+ Not certified: hidden sizes other than 2560 (the column block and warp count
+ step with it), gate layouts other than `[T, 1]` and the strided column, an
+ `output` buffer (including one aliasing `shared_expert_output`), row-padded
+ operands, host operands, torch.compile, operands on two devices, float8
+ activations, gate dtypes other than bf16, rows past 8192.
+ """
+
+ def __call__(
+ self,
+ final_hidden_states: torch.Tensor,
+ gate_logits: torch.Tensor,
+ shared_expert_output: torch.Tensor,
+ output: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ return _fused_sigmoid_gate_mul_add(
+ final_hidden_states,
+ gate_logits,
+ shared_expert_output,
+ output,
+ )
+
+ def reference(
+ self,
+ final_hidden_states: torch.Tensor,
+ gate_logits: torch.Tensor,
+ shared_expert_output: torch.Tensor,
+ output: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ """A fresh tensor in `final_hidden_states.dtype`; `output` is ignored.
+
+ Takes the operands as they were before the call: in place, the op
+ destroys `final_hidden_states`.
+ """
+ s = torch.sigmoid(gate_logits.reshape(-1, 1).float())
+ out = final_hidden_states.float() + s * shared_expert_output.float()
+ return out.to(final_hidden_states.dtype)
+
+
+fused_sigmoid_gate_mul_add = FusedSigmoidGateMulAdd()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/moe_sort.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/moe_sort.py
new file mode 100644
index 000000000000..a22c0614c1d0
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/moe_sort.py
@@ -0,0 +1,272 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""MoE sort: bucket the routed (token, expert) pairs into per-expert row tiles
+for the CuteDSL grouped GEMMs, in one trtllm call."""
+
+from __future__ import annotations
+
+import torch
+
+import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+
+def _ckpt(tokens: int, tile_tokens_dim: int, dist: str = "router") -> dict:
+ """The checkpoint's trunk MoE: 512 experts, top-10, all local (tp1), fp32 scales.
+
+ `dist` names how the test draws the expert ids: `router` = each row's top-10 of
+ uniform scores over the 512 experts; `tile_edges` = three experts hold exactly
+ tile-1, tile and tile+1 rows, the rest `router`. Both put 10 distinct experts on
+ every row, as the router does.
+ """
+ return dict(
+ tokens=tokens,
+ dist=dist,
+ num_experts=512,
+ top_k=10,
+ local_expert_offset=0,
+ local_num_experts=512,
+ tile_tokens_dim=tile_tokens_dim,
+ scales_dtype=torch.float32,
+ )
+
+
+class MoeSort(OpWrapper):
+ """`torch.ops.trtllm.moe_sort`: routed (token, slot) pairs -> per-expert row tiles.
+
+ Input: `token_selected_experts` `[T, top_k]` int32 global expert ids (the
+ router's output), `token_final_scales` `[T, top_k]` (its weights). Call it
+ with the op's own keyword names, as the CuteDSL MoE does.
+
+ Numbering: an *expanded* index is `t * top_k + k` (row-major over
+ `token_selected_experts`); a *permuted* index is a row of the grouped GEMM's
+ M dimension. Each local expert `e = id - local_expert_offset` owns a
+ contiguous run of `count_e` permuted rows starting on a tile boundary,
+ experts in increasing order, the run padded up to whole tiles of
+ `tile_tokens_dim` rows. Returns a list of six int32 tensors, in this order:
+
+ - `tile_idx_to_expert_idx` `[max_tiles]`: the local expert of tile `i`;
+ - `tile_idx_to_mn_limit` `[max_tiles]`: one past the last real row of tile
+ `i`, as an absolute permuted index: `min((i + 1) * tile, start_e + count_e)`;
+ - `expanded_idx_to_permuted_idx` `[T, top_k]`: the permuted row of each
+ pair, `-1` for a pair whose expert is not local;
+ - `permuted_idx_to_expanded_idx` `[max_tiles * tile]`: the inverse on real rows;
+ - `total_num_padded_tokens` `[1]` = `num_non_exiting_tiles * tile`;
+ - `num_non_exiting_tiles` `[1]`: tiles actually used.
+
+ `max_tiles = m + (T * top_k - m) // tile` with `m = min(local_num_experts,
+ T * top_k)` is computed on the host from the shapes and scalars alone, so
+ every output shape is fixed by `T`, which is what lets a CUDA graph capture
+ it. Only the first `num_non_exiting_tiles` entries of the tile arrays and
+ the rows below each tile's `mn_limit` are written; everything else is left
+ as allocated (see `note`).
+
+ Fusion boundary: sorting only. Routing (the top-k) is the caller's, and so
+ are the permutation of the activations (the FC1 gather does it) and the
+ routing-weighted unpermute; this op does not read `token_final_scales` into
+ any output.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # One cell per kernel path the target reaches at its tile, 128: block (T <= 4),
+ # cluster (T 5..8192) and cooperative (T >= 8193).
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="tile 128: decode, 1 request without MTP; block kernel",
+ spec=_ckpt(1, 128),
+ ),
+ Cell(
+ why="tile 128: expert counts tile-1, tile, tile+1; the mn_limit clamp; cluster kernel",
+ spec=_ckpt(300, 128, "tile_edges"),
+ ),
+ Cell(
+ why="tile 128: 8193 rows; first cooperative-kernel T (prefill)",
+ spec=_ckpt(8193, 128),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none. Stateless; allocates its six outputs plus scratch on
+ every call and runs on the current stream, with no host synchronization, so
+ it captures into a CUDA graph.
+
+ Call it as the engine's CuteDSL MoE does (`run_moe_nvfp4_impl`,
+ `_torch/moe/fused_moe/fused_moe_cute_dsl.py`): `num_experts` = the slot
+ count, `local_expert_offset` / `local_num_experts` = this rank's slots, and
+ `tile_tokens_dim` = the grouped GEMMs' tile, 128 or 256 on sm_103. The
+ downstream FC1 gather GEMM, output memset and FC2 take these six tensors as
+ they are.
+
+ Kernel path, by `T`:
+ T 1..4 routingCustom::routingIndicesBlockKernel
+ T 5..8192 routingCustom::routingIndicesClusterKernel
+ T 8193..471859 routingInitExpertCounts + routingIndicesCoopKernel
+ T >= 471860 routingInitExpertCounts + routingIndicesHistogramKernel
+ + routingIndicesOffsetsKernel
+ The cooperative ceiling is (SMs - 8) * 512 * 64 / top_k rows with 512
+ experts: 471859 on a 152-SM GB300, a different T on another SM count.
+
+ Order within an expert's run is unspecified. The block kernel (T <= 4) lays
+ each expert's rows out in token order and is bitwise repeatable; every other
+ path assigns them by atomics, so `expanded_idx_to_permuted_idx` differs
+ between identical calls. A test therefore compares each expert's rows as a
+ set and the two mappings as exact inverses; the tile arrays and counts are
+ deterministic.
+
+ Not written: the tile arrays past `num_non_exiting_tiles`, and the rows of
+ `permuted_idx_to_expanded_idx` from each tile's `mn_limit` to the tile's
+ end. Eager, those hold whatever the allocator handed back; under a CUDA
+ graph they hold the previous replay's indices, which look valid. Every
+ consumer must bound its reads by `num_non_exiting_tiles` and
+ `tile_idx_to_mn_limit`.
+
+ `token_final_scales`: shape-checked and never written; its values never
+ reach an output. Its pointer must be non-null, which is why T = 0 raises.
+
+ Data hazards `is_valid` cannot see (values): a row naming the same expert
+ twice keeps both pairs on the cluster kernel but loses one on the block
+ kernel (that pair gets -1). An id outside [local_expert_offset,
+ local_expert_offset + local_num_experts) gets -1, except that the cluster
+ kernel reads ids through int16, so an id of 65536 + e sorts as expert e. The
+ router produces neither.
+
+ Loud: int64 ids, 1-D ids, a `[T, top_k]` size mismatch on either input,
+ CPU tensors, top_k above 32, num_experts above 2048, T = 0.
+
+ Silent, refused by `is_valid`: a non-contiguous `token_selected_experts`
+ (read from `data_ptr()` as dense `[T, top_k]`); 2^31 or more elements (the
+ host counts pairs in int32: it raises a negative-dimension error from 2^31
+ and silently sorts only the wrapped remainder from 2^32); `tile_tokens_dim`
+ <= 0 (the host divides by it).
+
+ Certified only at the checkpoint's values: num_experts 512, top_k 10, all
+ experts local (offset 0), tile 128 (the target's), fp32 scales, and T 1,
+ 300 and 8193 (block, cluster and the first cooperative T); CUDA-graph replay
+ on new ids at T = 128; the contiguity guard. The order, unwritten-region,
+ data-hazard and loud-rejection statements above were measured at bring-up
+ and are not re-driven by the test. Expert parallelism, tile 256 and other
+ tiles, other expert counts and the histogram path are uncertified.
+
+ Hard limits (sm_103)
+ T * top_k boundary 2^31 pairs (int32 product)
+ max_tiles * tile boundary 2^31 padded rows (raises past it)
+ expert id int16 in the cluster kernel; ids < 2048 never reach it
+ """
+
+ def __call__(
+ self,
+ token_selected_experts: torch.Tensor,
+ token_final_scales: torch.Tensor,
+ num_experts: int,
+ top_k: int,
+ local_expert_offset: int,
+ local_num_experts: int,
+ tile_tokens_dim: int,
+ ) -> list[torch.Tensor]:
+ return torch.ops.trtllm.moe_sort(
+ token_selected_experts=token_selected_experts,
+ token_final_scales=token_final_scales,
+ num_experts=num_experts,
+ top_k=top_k,
+ local_expert_offset=local_expert_offset,
+ local_num_experts=local_num_experts,
+ tile_tokens_dim=tile_tokens_dim,
+ )
+
+ def reference(
+ self,
+ token_selected_experts: torch.Tensor,
+ token_final_scales: torch.Tensor,
+ num_experts: int,
+ top_k: int,
+ local_expert_offset: int,
+ local_num_experts: int,
+ tile_tokens_dim: int,
+ ) -> list[torch.Tensor]:
+ """Counting sort by local expert, stable in the expanded index.
+
+ Within an expert, rows are in increasing expanded index; the test's compare
+ accepts any order there. What the op leaves unwritten is `-1` here.
+ `token_final_scales`, `num_experts` and `top_k` are not read: the scales
+ never reach an output, and the ids' shape carries the rest.
+ """
+ ids = token_selected_experts
+ dev = ids.device
+ tile = tile_tokens_dim
+ num_tokens, width = ids.shape
+ num_expanded = num_tokens * width
+ local = ids.reshape(-1).long() - local_expert_offset
+ is_local = (local >= 0) & (local < local_num_experts)
+ counts = torch.bincount(local[is_local], minlength=local_num_experts)
+ tiles_per_expert = (counts + tile - 1) // tile
+ num_tiles = int(tiles_per_expert.sum())
+ first = min(local_num_experts, num_expanded)
+ max_tiles = first + (num_expanded - first) // tile
+
+ run_start = (torch.cumsum(tiles_per_expert, 0) - tiles_per_expert) * tile
+ tile_expert = torch.repeat_interleave(
+ torch.arange(local_num_experts, device=dev), tiles_per_expert
+ )
+ tile_idx = torch.arange(num_tiles, device=dev)
+ mn_limit = torch.minimum(
+ (tile_idx + 1) * tile, run_start[tile_expert] + counts[tile_expert]
+ )
+
+ expanded = torch.arange(num_expanded, device=dev)[is_local]
+ expert = local[is_local]
+ order = torch.sort(expert, stable=True).indices
+ expanded, expert = expanded[order], expert[order]
+ rank = torch.arange(expert.numel(), device=dev) - (torch.cumsum(counts, 0) - counts)[expert]
+ permuted = run_start[expert] + rank
+
+ def unwritten(n: int) -> torch.Tensor:
+ return torch.full((n,), -1, dtype=torch.long, device=dev)
+
+ tile_idx_to_expert_idx = unwritten(max_tiles)
+ tile_idx_to_expert_idx[:num_tiles] = tile_expert
+ tile_idx_to_mn_limit = unwritten(max_tiles)
+ tile_idx_to_mn_limit[:num_tiles] = mn_limit
+ permuted_idx_to_expanded_idx = unwritten(max_tiles * tile)
+ permuted_idx_to_expanded_idx[permuted] = expanded
+ expanded_idx_to_permuted_idx = unwritten(num_expanded)
+ expanded_idx_to_permuted_idx[expanded] = permuted
+ outs = (
+ tile_idx_to_expert_idx,
+ tile_idx_to_mn_limit,
+ expanded_idx_to_permuted_idx.view(num_tokens, width),
+ permuted_idx_to_expanded_idx,
+ torch.tensor([num_tiles * tile], device=dev),
+ torch.tensor([num_tiles], device=dev),
+ )
+ return [t.to(torch.int32) for t in outs]
+
+ def is_valid(
+ self,
+ token_selected_experts: torch.Tensor,
+ token_final_scales: torch.Tensor,
+ num_experts: int,
+ top_k: int,
+ local_expert_offset: int,
+ local_num_experts: int,
+ tile_tokens_dim: int,
+ ) -> None:
+ # The host computes T * top_k in int32: from 2^31 the op raises a
+ # negative-dimension error, from 2^32 it sorts only the wrapped remainder.
+ numel = token_selected_experts.numel()
+ assert numel < 2**31, (
+ f"token_selected_experts has {numel} >= 2^31 elements; moe_sort counts "
+ "pairs in int32 and silently sorts only the wrapped remainder past 2^32"
+ )
+ # The kernels read the ids from data_ptr() as a dense [T, top_k] buffer.
+ assert token_selected_experts.is_contiguous(), (
+ "token_selected_experts must be contiguous; a strided view is read as a "
+ "dense [T, top_k] buffer and silently sorts the wrong ids"
+ )
+ assert tile_tokens_dim > 0, (
+ f"tile_tokens_dim must be positive, got {tile_tokens_dim}; 0 returns a "
+ "negative num_non_exiting_tiles without raising"
+ )
+
+
+moe_sort = MoeSort()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/moe_unpermute_inplace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/moe_unpermute_inplace.py
new file mode 100644
index 000000000000..e82eaa5dfdc4
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/moe_unpermute_inplace.py
@@ -0,0 +1,246 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""MoE unpermute: each token's routing-weighted sum of its top-k expert rows, taken
+from the grouped GEMM's permuted output in a fixed slot order, in one trtllm call.
+"""
+
+from __future__ import annotations
+
+import torch
+
+import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+
+def _ckpt(tokens: int, values: str = "normal") -> dict:
+ """The checkpoint's routed-expert combine: top-10 of 512 experts, hidden 2560,
+ bf16 FC2 rows, fp32 routing weights.
+
+ The test lays `permuted_input` out exactly as the CuteDSL FC2 receives it: it
+ draws each row's top-10 of uniform scores over the 512 experts, sorts the ids
+ with `moe_sort` at `tile_tokens_dim` and writes each (token, slot) row at the
+ permuted index the sort gave it. `values` is how the FC2 rows are drawn:
+ `normal` = standard normal; `wide` = standard normal times 2^u, u uniform in
+ [-30, 30] per element, so one sum adds terms up to 2^60 apart in magnitude.
+ The scales are a softmax over each row's ten slots, as the renormalizing
+ router emits them.
+ """
+ return dict(
+ tokens=tokens,
+ tile_tokens_dim=128,
+ values=values,
+ num_experts=512,
+ top_k=10,
+ hidden_size=2560,
+ dtype=torch.bfloat16,
+ scales_dtype=torch.float32,
+ )
+
+
+def _fma_f32(acc: torch.Tensor, prod: torch.Tensor) -> torch.Tensor:
+ """`acc + prod` rounded once to fp32, as a fused multiply-add rounds it.
+
+ `acc` is fp32; `prod` is the fp64 product of a bf16 and an fp32 value, which is
+ exact (at most 8 + 24 significant bits). The fp64 sum is rounded to odd -- its
+ exact rounding error from TwoSum, then the last bit forced to 1 if the sum was
+ inexact -- so that the conversion to fp32 is one correct rounding of the exact
+ sum rather than two roundings in a row (53 >= 24 + 2 bits).
+ """
+ a = acc.double()
+ s = a + prod
+ bb = s - a
+ err = (a - (s - bb)) + (prod - bb)
+ inexact = (err != 0) & torch.isfinite(s)
+ # A sum rounded away from zero truncates to its neighbour one step toward zero.
+ toward_zero = inexact & ((err < 0) != (s < 0))
+ bits = s.view(torch.int64)
+ bits = torch.where(inexact, (bits - toward_zero.long()) | 1, bits)
+ return bits.view(torch.float64).float()
+
+
+class MoeUnpermuteInplace(OpWrapper):
+ """`torch.ops.trtllm.moe_unpermute_inplace`: the routing-weighted top-k combine.
+
+ For each token `t`, writes row `t` of `output`:
+
+ output[t] = bf16( sum over k = 0 .. top_k - 1, in that order, of
+ permuted_input[expanded_idx_to_permuted_idx[t, k]]
+ * topk_scales[t, k] )
+
+ accumulated in fp32 with one rounding per slot (each product fused into the
+ running sum) and rounded once to the output dtype. A slot whose permuted index
+ is negative is skipped; a token with every slot skipped gets a zero row.
+
+ Inputs, as the CuteDSL MoE hands them over: `permuted_input` `[M, H]` bf16 is
+ the FC2 grouped GEMM's output, one row per (token, expert) pair in permuted
+ order, `M >= T * top_k` (the sort's padded row count); `output` `[T, H]` bf16
+ is written in full; `expanded_idx_to_permuted_idx` `[T, top_k]` int32 is the
+ third output of `moe_sort`; `topk_scales` `[T, top_k]` fp32 are the routing
+ weights. Call it with these keyword names, the op's own. Returns `output`.
+
+ Fusion boundary: the gather, the weighting and the sum. The FC2 GEMM and the
+ sort are the caller's, and so is anything added to the routed output (the
+ shared expert). The op is deterministic: one thread block per token sums its
+ slots in slot order, so the result does not depend on where the sort placed
+ each row, on which other tokens share the batch, or on the call. That is what
+ separates it from the FC2 GEMM's fused finalize, whose reduction order varies
+ between launches.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="decode, 1 request without MTP; rows from the sort's block kernel",
+ spec=_ckpt(1),
+ ),
+ Cell(why="8192 rows, max_num_tokens prefill", spec=_ckpt(8192)),
+ Cell(
+ why="FC2 rows spanning 2^-30..2^30; terms 2^60 apart in one sum",
+ spec=_ckpt(8192, values="wide"),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none. Stateless; allocates nothing, writes `output` in place
+ on the current stream with one kernel, `moeUnpermuteKernel<__nv_bfloat16,
+ float, 256>` (`cpp/tensorrt_llm/kernels/moe/cuteDsl/moeUtils.cu`, registered
+ in `cpp/tensorrt_llm/thop/moe/cuteDsl/cuteDslMoeUtilsOp.cpp`), one thread block
+ per token. No host synchronization: it captures into a CUDA graph, and replay
+ on new rows, scales and ids (moe_sort captured with it) is bitwise equal to
+ eager and to `reference`.
+
+ Call sites: the engine's CuteDSL MoE calls it when the FC2 finalize is not
+ fused (`run_moe_nvfp4_impl`, `_torch/moe/fused_moe/fused_moe_cute_dsl.py`), and
+ this target's `_moe_block` calls it in the prefill and the decode forward with
+ the same four tensors: the FC2 output, the forward's dense `[num_tokens,
+ hidden]` routed-output buffer, the sort's third output and the router's fp32
+ weights. Prefill runs eagerly at up to max_num_tokens (8192) rows; decode runs
+ captured at requests x MTP verify width, up to 128.
+
+ What it reads: only the `T * top_k` rows `expanded_idx_to_permuted_idx` names,
+ and the scales. The padding rows of each tile and the tiles past the sort's
+ `num_non_exiting_tiles` are never read (every cell fills every other row of
+ `permuted_input` with NaN, and under a graph those rows hold the previous
+ replay's rows), so the FC2 output's unwritten regions need no memset before
+ this op. What it writes: every element of `output`, whatever it held, and none
+ of its inputs.
+
+ Determinism (bitwise, measured at bring-up; the test re-drives only graph
+ replay against eager): repeated calls; a second moe_sort of the same ids, which
+ places an expert's rows in another order; a random permutation of all permuted
+ rows; a prefix of a larger batch, with or without a fresh sort of the prefix
+ alone; a shuffle of the tokens. The FC2 GEMM's fused finalize
+ (`cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell`) differs from its
+ own first call in about half the output elements at decode, while the plain FC2
+ GEMM followed by this op does not, and is bitwise equal to the slot-order fp32
+ sum of the real FC2 rows.
+
+ Accuracy against the exact sum: the fp32 sum's error is bounded by its terms'
+ magnitudes, so an element whose terms nearly cancel can miss the float64 sum
+ rounded once to bf16 by many ulps of the element itself (827 of 20971520
+ elements at T = 8192, standard-normal rows and softmax weights, up to 137.5).
+
+ Silent, refused by `is_valid`: a non-contiguous `permuted_input`, `output`,
+ `expanded_idx_to_permuted_idx` or `topk_scales` (all four are read or written
+ from `data_ptr()` as dense row-major); an `output` whose dtype is not
+ `permuted_input`'s (an fp32 output receives the bf16 bit patterns in the first
+ half of its bytes); 2^31 or more tokens (nothing is written and nothing
+ raises). The test holds the `permuted_input` refusal.
+
+ Data hazards `is_valid` cannot see (values): a permuted index at or past
+ `permuted_input.size(0)` is read without a bounds check; moe_sort never emits
+ one.
+
+ Loud (measured, not re-driven): int64 indices; hidden size not a multiple of 8; fp32 rows, or fp16 scales
+ with bf16 rows; fewer than `T * top_k` permuted rows; a scales or index shape
+ that disagrees with `output`.
+
+ Accepted, not certified: T = 0 returns and writes nothing; bf16 scales; fp16
+ rows; a CPU `output`.
+
+ Certified only at the checkpoint's values: top_k 10, hidden 2560, bf16 rows and
+ output, fp32 scales, the sort's layout at tile 128, every slot local; T 1 and
+ 8192 eager, 128 under CUDA-graph replay. Slots marked -1 (another rank's
+ expert) are not certified.
+
+ Hard limits (sm_103)
+ tokens one thread block per token, the count narrowed to int32:
+ boundary 2^31, `is_valid` refuses it
+ output elements T * H, offset in int64: exact past 2^31 and 2^32
+ permuted rows int32 index, offset in int64; no index passes through fp32
+ """
+
+ def __call__(
+ self,
+ permuted_input: torch.Tensor,
+ output: torch.Tensor,
+ expanded_idx_to_permuted_idx: torch.Tensor,
+ topk_scales: torch.Tensor,
+ ) -> torch.Tensor:
+ torch.ops.trtllm.moe_unpermute_inplace(
+ permuted_input=permuted_input,
+ output=output,
+ expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx,
+ topk_scales=topk_scales,
+ )
+ return output
+
+ def reference(
+ self,
+ permuted_input: torch.Tensor,
+ output: torch.Tensor,
+ expanded_idx_to_permuted_idx: torch.Tensor,
+ topk_scales: torch.Tensor,
+ ) -> torch.Tensor:
+ """The slot-order fp32 sum, one rounding per slot, rounded once at the end.
+
+ Each product of a row and its weight is exact in float64; `_fma_f32` adds
+ it to the fp32 running sum with one rounding, as a fused multiply-add
+ does. A slot with a negative index leaves the sum untouched. `output` is
+ ignored: the result has `permuted_input`'s dtype, which is what the op
+ writes.
+ """
+ idx = expanded_idx_to_permuted_idx.long()
+ acc = torch.zeros(
+ idx.shape[0], permuted_input.shape[1], dtype=torch.float32, device=idx.device
+ )
+ for k in range(idx.shape[1]):
+ slot = idx[:, k]
+ rows = permuted_input.index_select(0, slot.clamp_min(0)).double()
+ prod = rows * topk_scales[:, k : k + 1].double()
+ acc = torch.where((slot >= 0).unsqueeze(1), _fma_f32(acc, prod), acc)
+ return acc.to(permuted_input.dtype)
+
+ def is_valid(
+ self,
+ permuted_input: torch.Tensor,
+ output: torch.Tensor,
+ expanded_idx_to_permuted_idx: torch.Tensor,
+ topk_scales: torch.Tensor,
+ ) -> None:
+ # One thread block per token, the count narrowed to int32: from 2^31
+ # tokens nothing is written and no error surfaces.
+ assert output.size(0) < 2**31, (
+ f"output has {output.size(0)} >= 2^31 rows; moe_unpermute_inplace launches one "
+ "block per token from an int32 count and silently writes nothing"
+ )
+ # The kernel reads and writes all four through data_ptr() as dense
+ # row-major buffers, and writes `output` with permuted_input's dtype.
+ for name, t in (
+ ("permuted_input", permuted_input),
+ ("output", output),
+ ("expanded_idx_to_permuted_idx", expanded_idx_to_permuted_idx),
+ ("topk_scales", topk_scales),
+ ):
+ assert t.is_contiguous(), (
+ f"{name} must be contiguous; moe_unpermute_inplace addresses it as a dense "
+ "row-major buffer and silently combines or writes the wrong elements"
+ )
+ assert output.dtype == permuted_input.dtype, (
+ f"output dtype {output.dtype} must be permuted_input's ({permuted_input.dtype}); "
+ "the op writes permuted_input's element type into it without checking"
+ )
+
+
+moe_unpermute_inplace = MoeUnpermuteInplace()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/renorm_moe_routing_op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/renorm_moe_routing_op.py
new file mode 100644
index 000000000000..ee9c4e5252c8
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/renorm_moe_routing_op.py
@@ -0,0 +1,198 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Renormalized MoE routing: top-k over the raw router logits, then a softmax
+over just those k, in one trtllm call.
+
+Called by the qwen4_exp target (`qwen3_8_flash_next_nvfp4__sm_103__default`) at
+every layout."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+
+def _ckpt(tokens: int, output_dtype: torch.dtype = torch.bfloat16) -> dict:
+ """The checkpoint's router: 512 experts, top-10, bf16 logits."""
+ return dict(
+ tokens=tokens,
+ num_experts=512,
+ topk=10,
+ logits_dtype=torch.bfloat16,
+ output_dtype=output_dtype,
+ dist="randn",
+ )
+
+
+class RenormMoeRoutingOp(OpWrapper):
+ """`norm_topk_prob=True` routing. Returns `(topk_ids, topk_weights)` -- IDS FIRST.
+
+ Per token: `ids` = the `topk` largest `router_logits`, ranked in the
+ logits' own dtype, by decreasing value, ties to the LOWER expert index
+ (a stable descending sort -- `torch.topk` breaks ties the other way, and at
+ 512 bf16 experts ties are common: it disagrees on ~30% of rows);
+ `weights` = an fp32 softmax over just those `topk` logits, cast to
+ `output_dtype` afterwards. Equal to HF's softmax-then-top-k-then-divide up
+ to rounding (5.96e-08 at 512 / top-10).
+
+ `ids` is `[T, topk]` int32; `weights` `[T, topk]` in `output_dtype`, which
+ defaults to fp32 -- not the logits' dtype. Sibling routing ops return
+ `(weights, ids)`; both are `[T, topk]`, so a swapped unpack is no shape
+ error.
+
+ Fusion boundary: routing only. The router GEMM is the caller's, and so is
+ expert parallelism -- every column of the row is ranked.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # The trunk pre-routes the trtllm-gen FP4 runner with bf16 weights at every
+ # member: tp4/tep4 route every row the rank holds; dep4 routes the
+ # all-gathered rows of the 4 ranks, then runs its expert window in chunks of
+ # at most 8192 rows.
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="MTPLayer._moe decode: one token, bf16 weights; the trunk's decode at batch 1",
+ spec=_ckpt(1),
+ ),
+ Cell(
+ why="the trunk's decode at the largest captured batch: 128 rows at tp4/tep4, "
+ "4 x 32 gathered at dep4, bf16 weights",
+ spec=_ckpt(128),
+ ),
+ Cell(
+ why="the trunk's MTP7 verify at tep4: 128 requests x 8 rows, bf16 weights",
+ spec=_ckpt(1024),
+ ),
+ Cell(
+ why="the trunk's prefill chunk at max_num_tokens 8192, bf16 weights",
+ spec=_ckpt(8192),
+ ),
+ Cell(
+ why="dep4 prefill over the gathered rows, 4 ranks x max_num_tokens 8192: "
+ "the envelope's largest call, bf16 weights",
+ spec=_ckpt(32768),
+ ),
+ Cell(
+ why="_moe_block at 16384 rows (max_num_tokens of the tuned config): fp32 weights, "
+ "as the CuteDSL experts take them; thousands of rows carry an exact bf16 tie",
+ spec=_ckpt(16384, torch.float32),
+ ),
+ Cell(
+ why="coarse logits (3 values): ties decide most of the ranking",
+ spec=dict(_ckpt(64), output_dtype=torch.float32, dist="coarse"),
+ ),
+ )
+
+ note = """
+ Source: compiled-in C++, no JIT path. `torch_ext::renorm_moe_routing_op`
+ in `cpp/tensorrt_llm/thop/moe/routing/customMoeRoutingOp.cpp` (a
+ `TORCH_LIBRARY_FRAGMENT(trtllm)` schema) launches `customMoeRoutingKernel`
+ from `cpp/tensorrt_llm/kernels/moe/routing/customMoeRoutingKernels.cu`.
+
+ Metadata consumed: none; stateless, deterministic, runs on the ambient
+ stream. The launch picks its kernel instance by
+ `nextPowerOfTwo(num_experts)` and `nextPowerOfTwo(topk)` -- 512 / top-10 is
+ the (512, 16) instance with a 256-thread block. PDL (`TRTLLM_ENABLE_PDL`)
+ was measured bitwise inert.
+
+ One warp per row, `ceil(T / 8)` blocks capped at 8192
+ (`invokeCustomMoeRouting`), so the row count changes only the grid: every
+ certified count is under the cap (32768 rows = 4096 blocks), and the
+ grid-stride loop past 65536 rows is not exercised.
+
+ The trunk's bf16-weight cells, measured on GB300: ids equal and weights
+ bit-exact at 1 / 128 / 1024 / 8192 / 32768 rows. Their randn logits tie
+ at the 10th / 11th place, so the ties-to-lower-index rule is exercised at
+ every multi-row cell. The bf16 weights of a row sum to 1 within 2.44e-3
+ (the rounding of ten weights).
+ CUDA-graph replay at 128 rows is bit-exact against the reference on
+ refilled logits (the test holds it).
+
+ Preconditions, enforced loudly by the op: `router_logits` 2-D CUDA bf16 /
+ fp16 / fp32; `output_dtype` fp32, bf16 or None (fp16 raises even for fp16
+ logits); `topk <= 16`; `num_experts <= 512` -- this checkpoint sits on
+ that ceiling.
+
+ Silent, refused by `is_valid`: a non-contiguous `router_logits` (read as a
+ 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
+ 0.09-0.27 here). `RenormalizeMoeRoutingMethod`'s PyTorch fallback (used past
+ 512 / 16) is not bit-compatible: it uses `torch.topk` and casts before the
+ softmax.
+
+ Measured bit-exact against `reference` over 2040 configurations (in dtype
+ x out dtype x num_experts {1..512} x topk {0..16} x tokens {0..8192}); the
+ cells hold the target's router (512 / top-10, bf16 logits,
+ bf16 and fp32 weights), and the test holds the weights to HF's
+ `norm_topk_prob` formulation at bf16 logits.
+
+ Hard limits
+ num_experts <= 512 (raises past it)
+ topk <= 16 (raises past it); > num_experts silent, guarded
+ T certified 1..32768 with bf16 weights (dep4's gathered
+ prefill, 4 x max_num_tokens 8192), 16384 with fp32 weights.
+ From the source, not probed: `numTokens` is narrowed to
+ int32 and the row offset `tokenId * numExperts` is uint32,
+ which wraps at 2^32 / 512 = 8388608 rows.
+
+ Not certified: other num_experts / topk / dtype combinations, `topk == 0`
+ and `num_tokens == 0` (both return empties), non-finite logits (`+inf`,
+ NaN and an all `-inf` row poison the row to NaN; a `-inf` among finite
+ logits is a soft mask). Row counts between the cells -- dep4's smallest
+ gathered decode (4 rows), the MTP3 verify (128 x 4) -- run the same kernel
+ instance on a different grid size and are not cells. Router logits from
+ the checkpoint's own router GEMM were not driven; the cells use randn and
+ coarse logits.
+ """
+
+ def __call__(
+ self,
+ router_logits: torch.Tensor,
+ topk: int,
+ output_dtype: Optional[torch.dtype] = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ return torch.ops.trtllm.renorm_moe_routing_op(router_logits, topk, output_dtype)
+
+ def reference(
+ self,
+ router_logits: torch.Tensor,
+ topk: int,
+ output_dtype: Optional[torch.dtype] = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """Stable descending sort in the logits' dtype, fp32 softmax over the
+ selected k, then the cast."""
+ order = torch.sort(router_logits, dim=-1, descending=True, stable=True).indices[:, :topk]
+ selected = torch.gather(router_logits, 1, order).float()
+ e = torch.exp(selected - selected.max(-1, keepdim=True).values)
+ weights = e / e.sum(-1, keepdim=True)
+ return order.to(torch.int32), weights.to(output_dtype or torch.float32)
+
+ def is_valid(
+ self,
+ router_logits: torch.Tensor,
+ topk: int,
+ output_dtype: Optional[torch.dtype] = None,
+ ) -> None:
+ # The kernel addresses router_logits from data_ptr() as dense
+ # [T, num_experts]: a column slice routes on the wider buffer's bytes.
+ assert router_logits.is_contiguous(), (
+ "router_logits must be contiguous; a strided view is read as a dense "
+ "[num_tokens, num_experts] buffer and silently gives wrong routing"
+ )
+ # Surplus slots come back as duplicated expert 0 at weight 0.
+ assert topk <= router_logits.shape[-1], (
+ f"topk ({topk}) exceeds num_experts ({router_logits.shape[-1]}); the "
+ "surplus slots come back as duplicates of expert 0 at weight 0"
+ )
+
+
+renorm_moe_routing_op = RenormMoeRoutingOp()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/__init__.py
new file mode 100644
index 000000000000..52a7a9daf028
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/flashinfer_gemma_rmsnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/flashinfer_gemma_rmsnorm.py
new file mode 100644
index 000000000000..5d4afe0c7dc4
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/flashinfer_gemma_rmsnorm.py
@@ -0,0 +1,139 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Gemma-style RMS normalization (weight applied as ``1 + w``) via the flashinfer kernel."""
+
+from __future__ import annotations
+
+import torch
+
+import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+#: The hyper-connection residual: hc_count 4 streams of hidden 2560, row stride 10240.
+_STREAM = dict(row_width=10240, dtype=torch.bfloat16, eps=1e-6)
+
+
+class FlashinferGemmaRmsnorm(OpWrapper):
+ """`x / sqrt(mean(x^2, -1) + eps) * (1 + weight)`, one op call, new tensor out.
+
+ The reduction, the `1 + w` and the product are all fp32 inside the kernel;
+ only the store rounds to `x.dtype` -- so `(1 + w)` is not first rounded to
+ bf16. `x` is 2-D `[T, H]` or 3-D `[B, heads, D]` (a separate per-head
+ kernel), normalized over the last dim; `x` is not mutated.
+
+ Fusion boundary: normalization and the biased weight only. No residual add
+ (`flashinfer_gemma_fused_add_rmsnorm` exists for that), no quantization; the
+ plain-`w` form is `flashinfer_rmsnorm`.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # The target's grouped norm is one call per 2560-wide stream of the
+ # 10240-wide residual, each a column slice at row stride 10240 -- that, not
+ # a contiguous row, is what the strided-kernel cells below certify.
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="decode batch 33: stream 1, row stride 10240 > width 2560",
+ spec=dict(shape=(33, 2560), col_offset=2560, **_STREAM),
+ ),
+ Cell(
+ why="prefill 2048 tokens: stream 2 at row stride 10240",
+ spec=dict(shape=(2048, 2560), col_offset=5120, **_STREAM),
+ ),
+ Cell(
+ why="MTP pre_e_norm: the contiguous [T, 2560] embedding",
+ spec=dict(shape=(4, 2560), dtype=torch.bfloat16, eps=1e-6),
+ ),
+ # The default target's grouped-norm sites, per rank: every tp4/tep4 rank
+ # holds all of a step's tokens, so its rows are the step's rows.
+ Cell(
+ why="decode, one token: a single stream-1 row, which torch reports contiguous",
+ spec=dict(shape=(1, 2560), col_offset=2560, **_STREAM),
+ ),
+ Cell(
+ why="decode at batch 128, the largest tp4/tep4 capture: stream 1",
+ spec=dict(shape=(128, 2560), col_offset=2560, **_STREAM),
+ ),
+ Cell(
+ why="MTP7 verify, 128 requests x 8 rows = 1024: stream 2",
+ spec=dict(shape=(1024, 2560), col_offset=5120, **_STREAM),
+ ),
+ Cell(
+ why="prefill chunk at max_num_tokens 8192, the envelope's largest row: stream 2",
+ spec=dict(shape=(8192, 2560), col_offset=5120, **_STREAM),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none. Registered only when flashinfer is importable.
+ `torch.ops.trtllm.flashinfer_gemma_rmsnorm` is a Python `custom_op`
+ (`custom_ops/flashinfer_custom_ops.py`) over `flashinfer.norm.gemma_rmsnorm`,
+ which JIT-compiles a CuTe DSL kernel per (dtype, H, weight_bias, PDL, SM,
+ contiguous) through TVM-FFI; on GB300 the DSL path ran
+ (`flashinfer.norm._use_cuda_norm` is False there).
+ Certified path: flashinfer's CuTe DSL `rmsnorm_cute` kernel for 2-D bf16
+ `x` (`weight_bias=1.0`) at hidden 2560: a stream slice at the residual's
+ 10240 row stride (33, 128, 1024, 2048 and 8192 rows; a strided view selects
+ the strided kernel variant), one stream row, and a contiguous `[4, 2560]`
+ (MTP `pre_e_norm`).
+
+ One row of a stream is contiguous to torch (a size-1 dim ignores its
+ stride), so it takes the compact kernel, compiled for 128-byte aligned
+ rows: every stream starts a multiple of 5120 bytes into the residual row,
+ so the slice is aligned. Its output keeps stride `(10240, 1)` (`empty_like`
+ preserves a dense view's strides); with one row nothing reads past it.
+ Streams 0 and 3 (column offsets 0 and 7680) run the same strided kernel at
+ another base pointer; probed once at 128 rows within the default
+ tolerance, they are not cells.
+
+ Not certified: other widths, the 3-D per-head kernel (`qk_rmsnorm_cute`)
+ and fp16/fp32 operands. The CUDA-JIT fallback (`FLASHINFER_USE_CUDA_NORM=1`,
+ or a DSL that cannot target the device) is a different implementation and
+ is not certified. PDL (`TRTLLM_ENABLE_PDL`, default on) changes scheduling,
+ not results. A CUDA graph captured at the 128-row stream slice replays
+ bitwise equal to the eager call (the test holds this).
+
+ Loud (measured), so not guarded: 1-D and 4-D `x`, `x.stride(-1) != 1`, a weight of
+ the wrong length, dtype or device, and a CPU `x` raise ValueError;
+ float64/int8/uint8 raise KeyError from the DSL dtype table. float8_e4m3fn
+ and float8_e5m2 are in that table and return a plausible fp8 tensor --
+ `is_valid` refuses them.
+
+ Row stride is free; the strided kernel declares it divisible by the vector
+ width (8 elements for bf16/fp16, 4 for fp32) and assumes 16-byte aligned
+ pointers, so odd row strides and misaligned slices were not measured.
+
+ Numerics: default tolerance holds for bf16, nothing loosened; the
+ largest deviation measured is one bf16 ulp (0.03125) at 2048 and 8192 rows
+ of 5120. At the stream-slice cells on GB300 the worst element is one bf16
+ ulp of its own magnitude at every row count (1 to 8192), on at most 1 in
+ 2560 elements; relative L2 at most 2.5e-5.
+ The `+ 1` is measured: against a plain-`w` reference the output is
+ 3.35 to 5.44 away, so the cells tell the two apart.
+
+ Hard limits
+ T * H contiguous inputs above 2^31 - 1 elements route to the strided
+ (int64-offset) kernel boundary 2^31 read from the dispatch, not probed
+ rows a stream slice at row stride 10240 first passes 2^31 elements of
+ offset at 209716 rows; the strided kernel's row stride and row count
+ are int64. The envelope's largest is 8192 rows (8.4e7 elements of
+ offset), so that boundary is not reached and was not probed
+ """
+
+ def __call__(self, x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+ return torch.ops.trtllm.flashinfer_gemma_rmsnorm(x, weight, eps)
+
+ def reference(self, x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+ """fp32 throughout, `(1 + w)` formed in fp32, one rounding at the end."""
+ xf = x.float()
+ normed = xf * torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps)
+ return (normed * (1.0 + weight.float())).to(x.dtype)
+
+ def is_valid(self, x: torch.Tensor, weight: torch.Tensor, eps: float) -> None:
+ assert x.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2), (
+ f"{x.dtype} is accepted by the kernel's dtype table and returns a "
+ "plausible fp8 result; dequantize before normalizing"
+ )
+
+
+flashinfer_gemma_rmsnorm = FlashinferGemmaRmsnorm()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/layer_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/layer_norm.py
new file mode 100644
index 000000000000..f5d9025f186d
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/layer_norm.py
@@ -0,0 +1,145 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""LayerNorm over the last dim -- mean subtraction, biased variance, weight and
+bias -- in one Triton kernel, through the FLA module callable
+`tensorrt_llm._torch.modules.fla.layernorm_gated.layernorm_fn`.
+
+Tier-2 module callable, not a `torch.ops.trtllm.*` op: the registered norm
+family is RMS-norm throughout, and `fused_adaptive_layernorm` refuses every
+width but 5120. One call reshapes to 2-D, allocates the output and launches
+`_layer_norm_fwd_1pass_kernel` once.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch.modules.fla.layernorm_gated import layernorm_fn as _layernorm_fn
+
+#: The vision tower's norm1 / norm2 / merger.norm: width 1152 with a bias, eps
+#: 1e-6. The merger norm runs before the 4-to-1 merge, so it too sees 1152.
+WIDTH = 1152
+EPS = 1e-6
+
+
+class LayerNorm(OpWrapper):
+ """`(x - mean) / sqrt(var + eps) * weight + bias` per row of the last dim,
+ variance over `N` (not `N - 1`), new contiguous tensor in `x`'s dtype.
+
+ Every operand is widened to fp32 on load, the reduction and the affine run
+ in fp32, and only the store rounds. `bias=None` drops the additive term (a
+ compile-time branch of the kernel, not a zero bias).
+
+ Fusion boundary: the norm and its affine only -- no residual, no gate, no
+ output quantization. The callable's `z`, `group_size`, `norm_before_gate`
+ and `is_rms_norm` are fixed at "no gate, one group, LayerNorm"; the gated
+ and RMS forms are different computations, not options of this entry.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: shape, dtype (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 = """
+ Source: the in-tree Python callable `layernorm_fn`
+ (`tensorrt_llm/_torch/modules/fla/layernorm_gated.py`: `LayerNormFn.apply`
+ -> `_layer_norm_fwd` -> `_layer_norm_fwd_1pass_kernel`). A Triton kernel,
+ JIT-compiled on first use, not in the wheel: BLOCK_N 2048 and 8 warps for
+ width 1152; `M` and both row strides are passed as `i32`. One program per
+ row.
+
+ Metadata consumed: none. Inputs are not mutated; two scratch fp32 [M]
+ tensors (mean, rstd) are allocated per call and discarded.
+
+ 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). `weight`/`bias` dtype is free.
+
+ Certified: contiguous bf16 `[rows, 1152]` `x` at 1, 1024 and 32768 rows
+ with a bf16 weight and bias, as the vision tower calls it. Not certified:
+ strided or 3-D `x`, `bias=None`, fp16/fp32 activations, weight/bias in
+ another dtype than `x`.
+
+ Loud failures (left to the callable): weight or bias length != N raises a
+ bare AssertionError with no message; `next_pow2(N) * x.element_size() >
+ 64KB` raises "This layer norm doesn't support feature dim >= 64KB" (N =
+ 65536 in bf16); CPU tensors fail before or at the Triton launch.
+
+ Hard limit: the row offsets `row * stride_x_row` and `row * stride_y_row`
+ are `i32` products. At width 1152 the first row whose offset reaches 2^31
+ is 1864136, so an `x` of >= 1864137 rows wraps. In a fresh process this
+ was measured as a CUDA illegal memory access. With other allocations
+ mapped below `x`, the wrapped read could land in them and be silent; that
+ was not measured and is not guarded. The envelope never reaches it:
+ `VisionTower` cuts a request's frames into calls of at most 32768 rows.
+
+ Accuracy at the cells against an fp64 reference, measured on GB300 beyond
+ the default gate:
+ rows differs from the correctly rounded fp64 worst |err| / row max
+ 1 0 / 1152 0.907 bf16 ulp
+ 1024 37 / 1179648 (3.1e-5) 0.949 bf16 ulp
+ 32768 1096 / 37748736 (2.9e-5) 0.995 bf16 ulp
+ Cosine 0.9999986 and relative L2 1.65e-3 at every cell (the bf16 output
+ rounding itself). The elements that differ are outputs where `x_hat * w`
+ and `b` cancel toward zero: up to 205 ulp of their own tiny magnitude,
+ and each within torch's bf16 tolerance. A constant row returns `bias`
+ bitwise. `eps` sits inside the sqrt.
+
+ What the default bf16 gate tells apart: an unbiased (N - 1) variance, a
+ missing mean subtraction and a missing bias each fail it at every cell
+ (`test_the_gate_discriminates_a_wrong_normalization` holds these). An eps
+ of 1e-5, or eps outside the sqrt, fails it only at some cells, so the gate
+ does not certify the eps placement.
+
+ Accepted but not certified: an fp8-e4m3 `x` returns fp8-e4m3 (output is
+ `empty_like(x)`), within one e4m3 ulp of the exact value; fp64 `x`.
+ Forward only: the autograd Function has no backward. Measured, not
+ certified: one kernel launch, runs unchanged under `inference_mode`,
+ CUDA-graph capturable after a warm-up call compiles the Triton kernel (a
+ 1024-row capture replayed on refilled input is bitwise the eager call).
+ The call site is eager-only.
+ """
+
+ def __call__(
+ self,
+ x: torch.Tensor,
+ weight: torch.Tensor,
+ bias: Optional[torch.Tensor],
+ eps: float,
+ ) -> torch.Tensor:
+ return _layernorm_fn(x, weight, bias, z=None, eps=eps)
+
+ def reference(
+ self,
+ x: torch.Tensor,
+ weight: torch.Tensor,
+ bias: Optional[torch.Tensor],
+ eps: float,
+ ) -> torch.Tensor:
+ """fp64 from primitive ops (not `F.layer_norm`), rounded once to `x.dtype`."""
+ xd = x.double()
+ mean = xd.mean(dim=-1, keepdim=True)
+ var = (xd - mean).square().mean(dim=-1, keepdim=True)
+ out = (xd - mean) / torch.sqrt(var + eps) * weight.double()
+ if bias is not None:
+ out = out + bias.double()
+ return out.to(x.dtype)
+
+
+layer_norm = LayerNorm()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/rms_norm_gated_token_major.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/rms_norm_gated_token_major.py
new file mode 100644
index 000000000000..b9ccf66104b1
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/rms_norm_gated_token_major.py
@@ -0,0 +1,249 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Gated RMS normalization with the gate read from a token-major 3-D view,
+with optional fused static-scale FP8 output quantization."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+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 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")
+#: One tp4/tep4 rank's share: 12 of the 48 value heads, `z` the column view of that rank's
+#: 4120-wide input projection `[q k v | z | a | b]` = 2560 | 1536 | 12 | 12.
+_TP4 = dict(_CKPT, heads=12, z="in_proj", z_row_stride=4120)
+#: A dep4 rank: all 48 heads, `z` in the full 16480-wide projection.
+_DEP4 = dict(_CKPT, z="in_proj", z_row_stride=16480)
+#: Prefill-sized launches (at least this many rows, i.e. 683 tokens of 48 heads) run
+#: the multi-row kernel at `_PREFILL_ROWS` rows a program.
+_PREFILL_MIN_ROWS = 32768 # capacity: fixed: a dispatch threshold, not a buffer size
+_PREFILL_ROWS = 8
+
+
+def _prefill_multirow(x: torch.Tensor, z: torch.Tensor, gate_activation: str) -> bool:
+ # A `z` that does not cover `x` goes to the registered op, which rejects it.
+ n = x.shape[-1]
+ return (
+ x.shape[0] >= _PREFILL_MIN_ROWS
+ and x.shape[0] == z.shape[0] * z.shape[1]
+ and z.shape[2] == n
+ and x.stride(-1) == 1
+ and z.stride(2) == 1
+ and z.stride(1) == n
+ and n <= 256
+ and n & (n - 1) == 0
+ and gate_activation in ("silu", "sigmoid")
+ )
+
+
+def _prefill_launch(x, z, weight, eps, gate_activation):
+ """The multi-row kernel at `_PREFILL_ROWS` rows a program (no PDL: the grid covers every SM)."""
+ rows, n = x.shape
+ out = torch.empty_like(x)
+ _layernorm_gated._rms_norm_gated_fwd_multirow_kernel[(triton.cdiv(rows, _PREFILL_ROWS),)](
+ x, out, weight, z, None, None, x.stride(0), out.stride(0), z.stride(0), rows, eps,
+ N=n, ROWS=_PREFILL_ROWS, HEADS_PER_TOK=z.shape[1], GATE_SIGMOID=gate_activation == "sigmoid",
+ WEIGHT_IS_DELTA=False, LAUNCH_WITH_PDL=False, num_warps=4,
+ ) # fmt: skip
+ return out
+
+
+class RmsNormGatedTokenMajor(OpWrapper):
+ """`rmsnorm(x) * weight * act(z)`, new tensor out, row `r` of `x` gated by `z[r // heads, r % heads]`.
+
+ `x` is `[T * heads, N]`; `z` is `[T, heads, N]` (heads = `z.shape[1]`),
+ its `(heads, N)` block contiguous per token and its token stride free, so a
+ column slice of a wider projection is read in place. `act(v)` is
+ `v * sigmoid(v)` for `"silu"` and `sigmoid(v)` for `"sigmoid"`. All fp32,
+ rounded once. Norm before gate, plain `w` (not `1 + w`).
+
+ With `fp8_scale` (one fp32 element, the downstream linear's static input
+ scale) the output is `float8_e4m3fn`: the gated value is rounded to
+ `x.dtype`, multiplied by `1 / fp8_scale[0]` and saturated at +-448 --
+ byte-identical to the unfused norm -> static-quantize pair.
+
+ Call sites (`PrefillTarget._gdn_epilogue`, and `DecodeTarget._gdn_epilogue`
+ at an MTP verify): `z` is the column view of the GDN layer's one
+ input-projection output, `[T, 6144]` of `[T, 16480]` reshaped `[T, 48, 128]`
+ on a dep4 rank, `[T, 1536]` of `[T, 4120]` reshaped `[T, 12, 128]` on a
+ tp4/tep4 rank, which holds 12 of the 48 value heads.
+
+ Fusion boundary: norm, weight, gate and the optional quantize. The caller
+ owns both projections producing `x` and `z`, and every reshape.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ 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="tp4/tep4, one token: 12 heads, z at token stride 4120; 4-row kernel, PDL",
+ spec=dict(tokens=1, **_TP4),
+ ),
+ Cell(
+ why="tp4/tep4, 4 tokens (an MTP-3 verify of one request): 4-row kernel, PDL",
+ spec=dict(tokens=4, **_TP4),
+ ),
+ Cell(
+ why="tp4/tep4, 128 tokens (the con-128 decode batch): grid past the SM count",
+ spec=dict(tokens=128, **_TP4),
+ ),
+ Cell(
+ why="tp4/tep4, 1024 tokens: the MTP-7 verify of 128 requests x 8 rows",
+ spec=dict(tokens=1024, **_TP4),
+ ),
+ Cell(
+ why="tp4/tep4 prefill at max_num_tokens 8192: 98304 rows on the 8-row kernel",
+ spec=dict(tokens=8192, **_TP4),
+ ),
+ Cell(
+ why="dep4 decode: 48 heads at token stride 16480, a rank's 32 tokens",
+ spec=dict(tokens=32, **_DEP4),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none. Two Triton kernels: a multi-row kernel (4 rows per
+ program; 8 from 32768 rows, without fp8 output) when `x.stride(-1) == 1`,
+ `z.stride(2) == 1`, `z.stride(1) == N` and N is a power of two <= 256 --
+ the checkpoint's N = 128 -- else the generic single-row kernel on a packed
+ copy of `z`. PDL on the 4-row path only when the grid is smaller than the
+ SM count and `TRTLLM_ENABLE_PDL != 0`; results are bitwise identical with
+ it off. The 8-row programs reduce a row's mean square in another order
+ than the 4-row ones (one-ulp differences on some outputs).
+
+ Loud (measured), so not guarded: 3-D `x` (ValueError), a `z` not covering `x`
+ (AssertionError "does not match x shape"; 2-D `z` ValueError), an unknown
+ `gate_activation` (ValueError before launch), and on the fallback a
+ non-unit `x.stride(-1)` (bare AssertionError). `is_valid` covers what the
+ multi-row kernel accepts silently: a weight longer than N (truncated to
+ `weight[:N]`) or shorter (read out of bounds), and an `fp8_scale` with
+ extra elements (only `[0]` read) or not fp32 (widened silently).
+
+ `weight` and `z` are widened to fp32 and need not match `x.dtype`. The
+ output is freshly allocated and contiguous; no input is mutated.
+
+ Both kernels are Triton, JIT-compiled at first call for the running GPU;
+ the cells ran that JIT path on sm_103 (GB300, 152 SMs).
+
+ Certified: bf16 `x` with the sigmoid gate over the in_proj column view of
+ `z`, N = 128.
+ 48 heads, token stride 16480 (tp1 and dep4): 128 and 32 tokens on the
+ 4-row kernel without PDL; 683 tokens (the first 8-row launch) and
+ 16384 tokens on the 8-row kernel.
+ 12 heads, token stride 4120 (a tp4/tep4 rank): 1 and 4 tokens on the
+ 4-row kernel with PDL (grids of 3 and 12 programs); 128 and 1024 tokens
+ on it without; 8192 tokens (98304 rows) on the 8-row kernel.
+ A one-token `z` view reports the dense token stride (1536 at 12 heads);
+ only token 0 is read, so it computes the same.
+ Not certified: fp8 output, the silu gate, the generic fallback kernel,
+ fp16/fp32 operands, a strided `x`. At 12 heads, the 8-row kernel's first
+ launch (2731 tokens) and its masked last program, which an odd token count
+ from 2731 up leaves half full (no 48-head row count does, and the 4-row
+ kernel never has one at either width): a prompt's last prefill chunk.
+
+ Numerics: bf16 inside the default tolerance against the fp32 reference.
+ Measured on sm_103: at most one bf16 ulp from the bf16-rounded reference,
+ on about 1 in 10^5 elements or fewer (none at 1, 4 and 32 tokens).
+ Against the unrounded fp32 value the output's relative L2 (1.56e-3 to
+ 1.66e-3) equals the rounded reference's, and two calls are bitwise equal.
+ fp8 output byte-identical (measured). Nothing loosened.
+
+ Hard limits
+ N * itemsize fallback holds a row in 64 KB boundary 65536 B
+ probed at bf16 N = 32768 (returns) and 40960 (RuntimeError)
+ offsets the multi-row kernel (4 and 8 rows a program) indexes x, z
+ and out in int64; the widest call the envelope makes (8192
+ tokens x 48 heads x 128, z at stride 16480) stays below
+ 2^31 elements on every operand anyway
+ """
+
+ def __call__(
+ self,
+ x: torch.Tensor,
+ z: torch.Tensor,
+ weight: torch.Tensor,
+ eps: float,
+ 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
+ )
+
+ def reference(
+ self,
+ x: torch.Tensor,
+ z: torch.Tensor,
+ weight: torch.Tensor,
+ eps: float,
+ fp8_scale: Optional[torch.Tensor] = None,
+ gate_activation: str = "silu",
+ ) -> torch.Tensor:
+ """fp32 throughout; row `t * heads + h` pairs with `z[t, h]`.
+
+ With `fp8_scale` the gated value is rounded to `x.dtype` first, then
+ multiplied by the fp32 reciprocal of the scale, as the kernel does.
+ """
+ num_rows, width = x.shape
+ xf = x.float()
+ y = xf * torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) * weight.float()
+ zf = z.reshape(num_rows, width).float()
+ gate = torch.sigmoid(zf)
+ if gate_activation != "sigmoid":
+ gate = gate * zf
+ y = y * gate
+ if fp8_scale is None:
+ return y.to(x.dtype)
+ return (y.to(x.dtype).float() * (1.0 / fp8_scale.float())).to(torch.float8_e4m3fn)
+
+ def is_valid(
+ self,
+ x: torch.Tensor,
+ z: torch.Tensor,
+ weight: torch.Tensor,
+ eps: float,
+ fp8_scale: Optional[torch.Tensor] = None,
+ gate_activation: str = "silu",
+ ) -> None:
+ assert weight.shape == (x.shape[-1],), (
+ f"weight must be 1-D of length x.shape[-1]={x.shape[-1]}; got "
+ f"{tuple(weight.shape)} -- the multi-row kernel reads weight[0:N] "
+ "unmasked, so a longer weight is silently truncated and a shorter "
+ "one is read out of bounds"
+ )
+ if fp8_scale is not None:
+ assert fp8_scale.numel() == 1, (
+ f"fp8_scale must hold exactly one element; got {fp8_scale.numel()} "
+ "-- the multi-row kernel loads fp8_scale[0] and ignores the rest"
+ )
+ assert fp8_scale.dtype == torch.float32, (
+ f"fp8_scale must be float32; got {fp8_scale.dtype} -- the "
+ "multi-row kernel widens it and quantizes against the "
+ "lower-precision scale"
+ )
+
+
+rms_norm_gated_token_major = RmsNormGatedTokenMajor()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/__init__.py
new file mode 100644
index 000000000000..52a7a9daf028
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ngram_token_history.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ngram_token_history.py
new file mode 100644
index 000000000000..f0d557a20929
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ngram_token_history.py
@@ -0,0 +1,357 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""PLE n-gram id history: a packed token stream becomes one n-gram id window
+per token, its out-of-slice predecessors read from -- and its tail published
+back to -- the per-slot n-gram context pool.
+
+A macro entry, composed from native torch integer ops (~20 small launches);
+exact, so gated bitwise. Called by the qwen4_exp target at every layout; its
+per-rank slot pools are among the cells. Nothing registered
+under ``torch.ops.trtllm`` assembles an n-gram id window, and
+``EngramHashProvider.compute_hashes`` / ``NgramHashMapping._get_ngram_hashes``
+(``tensorrt_llm/_torch/modules/engram/engram.py``) has no cross-step state: it
+left-pads every chunk from a pad id, so a decode step there loses the
+sequence's whole context.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+#: This checkpoint's PLE constants: `ngram_size` 3, so a 2-id context row
+#: (`Qwen4ExpPLECacheParams.ngram_context_len`).
+CKPT_NGRAM_SIZE = 3
+CKPT_EOS_TOKEN_ID = 248044
+CKPT_VOCAB_SIZE = 248320
+#: 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 takes the same path. The target hands the whole pool
+#: over unsliced.
+ENGINE_SLOTS = 461
+#: The per-rank `ple_ngram` slot counts of the default target's members, from the stock
+#: `Qwen4ExpHybridCacheManager` (`max_batch_size` 128, `max_seq_len` 9216, `max_num_tokens`
+#: 8192, n-gram table host-resident), the same on all four ranks. The pool shares the GDN
+#: recurrent pools' slot space, which the V2 manager sizes from the rank's memory budget,
+#: not from `max_batch_size`; a dep4 rank, which keeps every GDN head, holds fewer.
+TEP4_SLOTS = 1874
+DEP4_SLOTS = 850
+
+
+def _scattered(num_seqs: int) -> tuple[int, ...]:
+ """`num_seqs` distinct slots over the whole engine pool, both ends in."""
+ slots = ((0, ENGINE_SLOTS - 1) + tuple(i * 13 + 4 for i in range(num_seqs)))[:num_seqs]
+ assert len(set(slots)) == num_seqs and max(slots) < ENGINE_SLOTS
+ return slots
+
+
+def _spread(num_seqs: int, pool_slots: int) -> tuple[int, ...]:
+ """`num_seqs` distinct slots evenly over a `pool_slots` pool, both ends in, out of
+ order: batch row `s` is never slot `s`, as the scheduler's slot vector allows."""
+ even = sorted({round(i * (pool_slots - 1) / (num_seqs - 1)) for i in range(num_seqs)})
+ slots = tuple(even[1::2][::-1]) + tuple(even[0::2])
+ assert len(set(slots)) == num_seqs and {0, pool_slots - 1} <= set(slots)
+ return slots
+
+
+def _spec(lengths, slots, carries, pool_slots: int, **spec) -> dict:
+ return dict(
+ lengths=tuple(lengths),
+ slots=tuple(slots),
+ carries=tuple(carries),
+ pool_slots=pool_slots,
+ **spec,
+ )
+
+
+#: The target's own argument dtypes: int32 ids and slots, int64
+#: `query_start_loc_long` (or the int64 `cu_seqlens_from_seq_lens` rebuild).
+_CALL_SITE = dict(ids_dtype=torch.int32, index_dtype=torch.int32, qsl_dtype=torch.int64)
+
+
+class NgramTokenHistory(OpWrapper):
+ """Row `t` of the returned window is `[c[t-n+1], ..., c[t-1], c[t]]` --
+ **newest id last**, so column `ngram_size - 1` is the token's own id --
+ where `c` is the token's **sequence's** id stream: for sequence `s` over
+ packed tokens `[begin, end)` with `slot = state_indices[s]`, a predecessor
+ at shift `k` is `input_ids[t - k]` if `t - begin >= k`, else column
+ `ngram_size - 1 + (t - begin) - k` of the carry row. The carry row is
+ `context_pool[slot]` (oldest id first) when `has_initial_states[s]`, and
+ all `eos_token_id` when not.
+
+ Then, **in place**, each named slot's row becomes the last `ngram_size - 1`
+ ids of the sequence's history including the new tokens (`window[end - 1,
+ 1:]`); a zero-token sequence republishes its carry (so a fresh one is
+ reseeded to EOS). Unnamed slots are bitwise untouched.
+
+ The ids are **raw**: the only EOS written is the fresh-sequence seed, and
+ an in-stream EOS passes through -- the hash downstream owns the
+ segmentation, and `[eos, .., eos, id]` is what it treats as a sequence
+ start. Seeding with 0 would be silently wrong: 0 is a real token.
+
+ Returns a fresh contiguous int64 `[num_tokens, ngram_size]` window;
+ `context_pool` is the only argument modified.
+
+ Fusion boundary: token-to-sequence attribution, carry read, seed, the
+ per-shift select and the publish. The hash, the table gather and slot
+ assignment are the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: lengths, slots, carries, pool_slots (the cache manager's width, not
+ # the batch's), 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, ...] = (
+ Cell(
+ why="MTP verify p=8: 32 decodes x 8 rows, qsl rebuilt",
+ spec=_spec(
+ (8,) * 32,
+ _scattered(32),
+ (True,) * 32,
+ ENGINE_SLOTS,
+ qsl_from="seq_lens",
+ **_CALL_SITE,
+ ),
+ ),
+ Cell(
+ why="MTP verify mixed: 37-token context + 8 x 3 rows",
+ spec=_spec(
+ (37,) + (3,) * 8,
+ _scattered(9),
+ (False,) + (True,) * 8,
+ ENGINE_SLOTS,
+ qsl_from="seq_lens",
+ eos_density=0.05,
+ **_CALL_SITE,
+ ),
+ ),
+ Cell(
+ why="zero-token rows: carry kept, fresh reseeded",
+ spec=_spec((2, 0, 0), (0, 1, 2), (True, True, False), 3),
+ ),
+ Cell(
+ why=f"tep4 decode: 128 x 1 row at the rank's {TEP4_SLOTS}-slot pool",
+ spec=_spec(
+ (1,) * 128,
+ _spread(128, TEP4_SLOTS),
+ (True,) * 128,
+ TEP4_SLOTS,
+ qsl_from="seq_lens",
+ **_CALL_SITE,
+ ),
+ ),
+ Cell(
+ why=f"dep4 decode: a rank's 32 x 1 row at its {DEP4_SLOTS}-slot pool",
+ spec=_spec(
+ (1,) * 32,
+ _spread(32, DEP4_SLOTS),
+ (True,) * 32,
+ DEP4_SLOTS,
+ qsl_from="seq_lens",
+ **_CALL_SITE,
+ ),
+ ),
+ Cell(
+ why=f"tep4 MTP7 verify: 128 x 8 rows at the {TEP4_SLOTS}-slot pool",
+ spec=_spec(
+ (8,) * 128,
+ _spread(128, TEP4_SLOTS),
+ (True,) * 128,
+ TEP4_SLOTS,
+ qsl_from="seq_lens",
+ **_CALL_SITE,
+ ),
+ ),
+ Cell(
+ why="tep4 prefill: one fresh 8192-token context, max_num_tokens",
+ spec=_spec(
+ (8192,), (TEP4_SLOTS - 1,), (False,), TEP4_SLOTS, qsl_from="seq_lens", **_CALL_SITE
+ ),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none -- batch layout, pool, slot map and carry flags
+ are explicit. Call site, once per forward on the PLE layer:
+ `context_pool` the whole `MambaRole.PLE_NGRAM_CONTEXT` buffer
+ (`[slots, ngram_size - 1]` int64, allocated by the V2 hybrid cache
+ manager), `state_indices` the int32 slot vector the GDN mixer uses,
+ `has_initial_states` the context prefix padded with True for the
+ generation rows (the engine's own flags cover the context prefix only),
+ `query_start_loc` int64 -- `query_start_loc_long`, or under MTP verify the
+ `cu_seqlens_from_seq_lens` rebuild, because the metadata's vector has one
+ row per generation sequence and would mis-segment silently.
+
+ Preconditions:
+ - Every slot a real one, `0 <= slot < slots`. The carry read
+ bounds-checks, so -1 or an index past the pool is a device-side assert
+ that poisons the CUDA context. **A -1
+ pad sentinel is not supported**; give padded rows real slots.
+ - Slots distinct: two sequences on one slot is accepted and the publish
+ is unordered (data; not guarded).
+ - `query_start_loc[-1] == input_ids.shape[0]`: trailing tokens are
+ attributed to the last sequence (data; not guarded).
+ - `context_pool` int64 and `ngram_size - 1` wide, `has_initial_states`
+ bool, `state_indices` and `input_ids` 1-D, `ngram_size >= 2`,
+ `num_tokens >= 1`, `num_seqs >= 1`: each refused loudly.
+ - The two length checks in `is_valid` are the silent ones.
+ - Token ids are not range-checked (2^62 and -5 pass through verbatim).
+
+ A wrong window is a window of valid ids: newest-first order, a zero
+ seed, ignored carry flags, packing across sequence boundaries, a rotated
+ carry read and publishing the head all hash to valid rows. Chunked
+ prefill is bitwise equal to one-shot, window and pool, and the pool size
+ does not change the result (461 and 11 slots bitwise identical).
+
+ Certified at the tp1 forward's two verify-step shapes, both on the
+ `cu_seqlens_from_seq_lens` rebuild at the 461-slot pool (32 x 8 rows; a
+ 37-token fresh context with in-stream EOS beside 8 x 3 rows), and at
+ zero-token rows. At the default target's per-rank pools (`TEP4_SLOTS` 1874,
+ `DEP4_SLOTS` 850; every cell on the rebuild, with slots spread over the
+ whole pool, both ends in): a plain decode of 128 x 1 rows at tep4 and of
+ a dep4 rank's 32 x 1 rows, the MTP7 verify's 128 x 8 rows at tep4, and
+ one fresh 8192-token context at tep4 (`max_num_tokens`, the envelope's
+ largest row).
+
+ The pool the manager hands over is a contiguous `[slots, 2]` int64
+ (stride (2, 1), offset 0) at tp4, tep4 and dep4, with the GDN pools'
+ slot count. The call site is `ple_ngram_embed`, whose `reference`
+ composes this entry's `reference`; the tp1 forward launches the fused
+ entry, not this op.
+
+ Not certified: tp4's pool (pure tp4 is not served), the MTP and
+ multimodal members' pools, and the dep4 variants' pools (kv_fp8,
+ kv_nvfp4, chunked, reuse with its recurrent-state snapshots). The op
+ addresses a slot by index alone, so these differ from the cells in pool
+ length only.
+
+ CUDA-graph capture certified (no `repeat_interleave` on device lengths,
+ so no data-dependent shape).
+
+ Hard limits: none reached. `num_tokens` is certified to 8192; every index
+ is int64 and the largest tensor is the `[num_tokens, ngram_size]` window,
+ far from 2^31 elements. Pure torch: no JIT path and no compiled kernel of
+ its own, so nothing in it is architecture-specific.
+ """
+
+ def __call__(
+ self,
+ input_ids: torch.Tensor,
+ context_pool: torch.Tensor,
+ state_indices: torch.Tensor,
+ has_initial_states: torch.Tensor,
+ query_start_loc: torch.Tensor,
+ eos_token_id: int,
+ ngram_size: int,
+ ) -> torch.Tensor:
+ ids = input_ids.long()
+ num_tokens = ids.shape[0]
+ num_seqs = state_indices.shape[0]
+ context_len = ngram_size - 1
+
+ device = ids.device
+ starts = query_start_loc.long()
+ slots = state_indices.long()
+
+ # A copy, so the in-place publish cannot disturb rows still being read
+ # and a zero-token sequence has something to write back.
+ carried = context_pool.index_select(0, slots)
+ carried = torch.where(
+ has_initial_states.reshape(-1, 1),
+ carried,
+ carried.new_full((), eos_token_id),
+ )
+
+ # Device-side and static-shape, so capturable; `repeat_interleave` on
+ # device lengths would make the shape data-dependent and sync.
+ positions = torch.arange(num_tokens, device=device)
+ seq_id = torch.searchsorted(starts[1:].contiguous(), positions, right=True).clamp_max(
+ num_seqs - 1
+ )
+ offset = positions - starts[seq_id]
+
+ # One column per shift, oldest first: the packed vector once the token
+ # is `shift` into its slice, the carry row (virtual positions
+ # -context_len .. -1) before that.
+ columns = []
+ for shift in range(context_len, 0, -1):
+ from_packed = ids.index_select(0, (positions - shift).clamp_min(0))
+ pool_column = (context_len + offset - shift).clamp(0, context_len - 1)
+ from_pool = carried[seq_id, pool_column]
+ columns.append(torch.where(offset >= shift, from_packed, from_pool))
+ columns.append(ids)
+ window = torch.stack(columns, dim=1)
+
+ # The last token's window minus its oldest id, or the carry when empty.
+ ends = starts[1:]
+ last_token = (ends - 1).clamp_min(0)
+ published = torch.where(
+ (ends > starts[:-1]).reshape(-1, 1),
+ window.index_select(0, last_token)[:, 1:],
+ carried,
+ )
+ context_pool[slots] = published
+ return window
+
+ def reference(
+ self,
+ input_ids: torch.Tensor,
+ context_pool: torch.Tensor,
+ state_indices: torch.Tensor,
+ has_initial_states: torch.Tensor,
+ query_start_loc: torch.Tensor,
+ eos_token_id: int,
+ ngram_size: int,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """`(window, context_pool)` as the call leaves them, from each
+ sequence's own id stream materialized as a Python list (carry, then
+ its slice) and read straight out; mutates nothing."""
+ context_len = ngram_size - 1
+ ids = input_ids.tolist()
+ pool = context_pool.tolist()
+ starts = query_start_loc.tolist()
+ carries = has_initial_states.tolist()
+ window = [[0] * ngram_size for _ in ids]
+ published = [list(row) for row in pool]
+ for s, slot in enumerate(state_indices.tolist()):
+ carried = list(pool[slot]) if carries[s] else [eos_token_id] * context_len
+ begin, end = starts[s], starts[s + 1]
+ stream = carried + ids[begin:end]
+ for step in range(end - begin):
+ window[begin + step] = stream[step : step + ngram_size]
+ published[slot] = stream[len(stream) - context_len :]
+ device = input_ids.device
+ return (
+ torch.tensor(window, dtype=torch.long, device=device).reshape(len(ids), ngram_size),
+ torch.tensor(published, dtype=torch.long, device=device).reshape(context_pool.shape),
+ )
+
+ def is_valid(
+ self,
+ input_ids: torch.Tensor,
+ context_pool: torch.Tensor,
+ state_indices: torch.Tensor,
+ has_initial_states: torch.Tensor,
+ query_start_loc: torch.Tensor,
+ eos_token_id: int,
+ ngram_size: int,
+ ) -> None:
+ # Both measured silent; the over-long direction of each raises instead.
+ num_seqs = state_indices.shape[0]
+ assert has_initial_states.shape[0] == num_seqs, (
+ f"has_initial_states must be one flag per sequence; got "
+ f"{tuple(has_initial_states.shape)} against {num_seqs} sequences -- a "
+ f"length-1 flag tensor is broadcast over the whole batch rather than "
+ f"rejected"
+ )
+ assert query_start_loc.shape[0] == num_seqs + 1, (
+ f"query_start_loc must be [num_seqs + 1] = [{num_seqs + 1}] cumulative "
+ f"token counts; got {tuple(query_start_loc.shape)} -- a short one "
+ f"misattributes tokens to the wrong sequence rather than raising"
+ )
+
+
+ngram_token_history = NgramTokenHistory()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_gate_conv_inject_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_gate_conv_inject_norm.py
new file mode 100644
index 000000000000..f1eebf8224e8
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_gate_conv_inject_norm.py
@@ -0,0 +1,761 @@
+# 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.
+
+Called by the qwen4_exp target at every layout; uses the row-to-sequence
+search of ``ple_ngram_embed``. It computes the composition of three grouped
+norms, ``gated_ngram_injection``, ``dilated_short_conv_paged``, two adds and
+the norm-only ``qwen4_exp_hc_inject_norm`` after the PLE add. No
+``torch.ops.trtllm`` op or FlashInfer callable takes a dilation (see
+``dilated_short_conv_paged``).
+
+The three launches split where the parallelism changes shape. The norms and
+the gate are row reductions, one program per (row, stream). The conv and the
+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.
+"""
+
+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 tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ngram_token_history import (
+ DEP4_SLOTS,
+ TEP4_SLOTS,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ple_ngram_embed import (
+ PER_SEQ_MAX_ROWS,
+ _seq_of_rows,
+)
+
+
+@triton.jit
+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 at tp1 (see `dilated_short_conv_paged`).
+ENGINE_SLOTS = 461
+#: The default target's per-rank PLE-conv pools: the conv state and the n-gram context are
+#: bound over one slot space (the stock manager asserts their slot counts equal), so the
+#: counts are the n-gram pool's (`TEP4_SLOTS`, `DEP4_SLOTS`). Every rank holds the whole
+#: `[hc_count * hidden, 9]` conv state at every member. A tp4 rank holds the same per-rank
+#: PLE and GDN state as a tep4 rank, so tp4 is bounded by tep4's.
+TP4_SLOTS = TEP4_SLOTS
+
+
+def _spec(lengths, **spec) -> dict:
+ return dict(lengths=tuple(lengths), **spec)
+
+
+class PleGateConvInjectNorm(OpWrapper):
+ """With `C = hc_count`, `D = value.shape[-1]`, every grouped norm a gemma
+ RMS norm over each `D`-wide stream (`n(x, w) = x * rsqrt(mean(x^2) + eps)
+ * (1 + w)`), and each step rounded to the activation dtype where the
+ unfused composition rounds -- for row `t` of sequence `s`, stream `c`::
+
+ kn = round(n(key[t, c], key_norm[c]))
+ qn = round(n(hc[t, c], query_norm[c]))
+ 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)`, `[T, C*D]` in `hc`'s dtype: fresh, or the
+ two row-strided buffers of `out` (neither overlapping `hc`) written in place.
+ `conv_states` is advanced in place. When `conv_input` is given (a
+ contiguous `[T, C*D]` buffer in `hc`'s dtype) the conv input `xc` is
+ written there, for a caller that has to rebuild the conv state at another
+ row of the sequence later (the speculative rollback of a verify step).
+
+ Fusion boundary: everything from the key/query norms to the next site's
+ norm. The n-gram embedding and the key / value projections are the
+ caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: lengths (rows per sequence), pool_slots (default num_seqs + 2),
+ # max_rows (max(lengths)), carries (None -> the argument is None),
+ # conv_input (True -> a staged conv-input buffer is passed). The
+ # checkpoint's 4 x 2560 bf16 geometry throughout.
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="decode 128 x 1 row at the 461-slot pool",
+ spec=_spec((1,) * 128, pool_slots=ENGINE_SLOTS),
+ ),
+ Cell(
+ why="MTP3 verify: 64 sequences x 4 rows",
+ spec=_spec((4,) * 64, pool_slots=ENGINE_SLOTS),
+ ),
+ # Context-shaped (prefill / mixed) steps: the packed-row conv.
+ Cell(
+ why="prefill 8192 rows, fresh",
+ spec=_spec((8192,), carries=(False,), pool_slots=ENGINE_SLOTS),
+ ),
+ Cell(
+ why="mixed step: 2 contexts + 64 decode rows",
+ spec=_spec(
+ (3001, 5000) + (1,) * 64,
+ carries=(False, True) + (True,) * 64,
+ pool_slots=ENGINE_SLOTS,
+ ),
+ ),
+ # The default target's members, each at the per-rank conv pool its rank binds, as
+ # the forward calls it: decode and verify with `has_initial_states` None and
+ # `max_rows` the rows per sequence (a verify step also stages the conv
+ # input), prefill with one flag per sequence and `max_rows` the step's rows.
+ Cell(
+ why=f"tep4 decode: 128 x 1 row, {TEP4_SLOTS} slots",
+ spec=_spec((1,) * 128, pool_slots=TEP4_SLOTS),
+ ),
+ Cell(
+ why=f"dep4 decode: a rank's 32 x 1 row, {DEP4_SLOTS} slots",
+ spec=_spec((1,) * 32, pool_slots=DEP4_SLOTS),
+ ),
+ Cell(
+ why=f"tep4 MTP1 verify: 128 sequences x 2 rows, {TEP4_SLOTS} slots",
+ spec=_spec((2,) * 128, pool_slots=TEP4_SLOTS, conv_input=True),
+ ),
+ Cell(
+ why=f"tep4 MTP7 verify: 128 sequences x 8 rows (packed-row conv), {TEP4_SLOTS} slots",
+ spec=_spec((8,) * 128, pool_slots=TEP4_SLOTS, conv_input=True),
+ ),
+ Cell(
+ why=f"tep4 prefill: one fresh 8192-token context, {TEP4_SLOTS} slots",
+ spec=_spec((8192,), carries=(False,), pool_slots=TEP4_SLOTS),
+ ),
+ Cell(
+ why=f"envelope: 8192 rows over 128 sequences (1 context + 127 decode), {TEP4_SLOTS} slots",
+ spec=_spec(
+ (8192 - 127,) + (1,) * 127,
+ carries=(False,) + (True,) * 127,
+ max_rows=8192,
+ pool_slots=TEP4_SLOTS,
+ ),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none. Call site: the decode target's PLE block, once
+ per forward: `hc` the `[T, 10240]` stream, `key` the key projection
+ `[T, 10240]` and `value` the value projection `[T, 2560]` (both cuBLAS,
+ bf16), the three PLE norm weights and the next site's attention-side norm
+ weight as contiguous `[10240]`, the `[10240, 1, 4]` conv weight, the whole
+ `[slots, 10240, 9]` conv pool, int32 slots, int64 `query_start_loc_long`
+ (the rebuild under MTP verify), `max_rows = gen_p`.
+
+ Preconditions (data; not guarded): at most `max_rows` rows per sequence;
+ distinct live slots; `query_start_loc[-1] == T`. Operands 2-D with unit
+ column stride; weights contiguous. Certified at the checkpoint's 4 x 2560
+ bf16 geometry with every slot live (the target never passes `pad_slot_id`),
+ at a 128 x 1-row decode, a 64 x 4-row MTP3 verify, an 8192-row fresh
+ prefill and a mixed 2-context + 64-decode step, all on the 461-slot pool;
+ and at the default target's per-rank pools (`TEP4_SLOTS` 1874, `DEP4_SLOTS` 850,
+ slots spread over the whole pool with both ends taken): a 128 x 1 tep4 and
+ a 32 x 1 dep4 decode, 128 x 2 (MTP1) and 128 x 8 (MTP7, the packed-row conv
+ with no carry flags) verifies staging `conv_input`, an 8192-row fresh
+ prefill, and the envelope's largest step, 8192 rows over 128 sequences (one
+ 8065-row context, 127 carrying decode rows). A tp4 rank binds the same
+ conv state as a tep4 rank, so the tep4 cells cover it by bound.
+ Not certified: sequences shorter than `max_rows` on the per-sequence path
+ (a verify step gives every sequence `max_rows` rows); verify widths 3, 5,
+ 6 and 7 rows at the per-rank pools (4 rows only at 64 x 4 over 461 slots);
+ the MTP member's own pool (max_seq_len 16384; the verify cells run at the
+ tep4 count); a dep4 prefill or verify step; a chunked continuation at the
+ per-rank pools (a carrying context chunk is certified only by the 461-slot
+ mixed cell). Every pool count is the stock engine's with the n-gram table
+ host-resident.
+
+ Index range: slots and rows are widened to int64 before every offset, so
+ neither the pool (2^31 elements at 23301 slots of `[10240, 9]`) nor the
+ rows (2^31 elements at 209715 rows of 10240) is near a 32-bit limit at
+ the envelope (1874 slots, 8192 rows); nothing is indexed in fp32.
+
+ Source: the four Triton kernels in this file (three launches a call, the
+ packed-row conv a fourth on context-shaped steps), JIT-compiled for the
+ running device; every cell ran the JIT path on sm_103 (GB300).
+
+ Numerics: every rounding point of the composition is kept, so the only
+ differences are fp32 reduction order (the three norms' sums of squares
+ 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,
+ conv_input: Optional[torch.Tensor] = None,
+ out: Optional[Tuple[torch.Tensor, 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]
+ if out is None:
+ hc_out = torch.empty((rows, width), dtype=hc.dtype, device=hc.device)
+ normed = torch.empty((rows, width), dtype=hc.dtype, device=hc.device)
+ else:
+ hc_out, normed = out
+ 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 = (
+ conv_input
+ if conv_input is not None
+ else 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,
+ conv_input: Optional[torch.Tensor] = None,
+ out: Optional[Tuple[torch.Tensor, 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 unfused composition) 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,
+ conv_input: Optional[torch.Tensor] = None,
+ out: Optional[Tuple[torch.Tensor, 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"
+ )
+ assert conv_input is None or (
+ conv_input.shape == hc.shape
+ and conv_input.dtype == hc.dtype
+ and conv_input.is_contiguous()
+ ), "conv_input must be a contiguous buffer shaped and typed as hc"
+ for t in out or ():
+ assert t.shape == hc.shape and t.dtype == hc.dtype and t.stride(1) == 1, (
+ "out buffers must be shaped and typed as hc, unit column stride"
+ )
+
+
+ple_gate_conv_inject_norm = PleGateConvInjectNorm()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_ngram_embed.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_ngram_embed.py
new file mode 100644
index 000000000000..f727731c8386
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_ngram_embed.py
@@ -0,0 +1,624 @@
+# 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.
+
+Called by the qwen4_exp target at every layout. It computes the composition
+``ngram_token_history`` + ``ngram_hash`` + ``nn.functional.embedding`` +
+``dequantize_e4m3_per_tensor``, whose references it reads. Nothing registered
+under ``torch.ops.trtllm`` or in the installed FlashInfer gathers and
+dequantizes n-gram rows (``torch.ops.trtllm.qwen4_exp_ple_ngram_hash`` hashes a
+window only).
+
+One program per **sequence**, which loops over that sequence's rows: the carry
+read and the publish of one slot happen in one program, so no other program can
+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 tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.embedding.ngram_hash import (
+ ngram_hash,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ngram_token_history import (
+ DEP4_SLOTS,
+ TEP4_SLOTS,
+ ngram_token_history,
+)
+
+#: This checkpoint's PLE geometry: 3-gram, 16 heads of 160 fp8 columns.
+CKPT_NGRAM_SIZE = 3
+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 at tp1 (see `ngram_token_history`).
+ENGINE_SLOTS = 461
+#: A tp4 rank's pool, not counted on an engine. A tp4 rank holds the same per-rank PLE
+#: and GDN state as a tep4 rank, so the pool is bounded by tep4's (`TEP4_SLOTS`).
+TP4_SLOTS = TEP4_SLOTS
+#: Above this many rows per sequence (MTP verify carries at most 4) the call
+#: takes the context-shaped path: packed row tiles plus a separate publish.
+PER_SEQ_MAX_ROWS = 4
+_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 _spec(lengths, **spec) -> dict:
+ return dict(lengths=tuple(lengths), **spec)
+
+
+class PleNgramEmbed(OpWrapper):
+ """For sequence `s` over packed rows `[b, e)` with `slot = state_indices[s]`,
+ its id stream is `carry + input_ids[b:e]`, where `carry` is
+ `context_pool[slot]` (oldest first) when `has_initial_states[s]` -- or
+ always, when `has_initial_states` is None -- and `ngram_size - 1` copies of
+ `eos_token_id` when not. Row `t` of that slice gets the window
+ `stream[t : t + ngram_size]` (newest last), exactly `ngram_token_history`'s;
+ each window hashes to one global row per head exactly as `ngram_hash` with
+ `num_query_tokens=1`; and::
+
+ out[t, h * D : (h + 1) * D] = to(scale.dtype)(float(table[row[t, h]]) * float(scale))
+
+ with `D = table.shape[1]` -- `dequantize_e4m3_per_tensor`'s one rounding.
+ Then, **in place**, `context_pool[slot]` becomes the last `ngram_size - 1`
+ ids of the stream (a zero-row sequence republishes its carry). Unnamed slots
+ are untouched.
+
+ Returns a fresh contiguous `[num_tokens, heads * D]` in `scale`'s dtype.
+ `context_pool` is the only argument modified.
+
+ Fusion boundary: carry read and seed, window, hash, gather, dequantize,
+ publish. The key/value projections that consume the result are the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: lengths (rows per sequence), pool_slots (default num_seqs + 2),
+ # carries (None -> the argument is None), max_rows (default max(lengths)),
+ # eos_density (0.0), table "small" | "host" | "ckpt" (the full
+ # 320001536-row device table, rows filled where touched). The checkpoint's
+ # 3-gram of 8 heads per order, int32 ids.
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="decode 128 x 1 row, host-resident table, 461 slots",
+ spec=_spec((1,) * 128, pool_slots=ENGINE_SLOTS, table="host"),
+ ),
+ Cell(
+ why="decode 128 x 1 row, the full checkpoint-height table",
+ spec=_spec((1,) * 128, pool_slots=ENGINE_SLOTS, table="ckpt", eos_density=0.05),
+ ),
+ Cell(
+ why="MTP3 verify: 64 sequences x 4 rows, 461 slots",
+ spec=_spec((4,) * 64, pool_slots=ENGINE_SLOTS, table="host", eos_density=0.05),
+ ),
+ # Context-shaped (prefill / mixed) steps: the packed-row form.
+ Cell(
+ why="mixed step: 2 contexts + 64 decode rows",
+ spec=_spec(
+ (3001, 5000) + (1,) * 64,
+ carries=(False, True) + (True,) * 64,
+ pool_slots=ENGINE_SLOTS,
+ eos_density=0.05,
+ ),
+ ),
+ # The default target's members, each at the per-rank n-gram pool its rank binds,
+ # as the forward calls it: decode and verify with `has_initial_states`
+ # None and `max_rows` the rows per sequence, prefill with one flag per
+ # sequence and `max_rows` the step's row count.
+ Cell(
+ why=f"tep4 decode: 128 x 1 row, host-resident table, {TEP4_SLOTS} slots",
+ spec=_spec((1,) * 128, pool_slots=TEP4_SLOTS, table="host", eos_density=0.05),
+ ),
+ Cell(
+ why=f"dep4 decode: a rank's 32 x 1 row, host-resident table, {DEP4_SLOTS} slots",
+ spec=_spec((1,) * 32, pool_slots=DEP4_SLOTS, table="host", eos_density=0.05),
+ ),
+ Cell(
+ why=f"tp4 decode: 128 x 1 row, the full checkpoint-height table, {TP4_SLOTS} slots",
+ spec=_spec((1,) * 128, pool_slots=TP4_SLOTS, table="ckpt", eos_density=0.05),
+ ),
+ Cell(
+ why=f"tep4 MTP1 verify: 128 sequences x 2 rows, {TEP4_SLOTS} slots",
+ spec=_spec((2,) * 128, pool_slots=TEP4_SLOTS, table="host", eos_density=0.05),
+ ),
+ Cell(
+ why=f"tep4 MTP7 verify: 128 sequences x 8 rows (packed-row form), {TEP4_SLOTS} slots",
+ spec=_spec((8,) * 128, pool_slots=TEP4_SLOTS, table="host", eos_density=0.05),
+ ),
+ Cell(
+ why=f"tep4 prefill: one fresh 8192-token context, {TEP4_SLOTS} slots",
+ spec=_spec(
+ (8192,), carries=(False,), pool_slots=TEP4_SLOTS, table="host", eos_density=0.05
+ ),
+ ),
+ Cell(
+ why=f"envelope: 8192 rows over 128 sequences (1 context + 127 decode), {TEP4_SLOTS} slots",
+ spec=_spec(
+ (8192 - 127,) + (1,) * 127,
+ carries=(False,) + (True,) * 127,
+ max_rows=8192,
+ pool_slots=TEP4_SLOTS,
+ table="host",
+ eos_density=0.05,
+ ),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none -- batch layout, pool, slot map are explicit.
+ Call site: the decode target's PLE block, once per forward:
+ `input_ids` int32 `[T]`, the whole `[slots, 2]` int64 n-gram pool,
+ int32 `state_indices[:batch]`, `has_initial_states=None` (every decode
+ sequence carries), int64 `query_start_loc_long` (or the
+ `cu_seqlens_from_seq_lens` rebuild under MTP verify), the
+ `[320001536, 160]` fp8 table (host-resident by default), the `[1, 1, 1]`
+ bf16 scale, `max_rows = gen_p`.
+
+ Preconditions (data; not guarded -- each would need a sync):
+ - every sequence has at most `max_rows` rows; rows past it are neither
+ written nor published;
+ - slots real (`0 <= slot < slots`) and distinct; -1 is not supported;
+ - `query_start_loc[-1] == num_tokens`;
+ - token id * max multiplier < 2^63 (int64 wraps, as in `ngram_hash`).
+
+ Certified at a 128 x 1-row decode on a host-resident table and on the
+ full checkpoint-height device table, an MTP3 verify step of 64 x 4 rows
+ on the per-sequence path, and a mixed step of 2 contexts (one fresh, one
+ carrying) + 64 decode rows on the packed-row form. Not certified:
+ zero-row sequences.
+
+ At the default target's per-rank n-gram pools, with slots spread over the whole
+ pool and both ends taken (a lone sequence on the last slot):
+ - tep4 (`TEP4_SLOTS` 1874): 128 x 1 decode; MTP1 verify, 128 x 2 rows,
+ on the per-sequence path; MTP7 verify, 128 x 8 rows, on the packed-row
+ form with no carry flags (the decode call site's `None`); one fresh
+ 8192-token context; and the envelope's largest step, 8192 rows over
+ 128 sequences (one fresh 8065-token context + 127 carrying decode rows).
+ - dep4 (`DEP4_SLOTS` 850): a rank's 32 x 1 decode.
+ - tp4 (`TP4_SLOTS`, bounded by tep4's pool): 128 x 1 decode on the full
+ checkpoint-height device table.
+ The pools are the stock engine's with a host-resident table; an MTP
+ member reserves its own pool, not counted here. CUDA-graph replay
+ is certified at 16 x 1 rows over 461 slots and at 128 x 1 and 128 x 8
+ rows over 1874 slots.
+
+ The kernels are Triton, JIT-compiled on first call for the running
+ device: every cell above ran that path on sm_103. The host-resident
+ cells read a pinned host table through its mapped device address, the
+ stock PLE host offload's placement. The checkpoint-height cells allocate
+ the 47.7 GiB table on the device.
+
+ Bitwise equal to `ngram_token_history` + `ngram_hash` + embedding +
+ `dequantize_e4m3_per_tensor`, window, rows, values and pool alike.
+ CUDA-graph capturable: no host sync, one launch, output allocated per call.
+ 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 it fuses -- `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/ops/quantization/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/quantization/__init__.py
new file mode 100644
index 000000000000..52a7a9daf028
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/quantization/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/quantization/bf16_to_fp32.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/quantization/bf16_to_fp32.py
new file mode 100644
index 000000000000..61edc7eb2a1d
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/quantization/bf16_to_fp32.py
@@ -0,0 +1,80 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Widening bf16 -> fp32 of a contiguous tensor in one vectorized kernel: the logits
+processor's cast of the step's `[rows, vocab]` logits.
+
+Tier 5 (Triton written for this target). torch's `.float()` runs this as an unrolled,
+unvectorized `direct_copy` at under half the bandwidth; the values are the same bit for
+bit.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+#: This checkpoint's vocabulary: the width of the logits rows.
+VOCAB = 248320
+#: Elements a program (8 warps: 16 a thread, two 16-byte loads).
+_BLOCK = 4096
+_WARPS = 8
+
+
+@triton.jit
+def _bf16_to_fp32_kernel(x_ptr, y_ptr, n, BLOCK: tl.constexpr):
+ offs = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
+ mask = offs < n
+ tl.store(y_ptr + offs, tl.load(x_ptr + offs, mask=mask).to(tl.float32), mask=mask)
+
+
+class Bf16ToFp32(OpWrapper):
+ """`x.float()` for a contiguous bf16 `x`: a fresh contiguous fp32 tensor of `x`'s
+ shape. Widening is exact, so every element is `x`'s value; NaN payloads widen as
+ torch widens them (the 16 bits shifted up).
+
+ Fusion boundary: the cast only; the head's GEMM before it and the sampler after it
+ are the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="decode logits: 128 rows of the 248320 vocabulary", spec=dict(shape=(128, VOCAB))),
+ Cell(why="one row; a ragged last program", spec=dict(shape=(1, VOCAB))),
+ Cell(
+ why="every bf16 bit pattern: NaNs, infinities, subnormals, signed zeros",
+ spec=dict(shape=(1, 65536), source="all_codes"),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none; stateless. One Triton kernel, `cdiv(numel, 4096)` programs
+ of 8 warps over the flat tensor. Call site: the target's logits processor
+ (`_WideningLogitsProcessor`), on the head's bf16 logits; the shell's `.float()`
+ after it is then a no-op.
+
+ GB300, graph replay: 128 x 248320 31.3 us against 82.1 for `.float()`, 124 rows
+ 30.7 against 79.9, 33 rows 8.3 against 16.5, one row 4.1 against 4.2.
+
+ Preconditions: bf16, contiguous (the kernel reads the flat storage).
+ """
+
+ def __call__(self, x: torch.Tensor) -> torch.Tensor:
+ y = torch.empty(x.shape, dtype=torch.float32, device=x.device)
+ n = x.numel()
+ if n:
+ _bf16_to_fp32_kernel[(triton.cdiv(n, _BLOCK),)](x, y, n, BLOCK=_BLOCK, num_warps=_WARPS)
+ return y
+
+ def reference(self, x: torch.Tensor) -> torch.Tensor:
+ return x.float()
+
+ def is_valid(self, x: torch.Tensor) -> None:
+ assert x.dtype == torch.bfloat16, "bf16 input"
+ assert x.is_contiguous(), "contiguous input: the kernel reads the flat storage"
+
+
+bf16_to_fp32 = Bf16ToFp32()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/quantization/fp8_quantize_1x128.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/quantization/fp8_quantize_1x128.py
new file mode 100644
index 000000000000..1ac77fe895e5
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/quantization/fp8_quantize_1x128.py
@@ -0,0 +1,161 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Dynamic 1x128 block-wise FP8 quantization: bf16 -> e4m3 data + transposed fp32 block scales."""
+
+from __future__ import annotations
+
+import torch
+
+import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+#: Largest finite e4m3 magnitude.
+E4M3_MAX = 448.0
+
+#: The kernel's amax floor, as the bf16 constant it compares against:
+#: bfloat16(1e-10) == 1.000444171950221e-10.
+AMAX_FLOOR = float(torch.tensor(1e-10, dtype=torch.bfloat16))
+
+#: The two divisions in the kernel are approximate (`--use_fast_math`), each
+#: within CUDA's documented 2 ulp; the returned scale is two deep.
+SCALE_ULP_GATE = 4
+
+
+def block_scales(input: torch.Tensor, use_ue8m0: bool) -> tuple[torch.Tensor, torch.Tensor]:
+ """`(quant_scale, dequant_scale)`, each `[T, ceil(H/128)]` fp32, exactly rounded.
+
+ The amax runs in bf16, as the kernel's running maximum does (exact: |x| of
+ a bf16 is a bf16), floored at `AMAX_FLOOR`; a short tail block takes the
+ amax of its own elements. The kernel multiplies by the first before
+ rounding to e4m3 and returns the second.
+ """
+ t, n = input.shape
+ nb = (n + 127) // 128
+ mag = torch.nn.functional.pad(input.abs(), (0, nb * 128 - n))
+ amax = mag.reshape(t, nb, 128).amax(dim=2).float().clamp(min=AMAX_FLOOR)
+ qs = E4M3_MAX / amax
+ if not use_ue8m0:
+ return qs, 1.0 / qs
+ raw = 1.0 / qs
+ mant, exp = torch.frexp(raw)
+ ds = torch.ldexp(torch.ones_like(raw), torch.where(mant == 0.5, exp - 1, exp))
+ return 1.0 / ds, ds
+
+
+#: The MTP draft layer's hidden width, quantized with the plain scale the
+#: block-scale MoE runner wants.
+_CKPT = dict(hidden=2560, use_ue8m0=False)
+
+
+class Fp8Quantize1x128(OpWrapper):
+ """Quantize a bf16 `[T, H]` matrix to e4m3 with one fp32 scale per 1x128 row block.
+
+ Per token `t` and block `b` of 128 consecutive columns (the last one short
+ when `H % 128 != 0`): `amax = max(|x[t, block]|, bf16(1e-10))`,
+ `data = e4m3(x * (448 / amax))` saturating at +-448, `scale[b, t] =
+ 1 / (448 / amax)`. Returns `(data [T, H] float8_e4m3fn, scales [ceil(H/128),
+ T] fp32)`.
+
+ Three readings are each a silent wrong answer: `scales` is TRANSPOSED
+ (`[NB, T]`, the layout `fp8_block_scale_moe_runner` consumes); it
+ MULTIPLIES the code back (dividing is 6780 relative RMS off); and with
+ `use_ue8m0=False` it is the plain `amax / 448`, while `use_ue8m0=True`
+ rounds every scale up to a power of two -- a different quantization of the
+ same input (99% of scales, 96% of codes change), not a layout choice.
+
+ Fusion boundary: the amax reduction, scaling, rounding and (on the sm_100
+ family) the compaction into `[NB, T]` are inside; no GEMM is fused.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="MTPLayer._moe decode: one token at hidden 2560", spec=dict(tokens=1, **_CKPT)),
+ Cell(
+ why="20 tokens == H/128: the scale buffer is square, its transpose undetectable",
+ spec=dict(tokens=20, **_CKPT),
+ ),
+ Cell(
+ why="16384 rows: max_num_tokens of the tuned config",
+ spec=dict(tokens=16384, **_CKPT),
+ ),
+ Cell(
+ why="128 rows: the MTP drafter's decode step at max_batch_size 128",
+ spec=dict(tokens=128, **_CKPT),
+ ),
+ Cell(
+ why="1024 rows: the first draft over mtp7's verify rows, 128 requests x 8",
+ spec=dict(tokens=1024, **_CKPT),
+ ),
+ Cell(
+ why="8192 rows: the MTP layer over one max_num_tokens context chunk",
+ spec=dict(tokens=8192, **_CKPT),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none. The returned scale layout is chosen by an
+ architecture query (`isSM100Family()` in thop/fp8Quantize.cpp): off the
+ sm_100 family the op returns a 1-D padded buffer instead. That branch was
+ read, not run.
+
+ Preconditions, 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`. At a `T` that is a multiple of 4 (128, 1024, 8192) nothing is
+ padded and `scales` is the kernel's own buffer, not a compacting copy.
+ CUDA-graph capture and replay are certified at 64, 128 and 1024 rows.
+
+ Source: compiled-in C++, `torch.ops.trtllm.fp8_quantize_1x128` from
+ `thop/fp8Quantize.cpp` (`fp8_1x128_cs` -> `launch_scale_1x128`). The
+ launcher takes `scale_1x128_vec_kernel` when `can_vectorize_1x128` holds
+ (H % 8 == 0, 16-byte aligned input, 8-byte aligned output), as at every
+ cell; the element-wise `scale_1x128_kernel` it falls back to (an input
+ whose storage offset breaks the 16-byte alignment) is not certified. The
+ grid is persistent, `kNumDeviceSMs * 8` blocks of 8 warps quantizing 2
+ groups each (19456 groups a sweep on GB300); the cells run under one sweep
+ (128 rows), just past one (1024) and nine (8192).
+
+ Measured on GB300 at H 2560: codes bit-exact at 1 row and on >= 99.94% of
+ elements at every other row count, each miss within 1.122e-07 (relative)
+ of its rounding midpoint; scales within the 4-ulp gate.
+
+ Hard limits
+ T certified at 1, 20, 128, 1024, 8192 and 16384 at H 2560 (128: the
+ drafter's decode at max_batch_size 128; 1024: the first draft over
+ mtp7's 128 x 8 verify rows; 8192: a max_num_tokens context chunk;
+ 16384: max_num_tokens of the tuned config); above that not probed
+ H multiple of 16 (raises otherwise); certified at 2560, the draft layer's
+ width
+
+ Not certified: `T == 0` (the op accepts it and returns `[0, H]` /
+ `[NB, 0]`); a short tail block (`H % 128 != 0`); all-zero blocks (the
+ `AMAX_FLOOR` scale), `inf` / NaN inputs and saturation; `use_ue8m0=True`;
+ the element-wise fallback kernel; the registered fake.
+ """
+
+ def __call__(
+ self,
+ input: torch.Tensor,
+ use_ue8m0: bool = False,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ return torch.ops.trtllm.fp8_quantize_1x128(input, use_ue8m0)
+
+ def reference(
+ self,
+ input: torch.Tensor,
+ use_ue8m0: bool = False,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """`(data, scales[NB, T])` from exactly-rounded fp32 divisions."""
+ qs, ds = block_scales(input, use_ue8m0)
+ cols = qs.repeat_interleave(128, dim=1)[:, : input.shape[1]]
+ data = (input.float() * cols).to(torch.float8_e4m3fn)
+ return data, ds.t().contiguous()
+
+
+fp8_quantize_1x128 = Fp8Quantize1x128()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/__init__.py
new file mode 100644
index 000000000000..52a7a9daf028
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_triton_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_triton_update.py
new file mode 100644
index 000000000000..1bc5b6bf9d6a
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_triton_update.py
@@ -0,0 +1,388 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Multi-token causal depthwise conv1d step (Triton), with optional per-step
+window snapshots for speculative verification.
+
+Wraps ``tensorrt_llm._torch.modules.mamba.causal_conv1d_triton.causal_conv1d_update``
+-- one Triton launch (``_causal_conv1d_update_kernel``) after allocating the
+output. There is no torch-op boundary around this kernel. It is *not*
+``torch.ops.trtllm.causal_conv1d_update``: that op writes `x` in place, has a
+ring mode, takes `state_len` from the pool, has no snapshots, and multiplies in
+fp32 -- the two are not substitutable.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Union
+
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch.modules.mamba.causal_conv1d_triton import (
+ causal_conv1d_update as _causal_conv1d_triton_update,
+)
+
+#: This checkpoint's GDN `[q|k|v]` width; the conv is width 4 over a 3-column pool.
+CKPT_DIM = 10240
+#: 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
+#: Row stride of tp1's (and a dep4 rank's) input-projection output `[q|k|v|z|a|b]`.
+CKPT_X_STRIDE = 16480
+
+#: A tp4/tep4 rank's GDN conv width (`[q | k | v]` of 4 key and 12 value heads of 128) and
+#: the row stride of its input-projection output, `[q | k | v | z | a | b]` = 2560 + 1536 +
+#: 12 + 12.
+RANK_DIM = 2560
+RANK_X_STRIDE = 4120
+#: Conv-pool slots of one tep4 rank's engine in the V2 hybrid manager (memory-derived; the
+#: PLE and GDN pools share one slot space).
+TEP4_SLOTS = 1874
+#: GDN layers whose conv states share one slot's page in the V2 manager: a layer's view
+#: `[slots, dim, 3]` has slot stride 36 * dim * 3 (276480 at 2560) and layer `i` sits
+#: `i * dim * 3` elements into the page.
+GDN_PAGE_LAYERS = 36
+#: The stock GDN speculative state's snapshot buffer is `[36 layers, max_batch_size,
+#: tokens_per_gen_step, dim, 3]`, a layer's view `[max_batch_size, p, dim, 3]`, with
+#: `intermediate_state_indices = arange(max_batch_size)` (`GDNReplayState.bind`).
+MAX_BATCH_SIZE = 128
+
+
+def _pinned(batch: int) -> tuple[int, ...]:
+ """`batch` distinct slots of the engine pool, slot 0 and the last pinned in."""
+ return (0,) + tuple(i * 14 + 3 for i in range(1, batch - 1)) + (ENGINE_SLOTS - 1,)
+
+
+def _spread(batch: int, pool_slots: int) -> tuple[int, ...]:
+ """`batch` distinct slots of a `pool_slots` pool: the top slot (the page's highest
+ addresses) and slot 0 first, then a stride-37 walk over the interior (37 is coprime to
+ `pool_slots - 2`), so batch row `b` is not slot `b`."""
+ rest = tuple((i * 37 + 4) % (pool_slots - 2) + 1 for i in range(max(batch - 2, 0)))
+ slots = ((pool_slots - 1, 0) + rest)[:batch]
+ assert len(set(slots)) == batch
+ return slots
+
+
+def _tep4(batch: int, seqlen: int) -> dict:
+ """A tep4 rank's verify call: 2560 channels of x at row stride 4120, the layer's conv
+ view the last of the manager's 36-layer page of 1874 slots, and its snapshots into a
+ layer's `[128, seqlen, 2560, 3]` view of the stock buffer."""
+ return dict(
+ dim=RANK_DIM,
+ x_row_stride=RANK_X_STRIDE,
+ batch=batch,
+ seqlen=seqlen,
+ pool_slots=TEP4_SLOTS,
+ pool_layers=GDN_PAGE_LAYERS,
+ slots=_spread(batch, TEP4_SLOTS),
+ snapshots=MAX_BATCH_SIZE,
+ snapshot_layers=2,
+ )
+
+
+class CausalConv1dTritonUpdate(OpWrapper):
+ """Convolve each row's new tokens against its cached window, out of place,
+ and advance the window; optionally snapshot the window after every token.
+
+ Row `b`, slot `s = conv_state_indices[b]` (or `b`), `off = 0` (or
+ `num_accepted_tokens[b] - 1`): `window = [conv_state[s, :, off : off+W-1],
+ x[b]]`, `out[b, c, i] = act(bias[c] + sum_j fp32(weight[c, j] * window[i+j]))`
+ -- each product rounded to the operand dtype, accumulated in fp32, SiLU as
+ `acc / (1 + exp(-acc))`, rounded once. The state is then rewritten from
+ column 0: the last `W - 1` window columns (or, with `num_accepted_tokens`,
+ `W - 1 + seqlen - 1` columns rebased past the offset). Columns beyond that
+ are untouched -- `state_len` comes from `W`, never from the pool.
+
+ With `intermediate_conv_window`, snapshot `i` of row `b` --
+ `[intermediate_state_indices[b], i, :, :W-1]` -- is the window after token
+ `i`; after acceptance of `k` tokens the cache is snapshot `k - 1`.
+
+ Fusion boundary: window read, conv, bias, SiLU, state write and snapshots
+ are inside. The projection, the q/k/v split and the post-acceptance
+ promotion of a snapshot into the pool are the caller's.
+
+ Returns a fresh tensor with `x`'s shape (2-D in, 2-D out); `x` is read-only.
+ A row whose slot equals `pad_slot_id` is not written at all -- its result
+ rows are allocator garbage -- and its pool slot and snapshot rows stay
+ bitwise unchanged. `conv_state` and `intermediate_conv_window` are written
+ in place and not returned.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec keys: dim, x_row_stride (default 16480), batch, seqlen, pool_slots,
+ # pool_layers (the pool the last layer's view of a page this many layers
+ # deep, default 1), slots, snapshots (rows of the snapshot buffer),
+ # snapshot_steps (the buffer's step depth, default seqlen), snapshot_layers
+ # (the buffer the first layer's view of one this many layers deep, default
+ # 1), launch_dependent_kernels (default False). Every cell is the call site: `x` the
+ # token-major q|k|v columns of a
+ # [n * p, x_row_stride] projection seen as [n, C, p], bf16, W 4, no bias,
+ # SiLU, pad_slot_id -1, no accepted-token roll.
+ CELLS: tuple[Cell, ...] = tuple(
+ Cell(
+ 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,
+ ),
+ ),
+ # The default target's tep4 member (the MTP unit's): a rank's 4 k / 12 v heads.
+ Cell(
+ why="tep4 verify, 1 row x 4 tok (MTP3) at 2560 / stride 4120, the 1874-slot page",
+ spec=_tep4(1, 4),
+ ),
+ Cell(
+ why="tep4 verify, 1 row x 4 tok, launch_dependent_kernels: the PDL chain into the replay",
+ spec=dict(_tep4(1, 4), launch_dependent_kernels=True),
+ ),
+ Cell(
+ why="tep4 verify, 4 rows x 4 tok, launch_dependent_kernels (largest ll graph batch)",
+ spec=dict(_tep4(4, 4), launch_dependent_kernels=True),
+ ),
+ Cell(
+ why="tep4 verify, 128 rows (max_batch_size) x 4 tok (MTP3) at 2560 / stride 4120",
+ spec=_tep4(128, 4),
+ ),
+ Cell(
+ why="tep4 verify, 128 rows x 2 tok (draft length 1) at 2560 / stride 4120",
+ spec=_tep4(128, 2),
+ ),
+ Cell(
+ why="tep4 verify, 128 rows x 8 tok (draft length 7, the envelope's largest step)",
+ spec=_tep4(128, 8),
+ ),
+ )
+
+ note = """
+ 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, 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 the 461-slot pool
+ with a 32-row snapshot buffer and `x` as the target's q|k|v column view
+ of the projection: 32 rows at seqlen 3 and at seqlen 8 (the widest verify
+ step), and 8 rows at seqlen 5 into an 8-step buffer (a buffer deeper than
+ seqlen keeps its later steps untouched); bf16, W 4, no bias, SiLU only.
+
+ The default target's tep4 member (the MTP unit's verify, `DecodeTarget._gdn_verify`):
+ a rank's 4 k / 12 v heads, so `x` is the `[n, 2560, p]` view of the q|k|v
+ columns of the `[n * p, 4120]` in_proj output (strides `(p * 4120, 1,
+ 4120)`); the pool is one layer's view of the V2 manager's 36-layer conv page
+ at the measured 1874 slots (slot stride 276480, the last layer's view); the
+ snapshots go to a layer's `[128, p, 2560, 3]` view of the GDN speculative
+ state's `[36, max_batch_size, p, 2560, 3]` buffer with
+ `intermediate_state_indices = arange(n)`. Certified at 1 row x 4 tokens and
+ at 128 rows (max_batch_size) x 2, 4 and 8 tokens (draft lengths 1, 3 and 7;
+ 128 x 8 = 1024 rows is the envelope's largest verify step). The result
+ differs from the reference by one bf16 ulp in at most a few elements
+ (relative L2 <= 2.0e-5), the pool and the snapshots are bitwise the
+ reference, and no other layer's conv page or snapshots, nor x's buffer,
+ is written.
+
+ The kernel is Triton, JIT-compiled at the first call; the cells ran that
+ path on sm_103. The x strides, the seqlen and the snapshot strides are
+ `tl.constexpr`, so each (p, row stride) the target sends is its own
+ compile.
+
+ Not certified: a contiguous `x`, other seqlens, pad-sentinel rows (the
+ target never sends one), a dep4 rank (10240 channels at row stride 16480
+ over its 850-slot page) and tp4 (no MTP member). The tep4 slot count is
+ memory-derived and was not measured on an MTP-config engine; the kernel
+ addresses a slot by its index alone, so the pool's size does not change
+ the path.
+
+ Raise loudly, so not guarded: width 1 (triton CompilationError); a missing
+ `intermediate_state_indices` with snapshots and slot indices given.
+
+ Not guarded (content): duplicate slots or duplicate snapshot rows race;
+ `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.
+
+ `launch_dependent_kernels=True` (the target's verify step: the CuTe replay after it
+ launches with PDL and streams its state before `griddepcontrol.wait`) only adds
+ `griddepcontrol.launch_dependents` at the kernel's start: certified at the tep4 shape,
+ 1 and 4 rows x 4 tokens, held to the reference like every other cell.
+
+ Accepted but not exposed: `cache_seqlens` (ignored by the kernel --
+ there is no ring mode), the eagle-tree mask trio, `validate_data`,
+ `metadata`, the `_block_n`/`_num_warps`/`_num_stages` overrides. A
+ `conv_state` dtype different from `x` is accepted and uncertified.
+
+ Hard limits
+ slot offset int64 in the kernel (`conv_state_indices` loaded `.to(tl.int64)`):
+ at the 1874-slot page the last slot's offset is 517.8 M
+ elements, below 2^31 either way.
+ x offset `row * stride_x_seq` in int32: 128 x 8 x 4120 = 4.2 M
+ elements at the envelope, far below 2^31.
+ none probed beyond the cells.
+ """
+
+ def __call__(
+ self,
+ x: torch.Tensor,
+ conv_state: torch.Tensor,
+ weight: torch.Tensor,
+ bias: Optional[torch.Tensor],
+ activation: Union[bool, str, None],
+ conv_state_indices: Optional[torch.Tensor],
+ pad_slot_id: int,
+ num_accepted_tokens: Optional[torch.Tensor] = None,
+ intermediate_conv_window: Optional[torch.Tensor] = None,
+ intermediate_state_indices: Optional[torch.Tensor] = None,
+ launch_dependent_kernels: bool = False,
+ ) -> torch.Tensor:
+ return _causal_conv1d_triton_update(
+ x,
+ conv_state,
+ weight,
+ bias=bias,
+ activation=activation,
+ conv_state_indices=conv_state_indices,
+ num_accepted_tokens=num_accepted_tokens,
+ intermediate_conv_window=intermediate_conv_window,
+ intermediate_state_indices=intermediate_state_indices,
+ pad_slot_id=pad_slot_id,
+ launch_dependent_kernels=launch_dependent_kernels,
+ )
+
+ def reference(
+ self,
+ x: torch.Tensor,
+ conv_state: torch.Tensor,
+ weight: torch.Tensor,
+ bias: Optional[torch.Tensor],
+ activation: Union[bool, str, None],
+ conv_state_indices: Optional[torch.Tensor],
+ pad_slot_id: int,
+ num_accepted_tokens: Optional[torch.Tensor] = None,
+ intermediate_conv_window: Optional[torch.Tensor] = None,
+ intermediate_state_indices: Optional[torch.Tensor] = None,
+ launch_dependent_kernels: bool = False,
+ ) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+ """`(out, conv_state, intermediate_conv_window)` as the call leaves
+ them; mutates nothing. `launch_dependent_kernels` is ignored.
+
+ Rows the op leaves unwritten (padded) are NaN in `out`; the test
+ skips them. Products are formed in the operand dtype, as the kernel
+ forms them -- multiplying in fp32 instead misses the default band on
+ about a third of the elements.
+ """
+ x3 = x.unsqueeze(-1) if x.dim() == 2 else x
+ batch, dim, seqlen = x3.shape
+ width = weight.shape[1]
+ silu = activation in (True, "silu", "swish")
+ state_len = width - 1 + (seqlen - 1 if num_accepted_tokens is not None else 0)
+ out = torch.full(x3.shape, float("nan"), dtype=x.dtype, device=x.device)
+ new_state = conv_state.clone()
+ inter = None if intermediate_conv_window is None else intermediate_conv_window.clone()
+ slots = None if conv_state_indices is None else conv_state_indices.tolist()
+ accepted = None if num_accepted_tokens is None else num_accepted_tokens.tolist()
+ rows = None if intermediate_state_indices is None else intermediate_state_indices.tolist()
+ for b in range(batch):
+ slot = b if slots is None else slots[b]
+ if slots is not None and slot == pad_slot_id:
+ continue
+ off = 0 if accepted is None else accepted[b] - 1
+ state = conv_state[slot]
+ full = torch.cat([state[:, off : off + width - 1], x3[b]], dim=1)
+ acc = torch.zeros(dim, seqlen, dtype=torch.float32, device=x.device)
+ for j in range(width):
+ acc = acc + (weight[:, j].unsqueeze(1) * full[:, j : j + seqlen]).float()
+ if bias is not None:
+ acc = acc + bias.float().unsqueeze(1)
+ if silu:
+ acc = acc / (1.0 + torch.exp(-acc))
+ out[b] = acc.to(x.dtype)
+ if inter is not None:
+ row = b if rows is None else rows[b]
+ for i in range(seqlen):
+ inter[row, i, :dim, : width - 1] = full[:, i + 1 : i + width]
+ if accepted is None:
+ new_state[slot, :, : width - 1] = full[:, seqlen : seqlen + width - 1]
+ else:
+ new_state[slot, :, :state_len] = torch.cat(
+ [state[:, off + 1 : off + width - 1], x3[b]], dim=1
+ )
+ return (out.squeeze(-1) if x.dim() == 2 else out), new_state, inter
+
+ def is_valid(
+ self,
+ x: torch.Tensor,
+ conv_state: torch.Tensor,
+ weight: torch.Tensor,
+ bias: Optional[torch.Tensor],
+ activation: Union[bool, str, None],
+ conv_state_indices: Optional[torch.Tensor],
+ pad_slot_id: int,
+ num_accepted_tokens: Optional[torch.Tensor] = None,
+ intermediate_conv_window: Optional[torch.Tensor] = None,
+ intermediate_state_indices: Optional[torch.Tensor] = None,
+ launch_dependent_kernels: bool = False,
+ ) -> None:
+ dim = x.shape[1]
+ seqlen = 1 if x.dim() == 2 else x.shape[2]
+ width = weight.shape[1]
+ state_len = width - 1 + (seqlen - 1 if num_accepted_tokens is not None else 0)
+ # The tap loop is specialised for W 2..4; W 1 fails to compile.
+ assert 2 <= width <= 4, (
+ f"weight width must be 2, 3 or 4; got {width} -- width 5 is accepted "
+ "by the kernel and silently computes a wrong answer"
+ )
+ # Reads and writes `state_len` columns unmasked against the pool.
+ assert conv_state.shape[2] >= state_len, (
+ f"conv_state must hold at least {state_len} columns (width={width}, "
+ f"seqlen={seqlen}, num_accepted_tokens="
+ f"{'given' if num_accepted_tokens is not None else 'None'}); got "
+ f"{conv_state.shape[2]} -- a shorter pool is accepted and writes past "
+ "the slot into the next one"
+ )
+ # The channel axis is masked against x's dim only.
+ assert conv_state.shape[1] >= dim and weight.shape[0] >= dim, (
+ f"conv_state and weight must cover all {dim} channels of x; got "
+ f"{conv_state.shape[1]} and {weight.shape[0]} -- a narrower pool or "
+ "weight is accepted and reads/writes out of bounds"
+ )
+ if intermediate_conv_window is not None:
+ # No bound on the step axis: a shallower buffer is written into
+ # the next row.
+ assert (
+ intermediate_conv_window.shape[1] >= seqlen
+ and intermediate_conv_window.shape[2] >= dim
+ and intermediate_conv_window.shape[3] >= width - 1
+ ), (
+ f"intermediate_conv_window must be at least [rows, {seqlen}, {dim}, "
+ f"{width - 1}]; got {tuple(intermediate_conv_window.shape)} -- an "
+ "undersized buffer is accepted and writes past the row it was given"
+ )
+
+
+causal_conv1d_triton_update = CausalConv1dTritonUpdate()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_update.py
new file mode 100644
index 000000000000..f8f98d2a9807
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_update.py
@@ -0,0 +1,588 @@
+# 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.
+
+The decode step's shape (one token a row, SiLU, no bias, a width-4 conv over a
+3-column bf16 pool, slot indices, a channel count that is a multiple of 1024)
+from `_CUTE_MIN_ROWS` rows runs a CuTe DSL kernel that is bitwise the TRT-LLM
+kernel; every other call is `torch.ops.trtllm.causal_conv1d_update`.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass._mlir.dialects import llvm
+from cutlass.cute.runtime import from_dlpack
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
+
+#: This checkpoint's GDN `[q|k|v]` width; the conv is width 4 over a 3-column pool.
+CKPT_DIM = 10240
+#: Largest conv-pool slot count the cache manager allocates for this checkpoint at a batch
+#: of at most 32. The decode step hands the whole pool over unsliced.
+ENGINE_SLOTS = 461
+
+#: A tp4/tep4 rank's GDN conv width (`[q | k | v]` of 4 key and 12 value heads of 128) and
+#: the row stride of its input-projection output, `[q | k | v | z | a | b]` = 2560 + 1536 +
+#: 12 + 12; at tp1 and dep4 the projection is 16480 wide.
+RANK_DIM = 2560
+RANK_X_STRIDE = 4120
+CKPT_X_STRIDE = 16480
+#: Per-rank conv-pool slot counts of the V2 hybrid manager at bf16 KV and recurrent state,
+#: KV fraction 0.6, `max_batch_size` 128, `max_seq_len` 9216 (the PLE and GDN pools share
+#: one slot space). The count is memory-derived, so a dep4 rank, which keeps every GDN head,
+#: holds fewer than a tep4 rank. tp4 has no measured count; its rank holds tep4's conv
+#: geometry, so tep4's count stands in for it.
+TEP4_SLOTS = 1874
+DEP4_SLOTS = 850
+TP4_SLOTS = TEP4_SLOTS
+#: GDN layers whose conv states share one slot's page in the V2 manager: a layer's view
+#: `[slots, dim, 3]` has slot stride 36 * dim * 3 (276480 at 2560, 1105920 at 10240) and
+#: layer `i` sits `i * dim * 3` elements into the page (at tp4, tep4 and dep4).
+GDN_PAGE_LAYERS = 36
+
+
+# --------------------------------------------------------------------------------------
+# The decode shape in CuTe DSL. 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; below it the TRT-LLM kernel's
+#: single wave is faster.
+_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."""
+ return (0,) + tuple(i * 14 + 3 for i in range(1, batch - 1)) + (ENGINE_SLOTS - 1,)
+
+
+def _spread(batch: int, pool_slots: int, pad_every: int = 0) -> tuple[int, ...]:
+ """`batch` distinct slots of a `pool_slots` pool: the top slot (the page's highest
+ addresses) and slot 0 first, then a stride-37 walk over the interior (37 is coprime to
+ every `pool_slots - 2` here), so batch row `b` is not slot `b`. With `pad_every`, rows
+ `5, 5 + pad_every, ...` take the pad slot -1 instead."""
+ rest = tuple((i * 37 + 4) % (pool_slots - 2) + 1 for i in range(max(batch - 2, 0)))
+ slots = ((pool_slots - 1, 0) + rest)[:batch]
+ assert len(set(slots)) == batch
+ if pad_every:
+ slots = tuple(-1 if i % pad_every == 5 else s for i, s in enumerate(slots))
+ return slots
+
+
+def _rank(batch: int, pool_slots: int, pad_every: int = 0) -> dict:
+ """A tp4/tep4 rank's call: 2560 conv channels at x row stride 4120, the layer's view the
+ last of the manager's 36-layer conv page."""
+ return dict(
+ dim=RANK_DIM,
+ x_row_stride=RANK_X_STRIDE,
+ batch=batch,
+ pool_slots=pool_slots,
+ slots=_spread(batch, pool_slots, pad_every),
+ pool_layers=GDN_PAGE_LAYERS,
+ )
+
+
+def _dep4(batch: int) -> dict:
+ """A dep4 rank's call: tp1's 10240 channels at row stride 16480, the layer's view the
+ last of the manager's 36-layer conv page."""
+ return dict(
+ dim=CKPT_DIM,
+ x_row_stride=CKPT_X_STRIDE,
+ batch=batch,
+ pool_slots=DEP4_SLOTS,
+ slots=_spread(batch, DEP4_SLOTS),
+ pool_layers=GDN_PAGE_LAYERS,
+ )
+
+
+class CausalConv1dUpdate(OpWrapper):
+ """One causal depthwise conv step per batch row against that row's cached
+ window, then the window's advance -- `x` and the pool both written in place.
+
+ Row `b`, slot `s = conv_state_indices[b]` (or `b`), new tokens `i < seqlen`:
+ `history = [conv_state[s, c, :], x[b, c, :]]`, the window is its last
+ `W - 1 + seqlen` columns, `x[b, c, i] = act(bias[c] + sum_j weight[c, j] *
+ window[i + j])` in fp32 rounded once, and `conv_state[s] = history[-state_len:]`
+ with `state_len = conv_state.shape[2]`. With `cache_seqlens` the pool is a
+ ring instead: position `p` lives at `p % state_len`, `cache_seqlens[b]` is
+ the first new token's absolute position, and nothing shifts.
+
+ Fusion boundary: window read, convolution, bias, SiLU and state write are
+ inside. The projection producing `x`, slot allocation and the q/k/v split
+ are the caller's.
+
+ Returns `x` -- the same object, written in place. A 2-D `[batch, dim]` `x`
+ is handed to the op as an `unsqueeze(-1)` view (metadata only; the op
+ writes through it). `conv_state` is advanced in place and not returned.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec keys: dim, x_row_stride (default 16480), batch, pool_slots, slots
+ # ("random" | tuple), pool_layers (the pool the last layer's view of a buffer
+ # this many layers deep, default 1). Every cell is the call site: x the
+ # q|k|v column view of the decode rows of a [rows, x_row_stride] in_proj
+ # output, bf16, W 4, a 3-column pool, no bias, SiLU, no ring, pad_slot_id -1.
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="x the q|k|v column view of the [rows, 16480] in_proj, batch 32",
+ spec=dict(dim=CKPT_DIM, batch=32, pool_slots=ENGINE_SLOTS, slots=_pinned(32)),
+ ),
+ Cell(
+ why="x the q|k|v column view of the [rows, 16480] in_proj, batch 128",
+ spec=dict(dim=CKPT_DIM, batch=128, pool_slots=256, slots="random"),
+ ),
+ Cell(
+ why="64 rows (the CuTe kernel's first) over a layer view, pad -1 rows untouched",
+ spec=dict(
+ dim=CKPT_DIM,
+ batch=64,
+ pool_slots=ENGINE_SLOTS,
+ slots=tuple(-1 if i % 21 == 5 else i * 7 for i in range(64)),
+ pool_layers=2,
+ ),
+ ),
+ # The default target's members, each over its rank's measured conv page. Every
+ # 2560-channel call runs the TRT-LLM kernel (2560 is not a multiple of the
+ # CuTe kernel's 1024 channels a block row).
+ Cell(
+ why="tep4 decode, 1 row at 2560 / stride 4120: the top slot of the 1874-slot page",
+ spec=_rank(1, TEP4_SLOTS),
+ ),
+ Cell(
+ why="tep4, 64 rows at 2560 / stride 4120 over the 1874-slot page",
+ spec=_rank(64, TEP4_SLOTS),
+ ),
+ Cell(
+ why="tep4, 128 rows (max_batch_size) at 2560 / stride 4120 over the 1874-slot page",
+ spec=_rank(128, TEP4_SLOTS),
+ ),
+ Cell(
+ why="tp4, 128 rows at 2560 / stride 4120 over tp4's page, pad rows (-1) on the "
+ "TRT-LLM kernel",
+ spec=_rank(128, TP4_SLOTS, pad_every=21),
+ ),
+ Cell(
+ why="dep4 decode, 32 rows at 10240 / stride 16480: the top slot of the 850-slot page",
+ spec=_dep4(32),
+ ),
+ Cell(
+ why="dep4, 128 rows (max_batch_size, the CuTe kernel) over the 850-slot page",
+ spec=_dep4(128),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none.
+
+ Call site: the GDN decode step -- 2-D `[num_decodes, 10240]` x, the q|k|v
+ column view (row stride 16480) of the decode rows of the layer's one
+ `[T, 16480]` input-projection output (the op writes those columns only);
+ the layer's whole conv pool `[slots, 10240, 3]`, W 4, no bias, SiLU on,
+ `cache_seqlens` None, int32 slots, `pad_slot_id` -1.
+ That exact shape (silu, no ring, `state_len == W - 1`) runs TRT-LLM's
+ dedicated seqlen-1 kernel; any other shape runs its general one, which the
+ target never reaches and which is not certified. From `_CUTE_MIN_ROWS`
+ (64) rows, with bf16 operands, a contiguous weight, a pool whose slot is
+ `[dim, 3]` contiguous and 16-byte aligned rows, it runs
+ `_DecodeConvKernel` (CuTe DSL) instead: 8 channels a thread by 16-byte
+ loads, the TRT-LLM kernel's instructions, bitwise its result (x and pool,
+ pad rows untouched). Certified at that call: 32 rows (the TRT-LLM kernel)
+ over the 461-slot pool, 128 rows (the CuTe kernel) over a 256-slot pool,
+ and 64 rows (the CuTe kernel's first row count) with pad rows over the
+ 461-slot pool seen as one layer of a two-layer buffer.
+
+ The op checks metadata itself, loudly: width 2..4, `state_len >= W - 1`,
+ one dtype across x/conv_state/weight (bias matches weight), weight shape,
+ int32 contiguous length-batch `conv_state_indices`, int32
+ `cache_seqlens`, a batch-sized pool when indices are absent. So no guard.
+
+ Not guarded because it is content, not metadata: two rows naming one slot
+ race -- the slot ends up matching neither row's expected state while both
+ outputs match. Slot uniqueness is the caller's.
+ Out-of-range slot values are an illegal access, not probed.
+
+ A row whose index equals `pad_slot_id` (-1 at the call site) is skipped:
+ its x rows and the pool stay bitwise unchanged (certified on both kernels:
+ the CuTe one at 64 rows of 10240, the TRT-LLM one at tp4's 128 rows).
+
+ The default target's call site, per rank. tp4 / tep4: x `[num_decodes, 2560]`,
+ the q|k|v view (row stride 4120) of the rank's `[T, 4120]` projection; the
+ pool one layer's view `[slots, 2560, 3]` of the V2 manager's conv page,
+ 36 GDN layers a slot (slot stride 276480). dep4: tp1's 10240 channels at
+ stride 16480, page slot stride 1105920. Certified at the measured per-rank
+ slot counts (`TEP4_SLOTS` 1874, `DEP4_SLOTS` 850; tp4 has none measured and
+ takes tep4's, `TP4_SLOTS`), the view the page's last layer and the top slot
+ named, so a cell reaches the pool's highest address: tep4 / tp4 at 1, 64
+ and 128 rows, dep4 at 32 and 128 rows.
+
+ Which kernel runs: 2560 is not a multiple of `_CH * _THREADS` (1024), so
+ every tp4 / tep4 call runs TRT-LLM's seqlen-1 kernel
+ (`causal_conv1d_update_kernel_sl1`, one thread per (row, channel), 20
+ blocks of 128 channels a row), whatever the row count. dep4 runs the CuTe
+ kernel from 64 rows, as tp1 does.
+
+ Sources: the TRT-LLM kernel is compiled into the wheel
+ (`cpp/tensorrt_llm/kernels/causalConv1d/causalConv1d.cu`, registered by
+ `cpp/tensorrt_llm/thop/causalConv1dOp.cpp`); the CuTe kernel is JIT-compiled
+ by `cute.compile` on its first call for the running device, sm_103 here.
+
+ Hard limits (from source, not probed: the measured pools stay far below)
+ The TRT-LLM kernel forms a slot's offset as `int slot * uint32_t
+ conv_state_batch_stride`, 32-bit unsigned elements: it wraps from slot
+ 15535 at stride 276480 (tp4 / tep4) and slot 3884 at 1105920 (dep4).
+ The measured pools reach 518,123,519 (tep4) and 940,031,999 (dep4)
+ elements, under 2^31 elements and bytes. The CuTe kernel forms slot and
+ row offsets in 64 bits.
+ """
+
+ def __call__(
+ self,
+ x: torch.Tensor,
+ conv_state: torch.Tensor,
+ weight: torch.Tensor,
+ bias: Optional[torch.Tensor],
+ silu_activation: bool,
+ cache_seqlens: Optional[torch.Tensor],
+ conv_state_indices: Optional[torch.Tensor],
+ pad_slot_id: int,
+ ) -> torch.Tensor:
+ 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,
+ weight,
+ bias,
+ silu_activation,
+ cache_seqlens,
+ conv_state_indices,
+ pad_slot_id,
+ )
+ return x
+
+ def reference(
+ self,
+ x: torch.Tensor,
+ conv_state: torch.Tensor,
+ weight: torch.Tensor,
+ bias: Optional[torch.Tensor],
+ silu_activation: bool,
+ cache_seqlens: Optional[torch.Tensor],
+ conv_state_indices: Optional[torch.Tensor],
+ pad_slot_id: int,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """`(x, conv_state)` as the call leaves them; mutates nothing."""
+ x3 = x.unsqueeze(-1) if x.dim() == 2 else x
+ batch, dim, seqlen = x3.shape
+ width = weight.shape[1]
+ state_len = conv_state.shape[2]
+ out = x3.clone()
+ new_state = conv_state.clone()
+ slots = None if conv_state_indices is None else conv_state_indices.tolist()
+ starts = None if cache_seqlens is None else cache_seqlens.tolist()
+ for b in range(batch):
+ slot = b if slots is None else slots[b]
+ if slots is not None and slot == pad_slot_id:
+ continue
+ state = conv_state[slot].float()
+ tokens = x3[b].float()
+ if starts is None:
+ history = torch.cat([state, tokens], dim=1)
+ window = history[:, state_len - (width - 1) :]
+ new_state[slot] = history[:, -state_len:].to(conv_state.dtype)
+ else:
+ start = starts[b] % state_len
+ read = [(start - (width - 1) + i) % state_len for i in range(width - 1)]
+ window = torch.cat([state[:, read], tokens], dim=1)
+ advanced = state.clone()
+ for i in range(seqlen):
+ advanced[:, (start + i) % state_len] = tokens[:, i]
+ new_state[slot] = advanced.to(conv_state.dtype)
+ acc = torch.zeros(dim, seqlen, dtype=torch.float32, device=x.device)
+ for j in range(width):
+ acc = acc + weight[:, j].float().unsqueeze(1) * window[:, j : j + seqlen]
+ if bias is not None:
+ acc = acc + bias.float().unsqueeze(1)
+ if silu_activation:
+ acc = acc * torch.sigmoid(acc)
+ out[b] = acc.to(x.dtype)
+ return (out.squeeze(-1) if x.dim() == 2 else out), new_state
+
+
+causal_conv1d_update = CausalConv1dUpdate()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/cu_seqlens_from_seq_lens.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/cu_seqlens_from_seq_lens.py
new file mode 100644
index 000000000000..c94f9601b4d3
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/cu_seqlens_from_seq_lens.py
@@ -0,0 +1,180 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Prepend-zero cumulative sum of a per-request query-row count: the
+``cu_seqlens`` segmentation vector every varlen recurrent op takes.
+
+A macro entry composed from native torch (``torch.zeros`` + ``torch.cumsum``).
+"""
+
+from __future__ import annotations
+
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+#: Integer dtypes accepted on either side. ``torch.cumsum`` silently accepts
+#: bool and floating dtypes and answers wrongly with them.
+_INTEGER_DTYPES = (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64)
+
+#: Batch shapes this checkpoint's serving path produces 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 at tp1 reuses the
+#: recurrent metadata's vector instead, and ``decode_p1_128`` certifies the call
+#: there all the same.
+#: The ``_128`` batches are the serving envelope's widest step at every member of
+#: the default target (``max_batch_size`` 128 requests per rank, ``max_num_tokens``
+#: 8192 rows): ``mixed_p8_128`` is one context chunk beside 127 MTP-7 verifies,
+#: 7176 + 127 * 8 = 8192 rows.
+BATCHES = {
+ "mixed_p2": (1024, 37, 2, 2, 2, 2),
+ "padded_decode_p2": (2, 2, 2, 0, 0, 0),
+ "decode_p1_128": (1,) * 128,
+ "verify_p8_128": (8,) * 128,
+ "chunk_8192": (8192,),
+ "mixed_p8_128": (7176,) + (8,) * 127,
+}
+
+
+class CuSeqlensFromSeqLens(OpWrapper):
+ """`out[0] = 0`, `out[s] = seq_lens[0] + ... + seq_lens[s - 1]`: a fresh
+ contiguous `[num_seqs + 1]` tensor of `dtype` on `seq_lens.device`.
+
+ The scan accumulates in `dtype`, not in `seq_lens`'s dtype. Fusion boundary:
+ the leading zero, the scan and the dtype. Choosing which length vector to
+ scan, and keeping the packed activation exactly `out[-1]` rows long, are the
+ caller's. `seq_lens` is read-only.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: lengths, in_dtype, out_dtype; `layout` "contiguous" (default) or
+ # "prefix" (the first len(lengths) of a longer buffer -- the call site's
+ # `seq_lens_cuda[:batch]`).
+ CELLS: tuple[Cell, ...] = (
+ 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="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="padded_decode_p2: zero-length padded rows, int32->int64",
+ spec=dict(
+ lengths=BATCHES["padded_decode_p2"], in_dtype=torch.int32, out_dtype=torch.int64
+ ),
+ ),
+ Cell(
+ why="decode at con 128: 128 one-row requests, prefix int32->int64",
+ spec=dict(
+ lengths=BATCHES["decode_p1_128"],
+ in_dtype=torch.int32,
+ out_dtype=torch.int64,
+ layout="prefix",
+ ),
+ ),
+ Cell(
+ why="MTP-7 verify at con 128: 128 x 8 rows, prefix int32->int64",
+ spec=dict(
+ lengths=BATCHES["verify_p8_128"],
+ in_dtype=torch.int32,
+ out_dtype=torch.int64,
+ layout="prefix",
+ ),
+ ),
+ Cell(
+ why="one 8192-row prefill chunk (max_num_tokens), prefix int32->int64",
+ spec=dict(
+ lengths=BATCHES["chunk_8192"],
+ in_dtype=torch.int32,
+ out_dtype=torch.int64,
+ layout="prefix",
+ ),
+ ),
+ Cell(
+ why="widest step: 7176-row chunk + 127 MTP-7 verifies = 8192 rows, 128 requests",
+ spec=dict(
+ lengths=BATCHES["mixed_p8_128"],
+ in_dtype=torch.int32,
+ out_dtype=torch.int64,
+ layout="prefix",
+ ),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none. The call site reads `attn_metadata.seq_lens_cuda`
+ prefix, which still holds the true per-request query-row count during the
+ trunk's forward (the spec worker rewrites it only afterwards). Not the
+ recurrent metadata's `query_start_loc_long`: that one has one row per
+ generation sequence and is short under MTP; the two agree exactly when
+ every generation request contributes one row.
+
+ Preconditions, none checked (data, not metadata -- a check would sync):
+ entries non-negative (a negative length is silently wrong or a
+ device-side assert in the consumers); a padded row is length 0, never -1;
+ `out[-1]` equals the packed activation's row count.
+
+ `dtype` has no default: `causal_conv1d_fwd` requires int32 and
+ `chunk_gated_delta_rule` keys its autotuner on `cu_seqlens`' dtype.
+
+ Certified bitwise at the call site's prefix view of `seq_lens_cuda` (a
+ mixed step at p 2, and the p 8 verify of 32 requests) and at a decode
+ batch padded with zero-length rows, int32 in and int64 out. Also at a
+ rank's serving envelope (`max_batch_size` 128, `max_num_tokens` 8192, the
+ same at tp4, tep4 and dep4): 128 one-row requests, 128 MTP-7 verifies of
+ 8 rows, one 8192-row chunk, and a 7176-row chunk beside 127 verifies.
+ CUDA-graph capture is certified at 32 requests (int32 out) and 128
+ requests (int64 out): only `seq_lens.shape[0]` reaches Python, and a
+ captured graph follows lengths rewritten in place, up to a mixed step of
+ 8192 rows.
+
+ Composition (int32 in, int64 out): four kernels -- a fill, the int32 ->
+ int64 cast copy, and cub's scan-init and scan, all compiled into torch (no
+ JIT path).
+
+ Not certified: a `seq_lens` dtype other than int32 (the guard admits any
+ integer dtype), a non-contiguous `seq_lens` view, more than 128 requests,
+ and a step over 8192 rows (an int64 `out` holds any reachable total).
+
+ Hard limits
+ out[-1] accumulated in `dtype` boundary 2^31 - 1 for int32
+ int32 wraps silently ([2e9, 2e9] gives [0, 2000000000, -294967296]);
+ int64 does not. Not guarded (data).
+ """
+
+ def __call__(self, seq_lens: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+ out = torch.zeros(seq_lens.shape[0] + 1, dtype=dtype, device=seq_lens.device)
+ torch.cumsum(seq_lens, dim=0, dtype=dtype, out=out[1:])
+ return out
+
+ def reference(self, seq_lens: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+ """A scalar Python prefix sum; no torch op computes it."""
+ total, out = 0, [0]
+ for length in seq_lens.tolist():
+ total += int(length)
+ out.append(total)
+ return torch.tensor(out, dtype=dtype, device=seq_lens.device)
+
+ def is_valid(self, seq_lens: torch.Tensor, dtype: torch.dtype) -> None:
+ # Each of these answers wrongly without an error: a 2-D input gives a
+ # short or nonsense vector, a float input is truncated, a bool input is
+ # counted as nonzero rows, and a float `dtype` returns a float vector.
+ assert seq_lens.dim() == 1, (
+ f"seq_lens must be a 1-D per-request row count, got shape {tuple(seq_lens.shape)}"
+ )
+ assert seq_lens.dtype in _INTEGER_DTYPES, (
+ f"seq_lens must have an integer dtype, got {seq_lens.dtype}"
+ )
+ assert dtype in _INTEGER_DTYPES, f"dtype must be an integer dtype, got {dtype}"
+
+
+cu_seqlens_from_seq_lens = CuSeqlensFromSeqLens()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/flashinfer_chunk_gated_delta_rule.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/flashinfer_chunk_gated_delta_rule.py
new file mode 100644
index 000000000000..85a0d7c477f9
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/flashinfer_chunk_gated_delta_rule.py
@@ -0,0 +1,626 @@
+# 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 ``[N, HV, V, K]`` buffer in the pool's dtype (the sm_100/103 kernel keeps
+the state in fp32 inside and reads and writes it in that dtype), 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. With ``has_initial_state`` the call is the same
+sequence in this module, its gather a Triton kernel that writes zeros for a sequence
+carrying nothing instead of reading its slot, so the caller need not zero the slot first.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch.modules.fla.flashinfer_chunk import (
+ chunk_gated_delta_rule as _fi_chunk_gated_delta_rule,
+)
+from tensorrt_llm._torch.modules.fla.fused_state_io import cast_scatter_fp32_vk_to_vk
+
+#: 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 recurrent-pool slot count the cache manager allocates for
+#: qwen3.8-flash-next-nvfp4.
+ENGINE_SLOTS = 461
+
+
+#: A tp4/tep4 rank's Gated DeltaNet: tp1's 16 key / 48 value heads split four ways, the
+#: key-to-value head ratio (3) unchanged. A dep4 rank keeps `QWEN38_GEOMETRY`.
+RANK4_GEOMETRY = (4, 12, 128, 128)
+#: Per-rank recurrent-pool slot counts of `Qwen4ExpHybridCacheManager` on this checkpoint
+#: at bf16 KV and state, KV fraction 0.6, `max_batch_size` 128, `max_seq_len` 9216, TRTLLM
+#: MoE and the n-gram table host-resident (the PLE and GDN pools share one slot space). The
+#: V2 manager sizes the pool from the rank's memory budget, not from `max_batch_size`; a
+#: dep4 rank keeps every GDN head, so it holds fewer slots than a tep4 rank.
+TEP4_SLOTS = 1874
+DEP4_SLOTS = 850
+#: tp4 has no measured count. A tp4 rank's GDN state per slot is a tep4 rank's, so tep4's
+#: count stands in for it.
+TP4_SLOTS = TEP4_SLOTS
+#: GDN layers whose states share one slot page in the stock manager (48 layers, every
+#: 4th full attention): the pool one layer receives is a view whose slot stride is 36
+#: `[HV, V, K]` states, so a tep4 rank's last slot starts 1.3e10 elements past the base.
+PAGE_LAYERS = 36
+
+
+def _engine_slots(num_seqs: int) -> tuple[int, ...]:
+ """`num_seqs` distinct slots of the engine pool, slot 0 and the last live."""
+ return ((0, ENGINE_SLOTS - 1) + tuple(i * 13 + 4 for i in range(num_seqs)))[:num_seqs]
+
+
+def _rank_slots(num_seqs: int, slots: int) -> tuple[int, ...]:
+ """`num_seqs` distinct slots of a `slots`-slot pool: the last and the first, then a
+ stride-37 walk over the interior (37 is coprime to 1874 - 2 and 850 - 2)."""
+ rest = tuple((i * 37 + 4) % (slots - 2) + 1 for i in range(max(num_seqs - 2, 0)))
+ picked = ((slots - 1, 0) + rest)[:num_seqs]
+ assert len(set(picked)) == num_seqs
+ return picked
+
+
+def _rank_cell(why: str, geometry, seqlens, slots: int, initial, fi_path: str) -> Cell:
+ """The target's call (`PrefillTarget._gdn_core`) at one rank's pool: the stock
+ manager's bf16 state page at the member's slot count, the layer's view of it, a
+ prefix view of a wider output buffer, `has_initial_state`, `linear_gate`. `fi_path` is
+ the FlashInfer kernel the call reaches on sm_103: `cp` (the context-parallel pipeline)
+ or `chunked`."""
+ return Cell(
+ why=why,
+ spec=dict(
+ geometry=geometry,
+ seqlens=seqlens,
+ slots=slots,
+ page_layers=PAGE_LAYERS,
+ indices=_rank_slots(len(seqlens), slots),
+ pool_dtype=torch.bfloat16,
+ output="slice",
+ initial=initial,
+ linear_gate=True,
+ fi_path=fi_path,
+ ),
+ )
+
+
+#: Elements of a `[V, K]` state plane a gather program copies.
+_GATHER_BLOCK = 4096
+
+
+@triton.jit
+def _gather_carried_states_kernel(
+ pool_ptr,
+ dst_ptr,
+ slot_ptr,
+ init_ptr,
+ stride_ps,
+ stride_ph,
+ stride_pv,
+ stride_pk,
+ H: tl.constexpr,
+ K: tl.constexpr,
+ PLANE: tl.constexpr,
+ BLOCK: tl.constexpr,
+):
+ """One program: `BLOCK` elements of sequence `n`'s head `h` plane, into the packed
+ `[N, H, V, K]` buffer: the pool row `slot_ptr[n]` when `init_ptr[n]`, else +0.0 (the
+ pool row is not read)."""
+ n = tl.program_id(0)
+ h = tl.program_id(1)
+ i = tl.program_id(2) * BLOCK + tl.arange(0, BLOCK)
+ mask = i < PLANE
+ dst = dst_ptr + (n.to(tl.int64) * H + h) * PLANE + i
+ if tl.load(init_ptr + n) != 0:
+ slot = tl.load(slot_ptr + n).to(tl.int64)
+ src = (
+ pool_ptr + slot * stride_ps + h * stride_ph + (i // K) * stride_pv + (i % K) * stride_pk
+ )
+ tl.store(dst, tl.load(src, mask=mask), mask=mask)
+ else:
+ tl.store(dst, tl.zeros([BLOCK], dtype=dst_ptr.dtype.element_ty), mask=mask)
+
+
+def _gather_carried_states(
+ pool: torch.Tensor, indices: torch.Tensor, has_initial_state: torch.Tensor
+) -> torch.Tensor:
+ """`[N, H, V, K]` in the pool's dtype: `pool[indices[n]]` for a carried sequence, zeros
+ for a fresh one."""
+ num_seqs = indices.shape[0]
+ _, h, v, k = pool.shape
+ out = torch.empty(num_seqs, h, v, k, dtype=pool.dtype, device=pool.device)
+ _gather_carried_states_kernel[(num_seqs, h, triton.cdiv(v * k, _GATHER_BLOCK))](
+ pool, out, indices, has_initial_state, *pool.stride(),
+ H=h, K=k, PLANE=v * k, BLOCK=_GATHER_BLOCK, num_warps=4,
+ ) # fmt: skip
+ return out
+
+
+#: Context tokens up to which a `has_initial_state` call runs in place over the pool: one
+#: chunk (64 tokens) a sequence at most, where FlashInfer's chunked kernel is a single
+#: launch and its context-parallel pipeline (five launches plus the gather and the scatter)
+#: buys nothing.
+INPLACE_MAX_TOKENS = 64
+
+
+@triton.jit
+def _zero_fresh_states_kernel(
+ pool_ptr,
+ slot_ptr,
+ init_ptr,
+ stride_ps,
+ stride_ph,
+ stride_pv,
+ stride_pk,
+ K: tl.constexpr,
+ PLANE: tl.constexpr,
+ BLOCK: tl.constexpr,
+):
+ """One program: `BLOCK` elements of sequence `n`'s head `h` plane of the pool row
+ `slot_ptr[n]`, zeroed when `init_ptr[n]` is false; a carried sequence's row is not
+ touched."""
+ n = tl.program_id(0)
+ if tl.load(init_ptr + n) == 0:
+ h = tl.program_id(1)
+ i = tl.program_id(2) * BLOCK + tl.arange(0, BLOCK)
+ slot = tl.load(slot_ptr + n).to(tl.int64)
+ dst = (
+ pool_ptr + slot * stride_ps + h * stride_ph + (i // K) * stride_pv + (i % K) * stride_pk
+ )
+ tl.store(dst, tl.zeros([BLOCK], dtype=pool_ptr.dtype.element_ty), mask=i < PLANE)
+
+
+def _chunk_inplace(
+ q, k, v, g, beta, scale, pool, indices, has_initial_state, cu_seqlens, output, linear_gate
+) -> torch.Tensor:
+ """The short in-place indexed call: the fresh sequences' pool rows zeroed (one launch),
+ then FlashInfer's chunked kernel (`use_cp=False`) reading and writing the named rows of
+ the pool itself (`state_indices`, `output_state` the pool): no packed copy in or out."""
+ import flashinfer
+
+ num_seqs = indices.shape[0]
+ _, h, v_dim, k_dim = pool.shape
+ _zero_fresh_states_kernel[(num_seqs, h, triton.cdiv(v_dim * k_dim, _GATHER_BLOCK))](
+ pool, indices, has_initial_state, *pool.stride(),
+ K=k_dim, PLANE=v_dim * k_dim, BLOCK=_GATHER_BLOCK, num_warps=4,
+ ) # fmt: skip
+ q3, k3, v3 = q.squeeze(0).contiguous(), k.squeeze(0).contiguous(), v.squeeze(0).contiguous()
+ out = (
+ output.squeeze(0)
+ if output is not None
+ else q3.new_empty(q3.shape[0], v3.shape[1], v3.shape[2])
+ )
+ flashinfer.chunk_gated_delta_rule(
+ q=q3,
+ k=k3,
+ v=v3,
+ g=g.squeeze(0) if linear_gate else torch.exp(g.squeeze(0)),
+ beta=beta.squeeze(0).to(torch.float32),
+ scale=scale,
+ initial_state=pool,
+ output_final_state=True,
+ cu_seqlens=cu_seqlens,
+ use_qk_l2norm_in_kernel=False,
+ output=out,
+ output_state=pool,
+ use_cp=False,
+ state_indices=indices,
+ )
+ return out.unsqueeze(0)
+
+
+def _chunk_carried(
+ q, k, v, g, beta, scale, pool, indices, has_initial_state, cu_seqlens, output, linear_gate
+) -> torch.Tensor:
+ """The in-place indexed call of the adapter, its gather `_gather_carried_states`: the
+ packed states in the pool's dtype (the sm_100/103 kernel keeps fp32 inside), FlashInfer's
+ kernel (`g` exponentiated first unless `linear_gate`), and the fused scatter of the final
+ states into the same slots."""
+ import flashinfer
+
+ q3, k3, v3 = q.squeeze(0).contiguous(), k.squeeze(0).contiguous(), v.squeeze(0).contiguous()
+ packed = _gather_carried_states(pool, indices, has_initial_state)
+ out = (
+ output.squeeze(0)
+ if output is not None
+ else q3.new_empty(q3.shape[0], v3.shape[1], v3.shape[2])
+ )
+ final = torch.empty_like(packed)
+ flashinfer.chunk_gated_delta_rule(
+ q=q3,
+ k=k3,
+ v=v3,
+ g=g.squeeze(0) if linear_gate else torch.exp(g.squeeze(0)),
+ beta=beta.squeeze(0).to(torch.float32),
+ scale=scale,
+ initial_state=packed,
+ output_final_state=True,
+ cu_seqlens=cu_seqlens,
+ use_qk_l2norm_in_kernel=False,
+ output=out,
+ output_state=final,
+ )
+ cast_scatter_fp32_vk_to_vk(final, pool, indices)
+ return out.unsqueeze(0)
+
+
+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), or with `linear_gate` (the
+ `has_initial_state` call only) the factor `exp(g)` itself, as
+ `gdn_prefill_conv_post(linear_gate=True)` returns 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]]`, or, when `has_initial_state`
+ (bool `[N]`) is given and `has_initial_state[n]` is false, from zeros, its
+ slot not read (the in-place indexed call only). 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 (zeros for a fresh sequence with `has_initial_state`), 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); pool_dtype (fp32); output (None | "buffer" | "slice", a prefix
+ # view of a wider buffer); page_layers (the pool is the layer's view of the stock
+ # manager's state page, that many GDN states a slot); fi_path (the FlashInfer kernel
+ # reached). Every cell is in place: tp1's at 16/48, then one rank's of tp4/tep4/dep4.
+ CELLS: tuple[Cell, ...] = tuple(
+ Cell(
+ why=f"qwen3.8 prefill into output=, 461-slot pool, seqlens {list(s)}",
+ spec=dict(
+ geometry=QWEN38_GEOMETRY,
+ seqlens=s,
+ slots=ENGINE_SLOTS,
+ indices=_engine_slots(len(s)),
+ output="buffer",
+ ),
+ )
+ for s in ((200, 37, 91), (8192,))
+ ) + (
+ Cell(
+ why="qwen3.8 context chunk + 1-token rows [8, 1, 1, 1], 461-slot pool",
+ spec=dict(
+ geometry=QWEN38_GEOMETRY,
+ seqlens=(8, 1, 1, 1),
+ slots=ENGINE_SLOTS,
+ indices=_engine_slots(4),
+ ),
+ ),
+ Cell(
+ why="bf16 pool (mamba_ssm_cache_dtype), 461 slots, seqlens [200, 37, 91]",
+ spec=dict(
+ geometry=QWEN38_GEOMETRY,
+ seqlens=(200, 37, 91),
+ slots=ENGINE_SLOTS,
+ indices=_engine_slots(3),
+ pool_dtype=torch.bfloat16,
+ ),
+ ),
+ Cell(
+ why="target call: bf16 pool, has_initial_state [1, 0, 1], seqlens [300, 1024, 500]",
+ spec=dict(
+ geometry=QWEN38_GEOMETRY,
+ seqlens=(300, 1024, 500),
+ slots=ENGINE_SLOTS,
+ indices=_engine_slots(3),
+ pool_dtype=torch.bfloat16,
+ output="buffer",
+ initial=(True, False, True),
+ linear_gate=True,
+ ),
+ ),
+ ) + (
+ _rank_cell(
+ "tep4 4/12, fresh 8192-token prompt, CP path, 1874-slot bf16 page",
+ RANK4_GEOMETRY, (8192,), TEP4_SLOTS, (False,), "cp",
+ ),
+ _rank_cell(
+ "tep4 4/12, [200, 37, 91] carried/fresh/carried, CP path, 1874 slots",
+ RANK4_GEOMETRY, (200, 37, 91), TEP4_SLOTS, (True, False, True), "cp",
+ ),
+ _rank_cell(
+ "tp4 4/12, carried 8192-token chunk, CP path, 1874-slot bf16 page",
+ RANK4_GEOMETRY, (8192,), TP4_SLOTS, (True,), "cp",
+ ),
+ _rank_cell(
+ "dep4 16/48, carried 512-token prompt, 850-slot bf16 page",
+ QWEN38_GEOMETRY, (512,), DEP4_SLOTS, (True,), "chunked",
+ ),
+ _rank_cell(
+ "tep4 4/12 envelope, 8192 rows over 128 seqs of 63/65, chunked path",
+ RANK4_GEOMETRY, (63, 65) * 64, TEP4_SLOTS, (True, False) * 64, "chunked",
+ ),
+ _rank_cell(
+ "tep4 4/12, the carried 2-token snapshot chunk, in place, 1874-slot bf16 page",
+ RANK4_GEOMETRY, (2,), TEP4_SLOTS, (True,), "inplace",
+ ),
+ _rank_cell(
+ "tep4 4/12, fresh 5 + carried 2 tokens, in place, 1874 slots",
+ RANK4_GEOMETRY, (5, 2), TEP4_SLOTS, (False, True), "inplace",
+ ),
+ _rank_cell(
+ "tep4 4/12, one carried 64-token chunk (the in-place bound), 1874 slots",
+ RANK4_GEOMETRY, (64,), TEP4_SLOTS, (True,), "inplace",
+ ),
+ _rank_cell(
+ "tep4 4/12, one carried 65-token chunk, past the in-place bound: CP path",
+ RANK4_GEOMETRY, (65,), TEP4_SLOTS, (True,), "cp",
+ ),
+ _rank_cell(
+ "dep4 16/48, carried 2 + fresh 30 + carried 1 tokens, in place, 850 slots",
+ QWEN38_GEOMETRY, (2, 30, 1), DEP4_SLOTS, (True, False, True), "inplace",
+ ),
+ ) # fmt: skip
+
+ note = """
+ Metadata consumed: none.
+
+ 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, and belong to
+ FlashInfer's decode kernel (`gdn_decode.gated_delta_rule_decode_pretranspose`).
+
+ The target (`PrefillTarget._gdn_core`, `_gdn_core_verify`) passes the
+ context rows by keyword after `gdn_prefill_conv_post(linear_gate=True)`,
+ with `has_initial_state` and `linear_gate=True`.
+
+ Call site with `has_initial_state` (`hi_p` / `hi_all`, the flags
+ `reset_recurrent_state_rows` would zero the fresh slots by): a fresh
+ sequence's slot is never read, and the scatter writes every named slot,
+ so the reset of the recurrent pool before this call is not needed (the
+ conv pass reads no carry for such a sequence and its publish writes the
+ whole conv window). Bitwise the reset followed by the plain call; the
+ gather moves only the carried sequences' rows.
+
+ 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; and at the 461-slot bf16 pool: three
+ ragged context rows, and the target's call (300 + 1024 + 500 tokens,
+ the middle one fresh, `has_initial_state`, `linear_gate`, `output=`).
+
+ Certified per rank, as `PrefillTarget._gdn_core` calls it (bf16 pool,
+ `has_initial_state`, `linear_gate`, `output=` a prefix view of a wider
+ buffer whose tail stays untouched). The pool is the layer's view of the
+ stock manager's state page (`PAGE_LAYERS` states a slot), at the
+ member's slot count:
+ tp4/tep4 (`RANK4_GEOMETRY`, 4 k / 12 v heads), 1874 slots: one fresh
+ 8192-token prompt, a carried 8192-token chunk, [200, 37, 91]
+ carried/fresh/carried, and 8192 rows over 128 sequences of 63/65
+ tokens (the envelope: max_num_tokens over max_batch_size).
+ dep4 (16 / 48), 850 slots: a carried 512-token prompt.
+ The last slot is named at every one: offsets of 1.3e10 (tep4) and
+ 2.4e10 (dep4) elements, past int32; the gather and the scatter index the
+ slot in int64.
+
+ Short calls run in place (`fi_path` `inplace`): a `has_initial_state` call of at most
+ `INPLACE_MAX_TOKENS` (64) context tokens zeroes its fresh sequences' pool rows (one
+ Triton launch, `_zero_fresh_states_kernel`; a carried row is not touched) and runs
+ FlashInfer's chunked kernel with `use_cp=False`, `state_indices` the slots and the
+ pool as both `initial_state` and `output_state`: the kernel reads and writes the named
+ rows of the (non-compact) pool itself, so neither the packed gather nor the scatter
+ runs. That is the 2-token snapshot step of chunked prefill
+ (`additional_snapshot_offsets_from_end: [2]`) on every GDN layer. Numerics: the
+ chunked kernel's sum, the state read in the pool's
+ dtype, kept fp32 inside and rounded once to the pool's dtype on the write, as the
+ packed path rounds it on the scatter.
+
+ FlashInfer picks its kernel by `num_seqs x num_v_heads`: on sm_10x the
+ context-parallel pipeline (`cp_delta_rule_dsl_sm100`: t precompute, mn
+ precompute, fixup, prefill) when that times 4 is below the SM count
+ (`should_use_cp_host`), else the chunked kernel. GB300 has 152 SMs, so a
+ tp4/tep4 rank (12 v heads) runs CP for 1 to 3 context sequences and the
+ chunked kernel from 4; a dep4 rank (48) and tp1 always run the chunked
+ one. Both paths take the bf16 packed state. The cells hold the path
+ they reach (`fi_path`), and a CP fallback warning fails them.
+
+ Per rank cell against the fp64 reference (bf16 pool), the CP pipeline and
+ the chunked kernel alike: o rel-L2 <= 3.77e-3, max abs 7.3e-4; named
+ slots' states rel-L2 <= 3.47e-3, max abs 3.9e-3. A value head reading the
+ wrong key head at 4/12 gives o rel-L2 1.41, which the band refuses.
+
+ Memory per call: a packed copy of the named slots in and out in the pool's
+ dtype (2 x 3 MiB per sequence fp32, half that bf16), `exp(g)` (not with
+ `linear_gate`), and the output. q/k/v/beta are not copied when already
+ contiguous / fp32.
+
+ JIT and streams: the first call per process and path compiles the
+ CuTe-DSL kernels in-process for the device it runs on (the cells ran
+ both paths on sm_103), and the installed
+ FlashInfer keeps one TMA workspace per device (the CP pipeline also a
+ cached initial-state workspace), so two streams must not run this op
+ concurrently. Needs CUDA >= 13 torch (FlashInfer raises otherwise).
+
+ Numerics (vs an fp64 recurrence, 1 x 8192 tokens, qwen3.8 geometry): o
+ rel-L2 4.06e-3, max abs 1.26e-3; state rel-L2 3.31e-3 -- the Triton op's
+ are 4.00e-3 / 1.13e-3 / 3.26e-3. Not bitwise to it. A bf16 pool rounds each
+ sequence's state to bf16 once per call, on the scatter. No cold-call
+ defect: nothing in the call is autotuned.
+
+ 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=`, an fp16 pool, `scale` other than the
+ default, fp16 operands, K or V other than 128, head layouts other than
+ 16/48 and 4/12, and what the target never passes:
+ `inplace_indexed_state_update=False` (final states returned) and
+ `use_qk_l2norm_in_kernel`. At 4/12: 2 context sequences on CP and 4 to
+ 127 on the chunked kernel (the cells hold 1 and 3, and 128); a dep4 rank
+ at more than 512 carried tokens; tp4's own slot count (tep4's stands in)
+ and any pool counted on the default target's own engines or under MTP,
+ kv_fp8/kv_nvfp4, chunked or reuse;
+ a call inside a CUDA graph (the target calls it eagerly).
+
+ Hard limits
+ T: max_num_tokens (8192) rows a call is certified at; chunked
+ prefill of a longer prompt is one call per chunk.
+ slot offsets: past int32 at both members' last slots (certified).
+ """
+
+ def __call__(
+ 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,
+ has_initial_state: Optional[torch.Tensor] = None,
+ linear_gate: bool = False,
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ assert not linear_gate or has_initial_state is not None, (
+ "linear_gate is taken with has_initial_state only"
+ )
+ if has_initial_state is not None:
+ assert (
+ inplace_indexed_state_update
+ and initial_state_indices is not None
+ and not output_final_state
+ and not head_first
+ and not use_qk_l2norm_in_kernel
+ and state_workspace is None
+ and cu_seqlens is not None
+ and g.dtype == torch.float32
+ ), "has_initial_state: the in-place indexed call only, q/k normalized by the caller"
+ o = (_chunk_inplace if q.shape[1] <= INPLACE_MAX_TOKENS else _chunk_carried)(
+ q, k, v, g, beta, scale, initial_state, initial_state_indices,
+ has_initial_state, cu_seqlens, output, linear_gate,
+ ) # fmt: skip
+ return o, None
+ 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,
+ has_initial_state: Optional[torch.Tensor] = None,
+ linear_gate: bool = False,
+ ) -> 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()
+ carried = [True] * len(slots) if has_initial_state is None else has_initial_state.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)
+ if not carried[n]:
+ state = torch.zeros_like(state)
+ 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)
+ decay = g[0, t].to(dt) if linear_gate else g[0, t].to(dt).exp()
+ state = state * decay[:, None, None]
+ u = (v[0, t].to(dt) - (state * k_t[:, None, :]).sum(-1)) * beta[0, t].to(dt)[
+ :, None
+ ]
+ state = state + u[:, :, None] * k_t[:, None, :]
+ out[0, t] = (state * q_t[:, None, :]).sum(-1)
+ finals.append(state)
+ pool = initial_state.clone()
+ final_state = None
+ if inplace_indexed_state_update:
+ for slot, state in zip(slots, finals):
+ pool[slot] = state.to(pool.dtype)
+ elif output_final_state:
+ final_state = torch.stack(finals).to(initial_state.dtype)
+ return out.to(q.dtype), final_state, pool
+
+
+flashinfer_chunk_gated_delta_rule = FlashinferChunkGatedDeltaRule()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py
new file mode 100644
index 000000000000..ce6852817a35
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py
@@ -0,0 +1,616 @@
+# 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.
+
+Called by the qwen4_exp target at every layout: the cells of tp1 (all 16 key / 48 value
+heads), plus a tp4/tep4 rank's 4 key / 12 value heads over the tep4 engine's bf16 state page.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch.modules.fla.cached_replay import (
+ fused_recurrent_gated_delta_rule_cached_replay_update as _cached_replay_update,
+)
+
+#: This checkpoint's GDN geometry at tp1.
+NUM_K_HEADS, NUM_V_HEADS, HEAD_K_DIM, HEAD_V_DIM = 16, 48, 128, 128
+#: `[q | k | v]` row width of the post-conv activations a verify step packs.
+PACKED_QKV_WIDTH = 2 * NUM_K_HEADS * HEAD_K_DIM + NUM_V_HEADS * HEAD_V_DIM # 10240
+#: The cache manager's `replay_history_size`: max(MIN_REPLAY_HISTORY_SIZE=16, T).
+HISTORY_SIZE = 16
+#: From this batch size on the engine runs the partitioned replay (work items, no
+#: in-kernel commit, one all-layer commit after every layer has run):
+#: `CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE` in `modules/fla/cached_replay.py`.
+ALL_LAYER_COMMIT_MIN_BATCH = 16
+#: The recurrent pool's slot count of a full-size engine run.
+ENGINE_SLOTS = 461
+#: Layers per slot at which the last of `ENGINE_SLOTS` slots starts past 2^31 elements of
+#: the pool (`460 * L * 48 * 128 * 128 > 2**31` from L = 6).
+PAST_INT32_LAYERS = 8
+#: A tp4/tep4 rank's GDN heads: 16 key / 48 value split four ways.
+RANK_K_HEADS, RANK_V_HEADS = 4, 12
+#: Recurrent-state slots of a tep4 engine rank (bf16 state, host-resident PLE table).
+TEP4_SLOTS = 1874
+#: GDN layers whose states share one slot page in the stock manager (48 layers, every 4th
+#: full attention): the slot stride is 36 layers' `[HV, V, K]` states.
+PAGE_LAYERS = 36
+
+
+def _cell_spec(path, n, t, **spec) -> dict:
+ 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)
+ )
+
+
+def _tep4_cell(why, path, n, t, **spec) -> Cell:
+ """A tp4/tep4 rank's heads over the stock tep4 engine's bf16 state page."""
+ return Cell(
+ why=f"tep4 {why}",
+ spec=_cell_spec(
+ path,
+ n,
+ t,
+ state_dtype=torch.bfloat16,
+ H=RANK_K_HEADS,
+ HV=RANK_V_HEADS,
+ history_size=HISTORY_SIZE,
+ slots=TEP4_SLOTS,
+ page_layers=PAGE_LAYERS,
+ **spec,
+ ),
+ )
+
+
+class FusedRecurrentGatedDeltaRuleCachedReplayUpdate(OpWrapper):
+ """One verify step of `N` sequences x `T` tokens through one GDN layer, for
+ the cached-replay state layout. Per live row (`slot = state_indices[n] >= 0`),
+ with `p = prev_num_accepted_tokens[slot]`, `b = cache_buf_idx[slot]`:
+
+ 1. the state the step starts from is rebuilt from the checkpoint
+ `h0 = ssm_states[slot]` (`[HV, V, K]`, K innermost) and the `p` pending
+ history entries in half `b` -- update vectors `u_j = old_u[slot, b, j]`,
+ normalized keys `k_j = old_k[slot, b, j]`, cumulative log-decays
+ `G_j = old_G[slot, b, :, j]` measured from the checkpoint:
+ `h = exp(G_{p-1}) h0 + sum_j exp(G_{p-1} - G_j) u_j k_j^T`;
+ 2. the `T` new tokens run the gated delta rule in order (q/k head
+ `hv // (HV // H)`, so a v-head group shares one q/k head):
+ `h *= exp(g_t); u_t = beta_t (v_t - h k_t); h += u_t k_t^T; o_t = h (scale q_t)`,
+ q and k L2-normalized first (`x / sqrt(sum(x^2) + 1e-6)`) when
+ `use_qk_l2norm_in_kernel`, `scale=None` meaning `K ** -0.5`; with
+ `A_log` given, `g = -exp(A_log) * softplus(g + dt_bias)` and
+ `beta = sigmoid(beta)` are applied to the raw projections inside;
+ 3. the new tokens' `(u_t, k_t, G_t)` are appended to the history. When
+ `p + T <= history_size` they go to half `b` at `p .. p+T-1` with `G`
+ still cumulative from the checkpoint. Otherwise (a *wrap*) they go to
+ half `1 - b` at `0 .. T-1` with `G` restarted from zero, and the
+ checkpoint must advance to the state after the `p` pending tokens:
+ without work items the kernel writes that state into `ssm_states[slot]`
+ itself; with `replay_work_items` (`use_all_layer_commit=True`) it does
+ not, and the caller owes one `commit_gdn_cached_replay_history_layers`
+ over every layer before the next step.
+
+ Nothing about acceptance happens here: the caller advances `p`/`b` after
+ sampling (`p <- accepted` and `b <- 1 - b` on a wrap, else `p += accepted`,
+ `accepted` counting the bonus token), so rejected tokens' history entries
+ are simply never read.
+
+ Fusion boundary: L2 norms, scale, the gating of the raw `a`/`b`
+ projections, the history rebuild, the recurrence, the history append and
+ (small batches) the checkpoint commit are inside. The conv, the projections,
+ the all-layer commit and the `p`/`b` bookkeeping are not.
+
+ Returns `[N, T, HV, V]` in `output`'s dtype, else `q`'s; it *is* `output`
+ when one is given. `ssm_states`, `old_u`, `old_k`, `old_G` are written in
+ place and not returned. Rows with a negative slot are skipped entirely:
+ nothing is read or written for them, their output rows included.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: path ("small": no work items, the kernel commits a wrap itself |
+ # "all_layer": work items + use_all_layer_commit, the test then runs the
+ # manager's all-layer commit), n sequences x t tokens per verify step,
+ # steps (consecutive verify steps, each followed by an acceptance per row),
+ # accept ("random": 1..t per row, the default | "full": t every step, so from
+ # a fresh slot at t=8 the second step appends at `pnat + t == 16` without
+ # wrapping and every second step after wraps),
+ # state_dtype (the checkpoint pool's dtype: fp32 | bf16, as
+ # `mamba_ssm_cache_dtype` allocates it; fp32 by default),
+ # layers (each with its own activations, gates
+ # and history; one V2-style pool interleaving the layers per slot), slots
+ # (the pool's slot count), start ("warm": random pending histories in both
+ # halves, a wrap and an append forced at step 0 | "fresh": pnat 0 and
+ # zeroed history, as the manager leaves a slot after prefill), edge_slots
+ # (rows 0 and n-1 on the pool's last and first slot), dup (rows given the
+ # SAME slot, as CUDA-graph padding repeats one dummy slot; only the other
+ # rows and slots are held to the reference), H / HV (key / value heads; tp1's
+ # 16 / 48 where absent), history_size (16, the only one certified), page_layers
+ # (the stock manager's coalesced page: the slot stride is that many layers'
+ # states, `slots` slots of it a reserved virtual range with memory mapped under
+ # the touched slots only, the cell's `layers` the top sub-pages; rows on
+ # both ends of the pool, then spread over it). Every cell passes `packed_qkv`,
+ # separate raw `a`/`b` projections, no `output` and `launch_with_pdl=False`.
+ #
+ # tp1's four cells: the forward reaches two commit paths (in-kernel below 16
+ # requests, all-layer from 16) at two checkpoint dtypes, and two edges no
+ # main-shape cell can carry -- CUDA-graph padding rows sharing one slot, and slot
+ # offsets past 2^31 elements at the engine's pool. Then a tp4/tep4 rank's six:
+ # the verify widths T = 2, 4, 8 (max_draft_len 1, 3, 7), both commit
+ # paths, and N = 128, where a bf16 state turns on the large-workload mapping of
+ # the replay and the commit; all at the tep4 engine's page, whose top slot sits
+ # 1.3e10 elements past the base.
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="all_layer n=32 T=8 full accept: commit every second step",
+ spec=_cell_spec("all_layer", 32, 8, start="fresh", accept="full", steps=8),
+ ),
+ _bf16_cell(
+ "small n=4 T=8 full accept: pnat+T==16 appends, then wraps",
+ "small",
+ 4,
+ 8,
+ start="fresh",
+ accept="full",
+ steps=8,
+ ),
+ Cell(
+ why="small n=8 T=3 dup: rows 5..7 share one padding slot",
+ spec=_cell_spec("small", 8, 3, dup=(5, 6, 7), steps=6),
+ ),
+ Cell(
+ why="all_layer n=16 T=3 L=8 S=461: slot 460's offset is past 2^31 elements",
+ spec=_cell_spec(
+ "all_layer",
+ 16,
+ 3,
+ layers=PAST_INT32_LAYERS,
+ slots=ENGINE_SLOTS,
+ edge_slots=True,
+ steps=3,
+ ),
+ ),
+ _tep4_cell(
+ "small n=1 T=4: the in-kernel commit of the page's last slot (1873)",
+ "small",
+ 1,
+ 4,
+ steps=8,
+ ),
+ _tep4_cell(
+ "all_layer n=16 T=4: the smallest partitioned batch", "all_layer", 16, 4, steps=6
+ ),
+ _tep4_cell(
+ "all_layer n=32 T=8: the widest verify below N=128", "all_layer", 32, 8, steps=6
+ ),
+ _tep4_cell(
+ "all_layer n=128 T=4: bf16 at N=128, the large-workload mapping",
+ "all_layer",
+ 128,
+ 4,
+ steps=4,
+ ),
+ _tep4_cell(
+ "all_layer n=128 T=2: the BT=2 token tile (max_draft_len 1)",
+ "all_layer",
+ 128,
+ 2,
+ steps=6,
+ ),
+ _tep4_cell(
+ "all_layer n=128 T=8: the envelope's largest verify, 128 x 8 rows",
+ "all_layer",
+ 128,
+ 8,
+ steps=4,
+ ),
+ )
+
+ note = """
+ 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.
+
+ The default target at tp4/tep4 (`DecodeTarget._gdn_verify -> _replay_verify`): a rank's
+ 4 key / 12 value heads, q/k/v the `[N, T, 4|12, 128]` views of the post-conv
+ `[N*T, 2560]` rows, `g`/`beta` `[N, T, 12]`, and `ssm_states` the layer's view of the
+ stock manager's bf16 state page, `[slots, 12, 128, 128]` at a slot stride of 36
+ layers' states (7,077,888 elements): 1874 slots at the stock tep4 engine, whose top
+ slot sits 1.3e10 elements past the page's base. The history buffers are
+ `[slots, 2, 16, 12|4, 128]` (and `[slots, 2, 12, 16]` fp32); the manager's all-layer
+ commit stays on, so the call is partitioned from N = 16 exactly as above. At dep4 a
+ rank holds all 16 / 48 heads, tp1's geometry.
+
+ Metadata consumed: `state_indices` (one slot per row), and per slot
+ `prev_num_accepted_tokens` and `cache_buf_idx`. `replay_work_items` /
+ `n_writes` are only checked (shape, dtype); the kernel never reads them, their
+ 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 `modules/fla/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 kernels are Triton, JIT-compiled on first call for the running device; the cells
+ ran that JIT path on sm_103. The launch is one V tile, `num_warps=4`, grid
+ `(1, N * HV)`: the tuned mappings need HV = 2H or 4H and never engage at HV/H = 3
+ (16 / 48 and 4 / 12 alike). The checkpoint's dtype is the kernel's `STATE_FP32`
+ specialization and nothing else below N = 128: an fp32 checkpoint enters the history
+ rebuild's tensor-core dots as a bf16 hi/lo pair, a bf16 one as stored, and every
+ commit is computed in fp32 and rounded once to the pool's dtype. From N = 128 a bf16
+ state alone also turns on the large-workload mapping (`.cg` loads of the state and
+ the inputs, `main_num_stages=3`, and a 5-stage all-layer commit); the tep4 N = 128
+ cells drive it. `FLA_USE_FAST_OPS=1` at import switches the kernel's `exp` to the
+ approximate one; the cells run with it unset.
+
+ CUDA-graph replay on new data is bit for bit the eager call (the test captures
+ N = 32 at T = 8 over an fp32 checkpoint, with the all-layer commit). The two commit
+ 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 at tp1's 16 / 48 heads: the all-layer commit over an fp32 checkpoint at
+ N = 32, T = 8 and at N = 16, T = 3 with 8 interleaved layers over the tp1 engine's
+ 461-slot pool (the last slot's offset past 2^31 elements); the in-kernel commit over a
+ bf16 checkpoint at N = 4, T = 8 and over an fp32 one at N = 8, T = 3 with three
+ CUDA-graph padding rows on one shared slot. T = 5..8 share one compiled token tile
+ (`BT = next_pow2(T) = 8`); these cells drive the BT = 4 and BT = 8 tiles, not
+ BT = 2 (T = 2). At T = 8 and history
+ 16 a slot whose every draft is accepted appends rows 8..15 at `pnat + T == 16` without
+ wrapping and wraps on the next step, so it commits every second step; the
+ `accept="full"` cells drive that boundary (in-kernel and all-layer commit).
+
+ Certified at a tp4/tep4 rank's 4 / 12 heads over a bf16 checkpoint on the stock tep4
+ engine's 1874-slot page, each cell several consecutive verify steps from warm
+ histories with random acceptance, rows on both ends of the pool: the in-kernel commit
+ at N = 1, T = 4 on the page's last slot; the all-layer commit at N = 16, T = 4,
+ N = 32, T = 8, and N = 128 at T = 2, 4 and 8 (the large-workload mapping, the BT = 2,
+ 4 and 8 token tiles, and the envelope's largest verify, 128 x 8 rows). The gates:
+ output elements within 4 bf16 ulps of their token block and relative RMS within 1.5
+ ulps; the bf16 checkpoint within 2 ulps of its head's absmax; `old_u` within 4 ulps;
+ every cell lands at most 0.84 of its gate. A reference with the key heads reversed
+ (value heads 0..2 on key head 3) is refused by the output gate.
+
+ Loud refusals: work items without `use_all_layer_commit` (ValueError) and the
+ reverse (ValueError); one of `replay_work_items` / `n_writes` alone
+ (ValueError); a strided `output`; a strided head axis of `g`; `history_size <
+ 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; at 16 / 48 heads a bf16 one at
+ N >= 128, N above 32, T = 2, and any slot count but the cells' (dep4's own pool is not
+ driven); HV/H other than 3, K or V other than 128; T = 3, 5, 6 and 7 at 4 / 12 (they
+ share the BT = 4 and 8 tiles), T outside 2..8 (T = 9 raises: the kernel's bare
+ `assert T <= 8`); N above 128; `history_size` other than 16; `old_u` / `old_k`
+ other than bf16, `old_G` other than fp32; a `scale`; any launch-knob override
+ (`block_v`, `num_warps`, `use_l2_*`, `main_num_stages`); an output dtype other than
+ bf16; int64 `state_indices`; `g`/`beta` as strided column views of the input
+ projection (the cells pass them contiguous). At 4 / 12: the commit over the engine's
+ 36 layers (the cells commit 2 sub-pages of the 36-layer page; a 36-layer history puts
+ `old_u`'s layer offsets past 2^31 elements, addressed in int64 by the commit's source
+ but not driven); a tp4 engine's slot count and the MTP engine's at tep4 (neither
+ measured); CUDA-graph capture and padding rows sharing a slot (both certified at
+ 16 / 48 only).
+ """
+
+ def __call__(
+ 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/ops/ssm/gdn_decode_gated_norm_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_decode_gated_norm_pdl.py
new file mode 100644
index 000000000000..2112962a1b7a
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_decode_gated_norm_pdl.py
@@ -0,0 +1,559 @@
+# 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.
+
+A CuTe DSL kernel; the recurrence is the gated delta rule of FlashInfer's
+`gdn_decode.gated_delta_rule_decode_pretranspose`, reimplemented here.
+The state is not produced by anything in the step ahead of this kernel -- the
+pool row is this layer's own and was last written by this kernel one step ago, and
+the slot index is filled before the step's graph is launched -- so it is read
+before `griddepcontrol.wait`, while the input projection ahead of this kernel
+streams its weight.
+
+A v-head is split over a cluster of `CLUSTER` CTAs of `WARPS` warps; each warp owns
+`V / (CLUSTER * WARPS)` value rows of the head's `[V, K]` state and lane `l` the
+`K / 32` key columns from `l * K / 32` of each. Small CTAs fit beside the input
+projection's last wave, so their state reads start before it ends. The gated RMS
+norm (per v-head, 128 wide) needs the head's whole output: every warp writes the
+sum of squares of its rows into each CTA of the cluster through distributed shared
+memory, and every CTA then normalizes and gates its own rows.
+"""
+
+from __future__ import annotations
+
+import cuda.bindings.driver as cuda_drv
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass import const_expr
+from cutlass._mlir.dialects import llvm
+from cutlass.cute import experimental as cute_ext
+from cutlass.cute.runtime import from_dlpack
+from cutlass.cutlass_dsl import dsl_user_op
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.norm.rms_norm_gated_token_major import (
+ rms_norm_gated_token_major as _gated_norm,
+)
+from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+
+#: Decode rows per launch this entry is certified for.
+MAX_ROWS = 4
+#: CTAs per v-head (one cluster) and warps per CTA; their product must be <= 32.
+CLUSTER = 8
+WARPS = 4
+
+
+@dsl_user_op
+def _st_shared_cluster_f32(ptr, val, *, loc=None, ip=None):
+ """`st.shared::cluster.f32` to a pointer from `cute.arch.map_dsmem_ptr`."""
+ llvm.inline_asm(
+ None,
+ [
+ cutlass.Int32(ptr.toint(loc=loc, ip=ip)).ir_value(loc=loc, ip=ip),
+ cutlass.Float32(val).ir_value(loc=loc, ip=ip),
+ ],
+ "st.shared::cluster.f32 [$0], $1;",
+ "r,f",
+ has_side_effects=True,
+ is_align_stack=False,
+ asm_dialect=llvm.AsmDialect.AD_ATT,
+ loc=loc,
+ ip=ip,
+ )
+
+
+class _GdnKernel:
+ def __init__(self, *, k_heads, v_heads, kdim, vdim, pdl):
+ assert vdim % (CLUSTER * WARPS) == 0 and CLUSTER * WARPS <= 32
+ assert kdim % 32 == 0 and v_heads % k_heads == 0
+ self.k_heads, self.v_heads, self.kdim, self.vdim, self.pdl = (
+ k_heads,
+ v_heads,
+ kdim,
+ vdim,
+ pdl,
+ )
+ self.rows_per_cta = vdim // CLUSTER
+ self.rows_per_warp = vdim // (CLUSTER * WARPS)
+ self.per_lane = kdim // 32
+ self.k_off = k_heads * kdim
+ self.v_off = 2 * k_heads * kdim
+ self.z_off = self.v_off + v_heads * vdim
+ self.a_off = self.z_off + v_heads * vdim
+ self.b_off = self.a_off + v_heads
+
+ @cute.jit
+ def __call__(
+ self,
+ mProj,
+ mPool,
+ mSlot,
+ mALog,
+ mDt,
+ mW,
+ mOut,
+ eps: cutlass.Float32,
+ stream: cuda_drv.CUstream,
+ ):
+ rows = cute.size(mProj, mode=[0])
+ ld_state = cute.make_copy_atom(
+ cute.nvgpu.CopyG2ROp(),
+ cutlass.Float32,
+ num_bits_per_copy=32 * self.per_lane,
+ load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING,
+ )
+ st_state = cute.make_copy_atom(
+ cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=32 * self.per_lane
+ )
+ ld_act = cute.make_copy_atom(
+ cute.nvgpu.CopyG2ROp(),
+ cutlass.BFloat16,
+ num_bits_per_copy=16 * self.per_lane,
+ load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS,
+ )
+ self.kernel(
+ mProj, mPool, mSlot, mALog, mDt, mW, mOut, eps, ld_state, st_state, ld_act
+ ).launch(
+ grid=[self.v_heads * CLUSTER, rows, 1],
+ block=[WARPS * 32, 1, 1],
+ cluster=[CLUSTER, 1, 1],
+ smem=32 * 4 + 16,
+ stream=stream,
+ use_pdl=self.pdl,
+ # <= 64 registers a thread: the CTAs fit beside the producer's last wave.
+ min_blocks_per_mp=max(1, 8 * 4 // WARPS),
+ )
+
+ @cute.kernel
+ def kernel(
+ self,
+ mProj,
+ mPool,
+ mSlot,
+ mALog,
+ mDt,
+ mW,
+ mOut,
+ eps: cutlass.Float32,
+ ld_state: cute.CopyAtom,
+ st_state: cute.CopyAtom,
+ ld_act: cute.CopyAtom,
+ ):
+ bidx, n, _ = cute.arch.block_idx()
+ warp = cute.arch.warp_idx()
+ lane = cute.arch.lane_idx()
+ rpw: cutlass.Constexpr = self.rows_per_warp
+ pl: cutlass.Constexpr = self.per_lane
+ kdim: cutlass.Constexpr = self.kdim
+ vdim: cutlass.Constexpr = self.vdim
+ parts: cutlass.Constexpr = CLUSTER * WARPS
+ hv = bidx // CLUSTER
+ part = bidx % CLUSTER # the CTA's rank in its cluster: clusters tile the grid's x
+ kh = hv // (self.v_heads // self.k_heads)
+ v0 = part * self.rows_per_cta + warp * rpw
+ k0 = lane * pl
+ ss_parts = cutlass.utils.SmemAllocator().allocate_tensor(
+ cutlass.Float32, cute.make_layout((32,), stride=(1,)), byte_alignment=16
+ )
+ if const_expr(CLUSTER > 1):
+ # Every CTA of the cluster must have started before another writes into its
+ # shared memory; this arrival is waited on just before those writes.
+ cute.arch.cluster_arrive_relaxed()
+
+ # 1. Static and step-invariant operands, before the wait: the gates' constants,
+ # the norm weight of the warp's rows, the slot and the warp's state rows.
+ a_log_h = mALog[hv]
+ dtb = mDt[hv]
+ wn = cute.make_rmem_tensor(cute.make_layout((rpw,), stride=(1,)), cutlass.Float32)
+ for r in cutlass.range_constexpr(rpw):
+ wn[r] = mW[v0 + r].to(cutlass.Float32)
+ slot = mSlot[n]
+ # 64-bit: the cache manager coalesces every GDN layer's state of a slot into one
+ # page, so a layer's slot stride is ~28M elements and slot * stride overflows int32.
+ slot64 = cutlass.Int64(slot)
+ # (slot, hv, v, (pl, lane)): lane l's 16-byte vector of a state row.
+ pool_vec = cute.logical_divide(mPool, (None, None, None, pl))
+ h = cute.make_rmem_tensor(cute.make_layout((rpw, pl), stride=(pl, 1)), cutlass.Float32)
+ h.fill(0.0)
+ if slot >= 0:
+ for r in cutlass.range_constexpr(rpw):
+ cute.copy(ld_state, pool_vec[slot64, hv, v0 + r, (None, lane)], h[r, None])
+ if const_expr(self.pdl):
+ cute.arch.griddepcontrol_launch_dependents()
+ cute.arch.griddepcontrol_wait()
+
+ # 2. The projection's q, k, v, z, a, b for this row and head.
+ # (n, (pl, cols / pl)): 8-byte vectors of the row; every offset below is a multiple of pl.
+ proj_vec = cute.logical_divide(mProj, (None, pl))
+ raw = cute.make_rmem_tensor(cute.make_layout((2, pl), stride=(pl, 1)), cutlass.BFloat16)
+ cute.copy(ld_act, proj_vec[n, (None, (kh * kdim + k0) // pl)], raw[0, None])
+ cute.copy(ld_act, proj_vec[n, (None, (self.k_off + kh * kdim + k0) // pl)], raw[1, None])
+ qn = cute.make_rmem_tensor(cute.make_layout((pl,), stride=(1,)), cutlass.Float32)
+ kn = cute.make_rmem_tensor(cute.make_layout((pl,), stride=(1,)), cutlass.Float32)
+ zv = cute.make_rmem_tensor(cute.make_layout((rpw,), stride=(1,)), cutlass.Float32)
+ vv = cute.make_rmem_tensor(cute.make_layout((rpw,), stride=(1,)), cutlass.Float32)
+ for e in cutlass.range_constexpr(pl):
+ qn[e] = raw[0, e].to(cutlass.Float32)
+ kn[e] = raw[1, e].to(cutlass.Float32)
+ for r in cutlass.range_constexpr(rpw):
+ vv[r] = mProj[n, self.v_off + hv * vdim + v0 + r].to(cutlass.Float32)
+ zv[r] = mProj[n, self.z_off + hv * vdim + v0 + r].to(cutlass.Float32)
+ a = mProj[n, self.a_off + hv].to(cutlass.Float32)
+ b = mProj[n, self.b_off + hv].to(cutlass.Float32)
+
+ # q, k L2 norm (eps 1e-6 inside the root), the 1/sqrt(K) scale on q, the gates.
+ sq = cutlass.Float32(0.0)
+ sk = cutlass.Float32(0.0)
+ for e in cutlass.range_constexpr(pl):
+ sq = qn[e] * qn[e] + sq
+ sk = kn[e] * kn[e] + sk
+ sq = cute.arch.warp_reduction_sum(sq)
+ sk = cute.arch.warp_reduction_sum(sk)
+ rq = cute.math.rsqrt(sq + 1e-6) * (1.0 / cute.math.sqrt(cutlass.Float32(kdim)))
+ rk = cute.math.rsqrt(sk + 1e-6)
+ for e in cutlass.range_constexpr(pl):
+ qn[e] = qn[e] * rq
+ kn[e] = kn[e] * rk
+ x = a + dtb
+ sp = x
+ if x <= 20.0:
+ sp = cute.math.log1p(cute.math.exp(x))
+ decay = cute.math.exp(-cute.math.exp(a_log_h) * sp)
+ beta = 1.0 / (1.0 + cute.math.exp(-b))
+
+ # 3. The delta rule, one value row at a time; the new rows go straight back.
+ ov = cute.make_rmem_tensor(cute.make_layout((rpw,), stride=(1,)), cutlass.Float32)
+ ssw = cutlass.Float32(0.0)
+ for r in cutlass.range_constexpr(rpw):
+ hk = cutlass.Float32(0.0)
+ for e in cutlass.range_constexpr(pl):
+ hk = h[r, e] * kn[e] + hk
+ hk = cute.arch.warp_reduction_sum(hk)
+ u = beta * (vv[r] - decay * hk)
+ o = cutlass.Float32(0.0)
+ for e in cutlass.range_constexpr(pl):
+ h[r, e] = decay * h[r, e] + u * kn[e]
+ o = h[r, e] * qn[e] + o
+ o = cute.arch.warp_reduction_sum(o)
+ if slot >= 0:
+ cute.copy(st_state, h[r, None], pool_vec[slot64, hv, v0 + r, (None, lane)])
+ else:
+ o = cutlass.Float32(0.0)
+ # The recurrence's output is a bf16 tensor in the unfused composition.
+ ov[r] = o.to(cutlass.BFloat16).to(cutlass.Float32)
+ ssw = ov[r] * ov[r] + ssw
+
+ # 4. The head's sum of squares: warp `w` of CTA `c` writes its rows' share into
+ # slot `c * WARPS + w` of every CTA of the cluster (lane `i` into CTA `i`).
+ if const_expr(CLUSTER > 1):
+ cute.arch.cluster_wait()
+ if lane < CLUSTER:
+ _st_shared_cluster_f32(
+ cute.arch.map_dsmem_ptr(ss_parts.iterator + (part * WARPS + warp), lane), ssw
+ )
+ cute.arch.cluster_arrive()
+ cute.arch.cluster_wait()
+ else:
+ if lane == 0:
+ ss_parts[warp] = ssw
+ cute.arch.sync_threads()
+ ss = cutlass.Float32(0.0)
+ if lane < parts:
+ ss = ss_parts[lane]
+ ss = cute.arch.warp_reduction_sum(ss)
+ rstd = cute.math.rsqrt(ss / vdim + eps)
+
+ # 5. Gated RMS norm of the warp's rows: lane `r` writes row `r`.
+ o_r = ov[0]
+ z_r = zv[0]
+ w_r = wn[0]
+ for r in cutlass.range_constexpr(1, rpw):
+ o_r = cutlass.Float32(cutlass.select_(lane == r, ov[r], o_r))
+ z_r = cutlass.Float32(cutlass.select_(lane == r, zv[r], z_r))
+ w_r = cutlass.Float32(cutlass.select_(lane == r, wn[r], w_r))
+ y = o_r * rstd * w_r * (1.0 / (1.0 + cute.math.exp(-z_r)))
+ if lane < rpw:
+ mOut[n, hv * vdim + v0 + lane] = y.to(cutlass.BFloat16)
+
+
+def _dl(t: torch.Tensor, align: int = 4):
+ return from_dlpack(t, assumed_align=align)
+
+
+_COMPILED: dict = {}
+
+
+def _operands(proj, pool, slots, A_log, dt_bias, norm_weight, out, eps, stream):
+ return (
+ _dl(proj, 8),
+ _dl(pool, 16),
+ _dl(slots),
+ _dl(A_log),
+ _dl(dt_bias),
+ _dl(norm_weight, 2),
+ _dl(out, 2),
+ cutlass.Float32(eps),
+ stream,
+ )
+
+
+def launch(proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, pdl=True):
+ """The launch behind `GdnDecodeGatedNormPdl.__call__`; updates `pool` in place and
+ returns the gated-normed output `[rows, HV * V]` bf16."""
+ rows = proj.shape[0]
+ v_heads, vdim = pool.shape[1], pool.shape[2]
+ out = torch.empty((rows, v_heads * vdim), dtype=torch.bfloat16, device=proj.device)
+ if rows == 0:
+ return out
+ pdl = bool(pdl and get_env_enable_pdl())
+ stream = cuda_drv.CUstream(torch.cuda.current_stream(proj.device).cuda_stream)
+ operands = (proj, pool, slots, A_log, dt_bias, norm_weight, out, eps, stream)
+ key = (
+ rows,
+ k_heads,
+ v_heads,
+ kdim,
+ vdim,
+ proj.stride(0),
+ tuple(pool.shape),
+ pool.stride(),
+ pdl,
+ )
+ fn = _COMPILED.get(key)
+ if fn is None:
+ # Compiled from the call's own operands (a stand-in for the pool would span every
+ # layer of the coalesced state cache).
+ kern = _GdnKernel(k_heads=k_heads, v_heads=v_heads, kdim=kdim, vdim=vdim, pdl=pdl)
+ fn = _COMPILED[key] = cute_ext.compile(kern, *_operands(*operands))
+ fn(*_operands(*operands))
+ return out
+
+
+def _recurrence_reference(q, k, v, A_log, a, dt_bias, b, pool, slots):
+ """`(o, pool_after)` of one gated-delta-rule decode step per row in fp64, rounded once:
+ `o` to bf16, the state to the pool's dtype. Row `n` reads and writes `pool[slots[n]]`
+ (a negative slot skips the row: zero output, no state access); q and k are
+ L2-normalized, the decay is `exp(-exp(A_log) * softplus(a + dt_bias))`, `beta =
+ sigmoid(b)`, the scale `K^-0.5`. `q` / `k` `[n, 1, H, K]`, `v` `[n, 1, HV, V]`, `a` /
+ `b` `[n, 1, HV]`, `pool` `[slots, HV, V, K]`. Mutates nothing."""
+ dt = torch.float64
+ batch, _, num_k_heads, k_dim = q.shape
+ hv = v.shape[2]
+ rep = hv // num_k_heads
+ scale = k_dim**-0.5
+ reads = slots.tolist()
+ pool_after = pool.clone()
+ o = torch.zeros(batch, 1, hv, v.shape[3], dtype=dt, device=q.device)
+ neg_a = -torch.exp(A_log.to(dt))
+ for n in range(batch):
+ if reads[n] < 0:
+ continue
+ q_n = q[n, 0].to(dt).repeat_interleave(rep, dim=0)
+ k_n = k[n, 0].to(dt).repeat_interleave(rep, dim=0)
+ q_n = q_n * torch.rsqrt(q_n.pow(2).sum(-1, keepdim=True) + 1e-6)
+ k_n = k_n * torch.rsqrt(k_n.pow(2).sum(-1, keepdim=True) + 1e-6)
+ x = a[n, 0].to(dt) + dt_bias.to(dt)
+ softplus = torch.where(x <= 20.0, torch.log1p(torch.exp(x)), x)
+ beta = torch.sigmoid(b[n, 0].to(dt))
+ h = pool[reads[n]].to(dt) * torch.exp(neg_a * softplus)[:, None, None]
+ u = (v[n, 0].to(dt) - torch.einsum("hvk,hk->hv", h, k_n)) * beta[:, None]
+ h = h + u[:, :, None] * k_n[:, None, :]
+ o[n, 0] = torch.einsum("hvk,hk->hv", h, q_n * scale)
+ pool_after[reads[n]] = h.to(pool.dtype)
+ return o.to(torch.bfloat16), pool_after
+
+
+class GdnDecodeGatedNormPdl(OpWrapper):
+ """A single-token Gated-DeltaNet step per row over the fp32 slot pool, then the
+ sigmoid-gated RMS norm of its output per v-head. With `proj` the decode input
+ projection's output row `[q (H*K) | k (H*K) | v (HV*V) | z (HV*V) | a (HV) | b (HV)]`
+ (after the conv step: q, k, v are the conv's SiLU outputs)::
+
+ o, pool = gated_delta_rule_decode_pretranspose(q, k, v, A_log, a, dt_bias, b,
+ use_qk_l2norm=True, pool, slots)
+ out = rms_norm_gated_token_major(o (bf16), z, norm_weight, eps, "sigmoid")
+
+ Returns `out` `[rows, HV * V]` bf16 and updates `pool[slots[n]]` in place; a
+ negative slot skips the row's state (no read or write) and its output is zero.
+ `o` is rounded to bf16 between the two, as the composition stores it.
+
+ Fusion boundary: both gates, the q/k L2 norm, the recurrence, the indexed state
+ read and write, the gated RMS norm. The projection and the conv are the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: rows; k_heads / v_heads (default 16 / 48); pool (default a contiguous 461-slot
+ # pool | "coalesced": the cache manager's page, slot stride 36 layers' states -- 102 slots
+ # (slots 100 and 101), or with pool_slots the engine's measured slot count, the view the
+ # last layer's sub-page and the top slots named); pad (row 1 given slot -1; at one row, a
+ # second call over slot -1).
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="qwen3.8 decode, 1 row (con 1), 461-slot pool", spec=dict(rows=1)),
+ Cell(why="qwen3.8 decode, 4 rows (MAX_ROWS), scattered slots", spec=dict(rows=4)),
+ Cell(
+ why="2 rows, coalesced pool (slot stride 36 layers, slot 100), row 1 a pad slot (-1)",
+ spec=dict(rows=2, pool="coalesced", pad=True),
+ ),
+ Cell(
+ why="tp4/tep4 decode, 1 row at 4 k / 12 v heads, slot 1873 of tep4's 1874-slot "
+ "engine page, then a pad slot",
+ spec=dict(rows=1, k_heads=4, v_heads=12, pool="coalesced", pad=True, pool_slots=1874),
+ ),
+ Cell(
+ why="tp4/tep4, 2 rows at 4 / 12 heads, the engine page's top slots, row 1 a pad "
+ "slot (-1)",
+ spec=dict(rows=2, k_heads=4, v_heads=12, pool="coalesced", pad=True, pool_slots=1874),
+ ),
+ Cell(
+ why="tp4/tep4, 4 rows (MAX_ROWS) at 4 / 12 heads, scattered over the 1874-slot page",
+ spec=dict(rows=4, k_heads=4, v_heads=12, pool="coalesced", pool_slots=1874),
+ ),
+ Cell(
+ why="dep4 decode, 1 row at 16 / 48 heads, slot 849 of dep4's 850-slot engine page, "
+ "then a pad slot",
+ spec=dict(rows=1, k_heads=16, v_heads=48, pool="coalesced", pad=True, pool_slots=850),
+ ),
+ )
+
+ note = """
+ 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 the default target's decode steps (each
+ layer's pool is written only by its own recurrence, `slots` is filled before the
+ step's graph launches) and is why the entry is not for verify or prefill steps.
+
+ Launch: `HV * 8` CTAs a row, clusters of 8 (one v-head), 128 threads, <= 64
+ registers. Numerics: fp32 inside; `o` rounded to bf16 before the norm, the output
+ rounded once. Inside torch's default band for bf16 / fp32 of the fp64 composition.
+ Deterministic. Graph-capturable after the first call per geometry (compiled then).
+
+ Source: the CuTe DSL kernel in this file, JIT-compiled (`cute_ext.compile`) for the
+ running device at the first call per geometry; the cells ran that path on sm_103.
+
+ The default target: its routing admits a bf16 recurrent-state cache only, and this kernel takes
+ an fp32 pool only -- a bf16 pool fails its compile loudly ("`copy` currently only
+ supports equal source and destination element type bit width") -- so the default target's
+ decode step does not launch it. Its bf16 decode recurrence is
+ `gdn_decode_rows_gated_norm`, which reads this entry's `reference` and
+ `_recurrence_reference` (both round the state to the pool's own dtype) and not its
+ kernel.
+
+ Certified, sm_103:
+ - tp1's 16 / 48 heads: 1 and 4 rows over a contiguous 461-slot pool, and 2 rows over a
+ 102-slot coalesced page (slot stride 36 layers' states, slot 100 ~2.8G elements past
+ the base) with the second row a pad slot (-1).
+ - A tp4/tep4 rank's 4 / 12 heads (proj width 4120): 1 row, 2 rows (row 1 a pad) and 4
+ rows over tep4's engine page -- 1874 slots, slot stride 36 layers' states (7077888
+ elements), the view the last GDN layer's sub-page; slot 1873 sits ~1.3e10 elements
+ past the base.
+ - dep4's 16 / 48 heads: 1 row at slot 849 of its 850-slot engine page (~2.4e10 elements
+ past the base).
+ - A one-row call over a pad slot (-1) writes zeros and no state; under graph replay the
+ op reads the projection after `griddepcontrol.wait`, at 16 / 48 and at 4 / 12 heads.
+ Against the fp64 reference: at 4 / 12 heads every output element is bitwise its
+ bf16-rounded reference and the state is within 1.2e-7 (relative L2 <= 7.8e-8); at
+ 16 / 48 at most 5 of 24576 output elements differ (max 9.8e-4, relative L2 1.2e-5). On
+ the engine pages every other byte of a touched slot's page stays unchanged.
+
+ The engine slot counts are the stock engines' with a bf16 state cache at
+ `max_batch_size` 128 (the PLE / GDN slot space the manager sizes from memory); an fp32
+ pool in that memory holds about half as many, so the certified slots bound it.
+
+ Not certified: 3 rows; `pdl=False`; tp4's own engine slot count (tep4's 1874 stands in:
+ a tp4 rank's GDN pools are the same shape); a bf16 pool (refused, above).
+ """
+
+ def __call__(
+ self,
+ proj: torch.Tensor,
+ pool: torch.Tensor,
+ slots: torch.Tensor,
+ A_log: torch.Tensor,
+ dt_bias: torch.Tensor,
+ norm_weight: torch.Tensor,
+ eps: float,
+ k_heads: int,
+ kdim: int,
+ pdl: bool = True,
+ ) -> torch.Tensor:
+ return launch(proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, pdl)
+
+ def reference(
+ self,
+ proj: torch.Tensor,
+ pool: torch.Tensor,
+ slots: torch.Tensor,
+ A_log: torch.Tensor,
+ dt_bias: torch.Tensor,
+ norm_weight: torch.Tensor,
+ eps: float,
+ k_heads: int,
+ kdim: int,
+ pdl: bool = True,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """`(out, pool_after)`: the fp64 recurrence (`_recurrence_reference`; o rounded to
+ bf16, the pool to fp32), then the gated norm entry's fp32 reference."""
+ rows = proj.shape[0]
+ v_heads, vdim = pool.shape[1], pool.shape[2]
+ q, k, v, z, a, b = torch.split(
+ proj,
+ [k_heads * kdim, k_heads * kdim, v_heads * vdim, v_heads * vdim, v_heads, v_heads],
+ 1,
+ )
+ o, pool_after = _recurrence_reference(
+ q.reshape(rows, 1, k_heads, kdim),
+ k.reshape(rows, 1, k_heads, kdim),
+ v.reshape(rows, 1, v_heads, vdim),
+ A_log,
+ a.reshape(rows, 1, v_heads),
+ dt_bias,
+ b.reshape(rows, 1, v_heads),
+ pool,
+ slots,
+ )
+ out = _gated_norm.reference(
+ o.reshape(-1, vdim), z.reshape(rows, v_heads, vdim), norm_weight, eps, None, "sigmoid"
+ )
+ return out.reshape(rows, v_heads * vdim), pool_after
+
+ def is_valid(
+ self,
+ proj: torch.Tensor,
+ pool: torch.Tensor,
+ slots: torch.Tensor,
+ A_log: torch.Tensor,
+ dt_bias: torch.Tensor,
+ norm_weight: torch.Tensor,
+ eps: float,
+ k_heads: int,
+ kdim: int,
+ pdl: bool = True,
+ ) -> None:
+ rows = proj.shape[0]
+ v_heads, vdim, kd = pool.shape[1], pool.shape[2], pool.shape[3]
+ assert rows <= MAX_ROWS, f"at most {MAX_ROWS} rows"
+ assert proj.dtype == torch.bfloat16 and proj.stride(1) == 1, "bf16 proj, unit column stride"
+ assert proj.shape[1] == 2 * k_heads * kdim + 2 * v_heads * vdim + 2 * v_heads, "proj width"
+ assert pool.dtype == torch.float32 and pool.stride(3) == 1 and kd == kdim, (
+ "fp32 pool [S, HV, V, K]"
+ )
+ assert vdim % (CLUSTER * WARPS) == 0 and kdim % 32 == 0, "V % (CLUSTER * WARPS), K % 32"
+ assert v_heads % k_heads == 0, "HV a multiple of H"
+ assert slots.dtype == torch.int32 and slots.shape == (rows,), "int32 slots [rows]"
+ assert A_log.dtype == torch.float32 and dt_bias.dtype == torch.float32, (
+ "fp32 gate constants"
+ )
+ assert norm_weight.dtype == torch.bfloat16 and norm_weight.shape == (vdim,), (
+ "bf16 norm weight [V]"
+ )
+
+
+gdn_decode_gated_norm_pdl = GdnDecodeGatedNormPdl()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_decode_rows_gated_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_decode_rows_gated_norm.py
new file mode 100644
index 000000000000..c3c9d94b7133
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_decode_rows_gated_norm.py
@@ -0,0 +1,1155 @@
+# 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 or bf16 state pool: the decode step at batch sizes past the cluster kernel's
+over an fp32 pool, and at every batch size over a bf16 pool.
+
+A CuTe DSL kernel; the recurrence is the gated delta rule of FlashInfer's
+`gdn_decode.gated_delta_rule_decode_pretranspose`, reimplemented here. A CTA owns one
+(row, v-head) state, 64 KB fp32 or 32 KB bf16, read and written once, and streams it: warp w
+the value rows [w * RPW, (w + 1) * RPW), lane l the key columns [4l, 4l + 4) of each, so
+a warp's access is one contiguous state row. A bf16 state is widened to fp32 on the load
+and rounded back (nearest even) on the store. The state is loaded before
+`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.
+
+Two variants of that CTA, the same arithmetic instruction for instruction (bitwise the same
+output and state): the state rows loaded into registers (128 registers a thread, two CTAs
+an SM), or -- from the batch where the register variant needs a second wave -- each warp's
+rows brought into shared memory by one bulk copy (TMA) issued as soon as the slot is known,
+which holds no registers while in flight (64 a thread, three to four CTAs an SM).
+"""
+
+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.models.qwen4_exp.ops.ssm.gdn_decode_gated_norm_pdl import (
+ _recurrence_reference,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.gdn_decode_gated_norm_pdl import (
+ gdn_decode_gated_norm_pdl as _one_launch,
+)
+from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
+
+#: Key columns a lane (one 16-byte fp32 or 8-byte bf16 vector of a state row).
+_PL = 4
+#: Warps a CTA; a warp owns `V // _WARPS` value rows of the head's state.
+_WARPS = 8
+#: Largest recurrent-pool slot count the cache manager allocates for this checkpoint; the
+#: kernel addresses a slot by its index alone, any pool size alike.
+ENGINE_SLOTS = 461
+#: Recurrent-state slots one rank's engine allocates for this checkpoint (bf16 state cache,
+#: `max_batch_size` 128), counted at tep4 and dep4. A tp4 rank's GDN state has tep4's
+#: per-slot shape (4 k / 12 v heads), so tp4 takes tep4's count.
+TEP4_SLOTS = 1874
+DEP4_SLOTS = 850
+TP4_SLOTS = TEP4_SLOTS
+
+#: The rank geometries the engine cells take: a tp4/tep4 rank's 4 k / 12 v heads on its
+#: 4120-wide projection, a dep4 rank's whole 16 / 48 on 16480, each over its member's bf16
+#: engine page.
+_TEP4 = dict(
+ k_heads=4,
+ v_heads=12,
+ proj_width=4120,
+ state="bf16",
+ pool="coalesced",
+ pool_slots=TEP4_SLOTS,
+ member="tep4",
+)
+_DEP4 = dict(
+ k_heads=16,
+ v_heads=48,
+ proj_width=16480,
+ state="bf16",
+ pool="coalesced",
+ pool_slots=DEP4_SLOTS,
+ member="dep4",
+)
+
+
+def _iv(x, loc, ip):
+ 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_state_bf16(addr, pred, *, loc=None, ip=None):
+ """Four bf16 of a state row (streaming) as fp32, or zeros where `pred` is 0."""
+ r = _asm_n(
+ 4,
+ T.f32(),
+ [cutlass.Int64(addr), cutlass.Int32(pred)],
+ "{ .reg .pred p; .reg .b32 w0, w1; setp.ne.b32 p, $5, 0; mov.b32 w0, 0; "
+ "mov.b32 w1, 0; @p ld.global.cs.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,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+ return [cutlass.Float32(v) for v in r]
+
+
+@dsl_user_op
+def _st_state_bf16(addr, pred, a, b, c, d, *, loc=None, ip=None):
+ """Four fp32 rounded to bf16 into 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; .reg .b32 r0, r1; setp.ne.b32 p, $2, 0; "
+ "cvt.rn.bf16x2.f32 r0, $4, $3; cvt.rn.bf16x2.f32 r1, $6, $5; "
+ "@p st.global.cs.v2.b32 [$1], {r0, r1}; 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,
+ )
+ )
+
+
+@dsl_user_op
+def _lds_state(saddr, pred, *, loc=None, ip=None):
+ """Four fp32 of a staged state row (shared memory), or zeros where `pred` is 0."""
+ r = _asm_n(
+ 4,
+ T.f32(),
+ [cutlass.Int32(saddr), 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.f32 {$0, $1, $2, $3}, [$4]; }",
+ "=f,=f,=f,=f,r,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+ return [cutlass.Float32(v) for v in r]
+
+
+@dsl_user_op
+def _lds_state_bf16(saddr, pred, *, loc=None, ip=None):
+ """Four bf16 of a staged state row (shared memory) as fp32, or zeros where `pred` is 0."""
+ r = _asm_n(
+ 4,
+ T.f32(),
+ [cutlass.Int32(saddr), cutlass.Int32(pred)],
+ "{ .reg .pred p; .reg .b32 w0, w1; setp.ne.b32 p, $5, 0; mov.b32 w0, 0; "
+ "mov.b32 w1, 0; @p ld.shared.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,r,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+ return [cutlass.Float32(v) for v in r]
+
+
+@dsl_user_op
+def _fma(a, b, c, *, loc=None, ip=None):
+ """`a * b + c` rounded once (pinned: not left to the backend's contraction)."""
+ 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):
+ """`a * b` rounded (pinned: not contracted into a neighbouring add)."""
+ 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 _sts_f32(saddr, x, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [cutlass.Int32(saddr), cutlass.Float32(x)],
+ "{ st.shared.f32 [$1], $2; mov.u32 $0, 0; }",
+ "=r,r,f",
+ loc,
+ ip,
+ side_effects=True,
+ )
+
+
+@dsl_user_op
+def _lds_f32(saddr, *, loc=None, ip=None):
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [cutlass.Int32(saddr)],
+ "ld.shared.f32 $0, [$1];",
+ "=f,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+ )
+
+
+@dsl_user_op
+def _mbar_init(saddr, *, loc=None, ip=None):
+ """An mbarrier of arrival count 1, made visible to the async proxy."""
+ _asm(
+ T.i32(),
+ [cutlass.Int32(saddr)],
+ "{ mbarrier.init.shared::cta.b64 [$1], 1; fence.mbarrier_init.release.cluster; "
+ "mov.u32 $0, 0; }",
+ "=r,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+
+
+@dsl_user_op
+def _bulk_g2s(saddr, gaddr, nbytes, mbar, *, loc=None, ip=None):
+ """One 1-D bulk copy global -> shared of `nbytes`, completing on the mbarrier.
+
+ With `mbarrier.arrive.expect_tx`; the L2 policy is evict-first (the state is read once a
+ step).
+ """
+ _asm(
+ T.i32(),
+ [cutlass.Int32(saddr), cutlass.Int64(gaddr), cutlass.Int32(nbytes), cutlass.Int32(mbar)],
+ "{ .reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol, 1.0; "
+ "mbarrier.arrive.expect_tx.shared::cta.b64 _, [$4], $3; "
+ "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint "
+ "[$1], [$2], $3, [$4], pol; mov.u32 $0, 0; }",
+ "=r,r,l,r,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+
+
+@dsl_user_op
+def _mbar_wait(mbar, pred, *, loc=None, ip=None):
+ """Wait for phase 0 of the mbarrier where `pred` is not 0."""
+ _asm(
+ T.i32(),
+ [cutlass.Int32(mbar), cutlass.Int32(pred)],
+ "{ .reg .pred P1, P2; setp.ne.b32 P2, $2, 0; @!P2 bra DONE; LAB_WAIT: "
+ "mbarrier.try_wait.parity.shared::cta.b64 P1, [$1], 0; @P1 bra DONE; bra LAB_WAIT; "
+ "DONE: mov.u32 $0, 0; }",
+ "=r,r,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+
+
+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.
+
+ `bulk`: each warp's `RPW` state rows (contiguous: `s_row == K`) come into shared memory
+ by one 1-D bulk copy (TMA, an mbarrier a warp) issued as soon as the slot is known, and
+ both passes over the rows read them from there; the state then holds no registers while
+ it is in flight (64 registers a thread, three to four CTAs an SM, against two). Else the
+ rows are loaded into registers (any 4-element-aligned layout). The arithmetic is the same
+ instruction for instruction: the bulk variant pins the state update to the
+ `fma(decay, h, u * k)` the backend contracts the register variant's to."""
+
+ def __init__(self, k_heads, v_heads, kdim, vdim, pdl, norm=True, state_bf16=False, bulk=False):
+ 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
+ # The pool's element size and the lane's 4-column load and store of a state row.
+ self.sb = 2 if state_bf16 else 4
+ self.ld_state, self.st_state = (
+ (_ld_state_bf16, _st_state_bf16) if state_bf16 else (_ld_state, _st_state)
+ )
+ self.lds_state = _lds_state_bf16 if state_bf16 else _lds_state
+ self.bulk = bulk
+ # The register cap that lets the bulk variant hold four CTAs an SM (three with fp32).
+ self.minb = (4 if state_bf16 else 3) if bulk else 0
+ 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
+ # Shared memory: the norm's per-warp partials, the mbarriers, the staged rows.
+ self.wbytes = self.rpw * kdim * self.sb
+ self.o_mb = 64
+ self.o_st = 128
+ self.smem = self.o_st + (_WARPS * self.wbytes if bulk else 0)
+
+ @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=self.smem + 128 if self.bulk else _WARPS * 4 + 64,
+ stream=stream,
+ use_pdl=self.pdl,
+ min_blocks_per_mp=self.minb,
+ )
+
+ @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
+ if cutlass.const_expr(self.bulk):
+ sm = cutlass.utils.SmemAllocator().allocate_tensor(
+ cutlass.Int8, cute.make_layout((self.smem,), stride=(1,)), byte_alignment=128
+ )
+ sbase = cutlass.Int32(sm.iterator.toint())
+ s_red = sbase
+ s_mb = sbase + self.o_mb + warp * 8
+ s_st = sbase + self.o_st + warp * self.wbytes
+ else:
+ 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.
+ h = []
+ if cutlass.const_expr(self.bulk):
+ wbase = (
+ pool_base
+ + (
+ cutlass.Int64(cutlass.max(slot, cutlass.Int32(0))) * s_slot
+ + cutlass.Int64(hv * s_head + v0 * s_row)
+ )
+ * self.sb
+ )
+ base = wbase + cutlass.Int64(lane * _PL * self.sb)
+ if lane == 0:
+ _mbar_init(s_mb)
+ if slot >= 0:
+ _bulk_g2s(s_st, wbase, cutlass.Int32(self.wbytes), s_mb)
+ else:
+ base = (
+ pool_base
+ + (
+ cutlass.Int64(cutlass.max(slot, cutlass.Int32(0))) * s_slot
+ + cutlass.Int64(hv * s_head + v0 * s_row + lane * _PL)
+ )
+ * self.sb
+ )
+ h = [
+ self.ld_state(base + cutlass.Int64(j * s_row * self.sb), 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.
+ if cutlass.const_expr(self.bulk):
+ srow = s_st + lane * (_PL * self.sb)
+ cute.arch.sync_warp()
+ _mbar_wait(s_mb, live)
+ part = []
+ for j in cutlass.range_constexpr(self.rpw):
+ if cutlass.const_expr(self.bulk):
+ hj = self.lds_state(srow + j * (self.k * self.sb), live)
+ else:
+ hj = h[j]
+ acc = cutlass.Float32(0.0)
+ for e in cutlass.range_constexpr(_PL):
+ acc = hj[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)
+ if cutlass.const_expr(self.bulk):
+ hj = self.lds_state(srow + j * (self.k * self.sb), live)
+ else:
+ hj = h[j]
+ acc = cutlass.Float32(0.0)
+ for e in cutlass.range_constexpr(_PL):
+ if cutlass.const_expr(self.bulk):
+ # the register variant's contraction, pinned (rows from shared memory
+ # would let the backend contract the other product)
+ hj[e] = _fma(decay, hj[e], _mul(u, kn[e]))
+ else:
+ hj[e] = u * kn[e] + decay * hj[e]
+ acc = hj[e] * qn[e] + acc
+ part.append(acc)
+ self.st_state(base + cutlass.Int64(j * s_row * self.sb), live, *hj)
+ 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 cutlass.const_expr(self.bulk):
+ if lane == 0:
+ _sts_f32(s_red + warp * 4, ss)
+ cute.arch.sync_threads()
+ tot = _lds_f32(s_red)
+ for i in cutlass.range_constexpr(1, _WARPS):
+ tot = tot + _lds_f32(s_red + i * 4)
+ else:
+ 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 _bulk_eligible(pool, kdim) -> bool:
+ """The bulk copy's layout: a warp's rows contiguous, every copy 16-byte aligned."""
+ es = pool.element_size()
+ return (
+ pool.stride(2) == kdim
+ and pool.data_ptr() % 16 == 0
+ and all(s_ * es % 16 == 0 for s_ in pool.stride()[:2])
+ )
+
+
+def _use_bulk(rows, v_heads, device, state_bf16=True) -> bool:
+ """Whether a batch takes the bulk-copy variant.
+
+ Over a bf16 pool, from the batch the register variant cannot hold in one wave (two of its
+ CTAs an SM): below that the register variant's lead (its rows are in registers when the
+ wait returns) is worth more than the occupancy. An fp32 pool keeps the register variant:
+ its 64 state registers a thread already hold the CTA's 64 KB in flight.
+ """
+ sms = torch.cuda.get_device_properties(device).multi_processor_count
+ return state_bf16 and rows * v_heads > 2 * sms
+
+
+def _compile(key, k_heads, v_heads, kdim, vdim, pdl, norm, state_bf16, bulk, proj, ptrs):
+ gm = cute.AddressSpace.gmem
+ cargs = (
+ from_dlpack(proj, assumed_align=8).mark_layout_dynamic(leading_dim=1),
+ make_ptr(
+ cutlass.BFloat16 if state_bf16 else cutlass.Float32,
+ ptrs[1],
+ gm,
+ assumed_align=16 if (bulk or not state_bf16) else 8,
+ ),
+ 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, state_bf16, bulk),
+ *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
+ return fn
+
+
+def launch(
+ proj,
+ pool,
+ slots,
+ A_log,
+ dt_bias,
+ norm_weight,
+ eps,
+ k_heads,
+ kdim,
+ pdl=True,
+ out=None,
+ bulk=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.
+
+ `bulk`: None picks the variant by pool dtype and batch (`_use_bulk`), True / False force
+ it (True falls back to the register variant where the pool's layout does not admit the
+ bulk copy). Both variants are compiled together on the first call of a configuration, so
+ a batch size seen first under serving does not compile."""
+ rows = proj.shape[0]
+ v_heads, vdim = pool.shape[1], pool.shape[2]
+ norm = norm_weight is not None
+ if out is None:
+ out = torch.empty((rows, v_heads * vdim), dtype=torch.bfloat16, device=proj.device)
+ if rows == 0:
+ return out
+ pdl = bool(pdl and get_env_enable_pdl())
+ state_bf16 = pool.dtype == torch.bfloat16
+ eligible = _bulk_eligible(pool, kdim)
+ if bulk is None:
+ bulk = _use_bulk(rows, v_heads, proj.device, state_bf16)
+ bulk = bool(bulk) and eligible
+ 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(),
+ )
+ base_key = (k_heads, v_heads, kdim, vdim, pdl, norm, state_bf16)
+ fn = _COMPILED.get(base_key + (bulk,))
+ if fn is None:
+ for variant in (False, True) if eligible else (False,):
+ if base_key + (variant,) not in _COMPILED:
+ _compile(base_key + (variant,), *base_key, variant, proj, ptrs)
+ fn = _COMPILED[base_key + (bulk,)]
+ 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 or bf16 slot pool, then
+ the sigmoid-gated RMS norm of its output per v-head. With `proj` the decode input
+ projection's output row `[q (H*K) | k (H*K) | v (HV*V) | z (HV*V) | a (HV) | b (HV)]`
+ (after the conv step: q, k, v are the conv's SiLU outputs)::
+
+ o, pool = gated_delta_rule_decode_pretranspose(q, k, v, A_log, a, dt_bias, b,
+ use_qk_l2norm=True, pool, slots)
+ out = rms_norm_gated_token_major(o (bf16), z, norm_weight, eps, "sigmoid")
+
+ Returns `out` `[rows, HV * V]` bf16 and updates `pool[slots[n]]` in place; a
+ negative slot skips the row's state (no read or write) and its output is zero.
+ `o` is rounded to bf16 between the two, as the composition stores it.
+
+ Fusion boundary: both gates, the q/k L2 norm, the recurrence, the indexed state
+ read and write, the gated RMS norm. The projection and the conv are the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: rows; k_heads / v_heads (default 16 / 48) and proj_width (their projection's);
+ # state (default fp32 | "bf16"); pool (default a contiguous 461-slot pool |
+ # "coalesced": the cache manager's page, slot stride 36 layers' states -- 102 slots
+ # (slots 100 and 101), or with pool_slots an engine's slot count, the view the last
+ # layer's sub-page); pad (row 1 given slot -1); member (the layout the counts are from).
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="qwen3.8 decode, 128 rows (con 128), 461-slot pool", spec=dict(rows=128)),
+ Cell(why="qwen3.8 decode, 3 rows (past the conv-fused path)", spec=dict(rows=3)),
+ Cell(
+ why="2 rows, coalesced pool (slot stride 36 layers, slot 100), row 1 a pad slot (-1)",
+ spec=dict(rows=2, pool="coalesced", pad=True),
+ ),
+ Cell(why="bf16 pool (mamba_ssm_cache_dtype), 128 rows", spec=dict(rows=128, state="bf16")),
+ Cell(why="bf16 pool, 1 row (con 1)", spec=dict(rows=1, state="bf16")),
+ Cell(
+ why="bf16 pool, 33 rows, one pad slot (-1)", spec=dict(rows=33, pad=True, state="bf16")
+ ),
+ Cell(
+ why="bf16 pool, 2 rows, coalesced: slot stride 36 layers, slot 100",
+ spec=dict(rows=2, pool="coalesced", state="bf16"),
+ ),
+ Cell(
+ why="tep4 decode, 1 row (con 1) at 4 k / 12 v heads, slot 1873 of the 1874-slot "
+ "bf16 engine page",
+ spec=dict(_TEP4, rows=1),
+ ),
+ Cell(
+ why="tep4, 2 rows at 4 / 12 heads: slots 1873 and 0 of the 1874-slot bf16 page",
+ spec=dict(_TEP4, rows=2),
+ ),
+ Cell(
+ why="tep4, 3 rows at 4 / 12 heads over the 1874-slot bf16 engine page",
+ spec=dict(_TEP4, rows=3),
+ ),
+ Cell(
+ why="tep4, 128 rows (con 128; the bulk-copy variant) at 4 / 12 heads over the "
+ "1874-slot bf16 page",
+ spec=dict(_TEP4, rows=128),
+ ),
+ Cell(
+ why="tp4, 128 rows at 4 / 12 heads over its 1874-slot bf16 page (tep4's count), "
+ "row 1 a pad slot (-1)",
+ spec=dict(_TEP4, rows=128, pad=True, pool_slots=TP4_SLOTS, member="tp4"),
+ ),
+ Cell(
+ why="dep4 decode, 32 rows at 16 / 48 heads (proj 16480) over the 850-slot bf16 "
+ "engine page",
+ spec=dict(_DEP4, rows=32),
+ ),
+ Cell(
+ why="dep4, 128 rows (max_batch_size) at 16 / 48 heads over the 850-slot bf16 page",
+ spec=dict(_DEP4, rows=128),
+ ),
+ )
+
+ note = """
+ 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 or bf16 `[slots, HV, V, K]` with K innermost
+ (any outer strides, multiples of 4) and aligned to 4 elements; `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, the default target's decode
+ steps keep that (a layer's pool is written only by its own recurrence, `slots` is
+ filled before the step's graph launches), and verify and prefill steps do not use
+ the entry.
+
+ Launch: one CTA of 8 warps a (v-head, row), `HV * rows` CTAs, PDL (off with
+ TRTLLM_ENABLE_PDL=0); a bf16 pool's state by bulk copy into shared memory from
+ `rows * HV > 2 * SMs` (GB300's 152 SMs: 7 rows at 48 v heads, 26 at a tp4/tep4 rank's
+ 12; an fp32 pool, or a layout the copy cannot take, stays on registers), both variants
+ compiled on the first call.
+
+ Numerics: fp32 inside; each dot product over K is a lane's 4-term FMA chain and then
+ the transposed reduction's tree; `o` rounded to bf16 before the norm, the output
+ rounded once. A bf16 pool's state is rounded once per step, at the store, and `o`
+ comes from the state before that rounding. Inside torch's default band for bf16 and
+ the pool's dtype of the composition's float64 reference, not bitwise FlashInfer's
+ kernel. Batch-invariant: a row's output and state do not depend on the other rows.
+
+ Source: the CuTe DSL kernel in this file, JIT-compiled (`cute.compile`) for the running
+ device at the first call per geometry; the cells ran that path on sm_103 (GB300).
+
+ The default target: its routing admits a bf16 recurrent-state cache only, so its decode step
+ runs this entry over the bf16 pool at every batch size (the fp32-only one-launch kernel,
+ `gdn_decode_gated_norm_pdl`, is never launched here; its references are this entry's).
+ A tp4/tep4 rank holds 4 k / 12 v heads, its projection `[rows, 4120]`; a dep4 rank all
+ 16 / 48, `[rows, 16480]`. The pool is the stock manager's coalesced page: every GDN
+ layer's state of a slot in one page, slot stride 36 layers' states (7077888 elements at
+ 12 v heads, 28311552 at 48), a layer's view its sub-page.
+
+ Certified, sm_103:
+ - tp1's 16 / 48 heads. fp32 pool: 128 and 3 rows (the fewest a decode step sends here
+ once the conv-fused path stops) over a contiguous 461-slot pool, and 2 rows over a
+ 102-slot coalesced page (slot 100 ~2.8G elements past the base) with the second row a
+ pad slot (-1). bf16 pool: 128, 1 and 33 rows (one pad slot) over the contiguous pool,
+ and 2 rows over the coalesced page.
+ - A tp4/tep4 rank's 4 / 12 heads over tep4's bf16 engine page (1874 slots, the view the
+ last GDN layer's sub-page; slot 1873 sits ~1.3e10 elements past the base): 1, 2, 3
+ rows (register variant) and 128 rows (bulk copy); tp4 128 rows with a pad slot over
+ the same 1874 slots.
+ - dep4's 16 / 48 heads over its bf16 engine page (850 slots; slot 849 ~2.4e10 elements
+ past the base): 32 rows and 128 rows (`max_batch_size`), both by the bulk copy.
+ - Every other byte of a touched slot's page stays unchanged; a pad row's output is zero.
+ - Under graph replay the op reads the projection after `griddepcontrol.wait`: 16 rows at
+ 16 / 48, 3 rows (register) and 128 rows (bulk copy) at 4 / 12.
+ - Batch-invariant over a bf16 pool, bitwise, at 1-256 rows across the variant switch, at
+ 16 / 48 and 4 / 12.
+ - The bulk-copy and register variants bitwise equal (outputs and state), with and without
+ the norm, at 1-33 rows over contiguous bf16 and fp32 pools and the tp1 coalesced page,
+ and at 4 / 12 heads over a contiguous bf16 pool (3 and 128 rows) and tep4's engine page
+ (3 rows at slots 1873, 0, 5).
+ Against the fp64 reference, torch's default bands: at 4 / 12 heads at most 27 of 196608
+ output elements differ in value from the bf16-rounded reference (max 3.9e-3; relative
+ L2 <= 2.7e-5 at 128 rows, 5.2e-5 at 3 rows, where one element of 4608 differs by
+ 2.0e-3), state relative L2 <= 2.4e-5 (max 2.0e-3); dep4 at most
+ 93 of 786432 (max 7.8e-3, relative L2 3.1e-5), state relative L2 <= 1.7e-5.
+
+ The engine slot counts are at a bf16 state cache, `max_batch_size` 128 and
+ `kv_cache_free_gpu_memory_fraction` 0.6; a tp4 rank's GDN pools have tep4's per-slot
+ shape, so its cell takes tep4's count. The kernel addresses a slot as
+ `Int64(slot) * s_slot`: any slot count an int32 index can name.
+
+ Not certified: `pdl=False`; the slot counts of the members' variants (kv_fp8, kv_nvfp4,
+ chunked, reuse at dep4; MTP at tep4, `max_seq_len` 16384) and of tp4's own engine; this
+ target's own call, which its modeling does not make yet.
+ """
+
+ def __call__(
+ self,
+ proj: torch.Tensor,
+ pool: torch.Tensor,
+ slots: torch.Tensor,
+ A_log: torch.Tensor,
+ dt_bias: torch.Tensor,
+ norm_weight: torch.Tensor,
+ eps: float,
+ k_heads: int,
+ kdim: int,
+ pdl: bool = True,
+ ) -> torch.Tensor:
+ return launch(proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, pdl)
+
+ def reference(
+ self,
+ proj: torch.Tensor,
+ pool: torch.Tensor,
+ slots: torch.Tensor,
+ A_log: torch.Tensor,
+ dt_bias: torch.Tensor,
+ norm_weight: torch.Tensor,
+ eps: float,
+ k_heads: int,
+ kdim: int,
+ pdl: bool = True,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """`(out, pool_after)`: `gdn_decode_gated_norm_pdl`'s reference, the same contract --
+ the fp64 recurrence (o rounded to bf16, the pool to its dtype), then the gated norm
+ entry's fp32 reference."""
+ return _one_launch.reference(
+ proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, pdl
+ )
+
+ def is_valid(
+ self,
+ proj: torch.Tensor,
+ pool: torch.Tensor,
+ slots: torch.Tensor,
+ A_log: torch.Tensor,
+ dt_bias: torch.Tensor,
+ norm_weight: torch.Tensor,
+ eps: float,
+ k_heads: int,
+ kdim: int,
+ pdl: bool = True,
+ ) -> None:
+ rows = proj.shape[0]
+ v_heads, vdim, kd = pool.shape[1], pool.shape[2], pool.shape[3]
+ assert proj.dtype == torch.bfloat16 and proj.stride(1) == 1, "bf16 proj, unit column stride"
+ assert proj.shape[1] == 2 * k_heads * kdim + 2 * v_heads * vdim + 2 * v_heads, "proj width"
+ assert proj.data_ptr() % 8 == 0 and proj.stride(0) % 4 == 0, "8-byte aligned proj rows"
+ assert pool.dtype in (torch.float32, torch.bfloat16), "fp32 or bf16 pool"
+ assert pool.stride(3) == 1 and kd == kdim, "pool [S, HV, V, K]"
+ assert pool.data_ptr() % (_PL * pool.element_size()) == 0, "pool aligned to 4 elements"
+ assert all(s % _PL == 0 for s in pool.stride()[:3]), "pool strides multiples of 4"
+ 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"),
+ ),
+ Cell(
+ why="bf16 pool: 121 generation rows into a row range of the step output",
+ spec=dict(rows=121, out="row_range", state="bf16"),
+ ),
+ Cell(
+ why="bf16 pool: 3 rows, one pad slot (-1)",
+ spec=dict(rows=3, pad=True, out="row_range", state="bf16"),
+ ),
+ Cell(
+ why="tp4/tep4 mixed step: 3 generation rows at 4 k / 12 v heads (proj 4120), bf16 "
+ "pool, into a row range",
+ spec=dict(
+ rows=3, k_heads=4, v_heads=12, proj_width=4120, state="bf16", out="row_range"
+ ),
+ ),
+ Cell(
+ why="tp4/tep4 mixed step: 127 generation rows (max_batch_size with one context "
+ "row) at 4 / 12 heads",
+ spec=dict(
+ rows=127, k_heads=4, v_heads=12, proj_width=4120, state="bf16", out="row_range"
+ ),
+ ),
+ )
+
+ note = """
+ Metadata consumed: `slots`. Preconditions and launch: `GdnDecodeRowsGatedNorm`'s, plus
+ a bf16 `out` `[rows, HV * V]` with unit column stride. The pool rows and `slots` are
+ read before `griddepcontrol.wait`: in a mixed step nothing queued ahead of the call
+ writes the generation rows' slots (the context rows' reset and chunked recurrence touch
+ only theirs). Numerics: the recurrence of `GdnDecodeRowsGatedNorm` (the same kernel).
+
+ Certified, sm_103: at tp1's 16 / 48 heads a mixed step's 121 generation rows, 3 rows with
+ a pad slot (-1), and 2 rows over the cache manager's coalesced pool, each into a row range
+ of the step output (fp32 and bf16 pools). At a tp4/tep4 rank's 4 / 12 heads (proj 4120)
+ over a contiguous bf16 pool, 3 rows and 127 (`max_batch_size` with one context row), into
+ a row range: output at most 12 of 195072 elements off the bf16-rounded fp64 reference
+ (max 3.1e-5), state relative L2 <= 9.8e-6. The bitwise bulk / register agreement with
+ `norm` off covers this entry's launch at 4 / 12 over tep4's engine page.
+
+ Not certified: this entry's own call over an engine's page (a mixed step's eager path at
+ tp4/tep4 or dep4), `pdl=False`; at dep4 a mixed step runs the 16 / 48 cells above, over
+ tp1's pools.
+ """
+
+ def __call__(
+ self,
+ proj: torch.Tensor,
+ pool: torch.Tensor,
+ slots: torch.Tensor,
+ A_log: torch.Tensor,
+ dt_bias: torch.Tensor,
+ k_heads: int,
+ kdim: int,
+ out: torch.Tensor,
+ pdl: bool = True,
+ ) -> torch.Tensor:
+ return launch(proj, pool, slots, A_log, dt_bias, None, 0.0, k_heads, kdim, pdl, out=out)
+
+ def reference(
+ self,
+ proj: torch.Tensor,
+ pool: torch.Tensor,
+ slots: torch.Tensor,
+ A_log: torch.Tensor,
+ dt_bias: torch.Tensor,
+ k_heads: int,
+ kdim: int,
+ out: torch.Tensor,
+ pdl: bool = True,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """`(o, pool_after)`: the fp64 recurrence, o rounded to bf16 (the pool to its dtype)."""
+ rows = proj.shape[0]
+ v_heads, vdim = pool.shape[1], pool.shape[2]
+ q, k, v, _, a, b = torch.split(
+ proj,
+ [k_heads * kdim, k_heads * kdim, v_heads * vdim, v_heads * vdim, v_heads, v_heads],
+ 1,
+ )
+ o, pool_after = _recurrence_reference(
+ q.reshape(rows, 1, k_heads, kdim),
+ k.reshape(rows, 1, k_heads, kdim),
+ v.reshape(rows, 1, v_heads, vdim),
+ A_log,
+ a.reshape(rows, 1, v_heads),
+ dt_bias,
+ b.reshape(rows, 1, v_heads),
+ pool,
+ slots,
+ )
+ return o.reshape(rows, v_heads * vdim), pool_after
+
+ def is_valid(
+ self,
+ proj: torch.Tensor,
+ pool: torch.Tensor,
+ slots: torch.Tensor,
+ A_log: torch.Tensor,
+ dt_bias: torch.Tensor,
+ k_heads: int,
+ kdim: int,
+ out: torch.Tensor,
+ pdl: bool = True,
+ ) -> None:
+ v_heads, vdim = pool.shape[1], pool.shape[2]
+ gdn_decode_rows_gated_norm.is_valid(
+ proj,
+ pool,
+ slots,
+ A_log,
+ dt_bias,
+ torch.empty(vdim, dtype=torch.bfloat16),
+ 0.0,
+ k_heads,
+ kdim,
+ )
+ assert out.dtype == torch.bfloat16 and out.dim() == 2, "bf16 out [rows, HV * V]"
+ assert out.shape == (proj.shape[0], v_heads * vdim) and out.stride(1) == 1, (
+ "out [rows, HV * V], unit column stride"
+ )
+
+
+gdn_decode_rows_gated_norm = GdnDecodeRowsGatedNorm()
+gdn_decode_rows = GdnDecodeRows()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_prefill_conv_post.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_prefill_conv_post.py
new file mode 100644
index 000000000000..acb71d5577f5
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_prefill_conv_post.py
@@ -0,0 +1,1232 @@
+# 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.
+
+The conv pass is CuTe DSL, the gates and the window publish Triton. The conv pass
+reads a pool slot through a 64-bit base address (on a dep4 rank's stock conv page a
+slot's offset passes 2^31 elements from slot 1942).
+"""
+
+from __future__ import annotations
+
+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 triton.language.extra.cuda import libdevice
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
+
+#: This checkpoint's linear-attention geometry at tp1.
+NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS, HEAD_V_DIM = 16, 128, 48, 128
+QKV_DIM = 2 * NUM_K_HEADS * HEAD_K_DIM + NUM_V_HEADS * HEAD_V_DIM # 10240
+#: The target's one GDN input-projection output: [q|k|v | z | a | b].
+IN_PROJ_WIDTH = QKV_DIM + NUM_V_HEADS * HEAD_V_DIM + 2 * NUM_V_HEADS # 16480
+CONV_WIDTH = 4
+#: Largest conv-pool slot count the cache manager allocates for this checkpoint
+#: (461 slots on one GB300).
+ENGINE_SLOTS = 461
+#: One tp4/tep4 rank's share: the heads split four ways, its own `[q|k|v|z|a|b]`
+#: input projection 2560 + 1536 + 2 * 12 = 4120 wide.
+RANK4_K_HEADS, RANK4_V_HEADS = NUM_K_HEADS // 4, NUM_V_HEADS // 4
+RANK4_IN_PROJ_WIDTH = (2 * RANK4_K_HEADS + 2 * RANK4_V_HEADS) * HEAD_K_DIM + 2 * RANK4_V_HEADS
+#: Conv-pool slots the engine binds per rank (n-gram table host-resident): 1874 at
+#: tep4, 850 at dep4. A tp4 rank holds tep4's GDN state, so tep4's count bounds it.
+TEP4_SLOTS, DEP4_SLOTS = 1874, 850
+TP4_SLOTS = TEP4_SLOTS
+#: The stock manager's conv page: one slot holds all 36 GDN layers' `[C, 3]` states,
+#: so a layer's view has slot stride `36 * C * 3` and the last layer sits `35 * C * 3`
+#: elements into the page.
+PAGE_GDN_LAYERS = 36
+A_LOG_RANGE = (-5.5625, 3.5625)
+DT_BIAS_RANGE = (-8.9375, 19.25)
+SOFTPLUS_THRESHOLD = 20.0
+
+#: Channels per program of the window publish.
+_PUBLISH_BLOCK_C = 512
+#: Tokens per program of the gates.
+_GATES_BLOCK_T = 8
+
+
+@triton.jit
+def _gdn_prefill_conv_publish_kernel(
+ x_ptr,
+ stride_xt,
+ state_ptr,
+ stride_ss,
+ stride_sc,
+ stride_sw,
+ qsl_ptr,
+ slot_ptr,
+ init_ptr,
+ conv_dim,
+ pad_slot_id,
+ BLOCK_C: tl.constexpr,
+):
+ """One program: one sequence x BLOCK_C channels of its new 3-column window.
+
+ New column j holds position `e - 3 + j`: the raw input when it lies inside
+ the sequence, else the carried column `L + j` (or 0 without a carry). All
+ three are loaded before any is stored, and each channel's three columns
+ live in one thread, so a short sequence's shift reads the old window."""
+ n = tl.program_id(0)
+ c = tl.program_id(1) * BLOCK_C + tl.arange(0, BLOCK_C)
+ c_mask = c < conv_dim
+ slot = tl.load(slot_ptr + n)
+ if slot != pad_slot_id:
+ s0 = tl.load(qsl_ptr + n)
+ e = tl.load(qsl_ptr + n + 1)
+ carry = tl.load(init_ptr + n) != 0
+ base = state_ptr + slot.to(tl.int64) * stride_ss + c.to(tl.int64) * stride_sc
+ p0 = e - 3
+ p1 = e - 2
+ p2 = e - 1
+ x0 = tl.load(x_ptr + p0.to(tl.int64) * stride_xt + c, mask=c_mask & (p0 >= s0), other=0.0)
+ x1 = tl.load(x_ptr + p1.to(tl.int64) * stride_xt + c, mask=c_mask & (p1 >= s0), other=0.0)
+ x2 = tl.load(x_ptr + p2.to(tl.int64) * stride_xt + c, mask=c_mask & (p2 >= s0), other=0.0)
+ s_0 = tl.load(base + (3 + p0 - s0) * stride_sw, mask=c_mask & (p0 < s0) & carry, other=0.0)
+ s_1 = tl.load(base + (3 + p1 - s0) * stride_sw, mask=c_mask & (p1 < s0) & carry, other=0.0)
+ s_2 = tl.load(base + (3 + p2 - s0) * stride_sw, mask=c_mask & (p2 < s0) & carry, other=0.0)
+ n0 = tl.where(p0 >= s0, x0, s_0)
+ n1 = tl.where(p1 >= s0, x1, s_1)
+ n2 = tl.where(p2 >= s0, x2, s_2)
+ tl.store(base, n0, mask=c_mask)
+ tl.store(base + stride_sw, n1, mask=c_mask)
+ tl.store(base + 2 * stride_sw, n2, mask=c_mask)
+
+
+# --------------------------------------------------------------------------------------
+# The conv pass in CuTe DSL. A CTA is one 128-channel head x `_TT` tokens, 128 threads:
+# it first copies its `_TT + 3`-row tile (3 history rows) into shared memory with
+# 16-byte `cp.async` (rows outside `[0, T)` zero-filled), so the bytes in flight cost no
+# registers; then a thread (4 channels x `_TPT` consecutive tokens) slides the width-4
+# window over shared-memory rows, every element converted to fp32 once. The conv runs
+# as `fma.rn.f32x2` on channel pairs (first tap from zero, then in tap order: per element
+# an `fma.rn.f32` chain); the SiLU is `acc * rcp(1 + 2^(-acc * log2e))`
+# with `ex2.approx.ftz` and `rcp.approx.ftz`, rounded in pairs by `cvt.rn.bf16x2.f32`.
+# --------------------------------------------------------------------------------------
+
+_VEC = 4
+_TPT = 16
+_RG = 4
+_TT = _TPT * _RG
+_TILE_ROWS = _TT + 3
+_THREADS = 128
+_CHUNKS = _TILE_ROWS * 16
+
+
+def _iv(x, loc, ip):
+ return x.ir_value(loc=loc, ip=ip)
+
+
+def _asm1(ret, args, text, constraints, loc, ip, side_effects=False):
+ return llvm.inline_asm(
+ ret,
+ [_iv(a, loc, ip) for a in args],
+ text,
+ constraints,
+ has_side_effects=side_effects,
+ is_align_stack=False,
+ asm_dialect=llvm.AsmDialect.AD_ATT,
+ loc=loc,
+ ip=ip,
+ )
+
+
+def _asm2(ty, args, text, constraints, loc, ip, side_effects=False):
+ res = _asm1(
+ llvm.StructType.get_literal([ty, ty]), args, text, constraints, loc, ip, side_effects
+ )
+ return (
+ llvm.extractvalue(ty, res, [0], loc=loc, ip=ip),
+ llvm.extractvalue(ty, res, [1], loc=loc, ip=ip),
+ )
+
+
+@dsl_user_op
+def _ffma2(a0, a1, b0, b1, c0, c1, *, loc=None, ip=None):
+ """`fma.rn.f32x2`: two fp32 FMAs, each rounded once."""
+ r = _asm2(
+ T.f32(),
+ [cutlass.Float32(v) for v in (a0, a1, b0, b1, c0, c1)],
+ "{ .reg .b64 a, b, c, d; mov.b64 a, {$2, $3}; mov.b64 b, {$4, $5}; "
+ "mov.b64 c, {$6, $7}; fma.rn.f32x2 d, a, b, c; mov.b64 {$0, $1}, d; }",
+ "=f,=f,f,f,f,f,f,f",
+ loc,
+ ip,
+ )
+ return cutlass.Float32(r[0]), cutlass.Float32(r[1])
+
+
+@dsl_user_op
+def _fmul2(a0, a1, b0, b1, *, loc=None, ip=None):
+ """`mul.rn.f32x2`."""
+ r = _asm2(
+ T.f32(),
+ [cutlass.Float32(v) for v in (a0, a1, b0, b1)],
+ "{ .reg .b64 a, b, d; mov.b64 a, {$2, $3}; mov.b64 b, {$4, $5}; "
+ "mul.rn.f32x2 d, a, b; mov.b64 {$0, $1}, d; }",
+ "=f,=f,f,f,f,f",
+ loc,
+ ip,
+ )
+ return cutlass.Float32(r[0]), cutlass.Float32(r[1])
+
+
+@dsl_user_op
+def _silu2(a0, a1, *, loc=None, ip=None):
+ """`acc * rcp(1 + 2^(acc * -log2e))` for two values: `mul.rn.f32x2`, `ex2.approx.ftz`,
+ `add.rn.f32x2`, `rcp.approx.ftz`, `mul.rn`. For `1 + e` in [1, 2^126] this equals a
+ Triton `acc / (1 + exp(-acc))` (`div.full.f32`); past it (acc < -87.3) it gives -0
+ where that gives a value below 1e-36."""
+ r = _asm2(
+ T.f32(),
+ [cutlass.Float32(a0), cutlass.Float32(a1)],
+ "{ .reg .b64 a, k, t, o, d; .reg .f32 c, t0, t1, d0, d1; "
+ "mov.f32 c, 0fBFB8AA3B; mov.b64 k, {c, c}; mov.f32 c, 0f3F800000; mov.b64 o, {c, c}; "
+ "mov.b64 a, {$2, $3}; mul.rn.f32x2 t, a, k; mov.b64 {t0, t1}, t; "
+ "ex2.approx.ftz.f32 t0, t0; ex2.approx.ftz.f32 t1, t1; mov.b64 t, {t0, t1}; "
+ "add.rn.f32x2 d, t, o; mov.b64 {d0, d1}, d; "
+ "rcp.approx.ftz.f32 d0, d0; rcp.approx.ftz.f32 d1, d1; "
+ "mul.rn.f32 $0, $2, d0; mul.rn.f32 $1, $3, d1; }",
+ "=f,=f,f,f",
+ loc,
+ ip,
+ )
+ return cutlass.Float32(r[0]), cutlass.Float32(r[1])
+
+
+@dsl_user_op
+def _fma_rn(a, b, c, *, loc=None, ip=None):
+ return cutlass.Float32(
+ _asm1(
+ T.f32(),
+ [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)],
+ "fma.rn.f32 $0, $1, $2, $3;",
+ "=f,f,f,f",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _inv_norm(ss, eps, *, loc=None, ip=None):
+ """1 / sqrt(ss + eps): `add.f32`, `sqrt.approx.ftz.f32`, `div.full.f32` (a Triton
+ `1.0 / tl.sqrt(...)`)."""
+ return cutlass.Float32(
+ _asm1(
+ T.f32(),
+ [cutlass.Float32(ss), cutlass.Float32(eps)],
+ "{ .reg .f32 t; add.f32 t, $1, $2; sqrt.approx.ftz.f32 t, t; "
+ "div.full.f32 $0, 0f3F800000, t; }",
+ "=f,f,f",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _bf16x2(lo, hi, *, loc=None, ip=None):
+ """`cvt.rn.bf16x2.f32`: `lo` in the low half (the lower address)."""
+ return cutlass.Int32(
+ _asm1(
+ T.i32(),
+ [cutlass.Float32(hi), cutlass.Float32(lo)],
+ "cvt.rn.bf16x2.f32 $0, $1, $2;",
+ "=r,f,f",
+ loc,
+ ip,
+ )
+ )
+
+
+@dsl_user_op
+def _unpack2(w, *, loc=None, ip=None):
+ """A bf16x2 word as `(low, high)` fp32."""
+ r = _asm2(
+ T.f32(),
+ [cutlass.Int32(w)],
+ "{ .reg .b32 l, h; shl.b32 l, $2, 16; and.b32 h, $2, 0xFFFF0000; "
+ "mov.b32 $0, l; mov.b32 $1, h; }",
+ "=f,=f,r",
+ loc,
+ ip,
+ )
+ return cutlass.Float32(r[0]), cutlass.Float32(r[1])
+
+
+@dsl_user_op
+def _st_v2(addr, a, b, *, loc=None, ip=None):
+ _asm1(
+ T.i32(),
+ [cutlass.Int64(addr), cutlass.Int32(a), cutlass.Int32(b)],
+ "{ st.global.v2.b32 [$1], {$2, $3}; mov.u32 $0, 0; }",
+ "=r,l,r,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+
+
+@dsl_user_op
+def _cp_async16(smem_addr, gmem_addr, src_bytes, *, loc=None, ip=None):
+ """16 bytes global -> shared, the rest of the 16 zero-filled past `src_bytes`."""
+ _asm1(
+ T.i32(),
+ [cutlass.Int32(smem_addr), cutlass.Int64(gmem_addr), cutlass.Int32(src_bytes)],
+ "{ cp.async.cg.shared.global [$1], [$2], 16, $3; mov.u32 $0, 0; }",
+ "=r,r,l,r",
+ loc,
+ ip,
+ side_effects=True,
+ )
+
+
+@dsl_user_op
+def _ld_bf16(addr, *, loc=None, ip=None):
+ """One bf16 from global memory at a 64-bit byte address, as fp32."""
+ return cutlass.Float32(
+ _asm1(
+ T.f32(),
+ [cutlass.Int64(addr)],
+ "{ .reg .b32 r; ld.global.u16 r, [$1]; shl.b32 r, r, 16; mov.b32 $0, r; }",
+ "=f,l",
+ loc,
+ ip,
+ side_effects=True,
+ )
+ )
+
+
+@dsl_user_op
+def _cp_async_wait_all(*, loc=None, ip=None):
+ _asm1(T.i32(), [], "{ cp.async.wait_all; mov.u32 $0, 0; }", "=r", loc, ip, side_effects=True)
+
+
+@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 goes in as its base address and strides: a carried column is read at
+ # `base + slot * st_s0` in 64 bits (on the stock page that passes 2^31 elements
+ # from slot 1942 at 16 / 48 heads, 7768 at 4 / 12) plus its in-slot offset in 32.
+ qsl = cute.make_tensor(qsl_ptr, cute.make_layout((num_seqs + 1,), stride=(1,)))
+ slots = cute.make_tensor(slot_ptr, cute.make_layout((num_seqs,), stride=(1,)))
+ init = cute.make_tensor(init_ptr, cute.make_layout((num_seqs,), stride=(1,)))
+ 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,
+ st_ptr.toint(),
+ st_s0,
+ st_s1,
+ st_s2,
+ 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,
+ st_base: cutlass.Int64,
+ st_s0: cutlass.Int32,
+ st_s1: cutlass.Int32,
+ st_s2: cutlass.Int32,
+ mQsl: cute.Tensor,
+ mSlot: cute.Tensor,
+ mInit: cute.Tensor,
+ 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])
+ slot_base = st_base + cutlass.Int64(slot) * cutlass.Int64(st_s0) * 2
+ for v in cutlass.range_constexpr(_VEC):
+ a = cutlass.Float32(0.0)
+ for kk in cutlass.range_constexpr(4):
+ 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 = _ld_bf16(
+ slot_base
+ + cutlass.Int64((c0 + v) * st_s1 + (3 + pos - s0) * st_s2)
+ * 2
+ )
+ a = _fma_rn(xv, w[v, kk], a)
+ acc[v] = a
+ y0, y1 = _silu2(acc[0], acc[1])
+ 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,
+ LINEAR_GATE: tl.constexpr,
+):
+ """The decay/beta gates over `[T, NUM_V_HEADS]`, one row-coalesced tile a program;
+ with `LINEAR_GATE` the decay stored as `exp(g)` (libdevice `expf`, as `torch.exp`)."""
+ 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
+ g = -tl.exp(A_log) * softplus
+ if LINEAR_GATE:
+ g = libdevice.exp(g)
+ tl.store(g_ptr + g_off, g, 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 and an int32 offset inside a
+ pool slot (a slot's base is 64-bit)."""
+ in_slot = (conv_states.shape[1] - 1) * conv_states.stride(1) + 2 * conv_states.stride(2)
+ return (
+ x.dtype == torch.bfloat16
+ and head_k_dim == _VEC * 32
+ and x.data_ptr() % 16 == 0
+ and (x.stride(0) * 2) % 16 == 0
+ and max(*weight.stride(), *conv_states.stride(), in_slot) < 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,
+ linear_gate=False,
+):
+ assert head_k_dim == head_v_dim and weight.shape[-1] == CONV_WIDTH, (
+ "one program per 128-channel head: head_k_dim must equal head_v_dim, and the "
+ f"conv width must be {CONV_WIDTH}"
+ )
+ num_tokens = x.shape[0]
+ num_seqs = query_start_loc.shape[0] - 1
+ conv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim
+ dev, dt = x.device, x.dtype
+ q = torch.empty((1, num_tokens, num_k_heads, head_k_dim), dtype=dt, device=dev)
+ k = torch.empty_like(q)
+ v = torch.empty((1, num_tokens, num_v_heads, head_v_dim), dtype=dt, device=dev)
+ g = torch.empty((1, num_tokens, num_v_heads), dtype=torch.float32, device=dev)
+ beta = torch.empty((1, num_tokens, num_v_heads), dtype=torch.float32, device=dev)
+ if num_seqs == 0:
+ return q, k, v, g, beta
+ if num_tokens > 0:
+ assert _cute_conv_pass_ok(x, weight, conv_states, num_k_heads, head_k_dim, num_v_heads), (
+ "the conv pass takes bf16 x with 128-channel heads, 16-byte aligned rows and "
+ f"int32 weight/pool strides; got {x.dtype}, head_k_dim {head_k_dim}, x at "
+ f"{x.data_ptr() % 16} mod 16 with row stride {x.stride(0)}"
+ )
+ _cute_conv_pass(
+ x, weight, conv_states, query_start_loc, cache_indices, has_initial_state, q, k, v,
+ num_k_heads, num_v_heads, head_k_dim, pad_slot_id, l2_norm_eps,
+ ) # fmt: skip
+ _gdn_gates_kernel[(triton.cdiv(num_tokens, _GATES_BLOCK_T),)](
+ 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=_GATES_BLOCK_T,
+ LINEAR_GATE=linear_gate,
+ num_warps=4,
+ )
+ _gdn_prefill_conv_publish_kernel[(num_seqs, triton.cdiv(conv_dim, _PUBLISH_BLOCK_C))](
+ x,
+ x.stride(0),
+ conv_states,
+ *conv_states.stride(),
+ query_start_loc,
+ cache_indices,
+ has_initial_state,
+ conv_dim,
+ pad_slot_id,
+ BLOCK_C=_PUBLISH_BLOCK_C,
+ num_warps=4,
+ )
+ return q, k, v, g, beta
+
+
+def _engine_slots(num_seqs: int) -> tuple[int, ...]:
+ """`num_seqs` distinct slots of the engine pool, the last and slot 0 first."""
+ rest = tuple((i * 37 + 4) % (ENGINE_SLOTS - 2) + 1 for i in range(max(num_seqs - 2, 0)))
+ slots = ((ENGINE_SLOTS - 1, 0) + rest)[:num_seqs]
+ assert len(set(slots)) == num_seqs
+ return slots
+
+
+def _rank_cell(rank4: bool, lengths, initial, pool_slots, slots, linear_gate) -> dict:
+ """A rank's call: tp4/tep4 (`rank4`, 4 k / 12 v heads over its 4120-wide projection) or
+ dep4 (all 16 / 48 heads, 16480), over the stock conv page at `pool_slots`."""
+ heads = (RANK4_K_HEADS, RANK4_V_HEADS) if rank4 else (NUM_K_HEADS, NUM_V_HEADS)
+ return dict(
+ num_k_heads=heads[0],
+ num_v_heads=heads[1],
+ head_dim=HEAD_K_DIM,
+ x_row_stride=RANK4_IN_PROJ_WIDTH if rank4 else IN_PROJ_WIDTH,
+ lengths=lengths,
+ initial=initial,
+ pool_slots=pool_slots,
+ page_layers=PAGE_GDN_LAYERS,
+ slots=slots,
+ linear_gate=linear_gate,
+ )
+
+
+class GdnPrefillConvPost(OpWrapper):
+ """Causal conv + SiLU of the context rows, then the GDN q/k/v/g/beta.
+
+ `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; exp(g) with linear_gate)
+ beta = sigmoid(b) (fp32)
+ conv_states[slot] = xpad[e - b : e - b + W - 1] (raw input, exact)
+
+ -- the composition `extract_transpose_prefill_slice` -> `causal_conv1d_fwd`
+ (SiLU, no bias) -> `fused_gdn_post_conv(prefill, None, ...)`, in one pass
+ over the activation plus a per-sequence window publish.
+
+ Fusion boundary: the conv-state read, conv, SiLU, split, q/k norm, gates
+ and the state publish are inside. The projection, the generation rows'
+ conv and the recurrence are the caller's (`use_qk_l2norm_in_kernel=False`
+ downstream).
+
+ Returns `(q, k, v, g, beta)`: fresh contiguous `[1, T, num_k_heads,
+ head_k_dim]` x2, `[1, T, num_v_heads, head_v_dim]` in `x.dtype`, and
+ `[1, T, num_v_heads]` fp32 twice. `conv_states` is advanced in place for
+ every sequence whose slot is not `pad_slot_id`; nothing else is written.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: lengths; slots (pool rows used, None -> _engine_slots); pool_slots;
+ # initial; layout "in_proj" (x and a/b column views of one [T, 16480]
+ # buffer, the target's call) | "rows" (x/a/b the first rows of a longer
+ # projection, as a mixed step's context half), extra_rows. A rank's cell
+ # (`_rank_cell`) also names num_k_heads, num_v_heads, head_dim, x_row_stride
+ # (its own projection's width) and page_layers (the pool is the last layer's
+ # view of the stock page); without them, the tp1 geometry over a two-layer buffer.
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="target call, engine 461-slot pool: 37+1+128 tokens at slots 460, 0, 230",
+ spec=dict(
+ lengths=(37, 1, 128),
+ pool_slots=ENGINE_SLOTS,
+ slots=(ENGINE_SLOTS - 1, 0, ENGINE_SLOTS // 2),
+ initial=(True, False, True),
+ ),
+ ),
+ Cell(
+ why="mixed step's context half: first 8065 rows of a 8192-row projection",
+ spec=dict(
+ lengths=(8065,),
+ pool_slots=ENGINE_SLOTS,
+ initial=(False,),
+ layout="rows",
+ extra_rows=127,
+ ),
+ ),
+ Cell(
+ why="target call with linear_gate: 3000+8192+5000 tokens, the middle one fresh",
+ spec=dict(
+ lengths=(3000, 8192, 5000),
+ pool_slots=ENGINE_SLOTS,
+ initial=(True, False, True),
+ linear_gate=True,
+ ),
+ ),
+ Cell(
+ why="128 short contexts (the batch ceiling), some shorter than the window",
+ spec=dict(
+ lengths=tuple(1 + (i * 7) % 64 for i in range(128)),
+ pool_slots=ENGINE_SLOTS,
+ initial=tuple(i % 3 != 0 for i in range(128)),
+ ),
+ ),
+ Cell(
+ why="tep4 rank, 4/12 heads at stride 4120: fresh 8192-token context, last of 1874 slots",
+ spec=_rank_cell(True, (8192,), (False,), TEP4_SLOTS, (TEP4_SLOTS - 1,), True),
+ ),
+ Cell(
+ why="tep4 rank: 37+1+128 tokens, carry/fresh/carry, slots 1873, 0, 937 of 1874",
+ spec=_rank_cell(
+ True,
+ (37, 1, 128),
+ (True, False, True),
+ TEP4_SLOTS,
+ (TEP4_SLOTS - 1, 0, TEP4_SLOTS // 2),
+ False,
+ ),
+ ),
+ Cell(
+ why="tp4 rank (tep4's 4/12 heads and 1874-slot bound): fresh 8192 tokens at slot 0",
+ spec=_rank_cell(True, (8192,), (False,), TP4_SLOTS, (0,), True),
+ ),
+ Cell(
+ why="dep4 rank, 16/48 heads: two 512-token chunks, the first carrying, 850 slots",
+ spec=_rank_cell(
+ False,
+ (512, 512),
+ (True, False),
+ DEP4_SLOTS,
+ (DEP4_SLOTS - 1, DEP4_SLOTS // 2),
+ True,
+ ),
+ ),
+ Cell(
+ why="dep4 rank at max_num_tokens: one carrying 8192-token chunk, last of 850 slots",
+ spec=_rank_cell(False, (8192,), (True,), DEP4_SLOTS, (DEP4_SLOTS - 1,), True),
+ ),
+ # Pools past the int32 slot offset (slot 1942 at 16 / 48 heads, 7768 at 4 / 12),
+ # as larger dep4 pools have (kv_nvfp4, reuse, a 0.85 KV fraction).
+ Cell(
+ why="dep4 rank, 2048-slot pool: 37+1+128 tokens at slots 2047, 1942, 0, past int32",
+ spec=_rank_cell(False, (37, 1, 128), (True, False, True), 2048, (2047, 1942, 0), True),
+ ),
+ Cell(
+ why="dep4 rank, 4096-slot pool: a carrying 8192-token chunk at slot 4095, past 2^32",
+ spec=_rank_cell(False, (8192,), (True,), 4096, (4095,), True),
+ ),
+ Cell(
+ why="tp4/tep4 rank, 8192-slot pool: a fresh 8192-token context at slot 8191",
+ spec=_rank_cell(True, (8192,), (False,), 8192, (8191,), True),
+ ),
+ Cell(
+ why="tp4/tep4 rank, 8192-slot pool: carrying 8192 tokens at slot 8191, past int32",
+ spec=_rank_cell(True, (8192,), (True,), 8192, (8191,), True),
+ ),
+ )
+
+ note = """
+ 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` with `use_qk_l2norm_in_kernel=False`,
+ and `linear_gate=True` on both: the decay leaves this pass as `exp(g)`,
+ the factor FlashInfer's kernel takes, bitwise the `torch.exp` the chunk
+ op would otherwise launch.
+
+ The default target (tp4, tep4, dep4) makes the same call per rank, from its
+ prefill GDN core: a tp4 or tep4 rank over its own `[T, 4120]` projection
+ (4 key / 12 value heads, conv dim 2560: its conv weight rows, its
+ `A_log`/`dt_bias` entries and its conv pool are the rank's shares), a dep4
+ rank over the whole `[T, 16480]` projection. The pool is the layer's view
+ of the stock manager's conv page (`get_conv_states` at tp4, tep4 and
+ dep4): one slot holds all 36 GDN layers' `[C, 3]` states, so the view
+ has strides `(36 * C * 3, 3, 1)` --
+ (276480, 3, 1) at C 2560, (1105920, 3, 1) at C 10240 -- and layer `i`
+ starts `i * C * 3` elements into the page. The target's forward does
+ not call this variant yet; the rank cells mirror tp1's call at the
+ rank's widths.
+
+ Three launches. The conv pass (CuTe DSL, `_ConvPassKernel`): grid
+ `heads x cdiv(T, 64)` (80 heads at tp1 and dep4, 20 on a tp4/tep4
+ rank), 128 threads; a CTA copies its 67-row tile (3 history rows) of
+ the head's 128 channels into shared memory with
+ 16-byte `cp.async`, then a thread slides the width-4 window over 16
+ tokens of 4 channels, each element converted once; each thread finds its
+ first token's sequence by a binary search of `query_start_loc`
+ (sync-free, capture-safe), and a thread whose 16 tokens and windows lie
+ inside one sequence takes a straight-line branch, the rest read the pool
+ token by token (a carried column at the slot's 64-bit base address plus
+ its 32-bit in-slot offset, one `ld.global.u16`). The gates (Triton): one
+ row-coalesced `[8, 48]` tile a program (`exp(g)` with `linear_gate`,
+ libdevice `expf`). Then a per-sequence window publish (Triton), after
+ both, so no program reads a window another has already advanced.
+
+ Numerics: the conv result is rounded to bf16 before the norm reads it, as
+ the unfused path stores it. Conv taps: fp32 FMAs from zero in tap order
+ (`fma.rn.f32x2` on channel pairs). SiLU: `acc * rcp(1 + 2^(-acc *
+ log2e))` with `ex2.approx.ftz` and `rcp.approx.ftz` -- a Triton
+ `acc / (1 + exp(-acc))`'s `ex2.approx` / `div.full` values wherever
+ `1 + e <= 2^126`; below acc = -87.3 it returns -0 where those return a
+ value under 1e-36. q/k norm: the 128-channel fp32 sum of squares as each
+ thread's pairs plus a 32-lane butterfly, `sqrt.approx.ftz`, `div.full`.
+ Against the unfused path (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; 128 short contexts (1 to
+ 64 tokens, some shorter than the window); and the target's call with
+ `linear_gate` (3000 + 8192 + 5000 tokens, the middle one fresh).
+
+ Certified per rank, over the stock conv page's last-layer view at the
+ engine's per-rank slot counts (`TEP4_SLOTS` 1874, `DEP4_SLOTS` 850;
+ `TP4_SLOTS` is tep4's), each also checked to leave the page's other 35
+ layers untouched:
+ - 4 / 12 heads, stride 4120 (a tp4/tep4 rank, its own compile): a
+ fresh 8192-token context at the last slot and at slot 0 (with
+ `linear_gate`); 37 + 1 + 128 tokens carrying / fresh / carrying at
+ slots 1873, 0 and 937 (without).
+ - 16 / 48 heads, stride 16480 (a dep4 rank): two 512-token chunks, the
+ first carrying, at slots 849 and 425; one carrying 8192-token chunk
+ (max_num_tokens) at slot 849; both with `linear_gate`.
+ Against the reference at these cells: q/k/v differ by one bf16 ulp on
+ 1.2e-5 to 1.3e-4 of elements (largest 1.6e-2, on v near 4), within
+ torch's default bf16 band; g and beta inside the fp32 band; conv pool
+ bitwise.
+
+ Certified past the int32 slot offset, on the same page (larger dep4
+ pools: kv_fp8 binds 1217 slots, kv_nvfp4 1498; reuse and a 0.85-fraction
+ fp8 engine are uncounted), all with `linear_gate`:
+ - 16 / 48 heads, a 2048-slot pool: 37 + 1 + 128 tokens carrying /
+ fresh / carrying at slots 2047, 1942 and 0; a 4096-slot pool: one
+ carrying 8192-token chunk at slot 4095 (offset past 2^32).
+ - 4 / 12 heads, an 8192-slot pool: an 8192-token context at slot 8191,
+ fresh and carrying.
+ Same bands as above (q/k/v one-ulp flips on 2.9e-6 to 2.5e-5 of
+ elements). The carrying cells need the 64-bit slot base: an int32 slot
+ offset is an illegal memory access there, or at slot 4095 a wrapped
+ offset inside the page and a silently wrong q/k/v. The fresh 8192-slot
+ cell reads no carry.
+
+ A `pad_slot_id` context reads no carry and publishes nothing; its outputs
+ are computed from a zero window (the unfused conv leaves those columns
+ unconvolved instead -- nothing consumes either). The engine gives every
+ context row a real slot; a pad context is not certified.
+
+ Preconditions (`is_valid`): `x` 2-D, unit channel stride, at least
+ `qkv_dim` wide; `a`/`b` at least `T` rows (read through both strides);
+ A_log/dt_bias contiguous with `num_v_heads` entries; `weight` `[qkv_dim,
+ 4]` in `x`'s dtype (any strides); `conv_states` `[slots, qkv_dim, 3]` in
+ `x`'s dtype (any strides); int32 `query_start_loc`/`cache_indices`, bool
+ `has_initial_state`. Loud (`__call__`): `head_k_dim != head_v_dim`, conv
+ width other than 4, and what the conv pass cannot take: an `x` not bf16,
+ heads other than 128 channels, a row of `x` not 16-byte aligned, a weight
+ or pool stride, or an offset inside one pool slot, from 2^31 elements.
+
+ Not checked: `cache_indices` within the pool and distinct;
+ `query_start_loc` non-decreasing and ending at `T`.
+
+ Compiled at first call, per `(num_k_heads, num_v_heads, head_k_dim, row
+ stride of x)` (`_CONV_PASS`): the conv pass by CuTe DSL (TVM-FFI entry,
+ `--opt-level 2`), the gates and publish by Triton. The cells ran both JIT
+ paths on sm_103 (GB300).
+
+ Not certified: fp16/fp32 activations; head geometries other than 16 / 48
+ and 4 / 12; bias; a row stride of `x` other than the projections' 16480
+ and 4120 (the conv pass compiles one kernel per row stride); zero-token
+ contexts; `a + dt_bias` past the softplus threshold on every head; pools
+ past 4096 slots at 16 / 48 heads or 8192 at 4 / 12; tp4's own pool.
+
+ Hard limits
+ element offsets 64-bit for the rows of x, the outputs and a pool
+ slot's base (`slot * stride(0)`, in the conv pass and the publish);
+ weight and pool strides and the offset inside a slot int32 (refused
+ from 2^31) not probed beyond 8192 tokens.
+ pool slot `cache_indices` is int32, so a slot index below 2^31; driven
+ to offsets past 2^32 elements (slot 4095 at 16 / 48 heads, 4.5e9).
+ """
+
+ def __call__(
+ 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,
+ linear_gate: bool = False,
+ ) -> 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,
+ linear_gate,
+ )
+
+ 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,
+ linear_gate: bool = False,
+ ) -> 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 (with `linear_gate`, `exp` of that fp32 value in fp64, 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()
+ if linear_gate:
+ g = torch.exp(g.double()).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,
+ linear_gate: bool = False,
+ ) -> None:
+ # Every address below is computed from the arguments' own extents and
+ # strides only where stated; the rest is `ptr + index`, unmasked.
+ qkv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim
+ n_tok = x.shape[0]
+ assert x.dim() == 2 and x.stride(1) == 1 and x.shape[1] >= qkv_dim, (
+ f"x must be [T, >= {qkv_dim}] with unit channel stride; got {tuple(x.shape)} "
+ f"strides {x.stride()} -- the channel axis is read as `ptr + c`"
+ )
+ assert weight.dtype == x.dtype and weight.shape[0] == qkv_dim, (
+ f"weight must be [{qkv_dim}, 4] in x's dtype; got {tuple(weight.shape)} {weight.dtype}"
+ )
+ assert conv_states.dtype == x.dtype and tuple(conv_states.shape[1:]) == (qkv_dim, 3), (
+ f"conv_states must be [slots, {qkv_dim}, 3] in x's dtype; got "
+ f"{tuple(conv_states.shape)} {conv_states.dtype}"
+ )
+ assert a.shape[0] >= n_tok and b.shape[0] >= n_tok, (
+ f"a and b must have at least T={n_tok} rows; the token axis is unmasked"
+ )
+ assert A_log.stride(0) == 1 and dt_bias.stride(0) == 1, (
+ "A_log and dt_bias must be contiguous 1-D; they are read as `ptr + head`"
+ )
+ assert A_log.numel() >= num_v_heads and dt_bias.numel() >= num_v_heads, (
+ f"A_log and dt_bias must hold num_v_heads={num_v_heads} entries; the load is unmasked"
+ )
+ assert query_start_loc.dtype == torch.int32 and cache_indices.dtype == torch.int32, (
+ "query_start_loc and cache_indices must be int32"
+ )
+ assert has_initial_state.dtype == torch.bool, "has_initial_state must be bool"
+ n_seq = query_start_loc.shape[0] - 1
+ assert cache_indices.shape[0] >= n_seq and has_initial_state.shape[0] >= n_seq, (
+ f"cache_indices and has_initial_state must cover the {n_seq} sequences"
+ )
+
+
+gdn_prefill_conv_post = GdnPrefillConvPost()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/vision/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/vision/__init__.py
new file mode 100644
index 000000000000..52a7a9daf028
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/vision/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/vision/patch_embed_conv3d.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/vision/patch_embed_conv3d.py
new file mode 100644
index 000000000000..d0f2eabf34aa
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/vision/patch_embed_conv3d.py
@@ -0,0 +1,179 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Vision patch embedding: one 3-D convolution per pixel patch, made as the GEMM it is.
+
+With one whole patch per row and the stride equal to the kernel, the convolution
+has one output position per row: `out = rows @ filters.T + bias`, where `rows` and
+`filters` are views of `pixel_values` and `weight`. `__call__` is one
+`torch.ops.trtllm.cublas_mm` call (the C++ thop in `thop/cublasScaledMM.cpp`,
+cuBLASLt, bias fused into the epilogue) between views that launch nothing.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import (
+ Arch,
+ Cell,
+ OpWrapper,
+ true_fp32_matmul,
+)
+
+#: This checkpoint: in_channels 3, kernel (temporal_patch_size 2, patch_size
+#: 16, 16), out_channels 1152 -- weight [1152, 3, 2, 16, 16], a 1536-pixel row.
+IN_CHANNELS = 3
+KERNEL = (2, 16, 16)
+EMBED_DIM = 1152
+
+
+class PatchEmbedConv3d(OpWrapper):
+ """`out[n, k] = bias[k] + sum_{c,t,i,j} pixel[n, c, t, i, j] * weight[k, c, t, i, j]`.
+
+ One whole `[in_channels, kt, kh, kw]` patch per row of `pixel_values`, flat
+ `[N, C*kt*kh*kw]` (what the image processor emits, channel-slowest) or
+ explicit `[N, C, kt, kh, kw]`; the kernel extent is read from `weight`.
+ Cross-correlation (no kernel flip), like `nn.Conv3d`. Returns a fresh
+ contiguous `[N, out_channels]` in the input dtype.
+
+ Fusion boundary: the patch projection and its bias. Patchification, the
+ position embedding added after it, and the row order (rows pass through in
+ the order they arrive -- the merge-block order the merger needs) are the
+ caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: num_patches, dtype. The checkpoint's geometry, flat contiguous
+ # rows and a bias: what the vision tower passes. The tower is replicated,
+ # so these are every member's shapes (rank-invariant).
+ CELLS: tuple[Cell, ...] = tuple(
+ Cell(
+ why=f"checkpoint geometry, {n} patches bf16",
+ spec=dict(num_patches=n, dtype=torch.bfloat16),
+ )
+ for n in (64, 4096)
+ ) + (
+ Cell(
+ why="tower call ceiling: 8192 tokens x merge 2x2 = 32768 patches bf16",
+ spec=dict(num_patches=32768, dtype=torch.bfloat16),
+ ),
+ )
+
+ note = """
+ Source: `torch.ops.trtllm.cublas_mm`, compiled into the engine's torch-op
+ library (no JIT path). One kernel launch per call; cuBLASLt picks the
+ kernel by row count, each with the bias in its epilogue.
+
+ Metadata consumed: none. `pixel_values` is not mutated. Several images,
+ and a video's frame groups, are just more rows; there is no batch axis.
+
+ Measured on GB300 against `reference` at the three cells: 99.91-99.92% of
+ outputs bitwise equal, the rest one-ulp rounding flips of the element
+ (at most 4.4 bf16 ulps at the output RMS); relative L2 <= 9.4e-5; the worst
+ element at 0.49 of torch's default bf16 band. Each cell also holds
+ `reference`, in fp32, to an fp32 `F.conv3d` of the 5-D patches (TF32 off,
+ max |d| 2.0e-6) -- the flatten order `__call__` and `reference` share is
+ `nn.Conv3d`'s.
+
+ `torch.nn.functional.conv3d` (cuDNN) on the same patches adds the bias
+ separately in bf16, rounding a second time (relative L2 2.85e-3 against
+ `reference`).
+
+ Row-count invariance is not claimed: the first 64 rows of a 4096-row call
+ measured bitwise equal to a 64-row call, but the kernel changes with the
+ row count. Check rows against `reference`, never against a call of a
+ different size. CUDA-graph capturable (replayed at 64 and 32768 patches);
+ the output is a fresh allocation.
+
+ Silent, refused by `is_valid` (each measured unguarded): a row that is not
+ exactly one patch; row-strided pixel rows (2-D, or 5-D with a batch
+ stride) and a weight sliced along its output channels, read as if dense
+ (10^5 times outside the band); a bias in another dtype (fp32 or fp16 with
+ bf16 input) or strided (5000 times outside); a bias shorter than
+ `out_channels`, read past its end without an error; fp32 input with a
+ bias, which drops the bias. Loud, not guarded: a channels-last or
+ transposed operand fails the view, and a weight dtype unlike the input's
+ fails in cuBLASLt (CUBLAS_STATUS_NOT_SUPPORTED).
+
+ Not certified: dtypes other than bf16 (fp16 with an fp16 bias measured
+ inside the band, not a cell); other geometries; `bias=None`; patch counts
+ below 64, between the cells, and past 32768.
+ """
+
+ def __call__(
+ self,
+ pixel_values: torch.Tensor,
+ weight: torch.Tensor,
+ bias: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ out_channels, in_channels, kt, kh, kw = weight.shape
+ rows = pixel_values.view(-1, in_channels * kt * kh * kw)
+ filters = weight.view(out_channels, -1)
+ return torch.ops.trtllm.cublas_mm(rows, filters.t(), bias, None, 0, None)
+
+ def reference(
+ self,
+ pixel_values: torch.Tensor,
+ weight: torch.Tensor,
+ bias: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ """True-fp32 matmul of the flattened rows against the flattened filter."""
+ rows = pixel_values.reshape(pixel_values.shape[0], -1).float()
+ filters = weight.reshape(weight.shape[0], -1).float()
+ with true_fp32_matmul():
+ out = rows @ filters.t()
+ if bias is not None:
+ out = out + bias.float()
+ return out.to(pixel_values.dtype)
+
+ def is_valid(
+ self,
+ pixel_values: torch.Tensor,
+ weight: torch.Tensor,
+ bias: Optional[torch.Tensor] = None,
+ ) -> None:
+ # The view reinterprets anything whose element count divides: a whole
+ # [1, 3, 2, 64, 64] volume becomes 16 "patches" -- a [16, 1152] result
+ # of the right shape from the wrong pixels -- and [N, 3072] rows become
+ # 2N, without an error.
+ _, in_channels, kt, kh, kw = weight.shape
+ patch_numel = in_channels * kt * kh * kw
+ if pixel_values.dim() == 2:
+ one_patch_per_row = pixel_values.shape[1] == patch_numel
+ elif pixel_values.dim() == 5:
+ one_patch_per_row = tuple(pixel_values.shape[1:]) == (in_channels, kt, kh, kw)
+ else:
+ one_patch_per_row = False
+ assert one_patch_per_row, (
+ f"each row must be exactly one patch: pixel_values must be "
+ f"[num_patches, {patch_numel}] or "
+ f"[num_patches, {in_channels}, {kt}, {kh}, {kw}] for this weight; got "
+ f"{tuple(pixel_values.shape)}"
+ )
+ # cuBLASLt reads the rows as dense row-major and the filter as dense
+ # column-major: a row-strided view, or a weight sliced along its output
+ # channels, still views to 2-D and is read as if dense.
+ assert pixel_values.is_contiguous(), (
+ "pixel_values must be contiguous; a strided view is read as dense rows"
+ )
+ assert weight.is_contiguous(), "weight must be contiguous; a strided view is read as dense"
+ if bias is not None:
+ # The epilogue reads the bias as `out_channels` dense elements of
+ # the output dtype, and drops it altogether for fp32 input.
+ assert bias.dtype == pixel_values.dtype, (
+ f"bias dtype {bias.dtype} must equal the input dtype {pixel_values.dtype}"
+ )
+ assert bias.shape == (weight.shape[0],) and bias.is_contiguous(), (
+ f"bias must be a contiguous [{weight.shape[0]}] tensor; got "
+ f"{tuple(bias.shape)} with stride {bias.stride()}"
+ )
+ assert pixel_values.dtype != torch.float32, (
+ "the bias is silently dropped for fp32 input"
+ )
+
+
+patch_embed_conv3d = PatchEmbedConv3d()
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py
new file mode 100644
index 000000000000..bb60f45d720c
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py
@@ -0,0 +1,74 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the fused_sigmoid_mul_inplace catalog entry: every cell against
+the entry's reference with nothing written past `attention_output`'s last row."""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.activation.fused_sigmoid_mul_inplace import ( # noqa: E501
+ fused_sigmoid_mul_inplace as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "fused_sigmoid_mul_inplace requires a CUDA device"
+
+#: Elements past `attention_output`'s last row, inside the buffer it views: a
+#: masked lane that stored anyway would land here (one kernel block is 2048 wide).
+_TAIL = 2048
+
+
+def _seed(text: str) -> int:
+ return zlib.crc32(text.encode())
+
+
+def _build(spec, seed):
+ """`(out, out_buffer, gate_buffer, gate)` for one spec.
+
+ `out` is the compact `[T, hidden]` head of `out_buffer`, which runs `_TAIL`
+ elements past its last row. `gate` is a view of `gate_buffer`: the second
+ half of each `[T, heads, 2 * D]` q_proj head (`gate="interleaved"`), or the
+ 2-D gate window of the target's fused projection row (`gate="fused_row"`:
+ `row_stride` columns, the gate at `gate_col`; tp1's 14080 / 7168 where the
+ spec names neither).
+ """
+ g = torch.Generator(device="cuda").manual_seed(seed)
+ t, heads, d, dtype = spec["tokens"], spec["heads"], spec["head_dim"], spec["dtype"]
+ hidden = heads * d
+
+ def randn(*shape):
+ return torch.randn(*shape, generator=g, device="cuda").to(dtype)
+
+ out_buffer = randn(t * hidden + _TAIL)
+ out = out_buffer[: t * hidden].view(t, hidden)
+ if spec["gate"] == "interleaved":
+ gate_buffer = randn(t, heads, 2 * d)
+ gate = gate_buffer[:, :, d:]
+ else:
+ assert spec["gate"] == "fused_row", spec["gate"]
+ # The target's fused input projection row [q | k | v | gate | index scratch]:
+ # the gate a 2-D [T, hidden] window at column qkv_width.
+ row_stride, gate_col = spec.get("row_stride", 14080), spec.get("gate_col", 7168)
+ gate_buffer = randn(t, row_stride)
+ gate = gate_buffer[:, gate_col : gate_col + hidden]
+ return out, out_buffer, gate_buffer, gate
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name
+ certifying(op)
+ out, out_buffer, gate_buffer, gate = _build(cell.spec, seed=_seed(cell.why))
+ gate_before = gate_buffer.clone()
+ tail_before = out_buffer[out.numel() :].clone()
+ expected = op.reference(out, gate)
+
+ result = op(out, gate)
+
+ assert result.data_ptr() == out.data_ptr() and result.shape == out.shape
+ torch.testing.assert_close(result, expected)
+ assert torch.equal(gate_buffer, gate_before)
+ assert torch.equal(out_buffer[out.numel() :], tail_before), "a masked lane stored"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_erf.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_erf.py
new file mode 100644
index 000000000000..8e30429fc93f
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_erf.py
@@ -0,0 +1,93 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the gelu_erf catalog entry.
+
+The entry's identity is *which* GELU it is, and the wrong answer -- the tanh
+GELU this same tower's block MLPs use -- rounds to within about one bf16 ulp of
+it; only fp32 tells the two apart.
+"""
+
+import math
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.activation.gelu_erf import ( # noqa: E501
+ gelu_erf as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401
+
+assert torch.cuda.is_available(), "gelu_erf requires a CUDA device"
+
+MERGER_WIDTH = 4608
+TOWER_HIDDEN = 1152
+#: model.visual.merger, measured from the checkpoint: linear_fc1.weight std
+#: 0.010927, linear_fc1.bias std 0.002407, norm.weight mean 1.174581 / std
+#: 0.140536. The real tensors give an activand at std 0.875 / absmax 4.23.
+FC1_WEIGHT_STD = 0.010927
+FC1_BIAS_STD = 0.002407
+NORM_GAMMA_MEAN = 1.174581
+NORM_GAMMA_STD = 0.140536
+
+SQRT_2 = math.sqrt(2.0)
+SQRT_2_OVER_PI = math.sqrt(2.0 / math.pi)
+
+#: torch.testing's untouched per-dtype defaults, restated only so the
+#: discrimination test can say how far outside that gate each variant lands.
+_RTOL = {torch.bfloat16: 1.6e-2, torch.float16: 1e-3, torch.float32: 1.3e-6}
+_ATOL = 1e-5
+
+
+def _seed(text: str) -> int:
+ return zlib.crc32(text.encode())
+
+
+def _uniform(shape, lo, hi, dtype, seed):
+ g = torch.Generator(device="cuda").manual_seed(seed)
+ x = torch.empty(shape, device="cuda", dtype=torch.float32).uniform_(lo, hi, generator=g)
+ return x.to(dtype)
+
+
+def _build(spec, seed):
+ shape, dtype = spec["shape"], spec["dtype"]
+ if spec.get("dist", "uniform") == "merger":
+ g = torch.Generator(device="cuda").manual_seed(seed)
+ rows = shape[0]
+ gamma = NORM_GAMMA_MEAN + NORM_GAMMA_STD * torch.randn(
+ TOWER_HIDDEN, generator=g, device="cuda"
+ )
+ hidden = torch.randn(4 * rows, TOWER_HIDDEN, generator=g, device="cuda")
+ grouped = (hidden * gamma.to(dtype).float()).reshape(rows, MERGER_WIDTH)
+ weight = FC1_WEIGHT_STD * torch.randn(
+ MERGER_WIDTH, MERGER_WIDTH, generator=g, device="cuda"
+ )
+ bias = FC1_BIAS_STD * torch.randn(MERGER_WIDTH, generator=g, device="cuda")
+ return (grouped @ weight.to(dtype).float().t() + bias.to(dtype).float()).to(dtype)
+ lo, hi = spec["lo"], spec["hi"]
+ return _uniform(shape, lo, hi, dtype, seed)
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying): # noqa: F811
+ certifying(op)
+ x = _build(cell.spec, seed=_seed(cell.why))
+ before = x.clone()
+ out = op(x)
+ assert out.shape == x.shape and out.dtype == x.dtype and out.device == x.device
+ 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"
+ torch.testing.assert_close(out, op.reference(x))
+
+
+def _erf64(x):
+ x64 = x.double()
+ return 0.5 * x64 * (1.0 + torch.erf(x64 / SQRT_2))
+
+
+def _tanh64(x):
+ """The *other* GELU -- `gelu_pytorch_tanh`, this tower's block-MLP activation."""
+ x64 = x.double()
+ return 0.5 * x64 * (1.0 + torch.tanh(SQRT_2_OVER_PI * (x64 + 0.044715 * x64 * x64 * x64)))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_tanh.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_tanh.py
new file mode 100644
index 000000000000..14b83e9943ef
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_tanh.py
@@ -0,0 +1,84 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the qwen4_exp gelu_tanh entry.
+
+The entry's identity is *which* GELU it is: in bf16 the tanh and erf GELUs
+round to within about one ulp of each other, and only fp32 tells them apart.
+"""
+
+import math
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.activation.gelu_tanh import ( # noqa: E501
+ gelu_tanh as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401
+
+assert torch.cuda.is_available(), "gelu_tanh requires a CUDA device"
+
+INTERMEDIATE_SIZE = 4304
+HIDDEN_SIZE = 1152
+#: model.visual.blocks.*.mlp.linear_fc1, measured from the checkpoint: weight
+#: std 0.0197 (block 0), bias std 1.676 / absmax 9.44 -- the bias dominates.
+FC1_WEIGHT_STD = 0.0197
+FC1_BIAS_STD = 1.676
+
+SQRT_2_OVER_PI = math.sqrt(2.0 / math.pi)
+
+#: torch.testing's untouched per-dtype defaults, restated only so the
+#: discrimination test can say how far outside that gate each variant lands.
+_RTOL = {torch.bfloat16: 1.6e-2, torch.float16: 1e-3, torch.float32: 1.3e-6}
+_ATOL = 1e-5
+
+
+def _seed(text: str) -> int:
+ return zlib.crc32(text.encode())
+
+
+def _uniform(shape, lo, hi, dtype, seed):
+ g = torch.Generator(device="cuda").manual_seed(seed)
+ x = torch.empty(shape, device="cuda", dtype=torch.float32).uniform_(lo, hi, generator=g)
+ return x.to(dtype)
+
+
+def _build(spec, seed):
+ shape, dtype = spec["shape"], spec["dtype"]
+ if spec.get("dist", "uniform") == "fc1":
+ g = torch.Generator(device="cuda").manual_seed(seed)
+ rows = shape[0]
+ hidden = torch.randn(rows, HIDDEN_SIZE, generator=g, device="cuda").to(dtype)
+ weight = (
+ FC1_WEIGHT_STD * torch.randn(shape[-1], HIDDEN_SIZE, generator=g, device="cuda")
+ ).to(dtype)
+ bias = (FC1_BIAS_STD * torch.randn(shape[-1], generator=g, device="cuda")).to(dtype)
+ return (hidden.float() @ weight.float().t() + bias.float()).to(dtype)
+ lo, hi = spec["lo"], spec["hi"]
+ return _uniform(shape, lo, hi, dtype, seed)
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying): # noqa: F811
+ certifying(op)
+ x = _build(cell.spec, seed=_seed(cell.why))
+ before = x.clone()
+ out = op(x)
+ assert out.shape == x.shape and out.dtype == x.dtype and out.device == x.device
+ 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"
+ torch.testing.assert_close(out, op.reference(x))
+
+
+def _tanh64(x):
+ x64 = x.double()
+ return 0.5 * x64 * (1.0 + torch.tanh(SQRT_2_OVER_PI * (x64 + 0.044715 * x64 * x64 * x64)))
+
+
+def _erf64(x):
+ """The *other* GELU -- exact, erf-based; `torch.nn.GELU()` and `moe_gelu`."""
+ x64 = x.double()
+ return 0.5 * x64 * (1.0 + torch.erf(x64 / math.sqrt(2.0)))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_silu_and_mul_rows.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_silu_and_mul_rows.py
new file mode 100644
index 000000000000..7a4874047e69
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_silu_and_mul_rows.py
@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the silu_and_mul_rows entry: the cells bitwise `flashinfer_silu_and_mul` on a
+contiguous copy of the gate|up columns (and so inside its reference band) with the fused
+output's extra columns present."""
+
+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.ops.activation.silu_and_mul_rows import ( # noqa: E501
+ silu_and_mul_rows as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "silu_and_mul_rows requires a CUDA device"
+
+INTER = 640
+
+
+def _build(spec, seed):
+ g = torch.Generator(device="cuda").manual_seed(seed)
+ rows, width = spec["rows"], 2 * INTER + 8
+ x = (torch.randn(rows, width, generator=g, device="cuda") * 3.0).to(torch.bfloat16)
+ if rows:
+ x[0, :4] = torch.tensor([0.0, -0.0, 88.0, -100.0], device="cuda").to(torch.bfloat16)
+ return x
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
+ certifying(op)
+ x = _build(cell.spec, zlib.crc32(cell.why.encode()))
+ out = op(x, INTER)
+ assert out.shape == (x.shape[0], INTER) and out.is_contiguous()
+ torch.testing.assert_close(out, op.reference(x, INTER))
+ if x.shape[0]:
+ assert torch.equal(out, flashinfer_silu_and_mul(x[:, : 2 * INTER].contiguous()))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_canonical_selection.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_canonical_selection.py
new file mode 100644
index 000000000000..8701e8322c8c
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_canonical_selection.py
@@ -0,0 +1,282 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the qsa_canonical_selection entry.
+
+The rewrite fixes every slot of every row, so `_compare` is index equality against a
+composite-key top-k, and against `_torch_rewrite`, the composed-torch rewrite the fused
+kernel replaced, which also pins the rows outside the contract (NaN in a window, a
+window cut by the table's width). The input rows come from an adversary that keeps a
+random part of the tie class at the k-th value in a random order, or from the radix
+engine the target calls.
+"""
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_canonical_selection import ( # noqa: E501
+ _SPLIT_ROWS,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_canonical_selection import ( # noqa: E501
+ qsa_canonical_selection as op,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_topk_select import ( # noqa: E501
+ qsa_topk_select,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401
+
+assert torch.cuda.is_available(), "qsa_canonical_selection requires a CUDA device"
+
+DEV = "cuda"
+COMPRESS_RATIO = 4
+BLOCK_TOPK = 512
+COLS = 2080
+
+
+def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
+ """Index for index. The rule fixes every slot of every row, so nothing is tolerated: a
+ selection does no arithmetic."""
+ assert out.dtype == torch.int32, out.dtype
+ assert out.shape == ref.shape, (out.shape, ref.shape)
+ wrong = (out != ref).any(1)
+ if bool(wrong.any()):
+ row = int(wrong.nonzero()[0, 0])
+ raise AssertionError(
+ f"{int(wrong.sum())} of {out.shape[0]} rows differ from the canonical "
+ f"selection; first is row {row}"
+ )
+
+
+def _torch_rewrite(logits, row_starts, row_ends, indices) -> torch.Tensor:
+ """The rewrite as composed torch ops, in row tiles of 2^24 elements: the
+ implementation the fused kernel replaced, kept as a second oracle that is bitwise the
+ old rows wherever they disagree with `reference` (outside the contract)."""
+ num_rows, num_cols = logits.shape
+ index_topk = indices.shape[1]
+ tile = max(1, (1 << 24) // max(num_cols, 1))
+ for lo in range(0, num_rows, tile):
+ hi = min(lo + tile, num_rows)
+ scores, out = logits[lo:hi], indices[lo:hi]
+ starts, ends = row_starts[lo:hi, None], row_ends[lo:hi, None]
+ columns = torch.arange(num_cols, dtype=torch.int32, device=logits.device)[None, :]
+ windowed = torch.where((columns >= starts) & (columns < ends), scores, float("nan"))
+ picked = torch.gather(scores, 1, (out + starts).clamp_(0, num_cols - 1).long())
+ kth = picked.masked_fill_(out < 0, float("inf")).amin(1, keepdim=True)
+ at_least = windowed >= kth
+ tied = windowed == kth
+ surplus = at_least.sum(1, keepdim=True, dtype=torch.int32) - index_topk
+ dropped = tied & (torch.cumsum(tied, 1, dtype=torch.int32) <= surplus)
+ kept_seen = torch.cumsum(at_least ^ dropped, 1, dtype=torch.int32)
+ slots = torch.arange(1, index_topk + 1, dtype=torch.int32, device=logits.device)
+ slots = slots.repeat(hi - lo, 1)
+ torch.searchsorted(kept_seen, slots, out_int32=True, out=out)
+ out.sub_(starts)
+ out.masked_fill_(slots > kept_seen[:, -1:], -1)
+ return indices
+
+
+def _logits(kind: str, rows: int, cols: int, gen: torch.Generator) -> torch.Tensor:
+ if kind == "relu":
+ return torch.relu(torch.randn(rows, cols, generator=gen, device=DEV)) * 3.0
+ if kind == "relu_sparse":
+ x = torch.relu(torch.randn(rows, cols, generator=gen, device=DEV)) * 3.0
+ return x * (torch.rand(rows, cols, generator=gen, device=DEV) < 0.2)
+ if kind == "relu_rare":
+ # About one positive score in 750: a long window still has fewer than k, so
+ # the k-th value is 0 and the tie class is most of the window, half of it -0.0.
+ x = torch.relu(torch.randn(rows, cols, generator=gen, device=DEV) - 3.0)
+ flip = torch.rand(rows, cols, generator=gen, device=DEV) < 0.5
+ return torch.where(flip & (x == 0), -0.0, x)
+ if kind == "quantized":
+ # A handful of levels, so every k-th value is tied; rounding a small negative
+ # score gives -0.0.
+ return torch.round(torch.randn(rows, cols, generator=gen, device=DEV) * 2.0) / 2.0
+ raise ValueError(kind)
+
+
+def _windows(spec: dict, rows: int, cols: int, gen: torch.Generator):
+ kind = spec["windows"][0]
+ starts = torch.zeros(rows, dtype=torch.int32, device=DEV)
+ if kind == "chunk":
+ positions = spec["windows"][1] + torch.arange(rows, device=DEV)
+ else:
+ assert kind == "positions_random", kind
+ lo, hi = spec["windows"][1:]
+ positions = torch.randint(lo, hi + 1, (rows,), generator=gen, device=DEV)
+ ends = ((positions + 1) // COMPRESS_RATIO).clamp_max(cols).to(torch.int32)
+ return starts, ends
+
+
+def _poison(logits, starts, ends, kind) -> None:
+ columns = torch.arange(logits.shape[1], device=DEV)[None, :]
+ outside = (columns < starts[:, None]) | (columns >= ends[:, None])
+ if kind == "nan":
+ logits.masked_fill_(outside, float("nan"))
+ return
+ # Values that would win any comparison they were allowed into.
+ pattern = torch.tensor([float("nan"), float("inf"), 3e38, float("nan")], device=DEV)
+ logits.copy_(torch.where(outside, pattern[columns % 4].expand_as(logits), logits))
+
+
+def _order_key(logits: torch.Tensor) -> torch.Tensor:
+ """int64 that orders as the fp32 values do, -0.0 folded into +0.0."""
+ bits = (logits.float() + 0.0).view(torch.int32).long()
+ return torch.where(bits < 0, bits ^ 0x7FFFFFFF, bits)
+
+
+def _adversary(logits, starts, ends, index_topk, gen) -> torch.Tensor:
+ """An exact top-k of each window that keeps a random part of the tie class at the
+ k-th value, in a random order, `-1` last."""
+ rows, cols = logits.shape
+ columns = torch.arange(cols, device=DEV)[None, :]
+ in_window = (columns >= starts.long()[:, None]) & (columns < ends.long()[:, None])
+ # A random 32-bit low word: equal scores order at random (a collision only repeats
+ # an order), without sorting a `[rows, cols]` matrix.
+ tiebreak = torch.randint(0, 1 << 32, (rows, cols), generator=gen, device=DEV)
+ key = torch.where(
+ in_window, (_order_key(logits) << 32) | tiebreak, torch.iinfo(torch.int64).min
+ )
+ width = min(index_topk, cols)
+ top = torch.topk(key, width, dim=1).indices
+ local = torch.where(torch.gather(in_window, 1, top), top - starts.long()[:, None], -1)
+ shuffle = torch.where(local >= 0, torch.rand(local.shape, generator=gen, device=DEV), 2.0)
+ local = torch.gather(local, 1, shuffle.argsort(1))
+ out = torch.full((rows, index_topk), -1, dtype=torch.int32, device=DEV)
+ out[:, :width] = local.to(torch.int32)
+ return out
+
+
+def _radix(logits, starts, ends) -> torch.Tensor:
+ out = torch.empty(logits.shape[0], BLOCK_TOPK, dtype=torch.int32, device=DEV)
+ return qsa_topk_select(logits, starts, ends, out, index_topk=BLOCK_TOPK)
+
+
+def _build(spec: dict, rows: int, seed: int):
+ gen = torch.Generator(device=DEV).manual_seed(seed)
+ logits = _logits(spec["logits"], rows, spec["cols"], gen)
+ starts, ends = _windows(spec, rows, spec["cols"], gen)
+ _poison(logits, starts, ends, spec["tail"])
+ if spec["selection"] == "random":
+ indices = _adversary(logits, starts, ends, spec["index_topk"], gen)
+ else:
+ indices = _radix(logits, starts, ends)
+ return logits, starts, ends, indices
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811
+ certifying(op)
+ spec = cell.spec
+ seed = sum(map(ord, cell.why))
+ row_counts = spec["rows"] if isinstance(spec["rows"], tuple) else (spec["rows"],)
+ for rows in row_counts:
+ args = _build(spec, rows, seed + rows)
+ expected = op.reference(*args)
+ old = _torch_rewrite(args[0], args[1], args[2], args[3].clone())
+ out = op(*args)
+ assert out is args[3], "the rewrite writes `indices` and returns it"
+ _compare(out, expected)
+ _compare(out, old)
+
+
+def _random_windows(rows: int, cols: int, gen: torch.Generator):
+ """Windows anywhere in the table: a tenth empty, a fifth at most k wide, the rest of
+ any width the table leaves room for."""
+ starts = torch.randint(0, cols + 1, (rows,), generator=gen, device=DEV)
+ u = torch.rand(rows, generator=gen, device=DEV)
+ short = torch.randint(1, BLOCK_TOPK + 1, (rows,), generator=gen, device=DEV)
+ wide = torch.randint(0, cols + 1, (rows,), generator=gen, device=DEV)
+ length = torch.where(u < 0.1, 0, torch.where(u < 0.3, short, wide))
+ ends = (starts + length).clamp_max(cols)
+ return starts.to(torch.int32), ends.to(torch.int32)
+
+
+def _chunk_windows(rows: int, cols: int, last_position: int):
+ """A prefill chunk's causal windows, the last row at `last_position`."""
+ positions = last_position - rows + 1 + torch.arange(rows, device=DEV)
+ ends = ((positions + 1) // COMPRESS_RATIO).clamp(0, cols).to(torch.int32)
+ return torch.zeros(rows, dtype=torch.int32, device=DEV), ends
+
+
+def _long_windows(kind: str, rows: int, cols: int, gen: torch.Generator):
+ if kind == "random":
+ return _random_windows(rows, cols, gen)
+ if kind == "chunk_start":
+ return _chunk_windows(rows, cols, rows - 1)
+ assert kind == "chunk_end", kind
+ return _chunk_windows(rows, cols, cols * COMPRESS_RATIO - 1)
+
+
+#: (rows, cols, windows, logits) at long context: decode-sized to 16384-row prefill
+#: chunks, 2080- to 65536-wide tables (65536 is the 256k context), on both sides of the
+#: rows at which the op stops splitting a row into segments.
+_LONG_CASES = (
+ (1, 65536, "chunk_end", "relu_rare"),
+ (16, 32768, "random", "quantized"),
+ (256, 65536, "random", "relu_rare"),
+ (_SPLIT_ROWS - 1, 25000, "random", "quantized"),
+ (_SPLIT_ROWS, 25000, "random", "relu_rare"),
+ (4096, 8192, "random", "quantized"),
+ (4096, 32768, "chunk_end", "relu_sparse"),
+ (16384, 2080, "chunk_start", "relu_rare"),
+ (16384, 65536, "chunk_end", "quantized"),
+)
+
+
+@pytest.mark.parametrize(
+ "rows,cols,windows,kind", _LONG_CASES, ids=["-".join(map(str, c)) for c in _LONG_CASES]
+)
+def test_long_context_rows_match_reference_and_torch_rewrite(rows, cols, windows, kind) -> None:
+ """Tie classes of most of a window (relu zeros, half of them -0.0) or of a quantized
+ level, empty windows and windows of at most k, garbage past the windows."""
+ gen = torch.Generator(device=DEV).manual_seed(rows * 131 + cols)
+ logits = _logits(kind, rows, cols, gen)
+ starts, ends = _long_windows(windows, rows, cols, gen)
+ _poison(logits, starts, ends, "poison")
+ raw = _adversary(logits, starts, ends, BLOCK_TOPK, gen)
+ expected = op.reference(logits, starts, ends, raw)
+ _compare(_torch_rewrite(logits, starts, ends, raw.clone()), expected)
+ _compare(op(logits, starts, ends, raw), expected)
+
+
+@pytest.mark.parametrize("rows", (_SPLIT_ROWS - 1, 2048))
+def test_rows_outside_the_contract_match_the_torch_rewrite(rows) -> None:
+ """NaN and +-inf inside the windows, NaN picks, rows with no pick, picks past the
+ window, and windows cut short by a narrowed table: no canonical row to compare
+ against, so the rows are the torch rewrite's, bit for bit."""
+ gen = torch.Generator(device=DEV).manual_seed(5)
+ cols, live = 9000, 6000
+ clean = _logits("quantized", rows, cols, gen)
+ starts, ends = _random_windows(rows, cols, gen)
+ _poison(clean, starts, ends, "poison")
+ specials = torch.tensor([float("nan"), float("inf"), float("-inf")], device=DEV)
+ which = torch.randint(0, 3, (rows, cols), generator=gen, device=DEV)
+ sprinkle = torch.rand(rows, cols, generator=gen, device=DEV) < 0.002
+ dirty = torch.where(sprinkle, specials[which], clean)
+ # Even rows were selected over the dirty scores (NaN sorts first: NaN picks), odd
+ # rows over the clean ones.
+ seen = torch.where((torch.arange(rows, device=DEV) % 2 == 0)[:, None], dirty, clean)
+ raw = _adversary(seen, starts, ends, BLOCK_TOPK, gen)
+ raw[::7] = -1
+ raw[3::11, 0] = cols + 100
+ raw[5::13, 1] = -5
+ for scores in (dirty, dirty[:, :live]):
+ expected = _torch_rewrite(scores, starts, ends, raw.clone())
+ _compare(op(scores, starts, ends, raw.clone()), expected)
+
+
+def test_narrowing_the_table_to_a_bound_on_the_windows_changes_nothing() -> None:
+ """A caller that knows every window ends by column 1500 passes `logits[:, :1500]` of a
+ wide table: the rows are the full table's."""
+ gen = torch.Generator(device=DEV).manual_seed(21)
+ rows, cols, bound = 4096, 8192, 1500
+ logits = _logits("relu_sparse", rows, cols, gen)
+ starts = torch.zeros(rows, dtype=torch.int32, device=DEV)
+ ends = torch.randint(0, bound + 1, (rows,), generator=gen, device=DEV).to(torch.int32)
+ _poison(logits, starts, ends, "poison")
+ raw = _radix(logits, starts, ends)
+ wide = op(logits, starts, ends, raw.clone())
+ narrow = op(logits[:, :bound], starts, ends, raw.clone())
+ _compare(wide, op.reference(logits, starts, ends, raw))
+ assert torch.equal(wide, narrow)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_expand_block_indices.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_expand_block_indices.py
new file mode 100644
index 000000000000..b2e01fa46b19
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_expand_block_indices.py
@@ -0,0 +1,134 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the default target's qsa_expand_block_indices variant.
+
+Selections are built the way `qsa_topk_select` emits them: a random subset of
+the complete visible groups, prefix-packed, then `-1`. Every gate is bitwise.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_expand_block_indices import ( # noqa: E501
+ qsa_expand_block_indices as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401
+
+assert torch.cuda.is_available(), "qsa_expand_block_indices requires a CUDA device"
+
+
+def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
+ """Bitwise: these are KV-cache token indices, and one wrong entry reads
+ a different key. There is no rounding for a tolerance to absorb."""
+ assert out.dtype == ref.dtype == torch.int32, (out.dtype, ref.dtype)
+ assert out.shape == ref.shape, (out.shape, ref.shape)
+ torch.testing.assert_close(out, ref, rtol=0, atol=0)
+
+
+DEV = "cuda"
+CR, TOPK = 4, 2048
+# The tp1 engine's max_seq_len: where `_serving_rows` places its continuation chunks.
+ENGINE_MAX_SEQ_LEN = 8192
+# The default target's longest sequence: every cell's bound.
+ENVELOPE_MAX_SEQ_LEN = 262144
+
+
+def _serving_rows(requests: int, context_lens: tuple) -> tuple[list[int], list[int]]:
+ """Per-row positions and lengths of one packed step.
+
+ Chunk 0 starts at 0, so its rows span every visible-group count up to the
+ budget; later chunks are chunked-prefill continuations. Each remaining
+ request is an MTP generation request with two rows at consecutive positions.
+ """
+ positions: list[int] = []
+ lengths: list[int] = []
+ for chunk, length in enumerate(context_lens):
+ start = 0 if chunk == 0 else (907 * chunk) % (ENGINE_MAX_SEQ_LEN - length + 1)
+ positions.extend(range(start, start + length))
+ lengths.extend([start + length] * length)
+ for gen in range(requests - len(context_lens)):
+ position = 100 + (499 * gen) % 8000
+ positions.extend((position, position + 1))
+ lengths.extend((position + 2, position + 2))
+ assert max(lengths, default=0) <= ENGINE_MAX_SEQ_LEN
+ return positions, lengths
+
+
+def _decode_positions(rows: int, top: int) -> list[int]:
+ """One decode row per request, spread over [1, top): the first row sees one group,
+ the last is near `top`, and `i % 4` cycles the tail width through 0..3."""
+ step = (top - 4) // (rows - 1) // CR * CR
+ positions = [1 + i * step + i % CR for i in range(rows)]
+ assert positions[-1] < top
+ return positions
+
+
+def _rows(batch) -> tuple[list[int], list[int]]:
+ kind = batch[0]
+ if kind == "sweep":
+ positions = [4096 + 13 * i for i in range(batch[1])]
+ elif kind == "decode":
+ positions = _decode_positions(batch[1], batch[2])
+ assert max(positions) + 1 <= ENVELOPE_MAX_SEQ_LEN
+ elif kind == "serving":
+ return _serving_rows(batch[1], batch[2])
+ elif kind == "chunk":
+ start, rows = batch[1], batch[2]
+ assert start + rows <= ENVELOPE_MAX_SEQ_LEN
+ return list(range(start, start + rows)), [start + rows] * rows
+ else:
+ assert kind == "verify", kind
+ starts, per_request = batch[1], batch[2]
+ positions = [s + i for s in starts for i in range(per_request)]
+ lengths = [s + per_request for s in starts for _ in range(per_request)]
+ assert max(lengths) <= ENVELOPE_MAX_SEQ_LEN
+ return positions, lengths
+ return positions, [p + 1 for p in positions]
+
+
+def _packed_selection(positions: torch.Tensor, compress_ratio: int, token_topk: int, seed: int):
+ """A legal Top-K result: distinct complete visible groups, then -1 padding."""
+ block_topk = token_topk // compress_ratio
+ rows = positions.numel()
+ if rows == 0:
+ return torch.empty((0, block_topk), dtype=torch.int32, device=DEV)
+ available = (positions.long() + 1) // compress_ratio
+ width = max(block_topk, int(available.max()))
+ gen = torch.Generator(device=DEV).manual_seed(seed)
+ scores = torch.rand(rows, width, generator=gen, device=DEV)
+ groups = torch.arange(width, device=DEV)
+ scores = torch.where(groups[None, :] < available[:, None], scores, torch.inf)
+ chosen = scores.argsort(dim=1)[:, :block_topk]
+ taken = available.clamp(max=block_topk)
+ chosen = torch.where(torch.arange(block_topk, device=DEV)[None, :] < taken[:, None], chosen, -1)
+ return chosen.to(torch.int32)
+
+
+def _build(spec, seed: int):
+ positions, lengths = _rows(spec["batch"])
+ cr, topk = spec["compress_ratio"], spec["token_topk"]
+ 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)
+ return (blocks, query_positions, sequence_lengths), dict(compress_ratio=cr, token_topk=topk)
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying): # noqa: F811
+ certifying(op)
+ args, geometry = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
+ out = op(*args, **geometry)
+ rows = args[0].shape[0]
+ final = geometry["token_topk"] + geometry["compress_ratio"] - 1
+ assert out.dtype == torch.int32 and out.is_cuda and out.is_contiguous()
+ assert tuple(out.shape) == (rows, final)
+ _compare(out, op.reference(*args, **geometry))
+ if cell.spec["batch"][0] in ("serving", "chunk"):
+ positions, lengths = args[1], args[2]
+ assert bool((lengths.long() > positions.long() + 1).any()), (
+ "a context cell must drive sequence_lengths past position + 1"
+ )
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_mtp_reuse_topk.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_mtp_reuse_topk.py
new file mode 100644
index 000000000000..da6de8790424
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_mtp_reuse_topk.py
@@ -0,0 +1,255 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the default target's qsa_mtp_reuse_topk variant.
+
+Captured rows are built the way `qsa_topk_select` emits them: `min(K, visible)`
+distinct complete groups in an arbitrary order as a prefix, then `-1`. Inputs
+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.ops.attention.qsa_mtp_reuse_topk import ( # noqa: E501
+ qsa_mtp_reuse_topk as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401
+
+assert torch.cuda.is_available(), "qsa_mtp_reuse_topk requires a CUDA device"
+
+
+def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
+ """Bitwise: these are group ids, and one wrong entry attends to a different
+ group. There is no rounding for a tolerance to absorb."""
+ assert out.dtype == ref.dtype, (out.dtype, ref.dtype)
+ assert out.shape == ref.shape, (out.shape, ref.shape)
+ torch.testing.assert_close(out, ref, rtol=0, atol=0)
+
+
+DEV = "cuda"
+CR = 4
+#: The engine's `max_num_tokens`: height of the request-map and Top-K buffers.
+ENGINE_MAX_NUM_TOKENS = 8192
+#: 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 an 8192-token sequence allows. A cell
+#: with a longer `max_seq_len` also anchors its own last position.
+ANCHORS = (2, 3, 402, 2046, 2050, 8188, 2051, 2047, 1022, 4095, 5998, 7001, 37, 2044, 3001, 6143)
+
+
+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, max_seq_len: int) -> list[tuple[int, int]]:
+ """(captured_visible, current_visible) per request."""
+ kind = contexts[0]
+ if kind == "explicit":
+ positions = list(contexts[1])
+ else:
+ count = contexts[1]
+ top = max_seq_len - CR
+ low, high = {"mixed": (0, top), "full": (2047, top)}[kind]
+ fixed = {
+ "mixed": ANCHORS if top in ANCHORS else (*ANCHORS, top),
+ "full": (2047, 2050, top, 2051),
+ }[kind]
+ positions = list(fixed[:count])
+ positions += torch.randint(low, high + 1, (count - len(positions),), generator=gen).tolist()
+ out = []
+ for q, position in enumerate(positions):
+ d = drafted[q % len(drafted)]
+ position = min(position, 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, spec["max_seq_len"])
+ num_requests = len(pairs)
+ per_request = spec["rows_per_request"]
+ rows = num_requests * per_request
+ num_slots, slot = spec["slots"]
+ shared_rows = max(spec["shared_rows"], num_requests)
+ meta_dtype = spec["metadata_dtype"]
+
+ # Stale rows everywhere: other layers' slots and requests past this batch.
+ shared = torch.randint(-1, 3000, (num_slots, shared_rows, block_topk), generator=gen)
+ shared = shared.to(torch.int32)
+ shared_vis = torch.randint(0, 3000, (num_slots, shared_rows), generator=gen).to(meta_dtype)
+ for q, (captured, _) in enumerate(pairs):
+ shared[slot, q] = _captured_row(captured, block_topk, gen)
+ shared_vis[slot, q] = captured
+
+ request, visible = [], []
+ for q in range(num_requests):
+ captured, current = pairs[q]
+ for j in range(per_request):
+ request.append(q)
+ # A later row of the same request sees at least as many groups.
+ visible.append(current + j // CR)
+ request = torch.tensor(request, dtype=meta_dtype).reshape(rows)
+ visible = torch.tensor(visible, dtype=meta_dtype).reshape(rows)
+
+ height = max(ENGINE_MAX_NUM_TOKENS, rows + 37)
+ tall_out = torch.randint(-1, 3000, (height, block_topk), generator=gen).to(torch.int32)
+ tall_req = torch.randint(0, shared_rows, (height,), generator=gen).to(meta_dtype)
+ tall_vis = torch.randint(0, 3000, (height,), generator=gen).to(meta_dtype)
+ tall_req[:rows], tall_vis[:rows] = request, visible
+ buffers = dict(
+ shared=shared.to(DEV),
+ shared_vis=shared_vis.to(DEV),
+ request=tall_req.to(DEV),
+ visible=tall_vis.to(DEV),
+ output=tall_out.to(DEV),
+ )
+ kwargs = dict(
+ shared_indices=buffers["shared"][slot],
+ shared_visible_blocks=buffers["shared_vis"][slot],
+ request_indices=buffers["request"][:rows],
+ visible_blocks=buffers["visible"][:rows],
+ output=buffers["output"][:rows],
+ )
+ return dict(kwargs=kwargs, buffers=buffers, rows=rows)
+
+
+def _assert_legal_selection(out: torch.Tensor, kwargs: dict) -> None:
+ """What the producer's contract buys: no duplicate group, none past the horizon."""
+ if out.shape[0] == 0:
+ return
+ ids = out.long()
+ ordered = torch.sort(ids, dim=1).values
+ assert not bool(((ordered[:, 1:] == ordered[:, :-1]) & (ordered[:, 1:] >= 0)).any()), (
+ "a group id appears twice in a row"
+ )
+ captured = kwargs["shared_visible_blocks"].long()[kwargs["request_indices"].long()]
+ horizon = torch.maximum(captured, kwargs["visible_blocks"].long())
+ assert bool((ids < horizon[:, None]).all()), "a group id reaches past the row's horizon"
+ assert bool((ids >= -1).all())
+
+
+def _assert_widest_displacement_driven(spec, kwargs: dict) -> None:
+ """The cell reaches the riskiest path its offsets allow: a full row (every slot a
+ valid id) gaining `(CR - 1 + max(drafted)) // CR` fresh groups, which displaces that
+ many of its oldest ids. Seeds are fixed, so this holds the cell's coverage, not luck."""
+ block_topk = spec["block_topk"]
+ requests = kwargs["request_indices"].long()
+ filled = (kwargs["shared_indices"].long()[requests] >= 0).sum(dim=1)
+ captured = kwargs["shared_visible_blocks"].long()[requests]
+ fresh = (kwargs["visible_blocks"].long() - captured).clamp(0, block_topk)
+ widest = (CR - 1 + max(spec["drafted"])) // CR
+ assert bool(((filled == block_topk) & (fresh == widest)).any()), (
+ f"no full row gains {widest} fresh groups: the cell misses its displacement path"
+ )
+
+
+def _check_cell(spec, seed: int, requests: int | None = None) -> torch.Tensor:
+ built = _build(spec, seed, requests)
+ kwargs, buffers, rows = built["kwargs"], built["buffers"], built["rows"]
+ 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)
+ if requests is None:
+ _assert_widest_displacement_driven(spec, kwargs)
+ for name, tensor in buffers.items():
+ if name == "output":
+ assert torch.equal(tensor[rows:], before[name][rows:]), "rows past the view moved"
+ 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
+ """Eager at the cell's own request count (bitwise, legal selection, nothing outside
+ the views moved), then each of its CUDA-graph request counts."""
+ certifying(op)
+ seed = zlib.crc32(cell.why.encode())
+ _check_cell(cell.spec, seed)
+ if cell.spec["graph_batches"]:
+ _check_graph(cell.spec, seed)
+
+
+def _spec_of(prefix: str) -> dict:
+ (cell,) = [c for c in op.CELLS if c.why.startswith(prefix)]
+ return cell.spec
+
+
+def _guard_case():
+ """Legal arguments, then the whole `[slots, seqs]` visible-count buffer passed where
+ one slot's view belongs.
+
+ Returns (intended, kwargs): `intended` is the reference on the legal
+ arguments -- what a caller passing the bad layout meant to get.
+ """
+ spec = _spec_of("decode MTP")
+ kwargs = _build({**spec, "slots": (3, 2)}, seed=13, requests=6)["kwargs"]
+ intended = op.reference(**kwargs)
+ whole = torch.zeros((3, kwargs["shared_visible_blocks"].numel()), dtype=torch.int32)
+ whole = whole.to(DEV)
+ whole[2] = kwargs["shared_visible_blocks"]
+ kwargs["shared_visible_blocks"] = whole
+ return intended, kwargs
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_nvfp4_pages_to_bf16.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_nvfp4_pages_to_bf16.py
new file mode 100644
index 000000000000..61b73119dcc7
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_nvfp4_pages_to_bf16.py
@@ -0,0 +1,151 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the qsa_nvfp4_pages_to_bf16 entry.
+
+The paged state is real: one NVFP4 pool as the QSA cache manager serves it -- e2m1 data
+`[pages, role pages, Hkv, tpb, D/2]` and e4m3 block scales `[pages, role pages, Hkv, tpb,
+D/16]`, the layer's K and V two of the role planes -- and the engine's `[slots, width]`
+table, the listed requests on seeded distinct slots, their pages a seeded permutation,
+`-1` past each sequence. Scratch and table are pre-filled with sentinels, so a missing or
+a stray write shows. Every cell is bitwise `reference`, and the scratch is what the NVFP4
+sparse GQA kernels dequantize.
+"""
+
+import math
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention import (
+ qsa_paged_sparse_gqa_nvfp4 as _gqa,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_nvfp4_pages_to_bf16 import ( # noqa: E501
+ qsa_nvfp4_pages_to_bf16 as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401
+
+assert torch.cuda.is_available(), "qsa_nvfp4_pages_to_bf16 requires a CUDA device"
+
+DEV = "cuda"
+SCRATCH_FILL = 0x7FC1 # a bf16 NaN
+TABLE_FILL = -7
+
+
+def _build(spec, seed):
+ gen = torch.Generator(device=DEV).manual_seed(seed)
+ tpb, kv, d = spec["tokens_per_block"], spec["kv_heads"], spec["head_dim"]
+ counts = [math.ceil(n / tpb) for n in spec["kv_lens"]]
+ needed = sum(counts) + 7
+ pool_pages = needed + 5
+ role_pages, layer = spec["slot_pages"], spec["layer"]
+ data = torch.randint(
+ 0, 256, (pool_pages, role_pages, kv, tpb, d // 2), generator=gen, device=DEV
+ ).to(torch.uint8)
+ scales = torch.randint(
+ 0x08, 0x31, (pool_pages, role_pages, kv, tpb, d // 16), generator=gen, device=DEV
+ ).to(torch.uint8)
+ width, slots_total = spec["width"], spec["slots"]
+ table = torch.full((slots_total, width), -1, dtype=torch.int32, device=DEV)
+ slots = torch.randperm(slots_total, generator=torch.Generator().manual_seed(seed))
+ slots = slots[: len(counts)]
+ page_ids = torch.randperm(pool_pages, generator=gen, device=DEV).to(torch.int32)
+ cursor = 0
+ for slot, count in zip(slots.tolist(), counts):
+ table[slot, :count] = page_ids[cursor : cursor + count]
+ cursor += count
+ if spec["holes"]:
+ first = slots[0].item()
+ table[first, 1] = -1 # unmapped inside the sequence
+ table[first, 2] = pool_pages + 3 # past the pool
+ counts_t = torch.tensor(counts, dtype=torch.int64, device=DEV)
+ listed = int(counts_t.sum())
+ out_shape = (listed, kv, tpb, d)
+ args = dict(
+ 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=table,
+ slots=slots.to(DEV, torch.int32),
+ counts=counts_t,
+ offsets=torch.cumsum(counts_t, 0) - counts_t,
+ table=torch.full_like(table, TABLE_FILL),
+ k_out=torch.full(out_shape, SCRATCH_FILL, dtype=torch.int16, device=DEV).view(
+ torch.bfloat16
+ ),
+ v_out=torch.full(out_shape, SCRATCH_FILL, dtype=torch.int16, device=DEV).view(
+ torch.bfloat16
+ ),
+ )
+ return args
+
+
+def _bitwise(a: torch.Tensor, b: torch.Tensor) -> bool:
+ if a.dtype == torch.bfloat16:
+ a, b = a.view(torch.int16), b.view(torch.int16)
+ return a.shape == b.shape and torch.equal(a, b)
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 -- fixture by name
+ certifying(op)
+ args = _build(cell.spec, zlib.crc32(cell.why.encode()))
+ ref = op.reference(**{k: v.clone() if torch.is_tensor(v) else v for k, v in args.items()})
+ out = op(**args)
+ for name, got, want in zip(("k_out", "v_out", "table"), out, ref):
+ assert _bitwise(got, want), f"{name} differs from the reference"
+ # Every listed page was written, and only the listed requests' table entries.
+ assert not bool(out[0].isnan().any()) and not bool(out[1].isnan().any())
+ listed_rows = torch.zeros_like(args["block_table"], dtype=torch.bool)
+ for slot, count in zip(args["slots"].tolist(), args["counts"].tolist()):
+ listed_rows[slot, :count] = True
+ assert bool((out[2][~listed_rows] == TABLE_FILL).all())
+
+
+def test_scratch_is_what_the_sparse_gqa_kernels_read() -> None:
+ """The scratch holds `_dequantize_tokens`' values (the NVFP4 sparse GQA entry's own
+ dequantize, fp32 rounded once to bf16) for every listed page."""
+ spec = op.CELLS[4].spec
+ args = _build({**spec, "global_k": 0.37, "global_v": 2.5}, 11)
+ k_out, v_out, table = op(**args)
+ tpb = spec["tokens_per_block"]
+ for slot, count, base in zip(
+ args["slots"].tolist(), args["counts"].tolist(), args["offsets"].tolist()
+ ):
+ pages = args["block_table"][slot, :count].long()
+ tokens = torch.arange(tpb, device=DEV)[None, :].expand(count, -1)
+ for got, data, scales, index, grouped in (
+ (k_out, args["k_cache"], args["k_scale_cache"], 1, False),
+ (v_out, args["v_cache"], args["v_scale_cache"], 2, True),
+ ):
+ want = _gqa._dequantize_tokens(
+ data,
+ scales,
+ pages[:, None].expand(-1, tpb),
+ tokens,
+ args["kv_scale_quant_orig"][index],
+ grouped,
+ ) # [count, tpb, H, D]
+ want = want.transpose(1, 2).to(torch.bfloat16)
+ assert _bitwise(got[base : base + count], want)
+ assert torch.equal(
+ table[slot, :count], torch.arange(base, base + count, device=DEV, dtype=torch.int32)
+ )
+
+
+def test_listing_past_the_scratch_is_dropped() -> None:
+ """A request whose pages run past `listed` writes only those below it."""
+ spec = op.CELLS[2].spec
+ args = _build(spec, 5)
+ short = args["k_out"].shape[0] - 3
+ args["k_out"], args["v_out"] = args["k_out"][:short], args["v_out"][:short]
+ ref = op.reference(**{k: v.clone() if torch.is_tensor(v) else v for k, v in args.items()})
+ out = op(**args)
+ for got, want in zip(out, ref):
+ assert _bitwise(got, want)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_index_scores.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_index_scores.py
new file mode 100644
index 000000000000..84f02ccb594e
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_index_scores.py
@@ -0,0 +1,460 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the qwen4_exp qsa_paged_index_scores entry.
+
+The paged state is real: a `[pages, tokens_per_block, 1, dim]` index cache
+(a view of a wider page under the NVFP4 cells) sized from the pages the table maps, and an
+int32 block table whose rows are request *slots* -- more of them than the batch has requests, assigned by a
+seeded permutation -- with shuffled physical pages and `-1` past each
+sequence on the serving layouts, which is what the engine's table carries.
+
+The four-rank layouts' cells add one rank's measured pool: every attention layer's
+plane in one page (the last layer's plane taken), NaN wherever no request maps a
+page, with the pool's first and last page among those the cell reads -- the last
+one past 2^31 elements from the plane's base.
+
+The call has four kernels behind it: the entry's own visible-split kernel for
+ungrouped decode-sized visible-only calls and request-group kernel for verify-shaped
+ones, and the engine's row-wise and query-tile kernels. `_watch_launches` swaps the four Triton entry points
+for recording proxies that still launch the real kernel, so every cell asserts
+the branch it names.
+"""
+
+import math
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention import (
+ qsa_paged_index_scores as _entry,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_index_scores import ( # noqa: E501
+ qsa_paged_index_scores as op,
+)
+from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels as _qsa_kernels
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401
+
+assert torch.cuda.is_available(), "qsa_paged_index_scores requires a CUDA device"
+
+
+def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
+ """Torch's default fp32 band on every specified column; `-inf` exactly.
+
+ bf16 and fp16 products are exact in the kernel's fp32 accumulator, so
+ only summation order separates it from the reference, well inside the
+ default rtol 1.3e-6 / atol 1e-5 on every cell, the measured-pool ones
+ included. Wrong variants (no relu, mean over heads, first group member as
+ key, unscaled, 1/head_dim scale, block table ignored) land orders of
+ magnitude outside it. Columns the reference marks NaN are unspecified and
+ not read.
+ """
+ assert out.dtype == ref.dtype == torch.float32, (out.dtype, ref.dtype)
+ assert out.shape == ref.shape, (out.shape, ref.shape)
+ specified = ~torch.isnan(ref)
+ torch.testing.assert_close(out[specified], ref[specified])
+
+
+DEV = torch.device("cuda")
+#: The largest batch the visible-split kernel serves.
+DECODE_ROW_LIMIT = 256
+#: Index heads must divide this into more than one row for the query tile.
+DOT_TILE_WIDTH = 128
+#: Index heads must divide this for the visible-split kernel.
+SPLIT_DOT_WIDTH = 16
+#: A table width with columns enough for several of the request-group kernel's tiles.
+_WIDTH_PAGES = 260
+
+
+class _LaunchRecorder:
+ """Transparent proxy over a Triton entry point: records the grid, then launches."""
+
+ def __init__(self, inner, name, log):
+ self._inner, self._name, self._log = inner, name, log
+
+ def __getitem__(self, grid):
+ self._log.append((self._name, tuple(grid)))
+ return self._inner[grid]
+
+ def __getattr__(self, item):
+ return getattr(self._inner, item)
+
+
+class _watch_launches:
+ """Record which of the entry's four Triton kernels a call launches."""
+
+ _NAMES = (
+ (_qsa_kernels, "_qsa_paged_index_scores_kernel", "row_wise"),
+ (_qsa_kernels, "_qsa_paged_index_scores_qtile_kernel", "query_tile"),
+ (_entry, "_qsa_index_scores_group_kernel", "request_group"),
+ (_entry, "_qsa_index_scores_visible_split_kernel", "visible_split"),
+ )
+
+ def __enter__(self):
+ self.log = []
+ self._saved = [(module, name, getattr(module, name)) for module, name, _ in self._NAMES]
+ for (module, name, tag), (_, _, kernel) in zip(self._NAMES, self._saved):
+ setattr(module, name, _LaunchRecorder(kernel, tag, self.log))
+ return self
+
+ def __exit__(self, *exc):
+ for module, name, kernel in self._saved:
+ setattr(module, name, kernel)
+ return False
+
+ @property
+ def branch(self):
+ assert len(self.log) == 1, f"expected exactly one kernel launch, saw {self.log}"
+ return self.log[0][0]
+
+ @property
+ def grid(self):
+ assert len(self.log) == 1, f"expected exactly one kernel launch, saw {self.log}"
+ return self.log[0][1]
+
+
+def _case(rows, requests, **overrides):
+ """A narrow-table spec for the non-cell tests, same geometry as the entry's own."""
+ spec = dict(
+ layout="random",
+ rows=rows,
+ requests=requests,
+ slots=requests + 4,
+ pages=8,
+ heads=4,
+ head_dim=128,
+ tokens_per_block=32,
+ compress_ratio=4,
+ dtype=torch.bfloat16,
+ positions_dtype=torch.int32,
+ request_dtype=torch.int32,
+ q_layout="packed",
+ only_visible_blocks=False,
+ context_rows=False,
+ )
+ spec.update(overrides)
+ return spec
+
+
+def _rows_and_requests(spec, gen):
+ """Per-row positions, per-row request ordinal, and each request's mapped pages."""
+ layout, tpb, pages = spec["layout"], spec["tokens_per_block"], spec["pages"]
+
+ def pages_for(seq_len):
+ return math.ceil(seq_len / tpb)
+
+ if layout == "random":
+ rows, requests = spec["rows"], spec["requests"]
+ req = torch.arange(rows, device=DEV) * requests // max(rows, 1)
+ positions = torch.randint(0, pages * tpb, (rows,), generator=gen, device=DEV)
+ return positions, req, [pages] * requests
+
+ # (first position, row count) per request, in row order
+ if layout == "capture":
+ segments = [(0, n) for n in spec["lengths"]]
+ elif layout == "decode":
+ segments = [(p, spec["rows_per_request"]) for p in spec["gen_positions"]]
+ else:
+ raise AssertionError(f"unknown layout {layout!r}")
+ positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments])
+ req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)])
+ mapped = [pages_for(s + n) for s, n in segments]
+ assert max(mapped) <= pages
+ return positions, req, mapped
+
+
+def _build(spec, seed):
+ """Real paged state for one spec, keyed by `__call__`'s parameter names."""
+ gen = torch.Generator(device=DEV).manual_seed(seed)
+ heads, head_dim, dtype = spec["heads"], spec["head_dim"], spec["dtype"]
+ positions, req, mapped = _rows_and_requests(spec, gen)
+ rows, slots = positions.numel(), spec["slots"]
+ assert len(mapped) <= slots
+
+ slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(mapped)]
+ total = sum(mapped)
+ pool_pages = spec.get("pool_pages", total + 3)
+ page_ids = torch.randperm(pool_pages, generator=gen, device=DEV)
+ if "pool_pages" in spec:
+ # A measured pool: its first and last page go to the first request's first two
+ # logical pages, so the cell reads both ends.
+ ends = torch.tensor([pool_pages - 1, 0], device=DEV)
+ page_ids = torch.cat([ends, page_ids[(page_ids != 0) & (page_ids != pool_pages - 1)]])
+ page_ids = page_ids[:total].to(torch.int32)
+ block_table = torch.full((slots, spec["pages"]), -1, dtype=torch.int32, device=DEV)
+ cursor = 0
+ for slot, count in zip(slot_of.tolist(), mapped):
+ block_table[slot, :count] = page_ids[cursor : cursor + count]
+ cursor += count
+ tpb = spec["tokens_per_block"]
+ if "page_pad" in spec:
+ # The trunk pool's page: every attention layer's index buffer (the fp8 and NVFP4
+ # pools then the position page), so one layer's plane is a view at page stride
+ # buffers * tpb * dim + page_pad -- the last layer's here. Pages no request maps
+ # hold NaN: a read of one shows up in the scores.
+ buffers, buffer = spec["pool_buffers"], tpb * head_dim
+ page = torch.full(
+ (pool_pages, buffers * buffer + spec["page_pad"]), float("nan"), dtype=dtype, device=DEV
+ )
+ index_cache = page[:, (buffers - 1) * buffer : buffers * buffer].view(
+ pool_pages, tpb, 1, head_dim
+ )
+ index_cache[page_ids.long()] = torch.randn(
+ total, tpb, 1, head_dim, generator=gen, device=DEV
+ ).to(dtype)
+ assert index_cache.stride() == (page.stride(0), head_dim, head_dim, 1)
+ else:
+ index_cache = torch.randn(pool_pages, tpb, 1, head_dim, generator=gen, device=DEV).to(dtype)
+
+ if spec["q_layout"] == "scratch_slice":
+ # [T, index heads | 1 k | 1 v, D], q is the first `heads` slots.
+ scratch = torch.randn(rows, heads + 2, head_dim, generator=gen, device=DEV).to(dtype)
+ q = scratch[:, :heads]
+ else:
+ q = torch.randn(rows, heads, head_dim, generator=gen, device=DEV).to(dtype)
+ if spec["q_layout"] != "packed" and rows > 1: # one row is trivially contiguous
+ assert not q.is_contiguous()
+
+ return dict(
+ q=q,
+ index_cache=index_cache,
+ block_table=block_table,
+ query_positions=positions.to(spec["positions_dtype"]),
+ request_indices=slot_of.to(DEV)[req].to(spec["request_dtype"]),
+ tokens_per_block=spec["tokens_per_block"],
+ compress_ratio=spec["compress_ratio"],
+ only_visible_blocks=spec["only_visible_blocks"],
+ context_rows=spec["context_rows"],
+ use_pdl=spec.get("use_pdl", False),
+ prefetch=spec.get("prefetch", False),
+ rows_per_request=spec.get("call_rows_per_request", 1),
+ )
+
+
+def test_visible_split_metadata_offsets_reuse_the_compiled_kernel() -> None:
+ """Mixed-step metadata offsets preserve results without new alignment variants."""
+ cell = next(c for c in op.CELLS if c.spec["branch"] == "visible_split")
+ args = _build(cell.spec, seed=81)
+ rows = args["q"].shape[0]
+ position_storage = torch.zeros(rows + 4, dtype=args["query_positions"].dtype, device=DEV)
+ request_storage = torch.zeros(rows + 4, dtype=args["request_indices"].dtype, device=DEV)
+ compiled = []
+ for position_offset, request_offset in ((0, 0), (0, 1), (1, 1), (1, 3)):
+ positions = position_storage[position_offset : position_offset + rows]
+ requests = request_storage[request_offset : request_offset + rows]
+ positions.copy_(args["query_positions"])
+ requests.copy_(args["request_indices"])
+ shifted = {**args, "query_positions": positions, "request_indices": requests}
+ with _watch_launches() as watcher:
+ got = op(**shifted)
+ torch.cuda.synchronize()
+ assert watcher.branch == "visible_split"
+ _compare(got, op.reference(**shifted))
+ cache = _entry._qsa_index_scores_visible_split_kernel.device_caches[0][0]
+ compiled.append(len(cache))
+ assert len(set(compiled)) == 1, compiled
+
+
+def _assert_reaches_the_pool_ends(args, spec) -> None:
+ """A measured-pool cell scores keys on the pool's first and last page, and the last page
+ lies past 2^31 elements from the plane's base, so a 32-bit offset would wrap.
+ """
+ cache, pool_pages = args["index_cache"], spec["pool_pages"]
+ tpb, dim = spec["tokens_per_block"], spec["head_dim"]
+ page_stride = spec["pool_buffers"] * tpb * dim + spec["page_pad"]
+ assert cache.shape[0] == pool_pages and cache.stride() == (page_stride, dim, dim, 1)
+ assert (pool_pages - 1) * page_stride >= 2**31
+ cr = args["compress_ratio"]
+ logical = (torch.arange(_columns(args), device=DEV) * cr + cr - 1) // tpb
+ pages = args["block_table"][args["request_indices"].long()][:, logical]
+ read = pages[_visible(args)]
+ assert bool((read == 0).any()) and bool((read == pool_pages - 1).any()), "a pool end is unread"
+
+
+def _branch_rule(rows, spec):
+ heads = spec["heads"]
+ if (
+ spec.get("call_rows_per_request", 1) == 1
+ and spec["only_visible_blocks"]
+ and not spec.get("prefetch", False)
+ and rows <= DECODE_ROW_LIMIT
+ and SPLIT_DOT_WIDTH % heads == 0
+ ):
+ return "visible_split"
+ packs = DOT_TILE_WIDTH % heads == 0 and DOT_TILE_WIDTH // heads > 1
+ if spec["context_rows"] and rows > DECODE_ROW_LIMIT and packs:
+ return "query_tile"
+ grouped = (
+ spec.get("call_rows_per_request", 1) > 1
+ and spec["only_visible_blocks"]
+ and not spec.get("prefetch", False)
+ and heads in (2, 4, 8)
+ )
+ return "request_group" if grouped else "row_wise"
+
+
+def _columns(args):
+ return args["block_table"].shape[1] * args["tokens_per_block"] // args["compress_ratio"]
+
+
+def _engine_scores(**args):
+ """The engine's own launcher on the entry's arguments; `rows_per_request` is the entry's."""
+ args.pop("rows_per_request", None)
+ return _qsa_kernels.triton_qsa_paged_index_scores(**args)
+
+
+def _visible(args):
+ blocks = torch.arange(_columns(args), device=DEV)
+ bound = (args["query_positions"].long() + 1) // args["compress_ratio"]
+ return blocks[None, :] < bound[:, None]
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811
+ certifying(op)
+ args = _build(cell.spec, seed=sum(map(ord, cell.why)))
+ rows = args["q"].shape[0]
+ assert cell.spec["branch"] == _branch_rule(rows, cell.spec), "the cell names the wrong kernel"
+ if "pool_pages" in cell.spec:
+ _assert_reaches_the_pool_ends(args, cell.spec)
+ with _watch_launches() as watcher:
+ out = op(**args)
+ torch.cuda.synchronize()
+ assert watcher.branch == cell.spec["branch"]
+ assert out.shape == (rows, _columns(args)) and out.dtype == torch.float32
+ _compare(out, op.reference(**args))
+
+
+#: The decode and verify cells over a measured pool: the sites the step's CUDA graph captures.
+_CAPTURED = tuple(c for c in op.CELLS if "pool_pages" in c.spec and not c.spec["context_rows"])
+
+
+#: The visible-split cells on the tune's table: their grid is fixed at capture, the windows not.
+_SPLIT_CAPTURED = tuple(
+ c for c in _CAPTURED if c.spec["branch"] == "visible_split" and c.spec["pages"] > 516
+)
+
+
+@pytest.mark.parametrize("cell", _SPLIT_CAPTURED, ids=[c.why[:44] for c in _SPLIT_CAPTURED])
+def test_captured_visible_split_follows_the_positions(cell) -> None:
+ """The visible-split grid is sized by the device, and each program reads its rows'
+ positions when it runs: capture at the cell's positions, then replay with every row's
+ window shrunk by a different amount (to empty for some) and grown back, and match the
+ reference on every specified column each time.
+ """
+ args = _build(cell.spec, seed=sum(map(ord, cell.why)) + 2)
+ warmup = torch.cuda.Stream()
+ warmup.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(warmup):
+ op(**args)
+ torch.cuda.current_stream().wait_stream(warmup)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ captured = op(**args)
+ positions = args["query_positions"]
+ original = positions.clone()
+ gen = torch.Generator(device=DEV).manual_seed(5)
+ for trial in range(3):
+ if trial < 2:
+ keep = torch.rand(original.shape, generator=gen, device=DEV) ** (trial + 1)
+ positions.copy_((original.double() * keep).to(original.dtype))
+ else:
+ positions.copy_(original)
+ graph.replay()
+ torch.cuda.synchronize()
+ _compare(captured, op.reference(**args))
+
+
+def test_prefetch_reads_the_fresh_anchor_key_and_q_after_the_wait() -> None:
+ """`prefetch` reads pages and keys before the PDL wait. Under graph replay a producer
+ rewrites q and, at rows whose position ends a group, the last visible block's anchor
+ key -- what the decode prologue writes -- right before the op; both must be seen.
+ The other keys are static, as they are in the decode step. Without prefetch the
+ entry takes its visible-split kernel, so the plain row-wise launch it is held to
+ bitwise is the engine callable's.
+ """
+ args = _build(_case(4, 4, only_visible_blocks=True), seed=411)
+ cr, tpb = args["compress_ratio"], args["tokens_per_block"]
+ pos = args["query_positions"].clone()
+ assert bool((pos >= cr).all())
+ ends = torch.tensor([True, False, True, True], device=DEV)
+ args["query_positions"] = torch.where(ends, pos - (pos + 1) % cr, pos).to(pos.dtype)
+ pos = args["query_positions"].long()
+ assert bool(((pos + 1) % cr == 0).eq(ends).all())
+ pages = args["block_table"][args["request_indices"].long(), pos // tpb].long()
+ anchors = (pages[ends], (pos % tpb)[ends])
+ visible = _visible(args)
+ plain = _engine_scores(**args)
+ assert torch.equal(op(**{**args, "use_pdl": True, "prefetch": True})[visible], plain[visible])
+ src_q = args["q"].clone()
+ src_k = args["index_cache"][anchors].clone()
+
+ def step():
+ args["q"].copy_(src_q * 1.0)
+ args["index_cache"][anchors] = src_k * 1.0
+ return op(**{**args, "use_pdl": True, "prefetch": True})
+
+ warmup = torch.cuda.Stream()
+ warmup.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(warmup):
+ step()
+ torch.cuda.current_stream().wait_stream(warmup)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ graph_out = step()
+ for trial in range(3):
+ fresh = _build(_case(4, 4, only_visible_blocks=True), seed=412 + trial)
+ src_q.copy_(fresh["q"])
+ src_k.copy_(torch.randn(src_k.shape, device=DEV).to(src_k.dtype))
+ graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(graph_out[visible], _engine_scores(**args)[visible]), f"replay {trial}"
+ _compare(graph_out, op.reference(**args))
+
+
+@pytest.mark.parametrize(
+ "group, rows, heads, mixed, use_pdl",
+ [
+ (2, 16, 4, False, False),
+ (3, 13, 4, False, False),
+ (4, 64, 4, False, True),
+ (4, 62, 4, False, False),
+ (4, 64, 4, True, False),
+ (8, 40, 4, False, False),
+ (8, 36, 4, True, True),
+ (4, 16, 2, False, False),
+ (4, 16, 8, False, False),
+ (2, 16, 16, False, False), # 16 heads stay row-wise
+ (4, 256, 4, False, True),
+ ],
+)
+def test_request_group_is_bitwise_the_row_wise_kernel(group, rows, heads, mixed, use_pdl) -> None:
+ """`rows_per_request` picks the request-group kernel; its scores are the row-wise
+ kernel's bit patterns on every visible column. Runs of `group` rows of one request
+ (a partial last run when `group` does not divide the rows), or with `mixed` a
+ random request a row, so runs that name several requests gather row by row.
+ """
+ spec = _case(rows, 8, heads=heads, pages=_WIDTH_PAGES, only_visible_blocks=True)
+ spec.update(q_layout="scratch_slice", use_pdl=use_pdl)
+ args = _build(spec, seed=97 * group + rows + heads)
+ slots = torch.unique(args["request_indices"])
+ gen = torch.Generator(device=DEV).manual_seed(rows)
+ if mixed:
+ pick = torch.randint(0, slots.numel(), (rows,), generator=gen, device=DEV)
+ else:
+ pick = (torch.arange(rows, device=DEV) // group) % slots.numel()
+ args["request_indices"] = slots[pick].to(args["request_indices"].dtype)
+ # Every row sees at least one block (every request maps the whole table).
+ args["query_positions"] = args["query_positions"].clamp_min(args["compress_ratio"] - 1)
+ visible = _visible(args)
+ with _watch_launches() as watcher:
+ row_wise = _engine_scores(**args)
+ assert watcher.branch == "row_wise"
+ with _watch_launches() as watcher:
+ grouped = op(**{**args, "rows_per_request": group})
+ assert watcher.branch == ("request_group" if heads in (2, 4, 8) else "row_wise")
+ assert torch.equal(grouped[visible].view(torch.int32), row_wise[visible].view(torch.int32)), (
+ "the request-group kernel's bits differ from the row-wise kernel's"
+ )
+ _compare(grouped, op.reference(**args))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_kv_store.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_kv_store.py
new file mode 100644
index 000000000000..dc4581afb8c9
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_kv_store.py
@@ -0,0 +1,346 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the qsa_paged_kv_store entry.
+
+The paged state is real: an HND pool whose K and V role views are the caches, and the
+engine's `[slots, width]` table -- a row per request slot, each live row's pages drawn from a
+seeded permutation of the pool, `-1` past its allocation, and the rows past the batch holding a
+previous step's plausible page ids. The pool is pre-filled with a sentinel so a stray write
+shows. A cell's pool is one of two:
+
+- a contiguous `[pages, 2, Hkv, tpb, D]` pool, compared whole;
+- with `pool="engine"`, the stock cache manager's layout at an engine's measured page count:
+ every attention layer's K and V planes in one page (24 role pages), the views at layer 47's
+ planes. At ~257k pages that spans ~100 GB, so it is a reserved virtual range with physical
+ memory mapped under the pages the rows touch only (`_EnginePool`); the pool's last page is
+ always one of them. The reference runs on a dense copy of those pages, and every other mapped
+ byte is checked unchanged.
+"""
+
+import zlib
+
+import cuda.bindings.driver as cuda_drv
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_kv_store import ( # noqa: E501
+ qsa_paged_kv_store as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401
+
+assert torch.cuda.is_available(), "qsa_paged_kv_store requires a CUDA device"
+
+
+def _compare(out, ref) -> None:
+ """Bitwise over both role planes, compared as raw bits.
+
+ A store is a move with one conversion, and `reference` performs the
+ same conversion, so there is no rounding to spend a tolerance on. In
+ the target both planes together are the whole pool allocation, so this
+ is also the stray-write check.
+ """
+ for name, got, want in zip(("k_cache", "v_cache"), out, ref):
+ assert got.dtype == want.dtype, (name, got.dtype, want.dtype)
+ assert got.shape == want.shape, (name, got.shape, want.shape)
+ torch.testing.assert_close(
+ _bits(got), _bits(want), rtol=0, atol=0, msg=lambda m, n=name: f"{n}: {m}"
+ )
+
+
+DEV = "cuda"
+FILL = -2.0
+
+
+def _bits(tensor: torch.Tensor) -> torch.Tensor:
+ as_int = {1: torch.uint8, 2: torch.int16, 4: torch.int32}[tensor.element_size()]
+ return tensor.contiguous().view(as_int)
+
+
+def _check(result):
+ err, *value = result
+ assert err == cuda_drv.CUresult.CUDA_SUCCESS, err
+ return value[0] if len(value) == 1 else value
+
+
+class _ArrayInterface:
+ def __init__(self, ptr: int, nbytes: int):
+ self.__cuda_array_interface__ = dict(
+ shape=(nbytes,), strides=None, typestr="|u1", data=(ptr, False), version=2
+ )
+
+
+class _EnginePool:
+ """One attention layer's K and V views of the stock manager's QSA pool at `pages` pages.
+
+ A page holds `role_pages` role planes of `[hkv, tpb, D]` (K then V of every attention
+ layer); the views are layer `layer_slot`'s. The range is reserved whole and physical memory
+ is mapped under the K and V planes of `touched` pages only, and under page 0's (a tensor's
+ base pointer). Reading an unmapped address faults, so only those pages are ever read.
+ """
+
+ def __init__(self, spec, touched):
+ dtype = spec["cache_dtype"]
+ elem = torch.empty((), dtype=dtype).element_size()
+ hkv, tpb, dim = spec["hkv"], spec["tokens_per_block"], spec["head_dim"]
+ self.pages, self.touched = spec["pages"], sorted(set(touched))
+ self.role = hkv * tpb * dim
+ self.stride = spec["role_pages"] * self.role
+ self.own = spec["layer_slot"] * 2 * self.role
+
+ def planes(page):
+ """This layer's K and V planes of `page`, in bytes of the range."""
+ lo = (page * self.stride + self.own) * elem
+ return lo, lo + 2 * self.role * elem
+
+ self.written = [planes(p) for p in self.touched]
+
+ dev = torch.cuda.current_device()
+ prop = cuda_drv.CUmemAllocationProp()
+ prop.type = cuda_drv.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED
+ prop.location.type = cuda_drv.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
+ prop.location.id = dev
+ gran = _check(
+ cuda_drv.cuMemGetAllocationGranularity(
+ prop,
+ cuda_drv.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED,
+ )
+ )
+ self.size = -(-self.pages * self.stride * elem // gran) * gran
+ self.va = int(_check(cuda_drv.cuMemAddressReserve(self.size, gran, 0, 0)))
+ merged = []
+ for lo, hi in sorted({planes(0), *self.written}):
+ lo, hi = lo // gran * gran, min(-(-hi // gran) * gran, self.size)
+ if merged and lo <= merged[-1][1]:
+ merged[-1][1] = max(merged[-1][1], hi)
+ else:
+ merged.append([lo, hi])
+ access = cuda_drv.CUmemAccessDesc()
+ access.location.type = cuda_drv.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
+ access.location.id = dev
+ access.flags = cuda_drv.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE
+ self.maps, self.spans = [], merged
+ for lo, hi in merged:
+ handle = _check(cuda_drv.cuMemCreate(hi - lo, prop, 0))
+ ptr = cuda_drv.CUdeviceptr(self.va + lo)
+ _check(cuda_drv.cuMemMap(ptr, hi - lo, 0, handle, 0))
+ self.maps.append((ptr, hi - lo, handle))
+ _check(cuda_drv.cuMemSetAccess(ptr, hi - lo, [access], 1))
+
+ self.raw = torch.as_tensor(_ArrayInterface(self.va, self.size), device=DEV)
+ pool = self.raw.view(dtype).as_strided(
+ (self.pages, 2, hkv, tpb, dim),
+ (self.stride, self.role, tpb * dim, dim, 1),
+ self.own,
+ )
+ self.k, self.v = pool[:, 0], pool[:, 1]
+ self.index = torch.tensor(self.touched, dtype=torch.long, device=DEV)
+ self.fill()
+
+ def fill(self) -> None:
+ """Every mapped byte to the sentinel's pattern, and remember it."""
+ dtype = self.k.dtype
+ pattern = _bits(torch.full((2,), FILL, dtype=torch.bfloat16).to(dtype))
+ for lo, hi in self.spans:
+ self.raw[lo:hi].view(pattern.dtype).fill_(int(pattern[0]))
+ torch.cuda.synchronize()
+ self.before = [self.raw[lo:hi].clone() for lo, hi in self.spans]
+
+ def dense(self, table: torch.Tensor):
+ """`(k, v, table)`: copies of the touched pages, and `table` renumbered into them.
+
+ An entry naming no touched page becomes `-1`. No row resolves to one, so the
+ reference reads the same pages it would over the whole pool.
+ """
+ remap = torch.full((self.pages,), -1, dtype=torch.int32, device=DEV)
+ remap[self.index] = torch.arange(len(self.touched), dtype=torch.int32, device=DEV)
+ renamed = torch.where(table >= 0, remap[table.clamp_min(0).long()], -1).to(torch.int32)
+ return self.k[self.index].clone(), self.v[self.index].clone(), renamed.contiguous()
+
+ def gathered(self):
+ return self.k[self.index], self.v[self.index]
+
+ def assert_rest_unchanged(self) -> None:
+ """Every mapped byte outside the touched pages' K and V planes is what it was."""
+ torch.cuda.synchronize()
+ for (lo, hi), before in zip(self.spans, self.before):
+ keep = torch.ones(hi - lo, dtype=torch.bool, device=DEV)
+ for p_lo, p_hi in self.written:
+ if p_hi > lo and p_lo < hi:
+ keep[max(p_lo, lo) - lo : min(p_hi, hi) - lo] = False
+ assert torch.equal(self.raw[lo:hi][keep], before[keep]), (
+ f"bytes [{lo}, {hi}) of the pool changed outside the touched pages' planes"
+ )
+
+ def close(self) -> None:
+ torch.cuda.synchronize()
+ del self.k, self.v, self.raw
+ for ptr, size, handle in self.maps:
+ _check(cuda_drv.cuMemUnmap(ptr, size))
+ _check(cuda_drv.cuMemRelease(handle))
+ _check(cuda_drv.cuMemAddressFree(cuda_drv.CUdeviceptr(self.va), self.size))
+
+
+def _hole(spec, table, allocated_cols):
+ """One row that resolves to no slot: its (request, position)."""
+ assert spec["hole"] == "unallocated_page", spec["hole"]
+ assert int(table[0, allocated_cols[0]]) == -1
+ return 0, allocated_cols[0] * spec["tokens_per_block"]
+
+
+def _build(spec, seed: int):
+ """Real pool, table and rows for one cell: `(op kwargs, pool)`.
+
+ `pool` is the contiguous pool tensor, or the `_EnginePool` under a `pool="engine"` cell.
+ """
+ gen = torch.Generator(device="cpu").manual_seed(seed)
+ tpb, hkv, head_dim = spec["tokens_per_block"], spec["hkv"], spec["head_dim"]
+ slots, width, pages = spec["slots"], spec["width"], spec["pages"]
+ engine = spec.get("pool") == "engine"
+ segments = spec["segments"]
+ requests = [r for r, (_, n) in enumerate(segments) for _ in range(n)]
+ positions = [p for start, n in segments for p in range(start, start + n)]
+
+ table = torch.randint(0, pages, (slots, width), generator=gen, dtype=torch.int32)
+ order = torch.randperm(pages, generator=gen).to(torch.int32)
+ if engine:
+ # The first row's page is the pool's last one, whose planes are its highest bytes.
+ first = segments[0][0] // tpb
+ last = int((order == pages - 1).nonzero()[0])
+ order[[first, last]] = order[[last, first]]
+ used, allocated_cols = 0, []
+ for r, (start, n) in enumerate(segments):
+ cols = (start + n - 1) // tpb + 1
+ assert cols <= width, f"request {r} ends past the table's {width} columns"
+ table[r] = -1
+ table[r, :cols] = order[used : used + cols]
+ used += cols
+ allocated_cols.append(cols)
+ assert used <= pages and len(segments) <= slots, "pool or table too small for the batch"
+ if spec["hole"] is not None:
+ request, position = _hole(spec, table, allocated_cols)
+ requests.append(request)
+ positions.append(position)
+
+ rows = len(positions)
+ assert rows == spec.get("rows", rows), (rows, spec["rows"])
+ qkv = torch.randn(
+ rows,
+ spec["heads_q"] + 2 * hkv + spec.get("row_pad_heads", 0),
+ head_dim,
+ generator=gen,
+ dtype=torch.float32,
+ ).to(device=DEV, dtype=spec["row_dtype"])
+ assert qkv.stride(0) == spec.get("row_stride", qkv.stride(0)), qkv.stride(0)
+ k = qkv[:, spec["heads_q"] : spec["heads_q"] + hkv]
+ v = qkv[:, spec["heads_q"] + hkv : spec["heads_q"] + 2 * hkv]
+ touched = [int(table[r, p // tpb]) for r, p in zip(requests, positions)]
+ table = table.to(DEV)
+
+ if engine:
+ pool = _EnginePool(spec, [t for t in touched if t >= 0])
+ k_cache, v_cache = pool.k, pool.v
+ else:
+ pool = torch.full(
+ (pages, 2, hkv, tpb, head_dim), FILL, dtype=torch.bfloat16, device=DEV
+ ).to(spec["cache_dtype"])
+ k_cache, v_cache = pool[:, 0], pool[:, 1]
+ args = dict(
+ k=k,
+ v=v,
+ k_cache=k_cache,
+ v_cache=v_cache,
+ request_indices=torch.tensor(requests, dtype=spec["request_dtype"], device=DEV),
+ logical_positions=torch.tensor(positions, dtype=spec["position_dtype"], device=DEV),
+ block_table=table,
+ tokens_per_block=tpb,
+ )
+ return args, pool
+
+
+def _engine_reference(args, pool: _EnginePool):
+ """The entry's reference over a dense copy of the touched pages."""
+ dense_k, dense_v, renamed = pool.dense(args["block_table"])
+ return op.reference(**{**args, "k_cache": dense_k, "v_cache": dense_v, "block_table": renamed})
+
+
+def _assert_reaches_past_int32(pool: _EnginePool) -> None:
+ """The highest touched page's planes lie past 2^32 elements of the view's base: a page
+ offset computed in 32 bits would land elsewhere."""
+ top = max(pool.touched)
+ assert top == pool.pages - 1, (top, pool.pages)
+ assert top * pool.k.stride(0) >= 2**32, top * pool.k.stride(0)
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying): # noqa: F811
+ certifying(op)
+ args, pool = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
+ if isinstance(pool, _EnginePool):
+ try:
+ _assert_reaches_past_int32(pool)
+ before = pool.gathered()
+ before = (before[0].clone(), before[1].clone())
+ ref = _engine_reference(args, pool)
+ out = op(**args)
+ assert out[0] is args["k_cache"] and out[1] is args["v_cache"]
+ got = pool.gathered()
+ _compare(got, ref)
+ pool.assert_rest_unchanged()
+ wrote = any(not torch.equal(_bits(g), _bits(b)) for g, b in zip(got, before))
+ assert wrote, "a cell with rows must write"
+ print(
+ f"[{cell.why[:44]}] {len(pool.touched)} of {pool.pages} pages touched; "
+ f"top plane at byte {pool.written[-1][0]} of {pool.size}; "
+ f"{sum(hi - lo for lo, hi in pool.spans)} bytes mapped"
+ )
+ finally:
+ pool.close()
+ return
+ before = pool.clone()
+ ref = op.reference(**args)
+ out = op(**args)
+ assert out[0] is args["k_cache"] and out[1] is args["v_cache"]
+ _compare(out, ref)
+ wrote = not torch.equal(_bits(pool), _bits(before))
+ assert wrote == bool(cell.spec["segments"]), (
+ "a cell with rows must write, an empty one must not"
+ )
+
+
+def _decode_args(seed: int, **overrides):
+ """The decode cell's layout over a bf16 pool, without its skipped row."""
+ spec = {**op.CELLS[0].spec, "hole": None}
+ return _build({**spec, **overrides}, seed=seed)
+
+
+#: The cell of 32 MTP7 verify blocks (256 rows): the segments, pool and table it runs over.
+_MTP7_32 = next(c.spec for c in op.CELLS if c.why.startswith("fp8 decode MTP7: 32 requests"))
+_MTP7_BATCH = {key: _MTP7_32[key] for key in ("segments", "pages", "slots")}
+
+
+def test_cuda_graph_replay_follows_updated_positions() -> None:
+ """Captured once at 32 MTP7 verify blocks, replayed against new positions written
+ into the same buffer."""
+ args, pool = _decode_args(seed=9, **_MTP7_BATCH)
+ positions = args["logical_positions"]
+ saved = positions.clone()
+ positions.zero_()
+
+ def call():
+ op(**args)
+
+ for _ in range(3): # compile the Triton kernel before capture
+ call()
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ call()
+
+ pool.fill_(FILL)
+ positions.copy_(saved)
+ ref = op.reference(**args)
+ graph.replay()
+ torch.cuda.synchronize()
+ _compare((args["k_cache"], args["v_cache"]), ref)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py
new file mode 100644
index 000000000000..65066fdaf510
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py
@@ -0,0 +1,325 @@
+# 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 catalog entry.
+
+The paged state is real: one NVFP4 pool as the QSA cache manager serves it -- e2m1
+data `[pages, 2, Hkv, tpb, D/2]` and e4m3 block scales `[pages, 2, Hkv, tpb, D/16]`,
+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.ops.attention.qsa_paged_kv_store_nvfp4 import ( # noqa: E501
+ qsa_paged_kv_store_nvfp4 as op,
+)
+from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401
+
+assert torch.cuda.is_available(), "qsa_paged_kv_store_nvfp4 requires a CUDA device"
+
+DEV = "cuda"
+DATA_FILL, SCALE_FILL = 0xA5, 0x5A
+QUANT_MODE_NVFP4_KV = 8192
+
+
+def _compare(out, ref) -> None:
+ """Bitwise over all four planes: the store is a pure function of the row's
+ bits, and `reference` performs the same rounding steps."""
+ names = ("k_cache", "v_cache", "k_scale_cache", "v_scale_cache")
+ for name, got, want in zip(names, out, ref):
+ assert got.dtype == want.dtype and got.shape == want.shape, (
+ name,
+ got.shape,
+ want.shape,
+ )
+ differ = (got != want).nonzero()
+ if differ.numel():
+ first = tuple(differ[0].tolist())
+ raise AssertionError(
+ f"{name}: {differ.shape[0]} bytes differ from the reference, first at "
+ f"{first}: {int(got[first]):#04x} vs {int(want[first]):#04x}"
+ )
+
+
+def _hole(spec, table, allocated_cols):
+ """One row that resolves to no slot: its (request, position)."""
+ assert spec["hole"] == "unallocated_page", spec["hole"]
+ assert int(table[0, allocated_cols[0]]) == -1
+ return 0, allocated_cols[0] * spec["tokens_per_block"]
+
+
+def _fill_rows(spec, k, v, gen):
+ """Row values: random with per-block magnitudes over 2^-10..2^5, plus edge blocks
+ (zero, -0.0, an outlier) when the cell asks for them."""
+ hkv, head_dim = spec["hkv"], spec["head_dim"]
+ rows = k.shape[0]
+ for t in (k, v):
+ base = torch.randn(rows, hkv, head_dim // 16, 16, generator=gen)
+ exponent = torch.randint(-10, 6, (rows, hkv, head_dim // 16, 1), generator=gen)
+ t.copy_((base * torch.pow(2.0, exponent.float())).reshape(t.shape).to(DEV))
+ if spec["values"] == "edge":
+ k[0, 0, :16] = 0.0
+ v[0, 1, 16:32] = -0.0
+ k[1, 1, 32:48] = torch.tensor([-0.0] * 8 + [0.0] * 8)
+ k[2, 0, 5] = 900.0
+ v[3, 0, 200] = -900.0
+ k[4, 1, 64:80] = 1e-30 # far below the smallest e4m3 scale
+ else:
+ assert spec["values"] == "randn"
+
+
+def _pool(spec):
+ """The allocation the cell's planes view: data and scales of `slot_pages` role
+ pages per slot, every byte its sentinel."""
+ pages, slot_pages = spec["pages"], spec["slot_pages"]
+ hkv, tpb, head_dim = spec["hkv"], spec["tokens_per_block"], spec["head_dim"]
+ data = torch.full(
+ (pages, slot_pages, hkv, tpb, head_dim // 2), DATA_FILL, dtype=torch.uint8, device=DEV
+ )
+ scales = torch.full(
+ (pages, slot_pages, hkv, tpb, head_dim // 16), SCALE_FILL, dtype=torch.uint8, device=DEV
+ )
+ return data, scales
+
+
+def _planes(args):
+ return (args["k_cache"], args["v_cache"], args["k_scale_cache"], args["v_scale_cache"])
+
+
+def _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.get("last_page") is not None:
+ # The named request's last page becomes the pool's last slot; the live entry that held
+ # it, if any, takes the swapped-out page, so the live pages stay distinct.
+ r = spec["last_page"]
+ column = allocated_cols[r] - 1
+ live = table[: len(segments)]
+ holder = (live == pages - 1).nonzero()
+ if holder.numel():
+ live[holder[0, 0], holder[0, 1]] = live[r, column]
+ live[r, column] = pages - 1
+ if spec["hole"] is not None:
+ request, position = _hole(spec, table, allocated_cols)
+ requests.append(request)
+ 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 _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)
+ FallbackFmha.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=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)]
+
+
+def _wrote(planes) -> bool:
+ """Whether any byte of the call's four planes moved off its sentinel."""
+ fills = (DATA_FILL, DATA_FILL, SCALE_FILL, SCALE_FILL)
+ return any(bool((plane != fill).any()) for plane, fill in zip(planes, fills))
+
+
+def _assert_straddles(spec, args) -> None:
+ """The request named by `last_page` stores into the pool's last slot, which lies past
+ 2^31 bytes of both the data and the scale plane: every offset has to be formed in 64 bits."""
+ row = sum(n for _, n in spec["segments"][: spec["last_page"] + 1]) - 1
+ column = int(args["logical_positions"][row]) // spec["tokens_per_block"]
+ request = int(args["request_indices"][row])
+ assert int(args["block_table"][request, column]) == spec["pages"] - 1
+ for name in ("k_cache", "k_scale_cache"):
+ assert (spec["pages"] - 1) * args[name].stride(0) >= 2**31, (name, args[name].stride())
+
+
+# A pool sized to a member's measured page count does not fit a normal cell's 40 GiB.
+_CELL_PARAMS = [
+ pytest.param(c, marks=pytest.mark.large) if c.spec.get("large") else c for c in op.CELLS
+]
+
+
+@pytest.mark.parametrize("cell", _CELL_PARAMS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying): # noqa: F811
+ certifying(op)
+ seed = zlib.crc32(cell.why.encode())
+ if cell.spec["thop_append"]:
+ _thop_equality(cell.spec, seed=seed)
+ return
+ args, pool, _ = _build(cell.spec, seed=seed)
+ if cell.spec.get("last_page") is not None:
+ _assert_straddles(cell.spec, args)
+ ref = op.reference(**args)
+ out = op(**args)
+ assert out[0] is args["k_cache"] and out[3] is args["v_scale_cache"]
+ _compare(out, ref)
+ assert _wrote(_planes(args)) == bool(cell.spec["segments"]), (
+ "a cell with rows must write, an empty one must not"
+ )
+ _assert_other_roles_untouched(cell.spec, pool)
+
+
+def _decode_args(seed: int, **overrides):
+ return _build({**op.CELLS[0].spec, **overrides}, seed=seed)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa.py
new file mode 100644
index 000000000000..91ea1f705957
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa.py
@@ -0,0 +1,353 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the qsa_paged_sparse_gqa catalog entry.
+
+The paged state is real: one `[pages, slot_pages, kv, tpb, D]` HND pool whose
+layer planes are the call's K and V -- tp1's cells one layer's adjacent pair,
+the default target's cells the engine's 24-role-page trunk, sized so the
+requests' pages straddle 2^31 elements -- and the engine's slot table: a row
+per request slot, assigned by a seeded permutation, `-1` past each sequence.
+Selection rows are built the way QSA's top-k + expansion lay them out:
+complete groups expanded to tokens, then the trailing partial group, then
+`-1`, all packed at the front.
+
+The call has two kernels behind it (and a merge). The `launches` fixture
+records which ones a call indexed with a grid, and whether each launched with
+PDL, so every cell asserts its branch (and, where it names `splits`, the
+split count). A cell marked `capture` is also captured in a CUDA graph,
+refilled in place and replayed, as the engine's decode step is.
+"""
+
+import math
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa import ( # noqa: E501
+ qsa_paged_sparse_gqa as op,
+)
+from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels as _qsa_kernels
+
+__extra_import_path__ = ["../../.."]
+from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401
+
+assert torch.cuda.is_available(), "qsa_paged_sparse_gqa requires a CUDA device"
+
+
+def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
+ """4 ulp of each (row, head)'s magnitude, 2 ulp relative RMS, in q's dtype.
+
+ The kernel rounds its online-softmax probabilities to the cache dtype
+ before P.V, so the error tracks the magnitude of the values averaged,
+ not of the output: against an fp32 reference it measured 1.2e-3 on rows
+ averaging a thousand tokens (outputs 0.14 to 0.4) and 1.2e-2 on rows
+ averaging a handful (outputs to 5). Scaled by each head row's own
+ magnitude that is one band. Against this reference (fp32, rounded once)
+ every ungated cell measured at most 2.0 element ulp -- a single final-rounding
+ flip, since bf16's step is 2 of these units at the bottom of a binade --
+ and at most 0.62 ulp RMS: 2x and 3.2x margin. Six wrong references (no
+ scale, interleaved heads, padding attended, identity pages, K/V
+ swapped, page/slot transposed) were each measured more than 0.1 off.
+
+ The default target's cells (6/1 heads at 1 to 8192 rows, 24/2 over an
+ e4m3 pool) sit in the same band, and the table shifted by one page, the
+ softmax scale dropped and one selected token dropped a row each fall
+ outside it.
+
+ A gated output is rounded twice, the attention and then the gated
+ product, each by at most one bf16 ulp of the value rounded, so its bound
+ is near the band's 4 element ulp. The gate not applied, or read one head
+ early, falls far outside it; one selected token dropped at the 1-row
+ gated cell is refused on the RMS gate.
+ """
+ assert out.dtype == ref.dtype, (out.dtype, ref.dtype)
+ assert out.shape == ref.shape, (out.shape, ref.shape)
+ if out.numel() == 0:
+ return
+ assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0)
+
+
+DEV = "cuda"
+
+_KERNELS = {
+ "_qsa_paged_sparse_gqa_kernel": "prefill",
+ "_qsa_paged_sparse_gqa_splitk_kernel": "splitk",
+ "_qsa_merge_splitk_kernel": "merge",
+}
+
+
+class _LaunchRecorder:
+ """Records a Triton kernel's grid indexing and whether it launched with PDL, then
+ delegates to the kernel."""
+
+ def __init__(self, inner, name, log):
+ self._inner, self._name, self._log = inner, name, log
+
+ def __getitem__(self, grid):
+ entry = [self._name, tuple(grid), False]
+ self._log.append(entry)
+ launch = self._inner[grid]
+
+ def run(*args, **kwargs):
+ entry[2] = bool(kwargs.get("launch_pdl", False))
+ return launch(*args, **kwargs)
+
+ return run
+
+
+@pytest.fixture
+def launches(monkeypatch):
+ """Which branch kernels the next call launches, in order: [name, grid, PDL]."""
+ log: list = []
+ for name, tag in _KERNELS.items():
+ inner = getattr(_qsa_kernels, name)
+ monkeypatch.setattr(_qsa_kernels, name, _LaunchRecorder(inner, tag, log))
+ return log
+
+
+# --------------------------------------------------------------------------
+# building real state from a cell spec
+# --------------------------------------------------------------------------
+
+
+def _selected_tokens(positions, seed, compress_ratio, token_topk):
+ """QSA's expansion layout, built in one pass.
+
+ A random `min(available, block_topk)`-subset of the complete visible
+ groups, each expanded to its `compress_ratio` tokens, then the incomplete
+ trailing group, then `-1` -- packed at the front, as the producer emits.
+ """
+ block_topk = token_topk // compress_ratio
+ final_topk = token_topk + compress_ratio - 1
+ gen = torch.Generator(device=DEV).manual_seed(seed)
+ rows = positions.numel()
+ visible = positions.long() + 1
+ available = visible // compress_ratio
+ taken = available.clamp(max=block_topk)
+ width = max(block_topk, int(available.max().item()) if rows else 0)
+
+ scores = torch.rand(rows, width, generator=gen, device=DEV)
+ group_ids = torch.arange(width, device=DEV)[None, :]
+ scores = torch.where(group_ids < available[:, None], scores, torch.inf)
+ chosen = scores.argsort(dim=1)[:, :block_topk]
+
+ column = torch.arange(final_topk, device=DEV)[None, :]
+ group_of_column = (column // compress_ratio).clamp(max=block_topk - 1).expand(rows, -1)
+ body = torch.gather(chosen, 1, group_of_column) * compress_ratio + column % compress_ratio
+ in_body = column < (taken * compress_ratio)[:, None]
+ tail_column = column - (taken * compress_ratio)[:, None]
+ tail = (available * compress_ratio)[:, None] + tail_column
+ in_tail = ~in_body & (tail_column >= 0) & (tail < visible[:, None])
+ selected = torch.where(in_body, body, torch.where(in_tail, tail, torch.full_like(body, -1)))
+ return selected.to(torch.int32)
+
+
+#: NaN's bit pattern per pool dtype: filled through an integer view, since not every fill
+#: kernel takes an e4m3 tensor.
+_NAN_BITS = {torch.bfloat16: (torch.int16, 0x7FC0), torch.float8_e4m3fn: (torch.uint8, 0x7F)}
+
+
+def _nan_pool(shape, dtype):
+ bits_dtype, nan = _NAN_BITS[dtype]
+ return torch.full(shape, nan, dtype=bits_dtype, device=DEV).view(dtype)
+
+
+def _build(spec, seed):
+ """Real paged state for one spec, keyed by `__call__`'s parameter names.
+
+ The pool is `[pool_pages, slot_pages, kv, tpb, D]` and the call's K and V are layer
+ `layer`'s two role planes: tp1's cells keep one layer's adjacent pair (2 role pages a
+ page), the default target's cells the engine's trunk (24). The requests sit on the
+ pool's last pages; every other page, and every other plane, is NaN, so an address that
+ strayed could not come back finite.
+ """
+ gen = torch.Generator(device=DEV).manual_seed(seed)
+ tpb, kv, d, q_heads = (
+ spec["tokens_per_block"],
+ spec["kv_heads"],
+ spec["head_dim"],
+ spec["q_heads"],
+ )
+ segments = list(spec["segments"])
+ positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments])
+ req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)])
+ rows, slots = positions.numel(), spec["slots"]
+ assert len(segments) <= slots
+
+ slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(segments)]
+ mapped = [math.ceil((s + n) / tpb) for s, n in segments]
+ assert max(mapped) <= spec["pages"]
+ needed = sum(mapped) + 7
+ pool_pages = spec.get("pool_pages", needed)
+ assert pool_pages >= needed
+ first = pool_pages - needed
+ page_ids = first + torch.randperm(needed, generator=gen, device=DEV).to(torch.int32)
+ block_table = torch.full((slots, spec["pages"]), -1, dtype=torch.int32, device=DEV)
+ cursor = 0
+ for slot, count in zip(slot_of.tolist(), mapped):
+ block_table[slot, :count] = page_ids[cursor : cursor + count]
+ cursor += count
+
+ slot_pages, layer = spec.get("slot_pages", 2), spec.get("layer", 0)
+ planes = slice(2 * layer, 2 * layer + 2)
+ pool = _nan_pool((pool_pages, slot_pages, kv, tpb, d), spec["cache_dtype"])
+ pool[first:, planes] = torch.randn(needed, 2, kv, tpb, d, generator=gen, device=DEV).to(
+ spec["cache_dtype"]
+ )
+ k_cache, v_cache = pool[:, 2 * layer], pool[:, 2 * layer + 1]
+
+ q_dtype = spec["q_dtype"]
+ fused = spec["q_layout"] == "fused_row"
+ # The row q is read off: the packed `[q | k | v]` buffer, or the target's fused input
+ # projection row `[q | k | v | gate | index q, k | pad]` (`row_heads` heads of D).
+ row_heads = spec.get("row_heads", 55 if fused else q_heads + 2 * kv)
+ row = torch.randn(rows, row_heads * d, generator=gen, device=DEV).to(q_dtype)
+ q = row.view(rows, row_heads, d)[:, :q_heads]
+
+ selected = _selected_tokens(positions, seed + 3, spec["compress_ratio"], spec["token_topk"])
+ request_indices = slot_of.to(DEV)[req].to(spec["request_dtype"])
+
+ gate = None
+ if spec.get("gate") == "fused_row":
+ assert fused, "the gate is read off the fused row"
+ # The gate columns of the same row, after q, k and v.
+ start = (q_heads + 2 * kv) * d
+ gate = row[:, start : start + q_heads * d]
+
+ args = dict(
+ q=q,
+ k_cache=k_cache,
+ v_cache=v_cache,
+ block_table=block_table,
+ selected_tokens=selected,
+ request_indices=request_indices,
+ tokens_per_block=tpb,
+ softmax_scale=d**-0.5,
+ )
+ if spec["prefill_pair"]:
+ args.update(query_positions=positions, compress_ratio=spec["compress_ratio"])
+ if gate is not None:
+ args.update(output_gate=gate)
+ return args
+
+
+def _bound(positions, spec):
+ """The prefill kernel's per-row scan bound, from the positions alone."""
+ cr = spec["compress_ratio"]
+ visible = positions.long() + 1
+ return (visible // cr).clamp(max=spec["token_topk"] // cr) * cr + visible % cr
+
+
+def _assert_straddles(args, edge: int) -> None:
+ """The pages the cell's rows map put the K and V planes' element offsets on both sides
+ of `edge`: what the cell claims about its pool."""
+ table = args["block_table"]
+ pages = table[table >= 0].long()
+ for name in ("k_cache", "v_cache"):
+ plane = args[name]
+ ends = plane.storage_offset() + pages * plane.stride(0) + plane[0].numel() - 1
+ assert bool((ends >= edge).any()), f"{name}: no mapped page lies past {edge}"
+ assert bool((ends < edge).any()), f"{name}: every mapped page lies past {edge}"
+
+
+def _is_prefix_packed(table):
+ valid = table >= 0
+ columns = torch.arange(table.shape[1], device=table.device)[None, :]
+ return (valid == (columns < valid.sum(1)[:, None])).all(1)
+
+
+def _refill(args, spec, seed) -> None:
+ """The engine's next step, written into the call's own buffers: new q and gate values
+ (the fused row's columns), a new selection at the same positions, and new K/V in every
+ mapped page."""
+ gen = torch.Generator(device=DEV).manual_seed(seed)
+ for name in ("q", "output_gate"):
+ if name in args:
+ view = args[name]
+ view.copy_(torch.randn(view.shape, generator=gen, device=DEV).to(view.dtype))
+ positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in spec["segments"]])
+ args["selected_tokens"].copy_(
+ _selected_tokens(positions, seed + 3, spec["compress_ratio"], spec["token_topk"])
+ )
+ table = args["block_table"]
+ pages = table[table >= 0].long()
+ for name in ("k_cache", "v_cache"):
+ plane = args[name]
+ bits = _NAN_BITS[plane.dtype][0]
+ fresh = torch.randn((pages.numel(), *plane.shape[1:]), generator=gen, device=DEV)
+ plane.view(bits)[pages] = fresh.to(plane.dtype).view(bits)
+
+
+def _replay_in_a_graph(args, spec, seed, launches, first) -> None:
+ """The decode site's capture: the call captured into a CUDA graph (its split-K partials
+ allocated from the graph's pool), the inputs refilled in place, the graph replayed.
+ The replay must take the eager call's launches (`first` is that call's output), answer
+ the new inputs within the band, and match an eager call on them bitwise."""
+ eager = [tuple(entry) for entry in launches]
+ del launches[:]
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ captured = op(**args)
+ assert [tuple(entry) for entry in launches] == eager, launches
+ _refill(args, spec, seed)
+ graph.replay()
+ torch.cuda.synchronize()
+ replayed = captured.clone()
+ assert not torch.equal(replayed, first), "the replay did not read the refilled inputs"
+ _compare(replayed, op.reference(**args))
+ assert torch.equal(replayed, op(**args)), "the replay differs from an eager call"
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811
+ certifying(op)
+ spec = cell.spec
+ args = _build(spec, seed=sum(map(ord, cell.why)))
+ rows = args["q"].shape[0]
+ assert rows == spec.get("rows", rows)
+ if "row_heads" in spec:
+ # q, and the gate when fused, are views of one fused row `row_heads` heads wide.
+ q = args["q"]
+ assert q.stride(0) == spec["row_heads"] * spec["head_dim"]
+ if "output_gate" in args:
+ gate = args["output_gate"]
+ assert gate.stride(0) == q.stride(0)
+ assert gate.untyped_storage().data_ptr() == q.untyped_storage().data_ptr()
+ if "rows_per_request" in spec:
+ assert {n for _, n in spec["segments"]} == {spec["rows_per_request"]}
+ if "offset_edge" in spec:
+ _assert_straddles(args, spec["offset_edge"])
+ if spec["prefill_pair"]:
+ # The precondition the op cannot check: front-packed, and exactly as
+ # many valid entries as the prefill kernel will scan.
+ selected = args["selected_tokens"]
+ assert bool(_is_prefix_packed(selected).all())
+ valid = (selected >= 0).sum(1)
+ assert torch.equal(valid, _bound(args["query_positions"], spec))
+
+ out = op(**args)
+ assert [name for name, _, _ in launches] == list(spec["launches"])
+ assert launches[0][1][:2] == (rows, spec["kv_heads"])
+ if "splits" in spec:
+ assert launches[0][1][2] == spec["splits"], launches
+ assert all(pdl == spec.get("pdl", False) for _, _, pdl in launches), launches
+ if "output_gate" in args:
+ gated = (args["q"], args["k_cache"], args["selected_tokens"], args["output_gate"])
+ assert op.fuses_output_gate(*gated)
+ assert out.shape == (rows, spec["q_heads"], spec["head_dim"])
+ assert out.dtype == args["q"].dtype and out.is_contiguous()
+ assert not bool(out.isnan().any())
+ _compare(out, op.reference(**args))
+ if spec.get("capture"):
+ seed = sum(map(ord, cell.why)) + 1
+ _replay_in_a_graph(args, spec, seed=seed, launches=launches, first=out)
+
+
+#: The first cell is this checkpoint's own decode geometry; the tests below
+#: start from it.
+CKPT = op.CELLS[0].spec
+assert (CKPT["q_heads"], CKPT["kv_heads"], CKPT["head_dim"]) == (24, 2, 256)
+
+
+def _with(**overrides):
+ return {**CKPT, **overrides}
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py
new file mode 100644
index 000000000000..7ec9a3535fbc
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py
@@ -0,0 +1,110 @@
+# 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), and the selection's read-after-wait under graph replay with PDL."""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.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.ops.attention.qsa_topk_expand_decode import ( # noqa: E501
+ qsa_topk_expand_decode,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import ( # noqa: E402,F401 — path above; fixture by name
+ assert_within_ulp,
+ certifying,
+ validating,
+)
+
+assert torch.cuda.is_available(), "qsa_paged_sparse_gqa_decode requires a CUDA device"
+
+DEV = "cuda"
+Q_HEADS, KV_HEADS, D, TPB, CR, TOKEN_TOPK, PAGES, SLOTS = 24, 2, 256, 32, 4, 2048, 320, 6
+
+
+def _build(spec, seed):
+ """Paged K/V for one decode row per position, the selection QSA's expansion emits
+ (512 random complete groups, then the incomplete group, then -1), q and the gate as
+ column views of the target's 55-head fused projection row."""
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ pos = torch.tensor(spec["pos"], dtype=torch.int64, device=DEV)
+ rows = pos.numel()
+ pages_per = [int(p) // TPB + 1 for p in spec["pos"]]
+ pool_pages = sum(pages_per) + 5
+ page_ids = torch.randperm(pool_pages, generator=g, device=DEV).to(torch.int32)
+ block_table = torch.full((SLOTS, PAGES), -1, dtype=torch.int32, device=DEV)
+ slot_of = list(range(1, rows + 1))
+ cursor = 0
+ for slot, count in zip(slot_of, pages_per):
+ block_table[slot, :count] = page_ids[cursor : cursor + count]
+ cursor += count
+ pool = torch.randn(pool_pages, 2, KV_HEADS, TPB, D, generator=g, device=DEV).to(torch.bfloat16)
+ k_cache, v_cache = pool[:, 0], pool[:, 1]
+ fused = torch.randn(rows, 55 * D, generator=g, device=DEV).to(torch.bfloat16)
+ q = fused.view(rows, 55, D)[:, :Q_HEADS]
+ gate = fused[:, 28 * D : (28 + Q_HEADS) * D] if spec["gated"] else None
+ block_topk = TOKEN_TOPK // CR
+ selected = torch.full((rows, TOKEN_TOPK + CR - 1), -1, dtype=torch.int32, device=DEV)
+ for r, p in enumerate(spec["pos"]):
+ visible = p + 1
+ groups = visible // CR
+ chosen = torch.randperm(groups, generator=g, device=DEV)[: min(groups, block_topk)]
+ toks = (chosen[:, None] * CR + torch.arange(CR, device=DEV)).reshape(-1)
+ tail = torch.arange(groups * CR, visible, device=DEV)
+ row = torch.cat([toks, tail]).to(torch.int32)
+ selected[r, : row.numel()] = row
+ if spec.get("holes"):
+ selected[r, 1 : row.numel() : 2] = -1
+ request_indices = torch.tensor(slot_of, dtype=torch.int32, device=DEV)
+ return [q, k_cache, v_cache, block_table, selected, request_indices, TPB, D**-0.5, gate]
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
+ certifying(op)
+ args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31))
+ out = op(*args)
+ ref = op.reference(*args)
+ assert out.dtype == torch.bfloat16 and out.shape == ref.shape and out.is_contiguous()
+ assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0)
+
+
+def _replays_like_eager(step, refresh) -> None:
+ warmup = torch.cuda.Stream()
+ warmup.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(warmup):
+ step()
+ torch.cuda.current_stream().wait_stream(warmup)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ graph_out = step()
+ for trial in range(3):
+ refresh()
+ graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(graph_out, step()), f"replay {trial}"
+
+
+def test_reads_the_selection_after_the_wait_under_graph_replay() -> None:
+ """The target's top-k runs inside the graph right before the op: a PDL producer that
+ releases the op right after its own wait, before it writes the selection. A read of the
+ selection before the op's wait would see the previous replay's."""
+ args = _build(dict(pos=(8700,), gated=True), seed=23)
+ pos = torch.tensor([8700], dtype=torch.int64, device=DEV)
+ ends = ((pos + 1) // CR).to(torch.int32)
+ seqlen = (pos + 1).to(torch.int32)
+ scores = torch.rand(1, 2304, device=DEV)
+ topk = torch.empty(1, TOKEN_TOPK // CR, dtype=torch.int32, device=DEV)
+
+ def step():
+ sel = qsa_topk_expand_decode(scores, ends, pos, seqlen, topk, CR, TOKEN_TOPK)
+ return op(*args[:4], sel, *args[5:])
+
+ _replays_like_eager(step, lambda: scores.copy_(torch.rand(scores.shape, device=DEV)))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py
new file mode 100644
index 000000000000..f7db6f517d43
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py
@@ -0,0 +1,332 @@
+# 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 catalog entry.
+
+The paged state is real: one NVFP4 pool as the QSA cache manager serves it --
+e2m1 data `[pages, 2, kv, tpb, D/2]` and e4m3 block scales `[pages, 2, kv, tpb,
+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 pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention import (
+ qsa_paged_sparse_gqa_nvfp4 as _entry,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4 import ( # noqa: E501
+ qsa_paged_sparse_gqa_nvfp4 as op,
+)
+from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import triton_qsa_paged_sparse_gqa
+
+__extra_import_path__ = ["../../.."]
+from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401
+
+assert torch.cuda.is_available(), "qsa_paged_sparse_gqa_nvfp4 requires a CUDA device"
+
+DEV = "cuda"
+ENGINE_MAX_SEQ_LEN = 8192
+#: The e4m3fn NaN bit pattern.
+E4M3_NAN = 0x7F
+
+_KERNELS = {
+ "_nvfp4_sparse_gqa_prefill_kernel": "prefill",
+ "_nvfp4_sparse_gqa_splitk_kernel": "splitk",
+ "_nvfp4_merge_splitk_kernel": "merge",
+}
+
+
+class _LaunchRecorder:
+ """Records a Triton kernel's grid indexing, then delegates to the kernel."""
+
+ def __init__(self, inner, name, log):
+ self._inner, self._name, self._log = inner, name, log
+
+ def __getitem__(self, grid):
+ self._log.append((self._name, tuple(grid)))
+ return self._inner[grid]
+
+
+@pytest.fixture
+def launches(monkeypatch):
+ """Which branch kernels the next call launches, in order, with grids."""
+ log: list = []
+ for name, tag in _KERNELS.items():
+ monkeypatch.setattr(_entry, name, _LaunchRecorder(getattr(_entry, name), tag, log))
+ return log
+
+
+def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
+ """4 ulp of each (row, head)'s magnitude, 2 ulp relative RMS, in q's dtype.
+
+ After the dequantize the kernels are the bf16 sparse GQA kernels, whose band
+ this is (they round the online-softmax probabilities to the value dtype
+ before P.V); the dequantize itself is exact at unit scale and rounded once,
+ as `reference` rounds it, otherwise.
+ """
+ assert out.dtype == ref.dtype, (out.dtype, ref.dtype)
+ assert out.shape == ref.shape, (out.shape, ref.shape)
+ if out.numel() == 0:
+ return
+ assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0)
+
+
+def _selected_tokens(positions, seed, compress_ratio, token_topk):
+ """QSA's expansion layout: chosen complete groups, the trailing partial group, -1."""
+ block_topk = token_topk // compress_ratio
+ final_topk = token_topk + compress_ratio - 1
+ gen = torch.Generator(device=DEV).manual_seed(seed)
+ rows = positions.numel()
+ visible = positions.long() + 1
+ available = visible // compress_ratio
+ taken = available.clamp(max=block_topk)
+ width = max(block_topk, int(available.max().item()) if rows else 0)
+ scores = torch.rand(rows, width, generator=gen, device=DEV)
+ group_ids = torch.arange(width, device=DEV)[None, :]
+ scores = torch.where(group_ids < available[:, None], scores, torch.inf)
+ chosen = scores.argsort(dim=1)[:, :block_topk]
+ column = torch.arange(final_topk, device=DEV)[None, :]
+ group_of_column = (column // compress_ratio).clamp(max=block_topk - 1).expand(rows, -1)
+ body = torch.gather(chosen, 1, group_of_column) * compress_ratio + column % compress_ratio
+ in_body = column < (taken * compress_ratio)[:, None]
+ tail_column = column - (taken * compress_ratio)[:, None]
+ tail = (available * compress_ratio)[:, None] + tail_column
+ in_tail = ~in_body & (tail_column >= 0) & (tail < visible[:, None])
+ selected = torch.where(in_body, body, torch.where(in_tail, tail, torch.full_like(body, -1)))
+ return selected.to(torch.int32)
+
+
+def _random_pool(pages, kv, tpb, d, gen):
+ """Pool bytes: any e2m1 nibble, and positive e4m3 scales from 2^-6 to 0.5."""
+ data = torch.randint(0, 256, (pages, 2, kv, tpb, d // 2), generator=gen, device=DEV)
+ codes = torch.randint(0x08, 0x31, (pages, 2, kv, tpb, d // 16), generator=gen, device=DEV)
+ return data.to(torch.uint8), codes.to(torch.uint8)
+
+
+def _build(spec, seed):
+ """Real paged state for one spec, keyed by the op's parameter names."""
+ gen = torch.Generator(device=DEV).manual_seed(seed)
+ tpb, kv, d, q_heads = (spec[k] for k in ("tokens_per_block", "kv_heads", "head_dim", "q_heads"))
+ segments = list(spec["segments"])
+ positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments])
+ req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)])
+ rows, slots = positions.numel(), spec["slots"]
+ assert len(segments) <= slots
+ slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(segments)]
+ mapped = [math.ceil((s + n) / tpb) for s, n in segments]
+ assert not mapped or max(mapped) <= spec["pages"]
+ needed = sum(mapped) + 7
+ slot_pages, layer = spec["slot_pages"], spec["layer"]
+ planes = slice(2 * layer, 2 * layer + 2)
+ # The requests sit on the pool's last `needed` pages, and every scale outside them is
+ # NaN, so a page address that strayed could not come back finite.
+ pool_pages = spec.get("pool_pages", needed)
+ assert pool_pages >= needed
+ first = pool_pages - needed
+ page_ids = first + torch.randperm(needed, generator=gen, device=DEV).to(torch.int32)
+ data = torch.zeros(pool_pages, slot_pages, kv, tpb, d // 2, dtype=torch.uint8, device=DEV)
+ scales = torch.full(
+ (pool_pages, slot_pages, kv, tpb, d // 16), E4M3_NAN, dtype=torch.uint8, device=DEV
+ )
+ data[first:, planes], scales[first:, planes] = _random_pool(needed, kv, tpb, d, gen)
+ block_table = torch.full((slots, spec["pages"]), -1, dtype=torch.int32, device=DEV)
+ cursor = 0
+ for slot, count in zip(slot_of.tolist(), mapped):
+ block_table[slot, :count] = page_ids[cursor : cursor + count]
+ cursor += count
+
+ assert spec["q_layout"] == "qkv_slice", spec["q_layout"]
+ qkv = torch.randn(rows, q_heads + 2 * kv, d, generator=gen, device=DEV).to(spec["q_dtype"])
+ q = qkv[:, :q_heads]
+ selected = _selected_tokens(positions, seed + 3, spec["compress_ratio"], spec["token_topk"])
+ request_indices = slot_of.to(DEV)[req].to(spec["request_dtype"])
+
+ args = dict(
+ q=q,
+ k_cache=data[:, 2 * layer],
+ v_cache=data[:, 2 * layer + 1],
+ k_scale_cache=scales[:, 2 * layer],
+ v_scale_cache=scales[:, 2 * layer + 1],
+ kv_scale_quant_orig=torch.tensor(
+ [1.0, spec["global_k"], spec["global_v"]], dtype=torch.float32, device=DEV
+ ),
+ block_table=block_table,
+ selected_tokens=selected,
+ request_indices=request_indices,
+ tokens_per_block=tpb,
+ softmax_scale=d**-0.5,
+ )
+ if spec["prefill_pair"]:
+ args.update(query_positions=positions, compress_ratio=spec["compress_ratio"])
+ return args
+
+
+def _bf16_twin(args):
+ """QSA's bf16 sparse GQA over a bf16 pool holding the dequantized values.
+
+ The twin pool holds only the pages the table maps, renumbered in order, so a cell's
+ pool can reach past 2^31 bytes without the twin dequantizing all of it.
+ """
+ table = args["block_table"]
+ mapped = torch.unique(table[table >= 0]).long()
+ renumber = torch.full((args["k_cache"].shape[0],), -1, dtype=torch.int32, device=DEV)
+ renumber[mapped] = torch.arange(mapped.numel(), dtype=torch.int32, device=DEV)
+ twin_table = torch.where(table >= 0, renumber[table.clamp_min(0).long()], table)
+ tpb = args["tokens_per_block"]
+ pages = mapped[:, None].expand(-1, tpb)
+ slots = torch.arange(tpb, device=DEV)[None, :].expand(mapped.numel(), -1)
+ planes = []
+ for data, scales, index, grouped in (
+ (args["k_cache"], args["k_scale_cache"], 1, False),
+ (args["v_cache"], args["v_scale_cache"], 2, True),
+ ):
+ values = _entry._dequantize_tokens(
+ data, scales, pages, slots, args["kv_scale_quant_orig"][index], grouped
+ ) # [mapped pages, tpb, H, D]
+ planes.append(values.to(args["q"].dtype).transpose(1, 2).contiguous())
+ twin = {
+ name: args[name] for name in ("q", "selected_tokens", "request_indices", "tokens_per_block")
+ }
+ twin.update(
+ softmax_scale=args["softmax_scale"],
+ block_table=twin_table,
+ k_cache=planes[0],
+ v_cache=planes[1],
+ )
+ 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 _assert_straddles(args, edge: int) -> None:
+ """The pages the cell's rows read put the scale plane's byte offsets on both sides of
+ `edge` and the data plane's past it: what the cell claims about its pool."""
+ table = args["block_table"]
+ pages = table[table >= 0].long()
+ for name, side in (("k_scale_cache", "both"), ("v_scale_cache", "both"), ("k_cache", "past")):
+ plane = args[name]
+ base = plane.data_ptr() - plane.untyped_storage().data_ptr()
+ ends = base + pages * plane.stride(0) + plane[0].numel() - 1
+ assert bool((ends >= edge).any()), f"{name}: no page the rows read lies past {edge}"
+ if side == "both":
+ assert bool((ends < edge).any()), f"{name}: every page lies past {edge}"
+
+
+def _is_prefix_packed(table):
+ valid = table >= 0
+ columns = torch.arange(table.shape[1], device=table.device)[None, :]
+ return (valid == (columns < valid.sum(1)[:, None])).all(1)
+
+
+_OPERANDS = (
+ "q",
+ "k_cache",
+ "v_cache",
+ "k_scale_cache",
+ "v_scale_cache",
+ "block_table",
+ "selected_tokens",
+ "request_indices",
+)
+
+
+def _capture_sweep(spec, seed, launches):
+ """The decode site under a CUDA graph, at each G for each rows-per-request:
+ capture, write a second state's operands in place, replay, and require the
+ eager call on that second state bitwise."""
+ sweep = spec["capture"]
+ for per_request in sweep["rows_per_request"]:
+ for groups in sweep["groups"]:
+ _capture_once(
+ _decode_step(spec, groups, per_request),
+ seed + 100 * per_request + groups,
+ launches,
+ where=f"G={groups} x {per_request} rows",
+ )
+
+
+def _start(request: int, rows: int) -> int:
+ """A request's first position, across the budget."""
+ return min(100 + (499 * request) % 8000, ENGINE_MAX_SEQ_LEN - rows)
+
+
+def _decode_step(spec, groups, per_request):
+ """The decode site's step: `groups` requests x `per_request` rows."""
+ segments = tuple((_start(g, per_request), per_request) for g in range(groups))
+ return {**spec, "segments": segments}
+
+
+def _capture_once(shape, seed, launches, where) -> None:
+ """One call of `shape` captured, a second state's operands written in place,
+ replayed: the replay must be bitwise the eager call on the second state."""
+ where = f"{where}, [{shape['slots']}, {shape['pages']}] table"
+ args = _build(shape, seed)
+ fresh = _build(shape, seed + 50)
+ launches.clear()
+ op(**args)
+ assert [name for name, _ in launches] == list(shape["launches"]), where
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ try:
+ with torch.cuda.graph(graph):
+ captured = op(**args)
+ for name in _OPERANDS:
+ args[name].copy_(fresh[name])
+ graph.replay()
+ torch.cuda.synchronize()
+ eager = op(**fresh)
+ assert torch.equal(captured, eager), f"replay differs from eager at {where}"
+ _compare(captured, op.reference(**fresh))
+ finally:
+ del graph
+ torch.cuda.synchronize()
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811
+ certifying(op)
+ spec = cell.spec
+ if "capture" in spec:
+ _capture_sweep(spec, seed=sum(map(ord, cell.why)), launches=launches)
+ return
+ args = _build(spec, seed=sum(map(ord, cell.why)))
+ rows = args["q"].shape[0]
+ if spec["prefill_pair"]:
+ selected = args["selected_tokens"]
+ assert bool(_is_prefix_packed(selected).all())
+ valid = (selected >= 0).sum(1)
+ assert torch.equal(valid, _bound(args["query_positions"], spec))
+ if "offset_edge" in spec:
+ _assert_straddles(args, spec["offset_edge"])
+
+ 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}
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_decode.py
new file mode 100644
index 000000000000..93d998a93d4e
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_decode.py
@@ -0,0 +1,268 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the `qsa_paged_sparse_gqa_nvfp4_decode` catalog entry: against the reference of
+`qsa_paged_sparse_gqa_nvfp4`
+and QSA's bf16 kernel over the dequantized pool, inside the NVFP4 op's band (4 ulp of each
+(row, head)'s max, 2 ulp relative RMS), over its cells with and without the output gate; plus
+an unaligned row slice of the selection.
+
+The tp4-heads cells run one tp4 / tep4 rank's 6 q heads over 1 KV head, q a view of the packed
+`[q | k | v]` rows as the target hands it. The dep4 cells run at one rank's pool as the engine
+allocates it under `kv_nvfp4`: the requests' pages scattered over all 211797 pages, both planes'
+byte offsets on both sides of 2^31 and 2^32 (asserted), every scale outside the mapped pages
+NaN, and each captured and replayed at its own compile variant. The bf16 twin dequantizes only
+the mapped pages."""
+
+import math
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention import (
+ qsa_paged_sparse_gqa_nvfp4 as _nvfp4_module,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4_decode import ( # noqa: E501
+ qsa_paged_sparse_gqa_nvfp4_decode as op,
+)
+from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import triton_qsa_paged_sparse_gqa
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "qsa_paged_sparse_gqa_nvfp4_decode requires a CUDA device"
+
+DEV = "cuda"
+E4M3_NAN = 0x7F
+ULP = 2.0**-7 # bf16
+
+
+def within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0):
+ o = out.float().reshape(-1, out.shape[-1])
+ r = ref.float().reshape(-1, ref.shape[-1])
+ scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9)
+ el = ((o - r).abs() / scale).max().item() / ULP
+ rms = (((o - r) ** 2).sum() / (r**2).sum().clamp_min(1e-30)).sqrt().item() / ULP
+ return el, rms, el <= element_ulp and rms <= rms_ulp
+
+
+def _selected_tokens(positions, seed, compress_ratio, token_topk):
+ block_topk = token_topk // compress_ratio
+ final_topk = token_topk + compress_ratio - 1
+ gen = torch.Generator(device=DEV).manual_seed(seed)
+ rows = positions.numel()
+ visible = positions.long() + 1
+ available = visible // compress_ratio
+ taken = available.clamp(max=block_topk)
+ width = max(block_topk, int(available.max().item()) if rows else 0)
+ scores = torch.rand(rows, width, generator=gen, device=DEV)
+ group_ids = torch.arange(width, device=DEV)[None, :]
+ scores = torch.where(group_ids < available[:, None], scores, torch.inf)
+ chosen = scores.argsort(dim=1)[:, :block_topk]
+ column = torch.arange(final_topk, device=DEV)[None, :]
+ group_of_column = (column // compress_ratio).clamp(max=block_topk - 1).expand(rows, -1)
+ body = torch.gather(chosen, 1, group_of_column) * compress_ratio + column % compress_ratio
+ in_body = column < (taken * compress_ratio)[:, None]
+ tail_column = column - (taken * compress_ratio)[:, None]
+ tail = (available * compress_ratio)[:, None] + tail_column
+ in_tail = ~in_body & (tail_column >= 0) & (tail < visible[:, None])
+ selected = torch.where(in_body, body, torch.where(in_tail, tail, torch.full_like(body, -1)))
+ return selected.to(torch.int32)
+
+
+def build(
+ segments,
+ seed,
+ slots=32,
+ pages=260,
+ slot_pages=2,
+ layer=0,
+ global_k=1.0,
+ global_v=1.0,
+ request_dtype=torch.int32,
+ pool_pages=None,
+ spread=False,
+ q_heads=24,
+ kv_heads=2,
+):
+ gen = torch.Generator(device=DEV).manual_seed(seed)
+ tpb, kv, d = 32, kv_heads, 256
+ positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments])
+ req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)])
+ rows = positions.numel()
+ slots = max(slots, len(segments))
+ slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(segments)]
+ mapped = [math.ceil((s + n) / tpb) for s, n in segments]
+ assert max(mapped) <= pages, (mapped, pages)
+ needed = sum(mapped) + 7
+ planes = slice(2 * layer, 2 * layer + 2)
+ pool_pages = pool_pages or needed
+ data = torch.zeros(pool_pages, slot_pages, kv, tpb, d // 2, dtype=torch.uint8, device=DEV)
+ # Every scale outside the requests' pages is NaN, so a page address that strayed could
+ # not come back finite.
+ scales = torch.full(
+ (pool_pages, slot_pages, kv, tpb, d // 16), E4M3_NAN, dtype=torch.uint8, device=DEV
+ )
+ if spread:
+ # The requests' pages scattered over the whole pool; the pool's last page is the first
+ # request's first page, so it is always mapped.
+ picked = torch.randperm(pool_pages - 1, generator=gen, device=DEV)[: needed - 1]
+ last = torch.full((1,), pool_pages - 1, device=DEV)
+ page_ids = torch.cat([last, picked]).to(torch.int32)
+ resident = page_ids.long()
+ else:
+ first = pool_pages - needed
+ page_ids = first + torch.randperm(needed, generator=gen, device=DEV).to(torch.int32)
+ resident = slice(first, None)
+ data[resident, planes] = torch.randint(
+ 0, 256, (needed, 2, kv, tpb, d // 2), generator=gen, device=DEV
+ ).to(torch.uint8)
+ scales[resident, planes] = torch.randint(
+ 0x08, 0x31, (needed, 2, kv, tpb, d // 16), generator=gen, device=DEV
+ ).to(torch.uint8)
+ block_table = torch.full((slots, pages), -1, dtype=torch.int32, device=DEV)
+ cursor = 0
+ for slot, count in zip(slot_of.tolist(), mapped):
+ block_table[slot, :count] = page_ids[cursor : cursor + count]
+ cursor += count
+ qkv = torch.randn(rows, q_heads + 2 * kv, d, generator=gen, device=DEV).to(torch.bfloat16)
+ selected = _selected_tokens(positions, seed + 3, 4, 2048)
+ return dict(
+ q=qkv[:, :q_heads],
+ 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, global_k, global_v], dtype=torch.float32, device=DEV
+ ),
+ block_table=block_table,
+ selected_tokens=selected,
+ request_indices=slot_of.to(DEV)[req].to(request_dtype),
+ tokens_per_block=tpb,
+ softmax_scale=d**-0.5,
+ )
+
+
+def bf16_twin(args):
+ """QSA's bf16 sparse GQA over a bf16 pool holding the dequantized values of the pages the
+ table maps, renumbered in order, so a cell's pool can reach past 2^32 bytes without the
+ twin dequantizing all of it."""
+ table = args["block_table"]
+ mapped = torch.unique(table[table >= 0]).long()
+ renumber = torch.full((args["k_cache"].shape[0],), -1, dtype=torch.int32, device=DEV)
+ renumber[mapped] = torch.arange(mapped.numel(), dtype=torch.int32, device=DEV)
+ twin_table = torch.where(table >= 0, renumber[table.clamp_min(0).long()], table)
+ tpb = args["tokens_per_block"]
+ pages = mapped[:, None].expand(-1, tpb)
+ slots = torch.arange(tpb, device=DEV)[None, :].expand(mapped.numel(), -1)
+ planes = []
+ for data, scales, index, grouped in (
+ (args["k_cache"], args["k_scale_cache"], 1, False),
+ (args["v_cache"], args["v_scale_cache"], 2, True),
+ ):
+ values = _nvfp4_module._dequantize_tokens(
+ data, scales, pages, slots, args["kv_scale_quant_orig"][index], grouped
+ )
+ planes.append(values.to(args["q"].dtype).transpose(1, 2).contiguous())
+ twin = {n: args[n] for n in ("q", "selected_tokens", "request_indices", "tokens_per_block")}
+ twin.update(
+ softmax_scale=args["softmax_scale"],
+ block_table=twin_table,
+ k_cache=planes[0],
+ v_cache=planes[1],
+ )
+ return triton_qsa_paged_sparse_gqa(**twin)
+
+
+def _assert_straddles(args, edges) -> None:
+ """The pages the cell's rows read put both planes' byte offsets on both sides of every
+ edge: what the cell claims about its pool."""
+ table = args["block_table"]
+ pages = table[table >= 0].long()
+ for name in ("k_cache", "v_cache", "k_scale_cache", "v_scale_cache"):
+ plane = args[name]
+ base = plane.data_ptr() - plane.untyped_storage().data_ptr()
+ ends = base + pages * plane.stride(0) + plane[0].numel() - 1
+ for edge in edges:
+ assert bool((ends >= edge).any()), f"{name}: no page the rows read lies past {edge}"
+ assert bool((ends < edge).any()), f"{name}: every page the rows read lies past {edge}"
+ last = args["k_cache"].shape[0] - 1
+ assert bool((pages == last).any()), "the pool's last page is not among the rows' pages"
+
+
+#: Spec keys that describe the cell rather than the state `build` makes.
+_CELL_ONLY = ("segments", "gated", "capture", "large", "offset_edges")
+
+
+def _args(spec, seed):
+ kw = {k: v for k, v in spec.items() if k not in _CELL_ONLY}
+ args = build(spec["segments"], seed=seed, **kw)
+ if spec.get("gated"):
+ rows = args["q"].shape[0]
+ g = torch.Generator(device=DEV).manual_seed(seed + 1)
+ width = args["q"].shape[1] * args["q"].shape[2]
+ args["output_gate"] = torch.randn(rows, width, generator=g, device=DEV).to(torch.bfloat16)
+ return args
+
+
+def _replay_in_place(args, seed) -> None:
+ """The call captured at the cell's own compile variant, fresh q (and gate) written into
+ the captured buffers, replayed: bitwise the eager call on the fresh values."""
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ try:
+ with torch.cuda.graph(graph):
+ captured = op(**args)
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ for name in ("q", "output_gate"):
+ if name in args:
+ args[name].copy_(torch.randn(args[name].shape, generator=g, device=DEV))
+ graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(captured, op(**args)), "replay differs from eager"
+ finally:
+ del graph
+ torch.cuda.synchronize()
+
+
+def _marks(cell):
+ return (pytest.mark.large, pytest.mark.high_cuda_memory) if cell.spec.get("large") else ()
+
+
+@pytest.mark.parametrize(
+ "cell", [pytest.param(c, marks=_marks(c), id=c.why[:44]) for c in op.CELLS]
+)
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture
+ certifying(op)
+ seed = zlib.crc32(cell.why.encode()) % (2**31)
+ args = _args(cell.spec, seed)
+ if "offset_edges" in cell.spec:
+ _assert_straddles(args, cell.spec["offset_edges"])
+ out = op(**args)
+ assert out.dtype == torch.bfloat16 and not out.isnan().any()
+ el, rms, ok = within_ulp(out, op.reference(**args))
+ assert ok, f"vs reference: {el:.2f} ulp element, {rms:.2f} ulp RMS"
+ # QSA's bf16 kernel has no output gate, so a gated cell checks its ungated call over the
+ # same state against it.
+ gate = args.pop("output_gate", None)
+ plain = out if gate is None else op(**args)
+ el, rms, ok = within_ulp(plain, bf16_twin(args))
+ assert ok, f"vs QSA bf16 kernel: {el:.2f} ulp element, {rms:.2f} ulp RMS"
+ if gate is not None:
+ args["output_gate"] = gate
+ if cell.spec.get("capture"):
+ _replay_in_place(args, seed + 7)
+
+
+def test_unaligned_row_slice_of_the_selection() -> None:
+ """A mixed step hands its generation rows as a row slice of the step's 2051-wide
+ selection, 16-byte aligned only at every fourth row: an unaligned slice gives the same
+ rows as an aligned copy."""
+ args = build(tuple((3000 + 211 * i, 1) for i in range(9)), seed=4242)
+ aligned = dict(args)
+ for name in ("q", "selected_tokens", "request_indices"):
+ args[name] = args[name][1:]
+ aligned[name] = args[name].clone()
+ assert args["selected_tokens"].data_ptr() % 16
+ assert torch.equal(op(**args), op(**aligned))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py
new file mode 100644
index 000000000000..d13d08a9e60a
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py
@@ -0,0 +1,268 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test: `qsa_paged_sparse_gqa_nvfp4_verify` against the reference of
+`qsa_paged_sparse_gqa_nvfp4_decode` and against the per-row kernel, inside the NVFP4 op's band
+(4 ulp of each (row, head)'s max, 2 ulp relative RMS), over its cells; plus the routing to the
+per-row kernel, forced splits and other rows a request, the launch plan (a partial last wave's
+units split and merged), invalid request rows, and an unaligned row slice of the selection."""
+
+import math
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention import ( # noqa: E501
+ qsa_paged_sparse_gqa_nvfp4_verify as _verify_module,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4_decode import ( # noqa: E501
+ nvfp4_sparse_gqa_decode as per_row,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4_verify import ( # noqa: E501
+ qsa_paged_sparse_gqa_nvfp4_verify as op,
+)
+
+assert torch.cuda.is_available(), "qsa_paged_sparse_gqa_nvfp4_verify requires a CUDA device"
+
+DEV = "cuda"
+E4M3_NAN = 0x7F
+ULP = 2.0**-7 # bf16
+COMPRESS_RATIO, TOKEN_TOPK = 4, 2048
+WIDTH = TOKEN_TOPK + COMPRESS_RATIO - 1
+
+
+def within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0):
+ o = out.float().reshape(-1, out.shape[-1])
+ r = ref.float().reshape(-1, ref.shape[-1])
+ scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9)
+ el = ((o - r).abs() / scale).max().item() / ULP
+ rms = (((o - r) ** 2).sum() / (r**2).sum().clamp_min(1e-30)).sqrt().item() / ULP
+ return el, rms, el <= element_ulp and rms <= rms_ulp
+
+
+def _expand(blocks, taken, available, visible):
+ """Block ids `[rows, 512]` (the first `taken` valid) to the 2051-wide token selection:
+ each block's four tokens, then the causal tail, then -1."""
+ column = torch.arange(WIDTH, device=DEV)[None, :]
+ rows = blocks.shape[0]
+ block_of_column = (column // COMPRESS_RATIO).clamp(max=blocks.shape[1] - 1).expand(rows, -1)
+ body = torch.gather(blocks, 1, block_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 _selected_tokens(positions, requests, seed, overlap):
+ """Each row's top 512 visible four-token blocks, expanded. `overlap` None: independent
+ random scores a row, in random order (the decode op's test); else the request's scores plus
+ `overlap` x N(0, 1) a row, blocks ascending (as `qsa_topk_expand_decode` writes them)."""
+ block_topk = TOKEN_TOPK // COMPRESS_RATIO
+ 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)
+ block_ids = torch.arange(width, device=DEV)[None, :]
+ if overlap is None:
+ scores = torch.rand(rows, width, generator=gen, device=DEV)
+ scores = torch.where(block_ids < available[:, None], scores, torch.inf)
+ return _expand(scores.argsort(dim=1)[:, :block_topk], taken, available, visible)
+ base = torch.rand(int(requests.max().item()) + 1, width, generator=gen, device=DEV)
+ noise = torch.randn(rows, width, generator=gen, device=DEV)
+ scores = base[requests] + overlap * noise
+ scores = torch.where(block_ids < available[:, None], scores, -torch.inf)
+ top = scores.topk(block_topk, dim=1).indices
+ valid = torch.arange(block_topk, device=DEV)[None, :] < taken[:, None]
+ blocks = torch.where(valid, top, torch.full_like(top, 1 << 30)).sort(dim=1).values
+ return _expand(blocks, taken, available, visible)
+
+
+def build(
+ segments,
+ seed,
+ overlap=0.05,
+ slots=32,
+ pages=260,
+ slot_pages=2,
+ layer=0,
+ global_k=1.0,
+ global_v=1.0,
+ request_dtype=torch.int32,
+ pool_pages=None,
+):
+ """The decode op's test state for `segments` ((first position, rows) per request, in row
+ order) with `_selected_tokens(overlap)`."""
+ gen = torch.Generator(device=DEV).manual_seed(seed)
+ tpb, kv, d, q_heads = 32, 2, 256, 24
+ positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments])
+ req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)])
+ rows = positions.numel()
+ slots = max(slots, len(segments))
+ slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(segments)]
+ mapped = [math.ceil((s + n) / tpb) for s, n in segments]
+ needed = sum(mapped) + 7
+ planes = slice(2 * layer, 2 * layer + 2)
+ pool_pages = pool_pages or needed
+ first = pool_pages - needed
+ page_ids = first + torch.randperm(needed, generator=gen, device=DEV).to(torch.int32)
+ data = torch.zeros(pool_pages, slot_pages, kv, tpb, d // 2, dtype=torch.uint8, device=DEV)
+ scales = torch.full(
+ (pool_pages, slot_pages, kv, tpb, d // 16), E4M3_NAN, dtype=torch.uint8, device=DEV
+ )
+ data[first:, planes] = torch.randint(
+ 0, 256, (needed, 2, kv, tpb, d // 2), generator=gen, device=DEV
+ ).to(torch.uint8)
+ scales[first:, planes] = torch.randint(
+ 0x08, 0x31, (needed, 2, kv, tpb, d // 16), generator=gen, device=DEV
+ ).to(torch.uint8)
+ block_table = torch.full((slots, pages), -1, dtype=torch.int32, device=DEV)
+ cursor = 0
+ for slot, count in zip(slot_of.tolist(), mapped):
+ block_table[slot, :count] = page_ids[cursor : cursor + count]
+ cursor += count
+ qkv = torch.randn(rows, q_heads + 2 * kv, d, generator=gen, device=DEV).to(torch.bfloat16)
+ selected = _selected_tokens(positions, req.long(), seed + 3, overlap)
+ return dict(
+ q=qkv[:, :q_heads],
+ 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, global_k, global_v], dtype=torch.float32, device=DEV
+ ),
+ block_table=block_table,
+ selected_tokens=selected,
+ request_indices=slot_of.to(DEV)[req].to(request_dtype),
+ tokens_per_block=tpb,
+ softmax_scale=d**-0.5,
+ )
+
+
+def _args(spec, seed):
+ """A cell's arguments, `rows_per_request` and the gate included."""
+ kw = {k: v for k, v in spec.items() if k not in ("segments", "gated", "rows_per_request")}
+ args = build(spec["segments"], seed=seed, **kw)
+ args["rows_per_request"] = spec["rows_per_request"]
+ if spec.get("gated"):
+ rows = args["q"].shape[0]
+ g = torch.Generator(device=DEV).manual_seed(seed + 1)
+ args["output_gate"] = torch.randn(rows, 24 * 256, generator=g, device=DEV).to(
+ torch.bfloat16
+ )
+ return args
+
+
+def _takes_grouped(args):
+ return _verify_module.takes_grouped(
+ args["q"],
+ args["k_cache"],
+ args["block_table"],
+ args["selected_tokens"],
+ args["tokens_per_block"],
+ args["rows_per_request"],
+ )
+
+
+def _per_row(args):
+ return per_row(**{k: v for k, v in args.items() if k != "rows_per_request"})
+
+
+def _segments(requests, rows_each, first=4000, step=37):
+ return tuple((first + step * i, rows_each) for i in range(requests))
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell) -> None:
+ args = _args(cell.spec, zlib.crc32(cell.why.encode()) % (2**31))
+ op.is_valid(**args)
+ assert _takes_grouped(args), "every cell runs on the grouped kernel"
+ out = op(**args)
+ assert out.dtype == torch.bfloat16 and not out.isnan().any()
+ el, rms, ok = within_ulp(out, op.reference(**args))
+ assert ok, f"vs reference: {el:.2f} ulp element, {rms:.2f} ulp RMS"
+ el, rms, ok = within_ulp(out, _per_row(args))
+ assert ok, f"vs the per-row kernel: {el:.2f} ulp element, {rms:.2f} ulp RMS"
+
+
+@pytest.mark.parametrize(
+ "segments, rows_per_request",
+ [(_segments(1, 4), 4), (_segments(7, 4), 4), (_segments(128, 1, step=29), 1)],
+ ids=["4 rows", "7 groups of 4", "128 rows of 1"],
+)
+def test_small_batches_and_single_rows_take_the_per_row_kernel(segments, rows_per_request):
+ args = build(segments, seed=31)
+ args["rows_per_request"] = rows_per_request
+ assert not _takes_grouped(args)
+ assert torch.equal(op(**args), _per_row(args))
+
+
+@pytest.mark.parametrize("rows_per_request, splits", [(2, None), (4, 2)], ids=str)
+def test_other_rows_a_request_and_forced_splits(rows_per_request, splits) -> None:
+ args = build(_segments(8, rows_per_request), seed=500 + rows_per_request)
+ out = _verify_module.nvfp4_sparse_gqa_verify_grouped(
+ **args, rows_per_request=rows_per_request, splits=splits
+ )
+ el, rms, ok = within_ulp(out, op.reference(**args, rows_per_request=rows_per_request))
+ assert ok, f"vs reference: {el:.2f} ulp element, {rms:.2f} ulp RMS"
+
+
+def test_launch_plan_splits_only_a_partial_last_wave() -> None:
+ plan = _verify_module.plan_launch
+ # within one wave every unit (group, KV head) takes the same splits
+ assert plan(32, 2, WIDTH, 152) == (2, None)
+ # 512 units on 152 SMs: three full waves of one CTA a unit, the last 56 units split in two
+ assert plan(256, 2, WIDTH, 152) == (2, 456)
+ # a last wave too full to split (104 units): one CTA a unit
+ assert plan(128, 2, WIDTH, 152) == (1, None)
+
+
+@pytest.mark.parametrize("gated", [False, True], ids=["plain", "gated"])
+def test_split_last_wave_matches_one_cta_a_unit(gated) -> None:
+ """The units of the full waves are bitwise the one-CTA-a-unit launch; the split and
+ merged last wave stays inside the band of it and of the reference."""
+ args = build(tuple((2100 + (37 * i) % 1900, 4) for i in range(256)), seed=2718)
+ rows = args["q"].shape[0]
+ if gated:
+ g = torch.Generator(device=DEV).manual_seed(2719)
+ args["output_gate"] = torch.randn(rows, 24 * 256, generator=g, device=DEV).to(
+ torch.bfloat16
+ )
+ sms = torch.cuda.get_device_properties(0).multi_processor_count
+ splits, n_whole = _verify_module.plan_launch(rows // 4, 2, WIDTH, sms)
+ if n_whole is None:
+ pytest.skip(f"{sms} SMs: no partial last wave at {rows} rows")
+ out = op(**args, rows_per_request=4)
+ one = _verify_module.nvfp4_sparse_gqa_verify_grouped(**args, rows_per_request=4, splits=1)
+ head = (n_whole // 2) * 4
+ assert torch.equal(out[:head], one[:head])
+ el, rms, ok = within_ulp(out, one)
+ assert ok, f"vs one CTA a unit: {el:.2f} ulp element, {rms:.2f} ulp RMS"
+ el, rms, ok = within_ulp(out, op.reference(**args, rows_per_request=4))
+ assert ok, f"vs reference: {el:.2f} ulp element, {rms:.2f} ulp RMS"
+
+
+def test_rows_of_an_invalid_request_are_zero() -> None:
+ args = build(_segments(8, 4), seed=99)
+ args["request_indices"][5] = -1
+ out = op(**args, rows_per_request=4)
+ assert not out[5].any()
+ el, rms, ok = within_ulp(out, op.reference(**args, rows_per_request=4))
+ assert ok, f"vs reference: {el:.2f} ulp element, {rms:.2f} ulp RMS"
+
+
+def test_unaligned_row_slice_of_the_selection() -> None:
+ """A mixed step hands its generation rows as a row slice of the step's 2051-wide
+ selection, 16-byte aligned only at every fourth row: an unaligned slice gives the same
+ rows as an aligned copy."""
+ args = build(((3000, 1),) + _segments(8, 4), seed=4242)
+ aligned = dict(args)
+ for name in ("q", "selected_tokens", "request_indices"):
+ args[name] = args[name][1:]
+ aligned[name] = args[name].clone()
+ assert args["selected_tokens"].data_ptr() % 16
+ assert torch.equal(op(**args, rows_per_request=4), op(**aligned, rows_per_request=4))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_prefill_compress.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_prefill_compress.py
new file mode 100644
index 000000000000..f2426eb8d1ec
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_prefill_compress.py
@@ -0,0 +1,237 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the default target's qsa_prefill_compress variant.
+
+The paged state is real: an index-key pool of random raw keys, a position
+plane holding `(t, h, w)` for every allocated position, and the engine's
+`[slots, width]` table -- at least as many slots as requests, each live row's
+pages drawn from a seeded permutation of the pool, `-1` past its allocation,
+rows past the batch holding a previous step's plausible page ids. The cos|sin
+table is built the way the target builds it at load, `max_position_embeddings`
+rows tall.
+
+The `coalesced` cells hold the default target's pool: one rank's measured page count,
+the 12 sparse layers' index pages sharing each slot as KV cache manager V2 lays
+them out at bf16, the other 11 layers' planes holding a sentinel that must
+survive. Those pools are drawn on the device.
+
+A padding row at logical position -1 is deliberately absent: it is an illegal
+memory access that poisons the process for every test after it.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_prefill_compress import ( # noqa: E501
+ qsa_prefill_compress as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401
+
+assert torch.cuda.is_available(), "qsa_prefill_compress requires a CUDA device"
+
+
+def _compare(out: torch.Tensor, ref: torch.Tensor, 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. The 8192-row tep4 prefill cell has five of 262016, each an exact bf16 tie in
+ the reference's fp32 value; the nearest wrong variant measured there (sequential bf16
+ pooling) moves 94580.
+ """
+ assert out.dtype == ref.dtype, (out.dtype, ref.dtype)
+ assert out.shape == ref.shape, (out.shape, ref.shape)
+ # Only the differing elements are widened: a pool-sized int32 copy would not fit beside
+ # the target's pool.
+ a, b = out.view(torch.int16), ref.view(torch.int16)
+ differs = a != b
+ count = int(differs.sum())
+ if count == 0:
+ return
+ distance = int((a[differs].int() - b[differs].int()).abs().max())
+ assert distance <= 1 and count <= flips, (
+ f"{count} elements differ, by up to {distance} bf16 steps; at most {flips} may differ by one"
+ )
+
+
+DEV = "cuda"
+ROTARY_DIM = 64
+THETA = 1.0e7
+#: The target's table height: the checkpoint's `max_position_embeddings`.
+MAX_POSITIONS = 262144
+SENTINEL = 0x3C3C
+
+
+def _cos_sin() -> torch.Tensor:
+ """The target's `[max_positions, 2 * pairs]` float32 cos|sin table."""
+ exponent = torch.arange(0, ROTARY_DIM, 2, dtype=torch.float64, device=DEV)
+ inv_freq = (1.0 / (THETA ** (exponent / ROTARY_DIM))).to(torch.float32)
+ positions = torch.arange(MAX_POSITIONS, dtype=torch.float32, device=DEV)
+ angles = positions[:, None] * inv_freq[None, :]
+ return torch.cat([torch.cos(angles), torch.sin(angles)], dim=1).contiguous()
+
+
+def _build(spec, seed: int):
+ """Real caches, table and rows for one cell; returns the op's kwargs and extras."""
+ gen = torch.Generator(device="cpu").manual_seed(seed)
+ tpb, head_dim = spec["tokens_per_block"], spec["head_dim"]
+ slots, width, pages = spec["slots"], spec["width"], spec["pages"]
+ segments = spec["segments"]
+ requests = [r for r, (_, n) in enumerate(segments) for _ in range(n)]
+ positions = [p for start, n in segments for p in range(start, start + n)]
+
+ table = torch.randint(0, pages, (slots, width), generator=gen, dtype=torch.int32)
+ order = torch.randperm(pages, generator=gen).to(torch.int32)
+ if spec.get("top_page"):
+ # Request 0's last page is the pool's last, so its anchors address the top of the pool.
+ last = (segments[0][0] + segments[0][1] - 1) // tpb
+ where = int((order == pages - 1).nonzero())
+ order[[where, last]] = order[[last, where]]
+ coalesced = spec["layout"] == "coalesced"
+ index_cache = (
+ None # drawn on the device below
+ if coalesced
+ else torch.randn(pages, tpb, 1, head_dim, generator=gen).to(torch.bfloat16)
+ )
+ position_cache = torch.zeros(pages, tpb, 3, dtype=torch.int32)
+ used = 0
+ for r, (start, n) in enumerate(segments):
+ cols = (start + n - 1) // tpb + 1
+ table[r] = -1
+ table[r, :cols] = order[used : used + cols]
+ used += cols
+ p = torch.arange(cols * tpb).view(cols, tpb)
+ 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"])
+ table, position_cache = table.to(DEV), position_cache.to(DEV)
+
+ extras = {}
+ if coalesced:
+ # `roles` layers' index pages per slot; this layer's is `ordinal`, the rest a sentinel.
+ roles, ordinal = spec["roles"], spec["ordinal"]
+ pool = torch.empty(pages, roles, tpb, 1, head_dim, dtype=torch.bfloat16, device=DEV)
+ pool.view(torch.int16).fill_(SENTINEL)
+ index_cache = pool[:, ordinal]
+ device_gen = torch.Generator(device=DEV).manual_seed(seed)
+ index_cache.copy_(
+ torch.randn(
+ pages, tpb, 1, head_dim, generator=device_gen, device=DEV, dtype=torch.bfloat16
+ )
+ )
+ assert index_cache.stride() == (roles * tpb * head_dim, head_dim, head_dim, 1)
+ extras["other_layers"] = (pool, ordinal)
+ else:
+ index_cache = index_cache.to(DEV)
+ if spec["layout"] == "shared_page":
+ # An NVFP4 pool's page: `roles` index pages, then the position page. The other
+ # index roles hold a sentinel and must survive, as must the position page.
+ roles = spec["roles"]
+ index_bytes = tpb * head_dim * 2
+ page = torch.empty(pages, roles * index_bytes + tpb * 3 * 4, dtype=torch.uint8, device=DEV)
+ others = page[:, : (roles - 1) * index_bytes].view(torch.int16)
+ others.fill_(SENTINEL)
+ shared_index = page[:, (roles - 1) * index_bytes : roles * index_bytes]
+ shared_index = shared_index.view(torch.bfloat16).view(pages, tpb, 1, head_dim)
+ shared_position = page[:, roles * index_bytes :].view(torch.int32).view(pages, tpb, 3)
+ shared_index.copy_(index_cache)
+ shared_position.copy_(position_cache)
+ assert shared_index.stride(0) == roles * tpb * head_dim + tpb * 3 * 2
+ assert shared_position.stride() == (roles * tpb * head_dim // 2 + tpb * 3, 3, 1)
+ index_cache, position_cache = shared_index, shared_position
+ extras["neighbour"] = (others, others.clone())
+ extras["position_page"] = (position_cache, position_cache.clone())
+ else:
+ assert spec["layout"] in ("plain", "coalesced"), spec["layout"]
+
+ args = dict(
+ logical_positions=logical.to(DEV),
+ request_indices=request.to(DEV),
+ block_table=table,
+ index_cache=index_cache,
+ position_cache=position_cache,
+ k_norm_weight=k_norm_weight.to(DEV),
+ cos_sin=_cos_sin(),
+ eps=spec["eps"],
+ tokens_per_block=tpb,
+ compress_ratio=spec["compress_ratio"],
+ mrope_section=spec["mrope_section"],
+ )
+ return args, extras
+
+
+def _anchors(args) -> int:
+ """Distinct anchor slots the rows name."""
+ logical = args["logical_positions"].long()
+ keyed = args["request_indices"].long() * (1 << 32) + logical
+ return int(torch.unique(keyed[(logical + 1) % args["compress_ratio"] == 0]).numel())
+
+
+def _top_anchor_page(args) -> int:
+ """The highest page an anchor row writes."""
+ logical = args["logical_positions"].long()
+ anchored = (logical + 1) % args["compress_ratio"] == 0
+ request = args["request_indices"].long()[anchored]
+ pages = args["block_table"].long()[request, logical[anchored] // args["tokens_per_block"]]
+ return int(pages.max())
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying): # noqa: F811
+ certifying(op)
+ args, extras = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
+ before = args["index_cache"].clone()
+ ref = op.reference(**args)
+ out = op(**args)
+ assert out is args["index_cache"]
+ _compare(out, ref, 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 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"
+ if "other_layers" in extras:
+ pool, ordinal = extras["other_layers"]
+ for other in range(pool.shape[1]):
+ if other != ordinal:
+ plane = pool[:, other].view(torch.int16)
+ assert bool((plane == SENTINEL).all()), f"wrote into layer ordinal {other}'s plane"
+ if cell.spec.get("top_page"):
+ # The pool's last page is written, at an element offset past int32 from the base.
+ top = _top_anchor_page(args)
+ assert top == cell.spec["pages"] - 1, top
+ assert top * out.stride(0) >= 2**31, top * out.stride(0)
+
+
+#: The cell of 32 MTP7 verify blocks (256 rows, two anchors each): the segments, pool and table.
+_MTP7_32 = next(c.spec for c in op.CELLS if c.why.startswith("decode MTP7: 32 requests"))
+_MTP7_BATCH = {key: _MTP7_32[key] for key in ("segments", "pages", "slots", "rounding_flips")}
+
+
+def _strided(tensor: torch.Tensor, dim: int) -> torch.Tensor:
+ """The same values as a stride-2 view along `dim` of a zero-filled buffer."""
+ shape = list(tensor.shape)
+ shape[dim] *= 2
+ wide = torch.zeros(shape, dtype=tensor.dtype, device=tensor.device)
+ index = [slice(None)] * tensor.dim()
+ index[dim] = slice(None, None, 2)
+ wide[tuple(index)] = tensor
+ return wide[tuple(index)]
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_rollback_planes.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_rollback_planes.py
new file mode 100644
index 000000000000..1568de7ac079
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_rollback_planes.py
@@ -0,0 +1,219 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the qsa_rollback_planes variant: `qsa_rollback_gather` and
+`qsa_rollback_restore`, each cell through the entry's own reference, bitwise (both are copies).
+
+The pool is filled with a NaN pattern `randn` never draws, then every slot a cell names is
+written, so a whole-pool count of what no longer holds the fill shows that a call wrote nothing
+outside those slots. The tp1 cells run over separate contiguous planes and over the packed
+side-plane page of an FP8 / NVFP4 pool; the default target's cells over core's coalesced bf16 pool at a
+tep4 rank's measured page count, with the requests' pages drawn across it and the pool's first
+and last page among them, so the last is addressed past 2^31 elements from a plane's base.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_rollback_planes import ( # noqa: E501
+ ROW,
+ plane_table,
+ qsa_rollback_gather,
+ qsa_rollback_restore,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "qsa_rollback_planes requires a CUDA device"
+
+DEV = "cuda"
+TPB = 32
+#: The tp1 cells' pool pages.
+SMALL_POOL = 600
+#: One plane's page: `tokens_per_block` 128-wide rows.
+PLANE_PAGE = TPB * ROW
+#: The QSA position page an FP8 / NVFP4 pool packs after the 12 index pages (384 bytes).
+SIDE_PLANES_PAD = 192
+#: `max_seq_len` under MTP: the last position a verify row can hold is one less.
+MTP_MAX_SEQ_LEN = 16384
+#: bf16 NaN with a payload: the pool's fill.
+SENTINEL = 0x7FA5
+
+
+class _Pool:
+ """One rank's index planes, every element at `SENTINEL` until a slot is written.
+
+ `contiguous`: one `[pages, 32, 1, 128]` buffer a plane. `shared_page`: the planes side by
+ side in one page with the position page after them. `coalesced`: core's KV cache manager V2
+ pool, layer k's buffer at element 4096 k of a 49152-element page.
+ """
+
+ def __init__(self, spec: dict, layout: str):
+ count, pages = spec["planes"], spec.get("pool_pages", SMALL_POOL)
+ if layout == "contiguous":
+ self.buffers = [
+ torch.empty(pages, TPB, 1, ROW, dtype=torch.bfloat16, device=DEV)
+ for _ in range(count)
+ ]
+ self.planes = list(self.buffers)
+ else:
+ pad = SIDE_PLANES_PAD if layout == "shared_page" else 0
+ page = torch.empty(pages, count * PLANE_PAGE + pad, dtype=torch.bfloat16, device=DEV)
+ self.buffers = [page]
+ self.planes = [
+ page[:, k * PLANE_PAGE : (k + 1) * PLANE_PAGE].view(pages, TPB, 1, ROW)
+ for k in range(count)
+ ]
+ for b in self.buffers:
+ b.view(torch.int16).fill_(SENTINEL)
+ self.pages = pages
+
+ def written(self) -> int:
+ """Elements anywhere in the pool that no longer hold the fill."""
+ total = 0
+ for b in self.buffers:
+ for chunk in b.view(torch.int16).reshape(b.shape[0], -1).split(8192):
+ total += int((chunk != SENTINEL).sum())
+ return total
+
+
+def _layouts(cell):
+ return (cell.spec["layout"],) if "layout" in cell.spec else ("contiguous", "shared_page")
+
+
+def _cases(op):
+ return [(c, layout) for c in op.CELLS for layout in _layouts(c)]
+
+
+def _ids(cases):
+ return [f"{c.why[:44]}-{layout}" for c, layout in cases]
+
+
+GATHER_CASES = _cases(qsa_rollback_gather)
+RESTORE_CASES = _cases(qsa_rollback_restore)
+
+
+def _verify_slots(spec, seed):
+ """`requests` x `p` consecutive positions each, every request's QSA pages its own physical
+ pages of the pool. Request 0 straddles a page boundary from the pool's first page into its
+ last; the last request ends at position 16383."""
+ n, p, pool = spec["requests"], spec["p"], spec["pool_pages"]
+ starts = (
+ [TPB - p // 2]
+ + [(29 + 127 * s) % (MTP_MAX_SEQ_LEN - p) for s in range(1, n - 1)]
+ + [MTP_MAX_SEQ_LEN - p]
+ )
+ positions = torch.tensor(starts)[:, None] + torch.arange(p)[None, :]
+ logical = positions // TPB
+ rest = torch.randperm(pool - 2, generator=torch.Generator().manual_seed(seed)) + 1
+ physical = torch.cat([torch.tensor([0, pool - 1]), rest])
+ pages = torch.empty(n, p, dtype=torch.int64)
+ cursor = 0
+ for s in range(n):
+ for lp in logical[s].unique().tolist():
+ pages[s][logical[s] == lp] = physical[cursor]
+ cursor += 1
+ return pages.reshape(-1).to(DEV), (positions % TPB).reshape(-1).to(DEV)
+
+
+def _slots(spec, seed, g):
+ if "requests" in spec:
+ return _verify_slots(spec, seed)
+ rows = spec["rows"]
+ pages = torch.randperm(SMALL_POOL, generator=torch.Generator().manual_seed(rows))[:rows]
+ within = torch.randint(0, TPB, (rows,), generator=g, device=DEV)
+ return pages.to(DEV).to(torch.int64), within.to(torch.int64)
+
+
+def _rejected(spec, g):
+ """The call site's mask for a verify step (bool, `ordinal >= accepted`; one request rejects
+ every draft, one none), or the tp1 cells' int32 one rejecting every other row, so each
+ cell restores some slots and keeps the rest."""
+ if spec.get("rejected") != "bool":
+ return (torch.arange(spec["rows"], device=DEV) % 2).to(torch.int32)
+ n, p = spec["requests"], spec["p"]
+ accepted = torch.randint(1, p + 1, (n,), generator=g, device=DEV)
+ accepted[0], accepted[1] = 1, p
+ return (torch.arange(p, device=DEV)[None, :] >= accepted[:, None]).reshape(-1)
+
+
+def _fill(planes, pages, within, g) -> None:
+ """Write fresh values into every slot of every plane (a step's index store)."""
+ for t in planes:
+ t[pages, within] = torch.randn(pages.shape[0], 1, ROW, generator=g, device=DEV).to(t.dtype)
+
+
+def _slots_now(planes, pages, within) -> torch.Tensor:
+ return torch.stack([t[pages, within].reshape(-1, ROW) for t in planes])
+
+
+def _assert_far_page(spec, pool, far_pages) -> None:
+ """A measured-pool cell addresses the pool's last page, past 2^31 elements from the base of
+ every plane, where a 32-bit offset would wrap."""
+ if "pool_pages" not in spec:
+ return
+ last = int(far_pages.max())
+ assert last == pool.pages - 1, f"the pool's last page {pool.pages - 1} is unread ({last})"
+ assert last * pool.planes[-1].stride(0) >= 2**31
+
+
+def _setup(cell, layout):
+ seed = zlib.crc32(f"{cell.why}/{layout}".encode()) % 2**31
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ pool = _Pool(cell.spec, layout)
+ pages, within = _slots(cell.spec, seed, g)
+ assert pages.shape == (cell.spec["rows"],)
+ _fill(pool.planes, pages, within, g)
+ return g, pool, pages, within
+
+
+@pytest.mark.parametrize("cell, layout", GATHER_CASES, ids=_ids(GATHER_CASES))
+def test_gather_cells(cell, layout, certifying): # noqa: F811
+ certifying(qsa_rollback_gather)
+ _, pool, pages, within = _setup(cell, layout)
+ ptrs, strides = plane_table(pool.planes)
+ count, rows = len(pool.planes), pages.shape[0]
+
+ want = qsa_rollback_gather.reference(pool.planes, pages, within, count)
+ got = qsa_rollback_gather(ptrs, strides, pages, within, count)
+
+ assert got.shape == (count, rows, ROW) and got.dtype == torch.bfloat16
+ assert torch.equal(got.view(torch.int16), want.view(torch.int16))
+ assert torch.equal(
+ _slots_now(pool.planes, pages, within).view(torch.int16), want.view(torch.int16)
+ )
+ assert pool.written() == count * rows * ROW, "the gather wrote into the pool"
+ _assert_far_page(cell.spec, pool, pages)
+
+
+@pytest.mark.parametrize("cell, layout", RESTORE_CASES, ids=_ids(RESTORE_CASES))
+def test_restore_cells(cell, layout, certifying): # noqa: F811
+ certifying(qsa_rollback_restore)
+ g, pool, pages, within = _setup(cell, layout)
+ ptrs, strides = plane_table(pool.planes)
+ count, rows = len(pool.planes), pages.shape[0]
+ # What the begin step's gather captured (its own cells hold it to this), then the step
+ # overwrites every slot.
+ old = _slots_now(pool.planes, pages, within)
+ _fill(pool.planes, pages, within, g)
+ rejected = _rejected(cell.spec, g)
+ assert bool(rejected.bool().any()) and not bool(rejected.bool().all())
+
+ want = qsa_rollback_restore.reference(pool.planes, pages, within, rejected, old)
+ assert qsa_rollback_restore(ptrs, strides, pages, within, rejected, old) is None
+
+ got = _slots_now(pool.planes, pages, within)
+ assert torch.equal(got.view(torch.int16), want.view(torch.int16))
+ assert pool.written() == count * rows * ROW, "the restore wrote outside the rows' slots"
+ _assert_far_page(cell.spec, pool, pages[rejected.bool()])
+
+
+def _small():
+ spec = dict(planes=12, rows=16)
+ g = torch.Generator(device=DEV).manual_seed(5)
+ pool = _Pool(spec, "contiguous")
+ pages, within = _slots(spec, 5, g)
+ _fill(pool.planes, pages, within, g)
+ return g, pool, pages, within
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py
new file mode 100644
index 000000000000..cc612e23c742
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py
@@ -0,0 +1,238 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the 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 pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention import ( # noqa: E501
+ qsa_sparse_gqa_rowtile as _module,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.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_union_kernel",
+ "_rowtile_group_kernel",
+ "_rowtile_slots_kernel",
+ "_rowtile_lone_kernel",
+ "_rowtile_attention_cute",
+)
+
+
+@pytest.fixture
+def launches(monkeypatch):
+ log: list = []
+ for name in _KERNELS:
+ inner = getattr(_module, name)
+ monkeypatch.setattr(_module, name, _LaunchRecorder(inner, name, log))
+ return log
+
+
+def _selected_tokens(positions, seed, compress_ratio, token_topk, shared_scores=False):
+ """QSA's expansion layout: a random `min(available, block_topk)`-subset of
+ the complete visible groups, expanded, then the trailing partial group,
+ then `-1`, packed at the front. With `shared_scores` every row ranks the
+ groups by one score vector, so adjacent rows pick nearly the same set (a
+ sparse union, as on real selections)."""
+ block_topk = token_topk // compress_ratio
+ final_topk = token_topk + compress_ratio - 1
+ gen = torch.Generator(device=DEV).manual_seed(seed)
+ rows = positions.numel()
+ visible = positions.long() + 1
+ available = visible // compress_ratio
+ taken = available.clamp(max=block_topk)
+ width = max(block_topk, int(available.max().item()) if rows else 0)
+ out = torch.empty(rows, final_topk, dtype=torch.int32, device=DEV)
+ column = torch.arange(final_topk, device=DEV)[None, :]
+ for start in range(0, rows, 2048):
+ stop = min(start + 2048, rows)
+ avail, take, vis = available[start:stop], taken[start:stop], visible[start:stop]
+ if shared_scores:
+ scores = torch.rand(1, width, generator=gen, device=DEV).expand(stop - start, -1)
+ gen.manual_seed(seed)
+ else:
+ scores = torch.rand(stop - start, width, generator=gen, device=DEV)
+ group_ids = torch.arange(width, device=DEV)[None, :]
+ scores = torch.where(group_ids < avail[:, None], scores, torch.inf)
+ chosen = scores.argsort(dim=1)[:, :block_topk]
+ group_of_column = (column // compress_ratio).clamp(max=block_topk - 1)
+ body = torch.gather(chosen, 1, group_of_column.expand(stop - start, -1))
+ body = body * compress_ratio + column % compress_ratio
+ in_body = column < (take * compress_ratio)[:, None]
+ tail_column = column - (take * compress_ratio)[:, None]
+ tail = (avail * compress_ratio)[:, None] + tail_column
+ in_tail = ~in_body & (tail_column >= 0) & (tail < vis[:, None])
+ out[start:stop] = torch.where(in_body, body, torch.where(in_tail, tail, -1)).to(torch.int32)
+ return out
+
+
+def _build(spec, seed):
+ """Real paged state for one spec, keyed by `__call__`'s parameter names."""
+ gen = torch.Generator(device=DEV).manual_seed(seed)
+ tpb, kv, d, q_heads = (
+ spec["tokens_per_block"],
+ spec["kv_heads"],
+ spec["head_dim"],
+ spec["q_heads"],
+ )
+ segments = list(spec["segments"])
+ positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments])
+ req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)])
+ rows, slots = positions.numel(), spec["slots"]
+ assert len(segments) <= slots
+ slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(segments)]
+ mapped = [math.ceil((s + n) / tpb) for s, n in segments]
+ assert max(mapped) <= spec["pages"]
+ pool_pages = sum(mapped) + 7
+ page_ids = torch.randperm(pool_pages, generator=gen, device=DEV).to(torch.int32)
+ block_table = torch.full((slots, spec["pages"]), -1, dtype=torch.int32, device=DEV)
+ cursor = 0
+ for slot, count in zip(slot_of.tolist(), mapped):
+ block_table[slot, :count] = page_ids[cursor : cursor + count]
+ cursor += count
+
+ pool = torch.randn(pool_pages, 2, kv, tpb, d, generator=gen, device=DEV).to(spec["cache_dtype"])
+ k_cache, v_cache = pool[:, 0], pool[:, 1]
+ assert spec["cache_layout"] == "pool_planes" and spec["q_layout"] == "qkv_slice"
+ qkv = torch.randn(rows, q_heads + 2 * kv, d, generator=gen, device=DEV).to(spec["q_dtype"])
+ q = qkv[:, :q_heads]
+
+ selected = _selected_tokens(positions, seed + 3, spec["compress_ratio"], spec["token_topk"],
+ spec.get("shared_scores", False)) # fmt: skip
+ request_indices = slot_of.to(DEV)[req].to(spec["request_dtype"])
+ return dict(
+ q=q,
+ k_cache=k_cache,
+ v_cache=v_cache,
+ block_table=block_table,
+ selected_tokens=selected,
+ request_indices=request_indices,
+ tokens_per_block=tpb,
+ softmax_scale=d**-0.5,
+ query_positions=positions.to(spec["request_dtype"]),
+ compress_ratio=spec["compress_ratio"],
+ )
+
+
+def _bound(positions, spec):
+ cr = spec["compress_ratio"]
+ visible = positions.long() + 1
+ return (visible // cr).clamp(max=spec["token_topk"] // cr) * cr + visible % cr
+
+
+def _is_prefix_packed(table):
+ valid = table >= 0
+ columns = torch.arange(table.shape[1], device=table.device)[None, :]
+ return (valid == (columns < valid.sum(1)[:, None])).all(1)
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811
+ certifying(op)
+ spec = cell.spec
+ args = _build(spec, seed=sum(map(ord, cell.why)))
+ rows = args["q"].shape[0]
+ # The per-row kernel's precondition, which this op shares.
+ selected = args["selected_tokens"]
+ assert bool(_is_prefix_packed(selected).all())
+ assert torch.equal((selected >= 0).sum(1), _bound(args["query_positions"], spec))
+ out = op(**args)
+ assert [name for name, _ in launches] == list(_KERNELS)
+ groups = math.ceil(rows / 10)
+ # The plan's union and slot passes: a program per (group, 32 bitmap words of the table).
+ chunks = math.ceil(math.ceil(spec["pages"] * spec["tokens_per_block"] / 64) / 32)
+ assert launches[0][1] == (rows,)
+ assert launches[1][1] == launches[3][1] == (groups, chunks)
+ assert launches[2][1] == (groups,)
+ assert launches[4][1] == (min(rows, 256), spec["kv_heads"])
+ # The row-tile pass is persistent: one CTA per SM, never more than (group, KV head) items.
+ items = groups * spec["kv_heads"]
+ assert launches[5][1] == (_module._rowtile_ctas(items, args["q"].device),)
+ assert launches[5][1][0] == min(
+ items, torch.cuda.get_device_properties(0).multi_processor_count
+ )
+ assert out.shape == (rows, spec["q_heads"], spec["head_dim"])
+ assert out.dtype == args["q"].dtype and out.is_contiguous()
+ assert not bool(out.isnan().any())
+ _compare(out, op.reference(**args))
+
+
+CKPT = op.CELLS[0].spec
+assert (CKPT["q_heads"], CKPT["kv_heads"], CKPT["head_dim"]) == (24, 2, 256)
+
+
+def _with(segments, **overrides):
+ return {**CKPT, "segments": tuple(segments), **overrides}
+
+
+@pytest.mark.parametrize(
+ "segments",
+ [
+ # Fewer items (2 groups x 2 KV heads) than SMs: a 4-CTA grid, mostly one item a CTA.
+ [(5000, 17)],
+ # A lone group (rows of two requests) between row-tile groups of each request.
+ [(3000, 25), (7000, 25)],
+ ],
+ ids=["fewer items than SMs", "a lone group between row-tile groups"],
+)
+def test_small_launches_match_the_reference(segments) -> None:
+ """The persistent pass's edges: CTAs left without a second item, and an item the pass
+ skips (a lone group) between items it serves."""
+ args = _build(_with(segments, slots=8), seed=7700 + len(segments))
+ _compare(op(**args), op.reference(**args))
+
+
+def test_long_context_disjoint_groups_match_the_reference() -> None:
+ """Rows deep in a long context with independent selections: a group's list is about
+ ten times one row's columns (unions of up to 20480 tokens over 2690 pages, the plan's
+ words spanning many 32-word chunks). Every group of one request takes the row-tile
+ pass, however long its list."""
+ args = _build(_with([(86000, 197)], slots=4, pages=2700), seed=9300)
+ _compare(op(**args), op.reference(**args))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_expand_decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_expand_decode.py
new file mode 100644
index 000000000000..28f84c8d508f
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_expand_decode.py
@@ -0,0 +1,179 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the qsa_topk_expand_decode entry.
+
+Cells against the entry's reference (a stable descending sort, ties to the higher block,
+then the expansion entry's reference): the block list and the token selection exactly.
+Columns past each row's visible count carry garbage (huge values, NaN) that must never
+be read. Beyond the cells: the wide rows against the canonical composition, and a
+stale-read check under graph replay at each decode-path cell's shape.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_canonical_selection import ( # noqa: E501
+ qsa_canonical_selection,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_expand_block_indices import ( # noqa: E501
+ qsa_expand_block_indices,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_topk_expand_decode import ( # noqa: E501
+ qsa_topk_expand_decode as op,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_topk_select import ( # noqa: E501
+ qsa_topk_select,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "qsa_topk_expand_decode requires a CUDA device"
+
+DEV = "cuda"
+K, CR, TOKEN_TOPK = 512, 4, 2048
+
+
+def _build(spec, seed):
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ W = spec["width"]
+ if "chunk" in spec or "requests" in spec:
+ # Each request's consecutive rows (a context chunk, or a verify step's 1 + draft
+ # rows): every row's KV length is the request's last row + 1.
+ requests = spec["requests"] if "requests" in spec else (spec["chunk"],)
+ pos = torch.cat(
+ [torch.arange(s, s + c, dtype=torch.int64, device=DEV) for s, c in requests]
+ )
+ seqlen = torch.cat(
+ [torch.full((c,), s + c, dtype=torch.int32, device=DEV) for s, c in requests]
+ )
+ else:
+ pos = torch.tensor(spec["pos"], dtype=torch.int64, device=DEV)
+ seqlen = (pos + 1).to(torch.int32)
+ rows = pos.numel()
+ assert rows == spec["rows"], (rows, spec["rows"])
+ scores = torch.randn(rows, W, generator=g, device=DEV).abs() * 3.0
+ if spec.get("relu_sparse"):
+ # Fewer positive scores than the budget, so the zero class straddles it.
+ scores = torch.relu(torch.randn(rows, W, generator=g, device=DEV) - 2.2)
+ scores[:, ::7] = -0.0 # negative zero ties with zero
+ if spec.get("relu_sum"):
+ # The indexer's form: four heads' ReLUs summed, times a scale (zero in 1 of 16).
+ scores = torch.zeros(rows, W, device=DEV)
+ for _ in range(4):
+ scores += torch.relu(torch.randn(rows, W, generator=g, device=DEV))
+ scores *= 0.0884
+ if spec.get("equal"):
+ # Every score tied: the budget is the highest-index blocks.
+ scores = torch.full((rows, W), 0.5, device=DEV)
+ if spec.get("narrow"):
+ scores = 1.0 + 0.05 * torch.rand(rows, W, generator=g, device=DEV)
+ if spec.get("lognormal"):
+ scores = torch.exp(1.74 + 0.2 * torch.randn(rows, W, generator=g, device=DEV))
+ # bf16-like resolution, so equal scores tie; a zero in every other row.
+ scores = scores.to(torch.bfloat16).float()
+ scores[::2, 5] = 0.0
+ ends = ((pos + 1) // CR).to(torch.int32)
+ assert int(ends.max()) <= W, "a window wider than the score table"
+ # Garbage past the visible blocks: huge values, and a NaN three columns on.
+ cols = torch.arange(W, device=DEV)
+ scores = torch.where(cols[None, :] >= ends[:, None], 1.0e30, scores)
+ scores = torch.where(cols[None, :] == ends[:, None] + 3, float("nan"), scores)
+ topk = torch.full((rows, K), -7, dtype=torch.int32, device=DEV)
+ return [scores, ends, pos, seqlen, topk, CR, TOKEN_TOPK]
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
+ certifying(op)
+ args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31)
+ want_topk, want_sel = op.reference(*args)
+ margin = cell.spec.get("out_margin")
+ if margin is None:
+ sel = op(*args)
+ else:
+ # Rows [margin, margin + rows) of a step-sized selection; the rows around stay put.
+ rows = args[0].shape[0]
+ whole = torch.full(
+ (rows + 2 * margin, TOKEN_TOPK + CR - 1), -9, dtype=torch.int32, device=DEV
+ )
+ sel = op(*args, out=whole[margin : margin + rows])
+ assert sel.data_ptr() == whole[margin].data_ptr()
+ assert (whole[:margin] == -9).all() and (whole[margin + rows :] == -9).all(), (
+ "rows outside out"
+ )
+ assert sel.dtype == torch.int32 and sel.shape == want_sel.shape
+ assert torch.equal(args[4], want_topk), "block list"
+ assert torch.equal(sel, want_sel.to(torch.int32)), "token selection"
+
+
+@pytest.mark.parametrize("width", [65536, 65568, 65600])
+@pytest.mark.parametrize("distribution", ["relu_sum", "relu_sparse", "equal"])
+def test_wide_rows_equal_the_canonical_composition(width: int, distribution: str) -> None:
+ """Match the fused block order and expansion to exact top-k plus canonical rewrite.
+
+ The TP1 table includes eight pages of headroom; the default table includes four.
+ Ragged windows cross the block budget and slab boundaries and carry every causal
+ tail width. Both paths must ignore the poisoned score columns past each window.
+ """
+ positions = (0, 1, 2, 3, 2043, 2047, 2051, 16383, 84000, 99999, 262143)
+ spec = dict(rows=len(positions), pos=positions, width=width, **{distribution: True})
+ args = _build(spec, seed=width)
+ scores, ends, pos, lengths, fused_topk, cr, token_topk = args
+ lengths[8] += 3 # A verify row whose request continues past this query position.
+ starts = torch.zeros_like(ends)
+ composed_topk = torch.full_like(fused_topk, -1)
+ qsa_topk_select(scores, starts, ends, composed_topk, index_topk=K)
+ qsa_canonical_selection(scores, starts, ends, composed_topk)
+ composed = qsa_expand_block_indices(
+ composed_topk, pos, lengths, compress_ratio=cr, token_topk=token_topk
+ )
+ fused = op(*args)
+ expected_topk, expected = op.reference(*args)
+ assert torch.equal(composed_topk, expected_topk), "canonical block list"
+ assert torch.equal(composed, expected), "canonical token expansion"
+ assert torch.equal(fused_topk, composed_topk), "fused block order"
+ assert torch.equal(fused, composed), "fused token expansion"
+
+
+# The decode-path shapes, captured in the decode CUDA graph: tp1's one row at 2560 columns,
+# the default target's relu-sparse decode and verify cells (1 and 128 rows at core's width, 1024
+# at MTP's), and the wide tables' decode cells (`graph` in their spec).
+_GRAPH_SPECS = [dict(rows=1, pos=(8700,), width=2560)] + [
+ c.spec
+ for c in op.CELLS
+ if c.spec.get("graph") or (c.spec.get("relu_sparse") and "out_margin" not in c.spec)
+]
+
+
+@pytest.mark.parametrize(
+ "spec", _GRAPH_SPECS, ids=[f"{s['rows']}x{s['width']}" for s in _GRAPH_SPECS]
+)
+def test_reads_after_the_wait_under_graph_replay(spec) -> None:
+ args = _build(spec, seed=21)
+ src = args[0].clone()
+
+ def step():
+ args[0].copy_(src * 1.0)
+ return op(*args)
+
+ warmup = torch.cuda.Stream()
+ warmup.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(warmup):
+ step()
+ torch.cuda.current_stream().wait_stream(warmup)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ graph_sel = step()
+ for trial in range(3):
+ src.copy_(_build(spec, seed=100 + trial)[0])
+ graph.replay()
+ torch.cuda.synchronize()
+ got_topk = args[4].clone()
+ sel = op(src.clone(), *args[1:])
+ assert torch.equal(graph_sel, sel) and torch.equal(got_topk, args[4]), f"replay {trial}"
+ want_topk, want_sel = op.reference(src, *args[1:])
+ assert torch.equal(got_topk, want_topk), f"replay {trial}: block list"
+ assert torch.equal(graph_sel, want_sel), f"replay {trial}: token selection"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_select.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_select.py
new file mode 100644
index 000000000000..b13321c3b687
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_select.py
@@ -0,0 +1,105 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the qsa_topk_select catalog entry.
+
+The op selects; it computes nothing. `_compare` therefore checks the output's
+structure exactly and the *multiset of selected values* bit for bit, never the
+permutation, which the op leaves unspecified.
+"""
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_topk_select import ( # noqa: E501
+ TopkExpectation,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_topk_select import ( # noqa: E501
+ qsa_topk_select as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401
+
+assert torch.cuda.is_available(), "qsa_topk_select requires a CUDA device"
+
+
+def _compare(out: torch.Tensor, ref: TopkExpectation) -> None:
+ """Structure exactly, then the selected values bit for bit.
+
+ A selection does no arithmetic, so nothing is tolerated. Values rather
+ than offsets are compared so that a tied k-th score -- common here: the
+ relu cells are half zeros -- is a pass rather than a coin flip.
+ """
+ assert out.dtype == torch.int32, out.dtype
+ assert out.shape == ref.values.shape, (out.shape, ref.values.shape)
+ index_topk = out.shape[1]
+ valid = out >= 0
+ assert not bool((valid[:, 1:] & ~valid[:, :-1]).any()), "a valid offset follows a -1 pad"
+ assert bool((out >= -1).all()), "padding must be exactly -1"
+ assert torch.equal(valid.sum(1), ref.window.clamp(max=index_topk)), (
+ "each row must select exactly min(index_topk, window) offsets"
+ )
+ assert bool((out < ref.window[:, None]).all()), "an offset reaches past row_end"
+ descending = torch.sort(out, dim=1, descending=True).values
+ assert not bool(((descending[:, :-1] == descending[:, 1:]) & (descending[:, 1:] >= 0)).any()), (
+ "duplicate offset within a row"
+ )
+ absolute = (out.long() + ref.row_starts[:, None]).clamp(0, ref.logits.shape[1] - 1)
+ picked = torch.gather(ref.logits, 1, absolute).masked_fill(~valid, float("-inf"))
+ got = torch.sort(picked, dim=1, descending=True).values
+ assert torch.equal(got, ref.values), "selected values are not the window's top-k"
+
+
+DEV = "cuda"
+COMPRESS_RATIO = 4
+BLOCK_TOPK = 512
+#: Pre-fill for every output buffer, so a slot the op failed to write shows up
+#: as an offset past `row_end` in `_compare`.
+SENTINEL = 12345
+
+
+def _logits(kind: str, rows: int, cols: int, gen: torch.Generator) -> torch.Tensor:
+ x = torch.randn(rows, cols, generator=gen, device=DEV)
+ # QSA's scores are relu(q.k) summed over heads: non-negative with ties at 0.
+ return torch.relu(x) * 3.0 if kind == "relu" else x
+
+
+def _build(spec: dict, seed: int):
+ gen = torch.Generator(device=DEV).manual_seed(seed)
+ rows, cols, index_topk = spec["rows"], spec["cols"], spec["index_topk"]
+ logits = _logits(spec["logits"], rows, cols, gen)
+
+ kind, start = spec["windows"]
+ if kind == "positions":
+ positions = torch.tensor(start, device=DEV)
+ else:
+ 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, :]
+ outside = (columns < row_starts[:, None]) | (columns >= row_ends[:, None])
+ logits.masked_fill_(outside, float("nan"))
+ indices = torch.full((rows, index_topk), SENTINEL, dtype=torch.int32, device=DEV)
+ return logits, row_starts, row_ends, indices, index_topk
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811
+ certifying(op)
+ args = _build(cell.spec, seed=sum(map(ord, cell.why)))
+ out = op(*args)
+ assert out is args[3], "the op writes `indices` and returns it"
+ _compare(out, op.reference(*args))
+
+
+#: (rows, score columns, replay window): tp1's 256-row verify step, then the default
+#: target's decode batches 1 and 128 and its 1024-row MTP7 verify step at the
+#: 262144-token envelope's 65536 columns. Every capture runs full windows (the
+#: histogram path); the replays move the bound, 128 rows down to the budget's
+#: copy-through shortcut.
+GRAPH_CASES = ((256, 2048, 1500), (1, 65536, 65535), (128, 65536, 511), (1024, 65536, 30000))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue.py
new file mode 100644
index 000000000000..a49258c73753
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue.py
@@ -0,0 +1,219 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the qwen4_exp_attn_decode_prologue entry: every cell bitwise the
+composition it replaces (`qwen4_exp_qk_norm_rope`, `qsa_paged_kv_store` or
+`qsa_paged_kv_store_nvfp4`, `qsa_paged_index_store`, `qsa_prefill_compress`) on every
+allocation either writes -- the projection, the index query, the K/V pool (with its block
+scales over NVFP4), the index plane, the position cache; a stale-read check under graph
+replay with PDL over both pool kinds.
+
+The NVFP4 cells lay the pools out as the engine serves them: K/V the two role planes of a
+`[pages, 2, Hkv, tpb, D/2]` e2m1 pool and its `[pages, 2, Hkv, tpb, D/16]` scale pool,
+and the index plane and the position cache two views of one padded page (page stride
+12 x 4096 + 192 bf16 elements, the position page at its tail)."""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qwen4_exp_attn_decode_prologue import ( # noqa: E501
+ qwen4_exp_attn_decode_prologue as op,
+)
+
+__extra_import_path__ = ["../../.."]
+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
+#: The NVFP4 engine's index page: 12 index buffers of tpb x 128, then the position page.
+INDEX_PAGE, POSITION_PAD = 12 * TPB * ID + 192, 192
+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 _edge_rows(proj) -> None:
+ """Zero, -0.0 and outlier blocks in the k heads (past the rotary prefix, where the
+ norm keeps zeros zero) and in the v heads, over the first rows (up to four)."""
+ k0, v0 = HQ * D, (HQ + HK) * D
+ edges = (
+ (0, slice(k0 + 128, k0 + 144), 0.0),
+ (1, slice(k0 + D + 160, k0 + D + 176), -0.0),
+ (2, slice(k0 + 200, k0 + 201), 900.0),
+ (0, slice(v0 + 16, v0 + 32), 0.0),
+ (1, slice(v0 + D, v0 + D + 16), torch.tensor([-0.0] * 8 + [0.0] * 8)),
+ (3, slice(v0 + 40, v0 + 41), -900.0),
+ (2, slice(v0 + D + 64, v0 + D + 80), 1e-30),
+ )
+ for r, cols, value in edges:
+ if r < proj.shape[0]:
+ proj[r, cols] = value
+
+
+def _rows(spec):
+ """(request slot, position) per row: one row a request, or under `verify` the
+ `rows_per_request` consecutive positions of each request from its first."""
+ per = spec.get("rows_per_request", 1)
+ return [(1 + r, p + i) for r, p in enumerate(spec["pos"]) for i in range(per)]
+
+
+def _build(spec, seed):
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ layout = _rows(spec)
+ pos = [p for _, p in layout]
+ rows = len(pos)
+ last = {}
+ for req, p in layout:
+ last[req] = max(last.get(req, 0), p)
+ slots = max(SLOTS, max(last) + 2)
+ pages = max(SLOTS * WIDTH // 2, sum(p // TPB + 1 for p in last.values()) + 64)
+ perm = torch.randperm(pages, generator=g, device=DEV).to(torch.int32)
+ block_table = torch.full((slots, WIDTH), -1, dtype=torch.int32, device=DEV)
+ used = 0
+ for req, p in last.items():
+ n = p // TPB + 1
+ block_table[req, :n] = perm[used : used + n]
+ used += n
+ request = torch.tensor([req for req, _ in layout], dtype=torch.int32, device=DEV)
+ if spec.get("unmapped") is not None:
+ request[spec["unmapped"]] = slots + 5
+ logical = torch.tensor(pos, dtype=torch.int64, device=DEV)
+ ids = torch.stack([logical, logical + 3, logical + 7]).to(torch.int32).contiguous()
+ buf = torch.empty(3, 1, rows + 5, dtype=torch.int32, device=DEV)
+ buf[:, 0, :rows] = ids
+ coords = torch.transpose(buf[:, 0, :rows], 0, 1)
+ 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)
+ if spec.get("edge"):
+ _edge_rows(proj)
+ 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)
+ ]
+ if spec.get("nvfp4"):
+ data = torch.randint(
+ 0, 256, (pages, 2, HK, TPB, D // 2), generator=g, device=DEV, dtype=torch.uint8
+ )
+ scales = torch.randint(
+ 0, 256, (pages, 2, HK, TPB, D // 16), generator=g, device=DEV, dtype=torch.uint8
+ )
+ plane = torch.randn(pages, INDEX_PAGE, generator=g, device=DEV).to(torch.bfloat16)
+ index_cache = plane[:, : TPB * ID].view(pages, TPB, 1, ID)
+ position_cache = plane[:, INDEX_PAGE - POSITION_PAD :].view(torch.int32)
+ position_cache = position_cache.view(pages, TPB, 3)
+ position_cache.copy_(
+ torch.randint(0, TABLE, (pages, TPB, 3), generator=g, device=DEV, dtype=torch.int32)
+ )
+ global_scales = spec.get("global_scales", (1.0, 1.0, 1.0))
+ kv = dict(
+ k_cache=data[:, 0], v_cache=data[:, 1], k_scale_cache=scales[:, 0],
+ v_scale_cache=scales[:, 1],
+ kv_scale_orig_quant=torch.tensor(global_scales, dtype=torch.float32, device=DEV),
+ ) # fmt: skip
+ else:
+ pool = 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
+ )
+ kv = dict(k_cache=pool[:, 0], v_cache=pool[:, 1])
+ if spec.get("scores"):
+ # Columns past a row's visible blocks are not written: a sentinel shows a stray one.
+ kv["index_scores"] = torch.full(
+ (rows, WIDTH * TPB // CR), -7.0, dtype=torch.float32, device=DEV
+ )
+ 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],
+ index_cache=index_cache, position_cache=position_cache, request_indices=request,
+ logical_positions=logical, block_table=block_table,
+ rows_per_request=spec.get("rows_per_request", 1), **kv, **GEOM,
+ ) # fmt: skip
+
+
+def _bytes(storage, device) -> torch.Tensor:
+ return torch.empty(0, dtype=torch.uint8, device=device).set_(storage)
+
+
+def _clone(args):
+ """The arguments over fresh copies of every allocation they view, each view rebuilt at
+ its offset and strides; returns them and the copies, keyed by the first argument
+ viewing each allocation."""
+ out, copies, keys = dict(args), {}, {}
+ for name, t in args.items():
+ if not isinstance(t, torch.Tensor):
+ continue
+ key = keys.setdefault(t.untyped_storage().data_ptr(), name)
+ if key not in copies:
+ copies[key] = _bytes(t.untyped_storage(), t.device).clone()
+ view = torch.empty(0, dtype=t.dtype, device=t.device)
+ view.set_(copies[key].untyped_storage(), t.storage_offset(), t.shape, t.stride())
+ out[name] = view
+ return out, copies
+
+
+def _assert_same(got, want, where="") -> None:
+ """Every allocation bitwise equal: what the op writes and everything it must not."""
+ assert got[1].keys() == want[1].keys()
+ for key, copy in got[1].items():
+ differ = (copy != want[1][key]).nonzero()
+ assert differ.numel() == 0, f"{where}{key}: {differ.shape[0]} bytes differ"
+
+
+@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[0])
+ got = _clone(args)
+ op(**got[0])
+ torch.cuda.synchronize()
+ _assert_same(got, want)
+
+
+@pytest.mark.parametrize(
+ "spec",
+ [
+ dict(pos=(8699, 8700)),
+ dict(pos=(8697, 8700), rows_per_request=4, nvfp4=True),
+ dict(pos=(8699, 8700), nvfp4=True, scores=True),
+ ],
+ ids=["bf16_pool", "nvfp4_pool_verify", "nvfp4_pool_scores"],
+)
+def test_reads_after_the_wait_under_graph_replay(spec) -> None:
+ """The projection is rewritten inside the graph right before the op; with PDL a read of
+ it before the wait would see the previous replay's row (in a verify step, also the
+ raw index keys of a group's earlier rows)."""
+ args = _build(spec, seed=23)
+ src = args["proj"].clone()
+ work, work_copies = _clone(args)
+
+ def step():
+ work["proj"].copy_(src * 1.0)
+ op(**work)
+
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ step()
+ torch.cuda.current_stream().wait_stream(side)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ step()
+ 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[0])
+ want[0]["proj"].copy_(src)
+ op.reference(**want[0])
+ _assert_same((work, work_copies), want, where=f"replay {trial}: ")
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py
new file mode 100644
index 000000000000..306742125fc6
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py
@@ -0,0 +1,154 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the qwen4_exp_qk_norm_rope entry."""
+
+import functools
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qwen4_exp_qk_norm_rope import ( # noqa: E501
+ CKPT,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qwen4_exp_qk_norm_rope import ( # noqa: E501
+ qwen4_exp_qk_norm_rope as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import ( # noqa: E402,F401 — path above; fixture by name
+ assert_within_ulp,
+ certifying,
+ validating,
+)
+
+assert torch.cuda.is_available(), "qwen4_exp_qk_norm_rope requires a CUDA device"
+
+THETA = 1e7
+EPS = 1e-6
+TABLE_ROWS = 262144 # the checkpoint's max_position_embeddings
+#: Spec keys that shape the test's inputs rather than the op call.
+_NOT_OP_ARGS = ("tokens", "max_position", "proj_width", "index_scratch_heads")
+
+
+def _seed(text: str) -> int:
+ return zlib.crc32(text.encode())
+
+
+@functools.lru_cache(maxsize=1)
+def _table() -> torch.Tensor:
+ """The target's rope table, built the way `post_load_weights` builds it."""
+ rot = CKPT["rotary_dim"]
+ exponent = torch.arange(0, rot, 2, dtype=torch.float64, device="cuda")
+ inv_freq = (1.0 / (THETA ** (exponent / rot))).to(torch.float32)
+ positions = torch.arange(TABLE_ROWS, dtype=torch.float32, device="cuda")
+ angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0)
+ return torch.cat([torch.cos(angles), torch.sin(angles)], dim=1).contiguous()
+
+
+def _distinct_rows(tokens: int, max_position: int, g: torch.Generator) -> torch.Tensor:
+ """A contiguous int32 `[3, T]` whose three rows differ at every token."""
+ row0 = torch.randint(0, max_position, (tokens,), generator=g, device="cuda")
+ d1 = torch.randint(1, max_position, (tokens,), generator=g, device="cuda")
+ d2 = torch.randint(1, max_position - 1, (tokens,), generator=g, device="cuda")
+ d2 = d2 + (d2 >= d1).long()
+ rows = torch.stack([row0, (row0 + d1) % max_position, (row0 + d2) % max_position])
+ return rows.to(torch.int32).contiguous()
+
+
+def _build(spec, seed):
+ """`(proj, scratch, kwargs)`: the fused projection row, the index scratch
+ whose first `index_heads` heads are the `index_q` view, and the op's
+ keyword arguments. Norm weights near 1 put q/k at their largest magnitude
+ under the Gemma `1 + w`."""
+ g = torch.Generator(device="cuda").manual_seed(seed)
+ t = spec["tokens"]
+ proj = torch.randn(t, spec["proj_width"], generator=g, device="cuda").to(torch.bfloat16)
+ scratch = torch.randn(
+ t, spec["index_scratch_heads"], spec["index_head_dim"], generator=g, device="cuda"
+ ).to(torch.bfloat16)
+
+ def weight(n):
+ return (torch.randn(n, generator=g, device="cuda") * 0.1 + 1.0).to(torch.bfloat16)
+
+ positions = _distinct_rows(t, spec.get("max_position", 32768), g)
+ kwargs = {k: v for k, v in spec.items() if k not in _NOT_OP_ARGS}
+ kwargs.update(
+ index_q=scratch[:, : spec["index_heads"]],
+ position_ids=positions,
+ cos_sin=_table(),
+ q_weight=weight(spec["head_dim"]),
+ k_weight=weight(spec["head_dim"]),
+ index_q_weight=weight(spec["index_head_dim"]),
+ eps=EPS,
+ )
+ return proj, scratch, kwargs
+
+
+def _deviation(out: torch.Tensor, ref: torch.Tensor) -> str:
+ ulp = 2.0**-8
+ o, r = out.float(), ref.float()
+ row = r.abs().amax(dim=-1, keepdim=True).clamp_min(1e-9)
+ elt = ((o - r).abs() / row).max().item() / ulp
+ rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ulp
+ return f"max abs {(o - r).abs().max().item():.6f}, elt {elt:.2f} ulp, rms {rms:.3f} ulp"
+
+
+def _main_cols(spec) -> int:
+ return (spec["num_heads_q"] + spec["num_heads_k"]) * spec["head_dim"]
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811
+ certifying(op)
+ spec = cell.spec
+ proj, scratch, kwargs = _build(spec, seed=_seed(cell.why))
+ before, scratch_before = proj.clone(), scratch.clone()
+ ref_proj, ref_iq = op.reference(proj, **kwargs)
+ assert torch.equal(proj, before), "the reference mutated its input"
+ out_proj, out_iq = op(proj, **kwargs)
+ assert out_proj is proj and out_iq is kwargs["index_q"]
+ main = _main_cols(spec)
+ t = spec["tokens"]
+ flat_ref = ref_iq.reshape(t, -1)
+ flat_out = out_iq.reshape(t, -1)
+ print(f"\n {cell.why[:50]}: main {_deviation(proj[:, :main], ref_proj[:, :main])}")
+ print(f" {cell.why[:50]}: index {_deviation(flat_out, flat_ref)}")
+ # fp32 throughout and one rounding, against an fp64 reference over the same
+ # table: the two can differ by one output rounding step, which is up to 2
+ # ulp of the row scale (an element just above a power of two in a row whose
+ # largest magnitude sits just below the next one). Measured: max 1.48 ulp
+ # element-wise, relative RMS <= 0.011 ulp, so the RMS gate is tight.
+ assert_within_ulp(proj[:, :main], ref_proj[:, :main], element_ulp=2.0, rms_ulp=0.25)
+ assert_within_ulp(flat_out, flat_ref, element_ulp=2.0, rms_ulp=0.25)
+ assert torch.equal(proj[:, main:], before[:, main:]), "a column past the main k heads moved"
+ heads = spec["index_heads"]
+ assert torch.equal(scratch[:, heads:], scratch_before[:, heads:]), (
+ "index_q wrote past its heads"
+ )
+
+
+def test_cuda_graph_replay_reads_fresh_inputs() -> None:
+ """Captured once at M=4, replayed after the inputs are rewritten in place:
+ the replay equals an eager call on the new inputs, bit for bit."""
+ spec = dict(tokens=4, **CKPT)
+ proj, scratch, kwargs = _build(spec, seed=_seed("graph"))
+ proj_new, _, kwargs_new = _build(spec, seed=_seed("graph-new"))
+ eager_proj, eager_scratch = proj_new.clone(), scratch.clone()
+ eager_kwargs = dict(kwargs_new, index_q=eager_scratch[:, : CKPT["index_heads"]])
+ op(eager_proj, **eager_kwargs)
+ stream = torch.cuda.Stream()
+ stream.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(stream):
+ op(proj.clone(), **dict(kwargs, index_q=scratch.clone()[:, : CKPT["index_heads"]]))
+ torch.cuda.current_stream().wait_stream(stream)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ op(proj, **kwargs)
+ proj.copy_(proj_new)
+ for name in ("position_ids", "q_weight", "k_weight", "index_q_weight"):
+ kwargs[name].copy_(kwargs_new[name])
+ graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(proj, eager_proj)
+ assert torch.equal(kwargs["index_q"], eager_kwargs["index_q"])
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_spec_rollback_commit.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_spec_rollback_commit.py
new file mode 100644
index 000000000000..2631822136f5
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_spec_rollback_commit.py
@@ -0,0 +1,319 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU tests for the spec_rollback_commit variant: `spec_rollback_stage` and
+`spec_rollback_commit`, every cell through the entry's own reference, bitwise (both are data
+movement).
+
+A cell is one verify step at the call site's layout: `n` generation requests of `p` rows, each
+request's rows at consecutive positions on pages of their own, int32 `accepted` / `slots` /
+`ids` and int64 `pages` / `within` as `SpeculativeRollback` hands them. The index-K planes are
+`[pages, 32, 1, 128]` bf16 views addressed through `plane_table`: separate buffers in the tp1
+cells, the KV cache manager's coalesced pool (12 planes side by side in a 49152-element page) at
+a tep4 rank's measured page count in the default target's. Request 0's rows straddle the pool's first
+and last page and it accepts one row, so its rejected rows are restored on the last page, past
+2^31 elements from a plane's base. The pools are filled with a NaN pattern `randn` never draws
+and the trunk's writes are put at every row's slot, so a whole-pool count shows the commit
+wrote nothing else; the n-gram, conv and position pools are compared whole.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_rollback_planes import ( # noqa: E501
+ plane_table,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.spec_rollback_commit import ( # noqa: E501
+ ROW,
+ spec_rollback_commit,
+ spec_rollback_stage,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "spec_rollback_commit requires a CUDA device"
+
+DEV = "cuda"
+TPB = 32
+#: The tp1 cells' PLE geometry and plane pool, where a cell does not name its own.
+TP1_GEOMETRY = dict(context=2, channels=640, state_len=3, layout="separate", pool_pages=600)
+#: `max_seq_len` under MTP: the last position a verify row can hold is one less.
+MTP_MAX_SEQ_LEN = 16384
+#: The vocabulary the ids are drawn from.
+VOCAB = 248320
+#: bf16 NaN with a payload: the index pool's fill.
+SENTINEL = 0x7FA5
+#: The position plane's fill; positions stay below 2^24.
+POSITION_FILL = 0x7FA5A5A5
+#: One plane's page: `tokens_per_block` 128-wide rows.
+PLANE_PAGE = TPB * ROW
+
+
+def _seed(why: str) -> int:
+ return zlib.crc32(why.encode()) % 2**31
+
+
+def _geometry(spec: dict) -> dict:
+ g = dict(TP1_GEOMETRY, **spec)
+ g.setdefault("state_slots", 2 * spec["n"] + 4)
+ return g
+
+
+def _slots(n: int, state_slots: int, gen: torch.Generator) -> torch.Tensor:
+ """`n` distinct recurrent slots; at an engine-sized pool its first and last among them."""
+ if state_slots < n + 2 or n < 2:
+ picked = torch.randperm(state_slots - 1, generator=gen)[:n] + 1
+ else:
+ inner = torch.randperm(state_slots - 2, generator=gen)[: n - 2] + 1
+ picked = torch.cat([torch.tensor([0, state_slots - 1]), inner])
+ picked = picked[torch.randperm(n, generator=gen)]
+ return picked.to(DEV, torch.int32)
+
+
+def _ple_pools(g: dict, gen: torch.Generator):
+ ngram_pool = torch.randint(0, 1 << 40, (g["state_slots"], g["context"]), generator=gen)
+ conv_pool = torch.randn(g["state_slots"], g["channels"], g["state_len"], generator=gen)
+ return ngram_pool.to(DEV), conv_pool.to(DEV, torch.bfloat16)
+
+
+def _bits(t: torch.Tensor) -> torch.Tensor:
+ return t.view(torch.int16) if t.dtype == torch.bfloat16 else t
+
+
+def _equal(got: torch.Tensor, want: torch.Tensor, what: str) -> None:
+ assert got.dtype == want.dtype and got.shape == want.shape, (what, got.shape, want.shape)
+ differ = (_bits(got) != _bits(want)).reshape(got.shape[0], -1).any(dim=1)
+ assert not bool(differ.any()), f"{what}: {int(differ.sum())} of {got.shape[0]} rows differ"
+
+
+class _Planes:
+ """One rank's index-K planes, every element at `SENTINEL` until a slot is written."""
+
+ def __init__(self, count: int, pages: int, layout: str):
+ if layout == "separate":
+ self.buffers = [
+ torch.empty(pages, TPB, 1, ROW, dtype=torch.bfloat16, device=DEV)
+ for _ in range(count)
+ ]
+ self.views = list(self.buffers)
+ else:
+ assert layout == "coalesced", layout
+ page = torch.empty(pages, count * PLANE_PAGE, dtype=torch.bfloat16, device=DEV)
+ self.buffers = [page]
+ self.views = [
+ page[:, k * PLANE_PAGE : (k + 1) * PLANE_PAGE].view(pages, TPB, 1, ROW)
+ for k in range(count)
+ ]
+ self.fill()
+
+ def fill(self) -> None:
+ for b in self.buffers:
+ b.view(torch.int16).fill_(SENTINEL)
+
+ def written(self) -> int:
+ """Elements anywhere in the pool that no longer hold the fill."""
+ total = 0
+ for b in self.buffers:
+ for chunk in b.view(torch.int16).reshape(b.shape[0], -1).split(8192):
+ total += int((chunk != SENTINEL).sum())
+ return total
+
+ def rows(self, pages: torch.Tensor, within: torch.Tensor) -> torch.Tensor:
+ return torch.stack([v[pages, within].reshape(-1, ROW) for v in self.views])
+
+
+class _Step:
+ """One verify step's staged state, the pools it commits into, and the trunk's writes."""
+
+ def __init__(self, spec: dict, seed: int):
+ g = _geometry(spec)
+ self.n, self.p, self.count = spec["n"], spec["p"], spec["planes"]
+ self.planes = _Planes(self.count, g["pool_pages"], g["layout"])
+ self.position = torch.empty(g["pool_pages"], TPB, 3, dtype=torch.int32, device=DEV)
+ self.ptrs, self.strides = plane_table(self.planes.views)
+ gen = torch.Generator().manual_seed(seed)
+ self.ngram_pool, self.conv_pool = _ple_pools(g, gen)
+ self.g = g
+ n, p, rows = self.n, self.p, self.n * self.p
+ self.accepted = torch.empty(n, dtype=torch.int32, device=DEV)
+ self.slots = torch.empty(n, dtype=torch.int32, device=DEV)
+ self.pages = torch.empty(rows, dtype=torch.int64, device=DEV)
+ self.within = torch.empty(rows, dtype=torch.int64, device=DEV)
+ self.stream = torch.empty(n, g["context"] + p, dtype=torch.int64, device=DEV)
+ self.conv_prev = torch.empty(
+ n, g["channels"], g["state_len"], dtype=torch.bfloat16, device=DEV
+ )
+ self.conv_x = torch.empty(n, p, g["channels"], dtype=torch.bfloat16, device=DEV)
+ self.position_old = torch.empty(rows, 3, dtype=torch.int32, device=DEV)
+ self.index_old = torch.empty(self.count, rows, ROW, dtype=torch.bfloat16, device=DEV)
+ self.draw(gen)
+
+ def draw(self, gen: torch.Generator) -> None:
+ """Draw every staged input into its buffer (in place, so a captured graph re-reads it)
+ and put the trunk's writes at the rows' slots of a freshly filled index and position
+ pool."""
+ g, n, p = self.g, self.n, self.p
+ pool_pages = g["pool_pages"]
+ accepted = torch.randint(1, p + 1, (n,), generator=gen)
+ accepted[0] = 1
+ if n > 1:
+ accepted[1] = p
+ self.accepted.copy_(accepted)
+ self.slots.copy_(_slots(n, g["state_slots"], gen))
+
+ # Request 0 straddles its first two blocks, mapped to the pool's first and last page;
+ # every other block of every request gets a page of its own.
+ starts = torch.randint(0, MTP_MAX_SEQ_LEN - p + 1, (n,), generator=gen)
+ starts[0] = TPB - p // 2
+ positions = starts[:, None] + torch.arange(p)[None, :]
+ blocks = positions // TPB
+ free = torch.randperm(pool_pages - 2, generator=gen) + 1
+ page_of = {}
+ pages = torch.empty(n, p, dtype=torch.int64)
+ for i in range(n):
+ for t in range(p):
+ key = (i, int(blocks[i, t]))
+ if key not in page_of:
+ if i == 0:
+ page_of[key] = 0 if int(blocks[i, t]) == 0 else pool_pages - 1
+ else:
+ page_of[key] = int(free[len(page_of)])
+ pages[i, t] = page_of[key]
+ self.pages.copy_(pages.reshape(-1))
+ self.within.copy_((positions % TPB).reshape(-1))
+
+ rows = n * p
+ self.stream.copy_(torch.randint(0, 1 << 40, tuple(self.stream.shape), generator=gen))
+ self.conv_prev.copy_(torch.randn(tuple(self.conv_prev.shape), generator=gen))
+ self.conv_x.copy_(torch.randn(tuple(self.conv_x.shape), generator=gen))
+ self.position_old.copy_(torch.randint(0, 1 << 20, (rows, 3), generator=gen))
+ self.index_old.copy_(torch.randn(self.count, rows, ROW, generator=gen))
+
+ # What the trunk wrote at every row: kept where the request accepted the row.
+ self.planes.fill()
+ self.position.fill_(POSITION_FILL)
+ trunk = torch.randn(self.count, rows, ROW, generator=gen).to(DEV, torch.bfloat16)
+ for k, view in enumerate(self.planes.views):
+ view[self.pages, self.within] = trunk[k].reshape(rows, 1, ROW)
+ self.position[self.pages, self.within] = torch.randint(
+ 0, 1 << 20, (rows, 3), generator=gen
+ ).to(DEV, torch.int32)
+
+ def call_args(self, index_old=None, conv_prev=None) -> tuple:
+ return (
+ self.accepted,
+ self.slots,
+ self.ngram_pool,
+ self.stream,
+ self.conv_pool,
+ self.conv_prev if conv_prev is None else conv_prev,
+ self.conv_x,
+ self.position,
+ self.position_old,
+ self.pages,
+ self.within,
+ self.ptrs,
+ self.strides,
+ self.index_old if index_old is None else index_old,
+ )
+
+ def reference(self):
+ return spec_rollback_commit.reference(
+ self.accepted,
+ self.slots,
+ self.ngram_pool,
+ self.stream,
+ self.conv_pool,
+ self.conv_prev,
+ self.conv_x,
+ self.position,
+ self.position_old,
+ self.pages,
+ self.within,
+ self.planes.views,
+ self.index_old,
+ )
+
+ def snapshot(self):
+ return self.ngram_pool.clone(), self.conv_pool.clone(), self.position.clone()
+
+ def check(self, ref, before) -> None:
+ """The named slots hold `ref`, bit for bit, and nothing else in any pool moved."""
+ ngram, conv, pos, index = ref
+ ngram_before, conv_before, pos_before = before
+ slots = self.slots.long()
+ want = ngram_before.clone()
+ want[slots] = ngram
+ _equal(self.ngram_pool, want, "n-gram pool")
+ want = conv_before.clone()
+ want[slots] = conv
+ _equal(self.conv_pool, want, "PLE conv pool")
+ want = pos_before.clone()
+ want[self.pages, self.within] = pos
+ _equal(self.position, want, "position plane")
+ _equal(self.planes.rows(self.pages, self.within), index, "index-K planes")
+ rows = self.n * self.p
+ assert self.planes.written() == self.count * rows * ROW, (
+ "the commit wrote index-K pool elements outside the rows' slots"
+ )
+
+
+@pytest.mark.parametrize(
+ "cell", spec_rollback_stage.CELLS, ids=[c.why for c in spec_rollback_stage.CELLS]
+)
+def test_stage_cells(cell, certifying) -> None: # noqa: F811
+ certifying(spec_rollback_stage)
+ spec = cell.spec
+ n, p = spec["n"], spec["p"]
+ g = _geometry(spec)
+ gen = torch.Generator().manual_seed(_seed(cell.why))
+ ngram_pool, conv_pool = _ple_pools(g, gen)
+ slots = _slots(n, g["state_slots"], gen)
+ ids = torch.randint(0, VOCAB, (n * p,), generator=gen).to(DEV, torch.int32)
+ stream, conv_prev = spec_rollback_stage(ngram_pool, conv_pool, slots, ids, p)
+ want_stream, want_prev = spec_rollback_stage.reference(ngram_pool, conv_pool, slots, ids, p)
+ _equal(stream, want_stream, "stream")
+ _equal(conv_prev, want_prev, "conv_prev")
+
+
+@pytest.mark.parametrize(
+ "cell", spec_rollback_commit.CELLS, ids=[c.why for c in spec_rollback_commit.CELLS]
+)
+def test_commit_cells(cell, certifying) -> None: # noqa: F811
+ certifying(spec_rollback_commit)
+ step = _Step(cell.spec, _seed(cell.why))
+ ref = step.reference()
+ before = step.snapshot()
+ spec_rollback_commit(*step.call_args())
+ torch.cuda.synchronize()
+ step.check(ref, before)
+
+
+def test_commit_graph_replay_rereads_its_inputs() -> None:
+ """One graph captured at the call site's 1024 rows (128 requests x 8, the 10240 x 9 conv
+ state over 1874 slots), replayed after every staged input, slot and page is redrawn in place:
+ what the decode target's captured verify step does each iteration."""
+ spec = dict(spec_rollback_commit.CELLS[-1].spec, pool_pages=4096)
+ step = _Step(spec, 4242)
+ args = step.call_args()
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ spec_rollback_commit(*args)
+ torch.cuda.current_stream().wait_stream(side)
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ spec_rollback_commit(*args)
+ for seed in (7, 8):
+ step.draw(torch.Generator().manual_seed(seed))
+ ref = step.reference()
+ before = step.snapshot()
+ graph.replay()
+ torch.cuda.synchronize()
+ step.check(ref, before)
+
+
+_GUARD_SPEC = dict(n=8, p=4, planes=12)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_verify_step_setup.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_verify_step_setup.py
new file mode 100644
index 000000000000..ea58f90dda07
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_verify_step_setup.py
@@ -0,0 +1,139 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test: `verify_step_setup` against the composed ops of the decode target's step setup
+(`pos3` copy, `cu_seqlens_from_seq_lens`, `SpeculativeRollback.begin`'s slot arithmetic and
+`qsa_rollback_gather`), bitwise, eagerly and replayed from a CUDA graph."""
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_rollback_planes import ( # noqa: E501
+ plane_table,
+ qsa_rollback_gather,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.verify_step_setup import ( # noqa: E501
+ verify_step_setup,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.cu_seqlens_from_seq_lens import ( # noqa: E501
+ cu_seqlens_from_seq_lens,
+)
+
+assert torch.cuda.is_available(), "the verify step setup test requires a CUDA device"
+
+DEV = "cuda"
+TPB = 32
+PAGES = 96
+LAYERS = 12
+MAX_BLOCKS = 6
+
+
+def _case(batch, p, seed, negative=False):
+ g = torch.Generator().manual_seed(seed)
+ rows = batch * p
+ full = torch.randint(0, 1 << 20, (3, 1, rows + 5), generator=g).to(DEV, torch.int32)
+ position_ids = full[..., :rows]
+ seq_lens = torch.full((batch,), p, dtype=torch.int32, device=DEV)
+ slots = torch.randint(0, 64, (batch,), generator=g).to(DEV, torch.int32)
+ req = torch.arange(batch).repeat_interleave(p).to(DEV, torch.int32)
+ starts = torch.randint(0, MAX_BLOCKS * TPB - p - 20, (batch,), generator=g)
+ logical = (starts.repeat_interleave(p) + torch.arange(p).repeat(batch)).to(DEV, torch.int64)
+ if negative:
+ logical[:p] -= starts[0] + 40 # a padded request: negative positions, one page back
+ block_table = torch.randint(0, PAGES, (batch + 3, MAX_BLOCKS + 2), generator=g)
+ block_table = block_table.to(DEV, torch.int32)[:, :MAX_BLOCKS]
+ planes = [
+ torch.randn(PAGES, TPB, 1, 128, generator=g).to(DEV, torch.bfloat16) for _ in range(LAYERS)
+ ]
+ return position_ids, seq_lens, slots, req, logical, block_table, planes
+
+
+def _reference(position_ids, seq_lens, slots, req, logical, block_table, ptrs, strides):
+ pos3 = torch.empty([3, position_ids.shape[-1]], dtype=torch.int32, device=DEV)
+ pos3.copy_(torch.reshape(position_ids, [3, -1]))
+ qsl = cu_seqlens_from_seq_lens(seq_lens, torch.int64)
+ positions = logical.to(torch.long)
+ pages = block_table[req.to(torch.long), positions // TPB].to(torch.long)
+ within = positions % TPB
+ index = qsa_rollback_gather(ptrs, strides, pages, within, LAYERS)
+ return pos3, qsl, slots.to(torch.long), pages, within, index
+
+
+def _bits(t):
+ return t.view(torch.int16) if t.dtype == torch.bfloat16 else t
+
+
+@pytest.mark.parametrize("negative", [False, True], ids=["positions", "negative_positions"])
+@pytest.mark.parametrize("batch,p", [(1, 4), (17, 8), (255, 2)])
+def test_setup_matches_composed_ops(batch, p, negative):
+ position_ids, seq_lens, slots, req, logical, block_table, planes = _case(
+ batch, p, batch * 7 + p, negative
+ )
+ ptrs, strides = plane_table(planes)
+ ref = _reference(position_ids, seq_lens, slots, req, logical, block_table, ptrs, strides)
+ got = verify_step_setup(
+ position_ids, seq_lens, slots, req, logical, block_table, TPB, ptrs, strides, LAYERS
+ )
+ for a, b in zip(got, ref):
+ assert a.dtype == b.dtype and a.shape == b.shape
+ assert torch.equal(_bits(a), _bits(b))
+
+
+def test_setup_zero_length_rows_in_segmentation():
+ position_ids, seq_lens, slots, req, logical, block_table, planes = _case(6, 2, 5)
+ seq_lens.copy_(torch.tensor([2, 2, 2, 2, 2, 2], dtype=torch.int32))
+ seq_lens[3:] = 0
+ ptrs, strides = plane_table(planes)
+ got = verify_step_setup(
+ position_ids, seq_lens, slots, req, logical, block_table, TPB, ptrs, strides, LAYERS
+ )
+ assert torch.equal(got[1], cu_seqlens_from_seq_lens(seq_lens, torch.int64))
+
+
+def test_setup_cuda_graph_follows_inputs():
+ position_ids, seq_lens, slots, req, logical, block_table, planes = _case(8, 4, 11)
+ ptrs, strides = plane_table(planes)
+ args = (position_ids, seq_lens, slots, req, logical, block_table, TPB, ptrs, strides, LAYERS)
+ verify_step_setup(*args)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ out = verify_step_setup(*args)
+ g = torch.Generator().manual_seed(12)
+ for _ in range(3):
+ logical.add_(torch.randint(0, 3, logical.shape, generator=g).to(DEV))
+ position_ids.add_(1)
+ for plane in planes:
+ plane.copy_(torch.randn(plane.shape, generator=g).to(DEV, torch.bfloat16))
+ graph.replay()
+ torch.cuda.synchronize()
+ ref = _reference(position_ids, seq_lens, slots, req, logical, block_table, ptrs, strides)
+ for a, b in zip(out, ref):
+ assert torch.equal(_bits(a), _bits(b))
+
+
+@pytest.mark.parametrize(
+ "cell", verify_step_setup.CELLS, ids=[c.why for c in verify_step_setup.CELLS]
+)
+def test_setup_cells(cell) -> None:
+ """Each cell against the docstring's outputs written out in torch indexing."""
+ batch, p, num_planes = cell.spec["batch"], cell.spec["p"], cell.spec["planes"]
+ position_ids, seq_lens, slots, req, logical, block_table, planes = _case(batch, p, 21 * batch)
+ planes = planes[:num_planes]
+ ptrs, strides = plane_table(planes)
+ args = (position_ids, seq_lens, slots, req, logical, block_table, TPB, ptrs, strides,
+ num_planes) # fmt: skip
+ verify_step_setup.is_valid(*args)
+ pos3, qsl, slots_long, pages, within, index = verify_step_setup(*args)
+ rows = batch * p
+ assert torch.equal(pos3, position_ids.reshape(3, rows).to(torch.int32))
+ want_qsl = torch.zeros(batch + 1, dtype=torch.int64, device=DEV)
+ want_qsl[1:] = torch.cumsum(seq_lens.to(torch.int64), 0)
+ assert torch.equal(qsl, want_qsl)
+ assert torch.equal(slots_long, slots.to(torch.int64))
+ want_pages = block_table[req.long(), torch.div(logical, TPB, rounding_mode="floor")].long()
+ want_within = torch.remainder(logical, TPB)
+ assert torch.equal(pages, want_pages)
+ assert torch.equal(within, want_within)
+ want_index = torch.stack(
+ [plane[want_pages, want_within].reshape(rows, 128) for plane in planes]
+ )
+ assert torch.equal(_bits(index), _bits(want_index))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/embedding/test_modeling_v2_ngram_hash.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/embedding/test_modeling_v2_ngram_hash.py
new file mode 100644
index 000000000000..ce804c238fc9
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/embedding/test_modeling_v2_ngram_hash.py
@@ -0,0 +1,99 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the default target's ngram_hash variant (models/qwen4_exp/ops/embedding/ngram_hash).
+
+Exact integer arithmetic, gated bitwise against the entry's scalar reference.
+Beyond the cells: the EOS reset that `ngram_token_history`'s EOS seed relies
+on.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.embedding.ngram_hash import ( # noqa: E501
+ CKPT_CONTEXT_LEN,
+ CKPT_EOS_TOKEN_ID,
+ CKPT_HEAD_OFFSETS,
+ CKPT_HEAD_VOCAB_SIZES,
+ CKPT_MULTIPLIERS,
+ CKPT_VOCAB_SIZE,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.embedding.ngram_hash import ( # noqa: E501
+ ngram_hash as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "ngram_hash requires a CUDA device"
+
+
+def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
+ """Bitwise: exact integer arithmetic, no tolerance to loosen."""
+ assert out.dtype == ref.dtype == torch.int64, (out.dtype, ref.dtype)
+ assert out.shape == ref.shape, (out.shape, ref.shape)
+ assert out.is_contiguous()
+ assert torch.equal(out, ref), f"{int((out != ref).sum())} indices differ"
+
+
+DEV = "cuda"
+EOS = CKPT_EOS_TOKEN_ID
+NGRAM_HEADS = len(CKPT_HEAD_VOCAB_SIZES)
+PADDED_TABLE_ROWS = 320001536 # sum(head_vocab_sizes) rounded up to a multiple of 128
+
+
+def _constants():
+ return (
+ torch.tensor(CKPT_MULTIPLIERS, dtype=torch.long, device=DEV),
+ torch.tensor(CKPT_HEAD_VOCAB_SIZES, dtype=torch.long, device=DEV),
+ torch.tensor(CKPT_HEAD_OFFSETS, dtype=torch.long, device=DEV),
+ )
+
+
+def _tokens(batch, length, seed, eos_positions=(), lead_eos=True, eos_density=0.0):
+ g = torch.Generator().manual_seed(seed)
+ ids = torch.randint(0, CKPT_VOCAB_SIZE, (batch, length), generator=g, dtype=torch.long)
+ if eos_density > 0.0:
+ ids[torch.rand(batch, length, generator=g) < eos_density] = EOS
+ if lead_eos:
+ ids[:, : min(CKPT_CONTEXT_LEN, length)] = EOS
+ for position in eos_positions:
+ ids[:, position] = EOS
+ return ids.to(DEV)
+
+
+def _hash(ids, num_query_tokens):
+ return op(ids, *_constants(), EOS, num_query_tokens)
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
+ certifying(op)
+ spec = cell.spec
+ ids = _tokens(
+ spec["batch"],
+ spec["length"],
+ seed=zlib.crc32(cell.why.encode()),
+ eos_positions=spec.get("eos_positions", ()),
+ lead_eos=spec.get("lead_eos", True),
+ eos_density=spec.get("eos_density", 0.0),
+ )
+ before = ids.clone()
+ args = (ids, *_constants(), EOS, spec["num_query_tokens"])
+ out = op(*args)
+ assert out.shape == (spec["batch"], spec["num_query_tokens"], NGRAM_HEADS)
+ assert out.data_ptr() != ids.data_ptr()
+ assert torch.equal(ids, before), "the history is read-only"
+ _compare(out, op.reference(*args))
+
+
+def test_eos_restarts_the_window() -> None:
+ """One past an EOS hashes as a sequence's first token, and only there."""
+ token = 4242
+ fresh = torch.tensor([[EOS] * CKPT_CONTEXT_LEN + [token]], device=DEV)
+ after_eos = torch.tensor([[11, 22, 33, EOS, token]], device=DEV)
+ no_eos = torch.tensor([[11, 22, 33, 44, token]], device=DEV)
+ assert torch.equal(_hash(fresh, 1), _hash(after_eos, 1))
+ assert not torch.equal(_hash(fresh, 1), _hash(no_eos, 1))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_bf16_gemv_pdl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_bf16_gemv_pdl.py
new file mode 100644
index 000000000000..6f2cd3278b41
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_bf16_gemv_pdl.py
@@ -0,0 +1,102 @@
+# 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.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.bf16_gemv_pdl import ( # noqa: E501
+ bf16_gemv_pdl as op,
+)
+
+__extra_import_path__ = ["../../.."]
+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(
+ "rows, dims, last_n",
+ [
+ (1, [2560, 6144, 2560, 640, 2560], 16480),
+ # Four rows (an MTP-3 verify step) at a tp4/tep4 rank's K=1536.
+ (4, [1536, 1536, 1536, 1536], 2560),
+ ],
+ ids=["tp1-one-row", "tp4-four-rows"],
+)
+def test_dependent_chain_under_graph_replay(rows, dims, last_n) -> None:
+ """Decode runs this under CUDA graphs, back to back with other kernels. Each
+ link of the chain reads the output of the one before it: a GEMV that read
+ its activation before `griddepcontrol.wait` would see the previous replay's
+ values, and the comparison below would fail on the second replay."""
+ pdl = True
+ # K of each link; N is the next link's K.
+ weights = [
+ _build(rows, n, k, seed=i)[1] for i, (k, n) in enumerate(zip(dims, dims[1:] + [last_n]))
+ ]
+ x = torch.randn(rows, dims[0], device="cuda").to(torch.bfloat16)
+
+ def chain():
+ h = x
+ 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
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_gemv.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_gemv.py
new file mode 100644
index 000000000000..368839e703e2
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_gemv.py
@@ -0,0 +1,178 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test: `fp8_gemv` / `fp8_gemv_conv1d` / `fp8_gemv_conv1d_verify` against an fp32
+reference over the dequantized weight (128x128-block e4m3), per cell; the verify entry also
+bitwise against the two-launch path it replaces (`fp8_gemv` + `causal_conv1d_triton_update`);
+strided activations."""
+
+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.ops.gemm.fp8_gemv import ( # noqa: E501
+ dequantize_fp8_block,
+ fp8_gemv_conv1d_op,
+ fp8_gemv_conv1d_verify_op,
+ fp8_gemv_op,
+ pack_weight,
+ quantize_fp8_block,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.causal_conv1d_triton_update import ( # noqa: E501
+ causal_conv1d_triton_update,
+)
+
+assert torch.cuda.is_available(), "fp8_gemv requires a CUDA device"
+
+DEV = "cuda"
+
+
+def _weight(n, k, seed):
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ w = torch.randn(n, k, generator=g, device=DEV) * k**-0.5
+ w[:, 7] *= 30.0 # outlier column
+ w[: min(n, 128), 128:256] *= 1e-4 # a tiny block: subnormal codes
+ codes, scale = quantize_fp8_block(w.to(torch.bfloat16))
+ return codes, scale, pack_weight(codes, scale)
+
+
+def _close(y, ref, x, deq) -> None:
+ """Within one bf16 rounding of the fp32 reference (plus fp32 summation noise)."""
+ with true_fp32_matmul():
+ absref = x.float().abs() @ deq.abs().t()
+ ulp = torch.exp2(torch.floor(torch.log2(ref.abs().clamp_min(1e-30))) - 7)
+ excess = ((y.float() - ref).abs() - 0.5 * ulp).clamp_min(0) / absref.clamp_min(1e-30)
+ assert excess.max().item() < 1e-5, f"error beyond one rounding: {excess.max().item():.2e}"
+
+
+@pytest.mark.parametrize("cell", fp8_gemv_op.CELLS, ids=[c.why[:40] for c in fp8_gemv_op.CELLS])
+def test_gemv_cells(cell) -> None:
+ n, k, m = cell.spec["n"], cell.spec["k"], cell.spec["m"]
+ codes, scale, packed = _weight(n, k, n + m)
+ g = torch.Generator(device=DEV).manual_seed(m)
+ # A column slice of a wider row, as the target's split projections hand over.
+ wide = torch.randn(m, k + 256, generator=g, device=DEV).to(torch.bfloat16)
+ x = wide[:, :k]
+ y = fp8_gemv_op(x, packed)
+ assert y.shape == (m, n) and y.dtype == torch.bfloat16
+ _close(y, fp8_gemv_op.reference(x, codes, scale), x, dequantize_fp8_block(codes, scale))
+
+
+@pytest.mark.parametrize(
+ "cell", fp8_gemv_conv1d_op.CELLS, ids=[c.why[:40] for c in fp8_gemv_conv1d_op.CELLS]
+)
+def test_conv1d_cells(cell) -> None:
+ n, k, conv_dim, pool, pad = 16480, 2560, 10240, 64, -1
+ codes, scale, packed = _weight(n, k, 3)
+ deq = dequantize_fp8_block(codes, scale)
+ slots = list(cell.spec["slots"])
+ idx = torch.tensor(slots, dtype=torch.int32, device=DEV)
+ g = torch.Generator(device=DEV).manual_seed(11)
+ backing = torch.randn(pool, conv_dim, 4, generator=g, device=DEV).to(torch.bfloat16)
+ conv_state = backing[:, :, 1:] # a strided pool view
+ conv_weight = torch.randn(conv_dim, 4, generator=g, device=DEV).to(torch.bfloat16)
+ x = torch.randn(len(slots), k, generator=g, device=DEV).to(torch.bfloat16)
+ before = conv_state.clone()
+ with true_fp32_matmul():
+ proj = (x.float() @ deq.t()).to(torch.bfloat16)
+ want, want_state = proj.clone(), conv_state.clone()
+ for r, s in enumerate(slots):
+ if s == pad:
+ continue
+ window = torch.cat([conv_state[s].float(), proj[r, :conv_dim].float()[:, None]], 1)
+ c = (window * conv_weight.float()).sum(1)
+ want[r, :conv_dim] = (c * torch.sigmoid(c)).to(torch.bfloat16)
+ want_state[s] = window[:, 1:].to(torch.bfloat16)
+ out = fp8_gemv_conv1d_op(x, packed, conv_state, conv_weight, idx, conv_dim, pad)
+ # The projection may round one ulp differently than the reference's fp32 matmul; a flip
+ # carries through the conv. Nearly every output is bitwise the reference's.
+ assert (out == want).float().mean().item() > 0.999
+ assert ((out.float() - want.float()).abs() / want.float().abs().clamp_min(1e-2)).max() < 0.02
+ assert ((conv_state.float() - want_state.float()).abs()).max().item() < 0.02
+ others = [s for s in range(pool) if s not in slots]
+ assert torch.equal(conv_state[others], before[others]), "other slots untouched"
+
+
+def _verify_args(spec, seed):
+ """`fp8_gemv_conv1d_verify`'s arguments for a cell: the activation a column slice of a wider
+ row, the pool a strided view (every 4th column skipped), slot 0's window zero (a fresh
+ sequence), the snapshot buffer pre-filled so untouched steps and rows are visible."""
+ k, conv_dim, p = spec["k"], spec["conv_dim"], spec["p"]
+ slots = list(spec["slots"])
+ m = len(slots) * p
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ wide = torch.randn(m, k + 128, generator=g, device=DEV).to(torch.bfloat16)
+ backing = torch.randn(spec["pool_slots"], conv_dim, 4, generator=g, device=DEV)
+ conv_state = backing.to(torch.bfloat16)[:, :, 1:]
+ conv_state[0] = 0
+ conv_weight = torch.randn(conv_dim, 4, generator=g, device=DEV).to(torch.bfloat16)
+ window = torch.randn(
+ spec["snapshot_rows"], spec["snapshot_steps"], conv_dim, 3, generator=g, device=DEV
+ ).to(torch.bfloat16)
+ idx = torch.tensor(slots, dtype=torch.int32, device=DEV)
+ rows = torch.arange(len(slots), dtype=torch.int32, device=DEV)
+ return [wide[:, :k], conv_state, conv_weight, idx, window, rows, conv_dim, p, -1]
+
+
+def _two_launch(x, packed, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad):
+ """The target's verify path before the fusion: the projection, then the Triton conv over
+ the transposed q|k|v view; returns `(conv, rest)` in the fused entry's layout."""
+ proj = fp8_gemv_op(x, packed)
+ n = x.shape[0] // p
+ x3 = torch.transpose(torch.reshape(proj[:, :conv_dim], [n, p, conv_dim]), 1, 2)
+ out = causal_conv1d_triton_update(
+ x3, conv_state, conv_weight, None, "silu", idx, pad, None, window, rows, False
+ )
+ return torch.reshape(torch.transpose(out, 1, 2), [n * p, conv_dim]), proj[:, conv_dim:]
+
+
+@pytest.mark.parametrize(
+ "cell",
+ fp8_gemv_conv1d_verify_op.CELLS,
+ ids=[c.why[:40] for c in fp8_gemv_conv1d_verify_op.CELLS],
+)
+def test_conv1d_verify_cells(cell) -> None:
+ spec = cell.spec
+ codes, scale, packed = _weight(spec["n"], spec["k"], 3)
+ x, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad = _verify_args(spec, 7)
+ live = [r for r in range(x.shape[0]) if spec["slots"][r // p] != pad]
+
+ state_a, window_a = conv_state.clone(), window.clone()
+ conv, rest = fp8_gemv_conv1d_verify_op(
+ x, packed, state_a, conv_weight, idx, window_a, rows, conv_dim, p, pad
+ )
+ assert conv.shape == (x.shape[0], conv_dim) and conv.is_contiguous()
+ assert rest.shape == (x.shape[0], spec["n"] - conv_dim)
+
+ # Bitwise the two-launch path (same projection plan, the Triton conv's arithmetic).
+ state_b, window_b = conv_state.clone(), window.clone()
+ conv_b, rest_b = _two_launch(
+ x, packed, state_b, conv_weight, idx, window_b, rows, conv_dim, p, pad
+ )
+ assert torch.equal(rest.view(torch.int16), rest_b.view(torch.int16))
+ assert torch.equal(conv[live].view(torch.int16), conv_b[live].view(torch.int16))
+ assert torch.equal(state_a.view(torch.int16), state_b.view(torch.int16)), "whole pool"
+ assert torch.equal(window_a.view(torch.int16), window_b.view(torch.int16)), "whole buffer"
+
+ # Against the fp32 reference: the projection may round one ulp differently (fp32 summation
+ # order), and a flip carries through the conv (one tap times one ulp of the input: within
+ # 2^-5 of the row's peak), the snapshots and the pool (one ulp: within 2^-7 of the peak).
+ want, want_rest, want_state, want_window = fp8_gemv_conv1d_verify_op.reference(
+ x, codes, scale, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad
+ )
+ deq = dequantize_fp8_block(codes, scale)
+ with true_fp32_matmul():
+ ref = x.float() @ deq.t()
+ _close(rest, ref[:, conv_dim:], x, deq[conv_dim:])
+ assert (conv == want).float().mean().item() > 0.999
+ peak = want.float().abs().amax(1, keepdim=True)
+ assert ((conv.float() - want.float()).abs() <= 2.0**-5 * peak).all()
+ for got, ref_ in ((state_a, want_state), (window_a, want_window)):
+ assert (got == ref_).float().mean().item() > 0.9999
+ assert (got.float() - ref_.float()).abs().max() <= 2.0**-7 * ref_.float().abs().max()
+ # Non-vacuity: every live slot and snapshot row was rewritten.
+ for r, slot in enumerate(spec["slots"]):
+ if slot != pad:
+ assert not torch.equal(state_a[slot], conv_state[slot]), slot
+ assert not torch.equal(window_a[r, :p], window[r, :p]), r
+ if spec["snapshot_steps"] > p:
+ assert torch.equal(window_a[:, p:], window[:, p:]), "steps past p untouched"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_per_tensor_mm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_per_tensor_mm.py
new file mode 100644
index 000000000000..7329eb4f2890
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_per_tensor_mm.py
@@ -0,0 +1,70 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test: `fp8_per_tensor_mm` against its fp32 reference (the activation quantized as the
+op does, both operands dequantized) per cell, the activation quantization bitwise against
+torch's e4m3 cast, and a zero activation and strided rows."""
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.fp8_per_tensor_mm import ( # noqa: E501
+ E4M3_MAX,
+ fp8_per_tensor_mm,
+ fp8_per_tensor_weight,
+ fp8_quantize_per_tensor,
+)
+
+assert torch.cuda.is_available(), "fp8_per_tensor_mm requires a CUDA device"
+
+DEV = "cuda"
+
+
+def _operands(m, k, n, seed):
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ x = (torch.randn(m, k, generator=g, device=DEV) * 0.7).to(torch.bfloat16)
+ w = (torch.randn(n, k, generator=g, device=DEV) * k**-0.5).to(torch.bfloat16)
+ codes, scale = fp8_per_tensor_weight(w)
+ return x, w, codes.t(), scale
+
+
+@pytest.mark.parametrize(
+ "cell", fp8_per_tensor_mm.CELLS, ids=[c.why[:40] for c in fp8_per_tensor_mm.CELLS]
+)
+def test_cells(cell) -> None:
+ m, k, n = cell.spec["m"], cell.spec["k"], cell.spec["n"]
+ x, w, codes_t, scale = _operands(m, k, n, 1000 + m)
+ got = fp8_per_tensor_mm(x, codes_t, scale)
+ want = fp8_per_tensor_mm.reference(x, codes_t, scale)
+ assert got.dtype == torch.bfloat16 and got.shape == (m, n)
+ rel = (got.float() - want).norm() / want.norm()
+ # The bf16 output rounding and cuBLASLt's accumulation separate them: its small-M FP8
+ # kernels accumulate below fp32 precision (measured 8.0e-3 at one row, under 4e-3 from
+ # 128 rows) -- a fifth of the e4m3 quantization error the op is gated on below.
+ assert rel < 1.2e-2, f"relative error {rel:.2e} against the dequantized reference"
+ full = x.float() @ w.float().t()
+ lossy = (got.float() - full).norm() / full.norm()
+ assert lossy < 0.06, f"FP8 error against the bf16 projection {lossy:.3f}"
+
+
+def test_quantization_is_torchs_cast() -> None:
+ g = torch.Generator(device=DEV).manual_seed(5)
+ x = (torch.randn(37, 2560, generator=g, device=DEV) * 3).to(torch.bfloat16)
+ q, scale = fp8_quantize_per_tensor(x)
+ want_scale = x.abs().amax().float() / E4M3_MAX
+ assert torch.equal(scale, want_scale.reshape(1))
+ want = (x.float() * (E4M3_MAX / x.abs().amax().float())).to(torch.float8_e4m3fn)
+ assert (q.view(torch.uint8) != want.view(torch.uint8)).float().mean() < 1e-4
+
+
+def test_zero_activation_and_strided_rows() -> None:
+ x, _, codes_t, scale = _operands(8, 2560, 2560, 3)
+ zeros = torch.zeros_like(x)
+ assert torch.equal(
+ fp8_per_tensor_mm(zeros, codes_t, scale),
+ torch.zeros(8, 2560, dtype=torch.bfloat16, device=DEV),
+ )
+ wide = torch.zeros(8, 4096, dtype=torch.bfloat16, device=DEV)
+ wide[:, :2560] = x
+ assert torch.equal(
+ fp8_per_tensor_mm(wide[:, :2560], codes_t, scale), fp8_per_tensor_mm(x, codes_t, scale)
+ )
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_low_m_bf16_mm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_low_m_bf16_mm.py
new file mode 100644
index 000000000000..699442c6d0f3
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_low_m_bf16_mm.py
@@ -0,0 +1,113 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the low_m_bf16_mm entry.
+
+Every cell names the branch `plan` must pick, and a host-only check pins the
+branch of the table's other shapes and M buckets, so a table edit that moves a
+shape to another kernel fails here before it can change a number silently.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.low_m_bf16_mm import ( # noqa: E501
+ low_m_bf16_mm as op,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.low_m_bf16_mm import ( # noqa: E501
+ plan,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import ( # noqa: E402,F401 — path above; fixture by name
+ ULP,
+ assert_within_ulp,
+ certifying,
+ validating,
+)
+
+assert torch.cuda.is_available(), "low_m_bf16_mm requires a CUDA device"
+
+
+def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
+ """Accumulation order only: 4 bf16 ulp element-wise against the row scale,
+ 2 ulp of relative RMS -- the band `cublas_mm` uses, for the same reason.
+ Every branch sums the same bf16 products in fp32 and rounds once; what is
+ left between kernel and fp32 reference is the final bf16 rounding (half an
+ ulp) plus the summation order."""
+ assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0, ulp=ULP[torch.bfloat16])
+
+
+def _build(m, n, k, seed):
+ """`mat_b` is a `[N, K]` weight's `.t()`, as a target stores its projections."""
+ g = torch.Generator(device="cuda").manual_seed(seed)
+ mat_a = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16)
+ weight = (torch.randn(n, k, generator=g, device="cuda") * k**-0.5).to(torch.bfloat16)
+ return mat_a, weight.t()
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811
+ certifying(op)
+ s = cell.spec
+ assert plan(s["m"], s["n"], s["k"]).kind == s["kind"], plan(s["m"], s["n"], s["k"])
+ mat_a, mat_b = _build(s["m"], s["n"], s["k"], seed=zlib.crc32(cell.why.encode()) % 2**31)
+ out = op(mat_a, mat_b, s["pdl"])
+ assert out.shape == (s["m"], s["n"]) and out.dtype == torch.bfloat16 and out.is_contiguous()
+ _compare(out, op.reference(mat_a, mat_b, s["pdl"]))
+
+
+#: `(M, N, K, branch)` of the table's shapes and M buckets the cells do not drive:
+#: what `plan` picks for each, held without a launch.
+_PLAN_PINS = (
+ (1, 10240, 2560, "direct"),
+ (1, 48, 2560, "direct"),
+ (32, 48, 2560, "direct"),
+ (8, 2560, 6144, "splitk"),
+ (2, 16480, 2560, "cublas"),
+ (1, 2560, 6144, "direct"),
+ (3, 12288, 2560, "splitk"),
+ (8, 12288, 2560, "splitk"),
+ (16, 12288, 2560, "cublas"),
+ (16, 512, 2560, "direct"),
+ (24, 512, 2560, "splitk"),
+ (16, 640, 2560, "splitk"),
+ (4, 1280, 2560, "direct"),
+ (8, 1296, 2560, "splitk"),
+ (1, 1, 2560, "direct"),
+ (1, 16480, 2560, "direct"),
+ (2, 14080, 2560, "direct"),
+ (3, 14080, 2560, "splitk"),
+ (8, 14080, 2560, "splitk"),
+ (16, 14080, 2560, "cublas"),
+ (1, 2560, 2560, "direct"),
+ (4, 2560, 2560, "direct"),
+ (8, 2560, 2560, "splitk"),
+ (1, 2560, 640, "cublas"),
+ (8192, 2560, 6144, "cublas"),
+ (1, 1024, 2048, "cublas"),
+ (3, 4120, 2560, "splitk"),
+ (9, 4120, 2560, "cublas"),
+ (2, 4352, 2560, "direct"),
+ (3, 4352, 2560, "splitk"),
+ (9, 4352, 2560, "cublas"),
+ (2, 2560, 1536, "direct"),
+ (5, 2560, 1536, "splitk"),
+ (9, 2560, 1536, "cublas"),
+ (33, 512, 2560, "cublas"),
+)
+
+
+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 _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()
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_mid_rows_mm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_mid_rows_mm.py
new file mode 100644
index 000000000000..d6dfb627ad4f
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_mid_rows_mm.py
@@ -0,0 +1,102 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test: `mid_rows_mm` (FP8 with the per-row activation quantization, and bf16) against an
+fp32 reference over the operands it multiplies, per cell; `fp8_row_quant` bitwise against
+torch's e4m3 cast; strided activation rows."""
+
+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.ops.gemm.mid_rows_mm import ( # noqa: E501
+ AMAX_FLOOR,
+ E4M3_MAX,
+ MidWeight,
+ fp8_row_quant,
+ fp8_row_quant_op,
+ mid_rows_mm_op,
+)
+
+assert torch.cuda.is_available(), "mid_rows_mm requires a CUDA device"
+
+DEV = "cuda"
+
+
+def _weight(n, k, mode, seed):
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ w = torch.randn(n, k, generator=g, device=DEV) * k**-0.5
+ w[:, 7] *= 20.0 # outlier column
+ w = w.to(torch.bfloat16)
+ if mode == "bf16":
+ return w, MidWeight(w.contiguous())
+ if mode == "row":
+ scale = (w.float().abs().amax(1) / E4M3_MAX).clamp_min(1e-12)
+ return w, MidWeight((w.float() / scale[:, None]).to(torch.float8_e4m3fn), scale)
+ scale = (w.float().abs().amax() / E4M3_MAX).clamp_min(1e-12).reshape(1)
+ return w, MidWeight((w.float() / scale).to(torch.float8_e4m3fn), scale)
+
+
+def _activation(m, k, seed, pad=0):
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ wide = torch.randn(m, k + pad, generator=g, device=DEV) * 0.7
+ wide[:, 3] *= 12.0 # outlier column
+ if m > 2:
+ wide[1] = 0.0 # an all-zero row
+ return wide.to(torch.bfloat16)[:, :k]
+
+
+def _dequantized(x, mw):
+ """The fp32 operands the call multiplies."""
+ if not mw.fp8:
+ return x.float(), mw.w.float()
+ q, s = fp8_row_quant(x)
+ wd = mw.w.float() * (mw.scale[:, None] if mw.scale.numel() > 1 else mw.scale)
+ return q.float() * s[:, None], wd
+
+
+def _close(y, x, mw) -> None:
+ """Within one bf16 rounding of the fp32 product of the operands (plus fp32 summation
+ noise relative to sum |x w|)."""
+ xd, wd = _dequantized(x, mw)
+ with true_fp32_matmul():
+ ref = xd @ wd.t()
+ absref = xd.abs() @ wd.abs().t()
+ ulp = torch.exp2(torch.floor(torch.log2(ref.abs().clamp_min(1e-30))) - 7)
+ excess = ((y.float() - ref).abs() - 0.5 * ulp).clamp_min(0) / absref.clamp_min(1e-30)
+ assert torch.isfinite(y.float()).all()
+ assert excess.max().item() < 1e-5, f"error beyond one rounding: {excess.max().item():.2e}"
+
+
+@pytest.mark.parametrize(
+ "cell", mid_rows_mm_op.CELLS, ids=[c.why[:48] for c in mid_rows_mm_op.CELLS]
+)
+def test_cells(cell) -> None:
+ n, k, m, mode = cell.spec["n"], cell.spec["k"], cell.spec["m"], cell.spec["mode"]
+ w, mw = _weight(n, k, mode, n + m)
+ x = _activation(m, k, m, pad=64 if mode != "bf16" else 0)
+ mid_rows_mm_op.is_valid(x, mw)
+ y = mid_rows_mm_op(x, mw)
+ assert y.shape == (m, n) and y.dtype == torch.bfloat16 and y.is_contiguous()
+ _close(y, x, mw)
+ if mw.fp8:
+ # The FP8 product against the bf16 projection: e4m3 keeps 3 mantissa bits on both
+ # operands; per-row activation scales are at least as accurate as one per tensor.
+ full = x.float() @ w.float().t()
+ assert ((y.float() - full).norm() / full.norm()).item() < 0.06
+
+
+@pytest.mark.parametrize(
+ "cell", fp8_row_quant_op.CELLS, ids=[c.why[:48] for c in fp8_row_quant_op.CELLS]
+)
+def test_row_quant_is_torchs_cast(cell) -> None:
+ m, k, pad = cell.spec["m"], cell.spec["k"], cell.spec.get("pad", 0)
+ x = _activation(m, k, 5 + m, pad=pad) * 3
+ if m > 4:
+ x[4, 9] = 900.0 # a row whose maximum is far above the rest
+ q, s = fp8_row_quant_op(x)
+ # The kernel divides correctly rounded (torch divides by a scalar as a product with
+ # its reciprocal): the quotients in float64, rounded once.
+ amax = x.float().abs().amax(1).clamp_min(AMAX_FLOOR).double()
+ assert torch.equal(s, (amax / E4M3_MAX).float())
+ want = (x.float() * (E4M3_MAX / amax).float()[:, None]).to(torch.float8_e4m3fn)
+ assert torch.equal(q.view(torch.uint8), want.view(torch.uint8))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_nvfp4_gemv.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_nvfp4_gemv.py
new file mode 100644
index 000000000000..6fe44edb45fa
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_nvfp4_gemv.py
@@ -0,0 +1,288 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test: `nvfp4_gemv` against an fp32 reference over the dequantized NVFP4 weight, per
+cell; the quantizer against its definition; one compiled kernel per row
+octet. The split entries (NVFP4 units on `fp8_gemv`'s kernel) per cell:
+the projection against the fp32 reference, the conv entries also against the two-launch
+paths they replace (the split projection + `causal_conv1d_update` /
+`causal_conv1d_triton_update`)."""
+
+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.ops.gemm import ( # noqa: E501
+ nvfp4_gemv as nv,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.causal_conv1d_triton_update import ( # noqa: E501
+ causal_conv1d_triton_update,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.causal_conv1d_update import ( # noqa: E501
+ causal_conv1d_update,
+)
+
+assert torch.cuda.is_available(), "nvfp4_gemv requires a CUDA device"
+
+DEV = "cuda"
+nvfp4_gemv_op = nv.nvfp4_gemv_op
+
+
+def _weight(n, k, seed):
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ w = torch.randn(n, k, generator=g, device=DEV) * k**-0.5
+ w[:, 7] *= 30.0 # outlier column: its blocks' scales run up to e4m3's top
+ w[: min(n, 64), 128:256] *= 1e-4 # near-zero blocks: small (subnormal) e4m3 scales
+ w = w.to(torch.bfloat16)
+ codes, scales, gscale = nv.quantize_nvfp4(w)
+ return w, codes, scales, gscale, nv.pack_weight(codes, scales, gscale)
+
+
+def _x(m, k, seed):
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ # A column slice of a wider row, as the target's split activations hand over.
+ return torch.randn(m, k + 256, generator=g, device=DEV).to(torch.bfloat16)[:, :k]
+
+
+def _close(y, ref, x, deq) -> None:
+ """Within one bf16 rounding of the fp32 reference, plus fp32 summation noise."""
+ with true_fp32_matmul():
+ absref = x.float().abs() @ deq.abs().t()
+ ulp = torch.exp2(torch.floor(torch.log2(ref.abs().clamp_min(1e-30))) - 7)
+ excess = ((y.float() - ref).abs() - 0.5 * ulp).clamp_min(0) / absref.clamp_min(1e-30)
+ assert excess.max().item() < 1e-5, f"error beyond one rounding: {excess.max().item():.2e}"
+
+
+@pytest.mark.parametrize("cell", nvfp4_gemv_op.CELLS, ids=[c.why[:40] for c in nvfp4_gemv_op.CELLS])
+def test_gemv_cells(cell) -> None:
+ n, k, m = cell.spec["n"], cell.spec["k"], cell.spec["m"]
+ _, codes, scales, gscale, packed = _weight(n, k, n + m)
+ x = _x(m, k, m)
+ nvfp4_gemv_op.is_valid(x, packed)
+ y = nvfp4_gemv_op(x, packed)
+ assert y.shape == (m, n) and y.dtype == torch.bfloat16
+ ref = nvfp4_gemv_op.reference(x, codes, scales, gscale)
+ _close(y, ref, x, nv.dequantize_nvfp4(codes, scales, gscale))
+
+
+def test_quantizer_matches_its_definition() -> None:
+ """Codes are the nearest e2m1 value of w / (block scale x global scale); the largest block
+ scale is e4m3's 448; the relative error is NVFP4's (well under 2^-2 a row)."""
+ w, codes, scales, gscale, _ = _weight(512, 1024, 3)
+ assert codes.max().item() <= 15 and torch.isfinite(scales.float()).all()
+ assert scales.float().max().item() == 448.0
+ step = scales.float().repeat_interleave(nv.SF_VEC, 1) * gscale
+ v = torch.where(step > 0, w.float() / step.clamp_min(1e-38), 0.0)
+ grid = torch.tensor(nv.E2M1_VALUES, device=DEV)
+ best = (v.abs()[..., None] - grid).abs().min(-1).values
+ got = nv.dequantize_nvfp4(codes, scales, gscale) / step.clamp_min(1e-38)
+ err = (v - torch.where(step > 0, got, 0.0)).abs()
+ assert (err <= best + 1e-6).all(), "a code is not the nearest e2m1 value"
+ deq = nv.dequantize_nvfp4(codes, scales, gscale)
+ rel = (deq - w.float()).norm(dim=1) / w.float().norm(dim=1)
+ assert rel.max().item() < 0.2, rel.max().item()
+
+
+def test_chunked_pack_matches_whole() -> None:
+ """`pack_from_bf16` (row chunks, a partial last tile) is `pack_weight` of the whole."""
+ w = torch.randn(1000, 256, device=DEV).to(torch.bfloat16)
+ a = nv.pack_from_bf16(w, rows_per_chunk=192)
+ b = nv.pack_weight(*nv.quantize_nvfp4(w))
+ assert a.global_scale == b.global_scale and torch.equal(a.data, b.data)
+
+
+def test_one_kernel_per_row_octet() -> None:
+ """The row count is a launch argument: rows 1..16 build two kernels per shape."""
+ _, codes, scales, gscale, packed = _weight(2048, 512, 5)
+ before = nv._compiled.cache_info().currsize
+ for m in range(1, nv.MAX_ROWS + 1):
+ x = _x(m, 512, m)
+ _close(
+ nvfp4_gemv_op(x, packed),
+ nvfp4_gemv_op.reference(x, codes, scales, gscale),
+ x,
+ nv.dequantize_nvfp4(codes, scales, gscale),
+ )
+ assert nv._compiled.cache_info().currsize - before <= 2
+
+
+# ------------------------------------------------------------------------------------------
+# split entries: NVFP4 units on fp8_gemv's kernel
+# ------------------------------------------------------------------------------------------
+
+split_op = nv.nvfp4_split_gemv_op
+split_conv_op = nv.nvfp4_split_gemv_conv1d_op
+split_verify_op = nv.nvfp4_split_gemv_conv1d_verify_op
+GDN_N, GDN_K, GDN_CONV = 16480, 2560, 10240
+
+
+@pytest.mark.parametrize("cell", split_op.CELLS, ids=[c.why[:40] for c in split_op.CELLS])
+def test_split_gemv_cells(cell) -> None:
+ n, k, m = cell.spec["n"], cell.spec["k"], cell.spec["m"]
+ _, codes, scales, gscale, _ = _weight(n, k, n + m)
+ packed = nv.pack_split(codes, scales, gscale)
+ x = _x(m, k, m)
+ split_op.is_valid(x, packed)
+ y = split_op(x, packed)
+ assert y.shape == (m, n) and y.dtype == torch.bfloat16
+ _close(
+ y,
+ split_op.reference(x, codes, scales, gscale),
+ x,
+ nv.dequantize_nvfp4(codes, scales, gscale),
+ )
+
+
+def test_split_layouts_agree() -> None:
+ """Tile-major and k-block-major units give the same projection (same plan)."""
+ _, codes, scales, gscale, _ = _weight(1001, 384, 4)
+ x = _x(5, 384, 5)
+ a = split_op(x, nv.pack_split(codes, scales, gscale, tile_major=True))
+ b = split_op(x, nv.pack_split(codes, scales, gscale, tile_major=False))
+ assert torch.equal(a, b)
+
+
+@pytest.mark.parametrize("cell", split_conv_op.CELLS, ids=[c.why[:40] for c in split_conv_op.CELLS])
+def test_split_conv1d_cells(cell) -> None:
+ pool, pad = 64, -1
+ _, codes, scales, gscale, _ = _weight(GDN_N, GDN_K, 3)
+ packed = nv.pack_split(codes, scales, gscale)
+ deq = nv.dequantize_nvfp4(codes, scales, gscale)
+ slots = list(cell.spec["slots"])
+ idx = torch.tensor(slots, dtype=torch.int32, device=DEV)
+ g = torch.Generator(device=DEV).manual_seed(11)
+ backing = torch.randn(pool, GDN_CONV, 4, generator=g, device=DEV).to(torch.bfloat16)
+ conv_state = backing[:, :, 1:] # a strided pool view
+ conv_weight = torch.randn(GDN_CONV, 4, generator=g, device=DEV).to(torch.bfloat16)
+ x = torch.randn(len(slots), GDN_K, generator=g, device=DEV).to(torch.bfloat16)
+ before = conv_state.clone()
+ split_conv_op.is_valid(x, packed)
+ out = split_conv_op(x, packed, conv_state, conv_weight, idx, GDN_CONV, pad)
+
+ # Against the two-launch path over the same projection.
+ proj = split_op(x, packed)
+ state_b = before.clone()
+ two = proj.clone()
+ causal_conv1d_update(
+ x=two[:, :GDN_CONV],
+ conv_state=state_b,
+ weight=conv_weight,
+ bias=None,
+ silu_activation=True,
+ cache_seqlens=None,
+ conv_state_indices=idx,
+ pad_slot_id=pad,
+ )
+ assert torch.equal(out[:, GDN_CONV:], proj[:, GDN_CONV:])
+ assert (out == two).float().mean().item() > 0.999
+ assert ((out.float() - two.float()).abs() / two.float().abs().clamp_min(1e-2)).max() < 0.02
+ assert torch.equal(conv_state, state_b), "the pool as the conv kernel leaves it"
+
+ # Against the fp32 reference: the projection may round one ulp differently (fp32 summation
+ # order, the activation rounded to f16), and a flip carries through the conv (one tap times
+ # one ulp of the input: within 2^-5 of the row's peak).
+ with true_fp32_matmul():
+ ref = (x.float() @ deq.t()).to(torch.bfloat16)
+ want = ref.clone()
+ for r, s in enumerate(slots):
+ if s == pad:
+ continue
+ window = torch.cat([before[s].float(), ref[r, :GDN_CONV].float()[:, None]], 1)
+ c = (window * conv_weight.float()).sum(1)
+ want[r, :GDN_CONV] = (c * torch.sigmoid(c)).to(torch.bfloat16)
+ assert (out == want).float().mean().item() > 0.999
+ peak = want.float().abs().amax(1, keepdim=True)
+ assert ((out.float() - want.float()).abs() <= 2.0**-5 * peak).all()
+ others = [s for s in range(pool) if s not in slots]
+ assert torch.equal(conv_state[others], before[others]), "other slots untouched"
+
+
+def _verify_args(spec, seed):
+ """As the fp8 verify test: a column slice of a wider row, a strided pool (slot 0 zero), a
+ pre-filled snapshot buffer."""
+ k, conv_dim, p = spec["k"], spec["conv_dim"], spec["p"]
+ slots = list(spec["slots"])
+ m = len(slots) * p
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ wide = torch.randn(m, k + 128, generator=g, device=DEV).to(torch.bfloat16)
+ backing = torch.randn(spec["pool_slots"], conv_dim, 4, generator=g, device=DEV)
+ conv_state = backing.to(torch.bfloat16)[:, :, 1:]
+ conv_state[0] = 0
+ conv_weight = torch.randn(conv_dim, 4, generator=g, device=DEV).to(torch.bfloat16)
+ window = torch.randn(
+ spec["snapshot_rows"], spec["snapshot_steps"], conv_dim, 3, generator=g, device=DEV
+ ).to(torch.bfloat16)
+ idx = torch.tensor(slots, dtype=torch.int32, device=DEV)
+ rows = torch.arange(len(slots), dtype=torch.int32, device=DEV)
+ return [wide[:, :k], conv_state, conv_weight, idx, window, rows, conv_dim, p, -1]
+
+
+@pytest.mark.parametrize(
+ "cell", split_verify_op.CELLS, ids=[c.why[:40] for c in split_verify_op.CELLS]
+)
+def test_split_conv1d_verify_cells(cell) -> None:
+ spec = cell.spec
+ _, codes, scales, gscale, _ = _weight(spec["n"], spec["k"], 3)
+ packed = nv.pack_split(codes, scales, gscale)
+ x, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad = _verify_args(spec, 7)
+ live = [r for r in range(x.shape[0]) if spec["slots"][r // p] != pad]
+ split_verify_op.is_valid(
+ x, packed, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad
+ )
+ state_a, window_a = conv_state.clone(), window.clone()
+ conv, rest = split_verify_op(
+ x, packed, state_a, conv_weight, idx, window_a, rows, conv_dim, p, pad
+ )
+ # Bitwise the two-launch path over the same projection.
+ proj = split_op(x, packed)
+ n = x.shape[0] // p
+ state_b, window_b = conv_state.clone(), window.clone()
+ x3 = torch.transpose(torch.reshape(proj[:, :conv_dim], [n, p, conv_dim]), 1, 2)
+ two = causal_conv1d_triton_update(
+ x3, state_b, conv_weight, None, "silu", idx, pad, None, window_b, rows, False
+ )
+ two = torch.reshape(torch.transpose(two, 1, 2), [n * p, conv_dim])
+ assert torch.equal(rest, proj[:, conv_dim:])
+ assert torch.equal(conv[live], two[live])
+ assert torch.equal(state_a, state_b) and torch.equal(window_a, window_b)
+ # Against the fp32 reference.
+ want, want_rest, want_state, want_window = split_verify_op.reference(
+ x, codes, scales, gscale, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad
+ )
+ deq = nv.dequantize_nvfp4(codes, scales, gscale)
+ with true_fp32_matmul():
+ ref = x.float() @ deq.t()
+ _close(rest, ref[:, conv_dim:], x, deq[conv_dim:])
+ assert (conv == want).float().mean().item() > 0.999
+ peak = want.float().abs().amax(1, keepdim=True)
+ assert ((conv.float() - want.float()).abs() <= 2.0**-5 * peak).all()
+ for got, ref_ in ((state_a, want_state), (window_a, want_window)):
+ assert (got == ref_).float().mean().item() > 0.9999
+ assert (got.float() - ref_.float()).abs().max() <= 2.0**-7 * ref_.float().abs().max()
+
+
+def test_split_precompile_covers_the_decode_variants() -> None:
+ """After `precompile_split`, no projection / conv / verify call within its rows compiles."""
+ from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm import ( # noqa: E501
+ fp8_gemv as f8,
+ )
+
+ _, codes, scales, gscale, _ = _weight(GDN_N, GDN_K, 2)
+ packed = nv.pack_split(codes, scales, gscale)
+ nv.precompile_split(packed, 4, GDN_CONV, -1, (2, 4))
+ before = f8._compiled.cache_info().misses
+ conv_state = torch.zeros(8, GDN_CONV, 3, device=DEV).to(torch.bfloat16)
+ conv_weight = torch.randn(GDN_CONV, 4, device=DEV).to(torch.bfloat16)
+ window = torch.zeros(4, 4, GDN_CONV, 3, device=DEV).to(torch.bfloat16)
+ for m in range(1, 5):
+ x = torch.randn(m, GDN_K, device=DEV).to(torch.bfloat16)
+ idx = torch.arange(m, dtype=torch.int32, device=DEV)
+ split_op(x, packed)
+ split_conv_op(x, packed, conv_state, conv_weight, idx, GDN_CONV, -1)
+ for p in (2, 4):
+ if m % p == 0:
+ n = m // p
+ ids = torch.arange(n, dtype=torch.int32, device=DEV)
+ split_verify_op(
+ x, packed, conv_state, conv_weight, ids, window, ids, GDN_CONV, p, -1
+ )
+ assert f8._compiled.cache_info().misses == before
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_gated_residual.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_gated_residual.py
new file mode 100644
index 000000000000..a5f97c64c529
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_gated_residual.py
@@ -0,0 +1,126 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the qwen4_exp_gated_residual catalog entry.
+
+The thing under test is a formula -- which operand is gated, which axis is
+reduced, where each division by `hc_count` sits, the factor on the injection
+-- and every mis-reading of it keeps the shapes; the cells' band sits 32x..192x
+inside the nearest wrong formula (measured).
+"""
+
+import zlib
+from typing import Optional, Tuple
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501
+ gate_ratio,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501
+ qwen4_exp_gated_residual as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "qwen4_exp_gated_residual requires a CUDA device"
+
+
+def _compare(
+ out: Tuple[torch.Tensor, Optional[torch.Tensor]],
+ ref: Tuple[torch.Tensor, Optional[torch.Tensor]],
+) -> None:
+ """Each result within `rtol = eps`, `atol = eps * max|ref|` of the
+ activation dtype (`gate_ratio <= 1`).
+
+ The absolute term is not torch's 1e-5: `mixed` is a mean of signed
+ products, so its near-zero elements are cancellations whose error is
+ set by the whole stream's rounding, and a relative gate on them gates
+ nothing. The relative term is half torch's bf16 default. Measured on
+ sm_103 at 512 tokens: the correct result at 0.341 (`mixed`) and 0.404
+ (`inject`) of this band (0.198 / 0.210 against the unrounded float64,
+ the rest is the reference's own rounding), ten wrong formulas at
+ 32x..192x.
+ """
+ (mixed, inject), (ref_mixed, ref_inject) = out, ref
+ assert (inject is None) == (ref_inject is None), "inject must be None exactly when ref's is"
+ pairs = [(mixed, ref_mixed)] + ([] if inject is None else [(inject, ref_inject)])
+ for got, want in pairs:
+ assert got.dtype == want.dtype, (got.dtype, want.dtype)
+ assert got.shape == want.shape, (got.shape, want.shape)
+ assert got.is_contiguous()
+ ratio = gate_ratio(got, want)
+ assert ratio <= 1.0, f"{ratio:.4g}x the band"
+
+
+DEV = "cuda"
+HC_COUNT = 4
+HIDDEN = 2560
+WIDTH = HC_COUNT * HIDDEN
+LOWRANK = 320
+RMS_NORM_EPS = 1e-6
+
+# Element statistics of this checkpoint's layer-0 attn_hyper_connection tensors.
+# The normed stream is scaled by (1 + gamma), so gamma's tail -- |6.625| -- is
+# what puts it at ~+-30 rather than ~+-4, and the error floor scales with it.
+CKPT_DOWN_RMS = 0.0296
+CKPT_UP_RMS = 0.0681
+CKPT_INJECT_RMS = 0.0107
+CKPT_GAMMA_RMS = 0.477
+CKPT_GAMMA_ABSMAX = 6.625
+
+
+def _make_weights(seed, dtype=torch.bfloat16, hc_count=HC_COUNT):
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ width = hc_count * HIDDEN
+
+ def normal(*shape, scale):
+ return (torch.randn(*shape, generator=g, device=DEV) * scale).to(dtype)
+
+ down = normal(LOWRANK, width, scale=CKPT_DOWN_RMS)
+ up = normal(width, LOWRANK, scale=CKPT_UP_RMS)
+ inject = normal(hc_count, width, scale=CKPT_INJECT_RMS)
+ gamma = normal(width, scale=CKPT_GAMMA_RMS)
+ # A Gaussian does not reproduce the real gamma's ~14-sigma tail; plant it.
+ gamma[0] = CKPT_GAMMA_ABSMAX
+ gamma[width // 2] = -CKPT_GAMMA_ABSMAX
+ return down, up, inject, gamma
+
+
+def _grouped_gemma_rmsnorm(x, gamma, hc_count):
+ """The hc_norm the entry's input comes from -- here only to give the
+ operand its real distribution, not under test."""
+ hidden = x.shape[-1] // hc_count
+ xf = x.float().reshape(*x.shape[:-1], hc_count, hidden)
+ xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + RMS_NORM_EPS)
+ return (xf.flatten(-2) * (1.0 + gamma.float())).to(x.dtype)
+
+
+def _build(leading, seed, hc_count=HC_COUNT, dtype=torch.bfloat16, inject=True):
+ down, up, inject_w, gamma = _make_weights(seed, dtype, hc_count)
+ g = torch.Generator(device=DEV).manual_seed(seed + 1)
+ raw = torch.randn(*leading, hc_count * HIDDEN, generator=g, device=DEV).to(dtype)
+ normed = _grouped_gemma_rmsnorm(raw, gamma, hc_count)
+ return raw, (normed, down, up, inject_w if inject else None, hc_count)
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
+ certifying(op)
+ spec = cell.spec
+ hc_count = spec.get("hc_count", HC_COUNT)
+ _, args = _build(
+ spec["leading"],
+ seed=zlib.crc32(cell.why.encode()),
+ hc_count=hc_count,
+ dtype=spec.get("dtype", torch.bfloat16),
+ inject=spec.get("inject", True),
+ )
+ before = args[0].clone()
+ mixed, inject = op(*args)
+ assert mixed.shape == (*spec["leading"], HIDDEN)
+ if spec.get("inject", True):
+ assert inject.shape == (*spec["leading"], hc_count)
+ assert torch.equal(args[0], before), "the normed stream is read-only"
+ _compare((mixed, inject), op.reference(*args))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py
new file mode 100644
index 000000000000..877e8a866659
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py
@@ -0,0 +1,188 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the qwen4_exp_hc_decode_first_cute entry.
+
+Cells against the entry's float64 reference: `hc_out` within one ulp (plus the
+combine's and the finalize's own rounding at the MoE sites), `part` within `2^-16` of
+the reference relative to its scale, `ss` to fp32 accuracy. Beyond the cells: a
+stale-read check under CUDA-graph replay with PDL, where a kernel that read its
+activations before `griddepcontrol.wait` would see the previous replay's.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_decode_first_cute import ( # noqa: E501
+ qwen4_exp_hc_decode_first_cute as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "qwen4_exp_hc_decode_first_cute requires a CUDA device"
+
+DEV = "cuda"
+HC_COUNT, HIDDEN, RANK = 4, 2560, 320
+DOWN_RMS, INJECT_RMS, GAMMA_RMS, GAMMA_ABSMAX = 0.0296, 0.0107, 0.477, 6.625
+PART_EPS = 2.0**-16
+
+
+def _band(got, want, eps) -> float:
+ w = want.double()
+ return ((got.double() - w).abs() / (eps * w.abs().max().item() + eps * w.abs())).max().item()
+
+
+def _ulp(got, want) -> float:
+ eps, tiny = torch.finfo(want.dtype).eps, torch.finfo(want.dtype).tiny
+ w = want.double()
+ return ((got.double() - w).abs() / (eps * w.abs() + tiny)).max().item()
+
+
+def _compare(out, ref) -> None:
+ (hc, part, ss), (ref_hc, ref_part, ref_ss) = out, ref
+ assert hc.dtype == ref_hc.dtype and hc.shape == ref_hc.shape
+ assert part.dtype == torch.float32 and part.shape == ref_part.shape and part.stride(2) == 1
+ assert ss.dtype == torch.float32 and ss.shape == ref_ss.shape
+ assert _ulp(hc, ref_hc) <= 1.0, f"hc_out {_ulp(hc, ref_hc):.4g} ulp"
+ assert _band(part, ref_part, PART_EPS) <= 1.0, f"part {_band(part, ref_part, PART_EPS):.4g}x"
+ assert torch.allclose(ss, ref_ss, rtol=1e-5, atol=0), "ss"
+
+
+def _routed(block_out, topk_weights, topk_rows):
+ if topk_rows is None:
+ return block_out
+ valid = (topk_rows >= 0).unsqueeze(-1)
+ rows = block_out.double()[topk_rows.clamp(min=0).long()]
+ return (
+ torch.where(valid, rows * topk_weights.double().unsqueeze(-1), 0.0)
+ .sum(1)
+ .to(block_out.dtype)
+ )
+
+
+def _hc_band_ratio(
+ hc, ref_hc, block_out, inject, shared, gate, topk_weights=None, topk_rows=None
+) -> float:
+ """The `hc_out` band at the MoE sites: one ulp of `hc_out`, plus one
+ ulp of the combine (and with the fold, of the routed row) times `|inject|`."""
+ feps, tiny = torch.finfo(ref_hc.dtype).eps, torch.finfo(ref_hc.dtype).tiny
+ routed = _routed(block_out, topk_weights, topk_rows)
+ s = torch.sigmoid(gate.double().reshape(-1, 1))
+ b = (routed.double() + s * shared.double()).to(routed.dtype).double().abs()
+ if topk_rows is not None:
+ b = b + routed.double().abs()
+ band = feps * ref_hc.double().abs() + (
+ feps * b.unsqueeze(-2) * inject.double().abs().unsqueeze(-1)
+ ).flatten(-2)
+ return ((hc.double() - ref_hc.double()).abs() / (band + tiny)).max().item()
+
+
+def _compare_moe(args, out, ref) -> None:
+ hyper_input, block_out, inject, gamma, first, c, n, shared, gate = args[:9]
+ topk_weights, topk_rows = args[9:] if len(args) > 9 else (None, None)
+ (hc, part, ss), (ref_hc, _, _) = out, ref
+ ratio = _hc_band_ratio(hc, ref_hc, block_out, inject, shared, gate, topk_weights, topk_rows)
+ assert ratio <= 1.0, f"hc_out {ratio:.4g}x the band"
+ _compare((hc, part, ss), op.reference(hc, None, None, gamma, first, c, n))
+
+
+def _build(spec, seed):
+ """The test operands at one row."""
+ rows, hc_count, hidden, dtype = 1, HC_COUNT, HIDDEN, torch.bfloat16
+ width = hc_count * hidden
+ n_cols = spec.get("n_cols", RANK + hc_count)
+ g = torch.Generator(device=DEV).manual_seed(seed)
+
+ def normal(*shape, scale):
+ return torch.randn(*shape, generator=g, device=DEV) * scale
+
+ hyper_input = normal(rows, width, scale=3.0).to(dtype)
+ gamma = normal(width, scale=GAMMA_RMS)
+ gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX
+ gamma = gamma.to(dtype)
+ n_rows = -(-n_cols // 128) * 128
+ weight = torch.zeros(n_rows, width, device=DEV)
+ weight[:n_cols] = normal(n_cols, width, scale=DOWN_RMS)
+ weight[n_cols - hc_count : n_cols] *= INJECT_RMS / DOWN_RMS
+ first = weight.to(dtype).t()
+ block_out = inject = shared = gate_logits = None
+ if spec.get("inject", True):
+ block_out = normal(rows, hidden, scale=2.0).to(dtype)
+ inject = (2.0 * torch.sigmoid(normal(rows, hc_count, scale=1.0))).to(dtype)
+ if spec.get("moe", False):
+ shared = normal(rows, hidden, scale=2.0).to(dtype)
+ gate_logits = normal(rows, 1, scale=2.0).to(dtype)
+ top_k = spec.get("fold", 0)
+ if not top_k:
+ return hyper_input, block_out, inject, gamma, first, hc_count, n_cols, shared, gate_logits
+ n_perm = rows * top_k + 37
+ block_out = normal(n_perm, hidden, scale=2.0).to(dtype)
+ order = torch.randperm(n_perm, generator=g, device=DEV)[: rows * top_k]
+ topk_rows = order.reshape(rows, top_k).to(torch.int32)
+ if spec.get("drop"):
+ topk_rows[:, -3:] = -1
+ topk_weights = torch.softmax(normal(rows, top_k, scale=1.0), -1).to(dtype)
+ return (
+ hyper_input,
+ block_out,
+ inject,
+ gamma,
+ first,
+ hc_count,
+ n_cols,
+ shared,
+ gate_logits,
+ topk_weights,
+ topk_rows,
+ )
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
+ certifying(op)
+ args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31))
+ before = [a.clone() if isinstance(a, torch.Tensor) else a for a in args]
+ out = op(*args)
+ for a, b in zip(args, before):
+ assert not isinstance(a, torch.Tensor) or torch.equal(a, b), "every input is read-only"
+ if args[1] is None:
+ assert out[0] is args[0], "norm-only returns hyper_input itself"
+ if args[7] is None:
+ _compare(out, op.reference(*args))
+ else:
+ _compare_moe(args, out, op.reference(*args))
+
+
+def test_reads_after_the_wait_under_graph_replay() -> None:
+ """A producer kernel rewrites every activation inside the graph right before this op.
+ With PDL the op may start before the producer finishes; reading any activation before
+ `griddepcontrol.wait` would return the previous replay's value."""
+ args = list(_build(dict(moe=True, fold=10), seed=8))
+ sources = [a.clone() for a in (args[0], args[1], args[2], args[7], args[8], args[9])]
+ targets = (args[0], args[1], args[2], args[7], args[8], args[9])
+
+ def step():
+ for dst, src in zip(targets, sources):
+ dst.copy_(src * 1.0) # an elementwise producer kernel per operand
+ return op(*args)
+
+ warmup = torch.cuda.Stream()
+ warmup.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(warmup):
+ for _ in range(2):
+ step()
+ torch.cuda.current_stream().wait_stream(warmup)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ graph_out = step()
+ for trial in range(3):
+ fresh = _build(dict(moe=True, fold=10), seed=100 + trial)
+ for src, new in zip(sources, (fresh[0], fresh[1], fresh[2], fresh[7], fresh[8], fresh[9])):
+ src.copy_(new)
+ graph.replay()
+ torch.cuda.synchronize()
+ want = op(*[*args[:10], args[10]])
+ for a, b in zip(graph_out, want):
+ assert torch.equal(a, b), f"replay {trial}"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py
new file mode 100644
index 000000000000..831b6bfba981
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py
@@ -0,0 +1,152 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the 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), and a
+stale-read check under CUDA-graph replay with PDL for the NVFP4 build."""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import (
+ Fp4Quantize,
+ fp4_quantize,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_decode_mix_fp4_cute import ( # noqa: E501
+ qwen4_exp_hc_decode_mix_cute,
+ qwen4_exp_hc_decode_mix_fp4_cute,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "qwen4_exp_hc_decode_mix_fp4_cute requires a CUDA device"
+
+DEV = "cuda"
+HC_COUNT, HIDDEN, RANK, EPS = 4, 2560, 320, 1e-6
+UP_RMS, GAMMA_RMS, GAMMA_ABSMAX = 0.0681, 0.477, 6.625
+
+ENTRIES = {"fp4": qwen4_exp_hc_decode_mix_fp4_cute, "plain": qwen4_exp_hc_decode_mix_cute}
+CASES = [(name, cell) for name, entry in ENTRIES.items() for cell in entry.CELLS]
+
+
+def nvfp4_reference(x, gs):
+ """`fp4_quantize`'s native-torch NVFP4 of `x` `[T, K]` with linear scales:
+ `(packed [T, K/2] uint8, sf [T*K/16] uint8)`."""
+ return fp4_quantize.reference(x, gs, 16, False, False)
+
+
+def nvfp4_near_tie(x, gs):
+ """`[T, K]` bool: elements within `fp4_quantize`'s tie window of an e2m1 midpoint,
+ where the kernels' approximate reciprocals may round either way."""
+ return Fp4Quantize.near_tie(fp4_quantize.scaled(x, gs, 16))
+
+
+def _band(got, want) -> float:
+ eps = torch.finfo(want.dtype).eps
+ w = want.double()
+ return ((got.double() - w).abs() / (eps * w.abs().max().item() + eps * w.abs())).max().item()
+
+
+def _compare(out, ref) -> None:
+ (mixed, inj), (ref_mixed, ref_inj) = out[:2], ref[:2]
+ assert (inj is None) == (ref_inj is None)
+ for got, want in [(mixed, ref_mixed)] + ([] if inj is None else [(inj, ref_inj)]):
+ assert got.dtype == want.dtype and got.shape == want.shape and got.is_contiguous()
+ assert _band(got, want) <= 1.0, f"{_band(got, want):.4g}x the band"
+
+
+def _unpack(xq):
+ return torch.stack([xq & 0xF, xq >> 4], -1).flatten(-2)
+
+
+def _check_quant(mixed, xq, xsf, gs) -> None:
+ """The NVFP4 outputs exact on the returned `mixed`: scales bitwise, e2m1 codes
+ bitwise outside the window where the approximate reciprocals may pick either side."""
+ want_q, want_sf = nvfp4_reference(mixed, gs)
+ assert xq.dtype == torch.uint8 and xq.shape == want_q.shape and xq.is_contiguous()
+ assert xsf.dtype == torch.uint8 and xsf.shape == want_sf.shape
+ assert torch.equal(xsf, want_sf), f"{(xsf != want_sf).sum().item()} scale bytes differ"
+ stray = (_unpack(xq) != _unpack(want_q)) & ~nvfp4_near_tie(mixed, gs)
+ assert not stray.any(), f"{stray.sum().item()} e2m1 codes differ outside the tie window"
+
+
+def _build(spec, seed, fp4=True):
+ """The entry's arguments; `fp4` appends the NVFP4 build's global scale."""
+ rows, hc_count, hidden, rank, dtype = 1, HC_COUNT, HIDDEN, RANK, torch.bfloat16
+ with_inject = spec.get("with_inject", True)
+ width = hc_count * hidden
+ g = torch.Generator(device=DEV).manual_seed(seed)
+
+ def normal(*shape, scale):
+ return torch.randn(*shape, generator=g, device=DEV) * scale
+
+ hc_out = normal(rows, width, scale=3.0).to(dtype)
+ if spec.get("layout") == "row_slice":
+ buf = normal(2 * rows, width, scale=1.0).to(dtype)
+ buf[::2] = hc_out
+ hc_out = buf[::2]
+ ss = hc_out.double().unflatten(-1, (hc_count, hidden)).pow(2).sum(-1).t().float().contiguous()
+ gamma = normal(width, scale=GAMMA_RMS)
+ gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX
+ gamma = gamma.to(dtype)
+ n_pad = -(-(rank + hc_count) // 4) * 4
+ part = normal(hc_count, rows, n_pad + 4, scale=60.0)[:, :, :n_pad]
+ up = normal(width, rank, scale=UP_RMS).to(dtype).t()
+ args = [part, ss, up, hc_out, gamma, EPS, hc_count, with_inject]
+ if fp4:
+ # A calibrated activation global scale: 448 * 6 / amax of the mix (~|8|).
+ args.append(torch.tensor([448.0 * 6.0 / 8.0], device=DEV, dtype=torch.float32))
+ return args
+
+
+@pytest.mark.parametrize(
+ "name,cell", CASES, ids=[f"{name}-{cell.why[:44]}" for name, cell in CASES]
+)
+def test_certified_cells(name, cell, certifying) -> None: # noqa: F811 — fixture by name
+ op = ENTRIES[name]
+ certifying(op)
+ fp4 = name == "fp4"
+ args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31), fp4=fp4)
+ out = op(*args)
+ assert out[0].shape == (1, HIDDEN)
+ _compare(out, op.reference(*args))
+ if fp4:
+ _check_quant(out[0], out[2], out[3], args[8])
+
+
+def test_reads_after_the_wait_under_graph_replay() -> None:
+ """Producer kernels rewrite `part`, `ss` and `hc_out` inside the graph right before
+ the op; with PDL a read before `griddepcontrol.wait` would see the previous replay's.
+ The NVFP4 build: the MoE sites' variant, with the most outputs."""
+ op = qwen4_exp_hc_decode_mix_fp4_cute
+ args = _build(dict(), seed=8)
+ targets = (args[0], args[1], args[3])
+ sources = [t.clone() for t in targets]
+
+ def step():
+ for dst, src in zip(targets, sources):
+ dst.copy_(src * 1.0)
+ return op(*args)
+
+ warmup = torch.cuda.Stream()
+ warmup.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(warmup):
+ for _ in range(2):
+ step()
+ torch.cuda.current_stream().wait_stream(warmup)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ graph_out = step()
+ for trial in range(3):
+ fresh = _build(dict(), seed=100 + trial)
+ for src, new in zip(sources, (fresh[0], fresh[1], fresh[3])):
+ src.copy_(new)
+ graph.replay()
+ torch.cuda.synchronize()
+ for a, b in zip(graph_out, op(*args)):
+ assert torch.equal(a, b), f"replay {trial}"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_rows_tc.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_rows_tc.py
new file mode 100644
index 000000000000..c9a8f96ab984
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_rows_tc.py
@@ -0,0 +1,172 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test: the 2..16-row hyper-connection site (`qwen4_exp_hc_decode_rows_first` +
+`qwen4_exp_hc_decode_rows_mix`) against the 1-row GEMV kernels applied row by row, over the
+site kinds the forward sends (after a MoE block: combine + injection, with or without the
+runner's finalize folded in; MoE site: injection + NVFP4; norm-only); and the folded site
+bitwise against the runner's finalize (`fp4_block_scale_moe_renorm_permuted.finalize`, the
+runner's arithmetic) followed by the unfolded site. `hc_out` must be bitwise; `mixed` within
+1e-2 of the reference's largest magnitude (1-2 bf16 ulp: the projections take bf16 operands on
+tensor cores); the injection weights within one bf16 ulp; the NVFP4 codes may differ where
+`mixed` does."""
+
+import importlib
+
+import pytest
+import torch
+
+P = "tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc."
+rows_mod = importlib.import_module(P + "qwen4_exp_hc_decode_rows_tc")
+first_op = rows_mod.qwen4_exp_hc_decode_rows_first
+mix_op = rows_mod.qwen4_exp_hc_decode_rows_mix
+first1 = importlib.import_module(
+ P + "qwen4_exp_hc_decode_first_cute"
+).qwen4_exp_hc_decode_first_cute
+mix1 = importlib.import_module(P + "qwen4_exp_hc_decode_mix_fp4_cute")
+finalize = importlib.import_module(
+ P.replace(".ops.hc.", ".ops.moe.") + "fp4_block_scale_moe_renorm_permuted"
+).finalize
+
+assert torch.cuda.is_available(), "the multi-row hyper-connection site requires a CUDA device"
+
+DEV = "cuda"
+HIDDEN, C, RANK = 2560, 4, 320
+WIDTH = HIDDEN * C
+EPS = 1e-6
+TOPK = 10
+
+
+def _operands(rows, seed):
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ bf = torch.bfloat16
+ rnd = lambda *s, scale=1.0: (torch.randn(*s, generator=g, device=DEV) * scale).to(bf) # noqa: E731
+ return dict(
+ first=rnd(384, WIDTH, scale=0.01).t(),
+ up=rnd(WIDTH, RANK, scale=0.05).t(),
+ norm=rnd(WIDTH, scale=0.1),
+ hc=rnd(rows, WIDTH, scale=0.5),
+ block=rnd(rows, HIDDEN, scale=0.5),
+ shared=rnd(rows, HIDDEN, scale=0.5),
+ gate=rnd(rows, 1),
+ inject=(torch.rand(rows, C, generator=g, device=DEV) * 2).to(bf),
+ gscale=torch.tensor([448.0 * 6 / 4.0], device=DEV),
+ # The unfinalized runner's permuted FC2 rows (37 rows no slot maps), slot map
+ # (the last 3 slots of the odd rows -1) and routing weights.
+ perm=rnd(rows * TOPK + 37, HIDDEN, scale=0.5),
+ topk_rows=_slot_map(rows, g),
+ topk_w=torch.softmax(torch.randn(rows, TOPK, generator=g, device=DEV), -1).to(bf),
+ )
+
+
+def _slot_map(rows, g):
+ order = torch.randperm(rows * TOPK + 37, generator=g, device=DEV)[: rows * TOPK]
+ m = order.reshape(rows, TOPK).to(torch.int32)
+ m[1::2, -3:] = -1
+ return m.contiguous()
+
+
+def _site(o, site, rows_path: bool):
+ """The site through either path: `(hc_out, mixed, inject_out, xq)`."""
+ blk, inj = (None, None) if site == "norm" else (o["block"], o["inject"])
+ kw = dict(shared=o["shared"], gate_logits=o["gate"]) if site in ("combine", "fold") else {}
+ if site == "fold":
+ blk = o["perm"]
+ kw.update(topk_weights=o["topk_w"], topk_rows=o["topk_rows"])
+ n_cols = RANK if site == "norm" else RANK + C
+ gs = o["gscale"] if site == "inject" else None
+ if rows_path:
+ hc, low, inj_out, rstd = first_op(
+ o["hc"], blk, inj, o["norm"], o["first"], C, n_cols, EPS, **kw
+ )
+ mixed, xq, _ = mix_op(low, rstd, o["up"], hc, o["norm"], C, gs)
+ return hc, mixed, inj_out, xq
+ out = []
+ for r in range(o["hc"].shape[0]):
+ one = lambda t: None if t is None else t[r : r + 1] # noqa: E731
+ h, part, ss = first1(
+ one(o["hc"]),
+ blk if site == "fold" else one(blk),
+ one(inj),
+ o["norm"],
+ o["first"],
+ C,
+ n_cols,
+ **{k: one(v) for k, v in kw.items()},
+ )
+ if gs is not None:
+ m, i, q, _ = mix1.qwen4_exp_hc_decode_mix_fp4_cute(
+ part, ss, o["up"], h, o["norm"], EPS, C, True, gs
+ )
+ else:
+ m, i = mix1.qwen4_exp_hc_decode_mix_cute(
+ part, ss, o["up"], h, o["norm"], EPS, C, n_cols > RANK
+ )
+ q = None
+ out.append((h, m, i, q))
+ cat = lambda k: None if out[0][k] is None else torch.cat([x[k] for x in out]) # noqa: E731
+ return cat(0), cat(1), cat(2), cat(3)
+
+
+def _check(cell):
+ rows, site = cell.spec["rows"], cell.spec["site"]
+ o = _operands(rows, 1000 + rows)
+ got = _site(o, site, True)
+ want = _site(o, site, False)
+ assert torch.equal(got[0], want[0]), "hc_out bitwise"
+ rel = (got[1].float() - want[1].float()).abs().max() / want[1].float().abs().max()
+ assert rel <= 1e-2, f"mixed {rel:.2e} of the reference's max"
+ if want[2] is not None:
+ ulp = (got[2].float() - want[2].float()).abs() / want[2].float().abs().clamp_min(1e-6)
+ assert ulp.max() <= 2.0**-7 * 1.01, "injection weights within one bf16 ulp"
+ if want[3] is not None:
+ assert (got[3] != want[3]).float().mean() < 0.01, "NVFP4 codes"
+
+
+@pytest.mark.parametrize("cell", first_op.CELLS, ids=[c.why[:44] for c in first_op.CELLS])
+def test_first_cells(cell) -> None:
+ _check(cell)
+
+
+@pytest.mark.parametrize("cell", mix_op.CELLS, ids=[c.why[:44] for c in mix_op.CELLS])
+def test_mix_cells(cell) -> None:
+ _check(cell)
+
+
+@pytest.mark.parametrize("rows", [2, 7, 16])
+@pytest.mark.parametrize("site", ["combine", "fold"])
+def test_gate_column_of_the_fused_shared_gemm(rows, site) -> None:
+ """The target's gate logits are a column of the shared expert's fused gate|up GEMM output
+ (row stride 2 * 640 + 8): every output is bitwise the dense `[R, 1]` gate's, row by row,
+ so no row reads another's logit."""
+ o = _operands(rows, 500 + rows)
+ gug = torch.randn(rows, 2 * 640 + 8, device=DEV).to(torch.bfloat16)
+ gug[:, 1280:1281] = o["gate"]
+ want = _site(o, site, True)
+ got = _site(dict(o, gate=gug[:, 1280:1281]), site, True)
+ for a, b in zip(got, want):
+ assert (a is None and b is None) or torch.equal(a, b)
+
+
+@pytest.mark.parametrize("rows", [2, 16])
+@pytest.mark.parametrize("n_cols", [RANK + C, RANK])
+def test_fold_bitwise_vs_runner_finalize(rows, n_cols) -> None:
+ """Every output of the folded site equals the unfolded site fed the runner's finalize."""
+ o = _operands(rows, 300 + rows)
+ kw = dict(shared=o["shared"], gate_logits=o["gate"])
+ got = first_op(
+ o["hc"],
+ o["perm"],
+ o["inject"],
+ o["norm"],
+ o["first"],
+ C,
+ n_cols,
+ EPS,
+ topk_weights=o["topk_w"],
+ topk_rows=o["topk_rows"],
+ **kw,
+ )
+ routed = finalize(o["perm"], o["topk_w"], o["topk_rows"])
+ want = first_op(o["hc"], routed, o["inject"], o["norm"], o["first"], C, n_cols, EPS, **kw)
+ for a, b in zip(got, want):
+ assert (a is None and b is None) or torch.equal(a, b)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_first_rescale_tc.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_first_rescale_tc.py
new file mode 100644
index 000000000000..cd294ad6c6d8
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_first_rescale_tc.py
@@ -0,0 +1,92 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the qwen4_exp_hc_first_rescale_tc entry: `low` and `inject` in
+the gate band (`qwen4_exp_gated_residual`'s) of the float64 projection of the
+once-rounded normed stream, and the stack-final form.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501
+ gate_ratio,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_first_rescale_tc import ( # noqa: E501
+ qwen4_exp_hc_first_rescale_tc 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_first_rescale_tc requires a CUDA device"
+
+DEV = "cuda"
+HC_COUNT = 4
+HIDDEN = 2560
+RANK = 320
+EPS = 1e-6
+#: The padded first weight's rows (`_fuse_hc_first_projections`: 324 rounded up to 128).
+FIRST_ROWS = 384
+# This checkpoint's grouped-norm weight statistics (layer-0 hyper-connection norm).
+GAMMA_RMS = 0.477
+GAMMA_ABSMAX = 6.625
+
+
+def _compare(out, ref) -> None:
+ """gate_ratio <= 1 for `low` and `inject`: the band the mixer's outputs are
+ certified in (`rtol = eps`, `atol = eps * max|ref|`); the fp32 logits are a
+ 10240-term tensor-core sum against the float64 one.
+ """
+ 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 gate_ratio(got, want) <= 1.0, f"{gate_ratio(got, want):.4g}x the band"
+
+
+def _build(spec, seed):
+ rows, width, dtype = spec["rows"], HC_COUNT * HIDDEN, torch.bfloat16
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ hc_out = (torch.randn(rows, width, generator=g, device=DEV) * 3.0).to(dtype)
+ x = hc_out.double().unflatten(-1, (HC_COUNT, HIDDEN))
+ rstd = torch.rsqrt(x.pow(2).mean(-1) + EPS).float().contiguous()
+ gamma = torch.randn(width, generator=g, device=DEV) * GAMMA_RMS
+ gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX
+ # `[down; block_inject; 0]`, logits up to about +-60 (this checkpoint's reach |67|).
+ first = torch.zeros(FIRST_ROWS, width, device=DEV)
+ first[: RANK + HC_COUNT] = torch.randn(RANK + HC_COUNT, width, generator=g, device=DEV) * 0.15
+ return (
+ hc_out,
+ rstd,
+ gamma.to(dtype),
+ first.to(dtype).t(),
+ HC_COUNT,
+ RANK,
+ bool(spec.get("inject", True)),
+ )
+
+
+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))
+ low, inject = op(*args)
+ assert low.shape == (cell.spec["rows"], RANK)
+ assert (inject is None) == (not args[6])
+ _compare((low, inject), op.reference(*args))
+
+
+def test_stack_final_low_is_the_injecting_form_bitwise() -> None:
+ args = _build(dict(rows=300), seed=7)
+ low, inject = op(*args[:6], True)
+ low_f, none = op(*args[:6], False)
+ assert none is None and inject.shape == (300, HC_COUNT)
+ assert torch.equal(low, low_f)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py
new file mode 100644
index 000000000000..f87d5115878d
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py
@@ -0,0 +1,206 @@
+# 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 catalog entry.
+
+Cells are gated at one ulp of the activation dtype against the float64
+reference; the MoE-combine cells widen the residual's band by the combine's own
+rounding (see `_compare_moe`). Beyond them: the CuTe fold kernel bitwise the
+Triton kernel at the last row count before the switch.
+"""
+
+import sys
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_inject_norm import ( # noqa: E501
+ qwen4_exp_hc_inject_norm as op,
+)
+
+__extra_import_path__ = ["../../.."]
+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
+# The stack-final mixer's norm weight (`hyper_connection_mixer.hc_norm.weight`), which the
+# last layer's MoE site normalizes with: RMS 3.106, from -0.367 to 12.375, nearly all
+# positive -- so `1 + w` scales the normed stream by up to 13.4.
+FINAL_GAMMA_RMS = 3.106
+FINAL_GAMMA_MIN = -0.367
+FINAL_GAMMA_MAX = 12.375
+
+
+def _ulp_ratio(out: torch.Tensor, ref: torch.Tensor) -> float:
+ """Worst element's distance from `ref` in units of `eps(dtype) * |ref|`
+ (one ulp of the reference value, or a little more). 1 means one ulp."""
+ eps = torch.finfo(ref.dtype).eps
+ ref64 = ref.double()
+ tiny = torch.finfo(ref.dtype).tiny
+ return ((out.double() - ref64).abs() / (eps * ref64.abs() + tiny)).max().item()
+
+
+def _compare(out, ref) -> None:
+ """Both outputs within one ulp of the float64 result rounded once.
+
+ The op rounds each output once from an fp32 value whose own error is
+ ~1e-6 relative (a 16-bit product is exact in fp32; the RMS is a 2560-term
+ fp32 sum), so it can differ from the once-rounded float64 value only where
+ the two straddle a rounding midpoint: by exactly one ulp."""
+ (hc, normed), (ref_hc, ref_normed) = out, ref
+ for got, want in ((hc, ref_hc), (normed, ref_normed)):
+ assert got.dtype == want.dtype and got.shape == want.shape
+ if got.numel():
+ ratio = _ulp_ratio(got, want)
+ assert ratio <= 1.0, f"{ratio:.4g} ulp"
+
+
+def _combine(block_out, shared, gate_logits) -> torch.Tensor:
+ """The MoE block output in float64, rounded once to `block_out`'s dtype."""
+ s = torch.sigmoid(gate_logits.double().reshape(-1, 1))
+ return (block_out.double() + s * shared.double()).to(block_out.dtype)
+
+
+def _routed(block_out, topk_weights, topk_rows):
+ """The routed rows: `block_out` itself, or with the fold the float64 finalize of
+ the permuted rows rounded once."""
+ if topk_rows is None:
+ return block_out
+ valid = (topk_rows >= 0).unsqueeze(-1)
+ picked = block_out.double()[topk_rows.clamp(min=0).long()]
+ w = topk_weights.double().unsqueeze(-1)
+ return torch.where(valid, picked * w, 0.0).sum(1).to(block_out.dtype)
+
+
+def _compare_moe(args, out, ref) -> None:
+ """MoE-combine cells.
+
+ `hc_out`: the op's combine is an fp32 value rounded once, the reference's a
+ float64 value rounded once; where the two straddle a rounding midpoint they
+ differ by one ulp of the block output, and the accumulate carries that
+ scaled by `inject`. So the band is one ulp of `hc_out` plus one ulp of the
+ combine times `|inject|`, element by element -- still well under one ulp
+ of `hc_out` wherever the block term is small against the residual. With the
+ finalize fold the routed row has its own rounding (fp32 FMAs against the
+ float64 sum), one more ulp of the routed row times `|inject|`.
+
+ `normed`: one ulp of the float64 norm of the residual the op stored,
+ `reference(hc_out, None, None, ...)` -- the dataflow the op states."""
+ hyper_input, block_out, inject, gamma, eps, c, shared, gate = args[:8]
+ topk_weights, topk_rows = args[8:] if len(args) > 8 else (None, None)
+ (hc, normed), (ref_hc, _) = out, ref
+ assert hc.dtype == ref_hc.dtype and hc.shape == ref_hc.shape
+ if not hc.numel():
+ return
+ feps, tiny = torch.finfo(ref_hc.dtype).eps, torch.finfo(ref_hc.dtype).tiny
+ routed = _routed(block_out, topk_weights, topk_rows)
+ b = _combine(routed, shared, gate).double().abs()
+ if topk_rows is not None:
+ b = b + routed.double().abs()
+ band = feps * ref_hc.double().abs() + (
+ feps * b.unsqueeze(-2) * inject.double().abs().unsqueeze(-1)
+ ).flatten(-2)
+ ratio = ((hc.double() - ref_hc.double()).abs() / (band + tiny)).max().item()
+ assert ratio <= 1.0, f"hc_out {ratio:.4g}x the band"
+ own = op.reference(hc, None, None, gamma, eps, c)[1]
+ assert _ulp_ratio(normed, own) <= 1.0, f"normed {_ulp_ratio(normed, own):.4g} ulp"
+
+
+def _build(spec, seed):
+ rows, hc_count, hidden, dtype = spec["rows"], HC_COUNT, HIDDEN, torch.bfloat16
+ width = hc_count * hidden
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ hyper_input = (torch.randn(rows, width, generator=g, device=DEV) * 3.0).to(dtype)
+ if spec.get("site", "trunk") == "stack_final":
+ gamma = torch.randn(width, generator=g, device=DEV).abs() * FINAL_GAMMA_RMS
+ gamma = gamma.clamp(max=FINAL_GAMMA_MAX)
+ gamma[0], gamma[width // 2] = FINAL_GAMMA_MAX, FINAL_GAMMA_MIN
+ else:
+ gamma = torch.randn(width, generator=g, device=DEV) * GAMMA_RMS
+ gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX
+ gamma = gamma.to(dtype)
+ block_out = inject = shared = gate_logits = None
+ if spec.get("inject", True):
+ block_out = (torch.randn(rows, hidden, generator=g, device=DEV) * 2.0).to(dtype)
+ inject = (2.0 * torch.sigmoid(torch.randn(rows, hc_count, generator=g, device=DEV))).to(
+ dtype
+ )
+ if spec.get("moe", False):
+ shared = (torch.randn(rows, hidden, generator=g, device=DEV) * 2.0).to(dtype)
+ gate_logits = (torch.randn(rows, 1, generator=g, device=DEV) * 2.0).to(dtype)
+ top_k = spec.get("fold", 0)
+ if not top_k:
+ return hyper_input, block_out, inject, gamma, EPS, hc_count, shared, gate_logits
+ # The runner's unfinalized output: permuted FC2 rows (more rows than slots, as
+ # its padded buffer has), each slot's row and its routing weight.
+ n_perm = rows * top_k + 37
+ block_out = (torch.randn(n_perm, hidden, generator=g, device=DEV) * 2.0).to(dtype)
+ order = torch.randperm(n_perm, generator=g, device=DEV)[: rows * top_k]
+ topk_rows = order.reshape(rows, top_k).to(torch.int32)
+ if spec.get("drop"):
+ topk_rows[:, -3:] = -1
+ topk_weights = torch.softmax(torch.randn(rows, top_k, generator=g, device=DEV), -1).to(dtype)
+ return (
+ hyper_input,
+ block_out,
+ inject,
+ gamma,
+ EPS,
+ hc_count,
+ shared,
+ gate_logits,
+ topk_weights,
+ topk_rows,
+ )
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
+ certifying(op)
+ args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
+ before = [a.clone() if isinstance(a, torch.Tensor) else a for a in args]
+ hc, normed = op(*args)
+ for a, b in zip(args, before):
+ assert not isinstance(a, torch.Tensor) or torch.equal(a, b), "every input is read-only"
+ if args[1] is None:
+ assert hc is args[0], "norm-only returns hyper_input itself"
+ else:
+ assert hc.is_contiguous()
+ assert args[0].numel() == 0 or hc.data_ptr() != args[0].data_ptr(), "hc_out is fresh"
+ assert normed.is_contiguous()
+ if args[6] is None:
+ _compare((hc, normed), op.reference(*args))
+ else:
+ _compare_moe(args, (hc, normed), op.reference(*args))
+
+
+def test_fold_cute_is_the_triton_kernel_bitwise(monkeypatch) -> None:
+ """Below `_FEW_WARPS_MIN_ROWS` rows the fold runs the CuTe DSL kernel; its
+ `hc_out` and `normed` are bitwise the Triton kernel's (run by refusing the
+ CuTe guard) at the last row count before the switch."""
+ rows = M._FEW_WARPS_MIN_ROWS - 1
+ args = _build(dict(rows=rows, moe=True, fold=10), seed=500 + rows)
+ calls = []
+ cute_fold = M._fold_cute
+
+ def recording(*a):
+ calls.append(a[0].shape[0])
+ cute_fold(*a)
+
+ monkeypatch.setattr(M, "_fold_cute", recording)
+ hc, normed = op(*args)
+ assert calls == [rows], "the decode MoE site runs the CuTe kernel"
+ monkeypatch.setattr(M, "_fold_cute_ok", lambda *a: False)
+ want_hc, want_normed = op(*args)
+ assert torch.equal(hc, want_hc)
+ assert torch.equal(normed, want_normed)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_rstd.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_rstd.py
new file mode 100644
index 000000000000..5159e2281416
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_rstd.py
@@ -0,0 +1,139 @@
+# 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_rstd entry.
+
+Cells: `hc_out` within one ulp of the float64 reference (the MoE-combine cells
+widened by the combine's own rounding, as `qwen4_exp_hc_inject_norm`'s test),
+rstd within 2^-16 relative of the float64 statistic of the stored residual.
+Beyond them: `hc_out` and the normed stream rebuilt from rstd bitwise
+`qwen4_exp_hc_inject_norm`'s at the cells' sizes.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_inject_norm import ( # noqa: E501
+ qwen4_exp_hc_inject_norm,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_inject_rstd import ( # noqa: E501
+ qwen4_exp_hc_inject_rstd as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "qwen4_exp_hc_inject_rstd requires a CUDA device"
+
+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
+#: rstd band, relative: an fp32 2560-term sum of squares and an fp32 rsqrt.
+RSTD_RTOL = 2.0**-16
+
+
+def _ulp_ratio(out: torch.Tensor, ref: torch.Tensor) -> float:
+ """Worst element's distance from `ref` in units of `eps(dtype) * |ref|`."""
+ 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 _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 _compare(args, out, ref) -> None:
+ """`hc_out`: one ulp of the once-rounded float64 accumulate; with the MoE combine,
+ plus one ulp of the combine times `|inject|` (the combine is an fp32 value
+ rounded once, the reference's a float64 one). rstd: `RSTD_RTOL` of the float64
+ statistic of the residual the op stored -- the dataflow the op states.
+ """
+ hyper_input, block_out, inject, eps, c, shared, gate = args
+ (hc, rstd), (ref_hc, ref_rstd) = out, ref
+ assert hc.dtype == ref_hc.dtype and hc.shape == ref_hc.shape
+ assert rstd.dtype == torch.float32 and rstd.shape == ref_rstd.shape and rstd.is_contiguous()
+ if not hc.numel():
+ return
+ if shared is None:
+ assert _ulp_ratio(hc, ref_hc) <= 1.0, f"hc_out {_ulp_ratio(hc, ref_hc):.4g} ulp"
+ else:
+ feps, tiny = torch.finfo(ref_hc.dtype).eps, torch.finfo(ref_hc.dtype).tiny
+ b = _combine(block_out, shared, gate).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, eps, c)[1].double()
+ rel = ((rstd.double() - own).abs() / own).max().item()
+ assert rel <= RSTD_RTOL, f"rstd {rel:.3g} relative"
+
+
+def _gamma(seed):
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ width = HC_COUNT * HIDDEN
+ gamma = torch.randn(width, generator=g, device=DEV) * GAMMA_RMS
+ gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX
+ return gamma.to(torch.bfloat16)
+
+
+def _build(spec, seed):
+ rows, dtype = spec["rows"], 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)
+ 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)
+ return hyper_input, block_out, inject, EPS, HC_COUNT, shared, gate_logits
+
+
+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))
+ before = [a.clone() if isinstance(a, torch.Tensor) else a for a in args]
+ hc, rstd = 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], "statistic-only returns hyper_input itself"
+ else:
+ assert hc.is_contiguous() and hc.data_ptr() != args[0].data_ptr(), "hc_out is fresh"
+ _compare(args, (hc, rstd), op.reference(*args))
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_normed_rebuilt_from_rstd_is_inject_norm_bitwise(cell) -> None:
+ """The consumers' operand `bf16((hc_out * rstd) * (1 + w))` (fp32 products) is
+ bitwise `qwen4_exp_hc_inject_norm`'s `normed`, and `hc_out` its `hc_out`.
+ """
+ hyper_input, block_out, inject, eps, c, shared, gate = _build(cell.spec, seed=_seed(cell))
+ gamma = _gamma(_seed(cell) + 1)
+ hc, rstd = op(hyper_input, block_out, inject, eps, c, shared, gate)
+ want_hc, want_normed = qwen4_exp_hc_inject_norm(
+ hyper_input, block_out, inject, gamma, eps, c, shared, gate
+ )
+ assert torch.equal(hc, want_hc), f"{(hc != want_hc).sum().item()} hc_out elements differ"
+ x = hc.float().unflatten(-1, (c, -1)) * rstd.unsqueeze(-1)
+ normed = (x.flatten(-2) * (1.0 + gamma.float())).to(hc.dtype)
+ diff = (normed != want_normed).sum().item()
+ assert diff == 0, f"{diff} normed elements differ"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py
new file mode 100644
index 000000000000..07b816f9ae61
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py
@@ -0,0 +1,57 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the default target's qwen4_exp_hc_lowrank_act variant
+(models/qwen4_exp/ops/hc/qwen4_exp_hc_lowrank_act): cells within one ulp of the float64
+reference."""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_lowrank_act import ( # noqa: E501
+ qwen4_exp_hc_lowrank_act as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "qwen4_exp_hc_lowrank_act requires a CUDA device"
+
+DEV = "cuda"
+RANK = 320
+
+
+def _ulp_ratio(out, ref) -> float:
+ eps = torch.finfo(ref.dtype).eps
+ tiny = torch.finfo(ref.dtype).tiny
+ return ((out.double() - ref.double()).abs() / (eps * ref.double().abs() + tiny)).max().item()
+
+
+def _compare(out, ref) -> None:
+ """One ulp: each output is one fp32 activation of an fp32 input, rounded
+ once, so it can differ from the once-rounded float64 value only at a
+ rounding midpoint."""
+ for got, want in zip(out, ref):
+ assert (got is None) == (want is None)
+ if got is None:
+ continue
+ assert got.dtype == want.dtype and got.shape == want.shape and got.is_contiguous()
+ assert _ulp_ratio(got, want) <= 1.0, f"{_ulp_ratio(got, want):.4g} ulp"
+
+
+def _build(spec, seed):
+ hc_count = spec.get("hc_count", 4)
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ # The first projection's logits on this checkpoint reach |67| (down) / |48| (inject).
+ logits = torch.randn(spec["rows"], spec.get("logits_width", 384), generator=g, device=DEV) * 20
+ return logits, RANK, hc_count, spec.get("inject", True), spec.get("dtype", torch.bfloat16)
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
+ certifying(op)
+ args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31))
+ low, inject = op(*args)
+ assert low.shape == (cell.spec["rows"], RANK)
+ _compare((low, inject), op.reference(*args))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py
new file mode 100644
index 000000000000..55e12cf73622
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py
@@ -0,0 +1,67 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the qwen4_exp_hc_up_gate_mean entry: cells in the band
+`qwen4_exp_gated_residual`'s test derives for `mixed`."""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_up_gate_mean import ( # noqa: E501
+ _BLOCK_H,
+ _COMPILED,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.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 gate_ratio(out: torch.Tensor, ref: torch.Tensor) -> float:
+ """Worst element's distance from `ref` as a fraction of the test's band:
+ `rtol = eps`, `atol = eps * max|ref|`, both of `ref`'s dtype. Below 1 is
+ inside. (`qwen4_exp_gated_residual`'s band, restated so this test reads
+ no other target's code.)"""
+ eps = torch.finfo(ref.dtype).eps
+ ref64 = ref.double()
+ atol = eps * ref64.abs().max().item()
+ return ((out.double() - ref64).abs() / (atol + eps * ref64.abs())).max().item()
+
+
+def _compare(out, ref) -> None:
+ """`rtol = eps`, `atol = eps * max|ref|` (`gate_ratio <= 1`), the band
+ `qwen4_exp_gated_residual` derives for `mixed`."""
+ assert out.dtype == ref.dtype and out.shape == ref.shape and out.is_contiguous()
+ assert gate_ratio(out, ref) <= 1.0, f"{gate_ratio(out, ref):.4g}x the band"
+
+
+def _build(spec, seed):
+ hc, hidden, rank = spec.get("hc_count", 4), spec.get("hidden", 2560), spec.get("rank", 320)
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ rows, width = spec["rows"], hc * hidden
+ # This checkpoint's statistics: low = silu(.) of the bottleneck (|x| < ~3), the up
+ # weight ~N(0, 0.068) (gate logits within about +-8), the normed stream ~+-30.
+ low = torch.nn.functional.silu(torch.randn(rows, rank, generator=g, device=DEV) * 1.5).to(
+ torch.bfloat16
+ )
+ up_w = (torch.randn(width, rank, generator=g, device=DEV) * 0.068).to(torch.bfloat16)
+ normed = (torch.randn(rows, width, generator=g, device=DEV) * 6).to(torch.bfloat16)
+ return low, up_w.t(), normed, hc
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
+ certifying(op)
+ args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31))
+ mixed = op(*args)
+ assert mixed.shape == (cell.spec["rows"], cell.spec.get("hidden", 2560))
+ # The answer came from the CuTe DSL kernel, JIT-compiled for this device.
+ assert (_BLOCK_H, args[3], mixed.shape[1], args[0].shape[1], False) in _COMPILED
+ _compare(mixed, op.reference(*args))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_rstd.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_rstd.py
new file mode 100644
index 000000000000..cee012140e7b
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_rstd.py
@@ -0,0 +1,137 @@
+# 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_rstd entries (bf16 and NVFP4):
+`mixed` in the gate band of the float64 reference, the epilogue's fp32 operation
+sequence bitwise (gate at zero logits), the NVFP4 outputs against the native-torch
+quantization of the returned `mixed` outside the reciprocal tie window, the NVFP4
+build's `mixed` bitwise the bf16 build's.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import (
+ Fp4Quantize,
+ fp4_quantize,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501
+ gate_ratio,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_up_gate_mean_rstd import ( # noqa: E501
+ qwen4_exp_hc_up_gate_mean_rstd as op,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_up_gate_mean_rstd import ( # noqa: E501
+ qwen4_exp_hc_up_gate_mean_rstd_fp4 as op4,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "qwen4_exp_hc_up_gate_mean_rstd requires a CUDA device"
+
+DEV = "cuda"
+HC_COUNT = 4
+HIDDEN = 2560
+RANK = 320
+EPS = 1e-6
+GAMMA_RMS = 0.477
+GAMMA_ABSMAX = 6.625
+
+
+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 = fp4_quantize.reference(mixed, gs, 16, False, False)
+ assert xq.dtype == torch.uint8 and xq.shape == want_q.shape and xq.is_contiguous()
+ assert xsf.dtype == torch.uint8 and xsf.shape == want_sf.shape
+ assert torch.equal(xsf, want_sf), f"{(xsf != want_sf).sum().item()} scale bytes differ"
+ diff = _unpack(xq) != _unpack(want_q)
+ stray = diff & ~Fp4Quantize.near_tie(fp4_quantize.scaled(mixed, gs, 16))
+ assert not stray.any(), f"{stray.sum().item()} e2m1 codes differ outside the tie window"
+
+
+def _check_mixed(mixed, ref_mixed) -> None:
+ """gate_ratio <= 1, the band `qwen4_exp_hc_up_gate_mean` is certified in."""
+ assert mixed.dtype == ref_mixed.dtype and mixed.shape == ref_mixed.shape
+ assert mixed.is_contiguous()
+ assert gate_ratio(mixed, ref_mixed) <= 1.0, f"{gate_ratio(mixed, ref_mixed):.4g}x the band"
+
+
+def _build(spec, seed):
+ """`(low, up, hc_out, rstd, norm_w1, hc_count, global_scale)`."""
+ rows, width = spec["rows"], HC_COUNT * HIDDEN
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ # This checkpoint's statistics: low = silu(.) of the bottleneck (|x| < ~3), the up
+ # weight ~N(0, 0.068) (gate logits within about +-8).
+ 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)
+ hc_out = (torch.randn(rows, width, generator=g, device=DEV) * 3.0).to(torch.bfloat16)
+ x = hc_out.double().unflatten(-1, (HC_COUNT, HIDDEN))
+ rstd = torch.rsqrt(x.pow(2).mean(-1) + EPS).float().contiguous()
+ gamma = torch.randn(width, generator=g, device=DEV) * GAMMA_RMS
+ gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX
+ norm_w1 = (1.0 + gamma.to(torch.bfloat16).float()).contiguous()
+ # A calibrated activation global scale: 448 * 6 / amax over a typical batch.
+ gs = torch.tensor([448.0 * 6.0 / 4.0], device=DEV, dtype=torch.float32)
+ return low, up_w.t(), hc_out, rstd, norm_w1, HC_COUNT, 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))[:6]
+ mixed = op(*args)
+ assert mixed.shape == (cell.spec["rows"], HIDDEN)
+ _check_mixed(mixed, op.reference(*args))
+
+
+@pytest.mark.parametrize("cell", op4.CELLS, ids=[c.why[:44] for c in op4.CELLS])
+def test_certified_cells_fp4(cell, certifying) -> None: # noqa: F811 — fixture by name
+ certifying(op4)
+ args = _build(cell.spec, seed=_seed(cell))
+ mixed, xq, xsf = op4(*args)
+ assert mixed.shape == (cell.spec["rows"], HIDDEN)
+ ref_mixed = op4.reference(*args)[0]
+ _check_mixed(mixed, ref_mixed)
+ _check_quant(mixed, xq, xsf, args[6])
+
+
+@pytest.mark.parametrize("rows", [5, 300, 4100])
+def test_epilogue_operation_sequence_bitwise(rows) -> None:
+ """With `low = 0` every logit is 0 and `tanh.approx(0) = 0`, so the gate is the
+ folded `rstd / 2C` and `mixed` is the documented fp32 sequence alone:
+ `acc = (rstd_0 / 8) * (x_0 * w1_0)`, then `acc = fma(rstd_s / 8, x_s * w1_s, acc)`
+ in stream order, one bf16 rounding. Each step is emulated in float64 (products
+ of the fp32 operands are exact there) and rounded to fp32.
+ """
+ low, up, hc_out, rstd, w1, c, _ = _build(dict(rows=rows), seed=200 + rows)
+ low = torch.zeros_like(low)
+ mixed = op(low, up, hc_out, rstd, w1, c)
+ x = hc_out.double().unflatten(-1, (c, HIDDEN))
+ xw = (x * w1.double().unflatten(-1, (c, HIDDEN))).float().double()
+ hr = (rstd * (0.5 / c)).double()
+ acc = (hr[:, :1] * xw[:, 0]).float().double()
+ for s in range(1, c):
+ acc = (hr[:, s : s + 1] * xw[:, s] + acc).float().double()
+ want = acc.float().to(torch.bfloat16)
+ assert torch.equal(mixed, want), f"{(mixed != want).sum().item()} elements differ"
+
+
+@pytest.mark.parametrize("rows", [5, 2048, 4100])
+def test_fp4_mixed_is_the_bf16_entry_bitwise(rows) -> None:
+ args = _build(dict(rows=rows), seed=100 + rows)
+ assert torch.equal(op4(*args)[0], op(*args[:6]))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_post_inject.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_post_inject.py
new file mode 100644
index 000000000000..917c563af1b0
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_post_inject.py
@@ -0,0 +1,76 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the qwen4_exp_post_inject catalog entry.
+
+The cells are gated bitwise (16-bit) by `_compare`.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_post_inject import ( # noqa: E501
+ qwen4_exp_post_inject as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "qwen4_exp_post_inject requires a CUDA device"
+
+
+def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
+ """Bitwise for 16-bit operands: the product of two bf16 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()
+ assert torch.equal(out, ref), (out.float() - ref.float()).abs().max().item()
+
+
+DEV = "cuda"
+HC_COUNT = 4
+HIDDEN = 2560
+
+
+def _operands(leading, hc_count=HC_COUNT, hidden=HIDDEN, dtype=torch.bfloat16, seed=0):
+ """Operands in the ranges the model produces: `inject` is 2*sigmoid(.), so [0, 2]."""
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ hyper_input = torch.randn(*leading, hc_count * hidden, generator=g, device=DEV).to(dtype)
+ block_out = torch.randn(*leading, hidden, generator=g, device=DEV).to(dtype)
+ inject = (2.0 * torch.sigmoid(torch.randn(*leading, hc_count, generator=g, device=DEV))).to(
+ dtype
+ )
+ return hyper_input, block_out, inject
+
+
+def _build(spec, seed):
+ hc_count, hidden = spec["hc_count"], spec["hidden"]
+ dtype = spec.get("dtype", torch.bfloat16)
+ hyper_input, block_out, inject = _operands(spec["leading"], hc_count, hidden, dtype, seed)
+ layout = spec.get("layout", "contiguous")
+ width = hc_count * hidden
+ if layout == "col_slice":
+ g = torch.Generator(device=DEV).manual_seed(seed + 1)
+ buf = torch.randn(hyper_input.shape[0], 2 * width, generator=g, device=DEV).to(dtype)
+ buf[:, :width] = hyper_input
+ hyper_input = buf[:, :width]
+ # Strides rather than is_contiguous(): at one row torch ignores the
+ # size-1 dim's stride and calls the slice contiguous.
+ assert hyper_input.stride() == (2 * width, 1), hyper_input.stride()
+ assert hyper_input.shape[0] == 1 or not hyper_input.is_contiguous()
+ return hyper_input, block_out, inject
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
+ certifying(op)
+ hyper_input, block_out, inject = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
+ before = hyper_input.clone()
+ out = op(hyper_input, block_out, inject)
+ assert out.shape == hyper_input.shape and out.dtype == hyper_input.dtype
+ assert out.data_ptr() != hyper_input.data_ptr(), "the result is a fresh tensor"
+ assert torch.equal(hyper_input, before), "the residual it was given is read-only"
+ _compare(out, op.reference(hyper_input, block_out, inject))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py
new file mode 100644
index 000000000000..a75ead10c13d
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py
@@ -0,0 +1,366 @@
+# 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 claim no cell carries:
+per-pair invariance across tiles and tactics.
+"""
+
+import functools
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion import ( # noqa: E501
+ E2M1,
+ E2M1_MAX,
+ GATE_UP_GROUP,
+ Fc1Expected,
+ e2m1_codes,
+ e4m3_value,
+ unswizzle,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion import ( # noqa: E501
+ cute_dsl_nvfp4_gather_grouped_gemm_act_fusion as op,
+)
+from tensorrt_llm._torch.autotuner import AutoTuner
+from tensorrt_llm._torch.utils import ActivationType
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "the CuteDSL FC1 requires a CUDA device"
+
+DEV = "cuda"
+E, K, H, INTER, VEC = 512, 10, 2560, 640, 16
+SWIGLU = int(ActivationType.Swiglu)
+#: The activation's input scale: |x| = 5 maps to the e4m3 maximum, as the
+#: checkpoint's per-layer `input_scale` does for its own calibration maximum.
+ISC1 = 5.0 / (448 * 6)
+#: The SwiGLU output maximum `global_sf` is calibrated to (see `_ckpt` in the entry).
+ACT_AMAX = 48.0
+GSF = {
+ "calibrated": 448 * 6 / ACT_AMAX,
+ "subnormal": 2.0**-12 * 448 * 6 / ACT_AMAX,
+}
+
+#: Tie windows of `_compare`, derived in its docstring.
+SF_TIE = 2.0**-20
+CODE_TIE = 2.0**-17
+#: At most this fraction of block scales / codes may take the other side of a tie
+#: (4 of them when that is more).
+FLIP_BUDGET = 1e-5
+
+
+# ── the gate ──────────────────────────────────────────────────────────────
+
+
+def _e2m1_rank(packed: torch.Tensor) -> torch.Tensor:
+ """Packed codes -> signed magnitude rank in -7..7 per element; -0 and +0 are both 0."""
+ lo, hi = (packed & 0xF).long(), (packed >> 4).long()
+ codes = torch.stack([lo, hi], dim=-1).reshape(packed.shape[0], -1)
+ return torch.where(codes >= 8, -(codes & 7), codes & 7)
+
+
+def _e2m1_rank_of(v: torch.Tensor) -> torch.Tensor:
+ codes = e2m1_codes(v).long()
+ return torch.where(codes >= 8, -(codes & 7), codes & 7)
+
+
+def _rank_value(rank: torch.Tensor) -> torch.Tensor:
+ lut = torch.tensor(E2M1, dtype=torch.float64, device=rank.device)
+ return lut[rank.abs()] * torch.sign(rank).double()
+
+
+def _deviation(out: tuple[torch.Tensor, torch.Tensor], ref: Fc1Expected) -> dict:
+ """Every disagreement between the op and the reference on the rows it writes.
+
+ Block scales: equal, or the two neighbours of an e4m3 midpoint that the exact
+ pre-rounding scale lies within `SF_TIE` (relative) of -- a `sf_flip`; anything
+ else is `sf_outside`. Codes are re-derived under the op's OWN block scale (so
+ a flipped scale does not cascade into its 16 codes): equal, or the two
+ neighbours of an e2m1 midpoint that the exact `scaled / sf` lies within
+ `CODE_TIE` of -- a `code_flip`; anything else is `code_outside`. Blocks whose
+ scale is zero dequantize to zero whatever their codes hold and are not
+ code-checked. `*_far` counts the disagreements that are not between
+ neighbours at all.
+ """
+ c, c_sf = out
+ m, half = c.shape
+ inter = half * 2
+ nb = c_sf.numel() // m
+ vec = inter // nb
+ rows = ref.rows
+ s_op = unswizzle(c_sf.view(torch.uint8), m, nb)[rows].long()
+ s_ref = unswizzle(ref.c_sf.view(torch.uint8), m, nb)[rows].long()
+ pre = ref.scaled.view(-1, nb, vec).abs().amax(-1) / E2M1_MAX
+
+ v_op, v_ref = e4m3_value(s_op), e4m3_value(s_ref)
+ sf_diff = s_op != s_ref
+ finite = (s_op < 0x7F) & (s_ref < 0x7F)
+ sf_adj = sf_diff & finite & ((s_op - s_ref).abs() == 1)
+ sf_dist = ((pre - (v_op + v_ref) / 2).abs() / pre.clamp_min(1e-300)).masked_fill(~sf_adj, 0.0)
+ sf_flip = sf_adj & (sf_dist <= SF_TIE)
+ sf_outside = sf_diff & ~sf_flip
+
+ rank_op = _e2m1_rank(c.view(torch.uint8)[rows])
+ scale = v_op.repeat_interleave(vec, dim=1)
+ live = scale > 0
+ q = torch.where(live, ref.scaled / scale, torch.zeros_like(scale))
+ rank_exp = _e2m1_rank_of(q)
+ code_diff = live & (rank_op != rank_exp)
+ code_adj = code_diff & ((rank_op - rank_exp).abs() == 1)
+ mid = (_rank_value(rank_op) + _rank_value(rank_exp)) / 2
+ code_dist = (q - mid).abs().masked_fill(~code_adj, 0.0)
+ code_flip = code_adj & (code_dist <= CODE_TIE)
+ code_outside = code_diff & ~code_flip
+
+ return dict(
+ rows=int(rows.sum()),
+ blocks=s_op.numel(),
+ sf_flips=int(sf_flip.sum()),
+ sf_outside=int(sf_outside.sum()),
+ sf_far=int((sf_diff & ~sf_adj).sum()),
+ elements=int(live.sum()),
+ code_flips=int(code_flip.sum()),
+ code_outside=int(code_outside.sum()),
+ code_far=int((code_diff & ~code_adj).sum()),
+ )
+
+
+def _compare(out: tuple[torch.Tensor, torch.Tensor], ref: Fc1Expected) -> None:
+ """Bitwise on every written row, except at rounding ties of bounded width.
+
+ The op accumulates in fp32 on the tensor cores, evaluates SwiGLU with
+ `ex2.approx` / `rcp.approx`, and divides by the block scale through
+ `rcp.approx`; the reference is exact. The op's pre-rounding value is
+ therefore off by a handful of fp32 roundings (2^-24 relative each), and
+ only where the exact value lies that close to a rounding midpoint may the
+ op land on its other side. Everywhere else the two agree bit for bit.
+
+ - Block scale: a flip only within `SF_TIE` = 2^-20 (16 fp32 roundings)
+ relative of an e4m3 midpoint.
+ - Code: a flip only within `CODE_TIE` = 2^-17 of an e2m1 midpoint, in units
+ of the block scale -- 16 fp32 roundings of a block maximum of at most 6,
+ rounded up to a power of two.
+ - Flips: at most `FLIP_BUDGET` = 1e-5 of the scales and of the codes (4 when
+ that is fewer), so the op may not take the far side of every near-tie.
+
+ Every wrongly prepared operand puts most of the block scales outside the gate.
+ """
+ c, c_sf = out
+ assert c.dtype == torch.float4_e2m1fn_x2, c.dtype
+ assert c_sf.dtype == torch.uint8, c_sf.dtype
+ assert c.shape == ref.c.shape, (tuple(c.shape), tuple(ref.c.shape))
+ assert c_sf.shape == ref.c_sf.shape, (tuple(c_sf.shape), tuple(ref.c_sf.shape))
+ d = _deviation(out, ref)
+ assert d["rows"] > 0, d
+ assert d["sf_outside"] == 0, f"block scales off by more than a tie: {d}"
+ assert d["code_outside"] == 0, f"codes off by more than a tie: {d}"
+ assert d["sf_flips"] <= max(4, FLIP_BUDGET * d["blocks"]), f"too many scale ties flipped: {d}"
+ assert d["code_flips"] <= max(4, FLIP_BUDGET * d["elements"]), (
+ f"too many code ties flipped: {d}"
+ )
+
+
+# ── operands ──────────────────────────────────────────────────────────────
+
+
+@functools.lru_cache(maxsize=None)
+def _experts(num_experts: int = E):
+ """Checkpoint-format routed experts: per expert `[up; gate]` NVFP4 rows with
+ LINEAR e4m3 block scales and one `weight_scale_2` shared by both halves, as the
+ checkpoint stores them. Expert magnitudes spread over 2^-1..2^1 so a wrong
+ per-expert alpha shows."""
+ g = torch.Generator(device=DEV).manual_seed(20261013)
+ spread = torch.exp2(torch.rand(num_experts, generator=g, device=DEV) * 2 - 1)
+ w = torch.empty(num_experts, 2 * INTER, H // 2, dtype=torch.uint8, device=DEV)
+ s = torch.empty(num_experts, 2 * INTER, H // VEC, dtype=torch.uint8, device=DEV)
+ ws2 = torch.empty(num_experts, dtype=torch.float32, device=DEV)
+ for e in range(num_experts):
+ dense = torch.randn(2 * INTER, H, generator=g, device=DEV) * (float(spread[e]) / H**0.5)
+ dense = dense.to(torch.bfloat16)
+ gs = (448 * 6 / dense.abs().amax().float()).reshape(1)
+ q, sf = torch.ops.trtllm.fp4_quantize(dense, gs, VEC, False, False)
+ w[e] = q.view(2 * INTER, H // 2)
+ s[e] = sf.view(2 * INTER, H // VEC)
+ ws2[e] = 1.0 / gs
+ return w, s, ws2
+
+
+def _row_perm(n: int) -> torch.Tensor:
+ """Destination row -> source row of the `[up; gate]` stack, in `GATE_UP_GROUP`-row
+ groups."""
+ idx = torch.arange(n, device=DEV)
+ return idx.view(2, n // (2 * GATE_UP_GROUP), GATE_UP_GROUP).transpose(0, 1).reshape(-1)
+
+
+@functools.lru_cache(maxsize=None)
+def _weights():
+ """`(weight, weight_scale, alpha)` in the layout CuteDSL reads -- what the
+ target's load-time transform produces."""
+ w, s, ws2 = _experts()
+ perm = _row_perm(2 * INTER)
+ w_k = w[:, perm].contiguous()
+ s_k = torch.ops.trtllm.block_scale_interleave(s[:, perm].contiguous()).view(s.shape)
+ alpha = (ISC1 * ws2).contiguous()
+ return w_k.view(torch.float4_e2m1fn_x2), s_k, alpha
+
+
+def _activation(tokens: int, seed: int, zero_tokens: int = 0):
+ """The target's FC1 input: bf16 hidden rows through `fp4_quantize(x, 1/isc1,
+ 16, False, False)`, scales viewed `[T, H/16]`."""
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ x = torch.randn(tokens, H, generator=g, device=DEV).to(torch.bfloat16)
+ x[:zero_tokens] = 0
+ gs = torch.tensor([1.0 / ISC1], dtype=torch.float32, device=DEV)
+ q, sf = torch.ops.trtllm.fp4_quantize(x, gs, VEC, False, False)
+ return q.view(torch.float4_e2m1fn_x2), sf.view(tokens, H // VEC)
+
+
+def _ids(dist: str, tokens: int, seed: int) -> torch.Tensor:
+ """Expert ids `[tokens, K]` int32: each row's top-K of uniform scores."""
+ assert dist == "router", dist
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ scores = torch.rand(tokens, E, device=DEV, generator=g)
+ return scores.topk(K, dim=-1).indices.to(torch.int32)
+
+
+def _sort(ids: torch.Tensor, tile: int):
+ """moe_sort as the CuteDSL MoE calls it; the routing weights do not reach FC1."""
+ scales = torch.full(ids.shape, 0.1, dtype=torch.float32, device=DEV)
+ return torch.ops.trtllm.moe_sort(
+ token_selected_experts=ids,
+ token_final_scales=scales,
+ num_experts=E,
+ top_k=K,
+ local_expert_offset=0,
+ local_num_experts=E,
+ tile_tokens_dim=tile,
+ )
+
+
+def _args(x, x_sf, sort, tile, gsf):
+ """The keyword arguments the target passes, by the op's own names."""
+ weight, weight_scale, alpha = _weights()
+ t2e, t2mn, _e2p, p2e, _total, n_tiles = sort
+ return dict(
+ input=x,
+ weight=weight,
+ input_scale=x_sf,
+ weight_scale=weight_scale,
+ alpha=alpha,
+ tile_idx_to_group_idx=t2e,
+ tile_idx_to_mn_limit=t2mn,
+ permuted_idx_to_expanded_idx=p2e,
+ num_non_exiting_tiles=n_tiles,
+ global_sf=torch.tensor([gsf], dtype=torch.float32, device=DEV),
+ num_experts=E,
+ top_k=K,
+ num_local_experts=E,
+ local_expert_offset=0,
+ tile_size=tile,
+ scaling_vector_size=VEC,
+ activation_type=SWIGLU,
+ )
+
+
+def _problem(tokens, tile, seed, dist="router", gsf="calibrated", zero_tokens=0):
+ x, x_sf = _activation(tokens, seed, zero_tokens)
+ sort = _sort(_ids(dist, tokens, seed + 1), tile)
+ return _args(x, x_sf, sort, tile, GSF[gsf]), sort
+
+
+def _call(args, mma_n):
+ """One op call with the autotuner tactic whose MMA tile is `(tile_size, mma_n)`,
+ selected through the tuner's own capture/replay so the call still goes through
+ the op."""
+ tuner = AutoTuner.get()
+ with tuner.capture() as cap:
+ op(**args)
+ for config in cap:
+ ((_runner, tactic),) = config
+ if isinstance(tactic, tuple) and tactic[0] == (args["tile_size"], mma_n):
+ with tuner.replay(config):
+ return op(**args)
+ raise AssertionError(f"no tactic with MMA tile ({args['tile_size']}, {mma_n})")
+
+
+def _pairs(out, sort, tokens):
+ """Each (token, slot) pair's output row and its block-scale row, in expanded
+ order -- what FC2 and the combine see of it, whatever row the sort gave it."""
+ c, c_sf = out
+ e2p = sort[2].reshape(-1).long()[: tokens * K]
+ sf = unswizzle(c_sf, c.shape[0], c_sf.numel() // c.shape[0])
+ return c.view(torch.uint8)[e2p], sf[e2p]
+
+
+def _clone(v):
+ """A copy that does not depend on copy kernels for the fp4 dtype."""
+ if not isinstance(v, torch.Tensor):
+ return v
+ return v.view(torch.uint8).clone().view(v.dtype) if v.element_size() == 1 else v.clone()
+
+
+def _assert_pairs_equal(a, b, what):
+ assert torch.equal(a[0], b[0]), f"{what}: codes differ"
+ assert torch.equal(a[1], b[1]), f"{what}: block scales differ"
+
+
+def _cell_id(cell):
+ s = cell.spec
+ return (
+ f"t{s['tile_size']}-n{s['mma_n']}-T{s['tokens']}-{s['dist']}-{s['gsf']}-z{s['zero_tokens']}"
+ )
+
+
+# ── the certified cells ───────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[_cell_id(c) for c in op.CELLS])
+def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name
+ certifying(op)
+ s = cell.spec
+ assert (s["num_experts"], s["top_k"], s["hidden"], s["intermediate"]) == (E, K, H, INTER)
+ seed = zlib.crc32(cell.why.encode())
+ args, sort = _problem(s["tokens"], s["tile_size"], seed, s["dist"], s["gsf"], s["zero_tokens"])
+ before = {k: _clone(v) for k, v in args.items() if isinstance(v, torch.Tensor)}
+ out = _call(args, s["mma_n"])
+ _compare(out, op.reference(**args))
+ for k, v in before.items():
+ assert torch.equal(args[k].view(torch.uint8), v.view(torch.uint8)), f"{k} was written"
+ if s["zero_tokens"]:
+ zero = sort[2][: s["zero_tokens"]].reshape(-1).long()
+ c, c_sf = out
+ sf = unswizzle(c_sf, c.shape[0], INTER // VEC)
+ assert bool((sf[zero] == 0).all()), "an all-zero token's block scales are not zero"
+ assert bool((c.view(torch.uint8)[zero] & 0x77 == 0).all()), "a zero token's codes are not 0"
+
+
+# ── invariance: a pair's output depends on the pair alone ─────────────────
+
+
+def test_tiles_and_tactics_agree_per_pair(certifying): # noqa: F811
+ """Tile 128 / 256, MMA N 128 / 256 and the tuner's fallback all produce every
+ pair's codes and block scales bit for bit the same at the 8192-row prefill, so a
+ cell certifies the tactic the autotuner picks in serving as well as the one it
+ ran."""
+ certifying(op)
+ tokens = 8192
+ x, x_sf = _activation(tokens, seed=tokens)
+ ids = _ids("router", tokens, seed=tokens + 1)
+ seen = {}
+ for tile in (128, 256):
+ sort = _sort(ids, tile)
+ args = _args(x, x_sf, sort, tile, GSF["calibrated"])
+ for mma_n in (128, 256):
+ seen[(tile, mma_n)] = _pairs(_call(args, mma_n), sort, tokens)
+ seen[(tile, "fallback")] = _pairs(op(**args), sort, tokens)
+ if tile == 128:
+ _compare(op(**args), op.reference(**args))
+ first = seen[(128, 128)]
+ for key, got in seen.items():
+ _assert_pairs_equal(got, first, f"T={tokens} {key} vs (128, 128)")
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py
new file mode 100644
index 000000000000..b7f9d97efb50
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py
@@ -0,0 +1,348 @@
+# 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 claim no cell carries: per-pair invariance across tiles and tactics.
+"""
+
+import functools
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.cute_dsl_nvfp4_grouped_gemm import ( # noqa: E501
+ Fc2Expected,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.cute_dsl_nvfp4_grouped_gemm import ( # noqa: E501
+ cute_dsl_nvfp4_grouped_gemm as op,
+)
+from tensorrt_llm._torch.autotuner import AutoTuner
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "the CuteDSL FC2 requires a CUDA device"
+
+DEV = "cuda"
+E, K, H, INTER, VEC = 512, 10, 2560, 640, 16
+NB = INTER // VEC
+#: The SwiGLU output maximum FC1's `global_sf` is calibrated to (see `_ckpt` in the entry).
+ACT_AMAX = 48.0
+GSF = {
+ "calibrated": 448 * 6 / ACT_AMAX,
+ "subnormal": 2.0**-12 * 448 * 6 / ACT_AMAX,
+}
+
+#: Tie window of `_compare`, in units of the element's accumulation magnitude
+#: (derived in its docstring).
+TIE = 2.0**-20
+#: At most this fraction of the gated elements with a nonzero magnitude may take
+#: the other side of a tie (4 of them when that is more).
+FLIP_BUDGET = 1e-5
+
+
+# ── the gate ──────────────────────────────────────────────────────────────
+
+
+def _bf16_order(x: torch.Tensor) -> torch.Tensor:
+ """bf16 -> an integer that is monotonic in its value, consecutive between
+ neighbouring bf16 values, and equal for +0 and -0."""
+ bits = x.view(torch.int16).int() & 0xFFFF
+ mag = bits & 0x7FFF
+ return torch.where(bits >= 0x8000, -mag, mag)
+
+
+def _deviation(out: torch.Tensor, ref: Fc2Expected) -> dict:
+ """Every disagreement between the op and the reference on the rows it writes.
+
+ An element is `equal` when its value matches the reference's `out` (+0 and
+ -0 alike). Otherwise it is a `flip` when the two are neighbouring bf16
+ values and the exact value lies within `TIE * magnitude` of the midpoint
+ between them, and `outside` when it is anything else. `far` counts the
+ disagreements that are not between neighbours at all.
+ """
+ span = ref.exact.shape[0]
+ idx = torch.nonzero(ref.rows[:span]).squeeze(1)
+ o = out[idx]
+ r = ref.out[idx]
+ v = ref.exact[idx]
+ s = ref.magnitude[idx]
+ of, rf = o.double(), r.double()
+ finite = torch.isfinite(of)
+ diff = (of != rf) | ~finite
+ adj = diff & finite & ((_bf16_order(o) - _bf16_order(r)).abs() == 1)
+ dist = ((v - (of + rf) / 2).abs() / s.clamp_min(1e-300)).masked_fill(~adj, 0.0)
+ flip = adj & (dist <= TIE)
+ outside = diff & ~flip
+ return dict(
+ rows=int(idx.numel()),
+ elements=o.numel(),
+ nonzero=int((s > 0).sum()),
+ nonfinite=int((~finite).sum()),
+ flips=int(flip.sum()),
+ outside=int(outside.sum()),
+ far=int((diff & ~adj).sum()),
+ )
+
+
+def _compare(out: torch.Tensor, ref: Fc2Expected) -> None:
+ """Bitwise on every written row, except at bf16 rounding ties of derived width.
+
+ The op accumulates the 640 products of a row in fp32 on the tensor cores --
+ ten MMA k-steps of 64 -- multiplies by alpha in fp32 and rounds to bf16. Each
+ step rounds a partial sum no larger than the element's `magnitude` (`|alpha| *
+ sum |a_k w_k|`), so the op's fp32 value is within 11 roundings of `2^-24 *
+ magnitude` of the exact one; only where the exact value lies that close to a
+ bf16 midpoint may the op land on the other neighbour. Everywhere else the two
+ agree bit for bit.
+
+ - Flip: only to the neighbouring bf16 value, and only within `TIE` = 2^-20
+ (16 fp32 roundings) of `magnitude` from the midpoint.
+ - Flips: at most `FLIP_BUDGET` = 1e-5 of the gated elements whose magnitude is
+ nonzero (4 when that is fewer) -- padding rows do not widen the budget.
+ - No NaN or infinity.
+
+ Every wrongly prepared operand puts most of the pair rows' elements outside the
+ gate, most of them not even a bf16 neighbour away.
+ """
+ assert out.dtype == torch.bfloat16, out.dtype
+ assert out.shape == ref.out.shape, (tuple(out.shape), tuple(ref.out.shape))
+ d = _deviation(out, ref)
+ assert d["rows"] > 0, d
+ assert d["nonfinite"] == 0, f"non-finite output on written rows: {d}"
+ assert d["outside"] == 0, f"elements off by more than a tie: {d}"
+ assert d["flips"] <= max(4, FLIP_BUDGET * d["nonzero"]), f"too many ties flipped: {d}"
+
+
+# ── operands ──────────────────────────────────────────────────────────────
+
+
+def _swizzle(sf: torch.Tensor) -> torch.Tensor:
+ """`[rows, cols]` scale bytes -> the flat 128x4-swizzled buffer; inverse of `unswizzle`."""
+ rows, cols = sf.shape
+ return sf.reshape(rows // 128, 4, 32, cols // 4, 4).permute(0, 3, 2, 1, 4).reshape(-1)
+
+
+def _quantized_experts(num_experts: int, rows: int, cols: int, seed: int):
+ """Per expert, a Gaussian `[rows, cols]` matrix quantized the way the checkpoint
+ stores it: NVFP4 codes, LINEAR e4m3 block scales, one `weight_scale_2`. Expert
+ magnitudes spread over 2^-1..2^1 so a wrong per-expert alpha shows."""
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ spread = torch.exp2(torch.rand(num_experts, generator=g, device=DEV) * 2 - 1)
+ w = torch.empty(num_experts, rows, cols // 2, dtype=torch.uint8, device=DEV)
+ s = torch.empty(num_experts, rows, cols // VEC, dtype=torch.uint8, device=DEV)
+ ws2 = torch.empty(num_experts, dtype=torch.float32, device=DEV)
+ for e in range(num_experts):
+ dense = torch.randn(rows, cols, generator=g, device=DEV) * (float(spread[e]) / cols**0.5)
+ dense = dense.to(torch.bfloat16)
+ gs = (448 * 6 / dense.abs().amax().float()).reshape(1)
+ q, sf = torch.ops.trtllm.fp4_quantize(dense, gs, VEC, False, False)
+ w[e] = q.view(rows, cols // 2)
+ s[e] = sf.view(rows, cols // VEC)
+ ws2[e] = 1.0 / gs
+ return w, s, ws2
+
+
+@functools.lru_cache(maxsize=None)
+def _down_proj():
+ """The checkpoint's `down_proj` experts: `[E, H, I/2]` codes, `[E, H, I/16]`
+ linear scales, `weight_scale_2` per expert."""
+ return _quantized_experts(E, H, INTER, seed=20261013)
+
+
+@functools.lru_cache(maxsize=None)
+def _weights(gsf: str = "calibrated"):
+ """`(weight, weight_scale, alpha)` in the layout CuteDSL reads -- what the target's
+ load-time transform produces. `alpha = down_proj.input_scale * weight_scale_2`,
+ with the input scale the inverse of FC1's `global_sf`."""
+ w, s, ws2 = _down_proj()
+ s_k = torch.ops.trtllm.block_scale_interleave(s.contiguous()).view(s.shape)
+ alpha = (ws2 / GSF[gsf]).contiguous()
+ return w.contiguous().view(torch.float4_e2m1fn_x2), s_k, alpha
+
+
+def _pair_rows(pairs: int, seed: int, zero_pairs: int = 0) -> torch.Tensor:
+ """SwiGLU outputs of FC1, one bf16 `[I]` row per (token, slot) pair, in expanded
+ order: `up * silu(gate)` of Gaussian halves."""
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ up = torch.randn(pairs, INTER, generator=g, device=DEV) * 2
+ gate = torch.randn(pairs, INTER, generator=g, device=DEV) * 2
+ y = (up * gate * torch.sigmoid(gate)).to(torch.bfloat16)
+ y[:zero_pairs] = 0
+ return y
+
+
+def _fc2_input(y: torch.Tensor, sort, gsf: str, garbage_seed: int = 0):
+ """This op's input as the CuteDSL FC1 returns it: each pair's row NVFP4-quantized
+ with `global_sf` at its permuted row, block scales 128x4-swizzled, the padding
+ rows of live tiles zero, and the rows past the live tiles holding stale bytes."""
+ t2e, _t2mn, e2p, p2e, _total, n_tiles = sort
+ m = p2e.numel()
+ tile = m // t2e.numel()
+ span = int(n_tiles) * tile
+ x = torch.zeros(m, INTER, dtype=torch.bfloat16, device=DEV)
+ x[e2p.reshape(-1).long()[: y.shape[0]]] = y
+ gs = torch.tensor([GSF[gsf]], dtype=torch.float32, device=DEV)
+ q, sf = torch.ops.trtllm.fp4_quantize(x, gs, VEC, False, False)
+ q = q.view(torch.uint8).reshape(m, INTER // 2)
+ sf = sf.view(torch.uint8).reshape(m, NB)
+ g = torch.Generator(device=DEV).manual_seed(garbage_seed)
+ q[span:] = torch.randint(0, 256, q[span:].shape, generator=g, device=DEV, dtype=torch.uint8)
+ sf[span:] = torch.randint(0, 256, sf[span:].shape, generator=g, device=DEV, dtype=torch.uint8)
+ return q.view(torch.float4_e2m1fn_x2), _swizzle(sf).contiguous()
+
+
+def _ids(dist: str, tokens: int, seed: int) -> torch.Tensor:
+ """Expert ids `[tokens, K]` int32: each row's top-K of uniform scores."""
+ assert dist == "router", dist
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ scores = torch.rand(tokens, E, device=DEV, generator=g)
+ return scores.topk(K, dim=-1).indices.to(torch.int32)
+
+
+def _sort(ids: torch.Tensor, tile: int):
+ """moe_sort as the CuteDSL MoE calls it; the routing weights do not reach FC2."""
+ scales = torch.full(ids.shape, 0.1, dtype=torch.float32, device=DEV)
+ return torch.ops.trtllm.moe_sort(
+ token_selected_experts=ids,
+ token_final_scales=scales,
+ num_experts=E,
+ top_k=K,
+ local_expert_offset=0,
+ local_num_experts=E,
+ tile_tokens_dim=tile,
+ )
+
+
+def _args(x, x_sf, sort, tile, gsf="calibrated"):
+ """The keyword arguments the target passes, by the op's own names."""
+ weight, weight_scale, alpha = _weights(gsf)
+ return dict(
+ input=x,
+ weight=weight,
+ input_scale=x_sf,
+ weight_scale=weight_scale,
+ alpha=alpha,
+ tile_idx_to_group_idx=sort[0],
+ num_non_exiting_tiles=sort[5],
+ num_experts=E,
+ top_k=K,
+ num_local_experts=E,
+ local_expert_offset=0,
+ tile_size=tile,
+ output_dtype=torch.bfloat16,
+ scaling_vector_size=VEC,
+ )
+
+
+def _problem(tokens, tile, seed, dist="router", gsf="calibrated", zero_tokens=0):
+ sort = _sort(_ids(dist, tokens, seed + 1), tile)
+ y = _pair_rows(tokens * K, seed, zero_pairs=zero_tokens * K)
+ x, x_sf = _fc2_input(y, sort, gsf, garbage_seed=seed + 2)
+ return _args(x, x_sf, sort, tile, gsf), sort
+
+
+def _call(args, mma_n, cluster_n):
+ """One op call with the autotuner tactic whose MMA tile is `(tile_size, mma_n)`
+ and cluster `(tile_size // 128, cluster_n)`, selected through the tuner's own
+ capture/replay so the call still goes through the op."""
+ tuner = AutoTuner.get()
+ with tuner.capture() as cap:
+ op(**args)
+ tile = args["tile_size"]
+ want = ((tile, mma_n), (tile // 128, cluster_n))
+ for config in cap:
+ ((_runner, tactic),) = config
+ if tactic == want:
+ with tuner.replay(config):
+ return op(**args)
+ raise AssertionError(f"no tactic {want}")
+
+
+def _pairs(out, sort, tokens):
+ """Each (token, slot) pair's output row, in expanded order -- what the combine
+ reads, whatever row the sort gave it."""
+ return out[sort[2].reshape(-1).long()[: tokens * K]]
+
+
+def _clone(v):
+ """A copy that does not depend on copy kernels for the fp4 dtype."""
+ if not isinstance(v, torch.Tensor):
+ return v
+ return v.view(torch.uint8).clone().view(v.dtype) if v.element_size() == 1 else v.clone()
+
+
+def _bits(t: torch.Tensor) -> torch.Tensor:
+ return t.view(torch.int16)
+
+
+def _assert_pairs_equal(a, b, what):
+ assert torch.equal(_bits(a), _bits(b)), f"{what}: rows differ"
+
+
+def _cell_id(cell):
+ s = cell.spec
+ return (
+ f"t{s['tile_size']}-n{s['mma_n']}-c{s['cluster_n']}-T{s['tokens']}-{s['dist']}"
+ f"-{s['gsf']}-z{s['zero_tokens']}"
+ )
+
+
+# ── the certified cells ───────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[_cell_id(c) for c in op.CELLS])
+def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name
+ certifying(op)
+ s = cell.spec
+ assert (s["num_experts"], s["top_k"], s["hidden"], s["intermediate"]) == (E, K, H, INTER)
+ seed = zlib.crc32(cell.why.encode())
+ args, sort = _problem(s["tokens"], s["tile_size"], seed, s["dist"], s["gsf"], s["zero_tokens"])
+ before = {k: _clone(v) for k, v in args.items() if isinstance(v, torch.Tensor)}
+ out = _call(args, s["mma_n"], s["cluster_n"])
+ _compare(out, op.reference(**args))
+ for k, v in before.items():
+ assert torch.equal(_u8(args[k]), _u8(v)), f"{k} was written"
+ if s["zero_tokens"]:
+ zero = sort[2][: s["zero_tokens"]].reshape(-1).long()
+ assert bool((out[zero] == 0).all()), "an all-zero token's rows are not zero"
+
+
+# ── invariance: a pair's row depends on the pair alone ────────────────────
+
+
+def test_tiles_and_tactics_agree_per_pair(certifying): # noqa: F811
+ """Tile 128 / 256, MMA N 128 / 256, cluster N 1 / 2 and the tuner's fallback
+ all produce every pair's row bit for bit the same at the 8192-row prefill, so a
+ cell certifies the tactic the autotuner picks in serving as well as the one it
+ ran."""
+ certifying(op)
+ tokens = 8192
+ ids = _ids("router", tokens, seed=tokens + 1)
+ y = _pair_rows(tokens * K, seed=tokens)
+ seen = {}
+ for tile in (128, 256):
+ sort = _sort(ids, tile)
+ x, x_sf = _fc2_input(y, sort, "calibrated", garbage_seed=tile)
+ args = _args(x, x_sf, sort, tile)
+ for mma_n in (128, 256):
+ for cluster_n in (1, 2):
+ seen[(tile, mma_n, cluster_n)] = _pairs(_call(args, mma_n, cluster_n), sort, tokens)
+ fallback = op(**args)
+ seen[(tile, "fallback")] = _pairs(fallback, sort, tokens)
+ _compare(fallback, op.reference(**args))
+ first = seen[(128, 128, 1)]
+ for key, got in seen.items():
+ _assert_pairs_equal(got, first, f"T={tokens} {key} vs (128, 128, 1)")
+
+
+# ── helpers ───────────────────────────────────────────────────────────────
+
+
+def _u8(t: torch.Tensor) -> torch.Tensor:
+ return t.view(torch.uint8) if t.element_size() == 1 else t
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py
new file mode 100644
index 000000000000..34ce9252cc48
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py
@@ -0,0 +1,86 @@
+# 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 uses the shared NVFP4 test helpers.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_block_scale_moe_renorm_permuted import ( # noqa: E501
+ finalize,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.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.ops.moe.fp4_block_scale_moe_renorm_routed import ( # noqa: E501
+ fp4_block_scale_moe_renorm_routed as finalized,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.renorm_moe_routing_op import (
+ renorm_moe_routing_op,
+)
+
+__extra_import_path__ = ["../../.."]
+import _fp4_moe_test_utils as H # noqa: E402
+from _fp4_moe_test_utils import many_expert_layer # noqa: E402,F401 — fixture by name
+from _precision import true_fp32_matmul # noqa: E402
+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 _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, many_expert_layer): # 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 = H.renorm_operands(
+ many_expert_layer(), spec, seed, renorm_moe_routing_op
+ )
+ 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.
+ with true_fp32_matmul():
+ exp = op.reference(logits, x, k, ref["up"], ref["gate"], ref["down"], g2)
+ H._compare(y, exp.to(torch.bfloat16))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py
new file mode 100644
index 000000000000..78019fcf751d
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py
@@ -0,0 +1,199 @@
+# 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) uses ordinary shared test helpers. Beyond the cells: the prefill
+tactic resolution bitwise the custom op's, and per `DECODE_TACTICS` bucket
+that the pinned tactic is a valid runner config and answers bitwise what the
+AutoTuner path's untuned (runner-default) tactic answers.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_block_scale_moe_renorm_routed import ( # noqa: E501
+ ACT_SWIGLU,
+ DECODE_TACTICS,
+ DECODE_TACTICS_GEOMETRY,
+ ROUTING_RENORMALIZE,
+ _ckpt,
+ _pinned_tactic,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_block_scale_moe_renorm_routed import ( # noqa: E501
+ fp4_block_scale_moe_renorm_routed as op,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.renorm_moe_routing_op import ( # noqa: E501
+ renorm_moe_routing_op,
+)
+
+__extra_import_path__ = ["../../.."]
+import _fp4_moe_test_utils as H # noqa: E402
+from _fp4_moe_test_utils import many_expert_layer # noqa: E402,F401 — fixture by name
+from _precision import true_fp32_matmul # noqa: E402
+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 _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"],
+ )
+
+
+#: The whole-layer cells (tp1); the windows have their own test file.
+_WHOLE = [c for c in op.CELLS if "local_expert_offset" not in c.spec]
+
+
+@pytest.mark.parametrize("cell", _WHOLE, ids=[c.why[:44] for c in _WHOLE])
+def test_certified_cells(cell, certifying, many_expert_layer): # noqa: F811
+ certifying(op)
+ spec = cell.spec
+ seed = zlib.crc32(cell.why.encode()) % 2**31
+ args, ref, data, sf, x, logits, ids, wts, g2, (s1, sg, s2) = H.renorm_operands(
+ many_expert_layer(), spec, seed, renorm_moe_routing_op
+ )
+
+ 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.
+ with true_fp32_matmul():
+ exp = op.reference(logits, x, spec["top_k"], ref["up"], ref["gate"], ref["down"], g2)
+ H._compare(y, exp.to(torch.bfloat16))
+
+
+def test_prefill_tactic_resolution_matches_the_custom_op_bitwise(
+ monkeypatch,
+ many_expert_layer, # noqa: F811
+):
+ """A non-pinned prefill bucket keeps the direct op's routing and tactic."""
+ spec = _ckpt(2048)
+ args, _, data, sf, _, logits, _, _, _, scalars = H.renorm_operands(
+ many_expert_layer(), spec, 821, renorm_moe_routing_op
+ )
+ module = __import__(op.__class__.__module__, fromlist=["run_renorm_moe"])
+ calls = []
+ catalog = module._pre_routed
+
+ def recorded(*args, **kwargs):
+ calls.append(True)
+ return catalog(*args, **kwargs)
+
+ monkeypatch.setattr(module, "_pre_routed", recorded)
+ got = _routed(args, data, sf, logits, scalars, spec)
+ expected = torch.ops.trtllm.fp4_block_scale_moe_runner(
+ logits,
+ None,
+ data,
+ sf.view(torch.float8_e4m3fn),
+ args["gemm1_weights"],
+ args["gemm1_weights_scale"],
+ None,
+ None,
+ None,
+ None,
+ args["gemm2_weights"],
+ args["gemm2_weights_scale"],
+ None,
+ *scalars,
+ spec["num_experts"],
+ spec["top_k"],
+ None,
+ None,
+ spec["intermediate"],
+ 0,
+ spec["num_experts"],
+ None,
+ ROUTING_RENORMALIZE,
+ True,
+ ACT_SWIGLU,
+ )[0]
+ assert calls == [True]
+ assert torch.equal(got, expected)
+ assert torch.equal(_routed(args, data, sf, logits, scalars, spec), expected)
+
+
+@pytest.mark.parametrize("bucket", sorted(DECODE_TACTICS))
+def test_decode_tactic_valid_and_bitwise(bucket, many_expert_layer): # noqa: F811
+ """The pinned tactic of each decode bucket is what a call there runs, is a valid config of
+ the runner, and gives bitwise the output of the AutoTuner path outside tuning mode (the
+ runner's default tactic for the shape)."""
+ spec = _ckpt(bucket)
+ num_experts, top_k, hidden, inter = DECODE_TACTICS_GEOMETRY
+ assert (spec["num_experts"], spec["top_k"], spec["hidden"], spec["intermediate"]) == (
+ num_experts,
+ top_k,
+ hidden,
+ inter,
+ )
+ # None (the AutoTuner path, with a warning) if the runner does not list it as valid here.
+ assert _pinned_tactic(bucket, num_experts, top_k, hidden, inter) == DECODE_TACTICS[bucket]
+ args, _, data, sf, _, logits, _, _, _, scalars = H.renorm_operands(
+ many_expert_layer(), spec, 7000 + bucket, renorm_moe_routing_op
+ )
+ y = _routed(args, data, sf, logits, scalars, spec)
+ y_tuner = torch.ops.trtllm.fp4_block_scale_moe_runner(
+ logits,
+ None,
+ data,
+ sf.view(torch.float8_e4m3fn),
+ args["gemm1_weights"],
+ args["gemm1_weights_scale"],
+ None,
+ None,
+ None,
+ None,
+ args["gemm2_weights"],
+ args["gemm2_weights_scale"],
+ None,
+ *scalars,
+ num_experts,
+ top_k,
+ None,
+ None,
+ inter,
+ 0,
+ num_experts,
+ None,
+ ROUTING_RENORMALIZE,
+ True,
+ ACT_SWIGLU,
+ )[0]
+ differing = (y.view(torch.int16) != y_tuner.view(torch.int16)).any(dim=1).sum().item()
+ assert differing == 0, f"{differing} of {bucket} rows differ from the AutoTuner path"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed_windows.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed_windows.py
new file mode 100644
index 000000000000..80a89853c909
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed_windows.py
@@ -0,0 +1,207 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the fp4_block_scale_moe_renorm_routed entry's expert windows (the whole-layer cells
+are in test_modeling_v2_fp4_block_scale_moe_renorm_routed.py).
+
+Two claims per cell: the logits-routed call over a rank's expert window is **bitwise**
+the composition it replaces at the call site (`renorm_moe_routing_op` -> pre-routed
+`fp4_block_scale_moe_runner` over the same window), and it sits inside the pre-routed
+entry's band against this entry's own fp32 reference. The NVFP4 operand construction
+(layout, the 512-expert stack, windows, activations, scalars, the band) is the pre-routed
+entry's test's, loaded by path so the two tests cannot drift apart on what a valid operand
+is; the tp4 zero-padded stack is built here from the same helpers.
+"""
+
+import importlib.util
+import zlib
+from pathlib import Path
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_block_scale_moe_renorm_routed import ( # noqa: E501
+ fp4_block_scale_moe_renorm_routed as op,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.renorm_moe_routing_op import (
+ renorm_moe_routing_op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _fp4_moe_test_utils import many_expert_layer # noqa: E402,F401 — fixture by name
+from _precision import true_fp32_matmul # noqa: E402
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "fp4_block_scale_moe_renorm_routed requires a CUDA device"
+
+DEV = "cuda"
+G1 = 137.0 # activation global scale, as the pre-routed entry's test draws it
+
+
+def _helpers():
+ path = (
+ Path(__file__).resolve().parents[3]
+ / "moe"
+ / "test_modeling_v2_fp4_block_scale_moe_runner.py"
+ )
+ spec = importlib.util.spec_from_file_location("_fp4_moe_runner_r4_test_helpers", path)
+ mod = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(mod)
+ return mod
+
+
+H = _helpers()
+
+# tp4 holds each rank's 160-wide slice of the 640 intermediate, stored zero-padded to 256.
+_TP4_LOGICAL_INTER = 160
+_TP4_PADDED_INTER = 256
+_TP4 = None
+
+
+def _pad_zero(codes: torch.Tensor, sf: torch.Tensor, dim: int, size: int):
+ """Extend a `(codes, scales)` stack with zero codes and zero (e4m3 +0.0) block scales
+ along `dim` (1 = rows, 2 = the K columns) to `size`."""
+ pad_codes = list(codes.shape)
+ pad_codes[dim] = size - codes.shape[dim]
+ pad_sf = list(sf.shape)
+ pad_sf[dim] = pad_codes[dim] if dim == 1 else pad_codes[dim] // H.SV
+ codes = torch.cat([codes, torch.zeros(pad_codes, dtype=codes.dtype, device=DEV)], dim=dim)
+ sf_u8 = torch.cat(
+ [sf.view(torch.uint8), torch.zeros(pad_sf, dtype=torch.uint8, device=DEV)], dim=dim
+ )
+ return codes.contiguous(), sf_u8.view(torch.float8_e4m3fn).contiguous()
+
+
+def _pack(codes: torch.Tensor) -> torch.Tensor:
+ return (codes[..., 0::2] | (codes[..., 1::2] << 4)).contiguous()
+
+
+def _tp4_padded():
+ """The tp4 layer, built once: kernel-ready operands at I = 256 and reference operands
+ at I = 256 (the padded ones the kernel reads)."""
+ global _TP4
+ if _TP4 is None:
+ e, h, i, ip = 512, 2560, _TP4_LOGICAL_INTER, _TP4_PADDED_INTER
+ gen = torch.Generator(device=DEV).manual_seed(512160)
+ _, up_s, up_c = H._rand_nvfp4(e, i, h, gen)
+ _, gt_s, gt_c = H._rand_nvfp4(e, i, h, gen)
+ _, dn_s, dn_c = H._rand_nvfp4(e, h, i, gen)
+ up_cp, up_sp = _pad_zero(up_c, up_s, 1, ip)
+ gt_cp, gt_sp = _pad_zero(gt_c, gt_s, 1, ip)
+ dn_cp, dn_sp = _pad_zero(dn_c, dn_s, 2, ip)
+ up_p, gt_p, dn_p = _pack(up_cp), _pack(gt_cp), _pack(dn_cp)
+ w1, s1 = H._prep_fc1(up_p, gt_p, up_sp, gt_sp)
+ w2, s2 = H._prep_fc2(dn_p, dn_sp)
+ args = dict(
+ gemm1_weights=w1,
+ gemm1_weights_scale=s1,
+ gemm2_weights=w2,
+ gemm2_weights_scale=s2,
+ intermediate_size=ip,
+ )
+ ref = dict(
+ up=(up_cp, up_sp),
+ gate=(gt_cp, gt_sp),
+ down=(dn_cp, dn_sp),
+ raw=(up_p, gt_p, up_sp, gt_sp, dn_p, dn_sp),
+ hidden=h,
+ inter=ip,
+ num_experts=e,
+ gen=gen,
+ )
+ _TP4 = (args, ref)
+ return _TP4
+
+
+def _stack(spec, layer):
+ """`(args, ref)` of the whole 512-expert stack at the cell's intermediate width."""
+ if spec["intermediate"] == 256:
+ return _tp4_padded()
+ assert spec["intermediate"] == 640, spec["intermediate"]
+ return layer()
+
+
+def _operands(spec, seed, layer):
+ gen = torch.Generator(device=DEV).manual_seed(seed)
+ args, ref = _stack(spec, layer)
+ off, n, e = spec["local_expert_offset"], spec["local_num_experts"], spec["num_experts"]
+ data, sf, x = H._rand_activation(spec["tokens"], spec["hidden"], G1, gen)
+ logits = torch.randn(spec["tokens"], e, generator=gen, device=DEV).to(spec["logits_dtype"])
+ ids, wts = renorm_moe_routing_op(logits, spec["top_k"], torch.bfloat16)
+ g2 = H._calibrate_g2(x, ids, ref)
+ s1, sg, s2 = H._scalars(e, G1, g2)
+ win_args = args if n == e else H._window_args(args, off, n)
+ win_ref = ref if n == e else H._window_ref(ref, off, n)
+ scalars = H._window_scalars(s1, sg, s2, off, n)
+ return win_args, win_ref, data, sf, x, logits, ids, wts, g2, scalars
+
+
+def _routed(win_args, data, sf, logits, scalars, spec, top_k=None, weights=None):
+ w = win_args if weights is None else weights
+ return op(
+ logits,
+ data,
+ sf.view(torch.float8_e4m3fn),
+ w["gemm1_weights"],
+ w["gemm1_weights_scale"],
+ w["gemm2_weights"],
+ w["gemm2_weights_scale"],
+ scalars["output1_scale_scalar"],
+ scalars["output1_scale_gate_scalar"],
+ scalars["output2_scale_scalar"],
+ spec["num_experts"],
+ spec["top_k"] if top_k is None else top_k,
+ spec["intermediate"],
+ spec["local_expert_offset"],
+ spec["local_num_experts"],
+ )
+
+
+#: The window cells; the whole-layer cells have their own test file.
+_WINDOWS = [c for c in op.CELLS if "local_expert_offset" in c.spec]
+
+
+@pytest.mark.parametrize("cell", _WINDOWS, ids=[c.why[:44] for c in _WINDOWS])
+def test_certified_cells(cell, certifying, many_expert_layer): # noqa: F811
+ certifying(op)
+ spec = cell.spec
+ seed = zlib.crc32(cell.why.encode()) % 2**31
+ win_args, win_ref, data, sf, x, logits, ids, wts, g2, scalars = _operands(
+ spec, seed, many_expert_layer
+ )
+ off, n = spec["local_expert_offset"], spec["local_num_experts"]
+
+ y = _routed(win_args, data, sf, logits, scalars, spec)
+ assert y.shape == (spec["tokens"], spec["hidden"]) and y.dtype == torch.bfloat16
+
+ # 1. Bitwise the composition it replaces, on identical operands.
+ y_pre = H._call(
+ data,
+ sf,
+ win_args,
+ spec["num_experts"],
+ spec["top_k"],
+ local_expert_offset=off,
+ local_num_experts=n,
+ topk_ids=ids,
+ topk_weights=wts,
+ **scalars,
+ )[0]
+ differing = (y.view(torch.int16) != y_pre.view(torch.int16)).any(dim=1).sum().item()
+ assert differing == 0, f"{differing} of {spec['tokens']} rows differ from the composition"
+
+ # 2. Its own reference, through the pre-routed entry's gate.
+ with true_fp32_matmul():
+ exp = op.reference(
+ logits,
+ x,
+ spec["top_k"],
+ win_ref["up"],
+ win_ref["gate"],
+ win_ref["down"],
+ g2,
+ local_expert_offset=off,
+ local_num_experts=n,
+ )
+ H._compare(y, exp.to(torch.bfloat16))
+ elt, rms = H._dev(y, exp.to(torch.bfloat16))
+ print(f" cell {cell.why[:64]!r}: {elt:.2f} elt / {rms:.2f} rms ulp")
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp8_block_scale_moe_runner.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp8_block_scale_moe_runner.py
new file mode 100644
index 000000000000..7fc469633394
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp8_block_scale_moe_runner.py
@@ -0,0 +1,257 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the qwen4 default target's fp8_block_scale_moe_runner variant
+(models/qwen4_exp/ops/moe/fp8_block_scale_moe_runner).
+
+The entry owns the contract: `CELLS`, `reference`, `is_valid`. This
+file turns a cell into operands and holds the claims no cell carries: a warm
+autotuned tactic is bitwise the cold fallback the cells run (whole stack and
+window), and the default target's four windows sum to the 512-expert layer.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp8_block_scale_moe_runner import ( # noqa: E501
+ BLOCK,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp8_block_scale_moe_runner import ( # noqa: E501
+ fp8_block_scale_moe_runner as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "fp8_block_scale_moe_runner requires a CUDA device"
+
+
+def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
+ """16 ulp element-wise, 0.75 ulp RMS, both in bf16 ulp of each row's max.
+
+ Both statistics normalize each element by its reference row's largest
+ magnitude: `d = (out - ref) / (rowmax * 2^-8)`, gated on `max|d|` and
+ on `sqrt(mean(d^2))` -- the RMS of the row-normalized deviation, not
+ `assert_within_ulp`'s ratio of norms, which reads higher on the same
+ data (a row's max sits several sigma above its RMS). Kernel and
+ reference consume bit-identical operands and model both
+ requantizations, so they differ by accumulation order and by FC1 sums
+ that land on an e4m3 rounding boundary.
+
+ Measured worst over the cells: 6.65 element-wise (T=4096; the max over
+ T*H samples grows with T -- 0.91 at T=1) and 0.306 RMS (H=I=512), the
+ RMS sitting at 0.17-0.31 everywhere; earlier draws reached 10.03 and
+ 0.342. The default target's 128-expert window cells, and the four windows summed in
+ fp32, sit in the same range. Wrong references at the checkpoint geometry, RMS: gate/up swap
+ 65, transposed activation scales 43, strided hidden_states 94, no
+ FC1-output requant 3.5, neither requant 3.1, no activation requant 1.58
+ -- the closest.
+ """
+ assert out.dtype == ref.dtype == torch.bfloat16, (out.dtype, ref.dtype)
+ assert out.shape == ref.shape, (out.shape, ref.shape)
+ o, r = out.float(), ref.float()
+ rowmax = r.abs().amax(dim=-1, keepdim=True).clamp(min=1e-6)
+ d = (o - r) / (rowmax * 2.0**-8)
+ elt, rms = d.abs().max().item(), d.pow(2).mean().sqrt().item()
+ assert elt <= 16.0, f"{elt:.2f} ulp element-wise > 16"
+ assert rms <= 0.75, f"{rms:.3f} ulp RMS > 0.75"
+
+
+DEV = "cuda"
+ULP = 2.0**-8
+CKPT = dict(hidden=2560, inter=640, num_experts=512, top_k=10)
+
+# ── operands ──────────────────────────────────────────────────────────────
+
+_WEIGHTS: dict = {}
+
+
+def _weights(hidden: int, inter: int, experts: int, seed: int):
+ """`(w1, w1s, w2, w2s)` drawn directly in e4m3 / fp32, so every value is exact.
+
+ Cached: the checkpoint stack is 2.5 GB and every checkpoint-geometry cell
+ reads the same one. Drawn 32 experts at a time to cap the fp32 transient.
+ """
+ key = (hidden, inter, experts, seed)
+ if key not in _WEIGHTS:
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ w1 = torch.empty(experts, 2 * inter, hidden, dtype=torch.float8_e4m3fn, device=DEV)
+ w2 = torch.empty(experts, hidden, inter, dtype=torch.float8_e4m3fn, device=DEV)
+ for lo in range(0, experts, 32):
+ hi = min(lo + 32, experts)
+ w1[lo:hi] = torch.randn(hi - lo, 2 * inter, hidden, device=DEV, generator=g)
+ w2[lo:hi] = torch.randn(hi - lo, hidden, inter, device=DEV, generator=g)
+ w1s = 0.01 + 0.02 * torch.rand(
+ experts, 2 * inter // BLOCK, hidden // BLOCK, device=DEV, generator=g
+ )
+ w2s = 0.01 + 0.02 * torch.rand(
+ experts, hidden // BLOCK, inter // BLOCK, device=DEV, generator=g
+ )
+ _WEIGHTS[key] = (w1, w1s, w2, w2s)
+ return _WEIGHTS[key]
+
+
+def _args(
+ tokens, hidden, inter, num_experts, top_k, local_expert_offset=0, local_num_experts=None, seed=0
+):
+ """The op's arguments by name, pre-routed over all `num_experts`.
+
+ The call answers for `[local_expert_offset, local_expert_offset +
+ local_num_experts)` and its weights are that window's slice of the one
+ stack -- what a rank of the default target holds. The routing, activations
+ and stack depend on `seed` and the geometry only, so the windows of one
+ seed split one layer.
+ """
+ local = num_experts if local_num_experts is None else local_num_experts
+ window = slice(local_expert_offset, local_expert_offset + local)
+ w1, w1s, w2, w2s = _weights(hidden, inter, num_experts, seed=hidden + inter + num_experts)
+ w1, w1s, w2, w2s = w1[window], w1s[window], w2[window], w2s[window]
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ hs = (2.0 * torch.randn(tokens, hidden, device=DEV, generator=g)).to(torch.float8_e4m3fn)
+ # [H/128, T]: the transposed layout fp8_quantize_1x128 returns here.
+ hs_scale = 0.5 + torch.rand(hidden // BLOCK, tokens, device=DEV, generator=g)
+ logits = torch.randn(tokens, num_experts, device=DEV, generator=g)
+ weights, ids = torch.topk(torch.softmax(logits, dim=-1), top_k, dim=-1)
+ ids = ids.to(torch.int32)
+ return dict(
+ routing_logits=None,
+ routing_bias=None,
+ hidden_states=hs,
+ hidden_states_scale=hs_scale,
+ 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,
+ routed_scaling_factor=None,
+ routing_method_type=1,
+ topk_weights=weights.to(torch.bfloat16).contiguous(),
+ topk_ids=ids.contiguous(),
+ gemm1_clamp_limit=None,
+ )
+
+
+def _ulp(out: torch.Tensor, ref: torch.Tensor):
+ """(max, RMS) deviation in bf16 ulp of the row max -- `_compare`'s statistics."""
+ o, r = out.float(), ref.float()
+ d = (o - r) / (r.abs().amax(dim=-1, keepdim=True).clamp(min=1e-6) * ULP)
+ return d.abs().max().item(), d.pow(2).mean().sqrt().item()
+
+
+def _seed(cell) -> int:
+ return zlib.crc32(cell.why.encode())
+
+
+def _window_args(cell) -> tuple[dict, torch.Tensor]:
+ """The cell's arguments and which slots route into its window. A window cell certifies
+ the window only if some slot routes into it, so a draw whose routing misses the window
+ entirely (likely only at a row or two) is redrawn with the next seed."""
+ for attempt in range(8):
+ args = _args(**cell.spec, seed=_seed(cell) + attempt)
+ hit = _in_window(args)
+ if hit.any():
+ return args, hit
+ raise AssertionError(f"no slot routes into the window: {cell.why!r} would certify nothing")
+
+
+# ── the certified cells ───────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name
+ certifying(op)
+ args, hit = _window_args(cell)
+ out = op(**args)
+ assert out.shape == args["hidden_states"].shape and out.dtype == torch.bfloat16
+ # A row that routes nothing into the window comes back bitwise zero.
+ assert torch.equal(out[~hit.any(-1)], torch.zeros_like(out[~hit.any(-1)]))
+ ref = op.reference(**args)
+ elt, rms = _ulp(out, ref)
+ print(
+ f" {cell.why[:44]}: {elt:.2f} ulp element-wise, {rms:.3f} ulp RMS; "
+ f"{hit.sum().item()} of {hit.numel()} slots in the window"
+ )
+ _compare(out, ref)
+
+
+def _in_window(args) -> torch.Tensor:
+ """`[T, top_k]`: which slots the call's window answers for."""
+ local = args["topk_ids"].long() - args["local_expert_offset"]
+ return (local >= 0) & (local < args["local_num_experts"])
+
+
+# ── the default target's windows ──────────────────────────────────────────
+
+#: The qwen4 default target's split: four 128-expert windows of the 512, one per rank.
+WINDOWS = (0, 128, 256, 384)
+
+
+@pytest.mark.parametrize("tokens", [1, 128, 1024])
+def test_windows_sum_to_the_whole_stack(tokens):
+ """The four windows' bf16 outputs, added in fp32 (what the target's
+ all-reduce or reduce-scatter does), sit inside the gate against the
+ 512-expert reference -- and the sum with any one window left out does not.
+
+ Not bitwise to the 512-expert call: each window rounds its own partial to
+ bf16. The distance to it is printed.
+ """
+ seed = 900 + tokens
+ whole = _args(tokens, **CKPT, seed=seed)
+ parts = [
+ op(**_args(tokens, **CKPT, local_expert_offset=o, local_num_experts=128, seed=seed))
+ for o in WINDOWS
+ ]
+ total = torch.stack([p.float() for p in parts]).sum(0).to(torch.bfloat16)
+ ref = op.reference(**whole)
+ elt, rms = _ulp(total, ref)
+ elt_call, rms_call = _ulp(total, op(**whole))
+ print(
+ f" T={tokens}: windows summed vs reference {elt:.2f} / {rms:.3f} ulp; "
+ f"vs the 512-expert call {elt_call:.2f} / {rms_call:.3f} ulp"
+ )
+ _compare(total, ref)
+
+ # Negative control: drop each window that holds a slot.
+ ids = whole["topk_ids"].long()
+ for i, o in enumerate(WINDOWS):
+ if not ((ids >= o) & (ids < o + 128)).any():
+ continue
+ short = torch.stack([p.float() for j, p in enumerate(parts) if j != i]).sum(0)
+ with pytest.raises(AssertionError):
+ _compare(short.to(torch.bfloat16), ref)
+
+
+# ── what a cell cannot carry ──────────────────────────────────────────────
+
+
+@pytest.mark.parametrize(
+ "offset, local", [(0, 512), (256, 128)], ids=["whole stack", "a 128-expert window"]
+)
+def test_autotuned_tactic_matches_the_fallback_tactic(offset, local):
+ """A warm profiling cache changes the tactic, not the result -- for the
+ whole stack and for a window, which tunes under its own key
+ (`local_num_experts`) and its own config (`num_experts // local_num_experts`)."""
+ from tensorrt_llm._torch.autotuner import AutoTuner, autotune
+
+ a = _args(256, **CKPT, local_expert_offset=offset, local_num_experts=local, seed=42)
+ AutoTuner.get().clear_cache()
+ try:
+ cold = op(**a)
+ with autotune():
+ op(**a)
+ entries = len(AutoTuner.get().profiling_cache)
+ warm = op(**a)
+ print(f" warming wrote {entries} cache entries")
+ assert entries > 0, "warming profiled nothing -- the comparison would be vacuous"
+ assert torch.equal(warm, cold), "a tuned tactic changed the result"
+ _compare(warm, op.reference(**a))
+ finally:
+ AutoTuner.get().clear_cache()
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py
new file mode 100644
index 000000000000..9fb348eeb829
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py
@@ -0,0 +1,101 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the default target's fused_sigmoid_gate_mul_add variant
+(models/qwen4_exp/ops/moe/fused_sigmoid_gate_mul_add).
+
+The entry owns the contract. This file turns a cell into operands (a strided gate
+is the call site's `gug` column), locks both sides of the in-place write, and holds
+the claims no cell carries: the strided gate equal to the dense column, and the
+empty batch.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fused_sigmoid_gate_mul_add import ( # noqa: E501
+ fused_sigmoid_gate_mul_add as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "fused_sigmoid_gate_mul_add requires a CUDA device"
+
+DEV = "cuda"
+HIDDEN = 2560
+# The shared expert's fused GEMM output: gate|up (2 x 640) plus 8 gate-logit rows, the
+# logit in the first of them. A strided gate is that column, as the call site reads it.
+GUG_WIDTH = 2 * 640 + 8
+GATE_COL = 2 * 640
+
+
+def _operands(spec, seed):
+ """`(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)
+
+ final, shared = rand(t, h), rand(t, h)
+ gate = rand(t, 1, dtype=spec["gate_dtype"])
+ if spec["gate_layout"] == "strided":
+ gug = rand(t, GUG_WIDTH, dtype=spec["gate_dtype"])
+ gug[:, GATE_COL : GATE_COL + 1] = gate
+ gate = gug[:, GATE_COL : GATE_COL + 1]
+ assert not gate.is_contiguous() or t == 1
+ else:
+ assert spec["gate_layout"] == "column", spec["gate_layout"]
+ return final, gate, shared
+
+
+# ── the certified cells ───────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name
+ certifying(op)
+ final, gate, shared = _operands(cell.spec, zlib.crc32(cell.why.encode()))
+ before = {k: v.clone() for k, v in dict(final=final, gate=gate, shared=shared).items()}
+ ref = op.reference(before["final"], gate, shared)
+
+ out = op(final, gate, shared)
+
+ 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"])
+
+
+# ── what a cell cannot carry ──────────────────────────────────────────────
+
+
+def _strided_spec(tokens: int) -> dict:
+ """The certified strided-gate cell at `tokens` rows."""
+ (spec,) = [
+ c.spec
+ for c in op.CELLS
+ if c.spec["tokens"] == tokens and c.spec["gate_layout"] == "strided"
+ ]
+ return dict(spec)
+
+
+def test_strided_gate_is_the_dense_column_bit_for_bit() -> None:
+ """The call site's strided `gug` column and the same logits as a dense `[T, 1]`."""
+ spec = _strided_spec(1024)
+ final, gate, shared = _operands(spec, seed=11)
+ dense_final = final.clone()
+ op(final, gate, shared)
+ op(dense_final, gate.contiguous(), shared)
+ assert torch.equal(final, dense_final)
+
+
+def test_zero_tokens_is_a_no_op():
+ final = torch.empty(0, HIDDEN, dtype=torch.bfloat16, device=DEV)
+ shared = torch.empty(0, HIDDEN, dtype=torch.bfloat16, device=DEV)
+ gate = torch.empty(0, 1, dtype=torch.bfloat16, device=DEV)
+ returned = op(final, gate, shared)
+ assert returned.shape == (0, HIDDEN) and returned.data_ptr() == final.data_ptr()
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_sort.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_sort.py
new file mode 100644
index 000000000000..f240c17b8240
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_sort.py
@@ -0,0 +1,203 @@
+# 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 claim no cell carries: CUDA-graph replay on new ids at the largest captured
+row count.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.moe_sort import ( # noqa: E501
+ moe_sort as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "moe_sort requires a CUDA device"
+
+DEV = "cuda"
+E, K = 512, 10
+
+_NAMES = (
+ "tile_idx_to_expert_idx",
+ "tile_idx_to_mn_limit",
+ "expanded_idx_to_permuted_idx",
+ "permuted_idx_to_expanded_idx",
+ "total_num_padded_tokens",
+ "num_non_exiting_tiles",
+)
+
+
+def _compare(out: list[torch.Tensor], ref: list[torch.Tensor]) -> None:
+ """Integer outputs, so exact -- up to the one freedom the op has.
+
+ Counts and the tile arrays' written prefix are bitwise. The order of an
+ expert's rows is the kernel's atomic order (it varies between identical calls
+ from T = 5), so `permuted_idx_to_expanded_idx` is gated per expert run as a
+ set: each real row's value sorted within its run must equal the reference's
+ run. `expanded_idx_to_permuted_idx` must then be the exact inverse: -1
+ exactly where the expert is not local, and otherwise a real row of the right
+ run that maps back. Padding rows and tiles past `num_non_exiting_tiles` are
+ not written by the op and are not compared.
+ """
+ assert len(out) == len(ref) == len(_NAMES), (len(out), len(ref))
+ for name, o, r in zip(_NAMES, out, ref):
+ assert o.dtype == r.dtype == torch.int32, (name, o.dtype, r.dtype)
+ assert o.shape == r.shape, (name, tuple(o.shape), tuple(r.shape))
+ t2e, t2mn, e2p, p2e, total, n_tiles = out
+ r_t2e, r_t2mn, r_e2p, r_p2e, r_total, r_n_tiles = ref
+ assert torch.equal(n_tiles.cpu(), r_n_tiles.cpu()), (n_tiles, r_n_tiles)
+ assert torch.equal(total.cpu(), r_total.cpu()), (total, r_total)
+ n = int(r_n_tiles)
+ assert torch.equal(t2e[:n], r_t2e[:n]), "tile_idx_to_expert_idx differs"
+ assert torch.equal(t2mn[:n], r_t2mn[:n]), "tile_idx_to_mn_limit differs"
+
+ tile = p2e.numel() // t2e.numel()
+ num_expanded = e2p.numel()
+ real = r_p2e >= 0
+ rows = torch.nonzero(real).squeeze(1)
+ run = r_t2e[rows // tile].long()
+ got = p2e[real].long()
+ assert bool(((got >= 0) & (got < num_expanded)).all()), (
+ "permuted_idx_to_expanded_idx holds an out-of-range value on a real row"
+ )
+ assert torch.equal(
+ torch.sort(run * num_expanded + got).values, run * num_expanded + r_p2e[real].long()
+ ), "an expert's run of permuted_idx_to_expanded_idx holds the wrong pairs"
+
+ flat, r_flat = e2p.reshape(-1).long(), r_e2p.reshape(-1).long()
+ not_local = r_flat < 0
+ assert bool((flat[not_local] == -1).all()), "a non-local pair is not -1"
+ p = flat[~not_local]
+ assert bool(((p >= 0) & (p < p2e.numel())).all()), "a local pair points outside the rows"
+ assert bool(real[p].all()), "a local pair points at a padding row"
+ assert torch.equal(p2e[p].long(), torch.arange(num_expanded, device=p.device)[~not_local]), (
+ "expanded_idx_to_permuted_idx is not the inverse of permuted_idx_to_expanded_idx"
+ )
+
+
+def _ids(dist, tokens, tile, seed, num_experts=E, top_k=K):
+ """Expert ids, `[tokens, top_k]` int32, every row naming distinct experts."""
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ if dist == "router":
+ scores = torch.rand(tokens, num_experts, device=DEV, generator=g)
+ return scores.topk(top_k, dim=-1).indices.to(torch.int32)
+ if dist == "tile_edges":
+ edge = (7, 300, 511)
+ scores = torch.rand(tokens, num_experts, device=DEV, generator=g)
+ scores[:, list(edge)] = -1.0
+ ids = scores.topk(top_k, dim=-1).indices.to(torch.int32)
+ for slot, (expert, rows) in enumerate(zip(edge, (tile - 1, tile, tile + 1))):
+ ids[:rows, slot] = expert
+ return ids
+ raise ValueError(dist)
+
+
+def _args(ids, tile, seed=0, scales_dtype=torch.float32):
+ """The CuteDSL call site's keyword arguments at the checkpoint's values."""
+ g = torch.Generator(device=DEV).manual_seed(seed + 1)
+ scales = torch.rand(ids.shape, device=DEV, generator=g).to(scales_dtype)
+ return dict(
+ token_selected_experts=ids,
+ token_final_scales=scales,
+ num_experts=E,
+ top_k=K,
+ local_expert_offset=0,
+ local_num_experts=E,
+ tile_tokens_dim=tile,
+ )
+
+
+def _check(args):
+ """Drive the op and gate it through the entry's own reference and `_compare`."""
+ out = op(**args)
+ ref = op.reference(**args)
+ _compare(out, ref)
+ return out, ref
+
+
+def _cell_id(cell):
+ s = cell.spec
+ return f"t{s['tile_tokens_dim']}-T{s['tokens']}-{s['dist']}"
+
+
+# ── the certified cells ───────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[_cell_id(c) for c in op.CELLS])
+def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name
+ certifying(op)
+ s = cell.spec
+ seed = zlib.crc32(cell.why.encode())
+ ids = _ids(s["dist"], s["tokens"], s["tile_tokens_dim"], seed)
+ args = _args(ids, s["tile_tokens_dim"], seed, s["scales_dtype"])
+ assert (s["num_experts"], s["top_k"]) == (args["num_experts"], args["top_k"])
+ assert (s["local_expert_offset"], s["local_num_experts"]) == (0, E)
+ ids_before = ids.clone()
+ scales_before = args["token_final_scales"].clone()
+ out, ref = _check(args)
+ assert torch.equal(ids, ids_before), "token_selected_experts was written"
+ assert torch.equal(args["token_final_scales"], scales_before), "token_final_scales was written"
+ if s["dist"] == "tile_edges":
+ tile = s["tile_tokens_dim"]
+ n = int(ref[5])
+ limits = out[1][:n].long()
+ starts = torch.arange(n, device=DEV) * tile
+ fill = (limits - starts).tolist()
+ experts = out[0][:n].tolist()
+ per_expert = {e: [] for e in (7, 300, 511)}
+ for e, f in zip(experts, fill):
+ if e in per_expert:
+ per_expert[e].append(f)
+ assert per_expert == {7: [tile - 1], 300: [tile], 511: [tile, 1]}, per_expert
+
+
+# ── CUDA graphs: captured row counts, on new ids ──────────────────────────
+
+#: The largest decode capture: 32 requests at MTP verify width 4, cluster kernel.
+_CAPTURED_ROWS = 128
+
+
+def _capture(args):
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ for _ in range(2):
+ op(**args)
+ torch.cuda.current_stream().wait_stream(side)
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ static_out = op(**args)
+ return graph, static_out
+
+
+def test_captured_rows_replay_on_new_ids(certifying): # noqa: F811
+ """The graph is replayed on two fresh id sets and gated against the reference,
+ and against an eager call on the same ids: counts and tile arrays bitwise. At
+ the target's tile, 128."""
+ certifying(op)
+ tile, rows = 128, _CAPTURED_ROWS
+ args = _args(_ids("router", rows, tile, seed=rows), tile, seed=rows)
+ graph, static_out = _capture(args)
+ for trial in (1, 2):
+ fresh = _ids("router", rows, tile, seed=1000 * trial + rows)
+ args["token_selected_experts"].copy_(fresh)
+ graph.replay()
+ eager = op(**dict(args, token_selected_experts=fresh.clone()))
+ ref = op.reference(**args)
+ _compare(static_out, ref)
+ _compare(eager, ref)
+ n = int(ref[5])
+ for i in (0, 1):
+ assert torch.equal(static_out[i][:n], eager[i][:n]), (trial, i)
+ for i in (4, 5):
+ assert torch.equal(static_out[i], eager[i]), (trial, i)
+ del graph, static_out
+ torch.cuda.synchronize()
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_unpermute_inplace.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_unpermute_inplace.py
new file mode 100644
index 000000000000..723fdf39ca5a
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_unpermute_inplace.py
@@ -0,0 +1,242 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the moe_unpermute_inplace entry of the Qwen3.8-Flash-Next target.
+
+The entry owns the contract. This file lays a cell's FC2 rows out the way the
+CuteDSL FC2 hands them over -- placed by moe_sort, every row the op must not read
+filled with NaN, the output pre-filled with NaN -- and holds the claims no cell
+carries: CUDA-graph replay at the largest captured row count, bitwise equal to
+eager.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.moe_unpermute_inplace import ( # noqa: E501
+ moe_unpermute_inplace as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "moe_unpermute_inplace requires a CUDA device"
+
+DEV = "cuda"
+E, K, H = 512, 10, 2560
+NAN = float("nan")
+
+#: The largest decode capture: 32 requests at MTP verify width 4.
+_CAPTURED_ROWS = 128
+
+
+def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
+ """Bitwise.
+
+ `reference` is the computation the op performs, in its order and with its
+ roundings, so no band is needed and none is given. The gate sees the
+ accumulation order: the same sum taken in reverse slot order, with each
+ product rounded before its add, as the float64 sum rounded once, or
+ accumulated in bf16, differs in many elements and is rejected. A failure in a
+ few isolated one-ulp elements points at the kernel's rounding (a lost
+ multiply-add contraction, say); a wrong row or scale moves whole rows.
+ """
+ assert out.dtype == ref.dtype, (out.dtype, ref.dtype)
+ assert out.shape == ref.shape, (tuple(out.shape), tuple(ref.shape))
+ bits = {2: torch.int16, 4: torch.int32}[out.element_size()]
+ differ = out.view(bits) != ref.view(bits)
+ n = int(differ.sum())
+ if n:
+ first = torch.nonzero(differ)[0].tolist()
+ raise AssertionError(
+ f"{n} of {out.numel()} elements differ bitwise from the slot-order fp32 "
+ f"sum; first at {first}: {out[tuple(first)].item()} vs "
+ f"{ref[tuple(first)].item()}"
+ )
+
+
+def _spec(tokens, tile_tokens_dim=128, values="normal"):
+ """The fields `_problem` reads; the same names as a cell's spec."""
+ return dict(tokens=tokens, tile_tokens_dim=tile_tokens_dim, values=values)
+
+
+def _ids(tokens: int, seed: int) -> torch.Tensor:
+ """Expert ids, `[tokens, K]` int32: each row's top-10 of uniform scores."""
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ scores = torch.rand(tokens, E, device=DEV, generator=g)
+ return scores.topk(K, dim=-1).indices.to(torch.int32)
+
+
+def _scales(tokens: int, seed: int) -> torch.Tensor:
+ """fp32 routing weights, a softmax over each row's ten slots."""
+ g = torch.Generator(device=DEV).manual_seed(seed + 1)
+ return torch.softmax(torch.randn(tokens, K, device=DEV, generator=g), dim=-1)
+
+
+def _rows(tokens: int, seed: int, values: str = "normal") -> torch.Tensor:
+ """The FC2 output row of each (token, slot) pair, `[tokens, K, H]` bf16."""
+ g = torch.Generator(device=DEV).manual_seed(seed + 2)
+ v = torch.randn(tokens, K, H, device=DEV, generator=g)
+ if values == "wide":
+ v = v * torch.exp2(torch.randint(-30, 31, v.shape, device=DEV, generator=g).float())
+ elif values != "normal":
+ raise ValueError(values)
+ return v.to(torch.bfloat16)
+
+
+def _sort(ids: torch.Tensor, scales: torch.Tensor, tile: int):
+ """moe_sort as the CuteDSL MoE calls it: (expanded_idx_to_permuted_idx, padded rows)."""
+ out = torch.ops.trtllm.moe_sort(
+ token_selected_experts=ids,
+ token_final_scales=scales,
+ num_experts=E,
+ top_k=K,
+ local_expert_offset=0,
+ local_num_experts=E,
+ tile_tokens_dim=tile,
+ )
+ return out[2], out[3].numel()
+
+
+def _place(rows: torch.Tensor, e2p: torch.Tensor, num_rows: int) -> torch.Tensor:
+ """The FC2 output as the op receives it: each live pair's row at its permuted
+ index, NaN in every row the op has no business reading."""
+ perm = torch.full((num_rows, H), NAN, dtype=rows.dtype, device=DEV)
+ flat = e2p.reshape(-1)
+ live = flat >= 0
+ perm[flat[live].long()] = rows.reshape(-1, H)[live]
+ return perm
+
+
+def _output(tokens: int) -> torch.Tensor:
+ return torch.full((tokens, H), NAN, dtype=torch.bfloat16, device=DEV)
+
+
+def _problem(spec: dict, seed: int):
+ """A cell's arguments, plus the per-pair rows and ids they were built from."""
+ tokens = spec["tokens"]
+ ids = _ids(tokens, seed)
+ scales = _scales(tokens, seed)
+ e2p, num_rows = _sort(ids, scales, spec["tile_tokens_dim"])
+ rows = _rows(tokens, seed, spec["values"])
+ args = dict(
+ permuted_input=_place(rows, e2p, num_rows),
+ output=_output(tokens),
+ expanded_idx_to_permuted_idx=e2p,
+ topk_scales=scales,
+ )
+ return args
+
+
+def _bits(t: torch.Tensor) -> torch.Tensor:
+ return t.contiguous().view(torch.uint8)
+
+
+def _same_bits(a: torch.Tensor, b: torch.Tensor) -> bool:
+ return a.shape == b.shape and a.dtype == b.dtype and torch.equal(_bits(a), _bits(b))
+
+
+def _cell_id(cell):
+ s = cell.spec
+ return f"t{s['tile_tokens_dim']}-T{s['tokens']}-{s['values']}"
+
+
+# ── the certified cells ───────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[_cell_id(c) for c in op.CELLS])
+def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name
+ certifying(op)
+ s = cell.spec
+ assert (s["num_experts"], s["top_k"], s["hidden_size"]) == (E, K, H)
+ assert (s["dtype"], s["scales_dtype"]) == (torch.bfloat16, torch.float32)
+ args = _problem(s, zlib.crc32(cell.why.encode()))
+ before = {
+ n: args[n].clone()
+ for n in ("permuted_input", "expanded_idx_to_permuted_idx", "topk_scales")
+ }
+ out = op(**args)
+ assert out is args["output"], "the op must return the tensor it wrote"
+ _compare(out, op.reference(**args))
+ for name, snapshot in before.items():
+ assert _same_bits(args[name], snapshot), f"{name} was written"
+
+
+# ── CUDA graphs: captured row counts, on new data ─────────────────────────
+
+
+def _capture(fn):
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ for _ in range(2):
+ fn()
+ torch.cuda.current_stream().wait_stream(side)
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ static = fn()
+ return graph, static
+
+
+def test_captured_rows_replay_on_new_data(certifying): # noqa: F811
+ """The sort, the FC2 rows' placement and this op are captured together, as the
+ target runs them; the graph is replayed on two fresh sets of ids, scales and
+ rows. The permuted buffer is not cleared between replays, so the rows the new
+ sort does not use hold the previous placement's rows. Each replay is gated
+ against the reference, and is bitwise equal to an eager call on the same data
+ through its own sort. At the target's tile, 128.
+ """
+ certifying(op)
+ tile = 128
+ stale_rows = 0
+ tokens = _CAPTURED_ROWS
+ ids = _ids(tokens, seed=tokens)
+ scales = _scales(tokens, seed=tokens)
+ rows = _rows(tokens, seed=tokens)
+ num_rows = _sort(ids, scales, tile)[1]
+ perm = torch.full((num_rows, H), NAN, dtype=torch.bfloat16, device=DEV)
+ out = _output(tokens)
+
+ def step(ids=ids, scales=scales, rows=rows, perm=perm, out=out, tile=tile):
+ e2p = _sort(ids, scales, tile)[0]
+ perm.index_copy_(0, e2p.reshape(-1).long(), rows.reshape(-1, H))
+ op(
+ permuted_input=perm,
+ output=out,
+ expanded_idx_to_permuted_idx=e2p,
+ topk_scales=scales,
+ )
+ return e2p
+
+ graph, e2p = _capture(step)
+ for trial in (1, 2):
+ seed = 1000 * trial + tokens
+ ids.copy_(_ids(tokens, seed))
+ scales.copy_(_scales(tokens, seed))
+ rows.copy_(_rows(tokens, seed))
+ out.fill_(NAN)
+ graph.replay()
+ args = dict(
+ permuted_input=perm,
+ output=out,
+ expanded_idx_to_permuted_idx=e2p,
+ topk_scales=scales,
+ )
+ _compare(out, op.reference(**args))
+ used = torch.zeros(num_rows, dtype=torch.bool, device=DEV)
+ used[e2p.reshape(-1).long()] = True
+ stale_rows += int((~used & ~perm.isnan().any(1)).sum())
+
+ e2p_eager, n = _sort(ids.clone(), scales.clone(), tile)
+ eager = op(
+ permuted_input=_place(rows.clone(), e2p_eager, n),
+ output=_output(tokens),
+ expanded_idx_to_permuted_idx=e2p_eager,
+ topk_scales=scales.clone(),
+ )
+ assert _same_bits(out, eager), trial
+ del graph, e2p
+ torch.cuda.synchronize()
+ assert stale_rows > 0, "no replay ran over a stale row: the stale-row claim is untested"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_renorm_moe_routing_op.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_renorm_moe_routing_op.py
new file mode 100644
index 000000000000..c98408673a3e
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_renorm_moe_routing_op.py
@@ -0,0 +1,115 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the qwen4 default target's renorm_moe_routing_op variant
+(models/qwen4_exp/ops/moe/renorm_moe_routing_op).
+
+The entry owns the contract. This file turns a cell into logits and holds the
+claim no cell carries: equivalence to HF's `norm_topk_prob` routing.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.renorm_moe_routing_op import ( # noqa: E501
+ renorm_moe_routing_op as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "renorm_moe_routing_op requires a CUDA device"
+
+
+def _compare(
+ out: tuple[torch.Tensor, torch.Tensor],
+ ref: tuple[torch.Tensor, torch.Tensor],
+) -> None:
+ """Ids equal and weights bit-exact.
+
+ The reference is not an approximation of the kernel: over 2040
+ configurations and 16384-token runs at 512 / top-10, not one id and not
+ one weight bit differed. A one-ulp `assert_close` runs first only so a
+ drift reports a readable diff. The nearest wrong answers sit 0.4476
+ away (un-renormalized softmax, `1/topk`), and a `torch.topk` reference
+ differs on a third of the rows at 512 bf16 experts.
+ """
+ ids, weights = out
+ ref_ids, ref_weights = ref
+ assert ids.dtype == ref_ids.dtype == torch.int32, (ids.dtype, ref_ids.dtype)
+ assert weights.dtype == ref_weights.dtype, (weights.dtype, ref_weights.dtype)
+ assert ids.shape == ref_ids.shape == weights.shape, (ids.shape, ref_ids.shape)
+ torch.testing.assert_close(ids, ref_ids, rtol=0, atol=0)
+ eps = torch.finfo(weights.dtype).eps
+ torch.testing.assert_close(weights, ref_weights, rtol=eps, atol=0.0)
+ mismatched = int((weights != ref_weights).sum())
+ assert mismatched == 0, f"{mismatched} of {weights.numel()} weights are not bit-exact"
+
+
+DEV = "cuda"
+
+
+def _logits(tokens, num_experts, logits_dtype, dist="randn", magnitude=1.0, seed=0):
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ if dist == "coarse":
+ x = torch.randint(0, 3, (tokens, num_experts), device=DEV, generator=g).float()
+ else:
+ x = torch.randn(tokens, num_experts, device=DEV, generator=g) * magnitude
+ return x.to(logits_dtype)
+
+
+def _check(logits, topk, output_dtype=None):
+ """Drive the op and gate it through the entry's own reference and `_compare`."""
+ ids, weights = op(logits, topk, output_dtype)
+ assert ids.is_contiguous() and weights.is_contiguous()
+ assert ids.device == weights.device == logits.device
+ _compare((ids, weights), op.reference(logits, topk, output_dtype))
+ return ids, weights
+
+
+# ── the certified cells ───────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name
+ certifying(op)
+ s = cell.spec
+ logits = _logits(
+ s["tokens"],
+ s["num_experts"],
+ s["logits_dtype"],
+ s["dist"],
+ s.get("magnitude", 1.0),
+ seed=zlib.crc32(cell.why.encode()),
+ )
+ ids, weights = _check(logits, s["topk"], s["output_dtype"])
+ assert weights.dtype == s["output_dtype"]
+ # renormalized: every row sums to 1
+ torch.testing.assert_close(
+ weights.float().sum(-1),
+ torch.ones(s["tokens"], device=DEV),
+ rtol=8e-3 if s["output_dtype"] is torch.bfloat16 else 1e-6,
+ atol=0.0,
+ )
+
+
+# ── 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."""
+ 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(
+ 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
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py
new file mode 100644
index 000000000000..cb2e79386da1
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py
@@ -0,0 +1,60 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the flashinfer_gemma_rmsnorm entry."""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.norm.flashinfer_gemma_rmsnorm import ( # noqa: E501
+ flashinfer_gemma_rmsnorm as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "flashinfer_gemma_rmsnorm requires a CUDA device"
+
+
+def _seed(text: str) -> int:
+ return zlib.crc32(text.encode())
+
+
+def _build(spec, seed):
+ """`(buffer, x, weight, eps)`: `x` is a view of `buffer` laid out as the spec says."""
+ g = torch.Generator(device="cuda").manual_seed(seed)
+ shape, dtype = spec["shape"], spec["dtype"]
+ width = shape[-1]
+
+ def randn(*s):
+ return torch.randn(*s, generator=g, device="cuda").to(dtype)
+
+ if "row_width" in spec:
+ rows, off = shape[0], spec["col_offset"]
+ buffer = randn(rows, spec["row_width"])
+ x = buffer[:, off : off + width]
+ else:
+ buffer = x = randn(*shape)
+ weight = torch.randn(width, generator=g, device="cuda").to(dtype)
+ return buffer, x, weight, spec["eps"]
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
+ certifying(op)
+ buffer, x, weight, eps = _build(cell.spec, seed=_seed(cell.why))
+ assert tuple(x.shape) == cell.spec["shape"] and x.stride(-1) == 1
+ if "row_width" in cell.spec and x.shape[0] > 1:
+ assert x.stride(0) == cell.spec["row_width"] and not x.is_contiguous()
+ elif "row_width" in cell.spec:
+ # One stream row: torch calls it contiguous, so flashinfer takes the compact kernel.
+ assert x.is_contiguous()
+ buffer_before = buffer.clone()
+
+ out = op(x, weight, eps)
+
+ assert out.shape == x.shape and out.dtype == x.dtype
+ assert out.data_ptr() != x.data_ptr(), "the op must return a new tensor"
+ assert torch.equal(buffer, buffer_before), "x (and the stream it slices) is read-only"
+ torch.testing.assert_close(out, op.reference(x, weight, eps))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_layer_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_layer_norm.py
new file mode 100644
index 000000000000..b64080199d2b
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_layer_norm.py
@@ -0,0 +1,63 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the layer_norm catalog entry.
+
+Weight and bias are drawn over this checkpoint's own 110 vision LayerNorm
+tensors' ranges: the affine is wide enough (weights to 17.5, biases to -25.6)
+that a wrong normalization is not masked by a near-unit scale. The default
+bf16 gate against the entry's fp64 reference is the band.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.norm.layer_norm import ( # noqa: E501
+ layer_norm as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401
+
+assert torch.cuda.is_available(), "layer_norm requires a CUDA device"
+
+DEVICE = torch.device("cuda")
+WIDTH = 1152
+EPS = 1e-6
+
+#: Over all of model.visual.*.norm{1,2} and model.visual.merger.norm (bf16).
+WEIGHT_LO, WEIGHT_HI = -1.210938, 17.5
+BIAS_LO, BIAS_HI = -25.625, 14.875
+
+
+def _gen(seed: int) -> torch.Generator:
+ return torch.Generator(device=DEVICE).manual_seed(seed)
+
+
+def _affine(gen, width=WIDTH, dtype=torch.bfloat16):
+ weight = torch.rand(width, generator=gen, device=DEVICE) * (WEIGHT_HI - WEIGHT_LO) + WEIGHT_LO
+ bias = torch.rand(width, generator=gen, device=DEVICE) * (BIAS_HI - BIAS_LO) + BIAS_LO
+ return weight.to(dtype), bias.to(dtype)
+
+
+def _build(spec, seed):
+ gen = _gen(seed)
+ shape, dtype = spec["shape"], spec["dtype"]
+ 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)]
+ 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 = [x, weight, bias]
+ assert all(torch.equal(a, b) for a, b in zip(before, after)), "inputs must not be mutated"
+ assert torch.equal(out, op(x, weight, bias, eps)), "two launches must agree bitwise"
+ torch.testing.assert_close(out, op.reference(x, weight, bias, eps))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py
new file mode 100644
index 000000000000..6db54e25f309
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py
@@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the rms_norm_gated_token_major entry."""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.norm.rms_norm_gated_token_major import ( # noqa: E501
+ rms_norm_gated_token_major as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "rms_norm_gated_token_major requires a CUDA device"
+
+
+#: The width of the concatenated GDN input projection `[q k v | z | a | b]` a cell names no
+#: `z_row_stride` for: all 48 heads, 10240 | 6144 | 48 | 48. A tp4/tep4 rank's is 4120:
+#: 2560 | 1536 | 12 | 12.
+_IN_PROJ_WIDTH = 16480
+
+
+def _z_offset(row_stride: int, heads: int, n: int) -> int:
+ """First column of `z` in a `[q k v | z | a | b]` row: `a` and `b` are one column a head."""
+ offset = row_stride - heads * (n + 2)
+ assert offset > 0, (row_stride, heads, n)
+ return offset
+
+
+def _seed(text: str) -> int:
+ return zlib.crc32(text.encode())
+
+
+def _build(spec, seed):
+ """`__call__`'s positional arguments for one spec."""
+ g = torch.Generator(device="cuda").manual_seed(seed)
+ t, heads, n, dtype = spec["tokens"], spec["heads"], spec["width"], spec["dtype"]
+
+ def randn(*shape):
+ return torch.randn(*shape, generator=g, device="cuda").to(dtype)
+
+ x = randn(t * heads, n)
+
+ layout = spec.get("z", "packed")
+ if layout == "in_proj":
+ width = spec.get("z_row_stride", _IN_PROJ_WIDTH)
+ offset = _z_offset(width, heads, n)
+ z = randn(t, width)[:, offset : offset + heads * n].view(t, heads, n)
+ else:
+ assert layout == "packed", layout
+ z = randn(t, heads, n)
+
+ weight = randn(n)
+ 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 = args[:3]
+ assert z.stride()[1:] == (cell.spec["width"], 1)
+ if cell.spec["tokens"] > 1:
+ # A one-token view reports the dense stride for its size-1 token dim (and is then
+ # contiguous); the kernel never reads past token 0 there.
+ row_stride = cell.spec.get("z_row_stride", _IN_PROJ_WIDTH)
+ assert z.stride(0) == row_stride and not z.is_contiguous()
+ before = [x.clone(), z.clone(), weight.clone()]
+
+ out = op(*args)
+
+ 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"
+ torch.testing.assert_close(out, op.reference(*args))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ngram_token_history.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ngram_token_history.py
new file mode 100644
index 000000000000..a5e97badb438
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ngram_token_history.py
@@ -0,0 +1,134 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the ngram_token_history entry.
+
+Integer bookkeeping, gated bitwise -- window and the whole published pool --
+against the entry's scalar reference, which builds each sequence's stream as
+a list and shares no structure with the masked per-shift select under test.
+Beyond the cells: chunked prefill against one-shot (the cross-step carry
+the forward relies on).
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ngram_token_history import ( # noqa: E501
+ CKPT_EOS_TOKEN_ID,
+ CKPT_NGRAM_SIZE,
+ CKPT_VOCAB_SIZE,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ngram_token_history import ( # noqa: E501
+ ngram_token_history as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "ngram_token_history requires a CUDA device"
+
+
+def _compare(
+ out: tuple[torch.Tensor, torch.Tensor], ref: tuple[torch.Tensor, torch.Tensor]
+) -> None:
+ """`(window, context_pool)` -- the returned window with the pool as the
+ call left it -- against `reference`'s pair, both bitwise: integer
+ bookkeeping, no tolerance. The whole pool is compared, which is the
+ lock on every slot the batch did not name."""
+ (window, pool), (ref_window, ref_pool) = out, ref
+ assert window.dtype == ref_window.dtype == torch.int64, (window.dtype, ref_window.dtype)
+ assert window.shape == ref_window.shape, (window.shape, ref_window.shape)
+ assert window.is_contiguous()
+ assert torch.equal(window, ref_window), "window disagrees with the reference"
+ assert torch.equal(pool, ref_pool), "published context pool disagrees with the reference"
+
+
+DEV = "cuda"
+EOS = CKPT_EOS_TOKEN_ID
+NGRAM = CKPT_NGRAM_SIZE
+CONTEXT_LEN = NGRAM - 1
+
+
+def _qsl(lengths, dtype, qsl_from="literal"):
+ if qsl_from == "seq_lens":
+ # The call site under MTP verify: an int32 prefix of the metadata's
+ # seq_lens buffer, scanned into an int64 vector with a leading zero.
+ buf = torch.full((len(lengths) + 9,), 99, dtype=torch.int32, device=DEV)
+ buf[: len(lengths)] = torch.tensor(lengths, dtype=torch.int32, device=DEV)
+ seq_lens = buf[: len(lengths)]
+ out = torch.zeros(len(lengths) + 1, dtype=torch.int64, device=DEV)
+ torch.cumsum(seq_lens, dim=0, dtype=torch.int64, out=out[1:])
+ return out
+ starts = [0]
+ for length in lengths:
+ starts.append(starts[-1] + length)
+ return torch.tensor(starts, dtype=dtype, device=DEV)
+
+
+def _build(spec, seed):
+ """The op's positional arguments for one cell."""
+ g = torch.Generator().manual_seed(seed)
+ 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
+ pool = torch.randint(
+ 0, CKPT_VOCAB_SIZE, (spec["pool_slots"], CONTEXT_LEN), generator=g, dtype=torch.long
+ ).to(DEV)
+ index_dtype = spec.get("index_dtype", torch.int32)
+ 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,
+ )
+
+
+def _expect(args):
+ """The reference, driven from a snapshot of the pool before the call."""
+ return op.reference(args[0], args[1].clone(), *args[2:])
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
+ certifying(op)
+ args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
+ read_only = [a.clone() for a in (args[0], args[2], args[3], args[4])]
+ pool_before = args[1].clone()
+ expected = _expect(args)
+ window = op(*args)
+ _compare((window, args[1]), expected)
+ for a, b in zip((args[0], args[2], args[3], args[4]), read_only):
+ assert torch.equal(a, b), "only context_pool may be written"
+ # The whole-pool comparison locks the unnamed slots; this keeps it from
+ # being vacuous -- every cell has a sequence that publishes new ids.
+ assert not torch.equal(args[1], pool_before), "no slot moved"
+
+
+# ----------------------------------------------------------- multi-step
+
+
+def test_chunked_prefill_matches_one_shot() -> None:
+ """96 tokens split at 41 over two calls: window and pool bitwise one-shot."""
+ g = torch.Generator().manual_seed(5001)
+ ids = torch.randint(0, CKPT_VOCAB_SIZE, (96,), generator=g).to(DEV)
+ start = torch.randint(0, CKPT_VOCAB_SIZE, (3, CONTEXT_LEN), generator=g).to(DEV)
+ slot = torch.tensor([1], dtype=torch.int32, device=DEV)
+
+ def call(pool, chunk, carry):
+ flags = torch.tensor([carry], device=DEV)
+ qsl = torch.tensor([0, chunk.shape[0]], dtype=torch.int32, device=DEV)
+ return op(chunk, pool, slot, flags, qsl, EOS, NGRAM)
+
+ one_shot_pool = start.clone()
+ one_shot = call(one_shot_pool, ids, False)
+ pool = start.clone()
+ split = torch.cat([call(pool, ids[:41], False), call(pool, ids[41:], True)])
+ assert torch.equal(split, one_shot)
+ assert torch.equal(pool, one_shot_pool)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_gate_conv_inject_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_gate_conv_inject_norm.py
new file mode 100644
index 000000000000..4581ad5e52db
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_gate_conv_inject_norm.py
@@ -0,0 +1,199 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the default target's ple_gate_conv_inject_norm variant.
+
+Cells are gated against the float64 reference that rounds where the op rounds:
+the stream (`hc_out`) and the normed stream in units of their row's scale, the
+published conv state per element (see `_compare`); a verify cell's staged conv
+input is checked bitwise against the state it published. Beyond the cells: the
+staged conv input is what the conv consumed, and `out` writes only its rows of a
+larger buffer.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ple_gate_conv_inject_norm import ( # noqa: E501
+ CKPT_CLAMP_MIN,
+ CKPT_DILATION,
+ CKPT_HC_COUNT,
+ CKPT_HIDDEN,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ple_gate_conv_inject_norm import ( # noqa: E501
+ ple_gate_conv_inject_norm as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import ( # noqa: E402,F401 — path above; fixture by name
+ assert_within_ulp,
+ certifying,
+ validating,
+)
+
+assert torch.cuda.is_available(), "ple_gate_conv_inject_norm requires a CUDA device"
+
+DEV = "cuda"
+EPS = 1e-6
+PAD = -1
+GAMMA_RMS = 0.477
+
+
+def _ulp_ratio(out: torch.Tensor, ref: torch.Tensor) -> torch.Tensor:
+ """Per-element distance from `ref` in units of one ulp of `|ref|`."""
+ eps = torch.finfo(ref.dtype).eps
+ tiny = torch.finfo(ref.dtype).tiny
+ ref64 = ref.double()
+ return (out.double() - ref64).abs() / (eps * ref64.abs() + tiny)
+
+
+def _stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]:
+ ratio = _ulp_ratio(out, ref)
+ return ratio.max().item(), (out != ref).double().mean().item()
+
+
+#: The band. Every rounding point of the composition is kept, so the op and the
+#: float64 reference differ only where an fp32 reduction (three sums of squares
+#: and the gate's dot per stream) or fp32 `exp` lands a value across a bf16
+#: midpoint -- measured on 2e-5..2.3e-4 of the elements. A flipped upstream
+#: element then propagates, and through `hc + ple` it meets cancellation
+#: (`hc ~ -ple`), so a per-element ulp is unbounded there: 73 ulp of a
+#: 0.003-sized element in a row whose largest is ~10. The stream and its norm
+#: are therefore gated in units of their row's scale (`assert_within_ulp`):
+#: 2 ulp of the row maximum per element and 0.1 ulp relative RMS. The
+#: composition this replaces sits at 0.87 / 0.33 of a row-max ulp against the
+#: same reference (measured: normed / hc_out at 128 rows), the op at 0.72 /
+#: 0.38. The published conv state is one rounding of the conv-input norm,
+#: whose only perturbation is a one-ulp flip of a gated element: 2 ulp per
+#: element (measured worst 1.78, the old composition 1.56).
+ROW_ULP, RMS_ULP, STATE_ULP = 2.0, 0.1, 2.0
+#: The per-element row-max gate is a tail statistic, and a context-shaped cell
+#: draws ~100x more elements than a 128-row one. Against the same float64
+#: reference over three seeds at 16384 rows, the composition this replaces
+#: reaches 2.91 row-max ulp (normed) and the op 2.94; at 8192 rows 1.89 / 1.89.
+#: Cells above 1024 rows are therefore gated at 4; the RMS gate is unchanged.
+#: The default target's per-rank pools (1874 / 850 slots) sit inside the same gates; a
+#: reference reading each sequence's neighbouring slot falls far outside the RMS gate
+#: (measured).
+CONTEXT_ROWS, CONTEXT_ROW_ULP = 1024, 4.0
+
+
+def _compare(out, ref) -> None:
+ row_ulp = CONTEXT_ROW_ULP if out[0].shape[0] > CONTEXT_ROWS else ROW_ULP
+ for name, got, want in zip(("hc_out", "normed", "conv_states"), out, ref):
+ assert got.dtype == want.dtype and got.shape == want.shape, name
+ if not got.numel():
+ continue
+ if name == "conv_states":
+ worst = _ulp_ratio(got, want).max().item()
+ assert worst <= STATE_ULP, f"{name}: {worst:.4g} ulp"
+ else:
+ assert_within_ulp(got, want, row_ulp, RMS_ULP)
+
+
+def _build(spec, seed):
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ gc = torch.Generator().manual_seed(seed)
+ lengths = spec["lengths"]
+ c, d, dtype = CKPT_HC_COUNT, CKPT_HIDDEN, torch.bfloat16
+ w = c * d
+ rows = sum(lengths)
+ num_seqs = len(lengths)
+ pool_slots = spec.get("pool_slots", num_seqs + 2)
+
+ def rand(*shape, scale=1.0):
+ return (torch.randn(*shape, generator=g, device=DEV) * scale).to(dtype)
+
+ # Streams at different scales, as the trained residual's are.
+ scales = torch.tensor([0.5 * 2.0 ** (i % 4) for i in range(c)], device=DEV)
+ hc = (rand(rows, w, scale=3.0).float() * scales.repeat_interleave(d)).to(dtype)
+ key = rand(rows, w)
+ value = rand(rows, d, scale=0.5)
+ norms = [rand(w, scale=GAMMA_RMS) for _ in range(4)]
+ conv_weight = rand(w, 1, 4, scale=0.3)
+ conv_states = rand(pool_slots, w, 3 * CKPT_DILATION)
+ order = (torch.randperm(pool_slots - 2, generator=gc) + 1).tolist()
+ slots = ([0, pool_slots - 1] + order)[:num_seqs] if pool_slots >= 64 else list(range(num_seqs))
+ slots = torch.tensor(slots, dtype=torch.int32, device=DEV)
+ starts = [0]
+ for length in lengths:
+ starts.append(starts[-1] + length)
+ qsl = torch.tensor(starts, dtype=torch.int64, device=DEV)
+ max_rows = spec.get("max_rows", max(lengths))
+ carries = spec.get("carries")
+ carries = None if carries is None else torch.tensor(carries, dtype=torch.bool, device=DEV)
+ return [
+ hc, key, value, norms[0], norms[1], norms[2], conv_weight, conv_states, slots, qsl,
+ norms[3], EPS, c, CKPT_CLAMP_MIN, CKPT_DILATION, PAD, max_rows, carries,
+ ] # fmt: skip
+
+
+def _run(args, conv_input=None):
+ """`(hc_out, normed, conv_states)` from the op, on a copy of the pool."""
+ args = list(args)
+ args[7] = args[7].clone()
+ hc_out, normed = op(*args, conv_input=conv_input)
+ return hc_out, normed, args[7]
+
+
+def _assert_staged_conv_input(xc, conv_states, slots, lengths) -> None:
+ """With `L <= state_len` rows in every sequence, the published state's last `L`
+ columns are those rows' conv inputs: the staged buffer must hold exactly them."""
+ rows_per_seq = lengths[0]
+ assert set(lengths) == {rows_per_seq} and rows_per_seq <= conv_states.shape[-1]
+ tail = conv_states[slots.long()][:, :, -rows_per_seq:]
+ assert torch.equal(torch.reshape(torch.transpose(tail, 1, 2), xc.shape), xc), (
+ "the staged conv input is not what the conv consumed"
+ )
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
+ certifying(op)
+ args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
+ hc_before = args[0].clone()
+ ref = op.reference(*args)
+ xc = torch.empty_like(args[0]) if cell.spec.get("conv_input") else None
+ out = _run(args, conv_input=xc)
+ assert torch.equal(args[0], hc_before), "hc is read-only"
+ assert out[0].is_contiguous() and out[1].is_contiguous()
+ _compare(out, ref)
+ if xc is not None:
+ _assert_staged_conv_input(xc, out[2], args[8], cell.spec["lengths"])
+ for name, got, want in zip(("hc_out", "normed", "conv_states"), out, ref):
+ worst, frac = _stats(got, want)
+ print(
+ f"\n {cell.why[:40]:<40} {name:<12} worst {worst:.3g} ulp, differ {frac:.2e}", end=""
+ )
+
+
+def test_conv_input_is_what_the_conv_consumed() -> None:
+ """`conv_input` receives the conv input of every row, and changes nothing else: with
+ four rows a sequence, the published state's last four columns are those rows."""
+ args = _build(dict(lengths=(4,) * 16, pool_slots=461), seed=11)
+ plain = _run(args)
+ xc = torch.empty_like(args[0])
+ args[7] = args[7].clone()
+ hc_out, normed = op(*args, conv_input=xc)
+ assert torch.equal(hc_out, plain[0]) and torch.equal(normed, plain[1])
+ assert torch.equal(args[7], plain[2])
+ tail = args[7][args[8].long()][:, :, -4:]
+ assert torch.equal(torch.reshape(torch.transpose(tail, 1, 2), xc.shape), xc)
+
+
+def test_out_writes_the_rows_of_larger_buffers() -> None:
+ """`out` takes the leading rows of two padded `[T + 8, C*D]` buffers: the same bits as
+ the fresh outputs, the padding rows untouched -- the prefill target's in-place write."""
+ args = _build(dict(lengths=(3001, 5000) + (1,) * 64, carries=(False, True) + (True,) * 64,
+ pool_slots=461), seed=12) # fmt: skip
+ plain = _run(args)
+ rows = args[0].shape[0]
+ bufs = [torch.full((rows + 8, args[0].shape[1]), 7.0, dtype=args[0].dtype, device=DEV)
+ for _ in range(2)] # fmt: skip
+ args[7] = args[7].clone()
+ hc_out, normed = op(*args, out=(bufs[0][:rows], bufs[1][:rows]))
+ assert hc_out.data_ptr() == bufs[0].data_ptr() and normed.data_ptr() == bufs[1].data_ptr()
+ assert torch.equal(bufs[0][:rows], plain[0]) and torch.equal(bufs[1][:rows], plain[1])
+ assert torch.equal(args[7], plain[2])
+ assert bool((bufs[0][rows:] == 7.0).all()) and bool((bufs[1][rows:] == 7.0).all())
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_ngram_embed.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_ngram_embed.py
new file mode 100644
index 000000000000..7eb14dc43f01
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_ngram_embed.py
@@ -0,0 +1,181 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the ple_ngram_embed catalog entry.
+
+Integer bookkeeping plus a one-rounding dequantize, so gated bitwise -- the
+embedding and the whole published pool -- against the entry's reference, which
+is the scalar references of `ngram_token_history` and `ngram_hash` and a
+float64 gather-and-scale.
+"""
+
+import zlib
+from types import SimpleNamespace
+
+import pytest
+import torch
+from cuda.bindings import runtime as cudart
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.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.ops.ple.ngram_token_history import ( # noqa: E501
+ ngram_token_history,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ple_ngram_embed import ( # noqa: E501
+ CKPT_EOS_TOKEN_ID,
+ CKPT_HEAD_DIM,
+ CKPT_SCALE,
+ CKPT_TABLE_ROWS,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ple_ngram_embed import ( # noqa: E501
+ ple_ngram_embed as op,
+)
+from tensorrt_llm._utils import CUASSERT
+
+__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 _host_mapped(src: torch.Tensor) -> torch.Tensor:
+ """A device view of `src`'s bytes in page-locked host memory.
+
+ The stock engine's placement of the n-gram table under its PLE host offload
+ (`Qwen4ExpPinnedHostEmbedding` in `tensorrt_llm/_torch/modules/qwen4_exp/ple.py`):
+ a pinned CPU allocation reached through `cudaHostGetDevicePointer`, so the
+ kernel gathers rows over the host link and no byte of the table is in device
+ memory -- which the pointer's own attributes are asked to confirm.
+ """
+ rows, cols = src.shape
+ host = torch.empty(rows * cols * src.element_size(), dtype=torch.uint8, pin_memory=True)
+ host.copy_(src.contiguous().view(torch.uint8).reshape(-1))
+ _HOSTS.append(host)
+ (pointer,) = CUASSERT(cudart.cudaHostGetDevicePointer(host.data_ptr(), 0))
+ # No public torch API views a pointer the caching allocator did not hand out;
+ # `|u1` because the array interface has no float8 typestr.
+ interface = SimpleNamespace(
+ __cuda_array_interface__={
+ "shape": (rows, cols * src.element_size()),
+ "typestr": "|u1",
+ "data": (int(pointer), False),
+ "version": 3,
+ "strides": None,
+ }
+ )
+ table = torch.as_tensor(interface, device=DEV).view(src.dtype)
+ assert table.data_ptr() == int(pointer), "torch.as_tensor copied the host buffer"
+ (attributes,) = CUASSERT(cudart.cudaPointerGetAttributes(table.data_ptr()))
+ assert attributes.type == cudart.cudaMemoryType.cudaMemoryTypeHost, attributes.type
+ assert table.shape == src.shape and table.dtype == src.dtype
+ return table
+
+
+def _scattered(n: int, pool_slots: int, g) -> list:
+ """`n` distinct slots over the whole pool, both ends included; a lone
+ sequence takes the last slot, the farthest row of the pool."""
+ middle = (torch.randperm(pool_slots - 2, generator=g) + 1).tolist()
+ slots = ([pool_slots - 1, 0] + middle)[:n]
+ assert len(set(slots)) == n and max(slots) < pool_slots
+ return slots
+
+
+def _geometry(spec, g):
+ """`(multipliers, vocab, offsets)` for the cell: the checkpoint's own over
+ the full-height table, the checkpoint's multipliers over small stand-in head
+ vocabularies otherwise."""
+ if spec.get("table") == "ckpt":
+ return CKPT_MULTIPLIERS, CKPT_HEAD_VOCAB_SIZES, CKPT_HEAD_OFFSETS
+ mult = CKPT_MULTIPLIERS
+ vocab = tuple(
+ 1000 + 7 * h + int(x)
+ for h, x in enumerate(torch.randint(0, 50, (len(CKPT_HEAD_VOCAB_SIZES),), generator=g))
+ )
+ offsets, acc = [], 0
+ for v in vocab:
+ offsets.append(acc)
+ acc += v
+ return mult, vocab, tuple(offsets)
+
+
+def _table(kind: str, rows: int, touched: torch.Tensor, g) -> torch.Tensor:
+ if kind == "ckpt":
+ table = torch.empty(CKPT_TABLE_ROWS, CKPT_HEAD_DIM, dtype=FP8, device=DEV)
+ table[touched] = torch.randn(touched.numel(), CKPT_HEAD_DIM, generator=g).to(FP8).to(DEV)
+ return table
+ src = torch.randn(rows, CKPT_HEAD_DIM, generator=g).to(FP8)
+ if kind == "host":
+ return _host_mapped(src)
+ return src.to(DEV)
+
+
+def _with_eos(x: torch.Tensor, density: float, g) -> torch.Tensor:
+ if density > 0.0:
+ x = torch.where(torch.rand(x.shape, generator=g) < density, torch.full_like(x, EOS), x)
+ return x
+
+
+def _build(spec, seed):
+ g = torch.Generator().manual_seed(seed)
+ lengths = spec["lengths"]
+ n = len(CKPT_MULTIPLIERS)
+ num_seqs = len(lengths)
+ pool_slots = spec.get("pool_slots", num_seqs + 2)
+ density = spec.get("eos_density", 0.0)
+ ids = _with_eos(torch.randint(0, VOCAB, (sum(lengths),), generator=g), density, g)
+ ids = ids.to(torch.int32).to(DEV)
+ pool = _with_eos(torch.randint(0, VOCAB, (pool_slots, n - 1), generator=g), density, g).to(DEV)
+ slots = _scattered(num_seqs, pool_slots, g) if pool_slots >= 64 else list(range(num_seqs))
+ slots = torch.tensor(slots, dtype=torch.int32, device=DEV)
+ carries = spec.get("carries")
+ carries = None if carries is None else torch.tensor(carries, dtype=torch.bool, device=DEV)
+ starts = [0]
+ for length in lengths:
+ starts.append(starts[-1] + length)
+ qsl = torch.tensor(starts, dtype=torch.int64, device=DEV)
+ mult, vocab, offsets = _geometry(spec, g)
+ mult = torch.tensor(mult, dtype=torch.int64, device=DEV)
+ vocab = torch.tensor(vocab, dtype=torch.int64, device=DEV)
+ offsets = torch.tensor(offsets, dtype=torch.int64, device=DEV)
+ # The rows this batch touches, from the two scalar references, so that the
+ # full-height table needs bytes only where they are read.
+ seeded = torch.ones(num_seqs, dtype=torch.bool, device=DEV) if carries is None else carries
+ window, _ = ngram_token_history.reference(ids, pool, slots, seeded, qsl, EOS, n)
+ touched = ngram_hash.reference(window, mult, vocab, offsets, EOS, 1).reshape(-1)
+ rows = int(offsets[-1] + vocab[-1])
+ table = _table(spec.get("table", "small"), rows, touched, g)
+ scale = torch.tensor([CKPT_SCALE], dtype=torch.bfloat16, device=DEV).reshape(1, 1, 1)
+ max_rows = spec.get("max_rows", max(lengths))
+ return [ids, pool, slots, carries, qsl, table, scale, mult, vocab, offsets, EOS, max_rows]
+
+
+def _compare(out, ref) -> None:
+ (emb, pool), (ref_emb, ref_pool) = out, ref
+ assert emb.dtype == ref_emb.dtype and emb.shape == ref_emb.shape
+ assert emb.is_contiguous()
+ assert torch.equal(emb, ref_emb), (
+ f"embedding differs at {(emb != ref_emb).any(-1).nonzero().flatten()[:8].tolist()}"
+ )
+ assert torch.equal(pool, ref_pool), "published pool differs"
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
+ certifying(op)
+ args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
+ ref = op.reference(*args)
+ pool_before = args[1].clone()
+ out = op(*args)
+ _compare((out, args[1]), ref)
+ if cell.spec.get("table") == "host":
+ assert args[1].data_ptr() != 0 and not torch.equal(args[1], pool_before)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/quantization/test_modeling_v2_bf16_to_fp32.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/quantization/test_modeling_v2_bf16_to_fp32.py
new file mode 100644
index 000000000000..37cca3ab7333
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/quantization/test_modeling_v2_bf16_to_fp32.py
@@ -0,0 +1,44 @@
+# 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()`."""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.quantization.bf16_to_fp32 import ( # noqa: E501
+ bf16_to_fp32 as op,
+)
+
+__extra_import_path__ = ["../../.."]
+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"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/quantization/test_modeling_v2_fp8_quantize_1x128.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/quantization/test_modeling_v2_fp8_quantize_1x128.py
new file mode 100644
index 000000000000..0a502ed944fd
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/quantization/test_modeling_v2_fp8_quantize_1x128.py
@@ -0,0 +1,130 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the fp8_quantize_1x128 catalog entry.
+
+The entry owns the contract. This file turns a cell into an input and holds
+every code disagreement of a cell to a rounding tie.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.quantization.fp8_quantize_1x128 import ( # noqa: E501
+ SCALE_ULP_GATE,
+ block_scales,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.quantization.fp8_quantize_1x128 import ( # noqa: E501
+ fp8_quantize_1x128 as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "fp8_quantize_1x128 requires a CUDA device"
+
+
+def _ordinal(codes: torch.Tensor) -> torch.Tensor:
+ """Sign-magnitude ordinal of an e4m3 byte: adjacent codes differ by 1."""
+ b = codes.view(torch.uint8).int()
+ return torch.where(b >= 128, -(b & 0x7F), b & 0x7F)
+
+
+def _compare(
+ out: tuple[torch.Tensor, torch.Tensor],
+ ref: tuple[torch.Tensor, torch.Tensor],
+) -> None:
+ """Scales within 4 fp32 ulp; codes at most one e4m3 step apart, >=99.5% bit-exact.
+
+ This build compiles the kernel's two fp32 divisions approximately
+ (`--use_fast_math`), each within CUDA's documented 2 ulp, so neither
+ output is reproducible to the last bit with `use_ue8m0=False`. Measured
+ over the certified column: scales within 2 ulp of both `amax / 448` and
+ `1 / (448 / amax)` (~80% bit-exact at T=1024); codes at worst
+ 99.6875% bit-exact at H 2560 (99.8566% elsewhere), every miss exactly
+ one code away with the exact product within 1.122e-07 relative of the
+ midpoint between the two -- a rounding tie the scale's last ulp decides.
+ `use_ue8m0=True` is bit-exact (the power-of-two scale makes the
+ multiply exact), and passes this gate trivially.
+
+ Wrong recipes at T=256, H=2560 (bytes agreeing / worst relative scale
+ error): power-of-two 0.043 / 0.99, per-tensor 0.0002 / 2.0, per-row
+ 0.083 / 1.6, blocks down the token axis 0.19 / 1.0, blocks offset by 64
+ 0.48 / 0.93, e4m3 max as 240 0 / 0.87, signed max 0.58 / 0.58; the
+ correct one 0.9993 / 1.4e-07.
+ """
+ data, scales = out
+ ref_data, ref_scales = ref
+ assert data.dtype == ref_data.dtype == torch.float8_e4m3fn, (data.dtype, ref_data.dtype)
+ assert scales.dtype == ref_scales.dtype == torch.float32, (scales.dtype, ref_scales.dtype)
+ assert data.shape == ref_data.shape, (data.shape, ref_data.shape)
+ assert scales.shape == ref_scales.shape, (scales.shape, ref_scales.shape)
+ if scales.numel():
+ ulp = (
+ scales.contiguous().view(torch.int32).long()
+ - ref_scales.contiguous().view(torch.int32).long()
+ )
+ worst = int(ulp.abs().max())
+ assert worst <= SCALE_ULP_GATE, f"scale {worst} fp32 ulp off > {SCALE_ULP_GATE}"
+ if data.numel():
+ step = (_ordinal(data) - _ordinal(ref_data)).abs()
+ assert int(step.max()) <= 1, f"a code is {int(step.max())} e4m3 steps off"
+ exact = 1.0 - (step != 0).float().mean().item()
+ assert exact >= 0.995, f"only {100 * exact:.4f}% of codes bit-exact"
+
+
+DEV = "cuda"
+# A code can only disagree when the exact product sits within the first
+# division's 2 ulp of an e4m3 rounding midpoint.
+TIE_REL_GATE = 2.0 * 2.0**-23
+
+
+def _expand(v: torch.Tensor, n: int) -> torch.Tensor:
+ """`[T, nb]` block values -> `[T, n]` per-column values."""
+ return v.repeat_interleave(128, dim=1)[:, :n].contiguous()
+
+
+def _input(tokens: int, hidden: int, seed: int, magnitude: float = 3.0) -> torch.Tensor:
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ return (torch.randn(tokens, hidden, device=DEV, generator=g) * magnitude).to(torch.bfloat16)
+
+
+def _assert_misses_are_ties(x: torch.Tensor, data: torch.Tensor, ref_data: torch.Tensor) -> float:
+ """Every code disagreement is a rounding tie: the exact product sits within
+ the scale's 2-ulp band of the midpoint between the two codes."""
+ qs, _ = block_scales(x, False)
+ prod = x.float() * _expand(qs, x.shape[1])
+ diff = data.float() != ref_data.float() # -0 and +0 are the same code here
+ if not bool(diff.any()):
+ return 0.0
+ mid = ((data.float() + ref_data.float()) / 2).double()
+ rel = ((prod.double()[diff] - mid[diff]).abs() / mid[diff].abs()).max().item()
+ assert rel <= TIE_REL_GATE, f"a miss {rel:.3e} from its midpoint is not a tie"
+ return rel
+
+
+# ── the certified cells ───────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name
+ certifying(op)
+ spec = cell.spec
+ x = _input(spec["tokens"], spec["hidden"], seed=zlib.crc32(cell.why.encode()))
+ data, scales = op(x, spec["use_ue8m0"])
+ nb = (spec["hidden"] + 127) // 128
+ assert data.shape == x.shape and data.is_contiguous() and data.device == x.device
+ assert scales.shape == (nb, spec["tokens"]) and scales.is_contiguous()
+ ref = op.reference(x, spec["use_ue8m0"])
+ _compare((data, scales), ref)
+ rel = _assert_misses_are_ties(x, data, ref[0]) if not spec["use_ue8m0"] else 0.0
+ exact = (data.view(torch.uint8) == ref[0].view(torch.uint8)).float().mean().item()
+ sms = torch.cuda.get_device_properties(x.device).multi_processor_count
+ print(
+ f" {cell.why[:44]}: {100 * exact:.4f}% bit-exact, worst tie {rel:.3e}, "
+ f"{spec['tokens'] * nb} groups over a {sms * 8 * 8 * 2}-group grid sweep"
+ )
+
+
+# ── what a cell cannot carry ──────────────────────────────────────────────
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_triton_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_triton_update.py
new file mode 100644
index 000000000000..a289411fbafb
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_triton_update.py
@@ -0,0 +1,130 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the causal_conv1d_triton_update entry."""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.causal_conv1d_triton_update import ( # noqa: E501
+ CKPT_X_STRIDE,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.causal_conv1d_triton_update import ( # noqa: E501
+ causal_conv1d_triton_update as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "causal_conv1d_triton_update requires a CUDA device"
+
+
+def _compare(out: tuple, ref: tuple) -> None:
+ """`(out, conv_state, intermediate_conv_window)` triples: the returned
+ tensor with the two buffers as the call left them, against
+ `reference`'s triple.
+
+ The result at torch's default band over the rows the op writes (NaN in
+ the reference marks the rest): against the operand-dtype-product
+ reference bf16/fp16 were bit-exact at dim 512 and at dim 10240 off by
+ one ulp in at most 9 of 2.6M elements (the fp32 `exp`); fp32 is off by
+ 1.4e-6. At a tep4 rank's 2560 channels the same one-ulp differences
+ touch a handful of elements. The pool and the snapshots are copies, gated
+ bitwise over the whole buffer -- the lock on every unnamed slot and
+ snapshot row.
+ """
+ (act, state, inter), (ref_act, ref_state, ref_inter) = out, ref
+ assert act.dtype == ref_act.dtype and act.shape == ref_act.shape
+ written = ~torch.isnan(ref_act)
+ torch.testing.assert_close(act[written], ref_act[written])
+ assert torch.equal(state, ref_state), "the advanced state must be bitwise exact"
+ assert (inter is None) == (ref_inter is None)
+ if inter is not None:
+ assert torch.equal(inter, ref_inter), "the snapshots must be bitwise exact"
+
+
+PAD_SLOT_ID = -1
+
+
+def _build(spec, seed):
+ """`__call__`'s positional arguments for one spec, freshly allocated, plus the
+ buffers its views are cut from.
+
+ x is the q|k|v columns of an `[n * p, x_row_stride]` projection output
+ (z | a | b past them), reshaped `[n, p, C]` and seen as `[n, C, p]`; the
+ pool is the last layer's view of a `[slots, pool_layers, C, 3]` page, as the
+ V2 manager coalesces the GDN layers' conv states inside each slot; the
+ snapshot buffer is the first layer's view of a `[snapshot_layers, rows,
+ steps, C, 3]` buffer, as the GDN speculative state holds every layer's.
+ """
+ g = torch.Generator(device="cuda").manual_seed(seed)
+ dim, batch, seqlen, width = spec["dim"], spec["batch"], spec["seqlen"], 4
+ wide = spec.get("x_row_stride", CKPT_X_STRIDE)
+ layers = spec.get("pool_layers", 1)
+
+ def randn(*shape):
+ return torch.randn(*shape, generator=g, device="cuda").to(torch.bfloat16)
+
+ buf = randn(batch * seqlen, wide)
+ x = buf[:, :dim].reshape(batch, seqlen, dim).transpose(1, 2)
+ assert x.stride() == (seqlen * wide, 1, wide)
+ pool_buf = randn(spec["pool_slots"], layers, dim, width - 1)
+ conv_state = pool_buf[:, layers - 1]
+ assert conv_state.stride() == (layers * dim * (width - 1), width - 1, 1)
+ weight = randn(dim, width)
+ indices = torch.tensor(spec["slots"], dtype=torch.int32, device="cuda")
+ steps = spec.get("snapshot_steps", seqlen)
+ window_buf = torch.zeros(
+ spec.get("snapshot_layers", 1),
+ spec["snapshots"],
+ steps,
+ dim,
+ width - 1,
+ dtype=torch.bfloat16,
+ device="cuda",
+ )
+ rows_idx = torch.arange(batch, dtype=torch.int32, device="cuda")
+ args = [
+ x,
+ conv_state,
+ weight,
+ None,
+ "silu",
+ indices,
+ PAD_SLOT_ID,
+ None,
+ window_buf[0],
+ rows_idx,
+ spec.get("launch_dependent_kernels", False),
+ ]
+ return args, dict(buf=buf, pool_buf=pool_buf, window_buf=window_buf)
+
+
+def _run(args, extras):
+ """Drive the op on copies of the pool page and the snapshot buffer, through
+ the same views; return the triple compare takes."""
+ got = list(args)
+ pool_buf, window_buf = extras["pool_buf"].clone(), extras["window_buf"].clone()
+ got[1], got[8] = pool_buf[:, pool_buf.shape[1] - 1], window_buf[0]
+ buf = extras["buf"].clone()
+ returned = op(*got)
+ assert torch.equal(extras["buf"], buf), "x (and the columns past it) are read-only"
+ assert torch.equal(pool_buf[:, :-1], extras["pool_buf"][:, :-1]), "another layer written"
+ assert torch.equal(window_buf[1:], extras["window_buf"][1:]), "another layer's snapshots"
+ return returned, got[1], got[8]
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
+ certifying(op)
+ args, extras = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
+ weight = args[2].clone()
+ ref = op.reference(*args)
+ got = _run(args, extras)
+ assert got[0] is not args[0] and got[0].shape == args[0].shape
+ assert torch.equal(args[2], weight), "weight is read-only"
+ _compare(got, ref)
+ # Non-vacuity of the whole-pool gate: every slot the batch named moved.
+ for slot in cell.spec["slots"]:
+ assert not torch.equal(got[1][slot], args[1][slot]), slot
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_update.py
new file mode 100644
index 000000000000..d958f21cb0e3
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_update.py
@@ -0,0 +1,123 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the causal_conv1d_update entry."""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.causal_conv1d_update import ( # noqa: E501
+ CKPT_DIM,
+ CKPT_X_STRIDE,
+ RANK_DIM,
+ RANK_X_STRIDE,
+ _decode_conv_ok,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.causal_conv1d_update import ( # noqa: E501
+ causal_conv1d_update as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "causal_conv1d_update requires a CUDA device"
+
+
+def _compare(
+ out: tuple[torch.Tensor, torch.Tensor], ref: tuple[torch.Tensor, torch.Tensor]
+) -> None:
+ """`(x, conv_state)` pairs: the returned `x` with the pool as the call
+ left it, against `reference`'s pair.
+
+ Activation at torch's default band: bf16/fp16 were measured bitwise
+ equal to the fp32 reference in every cell; fp32 differs by tap order
+ (max 1.43e-6, inside the default fp32 atol). The pool is a copy of
+ input values, so it is gated bitwise -- over the whole pool, which is
+ the lock on every slot the batch did not name.
+ """
+ (act, state), (ref_act, ref_state) = out, ref
+ assert act.dtype == ref_act.dtype and act.shape == ref_act.shape
+ torch.testing.assert_close(act, ref_act)
+ assert torch.equal(state, ref_state), "the advanced state must be bitwise exact"
+
+
+PAD_SLOT_ID = -1
+#: Rows of prefill ahead of the decode rows in the in_proj output.
+_PREFILL_ROWS = 37
+
+
+def _build(spec, seed):
+ """`__call__`'s arguments for one spec, plus the buffers x and the pool view.
+
+ x is the first `dim` columns of the decode rows of a `[rows, x_row_stride]`
+ projection output (its z | a | b columns past them); the pool is the last
+ layer's view of a `[slots, pool_layers, dim, 3]` page, as the V2 manager
+ coalesces the GDN layers' conv states inside each slot.
+ """
+ g = torch.Generator(device="cuda").manual_seed(seed)
+ dim, batch, layers = spec["dim"], spec["batch"], spec.get("pool_layers", 1)
+ row_stride = spec.get("x_row_stride", CKPT_X_STRIDE)
+
+ def randn(*shape):
+ return torch.randn(*shape, generator=g, device="cuda").to(torch.bfloat16)
+
+ buf = randn(_PREFILL_ROWS + batch, row_stride)
+ x = buf[_PREFILL_ROWS:, :dim]
+ assert x.stride() == (row_stride, 1)
+ pool_buf = randn(spec["pool_slots"], layers, dim, 3)
+ conv_state = pool_buf[:, layers - 1]
+ assert conv_state.stride() == (layers * dim * 3, 3, 1)
+ weight = randn(dim, 4)
+ slots = spec["slots"]
+ if slots == "random":
+ indices = torch.randperm(spec["pool_slots"], generator=g, device="cuda")[:batch]
+ indices = indices.to(torch.int32)
+ else:
+ indices = torch.tensor(slots, dtype=torch.int32, device="cuda")
+ 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)
+ 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"
+ 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]
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
+ certifying(op)
+ args, extras = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
+ weight = args[2].clone()
+ ref = op.reference(*args)
+ returned, state = _run(args, extras)
+ assert torch.equal(args[2], weight), "weight is read-only"
+ _compare((returned, state), ref)
+ # Non-vacuity: every slot the batch named moved (the whole-pool gate above
+ # holds every other slot, and every pad row, unchanged).
+ for slot in args[6].tolist():
+ if slot != PAD_SLOT_ID:
+ assert not torch.equal(state[slot], args[1][slot]), slot
+
+
+def test_which_kernel_each_member_runs() -> None:
+ """The path the note states each member's cells certify: a tp4/tep4 rank's 2560
+ channels never reach the CuTe kernel (2560 is not a multiple of 1024), so its 128-row
+ cells certify the TRT-LLM kernel; a dep4 rank's 10240 do from 64 rows."""
+ for dim, row_stride, expect in (
+ (RANK_DIM, RANK_X_STRIDE, False),
+ (CKPT_DIM, CKPT_X_STRIDE, True),
+ ):
+ spec = dict(dim=dim, x_row_stride=row_stride, batch=128, pool_slots=256, slots="random")
+ args = _build(spec, seed=dim)[0]
+ assert _decode_conv_ok(*args[:7]) is expect, (dim, expect)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py
new file mode 100644
index 000000000000..cc65c056494b
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py
@@ -0,0 +1,72 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the cu_seqlens_from_seq_lens entry.
+
+Integer arithmetic, gated bitwise. A near-miss vector is monotonic and
+plausible, so each cell also proves the vector actually splits a packed
+buffer.
+"""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.cu_seqlens_from_seq_lens import ( # noqa: E501
+ cu_seqlens_from_seq_lens as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "cu_seqlens_from_seq_lens requires a CUDA device"
+
+
+def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
+ """Bitwise: integer arithmetic, no tolerance to loosen."""
+ assert out.dtype == ref.dtype, (out.dtype, ref.dtype)
+ assert out.shape == ref.shape, (out.shape, ref.shape)
+ assert out.is_contiguous()
+ assert torch.equal(out, ref), f"{out.tolist()} != {ref.tolist()}"
+
+
+DEV = "cuda"
+
+
+def _seq_lens(lengths, dtype=torch.int32):
+ return torch.tensor(list(lengths), dtype=dtype, device=DEV)
+
+
+def _build(spec):
+ lengths, dtype = spec["lengths"], spec["in_dtype"]
+ layout = spec.get("layout", "contiguous")
+ if layout == "prefix":
+ # The metadata buffer is sized for max_num_requests; the call site
+ # passes its first `batch` entries.
+ buf = torch.full((len(lengths) + 9,), 99, dtype=dtype, device=DEV)
+ buf[: len(lengths)] = _seq_lens(lengths, dtype)
+ return buf[: len(lengths)]
+ return _seq_lens(lengths, dtype)
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
+ certifying(op)
+ spec = cell.spec
+ seq_lens = _build(spec)
+ before = seq_lens.clone()
+ out = op(seq_lens, spec["out_dtype"])
+ assert torch.equal(seq_lens, before), "seq_lens is read-only"
+ assert out.device.type == "cuda"
+ _compare(out, op.reference(seq_lens, spec["out_dtype"]))
+
+ # The entry's purpose: slicing a packed buffer with the vector recovers
+ # exactly the sequences packed into it.
+ lengths = spec["lengths"]
+ g = torch.Generator(device=DEV).manual_seed(zlib.crc32(cell.why.encode()))
+ pieces = [torch.randint(0, 10_000, (n,), generator=g, device=DEV) for n in lengths]
+ packed = torch.cat(pieces) if pieces else torch.empty(0, dtype=torch.int64, device=DEV)
+ bounds = out.tolist()
+ assert bounds[-1] == packed.shape[0]
+ for s, piece in enumerate(pieces):
+ assert torch.equal(packed[bounds[s] : bounds[s + 1]], piece), f"sequence {s} mis-sliced"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py
new file mode 100644
index 000000000000..da97829e7f5d
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py
@@ -0,0 +1,377 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the flashinfer_chunk_gated_delta_rule entry (FlashInfer GDN prefill).
+
+A cell's pool is one of two:
+
+- a contiguous pool (tp1's cells), compared whole so slots the call did not name are
+ gated too;
+- with `page_layers`, the stock cache manager's bf16 state page at a member's measured
+ slot count: `page_layers` GDN layers' states a slot, the view at the last layer's
+ sub-page. A tep4 page spans 26.5 GB and a dep4 page 48 GB, so it is a reserved virtual
+ range with physical memory mapped under the named slots only (`_SparsePage`). The
+ reference runs on a dense copy of those slots, and every other byte of their pages is
+ checked unchanged; a write to any other slot would fault.
+
+Beyond the cells: the has_initial_state equivalence.
+"""
+
+import warnings
+import zlib
+
+import cuda.bindings.driver as cuda_drv
+import pytest
+import torch
+from flashinfer.gdn_kernels.delta_rule_dsl.varlen_helper import should_use_cp_host
+from flashinfer.utils import get_compute_capability, get_device_name, get_device_sm_count
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.flashinfer_chunk_gated_delta_rule import ( # noqa: E501
+ INPLACE_MAX_TOKENS,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.flashinfer_chunk_gated_delta_rule import ( # noqa: E501
+ flashinfer_chunk_gated_delta_rule as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "flashinfer_chunk_gated_delta_rule requires a CUDA device"
+
+DEV = "cuda"
+#: The band the Triton chunk op is certified at (rtol bf16 default; atol 2e-3
+#: on `o`, 2e-2 on states), derived for a chunked WY reformulation with bf16
+#: operands. This op is the same class of sum and sits inside it against fp64, CP
+#: pipeline and chunked kernel alike; a value head reading the wrong key head at 4/12
+#: lands far outside it (measured).
+RTOL, ATOL_OUTPUT, ATOL_STATE = 1.6e-2, 2e-3, 2e-2
+#: Rows past the call's own in the output buffer of a `slice` cell: the target passes
+#: `out[:, :num_tokens]` of a wider buffer, and the rows past it must stay as they were.
+_OUTPUT_TAIL = 16
+
+
+def _compare(out: tuple, ref: tuple) -> None:
+ """`(o, final_state, pool)` against `reference`'s triple; `None` must match."""
+ for got, want, atol in zip(out, ref, (ATOL_OUTPUT, ATOL_STATE, ATOL_STATE)):
+ assert (got is None) == (want is None), (got is None, want is None)
+ if got is None:
+ continue
+ assert got.dtype == want.dtype and got.shape == want.shape, (
+ (got.dtype, got.shape),
+ (want.dtype, want.shape),
+ )
+ assert torch.isfinite(got).all()
+ torch.testing.assert_close(got.float(), want.float(), rtol=RTOL, atol=atol)
+
+
+def _stats(name: str, got: torch.Tensor, want: torch.Tensor) -> str:
+ """Relative L2, cosine, p99.9 and max absolute error of `got` against `want`."""
+ g, w = got.double().flatten(), want.double().flatten()
+ err = (g - w).abs()
+ rel = (err.norm() / w.norm().clamp_min(1e-30)).item()
+ cos = torch.nn.functional.cosine_similarity(g, w, dim=0).item()
+ p999 = torch.quantile(err[:: max(1, err.numel() // 2_000_000)].float(), 0.999).item()
+ return f"{name} rel-L2 {rel:.3e} cos {cos:.6f} p99.9 {p999:.3e} max {err.max().item():.3e}"
+
+
+def _flashinfer_path(spec) -> str:
+ """Which FlashInfer kernel the call reaches: `inplace` (the chunked kernel over the pool
+ itself) for a `has_initial_state` call of at most `INPLACE_MAX_TOKENS` tokens, else `cp`
+ (the context-parallel pipeline) when FlashInfer's own heuristic picks it for
+ `num_seqs x num_v_heads` on this device."""
+ if "initial" in spec and sum(spec["seqlens"]) <= INPLACE_MAX_TOKENS:
+ return "inplace"
+ device = torch.device(DEV)
+ cp = should_use_cp_host(
+ len(spec["seqlens"]) * spec["geometry"][1],
+ get_device_sm_count(device),
+ get_device_name(device),
+ device_capability=get_compute_capability(device),
+ )
+ return "cp" if cp else "chunked"
+
+
+def _check(result):
+ err, *value = result
+ assert err == cuda_drv.CUresult.CUDA_SUCCESS, err
+ return value[0] if len(value) == 1 else value
+
+
+class _ArrayInterface:
+ def __init__(self, ptr: int, shape, strides_bytes):
+ self.__cuda_array_interface__ = dict(
+ shape=tuple(shape),
+ strides=tuple(strides_bytes),
+ typestr=" torch.Tensor:
+ """bf16 tensor at element `offset`, element `strides`."""
+ raw = torch.as_tensor(
+ _ArrayInterface(self.va + self.ELEM * offset, shape, [self.ELEM * s for s in strides]),
+ device=DEV,
+ )
+ return raw.view(torch.bfloat16)
+
+ def close(self) -> None:
+ torch.cuda.synchronize()
+ for ptr, size, handle in self.maps:
+ _check(cuda_drv.cuMemUnmap(ptr, size))
+ _check(cuda_drv.cuMemRelease(handle))
+ _check(cuda_drv.cuMemAddressFree(cuda_drv.CUdeviceptr(self.va), self.size))
+
+
+class _EnginePool:
+ """The stock manager's coalesced bf16 state page at `slots` slots (slot stride:
+ `layers` `[HV, V, K]` states), this layer's view the last layer's sub-page, mapped
+ under the named slots: each one's whole page and the next slot's first layer. The
+ view's own base (slot 0's sub-page) is mapped too: torch reads the device off a
+ tensor's base pointer."""
+
+ def __init__(self, slots: int, layers: int, used, v_heads: int, head_v: int, head_k: int, gen):
+ layer = v_heads * head_v * head_k
+ stride = layers * layer
+ total = slots * stride
+ base = (layers - 1) * layer
+ self.used = sorted(set(used))
+ self.spans = [(s * stride, min((s + 1) * stride + layer, total)) for s in self.used]
+ self.page = _SparsePage(total, [*self.spans, (base, base + layer)])
+ self.own = base
+ self.layer = layer
+ self.pool = self.page.view(
+ base, (slots, v_heads, head_v, head_k), (stride, head_v * head_k, head_k, 1)
+ )
+ self.flat = [self.page.view(lo, (hi - lo,), (1,)) for lo, hi in self.spans]
+ for f in self.flat:
+ f.copy_(torch.randn(f.numel(), generator=gen, device=DEV) * 0.1)
+ torch.cuda.synchronize()
+ self.before = [f.clone() for f in self.flat]
+
+ def dense(self, slots: torch.Tensor):
+ """A dense copy of the named slots and `slots` renumbered into it."""
+ index = {s: i for i, s in enumerate(self.used)}
+ compact = torch.stack([self.pool[s] for s in self.used]).contiguous()
+ renamed = [index[s] for s in slots.tolist()]
+ return compact, torch.tensor(renamed, dtype=slots.dtype, device=DEV)
+
+ def touched(self) -> torch.Tensor:
+ return torch.stack([self.pool[s] for s in self.used])
+
+ def assert_rest_unchanged(self) -> None:
+ """Every byte of a named slot's page but this layer's state, and the next slot's
+ first layer, is what it was."""
+ torch.cuda.synchronize()
+ for f, b, s in zip(self.flat, self.before, self.used):
+ keep = torch.ones(f.numel(), dtype=torch.bool, device=DEV)
+ keep[self.own : self.own + self.layer] = False
+ assert torch.equal(f[keep], b[keep]), f"slot {s}: the call wrote outside its state"
+
+ def close(self) -> None:
+ del self.pool, self.flat
+ self.page.close()
+
+
+def _build(spec, seed):
+ """`(args, engine)`: keyword arguments in the domain a GDN post-conv step produces
+ (unit-norm q/k, g < 0 in log space, beta a sigmoid), and the `_EnginePool` behind the
+ pool (or None)."""
+ gen = torch.Generator(device=DEV).manual_seed(seed)
+ num_k_heads, num_v_heads, head_k_dim, head_v_dim = spec["geometry"]
+ seqlens = spec["seqlens"]
+ total = sum(seqlens)
+ q = torch.randn(1, total, num_k_heads, head_k_dim, generator=gen, device=DEV)
+ k = torch.randn(1, total, num_k_heads, head_k_dim, generator=gen, device=DEV)
+ q = q / torch.sqrt(q.pow(2).sum(-1, keepdim=True) + 1e-6)
+ k = k / torch.sqrt(k.pow(2).sum(-1, keepdim=True) + 1e-6)
+ q, k = q.bfloat16(), k.bfloat16()
+ v = torch.randn(
+ 1, total, num_v_heads, head_v_dim, generator=gen, device=DEV, dtype=torch.bfloat16
+ )
+ a_log = torch.log(torch.rand(num_v_heads, generator=gen, device=DEV) * 15 + 1)
+ x = torch.randn(1, total, num_v_heads, generator=gen, device=DEV) + torch.randn(
+ num_v_heads, generator=gen, device=DEV
+ )
+ g = (-torch.exp(a_log) * torch.nn.functional.softplus(x)).float()
+ beta = torch.rand(1, total, num_v_heads, generator=gen, device=DEV).sigmoid().float()
+ engine = None
+ if spec.get("page_layers"):
+ engine = _EnginePool(
+ spec["slots"],
+ spec["page_layers"],
+ spec["indices"],
+ num_v_heads,
+ head_v_dim,
+ head_k_dim,
+ gen,
+ )
+ pool = engine.pool
+ else:
+ pool = (
+ torch.randn(
+ spec["slots"], num_v_heads, head_v_dim, head_k_dim, generator=gen, device=DEV
+ )
+ * 0.1
+ ).to(spec.get("pool_dtype", torch.float32))
+ cu = torch.tensor([0] + torch.tensor(seqlens).cumsum(0).tolist(), dtype=torch.long, device=DEV)
+ indices = torch.tensor(spec["indices"], dtype=torch.int32, device=DEV)
+ output = None
+ if spec.get("output") == "buffer":
+ output = torch.full(
+ (1, total, num_v_heads, head_v_dim), -7.0, device=DEV, dtype=torch.bfloat16
+ )
+ elif spec.get("output") == "slice":
+ output = torch.full(
+ (1, total + _OUTPUT_TAIL, num_v_heads, head_v_dim),
+ -7.0,
+ device=DEV,
+ dtype=torch.bfloat16,
+ )[:, :total]
+ extra = {}
+ if "initial" in spec:
+ extra["has_initial_state"] = torch.tensor(spec["initial"], dtype=torch.bool, device=DEV)
+ if spec.get("linear_gate"):
+ extra["linear_gate"] = True
+ g = torch.exp(g)
+ args = dict(
+ **extra,
+ q=q,
+ k=k,
+ v=v,
+ g=g,
+ beta=beta,
+ scale=None,
+ initial_state=pool,
+ initial_state_indices=indices,
+ inplace_indexed_state_update=True,
+ output_final_state=False,
+ cu_seqlens=cu,
+ head_first=False,
+ use_qk_l2norm_in_kernel=False,
+ output=output,
+ state_workspace=None,
+ )
+ return args, engine
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
+ certifying(op)
+ args, engine = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
+ try:
+ if engine is None:
+ pool = args["initial_state"]
+ pool_before = pool.clone()
+ ref = op.reference(**args)
+ else:
+ compact, renamed = engine.dense(args["initial_state_indices"])
+ ref = op.reference(**dict(args, initial_state=compact, initial_state_indices=renamed))
+ output = args["output"]
+ # A `slice` cell's output is a prefix view of the wider buffer the target owns.
+ full = None if output is None else output._base
+ tail_before = None if full is None else full[:, output.shape[1] :].clone()
+
+ with warnings.catch_warnings(record=True) as caught:
+ warnings.simplefilter("always")
+ o, final_state = op(**args)
+ # FlashInfer falls back from a CP heuristic match with a RuntimeWarning, so a CP
+ # cell certifies the CP pipeline only if none was raised.
+ fallback = [str(w.message) for w in caught if "falling back to non-CP" in str(w.message)]
+ path = "chunked" if fallback else _flashinfer_path(cell.spec)
+
+ got_pool = pool if engine is None else engine.touched()
+ _compare((o, final_state, got_pool), ref)
+ print(
+ f"\n[{cell.why}] path={path} sms={get_device_sm_count(torch.device(DEV))}; "
+ f"{_stats('o', o, ref[0])}; {_stats('state', got_pool, ref[2])}"
+ )
+ if "fi_path" in cell.spec:
+ assert path == cell.spec["fi_path"], (path, fallback)
+ if output is not None:
+ assert o.data_ptr() == output.data_ptr(), "output= must be written in place"
+ if tail_before is not None:
+ assert torch.equal(full[:, output.shape[1] :], tail_before), "rows past output written"
+ if engine is None:
+ named = set(args["initial_state_indices"].tolist())
+ unnamed = [s for s in range(pool.shape[0]) if s not in named]
+ assert torch.equal(pool[unnamed], pool_before[unnamed]), "unnamed slots written"
+ else:
+ engine.assert_rest_unchanged()
+ finally:
+ if engine is not None:
+ engine.close()
+
+
+_CARRIED = next(c for c in op.CELLS if "initial" in c.spec)
+
+
+def test_has_initial_state_is_the_reset_then_the_plain_call() -> None:
+ """The flags replace `reset_recurrent_state_rows` before the call: zeroing the fresh
+ slots and running the plain in-place call gives the same output and pool, bitwise; a
+ linear gate that is `torch.exp` of the plain call's log gate changes nothing."""
+ args, _ = _build(_CARRIED.spec, seed=5)
+ flags = args.pop("has_initial_state")
+ linear_gate = args.pop("linear_gate", False)
+ pool = args["initial_state"]
+ fresh = args["initial_state_indices"][~flags].long()
+ pool[fresh] = torch.randn_like(pool[fresh]) # stale rows the flags must not read
+ plain = dict(args, initial_state=pool.clone(), output=args["output"].clone())
+ plain["initial_state"][fresh] = 0.0
+ if linear_gate:
+ plain["g"] = torch.log(args["g"])
+ # The linear gate of the target call is torch.exp of the log gate, bitwise.
+ args["g"] = torch.exp(plain["g"])
+
+ o_plain, _ = op(**plain)
+ o, final_state = op(**args, has_initial_state=flags, linear_gate=linear_gate)
+
+ assert final_state is None
+ assert torch.equal(o, o_plain)
+ assert torch.equal(pool, plain["initial_state"])
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py
new file mode 100644
index 000000000000..d811aa537d70
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py
@@ -0,0 +1,759 @@
+# 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.
+
+A cell with `page_layers` runs over the stock manager's coalesced state page at an
+engine's slot count (36 layers' states a slot): 26.5 GB at tep4's 1874 bf16 slots, so a
+reserved virtual range with memory mapped under the touched slots only (`_SparsePage`).
+The reference and the twin run on a dense copy of those slots, and every byte of their
+pages outside the cell's sub-pages is checked unchanged.
+"""
+
+from __future__ import annotations
+
+import zlib
+
+import cuda.bindings.driver as cuda_drv
+import pytest
+import torch
+import torch.nn.functional as F
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.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.ops.ssm.fused_recurrent_gated_delta_rule_cached_replay_update import ( # noqa: E501
+ fused_recurrent_gated_delta_rule_cached_replay_update as op,
+)
+from tensorrt_llm._torch.modules.fla.cache_manager import (
+ ReplayStateUpdateMetadata,
+ _advance_replay_state,
+)
+from tensorrt_llm._torch.modules.fla.cached_replay import (
+ CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE,
+ commit_gdn_cached_replay_history_layers,
+)
+from tensorrt_llm._torch.modules.mamba.mamba2_metadata import _build_replay_work_items_triton
+
+__extra_import_path__ = ["../../.."]
+from _validating import ULP, certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), (
+ "fused_recurrent_gated_delta_rule_cached_replay_update requires a CUDA device"
+)
+
+DEV = "cuda"
+H, HV, K, V = NUM_K_HEADS, NUM_V_HEADS, HEAD_K_DIM, HEAD_V_DIM
+W = PACKED_QKV_WIDTH
+HIST = HISTORY_SIZE
+BF16_ULP = ULP[torch.bfloat16]
+
+#: This checkpoint's 36 GDN layers of `linear_attn.A_log` / `linear_attn.dt_bias`
+#: (bf16, 48 each): mean, standard deviation, range.
+A_LOG_STATS = (1.006, 1.487, (-7.125, 5.0625))
+DT_BIAS_STATS = (-4.411, 1.732, (-11.4375, 2.53125))
+
+#: Element gate of the bf16 members, in bf16 ulps of the block they belong to.
+_ELEMENT_ULPS = 4.0
+#: Relative-RMS gate of the output, in bf16 ulps.
+_OUTPUT_RMS_ULPS = 1.5
+#: Element gate of an fp32 checkpoint state, relative to its (slot, head) scale.
+_STATE_TOL = 2.0**-13
+#: Element gate of a bf16 checkpoint state, in bf16 ulps of its (slot, head) scale.
+_BF16_STATE_ULPS = 2.0
+#: Gate of the stored cumulative log-decays: |dG| <= atol + rtol * |G|.
+_G_ATOL, _G_RTOL = 2.0**-11, 2.0**-18
+
+#: The largest deviation each gated member reached in the running cell, in units of
+#: its gate's measure (printed per cell, so the margins are on record).
+_MEASURED: dict[str, float] = {}
+
+
+def _record(key: str, value: float) -> None:
+ _MEASURED[key] = max(_MEASURED.get(key, 0.0), value)
+
+
+def _gate(name, got, want, element, rms=None):
+ """`got` against `want` per block of the last two dims (a token's
+ `[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
+ _record(f"{name} element/gate", worst / element)
+ assert worst <= element, (
+ f"{name}: an element is {worst:.3e} of its block's absmax from the reference "
+ f"(gate {element:.3e})"
+ )
+ if rms is not None and w.numel():
+ rel = ((g - w).pow(2).mean().sqrt() / w.pow(2).mean().sqrt().clamp_min(1e-30)).item()
+ _record(f"{name} rms/gate", rel / rms)
+ assert rel <= rms, f"{name}: relative RMS {rel:.3e} exceeds {rms:.3e}"
+
+
+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":
+ bound = _G_ATOL + _G_RTOL * want.abs()
+ _record("old_G |dG|/gate", ((got - want).abs() / bound).max().item())
+ torch.testing.assert_close(got, want, rtol=_G_RTOL, atol=_G_ATOL, msg=name)
+ elif name == "ssm_states":
+ bf16 = got.dtype == torch.bfloat16
+ _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, hv=HV):
+ """Per-layer A_log / dt_bias drawn around the checkpoint's own statistics;
+ head 0 pinned to the corner that forgets everything every token
+ (g ~ -400), head 1 to the one that hardly decays (g ~ -1e-8)."""
+ (a_mean, a_std, a_range), (d_mean, d_std, d_range) = A_LOG_STATS, DT_BIAS_STATS
+ a_log = torch.randn(layers, hv, generator=gen, device=DEV) * a_std + a_mean
+ dt_bias = torch.randn(layers, hv, generator=gen, device=DEV) * d_std + d_mean
+ a_log, 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)
+
+
+def _scattered(n: int, slots: int) -> list[int]:
+ """`n` distinct slots of a `slots`-slot pool: both ends first, then a stride-37 walk
+ over the interior (37 is coprime to `slots - 2` for every pool here)."""
+ rest = [(i * 37 + 4) % (slots - 2) + 1 for i in range(max(n - 2, 0))]
+ picked = ([slots - 1, 0] + rest)[:n]
+ assert len(set(picked)) == n, (n, slots)
+ return picked
+
+
+def _check(result):
+ err, *value = result
+ assert err == cuda_drv.CUresult.CUDA_SUCCESS, err
+ return value[0] if len(value) == 1 else value
+
+
+#: Array-interface type of each pool dtype; bf16 has none, so it is read as int16 and viewed.
+_TYPESTR = {torch.float32: (" torch.Tensor:
+ """Tensor of the page's dtype at element `offset`, element `strides`."""
+ typestr, carrier = _TYPESTR[self.dtype]
+ es = self.es
+ t = torch.as_tensor(
+ _ArrayInterface(self.va + es * offset, shape, [es * s for s in strides], typestr),
+ device=DEV,
+ )
+ assert t.dtype == carrier
+ return t.view(self.dtype)
+
+ def close(self) -> None:
+ torch.cuda.synchronize()
+ for ptr, size, handle in self.maps:
+ _check(cuda_drv.cuMemUnmap(ptr, size))
+ _check(cuda_drv.cuMemRelease(handle))
+ _check(cuda_drv.cuMemAddressFree(cuda_drv.CUdeviceptr(self.va), self.size))
+
+
+class _EnginePage:
+ """The stock manager's coalesced state page at `slots` slots: a slot's page holds
+ `page_layers` layers' `[HV, V, K]` states, and the bench's `layers` are its top
+ sub-pages, the page's highest addresses. Memory is mapped under each touched slot's
+ page plus the next slot's first sub-page (where a write past the page would land),
+ and under slot 0's page, which holds the views' base."""
+
+ def __init__(self, slots, page_layers, layers, used, hv, dtype, gen):
+ state = hv * V * K
+ stride = page_layers * state
+ total = slots * stride
+ self.state, self.stride, self.layers = state, stride, layers
+ self.top = (page_layers - layers) * state
+ self.used = sorted(set(used))
+ self.mapped = sorted(set(self.used) | {0})
+ self.spans = [(s * stride, min((s + 1) * stride + state, total)) for s in self.mapped]
+ self.page = _SparsePage(total, self.spans, dtype)
+ self.pools = [
+ self.page.view(self.top + j * state, (slots, hv, V, K), (stride, V * K, K, 1))
+ for j in range(layers)
+ ]
+ self.flat = [self.page.view(lo, (hi - lo,), (1,)) for lo, hi in self.spans]
+ for f in self.flat:
+ f.normal_(0.0, 0.5, generator=gen)
+ torch.cuda.synchronize()
+ self.before = [f.clone() for f in self.flat]
+
+ def assert_rest_unchanged(self) -> None:
+ """Every mapped byte outside the touched slots' own sub-pages is what it was."""
+ torch.cuda.synchronize()
+ lo, hi = self.top, self.top + self.layers * self.state
+ for s, f, b in zip(self.mapped, self.flat, self.before):
+ if s in self.used:
+ same = torch.equal(f[:lo], b[:lo]) and torch.equal(f[hi:], b[hi:])
+ else:
+ same = torch.equal(f, b)
+ assert same, f"slot {s}: the page was written outside the cell's sub-pages"
+
+ def close(self) -> None:
+ del self.pools, self.flat, self.before
+ self.page.close()
+
+
+class _Bench:
+ """The state the cache manager keeps for `layers` GDN layers, laid out as
+ `MambaHybridCacheManagerV2` lays it out, driven the way the engine drives
+ 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).
+ Over an engine page (`page_layers`) the reference and the twin see a dense copy
+ of the touched slots: `idx` the slots in order, `position` a slot's row in it.
+ """
+
+ def __init__(self, spec, seed, exact=True):
+ self.spec = spec
+ self.path, self.n, self.t = spec["path"], spec["n"], spec["t"]
+ self.h, self.hv = spec.get("H", H), spec.get("HV", HV)
+ assert spec.get("history_size", HIST) == HIST, spec
+ self.width = 2 * self.h * K + self.hv * V
+ self.layers = spec.get("layers", 2)
+ self.slots = spec.get("slots", 64)
+ self.steps = spec.get("steps", 8)
+ self.gen = torch.Generator(device=DEV).manual_seed(seed)
+ self.cpu = torch.Generator().manual_seed(seed)
+ n, t, L, S = self.n, self.t, self.layers, self.slots
+ h, hv = self.h, self.hv
+ page_layers = spec.get("page_layers")
+
+ if page_layers:
+ chosen = torch.tensor(_scattered(n, S))
+ elif spec.get("edge_slots"):
+ inner = torch.randperm(S - 2, generator=self.cpu)[: n - 2] + 1
+ chosen = torch.cat([torch.tensor([S - 1]), inner, torch.tensor([0])])
+ else:
+ chosen = torch.randperm(S, generator=self.cpu)[:n]
+ chosen = chosen.to(torch.int32)
+ live_rows = torch.ones(n, dtype=torch.bool)
+ dup = tuple(spec.get("dup", ()))
+ if dup:
+ chosen[list(dup)] = int(chosen[dup[0]])
+ # The padding rows' shared slot is nobody's state: the bookkeeping never names it.
+ live_rows[list(dup)] = False
+ self.state_indices = chosen.to(DEV)
+ self.live = self.state_indices[live_rows.to(DEV)].clone()
+
+ state_dtype = spec.get("state_dtype", torch.float32)
+ self.page = self.idx = self.position = None
+ if page_layers:
+ self.page = _EnginePage(S, page_layers, L, chosen.tolist(), hv, state_dtype, self.gen)
+ self.pool = self.page.pools
+ self.layer_stride, self.slot_stride = self.page.state, self.page.stride
+ self.idx = torch.tensor(self.page.used, dtype=torch.long, device=DEV)
+ self.position = torch.full((S,), -1, dtype=torch.long, device=DEV)
+ self.position[self.idx] = torch.arange(self.idx.numel(), device=DEV)
+ else:
+ # [slots, layers, HV, V, K]: each layer's state is a slot-strided view,
+ # the affine layout whose strides the all-layer commit is handed.
+ self.backing = (self._randn(S, L, hv, V, K) * 0.5).to(state_dtype)
+ self.pool = [self.backing[:, i] for i in range(L)]
+ self.layer_stride, self.slot_stride = self.backing.stride(1), self.backing.stride(0)
+ self.old_x = torch.zeros(L, S, 2, HIST, hv, V, device=DEV, dtype=torch.bfloat16)
+ self.old_B = torch.zeros(L, S, 2, HIST, h, K, device=DEV, dtype=torch.bfloat16)
+ self.old_dt = torch.zeros(L, S, 2, hv, HIST, device=DEV)
+ self.old_dA = self._randn(L, S, 2, hv, HIST) # handed over as old_beta; never read
+ self.pnat = torch.zeros(S, dtype=torch.int32, device=DEV)
+ self.buf = torch.zeros(S, dtype=torch.int32, device=DEV)
+ self.a_log, self.dt_bias = _gating(self.gen, L, hv)
+ if spec.get("start", "warm") == "warm":
+ self.old_x.copy_(self._randn(L, S, 2, HIST, hv, V) * 0.5)
+ self.old_B.copy_(_unit(self._randn(L, S, 2, HIST, h, K)))
+ raw_a = self._randn(L, S, 2, hv, HIST)
+ 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 = [self.compact(p).clone() for p in self.pool]
+ rows = self.x_pool[0].shape[0]
+ self.x_u = torch.zeros(L, rows, 1, self.depth, hv, V, device=DEV)
+ self.x_k = torch.zeros(L, rows, 1, self.depth, h, K, device=DEV)
+ self.x_G = torch.zeros(L, rows, 1, hv, self.depth, device=DEV)
+ self.x_pnat = self.compact(self.pnat).clone()
+ self.x_buf = torch.zeros_like(self.x_pnat)
+ for s in self.live.tolist():
+ p, b = int(self.pnat[s]), int(self.buf[s])
+ c = s if self.position is None else int(self.position[s])
+ self.x_u[:, c, 0, :p] = self.old_x[:, s, b, :p].float()
+ self.x_k[:, c, 0, :p] = self.old_B[:, s, b, :p].float()
+ self.x_G[:, c, 0, :, :p] = self.old_dt[:, s, b, :, :p]
+
+ def _randn(self, *shape, dtype=torch.float32):
+ return torch.randn(*shape, generator=self.gen, device=DEV, dtype=dtype)
+
+ def compact(self, per_slot: torch.Tensor) -> torch.Tensor:
+ """A per-slot tensor as the reference sees it: whole over a dense pool, the touched
+ slots' rows (a copy, in `idx` order) over an engine page."""
+ return per_slot if self.idx is None else per_slot[self.idx]
+
+ def compact_indices(self, slots: torch.Tensor) -> torch.Tensor:
+ """Slot ids as rows of `compact`'s output (negative ids stay negative)."""
+ if self.position is None:
+ return slots
+ renamed = torch.where(slots >= 0, self.position[slots.long().clamp(min=0)], -1)
+ return renamed.to(slots.dtype)
+
+ def close(self) -> None:
+ if self.page is not None:
+ self.pool = []
+ self.page.close()
+
+ def inputs(self):
+ """One layer's fresh activations: post-conv `[q|k|v]` rows and the raw
+ `a`/`b` projections, as the verify step hands them over."""
+ n, t, h, hv = self.n, self.t, self.h, self.hv
+ packed = self._randn(n * t, self.width, dtype=torch.bfloat16)
+ rows = packed.view(n, t, self.width)
+ q = rows[..., : h * K].view(n, t, h, K)
+ k = rows[..., h * K : 2 * h * K].view(n, t, h, K)
+ v = rows[..., 2 * h * K :].view(n, t, hv, V)
+ a = self._randn(n, t, hv, dtype=torch.bfloat16)
+ b = self._randn(n, t, hv, dtype=torch.bfloat16)
+ return dict(q=q, k=k, v=v, a=a, b=b, packed=packed)
+
+ def kwargs(self, layer, act):
+ return dict(
+ q=act["q"],
+ k=act["k"],
+ v=act["v"],
+ g=act["a"],
+ beta=act["b"],
+ ssm_states=self.pool[layer],
+ state_indices=self.state_indices,
+ old_u=self.old_x[layer],
+ old_k=self.old_B[layer],
+ old_G=self.old_dt[layer],
+ old_beta=self.old_dA[layer],
+ cache_buf_idx=self.buf,
+ prev_num_accepted_tokens=self.pnat,
+ history_size=HIST,
+ scale=None,
+ use_qk_l2norm_in_kernel=True,
+ A_log=self.a_log[layer],
+ dt_bias=self.dt_bias[layer],
+ launch_with_pdl=False,
+ replay_work_items=self.work_items,
+ n_writes=self.n_writes,
+ packed_qkv=act["packed"],
+ use_all_layer_commit=self.path == "all_layer",
+ output=None,
+ )
+
+ def reference(self, kw):
+ """The entry's `reference` on a call's arguments, every per-slot buffer through
+ `compact`: the touched slots' rows, the rows' slots renumbered into them."""
+ if self.idx is None:
+ return op.reference(**kw)
+ per_slot = (
+ "ssm_states",
+ "old_u",
+ "old_k",
+ "old_G",
+ "old_beta",
+ "cache_buf_idx",
+ "prev_num_accepted_tokens",
+ )
+ compact = {name: self.compact(kw[name]) for name in per_slot}
+ return op.reference(
+ **dict(kw, **compact, state_indices=self.compact_indices(kw["state_indices"]))
+ )
+
+ def exact_kwargs(self, layer, kw):
+ return dict(
+ kw,
+ ssm_states=self.x_pool[layer],
+ state_indices=self.compact_indices(kw["state_indices"]),
+ old_u=self.x_u[layer],
+ old_k=self.x_k[layer],
+ old_G=self.x_G[layer],
+ 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 observed(self, layer):
+ """`state(layer)` through `compact`, as `reference`'s results are laid out."""
+ return tuple(self.compact(x) for x in self.state(layer))
+
+ def unwritten(self, wrapped):
+ """Rows of `observed`'s pool the call must leave as they were: all of them on the
+ partitioned path, all but the wrapping slots' when the call commits in-kernel."""
+ rows = self.slots if self.idx is None else self.idx.numel()
+ keep = torch.ones(rows, dtype=torch.bool, device=DEV)
+ if self.path == "small":
+ keep[self.compact_indices(wrapped).long()] = False
+ return keep
+
+ def prepare(self):
+ if self.path == "all_layer":
+ self.work_items.zero_()
+ _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.layer_stride,
+ ssm_state_slot_stride=self.slot_stride,
+ old_u=self.old_x,
+ old_k=self.old_B,
+ old_G=self.old_dt,
+ replay_work_items=self.work_items,
+ n_writes=self.n_writes,
+ history_size=HIST,
+ )
+
+ def draw_acceptance(self):
+ """Accepted tokens per live row, bonus token included, by `spec["accept"]`."""
+ mode = self.spec.get("accept", "random")
+ rows = self.live.numel()
+ if mode == "full":
+ return torch.full((rows,), self.t, dtype=torch.int32, device=DEV)
+ assert mode == "random", mode
+ return torch.randint(1, self.t + 1, (rows,), generator=self.cpu, dtype=torch.int32).to(DEV)
+
+ def accept(self, accepted):
+ _advance_replay_state(self.meta, self.live, accepted)
+ if self.exact:
+ self.x_pnat[self.compact_indices(self.live).long()] += accepted
+
+ def writes(self):
+ """Where this step's calls append, by the documented rule, which slots wrap, and how
+ many append exactly up to the history's end: `[slots, 2, HIST]` mask, wrapping slot
+ ids, the `pnat + t == HIST` count."""
+ s = self.live.long()
+ p, b = self.pnat[s].long(), self.buf[s].long()
+ wrap = p + self.t > HIST
+ half, start = torch.where(wrap, 1 - b, b), torch.where(wrap, 0, p)
+ mask = torch.zeros(self.slots, 2, HIST, dtype=torch.bool, device=DEV)
+ for t in range(self.t):
+ mask[s, half, start + t] = True
+ return mask, s[wrap], int((p + self.t == HIST).sum())
+
+
+def _check_duplicate_padding_slot(cell) -> None:
+ """The last rows of a batch share one slot, as CUDA-graph padding repeats the manager's
+ dummy slot: the rows on unique slots, and every slot but the shared one, are held to
+ the reference run on those rows alone -- the padding rows change nothing about them.
+ The shared slot's state and the padding rows' outputs are unspecified."""
+ bench = _Bench(cell.spec, seed=_seed(cell.why), exact=False)
+ dup = cell.spec["dup"]
+ unique = bench.n - len(dup)
+ assert dup == tuple(range(unique, bench.n)), "the padding rows trail the batch"
+ shared = int(bench.state_indices[unique])
+ assert (bench.state_indices[unique:] == shared).all()
+ assert shared not in bench.live.tolist()
+ others = torch.ones(bench.slots, dtype=torch.bool, device=DEV)
+ others[shared] = False
+ for _ in range(bench.steps):
+ for layer in range(bench.layers):
+ act = bench.inputs()
+ kw = bench.kwargs(layer, act)
+ alone = dict(
+ kw,
+ q=kw["q"][:unique],
+ k=kw["k"][:unique],
+ v=kw["v"][:unique],
+ g=kw["g"][:unique],
+ beta=kw["beta"][:unique],
+ state_indices=kw["state_indices"][:unique],
+ packed_qkv=act["packed"][: unique * bench.t],
+ )
+ ref = op.reference(**alone)
+ got = op(**kw)
+ pool, u, k, G = bench.state(layer)
+ _compare(
+ (got[:unique], pool[others], u[others], k[others], G[others]),
+ (ref[0], *(x[others] for x in ref[1:])),
+ )
+ bench.accept(bench.draw_acceptance())
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
+ certifying(op)
+ _MEASURED.clear()
+ if "dup" in cell.spec:
+ return _check_duplicate_padding_slot(cell)
+ bench = _Bench(cell.spec, seed=_seed(cell.why))
+ try:
+ _drive_cell(cell, bench)
+ finally:
+ bench.close()
+ print(f"\n[{cell.why}] " + ", ".join(f"{k} {v:.3f}" for k, v in sorted(_MEASURED.items())))
+
+
+def _drive_cell(cell, bench) -> None:
+ wraps = appends = to_the_end = 0
+ for _ in range(bench.steps):
+ mask, wrapped, at_end = bench.writes()
+ 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)
+ _, u, k, G = bench.state(layer)
+ before = dict(
+ pool=bench.observed(layer)[0].clone(),
+ u=u.clone(),
+ k=k.clone(),
+ G=G.clone(),
+ old_beta=bench.old_dA[layer].clone(),
+ )
+ ref = bench.reference(kw)
+ exact = op.reference(**bench.exact_kwargs(layer, kw))
+ committed = None
+ if bench.path == "all_layer":
+ committed = bench.reference(
+ dict(kw, replay_work_items=None, n_writes=None, use_all_layer_commit=False)
+ )
+ got = op(**kw)
+ _compare((got, *bench.observed(layer)), ref)
+ # The locks: in place, and nothing outside the documented writes.
+ _, 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"
+ kept = bench.unwritten(wrapped)
+ assert torch.equal(bench.observed(layer)[0][kept], before["pool"][kept]), (
+ "the pool moved where the call commits nothing (a partitioned call commits none)"
+ )
+ bench.adopt_exact(layer, exact)
+ 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.observed(layer)), committed)
+ _compare((got, None, None, None, None), (exact[0], None, None, None, None))
+ if bench.page is not None:
+ bench.page.assert_rest_unchanged()
+ bench.accept(bench.draw_acceptance())
+ assert wraps > 0 and appends > 0, (
+ f"the cell never wrapped or never appended ({wraps}, {appends})"
+ )
+ if cell.spec.get("accept") == "full":
+ live = bench.live.numel()
+ assert wraps >= 3 * live and to_the_end >= 3 * live, (
+ f"full acceptance should wrap every row at least 3 times and append up to the "
+ f"history's end as often ({wraps}, {to_the_end} over {live} rows)"
+ )
+
+
+# ── CUDA graphs ───────────────────────────────────────────────────────────
+
+
+def test_cuda_graph_replay_on_new_data_is_bitwise_eager() -> None:
+ """Two layers of 32 sequences x 8 tokens over an fp32 checkpoint captured with the
+ all-layer commit after them; replayed three times on new activations, with acceptance
+ in between, against the same steps run eagerly on a twin state."""
+ n, t = 32, 8
+ assert ALL_LAYER_COMMIT_MIN_BATCH == CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE
+ assert n >= ALL_LAYER_COMMIT_MIN_BATCH
+ spec = dict(path="all_layer", n=n, t=t, layers=2, slots=40, state_dtype=torch.float32)
+ seed = 7000 + 10 * n + t
+ eager, graph, warm = (_Bench(spec, seed, exact=False) for _ in range(3))
+ acts = [warm.inputs() for _ in range(2)]
+ warm.prepare()
+ for layer in range(2):
+ op(**warm.kwargs(layer, acts[layer]))
+ warm.commit()
+ torch.cuda.synchronize()
+
+ graph.prepare()
+ captured = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(captured):
+ outs_graph = [op(**graph.kwargs(layer, acts[layer])) for layer in range(2)]
+ graph.commit()
+ assert torch.equal(graph.backing, eager.backing) and torch.equal(graph.old_x, eager.old_x)
+
+ refill = torch.Generator(device=DEV).manual_seed(seed + 1)
+ for _ in range(3):
+ for act in acts:
+ for key in ("packed", "a", "b"):
+ act[key].copy_(
+ torch.randn(act[key].shape, generator=refill, device=DEV, dtype=act[key].dtype)
+ )
+ eager.prepare()
+ outs_eager = [op(**eager.kwargs(layer, acts[layer])) for layer in range(2)]
+ eager.commit()
+ graph.prepare()
+ captured.replay()
+ torch.cuda.synchronize()
+ for e, g in zip(outs_eager, outs_graph):
+ assert torch.equal(e, g), "replayed output differs from eager"
+ for name in ("backing", "old_x", "old_B", "old_dt", "pnat", "buf"):
+ assert torch.equal(getattr(eager, name), getattr(graph, name)), f"{name} differs"
+ accepted = eager.draw_acceptance()
+ eager.accept(accepted)
+ graph.accept(accepted)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py
new file mode 100644
index 000000000000..fafe33a39b78
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py
@@ -0,0 +1,311 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the gdn_decode_gated_norm_pdl entry.
+
+Cells against the composition's fp64 reference (the gated delta rule's recurrence,
+then the gated norm entry's) in torch's default band for the bf16 output and the fp32
+pool. A cell's pool is one of three:
+
+- a contiguous 461-slot pool, compared whole so slots the call did not name are gated too;
+- the tp1 coalesced page (36 layers' states a slot, 102 slots), the view at layer 0;
+- with `pool_slots`, the stock cache manager's page at an engine's measured slot count:
+ 36 GDN layers' states a slot, the view at the last layer's sub-page. Such a page spans
+ 53 GB (tep4) or 96 GB (dep4) in fp32, so it is a reserved virtual range with physical
+ memory mapped under the touched slots only (`_SparsePage`). The reference runs on a dense
+ copy of those slots, and every other byte of their pages is checked unchanged.
+
+Beyond the cells: a stale-read check under graph replay at both head layouts (a producer
+rewrites the projection inside the graph right before the op).
+"""
+
+import zlib
+
+import cuda.bindings.driver as cuda_drv
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.gdn_decode_gated_norm_pdl import ( # noqa: E501
+ gdn_decode_gated_norm_pdl as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "gdn_decode_gated_norm_pdl requires a CUDA device"
+
+DEV = "cuda"
+K, V = 128, 128
+SLOTS, EPS = 461, 1e-6
+#: GDN layers whose states share one slot page in the stock manager (48 layers, every 4th
+#: full attention); the view is at the last one's sub-page, the page's highest addresses.
+PAGE_LAYERS = 36
+
+
+def _heads(spec) -> tuple[int, int]:
+ """(k heads, v heads): tp1's 16 / 48 where the cell names none."""
+ return spec.get("k_heads", 16), spec.get("v_heads", 48)
+
+
+def _width(k_heads: int, v_heads: int) -> int:
+ """`[q | k | v | z | a | b]`: 16480 at 16 / 48, 4120 at a tp4/tep4 rank's 4 / 12."""
+ return 2 * k_heads * K + 2 * v_heads * V + 2 * v_heads
+
+
+def _scattered(n: int, slots: int) -> tuple[int, ...]:
+ """`n` distinct slots of a `slots`-row pool, both ends first, then a stride-37 walk
+ over the interior (37 is coprime to every pool size here)."""
+ rest = tuple((i * 37 + 4) % (slots - 2) + 1 for i in range(max(n - 2, 0)))
+ return ((slots - 1, 0) + rest)[:n]
+
+
+def _check(result):
+ err, *value = result
+ assert err == cuda_drv.CUresult.CUDA_SUCCESS, err
+ return value[0] if len(value) == 1 else value
+
+
+class _ArrayInterface:
+ def __init__(self, ptr: int, shape, strides_bytes):
+ self.__cuda_array_interface__ = dict(
+ shape=tuple(shape),
+ strides=tuple(strides_bytes),
+ typestr=" torch.Tensor:
+ """fp32 tensor at element `offset`, element `strides`."""
+ return torch.as_tensor(
+ _ArrayInterface(self.va + 4 * offset, shape, [4 * s for s in strides]), device=DEV
+ )
+
+ def close(self) -> None:
+ torch.cuda.synchronize()
+ for ptr, size, handle in self.maps:
+ _check(cuda_drv.cuMemUnmap(ptr, size))
+ _check(cuda_drv.cuMemRelease(handle))
+ _check(cuda_drv.cuMemAddressFree(cuda_drv.CUdeviceptr(self.va), self.size))
+
+
+class _EnginePool:
+ """The stock manager's coalesced state page at `slots` slots (measured slot stride: 36
+ layers' `[HV, V, K]` states), this layer's view the last layer's sub-page, mapped under
+ the touched slots: each one's whole page and the next slot's first layer. The view's own
+ base (slot 0's sub-page) is mapped too: torch reads the device off a tensor's base
+ pointer."""
+
+ def __init__(self, slots: int, used, v_heads: int, gen):
+ layer = v_heads * V * K
+ stride = PAGE_LAYERS * layer
+ total = slots * stride
+ base = (PAGE_LAYERS - 1) * layer
+ self.used = sorted(set(used))
+ self.spans = [(s * stride, min((s + 1) * stride + layer, total)) for s in self.used]
+ self.page = _SparsePage(total, [*self.spans, (base, base + layer)])
+ # This layer's state within each mapped span: every span starts at a slot's page.
+ self.own = base
+ self.layer = layer
+ self.pool = self.page.view(base, (slots, v_heads, V, K), (stride, V * K, K, 1))
+ self.flat = [self.page.view(lo, (hi - lo,), (1,)) for lo, hi in self.spans]
+ for f in self.flat:
+ f.normal_(0.0, 0.05, generator=gen)
+ torch.cuda.synchronize()
+ self.before = [f.clone() for f in self.flat]
+
+ def dense(self, slots: torch.Tensor):
+ """A dense copy of the touched slots and `slots` renumbered into it (-1 stays)."""
+ index = {s: i for i, s in enumerate(self.used)}
+ compact = torch.stack([self.pool[s] for s in self.used]).contiguous()
+ renamed = [index[s] if s >= 0 else -1 for s in slots.tolist()]
+ return compact, torch.tensor(renamed, dtype=torch.int32, device=DEV)
+
+ def touched(self) -> torch.Tensor:
+ return torch.stack([self.pool[s] for s in self.used])
+
+ def assert_rest_unchanged(self) -> None:
+ """Every byte of a touched slot's page but this layer's state, and the next slot's
+ first layer, is what it was."""
+ torch.cuda.synchronize()
+ for f, b, s in zip(self.flat, self.before, self.used):
+ keep = torch.ones(f.numel(), dtype=torch.bool, device=DEV)
+ keep[self.own : self.own + self.layer] = False
+ assert torch.equal(f[keep], b[keep]), f"slot {s}: the call wrote outside its state"
+
+ def close(self) -> None:
+ del self.pool, self.flat
+ self.page.close()
+
+
+def _build(spec, seed):
+ """`(args, engine)`: the call's arguments, and the `_EnginePool` behind them (or None)."""
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ rows = spec["rows"]
+ k_heads, v_heads = _heads(spec)
+ proj = torch.randn(rows, _width(k_heads, v_heads), generator=g, device=DEV).to(torch.bfloat16)
+ engine = None
+ if spec.get("pool_slots"):
+ n = spec["pool_slots"]
+ picked = (n - 1, n - 2) if rows == 2 else _scattered(rows, n)
+ engine = _EnginePool(n, picked, v_heads, g)
+ pool = engine.pool
+ slots = torch.tensor(picked, dtype=torch.int32, device=DEV)
+ elif spec.get("pool") == "coalesced":
+ # The KV-cache-v2 manager's view: every GDN layer's state of a slot in one page,
+ # this layer's state one sub-page of it, so the slot stride is 36 layers' states
+ # and slot 100 sits ~2.8G elements past the base (past int32).
+ layers, num = PAGE_LAYERS, 102
+ page = torch.empty(num * layers * v_heads * V * K, device=DEV)
+ pool = page.as_strided((num, v_heads, V, K), (layers * v_heads * V * K, V * K, K, 1))
+ # The pool is compared whole: no slot of the view is left uninitialized.
+ pool.zero_()
+ pool[0].normal_(0.0, 0.05, generator=g)
+ for s in (100, 101):
+ pool[s].copy_(torch.randn(v_heads, V, K, generator=g, device=DEV) * 0.05)
+ slots = torch.tensor([100, 101][:rows], dtype=torch.int32, device=DEV)
+ else:
+ pool = torch.randn(SLOTS, v_heads, V, K, generator=g, device=DEV) * 0.05
+ slots = torch.tensor(_scattered(rows, SLOTS), dtype=torch.int32, device=DEV)
+ if spec.get("pad") and rows > 1:
+ slots[1] = -1
+ A_log = torch.empty(v_heads, device=DEV).uniform_(-5.5, 1.0, generator=g)
+ dt_bias = torch.randn(v_heads, generator=g, device=DEV) * 4.0
+ w = (1.0 + 0.2 * torch.randn(V, generator=g, device=DEV)).to(torch.bfloat16)
+ return [proj, pool, slots, A_log, dt_bias, w, EPS, k_heads, K], engine
+
+
+def _compare(out, pool, ref_out, ref_pool) -> None:
+ assert out.dtype == torch.bfloat16 and out.shape == ref_out.shape and out.is_contiguous()
+ assert torch.isfinite(out).all()
+ torch.testing.assert_close(out, ref_out)
+ torch.testing.assert_close(pool, ref_pool)
+
+
+def _stats(out, ref_out, pool, ref_pool) -> str:
+ """What `note` quotes: the output against its bf16-rounded reference, the state against
+ the fp32-rounded one."""
+ o, r = out.float().flatten(), ref_out.float().flatten()
+ cos = torch.nn.functional.cosine_similarity(o, r, dim=0).item()
+ rel = ((o - r).norm() / r.norm().clamp_min(1e-30)).item()
+ differ = (out.view(torch.int16) != ref_out.view(torch.int16)).sum().item()
+ p, q = pool.flatten(), ref_pool.flatten()
+ p_rel = ((p - q).norm() / q.norm().clamp_min(1e-30)).item()
+ return (
+ f"out max|d| {(o - r).abs().max().item():.3e}, {differ}/{o.numel()} elements differ, "
+ f"cos {cos:.8f}, rel L2 {rel:.3e}; state max|d| {(p - q).abs().max().item():.3e}, "
+ f"rel L2 {p_rel:.3e}"
+ )
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
+ certifying(op)
+ spec = cell.spec
+ args, engine = _build(spec, seed=zlib.crc32(cell.why.encode()) % (2**31))
+ try:
+ if engine is None:
+ ref_out, ref_pool = op.reference(*args)
+ out = op(*args)
+ pool = args[1]
+ else:
+ dense, renamed = engine.dense(args[2])
+ ref_out, ref_pool = op.reference(args[0], dense, renamed, *args[3:])
+ out = op(*args)
+ torch.cuda.synchronize()
+ pool = engine.touched()
+ _compare(out, pool, ref_out, ref_pool)
+ if engine is not None:
+ engine.assert_rest_unchanged()
+ print(f"[{cell.why}] {_stats(out, ref_out, pool, ref_pool)}")
+ if spec.get("pad"):
+ if spec["rows"] > 1:
+ assert (out[1] == 0).all(), "a pad row's output is zero"
+ else:
+ # One row: a second call over the pad slot (-1) writes zeros and touches no state.
+ state = engine.touched().clone() if engine else args[1].clone()
+ pad = torch.full_like(args[2], -1)
+ out_pad = op(args[0], args[1], pad, *args[3:])
+ torch.cuda.synchronize()
+ assert (out_pad == 0).all(), "a pad row's output is zero"
+ after = engine.touched() if engine else args[1]
+ assert torch.equal(after, state), "a pad row touches no state"
+ if engine:
+ engine.assert_rest_unchanged()
+ finally:
+ if engine is not None:
+ engine.close()
+
+
+@pytest.mark.parametrize("k_heads, v_heads", [(16, 48), (4, 12)], ids=["16x48", "4x12"])
+def test_reads_after_the_wait_under_graph_replay(k_heads, v_heads) -> None:
+ """A producer rewrites the projection inside the graph right before the op; with PDL
+ a read of it before `griddepcontrol.wait` would see the previous replay's row. The
+ pool advances one step per replay, as in decode."""
+ spec = dict(rows=1, k_heads=k_heads, v_heads=v_heads)
+ args, _ = _build(spec, seed=41)
+ src = args[0].clone()
+
+ def step():
+ args[0].copy_(src * 1.0)
+ return op(*args)
+
+ warmup = torch.cuda.Stream()
+ warmup.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(warmup):
+ step()
+ torch.cuda.current_stream().wait_stream(warmup)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ graph_out = step()
+ for trial in range(3):
+ src.copy_(_build(spec, seed=100 + trial)[0][0])
+ before = args[1].clone()
+ graph.replay()
+ torch.cuda.synchronize()
+ after = args[1].clone()
+ args[1].copy_(before)
+ want = op(src.clone(), *args[1:])
+ assert torch.equal(graph_out, want), f"replay {trial} output"
+ assert torch.equal(after, args[1]), f"replay {trial} pool"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py
new file mode 100644
index 000000000000..d9688b4400f2
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py
@@ -0,0 +1,471 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the gdn_decode_rows_gated_norm entry.
+
+Cells against the composition's fp64 reference (the gated delta rule's recurrence,
+then the gated norm entry's) in torch's default band for the bf16 output and the pool's
+dtype (fp32 or bf16). A cell's pool is one of three:
+
+- a contiguous 461-slot pool, compared whole so slots the call did not name are gated too;
+- the tp1 coalesced page (36 layers' states a slot, 102 slots), the view at layer 0;
+- with `pool_slots`, the stock cache manager's page at an engine's slot count: 36 GDN
+ layers' states a slot, the view at the last layer's sub-page. Such a page spans 26 GB
+ (tep4) or 48 GB (dep4) in bf16, so it is a reserved virtual range with physical memory
+ mapped under the touched slots only (`_SparsePage`). The reference runs on a dense copy
+ of those slots, and every other byte of their pages is checked unchanged.
+
+Beyond the cells: a fixed row bitwise across batch sizes over a bf16 pool, a stale-read
+check under graph replay (a producer rewrites the projection inside the graph right
+before the op), the two launch variants (state rows by bulk copy into shared memory, or
+loaded into registers) bitwise each other, and the variant choice over an unaligned
+pool, each at tp1's 16 / 48 heads and a tp4/tep4 rank's 4 / 12. The `gdn_decode_rows`
+entry (the recurrence alone) has its own cells below.
+"""
+
+import zlib
+
+import cuda.bindings.driver as cuda_drv
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm import (
+ gdn_decode_rows_gated_norm as rows_mod,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.gdn_decode_rows_gated_norm import ( # noqa: E501
+ gdn_decode_rows as rows_op,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.gdn_decode_rows_gated_norm import ( # noqa: E501
+ gdn_decode_rows_gated_norm as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "gdn_decode_rows_gated_norm requires a CUDA device"
+
+DEV = "cuda"
+H, HV, K, V = 16, 48, 128, 128
+SLOTS, EPS = 461, 1e-6
+WIDTH = 2 * H * K + 2 * HV * V + 2 * HV
+#: GDN layers whose states share one slot page in the stock manager (48 layers, every 4th
+#: full attention); the view is at the last one's sub-page, the page's highest addresses.
+PAGE_LAYERS = 36
+#: (k heads, v heads) of a tp4/tep4 rank.
+RANK_HEADS = (4, 12)
+
+
+def _heads(spec) -> tuple[int, int]:
+ """(k heads, v heads): tp1's 16 / 48 where the cell names none."""
+ return spec.get("k_heads", H), spec.get("v_heads", HV)
+
+
+def _width(k_heads: int, v_heads: int) -> int:
+ """`[q | k | v | z | a | b]`: 16480 at 16 / 48, 4120 at a tp4/tep4 rank's 4 / 12."""
+ return 2 * k_heads * K + 2 * v_heads * V + 2 * v_heads
+
+
+def _scattered(n: int, slots: int) -> tuple[int, ...]:
+ """`n` distinct slots of a `slots`-row pool, both ends first, then a stride-37 walk
+ over the interior (37 is coprime to every pool size here)."""
+ rest = tuple((i * 37 + 4) % (slots - 2) + 1 for i in range(max(n - 2, 0)))
+ return ((slots - 1, 0) + rest)[:n]
+
+
+def _check(result):
+ err, *value = result
+ assert err == cuda_drv.CUresult.CUDA_SUCCESS, err
+ return value[0] if len(value) == 1 else value
+
+
+#: Array-interface type of each pool dtype; bf16 has none, so it is read as int16 and viewed.
+_TYPESTR = {torch.float32: (" torch.Tensor:
+ """Tensor of the page's dtype at element `offset`, element `strides`."""
+ typestr, carrier = _TYPESTR[self.dtype]
+ es = self.es
+ t = torch.as_tensor(
+ _ArrayInterface(self.va + es * offset, shape, [es * s for s in strides], typestr),
+ device=DEV,
+ )
+ assert t.dtype == carrier
+ return t.view(self.dtype)
+
+ def close(self) -> None:
+ torch.cuda.synchronize()
+ for ptr, size, handle in self.maps:
+ _check(cuda_drv.cuMemUnmap(ptr, size))
+ _check(cuda_drv.cuMemRelease(handle))
+ _check(cuda_drv.cuMemAddressFree(cuda_drv.CUdeviceptr(self.va), self.size))
+
+
+class _EnginePool:
+ """The stock manager's coalesced state page at `slots` slots (slot stride: 36 layers'
+ `[HV, V, K]` states), this layer's view the last layer's sub-page, mapped under the
+ touched slots: each one's whole page and the next slot's first layer. The view's own
+ base (slot 0's sub-page) is mapped too: torch reads the device off a tensor's base
+ pointer."""
+
+ def __init__(self, slots: int, used, v_heads: int, dtype: torch.dtype, gen):
+ layer = v_heads * V * K
+ stride = PAGE_LAYERS * layer
+ total = slots * stride
+ base = (PAGE_LAYERS - 1) * layer
+ self.used = sorted(set(used))
+ self.spans = [(s * stride, min((s + 1) * stride + layer, total)) for s in self.used]
+ self.page = _SparsePage(total, [*self.spans, (base, base + layer)], dtype)
+ # This layer's state within each mapped span: every span starts at a slot's page.
+ self.own = base
+ self.layer = layer
+ self.pool = self.page.view(base, (slots, v_heads, V, K), (stride, V * K, K, 1))
+ self.flat = [self.page.view(lo, (hi - lo,), (1,)) for lo, hi in self.spans]
+ for f in self.flat:
+ f.normal_(0.0, 0.05, generator=gen)
+ torch.cuda.synchronize()
+ self.before = [f.clone() for f in self.flat]
+
+ def dense(self, slots: torch.Tensor):
+ """A dense copy of the touched slots and `slots` renumbered into it (-1 stays)."""
+ index = {s: i for i, s in enumerate(self.used)}
+ compact = torch.stack([self.pool[s] for s in self.used]).contiguous()
+ renamed = [index[s] if s >= 0 else -1 for s in slots.tolist()]
+ return compact, torch.tensor(renamed, dtype=torch.int32, device=DEV)
+
+ def touched(self) -> torch.Tensor:
+ return torch.stack([self.pool[s] for s in self.used])
+
+ def assert_rest_unchanged(self) -> None:
+ """Every byte of a touched slot's page but this layer's state, and the next slot's
+ first layer, is what it was."""
+ torch.cuda.synchronize()
+ lo, hi = self.own, self.own + self.layer
+ for f, b, s in zip(self.flat, self.before, self.used):
+ assert torch.equal(f[:lo], b[:lo]) and torch.equal(f[hi:], b[hi:]), (
+ f"slot {s}: the call wrote outside its state"
+ )
+
+ def close(self) -> None:
+ del self.pool, self.flat
+ self.page.close()
+
+
+def _build(spec, seed):
+ """`(args, engine)`: the call's arguments, and the `_EnginePool` behind them (or None)."""
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ rows = spec["rows"]
+ k_heads, v_heads = _heads(spec)
+ width = _width(k_heads, v_heads)
+ assert spec.get("proj_width", width) == width, (spec, width)
+ state = torch.bfloat16 if spec.get("state") == "bf16" else torch.float32
+ proj = torch.randn(rows, width, generator=g, device=DEV).to(torch.bfloat16)
+ engine = None
+ if spec.get("pool_slots"):
+ n = spec["pool_slots"]
+ picked = _scattered(rows, n)
+ engine = _EnginePool(n, picked, v_heads, state, g)
+ pool = engine.pool
+ slots = torch.tensor(picked, dtype=torch.int32, device=DEV)
+ elif spec.get("pool") == "coalesced":
+ # The KV-cache-v2 manager's view: every GDN layer's state of a slot in one page,
+ # this layer's state one sub-page of it, so the slot stride is 36 layers' states
+ # and slot 100 sits ~2.8G elements past the base (past int32).
+ layers, num = PAGE_LAYERS, 102
+ page = torch.empty(num * layers * v_heads * V * K, device=DEV, dtype=state)
+ page[: v_heads * V * K * 4].normal_(0.0, 0.05, generator=g)
+ pool = page.as_strided((num, v_heads, V, K), (layers * v_heads * V * K, V * K, K, 1))
+ for s in (100, 101):
+ pool[s].copy_(torch.randn(v_heads, V, K, generator=g, device=DEV) * 0.05)
+ slots = torch.tensor([100, 101][:rows], dtype=torch.int32, device=DEV)
+ else:
+ pool = (torch.randn(SLOTS, v_heads, V, K, generator=g, device=DEV) * 0.05).to(state)
+ slots = torch.tensor(_scattered(rows, SLOTS), dtype=torch.int32, device=DEV)
+ if spec.get("pad"):
+ slots[1] = -1
+ A_log = torch.empty(v_heads, device=DEV).uniform_(-5.5, 1.0, generator=g)
+ dt_bias = torch.randn(v_heads, generator=g, device=DEV) * 4.0
+ w = (1.0 + 0.2 * torch.randn(V, generator=g, device=DEV)).to(torch.bfloat16)
+ return [proj, pool, slots, A_log, dt_bias, w, EPS, k_heads, K], engine
+
+
+def _compare(out, pool, ref_out, ref_pool) -> None:
+ assert out.dtype == torch.bfloat16 and out.shape == ref_out.shape and out.is_contiguous()
+ assert torch.isfinite(out).all()
+ torch.testing.assert_close(out, ref_out)
+ torch.testing.assert_close(pool, ref_pool)
+
+
+def _stats(out, ref_out, pool, ref_pool) -> str:
+ """What `note` quotes: the output against its bf16-rounded reference, the state against
+ the reference rounded to the pool's dtype. Elements are counted by value (a pad row's
+ -0.0 equals the reference's 0.0)."""
+ o, r = out.float().flatten(), ref_out.float().flatten()
+ cos = torch.nn.functional.cosine_similarity(o, r, dim=0).item()
+ rel = ((o - r).norm() / r.norm().clamp_min(1e-30)).item()
+ differ = (o != r).sum().item()
+ p, q = pool.float().flatten(), ref_pool.float().flatten()
+ p_rel = ((p - q).norm() / q.norm().clamp_min(1e-30)).item()
+ p_differ = (p != q).sum().item()
+ return (
+ f"out max|d| {(o - r).abs().max().item():.3e}, {differ}/{o.numel()} elements differ, "
+ f"cos {cos:.8f}, rel L2 {rel:.3e}; state max|d| {(p - q).abs().max().item():.3e}, "
+ f"{p_differ}/{p.numel()} differ, rel L2 {p_rel:.3e}"
+ )
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
+ certifying(op)
+ args, engine = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31))
+ try:
+ if engine is None:
+ ref_out, ref_pool = op.reference(*args)
+ out = op(*args)
+ pool = args[1]
+ else:
+ dense, renamed = engine.dense(args[2])
+ ref_out, ref_pool = op.reference(args[0], dense, renamed, *args[3:])
+ out = op(*args)
+ torch.cuda.synchronize()
+ pool = engine.touched()
+ _compare(out, pool, ref_out, ref_pool)
+ if engine is not None:
+ engine.assert_rest_unchanged()
+ print(f"[{cell.why}] {_stats(out, ref_out, pool, ref_pool)}")
+ if cell.spec.get("pad"):
+ assert (out[1] == 0).all(), "a pad row's output is zero"
+ finally:
+ if engine is not None:
+ engine.close()
+
+
+#: Decode batch sizes up to the engine's largest CUDA-graph batch.
+_BATCHES = (1, 2, 4, 5, 16, 33, 128, 256)
+_FIXED_SLOT = 301
+
+
+@pytest.mark.parametrize("k_heads, v_heads", [(H, HV), RANK_HEADS], ids=["16x48", "4x12"])
+def test_bf16_pool_is_batch_invariant(k_heads, v_heads) -> None:
+ """The property the bf16 decode step relies on: a row's output and state are bitwise
+ independent of the batch it rides in. One fixed row (its projection and slot 301 of a
+ bf16 pool) sits first and last in random batches of every size, across the batch where
+ the launch switches to the bulk-copy variant (7 rows at 48 v heads, 26 at 12)."""
+ width = _width(k_heads, v_heads)
+ gen = torch.Generator(device=DEV).manual_seed(3)
+ cpu = torch.Generator().manual_seed(3)
+ pool0 = (torch.randn(SLOTS, v_heads, V, K, generator=gen, device=DEV) * 0.05).to(torch.bfloat16)
+ row = torch.randn(width, generator=gen, device=DEV).to(torch.bfloat16)
+ rest = _build(dict(rows=1, k_heads=k_heads, v_heads=v_heads), seed=3)[0][3:]
+ first = None
+ for n in _BATCHES:
+ for pos in sorted({0, n - 1}):
+ g = torch.Generator(device=DEV).manual_seed(1000 * n + pos)
+ proj = torch.randn(n, width, generator=g, device=DEV).to(torch.bfloat16)
+ proj[pos] = row
+ others = [s for s in torch.randperm(SLOTS, generator=cpu).tolist() if s != _FIXED_SLOT]
+ others = others[:n]
+ others[pos] = _FIXED_SLOT
+ slots = torch.tensor(others, dtype=torch.int32, device=DEV)
+ pool = pool0.clone()
+ out = op(proj, pool, slots, *rest)
+ got = (out[pos].clone(), pool[_FIXED_SLOT].clone())
+ first = first or got
+ assert torch.equal(got[0], first[0]) and torch.equal(got[1], first[1]), (
+ f"the fixed row differs at batch {n}, position {pos}"
+ )
+
+
+@pytest.mark.parametrize(
+ "spec",
+ (
+ dict(rows=16),
+ dict(rows=3, k_heads=4, v_heads=12, state="bf16"),
+ dict(rows=128, k_heads=4, v_heads=12, state="bf16"),
+ ),
+ ids=["16x48-16rows", "4x12-3rows-bf16", "4x12-128rows-bf16"],
+)
+def test_reads_after_the_wait_under_graph_replay(spec) -> None:
+ """A producer rewrites the projection inside the graph right before the op; with PDL
+ a read of it before `griddepcontrol.wait` would see the previous replay's row. The
+ pool advances one step per replay, as in decode. At 4 / 12 heads 3 rows take the
+ register variant, 128 the bulk copy."""
+ args = _build(spec, seed=41)[0]
+ src = args[0].clone()
+
+ def step():
+ args[0].copy_(src * 1.0)
+ return op(*args)
+
+ warmup = torch.cuda.Stream()
+ warmup.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(warmup):
+ step()
+ torch.cuda.current_stream().wait_stream(warmup)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ graph_out = step()
+ for trial in range(3):
+ src.copy_(_build(spec, seed=100 + trial)[0][0])
+ before = args[1].clone()
+ graph.replay()
+ torch.cuda.synchronize()
+ after = args[1].clone()
+ args[1].copy_(before)
+ want = op(src.clone(), *args[1:])
+ assert torch.equal(graph_out, want), f"replay {trial} output"
+ assert torch.equal(after, args[1]), f"replay {trial} pool"
+
+
+_VARIANT_SPECS = (
+ dict(rows=1, state="bf16"),
+ dict(rows=33, state="bf16"),
+ dict(rows=33, pad=True),
+ dict(rows=2, pool="coalesced", state="bf16", pad=True),
+ dict(rows=3, k_heads=4, v_heads=12, state="bf16", pad=True),
+ dict(rows=128, k_heads=4, v_heads=12, state="bf16"),
+ dict(rows=3, k_heads=4, v_heads=12, state="bf16", pool_slots=rows_mod.TEP4_SLOTS),
+)
+
+
+@pytest.mark.parametrize("norm", [True, False], ids=["norm", "rows"])
+@pytest.mark.parametrize("spec", _VARIANT_SPECS, ids=lambda s: "-".join(map(str, s.values())))
+def test_bulk_and_register_variants_agree(spec, norm) -> None:
+ """The bulk-copy and the register variants: outputs and state bitwise equal.
+
+ The bulk-copy variant stages the state rows in shared memory; the register variant holds
+ them in registers; both compute the same instructions. Over contiguous bf16 and fp32 pools,
+ the tp1 coalesced page (slot 100, ~2.8G elements past the base, by the bulk copy) and
+ tep4's engine page (slot 1873, ~1.3e10 elements past the base), with pad rows, at 16 / 48
+ and 4 / 12 heads.
+ """
+ args, engine = _build(spec, seed=zlib.crc32(repr(spec).encode()) % (2**31))
+ try:
+ proj, pool, slots, k_heads = args[0], args[1], args[2], args[7]
+ assert rows_mod._bulk_eligible(pool, K)
+ live = slots[slots >= 0].long()
+ before = pool[live].clone()
+ got = []
+ for bulk in (False, True):
+ pool[live] = before
+ w = args[5] if norm else None
+ out = rows_mod.launch(
+ proj, pool, slots, args[3], args[4], w, EPS, k_heads, K, bulk=bulk
+ )
+ got.append((out.clone(), pool[live].clone()))
+ assert torch.equal(got[0][0], got[1][0]), "outputs differ between the variants"
+ assert torch.equal(got[0][1], got[1][1]), "states differ between the variants"
+ if engine is not None:
+ engine.assert_rest_unchanged()
+ finally:
+ if engine is not None:
+ engine.close()
+
+
+@pytest.mark.parametrize("v_heads", [HV, RANK_HEADS[1]], ids=["48", "12"])
+def test_variant_choice_and_unaligned_pool(v_heads) -> None:
+ """The variant choice, and the fallback for a pool the bulk copy cannot take.
+
+ The bulk variant from the batch the register variant needs a second wave for; a pool 8-
+ but not 16-byte aligned stays on the register variant whatever is asked, and matches the
+ aligned pool's result.
+ """
+ sms = torch.cuda.get_device_properties(DEV).multi_processor_count
+ edge = (2 * sms) // v_heads
+ assert not rows_mod._use_bulk(edge, v_heads, DEV) and rows_mod._use_bulk(edge + 1, v_heads, DEV)
+ print(f"[{v_heads} v heads] {sms} SMs: the bulk-copy variant from {edge + 1} rows")
+ k_heads = v_heads // 3
+ args = _build(dict(rows=33, k_heads=k_heads, v_heads=v_heads, state="bf16"), seed=19)[0]
+ pool = args[1]
+ flat = torch.empty(pool.numel() + 4, dtype=pool.dtype, device=DEV)
+ shifted = flat[4:].view(pool.shape)
+ shifted.copy_(pool)
+ assert not rows_mod._bulk_eligible(shifted, K)
+ want = rows_mod.launch(args[0], pool, *args[2:], bulk=True)
+ got = rows_mod.launch(args[0], shifted, *args[2:], bulk=True)
+ assert torch.equal(got, want) and torch.equal(shifted, pool)
+
+
+# ---- gdn_decode_rows: the recurrence alone, into a row range of a step's output ----------
+
+
+def _rows_args(spec, seed):
+ """`(args, buf)`: the entry's arguments, `out` rows 5.. of a `[1, rows + 9, HV, V]` buffer
+ (the step's output, the call's rows a range of it), and that buffer."""
+ (proj, pool, slots, A_log, dt_bias, _, _, h, k), _ = _build(spec, seed)
+ rows, v_heads = proj.shape[0], pool.shape[1]
+ g = torch.Generator(device=DEV).manual_seed(seed + 1)
+ buf = torch.randn(1, rows + 9, v_heads, V, generator=g, device=DEV).to(torch.bfloat16)
+ out = buf[0, 5 : 5 + rows].view(rows, v_heads * V)
+ return [proj, pool, slots, A_log, dt_bias, h, k, out], buf
+
+
+@pytest.mark.parametrize("cell", rows_op.CELLS, ids=[c.why[:44] for c in rows_op.CELLS])
+def test_rows_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
+ certifying(rows_op)
+ args, buf = _rows_args(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31))
+ outside = torch.cat([buf[0, :5], buf[0, 5 + args[0].shape[0] :]]).clone()
+ ref_o, ref_pool = rows_op.reference(*args)
+ got = rows_op(*args)
+ assert got.data_ptr() == args[-1].data_ptr(), "written into out"
+ torch.testing.assert_close(got, ref_o)
+ torch.testing.assert_close(args[1], ref_pool)
+ print(f"[{cell.why}] {_stats(got, ref_o, args[1], ref_pool)}")
+ assert torch.equal(torch.cat([buf[0, :5], buf[0, 5 + args[0].shape[0] :]]), outside), (
+ "rows outside the range untouched"
+ )
+ if cell.spec.get("pad"):
+ assert (got[1] == 0).all(), "a pad row's output is zero"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_prefill_conv_post.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_prefill_conv_post.py
new file mode 100644
index 000000000000..a49ba3fde83a
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_prefill_conv_post.py
@@ -0,0 +1,197 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the gdn_prefill_conv_post entry.
+
+Every cell is checked against the entry's fp32/fp64 reference, the conv pool
+bitwise. Geometry is a Qwen3.8-Flash-Next linear-attention layer, conv width 4, bf16
+activations: the whole layer (16 x 128 key heads, 48 x 128 value heads, qkv_dim 10240,
+a 16480-wide projection) at tp1 and on a dep4 rank, a quarter of it (4 / 12 heads,
+qkv_dim 2560, a 4120-wide projection) on a tp4 or tep4 rank. The pool is a view of a
+bigger buffer, which the test also checks the call leaves alone outside the view: a
+two-layer buffer for the tp1 cells, the stock manager's page (36 GDN layers a slot,
+the view at the last) for a rank's cells: at an engine's measured slot count, or at
+slots whose offset passes 2^31 elements.
+"""
+
+import math
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.gdn_prefill_conv_post import ( # noqa: E501
+ A_LOG_RANGE,
+ CONV_WIDTH,
+ DT_BIAS_RANGE,
+ HEAD_K_DIM,
+ HEAD_V_DIM,
+ IN_PROJ_WIDTH,
+ NUM_K_HEADS,
+ NUM_V_HEADS,
+ SOFTPLUS_THRESHOLD,
+ _engine_slots,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.gdn_prefill_conv_post import ( # noqa: E501
+ gdn_prefill_conv_post as op,
+)
+
+__extra_import_path__ = ["../../.."]
+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
+PAD = -1
+L2_NORM_EPS = 1e-6
+#: `g`'s band: two SASS
+#: approximate transcendentals (~2^-22 relative each), the softplus bounded by
+#: the threshold, every term scaled by exp(A_log).
+_APPROX_ULP = 2.0**-22
+G_ATOL = 4 * _APPROX_ULP * SOFTPLUS_THRESHOLD * math.exp(A_LOG_RANGE[1])
+
+
+def _geometry(spec):
+ """`(num_k_heads, num_v_heads, head_dim, qkv_dim, projection width)` of a cell: its
+ own, or the tp1 layer's."""
+ hk = spec.get("num_k_heads", NUM_K_HEADS)
+ hv = spec.get("num_v_heads", NUM_V_HEADS)
+ d = spec.get("head_dim", HEAD_K_DIM)
+ assert d == HEAD_K_DIM == HEAD_V_DIM
+ qkv = 2 * hk * d + hv * d
+ width = spec.get("x_row_stride", IN_PROJ_WIDTH)
+ assert width == qkv + hv * d + 2 * hv, (
+ f"x_row_stride {width} is not the [q|k|v|z|a|b] projection of {hk}/{hv} heads"
+ )
+ return hk, hv, d, qkv, width
+
+
+def _build(spec, seed):
+ """`__call__`'s positional arguments for one cell."""
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ lengths = spec["lengths"]
+ n_tok = sum(lengths)
+ hk, hv, d, qkv, width = _geometry(spec)
+ a_col = qkv + hv * d
+
+ def randn(*shape):
+ return torch.randn(*shape, generator=g, device=DEV).to(DTYPE)
+
+ # "in_proj" and "rows" both carve x, a and b out of one projection; "rows"
+ # leaves `extra_rows` trailing rows the call must not read.
+ proj = randn(n_tok + spec.get("extra_rows", 0), width)
+ x = proj[:n_tok, :qkv]
+ a = proj[:n_tok, a_col : a_col + hv]
+ b = proj[:n_tok, a_col + hv : a_col + 2 * hv]
+ assert x.stride(0) == a.stride(0) == width
+ weight = (torch.randn(qkv, CONV_WIDTH, generator=g, device=DEV) * 0.5).to(DTYPE)
+ pool_slots = spec["pool_slots"]
+ # The pool is a per-layer view of a multi-layer buffer, as a cache manager
+ # hands it out: two layers for the tp1 cells (slot stride twice the slot's own
+ # size); for a rank's cells the stock page, `page_layers` (36) GDN layers'
+ # states a slot, viewed at the last -- the strides and offset the manager's
+ # `get_conv_states` returns, measured.
+ layers = spec.get("page_layers", 2)
+ conv_states = randn(pool_slots, layers, qkv, CONV_WIDTH - 1)[:, layers - 1]
+ slots = spec.get("slots") or _engine_slots(len(lengths))
+ bounds = [0]
+ for n in lengths:
+ 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(hv, generator=g, device=DEV) * (hi_ - lo) + lo
+ lo, hi_ = DT_BIAS_RANGE
+ dt_bias = torch.rand(hv, generator=g, device=DEV) * (hi_ - lo) + lo
+ return [
+ x,
+ weight,
+ conv_states,
+ qsl,
+ si,
+ hi,
+ a,
+ b,
+ A_log,
+ dt_bias,
+ hk,
+ d,
+ hv,
+ d,
+ PAD,
+ L2_NORM_EPS,
+ SOFTPLUS_THRESHOLD,
+ spec.get("linear_gate", False),
+ ]
+
+
+def _compare(out, pool, want, what=""):
+ q, k, v, g, beta = out
+ q_r, k_r, v_r, g_r, beta_r, pool_r = want
+ for got, ref in zip(out, want[:5]):
+ assert got.dtype == ref.dtype and got.shape == ref.shape and got.is_contiguous()
+ for name, got, ref in (("q", q, q_r), ("k", k, k_r), ("v", v, v_r)):
+ d = (got.float() - ref.float()).abs()
+ flips = (d > 0).float().mean().item() if d.numel() else 0.0
+ print(
+ f"{what} {name}: max abs {d.max().item() if d.numel() else 0:.3e}, "
+ f"differing fraction {flips:.2e}"
+ )
+ torch.testing.assert_close(got, ref)
+ torch.testing.assert_close(g, g_r, rtol=8 * _APPROX_ULP, atol=G_ATOL)
+ torch.testing.assert_close(beta, beta_r)
+ assert torch.equal(pool, pool_r), "conv pool differs"
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
+ certifying(op)
+ args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
+ want = op.reference(*args)
+ inputs = [t.clone() if isinstance(t, torch.Tensor) else t for t in args]
+ pool = args[2]
+ page = pool._base
+ page_before = page.clone()
+ out = op(*args)
+ for i, (arg, before) in enumerate(zip(args, inputs)):
+ if isinstance(arg, torch.Tensor) and i != 2:
+ assert torch.equal(arg, before), f"input {i} must be read-only"
+ _compare(out, pool, want, what=f"{cell.why[:30]} vs reference")
+ # The rest of the buffer -- the other layers' sub-pages of every slot -- is untouched.
+ layer = (pool.storage_offset() - page.storage_offset()) // page.stride(1)
+ page_before[:, layer] = want[5]
+ assert torch.equal(page, page_before), "the call wrote outside its layer's view of the pool"
+
+
+def test_cells_carry_from_slots_past_the_int32_offset() -> None:
+ """The cells read a carried window at a slot whose element offset on the stock page
+ passes 2^31 at both rank geometries, and 2^32 at 16 / 48 heads -- where a 32-bit slot
+ offset faults or wraps back into the page -- so the 64-bit slot base stays exercised."""
+ reach = {}
+ for cell in op.CELLS:
+ spec = cell.spec
+ if "page_layers" not in spec:
+ continue
+ hk, _, _, qkv, _ = _geometry(spec)
+ slot_stride = spec["page_layers"] * qkv * (CONV_WIDTH - 1)
+ carried = [s for s, init in zip(spec["slots"], spec["initial"]) if init]
+ reach[hk] = max([reach.get(hk, 0)] + [s * slot_stride for s in carried])
+ assert reach[NUM_K_HEADS] >= 2**32 and reach[NUM_K_HEADS // 4] >= 2**31, reach
+
+
+def test_linear_gate_is_torch_exp_of_the_log_gate() -> None:
+ """`linear_gate=True` stores `torch.exp` of the log-space decay bitwise, and leaves
+ every other output and the conv pool as without it."""
+ spec = next(c for c in op.CELLS if c.spec.get("linear_gate")).spec
+ args = _build(spec, seed=77)
+ log_args = list(args)
+ log_args[-1] = False
+ log_args[2] = args[2].clone()
+ pool_lin = args[2]
+ log_out = op(*log_args)
+ lin_out = op(*args)
+ assert torch.equal(lin_out[3], torch.exp(log_out[3]))
+ for got, want in zip(lin_out[:3] + lin_out[4:], log_out[:3] + log_out[4:]):
+ assert torch.equal(got, want)
+ assert torch.equal(pool_lin, log_args[2])
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/vision/test_modeling_v2_patch_embed_conv3d.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/vision/test_modeling_v2_patch_embed_conv3d.py
new file mode 100644
index 000000000000..62f900fd6ed8
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/vision/test_modeling_v2_patch_embed_conv3d.py
@@ -0,0 +1,113 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the qwen4_exp default target's patch_embed_conv3d entry.
+
+The reference is a true-fp32 matmul of the flattened patch against the
+flattened filter. The op is the same contraction made by cuBLASLt, so each
+cell also holds the reference to an fp32 `F.conv3d` of the 5-D patches: that
+pins the flatten order the two share to `nn.Conv3d`'s rather than letting one
+restate the other.
+"""
+
+import contextlib
+import zlib
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.vision.patch_embed_conv3d import ( # noqa: E501
+ patch_embed_conv3d as op,
+)
+
+__extra_import_path__ = ["../../.."]
+from _validating import certifying, validating # noqa: E402,F401
+
+assert torch.cuda.is_available(), "patch_embed_conv3d requires a CUDA device"
+
+
+def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
+ """torch's default band for the dtype.
+
+ One fp32 accumulation and one rounding on both sides, so the difference is
+ accumulation order alone: nearly every output is bitwise equal and the rest
+ are one-ulp flips of the element, inside the default bf16 band. Wrong
+ variants (bias dropped, kh/kw transposed, channel-fastest rows, kt swapped)
+ land thousands of times outside it.
+ """
+ assert out.dtype == ref.dtype, (out.dtype, ref.dtype)
+ assert out.shape == ref.shape, (out.shape, ref.shape)
+ torch.testing.assert_close(out, ref)
+
+
+@contextlib.contextmanager
+def _cudnn_true_fp32():
+ prev = torch.backends.cudnn.allow_tf32
+ torch.backends.cudnn.allow_tf32 = False
+ try:
+ yield
+ finally:
+ torch.backends.cudnn.allow_tf32 = prev
+
+
+def _conv3d_fp32(pixel_values, weight, bias):
+ """`nn.Conv3d` semantics, independently: an fp32 cuDNN convolution of the
+ 5-D patches with the stride equal to the kernel."""
+ out_channels, in_channels, kt, kh, kw = weight.shape
+ with _cudnn_true_fp32():
+ out = F.conv3d(
+ pixel_values.float().view(-1, in_channels, kt, kh, kw),
+ weight.float(),
+ None if bias is None else bias.float(),
+ stride=(kt, kh, kw),
+ )
+ return out.view(-1, out_channels)
+
+
+IN_CHANNELS = 3
+KT, KH, KW = 2, 16, 16
+EMBED_DIM = 1152
+PATCH_NUMEL = IN_CHANNELS * KT * KH * KW # 1536
+
+
+def _operands(
+ num_patches,
+ dtype=torch.bfloat16,
+ seed=0,
+ in_channels=IN_CHANNELS,
+ kernel=(KT, KH, KW),
+ embed_dim=EMBED_DIM,
+ bias=True,
+):
+ """Pixels in [-1, 1] (processor mean/std 0.5), weights at the tower's
+ initializer_range 0.02."""
+ g = torch.Generator(device="cuda").manual_seed(seed)
+ kt, kh, kw = kernel
+ pixel_values = torch.empty(num_patches, in_channels * kt * kh * kw, device="cuda")
+ pixel_values = pixel_values.uniform_(-1.0, 1.0, generator=g).to(dtype)
+ weight = (
+ 0.02 * torch.randn(embed_dim, in_channels, kt, kh, kw, generator=g, device="cuda")
+ ).to(dtype)
+ bias_t = (0.02 * torch.randn(embed_dim, generator=g, device="cuda")).to(dtype) if bias else None
+ return pixel_values, weight, bias_t
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying): # noqa: F811
+ certifying(op)
+ spec = cell.spec
+ pixel_values, weight, bias = _operands(
+ spec["num_patches"], dtype=spec["dtype"], seed=zlib.crc32(cell.why.encode())
+ )
+ before = pixel_values.clone()
+ out = op(pixel_values, weight, bias)
+ assert out.shape == (pixel_values.shape[0], weight.shape[0])
+ assert out.dtype == pixel_values.dtype and out.is_contiguous()
+ assert out.data_ptr() != pixel_values.data_ptr()
+ assert torch.equal(pixel_values, before), "the pixels are read-only"
+
+ # The reference is the convolution: fp32 against an fp32 F.conv3d.
+ ref32 = op.reference(pixel_values.float(), weight.float(), bias.float())
+ torch.testing.assert_close(ref32, _conv3d_fp32(pixel_values, weight, bias))
+
+ _compare(out, op.reference(pixel_values, weight, bias))
From b9a4140724d2dcf31a562f420c5833e6f0b7c254 Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Tue, 3 Nov 2026 20:41:56 -0800
Subject: [PATCH 02/15] [None][feat] modeling_v2 qwen4_exp: one
Qwen3.8-Flash-Next NVFP4 target for tp1, tp4, tep4, dep4 and dep2
Replaces `qwen3_8_flash_next_nvfp4__sm_103__tp1` with `qwen3_8_flash_next_nvfp4__sm_103__default`,
which serves Qwen3.8-Flash-Next NVFP4 on GB300 (sm_103) at every layout the family routes: one GPU
(tp1); four in one node as tp4 (attention TP4, routed experts split by intermediate), tep4
(attention TP4, experts EP4) or dep4 (attention DP, experts EP4); and two as dep2 (attention DP,
experts EP2), the layout of disaggregated engines.
Routing: `layout()` names the member a mapping describes, tp1 included, and the engine checks are
per member. At tp1 the CuteDSL experts (`moe_config.backend: CUTEDSL`, up to three drafts) and an
fp32 recurrent state are served too. MTP needs the GDN cached replay: with
TRTLLM_USE_GDN_REPLAY=0 it is declined, so `auto` falls back to the built-in model.
At tp1 the target takes one-rank forms, fixed at construction: the decode projections' one-rank
dispatch; the hyper-connection sites' row-count forms (the decode GEMV and tcgen05 rows pairs, the
prefill rstd sites); a trtllm-gen decode MoE whose finalize the next site folds in; the CuteDSL
copy of the routed experts for long prefill steps; the fp32 state's one-launch decode recurrence;
the decode draft layer and its FP8 head; and the NVFP4 pool's context rows attended a group of
requests at a time. The other layouts keep one form at every row count, so a request's tokens do
not depend on what shares its steps. Opt-in at tp1, off by default: FP8 dense projections
(TRTLLM_MODELING_V2_DENSE_FP8), an NVFP4 copy of the GDN input projection
(TRTLLM_MODELING_V2_GDN_INPROJ_NVFP4), NVFP4 target and draft heads
(TRTLLM_MODELING_V2_TARGET_HEAD_NVFP4, TRTLLM_MODELING_V2_MTP_HEAD_NVFP4), NVFP4 draft experts
(TRTLLM_MODELING_V2_MTP_MOE_NVFP4) and FP8 draft dense projections
(TRTLLM_MODELING_V2_MTP_DENSE_FP8).
The target runs one-model MTP (MTP-Eagle, up to seven drafts) with IndexShare, MTP under attention
DP (dep4), FP8 and NVFP4 KV cache, block reuse, chunked prefill, the breakable prefill CUDA graph
under attention TP, disaggregated serving and the vision tower. Its own decode and prefill kernels
are in its `ops/`.
Tests: the target's unit tests in its directory (its cpu_only tests also in l0_cpu); the catalog
entries' cells keyed by the target; the routing tests. Accuracy is one class,
TestModelingV2Qwen38FlashNextNvfp4Sm103Default:
- tp1 (`test_tp1_*`): pre-merge MTP3 with IndexShare, gsm8k and then pass keys on the sparse path;
post-merge MTP1 with MMMU and video, CuteDSL experts over FP8 KV, MTP7 over NVFP4 KV, and the
AgentX engines on one context and one generation GPU; QA no speculation and the breakable
prefill graph against eager.
- four and two GPUs: pre-merge tep4 MTP3 with IndexShare and pass keys; post-merge tp4 over NVFP4
KV with the breakable prefill graph, dep4 over FP8 KV, dep2 disaggregated; QA MMMU at tep4.
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
.../__init__.py | 2 +-
.../modeling.py | 6715 ++++++++++++++
.../ops}/__init__.py | 2 +-
.../ops/attention}/__init__.py | 0
.../qsa_fp8_pages_to_bf16__prefill.py | 224 +
.../qsa_rowtile_attention_cute__prefill.py} | 3 +
.../qsa_sparse_gqa_rowtile__prefill.py} | 148 +-
.../qwen4_exp_attn_decode_prologue__decode.py | 719 ++
.../ops/attention/vision_rope_2d_packed_qk.py | 255 +
.../ops/gemm}/__init__.py | 0
.../ops/hc}/__init__.py | 0
.../qwen4_exp_hc_decode_tc_first__decode.py} | 106 +-
.../hc/qwen4_exp_hc_up_gate_mean_fp4__moe.py} | 81 +-
.../ops/moe}/__init__.py | 0
.../ops/ssm}/__init__.py | 0
.../ssm/gdn_cached_replay_cute__verify.py} | 1277 +--
.../ops/vision}/__init__.py | 0
.../vision/triton_pos_embed_interpolate.py | 210 +
.../pool.py | 488 +
.../vision.py | 694 ++
.../weights.py | 1861 ++++
...contiguous_grouped_gemm_finalize_fusion.py | 2738 ------
.../modeling.py | 8119 -----------------
.../ops/__init__.py | 3 -
.../activation/fused_sigmoid_mul_inplace.py | 100 -
.../ops/activation/gelu_erf.py | 91 -
.../ops/activation/gelu_tanh.py | 93 -
.../ops/activation/silu_and_mul_rows.py | 93 -
.../ops/attention/qsa_canonical_selection.py | 214 -
.../ops/attention/qsa_expand_block_indices.py | 178 -
.../ops/attention/qsa_mtp_reuse_topk.py | 246 -
.../ops/attention/qsa_nvfp4_pages_to_bf16.py | 390 -
.../ops/attention/qsa_nvfp4_verify_ptx.py | 248 -
.../ops/attention/qsa_paged_index_scores.py | 245 -
.../ops/attention/qsa_paged_index_store.py | 400 -
.../ops/attention/qsa_paged_kv_store.py | 219 -
.../ops/attention/qsa_paged_kv_store_nvfp4.py | 533 --
.../ops/attention/qsa_paged_sparse_gqa.py | 283 -
.../attention/qsa_paged_sparse_gqa_decode.py | 819 --
.../attention/qsa_paged_sparse_gqa_nvfp4.py | 1106 ---
.../qsa_paged_sparse_gqa_nvfp4_decode.py | 1547 ----
.../qsa_paged_sparse_gqa_nvfp4_verify.py | 2147 -----
.../ops/attention/qsa_prefill_compress.py | 276 -
.../ops/attention/qsa_rollback_planes.py | 144 -
.../ops/attention/qsa_topk_expand_decode.py | 651 --
.../ops/attention/qsa_topk_select.py | 213 -
.../qwen4_exp_attn_decode_prologue.py | 1009 --
.../ops/attention/qwen4_exp_qk_norm_rope.py | 1193 ---
.../ops/attention/spec_rollback_commit.py | 368 -
.../ops/attention/verify_step_setup.py | 215 -
.../ops/attention/vision_rope_2d.py | 262 -
.../ops/embedding/ngram_hash.py | 252 -
.../ops/gemm/bf16_gemv_pdl.py | 384 -
.../ops/gemm/fp8_gemv.py | 2016 ----
.../ops/gemm/fp8_per_tensor_mm.py | 192 -
.../ops/gemm/low_m_bf16_mm.py | 231 -
.../ops/gemm/mid_rows_mm.py | 1380 ---
.../ops/gemm/nvfp4_gemv.py | 992 --
.../ops/hc/qwen4_exp_gated_residual.py | 179 -
.../ops/hc/qwen4_exp_hc_decode_first_cute.py | 605 --
.../hc/qwen4_exp_hc_decode_mix_fp4_cute.py | 534 --
.../ops/hc/qwen4_exp_hc_decode_rows_tc.py | 1338 ---
.../ops/hc/qwen4_exp_hc_first_rescale_tc.py | 969 --
.../ops/hc/qwen4_exp_hc_inject_norm.py | 1058 ---
.../ops/hc/qwen4_exp_hc_inject_rstd.py | 258 -
.../ops/hc/qwen4_exp_hc_lowrank_act.py | 150 -
.../ops/hc/qwen4_exp_hc_up_gate_mean.py | 899 --
.../ops/hc/qwen4_exp_hc_up_gate_mean_rstd.py | 1119 ---
.../ops/hc/qwen4_exp_post_inject.py | 138 -
.../ops/moe/__init__.py | 2 -
...sl_nvfp4_gather_grouped_gemm_act_fusion.py | 484 -
.../ops/moe/cute_dsl_nvfp4_grouped_gemm.py | 367 -
.../fp4_block_scale_moe_renorm_permuted.py | 139 -
.../moe/fp4_block_scale_moe_renorm_routed.py | 388 -
.../ops/moe/fp4_moe_renorm_cute_dsl.py | 309 -
.../ops/moe/fp8_block_scale_moe_runner.py | 365 -
.../ops/moe/fused_sigmoid_gate_mul_add.py | 132 -
.../ops/moe/moe_sort.py | 272 -
.../ops/moe/moe_unpermute_inplace.py | 246 -
.../ops/moe/renorm_moe_routing_op.py | 144 -
.../ops/norm/__init__.py | 2 -
.../ops/norm/flashinfer_gemma_rmsnorm.py | 97 -
.../ops/norm/layer_norm.py | 116 -
.../ops/norm/rms_norm_gated_token_major.py | 199 -
.../ops/ple/__init__.py | 2 -
.../ops/ple/gated_ngram_injection.py | 187 -
.../ops/ple/ngram_token_history.py | 290 -
.../ops/ple/ple_gate_conv_inject_norm.py | 683 --
.../ops/ple/ple_ngram_embed.py | 560 --
.../ops/quantization/__init__.py | 2 -
.../ops/quantization/bf16_to_fp32.py | 80 -
.../dequantize_e4m3_per_tensor.py | 94 -
.../ops/quantization/fp8_quantize_1x128.py | 131 -
.../ops/ssm/__init__.py | 2 -
.../ops/ssm/causal_conv1d_fwd.py | 250 -
.../ops/ssm/causal_conv1d_triton_update.py | 283 -
.../ops/ssm/causal_conv1d_update.py | 481 -
.../ops/ssm/chunk_gated_delta_rule.py | 268 -
.../ops/ssm/cu_seqlens_from_seq_lens.py | 131 -
.../ops/ssm/dilated_short_conv_paged.py | 292 -
.../ssm/extract_transpose_prefill_slice.py | 129 -
.../ssm/flashinfer_chunk_gated_delta_rule.py | 414 -
.../ops/ssm/fused_gdn_gating.py | 153 -
.../ops/ssm/fused_gdn_post_conv.py | 232 -
...t_gated_delta_rule_cached_replay_update.py | 512 --
...fused_recurrent_gated_delta_rule_update.py | 298 -
.../fused_sigmoid_gating_delta_rule_update.py | 297 -
.../ops/ssm/gdn_decode_gated_norm_pdl.py | 509 --
.../ops/ssm/gdn_decode_rows_gated_norm.py | 1035 ---
.../ops/ssm/gdn_in_proj_conv1d_pdl.py | 432 -
.../ops/ssm/gdn_prefill_conv_post.py | 1074 ---
.../ops/ssm/pack_gdn_decode_qkv.py | 135 -
.../ops/ssm/reset_recurrent_state_rows.py | 163 -
.../ops/vision/__init__.py | 2 -
.../ops/vision/patch_embed_conv3d.py | 141 -
.../ops/vision/pos_embed_interpolate.py | 243 -
.../weights.py | 1438 ---
.../modeling_v2/models/qwen4_exp/routing.py | 186 +-
tests/integration/defs/.test_durations | 26 +-
.../accuracy/modeling_v2/test_qwen4_exp.py | 694 +-
.../references/acceptance_length.yaml | 49 +-
.../test_lists/qa/llm_function_core.txt | 16 +
.../test_lists/test-db/l0_b300.yml | 30 +-
.../integration/test_lists/test-db/l0_cpu.yml | 17 +-
.../test_lists/test-db/l0_gb300.yml | 4 +
.../test-db/l0_gb300_multi_gpus.yml | 11 +
.../test_modeling_v2_fused_qk_norm_rope.py | 2 +-
.../test_modeling_v2_thop_attention.py | 2 +-
.../test_modeling_v2_qsa_paged_index_store.py | 172 +-
...est_modeling_v2_fp4_moe_renorm_cute_dsl.py | 41 +-
.../__init__.py | 7 +
...ling_v2_qsa_fp8_pages_to_bf16__prefill.py} | 128 +-
...ing_v2_qsa_sparse_gqa_rowtile__prefill.py} | 144 +-
..._qwen4_exp_attn_decode_prologue__decode.py | 261 +
...st_modeling_v2_vision_rope_2d_packed_qk.py | 174 +
...2_qwen4_exp_hc_decode_tc_first__decode.py} | 32 +-
..._v2_qwen4_exp_hc_up_gate_mean_fp4__moe.py} | 86 +-
...eling_v2_gdn_cached_replay_cute__verify.py | 276 +
.../test_attention_plumbing.py | 386 +
.../test_core_layout.py | 769 ++
.../test_decode_paths.py | 295 +
.../test_modeling_v2_spec_rollback.py | 180 +-
.../test_mtp_draft_hc.py | 382 +
.../test_pool_contract.py} | 231 +-
.../test_pool_layout.py | 345 +
.../test_routed_experts.py | 671 ++
.../test_routing.py | 396 +
.../test_select_canonical.py | 69 +
.../test_thop_attention_calls.py} | 40 +-
.../test_vision.py | 568 ++
.../test_weight_shards.py | 799 ++
...odeling_v2_triton_pos_embed_interpolate.py | 116 +
...t_modeling_v2_fused_sigmoid_mul_inplace.py | 72 -
.../activation/test_modeling_v2_gelu_erf.py | 141 -
.../activation/test_modeling_v2_gelu_tanh.py | 133 -
.../test_modeling_v2_silu_and_mul_rows.py | 66 -
...est_modeling_v2_qsa_canonical_selection.py | 167 -
...st_modeling_v2_qsa_expand_block_indices.py | 156 -
.../test_modeling_v2_qsa_mtp_reuse_topk.py | 248 -
...test_modeling_v2_qsa_paged_index_scores.py | 284 -
.../test_modeling_v2_qsa_paged_kv_store.py | 190 -
...st_modeling_v2_qsa_paged_kv_store_nvfp4.py | 369 -
.../test_modeling_v2_qsa_paged_sparse_gqa.py | 250 -
...modeling_v2_qsa_paged_sparse_gqa_decode.py | 118 -
..._modeling_v2_qsa_paged_sparse_gqa_nvfp4.py | 321 -
...ng_v2_qsa_paged_sparse_gqa_nvfp4_decode.py | 209 -
...ng_v2_qsa_paged_sparse_gqa_nvfp4_verify.py | 316 -
.../test_modeling_v2_qsa_prefill_compress.py | 237 -
.../test_modeling_v2_qsa_rollback_planes.py | 84 -
...test_modeling_v2_qsa_topk_expand_decode.py | 108 -
.../test_modeling_v2_qsa_topk_select.py | 147 -
...eling_v2_qwen4_exp_attn_decode_prologue.py | 219 -
...test_modeling_v2_qwen4_exp_qk_norm_rope.py | 162 -
.../test_modeling_v2_spec_rollback_commit.py | 101 -
.../test_modeling_v2_verify_step_setup.py | 139 -
.../test_modeling_v2_vision_rope_2d.py | 120 -
.../embedding/test_modeling_v2_ngram_hash.py | 132 -
.../gemm/test_modeling_v2_bf16_gemv_pdl.py | 98 -
.../gemm/test_modeling_v2_fp8_gemv.py | 221 -
.../test_modeling_v2_fp8_per_tensor_mm.py | 83 -
.../gemm/test_modeling_v2_low_m_bf16_mm.py | 148 -
.../gemm/test_modeling_v2_mid_rows_mm.py | 119 -
.../gemm/test_modeling_v2_nvfp4_gemv.py | 326 -
...st_modeling_v2_qwen4_exp_gated_residual.py | 156 -
...eling_v2_qwen4_exp_hc_decode_first_cute.py | 196 -
...ing_v2_qwen4_exp_hc_decode_mix_fp4_cute.py | 162 -
...modeling_v2_qwen4_exp_hc_decode_rows_tc.py | 173 -
...deling_v2_qwen4_exp_hc_first_rescale_tc.py | 122 -
...st_modeling_v2_qwen4_exp_hc_inject_norm.py | 202 -
...st_modeling_v2_qwen4_exp_hc_inject_rstd.py | 165 -
...st_modeling_v2_qwen4_exp_hc_lowrank_act.py | 79 -
...t_modeling_v2_qwen4_exp_hc_up_gate_mean.py | 78 -
...eling_v2_qwen4_exp_hc_up_gate_mean_rstd.py | 175 -
.../test_modeling_v2_qwen4_exp_post_inject.py | 106 -
...sl_nvfp4_gather_grouped_gemm_act_fusion.py | 399 -
...modeling_v2_cute_dsl_nvfp4_grouped_gemm.py | 375 -
.../test_modeling_v2_cute_dsl_weight_path.py | 371 -
..._v2_fp4_block_scale_moe_renorm_permuted.py | 117 -
...ng_v2_fp4_block_scale_moe_renorm_routed.py | 155 -
..._modeling_v2_fp8_block_scale_moe_runner.py | 195 -
..._modeling_v2_fused_sigmoid_gate_mul_add.py | 83 -
.../moe/test_modeling_v2_moe_sort.py | 215 -
.../test_modeling_v2_moe_unpermute_inplace.py | 268 -
.../moe/test_modeling_v2_mtp_moe_nvfp4.py | 142 -
.../test_modeling_v2_renorm_moe_routing_op.py | 125 -
...st_modeling_v2_flashinfer_gemma_rmsnorm.py | 70 -
.../norm/test_modeling_v2_layer_norm.py | 71 -
...v2_qwen4_exp_rms_norm_gated_token_major.py | 80 -
.../test_modeling_v2_gated_ngram_injection.py | 164 -
.../test_modeling_v2_ngram_token_history.py | 212 -
...t_modeling_v2_ple_gate_conv_inject_norm.py | 212 -
.../ple/test_modeling_v2_ple_ngram_embed.py | 230 -
.../test_modeling_v2_bf16_to_fp32.py | 67 -
..._modeling_v2_dequantize_e4m3_per_tensor.py | 140 -
.../test_modeling_v2_fp8_quantize_1x128.py | 161 -
...modeling_v2_causal_conv1d_triton_update.py | 116 -
.../test_modeling_v2_causal_conv1d_update.py | 111 -
...test_modeling_v2_chunk_gated_delta_rule.py | 242 -
...st_modeling_v2_cu_seqlens_from_seq_lens.py | 105 -
...st_modeling_v2_dilated_short_conv_paged.py | 147 -
...ling_v2_extract_transpose_prefill_slice.py | 65 -
...ng_v2_flashinfer_chunk_gated_delta_rule.py | 143 -
.../ssm/test_modeling_v2_fused_gdn_gating.py | 115 -
.../test_modeling_v2_fused_gdn_post_conv.py | 143 -
...t_gated_delta_rule_cached_replay_update.py | 547 --
...fused_recurrent_gated_delta_rule_update.py | 250 -
..._fused_sigmoid_gating_delta_rule_update.py | 231 -
.../ssm/test_modeling_v2_gdn_block.py | 416 -
...test_modeling_v2_gdn_cached_replay_cute.py | 364 -
...t_modeling_v2_gdn_decode_gated_norm_pdl.py | 122 -
..._modeling_v2_gdn_decode_rows_gated_norm.py | 252 -
...test_modeling_v2_gdn_in_proj_conv1d_pdl.py | 103 -
.../test_modeling_v2_gdn_prefill_conv_post.py | 156 -
...test_modeling_v2_gdn_replay_verify_cute.py | 128 -
.../test_modeling_v2_gdn_verify_fused_conv.py | 100 -
.../test_modeling_v2_pack_gdn_decode_qkv.py | 117 -
...modeling_v2_qwen4_exp_causal_conv1d_fwd.py | 131 -
..._modeling_v2_reset_recurrent_state_rows.py | 117 -
.../test_modeling_v2_gdn_replay_verify.py | 357 -
.../test_modeling_v2_index_share.py | 168 -
.../test_modeling_v2_kv_range_probe.py | 82 -
.../test_modeling_v2_logits_widening.py | 100 -
.../test_modeling_v2_nvfp4_context_groups.py | 136 -
.../test_modeling_v2_patch_embed_conv3d.py | 118 -
.../test_modeling_v2_pos_embed_interpolate.py | 110 -
.../test_modeling_v2_vision_tower_calls.py | 165 -
.../test_modeling_v2_qwen4_exp_kv_quant.py | 229 -
.../test_modeling_v2_qwen4_exp_moe_scalars.py | 74 -
...modeling_v2_qwen4_exp_multimodal_embeds.py | 274 -
..._modeling_v2_qwen4_exp_validate_request.py | 209 -
...test_modeling_v2_qwen4_exp_weight_table.py | 401 -
..._modeling_v2_fp4_block_scale_moe_runner.py | 2 +-
.../test_modeling_v2_fp4_quantize.py | 2 +-
.../test_modeling_v2_reference_precision.py | 9 +-
.../modeling_v2/test_modeling_v2_routing.py | 24 +-
255 files changed, 18631 insertions(+), 70432 deletions(-)
rename tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/{qwen3_8_flash_next_nvfp4__sm_103__tp1 => qwen3_8_flash_next_nvfp4__sm_103__default}/__init__.py (64%)
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py
rename tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/{qwen3_8_flash_next_nvfp4__sm_103__tp1/kernels => qwen3_8_flash_next_nvfp4__sm_103__default/ops}/__init__.py (60%)
rename tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/{qwen3_8_flash_next_nvfp4__sm_103__tp1/kernels/moe => qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention}/__init__.py (100%)
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_fp8_pages_to_bf16__prefill.py
rename tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/{qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rowtile_attention_cute.py => qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_rowtile_attention_cute__prefill.py} (99%)
rename tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/{qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py => qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_sparse_gqa_rowtile__prefill.py} (86%)
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qwen4_exp_attn_decode_prologue__decode.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/vision_rope_2d_packed_qk.py
rename tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/{qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation => qwen3_8_flash_next_nvfp4__sm_103__default/ops/gemm}/__init__.py (100%)
rename tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/{qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention => qwen3_8_flash_next_nvfp4__sm_103__default/ops/hc}/__init__.py (100%)
rename 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 => qwen3_8_flash_next_nvfp4__sm_103__default/ops/hc/qwen4_exp_hc_decode_tc_first__decode.py} (89%)
rename 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 => qwen3_8_flash_next_nvfp4__sm_103__default/ops/hc/qwen4_exp_hc_up_gate_mean_fp4__moe.py} (55%)
rename tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/{qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding => qwen3_8_flash_next_nvfp4__sm_103__default/ops/moe}/__init__.py (100%)
rename tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/{qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm => qwen3_8_flash_next_nvfp4__sm_103__default/ops/ssm}/__init__.py (100%)
rename tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/{qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_cached_replay_cute.py => qwen3_8_flash_next_nvfp4__sm_103__default/ops/ssm/gdn_cached_replay_cute__verify.py} (63%)
rename tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/{qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc => qwen3_8_flash_next_nvfp4__sm_103__default/ops/vision}/__init__.py (100%)
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/vision/triton_pos_embed_interpolate.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/pool.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision.py
create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/weights.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/kernels/moe/blockscaled_contiguous_grouped_gemm_finalize_fusion.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/__init__.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_erf.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py
delete 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
delete 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
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_expand_block_indices.py
delete 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
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_nvfp4_pages_to_bf16.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_nvfp4_verify_ptx.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py
delete 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.py
delete 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
delete 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.py
delete 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
delete 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
delete 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_decode.py
delete 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_verify.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_prefill_compress.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rollback_planes.py
delete 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
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py
delete 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
delete 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
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/spec_rollback_commit.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/verify_step_setup.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/vision_rope_2d.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/ngram_hash.py
delete 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
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/fp8_gemv.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/fp8_per_tensor_mm.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/low_m_bf16_mm.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/mid_rows_mm.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/nvfp4_gemv.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_gated_residual.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_cute.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_cute.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_rows_tc.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_first_rescale_tc.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_inject_norm.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_inject_rstd.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_lowrank_act.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_up_gate_mean.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_up_gate_mean_rstd.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_post_inject.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/__init__.py
delete 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
delete 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
delete 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
delete 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
delete 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
delete mode 100644 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
delete mode 100644 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
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_sort.py
delete 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
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/renorm_moe_routing_op.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/__init__.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/flashinfer_gemma_rmsnorm.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py
delete mode 100644 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
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/__init__.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/gated_ngram_injection.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ngram_token_history.py
delete 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
delete 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
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/__init__.py
delete 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
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/dequantize_e4m3_per_tensor.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/fp8_quantize_1x128.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/__init__.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_fwd.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/chunk_gated_delta_rule.py
delete mode 100644 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
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/dilated_short_conv_paged.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/extract_transpose_prefill_slice.py
delete mode 100644 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
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_gating.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_post_conv.py
delete 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
delete 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_update.py
delete mode 100644 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
delete 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
delete 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
delete 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
delete 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
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/pack_gdn_decode_qkv.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/reset_recurrent_state_rows.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/__init__.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/patch_embed_conv3d.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/pos_embed_interpolate.py
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py
rename 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 (52%)
rename 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 (72%)
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/__init__.py
rename tests/unittest/_torch/modeling_v2/models/qwen4_exp/{qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_nvfp4_pages_to_bf16.py => qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_qsa_fp8_pages_to_bf16__prefill.py} (52%)
rename 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 => qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_qsa_sparse_gqa_rowtile__prefill.py} (58%)
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue__decode.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_vision_rope_2d_packed_qk.py
rename 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 => qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first__decode.py} (70%)
rename 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 => qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4__moe.py} (51%)
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ssm/test_modeling_v2_gdn_cached_replay_cute__verify.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_attention_plumbing.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_core_layout.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_decode_paths.py
rename tests/unittest/_torch/modeling_v2/models/qwen4_exp/{qwen3_8_flash_next_nvfp4__sm_103__tp1 => qwen3_8_flash_next_nvfp4__sm_103__default}/test_modeling_v2_spec_rollback.py (72%)
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_mtp_draft_hc.py
rename tests/unittest/_torch/modeling_v2/models/qwen4_exp/{test_modeling_v2_qwen4_exp_pool_contract.py => qwen3_8_flash_next_nvfp4__sm_103__default/test_pool_contract.py} (55%)
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_pool_layout.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routed_experts.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_select_canonical.py
rename tests/unittest/_torch/modeling_v2/models/qwen4_exp/{test_modeling_v2_qwen4_exp_thop_attention_calls.py => qwen3_8_flash_next_nvfp4__sm_103__default/test_thop_attention_calls.py} (75%)
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_weight_shards.py
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision/test_modeling_v2_triton_pos_embed_interpolate.py
delete mode 100644 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
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_erf.py
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_tanh.py
delete 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
delete 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
delete mode 100644 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
delete 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
delete mode 100644 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
delete 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.py
delete 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
delete 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.py
delete 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
delete 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
delete 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_decode.py
delete 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_verify.py
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_prefill_compress.py
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_rollback_planes.py
delete 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
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_select.py
delete 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
delete 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
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_spec_rollback_commit.py
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_verify_step_setup.py
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_vision_rope_2d.py
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/embedding/test_modeling_v2_ngram_hash.py
delete 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
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_fp8_gemv.py
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_fp8_per_tensor_mm.py
delete mode 100644 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
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_mid_rows_mm.py
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_nvfp4_gemv.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_gated_residual.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_cute.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_cute.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_rows_tc.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_first_rescale_tc.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_inject_norm.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_inject_rstd.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_lowrank_act.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_up_gate_mean.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_up_gate_mean_rstd.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_post_inject.py
delete 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
delete 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
delete 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
delete 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
delete 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
delete mode 100644 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
delete mode 100644 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
delete 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
delete 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
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_mtp_moe_nvfp4.py
delete mode 100644 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
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_layer_norm.py
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_gated_ngram_injection.py
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ngram_token_history.py
delete 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
delete 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
delete 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
delete mode 100644 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
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_fp8_quantize_1x128.py
delete mode 100644 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
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_update.py
delete mode 100644 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
delete mode 100644 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
delete mode 100644 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
delete mode 100644 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
delete mode 100644 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
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_gating.py
delete mode 100644 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
delete 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
delete 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_update.py
delete mode 100644 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
delete 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
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_cached_replay_cute.py
delete 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
delete 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
delete 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
delete 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
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_replay_verify_cute.py
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_verify_fused_conv.py
delete mode 100644 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
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_qwen4_exp_causal_conv1d_fwd.py
delete mode 100644 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
delete 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
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py
delete 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
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_logits_widening.py
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_nvfp4_context_groups.py
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_patch_embed_conv3d.py
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_pos_embed_interpolate.py
delete 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
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_kv_quant.py
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_moe_scalars.py
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_multimodal_embeds.py
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_validate_request.py
delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_weight_table.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/__init__.py
similarity index 64%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/__init__.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/__init__.py
index 8fd2a7f0c45b..eba053a6d445 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/__init__.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/__init__.py
@@ -1,3 +1,3 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
-"""Qwen3.8-Flash-Next-NVFP4 / sm_103 / tp1."""
+"""Qwen3.8-Flash-Next-NVFP4 / sm_103 / default: tp4, tep4, dep4 and dep2."""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py
new file mode 100644
index 000000000000..bb040f05c87b
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py
@@ -0,0 +1,6715 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""ModelingV2 target: Qwen3.8-Flash-Next-NVFP4 (hf-fc694b5_orig) / sm_103 / default.
+
+One target for the layouts `routing.py` admits in one node: one rank (tp1, every head and
+expert); at four ranks tp4 (attention TP4, routed experts split by intermediate), tep4
+(attention TP4, experts EP4) and dep4 (attention DP, experts EP4); at two ranks dep2 (attention
+DP, experts EP2), the layout of both disaggregated-serving engines, which is dep4 at half the
+group. The core holds the static
+structure (this rank's share of the weights, derived operands, per-manager pool views);
+`PrefillTarget` (any batch with context rows) and `DecodeTarget` (generation only) run a step
+over it, passing every argument to every op at the call site.
+
+The forward is cut per rank; what the layout settles is settled once, at construction:
+
+* **Attention and Gated-DeltaNet heads** are this rank's share (`weights.rank_dims`,
+ `pool.pool_geometry`): at attention TP4 six q heads, one of the two KV heads, four GDN key
+ and twelve value heads. The fused attention input projection row is
+ `[q | k | v | gate | index q, k | v slot]`, 4352 wide; the QSA indexer is replicated.
+* **The mixer output is a partial sum** under attention TP: o_proj and out_proj read this
+ rank's heads. One plain sum (`comm/allreduce`, op `NONE`) completes it before the
+ hyper-connection inject, which every rank runs on the same reduced rows. Under attention DP
+ a rank holds every head and serves its own requests, so there is nothing to sum.
+* **Routed experts** run on the trtllm-gen NVFP4 runner entered with the router logits
+ (`fp4_block_scale_moe_renorm_routed`: its own Renormalize routing over the whole routing
+ space). A rank holds an EP window of
+ `num_experts / moe_ep_size` whole experts (tep4, dep4), or every expert at its quarter of the
+ intermediate (tp4: 160 of 640 rows, stored zero-padded to 256, `weights.EXPERT_INTER_ALIGN`,
+ with zero codes and scales, so the padding adds nothing). Under attention TP every rank holds
+ every row and the partial output -- the window's, or the intermediate slice's -- is summed the
+ same way. Under attention DP each rank holds its own rows: they are padded to the group's
+ largest row count, gathered (`comm/allgather`), routed and run through the window, and
+ reduce-scattered back (`comm/reducescatter`), each row's window partials summed once
+ (`_expert_sum_op`). The router and the bf16 shared expert are replicated and run on this
+ rank's rows, so the block output `routed + sigmoid(gate) * shared` is formed by the next
+ hyper-connection site on the reduced routed rows.
+* **Hyper-connections, PLE, the embedding** run replicated on every rank's rows;
+ `lm_head` is the shell's (vocab-parallel under attention TP, gathered by the shell's head;
+ whole under attention DP, over this rank's own rows).
+
+Each hyper-connection site is `qwen4_exp_hc_inject_norm` (the previous block's injection
+fused with this site's grouped norm; after a MoE block it also forms that block's output)
+and then the mixer: the first projection with its activations (one tcgen05 launch,
+`qwen4_exp_hc_decode_tc_first`, up to `_HC_TC_FIRST_MAX_ROWS` rows; past them `cublas_mm`,
+fp32 out, then `qwen4_exp_hc_lowrank_act`) and `qwen4_exp_hc_up_gate_mean`. The MoE site's
+routed input is that mixer output quantized to NVFP4 with linear block scales at the layer's
+FC1 input global scale: in the mixer's own epilogue under attention TP
+(`qwen4_exp_hc_up_gate_mean_fp4`), on the gathered rows under attention DP (`fp4_quantize`).
+There is **no** `model.language_model.norm.weight`: the stack-final mixer is the final norm.
+
+**One rank (tp1).** With nothing to sum or exchange, a rank holding the whole model takes forms
+of its own, settled at construction (`one_rank`): the decode projections' one-rank dispatch
+(`_one_rank_projection`); the hyper-connection sites' row-count forms (`hc_fast_sites`: the
+decode GEMV and tcgen05 rows pairs, the prefill rstd sites), which the other layouts leave
+out so that a request's tokens do not depend on what its steps carry beside it; a trtllm-gen
+decode MoE whose finalize the next site folds in (`moe_fold`); the CuteDSL copy of the
+routed experts for prefill steps past `_CUTE_DSL_MOE_MIN_ROWS` (`moe_prefill_copy`), or the
+CuteDSL pipeline for every step (`moe_config.backend: CUTEDSL`, `_cute_dsl_routed`); an fp32
+recurrent state, whose captured single-token decode runs `gdn_decode_gated_norm_pdl`; the
+decode draft layer and its FP8 head; the NVFP4 pool's context rows attended a group of
+requests at a time (`_nvfp4_context_plan`); and the opt-in FP8 / NVFP4 copies of the dense
+projections and heads (`_DENSE_FP8_ENV`, `_GDN_INPROJ_NVFP4_ENV`, `_TARGET_HEAD_NVFP4_ENV`,
+`_MTP_HEAD_NVFP4_ENV`, `_MTP_MOE_NVFP4_ENV`).
+
+**QSA sparse attention.** The index plane is maintained on every step so a request can
+cross `dense_seq_len_threshold` later without rebuilding its prefix; above it the sparse
+core (index scores, block top-k and expansion, sparse GQA) replaces the dense op, and the
+main K/V append becomes this layer's own call. The raw index key is stored before anything
+normalizes it: the compression applies the norm and the rotation itself.
+
+**KV storage.** The paged K/V pool holds the checkpoint's bf16, fp8-e4m3 under
+`kv_cache_config.dtype: fp8` (`QM_FP8_KV_CACHE`) or NVFP4 under `nvfp4` (`QM_NVFP4_KV_CACHE`,
+e2m1 pairs with e4m3 block scales in a second pool, second-level scales 1.0). The NVFP4
+attention kernels exist for the paged-context path only. The QSA index and position planes
+keep their own types under all three.
+
+**Breakable prefill CUDA graph.** With `prefill_cuda_graph_backend: breakable` the engine
+captures the core's `forward` (the body, not the shell) for a step with context rows, pads
+the step up to the nearest bucket and replays it without calling the body. The shell
+therefore calls `begin_step` before the body on every step: the asserts, the cache-manager
+bind, the target choice and the prefill target's host bookkeeping all live there. The
+prefill forward is row-wise over the `T` rows it is handed, with four breaks
+(`catalog/graph/eager_on_graph`) that read the step: the vision splice, the PLE block, and
+each layer's attention core and Gated-DeltaNet recurrence. Each writes rows `[:n]` of a fixed
+output and zeroes rows `[n:]`.
+
+**Host offload.** By default the PLE n-gram table lives in host memory the device addresses
+(`weights._host_resident_empty`), one copy per rank; `TRTLLM_MODELING_V2_PLE_HOST_OFFLOAD=0`
+makes it device-resident. Output is identical.
+
+**One-model MTP-Eagle.** With a `speculative_config` the core also holds the checkpoint's MTP
+module (`MTPLayer`), replayed `max_draft_len` times by the engine's spec worker, and the trunk
+runs every generation request's `max_draft_len + 1` verify rows:
+
+* The draft layer is one more full-attention QSA layer, cut like a trunk attention layer (this
+ rank's q heads, its KV head, the replicated indexer; every head under attention DP), over the
+ manager the draft step is handed: under attention TP that is a separate draft KV cache, whose
+ K/V, index and position pages hold the draft layer alone; under attention DP the engine
+ builds none and the trunk's manager maps the draft layer too (`_bind_draft_cache_manager`).
+ Its front end, router, shared expert and hyper-connections are replicated; its FP8
+ block-scale experts run as this rank's window of `num_experts / tp_size` experts
+ (`fp8_block_scale_moe_runner`'s window). Under attention TP the o_proj and expert partials
+ are completed by the same plain sum as the trunk's. That window is the same at tp4 as at
+ tep4: the FP8 experts' 128x128 weight blocks do not take the trunk's 160-row intermediate
+ split, and a sum of whole-expert windows is the layer's output whichever way the trunk's
+ experts are cut.
+* Under attention DP each rank drafts for its own requests, and the draft MoE runs the trunk's
+ round trip (`_expert_sum_op`): the group's rows gathered, routed and run through this rank's
+ window, reduce-scattered back. The spec worker hands each draft step the group's row counts
+ (the trunk step's tokens on the first, one row a request on the rest, attention-DP dummy
+ requests included) without writing them into the metadata; `_dp_pad_rows` pads this rank's
+ rows by them, so every rank issues the same collectives at every draft step.
+* A verify step's Gated-DeltaNet rows take the multi-token conv with per-step window snapshots
+ and the stock Triton cached replay over the bf16 pool, which commits a full history in the
+ kernel below `CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE` requests and leaves it to the
+ manager's all-layer commit from there on, exactly as the stock layer hands it work items.
+* What a verify row writes that the runtime does not roll back -- the PLE n-gram history and
+ conv state, every attention layer's QSA index-K plane and the shared position plane -- is
+ staged during the forward and rewritten to the accepted prefix by `SpeculativeRollback`,
+ which the spec worker calls after its draft loop.
+* IndexShare (`index_share_for_mtp_iteration`) runs the draft indexer once per draft loop:
+ the draft-extend step captures each request's selection and the later draft steps rebuild
+ theirs from it (`qsa_mtp_reuse_topk`).
+
+Not assembled here: the legacy per-step recurrent verify (`TRTLLM_USE_GDN_REPLAY=0`), which
+routing declines and the bind refuses.
+
+Integration facts this file depends on:
+
+1. The engine rebuilds the composite config and normalizes its text tower
+ (`config_utils.load_pretrained_config`). The shell adopts that text config as
+ `model_config.pretrained_config` before building the core, because the recurrent-state and
+ PLE pool geometry are read from it.
+2. `position_ids` arrives in the mRoPE layout (`[3, 1, T]`) even for text-only batches. The
+ main rope reads all three rows; the QSA index plane takes coordinate 0 only.
+"""
+
+from __future__ import annotations
+
+import math
+import os
+import re
+from typing import NamedTuple
+
+import torch
+from torch import nn
+from transformers import PretrainedConfig
+
+from tensorrt_llm._torch._experimental.modeling_v2._router_index import step_contract_enabled
+from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import (
+ fused_qk_norm_rope,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import (
+ fused_qk_norm_rope as mtp_fused_qk_norm_rope,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import (
+ QM_FP8_KV_CACHE,
+ QM_NVFP4_KV_CACHE,
+ thop_attention,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.allgather import allgather
+from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.allreduce import (
+ allreduce,
+ allreduce_group_args,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.reducescatter import reducescatter
+from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm
+from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import (
+ cublas_mm as mtp_cublas_mm,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.catalog.graph.eager_on_graph import (
+ eager_on_graph,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import (
+ MAX_CERTIFIED_TOKENS as _MOE_MAX_ROWS,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import (
+ fp4_quantize,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.activation.fused_sigmoid_mul_inplace import (
+ fused_sigmoid_mul_inplace,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.activation.fused_sigmoid_mul_inplace import (
+ fused_sigmoid_mul_inplace as mtp_fused_sigmoid_mul_inplace,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.activation.silu_and_mul_rows import (
+ silu_and_mul_rows,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_canonical_selection import (
+ qsa_canonical_selection,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_expand_block_indices import (
+ qsa_expand_block_indices,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_mtp_reuse_topk import (
+ qsa_mtp_reuse_topk,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_nvfp4_pages_to_bf16 import (
+ qsa_nvfp4_pages_to_bf16,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_index_scores import (
+ qsa_paged_index_scores,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_index_store import (
+ qsa_paged_index_store,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_kv_store import (
+ qsa_paged_kv_store,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_kv_store_nvfp4 import (
+ qsa_paged_kv_store_nvfp4,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa import (
+ qsa_paged_sparse_gqa,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_decode import (
+ qsa_paged_sparse_gqa_decode,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_decode import (
+ supports as _sparse_decode_supports,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4 import (
+ qsa_paged_sparse_gqa_nvfp4,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4_decode import ( # noqa: E501
+ qsa_paged_sparse_gqa_nvfp4_decode,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4_verify import ( # noqa: E501
+ qsa_paged_sparse_gqa_nvfp4_verify,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4_verify import ( # noqa: E501
+ takes_grouped as _verify_takes_grouped,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_prefill_compress import (
+ qsa_prefill_compress,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_rollback_planes import (
+ plane_table as _plane_table,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_rollback_planes import (
+ qsa_rollback_gather,
+ qsa_rollback_restore,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_sparse_gqa_rowtile import (
+ qsa_sparse_gqa_rowtile as one_rank_sparse_gqa_rowtile,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_topk_expand_decode import (
+ qsa_topk_expand_decode,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_topk_expand_decode import (
+ supports as _topk_expand_supports,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_topk_select import (
+ qsa_topk_select,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qwen4_exp_attn_decode_prologue import ( # noqa: E501
+ SCORE_ROWS as _PROLOGUE_SCORE_ROWS,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qwen4_exp_attn_decode_prologue import ( # noqa: E501
+ qwen4_exp_attn_decode_prologue as one_rank_attn_decode_prologue,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qwen4_exp_qk_norm_rope import (
+ qwen4_exp_qk_norm_rope,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.spec_rollback_commit import (
+ spec_rollback_commit,
+ spec_rollback_stage,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.verify_step_setup import (
+ _BLOCK_S as _VERIFY_SETUP_MAX_SEQS,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.verify_step_setup import (
+ verify_step_setup,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.bf16_gemv_pdl import (
+ MAX_ROWS as _GEMV_MAX_ROWS,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.bf16_gemv_pdl import (
+ TILING as _GEMV_TILING,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.bf16_gemv_pdl import (
+ bf16_gemv_pdl,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.fp8_gemv import (
+ fp8_gemv_conv1d_op,
+ fp8_gemv_conv1d_verify_op,
+ fp8_gemv_op,
+ pack_from_bf16,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.fp8_per_tensor_mm import (
+ fp8_per_tensor_mm,
+ fp8_per_tensor_weight,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.low_m_bf16_mm import (
+ low_m_bf16_mm,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.mid_rows_mm import (
+ MidWeight,
+ mid_rows_mm_op,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.mid_rows_mm import (
+ precompile as mid_rows_precompile,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.nvfp4_gemv import (
+ nvfp4_gemv_op,
+ nvfp4_split_gemv_conv1d_op,
+ nvfp4_split_gemv_conv1d_verify_op,
+ nvfp4_split_gemv_op,
+ precompile_split,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.nvfp4_gemv import (
+ pack_from_bf16 as nvfp4_pack_from_bf16,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.nvfp4_gemv import (
+ pack_split_from_bf16 as nvfp4_pack_split_from_bf16,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.nvfp4_gemv import (
+ precompile as nvfp4_gemv_precompile,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_decode_first_cute import (
+ qwen4_exp_hc_decode_first_cute,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_decode_mix_fp4_cute import (
+ qwen4_exp_hc_decode_mix_cute,
+ qwen4_exp_hc_decode_mix_fp4_cute,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_decode_rows_tc import (
+ MAX_ROWS as _HC_DECODE_ROWS_MAX,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_decode_rows_tc import (
+ qwen4_exp_hc_decode_rows_first,
+ qwen4_exp_hc_decode_rows_mix,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_first_rescale_tc import (
+ qwen4_exp_hc_first_rescale_tc,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_inject_norm import (
+ qwen4_exp_hc_inject_norm,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_inject_rstd import (
+ qwen4_exp_hc_inject_rstd,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_lowrank_act import (
+ qwen4_exp_hc_lowrank_act,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_up_gate_mean import (
+ qwen4_exp_hc_up_gate_mean,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_up_gate_mean_rstd import (
+ qwen4_exp_hc_up_gate_mean_rstd,
+ qwen4_exp_hc_up_gate_mean_rstd_fp4,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_post_inject import (
+ qwen4_exp_post_inject,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion import ( # noqa: E501
+ cute_dsl_nvfp4_gather_grouped_gemm_act_fusion,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.cute_dsl_nvfp4_grouped_gemm import (
+ cute_dsl_nvfp4_grouped_gemm,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_block_scale_moe_renorm_permuted import (
+ fp4_block_scale_moe_renorm_permuted,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_block_scale_moe_renorm_routed import (
+ fp4_block_scale_moe_renorm_routed,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_moe_renorm_cute_dsl import (
+ fp4_moe_renorm_cute_dsl,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp8_block_scale_moe_runner import (
+ fp8_block_scale_moe_runner,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fused_sigmoid_gate_mul_add import (
+ fused_sigmoid_gate_mul_add,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.moe_sort import moe_sort
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.moe_unpermute_inplace import (
+ moe_unpermute_inplace,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.renorm_moe_routing_op import (
+ renorm_moe_routing_op,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.norm.flashinfer_gemma_rmsnorm import (
+ flashinfer_gemma_rmsnorm,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.norm.rms_norm_gated_token_major import (
+ rms_norm_gated_token_major,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ple_gate_conv_inject_norm import (
+ ple_gate_conv_inject_norm,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ple_ngram_embed import (
+ ple_ngram_embed,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.quantization.bf16_to_fp32 import (
+ bf16_to_fp32,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.quantization.fp8_quantize_1x128 import (
+ fp8_quantize_1x128,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.causal_conv1d_triton_update import (
+ causal_conv1d_triton_update,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.causal_conv1d_update import (
+ causal_conv1d_update,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.cu_seqlens_from_seq_lens import (
+ cu_seqlens_from_seq_lens,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.flashinfer_chunk_gated_delta_rule import (
+ INPLACE_MAX_TOKENS as _GDN_INPLACE_MAX_TOKENS,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.flashinfer_chunk_gated_delta_rule import (
+ flashinfer_chunk_gated_delta_rule,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.fused_recurrent_gated_delta_rule_cached_replay_update import ( # noqa: E501
+ fused_recurrent_gated_delta_rule_cached_replay_update,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.gdn_decode_gated_norm_pdl import (
+ MAX_ROWS as _GDN_DECODE_FUSED_MAX_ROWS,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.gdn_decode_gated_norm_pdl import (
+ gdn_decode_gated_norm_pdl,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.gdn_decode_rows_gated_norm import (
+ gdn_decode_rows,
+ gdn_decode_rows_gated_norm,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.gdn_prefill_conv_post import (
+ gdn_prefill_conv_post,
+)
+from tensorrt_llm._torch.attention.backends.fmha.phased import get_spec_decoding_position_offsets
+from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata
+from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata
+from tensorrt_llm._torch.model_config import ModelConfig
+from tensorrt_llm._torch.models.modeling_utils import DecoderModel, DecoderModelForCausalLM
+from tensorrt_llm._torch.modules.fla.cached_replay import CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE
+from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor
+from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks
+from tensorrt_llm._torch.speculative import get_spec_worker, should_use_separate_draft_kv_cache
+from tensorrt_llm._torch.utils import ActivationType
+from tensorrt_llm._utils import mpi_disabled
+from tensorrt_llm.functional import AllReduceFusionOp
+from tensorrt_llm.logger import logger
+from tensorrt_llm.quantization.mode import QuantAlgo
+
+from . import pool as _pool
+from . import vision as _vision
+from . import weights as _weights
+from .ops.attention.qsa_fp8_pages_to_bf16__prefill import qsa_fp8_pages_to_bf16
+from .ops.attention.qsa_sparse_gqa_rowtile__prefill import qsa_sparse_gqa_rowtile
+from .ops.attention.qwen4_exp_attn_decode_prologue__decode import qwen4_exp_attn_decode_prologue
+from .ops.hc.qwen4_exp_hc_decode_tc_first__decode import MAX_ROWS as _HC_TC_FIRST_ROWS
+from .ops.hc.qwen4_exp_hc_decode_tc_first__decode import precompile as hc_tc_first_precompile
+from .ops.hc.qwen4_exp_hc_decode_tc_first__decode import qwen4_exp_hc_decode_tc_first
+from .ops.hc.qwen4_exp_hc_up_gate_mean_fp4__moe import qwen4_exp_hc_up_gate_mean_fp4
+from .ops.ssm.gdn_cached_replay_cute__verify import gdn_cached_replay_cute
+
+# The GPU architecture this target is certified for; a direct instantiation could
+# bypass routing.
+_SM = (10, 3)
+
+# `moe_config.backend` request -> the routed-expert backend the trunk runs.
+_TRUNK_MOE_BACKENDS = {"AUTO": "TRTLLM", "TRTLLM": "TRTLLM", "CUTEDSL": "CUTEDSL"}
+# The CuteDSL grouped GEMMs' MMA tile M, which is also the tile `moe_sort` pads each expert's
+# rows to: the GEMMs walk the sort's tile tables. sm_103 admits 128 and 256.
+_CUTEDSL_TILE_SIZE = 128
+# The longest MTP draft the CuteDSL pipeline is certified at: its cells capture up to 128 rows
+# (32 requests of 4 verify rows).
+_CUTEDSL_MAX_DRAFT_LEN = 3
+
+# Rows per request of the widest verify step (`max_draft_len + 1`), and of the widest
+# one without the GDN cached replay (the per-step recurrent states).
+_MAX_VERIFY_WIDTH = _pool.MAX_VERIFY_WIDTH
+
+# Recurrent-state cache dtypes the target serves: at one rank the checkpoint's fp32 too.
+_SSM_STATE_DTYPES = (torch.bfloat16,)
+_ONE_RANK_SSM_STATE_DTYPES = (torch.float32, torch.bfloat16)
+
+# QSA backend metadata buffers rebuilt every step. Their absence means the engine did not
+# select the QSA metadata and the index plane would never be written.
+_QSA_FIELDS = (
+ "qsa_block_table",
+ "qsa_req_idx_per_token",
+ "qsa_logical_positions",
+ "qsa_sequence_lengths",
+ "qsa_visible_blocks",
+ "qsa_topk_indices",
+ "qsa_topk_row_starts",
+)
+
+# Fields read off `attn_metadata.mamba_metadata`. Absence is silent: a smoke run passes
+# with no recurrent state at all.
+_MAMBA_FIELDS = (
+ "state_indices",
+ "has_initial_states",
+ "query_start_loc",
+ "query_start_loc_long",
+)
+
+# thop_attention's constant arguments: in-kernel RoPE off (the rope runs before the op),
+# causal mask, no MLA, sparse or cross surface.
+_CALL_CONSTANTS = dict(
+ output_sf=None,
+ k=None,
+ v=None,
+ is_fused_qkv=True,
+ update_kv_cache=True,
+ attention_input_type=0,
+ is_mla_enable=False,
+ mask_type=1,
+ q_scaling=1.0,
+ attention_sinks=None,
+ position_embedding_type=0,
+ rotary_inv_freq=None,
+ rotary_cos_sin=None,
+ rope_dim=0,
+ rope_base=10000.0,
+ rope_scale_type=0,
+ rope_scale=1.0,
+ rope_short_m_scale=1.0,
+ rope_long_m_scale=1.0,
+ rope_max_positions=1024,
+ rope_original_max_positions=1024,
+ out_scale=None,
+ latent_cache=None,
+ q_pe=None,
+ q_lora_rank=None,
+ kv_lora_rank=None,
+ qk_nope_head_dim=None,
+ qk_rope_head_dim=None,
+ v_head_dim=None,
+ rope_append=None,
+ chunked_prefill_buffer_batch_size=1,
+ attention_chunk_size=None,
+ softmax_stats_tensor=None,
+ sparse_kv_indices=None,
+ sparse_kv_offsets=None,
+ sparse_attn_indices=None,
+ sparse_attn_offsets=None,
+ sparse_attn_indices_block_size=0,
+ mrope_rotary_cos_sin=None,
+ mrope_position_deltas=None,
+ helix_position_offsets=None,
+ helix_is_inactive_rank=None,
+ cross_kv=None,
+ relative_attention_bias=None,
+ relative_attention_max_distance=0,
+ kv_norm_weight=None,
+ kv_norm_eps=1e-6,
+ skip_correction_threshold=0.0,
+)
+
+_SF_VEC = _weights.SF_VEC
+#: Rows up to which the MoE router runs `low_m_bf16_mm` (FlashInfer's low-M direct
+#: and split-K kernels, PDL); more rows run cuBLAS. 32 is the top of both FlashInfer kernels'
+#: policy.
+_ROUTER_LOW_M_MAX_ROWS = 32 # capacity: fixed: a dispatch threshold, not a buffer size
+_ROUTING_METHOD_RENORM = 1
+_ACT_TYPE_SWIGLU = 0
+_PAD_SLOT_ID = -1
+#: Whether a captured decode step's projections on `low_m_bf16_mm` launch with programmatic
+#: dependent launch (their FlashInfer kernels stream the weight before waiting on the
+#: activation's producer; the cuBLAS branch ignores it).
+_DECODE_PROJ_PDL = True
+#: The K at which `_decode_projection` runs `bf16_gemv_pdl` up to `_GEMV_MAX_ROWS` rows: a
+#: tp4/tep4 rank's o and GDN out projections.
+_GEMV_TP4_K = 1536
+#: Widest decode projection that `_one_rank_projection` sends to `bf16_gemv_pdl` at 2..
+#: `_GEMV_MAX_ROWS` rows: the shared expert's one-column gate logit (~3 us a layer at 4 rows
+#: under `low_m_bf16_mm`).
+_GEMV_NARROW_N = 8 # capacity: fixed: a dispatch threshold, not a buffer size
+#: Whether a verify step's conv signals its dependents at its start and the CuTe replay after it
+#: launches with programmatic dependent launch: the replay then streams the checkpoint and the
+#: history (which nothing in the step writes ahead of it) while the conv runs, and reads the
+#: conv's output only after `griddepcontrol.wait`.
+_GDN_REPLAY_PDL = True
+_SOFTPLUS_THRESHOLD = 20.0
+_L2_NORM_EPS = 1e-6
+_PLE_CLAMP_MIN = 1e-6
+#: Row counts up to which a hyper-connection mixer runs its first half (projection and
+#: activations) as the one tcgen05 launch `qwen4_exp_hc_decode_tc_first` instead of cuBLAS
+#: split-K (+ its reduce) + `qwen4_exp_hc_lowrank_act`: the entry's MAX_ROWS (two 128-row
+#: MMA tiles). Every decode batch and MTP verify step a member captures is inside it
+#: (the HC is replicated, so the rows are this rank's tokens); larger steps keep cuBLAS.
+_HC_TC_FIRST_MAX_ROWS = _HC_TC_FIRST_ROWS
+#: At one rank (`hc_fast_sites`), decode rows up to which a hyper-connection site runs as
+#: `qwen4_exp_hc_decode_first_cute` + `qwen4_exp_hc_decode_mix_cute` (or its NVFP4 variant at an
+#: MoE site): two GEMV launches, which win at one row and lose from two (GB300 micro-bench). From
+#: there up to `_HC_DECODE_ROWS_MAX` the tcgen05 pair `qwen4_exp_hc_decode_rows_first` +
+#: `qwen4_exp_hc_decode_rows_mix`; larger steps, and every step at the other layouts,
+#: `qwen4_exp_hc_inject_norm` and the mixer. The two fast pairs form the mixer's low-rank
+#: activations within 1-2 bf16 ulp of that path, not bitwise, so a request's tokens would depend
+#: on how many rows its step carries; the other layouts keep one path at every batch size.
+_HC_DECODE_GEMV_MAX_ROWS = 1 # capacity: fixed: dispatch threshold, not a buffer bound
+#: Decode rows from which a folded MoE block (`DecodeTarget._moe`) launches its router before
+#: the shared expert: launched after the side stream's gate|up GEMM its CTAs waited for SMs
+#: behind that GEMM's. GB300 micro-bench (24 layers of cold weights in one graph), us a block,
+#: shared expert first vs router first: 2 rows 21.7 vs 22.4, 8 rows 46.5 vs 45.7, 32 rows
+#: 112.9 vs 112.1, 128 rows 206.7 vs 204.4, 256 rows 275.6 vs 271.7.
+_ROUTER_FIRST_MIN_ROWS = 8 # capacity: fixed: a dispatch threshold, not a buffer size
+#: At one rank (`hc_fast_sites`), prefill steps of at least this many rows run each trunk
+#: hyper-connection site on the un-normalized residual: `qwen4_exp_hc_inject_rstd` stores
+#: `hc_out` and a `[T, C]` rstd instead of the `[T, C*H]` normed stream, and the rstd mixer
+#: (`_hc_mixer_rstd`, `_hc_moe_mixer_rstd`) forms `normed` on chip. Smaller steps, and every step
+#: at the other layouts, keep `qwen4_exp_hc_inject_norm` and the stored-normed mixer, which the
+#: rstd mixer does not match bitwise (a generation row in a mixed step would round differently
+#: from the same row in a decode step). GB300 micro-bench, us a site
+#: (stored normed vs rstd): 2048 rows 57.0 vs 56.7, 3072 rows 82.4 vs 76.1, 4096 rows 104.0
+#: vs 90.8, 8192 rows 186.2 vs 164.9, 16384 rows 360.1 vs 293.7.
+_HC_RSTD_MIN_ROWS = 3072 # capacity: fixed: a dispatch threshold, not a buffer size
+#: Decode rows up to which a bf16 pool's gated sparse GQA runs the one-launch decode kernel
+#: (`qsa_paged_sparse_gqa_decode`) at one rank, and the index scores prefetch their block table.
+_SPARSE_DECODE_MAX_ROWS = 64 # capacity: fixed: a dispatch threshold
+_QSA_POSITION_AXES = _pool.QSA_POSITION_AXES
+# The compressed index keys rotate at coordinate 0 alone (no mRoPE section), as the index
+# query does.
+_QSA_MROPE_SECTION = None
+# The draft experts' activation scales are fp32, not the UE8M0 power-of-two form.
+_MTP_MOE_USE_UE8M0 = False
+# Rows per draft FP8 runner call: its 128-expert windows are certified up to one 8192-row
+# chunk (`fp8_block_scale_moe_runner`). The rows an attention-DP group gathers (up to
+# its size times a rank's) run in chunks of it.
+_MTP_MOE_MAX_ROWS = 8192
+
+# A context launch of at least this many rows over a bf16 or fp8 pool takes
+# `qsa_sparse_gqa_rowtile` (tensor cores, K/V tiles shared by a group of rows); below it the
+# per-row kernel is faster.
+_ROWTILE_MIN_ROWS = 2048 # capacity: fixed: a dispatch threshold, not a buffer size
+# Over an fp8 pool the launch reads its requests' pages widened into a bf16 scratch; a launch
+# whose scratch would pass this many bytes stays on the per-row kernel.
+_FP8_ROWTILE_MAX_SCRATCH = 1 << 30 # capacity: fixed: bounds a per-call scratch allocation
+# At one rank a step with more context rows than this over an NVFP4 pool attends them on the
+# row-tile pass over a bf16 scratch of their requests' pages (`_sparse_gqa_nvfp4_context`):
+# 16384 context rows 1.36 ms a layer there against 77-948 ms for the NVFP4 per-row kernel.
+_NVFP4_ROWTILE_MIN_ROWS = 256 # capacity: fixed: a dispatch threshold, not a buffer size
+# The bf16 K/V scratch those context rows read holds every page of each context request's KV
+# history. Consecutive context requests are dequantized and attended in groups whose pages fit
+# this many bytes; a request with more pages forms a group of its own.
+_NVFP4_CONTEXT_SCRATCH_BYTES = 128 << 20 # capacity: fixed: a scratch memory budget
+
+#: Set to 1 to run the GDN input / output and attention input / output projections in FP8
+#: (`fp8_per_tensor_mm`: e4m3 weights with a per-tensor scale, the activation quantized per
+#: call with a dynamic per-tensor scale) wherever a call reaches `_DENSE_FP8_MIN_ROWS` rows, and
+#: decode rows on `fp8_gemv` over packed FP8 copies with 128x128 block scales (`_FP8_GEMV_ROWS`)
+#: or `mid_rows_mm` (`_MID_MM_ROWS`). A lossy mode beside the checkpoint's bf16 projections
+#: (gsm8k-gated), certified at tp1; off by default and at every other layout.
+_DENSE_FP8_ENV = "TRTLLM_MODELING_V2_DENSE_FP8"
+
+#: Set to 1 to give each GDN layer an NVFP4 copy of its input projection (e2m1 values, one e4m3
+#: scale per 16 hidden elements of a row, an fp32 global scale; `nvfp4_gemv.pack_split`, built
+#: at load from the bf16 weight) and run its decode rows -- up to `_INPROJ_NVFP4_MAX_ROWS` -- on
+#: `fp8_gemv`'s kernel over it, the conv step as the GEMV's epilogue in a captured step
+#: (single-token and verify). Lossy (gsm8k-gated), certified at tp1; off by default.
+_GDN_INPROJ_NVFP4_ENV = "TRTLLM_MODELING_V2_GDN_INPROJ_NVFP4"
+
+#: Rows from which a projection takes the per-tensor FP8 GEMM, per projection kind: the input
+#: projections from 2 rows (one row stays on the GEMV entries, whose PDL overlap the two
+#: quantization launches would break), the K=6144 output projections only where the GEMM saving
+#: outgrows the activation quantization's two launches (GB300 micro-bench: cuBLASLt bf16 8-10 us
+#: vs FP8 6-7 us + ~4 us below ~512 rows).
+_DENSE_FP8_MIN_ROWS = {"in": 2, "out": 512} # capacity: fixed: dispatch thresholds
+
+#: Row ranges `(lo, hi)` in which a projection with a packed FP8 decode copy runs `fp8_gemv`
+#: (FP8 weights, bf16 activations, one PDL launch), per kind; each upper end is that kind's
+#: crossover against the next entry (GB300 micro, us): in 8.7 vs ~9 `mid_rows_mm` at 16 rows;
+#: out 6.4 vs 8.4 bf16 at 16; gate|up 2.6 vs 4.0 at 16; down 1.7 vs 2.9 at 8.
+_FP8_GEMV_ROWS = {"in": (1, 16), "out": (1, 16), "sh_gu": (1, 16), "sh_dn": (1, 8)} # dispatch
+
+#: Row ranges in which a projection runs `mid_rows_mm` (tcgen05, the per-tensor FP8 codes or the
+#: bf16 weight) between `fp8_gemv` and the GEMM entries.
+_MID_MM_ROWS = {
+ "in": (17, 256),
+ "out": (17, 128),
+ "sh_gu": (17, 64),
+ "sh_dn": (0, -1),
+ "router": (17, 64),
+} # capacity: fixed: dispatch thresholds
+
+#: `_mid` weight-name patterns -> the `_MID_MM_ROWS` kind their projections dispatch as.
+_MID_KINDS = (
+ (r"l\d+_(attn_in_w|in_proj_w)", "in"),
+ (r"l\d+_(o_w|out_w)", "out"),
+ (r"l\d+_sh_gu", "sh_gu"),
+ (r"l\d+_sh_dn", "sh_dn"),
+ (r"l\d+_router", "router"),
+)
+
+#: Rows up to which a GDN layer's decode in_proj reads the NVFP4 copy (`_GDN_INPROJ_NVFP4_ENV`).
+#: GB300 micro (us, conv included, behind a PDL predecessor): 6.4 / 6.1 at 1 / 4 rows against
+#: fp8_gemv_conv1d's 7.9 / 7.7; alone 7.1 against 7.9 at 8 rows, 9.1 against 8.6 at 16.
+_INPROJ_NVFP4_MAX_ROWS = 8 # capacity: fixed: a dispatch threshold, at most fp8_gemv's MAX_ROWS
+
+#: Rows up to which a captured verify step's conv runs in the in_proj GEMV's epilogue
+#: (`fp8_gemv_conv1d_verify`, certified at 2..16 rows); more rows take the Triton conv.
+_VERIFY_CONV_FUSED_MAX_ROWS = 16 # capacity: fixed: the entry's certified rows
+
+#: Rows up to which a captured single-token step runs the conv step as `fp8_gemv_conv1d`'s
+#: epilogue over the packed FP8 in_proj copy (dense FP8 mode).
+_FP8_CONV_FUSED_MAX_ROWS = 8 # capacity: fixed: a dispatch threshold
+
+
+def _switch(env: str) -> bool:
+ mode = os.environ.get(env, "0")
+ assert mode in ("0", "1"), f"{env}={mode!r}: expected 0 or 1"
+ return mode == "1"
+
+
+def _gemv_ready(x: torch.Tensor) -> bool:
+ """`fp8_gemv`'s activation layout: unit column stride, rows 16-byte aligned."""
+ return x.stride(1) == 1 and x.stride(0) % 8 == 0 and x.data_ptr() % 16 == 0
+
+
+def _mid_rows(kind: str, rows: int) -> bool:
+ lo, hi = _MID_MM_ROWS[kind]
+ return lo <= rows <= hi
+
+
+def _dense_mm(x, weight_t, fp8, gemm, kind: str, packed=None, mid=None) -> torch.Tensor:
+ """A dense projection: `fp8_gemv` over the packed decode copy (`packed`) within
+ `_FP8_GEMV_ROWS[kind]` rows, `mid_rows_mm` over `mid` (a `MidWeight`) within
+ `_MID_MM_ROWS[kind]` rows, the per-tensor FP8 GEMM (`fp8` = `(codes_t, scale)`) from
+ `_DENSE_FP8_MIN_ROWS[kind]` rows, else the step's bf16 entry `gemm` -- every call when the
+ dense FP8 mode is off (all copies None)."""
+ rows = x.shape[0]
+ lo, hi = _FP8_GEMV_ROWS.get(kind, (1, 0))
+ if packed is not None and lo <= rows <= hi and _gemv_ready(x):
+ return fp8_gemv_op(x, packed)
+ if mid is not None and _mid_rows(kind, rows) and (not mid.fp8 or _gemv_ready(x)):
+ return mid_rows_mm_op(x, mid)
+ if fp8 is not None and rows >= _DENSE_FP8_MIN_ROWS[kind]:
+ return fp8_per_tensor_mm(x, *fp8)
+ return gemm(x, weight_t)
+
+
+def _draft_mm(x: torch.Tensor, weight_t: torch.Tensor, packed, gemm) -> torch.Tensor:
+ """A draft-layer projection: `fp8_gemv` over its packed FP8 copy (`_MTP_DENSE_FP8_ENV`) at up
+ to `_MTP_DENSE_FP8_MAX_ROWS` rows, else `gemm` over the bf16 `weight_t`."""
+ if packed is not None and x.shape[0] <= _MTP_DENSE_FP8_MAX_ROWS and _gemv_ready(x):
+ return fp8_gemv_op(x, packed)
+ return gemm(x, weight_t)
+
+
+def _precompile_mid(mid: dict) -> None:
+ """Compile every `mid_rows_mm` variant the decode dispatch can reach, once per weight
+ geometry: a variant first met by a serving step would stall it for the compile."""
+ seen = {}
+ for name, mw in mid.items():
+ kind = next(k for rx, k in _MID_KINDS if re.fullmatch(rx, name))
+ lo, hi = _MID_MM_ROWS[kind]
+ key = (mw.n, mw.k, mw.fp8, None if mw.scale is None else mw.scale.numel())
+ if lo <= hi and key not in seen:
+ seen[key] = mid_rows_precompile(mw, max(lo, 1), hi)
+
+
+#: Set to 1 to give the FP8 draft head an NVFP4 copy of its lm_head (e2m1 values, one e4m3
+#: scale per 16 hidden elements of a vocabulary row, an fp32 global scale; built at load) and run
+#: its decode rows -- up to `_FP8_GEMV_HEAD_ROWS` -- on `nvfp4_gemv` (GB300, 248320 x 2560: 54-57
+#: us at 1-8 rows, 66-67 at 9-16, against fp8_gemv's 87-90 and 96-98). The verify head is
+#: untouched: outputs are unchanged and the acceptance length is the cost to watch. Certified at
+#: tp1; 0 (default): `fp8_gemv`.
+_MTP_HEAD_NVFP4_ENV = "TRTLLM_MODELING_V2_MTP_HEAD_NVFP4"
+
+#: Set to 1, under the dense FP8 mode, to run the target head's decode rows (up to
+#: `_FP8_GEMV_HEAD_ROWS`) on `nvfp4_gemv` over an NVFP4 copy of lm_head -- the draft head's under
+#: `_MTP_HEAD_NVFP4_ENV`, else one built at load -- instead of `fp8_gemv`. Lossy: the verify
+#: logits change. Certified at tp1; 0 (default): `fp8_gemv`.
+_TARGET_HEAD_NVFP4_ENV = "TRTLLM_MODELING_V2_TARGET_HEAD_NVFP4"
+
+#: Rows up to which the vocabulary projection (the draft head; under the dense FP8 mode the
+#: verify head too) runs `fp8_gemv` on the packed copy; more rows take cuBLASLt over the
+#: per-row-scaled copy. GB300, 248320 x 2560: fp8_gemv 86-98 us at 1-16 rows against 107-134,
+#: then 142 at 20 rows and 172 at 32 against a flat ~108.
+_FP8_GEMV_HEAD_ROWS = 16 # capacity: fixed: a dispatch threshold, at most fp8_gemv's MAX_ROWS
+
+#: Rows up to which the FP8 vocabulary projection's rows past `_FP8_GEMV_HEAD_ROWS` run
+#: `mid_rows_mm` over the per-row-scaled copy (activation quantized per row); more rows keep
+#: cuBLASLt. GB300, 248320 x 2560, us at 17 / 32 / 64 / 128 rows: 99.7 / 101.5 / 106.0 / 125.8
+#: against 110.0 / 111.9 / 107.7 / 137.5.
+_MID_HEAD_ROWS = 128 # capacity: fixed: a dispatch threshold, at most mid_rows_mm's MAX_ROWS
+
+#: Set to 1 to run the draft layer's dense projections (`fc_e`, `fc_h`, the fused attention
+#: input, `o`, the shared expert's gate|up and down) on `fp8_gemv` over packed FP8 copies at decode
+#: rows (up to `_MTP_DENSE_FP8_MAX_ROWS`). The draft only proposes tokens, so outputs are unchanged;
+#: only the acceptance length can move. Certified at tp1; 0 (default): bf16.
+_MTP_DENSE_FP8_ENV = "TRTLLM_MODELING_V2_MTP_DENSE_FP8"
+_MTP_DENSE_FP8_MAX_ROWS = 16 # capacity: fixed: a dispatch threshold, at most fp8_gemv's MAX_ROWS
+
+#: Set to 1 to run the draft layer's routed experts in NVFP4 on the trunk's trtllm-gen runner: at
+#: load the checkpoint's FP8 128x128 block-scaled experts are requantized (per-expert weight global
+#: scales, 16-element e4m3 block scales) into the trunk's layout and the FP8 copies are freed. The
+#: draft only proposes tokens, so this can change the acceptance length, not outputs. Certified at
+#: tp1; 0 (default): the FP8 block-scale runner.
+_MTP_MOE_NVFP4_ENV = "TRTLLM_MODELING_V2_MTP_MOE_NVFP4"
+
+#: The NVFP4 draft MoE's two activation global scales are static, `448 * 6 / A` with `A` this factor
+#: times the largest calibrated amax of the trunk's routed experts (its `input_scale`s, FC1 input and
+#: FC2 input; the draft layer carries none). A global scale only has to keep the e4m3 block scales
+#: inside their range, so the margin costs nothing until blocks are ~1e-3 of `A`.
+_MTP_MOE_NVFP4_FC1_MARGIN = 8.0
+_MTP_MOE_NVFP4_FC2_MARGIN = 4.0
+
+_PLE_HOST_OFFLOAD_ENV = "TRTLLM_MODELING_V2_PLE_HOST_OFFLOAD"
+_PLE_HOST_OFFLOAD_OFF = frozenset({"0", "false", "no", "off"})
+
+#: At one rank, context-step MoE launches of at least this many rows (and at most
+#: `_CUTE_DSL_MOE_MAX_ROWS`, the op's certified top) run on the CuTe DSL MoE
+#: (`fp4_moe_renorm_cute_dsl`, over the second expert copy `weights.py` keeps in its layout);
+#: smaller ones, every decode one among them, keep the trtllm-gen runner. GB300 micro-benches,
+#: us a layer, CuTe DSL vs trtllm-gen: 2048 rows 275 vs 273, 8192 rows 451 vs 555, 16384 rows
+#: 656 vs 882.
+_CUTE_DSL_MOE_MIN_ROWS = 4096 # capacity: fixed: a dispatch threshold, not a buffer size
+_CUTE_DSL_MOE_MAX_ROWS = 16384 # capacity: fixed: the op's certified rows
+# Whether one rank keeps that second copy (about 1.3 GiB a layer). No engine arg owns it, hence
+# an env var read at construction; "off" (the spellings above) drops the copy and every launch
+# runs on trtllm-gen.
+_MOE_PREFILL_COPY_ENV = "TRTLLM_MODELING_V2_MOE_PREFILL_COPY"
+
+
+def _moe_prefill_copy(one_rank: bool, trunk_moe_backend: str, max_num_tokens: int) -> bool:
+ """Whether the trunk keeps the CuteDSL-layout expert copy: at one rank only (every expert
+ is local, the layout the CuteDSL MoE is certified at), only under trtllm-gen, only when a
+ step's rows (at most `max_num_tokens`) can reach the dispatch threshold, and not when the
+ env var turns it off."""
+ return (
+ one_rank
+ and trunk_moe_backend == "TRTLLM"
+ and max_num_tokens >= _CUTE_DSL_MOE_MIN_ROWS
+ and os.environ.get(_MOE_PREFILL_COPY_ENV, "1").strip().lower() not in _PLE_HOST_OFFLOAD_OFF
+ )
+
+
+def _ple_host_offload_enabled() -> bool:
+ """Whether the n-gram table is declared in host memory (the default)."""
+ value = os.environ.get(_PLE_HOST_OFFLOAD_ENV, "1").strip().lower()
+ return value not in _PLE_HOST_OFFLOAD_OFF
+
+
+def _moe_layer_scalars(w, i: int, start: int, count: int) -> dict[str, torch.Tensor]:
+ """Layer `i`'s routed-expert scalars for the trtllm-gen runner over this rank's window
+ `[start, start + count)`, built from the checkpoint's `input_scale` / `weight_scale_2`
+ rows (declared over the whole routing space on every rank).
+
+ The checkpoint stores modelopt reciprocals (`input_scale = 1 / g_act`, `weight_scale_2 =
+ 1 / g_w`), so a quantizer's global scale is `1 / input_scale` and the runner's scalars
+ are products of the two."""
+ isc1, isc1_up = w[f"l{i}_e_isc1"], w[f"l{i}_e_isc1_up"]
+ assert torch.equal(isc1, isc1_up), (
+ f"layer {i}: expert gate/up input_scale differ; the fused FC1 reads one activation "
+ "quantization for both halves"
+ )
+ assert torch.equal(w[f"l{i}_e_ws2_1"], w[f"l{i}_e_ws2_1_up"]), (
+ f"layer {i}: expert gate/up weight_scale_2 differ; the fused FC1 carries one alpha "
+ "per expert for both halves"
+ )
+ # One activation quantization feeds all experts, on every rank.
+ assert torch.equal(isc1.min(), isc1.max()), (
+ f"layer {i}: the routed experts do not share one gate_proj.input_scale, so a single "
+ "activation quantization cannot feed them all"
+ )
+ e_isc2, ws2_2 = w[f"l{i}_e_isc2"], w[f"l{i}_e_ws2_2"]
+ gate1 = isc1 * w[f"l{i}_e_ws2_1"]
+ window = slice(start, start + count)
+ return dict(
+ # The routed FC1 activation global scale (one per layer).
+ global_scale=(1.0 / isc1[:1]).contiguous(),
+ # [count] each. Swapping the first two is finite and silent.
+ output1_scale_scalar=(gate1 / e_isc2)[window].contiguous(),
+ output1_scale_gate_scalar=gate1[window].contiguous(),
+ output2_scale_scalar=(e_isc2 * ws2_2)[window].contiguous(),
+ )
+
+
+def _moe_layer_scalars_cd(w, i: int) -> dict[str, torch.Tensor]:
+ """Layer `i`'s routed-expert scalars for the CuteDSL MoE over the second expert copy (one
+ rank, every expert; `_moe_layer_scalars` checks the checkpoint facts they rest on). FC1's
+ per-expert alpha is `input_scale * weight_scale_2`, its NVFP4 input the same activation
+ quantization as the trtllm-gen runner's; FC2's input is quantized at one global scale,
+ `1 / max(input_scale)` over the layer's experts, and its alpha carries that maximum."""
+ isc2_max = w[f"l{i}_e_isc2"].max()
+ return dict(
+ alpha1=(w[f"l{i}_e_isc1"] * w[f"l{i}_e_ws2_1"]).contiguous(),
+ alpha2=(isc2_max * w[f"l{i}_e_ws2_2"]).contiguous(),
+ global_sf=(1.0 / isc2_max).reshape(1).contiguous(),
+ )
+
+
+class _Replay(NamedTuple):
+ """The GDN cached replay's history buffers and bookkeeping of one manager."""
+
+ old: dict # GDN layer -> (old_x, old_B, old_dt, old_dA_cumsum)
+ cache_buf_idx: torch.Tensor # [slots] int32: the history half each slot writes next
+ prev_num_accepted: torch.Tensor # [slots] int32: tokens pending since the checkpoint
+ history: int # rows a history half holds
+ width: int # the one verify width a step may carry per generation request
+
+
+class _Bound(NamedTuple):
+ """What the core reads off one KV cache manager: the manager, this rank's pool views
+ (`pool.RankPools`) and, under MTP, the verify step's runtime-owned buffers: the per-step
+ conv windows (`int_conv`, by GDN layer), the rows `update_mamba_states` promotes them
+ from (`int_idx`) and the cached replay. Valid only for that manager."""
+
+ cm: object
+ pools: _pool.RankPools
+ int_idx: torch.Tensor | None = None
+ int_conv: dict | None = None
+ replay: _Replay | None = None
+
+
+class _DraftBound(NamedTuple):
+ """The draft layer's planes in the manager its draft step is handed: its QSA index plane,
+ that manager's position plane, its paged K/V and, under IndexShare, its row of the
+ metadata's shared selection buffers."""
+
+ cm: object
+ index_cache: torch.Tensor
+ position_cache: torch.Tensor
+ kv: _pool.KvPlanes
+ share_slot: int | None
+
+
+class Qwen4ExpDefaultModelingV2(DecoderModel):
+ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig):
+ """`composite_config` is passed in because the text sub-config adopted as
+ `model_config.pretrained_config` carries neither the vision half nor the media
+ token ids."""
+ super().__init__(model_config)
+ cfg = model_config.pretrained_config
+ mapping = model_config.mapping
+ assert torch.cuda.get_device_capability() == _SM, (
+ f"target certified on sm_{_SM[0]}{_SM[1]}, running on "
+ f"sm_{''.join(map(str, torch.cuda.get_device_capability()))}"
+ )
+ assert getattr(cfg, "model_type", None) == "qwen4_exp_text", (
+ f"the core must be built on the normalized text config; got model_type "
+ f"{getattr(cfg, 'model_type', None)!r}"
+ )
+ self.mapping = mapping
+ # The ranks attention heads and the recurrent heads are split over: the TP group,
+ # or one rank under attention DP.
+ self.attn_tp_size = _pool.attention_tp_size(mapping)
+ # Under attention DP each rank serves its own requests and the MoE blocks run over the
+ # rows the whole group holds (`_expert_sum_op`). The group is the TP group, which spans
+ # the DP ranks; its size is the world's, never a constant.
+ self.dp_size = mapping.tp_size if mapping.enable_attention_dp else 1
+ self.dp_rank = mapping.tp_rank
+ self.dp_group = list(mapping.tp_group)
+ # One-model MTP-Eagle (routing admits no other speculation): the checkpoint's one MTP
+ # module, built only when the engine asked for speculation.
+ spec_config = getattr(model_config, "spec_config", None)
+ self.mtp_enabled = spec_config is not None
+ # The draft layer's FP8 block-scale experts are whole-expert windows over the TP group
+ # at every layout, tp4 included (`weights.rank_dims`, `mtp_local_experts`).
+ # Rows per generation request of the widest verify step (`max_draft_len + 1`); 1
+ # without MTP.
+ self.max_verify_rows = spec_config.max_draft_len + 1 if self.mtp_enabled else 1
+ # The most rows one forward carries (a draft step's included): sizes the draft
+ # layer's unit injection weight (`_MtpWeights.unit_inject`).
+ self.max_num_tokens = model_config.max_num_tokens
+
+ self.num_layers = cfg.num_hidden_layers
+ self.hidden = cfg.hidden_size
+ self.vocab = cfg.vocab_size
+ self.eps = cfg.rms_norm_eps
+ layer_types = list(cfg.layer_types)
+ self.is_attn = [t == "full_attention" for t in layer_types]
+ self.attn_layers = [i for i, a in enumerate(self.is_attn) if a]
+ self.gdn_layers = [i for i, a in enumerate(self.is_attn) if not a]
+
+ # ---- hyper-connections ------------------------------------------
+ self.hc_count = cfg.hc_count
+ self.hc_lowrank = cfg.hc_lowrank
+ self.stream = self.hc_count * self.hidden
+ # Columns of the padded first-projection weight (`_fuse_hc_first_projections`) the
+ # first projection multiplies: the real `[down; block_inject]` rows rounded up to 32.
+ self.hc_first_cols = -(-(self.hc_lowrank + self.hc_count) // 32) * 32
+
+ # ---- RoPE (partial, neox pairs, interleaved mRoPE) -----------------
+ self.head_dim = cfg.head_dim
+ rope = getattr(cfg, "rope_parameters", None) or getattr(cfg, "rope_scaling", None)
+ self.theta = float(rope["rope_theta"])
+ partial = float(rope.get("partial_rotary_factor", cfg.partial_rotary_factor))
+ self.rotary_dim = int(self.head_dim * partial)
+ self.mrope_declared = str(rope.get("type", rope.get("rope_type"))) == "mrope"
+ section = list(rope["mrope_section"])
+ self.mrope_section1 = int(section[1])
+ self.mrope_section2 = int(section[2])
+ self.max_positions = cfg.max_position_embeddings
+
+ # ---- QSA sparse index (replicated on every rank) -------------------
+ self.index_heads = cfg.indexer_n_heads
+ self.index_kv_heads = cfg.indexer_kv_heads
+ self.index_head_dim = cfg.indexer_head_dim
+ self.token_topk = cfg.indexer_budget
+ self.compress_ratio = cfg.indexer_compress_ratio
+ self.block_topk = self.token_topk // self.compress_ratio
+ # Dense-attention cutoff: QSA resolves max(token_topk, seq_len_threshold); checked
+ # against the engine's value because the CUDA-graph runner splits its short and long
+ # decode families on it.
+ self.dense_threshold = self.token_topk
+ sparse_cfg = getattr(model_config, "sparse_attention_config", None)
+ assert getattr(sparse_cfg, "algorithm", None) == "qsa", (
+ f"this target is a QSA sparse-attention assembly and the engine resolved "
+ f"sparse_attention_config={sparse_cfg!r}; QSA is this architecture's own default"
+ )
+ qsa_params = sparse_cfg.to_sparse_params(pretrained_config=cfg)
+ assert (
+ qsa_params.index_n_heads,
+ qsa_params.index_kv_heads,
+ qsa_params.index_head_dim,
+ qsa_params.token_topk,
+ qsa_params.compress_ratio,
+ ) == (
+ self.index_heads,
+ self.index_kv_heads,
+ self.index_head_dim,
+ self.token_topk,
+ self.compress_ratio,
+ ), (
+ f"the engine lowered the QSA index geometry to {qsa_params!r}; the side caches "
+ "are sized from the engine's copy and the kernels are called with this one"
+ )
+ assert qsa_params.dense_seq_len_threshold == self.dense_threshold, (
+ f"the engine's resolved QSA cutoff is {qsa_params.dense_seq_len_threshold} and "
+ f"this forward branches at {self.dense_threshold}: the runner captures its short "
+ "decode graph family at the engine's number"
+ )
+ # IndexShare for the MTP draft loop (`index_share_for_mtp_iteration`), as the engine
+ # resolved it: the same lowering sizes the metadata's shared selection buffers.
+ self.mtp_index_share = bool(qsa_params.mtp_index_share)
+ self.index_width = (self.index_heads + self.index_kv_heads) * self.index_head_dim
+ # The index query's `[q | k | v]` scratch for its norm + rope; the spare head is the
+ # v slot the op needs but never reads.
+ self.index_scratch_width = self.index_width + self.index_head_dim
+ self.softmax_scale = 1.0 / math.sqrt(self.head_dim)
+
+ # ---- the MTP module ------------------------------------------------
+ # Global layer index of the draft layer (the KV managers map it there, and the
+ # weight table keys it there); its rope base is its own block's, else the trunk's.
+ self.mtp_layer_idx = self.num_layers
+ mtp_block = getattr(cfg, "mtp", None)
+ self.mtp_theta = (
+ float(mtp_block.get("rope_theta", self.theta))
+ if isinstance(mtp_block, dict)
+ else self.theta
+ )
+
+ # ---- gated deltanet ----------------------------------------------
+ self.head_k_dim = cfg.linear_key_head_dim
+ self.head_v_dim = cfg.linear_value_head_dim
+ self.conv_kernel = cfg.linear_conv_kernel_dim
+
+ # ---- MoE ----------------------------------------------------------
+ self.num_experts = cfg.num_experts
+ self.topk = cfg.num_experts_per_tok
+ self.moe_inter = cfg.moe_intermediate_size
+ self.shared_inter = cfg.shared_expert_intermediate_size
+ requested = model_config.moe_backend_requested.upper()
+ assert requested in _TRUNK_MOE_BACKENDS, (
+ f"moe_config.backend={model_config.moe_backend_requested!r} is not implemented by "
+ f"this target; the supported requests are {sorted(_TRUNK_MOE_BACKENDS)}"
+ )
+ # Fixes the load-time layout of the expert stacks (`weights.py`).
+ self.trunk_moe_backend = _TRUNK_MOE_BACKENDS[requested]
+ # The CuteDSL pipeline runs every expert at once: one rank only.
+ assert self.trunk_moe_backend != "CUTEDSL" or mapping.world_size == 1, (
+ "moe_config.backend=CUTEDSL is implemented at one rank (tp1) only; this layout "
+ f"has {mapping.world_size} ranks"
+ )
+ assert self.trunk_moe_backend != "CUTEDSL" or (
+ self.max_verify_rows - 1 <= _CUTEDSL_MAX_DRAFT_LEN
+ ), (
+ f"max_draft_len={self.max_verify_rows - 1} with moe_config.backend=CUTEDSL: the "
+ f"CuteDSL expert pipeline is certified at max_draft_len up to {_CUTEDSL_MAX_DRAFT_LEN}"
+ )
+ # At one rank the CuteDSL-layout copy of the experts for prefill-sized launches
+ # (`_moe_prefill_copy`); `weights.py` declares it from this.
+ self.moe_prefill_copy = _moe_prefill_copy(
+ mapping.world_size == 1, self.trunk_moe_backend, model_config.max_num_tokens
+ )
+ # The draft experts' NVFP4 copy at one rank (`_mtp_moe_nvfp4`), built once: it frees
+ # the FP8 experts, and the shell's post-load walk reaches `post_load_weights` again.
+ self._mtp_moe_nvfp4 = None
+
+ # ---- PLE -----------------------------------------------------------
+ # One-based in the config.
+ self.ple_layer = list(cfg.ple_layer_ids)[0] - 1
+ self.ngram_size = cfg.ngram_size
+ self.ple_dilation = self.ngram_size
+ self.ngram_vocab_divisor = cfg.make_ngram_vocab_size_divisible_by
+ eos = cfg.eos_token_id
+ self.eos_token_id = int(eos[0] if isinstance(eos, (list, tuple)) else eos)
+ self.ple_table_dtype = _weights.PLE_TABLE_DTYPE
+ self.ple_host_offload = _ple_host_offload_enabled()
+ # CPU allocation behind the offloaded table; the `ple_table` parameter is a
+ # device-address view that owns none of those pages.
+ self._ple_table_host = None
+
+ # ---- the vision half (the weight table declares the tower) ---------
+ vcfg = getattr(composite_config, "vision_config", None)
+ assert vcfg is not None, "the composite config carries no vision_config"
+ self.vision_config = vcfg
+ self.image_token_id = int(composite_config.image_token_id)
+ self.video_token_id = int(composite_config.video_token_id)
+
+ # ---- KV pool storage ------------------------------------------------
+ self.kv_cache_quant_algo = model_config.quant_config.kv_cache_quant_algo
+ assert self.kv_cache_quant_algo in (None, QuantAlgo.FP8, QuantAlgo.NVFP4), (
+ f"this target stores K/V as the checkpoint's bf16, as fp8-e4m3 or as NVFP4; the "
+ f"engine resolved kv_cache_quant_algo {self.kv_cache_quant_algo!r}"
+ )
+ self.quant_mode = {
+ None: 0,
+ QuantAlgo.FP8: QM_FP8_KV_CACHE,
+ QuantAlgo.NVFP4: QM_NVFP4_KV_CACHE,
+ }[self.kv_cache_quant_algo]
+
+ # ---- this rank's share ----------------------------------------------
+ d = _weights.MODEL_WEIGHTS.dims(self)
+ assert d.expert_split in _weights.LOADED_EXPERT_SPLITS, d.expert_split
+ self.heads_q = d.heads_q
+ self.heads_kv = d.heads_kv
+ self.kv_head_start = d.kv_head_start
+ self.q_width = d.q_width
+ self.kv_width = d.kv_width
+ self.qkv_width = self.q_width + 2 * self.kv_width
+ # The fused input projection's output row, `[q | k | v | gate | index q, k | v slot]`
+ # (`weights._attn_in_w`). Its leading columns are the packed QKV the main norm + rope
+ # rewrites in place; the rope addresses a row as `(Hq + Hk + Hv) * head_dim` and
+ # never touches the v heads, so everything after k is "v" to it.
+ self.attn_in_width = d.attn_in_rows
+ assert self.attn_in_width == self.qkv_width + self.q_width + self.index_scratch_width
+ assert self.attn_in_width % self.head_dim == 0, self.attn_in_width
+ self.attn_in_tail_heads = self.attn_in_width // self.head_dim - self.heads_q - self.heads_kv
+ self.index_col = self.qkv_width + self.q_width
+ self.num_k_heads = d.num_k_heads
+ self.num_v_heads = d.num_v_heads
+ self.key_dim = d.key_dim
+ self.value_dim = d.value_dim
+ self.conv_dim = d.conv_dim
+ # Column widths of the one GDN input projection's output: [q|k|v | z | a | b].
+ self.in_proj_split = d.in_proj_split
+ self.local_experts = d.local_experts
+ self.expert_offset = d.expert_offset
+ # The runner's `intermediate_size`: each held expert's intermediate as stored, this
+ # rank's slice zero-padded at tp4 (`weights.EXPERT_INTER_ALIGN`), whole otherwise.
+ self.expert_inter = d.expert_inter
+ # The draft layer's FP8 experts: this rank's window over the whole TP group.
+ self.mtp_local_experts = d.mtp_local_experts
+ self.mtp_expert_offset = d.mtp_expert_offset
+
+ # Every weight but the n-gram table; meta-init intercepts torch.empty.
+ self.w = _weights.MODEL_WEIGHTS.declare(d)
+
+ # Derived once after load (`post_load_weights`).
+ self._moe_scalars: list | None = None
+ self._gdn_gating: dict | None = None
+ self._cos_sin: torch.Tensor | None = None
+ self._hc_first: dict | None = None
+ # The attention op's KV scales: fp32 [1] ones over an fp8 pool, fp32 [3] ones (q, K,
+ # V) over an NVFP4 one, None over bf16.
+ self.kv_scale_orig_quant: torch.Tensor | None = None
+ self.kv_scale_quant_orig: torch.Tensor | None = None
+ # The shared expert's stream (see `_moe_block`); made after load, because a stream
+ # created inside a graph capture is not capturable.
+ self._side_stream: torch.cuda.Stream | None = None
+ # The tensor-parallel sum of a block's partial output (`_tp_sum_op`), and the routed
+ # experts' completion over this rank's rows (`_expert_sum_op`).
+ self.tp_sum = None
+ self.expert_sum = None
+ self._prefill: "PrefillTarget | None" = None
+ self._decode: "DecodeTarget | None" = None
+ # Everything derived from the KV cache manager (`_bind_cache_manager`).
+ self._bound: _Bound | None = None
+ self._binds = 0
+ # Set by `begin_step`, cleared by `forward`.
+ self._step_begun = False
+ # MTP module operands (`_mtp_views`) and its index rope table; None with speculation
+ # off.
+ self._mtp: "_MtpWeights | None" = None
+ self._mtp_cos_sin: torch.Tensor | None = None
+ # The verify step's side-cache rollback, registered with the spec worker.
+ self.spec_rollback: SpeculativeRollback | None = (
+ SpeculativeRollback() if self.mtp_enabled else None
+ )
+ # The draft layer's own bind, keyed on the manager its draft step sees.
+ self._draft_bound: _DraftBound | None = None
+
+ # ---------------------------------------------------------------- load ---
+
+ def _fuse_hc_first_projections(self) -> None:
+ """Concatenate each hyper-connection's `down` and `block_inject` weights into one
+ `[rows, stream]` tensor, zero-padded to a multiple of 128 rows, so a site's first
+ projection is one GEMM: the trunk's sites, keyed `(layer, tag)` and `"fhc"`, and with
+ MTP the draft layer's, keyed `"mtp_ahc"`, `"mtp_mhc"` and `"mtp_fhc"`. The
+ checkpoint's parameters are rebound as row views of the padded tensors."""
+ w = self.w
+
+ def fuse(*params):
+ rows = sum(t.shape[0] for t in params)
+ padded = -(-rows // 128) * 128
+ pad = torch.zeros(
+ [padded - rows, self.stream], dtype=params[0].dtype, device=params[0].device
+ )
+ both = torch.cat([t.data for t in params] + [pad], 0)
+ start = 0
+ for t in params:
+ n = t.shape[0]
+ t.data = both[start : start + n]
+ start += n
+ return both
+
+ self._hc_first = {
+ (i, tag): fuse(w[f"l{i}_{tag}_down"], w[f"l{i}_{tag}_inject"])
+ for i in range(self.num_layers)
+ for tag in ("ahc", "mhc")
+ }
+ self._hc_first["fhc"] = fuse(w["fhc_down"])
+ if self.mtp_enabled:
+ p = f"l{self.mtp_layer_idx}_mtp_"
+ self._hc_first["mtp_ahc"] = fuse(w[f"{p}ahc_down"], w[f"{p}ahc_inject"])
+ self._hc_first["mtp_mhc"] = fuse(w[f"{p}mhc_down"], w[f"{p}mhc_inject"])
+ self._hc_first["mtp_fhc"] = fuse(w[f"{p}fhc_down"])
+
+ def post_load_weights(self) -> None:
+ """Derive what the targets read, then build them, now that the weights are real.
+
+ Not an override (DecoderModel has no such hook); named for the shell hook."""
+ w = self.w
+ device = w["embed"].device
+ self._moe_scalars = [
+ _moe_layer_scalars(w, i, self.expert_offset, self.local_experts)
+ for i in range(self.num_layers)
+ ]
+ # The CuteDSL scalars: of the second expert copy (`moe_prefill_copy`), or of the
+ # experts themselves on the CuteDSL backend.
+ self._moe_scalars_cd = (
+ [_moe_layer_scalars_cd(w, i) for i in range(self.num_layers)]
+ if self.moe_prefill_copy or self.trunk_moe_backend == "CUTEDSL"
+ else None
+ )
+ # The delta-rule entries index A_log and dt_bias as `ptr + head` and certify fp32;
+ # the checkpoint is bf16.
+ self._gdn_gating = {
+ i: (
+ w[f"l{i}_A_log"].float().contiguous(),
+ w[f"l{i}_dt_bias"].float().contiguous(),
+ )
+ for i in self.gdn_layers
+ }
+ # Reconcile the table height against the checkpoint's id-space tensors: an index past
+ # the end is a device-side assert that poisons the process.
+ table_rows = w["ple_table"].shape[0]
+ used_rows = int(w["ple_offs"][-1]) + int(w["ple_vocab"][-1])
+ assert used_rows <= table_rows, (
+ f"the n-gram table has {table_rows} rows but its per-head offsets and vocabulary "
+ f"sizes address {used_rows}"
+ )
+ assert table_rows % self.ngram_vocab_divisor == 0, (table_rows, self.ngram_vocab_divisor)
+ self._cos_sin = self._index_rope_table(self.theta, device)
+ # Before the draft layer's views, which read its fused first projections.
+ self._fuse_hc_first_projections()
+ if self.mtp_enabled:
+ # The draft layer's index plane rotates at the module's own rope base, like its
+ # attention rope.
+ self._mtp_cos_sin = (
+ self._cos_sin
+ if self.mtp_theta == self.theta
+ else self._index_rope_table(self.mtp_theta, device)
+ )
+ self._mtp = _mtp_views(self)
+ self._build_draft_one_rank(device)
+ if self.quant_mode in (QM_FP8_KV_CACHE, QM_NVFP4_KV_CACHE):
+ # Built here, not at construction: that runs under meta init.
+ scales = 3 if self.quant_mode == QM_NVFP4_KV_CACHE else 1
+ self.kv_scale_quant_orig = torch.ones(scales, dtype=torch.float32, device=device)
+ self.kv_scale_orig_quant = 1.0 / self.kv_scale_quant_orig
+ self._build_dense_copies(device)
+ # The CPU-only load checks build the targets over host weights; they launch nothing.
+ if device.type == "cuda":
+ # Every first-half variant `_hc_mixer` launches up to `_HC_TC_FIRST_MAX_ROWS` rows,
+ # compiled here so that no serving step (or graph capture) meets a compile.
+ hc_tc_first_precompile(self.stream, self.hc_lowrank, self.hc_count, device)
+ self._report_ple_residency()
+ self._side_stream = torch.cuda.Stream(device=device)
+ # The prefill CuteDSL MoE's output memset, beside its FC1 (`fp4_moe_renorm_cute_dsl`):
+ # its own stream, so the join before FC2 does not also wait for the shared expert.
+ self._moe_memset_stream = torch.cuda.Stream(device=device)
+ self.tp_sum = self._tp_sum_op()
+ self.expert_sum = self._expert_sum_op()
+ # A decode step's bf16 projection entry, and its GDN in_proj's: at tp1 the one-rank
+ # dispatch (`_one_rank_projection`), at the other layouts the shapes' own.
+ one_rank = self.mapping.world_size == 1
+ self.decode_mm = _one_rank_projection if one_rank else _decode_projection
+ # At one rank a trtllm-gen decode MoE leaves its finalize to the next site
+ # (`DecodeTarget._moe`).
+ self.moe_fold = one_rank and self.trunk_moe_backend == "TRTLLM"
+ # At one rank the hyper-connection sites' fast forms by row count (the decode GEMV and
+ # tcgen05 rows pairs, the prefill rstd sites); at the other layouts one form serves every
+ # batch size, so a request's tokens do not depend on what its steps carry beside it.
+ self.hc_fast_sites = one_rank
+ self.decode_in_mm = (
+ _one_rank_projection
+ if one_rank
+ else (lambda mat_a, mat_b: low_m_bf16_mm(mat_a, mat_b, _DECODE_PROJ_PDL))
+ )
+ self._prefill = PrefillTarget(self)
+ self._decode = DecodeTarget(self)
+ # The FP8 draft head at tp1, set by the shell's `post_load_weights`.
+ self.draft_head = None
+
+ def _build_draft_one_rank(self, device) -> None:
+ """At one rank the draft layer's operands for `DecodeTarget._mtp_layer`: the fused
+ `[q | k | v | gate | index | pad]` attention input (`weights._attn_in_w`), with
+ `_MTP_MOE_NVFP4_ENV` its routed experts in NVFP4, and with `_MTP_DENSE_FP8_ENV` the
+ packed FP8 copies of its dense projections; None / empty at every other layout."""
+ self._mtp_attn_in = None
+ self._mtp_gemv8 = {}
+ if self.mapping.world_size != 1:
+ return
+ w = self.w
+ p = f"l{self.mtp_layer_idx}_mtp_"
+ attn_in = _weights._attn_in_w(
+ (w[f"{p}q_w"], w[f"{p}k_w"], w[f"{p}v_w"], w[f"{p}idx_w"]), self
+ )
+ self._mtp_attn_in = attn_in.t()
+ if device.type != "cuda":
+ return
+ if _switch(_MTP_MOE_NVFP4_ENV):
+ if self._mtp_moe_nvfp4 is None:
+ self._mtp_moe_nvfp4 = _mtp_moe_nvfp4(self)
+ self._mtp = self._mtp._replace(**self._mtp_moe_nvfp4)
+ if _switch(_MTP_DENSE_FP8_ENV):
+ self._mtp_gemv8 = {
+ "fc_e": pack_from_bf16(w[f"{p}fc_e"]),
+ "fc_h": pack_from_bf16(w[f"{p}fc_h"]),
+ "attn_in": pack_from_bf16(attn_in),
+ "o": pack_from_bf16(w[f"{p}o_w"]),
+ "sh_gu": pack_from_bf16(w[f"{p}sh_gu"][: 2 * self.shared_inter]),
+ "sh_dn": pack_from_bf16(w[f"{p}sh_dn"]),
+ }
+
+ def _build_dense_copies(self, device) -> None:
+ """The trunk's dense FP8 copies (`_DENSE_FP8_ENV`) beside the bf16 weights the one-row
+ GEMV and fused paths keep reading, and the GDN in_proj's NVFP4 copies
+ (`_GDN_INPROJ_NVFP4_ENV`), with every decode variant compiled now; both only at tp1,
+ the layout they are certified at, and empty otherwise."""
+ w = self.w
+ one_rank = self.mapping.world_size == 1
+ self._dense8, self._gemv8, self._mid, self._gemv_nv = {}, {}, {}, {}
+ if one_rank and _switch(_DENSE_FP8_ENV):
+ for name in w:
+ if re.fullmatch(r"l\d+_(attn_in_w|o_w|in_proj_w|out_w)", name):
+ codes, scale = fp8_per_tensor_weight(w[name])
+ self._dense8[name] = (codes.t(), scale)
+ self._gemv8[name] = pack_from_bf16(w[name])
+ self._mid[name] = MidWeight(codes, scale)
+ elif re.fullmatch(r"l\d+_sh_gu", name):
+ # The shared expert's gate|up without its gate-logit row, which stays bf16.
+ self._gemv8[name] = pack_from_bf16(w[name][: 2 * self.shared_inter])
+ self._mid[name] = MidWeight(w[name][: 2 * self.shared_inter])
+ elif re.fullmatch(r"l\d+_sh_dn", name):
+ self._gemv8[name] = pack_from_bf16(w[name])
+ self._mid[name] = MidWeight(w[name])
+ elif re.fullmatch(r"l\d+_router", name):
+ self._mid[name] = MidWeight(w[name])
+ if device.type == "cuda":
+ _precompile_mid(self._mid)
+ if one_rank and _switch(_GDN_INPROJ_NVFP4_ENV) and device.type == "cuda":
+ for name in w:
+ if re.fullmatch(r"l\d+_in_proj_w", name):
+ self._gemv_nv[name] = nvfp4_pack_split_from_bf16(w[name])
+ if self._gemv_nv:
+ precompile_split(
+ next(iter(self._gemv_nv.values())),
+ _INPROJ_NVFP4_MAX_ROWS,
+ self.conv_dim,
+ _PAD_SLOT_ID,
+ tuple(range(2, self.max_verify_rows + 1)),
+ )
+
+ def _tp_sum_op(self):
+ """The sum that completes a block's per-rank partial output, settled once: under
+ attention TP the plain all-reduce over the TP group (`comm/allreduce`, op `NONE`),
+ which every rank gets the same bits back from; otherwise nothing to sum."""
+ if self.attn_tp_size == 1:
+ return lambda x: x
+ group = allreduce_group_args(self.mapping)
+
+ def tp_sum(x: torch.Tensor) -> torch.Tensor:
+ return allreduce(
+ x,
+ None,
+ None,
+ scale=None,
+ bias=None,
+ op=int(AllReduceFusionOp.NONE),
+ eps=self.eps,
+ trigger_completion_at_end=False,
+ **group,
+ )
+
+ return tp_sum
+
+ def _expert_sum_op(self):
+ """How a MoE block's routed experts become the rows this rank owns, settled once.
+ The op is `expert_sum(run, x, pad_rows)`, where `run(rows)` is this rank's expert
+ window over `rows` and `pad_rows` is `_dp_pad_rows`'s count.
+
+ Under attention TP (or at one rank) every rank holds every row: the window's partial
+ over `x`, completed by `tp_sum`. Under attention DP each rank holds its own rows, as
+ gpt-oss's and deepseek's default targets run it: pad them to the group's largest row
+ count, gather the group's rows (`comm/allgather`, in rank order), run the window over
+ all of them, reduce-scatter (`comm/reducescatter`), which sums each row's window
+ partials once and hands this rank its slice, and drop the padding. Routing on the
+ gathered rows makes every rank's routing of a row the same, so each of its top-k
+ experts falls in exactly one window."""
+ if self.dp_size == 1:
+ tp_sum = self.tp_sum
+
+ def window_sum(run, x: torch.Tensor, pad_rows: int) -> torch.Tensor:
+ return tp_sum(run(x))
+
+ return window_sum
+ group = self.dp_group
+ # Under MPI the plain symbols; with MPI disabled the `_pg` ones over the TP process
+ # group, as `allreduce_group_args` routes them.
+ process_group = self.mapping.tp_group_pg.boxed() if mpi_disabled() else None
+
+ def round_trip(run, x: torch.Tensor, pad_rows: int) -> torch.Tensor:
+ rows = x.shape[0]
+ padded = x if not pad_rows else nn.functional.pad(x, [0, 0, 0, pad_rows])
+ gathered = allgather(padded, None, group, process_group)
+ mine = reducescatter(run(gathered), None, group, process_group)
+ return mine if not pad_rows else mine[:rows]
+
+ return round_trip
+
+ def _dp_pad_rows(self, md, num_tokens: int) -> int:
+ """Rows this rank pads its MoE input with before the attention-DP round trip: up to
+ the group's largest per-rank token count, so both collectives run in their uniform
+ form. Zero without attention DP.
+
+ A decode CUDA graph pads every rank to the captured batch, and a prefill CUDA graph
+ every rank to one bucket (the engine then reports the bucket for every rank), so a
+ captured step's counts are equal and its padding is zero rows."""
+ if self.dp_size == 1:
+ return 0
+ counts = [int(n) for n in md.all_rank_num_tokens]
+ assert len(counts) == self.dp_size, (counts, self.dp_size)
+ assert counts[self.dp_rank] == num_tokens, (
+ f"rank {self.dp_rank} holds {num_tokens} rows but all_rank_num_tokens says "
+ f"{counts[self.dp_rank]} ({counts}); the collectives would gather the wrong split"
+ )
+ return max(counts) - num_tokens
+
+ def _index_rope_table(self, theta: float, device) -> torch.Tensor:
+ """The index compression's rope table at `theta` in the QSA kernels' layout:
+ `[max_positions, 2 * rotary_pairs]` fp32, cos block then sin block."""
+ exponent = torch.arange(0, self.rotary_dim, 2, dtype=torch.float64, device=device)
+ inv_freq = (1.0 / (theta ** (exponent / self.rotary_dim))).to(torch.float32)
+ positions = torch.arange(self.max_positions, dtype=torch.float32, device=device)
+ angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0)
+ return torch.cat([torch.cos(angles), torch.sin(angles)], dim=1).contiguous()
+
+ def _report_ple_residency(self) -> None:
+ """Check where the n-gram table ended up, and report it once: offload is
+ output-identical, so the pointer's memory kind read back from the driver is the
+ signal that cannot be faked."""
+ table = self.w["ple_table"]
+ kind = _weights.memory_kind(table.data_ptr())
+ expected = (
+ _weights.CUDA_MEMORY_TYPE_HOST
+ if self.ple_host_offload
+ else _weights.CUDA_MEMORY_TYPE_DEVICE
+ )
+ assert kind == expected, (
+ f"the n-gram table was declared with {_PLE_HOST_OFFLOAD_ENV}="
+ f"{'1' if self.ple_host_offload else '0'}, but the driver reports cudaMemoryType "
+ f"{kind} for the pointer the forward gathers from"
+ )
+ mib = 1024.0 * 1024.0
+ free, total = torch.cuda.mem_get_info()
+ print(
+ f"ple n-gram table: {tuple(table.shape)} {table.dtype}, "
+ f"{table.numel() * table.element_size() / mib:.1f} MiB, "
+ f"{'HOST-resident (device-addressable)' if self.ple_host_offload else 'DEVICE-resident'}"
+ f"; after load: allocated {torch.cuda.memory_allocated() / mib:.1f} MiB, "
+ f"driver-used {(total - free) / mib:.1f} MiB of {total / mib:.1f} MiB",
+ flush=True,
+ )
+
+ # ---------------------------------------------------------------- bind ---
+
+ def _bind_cache_manager(self, md, position_ids) -> None:
+ """Bind this rank's pool views of the step's KV cache manager, checked against the
+ geometry the weights were cut to.
+
+ Keyed on the manager, not a first-forward flag: capacity estimation builds a
+ throwaway manager and then the real one on the same module, and a view held across
+ the swap dangles silently. The contract checks re-run on every bind."""
+ cm = md.kv_cache_manager
+ assert cm is not None, "a paged KV cache is required"
+ assert not (self._binds and torch.cuda.is_current_stream_capturing()), (
+ f"the KV cache manager changed to {type(cm).__name__} during CUDA graph capture"
+ )
+ self._binds += 1
+ missing = [name for name in _QSA_FIELDS if not hasattr(md, name)]
+ assert not missing, (
+ f"metadata fields missing: {missing}; the engine did not select QSA metadata"
+ )
+ assert position_ids.dtype == torch.int32, position_ids.dtype
+ assert self.mrope_declared, "this target takes the mRoPE position layout"
+ assert position_ids.dim() == 3 and position_ids.shape[0] == _QSA_POSITION_AXES, (
+ f"expected the [{_QSA_POSITION_AXES}, 1, T] mRoPE position buffer, got "
+ f"{tuple(position_ids.shape)}"
+ )
+ pools = {row[0] for row in md.host_kv_cache_pool_mapping.tolist()}
+ assert pools <= {0, 1}, (
+ f"KV addressing over {len(pools)} pools is not certified; layer->pool ids "
+ f"{sorted(pools)}"
+ )
+ assert self.quant_mode != QM_NVFP4_KV_CACHE or md.use_paged_context_fmha, (
+ "an NVFP4 KV cache is read by the paged-context attention kernels only, which the "
+ "engine enables for chunked prefill, KV cache block reuse or speculative decoding"
+ )
+ g = _pool.pool_geometry(
+ self.model_config.pretrained_config,
+ self.mapping,
+ tokens_per_block=md.tokens_per_block,
+ kv_cache_quant_algo=self.kv_cache_quant_algo,
+ dtype=self.w["embed"].dtype,
+ )
+ # The manager's per-rank share is the one the weights were cut to.
+ assert (g.heads_kv, g.kv_head_start, g.num_k_heads, g.num_v_heads, g.conv_dim) == (
+ self.heads_kv,
+ self.kv_head_start,
+ self.num_k_heads,
+ self.num_v_heads,
+ self.conv_dim,
+ ), (
+ f"the KV cache manager's share {g} disagrees with the weights' heads "
+ f"(kv {self.heads_kv} from {self.kv_head_start}, GDN {self.num_k_heads} / "
+ f"{self.num_v_heads}, conv {self.conv_dim})"
+ )
+ # `qsa_topk_indices` rows are addressed as `block_topk`-wide rows.
+ assert md.qsa_topk_indices.shape[1] == self.block_topk, (
+ f"the engine sized the QSA selection buffer {md.qsa_topk_indices.shape[1]} wide "
+ f"and this target selects {self.block_topk} blocks per row"
+ )
+ mm = md.mamba_metadata
+ assert mm is not None and mm is not False, (
+ "attn_metadata.mamba_metadata is absent: the runtime allocated no recurrent-state pool"
+ )
+ mm_missing = [name for name in _MAMBA_FIELDS if not hasattr(mm, name)]
+ assert not mm_missing, f"mamba metadata fields missing: {mm_missing}"
+ rank_pools = _pool.take_pools(cm, g, self.attn_layers, self.gdn_layers, self.ple_layer)
+ # No sliding window: the targets pass `max_seq_len` as the window.
+ windows = cm.max_attention_window_vec
+ attn_windows = {windows[off] for off in rank_pools.pool_idx.values()}
+ assert attn_windows <= {None, md.max_seq_len}, (
+ f"the attention layers' pool windows are {sorted(map(repr, attn_windows))}; a "
+ "sliding window is not implemented here"
+ )
+ assert md.effective_beam_width == 1, "beam search is not implemented"
+ assert md.cache_indirection is None, "beam search is not implemented"
+ assert not md.is_cross, "cross attention is not implemented"
+ # The pool sizes the entries' page- and slot-addressed cells have to cover, reported
+ # on every rank (the executor's workers keep the info log level off).
+ first_attn, first_gdn = self.attn_layers[0], self.gdn_layers[0]
+ print(
+ f"modeling_v2 qwen4_exp default target: rank {self.mapping.rank} binds "
+ f"{rank_pools.kv[first_attn].k.shape[0]} K/V pages, "
+ f"{rank_pools.index_cache[first_attn].shape[0]} QSA index pages, "
+ f"{rank_pools.conv[first_gdn].shape[0]} recurrent-state slots and "
+ f"{rank_pools.ple_state[0].shape[0]} PLE slots (bind {self._binds})",
+ flush=True,
+ )
+ bound = _Bound(cm, rank_pools)
+ if self.spec_rollback is not None:
+ self.spec_rollback.set_index_planes(
+ [rank_pools.index_cache[i] for i in self.attn_layers]
+ )
+ if self.mtp_enabled:
+ bound = bound._replace(**self._bind_step_buffers(cm, g, mm, rank_pools))
+ self._bound = bound
+
+ def _bind_step_buffers(self, cm, g: _pool.PoolGeometry, mm, pools: _pool.RankPools) -> dict:
+ """The verify step's runtime-owned buffers of one manager: every verify row's conv
+ window goes to a per-step buffer that `update_mamba_states` promotes from, and the
+ recurrent state is the cached replay (a checkpoint per slot plus a compact history of
+ the tokens accepted since it). Returns the `_Bound` fields it adds."""
+ _pool.require(
+ cm,
+ ("is_speculative", "mamba_layer_cache", "get_replay_state_update_metadata"),
+ "binds the verify step's per-step state buffers and replay metadata",
+ )
+ assert cm.is_speculative(), (
+ f"the KV cache manager {type(cm).__name__} reports is_speculative() False, so it "
+ "allocated none of the per-step state buffers the verify path writes"
+ )
+ layer_caches = {i: cm.mamba_layer_cache(i) for i in self.gdn_layers}
+ int_conv = {i: cache.intermediate_conv_window for i, cache in layer_caches.items()}
+ int_idx = getattr(cm, "intermediate_state_indices", None)
+ g0 = self.gdn_layers[0]
+ _pool.check_step_buffers(g, int_conv[g0], pools.conv[g0], int_idx)
+ meta = cm.get_replay_state_update_metadata()
+ if meta is None:
+ raise NotImplementedError(
+ "the KV cache manager runs the legacy per-step recurrent verify path "
+ "(TRTLLM_USE_GDN_REPLAY=0), which this target does not assemble; leave "
+ "TRTLLM_USE_GDN_REPLAY unset for the cached replay"
+ )
+ assert not getattr(cm, "use_gdn_flashinfer_ring_replay", False), (
+ "the manager laid its replay history out for the FlashInfer ring kernel; this "
+ "target runs the double-buffered Triton cached replay"
+ )
+ assert layer_caches[g0].intermediate_ssm is None, (
+ "the cache manager runs the cached-replay verify path and also allocated a "
+ "per-step recurrent buffer; the runtime allocates one or the other"
+ )
+ ssm0 = pools.ssm[g0]
+ history, width = meta.replay_history_size, meta.replay_step_width
+ _pool.check_replay_state(
+ history, width, meta.prev_num_accepted_tokens, meta.cache_buf_idx, ssm0
+ )
+ old = {}
+ for i, cache in layer_caches.items():
+ old[i] = (cache.old_x, cache.old_B, cache.old_dt, cache.old_dA_cumsum)
+ _pool.check_replay_history(g, i, history, ssm0.shape[0], *old[i][:3])
+ replay = _Replay(old, meta.cache_buf_idx, meta.prev_num_accepted_tokens, history, width)
+ return dict(int_idx=int_idx, int_conv=int_conv, replay=replay)
+
+ def _bind_draft_cache_manager(self, md) -> None:
+ """Bind the draft layer to the manager its draft step sees: its own index plane, that
+ manager's position plane and its paged K/V.
+
+ Under attention TP that is the separate draft KV cache the worker swaps in for the
+ draft loop, whose pools hold the draft layer's pages and none of the trunk's;
+ otherwise the trunk's manager, which maps the draft layer as well. Keyed on the
+ manager like `_bind_cache_manager`, and never first bound under capture: the runner
+ runs every captured batch shape eagerly first."""
+ cm = md.kv_cache_manager
+ assert not torch.cuda.is_current_stream_capturing(), (
+ f"the draft layer is not bound to KV cache manager {type(cm).__name__} (first use, "
+ "or the manager changed since its last bind) and CUDA graph capture is active"
+ )
+ _pool.require(
+ cm,
+ ("get_index_k_buffer", "get_qsa_position_buffer", "get_buffers", "layer_offsets"),
+ "reads the draft layer's QSA index plane, position plane and paged K/V",
+ )
+ g = _pool.pool_geometry(
+ self.model_config.pretrained_config,
+ self.mapping,
+ tokens_per_block=md.tokens_per_block,
+ kv_cache_quant_algo=self.kv_cache_quant_algo,
+ dtype=self.w["embed"].dtype,
+ )
+ layer = self.mtp_layer_idx
+ index_cache = cm.get_index_k_buffer(layer)
+ position_cache = cm.get_qsa_position_buffer()
+ _pool.check_position_plane(g, position_cache)
+ kv = _pool.kv_planes(cm, g, layer, index_cache)
+ share_slot = None
+ if self.mtp_index_share:
+ # The metadata sizes its shared selection buffers from the same resolved knob,
+ # one row per sparse layer of whichever manager is bound.
+ assert md.sparse_metadata_params.mtp_index_share, (
+ "this target resolved IndexShare on but the engine's QSA metadata did not"
+ )
+ slots = cm.qsa_shared_topk_slots
+ _pool.check_shared_selection(md.qsa_shared_topk_indices, slots, layer, self.block_topk)
+ share_slot = slots[layer]
+ print(
+ f"modeling_v2 qwen4_exp default target: rank {self.mapping.rank} binds the draft layer to "
+ f"{type(cm).__name__} ({'the trunk manager' if cm is self._bound.cm else 'separate'}"
+ f"): {kv.k.shape[0]} K/V pages, {index_cache.shape[0]} QSA index pages"
+ f"{'' if share_slot is None else f', IndexShare slot {share_slot}'}",
+ flush=True,
+ )
+ self._draft_bound = _DraftBound(cm, index_cache, position_cache, kv, share_slot)
+
+ def _draft_step_bind(self, md) -> tuple[_DraftBound, int]:
+ """The draft layer's bind for this draft step's manager, and the layer's pool offset
+ in it (what `thop_attention` addresses), checked on every call: the worker swaps the
+ manager and the page tables in for the draft loop, and a page table still decoding
+ the trunk's offsets would read and write another pool's pages."""
+ cm = md.kv_cache_manager
+ assert cm is not None, "a paged KV cache is required"
+ offsets = cm.layer_offsets
+ assert self.mtp_layer_idx in offsets, (
+ f"the MTP module's attention is global layer index {self.mtp_layer_idx}, which the "
+ f"cache manager {type(cm).__name__} does not map (it maps {len(offsets)} layers); "
+ "the engine allocated no KV layer for the draft module"
+ )
+ layer_idx = offsets[self.mtp_layer_idx]
+ window = cm.max_attention_window_vec[layer_idx]
+ assert window is None or window == md.max_seq_len, (
+ f"the MTP layer's pool window is {window!r} but this call passes max_seq_len "
+ f"{md.max_seq_len} as the attention window; a sliding window is not implemented"
+ )
+ if self._draft_bound is None or cm is not self._draft_bound.cm:
+ self._bind_draft_cache_manager(md)
+ assert cm is self._bound.cm or md.qsa_block_table is getattr(
+ md, "qsa_draft_block_table", None
+ ), (
+ "the draft step is bound to a separate draft KV cache manager but qsa_block_table "
+ "is not that manager's table: the engine's QSA metadata did not swap its page "
+ "table for the draft forward"
+ )
+ return self._draft_bound, layer_idx
+
+ # -------------------------------------------------------------- forward ---
+
+ def _select_target(self, attn_metadata) -> "PrefillTarget | DecodeTarget":
+ """Which target runs this step: decode if `num_contexts == 0`, else prefill."""
+ return self._decode if attn_metadata.num_contexts == 0 else self._prefill
+
+ def begin_step(
+ self,
+ attn_metadata: AttentionMetadata,
+ input_ids: torch.IntTensor | None,
+ position_ids: torch.IntTensor | None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ lora_params: dict | None = None,
+ spec_metadata=None,
+ vision: tuple | None = None,
+ ) -> None:
+ """Check what holds for every step, bind the cache manager, choose the target and
+ let the prefill target prepare the step.
+
+ The shell calls it before `forward`, on every step: a breakable prefill CUDA graph
+ captures `forward` and replays it without calling it, so whatever a step reads from
+ the host is read here. `vision` is `(rows, mm_token_indices)` when the step carries
+ media."""
+ assert self._prefill is not None and self._decode is not None, (
+ "load_weights must run before forward"
+ )
+ assert position_ids is not None
+ assert isinstance(attn_metadata, TrtllmAttentionMetadata)
+ assert lora_params is None, "LoRA is not implemented by this target"
+ # The PLE layer hashes token ids, so they are required on every step.
+ assert input_ids is not None, "this target's PLE layer hashes token ids"
+ assert inputs_embeds is None, (
+ "this target computes its token embedding itself; an engine-supplied "
+ "inputs_embeds is not implemented"
+ )
+ # `spec_metadata` is present exactly when drafting. One-model MTP-Eagle captures no
+ # trunk layer, so the trunk owes no capture hook.
+ assert (spec_metadata is not None) == self.mtp_enabled, (
+ f"spec_metadata {'arrived' if spec_metadata is not None else 'is absent'} but the "
+ f"model was built with mtp_enabled={self.mtp_enabled}; the draft path and the "
+ "weights it loads are decided at construction"
+ )
+ if self._bound is None or attn_metadata.kv_cache_manager is not self._bound.cm:
+ self._bind_cache_manager(attn_metadata, position_ids)
+ target = self._select_target(attn_metadata)
+ if target is self._prefill:
+ target.prepare(attn_metadata, input_ids, position_ids, vision, spec_metadata)
+ else:
+ assert vision is None, "a generation-only step carries vision rows"
+ self._step_begun = True
+
+ def forward(
+ self,
+ attn_metadata: AttentionMetadata,
+ input_ids: torch.IntTensor | None = None,
+ position_ids: torch.IntTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ lora_params: dict | None = None,
+ **kwargs,
+ ):
+ """Run the step `begin_step` prepared on the target it chose. Under a breakable
+ prefill CUDA graph this is the body the engine captures; nothing here may read the
+ step from the host. Returns the logits input, or under speculation the decode
+ target's `(final, hc)` pair and the prefill target's one packed tensor
+ (`split_spec_output`)."""
+ assert self._step_begun, (
+ "the shell calls begin_step before every forward; this forward has no prepared step"
+ )
+ self._step_begun = False
+ return self._select_target(attn_metadata).forward(
+ attn_metadata, input_ids=input_ids, position_ids=position_ids, **kwargs
+ )
+
+ def split_spec_output(self, output) -> tuple[torch.Tensor, torch.Tensor]:
+ """`(final, hc)` out of what `forward` returned under speculation: `final` is the
+ `[T, hidden]` mix `lm_head` reads and `hc` the `[T, hc_count * hidden]` stream the
+ MTP module reads. The decode target returns the pair; the prefill target returns one
+ `[T, hidden + stream]` tensor (a breakable prefill graph's body returns one tensor),
+ of which the two are column views."""
+ if isinstance(output, tuple):
+ return output
+ final, hc = torch.split(output, [self.hidden, self.stream], -1)
+ return final, hc
+
+
+# --------------------------------------------------------------- weight views ---
+
+
+class _LayerWeights(NamedTuple):
+ """One trunk layer's weights as the views the ops read (GEMM operands are `.t()` views).
+ Attention fields are None on a Gated-DeltaNet layer and vice versa."""
+
+ ahc_norm: torch.Tensor
+ ahc_first: torch.Tensor
+ ahc_up: torch.Tensor
+ mhc_norm: torch.Tensor
+ mhc_first: torch.Tensor
+ mhc_up: torch.Tensor
+ attn_in: torch.Tensor | None
+ o: torch.Tensor | None
+ q_norm: torch.Tensor | None
+ k_norm: torch.Tensor | None
+ idx_qn: torch.Tensor | None
+ idx_kn: torch.Tensor | None
+ in_proj: torch.Tensor | None
+ out: torch.Tensor | None
+ conv_w: torch.Tensor | None
+ a_log: torch.Tensor | None
+ dt_bias: torch.Tensor | None
+ gnorm: torch.Tensor | None
+ router: torch.Tensor
+ moe_gscale: torch.Tensor
+ fc1_w: torch.Tensor
+ fc1_s: torch.Tensor
+ fc2_w: torch.Tensor
+ fc2_s: torch.Tensor
+ out1_scale: torch.Tensor
+ out1_gate_scale: torch.Tensor
+ out2_scale: torch.Tensor
+ # `[gate | up | gate logit | pad]` (`weights._cat_rows`), and the down projection.
+ sh_gug: torch.Tensor
+ sh_dn: torch.Tensor
+ # Its gate|up rows and its gate-logit row, as row views of `sh_gug`.
+ sh_gu: torch.Tensor
+ sh_gate: torch.Tensor
+ # The dense FP8 mode's copies (`_DENSE_FP8_ENV`), None without it: per-tensor FP8
+ # `(codes_t, scale)`, packed `fp8_gemv` copies and `mid_rows_mm` views; and the GDN
+ # in_proj's NVFP4 copy (`_GDN_INPROJ_NVFP4_ENV`).
+ attn_in8: tuple | None = None
+ o8: tuple | None = None
+ in_proj8: tuple | None = None
+ out8: tuple | None = None
+ attn_in_g: object | None = None
+ o_g: object | None = None
+ in_proj_g: object | None = None
+ out_g: object | None = None
+ sh_gu_g: object | None = None
+ sh_dn_g: object | None = None
+ in_proj_nv: object | None = None
+ attn_in_m: object | None = None
+ o_m: object | None = None
+ in_proj_m: object | None = None
+ out_m: object | None = None
+ sh_gu_m: object | None = None
+ sh_dn_m: object | None = None
+ router_m: object | None = None
+ # The routed experts' CuteDSL-layout copy and its scalars (`_moe_layer_scalars_cd`), which
+ # prefill-sized launches run on at one rank (`moe_prefill_copy`); None otherwise.
+ fc1_w_cd: torch.Tensor | None = None
+ fc1_s_cd: torch.Tensor | None = None
+ fc2_w_cd: torch.Tensor | None = None
+ fc2_s_cd: torch.Tensor | None = None
+ alpha1_cd: torch.Tensor | None = None
+ alpha2_cd: torch.Tensor | None = None
+ global_sf_cd: torch.Tensor | None = None
+ # On the CuteDSL backend the experts' own scalars (`_moe_layer_scalars_cd`); None on
+ # trtllm-gen.
+ alpha1: torch.Tensor | None = None
+ alpha2: torch.Tensor | None = None
+ global_sf: torch.Tensor | None = None
+ # `1 + ahc_norm` / `1 + mhc_norm` in fp32, contiguous (`_norm_w1`): the rstd mixers'
+ # operands (`_HC_RSTD_MIN_ROWS`); set by `PrefillTarget` only.
+ ahc_norm_w1: torch.Tensor | None = None
+ mhc_norm_w1: torch.Tensor | None = None
+
+
+def _layer_views(core: Qwen4ExpDefaultModelingV2) -> list[_LayerWeights]:
+ """Every trunk layer's `_LayerWeights`, read off the realized weights."""
+ w = core.w
+ if getattr(core, "trunk_moe_backend", "TRTLLM") == "CUTEDSL":
+ w_dtype, s_dtype = torch.float4_e2m1fn_x2, torch.uint8
+ else:
+ w_dtype, s_dtype = torch.uint8, torch.float8_e4m3fn
+ layers = []
+ for i in range(core.num_layers):
+ attn: dict = dict(attn_in=None, o=None, q_norm=None, k_norm=None, idx_qn=None, idx_kn=None)
+ gdn: dict = dict(in_proj=None, out=None, conv_w=None, a_log=None, dt_bias=None, gnorm=None)
+ if core.is_attn[i]:
+ attn = dict(
+ attn_in=w[f"l{i}_attn_in_w"].t(),
+ o=w[f"l{i}_o_w"].t(),
+ q_norm=w[f"l{i}_q_norm"],
+ k_norm=w[f"l{i}_k_norm"],
+ idx_qn=w[f"l{i}_idx_qn"],
+ idx_kn=w[f"l{i}_idx_kn"],
+ )
+ else:
+ a_log, dt_bias = core._gdn_gating[i]
+ gdn = dict(
+ in_proj=w[f"l{i}_in_proj_w"].t(),
+ out=w[f"l{i}_out_w"].t(),
+ # [conv_dim, 1, width] -> [conv_dim, width]: a view (unit width stride).
+ conv_w=w[f"l{i}_conv_w"].reshape(core.conv_dim, core.conv_kernel),
+ a_log=a_log,
+ dt_bias=dt_bias,
+ gnorm=w[f"l{i}_gnorm"],
+ )
+ scalars = core._moe_scalars[i]
+ cute_dsl: dict = {}
+ if getattr(core, "trunk_moe_backend", "TRTLLM") == "CUTEDSL":
+ scalars_cd = core._moe_scalars_cd[i]
+ cute_dsl = dict(
+ alpha1=scalars_cd["alpha1"],
+ alpha2=scalars_cd["alpha2"],
+ global_sf=scalars_cd["global_sf"],
+ )
+ elif getattr(core, "moe_prefill_copy", False):
+ scalars_cd = core._moe_scalars_cd[i]
+ cute_dsl = dict(
+ fc1_w_cd=w[f"l{i}_fc1_w_cd"],
+ fc1_s_cd=w[f"l{i}_fc1_s_cd"],
+ fc2_w_cd=w[f"l{i}_fc2_w_cd"],
+ fc2_s_cd=w[f"l{i}_fc2_s_cd"],
+ alpha1_cd=scalars_cd["alpha1"],
+ alpha2_cd=scalars_cd["alpha2"],
+ global_sf_cd=scalars_cd["global_sf"],
+ )
+ layers.append(
+ _LayerWeights(
+ ahc_norm=w[f"l{i}_ahc_norm"],
+ ahc_first=core._hc_first[(i, "ahc")].t(),
+ ahc_up=w[f"l{i}_ahc_up"].t(),
+ mhc_norm=w[f"l{i}_mhc_norm"],
+ mhc_first=core._hc_first[(i, "mhc")].t(),
+ mhc_up=w[f"l{i}_mhc_up"].t(),
+ **attn,
+ **gdn,
+ router=w[f"l{i}_router"].t(),
+ moe_gscale=scalars["global_scale"],
+ # The expert stacks are uint8 storage; each backend reads the packed nibbles
+ # and the block scales as its own dtypes.
+ fc1_w=w[f"l{i}_fc1_w"].view(w_dtype),
+ fc1_s=w[f"l{i}_fc1_s"].view(s_dtype),
+ fc2_w=w[f"l{i}_fc2_w"].view(w_dtype),
+ fc2_s=w[f"l{i}_fc2_s"].view(s_dtype),
+ out1_scale=scalars["output1_scale_scalar"],
+ out1_gate_scale=scalars["output1_scale_gate_scalar"],
+ out2_scale=scalars["output2_scale_scalar"],
+ sh_gug=w[f"l{i}_sh_gu"].t(),
+ sh_dn=w[f"l{i}_sh_dn"].t(),
+ sh_gu=w[f"l{i}_sh_gu"][: 2 * core.shared_inter].t(),
+ sh_gate=w[f"l{i}_sh_gu"][2 * core.shared_inter : 2 * core.shared_inter + 1].t(),
+ **{
+ field: getattr(core, copies, {}).get(f"l{i}_{name}")
+ for field, copies, name in (
+ ("attn_in8", "_dense8", "attn_in_w"),
+ ("o8", "_dense8", "o_w"),
+ ("in_proj8", "_dense8", "in_proj_w"),
+ ("out8", "_dense8", "out_w"),
+ ("attn_in_g", "_gemv8", "attn_in_w"),
+ ("o_g", "_gemv8", "o_w"),
+ ("in_proj_g", "_gemv8", "in_proj_w"),
+ ("out_g", "_gemv8", "out_w"),
+ ("sh_gu_g", "_gemv8", "sh_gu"),
+ ("sh_dn_g", "_gemv8", "sh_dn"),
+ ("in_proj_nv", "_gemv_nv", "in_proj_w"),
+ ("attn_in_m", "_mid", "attn_in_w"),
+ ("o_m", "_mid", "o_w"),
+ ("in_proj_m", "_mid", "in_proj_w"),
+ ("out_m", "_mid", "out_w"),
+ ("sh_gu_m", "_mid", "sh_gu"),
+ ("sh_dn_m", "_mid", "sh_dn"),
+ ("router_m", "_mid", "router"),
+ )
+ },
+ **cute_dsl,
+ )
+ )
+ return layers
+
+
+class _StackWeights(NamedTuple):
+ """The weights outside the layer loop: the stack-final mixer (this checkpoint's final
+ norm) and the PLE side path."""
+
+ fhc_norm: torch.Tensor
+ fhc_first: torch.Tensor
+ fhc_up: torch.Tensor
+ ple_key: torch.Tensor
+ ple_val: torch.Tensor
+ ple_nk_w: torch.Tensor
+ ple_nq_w: torch.Tensor
+ ple_nc_w: torch.Tensor
+ ple_conv_w: torch.Tensor
+ ple_mult: torch.Tensor
+ ple_vocab: torch.Tensor
+ ple_offs: torch.Tensor
+ ple_tscale: torch.Tensor
+ # `1 + fhc_norm` in fp32 (see `_LayerWeights.ahc_norm_w1`); set by `PrefillTarget` only.
+ fhc_norm_w1: torch.Tensor | None = None
+
+
+def _norm_w1(norm: torch.Tensor) -> torch.Tensor:
+ """A grouped norm's `(1 + w)` in fp32, contiguous: what `qwen4_exp_hc_up_gate_mean_rstd`
+ multiplies the scaled residual by (the `1 + w` `qwen4_exp_hc_inject_norm` forms per
+ element, bitwise)."""
+ return (1.0 + norm.float()).contiguous()
+
+
+def _stack_views(core: Qwen4ExpDefaultModelingV2) -> _StackWeights:
+ w = core.w
+ return _StackWeights(
+ fhc_norm=w["fhc_norm"],
+ fhc_first=core._hc_first["fhc"].t(),
+ fhc_up=w["fhc_up"].t(),
+ ple_key=w["ple_key_w"].t(),
+ ple_val=w["ple_val_w"].t(),
+ # The three PLE norm weights whole: the fused tail normalizes each stream inside one
+ # program.
+ ple_nk_w=w["ple_nk"],
+ ple_nq_w=w["ple_nq"],
+ ple_nc_w=w["ple_nc"],
+ ple_conv_w=w["ple_conv_w"],
+ ple_mult=w["ple_mult"],
+ ple_vocab=w["ple_vocab"],
+ ple_offs=w["ple_offs"],
+ # The table's dequantize scale at the gathered rows' rank (3), all 1s.
+ ple_tscale=w["ple_tscale"].reshape(1, 1, 1),
+ )
+
+
+class _MtpWeights(NamedTuple):
+ """The MTP module's weights at this rank's share (keyed at its layer index,
+ `num_hidden_layers`), as the views the ops read: GEMM operands are `.t()` views. Each
+ hyper-connection site is the trunk's: its whole `[stream]` grouped-norm weight, its
+ `[down; block_inject; 0]` first projection (`_fuse_hc_first_projections`) and its up
+ projection; `unit_inject` is the front end's per-stream `x_h + x_e` as the first site's
+ injection at unit weight."""
+
+ pre_e_norm: torch.Tensor
+ pre_h_norm: torch.Tensor
+ fc_e: torch.Tensor
+ fc_h: torch.Tensor
+ ahc_norm: torch.Tensor
+ ahc_first: torch.Tensor
+ ahc_up: torch.Tensor
+ mhc_norm: torch.Tensor
+ mhc_first: torch.Tensor
+ mhc_up: torch.Tensor
+ unit_inject: torch.Tensor
+ q: torch.Tensor # [hidden, 2 * q_width]: per head its q, then its output gate
+ k: torch.Tensor
+ v: torch.Tensor
+ o: torch.Tensor
+ q_norm: torch.Tensor
+ k_norm: torch.Tensor
+ idx: torch.Tensor
+ idx_qn: torch.Tensor
+ idx_kn: torch.Tensor
+ router: torch.Tensor
+ fc1_w: torch.Tensor
+ fc1_s: torch.Tensor
+ fc2_w: torch.Tensor
+ fc2_s: torch.Tensor
+ # `[gate | up | gate logit | pad]` (`weights._cat_rows`), as the trunk's, and the down
+ # projection.
+ sh_gug: torch.Tensor
+ sh_dn: torch.Tensor
+ fhc_norm: torch.Tensor
+ fhc_first: torch.Tensor
+ fhc_up: torch.Tensor
+ # The shared expert's gate|up rows and gate-logit row, row views of `sh_gug`.
+ sh_gu: torch.Tensor
+ sh_gate: torch.Tensor
+ # With `_MTP_MOE_NVFP4_ENV` (one rank), `fc1_w` .. `fc2_s` are the NVFP4 copy in the trunk's
+ # layout and these its activation global scale and the runner's three scalars.
+ moe_gscale: torch.Tensor | None = None
+ out1_scale: torch.Tensor | None = None
+ out1_gate_scale: torch.Tensor | None = None
+ out2_scale: torch.Tensor | None = None
+
+
+def _mtp_views(core: Qwen4ExpDefaultModelingV2) -> _MtpWeights:
+ """Built after `_fuse_hc_first_projections`, whose `"mtp_*"` tensors it reads."""
+ w = core.w
+ p = f"l{core.mtp_layer_idx}_mtp_"
+ dt = w[f"{p}fc_h"].dtype
+ return _MtpWeights(
+ pre_e_norm=w[f"{p}pre_e_norm"],
+ pre_h_norm=w[f"{p}pre_h_norm"],
+ fc_e=w[f"{p}fc_e"].t(),
+ fc_h=w[f"{p}fc_h"].t(),
+ ahc_norm=w[f"{p}ahc_norm"],
+ ahc_first=core._hc_first["mtp_ahc"].t(),
+ ahc_up=w[f"{p}ahc_up"].t(),
+ mhc_norm=w[f"{p}mhc_norm"],
+ mhc_first=core._hc_first["mtp_mhc"].t(),
+ mhc_up=w[f"{p}mhc_up"].t(),
+ unit_inject=torch.ones(
+ [core.max_num_tokens, core.hc_count], dtype=dt, device=w[f"{p}fc_h"].device
+ ),
+ q=w[f"{p}q_w"].t(),
+ k=w[f"{p}k_w"].t(),
+ v=w[f"{p}v_w"].t(),
+ o=w[f"{p}o_w"].t(),
+ q_norm=w[f"{p}q_norm"],
+ k_norm=w[f"{p}k_norm"],
+ idx=w[f"{p}idx_w"].t(),
+ idx_qn=w[f"{p}idx_qn"],
+ idx_kn=w[f"{p}idx_kn"],
+ router=w[f"{p}router"].t(),
+ fc1_w=w[f"{p}fc1_w"],
+ fc1_s=w[f"{p}fc1_s"],
+ fc2_w=w[f"{p}fc2_w"],
+ fc2_s=w[f"{p}fc2_s"],
+ sh_gug=w[f"{p}sh_gu"].t(),
+ sh_dn=w[f"{p}sh_dn"].t(),
+ sh_gu=w[f"{p}sh_gu"][: 2 * core.shared_inter].t(),
+ sh_gate=w[f"{p}sh_gu"][2 * core.shared_inter : 2 * core.shared_inter + 1].t(),
+ fhc_norm=w[f"{p}fhc_norm"],
+ fhc_first=core._hc_first["mtp_fhc"].t(),
+ fhc_up=w[f"{p}fhc_up"].t(),
+ )
+
+
+# ------------------------------------------------------------ shared blocks ---
+
+
+def _verify_rows(core, md, spec_metadata, num_seqs: int, num_rows: int) -> tuple[int, bool]:
+ """`(p, is_verify)` of a step's generation part: `num_rows` rows over `num_seqs`
+ generation sequences, `p` per sequence -- 1 without speculation, else
+ `runtime_draft_len + 1`. Host ints of the metadata, constants of a captured step."""
+ p = num_rows // num_seqs if num_seqs else 1
+ assert p * num_seqs == num_rows, (
+ f"{num_rows} generation rows do not divide over {num_seqs} generation sequences; a "
+ "ragged per-sequence draft length can be expressed neither to the attention call nor "
+ "to the multi-token recurrence"
+ )
+ assert core.mtp_enabled or p == 1, (
+ f"{num_rows} generation rows over {num_seqs} generation sequences with speculation "
+ "off; the base path implements one query token per generation sequence"
+ )
+ if spec_metadata is not None and num_seqs > 0:
+ declared = spec_metadata.runtime_draft_len + 1
+ assert p == declared, (
+ f"the batch carries {p} query rows per generation sequence but the worker's "
+ f"runtime_draft_len says {declared}; the per-step conv snapshots and the replay "
+ "history are sized from the second and written from the first"
+ )
+ _pool.check_verify_width(p)
+ _pool.check_replay_width(core._bound.replay, p)
+ return p, spec_metadata is not None and num_seqs > 0 and p > 1
+
+
+def _verify_state(core, mm, num_decodes: int) -> tuple:
+ """What a verify step's `_gdn_verify` reads besides the layer: the rows
+ `update_mamba_states` promotes the conv windows from, and the replay's work items. Work
+ items go to the kernel from `CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE` requests on, and the
+ kernel leaves the history commit to the manager's all-layer commit exactly when the
+ manager runs one after the step -- the stock Gated-DeltaNet layer's two conditions."""
+ assert core._bound.int_idx is not None, "the cache-manager bind has not run"
+ int_idx = core._bound.int_idx[:num_decodes]
+ partitioned = num_decodes >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE
+ all_layer = partitioned and bool(
+ getattr(core._bound.cm, "use_gdn_cached_replay_all_layer_commit", False)
+ )
+ work_items = mm.replay_work_items[:num_decodes] if partitioned else None
+ n_writes = mm.replay_n_writes if partitioned else None
+ return int_idx, work_items, n_writes, all_layer
+
+
+def _gdn_verify(core, mq, a, b, i, L, si_d, n, p, verify, mq_out=None) -> torch.Tensor:
+ """The recurrent half of a verify step: `n` generation sequences of `p` rows, written so
+ that the accepted prefix, unknown until the sampler has run, can be kept.
+
+ The conv writes each row's window to the runtime's per-step buffer
+ (`update_mamba_states` promotes the accepted one) besides advancing the pool; the cached
+ replay rebuilds each slot's state from its checkpoint and pending history, runs the `p`
+ rows and appends them to the history, so rejected rows cost nothing to undo. `mq` is the
+ `[n * p, conv_dim]` q|k|v column view of the input projection, `a` and `b` its raw gate
+ columns (column views too): the kernel applies `A_log`, `dt_bias` and the q/k L2 norm
+ itself. Over a bf16 recurrent pool the replay is the CuTe kernel (`gdn_cached_replay_cute`),
+ which reads `a` / `b` through their row stride; over an fp32 one the Triton kernel, which is
+ handed contiguous copies. Both take the same work items and leave the same commits to the
+ manager. Returns `[1, n * p, num_v_heads, head_v_dim]`.
+
+ `mq_out` is the conv output when the input projection's GEMV ran the conv step as its
+ epilogue (`fp8_gemv_conv1d_verify`, `nvfp4_split_gemv_conv1d_verify`; `mq` is then None)."""
+ int_idx, work_items, n_writes, all_layer = verify
+ pools = core._bound.pools
+ replay = core._bound.replay
+ ssm = pools.ssm[i]
+ cute = ssm.dtype == torch.bfloat16
+ pdl = _GDN_REPLAY_PDL and cute
+ if mq_out is None:
+ # [n*p, C] -> [n, C, p]: the multi-token conv reads channel-major.
+ x3 = torch.transpose(torch.reshape(mq, [n, p, core.conv_dim]), 1, 2)
+ conv_out = causal_conv1d_triton_update(
+ x=x3,
+ conv_state=pools.conv[i],
+ weight=L.conv_w,
+ bias=None,
+ activation="silu",
+ conv_state_indices=si_d,
+ pad_slot_id=_PAD_SLOT_ID,
+ # Snapshots, not the accepted-token roll.
+ num_accepted_tokens=None,
+ intermediate_conv_window=core._bound.int_conv[i],
+ intermediate_state_indices=int_idx,
+ launch_dependent_kernels=pdl,
+ )
+ # Back to token-major: one contiguous copy, which the kernel reads q/k/v from directly.
+ mq_out = torch.reshape(torch.transpose(conv_out, 1, 2), [n * p, core.conv_dim])
+ old_x, old_B, old_dt, old_dA = replay.old[i]
+ qv, kv, vv = torch.split(mq_out, [core.key_dim, core.key_dim, core.value_dim], 1)
+ if cute:
+ replay_op = gdn_cached_replay_cute
+ else:
+ replay_op = fused_recurrent_gated_delta_rule_cached_replay_update
+ a, b = a.contiguous(), b.contiguous()
+ out = replay_op(
+ q=torch.reshape(qv, [n, p, core.num_k_heads, core.head_k_dim]),
+ k=torch.reshape(kv, [n, p, core.num_k_heads, core.head_k_dim]),
+ v=torch.reshape(vv, [n, p, core.num_v_heads, core.head_v_dim]),
+ g=torch.reshape(a, [n, p, core.num_v_heads]),
+ beta=torch.reshape(b, [n, p, core.num_v_heads]),
+ ssm_states=ssm,
+ state_indices=si_d,
+ old_u=old_x,
+ old_k=old_B,
+ old_G=old_dt,
+ old_beta=old_dA,
+ cache_buf_idx=replay.cache_buf_idx,
+ prev_num_accepted_tokens=replay.prev_num_accepted,
+ history_size=replay.history,
+ scale=None,
+ use_qk_l2norm_in_kernel=True,
+ A_log=L.a_log,
+ dt_bias=L.dt_bias,
+ launch_with_pdl=pdl,
+ replay_work_items=work_items,
+ n_writes=n_writes,
+ block_v=None,
+ num_warps=None,
+ use_l2_state_cache=None,
+ use_l2_streaming_inputs=None,
+ use_l2_shared_inputs=None,
+ main_num_stages=None,
+ packed_qkv=mq_out,
+ use_all_layer_commit=all_layer,
+ output=None,
+ )
+ return torch.reshape(out, [1, n * p, core.num_v_heads, core.head_v_dim])
+
+
+def _position_key(position_ids: torch.Tensor) -> tuple:
+ """What identifies a step's positions buffer: its storage, dtype and `[3, N]` layout."""
+ view = torch.reshape(position_ids, [_QSA_POSITION_AXES, -1])
+ return (view.data_ptr(), view.dtype, tuple(view.shape), view.stride())
+
+
+def _draft_positions(position_ids: torch.Tensor) -> tuple:
+ """`(pos3, pos, coords)` of a draft step's mRoPE positions (`[3, 1, N]` from the trunk,
+ `[3, N]` from the worker's later steps): a fresh contiguous int32 `[3, N]` copy, its
+ coordinate-0 row (the index rope's) and the `[N, 3]` transpose the position plane
+ stores."""
+ assert position_ids.shape[0] == _QSA_POSITION_AXES, (
+ f"expected the mRoPE position layout with {_QSA_POSITION_AXES} coordinate rows, got "
+ f"{tuple(position_ids.shape)}"
+ )
+ pos3 = _copy_positions(position_ids)
+ return pos3, pos3[0], torch.transpose(pos3, 0, 1)
+
+
+def _draft_embedding(input_ids: torch.Tensor, embed_tokens: torch.Tensor) -> torch.Tensor:
+ """The MTP front end's token embedding. A draft loop's first call embeds every slot of
+ each generation request's accepted tokens, and under rejection sampling
+ `chain_speculative_sampling` fills the slots past a request's accepted count with -1
+ (the in-tree `Embedding` masks such ids for the same reason). Those rows feed nothing the
+ step keeps: attention is causal, the worker continues from each request's last accepted
+ row, and the next step's first draft call rewrites their cache positions. So a negative
+ id reads row 0 instead of past the table; every valid id embeds as before."""
+ return nn.functional.embedding(input_ids.clamp_min(0), embed_tokens)
+
+
+def _mtp_front(core, M: _MtpWeights, embed_tokens, input_ids, hidden_states) -> tuple:
+ """The draft layer's front end and its attention-side hyper-connection. `fc_embedding`
+ over the next token's normalized embedding is added to every stream of `fc_hidden`
+ (one `[hidden, hidden]` matrix applied to each normalized stream of the trunk's
+ residual): the trunk's site with that sum as its injection at unit weight (bitwise the
+ bf16 add), then the trunk's mixer. Returns `(x, mixed, inject)`: the layer's residual,
+ the mixed attention input and its injection weights."""
+ rows = hidden_states.shape[0]
+ assert rows <= M.unit_inject.shape[0], (
+ f"a draft step of {rows} rows exceeds max_num_tokens {M.unit_inject.shape[0]}"
+ )
+ e = _draft_embedding(input_ids, embed_tokens)
+ x_e = mtp_cublas_mm(flashinfer_gemma_rmsnorm(e, M.pre_e_norm, core.eps), M.fc_e)
+ _, h_n = qwen4_exp_hc_inject_norm(
+ hidden_states, None, None, M.pre_h_norm, core.eps, core.hc_count
+ )
+ x_h = torch.reshape(
+ mtp_cublas_mm(torch.reshape(h_n, [rows * core.hc_count, core.hidden]), M.fc_h),
+ [rows, core.stream],
+ )
+ x, normed = qwen4_exp_hc_inject_norm(
+ x_h, x_e, M.unit_inject[:rows], M.ahc_norm, core.eps, core.hc_count
+ )
+ mixed, inject = _hc_mixer(core, normed, M.ahc_first, M.ahc_up, True)
+ return x, mixed, inject
+
+
+def _mtp_qkv(core, M: _MtpWeights, mixed, pos3) -> tuple:
+ """The draft layer's q (with its interleaved output gate), k and v projections at this
+ rank's heads, packed into one compact `[T, qkv_width]` row and normalized and rotated in
+ place at the module's rope base. Returns `(qkv, gate, k, v)`: `gate` is the
+ `[T, heads_q, head_dim]` gate half of the q projection, `k` and v` views of `qkv`."""
+ rows = mixed.shape[0]
+ qg = mtp_cublas_mm(mixed, M.q)
+ qkv = torch.empty([rows, core.qkv_width], dtype=mixed.dtype, device=mixed.device)
+ q_dst, k_dst, v_dst = torch.split(
+ torch.reshape(qkv, [rows, core.heads_q + 2 * core.heads_kv, core.head_dim]),
+ [core.heads_q, core.heads_kv, core.heads_kv],
+ 1,
+ )
+ q_src, gate = torch.split(
+ torch.reshape(qg, [rows, core.heads_q, 2 * core.head_dim]),
+ [core.head_dim, core.head_dim],
+ 2,
+ )
+ q_dst.copy_(q_src)
+ k_dst.copy_(torch.reshape(mtp_cublas_mm(mixed, M.k), [rows, core.heads_kv, core.head_dim]))
+ v_dst.copy_(torch.reshape(mtp_cublas_mm(mixed, M.v), [rows, core.heads_kv, core.head_dim]))
+ mtp_fused_qk_norm_rope(
+ qkv,
+ num_heads_q=core.heads_q,
+ num_heads_k=core.heads_kv,
+ num_heads_v=core.heads_kv,
+ head_dim=core.head_dim,
+ rotary_dim=core.rotary_dim,
+ eps=core.eps,
+ q_weight=M.q_norm,
+ k_weight=M.k_norm,
+ base=core.mtp_theta,
+ is_neox=True,
+ position_ids=pos3,
+ is_qk_norm=True,
+ use_gemma=True,
+ use_mrope=True,
+ mrope_section1=core.mrope_section1,
+ mrope_section2=core.mrope_section2,
+ )
+ return qkv, gate, k_dst, v_dst
+
+
+def _mtp_index_branch(core, M: _MtpWeights, mixed, draft: _DraftBound, qsa) -> torch.Tensor:
+ """The draft layer's QSA index branch over its own planes, as a trunk attention layer
+ runs it: the raw key stored before anything normalizes it, the query's norm and rotation
+ (coordinate 0 alone, at the module's rope base), then the compression of every group this
+ step completes. Returns the normalized, rotated index query."""
+ req_idx, logical, block_table, tokens_per_block, pos, coords = qsa
+ rows = mixed.shape[0]
+ qk_index = mtp_cublas_mm(mixed, M.idx)
+ token_k = torch.split(
+ torch.reshape(qk_index, [rows, core.index_heads + core.index_kv_heads, -1]),
+ [core.index_heads, core.index_kv_heads],
+ 1,
+ )[1]
+ qsa_paged_index_store(
+ token_k,
+ coords,
+ draft.index_cache,
+ draft.position_cache,
+ req_idx,
+ logical,
+ block_table,
+ tokens_per_block,
+ )
+ index_scratch = torch.empty(
+ [rows, core.index_scratch_width], dtype=mixed.dtype, device=mixed.device
+ )
+ index_scratch[:, : core.index_width].copy_(qk_index)
+ fused_qk_norm_rope(
+ index_scratch,
+ num_heads_q=core.index_heads,
+ num_heads_k=core.index_kv_heads,
+ num_heads_v=core.index_kv_heads,
+ head_dim=core.index_head_dim,
+ rotary_dim=core.rotary_dim,
+ eps=core.eps,
+ q_weight=M.idx_qn,
+ k_weight=M.idx_kn,
+ base=core.mtp_theta,
+ is_neox=True,
+ position_ids=pos,
+ is_qk_norm=True,
+ use_gemma=True,
+ use_mrope=False,
+ mrope_section1=0,
+ mrope_section2=0,
+ )
+ qsa_prefill_compress(
+ logical_positions=logical,
+ request_indices=req_idx,
+ block_table=block_table,
+ index_cache=draft.index_cache,
+ position_cache=draft.position_cache,
+ k_norm_weight=M.idx_kn,
+ cos_sin=core._mtp_cos_sin,
+ eps=core.eps,
+ tokens_per_block=tokens_per_block,
+ compress_ratio=core.compress_ratio,
+ mrope_section=_QSA_MROPE_SECTION,
+ )
+ return _index_query(core, index_scratch)
+
+
+def _mtp_moe_block(core, M: _MtpWeights, x, pad_rows: int = 0, routed_only=False) -> torch.Tensor:
+ """The draft layer's MoE: renormalized top-k over the whole routing space, the FP8
+ block-scale routed experts over this rank's window (1x128 activation scales, pre-routed),
+ and the replicated bf16 shared expert in the trunk's form -- one GEMM forms gate|up and
+ the per-token gate logit -- combined as `routed + sigmoid(gate logit) * shared`. The
+ activation scales come back transposed `[hidden / 128, tokens]` and go straight to the
+ runner.
+
+ The window's partial is completed as the trunk's is: under attention TP summed over the
+ TP group; under attention DP by the round trip over the group's gathered rows
+ (`core.expert_sum`), `pad_rows` padding this rank's (`_draft_pad_rows`). The routing,
+ the quantization and the runner are row-wise and run on the rows the window sees, in
+ chunks of at most `_MTP_MOE_MAX_ROWS`: one chunk at a TP rank's rows, several over the
+ rows an attention-DP group gathers."""
+
+ def run(rows: torch.Tensor) -> torch.Tensor:
+ parts = []
+ for lo in range(0, rows.shape[0], _MTP_MOE_MAX_ROWS):
+ chunk = rows[lo : lo + _MTP_MOE_MAX_ROWS]
+ logits = cublas_mm(chunk, M.router)
+ topk_ids, topk_w = renorm_moe_routing_op(
+ logits, topk=core.topk, output_dtype=chunk.dtype
+ )
+ xq, xsf = fp8_quantize_1x128(chunk, use_ue8m0=_MTP_MOE_USE_UE8M0)
+ parts.append(
+ fp8_block_scale_moe_runner(
+ routing_logits=None,
+ routing_bias=None,
+ hidden_states=xq,
+ hidden_states_scale=xsf,
+ gemm1_weights=M.fc1_w,
+ gemm1_weights_scale=M.fc1_s,
+ gemm2_weights=M.fc2_w,
+ gemm2_weights_scale=M.fc2_s,
+ num_experts=core.num_experts,
+ top_k=core.topk,
+ num_fused_shared_experts=None,
+ n_group=None,
+ topk_group=None,
+ intermediate_size=core.moe_inter,
+ local_expert_offset=core.mtp_expert_offset,
+ local_num_experts=core.mtp_local_experts,
+ routed_scaling_factor=None,
+ routing_method_type=_ROUTING_METHOD_RENORM,
+ topk_weights=topk_w,
+ topk_ids=topk_ids,
+ act_type=_ACT_TYPE_SWIGLU,
+ gemm1_clamp_limit=None,
+ )
+ )
+ return parts[0] if len(parts) == 1 else torch.cat(parts, 0)
+
+ if M.moe_gscale is not None:
+ # At one rank under `_MTP_MOE_NVFP4_ENV` the experts are their NVFP4 copy.
+ routed = _mtp_routed_nvfp4(core, M, x)
+ elif core.dp_size == 1:
+ routed = core.tp_sum(run(x))
+ else:
+ routed = core.expert_sum(run, x, pad_rows)
+ if routed_only:
+ return routed
+ gug = cublas_mm(x, M.sh_gug)
+ shared = cublas_mm(silu_and_mul_rows(gug, core.shared_inter), M.sh_dn)
+ gate_logits = gug[:, 2 * core.shared_inter : 2 * core.shared_inter + 1]
+ return fused_sigmoid_gate_mul_add(routed, gate_logits, shared)
+
+
+def _mtp_routed_nvfp4(core, M: _MtpWeights, x) -> torch.Tensor:
+ """The draft layer's routed experts on their NVFP4 copy (`_mtp_moe_nvfp4`, one rank):
+ the activation quantized at the static FC1 input global scale, then the trunk's
+ logits-routed trtllm-gen runner over every expert, finalized."""
+ xq, xsf = fp4_quantize(x.contiguous(), M.moe_gscale, _SF_VEC, False, False)
+ return fp4_block_scale_moe_renorm_routed(
+ cublas_mm(x, M.router),
+ xq,
+ torch.reshape(xsf, [-1]).view(torch.float8_e4m3fn),
+ M.fc1_w,
+ M.fc1_s,
+ M.fc2_w,
+ M.fc2_s,
+ M.out1_scale,
+ M.out1_gate_scale,
+ M.out2_scale,
+ core.num_experts,
+ core.topk,
+ core.moe_inter,
+ )
+
+
+def _mtp_moe_nvfp4(core) -> dict:
+ """The draft layer's routed experts requantized to NVFP4 (`_MTP_MOE_NVFP4_ENV`), as the
+ `_MtpWeights` fields `_mtp_routed_nvfp4` reads; frees the FP8 copies. The runner's scalars
+ follow `_moe_layer_scalars` with `1 / g_in` and `1 / g_mid` standing for the checkpoint's
+ `input_scale`s and `1 / g_w` for its `weight_scale_2`: `out1_gate = 1 / (g_in * g_w1[e])`,
+ `out1 = out1_gate * g_mid`, `out2 = 1 / (g_mid * g_w2[e])`. The activation scales are static
+ (`_MTP_MOE_NVFP4_FC1_MARGIN` / `_FC2_MARGIN`)."""
+ w = core.w
+ p = f"l{core.mtp_layer_idx}_mtp_"
+ keys = [f"{p}{k}" for k in ("fc1_w", "fc1_s", "fc2_w", "fc2_s")]
+ fc1_w, fc1_s, fc2_w, fc2_s, g_w1, g_w2 = _weights.mtp_experts_to_nvfp4(*(w[k] for k in keys))
+ for k in keys:
+ w[k].data = w[k].data.new_empty(0)
+ torch.cuda.empty_cache()
+ trunk = range(core.num_layers)
+ to_global = _weights.NVFP4_AMAX_TO_GLOBAL
+ in_amax = max(float(w[f"l{i}_e_isc1"].max()) for i in trunk) * to_global
+ mid_amax = max(float(w[f"l{i}_e_isc2"].max()) for i in trunk) * to_global
+ inv_g_in = _MTP_MOE_NVFP4_FC1_MARGIN * in_amax / to_global
+ inv_g_mid = _MTP_MOE_NVFP4_FC2_MARGIN * mid_amax / to_global
+ gate1 = (inv_g_in / g_w1).contiguous()
+ return dict(
+ fc1_w=fc1_w,
+ fc1_s=fc1_s.view(torch.float8_e4m3fn),
+ fc2_w=fc2_w,
+ fc2_s=fc2_s.view(torch.float8_e4m3fn),
+ moe_gscale=torch.full((1,), 1.0 / inv_g_in, dtype=torch.float32, device=fc1_w.device),
+ out1_scale=(gate1 / inv_g_mid).contiguous(),
+ out1_gate_scale=gate1,
+ out2_scale=(inv_g_mid / g_w2).contiguous(),
+ )
+
+
+def _mtp_tail(core, M: _MtpWeights, x, attn, inject, pad_rows: int) -> torch.Tensor:
+ """The draft layer after its attention core: the gated output projection's partial summed
+ over the TP group and injected, the MoE-side hyper-connection, the MoE and its injection.
+ `attn` is the `[T, q_width]` gated attention output; `pad_rows` is the MoE round trip's
+ padding under attention DP (`_draft_pad_rows`). The MoE-side site is the trunk's
+ (`qwen4_exp_hc_inject_norm`, `_hc_mixer`)."""
+ x, normed = qwen4_exp_hc_inject_norm(
+ x, core.tp_sum(cublas_mm(attn, M.o)), inject, M.mhc_norm, core.eps, core.hc_count
+ )
+ mixed, inject = _hc_mixer(core, normed, M.mhc_first, M.mhc_up, True)
+ return qwen4_exp_post_inject(x, _mtp_moe_block(core, M, mixed, pad_rows), inject)
+
+
+def _decode_projection(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor:
+ """A generation-only step's dense bf16 projection at every layout but tp1 (the attention
+ input and output projections and the GDN output projection of `DecodeTarget`, which the
+ engine captures in CUDA graphs; tp1 takes `_one_rank_projection`).
+
+ Up to `_GEMV_MAX_ROWS` rows of the o / GDN out projection at a tp4/tep4 rank (K=1536) run
+ `bf16_gemv_pdl`: its weight streams into registers before it waits on the producer.
+ Everything else runs `low_m_bf16_mm`, whose table picks FlashInfer's low-M kernels at the
+ rows and shapes it names (a tp4/tep4 rank's, and the whole layer's that a dep4 / dep2 rank
+ holds) and cuBLAS elsewhere. Both launch with PDL per `_DECODE_PROJ_PDL` and read the
+ activation only after `griddepcontrol.wait`."""
+ if mat_a.shape[0] <= _GEMV_MAX_ROWS and mat_b.shape[0] == _GEMV_TP4_K:
+ return bf16_gemv_pdl(mat_a, mat_b, _DECODE_PROJ_PDL)
+ return low_m_bf16_mm(mat_a, mat_b, _DECODE_PROJ_PDL)
+
+
+def _one_rank_projection(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor:
+ """`_decode_projection` at tp1, where every decode projection's K is one `bf16_gemv_pdl`
+ tiles: one row on that GEMV, as up to `_GEMV_MAX_ROWS` rows of a projection at most
+ `_GEMV_NARROW_N` wide (the shared expert's gate logit, which `low_m_bf16_mm` would run on one
+ CTA); every other call on `low_m_bf16_mm`. Both launch with PDL."""
+ rows = mat_a.shape[0]
+ if mat_b.shape[0] in _GEMV_TILING and (
+ rows == 1 or (rows <= _GEMV_MAX_ROWS and mat_b.shape[1] <= _GEMV_NARROW_N)
+ ):
+ return bf16_gemv_pdl(mat_a, mat_b, _DECODE_PROJ_PDL)
+ return low_m_bf16_mm(mat_a, mat_b, _DECODE_PROJ_PDL)
+
+
+class _DraftRowCounts(NamedTuple):
+ """The group's per-rank row counts of one draft step, in the field `_dp_pad_rows` reads."""
+
+ all_rank_num_tokens: list
+
+
+def _draft_pad_rows(core, all_rank_num_tokens, rows: int) -> int:
+ """Rows a draft step pads its MoE input with before the attention-DP round trip:
+ `_dp_pad_rows` over the counts the spec worker hands this draft step -- every rank's
+ trunk tokens on the loop's first step, its request count on the later ones (attention-DP
+ dummy requests included) -- which it passes alongside the metadata rather than writing
+ into it. Zero without attention DP, and under a decode CUDA graph, which pads every rank
+ to the captured batch."""
+ if core.dp_size == 1:
+ return 0
+ assert all_rank_num_tokens is not None, (
+ "a draft step under attention DP came without the group's row counts; the MoE round "
+ "trip cannot pad this rank's rows to the group's largest count"
+ )
+ return core._dp_pad_rows(_DraftRowCounts(list(all_rank_num_tokens)), rows)
+
+
+def _hc_first(core, normed, first, with_inject):
+ """A hyper-connection mixer's first half: the first projection over
+ `[down; block_inject; 0]` (fp32 logits), then the SiLU and injection activations.
+ Returns `(low, inject)`; `inject` is None without `with_inject`. Up to
+ `_HC_TC_FIRST_MAX_ROWS` rows (a per-capture constant) it is one tcgen05 launch; larger
+ steps run cuBLAS's fp32 logits and the activation kernel."""
+ if normed.shape[0] <= _HC_TC_FIRST_MAX_ROWS:
+ return qwen4_exp_hc_decode_tc_first(
+ normed, first, core.hc_lowrank, core.hc_count, with_inject
+ )
+ return qwen4_exp_hc_lowrank_act(
+ cublas_mm(normed, first[:, : core.hc_first_cols], None, torch.float32),
+ core.hc_lowrank,
+ core.hc_count,
+ with_inject,
+ normed.dtype,
+ )
+
+
+def _hc_mixer(core, normed, first, up, with_inject):
+ """One hyper-connection mixer from the site's normed stream: `_hc_first`, then the up
+ projection with the sigmoid gate and the stream mean in its epilogue. Returns
+ `(mixed, inject)`; `inject` is None without `with_inject` (the stack-final mixer)."""
+ low, inject = _hc_first(core, normed, first, with_inject)
+ return qwen4_exp_hc_up_gate_mean(low, up, normed, core.hc_count), inject
+
+
+def _hc_moe_mixer(core, normed, L: _LayerWeights):
+ """The MoE-site mixer (`_hc_mixer` with injection). Under attention TP its up
+ projection's epilogue also writes the routed experts' NVFP4 input -- linear block
+ scales at the layer's FC1 input global scale -- which the expert window reads as is.
+ Under attention DP the window runs on the group's gathered bf16 rows and quantizes
+ them itself (`_routed_experts`). Returns `(mixed, inject, fp4)`, `fp4` the
+ `(xq, xsf)` pair or None."""
+ if core.dp_size > 1:
+ mixed, inject = _hc_mixer(core, normed, L.mhc_first, L.mhc_up, True)
+ return mixed, inject, None
+ low, inject = _hc_first(core, normed, L.mhc_first, True)
+ mixed, xq, xsf = qwen4_exp_hc_up_gate_mean_fp4(
+ low, L.mhc_up, normed, core.hc_count, L.moe_gscale
+ )
+ return mixed, inject, (xq, xsf)
+
+
+def _hc_mixer_rstd(core, hc, rstd, norm, norm_w1, first, up, with_inject):
+ """`_hc_mixer` from the site's residual `hc` and its `[T, C]` rstd
+ (`qwen4_exp_hc_inject_rstd`) instead of a stored normed stream: the first projection forms
+ `bf16(hc * rstd * (1 + norm))` on its A tiles, the up projection's epilogue
+ `hc * rstd * norm_w1` in fp32. Returns `(mixed, inject)`."""
+ low, inject = qwen4_exp_hc_first_rescale_tc(
+ hc, rstd, norm, first, core.hc_count, core.hc_lowrank, with_inject
+ )
+ return qwen4_exp_hc_up_gate_mean_rstd(low, up, hc, rstd, norm_w1, core.hc_count), inject
+
+
+def _hc_moe_mixer_rstd(core, hc, rstd, L: _LayerWeights):
+ """`_hc_moe_mixer` on the rstd form (see `_hc_mixer_rstd`): under attention TP the up
+ projection's epilogue also writes the routed experts' NVFP4 input; under attention DP the
+ window quantizes the gathered rows itself. Returns `(mixed, inject, fp4)`."""
+ if core.dp_size > 1:
+ mixed, inject = _hc_mixer_rstd(
+ core, hc, rstd, L.mhc_norm, L.mhc_norm_w1, L.mhc_first, L.mhc_up, True
+ )
+ return mixed, inject, None
+ low, inject = qwen4_exp_hc_first_rescale_tc(
+ hc, rstd, L.mhc_norm, L.mhc_first, core.hc_count, core.hc_lowrank, True
+ )
+ mixed, xq, xsf = qwen4_exp_hc_up_gate_mean_rstd_fp4(
+ low, L.mhc_up, hc, rstd, L.mhc_norm_w1, core.hc_count, L.moe_gscale
+ )
+ return mixed, inject, (xq, xsf)
+
+
+def _decode_router_mid(core, L: _LayerWeights):
+ """A decode step's router projection for layer `L`: `mid_rows_mm` over its view in the
+ dense FP8 mode's rows (`_MID_MM_ROWS["router"]`), else at tp1 the step's projection entry
+ (`core.decode_mm`) and at the other layouts `_decode_router`."""
+ fallback = core.decode_mm if core.mapping.world_size == 1 else _decode_router
+ if L.router_m is None:
+ return fallback
+ return lambda x, router: _dense_mm(x, router, None, fallback, "router", mid=L.router_m)
+
+
+def _decode_router(x, router):
+ """A decode step's router projection: up to `_ROUTER_LOW_M_MAX_ROWS` rows FlashInfer's
+ low-M kernels with PDL (`low_m_bf16_mm`), more rows cuBLAS. Decode only: those
+ kernels compile on the first call per row count, which the engine's warmup forward at
+ each captured batch makes before capture; a prefill step's row count is any, so prefill
+ keeps `cublas_mm`. Both allocate the logits fresh and dense, as the runner's routing
+ reads them from `data_ptr()`."""
+ if x.shape[0] <= _ROUTER_LOW_M_MAX_ROWS:
+ return low_m_bf16_mm(x, router, _DECODE_PROJ_PDL)
+ return cublas_mm(x, router)
+
+
+def _routed_experts(core, x, L: _LayerWeights, fp4=None, router_mm=cublas_mm) -> torch.Tensor:
+ """This rank's partial of the routed experts over every row of `x`: the router logits
+ over the whole routing space, and the trtllm-gen runner over this rank's expert window
+ entered with them -- its own Renormalize routing (bitwise `renorm_moe_routing_op` then
+ the pre-routed runner) -- on the NVFP4 activation with linear block scales at the
+ layer's FC1 input global scale, finalized. `fp4` is that activation as the MoE-site
+ mixer wrote it (`_hc_moe_mixer`, the rows of `x`), or None to quantize `x` here;
+ `router_mm` the step's router projection (`_decode_router` at decode). A slot
+ routed outside the window adds nothing, so the window outputs sum to the layer's. Under
+ the intermediate split (tp4) the window is every expert at this rank's
+ `expert_inter`-wide slice, and the slices' outputs sum to the layer's the same way.
+
+ Every step is row-wise, so `x` runs in chunks of at most the runner's certified rows (a
+ host decision on `T`): one chunk at a TP rank's rows, several over the rows an
+ attention-DP group gathers.
+
+ At one rank with the CuteDSL-layout copy (`moe_prefill_copy`) a launch of
+ `_CUTE_DSL_MOE_MIN_ROWS` to `_CUTE_DSL_MOE_MAX_ROWS` rows runs the CuteDSL MoE over it
+ instead, in one call. On the CuteDSL backend (one rank) every launch is
+ `_cute_dsl_routed`, in one call."""
+ rows = x.shape[0]
+ if core.trunk_moe_backend == "CUTEDSL":
+ xq, xsf = fp4_quantize(x, L.moe_gscale, _SF_VEC, False, False) if fp4 is None else fp4
+ routed_out = torch.empty([rows, core.hidden], dtype=x.dtype, device=x.device)
+ return _cute_dsl_routed(core, L, router_mm(x, L.router), xq, xsf, routed_out)
+ if L.fc1_w_cd is not None and _CUTE_DSL_MOE_MIN_ROWS <= rows <= _CUTE_DSL_MOE_MAX_ROWS:
+ xq, xsf = fp4_quantize(x, L.moe_gscale, _SF_VEC, False, False) if fp4 is None else fp4
+ return fp4_moe_renorm_cute_dsl(
+ router_mm(x, L.router),
+ xq,
+ xsf,
+ L.fc1_w_cd,
+ L.fc1_s_cd,
+ L.alpha1_cd,
+ L.fc2_w_cd,
+ L.fc2_s_cd,
+ L.alpha2_cd,
+ L.global_sf_cd,
+ core.num_experts,
+ core.topk,
+ memset_stream=core._moe_memset_stream,
+ )
+ sf_cols = x.shape[1] // _SF_VEC
+ parts = []
+ for lo in range(0, rows, _MOE_MAX_ROWS):
+ hi = min(rows, lo + _MOE_MAX_ROWS)
+ chunk = x[lo:hi]
+ logits = router_mm(chunk, L.router)
+ if fp4 is None:
+ xq, xsf = fp4_quantize(chunk, L.moe_gscale, _SF_VEC, False, False)
+ else:
+ xq, xsf = fp4[0][lo:hi], fp4[1][lo * sf_cols : hi * sf_cols]
+ parts.append(
+ fp4_block_scale_moe_renorm_routed(
+ logits,
+ xq,
+ # The runner reads the chunk's linear block scales flat.
+ torch.reshape(xsf, [-1]).view(torch.float8_e4m3fn),
+ L.fc1_w,
+ L.fc1_s,
+ L.fc2_w,
+ L.fc2_s,
+ L.out1_scale,
+ L.out1_gate_scale,
+ L.out2_scale,
+ core.num_experts,
+ core.topk,
+ core.expert_inter,
+ core.expert_offset,
+ core.local_experts,
+ )
+ )
+ return parts[0] if len(parts) == 1 else torch.cat(parts, 0)
+
+
+def _cute_dsl_routed(core, L, logits, xq, xsf, routed_out) -> torch.Tensor:
+ """The trunk's routed experts on the CuteDSL pipeline (one rank, every expert):
+ Renormalize routing, the sort into per-expert row tiles, FC1 gathering its rows from the
+ unpermuted NVFP4 activation `xq` / `xsf` (linear scales) with SwiGLU and the requantize to
+ the layer's one FC2 input scale fused, FC2, and the routing-weighted unpermute in a fixed
+ slot order, which writes every row of `routed_out` `[T, hidden]` and returns it."""
+ # fp32 weights: the dtype the unpermute scales by.
+ topk_ids, topk_w = renorm_moe_routing_op(logits, topk=core.topk, output_dtype=torch.float32)
+ # The sort pads each expert's (token, expert) pairs to whole tiles and emits the tile
+ # tables both GEMMs walk.
+ tile_expert, tile_limit, expanded_to_permuted, permuted_to_expanded, _, live_tiles = moe_sort(
+ token_selected_experts=topk_ids,
+ token_final_scales=topk_w,
+ num_experts=core.num_experts,
+ top_k=core.topk,
+ local_expert_offset=0,
+ local_num_experts=core.num_experts,
+ tile_tokens_dim=_CUTEDSL_TILE_SIZE,
+ )
+ h, h_sf = cute_dsl_nvfp4_gather_grouped_gemm_act_fusion(
+ input=xq.view(torch.float4_e2m1fn_x2),
+ weight=L.fc1_w,
+ input_scale=xsf.view(torch.uint8).view(xq.shape[0], core.hidden // _SF_VEC),
+ weight_scale=L.fc1_s,
+ alpha=L.alpha1,
+ tile_idx_to_group_idx=tile_expert,
+ tile_idx_to_mn_limit=tile_limit,
+ permuted_idx_to_expanded_idx=permuted_to_expanded,
+ num_non_exiting_tiles=live_tiles,
+ global_sf=L.global_sf,
+ num_experts=core.num_experts,
+ top_k=core.topk,
+ num_local_experts=core.num_experts,
+ local_expert_offset=0,
+ tile_size=_CUTEDSL_TILE_SIZE,
+ scaling_vector_size=_SF_VEC,
+ activation_type=int(ActivationType.Swiglu),
+ )
+ # FC2 writes one bf16 row per permuted row.
+ y = cute_dsl_nvfp4_grouped_gemm(
+ input=h,
+ weight=L.fc2_w,
+ input_scale=h_sf,
+ weight_scale=L.fc2_s,
+ alpha=L.alpha2,
+ tile_idx_to_group_idx=tile_expert,
+ num_non_exiting_tiles=live_tiles,
+ num_experts=core.num_experts,
+ top_k=core.topk,
+ num_local_experts=core.num_experts,
+ local_expert_offset=0,
+ tile_size=_CUTEDSL_TILE_SIZE,
+ output_dtype=torch.bfloat16,
+ scaling_vector_size=_SF_VEC,
+ )
+ return moe_unpermute_inplace(
+ permuted_input=y,
+ output=routed_out,
+ expanded_idx_to_permuted_idx=expanded_to_permuted,
+ topk_scales=topk_w,
+ )
+
+
+def _shared_expert(core, x, L: _LayerWeights, decode_gemm=None) -> tuple:
+ """The shared expert and its per-token gate logit over `x`: `(shared, gate_logits)`. A
+ decode step under the dense FP8 mode (`decode_gemm` its bf16 projection entry, the copies
+ present) runs gate|up and down on their FP8 / mid-rows copies and the gate logit on its own
+ bf16 row; otherwise one GEMM forms gate|up and the gate logit (`sh_gu`'s extra rows)."""
+ if decode_gemm is not None and L.sh_gu_g is not None:
+ gu = _dense_mm(x, L.sh_gu, None, decode_gemm, "sh_gu", L.sh_gu_g, mid=L.sh_gu_m)
+ shared = _dense_mm(
+ silu_and_mul_rows(gu, core.shared_inter),
+ L.sh_dn,
+ None,
+ decode_gemm,
+ "sh_dn",
+ L.sh_dn_g,
+ mid=L.sh_dn_m,
+ )
+ return shared, decode_gemm(x, L.sh_gate)
+ gug = cublas_mm(x, L.sh_gug)
+ shared = cublas_mm(silu_and_mul_rows(gug, core.shared_inter), L.sh_dn)
+ return shared, gug[:, 2 * core.shared_inter : 2 * core.shared_inter + 1]
+
+
+def _moe_block(
+ core, x, L: _LayerWeights, pad_rows: int, fp4=None, router_mm=cublas_mm, decode_gemm=None
+) -> tuple:
+ """One MoE layer: NVFP4 routed experts plus the replicated bf16 shared expert, left
+ uncombined. Returns `(routed, gate_logits, shared)`: the next `qwen4_exp_hc_inject_norm`
+ forms `routed + sigmoid(gate_logits) * shared` inside its kernel. `routed` is every
+ window's contribution to this rank's rows (`core.expert_sum`): under attention TP the
+ window partials' sum over the TP group, the same bits on every rank; under attention DP
+ the round trip over the group's gathered rows, `pad_rows` padding this rank's.
+
+ The shared expert and its per-token gate logit run on the core's side stream,
+ concurrently with the routed path, and join it before this returns. `x` is the bf16
+ MoE-site mixer output and `fp4` its NVFP4 form (`_hc_moe_mixer`): present only under
+ attention TP, where the window runs on `x`'s own rows. `router_mm` is the step's
+ router projection (`_routed_experts`). `decode_gemm` is a decode step's bf16 projection
+ entry, under which the shared expert reads the dense FP8 mode's copies when they exist."""
+ main = torch.cuda.current_stream()
+ side = core._side_stream
+ # Fork: the side stream reads `x` only after everything queued on `main` so far. What it
+ # allocates is consumed on `main` after the join below.
+ side.wait_stream(main)
+ with torch.cuda.stream(side):
+ shared, gate_logits = _shared_expert(core, x, L, decode_gemm)
+ routed = core.expert_sum(
+ lambda rows: _routed_experts(core, rows, L, fp4, router_mm), x, pad_rows
+ )
+ main.wait_stream(side)
+ return routed, gate_logits, shared
+
+
+def _attn_in_heads(core, proj):
+ """The main q, k and v heads of the fused attention projection's output `proj`
+ `[T, attn_in_width]`, as `[T, heads, head_dim]` views."""
+ q, k, v, _ = torch.split(
+ torch.reshape(proj, [proj.shape[0], core.attn_in_width // core.head_dim, core.head_dim]),
+ [core.heads_q, core.heads_kv, core.heads_kv, core.attn_in_tail_heads - core.heads_kv],
+ 1,
+ )
+ return q, k, v
+
+
+def _packed_q(core, qkv):
+ """The query heads `[T, heads_q, head_dim]` of the compact `[T, qkv_width]` packed QKV."""
+ heads = core.heads_q + 2 * core.heads_kv
+ return torch.reshape(qkv, [qkv.shape[0], heads, core.head_dim])[:, : core.heads_q]
+
+
+def _attention_rope(core, proj, index_scratch, L: _LayerWeights, pos3, pos) -> None:
+ """Both planes' Gemma norm + rope. The main q/k heads in place on the fused projection
+ row (interleaved mRoPE over the three position rows; everything after k is passed as v
+ heads, which only sets the row stride). The index projection's query and key are copied
+ into the index scratch and normalized and rotated there (coordinate 0 alone), so the raw
+ index key stays in `proj` for the index plane's store."""
+ fused_qk_norm_rope(
+ proj,
+ num_heads_q=core.heads_q,
+ num_heads_k=core.heads_kv,
+ num_heads_v=core.attn_in_tail_heads,
+ head_dim=core.head_dim,
+ rotary_dim=core.rotary_dim,
+ eps=core.eps,
+ q_weight=L.q_norm,
+ k_weight=L.k_norm,
+ base=core.theta,
+ is_neox=True,
+ position_ids=pos3,
+ is_qk_norm=True,
+ use_gemma=True,
+ use_mrope=True,
+ mrope_section1=core.mrope_section1,
+ mrope_section2=core.mrope_section2,
+ )
+ index_scratch[:, : core.index_width].copy_(
+ proj[:, core.index_col : core.index_col + core.index_width]
+ )
+ fused_qk_norm_rope(
+ index_scratch,
+ num_heads_q=core.index_heads,
+ num_heads_k=core.index_kv_heads,
+ num_heads_v=core.index_kv_heads,
+ head_dim=core.index_head_dim,
+ rotary_dim=core.rotary_dim,
+ eps=core.eps,
+ q_weight=L.idx_qn,
+ k_weight=L.idx_kn,
+ base=core.theta,
+ is_neox=True,
+ position_ids=pos,
+ is_qk_norm=True,
+ use_gemma=True,
+ use_mrope=False,
+ mrope_section1=0,
+ mrope_section2=0,
+ )
+
+
+def _index_query(core, index_scratch):
+ """The normalized, rotated index query `[T, index_heads, index_head_dim]`: the q slots
+ of the index scratch."""
+ rows = index_scratch.shape[0]
+ return torch.split(
+ torch.reshape(
+ index_scratch,
+ [rows, core.index_heads + 2 * core.index_kv_heads, core.index_head_dim],
+ ),
+ [core.index_heads, core.index_kv_heads, core.index_kv_heads],
+ 1,
+ )[0]
+
+
+def _index_plane(core, proj, i, L: _LayerWeights, qsa) -> None:
+ """The QSA index branch of attention layer `i`: store the raw index key and the
+ positions of `proj`'s rows into the paged index plane, then compress every group this
+ step completes. The raw key is stored before anything normalizes or rotates it, because
+ the compression applies both itself."""
+ req_idx, logical, block_table, tokens_per_block, coords = qsa
+ rows = proj.shape[0]
+ token_k = torch.split(
+ torch.reshape(
+ proj[:, core.index_col : core.index_col + core.index_scratch_width],
+ [rows, core.index_scratch_width // core.index_head_dim, core.index_head_dim],
+ ),
+ [core.index_heads, core.index_kv_heads, 1],
+ 1,
+ )[1]
+ pools = core._bound.pools
+ index_cache = pools.index_cache[i]
+ qsa_paged_index_store(
+ token_k,
+ coords,
+ index_cache,
+ pools.position_cache,
+ req_idx,
+ logical,
+ block_table,
+ tokens_per_block,
+ )
+ qsa_prefill_compress(
+ logical_positions=logical,
+ request_indices=req_idx,
+ block_table=block_table,
+ index_cache=index_cache,
+ position_cache=pools.position_cache,
+ k_norm_weight=L.idx_kn,
+ cos_sin=core._cos_sin,
+ eps=core.eps,
+ tokens_per_block=tokens_per_block,
+ compress_ratio=core.compress_ratio,
+ mrope_section=_QSA_MROPE_SECTION,
+ )
+
+
+def _select(core, scores, row_starts, row_ends, topk_indices, live_cols=None) -> None:
+ """Each row's canonical top-`block_topk` groups of its window into `topk_indices`: the
+ groups strictly above the budget's last score, then the highest-index ties, ascending,
+ `-1` last -- a function of the scores alone, which the expansion and the sparse GQA
+ reduce in. QSA's row starts are zero, so window-local and table indices agree.
+
+ The exact top-k and the canonical rewrite, two launches, which serve a draft step whose
+ rows are the rejection kernel's accepted-token runs at every width the engine's graph
+ buckets use.
+
+ The rewrite reads only the columns of each row's window, and `scores` is as wide as the
+ block table; a caller that knows a host-side bound on the windows passes it as
+ `live_cols` so the rewrite reads only those columns and returns the same rows."""
+ qsa_topk_select(scores, row_starts, row_ends, topk_indices, index_topk=core.block_topk)
+ if live_cols is not None:
+ scores = scores[:, :live_cols]
+ qsa_canonical_selection(scores, row_starts, row_ends, topk_indices)
+
+
+def _live_score_columns(core, md, batch: int) -> int:
+ """A host-side bound on every row's selection window in a step with context rows: the
+ longest sequence's compressed length, with the verify width as slack for the draft rows
+ that sit past the host's KV lengths."""
+ longest = int(md.kv_lens_runtime[:batch].max()) if batch > 0 else 0
+ return (longest + _MAX_VERIFY_WIDTH) // core.compress_ratio + 1
+
+
+def _selection(
+ core, scores, visible_blocks, logical, seq_lens_q, topk_indices, row_starts, out, live_cols=None
+):
+ """The selected token indices of `scores`' rows into `out` `[rows, token_topk +
+ compress_ratio - 1]`: the canonical block top-k and its expansion in one launch, or, for
+ a score table that kernel does not serve, the selection and the expansion apart (the
+ rewrite bounded by `live_cols`, `_select`)."""
+ if _topk_expand_supports(scores):
+ qsa_topk_expand_decode(
+ scores,
+ visible_blocks,
+ logical,
+ seq_lens_q,
+ topk_indices,
+ core.compress_ratio,
+ core.token_topk,
+ out=out,
+ )
+ return
+ _select(core, scores, row_starts, visible_blocks, topk_indices, live_cols)
+ out.copy_(
+ qsa_expand_block_indices(
+ topk_indices,
+ logical,
+ seq_lens_q,
+ compress_ratio=core.compress_ratio,
+ token_topk=core.token_topk,
+ )
+ )
+
+
+def _store_kv(core, kv, k, v, request_indices, logical_positions, block_table, tokens_per_block):
+ """Append this step's K/V rows to the layer's paged pool: a plain move over a bf16 or fp8
+ pool (an fp8 pool takes e4m3 unscaled, as the dense op does at the unit KV scales), the
+ quantizing store over an NVFP4 one, at the dense op's second-level scales."""
+ if kv.k_scale is None:
+ qsa_paged_kv_store(
+ k, v, kv.k, kv.v, request_indices, logical_positions, block_table, tokens_per_block
+ )
+ return
+ qsa_paged_kv_store_nvfp4(
+ k,
+ v,
+ kv.k,
+ kv.v,
+ kv.k_scale,
+ kv.v_scale,
+ request_indices,
+ logical_positions,
+ block_table,
+ tokens_per_block,
+ core.kv_scale_orig_quant,
+ )
+
+
+def _nvfp4_context_groups(pages: list, rows: list, max_pages: int) -> list:
+ """Consecutive context requests in groups of at most `max_pages` pages, a request with more
+ pages alone: `(first request, end request, first row, end row)` a group, where request `r`
+ has `pages[r]` pages and the next `rows[r]` rows."""
+ groups = []
+ first = first_row = row = listed = 0
+ for r, (n_pages, n_rows) in enumerate(zip(pages, rows)):
+ if r > first and listed + n_pages > max_pages:
+ groups.append((first, r, first_row, row))
+ first, first_row, listed = r, row, 0
+ listed += n_pages
+ row += n_rows
+ if pages:
+ groups.append((first, len(pages), first_row, row))
+ return groups
+
+
+def _nvfp4_context_plan(core, md, req_idx, block_table, tpb: int) -> tuple:
+ """The pages a step's context requests read, for the bf16 row-tile pass over an NVFP4 pool:
+ request `r` lists its first `ceil(kv_len / tpb)` table pages. The requests come in groups
+ that fit `_NVFP4_CONTEXT_SCRATCH_BYTES` (`_nvfp4_context_groups`); within a group request `r`
+ takes scratch pages `offsets[r]` on. Returns `(groups, table, k_scratch, v_scratch)`: per
+ group `(first row, end row, slots, counts, offsets)`, the operands `qsa_nvfp4_pages_to_bf16`
+ fills per layer and group; `table` (all -1 here) gets the requests' slot rows, the scratch
+ pools `[pages, Hkv, tpb, D]` bf16 the largest group."""
+ num_contexts = md.num_contexts
+ device = block_table.device
+ pages = ((md.kv_lens_runtime[:num_contexts] + (tpb - 1)) // tpb).tolist()
+ rows = md.seq_lens[:num_contexts].tolist()
+ page_bytes = 2 * core.heads_kv * tpb * core.head_dim * 2 # K and V, bf16
+ groups = _nvfp4_context_groups(pages, rows, max(1, _NVFP4_CONTEXT_SCRATCH_BYTES // page_bytes))
+ lens = md.seq_lens_cuda[:num_contexts].long()
+ slots = req_idx[torch.cumsum(lens, 0) - lens]
+ counts = torch.div(
+ md.kv_lens_cuda_runtime[:num_contexts].long() + (tpb - 1), tpb, rounding_mode="floor"
+ )
+ planned = []
+ for first, end, first_row, end_row in groups:
+ group_counts = counts[first:end]
+ planned.append(
+ (
+ first_row,
+ end_row,
+ slots[first:end],
+ group_counts,
+ torch.cumsum(group_counts, 0) - group_counts,
+ )
+ )
+ listed = max(sum(pages[first:end]) for first, end, _, _ in groups)
+ table = torch.full_like(block_table, -1)
+ shape = (listed, core.heads_kv, tpb, core.head_dim)
+ k_scratch = torch.empty(shape, dtype=torch.bfloat16, device=device)
+ v_scratch = torch.empty(shape, dtype=torch.bfloat16, device=device)
+ return planned, table, k_scratch, v_scratch
+
+
+def _sparse_gqa_nvfp4_context(
+ core, kv, plan, q, block_table, selected, req_idx, logical, tpb, num_ctx, out, rows_per_request
+) -> torch.Tensor:
+ """Sparse GQA of a step with context rows over an NVFP4 pool at one rank, into `out` `[rows,
+ q_width]`. With a plan (`_nvfp4_context_plan`) the first `num_ctx` rows, the context rows,
+ take the row-tile pass over the plan's scratch, filled here from this layer's pool a group of
+ requests at a time; the other rows, and every row without a plan, go split-K over the pool
+ itself. The generation rows after the context rows come `rows_per_request` a request: more
+ than one (a verify step's) takes the request-grouped entry for them on their own."""
+ rows = q.shape[0]
+ head = 0
+ if plan is not None:
+ groups, table, k_scratch, v_scratch = plan
+ assert groups[-1][1] == num_ctx, (groups[-1][1], num_ctx)
+ for first_row, end_row, slots, counts, offsets in groups:
+ qsa_nvfp4_pages_to_bf16(
+ kv.k,
+ kv.v,
+ kv.k_scale,
+ kv.v_scale,
+ core.kv_scale_quant_orig,
+ block_table,
+ slots,
+ counts,
+ offsets,
+ table,
+ k_scratch,
+ v_scratch,
+ )
+ context = one_rank_sparse_gqa_rowtile(
+ q=q[first_row:end_row],
+ k_cache=k_scratch,
+ v_cache=v_scratch,
+ block_table=table,
+ selected_tokens=selected[first_row:end_row],
+ request_indices=req_idx[first_row:end_row],
+ tokens_per_block=tpb,
+ softmax_scale=core.softmax_scale,
+ query_positions=logical[first_row:end_row],
+ compress_ratio=core.compress_ratio,
+ )
+ out[first_row:end_row].copy_(
+ torch.reshape(context, [end_row - first_row, core.q_width])
+ )
+ head = num_ctx
+ elif (
+ rows_per_request > 1
+ and 0 < num_ctx < rows
+ and _verify_takes_grouped(
+ q[num_ctx:], kv.k, block_table, selected[num_ctx:], tpb, rows_per_request
+ )
+ ):
+ context = _sparse_gqa(
+ core, kv, q[:num_ctx], block_table, selected[:num_ctx], req_idx[:num_ctx], tpb, None
+ )
+ out[:num_ctx].copy_(torch.reshape(context, [num_ctx, core.q_width]))
+ head = num_ctx
+ if head < rows:
+ rest = _sparse_gqa(
+ core,
+ kv,
+ q[head:],
+ block_table,
+ selected[head:],
+ req_idx[head:],
+ tpb,
+ None,
+ rows_per_request=rows_per_request if head == num_ctx else 1,
+ )
+ out[head:rows].copy_(torch.reshape(rest, [rows - head, core.q_width]))
+ return out
+
+
+def _fp8_context_pages(md, tpb: int) -> int:
+ """Pool pages a step's context requests hold, `ceil(kv_len / tpb)` each, from the host's
+ KV lengths: the bf16 scratch `_sparse_gqa_rowtile` widens them into over an fp8 pool."""
+ lens = md.kv_lens_runtime[: md.num_contexts]
+ return int(((lens + (tpb - 1)) // tpb).sum())
+
+
+def _sparse_gqa_rowtile(
+ core, kv, q, block_table, selected, request_indices, tpb, query_positions, fp8_pages
+):
+ """The row-tile pass over a context launch, or None where it does not serve it. A bf16
+ pool is read in place. Over an fp8 pool, request slot `s`'s first `ceil((p_s + 1) / tpb)`
+ pages (`p_s` the highest position the launch's rows hold there) are widened, exactly, into
+ a bf16 scratch pool addressed through a table of `block_table`'s shape; `fp8_pages` is
+ their host count (`_fp8_context_pages`), read back from the device when not given. A
+ scratch past `_FP8_ROWTILE_MAX_SCRATCH` bytes declines (None)."""
+ k_cache, v_cache, table = kv.k, kv.v, block_table
+ if kv.k.dtype == torch.float8_e4m3fn:
+ # Each slot's highest position among the launch's rows, -1 for a slot it does not
+ # hold: a compare-and-max over `[slots, rows]` (a scatter max would serialize every
+ # row of a request on one address).
+ slots = torch.arange(block_table.shape[0], dtype=request_indices.dtype, device=q.device)
+ held = request_indices[None, :] == slots[:, None]
+ last = torch.where(held, query_positions[None, :].long(), -1).amax(1)
+ counts = torch.div(last + tpb, tpb, rounding_mode="floor")
+ if fp8_pages is None:
+ fp8_pages = int(counts.sum())
+ # K and V pages, bf16.
+ if fp8_pages * kv.k[0].numel() * 4 > _FP8_ROWTILE_MAX_SCRATCH:
+ return None
+ k_cache = torch.empty([fp8_pages, *kv.k.shape[1:]], dtype=q.dtype, device=q.device)
+ v_cache = torch.empty_like(k_cache)
+ table = torch.full_like(block_table, -1)
+ qsa_fp8_pages_to_bf16(
+ kv.k,
+ kv.v,
+ block_table,
+ slots.int(),
+ counts,
+ torch.cumsum(counts, 0) - counts,
+ table,
+ k_cache,
+ v_cache,
+ )
+ elif kv.k.dtype != q.dtype:
+ return None
+ return qsa_sparse_gqa_rowtile(
+ q=q,
+ k_cache=k_cache,
+ v_cache=v_cache,
+ block_table=table,
+ selected_tokens=selected,
+ request_indices=request_indices,
+ tokens_per_block=tpb,
+ softmax_scale=core.softmax_scale,
+ query_positions=query_positions,
+ compress_ratio=core.compress_ratio,
+ )
+
+
+def _sparse_gqa(
+ core,
+ kv,
+ q,
+ block_table,
+ selected,
+ request_indices,
+ tpb,
+ query_positions,
+ fp8_pages=None,
+ output_gate=None,
+ rows_per_request=1,
+):
+ """Sparse GQA over the layer's paged pool. `query_positions` is a context row's bound on
+ its scan (the prefill pair with `compress_ratio`), or None for decode rows. A context
+ launch of `_ROWTILE_MIN_ROWS` rows or more over a bf16 or fp8 pool takes the row-tile
+ pass (`fp8_pages`: see `_sparse_gqa_rowtile`). Over an NVFP4 pool context rows take the
+ prefill kernel and decode rows the split-K decode kernel."""
+ if kv.k_scale is None:
+ if query_positions is not None and q.shape[0] >= _ROWTILE_MIN_ROWS:
+ out = _sparse_gqa_rowtile(
+ core, kv, q, block_table, selected, request_indices, tpb, query_positions, fp8_pages
+ )
+ if out is not None:
+ return out
+ return qsa_paged_sparse_gqa(
+ q=q,
+ k_cache=kv.k,
+ v_cache=kv.v,
+ block_table=block_table,
+ selected_tokens=selected,
+ request_indices=request_indices,
+ tokens_per_block=tpb,
+ softmax_scale=core.softmax_scale,
+ query_positions=query_positions,
+ compress_ratio=None if query_positions is None else core.compress_ratio,
+ output_gate=None,
+ )
+ if query_positions is None:
+ # `output_gate` (`[rows, Hq * D]`) returns the gated output; `rows_per_request` > 1 is a
+ # verify step's rows, consecutive per request, whose K/V the verify kernel loads once.
+ if rows_per_request > 1:
+ return qsa_paged_sparse_gqa_nvfp4_verify(
+ q,
+ kv.k,
+ kv.v,
+ kv.k_scale,
+ kv.v_scale,
+ core.kv_scale_quant_orig,
+ block_table,
+ selected,
+ request_indices,
+ tpb,
+ core.softmax_scale,
+ rows_per_request,
+ output_gate=output_gate,
+ )
+ return qsa_paged_sparse_gqa_nvfp4_decode(
+ q,
+ kv.k,
+ kv.v,
+ kv.k_scale,
+ kv.v_scale,
+ core.kv_scale_quant_orig,
+ block_table,
+ selected,
+ request_indices,
+ tpb,
+ core.softmax_scale,
+ output_gate,
+ )
+ return qsa_paged_sparse_gqa_nvfp4(
+ q,
+ kv.k,
+ kv.v,
+ kv.k_scale,
+ kv.v_scale,
+ core.kv_scale_quant_orig,
+ block_table,
+ selected,
+ request_indices,
+ tpb,
+ core.softmax_scale,
+ query_positions=query_positions,
+ compress_ratio=core.compress_ratio,
+ )
+
+
+def _copy_positions(position_ids: torch.Tensor) -> torch.Tensor:
+ """A fresh contiguous int32 `[3, N]` copy of the step's mRoPE positions (the engine hands
+ a slice of a larger buffer, whose row stride the rope rejects)."""
+ pos3 = torch.empty(
+ [_QSA_POSITION_AXES, position_ids.numel() // _QSA_POSITION_AXES],
+ dtype=torch.int32,
+ device=position_ids.device,
+ )
+ pos3.copy_(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1]))
+ return pos3
+
+
+def _step_args(md: TrtllmAttentionMetadata, num_contexts, num_ctx_tokens) -> dict:
+ """thop_attention's batch state off the prepared metadata. Tensors are engine-owned
+ persistent buffers refreshed in place; Python ints are per-capture constants."""
+ return dict(
+ sequence_length=md.kv_lens_cuda_runtime,
+ host_past_key_value_lengths=md.kv_lens_runtime,
+ host_total_kv_lens=md.host_total_kv_lens,
+ context_lengths=md.prompt_lens_cuda_runtime,
+ host_context_lengths=md.prompt_lens_cpu_runtime,
+ host_request_types=md.host_request_types_runtime,
+ kv_cache_block_offsets=md.kv_cache_block_offsets,
+ host_kv_cache_pool_pointers=md.host_kv_cache_pool_pointers,
+ host_kv_cache_pool_mapping=md.host_kv_cache_pool_mapping,
+ workspace=md.effective_workspace,
+ tokens_per_block=md.tokens_per_block,
+ max_num_requests=md.max_num_requests,
+ max_context_length=md.max_context_length,
+ max_seq_len=md.max_seq_len,
+ # No sliding window: the full admitted length on every attention layer.
+ attention_window_size=md.max_seq_len,
+ num_contexts=num_contexts,
+ num_ctx_tokens=num_ctx_tokens,
+ trtllm_gen_jit_warmup=md.trtllm_gen_jit_warmup,
+ use_paged_context_fmha=md.use_paged_context_fmha,
+ beam_width=md.effective_beam_width,
+ cache_indirection=md.cache_indirection,
+ block_ids_per_seq=md.block_ids_per_seq,
+ max_context_q_len_override=md.max_context_q_len_override,
+ is_cross=md.is_cross,
+ is_spec_decoding_enabled=md.is_spec_decoding_enabled,
+ use_spec_decoding=md.use_spec_decoding,
+ is_spec_dec_tree=md.is_spec_dec_tree,
+ spec_decoding_generation_lengths=md.spec_decoding_generation_lengths,
+ spec_decoding_position_offsets=get_spec_decoding_position_offsets(md),
+ spec_decoding_packed_mask=md.spec_decoding_packed_mask,
+ spec_decoding_bl_tree_mask_offset=md.spec_decoding_bl_tree_mask_offset,
+ spec_decoding_bl_tree_mask=md.spec_decoding_bl_tree_mask,
+ spec_decoding_target_max_draft_tokens=md.max_total_draft_tokens,
+ spec_bl_tree_first_sparse_mask_offset_kv=md.spec_bl_tree_first_sparse_mask_offset_kv,
+ num_sparse_topk=md.num_sparse_topk,
+ flash_mla_tile_scheduler_metadata=md.flash_mla_tile_scheduler_metadata,
+ flash_mla_num_splits=md.flash_mla_num_splits,
+ max_num_sequences=md.max_num_sequences,
+ force_prepare_spec_dec_tree_mask=md.force_prepare_spec_dec_tree_mask,
+ )
+
+
+class _KvRangeProbe:
+ """`TRTLLM_MODELING_V2_VALIDATE` over an fp8 KV pool: reports, once, the largest K/V
+ magnitude the first eager forward's attention layers store.
+
+ The pool takes K/V at the unit scale, so anything past the e4m3 maximum saturates
+ silently (a bf16 inf included). One host sync per attention layer of that one forward;
+ a forward under CUDA graph capture is passed over without a read."""
+
+ def __init__(self, last_layer: int) -> None:
+ self._last = last_layer
+ self._worst = 0.0
+
+ @classmethod
+ def for_target(cls, core, enabled: bool) -> "_KvRangeProbe | None":
+ """A probe over `core`'s attention layers when `enabled` and the pool is fp8."""
+ if not enabled or core.quant_mode != QM_FP8_KV_CACHE:
+ return None
+ return cls(core.attn_layers[-1])
+
+ def observe(self, layer: int, k: torch.Tensor, v: torch.Tensor) -> bool:
+ """Fold one attention layer's K/V in; False once the forward's last layer has
+ reported."""
+ if torch.cuda.is_current_stream_capturing():
+ return True
+ peak = float(torch.maximum(k.abs().amax(), v.abs().amax()))
+ # A NaN stays: `max` would drop it.
+ self._worst = math.nan if math.isnan(peak + self._worst) else max(self._worst, peak)
+ if layer != self._last:
+ return True
+ limit = torch.finfo(torch.float8_e4m3fn).max
+ # Written so that NaN fails too.
+ if not self._worst <= limit:
+ logger.warning(
+ f"FP8 KV cache: the first forward stores a K/V magnitude of {self._worst:g}, "
+ f"past the e4m3 maximum {limit:g}; the pool is written at the unit scale, "
+ "so those values saturate"
+ )
+ return False
+
+
+# ------------------------------------------------------------------ prefill ---
+
+
+class _PrefillStep(NamedTuple):
+ """One step's host-side state, built by `PrefillTarget.prepare` and read by the forward's
+ breaks. `num_tokens` is the batch's row count `n` and `physical_tokens` the `T` the
+ forward runs over (`n`, or the prefill CUDA graph bucket it was padded to), and
+ `dp_pad_rows` the rows the MoE blocks pad them with under attention DP (zero otherwise,
+ and on every step a prefill CUDA graph runs). A new one replaces it every step, so no
+ break reads a previous step's."""
+
+ num_tokens: int
+ physical_tokens: int
+ dp_pad_rows: int
+ input_ids: torch.Tensor
+ attn: dict
+ gdn: tuple
+ ple: tuple
+ qsa: tuple
+ vision: tuple | None
+ # A verify step: the generation rows' side-cache writes are staged for the rollback.
+ stage: bool = False
+ # At one rank over an NVFP4 pool, `_nvfp4_context_plan`'s row-tile plan or None.
+ nvfp4_context: tuple | None = None
+ # Over an fp8 pool, the context requests' pages the row-tile pass widens
+ # (`_fp8_context_pages`); None otherwise.
+ fp8_pages: int | None = None
+ # `_live_score_columns`: a bound on every row's selection window.
+ live_cols: int | None = None
+
+
+class PrefillTarget:
+ """Any batch with context rows (`num_contexts > 0`), possibly with generation rows
+ beside them."""
+
+ def __init__(self, core: Qwen4ExpDefaultModelingV2) -> None:
+ self.core = core
+ # Off unless TRTLLM_MODELING_V2_VALIDATE asks.
+ self._contract_pending = step_contract_enabled()
+ # Dropped once it has reported.
+ self._kv_probe = _KvRangeProbe.for_target(core, self._contract_pending)
+ # The rstd sites' `(1 + w)` operands (`_HC_RSTD_MIN_ROWS`), formed once here.
+ self._layers = [
+ L._replace(ahc_norm_w1=_norm_w1(L.ahc_norm), mhc_norm_w1=_norm_w1(L.mhc_norm))
+ for L in _layer_views(core)
+ ]
+ stack = _stack_views(core)
+ self._stack = stack._replace(fhc_norm_w1=_norm_w1(stack.fhc_norm))
+ # The step `prepare` built, read by the breaks.
+ self._step: _PrefillStep | None = None
+ # Whether this rank reported the GDN in-place short-context path yet.
+ self._gdn_inplace_reported = False
+
+ def _check_step_contract(self, md) -> None:
+ """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): calling `step_args`
+ surfaces an upstream metadata rename here, not mid-forward."""
+ if not self._contract_pending:
+ return
+ try:
+ self.step_args(md)
+ except AttributeError as exc:
+ raise AssertionError(f"attention metadata surface drifted: {exc}") from exc
+ self._contract_pending = False
+
+ def step_args(self, md: TrtllmAttentionMetadata) -> dict:
+ """Project the prepared metadata onto thop_attention's batch state."""
+ return _step_args(md, md.num_contexts, md.num_ctx_tokens)
+
+ # --------------------------------------------------------------- blocks ---
+
+ def _attention_block(self, x, i, L, proj_qkv, index_scratch, attn_out, pos3, pos):
+ """One gated GQA layer plus its QSA index branch, at this rank's heads.
+
+ The four input projections are one GEMM against the fused operand; the norm + rope
+ of both planes is row-wise and runs over all `T` rows here; the cache writes, the
+ selection and the attention core read the step's batch and run in the
+ `_attention_core` break, which leaves the attention output in `attn_out`. The gate
+ is read through its row stride, and the output projection's per-rank partial is
+ summed over the TP group."""
+ core = self.core
+ proj = _dense_mm(x, L.attn_in, L.attn_in8, cublas_mm, "in")
+ gate = torch.split(proj, [core.qkv_width, core.q_width, core.index_scratch_width], 1)[1]
+ _attention_rope(core, proj, index_scratch, L, pos3, pos)
+ eager_on_graph(proj, proj_qkv, index_scratch, attn_out, i, inner=self._attention_core)
+ fused_sigmoid_mul_inplace(attn_out, gate)
+ return core.tp_sum(_dense_mm(attn_out, L.o, L.o8, cublas_mm, "out"))
+
+ def _attention_core(self, proj, proj_qkv, index_scratch, attn_out, i) -> None:
+ """Break: the index plane's store and compression, the dense/sparse decision and the
+ attention core over the step's `n` rows. Writes rows `[:n]` of `attn_out`
+ `[T, q_width]` in place and zeroes rows `[n:]`.
+
+ The index plane advances on every step, including below the threshold, so a request
+ growing past it finds its prefix compressed; above it the main K/V append becomes
+ this layer's own call. On a verify step the generation rows' index slots (and, on
+ the first attention layer, their position slots) are staged for `SpeculativeRollback`
+ before the store and the compression overwrite them."""
+ core = self.core
+ L = self._layers[i]
+ step = self._step
+ (
+ req_idx,
+ logical,
+ seq_lens_q,
+ visible_blocks,
+ topk_indices,
+ topk_row_starts,
+ block_table,
+ tokens_per_block,
+ coords,
+ sparse,
+ ) = step.qsa
+ num_tokens = step.num_tokens
+ proj = proj[:num_tokens]
+ _, k_dst, v_dst = _attn_in_heads(core, proj)
+ if self._kv_probe is not None and not self._kv_probe.observe(i, k_dst, v_dst):
+ self._kv_probe = None
+ # The packed `[q | k | v]` as one compact copy: the dense op reads it flat, and the
+ # sparse GQA reads its query heads from it.
+ qkv = proj_qkv[:num_tokens]
+ qkv.copy_(proj[:, : core.qkv_width])
+ if step.stage:
+ core.spec_rollback.stage_index(
+ core._bound.pools.index_cache[i], core._bound.pools.position_cache
+ )
+ _index_plane(core, proj, i, L, (req_idx, logical, block_table, tokens_per_block, coords))
+
+ if sparse:
+ kv = core._bound.pools.kv[i]
+ # Append the normalized, rotated K/V: nothing else writes the pool on this path,
+ # and the selection reads the current token from it.
+ _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block)
+ q_index = _index_query(core, index_scratch[:num_tokens])
+ # The step's context rows come first, consecutive rows of a request, and score as
+ # query tiles; its generation rows, one a request, score row by row. Each part
+ # selects into its own rows.
+ num_ctx = step.attn["num_ctx_tokens"]
+ parts = [(0, num_ctx, True)]
+ if num_ctx < num_tokens:
+ parts.append((num_ctx, num_tokens, False))
+ selected = torch.empty(
+ [num_tokens, core.token_topk + core.compress_ratio - 1],
+ dtype=torch.int32,
+ device=proj.device,
+ )
+ for lo, hi, context_rows in parts:
+ scores = qsa_paged_index_scores(
+ q=q_index[lo:hi],
+ index_cache=core._bound.pools.index_cache[i],
+ block_table=block_table,
+ query_positions=logical[lo:hi],
+ request_indices=req_idx[lo:hi],
+ tokens_per_block=tokens_per_block,
+ compress_ratio=core.compress_ratio,
+ only_visible_blocks=True,
+ context_rows=context_rows,
+ )
+ _selection(
+ core,
+ scores,
+ visible_blocks[lo:hi],
+ logical[lo:hi],
+ seq_lens_q[lo:hi],
+ topk_indices[lo:hi],
+ topk_row_starts[lo:hi],
+ selected[lo:hi],
+ step.live_cols,
+ )
+ # Context rows with the prefill pair (a bound on a row's scan), generation rows
+ # without.
+ if core.mapping.world_size == 1 and kv.k_scale is not None:
+ _sparse_gqa_nvfp4_context(
+ core,
+ kv,
+ step.nvfp4_context,
+ _packed_q(core, qkv),
+ block_table,
+ selected,
+ req_idx,
+ logical,
+ tokens_per_block,
+ num_ctx,
+ attn_out[:num_tokens],
+ step.attn["predicted_tokens_per_seq"],
+ )
+ parts = []
+ for lo, hi, context_rows in parts:
+ out = _sparse_gqa(
+ core,
+ kv,
+ _packed_q(core, qkv)[lo:hi],
+ block_table,
+ selected[lo:hi],
+ req_idx[lo:hi],
+ tokens_per_block,
+ logical[lo:hi] if context_rows else None,
+ step.fp8_pages,
+ )
+ attn_out[lo:hi].copy_(torch.reshape(out, [hi - lo, core.q_width]))
+ else:
+ thop_attention(
+ q=qkv,
+ output=attn_out[:num_tokens],
+ local_layer_idx=core._bound.pools.pool_idx[i],
+ num_heads=core.heads_q,
+ num_kv_heads=core.heads_kv,
+ head_size=core.head_dim,
+ quant_mode=core.quant_mode,
+ kv_scale_orig_quant=core.kv_scale_orig_quant,
+ kv_scale_quant_orig=core.kv_scale_quant_orig,
+ **step.attn,
+ **_CALL_CONSTANTS,
+ )
+ attn_out[num_tokens:].zero_()
+
+ def _gdn_block(self, x, i, L, core_out) -> torch.Tensor:
+ """One Gated-DeltaNet layer over a batch with a context prefix, at this rank's heads:
+ the one input projection GEMM (q|k|v, z, a and b are column views of its output) and
+ the gated-norm epilogue run over all `T` rows here, the recurrence in the `_gdn_core`
+ break, which leaves its output in `core_out` `[1, T, num_v_heads, head_v_dim]`."""
+ core = self.core
+ proj = _dense_mm(x, L.in_proj, L.in_proj8, cublas_mm, "in")
+ z = torch.split(proj, core.in_proj_split, 1)[1]
+ eager_on_graph(proj, core_out, i, inner=self._gdn_core)
+ return self._gdn_epilogue(core_out, z, L, x.shape[0])
+
+ def _gdn_core(self, proj, out, i) -> None:
+ """Break: the recurrence over the step's `n` rows. Writes rows `[:n]` of `out` in place
+ and zeroes rows `[n:]`. Every branch is on a host int of the step.
+
+ The context rows take one fused conv + post-conv pass (`_gdn_prefill_prep`) and
+ FlashInfer's chunked recurrence, and a mixed step's one-token generation rows the
+ decode step's conv and recurrence. A context sequence carrying nothing starts from
+ zeros without its pool slots being cleared: the conv pass reads no carried window for
+ it, and the recurrence (`has_initial_state`) gathers zeros for it.
+
+ On a verify step the generation rows are draft chains of unknown accepted length:
+ they take `_gdn_verify` (per-step conv windows and the cached replay) instead."""
+ core = self.core
+ L = self._layers[i]
+ step = self._step
+ (
+ num_prefill_tokens,
+ num_decode_tokens,
+ si_p,
+ si_all,
+ si_d,
+ hi_p,
+ hi_all,
+ qsl_p,
+ qsl_all,
+ qsl_long,
+ qsl_long_p,
+ num_decodes,
+ gen_p,
+ verify,
+ ) = step.gdn
+ num_tokens = step.num_tokens
+ pools = core._bound.pools
+ ssm_states = pools.ssm[i]
+ mixed_qkv, _, a, b = torch.split(proj[:num_tokens], core.in_proj_split, 1)
+
+ if verify is not None:
+ mq_p, mq_d = torch.split(mixed_qkv, [num_prefill_tokens, num_decode_tokens], 0)
+ a_p, a_d = torch.split(a, [num_prefill_tokens, num_decode_tokens], 0)
+ b_p, b_d = torch.split(b, [num_prefill_tokens, num_decode_tokens], 0)
+ # The context rows as on a step without verify rows.
+ q, k, v, g, beta = self._gdn_prefill_prep(mq_p, a_p, b_p, L, i, qsl_p, si_p, hi_p)
+ flashinfer_chunk_gated_delta_rule(
+ q,
+ k,
+ v,
+ g,
+ beta,
+ scale=None,
+ initial_state=ssm_states,
+ initial_state_indices=si_p,
+ inplace_indexed_state_update=True,
+ output_final_state=False,
+ cu_seqlens=qsl_long_p,
+ head_first=False,
+ use_qk_l2norm_in_kernel=False,
+ output=out[:, :num_prefill_tokens],
+ state_workspace=None,
+ has_initial_state=hi_p,
+ linear_gate=True,
+ )
+ out[:, num_prefill_tokens:num_tokens].copy_(
+ _gdn_verify(
+ core,
+ mq_d,
+ a_d,
+ b_d,
+ i,
+ L,
+ si_d,
+ num_decodes,
+ gen_p,
+ verify,
+ )
+ )
+ elif num_decode_tokens > 0:
+ mq_p, mq_d = torch.split(mixed_qkv, [num_prefill_tokens, num_decode_tokens], 0)
+ a_p, b_p = a[:num_prefill_tokens], b[:num_prefill_tokens]
+ q, k, v, g, beta = self._gdn_prefill_prep(mq_p, a_p, b_p, L, i, qsl_p, si_p, hi_p)
+ causal_conv1d_update(
+ x=mq_d,
+ conv_state=pools.conv[i],
+ weight=L.conv_w,
+ bias=None,
+ silu_activation=True,
+ cache_seqlens=None,
+ conv_state_indices=si_d,
+ pad_slot_id=_PAD_SLOT_ID,
+ )
+ # The two halves run separate recurrences over disjoint slots and write disjoint
+ # row ranges of the output.
+ flashinfer_chunk_gated_delta_rule(
+ q,
+ k,
+ v,
+ g,
+ beta,
+ scale=None,
+ initial_state=ssm_states,
+ initial_state_indices=si_p,
+ inplace_indexed_state_update=True,
+ output_final_state=False,
+ cu_seqlens=qsl_long_p,
+ head_first=False,
+ use_qk_l2norm_in_kernel=False,
+ output=out[:, :num_prefill_tokens],
+ state_workspace=None,
+ has_initial_state=hi_p,
+ linear_gate=True,
+ )
+ # The generation rows: the decode step's recurrence (without the gated norm,
+ # which runs with the context rows' below) over their rows of the projection,
+ # whose q/k/v columns the conv step rewrote; it computes both gates and the q/k
+ # norm itself.
+ gdn_decode_rows(
+ proj[num_prefill_tokens:num_tokens],
+ ssm_states,
+ si_d,
+ L.a_log,
+ L.dt_bias,
+ core.num_k_heads,
+ core.head_k_dim,
+ out[0, num_prefill_tokens:num_tokens].view(num_decode_tokens, core.value_dim),
+ )
+ else:
+ q, k, v, g, beta = self._gdn_prefill_prep(
+ mixed_qkv, a, b, L, i, qsl_all, si_all, hi_all
+ )
+ # q and k arrive L2-normalized; the recurrence must not repeat it.
+ flashinfer_chunk_gated_delta_rule(
+ q,
+ k,
+ v,
+ g,
+ beta,
+ scale=None,
+ initial_state=ssm_states,
+ initial_state_indices=si_all,
+ inplace_indexed_state_update=True,
+ output_final_state=False,
+ cu_seqlens=qsl_long,
+ head_first=False,
+ use_qk_l2norm_in_kernel=False,
+ output=out[:, :num_tokens],
+ state_workspace=None,
+ has_initial_state=hi_all,
+ linear_gate=True,
+ )
+ out[:, num_tokens:].zero_()
+
+ def _gdn_prefill_prep(self, mq, a, b, L, i, qsl, si, hi):
+ """The context rows' recurrence inputs in one pass over `mq`, the token-major q|k|v
+ column view of the projection: the causal conv against each sequence's carried
+ window (published back to the pool), SiLU, the q|k|v split with q/k L2-normalized,
+ and the decay/beta gates (the decay as the factor `exp(g)`, `linear_gate=True`)."""
+ core = self.core
+ return gdn_prefill_conv_post(
+ x=mq,
+ weight=L.conv_w,
+ conv_states=core._bound.pools.conv[i],
+ query_start_loc=qsl,
+ cache_indices=si,
+ has_initial_state=hi,
+ a=a,
+ b=b,
+ A_log=L.a_log,
+ dt_bias=L.dt_bias,
+ num_k_heads=core.num_k_heads,
+ head_k_dim=core.head_k_dim,
+ num_v_heads=core.num_v_heads,
+ head_v_dim=core.head_v_dim,
+ pad_slot_id=_PAD_SLOT_ID,
+ l2_norm_eps=_L2_NORM_EPS,
+ softplus_threshold=_SOFTPLUS_THRESHOLD,
+ linear_gate=True,
+ )
+
+ def _gdn_epilogue(self, out, z, L, num_tokens) -> torch.Tensor:
+ """The gated RMS norm (sigmoid gate, plain non-gemma weight) and the output
+ projection, whose per-rank partial is summed over the TP group."""
+ core = self.core
+ o = rms_norm_gated_token_major(
+ torch.reshape(out, [-1, core.head_v_dim]),
+ torch.reshape(z, [num_tokens, core.num_v_heads, core.head_v_dim]),
+ weight=L.gnorm,
+ eps=core.eps,
+ fp8_scale=None,
+ gate_activation="sigmoid",
+ )
+ o = torch.reshape(o, [num_tokens, core.value_dim])
+ return core.tp_sum(_dense_mm(o, L.out, L.out8, cublas_mm, "out"))
+
+ def _ple_block(self, hc, hc_out, normed_out) -> None:
+ """Break: the n-gram side path on its decoder layer, over the step's `n` rows, added
+ into the stream and normalized for the layer's attention-side site. Writes rows
+ `[:n]` of `hc_out` (the stream) and `normed_out` (its grouped norm) in place and zeroes
+ rows `[n:]`.
+
+ Per-request stateful: ids become one `ngram_size`-wide window per token
+ (predecessors from the paged id history); each window hashes to one row per head;
+ rows gather fp8 vectors from one table (`ple_ngram_embed`); the key/value projections,
+ the gate, a dilated depthwise short conv over a paged state, the residual add and the
+ next site's grouped norm are `ple_gate_conv_inject_norm`. A sequence carrying nothing
+ is seeded with `eos_token_id`.
+
+ On a verify step the generation rows' pre-step pool rows and ids are staged before
+ the first launch publishes, and the conv input the fused launch writes out after it,
+ for `SpeculativeRollback` to rewrite both pools to the accepted prefix."""
+ core = self.core
+ S = self._stack
+ step = self._step
+ conv_states, ngram_pool, si_all, carry_all, qsl = step.ple
+ num_tokens = step.num_tokens
+ next_norm = self._layers[core.ple_layer].ahc_norm
+ conv_input = None
+ if step.stage:
+ core.spec_rollback.stage_ple_verify(ngram_pool, conv_states, step.input_ids)
+ conv_input = torch.empty_like(hc[:num_tokens])
+ ngram = ple_ngram_embed(
+ step.input_ids,
+ ngram_pool,
+ si_all,
+ carry_all,
+ qsl,
+ core.w["ple_table"],
+ S.ple_tscale,
+ S.ple_mult,
+ S.ple_vocab,
+ S.ple_offs,
+ core.eos_token_id,
+ num_tokens,
+ )
+ ple_gate_conv_inject_norm(
+ hc[:num_tokens],
+ cublas_mm(ngram, S.ple_key),
+ cublas_mm(ngram, S.ple_val),
+ S.ple_nk_w,
+ S.ple_nq_w,
+ S.ple_nc_w,
+ S.ple_conv_w,
+ conv_states,
+ si_all,
+ qsl,
+ next_norm,
+ core.eps,
+ core.hc_count,
+ _PLE_CLAMP_MIN,
+ core.ple_dilation,
+ _PAD_SLOT_ID,
+ num_tokens,
+ carry_all,
+ conv_input=conv_input,
+ out=(hc_out[:num_tokens], normed_out[:num_tokens]),
+ )
+ if step.stage:
+ core.spec_rollback.stage_ple_conv_input(conv_input)
+ hc_out[num_tokens:].zero_()
+ normed_out[num_tokens:].zero_()
+
+ def _splice(self, h) -> None:
+ """Break: write the step's vision rows over their placeholder rows of the token
+ embedding `h` `[T, hidden]` in place, and zero rows `[n:]`."""
+ step = self._step
+ if step.vision is not None:
+ rows, mm_token_indices = step.vision
+ h.index_copy_(0, mm_token_indices, rows)
+ h[step.num_tokens :].zero_()
+
+ # -------------------------------------------------------------- forward ---
+
+ def prepare(
+ self,
+ attn_metadata: AttentionMetadata,
+ input_ids: torch.IntTensor,
+ position_ids: torch.IntTensor,
+ vision: tuple | None,
+ spec_metadata=None,
+ ) -> None:
+ """Everything this step reads from the host, kept for the forward's breaks as one
+ `_PrefillStep`: the contract check, the attention op's per-step arguments and the
+ batch's bookkeeping, built once and shared by every layer.
+
+ Called by the core's `begin_step` on every step with context rows, before the
+ forward. On a step a breakable prefill graph replays, this and the breaks are all of
+ this target that runs."""
+ core = self.core
+ md = attn_metadata
+ self._check_step_contract(md)
+ attn_args = self.step_args(md)
+ mm = md.mamba_metadata
+
+ # The activation rows are the batch's rows, or a prefill CUDA graph bucket's.
+ physical_tokens = input_ids.shape[0]
+ assert physical_tokens in (md.num_tokens, md.padded_num_tokens), (
+ f"the engine handed {physical_tokens} activation rows for a batch its metadata "
+ f"describes as {md.num_tokens} tokens (padded_num_tokens={md.padded_num_tokens})"
+ )
+ num_tokens = md.num_tokens
+ # Under attention DP the MoE blocks pad the `T` rows to the group's largest count. A
+ # prefill CUDA graph pads every rank to one bucket, so the body it captures pads
+ # nothing; that zero is what makes the count a per-capture constant.
+ dp_pad_rows = core._dp_pad_rows(md, physical_tokens)
+ assert physical_tokens == num_tokens or dp_pad_rows == 0, (
+ f"a prefill CUDA graph bucket of {physical_tokens} rows with {dp_pad_rows} attention-DP "
+ "padding rows: the graph would freeze the group's row counts"
+ )
+ batch = md.num_seqs
+ num_prefills = md.num_contexts
+ num_decodes = batch - num_prefills
+ num_prefill_tokens = md.num_ctx_tokens
+ num_decode_tokens = num_tokens - num_prefill_tokens
+ if not self._gdn_inplace_reported and num_prefill_tokens <= _GDN_INPLACE_MAX_TOKENS:
+ # The recurrence of a context this short runs in place over the pool rows
+ # (`flashinfer_chunk_gated_delta_rule`'s short path).
+ self._gdn_inplace_reported = True
+ print(
+ f"modeling_v2 qwen4_exp default target: rank {core.mapping.rank} runs short GDN "
+ "context chunks in place over the recurrent pool"
+ )
+ # Rows per generation sequence: 1 without speculation, else `runtime_draft_len + 1`.
+ gen_p, is_verify = _verify_rows(core, md, spec_metadata, num_decodes, num_decode_tokens)
+ attn_args["predicted_tokens_per_seq"] = gen_p
+
+ si_all = mm.state_indices[:batch]
+ hi_all = mm.has_initial_states[:batch]
+ si_p, si_d = torch.split(si_all, [num_prefills, num_decodes], 0)
+ # The single-token decode kernel takes its slot vector through a 16-byte-aligned
+ # pointer; the generation half of the int32 slot vector starts `num_prefills` rows
+ # in, so one copy per step makes it aligned for every layer.
+ si_d = torch.clone(si_d)
+ hi_p = hi_all[:num_prefills]
+ # `query_start_loc` covers the context prefix; `_long` the whole batch, and its prefix
+ # the context rows' chunked recurrence when generation rows run beside them.
+ qsl_all = mm.query_start_loc
+ qsl_p = None if qsl_all is None else qsl_all[: num_prefills + 1]
+ qsl_long = mm.query_start_loc_long
+ qsl_long_p = qsl_long[: num_prefills + 1]
+ gdn_state = (
+ num_prefill_tokens,
+ num_decode_tokens,
+ si_p,
+ si_all,
+ si_d,
+ hi_p,
+ hi_all,
+ qsl_p,
+ qsl_all,
+ qsl_long,
+ qsl_long_p,
+ num_decodes,
+ gen_p,
+ _verify_state(core, mm, num_decodes) if is_verify else None,
+ )
+
+ # PLE carry flags must cover the whole batch, but `has_initial_states` is filled for
+ # the context prefix only; every generation sequence carries, so pad with True.
+ ple_carry = nn.functional.pad(hi_p, [0, num_decodes], mode="constant", value=1.0)
+ # Under speculation `query_start_loc_long` has one row per generation sequence, so
+ # with `gen_p > 1` both PLE stateful calls would be mis-segmented: rebuilt from
+ # `seq_lens_cuda`, which holds the true counts during the trunk forward.
+ ple_qsl = (
+ qsl_long
+ if gen_p == 1
+ else cu_seqlens_from_seq_lens(md.seq_lens_cuda[:batch], torch.int64)
+ )
+ pools = core._bound.pools
+ ple_state = (pools.ple_state[0], pools.ple_state[1], si_all, ple_carry, ple_qsl)
+
+ # Dense/sparse decision once per forward, from a host-side metadata read: a
+ # per-capture constant.
+ sparse = batch > 0 and int(md.kv_lens_runtime[:batch].max()) > core.dense_threshold
+ # The `[T, 3]` transpose of the `[3, 1, T]` position buffer, which the QSA position
+ # plane stores.
+ coords = torch.transpose(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1]), 0, 1)
+ qsa = (
+ md.qsa_req_idx_per_token[:num_tokens],
+ md.qsa_logical_positions[:num_tokens],
+ md.qsa_sequence_lengths[:num_tokens],
+ md.qsa_visible_blocks[:num_tokens],
+ md.qsa_topk_indices[:num_tokens],
+ md.qsa_topk_row_starts[:num_tokens],
+ md.qsa_block_table,
+ md.tokens_per_block,
+ coords[:num_tokens],
+ sparse,
+ )
+
+ # A verify step's generation rows write the PLE pools and the QSA side caches for
+ # every draft row, accepted or not; `SpeculativeRollback` keeps what its commit needs
+ # to undo the rejected ones. Reopened on every forward, so a step whose commit never
+ # ran leaves nothing for the next one. Context rows are never staged.
+ rollback = core.spec_rollback
+ if rollback is not None:
+ rollback.discard()
+ if is_verify:
+ rollback.begin(
+ num_contexts=num_prefills,
+ prefix_tokens=num_prefill_tokens,
+ num_gens=num_decodes,
+ rows_per_gen=gen_p,
+ slots=si_d,
+ request_indices=md.qsa_req_idx_per_token[num_prefill_tokens:num_tokens],
+ logical_positions=md.qsa_logical_positions[num_prefill_tokens:num_tokens],
+ block_table=md.qsa_block_table,
+ tokens_per_block=md.tokens_per_block,
+ )
+ fp8_pages = None
+ if (
+ sparse
+ and core.quant_mode == QM_FP8_KV_CACHE
+ and num_prefill_tokens >= _ROWTILE_MIN_ROWS
+ ):
+ fp8_pages = _fp8_context_pages(md, md.tokens_per_block)
+ nvfp4_context = None
+ if (
+ core.mapping.world_size == 1
+ and sparse
+ and core.quant_mode == QM_NVFP4_KV_CACHE
+ and num_prefill_tokens > _NVFP4_ROWTILE_MIN_ROWS
+ ):
+ nvfp4_context = _nvfp4_context_plan(
+ core, md, qsa[0], md.qsa_block_table, md.tokens_per_block
+ )
+ self._step = _PrefillStep(
+ num_tokens=num_tokens,
+ physical_tokens=physical_tokens,
+ dp_pad_rows=dp_pad_rows,
+ input_ids=input_ids[:num_tokens],
+ attn=attn_args,
+ gdn=gdn_state,
+ ple=ple_state,
+ qsa=qsa,
+ vision=vision,
+ stage=is_verify,
+ fp8_pages=fp8_pages,
+ nvfp4_context=nvfp4_context,
+ live_cols=_live_score_columns(core, md, batch),
+ )
+
+ def forward(
+ self,
+ attn_metadata: AttentionMetadata,
+ input_ids: torch.IntTensor,
+ position_ids: torch.IntTensor,
+ **kwargs,
+ ):
+ """The step `prepare` built, over the `T` activation rows the engine handed in: the
+ batch's `n` rows, or the prefill CUDA graph bucket they were padded to.
+
+ **This is the body a breakable prefill CUDA graph captures**, once per bucket, and
+ replays for every step it pads to that bucket without running this Python. So nothing
+ here reads the step: every op is row-wise and sized by `T` alone, and what depends on
+ the step's batch runs in one of the four breaks (`_splice`, `_ple_block`,
+ `_attention_core`, `_gdn_core`), each writing rows `[:n]` of a fixed output allocated
+ here and zeroing rows `[n:]`. The tensor-parallel sums run on all `T` rows, which
+ every rank holds alike; under attention DP the MoE round trip pads them by the step's
+ `dp_pad_rows`, zero whenever a graph captures this body."""
+ core = self.core
+ num_tokens = input_ids.shape[0]
+ assert num_tokens == self._step.physical_tokens, (
+ f"the forward got {num_tokens} activation rows and prepare saw "
+ f"{self._step.physical_tokens}"
+ )
+ pos3 = _copy_positions(position_ids)
+ pos = pos3[0]
+
+ # The token embedding (replicated) with this step's vision rows spliced in.
+ h = nn.functional.embedding(input_ids, core.w["embed"])
+ eager_on_graph(h, inner=self._splice)
+ dt = h.dtype
+ dev = h.device
+
+ # Scratch shared by every layer that uses it; each break rewrites its whole output.
+ proj_qkv = torch.empty([num_tokens, core.qkv_width], dtype=dt, device=dev)
+ index_scratch = torch.empty([num_tokens, core.index_scratch_width], dtype=dt, device=dev)
+ attn_out = torch.empty([num_tokens, core.q_width], dtype=dt, device=dev)
+ gdn_out = torch.empty(
+ [1, num_tokens, core.num_v_heads, core.head_v_dim], dtype=dt, device=dev
+ )
+ ple_hc = torch.empty([num_tokens, core.stream], dtype=dt, device=dev)
+ ple_normed = torch.empty([num_tokens, core.stream], dtype=dt, device=dev)
+
+ # The embedding repeated across the hc_count residual streams.
+ hc_count = core.hc_count
+ hc = torch.cat([h] * hc_count, -1)
+ layers, S = self._layers, self._stack
+ pad_rows = self._step.dp_pad_rows
+ # A host decision on `T` alone (each captured bucket takes one branch): at one rank, from
+ # `_HC_RSTD_MIN_ROWS` rows, every site but the PLE layer's attention-side one is
+ # `qwen4_exp_hc_inject_rstd` and the rstd mixer.
+ by_rstd = core.hc_fast_sites and num_tokens >= _HC_RSTD_MIN_ROWS
+ normed = rstd = None
+ if by_rstd:
+ _, rstd = qwen4_exp_hc_inject_rstd(hc, None, None, core.eps, hc_count)
+ else:
+ _, normed = qwen4_exp_hc_inject_norm(
+ hc, None, None, layers[0].ahc_norm, core.eps, hc_count
+ )
+ for i, L in enumerate(layers):
+ if i == core.ple_layer:
+ eager_on_graph(hc, ple_hc, ple_normed, inner=self._ple_block)
+ hc, normed = ple_hc, ple_normed
+ # The PLE kernel stores this site's normed stream: the stored-normed mixer
+ # reads it (cheaper than a second pass for its rstd).
+ mixed, inject = _hc_mixer(core, normed, L.ahc_first, L.ahc_up, True)
+ elif by_rstd:
+ mixed, inject = _hc_mixer_rstd(
+ core, hc, rstd, L.ahc_norm, L.ahc_norm_w1, L.ahc_first, L.ahc_up, True
+ )
+ else:
+ mixed, inject = _hc_mixer(core, normed, L.ahc_first, L.ahc_up, True)
+ if core.is_attn[i]:
+ block_out = self._attention_block(
+ mixed, i, L, proj_qkv, index_scratch, attn_out, pos3, pos
+ )
+ else:
+ block_out = self._gdn_block(mixed, i, L, gdn_out)
+ if by_rstd:
+ hc, rstd = qwen4_exp_hc_inject_rstd(hc, block_out, inject, core.eps, hc_count)
+ mixed, inject, fp4 = _hc_moe_mixer_rstd(core, hc, rstd, L)
+ else:
+ hc, normed = qwen4_exp_hc_inject_norm(
+ hc, block_out, inject, L.mhc_norm, core.eps, hc_count
+ )
+ mixed, inject, fp4 = _hc_moe_mixer(core, normed, L)
+ next_norm = layers[i + 1].ahc_norm if i + 1 < len(layers) else S.fhc_norm
+ # The MoE block output is formed inside the injection: routed + sigmoid(gate) *
+ # shared, rounded once.
+ routed, gate_logits, shared = _moe_block(core, mixed, L, pad_rows, fp4)
+ if by_rstd:
+ hc, rstd = qwen4_exp_hc_inject_rstd(
+ hc, routed, inject, core.eps, hc_count, shared, gate_logits
+ )
+ else:
+ hc, normed = qwen4_exp_hc_inject_norm(
+ hc, routed, inject, next_norm, core.eps, hc_count, shared, gate_logits
+ )
+ del routed, gate_logits, shared, fp4
+ # The stack-final mixer with no injection weight is the final norm.
+ if by_rstd:
+ final, _ = _hc_mixer_rstd(
+ core, hc, rstd, S.fhc_norm, S.fhc_norm_w1, S.fhc_first, S.fhc_up, False
+ )
+ else:
+ final, _ = _hc_mixer(core, normed, S.fhc_first, S.fhc_up, False)
+ if not core.mtp_enabled:
+ return final
+ # `lm_head` reads the `[T, hidden]` mix; the MTP module reads the `[T, stream]` stream
+ # it was collapsed from. They leave packed in one tensor, because a breakable prefill
+ # graph's body returns one tensor (`split_spec_output`).
+ return torch.cat([final, hc], -1)
+
+ # ------------------------------------------------------------- the draft ---
+
+ def _capture_selection(self, md, topk_indices, visible_blocks, compact=False) -> None:
+ """IndexShare: stash, per request, the selection of the row the draft loop continues
+ from, with that row's complete-group count, in the draft layer's row of the shared
+ buffers: a context request's last row of this chunk, and a generation request's last
+ accepted row (`num_accepted - 1` past its first). A `compact` selection holds each
+ context request's last row, then every generation row (`mtp_forward`'s `keep`)."""
+ batch = md.num_seqs
+ assert md.qsa_mtp_num_accepted is not None, (
+ "IndexShare needs the worker's accepted counts to pick each request's capture row; "
+ "the draft loop did not supply them"
+ )
+ seq_lens = md.seq_lens_cuda[:batch].to(torch.int64)
+ starts = torch.cumsum(seq_lens, 0) - seq_lens
+ accepted = md.qsa_mtp_num_accepted[:batch].to(torch.int64)
+ offsets = torch.minimum(torch.clamp(accepted - 1, min=0), seq_lens - 1)
+ offsets[: md.num_contexts] = seq_lens[: md.num_contexts] - 1
+ rows = starts + offsets
+ if compact:
+ num_contexts = md.num_contexts
+ rows[:num_contexts] = torch.arange(num_contexts, device=rows.device)
+ rows[num_contexts:] -= md.num_ctx_tokens - num_contexts
+ slot = self.core._draft_bound.share_slot
+ md.qsa_shared_topk_indices[slot, :batch].copy_(torch.index_select(topk_indices, 0, rows))
+ md.qsa_shared_topk_visible_blocks[slot, :batch].copy_(
+ torch.index_select(visible_blocks, 0, rows)
+ )
+ # Host-side readiness, recorded in program order.
+ md.qsa_shared_topk_captured_slots.add(slot)
+
+ def mtp_forward(
+ self,
+ embed_tokens: torch.Tensor,
+ all_rank_num_tokens,
+ input_ids: torch.Tensor,
+ position_ids: torch.Tensor,
+ hidden_states: torch.Tensor,
+ attn_metadata: AttentionMetadata,
+ **kwargs,
+ ) -> torch.Tensor:
+ """One draft step whose metadata still carries context rows: the draft loop's first
+ (draft-extend) step of a batch with context rows, eager. Every row writes the layer's
+ KV and index planes; the worker gathers one row per request from the result. Under
+ attention TP only `keep` -- each context request's last row, then every generation
+ row -- goes on past those writes (the query side included on the sparse path; the
+ dense call's metadata describes every row), and the other returned rows keep the
+ layer's input stream. Under attention DP every row runs the whole layer.
+
+ The draft layer's attention is the trunk prefill target's: the context rows score
+ as query tiles and bound their sparse scan by their positions, the generation rows
+ (`p` a request) score row by row. `all_rank_num_tokens` is the group's row counts of
+ this draft step, which pad the MoE round trip under attention DP (`_draft_pad_rows`)."""
+ core = self.core
+ assert isinstance(attn_metadata, TrtllmAttentionMetadata)
+ assert core._bound is not None, "the trunk's cache-manager bind has not run"
+ md = attn_metadata
+ rows = hidden_states.shape[0]
+ assert input_ids.shape[0] == rows and hidden_states.shape[1] == core.stream, (
+ f"the draft step got {input_ids.shape[0]} token ids and {tuple(hidden_states.shape)} "
+ f"hidden states; the MTP module reads one {core.stream}-wide stream row per token"
+ )
+ # Projected on every call: the worker rewrites `attn_metadata` in place between
+ # draft steps.
+ step = self.step_args(md)
+ num_ctx = md.num_ctx_tokens
+ gen_seqs = md.num_seqs - md.num_contexts
+ p, _ = _verify_rows(core, md, None, gen_seqs, rows - num_ctx)
+ step["predicted_tokens_per_seq"] = p
+ draft, layer_idx = core._draft_step_bind(md)
+ assert rows == md.num_tokens, (
+ f"the draft step carries {rows} rows for a batch its metadata describes as "
+ f"{md.num_tokens} tokens; the QSA row mapping covers the second"
+ )
+ pad_rows = _draft_pad_rows(core, all_rank_num_tokens, rows)
+ # The trunk's dense/sparse decision from the same host buffer: every draft step of
+ # one engine forward takes the branch its trunk took.
+ batch = md.num_seqs
+ sparse = batch > 0 and int(md.kv_lens_runtime[:batch].max()) > core.dense_threshold
+ # IndexShare: the only draft step with context rows is the loop's first, which always
+ # runs the indexer and, with sharing on, captures the selection the later steps reuse.
+ share = core.mtp_index_share and md.qsa_in_mtp_draft_loop
+ assert not (share and md.qsa_indexer_skip_topk), (
+ "the worker asked a draft step with context rows to reuse a selection"
+ )
+ pos3, pos, coords = _draft_positions(position_ids)
+ M = core._mtp
+ req_idx = md.qsa_req_idx_per_token[:rows]
+ logical = md.qsa_logical_positions[:rows]
+ block_table = md.qsa_block_table
+ tokens_per_block = md.tokens_per_block
+
+ x, mixed, inject = _mtp_front(core, M, embed_tokens, input_ids, hidden_states)
+ qkv, gate, k_dst, v_dst = _mtp_qkv(core, M, mixed, pos3)
+ q_index = _mtp_index_branch(
+ core, M, mixed, draft, (req_idx, logical, block_table, tokens_per_block, pos, coords)
+ )
+ keep = None
+ if core.dp_size == 1 and num_ctx > md.num_contexts:
+ ctx_lens = md.seq_lens_cuda[: md.num_contexts].to(torch.long)
+ keep = torch.cat(
+ [
+ torch.cumsum(ctx_lens, 0) - 1,
+ torch.arange(num_ctx, rows, dtype=torch.long, device=x.device),
+ ]
+ )
+ if sparse and keep is not None:
+ kv = draft.kv
+ _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block)
+ q_rows = keep.shape[0]
+ q_req, q_logical = req_idx[keep], logical[keep]
+ visible_blocks = md.qsa_visible_blocks[:rows][keep]
+ seq_lens_q = md.qsa_sequence_lengths[:rows][keep]
+ row_starts = md.qsa_topk_row_starts[:rows][keep]
+ topk_indices = md.qsa_topk_indices[:q_rows]
+ scores = qsa_paged_index_scores(
+ q=q_index[keep],
+ index_cache=draft.index_cache,
+ block_table=block_table,
+ query_positions=q_logical,
+ request_indices=q_req,
+ tokens_per_block=tokens_per_block,
+ compress_ratio=core.compress_ratio,
+ only_visible_blocks=True,
+ # One row per context request: no run of rows to share a key gather.
+ context_rows=False,
+ )
+ selected = torch.empty(
+ [q_rows, core.token_topk + core.compress_ratio - 1],
+ dtype=torch.int32,
+ device=qkv.device,
+ )
+ _selection(
+ core,
+ scores,
+ visible_blocks,
+ q_logical,
+ seq_lens_q,
+ topk_indices,
+ row_starts,
+ selected,
+ _live_score_columns(core, md, batch),
+ )
+ if share:
+ self._capture_selection(md, topk_indices, visible_blocks, compact=True)
+ q_att = _packed_q(core, qkv)[keep]
+ attn = torch.empty([q_rows, core.q_width], dtype=qkv.dtype, device=qkv.device)
+ if kv.k_scale is not None and core.mapping.world_size == 1:
+ # One row per context request, then `p` per generation request: the split-K
+ # kernel for the first, the verify kernel for the rest where it serves.
+ _sparse_gqa_nvfp4_context(
+ core,
+ kv,
+ None,
+ q_att,
+ block_table,
+ selected,
+ q_req,
+ q_logical,
+ tokens_per_block,
+ md.num_contexts,
+ attn,
+ p,
+ )
+ else:
+ out = _sparse_gqa(
+ core, kv, q_att, block_table, selected, q_req, tokens_per_block, q_logical
+ )
+ attn.copy_(torch.reshape(out, [q_rows, core.q_width]))
+ gate = torch.reshape(gate[keep], [q_rows, core.q_width])
+ elif sparse:
+ kv = draft.kv
+ _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block)
+ topk_indices = md.qsa_topk_indices[:rows]
+ visible_blocks = md.qsa_visible_blocks[:rows]
+ seq_lens_q = md.qsa_sequence_lengths[:rows]
+ row_starts = md.qsa_topk_row_starts[:rows]
+ parts = [(0, num_ctx, True)]
+ if num_ctx < rows:
+ parts.append((num_ctx, rows, False))
+ selected = torch.empty(
+ [rows, core.token_topk + core.compress_ratio - 1],
+ dtype=torch.int32,
+ device=qkv.device,
+ )
+ for lo, hi, context_rows in parts:
+ scores = qsa_paged_index_scores(
+ q=q_index[lo:hi],
+ index_cache=draft.index_cache,
+ block_table=block_table,
+ query_positions=logical[lo:hi],
+ request_indices=req_idx[lo:hi],
+ tokens_per_block=tokens_per_block,
+ compress_ratio=core.compress_ratio,
+ only_visible_blocks=True,
+ context_rows=context_rows,
+ )
+ _selection(
+ core,
+ scores,
+ visible_blocks[lo:hi],
+ logical[lo:hi],
+ seq_lens_q[lo:hi],
+ topk_indices[lo:hi],
+ row_starts[lo:hi],
+ selected[lo:hi],
+ _live_score_columns(core, md, batch),
+ )
+ if share:
+ self._capture_selection(md, topk_indices, visible_blocks)
+ attn = torch.empty([rows, core.q_width], dtype=qkv.dtype, device=qkv.device)
+ for lo, hi, context_rows in parts:
+ out = _sparse_gqa(
+ core,
+ kv,
+ _packed_q(core, qkv)[lo:hi],
+ block_table,
+ selected[lo:hi],
+ req_idx[lo:hi],
+ tokens_per_block,
+ logical[lo:hi] if context_rows else None,
+ )
+ attn[lo:hi].copy_(torch.reshape(out, [hi - lo, core.q_width]))
+ else:
+ attn = thop_attention(
+ q=qkv,
+ output=torch.empty([rows, core.q_width], dtype=qkv.dtype, device=qkv.device),
+ local_layer_idx=layer_idx,
+ num_heads=core.heads_q,
+ num_kv_heads=core.heads_kv,
+ head_size=core.head_dim,
+ quant_mode=core.quant_mode,
+ kv_scale_orig_quant=core.kv_scale_orig_quant,
+ kv_scale_quant_orig=core.kv_scale_quant_orig,
+ **step,
+ **_CALL_CONSTANTS,
+ )
+ mtp_fused_sigmoid_mul_inplace(attn, gate)
+ out = x
+ if keep is not None:
+ # The sparse path attended the kept rows only; the dense call wrote every row.
+ if not sparse:
+ attn = attn[keep]
+ x, inject = x[keep], inject[keep]
+ x = _mtp_tail(core, M, x, attn, inject, pad_rows)
+ return x if keep is None else out.index_copy_(0, keep, x)
+
+
+# ------------------------------------------------------------------- decode ---
+
+
+class DecodeTarget:
+ """Generation-only batches (no context rows): one token per sequence, every op inside
+ the captured decode graph."""
+
+ def __init__(self, core: Qwen4ExpDefaultModelingV2) -> None:
+ self.core = core
+ # See PrefillTarget.__init__.
+ self._contract_pending = step_contract_enabled()
+ self._kv_probe = _KvRangeProbe.for_target(core, self._contract_pending)
+ self._layers = _layer_views(core)
+ self._stack = _stack_views(core)
+ # The trunk step's positions buffer (`_position_key`) and its `pos3`, which the first
+ # draft step of the same forward reads instead of a copy of its own.
+ self._trunk_pos3: tuple | None = None
+
+ def _check_step_contract(self, md) -> None:
+ """Opt-in first-forward fail-fast; see PrefillTarget._check_step_contract."""
+ if not self._contract_pending:
+ return
+ try:
+ self.step_args(md)
+ except AttributeError as exc:
+ raise AssertionError(f"attention metadata surface drifted: {exc}") from exc
+ self._contract_pending = False
+
+ def step_args(self, md: TrtllmAttentionMetadata) -> dict:
+ """See PrefillTarget.step_args. Context counts are stated, not read: routing
+ guarantees them, and a decode-only CUDA graph would freeze whatever it read."""
+ return _step_args(md, num_contexts=0, num_ctx_tokens=0)
+
+ # --------------------------------------------------------------- blocks ---
+
+ def _attention_prologue(self, proj, q_index, index_view, index_cache, L, qsa, side, stage):
+ """At one rank, without the fused prologue: the norm + rope of both planes on this stream
+ and the QSA index branch (the verify step's staging, the raw-key store, then compression)
+ on the side stream, which reads only the projection's index-key columns and the index
+ pools; the attention core joins it before reading either."""
+ core = self.core
+ pools = core._bound.pools
+ num_tokens = proj.shape[0]
+ req_idx, logical, _, _, _, _, block_table, tokens_per_block, pos3, _, coords, _ = qsa
+ side.wait_stream(torch.cuda.current_stream())
+ qwen4_exp_qk_norm_rope(
+ proj,
+ q_index,
+ pos3,
+ core._cos_sin,
+ L.q_norm,
+ L.k_norm,
+ L.idx_qn,
+ num_heads_q=core.heads_q,
+ num_heads_k=core.heads_kv,
+ head_dim=core.head_dim,
+ index_heads=core.index_heads,
+ index_head_dim=core.index_head_dim,
+ index_q_col=core.qkv_width + core.q_width,
+ rotary_dim=core.rotary_dim,
+ eps=core.eps,
+ mrope_section1=core.mrope_section1,
+ mrope_section2=core.mrope_section2,
+ index_mrope_section1=0,
+ index_mrope_section2=0,
+ )
+ with torch.cuda.stream(side):
+ token_k = torch.split(
+ torch.reshape(
+ index_view,
+ [
+ num_tokens,
+ core.index_scratch_width // core.index_head_dim,
+ core.index_head_dim,
+ ],
+ ),
+ [core.index_heads, core.index_kv_heads, 1],
+ 1,
+ )[1]
+ if stage:
+ core.spec_rollback.stage_index(index_cache, pools.position_cache)
+ qsa_paged_index_store(
+ token_k,
+ coords,
+ index_cache,
+ pools.position_cache,
+ req_idx,
+ logical,
+ block_table,
+ tokens_per_block,
+ )
+ qsa_prefill_compress(
+ logical_positions=logical,
+ request_indices=req_idx,
+ block_table=block_table,
+ index_cache=index_cache,
+ position_cache=pools.position_cache,
+ k_norm_weight=L.idx_kn,
+ cos_sin=core._cos_sin,
+ eps=core.eps,
+ tokens_per_block=tokens_per_block,
+ compress_ratio=core.compress_ratio,
+ mrope_section=_QSA_MROPE_SECTION,
+ )
+
+ def _attention_one_rank(self, x, i, L, qkv, index_scratch, attn_out, qsa, step, stage):
+ """`_attention_block` at one rank, where a layer holds every head: the input
+ projection's output row is `[q | k | v | gate | index q, k | v slot]`. Over a bf16 or
+ NVFP4 pool above the dense threshold everything up to the index scores is one launch,
+ `one_rank_attn_decode_prologue` (over an NVFP4 pool it quantizes the K/V append; on a
+ single-token step of up to `_PROLOGUE_SCORE_ROWS` rows it writes the index scores too);
+ otherwise `_attention_prologue`. The sparse path reads q/k/v as views of the
+ projection, selects with `qsa_topk_expand_decode` where it serves, and fuses the
+ output gate into the sparse GQA; the dense path takes one compact copy of `[q | k | v]`.
+ A verify step stages the index slots for `SpeculativeRollback` first."""
+ core = self.core
+ pools = core._bound.pools
+ (
+ req_idx,
+ logical,
+ seq_lens_q,
+ visible_blocks,
+ topk_indices,
+ topk_row_starts,
+ block_table,
+ tokens_per_block,
+ pos3,
+ _,
+ coords,
+ sparse,
+ ) = qsa
+ num_tokens = x.shape[0]
+ main = torch.cuda.current_stream()
+ side = core._side_stream
+ index_cache = pools.index_cache[i]
+ kv = pools.kv[i]
+ rows_per_request = step["predicted_tokens_per_seq"]
+ fused_prologue = sparse and (kv.k_scale is not None or kv.k.dtype == x.dtype)
+ if fused_prologue and stage:
+ side.wait_stream(main)
+ with torch.cuda.stream(side):
+ core.spec_rollback.stage_index(index_cache, pools.position_cache)
+ proj = _dense_mm(
+ x, L.attn_in, L.attn_in8, core.decode_mm, "in", L.attn_in_g, mid=L.attn_in_m
+ )
+ qkv_view, gate, index_view = torch.split(
+ proj, [core.qkv_width, core.q_width, core.index_scratch_width], 1
+ )
+ q_dst, k_dst, v_dst = _attn_in_heads(core, proj)
+ q_index = _index_query(core, index_scratch)
+ scores = None
+ if fused_prologue and not stage and num_tokens <= _PROLOGUE_SCORE_ROWS:
+ scores = torch.empty(
+ (num_tokens, block_table.shape[1] * tokens_per_block // core.compress_ratio),
+ dtype=torch.float32,
+ device=proj.device,
+ )
+ if fused_prologue:
+ if stage:
+ main.wait_stream(side)
+ one_rank_attn_decode_prologue(
+ proj,
+ q_index,
+ pos3,
+ coords,
+ core._cos_sin,
+ L.q_norm,
+ L.k_norm,
+ L.idx_qn,
+ L.idx_kn,
+ kv.k,
+ kv.v,
+ index_cache,
+ pools.position_cache,
+ req_idx,
+ logical,
+ block_table,
+ tokens_per_block,
+ core.compress_ratio,
+ num_heads_q=core.heads_q,
+ num_heads_k=core.heads_kv,
+ head_dim=core.head_dim,
+ index_heads=core.index_heads,
+ index_head_dim=core.index_head_dim,
+ index_q_col=core.qkv_width + core.q_width,
+ rotary_dim=core.rotary_dim,
+ eps=core.eps,
+ mrope_section1=core.mrope_section1,
+ mrope_section2=core.mrope_section2,
+ index_mrope_section1=0,
+ index_mrope_section2=0,
+ compress_mrope_section=_QSA_MROPE_SECTION,
+ k_scale_cache=kv.k_scale,
+ v_scale_cache=kv.v_scale,
+ kv_scale_orig_quant=None if kv.k_scale is None else core.kv_scale_orig_quant,
+ rows_per_request=rows_per_request,
+ index_scores=scores,
+ )
+ else:
+ self._attention_prologue(proj, q_index, index_view, index_cache, L, qsa, side, stage)
+ if self._kv_probe is not None and not self._kv_probe.observe(i, k_dst, v_dst):
+ self._kv_probe = None
+ gated = False
+ if sparse:
+ if not fused_prologue:
+ main.wait_stream(side)
+ _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block)
+ if scores is None:
+ scores = qsa_paged_index_scores(
+ q=q_index,
+ index_cache=index_cache,
+ block_table=block_table,
+ query_positions=logical,
+ request_indices=req_idx,
+ tokens_per_block=tokens_per_block,
+ compress_ratio=core.compress_ratio,
+ only_visible_blocks=True,
+ context_rows=False,
+ use_pdl=fused_prologue,
+ prefetch=fused_prologue and not stage and num_tokens <= _SPARSE_DECODE_MAX_ROWS,
+ rows_per_request=rows_per_request,
+ )
+ if _topk_expand_supports(scores):
+ selected = qsa_topk_expand_decode(
+ scores,
+ visible_blocks,
+ logical,
+ seq_lens_q,
+ topk_indices,
+ core.compress_ratio,
+ core.token_topk,
+ )
+ else:
+ selected = torch.empty(
+ [num_tokens, core.token_topk + core.compress_ratio - 1],
+ dtype=torch.int32,
+ device=proj.device,
+ )
+ _selection(
+ core,
+ scores,
+ visible_blocks,
+ logical,
+ seq_lens_q,
+ topk_indices,
+ topk_row_starts,
+ selected,
+ )
+ if kv.k_scale is not None:
+ attn = _sparse_gqa(
+ core,
+ kv,
+ q_dst,
+ block_table,
+ selected,
+ req_idx,
+ tokens_per_block,
+ None,
+ output_gate=gate,
+ rows_per_request=rows_per_request,
+ )
+ gated = True
+ else:
+ gated = qsa_paged_sparse_gqa.fuses_output_gate(q_dst, kv.k, selected, gate)
+ if (
+ gated
+ and num_tokens <= _SPARSE_DECODE_MAX_ROWS
+ and _sparse_decode_supports(q_dst, kv.k, selected)
+ ):
+ attn = qsa_paged_sparse_gqa_decode(
+ q_dst,
+ kv.k,
+ kv.v,
+ block_table,
+ selected,
+ req_idx,
+ tokens_per_block,
+ core.softmax_scale,
+ gate,
+ )
+ else:
+ attn = qsa_paged_sparse_gqa(
+ q=q_dst,
+ k_cache=kv.k,
+ v_cache=kv.v,
+ block_table=block_table,
+ selected_tokens=selected,
+ request_indices=req_idx,
+ tokens_per_block=tokens_per_block,
+ softmax_scale=core.softmax_scale,
+ query_positions=None,
+ compress_ratio=None,
+ output_gate=gate if gated else None,
+ )
+ out = torch.reshape(attn, [num_tokens, core.q_width])
+ else:
+ main.wait_stream(side)
+ qkv.copy_(qkv_view)
+ out = thop_attention(
+ q=qkv,
+ output=attn_out,
+ local_layer_idx=pools.pool_idx[i],
+ num_heads=core.heads_q,
+ num_kv_heads=core.heads_kv,
+ head_size=core.head_dim,
+ quant_mode=core.quant_mode,
+ kv_scale_orig_quant=core.kv_scale_orig_quant,
+ kv_scale_quant_orig=core.kv_scale_quant_orig,
+ **step,
+ **_CALL_CONSTANTS,
+ )
+ if not gated:
+ fused_sigmoid_mul_inplace(out, gate)
+ return _dense_mm(out, L.o, L.o8, core.decode_mm, "out", L.o_g, mid=L.o_m)
+
+ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step, stage):
+ """One gated GQA layer plus its QSA index branch for a generation-only batch, at this
+ rank's heads.
+
+ The four input projections are one GEMM against the fused operand; the norm + rope of
+ both planes runs on it; one compact copy of its leading `[q | k | v]` feeds the dense
+ op and the sparse GQA's query heads, the sparse K/V append reads k and v as views of
+ it, and the index branch reads the raw key from it. The index plane is maintained on
+ every step, at both ends of the threshold; above it the dense core is replaced and the
+ main K/V append becomes this layer's own call; over a bf16 or e4m3 pool all of this up
+ to the index scores is one launch (`qwen4_exp_attn_decode_prologue`), which writes the
+ packed buffer and the index query itself. A verify step's rows (`stage`) come
+ `predicted_tokens_per_seq` a request, each scored and attended on its own, and their
+ index slots are staged for `SpeculativeRollback` before the store. At one rank
+ `_attention_one_rank` runs instead."""
+ core = self.core
+ if core.mapping.world_size == 1:
+ return self._attention_one_rank(x, i, L, qkv, index_scratch, attn_out, qsa, step, stage)
+ (
+ req_idx,
+ logical,
+ seq_lens_q,
+ visible_blocks,
+ topk_indices,
+ topk_row_starts,
+ block_table,
+ tokens_per_block,
+ pos3,
+ pos,
+ coords,
+ sparse,
+ ) = qsa
+ num_tokens = x.shape[0]
+ proj = _dense_mm(
+ x, L.attn_in, L.attn_in8, core.decode_mm, "in", L.attn_in_g, mid=L.attn_in_m
+ )
+ qkv_view, gate, _ = torch.split(
+ proj, [core.qkv_width, core.q_width, core.index_scratch_width], 1
+ )
+ pools = core._bound.pools
+ # Over a bf16 or e4m3 pool the sparse path's norm + rope, packed copy, K/V append,
+ # index store and compression are one launch; it writes the packed buffer itself.
+ prologue = sparse and pools.kv[i].k_scale is None
+ if stage:
+ core.spec_rollback.stage_index(pools.index_cache[i], pools.position_cache)
+ if prologue:
+ kv = pools.kv[i]
+ qwen4_exp_attn_decode_prologue(
+ proj,
+ qkv,
+ _index_query(core, index_scratch),
+ pos3,
+ coords,
+ core._cos_sin,
+ L.q_norm,
+ L.k_norm,
+ L.idx_qn,
+ L.idx_kn,
+ kv.k,
+ kv.v,
+ pools.index_cache[i],
+ pools.position_cache,
+ req_idx,
+ logical,
+ block_table,
+ tokens_per_block,
+ core.compress_ratio,
+ core.heads_q,
+ core.heads_kv,
+ core.head_dim,
+ core.index_heads,
+ core.index_head_dim,
+ core.index_col,
+ core.rotary_dim,
+ core.eps,
+ core.mrope_section1,
+ core.mrope_section2,
+ rows_per_request=step["predicted_tokens_per_seq"],
+ )
+ else:
+ _, k_dst, v_dst = _attn_in_heads(core, proj)
+ _attention_rope(core, proj, index_scratch, L, pos3, pos)
+ _index_plane(
+ core, proj, i, L, (req_idx, logical, block_table, tokens_per_block, coords)
+ )
+ # The packed `[q | k | v]` as one compact copy: the dense op reads it flat, and
+ # the sparse GQA reads its query heads from it.
+ qkv.copy_(qkv_view)
+ if self._kv_probe is not None:
+ # Both branches leave the normalized, rotated K/V in the packed copy.
+ k_rows, v_rows = torch.split(
+ torch.reshape(qkv, [qkv.shape[0], -1, core.head_dim])[:, core.heads_q :],
+ core.heads_kv,
+ 1,
+ )
+ if not self._kv_probe.observe(i, k_rows, v_rows):
+ self._kv_probe = None
+
+ if sparse:
+ kv = pools.kv[i]
+ if not prologue:
+ _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block)
+ scores = qsa_paged_index_scores(
+ q=_index_query(core, index_scratch),
+ index_cache=core._bound.pools.index_cache[i],
+ block_table=block_table,
+ query_positions=logical,
+ request_indices=req_idx,
+ tokens_per_block=tokens_per_block,
+ compress_ratio=core.compress_ratio,
+ only_visible_blocks=True,
+ context_rows=False,
+ )
+ selected = torch.empty(
+ [num_tokens, core.token_topk + core.compress_ratio - 1],
+ dtype=torch.int32,
+ device=proj.device,
+ )
+ _selection(
+ core,
+ scores,
+ visible_blocks,
+ logical,
+ seq_lens_q,
+ topk_indices,
+ topk_row_starts,
+ selected,
+ )
+ attn = _sparse_gqa(
+ core,
+ kv,
+ _packed_q(core, qkv),
+ block_table,
+ selected,
+ req_idx,
+ tokens_per_block,
+ None,
+ )
+ out = torch.reshape(attn, [num_tokens, core.q_width])
+ else:
+ out = thop_attention(
+ q=qkv,
+ output=attn_out,
+ local_layer_idx=core._bound.pools.pool_idx[i],
+ num_heads=core.heads_q,
+ num_kv_heads=core.heads_kv,
+ head_size=core.head_dim,
+ quant_mode=core.quant_mode,
+ kv_scale_orig_quant=core.kv_scale_orig_quant,
+ kv_scale_quant_orig=core.kv_scale_quant_orig,
+ **step,
+ **_CALL_CONSTANTS,
+ )
+ fused_sigmoid_mul_inplace(out, gate)
+ return core.tp_sum(_dense_mm(out, L.o, L.o8, core.decode_mm, "out", L.o_g, mid=L.o_m))
+
+ def _hc_site(self, hc, block_out, inject, norm, first, up, with_inject, normed=None, moe=None):
+ """One hyper-connection site: the previous block's injection (none at the stack's first
+ site and after the PLE add), this site's grouped norm and its mixer. Returns `(hc, mixed,
+ inject)`. At one rank (`hc_fast_sites`) the row count, a per-capture constant, picks the
+ form: one row runs the CuTe GEMV pair, up to `_HC_DECODE_ROWS_MAX` rows the tcgen05 rows
+ pair, larger batches `qwen4_exp_hc_inject_norm` and `_hc_mixer`, which every batch takes
+ at the other layouts. `normed`, when given, is this site's grouped
+ norm of `hc` already computed (the PLE tail's, no injection pending), which only the
+ inject-norm route reads (the first kernels of the other two compute it themselves).
+
+ After a MoE block `block_out` is its routed output and `moe` the `(gate_logits, shared,
+ topk_weights, topk_rows)` `_moe` left uncombined: the site's first kernel forms the
+ block output itself; with the top-k pair `block_out` is the runner's permuted FC2 rows
+ and that kernel also forms their finalize."""
+ core = self.core
+ gate_logits, shared, topk_w, topk_rows = moe if moe is not None else (None,) * 4
+ n_cols = core.hc_lowrank + (core.hc_count if with_inject else 0)
+ fast = core.hc_fast_sites
+ if fast and _HC_DECODE_GEMV_MAX_ROWS < hc.shape[0] <= _HC_DECODE_ROWS_MAX:
+ hc, low, inject, rstd = qwen4_exp_hc_decode_rows_first(
+ hc,
+ block_out,
+ inject,
+ norm,
+ first,
+ core.hc_count,
+ n_cols,
+ core.eps,
+ shared,
+ gate_logits,
+ topk_w,
+ topk_rows,
+ )
+ mixed, _, _ = qwen4_exp_hc_decode_rows_mix(low, rstd, up, hc, norm, core.hc_count)
+ return hc, mixed, inject
+ if fast and hc.shape[0] <= _HC_DECODE_GEMV_MAX_ROWS:
+ hc, part, ss = qwen4_exp_hc_decode_first_cute(
+ hc,
+ block_out,
+ inject,
+ norm,
+ first,
+ core.hc_count,
+ n_cols,
+ shared,
+ gate_logits,
+ topk_w,
+ topk_rows,
+ )
+ mixed, inject = qwen4_exp_hc_decode_mix_cute(
+ part, ss, up, hc, norm, core.eps, core.hc_count, with_inject
+ )
+ return hc, mixed, inject
+ if normed is None:
+ hc, normed = qwen4_exp_hc_inject_norm(
+ hc,
+ block_out,
+ inject,
+ norm,
+ core.eps,
+ core.hc_count,
+ shared,
+ gate_logits,
+ topk_w,
+ topk_rows,
+ )
+ mixed, inject = _hc_mixer(core, normed, first, up, with_inject)
+ return hc, mixed, inject
+
+ def _mtp_layer(
+ self,
+ embed_tokens,
+ input_ids,
+ hidden_states,
+ md,
+ step,
+ draft,
+ layer_idx,
+ rows,
+ gen_seqs,
+ p,
+ sparse,
+ share,
+ reuse,
+ pos3,
+ coords,
+ ) -> torch.Tensor:
+ """The draft layer at one rank on the decode machinery: the front end's two projections,
+ both hyper-connection sites as `_hc_site` (the front end's per-stream `x_h + x_e` is the
+ first site's injection at unit weight, the attention output the second's), the attention
+ as `_mtp_attention`, the MoE as `_mtp_moe`. Returns the stream with the MoE output
+ injected."""
+ core = self.core
+ M, G = core._mtp, core._mtp_gemv8
+ gemm = core.decode_mm if md.is_cuda_graph else cublas_mm
+ e = _draft_embedding(input_ids, embed_tokens)
+ x_e = _draft_mm(
+ flashinfer_gemma_rmsnorm(e, M.pre_e_norm, core.eps), M.fc_e, G.get("fc_e"), gemm
+ )
+ _, h_n = qwen4_exp_hc_inject_norm(
+ hidden_states, None, None, M.pre_h_norm, core.eps, core.hc_count
+ )
+ x_h = torch.reshape(
+ _draft_mm(
+ torch.reshape(h_n, [rows * core.hc_count, core.hidden]), M.fc_h, G.get("fc_h"), gemm
+ ),
+ [rows, core.stream],
+ )
+ x, mixed, inject = self._hc_site(
+ x_h, x_e, M.unit_inject[:rows], M.ahc_norm, M.ahc_first, M.ahc_up, True
+ )
+ attn = self._mtp_attention(
+ mixed,
+ md,
+ step,
+ draft,
+ layer_idx,
+ rows,
+ gen_seqs,
+ p,
+ sparse,
+ share,
+ reuse,
+ pos3,
+ coords,
+ gemm,
+ )
+ x, mixed, inject = self._hc_site(x, attn, inject, M.mhc_norm, M.mhc_first, M.mhc_up, True)
+ return qwen4_exp_post_inject(x, self._mtp_moe(mixed), inject)
+
+ def _mtp_attention(
+ self,
+ x,
+ md,
+ step,
+ draft,
+ layer_idx,
+ rows,
+ gen_seqs,
+ p,
+ sparse,
+ share,
+ reuse,
+ pos3,
+ coords,
+ gemm,
+ ) -> torch.Tensor:
+ """The draft layer's gated GQA and QSA index branch over the draft manager's planes
+ (`draft`), as `_attention_one_rank` runs a trunk layer: one input projection, the norm +
+ rope of both planes in one launch at the module's rope base, the index store and
+ compression on the side stream, the K/V append, the index scores and the one-launch
+ top-k + expansion, the gated sparse GQA, the output projection. Under IndexShare a draft
+ decode step (`reuse`) skips the index branch and rebuilds its selection from the
+ draft-extend step's capture; the draft-extend step (`share`) captures it."""
+ core = self.core
+ M = core._mtp
+ dt, dev = x.dtype, x.device
+ proj = _draft_mm(x, core._mtp_attn_in, core._mtp_gemv8.get("attn_in"), gemm)
+ qkv_view, gate, index_view = torch.split(
+ proj, [core.qkv_width, core.q_width, core.index_scratch_width], 1
+ )
+ q_dst, k_dst, v_dst = _attn_in_heads(core, proj)
+ index_scratch = torch.empty([rows, core.index_scratch_width], dtype=dt, device=dev)
+ q_index = _index_query(core, index_scratch)
+ req_idx = md.qsa_req_idx_per_token[:rows]
+ logical = md.qsa_logical_positions[:rows]
+ block_table = md.qsa_block_table
+ tokens_per_block = md.tokens_per_block
+ main = torch.cuda.current_stream()
+ side = core._side_stream
+ side.wait_stream(main)
+ qwen4_exp_qk_norm_rope(
+ proj,
+ q_index,
+ pos3,
+ core._mtp_cos_sin,
+ M.q_norm,
+ M.k_norm,
+ M.idx_qn,
+ num_heads_q=core.heads_q,
+ num_heads_k=core.heads_kv,
+ head_dim=core.head_dim,
+ index_heads=core.index_heads,
+ index_head_dim=core.index_head_dim,
+ index_q_col=core.qkv_width + core.q_width,
+ rotary_dim=core.rotary_dim,
+ eps=core.eps,
+ mrope_section1=core.mrope_section1,
+ mrope_section2=core.mrope_section2,
+ index_mrope_section1=0,
+ index_mrope_section2=0,
+ )
+ if not reuse:
+ with torch.cuda.stream(side):
+ token_k = torch.split(
+ torch.reshape(
+ index_view,
+ [
+ rows,
+ core.index_scratch_width // core.index_head_dim,
+ core.index_head_dim,
+ ],
+ ),
+ [core.index_heads, core.index_kv_heads, 1],
+ 1,
+ )[1]
+ qsa_paged_index_store(
+ token_k,
+ coords,
+ draft.index_cache,
+ draft.position_cache,
+ req_idx,
+ logical,
+ block_table,
+ tokens_per_block,
+ )
+ qsa_prefill_compress(
+ logical_positions=logical,
+ request_indices=req_idx,
+ block_table=block_table,
+ index_cache=draft.index_cache,
+ position_cache=draft.position_cache,
+ k_norm_weight=M.idx_kn,
+ cos_sin=core._mtp_cos_sin,
+ eps=core.eps,
+ tokens_per_block=tokens_per_block,
+ compress_ratio=core.compress_ratio,
+ mrope_section=_QSA_MROPE_SECTION,
+ )
+ main.wait_stream(side)
+ gated = False
+ if sparse:
+ kv = draft.kv
+ _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block)
+ topk_indices = md.qsa_topk_indices[:rows]
+ visible_blocks = md.qsa_visible_blocks[:rows]
+ seq_lens_q = md.qsa_sequence_lengths[:rows]
+ if reuse:
+ slot = draft.share_slot
+ assert slot in md.qsa_shared_topk_captured_slots, (
+ "a draft decode step would reuse a selection that no draft-extend step of "
+ "this loop captured"
+ )
+ qsa_mtp_reuse_topk(
+ shared_indices=md.qsa_shared_topk_indices[slot],
+ shared_visible_blocks=md.qsa_shared_topk_visible_blocks[slot],
+ request_indices=req_idx,
+ visible_blocks=visible_blocks,
+ output=topk_indices,
+ )
+ selected = qsa_expand_block_indices(
+ topk_indices,
+ logical,
+ seq_lens_q,
+ compress_ratio=core.compress_ratio,
+ token_topk=core.token_topk,
+ )
+ else:
+ scores = qsa_paged_index_scores(
+ q=q_index,
+ index_cache=draft.index_cache,
+ block_table=block_table,
+ query_positions=logical,
+ request_indices=req_idx,
+ tokens_per_block=tokens_per_block,
+ compress_ratio=core.compress_ratio,
+ only_visible_blocks=True,
+ context_rows=False,
+ rows_per_request=p,
+ )
+ if _topk_expand_supports(scores):
+ selected = qsa_topk_expand_decode(
+ scores,
+ visible_blocks,
+ logical,
+ seq_lens_q,
+ topk_indices,
+ core.compress_ratio,
+ core.token_topk,
+ )
+ else:
+ selected = torch.empty(
+ [rows, core.token_topk + core.compress_ratio - 1],
+ dtype=torch.int32,
+ device=dev,
+ )
+ _selection(
+ core,
+ scores,
+ visible_blocks,
+ logical,
+ seq_lens_q,
+ topk_indices,
+ md.qsa_topk_row_starts[:rows],
+ selected,
+ )
+ if share:
+ self._capture_selection(md, topk_indices, visible_blocks, gen_seqs, p)
+ if kv.k_scale is not None:
+ attn = _sparse_gqa(
+ core,
+ kv,
+ q_dst,
+ block_table,
+ selected,
+ req_idx,
+ tokens_per_block,
+ None,
+ output_gate=gate,
+ )
+ gated = True
+ else:
+ gated = qsa_paged_sparse_gqa.fuses_output_gate(q_dst, kv.k, selected, gate)
+ attn = qsa_paged_sparse_gqa(
+ q=q_dst,
+ k_cache=kv.k,
+ v_cache=kv.v,
+ block_table=block_table,
+ selected_tokens=selected,
+ request_indices=req_idx,
+ tokens_per_block=tokens_per_block,
+ softmax_scale=core.softmax_scale,
+ query_positions=None,
+ compress_ratio=None,
+ output_gate=gate if gated else None,
+ )
+ out = torch.reshape(attn, [rows, core.q_width])
+ else:
+ qkv = torch.empty([rows, core.qkv_width], dtype=dt, device=dev)
+ qkv.copy_(qkv_view)
+ out = thop_attention(
+ q=qkv,
+ output=torch.empty([rows, core.q_width], dtype=dt, device=dev),
+ local_layer_idx=layer_idx,
+ num_heads=core.heads_q,
+ num_kv_heads=core.heads_kv,
+ head_size=core.head_dim,
+ quant_mode=core.quant_mode,
+ kv_scale_orig_quant=core.kv_scale_orig_quant,
+ kv_scale_quant_orig=core.kv_scale_quant_orig,
+ **step,
+ **_CALL_CONSTANTS,
+ )
+ if not gated:
+ fused_sigmoid_mul_inplace(out, gate)
+ return _draft_mm(out, M.o, core._mtp_gemv8.get("o"), gemm)
+
+ def _mtp_moe(self, x) -> torch.Tensor:
+ """The draft layer's MoE at one rank: the FP8 block-scale routed experts (or with
+ `_MTP_MOE_NVFP4_ENV` the NVFP4 ones on the trunk's logits-routed runner) and the bf16
+ shared expert (its dense projections on `_draft_mm`), combined as `routed +
+ sigmoid(gate logit) * shared`; the NVFP4 path runs the shared expert on the side
+ stream."""
+ core = self.core
+ M, G = core._mtp, core._mtp_gemv8
+
+ def shared_expert():
+ gu = _draft_mm(x, M.sh_gu, G.get("sh_gu"), cublas_mm)
+ shared = _draft_mm(
+ silu_and_mul_rows(gu, core.shared_inter), M.sh_dn, G.get("sh_dn"), cublas_mm
+ )
+ return shared, cublas_mm(x, M.sh_gate)
+
+ if M.moe_gscale is None:
+ routed = _mtp_moe_block(core, M, x, routed_only=True)
+ shared, gate_logits = shared_expert()
+ return fused_sigmoid_gate_mul_add(routed, gate_logits, shared)
+ main = torch.cuda.current_stream()
+ side = core._side_stream
+ side.wait_stream(main)
+ with torch.cuda.stream(side):
+ shared, gate_logits = shared_expert()
+ routed = _mtp_routed_nvfp4(core, M, x)
+ main.wait_stream(side)
+ return fused_sigmoid_gate_mul_add(routed, gate_logits, shared)
+
+ def _hc_moe_site(self, hc, block_out, inject, L: _LayerWeights):
+ """`_hc_site` at an MoE site (always with injection). Under attention TP (or at one
+ rank) it also returns the routed experts' NVFP4 input (linear scales at the layer's FC1
+ input global scale), written by the mixer's own last launch; under attention DP the
+ window quantizes the group's gathered rows itself and `fp4` is None. Returns `(hc,
+ mixed, inject, fp4)`."""
+ core = self.core
+ if core.dp_size > 1:
+ hc, mixed, inject = self._hc_site(
+ hc, block_out, inject, L.mhc_norm, L.mhc_first, L.mhc_up, True
+ )
+ return hc, mixed, inject, None
+ n_cols = core.hc_lowrank + core.hc_count
+ fast = core.hc_fast_sites
+ if fast and _HC_DECODE_GEMV_MAX_ROWS < hc.shape[0] <= _HC_DECODE_ROWS_MAX:
+ hc, low, inject, rstd = qwen4_exp_hc_decode_rows_first(
+ hc, block_out, inject, L.mhc_norm, L.mhc_first, core.hc_count, n_cols, core.eps
+ )
+ mixed, xq, xsf = qwen4_exp_hc_decode_rows_mix(
+ low, rstd, L.mhc_up, hc, L.mhc_norm, core.hc_count, L.moe_gscale
+ )
+ return hc, mixed, inject, (xq, xsf)
+ if fast and hc.shape[0] <= _HC_DECODE_GEMV_MAX_ROWS:
+ hc, part, ss = qwen4_exp_hc_decode_first_cute(
+ hc, block_out, inject, L.mhc_norm, L.mhc_first, core.hc_count, n_cols
+ )
+ mixed, inject, xq, xsf = qwen4_exp_hc_decode_mix_fp4_cute(
+ part, ss, L.mhc_up, hc, L.mhc_norm, core.eps, core.hc_count, True, L.moe_gscale
+ )
+ return hc, mixed, inject, (xq, xsf)
+ hc, normed = qwen4_exp_hc_inject_norm(
+ hc, block_out, inject, L.mhc_norm, core.eps, core.hc_count
+ )
+ mixed, inject, fp4 = _hc_moe_mixer(core, normed, L)
+ return hc, mixed, inject, fp4
+
+ def _moe(self, x, L: _LayerWeights, pad_rows: int, fp4, fold: bool) -> tuple:
+ """One decode MoE block, left uncombined for the next site: `(routed, (gate_logits,
+ shared, topk_weights, topk_rows))`. With `fold` (one rank, where the window is the whole
+ layer and needs no sum) the runner skips its finalize: `routed` is its permuted FC2 rows
+ and the top-k pair its routing weights and index map, which the next site's first kernel
+ gathers; otherwise the pair is None and `routed` the finalized rows (`_moe_block`)."""
+ core = self.core
+ router_mm = _decode_router_mid(core, L)
+ if not fold:
+ routed, gate_logits, shared = _moe_block(
+ core, x, L, pad_rows, fp4, router_mm, core.decode_mm
+ )
+ return routed, (gate_logits, shared, None, None)
+ main = torch.cuda.current_stream()
+ side = core._side_stream
+ side.wait_stream(main)
+ # From `_ROUTER_FIRST_MIN_ROWS` rows the router launches before the shared expert (the
+ # fork above precedes it, so the side stream does not wait for it).
+ router_first = x.shape[0] >= _ROUTER_FIRST_MIN_ROWS
+ if router_first:
+ logits = router_mm(x, L.router)
+ with torch.cuda.stream(side):
+ shared, gate_logits = _shared_expert(core, x, L, core.decode_mm)
+ if not router_first:
+ logits = router_mm(x, L.router)
+ xq, xsf = fp4
+ rows, topk_w, topk_rows = fp4_block_scale_moe_renorm_permuted(
+ logits,
+ xq,
+ torch.reshape(xsf, [-1]).view(torch.float8_e4m3fn),
+ L.fc1_w,
+ L.fc1_s,
+ L.fc2_w,
+ L.fc2_s,
+ L.out1_scale,
+ L.out1_gate_scale,
+ L.out2_scale,
+ core.num_experts,
+ core.topk,
+ core.expert_inter,
+ )
+ main.wait_stream(side)
+ return rows, (gate_logits, shared, topk_w, topk_rows)
+
+ def _gdn_block(self, x, i, L, si_all, verify, gen_p, captured=False) -> torch.Tensor:
+ """One Gated-DeltaNet layer over a generation-only batch, one token per sequence, at
+ this rank's heads: one GEMM for all four input projections (q|k|v, z, a and b are
+ column views of its output), the conv step over the q|k|v columns in place, then the
+ recurrence and the gated norm as one streaming launch that computes each (row,
+ v-head) on its own, so over a bf16 pool a row's output and state do not depend on its
+ batch. The output projection's per-rank partial is summed over the TP group.
+
+ A verify step (`verify`, `gen_p` rows a sequence) takes `_gdn_verify` and the gated
+ norm epilogue instead.
+
+ With the in_proj's NVFP4 copy (`_GDN_INPROJ_NVFP4_ENV`) every step of up to
+ `_INPROJ_NVFP4_MAX_ROWS` rows reads it; with its packed FP8 copy (`_DENSE_FP8_ENV`) the
+ rows `_dense_mm` gives it. In a captured step (`captured`) those GEMVs run the conv step
+ as their epilogue -- a single-token step's and, up to `_VERIFY_CONV_FUSED_MAX_ROWS`
+ rows, a verify step's with its window snapshots and the pool advance -- reading the
+ pool, taps and slots before their PDL wait, bitwise the projection + conv they
+ replace."""
+ core = self.core
+ pools = core._bound.pools
+ num_tokens = x.shape[0]
+ is_verify = verify is not None
+ nv = (
+ L.in_proj_nv
+ if L.in_proj_nv is not None and num_tokens <= _INPROJ_NVFP4_MAX_ROWS and _gemv_ready(x)
+ else None
+ )
+ nv_conv = nv is not None and captured and not is_verify
+ fp8_conv = (
+ not nv_conv
+ and captured
+ and not is_verify
+ and L.in_proj_g is not None
+ and _gemv_ready(x)
+ and num_tokens <= _FP8_CONV_FUSED_MAX_ROWS
+ )
+ verify_fused = (
+ captured
+ and is_verify
+ and (L.in_proj_g is not None or nv is not None)
+ and _gemv_ready(x)
+ and num_tokens <= min(_FP8_GEMV_ROWS["in"][1], _VERIFY_CONV_FUSED_MAX_ROWS)
+ )
+ if verify_fused:
+ verify_op = (
+ fp8_gemv_conv1d_verify_op if nv is None else nvfp4_split_gemv_conv1d_verify_op
+ )
+ mq_out, zab = verify_op(
+ x,
+ L.in_proj_g if nv is None else nv,
+ pools.conv[i],
+ L.conv_w,
+ si_all,
+ core._bound.int_conv[i],
+ verify[0],
+ core.conv_dim,
+ gen_p,
+ _PAD_SLOT_ID,
+ )
+ z, a, b = torch.split(zab, core.in_proj_split[1:], 1)
+ out = _gdn_verify(
+ core, None, a, b, i, L, si_all, num_tokens // gen_p, gen_p, verify, mq_out
+ )
+ return self._gdn_epilogue(out, z, L, num_tokens)
+ if nv_conv:
+ proj = nvfp4_split_gemv_conv1d_op(
+ x, nv, pools.conv[i], L.conv_w, si_all, core.conv_dim, _PAD_SLOT_ID
+ )
+ elif fp8_conv:
+ proj = fp8_gemv_conv1d_op(
+ x, L.in_proj_g, pools.conv[i], L.conv_w, si_all, core.conv_dim, _PAD_SLOT_ID
+ )
+ elif nv is not None:
+ proj = nvfp4_split_gemv_op(x, nv)
+ else:
+ # FlashInfer's low-M kernels at the rows its table names (a verify step's few rows
+ # a request), cuBLAS past them; the dense FP8 copies where `_dense_mm` takes them.
+ proj = _dense_mm(
+ x,
+ L.in_proj,
+ L.in_proj8,
+ core.decode_in_mm,
+ "in",
+ L.in_proj_g,
+ mid=L.in_proj_m,
+ )
+ if is_verify:
+ mixed_qkv, z, a, b = torch.split(proj, core.in_proj_split, 1)
+ out = _gdn_verify(
+ core,
+ mixed_qkv,
+ a,
+ b,
+ i,
+ L,
+ si_all,
+ num_tokens // gen_p,
+ gen_p,
+ verify,
+ )
+ return self._gdn_epilogue(out, z, L, num_tokens)
+ if not (nv_conv or fp8_conv):
+ mixed_qkv = torch.split(proj, core.in_proj_split, 1)[0]
+ causal_conv1d_update(
+ x=mixed_qkv,
+ conv_state=pools.conv[i],
+ weight=L.conv_w,
+ bias=None,
+ silu_activation=True,
+ cache_seqlens=None,
+ conv_state_indices=si_all,
+ pad_slot_id=_PAD_SLOT_ID,
+ )
+ if (
+ captured
+ and core.mapping.world_size == 1
+ and pools.ssm[i].dtype == torch.float32
+ and num_tokens <= _GDN_DECODE_FUSED_MAX_ROWS
+ ):
+ # At one rank over an fp32 pool, the recurrence and the gated norm as one launch
+ # that reads this layer's pool rows before its PDL wait: in a captured
+ # single-token step nothing ahead of it writes them.
+ o = gdn_decode_gated_norm_pdl(
+ proj,
+ pools.ssm[i],
+ si_all,
+ L.a_log,
+ L.dt_bias,
+ L.gnorm,
+ core.eps,
+ core.num_k_heads,
+ core.head_k_dim,
+ )
+ return _dense_mm(o, L.out, L.out8, core.decode_mm, "out", L.out_g, mid=L.out_m)
+ # It computes both gates and the q/k L2 norm itself, at the softplus pair (1, 20) and
+ # epsilon 1e-6 this target's constants name.
+ o = gdn_decode_rows_gated_norm(
+ proj,
+ pools.ssm[i],
+ si_all,
+ L.a_log,
+ L.dt_bias,
+ L.gnorm,
+ core.eps,
+ core.num_k_heads,
+ core.head_k_dim,
+ )
+ return core.tp_sum(_dense_mm(o, L.out, L.out8, core.decode_mm, "out", L.out_g, mid=L.out_m))
+
+ def _gdn_epilogue(self, out, z, L, num_tokens) -> torch.Tensor:
+ """A verify step's gated RMS norm (sigmoid gate, plain non-gemma weight) and the
+ output projection, whose per-rank partial is summed over the TP group."""
+ core = self.core
+ o = rms_norm_gated_token_major(
+ torch.reshape(out, [-1, core.head_v_dim]),
+ torch.reshape(z, [num_tokens, core.num_v_heads, core.head_v_dim]),
+ weight=L.gnorm,
+ eps=core.eps,
+ fp8_scale=None,
+ gate_activation="sigmoid",
+ )
+ o = torch.reshape(o, [num_tokens, core.value_dim])
+ return core.tp_sum(_dense_mm(o, L.out, L.out8, core.decode_mm, "out", L.out_g, mid=L.out_m))
+
+ def _ple_block(self, hc, input_ids, ple_state, next_norm, stage) -> tuple:
+ """The n-gram side path on its decoder layer, added into the stream and normalized
+ for the layer's attention-side site. Returns `(hc, normed)`. The same steps as the
+ prefill target's `_ple_block`, over up to `rows` rows a sequence; every generation
+ sequence carries its history by construction, so no carry flags are passed. A verify
+ step (`stage`) stages what `SpeculativeRollback` rewrites the two pools from: both
+ pools' pre-step rows before either launch publishes, and the conv input the second
+ launch writes out."""
+ core = self.core
+ S = self._stack
+ conv_states, ngram_pool, si_all, qsl, rows = ple_state
+ if stage:
+ core.spec_rollback.stage_ple_verify(ngram_pool, conv_states, input_ids)
+ ngram = ple_ngram_embed(
+ input_ids,
+ ngram_pool,
+ si_all,
+ None,
+ qsl,
+ core.w["ple_table"],
+ S.ple_tscale,
+ S.ple_mult,
+ S.ple_vocab,
+ S.ple_offs,
+ core.eos_token_id,
+ rows,
+ )
+ conv_input = torch.empty_like(hc) if stage else None
+ out = ple_gate_conv_inject_norm(
+ hc,
+ cublas_mm(ngram, S.ple_key),
+ cublas_mm(ngram, S.ple_val),
+ S.ple_nk_w,
+ S.ple_nq_w,
+ S.ple_nc_w,
+ S.ple_conv_w,
+ conv_states,
+ si_all,
+ qsl,
+ next_norm,
+ core.eps,
+ core.hc_count,
+ _PLE_CLAMP_MIN,
+ core.ple_dilation,
+ _PAD_SLOT_ID,
+ rows,
+ conv_input=conv_input,
+ )
+ if stage:
+ core.spec_rollback.stage_ple_conv_input(conv_input)
+ return out
+
+ # -------------------------------------------------------------- forward ---
+
+ def forward(
+ self,
+ attn_metadata: AttentionMetadata,
+ input_ids: torch.IntTensor | None = None,
+ position_ids: torch.IntTensor | None = None,
+ **kwargs,
+ ):
+ core = self.core
+ md = attn_metadata
+ self._check_step_contract(md)
+ step = self.step_args(md)
+ spec_metadata = kwargs.get("spec_metadata")
+ mm = md.mamba_metadata
+ # `coords` is the `[N, 3]` view of the engine's positions the position plane stores;
+ # `pos3` the contiguous `[3, N]` copy, built below (by `verify_step_setup` on a
+ # captured verify step).
+ coords = torch.transpose(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1]), 0, 1)
+
+ h = nn.functional.embedding(input_ids, core.w["embed"])
+ num_tokens = h.shape[0]
+ dt = h.dtype
+ dev = h.device
+ batch = md.num_seqs
+ # Decode CUDA graphs pad with real dummy requests, so they are in the metadata.
+ assert num_tokens == md.num_tokens, (
+ f"the engine handed {num_tokens} activation rows for a generation batch its "
+ f"metadata describes as {md.num_tokens} tokens"
+ )
+ step["predicted_tokens_per_seq"] = 1
+ # Zero rows under a decode CUDA graph, which pads every attention-DP rank to the
+ # captured batch.
+ pad_rows = core._dp_pad_rows(md, num_tokens)
+ # Rows per generation sequence: 1 without speculation, else `runtime_draft_len + 1`.
+ gen_p, is_verify = _verify_rows(core, md, spec_metadata, batch, num_tokens)
+ step["predicted_tokens_per_seq"] = gen_p
+ si_all = mm.state_indices[:batch]
+ verify = _verify_state(core, mm, batch) if is_verify else None
+ pools = core._bound.pools
+ # A captured verify step builds `pos3`, the PLE segmentation vector and the rollback's
+ # staging in one launch (`verify_step_setup`).
+ rollback = core.spec_rollback
+ fused_setup = (
+ is_verify
+ and md.is_cuda_graph
+ and batch < _VERIFY_SETUP_MAX_SEQS
+ and position_ids.stride(-1) == 1
+ and rollback.index_planes is not None
+ )
+ if fused_setup:
+ pos3, ple_qsl, staged_slots, staged_pages, staged_within, staged_index = (
+ verify_step_setup(
+ position_ids,
+ md.seq_lens_cuda[:batch],
+ si_all,
+ md.qsa_req_idx_per_token[:num_tokens],
+ md.qsa_logical_positions[:num_tokens],
+ md.qsa_block_table,
+ md.tokens_per_block,
+ *rollback.index_planes,
+ )
+ )
+ else:
+ pos3 = _copy_positions(position_ids)
+ # Under speculation `query_start_loc_long` has one row per generation sequence:
+ # with `gen_p > 1` the PLE segmentation is rebuilt from `seq_lens_cuda`.
+ ple_qsl = (
+ mm.query_start_loc_long
+ if gen_p == 1
+ else cu_seqlens_from_seq_lens(md.seq_lens_cuda[:batch], torch.int64)
+ )
+ pos = pos3[0]
+ self._trunk_pos3 = (_position_key(position_ids), pos3)
+ ple_state = (pools.ple_state[0], pools.ple_state[1], si_all, ple_qsl, gen_p)
+
+ # Staging for the speculative rollback (see PrefillTarget.prepare); every row here is
+ # a generation row. A captured step stages and commits the PLE pools in one launch
+ # each.
+ if rollback is not None:
+ rollback.discard()
+ if fused_setup:
+ rollback.begin_staged(
+ num_gens=batch,
+ rows_per_gen=gen_p,
+ slots=staged_slots,
+ pages=staged_pages,
+ within=staged_within,
+ index_all=staged_index,
+ )
+ elif is_verify:
+ rollback.begin(
+ num_contexts=0,
+ prefix_tokens=0,
+ num_gens=batch,
+ rows_per_gen=gen_p,
+ slots=si_all,
+ request_indices=md.qsa_req_idx_per_token[:num_tokens],
+ logical_positions=md.qsa_logical_positions[:num_tokens],
+ block_table=md.qsa_block_table,
+ tokens_per_block=md.tokens_per_block,
+ fused=md.is_cuda_graph,
+ )
+
+ # Dense/sparse decision once per forward, from a host-side metadata read: a
+ # per-capture constant matching the engine's short/long decode graph families.
+ sparse = batch > 0 and int(md.kv_lens_runtime[:batch].max()) > core.dense_threshold
+ qsa = (
+ md.qsa_req_idx_per_token[:num_tokens],
+ md.qsa_logical_positions[:num_tokens],
+ md.qsa_sequence_lengths[:num_tokens],
+ md.qsa_visible_blocks[:num_tokens],
+ md.qsa_topk_indices[:num_tokens],
+ md.qsa_topk_row_starts[:num_tokens],
+ md.qsa_block_table,
+ md.tokens_per_block,
+ pos3,
+ pos,
+ coords,
+ sparse,
+ )
+
+ # Scratch shared by all attention layers: thop_attention overwrites rows
+ # [:num_tokens] of `output` in place, and the packed buffers are fully rewritten every
+ # layer (the QKV one on the dense path only).
+ qkv = torch.empty([num_tokens, core.qkv_width], dtype=dt, device=dev)
+ index_scratch = torch.empty([num_tokens, core.index_scratch_width], dtype=dt, device=dev)
+ attn_out = torch.empty([num_tokens, core.q_width], dtype=dt, device=dev)
+
+ hc_count = core.hc_count
+ hc = torch.cat([h] * hc_count, -1)
+ layers, S = self._layers, self._stack
+ # Each block's output and injection go into the next site's first kernel (`_hc_site`),
+ # after an MoE block with its combine (and at one rank its finalize, `fold`); before
+ # the PLE layer the pending block is added here, the PLE tail taking the site's norm.
+ block_out = inject = moe = None
+ for i, L in enumerate(layers):
+ normed = None
+ if i == core.ple_layer:
+ if block_out is not None:
+ if moe is not None:
+ block_out = fused_sigmoid_gate_mul_add(block_out, moe[0], moe[1])
+ hc = qwen4_exp_post_inject(hc, block_out, inject)
+ block_out = inject = moe = None
+ hc, normed = self._ple_block(hc, input_ids, ple_state, L.ahc_norm, is_verify)
+ hc, mixed, inject = self._hc_site(
+ hc, block_out, inject, L.ahc_norm, L.ahc_first, L.ahc_up, True, normed, moe
+ )
+ if core.is_attn[i]:
+ block_out = self._attention_block(
+ mixed, i, L, qkv, index_scratch, attn_out, qsa, step, is_verify
+ )
+ else:
+ block_out = self._gdn_block(mixed, i, L, si_all, verify, gen_p, md.is_cuda_graph)
+ hc, mixed, inject, fp4 = self._hc_moe_site(hc, block_out, inject, L)
+ block_out, moe = self._moe(
+ mixed, L, pad_rows, fp4, core.moe_fold and i + 1 != core.ple_layer
+ )
+ # The stack-final site with no injection weight: its mix *is* the final norm.
+ hc, final, _ = self._hc_site(
+ hc, block_out, inject, S.fhc_norm, S.fhc_first, S.fhc_up, False, moe=moe
+ )
+ if spec_metadata is None:
+ return final
+ # `lm_head` reads the `[T, hidden]` mix; the MTP module reads the `[T, stream]` stream
+ # it was collapsed from.
+ return final, hc
+
+ # ------------------------------------------------------------- the draft ---
+
+ def _capture_selection(self, md, topk_indices, visible_blocks, batch, p) -> None:
+ """IndexShare: stash, per request, the selection of its last accepted draft-extend
+ row with that row's complete-group count, in the draft layer's row of the shared
+ buffers. Every request carries `p` rows here, and the worker continues the draft loop
+ from row `num_accepted - 1` of each (the golden token counts as accepted)."""
+ assert md.qsa_mtp_num_accepted is not None, (
+ "IndexShare needs the worker's accepted counts to pick each request's capture row; "
+ "the draft loop did not supply them"
+ )
+ accepted = md.qsa_mtp_num_accepted[:batch].to(torch.int64)
+ rows = torch.arange(batch, dtype=torch.int64, device=topk_indices.device) * p
+ rows = rows + torch.clamp(accepted - 1, min=0, max=p - 1)
+ slot = self.core._draft_bound.share_slot
+ md.qsa_shared_topk_indices[slot, :batch].copy_(torch.index_select(topk_indices, 0, rows))
+ md.qsa_shared_topk_visible_blocks[slot, :batch].copy_(
+ torch.index_select(visible_blocks, 0, rows)
+ )
+ # Host-side readiness, recorded in program order: under capture this runs while the
+ # graph is recorded, so a graph holds the reuse only after the capture it reads.
+ md.qsa_shared_topk_captured_slots.add(slot)
+
+ def mtp_forward(
+ self,
+ embed_tokens: torch.Tensor,
+ all_rank_num_tokens,
+ input_ids: torch.Tensor,
+ position_ids: torch.Tensor,
+ hidden_states: torch.Tensor,
+ attn_metadata: AttentionMetadata,
+ **kwargs,
+ ) -> torch.Tensor:
+ """One draft step over generation rows only: every draft step past the first in an
+ engine forward (one row a request), and the first (`p` rows a request) when the
+ batch had no context. Inside the captured decode graph when the trunk step is.
+
+ Under IndexShare the loop's draft-extend step (`p` rows) runs the draft indexer and
+ captures each request's selection; every draft decode step after it skips the index
+ branch and rebuilds its selection from that capture with `qsa_mtp_reuse_topk`, which
+ appends the groups completed since (the expansion adds the incomplete tail, so the
+ tokens drafted since stay visible). What a skipped step would have written to the
+ index plane sits at or past the newly committed length, and the next engine step's
+ draft-extend rewrites it before anything reads it. `all_rank_num_tokens` is the
+ group's row counts of this draft step, which pad the MoE round trip under attention DP
+ (`_draft_pad_rows`)."""
+ core = self.core
+ assert isinstance(attn_metadata, TrtllmAttentionMetadata)
+ assert core._bound is not None, "the trunk's cache-manager bind has not run"
+ md = attn_metadata
+ rows = hidden_states.shape[0]
+ assert input_ids.shape[0] == rows and hidden_states.shape[1] == core.stream, (
+ f"the draft step got {input_ids.shape[0]} token ids and {tuple(hidden_states.shape)} "
+ f"hidden states; the MTP module reads one {core.stream}-wide stream row per token"
+ )
+ # Projected on every call: the worker rewrites `attn_metadata` in place between
+ # draft steps.
+ step = self.step_args(md)
+ gen_seqs = md.num_seqs
+ p, _ = _verify_rows(core, md, None, gen_seqs, rows)
+ step["predicted_tokens_per_seq"] = p
+ draft, layer_idx = core._draft_step_bind(md)
+ assert rows == md.num_tokens, (
+ f"the draft step carries {rows} rows for a batch its metadata describes as "
+ f"{md.num_tokens} tokens; the QSA row mapping covers the second"
+ )
+ pad_rows = _draft_pad_rows(core, all_rank_num_tokens, rows)
+ # The trunk's dense/sparse decision from the same host buffer: the worker advances
+ # only the device KV lengths through the draft loop, so every draft step of one engine
+ # forward takes the branch its trunk took, in the graph family the runner chose.
+ sparse = gen_seqs > 0 and int(md.kv_lens_runtime[:gen_seqs].max()) > core.dense_threshold
+ share = core.mtp_index_share and md.qsa_in_mtp_draft_loop
+ reuse = share and md.qsa_indexer_skip_topk
+ assert not share or reuse == (p == 1), (
+ f"the worker's reuse flag ({md.qsa_indexer_skip_topk}) disagrees with the step's "
+ f"shape ({p} rows per request): only the draft decode steps after the draft-extend "
+ "step reuse its selection"
+ )
+ # The first draft step reads the trunk's positions buffer, whose int32 copy the trunk
+ # made (taken once); a contiguous int32 `position_ids` (every later draft step) is read
+ # in place. The step's ops only read positions.
+ trunk, self._trunk_pos3 = self._trunk_pos3, None
+ if trunk is not None and trunk[0] == _position_key(position_ids):
+ pos3 = trunk[1]
+ pos, coords = pos3[0], torch.transpose(pos3, 0, 1)
+ elif position_ids.dtype == torch.int32 and position_ids.is_contiguous():
+ pos3 = torch.reshape(position_ids, [_QSA_POSITION_AXES, -1])
+ pos, coords = pos3[0], torch.transpose(pos3, 0, 1)
+ else:
+ pos3, pos, coords = _draft_positions(position_ids)
+ M = core._mtp
+ req_idx = md.qsa_req_idx_per_token[:rows]
+ logical = md.qsa_logical_positions[:rows]
+ block_table = md.qsa_block_table
+ tokens_per_block = md.tokens_per_block
+
+ if core.mapping.world_size == 1:
+ return self._mtp_layer(
+ embed_tokens,
+ input_ids,
+ hidden_states,
+ md,
+ step,
+ draft,
+ layer_idx,
+ rows,
+ gen_seqs,
+ p,
+ sparse,
+ share,
+ reuse,
+ pos3,
+ coords,
+ )
+ x, mixed, inject = _mtp_front(core, M, embed_tokens, input_ids, hidden_states)
+ qkv, gate, k_dst, v_dst = _mtp_qkv(core, M, mixed, pos3)
+ q_index = None
+ if not reuse:
+ q_index = _mtp_index_branch(
+ core,
+ M,
+ mixed,
+ draft,
+ (req_idx, logical, block_table, tokens_per_block, pos, coords),
+ )
+ if sparse:
+ kv = draft.kv
+ _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block)
+ topk_indices = md.qsa_topk_indices[:rows]
+ visible_blocks = md.qsa_visible_blocks[:rows]
+ seq_lens_q = md.qsa_sequence_lengths[:rows]
+ selected = torch.empty(
+ [rows, core.token_topk + core.compress_ratio - 1],
+ dtype=torch.int32,
+ device=qkv.device,
+ )
+ if reuse:
+ slot = draft.share_slot
+ assert slot in md.qsa_shared_topk_captured_slots, (
+ "a draft decode step would reuse a selection that no draft-extend step of "
+ "this loop captured"
+ )
+ qsa_mtp_reuse_topk(
+ shared_indices=md.qsa_shared_topk_indices[slot],
+ shared_visible_blocks=md.qsa_shared_topk_visible_blocks[slot],
+ request_indices=req_idx,
+ visible_blocks=visible_blocks,
+ output=topk_indices,
+ )
+ selected.copy_(
+ qsa_expand_block_indices(
+ topk_indices,
+ logical,
+ seq_lens_q,
+ compress_ratio=core.compress_ratio,
+ token_topk=core.token_topk,
+ )
+ )
+ else:
+ scores = qsa_paged_index_scores(
+ q=q_index,
+ index_cache=draft.index_cache,
+ block_table=block_table,
+ query_positions=logical,
+ request_indices=req_idx,
+ tokens_per_block=tokens_per_block,
+ compress_ratio=core.compress_ratio,
+ only_visible_blocks=True,
+ context_rows=False,
+ )
+ _selection(
+ core,
+ scores,
+ visible_blocks,
+ logical,
+ seq_lens_q,
+ topk_indices,
+ md.qsa_topk_row_starts[:rows],
+ selected,
+ )
+ if share:
+ self._capture_selection(md, topk_indices, visible_blocks, gen_seqs, p)
+ attn = torch.reshape(
+ _sparse_gqa(
+ core,
+ kv,
+ _packed_q(core, qkv),
+ block_table,
+ selected,
+ req_idx,
+ tokens_per_block,
+ None,
+ ),
+ [rows, core.q_width],
+ )
+ else:
+ attn = thop_attention(
+ q=qkv,
+ output=torch.empty([rows, core.q_width], dtype=qkv.dtype, device=qkv.device),
+ local_layer_idx=layer_idx,
+ num_heads=core.heads_q,
+ num_kv_heads=core.heads_kv,
+ head_size=core.head_dim,
+ quant_mode=core.quant_mode,
+ kv_scale_orig_quant=core.kv_scale_orig_quant,
+ kv_scale_quant_orig=core.kv_scale_quant_orig,
+ **step,
+ **_CALL_CONSTANTS,
+ )
+ mtp_fused_sigmoid_mul_inplace(attn, gate)
+ return _mtp_tail(core, M, x, attn, inject, pad_rows)
+
+
+# ------------------------------------------------------------ the MTP module ---
+
+
+class MTPLayer:
+ """The checkpoint's MTP module, replayed once per draft step by the spec worker: one more
+ decoder layer (two hyper-connections, gated GQA over its own QSA indexer, a 512-expert
+ FP8 MoE) behind a two-branch front end (`fc_embedding` over the next token's embedding,
+ `fc_hidden` over the trunk's residual per stream). It returns the residual un-collapsed;
+ `shared_head` owns the module's mixer. The step runs on the targets (`mtp_forward`),
+ routed by the phase each call's metadata shows: the worker rewrites `attn_metadata`
+ between draft steps, so one engine step's draft calls cross from prefill to decode.
+
+ Not an `nn.Module`: the weights live in the core's `ParameterDict`, and a second
+ registration would break the weight load's both-directions coverage."""
+
+ def __init__(self, core: Qwen4ExpDefaultModelingV2, logits_processor) -> None:
+ self.core = core
+ self.logits_processor = logits_processor
+
+ def __call__(self, *args, **kwargs) -> torch.Tensor:
+ return self.forward(*args, **kwargs)
+
+ def shared_head(self, hidden_states, lm_head, attn_metadata, return_context_logits):
+ """The module's own head: its stack-final hyper-connection mixer (a gated mixer with
+ its own parameters, not a plain RMS norm) over the grouped norm, then the trunk's
+ `lm_head` through the shell's logits processor (vocab-parallel and gathered under
+ attention TP). The norm and the mixer are the trunk's stack-final site's."""
+ core = self.core
+ M = core._mtp
+ _, normed = qwen4_exp_hc_inject_norm(
+ hidden_states, None, None, M.fhc_norm, core.eps, core.hc_count
+ )
+ final, _ = _hc_mixer(core, normed, M.fhc_first, M.fhc_up, False)
+ draft_head = getattr(core, "draft_head", None)
+ if draft_head is None:
+ return self.logits_processor.forward(
+ final, lm_head, attn_metadata, return_context_logits
+ )
+ # At tp1 the FP8 draft head, its logits bf16 when only the draft sampler's argmax reads
+ # them.
+ return self.logits_processor.forward(
+ final,
+ draft_head,
+ attn_metadata,
+ return_context_logits,
+ widen=not self.logits_processor.draft_argmax_only,
+ )
+
+ def forward(
+ self,
+ embed_tokens: torch.Tensor,
+ all_rank_num_tokens,
+ input_ids: torch.Tensor,
+ position_ids: torch.Tensor,
+ hidden_states: torch.Tensor,
+ attn_metadata: AttentionMetadata,
+ **kwargs,
+ ) -> torch.Tensor:
+ return self.core._select_target(attn_metadata).mtp_forward(
+ embed_tokens,
+ all_rank_num_tokens,
+ input_ids,
+ position_ids,
+ hidden_states,
+ attn_metadata,
+ **kwargs,
+ )
+
+
+class DraftModel:
+ """The container the spec worker reaches the drafter through: `mtp_layers` (`[0]`,
+ MTP-Eagle replays one layer), `embed_tokens` and `lm_head`. `embed_tokens` resolves
+ through the core on every read: this container is built while parameters are meta
+ tensors, and materialization replaces the tensor objects."""
+
+ def __init__(self, core: Qwen4ExpDefaultModelingV2, lm_head, logits_processor) -> None:
+ self.core = core
+ self.mtp_layers = [MTPLayer(core, logits_processor)]
+ self.lm_head = lm_head
+
+ @property
+ def embed_tokens(self) -> torch.Tensor:
+ return self.core.w["embed"]
+
+
+class SpeculativeRollback:
+ """The target's auxiliary speculative-state handler: it undoes what a verify step's
+ rejected draft rows wrote into the side-cache planes the runtime's own state promotion
+ does not cover: the PLE n-gram id history and dilated-conv state, each attention layer's
+ QSA index-K plane and the shared QSA position plane. The Gated-DeltaNet state is not
+ here: the runtime promotes it.
+
+ For a generation request that accepted `a` rows (golden token included, `1 <= a <= p`)
+ the commit rewrites the n-gram pool row to the id window of row `a - 1` minus its oldest
+ id, the conv pool row to columns `[a, a + state_len)` of `[pre-step state | conv input of
+ the request's p rows]`, and every index-K and position slot written by a row of ordinal
+ `>= a` to its pre-verify contents.
+
+ The targets stage during the forward (`begin`, then `stage_ple_verify` before the fused PLE
+ launches and `stage_ple_conv_input` after them, and `stage_index` before each attention
+ layer's store) and the spec worker calls `commit_speculative_states` after its draft loop,
+ inside the same model forward. Everything staged is allocated by the forward and consumed
+ by that forward's commit, so a captured graph holds both ends. A step opened with `fused`
+ (the decode target's captured steps) stages the PLE pools with one `spec_rollback_stage`
+ launch and commits with one `spec_rollback_commit` launch."""
+
+ def __init__(self) -> None:
+ self._pending: dict | None = None
+ # (base addresses, page strides, plane count) of every attention layer's index-K
+ # plane, once the cache manager is bound.
+ self._planes: tuple | None = None
+
+ def set_index_planes(self, planes: list) -> None:
+ """Bind the attention layers' index-K planes (static for the manager's lifetime)."""
+ ptrs, strides = _plane_table(planes)
+ self._planes = (ptrs, strides, len(planes))
+
+ # ------------------------------------------------------------- staging ---
+
+ def discard(self) -> None:
+ """Drop whatever an earlier forward staged."""
+ self._pending = None
+
+ @property
+ def index_planes(self) -> tuple | None:
+ """`(base addresses, page strides, plane count)` of the bound index-K planes."""
+ return self._planes
+
+ def begin_staged(
+ self,
+ *,
+ num_gens: int,
+ rows_per_gen: int,
+ slots: torch.Tensor,
+ pages: torch.Tensor,
+ within: torch.Tensor,
+ index_all: torch.Tensor,
+ ) -> None:
+ """Open a generation-only verify step's staging from what `verify_step_setup`
+ resolved: the requests' int64 recurrent slots, the rows' int64 QSA pages and offsets,
+ and every bound index-K plane at them. The PLE staging and the commit are the
+ one-launch ones (`begin(..., fused=True)`)."""
+ n, p = num_gens, rows_per_gen
+ assert n > 0 and p > 1, (
+ f"a verify step stages generation requests carrying drafts; got {n} requests of "
+ f"{p} rows"
+ )
+ assert pages.shape[0] == n * p and index_all.shape[1] == n * p, (
+ f"{n} generation requests of {p} rows are {n * p} rows; got {pages.shape[0]} slots"
+ )
+ self._pending = dict(
+ num_contexts=0,
+ prefix=0,
+ n=n,
+ p=p,
+ slots=slots,
+ pages=pages,
+ within=within,
+ staged_layers=0,
+ position=None,
+ index_all=index_all,
+ fused=True,
+ ngram_stream=None,
+ )
+
+ def begin(
+ self,
+ *,
+ num_contexts: int,
+ prefix_tokens: int,
+ num_gens: int,
+ rows_per_gen: int,
+ slots: torch.Tensor,
+ request_indices: torch.Tensor,
+ logical_positions: torch.Tensor,
+ block_table: torch.Tensor,
+ tokens_per_block: int,
+ fused: bool = False,
+ ) -> None:
+ """Open a verify step's staging. The generation rows follow the `prefix_tokens`
+ context rows as `num_gens` runs of `rows_per_gen`, so row `s * p + t` is ordinal `t`
+ of generation request `s`. `slots` are those requests' recurrent slots;
+ `request_indices` and `logical_positions` the generation rows' own. Each row's QSA
+ slot is resolved here, as `qsa_paged_index_store` resolves it, and every index plane's
+ slots are gathered before the trunk writes any of them."""
+ n, p = num_gens, rows_per_gen
+ assert n > 0 and p > 1, (
+ f"a verify step stages generation requests carrying drafts; got {n} requests of "
+ f"{p} rows"
+ )
+ assert logical_positions.shape[0] == n * p, (
+ f"{n} generation requests of {p} rows are {n * p} rows; got "
+ f"{logical_positions.shape[0]} positions"
+ )
+ assert self._planes is not None, "the attention layers' index planes are not bound"
+ positions = logical_positions.to(torch.long)
+ pages = block_table[request_indices.to(torch.long), positions // tokens_per_block]
+ pages = pages.to(torch.long)
+ within = positions % tokens_per_block
+ ptrs, strides, count = self._planes
+ self._pending = dict(
+ num_contexts=num_contexts,
+ prefix=prefix_tokens,
+ n=n,
+ p=p,
+ slots=slots.to(torch.long),
+ pages=pages,
+ within=within,
+ staged_layers=0,
+ position=None,
+ index_all=qsa_rollback_gather(ptrs, strides, pages, within, count),
+ fused=fused,
+ ngram_stream=None,
+ )
+
+ def stage_ple_verify(
+ self, ngram_pool: torch.Tensor, conv_pool: torch.Tensor, input_ids: torch.Tensor
+ ) -> None:
+ """A verify step's PLE staging, before either pool is written: both pools, the
+ generation slots' pre-step conv rows, and the n-gram candidates rebuilt from the slots'
+ pre-step id history and the rows' ids. Row `t`'s history after the step is the last
+ `context` ids of `[history | ids of rows 0..t]`."""
+ pending = self._pending
+ prefix, n, p = pending["prefix"], pending["n"], pending["p"]
+ slots = pending["slots"]
+ pending["ngram_pool"] = ngram_pool
+ pending["conv_pool"] = conv_pool
+ if pending["fused"]:
+ pending["ngram_stream"], pending["conv_prev"] = spec_rollback_stage(
+ ngram_pool, conv_pool, slots, input_ids[prefix : prefix + n * p], p
+ )
+ return
+ pending["conv_prev"] = conv_pool[slots]
+ context = ngram_pool.shape[1]
+ ids = torch.reshape(input_ids[prefix : prefix + n * p], [n, p]).to(ngram_pool.dtype)
+ stream = torch.cat([ngram_pool[slots], ids], 1)
+ # Window k of the unfold is the history after k rows.
+ pending["ngram_candidates"] = stream.unfold(1, context, 1)[:, 1:]
+
+ def stage_ple_conv_input(self, conv_input: torch.Tensor) -> None:
+ """The generation rows' conv input (`[T, C]`, token-major), as the fused PLE launch
+ wrote it into a buffer the forward allocated for this staging alone. A mixed batch's
+ rows are copied out, so its context prefix is not kept alive through the draft loop."""
+ pending = self._pending
+ prefix, n, p = pending["prefix"], pending["n"], pending["p"]
+ conv_x = torch.reshape(conv_input[prefix:], [n, p, -1])
+ pending["conv_x"] = conv_x.clone() if prefix else conv_x
+
+ def stage_index(self, index_cache: torch.Tensor, position_cache: torch.Tensor) -> None:
+ """Before an attention layer's index store: on the step's first attention layer (the
+ first to write the shared plane) the generation rows' position slots. The index-K
+ slots of every layer were gathered when the step was opened."""
+ pending = self._pending
+ if pending["position"] is None:
+ pages, within = pending["pages"], pending["within"]
+ pending["position"] = (position_cache, position_cache[pages, within])
+ pending["staged_layers"] += 1
+
+ # -------------------------------------------------------------- commit ---
+
+ def commit_speculative_states(
+ self,
+ num_accepted_tokens: torch.Tensor,
+ state_indices: torch.Tensor,
+ num_contexts: int,
+ ) -> None:
+ """Rewrite the staged planes to the accepted prefix. Called by the spec worker."""
+ pending, self._pending = self._pending, None
+ if pending is None:
+ return
+ n, p = pending["n"], pending["p"]
+ assert num_contexts == pending["num_contexts"], (
+ f"the worker commits with {num_contexts} context requests and the forward staged "
+ f"{pending['num_contexts']}"
+ )
+ assert state_indices.shape[0] == num_contexts + n, (
+ f"the worker hands {state_indices.shape[0]} slots for a batch the forward staged "
+ f"as {num_contexts} + {n} requests"
+ )
+ assert pending.get("conv_x") is not None and pending["staged_layers"] > 0, (
+ "a verify step staged neither its PLE rows nor any attention layer; the forward "
+ "did not run the staging it opened"
+ )
+ ptrs, strides, _ = self._planes
+ position, position_old = pending["position"]
+ pages, within = pending["pages"], pending["within"]
+ if pending["fused"]:
+ spec_rollback_commit(
+ num_accepted_tokens[num_contexts : num_contexts + n],
+ state_indices[num_contexts:],
+ pending["ngram_pool"],
+ pending["ngram_stream"],
+ pending["conv_pool"],
+ pending["conv_prev"],
+ pending["conv_x"],
+ position,
+ position_old,
+ pages,
+ within,
+ ptrs,
+ strides,
+ pending["index_all"],
+ )
+ return
+ accepted = num_accepted_tokens[num_contexts : num_contexts + n].to(torch.long)
+ slots = state_indices[num_contexts:].to(torch.long)
+ rows = torch.arange(n, device=accepted.device)
+ ngram_pool, conv_pool = pending["ngram_pool"], pending["conv_pool"]
+ state_len = conv_pool.shape[-1]
+ ngram_pool[slots] = pending["ngram_candidates"][rows, accepted - 1]
+ # [n, C, state_len + p]; window k of the unfold is the state after k tokens.
+ history = torch.cat([pending["conv_prev"], torch.transpose(pending["conv_x"], 1, 2)], 2)
+ conv_pool[slots] = history.unfold(2, state_len, 1)[rows, :, accepted]
+ # [n * p]: True at every ordinal the sampler rejected.
+ rejected = torch.reshape(
+ torch.arange(p, device=accepted.device).unsqueeze(0) >= accepted.unsqueeze(1), [-1]
+ )
+ qsa_rollback_restore(ptrs, strides, pages, within, rejected, pending["index_all"])
+ position[pages, within] = torch.where(
+ rejected[:, None], position_old, position[pages, within]
+ )
+
+
+# -------------------------------------------------------------------- shell ---
+
+
+class _WidenedHead:
+ """`lm_head` whose bf16 logits come back as fp32 from `bf16_to_fp32` (one vectorized
+ kernel; torch's `.float()` is an unvectorized copy at under half the bandwidth), so the
+ logits processor's own `.float()` returns them as they are. The same values bit for
+ bit."""
+
+ def __init__(self, lm_head) -> None:
+ self.lm_head = lm_head
+
+ def __call__(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ logits = self.lm_head(hidden_states)
+ if logits.dtype == torch.bfloat16 and logits.is_contiguous():
+ return bf16_to_fp32(logits)
+ return logits
+
+
+def _fp8_rowwise(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ """`weight` [rows, cols] bf16 -> (e4m3 [rows, cols], fp32 [1, rows]): one scale a row,
+ amax / 448, so each row's largest magnitude maps to e4m3's largest finite value."""
+ scale = torch.empty(weight.shape[0], dtype=torch.float32, device=weight.device)
+ w8 = torch.empty_like(weight, dtype=torch.float8_e4m3fn)
+ for lo in range(0, weight.shape[0], 16384):
+ hi = min(lo + 16384, weight.shape[0])
+ chunk = weight[lo:hi].float()
+ scale[lo:hi] = (chunk.abs().amax(dim=1) / 448.0).clamp(min=1e-12)
+ w8[lo:hi] = (chunk / scale[lo:hi, None]).to(torch.float8_e4m3fn)
+ return w8, scale.reshape(1, -1).contiguous()
+
+
+def _head_mid(weight_fp8: torch.Tensor, weight_scale: torch.Tensor) -> MidWeight:
+ """`mid_rows_mm`'s view of a per-row-scaled e4m3 head copy (`_fp8_rowwise`)."""
+ return MidWeight(weight_fp8, weight_scale.reshape(-1))
+
+
+class _Fp8Head:
+ """A vocabulary projection over an FP8 copy of `lm_head`: up to `_FP8_GEMV_HEAD_ROWS` rows on
+ `fp8_gemv` over the packed block-scaled copy -- or, given the NVFP4 copy `nv`, `nvfp4_gemv`
+ over it --, up to `_MID_HEAD_ROWS` on `mid_rows_mm` over the per-row-scaled copy, more rows as
+ per-token e4m3 activations times that copy, scaled in cuBLASLt's epilogue; bf16 out
+ (`_WidenedHead` widens it as it does the bf16 head's). The draft head always (at tp1), the
+ verify head under the dense FP8 mode."""
+
+ def __init__(self, weight_fp8: torch.Tensor, weight_scale: torch.Tensor, packed, nv) -> None:
+ self.weight_fp8, self.weight_scale, self.packed, self.nv = (
+ weight_fp8,
+ weight_scale,
+ packed,
+ nv,
+ )
+ self.mid = _head_mid(weight_fp8, weight_scale)
+
+ def __call__(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ rows = hidden_states.shape[0]
+ if rows <= _FP8_GEMV_HEAD_ROWS and _gemv_ready(hidden_states):
+ if self.nv is not None:
+ return nvfp4_gemv_op(hidden_states, self.nv)
+ if self.packed is not None:
+ return fp8_gemv_op(hidden_states, self.packed)
+ if rows <= _MID_HEAD_ROWS and _gemv_ready(hidden_states):
+ return mid_rows_mm_op(hidden_states, self.mid)
+ x8, xs = torch.ops.tensorrt_llm.quantize_e4m3_activation(hidden_states)
+ return torch._scaled_mm(
+ x8,
+ self.weight_fp8.t(),
+ scale_a=xs.float().reshape(-1, 1),
+ scale_b=self.weight_scale,
+ out_dtype=hidden_states.dtype,
+ )
+
+
+class _WideningLogitsProcessor(LogitsProcessor):
+ """The shell's logits processor (its row gather, its fp32 logits) over `_WidenedHead`, at
+ tp1; with `verify_head` set (the dense FP8 mode) the target's own head is that `_Fp8Head`.
+
+ `widen=False` (with `return_context_logits`, the speculative paths) returns the head's own
+ bf16 logits: for a consumer that only takes their row argmax, which is the same token over
+ bf16 as over their exact fp32 widening (same order, ties, NaN and +-0)."""
+
+ verify_head = None
+ #: Set by the shell's forward each engine step: the step's draft logits feed only the
+ #: greedy draft sampler's argmax (`_spec_argmax_only`).
+ draft_argmax_only = False
+
+ def forward(
+ self,
+ hidden_states,
+ lm_head,
+ attn_metadata,
+ return_context_logits: bool = False,
+ widen: bool = True,
+ ) -> torch.Tensor:
+ if self.verify_head is not None and not isinstance(lm_head, _Fp8Head):
+ lm_head = self.verify_head
+ if not widen and return_context_logits:
+ return lm_head(hidden_states)
+ return super().forward(
+ hidden_states, _WidenedHead(lm_head), attn_metadata, return_context_logits
+ )
+
+
+def _spec_argmax_only(spec_worker, spec_metadata) -> tuple[bool, bool]:
+ """`(verify, draft)`: whether this step's verify / draft logits are consumed only by a row
+ argmax, so they can stay bf16.
+
+ Both need no guided decoder (it edits the logits). The draft logits then feed only the draft
+ sampler, an argmax unless rejection sampling runs on a non-greedy batch. The verify logits
+ are also the step's `logits` output and the occurrence penalties' and relaxed acceptance's
+ input; the one-model speculative sampler admits no request asking for logits or log
+ probabilities, so they stay bf16 only for an all-greedy batch with the penalties off (a
+ deploy-time flag) and strict acceptance. Every input is a deploy-time constant or
+ `is_all_greedy_sample`, part of the CUDA graph key."""
+ if getattr(spec_worker, "guided_decoder", None) is not None:
+ return False, False
+ draft = not spec_metadata.wants_advanced_draft_sampling
+ verify = (
+ spec_metadata.is_all_greedy_sample
+ and not getattr(spec_metadata, "enable_penalty", False)
+ and not getattr(
+ getattr(spec_worker, "spec_config", None), "use_relaxed_acceptance_for_thinking", False
+ )
+ )
+ return verify, draft
+
+
+class ModelingV2Qwen38FlashNextNvfp4Sm103Default(
+ DecoderModelForCausalLM[Qwen4ExpDefaultModelingV2, PretrainedConfig]
+):
+ """Qwen3.8-Flash-Next-NVFP4 at tp1, tp4, tep4, dep4 or dep2 on sm_103."""
+
+ # At tp1 one-engine speculative decoding takes the engine's fast paths (greedy and
+ # temperature-only sampling CUDA graph variants, the overlap loop on the execution stream,
+ # hoisted position ids, the MTP-Eagle draft loop's step-0 epilogue, the recorded steady
+ # decode-step prepare), validated there; off at every other layout.
+ spec_decode_fast_paths: bool = False
+ runtime_hooks: RuntimeHooks | None = None
+
+ @classmethod
+ def get_preferred_kv_cache_manager_version(cls, pretrained_config=None) -> str:
+ """Prefer KV cache manager V2: PLE state (`ple_layer_cache`) and QSA's index and
+ position planes exist only there."""
+ return "V2"
+
+ def __init__(self, model_config: ModelConfig):
+ cfg = model_config.pretrained_config
+ assert cfg is not None
+ # Adopt the engine-normalized text tower (`Qwen4ExpTextConfig`) as the model's config:
+ # everything that sizes the pools reads `model_engine.model.model_config.
+ # pretrained_config`. The swap is in place because the engine keeps its own reference
+ # to this `ModelConfig`.
+ text = getattr(cfg, "text_config", None)
+ assert text is not None, (
+ "this target is the language tower of the composite Qwen4-Exp export; "
+ "pretrained_config carries no text_config"
+ )
+ assert getattr(text, "architectures", None) == ["Qwen4ExpForCausalLM"], (
+ f"the text tower must be the normalized Qwen4-Exp text config; got architectures "
+ f"{getattr(text, 'architectures', None)!r}"
+ )
+ composite = cfg
+ model_config.pretrained_config = text
+ super().__init__(
+ Qwen4ExpDefaultModelingV2(model_config, composite),
+ config=model_config,
+ hidden_size=text.hidden_size,
+ vocab_size=text.vocab_size,
+ )
+ # `lm_head` is outside the quantization scheme: a dense bf16 `LMHead`.
+ assert self.lm_head.weight.dtype == torch.bfloat16, (
+ f"expected a bf16 lm_head; the shell built {self.lm_head.weight.dtype}"
+ )
+ self.one_rank = model_config.mapping.world_size == 1
+ if self.one_rank:
+ self.logits_processor = _WideningLogitsProcessor()
+ # No separate draft checkpoint: the MTP module's weights are the target's own.
+ self.draft_config = None
+ # The speculative branch, built only when `speculative_config` asked for it.
+ self.spec_config = getattr(model_config, "spec_config", None)
+ self.draft_model = None
+ self.spec_worker = None
+ if self.spec_config is not None:
+ self.draft_model = DraftModel(self.model, self.lm_head, self.logits_processor)
+ # Must match what the resource manager decided: left at `get_spec_worker`'s
+ # default, the draft step would be handed a manager that maps no draft layer.
+ self.use_separate_draft_kv_cache = should_use_separate_draft_kv_cache(self.spec_config)
+ if self.one_rank:
+ self.spec_decode_fast_paths = True
+ self.runtime_hooks = RuntimeHooks(loop_on_execution_stream=True)
+ self.spec_worker = get_spec_worker(
+ self.spec_config,
+ model_config,
+ model_config.mapping,
+ use_separate_draft_kv_cache=self.use_separate_draft_kv_cache,
+ spec_decode_fast_paths=self.spec_decode_fast_paths,
+ )
+ assert self.spec_worker is not None, (
+ f"the engine resolved no spec worker for spec_dec_mode "
+ f"{self.spec_config.spec_dec_mode!r}"
+ )
+ self.spec_worker.set_draft_model(self.draft_model)
+ # The PLE pools and the QSA side caches are written per verify row and nothing the
+ # runtime owns rolls them back; the worker calls this after its draft loop, inside
+ # the same forward.
+ self.spec_worker.register_auxiliary_state_handler(self.model.spec_rollback)
+ # Media placeholder ids for the executor; built lazily (`__init__` runs under meta
+ # init).
+ self._mm_token_ids_value = (self.model.image_token_id, self.model.video_token_id)
+ self._mm_token_ids_cache: torch.Tensor | None = None
+ # Replicated on every rank; built once the weights are real (`post_load_weights`).
+ self.vision_tower: _vision.VisionTower | None = None
+
+ @property
+ def multimodal_token_ids(self) -> torch.Tensor:
+ """Placeholder ids whose rows vision embeddings replace. This checkpoint's
+ placeholders are in vocabulary, so the executor's out-of-vocabulary fallback would
+ select nothing."""
+ if self._mm_token_ids_cache is None:
+ self._mm_token_ids_cache = torch.tensor(self._mm_token_ids_value, dtype=torch.int32)
+ return self._mm_token_ids_cache
+
+ @property
+ def multimodal_data_device_paths(self) -> list[str]:
+ """Tensors under `multimodal_data` the executor moves to the device. Absent, it moves
+ everything, putting the `(t, h, w)` grids the tower reads on the host there too."""
+ return list(_vision.DEVICE_PATHS)
+
+ def validate_request(self, request) -> None:
+ """Admission check, run by the executor on every new request: a media request the
+ vision tower would assert on fails here as its own error (`vision.validate_request`);
+ a text request returns at once."""
+ _vision.validate_request(self.vision_tower, request)
+
+ def forward(
+ self,
+ attn_metadata: AttentionMetadata,
+ input_ids: torch.IntTensor | None = None,
+ position_ids: torch.IntTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ return_context_logits: bool = False,
+ spec_metadata=None,
+ lora_params: dict | None = None,
+ **kwargs,
+ ):
+ """One engine step: the core's `begin_step`, then the base forward with every
+ argument it was handed (`resource_manager` stays in `**kwargs`).
+
+ With a spec worker the shell owns the logits gather (one-model modes gather the rows
+ the sampler verifies, `spec_metadata.gather_ids`) and hands the worker the trunk's
+ ungathered stream rows, which the MTP module reads; `position_ids` goes in the
+ engine's layout."""
+ self.model.begin_step(
+ attn_metadata,
+ input_ids,
+ position_ids,
+ inputs_embeds=inputs_embeds,
+ lora_params=lora_params,
+ spec_metadata=spec_metadata,
+ vision=_vision.multimodal_embeds(self.vision_tower, kwargs),
+ )
+ if self.spec_worker is None:
+ assert spec_metadata is None, (
+ "spec_metadata arrived without a spec worker; this engine was built without a "
+ "speculative_config"
+ )
+ return super().forward(
+ attn_metadata,
+ input_ids,
+ position_ids,
+ inputs_embeds,
+ return_context_logits,
+ spec_metadata,
+ lora_params,
+ **kwargs,
+ )
+ assert spec_metadata is not None, "the spec worker is built but no spec_metadata came"
+ final, hc = self.model.split_spec_output(
+ self.model(
+ attn_metadata=attn_metadata,
+ input_ids=input_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ lora_params=lora_params,
+ spec_metadata=spec_metadata,
+ **kwargs,
+ )
+ )
+ if isinstance(self.logits_processor, _WideningLogitsProcessor):
+ verify_argmax_only, self.logits_processor.draft_argmax_only = _spec_argmax_only(
+ self.spec_worker, spec_metadata
+ )
+ logits = self.logits_processor.forward(
+ nn.functional.embedding(spec_metadata.gather_ids, final),
+ self.lm_head,
+ attn_metadata,
+ True,
+ widen=not verify_argmax_only,
+ )
+ else:
+ logits = self.logits_processor.forward(
+ nn.functional.embedding(spec_metadata.gather_ids, final),
+ self.lm_head,
+ attn_metadata,
+ True,
+ )
+ # A prefill CUDA graph pads the step's rows; the worker reshapes generation rows per
+ # request, so it takes the batch's own.
+ if attn_metadata.padded_num_tokens is not None:
+ n = attn_metadata.num_tokens
+ input_ids, hc = input_ids[:n], hc[:n]
+ position_ids = position_ids[..., :n]
+ return self.spec_worker(
+ input_ids=input_ids,
+ position_ids=position_ids,
+ hidden_states=hc,
+ logits=logits,
+ attn_metadata=attn_metadata,
+ spec_metadata=spec_metadata,
+ draft_model=self.draft_model,
+ resource_manager=kwargs.get("resource_manager"),
+ )
+
+ def load_weights(self, weights, *args, **kwargs):
+ _weights.MODEL_WEIGHTS.load(self, weights)
+
+ def post_load_weights(self):
+ super().post_load_weights()
+ self.model.post_load_weights()
+ self.vision_tower = _vision.VisionTower(self.model.w, self.model.vision_config)
+ if self.one_rank:
+ self._build_heads(self.lm_head.weight.data)
+
+ def _build_heads(self, weight: torch.Tensor) -> None:
+ """The tp1 heads' copies of `weight`: the draft head's FP8 copies (per row, and packed
+ for `fp8_gemv` unless its NVFP4 copy replaces it, `_MTP_HEAD_NVFP4_ENV`), and under the
+ dense FP8 mode the verify head over the same per-row copy with its packed copy or its
+ NVFP4 one (`_TARGET_HEAD_NVFP4_ENV`); every decode variant compiled now."""
+ dense_fp8 = _switch(_DENSE_FP8_ENV)
+ draft = self.model.mtp_enabled
+ if not (draft or dense_fp8):
+ return
+ rowwise = _fp8_rowwise(weight)
+ draft_nv = draft and _switch(_MTP_HEAD_NVFP4_ENV)
+ target_nv = dense_fp8 and _switch(_TARGET_HEAD_NVFP4_ENV)
+ packed = (
+ pack_from_bf16(weight)
+ if (draft and not draft_nv) or (dense_fp8 and not target_nv)
+ else None
+ )
+ nv = None
+ if draft_nv or target_nv:
+ nv = nvfp4_pack_from_bf16(weight)
+ if weight.is_cuda:
+ nvfp4_gemv_precompile(nv)
+ if weight.is_cuda:
+ mid_rows_precompile(_head_mid(*rowwise), _FP8_GEMV_HEAD_ROWS + 1, _MID_HEAD_ROWS)
+ if draft:
+ self.model.draft_head = _Fp8Head(*rowwise, packed, nv if draft_nv else None)
+ if dense_fp8:
+ self.logits_processor.verify_head = _Fp8Head(
+ *rowwise, packed, nv if target_nv else None
+ )
+ if target_nv:
+ print(
+ f"target head NVFP4 ({_TARGET_HEAD_NVFP4_ENV}=1): decode rows <= "
+ f"{_FP8_GEMV_HEAD_ROWS} on nvfp4_gemv over lm_head {tuple(nv.data.shape)}",
+ flush=True,
+ )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/kernels/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/__init__.py
similarity index 60%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/kernels/__init__.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/__init__.py
index 3811a6f831f0..777b6b725c00 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/kernels/__init__.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/__init__.py
@@ -1,3 +1,3 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
-"""Kernels used only by this target, called through its ops."""
+"""Ops used only by this target; the family's shared ops are in `models/qwen4_exp/ops`."""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/kernels/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/__init__.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/kernels/moe/__init__.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/__init__.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_fp8_pages_to_bf16__prefill.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_fp8_pages_to_bf16__prefill.py
new file mode 100644
index 000000000000..e993fd24eae7
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_fp8_pages_to_bf16__prefill.py
@@ -0,0 +1,224 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""A batch's pages of an e4m3 paged K/V pool, widened into a bf16 scratch pool, and the block
+table that addresses them there.
+
+Tier 5 (Triton written for this target): the target's context launches over a kv_fp8 pool
+read their requests' pages through this scratch with the bf16 row-tile kernel.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+#: Elements a program widens per loop step.
+_CHUNK = 4096
+
+#: One rank's K/V geometry and the engine's QSA table at max_seq_len 262144
+#: (`ceil(262145 / 32)` columns padded to a multiple of 4): a dep4 rank's two KV heads and
+#: 32-row table; `slot_pages` role pages of one lifecycle slot, layer `layer`'s K and V
+#: planes among them.
+_CHECKPOINT = dict(
+ kv_heads=2,
+ head_dim=256,
+ tokens_per_block=32,
+ width=8196,
+ slots=32,
+ slot_pages=24,
+ layer=11,
+ holes=False,
+ straddle=False,
+)
+
+
+def _spec(kv_lens, **overrides) -> dict:
+ """`kv_lens` is one KV length per listed request, in order."""
+ return {**_CHECKPOINT, "kv_lens": tuple(kv_lens), **overrides}
+
+
+@triton.jit
+def _fp8_pages_to_bf16_kernel(
+ k_cache,
+ v_cache,
+ block_table,
+ slots,
+ counts,
+ offsets,
+ table,
+ k_out,
+ v_out,
+ num_pool_pages,
+ listed,
+ stride_page,
+ bt_stride,
+ table_stride,
+ PAGE_ELEMS: tl.constexpr,
+ CHUNK: tl.constexpr,
+):
+ request = tl.program_id(0)
+ count = tl.load(counts + request)
+ base = tl.load(offsets + request).to(tl.int64)
+ slot = tl.load(slots + request).to(tl.int64)
+ for page in range(tl.program_id(1), count, tl.num_programs(1)):
+ dst = base + page
+ if dst < listed:
+ src = tl.load(block_table + slot * bt_stride + page).to(tl.int64)
+ resident = (src >= 0) & (src < num_pool_pages)
+ tl.store(table + slot * table_stride + page, tl.where(resident, dst, -1).to(tl.int32))
+ safe = tl.where(resident, src, 0)
+ for start in tl.static_range(0, PAGE_ELEMS, CHUNK):
+ elems = start + tl.arange(0, CHUNK)
+ k = tl.load(k_cache + safe * stride_page + elems, mask=resident, other=0.0)
+ v = tl.load(v_cache + safe * stride_page + elems, mask=resident, other=0.0)
+ tl.store(k_out + dst * PAGE_ELEMS + elems, k.to(tl.bfloat16))
+ tl.store(v_out + dst * PAGE_ELEMS + elems, v.to(tl.bfloat16))
+
+
+#: Programs per request; each strides the request's pages.
+_PAGE_PROGRAMS = 1024 # capacity: fixed: a launch width; a request may list any page count
+
+
+class QsaFp8PagesToBf16(OpWrapper):
+ """K and V of listed requests' pages of an e4m3 HND paged pool, as bf16 scratch pools,
+ and the block table that addresses them there.
+
+ Request `r` lists its first `counts[r]` table pages, `block_table[slots[r], j]` for `j <
+ counts[r]`, at scratch pages `offsets[r] + j` (those below `listed`, the scratch's page
+ count)::
+
+ k_out[offsets[r] + j] = bf16(k_cache[block_table[slots[r], j]])
+
+ and `v_out` likewise. Each e4m3 code widens exactly to bf16 (no scale: the pool holds
+ K/V at unit scale, as `qsa_paged_sparse_gqa` reads it), so a bf16 kernel over the scratch
+ reads the values the e4m3 kernels read. `table[slots[r], j]` becomes `offsets[r] + j`,
+ or -1 where the pool's entry is unmapped or out of the pool (whose scratch page is then
+ zeros); the call writes no other entry of `table`, which the caller fills with -1.
+ Returns `(k_out, v_out, table)`.
+
+ Fusion boundary: the listing's widening copy and table; choosing the requests and their
+ page counts (and `offsets`, their exclusive prefix sum) is the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="one request, one page", spec=_spec((17,))),
+ Cell(why="dep4 cold context: one 8192-token request", spec=_spec((8192,))),
+ Cell(
+ why="dep4 deep chunk: one request at 198160 tokens (6193 pages) over a pool "
+ "straddling 2^31 elements of the trunk plane",
+ spec=_spec((198160,), straddle=True),
+ ),
+ Cell(
+ why="three requests, a chunked continuation and two short ones",
+ spec=_spec((40000 + 3192, 3000, 900)),
+ ),
+ Cell(why="holes: unmapped and out-of-pool entries", spec=_spec((4000, 2000), holes=True)),
+ Cell(
+ why="a tp4 / tep4 rank's single KV head, [128, 8196] table",
+ spec=_spec((9000, 2100), kv_heads=1, slots=128),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none -- every input is an argument. Call site: the target's
+ context/mixed forward over a kv_fp8 pool, once per sparse attention layer, for the
+ step's context requests (`_fp8_context_plan`: slots from the context rows' request
+ indices, counts `ceil(kv_len / tpb)` from the device KV lengths, `listed` from the host
+ ones); the row-tile kernel then reads the scratch through `table`.
+
+ One Triton kernel, `(requests, min(width, 1024))` programs, each striding its request's
+ pages; a page (all KV heads, contiguous in the role plane) is copied in 4096-element
+ steps, K and V together. Page ids are widened to int64 before the stride product: the
+ engine's trunk planes reach 6.7e10 elements (dep4 kv_fp8 171362 pages x 24 role pages).
+
+ Preconditions: e4m3 role planes `[pages, Hkv, tpb, D]` with a page's heads, tokens and
+ dims contiguous (`stride(1) == tpb * D`, `stride(2) == D`, unit last stride; any page
+ stride), K's and V's at equal strides; `k_out` / `v_out` contiguous bf16 `[listed, Hkv,
+ tpb, D]`; `table` int32 of `block_table`'s shape. Bitwise deterministic (a copy).
+ """
+
+ def __call__(
+ self,
+ k_cache: torch.Tensor,
+ v_cache: torch.Tensor,
+ block_table: torch.Tensor,
+ slots: torch.Tensor,
+ counts: torch.Tensor,
+ offsets: torch.Tensor,
+ table: torch.Tensor,
+ k_out: torch.Tensor,
+ v_out: torch.Tensor,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ requests, listed = slots.shape[0], k_out.shape[0]
+ _, heads, tpb, dim = k_cache.shape
+ if k_cache.dtype != torch.float8_e4m3fn or v_cache.dtype != k_cache.dtype:
+ raise ValueError("the widening copy reads e4m3 K/V planes")
+ if k_cache.stride() != v_cache.stride() or k_cache.stride()[1:] != (tpb * dim, dim, 1):
+ raise ValueError("K/V planes must share strides with contiguous head pages")
+ if not (k_out.is_contiguous() and v_out.is_contiguous()):
+ raise ValueError("the bf16 scratch pools must be contiguous")
+ if k_out.dtype != torch.bfloat16 or k_out.shape[1:] != k_cache.shape[1:]:
+ raise ValueError("the scratch pools must be bf16 pages of the pool's geometry")
+ page_elems = heads * tpb * dim
+ if requests and listed:
+ _fp8_pages_to_bf16_kernel[(requests, min(block_table.shape[1], _PAGE_PROGRAMS))](
+ k_cache,
+ v_cache,
+ block_table,
+ slots,
+ counts,
+ offsets,
+ table,
+ k_out,
+ v_out,
+ k_cache.shape[0],
+ listed,
+ k_cache.stride(0),
+ block_table.stride(0),
+ table.stride(0),
+ PAGE_ELEMS=page_elems,
+ CHUNK=min(_CHUNK, page_elems),
+ num_warps=8,
+ )
+ return k_out, v_out, table
+
+ def reference(
+ self,
+ k_cache: torch.Tensor,
+ v_cache: torch.Tensor,
+ block_table: torch.Tensor,
+ slots: torch.Tensor,
+ counts: torch.Tensor,
+ offsets: torch.Tensor,
+ table: torch.Tensor,
+ k_out: torch.Tensor,
+ v_out: torch.Tensor,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ k_ref, v_ref, table_ref = k_out.clone(), v_out.clone(), table.clone()
+ listed, pool_pages = k_out.shape[0], k_cache.shape[0]
+ for r in range(slots.shape[0]):
+ slot, base = int(slots[r]), int(offsets[r])
+ count = max(0, min(int(counts[r]), block_table.shape[1], listed - base))
+ if count == 0:
+ continue
+ src = block_table[slot, :count].long()
+ resident = (src >= 0) & (src < pool_pages)
+ dst = torch.arange(base, base + count, device=src.device)
+ table_ref[slot, :count] = torch.where(resident, dst, -1).to(table.dtype)
+ safe = torch.where(resident, src, 0)
+ keep = resident[:, None, None, None]
+ k_ref[base : base + count] = torch.where(keep, k_cache[safe].float(), 0.0).to(
+ k_out.dtype
+ )
+ v_ref[base : base + count] = torch.where(keep, v_cache[safe].float(), 0.0).to(
+ v_out.dtype
+ )
+ return k_ref, v_ref, table_ref
+
+
+qsa_fp8_pages_to_bf16 = QsaFp8PagesToBf16()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__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__default/ops/attention/qsa_rowtile_attention_cute__prefill.py
similarity index 99%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rowtile_attention_cute.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_rowtile_attention_cute__prefill.py
index 1a9cac72e24e..4fa6fef27d45 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__default/ops/attention/qsa_rowtile_attention_cute__prefill.py
@@ -101,6 +101,9 @@ def __init__(
"the TMEM plan (O 256 + 2 x S 64 + Q 128 columns) is laid out for D = 256"
)
assert group_size * rows_per_tile <= 128
+ # The O staging's four SW128 boxes sit `64 G R` elements apart: each must start on
+ # a 1024-byte swizzle atom, so the live MMA rows come in whole 8-row atoms.
+ assert group_size * rows_per_tile % 8 == 0, (group_size, rows_per_tile)
self.group_size = group_size
self.rows_per_tile = rows_per_tile
self.d = head_dim
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_sparse_gqa_rowtile__prefill.py
similarity index 86%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_sparse_gqa_rowtile__prefill.py
index b91eee94fbeb..2d7cb622703f 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_sparse_gqa_rowtile__prefill.py
@@ -12,6 +12,10 @@
group, each row masked to its own bits. A group whose rows span requests
(generation rows of a mixed batch) walks each row's own selected columns
instead, as the per-row kernel does.
+
+A row group is sized to whole 8-row swizzle atoms (10 rows of 12 heads at a dep4 rank, 20 of 6
+at a tp4 / tep4 rank); scratch is bounded by the group's selections rather than the table's
+width.
"""
from __future__ import annotations
@@ -28,14 +32,16 @@
from cutlass.cutlass_dsl import T, dsl_user_op
from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa import (
+ qsa_paged_sparse_gqa as _per_row_entry,
+)
from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
-from .qsa_paged_sparse_gqa import qsa_paged_sparse_gqa as _per_row_entry
-from .qsa_rowtile_attention_cute import BLOCK_M, BLOCK_N
-from .qsa_rowtile_attention_cute import eligible as _cute_eligible
-from .qsa_rowtile_attention_cute import num_ctas as _rowtile_ctas
-from .qsa_rowtile_attention_cute import rowtile_attention_cute as _rowtile_attention_cute
+from .qsa_rowtile_attention_cute__prefill import BLOCK_M, BLOCK_N
+from .qsa_rowtile_attention_cute__prefill import eligible as _cute_eligible
+from .qsa_rowtile_attention_cute__prefill import num_ctas as _rowtile_ctas
+from .qsa_rowtile_attention_cute__prefill import rowtile_attention_cute as _rowtile_attention_cute
_INTEGER_DTYPES = (torch.int32, torch.int64)
# The per-row kernel's scan tile: the prefill bound is rounded up to it, so the
@@ -43,12 +49,15 @@
_BOUND_TILE = 64
# Programs of the lone-row pass per KV head; each walks the lone list with this stride.
_LONE_PROGRAMS = 256 # capacity: fixed: a launch width; the lone list may be any length
-# Bitmap words the plan pass unions and packs per loop step.
-_PLAN_WORDS = 16 # capacity: fixed: a loop step; the plan loops over every word
+# Bitmap words the plan pass unions and packs per loop step: wider for the 10-row groups of a
+# dep4 rank, narrower for the 20-row groups of a tp4 / tep4 rank.
+_PLAN_WORDS = {16: 64, 32: 16} # capacity: fixed: a loop step; the plan loops over every word
# A one-request group takes the row-tile pass when its list's 64-slot tiles, each costing
-# this many quarters of a per-row tile (2.5x), cost no more than its rows' own columns.
-_SHARED_TILE_COST = 10
+# this many quarters of a per-row tile (1.5x), cost no more than its rows' own columns.
+_SHARED_TILE_COST = 6
_PLAN_WARPS = 4
+#: The O staging's swizzle atom in MMA rows: a group's rows x heads come in whole atoms.
+_ROW_ATOM = 8
_LONE_WARPS = 2
_LONE_STAGES = 2
@@ -262,8 +271,8 @@ def _compress_bits(x, m):
def _rowtile_plan_kernel(
request_indices, row_info, bits, block_table, slots, packed_bits, group_plan,
lone_rows, lone_count, rows,
- bits_stride_row, block_table_stride_request, block_table_stride_page, slots_stride_group,
- NUM_REQUESTS, NUM_CACHE_PAGES,
+ bits_stride_row, packed_stride_row, block_table_stride_request, block_table_stride_page,
+ slots_stride_group, NUM_REQUESTS, NUM_CACHE_PAGES,
ROWS_PER_TILE: tl.constexpr, ROWS_POW2: tl.constexpr,
BLOCK_N: tl.constexpr, SHARED_TILE_COST: tl.constexpr,
PAGE_TABLE_WIDTH: tl.constexpr, TOKENS_PER_BLOCK: tl.constexpr, PLAN_WORDS: tl.constexpr,
@@ -327,7 +336,7 @@ def _rowtile_plan_kernel(
resident = tl.reduce(tl.where(kept, lane_bits[None, :], 0), 1, _or_combine)
packed = _compress_bits(row_words, resident[None, :])
shift = (first % 64).to(tl.uint64)
- destination = packed_bits + members[:, None] * bits_stride_row + (first // 64)[None, :]
+ destination = packed_bits + members[:, None] * packed_stride_row + (first // 64)[None, :]
writes = live[:, None] & in_range[None, :] & (packed != 0)
tl.atomic_or(destination, (packed << shift[None, :]).to(tl.int64, bitcast=True),
mask=writes, sem="relaxed") # fmt: skip
@@ -527,6 +536,15 @@ def _check_inputs(q, k_cache, v_cache, block_table, selected_tokens, request_ind
)
+def rows_per_group(group_size: int) -> int:
+ """Rows of one request a row-tile item serves: the most whose `rows x group_size` MMA
+ rows fill whole 8-row swizzle atoms of the 128-row tile (10 at 12 heads, 20 at 6)."""
+ rows = BLOCK_M // group_size
+ while rows > 1 and rows * group_size % _ROW_ATOM:
+ rows -= 1
+ return rows
+
+
def qsa_rowtile_sparse_gqa(
*,
q: torch.Tensor,
@@ -555,7 +573,7 @@ def qsa_rowtile_sparse_gqa(
bits, row_info = _row_bits(
selected_tokens, query_positions, compress_ratio, tokens_per_block, table_width, num_words
)
- rows_per_tile = BLOCK_M // group_size
+ rows_per_tile = rows_per_group(group_size)
rows_pow2 = triton.next_power_of_2(rows_per_tile)
kv_strides = (
k_cache.stride(0), k_cache.stride(1), k_cache.stride(2), k_cache.stride(3),
@@ -571,8 +589,12 @@ def qsa_rowtile_sparse_gqa(
) # fmt: skip
out_strides = (output.stride(0), output.stride(1), output.stride(2))
num_groups = triton.cdiv(rows, rows_per_tile)
- slots = torch.empty((num_groups, num_words * 64), dtype=torch.int32, device=q.device)
- packed_bits = torch.zeros((rows, num_words), dtype=torch.int64, device=q.device)
+ # A group lists at most its rows' columns: the bound on its slots and re-indexed bits,
+ # far below the table's tokens at a deep position.
+ listed_cap = min(num_words * 64, rows_per_tile * selected_tokens.shape[1])
+ packed_words = triton.cdiv(listed_cap, 64)
+ slots = torch.empty((num_groups, packed_words * 64), dtype=torch.int32, device=q.device)
+ packed_bits = torch.zeros((rows, packed_words), dtype=torch.int64, device=q.device)
group_plan = torch.empty((num_groups, 2), dtype=torch.int32, device=q.device)
lone_rows = torch.empty((rows,), dtype=torch.int32, device=q.device)
# [0]: the lone list's length; [1]: the row-tile pass's work-ticket counter.
@@ -581,11 +603,12 @@ def qsa_rowtile_sparse_gqa(
_rowtile_plan_kernel[(num_groups,)](
request_indices, row_info, bits, block_table, slots, packed_bits, group_plan,
lone_rows, lone_count, rows,
- bits.stride(0), block_table.stride(0), block_table.stride(1), slots.stride(0),
- block_table.shape[0], k_cache.shape[0],
+ bits.stride(0), packed_bits.stride(0), block_table.stride(0), block_table.stride(1),
+ slots.stride(0), block_table.shape[0], k_cache.shape[0],
ROWS_PER_TILE=rows_per_tile, ROWS_POW2=rows_pow2, BLOCK_N=BLOCK_N,
SHARED_TILE_COST=_SHARED_TILE_COST, PAGE_TABLE_WIDTH=table_width,
- TOKENS_PER_BLOCK=tokens_per_block, PLAN_WORDS=_PLAN_WORDS, num_warps=_PLAN_WARPS,
+ TOKENS_PER_BLOCK=tokens_per_block, PLAN_WORDS=_PLAN_WORDS.get(rows_pow2, 16),
+ num_warps=_PLAN_WARPS,
) # fmt: skip
use_pdl = get_env_enable_pdl()
_rowtile_lone_kernel[(min(rows, _LONE_PROGRAMS), num_kv_heads)](
@@ -614,23 +637,30 @@ def qsa_rowtile_sparse_gqa(
token_topk=2048,
cache_dtype=torch.bfloat16,
q_dtype=torch.bfloat16,
- # K and V are the two role planes of one [pages, 2, kv, tpb, D] pool.
- cache_layout="pool_planes",
- # q is the first 24 heads of the packed [T, 24 + 2 + 2, 256] QKV buffer.
+ # q is the first heads of the packed [T, Hq + 2 Hkv, 256] QKV buffer.
q_layout="qkv_slice",
request_dtype=torch.int32,
- # The engine's table at max_seq_len 10240: 320 pages per slot, -1 past each sequence.
- pages=320,
+ # The engine's table at max_seq_len 262144: `ceil(262145 / 32)` columns padded
+ # to a multiple of 4, -1 past each sequence.
+ pages=8196,
)
-
-def _spec(segments, slots=32, **overrides) -> dict:
- """`segments` is `(first position, rows)` per request, in row order."""
- return {**_CHECKPOINT, "segments": tuple(segments), "slots": slots, **overrides}
+#: The engine's trunk (`Qwen4ExpHybridCacheManager.get_buffers`): a lifecycle slot holds the 24
+#: role pages of the 12 attention layers, K then V each; the call's K and V are layer 11's two
+#: planes, and the requests' pages straddle 2^31 elements of them.
+_TRUNK = dict(cache_layout="trunk", slot_pages=24, layer=11, offset_edge=2**31)
+#: A tp4 / tep4 rank (attention TP 4): 6 q heads and one of the 2 KV heads, a 128-row table.
+_TP4 = dict(q_heads=6, kv_heads=1, slots=128, **_TRUNK)
+#: A dep4 rank (attention DP): all 24 / 2 heads, a 32-row table; core's bf16 trunk, or the
+#: kv_fp8 variant's pages widened into a contiguous `[pages, 2, 32, 256]` bf16 scratch
+#: (`qsa_fp8_pages_to_bf16`) addressed through its own table.
+_DEP4 = dict(q_heads=24, kv_heads=2, slots=32, **_TRUNK)
+_DEP4_SCRATCH = dict(q_heads=24, kv_heads=2, slots=32, cache_layout="scratch")
-def _decodes(count, rows_each=1, first=8200):
- return [(first + (37 * g) % 1000, rows_each) for g in range(count)]
+def _spec(segments, member, **overrides) -> dict:
+ """`segments` is `(first position, rows)` per request, in row order."""
+ return {**_CHECKPOINT, **member, "segments": tuple(segments), **overrides}
class QsaSparseGqaRowtile(OpWrapper):
@@ -654,18 +684,19 @@ class QsaSparseGqaRowtile(OpWrapper):
fresh, contiguous `[rows, Hq, D]` tensor of `q.dtype`.
Four launches. A bitmap pass: one bit per selected logical token per row.
- A plan pass per group of 10 consecutive rows of one request: the union of
+ A plan pass per group of `rows_per_group` consecutive rows of one request
+ (10 at a dep4 rank's 12 heads a KV head, 20 at a tp4 / tep4 rank's 6): the union of
the rows' selected tokens whose page is mapped, listed in logical order
as physical cache slots, and each row's bitmap re-indexed onto that list
(a per-word parallel bit extract). The group goes to the row-tile pass
- when walking its list in 64-slot tiles costs at most 2.5x per-row tiles
+ when walking its list in 64-slot tiles costs at most 1.5x per-row tiles
less than the rows' own columns would; otherwise -- or when its rows span
requests -- its rows join the lone list. A lone-row pass (256 programs
per KV head striding the list) serves listed rows over their own columns,
as the per-row kernel does. The row-tile pass is persistent (one program
per SM drawing (group, KV head) items from a ticket counter, the plan's
last groups first); per item it loads each listed tile's K/V once for the
- 120 (row, head) pairs and masks each row to its own re-indexed bits, and
+ group's 120 (row, head) pairs and masks each row to its own re-indexed bits, and
one item's epilogue overlaps the next item's first MMAs. It is launched
programmatically dependent on the lone pass, whose rows are disjoint, so
the two overlap. That pass is `qsa_rowtile_attention_cute` (tcgen05 MMAs,
@@ -681,16 +712,28 @@ class QsaSparseGqaRowtile(OpWrapper):
CELLS: tuple[Cell, ...] = (
Cell(
- why="con 1 prefill: 8192 rows of one request from 0, across the budget",
- spec=_spec([(0, 8192)]),
+ why="tp4/tep4 cold prefill: 8192 rows of one request from 0, 6 q / 1 kv heads "
+ "(20-row groups), trunk pool straddling 2^31, [128, 8196] table",
+ spec=_spec([(0, 8192)], _TP4),
+ ),
+ Cell(
+ why="tp4/tep4 deep chunk: 8192 rows from 188416 ranking one score vector "
+ "(a sparse union), 6 q / 1 kv heads, trunk pool straddling 2^31",
+ spec=_spec([(188416, 8192)], _TP4, shared_scores=True),
+ ),
+ Cell(
+ why="tp4/tep4 2049 rows from 9000: a last group of 9 rows, 6 q / 1 kv heads",
+ spec=_spec([(9000, 2049)], _TP4),
),
Cell(
- why="sparse union: 4096 rows from 4096 ranking one score vector",
- spec=_spec([(4096, 4096)], shared_scores=True),
+ why="dep4 kv_fp8 deep chunk: 8160 rows from 188416 over the widened bf16 "
+ "scratch, 24 q / 2 kv heads (10-row groups), [32, 8196] table",
+ spec=_spec([(188416, 8160)], _DEP4_SCRATCH, shared_scores=True),
),
Cell(
- why="con 128 mixed: 3 chunks (16240 rows) + 52 decodes on the lone pass",
- spec=_spec([(0, 8192), (2192, 6000), (6144, 2048)] + _decodes(52), slots=128),
+ why="dep4 core: three context requests (0+3000, 40000+4000, 120000+1160) in one "
+ "launch, 24 q / 2 kv heads, trunk pool straddling 2^31",
+ spec=_spec([(0, 3000), (40000, 4000), (120000, 1160)], _DEP4),
),
)
@@ -706,8 +749,7 @@ class QsaSparseGqaRowtile(OpWrapper):
pass attends to the *set* of tokens a row names; the per-row kernel
attends to every column, so a repeated index would weigh twice there
and once here. QSA's top-k + expansion never repeats one (distinct
- blocks, and a tail outside every complete block): measured zero
- repeats on four layers of the target's own 8192-token prefill.
+ blocks, and a tail outside every complete block).
Rows served by the lone-row pass keep per-column semantics.
- As for the per-row kernel: each row is prefix-packed with exactly the
bound's count of valid entries; entries past the 64-rounded bound
@@ -715,29 +757,19 @@ class QsaSparseGqaRowtile(OpWrapper):
Work: a row-tile group walks its packed list, so its MMA work grows
with the union of its rows' selections rather than with the group's
- highest position. On the target's own selections (8192-token prefill,
- layers 0/3/6/11) that union is 13-32% shorter than every 64-token tile up
- to the highest token, and the walk loads ~5x fewer K/V tiles than the
- per-row kernel. Measured on GB300 (eager, whole call) at 8192 context
- rows of real selections: 1544-1565 us on layers 0-6 and 1274 us on layer
- 11, against 1711-1715 us for a walk of every tile up to the highest token
- and 2953 us for `qsa_paged_sparse_gqa`. Uniformly random selections
- leave almost no gap in the union, and the call costs what the full walk
- does (16292-row mixed step: 3859 vs 3836 us). Below a few hundred rows
- it is slower than the per-row kernel (generation-only launches 1.4-2x);
- the target's served context launches are thousands of rows. The tensor-core
- row-tile pass alone, on six dumped 16384-row context launches of the
- target (GB300, layers 0-5 of one step): 581-775 us, against 694-886 us for
- the one-program-per-item pass it replaced (the same output, bit for bit).
+ highest position. Below a few hundred rows it is slower than the per-row
+ kernel; the target's served context launches are thousands of rows.
Compiled once per (group size, rows per tile, tokens per block, PDL):
nothing in the row-tile pass's compile key depends on the row count or
the plan.
- Scratch per call: the bitmap and the re-indexed bitmap, each `rows x
- ceil(table_width * tpb / 64)` int64 words (21 MB at 16384 rows and 320
- pages), the slot lists, `ceil(rows / 10) x` that many words' tokens as
- int32 (67 MB at 16384 rows), the plan and the lone list. Calls allocate,
+ Scratch per call: the bitmap, `rows x ceil(table_width * tpb / 64)` int64
+ words (267 MB at 8160 rows and an 8196-page table), the re-indexed
+ bitmap and the slot lists, bounded by a group's selections (`rows x
+ ceil(rows_per_group * 2051 / 64)` words, 42 MB, and `ceil(rows /
+ rows_per_group) x` that many words' tokens as int32, 134 MB at 8160 rows),
+ the plan and the lone list. Calls allocate,
so capture in a CUDA graph is not certified; this entry serves the
prefill path, which runs eager. Results are deterministic: the bitmap
OR, the re-indexed bitmap OR and the lone list order do not change any
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qwen4_exp_attn_decode_prologue__decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qwen4_exp_attn_decode_prologue__decode.py
new file mode 100644
index 000000000000..ff83f898bf92
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qwen4_exp_attn_decode_prologue__decode.py
@@ -0,0 +1,719 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""A sparse decode attention layer's work between its input projection and its index
+scores, one PDL launch, at one rank's heads over a bf16 or e4m3 K/V pool.
+
+The rotated q and k and the raw v go into the compact packed `[q | k | v]` buffer the sparse
+GQA and the dense op read.
+
+It fuses the composition's seven launches: the main q/k norm + rope in place on the
+projection (`fused_qk_norm_rope`), the index-scratch copy and the index plane's norm + rope
+(a second `fused_qk_norm_rope`), the index key and coordinate store
+(`qsa_paged_index_store`), the group compression (`qsa_prefill_compress`), the packed-QKV
+copy and the K/V append (`qsa_paged_kv_store`). Per row: programs of `HPB` main heads
+normalize and rotate q and k into the packed buffer, and those holding k heads also store
+them into the K pool; one program normalizes and rotates the index query into the index
+scratch; one program stores the raw index key and the coordinates and -- at a group's last
+token -- pools the group with this token's key taken from the projection row, normalizes,
+rotates and stores it at the anchor slot, as the compression kernel does; one program per v
+head stores v into the V pool and the packed buffer.
+
+A verify step (`rows_per_request` > 1, each request's rows consecutive at consecutive
+positions) writes several raw keys of one group in the launch: at a group's last token the
+earlier members held by the same request's preceding rows are taken from their projection
+rows, the others from the index plane.
+
+With PDL every operand that is static or host-filled -- the page and slot, the rotation
+factors at the row's positions, the norm weights, and at a group end the group's earlier
+keys from the index plane with their rotation -- is read before `griddepcontrol.wait`;
+after it a program reads only projection rows.
+"""
+
+from __future__ import annotations
+
+from typing import Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
+
+
+@triton.jit
+def _rope_factors(positions, num_tokens, token, cos_sin, stride_cs, D: tl.constexpr, ROT: tl.constexpr,
+ S1: tl.constexpr, S2: tl.constexpr): # fmt: skip
+ """A D-wide head's rotation factors at the row's positions (interleaved mRoPE: pair j
+ reads coordinate 1 when j % 3 == 1 and j < 3 * S1, coordinate 2 when j % 3 == 2 and
+ j < 3 * S2, else 0): cos, and sin signed for the NeoX halves; identity past ROT."""
+ HALF: tl.constexpr = ROT // 2
+ d = tl.arange(0, D)
+ rot = d < ROT
+ j = d % HALF
+ p0 = tl.load(positions + token)
+ p1 = tl.load(positions + num_tokens + token)
+ p2 = tl.load(positions + 2 * num_tokens + token)
+ pos = tl.where((j % 3 == 1) & (j < 3 * S1), p1, tl.where((j % 3 == 2) & (j < 3 * S2), p2, p0))
+ row = cos_sin + pos.to(tl.int64) * stride_cs
+ c = tl.load(row + j, mask=rot, other=1.0)
+ s = tl.load(row + HALF + j, mask=rot, other=0.0)
+ s = tl.where(d < HALF, -s, s)
+ return c, s
+
+
+@triton.jit
+def _norm_weights(
+ heads, weight, weight_k, D: tl.constexpr, ROT: tl.constexpr, WEIGHT_SPLIT: tl.constexpr
+):
+ """`1 + w` for a tile's elements and for their rotation partners; heads at or past
+ WEIGHT_SPLIT (when > 0) take the second weight (the k heads)."""
+ HALF: tl.constexpr = ROT // 2
+ d = tl.arange(0, D)
+ rot = d < ROT
+ partner = tl.where(d < HALF, d + HALF, tl.where(rot, d - HALF, d))
+ w = tl.load(weight + d).to(tl.float32)
+ wp = tl.load(weight + partner).to(tl.float32)
+ if WEIGHT_SPLIT > 0:
+ is_k = (heads >= WEIGHT_SPLIT)[:, None]
+ wk = tl.load(weight_k + d).to(tl.float32)
+ wkp = tl.load(weight_k + partner).to(tl.float32)
+ w2 = tl.where(is_k, wk[None, :], w[None, :]) + 1.0
+ wp2 = tl.where(is_k, wkp[None, :], wp[None, :]) + 1.0
+ else:
+ w2 = w[None, :] + 1.0
+ wp2 = wp[None, :] + 1.0
+ return w2, wp2
+
+
+@triton.jit
+def _norm_rope_apply(src, dst, heads, head_mask, w2, wp2, c, s, eps, D: tl.constexpr, ROT: tl.constexpr,
+ dst_head_stride): # fmt: skip
+ """rms over each head, `* (1 + w)`, rotate; fp32 throughout, one rounding to `dst`'s
+ dtype, stored at `dst`. Returns the stored (rounded) tile."""
+ HALF: tl.constexpr = ROT // 2
+ d = tl.arange(0, D)
+ rot = d < ROT
+ partner = tl.where(d < HALF, d + HALF, tl.where(rot, d - HALF, d))
+ hm = head_mask[:, None]
+ x = tl.load(src + heads[:, None] * D + d[None, :], mask=hm, other=0.0).to(tl.float32)
+ xp = tl.load(src + heads[:, None] * D + partner[None, :], mask=hm & rot[None, :], other=0.0).to(
+ tl.float32
+ )
+ rstd = tl.rsqrt(tl.sum(x * x, axis=1) / D + eps)
+ y = x * rstd[:, None] * w2
+ yp = xp * rstd[:, None] * wp2
+ out = (y * c[None, :] + yp * s[None, :]).to(dst.dtype.element_ty)
+ tl.store(dst + heads[:, None] * dst_head_stride + d[None, :], out, mask=hm)
+ return out
+
+
+@triton.jit
+def _compress_norm_rope(x, weight, cosine, sine, eps, HEAD_DIM: tl.constexpr, ROTARY_PAIRS: tl.constexpr): # fmt: skip
+ """QSA's `_qsa_gemma_norm_rope` for one row with its static factors given (`weight` is
+ `1 + w`; the fp32 table makes its intermediate rounding to the table's dtype a no-op)."""
+ x_f32 = x.to(tl.float32)
+ reciprocal_rms = tl.rsqrt(tl.sum(x_f32 * x_f32, axis=1) / HEAD_DIM + eps)
+ normalized = x_f32 * reciprocal_rms[:, None] * weight[None, :]
+ rotated, passthrough = tl.split(
+ tl.permute(tl.reshape(normalized, (1, 2, HEAD_DIM // 2)), (0, 2, 1))
+ )
+ first, second = tl.split(tl.permute(tl.reshape(rotated, (1, 2, ROTARY_PAIRS)), (0, 2, 1)))
+ output_first = first * cosine[None, :] - second * sine[None, :]
+ output_second = second * cosine[None, :] + first * sine[None, :]
+ rotated = tl.reshape(
+ tl.permute(tl.join(output_first, output_second), (0, 2, 1)), (1, HEAD_DIM // 2)
+ )
+ return tl.reshape(
+ tl.permute(tl.join(rotated, passthrough), (0, 2, 1)),
+ (1, HEAD_DIM),
+ ).to(tl.bfloat16)
+
+
+@triton.jit(do_not_specialize=["num_tokens", "num_requests", "table_width", "num_pages"])
+def _attn_decode_prologue_kernel(
+ proj,
+ stride_proj,
+ out,
+ stride_out,
+ index_q,
+ stride_iq_token,
+ stride_iq_head,
+ positions,
+ num_tokens,
+ cos_sin,
+ stride_cs,
+ q_weight,
+ k_weight,
+ index_q_weight,
+ index_k_weight,
+ eps,
+ k_cache,
+ v_cache,
+ s_kc_page,
+ s_kc_head,
+ s_kc_tok,
+ s_vc_page,
+ s_vc_head,
+ s_vc_tok,
+ request_indices,
+ logical_positions,
+ block_table,
+ s_bt_req,
+ s_bt_col,
+ num_requests,
+ table_width,
+ num_pages,
+ index_cache,
+ s_ic_page,
+ s_ic_slot,
+ position_cache,
+ s_pc_page,
+ s_pc_slot,
+ s_pc_axis,
+ coords,
+ s_c_row,
+ s_c_axis,
+ HQ: tl.constexpr,
+ HK: tl.constexpr,
+ D: tl.constexpr,
+ IH: tl.constexpr,
+ ID: tl.constexpr,
+ IQ_COL: tl.constexpr,
+ ROT: tl.constexpr,
+ S1: tl.constexpr,
+ S2: tl.constexpr,
+ HPB: tl.constexpr,
+ MAIN_PROGRAMS: tl.constexpr,
+ TPB: tl.constexpr,
+ CR: tl.constexpr,
+ ROTARY_PAIRS: tl.constexpr,
+ IN_LAUNCH_GROUPS: tl.constexpr,
+ USE_PDL: tl.constexpr,
+):
+ token = tl.program_id(0).to(tl.int64)
+ program = tl.program_id(1)
+ if USE_PDL:
+ # The consumer (the index scores) reads nothing of this launch before its own wait.
+ tl.extra.cuda.gdc_launch_dependents()
+ row = proj + token * stride_proj
+ out_row = out + token * stride_out
+
+ # The row's page and slot (the K/V pool, the index plane and the coordinates share
+ # the table); an unmapped row writes no cache. Host-filled: read before the wait.
+ request = tl.load(request_indices + token).to(tl.int64)
+ logical = tl.load(logical_positions + token).to(tl.int64)
+ column = logical // TPB
+ slot = logical % TPB
+ live = (request >= 0) & (request < num_requests) & (logical >= 0) & (column < table_width)
+ page = tl.load(block_table + request * s_bt_req + column * s_bt_col, mask=live, other=-1).to(
+ tl.int64
+ )
+ live = live & (page >= 0) & (page < num_pages)
+
+ if program < MAIN_PROGRAMS:
+ main_heads = program * HPB + tl.arange(0, HPB)
+ c, s = _rope_factors(positions, num_tokens, token, cos_sin, stride_cs, D, ROT, S1, S2)
+ w2, wp2 = _norm_weights(main_heads, q_weight, k_weight, D, ROT, HQ)
+ if USE_PDL:
+ tl.extra.cuda.gdc_wait()
+ rotated = _norm_rope_apply(
+ row, out_row, main_heads, main_heads < HQ + HK, w2, wp2, c, s, eps, D, ROT, D
+ )
+ if program * HPB + HPB > HQ:
+ # The rounded rows the composition's store read back from the packed buffer.
+ is_k = (main_heads >= HQ) & (main_heads < HQ + HK) & live
+ k_heads = (main_heads - HQ).to(tl.int64)[:, None]
+ d = tl.arange(0, D)
+ tl.store(
+ k_cache + page * s_kc_page + k_heads * s_kc_head + slot * s_kc_tok + d[None, :],
+ rotated.to(k_cache.dtype.element_ty),
+ mask=is_k[:, None],
+ )
+ elif program == MAIN_PROGRAMS:
+ index_heads = tl.arange(0, IH)
+ c, s = _rope_factors(positions, num_tokens, token, cos_sin, stride_cs, ID, ROT, 0, 0)
+ w2, wp2 = _norm_weights(index_heads, index_q_weight, index_q_weight, ID, ROT, 0)
+ if USE_PDL:
+ tl.extra.cuda.gdc_wait()
+ _norm_rope_apply(
+ row + IQ_COL, index_q + token * stride_iq_token, index_heads, index_heads < IH,
+ w2, wp2, c, s, eps, ID, ROT, stride_iq_head,
+ ) # fmt: skip
+ elif program >= MAIN_PROGRAMS + 2:
+ # One v head a program: the V pool, as the K/V append stores it, and the packed buffer.
+ h = program - MAIN_PROGRAMS - 2
+ d = tl.arange(0, D)
+ if USE_PDL:
+ tl.extra.cuda.gdc_wait()
+ v = tl.load(row + (HQ + HK + h) * D + d)
+ tl.store(out_row + (HQ + HK + h) * D + d, v)
+ if live:
+ tl.store(
+ v_cache + page * s_vc_page + h * s_vc_head + slot * s_vc_tok + d,
+ v.to(v_cache.dtype.element_ty),
+ )
+ else:
+ # Before the wait: the coordinates; at a group's last token, the group's earlier
+ # raw keys (written by earlier steps), the first token's coordinate 0, its rotation
+ # and the key norm weight.
+ axes = tl.arange(0, 4)
+ coordinate = tl.load(coords + token * s_c_row + axes * s_c_axis, mask=axes < 3, other=0)
+ dims = tl.arange(0, ID)
+ group_end = live & ((logical + 1) % CR == 0)
+ offsets = tl.arange(0, CR)
+ members = logical - (CR - 1) + offsets
+ member_pages = tl.load(
+ block_table + request * s_bt_req + (tl.maximum(members, 0) // TPB) * s_bt_col,
+ mask=group_end,
+ other=0,
+ ).to(tl.int64)
+ earlier = group_end & (offsets < CR - 1)
+ if IN_LAUNCH_GROUPS:
+ # A request's rows come consecutive at consecutive positions: the earlier
+ # member at row `token - (CR - 1) + o`, when that row is this request's at
+ # that position, is written by this launch -- its raw key is taken from its
+ # projection row after the wait, not from the index plane, and as the first
+ # member its coordinates from its row of `coords`.
+ member_rows = token - (CR - 1) + offsets
+ maybe = earlier & (member_rows >= 0)
+ member_requests = tl.load(request_indices + member_rows, mask=maybe, other=-1)
+ member_logical = tl.load(logical_positions + member_rows, mask=maybe, other=-1)
+ in_launch = (
+ maybe
+ & (member_requests.to(tl.int64) == request)
+ & (member_logical.to(tl.int64) == members)
+ )
+ else:
+ in_launch = offsets < 0
+ group = tl.load(
+ index_cache
+ + member_pages[:, None] * s_ic_page
+ + (tl.maximum(members, 0) % TPB)[:, None] * s_ic_slot
+ + dims[None, :],
+ mask=(earlier & ~in_launch)[:, None],
+ other=0.0,
+ )
+ first = tl.maximum(logical - (CR - 1), 0)
+ first_page = tl.load(
+ block_table + request * s_bt_req + (first // TPB) * s_bt_col, mask=group_end, other=0
+ ).to(tl.int64)
+ first_base = position_cache + first_page * s_pc_page + (first % TPB) * s_pc_slot
+ if IN_LAUNCH_GROUPS:
+ # A first member written by this launch has its coordinates in its row.
+ first_in_launch = tl.sum(tl.where(offsets == 0, in_launch, False).to(tl.int32)) > 0
+ cached = group_end & ~first_in_launch
+ pos_t = tl.load(coords + (token - (CR - 1)) * s_c_row, mask=first_in_launch, other=0)
+ pos_t = tl.where(cached, tl.load(first_base, mask=cached, other=0), pos_t)
+ else:
+ pos_t = tl.load(first_base, mask=group_end, other=0)
+ # The compressed key rotates on the first token's coordinate 0 (QSA's plain rope).
+ pairs = tl.arange(0, ROTARY_PAIRS)
+ cs_row = cos_sin + pos_t.to(tl.int64) * stride_cs
+ cosine = tl.load(cs_row + pairs, mask=group_end, other=0.0)
+ sine = tl.load(cs_row + ROTARY_PAIRS + pairs, mask=group_end, other=0.0)
+ key_weight = tl.load(index_k_weight + dims).to(tl.float32) + 1.0
+ if USE_PDL:
+ tl.extra.cuda.gdc_wait()
+ if live:
+ # The coordinates, as the index store writes them.
+ tl.store(
+ position_cache + page * s_pc_page + slot * s_pc_slot + axes * s_pc_axis,
+ coordinate.to(position_cache.dtype.element_ty),
+ mask=axes < 3,
+ )
+ # The raw index key -- or, at a group's last token, the group compressed in
+ # its place (the compression kernel overwrites the same slot).
+ current_k = tl.load(row + IQ_COL + IH * ID + dims)
+ stored_k = current_k.to(index_cache.dtype.element_ty)
+ if (logical + 1) % CR == 0:
+ if IN_LAUNCH_GROUPS:
+ launch_k = tl.load(
+ proj
+ + member_rows[:, None] * stride_proj
+ + IQ_COL
+ + IH * ID
+ + dims[None, :],
+ mask=in_launch[:, None],
+ other=0.0,
+ )
+ group = tl.where(
+ in_launch[:, None], launch_k.to(index_cache.dtype.element_ty), group
+ )
+ group = tl.where(offsets[:, None] == CR - 1, stored_k[None, :], group)
+ pooled = (tl.sum(group.to(tl.float32), axis=0) / CR).to(tl.bfloat16)
+ stored_k = tl.reshape(
+ _compress_norm_rope(
+ pooled[None, :], key_weight, cosine, sine, eps, ID, ROTARY_PAIRS
+ ),
+ (ID,),
+ )
+ tl.store(index_cache + page * s_ic_page + slot * s_ic_slot + dims, stored_k)
+
+
+def _launch(num_tokens: int) -> Tuple[int, int]:
+ """Main-plane heads per program and warps per program, by row count: 1/2 at up to 16
+ rows, 2/2 to 64, 8/8 past."""
+ if num_tokens <= 16:
+ return 1, 2
+ if num_tokens <= 64:
+ return 2, 2
+ return 8, 8
+
+
+#: One tp4 / tep4 rank (attention TP 4): 6 q heads and one of the 2 KV heads, the fused
+#: projection row `[q 6 | k 1 | v 1 | gate 6 | index q 4, k 1 | v slot]` in heads of 256 and
+#: 128 (4352 wide, the index query at column 3584), the packed `[q | k | v]` 2048 wide.
+_TP4 = dict(geom="tp4")
+#: One dep4 rank (attention DP): every head, `[q 24 | k 2 | v 2 | gate 24 | index ...]`
+#: (14080 wide, the index query at 13312), the packed buffer 7168 wide.
+_DEP4 = dict(geom="dep4")
+#: ht_dep4's decode replay: 16 requests at contexts spread 20k..94k.
+_HT16 = tuple(int(20000 + r * (94000 - 20000) / 15) for r in range(16))
+
+
+class Qwen4ExpAttnDecodePrologue(OpWrapper):
+ """The main and index planes' Gemma norm + rope, the packed-QKV copy, the paged K/V append,
+ the index key and coordinate store and the group compression of one sparse decode
+ attention layer, in one launch.
+
+ `proj` `[T, W]` is the layer's fused input projection: q heads `[0, Hq)`, k heads
+ `[Hq, Hq + Hk)` and v heads `[Hq + Hk, Hq + 2 Hk)` of `head_dim`, the index query heads
+ at `index_q_col` and the raw index key right after them; it is not written. Writes:
+ `out` `[T, (Hq + 2 Hk) * head_dim]` (any row stride), q and k normalized (`* (1 + w)`)
+ and rotated, v copied; `index_q` `[T, index_heads, index_head_dim]` (any token and head
+ stride), the index query normalized and rotated on coordinate 0; the rotated k and the
+ raw v into the paged K/V pool (`k_cache` / `v_cache` `[pages, Hk, tpb, head_dim]`, bf16
+ or e4m3, the bf16 values converted as the store converts them); the raw index key (or,
+ at a group's last token, the compressed group) and the coordinates into the index plane
+ and the position cache, at the row's page and slot. Returns `(out, index_q)`.
+
+ The rotation is NeoX over the first `rotary_dim` dims of a head, the factors read from
+ `cos_sin` (fp32 `[P, rotary_dim]`, cos block then sin block) at the row's position:
+ `position_ids` is the contiguous int32 `[3, T]`, pair `j` of the main plane reads
+ coordinate 1 when `j % 3 == 1 and j < 3 * mrope_section1`, coordinate 2 when
+ `j % 3 == 2 and j < 3 * mrope_section2`, coordinate 0 otherwise (interleaved mRoPE); the
+ index plane reads coordinate 0. The compressed key is `qsa_prefill_compress`'s: the
+ group's four raw keys averaged (fp32, rounded to bf16), normalized with
+ `index_k_weight` and rotated at its first token's coordinate 0 in the position cache.
+
+ `rows_per_request` > 1 declares a verify step's layout: each request's rows
+ consecutive, at consecutive positions.
+
+ Fusion boundary: the projection GEMM before it, and the index scores, the selection and
+ the sparse GQA after it, are the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="tp4/tep4 rank, MTP-3 verify: 1 request x 4 rows at 84000..84003 (ll's step; "
+ "a group end with its three earlier members in the launch), bf16 pool",
+ spec=dict(_TP4, pos=(84000,), rows_per_request=4),
+ ),
+ Cell(
+ why="tp4/tep4 rank, verify 2 requests x 4 rows: a group end on a page's last slot "
+ "with one member in the plane, one with all three in the launch; zero, -0.0 and "
+ "outlier blocks in k and v",
+ spec=dict(_TP4, pos=(4093, 8696), rows_per_request=4, edge=True),
+ ),
+ Cell(
+ why="tp4/tep4 rank, MTP-7 verify 2 requests x 8 rows, one up to position 262143 "
+ "(the table's last row)",
+ spec=dict(_TP4, pos=(1001, 262136), rows_per_request=8),
+ ),
+ Cell(
+ why="tp4/tep4 rank, decode 3 rows: a group end, a mid-group row, an unmapped row",
+ spec=dict(_TP4, pos=(8699, 1201, 31), unmapped=2),
+ ),
+ Cell(
+ why="tp4/tep4 rank, decode 72 rows (eight-head programs: the one main program "
+ "masks its eighth head)",
+ spec=dict(_TP4, pos=tuple(1000 + 113 * r for r in range(72))),
+ ),
+ Cell(
+ why="dep4 rank, decode 16 requests at 20k..94k (ht's step), e4m3 pool",
+ spec=dict(_DEP4, pos=_HT16, fp8=True),
+ ),
+ Cell(
+ why="dep4 rank, decode 32 requests (the core's bf16 pool, its graph batch)",
+ spec=dict(_DEP4, pos=tuple(2000 + 151 * r for r in range(32))),
+ ),
+ Cell(
+ why="dep4 rank, e4m3 pool, verify 10 requests x 4 rows (two-head programs)",
+ spec=dict(
+ _DEP4, pos=tuple(3000 + 211 * r for r in range(10)), rows_per_request=4, fp8=True
+ ), # fmt: skip
+ ),
+ )
+
+ note = """
+ Metadata consumed: the QSA request indices, logical positions and block table; the
+ position ids as `[3, T]` (the rope) and as `[T, 3]` (the coordinates stored).
+ Preconditions: the index plane one KV head of `index_head_dim` (bf16), the position
+ cache int32 `[pages, tpb, >=3]`, `compress_ratio` dividing `tokens_per_block`;
+ generation rows: one per request (`rows_per_request` 1: a group's earlier tokens are
+ already in the index plane), or a verify step's `rows_per_request` consecutive rows a
+ request at consecutive positions (a group's earlier tokens in the index plane or in the
+ request's preceding rows). A row whose page lookup fails writes no cache, as the stores
+ and compression skip it; its packed row and index query are still written.
+
+ Launch: one Triton kernel, grid `(T, cdiv(Hq + Hk, HPB) + 2 + Hk)`, HPB and warps by row
+ count (`_launch`); PDL when the engine enables it.
+
+ Numerics: fp32 throughout, one rounding at each store. Against the fp32 `reference`
+ (the same formulas in torch) the rotated q, k, index query and compressed keys differ
+ only where the reduction order of the rms moves a value across a bf16 rounding
+ boundary; v, the raw keys and the coordinates are bitwise. Against the composition the
+ rotation differs: `fused_qk_norm_rope` derives cos/sin with `__sincosf`, this entry reads
+ the target's fp32 table (the one the compression and the stock implementation use),
+ which keeps long positions within bf16 rounding of float64.
+
+ Not certified: NVFP4 pools (the forward keeps the composition there); in-launch index
+ scores; non-zero index-plane mRoPE sections; mRoPE on the compressed key; `T == 0` returns
+ at once; head geometries other than 6/1 and 24/2 x 256 with a 4 x 128 index plane.
+ """
+
+ def __call__(
+ self,
+ proj: torch.Tensor,
+ out: torch.Tensor,
+ index_q: torch.Tensor,
+ position_ids: torch.Tensor,
+ coords: torch.Tensor,
+ cos_sin: torch.Tensor,
+ q_weight: torch.Tensor,
+ k_weight: torch.Tensor,
+ index_q_weight: torch.Tensor,
+ index_k_weight: torch.Tensor,
+ k_cache: torch.Tensor,
+ v_cache: torch.Tensor,
+ index_cache: torch.Tensor,
+ position_cache: torch.Tensor,
+ request_indices: torch.Tensor,
+ logical_positions: torch.Tensor,
+ block_table: torch.Tensor,
+ tokens_per_block: int,
+ compress_ratio: int,
+ num_heads_q: int,
+ num_heads_k: int,
+ head_dim: int,
+ index_heads: int,
+ index_head_dim: int,
+ index_q_col: int,
+ rotary_dim: int,
+ eps: float,
+ mrope_section1: int,
+ mrope_section2: int,
+ rows_per_request: int = 1,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ num_tokens = proj.shape[0]
+ if num_tokens == 0:
+ return out, index_q
+ hpb, warps = _launch(num_tokens)
+ main_programs = triton.cdiv(num_heads_q + num_heads_k, hpb)
+ pdl = get_env_enable_pdl()
+ _attn_decode_prologue_kernel[(num_tokens, main_programs + 2 + num_heads_k)](
+ proj,
+ proj.stride(0),
+ out,
+ out.stride(0),
+ index_q,
+ index_q.stride(0),
+ index_q.stride(1),
+ position_ids,
+ num_tokens,
+ cos_sin,
+ cos_sin.stride(0),
+ q_weight,
+ k_weight,
+ index_q_weight,
+ index_k_weight,
+ eps,
+ k_cache,
+ v_cache,
+ k_cache.stride(0),
+ k_cache.stride(1),
+ k_cache.stride(2),
+ v_cache.stride(0),
+ v_cache.stride(1),
+ v_cache.stride(2),
+ request_indices,
+ logical_positions,
+ block_table,
+ block_table.stride(0),
+ block_table.stride(1),
+ block_table.shape[0],
+ block_table.shape[1],
+ k_cache.shape[0],
+ index_cache,
+ index_cache.stride(0),
+ index_cache.stride(1),
+ position_cache,
+ position_cache.stride(0),
+ position_cache.stride(1),
+ position_cache.stride(2),
+ coords,
+ coords.stride(0),
+ coords.stride(1),
+ HQ=num_heads_q,
+ HK=num_heads_k,
+ D=head_dim,
+ IH=index_heads,
+ ID=index_head_dim,
+ IQ_COL=index_q_col,
+ ROT=rotary_dim,
+ S1=mrope_section1,
+ S2=mrope_section2,
+ HPB=hpb,
+ MAIN_PROGRAMS=main_programs,
+ TPB=tokens_per_block,
+ CR=compress_ratio,
+ ROTARY_PAIRS=cos_sin.shape[1] // 2,
+ IN_LAUNCH_GROUPS=rows_per_request > 1,
+ USE_PDL=pdl,
+ num_warps=warps,
+ launch_pdl=pdl,
+ )
+ return out, index_q
+
+ def reference(
+ self,
+ proj: torch.Tensor,
+ out: torch.Tensor,
+ index_q: torch.Tensor,
+ position_ids: torch.Tensor,
+ coords: torch.Tensor,
+ cos_sin: torch.Tensor,
+ q_weight: torch.Tensor,
+ k_weight: torch.Tensor,
+ index_q_weight: torch.Tensor,
+ index_k_weight: torch.Tensor,
+ k_cache: torch.Tensor,
+ v_cache: torch.Tensor,
+ index_cache: torch.Tensor,
+ position_cache: torch.Tensor,
+ request_indices: torch.Tensor,
+ logical_positions: torch.Tensor,
+ block_table: torch.Tensor,
+ tokens_per_block: int,
+ compress_ratio: int,
+ num_heads_q: int,
+ num_heads_k: int,
+ head_dim: int,
+ index_heads: int,
+ index_head_dim: int,
+ index_q_col: int,
+ rotary_dim: int,
+ eps: float,
+ mrope_section1: int,
+ mrope_section2: int,
+ rows_per_request: int = 1,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Plain torch, fp32, on the given buffers: the norm + rope of both planes, the
+ packed buffer, then every row's K/V, raw index key and coordinate stores, then the
+ compression of every group a row ends (reading the plane those stores left)."""
+ rows, dev = proj.shape[0], proj.device
+ half = rotary_dim // 2
+ pos = position_ids.long()
+ table = cos_sin.float()
+ j = torch.arange(half, device=dev)
+
+ def factors(s1: int, s2: int):
+ coord = torch.zeros(half, dtype=torch.long, device=dev)
+ coord[(j % 3 == 1) & (j < 3 * s1)] = 1
+ coord[(j % 3 == 2) & (j < 3 * s2)] = 2
+ p = pos[coord, :].t() # [T, half]
+ return table[p, j][:, None, :], table[p, half + j][:, None, :]
+
+ def plane(x: torch.Tensor, w: torch.Tensor, cos, sin) -> torch.Tensor:
+ x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * (1.0 + w)
+ x1, x2 = x[..., :half], x[..., half:rotary_dim]
+ return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin, x[..., rotary_dim:]], -1)
+
+ main_heads = num_heads_q + num_heads_k
+ width = main_heads * head_dim
+ v_width = num_heads_k * head_dim
+ weights = torch.cat(
+ [
+ q_weight.float().expand(num_heads_q, head_dim),
+ k_weight.float().expand(num_heads_k, head_dim),
+ ]
+ )
+ main = plane(
+ proj[:, :width].float().reshape(rows, main_heads, head_dim),
+ weights,
+ *factors(mrope_section1, mrope_section2),
+ ).to(out.dtype)
+ out[:, :width] = main.reshape(rows, width)
+ v = proj[:, width : width + v_width]
+ out[:, width : width + v_width] = v
+ iq = proj[:, index_q_col : index_q_col + index_heads * index_head_dim].float()
+ iq = plane(
+ iq.reshape(rows, index_heads, index_head_dim), index_q_weight.float(), *factors(0, 0)
+ )
+ index_q.copy_(iq.to(index_q.dtype))
+ key_col = index_q_col + index_heads * index_head_dim
+ raw = proj[:, key_col : key_col + index_head_dim]
+
+ requests = request_indices.long().tolist()
+ logical = logical_positions.long().tolist()
+ num_requests, table_width = block_table.shape
+
+ def page_of(request: int, position: int):
+ column = position // tokens_per_block
+ if not (0 <= request < num_requests and position >= 0 and column < table_width):
+ return None
+ page = int(block_table[request, column])
+ return page if 0 <= page < k_cache.shape[0] else None
+
+ k_rows = main[:, num_heads_q:]
+ v_rows = v.reshape(rows, num_heads_k, head_dim)
+ slots = [page_of(r, p) for r, p in zip(requests, logical)]
+ for t, page in enumerate(slots):
+ if page is None:
+ continue
+ slot = logical[t] % tokens_per_block
+ k_cache[page, :, slot] = k_rows[t].to(k_cache.dtype)
+ v_cache[page, :, slot] = v_rows[t].to(v_cache.dtype)
+ index_cache[page, slot, 0] = raw[t].to(index_cache.dtype)
+ position_cache[page, slot, :3] = coords[t].to(position_cache.dtype)
+
+ key_weight = 1.0 + index_k_weight.float()
+ rot = 2 * (cos_sin.shape[1] // 2)
+ for t, page in enumerate(slots):
+ if page is None or (logical[t] + 1) % compress_ratio:
+ continue
+ members = [logical[t] - compress_ratio + 1 + o for o in range(compress_ratio)]
+ keys = torch.stack(
+ [
+ index_cache[page_of(requests[t], m), m % tokens_per_block, 0].float()
+ for m in members
+ ]
+ )
+ pooled = (keys.sum(0) / compress_ratio).to(torch.bfloat16).float()
+ first_page = page_of(requests[t], members[0])
+ first = int(position_cache[first_page, members[0] % tokens_per_block, 0])
+ x = pooled * torch.rsqrt(pooled.pow(2).mean() + eps) * key_weight
+ cos, sin = table[first, : rot // 2], table[first, rot // 2 : rot]
+ x1, x2 = x[: rot // 2], x[rot // 2 : rot]
+ key = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin, x[rot:]])
+ index_cache[page, logical[t] % tokens_per_block, 0] = key.to(index_cache.dtype)
+ return out, index_q
+
+ def is_valid(self, proj, out, index_q, position_ids, coords, cos_sin, *args, **kwargs) -> None:
+ assert proj.stride(1) == 1 and proj.dtype == torch.bfloat16, "bf16 proj, unit column stride"
+ assert out.stride(1) == 1 and out.dtype == torch.bfloat16, "bf16 out, unit column stride"
+ assert index_q.stride(2) == 1, "index_q must have a unit dim stride"
+ assert position_ids.is_contiguous() and position_ids.shape == (3, proj.shape[0]), (
+ "contiguous [3, T] positions (the kernel reads coordinate a of row t at a * T + t)"
+ )
+ assert position_ids.dtype == torch.int32, "int32 positions"
+ assert coords.shape == (proj.shape[0], 3), "coords [T, 3]"
+ assert cos_sin.dtype == torch.float32 and cos_sin.stride(1) == 1, (
+ "an fp32 cos|sin table with a unit column stride"
+ )
+
+
+qwen4_exp_attn_decode_prologue = Qwen4ExpAttnDecodePrologue()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/vision_rope_2d_packed_qk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/vision_rope_2d_packed_qk.py
new file mode 100644
index 000000000000..6b8bdc160131
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/vision_rope_2d_packed_qk.py
@@ -0,0 +1,255 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, Tri Dao.
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# * Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# * Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+# Adapted from flash-attn 2 (https://github.com/Dao-AILab/flash-attention),
+# `flash_attn/ops/triton/rotary.py`: `rotary_kernel` and `apply_rotary`.
+"""Vision-tower 2-D rotary, in place on the q|k prefix of the packed qkv projection.
+
+The rotation half of the vision tower's 2-D axial rotary (`vision_rope_2d__r52` is the
+whole op, angle table included). The angle table depends only on the request's
+`grid_thw`, so the caller builds it once per tower forward -- spatial-merge `(row, col)`
+per token, fp64 angles rounded once to fp32, then `cos` / `sin` -- and every one of the 27
+blocks hands it here with that block's qkv projection output.
+
+Tier 5: one Triton launch, flash-attn 2's `rotary_kernel` (BSD-3-Clause, header above)
+cut down to the path this call takes. flash-attn 2.x is not a pinned dependency
+(`requirements.txt` pins only flash-attn 4, and the stock Qwen VL towers guard the 2.x
+import), so the kernel is vendored here rather than imported. What changed from the
+origin: one sequence (no batch axis, `cu_seqlens` or `seqlen_offsets`), NeoX pairing only
+(no `INTERLEAVED` or `CONJUGATE`), int64 token offsets, in place only, a direct
+`kernel[grid]` launch instead of `torch.library.wrap_triton`, and 4 heads per program
+where the origin takes 2. The arithmetic, `BLOCK_M` 8 and the masked half-width loads are
+the origin's.
+
+Nothing in the tree serves it in one launch: the registered rope ops take 1-D
+positions, `fused_dit_qk_norm_rope` always RMS-normalizes q and k, and FlashInfer's
+`apply_rope_with_cos_sin_cache_inplace` raises "misaligned address" at head size 72.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+
+#: This checkpoint's tower: 1152 / 16 heads = a 72-wide head, spatial_merge_size 2, theta 1e4
+#: (`model.visual.blocks.N.attn.qkv.weight` is [3456, 1152], so a qkv row is 3 * 16 * 72).
+NUM_HEADS = 16
+HEAD_DIM = 72
+
+#: Heads x tokens per program: the origin's 8 tokens for a rotary width <= 128, and 4 heads
+#: (faster than the origin's 2).
+_BLOCK_H = 4
+_BLOCK_M = 8
+
+
+@triton.jit
+def _rotary_neox_inplace_kernel(
+ X,
+ COS,
+ SIN,
+ seqlen,
+ nheads,
+ seqlen_ro,
+ stride_x_seqlen,
+ stride_x_nheads,
+ stride_x_headdim,
+ # ROTARY_DIM is constexpr so the compiler sees the mask is constant every 8 elements
+ # and issues 128-bit loads (the origin's comment).
+ ROTARY_DIM: tl.constexpr,
+ BLOCK_K: tl.constexpr,
+ BLOCK_H: tl.constexpr,
+ BLOCK_M: tl.constexpr,
+):
+ ROTARY_DIM_HALF: tl.constexpr = ROTARY_DIM // 2
+ pid_head = tl.program_id(axis=0)
+ pid_m = tl.program_id(axis=1)
+ if pid_m * BLOCK_M >= seqlen:
+ return
+
+ rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H)
+ rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+ # Token offsets in int64: rows x row stride passes 2^31 at 621379 rows of a 3456-wide row.
+ rm64 = rm.to(tl.int64)
+ rk_half = tl.arange(0, BLOCK_K // 2)
+
+ COS = COS + (rm64[:, None] * ROTARY_DIM_HALF + rk_half[None, :])
+ SIN = SIN + (rm64[:, None] * ROTARY_DIM_HALF + rk_half[None, :])
+ mask_cs = (rm[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF)
+ cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32)
+ sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32)
+
+ # The 1st and 2nd halves of each head, rotated as pairs (i, i + ROTARY_DIM_HALF), stored
+ # back where they were read.
+ X = X + (
+ rh[:, None, None] * stride_x_nheads
+ + rm64[None, :, None] * stride_x_seqlen
+ + rk_half[None, None, :] * stride_x_headdim
+ )
+ mask = (
+ (rh[:, None, None] < nheads)
+ & (rm[None, :, None] < seqlen)
+ & (rk_half[None, None, :] < ROTARY_DIM_HALF)
+ )
+ x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32)
+ x1 = tl.load(X + ROTARY_DIM_HALF * stride_x_headdim, mask=mask, other=0.0).to(tl.float32)
+ o0 = x0 * cos - x1 * sin
+ o1 = x0 * sin + x1 * cos
+ tl.store(X, o0, mask=mask)
+ tl.store(X + ROTARY_DIM_HALF * stride_x_headdim, o1, mask=mask)
+
+
+class VisionRope2dPackedQk(OpWrapper):
+ """Rotate every head of `qk` by its token's angles, in place; return `qk`.
+
+ half = head_dim // 2
+ qk[t, h, i] <- x[i] cos[t, i] - x[i + half] sin[t, i]
+ qk[t, h, i + half] <- x[i] sin[t, i] + x[i + half] cos[t, i] i < half
+
+ NeoX half-split pairing `(i, i + half)` over the whole head, in fp32, rounded once on
+ the store. `qk` is the `[T, 2 * heads, head_dim]` q|k prefix of the tower's packed
+ `[T, 3 * heads, head_dim]` qkv projection, read and written through its strides; the
+ v heads that follow in each row are never addressed. `cos` / `sin` are the caller's
+ per-token fp32 tables `[T, half]`: the height angles, then the width angles.
+
+ Fusion boundary: the rotation. The angle table, the qkv projection, the split into
+ q / k / v and the attention are the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: grid -- `grid_thw` rows (frames, height, width) in patch units, from which
+ # the test builds the caller's angle table the way the target does; grid_dtype --
+ # the integer dtype of that host `grid_thw` (it reaches only the table build);
+ # dtype -- the qkv projection's dtype. Heads 16 x 72, merge 2, theta 1e4 throughout.
+ CELLS: tuple[Cell, ...] = (
+ Cell(
+ why="62x90 serving image, int32 host grid: 5580 rows",
+ spec=dict(grid=((1, 62, 90),), grid_dtype=torch.int32, dtype=torch.bfloat16),
+ ),
+ Cell(
+ why="62x90 with the processor's int64 grid_thw: 5580 rows",
+ spec=dict(grid=((1, 62, 90),), grid_dtype=torch.int64, dtype=torch.bfloat16),
+ ),
+ Cell(
+ why="video (3,8,8) + image (1,16,24), two grid rows: 576 rows",
+ spec=dict(grid=((3, 8, 8), (1, 16, 24)), grid_dtype=torch.int64, dtype=torch.bfloat16),
+ ),
+ Cell(
+ why="video (64,16,32): 32768 rows, the tower's per-call ceiling",
+ spec=dict(grid=((64, 16, 32),), grid_dtype=torch.int64, dtype=torch.bfloat16),
+ ),
+ )
+
+ note = """
+ Metadata consumed: none. Call site: the VisionTower's 27 blocks (multimodal unit),
+ eager only -- the table bakes in the request's grid. Rank-invariant: every member
+ (tp4, tep4, dep4) runs the tower at these shapes.
+
+ Source: Triton, JIT-compiled at first call for the running device; the cells ran that
+ path on sm_103 (GB300). Launch: `(cdiv(heads, 4), cdiv(T, 8))` programs, one per 4
+ heads x 8 tokens, each loading its 8 tokens' cos / sin once and both half-heads of
+ its 4 heads with 64-lane masked loads (36 live). Token offsets are int64, so no
+ address product overflows (the origin's int32 product would pass 2^31 at 621379 rows
+ of the 3456-wide qkv row); the cos / sin index is `t * half`.
+
+ Preconditions: `qk` 3-D, any strides (the kernel reads all three); `cos` and `sin`
+ fp32, contiguous, both `[T, head_dim // 2]`. In place: the return value is `qk`
+ itself, and the caller's qkv buffer holds the rotated q|k on return.
+
+ Numerics: bitwise equal to `reference` (a complex multiply in fp32, rounded once) and
+ to `vision_rope_2d__r52.reference` (positions by nested loops, fp64 angles, `polar`)
+ at every cell, q and k alike: relative L2 0, cosine 1.
+
+ Silent, refused by `is_valid` (each demonstrated unguarded in the test):
+ - a bf16 (or any non-fp32) table: the kernel widens it and rotates by rounded
+ angles;
+ - a table with fewer rows than `qk`: rows past it load cos 1 / sin 0, so the tail
+ tokens come back unrotated;
+ - a table narrower than half the head: it rotates a prefix as pairs
+ `(i, i + width)` and leaves the rest of the head untouched (partial rotary);
+ wider would write across the next head;
+ - a non-contiguous table (e.g. the first half of a `[T, head_dim]` duplicated
+ table): it is indexed as `t * half + i` and reads the wrong angles.
+
+ Not certified: float16 / float32 / float8 q|k; head geometries other than 16 x 72;
+ partial rotary; grids past the cells (sides past 90, calls past 32768 rows); T = 0;
+ CUDA-graph capture; out-of-place use.
+ """
+
+ def __call__(self, qk: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+ seqlen, nheads, _ = qk.shape
+ rotary_dim = 2 * cos.shape[-1]
+ grid = (triton.cdiv(nheads, _BLOCK_H), triton.cdiv(seqlen, _BLOCK_M))
+ _rotary_neox_inplace_kernel[grid](
+ qk,
+ cos,
+ sin,
+ seqlen,
+ nheads,
+ cos.shape[0],
+ qk.stride(0),
+ qk.stride(1),
+ qk.stride(2),
+ ROTARY_DIM=rotary_dim,
+ BLOCK_K=triton.next_power_of_2(rotary_dim),
+ BLOCK_H=_BLOCK_H,
+ BLOCK_M=_BLOCK_M,
+ )
+ return qk
+
+ def reference(self, qk: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+ """A fresh tensor: each head's halves as one complex number, times `cos + i sin`, in
+ fp32; rounded once to `qk`'s dtype. `qk` is not mutated."""
+ half = qk.shape[-1] // 2
+ x = qk.float()
+ rotor = torch.complex(cos.float(), sin.float()).unsqueeze(-2)
+ rotated = torch.complex(x[..., :half], x[..., half:]) * rotor
+ return torch.cat((rotated.real, rotated.imag), dim=-1).to(qk.dtype)
+
+ def is_valid(self, qk: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> None:
+ assert cos.dtype == torch.float32 and sin.dtype == torch.float32, (
+ f"cos and sin must be fp32; got {cos.dtype} / {sin.dtype} -- the kernel widens a "
+ "narrower table and rotates by its rounded angles"
+ )
+ half = qk.shape[-1] // 2
+ assert cos.shape == (qk.shape[0], half) and sin.shape == cos.shape, (
+ f"cos and sin must be [T, head_dim // 2] = [{qk.shape[0]}, {half}]; got "
+ f"{tuple(cos.shape)} / {tuple(sin.shape)} -- a short table leaves the tail tokens "
+ "unrotated, a narrow one rotates only a prefix of each head"
+ )
+ assert cos.is_contiguous() and sin.is_contiguous(), (
+ "cos and sin must be contiguous; the kernel indexes them as t * half + i and a "
+ "strided table reads the wrong angles"
+ )
+
+
+vision_rope_2d_packed_qk = VisionRope2dPackedQk()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/gemm/__init__.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/__init__.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/gemm/__init__.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/hc/__init__.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/__init__.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/hc/__init__.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/hc/qwen4_exp_hc_decode_tc_first__decode.py
similarity index 89%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/hc/qwen4_exp_hc_decode_tc_first__decode.py
index 834196008834..80e12f692388 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/hc/qwen4_exp_hc_decode_tc_first__decode.py
@@ -1,11 +1,14 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
-"""Hyper-connection mixer first half at decode row counts (2..256 rows) as one
+"""Hyper-connection mixer first half at decode row counts (1..256 rows) as one
tcgen05 split-K GEMM: the first projection `normed @ [down; block_inject]^T`
with the SiLU and injection activations in its epilogue, so neither the fp32
logits nor a split-K reduce kernel nor `qwen4_exp_hc_lowrank_act` exists.
-Tier 6 (CuTe DSL written for this target). One CTA computes one
+The hyper-connection is replicated on every rank, so its shapes are the checkpoint's at
+every member (tp4, tep4, dep4); only the rows move.
+
+Tier 6 (CuTe DSL). One CTA computes one
`[128 rows, BN columns]` partial over `K / KS` of the 10240-wide stream (grid z
is the 128-row M tile): a TMA warp streams the normed rows (rows past `T`
zero-filled by TMA, never read)
@@ -547,10 +550,12 @@ def hc_decode_tc_first(normed, first, rank, hc_count, with_inject, bn, ks):
#: Tile `(BN, KS)` by row count: `BN` weight rows per CTA (48: 7 tiles cover the
#: 324 real rows of 336) and a `KS`-way K split, one cluster per column tile. One
#: M tile (<= 128 rows) runs 16-way (7 x 16 = 112 CTAs); two M tiles run 8-way, so
-#: the grid (7 x 8 x 2 = 112 CTAs) stays one wave of one CTA an SM.
-_TILES = ((128, (48, 16)), (MAX_ROWS, (48, 8)))
-#: Every K split `_TILES` uses divides K into 64-wide tiles per CTA when K does by this.
-_K_ALIGN = 64 * 16
+#: the grid (7 x 8 x 2 = 112 CTAs) stays one wave of one CTA an SM. Up to 4 rows the
+#: split is 10-way (7 x 10 = 70 CTAs).
+_TILES = ((4, (48, 10)), (128, (48, 16)), (MAX_ROWS, (48, 8)))
+#: Every K split `_TILES` uses (16, 10, 8) divides K into 64-wide tiles per CTA when K
+#: does by this.
+_K_ALIGN = 64 * 80
def _tile(rows: int) -> tuple[int, int]:
@@ -590,68 +595,91 @@ class Qwen4ExpHcDecodeTcFirst(OpWrapper):
ARCHS = frozenset({Arch.SM_103})
# spec: rows, inject (True). The checkpoint's geometry throughout: hc_count 4,
- # rank 320, hidden 2560, the padded [384, K] weight, bf16.
+ # rank 320, hidden 2560, the padded [384, K] weight, bf16. HC is replicated, so the
+ # rows are this rank's tokens at every member.
CELLS: tuple[Cell, ...] = (
Cell(
- why="checkpoint 4x2560, 324 cols of 384, inject, decode 128 rows (one M tile)",
+ why="1 row: a dep4 rank's smallest captured decode batch, 1 of a tile's 128 rows",
+ spec=dict(rows=1),
+ ),
+ Cell(why="2 rows: one request's MTP-1 verify step at tep4", spec=dict(rows=2)),
+ Cell(
+ why="4 rows: one request's MTP-3 verify step (ll_tp4 / ll_tep4, concurrency 1)",
+ spec=dict(rows=4),
+ ),
+ Cell(
+ why="16 rows: a dep4 rank's padded decode batch at concurrency 128",
+ spec=dict(rows=16),
+ ),
+ Cell(
+ why="32 rows: dep4's per-rank graph cap; 4 requests x MTP-7 verify at tep4",
+ spec=dict(rows=32),
+ ),
+ Cell(
+ why="128 rows: one full M tile (a short prefill step's rows)",
spec=dict(rows=128),
),
+ Cell(why="129 rows: a 1-row second M tile", spec=dict(rows=129)),
+ Cell(why="256 rows: MAX_ROWS, two full M tiles", spec=dict(rows=256)),
Cell(
- why="checkpoint 4x2560, 324 cols of 384, inject, 256 rows (MAX_ROWS, two M tiles)",
- spec=dict(rows=256),
+ why="stack-final mixer: 320 cols, no inject, 4 rows (MTP-3 verify)",
+ spec=dict(rows=4, inject=False),
),
Cell(
- why="checkpoint 4x2560, 324 cols of 384, inject, 129 rows (a 1-row second M tile)",
- spec=dict(rows=129),
+ why="stack-final mixer: 320 cols, no inject, 16 rows (dep4 decode)",
+ spec=dict(rows=16, inject=False),
),
Cell(
why="stack-final mixer: 320 cols, no inject, 200 rows (two M tiles)",
spec=dict(rows=200, inject=False),
),
Cell(
- why="stack-final mixer: 320 cols, no inject, 128 rows",
- spec=dict(rows=128, inject=False),
+ why="MTP draft head mixer: 320 cols, no inject, 1 row (ll draft step)",
+ spec=dict(rows=1, inject=False),
+ ),
+ Cell(
+ why="MTP draft head mixer: 320 cols, no inject, 32 rows (dep4 rank batch)",
+ spec=dict(rows=32, inject=False),
),
- Cell(why="checkpoint 4x2560, 324 cols of 384, inject, decode 2 rows", spec=dict(rows=2)),
)
-
note = """
- Tier 6 (CuTe DSL written for this target): tiers 1-5 have no GEMM whose
- epilogue applies an activation after a split-K reduction; cuBLAS at
- N = 384 runs split-K plus its own reduce kernel and writes fp32 logits
- that `qwen4_exp_hc_lowrank_act` reads back (three launches), and the
- Triton `tl.dot` K-loop designs measured 36-60 us at 128 rows.
- Metadata consumed: none. Call sites: `DecodeTarget`'s trunk
- hyper-connection mixers past `_HC_DECODE_GEMV_MAX_ROWS` and up to
- `_HC_TC_FIRST_MAX_ROWS` (= MAX_ROWS) rows, 97 per forward, between
- `qwen4_exp_hc_inject_norm` (or the PLE tail) and
- `qwen4_exp_hc_up_gate_mean(_fp4)`. Preconditions (bf16): `normed`
+ Tier 6 (CuTe DSL): tiers 1-5 have no GEMM whose epilogue applies an
+ activation after a split-K reduction; cuBLAS at N = 384 runs split-K plus
+ its own reduce kernel and writes fp32 logits that `qwen4_exp_hc_lowrank_act`
+ reads back (three launches).
+ Metadata consumed: none. Call sites: `_hc_mixer`, every trunk
+ hyper-connection mixer of a step of at most `_HC_TC_FIRST_MAX_ROWS`
+ (= MAX_ROWS) rows (every captured decode / MTP verify batch of tp4,
+ tep4 and dep4, and short prefill steps), 97 per forward, and with MTP
+ the draft layer's three sites per draft step (`_mtp_front`'s and
+ `_mtp_tail`'s with the injection, `MTPLayer.shared_head`'s without:
+ 1 and 4 rows at ll, a rank's batch and its verify rows at dep4), each
+ between `qwen4_exp_hc_inject_norm` (or the PLE tail) and
+ `qwen4_exp_hc_up_gate_mean`. Preconditions (bf16): `normed`
`[T, K]` contiguous, T <= 256; `first` the `.t()` view of a row-major
`[N, K]` weight with N >= R (+ C); 16-byte aligned; K a multiple of
- 64 * 16. Only the R (+ C) real weight rows are read (TMA zero-fills
+ 64 * 80. Only the R (+ C) real weight rows are read (TMA zero-fills
the tile past them), never the load-time padding.
Launch: grid (KS, cdiv(R + C, BN), cdiv(T, 128)) CTAs of 6 warps,
cluster (KS, 1, 1), one CTA per SM (~227 KB shared memory: the K tiles
of the slice in flight, as many as fit, plus the padded fp32 partial);
- `(BN, KS)` = (48, 16) up to 128 rows (16 x 7 = 112 CTAs), (48, 8) above
- (8 x 7 x 2 = 112 CTAs). The partials are reduced through distributed
- shared memory between two cluster barriers; no global workspace.
+ `(BN, KS)` = (48, 10) up to 4 rows (10 x 7 = 70 CTAs), (48, 16) up to 128
+ rows (16 x 7 = 112 CTAs), (48, 8) above (8 x 7 x 2 = 112 CTAs). The
+ partials are reduced through distributed shared memory between two
+ cluster barriers; no global workspace.
Deterministic (fixed K-slice order). Programmatic dependent launch
(TRTLLM_ENABLE_PDL=0 turns it off): weight stages before
`griddepcontrol.wait`, `normed` after it; dependents released at kernel
start. Graph-capturable (current stream, no host sync, no persistent
state). Compiles once per (BN, KS, C, R, N, K, inject); `precompile`
builds every variant at load.
- Measured on GB300 (graph replay, 96 cold weights, us a site, first half
- alone): 128 rows 5.5 (48 x 16) against 7.9 for cuBLAS split-K + reduce
- + lowrank_act; 256 rows 6.5 (48 x 8) against 9.2. Tile sweep at 129-256
- rows: 48 x 16 (224 CTAs, two waves) 9.6-10.4, 48 x 10 9.5-11.6, 64 x 8
- 6.6-7.4, 112 x 16 7.4-7.9, 32 x 16 16-18; 48 x 8 fastest.
- Certified: bf16, 324 columns with the injection at 2, 128, 129 and 256
- rows, 320 columns without it at 128 and 200 rows; graph replay.
+ Certified: bf16, 324 columns with the injection at 1, 2, 4, 16, 32,
+ 128, 129 and 256 rows, 320 columns without it at 1, 4, 16, 32 and 200
+ rows;
+ graph replay at 1, 4, 16 and 256 rows.
Not certified: hc_count other than 4, an unpadded weight (N = R + C), zero
- rows, a single row (below the dispatch range; accepted, no cell drives
- them), other row counts without the injection.
+ rows (no launch; accepted, no cell drives them), other row counts without
+ the injection.
"""
def __call__(
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/hc/qwen4_exp_hc_up_gate_mean_fp4__moe.py
similarity index 55%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/hc/qwen4_exp_hc_up_gate_mean_fp4__moe.py
index 6d4cfb844b64..5c2968d29a64 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/hc/qwen4_exp_hc_up_gate_mean_fp4__moe.py
@@ -4,13 +4,12 @@
(the up projection as a tcgen05 GEMM, sigmoid gate and stream mean in its
epilogue) whose epilogue also writes the routed experts' NVFP4 input.
-Tier 6 (CuTe DSL written for this target): the `fp4=True` build of
-`qwen4_exp_hc_up_gate_mean`'s kernel. The routed experts read the mixer's output
-only as NVFP4 (packed e2m1 + one e4m3 scale per 16 elements, linear scale
-layout); the bf16 mix is still written for the router and the shared expert.
-Each epilogue thread already holds one row's 32 consecutive mix values, i.e.
-two whole scale blocks, so the quantization needs no data exchange and removes
-the `fp4_quantize` launch and its re-read of the bf16 activation.
+It runs the `fp4=True` build of that kernel (`qwen4_exp_hc_up_gate_mean.hc_up_gate_mean_fp4`).
+The routed experts read the mixer's output only as NVFP4 (packed e2m1 + one
+e4m3 scale per 16 elements, linear scale layout); the bf16 mix is still written
+for the router and the shared expert. Each epilogue thread already holds one
+row's 32 consecutive mix values, i.e. two whole scale blocks, so the
+quantization needs no data exchange.
"""
from __future__ import annotations
@@ -23,8 +22,11 @@
from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import (
fp4_quantize,
)
-
-from .qwen4_exp_hc_up_gate_mean import _BLOCK_H, hc_up_gate_mean_fp4, qwen4_exp_hc_up_gate_mean
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_up_gate_mean import (
+ _BLOCK_H,
+ hc_up_gate_mean_fp4,
+ qwen4_exp_hc_up_gate_mean,
+)
#: NVFP4 block: one e4m3 scale per this many contiguous elements.
SF_VEC = 16
@@ -39,51 +41,58 @@ class Qwen4ExpHcUpGateMeanFp4(OpWrapper):
`mixed` is bitwise `qwen4_exp_hc_up_gate_mean`'s (same kernel, same
epilogue arithmetic); the quantization reads that bf16-rounded value, so
- `xq` / `xsf` are bitwise `torch.ops.trtllm.fp4_quantize(mixed, gs, 16,
- False, False)`. Returns fresh `(mixed [T, H], xq [T, H/2] uint8, xsf
- [T*H/16] uint8)`. `global_scale` is the 1-element fp32 NVFP4 global scale
- (`1 / input_scale`). Fusion boundary: the up projection, the sigmoid
- gate, the stream mean and the NVFP4 quantization of the mean.
+ `xq` / `xsf` are `torch.ops.trtllm.fp4_quantize(mixed, gs, 16, False,
+ False)` outside the reciprocal tie window. Returns fresh `(mixed [T, H],
+ xq [T, H/2] uint8, xsf [T*H/16] uint8)`. `global_scale` is the 1-element
+ fp32 NVFP4 global scale (`1 / input_scale`). Fusion boundary: the up
+ projection, the sigmoid gate, the stream mean and the NVFP4 quantization
+ of the mean.
"""
ARCHS = frozenset({Arch.SM_103})
- # spec: rows, hc_count (4), hidden (2560), rank (320); bf16 only.
+ # spec: rows; hc_count 4, hidden 2560, rank 320 (the checkpoint); bf16 only. Under
+ # attention TP (tp4, tep4) every rank mixes the full 4 x 2560 residual of every row,
+ # so only the row count moves per member.
CELLS: tuple[Cell, ...] = tuple(
Cell(why=f"checkpoint 4x2560, rank 320, {why}", spec=dict(rows=t))
for t, why in (
- (128, "decode 128 rows (max_con)"),
- (16384, "prefill chunk 16384 (max_num_tokens)"),
+ (1, "1 row: one decode token, the smallest captured batch"),
+ (4, "4 rows: the ll MTP-3 verify step (con1, tp4/tep4)"),
+ (32, "32 rows: a quarter of one tile"),
+ (128, "128 rows: one full tile, the largest captured decode batch"),
+ (1024, "1024 rows: mtp7 verify at tep4, 128 requests x 8 rows"),
+ (8192, "8192 rows: one prefill chunk at max_num_tokens"),
)
- ) + (Cell(why="ragged tail, 16320 + 70 rows (mixed step)", spec=dict(rows=16390)),)
+ ) + (Cell(why="ragged tail, 8192 + 70 rows (mixed step)", spec=dict(rows=8262)),)
note = """
- Tier 6 (CuTe DSL written for this target): `qwen4_exp_hc_up_gate_mean`'s
- kernel compiled with its NVFP4 epilogue; tiers 1-5 have no GEMM with this
- multi-input epilogue, and the fused up projection itself is tier 6.
- Metadata consumed: none. Call sites: the 48 MoE-site mixers (prefill, and
- decode past `_HC_DECODE_GEMV_MAX_ROWS` rows); `xq`/`xsf` feed
- `fp4_block_scale_moe_renorm_routed` (linear scales), `mixed` the router and
- the shared expert. Preconditions: those of `qwen4_exp_hc_up_gate_mean`
- (bf16, `R % 64 == 0`, `H % 32 == 0`, `C * 32 <= 256`, 16-byte aligned
- operands) and `global_scale` one fp32 element on the device. Compiles once
- per `(C, H, R)` on first call; graph-capturable (current stream, no host
- sync).
+ Source: the tier-6 CuTe DSL kernel of `qwen4_exp_hc_up_gate_mean` with its
+ NVFP4 epilogue; no trtllm op or FlashInfer call fuses this multi-input
+ epilogue. Metadata consumed:
+ none. Call sites: the 48 MoE-site mixers under attention TP (tp4, tep4),
+ at decode, MTP verify and prefill; `xq`/`xsf` feed
+ `fp4_block_scale_moe_renorm_routed` (linear scales), `mixed` the router
+ and the shared expert. Under attention DP the routed experts run on the
+ group's gathered bf16 rows, so this entry is not called there.
+ Preconditions: those of `qwen4_exp_hc_up_gate_mean` (bf16, `R % 64 == 0`,
+ `H % 32 == 0`, `C * 32 <= 256`, 16-byte aligned operands) and
+ `global_scale` one fp32 element on the device. Compiles once per
+ `(C, H, R)` on first call; graph-capturable (current stream, no host sync).
Numerics: the quantization is the trtllm NVFP4 recipe instruction for
instruction (`rcp.approx.ftz.f32` reciprocals, `cvt.rn.satfinite` to e4m3
and to e2m1, every multiply `mul.rn.f32`); the GPU test checks it against
`fp4_quantize`'s native-torch reference on the returned `mixed` (scales
- bitwise, e2m1 codes bitwise outside the reciprocal tie window). Bitwise
- `torch.ops.trtllm.fp4_quantize`, ties included, and across tiny and
- saturating global scales: measured, not certified.
+ bitwise, e2m1 codes bitwise outside the reciprocal tie window) and holds
+ `mixed` bitwise to the non-NVFP4 build.
Certified: bf16 at the checkpoint's 4 x 2560, rank 320, a calibrated global
- scale: 128 rows, 16384 rows, 16390 rows (a partial last tile); graph replay.
+ scale: 1, 4, 32, 128, 1024, 8192 and 8262 rows (a partial last tile); graph
+ replay.
- Not certified: zero rows, fewer rows than a tile, hc_count other than 4,
- hidden other than 2560, rank other than 320 (the kernel compiles for them;
- no cell drives them).
+ Not certified: zero rows, hc_count other than 4, hidden other than 2560,
+ rank other than 320 (the kernel compiles for them; no cell drives them).
"""
def __call__(
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/moe/__init__.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/__init__.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/moe/__init__.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/ssm/__init__.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/__init__.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/ssm/__init__.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_cached_replay_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/ssm/gdn_cached_replay_cute__verify.py
similarity index 63%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_cached_replay_cute.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/ssm/gdn_cached_replay_cute__verify.py
index ff7b5bf2f6a6..9c4ceeef43f3 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_cached_replay_cute.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/ssm/gdn_cached_replay_cute__verify.py
@@ -2,8 +2,11 @@
# SPDX-License-Identifier: Apache-2.0
"""GDN cached-replay verify step on warp tensor cores (CuTe DSL), a drop-in for the Triton
`_cached_replay_kernel` (`fused_recurrent_gated_delta_rule_cached_replay_update`) at this
-checkpoint's geometry (H=16 q/k heads, HV=48 v heads, K=V=128,
-history 16, bf16 checkpoint state, T=1..8 new tokens).
+checkpoint's geometry (K=V=128, history 16, bf16 checkpoint state, T=1..8 new tokens), over
+H q/k heads and HV = 3 H value heads: a whole layer's 16 / 48, or a tp4/tep4 rank's 4 / 12.
+
+Tier 5 (CuTe DSL written for this checkpoint); the head counts are compile-time parameters of
+`_Replay` (`h`, `hv`).
One CTA serves one (request, q/k head, split): the 3 v heads of a q/k head hold 3 x 128 = 384
state rows of 128 (K innermost), cut into 24 m-tiles of 16 rows; a CTA owns `tpc` consecutive
@@ -30,7 +33,7 @@
are issued first. Small batches (the launch is latency-bound there) bring every state tile in by
bulk copy as soon as the slot is known (`pf=0`: no tile held in registers through the prologue) and
take the tiles' state dots before the first barrier, while the history is still in flight
-(`early_acc`). `_Replay.kernel_p` is an experimental persistent variant (slower; not used).
+(`early_acc`).
"""
from __future__ import annotations
@@ -40,6 +43,7 @@
import cutlass
import cutlass.cute as cute
import torch
+import torch.nn.functional as F
from cutlass import const_expr
from cutlass._mlir.dialects import llvm
from cutlass.cute.runtime import from_dlpack
@@ -48,9 +52,13 @@
from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-H, HV, K, V, HIST = 16, 48, 128, 128, 16
-GRP = HV // H
-QKV_WIDTH = 2 * H * K + HV * V
+K, V, HIST = 128, 128, 16
+#: Value heads a q/k head serves. The q/k and value head counts themselves are compile-time
+#: parameters (`_Replay(h=..., hv=...)`, from the call's shapes): `GEOMETRIES`.
+GRP = 3
+#: `(q/k heads, value heads)` the kernel is compiled for: a whole layer (an attention-DP
+#: rank's), and a tp4/tep4 rank's quarter.
+GEOMETRIES = ((16, 48), (4, 12))
TILES = GRP * (V // 16) # m-tiles of one (request, q/k head)
PU = 272 # shared row pitch (bytes) of 128-element bf16 rows read by ldmatrix (conflict-free)
PK = 320 # shared row pitch of the history keys read as permuted 16-byte fragments
@@ -790,15 +798,17 @@ def __init__(
lo_split=True,
state_before_wait=True,
smem_tiles=True,
- persistent=False,
- gmax=0,
spec_hist=False,
qk_first=True,
use_wi=False,
evict_first=True,
early_acc=False,
+ h=16,
+ hv=48,
):
assert 1 <= T <= 8
+ assert (h, hv) in GEOMETRIES, (h, hv)
+ self.H, self.HV = h, hv
assert tpc in (24, 12, 8, 4, 2, 1) and tpc % warps == 0
self.T = T
self.TP = 4 if T <= 4 else 8
@@ -812,40 +822,33 @@ def __init__(
# as the slot is known, and no state tile is held in registers through the prologue
self.ld, self.pf, self.minb = ld, max(0 if smem_tiles else 1, min(pf, self.tpw)), minb
# the tiles' state dots before the first barrier, while the history is in flight
- self.use_early = bool(early_acc) and not persistent
+ self.use_early = bool(early_acc)
self.l2norm, self.gating, self.g_f32, self.b_f32 = l2norm, gating, g_f32, b_f32
self.early_trigger, self.lo_split = early_trigger, lo_split
self.state_before_wait = state_before_wait
# tiles past the first `pf` of a warp: bulk-copied into shared memory at the start
- self.persistent, self.gmax = persistent, gmax
- if persistent:
- assert not (g_f32 or b_f32), "the persistent variant takes bf16 raw g / beta"
- self.pf, smem_tiles = 1, True
self.nsm = self.tpw - self.pf if smem_tiles else 0
- self.spec_hist = spec_hist and not persistent
+ self.spec_hist = spec_hist
self.qk_first = qk_first
# (row, slot, pending count, half) per item from the engine's replay work items
- self.use_wi = use_wi and not persistent
+ self.use_wi = use_wi
self.evict_first = evict_first
- nset = 2 if (persistent or self.spec_hist) else 1
+ nset = 2 if self.spec_hist else 1
self.nset = nset
self.RP = (2 * self.TP + 4) * 4 # per-warp scratch row pitch (bytes)
o = 0
- tset = 2 if persistent else 1
self.o_st = o
- o += tset * warps * self.nsm * 4096
+ o += warps * self.nsm * 4096
self.o_mb = o
- o += -(-(tset * warps * self.nsm * 8) // 16) * 16
+ o += -(-(warps * self.nsm * 8) // 16) * 16
self.o_u = o
o += nset * self.nhv * 16 * PU
self.o_kh = o
o += nset * 16 * PK
- self.o_meta = o
- o += 256 * 16 if persistent else 0
self.o_kp = o
o += 16 * PU if commit else 0
self.o_v = o
- o += (2 if persistent else 1) * self.nhv * T * 256
+ o += self.nhv * T * 256
self.o_x = o
o += warps * 16 * self.RP
self.o_kk = o
@@ -902,79 +905,41 @@ def __call__(
scale: cutlass.Float32,
stream,
):
- if const_expr(self.persistent):
- n_items = n_rows * (H * self.S)
- self.kernel_p(
- sidx_t.iterator.toint(),
- q_ptr,
- k_ptr,
- v_ptr,
- q_sn,
- q_st,
- k_sn,
- k_st,
- v_sn,
- v_st,
- g_ptr,
- g_st,
- b_ptr,
- b_st,
- alog_ptr,
- dtb_ptr,
- o_ptr,
- h_ptr,
- h_ss,
- pnat_ptr,
- buf_ptr,
- ou_ptr,
- ok_ptr,
- og_ptr,
- scale,
- n_items,
- ).launch(
- grid=[cutlass.min(n_items, self.gmax), 1, 1],
- block=[32 * self.W, 1, 1],
- smem=self.smem,
- stream=stream,
- use_pdl=self.pdl,
- min_blocks_per_mp=self.minb,
- )
- else:
- self.kernel(
- sidx_t.iterator.toint(),
- q_ptr,
- k_ptr,
- v_ptr,
- q_sn,
- q_st,
- k_sn,
- k_st,
- v_sn,
- v_st,
- g_ptr,
- g_st,
- b_ptr,
- b_st,
- alog_ptr,
- dtb_ptr,
- o_ptr,
- h_ptr,
- h_ss,
- pnat_ptr,
- buf_ptr,
- ou_ptr,
- ok_ptr,
- og_ptr,
- wi_ptr,
- scale,
- ).launch(
- grid=[n_rows * (H * self.S), 1, 1],
- block=[32 * self.W, 1, 1],
- smem=self.smem,
- stream=stream,
- use_pdl=self.pdl,
- min_blocks_per_mp=self.minb,
- )
+ self.kernel(
+ sidx_t.iterator.toint(),
+ q_ptr,
+ k_ptr,
+ v_ptr,
+ q_sn,
+ q_st,
+ k_sn,
+ k_st,
+ v_sn,
+ v_st,
+ g_ptr,
+ g_st,
+ b_ptr,
+ b_st,
+ alog_ptr,
+ dtb_ptr,
+ o_ptr,
+ h_ptr,
+ h_ss,
+ pnat_ptr,
+ buf_ptr,
+ ou_ptr,
+ ok_ptr,
+ og_ptr,
+ wi_ptr,
+ scale,
+ ).launch(
+ grid=[n_rows * (self.H * self.S), 1, 1],
+ block=[32 * self.W, 1, 1],
+ smem=self.smem,
+ stream=stream,
+ use_pdl=self.pdl,
+ min_blocks_per_mp=self.minb,
+ )
# -- straight-line helpers (traced inline) ------------------------------------------------
@@ -1170,6 +1135,8 @@ def kernel(
S: cutlass.Constexpr = self.S
NTHR: cutlass.Constexpr = 32 * self.W
RP: cutlass.Constexpr = self.RP
+ H: cutlass.Constexpr = self.H
+ HV: cutlass.Constexpr = self.HV
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
@@ -1781,639 +1748,6 @@ def kernel(
if const_expr(self.pdl and not self.early_trigger):
cute.arch.griddepcontrol_launch_dependents()
- # -- the persistent variant: CTAs loop over (request, head, split) items, the next item's
- # state tiles and history in flight while the current one computes -------------------------
-
- @cute.kernel
- def kernel_p(
- self,
- sidx_ptr: cutlass.Int64,
- q_ptr: cutlass.Int64,
- k_ptr: cutlass.Int64,
- v_ptr: cutlass.Int64,
- q_sn: cutlass.Int64,
- q_st: cutlass.Int64,
- k_sn: cutlass.Int64,
- k_st: cutlass.Int64,
- v_sn: cutlass.Int64,
- v_st: cutlass.Int64,
- g_ptr: cutlass.Int64,
- g_st: cutlass.Int64,
- b_ptr: cutlass.Int64,
- b_st: cutlass.Int64,
- alog_ptr: cutlass.Int64,
- dtb_ptr: cutlass.Int64,
- o_ptr: cutlass.Int64,
- h_ptr: cutlass.Int64,
- h_ss: cutlass.Int64,
- pnat_ptr: cutlass.Int64,
- buf_ptr: cutlass.Int64,
- ou_ptr: cutlass.Int64,
- ok_ptr: cutlass.Int64,
- og_ptr: cutlass.Int64,
- scale: cutlass.Float32,
- n_items: cutlass.Int32,
- ):
- T: cutlass.Constexpr = self.T
- TP: cutlass.Constexpr = self.TP
- NT: cutlass.Constexpr = self.NT
- W: cutlass.Constexpr = self.W
- TPW: cutlass.Constexpr = self.tpw
- TPC: cutlass.Constexpr = self.tpc
- NHV: cutlass.Constexpr = self.nhv
- S: cutlass.Constexpr = self.S
- NTHR: cutlass.Constexpr = 32 * self.W
- RP: cutlass.Constexpr = self.RP
- NSM: cutlass.Constexpr = self.nsm
- SU: cutlass.Constexpr = NHV * 16 * PU
- SK: cutlass.Constexpr = 16 * PK
- SG: cutlass.Constexpr = NHV * 64
- ST: cutlass.Constexpr = W * NSM * 4096
- SMB: cutlass.Constexpr = W * NSM * 8
-
- tidx, _, _ = cute.arch.thread_idx()
- bidx, _, _ = cute.arch.block_idx()
- gsz, _, _ = cute.arch.grid_dim()
- warp = tidx // 32
- lane = tidx % 32
- g = lane // 4
- c = lane % 4
- zf = cutlass.Float32(0.0)
- zi = cutlass.Int32(0)
-
- smem = cutlass.utils.SmemAllocator()
- base_t = smem.allocate_tensor(
- cutlass.Int8, cute.make_layout((self.smem - 128,), stride=(1,)), byte_alignment=128
- )
- sbase = cutlass.Int32(base_t.iterator.toint())
- s_u0 = sbase + self.o_u
- s_kh0 = sbase + self.o_kh
- s_gh0 = sbase + self.o_gh
- s_meta = sbase + self.o_meta
- s_kp = sbase + self.o_kp
- s_v = sbase + self.o_v
- s_x = sbase + self.o_x + warp * (16 * RP)
- s_kk = sbase + self.o_kk
- s_cf = sbase + self.o_cf
- s_dec = sbase + self.o_dec
- s_g = sbase + self.o_g
- s_eg = sbase + self.o_eg
- s_beta = sbase + self.o_beta
- s_a = sbase + self.o_a
- s_cq = sbase + self.o_cq
- s_egs = sbase + self.o_egs
- s_st0 = sbase + self.o_st + warp * (NSM * 4096)
- s_mb0 = sbase + self.o_mb + warp * (NSM * 8)
-
- nk = (n_items - bidx + gsz - 1) // gsz
- # ---- the items' slot / pending count / half, once --------------------------------------
- for m_t in cutlass.range(tidx, nk, NTHR):
- m_it = bidx + m_t * gsz
- m_n = m_it // (H * S)
- m_sl = _ldg_i32(sidx_ptr + cutlass.Int64(m_n) * 4)
- m_p = zi
- m_b = zi
- if m_sl >= 0:
- m_p = _ldg_i32(pnat_ptr + cutlass.Int64(m_sl) * 4)
- m_b = _ldg_i32(buf_ptr + cutlass.Int64(m_sl) * 4)
- _sts_i32(s_meta + m_t * 16, m_sl)
- _sts_i32(s_meta + m_t * 16 + 4, m_p)
- _sts_i32(s_meta + m_t * 16 + 8, m_b)
- if const_expr(NSM > 0):
- if lane == 0:
- for k_ in cutlass.range_constexpr(2 * NSM):
- _mbar_init(s_mb0 + (k_ // NSM) * SMB + (k_ % NSM) * 8, cutlass.Int32(1))
- _fence_mbar_init()
- cute.arch.sync_threads()
-
- cur0 = cute.make_rmem_tensor(cute.make_layout((32,), stride=(1,)), cutlass.Int32)
- nxt0 = cute.make_rmem_tensor(cute.make_layout((32,), stride=(1,)), cutlass.Int32)
- NQ: cutlass.Constexpr = 16 * NT
- GN: cutlass.Constexpr = 2 * T
- qkc = cute.make_rmem_tensor(cute.make_layout((NQ,), stride=(1,)), cutlass.Int32)
- qkn = cute.make_rmem_tensor(cute.make_layout((NQ,), stride=(1,)), cutlass.Int32)
- gtc = cute.make_rmem_tensor(cute.make_layout((GN,), stride=(1,)), cutlass.Int32)
- gtn = cute.make_rmem_tensor(cute.make_layout((GN,), stride=(1,)), cutlass.Int32)
- gfc = cute.make_rmem_tensor(cute.make_layout((2,), stride=(1,)), cutlass.Float32)
- gfn = cute.make_rmem_tensor(cute.make_layout((2,), stride=(1,)), cutlass.Float32)
- SVB: cutlass.Constexpr = NHV * T * 256
- jrow = (lane % 8) + 8 * (lane // 16)
- v8 = 8 * ((lane // 8) % 2)
- r = lane % 16
- hf = lane // 16
- godd = (g % 2) == 1
- gi = tidx - (32 * (W - 1) if W > 1 else 0)
- gpred = (gi >= 0) & (gi < NHV)
- gic = cutlass.min(cutlass.max(gi, zi), NHV - 1)
-
- for k in cutlass.range(nk):
- # ---- issue: item 0 (first pass only), then item k + 1; everything an item reads
- # from global memory is in flight one item ahead ----------------------------------
- for q2 in cutlass.range_constexpr(2):
- kk = k + q2
- if const_expr(q2 == 0):
- do_it = k == 0
- else:
- do_it = kk < nk
- if do_it:
- st_ = kk % 2
- it_ = bidx + kk * gsz
- sp_ = it_ % S
- nh_ = it_ // S
- h_ = nh_ % H
- n64_ = cutlass.Int64(nh_ // H)
- sl = _lds_i32(s_meta + kk * 16)
- pp = _lds_i32(s_meta + kk * 16 + 4)
- bb = _lds_i32(s_meta + kk * 16 + 8)
- hvf = (sp_ * TPC) // 8
- hv0_ = h_ * GRP + hvf
- if const_expr(q2 == 0):
- if const_expr(self.pdl):
- cute.arch.griddepcontrol_wait()
- if const_expr(self.early_trigger):
- cute.arch.griddepcontrol_launch_dependents()
- # the item's own inputs first (they are needed first), then its history, then its
- # state: a lane's requests come back roughly in order
- if sl >= 0:
- sl64 = cutlass.Int64(sl)
- n_v: cutlass.Constexpr = NHV * T * 16
- for it2 in cutlass.range_constexpr(-(-n_v // NTHR)):
- e = tidx + it2 * NTHR
- if e < n_v:
- hl = e // (T * 16)
- t = (e // 16) % T
- qq = e % 16
- src = (
- v_ptr
- + (
- n64_ * v_sn
- + cutlass.Int64(t) * v_st
- + cutlass.Int64((hv0_ + hl) * V + qq * 8)
- )
- * 2
- )
- _cp16z(
- s_v + st_ * SVB + (hl * T + t) * 256 + qq * 16,
- src,
- cutlass.Int32(16),
- )
- for nt in cutlass.range_constexpr(NT):
- rq = self._qk_raw(
- nt,
- q_ptr,
- k_ptr,
- q_sn,
- q_st,
- k_sn,
- k_st,
- n64_,
- h_,
- g,
- c,
- scale,
- tidx >= 0,
- )
- for e in cutlass.range_constexpr(16):
- if const_expr(q2 == 0):
- qkc[nt * 16 + e] = rq[0][e]
- else:
- qkn[nt * 16 + e] = rq[0][e]
- gr = self._gate_raw(
- gpred, hv0_ + gic, n64_, g_ptr, g_st, b_ptr, b_st, alog_ptr, dtb_ptr
- )
- for t in cutlass.range_constexpr(T):
- if const_expr(q2 == 0):
- gtc[t] = gr[0][t]
- gtc[T + t] = gr[1][t]
- else:
- gtn[t] = gr[0][t]
- gtn[T + t] = gr[1][t]
- if const_expr(q2 == 0):
- gfc[0] = gr[2]
- gfc[1] = gr[3]
- else:
- gfn[0] = gr[2]
- gfn[1] = gr[3]
- rsl = sl64 * 2 + cutlass.Int64(bb)
- u_rd_ = ou_ptr + rsl * (HIST * HV * V * 2)
- k_rd_ = ok_ptr + rsl * (HIST * H * K * 2)
- g_rd_ = og_ptr + rsl * (HV * HIST * 4)
- n_u: cutlass.Constexpr = NHV * 256
- for it2 in cutlass.range_constexpr(-(-n_u // NTHR)):
- e = tidx + it2 * NTHR
- if e < n_u:
- hl = e // 256
- j = (e // 16) % 16
- qq = e % 16
- jc = cutlass.min(j, cutlass.max(pp - 1, zi))
- src = u_rd_ + cutlass.Int64(
- ((jc * HV + hv0_ + hl) * V + qq * 8) * 2
- )
- _cp16z(
- s_u0 + st_ * SU + hl * (16 * PU) + j * PU + qq * 16,
- src,
- cutlass.Int32(cutlass.select_(j < pp, cutlass.Int32(16), zi)),
- )
- for it2 in cutlass.range_constexpr(-(-256 // NTHR)):
- e = tidx + it2 * NTHR
- if e < 256:
- j = e // 16
- qq = e % 16
- jc = cutlass.min(j, cutlass.max(pp - 1, zi))
- src = k_rd_ + cutlass.Int64(((jc * H + h_) * K + qq * 8) * 2)
- _cp16z(
- s_kh0 + st_ * SK + j * PK + qq * 16,
- src,
- cutlass.Int32(cutlass.select_(j < pp, cutlass.Int32(16), zi)),
- )
- if gpred:
- # a v head's G row by the thread that gates that head (read back by it)
- for qq in cutlass.range_constexpr(4):
- src = g_rd_ + cutlass.Int64(((hv0_ + gi) * HIST + qq * 4) * 4)
- _cp16z(s_gh0 + st_ * SG + gi * 64 + qq * 16, src, cutlass.Int32(16))
- vals = self._load_tile(
- self._tile_addr(h_ptr, sl64, h_ss, h_, sp_, warp, 0, g, c)
- )
- for e in cutlass.range_constexpr(32):
- if const_expr(q2 == 0):
- cur0[e] = vals[e]
- else:
- nxt0[e] = vals[e]
- if const_expr(NSM > 0):
- if lane == 0:
- for k_ in cutlass.range_constexpr(NSM):
- _bulk_g2s(
- s_st0 + st_ * ST + k_ * 4096,
- self._tile_addr(
- h_ptr, sl64, h_ss, h_, sp_, warp, 1 + k_, 0, 0
- ),
- cutlass.Int32(4096),
- s_mb0 + st_ * SMB + k_ * 8,
- )
- else:
- if const_expr(NSM > 0):
- if lane == 0:
- for k_ in cutlass.range_constexpr(NSM):
- _mbar_arrive(s_mb0 + st_ * SMB + k_ * 8)
- _cp_commit() # item 0 (empty later) / item k + 1
-
- # ---- item k ---------------------------------------------------------------------
- st = k % 2
- it = bidx + k * gsz
- sp = it % S
- nh = it // S
- h = nh % H
- n = nh // H
- n64 = cutlass.Int64(n)
- slot = _lds_i32(s_meta + k * 16)
- p = _lds_i32(s_meta + k * 16 + 4)
- rb = _lds_i32(s_meta + k * 16 + 8)
- s_u = s_u0 + st * SU
- s_kh = s_kh0 + st * SK
- s_gh = s_gh0 + st * SG
- s_st = s_st0 + st * ST
- s_mb = s_mb0 + st * SMB
- s_vk = s_v + st * SVB
- ph = (k // 2) % 2
- slot64 = cutlass.Int64(cutlass.max(slot, zi))
- wrap = (p + T) > HIST
- wb = cutlass.Int32(cutlass.select_(wrap, 1 - rb, rb))
- woff = cutlass.Int32(cutlass.select_(wrap, zi, p))
- hv_first = (sp * TPC) // 8
- hv0 = h * GRP + hv_first
- wslot = slot64 * 2 + cutlass.Int64(wb)
- u_wr = ou_ptr + wslot * (HIST * HV * V * 2)
- k_wr = ok_ptr + wslot * (HIST * H * K * 2)
- g_wr = og_ptr + wslot * (HV * HIST * 4)
-
- live = slot >= 0
- if live:
- bq = []
- for nt in cutlass.range_constexpr(NT):
- rw = [qkc[nt * 16 + e] for e in range(16)]
- if const_expr(TP == 4):
- is_q = g >= 4
- mult = cutlass.Float32(cutlass.select_(is_q, scale, cutlass.Float32(1.0)))
- elif const_expr(nt == 1):
- is_q = g >= 0
- mult = scale
- else:
- is_q = g < 0
- mult = cutlass.Float32(1.0)
- bq.append(self._qk_frag(rw, is_q, mult))
-
- # the new keys into the history (one CTA of the head, warp 0)
- if sp == 0:
- if warp == 0:
- if g < T:
- kfr = bq[0]
- dst = k_wr + cutlass.Int64((((woff + g) * H + h) * K + c * 8) * 2)
- for jj in cutlass.range_constexpr(4):
- _stg4(
- dst + jj * 64,
- kfr[4 * jj],
- kfr[4 * jj + 1],
- kfr[4 * jj + 2],
- kfr[4 * jj + 3],
- )
-
- # new-token dot products: rows (keys | queries) x key columns
- if warp == 0:
- d = [zf, zf, zf, zf]
- for jj in cutlass.range_constexpr(4):
- for s in cutlass.range_constexpr(2):
- b0 = bq[0][4 * jj + 2 * s]
- b1 = bq[0][4 * jj + 2 * s + 1]
- if const_expr(NT == 1):
- a = (b0, zi, b1, zi)
- else:
- a = (b0, bq[1][4 * jj + 2 * s], b1, bq[1][4 * jj + 2 * s + 1])
- d = _mma(a, (b0, b1), d)
- _sts2_f32(s_kk + (g * 8 + 2 * c) * 4, d[0], d[1])
- _sts2_f32(s_kk + ((g + 8) * 8 + 2 * c) * 4, d[2], d[3])
-
- _cp_wait1()
- # gating, cumulative decays, the G history (a thread per v head, off warp 0); the
- # thread's own G row copy has landed
- if gpred:
- hl = gi
- hv = hv0 + hl
- gs = cutlass.Float32(
- cutlass.select_(
- p > 0, _lds_f32(s_gh + (hl * 16 + cutlass.max(p - 1, zi)) * 4), zf
- )
- )
- graws = [_b16_f32(gtc[t]) for t in range(T)]
- braws = [_b16_f32(gtc[T + t]) for t in range(T)]
- alv = gfc[0]
- dtb = gfc[1]
- aexp = _exp(alv)
- owns = (hv_first + hl) * 8 >= sp * TPC
- gl = zf
- for t in cutlass.range_constexpr(T):
- graw = graws[t]
- braw = braws[t]
- if const_expr(self.gating):
- x = graw + dtb
- spl = cutlass.Float32(
- cutlass.select_(x <= 20.0, LN2 * _lg2(1.0 + _ex2(x * LOG2E)), x)
- )
- gt = -aexp * spl
- beta = _div_approx(cutlass.Float32(1.0), 1.0 + _ex2(-braw * LOG2E))
- else:
- gt = graw
- beta = braw
- gl = gl + gt
- gabs = gs + gl
- _sts_f32(s_g + (hl * 8 + t) * 4, gabs)
- _sts_f32(s_eg + (hl * 8 + t) * 4, _exp(gabs))
- _sts_f32(s_beta + (hl * 8 + t) * 4, beta)
- if owns:
- _stg_f32(
- g_wr + cutlass.Int64((hv * HIST + woff + t) * 4),
- cutlass.Float32(cutlass.select_(wrap, gl, gabs)),
- )
- _sts_f32(s_egs + hl * 4, _exp(gs))
- cute.arch.sync_threads()
-
- hk_warp: cutlass.Constexpr = 1 if W > 1 else 0
- if warp == hk_warp:
- dh = [[zf, zf, zf, zf] for _ in range(NT)]
- for jj in cutlass.range_constexpr(4):
- rg = _lds4_i32(s_kh + g * PK + jj * 64 + c * 16)
- rh = _lds4_i32(s_kh + (g + 8) * PK + jj * 64 + c * 16)
- for s in cutlass.range_constexpr(2):
- a = (rg[2 * s], rh[2 * s], rg[2 * s + 1], rh[2 * s + 1])
- for nt in cutlass.range_constexpr(NT):
- dh[nt] = _mma(
- a, (bq[nt][4 * jj + 2 * s], bq[nt][4 * jj + 2 * s + 1]), dh[nt]
- )
- for hl in cutlass.range_constexpr(NHV):
- gh_g = _lds_f32(s_gh + (hl * 16 + g) * 4)
- gh_h = _lds_f32(s_gh + (hl * 16 + g + 8) * 4)
- for nt in cutlass.range_constexpr(NT):
- for e in cutlass.range_constexpr(2):
- col = nt * 8 + 2 * c + e
- if const_expr(TP == 4):
- tt = (2 * c + e) % 4
- else:
- tt = 2 * c + e
- tv = tt < T
- gtt = _lds_f32(s_g + (hl * 8 + tt) * 4)
- vg = cutlass.Float32(
- cutlass.select_((g < p) & tv, dh[nt][e] * _exp(gtt - gh_g), zf)
- )
- vh = cutlass.Float32(
- cutlass.select_(
- ((g + 8) < p) & tv, dh[nt][2 + e] * _exp(gtt - gh_h), zf
- )
- )
- _sts_bf16(s_cf + ((hl * 16 + col) * 16 + g) * 2, vg)
- _sts_bf16(s_cf + ((hl * 16 + col) * 16 + g + 8) * 2, vh)
-
- if const_expr(self.commit):
- if wrap:
- for it2 in cutlass.range_constexpr(-(-256 // NTHR)):
- e = tidx + it2 * NTHR
- if e < 256:
- j = e // 16
- jj = (e // 4) % 4
- aa = e % 4
- w4 = _lds4_i32(s_kh + j * PK + jj * 64 + aa * 16)
- for s4 in cutlass.range_constexpr(4):
- _sts_i32(s_kp + j * PU + (16 * jj + 4 * s4 + aa) * 4, w4[s4])
-
- n_ts: cutlass.Constexpr = NHV * T * T
- for it2 in cutlass.range_constexpr(-(-n_ts // NTHR)):
- e = tidx + it2 * NTHR
- if e < n_ts:
- hl = e // (T * T)
- t = (e // T) % T
- s = e % T
- gt = _lds_f32(s_g + (hl * 8 + t) * 4)
- gsv = _lds_f32(s_g + (hl * 8 + s) * 4)
- bt = _lds_f32(s_beta + (hl * 8 + t) * 4)
- dec = _exp(gt - gsv)
- kkv = _lds_f32(s_kk + (t * 8 + s) * 4)
- qkv = _lds_f32(s_kk + ((TP + t) * 8 + s) * 4)
- av = cutlass.Float32(cutlass.select_(s < t, bt * kkv * dec, zf))
- cv = cutlass.Float32(cutlass.select_(s <= t, qkv * dec, zf))
- _sts_f32(s_a + (hl * 64 + t * 8 + s) * 4, av)
- _sts_f32(s_cq + (hl * 64 + t * 8 + s) * 4, cv)
- if const_expr(self.commit):
- if tidx < NHV * 16:
- hl = tidx // 16
- j = tidx % 16
- gsv = cutlass.Float32(0.0)
- if p > 0:
- gsv = _lds_f32(s_gh + (hl * 16 + p - 1) * 4)
- ghj = _lds_f32(s_gh + (hl * 16 + j) * 4)
- _sts_f32(
- s_dec + (hl * 16 + j) * 4,
- cutlass.Float32(cutlass.select_(j < p, _exp(gsv - ghj), zf)),
- )
- cute.arch.sync_threads()
-
- for i in cutlass.range_constexpr(TPW):
- if const_expr(i == 0):
- cur = [cur0[e] for e in range(32)]
- else:
- kslot = i - 1
- _mbar_wait(s_mb + kslot * 8, ph)
- tb = s_st + kslot * 4096 + g * 256 + c * 16
- raw = []
- for ii in cutlass.range_constexpr(4):
- jsw = (g % 2) ^ ii
- raw.append(_lds4_i32(tb + jsw * 64))
- raw.append(_lds4_i32(tb + 2048 + jsw * 64))
- cur = []
- for jj in cutlass.range_constexpr(4):
- for rr in cutlass.range_constexpr(2):
- for e in cutlass.range_constexpr(4):
- cur.append(
- cutlass.Int32(
- cutlass.select_(
- godd, raw[2 * (jj ^ 1) + rr][e], raw[2 * jj + rr][e]
- )
- )
- )
- mt = sp * TPC + warp + W * i
- hvr = mt // 8
- hv = h * GRP + hvr
- hl = hvr - hv_first
- v0 = (mt % 8) * 16
-
- acc0 = [[zf, zf, zf, zf] for _ in range(NT)]
- acc1 = [[zf, zf, zf, zf] for _ in range(NT)]
- for jj in cutlass.range_constexpr(4):
- for s in cutlass.range_constexpr(2):
- a = (
- cur[8 * jj + 2 * s],
- cur[8 * jj + 4 + 2 * s],
- cur[8 * jj + 2 * s + 1],
- cur[8 * jj + 4 + 2 * s + 1],
- )
- for nt in cutlass.range_constexpr(NT):
- bfr = (bq[nt][4 * jj + 2 * s], bq[nt][4 * jj + 2 * s + 1])
- if const_expr(s == 0):
- acc0[nt] = _mma(a, bfr, acc0[nt])
- else:
- acc1[nt] = _mma(a, bfr, acc1[nt])
- acc = [[acc0[nt][e] + acc1[nt][e] for e in range(4)] for nt in range(NT)]
-
- au = _ldsm4t(s_u + hl * (16 * PU) + jrow * PU + (v0 + v8) * 2)
- hacc = []
- for nt in cutlass.range_constexpr(NT):
- cfa = s_cf + ((hl * 16 + nt * 8 + g) * 16 + 2 * c) * 2
- bc = (_lds_i32(cfa), _lds_i32(cfa + 16))
- hacc.append(_mma(au, bc, [zf, zf, zf, zf]))
-
- for nt in cutlass.range_constexpr(NT):
- col0 = nt * 8 + 2 * c
- if const_expr(TP == 4):
- t0 = (2 * c) % 4
- else:
- t0 = 2 * c
- eg0 = _lds_f32(s_eg + (hl * 8 + t0) * 4)
- eg1 = _lds_f32(s_eg + (hl * 8 + t0 + 1) * 4)
- _sts2_f32(
- s_x + g * RP + col0 * 4,
- eg0 * acc[nt][0] + hacc[nt][0],
- eg1 * acc[nt][1] + hacc[nt][1],
- )
- _sts2_f32(
- s_x + (g + 8) * RP + col0 * 4,
- eg0 * acc[nt][2] + hacc[nt][2],
- eg1 * acc[nt][3] + hacc[nt][3],
- )
- cute.arch.sync_warp()
-
- xs = []
- ys = []
- for q4 in cutlass.range_constexpr(TP // 4):
- xs.extend(_lds4_f32(s_x + r * RP + q4 * 16))
- ys.extend(_lds4_f32(s_x + r * RP + (TP + 4 * q4) * 4))
- us = []
- for t in cutlass.range_constexpr(T):
- vt = _lds_bf16(s_vk + ((hl * T + t) * 128 + v0 + r) * 2)
- bt = _lds_f32(s_beta + (hl * 8 + t) * 4)
- u = bt * (vt - xs[t])
- for s in cutlass.range_constexpr(t):
- u = u - _lds_f32(s_a + (hl * 64 + t * 8 + s) * 4) * us[s]
- us.append(u)
- if hf == 0:
- for t in cutlass.range_constexpr(T):
- o = ys[t]
- for s in cutlass.range_constexpr(t + 1):
- o = o + _lds_f32(s_cq + (hl * 64 + t * 8 + s) * 4) * us[s]
- _stg_bf16(o_ptr + (((n64 * T + t) * HV + hv) * V + v0 + r) * 2, o)
- else:
- for t in cutlass.range_constexpr(T):
- _stg_bf16(
- u_wr + cutlass.Int64((((woff + t) * HV + hv) * V + v0 + r) * 2),
- us[t],
- )
- cute.arch.sync_warp()
-
- if const_expr(self.commit):
- if wrap:
- dc = [
- _lds_f32(s_dec + (hl * 16 + 2 * c + jx) * 4) for jx in [0, 1, 8, 9]
- ]
- egs = _lds_f32(s_egs + hl * 4)
- a_hi = []
- a_lo = []
- for q4 in cutlass.range_constexpr(4):
- dl = dc[0] if q4 < 2 else dc[2]
- dhh = dc[1] if q4 < 2 else dc[3]
- x0 = _bf2_lo(au[q4]) * dl
- x1 = _bf2_hi(au[q4]) * dhh
- hw = _bf2_pack(x0, x1)
- a_hi.append(hw)
- a_lo.append(_bf2_pack(x0 - _bf2_lo(hw), x1 - _bf2_hi(hw)))
- taddr = self._tile_addr(h_ptr, slot64, h_ss, h, sp, warp, i, g, c)
- kp_row = (
- s_kp + ((lane % 8) + 8 * ((lane // 8) % 2)) * PU + 16 * (lane // 16)
- )
- for jj in cutlass.range_constexpr(4):
- bb0 = _ldsm4t(kp_row + (4 * jj) * 16)
- bb1 = _ldsm4t(kp_row + (4 * jj + 2) * 16)
- ng = []
- nh8 = []
- for s4 in cutlass.range_constexpr(4):
- bsrc = bb0 if s4 < 2 else bb1
- bf = (bsrc[2 * (s4 % 2)], bsrc[2 * (s4 % 2) + 1])
- dcm = _mma(a_hi, bf, [zf, zf, zf, zf])
- if const_expr(self.lo_split):
- dcm = _mma(a_lo, bf, dcm)
- wg = cur[8 * jj + s4]
- wh = cur[8 * jj + 4 + s4]
- ng.append(
- _bf2_pack(
- egs * _bf2_lo(wg) + dcm[0], egs * _bf2_hi(wg) + dcm[1]
- )
- )
- nh8.append(
- _bf2_pack(
- egs * _bf2_lo(wh) + dcm[2], egs * _bf2_hi(wh) + dcm[3]
- )
- )
- _stg4(taddr + jj * 64, ng[0], ng[1], ng[2], ng[3])
- _stg4(taddr + (2048 + jj * 64), nh8[0], nh8[1], nh8[2], nh8[3])
-
- for ce in cutlass.range_constexpr(32):
- cur0[ce] = nxt0[ce]
- for ce in cutlass.range_constexpr(NQ):
- qkc[ce] = qkn[ce]
- for ce in cutlass.range_constexpr(GN):
- gtc[ce] = gtn[ce]
- for ce in cutlass.range_constexpr(2):
- gfc[ce] = gfn[ce]
- cute.arch.sync_threads()
-
- if const_expr(self.pdl and not self.early_trigger):
- cute.arch.griddepcontrol_launch_dependents()
-
# ---------------------------------------------------------------------------------------
# host side
@@ -2422,49 +1756,12 @@ def kernel_p(
_COMPILED: dict = {}
-def parse_config(text: str) -> Optional[dict]:
- """'default', positional 'tpc,warps,pf,minb[,smem_tiles[,persistent[,occ]]]' or 'key=value,...'."""
- if text in ("", "default"):
- return None
- if "=" in text:
- out = {}
- for kv in text.split(","):
- k_, v_ = kv.split("=")
- if v_.lstrip("-").isdigit():
- out[k_] = int(v_)
- elif v_ in ("True", "False"):
- out[k_] = v_ == "True"
- else:
- out[k_] = v_
- for b_ in (
- "smem_tiles",
- "persistent",
- "spec_hist",
- "qk_first",
- "lo_split",
- "early_trigger",
- "state_before_wait",
- "early_acc",
- ):
- if b_ in out:
- out[b_] = bool(out[b_])
- return out
- f = [int(x) for x in text.split(",")]
- names = ("tpc", "warps", "pf", "minb", "smem_tiles", "persistent", "occ")
- out = dict(zip(names, f))
- for b_ in ("smem_tiles", "persistent"):
- if b_ in out:
- out[b_] = bool(out[b_])
- return out
-
-
def default_config(n: int, t: int, commit: bool = False) -> dict:
- """Launch mapping by request count, measured on GB300 at T = 4 (36 cold layers): small
- batches split a head's 24 m-tiles over more, smaller CTAs, every state tile by bulk copy
- issued as soon as the slot is known and the state dots taken before the history lands (a
- latency-bound launch: 1 request 3.3 -> 2.8 us); from 32 requests (or with the in-kernel
- commit from 16) CTAs of 8 m-tiles x 4 warps, capped at 5 (4 with the commit) resident per
- SM's register file, keep the most state in flight."""
+ """Launch mapping by request count: small batches split a head's 24 m-tiles over more,
+ smaller CTAs, every state tile by bulk copy issued as soon as the slot is known and the
+ state dots taken before the history lands (a latency-bound launch); from 32 requests (or
+ with the in-kernel commit from 16) CTAs of 8 m-tiles x 4 warps, capped at 5 (4 with the
+ commit) resident per SM's register file, keep the most state in flight."""
if n <= 2:
return dict(tpc=2, warps=2, pf=0, minb=4, early_acc=True)
if n <= 8:
@@ -2536,8 +1833,12 @@ def cached_replay_update(
The Triton launch knobs (`block_v`, `num_warps`, `use_l2_*`, `main_num_stages`) are ignored.
"""
del old_beta
- N, T, Hq, Kd = k.shape
- assert (Hq, Kd) == (H, K) and v.shape[2:] == (HV, V) and history_size == HIST, "geometry"
+ N, T, H, Kd = k.shape
+ HV = v.shape[2]
+ assert (H, HV) in GEOMETRIES and Kd == K and v.shape[3] == V and history_size == HIST, (
+ f"geometry: {H} q/k x {Kd}, {HV} v x {v.shape[3]}, history {history_size}"
+ )
+ qkv_width = 2 * H * K + HV * V
assert 1 <= T <= 8
assert ssm_states.dtype == torch.bfloat16, "bf16 checkpoint state only"
assert ssm_states.stride()[1:] == (V * K, K, 1)
@@ -2567,11 +1868,11 @@ def cached_replay_update(
if N == 0:
return output
if packed_qkv is not None:
- assert packed_qkv.shape == (N * T, QKV_WIDTH) and packed_qkv.is_contiguous()
+ assert packed_qkv.shape == (N * T, qkv_width) and packed_qkv.is_contiguous()
base = packed_qkv.data_ptr()
es = packed_qkv.element_size()
qp, kp, vp = base, base + H * K * es, base + 2 * H * K * es
- qs = ks = vs = (T * QKV_WIDTH, QKV_WIDTH)
+ qs = ks = vs = (T * qkv_width, qkv_width)
else:
for x, nh, d in ((q, H, K), (k, H, K), (v, HV, V)):
assert x.stride(3) == 1 and x.stride(2) == d, "head rows must be dense"
@@ -2591,23 +1892,13 @@ def cached_replay_update(
# with PDL, issue the state / history (not written by the previous kernel) before the wait;
# without it, the step's own small inputs first
cfg.setdefault("qk_first", not bool(launch_with_pdl))
- occ = cfg.pop("occ", 1)
- use_wi = (
- bool(cfg.pop("use_wi", True))
- and replay_work_items is not None
- and not cfg.get("persistent")
- )
+ use_wi = bool(cfg.pop("use_wi", True)) and replay_work_items is not None
cfg["use_wi"] = use_wi
- if cfg.get("persistent"):
- sms = torch.cuda.get_device_properties(q.device).multi_processor_count
- cfg["gmax"] = occ * sms
- items = N * H * (TILES // cfg["tpc"])
- assert items <= 256 * min(items, cfg["gmax"]), "more than 256 items per CTA"
- else:
- cfg.pop("gmax", None)
pdl = bool(launch_with_pdl)
kw = dict(
T=T,
+ h=H,
+ hv=HV,
commit=commit,
pdl=pdl,
l2norm=bool(use_qk_l2norm_in_kernel),
@@ -2653,65 +1944,399 @@ def cached_replay_update(
return output
+#: A tp4/tep4 rank's GDN heads: 16 key / 48 value split four ways.
+RANK_K_HEADS, RANK_V_HEADS = 4, 12
+#: The stock tep4 engine's recurrent-state slots a rank (bf16 state), and the GDN layers whose
+#: states share one slot page in the stock manager: the slot stride is 36 layers' states.
+TEP4_SLOTS, PAGE_LAYERS = 1874, 36
+
+
+def _tep4_cell(why, path, n, t, **spec) -> Cell:
+ """A tp4/tep4 rank's heads over the stock tep4 engine's bf16 state page."""
+ return Cell(
+ why=f"tep4 {why}",
+ spec=dict(
+ path=path,
+ n=n,
+ t=t,
+ state_dtype=torch.bfloat16,
+ H=RANK_K_HEADS,
+ HV=RANK_V_HEADS,
+ history_size=HIST,
+ slots=TEP4_SLOTS,
+ page_layers=PAGE_LAYERS,
+ **spec,
+ ),
+ )
+
+
class GdnCachedReplayCute(OpWrapper):
"""`cached_replay_update` as a catalog entry: one verify step of `N` sequences x `T` tokens
- through one GDN layer under the cached-replay state layout, bf16 checkpoint state only,
- with the same contract as `fused_recurrent_gated_delta_rule_cached_replay_update` (see that
- entry's docstring: history rebuild, gated delta rule over the T new tokens, history append,
- and -- `in_kernel_commit` -- the checkpoint advance of a wrapping row written in place).
+ through one GDN layer under the cached-replay state layout, bf16 checkpoint state only, with
+ the contract of `fused_recurrent_gated_delta_rule_cached_replay_update` (history rebuild,
+ gated delta rule over the T new tokens, history append, and the checkpoint advance of a
+ wrapping row written in place when no work items are given -- or, `in_kernel_commit=True`,
+ always). With work items and `use_all_layer_commit` it commits nothing (the manager's
+ all-layer commit does), exactly as the Triton entry.
- With `in_kernel_commit=True` every wrapping row commits here at any `N`: the caller must
- not also run the manager's all-layer commit (the checkpoint would advance twice).
+ The q/k and value head counts are compile-time parameters taken from the call's shapes
+ (`GEOMETRIES`: a whole layer's 16 / 48 and a tp4/tep4 rank's 4 / 12), and `g` / `beta` may
+ be strided column views of the input projection (any row stride, unit head stride).
Fusion boundary: as the Triton entry's.
"""
ARCHS = frozenset({Arch.SM_103})
- # spec: the Triton entry's test harness (`_Bench`) spec -- path, n sequences x t tokens,
- # consecutive steps with random acceptance, pool slots, layers -- run in the commit mode
- # the target uses at that size: in-kernel on the small path, the manager's all-layer commit
- # on the partitioned one.
+ # spec: the Triton entry's test harness (`_Bench`) spec -- path, n sequences x t
+ # tokens, consecutive steps with random acceptance, start, accept, dup, H / HV, pool slots,
+ # page_layers (the stock manager's coalesced state page) -- plus `strided_gates` (g / beta
+ # handed over as column views of a `[N*T, conv + z + 2 HV]` projection, as the target does).
CELLS: tuple[Cell, ...] = (
- Cell(
- why="MTP3 verify at con1: N=1, T=4, in-kernel commit",
- spec=dict(path="small", n=1, t=4, steps=6),
+ _tep4_cell(
+ "small n=1 T=4: MTP3 verify at con1, in-kernel commit on the page's last slot",
+ "small",
+ 1,
+ 4,
+ steps=8,
+ strided_gates=True,
+ ),
+ _tep4_cell(
+ "small n=4 T=4: the largest captured ll batch",
+ "small",
+ 4,
+ 4,
+ steps=6,
+ strided_gates=True,
+ ),
+ _tep4_cell(
+ "small n=4 T=4 dup: rows 2..3 share one CUDA-graph padding slot",
+ "small",
+ 4,
+ 4,
+ steps=6,
+ dup=(2, 3),
+ strided_gates=True,
),
- Cell(
- why="MTP3 verify at con32: N=32, T=4, work items + the manager's all-layer commit",
- spec=dict(path="all_layer", n=32, t=4, steps=4, slots=72),
+ _tep4_cell("small n=2 T=2: MTP1 verify", "small", 2, 2, steps=8, strided_gates=True),
+ _tep4_cell(
+ "small n=1 T=8 full accept: MTP7 appends to the history's end, then wraps",
+ "small",
+ 1,
+ 8,
+ steps=8,
+ start="fresh",
+ accept="full",
+ strided_gates=True,
),
- Cell(
- why="MTP3 verify at con256: N=256, T=4, work items + the manager's all-layer commit",
- spec=dict(path="all_layer", n=256, t=4, steps=3, slots=300, layers=1),
+ _tep4_cell(
+ "all_layer n=16 T=4: the smallest partitioned batch, all-layer commit",
+ "all_layer",
+ 16,
+ 4,
+ steps=6,
+ strided_gates=True,
),
+ _tep4_cell("all_layer n=128 T=4: a large partitioned batch", "all_layer", 128, 4, steps=4),
)
note = """
- Metadata consumed: as the Triton entry (slots, per-slot pending count and history half),
- or the engine's replay work items (row, slot, pending count, half) over every row in one
- load each. Call site: `_replay_verify` over a bf16 recurrent pool, under the Triton
- entry's contract: no work items below `CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE` requests
- (wrapping rows commit here), work items and the manager's all-layer commit from there on.
- `g` / `beta` may be the in_proj's column views (any row stride). Compiles one kernel per
- (T, launch mapping, commit, PDL, ...) on first use (5-10 s): warm-up runs before CUDA-graph
- capture cover it.
- Not supported: fp32 state, T > 8, another geometry. Optimized for T = 4; T = 1..8 are
- tested for correctness. Every launch mapping is bitwise every other (the per-row
- arithmetic does not depend on it).
-
- GB300, T = 4, in-kernel commit, 25 % of the rows wrapping, PDL off, a CUDA graph of 36
- layers each streaming its own state, us a layer: 1 request 2.8 (2.76 with the work
- items), 4: 4.8 (4.4), 32: 17.0 (16.2), 128: 52.3 (51.9), 256: 100.9 (99.6).
+ Call site: `_gdn_verify` (`DecodeTarget._gdn_block` / `PrefillTarget._gdn_block` of a verify
+ step) over the bf16 recurrent pool, at a tp4/tep4 rank's 4 / 12 heads: q/k/v from the conv's
+ fresh contiguous `[N*T, 2560]` rows (`packed_qkv`), `g` / `beta` the input projection's raw
+ a / b columns read through its row stride (no copies), `A_log` / `dt_bias` fp32, the
+ layer's view of the stock manager's bf16 state page (slot stride 36 layers' states), the
+ history buffers `[slots, 2, 16, 12|4, 128]` / `[slots, 2, 12, 16]`, and work items +
+ `use_all_layer_commit` exactly when the Triton entry gets them (N >= 16 with the manager's
+ all-layer commit on). `launch_with_pdl=True`: the Triton conv before it signals its
+ dependents at its start (`launch_dependent_kernels`), so the kernel streams the checkpoint
+ and the history -- which nothing in the step writes ahead of it -- while the conv runs, and
+ reads q/k/v, a and b only after `griddepcontrol.wait`. Compiles one kernel per (geometry,
+ T, launch mapping, commit, PDL, ...) on first use (5-10 s): the warm-up before CUDA-graph
+ capture covers it.
+
+ Certified at 4 / 12 heads over a bf16 checkpoint (cells: the stock tep4 engine's 36-layer
+ state page at 1874 slots, rows on both ends of it, several consecutive verify steps from
+ warm histories with random acceptance): the in-kernel commit at N = 1, 2, 4 and T = 2, 4,
+ 8 (T = 8 with full acceptance appending to the history's end, then wrapping), CUDA-graph
+ padding rows sharing one slot, and the all-layer commit at N = 16 and 128 (T = 4); g / beta
+ as column views of a `[N*T, 4120]` projection. Against the Triton entry's bands (fractions
+ of each gate): output elements <= 0.50 of 4 bf16 ulps of their token block, relative RMS
+ <= 0.71 of 1.5 ulps, bf16 checkpoint <= 0.68 of 2 ulps of its head's absmax, `old_u` <=
+ 0.50 of 4 ulps, `old_G` <= 0.06 -- the Triton entry's own margins (0.52 / 0.64 / 0.84 /
+ 0.50 / 0.08). A reference with the key heads reversed is refused by the output gate.
+ CUDA-graph replay (PDL on and off, N = 1, 4 and 16) is bitwise eager, and the launch
+ mappings are bitwise each other at N = 1, 2, 4 and T = 2, 4, 8.
+
+ Numerics: the Triton kernel's rounding points (bf16 L2-normalized q / k and history
+ coefficients into the tensor-core dots, fp32 recurrence, fp32 commit rounded once to the
+ bf16 pool, bf16 history u / k, fp32 G); the dot accumulation order differs, so outputs
+ match the Triton kernel to a few bf16 ulps, not bitwise. Every launch mapping is bitwise
+ every other (the per-row arithmetic does not depend on it).
+
+ Not supported: fp32 state, T > 8, a geometry outside `GEOMETRIES`, K or V other than 128,
+ history other than 16, pre-gated g / beta (`A_log=None` is compiled but not certified).
+ The Triton launch knobs (`block_v`, `num_warps`, `use_l2_*`, `main_num_stages`) are ignored.
"""
- def __call__(self, *args, **kwargs) -> torch.Tensor:
- return cached_replay_update(*args, **kwargs)
+ def __call__(
+ self,
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ g: torch.Tensor,
+ beta: torch.Tensor,
+ ssm_states: torch.Tensor,
+ state_indices: torch.Tensor,
+ old_u: torch.Tensor,
+ old_k: torch.Tensor,
+ old_G: torch.Tensor,
+ old_beta: torch.Tensor,
+ cache_buf_idx: torch.Tensor,
+ prev_num_accepted_tokens: torch.Tensor,
+ history_size: int,
+ scale: Optional[float] = None,
+ use_qk_l2norm_in_kernel: bool = False,
+ A_log: Optional[torch.Tensor] = None,
+ dt_bias: Optional[torch.Tensor] = None,
+ launch_with_pdl: bool = False,
+ replay_work_items: Optional[torch.Tensor] = None,
+ n_writes: Optional[torch.Tensor] = None,
+ block_v: Optional[int] = None,
+ num_warps: Optional[int] = None,
+ use_l2_state_cache: Optional[bool] = None,
+ use_l2_streaming_inputs: Optional[bool] = None,
+ use_l2_shared_inputs: Optional[bool] = None,
+ main_num_stages: Optional[int] = None,
+ packed_qkv: Optional[torch.Tensor] = None,
+ use_all_layer_commit: bool = False,
+ output: Optional[torch.Tensor] = None,
+ in_kernel_commit: Optional[bool] = None,
+ config: Optional[dict] = None,
+ ) -> torch.Tensor:
+ return cached_replay_update(
+ q=q,
+ k=k,
+ v=v,
+ g=g,
+ beta=beta,
+ ssm_states=ssm_states,
+ state_indices=state_indices,
+ old_u=old_u,
+ old_k=old_k,
+ old_G=old_G,
+ old_beta=old_beta,
+ cache_buf_idx=cache_buf_idx,
+ prev_num_accepted_tokens=prev_num_accepted_tokens,
+ history_size=history_size,
+ scale=scale,
+ use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+ A_log=A_log,
+ dt_bias=dt_bias,
+ launch_with_pdl=launch_with_pdl,
+ replay_work_items=replay_work_items,
+ n_writes=n_writes,
+ packed_qkv=packed_qkv,
+ use_all_layer_commit=use_all_layer_commit,
+ output=output,
+ in_kernel_commit=in_kernel_commit,
+ config=config,
+ )
+
+ def reference(
+ self,
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ g: torch.Tensor,
+ beta: torch.Tensor,
+ ssm_states: torch.Tensor,
+ state_indices: torch.Tensor,
+ old_u: torch.Tensor,
+ old_k: torch.Tensor,
+ old_G: torch.Tensor,
+ old_beta: torch.Tensor,
+ cache_buf_idx: torch.Tensor,
+ prev_num_accepted_tokens: torch.Tensor,
+ history_size: int,
+ scale: Optional[float] = None,
+ use_qk_l2norm_in_kernel: bool = False,
+ A_log: Optional[torch.Tensor] = None,
+ dt_bias: Optional[torch.Tensor] = None,
+ launch_with_pdl: bool = False,
+ replay_work_items: Optional[torch.Tensor] = None,
+ n_writes: Optional[torch.Tensor] = None,
+ block_v: Optional[int] = None,
+ num_warps: Optional[int] = None,
+ use_l2_state_cache: Optional[bool] = None,
+ use_l2_streaming_inputs: Optional[bool] = None,
+ use_l2_shared_inputs: Optional[bool] = None,
+ main_num_stages: Optional[int] = None,
+ packed_qkv: Optional[torch.Tensor] = None,
+ use_all_layer_commit: bool = False,
+ output: Optional[torch.Tensor] = None,
+ in_kernel_commit: Optional[bool] = None,
+ config: Optional[dict] = None,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ """`(out, ssm_states, old_u, old_k, old_G)` as the call leaves them;
+ mutates nothing. The state is rebuilt from the checkpoint and history
+ in fp32 and the new tokens are stepped one at a time in fp32 -- no
+ matmul, so no TF32 -- and every store is rounded once to its buffer's
+ dtype. The history is read at whatever dtype and depth its buffers
+ have, so the same function driven on fp32 history deeper than any run
+ ever fills is the exact token-by-token recurrence.
+
+ Rows with a negative slot keep `output`'s prior content, or are NaN
+ (unspecified) when no `output` is given. `old_beta`, `launch_with_pdl`,
+ `n_writes` and the launch-tuning knobs are ignored; `replay_work_items`
+ only decides whether a wrap commits the checkpoint here, unless `in_kernel_commit`
+ says (True: every wrap commits here; False: none does); `config` is ignored."""
+ N, T, H, K = k.shape
+ HV, V = v.shape[2], v.shape[3]
+ rep = HV // H
+ dev = v.device
+ scale = K**-0.5 if scale is None else scale
+ if packed_qkv is not None:
+ rows = packed_qkv.view(N, T, -1)
+ q = rows[..., : H * K].view(N, T, H, K)
+ k = rows[..., H * K : 2 * H * K].view(N, T, H, K)
+ v = rows[..., 2 * H * K :].view(N, T, HV, V)
+ qf, kf, vf = q.float(), k.float(), v.float()
+ if use_qk_l2norm_in_kernel:
+ qf = qf / torch.sqrt(qf.square().sum(-1, keepdim=True) + 1e-6)
+ kf = kf / torch.sqrt(kf.square().sum(-1, keepdim=True) + 1e-6)
+ qf = (qf * scale).repeat_interleave(rep, dim=2) # [N, T, HV, K]
+ kv = kf.repeat_interleave(rep, dim=2) # [N, T, HV, K]
+ if A_log is not None:
+ gf = -torch.exp(A_log.float()) * F.softplus(g.float() + dt_bias.float())
+ bf = torch.sigmoid(beta.float())
+ else:
+ gf, bf = g.float(), beta.float()
- def reference(self, *args, **kwargs):
- raise NotImplementedError(
- "the multi-step cell harness drives the Triton entry's reference (see the test)"
+ out_dtype = q.dtype if output is None else output.dtype
+ out = (
+ torch.full((N, T, HV, V), float("nan"), device=dev)
+ if output is None
+ else output.float().clone()
+ )
+ pool, u_buf, k_buf, G_buf = (x.clone() for x in (ssm_states, old_u, old_k, old_G))
+ live = (state_indices[:N] >= 0).nonzero().flatten()
+ if live.numel() == 0:
+ return out.to(out_dtype), pool, u_buf, k_buf, G_buf
+ slot = state_indices[:N][live].long()
+ p = prev_num_accepted_tokens[slot].long()
+ b = cache_buf_idx[slot].long()
+ M, depth = live.numel(), old_u.shape[2]
+
+ # 1. the pending state: checkpoint plus the p history entries of half b
+ pending = torch.arange(depth, device=dev)[None, :] < p[:, None] # [M, depth]
+ Gh = torch.where(pending[:, None, :], old_G[slot, b].float(), 0.0) # [M, HV, depth]
+ Uh = torch.where(pending[:, :, None, None], old_u[slot, b].float(), 0.0)
+ Kh = torch.where(pending[:, :, None, None], old_k[slot, b].float(), 0.0)
+ Kh = Kh.repeat_interleave(rep, dim=2) # [M, depth, HV, K]
+ last = (p - 1).clamp(min=0)[:, None, None].expand(M, HV, 1)
+ g_start = torch.where(p[:, None] > 0, Gh.gather(2, last).squeeze(2), 0.0) # [M, HV]
+ weight = torch.where(pending[:, None, :], torch.exp(g_start[:, :, None] - Gh), 0.0)
+ h = torch.exp(g_start)[:, :, None, None] * ssm_states[slot].float() # [M, HV, V, K]
+ for j in range(depth):
+ h = h + weight[:, :, j, None, None] * Uh[:, j, :, :, None] * Kh[:, j, :, None, :]
+ checkpoint = h
+
+ # 2. the new tokens, in order
+ gl, bl = gf[live], bf[live] # [M, T, HV]
+ ql, kl, vl = qf[live], kv[live], vf[live]
+ G_local = torch.zeros(M, HV, device=dev)
+ u_new, G_new = [], []
+ for t in range(T):
+ h = h * torch.exp(gl[:, t])[:, :, None, None]
+ u = bl[:, t, :, None] * (vl[:, t] - (h * kl[:, t, :, None, :]).sum(-1)) # [M, HV, V]
+ h = h + u[:, :, :, None] * kl[:, t, :, None, :]
+ out[live, t] = (h * ql[:, t, :, None, :]).sum(-1)
+ G_local = G_local + gl[:, t]
+ u_new.append(u)
+ G_new.append(G_local.clone())
+
+ # 3. the history append, and the checkpoint commit on a wrap
+ wrap = p + T > history_size
+ w_buf = torch.where(wrap, 1 - b, b)
+ w_off = torch.where(wrap, 0, p)
+ for t in range(T):
+ pos = w_off + t
+ u_buf[slot, w_buf, pos] = u_new[t].to(u_buf.dtype)
+ k_buf[slot, w_buf, pos] = kf[live, t].to(k_buf.dtype)
+ G_store = torch.where(wrap[:, None], G_new[t], g_start + G_new[t])
+ G_buf[slot, w_buf, :, pos] = G_store.to(G_buf.dtype)
+ commit = replay_work_items is None if in_kernel_commit is None else bool(in_kernel_commit)
+ if commit and bool(wrap.any()):
+ pool[slot[wrap]] = checkpoint[wrap].to(pool.dtype)
+ return out.to(out_dtype), pool, u_buf, k_buf, G_buf
+
+ def is_valid(
+ self,
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ g: torch.Tensor,
+ beta: torch.Tensor,
+ ssm_states: torch.Tensor,
+ state_indices: torch.Tensor,
+ old_u: torch.Tensor,
+ old_k: torch.Tensor,
+ old_G: torch.Tensor,
+ old_beta: torch.Tensor,
+ cache_buf_idx: torch.Tensor,
+ prev_num_accepted_tokens: torch.Tensor,
+ history_size: int,
+ scale: Optional[float] = None,
+ use_qk_l2norm_in_kernel: bool = False,
+ A_log: Optional[torch.Tensor] = None,
+ dt_bias: Optional[torch.Tensor] = None,
+ launch_with_pdl: bool = False,
+ replay_work_items: Optional[torch.Tensor] = None,
+ n_writes: Optional[torch.Tensor] = None,
+ block_v: Optional[int] = None,
+ num_warps: Optional[int] = None,
+ use_l2_state_cache: Optional[bool] = None,
+ use_l2_streaming_inputs: Optional[bool] = None,
+ use_l2_shared_inputs: Optional[bool] = None,
+ main_num_stages: Optional[int] = None,
+ packed_qkv: Optional[torch.Tensor] = None,
+ use_all_layer_commit: bool = False,
+ output: Optional[torch.Tensor] = None,
+ in_kernel_commit: Optional[bool] = None,
+ config: Optional[dict] = None,
+ ) -> None:
+ N, _, H, K = k.shape
+ HV, V = v.shape[2], v.shape[3]
+ assert state_indices.dim() == 1 and state_indices.shape[0] >= N, (
+ f"state_indices has shape {tuple(state_indices.shape)} for {N} sequences -- the "
+ "kernel reads one slot per sequence with no bound, so a shorter tensor hands the "
+ "trailing rows whatever follows it in memory, and they write that slot's history"
+ )
+ pool_slots, slots = ssm_states.shape[0], old_u.shape[0]
+ assert slots >= pool_slots, (
+ f"the history buffers have {slots} slots for a {pool_slots}-slot pool -- a slot "
+ "past the history's end is read and written out of bounds"
)
+ for name, tensor in (
+ ("prev_num_accepted_tokens", prev_num_accepted_tokens),
+ ("cache_buf_idx", cache_buf_idx),
+ ):
+ assert tensor.dim() == 1 and tensor.is_contiguous() and tensor.shape[0] >= pool_slots, (
+ f"{name} has shape {tuple(tensor.shape)} and stride {tensor.stride()}; the kernel "
+ f"reads one element per slot of the {pool_slots}-slot pool, at the slot's index"
+ )
+ expected = {
+ "old_u": (old_u, (slots, 2, history_size, HV, V)),
+ "old_k": (old_k, (slots, 2, history_size, H, K)),
+ "old_G": (old_G, (slots, 2, HV, history_size)),
+ }
+ for name, (tensor, shape) in expected.items():
+ assert tuple(tensor.shape) == shape, (
+ f"{name} has shape {tuple(tensor.shape)}; history_size={history_size} and this "
+ f"geometry need {shape} -- the kernel addresses the history by history_size and "
+ "the head geometry, never by the buffer's own shape"
+ )
gdn_cached_replay_cute = GdnCachedReplayCute()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/vision/__init__.py
similarity index 100%
rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/__init__.py
rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/vision/__init__.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/vision/triton_pos_embed_interpolate.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/vision/triton_pos_embed_interpolate.py
new file mode 100644
index 000000000000..5b395305068b
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/vision/triton_pos_embed_interpolate.py
@@ -0,0 +1,210 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Vision position-embedding bilinear interpolation onto one image's patch
+grid, emitted in spatial-merge order, by the stock Qwen3-VL tower's fused
+Triton kernel.
+
+`__call__` is one call to `_triton_pos_embed_interpolate` in
+`tensorrt_llm/_torch/models/modeling_qwen3vl.py`, which allocates the output
+and makes one launch of the Triton JIT kernel `_bilinear_pos_embed_kernel`
+beside it: the four corner gathers, the bilinear weights and the merge-order
+row permutation in one pass. It is the kernel the stock `modeling_qwen4_exp`
+vision tower runs (`Qwen3VisionModel.fast_pos_embed_interpolate`), one call per
+image. Concatenating several images' rows is the caller's.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
+from tensorrt_llm._torch.models.modeling_qwen3vl import _triton_pos_embed_interpolate
+
+#: This checkpoint: num_position_embeddings 2304 (a 48 x 48 grid), hidden 1152,
+#: spatial_merge_size 2.
+NUM_GRID = 48
+HIDDEN = 1152
+MERGE = 2
+
+
+def _ckpt(t: int, h: int, w: int) -> dict:
+ """The checkpoint's bf16 table at one `(t, h, w)` grid."""
+ return dict(
+ t=t,
+ h=h,
+ w=w,
+ num_grid_per_side=NUM_GRID,
+ m_size=MERGE,
+ dtype=torch.bfloat16,
+ hidden=HIDDEN,
+ )
+
+
+class TritonPosEmbedInterpolate(OpWrapper):
+ """Resample a square row-major `[N*N, hidden]` table onto one `(t, h, w)`
+ patch grid; return a fresh `[t*h*w, hidden]` tensor in `dtype`.
+
+ Corner-aligned bilinear (`align_corners=True`): row `i` of `h` samples the
+ table at `i * (N - 1) / (h - 1)` (0 when `h == 1`), `lo = floor`,
+ `hi = min(lo + 1, N - 1)`; `h` indexes table rows (stride `N`), `w`
+ columns. Rows come out in spatial-merge order -- whole `m_size x m_size`
+ blocks row-major, row-major inside a block -- which is what the patch
+ merger's 4-to-1 reshape consumes. `t` repeats the frame's rows.
+
+ Arithmetic is the HF reference's (`fast_pos_embed_interpolate`): the four
+ corner weights are computed in fp32, rounded to `dtype`, and the
+ multiply-adds run in `dtype`. Not an fp32 accumulation; see `note`.
+
+ Fusion boundary: one image's position embeddings. Adding them to the patch
+ embeddings, the 2-D rope, and joining images are the caller's.
+ """
+
+ ARCHS = frozenset({Arch.SM_103})
+
+ # spec: t, h, w, num_grid_per_side, m_size, dtype, hidden. The tower is
+ # replicated on every rank, so these are every member's shapes.
+ CELLS: tuple[Cell, ...] = (
+ Cell(why="one 4x8 image: 32 rows, both sides below the table's 48", spec=_ckpt(1, 4, 8)),
+ Cell(why="the 62x90 serving grid: 5580 rows, one image per call", spec=_ckpt(1, 62, 90)),
+ Cell(why="a 3-frame 6x10 video: t repeats one frame's 60 rows", spec=_ckpt(3, 6, 10)),
+ Cell(
+ why="128x256: 32768 rows, the tower's per-call cap and the envelope's largest",
+ spec=_ckpt(1, 128, 256),
+ ),
+ )
+
+ note = """
+ Source: a Triton JIT kernel, compiled at first call for the device it runs
+ on. `H` and `W` are int32, the scales fp32, and `NUM_GRID`, `M_SIZE`,
+ `HIDDEN_DIM` (1152) and `BLOCK_D` (2048) are constexprs, so another table
+ geometry compiles another kernel. Tier 2: a callable already in the tree (a
+ private function of the stock Qwen3-VL modeling file), not a registered
+ `torch.ops.trtllm` op. Importing this entry imports `modeling_qwen3vl`.
+
+ Accuracy, against the exact (fp64) corner-aligned bilinear value, as
+ row-scale element error / relative RMS in bf16 ulp (`assert_within_ulp`):
+ at most 3.522 / 0.752 over 50 seeds of the test's smooth table (a bicubic
+ upsampled normal field, neighbour cosine 0.911 against the checkpoint
+ table's 0.907), and at most 3.349 / 0.745 on the checkpoint's own
+ `model.visual.pos_embed.weight`. The test gates 4 / 1. The closest wrong
+ variant (align_corners=False) lands at 41.3 / 15.2 or more; raster row
+ order and swapped table axes above 400 / 300.
+
+ Not certified on tables whose neighbouring rows are uncorrelated: a normal
+ table with independent rows reaches 4.182 / 0.879 over 20 seeds.
+ Averaging uncorrelated corners shrinks the output's row maximum below the
+ corners' magnitude, and the bf16 rounding error scales with the corners. A
+ learned 2-D position table is not such a table.
+
+ Silent, refused by `is_valid` (each measured on GB300):
+ - a strided table: the kernel indexes `embed_weight` as dense rows of
+ `hidden`. A column slice of a wider buffer returned, differing from the
+ same values made contiguous by up to 5.05.
+ - `num_grid_per_side ** 2 != embed_weight.shape[0]`: the grid side is an
+ argument, not read off the table. 47 on the 2304-row table returned a
+ correctly shaped result 564 / 340 ulp off.
+ - an empty side (`t`, `h` or `w` of 0): returns `[0, hidden]`, no error.
+ - `h` or `w` not divisible by `m_size`: the wrapper's own check is a bare
+ `assert`. Under `python -O` a 6x9 grid returned a finite `[54, 1152]`
+ tensor, its rows read from positions outside the grid and the table.
+
+ Hard limits: the output offset is `program_id * hidden` in int32, so
+ `t*h*w*hidden` must stay below 2^31 (1,864,135 rows at hidden 1152; the
+ tower's per-call cap is 32768). Positions are `row * scale` in fp32,
+ exact up to 2^24 rows per side.
+
+ Not certified: fp16 or fp32 tables, a `dtype` other than the table's,
+ `m_size` other than 2, a hidden size other than 1152, a CPU table, and
+ CUDA-graph capture (the call site, the vision tower, is eager only; the
+ grid is a host argument, frozen at capture).
+ """
+
+ def __call__(
+ self,
+ embed_weight: torch.Tensor,
+ t: int,
+ h: int,
+ w: int,
+ num_grid_per_side: int,
+ m_size: int,
+ dtype: torch.dtype,
+ ) -> torch.Tensor:
+ return _triton_pos_embed_interpolate(
+ embed_weight, t, h, w, num_grid_per_side, m_size, dtype
+ )
+
+ def reference(
+ self,
+ embed_weight: torch.Tensor,
+ t: int,
+ h: int,
+ w: int,
+ num_grid_per_side: int,
+ m_size: int,
+ dtype: torch.dtype,
+ ) -> torch.Tensor:
+ """The explicit four-corner weighted sum in fp64, returned as fp32.
+
+ Sample positions are the exact `i * (N - 1) / (size - 1)` in fp64.
+ `dtype` is not applied: the result is the exact value rounded once to
+ fp32, so a test's band in `dtype` ulp measures the op's error alone
+ rather than adding the reference's own rounding to `dtype`.
+ """
+ n = num_grid_per_side
+ hidden = embed_weight.shape[1]
+ device = embed_weight.device
+ grid = embed_weight.double().reshape(n, n, hidden)
+
+ def positions(size: int) -> torch.Tensor:
+ index = torch.arange(size, dtype=torch.float64, device=device)
+ return index * ((n - 1) / (size - 1)) if size > 1 else torch.zeros_like(index)
+
+ h_pos, w_pos = positions(h), positions(w)
+ h_lo, w_lo = h_pos.floor().long(), w_pos.floor().long()
+ h_hi = (h_lo + 1).clamp(max=n - 1)
+ w_hi = (w_lo + 1).clamp(max=n - 1)
+ d_h = (h_pos - h_lo)[:, None, None]
+ d_w = (w_pos - w_lo)[None, :, None]
+ raster = (
+ grid[h_lo][:, w_lo] * ((1 - d_h) * (1 - d_w))
+ + grid[h_lo][:, w_hi] * ((1 - d_h) * d_w)
+ + grid[h_hi][:, w_lo] * (d_h * (1 - d_w))
+ + grid[h_hi][:, w_hi] * (d_h * d_w)
+ )
+ rows = (
+ raster.reshape(h // m_size, m_size, w // m_size, m_size, hidden)
+ .permute(0, 2, 1, 3, 4)
+ .reshape(h * w, hidden)
+ )
+ return (rows.repeat(t, 1) if t > 1 else rows).float()
+
+ def is_valid(
+ self,
+ embed_weight: torch.Tensor,
+ t: int,
+ h: int,
+ w: int,
+ num_grid_per_side: int,
+ m_size: int,
+ dtype: torch.dtype,
+ ) -> None:
+ assert embed_weight.is_contiguous(), (
+ "embed_weight must be contiguous; the kernel indexes it as dense rows "
+ "and reads a strided table wrongly without raising"
+ )
+ assert num_grid_per_side * num_grid_per_side == embed_weight.shape[0], (
+ f"num_grid_per_side {num_grid_per_side} does not describe a "
+ f"{embed_weight.shape[0]}-row table; the kernel resamples the grid it is "
+ f"told, in or out of the table"
+ )
+ assert t > 0 and h > 0 and w > 0, (
+ f"empty grid side (t, h, w) = ({t}, {h}, {w}); the kernel returns no rows "
+ f"rather than raising"
+ )
+ assert h % m_size == 0 and w % m_size == 0, (
+ f"grid sides h={h}, w={w} are not whole m_size={m_size} merge blocks; "
+ f"the kernel's bare assert is all that stops an out-of-grid read"
+ )
+
+
+triton_pos_embed_interpolate = TritonPosEmbedInterpolate()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/pool.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/pool.py
new file mode 100644
index 000000000000..4d88a74eaeb0
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/pool.py
@@ -0,0 +1,488 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The KV cache manager's pools as one rank of this target sees them, and their contract.
+
+Every rank builds the stock ``Qwen4ExpHybridCacheManager`` (``modules/qwen4_exp/cache_manager.py``,
+a ``MambaHybridCacheManagerV2``): QSA is this architecture's default sparse attention, so
+``get_kv_cache_manager_cls`` resolves the manager through ``get_sparse_attn_kv_cache_manager``,
+which requires KV cache manager V2. The manager sizes each pool from the checkpoint's text config
+and the rank's ``Mapping``; `pool_geometry` derives the same per-rank shapes the same way:
+
+=================== ============================================ =============================
+pool shape after the slot dim split over ranks by
+=================== ============================================ =============================
+paged K/V (HND) ``[2, H_kv, tokens_per_block, kv_width]`` attention TP (a)
+NVFP4 block scales ``[2, H_kv, tokens_per_block, head_dim/16]`` attention TP (a)
+QSA index plane ``[tokens_per_block, 1, index_head_dim]`` replicated (b)
+QSA position plane ``[tokens_per_block, 3]`` int32 replicated (b)
+GDN conv state ``[2 * H_k * D_k + H_v * D_v, kernel - 1]`` recurrent-state TP (c)
+GDN recurrent state ``[H_v, D_v, D_k]`` recurrent-state TP (c)
+PLE conv state ``[hc_count * hidden, (ple_kernel-1) * n]`` replicated (d)
+PLE n-gram history ``[n - 1]`` int64 replicated (d)
+=================== ============================================ =============================
+
+``H_kv`` is ``num_key_value_heads``; ``H_k``/``D_k`` and ``H_v``/``D_v`` are the GDN key and value
+heads (``linear_num_key_heads``, ``linear_key_head_dim``, ...); ``kernel`` is
+``linear_conv_kernel_dim``, ``ple_kernel`` ``ple_conv_kernel_size`` and ``n`` ``ngram_size``. Head
+counts are the rank's share.
+
+(a) ``KVCacheManagerV2.__init__`` splits ``num_key_value_heads`` over ``tp_size``, or over 1 under
+ attention DP, rounding up: with more TP ranks than KV heads each head is repeated, and rank
+ ``r`` holds global KV head ``r * H / tp`` (the order ``duplicate_kv_weight`` repeats them in).
+(b) ``Qwen4ExpHybridCacheManager._extra_buffers_per_layer`` sizes them from ``QSASparseParams``
+ (``index_kv_heads``, ``index_head_dim``), which no mapping divides: the indexer is replicated.
+ Under an FP8 or NVFP4 pool both are packed into one ``QSA_SIDE_PLANES`` buffer; the views
+ keep these shapes either way.
+(c) ``MambaHybridCacheManagerV2.__init__`` divides the GDN value heads, key heads (``n_groups``)
+ and ``conv_dim`` by ``mamba_effective_tp_size``: ``tp_size``, or 1 under attention DP. The
+ conv channels of a rank are ``[q | k | v]`` of its own heads (``conv_state_layout`` "q_k_v",
+ the manager's ``conv_section_dims``), and rank ``r`` holds key heads ``[r * H_k / tp, ...)``
+ and value heads ``[r * H_v / tp, ...)``: the contiguous head ranges a disaggregated transfer
+ between layouts moves (``get_disagg_role_mapper_kinds``: conv ``SECTIONED``, recurrent state
+ by head).
+(d) ``extract_qwen4_exp_ple_cache_params``: the PLE layer reads the whole hyper-connection
+ stream on every rank.
+
+At tp4 and tep4 a rank therefore holds one of the two KV heads, four key and twelve value GDN
+heads; at dep4 it holds all of them, and the recurrent pools carry one more slot (the
+attention-DP dummy request's). With MTP, attention TP puts the draft layer in a separate draft
+KV cache (``should_use_separate_draft_kv_cache``); under attention DP the draft layer is a layer
+of the trunk's manager. Its pools have the trunk's per-rank shapes in either.
+
+The ``check_*`` functions read only ``.shape``, ``.dtype`` and ``.stride()``, so they run on meta
+tensors as well as on a manager's views. Each one is a wrong pool its consuming kernel would
+accept and mis-address.
+"""
+
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import torch
+
+from tensorrt_llm._torch.pyexecutor.config_utils import mamba_effective_tp_size
+from tensorrt_llm.quantization.mode import QuantAlgo
+
+#: Coordinates per token in the QSA position plane (the mRoPE layout's three rows).
+QSA_POSITION_AXES = 3
+
+#: Values per e4m3 block scale of an NVFP4 K/V pool.
+NVFP4_BLOCK = 16
+
+#: The recurrent-state pool dtypes the target serves (`kv_cache_config.mamba_ssm_cache_dtype`): the
+#: checkpoint's fp32, or bf16.
+SSM_STATE_DTYPES = (torch.float32, torch.bfloat16)
+
+#: The GDN cached replay's history length the replay entry is certified at.
+GDN_REPLAY_HISTORY = 16
+
+#: The most query rows a verify step carries per generation request, `max_draft_len + 1`: the
+#: replay kernel asserts `T <= 8` and `thop_attention`'s `predicted_tokens_per_seq` is certified
+#: up to it.
+MAX_VERIFY_WIDTH = 8
+
+
+class PoolGeometry(NamedTuple):
+ """The pool shapes one rank binds. Head counts and channel widths are this rank's share."""
+
+ tokens_per_block: int
+ heads_kv: int # KV heads in this rank's K/V pages
+ kv_head_start: int # the global KV head the rank's first page head holds
+ head_dim: int
+ index_kv_heads: int
+ index_head_dim: int
+ conv_dim: int # this rank's GDN conv channels, [q | k | v] of its own heads
+ conv_kernel: int
+ num_k_heads: int # this rank's GDN key heads
+ k_head_start: int # the global key head the rank's first one is
+ num_v_heads: int # this rank's GDN value heads
+ v_head_start: int # the global value head the rank's first one is
+ head_v_dim: int
+ head_k_dim: int
+ stream: int # the hyper-connection stream width, the PLE conv channels
+ ple_state_len: int
+ ngram_context_len: int
+ dtype: torch.dtype # the activation dtype: GDN conv, QSA index and PLE conv pools
+ kv_dtype: torch.dtype # the paged K/V pool's dtype: bf16, e4m3 (FP8) or uint8 (NVFP4)
+ kv_width: int # its innermost extent: head_dim, or head_dim // 2 for NVFP4's e2m1 pairs
+
+ @property
+ def conv_sections(self) -> tuple[int, int, int]:
+ """The q, k and v channel widths of this rank's conv state, in that order."""
+ key = self.num_k_heads * self.head_k_dim
+ return key, key, self.num_v_heads * self.head_v_dim
+
+ @property
+ def kv_block_scales(self) -> bool:
+ """Whether the K/V pool is NVFP4, whose e4m3 block scales sit in a second pool."""
+ return self.kv_dtype == torch.uint8
+
+
+def attention_tp_size(mapping) -> int:
+ """The ranks the KV heads are split over (`KVCacheManagerV2.__init__`)."""
+ return 1 if mapping.enable_attention_dp else mapping.tp_size
+
+
+def _share(total: int, ranks: int, rank: int) -> tuple[int, int]:
+ """(heads on `rank`, its first global head) when `total` heads are split over `ranks`: an even
+ split, or each head repeated on `ranks // total` neighbouring ranks when there are more
+ ranks."""
+ return -(-total // ranks), rank * total // ranks
+
+
+#: The paged K/V pool's storage per `kv_cache_quant_algo`: (dtype or None for the activation
+#: dtype, values per stored element).
+_KV_STORAGE = {
+ None: (None, 1),
+ QuantAlgo.FP8: (torch.float8_e4m3fn, 1),
+ QuantAlgo.NVFP4: (torch.uint8, 2),
+}
+
+
+def pool_geometry(
+ cfg, mapping, *, tokens_per_block: int, kv_cache_quant_algo, dtype: torch.dtype
+) -> PoolGeometry:
+ """The pools rank `mapping.rank` binds.
+
+ Args:
+ cfg: the engine-normalized text config (`Qwen4ExpTextConfig`) the manager is sized from.
+ mapping: this rank's `Mapping`.
+ tokens_per_block: the manager's page size (`kv_cache_config.tokens_per_block`).
+ kv_cache_quant_algo: the KV-cache algorithm the engine resolved: None (the checkpoint's
+ bf16), `QuantAlgo.FP8` or `QuantAlgo.NVFP4`.
+ dtype: the activation dtype.
+ """
+ if kv_cache_quant_algo not in _KV_STORAGE:
+ raise ValueError(
+ f"kv_cache_quant_algo {kv_cache_quant_algo!r}; the K/V pool stores the checkpoint's "
+ "bf16, FP8 or NVFP4"
+ )
+ attn_tp = attention_tp_size(mapping)
+ heads_kv, kv_head_start = _share(
+ cfg.num_key_value_heads, attn_tp, mapping.tp_rank if attn_tp > 1 else 0
+ )
+ state_tp = mamba_effective_tp_size(mapping)
+ state_rank = mapping.tp_rank if state_tp > 1 else 0
+ num_k_heads, k_head_start = _share(cfg.linear_num_key_heads, state_tp, state_rank)
+ num_v_heads, v_head_start = _share(cfg.linear_num_value_heads, state_tp, state_rank)
+ kv_dtype, per_element = _KV_STORAGE[kv_cache_quant_algo]
+ return PoolGeometry(
+ tokens_per_block=tokens_per_block,
+ heads_kv=heads_kv,
+ kv_head_start=kv_head_start,
+ head_dim=cfg.head_dim,
+ index_kv_heads=cfg.indexer_kv_heads,
+ index_head_dim=cfg.indexer_head_dim,
+ conv_dim=2 * num_k_heads * cfg.linear_key_head_dim
+ + num_v_heads * cfg.linear_value_head_dim,
+ conv_kernel=cfg.linear_conv_kernel_dim,
+ num_k_heads=num_k_heads,
+ k_head_start=k_head_start,
+ num_v_heads=num_v_heads,
+ v_head_start=v_head_start,
+ head_v_dim=cfg.linear_value_head_dim,
+ head_k_dim=cfg.linear_key_head_dim,
+ stream=cfg.hc_count * cfg.hidden_size,
+ ple_state_len=(cfg.ple_conv_kernel_size - 1) * cfg.ngram_size,
+ ngram_context_len=cfg.ngram_size - 1,
+ dtype=dtype,
+ kv_dtype=kv_dtype or dtype,
+ kv_width=cfg.head_dim // per_element,
+ )
+
+
+# ------------------------------------------------------------------ contract ---
+
+
+def require(cm, names, what: str) -> None:
+ """`cm` exposes every attribute in `names`, which this target reads to `what`."""
+ for name in names:
+ assert hasattr(cm, name), (
+ f"the KV cache manager {type(cm).__name__} does not expose {name}; this target {what}"
+ )
+
+
+def _check_pool(t, what, lead, tail, dtype=None) -> None:
+ """`t` has `lead` leading dims of any size, then exactly `tail`, at `dtype`."""
+ assert t is not None, f"{what}: the runtime did not allocate it"
+ expected = ["*"] * lead + [str(n) for n in tail]
+ assert t.dim() == lead + len(tail) and tuple(t.shape[lead:]) == tuple(tail), (
+ f"{what} is {tuple(t.shape)}, expected [{', '.join(expected)}]"
+ )
+ assert dtype is None or t.dtype == dtype, f"{what} is {t.dtype}, expected {dtype}"
+
+
+def check_position_plane(g: PoolGeometry, position) -> None:
+ shape = (g.tokens_per_block, QSA_POSITION_AXES)
+ _check_pool(position, "QSA position plane", 1, shape, torch.int32)
+
+
+def check_attention_pools(g: PoolGeometry, layer: int, index_cache, pool) -> None:
+ """QSA index plane and paged K/V for one attention layer."""
+ index_shape = (g.tokens_per_block, g.index_kv_heads, g.index_head_dim)
+ _check_pool(index_cache, f"layer {layer}: QSA index plane", 1, index_shape, g.dtype)
+ # The sparse path needs adjacent HND K/V role planes: [pages, 2, H, block, D].
+ kv_shape = (g.heads_kv, g.tokens_per_block, g.kv_width)
+ _check_pool(pool, f"layer {layer}: paged K/V", 2, kv_shape, g.kv_dtype)
+ assert pool.shape[1] == 2, f"layer {layer}: paged K/V has {pool.shape[1]} role planes, not 2"
+
+
+def check_kv_scale_planes(g: PoolGeometry, layer: int, pool, scales) -> None:
+ """An NVFP4 layer's e4m3 block scales: one uint8 per `NVFP4_BLOCK` values, laid out like the
+ data pages (`[slots, 2, H, block, D / 16]`), and holding as many pages per slot as the data
+ pool does.
+
+ The attention op addresses a layer's data pages and its scale pages with one block-offset
+ table, so a scale pool whose pages-per-slot differs from the data pool's reads and writes other
+ slots' scales without any error. Both counts are read off the views' slot strides."""
+ scale_shape = (g.heads_kv, g.tokens_per_block, g.head_dim // NVFP4_BLOCK)
+ _check_pool(scales, f"layer {layer}: NVFP4 block scales", 2, scale_shape, torch.uint8)
+ assert scales.shape[1] == 2 and scales.shape[0] == pool.shape[0], (
+ f"layer {layer}: the block scales are {tuple(scales.shape)} against paged K/V "
+ f"{tuple(pool.shape)}; they must hold the same slots, K then V"
+ )
+ pages = []
+ for what, t in (("paged K/V", pool), ("block scales", scales)):
+ page = t[0, 0].numel()
+ assert t.stride(0) % page == 0, (
+ f"layer {layer}: the {what} slot stride {t.stride(0)} is not a whole number of "
+ f"{page}-element pages"
+ )
+ pages.append(t.stride(0) // page)
+ assert pages[0] == pages[1], (
+ f"layer {layer}: the paged K/V pool holds {pages[0]} pages per slot and the block-scale "
+ f"pool {pages[1]}; the attention op would address the scales of other slots"
+ )
+
+
+def check_recurrent_pools(g: PoolGeometry, conv, ssm) -> None:
+ _check_pool(conv, "conv state pool", 1, (g.conv_dim, g.conv_kernel - 1), g.dtype)
+ # Key dim innermost.
+ _check_pool(ssm, "recurrent state pool", 1, (g.num_v_heads, g.head_v_dim, g.head_k_dim))
+ assert ssm.dtype in SSM_STATE_DTYPES, (
+ f"the recurrent state pool is {ssm.dtype}; this target serves "
+ f"{' or '.join(str(d) for d in SSM_STATE_DTYPES)} (kv_cache_config.mamba_ssm_cache_dtype)"
+ )
+
+
+def check_conv_layout(g: PoolGeometry, section_dims) -> None:
+ """The manager splits a rank's conv channels as `[q | k | v]` of the rank's own heads, the
+ order this target's input projection writes them in. Another order (`x_b_c`) or another split
+ has the same total width, so the conv kernel would read every section from the wrong place."""
+ assert tuple(section_dims) == g.conv_sections, (
+ f"the manager lays a rank's conv state out as sections {tuple(section_dims)}; this target "
+ f"writes [q | k | v] = {g.conv_sections} (conv_state_layout 'q_k_v')"
+ )
+
+
+def check_ple_pools(g: PoolGeometry, ple_conv, ple_ngram, conv) -> None:
+ _check_pool(ple_conv, "PLE conv state", 1, (g.stream, g.ple_state_len), g.dtype)
+ # The id-history entry writes int64 into it; an int32 pool raises.
+ _check_pool(ple_ngram, "PLE n-gram context pool", 1, (g.ngram_context_len,), torch.int64)
+ # The forward addresses all three with one scheduler slot vector.
+ assert ple_conv.shape[0] == conv.shape[0] == ple_ngram.shape[0], (
+ "the PLE and Gated-DeltaNet pools must share one slot space"
+ )
+
+
+def check_step_buffers(g: PoolGeometry, int_conv, conv, int_idx) -> None:
+ """MTP verify buffers both recurrent modes use: the per-step conv windows
+ `[slots, steps, ...]` at the pool's dtype (the promotion copies between them without
+ converting) and the rows `update_mamba_states` reads its sources from."""
+ assert int_conv is not None, (
+ "the runtime allocated no per-step conv buffer; a partially accepted draft "
+ "cannot be rolled back"
+ )
+ _check_pool(int_conv, "per-step conv buffer", 2, (g.conv_dim, g.conv_kernel - 1), conv.dtype)
+ assert int_idx is not None and int_idx.dtype == torch.int32, (
+ "the cache manager exposes no int32 intermediate_state_indices"
+ )
+
+
+def check_replay_state(history, width, prev_accepted, cache_buf_idx, ssm) -> None:
+ """Cached replay, manager-wide: the geometry the replay entry is certified at, and the
+ per-slot bookkeeping, which must cover every pool slot."""
+ assert ssm.dtype in SSM_STATE_DTYPES, (
+ f"the recurrent pool is {ssm.dtype}; the replay entry is certified at an fp32 or a "
+ "bf16 checkpoint"
+ )
+ assert history == GDN_REPLAY_HISTORY, (
+ f"replay_history_size {history}; the replay entry is certified at {GDN_REPLAY_HISTORY}"
+ )
+ assert 2 <= width <= min(MAX_VERIFY_WIDTH, history), (
+ f"replay_step_width {width}; the replay entry is certified at 2..{MAX_VERIFY_WIDTH} "
+ f"verify rows and never above the history ({history})"
+ )
+ for name, t in (("prev_num_accepted_tokens", prev_accepted), ("cache_buf_idx", cache_buf_idx)):
+ assert t.dim() == 1 and t.shape[0] >= ssm.shape[0] and t.dtype == torch.int32, (
+ f"{name} is {tuple(t.shape)} {t.dtype}; the kernel reads one int32 per pool slot "
+ f"({ssm.shape[0]})"
+ )
+
+
+def check_replay_history(
+ g: PoolGeometry, layer: int, history, pool_slots, old_x, old_B, old_dt
+) -> None:
+ """One GDN layer's replay history: `[slots, 2 halves, ...]` over this rank's heads, addressed
+ by `history` and the head geometry, never by the buffer's own shape."""
+ for name, t in (("old_x", old_x), ("old_B", old_B), ("old_dt", old_dt)):
+ assert t is not None, f"layer {layer}: the manager's layer cache carries no {name}"
+ assert old_x.shape[0] >= pool_slots, (
+ f"layer {layer}: the replay history has {old_x.shape[0]} slots for a {pool_slots}-slot pool"
+ )
+ bf16, fp32 = torch.bfloat16, torch.float32
+ x_tail, b_tail = (history, g.num_v_heads, g.head_v_dim), (history, g.num_k_heads, g.head_k_dim)
+ _check_pool(old_x, f"layer {layer}: old_x", 2, x_tail, bf16)
+ _check_pool(old_B, f"layer {layer}: old_B", 2, b_tail, bf16)
+ _check_pool(old_dt, f"layer {layer}: old_dt", 2, (g.num_v_heads, history), fp32)
+ assert old_x.shape[1] == old_B.shape[1] == old_dt.shape[1] == 2, (
+ f"layer {layer}: the replay history is not double-buffered"
+ )
+
+
+def check_verify_width(rows: int) -> None:
+ """The rows per generation request one step carries: `max_draft_len + 1` (`runtime_draft_len
+ + 1` at run time), a `thop_attention` `predicted_tokens_per_seq` the entry is certified at."""
+ assert 1 <= rows <= MAX_VERIFY_WIDTH, (
+ f"predicted_tokens_per_seq {rows} is outside the 1..{MAX_VERIFY_WIDTH} range "
+ "attention/thop_attention is certified at"
+ )
+
+
+def check_replay_width(replay, rows: int) -> None:
+ """Under the cached replay a single-row generation step cannot take the ordinary
+ recurrence: the pool row is a checkpoint whose pending history only the replay kernel
+ applies. The manager's step width (`replay.width`) is the one row count per sequence a step
+ may carry."""
+ assert replay is None or rows == replay.width, (
+ f"the batch carries {rows} query rows per generation sequence and the GDN cached "
+ f"replay runs at step width {replay.width}; a runtime_draft_len of {rows - 1} (a "
+ "draft-length schedule, or skipped drafting) cannot run under it. Export "
+ "TRTLLM_USE_GDN_REPLAY=0 for variable draft lengths"
+ )
+
+
+def check_shared_selection(shared, slots, layer: int, block_topk: int) -> None:
+ """IndexShare: the metadata's shared selection buffers have a row for `layer` and are
+ `block_topk` wide."""
+ assert shared is not None, "IndexShare is on but the QSA metadata allocated no shared buffers"
+ assert layer in slots and slots[layer] < shared.shape[0], (
+ f"the KV cache manager gives layer {layer} no row of the {shared.shape[0]}-row shared "
+ f"selection buffers (qsa_shared_topk_slots={slots})"
+ )
+ assert shared.shape[2] == block_topk, (
+ f"the shared selection rows are {shared.shape[2]} wide and this target selects "
+ f"{block_topk} blocks per row"
+ )
+
+
+# ------------------------------------------------------------- the views ---
+
+
+class KvPlanes(NamedTuple):
+ """One layer's paged K/V views, and under an NVFP4 pool its e4m3 block-scale views:
+ `[slots, H_kv, tokens_per_block, kv_width]` and `[slots, H_kv, tokens_per_block, D / 16]`."""
+
+ k: torch.Tensor
+ v: torch.Tensor
+ k_scale: torch.Tensor | None = None
+ v_scale: torch.Tensor | None = None
+
+
+class RankPools(NamedTuple):
+ """The trunk's pool views on one rank, from one manager. A view is a `TensorWrapper` over a
+ raw device address that keeps no memory alive, so it is valid only for that manager."""
+
+ pool_idx: dict # attention layer -> pool offset (what thop_attention addresses)
+ index_cache: dict # attention layer -> QSA index plane
+ position_cache: torch.Tensor # the shared QSA position plane
+ kv: dict # attention layer -> `KvPlanes`
+ conv: dict # GDN layer -> conv state pool
+ ssm: dict # GDN layer -> recurrent state pool
+ ple_state: tuple # (PLE conv state, PLE n-gram id history)
+
+
+def kv_planes(cm, g: PoolGeometry, layer: int, index_cache) -> KvPlanes:
+ """One layer's paged K/V views from `cm`, checked against the contract: the data planes, and
+ under NVFP4 the block-scale planes that go with them."""
+ pool = cm.get_buffers(layer, kv_layout="HND")
+ check_attention_pools(g, layer, index_cache, pool)
+ if not g.kv_block_scales:
+ return KvPlanes(pool[:, 0], pool[:, 1])
+ require(cm, ("get_block_scale_buffers",), "reads an NVFP4 pool's e4m3 block-scale planes")
+ scales = cm.get_block_scale_buffers(layer, kv_layout="HND")
+ check_kv_scale_planes(g, layer, pool, scales)
+ return KvPlanes(pool[:, 0], pool[:, 1], scales[:, 0], scales[:, 1])
+
+
+def take_pools(cm, g: PoolGeometry, attn_layers, gdn_layers, ple_layer: int) -> RankPools:
+ """Every trunk pool view of `cm` this rank binds, each checked against `g`.
+
+ `attn_layers` and `gdn_layers` are the full-attention and Gated-DeltaNet decoder layers, and
+ `ple_layer` the zero-based PLE layer. Reads no attention metadata: what depends on the step
+ (pool ids, windows, the QSA selection width) is the binder's to check."""
+ require(
+ cm,
+ ("get_index_k_buffer", "get_qsa_position_buffer", "get_buffers", "layer_offsets"),
+ "reads the QSA index plane, the shared position plane and the paged K/V",
+ )
+ # `ple_layer_cache` exists only on the V2 hybrid manager; under V1 the PLE state silently
+ # does not exist.
+ require(
+ cm,
+ ("get_conv_states", "get_ssm_states", "conv_section_dims", "ple_layer_cache"),
+ "reads the recurrent, conv and PLE pools",
+ )
+ # thop_attention addresses the pool by the layer's pool offset, not its ordinal among
+ # attention layers.
+ pool_idx = {i: cm.layer_offsets[i] for i in attn_layers}
+ position_cache = cm.get_qsa_position_buffer()
+ check_position_plane(g, position_cache)
+ index_cache, kv = {}, {}
+ for i in attn_layers:
+ index_cache[i] = cm.get_index_k_buffer(i)
+ kv[i] = kv_planes(cm, g, i, index_cache[i])
+
+ check_conv_layout(g, cm.conv_section_dims)
+ conv = {i: cm.get_conv_states(i) for i in gdn_layers}
+ ssm = {i: cm.get_ssm_states(i) for i in gdn_layers}
+ for i in gdn_layers:
+ check_recurrent_pools(g, conv[i], ssm[i])
+
+ ple_cache = cm.ple_layer_cache(ple_layer)
+ assert ple_cache is not None, (
+ f"no PLE lifecycle state for decoder layer {ple_layer}: the runtime allocated neither "
+ "the dilated conv state nor the n-gram id history"
+ )
+ ple_conv, ple_ngram = ple_cache
+ check_ple_pools(g, ple_conv, ple_ngram, conv[gdn_layers[0]])
+ return RankPools(pool_idx, index_cache, position_cache, kv, conv, ssm, (ple_conv, ple_ngram))
+
+
+__all__ = [
+ "GDN_REPLAY_HISTORY",
+ "MAX_VERIFY_WIDTH",
+ "NVFP4_BLOCK",
+ "QSA_POSITION_AXES",
+ "SSM_STATE_DTYPES",
+ "KvPlanes",
+ "PoolGeometry",
+ "RankPools",
+ "attention_tp_size",
+ "check_attention_pools",
+ "check_conv_layout",
+ "check_kv_scale_planes",
+ "check_ple_pools",
+ "check_position_plane",
+ "check_recurrent_pools",
+ "check_replay_history",
+ "check_replay_state",
+ "check_replay_width",
+ "check_shared_selection",
+ "check_step_buffers",
+ "check_verify_width",
+ "kv_planes",
+ "pool_geometry",
+ "require",
+ "take_pools",
+]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision.py
new file mode 100644
index 000000000000..cc2e590770f1
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision.py
@@ -0,0 +1,694 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The vision tower and the multimodal splice's host bookkeeping.
+
+The checkpoint's 27-block ViT (no KV cache, deepstack or quantization) and its 2x2 patch
+merger turn a request's images and videos into one `[rows, hidden]` embedding row per
+placeholder token, which the prefill target's `_splice` break writes over the token
+embedding. The tower is replicated: every rank holds all of its weights and runs it on the
+media of the requests it serves -- under attention TP every rank the whole batch's, and gets
+the same rows back, since every op is a deterministic function of its operands; under
+attention DP a rank its own requests'.
+
+The shell calls `multimodal_embeds` on every step, before the core's `begin_step`, and the
+executor calls `validate_request` on every new request. Neither runs inside a captured
+graph: the tower is eager and reads the request's grid on the host.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import NamedTuple
+
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import (
+ thop_attention,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.activation.gelu_erf import (
+ gelu_erf,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.activation.gelu_tanh import (
+ gelu_tanh,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.norm.layer_norm import (
+ layer_norm,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.vision.patch_embed_conv3d import (
+ patch_embed_conv3d,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.attention.vision_rope_2d_packed_qk import ( # noqa: E501
+ vision_rope_2d_packed_qk,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.vision.triton_pos_embed_interpolate import ( # noqa: E501
+ triton_pos_embed_interpolate,
+)
+from tensorrt_llm.inputs.multimodal import get_multimodal_item_metadata
+
+#: Tower rope base (`vision_config` has no rope field); unrelated to the text 1e7.
+VISION_ROPE_THETA = 10000.0
+
+#: LayerNorm epsilon of `norm1`, `norm2`, `merger.norm`: a literal in the reference.
+VISION_NORM_EPS = 1e-6
+
+#: The most patch rows one tower call carries, the ceiling the vision entries' cells reach:
+#: 8192 placeholder tokens x `spatial_merge_size`**2 patch rows each. `VisionTower` splits a
+#: request's frames into calls under it, so the bound a request meets is per frame: one
+#: frame's `h * w` patches.
+VISION_CALL_MAX_PATCHES = 32768
+
+#: `mask_type` 0 = padding mask: bidirectional within a batch segment. Causal (1) would make
+#: a patch depend on raster order.
+_VISION_MASK_PADDING = 0
+
+# Per-call constants of the tower's attention: the encoder-only mode of `thop_attention`.
+# Omitting `kv_cache_block_offsets`, `host_kv_cache_pool_pointers` and
+# `host_kv_cache_pool_mapping` selects that mode. `update_kv_cache=False` is rejected ("KV
+# cache update cannot be disabled now"), so it stays True.
+_VISION_CALL_CONSTANTS = dict(
+ k=None,
+ v=None,
+ output_sf=None,
+ kv_cache_block_offsets=None,
+ host_kv_cache_pool_pointers=None,
+ host_kv_cache_pool_mapping=None,
+ tokens_per_block=None,
+ use_paged_context_fmha=False,
+ is_fused_qkv=True,
+ update_kv_cache=True,
+ predicted_tokens_per_seq=1,
+ # No pool, so no layer to address.
+ local_layer_idx=0,
+ max_context_q_len_override=None,
+ cache_indirection=None,
+ kv_scale_orig_quant=None,
+ kv_scale_quant_orig=None,
+ out_scale=None,
+ quant_mode=0,
+ beam_width=1,
+ mask_type=_VISION_MASK_PADDING,
+ q_scaling=1.0,
+ # Rotation is applied before the op (`vision_rope_2d_packed_qk`); rope_* inert.
+ position_embedding_type=0,
+ rotary_inv_freq=None,
+ rotary_cos_sin=None,
+ rope_dim=0,
+ rope_base=10000.0,
+ rope_scale_type=0,
+ rope_scale=1.0,
+ rope_short_m_scale=1.0,
+ rope_long_m_scale=1.0,
+ rope_max_positions=1024,
+ rope_original_max_positions=1024,
+ attention_input_type=0,
+ is_mla_enable=False,
+ chunked_prefill_buffer_batch_size=1,
+ attention_sinks=None,
+ block_ids_per_seq=None,
+ latent_cache=None,
+ q_pe=None,
+ q_lora_rank=None,
+ kv_lora_rank=None,
+ qk_nope_head_dim=None,
+ qk_rope_head_dim=None,
+ v_head_dim=None,
+ rope_append=None,
+ attention_chunk_size=None,
+ softmax_stats_tensor=None,
+ is_spec_decoding_enabled=False,
+ use_spec_decoding=False,
+ is_spec_dec_tree=False,
+ spec_decoding_generation_lengths=None,
+ spec_decoding_position_offsets=None,
+ spec_decoding_packed_mask=None,
+ spec_decoding_bl_tree_mask_offset=None,
+ spec_decoding_bl_tree_mask=None,
+ spec_bl_tree_first_sparse_mask_offset_kv=None,
+ sparse_kv_indices=None,
+ sparse_kv_offsets=None,
+ sparse_attn_indices=None,
+ sparse_attn_offsets=None,
+ sparse_attn_indices_block_size=0,
+ mrope_rotary_cos_sin=None,
+ mrope_position_deltas=None,
+ helix_position_offsets=None,
+ helix_is_inactive_rank=None,
+ is_cross=False,
+ cross_kv=None,
+ relative_attention_bias=None,
+ relative_attention_max_distance=0,
+ kv_norm_weight=None,
+ kv_norm_eps=1e-6,
+ skip_correction_threshold=0.0,
+)
+
+#: Per modality, the keys under `multimodal_data[modality]`: the flat patch rows and the
+#: host-side `(t, h, w)` grid.
+MEDIA_KEYS = {
+ "image": ("pixel_values", "image_grid_thw"),
+ "video": ("pixel_values_videos", "video_grid_thw"),
+}
+
+#: Tensors under `multimodal_data` the executor moves to the device. Everything else stays
+#: on the host: the grids are read there, and moved they would cost a synchronization each.
+DEVICE_PATHS = (
+ "image.pixel_values",
+ "video.pixel_values_videos",
+ "mrope_config.mrope_position_ids",
+ "mrope_config.mrope_position_deltas",
+)
+
+#: The position-ids axes of the mRoPE layout a media request carries.
+_POSITION_AXES = 3
+
+
+# --------------------------------------------------------------- weight views ---
+
+
+class _VisionBlockWeights(NamedTuple):
+ """One vision block's weights (`.t()` GEMM views). LayerNorm has a bias, unlike every
+ norm in the language tower."""
+
+ n1_w: torch.Tensor
+ n1_b: torch.Tensor
+ qkv: torch.Tensor
+ qkv_b: torch.Tensor
+ proj: torch.Tensor
+ proj_b: torch.Tensor
+ n2_w: torch.Tensor
+ n2_b: torch.Tensor
+ fc1: torch.Tensor
+ fc1_b: torch.Tensor
+ fc2: torch.Tensor
+ fc2_b: torch.Tensor
+
+
+class _VisionWeights(NamedTuple):
+ """The tower's weights: patch embedding, position table, blocks and patch merger."""
+
+ patch_w: torch.Tensor
+ patch_b: torch.Tensor
+ pos: torch.Tensor
+ blocks: list
+ m_n_w: torch.Tensor
+ m_n_b: torch.Tensor
+ m_fc1: torch.Tensor
+ m_fc1_b: torch.Tensor
+ m_fc2: torch.Tensor
+ m_fc2_b: torch.Tensor
+
+
+def _vision_views(w: dict, depth: int) -> _VisionWeights:
+ """The weight table's `vis_*` entries (`weights.py`) as the views the ops read."""
+ return _VisionWeights(
+ # Not transposed: the entry takes the checkpoint's `[out, in, kt, kh, kw]`.
+ patch_w=w["vis_patch_w"],
+ patch_b=w["vis_patch_b"],
+ pos=w["vis_pos"],
+ blocks=[
+ _VisionBlockWeights(
+ n1_w=w[f"l{b}_vis_n1_w"],
+ n1_b=w[f"l{b}_vis_n1_b"],
+ # Row layout (q heads, k, v) is already thop_attention's packed form.
+ qkv=w[f"l{b}_vis_qkv_w"].t(),
+ qkv_b=w[f"l{b}_vis_qkv_b"],
+ proj=w[f"l{b}_vis_proj_w"].t(),
+ proj_b=w[f"l{b}_vis_proj_b"],
+ n2_w=w[f"l{b}_vis_n2_w"],
+ n2_b=w[f"l{b}_vis_n2_b"],
+ fc1=w[f"l{b}_vis_fc1_w"].t(),
+ fc1_b=w[f"l{b}_vis_fc1_b"],
+ fc2=w[f"l{b}_vis_fc2_w"].t(),
+ fc2_b=w[f"l{b}_vis_fc2_b"],
+ )
+ for b in range(depth)
+ ],
+ m_n_w=w["vis_m_n_w"],
+ m_n_b=w["vis_m_n_b"],
+ m_fc1=w["vis_m_fc1_w"].t(),
+ m_fc1_b=w["vis_m_fc1_b"],
+ m_fc2=w["vis_m_fc2_w"].t(),
+ m_fc2_b=w["vis_m_fc2_b"],
+ )
+
+
+# ------------------------------------------------------------------- the tower ---
+
+
+def _frame_calls(
+ grid: list[tuple[int, int, int]], max_rows: int
+) -> list[list[tuple[int, int, int]]]:
+ """Cut `grid` (one `(frames, height, width)` row per image or video) into tower calls of
+ whole frames, in order, each at most `max_rows` patch rows. A call is its grid rows: the
+ frames of one item that fall in it are one `(frames, height, width)` row."""
+ calls: list[list[tuple[int, int, int]]] = [[]]
+ room = max_rows
+ for frames, height, width in grid:
+ rows = height * width
+ assert rows <= max_rows, (
+ f"one frame is {height}x{width} = {rows} patch rows; the vision ops are "
+ f"certified up to {max_rows} rows in one call"
+ )
+ while frames:
+ take = min(frames, room // rows)
+ if take == 0:
+ calls.append([])
+ room = max_rows
+ continue
+ calls[-1].append((take, height, width))
+ room -= take * rows
+ frames -= take
+ return calls
+
+
+def _merge_order_positions(grid: list[tuple[int, int, int]], merge: int) -> torch.Tensor:
+ """Each patch row's `(row, col)` in its frame, int64 `[rows, 2]` on the host, in
+ spatial-merge order: whole `merge x merge` blocks row-major, row-major inside a block,
+ the order the patch merger's 4-to-1 reshape consumes. A frame group repeats its frame."""
+ intra = torch.arange(merge)
+ parts = []
+ for frames, height, width in grid:
+ blocks_h, blocks_w = height // merge, width // merge
+ shape = (blocks_h, blocks_w, merge, merge)
+ rows = torch.arange(blocks_h)[:, None, None, None] * merge + intra[None, None, :, None]
+ cols = torch.arange(blocks_w)[None, :, None, None] * merge + intra[None, None, None, :]
+ coords = torch.stack((rows.expand(shape).reshape(-1), cols.expand(shape).reshape(-1)), 1)
+ parts.append(coords.repeat(frames, 1) if frames > 1 else coords)
+ return parts[0] if len(parts) == 1 else torch.cat(parts, 0)
+
+
+def _rope_tables(
+ grid: list[tuple[int, int, int]], merge: int, theta: float, head_dim: int, device
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """The 2-D rotary's per-row fp32 `cos` / `sin`, `[rows, head_dim // 2]` each and
+ contiguous: the height angles, then the width angles, as `vision_rope_2d_packed_qk`
+ takes them. Built once per tower call and shared by every block.
+
+ Each axis owns half of the `head_dim // 2` rotation pairs. The angles are formed in
+ fp64 and rounded once to fp32: the largest is `(side - 1)` radians, and fp32 throughout
+ rounds twice there."""
+ rope_dim = head_dim // 2
+ inv_freq = 1.0 / (theta ** (torch.arange(0, rope_dim, 2, dtype=torch.float64) / rope_dim))
+ pos = _merge_order_positions(grid, merge).to(torch.float64)
+ angles = (pos[:, :, None] * inv_freq).flatten(1).to(device=device, dtype=torch.float32)
+ return angles.cos().contiguous(), angles.sin().contiguous()
+
+
+class VisionTower:
+ """One request's images or videos through the 27-block ViT and its patch merger.
+
+ Not a target: runs only for a context request carrying media, from the shell, before
+ the step. Attention is bidirectional per frame (a video's frame group is a frame),
+ expressed to `thop_attention` as one context "sequence" per frame; running the rows as
+ one sequence is far outside tolerance. The position-embedding interpolation and the 2-D
+ rope both emit in spatial-merge order (what the merger's 4-to-1 reshape consumes); a
+ raster-order emitter gives a correctly shaped, silently wrong tower.
+
+ The frames run in calls of at most `VISION_CALL_MAX_PATCHES` rows, cut only between
+ frames: attention is per frame, the position embedding and rope are per grid row and
+ repeat per frame, the norms and MLPs are row-wise, and a frame's `h * w` is a multiple
+ of `merge**2`, so a cut keeps every merger block whole. One frame past the bound cannot
+ be served; `validate_request` refuses it at admission."""
+
+ def __init__(self, w: dict, vision_config) -> None:
+ """`w` is the core's weight table after load; `vision_config` the composite
+ config's vision half, which `routing.py` matched against this checkpoint's."""
+ vcfg = vision_config
+ self.depth = vcfg.depth
+ self.hidden = vcfg.hidden_size
+ self.heads = vcfg.num_heads
+ self.head_dim = self.hidden // self.heads
+ self.merge = vcfg.spatial_merge_size
+ self.patch = vcfg.patch_size
+ # One flat pixel row per patch, `(in_channels, kt, kh, kw)` channel-slowest.
+ self.pixel_width = vcfg.in_channels * vcfg.temporal_patch_size * self.patch * self.patch
+ # `spatial_merge_size ** 2` patches become one token at the merger.
+ self.merge_dim = self.hidden * self.merge * self.merge
+ # The position table is a square grid of this side.
+ self.grid_side = math.isqrt(vcfg.num_position_embeddings)
+ self.dtype = w["embed"].dtype
+ self._w = _vision_views(w, self.depth)
+ # Persistent scratch: `thop_attention` grows `workspace` in place; never the
+ # engine's own workspace.
+ self._ws: torch.Tensor | None = None
+
+ def __call__(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor:
+ """`pixel_values` is `[sum(t*h*w), in_channels*kt*kh*kw]`, one flat patch per row,
+ channel-slowest. `grid_thw` is one `(frames, height, width)` row per image or video
+ and stays on the host (no synchronization). Returns
+ `[sum(t*h*w) / merge**2, out_hidden_size]`, one row per placeholder, in grid row
+ order."""
+ assert pixel_values.dtype == self.dtype, (
+ f"the tower runs at the activation dtype {self.dtype}; the runtime handed "
+ f"pixel_values as {pixel_values.dtype}. The input processor's `dtype` property is "
+ "what decides this, and no cast happens here"
+ )
+ assert pixel_values.dim() == 2 and pixel_values.shape[1] == self.pixel_width, (
+ f"pixel_values must be [patches, {self.pixel_width}] flat patch rows; got "
+ f"{tuple(pixel_values.shape)}"
+ )
+ grid: list[tuple[int, int, int]] = []
+ for row in grid_thw.tolist():
+ assert len(row) == 3, f"grid_thw rows are (t, h, w); got {row}"
+ frames, height, width = (int(v) for v in row)
+ assert frames >= 1 and height >= 1 and width >= 1, row
+ assert height % self.merge == 0 and width % self.merge == 0, (
+ f"grid sides must divide by spatial_merge_size {self.merge}; got "
+ f"height={height}, width={width}"
+ )
+ grid.append((frames, height, width))
+ assert grid, "grid_thw describes no image or video"
+ described = sum(t * h * w for t, h, w in grid)
+ assert described == pixel_values.shape[0], (
+ f"grid_thw describes {described} patch rows and pixel_values carries "
+ f"{pixel_values.shape[0]}"
+ )
+
+ outputs = []
+ start = 0
+ for call in _frame_calls(grid, VISION_CALL_MAX_PATCHES):
+ rows = sum(t * h * w for t, h, w in call)
+ outputs.append(self._run(pixel_values[start : start + rows], call))
+ start += rows
+ return outputs[0] if len(outputs) == 1 else torch.cat(outputs, 0)
+
+ def _run(self, pixel_values: torch.Tensor, grid: list[tuple[int, int, int]]) -> torch.Tensor:
+ """One tower call over whole frames: `grid` describes exactly the rows of
+ `pixel_values`, at most `VISION_CALL_MAX_PATCHES` of them."""
+ dt, dev = pixel_values.dtype, pixel_values.device
+ assert pixel_values.is_cuda, "pixel_values must reach the tower on the device"
+
+ # ---- the frame segmentation, host-side ------------------------------
+ # One attention segment per frame: a video's frames do not attend to each other.
+ seg_lens: list[int] = []
+ for frames, height, width in grid:
+ seg_lens.extend([height * width] * frames)
+ num_patches = pixel_values.shape[0]
+ num_segments = len(seg_lens)
+ max_seg = max(seg_lens)
+
+ # Batch-state tensors for the encoder-only call; they describe the frames.
+ host_lens = torch.tensor(seg_lens, dtype=torch.int32)
+ host_types = torch.zeros([num_segments], dtype=torch.int32)
+ host_total = torch.tensor([num_patches, 0], dtype=torch.int32)
+ dev_lens = host_lens.to(dev)
+ if self._ws is None:
+ self._ws = torch.empty([0], dtype=torch.int8, device=dev)
+ # The rotary's angles, once for all blocks.
+ cos, sin = _rope_tables(grid, self.merge, VISION_ROPE_THETA, self.head_dim, dev)
+
+ # ---- patch embedding + absolute position ----------------------------
+ V = self._w
+ pos = [
+ triton_pos_embed_interpolate(V.pos, t, h, w, self.grid_side, self.merge, dt)
+ for t, h, w in grid
+ ]
+ h = torch.add(
+ patch_embed_conv3d(pixel_values, V.patch_w, V.patch_b),
+ pos[0] if len(pos) == 1 else torch.cat(pos, 0),
+ )
+ attn_out = torch.empty([num_patches, self.hidden], dtype=dt, device=dev)
+
+ for B in V.blocks:
+ qkv = cublas_mm(layer_norm(h, B.n1_w, B.n1_b, VISION_NORM_EPS), B.qkv, B.qkv_b)
+ # The q|k prefix of each packed row, rotated in place; v is never addressed.
+ qkv_heads = torch.reshape(qkv, [num_patches, 3 * self.heads, self.head_dim])
+ vision_rope_2d_packed_qk(qkv_heads[:, : 2 * self.heads], cos, sin)
+ attn_out = thop_attention(
+ q=qkv,
+ output=attn_out,
+ workspace=self._ws,
+ sequence_length=dev_lens,
+ host_past_key_value_lengths=host_lens,
+ host_total_kv_lens=host_total,
+ context_lengths=dev_lens,
+ host_context_lengths=host_lens,
+ host_request_types=host_types,
+ max_num_requests=num_segments,
+ max_context_length=max_seg,
+ max_seq_len=max_seg,
+ attention_window_size=max_seg,
+ num_contexts=num_segments,
+ num_ctx_tokens=num_patches,
+ num_heads=self.heads,
+ num_kv_heads=self.heads,
+ head_size=self.head_dim,
+ **_VISION_CALL_CONSTANTS,
+ )
+ h = torch.add(h, cublas_mm(attn_out, B.proj, B.proj_b))
+ mlp = cublas_mm(layer_norm(h, B.n2_w, B.n2_b, VISION_NORM_EPS), B.fc1, B.fc1_b)
+ # Non-gated MLP with the tanh GELU (not the merger's erf).
+ h = torch.add(h, cublas_mm(gelu_tanh(mlp), B.fc2, B.fc2_b))
+
+ # ---- the patch merger -----------------------------------------------
+ # The norm runs over `hidden` before the 4-to-1 shuffle; the reshape is legal only
+ # because both order-bearing steps emitted in spatial-merge order.
+ merged = torch.reshape(
+ layer_norm(h, V.m_n_w, V.m_n_b, VISION_NORM_EPS), [-1, self.merge_dim]
+ )
+ # The erf GELU (`nn.GELU()`); within one bf16 ulp of tanh, hence a separate entry.
+ return cublas_mm(gelu_erf(cublas_mm(merged, V.m_fc1, V.m_fc1_b)), V.m_fc2, V.m_fc2_b)
+
+
+# ------------------------------------------------------------ the shell's side ---
+
+
+def multimodal_embeds(tower: VisionTower | None, kwargs: dict):
+ """Run the vision tower for this batch; returns `(rows, mm_token_indices)` (the tower's
+ `[R, hidden]` rows in packed-batch order and the `[R]` placeholder rows they replace,
+ which the prefill target splices into its embedding), or None when the step carries no
+ media.
+
+ Called by the shell, because the inherited `DecoderModelForCausalLM.forward` drops
+ `multimodal_params` from `**kwargs` on the no-speculation path. Host-side bookkeeping:
+ `multimodal_params` is a subset of the context requests, so never zip it positionally;
+ `multimodal_runtime` (chunked prefill / KV reuse only) gives the cached and in-chunk row
+ counts over the request's images and videos together, in prompt order; `mm_token_indices`
+ names every placeholder row.
+
+ The tower runs once per modality a request carries, and `in_prompt_order` merges the
+ two outputs. Positions are not derived here: every coordinate of a vision span, a
+ video's temporal one included, is in the payload's `mrope_config.mrope_position_ids`,
+ which the engine slices into `position_ids`."""
+ mm_params = kwargs.get("multimodal_params") or []
+ if not mm_params:
+ return None
+ rows = []
+ for params in mm_params:
+ data = getattr(params, "multimodal_data", None) or {}
+ image, video = data.get("image"), data.get("video")
+ if image is None and video is None:
+ continue
+ mrope_ids = (data.get("mrope_config") or {}).get("mrope_position_ids")
+ assert (
+ mrope_ids is not None and mrope_ids.dim() == 3 and mrope_ids.shape[0] == _POSITION_AXES
+ ), (
+ f"a media request needs mrope_config.mrope_position_ids of shape "
+ f"[{_POSITION_AXES}, 1, L]; got {None if mrope_ids is None else tuple(mrope_ids.shape)}"
+ )
+ runtime = getattr(params, "multimodal_runtime", None)
+ if runtime is not None and int(runtime.num_mm_tokens_in_chunk) == 0:
+ # Every placeholder row of this request lies in another chunk.
+ continue
+ assert tower is not None, "post_load_weights builds the vision tower before any step"
+ towers = {}
+ if image is not None:
+ towers["image"] = run_tower(tower, image, *MEDIA_KEYS["image"])
+ if video is not None:
+ towers["video"] = run_tower(tower, video, *MEDIA_KEYS["video"])
+ embeds = in_prompt_order(tower.merge, params, towers)
+ if runtime is not None:
+ # Chunked prefill or reused prefix: splice only this chunk's slice.
+ start = int(runtime.num_cached_mm_tokens)
+ count = int(runtime.num_mm_tokens_in_chunk)
+ total = embeds.shape[0]
+ assert int(runtime.total_embeds_in_request) == total, (
+ f"the executor counted {runtime.total_embeds_in_request} placeholder rows for "
+ f"this request and the tower produced {total} over its {sorted(towers)}; the "
+ "two are derived from the same grids and a disagreement means the merger's "
+ "row count is wrong"
+ )
+ assert 0 <= start <= start + count <= total, (start, count, total)
+ embeds = embeds[start : start + count]
+ rows.append(embeds)
+ if not rows:
+ return None
+ vision = rows[0] if len(rows) == 1 else torch.cat(rows, 0)
+ mm_token_indices = kwargs.get("mm_token_indices")
+ assert mm_token_indices is not None, (
+ "the executor shipped multimodal params but no mm_token_indices; that pairing is what "
+ "the shell's `multimodal_token_ids` buys, and without it there is no sync-free way to "
+ "find the placeholder rows"
+ )
+ assert mm_token_indices.shape[0] == vision.shape[0], (
+ f"the batch has {mm_token_indices.shape[0]} placeholder rows and the tower produced "
+ f"{vision.shape[0]} embedding rows"
+ )
+ return vision, mm_token_indices
+
+
+def run_tower(tower, media: dict, pixel_key: str, grid_key: str):
+ """The tower over one modality's payload: `(rows, grid_thw)`, the rows one per
+ placeholder, item by item in the order the processor emitted them."""
+ pixel_values, grid_thw = media.get(pixel_key), media.get(grid_key)
+ assert pixel_values is not None and grid_thw is not None, (
+ f"the multimodal payload carries {sorted(media)}; this target reads {pixel_key!r} and "
+ f"{grid_key!r} from it"
+ )
+ assert not grid_thw.is_cuda, f"{grid_key} must stay on the host; see DEVICE_PATHS"
+ return tower(pixel_values, grid_thw), grid_thw
+
+
+def in_prompt_order(merge: int, params, towers: dict) -> torch.Tensor:
+ """One request's tower rows in the order its placeholders appear.
+
+ `towers` maps each modality the request carries to `(rows, grid_thw)`; each grid row is
+ one item of `t*h*w / merge**2` rows. One modality is already in prompt order. Two are
+ reassembled item by item from `item_metadata.item_refs` (`(modality, index)` per item in
+ prompt order, `index` counting that modality's items); without it there is no host-side
+ way to tell which placeholder run is an image and which a video."""
+ if len(towers) == 1:
+ return next(iter(towers.values()))[0]
+ metadata = params.item_metadata
+ order = metadata.item_refs if metadata is not None else None
+ assert order, (
+ f"a request carrying {sorted(towers)} needs item metadata to splice them in prompt "
+ "order; the payload holds each modality's rows apart"
+ )
+ merge_unit = merge * merge
+ items = {}
+ for modality, (embeds, grid_thw) in towers.items():
+ counts = [t * h * w // merge_unit for t, h, w in grid_thw.tolist()]
+ assert sum(counts) == embeds.shape[0], (
+ f"the {modality} grid describes {sum(counts)} rows and the tower produced "
+ f"{embeds.shape[0]}"
+ )
+ items[modality] = torch.split(embeds, counts, 0)
+ listed = {modality: [] for modality in items}
+ for modality, index in order:
+ assert modality in items, (
+ f"item_refs lists a {modality!r} item; this request carries {sorted(items)}"
+ )
+ listed[modality].append(int(index))
+ for modality, parts in items.items():
+ assert sorted(listed[modality]) == list(range(len(parts))), (
+ f"item_refs lists {modality} items {listed[modality]}; the payload carries {len(parts)}"
+ )
+ return torch.cat([items[modality][int(index)] for modality, index in order], 0)
+
+
+# ---------------------------------------------------------------- admission ---
+
+
+def validate_request(tower, request) -> None:
+ """Admission check, run by the executor on every new request. A media request the
+ forward would assert on fails here as that request's own error; an assert inside the
+ step ends the executor under the overlap scheduler. Reads only the request's host-side
+ payload, and a request without media returns at once. `tower` supplies the geometry
+ (`merge`, `patch`, `pixel_width`).
+
+ Refused: a frame past `VISION_CALL_MAX_PATCHES` patches (the one bound the tower keeps),
+ a malformed grid or pixel payload, media whose rows disagree with the prompt's
+ placeholder count, missing or mis-sized `mrope_position_ids`, and a request with both
+ images and videos whose item metadata is missing or does not list its items."""
+ data = request.py_multimodal_data
+ if not data:
+ return
+ grids = {}
+ for modality, (pixel_key, grid_key) in MEDIA_KEYS.items():
+ media = data.get(modality)
+ if media is not None:
+ grids[modality] = _media_grid(tower, modality, media, pixel_key, grid_key)
+ if not grids:
+ return
+
+ mrope_ids = (data.get("mrope_config") or {}).get("mrope_position_ids")
+ if (
+ mrope_ids is None
+ or mrope_ids.dim() != 3
+ or mrope_ids.shape[0] != _POSITION_AXES
+ or mrope_ids.shape[2] != request.prompt_len
+ ):
+ raise ValueError(
+ f"a media request needs the [{_POSITION_AXES}, 1, {request.prompt_len}] "
+ "mrope_config.mrope_position_ids from the input processor, one position per prompt "
+ f"token; got {None if mrope_ids is None else tuple(mrope_ids.shape)}"
+ )
+
+ if len(grids) > 1:
+ metadata = get_multimodal_item_metadata(data)
+ _check_item_order(grids, metadata.item_refs if metadata is not None else None)
+
+ cumsum = data.get("multimodal_embed_mask_cumsum")
+ if cumsum is not None:
+ merge_unit = tower.merge * tower.merge
+ rows = sum(t * h * w for grid in grids.values() for t, h, w in grid) // merge_unit
+ placeholders = int(cumsum[-1])
+ if placeholders != rows:
+ raise ValueError(
+ f"the prompt holds {placeholders} image/video placeholder tokens and its "
+ f"{' and '.join(sorted(grids))} payload produces {rows} embedding rows"
+ )
+
+
+def _media_grid(tower, modality, media, pixel_key, grid_key):
+ """One modality's `(frames, height, width)` rows, checked on the host."""
+ pixel_values, grid_thw = media.get(pixel_key), media.get(grid_key)
+ if pixel_values is None or grid_thw is None:
+ raise ValueError(
+ f"the {modality} payload carries {sorted(media)}; this target reads {pixel_key!r} "
+ f"and {grid_key!r} from it"
+ )
+ if grid_thw.dim() != 2 or grid_thw.shape[0] == 0 or grid_thw.shape[1] != 3:
+ raise ValueError(
+ f"{grid_key} must hold one (t, h, w) row per {modality}; got shape "
+ f"{tuple(grid_thw.shape)}"
+ )
+ merge = tower.merge
+ grid = [tuple(int(v) for v in row) for row in grid_thw.tolist()]
+ for item, (frames, height, width) in enumerate(grid):
+ if min(frames, height, width) < 1 or height % merge or width % merge:
+ raise ValueError(
+ f"{modality} {item} has grid {(frames, height, width)}; every entry must be "
+ f"positive and the sides divisible by spatial_merge_size {merge}"
+ )
+ if height * width > VISION_CALL_MAX_PATCHES:
+ raise ValueError(
+ f"{modality} {item} has frames of {height}x{width} = {height * width} patches "
+ f"({height * tower.patch}x{width * tower.patch} pixels); the vision tower "
+ f"serves frames of at most {VISION_CALL_MAX_PATCHES} patches, so send it at a "
+ "lower resolution"
+ )
+ patches = sum(t * h * w for t, h, w in grid)
+ if tuple(pixel_values.shape) != (patches, tower.pixel_width):
+ raise ValueError(
+ f"{pixel_key} is {tuple(pixel_values.shape)}; its grids describe {patches} patch "
+ f"rows of {tower.pixel_width}"
+ )
+ return grid
+
+
+def _check_item_order(grids, order) -> None:
+ """A request with several modalities lists its items in prompt order."""
+ if not order:
+ raise ValueError(
+ f"a request carrying {sorted(grids)} needs item metadata to splice them in prompt "
+ "order; the payload holds each modality's rows apart"
+ )
+ listed = {modality: [] for modality in grids}
+ for entry in order:
+ try:
+ modality, index = entry
+ listed[modality].append(int(index))
+ except (KeyError, TypeError, ValueError) as error:
+ raise ValueError(
+ f"item_refs entry {entry!r} is not an item of this request, which carries "
+ f"{sorted(grids)}; entries are (modality, index) pairs"
+ ) from error
+ for modality, grid in grids.items():
+ if sorted(listed[modality]) != list(range(len(grid))):
+ raise ValueError(
+ f"item_refs lists {modality} items {listed[modality]}; the payload carries "
+ f"{len(grid)}"
+ )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/weights.py
new file mode 100644
index 000000000000..ffa1185ca5e0
--- /dev/null
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/weights.py
@@ -0,0 +1,1861 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Weight table and loader: Qwen3.8-Flash-Next-NVFP4 (hf-fc694b5_orig) / sm_103 / default.
+
+One `W` per weight role (shape, checkpoint source, the rank's cut, dtype, load-time
+transform); `Qwen4ExpDefaultWeights` unrolls the table into declaration, manifest and load.
+Storage is HF `[out, in]`; GEMM views, grouped-norm slices, fp32 gating params and NVFP4
+scalars are derived in `modeling.py` after load.
+
+**Every weight is declared at this rank's share.** The shares are computed once, in
+`rank_dims`, from the engine's `Mapping`:
+
+* **Attention heads** are split over the attention TP group: the TP group, or one rank
+ under attention DP. Rank `r` of four holds q heads `[6r, 6r + 6)`: their rows of
+ `q_proj` (q and the output gate interleave per head, so a head's `2 * head_dim` rows
+ are contiguous) and their input columns of `o_proj`. The two KV heads are repeated on
+ neighbouring ranks: rank `r` holds head `r * 2 // 4`, the head `duplicate_kv_weight`
+ gives it and the K/V pool it binds holds (`pool.pool_geometry`). The QSA indexer and
+ the q/k norms are replicated. The MTP layer's attention is cut the same way.
+* **Gated-DeltaNet heads** are split over `mamba_effective_tp_size` (the TP group, or
+ one rank under attention DP), as the recurrent pools are: rank `r` holds key heads
+ `[4r, 4r + 4)` and value heads `[12r, 12r + 12)`. Of `in_proj_qkv` and the depthwise
+ `conv1d` it reads three row blocks, the q and k rows of its key heads and the v rows
+ of its value heads, into one `[q | k | v]` operand: the channel order of the conv pool
+ (`pool.check_conv_layout`). `in_proj_z`, `in_proj_a`, `in_proj_b`, `A_log` and
+ `dt_bias` are cut by value head, `out_proj` by its value-head input columns.
+* **Routed experts** are split over the MoE EP group (tep4, dep4): rank `r` holds the
+ contiguous window `[128r, 128r + 128)`, NVFP4 codes and block scales. tp4 splits each
+ expert's intermediate instead: rank `r` holds rows `[160r, 160r + 160)` of every expert's
+ `up_proj` and `gate_proj` and the same columns of its `down_proj`, all 512 experts. 160
+ is short of the layout the runner's operands need (the FC1 row shuffle and scale swizzle
+ work in 128-row blocks), so each half is zero-extended to `expert_inter` = 256 rows (FC2
+ to 256 columns) with zero codes and zero block scales before the relayout: SwiGLU of a
+ zero row pair is zero and a zero FC2 column adds nothing, so the padding is exact.
+* **The MTP layer's FP8 block-scale experts** are split by expert over the whole TP
+ group at every member, `[128r, 128r + 128)`: the FP8 runner takes `H` and `I` in
+ multiples of 128, which 640 / 4 is not.
+
+Everything else is replicated: hyper-connections, the PLE block and its n-gram table,
+the router, the shared expert, the embedding, the vision tower, the MTP layer's front
+end, router, shared expert and mixer, and every routed expert's per-tensor scalars over
+the whole routing space (one activation quantization feeds all 512 experts, so the core
+checks their FC1 input scale across all of them, not across its window). `lm_head` is
+the shell's: `load` reads the rows the shell built, vocab-parallel under attention TP,
+whole under attention DP.
+
+**A rank reads only its share.** Each split weight carries a cut, the rows or columns of
+the checkpoint tensor this rank holds, applied to the checkpoint entry before it is
+materialized: on a lazy safetensors slice that reads those bytes and nothing else, on an
+mmapped tensor it is a view, and either way only the share is copied to the device. An
+off-window expert is never read. The kernel relayouts then run on the share, never
+before the cut: the 32-row block shuffle and the 128x4 scale swizzle permute within
+blocks of rows, so cutting after them would cut through blocks.
+
+Transforms:
+
+* trunk attention input projection: `q_proj`, `k_proj`, `v_proj` and the indexer's
+ `index_qk_proj` read the same activation, so they are one
+ `[q | k | v | gate | index | pad, hidden]` operand and one GEMM (`_attn_in_w`), at this
+ rank's heads: 4352 rows at attention TP, 14080 under attention DP. The q halves go
+ first so the GEMM output's leading columns are the packed `[q | k | v]` the fused
+ norm+rope rewrites in place, then the gate halves, the replicated index projection and
+ `index_head_dim` zero rows (a row is a whole number of `head_dim` heads). The draft
+ (MTP) layer keeps its separate projections;
+* fused shared-expert gate|up: `[gate ; up ; gate logit ; pad]` (`_cat_rows`);
+* NVFP4 routed-expert stacks, laid out for the trunk's routed backend
+ (`trunk_moe_backend`): `[up ; gate]`, then 128x4 scale swizzle. trtllm-gen also
+ interleaves FC1 row by row and 32-row block shuffles weight and scale bytes (FC2:
+ shuffle only); CuteDSL interleaves FC1 in 64-row up/gate groups and leaves FC2 in
+ checkpoint row order. Under trtllm-gen a second copy of the window in the CuteDSL
+ layout (`fc*_cd`) is declared when the core keeps it (`moe_prefill_copy`);
+* n-gram table: 128 row chunks concatenated into one `[rows, 160]` buffer (the hash
+ emits global row indices). It is replicated, and host-resident by default
+ (`_host_resident_empty`): one copy per rank.
+
+The n-gram table is not in the table: its height is a checkpoint fact, so `load`
+registers it with `torch.empty` (`_declare_ngram_table`) and adds its rows to the
+manifest (`ngram_table_sources`).
+
+Not loaded: `mtp.*` when there is no `speculative_config`. With one, the MTP module is
+loaded at this rank's share, including the draft layer's own QSA indexer. MTP experts
+are plain row-major FP8 128x128: `[up ; gate]` concat, bf16 `weight_scale_inv` cast to
+fp32; `mtp_experts_to_nvfp4` requantizes a rank's window to NVFP4 after load.
+
+`load(model, weights)` fills every declared parameter once and asserts coverage both
+ways: every parameter written, and every checkpoint key consumed (a key read through a
+cut is consumed) or predicted by `expected_unconsumed` -- the experts outside this
+rank's windows, and the MTP module when there is no speculation.
+"""
+
+import os
+from dataclasses import dataclass
+from types import SimpleNamespace
+from typing import Any, Callable
+
+import torch
+from cuda.bindings import runtime as cudart
+from torch import nn
+
+from tensorrt_llm._torch.pyexecutor.config_utils import mamba_effective_tp_size
+from tensorrt_llm._utils import CUASSERT
+
+from .pool import attention_tp_size
+
+#: `layers=ALL_LAYERS`, the default: the weight is on every layer.
+ALL_LAYERS = object()
+
+#: An index into a checkpoint tensor: a tuple of slices, one per leading axis.
+Cut = tuple[slice, ...]
+
+
+@dataclass(frozen=True)
+class W:
+ """One weight role: what it is, where its bytes come from, what happens on the way.
+
+ `shape` takes the dimension bundle and returns the declared shape: this rank's
+ share, every width derived from the config and the mapping.
+
+ `src` says where the bytes come from, in one of two forms:
+
+ * a checkpoint key template, with `{p}` standing for the layer prefix (a key with no
+ `{p}` is itself);
+ * a callable `(d, i) -> [(key, index), ...]` or `[(key, index, cut), ...]`, taking
+ the dimension bundle and the layer index. `index` addresses the destination
+ parameter (None for all of it). A `key` may itself be a tuple, meaning the
+ transform is handed those tensors together. A source that names its own `cut`
+ reads one block of a tensor several of whose blocks land in the parameter.
+
+ `cut` takes the dimension bundle and returns this rank's rows or columns of the
+ checkpoint tensor, read from the entry before it is materialized; for a tuple key
+ it returns one cut (or None) per key. `None` reads the tensor whole.
+
+ `layers` says which layers carry this weight: `ALL_LAYERS` for every one,
+ `None` for a weight that is not per-layer at all, or a callable
+ `d -> iterable[int]` when only some layers have it.
+
+ `transform` runs on the cut tensor(s) before the copy, on the destination's
+ device, with the dimension bundle: the kernels' relayouts, with no slicing of
+ their own. `None` means the cut is already the declared shape.
+ """
+
+ name: str
+ shape: Callable[[Any], tuple[int, ...]]
+ src: str | Callable
+ dtype: torch.dtype | None = None
+ cut: Callable[[Any], Any] | None = None
+ transform: Callable | None = None
+ layers: object | None = ALL_LAYERS
+
+ def layer_indices(self, d: SimpleNamespace) -> list[int] | None:
+ """The layers this weight exists on, or None if it is not per-layer."""
+ if self.layers is None:
+ return None
+ if self.layers is ALL_LAYERS:
+ return list(range(d.num_layers))
+ return list(self.layers(d))
+
+ def sources(self, d: SimpleNamespace, i: int | None, prefix: str) -> list[tuple]:
+ """`[(checkpoint key or key tuple, index into the param | None, cut | None), ...]`."""
+ rows = self.src(d, i) if callable(self.src) else [(self.src.format(p=prefix), None)]
+ cut = self.cut(d) if self.cut is not None else None
+ out = []
+ for row in rows:
+ if len(row) == 3:
+ assert self.cut is None, f"{self.name}: a source with its own cut takes no W.cut"
+ out.append(tuple(row))
+ else:
+ out.append((*row, cut))
+ return out
+
+
+# ----------------------------------------------------------- host residency --
+#
+# Allocates the one PLE n-gram table. The table stays a `cuda` tensor of the same
+# shape/dtype. Registered, device-mapped pages (not a host gather) because the gather's
+# row indices are produced on device inside the captured decode graph.
+
+#: `cudaMemoryType` values; the core asserts on them after load.
+CUDA_MEMORY_TYPE_HOST = 1
+CUDA_MEMORY_TYPE_DEVICE = 2
+
+
+def memory_kind(ptr: int) -> int:
+ """The driver's own answer to "what kind of memory is this pointer": a `cudaMemoryType`."""
+ (attributes,) = CUASSERT(cudart.cudaPointerGetAttributes(ptr))
+ return int(attributes.type)
+
+
+#: `cudaHostRegister` flags (`driver_types.h`): mapped, so a kernel may dereference the
+#: pages, and portable, for every context in the process.
+_CUDA_HOST_REGISTER_PORTABLE = 0x01
+_CUDA_HOST_REGISTER_MAPPED = 0x02
+
+
+def _local_cpus(device) -> set | None:
+ """The CPUs sysfs lists as local to `device`'s PCI function and this process may run on,
+ or None when that is unknown."""
+ try:
+ props = torch.cuda.get_device_properties(device)
+ path = (
+ f"/sys/bus/pci/devices/{props.pci_domain_id:04x}:{props.pci_bus_id:02x}:"
+ f"{props.pci_device_id:02x}.0/local_cpulist"
+ )
+ with open(path) as fh:
+ text = fh.read().strip()
+ except (OSError, AttributeError, RuntimeError):
+ return None
+ cpus = set()
+ for part in filter(None, text.split(",")):
+ lo, _, hi = part.partition("-")
+ cpus.update(range(int(lo), int(hi or lo) + 1))
+ return (cpus & os.sched_getaffinity(0)) or None
+
+
+class _RegisteredHostBuffer:
+ """A pageable CPU allocation registered (pinned and mapped) for the device, and
+ unregistered before its storage is released: at this size glibc serves the
+ allocation with `mmap` and returns it with `munmap`, so freeing registered pages
+ would leave the driver mapping memory the OS may hand out again.
+
+ Not torch's pinned allocator: it rounds the 47.7 GiB table up to 64 GiB, and the
+ device's random row gathers from its pages are slower than from registered pageable
+ memory."""
+
+ def __init__(self, nbytes: int, device=None) -> None:
+ # The registration faults the pages in on this thread, so under the default local
+ # allocation policy they land on the NUMA node of the CPU it runs on. Pinned to the
+ # CPUs local to `device` meanwhile, the table sits next to the GPU that gathers from
+ # it: with a process per GPU, half the ranks would otherwise gather from the other
+ # socket's memory.
+ local = _local_cpus(device) if device is not None else None
+ previous = os.sched_getaffinity(0) if local else None
+ if local:
+ os.sched_setaffinity(0, local)
+ try:
+ self._storage = torch.empty(nbytes, dtype=torch.uint8)
+ self.pointer = self._storage.data_ptr()
+ CUASSERT(
+ cudart.cudaHostRegister(
+ self.pointer, nbytes, _CUDA_HOST_REGISTER_PORTABLE | _CUDA_HOST_REGISTER_MAPPED
+ )
+ )
+ finally:
+ if previous:
+ os.sched_setaffinity(0, previous)
+ # Bound now: at interpreter shutdown the module may already be torn down.
+ self._unregister = cudart.cudaHostUnregister
+
+ def __del__(self) -> None:
+ pointer, self.pointer = self.pointer, 0
+ if pointer:
+ self._unregister(pointer)
+
+
+def _host_resident_empty(
+ rows: int, cols: int, dtype, device
+) -> tuple[torch.Tensor, _RegisteredHostBuffer]:
+ """`[rows, cols]` of `dtype` on `device`, whose bytes live in host memory.
+
+ Returns the tensor and the registered CPU allocation that owns the pages. The
+ caller **must** keep the second alive for as long as the first: the tensor
+ is a view through a device address and owns nothing.
+ """
+ itemsize = torch.empty(0, dtype=dtype).element_size()
+ assert itemsize == 1, (
+ f"the host-resident allocation reinterprets a `|u1` buffer, so its "
+ f"element must be one byte; {dtype} is {itemsize}"
+ )
+ assert device.type == "cuda", f"host offload needs a cuda device to map into; got {device}"
+
+ host = _RegisteredHostBuffer(rows * cols, device)
+ (device_pointer,) = CUASSERT(cudart.cudaHostGetDevicePointer(host.pointer, 0))
+
+ # No public torch API builds a tensor over a pointer the caching allocator
+ # did not hand out; `__cuda_array_interface__` is the numba/cupy path.
+ # `|u1` because the protocol has no float8 typestr; reinterpreted below.
+ interface = SimpleNamespace(
+ __cuda_array_interface__={
+ "shape": (rows, cols),
+ "typestr": "|u1",
+ "data": (int(device_pointer), False),
+ "version": 3,
+ "strides": None,
+ }
+ )
+ table = torch.as_tensor(interface, device=device).view(dtype)
+ # `torch.as_tensor` copies when it decides the buffer is not already on the
+ # requested device; a copy here would be a silent 48 GiB device allocation.
+ assert table.data_ptr() == int(device_pointer), (
+ "torch.as_tensor copied the host-mapped buffer instead of viewing it, "
+ "so the table would be device-resident after all"
+ )
+ assert table.shape == (rows, cols) and table.dtype == dtype, (table.shape, table.dtype)
+ return table, host
+
+
+# ------------------------------------------------------------------ the table --
+
+_LM = "model.language_model"
+_VISUAL = "model.visual"
+_MTP = "mtp"
+_MTP_LAYER = "mtp.layers.0"
+
+# NVFP4 block size: one e4m3 scale per 16 contiguous elements along K.
+SF_VEC = 16
+
+# The FP8 block-scale family's block edge, on both axes of the MTP experts'
+# weight scales and along the activation row.
+FP8_BLOCK = 128
+
+#: The PLE n-gram table's storage: the checkpoint's FP8 per-tensor quantization, kept as
+#: stored (the gather dequantizes the rows it reads).
+PLE_TABLE_DTYPE = torch.float8_e4m3fn
+
+#: How the trunk's routed experts are split (`rank_dims`'s `expert_split`) that this table
+#: declares and `load` fills: by expert (`ep`, a window of whole experts) or by each
+#: expert's intermediate (`tp`, every expert at a zero-padded slice).
+LOADED_EXPERT_SPLITS = ("ep", "tp")
+
+#: The multiple a rank's stored routed-expert intermediate is padded up to. It is the
+#: engine's own storage rule for an NVFP4 trtllm-gen shard
+#: (`NVFP4TRTLLMGenFusedMoEMethod.resolve_alignments`: a padded intermediate that is not a
+#: multiple of 128 drops to a 128 weight alignment), which turns tp4's 160 into 256, the
+#: width `fp4_block_scale_moe_runner` certifies the padded 512-expert stack at. 640
+#: (tep4, dep4) is already a multiple.
+EXPERT_INTER_ALIGN = 128
+
+# Parameter-key suffixes after which to release transform scratch. The expert
+# relayouts allocate per-expert scratch once per local expert and the table
+# fill stages one shard at a time; without this, peak load memory is the whole
+# model rather than one parameter.
+_RELEASE_AFTER = (
+ "_fc2_s",
+ "_fc1_s",
+ "_fc2_s_cd",
+ "_fc1_s_cd",
+ "ple_table",
+ "mtp_fc1_w",
+ "mtp_fc2_w",
+)
+
+
+def _read(entry, cut: Cut | None) -> torch.Tensor:
+ """Realize one checkpoint entry, reading only `cut` of it.
+
+ `entry[cut]` on a lazy safetensors slice reads those rows or columns and nothing
+ else; on a plain tensor it is a view. Uncut, `[:]` reads it whole -- but a 0-dim
+ tensor (every NVFP4 per-tensor scalar on this checkpoint) rejects that index, and no
+ 0-dim entry has a cut."""
+ if cut is None:
+ return entry if getattr(entry, "ndim", 1) == 0 else entry[:]
+ return entry[cut[0]] if len(cut) == 1 else entry[cut]
+
+
+# --- cuts: this rank's block of rows or columns of a checkpoint tensor ---
+
+
+def _span(start: int, width: int) -> slice:
+ return slice(start, start + width)
+
+
+def _q_rows(d) -> Cut:
+ """`q_proj` rows of this rank's q heads: each head's q rows, then its gate rows."""
+ return (_span(d.q_head_start * 2 * d.head_dim, 2 * d.q_width),)
+
+
+def _kv_rows(d) -> Cut:
+ """`k_proj` / `v_proj` rows of this rank's KV heads."""
+ return (_span(d.kv_head_start * d.head_dim, d.kv_width),)
+
+
+def _o_cols(d) -> Cut:
+ """`o_proj` input columns of this rank's q heads."""
+ return (slice(None), _span(d.q_head_start * d.head_dim, d.q_width))
+
+
+def _key_rows(d, section: int) -> Cut:
+ """Rows of this rank's key heads in section `section` (0 q, 1 k) of a `[q | k | v]`
+ Gated-DeltaNet tensor."""
+ return (_span(section * d.key_dim_total + d.k_head_start * d.head_k_dim, d.key_dim),)
+
+
+def _value_rows(d, base: int = 0) -> Cut:
+ """Rows of this rank's value heads, `base` rows into the tensor."""
+ return (_span(base + d.v_head_start * d.head_v_dim, d.value_dim),)
+
+
+def _value_head_rows(d) -> Cut:
+ """One row per value head (`in_proj_a`, `in_proj_b`, `A_log`, `dt_bias`)."""
+ return (_span(d.v_head_start, d.num_v_heads),)
+
+
+def _value_cols(d) -> Cut:
+ """`out_proj` input columns of this rank's value heads."""
+ return (slice(None), _span(d.v_head_start * d.head_v_dim, d.value_dim))
+
+
+def _qkv_blocks(d) -> tuple[tuple[slice, Cut], ...]:
+ """(destination rows, checkpoint cut) of the three blocks of a `[q | k | v]`
+ Gated-DeltaNet tensor this rank holds, in the `[q | k | v]` order of its conv pool."""
+ return (
+ (slice(0, d.key_dim), _key_rows(d, 0)),
+ (slice(d.key_dim, 2 * d.key_dim), _key_rows(d, 1)),
+ (slice(2 * d.key_dim, d.conv_dim), _value_rows(d, 2 * d.key_dim_total)),
+ )
+
+
+# --- relayouts, on the rank's cut ---
+
+
+def _block32_perm(rows: int, device) -> torch.Tensor:
+ """Gather index of the 32-row block shuffle: inside each aligned block of
+ 32 rows, source row `4u + v` (0<=u<8, 0<=v<4) lands at destination row
+ `8v + u`."""
+ assert rows % 32 == 0, rows
+ src = torch.arange(32)
+ dst = (src % 4) * 8 + src // 4
+ idx = torch.empty(32, dtype=torch.long)
+ idx[dst] = src
+ blocks = rows // 32
+ return (idx.repeat(blocks) + torch.arange(blocks).repeat_interleave(32) * 32).to(device)
+
+
+def _interleave_perm(rows: int, device) -> torch.Tensor:
+ """Gather index that interleaves the `[up | gate]` halves of a `2*I` row
+ stack into up0, gate0, up1, gate1, ..."""
+ p = torch.empty(rows, dtype=torch.long)
+ p[0::2] = torch.arange(0, rows // 2)
+ p[1::2] = torch.arange(rows // 2, rows)
+ return p.to(device)
+
+
+#: The group the CuteDSL FC1 epilogue pairs up and gate rows in: rows `[128k, 128k + 64)`
+#: of the fused operand are up rows and the next 64 their gate rows.
+_CUTEDSL_GATE_UP_GROUP = 64
+
+
+def _group_interleave_perm(rows: int, group: int, device) -> torch.Tensor:
+ """Gather index that interleaves the `[up | gate]` halves of a `2*I` row stack in
+ `group`-row blocks: up[0:g], gate[0:g], up[g:2g], gate[g:2g], ..."""
+ assert rows % (2 * group) == 0, (rows, group)
+ idx = torch.arange(rows, device=device).view(2, rows // (2 * group), group)
+ return idx.transpose(0, 1).reshape(-1)
+
+
+def _fc1_perm(rows: int, device, d) -> torch.Tensor:
+ """The FC1 row gather of the trunk's routed backend. trtllm-gen: interleave then
+ 32-row block shuffle, composed into one gather. CuteDSL: the 64-row group
+ interleave alone."""
+ if d.trunk_moe_backend == "CUTEDSL":
+ return _group_interleave_perm(rows, _CUTEDSL_GATE_UP_GROUP, device)
+ return _interleave_perm(rows, device)[_block32_perm(rows, device)]
+
+
+def _swizzle(x: torch.Tensor) -> torch.Tensor:
+ """128x4 block-scale swizzle of an `[M, C]` uint8 scale matrix, returned
+ flat (`M % 128 == 0` and `C % 4 == 0` hold at this geometry, so the result
+ is exactly `M * C` bytes)."""
+ return torch.ops.trtllm.block_scale_interleave(x.contiguous())
+
+
+#: Rows after the shared expert's gate|up rows in its fused operand: the per-token gate's
+#: one row, then zeros, so the GEMM's output row stays a whole number of 16-byte vectors.
+SH_GATE_ROWS = 8
+
+
+def _cat_rows(t: tuple[torch.Tensor, ...], d) -> torch.Tensor:
+ """Fused shared-expert weight bytes: **gate** rows first, the half order
+ `flashinfer_silu_and_mul` reads, then up, then the per-token gate's row and
+ `SH_GATE_ROWS - 1` zero rows -- one GEMM forms gate|up and the gate logit."""
+ gate, up, sh_gate = t
+ pad = sh_gate.new_zeros(SH_GATE_ROWS - 1, sh_gate.shape[-1])
+ return torch.cat((gate, up, sh_gate, pad), dim=0).contiguous()
+
+
+def _attn_in_w(t: tuple[torch.Tensor, ...], d) -> torch.Tensor:
+ """(q_proj, k_proj, v_proj, index_qk_proj), each at this rank's cut -> the fused
+ `[q | k | v | gate | index | pad]` input projection of one trunk attention layer (see
+ the module docstring).
+
+ `q_proj`'s rows are `[heads, 2, head_dim]`: head `h`'s q rows, then its gate rows.
+ """
+ q, k, v, idx = t
+ heads, head_dim, hidden = d.heads_q, d.head_dim, q.shape[-1]
+ qg = q.reshape(heads, 2, head_dim, hidden)
+ pad = q.new_zeros(d.index_head_dim, hidden)
+ rows = (qg[:, 0].reshape(heads * head_dim, hidden), k, v)
+ rows += (qg[:, 1].reshape(heads * head_dim, hidden), idx, pad)
+ return torch.cat(rows, dim=0).contiguous()
+
+
+def _expert_fc1_w(t: tuple[torch.Tensor, ...], d) -> torch.Tensor:
+ """(up, gate) packed nibbles -> one expert's FC1 operand."""
+ up, gate = t
+ w = torch.cat([up, gate], dim=0)
+ return torch.index_select(w, 0, _fc1_perm(w.shape[0], w.device, d)).contiguous()
+
+
+def _expert_fc1_s(t: tuple[torch.Tensor, ...], d) -> torch.Tensor:
+ """(up, gate) e4m3 scales -> one expert's FC1 scale operand: the weight's own row
+ permutation, then the swizzle."""
+ up, gate = t
+ s = torch.cat([up.view(torch.uint8), gate.view(torch.uint8)], dim=0)
+ rows, cols = s.shape
+ s = torch.index_select(s, 0, _fc1_perm(rows, s.device, d))
+ return _swizzle(s).reshape(rows, cols)
+
+
+def _expert_fc2_w(t: torch.Tensor, d) -> torch.Tensor:
+ """`down_proj` packed nibbles -> one expert's FC2 operand (no interleave): the
+ 32-row block shuffle on trtllm-gen, checkpoint row order on CuteDSL."""
+ if d.trunk_moe_backend == "CUTEDSL":
+ return t.contiguous()
+ return torch.index_select(t, 0, _block32_perm(t.shape[0], t.device)).contiguous()
+
+
+def _expert_fc2_s(t: torch.Tensor, d) -> torch.Tensor:
+ """`down_proj` e4m3 scales -> one expert's FC2 scale operand: the weight's own row
+ order, then the swizzle."""
+ s = t.view(torch.uint8)
+ rows, cols = s.shape
+ if d.trunk_moe_backend != "CUTEDSL":
+ s = torch.index_select(s, 0, _block32_perm(rows, s.device))
+ return _swizzle(s).reshape(rows, cols)
+
+
+def _zero_extend(x: torch.Tensor, dim: int, size: int) -> torch.Tensor:
+ """`x` extended along `dim` to `size` with zero bytes: a zero packed byte is two +0.0
+ e2m1 codes, a zero scale byte is an e4m3 +0.0 block scale."""
+ if x.shape[dim] == size:
+ return x
+ assert x.shape[dim] < size, (tuple(x.shape), dim, size)
+ raw = x.view(torch.uint8)
+ pad = list(raw.shape)
+ pad[dim] = size - raw.shape[dim]
+ return torch.cat((raw, raw.new_zeros(pad)), dim).view(x.dtype)
+
+
+def _zero_padded(transform: Callable | None, dim: int, unit: int) -> Callable:
+ """`transform` run on its source(s) zero-extended along `dim` to the rank's stored
+ routed-expert intermediate, `d.expert_inter // unit` there: unit 1 for the rows of an
+ `up_proj` / `gate_proj` and of their scales, 2 for `down_proj`'s packed code columns,
+ `SF_VEC` for its scale columns. Each half of FC1 is extended before the `[up ; gate]`
+ concatenation, so the pad rows sit at the end of each half, where the runner's
+ `intermediate_size` puts them. A no-op where the rank holds whole experts."""
+
+ def padded(t, d):
+ size = d.expert_inter // unit
+ if isinstance(t, tuple):
+ t = tuple(_zero_extend(x, dim, size) for x in t)
+ else:
+ t = _zero_extend(t, dim, size)
+ return t if transform is None else transform(t, d)
+
+ return padded
+
+
+def _expert_inter_rows(d) -> Cut | None:
+ """This rank's rows of one expert's `up_proj` or `gate_proj` (or their block scales):
+ its slice of the intermediate, or None (whole) when the experts are split by expert."""
+ if d.expert_inter_share == d.moe_inter:
+ return None
+ return (_span(d.expert_inter_start, d.expert_inter_share),)
+
+
+def _fc1_cut(d):
+ """The `(up, gate)` pair's cuts: the same intermediate rows of both halves."""
+ rows = _expert_inter_rows(d)
+ return None if rows is None else (rows, rows)
+
+
+def _fc2_cut(unit: int) -> Callable:
+ """`down_proj`'s input columns of this rank's intermediate slice, `unit` intermediate
+ elements to a stored column (2 for the packed codes, `SF_VEC` for the block scales)."""
+
+ def cut(d):
+ if _expert_inter_rows(d) is None:
+ return None
+ return (
+ slice(None),
+ _span(d.expert_inter_start // unit, d.expert_inter_share // unit),
+ )
+
+ return cut
+
+
+#: The layout argument of the second, CuteDSL-layout expert copy (`fc*_cd`).
+_CUTEDSL_LAYOUT = SimpleNamespace(trunk_moe_backend="CUTEDSL")
+
+
+def _expert_fc1_w_cd(t: tuple[torch.Tensor, ...], d) -> torch.Tensor:
+ """(up, gate) packed nibbles -> one expert's FC1 operand in the CuteDSL layout."""
+ return _expert_fc1_w(t, _CUTEDSL_LAYOUT)
+
+
+def _expert_fc1_s_cd(t: tuple[torch.Tensor, ...], d) -> torch.Tensor:
+ """(up, gate) e4m3 scales -> one expert's FC1 scale operand in the CuteDSL layout."""
+ return _expert_fc1_s(t, _CUTEDSL_LAYOUT)
+
+
+def _expert_fc2_s_cd(t: torch.Tensor, d) -> torch.Tensor:
+ """`down_proj` e4m3 scales -> one expert's FC2 scale operand in the CuteDSL layout."""
+ return _expert_fc2_s(t, _CUTEDSL_LAYOUT)
+
+
+def _mtp_fc1_w(t: tuple[torch.Tensor, ...], d) -> torch.Tensor:
+ """(up, gate) e4m3 codes -> one draft expert's FC1 operand, `[up ; gate]` on
+ the row axis and nothing else. The FP8 block-scale runner reads the
+ checkpoint's plain row-major layout."""
+ up, gate = t
+ return torch.cat([up, gate], dim=0).contiguous()
+
+
+def _mtp_fc1_s(t: tuple[torch.Tensor, ...], d) -> torch.Tensor:
+ """(up, gate) 128x128 block scales -> one draft expert's FC1 scale operand:
+ the same `[up ; gate]` concatenation, cast bf16 -> fp32."""
+ up, gate = t
+ return torch.cat([up, gate], dim=0).to(torch.float32).contiguous()
+
+
+def _mtp_fc2_s(t: torch.Tensor, d) -> torch.Tensor:
+ """`down_proj.weight_scale_inv` -> the FC2 scale operand: a bf16 -> fp32
+ cast, no relayout."""
+ return t.to(torch.float32).contiguous()
+
+
+#: The layout argument of the draft experts' NVFP4 copy: the draft MoE runs it on trtllm-gen
+#: whatever the trunk's routed backend is.
+_TRTLLM_GEN_LAYOUT = SimpleNamespace(trunk_moe_backend="TRTLLM")
+
+#: `448 * 6`: an NVFP4 global scale is this over the tensor's amax (modelopt's recipe).
+NVFP4_AMAX_TO_GLOBAL = 448.0 * 6.0
+
+
+def _fp8_block_dequant(codes: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
+ """`[E, N, K]` e4m3 codes times their `[E, N/128, K/128]` fp32 block scales, in fp32."""
+ s = scales.repeat_interleave(FP8_BLOCK, 1).repeat_interleave(FP8_BLOCK, 2)
+ return codes.float() * s
+
+
+def mtp_experts_to_nvfp4(fc1_w, fc1_s, fc2_w, fc2_s, chunk: int = 32) -> tuple:
+ """Requantize the draft layer's FP8 128x128 block-scaled experts (this rank's window)
+ to NVFP4 in the trunk's trtllm-gen layout (`_expert_fc1_w/_s`, `_expert_fc2_w/_s`).
+
+ Per expert and projection (FC1's up|gate share one, as the fused FC1's alpha
+ requires): dequantize with the block scales, take the global scale
+ `448 * 6 / amax`, then `fp4_quantize` (16-element blocks, e4m3 block scales,
+ linear order), which then gets the trunk's row permutation and 128x4 swizzle.
+ Returns `(fc1_w [E, 2I, H/2], fc1_s [E, 2I, H/16], fc2_w [E, H, I/2],
+ fc2_s [E, H, I/16])` uint8 and the fp32 `[E]` weight global scales `(g1, g2)`;
+ a weight's real value is its NVFP4 reconstruction over its global scale.
+ `fc1_w` / `fc1_s` are the `[up ; gate]` stacks `_mtp_fc1_w` / `_mtp_fc1_s` build."""
+ experts, two_i, hidden = fc1_w.shape
+ inter = two_i // 2
+ dev = fc1_w.device
+ u8 = torch.uint8
+ out = (
+ torch.empty(experts, two_i, hidden // 2, dtype=u8, device=dev),
+ torch.empty(experts, two_i, hidden // SF_VEC, dtype=u8, device=dev),
+ torch.empty(experts, hidden, inter // 2, dtype=u8, device=dev),
+ torch.empty(experts, hidden, inter // SF_VEC, dtype=u8, device=dev),
+ )
+ g1 = torch.empty(experts, dtype=torch.float32, device=dev)
+ g2 = torch.empty(experts, dtype=torch.float32, device=dev)
+ layout = _TRTLLM_GEN_LAYOUT
+ for e0 in range(0, experts, chunk):
+ e1 = min(e0 + chunk, experts)
+ w1 = _fp8_block_dequant(fc1_w[e0:e1], fc1_s[e0:e1])
+ w2 = _fp8_block_dequant(fc2_w[e0:e1], fc2_s[e0:e1])
+ g1[e0:e1] = NVFP4_AMAX_TO_GLOBAL / w1.abs().amax(dim=(1, 2)).clamp_min(1e-30)
+ g2[e0:e1] = NVFP4_AMAX_TO_GLOBAL / w2.abs().amax(dim=(1, 2)).clamp_min(1e-30)
+ w1, w2 = w1.to(torch.bfloat16), w2.to(torch.bfloat16)
+ for j, e in enumerate(range(e0, e1)):
+ q, s = torch.ops.trtllm.fp4_quantize(w1[j], g1[e : e + 1], SF_VEC, False, False)
+ s = s.view(two_i, hidden // SF_VEC)
+ out[0][e] = _expert_fc1_w((q[:inter], q[inter:]), layout)
+ out[1][e] = _expert_fc1_s((s[:inter], s[inter:]), layout)
+ q, s = torch.ops.trtllm.fp4_quantize(w2[j], g2[e : e + 1], SF_VEC, False, False)
+ out[2][e] = _expert_fc2_w(q, layout)
+ out[3][e] = _expert_fc2_s(s.view(hidden, inter // SF_VEC), layout)
+ del w1, w2
+ return (*out, g1, g2)
+
+
+# --- which layers carry a weight ---
+
+
+def _attn_layers(d):
+ """The full-attention layers (`i % 4 == 3` on this checkpoint)."""
+ return [i for i, a in enumerate(d.is_attn) if a]
+
+
+def _gdn_layers(d):
+ """The Gated-DeltaNet (linear-attention) layers."""
+ return [i for i, a in enumerate(d.is_attn) if not a]
+
+
+def _trunk_expert_layers(d):
+ """Every layer when the routed experts are split in a way this table declares
+ (`LOADED_EXPERT_SPLITS`); none otherwise."""
+ return range(d.num_layers) if d.expert_split in LOADED_EXPERT_SPLITS else ()
+
+
+def _cute_dsl_copy_layers(d):
+ """The trunk expert layers when the core keeps the second, CuteDSL-layout copy that
+ prefill-sized launches read under trtllm-gen (`moe_prefill_copy`); none otherwise."""
+ return _trunk_expert_layers(d) if d.moe_prefill_copy else ()
+
+
+def _vision_blocks(d):
+ """The vision tower's blocks. Keyed `l{b}_vis_*`: block `b`, not decoder layer `b`."""
+ return range(d.v_depth)
+
+
+def _mtp_only(d):
+ """The MTP module, at its own layer index `num_hidden_layers`, only when it is live."""
+ return [d.num_layers] if d.mtp_enabled else []
+
+
+# --- where the bytes come from ---
+
+
+def _ple(suffix):
+ """A PLE weight. The PLE layer's index is config, so the key is computed."""
+ return lambda d, i: [(f"{_LM}.layers.{d.ple_layer}.ple.{suffix}", None)]
+
+
+def _ngram_prefix(ple_layer: int) -> str:
+ """The quantized n-gram table submodule: the 128 shards and the one
+ per-tensor scale that dequantizes them both hang off this, one level below
+ `ple_embedding` itself."""
+ return f"{_LM}.layers.{ple_layer}.ple.ple_embedding.ngram_embedding"
+
+
+def _visual(suffix):
+ """A vision-tower weight that is not per-block."""
+ return lambda d, i: [(f"{_VISUAL}.{suffix}", None)]
+
+
+def _vision_block(suffix):
+ """One vision block's weight."""
+ return lambda d, i: [(f"{_VISUAL}.blocks.{i}.{suffix}", None)]
+
+
+def _trunk_window(d) -> tuple[int, int]:
+ """(first expert, count) of this rank's routed-expert window."""
+ return d.expert_offset, d.local_experts
+
+
+def _mtp_window(d) -> tuple[int, int]:
+ """(first expert, count) of this rank's window of the draft layer's experts."""
+ return d.mtp_expert_offset, d.mtp_local_experts
+
+
+def _routing_space(d) -> tuple[int, int]:
+ """Every expert: what the replicated per-tensor scalars cover."""
+ return 0, d.num_experts
+
+
+def _experts(prefix, *suffixes, window):
+ """One source per expert of `window(d)`, written at its local index; a pair of
+ suffixes hands the transform both tensors together. The checkpoint key depends on
+ the rank, which is why the source is computed rather than templated."""
+
+ def build(d, i):
+ start, count = window(d)
+ base = prefix(i)
+ out = []
+ for local in range(count):
+ keys = tuple(f"{base}.experts.{start + local}.{s}" for s in suffixes)
+ out.append((keys if len(keys) > 1 else keys[0], (local,)))
+ return out
+
+ return build
+
+
+def _trunk_mlp(i):
+ return f"{_LM}.layers.{i}.mlp"
+
+
+def _mtp_mlp(i):
+ return f"{_MTP_LAYER}.mlp"
+
+
+def _in_proj(d, i):
+ """The Gated-DeltaNet input projections at this rank's heads, in the row order of
+ the concatenated `in_proj_w` (and so the column order of its GEMM output):
+ `[q | k | v (conv_dim) | z (value_dim) | a | b]`, `d.in_proj_split`."""
+ p = f"{_LM}.layers.{i}.linear_attn"
+ out = [(f"{p}.in_proj_qkv.weight", (rows,), cut) for rows, cut in _qkv_blocks(d)]
+ row = d.conv_dim
+ out.append((f"{p}.in_proj_z.weight", (_span(row, d.value_dim),), _value_rows(d)))
+ row += d.value_dim
+ for name in ("in_proj_a", "in_proj_b"):
+ out.append((f"{p}.{name}.weight", (_span(row, d.num_v_heads),), _value_head_rows(d)))
+ row += d.num_v_heads
+ return out
+
+
+def _conv(d, i):
+ """The depthwise conv's `[q | k | v]` channels of this rank's heads."""
+ key = f"{_LM}.layers.{i}.linear_attn.conv1d.weight"
+ return [(key, (rows,), cut) for rows, cut in _qkv_blocks(d)]
+
+
+#: The six per-tensor NVFP4 scalars every routed expert carries, as
+#: (parameter suffix, checkpoint suffix). The checkpoint stores modelopt
+#: reciprocals; the core folds them into the runner's scalars after load.
+_EXPERT_SCALARS = (
+ ("e_isc1", "gate_proj.input_scale"),
+ ("e_isc1_up", "up_proj.input_scale"),
+ ("e_ws2_1", "gate_proj.weight_scale_2"),
+ ("e_ws2_1_up", "up_proj.weight_scale_2"),
+ ("e_isc2", "down_proj.input_scale"),
+ ("e_ws2_2", "down_proj.weight_scale_2"),
+)
+
+#: The checkpoint tensors of one routed expert that are not per-tensor scalars: what a
+#: rank outside its window leaves unread. The draft layer's FP8 experts carry
+#: `weight_scale_inv` in place of the NVFP4 `weight_scale`.
+_PROJECTIONS = ("gate_proj", "up_proj", "down_proj")
+_TRUNK_EXPERT_TENSORS = tuple(f"{p}.{s}" for p in _PROJECTIONS for s in ("weight", "weight_scale"))
+_MTP_EXPERT_TENSORS = tuple(
+ f"{p}.{s}" for p in _PROJECTIONS for s in ("weight", "weight_scale_inv")
+)
+
+
+#: (parameter tag, checkpoint module) of the two per-layer hyper-connection gates.
+_HC_MODULES = (("ahc", "attn_hyper_connection"), ("mhc", "mlp_hyper_connection"))
+
+
+def _hyper_connection(tag, src, layers):
+ """The four weights of one gated-residual hyper-connection; `src(name)` names the source."""
+ return (
+ W(f"{tag}_norm", shape=lambda d: (d.stream,), src=src("hc_norm"), layers=layers),
+ W(
+ f"{tag}_down",
+ shape=lambda d: (d.hc_lowrank, d.stream),
+ src=src("input_mix_weight_down"),
+ layers=layers,
+ ),
+ W(
+ f"{tag}_up",
+ shape=lambda d: (d.stream, d.hc_lowrank),
+ src=src("input_mix_weight_up"),
+ layers=layers,
+ ),
+ W(
+ f"{tag}_inject",
+ shape=lambda d: (d.hc_count, d.stream),
+ src=src("block_inject_weight"),
+ layers=layers,
+ ),
+ )
+
+
+def _mixer(tag, prefix, layers):
+ """A stack-final hyper-connection mixer: a layer's gate minus the injection weight."""
+ return (
+ W(
+ f"{tag}_norm",
+ shape=lambda d: (d.stream,),
+ src=f"{prefix}.hc_norm.weight",
+ layers=layers,
+ ),
+ W(
+ f"{tag}_down",
+ shape=lambda d: (d.hc_lowrank, d.stream),
+ src=f"{prefix}.input_mix_weight_down.weight",
+ layers=layers,
+ ),
+ W(
+ f"{tag}_up",
+ shape=lambda d: (d.stream, d.hc_lowrank),
+ src=f"{prefix}.input_mix_weight_up.weight",
+ layers=layers,
+ ),
+ )
+
+
+def _attn_in_fused(d, i):
+ """The trunk attention layers' one input projection: q_proj, k_proj and v_proj at this
+ rank's heads and the QSA indexer's replicated index_qk_proj, re-laid by `_attn_in_w`."""
+ p = f"{_LM}.layers.{i}.self_attn"
+ keys = (
+ f"{p}.q_proj.weight",
+ f"{p}.k_proj.weight",
+ f"{p}.v_proj.weight",
+ f"{p}.indexer.index_qk_proj.weight",
+ )
+ return [(keys, None, (_q_rows(d), _kv_rows(d), _kv_rows(d), None))]
+
+
+def _attn_out_and_norms(tag, prefix, layers):
+ """o_proj at this rank's q heads, and the replicated per-head q/k norms."""
+ return (
+ W(
+ f"{tag}o_w",
+ shape=lambda d: (d.hidden, d.q_width),
+ src=f"{prefix}.o_proj.weight",
+ cut=_o_cols,
+ layers=layers,
+ ),
+ W(
+ f"{tag}q_norm",
+ shape=lambda d: (d.head_dim,),
+ src=f"{prefix}.q_norm.weight",
+ layers=layers,
+ ),
+ W(
+ f"{tag}k_norm",
+ shape=lambda d: (d.head_dim,),
+ src=f"{prefix}.k_norm.weight",
+ layers=layers,
+ ),
+ )
+
+
+def _attn_in_separate(tag, prefix, layers):
+ """The draft layer's separate q (with its interleaved gate), k and v projections, at
+ this rank's heads."""
+ return (
+ W(
+ f"{tag}q_w",
+ shape=lambda d: (2 * d.q_width, d.hidden),
+ src=f"{prefix}.q_proj.weight",
+ cut=_q_rows,
+ layers=layers,
+ ),
+ W(
+ f"{tag}k_w",
+ shape=lambda d: (d.kv_width, d.hidden),
+ src=f"{prefix}.k_proj.weight",
+ cut=_kv_rows,
+ layers=layers,
+ ),
+ W(
+ f"{tag}v_w",
+ shape=lambda d: (d.kv_width, d.hidden),
+ src=f"{prefix}.v_proj.weight",
+ cut=_kv_rows,
+ layers=layers,
+ ),
+ )
+
+
+def _index_branch(tag, prefix, layers, fused_in=False):
+ """The QSA index branch's three weights: one replicated, unquantized projection
+ producing `[index_heads q | 1 k]`, and two per-head gemma norms. With `fused_in`, the
+ projection is part of `_attn_in_fused`'s one operand and only the norms are declared."""
+ projection = W(
+ f"{tag}idx_w",
+ shape=lambda d: (d.index_width, d.hidden),
+ src=f"{prefix}.index_qk_proj.weight",
+ layers=layers,
+ )
+ return (
+ *(() if fused_in else (projection,)),
+ W(
+ f"{tag}idx_qn",
+ shape=lambda d: (d.index_head_dim,),
+ src=f"{prefix}.q_layernorm.weight",
+ layers=layers,
+ ),
+ W(
+ f"{tag}idx_kn",
+ shape=lambda d: (d.index_head_dim,),
+ src=f"{prefix}.k_layernorm.weight",
+ layers=layers,
+ ),
+ )
+
+
+def _shared_expert(tag, prefix, layers):
+ """The replicated bf16 shared expert beside the routed ones, its per-token gate in the
+ same operand as its gate|up (`_cat_rows`)."""
+ return (
+ W(
+ f"{tag}sh_gu",
+ shape=lambda d: (2 * d.shared_inter + SH_GATE_ROWS, d.hidden),
+ src=lambda d, i: [
+ (
+ (
+ f"{prefix(i)}.shared_expert.gate_proj.weight",
+ f"{prefix(i)}.shared_expert.up_proj.weight",
+ f"{prefix(i)}.shared_expert_gate.weight",
+ ),
+ None,
+ )
+ ],
+ transform=_cat_rows,
+ layers=layers,
+ ),
+ W(
+ f"{tag}sh_dn",
+ shape=lambda d: (d.hidden, d.shared_inter),
+ src=lambda d, i: [(f"{prefix(i)}.shared_expert.down_proj.weight", None)],
+ layers=layers,
+ ),
+ )
+
+
+u8, f8, f32, i64 = torch.uint8, torch.float8_e4m3fn, torch.float32, torch.int64
+
+
+class Qwen4ExpDefaultWeights:
+ """Qwen3.8-Flash-Next-NVFP4 / sm_103 / default, at one rank's share.
+
+ The table, and the passes over it: `declare` allocates every weight it names,
+ `manifest` maps checkpoint keys and cuts onto them, `load` fills them. `declare` and
+ `manifest` share one "which parameters exist" decision (`_targets`), so they
+ cannot disagree about a weight.
+ """
+
+ WEIGHTS: tuple[W, ...] = (
+ # --- the two per-layer hyper-connection gates: every trunk layer, replicated ---
+ *(
+ w
+ for tag, module in _HC_MODULES
+ for w in _hyper_connection(
+ tag, lambda name, m=module: f"{{p}}.{m}.{name}.weight", ALL_LAYERS
+ )
+ ),
+ # --- the gated GQA layers and their QSA index branch ---
+ # q/k/v at this rank's heads and the indexer's replicated [index_heads q | 1 k]
+ # projection as one operand, then o_proj at this rank's heads, the q/k norms and
+ # the indexer's two per-head gemma norms.
+ W(
+ "attn_in_w",
+ shape=lambda d: (d.attn_in_rows, d.hidden),
+ src=_attn_in_fused,
+ transform=_attn_in_w,
+ layers=_attn_layers,
+ ),
+ *_attn_out_and_norms("", "{p}.self_attn", _attn_layers),
+ *_index_branch("", "{p}.self_attn.indexer", _attn_layers, fused_in=True),
+ # --- the Gated-DeltaNet layers, at this rank's key and value heads ---
+ # The four input projections read the same activation, so they are one
+ # weight and one GEMM: rows `[q|k|v (conv_dim) | z (value_dim) | a | b]`
+ # (`d.in_proj_split`), each checkpoint block copied into its own row range.
+ W(
+ "in_proj_w",
+ shape=lambda d: (sum(d.in_proj_split), d.hidden),
+ src=_in_proj,
+ layers=_gdn_layers,
+ ),
+ W(
+ "A_log",
+ shape=lambda d: (d.num_v_heads,),
+ src="{p}.linear_attn.A_log",
+ cut=_value_head_rows,
+ layers=_gdn_layers,
+ ),
+ W(
+ "dt_bias",
+ shape=lambda d: (d.num_v_heads,),
+ src="{p}.linear_attn.dt_bias",
+ cut=_value_head_rows,
+ layers=_gdn_layers,
+ ),
+ # Stored [conv_dim, 1, width]; the conv entries want [conv_dim, width], which is
+ # a zero-copy view the targets take once, after load.
+ W(
+ "conv_w",
+ shape=lambda d: (d.conv_dim, 1, d.conv_kernel),
+ src=_conv,
+ layers=_gdn_layers,
+ ),
+ W(
+ "gnorm",
+ shape=lambda d: (d.head_v_dim,),
+ src="{p}.linear_attn.norm.weight",
+ layers=_gdn_layers,
+ ),
+ W(
+ "out_w",
+ shape=lambda d: (d.hidden, d.value_dim),
+ src="{p}.linear_attn.out_proj.weight",
+ cut=_value_cols,
+ layers=_gdn_layers,
+ ),
+ # --- MoE on every layer: this rank's window of NVFP4 routed experts, a replicated
+ # bf16 router and shared expert ---
+ W("router", shape=lambda d: (d.num_experts, d.hidden), src="{p}.mlp.gate.weight"),
+ # Each expert at this rank's intermediate share (`_fc1_cut`, `_fc2_cut`), stored
+ # zero-padded to `expert_inter` (`_zero_padded`); whole experts under an EP window.
+ W(
+ "fc1_w",
+ shape=lambda d: (d.local_experts, 2 * d.expert_inter, d.hidden // 2),
+ dtype=u8,
+ src=_experts(_trunk_mlp, "up_proj.weight", "gate_proj.weight", window=_trunk_window),
+ cut=_fc1_cut,
+ transform=_zero_padded(_expert_fc1_w, 0, 1),
+ layers=_trunk_expert_layers,
+ ),
+ W(
+ "fc1_s",
+ shape=lambda d: (d.local_experts, 2 * d.expert_inter, d.hidden // SF_VEC),
+ dtype=u8,
+ src=_experts(
+ _trunk_mlp, "up_proj.weight_scale", "gate_proj.weight_scale", window=_trunk_window
+ ),
+ cut=_fc1_cut,
+ transform=_zero_padded(_expert_fc1_s, 0, 1),
+ layers=_trunk_expert_layers,
+ ),
+ W(
+ "fc2_w",
+ shape=lambda d: (d.local_experts, d.hidden, d.expert_inter // 2),
+ dtype=u8,
+ src=_experts(_trunk_mlp, "down_proj.weight", window=_trunk_window),
+ cut=_fc2_cut(2),
+ transform=_zero_padded(_expert_fc2_w, 1, 2),
+ layers=_trunk_expert_layers,
+ ),
+ W(
+ "fc2_s",
+ shape=lambda d: (d.local_experts, d.hidden, d.expert_inter // SF_VEC),
+ dtype=u8,
+ src=_experts(_trunk_mlp, "down_proj.weight_scale", window=_trunk_window),
+ cut=_fc2_cut(SF_VEC),
+ transform=_zero_padded(_expert_fc2_s, 1, SF_VEC),
+ layers=_trunk_expert_layers,
+ ),
+ # Under trtllm-gen, the same window again in the CuteDSL layout, when the core keeps
+ # it for prefill-sized launches.
+ W(
+ "fc1_w_cd",
+ shape=lambda d: (d.local_experts, 2 * d.expert_inter, d.hidden // 2),
+ dtype=u8,
+ src=_experts(_trunk_mlp, "up_proj.weight", "gate_proj.weight", window=_trunk_window),
+ cut=_fc1_cut,
+ transform=_zero_padded(_expert_fc1_w_cd, 0, 1),
+ layers=_cute_dsl_copy_layers,
+ ),
+ W(
+ "fc1_s_cd",
+ shape=lambda d: (d.local_experts, 2 * d.expert_inter, d.hidden // SF_VEC),
+ dtype=u8,
+ src=_experts(
+ _trunk_mlp, "up_proj.weight_scale", "gate_proj.weight_scale", window=_trunk_window
+ ),
+ cut=_fc1_cut,
+ transform=_zero_padded(_expert_fc1_s_cd, 0, 1),
+ layers=_cute_dsl_copy_layers,
+ ),
+ W(
+ "fc2_w_cd",
+ shape=lambda d: (d.local_experts, d.hidden, d.expert_inter // 2),
+ dtype=u8,
+ src=_experts(_trunk_mlp, "down_proj.weight", window=_trunk_window),
+ cut=_fc2_cut(2),
+ transform=_zero_padded(None, 1, 2),
+ layers=_cute_dsl_copy_layers,
+ ),
+ W(
+ "fc2_s_cd",
+ shape=lambda d: (d.local_experts, d.hidden, d.expert_inter // SF_VEC),
+ dtype=u8,
+ src=_experts(_trunk_mlp, "down_proj.weight_scale", window=_trunk_window),
+ cut=_fc2_cut(SF_VEC),
+ transform=_zero_padded(_expert_fc2_s_cd, 1, SF_VEC),
+ layers=_cute_dsl_copy_layers,
+ ),
+ # Replicated over the whole routing space, at every split.
+ *(
+ W(
+ name,
+ shape=lambda d: (d.num_experts,),
+ dtype=f32,
+ src=_experts(_trunk_mlp, suffix, window=_routing_space),
+ )
+ for name, suffix in _EXPERT_SCALARS
+ ),
+ *_shared_expert("", _trunk_mlp, ALL_LAYERS),
+ # --- not per-layer: the stack-final mixer, which is also the final norm ---
+ *_mixer("fhc", f"{_LM}.hyper_connection_mixer", None),
+ W(
+ "embed",
+ shape=lambda d: (d.vocab, d.hidden),
+ src=f"{_LM}.embed_tokens.weight",
+ layers=None,
+ ),
+ # --- PLE, on its one decoder layer (the n-gram table itself is declared at load) ---
+ W(
+ "ple_key_w",
+ shape=lambda d: (d.stream, d.hidden),
+ src=_ple("key_proj.weight"),
+ layers=None,
+ ),
+ W(
+ "ple_val_w",
+ shape=lambda d: (d.ple_embed_dim, d.hidden),
+ src=_ple("value_proj.weight"),
+ layers=None,
+ ),
+ W("ple_nk", shape=lambda d: (d.stream,), src=_ple("norm_key.weight"), layers=None),
+ W("ple_nq", shape=lambda d: (d.stream,), src=_ple("norm_query.weight"), layers=None),
+ W("ple_nc", shape=lambda d: (d.stream,), src=_ple("norm_conv.weight"), layers=None),
+ W(
+ "ple_conv_w",
+ shape=lambda d: (d.stream, 1, d.ple_conv_kernel),
+ src=_ple("conv1d.weight"),
+ layers=None,
+ ),
+ W(
+ "ple_mult",
+ shape=lambda d: (d.ngram_size,),
+ dtype=i64,
+ src=_ple("ple_embedding.layer_multipliers"),
+ layers=None,
+ ),
+ W(
+ "ple_vocab",
+ shape=lambda d: (d.ngram_heads,),
+ dtype=i64,
+ src=_ple("ple_embedding.ngram_heads_vocab_sizes"),
+ layers=None,
+ ),
+ W(
+ "ple_offs",
+ shape=lambda d: (d.ngram_heads,),
+ dtype=i64,
+ src=_ple("ple_embedding.ngram_heads_offsets"),
+ layers=None,
+ ),
+ # The table's per-tensor scale sits with the **table**, one level below the three
+ # metadata tensors above: the shards and this scale are the quantized
+ # `ngram_embedding` submodule, while the multipliers, offsets and vocabulary sizes
+ # describe the id space and stay on `ple_embedding` itself.
+ W(
+ "ple_tscale",
+ shape=lambda d: (1,),
+ src=lambda d, i: [(f"{_ngram_prefix(d.ple_layer)}.weight_scale", None)],
+ layers=None,
+ ),
+ # --- the vision tower: all bf16, replicated, one-to-one, no transform ---
+ W(
+ "vis_patch_w",
+ shape=lambda d: (d.v_hidden, d.v_channels, d.v_tpatch, d.v_patch, d.v_patch),
+ src=_visual("patch_embed.proj.weight"),
+ layers=None,
+ ),
+ W(
+ "vis_patch_b",
+ shape=lambda d: (d.v_hidden,),
+ src=_visual("patch_embed.proj.bias"),
+ layers=None,
+ ),
+ W(
+ "vis_pos",
+ shape=lambda d: (d.v_pos_rows, d.v_hidden),
+ src=_visual("pos_embed.weight"),
+ layers=None,
+ ),
+ # LayerNorm **with a bias**, unlike every norm in the language tower.
+ W(
+ "vis_n1_w",
+ shape=lambda d: (d.v_hidden,),
+ src=_vision_block("norm1.weight"),
+ layers=_vision_blocks,
+ ),
+ W(
+ "vis_n1_b",
+ shape=lambda d: (d.v_hidden,),
+ src=_vision_block("norm1.bias"),
+ layers=_vision_blocks,
+ ),
+ # One fused `[q | k | v]` projection, MHA, with a bias.
+ W(
+ "vis_qkv_w",
+ shape=lambda d: (3 * d.v_hidden, d.v_hidden),
+ src=_vision_block("attn.qkv.weight"),
+ layers=_vision_blocks,
+ ),
+ W(
+ "vis_qkv_b",
+ shape=lambda d: (3 * d.v_hidden,),
+ src=_vision_block("attn.qkv.bias"),
+ layers=_vision_blocks,
+ ),
+ W(
+ "vis_proj_w",
+ shape=lambda d: (d.v_hidden, d.v_hidden),
+ src=_vision_block("attn.proj.weight"),
+ layers=_vision_blocks,
+ ),
+ W(
+ "vis_proj_b",
+ shape=lambda d: (d.v_hidden,),
+ src=_vision_block("attn.proj.bias"),
+ layers=_vision_blocks,
+ ),
+ W(
+ "vis_n2_w",
+ shape=lambda d: (d.v_hidden,),
+ src=_vision_block("norm2.weight"),
+ layers=_vision_blocks,
+ ),
+ W(
+ "vis_n2_b",
+ shape=lambda d: (d.v_hidden,),
+ src=_vision_block("norm2.bias"),
+ layers=_vision_blocks,
+ ),
+ # A **non-gated** MLP: nothing here is split in half, unlike every SwiGLU in the
+ # language tower.
+ W(
+ "vis_fc1_w",
+ shape=lambda d: (d.v_inter, d.v_hidden),
+ src=_vision_block("mlp.linear_fc1.weight"),
+ layers=_vision_blocks,
+ ),
+ W(
+ "vis_fc1_b",
+ shape=lambda d: (d.v_inter,),
+ src=_vision_block("mlp.linear_fc1.bias"),
+ layers=_vision_blocks,
+ ),
+ W(
+ "vis_fc2_w",
+ shape=lambda d: (d.v_hidden, d.v_inter),
+ src=_vision_block("mlp.linear_fc2.weight"),
+ layers=_vision_blocks,
+ ),
+ W(
+ "vis_fc2_b",
+ shape=lambda d: (d.v_hidden,),
+ src=_vision_block("mlp.linear_fc2.bias"),
+ layers=_vision_blocks,
+ ),
+ # The patch merger. Its norm is applied **before** the 4-to-1 shuffle, so it is
+ # `hidden` wide and not `merge_dim` wide.
+ W(
+ "vis_m_n_w",
+ shape=lambda d: (d.v_hidden,),
+ src=_visual("merger.norm.weight"),
+ layers=None,
+ ),
+ W("vis_m_n_b", shape=lambda d: (d.v_hidden,), src=_visual("merger.norm.bias"), layers=None),
+ W(
+ "vis_m_fc1_w",
+ shape=lambda d: (d.v_merge_dim, d.v_merge_dim),
+ src=_visual("merger.linear_fc1.weight"),
+ layers=None,
+ ),
+ W(
+ "vis_m_fc1_b",
+ shape=lambda d: (d.v_merge_dim,),
+ src=_visual("merger.linear_fc1.bias"),
+ layers=None,
+ ),
+ W(
+ "vis_m_fc2_w",
+ shape=lambda d: (d.hidden, d.v_merge_dim),
+ src=_visual("merger.linear_fc2.weight"),
+ layers=None,
+ ),
+ W(
+ "vis_m_fc2_b",
+ shape=lambda d: (d.hidden,),
+ src=_visual("merger.linear_fc2.bias"),
+ layers=None,
+ ),
+ # --- the MTP draft module, at layer index num_hidden_layers, only when live ---
+ #
+ # Structurally one more decoder layer of this architecture with a two-branch
+ # front end; every shape equals a trunk full-attention layer's (cut the same way)
+ # except the front end and the expert quantization (FP8 128x128 block scales, not
+ # NVFP4). Its keys carry no `{p}`: the module's prefix is not the trunk's.
+ W(
+ "mtp_pre_e_norm",
+ shape=lambda d: (d.hidden,),
+ src=f"{_MTP}.pre_fc_norm_embedding.weight",
+ layers=_mtp_only,
+ ),
+ # `[hc_count * hidden]`: the module's input is the trunk's four-stream residual.
+ W(
+ "mtp_pre_h_norm",
+ shape=lambda d: (d.stream,),
+ src=f"{_MTP}.pre_fc_norm_hidden.weight",
+ layers=_mtp_only,
+ ),
+ W(
+ "mtp_fc_e",
+ shape=lambda d: (d.hidden, d.hidden),
+ src=f"{_MTP}.fc_embedding.weight",
+ layers=_mtp_only,
+ ),
+ W(
+ "mtp_fc_h",
+ shape=lambda d: (d.hidden, d.hidden),
+ src=f"{_MTP}.fc_hidden.weight",
+ layers=_mtp_only,
+ ),
+ *(
+ w
+ for tag, module in _HC_MODULES
+ for w in _hyper_connection(
+ f"mtp_{tag}",
+ lambda name, m=module: f"{_MTP_LAYER}.{m}.{name}.weight",
+ _mtp_only,
+ )
+ ),
+ *_attn_in_separate("mtp_", f"{_MTP_LAYER}.self_attn", _mtp_only),
+ *_attn_out_and_norms("mtp_", f"{_MTP_LAYER}.self_attn", _mtp_only),
+ *_index_branch("mtp_", f"{_MTP_LAYER}.self_attn.indexer", _mtp_only),
+ W(
+ "mtp_router",
+ shape=lambda d: (d.num_experts, d.hidden),
+ src=f"{_MTP_LAYER}.mlp.gate.weight",
+ layers=_mtp_only,
+ ),
+ # e4m3 codes stored row-major with one fp32 scale per 128x128 weight block, the
+ # FC1 halves concatenated `[up ; gate]` on the row axis; this rank's window.
+ W(
+ "mtp_fc1_w",
+ shape=lambda d: (d.mtp_local_experts, 2 * d.moe_inter, d.hidden),
+ dtype=f8,
+ src=_experts(_mtp_mlp, "up_proj.weight", "gate_proj.weight", window=_mtp_window),
+ transform=_mtp_fc1_w,
+ layers=_mtp_only,
+ ),
+ W(
+ "mtp_fc1_s",
+ shape=lambda d: (
+ d.mtp_local_experts,
+ (2 * d.moe_inter) // FP8_BLOCK,
+ d.hidden // FP8_BLOCK,
+ ),
+ dtype=f32,
+ src=_experts(
+ _mtp_mlp,
+ "up_proj.weight_scale_inv",
+ "gate_proj.weight_scale_inv",
+ window=_mtp_window,
+ ),
+ transform=_mtp_fc1_s,
+ layers=_mtp_only,
+ ),
+ # FC2 needs no relayout at all: the checkpoint's `[hidden, I]` e4m3 block is
+ # exactly the operand.
+ W(
+ "mtp_fc2_w",
+ shape=lambda d: (d.mtp_local_experts, d.hidden, d.moe_inter),
+ dtype=f8,
+ src=_experts(_mtp_mlp, "down_proj.weight", window=_mtp_window),
+ layers=_mtp_only,
+ ),
+ W(
+ "mtp_fc2_s",
+ shape=lambda d: (
+ d.mtp_local_experts,
+ d.hidden // FP8_BLOCK,
+ d.moe_inter // FP8_BLOCK,
+ ),
+ dtype=f32,
+ src=_experts(_mtp_mlp, "down_proj.weight_scale_inv", window=_mtp_window),
+ transform=_mtp_fc2_s,
+ layers=_mtp_only,
+ ),
+ *_shared_expert("mtp_", _mtp_mlp, _mtp_only),
+ # The module's own stack-final mixer, which is its `shared_head`.
+ *_mixer("mtp_fhc", f"{_MTP}.hyper_connection_mixer", _mtp_only),
+ )
+
+ def dims(self, core) -> SimpleNamespace:
+ """This rank's dimension bundle, from the core's normalized text config, the
+ composite's vision config (the core holds it as `vision_config` before it
+ declares anything), the engine's `Mapping` and the core's engine choices."""
+ return rank_dims(
+ core.model_config.pretrained_config,
+ core.vision_config,
+ core.model_config.mapping,
+ mtp_enabled=getattr(core.model_config, "spec_config", None) is not None,
+ trunk_moe_backend=core.trunk_moe_backend,
+ moe_prefill_copy=core.moe_prefill_copy,
+ dtype=core.model_config.torch_dtype,
+ )
+
+ def expected_unconsumed(self, d: SimpleNamespace, keys) -> set:
+ """The checkpoint keys this rank predicts it leaves unread: the tensors of every
+ routed expert outside its window (its per-tensor scalars are read on every rank;
+ under the intermediate split the window is every expert, each read at this rank's
+ slice), and, with MTP live, the draft experts outside its draft window; with MTP off, the
+ whole `mtp.` subtree. The windows are enumerated from `d`, not read off the
+ checkpoint, so a key that belongs to a real weight can never land here."""
+ out = set()
+ if d.expert_split in LOADED_EXPERT_SPLITS:
+ start, count = _trunk_window(d)
+ for i in range(d.num_layers):
+ for e in range(d.num_experts):
+ if not start <= e < start + count:
+ out.update(
+ f"{_trunk_mlp(i)}.experts.{e}.{t}" for t in _TRUNK_EXPERT_TENSORS
+ )
+ if not d.mtp_enabled:
+ return out | {k for k in keys if k.startswith(f"{_MTP}.")}
+ start, count = _mtp_window(d)
+ for e in range(d.num_experts):
+ if not start <= e < start + count:
+ out.update(f"{_mtp_mlp(None)}.experts.{e}.{t}" for t in _MTP_EXPERT_TENSORS)
+ return out
+
+ def declare(self, d: SimpleNamespace) -> nn.ParameterDict:
+ """Allocate every weight the table declares, at this rank's share.
+
+ Meta-init intercepts `torch.empty` here -- real CUDA storage arrives
+ when the engine materializes the registry.
+ """
+ w = nn.ParameterDict()
+ for entry in self.WEIGHTS:
+ dtype = entry.dtype or d.dtype
+ for key, _, _ in self._targets(entry, d):
+ w[key] = nn.Parameter(
+ torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False
+ )
+ return w
+
+ def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None, str]]:
+ """`[(param key, layer index | None, checkpoint layer prefix), ...]`.
+
+ One place decides which parameters an entry declares, so `declare` and
+ `manifest` cannot disagree about whether a weight exists.
+ """
+ layers = entry.layer_indices(d)
+ if layers is None:
+ return [(entry.name, None, "")]
+ return [(f"l{i}_{entry.name}", i, f"{_LM}.layers.{i}") for i in layers]
+
+ def manifest(self, d: SimpleNamespace) -> dict[str, list[tuple]]:
+ """param key -> [(checkpoint key, index into the param | None, cut | None,
+ transform | None)]."""
+ rows: dict[str, list[tuple]] = {}
+ for entry in self.WEIGHTS:
+ for key, i, prefix in self._targets(entry, d):
+ rows[key] = [
+ (ckpt, index, cut, entry.transform)
+ for ckpt, index, cut in entry.sources(d, i, prefix)
+ ]
+ return rows
+
+ def read(self, weights, ckpt_key, cut, transform, d, device=None) -> torch.Tensor:
+ """One manifest source -> the tensor its destination receives.
+
+ Each key of the source is read through its cut -- only this rank's rows or
+ columns leave the checkpoint -- moved to `device`, and handed to `transform`.
+ The one path `load` takes, so the shard test drives it directly.
+ """
+ multi = isinstance(ckpt_key, tuple)
+ keys = ckpt_key if multi else (ckpt_key,)
+ cuts = (cut if cut is not None else (None,) * len(keys)) if multi else (cut,)
+ assert len(cuts) == len(keys), (ckpt_key, cut)
+ for key in keys:
+ assert key in weights, f"checkpoint key missing: {key}"
+ src_all = tuple(_read(weights[key], c) for key, c in zip(keys, cuts))
+ if device is not None:
+ src_all = tuple(t.to(device, non_blocking=True) for t in src_all)
+ if transform is not None:
+ return transform(src_all if multi else src_all[0], d)
+ assert len(src_all) == 1, "a multi-key row needs a transform"
+ return src_all[0]
+
+ def load(self, model, weights) -> None:
+ """Fill every declared parameter. The n-gram table is declared here first, each
+ source is read through its cut, multi-key rows feed their transform several
+ tensors at once, and coverage is asserted in both directions."""
+ core = model.model
+ d = self.dims(core)
+ assert d.expert_split in LOADED_EXPERT_SPLITS, (
+ f"routed experts split by {d.expert_split!r}: this table declares the splits "
+ f"{LOADED_EXPERT_SPLITS}"
+ )
+ # The one load-time declaration, before the manifest is built: the table's
+ # height comes from the shards, and the manifest addresses it by row range.
+ rows_per_shard = _declare_ngram_table(core, d, weights)
+ manifest = self.manifest(d)
+ manifest["ple_table"] = ngram_table_sources(d, rows_per_shard)
+ consumed: set = set()
+
+ def fill(param: nn.Parameter, ckpt_key, index, cut, transform) -> None:
+ dst = param.data if index is None else param.data[index]
+ # Materialize where the destination lives: the expert relayouts are row
+ # gathers and the scale swizzle is a kernel, both on the device.
+ src = self.read(weights, ckpt_key, cut, transform, d, dst.device)
+ assert dst.shape == src.shape, (ckpt_key, tuple(dst.shape), tuple(src.shape))
+ assert src.dtype == dst.dtype, (ckpt_key, src.dtype, dst.dtype)
+ dst.copy_(src, non_blocking=True)
+ consumed.update(ckpt_key if isinstance(ckpt_key, tuple) else (ckpt_key,))
+
+ assert set(manifest.keys()) == set(core.w.keys()), (
+ "manifest/parameter drift",
+ set(manifest.keys()) ^ set(core.w.keys()),
+ )
+ for param_key, sources in manifest.items():
+ for ckpt_key, index, cut, transform in sources:
+ fill(core.w[param_key], ckpt_key, index, cut, transform)
+ if param_key.endswith(_RELEASE_AFTER):
+ torch.cuda.empty_cache()
+
+ # The shell owns lm_head (untied, bf16 on this checkpoint); read the rows it built
+ # (`lm_head_share`) and nothing else.
+ cut, pad = lm_head_share(model.lm_head.weight.shape[0], d.vocab, d.tp_rank)
+ assert tuple(model.lm_head.weight.shape[1:]) == (d.hidden,), (
+ f"the shell built lm_head {tuple(model.lm_head.weight.shape)}; this checkpoint's "
+ f"is a dense [{d.vocab}, {d.hidden}] weight"
+ )
+ fill(model.lm_head.weight, "lm_head.weight", None, cut, pad)
+ # Parameter-side coverage, the other direction: name every registered parameter
+ # outside the target's own ParameterDict instead of assuming lm_head is the only
+ # one.
+ shell = {n for n, _ in model.named_parameters() if not n.startswith("model.w.")}
+ assert shell == {"lm_head.weight"}, f"unfed shell parameters: {sorted(shell)}"
+
+ torch.cuda.synchronize()
+ leftover = set(weights.keys()) - consumed
+ expected = self.expected_unconsumed(d, weights.keys())
+ assert leftover == expected, (
+ "checkpoint coverage: leftover keys are not exactly this rank's off-window "
+ f"experts{'' if d.mtp_enabled else ' and the MTP module'}; unexpected "
+ f"{sorted(leftover - expected)[:8]}, missed {sorted(expected - leftover)[:8]}"
+ )
+
+
+MODEL_WEIGHTS = Qwen4ExpDefaultWeights()
+
+
+def _heads(total: int, ranks: int, rank: int) -> tuple[int, int]:
+ """(heads on `rank`, its first global head) when `total` heads are split over `ranks`:
+ an even split, or each head repeated on `ranks // total` neighbouring ranks when there
+ are more ranks than heads -- the rule the K/V and recurrent pools follow
+ (`pool.pool_geometry`)."""
+ assert total % ranks == 0 or ranks % total == 0, (total, ranks)
+ return -(-total // ranks), rank * total // ranks
+
+
+def rank_dims(
+ cfg,
+ vcfg,
+ mapping,
+ *,
+ mtp_enabled: bool,
+ trunk_moe_backend: str,
+ moe_prefill_copy: bool,
+ dtype: torch.dtype,
+) -> SimpleNamespace:
+ """Every width the table declares at rank `mapping.rank`, and the shares behind them.
+
+ Args:
+ cfg: the engine-normalized text config (`Qwen4ExpTextConfig`).
+ vcfg: the composite's vision config.
+ mapping: this rank's `Mapping`.
+ mtp_enabled: whether the engine has a `speculative_config` (the MTP module is live).
+ trunk_moe_backend: the routed backend the trunk runs, which picks the expert layout.
+ moe_prefill_copy: whether the core keeps the CuteDSL-layout copy of its window.
+ dtype: the activation dtype, which every unquantized weight is declared in.
+ """
+ hidden = cfg.hidden_size
+ head_dim = cfg.head_dim
+ index_heads, index_kv_heads = cfg.indexer_n_heads, cfg.indexer_kv_heads
+
+ attn_tp = attention_tp_size(mapping)
+ attn_rank = mapping.tp_rank if attn_tp > 1 else 0
+ heads_q, q_head_start = _heads(cfg.num_attention_heads, attn_tp, attn_rank)
+ heads_kv, kv_head_start = _heads(cfg.num_key_value_heads, attn_tp, attn_rank)
+
+ state_tp = mamba_effective_tp_size(mapping)
+ state_rank = mapping.tp_rank if state_tp > 1 else 0
+ num_k_heads, k_head_start = _heads(cfg.linear_num_key_heads, state_tp, state_rank)
+ num_v_heads, v_head_start = _heads(cfg.linear_num_value_heads, state_tp, state_rank)
+ key_dim = num_k_heads * cfg.linear_key_head_dim
+ value_dim = num_v_heads * cfg.linear_value_head_dim
+
+ num_experts = cfg.num_experts
+ # The routed experts' MoE group is the TP group: tp4 splits it by intermediate,
+ # tep4 and dep4 by expert.
+ assert mapping.moe_tp_size * mapping.moe_ep_size == mapping.tp_size, (
+ mapping.moe_tp_size,
+ mapping.moe_ep_size,
+ mapping.tp_size,
+ )
+ assert num_experts % mapping.moe_ep_size == 0 and num_experts % mapping.tp_size == 0, (
+ num_experts,
+ mapping.moe_ep_size,
+ mapping.tp_size,
+ )
+ local_experts = num_experts // mapping.moe_ep_size
+ mtp_local_experts = num_experts // mapping.tp_size
+ # Each expert's intermediate over the MoE TP ranks (tp4; one rank, the whole of it,
+ # otherwise). A share holds whole group-16 scale vectors, so `down_proj`'s scale
+ # columns cut on a vector boundary; it is stored zero-padded to EXPERT_INTER_ALIGN.
+ moe_inter = cfg.moe_intermediate_size
+ assert moe_inter % mapping.moe_tp_size == 0, (moe_inter, mapping.moe_tp_size)
+ expert_inter_share = moe_inter // mapping.moe_tp_size
+ assert expert_inter_share % SF_VEC == 0, (expert_inter_share, SF_VEC)
+ expert_inter = -(-expert_inter_share // EXPERT_INTER_ALIGN) * EXPERT_INTER_ALIGN
+
+ ngram_heads = cfg.heads_per_ngram * (cfg.ngram_size - 1)
+ v_hidden = vcfg.hidden_size
+ return SimpleNamespace(
+ num_layers=cfg.num_hidden_layers,
+ is_attn=[t == "full_attention" for t in cfg.layer_types],
+ hidden=hidden,
+ hc_count=cfg.hc_count,
+ hc_lowrank=cfg.hc_lowrank,
+ stream=cfg.hc_count * hidden,
+ tp_rank=mapping.tp_rank,
+ # --- attention, at this rank's heads ---
+ head_dim=head_dim,
+ heads_q=heads_q,
+ q_head_start=q_head_start,
+ heads_kv=heads_kv,
+ kv_head_start=kv_head_start,
+ q_width=heads_q * head_dim,
+ kv_width=heads_kv * head_dim,
+ index_head_dim=cfg.indexer_head_dim,
+ index_width=(index_heads + index_kv_heads) * cfg.indexer_head_dim,
+ # [q | k | v | gate | index q, k | pad of one index head]
+ attn_in_rows=2 * heads_q * head_dim
+ + 2 * heads_kv * head_dim
+ + (index_heads + index_kv_heads + 1) * cfg.indexer_head_dim,
+ # --- Gated-DeltaNet, at this rank's key and value heads ---
+ num_k_heads=num_k_heads,
+ k_head_start=k_head_start,
+ num_v_heads=num_v_heads,
+ v_head_start=v_head_start,
+ head_k_dim=cfg.linear_key_head_dim,
+ head_v_dim=cfg.linear_value_head_dim,
+ key_dim=key_dim,
+ key_dim_total=cfg.linear_num_key_heads * cfg.linear_key_head_dim,
+ value_dim=value_dim,
+ conv_dim=2 * key_dim + value_dim,
+ in_proj_split=(2 * key_dim + value_dim, value_dim, num_v_heads, num_v_heads),
+ conv_kernel=cfg.linear_conv_kernel_dim,
+ # --- MoE ---
+ num_experts=num_experts,
+ expert_split="ep" if mapping.moe_tp_size == 1 else "tp",
+ local_experts=local_experts,
+ expert_offset=local_experts * mapping.moe_ep_rank,
+ # This rank's intermediate rows `[expert_inter_start, + expert_inter_share)` of
+ # every expert it holds, stored `expert_inter` wide (the runner's
+ # `intermediate_size`).
+ expert_inter_share=expert_inter_share,
+ expert_inter_start=expert_inter_share * mapping.moe_tp_rank,
+ expert_inter=expert_inter,
+ mtp_local_experts=mtp_local_experts,
+ mtp_expert_offset=mtp_local_experts * mapping.tp_rank,
+ moe_inter=moe_inter,
+ shared_inter=cfg.shared_expert_intermediate_size,
+ vocab=cfg.vocab_size,
+ # --- PLE ---
+ ple_layer=list(cfg.ple_layer_ids)[0] - 1, # one-based in the config
+ ple_embed_dim=cfg.ple_embed_dim,
+ ple_conv_kernel=cfg.ple_conv_kernel_size,
+ ngram_size=cfg.ngram_size,
+ ngram_heads=ngram_heads,
+ ple_head_dim=cfg.ple_embed_dim // ngram_heads,
+ ple_shards=cfg.split_ngram_parts,
+ ple_table_dtype=PLE_TABLE_DTYPE,
+ # --- the vision tower ---
+ v_depth=vcfg.depth,
+ v_hidden=v_hidden,
+ v_inter=vcfg.intermediate_size,
+ v_channels=vcfg.in_channels,
+ v_tpatch=vcfg.temporal_patch_size,
+ v_patch=vcfg.patch_size,
+ v_pos_rows=vcfg.num_position_embeddings,
+ v_merge_dim=v_hidden * vcfg.spatial_merge_size * vcfg.spatial_merge_size,
+ # --- engine choices ---
+ mtp_enabled=mtp_enabled,
+ trunk_moe_backend=trunk_moe_backend,
+ moe_prefill_copy=moe_prefill_copy,
+ dtype=dtype,
+ )
+
+
+def lm_head_share(rows: int, vocab: int, tp_rank: int) -> tuple[Cut, Callable]:
+ """The cut and the pad that fill an lm_head the shell built with `rows` rows.
+
+ Whole under attention DP (`rows == vocab`); a vocab-parallel shard under attention TP.
+ A vocab that does not split evenly leaves the last shard short, and its missing rows
+ are zero, as the built-in LMHead pads them -- so the cut stops at `vocab` and the pad
+ fills the rest.
+ """
+ start = 0 if rows == vocab else tp_rank * rows
+ stop = min(start + rows, vocab)
+ assert start < stop, (rows, vocab, tp_rank)
+
+ def pad(t: torch.Tensor, d) -> torch.Tensor:
+ if t.shape[0] == rows:
+ return t.contiguous()
+ out = torch.zeros(rows, t.shape[1], dtype=t.dtype, device=t.device)
+ out[: t.shape[0]] = t
+ return out
+
+ return (slice(start, stop),), pad
+
+
+def ngram_shard_keys(d: SimpleNamespace) -> list[str]:
+ """The n-gram table's shard keys, in concatenation order."""
+ prefix = _ngram_prefix(d.ple_layer)
+ return [f"{prefix}.shard_{s}.weight" for s in range(d.ple_shards)]
+
+
+def ngram_table_sources(d: SimpleNamespace, rows_per_shard: int) -> list[tuple]:
+ """The `ple_table` manifest rows: shard `s` fills rows `[s * rows_per_shard, ...)`,
+ read whole on every rank (the table is replicated)."""
+ return [
+ (key, (_span(s * rows_per_shard, rows_per_shard),), None, None)
+ for s, key in enumerate(ngram_shard_keys(d))
+ ]
+
+
+def _declare_ngram_table(core, d: SimpleNamespace, weights) -> int:
+ """Register the n-gram table with the height its shards concatenate to, and return
+ the rows each shard holds.
+
+ The shards are **contiguous row chunks**, not an interleaving: global row `r` lives in
+ shard `r // rows_per_shard` at local row `r % rows_per_shard`. That is what makes the
+ flat buffer indexable by the hash's global rows.
+
+ **This is the whole of PLE host offload.** `core.ple_host_offload` picks between an
+ ordinary device allocation and host memory the device addresses; shape, dtype,
+ `ParameterDict` key, the shard-by-shard fill and the forward's gather are the same either
+ way, so an offloaded table is loaded and covered in both directions like any other. The
+ core asserts where it landed against the driver.
+ """
+ keys = ngram_shard_keys(d)
+ for key in keys:
+ assert key in weights, f"checkpoint key missing: {key}"
+ shapes = {tuple(weights[key].shape) for key in keys}
+ assert len(shapes) == 1, (
+ f"the n-gram table's {len(keys)} shards do not share one shape: "
+ f"{sorted(shapes)}; they are concatenated as equal contiguous row "
+ f"chunks and an uneven one would shift every row above it"
+ )
+ rows_per_shard, head_dim = shapes.pop()
+ assert head_dim == d.ple_head_dim, (
+ f"the n-gram table is {head_dim} wide but {d.ngram_heads} heads must "
+ f"tile ple_embed_dim {d.ple_embed_dim}, i.e. {d.ple_head_dim}"
+ )
+ dtype = weights[keys[0]].dtype
+ assert dtype == d.ple_table_dtype, (
+ f"the n-gram table is stored {dtype}; this target declares "
+ f"{d.ple_table_dtype} and does not convert it"
+ )
+ device = core.w["embed"].device
+ rows = d.ple_shards * rows_per_shard
+ if core.ple_host_offload:
+ table, core._ple_table_host = _host_resident_empty(rows, head_dim, dtype, device)
+ else:
+ table = torch.empty(rows, head_dim, dtype=dtype, device=device)
+ core._ple_table_host = None
+ core.w["ple_table"] = nn.Parameter(table, requires_grad=False)
+ # `nn.Parameter` (`_make_subclass`) shares storage rather than copying; host offload
+ # depends on that, so it is asserted.
+ assert core.w["ple_table"].data_ptr() == table.data_ptr(), (
+ "wrapping the n-gram table in nn.Parameter moved its storage"
+ )
+ core.ple_rows_per_shard = rows_per_shard
+ return rows_per_shard
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/kernels/moe/blockscaled_contiguous_grouped_gemm_finalize_fusion.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/kernels/moe/blockscaled_contiguous_grouped_gemm_finalize_fusion.py
deleted file mode 100644
index 078eced39888..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/kernels/moe/blockscaled_contiguous_grouped_gemm_finalize_fusion.py
+++ /dev/null
@@ -1,2738 +0,0 @@
-# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: BSD-3-Clause
-
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-
-# 1. Redistributions of source code must retain the above copyright notice, this
-# list of conditions and the following disclaimer.
-
-# 2. Redistributions in binary form must reproduce the above copyright notice,
-# this list of conditions and the following disclaimer in the documentation
-# and/or other materials provided with the distribution.
-
-# 3. Neither the name of the copyright holder nor the names of its
-# contributors may be used to endorse or promote products derived from
-# this software without specific prior written permission.
-
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-"""The target's copy of TRT-LLM's CuTe DSL finalize-fused grouped GEMM (the CuteDSL MoE
-FC2), with three options the target's prefill MoE sets (`fp4_moe_renorm_cute_dsl`): the
-MMA instructions per K stage (`mma_inst_tile_k`), an epilogue in which each warp owns
-whole output rows (`row_owned_epilogue`), and L2 cache hints on the operand loads and the
-scatter-add (`l2_hints`). Their defaults keep the shared kernel's behavior."""
-
-from typing import Tuple, Type, Union
-
-import cuda.bindings.driver as cuda
-import cutlass
-import cutlass.cute as cute
-import cutlass.pipeline as pipeline
-import cutlass.utils as utils
-import cutlass.utils.blackwell_helpers as sm100_utils
-import cutlass.utils.blockscaled_layout as blockscaled_utils
-from cutlass._mlir.dialects import llvm
-from cutlass.cute.nvgpu import cpasync, tcgen05
-from cutlass.cutlass_dsl import dsl_user_op
-
-from tensorrt_llm._torch.kernels.blackwell.utils import (
- TRTLLM_ENABLE_PDL,
- blk_reduce_bf16,
- blk_reduce_fp16,
- blk_reduce_fp32,
- griddepcontrol_launch_dependents,
- griddepcontrol_wait,
- is_power_of_2,
-)
-
-#: L2 cache-hint policies (CUTLASS CacheHintSm90 encodings).
-_EVICT_FIRST = 0x12F0000000000000
-_EVICT_LAST = 0x14F0000000000000
-
-
-@dsl_user_op
-def _blk_reduce_bf16_evict_last(dst_gemm, src_smem, size, loc=None, ip=None):
- llvm.inline_asm(
- None,
- [
- dst_gemm.iterator.llvm_ptr,
- src_smem.iterator.llvm_ptr,
- size.ir_value(),
- cutlass.Int64(_EVICT_LAST).ir_value(),
- ],
- "cp.reduce.async.bulk.global.shared::cta.bulk_group.L2::cache_hint.add.noftz.bf16 "
- "[$0], [$1], $2, $3;",
- "l,l,r,l",
- has_side_effects=True,
- loc=loc,
- ip=ip,
- )
-
-
-"""
-High-performance persistent blockscaled contiguous grouped dense GEMM (C = alpha * (SFA * A) * (SFB * B)) example for
-the NVIDIA Blackwell architecture using CUTE DSL.
-- Matrix A is MxKx1, A can be row-major("K"), ValidM is composed of valid m in different groups
-- Matrix B is NxKxL, B can be column-major("K"), L is grouped dimension
-- Matrix C is SxNX1, C can be row-major("N"), ValidM is composed of valid m in different groups
-- Matrix SFA layout is filled internally according to A shape and BlockScaledBasicChunk, which has
- M x ceil_div(K, sf_vec_size) x L elements respectively
-- Matrix SFB layout is filled internally according to B shape and BlockScaledBasicChunk, which has
- N x ceil_div(K, sf_vec_size) x L elements respectively
-
-Matrix A/C Memory Layout Diagrams:
-
- ```
- Group 0 Group 1 Group 2
- -+---------+---------+---------+
- | | | |
- K| ValidM0 | ValidM1 | ValidM2 |
- | | | |
- -+---------+---------+---------+
- |<- ValidM ->|
- ```
- Note: the Group(L) dimension will be flatted into M dimension, and the rest Group(L) size is 1.
- each ValidM will be aligned to 256 or 128. The alignment is determined by the mma_tiler_mn parameter.
- For NVFP4, 2CTA, the alignment is 256. For NVFP4, 1CTA, the alignment is 128.
-
-This GEMM kernel supports the following features:
- - Utilizes Tensor Memory Access (TMA) for efficient memory operations
- - Utilizes Blackwell's tcgen05.mma for matrix multiply-accumulate (MMA) operations
- - Implements TMA multicast with cluster to reduce L2 memory traffic
- - Support persistent tile scheduling to better overlap memory load/store with mma between tiles
- - Support warp specialization to avoid explicit pipelining between mainloop load and mma
-
-This GEMM works as follows:
-1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations.
-2. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using non-TMA
- operations.
-2. MMA warp:
- - Load scale factor A/B from shared memory (SMEM) to tensor memory (TMEM) using tcgen05.cp instruction.
- - Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction.
-3. EPILOGUE warp (with Fused Finalize for MoE):
- - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld.
- - Apply alpha scaling: acc_scaled = alpha * acc
- - **Fused Finalize Logic** (following TensorRT-LLM's sm90_visitor_scatter.hpp pattern):
- a) Use permuted_idx_to_expanded_idx to map from permuted row to token/topk indices
- b) Load router_scale directly from global memory to register (no shared memory)
- c) Apply router_scale: Final = router_scale * acc_scaled
- - Type convert Final matrix to output type.
- - Store C matrix from registers (RMEM) to shared memory (SMEM) to global memory (GMEM) with TMA operations.
-
-SM100 tcgen05.mma.kind.block_scale instructions operate as follows:
-- Read matrix A from SMEM
-- Read matrix B from SMEM
-- Read scalefactor A from TMEM
-- Read scalefactor B from TMEM
-- Write accumulator to TMEM
-The accumulator in TMEM must then be loaded to registers before writing back to GMEM.
-
-.. code-block:: bash
-
- python blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py \
- --ab_dtype Float4E2M1FN --out_dtype BFloat16 \
- --sf_dtype Float8E4M3FN --sf_vec_size 16 \
- --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \
- --benchmark 1024x7168x2048x64
-
-To collect performance with NCU profiler:
-
-.. code-block:: bash
-
- ncu python blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py \\ \
- --ab_dtype Float4E2M1FN --out_dtype BFloat16 \
- --sf_dtype Float8E4M3FN --sf_vec_size 16 \
- --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \
- --benchmark [80,120,160]x7168x2048x64
-
-Constraints:
-* Supported input data types: mxf8, mxf4, nvf4
- see detailed valid dtype combinations in below Sm100BlockScaledPersistentDenseGemmKernel class documentation
-* A/B tensor must have the same data type, mixed data type is not supported (e.g., mxf8 x mxf4)
-* Mma tiler M must be 128 or 256(use_2cta_instrs)
-* Mma tiler N must be 64/128/192/256
-* Cluster shape M/N must be positive and power of 2, total cluster size <= 16
-* Cluster shape M must be multiple of 2 if Mma tiler M is 256(use_2cta_instrs)
-* The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned,
- i.e, number of elements is a multiple of 16 and 32 for Float8 and Float4, respectively.
-
-CUDA Graph Support:
-* For CUDA graph support, the tile_idx_to_expert_idx, A/C matrices, and scale factor A can be padded to a larger size
- (e.g., permuted_m = m*topK + num_local_experts*(256-1), example: 4096*8 + (256/32)*255 = 34808)
-* Use create_tensors() with permuted_m parameter to automatically pad:
- - tile_idx_to_expert_idx: padded for invalid tiles
- - A matrix: padded to permuted_m rows (padding rows contain dummy data)
- - C matrix: padded to permuted_m rows (output buffer for cuda_graph)
- - Scale factor A: padded to match A matrix dimensions
-* Kernel handling of padding (similar to masked_grouped_gemm.py):
- - Scheduler warp checks if tile_idx >= num_non_exiting_tiles to exit
- - Only valid tiles (tile_idx < num_non_exiting_tiles) are written to tile_info pipeline
- - When no more valid tiles exist, outer loop exits and calls producer_tail()
- - Consumer warps process only valid tiles from pipeline
- - No deadlock or synchronization issues
-* Consumer warps check initial tile against num_non_exiting_tiles and set is_valid_tile=False if
- tile_idx >= num_non_exiting_tiles
-* Only rows within (aligned_groupm[0]+aligned_groupm[1]+...) contain valid data
-* Padding rows in C matrix will not be written by the kernel
-"""
-
-
-class Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel:
- """This class implements batched matrix multiplication (C = A x SFA x B x SFB) with support for various data types
- and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization.
-
- :param sf_vec_size: Scalefactor vector size.
- :type sf_vec_size: int
- :param mma_tiler_mn: Shape of the Matrix Multiply-Accumulate (MMA) tile (M,N)
- :type mma_tiler_mn: Tuple[int, int]
- :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing
- :type cluster_shape_mn: Tuple[int, int]
-
- :note: In current version, A and B tensor must have the same data type
- - i.e., Float8E4M3FN for A and Float8E5M2 for B is not supported
-
- :note: Supported combinations of A/B data types, SF data typs and SF vector size:
- - MXF8: A/B: Float8E5M2/Float8E4M3FN + SF: Float8E8M0FNU + sf_vec_size: 32
- - MXF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU + sf_vec_size: 32
- - NVF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU/Float8E4M3FN + sf_vec_size: 16
-
- :note: Supported accumulator data types:
- - Float32
-
- :note: Supported C data types:
- - Float32
- - Float16/BFloat16
- - Float8E4M3FN/Float8E5M2
-
- :note: Constraints:
- - MMA tiler M must be 128 or 256 (use_2cta_instrs)
- - MMA tiler N must be 64/128/192/256
- - Cluster shape M must be multiple of 2 if Mma tiler M is 256
- - Cluster shape M/N must be positive and power of 2, total cluster size <= 16
- - Also, Cluster shape M/N must be <= 4 for scale factor multicasts due to limited size of scale factors
-
- Example:
- >>> gemm = Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel(
- ... sf_vec_size=16, mma_tiler_mn=(256, 128), cluster_shape_mn=(2, 1)
- ... )
- >>> gemm(
- ... a_tensor, b_tensor, sfa_tensor, sfb_tensor, out_tensor, max_active_clusters, stream
- ... )
- """
-
- def __init__(
- self,
- sf_vec_size: int,
- mma_tiler_mn: Tuple[int, int],
- cluster_shape_mn: Tuple[int, int],
- raster_along_m: bool = False,
- use_expert_counts: bool = False,
- num_local_experts: int = 0,
- mma_inst_tile_k: int = 4,
- l2_hints: bool = False,
- row_owned_epilogue: bool = False,
- ):
- """Initializes the configuration for a Blackwell blockscaled dense GEMM kernel.
-
- This configuration includes several key aspects:
-
- 1. MMA Instruction Settings (tcgen05):
- - mma_tiler_mn: The (M, N) shape of the MMA instruction tiler.
-
- 2. Cluster Shape:
- - cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster.
-
- :param mma_tiler_mn: Tuple (M, N) shape of the MMA instruction.
- :type mma_tiler_mn: Tuple[int, int]
- :param cluster_shape_mn: Tuple (ClusterM, ClusterN) shape of the cluster.
- :type cluster_shape_mn: Tuple[int, int]
- :param raster_along_m: Boolean, True to use raster along M.
- :type raster_along_m: bool
- """
-
- self.sf_vec_size = sf_vec_size
- self.acc_dtype = cutlass.Float32
- self.use_2cta_instrs = mma_tiler_mn[0] == 256
- self.cluster_shape_mn = cluster_shape_mn
- self.raster_along_m = raster_along_m
- # MMA instructions (K 64 each for NVFP4) per A/B pipeline stage. At N 256 the 66 KB C
- # staging buffer leaves room for only 2 stages of the default K 256, so with a short K
- # (an FC2 of K 640 is 2.5 such tiles) the TMA producer cannot run ahead of the MMA;
- # 2 (K 128) fits 5 stages.
- self.mma_inst_tile_k = mma_inst_tile_k
- # Operand TMA loads EVICT_FIRST and the bulk scatter-add EVICT_LAST, so the [T, H]
- # output the reductions hit is kept in L2 over the streamed operands.
- self.l2_hints = l2_hints
- # Each epilogue thread writes only its own row of the C staging buffer and bulk-reduces
- # only that row, so it can wait for its own previous bulk read right before overwriting
- # the row instead of a per-tile bulk wait + CTA barrier.
- self.row_owned_epilogue = row_owned_epilogue
- self.use_expert_counts = use_expert_counts
- self.num_local_experts = num_local_experts
- self.num_tile_info_fields = 6 if self.use_expert_counts else 5
- if self.use_expert_counts:
- assert self.num_local_experts > 0
- # K dimension is deferred in _setup_attributes
- self.mma_tiler = (*mma_tiler_mn, 1)
-
- self.cta_group = tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE
-
- self.occupancy = 1
- self.epilog_warp_id = (0, 1, 2, 3)
- self.mma_warp_id = 4
- self.tma_warp_id = 5
- self.sched_warp_id = 6
- self.meta_load_warp_id = 7
- self.threads_per_warp = 32
- self.threads_per_cta = self.threads_per_warp * len(
- (
- *self.epilog_warp_id,
- self.mma_warp_id,
- self.tma_warp_id,
- self.sched_warp_id,
- self.meta_load_warp_id,
- )
- )
- self.threads_wo_sched = self.threads_per_warp * len(
- (
- *self.epilog_warp_id,
- self.mma_warp_id,
- self.tma_warp_id,
- self.meta_load_warp_id,
- )
- )
-
- # Set barrier for cta sync, epilogue sync and tmem ptr sync
- self.cta_sync_barrier = pipeline.NamedBarrier(
- barrier_id=1,
- num_threads=self.threads_per_cta,
- )
- self.epilog_sync_barrier = pipeline.NamedBarrier(
- barrier_id=2,
- num_threads=32 * len(self.epilog_warp_id),
- )
- self.tmem_alloc_barrier = pipeline.NamedBarrier(
- barrier_id=3,
- num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)),
- )
- self.sched_sync_barrier = pipeline.NamedBarrier(
- barrier_id=4,
- num_threads=self.threads_per_warp,
- )
- self.num_smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
- # TMEM offset for final accumulator
- self.tmem_final_offset = 384
-
- def _setup_attributes(self):
- """Set up configurations that are dependent on GEMM inputs
-
- This method configures various attributes based on the input tensor properties
- (data types, leading dimensions) and kernel settings:
- - Configuring tiled MMA
- - Computing MMA/cluster/tile shapes
- - Computing cluster layout
- - Computing multicast CTAs for A/B
- - Computing epilogue subtile
- - Setting up A/B/C stage counts in shared memory
- - Computing A/B/C shared memory layout
- - Computing tensor memory allocation columns
- """
- self.mma_inst_shape_mn = (
- self.mma_tiler[0],
- self.mma_tiler[1],
- )
- # (CTA_Tile_Shape_M, Round_Up(MMA_Tile_Shape_N, 128), MMA_Inst_Shape_K)
- self.mma_inst_shape_mn_sfb = (
- self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1),
- cute.round_up(self.mma_inst_shape_mn[1], 128),
- )
-
- # Configure tiled mma
- tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
- self.a_dtype,
- self.a_major_mode,
- self.b_major_mode,
- self.sf_dtype,
- self.sf_vec_size,
- self.cta_group,
- self.mma_inst_shape_mn,
- )
-
- tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
- self.a_dtype,
- self.a_major_mode,
- self.b_major_mode,
- self.sf_dtype,
- self.sf_vec_size,
- cute.nvgpu.tcgen05.CtaGroup.ONE,
- self.mma_inst_shape_mn_sfb,
- )
-
- # Compute mma/cluster/tile shapes
- mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
- mma_inst_tile_k = self.mma_inst_tile_k
- self.mma_tiler = (
- self.mma_tiler[0],
- self.mma_tiler[1],
- mma_inst_shape_k * mma_inst_tile_k,
- )
-
- self.mma_tiler_sfb = (
- self.mma_inst_shape_mn_sfb[0],
- self.mma_inst_shape_mn_sfb[1],
- mma_inst_shape_k * mma_inst_tile_k,
- )
-
- self.cta_tile_shape_mnk = (
- self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
- self.mma_tiler[1],
- self.mma_tiler[2],
- )
-
- self.cta_tile_shape_mnk_sfb = (
- self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape),
- self.mma_tiler_sfb[1],
- self.mma_tiler_sfb[2],
- )
-
- # Compute cluster layout
- self.cluster_layout_vmnk = cute.tiled_divide(
- cute.make_layout((*self.cluster_shape_mn, 1)),
- (tiled_mma.thr_id.shape,),
- )
-
- self.cluster_layout_sfb_vmnk = cute.tiled_divide(
- cute.make_layout((*self.cluster_shape_mn, 1)),
- (tiled_mma_sfb.thr_id.shape,),
- )
-
- # Compute number of multicast CTAs for A/B
- self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2])
- self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1])
- self.is_a_mcast = self.num_mcast_ctas_a > 1
- self.is_b_mcast = self.num_mcast_ctas_b > 1
-
- # Compute epilogue subtile
- self.epi_tile = sm100_utils.compute_epilogue_tile_shape(
- self.cta_tile_shape_mnk,
- self.use_2cta_instrs,
- self.gemm_output_layout,
- self.out_dtype,
- )
-
- self.epi_tile_n = cute.size(self.epi_tile[1])
-
- # Setup A/B/C/Scale stage count in shared memory and ACC stage count in tensor memory
- (
- self.num_acc_stage,
- self.num_ab_stage,
- self.num_c_stage,
- self.num_tile_stage,
- self.num_meta_stage,
- ) = self._compute_stages(
- tiled_mma,
- self.mma_tiler,
- self.a_dtype,
- self.b_dtype,
- self.out_dtype,
- self.cta_tile_shape_mnk,
- self.sf_dtype,
- self.sf_vec_size,
- self.final_scale_dtype,
- self.num_smem_capacity,
- self.occupancy,
- )
-
- # Compute A/B/C/Scale shared memory layout
- self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
- tiled_mma,
- self.mma_tiler,
- self.a_dtype,
- self.num_ab_stage,
- )
- self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
- tiled_mma,
- self.mma_tiler,
- self.b_dtype,
- self.num_ab_stage,
- )
- self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
- tiled_mma,
- self.mma_tiler,
- self.sf_vec_size,
- self.num_ab_stage,
- )
- self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(
- tiled_mma,
- self.mma_tiler,
- self.sf_vec_size,
- self.num_ab_stage,
- )
-
- swizzled_pad = 16 // (self.out_dtype.width // 8)
- self.c_smem_layout_staged = cute.make_layout(
- (self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[1], self.num_c_stage),
- stride=(
- self.cta_tile_shape_mnk[1] + swizzled_pad,
- 1,
- self.cta_tile_shape_mnk[0] * (self.cta_tile_shape_mnk[1] + 8),
- ),
- )
-
- # Overlap and double buffer accumulator when num_acc_stage == 1 for cta_tile_n = 256 case
- self.overlapping_accum = self.num_acc_stage == 1
-
- sf_atom_mn = 32
-
- self.num_sfa_tmem_cols = (self.cta_tile_shape_mnk[0] // sf_atom_mn) * mma_inst_tile_k
- self.num_sfb_tmem_cols = (self.cta_tile_shape_mnk_sfb[1] // sf_atom_mn) * mma_inst_tile_k
- self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols
- self.num_accumulator_tmem_cols = (
- self.cta_tile_shape_mnk[1] * self.num_acc_stage
- if not self.overlapping_accum
- else self.cta_tile_shape_mnk[1] * 2 - self.num_sf_tmem_cols
- )
-
- self.iter_acc_early_release_in_epilogue = self.num_sf_tmem_cols // self.epi_tile_n
-
- # Compute the number of tensor memory allocation columns
- self.num_tmem_alloc_cols = 512
-
- @cute.jit
- def _expert_count_tile_info(
- self,
- expert_counts: cute.Tensor,
- tile_idx: cutlass.Int32,
- expert_capacity: cutlass.Int32,
- ):
- """Resolve a compact tile directly from expert-major receive counts."""
- tile_size = cutlass.Int32(self.mma_tiler[0])
- num_valid_tiles = cutlass.Int32(0)
- expert_idx = cutlass.Int32(-1)
- expert_tile_start = cutlass.Int32(0)
- expert_count = cutlass.Int32(0)
- for candidate_expert in cutlass.range_constexpr(self.num_local_experts):
- count = cutlass.min(
- cutlass.max(expert_counts[candidate_expert], cutlass.Int32(0)),
- expert_capacity,
- )
- num_expert_tiles = (count + tile_size - 1) // tile_size
- is_target = (tile_idx >= num_valid_tiles) and (
- tile_idx < num_valid_tiles + num_expert_tiles
- )
- if is_target:
- expert_idx = cutlass.Int32(candidate_expert)
- expert_tile_start = num_valid_tiles
- expert_count = count
- num_valid_tiles += num_expert_tiles
-
- row_in_expert = (tile_idx - expert_tile_start) * tile_size
- rows_in_tile = cutlass.min(
- tile_size,
- cutlass.max(expert_count - row_in_expert, cutlass.Int32(0)),
- )
- mn_limit = tile_idx * tile_size + rows_in_tile
- expanded_row_start = expert_idx * expert_capacity + row_in_expert
- return num_valid_tiles, expert_idx, mn_limit, expanded_row_start
-
- @cute.jit
- def __call__(
- self,
- a: cute.Tensor,
- b: Union[cute.Tensor, Tuple[cute.Tensor, ...]],
- out: cute.Tensor,
- sfa: cute.Tensor,
- sfb: Union[cute.Tensor, Tuple[cute.Tensor, ...]],
- tile_idx_to_expert_idx: cute.Tensor,
- num_non_exiting_tiles: cute.Tensor,
- tile_idx_to_mn_limit: cute.Tensor,
- alpha: Union[cute.Tensor, Tuple[cute.Tensor, ...]],
- expert_capacity: cutlass.Int32,
- max_active_clusters: cutlass.Constexpr,
- stream: cuda.CUstream,
- permuted_idx_to_expanded_idx: cute.Tensor,
- token_final_scales: cute.Tensor,
- epilogue_op: cutlass.Constexpr = lambda x: x,
- ):
- """Execute the GEMM operation in steps:
- - Setup static attributes before smem/grid/tma computation
- - Setup TMA load/store atoms and tensors
- - Compute grid size with regard to hardware constraints
- - Define shared storage for kernel
- - Launch the kernel synchronously
-
- :param a: Input tensor A
- :type a: cute.Tensor
- :param b: Input tensor B
- :type b: cute.Tensor
- :param out: Finalized output tensor (shape [seq_len, n])
- :type out: cute.Tensor
- :param sfa: Scale factor tensor A
- :type sfa: cute.Tensor
- :param sfb: Scale factor tensor B
- :type sfb: cute.Tensor
- :param tile_idx_to_expert_idx: Mapping from tile index to expert ID,
- shape (permuted_m/cta_tile_m,) where cta_tile_m is the CTA tile M size
- :type tile_idx_to_expert_idx: cute.Tensor
- :param num_non_exiting_tiles: Number of valid tiles (valid_m/cta_tile_m), shape (1,)
- :type num_non_exiting_tiles: cute.Tensor
- :param alpha: Alpha tensor for each group
- :type alpha: Union[cute.Tensor, Tuple[cute.Tensor, ...]]
- :param max_active_clusters: Maximum number of active clusters
- :type max_active_clusters: cutlass.Constexpr
- :param stream: CUDA stream for asynchronous execution
- :type stream: cuda.CUstream
- :param permuted_idx_to_expanded_idx: Mapping from permuted index to expanded index, shape (permuted_m,)
- :type permuted_idx_to_expanded_idx: cute.Tensor
- :param token_final_scales: Token-wise scaling factors, shape (m, topK)
- :type token_final_scales: cute.Tensor
- :param epilogue_op: Optional elementwise lambda function to apply to the output tensor
- :type epilogue_op: cutlass.Constexpr
- :raises TypeError: If input data types are incompatible with the MMA instruction.
- """
- # Setup static attributes before smem/grid/tma computation
- self.a_dtype: Type[cutlass.Numeric] = a.element_type
- # Handle tuple of B tensors
- b_tuple = b if isinstance(b, tuple) else (b,)
- sfb_tuple = sfb if isinstance(sfb, tuple) else (sfb,)
- alpha_tuple = alpha if isinstance(alpha, tuple) else (alpha,)
- self.b_dtype: Type[cutlass.Numeric] = b_tuple[0].element_type
- self.out_dtype: Type[cutlass.Numeric] = out.element_type
- self.sf_dtype: Type[cutlass.Numeric] = sfa.element_type
- self.final_scale_dtype: Type[cutlass.Numeric] = token_final_scales.element_type
- self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode()
- self.b_major_mode = utils.LayoutEnum.from_tensor(b_tuple[0]).mma_major_mode()
- self.gemm_output_layout = utils.LayoutEnum.ROW_MAJOR
-
- self.topK = token_final_scales.shape[1]
- # Check if input data types are compatible with MMA instruction
- if cutlass.const_expr(self.a_dtype != self.b_dtype):
- raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}")
-
- # Setup attributes that dependent on gemm inputs
- self._setup_attributes()
- # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout
- # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL)
- sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(a.shape, self.sf_vec_size)
- sfa = cute.make_tensor(sfa.iterator, sfa_layout)
-
- # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL)
- sfb_layout_0 = blockscaled_utils.tile_atom_to_shape_SF(b_tuple[0].shape, self.sf_vec_size)
- sfb_tensor_0 = cute.make_tensor(sfb_tuple[0].iterator, sfb_layout_0)
- sfb_tensors = [sfb_tensor_0]
- sfb_tuple = tuple(sfb_tensors)
- # Backward compat alias
- sfb = sfb_tuple[0]
-
- tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
- self.a_dtype,
- self.a_major_mode,
- self.b_major_mode,
- self.sf_dtype,
- self.sf_vec_size,
- self.cta_group,
- self.mma_inst_shape_mn,
- )
-
- tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
- self.a_dtype,
- self.a_major_mode,
- self.b_major_mode,
- self.sf_dtype,
- self.sf_vec_size,
- cute.nvgpu.tcgen05.CtaGroup.ONE,
- self.mma_inst_shape_mn_sfb,
- )
- atom_thr_size = cute.size(tiled_mma.thr_id.shape)
-
- # Setup TMA load for A
- a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id)
- a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
- tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
- a_op,
- a,
- a_smem_layout,
- self.mma_tiler,
- tiled_mma,
- self.cluster_layout_vmnk.shape,
- )
-
- # Setup TMA load for B
- b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id)
- b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
-
- # Setup TMA load for SFA
- sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id)
- sfa_smem_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0))
- tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A(
- sfa_op,
- sfa,
- sfa_smem_layout,
- self.mma_tiler,
- tiled_mma,
- self.cluster_layout_vmnk.shape,
- internal_type=cutlass.Int16,
- )
-
- # Setup TMA load for SFB
- sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, tiled_mma.thr_id)
- sfb_smem_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0))
-
- # Helper to create TMA for one B tensor
- def _make_tma_b(b_tensor, sfb_tensor):
- atom_b, tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
- b_op,
- b_tensor,
- b_smem_layout,
- self.mma_tiler,
- tiled_mma,
- self.cluster_layout_vmnk.shape,
- )
- atom_sfb, tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B(
- sfb_op,
- sfb_tensor,
- sfb_smem_layout,
- self.mma_tiler_sfb,
- tiled_mma_sfb,
- self.cluster_layout_sfb_vmnk.shape,
- internal_type=cutlass.Int16,
- )
- if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192):
- x = tensor_sfb.stride[0][1]
- y = cute.ceil_div(tensor_sfb.shape[0][1], 4)
-
- new_shape = (
- (tensor_sfb.shape[0][0], ((2, 2), y)),
- tensor_sfb.shape[1],
- tensor_sfb.shape[2],
- )
- # Use right multiplication for ScaledBasis (3 * x instead of x * 3)
- x_times_3 = 3 * x
- new_stride = (
- (tensor_sfb.stride[0][0], ((x, x), x_times_3)),
- tensor_sfb.stride[1],
- tensor_sfb.stride[2],
- )
- tensor_sfb_new_layout = cute.make_layout(new_shape, stride=new_stride)
- tensor_sfb = cute.make_tensor(tensor_sfb.iterator, tensor_sfb_new_layout)
- return atom_b, tensor_b, atom_sfb, tensor_sfb
-
- # Create TMA for the B tensor
- atom_b_0, tensor_b_0, atom_sfb_0, tensor_sfb_0 = _make_tma_b(b_tuple[0], sfb_tuple[0])
- tma_atoms_b = [atom_b_0]
- tma_tensors_b = [tensor_b_0]
- tma_atoms_sfb = [atom_sfb_0]
- tma_tensors_sfb = [tensor_sfb_0]
- tma_atoms_b = tuple(tma_atoms_b)
- tma_tensors_b = tuple(tma_tensors_b)
- tma_atoms_sfb = tuple(tma_atoms_sfb)
- tma_tensors_sfb = tuple(tma_tensors_sfb)
-
- a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout)
- b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout)
- sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout)
- sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout)
- self.num_tma_load_bytes = (
- a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size
- ) * atom_thr_size
-
- self.tile_sched_params, grid = self._compute_grid(
- (a.shape[0], b_tuple[0].shape[0], a.shape[2]),
- self.cta_tile_shape_mnk,
- self.cluster_shape_mn,
- max_active_clusters,
- self.raster_along_m,
- )
-
- self.buffer_align_bytes = 1024
-
- #### finalized epi layout ####
- epi_tile_m = cute.size(self.epi_tile[0])
- epi_tile_n = cute.size(self.epi_tile[1])
- epi_tile_size = epi_tile_m * epi_tile_n
- num_epilogue_threads = 32 * len(self.epilog_warp_id)
- self.ttr_racc_size = epi_tile_size // num_epilogue_threads
- self.copy_size = self.cta_tile_shape_mnk[1] * (self.out_dtype.width // 8)
-
- if cutlass.const_expr(self.out_dtype == cutlass.BFloat16):
- # 8-element vectorization for BF16
- self.epi_layout = cute.make_layout(
- shape=(self.ttr_racc_size // 8, 4, 2), stride=(8, 2, 1)
- )
- self.epi_loop_size = self.ttr_racc_size // 8
- self.element_offset = 8
-
- elif cutlass.const_expr(self.out_dtype == cutlass.Float32):
- # 2-element vectorization for FP32
- self.epi_layout = cute.make_layout(shape=(self.ttr_racc_size // 2, 2), stride=(2, 1))
- self.epi_loop_size = self.ttr_racc_size // 2
- self.element_offset = 2
- else:
- # Scalar fallback
- self.epi_layout = cute.make_layout(shape=(self.ttr_racc_size,), stride=(1,))
- self.epi_loop_size = self.ttr_racc_size
- self.element_offset = 1
-
- # Define shared storage for kernel
- @cute.struct
- class SharedStorage:
- # (bidx, bidy, expert, valid, mn_limit, expanded_row_start)
- sInfo: cute.struct.Align[
- cute.struct.MemRange[
- cutlass.Int32, self.num_tile_info_fields * self.num_tile_stage
- ],
- # 1 byte alignment
- 1,
- ]
- ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]
- acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]
- tile_info_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_tile_stage * 2]
- meta_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_meta_stage * 2]
- tmem_dealloc_mbar_ptr: cutlass.Int64
- tmem_holding_buf: cutlass.Int32
- # (MMA, MMA_M, MMA_K, STAGE)
- sA: cute.struct.Align[
- cute.struct.MemRange[self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)],
- self.buffer_align_bytes,
- ]
- # (MMA, MMA_N, MMA_K, STAGE)
- sB: cute.struct.Align[
- cute.struct.MemRange[self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)],
- self.buffer_align_bytes,
- ]
- # (granularity_m, repeat_m), (granularity_k, repeat_k), num_scale_stage)
- sSFA: cute.struct.Align[
- cute.struct.MemRange[self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged)],
- self.buffer_align_bytes,
- ]
- # (granularity_n, repeat_n), (granularity_k, repeat_k), num_scale_stage)
- sSFB: cute.struct.Align[
- cute.struct.MemRange[self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)],
- self.buffer_align_bytes,
- ]
-
- sC: cute.struct.Align[
- cute.struct.MemRange[self.out_dtype, cute.cosize(self.c_smem_layout_staged)],
- self.buffer_align_bytes,
- ]
- meta_token_idx: cute.struct.Align[
- cute.struct.MemRange[
- cutlass.Int32, self.cta_tile_shape_mnk[0] * self.num_meta_stage
- ],
- 1,
- ]
- meta_scale: cute.struct.Align[
- cute.struct.MemRange[
- self.final_scale_dtype, self.cta_tile_shape_mnk[0] * self.num_meta_stage
- ],
- 1,
- ]
-
- self.shared_storage = SharedStorage
-
- # Launch the kernel synchronously
- self.kernel(
- tiled_mma,
- tiled_mma_sfb,
- tma_atom_a,
- tma_tensor_a,
- tma_atoms_b,
- tma_tensors_b,
- tma_atom_sfa,
- tma_tensor_sfa,
- tma_atoms_sfb,
- tma_tensors_sfb,
- out,
- tile_idx_to_expert_idx,
- num_non_exiting_tiles,
- tile_idx_to_mn_limit,
- alpha_tuple,
- permuted_idx_to_expanded_idx,
- token_final_scales,
- expert_capacity,
- self.cluster_layout_vmnk,
- self.cluster_layout_sfb_vmnk,
- self.a_smem_layout_staged,
- self.b_smem_layout_staged,
- self.sfa_smem_layout_staged,
- self.sfb_smem_layout_staged,
- self.c_smem_layout_staged,
- self.epi_tile,
- self.epi_layout,
- self.topK,
- self.tile_sched_params,
- epilogue_op,
- ).launch(
- grid=grid,
- block=[self.threads_per_cta, 1, 1],
- cluster=(*self.cluster_shape_mn, 1),
- smem=self.shared_storage.size_in_bytes(),
- stream=stream,
- min_blocks_per_mp=1,
- use_pdl=TRTLLM_ENABLE_PDL,
- )
- return
-
- def mainloop_s2t_copy_and_partition(
- self,
- sSF: cute.Tensor,
- tSF: cute.Tensor,
- ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
- """
- Make tiledCopy for smem to tmem load for scale factor tensor, then use it to partition smem memory (source) and
- tensor memory (destination).
-
- :param sSF: The scale factor tensor in smem
- :type sSF: cute.Tensor
- :param tSF: The scale factor tensor in tmem
- :type tSF: cute.Tensor
-
- :return: A tuple containing (tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t) where:
- - tiled_copy_s2t: The tiled copy operation for smem to tmem load for scale factor tensor(s2t)
- - tCsSF_compact_s2t: The partitioned scale factor tensor in smem
- - tSF_compact_s2t: The partitioned scale factor tensor in tmem
- :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]
- """
- # (MMA, MMA_MN, MMA_K, STAGE)
- tCsSF_compact = cute.filter_zeros(sSF)
- # (MMA, MMA_MN, MMA_K)
- tCtSF_compact = cute.filter_zeros(tSF)
-
- # Make S2T CopyAtom and tiledCopy
- copy_atom_s2t = cute.make_copy_atom(
- tcgen05.Cp4x32x128bOp(self.cta_group),
- self.sf_dtype,
- )
- tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact)
- thr_copy_s2t = tiled_copy_s2t.get_slice(0)
-
- # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
- tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact)
- # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
- tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy_s2t, tCsSF_compact_s2t_)
- # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
- tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact)
-
- return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t
-
- @cute.kernel
- def kernel(
- self,
- tiled_mma: cute.TiledMma,
- tiled_mma_sfb: cute.TiledMma,
- tma_atom_a: cute.CopyAtom,
- mA_mkl: cute.Tensor,
- tma_atoms_b: Tuple[cute.CopyAtom, ...],
- mB_nkl_tuple: Tuple[cute.Tensor, ...],
- tma_atom_sfa: cute.CopyAtom,
- mSFA_mkl: cute.Tensor,
- tma_atoms_sfb: Tuple[cute.CopyAtom, ...],
- mSFB_nkl_tuple: Tuple[cute.Tensor, ...],
- out: cute.Tensor,
- tile_idx_to_expert_idx: cute.Tensor,
- num_non_exiting_tiles: cute.Tensor,
- tile_idx_to_mn_limit: cute.Tensor,
- alpha_tuple: Tuple[cute.Tensor, ...],
- permuted_idx_to_expanded_idx: cute.Tensor,
- token_final_scales: cute.Tensor,
- expert_capacity: cutlass.Int32,
- cluster_layout_vmnk: cute.Layout,
- cluster_layout_sfb_vmnk: cute.Layout,
- a_smem_layout_staged: cute.ComposedLayout,
- b_smem_layout_staged: cute.ComposedLayout,
- sfa_smem_layout_staged: cute.Layout,
- sfb_smem_layout_staged: cute.Layout,
- c_smem_layout_staged: cute.Layout,
- epi_tile: cute.Tile,
- epi_layout: cute.Layout,
- topK: cutlass.Int32,
- tile_sched_params: utils.PersistentTileSchedulerParams,
- epilogue_op: cutlass.Constexpr,
- ):
- """
- GPU device kernel performing the Persistent batched GEMM computation.
- """
- warp_idx = cute.arch.warp_idx()
- warp_idx = cute.arch.make_warp_uniform(warp_idx)
-
- #
- # Prefetch tma desc
- #
- if warp_idx == self.tma_warp_id:
- cpasync.prefetch_descriptor(tma_atom_a)
- cpasync.prefetch_descriptor(tma_atom_sfa)
- cpasync.prefetch_descriptor(tma_atoms_b[0])
- cpasync.prefetch_descriptor(tma_atoms_sfb[0])
-
- use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2
-
- #
- # Setup cta/thread coordinates
- #
- # Coords inside cluster
- bidx, bidy, bidz = cute.arch.block_idx()
- mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape)
- is_leader_cta = mma_tile_coord_v == 0
- cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
- block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
-
- block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord(
- cta_rank_in_cluster
- )
-
- # Coord inside cta
- tidx, _, _ = cute.arch.thread_idx()
-
- #
- # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier
- #
- smem = utils.SmemAllocator()
- storage = smem.allocate(self.shared_storage)
-
- # Initialize mainloop ab_pipeline (barrier) and states
- ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
- num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
- ab_pipeline_consumer_group = pipeline.CooperativeGroup(
- pipeline.Agent.Thread, num_tma_producer
- )
- ab_pipeline = pipeline.PipelineTmaUmma.create(
- barrier_storage=storage.ab_mbar_ptr.data_ptr(),
- num_stages=self.num_ab_stage,
- producer_group=ab_pipeline_producer_group,
- consumer_group=ab_pipeline_consumer_group,
- tx_count=self.num_tma_load_bytes,
- cta_layout_vmnk=cluster_layout_vmnk,
- )
-
- # Initialize acc_pipeline (barrier) and states
- acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
- num_acc_consumer_threads = (
- len(self.epilog_warp_id) * self.threads_per_warp * (2 if use_2cta_instrs else 1)
- )
- acc_pipeline_consumer_group = pipeline.CooperativeGroup(
- pipeline.Agent.Thread, num_acc_consumer_threads
- )
- acc_pipeline = pipeline.PipelineUmmaAsync.create(
- barrier_storage=storage.acc_mbar_ptr.data_ptr(),
- num_stages=self.num_acc_stage,
- producer_group=acc_pipeline_producer_group,
- consumer_group=acc_pipeline_consumer_group,
- cta_layout_vmnk=cluster_layout_vmnk,
- )
-
- # Initialize tile info pipeline (barrier) and states
- tile_info_pipeline_producer_group = pipeline.CooperativeGroup(
- pipeline.Agent.Thread,
- self.threads_per_warp * 1,
- )
- tile_info_pipeline_consumer_group = pipeline.CooperativeGroup(
- pipeline.Agent.Thread,
- self.threads_wo_sched,
- )
- tile_info_pipeline = pipeline.PipelineAsync.create(
- barrier_storage=storage.tile_info_mbar_ptr.data_ptr(),
- num_stages=self.num_tile_stage,
- producer_group=tile_info_pipeline_producer_group,
- consumer_group=tile_info_pipeline_consumer_group,
- )
-
- # Initialize metadata pipeline (meta loader warp -> epilogue warps)
- meta_pipeline_producer_group = pipeline.CooperativeGroup(
- pipeline.Agent.Thread,
- self.threads_per_warp * 1,
- )
- meta_pipeline_consumer_group = pipeline.CooperativeGroup(
- pipeline.Agent.Thread,
- self.threads_per_warp * len(self.epilog_warp_id),
- )
- meta_pipeline = pipeline.PipelineAsync.create(
- barrier_storage=storage.meta_mbar_ptr.data_ptr(),
- num_stages=self.num_meta_stage,
- producer_group=meta_pipeline_producer_group,
- consumer_group=meta_pipeline_consumer_group,
- )
-
- # Tensor memory dealloc barrier init
- tmem = utils.TmemAllocator(
- storage.tmem_holding_buf,
- barrier_for_retrieve=self.tmem_alloc_barrier,
- allocator_warp_id=self.epilog_warp_id[0],
- is_two_cta=use_2cta_instrs,
- two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
- )
-
- # Cluster arrive after barrier init
- if cute.size(self.cluster_shape_mn) > 1:
- cute.arch.cluster_arrive_relaxed()
-
- #
- # Setup smem tensor A/B/C/Scale/ExpandedIdx
- #
- # (MMA, MMA_M, MMA_K, STAGE)
- sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner)
- # (MMA, MMA_N, MMA_K, STAGE)
- sB = storage.sB.get_tensor(b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner)
- # (granularity_m, repeat_m), (granularity_k, repeat_k), num_scale_stage)
- sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged)
- # (granularity_n, repeat_n), (granularity_k, repeat_k), num_scale_stage)
- sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged)
-
- sC = storage.sC.get_tensor(c_smem_layout_staged)
-
- # (bidx, bidy, expert, valid, mn_limit, expanded_row_start)
- info_layout = cute.make_layout(
- (self.num_tile_info_fields, self.num_tile_stage),
- stride=(1, self.num_tile_info_fields),
- )
- sInfo = storage.sInfo.get_tensor(info_layout)
-
- # Per-row finalize metadata staged by the meta loader warp: (row, stage)
- meta_layout = cute.make_layout(
- (self.cta_tile_shape_mnk[0], self.num_meta_stage),
- stride=(1, self.cta_tile_shape_mnk[0]),
- )
- sMetaTokenIdx = storage.meta_token_idx.get_tensor(meta_layout)
- sMetaScale = storage.meta_scale.get_tensor(meta_layout)
-
- #
- # Compute multicast mask for A/B buffer full
- #
- a_full_mcast_mask = None
- b_full_mcast_mask = None
- sfa_full_mcast_mask = None
- sfb_full_mcast_mask = None
- if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs):
- a_full_mcast_mask = cpasync.create_tma_multicast_mask(
- cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
- )
- b_full_mcast_mask = cpasync.create_tma_multicast_mask(
- cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1
- )
- sfa_full_mcast_mask = cpasync.create_tma_multicast_mask(
- cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
- )
- sfb_full_mcast_mask = cpasync.create_tma_multicast_mask(
- cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_vmnk, mcast_mode=1
- )
-
- #
- # Local_tile partition global tensors
- #
- # (bM, bK, loopM, loopK, loopL)
- gA_mkl = cute.local_tile(
- mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
- )
- # (bN, bK, loopN, loopK, loopL)
- gB_nkl_0 = cute.local_tile(
- mB_nkl_tuple[0],
- cute.slice_(self.mma_tiler, (0, None, None)),
- (None, None, None),
- )
-
- # (bM, bK, RestM, RestK, RestL)
- gSFA_mkl = cute.local_tile(
- mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
- )
-
- # (bN, bK, RestN, RestK, RestL)
- gSFB_nkl_0 = cute.local_tile(
- mSFB_nkl_tuple[0],
- cute.slice_(self.mma_tiler_sfb, (0, None, None)),
- (None, None, None),
- )
-
- k_tile_cnt = cutlass.Int32(cute.size(gA_mkl, mode=[3]))
-
- #
- # Partition global tensor for TiledMMA_A/B
- #
- thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
- thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v)
- # (MMA, MMA_M, MMA_K, loopM, loopK, loopL)
- tCgA = thr_mma.partition_A(gA_mkl)
- # (MMA, MMA_N, MMA_K, loopN, loopK, loopL)
- tCgB_0 = thr_mma.partition_B(gB_nkl_0)
- # (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
- tCgSFA = thr_mma.partition_A(gSFA_mkl)
- # (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
- tCgSFB_0 = thr_mma_sfb.partition_B(gSFB_nkl_0)
-
- #
- # Partition global/shared tensor for TMA load A/B
- #
- # TMA load A partition_S/D
- a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape)
- # ((atom_v, rest_v), STAGE)
- # ((atom_v, rest_v), loopM, loopK, loopL)
- tAsA, tAgA = cpasync.tma_partition(
- tma_atom_a,
- block_in_cluster_coord_vmnk[2],
- a_cta_layout,
- cute.group_modes(sA, 0, 3),
- cute.group_modes(tCgA, 0, 3),
- )
- # TMA load B partition_S/D
- b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape)
- # ((atom_v, rest_v), STAGE)
- # ((atom_v, rest_v), loopM, loopK, loopL)
- tBsB_0, tBgB_0 = cpasync.tma_partition(
- tma_atoms_b[0],
- block_in_cluster_coord_vmnk[1],
- b_cta_layout,
- cute.group_modes(sB, 0, 3),
- cute.group_modes(tCgB_0, 0, 3),
- )
-
- # TMA load SFA partition_S/D
- sfa_cta_layout = a_cta_layout
- # ((atom_v, rest_v), STAGE)
- # ((atom_v, rest_v), RestM, RestK, RestL)
-
- tAsSFA, tAgSFA = cute.nvgpu.cpasync.tma_partition(
- tma_atom_sfa,
- block_in_cluster_coord_vmnk[2],
- sfa_cta_layout,
- cute.group_modes(sSFA, 0, 3),
- cute.group_modes(tCgSFA, 0, 3),
- )
-
- tAsSFA = cute.filter_zeros(tAsSFA)
- tAgSFA = cute.filter_zeros(tAgSFA)
-
- # TMA load SFB partition_S/D
- sfb_cta_layout = cute.make_layout(
- cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape
- )
- # ((atom_v, rest_v), STAGE)
- # ((atom_v, rest_v), RestN, RestK, RestL)
- tBsSFB_0, tBgSFB_0 = cute.nvgpu.cpasync.tma_partition(
- tma_atoms_sfb[0],
- block_in_cluster_coord_sfb_vmnk[1],
- sfb_cta_layout,
- cute.group_modes(sSFB, 0, 3),
- cute.group_modes(tCgSFB_0, 0, 3),
- )
- tBsSFB_0 = cute.filter_zeros(tBsSFB_0)
- tBgSFB_0 = cute.filter_zeros(tBgSFB_0)
-
- #
- # Partition shared/tensor memory tensor for TiledMMA_A/B/C
- #
- # (MMA, MMA_M, MMA_K, STAGE)
- tCrA = tiled_mma.make_fragment_A(sA)
- # (MMA, MMA_N, MMA_K, STAGE)
- tCrB = tiled_mma.make_fragment_B(sB)
- # (MMA, MMA_M, MMA_N)
- acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
-
- if cutlass.const_expr(self.overlapping_accum):
- num_acc_stage_overlapped = 2
- tCtAcc_fake = tiled_mma.make_fragment_C(
- cute.append(acc_shape, num_acc_stage_overlapped)
- )
- # (MMA, MMA_M, MMA_N, STAGE)
- tCtAcc_fake = cute.make_tensor(
- tCtAcc_fake.iterator,
- cute.make_layout(
- tCtAcc_fake.shape,
- stride=(
- tCtAcc_fake.stride[0],
- tCtAcc_fake.stride[1],
- tCtAcc_fake.stride[2],
- (256 - self.num_sf_tmem_cols) * tCtAcc_fake.stride[0][1],
- ),
- ),
- )
- else:
- # (MMA, MMA_M, MMA_N, STAGE)
- tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage))
-
- gC_mnl = cute.local_tile(
- out, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
- )
-
- # (MMA, MMA_M, MMA_N, loopM, loopN, loopL)
- tCgC = thr_mma.partition_C(gC_mnl)
-
- #
- # Cluster wait before tensor memory alloc
- #
- if cute.size(self.cluster_shape_mn) > 1:
- cute.arch.cluster_wait()
- else:
- self.cta_sync_barrier.arrive_and_wait()
-
- griddepcontrol_wait()
-
- #
- # Specialized Schedule warp
- #
- if warp_idx == self.sched_warp_id:
- #
- # Persistent tile scheduling loop
- #
- tile_sched = utils.StaticPersistentTileScheduler.create(
- tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
- )
- # First tile
- work_tile = tile_sched.initial_work_tile_info()
-
- tile_info_producer_state = pipeline.make_pipeline_state(
- pipeline.PipelineUserType.Producer, self.num_tile_stage
- )
-
- if cutlass.const_expr(self.use_expert_counts):
- num_valid_tiles, _, _, _ = self._expert_count_tile_info(
- tile_idx_to_expert_idx, cutlass.Int32(0), expert_capacity
- )
- else:
- num_valid_tiles = num_non_exiting_tiles[0]
-
- if cutlass.const_expr(self.raster_along_m):
- while work_tile.is_valid_tile:
- cur_tile_coord = work_tile.tile_idx
- mma_tile_coord_m = cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape)
- tile_idx = mma_tile_coord_m
- if tile_idx < num_valid_tiles:
- tile_info_pipeline.producer_acquire(tile_info_producer_state)
- expanded_row_start = cutlass.Int32(-1)
- if cutlass.const_expr(self.use_expert_counts):
- _, expert_idx, mn_limit, expanded_mma_row_start = (
- self._expert_count_tile_info(
- tile_idx_to_expert_idx, mma_tile_coord_m, expert_capacity
- )
- )
- expanded_row_start = (
- expanded_mma_row_start
- + mma_tile_coord_v * self.cta_tile_shape_mnk[0]
- )
- else:
- expert_idx = tile_idx_to_expert_idx[mma_tile_coord_m]
- mn_limit = tile_idx_to_mn_limit[tile_idx]
- with cute.arch.elect_one():
- sInfo[(0, tile_info_producer_state.index)] = cur_tile_coord[0]
- sInfo[(1, tile_info_producer_state.index)] = cur_tile_coord[1]
- sInfo[(2, tile_info_producer_state.index)] = expert_idx
- sInfo[(3, tile_info_producer_state.index)] = cutlass.Int32(
- work_tile.is_valid_tile
- )
- sInfo[(4, tile_info_producer_state.index)] = mn_limit
- if cutlass.const_expr(self.use_expert_counts):
- sInfo[(5, tile_info_producer_state.index)] = expanded_row_start
- # fence view async shared
- cute.arch.fence_proxy(
- "async.shared",
- space="cta",
- )
-
- self.sched_sync_barrier.arrive_and_wait()
- tile_info_pipeline.producer_commit(tile_info_producer_state)
- tile_info_producer_state.advance()
-
- tile_sched.advance_to_next_work()
- work_tile = tile_sched.get_current_work()
- else:
- is_continue = cutlass.Boolean(1)
- while work_tile.is_valid_tile and is_continue:
- cur_tile_coord = work_tile.tile_idx
- mma_tile_coord_m = cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape)
- tile_idx = mma_tile_coord_m
- if tile_idx < num_valid_tiles:
- tile_info_pipeline.producer_acquire(tile_info_producer_state)
- expanded_row_start = cutlass.Int32(-1)
- if cutlass.const_expr(self.use_expert_counts):
- _, expert_idx, mn_limit, expanded_mma_row_start = (
- self._expert_count_tile_info(
- tile_idx_to_expert_idx, mma_tile_coord_m, expert_capacity
- )
- )
- expanded_row_start = (
- expanded_mma_row_start
- + mma_tile_coord_v * self.cta_tile_shape_mnk[0]
- )
- else:
- expert_idx = tile_idx_to_expert_idx[mma_tile_coord_m]
- mn_limit = tile_idx_to_mn_limit[tile_idx]
- with cute.arch.elect_one():
- sInfo[(0, tile_info_producer_state.index)] = cur_tile_coord[0]
- sInfo[(1, tile_info_producer_state.index)] = cur_tile_coord[1]
- sInfo[(2, tile_info_producer_state.index)] = expert_idx
- sInfo[(3, tile_info_producer_state.index)] = cutlass.Int32(
- work_tile.is_valid_tile
- )
- sInfo[(4, tile_info_producer_state.index)] = mn_limit
- if cutlass.const_expr(self.use_expert_counts):
- sInfo[(5, tile_info_producer_state.index)] = expanded_row_start
- # fence view async shared
- cute.arch.fence_proxy(
- "async.shared",
- space="cta",
- )
-
- self.sched_sync_barrier.arrive_and_wait()
- tile_info_pipeline.producer_commit(tile_info_producer_state)
- tile_info_producer_state.advance()
-
- else:
- is_continue = cutlass.Boolean(0)
-
- tile_sched.advance_to_next_work()
- work_tile = tile_sched.get_current_work()
-
- tile_info_pipeline.producer_acquire(tile_info_producer_state)
- with cute.arch.elect_one():
- sInfo[(0, tile_info_producer_state.index)] = work_tile.tile_idx[0]
- sInfo[(1, tile_info_producer_state.index)] = work_tile.tile_idx[1]
- sInfo[(2, tile_info_producer_state.index)] = -1
- sInfo[(3, tile_info_producer_state.index)] = cutlass.Int32(0)
- sInfo[(4, tile_info_producer_state.index)] = cutlass.Int32(0)
- if cutlass.const_expr(self.use_expert_counts):
- sInfo[(5, tile_info_producer_state.index)] = cutlass.Int32(-1)
- cute.arch.fence_proxy(
- "async.shared",
- space="cta",
- )
- self.sched_sync_barrier.arrive_and_wait()
- tile_info_pipeline.producer_commit(tile_info_producer_state)
- tile_info_producer_state.advance()
- tile_info_pipeline.producer_tail(tile_info_producer_state)
-
- #
- # Specialized TMA load warp
- #
- if warp_idx == self.tma_warp_id:
- ab_producer_state = pipeline.make_pipeline_state(
- pipeline.PipelineUserType.Producer, self.num_ab_stage
- )
-
- tile_info_consumer_state = pipeline.make_pipeline_state(
- pipeline.PipelineUserType.Consumer, self.num_tile_stage
- )
-
- # Get the first tile info from pipeline (scheduler has filtered out tiles >= num_non_exiting_tiles)
- tile_info = cute.make_rmem_tensor((self.num_tile_info_fields,), cutlass.Int32)
- tile_info_pipeline.consumer_wait(tile_info_consumer_state)
- tile_info[0] = sInfo[(0, tile_info_consumer_state.index)]
- tile_info[1] = sInfo[(1, tile_info_consumer_state.index)]
- tile_info[2] = sInfo[(2, tile_info_consumer_state.index)]
- tile_info[3] = sInfo[(3, tile_info_consumer_state.index)]
- is_valid_tile = tile_info[3] == 1
- cute.arch.fence_proxy(
- "async.shared",
- space="cta",
- )
- tile_info_pipeline.consumer_release(tile_info_consumer_state)
- tile_info_consumer_state.advance()
-
- while is_valid_tile:
- mma_tile_coord_mnl = (
- tile_info[0] // cute.size(tiled_mma.thr_id.shape),
- tile_info[1],
- tile_info[2],
- )
- #
- # Slice to per mma tile index
- #
- # ((atom_v, rest_v), loopK)
- tAgA_slice = tAgA[(None, mma_tile_coord_mnl[0], None, 0)]
-
- # ((atom_v, rest_v), RestK)
- tAgSFA_slice = tAgSFA[(None, mma_tile_coord_mnl[0], None, 0)]
-
- expert_idx = mma_tile_coord_mnl[2]
- slice_n = mma_tile_coord_mnl[1]
- if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
- slice_n = mma_tile_coord_mnl[1] // 2
-
- # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt
- ab_producer_state.reset_count()
- peek_ab_empty_status = cutlass.Boolean(1)
- if ab_producer_state.count < k_tile_cnt:
- peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state)
- #
- # Tma load loop
- #
- for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):
- tAgA_k = tAgA_slice[(None, ab_producer_state.count)]
- tAgSFA_k = tAgSFA_slice[(None, ab_producer_state.count)]
- tAsA_pipe = tAsA[(None, ab_producer_state.index)]
- tBsB_pipe = tBsB_0[(None, ab_producer_state.index)]
- tAsSFA_pipe = tAsSFA[(None, ab_producer_state.index)]
- tBsSFB_pipe = tBsSFB_0[(None, ab_producer_state.index)]
-
- tma_bar = ab_pipeline.producer_get_barrier(ab_producer_state)
-
- # Conditionally wait for AB buffer empty
- ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status)
-
- # TMA load A/B
- tBgB_slice = tBgB_0[(None, mma_tile_coord_mnl[1], None, expert_idx)]
- tBgSFB_slice = tBgSFB_0[(None, slice_n, None, expert_idx)]
- if cutlass.const_expr(self.l2_hints):
- cute.copy(
- tma_atom_a,
- tAgA_k,
- tAsA_pipe,
- tma_bar_ptr=tma_bar,
- mcast_mask=a_full_mcast_mask,
- cache_policy=cutlass.Int64(_EVICT_FIRST),
- )
- cute.copy(
- tma_atom_sfa,
- tAgSFA_k,
- tAsSFA_pipe,
- tma_bar_ptr=tma_bar,
- mcast_mask=sfa_full_mcast_mask,
- cache_policy=cutlass.Int64(_EVICT_FIRST),
- )
- cute.copy(
- tma_atoms_b[0],
- tBgB_slice[(None, ab_producer_state.count)],
- tBsB_pipe,
- tma_bar_ptr=tma_bar,
- mcast_mask=b_full_mcast_mask,
- cache_policy=cutlass.Int64(_EVICT_FIRST),
- )
- cute.copy(
- tma_atoms_sfb[0],
- tBgSFB_slice[(None, ab_producer_state.count)],
- tBsSFB_pipe,
- tma_bar_ptr=tma_bar,
- mcast_mask=sfb_full_mcast_mask,
- cache_policy=cutlass.Int64(_EVICT_FIRST),
- )
- else:
- cute.copy(
- tma_atom_a,
- tAgA_k,
- tAsA_pipe,
- tma_bar_ptr=tma_bar,
- mcast_mask=a_full_mcast_mask,
- )
- cute.copy(
- tma_atom_sfa,
- tAgSFA_k,
- tAsSFA_pipe,
- tma_bar_ptr=tma_bar,
- mcast_mask=sfa_full_mcast_mask,
- )
- cute.copy(
- tma_atoms_b[0],
- tBgB_slice[(None, ab_producer_state.count)],
- tBsB_pipe,
- tma_bar_ptr=tma_bar,
- mcast_mask=b_full_mcast_mask,
- )
- cute.copy(
- tma_atoms_sfb[0],
- tBgSFB_slice[(None, ab_producer_state.count)],
- tBsSFB_pipe,
- tma_bar_ptr=tma_bar,
- mcast_mask=sfb_full_mcast_mask,
- )
-
- # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1
- ab_producer_state.advance()
- peek_ab_empty_status = cutlass.Boolean(1)
- if ab_producer_state.count < k_tile_cnt:
- peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state)
-
- #
- # Advance to next tile
- #
- tile_info_pipeline.consumer_wait(tile_info_consumer_state)
- tile_info[0] = sInfo[(0, tile_info_consumer_state.index)]
- tile_info[1] = sInfo[(1, tile_info_consumer_state.index)]
- tile_info[2] = sInfo[(2, tile_info_consumer_state.index)]
- tile_info[3] = sInfo[(3, tile_info_consumer_state.index)]
- is_valid_tile = tile_info[3] == 1
- cute.arch.fence_proxy(
- "async.shared",
- space="cta",
- )
- tile_info_pipeline.consumer_release(tile_info_consumer_state)
- tile_info_consumer_state.advance()
- #
- # Wait A/B buffer empty
- #
- ab_pipeline.producer_tail(ab_producer_state)
-
- #
- # Specialized MMA warp
- #
- if warp_idx == self.mma_warp_id:
- #
- # Bar sync for retrieve tensor memory ptr from shared mem
- #
- tmem.wait_for_alloc()
-
- #
- # Retrieving tensor memory ptr and make accumulator tensor
- #
- acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
- # (MMA, MMA_M, MMA_N, STAGE)
- tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
-
- # Make SFA tmem tensor
- sfa_tmem_ptr = cute.recast_ptr(
- acc_tmem_ptr + self.num_accumulator_tmem_cols,
- dtype=self.sf_dtype,
- )
- # (MMA, MMA_M, MMA_K)
- tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
- tiled_mma,
- self.mma_tiler,
- self.sf_vec_size,
- cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)),
- )
- tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout)
-
- # Make SFB tmem tensor
- # Make SFB tmem tensor
- sfb_tmem_ptr = cute.recast_ptr(
- acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols,
- dtype=self.sf_dtype,
- )
- # (MMA, MMA_N, MMA_K)
- tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
- tiled_mma,
- self.mma_tiler,
- self.sf_vec_size,
- cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)),
- )
- tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout)
-
- # Partition for S2T copy of SFA/SFB
- #
- (
- tiled_copy_s2t_sfa,
- tCsSFA_compact_s2t,
- tCtSFA_compact_s2t,
- ) = self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA)
- (
- tiled_copy_s2t_sfb,
- tCsSFB_compact_s2t,
- tCtSFB_compact_s2t,
- ) = self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB)
-
- ab_consumer_state = pipeline.make_pipeline_state(
- pipeline.PipelineUserType.Consumer, self.num_ab_stage
- )
- acc_producer_state = pipeline.make_pipeline_state(
- pipeline.PipelineUserType.Producer, self.num_acc_stage
- )
-
- tile_info_consumer_state = pipeline.make_pipeline_state(
- pipeline.PipelineUserType.Consumer, self.num_tile_stage
- )
-
- # Get the first tile info from pipeline (scheduler has filtered out tiles >= num_non_exiting_tiles)
- tile_info = cute.make_rmem_tensor((5,), cutlass.Int32)
- tile_info_pipeline.consumer_wait(tile_info_consumer_state)
- tile_info[0] = sInfo[(0, tile_info_consumer_state.index)]
- tile_info[1] = sInfo[(1, tile_info_consumer_state.index)]
- tile_info[2] = sInfo[(2, tile_info_consumer_state.index)]
- tile_info[3] = sInfo[(3, tile_info_consumer_state.index)]
- is_valid_tile = tile_info[3] == 1
- cute.arch.fence_proxy(
- "async.shared",
- space="cta",
- )
- tile_info_pipeline.consumer_release(tile_info_consumer_state)
- tile_info_consumer_state.advance()
-
- while is_valid_tile:
- # Peek (try_wait) AB buffer full for k_tile = 0
- ab_consumer_state.reset_count()
- peek_ab_full_status = cutlass.Boolean(1)
- if ab_consumer_state.count < k_tile_cnt and is_leader_cta:
- peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state)
-
- mma_tile_coord_mnl = (
- tile_info[0] // cute.size(tiled_mma.thr_id.shape),
- tile_info[1],
- tile_info[2],
- )
-
- # Get accumulator stage index
- if cutlass.const_expr(self.overlapping_accum):
- acc_stage_index = acc_producer_state.phase ^ 1
- else:
- acc_stage_index = acc_producer_state.index
-
- tCtAcc = tCtAcc_base[(None, None, None, acc_stage_index)]
-
- tCtSFB_mma = tCtSFB
- if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192):
- # If this is an ODD tile, shift the TMEM start address for cta_tile_shape_n=192 case by two words
- # (ignores first 64 columns of SFB)
- offset = (
- cutlass.Int32(2) if mma_tile_coord_mnl[1] % 2 == 1 else cutlass.Int32(0)
- )
- shifted_ptr = cute.recast_ptr(
- acc_tmem_ptr
- + self.num_accumulator_tmem_cols
- + self.num_sfa_tmem_cols
- + offset,
- dtype=self.sf_dtype,
- )
- tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout)
- elif cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
- # Move in increments of 64 columns of SFB
- offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2)
- shifted_ptr = cute.recast_ptr(
- acc_tmem_ptr
- + self.num_accumulator_tmem_cols
- + self.num_sfa_tmem_cols
- + offset,
- dtype=self.sf_dtype,
- )
- tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout)
- #
- # Wait for accumulator buffer empty
- #
- if is_leader_cta:
- acc_pipeline.producer_acquire(acc_producer_state)
- #
- # Reset the ACCUMULATE field for each tile
- #
- tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
-
- #
- # Mma mainloop
- #
- for k_tile in cutlass.range(k_tile_cnt):
- if is_leader_cta:
- # Conditionally wait for AB buffer full
- ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status)
-
- # Copy SFA/SFB from smem to tmem
- s2t_stage_coord = (
- None,
- None,
- None,
- None,
- ab_consumer_state.index,
- )
- tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord]
- tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord]
- cute.copy(
- tiled_copy_s2t_sfa,
- tCsSFA_compact_s2t_staged,
- tCtSFA_compact_s2t,
- )
- cute.copy(
- tiled_copy_s2t_sfb,
- tCsSFB_compact_s2t_staged,
- tCtSFB_compact_s2t,
- )
-
- # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB
- num_kblocks = cute.size(tCrA, mode=[2])
-
- for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):
- kblock_coord = (
- None,
- None,
- kblock_idx,
- ab_consumer_state.index,
- )
-
- # Set SFA/SFB tensor to tiled_mma
- sf_kblock_coord = (None, None, kblock_idx)
- tiled_mma.set(
- tcgen05.Field.SFA,
- tCtSFA[sf_kblock_coord].iterator,
- )
- tiled_mma.set(
- tcgen05.Field.SFB,
- tCtSFB_mma[sf_kblock_coord].iterator,
- )
-
- cute.gemm(
- tiled_mma,
- tCtAcc,
- tCrA[kblock_coord],
- tCrB[kblock_coord],
- tCtAcc,
- )
-
- # Enable accumulate on tCtAcc after first kblock
- tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
-
- # Async arrive AB buffer empty
- ab_pipeline.consumer_release(ab_consumer_state)
-
- # Peek (try_wait) AB buffer full for k_tile = k_tile + 1
- ab_consumer_state.advance()
- peek_ab_full_status = cutlass.Boolean(1)
- if ab_consumer_state.count < k_tile_cnt:
- if is_leader_cta:
- peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state)
-
- #
- # Async arrive accumulator buffer full(each kblock)
- #
- if is_leader_cta:
- acc_pipeline.producer_commit(acc_producer_state)
-
- # Peek (try_wait) Acc buffer empty for k_tile = k_tile + 1
- acc_producer_state.advance()
- #
- # Advance to next tile
- #
- tile_info_pipeline.consumer_wait(tile_info_consumer_state)
- tile_info[0] = sInfo[(0, tile_info_consumer_state.index)]
- tile_info[1] = sInfo[(1, tile_info_consumer_state.index)]
- tile_info[2] = sInfo[(2, tile_info_consumer_state.index)]
- tile_info[3] = sInfo[(3, tile_info_consumer_state.index)]
- is_valid_tile = tile_info[3] == 1
- cute.arch.fence_proxy(
- "async.shared",
- space="cta",
- )
- tile_info_pipeline.consumer_release(tile_info_consumer_state)
- tile_info_consumer_state.advance()
- #
- # Wait for accumulator buffer empty
- #
- acc_pipeline.producer_tail(acc_producer_state)
-
- #
- # Specialized metadata loader warp
- #
- if warp_idx == self.meta_load_warp_id:
- meta_lane = tidx % self.threads_per_warp
-
- tile_info_consumer_state = pipeline.make_pipeline_state(
- pipeline.PipelineUserType.Consumer, self.num_tile_stage
- )
- meta_producer_state = pipeline.make_pipeline_state(
- pipeline.PipelineUserType.Producer, self.num_meta_stage
- )
- tile_info = cute.make_rmem_tensor((self.num_tile_info_fields,), cutlass.Int32)
- tile_info_pipeline.consumer_wait(tile_info_consumer_state)
- tile_info[0] = sInfo[(0, tile_info_consumer_state.index)]
- tile_info[2] = sInfo[(2, tile_info_consumer_state.index)]
- tile_info[3] = sInfo[(3, tile_info_consumer_state.index)]
- tile_info[4] = sInfo[(4, tile_info_consumer_state.index)]
- if cutlass.const_expr(self.use_expert_counts):
- tile_info[5] = sInfo[(5, tile_info_consumer_state.index)]
- is_valid_tile = tile_info[3] == 1
- cute.arch.fence_proxy("async.shared", space="cta")
- tile_info_pipeline.consumer_release(tile_info_consumer_state)
- tile_info_consumer_state.advance()
-
- while is_valid_tile:
- tile_m_start = tile_info[0] * self.cta_tile_shape_mnk[0]
- expert_idx = tile_info[2]
- alpha_val = alpha_tuple[0][expert_idx]
-
- meta_pipeline.producer_acquire(meta_producer_state)
- meta_stage = meta_producer_state.index
- # Strided row assignment keeps the permuted_idx loads and smem
- # stores coalesced (each fixed j: 32 lanes touch 32 contiguous
- # rows). Rows beyond mn_limit are padding: their expanded_idx is
- # not guaranteed to be in range, so gate the token_final_scales
- # gather to token 0 for them (branchless) to keep it in-bounds.
- # The epilogue ignores padding rows via its own is_valid_row, so
- # the value staged for them is irrelevant.
- for j in cutlass.range(
- self.cta_tile_shape_mnk[0] // self.threads_per_warp,
- unroll_full=True,
- ):
- r = meta_lane + j * self.threads_per_warp
- permuted_row = tile_m_start + r
- if cutlass.const_expr(self.use_expert_counts):
- expanded_idx = tile_info[5] + r
- else:
- expanded_idx = permuted_idx_to_expanded_idx[permuted_row]
- # max(., 0) keeps topk_idx in-bounds for padding rows (-1);
- # the is_valid gate keeps token_idx in-bounds (padding rows
- # may hold out-of-range garbage, not just -1).
- safe_idx = cutlass.max(expanded_idx, cutlass.Int32(0))
- token_idx = safe_idx // topK
- topk_idx = safe_idx % topK
- is_valid_row = cutlass.Int32(permuted_row < tile_info[4])
- gather_tok = token_idx * is_valid_row
- token_scale = token_final_scales[(gather_tok, topk_idx)]
- sMetaTokenIdx[(r, meta_stage)] = token_idx
- sMetaScale[(r, meta_stage)] = alpha_val * token_scale
- cute.arch.fence_proxy("async.shared", space="cta")
- meta_pipeline.producer_commit(meta_producer_state)
- meta_producer_state.advance()
-
- tile_info_pipeline.consumer_wait(tile_info_consumer_state)
- tile_info[0] = sInfo[(0, tile_info_consumer_state.index)]
- tile_info[2] = sInfo[(2, tile_info_consumer_state.index)]
- tile_info[3] = sInfo[(3, tile_info_consumer_state.index)]
- tile_info[4] = sInfo[(4, tile_info_consumer_state.index)]
- if cutlass.const_expr(self.use_expert_counts):
- tile_info[5] = sInfo[(5, tile_info_consumer_state.index)]
- is_valid_tile = tile_info[3] == 1
- cute.arch.fence_proxy("async.shared", space="cta")
- tile_info_pipeline.consumer_release(tile_info_consumer_state)
- tile_info_consumer_state.advance()
- meta_pipeline.producer_tail(meta_producer_state)
-
- #
- # Specialized epilogue warps
- #
- if warp_idx < self.mma_warp_id:
- #
- # Alloc tensor memory buffer
- #
- tmem.allocate(self.num_tmem_alloc_cols)
-
- #
- # Bar sync for retrieve tensor memory ptr from shared memory
- #
- tmem.wait_for_alloc()
-
- #
- # Retrieving tensor memory ptr and make accumulator tensor
- #
- tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
- # (MMA, MMA_M, MMA_N, STAGE)
- tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
-
- #
- # Partition for epilogue
- #
- epi_tidx = tidx % 128
- (
- tiled_copy_t2r,
- tTR_tAcc_base,
- tTR_rAcc,
- ) = self.epilog_tmem_copy_and_partition(
- epi_tidx, tCtAcc_base, tCgC, epi_tile, use_2cta_instrs
- )
-
- tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.out_dtype)
- tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition(
- epi_tidx, tTR_rC, sC, tiled_copy_t2r
- )
-
- acc_consumer_state = pipeline.make_pipeline_state(
- pipeline.PipelineUserType.Consumer, self.num_acc_stage
- )
-
- tile_info_consumer_state = pipeline.make_pipeline_state(
- pipeline.PipelineUserType.Consumer, self.num_tile_stage
- )
- meta_consumer_state = pipeline.make_pipeline_state(
- pipeline.PipelineUserType.Consumer, self.num_meta_stage
- )
-
- # Get the first tile info
- tile_info = cute.make_rmem_tensor((5,), cutlass.Int32)
-
- tile_info_pipeline.consumer_wait(tile_info_consumer_state)
- tile_info[0] = sInfo[(0, tile_info_consumer_state.index)]
- tile_info[1] = sInfo[(1, tile_info_consumer_state.index)]
- tile_info[3] = sInfo[(3, tile_info_consumer_state.index)]
- tile_info[4] = sInfo[(4, tile_info_consumer_state.index)]
- is_valid_tile = tile_info[3] == 1
- cute.arch.fence_proxy(
- "async.shared",
- space="cta",
- )
- tile_info_pipeline.consumer_release(tile_info_consumer_state)
- tile_info_consumer_state.advance()
-
- while is_valid_tile:
- tile_m_start = tile_info[0] * self.cta_tile_shape_mnk[0]
- permuted_row = tile_m_start + epi_tidx
- is_valid_row = permuted_row < tile_info[4]
- if cutlass.const_expr(self.row_owned_epilogue):
- cute.arch.cp_async_bulk_wait_group(0, read=True)
-
- # Read per-row finalize metadata prefetched by the meta loader
- # warp (token_idx for scatter, combined_scale = alpha * token_scale).
- meta_pipeline.consumer_wait(meta_consumer_state)
- token_idx = sMetaTokenIdx[(epi_tidx, meta_consumer_state.index)]
- meta_scale = sMetaScale[(epi_tidx, meta_consumer_state.index)]
-
- # Get accumulator stage index
- if cutlass.const_expr(self.overlapping_accum):
- acc_stage_index = acc_consumer_state.phase
- reverse_subtile = (
- cutlass.Boolean(True) if acc_stage_index == 0 else cutlass.Boolean(False)
- )
- else:
- acc_stage_index = acc_consumer_state.index
-
- # Set tensor memory buffer for current tile
- # (T2R, T2R_M, T2R_N, EPI_M, EPI_M)
- tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_stage_index)]
-
- #
- # Wait for accumulator buffer full
- #
- acc_pipeline.consumer_wait(acc_consumer_state)
-
- tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
- #
- # Process sub-tiles with vectorized scatter-add
- # Following TensorRT-LLM's direct G2R (global to register) approach
- #
- subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
-
- for subtile_idx in cutlass.range(subtile_cnt):
- real_subtile_idx = subtile_idx
- if cutlass.const_expr(self.overlapping_accum):
- if reverse_subtile:
- real_subtile_idx = subtile_cnt - 1 - subtile_idx
- #
- # Load accumulator from tensor memory buffer to register
- #
- tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)]
-
- cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc)
-
- #
- # Async arrive accumulator buffer empty earlier when overlapping_accum is enabled
- #
- if cutlass.const_expr(self.overlapping_accum):
- if subtile_idx == self.iter_acc_early_release_in_epilogue:
- # Fence for TMEM load
- cute.arch.fence_view_async_tmem_load()
- acc_pipeline.consumer_release(acc_consumer_state)
- acc_consumer_state.advance()
-
- # Get vectorized accumulator and apply the combined finalize scale
- acc_vec = tTR_rAcc.load()
- acc_vec_final = meta_scale * acc_vec
-
- tRS_rC.store(acc_vec_final.to(self.out_dtype))
- if is_valid_row:
- cute.copy(
- tiled_copy_r2s,
- tRS_rC,
- tRS_sC[(None, None, real_subtile_idx, None)],
- )
-
- # Make all R2S smem writes visible to the async bulk-reduce proxy.
- cute.arch.fence_proxy(
- "async.shared",
- space="cta",
- )
- #
- # Async arrive accumulator buffer empty
- #
- if cutlass.const_expr(not self.overlapping_accum):
- cute.arch.fence_view_async_tmem_load()
- acc_pipeline.consumer_release(acc_consumer_state)
- acc_consumer_state.advance()
-
- # Whole-row async bulk reduce (smem -> global scatter-add).
- if is_valid_row:
- coord_n = tile_info[1] * self.cta_tile_shape_mnk[1]
- scatter_out_offset = cute.domain_offset((token_idx, coord_n, 0), out)
- if cutlass.const_expr(self.out_dtype == cutlass.BFloat16 and self.l2_hints):
- _blk_reduce_bf16_evict_last(
- scatter_out_offset,
- sC[epi_tidx, None, 0],
- cutlass.Int32(self.copy_size),
- )
- elif cutlass.const_expr(self.out_dtype == cutlass.BFloat16):
- blk_reduce_bf16(
- scatter_out_offset,
- sC[epi_tidx, None, 0],
- cutlass.Int32(self.copy_size),
- )
- elif cutlass.const_expr(self.out_dtype == cutlass.Float32):
- blk_reduce_fp32(
- scatter_out_offset,
- sC[epi_tidx, None, 0],
- cutlass.Int32(self.copy_size),
- )
- elif cutlass.const_expr(self.out_dtype == cutlass.Float16):
- blk_reduce_fp16(
- scatter_out_offset,
- sC[epi_tidx, None, 0],
- cutlass.Int32(self.copy_size),
- )
-
- cute.arch.cp_async_bulk_commit_group()
- if cutlass.const_expr(not self.row_owned_epilogue):
- cute.arch.cp_async_bulk_wait_group(0, read=True)
- self.epilog_sync_barrier.arrive_and_wait()
-
- # Release the prefetched metadata slot for this tile.
- meta_pipeline.consumer_release(meta_consumer_state)
- meta_consumer_state.advance()
-
- # Advance to next tile
- #
- tile_info_pipeline.consumer_wait(tile_info_consumer_state)
- tile_info[0] = sInfo[(0, tile_info_consumer_state.index)]
- tile_info[1] = sInfo[(1, tile_info_consumer_state.index)]
- tile_info[3] = sInfo[(3, tile_info_consumer_state.index)]
- tile_info[4] = sInfo[(4, tile_info_consumer_state.index)]
- is_valid_tile = tile_info[3] == 1
- cute.arch.fence_proxy(
- "async.shared",
- space="cta",
- )
- tile_info_pipeline.consumer_release(tile_info_consumer_state)
- tile_info_consumer_state.advance()
- if cutlass.const_expr(self.row_owned_epilogue):
- cute.arch.cp_async_bulk_wait_group(0, read=True)
- #
- # Dealloc the tensor memory buffer
- #
- tmem.relinquish_alloc_permit()
- self.epilog_sync_barrier.arrive_and_wait()
- tmem.free(tmem_ptr)
-
- griddepcontrol_launch_dependents()
-
- def epilog_tmem_copy_and_partition(
- self,
- tidx: cutlass.Int32,
- tAcc: cute.Tensor,
- gC_mnl: cute.Tensor,
- epi_tile: cute.Tile,
- use_2cta_instrs: Union[cutlass.Boolean, bool],
- ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
- """
- Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array
- (destination).
-
- :param tidx: The thread index in epilogue warp groups
- :type tidx: cutlass.Int32
- :param tAcc: The accumulator tensor to be copied and partitioned
- :type tAcc: cute.Tensor
- :param gC_mnl: The global tensor C
- :type gC_mnl: cute.Tensor
- :param epi_tile: The epilogue tiler
- :type epi_tile: cute.Tile
- :param use_2cta_instrs: Whether use_2cta_instrs is enabled
- :type use_2cta_instrs: bool
-
- :return: A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where:
- - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r)
- - tTR_tAcc: The partitioned accumulator tensor
- - tTR_rAcc: The accumulated tensor in register used to hold t2r results
- :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]
- """
- # Make tiledCopy for tensor memory load
- copy_atom_t2r = sm100_utils.get_tmem_load_op(
- self.cta_tile_shape_mnk,
- self.gemm_output_layout,
- self.out_dtype,
- self.acc_dtype,
- epi_tile,
- use_2cta_instrs,
- )
-
- # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE)
- tAcc_epi = cute.flat_divide(
- tAcc[((None, None), 0, 0, None)],
- epi_tile,
- )
- # (EPI_TILE_M, EPI_TILE_N)
- 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)
- # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE)
- tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi)
-
- # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, loopM, loopN, loopL)
- gC_mnl_epi = cute.flat_divide(gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile)
-
- # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, loopM, loopN, loopL)
- tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi)
- # (T2R, T2R_M, T2R_N)
- tTR_rAcc = cute.make_rmem_tensor(
- tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype
- )
-
- return tiled_copy_t2r, tTR_tAcc, tTR_rAcc
-
- def epilog_smem_copy_and_partition(
- self,
- tidx: cutlass.Int32,
- tTR_rC: cute.Tensor,
- sC: cute.Tensor,
- tiled_copy_t2r: cute.TiledCopy,
- ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
- """
- Create tiled copy for register to shared memory (R2S).
- """
- atom = cute.make_copy_atom(
- cute.nvgpu.CopyUniversalOp(),
- self.out_dtype,
- )
-
- tiled_copy_r2s = cute.make_tiled_copy_D(atom, tiled_copy_t2r)
- # (R2S, R2S_M, R2S_N, PIPE_D)
- thr_copy_r2s = tiled_copy_r2s.get_slice(tidx)
- tRS_sC = thr_copy_r2s.partition_D(sC)
- tRS_rC = tiled_copy_r2s.retile(tTR_rC)
- return tiled_copy_r2s, tRS_rC, tRS_sC
-
- @staticmethod
- def _compute_stages(
- tiled_mma: cute.TiledMma,
- mma_tiler_mnk: Tuple[int, int, int],
- a_dtype: Type[cutlass.Numeric],
- b_dtype: Type[cutlass.Numeric],
- out_dtype: Type[cutlass.Numeric],
- cta_tile: cute.Tile,
- sf_dtype: Type[cutlass.Numeric],
- sf_vec_size: int,
- final_scale_dtype: Type[cutlass.Numeric],
- num_smem_capacity: int,
- occupancy: int,
- ) -> Tuple[int, int, int, int, int]:
- """Computes the number of stages for A/B/C operands based on heuristics.
-
- :param tiled_mma: The tiled MMA object defining the core computation.
- :type tiled_mma: cute.TiledMma
- :param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler.
- :type mma_tiler_mnk: tuple[int, int, int]
- :param a_dtype: Data type of operand A.
- :type a_dtype: type[cutlass.Numeric]
- :param b_dtype: Data type of operand B.
- :type b_dtype: type[cutlass.Numeric]
- :param out_dtype: Data type of operand C (output).
- :type out_dtype: type[cutlass.Numeric]
- :param cta_tile: The CTA tile shape.
- :type cta_tile: cute.Tile
- :param sf_dtype: Data type of scale factor.
- :type sf_dtype: type[cutlass.Numeric]
- :param sf_vec_size: Vector size of scale factor.
- :type sf_vec_size: int
- :param num_smem_capacity: Total available shared memory capacity in bytes.
- :type num_smem_capacity: int
- :param occupancy: Target number of CTAs per SM (occupancy).
- :type occupancy: int
-
- :return: A tuple containing the computed number of stages for:
- (ACC stages, A/B operand stages, C stages)
- :rtype: tuple[int, int, int]
- """
- # Default ACC stages
- num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2
-
- # Default C stages
- num_c_stage = 1
-
- # Default Tile info stages
- num_tile_stage = 2
-
- # Metadata loader pipeline depth (meta warp -> epilogue). Lets the loader
- # prefetch the per-row {token_idx, combined_scale} ahead of the epilogue.
- num_meta_stage = 2
- # Per-row metadata smem held by the loader: token_idx (int32) +
- # combined_scale (final_scale_dtype), cta_M rows per stage.
- token_idx_bytes = cutlass.Int32.width // 8
- scale_bytes = final_scale_dtype.width // 8
- meta_smem_bytes = cta_tile[0] * (token_idx_bytes + scale_bytes) * num_meta_stage
-
- # Calculate smem layout and size for one stage of A, B, and C
- a_smem_layout_stage_one = sm100_utils.make_smem_layout_a(
- tiled_mma,
- mma_tiler_mnk,
- a_dtype,
- 1, # a tmp 1 stage is provided
- )
- b_smem_layout_staged_one = sm100_utils.make_smem_layout_b(
- tiled_mma,
- mma_tiler_mnk,
- b_dtype,
- 1, # a tmp 1 stage is provided
- )
-
- sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa(
- tiled_mma,
- mma_tiler_mnk,
- sf_vec_size,
- 1, # a tmp 1 stage is provided
- )
-
- sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb(
- tiled_mma,
- mma_tiler_mnk,
- sf_vec_size,
- 1, # a tmp 1 stage is provided
- )
-
- # satisfy 16B alignment for the output tensor
- swizzled_pad = 16 // (out_dtype.width // 8)
- c_smem_layout_staged_one = cute.make_layout(
- (cta_tile[0], cta_tile[1]), stride=(cta_tile[1] + swizzled_pad, 1)
- )
-
- ab_bytes_per_stage = (
- cute.size_in_bytes(a_dtype, a_smem_layout_stage_one)
- + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one)
- + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one)
- + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one)
- )
- # 1024B alignment for mbar
- mbar_helpers_bytes = 1024
-
- c_bytes_per_stage = cute.size_in_bytes(out_dtype, c_smem_layout_staged_one)
- c_bytes = c_bytes_per_stage * num_c_stage
-
- # Calculate A/B stages:
- # Start with total smem per CTA (capacity / occupancy)
- # Subtract reserved bytes, initial C stages bytes, and the per-row
- # metadata smem held by the loader warp
- # Divide remaining by bytes needed per A/B stage
- num_ab_stage = (
- num_smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes + meta_smem_bytes)
- ) // ab_bytes_per_stage
-
- return num_acc_stage, num_ab_stage, num_c_stage, num_tile_stage, num_meta_stage
-
- @staticmethod
- def _compute_grid(
- gemm_shape: Tuple[int, int, int],
- cta_tile_shape_mnk: Tuple[int, int, int],
- cluster_shape_mn: Tuple[int, int],
- max_active_clusters: cutlass.Constexpr,
- raster_along_m: bool,
- ) -> Tuple[utils.PersistentTileSchedulerParams, Tuple[int, int, int]]:
- """Use persistent tile scheduler to compute the grid size based on GEMM shape.
-
- :param gemm_shape: The GEMM computation shape (M, N, L)
- :type gemm_shape: tuple[int, int, int]
- :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile.
- :type cta_tile_shape_mnk: tuple[int, int, int]
- :param cluster_shape_mn: Shape of each cluster in M, N dimensions.
- :type cluster_shape_mn: tuple[int, int]
- :param max_active_clusters: Maximum number of active clusters.
- :type max_active_clusters: cutlass.Constexpr
- :param raster_along_m: Boolean, True to use raster along M.
- :type raster_along_m: bool
-
- :return: A tuple containing:
- - tile_sched_params: Parameters for the persistent tile scheduler.
- - grid: Grid shape for kernel launch.
- :rtype: Tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]]
- """
- (m, n, l) = gemm_shape # noqa: E741
-
- num_ctas_m = cute.ceil_div(m, cta_tile_shape_mnk[0])
- num_ctas_n = cute.ceil_div(n, cta_tile_shape_mnk[1])
- num_ctas_l = l
-
- num_ctas_mnl = (num_ctas_m, num_ctas_n, num_ctas_l)
- cluster_shape_mnl = (*cluster_shape_mn, 1)
-
- tile_sched_params = utils.PersistentTileSchedulerParams(
- num_ctas_mnl, cluster_shape_mnl, raster_along_m=raster_along_m
- )
- grid = utils.StaticPersistentTileScheduler.get_grid_shape(
- tile_sched_params, max_active_clusters
- )
-
- return tile_sched_params, grid
-
- @staticmethod
- def _get_tma_atom_kind(
- atom_sm_cnt: cutlass.Int32, mcast: cutlass.Boolean
- ) -> Union[cpasync.CopyBulkTensorTileG2SMulticastOp, cpasync.CopyBulkTensorTileG2SOp]:
- """
- Select the appropriate TMA copy atom based on the number of SMs and the multicast flag.
-
- :param atom_sm_cnt: The number of SMs
- :type atom_sm_cnt: cutlass.Int32
- :param mcast: The multicast flag
- :type mcast: cutlass.Boolean
-
- :return: The appropriate TMA copy atom kind
- :rtype: cpasync.CopyBulkTensorTileG2SMulticastOp or cpasync.CopyBulkTensorTileG2SOp
-
- :raise ValueError: If the atom_sm_cnt is invalid
- """
- if atom_sm_cnt == 2 and mcast:
- return cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.TWO)
- elif atom_sm_cnt == 2 and not mcast:
- return cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.TWO)
- elif atom_sm_cnt == 1 and mcast:
- return cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.ONE)
- elif atom_sm_cnt == 1 and not mcast:
- return cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE)
-
- raise ValueError(f"Invalid atom_sm_cnt: {atom_sm_cnt} and {mcast}")
-
- @staticmethod
- def is_valid_dtypes_and_scale_factor_vec_size(
- ab_dtype: Type[cutlass.Numeric],
- sf_dtype: Type[cutlass.Numeric],
- sf_vec_size: int,
- out_dtype: Type[cutlass.Numeric],
- ) -> bool:
- """
- Check if the dtypes are valid
-
- :param ab_dtype: The data type of the A and B operands
- :type ab_dtype: Type[cutlass.Numeric]
- :param sf_dtype: The data type of the scale factor
- :type sf_dtype: Type[cutlass.Numeric]
- :param sf_vec_size: The vector size of the scale factor
- :type sf_vec_size: int
- :param out_dtype: The data type of the output tensor
- :type out_dtype: Type[cutlass.Numeric]
-
- :return: True if the dtypes are valid, False otherwise
- :rtype: bool
- """
- is_valid = True
- if ab_dtype not in {
- cutlass.Float4E2M1FN,
- cutlass.Float8E5M2,
- cutlass.Float8E4M3FN,
- }:
- is_valid = False
-
- # Check valid sf_vec_size
- if sf_vec_size not in {16, 32}:
- is_valid = False
-
- # Check valid sf_dtype
- if sf_dtype not in {cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN}:
- is_valid = False
-
- # Check valid sf_dtype and sf_vec_size combinations
- if sf_dtype == cutlass.Float8E4M3FN and sf_vec_size == 32:
- is_valid = False
- if ab_dtype in {cutlass.Float8E5M2, cutlass.Float8E4M3FN} and sf_vec_size == 16:
- is_valid = False
-
- if out_dtype not in {cutlass.Float32, cutlass.Float16, cutlass.BFloat16}:
- is_valid = False
-
- return is_valid
-
- @staticmethod
- def is_valid_layouts(
- ab_dtype: Type[cutlass.Numeric],
- out_dtype: Type[cutlass.Numeric],
- a_major: str,
- b_major: str,
- out_major: str,
- ) -> bool:
- """
- Check if layouts and dtypes are valid combinations
-
- :param ab_dtype: The data type of the A and B operands
- :type ab_dtype: Type[cutlass.Numeric]
- :param out_dtype: The data type of the output tensor
- :type out_dtype: Type[cutlass.Numeric]
- :param a_major: The major dimension of the A tensor
- :type a_major: str
- :param b_major: The major dimension of the B tensor
- :type b_major: str
- :param out_major: The major dimension of the C tensor
- :type out_major: str
-
- :return: True if the layouts are valid, False otherwise
- :rtype: bool
- """
- is_valid = True
-
- if ab_dtype is cutlass.Float4E2M1FN and not (a_major == "k" and b_major == "k"):
- is_valid = False
- if out_dtype is cutlass.Float4E2M1FN and out_major == "m":
- is_valid = False
- return is_valid
-
- @staticmethod
- def is_valid_mma_tiler_and_cluster_shape(
- mma_tiler_mn: Tuple[int, int],
- cluster_shape_mn: Tuple[int, int],
- ) -> bool:
- """
- Check if the mma tiler and cluster shape are valid
-
- :param use_2cta_instrs: Whether to use 2 CTA groups
- :type use_2cta_instrs: bool
- :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler
- :type mma_tiler_mn: Tuple[int, int]
- :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster
- :type cluster_shape_mn: Tuple[int, int]
-
- :return: True if the mma tiler and cluster shape are valid, False otherwise
- :rtype: bool
- """
- is_valid = True
-
- # Skip invalid mma tile shape
- if mma_tiler_mn[0] not in (128, 256):
- is_valid = False
- # Skip invalid mma tile n
- if mma_tiler_mn[1] not in (64, 128, 192, 256):
- is_valid = False
-
- # Skip illegal cluster shape
- if (mma_tiler_mn[0] // cluster_shape_mn[0]) != 128:
- is_valid = False
-
- if (
- cluster_shape_mn[0] * cluster_shape_mn[1] > 16
- or cluster_shape_mn[0] <= 0
- or cluster_shape_mn[1] <= 0
- # Special cluster shape check for scale factor multicasts.
- # Due to limited size of scale factors, we can't multicast among more than 4 CTAs.
- or cluster_shape_mn[0] > 4
- or cluster_shape_mn[1] > 4
- or not is_power_of_2(cluster_shape_mn[0])
- or not is_power_of_2(cluster_shape_mn[1])
- ):
- is_valid = False
-
- return is_valid
-
- @staticmethod
- def is_valid_tensor_alignment(
- m: cutlass.Int64,
- n: cutlass.Int64,
- k: cutlass.Int64,
- l: cutlass.Int64, # noqa: E741
- ab_dtype: Type[cutlass.Numeric],
- out_dtype: Type[cutlass.Numeric],
- a_major: str,
- b_major: str,
- out_major: str,
- ) -> bool:
- """
- Check if the tensor alignment is valid
-
- :param m: The number of rows in the A tensor
- :type m: cutlass.Int64
- :param n: The number of columns in the B tensor
- :type n: cutlass.Int64
- :param k: The number of columns in the A tensor
- :type k: cutlass.Int64
- :param l: The number of columns in the C tensor
- :type l: cutlass.Int64
- :param ab_dtype: The data type of the A and B operands
- :type ab_dtype: Type[cutlass.Numeric]
- :param out_dtype: The data type of the output tensor
- :type out_dtype: Type[cutlass.Numeric]
- :param a_major: The major axis of the A tensor
- :type a_major: str
- :param b_major: The major axis of the B tensor
- :type b_major: str
- :param out_major: The major axis of the C tensor
- :type out_major: str
-
- :return: True if the problem shape is valid, False otherwise
- :rtype: bool
- """
- is_valid = True
-
- def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape):
- major_mode_idx = 0 if is_mode0_major else 1
- num_major_elements = tensor_shape[major_mode_idx]
- num_contiguous_elements = 16 * 8 // dtype.width
- return num_major_elements % num_contiguous_elements == 0
-
- if (
- not check_contigous_16B_alignment(ab_dtype, a_major == "m", (m, k, l))
- or not check_contigous_16B_alignment(ab_dtype, b_major == "n", (n, k, l))
- or not check_contigous_16B_alignment(out_dtype, out_major == "m", (m, n, l))
- ):
- is_valid = False
- return is_valid
-
- @classmethod
- def can_implement(
- cls,
- ab_dtype: Type[cutlass.Numeric],
- sf_dtype: Type[cutlass.Numeric],
- sf_vec_size: int,
- out_dtype: Type[cutlass.Numeric],
- mma_tiler_mn: Tuple[int, int],
- cluster_shape_mn: Tuple[int, int],
- m: cutlass.Int64,
- n: cutlass.Int64,
- k: cutlass.Int64,
- l: cutlass.Int64, # noqa: E741
- a_major: str,
- b_major: str,
- out_major: str,
- ) -> bool:
- """
- Check if the gemm can be implemented
-
- :param ab_dtype: The data type of the A and B operands
- :type ab_dtype: Type[cutlass.Numeric]
- :param sf_dtype: The data type of the scale factor
- :type sf_dtype: Type[cutlass.Numeric]
- :param sf_vec_size: The vector size of the scale factor
- :type sf_vec_size: int
- :param out_dtype: The data type of the output tensor
- :type out_dtype: Type[cutlass.Numeric]
- :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler
- :type mma_tiler_mn: Tuple[int, int]
- :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster
- :type cluster_shape_mn: Tuple[int, int]
- :param m: The number of rows in the A tensor
- :type m: cutlass.Int64
- :param n: The number of columns in the B tensor
- :type n: cutlass.Int64
- :param k: The number of columns in the A tensor
- :type k: cutlass.Int64
- :param l: The number of columns in the C tensor
- :type l: cutlass.Int64
- :param a_major: The major axis of the A tensor
- :type a_major: str
- :param b_major: The major axis of the B tensor
- :type b_major: str
- :param out_major: The major axis of the C tensor
- :type out_major: str
-
- :return: True if the gemm can be implemented, False otherwise
- :rtype: bool
- """
- can_implement = True
- # Skip unsupported types
- if not cls.is_valid_dtypes_and_scale_factor_vec_size(
- ab_dtype, sf_dtype, sf_vec_size, out_dtype
- ):
- can_implement = False
-
- # Skip unsupported layouts
- if not cls.is_valid_layouts(ab_dtype, out_dtype, a_major, b_major, out_major):
- can_implement = False
-
- # Skip invalid mma tile shape and cluster shape
- if not cls.is_valid_mma_tiler_and_cluster_shape(mma_tiler_mn, cluster_shape_mn):
- can_implement = False
- # Skip illegal problem shape for load/store alignment
- if not cls.is_valid_tensor_alignment(
- m, n, k, l, ab_dtype, out_dtype, a_major, b_major, out_major
- ):
- can_implement = False
- # Skip unsupported A/B layout
- if not (a_major == "k" and b_major == "k"):
- can_implement = False
- return can_implement
-
- @cute.jit
- def wrapper(
- self,
- a_ptr: cute.Pointer,
- b_ptr: cute.Pointer,
- a_sf_ptr: cute.Pointer,
- b_sf_ptr: cute.Pointer,
- c_ptr: cute.Pointer,
- alpha_ptr: cute.Pointer,
- tile_idx_to_group_idx_ptr: cute.Pointer,
- tile_idx_to_mn_limit_ptr: cute.Pointer,
- permuted_idx_to_expanded_idx_ptr: cute.Pointer,
- num_non_exiting_tiles_ptr: cute.Pointer,
- token_final_scales_ptr: cute.Pointer,
- m: cutlass.Int64,
- n: cutlass.Int64,
- k: cutlass.Int64,
- l: cutlass.Int64, # noqa: E741
- num_tokens: cutlass.Int64,
- top_k: cutlass.Int64,
- expert_capacity: cutlass.Int32,
- tile_size: cutlass.Constexpr,
- scaling_vector_size: cutlass.Constexpr,
- max_active_clusters: cutlass.Constexpr,
- stream: cuda.CUstream,
- epilogue_op: cutlass.Constexpr = lambda x: x,
- c_stride_row: cutlass.Int64 = cutlass.Int64(0),
- ):
- """Single-B wrapper.
-
- ``l`` is the number of experts in the (sole) B tensor.
- """
- scale_k = k // scaling_vector_size
- num_tiles = m // tile_size
-
- a = cute.make_tensor(a_ptr, layout=cute.make_ordered_layout((m, k, 1), order=(1, 0, 2)))
- a_sf = cute.make_tensor(
- a_sf_ptr,
- layout=cute.make_ordered_layout(
- (32, 4, m // 128, 4, scale_k // 4, 1), order=(2, 1, 4, 0, 3, 5)
- ),
- )
- # c supports strided output for locality domain shared buffers.
- # c_stride_row: row stride in output elements (0 = default = n)
- actual_c_stride_row = n if c_stride_row == 0 else c_stride_row
- c = cute.make_tensor(
- c_ptr,
- layout=cute.make_layout(
- (num_tokens, n, 1),
- stride=(actual_c_stride_row, 1, num_tokens * actual_c_stride_row),
- ),
- )
-
- alpha = cute.make_tensor(alpha_ptr, layout=cute.make_layout((l,)))
- b = cute.make_tensor(b_ptr, layout=cute.make_ordered_layout((n, k, l), order=(1, 0, 2)))
- b_sf = cute.make_tensor(
- b_sf_ptr,
- layout=cute.make_ordered_layout(
- (32, 4, n // 128, 4, scale_k // 4, l), order=(2, 1, 4, 0, 3, 5)
- ),
- )
-
- tile_idx_to_group_idx = cute.make_tensor(
- tile_idx_to_group_idx_ptr, layout=cute.make_layout((num_tiles,))
- )
- tile_idx_to_mn_limit = cute.make_tensor(
- tile_idx_to_mn_limit_ptr, layout=cute.make_layout((num_tiles,))
- )
- permuted_idx_to_expanded_idx = cute.make_tensor(
- permuted_idx_to_expanded_idx_ptr, layout=cute.make_layout((m,))
- )
- num_non_exiting_tiles = cute.make_tensor(
- num_non_exiting_tiles_ptr, layout=cute.make_layout((1,))
- )
- token_final_scales = cute.make_tensor(
- token_final_scales_ptr,
- layout=cute.make_ordered_layout((num_tokens, top_k), order=(1, 0)),
- )
-
- return self(
- a,
- b,
- c,
- a_sf,
- b_sf,
- tile_idx_to_group_idx,
- num_non_exiting_tiles,
- tile_idx_to_mn_limit,
- alpha,
- expert_capacity,
- max_active_clusters=max_active_clusters,
- stream=stream,
- permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx,
- token_final_scales=token_final_scales,
- epilogue_op=epilogue_op,
- )
-
-
-@cute.jit
-def cvt_sf_MKL_to_M32x4xrm_K4xrk_L(
- sf_ref_tensor: cute.Tensor,
- sf_mma_tensor: cute.Tensor,
-):
- """Convert scale factor tensor from MKL layout to mma specification M(32x4xrest_m)xK(4xrest_k)xL layout"""
- # sf_mma_tensor has flatten shape (32, 4, rest_m, 4, rest_k, l)
- # group to ((32, 4, rest_m), (4, rest_k), l)
- sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3)
- sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3)
- for i in cutlass.range(cute.size(sf_ref_tensor)):
- mkl_coord = sf_ref_tensor.layout.get_hier_coord(i)
- sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord]
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py
deleted file mode 100644
index 6f48df505098..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py
+++ /dev/null
@@ -1,8119 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""ModelingV2 target: Qwen3.8-Flash-Next-NVFP4 / sm_103 / tp1 — self-contained modeling code.
-
-The core holds the static structure (weights, derived operands, per-manager pool
-views); `PrefillTarget` (any batch with context rows) and `DecodeTarget`
-(generation-only) run a step over it, passing every argument to every op at the
-call site.
-
-The checkpoint is the composite Qwen4-Exp VLM export. The MTP module is built only
-when the engine has a ``speculative_config`` (else ``mtp.*`` stays in
-``weights.py``'s non-load set). The vision tower is always built; a text-only
-request never reaches it. Every speculation-dependent change is gated on a host
-int that is 1 when speculation is off. The PLE pools and the QSA side caches are
-written for every verify row, so `SpeculativeRollback` (registered with the spec
-worker) rewrites them to the accepted prefix after each verify step.
-
-Differences from other targets:
-
-* **Hyper-connections.** The inter-block residual is ``4 * 2560 = 10240`` wide.
- There is **no** ``model.language_model.norm.weight``: the stack-final
- ``hyper_connection_mixer`` is the final norm and ``lm_head`` reads its output.
- In the trunk each site is ``ops/hc/qwen4_exp_hc_inject_norm`` (the previous block's
- accumulate fused with this site's grouped norm; after a MoE block it also forms that
- block's output from the routed and shared outputs and the gate logit) and then the
- mixer: ``catalog/gemm/cublas_mm`` over the concatenated ``[down; block_inject]``
- weight, ``ops/hc/qwen4_exp_hc_lowrank_act`` and ``ops/hc/qwen4_exp_hc_up_gate_mean``
- (the up projection as a tcgen05 GEMM whose epilogue applies the sigmoid gate and the
- stream mean). ``DecodeTarget`` mixers run the first projection as one tcgen05 launch
- (``ops/hc/qwen4_exp_hc_decode_tc_first``, up to ``_HC_TC_FIRST_MAX_ROWS`` rows) or, at
- ``_HC_DECODE_GEMV_MAX_ROWS`` rows, the whole site as two GEMV launches
- (``ops/hc/qwen4_exp_hc_decode_first_cute``, ``ops/hc/qwen4_exp_hc_decode_mix_fp4_cute``);
- at the 48 MoE sites the mixer's last launch also writes the routed experts' NVFP4
- input. The MTP layer's sites run the same fused site; its front end's grouped norm
- (``_grouped_norm``) and the accumulates before the PLE block and after the MTP layer's
- MoE block (``qwen4_exp_post_inject``) run alone.
-* **QSA sparse attention.** The index plane is maintained on every step so a
- request can cross ``dense_seq_len_threshold`` later without rebuilding its
- prefix; sparse selection replaces the dense core only above it. The MTP draft
- layer runs the same branch over its own indexer and its own manager's planes.
-* **PLE.** Decoder layer 1 only (``ple_layer_ids: [2]`` is one-based): an n-gram
- side path over a 320,001,536-row fp8 table (47.68 GiB of a 123.5 GiB checkpoint).
-* **NVFP4 MoE on every layer** (routed experts only); everything else is bf16. The
- trunk's routed experts run on trtllm-gen's block-scale runner (``TRTLLM``, and what
- ``AUTO`` selects) or on the CuteDSL sort / grouped-GEMM pipeline (``CUTEDSL``),
- chosen by ``moe_config.backend``; any other request is refused. The MTP draft
- layer's FP8 block-scale experts run on trtllm-gen under either. Under ``TRTLLM`` the
- prefill-sized launches (``_CUTE_DSL_MOE_MIN_ROWS``) run the CuteDSL pipeline over a
- second copy of the experts in its layout, about 1.3 GiB a layer, kept only when a step
- can reach that size and ``TRTLLM_MODELING_V2_MOE_PREFILL_COPY`` does not turn it off
- (``_moe_prefill_copy``). A graph-captured decode step routes its block
- projections through ``ops/gemm/low_m_bf16_mm`` (FlashInfer low-M kernels at
- M <= 32 on the shapes it lists, cuBLAS otherwise); see ``DecodeTarget.forward``.
-
-Integration facts this file depends on:
-
-1. The engine rebuilds the composite config and normalizes its text tower
- (``config_utils.load_pretrained_config``). The shell adopts that text config as
- ``model_config.pretrained_config`` before building the core, because the
- recurrent-state and PLE pool geometry are read from it.
-2. ``position_ids`` arrives in the mRoPE layout (``[3, 1, T]``) even for text-only
- batches, because the normalized text config carries ``rope_scaling['type'] ==
- 'mrope'``. The main and MTP rope read all three rows; the QSA index plane takes
- coordinate 0 only (see ``_QSA_MROPE_SECTION``).
-3. The engine owns the multimodal input processor (resolved off the checkpoint's
- ``architectures[0]``; modeling_v2 has no hook), so it is the built-in
- ``Qwen3VLInputProcessorBase``. This target asserts every ``multimodal_data`` key
- where it is read.
-
-**KV storage.** The paged K/V pool holds the checkpoint's bf16, fp8-e4m3 when the engine
-resolves ``kv_cache_quant_algo: FP8`` (``kv_cache_config.dtype: fp8``), or NVFP4 when it resolves
-``NVFP4`` (``kv_cache_config.dtype: nvfp4``): a lossy substitution, since the checkpoint's KV is
-bf16 and the fp8 and NVFP4 attention kernels run q, K, V and the softmax probabilities through
-e4m3. An NVFP4 pool stores e2m1 pairs (uint8, half a head wide) and keeps one e4m3 scale per 16
-values in a second pool; the second-level scales are 1.0. The NVFP4 attention kernels exist for
-the paged-context path only, so an NVFP4 engine needs chunked prefill, block reuse or
-speculative decoding. The QSA index and position planes keep their own types under all three.
-
-**Recurrent-state storage.** The Gated-DeltaNet recurrent pool holds the checkpoint's fp32
-(``mamba_ssm_dtype``) or bf16 when ``kv_cache_config.mamba_ssm_cache_dtype`` asks for it, which
-is also what ``auto`` resolves to for this checkpoint (`_SSM_STATE_DTYPES`): a lossy
-substitution, since the state is rounded to bf16 once per decode token, once per prefill call and
-at every cached-replay checkpoint commit. The conv, PLE and replay-history buffers keep their own
-types under both. Over a bf16 pool a decode step takes one recurrence kernel at every batch size,
-so a row's tokens do not depend on its batch; see ``DecodeTarget._gdn_block``.
-
-**Breakable prefill CUDA graph.** With ``prefill_cuda_graph_backend: breakable`` the engine
-captures the core's ``forward`` (the body, not the shell) for a step with context rows at each of
-``prefill_capture_num_tokens``, pads any such step up to the nearest bucket and replays it without
-calling the body. The shell therefore calls the core's ``begin_step`` before the body on every
-step: the asserts, the cache-manager bind, the target choice and the prefill target's host
-bookkeeping (``PrefillTarget.prepare``, one immutable ``_PrefillStep``) all live there. The prefill
-``forward`` is row-wise over the ``T`` rows it is handed, with four breaks
-(``catalog/graph/eager_on_graph``) that run eagerly between the captured segments and read the
-step: the vision splice, the PLE block, and each layer's attention core and Gated-DeltaNet
-recurrence. Each writes rows ``[:n]`` of a fixed output and zeroes rows ``[n:]``. Under MTP the
-prefill body returns one packed ``[T, hidden + stream]`` tensor that ``split_spec_output`` splits,
-and the shell hands the spec worker the batch's own ``n`` rows. With the backend off every step is
-eager and unpadded and the breaks are plain calls.
-
-**Host offload.** By default the PLE table lives in host memory the device
-addresses (`weights._host_resident_empty`), and the gather reads it unchanged;
-``TRTLLM_MODELING_V2_PLE_HOST_OFFLOAD=0`` restores the resident declaration. Output
-is identical, so only the device-memory delta sees it. The gather stays
-device-driven because row indices are produced on device inside a captured decode
-forward, which cannot express a host readback.
-"""
-
-import math
-import os
-import re
-from typing import NamedTuple
-
-import torch
-from torch import nn
-from transformers import PretrainedConfig
-
-from tensorrt_llm._torch._experimental.modeling_v2._router_index import step_contract_enabled
-from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.flashinfer_silu_and_mul import ( # noqa: E501
- flashinfer_silu_and_mul,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import (
- fused_qk_norm_rope,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import (
- QM_FP8_KV_CACHE,
- QM_NVFP4_KV_CACHE,
- thop_attention,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm
-from tensorrt_llm._torch._experimental.modeling_v2.catalog.graph.eager_on_graph import (
- eager_on_graph,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import (
- fp4_quantize,
-)
-from tensorrt_llm._torch.attention.backends.fmha.phased import get_spec_decoding_position_offsets
-from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata
-from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata
-from tensorrt_llm._torch.model_config import ModelConfig
-from tensorrt_llm._torch.models.modeling_utils import DecoderModel, DecoderModelForCausalLM
-from tensorrt_llm._torch.modules.fla.cached_replay import CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE
-from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor
-from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks
-from tensorrt_llm._torch.speculative import get_spec_worker, should_use_separate_draft_kv_cache
-from tensorrt_llm._torch.utils import ActivationType
-from tensorrt_llm.inputs.multimodal import get_multimodal_item_metadata
-from tensorrt_llm.logger import logger
-from tensorrt_llm.quantization.mode import QuantAlgo
-
-from . import weights as _weights
-from .ops.activation.fused_sigmoid_mul_inplace import fused_sigmoid_mul_inplace
-from .ops.activation.gelu_erf import gelu_erf
-from .ops.activation.gelu_tanh import gelu_tanh
-from .ops.activation.silu_and_mul_rows import silu_and_mul_rows
-from .ops.attention.qsa_canonical_selection import qsa_canonical_selection
-from .ops.attention.qsa_expand_block_indices import qsa_expand_block_indices
-from .ops.attention.qsa_mtp_reuse_topk import qsa_mtp_reuse_topk
-from .ops.attention.qsa_nvfp4_pages_to_bf16 import qsa_nvfp4_pages_to_bf16
-from .ops.attention.qsa_paged_index_scores import qsa_paged_index_scores
-from .ops.attention.qsa_paged_index_store import qsa_paged_index_store
-from .ops.attention.qsa_paged_kv_store import qsa_paged_kv_store
-from .ops.attention.qsa_paged_kv_store_nvfp4 import qsa_paged_kv_store_nvfp4
-from .ops.attention.qsa_paged_sparse_gqa import qsa_paged_sparse_gqa
-from .ops.attention.qsa_paged_sparse_gqa_decode import MAX_ROWS as _SPARSE_DECODE_MAX_ROWS
-from .ops.attention.qsa_paged_sparse_gqa_decode import qsa_paged_sparse_gqa_decode
-from .ops.attention.qsa_paged_sparse_gqa_decode import supports as _sparse_decode_supports
-from .ops.attention.qsa_paged_sparse_gqa_nvfp4 import qsa_paged_sparse_gqa_nvfp4
-from .ops.attention.qsa_paged_sparse_gqa_nvfp4_decode import qsa_paged_sparse_gqa_nvfp4_decode
-from .ops.attention.qsa_paged_sparse_gqa_nvfp4_verify import qsa_paged_sparse_gqa_nvfp4_verify
-from .ops.attention.qsa_paged_sparse_gqa_nvfp4_verify import takes_grouped as _verify_takes_grouped
-from .ops.attention.qsa_prefill_compress import qsa_prefill_compress
-from .ops.attention.qsa_rollback_planes import plane_table as _plane_table
-from .ops.attention.qsa_rollback_planes import qsa_rollback_gather, qsa_rollback_restore
-from .ops.attention.qsa_sparse_gqa_rowtile import qsa_sparse_gqa_rowtile
-from .ops.attention.qsa_topk_expand_decode import qsa_topk_expand_decode
-from .ops.attention.qsa_topk_expand_decode import supports as _topk_expand_supports
-from .ops.attention.qsa_topk_select import qsa_topk_select
-from .ops.attention.qwen4_exp_attn_decode_prologue import SCORE_ROWS as _PROLOGUE_SCORE_ROWS
-from .ops.attention.qwen4_exp_attn_decode_prologue import qwen4_exp_attn_decode_prologue
-from .ops.attention.qwen4_exp_qk_norm_rope import qwen4_exp_qk_norm_rope
-from .ops.attention.spec_rollback_commit import spec_rollback_commit, spec_rollback_stage
-from .ops.attention.verify_step_setup import _BLOCK_S as _VERIFY_SETUP_MAX_SEQS
-from .ops.attention.verify_step_setup import verify_step_setup
-from .ops.attention.vision_rope_2d import vision_rope_2d
-from .ops.gemm.bf16_gemv_pdl import MAX_ROWS as _GEMV_MAX_ROWS
-from .ops.gemm.bf16_gemv_pdl import TILING as _GEMV_TILED_K
-from .ops.gemm.bf16_gemv_pdl import bf16_gemv_pdl
-from .ops.gemm.fp8_gemv import (
- fp8_gemv_conv1d_op,
- fp8_gemv_conv1d_verify_op,
- fp8_gemv_op,
- pack_from_bf16,
-)
-from .ops.gemm.fp8_per_tensor_mm import fp8_per_tensor_mm, fp8_per_tensor_weight
-from .ops.gemm.low_m_bf16_mm import low_m_bf16_mm
-from .ops.gemm.mid_rows_mm import MidWeight, mid_rows_mm_op
-from .ops.gemm.mid_rows_mm import precompile as mid_rows_precompile
-from .ops.gemm.nvfp4_gemv import (
- nvfp4_gemv_op,
- nvfp4_split_gemv_conv1d_op,
- nvfp4_split_gemv_conv1d_verify_op,
- nvfp4_split_gemv_op,
- precompile_split,
-)
-from .ops.gemm.nvfp4_gemv import pack_from_bf16 as nvfp4_pack_from_bf16
-from .ops.gemm.nvfp4_gemv import pack_split_from_bf16 as nvfp4_pack_split_from_bf16
-from .ops.gemm.nvfp4_gemv import precompile as nvfp4_gemv_precompile
-from .ops.hc.qwen4_exp_hc_decode_first_cute import qwen4_exp_hc_decode_first_cute
-from .ops.hc.qwen4_exp_hc_decode_mix_fp4_cute import (
- qwen4_exp_hc_decode_mix_cute,
- qwen4_exp_hc_decode_mix_fp4_cute,
-)
-from .ops.hc.qwen4_exp_hc_decode_rows_tc import MAX_ROWS as _HC_DECODE_ROWS_MAX
-from .ops.hc.qwen4_exp_hc_decode_rows_tc import (
- qwen4_exp_hc_decode_rows_first,
- qwen4_exp_hc_decode_rows_mix,
-)
-from .ops.hc.qwen4_exp_hc_decode_tc_first import precompile as hc_tc_first_precompile
-from .ops.hc.qwen4_exp_hc_decode_tc_first import qwen4_exp_hc_decode_tc_first
-from .ops.hc.qwen4_exp_hc_first_rescale_tc import qwen4_exp_hc_first_rescale_tc
-from .ops.hc.qwen4_exp_hc_inject_norm import qwen4_exp_hc_inject_norm
-from .ops.hc.qwen4_exp_hc_inject_rstd import qwen4_exp_hc_inject_rstd
-from .ops.hc.qwen4_exp_hc_lowrank_act import qwen4_exp_hc_lowrank_act
-from .ops.hc.qwen4_exp_hc_up_gate_mean import qwen4_exp_hc_up_gate_mean
-from .ops.hc.qwen4_exp_hc_up_gate_mean_fp4 import qwen4_exp_hc_up_gate_mean_fp4
-from .ops.hc.qwen4_exp_hc_up_gate_mean_rstd import (
- qwen4_exp_hc_up_gate_mean_rstd,
- qwen4_exp_hc_up_gate_mean_rstd_fp4,
-)
-from .ops.hc.qwen4_exp_post_inject import qwen4_exp_post_inject
-from .ops.moe.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion import (
- cute_dsl_nvfp4_gather_grouped_gemm_act_fusion,
-)
-from .ops.moe.cute_dsl_nvfp4_grouped_gemm import cute_dsl_nvfp4_grouped_gemm
-from .ops.moe.fp4_block_scale_moe_renorm_permuted import fp4_block_scale_moe_renorm_permuted
-from .ops.moe.fp4_block_scale_moe_renorm_routed import fp4_block_scale_moe_renorm_routed
-from .ops.moe.fp4_moe_renorm_cute_dsl import fp4_moe_renorm_cute_dsl
-from .ops.moe.fp8_block_scale_moe_runner import fp8_block_scale_moe_runner
-from .ops.moe.fused_sigmoid_gate_mul_add import fused_sigmoid_gate_mul_add
-from .ops.moe.moe_sort import moe_sort
-from .ops.moe.moe_unpermute_inplace import moe_unpermute_inplace
-from .ops.moe.renorm_moe_routing_op import renorm_moe_routing_op
-from .ops.norm.flashinfer_gemma_rmsnorm import flashinfer_gemma_rmsnorm
-from .ops.norm.layer_norm import layer_norm
-from .ops.norm.rms_norm_gated_token_major import rms_norm_gated_token_major
-from .ops.ple.ple_gate_conv_inject_norm import ple_gate_conv_inject_norm
-from .ops.ple.ple_ngram_embed import ple_ngram_embed
-from .ops.quantization.bf16_to_fp32 import bf16_to_fp32
-from .ops.quantization.fp8_quantize_1x128 import fp8_quantize_1x128
-from .ops.ssm.causal_conv1d_triton_update import causal_conv1d_triton_update
-from .ops.ssm.causal_conv1d_update import causal_conv1d_update
-from .ops.ssm.cu_seqlens_from_seq_lens import cu_seqlens_from_seq_lens
-from .ops.ssm.flashinfer_chunk_gated_delta_rule import flashinfer_chunk_gated_delta_rule
-from .ops.ssm.fused_gdn_gating import fused_gdn_gating
-from .ops.ssm.fused_recurrent_gated_delta_rule_cached_replay_update import (
- fused_recurrent_gated_delta_rule_cached_replay_update,
-)
-from .ops.ssm.fused_recurrent_gated_delta_rule_update import fused_recurrent_gated_delta_rule_update
-from .ops.ssm.gdn_cached_replay_cute import gdn_cached_replay_cute
-from .ops.ssm.gdn_decode_gated_norm_pdl import MAX_ROWS as _GDN_DECODE_FUSED_MAX_ROWS
-from .ops.ssm.gdn_decode_gated_norm_pdl import gdn_decode_gated_norm_pdl
-from .ops.ssm.gdn_decode_rows_gated_norm import gdn_decode_rows, gdn_decode_rows_gated_norm
-from .ops.ssm.gdn_in_proj_conv1d_pdl import MAX_ROWS as _GDN_CONV_FUSED_MAX_ROWS
-from .ops.ssm.gdn_in_proj_conv1d_pdl import gdn_in_proj_conv1d_pdl
-from .ops.ssm.gdn_prefill_conv_post import gdn_prefill_conv_post
-from .ops.ssm.pack_gdn_decode_qkv import pack_gdn_decode_qkv
-from .ops.vision.patch_embed_conv3d import patch_embed_conv3d
-from .ops.vision.pos_embed_interpolate import pos_embed_interpolate
-
-# The GPU architecture this target is certified for; a direct instantiation could
-# bypass routing.
-_SM = (10, 3)
-
-# QSA backend metadata buffers rebuilt by prepare(). Their absence means the engine
-# did not select the QSA metadata and the index plane would never be written.
-_QSA_FIELDS = (
- "qsa_block_table",
- "qsa_req_idx_per_token",
- "qsa_logical_positions",
- "qsa_sequence_lengths",
- "qsa_visible_blocks",
- "qsa_topk_indices",
- "qsa_topk_row_starts",
-)
-
-# Fields read off `attn_metadata.mamba_metadata`. Absence is silent: a smoke run
-# passes with no recurrent state at all.
-_MAMBA_FIELDS = (
- "state_indices",
- "has_initial_states",
- "query_start_loc",
- "query_start_loc_long",
-)
-
-
-# A context launch at or above this many rows takes `qsa_sparse_gqa_rowtile`;
-# measured on GB300 it loses at 543 rows and wins from 4096 rows up.
-_ROWTILE_MIN_ROWS = 2048 # capacity: fixed: a dispatch threshold, not a buffer size
-# Over an NVFP4 pool, a step with more context rows than this reads its context requests'
-# pages dequantized into a bf16 scratch through `qsa_sparse_gqa_rowtile`. GB300, a layer:
-# 16384 context rows 1.36 ms there against 77-948 ms for the NVFP4 per-row kernel.
-_NVFP4_ROWTILE_MIN_ROWS = 256 # capacity: fixed: a dispatch threshold, not a buffer size
-# The bf16 K/V scratch those context rows read holds every page of each context request's KV
-# history, not just the step's rows. Consecutive context requests are dequantized and attended
-# in groups whose pages fit this many bytes; a request with more pages forms a group of its
-# own, so the scratch peaks at the larger of this and the longest request's history.
-_NVFP4_CONTEXT_SCRATCH_BYTES = 128 << 20 # capacity: fixed: a scratch memory budget
-
-#: Context-step MoE launches of at least this many rows run on the CuTe DSL MoE
-#: (`fp4_moe_renorm_cute_dsl`, from the second expert copy `weights.py` keeps in
-#: its layout); smaller ones, every decode one among them, keep the trtllm-gen
-#: runner. From GB300 micro-benches, us a layer, CuTe DSL vs trtllm-gen: 2048
-#: rows 275 vs 273, 8192 rows 451 vs 555, 16384 rows 656 vs 882.
-_CUTE_DSL_MOE_MIN_ROWS = 4096 # capacity: fixed: a dispatch threshold, not a buffer size
-
-#: Widest decode projection that `bf16_gemv_pdl` takes at 2..`_GEMV_MAX_ROWS` rows: the shared
-#: expert's one-column gate logit, on one CTA under `low_m_bf16_mm` (~3 us a layer at 4 rows).
-_GEMV_NARROW_N = 8 # capacity: fixed: a dispatch threshold, not a buffer size
-
-
-def _decode_projection(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor:
- """A graph-captured decode step's bf16 projection, always a PDL launch.
-
- One row runs `bf16_gemv_pdl` on every K it tiles (each decode projection of this
- checkpoint): its weight streams into registers before it waits on the producer, and it
- triggers its own dependents. So do up to `_GEMV_MAX_ROWS` rows of a projection at most
- `_GEMV_NARROW_N` wide (the shared expert's gate logit), which `low_m_bf16_mm` would run
- on one CTA. Other calls run `low_m_bf16_mm` with its FlashInfer branches' PDL on (its
- cuBLAS branch, above M=32 and off its table, ignores the flag). Every kernel either
- path launches reads the activation only after `griddepcontrol.wait`.
- """
- rows = mat_a.shape[0]
- if mat_b.shape[0] in _GEMV_TILED_K and (
- rows == 1 or (rows <= _GEMV_MAX_ROWS and mat_b.shape[1] <= _GEMV_NARROW_N)
- ):
- return bf16_gemv_pdl(mat_a, mat_b, True)
- return low_m_bf16_mm(mat_a, mat_b, True)
-
-
-# Per-call constants of the dense attention call shape: packed-QKV causal GQA, RoPE
-# applied outside (fused_qk_norm_rope), bf16, fp8-e4m3 or NVFP4 pool. The sparse / skip-softmax
-# group stays None/0: QSA calls selection and sparse attention as separate kernels and
-# `num_sparse_topk` is 0. `quant_mode`, the two KV scales (all three follow the pool's
-# storage) and `predicted_tokens_per_seq` are passed beside these.
-_CALL_CONSTANTS = dict(
- output_sf=None,
- k=None,
- v=None,
- is_fused_qkv=True,
- update_kv_cache=True,
- attention_input_type=0,
- is_mla_enable=False,
- mask_type=1,
- q_scaling=1.0,
- attention_sinks=None,
- # In-kernel RoPE disabled; rope_* values are inert placeholders.
- position_embedding_type=0,
- rotary_inv_freq=None,
- rotary_cos_sin=None,
- rope_dim=0,
- rope_base=10000.0,
- rope_scale_type=0,
- rope_scale=1.0,
- rope_short_m_scale=1.0,
- rope_long_m_scale=1.0,
- rope_max_positions=1024,
- rope_original_max_positions=1024,
- out_scale=None,
- latent_cache=None,
- q_pe=None,
- q_lora_rank=None,
- kv_lora_rank=None,
- qk_nope_head_dim=None,
- qk_rope_head_dim=None,
- v_head_dim=None,
- rope_append=None,
- chunked_prefill_buffer_batch_size=1,
- attention_chunk_size=None,
- softmax_stats_tensor=None,
- sparse_kv_indices=None,
- sparse_kv_offsets=None,
- sparse_attn_indices=None,
- sparse_attn_offsets=None,
- sparse_attn_indices_block_size=0,
- mrope_rotary_cos_sin=None,
- mrope_position_deltas=None,
- helix_position_offsets=None,
- helix_is_inactive_rank=None,
- cross_kv=None,
- relative_attention_bias=None,
- relative_attention_max_distance=0,
- # MLA-only surface, unused here.
- kv_norm_weight=None,
- kv_norm_eps=1e-6,
- skip_correction_threshold=0.0,
-)
-
-# ------------------------------------------------------------ vision ------
-
-#: Tower rope base (`vision_config` has no rope field); unrelated to the text 1e7.
-_VISION_ROPE_THETA = 10000.0
-
-#: LayerNorm epsilon of `norm1`, `norm2`, `merger.norm`: a literal in the reference.
-_VISION_NORM_EPS = 1e-6
-
-#: `mask_type` 0 = padding mask: bidirectional within a batch segment. Causal (1)
-#: would make a patch depend on raster order.
-_VISION_MASK_PADDING = 0
-
-# Per-call constants of the tower's attention: the encoder-only mode of
-# `attention/thop_attention`. Omitting `kv_cache_block_offsets`,
-# `host_kv_cache_pool_pointers` and `host_kv_cache_pool_mapping` selects that mode.
-# `update_kv_cache=False` is rejected ("KV cache update cannot be disabled now"),
-# so it stays True.
-_VISION_CALL_CONSTANTS = dict(
- k=None,
- v=None,
- output_sf=None,
- kv_cache_block_offsets=None,
- host_kv_cache_pool_pointers=None,
- host_kv_cache_pool_mapping=None,
- tokens_per_block=None,
- use_paged_context_fmha=False,
- is_fused_qkv=True,
- update_kv_cache=True,
- predicted_tokens_per_seq=1,
- # No pool, so no layer to address.
- local_layer_idx=0,
- max_context_q_len_override=None,
- cache_indirection=None,
- kv_scale_orig_quant=None,
- kv_scale_quant_orig=None,
- out_scale=None,
- quant_mode=0,
- beam_width=1,
- mask_type=_VISION_MASK_PADDING,
- q_scaling=1.0,
- # Rotation is applied outside (`ops/attention/vision_rope_2d`); rope_* inert.
- position_embedding_type=0,
- rotary_inv_freq=None,
- rotary_cos_sin=None,
- rope_dim=0,
- rope_base=10000.0,
- rope_scale_type=0,
- rope_scale=1.0,
- rope_short_m_scale=1.0,
- rope_long_m_scale=1.0,
- rope_max_positions=1024,
- rope_original_max_positions=1024,
- attention_input_type=0,
- is_mla_enable=False,
- chunked_prefill_buffer_batch_size=1,
- attention_sinks=None,
- block_ids_per_seq=None,
- latent_cache=None,
- q_pe=None,
- q_lora_rank=None,
- kv_lora_rank=None,
- qk_nope_head_dim=None,
- qk_rope_head_dim=None,
- v_head_dim=None,
- rope_append=None,
- attention_chunk_size=None,
- softmax_stats_tensor=None,
- is_spec_decoding_enabled=False,
- use_spec_decoding=False,
- is_spec_dec_tree=False,
- spec_decoding_generation_lengths=None,
- spec_decoding_position_offsets=None,
- spec_decoding_packed_mask=None,
- spec_decoding_bl_tree_mask_offset=None,
- spec_decoding_bl_tree_mask=None,
- spec_bl_tree_first_sparse_mask_offset_kv=None,
- sparse_kv_indices=None,
- sparse_kv_offsets=None,
- sparse_attn_indices=None,
- sparse_attn_offsets=None,
- sparse_attn_indices_block_size=0,
- mrope_rotary_cos_sin=None,
- mrope_position_deltas=None,
- helix_position_offsets=None,
- helix_is_inactive_rank=None,
- is_cross=False,
- cross_kv=None,
- relative_attention_bias=None,
- relative_attention_max_distance=0,
- kv_norm_weight=None,
- kv_norm_eps=1e-6,
- skip_correction_threshold=0.0,
-)
-
-#: Head geometries the encoder-only mode is certified at: MHA, `head_size` 72 or
-#: 128, context rows only. `head_size` 72 is context-phase only ("unsupported
-#: head_size 72" in generation).
-_VISION_CERTIFIED_HEADS = {(16, 16, 72), (16, 16, 128)}
-
-#: Head widths `ops/attention/vision_rope_2d` is certified at.
-_VISION_ROPE_HEAD_DIMS = (40, 64, 72, 128)
-
-#: The most patch rows one tower call carries, the ceiling the vision ops' cells reach:
-#: 8192 placeholder tokens x `spatial_merge_size`**2 patch rows each. `VisionTower` splits
-#: a request's frames into calls under it, so the bound a request meets is per frame: one
-#: frame's `h * w` patches.
-_VISION_CALL_MAX_PATCHES = 32768
-
-# NVFP4 block size: one e4m3 scale per 16 contiguous elements along K.
-_SF_VEC = _weights.SF_VEC
-# Routing happens outside (renorm_moe_routing_op), so the runner's routing method is
-# inert, but it must be 1 (Renormalize): the tier table covers (512 experts, 16
-# top-k) only for it, and a tier miss is a logged error with no exception and no
-# kernel launch.
-_ROUTING_METHOD_RENORM = 1
-_ACT_TYPE_SWIGLU = 0
-
-# The trunk's routed-expert backend by the requested `moe_config.backend`
-# (`ModelConfig.moe_backend_requested`; `moe_backend` itself is already resolved,
-# AUTO -> CUTLASS). Any other request is refused at construction. The MTP draft
-# layer's FP8 block-scale experts run on trtllm-gen whatever this says.
-_TRUNK_MOE_BACKENDS = {"AUTO": "TRTLLM", "TRTLLM": "TRTLLM", "CUTEDSL": "CUTEDSL"}
-# The CuteDSL grouped GEMMs' MMA tile M, which is also the tile `moe_sort` pads each
-# expert's rows to: the GEMMs walk the sort's tile tables. sm_103 admits 128 and 256.
-_CUTEDSL_TILE_SIZE = 128
-# The longest MTP draft the CuteDSL pipeline is certified at: its cells capture up to 128 rows
-# (32 requests of 4 verify rows).
-_CUTEDSL_MAX_DRAFT_LEN = 3
-
-# Skip sentinel the conv ops and paged short conv compare `cache_indices` against
-# (the runtime's mamba `PAD_SLOT_ID`). Never a substitution: the scheduler's slot
-# vector gives every row a real slot, and `ops/ple/ngram_token_history` reads its
-# carry via `index_select`, so a -1 there is a device-side assert that poisons the
-# process.
-_PAD_SLOT_ID = -1
-
-# Gated-DeltaNet gating form: the pair the decode kernels compute with and
-# `gdn_prefill_conv_post`'s default, so prefill and decode compute the same `g`.
-_SOFTPLUS_BETA = 1.0
-_SOFTPLUS_THRESHOLD = 20.0
-
-# Gated-DeltaNet Q/K L2-norm epsilon: a literal, not `rms_norm_eps`.
-_L2_NORM_EPS = 1e-6
-
-# PLE gate magnitude floor before its signed square root (a model constant).
-_PLE_CLAMP_MIN = 1e-6
-
-# --------------------------------------------------------------- host offload
-#
-# Where the PLE n-gram table lives (module docstring, *Host offload*). No engine
-# arg owns this, hence an env var, read once at construction (the table is
-# allocated once, during `load_weights`). Default on: 47.68 GiB otherwise comes
-# out of the KV pool.
-_PLE_HOST_OFFLOAD_ENV = "TRTLLM_MODELING_V2_PLE_HOST_OFFLOAD"
-
-#: Spellings of "off"; anything else, including empty, leaves offload on.
-_PLE_HOST_OFFLOAD_OFF = frozenset({"0", "false", "no", "off"})
-
-
-def _ple_host_offload_enabled() -> bool:
- """Whether the n-gram table is allocated in host memory."""
- return os.environ.get(_PLE_HOST_OFFLOAD_ENV, "1").strip().lower() not in _PLE_HOST_OFFLOAD_OFF
-
-
-# Whether trtllm-gen keeps the second, CuteDSL-layout copy of the routed experts for the
-# launches of at least `_CUTE_DSL_MOE_MIN_ROWS` rows (module docstring). No engine arg owns
-# it, hence an env var read at construction; "off" (the spellings above) drops the copy and
-# every launch runs on trtllm-gen.
-_MOE_PREFILL_COPY_ENV = "TRTLLM_MODELING_V2_MOE_PREFILL_COPY"
-
-
-def _moe_prefill_copy(trunk_moe_backend: str, max_num_tokens: int) -> bool:
- """Whether the trunk keeps the second expert copy: only under trtllm-gen, only when a
- step's rows (at most `max_num_tokens`) can reach the CuteDSL dispatch threshold, and not
- when the env var turns it off."""
- return (
- trunk_moe_backend == "TRTLLM"
- and max_num_tokens >= _CUTE_DSL_MOE_MIN_ROWS
- and os.environ.get(_MOE_PREFILL_COPY_ENV, "1").strip().lower() not in _PLE_HOST_OFFLOAD_OFF
- )
-
-
-# QSA's fixed constants (backend `constants.py`): one index KV head, three position
-# axes, index head exactly twice the rotary width (128 == 2 * 64).
-_QSA_INDEX_KV_HEADS = 1
-_QSA_POSITION_AXES = 3
-_QSA_INDEX_HEAD_TO_ROTARY_RATIO = 2
-
-# The QSA index plane takes coordinate 0 alone, a deliberate departure from the
-# checkpoint. `qsa_prefill_compress` (key side, `mrope_section`) and
-# `fused_qk_norm_rope` on `index_scratch` (query side, `use_mrope`) must agree:
-# both read the three coordinates or neither, else scores mix two position
-# encodings. Both-on is certified (`mrope_section` (11, 11, 10)); None is a
-# modeling choice. No cost on text; on an image or video batch the index plane sees
-# a frame at one position, changing block selection only above
-# `dense_seq_len_threshold`.
-# To flip, set the config's `(s0, s1, s2)` and `_QSA_INDEX_USE_MROPE = True`.
-_QSA_MROPE_SECTION = None
-_QSA_INDEX_USE_MROPE = False
-
-# ---------------------------------------------------------------- MTP ------
-
-#: The checkpoint's MTP module count (`mtp_num_hidden_layers`), replayed
-#: `max_draft_len` times; NOT the draft length.
-_MTP_NUM_LAYERS = 1
-
-#: Set to 1 to give the FP8 draft head an NVFP4 copy of its lm_head (e2m1 values, one e4m3
-#: scale per 16 hidden elements of a vocabulary row, an fp32 global scale; built at load, 357
-#: MB in place of the 636 MB packed FP8 copy unless the dense FP8 mode's verify head needs that)
-#: and run its decode rows -- up to `_FP8_GEMV_HEAD_ROWS` -- on `nvfp4_gemv` (GB300 micro,
-#: 248320 x 2560: 54-57 us at 1-8 rows and 66-67 at 9-16, against fp8_gemv's 87-90 and 96-98);
-#: more rows keep the per-row FP8 copy through cuBLASLt. The verify head is untouched: outputs
-#: are unchanged and the acceptance length is the cost to watch. 0 (default): `fp8_gemv`.
-_MTP_HEAD_NVFP4_ENV = "TRTLLM_MODELING_V2_MTP_HEAD_NVFP4"
-
-
-def _mtp_head_nvfp4_enabled() -> bool:
- mode = os.environ.get(_MTP_HEAD_NVFP4_ENV, "0")
- assert mode in ("0", "1"), f"{_MTP_HEAD_NVFP4_ENV}={mode!r}: expected 0 or 1"
- return mode == "1"
-
-
-#: Set to 1 to give each GDN layer an NVFP4 copy of its input projection (e2m1 values, one e4m3
-#: scale per 16 hidden elements of a row, an fp32 global scale; `nvfp4_gemv.pack_split`, 23.7 MB
-#: a layer, built at load from the bf16 weight) and run its decode rows -- up to
-#: `_INPROJ_NVFP4_MAX_ROWS` -- on `fp8_gemv`'s kernel over it, the conv step as the GEMV's
-#: epilogue in a captured step (single-token and verify). Prefill, mixed steps and larger
-#: batches keep their entries. Lossy (gsm8k-gated); 0 (default): off.
-_GDN_INPROJ_NVFP4_ENV = "TRTLLM_MODELING_V2_GDN_INPROJ_NVFP4"
-
-#: Rows up to which a decode step's GDN in_proj reads the NVFP4 copy (`_GDN_INPROJ_NVFP4_ENV`).
-#: GB300 micro (us, conv included, behind a PDL predecessor): 6.4 / 6.1 at 1 / 4 rows against
-#: fp8_gemv_conv1d's 7.9 / 7.7; alone 7.1 against 7.9 at 8 rows, 9.1 against 8.6 at 16.
-_INPROJ_NVFP4_MAX_ROWS = 8 # capacity: fixed: a dispatch threshold, at most fp8_gemv's MAX_ROWS
-
-
-def _gdn_inproj_nvfp4_enabled() -> bool:
- mode = os.environ.get(_GDN_INPROJ_NVFP4_ENV, "0")
- assert mode in ("0", "1"), f"{_GDN_INPROJ_NVFP4_ENV}={mode!r}: expected 0 or 1"
- return mode == "1"
-
-
-#: Set to 1 to run the GDN input / output and attention input / output projections in FP8
-#: (`fp8_per_tensor_mm`: e4m3 weights with a per-tensor scale, the activation quantized per
-#: call with a dynamic per-tensor scale) wherever a call reaches `_DENSE_FP8_MIN_ROWS` rows. A
-#: lossy mode beside the checkpoint's bf16 projections (gsm8k-gated); off by default.
-_DENSE_FP8_ENV = "TRTLLM_MODELING_V2_DENSE_FP8"
-
-#: Rows from which a projection takes the FP8 entry, per projection kind: the input
-#: projections from 2 rows (one row stays on the bf16 GEMV / fused launches, whose PDL
-#: overlap the two quantization launches would break), the K=6144 output projections only
-#: where the GEMM saving outgrows the activation quantization's two launches (GB300
-#: micro-bench: cuBLASLt bf16 8-10 us vs FP8 6-7 us + ~4 us below ~512 rows).
-_DENSE_FP8_MIN_ROWS = {"in": 2, "out": 512} # capacity: fixed: dispatch thresholds
-
-
-def _dense_fp8_enabled() -> bool:
- mode = os.environ.get(_DENSE_FP8_ENV, "0")
- assert mode in ("0", "1"), f"{_DENSE_FP8_ENV}={mode!r}: expected 0 or 1"
- return mode == "1"
-
-
-#: Row ranges `(lo, hi)` in which a projection with a packed FP8 decode copy runs `fp8_gemv`
-#: (FP8 weights, bf16 activations, one PDL launch), per projection kind, every kind from one row:
-#: the input projections are 1.6-1.8x the bf16 GEMV through 16 rows; at one row the output and
-#: shared-expert projections on fp8_gemv serve STP concurrency 1 1.5% faster than on
-#: `bf16_gemv_pdl` (two same-node pairs), from two rows the bf16 entry is the slower low-M GEMM
-#: (micro: out/o 5.8 vs 4.3 us, gate|up 4.4 vs 2.1, down 3.2 vs 1.5 at 4 rows). The upper ends are each
-#: kind's crossover (GB300 micro, us, fp8_gemv against the next entry): in/attention-in 8.7 vs
-#: ~9 `mid_rows_mm` at 16 rows, 10.9 vs 9.1 at 17; out/o 6.4 vs 8.4 bf16 at 16 rows, 9.3 vs 8.8
-#: at 20; gate|up 2.6 vs 4.0 at 16, 4.6 vs 4.1 at 20; down 1.7 vs 2.9 at 8, 2.3 vs 2.0 at 12.
-_FP8_GEMV_ROWS = {"in": (1, 16), "out": (1, 16), "sh_gu": (1, 16), "sh_dn": (1, 8)} # dispatch
-
-#: Rows up to which a verify step's conv runs in the in_proj GEMV's epilogue
-#: (`fp8_gemv_conv1d_verify`, certified at 2..16 rows); more rows take the Triton conv.
-_VERIFY_CONV_FUSED_MAX_ROWS = 16 # capacity: fixed: the entry's certified rows
-
-#: Rows up to which a captured single-token step runs the conv step as `fp8_gemv_conv1d`'s
-#: epilogue over the packed FP8 in_proj copy (dense FP8 mode).
-_FP8_CONV_FUSED_MAX_ROWS = 8 # capacity: fixed: a dispatch threshold
-
-
-#: Rows up to which the vocabulary projection (the draft head; under the dense FP8 mode the
-#: verify head too) runs `fp8_gemv` on the packed copy; more rows take cuBLASLt over the
-#: per-row-scaled copy. GB300, 248320 x 2560: fp8_gemv 86-98 us at 1-16 rows against 107-134,
-#: then 142 at 20 rows and 172 at 32 against a flat ~108.
-_FP8_GEMV_HEAD_ROWS = 16 # capacity: fixed: a dispatch threshold, at most fp8_gemv's MAX_ROWS
-
-#: Rows up to which the FP8 vocabulary projection's rows past `_FP8_GEMV_HEAD_ROWS` run
-#: `mid_rows_mm` over the per-row-scaled copy (activation quantized per row); more rows keep
-#: cuBLASLt. GB300, 248320 x 2560, us at 17 / 32 / 64 / 128 rows: 99.7 / 101.5 / 106.0 / 125.8
-#: against 110.0 / 111.9 / 107.7 / 137.5; at 256 rows 158 against 150.
-_MID_HEAD_ROWS = 128 # capacity: fixed: a dispatch threshold, at most mid_rows_mm's MAX_ROWS
-
-
-#: Row ranges `(lo, hi)` in which a decode projection runs `mid_rows_mm` (tcgen05, the weight
-#: streamed before the PDL wait; at most its MAX_ROWS): the dense FP8 projections over their
-#: per-tensor e4m3 codes with the activation quantized per row (`fp8_row_quant`, one launch),
-#: the shared expert and the router over their bf16 weights. Each range is where it beat the
-#: kind's previous entry in the GB300 micro-bench (CUDA-graph chain over HBM-cold weights,
-#: quantization launch included; us at 17 / 32 / 64 / 128 / 256 rows): GDN in_proj 9.1 / 9.3
-#: / 9.9 / 11.4 / 13.8 against 10.9 (fp8_gemv) / 12.4 / 12.8 / 13.7 / 15.4 (per-tensor FP8),
-#: attention in-proj 8.9 / 9.1 / 9.2 / 10.3 / 13.4 against 9.7 / 12.0 / 12.4 / 13.2 / 14.0,
-#: out/o 5.6 / 6.1 / 7.0 / 8.7 against 6.7 / 6.7 / 9.3 / 9.5 bf16 (10.1 vs 12.5 at 256), the
-#: router 2.3 / 2.5 / 2.9 against 2.4 / 2.6 / 4.8 (equal at 128), gate|up 2.9 / 3.3 / 3.7
-#: against 4.4 / 4.5 / 4.5 (behind from 128); the shared expert's down projection never won by
-#: more than noise. An empty range `(0, -1)` keeps the kind off it.
-_MID_MM_ROWS = {
- "in": (17, 256),
- "out": (17, 128),
- "sh_gu": (17, 64),
- "sh_dn": (0, -1),
- "router": (17, 64),
-} # capacity: fixed: dispatch thresholds
-
-
-#: Decode rows from which `DecodeTarget._moe_block` launches the router before the shared
-#: expert's side-stream work: the router heads the routed chain (router, routing, FC1, FC2),
-#: and launched after the side stream's gate|up GEMM its CTAs waited for SMs behind that
-#: GEMM's (STP con256 nsys: a router call 9.2 us launched second, 5.8 launched first). GB300
-#: micro-bench (`DecodeTarget`'s MoE block, 24 layers of cold weights in one graph), us a
-#: block, shared expert first vs router first: 2 rows 21.7 vs 22.4, 4 rows 30.8 vs 30.9, 8
-#: rows 46.5 vs 45.7, 16 rows 76.0 vs 74.6, 32 rows 112.9 vs 112.1, 128 rows 206.7 vs 204.4,
-#: 256 rows 275.6 vs 271.7, 1024 rows 366.6 vs 364.7.
-_ROUTER_FIRST_MIN_ROWS = 8 # capacity: fixed: a dispatch threshold, not a buffer size
-
-
-def _mid_rows(kind: str, rows: int) -> bool:
- lo, hi = _MID_MM_ROWS[kind]
- return lo <= rows <= hi
-
-
-def _gemv_ready(x: torch.Tensor) -> bool:
- """`fp8_gemv`'s activation layout: unit column stride, rows 16-byte aligned."""
- return x.stride(1) == 1 and x.stride(0) % 8 == 0 and x.data_ptr() % 16 == 0
-
-
-def _dense_mm(
- x: torch.Tensor,
- weight_t: torch.Tensor,
- fp8,
- gemm,
- kind: str,
- packed=None,
- mid=None,
-) -> torch.Tensor:
- """A dense projection under the dense FP8 mode: `fp8_gemv` over the packed decode copy
- (`packed`) within `_FP8_GEMV_ROWS[kind]` rows, `mid_rows_mm` over `mid` (a `MidWeight`)
- within `_MID_MM_ROWS[kind]` rows, the per-tensor FP8 GEMM (`fp8` = `(codes_t, scale)`) from
- `_DENSE_FP8_MIN_ROWS[kind]` rows, else the step's bf16 entry `gemm` -- which is every call
- with the mode off (all copies None)."""
- rows = x.shape[0]
- lo, hi = _FP8_GEMV_ROWS.get(kind, (1, 0))
- if packed is not None and lo <= rows <= hi and _gemv_ready(x):
- return fp8_gemv_op(x, packed)
- if mid is not None and _mid_rows(kind, rows) and (not mid.fp8 or _gemv_ready(x)):
- return mid_rows_mm_op(x, mid)
- if fp8 is not None and rows >= _DENSE_FP8_MIN_ROWS[kind]:
- return fp8_per_tensor_mm(x, *fp8)
- return gemm(x, weight_t)
-
-
-#: `_mid` weight-name patterns -> the `_MID_MM_ROWS` kind their projections dispatch as.
-_MID_KINDS = (
- (r"l\d+_(attn_in_w|in_proj_w)", "in"),
- (r"l\d+_(o_w|out_w)", "out"),
- (r"l\d+_sh_gu", "sh_gu"),
- (r"l\d+_sh_dn", "sh_dn"),
- (r"l\d+_router", "router"),
-)
-
-
-def _precompile_mid(mid: dict) -> None:
- """Compile every `mid_rows_mm` variant the decode dispatch can reach now, once per weight
- geometry: a variant first met by a serving step would stall it for the compile."""
- seen = {}
- for name, mw in mid.items():
- kind = next(k for rx, k in _MID_KINDS if re.fullmatch(rx, name))
- lo, hi = _MID_MM_ROWS[kind]
- key = (mw.n, mw.k, mw.fp8, None if mw.scale is None else mw.scale.numel())
- if lo <= hi and key not in seen:
- seen[key] = mid_rows_precompile(mw, max(lo, 1), hi)
-
-
-#: Decode batches up to this many rows run each trunk hyper-connection site as
-#: `qwen4_exp_hc_decode_first_cute` + `qwen4_exp_hc_decode_mix_cute` (or its NVFP4
-#: variant at an MoE site; two GEMV launches); larger ones as `qwen4_exp_hc_inject_norm`
-#: and the mixer. A dispatch threshold from the GB300 micro-bench (the GEMVs win at 1
-#: row and lose from 2), at most the entries' MAX_ROWS.
-_HC_DECODE_GEMV_MAX_ROWS = 1 # capacity: fixed: dispatch threshold, not a buffer bound
-
-#: Decode row counts up to which a hyper-connection mixer (so from
-#: `_HC_DECODE_GEMV_MAX_ROWS + 1` rows) runs its first half (projection and
-#: activations) as the one tcgen05 launch `qwen4_exp_hc_decode_tc_first` instead of
-#: cuBLAS split-K (+ its reduce) + `qwen4_exp_hc_lowrank_act`: the entry's MAX_ROWS
-#: (two 128-row MMA tiles); larger batches keep cuBLAS. GB300 micro-bench (96 sites,
-#: cold weights), us a site for the whole site at 256 rows: 17.8 against 22.4.
-_HC_TC_FIRST_MAX_ROWS = (
- 256 # capacity: fixed: dispatch threshold (two MMA tiles), not a buffer bound
-)
-
-#: Prefill steps of at least this many rows run each trunk hyper-connection site on the
-#: un-normalized residual: `qwen4_exp_hc_inject_rstd` stores `hc_out` and a `[T, C]` rstd
-#: instead of the `[T, C*H]` normed stream, and the mixer
-#: (`qwen4_exp_hc_first_rescale_tc` + `qwen4_exp_hc_up_gate_mean_rstd(_fp4)`) forms
-#: `normed` on chip. Smaller steps keep `qwen4_exp_hc_inject_norm` and the stored-normed
-#: mixer. GB300 micro-bench, us a site (stored normed vs rstd): 2048 rows 57.0 vs 56.7,
-#: 3072 rows 82.4 vs 76.1, 4096 rows 104.0 vs 90.8, 8192 rows 186.2 vs 164.9, 16384 rows
-#: 360.1 vs 293.7.
-_HC_RSTD_MIN_ROWS = 3072 # capacity: fixed: a dispatch threshold, not a buffer size
-
-#: FP8 block-scale block edge, on both weight-scale axes and the activation row.
-_FP8_BLOCK = _weights.FP8_BLOCK
-#: Values per e4m3 block scale in an NVFP4 KV pool.
-_NVFP4_BLOCK = 16
-
-#: False: plain `amax / 448` scale (the runner's recipe). True: power-of-two
-#: (UE8M0), a different quantization of the same input.
-_MTP_MOE_USE_UE8M0 = False
-
-#: Set to 1 to run the draft layer's routed experts in NVFP4 on the trunk's trtllm-gen runner:
-#: at load the checkpoint's FP8 128x128 block-scaled experts are requantized (per-expert weight
-#: global scales, 16-element e4m3 block scales) into the trunk's layout and the FP8 copies are
-#: freed. The draft only proposes tokens -- the verify step decides what is emitted -- so this
-#: can change the acceptance length, not outputs. 0 (default): the FP8 block-scale runner.
-_MTP_MOE_NVFP4_ENV = "TRTLLM_MODELING_V2_MTP_MOE_NVFP4"
-
-#: The NVFP4 draft MoE's two activation global scales are static, `448 * 6 / A` with `A` this
-#: factor times the largest calibrated amax of the trunk's routed experts (its `input_scale`s,
-#: FC1 input and FC2 input respectively; the draft layer carries none). A global scale only has
-#: to keep the e4m3 block scales inside their range, whose 3-bit precision is the same from
-#: 2^-6 to 448: a block whose amax exceeds `A` is clipped, one below `A / 28672` loses scale
-#: precision, so the margin costs nothing until blocks are ~1e-3 of `A`.
-_MTP_MOE_NVFP4_FC1_MARGIN = 8.0
-_MTP_MOE_NVFP4_FC2_MARGIN = 4.0
-
-
-def _mtp_moe_nvfp4_enabled() -> bool:
- mode = os.environ.get(_MTP_MOE_NVFP4_ENV, "0")
- assert mode in ("0", "1"), f"{_MTP_MOE_NVFP4_ENV}={mode!r}: expected 0 or 1"
- return mode == "1"
-
-
-#: The history depth of the GDN cached-replay verify path (the manager's
-#: `replay_history_size`). The replay entry is certified at 16 only, and a manager that
-#: reports another depth is refused at bind.
-_GDN_REPLAY_HISTORY = 16
-
-#: The most query rows a verify step carries per generation request, `max_draft_len + 1`.
-#: It reaches `thop_attention` as `predicted_tokens_per_seq` and the cached replay as its step
-#: width `T`; the replay kernel asserts `T <= 8`, and above it the engine would verify with the
-#: legacy per-step states, which this target refuses past `_MAX_LEGACY_VERIFY_WIDTH` rows.
-_MAX_VERIFY_WIDTH = 8
-
-#: The same bound on the legacy verify path (`TRTLLM_USE_GDN_REPLAY=0`): the entries only it
-#: calls (`pack_gdn_decode_qkv`, `fused_gdn_gating`, the multi-token
-#: `fused_recurrent_gated_delta_rule_update`) are certified at up to 4 rows per request, and its
-#: per-step recurrent buffer holds one state per verify row, request and layer (6.75 GiB at 8
-#: rows and 8 requests in fp32).
-_MAX_LEGACY_VERIFY_WIDTH = 4
-
-#: The recurrent-state pool dtypes this target serves (`kv_cache_config.mamba_ssm_cache_dtype`):
-#: the checkpoint's fp32, or bf16, which rounds the state to bf16 once per decode token, once
-#: per prefill call and at every cached-replay checkpoint commit. The engine also accepts fp16.
-_SSM_STATE_DTYPES = (torch.float32, torch.bfloat16)
-
-
-# ----------------------------------------------- KV cache manager binding ---
-#
-# The ops are plain functions over raw tensors, and the pool views taken from the
-# KV cache manager are `TensorWrapper`s over raw device addresses that do not keep
-# memory alive, so a view is only valid for the manager that produced it. `_Bound`
-# is everything derived from one manager; the core replaces it as a unit
-# (`_bind_cache_manager`) and the hot path pays one identity comparison. The MTP
-# draft layer's step can see a second manager (the separate draft KV cache), bound
-# the same way into `_DraftBound` (`_bind_draft_cache_manager`).
-#
-# The `check_*` functions hold the shape/dtype contract. They read only
-# `.shape` / `.dtype`, so unit tests run them on meta tensors with no engine.
-#
-# For a later target with pools that are not in `attn_metadata` (mamba/GDN state,
-# side caches): the reusable parts are (1) one immutable bound object per manager,
-# replaced atomically and keyed on manager identity, (2) pure contract checks,
-# (3) refusing a bind during CUDA graph capture that no eager step validated first
-# (a re-bind for the trunk's, a first bind for the draft layer's). Once a second
-# target needs it, the next step is a declarative
-# `PoolSpec(name, layers, getter, tail_shape, dtype)` list consumed by one shared
-# binder under `modeling_v2/` (not `catalog/`), leaving each target only its specs
-# and shape constants. Pure-attention targets that read the pool pointers from
-# `attn_metadata` on every call (deepseek_v3) do not need this.
-
-
-class _PoolGeometry(NamedTuple):
- """The pool shapes this target expects, taken from its config and the metadata."""
-
- tokens_per_block: int
- heads_kv: int
- head_dim: int
- index_kv_heads: int
- index_head_dim: int
- conv_dim: int
- conv_kernel: int
- num_k_heads: int
- num_v_heads: int
- head_v_dim: int
- head_k_dim: int
- stream: int
- ple_state_len: int
- ngram_context_len: int
- dtype: torch.dtype # the activation dtype
- kv_dtype: torch.dtype # the paged K/V pool's dtype: bf16, e4m3 (fp8) or uint8 (NVFP4)
- kv_width: int # its innermost extent: head_dim, or head_dim // 2 for NVFP4's e2m1 pairs
-
-
-class _KvPlanes(NamedTuple):
- """One layer's paged K/V views, and under an NVFP4 pool its e4m3 block-scale views."""
-
- k: torch.Tensor
- v: torch.Tensor
- k_scale: torch.Tensor | None = None
- v_scale: torch.Tensor | None = None
-
-
-class _Replay(NamedTuple):
- """The GDN cached-replay verify state of one manager: per slot, the recurrent pool
- row is a checkpoint and a compact history holds the tokens accepted since it."""
-
- old: dict # GDN layer -> (old_x, old_B, old_dt, old_dA_cumsum) history buffers
- cache_buf_idx: torch.Tensor # per slot: which history half is pending
- prev_num_accepted: torch.Tensor # per slot: pending history entries
- history: int
- width: int # verify rows per generation request
- all_layer_commit: bool # the manager commits partitioned steps over every layer
- cute: bool # bf16 pool: the CuTe replay kernel (the Triton kernel's commit contract)
-
-
-class _Bound(NamedTuple):
- """Everything derived from one KV cache manager."""
-
- cm: object
- pool_idx: dict # attention layer -> pool offset (what thop_attention addresses)
- index_cache: dict # attention layer -> QSA index plane
- position_cache: torch.Tensor # shared QSA position plane
- kv: dict # attention layer -> `_KvPlanes`
- conv: dict # GDN layer -> conv state pool
- ssm: dict # GDN layer -> recurrent state pool
- ple_state: tuple # (PLE conv state, PLE n-gram id history)
- # MTP only: the verify step's per-step conv buffers, and one recurrent mode. Legacy:
- # per-step recurrent buffers (`int_ssm`). Cached replay: `replay`, no `int_ssm`.
- int_idx: torch.Tensor | None = None
- int_conv: dict | None = None
- int_ssm: dict | None = None
- replay: _Replay | None = None
-
-
-class _DraftBound(NamedTuple):
- """The MTP draft layer's views of the manager its draft step sees: the separate
- draft manager when the engine built one, else the trunk's."""
-
- cm: object
- index_cache: torch.Tensor # the draft layer's QSA index plane
- position_cache: torch.Tensor # that manager's QSA position plane
- kv: _KvPlanes # the draft layer's paged K/V views
- share_slot: int | None # its row in the metadata's shared selection buffers (IndexShare)
-
-
-def _require(cm, names, what: str) -> None:
- for name in names:
- assert hasattr(cm, name), (
- f"the KV cache manager {type(cm).__name__} does not expose {name}; this target {what}"
- )
-
-
-def _check_pool(t, what, lead, tail, dtype=None) -> None:
- """`t` has `lead` leading dims of any size, then exactly `tail`, at `dtype`."""
- assert t is not None, f"{what}: the runtime did not allocate it"
- expected = ["*"] * lead + [str(n) for n in tail]
- assert t.dim() == lead + len(tail) and tuple(t.shape[lead:]) == tuple(tail), (
- f"{what} is {tuple(t.shape)}, expected [{', '.join(expected)}]"
- )
- assert dtype is None or t.dtype == dtype, f"{what} is {t.dtype}, expected {dtype}"
-
-
-def check_position_plane(g: _PoolGeometry, position) -> None:
- shape = (g.tokens_per_block, _QSA_POSITION_AXES)
- _check_pool(position, "QSA position plane", 1, shape, torch.int32)
-
-
-def check_attention_pools(g: _PoolGeometry, layer: int, index_cache, pool) -> None:
- """QSA index plane and paged K/V for one attention layer."""
- index_shape = (g.tokens_per_block, g.index_kv_heads, g.index_head_dim)
- _check_pool(index_cache, f"layer {layer}: QSA index plane", 1, index_shape, g.dtype)
- # The sparse path needs adjacent HND K/V role planes: [pages, 2, H, block, D].
- kv_shape = (g.heads_kv, g.tokens_per_block, g.kv_width)
- _check_pool(pool, f"layer {layer}: paged K/V", 2, kv_shape, g.kv_dtype)
- assert pool.shape[1] == 2, f"layer {layer}: paged K/V has {pool.shape[1]} role planes, not 2"
-
-
-def check_kv_scale_planes(g: _PoolGeometry, layer: int, pool, scales) -> None:
- """An NVFP4 layer's e4m3 block scales: one uint8 per `_NVFP4_BLOCK` values, laid out like
- the data pages (`[slots, 2, H, block, D / 16]`), and holding as many pages per slot as the
- data pool does.
-
- The attention op addresses a layer's data pages and its scale pages with one block-offset
- table, so a scale pool whose pages-per-slot differs from the data pool's reads and writes
- other slots' scales without any error. Both counts are read off the views' slot strides."""
- scale_shape = (g.heads_kv, g.tokens_per_block, g.head_dim // _NVFP4_BLOCK)
- _check_pool(scales, f"layer {layer}: NVFP4 block scales", 2, scale_shape, torch.uint8)
- assert scales.shape[1] == 2 and scales.shape[0] == pool.shape[0], (
- f"layer {layer}: the block scales are {tuple(scales.shape)} against paged K/V "
- f"{tuple(pool.shape)}; they must hold the same slots, K then V"
- )
- pages = []
- for what, t in (("paged K/V", pool), ("block scales", scales)):
- page = t[0, 0].numel()
- assert t.stride(0) % page == 0, (
- f"layer {layer}: the {what} slot stride {t.stride(0)} is not a whole number of "
- f"{page}-element pages"
- )
- pages.append(t.stride(0) // page)
- assert pages[0] == pages[1], (
- f"layer {layer}: the paged K/V pool holds {pages[0]} pages per slot and the block-scale "
- f"pool {pages[1]}; the attention op would address the scales of other slots"
- )
-
-
-def check_recurrent_pools(g: _PoolGeometry, conv, ssm) -> None:
- _check_pool(conv, "conv state pool", 1, (g.conv_dim, g.conv_kernel - 1), g.dtype)
- # Key dim innermost.
- _check_pool(ssm, "recurrent state pool", 1, (g.num_v_heads, g.head_v_dim, g.head_k_dim))
- assert ssm.dtype in _SSM_STATE_DTYPES, (
- f"the recurrent state pool is {ssm.dtype}; this target serves "
- f"{' or '.join(str(d) for d in _SSM_STATE_DTYPES)} (kv_cache_config.mamba_ssm_cache_dtype)"
- )
-
-
-def check_ple_pools(g: _PoolGeometry, ple_conv, ple_ngram, conv) -> None:
- _check_pool(ple_conv, "PLE conv state", 1, (g.stream, g.ple_state_len), g.dtype)
- # The id-history entry writes int64 into it; an int32 pool raises.
- _check_pool(ple_ngram, "PLE n-gram context pool", 1, (g.ngram_context_len,), torch.int64)
- # The forward addresses all three with one scheduler slot vector.
- assert ple_conv.shape[0] == conv.shape[0] == ple_ngram.shape[0], (
- "the PLE and Gated-DeltaNet pools must share one slot space"
- )
-
-
-def check_step_buffers(g: _PoolGeometry, int_conv, conv, int_idx) -> None:
- """MTP verify buffers both recurrent modes use: the per-step conv windows
- `[slots, steps, ...]` at the pool's dtype (the promotion copies between them without
- converting) and the rows `update_mamba_states` reads its sources from."""
- assert int_conv is not None, (
- "the runtime allocated no per-step conv buffer; a partially accepted draft "
- "cannot be rolled back"
- )
- _check_pool(int_conv, "per-step conv buffer", 2, (g.conv_dim, g.conv_kernel - 1), conv.dtype)
- assert int_idx is not None and int_idx.dtype == torch.int32, (
- "the cache manager exposes no int32 intermediate_state_indices"
- )
-
-
-def check_legacy_state(g: _PoolGeometry, int_ssm, ssm) -> None:
- """Legacy verify: per-step recurrent buffer `[slots, steps, ...]` at the pool's dtype."""
- assert int_ssm is not None, (
- "the cache manager runs the legacy verify path but allocated no per-step recurrent "
- "buffer; a partially accepted draft cannot be rolled back"
- )
- ssm_shape = (g.num_v_heads, g.head_v_dim, g.head_k_dim)
- _check_pool(int_ssm, "per-step recurrent buffer", 2, ssm_shape, ssm.dtype)
- assert int_ssm.shape[1] <= _MAX_LEGACY_VERIFY_WIDTH, (
- f"the legacy verify path runs {int_ssm.shape[1]} rows per generation request; its entries "
- f"are certified at up to {_MAX_LEGACY_VERIFY_WIDTH} (max_draft_len "
- f"{_MAX_LEGACY_VERIFY_WIDTH - 1}). Leave TRTLLM_USE_GDN_REPLAY unset (cached replay) for "
- "a longer draft"
- )
-
-
-def check_replay_state(history, width, prev_accepted, cache_buf_idx, ssm) -> None:
- """Cached replay, manager-wide: the geometry the replay entry is certified at, and the
- per-slot bookkeeping, which must cover every pool slot."""
- assert ssm.dtype in _SSM_STATE_DTYPES, (
- f"the recurrent pool is {ssm.dtype}; the replay entry is certified at an fp32 or a "
- "bf16 checkpoint"
- )
- assert history == _GDN_REPLAY_HISTORY, (
- f"replay_history_size {history}; the replay entry is certified at {_GDN_REPLAY_HISTORY}"
- )
- assert 2 <= width <= min(_MAX_VERIFY_WIDTH, history), (
- f"replay_step_width {width}; the replay entry is certified at 2..{_MAX_VERIFY_WIDTH} "
- f"verify rows and never above the history ({history})"
- )
- for name, t in (("prev_num_accepted_tokens", prev_accepted), ("cache_buf_idx", cache_buf_idx)):
- assert t.dim() == 1 and t.shape[0] >= ssm.shape[0] and t.dtype == torch.int32, (
- f"{name} is {tuple(t.shape)} {t.dtype}; the kernel reads one int32 per pool slot "
- f"({ssm.shape[0]})"
- )
-
-
-def check_replay_history(
- g: _PoolGeometry, layer: int, history, pool_slots, old_x, old_B, old_dt
-) -> None:
- """One GDN layer's replay history: `[slots, 2 halves, ...]`, addressed by `history` and
- the head geometry, never by the buffer's own shape."""
- for name, t in (("old_x", old_x), ("old_B", old_B), ("old_dt", old_dt)):
- assert t is not None, f"layer {layer}: the manager's layer cache carries no {name}"
- assert old_x.shape[0] >= pool_slots, (
- f"layer {layer}: the replay history has {old_x.shape[0]} slots for a {pool_slots}-slot pool"
- )
- bf16, fp32 = torch.bfloat16, torch.float32
- x_tail, b_tail = (history, g.num_v_heads, g.head_v_dim), (history, g.num_k_heads, g.head_k_dim)
- _check_pool(old_x, f"layer {layer}: old_x", 2, x_tail, bf16)
- _check_pool(old_B, f"layer {layer}: old_B", 2, b_tail, bf16)
- _check_pool(old_dt, f"layer {layer}: old_dt", 2, (g.num_v_heads, history), fp32)
- assert old_x.shape[1] == old_B.shape[1] == old_dt.shape[1] == 2, (
- f"layer {layer}: the replay history is not double-buffered"
- )
-
-
-def check_verify_width(rows: int) -> None:
- """The rows per generation request one step carries: `max_draft_len + 1` (`runtime_draft_len
- + 1` at run time), a `thop_attention` `predicted_tokens_per_seq` the entry is certified at."""
- assert 1 <= rows <= _MAX_VERIFY_WIDTH, (
- f"predicted_tokens_per_seq {rows} is outside the 1..{_MAX_VERIFY_WIDTH} range "
- "attention/thop_attention is certified at"
- )
-
-
-def check_replay_width(replay, rows: int) -> None:
- """Under the cached replay a single-row generation step cannot take the ordinary
- recurrence: the pool row is a checkpoint whose pending history only the replay kernel
- applies. The manager's step width is the one row count per sequence a step may carry."""
- assert replay is None or rows == replay.width, (
- f"the batch carries {rows} query rows per generation sequence and the GDN cached "
- f"replay runs at step width {replay.width}; a runtime_draft_len of {rows - 1} (a "
- "draft-length schedule, or skipped drafting) cannot run under it. Export "
- "TRTLLM_USE_GDN_REPLAY=0 for variable draft lengths"
- )
-
-
-def check_shared_selection(shared, slots, layer: int, block_topk: int) -> None:
- """IndexShare: the metadata's shared selection buffers have a row for `layer` and are
- `block_topk` wide."""
- assert shared is not None, "IndexShare is on but the QSA metadata allocated no shared buffers"
- assert layer in slots and slots[layer] < shared.shape[0], (
- f"the KV cache manager gives layer {layer} no row of the {shared.shape[0]}-row shared "
- f"selection buffers (qsa_shared_topk_slots={slots})"
- )
- assert shared.shape[2] == block_topk, (
- f"the shared selection rows are {shared.shape[2]} wide and this target selects "
- f"{block_topk} blocks per row"
- )
-
-
-def _moe_layer_scalars(w, i: int, backend: str) -> dict[str, torch.Tensor]:
- """Layer `i`'s routed-expert scalars for `backend`, built from the checkpoint's
- `input_scale` / `weight_scale_2` rows.
-
- The checkpoint stores modelopt reciprocals (`input_scale = 1 / g_act`,
- `weight_scale_2 = 1 / g_w`), so a quantizer's global scale is `1 / input_scale` and
- the runner's scalars are products of the two.
- """
- isc1, isc1_up = w[f"l{i}_e_isc1"], w[f"l{i}_e_isc1_up"]
- assert torch.equal(isc1, isc1_up), (
- f"layer {i}: expert gate/up input_scale differ; the fused FC1 "
- "reads one activation quantization for both halves"
- )
- assert torch.equal(w[f"l{i}_e_ws2_1"], w[f"l{i}_e_ws2_1_up"]), (
- f"layer {i}: expert gate/up weight_scale_2 differ; the fused "
- "FC1 carries one alpha per expert for both halves"
- )
- # One activation quantization feeds all 512 experts, so FC1's activation
- # global scale must be one number (the shared expert is bf16, no scale).
- assert torch.equal(isc1.min(), isc1.max()), (
- f"layer {i}: the routed experts do not share one "
- "gate_proj.input_scale, so a single activation quantization "
- "cannot feed them all and this assembly is wrong for this "
- "checkpoint"
- )
- e_isc2, ws2_2 = w[f"l{i}_e_isc2"], w[f"l{i}_e_ws2_2"]
- gate1 = (isc1 * w[f"l{i}_e_ws2_1"]).contiguous()
- # Routed FC1 activation global scale (one per layer).
- scalars = dict(global_scale=(1.0 / isc1[:1]).contiguous())
- if backend == "CUTEDSL":
- # The CuteDSL FC1 epilogue requantizes with one global scale for the whole layer,
- # where trtllm-gen folds each expert's `down_proj.input_scale` into its FC1
- # scalar. Quantizing at `1 / max(input_scale)` (the smallest per-expert global
- # scale, so no block scale saturates) and putting the same factor into FC2's alpha
- # makes `global_sf * alpha2[e] == ws2_2[e]` for every expert, whatever the spread.
- isc2_max = e_isc2.max()
- scalars.update(
- # [num_experts] each: activation scale times weight scale.
- alpha1=gate1,
- alpha2=(isc2_max * ws2_2).contiguous(),
- # The FC2 activation's quantization scale, numel 1.
- global_sf=(1.0 / isc2_max).reshape(1).contiguous(),
- )
- else:
- scalars.update(
- # [num_experts] each. Swapping the first two is finite and silent.
- output1_scale_scalar=(gate1 / e_isc2).contiguous(),
- output1_scale_gate_scalar=gate1,
- output2_scale_scalar=(e_isc2 * ws2_2).contiguous(),
- )
- return scalars
-
-
-class Qwen4ExpTp1ModelingV2(DecoderModel):
- def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig):
- """`composite_config` is passed in because the text sub-config adopted as
- `model_config.pretrained_config` carries neither the vision half nor the
- media token ids."""
- super().__init__(model_config)
- cfg = model_config.pretrained_config
- assert cfg is not None
-
- # This target IS this geometry — assert, never adapt.
- assert torch.cuda.get_device_capability() == _SM, (
- f"target certified on sm_{_SM[0]}{_SM[1]}, running on "
- f"sm_{''.join(map(str, torch.cuda.get_device_capability()))}"
- )
- assert model_config.mapping.tp_size == 1, "tp1 target"
- assert model_config.mapping.pp_size == 1, "tp1 target"
- assert not model_config.mapping.enable_attention_dp, "tp1 target"
-
- # `cfg` is the engine-normalized Qwen4ExpTextConfig (pool geometry is read
- # from it).
- assert getattr(cfg, "model_type", None) == "qwen4_exp_text", (
- f"the core must be built on the normalized text config; got "
- f"model_type {getattr(cfg, 'model_type', None)!r}"
- )
-
- dt = model_config.torch_dtype
- assert dt == torch.bfloat16, f"bf16 target, engine resolved {dt}"
- assert cfg.torch_dtype == torch.bfloat16, cfg.torch_dtype
- assert not cfg.tie_word_embeddings, "untied lm_head"
- assert not cfg.attention_bias, "no q/k/v/o bias anywhere"
- assert cfg.hidden_act == "silu", "SwiGLU over a silu gate"
-
- self.num_layers = cfg.num_hidden_layers
- self.hidden = cfg.hidden_size
- self.vocab = cfg.vocab_size
- self.eps = cfg.rms_norm_eps
-
- # `layer_types` is authoritative; `full_attention_interval` is the runtime
- # helper's fallback, so both must agree.
- layer_types = list(cfg.layer_types)
- assert len(layer_types) == self.num_layers, (len(layer_types), self.num_layers)
- assert set(layer_types) <= {"full_attention", "linear_attention"}, set(layer_types)
- self.is_attn = [t == "full_attention" for t in layer_types]
- interval = cfg.full_attention_interval
- assert self.is_attn == [i % interval == interval - 1 for i in range(self.num_layers)], (
- "layer_types disagrees with full_attention_interval"
- )
- self.attn_layers = [i for i, a in enumerate(self.is_attn) if a]
- self.gdn_layers = [i for i, a in enumerate(self.is_attn) if not a]
- assert self.attn_layers and self.gdn_layers, "this target is the hybrid assembly"
-
- # The checkpoint picks the speculative mode (`MTP_EAGLE_ONE_MODEL`, one module
- # replayed `max_draft_len` times); the module is built only when the engine
- # asked for speculation.
- spec_config = getattr(model_config, "spec_config", None)
- self.mtp_enabled = spec_config is not None
- # Rows per request of the widest verify step (`max_draft_len + 1`); 1 without MTP.
- self.max_verify_rows = spec_config.max_draft_len + 1 if self.mtp_enabled else 1
- declared = getattr(cfg, "mtp_num_hidden_layers", None)
- assert declared in (None, _MTP_NUM_LAYERS), (
- f"this target implements the one-layer MTP module this checkpoint "
- f"ships; the text config declares mtp_num_hidden_layers {declared}"
- )
- # The module shares the trunk's embedding and the shell's lm_head.
- assert not getattr(cfg, "mtp_use_dedicated_embeddings", False), (
- "this checkpoint's MTP module shares the decoder's embedding "
- "(`mtp_use_dedicated_embeddings: false`) and ships no "
- "`mtp.embed_tokens` tensor; a dedicated-embedding module is a "
- "different assembly"
- )
- # The nested `mtp` block, when kept: one `full_attention` layer at the trunk's
- # rope base.
- mtp_block = getattr(cfg, "mtp", None)
- self.mtp_theta = None
- if mtp_block is not None:
- assert isinstance(mtp_block, dict), (
- f"the text config's `mtp` block is {type(mtp_block).__name__}; "
- "this target reads it as the published mapping"
- )
- assert mtp_block.get("num_hidden_layers", _MTP_NUM_LAYERS) == _MTP_NUM_LAYERS, (
- f"mtp.num_hidden_layers is {mtp_block.get('num_hidden_layers')}"
- )
- assert list(mtp_block.get("layer_types", ["full_attention"])) == ["full_attention"], (
- f"the MTP module's layer_types is {mtp_block.get('layer_types')}; this "
- "target implements the full-attention form, which is the only "
- "one whose weights the checkpoint carries"
- )
- # `null` means the trunk's own output.
- assert mtp_block.get("mtp_use_hidden_state_from_layer") is None, (
- f"the MTP module asks for the hidden state of layer "
- f"{mtp_block.get('mtp_use_hidden_state_from_layer')}; this "
- "target hands it the end of the stack, which is what the "
- "runtime passes and all this file can reach"
- )
- if "rope_theta" in mtp_block:
- self.mtp_theta = float(mtp_block["rope_theta"])
-
- # ---- hyper-connections ------------------------------------------
- self.hc_count = cfg.hc_count
- self.hc_lowrank = cfg.hc_lowrank
- assert self.hc_count >= 2, (
- f"the hyper-connection residual carries hc_count streams; this "
- f"config declares {self.hc_count}, which is not a multi-stream "
- f"residual and is a different assembly"
- )
- self.stream = self.hc_count * self.hidden
- # Columns of the padded first-projection weight (`_fuse_hc_first_projections`)
- # the cuBLAS first projection multiplies: the real `[down; block_inject]` rows
- # rounded up to 32 (352 here). GB300, fp32 out, K 10240: 16384 rows 81.5 us
- # against 91.5 at the padded 384, 8193 rows 41.6 against 51.2, never slower
- # below; the real columns are bitwise the 384-wide product's at every row
- # count measured (1-16384; at 320 columns they are not).
- self.hc_first_cols = -(-(self.hc_lowrank + self.hc_count) // 32) * 32
-
- # ---- full attention ---------------------------------------------
- self.heads_q = cfg.num_attention_heads
- self.heads_kv = cfg.num_key_value_heads
- self.head_dim = cfg.head_dim
- assert self.head_dim in (64, 128, 256), "fused_qk_norm_rope head_dim set"
- assert self.heads_q % self.heads_kv == 0, "thop_attention GQA rule"
- self.q_width = self.heads_q * self.head_dim
- self.kv_width = self.heads_kv * self.head_dim
- self.qkv_width = self.q_width + 2 * self.kv_width
- # The per-head interleaved output gate makes q_proj twice as wide as q. No
- # config flag exists, so it is derived from the weight (asserted in weights.py);
- # the per-head split below would otherwise address the wrong columns.
- assert not hasattr(cfg, "attn_output_gate"), (
- "this config gained an attn_output_gate flag; the gate is derived "
- "from q_proj's width here, and the two must be reconciled"
- )
-
- # ---- RoPE --------------------------------------------------------
- # Partial, neox half-split pairs, no scaling; `rope_theta` is inside
- # `rope_parameters` (mirrored as `rope_scaling`).
- rope = getattr(cfg, "rope_parameters", None) or getattr(cfg, "rope_scaling", None)
- assert isinstance(rope, dict), f"rope parameters must be a dict, got {rope!r}"
- self.theta = float(rope["rope_theta"])
- assert self.theta > 0.0, "rope theta"
- partial = float(rope.get("partial_rotary_factor", cfg.partial_rotary_factor))
- self.rotary_dim = int(self.head_dim * partial)
- assert self.rotary_dim % 2 == 0 and 0 < self.rotary_dim <= self.head_dim, (
- f"rotary_dim {self.rotary_dim} must be even and at most head_dim "
- f"{self.head_dim}; the op enforces only the evenness half"
- )
- # The mRoPE declaration is what makes the engine hand this forward a
- # [3, 1, T] position buffer.
- self.mrope_declared = str(rope.get("type", rope.get("rope_type"))) == "mrope"
- # `fused_qk_norm_rope` implements the interleaved mRoPE layout only (pair `j`
- # reads coordinate 1 when `j % 3 == 1 and j < 3*s1`, 2 when `j % 3 == 2 and
- # j < 3*s2`, else 0); the contiguous variant is a different operator. A text
- # gate cannot see the split.
- assert rope.get("mrope_interleaved") is True, (
- f"this target passes the interleaved mRoPE layout, which is the only "
- f"one attention/fused_qk_norm_rope implements; the config declares "
- f"mrope_interleaved={rope.get('mrope_interleaved')!r}"
- )
- section = list(rope["mrope_section"])
- assert len(section) == _QSA_POSITION_AXES, (
- f"mrope_section must carry one entry per position coordinate; got {section}"
- )
- rotary_pairs = self.rotary_dim // 2
- assert sum(section) == rotary_pairs, (
- f"mrope_section {section} sums to {sum(section)} but rotary_dim "
- f"{self.rotary_dim} makes {rotary_pairs} rotation pairs; the split "
- f"partitions pairs, not dimensions"
- )
- self.mrope_section1 = int(section[1])
- self.mrope_section2 = int(section[2])
- # Row-1 section is load-bearing here; row-2 saturates (10, 11, 12 are
- # bit-identical at 32 pairs). `use_gemma` with `use_mrope` is certified only
- # at the geometry asserted below.
- assert (
- self.heads_q,
- self.heads_kv,
- self.head_dim,
- self.rotary_dim,
- self.theta,
- ) == (24, 2, 256, 64, 1e7), (
- f"the main attention rope runs use_gemma together with use_mrope, "
- f"which attention/fused_qk_norm_rope certifies only at "
- f"(24, 2, 256, 64, 1e7); this config gives "
- f"{(self.heads_q, self.heads_kv, self.head_dim, self.rotary_dim, self.theta)}"
- )
- # No YaRN or other scaling.
- assert "factor" not in rope, (
- f"this target builds a plain (unscaled) rope; rope carries {rope!r}"
- )
- self.max_positions = cfg.max_position_embeddings
-
- # ---- QSA sparse index --------------------------------------------
- self.index_heads = cfg.indexer_n_heads
- self.index_kv_heads = cfg.indexer_kv_heads
- self.index_head_dim = cfg.indexer_head_dim
- self.token_topk = cfg.indexer_budget
- self.compress_ratio = cfg.indexer_compress_ratio
- assert self.index_kv_heads == _QSA_INDEX_KV_HEADS, (
- f"QSA requires exactly one index KV head; this config declares {self.index_kv_heads}"
- )
- assert self.compress_ratio >= 2 and self.token_topk % self.compress_ratio == 0, (
- f"QSA selection geometry: token_topk {self.token_topk} must divide "
- f"by compress_ratio {self.compress_ratio}, which must be at least 2"
- )
- # Fused index kernels: both powers of two, index head twice the rotary width.
- assert self.index_head_dim == self.rotary_dim * _QSA_INDEX_HEAD_TO_ROTARY_RATIO, (
- f"the fused QSA index kernels require index_head_dim "
- f"({self.index_head_dim}) == 2 * rotary_dim ({self.rotary_dim})"
- )
- assert self.index_head_dim & (self.index_head_dim - 1) == 0, "index_head_dim power of two"
- assert self.compress_ratio & (self.compress_ratio - 1) == 0, "compress_ratio power of two"
- self.block_topk = self.token_topk // self.compress_ratio
- self.expanded_topk = self.token_topk + self.compress_ratio - 1
- # Dense-attention cutoff: QSA resolves max(token_topk, seq_len_threshold);
- # checked against the engine's value below because the CUDA-graph runner
- # splits its short and long decode families on it.
- self.dense_threshold = self.token_topk
-
- # ---- the engine's sparse configuration, checked at construction -----
- # The index/position side caches, `qsa_*` metadata and sparse kernels all hang
- # off `sparse_attention_config`; checking here fails at cold start rather than
- # at the first forward.
- sparse_cfg = getattr(model_config, "sparse_attention_config", None)
- assert getattr(sparse_cfg, "algorithm", None) == "qsa", (
- f"this target is a QSA sparse-attention assembly and the engine "
- f"resolved sparse_attention_config={sparse_cfg!r}. QSA is this "
- f"architecture's own model default, so an engine that reaches here "
- f"without it dropped the default on the way in rather than being "
- f"asked for dense attention"
- )
- # The engine's own lowering on this core's config: a mismatch with the side
- # cache sizing would otherwise be a silent mis-index.
- qsa_params = sparse_cfg.to_sparse_params(pretrained_config=cfg)
- assert (
- qsa_params.index_n_heads,
- qsa_params.index_kv_heads,
- qsa_params.index_head_dim,
- qsa_params.token_topk,
- qsa_params.compress_ratio,
- ) == (
- self.index_heads,
- self.index_kv_heads,
- self.index_head_dim,
- self.token_topk,
- self.compress_ratio,
- ), (
- f"the engine lowered the QSA index geometry to {qsa_params!r} and "
- f"this target derived heads {self.index_heads}, kv heads "
- f"{self.index_kv_heads}, head_dim {self.index_head_dim}, token_topk "
- f"{self.token_topk}, compress_ratio {self.compress_ratio}; the side "
- f"caches are sized from the engine's copy and the kernels are "
- f"called with this one"
- )
- assert qsa_params.dense_seq_len_threshold == self.dense_threshold, (
- f"the engine's resolved QSA cutoff is "
- f"{qsa_params.dense_seq_len_threshold} and this forward branches at "
- f"{self.dense_threshold}. The runner captures its short decode "
- f"graph family at the engine's number, so a disagreement freezes "
- f"the dense body into the graph that long batches replay — which "
- f"is finite, plausible output. A serving config that raises "
- f"seq_len_threshold above token_topk is the way to get here, and "
- f"it is a configuration this target does not implement"
- )
- # IndexShare for the MTP draft loop (`sparse_attention_config.
- # index_share_for_mtp_iteration`), as the engine resolved it: the same lowering
- # sizes the metadata's shared selection buffers.
- self.mtp_index_share = bool(qsa_params.mtp_index_share)
- self.index_width = (self.index_heads + self.index_kv_heads) * self.index_head_dim
- # Packed [q | k | v] scratch for the fused norm+rope call; the spare KV head
- # is the v slot the op needs but never reads.
- self.index_scratch_width = self.index_width + self.index_head_dim
- # The fused input projection's output row (`weights._attn_in_w`):
- # [q | k | v | gate | index q, k | v slot]. Its leading columns are the
- # packed QKV the main norm+rope rewrites in place, and its tail is the
- # index scratch's layout, so the row is `qkv_width + q_width +
- # index_scratch_width` wide. The norm+rope addresses a row as
- # `(Hq + Hk + Hv) * head_dim` and never touches the v heads, so the
- # row must be a whole number of heads: everything after k is "v" to it.
- self.attn_in_width = self.qkv_width + self.q_width + self.index_scratch_width
- assert self.attn_in_width % self.head_dim == 0, (
- f"the fused projection row ({self.attn_in_width}) is not a whole number of "
- f"{self.head_dim}-wide heads, so the main norm+rope cannot address it in place"
- )
- self.attn_in_tail_heads = self.attn_in_width // self.head_dim - self.heads_q - self.heads_kv
- self.softmax_scale = 1.0 / math.sqrt(self.head_dim)
-
- # ---- gated deltanet ----------------------------------------------
- self.num_k_heads = cfg.linear_num_key_heads
- self.num_v_heads = cfg.linear_num_value_heads
- self.head_k_dim = cfg.linear_key_head_dim
- self.head_v_dim = cfg.linear_value_head_dim
- self.conv_kernel = cfg.linear_conv_kernel_dim
- assert self.num_v_heads % self.num_k_heads == 0, "GDN key-head broadcast"
- assert self.head_k_dim == 128 and self.head_v_dim == 128, (
- f"the delta-rule entries certify head dims 128/128; this config "
- f"declares {self.head_k_dim}/{self.head_v_dim}"
- )
- assert self.conv_kernel in (2, 3, 4), "causal_conv1d width domain"
- self.key_dim = self.num_k_heads * self.head_k_dim
- self.value_dim = self.num_v_heads * self.head_v_dim
- self.conv_dim = 2 * self.key_dim + self.value_dim
- # Column widths of the one input-projection GEMM's output: [q|k|v | z | a | b].
- self.in_proj_split = (self.conv_dim, self.value_dim, self.num_v_heads, self.num_v_heads)
- # Gate activation of the gated RMS norm (`silu` or `sigmoid`).
- gate_type = getattr(cfg, "output_gate_type", None)
- assert gate_type == "sigmoid", (
- f"gate activation {gate_type!r} is not the sigmoid form this "
- f"target passes to rms_norm_gated_token_major"
- )
-
- # ---- MoE ----------------------------------------------------------
- self.num_experts = cfg.num_experts
- self.topk = cfg.num_experts_per_tok
- self.moe_inter = cfg.moe_intermediate_size
- self.shared_inter = cfg.shared_expert_intermediate_size
- assert cfg.norm_topk_prob, "renorm_moe_routing_op always renormalizes"
- assert self.num_experts <= 512 and self.topk <= 16, (
- f"renorm_moe_routing_op takes at most 512 experts and top-16; this "
- f"config declares {self.num_experts} / {self.topk}"
- )
- # Every layer is MoE; the dense `intermediate_size` is dead config.
- assert getattr(cfg, "decoder_sparse_step", 1) == 1, "every layer must be MoE"
- assert not getattr(cfg, "mlp_only_layers", []), "every layer must be MoE"
- # trtllm-gen NVFP4 runner operand-width rules and fp4_quantize's block rule; they
- # also cover the CuteDSL FC1 / FC2 rules (K and I in 64-element scale groups, the
- # gated N in 128-row pairs of 64 up rows and 64 gate rows).
- assert self.hidden % 256 == 0, "fp4_block_scale_moe_runner hidden alignment"
- assert (2 * self.moe_inter) % 128 == 0, "fp4_block_scale_moe_runner 2I alignment"
- assert (self.hidden // _SF_VEC) % 4 == 0 and (self.moe_inter // _SF_VEC) % 4 == 0, (
- "the 128x4 block-scale swizzle needs both scale widths to be a multiple of 4"
- )
- # flashinfer_silu_and_mul's precondition is on the half width.
- assert (self.shared_inter * 2) % 16 == 0, "flashinfer_silu_and_mul half alignment"
- # Which kernels run the trunk's routed experts; it also fixes the load-time layout
- # of the expert stacks (`weights.py`), so it is decided before anything is declared.
- requested = model_config.moe_backend_requested.upper()
- assert requested in _TRUNK_MOE_BACKENDS, (
- f"moe_config.backend={model_config.moe_backend_requested!r} is not implemented by "
- f"this target; the supported requests are {sorted(_TRUNK_MOE_BACKENDS)}"
- )
- self.trunk_moe_backend = _TRUNK_MOE_BACKENDS[requested]
- self.moe_prefill_copy = _moe_prefill_copy(
- self.trunk_moe_backend, model_config.max_num_tokens
- )
- self.max_num_tokens = model_config.max_num_tokens
-
- # ---- PLE -----------------------------------------------------------
- ple_ids = list(getattr(cfg, "ple_layer_ids", None) or [])
- assert len(ple_ids) == 1, (
- f"this target implements the single PLE decoder layer this "
- f"checkpoint ships; the config declares ple_layer_ids={ple_ids}"
- )
- # One-based in the config.
- self.ple_layer = ple_ids[0] - 1
- assert 0 <= self.ple_layer < self.num_layers, self.ple_layer
- assert not self.is_attn[self.ple_layer], (
- f"the PLE lifecycle state rides the recurrent pools, so its layer "
- f"must be a linear-attention layer; layer {self.ple_layer} is not"
- )
- self.ngram_size = cfg.ngram_size
- self.heads_per_ngram = cfg.heads_per_ngram
- self.ngram_heads = self.heads_per_ngram * (self.ngram_size - 1)
- self.ple_embed_dim = cfg.ple_embed_dim
- self.ple_head_dim = self.ple_embed_dim // self.ngram_heads
- assert self.ple_head_dim * self.ngram_heads == self.ple_embed_dim, (
- f"{self.ngram_heads} n-gram heads must tile ple_embed_dim {self.ple_embed_dim}"
- )
- self.ple_conv_kernel = cfg.ple_conv_kernel_size
- # Length of the runtime's paged conv state.
- self.ple_dilation = self.ngram_size
- self.ple_state_len = (self.ple_conv_kernel - 1) * self.ple_dilation
- self.ngram_context_len = self.ngram_size - 1
- # Shard count and padding rule are config; the table height is not, so the
- # table is declared at load time (`weights.py`).
- self.ple_shards = cfg.split_ngram_parts
- self.ngram_vocab_divisor = cfg.make_ngram_vocab_size_divisible_by
- assert self.ple_shards > 0 and self.ngram_vocab_divisor > 0, (
- self.ple_shards,
- self.ngram_vocab_divisor,
- )
- # Fresh-sequence seed for the n-gram id history and the hash's segment token
- # (element 0 when a list, as in the HF reference).
- eos = cfg.eos_token_id
- self.eos_token_id = int(eos[0] if isinstance(eos, (list, tuple)) else eos)
-
- # ---- the vision tower ------------------------------------------------
- # A 27-block ViT (no KV cache, deepstack or quantization) feeding a 2x2 patch
- # merger. Constants the config omits: `_VISION_ROPE_THETA`,
- # `_VISION_NORM_EPS`, and the merger's erf GELU.
- vcfg = getattr(composite_config, "vision_config", None)
- assert vcfg is not None, (
- "this target implements the composite export's vision tower; the "
- "composite config carries no vision_config"
- )
- assert getattr(vcfg, "model_type", None) == "qwen4_exp_vision", (
- f"vision_config.model_type is {getattr(vcfg, 'model_type', None)!r}; "
- "this target is the Qwen4-Exp vision tower"
- )
- self.v_depth = vcfg.depth
- self.v_hidden = vcfg.hidden_size
- self.v_heads = vcfg.num_heads
- self.v_inter = vcfg.intermediate_size
- self.v_merge = vcfg.spatial_merge_size
- self.v_patch = vcfg.patch_size
- self.v_tpatch = vcfg.temporal_patch_size
- self.v_channels = vcfg.in_channels
- self.v_pos_rows = vcfg.num_position_embeddings
- assert self.v_hidden % self.v_heads == 0, (self.v_hidden, self.v_heads)
- self.v_head_dim = self.v_hidden // self.v_heads
- # vision_rope_2d needs a head divisible by 4 (two axes, one rotation pair
- # each).
- assert self.v_head_dim in _VISION_ROPE_HEAD_DIMS, (
- f"ops/attention/vision_rope_2d is certified at head widths "
- f"{_VISION_ROPE_HEAD_DIMS}; this tower's is {self.v_head_dim}"
- )
- assert (self.v_heads, self.v_heads, self.v_head_dim) in _VISION_CERTIFIED_HEADS, (
- f"the tower's attention is the encoder-only (no KV cache) mode of "
- f"attention/thop_attention, certified at {sorted(_VISION_CERTIFIED_HEADS)}; "
- f"this tower is {(self.v_heads, self.v_heads, self.v_head_dim)}"
- )
- # `hidden_act` names the block MLP's activation only; the merger runs the erf
- # `nn.GELU()`. Applying it to both is wrong by under one bf16 ulp, so there
- # are two entries (`gelu_tanh`, `gelu_erf`).
- assert vcfg.hidden_act == "gelu_pytorch_tanh", (
- f"the vision blocks' MLP activation is {vcfg.hidden_act!r}; this "
- "target calls ops/activation/gelu_tanh for it"
- )
- # No deepstack (`routing.py` refuses it; this is the second lock).
- assert not list(getattr(vcfg, "deepstack_visual_indexes", []) or []), (
- f"deepstack_visual_indexes={list(vcfg.deepstack_visual_indexes)} asks "
- "for per-layer visual feature injection, which this target does not "
- "implement"
- )
- assert vcfg.out_hidden_size == self.hidden, (
- f"the merger projects into the language model's hidden width; "
- f"out_hidden_size {vcfg.out_hidden_size} against hidden {self.hidden}"
- )
- # pos_embed_interpolate asserts a square table; fail before the first image.
- side = int(math.isqrt(self.v_pos_rows))
- assert side * side == self.v_pos_rows, (
- f"num_position_embeddings {self.v_pos_rows} is not a perfect square, "
- "so it is not a square position grid"
- )
- self.v_grid_side = side
- # `spatial_merge_size ** 2` patches become one token at the merger.
- self.v_merge_dim = self.v_hidden * self.v_merge * self.v_merge
- # One flat pixel row per patch, `(in_channels, kt, kh, kw)` channel-slowest.
- self.v_pixel_width = self.v_channels * self.v_tpatch * self.v_patch * self.v_patch
- # Placeholder token ids live on the composite config's top level.
- self.image_token_id = int(composite_config.image_token_id)
- self.video_token_id = int(composite_config.video_token_id)
- assert self.image_token_id < self.vocab and self.video_token_id < self.vocab, (
- f"the media placeholders ({self.image_token_id}, {self.video_token_id}) "
- f"are in-vocabulary against vocab_size {self.vocab}, so the engine "
- "cannot find them with its out-of-vocabulary fallback and this target "
- "must publish them as `multimodal_token_ids`"
- )
-
- # ---- KV pool storage ------------------------------------------------
- # `None` keeps the checkpoint's bf16 K/V (`quant_mode` 0); `FP8`, what
- # `kv_cache_config.dtype: fp8` resolves to, stores them as e4m3 under
- # QM_FP8_KV_CACHE; `NVFP4` as e2m1 pairs with e4m3 block scales under
- # QM_NVFP4_KV_CACHE. Anything else is refused: the QSA backend serves its sparse path
- # over half, bf16, fp8 and (with its own kernels here) NVFP4 pools, and silently falls
- # back to dense otherwise. The index and position planes are kernel contracts, not KV
- # storage, and keep their types under all three.
- kv_algo = model_config.quant_config.kv_cache_quant_algo
- assert kv_algo in (None, QuantAlgo.FP8, QuantAlgo.NVFP4), (
- f"this target stores K/V as the checkpoint's bf16 (kv_cache_quant_algo None), "
- f"as fp8-e4m3 (FP8) or as NVFP4; the engine resolved kv_cache_quant_algo {kv_algo!r}"
- )
- self.quant_mode = {
- None: 0,
- QuantAlgo.FP8: QM_FP8_KV_CACHE,
- QuantAlgo.NVFP4: QM_NVFP4_KV_CACHE,
- }[kv_algo]
- # What `check_attention_pools` holds every manager's K/V pool to, the draft layer's
- # included: the dtype and the innermost extent (an NVFP4 byte holds two values).
- self.kv_pool_dtype = {
- None: dt,
- QuantAlgo.FP8: torch.float8_e4m3fn,
- QuantAlgo.NVFP4: torch.uint8,
- }[kv_algo]
- self.kv_pool_width = self.head_dim // 2 if kv_algo == QuantAlgo.NVFP4 else self.head_dim
- # The weight table reads the tower's geometry off this.
- self.vision_config = vcfg
-
- self.mtp_layer_idx = self.num_layers
- if self.mtp_theta is None:
- self.mtp_theta = self.theta
- assert self.mtp_theta > 0.0, self.mtp_theta
- if self.mtp_enabled:
- mode = spec_config.spec_dec_mode
- assert mode.is_mtp_eagle_one_model(), (
- f"this target implements the one-model MTP-Eagle draft loop, "
- f"where one module is replayed; the engine resolved "
- f"spec_dec_mode {mode!r}"
- )
- resolved = getattr(spec_config, "num_nextn_predict_layers", None)
- assert resolved == _MTP_NUM_LAYERS, (
- f"the engine resolved num_nextn_predict_layers={resolved} from "
- f"the checkpoint; this target carries the weights of "
- f"{_MTP_NUM_LAYERS} MTP module"
- )
- # `max_draft_len + 1` becomes `predicted_tokens_per_seq` on every attention call
- # and the cached replay's step width, certified at up to `_MAX_VERIFY_WIDTH`.
- max_draft = spec_config.max_draft_len
- assert 1 <= max_draft <= _MAX_VERIFY_WIDTH - 1, (
- f"max_draft_len={max_draft} needs {max_draft + 1} verify rows per request, and "
- f"attention/thop_attention and the GDN cached replay are certified at up to "
- f"{_MAX_VERIFY_WIDTH} (max_draft_len {_MAX_VERIFY_WIDTH - 1}); 8 is also the "
- "replay kernel's ceiling"
- )
- assert self.trunk_moe_backend != "CUTEDSL" or max_draft <= _CUTEDSL_MAX_DRAFT_LEN, (
- f"max_draft_len={max_draft} with moe_config.backend=CUTEDSL: the CuteDSL expert "
- f"pipeline is certified at max_draft_len up to {_CUTEDSL_MAX_DRAFT_LEN}"
- )
- # A linear tree keeps the Blackwell spec-dec mask surface inert.
- assert not getattr(spec_config, "use_dynamic_tree", False), (
- "dynamic-tree drafting is not implemented by this target"
- )
- assert getattr(spec_config, "eagle_choices", None) is None, (
- "static-tree drafting is not implemented by this target"
- )
- assert spec_config.max_total_draft_tokens == max_draft, (
- f"a linear tree mirrors max_total_draft_tokens from "
- f"max_draft_len; the engine resolved "
- f"{spec_config.max_total_draft_tokens} against {max_draft}"
- )
-
- # The recurrent half of the verify step has two forms, and the cache manager
- # picks one when it is built (`TRTLLM_USE_GDN_REPLAY`, default on): the cached
- # replay or the legacy per-step recurrent buffers. Both are implemented;
- # `_bind_step_buffers` reads which one a manager allocated.
- # The FP8 block-scale runner's own shape rules.
- assert self.hidden % _FP8_BLOCK == 0 and self.moe_inter % _FP8_BLOCK == 0, (
- f"fp8_block_scale_moe_runner needs hidden ({self.hidden}) and "
- f"intermediate ({self.moe_inter}) to be multiples of {_FP8_BLOCK}"
- )
- assert self.topk < self.num_experts and self.topk <= 32, (
- f"the no-groups renormalize routing kernel takes 0 < top_k <= 32 "
- f"and top_k < num_experts; this config declares {self.topk} of "
- f"{self.num_experts}"
- )
-
- # The dtype the table is declared at in weights.py.
- self.ple_table_dtype = torch.float8_e4m3fn
- # Resolved at construction; read by `weights._declare_ngram_table`.
- self.ple_host_offload = _ple_host_offload_enabled()
- # CPU allocation behind the offloaded table. The `ple_table` parameter is a
- # device-address view that owns none of those pages; dropping this reference
- # frees memory a kernel still points into. None when resident.
- self._ple_table_host = None
-
- # Every weight but the n-gram table; meta-init intercepts torch.empty.
- self.w = _weights.MODEL_WEIGHTS.declare(_weights.MODEL_WEIGHTS.dims(self))
-
- # Derived once after load (`post_load_weights`).
- self._moe_scalars: list | None = None
- self._moe_scalars_cd: list | None = None
- self._gdn_gating: dict | None = None
- self._cos_sin: torch.Tensor | None = None
- # The attention op's KV scales: fp32 [1] ones over an fp8 pool (never None there: the
- # generation kernels read `kvCacheScale[0]`), fp32 [3] ones (q, K, V) over an NVFP4
- # one, None over bf16. Unit, because the fp8 QSA store and sparse GQA take no scale
- # while a request's history is written by the dense and the sparse path alike, so the
- # two agree on what a slot holds only at 1.0 (the NVFP4 entries take the [3] tensor
- # and agree at any value, but the checkpoint has no KV scales to calibrate with).
- # Lossy: fp8 K/V are stored by a plain cast, so a magnitude past the e4m3 maximum (448)
- # saturates (`_KvRangeProbe` reports it under TRTLLM_MODELING_V2_VALIDATE).
- self.kv_scale_orig_quant: torch.Tensor | None = None
- self.kv_scale_quant_orig: torch.Tensor | None = None
- # MTP module operands; None with speculation off.
- self._mtp: _MtpWeights | None = None
- # The draft experts' NVFP4 copy (`_mtp_moe_nvfp4`), built once: it frees the FP8 ones.
- self._mtp_moe_nvfp4: dict | None = None
- self._mtp_cos_sin: torch.Tensor | None = None
- # The verify step's side-cache rollback, registered with the spec worker; None
- # with speculation off.
- self.spec_rollback: SpeculativeRollback | None = (
- SpeculativeRollback() if self.mtp_enabled else None
- )
- self._prefill: "PrefillTarget | None" = None
- self._decode: "DecodeTarget | None" = None
- self._mtp_fused: "_MtpFused | None" = None
- # The draft head's FP8 copy of lm_head: (e4m3 [vocab, hidden], fp32 [1, vocab] row
- # scales).
- self._mtp_head_fp8: "tuple[torch.Tensor, torch.Tensor] | None" = None
- # A packed block-scaled e4m3 copy of lm_head for decode rows (`fp8_gemv`): the draft
- # head's, and the target head's under the dense FP8 mode.
- self._lm_head_g = None
- # The draft head's NVFP4 copy (`_MTP_HEAD_NVFP4_ENV`), read by `nvfp4_gemv` at decode
- # rows in place of the packed FP8 copy.
- self._mtp_head_nv = None
- # The shared expert's stream; see `_moe_block`. Made in `post_load_weights`:
- # a stream created inside a graph capture is not capturable.
- self._side_stream: torch.cuda.Stream | None = None
- self._moe_memset_stream: torch.cuda.Stream | None = None
- self.vision_tower: VisionTower | None = None
- # Everything derived from the KV cache manager (`_bind_cache_manager`);
- # None until the first forward.
- self._bound: _Bound | None = None
- self._binds = 0
- # Set by `begin_step`, cleared by `forward`: a forward without a prepared step is a bug.
- self._step_begun = False
- # The draft layer's own bind, keyed on the manager its draft step sees.
- self._draft_bound: _DraftBound | None = None
-
- # ---------------------------------------------------------------- load ---
-
- def _norm_slices(self, key: str) -> tuple:
- """Split a grouped-norm weight into its per-group slices (taken once)."""
- return tuple(torch.split(self.w[key], self.hidden, dim=0))
-
- def _fuse_hc_first_projections(self) -> None:
- """Concatenate each trunk hyper-connection's `down` and `block_inject`
- weights into one `[rows, stream]` tensor, zero-padded to a multiple of
- 128 rows, so a site's first projection is one GEMM
- (`_hc_first[(layer, tag)]`; `_hc_first["fhc"]` is the stack-final
- mixer's `down`, padded the same way). cuBLAS was measured 1.8x (decode)
- to 4.5x (8192 rows) faster at N = 384 than at the unpadded 324; its calls
- multiply the first `hc_first_cols` (352) columns, faster still.
-
- The checkpoint's parameters are rebound as row views of the padded
- tensors, so this costs only the padding rows (~1.2 MB per site), and
- anything still reading them separately reads the same values."""
- w = self.w
-
- def fuse(*params):
- rows = sum(t.shape[0] for t in params)
- padded = -(-rows // 128) * 128
- pad = torch.zeros(
- [padded - rows, self.stream], dtype=params[0].dtype, device=params[0].device
- )
- both = torch.cat([t.data for t in params] + [pad], 0)
- start = 0
- for t in params:
- n = t.shape[0]
- t.data = both[start : start + n]
- start += n
- return both
-
- self._hc_first = {
- (i, tag): fuse(w[f"l{i}_{tag}_down"], w[f"l{i}_{tag}_inject"])
- for i in range(self.num_layers)
- for tag in ("ahc", "mhc")
- }
- self._hc_first["fhc"] = fuse(w["fhc_down"])
- if self.mtp_enabled:
- p = f"l{self.mtp_layer_idx}_mtp_"
- self._hc_first["mtp_ahc"] = fuse(w[f"{p}ahc_down"], w[f"{p}ahc_inject"])
- self._hc_first["mtp_mhc"] = fuse(w[f"{p}mhc_down"], w[f"{p}mhc_inject"])
- self._hc_first["mtp_fhc"] = fuse(w[f"{p}fhc_down"])
-
- def post_load_weights(self) -> None:
- """Derive what the targets read, then build them, now that the weights are real.
-
- Not an override (DecoderModel has no such hook); named for the shell hook."""
- w = self.w
- device = w["embed"].device
- moe_scalars = [
- _moe_layer_scalars(w, i, self.trunk_moe_backend) for i in range(self.num_layers)
- ]
- # With the second expert copy (`moe_prefill_copy`), its CuteDSL scalars, which
- # prefill-sized launches run on (`_CUTE_DSL_MOE_MIN_ROWS`).
- moe_scalars_cd = (
- [_moe_layer_scalars(w, i, "CUTEDSL") for i in range(self.num_layers)]
- if self.moe_prefill_copy
- else None
- )
- self._moe_scalars = moe_scalars
- self._moe_scalars_cd = moe_scalars_cd
- # Both delta-rule entries index A_log and dt_bias as `ptr + head` and certify
- # fp32; the checkpoint is bf16.
- self._gdn_gating = {
- i: (
- w[f"l{i}_A_log"].float().contiguous(),
- w[f"l{i}_dt_bias"].float().contiguous(),
- )
- for i in self.gdn_layers
- }
- # Reconcile the table height against the checkpoint's id-space tensors: an
- # index past the end is a device-side assert that poisons the process.
- table_rows = w["ple_table"].shape[0]
- used_rows = int(w["ple_offs"][-1]) + int(w["ple_vocab"][-1])
- assert used_rows <= table_rows, (
- f"the n-gram table has {table_rows} rows but its per-head offsets "
- f"and vocabulary sizes address {used_rows}"
- )
- assert table_rows % self.ngram_vocab_divisor == 0, (
- f"the n-gram table's {table_rows} rows are not a multiple of "
- f"make_ngram_vocab_size_divisible_by={self.ngram_vocab_divisor}"
- )
- assert table_rows - used_rows < self.ngram_vocab_divisor, (
- f"the n-gram table carries {table_rows - used_rows} padding rows, "
- f"more than one divisor of {self.ngram_vocab_divisor}: the shards "
- f"do not concatenate to the padded id space this checkpoint "
- f"declares"
- )
- self._cos_sin = self._index_rope_table(self.theta, device)
- if self.mtp_enabled:
- # The draft layer's index plane rotates at the module's own rope base, like its
- # attention rope.
- self._mtp_cos_sin = (
- self._cos_sin
- if self.mtp_theta == self.theta
- else self._index_rope_table(self.mtp_theta, device)
- )
- if self.quant_mode in (QM_FP8_KV_CACHE, QM_NVFP4_KV_CACHE):
- # Built here, not at construction: that runs under meta init.
- scales = 3 if self.quant_mode == QM_NVFP4_KV_CACHE else 1
- self.kv_scale_quant_orig = torch.ones(scales, dtype=torch.float32, device=device)
- self.kv_scale_orig_quant = 1.0 / self.kv_scale_quant_orig
- self._report_ple_residency()
- if self.mtp_enabled:
- self._mtp = _mtp_views(self)
- if _mtp_moe_nvfp4_enabled() and device.type == "cuda":
- if self._mtp_moe_nvfp4 is None:
- self._mtp_moe_nvfp4 = _mtp_moe_nvfp4(self)
- self._mtp = self._mtp._replace(**self._mtp_moe_nvfp4)
-
- self._fuse_hc_first_projections()
- if device.type == "cuda":
- # Every first-half tile `DecodeTarget` reaches up to `_HC_TC_FIRST_MAX_ROWS` rows.
- hc_tc_first_precompile(self.stream, self.hc_lowrank, self.hc_count, device)
- if self.mtp_enabled:
- self._mtp_fused = _mtp_fused_views(self)
- # FP8 copies of the trunk's dense projections (`_DENSE_FP8_ENV`), beside the bf16
- # weights the 1-row GEMV / fused paths and the draft layer keep reading.
- self._dense8 = {}
- # ... and packed FP8 copies with 128x128 block scales for the decode-row `fp8_gemv`
- # (the shared expert's gate|up without its gate-logit row, which stays bf16).
- self._gemv8 = {}
- # ... and `mid_rows_mm`'s view of the same weights for decode rows in `_MID_MM_ROWS`:
- # the per-tensor e4m3 codes above, the shared expert's and the router's bf16 weights
- # (no new copies).
- self._mid = {}
- if _dense_fp8_enabled():
- for name in w:
- if re.fullmatch(r"l\d+_(attn_in_w|o_w|in_proj_w|out_w)", name):
- codes, scale = fp8_per_tensor_weight(w[name])
- self._dense8[name] = (codes.t(), scale)
- self._gemv8[name] = pack_from_bf16(w[name])
- self._mid[name] = MidWeight(codes, scale)
- elif re.fullmatch(r"l\d+_sh_gu", name):
- self._gemv8[name] = pack_from_bf16(w[name][: 2 * self.shared_inter])
- self._mid[name] = MidWeight(w[name][: 2 * self.shared_inter])
- elif re.fullmatch(r"l\d+_sh_dn", name):
- self._gemv8[name] = pack_from_bf16(w[name])
- self._mid[name] = MidWeight(w[name])
- elif re.fullmatch(r"l\d+_router", name):
- self._mid[name] = MidWeight(w[name])
- if device.type == "cuda":
- _precompile_mid(self._mid)
- # NVFP4 copies of the GDN input projections (`_GDN_INPROJ_NVFP4_ENV`) for decode rows,
- # with every variant a decode step can reach compiled now.
- self._gemv_nv = {}
- if _gdn_inproj_nvfp4_enabled() and device.type == "cuda":
- for name in w:
- if re.fullmatch(r"l\d+_in_proj_w", name):
- self._gemv_nv[name] = nvfp4_pack_split_from_bf16(w[name])
- if self._gemv_nv:
- precompile_split(
- next(iter(self._gemv_nv.values())),
- _INPROJ_NVFP4_MAX_ROWS,
- self.conv_dim,
- _PAD_SLOT_ID,
- tuple(range(2, self.max_verify_rows + 1)),
- )
-
- # The CPU-only load checks build the targets over host weights; they launch nothing.
- self._side_stream = torch.cuda.Stream(device=device) if device.type == "cuda" else None
- # The prefill CuteDSL MoE's output memset, beside its FC1 (`fp4_moe_renorm_cute_dsl`):
- # its own stream, so the join before FC2 does not also wait for the shared expert.
- self._moe_memset_stream = (
- torch.cuda.Stream(device=device) if device.type == "cuda" else None
- )
- self._prefill = PrefillTarget(self)
- self._decode = DecodeTarget(self)
- self.vision_tower = VisionTower(self)
-
- def _index_rope_table(self, theta: float, device) -> torch.Tensor:
- """The index branch's rope table at `theta` in the QSA kernels' layout:
- `[max_positions, 2 * rotary_pairs]` fp32, cos block then sin block."""
- pairs = self.rotary_dim // 2
- exponent = torch.arange(0, self.rotary_dim, 2, dtype=torch.float64, device=device)
- inv_freq = (1.0 / (theta ** (exponent / self.rotary_dim))).to(torch.float32)
- positions = torch.arange(self.max_positions, dtype=torch.float32, device=device)
- angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0)
- table = torch.cat([torch.cos(angles), torch.sin(angles)], dim=1).contiguous()
- assert table.shape == (self.max_positions, 2 * pairs), table.shape
- return table
-
- def _report_ple_residency(self) -> None:
- """Check where the n-gram table ended up, and report it once.
-
- Offload is output-identical, so a stray `.to()` or a driver that declined
- the mapping would pass every smoke test; the pointer's memory kind read
- back from the driver is the signal that cannot be faked. The report is
- printed here, before the KV pool is sized from what is left."""
- table = self.w["ple_table"]
- kind = _weights.memory_kind(table.data_ptr())
- expected = (
- _weights.CUDA_MEMORY_TYPE_HOST
- if self.ple_host_offload
- else _weights.CUDA_MEMORY_TYPE_DEVICE
- )
- assert kind == expected, (
- f"the n-gram table was declared with {_PLE_HOST_OFFLOAD_ENV}="
- f"{'1' if self.ple_host_offload else '0'}, so its pages must be "
- f"{'host' if self.ple_host_offload else 'device'} memory, but the "
- f"driver reports cudaMemoryType {kind} for the pointer the forward "
- f"gathers from. The offload is output-identical, so nothing "
- f"downstream of this would notice"
- )
- mib = 1024.0 * 1024.0
- table_mib = table.numel() * table.element_size() / mib
- free, total = torch.cuda.mem_get_info()
- print(
- f"ple n-gram table: {tuple(table.shape)} {table.dtype}, "
- f"{table_mib:.1f} MiB, "
- f"{'HOST-resident (device-addressable)' if self.ple_host_offload else 'DEVICE-resident'}"
- f" [{_PLE_HOST_OFFLOAD_ENV}="
- f"{os.environ.get(_PLE_HOST_OFFLOAD_ENV, '1 (default)')}]; "
- f"after load: allocated {torch.cuda.memory_allocated() / mib:.1f} MiB, "
- f"reserved {torch.cuda.memory_reserved() / mib:.1f} MiB, "
- f"driver-used {(total - free) / mib:.1f} MiB of {total / mib:.1f} MiB",
- flush=True,
- )
-
- # ------------------------------------------------------------ fail-fast ---
-
- def _pool_geometry(self, tokens_per_block: int) -> _PoolGeometry:
- return _PoolGeometry(
- tokens_per_block=tokens_per_block,
- heads_kv=self.heads_kv,
- head_dim=self.head_dim,
- index_kv_heads=self.index_kv_heads,
- index_head_dim=self.index_head_dim,
- conv_dim=self.conv_dim,
- conv_kernel=self.conv_kernel,
- num_k_heads=self.num_k_heads,
- num_v_heads=self.num_v_heads,
- head_v_dim=self.head_v_dim,
- head_k_dim=self.head_k_dim,
- stream=self.stream,
- ple_state_len=self.ple_state_len,
- ngram_context_len=self.ngram_context_len,
- dtype=self.w["embed"].dtype,
- kv_dtype=self.kv_pool_dtype,
- kv_width=self.kv_pool_width,
- )
-
- def _kv_planes(self, cm, g: _PoolGeometry, layer: int, index_cache) -> _KvPlanes:
- """One layer's paged K/V views from `cm`, checked against the contract: the data
- planes, and under NVFP4 the block-scale planes that go with them."""
- pool = cm.get_buffers(layer, kv_layout="HND")
- check_attention_pools(g, layer, index_cache, pool)
- if self.quant_mode != QM_NVFP4_KV_CACHE:
- return _KvPlanes(pool[:, 0], pool[:, 1])
- _require(cm, ("get_block_scale_buffers",), "reads an NVFP4 pool's e4m3 block-scale planes")
- scales = cm.get_block_scale_buffers(layer, kv_layout="HND")
- check_kv_scale_planes(g, layer, pool, scales)
- return _KvPlanes(pool[:, 0], pool[:, 1], scales[:, 0], scales[:, 1])
-
- def _bind_cache_manager(self, md, position_ids) -> None:
- """Bind this module to one KV cache manager: fail-fast on the contract, then
- take every pool view from it.
-
- Called once per manager, not per model: with capacity estimation on, the
- engine builds a throwaway manager, warms up, shuts it down and builds the
- real one on the same module, and a view held across the swap dangles
- silently. The result is built in locals and installed as one `_Bound`, so a
- failed bind leaves the previous binding untouched. The contract checks
- re-run on every bind."""
- # Checked first: some fields below are properties that assert on a
- # manager-less metadata instead of raising AttributeError.
- cm = md.kv_cache_manager
- assert cm is not None, "a paged KV cache is required"
- # A re-bind during capture would bake two managers' pointers into one graph;
- # the engine releases graphs before its own swap, so this is an ordering
- # violation.
- assert not (self._binds and torch.cuda.is_current_stream_capturing()), (
- f"the KV cache manager changed to {type(cm).__name__} during CUDA graph capture"
- )
- self._binds += 1
-
- missing = [name for name in _QSA_FIELDS if not hasattr(md, name)]
- assert not missing, (
- f"metadata fields missing: {missing}; the engine did not select QSA "
- "metadata for this model"
- )
-
- # `position_ids` is the int32 [3, 1, T] buffer the engine builds on an mRoPE
- # config even for text-only batches; row 0 is the text position.
- assert position_ids.dtype == torch.int32, position_ids.dtype
- assert self.mrope_declared, "this target takes the mRoPE position layout"
- assert position_ids.dim() == 3 and position_ids.shape[0] == _QSA_POSITION_AXES, (
- f"expected the [{_QSA_POSITION_AXES}, 1, T] mRoPE position buffer, got "
- f"{tuple(position_ids.shape)}"
- )
-
- # thop_attention addresses the pool by the layer's pool offset, not its
- # ordinal among attention layers; counting them aliases layers with finite,
- # plausible output.
- pool_idx = {i: cm.layer_offsets[i] for i in self.attn_layers}
- pools = {row[0] for row in md.host_kv_cache_pool_mapping.tolist()}
- assert pools <= {0, 1}, (
- f"KV addressing over {len(pools)} pools is not certified "
- f"(thop_attention's cells cover one and two); layer->pool ids {sorted(pools)}"
- )
-
- # No sliding window: both targets pass `max_seq_len` as the window. The vector
- # is indexed by pool offset and recurrent layers carry a large negative
- # sentinel, so only attention offsets are read.
- windows = cm.max_attention_window_vec
- attn_windows = {windows[off] for off in pool_idx.values()}
- assert attn_windows <= {None, md.max_seq_len}, (
- f"the attention layers' pool windows are {sorted(map(repr, attn_windows))} "
- f"but this target passes max_seq_len {md.max_seq_len} as the window; "
- "a sliding window is not implemented here"
- )
-
- # The NVFP4 attention kernels exist for the paged-context path only; without it the
- # dense op would not read the pool the sparse path wrote.
- assert self.quant_mode != QM_NVFP4_KV_CACHE or md.use_paged_context_fmha, (
- "an NVFP4 KV cache is read by the paged-context attention kernels only, which the "
- "engine enables for chunked prefill, KV cache block reuse or speculative decoding; "
- "none of the three is on"
- )
- g = self._pool_geometry(md.tokens_per_block)
- _require(
- cm,
- ("get_index_k_buffer", "get_qsa_position_buffer", "get_buffers"),
- "reads the QSA index plane, the shared position plane and the paged K/V",
- )
- # `qsa_topk_select` checks only that the buffer is at least `index_topk` wide,
- # then addresses it as flat `index_topk`-wide rows ignoring the real stride, so
- # a width mismatch silently mis-writes every row but the first.
- assert md.qsa_topk_indices.shape[1] == self.block_topk, (
- f"the engine sized the QSA selection buffer {md.qsa_topk_indices.shape[1]} "
- f"wide and this target selects {self.block_topk} blocks per row"
- )
- position_cache = cm.get_qsa_position_buffer()
- check_position_plane(g, position_cache)
- index_cache, kv = {}, {}
- for i in self.attn_layers:
- index_cache[i] = cm.get_index_k_buffer(i)
- kv[i] = self._kv_planes(cm, g, i, index_cache[i])
-
- # A missed hybrid predicate would still build and pass a smoke run with 36
- # layers carrying no state.
- mm = md.mamba_metadata
- assert mm is not None and mm is not False, (
- "attn_metadata.mamba_metadata is absent: the runtime allocated no "
- f"recurrent-state pool, so the {len(self.gdn_layers)} linear-attention "
- "layers would run with no state"
- )
- mm_missing = [name for name in _MAMBA_FIELDS if not hasattr(mm, name)]
- assert not mm_missing, f"mamba metadata fields missing: {mm_missing}"
- # `ple_layer_cache` exists only on the V2 hybrid manager; under V1 the PLE
- # state silently does not exist.
- _require(
- cm,
- ("get_conv_states", "get_ssm_states", "ple_layer_cache"),
- "reads the recurrent, conv and PLE pools",
- )
- conv = {i: cm.get_conv_states(i) for i in self.gdn_layers}
- ssm = {i: cm.get_ssm_states(i) for i in self.gdn_layers}
- conv0, ssm0 = conv[self.gdn_layers[0]], ssm[self.gdn_layers[0]]
- check_recurrent_pools(g, conv0, ssm0)
-
- ple_cache = cm.ple_layer_cache(self.ple_layer)
- assert ple_cache is not None, (
- f"no PLE lifecycle state for decoder layer {self.ple_layer}: the runtime "
- "allocated neither the dilated conv state nor the n-gram id history"
- )
- ple_conv, ple_ngram = ple_cache
- check_ple_pools(g, ple_conv, ple_ngram, conv0)
-
- assert md.effective_beam_width == 1, "beam search is not implemented"
- assert md.cache_indirection is None, "beam search is not implemented"
- assert not md.is_cross, "cross attention is not implemented"
- # With speculation on, the Blackwell trtllm-gen spec-dec mask machinery is off
- # for a linear tree at sm_103: every `spec_decoding_*` tensor stays None and
- # `predicted_tokens_per_seq` alone produces the within-block mask.
- assert not md.is_spec_decoding_enabled, (
- "the attention metadata reports is_spec_decoding_enabled; on sm_103 a "
- "linear-tree MTP forces that machinery off"
- )
- bound = _Bound(
- cm, pool_idx, index_cache, position_cache, kv, conv, ssm, (ple_conv, ple_ngram)
- )
- if self.spec_rollback is not None:
- self.spec_rollback.set_index_planes([index_cache[i] for i in self.attn_layers])
- if self.mtp_enabled:
- bound = bound._replace(**self._bind_step_buffers(cm, g, mm, conv0, ssm0))
- self._bound = bound
-
- def _bind_step_buffers(self, cm, g: _PoolGeometry, mm, conv0, ssm0) -> dict:
- """Per-manager fail-fast for the MTP increment: the verify step needs
- runtime-owned buffers the base path never touches. Every row's conv window goes
- to a per-step buffer in both recurrent modes; the recurrent state is either the
- cached replay (a checkpoint per slot plus a compact history of accepted tokens,
- no per-step recurrent buffer) or, with `TRTLLM_USE_GDN_REPLAY=0`, per-step
- recurrent buffers. Returns the `_Bound` fields it adds."""
- assert cm.is_speculative(), (
- f"the KV cache manager {type(cm).__name__} reports is_speculative() False, "
- "so it allocated none of the per-step state buffers the verify path writes"
- )
- _require(
- cm,
- ("mamba_layer_cache", "get_replay_state_update_metadata"),
- "binds its per-step state buffers and replay metadata",
- )
- # Bind the per-layer views once for subsequent forwards.
- layer_caches = {i: cm.mamba_layer_cache(i) for i in self.gdn_layers}
- int_conv = {i: cache.intermediate_conv_window for i, cache in layer_caches.items()}
- int_idx = getattr(cm, "intermediate_state_indices", None)
- g0 = self.gdn_layers[0]
- check_step_buffers(g, int_conv[g0], conv0, int_idx)
- int_ssm0 = layer_caches[g0].intermediate_ssm
- if cm.get_replay_state_update_metadata() is None:
- check_legacy_state(g, int_ssm0, ssm0)
- int_ssm = {i: cache.intermediate_ssm for i, cache in layer_caches.items()}
- return dict(int_idx=int_idx, int_conv=int_conv, int_ssm=int_ssm)
- assert int_ssm0 is None, (
- "the cache manager runs the cached-replay verify path and also allocated a "
- "per-step recurrent buffer; the runtime allocates one or the other"
- )
- return dict(int_idx=int_idx, int_conv=int_conv, replay=self._bind_replay(cm, g, mm, ssm0))
-
- def _bind_replay(self, cm, g: _PoolGeometry, mm, ssm0) -> _Replay:
- """The cached-replay history buffers and bookkeeping of one manager."""
- meta = cm.get_replay_state_update_metadata()
- assert meta is not None, (
- "the cached-replay binding requires replay metadata "
- "(prev_num_accepted_tokens / cache_buf_idx / widths)"
- )
- history, width = meta.replay_history_size, meta.replay_step_width
- check_replay_state(history, width, meta.prev_num_accepted_tokens, meta.cache_buf_idx, ssm0)
- old = {}
- for i in self.gdn_layers:
- layer_cache = cm.mamba_layer_cache(i)
- old[i] = (
- layer_cache.old_x,
- layer_cache.old_B,
- layer_cache.old_dt,
- layer_cache.old_dA_cumsum,
- )
- check_replay_history(g, i, history, ssm0.shape[0], *old[i][:3])
- # The partitioned replay is the manager's to commit: it commits exactly when this
- # flag holds and a step carries `CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE` or more
- # generation requests, and the forward hands the kernel the work items under that
- # same condition.
- all_layer = bool(getattr(cm, "use_gdn_cached_replay_all_layer_commit", False))
- # A bf16 pool replays on the CuTe kernel, an fp32 pool on the Triton one; both commit a
- # wrapping checkpoint in-kernel when given no work items and leave it to the manager's
- # all-layer commit when given them.
- cute = ssm0.dtype == torch.bfloat16
- if all_layer:
- missing = [n for n in ("replay_work_items", "replay_n_writes") if not hasattr(mm, n)]
- assert not missing, (
- f"the manager's all-layer replay commit is on but the mamba metadata lacks "
- f"{missing}; the partitioned verify call needs both"
- )
- return _Replay(
- old, meta.cache_buf_idx, meta.prev_num_accepted_tokens, history, width, all_layer, cute
- )
-
- def _bind_draft_cache_manager(self, md) -> None:
- """Bind the draft layer to the manager its draft step sees: its own index plane,
- that manager's position plane and its paged K/V.
-
- With a separate draft KV cache that is the draft manager the worker swaps in for
- the draft loop, whose pools hold the draft layer's pages and none of the trunk's;
- otherwise the trunk's manager, which maps the draft layer as well. Keyed on the
- manager like `_bind_cache_manager`. Never first bound under capture: the runner
- runs every captured batch shape eagerly first, so a bind under capture would bake
- views no eager step validated into the graph."""
- cm = md.kv_cache_manager
- assert not torch.cuda.is_current_stream_capturing(), (
- f"the draft layer is not bound to KV cache manager {type(cm).__name__} (first use, "
- "or the manager changed since its last bind) and CUDA graph capture is active; "
- "an eager step must take and check its views first"
- )
- _require(
- cm,
- ("get_index_k_buffer", "get_qsa_position_buffer", "get_buffers"),
- "reads the draft layer's QSA index plane, position plane and paged K/V",
- )
- g = self._pool_geometry(md.tokens_per_block)
- layer = self.mtp_layer_idx
- index_cache = cm.get_index_k_buffer(layer)
- position_cache = cm.get_qsa_position_buffer()
- check_position_plane(g, position_cache)
- kv = self._kv_planes(cm, g, layer, index_cache)
- share_slot = None
- if self.mtp_index_share:
- # The metadata sizes its shared selection buffers from the same resolved knob,
- # one row per sparse layer of whichever manager is bound.
- assert md.sparse_metadata_params.mtp_index_share, (
- "this target resolved IndexShare on but the engine's QSA metadata did not"
- )
- slots = cm.qsa_shared_topk_slots
- check_shared_selection(md.qsa_shared_topk_indices, slots, layer, self.block_topk)
- share_slot = slots[layer]
- self._draft_bound = _DraftBound(cm, index_cache, position_cache, kv, share_slot)
-
- # -------------------------------------------------------------- forward ---
-
- def _select_target(self, attn_metadata) -> "PrefillTarget | DecodeTarget":
- """Which target runs this step: decode if `num_contexts == 0`, else prefill.
-
- A draft loop rewrites `num_contexts` in place between draft steps, so the
- draft step routes through here per call."""
- return self._decode if attn_metadata.num_contexts == 0 else self._prefill
-
- def begin_step(
- self,
- attn_metadata: AttentionMetadata,
- input_ids: torch.IntTensor | None,
- position_ids: torch.IntTensor | None,
- inputs_embeds: torch.FloatTensor | None = None,
- lora_params: dict | None = None,
- spec_metadata=None,
- vision: tuple | None = None,
- ) -> None:
- """Check what holds for every step, bind the cache manager, choose the target and
- let the prefill target prepare the step.
-
- The shell calls it before `forward`, on every step. It is separate because a
- breakable prefill CUDA graph captures `forward` and replays it without calling it:
- whatever a step reads from the host (metadata, host buffers, the manager) is read
- here, where it runs eagerly on a replayed step too. `vision` is `(rows,
- mm_token_indices)` when the step carries media: the tower's `[R, hidden]` rows and
- the `[R]` packed-batch rows they replace; the prefill target splices them inside
- its forward."""
- assert self._prefill is not None and self._decode is not None, (
- "load_weights must run before forward"
- )
- assert position_ids is not None
- assert isinstance(attn_metadata, TrtllmAttentionMetadata)
- # Unimplemented inputs must fail loudly, not be dropped.
- assert lora_params is None, "LoRA is not implemented by this target"
- # `input_ids` is required on every step: the PLE block hashes token ids whether or
- # not vision rows replace some embeddings, and text gates would not catch a mis-hash.
- assert input_ids is not None, (
- "this target's PLE layer hashes token ids, so input_ids is required "
- "on every step, including one that carries vision rows"
- )
- # The token embedding is this target's own and vision rows are spliced into it
- # inside the forward; an embedding handed in from outside would be dropped on every
- # replay of a breakable prefill graph.
- assert inputs_embeds is None, (
- "this target computes its token embedding itself and splices vision rows into "
- "it; an engine-supplied inputs_embeds is not implemented"
- )
- # `spec_metadata` is present exactly when drafting. On MTP_EAGLE_ONE_MODEL it
- # sets `layers_to_capture = ()`, so the trunk owes no capture hook.
- assert (spec_metadata is not None) == self.mtp_enabled, (
- f"spec_metadata {'arrived' if spec_metadata is not None else 'is absent'} "
- f"but the model was built with mtp_enabled={self.mtp_enabled}; the "
- "draft path and the weights it loads are decided at construction"
- )
- # Keyed on the manager, not a first-forward flag: capacity estimation swaps the
- # manager under the same module (see `_bind_cache_manager`).
- if self._bound is None or attn_metadata.kv_cache_manager is not self._bound.cm:
- self._bind_cache_manager(attn_metadata, position_ids)
- target = self._select_target(attn_metadata)
- if target is self._prefill:
- target.prepare(attn_metadata, input_ids, position_ids, vision, spec_metadata)
- else:
- assert vision is None, (
- "a generation-only step carries vision rows; the executor attaches media "
- "to context requests only"
- )
- self._step_begun = True
-
- def forward(
- self,
- attn_metadata: AttentionMetadata,
- input_ids: torch.IntTensor | None = None,
- position_ids: torch.IntTensor | None = None,
- inputs_embeds: torch.FloatTensor | None = None,
- lora_params: dict | None = None,
- **kwargs,
- ):
- """Run the step `begin_step` prepared on the target it chose.
-
- Under a breakable prefill CUDA graph this is the body the engine captures for a
- step with context rows; a replayed step never enters it, which is why nothing here
- may read the step from the host. Returns the logits input, or under speculation the
- decode target's `(final, hc)` pair and the prefill target's one packed tensor
- (`split_spec_output`)."""
- assert self._step_begun, (
- "the shell calls begin_step before every forward; this forward has no prepared step"
- )
- self._step_begun = False
- return self._select_target(attn_metadata).forward(
- attn_metadata, input_ids=input_ids, position_ids=position_ids, **kwargs
- )
-
- def split_spec_output(self, output) -> tuple[torch.Tensor, torch.Tensor]:
- """`(final, hc)` out of what `forward` returned under speculation: `final` is the
- `[T, hidden]` mix `lm_head` reads and `hc` the `[T, hc_count * hidden]` stream the
- MTP module reads. The decode target returns the pair; the prefill target returns one
- `[T, hidden + stream]` tensor, of which the two are column views."""
- if isinstance(output, tuple):
- return output
- final, hc = torch.split(output, [self.hidden, self.stream], -1)
- return final, hc
-
-
-class _LayerWeights(NamedTuple):
- """One trunk layer's weights as the views the ops read (GEMM operands are `.t()`
- views). Attention fields are None on a Gated-DeltaNet layer and vice versa."""
-
- ahc_norm: torch.Tensor
- ahc_first: torch.Tensor
- ahc_up: torch.Tensor
- mhc_norm: torch.Tensor
- mhc_first: torch.Tensor
- mhc_up: torch.Tensor
- attn_in: torch.Tensor | None
- o: torch.Tensor | None
- q_norm: torch.Tensor | None
- k_norm: torch.Tensor | None
- idx_qn: torch.Tensor | None
- idx_kn: torch.Tensor | None
- in_proj: torch.Tensor | None
- out: torch.Tensor | None
- conv_w: torch.Tensor | None
- a_log: torch.Tensor | None
- dt_bias: torch.Tensor | None
- gnorm: torch.Tensor | None
- router: torch.Tensor
- moe_gscale: torch.Tensor
- fc1_w: torch.Tensor
- fc1_s: torch.Tensor
- fc2_w: torch.Tensor
- fc2_s: torch.Tensor
- # trtllm-gen scalars; None on CuteDSL.
- out1_scale: torch.Tensor | None
- out1_gate_scale: torch.Tensor | None
- out2_scale: torch.Tensor | None
- # CuteDSL scalars; None on trtllm-gen.
- alpha1: torch.Tensor | None
- alpha2: torch.Tensor | None
- global_sf: torch.Tensor | None
- # Under trtllm-gen, the second expert copy in the CuteDSL layout and its scalars,
- # which prefill-sized launches run on (`_CUTE_DSL_MOE_MIN_ROWS`); None on CuteDSL and
- # without the copy (`moe_prefill_copy`).
- fc1_w_cd: torch.Tensor | None
- fc1_s_cd: torch.Tensor | None
- fc2_w_cd: torch.Tensor | None
- fc2_s_cd: torch.Tensor | None
- alpha1_cd: torch.Tensor | None
- alpha2_cd: torch.Tensor | None
- global_sf_cd: torch.Tensor | None
- sh_gu: torch.Tensor
- sh_dn: torch.Tensor
- sh_gate: torch.Tensor
- # `[gate | up | gate logit | pad]` (`weights._cat_rows`): `sh_gu` and `sh_gate` are its
- # column blocks.
- sh_gug: torch.Tensor
- # FP8 copies `(codes_t, scale)` of the dense projections (`_DENSE_FP8_ENV`); None when off
- # and on the other layer kind.
- attn_in8: tuple | None = None
- o8: tuple | None = None
- in_proj8: tuple | None = None
- out8: tuple | None = None
- # Packed FP8 decode copies (`fp8_gemv`, 128x128 block scales) under the same mode.
- attn_in_g: object | None = None
- o_g: object | None = None
- in_proj_g: object | None = None
- out_g: object | None = None
- sh_gu_g: object | None = None
- sh_dn_g: object | None = None
- # The GDN in_proj's NVFP4 decode copy (`_GDN_INPROJ_NVFP4_ENV`, `nvfp4_gemv.SplitNvfp4Weight`).
- in_proj_nv: object | None = None
- # `mid_rows_mm` views (`MidWeight`) under the same mode, for `_MID_MM_ROWS` decode rows.
- attn_in_m: object | None = None
- o_m: object | None = None
- in_proj_m: object | None = None
- out_m: object | None = None
- sh_gu_m: object | None = None
- sh_dn_m: object | None = None
- router_m: object | None = None
- # `1 + ahc_norm` / `1 + mhc_norm` in fp32, the rstd mixers' operand
- # (`qwen4_exp_hc_up_gate_mean_rstd`); set by `PrefillTarget` only.
- ahc_norm_w1: torch.Tensor | None = None
- mhc_norm_w1: torch.Tensor | None = None
-
-
-def _layer_views(core: Qwen4ExpTp1ModelingV2) -> list[_LayerWeights]:
- """Every trunk layer's `_LayerWeights`, read off the realized weights."""
- w = core.w
- # The expert stacks are uint8 storage; each backend reads the packed nibbles and the
- # block scales as its own dtypes.
- if core.trunk_moe_backend == "CUTEDSL":
- w_dtype, s_dtype = torch.float4_e2m1fn_x2, torch.uint8
- else:
- w_dtype, s_dtype = torch.uint8, torch.float8_e4m3fn
- layers = []
- for i in range(core.num_layers):
- attn: dict = dict(
- attn_in=None,
- o=None,
- q_norm=None,
- k_norm=None,
- idx_qn=None,
- idx_kn=None,
- )
- gdn: dict = dict(
- in_proj=None,
- out=None,
- conv_w=None,
- a_log=None,
- dt_bias=None,
- gnorm=None,
- )
- if core.is_attn[i]:
- attn = dict(
- attn_in=w[f"l{i}_attn_in_w"].t(),
- o=w[f"l{i}_o_w"].t(),
- q_norm=w[f"l{i}_q_norm"],
- k_norm=w[f"l{i}_k_norm"],
- idx_qn=w[f"l{i}_idx_qn"],
- idx_kn=w[f"l{i}_idx_kn"],
- )
- else:
- a_log, dt_bias = core._gdn_gating[i]
- gdn = dict(
- in_proj=w[f"l{i}_in_proj_w"].t(),
- out=w[f"l{i}_out_w"].t(),
- # [conv_dim, 1, width] -> [conv_dim, width]: a view (unit width stride).
- conv_w=w[f"l{i}_conv_w"].reshape(core.conv_dim, core.conv_kernel),
- a_log=a_log,
- dt_bias=dt_bias,
- gnorm=w[f"l{i}_gnorm"],
- )
- scalars = core._moe_scalars[i]
- cute_dsl_copy: dict = dict(
- fc1_w_cd=None,
- fc1_s_cd=None,
- fc2_w_cd=None,
- fc2_s_cd=None,
- alpha1_cd=None,
- alpha2_cd=None,
- global_sf_cd=None,
- )
- if core._moe_scalars_cd is not None:
- scalars_cd = core._moe_scalars_cd[i]
- cute_dsl_copy = dict(
- fc1_w_cd=w[f"l{i}_fc1_w_cd"],
- fc1_s_cd=w[f"l{i}_fc1_s_cd"],
- fc2_w_cd=w[f"l{i}_fc2_w_cd"],
- fc2_s_cd=w[f"l{i}_fc2_s_cd"],
- alpha1_cd=scalars_cd["alpha1"],
- alpha2_cd=scalars_cd["alpha2"],
- global_sf_cd=scalars_cd["global_sf"],
- )
- layers.append(
- _LayerWeights(
- ahc_norm=w[f"l{i}_ahc_norm"],
- ahc_first=core._hc_first[(i, "ahc")].t(),
- ahc_up=w[f"l{i}_ahc_up"].t(),
- mhc_norm=w[f"l{i}_mhc_norm"],
- mhc_first=core._hc_first[(i, "mhc")].t(),
- mhc_up=w[f"l{i}_mhc_up"].t(),
- **attn,
- **gdn,
- router=w[f"l{i}_router"].t(),
- moe_gscale=scalars["global_scale"],
- fc1_w=w[f"l{i}_fc1_w"].view(w_dtype),
- fc1_s=w[f"l{i}_fc1_s"].view(s_dtype),
- fc2_w=w[f"l{i}_fc2_w"].view(w_dtype),
- fc2_s=w[f"l{i}_fc2_s"].view(s_dtype),
- out1_scale=scalars.get("output1_scale_scalar"),
- out1_gate_scale=scalars.get("output1_scale_gate_scalar"),
- out2_scale=scalars.get("output2_scale_scalar"),
- alpha1=scalars.get("alpha1"),
- alpha2=scalars.get("alpha2"),
- global_sf=scalars.get("global_sf"),
- **cute_dsl_copy,
- sh_gu=w[f"l{i}_sh_gu"][: 2 * core.shared_inter].t(),
- sh_dn=w[f"l{i}_sh_dn"].t(),
- sh_gate=w[f"l{i}_sh_gu"][2 * core.shared_inter : 2 * core.shared_inter + 1].t(),
- sh_gug=w[f"l{i}_sh_gu"].t(),
- attn_in8=getattr(core, "_dense8", {}).get(f"l{i}_attn_in_w"),
- o8=getattr(core, "_dense8", {}).get(f"l{i}_o_w"),
- in_proj8=getattr(core, "_dense8", {}).get(f"l{i}_in_proj_w"),
- out8=getattr(core, "_dense8", {}).get(f"l{i}_out_w"),
- attn_in_g=getattr(core, "_gemv8", {}).get(f"l{i}_attn_in_w"),
- o_g=getattr(core, "_gemv8", {}).get(f"l{i}_o_w"),
- in_proj_g=getattr(core, "_gemv8", {}).get(f"l{i}_in_proj_w"),
- out_g=getattr(core, "_gemv8", {}).get(f"l{i}_out_w"),
- sh_gu_g=getattr(core, "_gemv8", {}).get(f"l{i}_sh_gu"),
- sh_dn_g=getattr(core, "_gemv8", {}).get(f"l{i}_sh_dn"),
- in_proj_nv=getattr(core, "_gemv_nv", {}).get(f"l{i}_in_proj_w"),
- attn_in_m=getattr(core, "_mid", {}).get(f"l{i}_attn_in_w"),
- o_m=getattr(core, "_mid", {}).get(f"l{i}_o_w"),
- in_proj_m=getattr(core, "_mid", {}).get(f"l{i}_in_proj_w"),
- out_m=getattr(core, "_mid", {}).get(f"l{i}_out_w"),
- sh_gu_m=getattr(core, "_mid", {}).get(f"l{i}_sh_gu"),
- sh_dn_m=getattr(core, "_mid", {}).get(f"l{i}_sh_dn"),
- router_m=getattr(core, "_mid", {}).get(f"l{i}_router"),
- )
- )
- return layers
-
-
-class _StackWeights(NamedTuple):
- """The weights outside the layer loop: the stack-final mixer (this checkpoint's
- final norm) and the PLE side path."""
-
- fhc_norm: torch.Tensor
- fhc_first: torch.Tensor
- fhc_up: torch.Tensor
- ple_key: torch.Tensor
- ple_val: torch.Tensor
- ple_nk_w: torch.Tensor
- ple_nq_w: torch.Tensor
- ple_nc_w: torch.Tensor
- ple_conv_w: torch.Tensor
- ple_mult: torch.Tensor
- ple_vocab: torch.Tensor
- ple_offs: torch.Tensor
- ple_tscale: torch.Tensor
- # `1 + fhc_norm` in fp32 (see `_LayerWeights.ahc_norm_w1`); set by `PrefillTarget` only.
- fhc_norm_w1: torch.Tensor | None = None
-
-
-def _norm_w1(norm: torch.Tensor) -> torch.Tensor:
- """A grouped norm's `(1 + w)` in fp32, contiguous: what
- `qwen4_exp_hc_up_gate_mean_rstd` multiplies the scaled residual by (the `1 + w`
- `qwen4_exp_hc_inject_norm` forms per element, bitwise)."""
- return (1.0 + norm.float()).contiguous()
-
-
-def _stack_views(core: Qwen4ExpTp1ModelingV2) -> _StackWeights:
- w = core.w
- return _StackWeights(
- fhc_norm=w["fhc_norm"],
- fhc_first=core._hc_first["fhc"].t(),
- fhc_up=w["fhc_up"].t(),
- ple_key=w["ple_key_w"].t(),
- ple_val=w["ple_val_w"].t(),
- # The three PLE norm weights whole: the fused tail normalizes each stream inside
- # one program.
- ple_nk_w=w["ple_nk"],
- ple_nq_w=w["ple_nq"],
- ple_nc_w=w["ple_nc"],
- ple_conv_w=w["ple_conv_w"],
- ple_mult=w["ple_mult"],
- ple_vocab=w["ple_vocab"],
- ple_offs=w["ple_offs"],
- # `dequantize_e4m3_per_tensor` needs the scale at the input's rank (3), all 1s.
- ple_tscale=w["ple_tscale"].reshape(1, 1, 1),
- )
-
-
-class _MtpWeights(NamedTuple):
- """The MTP module's weights (keyed at its layer index, `num_hidden_layers`)."""
-
- pre_e_norm: torch.Tensor
- pre_h_norm: tuple
- fc_e: torch.Tensor
- fc_h: torch.Tensor
- ahc_norm: tuple
- ahc_down: torch.Tensor
- ahc_up: torch.Tensor
- ahc_inject: torch.Tensor
- mhc_norm: tuple
- mhc_down: torch.Tensor
- mhc_up: torch.Tensor
- mhc_inject: torch.Tensor
- q: torch.Tensor
- k: torch.Tensor
- v: torch.Tensor
- o: torch.Tensor
- q_norm: torch.Tensor
- k_norm: torch.Tensor
- idx: torch.Tensor
- idx_qn: torch.Tensor
- idx_kn: torch.Tensor
- router: torch.Tensor
- fc1_w: torch.Tensor
- fc1_s: torch.Tensor
- fc2_w: torch.Tensor
- fc2_s: torch.Tensor
- sh_gu: torch.Tensor
- sh_dn: torch.Tensor
- sh_gate: torch.Tensor
- fhc_norm: tuple
- fhc_down: torch.Tensor
- fhc_up: torch.Tensor
- # With `_MTP_MOE_NVFP4_ENV`, `fc1_w` .. `fc2_s` are the NVFP4 copy in the trunk's
- # trtllm-gen layout and these are its static activation global scale `[1]` and the
- # runner's three per-expert scalars; None on the FP8 block-scale path.
- moe_gscale: torch.Tensor | None = None
- out1_scale: torch.Tensor | None = None
- out1_gate_scale: torch.Tensor | None = None
- out2_scale: torch.Tensor | None = None
-
-
-def _mtp_views(core: Qwen4ExpTp1ModelingV2) -> _MtpWeights:
- w = core.w
- p = f"l{core.mtp_layer_idx}_mtp_"
- return _MtpWeights(
- pre_e_norm=w[f"{p}pre_e_norm"],
- pre_h_norm=core._norm_slices(f"{p}pre_h_norm"),
- fc_e=w[f"{p}fc_e"].t(),
- fc_h=w[f"{p}fc_h"].t(),
- ahc_norm=core._norm_slices(f"{p}ahc_norm"),
- ahc_down=w[f"{p}ahc_down"],
- ahc_up=w[f"{p}ahc_up"],
- ahc_inject=w[f"{p}ahc_inject"],
- mhc_norm=core._norm_slices(f"{p}mhc_norm"),
- mhc_down=w[f"{p}mhc_down"],
- mhc_up=w[f"{p}mhc_up"],
- mhc_inject=w[f"{p}mhc_inject"],
- q=w[f"{p}q_w"].t(),
- k=w[f"{p}k_w"].t(),
- v=w[f"{p}v_w"].t(),
- o=w[f"{p}o_w"].t(),
- q_norm=w[f"{p}q_norm"],
- k_norm=w[f"{p}k_norm"],
- idx=w[f"{p}idx_w"].t(),
- idx_qn=w[f"{p}idx_qn"],
- idx_kn=w[f"{p}idx_kn"],
- router=w[f"{p}router"].t(),
- fc1_w=w[f"{p}fc1_w"],
- fc1_s=w[f"{p}fc1_s"],
- fc2_w=w[f"{p}fc2_w"],
- fc2_s=w[f"{p}fc2_s"],
- sh_gu=w[f"{p}sh_gu"][: 2 * core.shared_inter].t(),
- sh_dn=w[f"{p}sh_dn"].t(),
- sh_gate=w[f"{p}sh_gu"][2 * core.shared_inter : 2 * core.shared_inter + 1].t(),
- fhc_norm=core._norm_slices(f"{p}fhc_norm"),
- fhc_down=w[f"{p}fhc_down"],
- fhc_up=w[f"{p}fhc_up"],
- )
-
-
-def _mtp_moe_nvfp4(core: Qwen4ExpTp1ModelingV2) -> dict:
- """The draft layer's routed experts requantized to NVFP4 (`_MTP_MOE_NVFP4_ENV`), as the
- `_MtpWeights` fields `_mtp_routed_nvfp4` reads; frees the FP8 copies.
-
- The runner's scalars follow `_moe_layer_scalars` (trtllm-gen) with `1 / g_in` and
- `1 / g_mid` standing for the checkpoint's `input_scale`s and `1 / g_w` for its
- `weight_scale_2`: `out1_gate = 1 / (g_in * g_w1[e])`, `out1 = out1_gate * g_mid`,
- `out2 = 1 / (g_mid * g_w2[e])`. The activation scales are static
- (`_MTP_MOE_NVFP4_FC1_MARGIN` / `_FC2_MARGIN`)."""
- w = core.w
- p = f"l{core.mtp_layer_idx}_mtp_"
- keys = [f"{p}{k}" for k in ("fc1_w", "fc1_s", "fc2_w", "fc2_s")]
- fc1_w, fc1_s, fc2_w, fc2_s, g_w1, g_w2 = _weights.mtp_experts_to_nvfp4(*(w[k] for k in keys))
- # Nothing else reads the FP8 experts: `_mtp_views` handed them to `_MtpWeights` only, and
- # the caller replaces those fields with the NVFP4 copy.
- for k in keys:
- w[k].data = w[k].data.new_empty(0)
- torch.cuda.empty_cache()
- trunk = range(core.num_layers)
- in_amax = max(float(w[f"l{i}_e_isc1"].max()) for i in trunk) * _weights.NVFP4_AMAX_TO_GLOBAL
- mid_amax = max(float(w[f"l{i}_e_isc2"].max()) for i in trunk) * _weights.NVFP4_AMAX_TO_GLOBAL
- # `1 / g` of each activation quantization.
- inv_g_in = _MTP_MOE_NVFP4_FC1_MARGIN * in_amax / _weights.NVFP4_AMAX_TO_GLOBAL
- inv_g_mid = _MTP_MOE_NVFP4_FC2_MARGIN * mid_amax / _weights.NVFP4_AMAX_TO_GLOBAL
- gate1 = (inv_g_in / g_w1).contiguous()
- print(
- f"draft MoE: {len(g_w1)} experts requantized FP8 -> NVFP4 ({_MTP_MOE_NVFP4_ENV}=1), "
- f"static activation amax FC1 {inv_g_in * _weights.NVFP4_AMAX_TO_GLOBAL:.1f}, "
- f"FC2 {inv_g_mid * _weights.NVFP4_AMAX_TO_GLOBAL:.1f}",
- flush=True,
- )
- return dict(
- fc1_w=fc1_w,
- fc1_s=fc1_s.view(torch.float8_e4m3fn),
- fc2_w=fc2_w,
- fc2_s=fc2_s.view(torch.float8_e4m3fn),
- moe_gscale=torch.full((1,), 1.0 / inv_g_in, dtype=torch.float32, device=fc1_w.device),
- out1_scale=(gate1 / inv_g_mid).contiguous(),
- out1_gate_scale=gate1,
- out2_scale=(inv_g_mid / g_w2).contiguous(),
- )
-
-
-def _mtp_routed_nvfp4(core: Qwen4ExpTp1ModelingV2, M: _MtpWeights, x: torch.Tensor):
- """The draft layer's routed experts on their NVFP4 copy: the activation quantized at the
- static global scale (linear block scales, as the trunk's MoE sites write them), then the
- trunk's logits-routed trtllm-gen call (its Renormalize routing is the FP8 path's
- `renorm_moe_routing_op`). Returns the combined `[T, H]` bf16 rows."""
- logits = cublas_mm(x, M.router)
- xq, xsf = fp4_quantize(x.contiguous(), M.moe_gscale, _SF_VEC, False, False)
- return fp4_block_scale_moe_renorm_routed(
- logits,
- xq,
- xsf.view(torch.float8_e4m3fn),
- M.fc1_w,
- M.fc1_s,
- M.fc2_w,
- M.fc2_s,
- M.out1_scale,
- M.out1_gate_scale,
- M.out2_scale,
- core.num_experts,
- core.topk,
- core.moe_inter,
- )
-
-
-def _mtp_moe_nvfp4_block(core: Qwen4ExpTp1ModelingV2, M: _MtpWeights, x: torch.Tensor):
- """The draft layer's MoE on the NVFP4 experts (`_mtp_routed_nvfp4`) plus its bf16 shared
- expert, combined as on the FP8 path."""
- routed = _mtp_routed_nvfp4(core, M, x)
- shared = cublas_mm(flashinfer_silu_and_mul(cublas_mm(x, M.sh_gu)), M.sh_dn)
- gate_logits = cublas_mm(x, M.sh_gate)
- return fused_sigmoid_gate_mul_add(routed, gate_logits, shared)
-
-
-class _MtpFused(NamedTuple):
- """The draft layer's operands in the trunk's fused layouts, built after load: the
- `[q | k | v | gate | index | pad]` input projection (`weights._attn_in_w`), each
- hyper-connection's `[down; block_inject; 0]` first projection (`_fuse_hc_first_projections`)
- with its up projection and whole grouped-norm weight, and a unit injection weight (the
- front end's per-stream `x_h + x_e` as the first site's injection). The embedding branch
- enters all four streams; nothing in the checkpoint or config says so, and entering stream
- 0 alone would leave the text bit-correct with every draft rejected, which only the
- acceptance gate detects."""
-
- attn_in: torch.Tensor
- ahc_norm: torch.Tensor
- ahc_first: torch.Tensor
- ahc_up: torch.Tensor
- mhc_norm: torch.Tensor
- mhc_first: torch.Tensor
- mhc_up: torch.Tensor
- fhc_norm: torch.Tensor
- fhc_first: torch.Tensor
- fhc_up: torch.Tensor
- unit_inject: torch.Tensor
-
-
-def _mtp_fused_views(core: Qwen4ExpTp1ModelingV2) -> _MtpFused:
- w = core.w
- p = f"l{core.mtp_layer_idx}_mtp_"
- attn_in = _weights._attn_in_w((w[f"{p}q_w"], w[f"{p}k_w"], w[f"{p}v_w"], w[f"{p}idx_w"]), core)
- assert attn_in.shape == (core.attn_in_width, core.hidden), attn_in.shape
- return _MtpFused(
- attn_in=attn_in.t(),
- ahc_norm=w[f"{p}ahc_norm"],
- ahc_first=core._hc_first["mtp_ahc"].t(),
- ahc_up=w[f"{p}ahc_up"].t(),
- mhc_norm=w[f"{p}mhc_norm"],
- mhc_first=core._hc_first["mtp_mhc"].t(),
- mhc_up=w[f"{p}mhc_up"].t(),
- fhc_norm=w[f"{p}fhc_norm"],
- fhc_first=core._hc_first["mtp_fhc"].t(),
- fhc_up=w[f"{p}fhc_up"].t(),
- unit_inject=torch.ones(
- [core.max_num_tokens, core.hc_count], dtype=attn_in.dtype, device=attn_in.device
- ),
- )
-
-
-def _select(core, scores, row_starts, row_ends, topk_indices, live_cols=None) -> None:
- """Select each row's top-`block_topk` groups into `topk_indices` with `qsa_topk_select`,
- then rewrite the row canonically, in place.
-
- The radix engine compacts its picks through a shared-memory atomic, so the order is
- launch-dependent, and where the tie class at the budget's last score straddles the
- budget -- every zero-score group of a row with fewer positive scores than the budget,
- on this checkpoint's relu scores -- which tied groups are kept can be too. The
- expansion keeps the row's order and the sparse GQA reduces in it, so a raw row past the
- dense threshold is not reproducible run to run. `qsa_canonical_selection` rewrites it
- into the window's canonical selection -- the groups strictly above the budget's last
- score, then the highest-index (most recent) ties, ascending, `-1` last -- a function of
- the scores alone, and the engine's own set in ascending order wherever the budget's
- last score is untied. A canonical row is prefix-packed, which the expansion and the
- IndexShare capture and reuse require.
-
- `index_topk` is the block budget: the entry's default (2048) is 4x the buffer width and
- the op ignores the row stride, so it mis-writes silently.
-
- The rewrite reads every column it is given, and `scores` is as wide as the block table,
- not the live window. A caller that knows a host-side bound on the rows' windows passes
- it as `live_cols`: the rewrite then reads only those columns and returns the same rows."""
- qsa_topk_select(scores, row_starts, row_ends, topk_indices, index_topk=core.block_topk)
- if live_cols is not None:
- scores = scores[:, :live_cols]
- qsa_canonical_selection(scores, row_starts, row_ends, topk_indices)
-
-
-def _store_kv(core, kv, k, v, request_indices, logical_positions, block_table, tokens_per_block):
- """Append this step's K/V rows to the layer's paged pool: a plain move over a bf16 or fp8
- pool (an fp8 pool takes e4m3 unscaled, as the dense op does at the unit KV scales), the
- quantizing store over an NVFP4 one, at the same second-level scales as the dense op."""
- if kv.k_scale is None:
- qsa_paged_kv_store(
- k, v, kv.k, kv.v, request_indices, logical_positions, block_table, tokens_per_block
- )
- return
- qsa_paged_kv_store_nvfp4(
- k,
- v,
- kv.k,
- kv.v,
- kv.k_scale,
- kv.v_scale,
- request_indices,
- logical_positions,
- block_table,
- tokens_per_block,
- core.kv_scale_orig_quant,
- )
-
-
-def _sparse_gqa(
- core,
- kv,
- q,
- block_table,
- selected_tokens,
- request_indices,
- tokens_per_block,
- query_positions,
- compress_ratio,
- output_gate=None,
- rows_per_request=1,
-):
- """Sparse GQA over the layer's paged pool (dequantizing an NVFP4 one). The prefill pair
- `query_positions` / `compress_ratio` is a bound on a context row's scan, or None, None.
- `output_gate` (NVFP4 pool, no prefill pair only) returns the gated output.
- `rows_per_request` (NVFP4 pool, no prefill pair only): the rows come that many consecutive
- rows a request, a verify step's generation rows; 1 for any other batch."""
- if kv.k_scale is None:
- assert output_gate is None, "the output gate is fused on the NVFP4 decode-row path only"
- return qsa_paged_sparse_gqa(
- q=q,
- k_cache=kv.k,
- v_cache=kv.v,
- block_table=block_table,
- selected_tokens=selected_tokens,
- request_indices=request_indices,
- tokens_per_block=tokens_per_block,
- softmax_scale=core.softmax_scale,
- query_positions=query_positions,
- compress_ratio=compress_ratio,
- )
- if query_positions is None:
- # Rows without the prefill pair (decode, draft, a mixed step's generation rows): the
- # CuTe split-K + merge, with the output gate fused when the caller hands it over.
- if rows_per_request > 1:
- # A verify step's rows, `rows_per_request` a request: the request-grouped entry
- # loads each selected K/V token once for a group of a request's rows (from eight
- # groups up; below that it is the per-row kernel).
- return qsa_paged_sparse_gqa_nvfp4_verify(
- q,
- kv.k,
- kv.v,
- kv.k_scale,
- kv.v_scale,
- core.kv_scale_quant_orig,
- block_table,
- selected_tokens,
- request_indices,
- tokens_per_block,
- core.softmax_scale,
- rows_per_request,
- output_gate=output_gate,
- )
- return qsa_paged_sparse_gqa_nvfp4_decode(
- q,
- kv.k,
- kv.v,
- kv.k_scale,
- kv.v_scale,
- core.kv_scale_quant_orig,
- block_table,
- selected_tokens,
- request_indices,
- tokens_per_block,
- core.softmax_scale,
- output_gate=output_gate,
- )
- assert output_gate is None, "the output gate is fused on the decode-row path only"
- return qsa_paged_sparse_gqa_nvfp4(
- q=q,
- k_cache=kv.k,
- v_cache=kv.v,
- k_scale_cache=kv.k_scale,
- v_scale_cache=kv.v_scale,
- kv_scale_quant_orig=core.kv_scale_quant_orig,
- block_table=block_table,
- selected_tokens=selected_tokens,
- request_indices=request_indices,
- tokens_per_block=tokens_per_block,
- softmax_scale=core.softmax_scale,
- query_positions=query_positions,
- compress_ratio=compress_ratio,
- )
-
-
-def _attn_in_heads(core, proj):
- """The main q, k and v heads of the fused attention projection's output `proj`
- `[T, attn_in_width]` (`weights._attn_in_w`), as `[T, heads, head_dim]` views."""
- q, k, v, _ = torch.split(
- torch.reshape(proj, [proj.shape[0], core.attn_in_width // core.head_dim, core.head_dim]),
- [core.heads_q, core.heads_kv, core.heads_kv, core.attn_in_tail_heads - core.heads_kv],
- 1,
- )
- return q, k, v
-
-
-def _cute_dsl_routed(core, L, logits, xq, xsf, routed_out) -> torch.Tensor:
- """The trunk's routed experts on the CuteDSL pipeline: Renormalize routing, the sort
- into per-expert row tiles, FC1 gathering its rows from the unpermuted NVFP4 activation
- `xq` / `xsf` (linear scales) with SwiGLU and the requantize to the layer's one FC2 input
- scale fused, FC2, and the routing-weighted unpermute in a fixed slot order, which
- writes every row of `routed_out` `[T, hidden]` and returns it."""
- # fp32 weights: the dtype the unpermute scales by.
- topk_ids, topk_w = renorm_moe_routing_op(logits, topk=core.topk, output_dtype=torch.float32)
- # The sort pads each expert's (token, expert) pairs to whole tiles and emits the tile
- # tables both GEMMs walk.
- tile_expert, tile_limit, expanded_to_permuted, permuted_to_expanded, _, live_tiles = moe_sort(
- token_selected_experts=topk_ids,
- token_final_scales=topk_w,
- num_experts=core.num_experts,
- top_k=core.topk,
- local_expert_offset=0,
- local_num_experts=core.num_experts,
- tile_tokens_dim=_CUTEDSL_TILE_SIZE,
- )
- h, h_sf = cute_dsl_nvfp4_gather_grouped_gemm_act_fusion(
- input=xq.view(torch.float4_e2m1fn_x2),
- weight=L.fc1_w,
- input_scale=xsf.view(torch.uint8).view(xq.shape[0], core.hidden // _SF_VEC),
- weight_scale=L.fc1_s,
- alpha=L.alpha1,
- tile_idx_to_group_idx=tile_expert,
- tile_idx_to_mn_limit=tile_limit,
- permuted_idx_to_expanded_idx=permuted_to_expanded,
- num_non_exiting_tiles=live_tiles,
- global_sf=L.global_sf,
- num_experts=core.num_experts,
- top_k=core.topk,
- num_local_experts=core.num_experts,
- local_expert_offset=0,
- tile_size=_CUTEDSL_TILE_SIZE,
- scaling_vector_size=_SF_VEC,
- activation_type=int(ActivationType.Swiglu),
- )
- # FC2 writes one bf16 row per permuted row.
- y = cute_dsl_nvfp4_grouped_gemm(
- input=h,
- weight=L.fc2_w,
- input_scale=h_sf,
- weight_scale=L.fc2_s,
- alpha=L.alpha2,
- tile_idx_to_group_idx=tile_expert,
- num_non_exiting_tiles=live_tiles,
- num_experts=core.num_experts,
- top_k=core.topk,
- num_local_experts=core.num_experts,
- local_expert_offset=0,
- tile_size=_CUTEDSL_TILE_SIZE,
- output_dtype=torch.bfloat16,
- scaling_vector_size=_SF_VEC,
- )
- return moe_unpermute_inplace(
- permuted_input=y,
- output=routed_out,
- expanded_idx_to_permuted_idx=expanded_to_permuted,
- topk_scales=topk_w,
- )
-
-
-def _nvfp4_context_groups(pages: list, rows: list, max_pages: int) -> list:
- """Consecutive context requests in groups of at most `max_pages` pages, a request with
- more pages alone: `(first request, end request, first row, end row)` a group, where
- request `r` has `pages[r]` pages and the next `rows[r]` rows."""
- groups = []
- first = first_row = row = listed = 0
- for r, (n_pages, n_rows) in enumerate(zip(pages, rows)):
- if r > first and listed + n_pages > max_pages:
- groups.append((first, r, first_row, row))
- first, first_row, listed = r, row, 0
- listed += n_pages
- row += n_rows
- if pages:
- groups.append((first, len(pages), first_row, row))
- return groups
-
-
-def _nvfp4_context_plan(
- core, md, req_idx, block_table, tpb: int, one_row_per_request: bool = False
-) -> tuple:
- """The pages a step's context requests read, for the bf16 row-tile pass over an NVFP4
- pool: request `r` lists its first `ceil(kv_len / tpb)` table pages. The requests come in
- groups that fit `_NVFP4_CONTEXT_SCRATCH_BYTES` (`_nvfp4_context_groups`); within a group
- request `r` takes scratch pages `offsets[r]` on. Returns `(groups, table, k_scratch,
- v_scratch)`: per group `(first row, end row, slots, counts, offsets)`, the operands
- `qsa_nvfp4_pages_to_bf16` fills per layer and group, the rows those of its requests (one
- a request with `one_row_per_request`, else their step rows); `table` (all -1 here) gets
- the requests' slot rows, the scratch pools `[pages, Hkv, tpb, D]` bf16 the largest group."""
- num_contexts = md.num_contexts
- device = block_table.device
- # The host's and the device's KV lengths are the same per-step values: the host's
- # group the requests and size the scratch, the device's place the pages.
- pages = ((md.kv_lens_runtime[:num_contexts] + (tpb - 1)) // tpb).tolist()
- rows = [1] * num_contexts if one_row_per_request else md.seq_lens[:num_contexts].tolist()
- page_bytes = 2 * core.heads_kv * tpb * core.head_dim * 2 # K and V, bf16
- groups = _nvfp4_context_groups(pages, rows, max(1, _NVFP4_CONTEXT_SCRATCH_BYTES // page_bytes))
- lens = md.seq_lens_cuda[:num_contexts].long()
- slots = req_idx[torch.cumsum(lens, 0) - lens]
- counts = torch.div(
- md.kv_lens_cuda_runtime[:num_contexts].long() + (tpb - 1), tpb, rounding_mode="floor"
- )
- planned = []
- for first, end, first_row, end_row in groups:
- group_counts = counts[first:end]
- planned.append(
- (
- first_row,
- end_row,
- slots[first:end],
- group_counts,
- torch.cumsum(group_counts, 0) - group_counts,
- )
- )
- listed = max(sum(pages[first:end]) for first, end, _, _ in groups)
- table = torch.full_like(block_table, -1)
- shape = (listed, core.heads_kv, tpb, core.head_dim)
- k_scratch = torch.empty(shape, dtype=torch.bfloat16, device=device)
- v_scratch = torch.empty(shape, dtype=torch.bfloat16, device=device)
- return planned, table, k_scratch, v_scratch
-
-
-def _sparse_gqa_nvfp4_context(
- core,
- kv,
- plan,
- q,
- block_table,
- selected,
- req_idx,
- logical,
- tpb: int,
- num_ctx: int,
- out,
- rows_per_request: int = 1,
-) -> torch.Tensor:
- """Sparse GQA of a step with context rows over an NVFP4 pool, into `out` `[rows,
- q_width]`. With a plan (`_nvfp4_context_plan`) the first `num_ctx` rows, the context
- rows, take the row-tile pass over the plan's scratch, filled here from this layer's
- pool a group of requests at a time; the other rows, and every row without a plan, go
- split-K over the pool itself (with the prefill pair the entry would take its per-row
- prefill kernel). The generation
- rows after the context rows come `rows_per_request` a request: more than one (a verify
- step's) takes the request-grouped entry for them on their own."""
- rows = q.shape[0]
- head = 0
- if plan is not None:
- groups, table, k_scratch, v_scratch = plan
- assert groups[-1][1] == num_ctx, (
- f"the plan's groups end at row {groups[-1][1]}, not at the {num_ctx} context rows"
- )
- for first_row, end_row, slots, counts, offsets in groups:
- qsa_nvfp4_pages_to_bf16(
- kv.k,
- kv.v,
- kv.k_scale,
- kv.v_scale,
- core.kv_scale_quant_orig,
- block_table,
- slots,
- counts,
- offsets,
- table,
- k_scratch,
- v_scratch,
- )
- context = qsa_sparse_gqa_rowtile(
- q=q[first_row:end_row],
- k_cache=k_scratch,
- v_cache=v_scratch,
- block_table=table,
- selected_tokens=selected[first_row:end_row],
- request_indices=req_idx[first_row:end_row],
- tokens_per_block=tpb,
- softmax_scale=core.softmax_scale,
- query_positions=logical[first_row:end_row],
- compress_ratio=core.compress_ratio,
- )
- out[first_row:end_row].copy_(
- torch.reshape(context, [end_row - first_row, core.q_width])
- )
- head = num_ctx
- elif (
- rows_per_request > 1
- and 0 < num_ctx < rows
- and _verify_takes_grouped(
- q[num_ctx:], kv.k, block_table, selected[num_ctx:], tpb, rows_per_request
- )
- ):
- # Without a plan the context rows would share the generation rows' split-K launch;
- # when those are enough for the request-grouped kernel the context rows go on their own.
- context = _sparse_gqa(
- core,
- kv,
- q=q[:num_ctx],
- block_table=block_table,
- selected_tokens=selected[:num_ctx],
- request_indices=req_idx[:num_ctx],
- tokens_per_block=tpb,
- query_positions=None,
- compress_ratio=None,
- )
- out[:num_ctx].copy_(torch.reshape(context, [num_ctx, core.q_width]))
- head = num_ctx
- if head < rows:
- rest = _sparse_gqa(
- core,
- kv,
- q=q[head:],
- block_table=block_table,
- selected_tokens=selected[head:],
- request_indices=req_idx[head:],
- tokens_per_block=tpb,
- query_positions=None,
- compress_ratio=None,
- # the generation rows alone, or every row of the step
- rows_per_request=rows_per_request if head == num_ctx else 1,
- )
- out[head:rows].copy_(torch.reshape(rest, [rows - head, core.q_width]))
- return out
-
-
-def _live_score_columns(core, md, batch: int) -> int:
- """A host-side bound on every row's selection window in a step with context rows: the
- longest sequence's compressed length, with the verify width as slack for the draft
- rows that sit past the host's KV lengths."""
- longest = int(md.kv_lens_runtime[:batch].max()) if batch > 0 else 0
- return (longest + _MAX_VERIFY_WIDTH) // core.compress_ratio + 1
-
-
-def _replay_verify(
- core, mq_out, a_d, b_d, ssm_states, i, L, si_d, n, p, replay, pdl=False
-) -> torch.Tensor:
- """One GDN layer's verify recurrence under the cached replay: `n` generation sequences
- of `p` rows over the conv output `mq_out` (`[n * p, conv_dim]`, token-major) and the
- raw `a`/`b` projections, with the fp32 gating parameters fused into the kernel.
-
- `replay` is `(work_items, n_writes, partitioned)`: the work items are given exactly when
- the manager runs its all-layer commit after the step (without that commit every history
- wrap goes uncommitted, with it a wrap committed here is committed twice); without them
- the kernel commits every wrap itself. An fp32 pool runs the Triton kernel, a bf16 pool
- the CuTe one (`gdn_cached_replay_cute`) under the same contract. `mq_out`
- is the conv's fresh, contiguous output, so the kernel reads q/k/v from it directly
- (`packed_qkv`) and never from the views; the CuTe kernel reads `a`/`b` through their
- row stride (column views of the input projection). `pdl`: the kernel right before is
- the in_proj GEMV with the conv as its epilogue, which signals its dependents at its
- start -- the CuTe kernel then reads the checkpoint and the history (which nothing in
- the step writes ahead of it) before `griddepcontrol.wait`. The Triton conv does not
- signal dependents, so after it the kernel launches without PDL."""
- work_items, n_writes, partitioned = replay
- state = core._bound.replay
- old_x, old_B, old_dt, old_dA = state.old[i]
- qv, kv, vv = torch.split(mq_out, [core.key_dim, core.key_dim, core.value_dim], 1)
- entry = (
- gdn_cached_replay_cute
- if state.cute
- else fused_recurrent_gated_delta_rule_cached_replay_update
- )
- out = entry(
- q=torch.reshape(qv, [n, p, core.num_k_heads, core.head_k_dim]),
- k=torch.reshape(kv, [n, p, core.num_k_heads, core.head_k_dim]),
- v=torch.reshape(vv, [n, p, core.num_v_heads, core.head_v_dim]),
- g=torch.reshape(a_d, [n, p, core.num_v_heads]),
- beta=torch.reshape(b_d, [n, p, core.num_v_heads]),
- ssm_states=ssm_states,
- state_indices=si_d,
- old_u=old_x,
- old_k=old_B,
- old_G=old_dt,
- old_beta=old_dA,
- cache_buf_idx=state.cache_buf_idx,
- prev_num_accepted_tokens=state.prev_num_accepted,
- history_size=state.history,
- scale=None,
- use_qk_l2norm_in_kernel=True,
- A_log=L.a_log,
- dt_bias=L.dt_bias,
- launch_with_pdl=bool(pdl and state.cute),
- replay_work_items=work_items,
- n_writes=n_writes,
- block_v=None,
- num_warps=None,
- use_l2_state_cache=None,
- use_l2_streaming_inputs=None,
- use_l2_shared_inputs=None,
- main_num_stages=None,
- packed_qkv=mq_out,
- use_all_layer_commit=partitioned,
- output=None,
- )
- return torch.reshape(out, [1, n * p, core.num_v_heads, core.head_v_dim])
-
-
-def _replay_work(core, mm, num_decodes: int) -> tuple:
- """`_replay_verify`'s `replay` argument for a verify step of `num_decodes` generation
- requests: the partitioned replay exactly when `update_mamba_states` runs the manager's
- all-layer commit after the step."""
- partitioned = (
- core._bound.replay.all_layer_commit
- and num_decodes >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE
- )
- return (
- mm.replay_work_items[:num_decodes] if partitioned else None,
- mm.replay_n_writes if partitioned else None,
- partitioned,
- )
-
-
-class _KvRangeProbe:
- """`TRTLLM_MODELING_V2_VALIDATE` over an fp8 KV pool: reports, once, the largest K/V
- magnitude the first eager forward's attention layers store.
-
- The pool takes K/V at the unit scale, so anything past the e4m3 maximum saturates
- silently (a bf16 inf included). One host sync per attention layer of that one forward;
- a forward under CUDA graph capture is passed over without a read."""
-
- def __init__(self, last_layer: int) -> None:
- self._last = last_layer
- self._worst = 0.0
-
- def observe(self, layer: int, k: torch.Tensor, v: torch.Tensor) -> bool:
- """Fold one attention layer's K/V in; False once the forward's last layer has
- reported."""
- if torch.cuda.is_current_stream_capturing():
- return True
- peak = float(torch.maximum(k.abs().amax(), v.abs().amax()))
- # A NaN stays: `max` would drop it.
- self._worst = math.nan if math.isnan(peak + self._worst) else max(self._worst, peak)
- if layer != self._last:
- return True
- limit = torch.finfo(torch.float8_e4m3fn).max
- # Written so that NaN fails too.
- if not self._worst <= limit:
- logger.warning(
- f"FP8 KV cache: the first forward stores a K/V magnitude of {self._worst:g}, "
- f"past the e4m3 maximum {limit:g}; the pool is written at the unit scale, "
- "so those values saturate"
- )
- return False
-
-
-class _PrefillStep(NamedTuple):
- """One step's host-side state, built by `PrefillTarget.prepare` and read by the forward's
- breaks.
-
- `num_tokens` is the batch's row count `n` and `physical_tokens` the `T` the forward runs
- over (`n`, or the prefill CUDA graph bucket it was padded to); `input_ids` is the batch's
- `n` ids. `attn` is thop_attention's per-step arguments; `gdn`, `ple` and `qsa` are the
- recurrent, PLE and QSA bookkeeping tuples their breaks unpack, each cut to the batch's
- rows; `stage` says the speculative rollback stages this verify step; `vision` is the
- shell's `(rows, mm_token_indices)` or None; `nvfp4_context` is `_nvfp4_context_plan`'s
- scratch for the context rows over an NVFP4 pool, or None. A new one replaces it every
- step, so no break reads a previous step's."""
-
- num_tokens: int
- physical_tokens: int
- input_ids: torch.Tensor
- attn: dict
- gdn: tuple
- ple: tuple
- qsa: tuple
- stage: bool
- vision: tuple | None
- nvfp4_context: tuple | None = None
-
-
-class PrefillTarget:
- """Any batch with context rows (`num_contexts > 0`), possibly with generation rows
- beside them; the recurrent layers' decode-only branches never run here."""
-
- def __init__(self, core: Qwen4ExpTp1ModelingV2) -> None:
- """Take the per-layer weight views once; MTP operands are read off `core._mtp`."""
- self.core = core
- # Off unless TRTLLM_MODELING_V2_VALIDATE asks; each target owns its flag (the
- # two `step_args` projections differ).
- self._contract_pending = step_contract_enabled()
- # Dropped once it has reported.
- self._kv_probe = (
- _KvRangeProbe(core.attn_layers[-1])
- if self._contract_pending and core.quant_mode == QM_FP8_KV_CACHE
- else None
- )
- # The rstd sites' `(1 + w)` operands (`_HC_RSTD_MIN_ROWS`), formed once here.
- self._layers = [
- L._replace(ahc_norm_w1=_norm_w1(L.ahc_norm), mhc_norm_w1=_norm_w1(L.mhc_norm))
- for L in _layer_views(core)
- ]
- stack = _stack_views(core)
- self._stack = stack._replace(fhc_norm_w1=_norm_w1(stack.fhc_norm))
- # The draft sites' `(1 + w)` operands, by norm weight, formed on first use.
- self._mtp_norm_w1: dict = {}
- # The step `prepare` built, read by the breaks.
- self._step: _PrefillStep | None = None
-
- def _check_step_contract(self, md) -> None:
- """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): calling
- `step_args` surfaces an upstream metadata rename here, not mid-forward."""
- if not self._contract_pending:
- return
- try:
- self.step_args(md)
- except AttributeError as exc:
- raise AssertionError(f"attention metadata surface drifted: {exc}") from exc
- self._contract_pending = False
-
- def step_args(self, md: TrtllmAttentionMetadata) -> dict:
- """Project the prepared metadata onto thop_attention's batch state, once per
- forward. Tensors are engine-owned persistent buffers refreshed in place;
- Python ints are per-capture constants."""
- return dict(
- sequence_length=md.kv_lens_cuda_runtime,
- host_past_key_value_lengths=md.kv_lens_runtime,
- host_total_kv_lens=md.host_total_kv_lens,
- context_lengths=md.prompt_lens_cuda_runtime,
- host_context_lengths=md.prompt_lens_cpu_runtime,
- host_request_types=md.host_request_types_runtime,
- kv_cache_block_offsets=md.kv_cache_block_offsets,
- host_kv_cache_pool_pointers=md.host_kv_cache_pool_pointers,
- host_kv_cache_pool_mapping=md.host_kv_cache_pool_mapping,
- workspace=md.effective_workspace,
- tokens_per_block=md.tokens_per_block,
- max_num_requests=md.max_num_requests,
- max_context_length=md.max_context_length,
- max_seq_len=md.max_seq_len,
- # No sliding window: the full admitted length on every attention layer.
- attention_window_size=md.max_seq_len,
- num_contexts=md.num_contexts,
- num_ctx_tokens=md.num_ctx_tokens,
- trtllm_gen_jit_warmup=md.trtllm_gen_jit_warmup,
- use_paged_context_fmha=md.use_paged_context_fmha,
- beam_width=md.effective_beam_width,
- cache_indirection=md.cache_indirection,
- block_ids_per_seq=md.block_ids_per_seq,
- max_context_q_len_override=md.max_context_q_len_override,
- is_cross=md.is_cross,
- is_spec_decoding_enabled=md.is_spec_decoding_enabled,
- use_spec_decoding=md.use_spec_decoding,
- is_spec_dec_tree=md.is_spec_dec_tree,
- spec_decoding_generation_lengths=md.spec_decoding_generation_lengths,
- spec_decoding_position_offsets=get_spec_decoding_position_offsets(md),
- spec_decoding_packed_mask=md.spec_decoding_packed_mask,
- spec_decoding_bl_tree_mask_offset=md.spec_decoding_bl_tree_mask_offset,
- spec_decoding_bl_tree_mask=md.spec_decoding_bl_tree_mask,
- spec_decoding_target_max_draft_tokens=md.max_total_draft_tokens,
- spec_bl_tree_first_sparse_mask_offset_kv=md.spec_bl_tree_first_sparse_mask_offset_kv,
- num_sparse_topk=md.num_sparse_topk,
- flash_mla_tile_scheduler_metadata=md.flash_mla_tile_scheduler_metadata,
- flash_mla_num_splits=md.flash_mla_num_splits,
- max_num_sequences=md.max_num_sequences,
- force_prepare_spec_dec_tree_mask=md.force_prepare_spec_dec_tree_mask,
- )
-
- # --------------------------------------------------------------- blocks ---
-
- def _hc_mixer(self, normed, first, up, with_inject):
- """One hyper-connection mixer from the site's normed stream: the first
- projection over `[down; block_inject; 0]` (fp32 logits), the SiLU and
- injection activations, then the up projection with the sigmoid gate and
- the stream mean in its epilogue (the gate logits never leave the GEMM).
- Returns `(mixed, inject)`; `inject` is None without `with_inject` (the
- stack-final mixer)."""
- core = self.core
- low, inject = qwen4_exp_hc_lowrank_act(
- cublas_mm(normed, first[:, : core.hc_first_cols], None, torch.float32),
- core.hc_lowrank,
- core.hc_count,
- with_inject,
- normed.dtype,
- )
- return qwen4_exp_hc_up_gate_mean(low, up, normed, core.hc_count), inject
-
- def _hc_moe_mixer(self, normed, first, up, gscale):
- """`_hc_mixer` at an MoE site (always with injection) whose fused up
- projection also writes the routed experts' NVFP4 input (linear scales,
- the layer's global scale `gscale`). Returns `(mixed, inject, xq, xsf)`."""
- core = self.core
- low, inject = qwen4_exp_hc_lowrank_act(
- cublas_mm(normed, first[:, : core.hc_first_cols], None, torch.float32),
- core.hc_lowrank,
- core.hc_count,
- True,
- normed.dtype,
- )
- mixed, xq, xsf = qwen4_exp_hc_up_gate_mean_fp4(low, up, normed, core.hc_count, gscale)
- return mixed, inject, xq, xsf
-
- def _hc_mixer_rstd(self, hc, rstd, norm, norm_w1, first, up, with_inject):
- """`_hc_mixer` from the site's residual `hc` and its `[T, C]` rstd
- (`qwen4_exp_hc_inject_rstd`) instead of a stored normed stream: the first
- projection forms `bf16(hc * rstd * (1 + norm))` on its A tiles, the up
- projection's epilogue `hc * rstd * norm_w1` in fp32."""
- core = self.core
- low, inject = qwen4_exp_hc_first_rescale_tc(
- hc, rstd, norm, first, core.hc_count, core.hc_lowrank, with_inject
- )
- mixed = qwen4_exp_hc_up_gate_mean_rstd(low, up, hc, rstd, norm_w1, core.hc_count)
- return mixed, inject
-
- def _hc_moe_mixer_rstd(self, hc, rstd, norm, norm_w1, first, up, gscale):
- """`_hc_moe_mixer` on the rstd form (see `_hc_mixer_rstd`). Returns
- `(mixed, inject, xq, xsf)`."""
- core = self.core
- low, inject = qwen4_exp_hc_first_rescale_tc(
- hc, rstd, norm, first, core.hc_count, core.hc_lowrank, True
- )
- mixed, xq, xsf = qwen4_exp_hc_up_gate_mean_rstd_fp4(
- low, up, hc, rstd, norm_w1, core.hc_count, gscale
- )
- return mixed, inject, xq, xsf
-
- def _grouped_norm(self, x: torch.Tensor, weights: tuple) -> torch.Tensor:
- """`Qwen4ExpTextRMSNorm(group_size=hidden)`: an independent RMS per `hidden`-wide
- group, then `(1 + w)` over the full width; one gemma-RMS-norm per column
- slice (row stride certified by that entry) plus a concat."""
- core = self.core
- parts = torch.split(x, [core.hidden] * core.hc_count, 1)
- return torch.cat(
- [flashinfer_gemma_rmsnorm(part, w, core.eps) for w, part in zip(weights, parts)], 1
- )
-
- def _moe_block(self, x, xq, xsf, L: _LayerWeights, routed_out) -> tuple:
- """One MoE layer: NVFP4 routed experts plus a bf16 shared expert, left
- uncombined. Returns `(routed, (gate_logits, shared))`: the next
- `qwen4_exp_hc_inject_norm` forms `routed + sigmoid(gate_logits) * shared`
- inside its kernel.
-
- The shared expert and its per-token gate logit run on the core's side
- stream, concurrently with the routed path, and join it before this
- returns, so every consumer on the main stream is ordered after both.
- `x` is the bf16 hidden states (router and shared expert); `xq` / `xsf`
- are the same states in NVFP4 with **linear** block scales, written by the
- MoE-site mixer (`_hc_moe_mixer`). On trtllm-gen the routed path hands the
- runner the router logits (its Renormalize routing is bitwise the
- `renorm_moe_routing_op` + pre-routed composition), and a launch of at
- least `_CUTE_DSL_MOE_MIN_ROWS` rows runs the CuteDSL MoE over the second
- expert copy instead when the core keeps it (`moe_prefill_copy`); on CuteDSL
- every launch is `_cute_dsl_routed`, which writes `routed_out`."""
- core = self.core
- main = torch.cuda.current_stream()
- side = core._side_stream
- # Fork: the side stream reads `x` only after everything queued on `main`
- # so far. What it allocates is consumed on `main` after the join below,
- # and the next layer's fork orders any reuse of those blocks after that.
- side.wait_stream(main)
- with torch.cuda.stream(side):
- # One GEMM forms gate|up and the gate logit (`sh_gug`'s extra rows): a separate
- # gate GEMV at prefill sizes keeps the SMs from the routed path's persistent
- # FC1 until it drains.
- gug = cublas_mm(x, L.sh_gug)
- shared = cublas_mm(silu_and_mul_rows(gug, core.shared_inter), L.sh_dn)
- gate_logits = gug[:, 2 * core.shared_inter : 2 * core.shared_inter + 1]
- # The routing reads its logits as a dense row-major buffer from data_ptr()
- # and ignores strides; cublas_mm allocates them contiguous.
- logits = cublas_mm(x, L.router)
- if core.trunk_moe_backend == "CUTEDSL":
- routed = _cute_dsl_routed(core, L, logits, xq, xsf, routed_out)
- elif core.moe_prefill_copy and x.shape[0] >= _CUTE_DSL_MOE_MIN_ROWS:
- routed = fp4_moe_renorm_cute_dsl(
- logits,
- xq,
- xsf,
- L.fc1_w_cd,
- L.fc1_s_cd,
- L.alpha1_cd,
- L.fc2_w_cd,
- L.fc2_s_cd,
- L.alpha2_cd,
- L.global_sf_cd,
- core.num_experts,
- core.topk,
- memset_stream=core._moe_memset_stream,
- )
- else:
- routed = fp4_block_scale_moe_renorm_routed(
- logits,
- xq,
- xsf.view(torch.float8_e4m3fn),
- L.fc1_w,
- L.fc1_s,
- L.fc2_w,
- L.fc2_s,
- L.out1_scale,
- L.out1_gate_scale,
- L.out2_scale,
- core.num_experts,
- core.topk,
- core.moe_inter,
- )
- main.wait_stream(side)
- return routed, (gate_logits, shared)
-
- def _attention_block(self, x, i, L, proj_qkv, index_scratch, attn_out, pos3):
- """One gated GQA layer plus its QSA index branch.
-
- The four input projections are one GEMM against the load-time fused
- operand (`weights._attn_in_w`), whose output row is
- `[q | k | v | gate | index q, k | v slot]`. One norm+rope launch rewrites
- the main q/k heads in place on that output and writes the normalized,
- rotated index query into the scratch; v, the gate and the raw index key
- are not touched. Both are row-wise and run over all `T` rows here; the
- cache writes, the selection and the attention core read the step's batch
- and run in the `_attention_core` break, which leaves the attention output
- in `attn_out`. The gate is read through its row stride."""
- core = self.core
- num_tokens = x.shape[0]
-
- proj = _dense_mm(x, L.attn_in, L.attn_in8, cublas_mm, "in")
- gate = torch.split(proj, [core.qkv_width, core.q_width, core.index_scratch_width], 1)[1]
- # The index query lands in the q slots of the scratch; its k and v slots
- # are layout only and nothing reads them.
- q_index = torch.split(
- torch.reshape(
- index_scratch,
- [num_tokens, core.index_heads + 2 * core.index_kv_heads, core.index_head_dim],
- ),
- [core.index_heads, core.index_kv_heads, core.index_kv_heads],
- 1,
- )[0]
- # Both planes' Gemma norm + rope in one launch, over the target's fp32
- # cos|sin table: the main q/k heads in place (the real interleaved mRoPE;
- # three agreeing rows reduce exactly to the plain rope on row 0), the index
- # query from the projection's index columns into `q_index` (coordinate 0
- # alone, matching `_QSA_MROPE_SECTION` on the key side).
- index_section = (_QSA_MROPE_SECTION or (0, 0, 0)) if _QSA_INDEX_USE_MROPE else (0, 0, 0)
- qwen4_exp_qk_norm_rope(
- proj,
- q_index,
- pos3,
- core._cos_sin,
- L.q_norm,
- L.k_norm,
- L.idx_qn,
- num_heads_q=core.heads_q,
- num_heads_k=core.heads_kv,
- head_dim=core.head_dim,
- index_heads=core.index_heads,
- index_head_dim=core.index_head_dim,
- index_q_col=core.qkv_width + core.q_width,
- rotary_dim=core.rotary_dim,
- eps=core.eps,
- mrope_section1=core.mrope_section1,
- mrope_section2=core.mrope_section2,
- index_mrope_section1=index_section[1],
- index_mrope_section2=index_section[2],
- )
- if self._kv_probe is not None:
- _, k_dst, v_dst = _attn_in_heads(core, proj)
- if not self._kv_probe.observe(i, k_dst, v_dst):
- self._kv_probe = None
- eager_on_graph(proj, proj_qkv, index_scratch, attn_out, i, inner=self._attention_core)
- fused_sigmoid_mul_inplace(attn_out, gate)
- return _dense_mm(attn_out, L.o, L.o8, cublas_mm, "out")
-
- def _attention_core(self, proj, proj_qkv, index_scratch, attn_out, i) -> None:
- """Break: the index plane's store and compression, the dense/sparse decision and the
- attention core over the step's `n` rows. Writes rows `[:n]` of `attn_out`
- `[T, q_width]` in place and zeroes rows `[n:]`. `proj` is the layer's fused
- projection output after the norm+rope, `proj_qkv` the dense op's packed QKV
- scratch.
-
- The index plane advances on every step, including below the threshold, so a
- request growing past it finds its prefix compressed; above it the main K/V append
- becomes this layer's own call. The raw index key is stored before any
- norm/rotation: the compression step applies both itself, so a pre-normalized key
- is double-normalized (invisible below the sparse threshold).
-
- On a verify step (`stage`) the generation rows' index slots (and, on the first
- attention layer, their position slots) are snapshotted for `SpeculativeRollback`
- before the store and compression overwrite them."""
- core = self.core
- L = self._layers[i]
- step = self._step
- (
- req_idx,
- logical,
- seq_lens_q,
- visible_blocks,
- topk_indices,
- topk_row_starts,
- block_table,
- tokens_per_block,
- coords,
- sparse,
- live_cols,
- ) = step.qsa
- num_tokens = step.num_tokens
- proj = proj[:num_tokens]
- qkv_view, _, index_view = torch.split(
- proj, [core.qkv_width, core.q_width, core.index_scratch_width], 1
- )
- q_dst, k_dst, v_dst = _attn_in_heads(core, proj)
-
- index_cache = core._bound.index_cache[i]
- token_k = torch.split(
- torch.reshape(
- index_view,
- [num_tokens, core.index_scratch_width // core.index_head_dim, core.index_head_dim],
- ),
- [core.index_heads, core.index_kv_heads, 1],
- 1,
- )[1]
- if step.stage:
- core.spec_rollback.stage_index(index_cache, core._bound.position_cache)
- qsa_paged_index_store(
- token_k,
- coords,
- index_cache,
- core._bound.position_cache,
- req_idx,
- logical,
- block_table,
- tokens_per_block,
- )
- q_index = torch.split(
- torch.reshape(
- index_scratch[:num_tokens],
- [num_tokens, core.index_heads + 2 * core.index_kv_heads, core.index_head_dim],
- ),
- [core.index_heads, core.index_kv_heads, core.index_kv_heads],
- 1,
- )[0]
- qsa_prefill_compress(
- logical_positions=logical,
- request_indices=req_idx,
- block_table=block_table,
- index_cache=index_cache,
- position_cache=core._bound.position_cache,
- k_norm_weight=L.idx_kn,
- cos_sin=core._cos_sin,
- eps=core.eps,
- tokens_per_block=tokens_per_block,
- compress_ratio=core.compress_ratio,
- mrope_section=_QSA_MROPE_SECTION,
- )
-
- if sparse:
- kv = core._bound.kv[i]
- # Append the fused-norm+rope'd K/V here: nothing else writes the pool on
- # this path, and selection reads the current token from it. Over an fp8 pool the
- # store and the sparse GQA below take e4m3 unscaled, as the dense op does at the
- # unit KV scales; over an NVFP4 pool they quantize and dequantize at the dense op's
- # second-level scales.
- _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block)
- # The step's context rows come first, consecutive rows of a request, and score
- # as query tiles; its generation rows, one a request, score row by row (a query
- # tile of them would walk every request's keys once per row). Each part selects
- # into its own rows of the top-k.
- num_ctx = step.attn["num_ctx_tokens"]
- parts = [(0, num_ctx, True)]
- if num_ctx < num_tokens:
- parts.append((num_ctx, num_tokens, False))
- # Each part's canonical block top-k and its expansion are one launch writing its
- # rows of the step's selection (QSA's row starts are zero); a score table that
- # kernel does not serve (wider than its 4096 columns) takes the composition.
- selected = torch.empty(
- [num_tokens, core.token_topk + core.compress_ratio - 1],
- dtype=torch.int32,
- device=proj.device,
- )
- for lo, hi, context_rows in parts:
- scores = qsa_paged_index_scores(
- q=q_index[lo:hi],
- index_cache=index_cache,
- block_table=block_table,
- query_positions=logical[lo:hi],
- request_indices=req_idx[lo:hi],
- tokens_per_block=tokens_per_block,
- compress_ratio=core.compress_ratio,
- only_visible_blocks=True,
- context_rows=context_rows,
- )
- if _topk_expand_supports(scores):
- qsa_topk_expand_decode(
- scores,
- visible_blocks[lo:hi],
- logical[lo:hi],
- seq_lens_q[lo:hi],
- topk_indices[lo:hi],
- core.compress_ratio,
- core.token_topk,
- out=selected[lo:hi],
- )
- else:
- _select(
- core,
- scores,
- topk_row_starts[lo:hi],
- visible_blocks[lo:hi],
- topk_indices[lo:hi],
- live_cols,
- )
- selected[lo:hi].copy_(
- qsa_expand_block_indices(
- topk_indices[lo:hi],
- logical[lo:hi],
- seq_lens_q[lo:hi],
- compress_ratio=core.compress_ratio,
- token_topk=core.token_topk,
- )
- )
- if kv.k_scale is not None:
- _sparse_gqa_nvfp4_context(
- core,
- kv,
- step.nvfp4_context,
- q_dst,
- block_table,
- selected,
- req_idx,
- logical,
- tokens_per_block,
- num_ctx,
- attn_out[:num_tokens],
- rows_per_request=step.attn["predicted_tokens_per_seq"],
- )
- elif num_tokens >= _ROWTILE_MIN_ROWS and kv.k.dtype == q_dst.dtype:
- # A context launch of thousands of rows over a bf16 pool takes the row-tile
- # pass: consecutive rows of one request share their K/V tiles; rows of
- # groups that span requests, the generation rows of a mixed step among them,
- # are served row by row. Below the threshold the per-row kernel is faster.
- out = qsa_sparse_gqa_rowtile(
- q=q_dst,
- k_cache=kv.k,
- v_cache=kv.v,
- block_table=block_table,
- selected_tokens=selected,
- request_indices=req_idx,
- tokens_per_block=tokens_per_block,
- softmax_scale=core.softmax_scale,
- query_positions=logical,
- compress_ratio=core.compress_ratio,
- )
- attn_out[:num_tokens].copy_(torch.reshape(out, [num_tokens, core.q_width]))
- else:
- out = _sparse_gqa(
- core,
- kv,
- q=q_dst,
- block_table=block_table,
- selected_tokens=selected,
- request_indices=req_idx,
- tokens_per_block=tokens_per_block,
- # A bound, not a mask: lets the kernel stop scanning a row early
- # (context-batch branch only).
- query_positions=logical,
- compress_ratio=core.compress_ratio,
- )
- attn_out[:num_tokens].copy_(torch.reshape(out, [num_tokens, core.q_width]))
- else:
- # The dense op reads its packed QKV flat: one compact copy.
- qkv = proj_qkv[:num_tokens]
- qkv.copy_(qkv_view)
- thop_attention(
- q=qkv,
- output=attn_out[:num_tokens],
- local_layer_idx=core._bound.pool_idx[i],
- num_heads=core.heads_q,
- num_kv_heads=core.heads_kv,
- head_size=core.head_dim,
- quant_mode=core.quant_mode,
- kv_scale_orig_quant=core.kv_scale_orig_quant,
- kv_scale_quant_orig=core.kv_scale_quant_orig,
- **step.attn,
- **_CALL_CONSTANTS,
- )
- attn_out[num_tokens:].zero_()
-
- def _gdn_verify(
- self, mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, n, p, replay
- ) -> torch.Tensor:
- """The recurrent half of a speculative verify step: `n` generation sequences of
- `p = runtime_draft_len + 1` tokens, written so that the accepted prefix, unknown
- until the sampler has run, can be kept and the rest dropped.
-
- The conv writes each row's window to the runtime's per-step buffer and leaves the
- conv pool alone; `update_mamba_states` promotes window `k - 1`.
- `causal_conv1d_triton_update` is the only conv with per-step snapshots (its result
- for a pad-sentinel slot is uninitialized, safe because every row is a real
- sequence). The recurrence has two forms, bound per manager
- (`_bind_step_buffers`):
-
- * cached replay (`replay` given): the pool row is a checkpoint; one call rebuilds
- the current state from it and the slot's pending history, runs the `p` rows and
- appends them to the history. Rejected rows cost nothing to undo, because the
- manager advances the pending count by the accepted tokens only. A full history
- advances the checkpoint inside the kernel below
- `CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE` requests and in the manager's
- all-layer commit from there on, hence `replay` carries the work items exactly
- under the manager's own condition.
- * legacy (`replay` None): every row's recurrent state goes to the per-step buffer
- instead of the pool, and the promotion copies step `k - 1` into it. Advancing
- the pool directly would leave rejected tokens' state in every recurrent layer,
- invisible to the acceptance rate. `pack_gdn_decode_qkv` does no Q/K L2 norm,
- hence `use_qk_l2norm_in_kernel=True` below."""
- core = self.core
- # [n*p, C] -> [n, C, p]: the multi-token conv reads channel-major; its cells
- # certify this transposed view directly.
- x3 = torch.transpose(torch.reshape(mq_d, [n, p, core.conv_dim]), 1, 2)
- conv_out = causal_conv1d_triton_update(
- x=x3,
- conv_state=core._bound.conv[i],
- weight=L.conv_w,
- bias=None,
- activation="silu",
- conv_state_indices=si_d,
- pad_slot_id=_PAD_SLOT_ID,
- # Snapshots, not the accepted-token roll.
- num_accepted_tokens=None,
- intermediate_conv_window=core._bound.int_conv[i],
- intermediate_state_indices=int_idx,
- launch_dependent_kernels=False,
- )
- # Back to token-major (a view).
- mq_out = torch.reshape(torch.transpose(conv_out, 1, 2), [n * p, core.conv_dim])
- if replay is not None:
- return _replay_verify(core, mq_out, a_d, b_d, ssm_states, i, L, si_d, n, p, replay)
- q, k, v = pack_gdn_decode_qkv(
- mq_out, core.num_k_heads, core.head_k_dim, core.num_v_heads, core.head_v_dim
- )
- g = torch.reshape(
- fused_gdn_gating(
- A_log=L.a_log,
- a=a_d,
- dt_bias=L.dt_bias,
- beta=_SOFTPLUS_BETA,
- threshold=_SOFTPLUS_THRESHOLD,
- ),
- [n, p, core.num_v_heads],
- )
- beta = torch.reshape(torch.sigmoid(b_d), [n, p, core.num_v_heads])
- # Initial states are gathered from the live pool into a batch-scoped copy. The
- # per-step buffer is indexed by batch row (`intermediate_state_indices[:n]`),
- # not slot id: slot ids would write rows nothing promotes, or past the end.
- src = torch.reshape(
- nn.functional.embedding(si_d, torch.reshape(ssm_states, [ssm_states.shape[0], -1])),
- [n, core.num_v_heads, core.head_v_dim, core.head_k_dim],
- )
- out = fused_recurrent_gated_delta_rule_update(
- torch.reshape(q, [n, p, core.num_k_heads, core.head_k_dim]),
- torch.reshape(k, [n, p, core.num_k_heads, core.head_k_dim]),
- torch.reshape(v, [n, p, core.num_v_heads, core.head_v_dim]),
- g,
- beta,
- initial_state_source=src,
- initial_state_indices=int_idx,
- scale=None,
- use_qk_l2norm_in_kernel=True,
- cu_seqlens=None,
- disable_state_update=True,
- disable_output_calculation=False,
- intermediate_states_buffer=core._bound.int_ssm[i],
- cache_steps=p,
- output=None,
- )
- return torch.reshape(out, [1, n * p, core.num_v_heads, core.head_v_dim])
-
- def _gdn_block(self, x, i, L, core_out) -> torch.Tensor:
- """One Gated-DeltaNet layer over a batch with a context prefix: the one input
- projection GEMM (q|k|v, z, a and b are column views of its output) and the
- gated-norm epilogue run over all `T` rows here, the recurrence in the `_gdn_core`
- break, which leaves its output in `core_out` `[1, T, num_v_heads, head_v_dim]`."""
- core = self.core
- proj = _dense_mm(x, L.in_proj, L.in_proj8, cublas_mm, "in")
- z = torch.split(proj, core.in_proj_split, 1)[1]
- eager_on_graph(proj, core_out, i, inner=self._gdn_core)
- return self._gdn_epilogue(core_out, z, L, x.shape[0])
-
- def _gdn_core(self, proj, out, i) -> None:
- """Break: the recurrence over the step's `n` rows. Writes rows `[:n]` of `out`
- `[1, T, num_v_heads, head_v_dim]` in place and zeroes rows `[n:]`. Every branch
- is on a host int of the step, never a tensor value.
-
- Outside a verify step, over either pool dtype, the context rows take one fused
- conv + post-conv pass (`_gdn_prefill_prep`) and FlashInfer's chunked recurrence, and
- a mixed step's one-token generation rows the decode step's recurrence
- (`gdn_decode_rows`); a verify step takes `_gdn_core_verify`. A context sequence
- carrying nothing starts from zeros without its pool slots being cleared: the conv
- pass reads no carried window for it and publishes the whole window, and the
- recurrence (`has_initial_state`) gathers zeros for it and scatters its final state."""
- core = self.core
- L = self._layers[i]
- step = self._step
- (
- num_prefill_tokens,
- num_decode_tokens,
- si_p,
- si_all,
- si_d,
- hi_p,
- hi_all,
- qsl_p,
- qsl_all,
- qsl_long,
- is_verify,
- num_decodes,
- gen_p,
- int_idx,
- qsl_long_p,
- replay,
- ) = step.gdn
- num_tokens = step.num_tokens
- if is_verify:
- self._gdn_core_verify(proj, out, i)
- return
- ssm_states = core._bound.ssm[i]
- mixed_qkv, _, a, b = torch.split(proj[:num_tokens], core.in_proj_split, 1)
- conv_states = core._bound.conv[i]
-
- if num_decode_tokens > 0:
- mq_p, mq_d = torch.split(mixed_qkv, [num_prefill_tokens, num_decode_tokens], 0)
- a_p, b_p = a[:num_prefill_tokens], b[:num_prefill_tokens]
- q, k, v, g, beta = self._gdn_prefill_prep(mq_p, a_p, b_p, L, i, qsl_p, si_p, hi_p)
- causal_conv1d_update(
- x=mq_d,
- conv_state=conv_states,
- weight=L.conv_w,
- bias=None,
- silu_activation=True,
- cache_seqlens=None,
- conv_state_indices=si_d,
- pad_slot_id=_PAD_SLOT_ID,
- )
- # The two halves run separate recurrences over disjoint slots and write
- # disjoint row ranges of the output.
- out_p = out[:, :num_prefill_tokens]
- flashinfer_chunk_gated_delta_rule(
- q,
- k,
- v,
- g,
- beta,
- scale=None,
- initial_state=ssm_states,
- initial_state_indices=si_p,
- inplace_indexed_state_update=True,
- output_final_state=False,
- cu_seqlens=qsl_long_p,
- head_first=False,
- use_qk_l2norm_in_kernel=False,
- output=out_p,
- state_workspace=None,
- has_initial_state=hi_p,
- linear_gate=True,
- )
- # The generation rows: the decode step's recurrence (its kernel without the
- # gated norm, which runs with the context rows' below) over their rows of the
- # projection, whose q/k/v columns the conv step rewrote; it computes both gates
- # and the q/k norm itself and reads the rows' pool slots before its PDL wait
- # (the context rows' recurrence touches only theirs).
- gdn_decode_rows(
- proj[num_prefill_tokens:num_tokens],
- ssm_states,
- si_d,
- L.a_log,
- L.dt_bias,
- core.num_k_heads,
- core.head_k_dim,
- out[0, num_prefill_tokens:num_tokens].view(num_decode_tokens, core.value_dim),
- )
- else:
- q, k, v, g, beta = self._gdn_prefill_prep(
- mixed_qkv, a, b, L, i, qsl_all, si_all, hi_all
- )
- # q and k arrive L2-normalized; the recurrence must not repeat it.
- flashinfer_chunk_gated_delta_rule(
- q,
- k,
- v,
- g,
- beta,
- scale=None,
- initial_state=ssm_states,
- initial_state_indices=si_all,
- inplace_indexed_state_update=True,
- output_final_state=False,
- cu_seqlens=qsl_long,
- head_first=False,
- use_qk_l2norm_in_kernel=False,
- output=out[:, :num_tokens],
- state_workspace=None,
- has_initial_state=hi_all,
- linear_gate=True,
- )
- out[:, num_tokens:].zero_()
-
- def _gdn_core_verify(self, proj, out, i) -> None:
- """`_gdn_core` on a verify step: the context rows' fused conv + post-conv pass and
- FlashInfer's chunked delta rule (as without verify rows), and the verify path for the
- generation rows. A context row
- advances its slot as on the base path; a generation row is a draft chain of unknown
- accepted length, so it goes to the rollback buffers. The q|k|v columns are read in
- place, as the column view of the projection every consumer here is certified on;
- the a and b columns are copied out first, as the verify path reads them contiguous.
- The chunked delta rule writes its rows of `out` directly."""
- core = self.core
- L = self._layers[i]
- step = self._step
- (
- num_prefill_tokens,
- num_decode_tokens,
- si_p,
- si_all,
- si_d,
- hi_p,
- hi_all,
- qsl_p,
- qsl_all,
- qsl_long,
- is_verify,
- num_decodes,
- gen_p,
- int_idx,
- qsl_long_p,
- replay,
- ) = step.gdn
- num_tokens = step.num_tokens
- mixed_qkv, _, a, b = torch.split(proj[:num_tokens], core.in_proj_split, 1)
- a, b = a.contiguous(), b.contiguous()
-
- ssm_states = core._bound.ssm[i]
-
- # A context sequence carrying nothing starts from zeros (`has_initial_state`
- # below; its slots are not cleared first, see `_gdn_core`).
- mq_p, mq_d = torch.split(mixed_qkv, [num_prefill_tokens, num_decode_tokens], 0)
- a_p, a_d = torch.split(a, [num_prefill_tokens, num_decode_tokens], 0)
- b_p, b_d = torch.split(b, [num_prefill_tokens, num_decode_tokens], 0)
- # The context rows as on a step without verify rows: one fused conv + post-conv
- # pass and FlashInfer's chunked recurrence (q and k arrive L2-normalized).
- q, k, v, g, beta = self._gdn_prefill_prep(mq_p, a_p, b_p, L, i, qsl_p, si_p, hi_p)
- flashinfer_chunk_gated_delta_rule(
- q,
- k,
- v,
- g,
- beta,
- scale=None,
- initial_state=ssm_states,
- initial_state_indices=si_p,
- inplace_indexed_state_update=True,
- output_final_state=False,
- cu_seqlens=qsl_long_p,
- head_first=False,
- use_qk_l2norm_in_kernel=False,
- output=out[:, :num_prefill_tokens],
- state_workspace=None,
- has_initial_state=hi_p,
- linear_gate=True,
- )
- out_d = self._gdn_verify(
- mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, num_decodes, gen_p, replay
- )
- out[:, num_prefill_tokens:num_tokens].copy_(out_d)
- out[:, num_tokens:].zero_()
-
- def _gdn_prefill_prep(self, mq, a, b, L, i, qsl, si, hi):
- """The context rows' recurrence inputs: the causal conv against each sequence's
- carried window (published back to the pool), SiLU, the q|k|v split with q/k
- L2-normalized, and the decay/beta gates (the decay as the factor `exp(g)`, which
- the chunked recurrence takes with `linear_gate=True`) -- one pass over `mq`, the
- token-major q|k|v column view of the projection."""
- core = self.core
- return gdn_prefill_conv_post(
- x=mq,
- weight=L.conv_w,
- conv_states=core._bound.conv[i],
- query_start_loc=qsl,
- cache_indices=si,
- has_initial_state=hi,
- a=a,
- b=b,
- A_log=L.a_log,
- dt_bias=L.dt_bias,
- num_k_heads=core.num_k_heads,
- head_k_dim=core.head_k_dim,
- num_v_heads=core.num_v_heads,
- head_v_dim=core.head_v_dim,
- pad_slot_id=_PAD_SLOT_ID,
- l2_norm_eps=_L2_NORM_EPS,
- softplus_threshold=_SOFTPLUS_THRESHOLD,
- linear_gate=True,
- )
-
- def _gdn_epilogue(self, out, z, L, num_tokens) -> torch.Tensor:
- """Gated RMS norm and output projection shared by every recurrence."""
- core = self.core
- # Sigmoid gate and a plain (non-gemma) weight (`output_gate_type`).
- o = rms_norm_gated_token_major(
- torch.reshape(out, [-1, core.head_v_dim]),
- torch.reshape(z, [num_tokens, core.num_v_heads, core.head_v_dim]),
- weight=L.gnorm,
- eps=core.eps,
- fp8_scale=None,
- gate_activation="sigmoid",
- )
- o = torch.reshape(o, [num_tokens, core.value_dim])
- return _dense_mm(o, L.out, L.out8, cublas_mm, "out")
-
- def _ple_block(self, hc, hc_out, normed_out) -> None:
- """Break: the n-gram side path on one decoder layer, over the step's `n` rows,
- added into the stream and normalized for the layer's attention-side site. Writes
- rows `[:n]` of `hc_out` (the stream) and `normed_out` (its grouped norm), both
- `[T, stream]`, in place and zeroes rows `[n:]`.
-
- Per-request stateful: ids become one `ngram_size`-wide window per token
- (predecessors from the paged id history); each window hashes to one row per
- head; rows gather fp8 vectors from one table and dequantize by one per-tensor
- scale; key/value projections and the gate scale one value into each stream;
- a dilated depthwise short conv over a paged state is added back. A sequence
- carrying nothing is seeded with `eos_token_id`, not 0: id 0 is a real token
- the hash cannot distinguish. The id history and the conv state publish the
- state after the whole chunk.
-
- The window, hash, gather and dequantize are `ple_ngram_embed`, and the gate, the
- conv, the residual add and the next site's grouped norm `ple_gate_conv_inject_norm`,
- which writes both outputs in place. On a verify step (`stage`) the generation rows'
- pre-step pool rows and ids are staged before the first launch publishes, and the
- conv input the fused launch writes out after it, for `SpeculativeRollback` to
- rewrite both pools to the accepted prefix once the sampler has decided it."""
- core = self.core
- S = self._stack
- step = self._step
- conv_states, ngram_pool, si_all, carry_all, qsl = step.ple
- num_tokens = step.num_tokens
- next_norm = self._layers[core.ple_layer].ahc_norm
- hc = hc[:num_tokens]
- conv_input = None
- if step.stage:
- core.spec_rollback.stage_ple_verify(ngram_pool, conv_states, step.input_ids)
- conv_input = torch.empty_like(hc)
- ngram = ple_ngram_embed(
- step.input_ids,
- ngram_pool,
- si_all,
- carry_all,
- qsl,
- core.w["ple_table"],
- S.ple_tscale,
- S.ple_mult,
- S.ple_vocab,
- S.ple_offs,
- core.eos_token_id,
- num_tokens,
- )
- ple_gate_conv_inject_norm(
- hc,
- cublas_mm(ngram, S.ple_key),
- cublas_mm(ngram, S.ple_val),
- S.ple_nk_w,
- S.ple_nq_w,
- S.ple_nc_w,
- S.ple_conv_w,
- conv_states,
- si_all,
- qsl,
- next_norm,
- core.eps,
- core.hc_count,
- _PLE_CLAMP_MIN,
- core.ple_dilation,
- _PAD_SLOT_ID,
- num_tokens,
- carry_all,
- conv_input=conv_input,
- out=(hc_out[:num_tokens], normed_out[:num_tokens]),
- )
- if step.stage:
- core.spec_rollback.stage_ple_conv_input(conv_input)
- hc_out[num_tokens:].zero_()
- normed_out[num_tokens:].zero_()
-
- def _splice(self, h) -> None:
- """Break: write the step's vision rows over their placeholder rows of the token
- embedding `h` `[T, hidden]` in place, and zero rows `[n:]`.
-
- `mm_token_indices` names every placeholder row of the packed batch in the order the
- per-request rows concatenate in. An indexed write, not a slice copy: two images in
- one prompt are separated by framing tokens, so the placeholder rows are several
- disjoint runs. A text step writes nothing but the padding. The rows are the step's:
- a replayed graph never sees the previous step's."""
- step = self._step
- if step.vision is not None:
- rows, mm_token_indices = step.vision
- h.index_copy_(0, mm_token_indices, rows)
- h[step.num_tokens :].zero_()
-
- # -------------------------------------------------------------- forward ---
-
- def prepare(
- self,
- attn_metadata: AttentionMetadata,
- input_ids: torch.IntTensor,
- position_ids: torch.IntTensor,
- vision: tuple | None,
- spec_metadata,
- ) -> None:
- """Everything this step reads from the host, kept for the forward's breaks as one
- `_PrefillStep`: the contract check, the attention op's per-step arguments and the
- batch's bookkeeping, built once and shared by every layer.
-
- Called by the core's `begin_step` on every step with context rows, before the
- forward. On a step a breakable prefill graph replays, this and the breaks are all of
- this target that runs."""
- core = self.core
- md = attn_metadata
- self._check_step_contract(md)
- attn_args = self.step_args(md)
- mm = md.mamba_metadata
-
- # **The activation rows are the batch's rows, or a prefill CUDA graph bucket's.** A
- # breakable prefill graph pads the step's rows up to `padded_num_tokens` (zero ids,
- # zero positions), and so does a padded step that runs eagerly; the forward runs over
- # all of them and every break slices the batch's own rows back out. Anything else is
- # an activation buffer longer than the batch, which would push padding rows through
- # the recurrences into the state pools.
- physical_tokens = input_ids.shape[0]
- assert physical_tokens in (md.num_tokens, md.padded_num_tokens), (
- f"the engine handed {physical_tokens} activation rows for a batch its "
- f"metadata describes as {md.num_tokens} tokens "
- f"(padded_num_tokens={md.padded_num_tokens})"
- )
-
- # The batch's bookkeeping, built once and shared by every layer.
- num_tokens = md.num_tokens
- batch = md.num_seqs
- num_prefills = md.num_contexts
- num_decodes = batch - num_prefills
- num_prefill_tokens = md.num_ctx_tokens
- num_decode_tokens = num_tokens - num_prefill_tokens
- # Rows per generation sequence: 1 without speculation, else
- # `runtime_draft_len + 1`; read from the metadata because the worker rewrites
- # it between draft steps.
- gen_p = num_decode_tokens // num_decodes if num_decodes else 1
- assert gen_p * num_decodes == num_decode_tokens, (
- f"{num_decode_tokens} generation rows do not divide over "
- f"{num_decodes} generation sequences; a ragged per-sequence draft "
- "length can be expressed neither to the attention call nor to the "
- "multi-token recurrence"
- )
- assert core.mtp_enabled or gen_p == 1, (
- f"{num_decode_tokens} generation rows over {num_decodes} generation "
- "sequences with speculation off; the base path implements one query "
- "token per generation sequence, and both the attention call's "
- "within-block mask and the recurrent path's cumulative-length "
- "vector assume it"
- )
- # `gen_p > 1` is load-bearing: verify parks a generation row in the rollback
- # buffers, which `update_mamba_states` promotes only if the worker drafts.
- # With `runtime_draft_len == 0` (a `draft_len_schedule` can do this) it skips
- # drafting before the promotion, leaving every recurrent layer one token
- # stale; that row is the always-accepted golden token, so the ordinary decode
- # path is right.
- if spec_metadata is not None and num_decodes > 0:
- declared_p = spec_metadata.runtime_draft_len + 1
- assert gen_p == declared_p, (
- f"the batch carries {gen_p} query rows per generation sequence "
- f"but the worker's runtime_draft_len says {declared_p}; the "
- "per-step conv snapshots and the recurrent rollback buffer are "
- "sized from the second and written from the first"
- )
- # Per-iteration counterpart of the core's `max_draft_len` bound.
- check_verify_width(gen_p)
- check_replay_width(core._bound.replay, gen_p)
- is_verify = spec_metadata is not None and num_decodes > 0 and gen_p > 1
- attn_args["predicted_tokens_per_seq"] = gen_p
-
- si_all = mm.state_indices[:batch]
- hi_all = mm.has_initial_states[:batch]
- si_p, si_d = torch.split(si_all, [num_prefills, num_decodes], 0)
- # The single-token decode kernel takes its slot vector through a 16-byte-aligned
- # pointer, and the generation half of the int32 slot vector starts `num_prefills`
- # rows in, aligned only when that count is a multiple of four: one copy per step
- # makes it aligned for every layer.
- si_d = torch.clone(si_d)
- hi_p = hi_all[:num_prefills]
- # `query_start_loc` covers the context prefix; `_long` the whole batch, and its
- # prefix the context rows' chunked recurrence when generation rows run beside them.
- qsl_all = mm.query_start_loc
- qsl_p = None if qsl_all is None else qsl_all[: num_prefills + 1]
- qsl_long = mm.query_start_loc_long
- qsl_long_p = qsl_long[: num_prefills + 1]
- int_idx = None
- replay = None
- if is_verify:
- # The runtime's arange over the slot pool, from the cache manager (what
- # `update_mamba_states` reads).
- assert core._bound.int_idx is not None, "the cache-manager bind has not run"
- int_idx = core._bound.int_idx[:num_decodes]
- if core._bound.replay is not None:
- replay = _replay_work(core, mm, num_decodes)
- gdn_state = (
- num_prefill_tokens,
- num_decode_tokens,
- si_p,
- si_all,
- si_d,
- hi_p,
- hi_all,
- qsl_p,
- qsl_all,
- qsl_long,
- is_verify,
- num_decodes,
- gen_p,
- int_idx,
- qsl_long_p,
- replay,
- )
-
- # PLE carry flags must cover the whole batch, but `has_initial_states` is
- # filled for the context prefix only (generation entries are stale); every
- # generation sequence carries, so pad with True (host-int sized, graph-safe).
- ple_carry = nn.functional.pad(hi_p, [0, num_decodes], mode="constant", value=1.0)
- # Under speculation `Mamba2Metadata.prepare`'s `query_start_loc_long` has one
- # row per generation sequence, so with `gen_p > 1` it is short and both PLE
- # stateful calls are silently mis-segmented; rebuild it from `seq_lens_cuda`,
- # which still holds the true counts during the trunk forward.
- if gen_p == 1:
- ple_qsl = qsl_long
- else:
- ple_qsl = cu_seqlens_from_seq_lens(md.seq_lens_cuda[:batch], torch.int64)
- ple_state = (core._bound.ple_state[0], core._bound.ple_state[1], si_all, ple_carry, ple_qsl)
-
- # Dense/sparse decision once per forward, from a host-side metadata read (as
- # QSA's backend does): a per-capture constant, matching the engine's separate
- # short/long decode graph families.
- sparse = batch > 0 and int(md.kv_lens_runtime[:batch].max()) > core.dense_threshold
-
- # The `[T, 3]` transpose of the `[3, 1, T]` position buffer, which the QSA position
- # plane stores.
- coords = torch.transpose(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1]), 0, 1)
- qsa = (
- md.qsa_req_idx_per_token[:num_tokens],
- md.qsa_logical_positions[:num_tokens],
- md.qsa_sequence_lengths[:num_tokens],
- md.qsa_visible_blocks[:num_tokens],
- md.qsa_topk_indices[:num_tokens],
- md.qsa_topk_row_starts[:num_tokens],
- md.qsa_block_table,
- md.tokens_per_block,
- coords[:num_tokens],
- sparse,
- _live_score_columns(core, md, batch),
- )
-
- # A verify step's generation rows write the PLE pools and the QSA side caches for
- # every draft row, accepted or not; `SpeculativeRollback` keeps what its commit
- # needs to undo the rejected ones. Reopened on every forward, so a step whose
- # commit never ran (a failed forward) leaves nothing for the next one. Context
- # rows are never staged.
- stage = is_verify
- rollback = core.spec_rollback
- if rollback is not None:
- rollback.discard()
- if stage:
- rollback.begin(
- num_contexts=num_prefills,
- prefix_tokens=num_prefill_tokens,
- num_gens=num_decodes,
- rows_per_gen=gen_p,
- slots=si_d,
- request_indices=md.qsa_req_idx_per_token[num_prefill_tokens:num_tokens],
- logical_positions=md.qsa_logical_positions[num_prefill_tokens:num_tokens],
- block_table=md.qsa_block_table,
- tokens_per_block=md.tokens_per_block,
- )
-
- nvfp4_context = None
- if (
- sparse
- and core.quant_mode == QM_NVFP4_KV_CACHE
- and num_prefill_tokens > _NVFP4_ROWTILE_MIN_ROWS
- ):
- nvfp4_context = _nvfp4_context_plan(
- core, md, qsa[0], md.qsa_block_table, md.tokens_per_block
- )
-
- self._step = _PrefillStep(
- num_tokens=num_tokens,
- physical_tokens=physical_tokens,
- input_ids=input_ids[:num_tokens],
- attn=attn_args,
- gdn=gdn_state,
- ple=ple_state,
- qsa=qsa,
- stage=stage,
- vision=vision,
- nvfp4_context=nvfp4_context,
- )
-
- def forward(
- self,
- attn_metadata: AttentionMetadata,
- input_ids: torch.IntTensor,
- position_ids: torch.IntTensor,
- **kwargs,
- ):
- """The step `prepare` built, over the `T` activation rows the engine handed in: the
- batch's `n` rows, or the prefill CUDA graph bucket they were padded to.
-
- **This is the body a breakable prefill CUDA graph captures**, once per bucket, and
- replays for every step it pads to that bucket without running this Python. So
- nothing here reads the step: every op is row-wise and sized by `T` alone, its inputs
- are the engine's persistent `input_ids` / `position_ids` buffers, and every buffer is
- allocated here. What depends on the step's batch runs in one of four breaks that
- run eagerly between the captured segments and read `self._step`: `_splice` (once),
- `_ple_block` (once), `_attention_core` (every attention layer) and `_gdn_core` (every
- Gated-DeltaNet layer). Each writes rows `[:n]` of a fixed output allocated here and
- zeroes rows `[n:]`, so the padding rows stay finite through the row-wise ops and
- reach neither a state pool nor a real row. With the backend off every step is
- eager and unpadded, and the breaks are plain calls.
-
- Every hyper-connection site is `qwen4_exp_hc_inject_norm` (the previous block's
- injection fused with this site's grouped norm; after a MoE block it also forms that
- block's output) and then the mixer (`_hc_mixer`, `_hc_moe_mixer`). From
- `_HC_RSTD_MIN_ROWS` rows the site is `qwen4_exp_hc_inject_rstd` (the same
- injection, the norm's `[T, C]` rstd instead of the normed stream) and the rstd
- mixer (`_hc_mixer_rstd`, `_hc_moe_mixer_rstd`); the PLE layer's attention-side
- site keeps the stored-normed mixer either way, on the normed stream the PLE
- kernel writes."""
- core = self.core
- num_tokens = input_ids.shape[0]
- assert num_tokens == self._step.physical_tokens, (
- f"the forward got {num_tokens} activation rows and prepare saw "
- f"{self._step.physical_tokens}"
- )
-
- # `pos3`: contiguous `[3, T]` for `use_mrope=True` (the engine hands a slice of a
- # larger buffer, whose row stride the ops reject, hence the copy).
- pos3 = torch.empty(
- [_QSA_POSITION_AXES, position_ids.shape[-1]],
- dtype=torch.int32,
- device=position_ids.device,
- )
- pos3.copy_(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1]))
-
- # The token embedding with this step's vision rows spliced over their placeholders.
- h = nn.functional.embedding(input_ids, core.w["embed"])
- eager_on_graph(h, inner=self._splice)
- dt = h.dtype
- dev = h.device
-
- # Scratch shared by every layer that uses it; each break rewrites its whole output.
- proj_qkv = torch.empty([num_tokens, core.qkv_width], dtype=dt, device=dev)
- index_scratch = torch.empty([num_tokens, core.index_scratch_width], dtype=dt, device=dev)
- attn_out = torch.empty([num_tokens, core.q_width], dtype=dt, device=dev)
- gdn_out = torch.empty(
- [1, num_tokens, core.num_v_heads, core.head_v_dim], dtype=dt, device=dev
- )
- ple_hc = torch.empty([num_tokens, core.stream], dtype=dt, device=dev)
- ple_normed = torch.empty([num_tokens, core.stream], dtype=dt, device=dev)
- # The CuteDSL routed output, shared by all layers: each layer's unpermute rewrites
- # every row, and the next site's injection consumes it before the next layer's MoE.
- routed_out = (
- torch.empty([num_tokens, core.hidden], dtype=dt, device=dev)
- if core.trunk_moe_backend == "CUTEDSL"
- else None
- )
-
- # The embedding repeated across the hc_count residual streams.
- hc_count = core.hc_count
- hc = torch.cat([h] * hc_count, -1)
- layers, S = self._layers, self._stack
- # A host decision on `T` alone (each captured bucket takes one branch).
- by_rstd = num_tokens >= _HC_RSTD_MIN_ROWS
- normed = rstd = None
- if by_rstd:
- _, rstd = qwen4_exp_hc_inject_rstd(hc, None, None, core.eps, hc_count)
- else:
- _, normed = qwen4_exp_hc_inject_norm(
- hc, None, None, layers[0].ahc_norm, core.eps, hc_count
- )
- for i, L in enumerate(layers):
- if i == core.ple_layer:
- eager_on_graph(hc, ple_hc, ple_normed, inner=self._ple_block)
- hc, normed = ple_hc, ple_normed
- # The PLE kernel stores this site's normed stream: the stored-normed
- # mixer reads it (cheaper than a second pass for its rstd).
- mixed, inject = self._hc_mixer(normed, L.ahc_first, L.ahc_up, True)
- elif by_rstd:
- mixed, inject = self._hc_mixer_rstd(
- hc, rstd, L.ahc_norm, L.ahc_norm_w1, L.ahc_first, L.ahc_up, True
- )
- else:
- mixed, inject = self._hc_mixer(normed, L.ahc_first, L.ahc_up, True)
- if core.is_attn[i]:
- block_out = self._attention_block(
- mixed, i, L, proj_qkv, index_scratch, attn_out, pos3
- )
- else:
- block_out = self._gdn_block(mixed, i, L, gdn_out)
- if by_rstd:
- hc, rstd = qwen4_exp_hc_inject_rstd(hc, block_out, inject, core.eps, hc_count)
- mixed, inject, xq, xsf = self._hc_moe_mixer_rstd(
- hc, rstd, L.mhc_norm, L.mhc_norm_w1, L.mhc_first, L.mhc_up, L.moe_gscale
- )
- else:
- hc, normed = qwen4_exp_hc_inject_norm(
- hc, block_out, inject, L.mhc_norm, core.eps, hc_count
- )
- mixed, inject, xq, xsf = self._hc_moe_mixer(
- normed, L.mhc_first, L.mhc_up, L.moe_gscale
- )
- next_norm = layers[i + 1].ahc_norm if i + 1 < len(layers) else S.fhc_norm
- # The MoE block output is formed inside the injection: routed +
- # sigmoid(gate) * shared, rounded once. Both [T, hidden] operands are
- # released here, not at the next layer's MoE block, whose own peak they
- # would otherwise add to.
- routed, (gate_logits, shared) = self._moe_block(mixed, xq, xsf, L, routed_out)
- if by_rstd:
- hc, rstd = qwen4_exp_hc_inject_rstd(
- hc, routed, inject, core.eps, hc_count, shared, gate_logits
- )
- else:
- hc, normed = qwen4_exp_hc_inject_norm(
- hc, routed, inject, next_norm, core.eps, hc_count, shared, gate_logits
- )
- del routed, gate_logits, shared
- # The stack-final mixer with no injection weight is the final norm.
- if by_rstd:
- final, _ = self._hc_mixer_rstd(
- hc, rstd, S.fhc_norm, S.fhc_norm_w1, S.fhc_first, S.fhc_up, False
- )
- else:
- final, _ = self._hc_mixer(normed, S.fhc_first, S.fhc_up, False)
- if not core.mtp_enabled:
- return final
- # `lm_head` reads the `[T, hidden]` mix; the MTP module reads the
- # `[T, hc_count * hidden]` stream it was collapsed from. They leave packed in one
- # tensor, because a breakable prefill graph's body returns one tensor; the shell
- # takes the two column views (`split_spec_output`).
- return torch.cat([final, hc], -1)
-
- # ------------------------------------------------------------- the draft ---
-
- def _mtp_site(self, hc, block_out, inject, norm, first, up) -> tuple:
- """One of the draft layer's hyper-connection sites on the trunk's prefill kernels:
- `block_out` injected into the stream `hc`, then the mixer (with injection) from the
- stored normed stream, or from its rstd from `_HC_RSTD_MIN_ROWS` rows. Returns
- `(hc, mixed, inject)`."""
- core = self.core
- if hc.shape[0] < _HC_RSTD_MIN_ROWS:
- hc, normed = qwen4_exp_hc_inject_norm(
- hc, block_out, inject, norm, core.eps, core.hc_count
- )
- return (hc, *self._hc_mixer(normed, first, up, True))
- w1 = self._mtp_norm_w1.get(norm.data_ptr())
- if w1 is None:
- w1 = self._mtp_norm_w1[norm.data_ptr()] = _norm_w1(norm)
- hc, rstd = qwen4_exp_hc_inject_rstd(hc, block_out, inject, core.eps, core.hc_count)
- return (hc, *self._hc_mixer_rstd(hc, rstd, norm, w1, first, up, True))
-
- def _mtp_moe_block(self, x, dt) -> torch.Tensor:
- """The draft layer's MoE: FP8 block-scale routed experts plus a bf16 shared expert.
-
- The scale buffer comes back transposed `[hidden/128, tokens]` and is passed
- straight through; never reshape it. At `tokens == hidden/128` (20, an
- ordinary decode batch) the shapes coincide, so nothing would catch a
- re-materialization. Pre-routed; `routing_bias` must be `None` (the entry
- refuses one). With `_MTP_MOE_NVFP4_ENV` the routed experts run NVFP4
- (`_mtp_moe_nvfp4_block`)."""
- core = self.core
- M = core._mtp
- if M.moe_gscale is not None:
- return _mtp_moe_nvfp4_block(core, M, x)
- logits = cublas_mm(x, M.router)
- topk_ids, topk_w = renorm_moe_routing_op(logits, topk=core.topk, output_dtype=dt)
- xq, xsf = fp8_quantize_1x128(x, use_ue8m0=_MTP_MOE_USE_UE8M0)
- routed = fp8_block_scale_moe_runner(
- routing_logits=None,
- routing_bias=None,
- hidden_states=xq,
- hidden_states_scale=xsf,
- gemm1_weights=M.fc1_w,
- gemm1_weights_scale=M.fc1_s,
- gemm2_weights=M.fc2_w,
- gemm2_weights_scale=M.fc2_s,
- num_experts=core.num_experts,
- top_k=core.topk,
- num_fused_shared_experts=None,
- n_group=None,
- topk_group=None,
- intermediate_size=core.moe_inter,
- local_expert_offset=0,
- local_num_experts=core.num_experts,
- routed_scaling_factor=None,
- routing_method_type=_ROUTING_METHOD_RENORM,
- topk_weights=topk_w,
- topk_ids=topk_ids,
- act_type=_ACT_TYPE_SWIGLU,
- gemm1_clamp_limit=None,
- )
- shared = cublas_mm(flashinfer_silu_and_mul(cublas_mm(x, M.sh_gu)), M.sh_dn)
- gate_logits = cublas_mm(x, M.sh_gate)
- return fused_sigmoid_gate_mul_add(routed, gate_logits, shared)
-
- def _capture_selection(self, md, topk_indices, visible_blocks, compact=False) -> None:
- """IndexShare: stash, per request, the selection of the row the draft loop continues
- from, with that row's complete-group count, in the draft layer's row of the shared
- buffers.
-
- That row is the one the worker gathers for the next draft step: a context
- request's last row of this chunk, and a generation request's last accepted row
- (`num_accepted - 1` past its first). Capturing here rather than at the first draft
- decode step keeps the indexer at one call per draft loop for a batch with context
- rows too. A `compact` selection holds each context request's last row, then every
- generation row (the draft step's `keep`)."""
- batch = md.num_seqs
- assert md.qsa_mtp_num_accepted is not None, (
- "IndexShare needs the worker's accepted counts to pick each request's capture "
- "row; the draft loop did not supply them"
- )
- seq_lens = md.seq_lens_cuda[:batch].to(torch.int64)
- starts = torch.cumsum(seq_lens, 0) - seq_lens
- accepted = md.qsa_mtp_num_accepted[:batch].to(torch.int64)
- offsets = torch.minimum(torch.clamp(accepted - 1, min=0), seq_lens - 1)
- offsets[: md.num_contexts] = seq_lens[: md.num_contexts] - 1
- rows = starts + offsets
- if compact:
- num_contexts = md.num_contexts
- rows[:num_contexts] = torch.arange(num_contexts, device=rows.device)
- rows[num_contexts:] -= md.num_ctx_tokens - num_contexts
- slot = self.core._draft_bound.share_slot
- md.qsa_shared_topk_indices[slot, :batch].copy_(torch.index_select(topk_indices, 0, rows))
- md.qsa_shared_topk_visible_blocks[slot, :batch].copy_(
- torch.index_select(visible_blocks, 0, rows)
- )
- # Host-side readiness, recorded in program order: under capture this runs while the
- # graph is recorded, so a graph holds the reuse only after the capture it reads.
- md.qsa_shared_topk_captured_slots.add(slot)
-
- def mtp_forward(
- self,
- embed_tokens: torch.Tensor,
- all_rank_num_tokens,
- input_ids: torch.Tensor,
- position_ids: torch.Tensor,
- hidden_states: torch.Tensor,
- attn_metadata: AttentionMetadata,
- **kwargs,
- ) -> torch.Tensor:
- """One draft step whose metadata still carries context rows.
-
- `spec_metadata` and `all_rank_num_tokens` are accepted and unused."""
- core = self.core
- assert isinstance(attn_metadata, TrtllmAttentionMetadata)
- assert core._bound is not None, (
- "the trunk's cache-manager bind has not run; the shell calls the "
- "core before the worker, so this cannot be reached first"
- )
- md = attn_metadata
- rows = hidden_states.shape[0]
- assert input_ids.shape[0] == rows, (
- f"the draft step's {input_ids.shape[0]} token ids and {rows} "
- "hidden-state rows must describe the same tokens"
- )
- assert hidden_states.shape[1] == core.stream, (
- f"the MTP module reads the trunk's {core.stream}-wide four-stream "
- f"residual; the runtime handed it {hidden_states.shape[1]} columns"
- )
- dt = hidden_states.dtype
- dev = hidden_states.device
-
- # Projected on every call: the worker rewrites `attn_metadata` in place
- # between draft step 0 and 1+ (`_seq_lens` to 1, `num_contexts` and
- # `num_ctx_tokens` to 0), so a first-call value is wrong on the rest.
- step = self.step_args(md)
- M = core._mtp
- num_ctx = md.num_ctx_tokens
- gen_rows = rows - num_ctx
- gen_seqs = md.num_seqs - md.num_contexts
- p = gen_rows // gen_seqs if gen_seqs else 1
- assert p * gen_seqs == gen_rows, (
- f"{gen_rows} generation rows do not divide over {gen_seqs} "
- "generation sequences; a ragged per-sequence draft length cannot be "
- "expressed to the attention call"
- )
-
- # The KV cache written here is not the trunk's: with a separate draft KV
- # manager the worker swaps `attn_metadata.kv_cache_manager` and
- # `kv_cache_block_offsets` for the draft loop. The layer index must be taken
- # from the current manager every call; caching the trunk's addresses another
- # layer's pages with finite, plausible drafts.
- cm = md.kv_cache_manager
- assert cm is not None, "a paged KV cache is required"
- offsets = cm.layer_offsets
- assert core.mtp_layer_idx in offsets, (
- f"the MTP module's attention is global layer index "
- f"{core.mtp_layer_idx}, which the cache manager "
- f"{type(cm).__name__} does not map (it maps {len(offsets)} layers, "
- f"is_draft={getattr(cm, 'is_draft', None)}); the engine allocated "
- "no KV layer for the draft module"
- )
- layer_idx = offsets[core.mtp_layer_idx]
- window = cm.max_attention_window_vec[layer_idx]
- assert window is None or window == md.max_seq_len, (
- f"the MTP layer's pool window is {window!r} but this call passes "
- f"max_seq_len {md.max_seq_len} as the attention window; a sliding "
- "window is not implemented here"
- )
-
- # The draft layer's QSA planes and the page table that addresses them. With a
- # separate draft manager the engine's QSA metadata swaps `qsa_block_table` for the
- # table it decodes from the draft manager's offsets (`prepare_for_draft_forward`);
- # a table still decoding the trunk's offsets would read and write another pool's
- # pages.
- draft = core._draft_bound
- if draft is None or cm is not draft.cm:
- core._bind_draft_cache_manager(md)
- draft = core._draft_bound
- assert cm is core._bound.cm or md.qsa_block_table is getattr(
- md, "qsa_draft_block_table", None
- ), (
- "the draft step is bound to a separate draft KV cache manager but "
- "qsa_block_table is not that manager's table: the engine's QSA metadata did "
- "not swap its page table for the draft forward"
- )
- # The `qsa_*` row mapping is sized from the metadata's token count, which the
- # worker rewrites together with it between draft steps.
- assert rows == md.num_tokens, (
- f"the draft step carries {rows} rows for a batch its metadata describes as "
- f"{md.num_tokens} tokens; the QSA row mapping covers the second"
- )
- # The trunk's dense/sparse decision from the same host buffer: the worker advances
- # only the device KV lengths through the draft loop, so every draft step of one
- # engine forward takes the branch its trunk took, in the CUDA-graph family the
- # runner chose for that trunk.
- batch = md.num_seqs
- sparse = batch > 0 and int(md.kv_lens_runtime[:batch].max()) > core.dense_threshold
- # IndexShare: the only draft step with context rows is the loop's first, its
- # draft-extend step. It always runs the indexer and, with sharing on, captures the
- # selection the draft decode steps reuse.
- share = core.mtp_index_share and md.qsa_in_mtp_draft_loop
- assert not (share and md.qsa_indexer_skip_topk), (
- "the worker asked a draft step with context rows to reuse a selection; only "
- "the draft decode steps after it reuse"
- )
-
- # Copy into a fresh contiguous `[3, N]`: a sliced view keeps the source row
- # stride and the op rejects it under `use_mrope`.
- assert position_ids.shape[0] == _QSA_POSITION_AXES, (
- f"expected the mRoPE position layout with {_QSA_POSITION_AXES} "
- f"coordinate rows, got {tuple(position_ids.shape)}"
- )
- pos3 = torch.empty(
- [_QSA_POSITION_AXES, position_ids.numel() // _QSA_POSITION_AXES],
- dtype=torch.int32,
- device=position_ids.device,
- )
- pos3.copy_(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1]))
- # The trunk's two other views, from the copy: `pos` (coordinate 0 alone, which the
- # index rope reads) and `coords` (the `[N, 3]` transpose the position plane stores).
- pos = torch.reshape(torch.split(pos3, [1, _QSA_POSITION_AXES - 1], 0)[0], [-1])
- coords = torch.transpose(pos3, 0, 1)
-
- # ---- the front end --------------------------------------------------
- e = nn.functional.embedding(input_ids, embed_tokens)
- x_e = cublas_mm(flashinfer_gemma_rmsnorm(e, M.pre_e_norm, core.eps), M.fc_e)
- # `fc_hidden` is one `[hidden, hidden]` matrix applied to each of the four
- # streams (the checkpoint's norm statistics say they are distinct streams).
- h_n = self._grouped_norm(hidden_states, M.pre_h_norm)
- x_h = torch.reshape(
- cublas_mm(torch.reshape(h_n, [rows * core.hc_count, core.hidden]), M.fc_h),
- [rows, core.stream],
- )
- # Both hyper-connection sites run the trunk's site and mixers, as
- # `DecodeTarget._mtp_layer` does: the front end's per-stream `x_h + x_e` is the first
- # site's injection at unit weight.
- F = core._mtp_fused
- x, mixed, inject = self._mtp_site(
- x_h, x_e, F.unit_inject[:rows], F.ahc_norm, F.ahc_first, F.ahc_up
- )
-
- # ---- the decoder layer ----------------------------------------------
- # The draft loop reads one row per request from this step: a context request's last
- # row and one row of each generation run (the worker's `gather_ids`). Every row writes
- # the layer's KV and index planes; past those writes only `keep` -- each context's
- # last row, then every generation row -- goes on, the query side included on the
- # sparse path (the dense call's metadata describes every row). The other returned
- # rows keep the layer's input stream.
- keep = None
- if num_ctx > md.num_contexts:
- ctx_lens = md.seq_lens_cuda[: md.num_contexts].to(torch.long)
- keep = torch.cat(
- [torch.cumsum(ctx_lens, 0) - 1, torch.arange(num_ctx, rows, device=dev)]
- )
- q_keep = keep if sparse else None
- q_rows = rows if q_keep is None else q_keep.shape[0]
- qkv = torch.empty([rows, core.qkv_width], dtype=dt, device=dev)
- attn_out = torch.empty([q_rows, core.q_width], dtype=dt, device=dev)
- qg = cublas_mm(mixed if q_keep is None else mixed[q_keep], M.q)
- k = cublas_mm(mixed, M.k)
- v = cublas_mm(mixed, M.v)
- qkv_h = torch.reshape(qkv, [rows, core.heads_q + 2 * core.heads_kv, core.head_dim])
- q_dst, k_dst, v_dst = torch.split(qkv_h, [core.heads_q, core.heads_kv, core.heads_kv], 1)
- q_src, gate = torch.split(
- torch.reshape(qg, [q_rows, core.heads_q, 2 * core.head_dim]),
- [core.head_dim, core.head_dim],
- 2,
- )
- if q_keep is None:
- q_dst.copy_(q_src)
- else:
- # The other rows' query heads stay unwritten: only the rope reads them, row-local.
- q_dst.index_copy_(0, q_keep, q_src)
- k_dst.copy_(torch.reshape(k, [rows, core.heads_kv, core.head_dim]))
- v_dst.copy_(torch.reshape(v, [rows, core.heads_kv, core.head_dim]))
- # The trunk rope's certified cell at the module's own base (`mtp.rope_theta`).
- fused_qk_norm_rope(
- qkv,
- num_heads_q=core.heads_q,
- num_heads_k=core.heads_kv,
- num_heads_v=core.heads_kv,
- head_dim=core.head_dim,
- rotary_dim=core.rotary_dim,
- eps=core.eps,
- q_weight=M.q_norm,
- k_weight=M.k_norm,
- base=core.mtp_theta,
- is_neox=True,
- position_ids=pos3,
- is_qk_norm=True,
- use_gemma=True,
- use_mrope=True,
- mrope_section1=core.mrope_section1,
- mrope_section2=core.mrope_section2,
- )
- # The QSA index branch, as in the trunk's `_attention_block`: the raw key stored
- # before anything normalizes or rotates it, the query's own norm and rotation,
- # then the compression of every group this step completes. Maintained on every
- # draft step at both ends of the threshold, so the draft plane is complete when a
- # request crosses it.
- req_idx = md.qsa_req_idx_per_token[:rows]
- logical = md.qsa_logical_positions[:rows]
- block_table = md.qsa_block_table
- tokens_per_block = md.tokens_per_block
- index_cache = draft.index_cache
- qk_index = cublas_mm(mixed, M.idx)
- token_k = torch.split(
- torch.reshape(qk_index, [rows, core.index_heads + core.index_kv_heads, -1]),
- [core.index_heads, core.index_kv_heads],
- 1,
- )[1]
- qsa_paged_index_store(
- token_k,
- coords,
- index_cache,
- draft.position_cache,
- req_idx,
- logical,
- block_table,
- tokens_per_block,
- )
- index_scratch = torch.empty([q_rows, core.index_scratch_width], dtype=dt, device=dev)
- index_scratch[:, : core.index_width].copy_(qk_index if q_keep is None else qk_index[q_keep])
- q_pos3, q_pos, q_req, q_logical = pos3, pos, req_idx, logical
- if q_keep is not None:
- q_pos3 = torch.index_select(pos3, 1, q_keep)
- q_pos, q_req, q_logical = pos[q_keep], req_idx[q_keep], logical[q_keep]
- fused_qk_norm_rope(
- index_scratch,
- num_heads_q=core.index_heads,
- num_heads_k=core.index_kv_heads,
- num_heads_v=core.index_kv_heads,
- head_dim=core.index_head_dim,
- rotary_dim=core.rotary_dim,
- eps=core.eps,
- q_weight=M.idx_qn,
- k_weight=M.idx_kn,
- base=core.mtp_theta,
- is_neox=True,
- position_ids=q_pos3 if _QSA_INDEX_USE_MROPE else q_pos,
- is_qk_norm=True,
- use_gemma=True,
- use_mrope=_QSA_INDEX_USE_MROPE,
- mrope_section1=(_QSA_MROPE_SECTION or (0, 0, 0))[1],
- mrope_section2=(_QSA_MROPE_SECTION or (0, 0, 0))[2],
- )
- q_index = torch.split(
- torch.reshape(
- index_scratch,
- [q_rows, core.index_heads + 2 * core.index_kv_heads, core.index_head_dim],
- ),
- [core.index_heads, core.index_kv_heads, core.index_kv_heads],
- 1,
- )[0]
- qsa_prefill_compress(
- logical_positions=logical,
- request_indices=req_idx,
- block_table=block_table,
- index_cache=index_cache,
- position_cache=draft.position_cache,
- k_norm_weight=M.idx_kn,
- cos_sin=core._mtp_cos_sin,
- eps=core.eps,
- tokens_per_block=tokens_per_block,
- compress_ratio=core.compress_ratio,
- mrope_section=_QSA_MROPE_SECTION,
- )
-
- if sparse:
- kv = draft.kv
- # The append the dense op would have owned, then selection over the draft
- # layer's own compressed keys and the sparse GQA over its pool.
- _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block)
- row_starts = md.qsa_topk_row_starts[:rows]
- visible_blocks = md.qsa_visible_blocks[:rows]
- seq_lengths = md.qsa_sequence_lengths[:rows]
- if q_keep is not None:
- row_starts, visible_blocks = row_starts[q_keep], visible_blocks[q_keep]
- seq_lengths = seq_lengths[q_keep]
- scores = qsa_paged_index_scores(
- q=q_index,
- index_cache=index_cache,
- block_table=block_table,
- query_positions=q_logical,
- request_indices=q_req,
- tokens_per_block=tokens_per_block,
- compress_ratio=core.compress_ratio,
- only_visible_blocks=True,
- # A context's run of rows shares one key gather; `keep` has none.
- context_rows=q_keep is None,
- )
- topk_indices = md.qsa_topk_indices[:q_rows]
- _select(
- core,
- scores,
- row_starts,
- visible_blocks,
- topk_indices,
- _live_score_columns(core, md, batch),
- )
- if share:
- self._capture_selection(md, topk_indices, visible_blocks, q_keep is not None)
- selected = qsa_expand_block_indices(
- topk_indices,
- q_logical,
- seq_lengths,
- compress_ratio=core.compress_ratio,
- token_topk=core.token_topk,
- )
- q_att = q_dst if q_keep is None else q_dst[q_keep]
- if kv.k_scale is not None:
- # `keep` holds one row per context request.
- q_ctx = num_ctx if q_keep is None else md.num_contexts
- plan = None
- if q_ctx > _NVFP4_ROWTILE_MIN_ROWS:
- plan = _nvfp4_context_plan(
- core,
- md,
- req_idx,
- block_table,
- tokens_per_block,
- one_row_per_request=q_keep is not None,
- )
- _sparse_gqa_nvfp4_context(
- core,
- kv,
- plan,
- q_att,
- block_table,
- selected,
- q_req,
- q_logical,
- tokens_per_block,
- q_ctx,
- attn_out,
- rows_per_request=p,
- )
- else:
- attn_out = torch.reshape(
- _sparse_gqa(
- core,
- kv,
- q=q_att,
- block_table=block_table,
- selected_tokens=selected,
- request_indices=q_req,
- tokens_per_block=tokens_per_block,
- query_positions=q_logical,
- compress_ratio=core.compress_ratio,
- ),
- [q_rows, core.q_width],
- )
- else:
- attn_out = thop_attention(
- q=qkv,
- output=attn_out,
- local_layer_idx=layer_idx,
- predicted_tokens_per_seq=p,
- num_heads=core.heads_q,
- num_kv_heads=core.heads_kv,
- head_size=core.head_dim,
- quant_mode=core.quant_mode,
- kv_scale_orig_quant=core.kv_scale_orig_quant,
- kv_scale_quant_orig=core.kv_scale_quant_orig,
- **step,
- **_CALL_CONSTANTS,
- )
- fused_sigmoid_mul_inplace(attn_out, gate)
- if keep is not None:
- if q_keep is None:
- attn_out = attn_out[keep]
- out, inject = x, inject[keep]
- x = x[keep]
- x, mixed, inject = self._mtp_site(
- x, cublas_mm(attn_out, M.o), inject, F.mhc_norm, F.mhc_first, F.mhc_up
- )
- x = qwen4_exp_post_inject(x, self._mtp_moe_block(mixed, dt), inject)
- return x if keep is None else out.index_copy_(0, keep, x)
-
-
-def _position_key(position_ids: torch.Tensor) -> tuple:
- """What identifies a step's positions buffer: its storage, dtype and `[3, N]` layout."""
- view = torch.reshape(position_ids, [_QSA_POSITION_AXES, -1])
- return (view.data_ptr(), view.dtype, tuple(view.shape), view.stride())
-
-
-def _copy_positions(position_ids: torch.Tensor) -> torch.Tensor:
- """A fresh contiguous int32 `[3, N]` copy of the step's mRoPE positions."""
- pos3 = torch.empty(
- [_QSA_POSITION_AXES, position_ids.numel() // _QSA_POSITION_AXES],
- dtype=torch.int32,
- device=position_ids.device,
- )
- pos3.copy_(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1]))
- return pos3
-
-
-class DecodeTarget:
- """Generation-only batches (no context rows): drops the context-row recurrent
- ops and adds the fused single-token recurrence."""
-
- def __init__(self, core: Qwen4ExpTp1ModelingV2) -> None:
- """Take the per-layer weight views once; MTP operands are read off `core._mtp`."""
- self.core = core
- # See PrefillTarget.__init__.
- self._contract_pending = step_contract_enabled()
- self._kv_probe = (
- _KvRangeProbe(core.attn_layers[-1])
- if self._contract_pending and core.quant_mode == QM_FP8_KV_CACHE
- else None
- )
- self._layers = _layer_views(core)
- self._stack = _stack_views(core)
- # `(position key, pos3)` of the last trunk step: the first draft step of the same
- # forward reads the trunk's positions buffer and reuses its int32 copy.
- self._trunk_pos3: tuple | None = None
-
- def _check_step_contract(self, md) -> None:
- """Opt-in first-forward fail-fast; see PrefillTarget._check_step_contract."""
- if not self._contract_pending:
- return
- try:
- self.step_args(md)
- except AttributeError as exc:
- raise AssertionError(f"attention metadata surface drifted: {exc}") from exc
- self._contract_pending = False
-
- def step_args(self, md: TrtllmAttentionMetadata) -> dict:
- """See PrefillTarget.step_args. Context counts are stated, not read: routing
- guarantees them, and a decode-only CUDA graph would freeze whatever it read."""
- return dict(
- sequence_length=md.kv_lens_cuda_runtime,
- host_past_key_value_lengths=md.kv_lens_runtime,
- host_total_kv_lens=md.host_total_kv_lens,
- context_lengths=md.prompt_lens_cuda_runtime,
- host_context_lengths=md.prompt_lens_cpu_runtime,
- host_request_types=md.host_request_types_runtime,
- kv_cache_block_offsets=md.kv_cache_block_offsets,
- host_kv_cache_pool_pointers=md.host_kv_cache_pool_pointers,
- host_kv_cache_pool_mapping=md.host_kv_cache_pool_mapping,
- workspace=md.effective_workspace,
- tokens_per_block=md.tokens_per_block,
- max_num_requests=md.max_num_requests,
- max_context_length=md.max_context_length,
- max_seq_len=md.max_seq_len,
- # No sliding window: the full admitted length on every attention layer.
- attention_window_size=md.max_seq_len,
- num_contexts=0,
- num_ctx_tokens=0,
- trtllm_gen_jit_warmup=md.trtllm_gen_jit_warmup,
- use_paged_context_fmha=md.use_paged_context_fmha,
- beam_width=md.effective_beam_width,
- cache_indirection=md.cache_indirection,
- block_ids_per_seq=md.block_ids_per_seq,
- max_context_q_len_override=md.max_context_q_len_override,
- is_cross=md.is_cross,
- is_spec_decoding_enabled=md.is_spec_decoding_enabled,
- use_spec_decoding=md.use_spec_decoding,
- is_spec_dec_tree=md.is_spec_dec_tree,
- spec_decoding_generation_lengths=md.spec_decoding_generation_lengths,
- spec_decoding_position_offsets=get_spec_decoding_position_offsets(md),
- spec_decoding_packed_mask=md.spec_decoding_packed_mask,
- spec_decoding_bl_tree_mask_offset=md.spec_decoding_bl_tree_mask_offset,
- spec_decoding_bl_tree_mask=md.spec_decoding_bl_tree_mask,
- spec_decoding_target_max_draft_tokens=md.max_total_draft_tokens,
- spec_bl_tree_first_sparse_mask_offset_kv=md.spec_bl_tree_first_sparse_mask_offset_kv,
- num_sparse_topk=md.num_sparse_topk,
- flash_mla_tile_scheduler_metadata=md.flash_mla_tile_scheduler_metadata,
- flash_mla_num_splits=md.flash_mla_num_splits,
- max_num_sequences=md.max_num_sequences,
- force_prepare_spec_dec_tree_mask=md.force_prepare_spec_dec_tree_mask,
- )
-
- # --------------------------------------------------------------- blocks ---
-
- def _hc_first(self, normed, first, with_inject):
- """The mixer's first half from the site's normed stream: the projection
- over `[down; block_inject; 0]` and the SiLU / injection activations.
- Returns `(low, inject)`. Up to `_HC_TC_FIRST_MAX_ROWS` rows (a per-capture
- constant) one tcgen05 launch, `qwen4_exp_hc_decode_tc_first`; larger batches
- cuBLAS's fp32 logits and `qwen4_exp_hc_lowrank_act`. Only sites past
- `_HC_DECODE_GEMV_MAX_ROWS` rows reach a mixer."""
- core = self.core
- if normed.shape[0] <= _HC_TC_FIRST_MAX_ROWS:
- return qwen4_exp_hc_decode_tc_first(
- normed, first, core.hc_lowrank, core.hc_count, with_inject
- )
- return qwen4_exp_hc_lowrank_act(
- cublas_mm(normed, first[:, : core.hc_first_cols], None, torch.float32),
- core.hc_lowrank,
- core.hc_count,
- with_inject,
- normed.dtype,
- )
-
- def _hc_mixer(self, normed, first, up, with_inject):
- """One hyper-connection mixer from the site's normed stream: the first
- projection over `[down; block_inject; 0]` (fp32 logits), the SiLU and
- injection activations, then the up projection with the sigmoid gate and the
- stream mean in its epilogue. Returns `(mixed, inject)`; `inject` is None
- without `with_inject` (the stack-final mixer).
-
- Neither projection takes the step's `gemm` (`low_m_bf16_mm` under CUDA
- graphs): neither FlashInfer low-M kernel writes the first projection's
- fp32 logits, and the up projection's K = 320 tiles neither of them."""
- core = self.core
- low, inject = self._hc_first(normed, first, with_inject)
- return qwen4_exp_hc_up_gate_mean(low, up, normed, core.hc_count), inject
-
- def _hc_site(self, hc, block_out, inject, norm, first, up, with_inject, normed=None, moe=None):
- """One trunk hyper-connection site: the previous block's injection (none at
- the stack's first site and after the PLE add), this site's grouped norm and
- its mixer. Returns `(hc, mixed, inject)`. The row count is a per-capture
- constant: at most `_HC_DECODE_GEMV_MAX_ROWS` rows run two GEMV launches,
- larger batches `qwen4_exp_hc_inject_norm` and `_hc_mixer`. `normed`, when
- given, is this site's grouped norm of `hc` already computed (the PLE tail's,
- with no injection pending); the cuBLAS path starts from it.
-
- After a MoE block, `block_out` is its routed output and `moe` the
- `(gate_logits, shared, topk_weights, topk_rows)` `_moe_block` left
- uncombined: the site's first kernel forms the block output itself. With
- the top-k pair `block_out` is the runner's permuted FC2 rows and that
- kernel (on every route) also forms their finalize. A site handed
- `normed` has no pending block, so it never carries `moe`."""
- core = self.core
- gate_logits, shared, topk_w, topk_rows = moe if moe is not None else (None,) * 4
- if _HC_DECODE_GEMV_MAX_ROWS < hc.shape[0] <= _HC_DECODE_ROWS_MAX:
- # Two tensor-core launches; with the top-k pair the first also forms the finalize.
- n_cols = core.hc_lowrank + (core.hc_count if with_inject else 0)
- hc, low, inject, rstd = qwen4_exp_hc_decode_rows_first(
- hc,
- block_out,
- inject,
- norm,
- first,
- core.hc_count,
- n_cols,
- core.eps,
- shared,
- gate_logits,
- topk_w,
- topk_rows,
- )
- mixed, _, _ = qwen4_exp_hc_decode_rows_mix(low, rstd, up, hc, norm, core.hc_count)
- return hc, mixed, inject
- if hc.shape[0] <= _HC_DECODE_GEMV_MAX_ROWS:
- n_cols = core.hc_lowrank + (core.hc_count if with_inject else 0)
- hc, part, ss = qwen4_exp_hc_decode_first_cute(
- hc,
- block_out,
- inject,
- norm,
- first,
- core.hc_count,
- n_cols,
- shared,
- gate_logits,
- topk_w,
- topk_rows,
- )
- mixed, inject = qwen4_exp_hc_decode_mix_cute(
- part, ss, up, hc, norm, core.eps, core.hc_count, with_inject
- )
- return hc, mixed, inject
- if normed is None:
- hc, normed = qwen4_exp_hc_inject_norm(
- hc,
- block_out,
- inject,
- norm,
- core.eps,
- core.hc_count,
- shared,
- gate_logits,
- topk_w,
- topk_rows,
- )
- mixed, inject = self._hc_mixer(normed, first, up, with_inject)
- return hc, mixed, inject
-
- def _hc_moe_site(self, hc, block_out, inject, norm, first, up, gscale):
- """`_hc_site` at an MoE site (always with injection), returning the
- routed experts' NVFP4 input (linear scales, the layer's global scale
- `gscale`) beside the mix, written by the mixer's own last launch: on the
- GEMV route `qwen4_exp_hc_decode_mix_fp4_cute`, otherwise the fused up
- projection `qwen4_exp_hc_up_gate_mean_fp4`. Returns
- `(hc, mixed, inject, xq, xsf)`."""
- core = self.core
- if _HC_DECODE_GEMV_MAX_ROWS < hc.shape[0] <= _HC_DECODE_ROWS_MAX:
- hc, low, inject, rstd = qwen4_exp_hc_decode_rows_first(
- hc,
- block_out,
- inject,
- norm,
- first,
- core.hc_count,
- core.hc_lowrank + core.hc_count,
- core.eps,
- )
- mixed, xq, xsf = qwen4_exp_hc_decode_rows_mix(
- low, rstd, up, hc, norm, core.hc_count, gscale
- )
- return hc, mixed, inject, xq, xsf
- if hc.shape[0] <= _HC_DECODE_GEMV_MAX_ROWS:
- hc, part, ss = qwen4_exp_hc_decode_first_cute(
- hc, block_out, inject, norm, first, core.hc_count, core.hc_lowrank + core.hc_count
- )
- mixed, inject, xq, xsf = qwen4_exp_hc_decode_mix_fp4_cute(
- part, ss, up, hc, norm, core.eps, core.hc_count, True, gscale
- )
- return hc, mixed, inject, xq, xsf
- hc, normed = qwen4_exp_hc_inject_norm(hc, block_out, inject, norm, core.eps, core.hc_count)
- low, inject = self._hc_first(normed, first, True)
- mixed, xq, xsf = qwen4_exp_hc_up_gate_mean_fp4(low, up, normed, core.hc_count, gscale)
- return hc, mixed, inject, xq, xsf
-
- def _grouped_norm(self, x: torch.Tensor, weights: tuple) -> torch.Tensor:
- """See `PrefillTarget._grouped_norm`."""
- core = self.core
- parts = torch.split(x, [core.hidden] * core.hc_count, 1)
- return torch.cat(
- [flashinfer_gemma_rmsnorm(part, w, core.eps) for w, part in zip(weights, parts)], 1
- )
-
- def _moe_block(
- self, x, xq, xsf, L: _LayerWeights, routed_out, gemm=cublas_mm, fold=False
- ) -> tuple:
- """One MoE layer: NVFP4 routed experts plus a bf16 shared expert, left
- uncombined. Returns `(routed, (gate_logits, shared, topk_weights,
- topk_rows))`: the next hyper-connection site forms `routed +
- sigmoid(gate_logits) * shared` inside its first kernel (`_hc_site`'s
- `moe`). With `fold` (trtllm-gen, and the next site takes the combine; a
- per-capture constant) the runner also skips its finalize: `routed` is then
- its permuted FC2 rows and `topk_weights` / `topk_rows` its routing weights
- and index map, which that site's first kernel
- (`qwen4_exp_hc_decode_first_cute`, `qwen4_exp_hc_decode_rows_first` or
- `qwen4_exp_hc_inject_norm`) gathers;
- without it both are None and `routed` is the finalized `[T, H]` rows.
-
- The shared expert and its per-token gate logit run on the core's side
- stream, concurrently with the routed path, and join it before this
- returns; from `_ROUTER_FIRST_MIN_ROWS` rows they launch after the router.
- `x` is the bf16 hidden states (router and shared expert); `xq` /
- `xsf` are the same states in NVFP4 with **linear** block scales, from
- `_hc_moe_site`. On trtllm-gen the routed path hands the runner the router
- logits (its Renormalize routing is bitwise the `renorm_moe_routing_op` +
- pre-routed composition); on CuteDSL it is `_cute_dsl_routed`, which writes
- `routed_out`. `gemm` is the step's bf16 projection entry (see `forward`); it
- launches on the current stream, so the shared expert's GEMMs run on the side
- stream."""
- core = self.core
- main = torch.cuda.current_stream()
- side = core._side_stream
- # Fork: the side stream reads `x` only after everything queued on `main`
- # so far. What it allocates is consumed on `main` after the join below,
- # and the next layer's fork orders any reuse of those blocks after that.
- side.wait_stream(main)
- # The routing reads its logits as a dense row-major buffer from data_ptr() and
- # ignores strides; every projection entry allocates them contiguous. From
- # `_ROUTER_FIRST_MIN_ROWS` rows the router launches before the shared expert (the
- # fork above precedes it, so the side stream does not wait for it).
- router_first = x.shape[0] >= _ROUTER_FIRST_MIN_ROWS
- if router_first:
- logits = _dense_mm(x, L.router, None, gemm, "router", mid=L.router_m)
- with torch.cuda.stream(side):
- gu = _dense_mm(x, L.sh_gu, None, gemm, "sh_gu", L.sh_gu_g, mid=L.sh_gu_m)
- shared = _dense_mm(
- flashinfer_silu_and_mul(gu), L.sh_dn, None, gemm, "sh_dn", L.sh_dn_g, mid=L.sh_dn_m
- )
- gate_logits = gemm(x, L.sh_gate)
- if not router_first:
- logits = _dense_mm(x, L.router, None, gemm, "router", mid=L.router_m)
- if core.trunk_moe_backend == "CUTEDSL":
- routed = _cute_dsl_routed(core, L, logits, xq, xsf, routed_out)
- main.wait_stream(side)
- return routed, (gate_logits, shared, None, None)
- if fold:
- moe_op = fp4_block_scale_moe_renorm_permuted
- else:
- moe_op = fp4_block_scale_moe_renorm_routed
- out = moe_op(
- logits,
- xq,
- xsf.view(torch.float8_e4m3fn),
- L.fc1_w,
- L.fc1_s,
- L.fc2_w,
- L.fc2_s,
- L.out1_scale,
- L.out1_gate_scale,
- L.out2_scale,
- core.num_experts,
- core.topk,
- core.moe_inter,
- )
- main.wait_stream(side)
- if fold:
- routed, topk_w, topk_rows = out
- return routed, (gate_logits, shared, topk_w, topk_rows)
- return out, (gate_logits, shared, None, None)
-
- def _attention_prologue(
- self, proj, q_index, index_view, index_cache, index_section, L, qsa, side, stage
- ) -> None:
- """The norm + rope of both planes on this stream and the QSA index branch (the
- verify step's staging, the raw-key store, then compression) on the side stream,
- which reads only the projection's index-key columns and the index pools; the
- attention core joins it before reading either."""
- core = self.core
- num_tokens = proj.shape[0]
- req_idx, logical, _, _, _, _, block_table, tokens_per_block, pos3, coords, _ = qsa
- side.wait_stream(torch.cuda.current_stream())
- qwen4_exp_qk_norm_rope(
- proj,
- q_index,
- pos3,
- core._cos_sin,
- L.q_norm,
- L.k_norm,
- L.idx_qn,
- num_heads_q=core.heads_q,
- num_heads_k=core.heads_kv,
- head_dim=core.head_dim,
- index_heads=core.index_heads,
- index_head_dim=core.index_head_dim,
- index_q_col=core.qkv_width + core.q_width,
- rotary_dim=core.rotary_dim,
- eps=core.eps,
- mrope_section1=core.mrope_section1,
- mrope_section2=core.mrope_section2,
- index_mrope_section1=index_section[1],
- index_mrope_section2=index_section[2],
- )
-
- # ---- the QSA index branch ------------------------------------------
- # The raw key is stored **before** anything normalizes or rotates it, because
- # the compression step pools raw keys and applies the norm and the rotation
- # itself at the group's first member's coordinates; a key normalized on the way
- # in double-normalizes every anchor that reads it, which is finite, silent, and
- # invisible below the sparse threshold. The norm + rope above never writes the
- # index key columns.
- with torch.cuda.stream(side):
- token_k = torch.split(
- torch.reshape(
- index_view,
- [
- num_tokens,
- core.index_scratch_width // core.index_head_dim,
- core.index_head_dim,
- ],
- ),
- [core.index_heads, core.index_kv_heads, 1],
- 1,
- )[1]
- if stage:
- core.spec_rollback.stage_index(index_cache, core._bound.position_cache)
- qsa_paged_index_store(
- token_k,
- coords,
- index_cache,
- core._bound.position_cache,
- req_idx,
- logical,
- block_table,
- tokens_per_block,
- )
- qsa_prefill_compress(
- logical_positions=logical,
- request_indices=req_idx,
- block_table=block_table,
- index_cache=index_cache,
- position_cache=core._bound.position_cache,
- k_norm_weight=L.idx_kn,
- cos_sin=core._cos_sin,
- eps=core.eps,
- tokens_per_block=tokens_per_block,
- compress_ratio=core.compress_ratio,
- mrope_section=_QSA_MROPE_SECTION,
- )
-
- def _attention_block(
- self, x, i, L, qkv, index_scratch, attn_out, qsa, step, stage, gemm=cublas_mm
- ):
- """One gated GQA layer plus its QSA index branch for a generation-only batch.
-
- The four input projections are one GEMM against the load-time fused
- operand (`weights._attn_in_w`), whose output row is
- `[q | k | v | gate | index q, k | v slot]`. One norm+rope launch rewrites
- the main q/k heads in place on that output and writes the normalized,
- rotated index query into the scratch; the sparse path reads q/k/v as views
- of it, the dense path takes one compact copy of its leading `[q | k | v]`
- (the dense op reads its QKV flat), the gate is read through its row stride,
- and the index branch reads its raw key from it.
-
- **The index plane is maintained on every step, at both ends of the
- threshold.** Below it the dense core runs and the index cache is still
- advanced, so a request that grows past the threshold finds its prefix
- already compressed; above it the dense core is replaced outright and the
- main K/V append becomes this layer's own call, because the op that would
- have done it is not the one running. Above the threshold over a bf16 or NVFP4 pool,
- everything up to the index scores is one launch, `qwen4_exp_attn_decode_prologue`
- (over an NVFP4 pool it quantizes the K/V append). On a verify step the rollback
- stages the index slots first, and at a group's last token the launch takes the
- raw keys of the group's earlier rows in the same step from their projection
- rows, as the index plane holds them only after the launch."""
- core = self.core
- (
- req_idx,
- logical,
- seq_lens_q,
- visible_blocks,
- topk_indices,
- topk_row_starts,
- block_table,
- tokens_per_block,
- pos3,
- coords,
- sparse,
- ) = qsa
- num_tokens = x.shape[0]
-
- main = torch.cuda.current_stream()
- side = core._side_stream
- index_cache = core._bound.index_cache[i]
- kv = core._bound.kv[i]
- # A bf16 or an NVFP4 pool (an fp8 one stays on the composed prologue).
- fused_prologue = sparse and (kv.k_scale is not None or kv.k.dtype == x.dtype)
- if fused_prologue and stage:
- # The rollback's copy of the slots the prologue launch overwrites, taken on
- # the side stream under the input projection.
- side.wait_stream(main)
- with torch.cuda.stream(side):
- core.spec_rollback.stage_index(index_cache, core._bound.position_cache)
- proj = _dense_mm(x, L.attn_in, L.attn_in8, gemm, "in", L.attn_in_g, mid=L.attn_in_m)
- qkv_view, gate, index_view = torch.split(
- proj, [core.qkv_width, core.q_width, core.index_scratch_width], 1
- )
- q_dst, k_dst, v_dst = _attn_in_heads(core, proj)
- # The index query lands in the q slots of the scratch; its k and v slots
- # are layout only and nothing reads them.
- q_index = torch.split(
- torch.reshape(
- index_scratch,
- [num_tokens, core.index_heads + 2 * core.index_kv_heads, core.index_head_dim],
- ),
- [core.index_heads, core.index_kv_heads, core.index_kv_heads],
- 1,
- )[0]
- index_section = (_QSA_MROPE_SECTION or (0, 0, 0)) if _QSA_INDEX_USE_MROPE else (0, 0, 0)
- # Up to SCORE_ROWS decode rows the index scores ride the prologue launch too.
- scores = None
- if fused_prologue and not stage and num_tokens <= _PROLOGUE_SCORE_ROWS:
- scores = torch.empty(
- (num_tokens, block_table.shape[1] * tokens_per_block // core.compress_ratio),
- dtype=torch.float32,
- device=proj.device,
- )
- if fused_prologue:
- if stage:
- main.wait_stream(side)
- # The norm + rope of both planes, the K/V append the dense op would have
- # owned (quantizing over an NVFP4 pool), and the index branch's raw-key and
- # coordinate store and group compression in one launch on this stream.
- qwen4_exp_attn_decode_prologue(
- proj,
- q_index,
- pos3,
- coords,
- core._cos_sin,
- L.q_norm,
- L.k_norm,
- L.idx_qn,
- L.idx_kn,
- kv.k,
- kv.v,
- index_cache,
- core._bound.position_cache,
- req_idx,
- logical,
- block_table,
- tokens_per_block,
- core.compress_ratio,
- num_heads_q=core.heads_q,
- num_heads_k=core.heads_kv,
- head_dim=core.head_dim,
- index_heads=core.index_heads,
- index_head_dim=core.index_head_dim,
- index_q_col=core.qkv_width + core.q_width,
- rotary_dim=core.rotary_dim,
- eps=core.eps,
- mrope_section1=core.mrope_section1,
- mrope_section2=core.mrope_section2,
- index_mrope_section1=index_section[1],
- index_mrope_section2=index_section[2],
- compress_mrope_section=_QSA_MROPE_SECTION,
- k_scale_cache=kv.k_scale,
- v_scale_cache=kv.v_scale,
- kv_scale_orig_quant=None if kv.k_scale is None else core.kv_scale_orig_quant,
- rows_per_request=step["predicted_tokens_per_seq"],
- index_scores=scores,
- )
- else:
- self._attention_prologue(
- proj, q_index, index_view, index_cache, index_section, L, qsa, side, stage
- )
- if self._kv_probe is not None and not self._kv_probe.observe(i, k_dst, v_dst):
- self._kv_probe = None
-
- gated = False
- if sparse:
- if not fused_prologue:
- # Selection reads the index plane the side stream wrote.
- main.wait_stream(side)
- # Append the fused-norm+rope'd K/V; the selection reads the current
- # token from the pool. Over an fp8 or NVFP4 pool as in the prefill target.
- _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block)
- if scores is None:
- scores = qsa_paged_index_scores(
- q=q_index,
- index_cache=index_cache,
- block_table=block_table,
- query_positions=logical,
- request_indices=req_idx,
- tokens_per_block=tokens_per_block,
- compress_ratio=core.compress_ratio,
- only_visible_blocks=True,
- context_rows=False,
- # Its predecessor is the prologue launch and its consumer the top-k,
- # which waits before reading the scores. With one row per request the
- # prologue writes no compressed key but the one the entry re-reads
- # after its wait, so at decode-sized batches the pages and keys are
- # read before it; a verify step's later rows also see the key its
- # group-end row compresses, so they read every key after the wait.
- use_pdl=fused_prologue,
- prefetch=fused_prologue and not stage and num_tokens <= _SPARSE_DECODE_MAX_ROWS,
- )
- if _topk_expand_supports(scores):
- # The canonical block top-k and its expansion to tokens in one launch
- # (QSA's row starts are zero).
- selected = qsa_topk_expand_decode(
- scores,
- visible_blocks,
- logical,
- seq_lens_q,
- topk_indices,
- core.compress_ratio,
- core.token_topk,
- )
- else:
- _select(core, scores, topk_row_starts, visible_blocks, topk_indices)
- selected = qsa_expand_block_indices(
- topk_indices,
- logical,
- seq_lens_q,
- compress_ratio=core.compress_ratio,
- token_topk=core.token_topk,
- )
- if kv.k_scale is not None:
- # The output gate rides the CuTe kernel's store (one split) or merge. A verify
- # step's rows come `predicted_tokens_per_seq` a request.
- attn = _sparse_gqa(
- core,
- kv,
- q=q_dst,
- block_table=block_table,
- selected_tokens=selected,
- request_indices=req_idx,
- tokens_per_block=tokens_per_block,
- # The prefill pair bounds a context row's scan; none here.
- query_positions=None,
- compress_ratio=None,
- output_gate=gate,
- rows_per_request=step["predicted_tokens_per_seq"],
- )
- gated = True
- else:
- # The output gate rides the split-K merge whenever the launch has
- # one (more than one split); otherwise it is its own kernel below.
- gated = qsa_paged_sparse_gqa.fuses_output_gate(q_dst, kv.k, selected, gate)
- if (
- gated
- and num_tokens <= _SPARSE_DECODE_MAX_ROWS
- and _sparse_decode_supports(q_dst, kv.k, selected)
- ):
- # The split-K pass, then the merge and the gate as one launch spread
- # over the SMs.
- attn = qsa_paged_sparse_gqa_decode(
- q_dst,
- kv.k,
- kv.v,
- block_table,
- selected,
- req_idx,
- tokens_per_block,
- core.softmax_scale,
- gate,
- )
- else:
- attn = qsa_paged_sparse_gqa(
- q=q_dst,
- k_cache=kv.k,
- v_cache=kv.v,
- block_table=block_table,
- selected_tokens=selected,
- request_indices=req_idx,
- tokens_per_block=tokens_per_block,
- softmax_scale=core.softmax_scale,
- query_positions=None,
- compress_ratio=None,
- output_gate=gate if gated else None,
- )
- out = torch.reshape(attn, [num_tokens, core.q_width])
- else:
- main.wait_stream(side)
- # The dense op reads its packed QKV flat: one compact copy.
- qkv.copy_(qkv_view)
- out = thop_attention(
- q=qkv,
- output=attn_out,
- local_layer_idx=core._bound.pool_idx[i],
- num_heads=core.heads_q,
- num_kv_heads=core.heads_kv,
- head_size=core.head_dim,
- quant_mode=core.quant_mode,
- kv_scale_orig_quant=core.kv_scale_orig_quant,
- kv_scale_quant_orig=core.kv_scale_quant_orig,
- **step,
- **_CALL_CONSTANTS,
- )
- if not gated:
- fused_sigmoid_mul_inplace(out, gate)
- return _dense_mm(out, L.o, L.o8, gemm, "out", L.o_g, mid=L.o_m)
-
- def _gdn_verify(
- self, mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, n, p, replay, mq_out=None
- ) -> torch.Tensor:
- """Recurrent half of a speculative verify step; see
- `PrefillTarget._gdn_verify`. `mq_out` is the conv output when the in_proj ran the
- conv step as its epilogue (`fp8_gemv_conv1d_verify`, then `mq_d` is None); else the
- conv runs here over the projection's q|k|v columns `mq_d`."""
- core = self.core
- # The in_proj GEMV ran the conv as its epilogue and signals its dependents early.
- after_gemv = mq_out is not None
- if mq_out is None:
- # [n*p, C] -> [n, C, p]: the multi-token conv reads channel-major.
- x3 = torch.transpose(torch.reshape(mq_d, [n, p, core.conv_dim]), 1, 2)
- conv_out = causal_conv1d_triton_update(
- x=x3,
- conv_state=core._bound.conv[i],
- weight=L.conv_w,
- bias=None,
- activation="silu",
- conv_state_indices=si_d,
- pad_slot_id=_PAD_SLOT_ID,
- # Snapshots, not the accepted-token roll.
- num_accepted_tokens=None,
- intermediate_conv_window=core._bound.int_conv[i],
- intermediate_state_indices=int_idx,
- launch_dependent_kernels=False,
- )
- # Back to token-major (a view).
- mq_out = torch.reshape(torch.transpose(conv_out, 1, 2), [n * p, core.conv_dim])
- if replay is not None:
- return _replay_verify(
- core, mq_out, a_d, b_d, ssm_states, i, L, si_d, n, p, replay, pdl=after_gemv
- )
- q, k, v = pack_gdn_decode_qkv(
- mq_out, core.num_k_heads, core.head_k_dim, core.num_v_heads, core.head_v_dim
- )
- g = torch.reshape(
- fused_gdn_gating(
- A_log=L.a_log,
- a=a_d,
- dt_bias=L.dt_bias,
- beta=_SOFTPLUS_BETA,
- threshold=_SOFTPLUS_THRESHOLD,
- ),
- [n, p, core.num_v_heads],
- )
- beta = torch.reshape(torch.sigmoid(b_d), [n, p, core.num_v_heads])
- # Initial states are gathered from the live pool; the per-step buffer is
- # indexed by batch row (`intermediate_state_indices[:n]`), not slot id.
- src = torch.reshape(
- nn.functional.embedding(si_d, torch.reshape(ssm_states, [ssm_states.shape[0], -1])),
- [n, core.num_v_heads, core.head_v_dim, core.head_k_dim],
- )
- out = fused_recurrent_gated_delta_rule_update(
- torch.reshape(q, [n, p, core.num_k_heads, core.head_k_dim]),
- torch.reshape(k, [n, p, core.num_k_heads, core.head_k_dim]),
- torch.reshape(v, [n, p, core.num_v_heads, core.head_v_dim]),
- g,
- beta,
- initial_state_source=src,
- initial_state_indices=int_idx,
- scale=None,
- use_qk_l2norm_in_kernel=True,
- cu_seqlens=None,
- disable_state_update=True,
- disable_output_calculation=False,
- intermediate_states_buffer=core._bound.int_ssm[i],
- cache_steps=p,
- output=None,
- )
- return torch.reshape(out, [1, n * p, core.num_v_heads, core.head_v_dim])
-
- def _gdn_block(self, x, i, L, gdn_state, gemm=cublas_mm, fused_conv=False) -> torch.Tensor:
- """One Gated-DeltaNet layer over a generation-only batch: one token per
- sequence, or under speculation the whole draft chain through the verify
- path. Every branch is on a host int the engine fixes for a capture.
-
- One GEMM for all four input projections: q|k|v, z, a and b are column views of
- its output. A graph-captured single-token step at up to `_GDN_CONV_FUSED_MAX_ROWS`
- rows runs the conv step as that GEMM's epilogue (`gdn_in_proj_conv1d_pdl`), and
- over an fp32 recurrent pool at up to `_GDN_DECODE_FUSED_MAX_ROWS` rows the
- recurrence and the gated norm as one launch (`gdn_decode_gated_norm_pdl`); past
- them, and at every batch size over a bf16 pool, after the conv step, as one
- streaming launch (`gdn_decode_rows_gated_norm`). That kernel computes each (row,
- v-head) on its own, so over a bf16 pool a row's output and state do not depend on
- its batch.
-
- A graph-captured verify step whose in_proj takes `fp8_gemv` (dense FP8 mode, up to
- `_VERIFY_CONV_FUSED_MAX_ROWS` rows) runs the multi-token conv step -- window snapshots and
- the pool advance included -- as that GEMM's epilogue too (`fp8_gemv_conv1d_verify`,
- bitwise the projection + `causal_conv1d_triton_update` it replaces); the conv output
- is its own contiguous buffer, and z, a and b are column views of the rest.
-
- With the in_proj's NVFP4 copy (`_GDN_INPROJ_NVFP4_ENV`) every step of up to
- `_INPROJ_NVFP4_MAX_ROWS` rows reads it on the same kernel instead, a captured
- single-token step at all of those rows with the conv step as the epilogue."""
- core = self.core
- si_all, is_verify, num_decodes, gen_p, int_idx, replay = gdn_state
- num_tokens = x.shape[0]
- ssm_states = core._bound.ssm[i]
-
- conv_fused = fused_conv and not is_verify and num_tokens <= _GDN_CONV_FUSED_MAX_ROWS
- # The NVFP4 copy takes every decode row count within its range, the conv step fused
- # at each of them in a captured step.
- nv = (
- L.in_proj_nv
- if L.in_proj_nv is not None and num_tokens <= _INPROJ_NVFP4_MAX_ROWS and _gemv_ready(x)
- else None
- )
- nv_conv = nv is not None and fused_conv and not is_verify
- fp8_conv = (
- not nv_conv
- and fused_conv
- and not is_verify
- and L.in_proj_g is not None
- and _gemv_ready(x)
- and num_tokens <= _FP8_CONV_FUSED_MAX_ROWS
- )
- # The verify conv as the in_proj's epilogue reads the pool, taps and slots before its
- # PDL wait, like the single-token fusion: captured steps only.
- verify_fused = (
- fused_conv
- and is_verify
- and (L.in_proj_g is not None or nv is not None)
- and _gemv_ready(x)
- and num_tokens <= min(_FP8_GEMV_ROWS["in"][1], _VERIFY_CONV_FUSED_MAX_ROWS)
- )
- mq_out = None
- if verify_fused:
- verify_op = (
- fp8_gemv_conv1d_verify_op if nv is None else nvfp4_split_gemv_conv1d_verify_op
- )
- mq_out, zab = verify_op(
- x,
- L.in_proj_g if nv is None else nv,
- core._bound.conv[i],
- L.conv_w,
- si_all,
- core._bound.int_conv[i],
- int_idx,
- core.conv_dim,
- gen_p,
- _PAD_SLOT_ID,
- )
- elif nv_conv:
- proj = nvfp4_split_gemv_conv1d_op(
- x, nv, core._bound.conv[i], L.conv_w, si_all, core.conv_dim, _PAD_SLOT_ID
- )
- elif fp8_conv:
- proj = fp8_gemv_conv1d_op(
- x,
- L.in_proj_g,
- core._bound.conv[i],
- L.conv_w,
- si_all,
- core.conv_dim,
- _PAD_SLOT_ID,
- )
- elif conv_fused:
- proj = gdn_in_proj_conv1d_pdl(
- x,
- L.in_proj,
- core._bound.conv[i],
- L.conv_w,
- si_all,
- core.conv_dim,
- _PAD_SLOT_ID,
- True,
- )
- elif nv is not None:
- proj = nvfp4_split_gemv_op(x, nv)
- else:
- proj = _dense_mm(x, L.in_proj, L.in_proj8, gemm, "in", L.in_proj_g, mid=L.in_proj_m)
- if verify_fused:
- mixed_qkv = None
- z, a, b = torch.split(zab, core.in_proj_split[1:], 1)
- else:
- mixed_qkv, z, a, b = torch.split(proj, core.in_proj_split, 1)
-
- if (
- ssm_states.dtype == torch.float32
- and conv_fused
- and num_tokens <= _GDN_DECODE_FUSED_MAX_ROWS
- ):
- # The recurrence and the gated norm as one launch that reads this layer's
- # pool rows before its PDL wait: in a captured single-token step nothing
- # ahead of it writes them, and the slot indices are filled before the
- # graph is launched.
- o = gdn_decode_gated_norm_pdl(
- proj,
- ssm_states,
- si_all,
- L.a_log,
- L.dt_bias,
- L.gnorm,
- core.eps,
- core.num_k_heads,
- core.head_k_dim,
- )
- return _dense_mm(o, L.out, L.out8, gemm, "out", L.out_g, mid=L.out_m)
- if is_verify:
- # The CuTe replay kernel reads `a`/`b` through their row stride: no copies.
- strided_ab = replay is not None and core._bound.replay.cute
- out = self._gdn_verify(
- mixed_qkv,
- a if strided_ab else a.contiguous(),
- b if strided_ab else b.contiguous(),
- ssm_states,
- i,
- L,
- si_all,
- int_idx,
- num_decodes,
- gen_p,
- replay,
- mq_out,
- )
- return self._gdn_epilogue(out, z, L, num_tokens, gemm)
- if not (conv_fused or fp8_conv or nv_conv):
- causal_conv1d_update(
- x=mixed_qkv,
- conv_state=core._bound.conv[i],
- weight=L.conv_w,
- bias=None,
- silu_activation=True,
- cache_seqlens=None,
- conv_state_indices=si_all,
- pad_slot_id=_PAD_SLOT_ID,
- )
- # The recurrence and the gated norm as one launch over `proj`, whose q/k/v
- # columns the conv step rewrote in place; it computes both gates and the q/k L2
- # norm itself, at the softplus pair (1, 20) and epsilon 1e-6 this target's
- # constants name, and reads this layer's pool rows before its PDL wait (nothing
- # in a single-token step writes them ahead of it).
- o = gdn_decode_rows_gated_norm(
- proj,
- ssm_states,
- si_all,
- L.a_log,
- L.dt_bias,
- L.gnorm,
- core.eps,
- core.num_k_heads,
- core.head_k_dim,
- )
- return _dense_mm(o, L.out, L.out8, gemm, "out", L.out_g, mid=L.out_m)
-
- def _gdn_epilogue(self, out, z, L, num_tokens, gemm=cublas_mm) -> torch.Tensor:
- """Gated RMS norm (sigmoid gate, plain non-gemma weight) and output projection
- shared by every recurrence."""
- core = self.core
- o = rms_norm_gated_token_major(
- torch.reshape(out, [-1, core.head_v_dim]),
- torch.reshape(z, [num_tokens, core.num_v_heads, core.head_v_dim]),
- weight=L.gnorm,
- eps=core.eps,
- fp8_scale=None,
- gate_activation="sigmoid",
- )
- o = torch.reshape(o, [num_tokens, core.value_dim])
- return _dense_mm(o, L.out, L.out8, gemm, "out", L.out_g, mid=L.out_m)
-
- def _ple_block(self, hc, input_ids, ple_state, next_norm, stage, gemm=cublas_mm) -> tuple:
- """The n-gram side path on one decoder layer, added into the stream and
- normalized for the next site. Returns `(hc, normed)`.
-
- The same steps as the prefill target's `_ple_block` -- id history, hash, table
- gather and dequant, the gated injection, the dilated short conv -- in three
- launches plus the two projections, because a decode batch carries at most
- `gen_p` rows per sequence:
-
- 1. `ple_ngram_embed`: window, hash, gather, dequantize, and the id
- history published back, one program per sequence;
- 2. the key and value projections, on the step's bf16 projection entry
- `gemm` (see `forward`);
- 3. `ple_gate_conv_inject_norm`: the key / query grouped norms, the gate,
- the conv-input grouped norm, the conv against the paged conv state
- (published back), the residual add into the stream and this
- layer's attention-side grouped norm, one program per
- (sequence, stream).
-
- Every generation sequence carries its history by construction, so no carry
- flags are passed. A verify step (`stage`) stages what `SpeculativeRollback`
- rewrites the two pools from after sampling: both pools' pre-step rows before
- either launch publishes, and the conv input the third launch writes out."""
- core = self.core
- S = self._stack
- conv_states, ngram_pool, si_all, carry_all, qsl, gen_p = ple_state
- if stage:
- core.spec_rollback.stage_ple_verify(ngram_pool, conv_states, input_ids)
- ngram = ple_ngram_embed(
- input_ids,
- ngram_pool,
- si_all,
- None,
- qsl,
- core.w["ple_table"],
- S.ple_tscale,
- S.ple_mult,
- S.ple_vocab,
- S.ple_offs,
- core.eos_token_id,
- gen_p,
- )
- conv_input = torch.empty_like(hc) if stage else None
- out = ple_gate_conv_inject_norm(
- hc,
- gemm(ngram, S.ple_key),
- gemm(ngram, S.ple_val),
- S.ple_nk_w,
- S.ple_nq_w,
- S.ple_nc_w,
- S.ple_conv_w,
- conv_states,
- si_all,
- qsl,
- next_norm,
- core.eps,
- core.hc_count,
- _PLE_CLAMP_MIN,
- core.ple_dilation,
- _PAD_SLOT_ID,
- gen_p,
- conv_input=conv_input,
- )
- if stage:
- core.spec_rollback.stage_ple_conv_input(conv_input)
- return out
-
- # -------------------------------------------------------------- forward ---
-
- def forward(
- self,
- attn_metadata: AttentionMetadata,
- input_ids: torch.IntTensor | None = None,
- position_ids: torch.IntTensor | None = None,
- **kwargs,
- ):
- core = self.core
- md = attn_metadata
- self._check_step_contract(md)
- step = self.step_args(md)
- spec_metadata = kwargs.get("spec_metadata")
- mm = md.mamba_metadata
- # The bf16 block projections of a graph-captured step run as PDL launches
- # (`_decode_projection`); an eager step keeps cuBLAS.
- gemm = _decode_projection if md.is_cuda_graph else cublas_mm
-
- # A captured verify step builds `pos3`, the PLE segmentation vector and the
- # rollback's staging in one launch (`verify_step_setup`, below).
- rollback = core.spec_rollback
- fused_setup = (
- md.is_cuda_graph
- and spec_metadata is not None
- and 0 < md.num_seqs < _VERIFY_SETUP_MAX_SEQS
- and position_ids.shape[-1] > md.num_seqs
- and position_ids.stride(-1) == 1
- and rollback is not None
- and rollback.index_planes is not None
- )
- # Two views of one `[3, 1, T]` buffer (see PrefillTarget.forward): contiguous
- # `pos3` and `coords`.
- coords = torch.transpose(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1]), 0, 1)
- pos3 = None if fused_setup else _copy_positions(position_ids)
-
- h = nn.functional.embedding(input_ids, core.w["embed"])
- num_tokens = h.shape[0]
- dt = h.dtype
- dev = h.device
-
- # The batch's bookkeeping, built once; every sequence is a generation sequence.
- batch = md.num_seqs
- # Host-int check (graph-safe); see PrefillTarget.prepare.
- assert num_tokens == md.num_tokens, (
- f"the engine handed {num_tokens} activation rows for a batch its "
- f"metadata describes as {md.num_tokens} tokens "
- f"(padded_num_tokens={getattr(md, 'padded_num_tokens', None)}). A "
- "prefill/piecewise CUDA-graph backend pads the activation buffer "
- "past the batch and expects the model to slice it back; this "
- "target does not implement that configuration. Decode CUDA graphs "
- "pad with real dummy requests instead, so they are in the metadata "
- "and this assert is silent for them."
- )
- # Rows per generation sequence: 1 without speculation, else
- # `runtime_draft_len + 1`.
- gen_p = num_tokens // batch if batch else 1
- assert gen_p * batch == num_tokens, (
- f"{num_tokens} generation rows do not divide over {batch} generation "
- "sequences; a ragged per-sequence draft length can be expressed "
- "neither to the attention call nor to the multi-token recurrence"
- )
- assert core.mtp_enabled or gen_p == 1, (
- f"{num_tokens} generation rows over {batch} generation sequences with "
- "speculation off; the base path implements one query token per "
- "generation sequence"
- )
- # `gen_p > 1` is load-bearing for verify (see PrefillTarget.prepare).
- if spec_metadata is not None and batch > 0:
- declared_p = spec_metadata.runtime_draft_len + 1
- assert gen_p == declared_p, (
- f"the batch carries {gen_p} query rows per generation sequence "
- f"but the worker's runtime_draft_len says {declared_p}; the "
- "per-step conv snapshots and the recurrent rollback buffer are "
- "sized from the second and written from the first"
- )
- check_verify_width(gen_p)
- check_replay_width(core._bound.replay, gen_p)
- is_verify = spec_metadata is not None and batch > 0 and gen_p > 1
- step["predicted_tokens_per_seq"] = gen_p
-
- si_all = mm.state_indices[:batch]
- qsl_long = mm.query_start_loc_long
- int_idx = None
- replay = None
- if is_verify:
- # The runtime's arange over the slot pool (what `update_mamba_states` reads).
- assert core._bound.int_idx is not None, "the cache-manager bind has not run"
- int_idx = core._bound.int_idx[:batch]
- if core._bound.replay is not None:
- replay = _replay_work(core, mm, batch)
- gdn_state = (si_all, is_verify, batch, gen_p, int_idx, replay)
-
- # PLE segmentation vector: rebuilt from `seq_lens_cuda` when `gen_p > 1` (see
- # PrefillTarget.prepare). Every generation sequence carries its n-gram history, so
- # the PLE block takes no carry flags.
- fused_setup = fused_setup and is_verify
- if pos3 is None and not fused_setup:
- pos3 = _copy_positions(position_ids)
- if fused_setup:
- planes = rollback.index_planes
- (pos3, ple_qsl, staged_slots, staged_pages, staged_within, staged_index) = (
- verify_step_setup(
- position_ids,
- md.seq_lens_cuda[:batch],
- si_all,
- md.qsa_req_idx_per_token[:num_tokens],
- md.qsa_logical_positions[:num_tokens],
- md.qsa_block_table,
- md.tokens_per_block,
- *planes,
- )
- )
- elif gen_p == 1:
- ple_qsl = qsl_long
- else:
- ple_qsl = cu_seqlens_from_seq_lens(md.seq_lens_cuda[:batch], torch.int64)
- ple_state = (
- core._bound.ple_state[0],
- core._bound.ple_state[1],
- si_all,
- None,
- ple_qsl,
- gen_p,
- )
- self._trunk_pos3 = (_position_key(position_ids), pos3)
-
- # Dense/sparse decision once per forward, from a host-side metadata read: a
- # per-capture constant matching the engine's short/long decode graph families.
- sparse = batch > 0 and int(md.kv_lens_runtime[:batch].max()) > core.dense_threshold
-
- qsa = (
- md.qsa_req_idx_per_token[:num_tokens],
- md.qsa_logical_positions[:num_tokens],
- md.qsa_sequence_lengths[:num_tokens],
- md.qsa_visible_blocks[:num_tokens],
- md.qsa_topk_indices[:num_tokens],
- md.qsa_topk_row_starts[:num_tokens],
- md.qsa_block_table,
- md.tokens_per_block,
- pos3,
- coords,
- sparse,
- )
-
- # Staging for the speculative rollback, as on the prefill target (see
- # PrefillTarget.prepare); every row here is a generation row.
- stage = is_verify
- if rollback is not None:
- rollback.discard()
- if fused_setup:
- rollback.begin_staged(
- num_gens=batch,
- rows_per_gen=gen_p,
- slots=staged_slots,
- pages=staged_pages,
- within=staged_within,
- index_all=staged_index,
- )
- elif stage:
- rollback.begin(
- num_contexts=0,
- prefix_tokens=0,
- num_gens=batch,
- rows_per_gen=gen_p,
- slots=si_all,
- request_indices=md.qsa_req_idx_per_token[:num_tokens],
- logical_positions=md.qsa_logical_positions[:num_tokens],
- block_table=md.qsa_block_table,
- tokens_per_block=md.tokens_per_block,
- fused=md.is_cuda_graph,
- )
-
- # Scratch shared by all attention layers: thop_attention overwrites rows
- # [:num_tokens] of `output` in place, and the packed buffers are fully
- # rewritten every layer (the QKV one on the dense path only).
- qkv = torch.empty([num_tokens, core.qkv_width], dtype=dt, device=dev)
- index_scratch = torch.empty([num_tokens, core.index_scratch_width], dtype=dt, device=dev)
- attn_out = torch.empty([num_tokens, core.q_width], dtype=dt, device=dev)
- # The CuteDSL routed output, shared by all layers: each layer's unpermute rewrites
- # every row, and the next site's injection consumes it before the next layer's MoE.
- routed_out = (
- torch.empty([num_tokens, core.hidden], dtype=dt, device=dev)
- if core.trunk_moe_backend == "CUTEDSL"
- else None
- )
-
- # The embedding repeated across the hc_count residual streams.
- hc_count = core.hc_count
- hc = torch.cat([h] * hc_count, -1)
- layers, S = self._layers, self._stack
- # Every hyper-connection site (`_hc_site`): the previous block's pending
- # injection, this site's grouped norm and its mixer. The stack's first
- # site has no pending injection; at the PLE layer it is applied alone
- # (the PLE block reads the injected stream), and the PLE tail returns the
- # stream with its output added plus this site's grouped norm of it, which
- # the site's cuBLAS mixer starts from (a GEMV site normalizes inside its
- # first launch and reads the stream alone). A MoE block's output is
- # pending uncombined (`moe`): the next site's first kernel forms it,
- # except before the PLE add, where the combine runs alone.
- block_out = inject = moe = None
- for i, L in enumerate(layers):
- if i == core.ple_layer:
- if block_out is not None:
- if moe is not None:
- # `_moe_block` never folds the finalize before this layer.
- block_out = fused_sigmoid_gate_mul_add(block_out, *moe[:2])
- hc = qwen4_exp_post_inject(hc, block_out, inject)
- block_out = inject = moe = None
- hc, normed = self._ple_block(hc, input_ids, ple_state, L.ahc_norm, stage, gemm)
- else:
- normed = None
- hc, mixed, inject = self._hc_site(
- hc,
- block_out,
- inject,
- L.ahc_norm,
- L.ahc_first,
- L.ahc_up,
- True,
- normed=normed,
- moe=moe,
- )
- if core.is_attn[i]:
- block_out = self._attention_block(
- mixed, i, L, qkv, index_scratch, attn_out, qsa, step, stage, gemm
- )
- else:
- block_out = self._gdn_block(mixed, i, L, gdn_state, gemm, md.is_cuda_graph)
- hc, mixed, inject, xq, xsf = self._hc_moe_site(
- hc, block_out, inject, L.mhc_norm, L.mhc_first, L.mhc_up, L.moe_gscale
- )
- # The trtllm-gen runner's finalize folds into the next site's first kernel
- # when that site takes the combine (every one but the PLE layer's).
- fold = core.trunk_moe_backend == "TRTLLM" and i + 1 != core.ple_layer
- block_out, moe = self._moe_block(mixed, xq, xsf, L, routed_out, gemm, fold=fold)
- # The stack-final mixer with no injection weight: it returns the mix
- # alone, and that mix *is* this checkpoint's final norm.
- hc, final, _ = self._hc_site(
- hc, block_out, inject, S.fhc_norm, S.fhc_first, S.fhc_up, False, moe=moe
- )
- if spec_metadata is None:
- return final
- # `lm_head` reads the `[T, hidden]` mix; the MTP module reads the
- # `[T, hc_count * hidden]` stream it was collapsed from.
- return final, hc
-
- # ------------------------------------------------------------- the draft ---
-
- def _mtp_moe_block(self, x, dt) -> torch.Tensor:
- """FP8 block-scale (or, with `_MTP_MOE_NVFP4_ENV`, NVFP4) draft MoE; see
- `PrefillTarget._mtp_moe_block`."""
- core = self.core
- M = core._mtp
- if M.moe_gscale is not None:
- return _mtp_moe_nvfp4_block(core, M, x)
- logits = cublas_mm(x, M.router)
- topk_ids, topk_w = renorm_moe_routing_op(logits, topk=core.topk, output_dtype=dt)
- xq, xsf = fp8_quantize_1x128(x, use_ue8m0=_MTP_MOE_USE_UE8M0)
- routed = fp8_block_scale_moe_runner(
- routing_logits=None,
- routing_bias=None,
- hidden_states=xq,
- hidden_states_scale=xsf,
- gemm1_weights=M.fc1_w,
- gemm1_weights_scale=M.fc1_s,
- gemm2_weights=M.fc2_w,
- gemm2_weights_scale=M.fc2_s,
- num_experts=core.num_experts,
- top_k=core.topk,
- num_fused_shared_experts=None,
- n_group=None,
- topk_group=None,
- intermediate_size=core.moe_inter,
- local_expert_offset=0,
- local_num_experts=core.num_experts,
- routed_scaling_factor=None,
- routing_method_type=_ROUTING_METHOD_RENORM,
- topk_weights=topk_w,
- topk_ids=topk_ids,
- act_type=_ACT_TYPE_SWIGLU,
- gemm1_clamp_limit=None,
- )
- shared = cublas_mm(flashinfer_silu_and_mul(cublas_mm(x, M.sh_gu)), M.sh_dn)
- gate_logits = cublas_mm(x, M.sh_gate)
- return fused_sigmoid_gate_mul_add(routed, gate_logits, shared)
-
- def _mtp_layer(
- self,
- embed_tokens,
- input_ids,
- hidden_states,
- md,
- step,
- draft,
- layer_idx,
- rows,
- gen_seqs,
- p,
- sparse,
- share,
- reuse,
- pos3,
- pos,
- coords,
- ) -> torch.Tensor:
- """The draft layer on the trunk's decode machinery: the front end's two projections,
- both hyper-connection sites as `_hc_site` (the front end's per-stream `x_h + x_e` is
- the first site's injection at unit weight, the attention output the second's), the
- attention as `_mtp_attention`, the FP8 MoE as `_mtp_moe_block`. Returns the stream
- with the MoE output injected, as the composed step does. The hyper-connection
- mixers compute the composed `qwen4_exp_gated_residual`'s formulas with fp32 logits
- and an fp32 bottleneck (the trunk sites' numerics), so the draft's tokens may differ
- from the composed step's in the last place; the verify step decides acceptance."""
- core = self.core
- M, F = core._mtp, core._mtp_fused
- gemm = _decode_projection if md.is_cuda_graph else cublas_mm
- e = nn.functional.embedding(input_ids, embed_tokens)
- x_e = gemm(flashinfer_gemma_rmsnorm(e, M.pre_e_norm, core.eps), M.fc_e)
- h_n = self._grouped_norm(hidden_states, M.pre_h_norm)
- x_h = torch.reshape(
- gemm(torch.reshape(h_n, [rows * core.hc_count, core.hidden]), M.fc_h),
- [rows, core.stream],
- )
- x, mixed, inject = self._hc_site(
- x_h, x_e, F.unit_inject[:rows], F.ahc_norm, F.ahc_first, F.ahc_up, True
- )
- attn = self._mtp_attention(
- mixed,
- md,
- step,
- draft,
- layer_idx,
- rows,
- gen_seqs,
- p,
- sparse,
- share,
- reuse,
- pos3,
- pos,
- coords,
- gemm,
- )
- x, mixed, inject = self._hc_site(x, attn, inject, F.mhc_norm, F.mhc_first, F.mhc_up, True)
- return qwen4_exp_post_inject(x, self._mtp_moe_block(mixed, x.dtype), inject)
-
- def _mtp_attention(
- self,
- x,
- md,
- step,
- draft,
- layer_idx,
- rows,
- gen_seqs,
- p,
- sparse,
- share,
- reuse,
- pos3,
- pos,
- coords,
- gemm,
- ) -> torch.Tensor:
- """The draft layer's gated GQA and QSA index branch, as `_attention_block` runs a
- trunk layer over the draft manager's planes (`draft`): one input projection
- against `_MtpFused.attn_in`, the composed prologue (norm + rope of both planes at
- the module's rope base, the index store and compression on the side stream), the
- K/V append, the index scores and the one-launch top-k + expansion, sparse GQA, the
- output gate and projection. Under IndexShare a draft decode step (`reuse`) skips the
- index branch and rebuilds its selection from the draft-extend step's capture; the
- draft-extend step (`share`) captures it."""
- core = self.core
- M, F = core._mtp, core._mtp_fused
- dt, dev = x.dtype, x.device
- proj = gemm(x, F.attn_in)
- qkv_view, gate, index_view = torch.split(
- proj, [core.qkv_width, core.q_width, core.index_scratch_width], 1
- )
- q_dst, k_dst, v_dst = _attn_in_heads(core, proj)
- index_scratch = torch.empty([rows, core.index_scratch_width], dtype=dt, device=dev)
- q_index = torch.split(
- torch.reshape(
- index_scratch,
- [rows, core.index_heads + 2 * core.index_kv_heads, core.index_head_dim],
- ),
- [core.index_heads, core.index_kv_heads, core.index_kv_heads],
- 1,
- )[0]
- index_section = (_QSA_MROPE_SECTION or (0, 0, 0)) if _QSA_INDEX_USE_MROPE else (0, 0, 0)
- req_idx = md.qsa_req_idx_per_token[:rows]
- logical = md.qsa_logical_positions[:rows]
- block_table = md.qsa_block_table
- tokens_per_block = md.tokens_per_block
- main = torch.cuda.current_stream()
- side = core._side_stream
- side.wait_stream(main)
- qwen4_exp_qk_norm_rope(
- proj,
- q_index,
- pos3,
- core._mtp_cos_sin,
- M.q_norm,
- M.k_norm,
- M.idx_qn,
- num_heads_q=core.heads_q,
- num_heads_k=core.heads_kv,
- head_dim=core.head_dim,
- index_heads=core.index_heads,
- index_head_dim=core.index_head_dim,
- index_q_col=core.qkv_width + core.q_width,
- rotary_dim=core.rotary_dim,
- eps=core.eps,
- mrope_section1=core.mrope_section1,
- mrope_section2=core.mrope_section2,
- index_mrope_section1=index_section[1],
- index_mrope_section2=index_section[2],
- )
- if not reuse:
- # Raw key stored before anything normalizes it (see `_attention_prologue`).
- with torch.cuda.stream(side):
- token_k = torch.split(
- torch.reshape(
- index_view,
- [
- rows,
- core.index_scratch_width // core.index_head_dim,
- core.index_head_dim,
- ],
- ),
- [core.index_heads, core.index_kv_heads, 1],
- 1,
- )[1]
- qsa_paged_index_store(
- token_k,
- coords,
- draft.index_cache,
- draft.position_cache,
- req_idx,
- logical,
- block_table,
- tokens_per_block,
- )
- qsa_prefill_compress(
- logical_positions=logical,
- request_indices=req_idx,
- block_table=block_table,
- index_cache=draft.index_cache,
- position_cache=draft.position_cache,
- k_norm_weight=M.idx_kn,
- cos_sin=core._mtp_cos_sin,
- eps=core.eps,
- tokens_per_block=tokens_per_block,
- compress_ratio=core.compress_ratio,
- mrope_section=_QSA_MROPE_SECTION,
- )
- main.wait_stream(side)
- gated = False
- if sparse:
- kv = draft.kv
- _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block)
- topk_indices = md.qsa_topk_indices[:rows]
- visible_blocks = md.qsa_visible_blocks[:rows]
- seq_lens_q = md.qsa_sequence_lengths[:rows]
- if reuse:
- slot = draft.share_slot
- assert slot in md.qsa_shared_topk_captured_slots, (
- "a draft decode step would reuse a selection that no draft-extend step "
- "of this loop captured"
- )
- qsa_mtp_reuse_topk(
- shared_indices=md.qsa_shared_topk_indices[slot],
- shared_visible_blocks=md.qsa_shared_topk_visible_blocks[slot],
- request_indices=req_idx,
- visible_blocks=visible_blocks,
- output=topk_indices,
- )
- selected = qsa_expand_block_indices(
- topk_indices,
- logical,
- seq_lens_q,
- compress_ratio=core.compress_ratio,
- token_topk=core.token_topk,
- )
- else:
- scores = qsa_paged_index_scores(
- q=q_index,
- index_cache=draft.index_cache,
- block_table=block_table,
- query_positions=logical,
- request_indices=req_idx,
- tokens_per_block=tokens_per_block,
- compress_ratio=core.compress_ratio,
- only_visible_blocks=True,
- context_rows=False,
- )
- if _topk_expand_supports(scores):
- selected = qsa_topk_expand_decode(
- scores,
- visible_blocks,
- logical,
- seq_lens_q,
- topk_indices,
- core.compress_ratio,
- core.token_topk,
- )
- else:
- _select(
- core, scores, md.qsa_topk_row_starts[:rows], visible_blocks, topk_indices
- )
- selected = qsa_expand_block_indices(
- topk_indices,
- logical,
- seq_lens_q,
- compress_ratio=core.compress_ratio,
- token_topk=core.token_topk,
- )
- if share:
- self._capture_selection(md, topk_indices, visible_blocks, gen_seqs, p)
- if kv.k_scale is not None:
- attn = _sparse_gqa(
- core,
- kv,
- q=q_dst,
- block_table=block_table,
- selected_tokens=selected,
- request_indices=req_idx,
- tokens_per_block=tokens_per_block,
- query_positions=None,
- compress_ratio=None,
- output_gate=gate,
- )
- gated = True
- else:
- gated = qsa_paged_sparse_gqa.fuses_output_gate(q_dst, kv.k, selected, gate)
- attn = qsa_paged_sparse_gqa(
- q=q_dst,
- k_cache=kv.k,
- v_cache=kv.v,
- block_table=block_table,
- selected_tokens=selected,
- request_indices=req_idx,
- tokens_per_block=tokens_per_block,
- softmax_scale=core.softmax_scale,
- query_positions=None,
- compress_ratio=None,
- output_gate=gate if gated else None,
- )
- out = torch.reshape(attn, [rows, core.q_width])
- else:
- qkv = torch.empty([rows, core.qkv_width], dtype=dt, device=dev)
- qkv.copy_(qkv_view)
- attn_out = torch.empty([rows, core.q_width], dtype=dt, device=dev)
- out = thop_attention(
- q=qkv,
- output=attn_out,
- local_layer_idx=layer_idx,
- predicted_tokens_per_seq=p,
- num_heads=core.heads_q,
- num_kv_heads=core.heads_kv,
- head_size=core.head_dim,
- quant_mode=core.quant_mode,
- kv_scale_orig_quant=core.kv_scale_orig_quant,
- kv_scale_quant_orig=core.kv_scale_quant_orig,
- **step,
- **_CALL_CONSTANTS,
- )
- if not gated:
- fused_sigmoid_mul_inplace(out, gate)
- return gemm(out, M.o)
-
- def _capture_selection(self, md, topk_indices, visible_blocks, batch, p) -> None:
- """IndexShare: stash, per request, the selection of its last accepted draft-extend
- row with that row's complete-group count, in the draft layer's row of the shared
- buffers. Every request carries `p` rows here, and the worker continues the draft
- loop from row `num_accepted - 1` of each (the golden token counts as accepted)."""
- assert md.qsa_mtp_num_accepted is not None, (
- "IndexShare needs the worker's accepted counts to pick each request's capture "
- "row; the draft loop did not supply them"
- )
- accepted = md.qsa_mtp_num_accepted[:batch].to(torch.int64)
- rows = torch.arange(batch, dtype=torch.int64, device=topk_indices.device) * p
- rows = rows + torch.clamp(accepted - 1, min=0, max=p - 1)
- slot = self.core._draft_bound.share_slot
- md.qsa_shared_topk_indices[slot, :batch].copy_(torch.index_select(topk_indices, 0, rows))
- md.qsa_shared_topk_visible_blocks[slot, :batch].copy_(
- torch.index_select(visible_blocks, 0, rows)
- )
- # Host-side readiness; see PrefillTarget._capture_selection.
- md.qsa_shared_topk_captured_slots.add(slot)
-
- def mtp_forward(
- self,
- embed_tokens: torch.Tensor,
- all_rank_num_tokens,
- input_ids: torch.Tensor,
- position_ids: torch.Tensor,
- hidden_states: torch.Tensor,
- attn_metadata: AttentionMetadata,
- **kwargs,
- ) -> torch.Tensor:
- """One draft step over generation rows only (every step past the first in an
- engine forward, and the first when the batch had no context).
-
- `spec_metadata` and `all_rank_num_tokens` are accepted and unused."""
- core = self.core
- assert isinstance(attn_metadata, TrtllmAttentionMetadata)
- assert core._bound is not None, (
- "the trunk's cache-manager bind has not run; the shell calls the "
- "core before the worker, so this cannot be reached first"
- )
- md = attn_metadata
- rows = hidden_states.shape[0]
- assert input_ids.shape[0] == rows, (
- f"the draft step's {input_ids.shape[0]} token ids and {rows} "
- "hidden-state rows must describe the same tokens"
- )
- assert hidden_states.shape[1] == core.stream, (
- f"the MTP module reads the trunk's {core.stream}-wide four-stream "
- f"residual; the runtime handed it {hidden_states.shape[1]} columns"
- )
-
- # Projected on every call: the worker rewrites `attn_metadata` in place
- # between draft steps (see PrefillTarget.mtp_forward).
- step = self.step_args(md)
- gen_seqs = md.num_seqs
- p = rows // gen_seqs if gen_seqs else 1
- assert p * gen_seqs == rows, (
- f"{rows} generation rows do not divide over {gen_seqs} "
- "generation sequences; a ragged per-sequence draft length cannot be "
- "expressed to the attention call"
- )
-
- # The KV cache written here is the draft manager's, swapped in by the worker
- # per call (see PrefillTarget.mtp_forward); take the layer index from the
- # current manager every time.
- cm = md.kv_cache_manager
- assert cm is not None, "a paged KV cache is required"
- offsets = cm.layer_offsets
- assert core.mtp_layer_idx in offsets, (
- f"the MTP module's attention is global layer index "
- f"{core.mtp_layer_idx}, which the cache manager "
- f"{type(cm).__name__} does not map (it maps {len(offsets)} layers, "
- f"is_draft={getattr(cm, 'is_draft', None)}); the engine allocated "
- "no KV layer for the draft module"
- )
- layer_idx = offsets[core.mtp_layer_idx]
- window = cm.max_attention_window_vec[layer_idx]
- assert window is None or window == md.max_seq_len, (
- f"the MTP layer's pool window is {window!r} but this call passes "
- f"max_seq_len {md.max_seq_len} as the attention window; a sliding "
- "window is not implemented here"
- )
-
- # The draft layer's QSA planes and the page table that addresses them. With a
- # separate draft manager the engine's QSA metadata swaps `qsa_block_table` for the
- # table it decodes from the draft manager's offsets (`prepare_for_draft_forward`);
- # a table still decoding the trunk's offsets would read and write another pool's
- # pages.
- draft = core._draft_bound
- if draft is None or cm is not draft.cm:
- core._bind_draft_cache_manager(md)
- draft = core._draft_bound
- assert cm is core._bound.cm or md.qsa_block_table is getattr(
- md, "qsa_draft_block_table", None
- ), (
- "the draft step is bound to a separate draft KV cache manager but "
- "qsa_block_table is not that manager's table: the engine's QSA metadata did "
- "not swap its page table for the draft forward"
- )
- # The `qsa_*` row mapping is sized from the metadata's token count, which the
- # worker rewrites together with it between draft steps.
- assert rows == md.num_tokens, (
- f"the draft step carries {rows} rows for a batch its metadata describes as "
- f"{md.num_tokens} tokens; the QSA row mapping covers the second"
- )
- # The trunk's dense/sparse decision from the same host buffer: the worker advances
- # only the device KV lengths through the draft loop, so every draft step of one
- # engine forward takes the branch its trunk took, in the CUDA-graph family the
- # runner chose for that trunk.
- sparse = gen_seqs > 0 and int(md.kv_lens_runtime[:gen_seqs].max()) > core.dense_threshold
- # IndexShare: the loop's draft-extend step (`p` rows per request) runs the indexer
- # and, with sharing on, captures each request's selection; every draft decode step
- # after it (one row per request) reuses that capture instead of running the indexer.
- share = core.mtp_index_share and md.qsa_in_mtp_draft_loop
- reuse = share and md.qsa_indexer_skip_topk
- assert not share or reuse == (p == 1), (
- f"the worker's reuse flag ({md.qsa_indexer_skip_topk}) disagrees with the "
- f"step's shape ({p} rows per request): only the draft decode steps after the "
- "draft-extend step reuse its selection"
- )
-
- # Copy into a fresh contiguous `[3, N]`: a sliced view keeps the source row
- # stride and the op rejects it under `use_mrope`.
- assert position_ids.shape[0] == _QSA_POSITION_AXES, (
- f"expected the mRoPE position layout with {_QSA_POSITION_AXES} "
- f"coordinate rows, got {tuple(position_ids.shape)}"
- )
- # The first draft step reads the trunk's positions buffer, whose int32 copy the trunk
- # made (taken once); a contiguous int32 `position_ids` (every later draft step) is
- # read in place. The step's ops only read positions.
- trunk, self._trunk_pos3 = self._trunk_pos3, None
- if trunk is not None and trunk[0] == _position_key(position_ids):
- pos3 = trunk[1]
- elif position_ids.dtype == torch.int32 and position_ids.is_contiguous():
- pos3 = torch.reshape(position_ids, [_QSA_POSITION_AXES, -1])
- else:
- pos3 = _copy_positions(position_ids)
- # The trunk's two other views, from `pos3`: `pos` (coordinate 0 alone, which the
- # index rope reads) and `coords` (the `[N, 3]` transpose the position plane stores).
- pos = torch.reshape(torch.split(pos3, [1, _QSA_POSITION_AXES - 1], 0)[0], [-1])
- coords = torch.transpose(pos3, 0, 1)
-
- return self._mtp_layer(
- embed_tokens,
- input_ids,
- hidden_states,
- md,
- step,
- draft,
- layer_idx,
- rows,
- gen_seqs,
- p,
- sparse,
- share,
- reuse,
- pos3,
- pos,
- coords,
- )
-
-
-class _VisionBlockWeights(NamedTuple):
- """One vision block's weights (`.t()` GEMM views). LayerNorm has a bias, unlike
- every norm in the language tower."""
-
- n1_w: torch.Tensor
- n1_b: torch.Tensor
- qkv: torch.Tensor
- qkv_b: torch.Tensor
- proj: torch.Tensor
- proj_b: torch.Tensor
- n2_w: torch.Tensor
- n2_b: torch.Tensor
- fc1: torch.Tensor
- fc1_b: torch.Tensor
- fc2: torch.Tensor
- fc2_b: torch.Tensor
-
-
-class _VisionWeights(NamedTuple):
- """The tower's weights: patch embedding, position table, blocks and patch merger."""
-
- patch_w: torch.Tensor
- patch_b: torch.Tensor
- pos: torch.Tensor
- blocks: list
- m_n_w: torch.Tensor
- m_n_b: torch.Tensor
- m_fc1: torch.Tensor
- m_fc1_b: torch.Tensor
- m_fc2: torch.Tensor
- m_fc2_b: torch.Tensor
-
-
-def _vision_views(core: Qwen4ExpTp1ModelingV2) -> _VisionWeights:
- w = core.w
- return _VisionWeights(
- # Not transposed: the entry takes the checkpoint's `[out, in, kt, kh, kw]`.
- patch_w=w["vis_patch_w"],
- patch_b=w["vis_patch_b"],
- pos=w["vis_pos"],
- blocks=[
- _VisionBlockWeights(
- n1_w=w[f"l{b}_vis_n1_w"],
- n1_b=w[f"l{b}_vis_n1_b"],
- # Row layout (q heads, k, v) is already thop_attention's packed form.
- qkv=w[f"l{b}_vis_qkv_w"].t(),
- qkv_b=w[f"l{b}_vis_qkv_b"],
- proj=w[f"l{b}_vis_proj_w"].t(),
- proj_b=w[f"l{b}_vis_proj_b"],
- n2_w=w[f"l{b}_vis_n2_w"],
- n2_b=w[f"l{b}_vis_n2_b"],
- fc1=w[f"l{b}_vis_fc1_w"].t(),
- fc1_b=w[f"l{b}_vis_fc1_b"],
- fc2=w[f"l{b}_vis_fc2_w"].t(),
- fc2_b=w[f"l{b}_vis_fc2_b"],
- )
- for b in range(core.v_depth)
- ],
- m_n_w=w["vis_m_n_w"],
- m_n_b=w["vis_m_n_b"],
- m_fc1=w["vis_m_fc1_w"].t(),
- m_fc1_b=w["vis_m_fc1_b"],
- m_fc2=w["vis_m_fc2_w"].t(),
- m_fc2_b=w["vis_m_fc2_b"],
- )
-
-
-def _frame_calls(
- grid: list[tuple[int, int, int]], max_rows: int
-) -> list[list[tuple[int, int, int]]]:
- """Cut `grid` (one `(frames, height, width)` row per image or video) into tower calls
- of whole frames, in order, each at most `max_rows` patch rows. A call is its grid
- rows: the frames of one item that fall in it are one `(frames, height, width)` row."""
- calls: list[list[tuple[int, int, int]]] = [[]]
- room = max_rows
- for frames, height, width in grid:
- rows = height * width
- assert rows <= max_rows, (
- f"one frame is {height}x{width} = {rows} patch rows; the vision ops are "
- f"certified up to {max_rows} rows in one call"
- )
- while frames:
- take = min(frames, room // rows)
- if take == 0:
- calls.append([])
- room = max_rows
- continue
- calls[-1].append((take, height, width))
- room -= take * rows
- frames -= take
- return calls
-
-
-class VisionTower:
- """One request's images or videos through the 27-block ViT and its patch merger.
-
- Not a target: runs only for a context request carrying media, from the shell's
- splice. Attention is bidirectional per frame (a video's frame group is a frame),
- expressed to `thop_attention` as one context "sequence" per frame; running the
- rows as one sequence is 315x-582x outside tolerance. The position-embedding
- interpolation and 2-D rope must both emit in spatial-merge order (what the
- merger's 4-to-1 reshape consumes); a raster-order emitter gives a correctly
- shaped, silently wrong tower.
-
- The frames run in calls of at most `_VISION_CALL_MAX_PATCHES` rows, cut only between
- frames: attention is per frame, the position embedding and rope are per grid row
- and repeat per frame, the norms and MLPs are row-wise, and a frame's `h * w` is a
- multiple of `merge**2`, so a cut keeps every merger block whole. One frame past the
- bound cannot be served."""
-
- def __init__(self, core: Qwen4ExpTp1ModelingV2) -> None:
- self.core = core
- self._w = _vision_views(core)
- # Persistent scratch: `thop_attention` grows `workspace` in place; never
- # the engine's own workspace.
- self._ws: torch.Tensor | None = None
-
- def __call__(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor:
- """`pixel_values` is `[sum(t*h*w), in_channels*kt*kh*kw]`, one flat patch per
- row, channel-slowest. `grid_thw` is one `(frames, height, width)` row per
- image or video and stays on the host (no synchronization). Returns
- `[sum(t*h*w) / merge**2, out_hidden_size]`, one row per placeholder, in grid
- row order."""
- core = self.core
- dt = pixel_values.dtype
- assert dt == core.w["embed"].dtype, (
- f"the tower runs at the activation dtype {core.w['embed'].dtype}; the "
- f"runtime handed pixel_values as {dt}. The input processor's `dtype` "
- "property is what decides this, and no cast happens here"
- )
- assert pixel_values.dim() == 2 and pixel_values.shape[1] == core.v_pixel_width, (
- f"pixel_values must be [patches, {core.v_pixel_width}] flat patch rows; "
- f"got {tuple(pixel_values.shape)}"
- )
-
- grid: list[tuple[int, int, int]] = []
- for row in grid_thw.tolist():
- assert len(row) == 3, f"grid_thw rows are (t, h, w); got {row}"
- frames, height, width = (int(v) for v in row)
- assert frames >= 1 and height >= 1 and width >= 1, row
- assert height % core.v_merge == 0 and width % core.v_merge == 0, (
- f"grid sides must divide by spatial_merge_size {core.v_merge}; got "
- f"height={height}, width={width}"
- )
- grid.append((frames, height, width))
- assert grid, "grid_thw describes no image or video"
- described = sum(t * h * w for t, h, w in grid)
- assert described == pixel_values.shape[0], (
- f"grid_thw describes {described} patch rows and pixel_values carries "
- f"{pixel_values.shape[0]}"
- )
-
- outputs = []
- start = 0
- for call in _frame_calls(grid, _VISION_CALL_MAX_PATCHES):
- rows = sum(t * h * w for t, h, w in call)
- outputs.append(
- self._run(pixel_values[start : start + rows], call, grid_thw.new_tensor(call))
- )
- start += rows
- return outputs[0] if len(outputs) == 1 else torch.cat(outputs, 0)
-
- def _run(
- self, pixel_values: torch.Tensor, grid: list[tuple[int, int, int]], grid_thw: torch.Tensor
- ) -> torch.Tensor:
- """One tower call over whole frames: `grid` and `grid_thw` (the same rows, as a
- host list and a host tensor) describe exactly the rows of `pixel_values`."""
- core = self.core
- dt, dev = pixel_values.dtype, pixel_values.device
- assert pixel_values.is_cuda, "pixel_values must reach the tower on the device"
-
- # ---- the frame segmentation, host-side -------------------------------
- # One attention segment per frame: a video's frames do not attend to each other.
- seg_lens: list[int] = []
- for frames, height, width in grid:
- seg_lens.extend([height * width] * frames)
- num_patches = pixel_values.shape[0]
- num_segments = len(seg_lens)
- max_seg = max(seg_lens)
-
- # Batch-state tensors for the encoder-only call; they describe the frames.
- host_lens = torch.tensor(seg_lens, dtype=torch.int32)
- host_types = torch.zeros([num_segments], dtype=torch.int32)
- host_total = torch.tensor([num_patches, 0], dtype=torch.int32)
- dev_lens = host_lens.to(dev)
- if self._ws is None:
- self._ws = torch.empty([0], dtype=torch.int8, device=dev)
-
- # ---- patch embedding + absolute position -----------------------------
- V = self._w
- h = torch.add(
- patch_embed_conv3d(pixel_values, V.patch_w, V.patch_b),
- pos_embed_interpolate(V.pos, grid, core.v_merge),
- )
- attn_out = torch.empty([num_patches, core.v_hidden], dtype=dt, device=dev)
-
- for B in V.blocks:
- qkv = cublas_mm(layer_norm(h, B.n1_w, B.n1_b, _VISION_NORM_EPS), B.qkv, B.qkv_b)
- qkv_h = torch.reshape(qkv, [num_patches, 3 * core.v_heads, core.v_head_dim])
- q_view, k_view, _ = torch.split(qkv_h, [core.v_heads] * 3, 1)
- q_rot, k_rot = vision_rope_2d(
- q_view, k_view, grid_thw, core.v_merge, theta=_VISION_ROPE_THETA
- )
- q_view.copy_(q_rot)
- k_view.copy_(k_rot)
- attn_out = thop_attention(
- q=qkv,
- output=attn_out,
- workspace=self._ws,
- sequence_length=dev_lens,
- host_past_key_value_lengths=host_lens,
- host_total_kv_lens=host_total,
- context_lengths=dev_lens,
- host_context_lengths=host_lens,
- host_request_types=host_types,
- max_num_requests=num_segments,
- max_context_length=max_seg,
- max_seq_len=max_seg,
- attention_window_size=max_seg,
- num_contexts=num_segments,
- num_ctx_tokens=num_patches,
- num_heads=core.v_heads,
- num_kv_heads=core.v_heads,
- head_size=core.v_head_dim,
- **_VISION_CALL_CONSTANTS,
- )
- h = torch.add(h, cublas_mm(attn_out, B.proj, B.proj_b))
- mlp = cublas_mm(layer_norm(h, B.n2_w, B.n2_b, _VISION_NORM_EPS), B.fc1, B.fc1_b)
- # Non-gated MLP with the tanh GELU (not the merger's erf).
- h = torch.add(h, cublas_mm(gelu_tanh(mlp), B.fc2, B.fc2_b))
-
- # ---- the patch merger ------------------------------------------------
- # The norm runs over `hidden` before the 4-to-1 shuffle; the reshape is legal
- # only because both order-bearing steps emitted in spatial-merge order.
- merged = torch.reshape(
- layer_norm(h, V.m_n_w, V.m_n_b, _VISION_NORM_EPS), [-1, core.v_merge_dim]
- )
- # The erf GELU (`nn.GELU()`); within one bf16 ulp of tanh, hence a separate entry.
- return cublas_mm(gelu_erf(cublas_mm(merged, V.m_fc1, V.m_fc1_b)), V.m_fc2, V.m_fc2_b)
-
-
-class MTPLayer:
- """The checkpoint's MTP module, replayed once per draft step: one more decoder
- layer (two hyper-connections, gated GQA, 512-expert MoE) behind a two-branch
- front end (`fc_embedding` over the NEXT token's embedding, `fc_hidden` over the
- hidden state per stream). It returns `x` un-collapsed; `shared_head` owns the
- module's mixer. The step runs on the targets (`mtp_forward`) over `core._mtp`.
-
- The MoE is FP8 128x128 weight blocks (the trunk's is NVFP4); re-quantizing it
- would change the checkpoint's meaning, visible only as lower acceptance.
-
- Attention is the trunk's QSA layer over the module's own indexer. Every draft step
- maintains the layer's index plane, then runs the dense core while the trunk's host KV
- lengths stay within `dense_seq_len_threshold`, and selection plus sparse GQA above it.
- With a separate draft KV manager the layer's K/V, index and position pages live in
- that manager's pools: the worker swaps `attn_metadata.kv_cache_manager` and the block
- offsets for the draft loop, `QSAAttentionMetadata.prepare_for_draft_forward` swaps
- `qsa_block_table` to the table it decodes from the draft manager's offsets, and the
- core binds the draft manager's planes once per manager (`_DraftBound`). The planes need
- no rollback: each engine step's first draft call rewrites every position from the
- newly committed length on.
-
- IndexShare (`sparse_attention_config.index_share_for_mtp_iteration`) runs the draft
- indexer once per MTP iteration instead of once per draft step. The draft-extend step,
- the loop's first, runs it and captures per request the selection of the row the loop
- continues from, into the metadata's shared buffers; every draft decode step after it
- skips the index branch and rebuilds its selection from that capture with
- `qsa_mtp_reuse_topk`, which appends the groups completed since; a full row gives up
- its oldest ids to them. The expansion adds the incomplete tail, so the tokens drafted
- since the capture stay visible. Every selection is canonical first (ascending,
- padding last), so the reuse keeps the order.
-
- Not an `nn.Module`: the weights live in the core's `ParameterDict`, and a second
- registration would break the weight load's both-directions coverage assert."""
-
- def __init__(self, core: Qwen4ExpTp1ModelingV2, logits_processor) -> None:
- self.core = core
- self.logits_processor = logits_processor
-
- def __call__(self, *args, **kwargs) -> torch.Tensor:
- return self.forward(*args, **kwargs)
-
- def shared_head(
- self, hidden_states, lm_head, attn_metadata, return_context_logits
- ) -> torch.Tensor:
- """The module's own head: its stack-final hyper-connection mixer (a full gated
- mixer with its own parameters, not a plain RMSNorm), then the trunk's
- `lm_head` through the shell's `logits_processor`."""
- core = self.core
- F = core._mtp_fused
- # The decode target's norm-only site: the GEMV pair at one row, the fused norm and
- # mixer above (fp32 logits, as the trunk's stack-final mixer).
- _, final, _ = core._decode._hc_site(
- hidden_states, None, None, F.fhc_norm, F.fhc_first, F.fhc_up, False
- )
- # The draft head reads its FP8 copy of `lm_head` (`post_load_weights`).
- lm_head = _Fp8DraftHead(*core._mtp_head_fp8, core._lm_head_g, core._mtp_head_nv)
- return self.logits_processor.forward(
- final,
- lm_head,
- attn_metadata,
- return_context_logits,
- widen=not getattr(self.logits_processor, "draft_argmax_only", False),
- )
-
- def forward(
- self,
- embed_tokens: torch.Tensor,
- all_rank_num_tokens,
- input_ids: torch.Tensor,
- position_ids: torch.Tensor,
- hidden_states: torch.Tensor,
- attn_metadata: AttentionMetadata,
- **kwargs,
- ) -> torch.Tensor:
- """Route the draft step by the phase this call's metadata shows, per call: the
- worker rewrites `attn_metadata` between draft steps, so one engine step's
- draft calls can cross from prefill to decode."""
- return self.core._select_target(attn_metadata).mtp_forward(
- embed_tokens,
- all_rank_num_tokens,
- input_ids,
- position_ids,
- hidden_states,
- attn_metadata,
- **kwargs,
- )
-
-
-class DraftModel:
- """The container the spec worker reaches the drafter through. The runtime
- inspects only `mtp_layers` (`[0]`: MTP-Eagle replays one layer),
- `embed_tokens`, `lm_head`, and `model.d2t` (correctly absent: shared
- vocabulary).
-
- `embed_tokens` resolves through the core on every read: this container is
- built while parameters are meta tensors, and materialization replaces the
- tensor objects, so a captured reference stays on meta."""
-
- def __init__(self, core: Qwen4ExpTp1ModelingV2, lm_head, logits_processor) -> None:
- self.core = core
- self.mtp_layers = [MTPLayer(core, logits_processor)]
- self.lm_head = lm_head
-
- @property
- def embed_tokens(self) -> torch.Tensor:
- return self.core.w["embed"]
-
-
-class SpeculativeRollback:
- """The target's auxiliary speculative-state handler: it undoes what a verify step's
- rejected draft rows wrote into the four side-cache planes the runtime's own state
- promotion does not cover.
-
- A verify step runs `p = runtime_draft_len + 1` rows per generation request through the
- trunk, and four ops write per-request state for every row without knowing how many the
- sampler keeps: the PLE n-gram id history (`ops/ple/ngram_token_history`), the PLE
- dilated-conv state (`ops/ssm/dilated_short_conv_paged`), each attention layer's QSA
- index-K plane (`qsa_paged_index_store`, `qsa_prefill_compress`) and the shared QSA
- position plane. The Gated-DeltaNet state is not here: the runtime promotes it.
-
- For a generation request that accepted `a` rows (golden token included, so
- `1 <= a <= p`) the commit rewrites:
-
- * the n-gram pool row to the id window of row `a - 1` minus its oldest id;
- * the conv pool row to columns `[a, a + state_len)` of `[pre-step state | conv input
- of the request's p rows]`;
- * every index-K and position slot written by a row of ordinal `>= a` to its
- pre-verify contents.
-
- The targets stage during the forward (`begin`, then `stage_ple_verify` before the
- fused PLE launches and `stage_ple_conv_input` after them in the PLE block, and
- `stage_index` before each attention layer's store) and the spec worker calls
- `commit_speculative_states` after its draft loop, inside the same model forward.
- Context rows and non-verify steps stage nothing, and a commit with nothing staged
- returns. Staging and commit use metadata-sized
- allocations, gathers, `torch.where` and indexed writes only, and everything staged is
- allocated by the forward and consumed by that forward's commit, so a captured graph
- holds both ends and replay re-reads the engine's buffers in place.
-
- A step opened with `fused` (the decode target's graph-captured steps) stages the PLE
- pools with one `spec_rollback_stage` launch and, with the index planes bound, commits
- with one `spec_rollback_commit` launch instead of the gathers and indexed writes."""
-
- def __init__(self) -> None:
- self._pending: dict | None = None
- # (base addresses, page strides, plane count) of every attention layer's index-K
- # plane, once the cache manager is bound: `begin` then gathers all of them at once and
- # the commit restores them in one launch (`ops/attention/qsa_rollback_planes`).
- self._planes: tuple | None = None
-
- @property
- def index_planes(self) -> tuple | None:
- """`(base addresses, page strides, plane count)` of the bound index-K planes."""
- return self._planes
-
- def begin_staged(
- self,
- *,
- num_gens: int,
- rows_per_gen: int,
- slots: torch.Tensor,
- pages: torch.Tensor,
- within: torch.Tensor,
- index_all: torch.Tensor,
- ) -> None:
- """Open a generation-only verify step's staging from what `verify_step_setup`
- resolved: the requests' int64 recurrent slots, the rows' int64 QSA pages and offsets,
- and every bound index-K plane at them. The PLE staging and the commit are the
- one-launch ones (`begin(..., fused=True)`)."""
- n, p = num_gens, rows_per_gen
- assert n > 0 and p > 1, (
- f"a verify step stages generation requests carrying drafts; got {n} "
- f"requests of {p} rows"
- )
- assert pages.shape[0] == n * p and index_all.shape[1] == n * p, (
- f"{n} generation requests of {p} rows are {n * p} rows; got {pages.shape[0]} slots"
- )
- self._pending = dict(
- num_contexts=0,
- prefix=0,
- n=n,
- p=p,
- slots=slots,
- pages=pages,
- within=within,
- index=[],
- position=None,
- index_all=index_all,
- fused=True,
- ngram_stream=None,
- )
-
- def set_index_planes(self, planes: list) -> None:
- """Bind the attention layers' index-K planes (static for the manager's lifetime)."""
- ptrs, strides = _plane_table(planes)
- self._planes = (ptrs, strides, len(planes))
-
- # ------------------------------------------------------------- staging ---
-
- def discard(self) -> None:
- """Drop whatever an earlier forward staged."""
- self._pending = None
-
- def begin(
- self,
- *,
- num_contexts: int,
- prefix_tokens: int,
- num_gens: int,
- rows_per_gen: int,
- slots: torch.Tensor,
- request_indices: torch.Tensor,
- logical_positions: torch.Tensor,
- block_table: torch.Tensor,
- tokens_per_block: int,
- fused: bool = False,
- ) -> None:
- """Open a verify step's staging. The generation rows follow the `prefix_tokens`
- context rows as `num_gens` runs of `rows_per_gen`, so row `s * p + t` is ordinal `t`
- of generation request `s`. `slots` are those requests' recurrent slots;
- `request_indices` and `logical_positions` are the generation rows' own. Each row's
- QSA slot is resolved here, as `qsa_paged_index_store` resolves it. `fused` selects
- the one-launch PLE staging and commit."""
- n, p = num_gens, rows_per_gen
- assert n > 0 and p > 1, (
- f"a verify step stages generation requests carrying drafts; got {n} "
- f"requests of {p} rows"
- )
- assert logical_positions.shape[0] == n * p, (
- f"{n} generation requests of {p} rows are {n * p} rows; got "
- f"{logical_positions.shape[0]} positions"
- )
- positions = logical_positions.to(torch.long)
- pages = block_table[request_indices.to(torch.long), positions // tokens_per_block]
- self._pending = dict(
- num_contexts=num_contexts,
- prefix=prefix_tokens,
- n=n,
- p=p,
- slots=slots.to(torch.long),
- pages=pages.to(torch.long),
- within=positions % tokens_per_block,
- index=[],
- position=None,
- index_all=None,
- fused=fused,
- ngram_stream=None,
- )
- if self._planes is not None:
- # Every plane's slots before the trunk writes any of them: what each layer's
- # `stage_index` would capture before its own store, in one launch.
- ptrs, strides, count = self._planes
- pending = self._pending
- pending["index_all"] = qsa_rollback_gather(
- ptrs, strides, pending["pages"], pending["within"], count
- )
-
- def stage_ple_pools(self, ngram_pool: torch.Tensor, conv_pool: torch.Tensor) -> None:
- """Before the PLE block's first write: both pools, and the generation slots'
- pre-step conv rows."""
- pending = self._pending
- slots = pending["slots"]
- pending["ngram_pool"] = ngram_pool
- pending["conv_pool"] = conv_pool
- pending["conv_prev"] = conv_pool[slots]
-
- def stage_ple_verify(
- self, ngram_pool: torch.Tensor, conv_pool: torch.Tensor, input_ids: torch.Tensor
- ) -> None:
- """A verify step's PLE staging, before either pool is written:
- `stage_ple_pools`, plus the n-gram candidates rebuilt from the slots' pre-step id
- history and the rows' ids. Row `t`'s history after the step is the last
- `context` ids of `[history | ids of rows 0..t]` -- what the id window of row `t`
- minus its oldest id is, every generation sequence carrying its history."""
- pending = self._pending
- prefix, n, p = pending["prefix"], pending["n"], pending["p"]
- if pending["fused"]:
- pending["ngram_pool"] = ngram_pool
- pending["conv_pool"] = conv_pool
- pending["ngram_stream"], pending["conv_prev"] = spec_rollback_stage(
- ngram_pool, conv_pool, pending["slots"], input_ids[prefix : prefix + n * p], p
- )
- return
- self.stage_ple_pools(ngram_pool, conv_pool)
- context = ngram_pool.shape[1]
- ids = torch.reshape(input_ids[prefix : prefix + n * p], [n, p]).to(ngram_pool.dtype)
- stream = torch.cat([ngram_pool[pending["slots"]], ids], 1)
- # Window k of the unfold is the history after k rows.
- pending["ngram_candidates"] = stream.unfold(1, context, 1)[:, 1:]
-
- def stage_ple_conv_input(self, conv_input: torch.Tensor) -> None:
- """The generation rows' conv input (`[T, C]`, token-major), as the fused PLE
- launch wrote it into a buffer the forward allocated for this staging alone. A
- mixed batch's rows are copied out, so its context prefix is not kept alive
- through the draft loop."""
- pending = self._pending
- prefix, n, p = pending["prefix"], pending["n"], pending["p"]
- conv_x = torch.reshape(conv_input[prefix:], [n, p, -1])
- pending["conv_x"] = conv_x.clone() if prefix else conv_x
-
- def stage_index(self, index_cache: torch.Tensor, position_cache: torch.Tensor) -> None:
- """Before an attention layer's index store: the generation rows' slots of its
- index-K plane, and, on the step's first attention layer (the first to write the
- shared plane), of the position plane."""
- pending = self._pending
- pages, within = pending["pages"], pending["within"]
- if pending["position"] is None:
- pending["position"] = (position_cache, position_cache[pages, within])
- if pending["index_all"] is not None:
- pending["index"].append(None)
- return
- pending["index"].append((index_cache, index_cache[pages, within]))
-
- # -------------------------------------------------------------- commit ---
-
- def commit_speculative_states(
- self,
- num_accepted_tokens: torch.Tensor,
- state_indices: torch.Tensor,
- num_contexts: int,
- ) -> None:
- """Rewrite the staged planes to the accepted prefix. Called by the spec worker."""
- pending, self._pending = self._pending, None
- if pending is None:
- return
- n, p = pending["n"], pending["p"]
- assert num_contexts == pending["num_contexts"], (
- f"the worker commits with {num_contexts} context requests and the "
- f"forward staged {pending['num_contexts']}"
- )
- assert state_indices.shape[0] == num_contexts + n, (
- f"the worker hands {state_indices.shape[0]} slots for a batch the "
- f"forward staged as {num_contexts} + {n} requests"
- )
- assert pending.get("conv_x") is not None and len(pending["index"]) > 0, (
- "a verify step staged neither its PLE rows nor any attention layer; "
- "the forward did not run the staging it opened"
- )
- if pending["ngram_stream"] is not None and pending["index_all"] is not None:
- ptrs, strides, _ = self._planes
- cache, old = pending["position"]
- spec_rollback_commit(
- num_accepted_tokens[num_contexts : num_contexts + n],
- state_indices[num_contexts:],
- pending["ngram_pool"],
- pending["ngram_stream"],
- pending["conv_pool"],
- pending["conv_prev"],
- pending["conv_x"],
- cache,
- old,
- pending["pages"],
- pending["within"],
- ptrs,
- strides,
- pending["index_all"],
- )
- return
- accepted = num_accepted_tokens[num_contexts : num_contexts + n].to(torch.long)
- slots = state_indices[num_contexts:].to(torch.long)
- rows = torch.arange(n, device=accepted.device)
- ngram_pool, conv_pool = pending["ngram_pool"], pending["conv_pool"]
- state_len = conv_pool.shape[-1]
- ngram_pool[slots] = pending["ngram_candidates"][rows, accepted - 1]
- # [n, C, state_len + p]; window k of the unfold is the state after k tokens.
- history = torch.cat([pending["conv_prev"], torch.transpose(pending["conv_x"], 1, 2)], 2)
- conv_pool[slots] = history.unfold(2, state_len, 1)[rows, :, accepted]
-
- # [n * p]: True at every ordinal the sampler rejected.
- rejected = torch.reshape(
- torch.arange(p, device=accepted.device).unsqueeze(0) >= accepted.unsqueeze(1), [-1]
- )
- pages, within = pending["pages"], pending["within"]
- if pending["index_all"] is not None:
- ptrs, strides, _ = self._planes
- qsa_rollback_restore(ptrs, strides, pages, within, rejected, pending["index_all"])
- else:
- for cache, old in pending["index"]:
- cache[pages, within] = torch.where(
- rejected[:, None, None], old, cache[pages, within]
- )
- cache, old = pending["position"]
- cache[pages, within] = torch.where(rejected[:, None], old, cache[pages, within])
-
-
-#: Per modality, the keys under `multimodal_data[modality]`: the flat patch rows and the
-#: host-side `(t, h, w)` grid.
-_MEDIA_KEYS = {
- "image": ("pixel_values", "image_grid_thw"),
- "video": ("pixel_values_videos", "video_grid_thw"),
-}
-
-
-class _WidenedHead:
- """`lm_head` whose bf16 logits come back as fp32 from `bf16_to_fp32` (one vectorized
- kernel; torch's `.float()` is an unvectorized copy at under half the bandwidth), so the
- logits processor's own `.float()` returns them as they are. The same values bit for
- bit."""
-
- def __init__(self, lm_head) -> None:
- self.lm_head = lm_head
-
- def __call__(self, hidden_states: torch.Tensor) -> torch.Tensor:
- logits = self.lm_head(hidden_states)
- if logits.dtype == torch.bfloat16 and logits.is_contiguous():
- return bf16_to_fp32(logits)
- return logits
-
-
-class _Fp8DraftHead:
- """The draft head's projection over an FP8 copy of `lm_head`: up to `fp8_gemv`'s rows on
- the packed block-scaled copy (bf16 activations, one launch at HBM rate) -- or, given the
- NVFP4 copy `nv` (`_MTP_HEAD_NVFP4_ENV`), `nvfp4_gemv` over it --, more rows as per-token
- e4m3 activations (`quantize_e4m3_activation`) times the per-row e4m3 weight, scaled in
- cuBLASLt's epilogue; bf16 out (`_WidenedHead` widens it as it does the bf16 head's)."""
-
- def __init__(
- self, weight_fp8: torch.Tensor, weight_scale: torch.Tensor, packed=None, nv=None
- ) -> None:
- self.weight_fp8, self.weight_scale, self.packed = weight_fp8, weight_scale, packed
- self.nv = nv
- # The same per-row copy as `mid_rows_mm` reads it (views).
- self.mid = _head_mid(weight_fp8, weight_scale)
-
- def __call__(self, hidden_states: torch.Tensor) -> torch.Tensor:
- rows = hidden_states.shape[0]
- if rows <= _FP8_GEMV_HEAD_ROWS and _gemv_ready(hidden_states):
- if self.nv is not None:
- return nvfp4_gemv_op(hidden_states, self.nv)
- if self.packed is not None:
- return fp8_gemv_op(hidden_states, self.packed)
- if rows <= _MID_HEAD_ROWS and _gemv_ready(hidden_states):
- return mid_rows_mm_op(hidden_states, self.mid)
- x8, xs = torch.ops.tensorrt_llm.quantize_e4m3_activation(hidden_states)
- return torch._scaled_mm(
- x8,
- self.weight_fp8.t(),
- scale_a=xs.float().reshape(-1, 1),
- scale_b=self.weight_scale,
- out_dtype=hidden_states.dtype,
- )
-
-
-def _head_mid(weight_fp8: torch.Tensor, weight_scale: torch.Tensor) -> MidWeight:
- """`mid_rows_mm`'s view of a per-row-scaled e4m3 head copy (`_fp8_rowwise`)."""
- return MidWeight(weight_fp8, weight_scale.reshape(-1))
-
-
-def _fp8_rowwise(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
- """`weight` [rows, cols] bf16 -> (e4m3 [rows, cols], fp32 [1, rows]): one scale a row,
- amax / 448, so each row's largest magnitude maps to e4m3's largest finite value."""
- scale = torch.empty(weight.shape[0], dtype=torch.float32, device=weight.device)
- w8 = torch.empty_like(weight, dtype=torch.float8_e4m3fn)
- for lo in range(0, weight.shape[0], 16384):
- hi = min(lo + 16384, weight.shape[0])
- chunk = weight[lo:hi].float()
- scale[lo:hi] = (chunk.abs().amax(dim=1) / 448.0).clamp(min=1e-12)
- w8[lo:hi] = (chunk / scale[lo:hi, None]).to(torch.float8_e4m3fn)
- return w8, scale.reshape(1, -1).contiguous()
-
-
-class _Fp8VerifyHead:
- """The target's `lm_head` under the dense FP8 mode: `_Fp8DraftHead` over the same copies --
- decode rows up to `fp8_gemv`'s on the packed block-scaled copy, more rows (a large verify
- step) on the per-row-scaled copy through cuBLASLt with per-token activation scales."""
-
- def __init__(self, rowwise: tuple, packed) -> None:
- self.head = _Fp8DraftHead(*rowwise, packed)
-
- def __call__(self, hidden_states: torch.Tensor) -> torch.Tensor:
- return self.head(hidden_states)
-
-
-class _WideningLogitsProcessor(LogitsProcessor):
- """The shell's logits processor (its row gather, its fp32 logits) over `_WidenedHead`;
- with `fp8_head` set (the dense FP8 mode) the target's own head runs `_Fp8VerifyHead`.
-
- `widen=False` (with `return_context_logits`, the speculative paths) returns the head's
- own bf16 logits: for a consumer that only takes their row argmax, which is the same
- token over bf16 as over their exact fp32 widening (same order, ties, NaN and +-0)."""
-
- fp8_head = None # (rowwise copy, packed copy) under the dense FP8 mode
- #: Set by the shell's forward each engine step: the step's draft logits feed only the
- #: greedy draft sampler's argmax (see `_spec_argmax_only`).
- draft_argmax_only = False
-
- def forward(
- self,
- hidden_states,
- lm_head,
- attn_metadata,
- return_context_logits: bool = False,
- widen: bool = True,
- ) -> torch.Tensor:
- if self.fp8_head is not None and not isinstance(lm_head, _Fp8DraftHead):
- lm_head = _Fp8VerifyHead(*self.fp8_head)
- if not widen and return_context_logits:
- return lm_head(hidden_states)
- return super().forward(
- hidden_states, _WidenedHead(lm_head), attn_metadata, return_context_logits
- )
-
-
-def _spec_argmax_only(spec_worker, spec_metadata) -> tuple[bool, bool]:
- """`(verify, draft)`: whether this step's verify / draft logits are consumed only by a
- row argmax, so they can stay bf16.
-
- Both need no guided decoder (it edits the logits). The draft logits then feed only
- the draft sampler, an argmax unless rejection sampling runs on a non-greedy batch. The
- verify logits are also the step's `logits` output and the occurrence penalties' and
- relaxed acceptance's input; the one-model speculative sampler admits no request
- asking for logits or log probabilities, so they stay bf16 only for an all-greedy batch
- with the penalties off (a deploy-time flag) and strict acceptance. Every input is a
- deploy-time constant or `is_all_greedy_sample`, part of the CUDA graph key."""
- if getattr(spec_worker, "guided_decoder", None) is not None:
- return False, False
- draft = not spec_metadata.wants_advanced_draft_sampling
- verify = (
- spec_metadata.is_all_greedy_sample
- and not getattr(spec_metadata, "enable_penalty", False)
- and not getattr(
- getattr(spec_worker, "spec_config", None), "use_relaxed_acceptance_for_thinking", False
- )
- )
- return verify, draft
-
-
-class ModelingV2Qwen38FlashNextNvfp4Sm103Tp1(
- DecoderModelForCausalLM[Qwen4ExpTp1ModelingV2, PretrainedConfig]
-):
- # Whether the engine takes the one-engine speculative fast paths for this model; set in
- # `__init__` when it builds its speculative branch.
- spec_decode_fast_paths: bool = False
- # The step-loop hooks this model opts into (`runtime_hooks_of`); None: the defaults.
- runtime_hooks: RuntimeHooks | None = None
-
- @classmethod
- def get_preferred_kv_cache_manager_version(cls, pretrained_config=None) -> str:
- """Prefer KV cache manager V2: PLE state (`ple_layer_cache`) and QSA's index and
- position planes exist only there. Under V1 the model builds and passes a
- smoke run with the PLE layer stateless. The engine resolves this from the
- built-in class; the declaration serves a direct instantiation."""
- return "V2"
-
- def __init__(self, model_config: ModelConfig):
- cfg = model_config.pretrained_config
- assert cfg is not None
-
- # Adopt the engine-normalized text tower (`Qwen4ExpTextConfig`) as the model's
- # config: everything that sizes the pools reads
- # `model_engine.model.model_config.pretrained_config`, and the composite
- # parent mirrors text geometry through a pass that silently skips keys the
- # parent already defaults. The swap is in place because the engine keeps its
- # own reference to this `ModelConfig`.
- text = getattr(cfg, "text_config", None)
- assert text is not None, (
- "this target is the language tower of the composite Qwen4-Exp "
- "export; pretrained_config carries no text_config"
- )
- assert getattr(text, "architectures", None) == ["Qwen4ExpForCausalLM"], (
- f"the text tower must be the normalized Qwen4-Exp text config; got "
- f"architectures {getattr(text, 'architectures', None)!r}"
- )
- assert text.torch_dtype is not None, "the text tower must declare a dtype"
- # Captured before the swap: the core needs the vision half and media token ids.
- composite = cfg
- model_config.pretrained_config = text
-
- super().__init__(
- Qwen4ExpTp1ModelingV2(model_config, composite),
- config=model_config,
- hidden_size=text.hidden_size,
- vocab_size=text.vocab_size,
- )
- # The logits' bf16 -> fp32 cast as one vectorized kernel (the draft head shares it).
- self.logits_processor = _WideningLogitsProcessor()
- # Media placeholder ids for the executor; built lazily (`__init__` runs under
- # meta init).
- self._mm_token_ids_value = (self.model.image_token_id, self.model.video_token_id)
- self._mm_token_ids_cache: torch.Tensor | None = None
- # `lm_head` is outside the quantization scheme: a dense bf16 `LMHead`; a
- # quantized head would register a different parameter set.
- assert self.lm_head.weight.dtype == torch.bfloat16, (
- f"expected a bf16 lm_head; the shell built {self.lm_head.weight.dtype}, "
- "which means the engine resolved a quantized output head for a "
- "checkpoint that stores it unquantized"
- )
-
- # The speculative branch, built only when `speculative_config` asked for it.
- # Written out rather than inherited: the in-tree one-engine shell builds its
- # drafter via a module-level function dispatching on `model_type` (still the
- # upstream family name), which would construct trtllm's own MTP layer.
- self.spec_config = getattr(model_config, "spec_config", None)
- self.draft_model = None
- self.spec_worker = None
- # No separate draft checkpoint; the loader reads this, so None is stated.
- self.draft_config = None
- if self.spec_config is not None:
- assert self.model.mtp_enabled, "core built without the MTP module"
- assert hasattr(self, "logits_processor"), (
- "the inherited shell no longer exposes logits_processor; the "
- "draft head and the shell's own gather both project through it"
- )
- self.draft_model = DraftModel(self.model, self.lm_head, self.logits_processor)
- # Must not be left at `get_spec_worker`'s default False: the resource
- # manager decides separately (`_should_create_separate_draft_kv_cache`) and
- # builds the target manager from target layer masks alone, so left False
- # the MTP layer is handed a manager with no `layer_offsets` entry for it.
- self.use_separate_draft_kv_cache = should_use_separate_draft_kv_cache(self.spec_config)
- # The one-engine speculative fast paths, as `SpecDecOneEngineForCausalLM` takes them
- # where validated: the draft loop's step-0 epilogue, the recorded steady decode-step
- # prepare and the overlap loop on the execution stream. At TP1 no logits are
- # vocabulary-sharded, so the paths for those do not apply.
- self.spec_decode_fast_paths = True
- self.runtime_hooks = RuntimeHooks(loop_on_execution_stream=True)
- self.spec_worker = get_spec_worker(
- self.spec_config,
- model_config,
- model_config.mapping,
- use_separate_draft_kv_cache=self.use_separate_draft_kv_cache,
- spec_decode_fast_paths=self.spec_decode_fast_paths,
- )
- assert self.spec_worker is not None, (
- f"the engine resolved no spec worker for spec_dec_mode "
- f"{self.spec_config.spec_dec_mode!r}"
- )
- # Caches the draft->target vocab offset map (correctly None here).
- self.spec_worker.set_draft_model(self.draft_model)
- # The PLE pools and the QSA side caches are written per verify row and nothing
- # the runtime owns rolls them back; the worker calls this handler after its
- # draft loop, inside the same forward.
- self.spec_worker.register_auxiliary_state_handler(self.model.spec_rollback)
-
- # ----------------------------------------------------------- multimodal ---
- # Names the executor reads off a model class by `getattr`; each is silently
- # optional and its absence costs correctness rather than raising.
-
- @property
- def multimodal_data_device_paths(self) -> list[str]:
- """Tensors under ``multimodal_data`` the executor moves to CUDA. Absent, it moves
- everything, putting `image_grid_thw` and `video_grid_thw` on the device where
- `pos_embed_interpolate` and `vision_rope_2d` would each pay a D2H sync."""
- return [
- "image.pixel_values",
- "video.pixel_values_videos",
- "mrope_config.mrope_position_ids",
- "mrope_config.mrope_position_deltas",
- ]
-
- @property
- def multimodal_token_ids(self) -> torch.Tensor:
- """Placeholder ids whose rows the vision embeddings replace. Absent, the
- executor uses `input_ids >= vocab_size`, but this checkpoint's placeholders
- (248056, 248057) are in vocabulary, so nothing is selected and the engine
- generates fluent text while blind. Built on first read (meta init); not a
- parameter or buffer."""
- if self._mm_token_ids_cache is None:
- self._mm_token_ids_cache = torch.tensor(self._mm_token_ids_value, dtype=torch.int32)
- return self._mm_token_ids_cache
-
- def validate_request(self, request) -> None:
- """Admission check, run by the executor on every new request. A media request
- the forward would assert on fails here as that request's own error; an assert
- inside the step ends the executor under the overlap scheduler. Reads only the
- request's host-side payload, and a request without media returns at once.
-
- Refused: a frame past `_VISION_CALL_MAX_PATCHES` patches (the one bound the
- tower keeps), a malformed grid or pixel payload, media whose rows disagree with
- the prompt's placeholder count, missing or mis-sized `mrope_position_ids`, and a
- request with both images and videos whose item metadata is missing or does not
- list its items."""
- data = request.py_multimodal_data
- if not data:
- return
- grids = {}
- for modality, (pixel_key, grid_key) in _MEDIA_KEYS.items():
- media = data.get(modality)
- if media is not None:
- grids[modality] = self._media_grid(modality, media, pixel_key, grid_key)
- if not grids:
- return
-
- mrope_ids = (data.get("mrope_config") or {}).get("mrope_position_ids")
- if (
- mrope_ids is None
- or mrope_ids.dim() != 3
- or mrope_ids.shape[0] != _QSA_POSITION_AXES
- or mrope_ids.shape[2] != request.prompt_len
- ):
- raise ValueError(
- f"a media request needs the [{_QSA_POSITION_AXES}, 1, {request.prompt_len}] "
- "mrope_config.mrope_position_ids from the input processor, one position per "
- "prompt token; got "
- f"{None if mrope_ids is None else tuple(mrope_ids.shape)}"
- )
-
- if len(grids) > 1:
- metadata = get_multimodal_item_metadata(data)
- self._check_item_order(grids, metadata.item_refs if metadata is not None else None)
-
- cumsum = data.get("multimodal_embed_mask_cumsum")
- if cumsum is not None:
- merge_unit = self.model.v_merge * self.model.v_merge
- rows = sum(t * h * w for grid in grids.values() for t, h, w in grid) // merge_unit
- placeholders = int(cumsum[-1])
- if placeholders != rows:
- raise ValueError(
- f"the prompt holds {placeholders} image/video placeholder tokens and its "
- f"{' and '.join(sorted(grids))} payload produces {rows} embedding rows"
- )
-
- def _media_grid(self, modality, media, pixel_key, grid_key):
- """One modality's `(frames, height, width)` rows, checked on the host."""
- core = self.model
- pixel_values, grid_thw = media.get(pixel_key), media.get(grid_key)
- if pixel_values is None or grid_thw is None:
- raise ValueError(
- f"the {modality} payload carries {sorted(media)}; this target reads "
- f"{pixel_key!r} and {grid_key!r} from it"
- )
- if grid_thw.dim() != 2 or grid_thw.shape[0] == 0 or grid_thw.shape[1] != 3:
- raise ValueError(
- f"{grid_key} must hold one (t, h, w) row per {modality}; got shape "
- f"{tuple(grid_thw.shape)}"
- )
- grid = [tuple(int(v) for v in row) for row in grid_thw.tolist()]
- for item, (frames, height, width) in enumerate(grid):
- if min(frames, height, width) < 1 or height % core.v_merge or width % core.v_merge:
- raise ValueError(
- f"{modality} {item} has grid {(frames, height, width)}; every entry must "
- f"be positive and the sides divisible by spatial_merge_size {core.v_merge}"
- )
- if height * width > _VISION_CALL_MAX_PATCHES:
- raise ValueError(
- f"{modality} {item} has frames of {height}x{width} = {height * width} "
- f"patches ({height * core.v_patch}x{width * core.v_patch} pixels); the "
- f"vision tower serves frames of at most {_VISION_CALL_MAX_PATCHES} "
- "patches, so send it at a lower resolution"
- )
- patches = sum(t * h * w for t, h, w in grid)
- if tuple(pixel_values.shape) != (patches, core.v_pixel_width):
- raise ValueError(
- f"{pixel_key} is {tuple(pixel_values.shape)}; its grids describe "
- f"{patches} patch rows of {core.v_pixel_width}"
- )
- return grid
-
- @staticmethod
- def _check_item_order(grids, order) -> None:
- """A request with several modalities lists its items in prompt order."""
- if not order:
- raise ValueError(
- f"a request carrying {sorted(grids)} needs item metadata to splice them in "
- "prompt order; the payload holds each modality's rows apart"
- )
- listed = {modality: [] for modality in grids}
- for entry in order:
- try:
- modality, index = entry
- listed[modality].append(int(index))
- except (KeyError, TypeError, ValueError) as error:
- raise ValueError(
- f"item_refs entry {entry!r} is not an item of this request, which "
- f"carries {sorted(grids)}; entries are (modality, index) pairs"
- ) from error
- for modality, grid in grids.items():
- if sorted(listed[modality]) != list(range(len(grid))):
- raise ValueError(
- f"item_refs lists {modality} items {listed[modality]}; the payload "
- f"carries {len(grid)}"
- )
-
- def _multimodal_embeds(self, kwargs):
- """Run the vision tower for this batch; returns `(rows, mm_token_indices)` (the
- tower's `[R, hidden]` rows in packed-batch order and the `[R]` placeholder rows they
- replace, which the prefill target splices into its embedding), or None when the step
- carries no media.
-
- Lives in the shell because the inherited `DecoderModelForCausalLM.forward`
- drops `multimodal_params` from `**kwargs` on the no-speculation path (a
- booting, fluent, blind engine). Host-side bookkeeping: `multimodal_params`
- is a subset of the context requests, so never zip it positionally;
- `multimodal_runtime` (chunked prefill / KV reuse only) gives the cached and
- in-chunk row counts over the request's images and videos together, in prompt
- order; `mm_token_indices` names every placeholder row.
-
- The tower runs once per modality a request carries, and `_in_prompt_order`
- merges the two outputs. Positions are not derived here: every coordinate of
- a vision span, a video's temporal one included, is in the payload's
- `mrope_config.mrope_position_ids`, which the engine slices into
- `position_ids`."""
- mm_params = kwargs.get("multimodal_params") or []
- if not mm_params:
- return None
- rows = []
- for params in mm_params:
- data = getattr(params, "multimodal_data", None) or {}
- image, video = data.get("image"), data.get("video")
- if image is None and video is None:
- continue
- mrope_ids = (data.get("mrope_config") or {}).get("mrope_position_ids")
- assert (
- mrope_ids is not None
- and mrope_ids.dim() == 3
- and mrope_ids.shape[0] == _QSA_POSITION_AXES
- ), (
- f"a media request needs mrope_config.mrope_position_ids of shape "
- f"[{_QSA_POSITION_AXES}, 1, L]; got "
- f"{None if mrope_ids is None else tuple(mrope_ids.shape)}"
- )
- runtime = getattr(params, "multimodal_runtime", None)
- if runtime is not None and int(runtime.num_mm_tokens_in_chunk) == 0:
- # Every placeholder row of this request lies in another chunk.
- continue
- towers = {}
- if image is not None:
- towers["image"] = self._run_tower(image, *_MEDIA_KEYS["image"])
- if video is not None:
- towers["video"] = self._run_tower(video, *_MEDIA_KEYS["video"])
- embeds = self._in_prompt_order(params, towers)
- if runtime is not None:
- # Chunked prefill or reused prefix: splice only this chunk's slice.
- start = int(runtime.num_cached_mm_tokens)
- count = int(runtime.num_mm_tokens_in_chunk)
- total = embeds.shape[0]
- assert int(runtime.total_embeds_in_request) == total, (
- f"the executor counted {runtime.total_embeds_in_request} "
- f"placeholder rows for this request and the tower produced "
- f"{total} over its {sorted(towers)}; the two are derived from the "
- "same grids and a disagreement means the merger's row count is wrong"
- )
- assert 0 <= start <= start + count <= total, (start, count, total)
- embeds = embeds[start : start + count]
- rows.append(embeds)
- if not rows:
- return None
- vision = rows[0] if len(rows) == 1 else torch.cat(rows, 0)
- mm_token_indices = kwargs.get("mm_token_indices")
- assert mm_token_indices is not None, (
- "the executor shipped multimodal params but no mm_token_indices; "
- "that pairing is what `multimodal_token_ids` on this class buys, and "
- "without it there is no sync-free way to find the placeholder rows"
- )
- assert mm_token_indices.shape[0] == vision.shape[0], (
- f"the batch has {mm_token_indices.shape[0]} placeholder rows and the "
- f"tower produced {vision.shape[0]} embedding rows"
- )
- return vision, mm_token_indices
-
- def _run_tower(self, media, pixel_key, grid_key):
- """The tower over one modality's payload: `(rows, grid_thw)`, the rows one per
- placeholder, item by item in the order the processor emitted them."""
- pixel_values, grid_thw = media.get(pixel_key), media.get(grid_key)
- assert pixel_values is not None and grid_thw is not None, (
- f"the multimodal payload carries {sorted(media)}; this target "
- f"reads {pixel_key!r} and {grid_key!r} from it"
- )
- assert not grid_thw.is_cuda, (
- f"{grid_key} must stay on the host; see multimodal_data_device_paths"
- )
- return self.model.vision_tower(pixel_values, grid_thw), grid_thw
-
- def _in_prompt_order(self, params, towers):
- """One request's tower rows in the order its placeholders appear.
-
- `towers` maps each modality the request carries to `(rows, grid_thw)`; each
- grid row is one item of `t*h*w / merge**2` rows. One modality is already in
- prompt order. Two are reassembled item by item from `item_metadata.item_refs`
- (`(modality, index)` per item in prompt order, `index` counting that
- modality's items); without it there is no host-side way to tell which
- placeholder run is an image and which a video."""
- if len(towers) == 1:
- return next(iter(towers.values()))[0]
- metadata = params.item_metadata
- order = metadata.item_refs if metadata is not None else None
- assert order, (
- f"a request carrying {sorted(towers)} needs item metadata to splice them "
- "in prompt order; the payload holds each modality's rows apart"
- )
- merge_unit = self.model.v_merge * self.model.v_merge
- items = {}
- for modality, (embeds, grid_thw) in towers.items():
- counts = [t * h * w // merge_unit for t, h, w in grid_thw.tolist()]
- assert sum(counts) == embeds.shape[0], (
- f"the {modality} grid describes {sum(counts)} rows and the tower "
- f"produced {embeds.shape[0]}"
- )
- items[modality] = torch.split(embeds, counts, 0)
- listed = {modality: [] for modality in items}
- for modality, index in order:
- assert modality in items, (
- f"item_refs lists a {modality!r} item; this request carries {sorted(items)}"
- )
- listed[modality].append(int(index))
- for modality, parts in items.items():
- assert sorted(listed[modality]) == list(range(len(parts))), (
- f"item_refs lists {modality} items {listed[modality]}; the "
- f"payload carries {len(parts)}"
- )
- return torch.cat([items[modality][int(index)] for modality, index in order], 0)
-
- def forward(
- self,
- attn_metadata: AttentionMetadata,
- input_ids: torch.IntTensor | None = None,
- position_ids: torch.IntTensor | None = None,
- inputs_embeds: torch.FloatTensor | None = None,
- return_context_logits: bool = False,
- spec_metadata=None,
- lora_params: dict | None = None,
- **kwargs,
- ):
- """One engine step.
-
- Without a spec worker this delegates to the base forward with every argument
- it was handed (`resource_manager` stays in `**kwargs` deliberately: naming
- it would drop it from what the base receives). With one, the shell owns the
- logits gather (one-model modes set `without_logits`) and the worker needs
- ungathered hidden states; `position_ids` is passed in the engine's layout
- (flattening gives silently wrong draft positions). The vision tower runs here
- because the base forward would not pass `multimodal_params` through; its rows reach
- the core through `begin_step`, and the prefill target splices them."""
- self.model.begin_step(
- attn_metadata,
- input_ids,
- position_ids,
- inputs_embeds=inputs_embeds,
- lora_params=lora_params,
- spec_metadata=spec_metadata,
- vision=self._multimodal_embeds(kwargs),
- )
- if self.spec_worker is None:
- assert spec_metadata is None, (
- "spec_metadata arrived without a spec worker; this engine was "
- "built without a speculative_config"
- )
- return super().forward(
- attn_metadata,
- input_ids,
- position_ids,
- inputs_embeds,
- return_context_logits,
- spec_metadata,
- lora_params,
- **kwargs,
- )
- assert spec_metadata is not None, (
- "the spec worker is built but the engine passed no spec_metadata"
- )
- # `final` is the `[T, hidden]` mix for `lm_head`, `hc` the stream the MTP reads.
- final, hc = self.model.split_spec_output(
- self.model(
- attn_metadata=attn_metadata,
- input_ids=input_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- lora_params=lora_params,
- spec_metadata=spec_metadata,
- **kwargs,
- )
- )
- # `gather_ids`: one row per context request, `runtime_draft_len + 1` per
- # generation request. Logits only an argmax reads stay bf16.
- verify_argmax_only, self.logits_processor.draft_argmax_only = _spec_argmax_only(
- self.spec_worker, spec_metadata
- )
- logits = self.logits_processor.forward(
- nn.functional.embedding(spec_metadata.gather_ids, final),
- self.lm_head,
- attn_metadata,
- True,
- widen=not verify_argmax_only,
- )
- # A prefill CUDA graph pads the step's rows; the worker reshapes generation rows per
- # request, so it takes the batch's own.
- if attn_metadata.padded_num_tokens is not None:
- n = attn_metadata.num_tokens
- input_ids, hc = input_ids[:n], hc[:n]
- position_ids = position_ids[..., :n]
- return self.spec_worker(
- input_ids=input_ids,
- position_ids=position_ids,
- hidden_states=hc,
- logits=logits,
- attn_metadata=attn_metadata,
- spec_metadata=spec_metadata,
- draft_model=self.draft_model,
- resource_manager=kwargs.get("resource_manager"),
- )
-
- def load_weights(self, weights, *args, **kwargs):
- _weights.MODEL_WEIGHTS.load(self, weights)
-
- def post_load_weights(self):
- super().post_load_weights()
- self.model.post_load_weights()
- # The draft head reads a per-vocabulary-row e4m3 copy of lm_head (cuBLASLt with
- # per-token activation scales). It only proposes tokens -- the verify step's bf16 head
- # decides what is emitted -- so the copy can change acceptance, not outputs (measured
- # unchanged: gsm8k and pass-key acceptance).
- draft_fp8 = self.model.mtp_enabled
- rowwise = (
- _fp8_rowwise(self.lm_head.weight.data) if draft_fp8 or _dense_fp8_enabled() else None
- )
- if draft_fp8:
- self.model._mtp_head_fp8 = rowwise
- # The draft head's decode rows read an NVFP4 copy instead of the packed FP8 one.
- draft_nv = draft_fp8 and _mtp_head_nvfp4_enabled()
- # One packed block-scaled copy serves the draft head's decode rows and, under the dense
- # FP8 mode, the target's own head's.
- if (draft_fp8 and not draft_nv) or _dense_fp8_enabled():
- self.model._lm_head_g = pack_from_bf16(self.lm_head.weight.data)
- if _dense_fp8_enabled():
- self.logits_processor.fp8_head = (rowwise, self.model._lm_head_g)
- if rowwise is not None and self.lm_head.weight.is_cuda:
- # Every head variant decode rows past the GEMV reach, compiled now.
- mid_rows_precompile(_head_mid(*rowwise), _FP8_GEMV_HEAD_ROWS + 1, _MID_HEAD_ROWS)
- if draft_nv:
- self._build_draft_head_nvfp4(self.lm_head.weight.data)
-
- def _build_draft_head_nvfp4(self, weight: torch.Tensor) -> None:
- """The draft head's NVFP4 copy of `weight` (`_MTP_HEAD_NVFP4_ENV`), with both of
- `nvfp4_gemv`'s row-octet kernels compiled now rather than at a first serving call."""
- nv = nvfp4_pack_from_bf16(weight)
- nvfp4_gemv_precompile(nv)
- self.model._mtp_head_nv = nv
- print(
- f"draft head NVFP4 ({_MTP_HEAD_NVFP4_ENV}=1): lm_head {tuple(weight.shape)} -> "
- f"{nv.data.numel() / 2**20:.1f} MiB (e2m1 + 16-element e4m3 scales, global scale "
- f"{nv.global_scale:.3e}), decode rows <= {_FP8_GEMV_HEAD_ROWS} on nvfp4_gemv",
- flush=True,
- )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/__init__.py
deleted file mode 100644
index 1059445198db..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/__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
-"""Ops used only by this target. Shared ops stay in the catalog."""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py
deleted file mode 100644
index 7987c692ddfd..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py
+++ /dev/null
@@ -1,100 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""In-place sigmoid output gating: ``attention_output *= sigmoid(gate)``."""
-
-from __future__ import annotations
-
-import torch
-
-# Registered by the module that defines it: after `import
-# tensorrt_llm._torch.custom_ops` alone, torch.ops.trtllm has no
-# `fused_sigmoid_mul_inplace` (measured).
-import tensorrt_llm._torch.attention.kernels.fused_qk_norm_rope_gate # noqa: F401
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: The checkpoint's gated attention: 24 q heads of 256, q projection [T, 24, 512].
-_CKPT = dict(heads=24, head_dim=256, dtype=torch.bfloat16, gate="interleaved")
-
-
-class FusedSigmoidMulInplace(OpWrapper):
- """`attention_output[t, i] *= sigmoid(gate[t, i // D, i % D])` in place, `D = gate.shape[-1]`.
-
- Returns `attention_output`, the tensor it wrote. Both operands are widened to
- fp32 and the product is rounded once to `attention_output.dtype`. `gate` is
- read through its own token and head strides, so the gate half of an
- interleaved `[T, heads, 2 * D]` q projection is passed as a strided view with
- no repack (bitwise equal to its compacted copy, measured). A 2-D `gate` of
- `attention_output`'s shape is one head spanning the row.
-
- Fusion boundary: sigmoid and multiply only. The caller owns producing both
- operands and everything downstream (the o projection). `gate` is read-only.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="decode 128 tokens: 2-D gate off the fused projection row (stride 14080)",
- spec=dict(tokens=128, heads=24, head_dim=256, dtype=torch.bfloat16, gate="fused_row"),
- ),
- Cell(
- why="prefill 8192 tokens: 2-D gate off the fused projection row, 2048-wide blocks",
- spec=dict(tokens=8192, heads=24, head_dim=256, dtype=torch.bfloat16, gate="fused_row"),
- ),
- Cell(
- why="MTP decode: 1 token, 24x256 gate as the strided half of q_proj",
- spec=dict(tokens=1, **_CKPT),
- ),
- )
-
- note = """
- Metadata consumed: none. One Triton kernel, JIT-compiled on the first call,
- no autotuner: grid `(T, cdiv(hidden, block))`, `block = min(next_pow2(hidden),
- 1024 if T < 1024 else 2048)`, 4 warps.
-
- Preconditions, every one enforced loudly by the op's own bare `assert`
- (an AssertionError with no message): `attention_output` 2-D with
- `stride(-1) == 1` (row stride free); `gate` 2-D of `attention_output`'s
- shape, or 3-D with `shape[0] == T` and `shape[1] * shape[2] == hidden`,
- `stride(-1) == 1`. A CPU tensor raises ValueError from Triton (measured).
- Hence no `is_valid`. Cross-device operands were not probed.
-
- The pre-call contents of `attention_output` are destroyed; the gate is
- left bitwise unchanged.
-
- Numerics: bf16 bitwise equal to the fp32 reference except ~1e-6 of
- elements one ulp off (`tl.sigmoid`'s fast exp); default tolerance holds,
- nothing loosened. Feeding the q half instead of the gate half moves the
- result by ~2.15 max abs at [17, 6144] (measured), so the per-head indexing
- is under test.
-
- Certified: bf16 operands at the checkpoint's hidden 6144, a compact
- `attention_output`; the 2-D window of the fused projection row at 128
- tokens (decode, 1024-wide blocks) and 8192 tokens (prefill, 2048-wide
- blocks), and the strided q_proj half (MTP) at 1 token.
-
- Hard limits
- T * row stride token offset computed in int64 boundary 2^63 not probed
- (largest run 8192 rows)
-
- Not certified: 0 tokens, the strided q_proj half above 1 token, a strided
- `attention_output`, a gate dtype other than `attention_output.dtype`,
- fp16/fp32 operands, other hidden sizes, CUDA-graph capture, torch.compile.
- """
-
- def __call__(self, attention_output: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
- torch.ops.trtllm.fused_sigmoid_mul_inplace(attention_output, gate)
- return attention_output
-
- def reference(self, attention_output: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
- """What `attention_output` holds after the call; neither input is mutated.
-
- `gate.reshape(T, hidden)` is the per-head `(t, i // D, i % D)` read, for
- the 2-D and 3-D forms alike.
- """
- num_tokens, hidden = attention_output.shape
- g = gate.reshape(num_tokens, hidden).float()
- return (attention_output.float() * torch.sigmoid(g)).to(attention_output.dtype)
-
-
-fused_sigmoid_mul_inplace = FusedSigmoidMulInplace()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/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
deleted file mode 100644
index 07150279c931..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_erf.py
+++ /dev/null
@@ -1,91 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Non-gated exact (erf) GELU -- what `torch.nn.GELU()` is -- the vision patch
-merger's activation, which no config key names.
-
-Tier-3 composed: one `torch.nn.functional.gelu(x)` call, `approximate="none"`.
-`torch.ops.trtllm.moe_gelu` computes this function bitwise but is not a
-substitute: its `tile_idx_to_mn_limit` / `num_non_exiting_tiles` operands must
-be int32 *device* tensors (on the host it raises nothing and returns garbage,
-nan, or an illegal memory access), so wrapping it means launching two fills
-before the call; it also rejects fp32 and non-2-D inputs.
-"""
-
-from __future__ import annotations
-
-import math
-
-import torch
-import torch.nn.functional as F
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: `merger.linear_fc1` is `[4608, 4608]`: spatial merge 2 over a 1152-wide tower.
-MERGER_WIDTH = 4608
-
-_SQRT_2 = math.sqrt(2.0)
-
-
-class GeluErf(OpWrapper):
- """`0.5 * x * (1 + erf(x / sqrt(2)))`, elementwise, any shape, new tensor
- out in `x`'s dtype.
-
- Fusion boundary: the activation only. `merger.linear_fc1` before it and
- `merger.linear_fc2` after it are the caller's; there is no gate half and no
- multiply. This is the erf GELU and not the tanh one (`GeluTanh`, the block
- MLPs' `hidden_act`); both appear in this one tower.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: shape, dtype; `dist` "uniform" over [lo, hi] (default) or "merger"
- # (a tower hidden state scaled by the merger norm's measured gamma, grouped
- # 4-to-1, through an fc1 at the checkpoint's measured weight/bias scale).
- 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="merger activand, 8192 merged rows x 4608 bf16: the most one tower call carries",
- spec=dict(shape=(8192, MERGER_WIDTH), dtype=torch.bfloat16, lo=-16.0, hi=16.0),
- ),
- Cell(
- why="merger activand, 1 merged row x 4608 bf16",
- spec=dict(shape=(1, MERGER_WIDTH), dtype=torch.bfloat16, lo=-16.0, hi=16.0),
- ),
- )
-
- note = """
- Metadata consumed: none. Out-of-place; `x` is not mutated.
-
- Certified on contiguous bf16 `[rows, 4608]` activands (what the merger's fc1
- returns) at 1, 512 and 8192 rows (8192, from 32768 patch rows, the most one
- tower call carries).
-
- Identity is pinned by fp32 only (the identity test).
- Max |erf GELU - tanh GELU| is 4.73e-4 at x = -2.699; in fp32 the tanh form
- lands ~1400x the correct path's deviation away, 31x outside the default
- gate. In bf16 one ulp at the output scale exceeds that gap (tanh variant
- 8.24e-3 vs correct 7.77e-3 max abs), so a bf16 match does not say which GELU
- was computed.
-
- 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: 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:
- return F.gelu(x)
-
- def reference(self, x: torch.Tensor) -> torch.Tensor:
- """The closed form in fp64 from primitive ops, rounded once to `x.dtype`."""
- x64 = x.double()
- return (0.5 * x64 * (1.0 + torch.erf(x64 / _SQRT_2))).to(x.dtype)
-
-
-gelu_erf = GeluErf()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py
deleted file mode 100644
index e9cd2ebec23a..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py
+++ /dev/null
@@ -1,93 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Non-gated tanh GELU -- HuggingFace's `gelu_pytorch_tanh` -- the vision
-tower's block-MLP activation.
-
-Tier-3 composed: one `torch.nn.functional.gelu(x, approximate="tanh")` call.
-No registered op computes a plain non-gated GELU over an unquantized tensor:
-`flashinfer_gelu_tanh_and_mul` is the gated form (and silently returns a
-half-width `gelu(a) * b` when handed this tower's non-gated activand),
-`moe_gelu` is the erf GELU behind MoE tile metadata, and the
-`cute_dsl_nvfp4_dense_gemm_gelu*` ops are NVFP4 GEMM epilogues.
-"""
-
-from __future__ import annotations
-
-import math
-
-import torch
-import torch.nn.functional as F
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: The block MLP's `linear_fc1` is `[4304, 1152]`: one width, nothing split.
-INTERMEDIATE_SIZE = 4304
-
-_SQRT_2_OVER_PI = math.sqrt(2.0 / math.pi)
-
-
-class GeluTanh(OpWrapper):
- """`0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x**3)))`, elementwise,
- any shape, new tensor out in `x`'s dtype.
-
- Fusion boundary: the activation only. `linear_fc1` before it and
- `linear_fc2` after it are the caller's; there is no gate half and no
- multiply. This is the tanh GELU and not the erf one (`GeluErf`, the patch
- merger's); which one a module wants is a property of the checkpoint.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: shape, dtype; `dist` "uniform" over [lo, hi] (default) or "fc1" (a
- # unit-normal hidden state through an fc1 at the checkpoint's measured
- # weight std 0.0197 / bias std 1.676).
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="activand built through fc1 at the checkpoint's measured scale",
- spec=dict(shape=(512, INTERMEDIATE_SIZE), dtype=torch.bfloat16, dist="fc1"),
- ),
- Cell(
- why="block MLP activand at the tower call ceiling, 32768 patch rows x 4304 bf16",
- spec=dict(shape=(32768, INTERMEDIATE_SIZE), dtype=torch.bfloat16, lo=-16.0, hi=16.0),
- ),
- Cell(
- why="block MLP activand, 1 patch row x 4304 bf16",
- spec=dict(shape=(1, INTERMEDIATE_SIZE), dtype=torch.bfloat16, lo=-16.0, hi=16.0),
- ),
- )
-
- note = """
- Metadata consumed: none. Out-of-place; `x` is not mutated.
-
- Certified on contiguous bf16 `[rows, 4304]` activands (what the block MLP's
- fc1 returns) at 1, 512 and 32768 rows (32768, the tower call ceiling).
-
- Identity is pinned by 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.
-
- 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: 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:
- return F.gelu(x, approximate="tanh")
-
- def reference(self, x: torch.Tensor) -> torch.Tensor:
- """The closed form in fp64 from primitive ops, rounded once to `x.dtype`."""
- x64 = x.double()
- inner = _SQRT_2_OVER_PI * (x64 + 0.044715 * x64 * x64 * x64)
- return (0.5 * x64 * (1.0 + torch.tanh(inner))).to(x.dtype)
-
-
-gelu_tanh = GeluTanh()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/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
deleted file mode 100644
index cae68c64e1a9..000000000000
--- 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
+++ /dev/null
@@ -1,93 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""SiLU-and-mul over the leading `2 * inter` columns of each row, read through the row stride:
-the shared expert's activation taken out of its fused gate|up|gate-logit GEMM output.
-
-Tier 5 (Triton written for this target). `flashinfer_silu_and_mul` addresses row `t` at
-`t * 2 * inter`, so its input must be exactly the gate|up columns, contiguous; here they are
-the head of a wider row (the same GEMM also forms the per-token gate logit) and are read in
-place.
-"""
-
-from __future__ import annotations
-
-import torch
-import torch.nn.functional as F
-import triton
-import triton.language as tl
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: This checkpoint's shared expert: intermediate 640, the fused GEMM output row 2 * 640 + 8.
-_INTER = 640
-_WIDTH = 2 * _INTER + 8
-_BLOCK = 1024
-
-
-@triton.jit
-def _silu_and_mul_rows_kernel(x_ptr, out_ptr, stride_x, INTER: tl.constexpr, BLOCK: tl.constexpr):
- row = tl.program_id(0).to(tl.int64)
- cols = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
- mask = cols < INTER
- g = tl.load(x_ptr + row * stride_x + cols, mask=mask, other=0.0).to(tl.float32)
- u = tl.load(x_ptr + row * stride_x + INTER + cols, mask=mask, other=0.0).to(tl.float32)
- y = g / (1.0 + tl.exp(-g)) * u
- tl.store(out_ptr + row * INTER + cols, y.to(out_ptr.dtype.element_ty), mask=mask)
-
-
-class SiluAndMulRows(OpWrapper):
- """`out[t, i] = silu(x[t, i]) * x[t, inter + i]` for `i < inter`: the gate half first, as
- `flashinfer_silu_and_mul` reads it, in fp32 and rounded once to `x`'s dtype. Columns of
- `x` past `2 * inter` are not read. Returns a fresh contiguous `[T, inter]` tensor.
-
- Fusion boundary: the activation only; the GEMM that forms `x` and the one that reads
- `out` are the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(why="prefill chunk 16384 rows of the fused 1288-wide output", spec=dict(rows=16384)),
- Cell(why="short prefill step, 128 rows of the fused output", spec=dict(rows=128)),
- Cell(why="zero rows: no launch", spec=dict(rows=0)),
- )
-
- note = """
- Metadata consumed: none; every row independent. Call site: the prefill target's MoE
- block, on the shared expert's fused `[gate | up | gate logit | pad]` GEMM output.
-
- Preconditions: 2-D `x` with unit column stride, `x.shape[1] >= 2 * inter`; bf16 or fp16.
-
- Launch: `(rows, cdiv(inter, 1024))` programs of 4 warps. GB300, 16384 rows of the
- 1288-wide output: 16.4 us (`flashinfer_silu_and_mul` on a contiguous copy of the gate|up
- columns: 20.5 us for the kernel, 62.1 with the copy).
-
- Numerics: `g / (1 + exp(-g)) * u` in fp32 (Triton's `exp` and division), rounded once;
- bitwise `flashinfer_silu_and_mul` on the same values (16384 x 640 random bf16 values,
- the test's cells).
-
- Certified: bf16 at inter 640 on the 1288-wide fused output, 0, 128 and 16384 rows;
- CUDA-graph replay at 128 rows. Not certified: fp16, other widths or row counts.
- """
-
- def __call__(self, x: torch.Tensor, inter: int) -> torch.Tensor:
- rows = x.shape[0]
- out = torch.empty((rows, inter), dtype=x.dtype, device=x.device)
- if rows > 0:
- _silu_and_mul_rows_kernel[(rows, triton.cdiv(inter, _BLOCK))](
- x, out, x.stride(0), INTER=inter, BLOCK=_BLOCK, num_warps=4
- )
- return out
-
- def reference(self, x: torch.Tensor, inter: int) -> torch.Tensor:
- """fp32; the gate half first."""
- gate, up = x[:, : 2 * inter].float().chunk(2, dim=-1)
- return (F.silu(gate) * up).to(x.dtype)
-
- def is_valid(self, x: torch.Tensor, inter: int) -> None:
- assert x.dim() == 2 and x.stride(1) == 1, "x [T, width] with unit column stride"
- assert x.shape[1] >= 2 * inter, "x carries the gate and up halves"
- assert x.dtype in (torch.bfloat16, torch.float16), "bf16 / fp16"
-
-
-silu_and_mul_rows = SiluAndMulRows()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_canonical_selection.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_canonical_selection.py
deleted file mode 100644
index b2736198b176..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_canonical_selection.py
+++ /dev/null
@@ -1,214 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""QSA block selection, made canonical: one answer per row, whichever tied columns the
-selection engine happened to keep.
-
-Composed from native torch ops: no registered op rewrites a selection, and none of the
-top-k engines (`indexer_topk_prefill`, the CuTe DSL kernels, FlashInfer's top-k) takes a
-tie-break rule or fixes the order of its output.
-"""
-
-from __future__ import annotations
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: This checkpoint's block budget: `indexer_budget` 2048 tokens over
-#: `indexer_compress_ratio` 4.
-_BLOCK_TOPK = 512
-#: Score-row width the target hands over: a 260-page block table at 32 tokens
-#: per block, compressed 4:1.
-_SCORE_COLUMNS = 260 * 32 // 4
-#: The most `[rows, cols]` elements one rewrite pass holds at once; a taller matrix is
-#: rewritten in row tiles so the fp32 / int32 temporaries stay a few hundred MiB.
-_TILE_ELEMENTS = 1 << 24
-
-
-def _cell(why: str, rows, windows, logits: str, tail: str = "nan", **extra) -> Cell:
- return Cell(
- why=why,
- spec=dict(
- rows=rows,
- cols=extra.pop("cols", _SCORE_COLUMNS),
- index_topk=_BLOCK_TOPK,
- windows=windows,
- logits=logits,
- tail=tail,
- selection=extra.pop("selection", "random"),
- **extra,
- ),
- )
-
-
-def _rewrite_tile(logits, row_starts, row_ends, indices) -> None:
- num_rows, num_cols = logits.shape
- index_topk = indices.shape[1]
- starts = row_starts[:, None]
- columns = torch.arange(num_cols, dtype=torch.int32, device=logits.device)[None, :]
- # The window's scores, NaN elsewhere: NaN compares false, so nothing outside the
- # window is ever counted or kept.
- windowed = torch.where(
- (columns >= starts) & (columns < row_ends[:, None]), logits, float("nan")
- )
- # The k-th value: the smallest score the engine kept.
- picked = torch.gather(logits, 1, (indices + starts).clamp_(0, num_cols - 1).long())
- kth = picked.masked_fill_(indices < 0, float("inf")).amin(1, keepdim=True)
- at_least = windowed >= kth
- tied = windowed == kth
- # Keep every column at or above the k-th value except the lowest-index ties past
- # k: a tied column is dropped while no more than `surplus` ties lie at or before
- # it. A window of at most k has no surplus.
- surplus = at_least.sum(1, keepdim=True, dtype=torch.int32) - index_topk
- dropped = tied & (torch.cumsum(tied, 1, dtype=torch.int32) <= surplus)
- kept_seen = torch.cumsum(at_least ^ dropped, 1, dtype=torch.int32)
- # The j-th kept column is the first whose running count reaches j.
- slots = torch.arange(1, index_topk + 1, dtype=torch.int32, device=logits.device)
- slots = slots.repeat(num_rows, 1)
- torch.searchsorted(kept_seen, slots, out_int32=True, out=indices)
- indices.sub_(starts)
- indices.masked_fill_(slots > kept_seen[:, -1:], -1)
-
-
-class QsaCanonicalSelection(OpWrapper):
- """Rewrite an exact top-k selection, in place, into the window's canonical selection.
-
- For row `r` with window `[s, e) = [row_starts[r], row_ends[r])` of `logits` and
- `k = indices.shape[1]`, `indices[r]` holds on entry what a top-k engine wrote:
- window-local offsets of an exact top-k of the window, in any order, `-1` past
- `min(k, e - s)`. Let `v` be the smallest selected score. On return the row is
-
- * every window column scoring strictly above `v`, then the highest-index window
- columns scoring exactly `v` until there are `k` -- the k largest
- `(score, column)` pairs; or the whole window when `e - s <= k`;
- * as window-local offsets in ascending order, `-1` in every remaining slot.
-
- So the row is a function of `(logits, row_starts, row_ends, k)` alone: two engines,
- or two launches of one engine, that broke a tie at the k-th value differently produce
- the same row. On a row with no tie at the k-th value this is the engine's own set,
- sorted. Ties are float equality: `-0.0` and `+0.0` are one tie class. Returns
- `indices`.
-
- Fusion boundary: the rewrite only. Selecting is the engine's call before it; the
- engine must be exact -- `v` is read off its picks, and a pick below the true k-th
- value would admit more than `k` columns above it.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # The target calls this right after `qsa_topk_select`, at `k = 512` over the
- # `[rows, 2080]` score matrix whose columns past each `row_end` are allocator
- # garbage (the cells poison them). `selection` names what produced the input row:
- # `random` is an exact top-k that keeps a random subset of the tie class at the k-th
- # value, in a random order (the adversary); `radix` is the engine the target calls.
- CELLS: tuple[Cell, ...] = (
- _cell(
- "decode: 1, 128 and 256 rows (MTP3 and MTP7 at batch 32)",
- (1, 128, 256),
- ("positions_random", 2051, 8191),
- "relu",
- ),
- _cell(
- "prefill: 8192 rows from position 0, rewritten in row tiles",
- 8192,
- ("chunk", 0),
- "relu",
- ),
- _cell(
- "input from the radix engine (qsa_topk_select), zero ties",
- 64,
- ("positions_random", 2051, 8191),
- "relu_sparse",
- selection="radix",
- ),
- )
-
- note = """
- Metadata consumed: none. Stateless; every input is an argument, and `indices` is
- both the engine's selection read and the canonical row written.
-
- Preconditions (none checked):
- - `indices` holds an EXACT top-k of each window: the smallest selected score is
- the window's k-th largest. `qsa_topk_select` is exact; an approximate selection
- would be rewritten into a row with more than `k` columns at or above its k-th
- pick, truncated to the lowest `k` column indices -- wrong, silently.
- - `0 <= row_starts[r] <= row_ends[r] <= logits.shape[1]`. Anything outside the
- window is never read for a decision, so garbage, NaN and +inf there are fine
- (the cells poison it).
- - Scores inside the window are not NaN: a NaN column is never kept, and a NaN pick
- empties the row. (The engine returns the identity for a window of at most k
- without reading it; this rewrite reads it.)
-
- Cost: about 30 kernels -- elementwise passes, two int32 prefix sums over
- `[rows, cols]` and one sorted search over `[rows, k]` -- with no host read and no
- data-dependent shape, so it can be captured into a CUDA graph (not certified: no
- cell replays one). It reads the whole `cols`-wide score row, not the live window: a
- caller that knows a bound on the windows passes `logits[:, :bound]`, which gives the
- same rows. A matrix of more than 2^24 elements is rewritten in row tiles.
- """
-
- def __call__(
- self,
- logits: torch.Tensor,
- row_starts: torch.Tensor,
- row_ends: torch.Tensor,
- indices: torch.Tensor,
- ) -> torch.Tensor:
- num_rows, num_cols = logits.shape
- tile = max(1, _TILE_ELEMENTS // max(num_cols, 1))
- for lo in range(0, num_rows, tile):
- hi = min(lo + tile, num_rows)
- _rewrite_tile(logits[lo:hi], row_starts[lo:hi], row_ends[lo:hi], indices[lo:hi])
- return indices
-
- def reference(
- self,
- logits: torch.Tensor,
- row_starts: torch.Tensor,
- row_ends: torch.Tensor,
- indices: torch.Tensor,
- ) -> torch.Tensor:
- """The k largest `(score, column)` pairs of each window, ascending by column,
- `-1`-padded: a top-k over one int64 key per column whose high word orders the
- scores as fp32 does and whose low word is the column. `indices` supplies only the
- width `k`; its contents are ignored."""
- num_rows, num_cols = logits.shape
- index_topk = indices.shape[1]
- columns = torch.arange(num_cols, device=logits.device)[None, :]
- starts, ends = row_starts.long()[:, None], row_ends.long()[:, None]
- in_window = (columns >= starts) & (columns < ends)
- # `+ 0.0` folds -0.0 into +0.0, so the two are one tie class.
- bits = (logits.float() + 0.0).view(torch.int32).long()
- order = torch.where(bits < 0, bits ^ 0x7FFFFFFF, bits)
- key = torch.where(in_window, (order << 32) | columns, torch.iinfo(torch.int64).min)
- width = min(index_topk, num_cols)
- top = torch.topk(key, width, dim=1).indices
- taken = torch.gather(in_window, 1, top)
- pad = torch.iinfo(torch.int64).max
- ordered = torch.sort(torch.where(taken, top, pad), dim=1).values
- out = torch.where(ordered == pad, -1, ordered - starts).to(torch.int32)
- if width < index_topk:
- out = torch.nn.functional.pad(out, (0, index_topk - width), value=-1)
- return out
-
- def is_valid(
- self,
- logits: torch.Tensor,
- row_starts: torch.Tensor,
- row_ends: torch.Tensor,
- indices: torch.Tensor,
- ) -> None:
- # A one-element `row_starts` / `row_ends`, or a one-row `indices`, broadcasts over
- # every row of `logits` without an error: a [1] `row_starts` against 4 rows
- # returns every row framed by row 0's start.
- num_rows = logits.shape[0]
- assert row_starts.shape == (num_rows,) and row_ends.shape == (num_rows,), (
- f"row_starts and row_ends need one entry per logits row ({num_rows}); got "
- f"{tuple(row_starts.shape)} and {tuple(row_ends.shape)}"
- )
- assert indices.dim() == 2 and indices.shape[0] == num_rows, (
- f"indices must be [{num_rows}, k]; got {tuple(indices.shape)}"
- )
-
-
-qsa_canonical_selection = QsaCanonicalSelection()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_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
deleted file mode 100644
index ba1cf5acd32e..000000000000
--- 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
+++ /dev/null
@@ -1,178 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""QSA Top-K group ids expanded into raw KV token indices, plus the causal tail."""
-
-from __future__ import annotations
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-# Tier-2 sourcing: a vendored Triton callable reached as ordinary Python. QSA's
-# kernel plane has no torch.ops surface.
-from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels as _qsa_kernels
-
-#: This checkpoint's selection geometry: `indexer_compress_ratio` and
-#: `indexer_budget`, so `block_topk` 512 and an output 2051 wide.
-_CKPT = dict(compress_ratio=4, token_topk=2048)
-
-
-#: 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.
-
- `batch` is one of:
- ("sweep", n) n rows at 4096 + 13 i, length = position + 1
- ("serving", requests, context_lens) one packed step: context chunks, then
- two MTP rows per remaining request
- ("verify", starts, p) one MTP verify block of p rows per start,
- every row's length the block's end
- """
- base = dict(
- batch=batch,
- position_dtype=torch.int64, # qsa_logical_positions
- length_dtype=torch.int32, # qsa_sequence_lengths
- **_CKPT,
- )
- return {**base, **overrides}
-
-
-class QsaExpandBlockIndices(OpWrapper):
- """Expand each row's selected compressed groups into token ids, then append
- the incomplete group at the causal frontier.
-
- Wraps `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels.
- triton_expand_qsa_block_indices`: one Triton launch,
- `_expand_qsa_block_indices_kernel`. Per row (one per query token), with
- `final_topk = token_topk + compress_ratio - 1`:
-
- 1. column `c < token_topk` holds member `c % cr` of group
- `block_indices[row, c // cr]`; `-1` if that group is `-1` or the
- token is `>= sequence_lengths[row]`;
- 2. the last `(query_positions[row] + 1) % cr` visible tokens -- the
- group the indexer never scored -- go in the columns right after the
- selected groups, i.e. from column `(# non-negative groups) * cr`;
- 3. everything else is `-1`.
-
- Selection order is kept (nothing sorted, compacted or de-duplicated).
- `query_positions` feeds the tail only; the expansion is masked by
- `sequence_lengths` alone, so the causal bound on the *selection* is the
- producer's. Returns a fresh contiguous int32 `[rows, final_topk]`; inputs
- are read-only.
-
- Fusion boundary: index arithmetic and the tail only. Scoring, Top-K, the
- per-row causal group bound and the geometry are the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- # ---- decode rows: the column-split launch below 128 rows, whole rows from 128
- Cell(why="decode: one row, columns split across programs", spec=_spec(("sweep", 1))),
- # ---- 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; 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))),
- ),
- )
-
- note = """
- Metadata consumed: none -- every input is an argument. The target passes
- `qsa_topk_indices[:rows]` (int32, written by `qsa_topk_select`),
- `qsa_logical_positions` (int64) and `qsa_sequence_lengths` (int32: the
- request's whole KV length broadcast over its rows, so inside a context
- chunk it exceeds `position + 1` on every row but the last).
-
- Preconditions that are values, so not `is_valid` guards (both 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
- 12..15 lost. `qsa_topk_select` emits packed rows.
- - selected groups are complete and visible (`g < (pos + 1) // cr`). An
- incomplete group comes back with `-1` holes in place and its real
- tokens repeated in the tail; a group past the query's horizon but
- inside `sequence_lengths` is expanded verbatim.
- Duplicate group ids expand to duplicate tokens.
-
- 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. No autotune. CUDA-graph capturable (captured at 256 rows); a
- replay recomputes from the live buffers, and the captured row count fixes
- the launch shape.
-
- Hard limits (read from the kernel, not probed): the row offset is
- `program_id * stride(0)` in int32, so `rows * stride(0) >= 2^31` (4 Mi
- rows at width 512) would wrap. Certified extent: 8192 rows x 2051 columns.
-
- Not certified: multiple devices, `torch.compile`, rows above 8192.
- """
-
- def __call__(
- self,
- block_indices: torch.Tensor,
- query_positions: torch.Tensor,
- sequence_lengths: torch.Tensor,
- *,
- compress_ratio: int,
- token_topk: int,
- ) -> torch.Tensor:
- return _qsa_kernels.triton_expand_qsa_block_indices(
- block_indices,
- query_positions,
- sequence_lengths,
- compress_ratio=compress_ratio,
- token_topk=token_topk,
- )
-
- def reference(
- self,
- block_indices: torch.Tensor,
- query_positions: torch.Tensor,
- sequence_lengths: torch.Tensor,
- *,
- compress_ratio: int,
- token_topk: int,
- ) -> torch.Tensor:
- """The expansion as the contract states it, in int64 torch."""
- rows, block_topk = block_indices.shape
- final_topk = token_topk + compress_ratio - 1
- device = block_indices.device
- blocks = block_indices.long()
- lengths = sequence_lengths.long()
- offsets = torch.arange(compress_ratio, device=device)
-
- expanded = (blocks[:, :, None] * compress_ratio + offsets).reshape(rows, token_topk)
- chosen = (blocks >= 0).repeat_interleave(compress_ratio, dim=1)
- keep = chosen & (expanded < lengths[:, None])
- out = torch.full((rows, final_topk), -1, dtype=torch.int64, device=device)
- out[:, :token_topk] = torch.where(keep, expanded, -1)
-
- first_free = (blocks >= 0).sum(dim=1) * compress_ratio
- visible = query_positions.long() + 1
- tail_start = visible // compress_ratio * compress_ratio
- row = torch.arange(rows, device=device)
- for offset in range(compress_ratio - 1):
- column = first_free + offset
- token = tail_start + offset
- live = (offset < visible - tail_start) & (token < lengths) & (column < final_topk)
- out[row[live], column[live]] = token[live]
- return out.to(torch.int32)
-
-
-qsa_expand_block_indices = QsaExpandBlockIndices()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_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
deleted file mode 100644
index 582e514682ea..000000000000
--- 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
+++ /dev/null
@@ -1,246 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""QSA MTP IndexShare: draft-decode selections rebuilt from the row captured at draft-extend."""
-
-from __future__ import annotations
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-# Tier-2 sourcing: a vendored Triton callable reached as ordinary Python. QSA's
-# kernel plane has no torch.ops surface.
-from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels as _qsa_kernels
-
-#: This checkpoint's selection width: `indexer_budget` 2048 tokens over
-#: `indexer_compress_ratio` 4.
-_BLOCK_TOPK = 512
-#: `(slots, this layer's slot)` of the shared buffers: with the MTP layer in the
-#: target's cache manager, 12 trunk full-attention layers plus the MTP layer.
-_COMBINED_SLOTS = (13, 12)
-
-
-def _spec(contexts, **overrides) -> dict:
- """A cell spec at this checkpoint's geometry and the engine's metadata dtypes.
-
- `contexts` says what each request's captured row and current row are:
- ("mixed", B) B requests; the first 16 at fixed boundary positions
- (empty, short, 511, 512 and full rows), the rest seeded
- ("full", B) captured positions >= 2047: all 512 slots valid
- ("explicit", P) one request per captured position in P
- A request captured at position `p` (its last accepted row) holds
- `(p + 1) // 4` complete groups; its draft row `d` positions later sees
- `(p + d + 1) // 4`. `drafted` cycles over the requests.
- """
- base = dict(
- contexts=contexts,
- drafted=(1, 2, 3),
- rows_per_request=1,
- shared_rows=64,
- slots=_COMBINED_SLOTS,
- block_topk=_BLOCK_TOPK,
- metadata_dtype=torch.int32,
- graph_batches=(),
- )
- return {**base, **overrides}
-
-
-class QsaMtpReuseTopk(OpWrapper):
- """Rebuild each query row's group selection from its request's captured row,
- with the groups that completed since the capture written into it.
-
- Wraps `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels.
- triton_qsa_mtp_reuse_topk`: one Triton launch, one program per output row.
- For row `r`, with `q = request_indices[r]` and `K = output.shape[1]`
- (`block_topk`):
-
- 1. start from the captured row `shared_indices[q]`: valid group ids as a
- prefix, then `-1`;
- 2. `captured = shared_visible_blocks[q]`, `current = visible_blocks[r]`,
- `fresh = clamp(current - captured, 0, K)`: the groups
- `captured .. captured + fresh - 1` completed after the capture;
- 3. `filled` = the count of non-negative ids in the captured row; the
- fresh ids go, ascending, into slots `start .. start + fresh - 1` with
- `start = min(filled, K - fresh)` -- right after the prefix when there
- is room, otherwise after the captured row's newest `K - fresh` ids:
- the oldest `filled + fresh - K` ids are displaced and the survivors
- shift to the front;
- 4. every other slot is the captured row's survivors in their order, `-1`
- padding included.
-
- Every slot of `output[:rows]` is written and nothing else; returns `output`.
- The inputs are read-only. Nothing is sorted, compacted or de-duplicated and
- the ranking is not recomputed: the captured row is taken as ascending (the
- target makes it canonical), so the displaced ids are the oldest groups, not the
- lowest-scored ones, and the result is ascending again.
-
- Fusion boundary: the copy-and-append only. Capturing each request's row
- into the shared buffers, the per-row request map and visible counts, and
- the expansion into token indices plus the incomplete causal tail
- (`qsa_expand_block_indices`) are the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- # ---- decode batch: one row per request; the first 16 requests sit at the
- # boundaries (empty, short, 511-id, 512-id and full rows), fresh 0 and 1.
- Cell(
- why="decode MTP: 32 requests (envelope), 256-row shared buffer",
- spec=_spec(("mixed", 32), shared_rows=256),
- ),
- # ---- the riskiest edge: two fresh groups displace a full row's two oldest
- Cell(
- why="8 requests at fresh 2: full, 511-id and short rows gain two groups",
- spec=_spec(
- ("explicit", (2050, 402, 2046, 8186, 6, 2047, 1022, 2049)),
- drafted=(5, 6, 5, 6, 6, 4, 5, 6),
- shared_rows=8,
- ),
- ),
- # ---- CUDA graph -----------------------------------------------------
- Cell(
- why="CUDA graph, drafted 1..6: captured at 1, 8 and 32 requests, new-data replay",
- spec=_spec(("mixed", 32), drafted=(1, 2, 3, 4, 5, 6), graph_batches=(1, 8, 32)),
- ),
- )
-
- note = """
- Metadata consumed: none; every input is an argument, all int32 at the call
- site:
- shared_indices `qsa_shared_topk_indices[slot]`, [max_num_sequences, block_topk]
- shared_visible_blocks `qsa_shared_topk_visible_blocks[slot]`, [max_num_sequences]
- request_indices `qsa_req_idx_per_token[:rows]`
- visible_blocks `qsa_visible_blocks[:rows]`
- output `qsa_topk_indices[:rows]`, the [max_num_tokens, block_topk] buffer
- `slot` is the layer's row in the cache manager's `qsa_shared_topk_slots`.
-
- It runs in the MTP draft loop's decode steps after the first: no context
- rows, one row per request, so rows = requests, and the request map is the
- row index. Draft step `i` sits `i` positions past the captured row, so
- `max_draft_len` 1 never reaches this op, 3 reaches 1 and 2 positions, and 7
- (the target's ceiling) reaches 1..6. With compress_ratio 4, `fresh =
- ((p + 1) % 4 + i) // 4` for a capture at position `p`: 0 or 1 up to i = 3
- and 0, 1 or 2 from i = 4. In a full row
- the fresh ids displace the row's oldest ids (one at fresh 1, two at fresh
- 2), and a 511-id row at fresh 2 loses its oldest id.
-
- Preconditions that are values, so not `is_valid` guards:
- - each captured row is prefix-packed (every valid id before every -1), as
- `qsa_topk_select` emits it. The append slot comes from the COUNT of
- non-negative ids, so a -1 hole puts a fresh id over a valid one and
- loses a group silently;
- - captured ids are below that row's `shared_visible_blocks`, so a fresh id
- never duplicates a captured one;
- - every `request_indices[r]` is a row of both shared tensors;
- - `shared_visible_blocks[q]` is the visible count of the very row that was
- captured, and `visible_blocks[r]` is `(position + 1) // cr` of row `r`.
-
- Loud rejections (not pinned by a test): `ValueError` when the shared width differs from the output
- width, when `request_indices` or `visible_blocks` does not hold one element
- per output row, or on CPU tensors; a 1-D `output` or `shared_indices`; a
- width that is not a power of two (Triton compilation).
-
- Accepted and silently wrong, hence `is_valid`: any non-contiguous operand
- (the kernel addresses raw pointers densely and ignores strides); a 3-D
- `shared_indices` whose middle extent equals `block_topk` (read as slot 0);
- a 2-D `shared_visible_blocks` (the whole `[slots, seqs]` buffer is read as
- slot 0's counts). The test pins the last.
-
- Certified: int32 ids and metadata, `block_topk` 512, up to 32 requests, the
- identity request map over the target cache manager's 13-slot buffers, CUDA
- graph capture at 1, 8 and 32 requests with a new-data replay.
- Not certified: a permuted request map, a separate draft manager's one-slot
- buffers.
- """
-
- def __call__(
- self,
- *,
- shared_indices: torch.Tensor,
- shared_visible_blocks: torch.Tensor,
- request_indices: torch.Tensor,
- visible_blocks: torch.Tensor,
- output: torch.Tensor,
- ) -> torch.Tensor:
- return _qsa_kernels.triton_qsa_mtp_reuse_topk(
- shared_indices=shared_indices,
- shared_visible_blocks=shared_visible_blocks,
- request_indices=request_indices,
- visible_blocks=visible_blocks,
- output=output,
- displace_oldest=True,
- )
-
- def reference(
- self,
- *,
- shared_indices: torch.Tensor,
- shared_visible_blocks: torch.Tensor,
- request_indices: torch.Tensor,
- visible_blocks: torch.Tensor,
- output: torch.Tensor,
- ) -> torch.Tensor:
- """The rebuilt rows as the contract states them, in int64 torch.
-
- `output` supplies only the shape and dtype; its contents are ignored.
- """
- rows, block_topk = output.shape
- device = output.device
- requests = request_indices.long().reshape(-1)
- captured = shared_visible_blocks.long()[requests]
- current = visible_blocks.long().reshape(-1)
- rebuilt = shared_indices.long()[requests].clone()
-
- filled = (rebuilt >= 0).sum(dim=1)
- fresh = (current - captured).clamp(min=0, max=block_topk)
- start = torch.minimum(filled, block_topk - fresh)
- slot = torch.arange(block_topk, device=device)[None, :]
- # The oldest `filled - start` ids go; the survivors move to the front.
- source = slot + (filled - start)[:, None]
- rebuilt = torch.where(
- source < block_topk,
- rebuilt.gather(1, source.clamp(max=block_topk - 1)),
- torch.full_like(rebuilt, -1),
- )
- rank = slot - start[:, None]
- appended = (rank >= 0) & (rank < fresh[:, None])
- rebuilt[appended] = (captured[:, None] + rank)[appended]
- return rebuilt.to(output.dtype)
-
- def is_valid(
- self,
- *,
- shared_indices: torch.Tensor,
- shared_visible_blocks: torch.Tensor,
- request_indices: torch.Tensor,
- visible_blocks: torch.Tensor,
- output: torch.Tensor,
- ) -> None:
- # The kernel takes raw pointers and addresses them densely; each case
- # below returns a plausible wrong answer rather than an error.
- if shared_indices.dim() != 2 or not shared_indices.is_contiguous():
- raise ValueError(
- "shared_indices must be one slot's contiguous [num_seqs, block_topk] "
- f"rows; got shape {tuple(shared_indices.shape)}, "
- f"stride {shared_indices.stride()}"
- )
- if shared_visible_blocks.dim() != 1 or not shared_visible_blocks.is_contiguous():
- raise ValueError(
- "shared_visible_blocks must be one slot's contiguous [num_seqs] counts; "
- f"got shape {tuple(shared_visible_blocks.shape)}, "
- f"stride {shared_visible_blocks.stride()}"
- )
- for name, tensor in (
- ("request_indices", request_indices),
- ("visible_blocks", visible_blocks),
- ("output", output),
- ):
- if not tensor.is_contiguous():
- raise ValueError(
- f"{name} must be contiguous; got shape {tuple(tensor.shape)}, "
- f"stride {tensor.stride()}"
- )
-
-
-qsa_mtp_reuse_topk = QsaMtpReuseTopk()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_nvfp4_pages_to_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_nvfp4_pages_to_bf16.py
deleted file mode 100644
index 03f261bf15ca..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_nvfp4_pages_to_bf16.py
+++ /dev/null
@@ -1,390 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""A batch's pages of an NVFP4 paged K/V pool, dequantized into a bf16 scratch pool, and the
-block table that addresses them there.
-
-Tier 5 (Triton written for this target). The target's context launches over an NVFP4 pool
-read their requests' pages through this scratch with the bf16 row-tile kernel.
-"""
-
-from __future__ import annotations
-
-import math
-
-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_sparse_gqa_nvfp4 import _e2m1_to_f32
-
-#: 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_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
-
-#: This checkpoint's K/V geometry and the engine's QSA table (`width` pages a slot);
-#: `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(
- kv_heads=2,
- head_dim=256,
- tokens_per_block=32,
- width=320,
- slots=40,
- global_k=1.0,
- global_v=1.0,
- slot_pages=2,
- layer=0,
- holes=False,
-)
-
-
-def _spec(kv_lens, **overrides) -> dict:
- """`kv_lens` is one KV length per listed request, in order."""
- return {**_CHECKPOINT, "kv_lens": tuple(kv_lens), **overrides}
-
-
-@triton.jit
-def _dequantize_head_page(
- data,
- scales,
- page,
- head,
- global_scale,
- valid,
- d_stride_page,
- d_stride_head,
- d_stride_token,
- s_stride_page,
- s_stride_head,
- GROUPED: tl.constexpr,
- TOKENS_PER_BLOCK: tl.constexpr,
- HEAD_DIM: tl.constexpr,
-):
- """[tpb, D] fp32 values of one head page: e2m1 x e4m3 block scale x second-level
- scale, in the sparse GQA kernels' order."""
- tokens = tl.arange(0, TOKENS_PER_BLOCK)[:, None]
- dims = tl.arange(0, HEAD_DIM)[None, :]
- packed = tl.load(
- data + page * d_stride_page + head * d_stride_head + tokens * d_stride_token + dims // 2,
- mask=valid,
- other=0,
- ).to(tl.int32)
- codes = (packed >> ((dims % 2) * 4)) & 0xF
- if GROUPED:
- # V: `[tpb / 4, D / 16, 4]` per head page.
- offsets = (tokens // 4) * (HEAD_DIM // 4) + (dims // 16) * 4 + tokens % 4
- else:
- # K: token-major `[tpb, D / 16]` per head page.
- offsets = tokens * (HEAD_DIM // 16) + dims // 16
- block_scales = (
- tl.load(scales + page * s_stride_page + head * s_stride_head + offsets, mask=valid, other=0)
- .to(tl.float8e4nv, bitcast=True)
- .to(tl.float32)
- )
- return _e2m1_to_f32(codes) * block_scales * tl.load(global_scale)
-
-
-@triton.jit
-def _nvfp4_pages_to_bf16_kernel(
- k_cache,
- v_cache,
- k_scales,
- v_scales,
- kv_scale_quant_orig,
- block_table,
- slots,
- counts,
- offsets,
- table,
- k_out,
- v_out,
- num_pool_pages,
- listed,
- d_stride_page,
- d_stride_head,
- d_stride_token,
- s_stride_page,
- s_stride_head,
- bt_stride,
- table_stride,
- o_stride_page,
- o_stride_head,
- o_stride_token,
- NUM_KV_HEADS: tl.constexpr,
- TOKENS_PER_BLOCK: tl.constexpr,
- HEAD_DIM: tl.constexpr,
-):
- request = tl.program_id(0)
- page = tl.program_id(1)
- slot = tl.load(slots + request).to(tl.int64)
- dst = tl.load(offsets + request).to(tl.int64) + page
- listed_page = (page < tl.load(counts + request)) & (dst < listed)
- src = tl.load(block_table + slot * bt_stride + page, mask=listed_page, other=-1).to(tl.int64)
- resident = listed_page & (src >= 0) & (src < num_pool_pages)
- safe = tl.where(resident, src, 0)
- tl.store(
- table + slot * table_stride + page,
- tl.where(resident, dst, -1).to(tl.int32),
- mask=listed_page,
- )
- tokens = tl.arange(0, TOKENS_PER_BLOCK)[:, None]
- dims = tl.arange(0, HEAD_DIM)[None, :]
- out_offsets = dst * o_stride_page + tokens * o_stride_token + dims
- for head in tl.static_range(NUM_KV_HEADS):
- k = _dequantize_head_page(
- k_cache,
- k_scales,
- safe,
- head,
- kv_scale_quant_orig + 1,
- resident,
- d_stride_page,
- d_stride_head,
- d_stride_token,
- s_stride_page,
- s_stride_head,
- False,
- TOKENS_PER_BLOCK,
- HEAD_DIM,
- )
- v = _dequantize_head_page(
- v_cache,
- v_scales,
- safe,
- head,
- kv_scale_quant_orig + 2,
- resident,
- d_stride_page,
- d_stride_head,
- d_stride_token,
- s_stride_page,
- s_stride_head,
- True,
- TOKENS_PER_BLOCK,
- HEAD_DIM,
- )
- out = out_offsets + head * o_stride_head
- tl.store(k_out + out, k.to(k_out.dtype.element_ty), mask=listed_page)
- tl.store(v_out + out, v.to(v_out.dtype.element_ty), mask=listed_page)
-
-
-def _dequantize_pages(data, scales, pages, global_scale, grouped: bool) -> torch.Tensor:
- """[n, Hkv, tpb, D] fp32 values of `pages` (in the pool) of one role plane."""
- heads, tokens_per_block, packed = data.shape[1], data.shape[2], data.shape[3]
- head_dim, groups = packed * _PER_BYTE, packed * _PER_BYTE // _SF_VEC
- raw = data[pages].long() # [n, H, tpb, 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)
- flat = scales[pages].reshape(pages.shape[0], heads, tokens_per_block * groups)
- tokens = torch.arange(tokens_per_block, device=data.device)[:, None]
- dims = torch.arange(head_dim, device=data.device)[None, :]
- if grouped:
- columns = (tokens // _V_SF_QUAD) * (_V_SF_QUAD * groups) + (dims // _SF_VEC) * _V_SF_QUAD
- columns = columns + tokens % _V_SF_QUAD
- else:
- columns = tokens * groups + dims // _SF_VEC
- sf = flat[:, :, columns.flatten()].reshape(pages.shape[0], heads, tokens_per_block, head_dim)
- return values * sf.view(torch.float8_e4m3fn).float() * global_scale
-
-
-class QsaNvfp4PagesToBf16(OpWrapper):
- """K and V of listed requests' pages of an NVFP4 HND paged pool, as bf16 scratch pools,
- and the block table that addresses them there.
-
- Request `r` lists its first `counts[r]` table pages, `block_table[slots[r], j]` for `j
- < counts[r]`, at scratch pages `offsets[r] + j` (those below `listed`, the scratch's
- page count)::
-
- k_out[offsets[r] + j, h, t, d] = bf16(e2m1(K code at (s, h, t, d)) x
- e4m3(K scale of (t, d // 16)) x g_k)
-
- with `s = block_table[slots[r], j]` and `g_k = kv_scale_quant_orig[1]`, `v_out` likewise
- from the V planes and `g_v = kv_scale_quant_orig[2]`. K's block scales are token-major
- `[tpb, D/16]` per head page, V's in four-token groups `[tpb/4, D/16, 4]`. The product is
- formed in fp32 in the sparse GQA kernels' order and rounded once to bf16, so a bf16
- kernel over the scratch reads the values those kernels dequantize. `table[slots[r], j]`
- becomes `offsets[r] + j`, or -1 where the pool's entry is unmapped or out of the pool
- (whose scratch page is then zeros); the call writes no other entry of `table`, which the
- caller fills with -1. Returns `(k_out, v_out, table)`.
-
- Fusion boundary: the listing's dequantize and table; choosing the requests and their
- page counts (and `offsets`, their exclusive prefix sum) is the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(why="one request, one page", spec=_spec((17,))),
- Cell(why="con 1 prefill: one 8192-token request", spec=_spec((8192,))),
- Cell(why="max_num_tokens: two 8192-token requests", spec=_spec((8192, 8192))),
- Cell(
- why="three requests, a chunked continuation and one filling the table",
- spec=_spec((5000 + 3192, 3000, 10240)),
- ),
- Cell(why="many short requests", spec=_spec(tuple(33 + 61 * i for i in range(30)))),
- Cell(why="holes: unmapped and out-of-pool entries", spec=_spec((4000, 2000), holes=True)),
- Cell(
- why="second-level scales g_k 0.37, g_v 2.5: the product rounds to bf16",
- spec=_spec((1500, 700), global_k=0.37, global_v=2.5),
- ),
- Cell(
- why="the engine's trunk layout: layer 11's planes among the 24 role pages",
- spec=_spec((3000, 900), slot_pages=24, layer=11),
- ),
- )
-
- note = """
- Metadata consumed: none -- every input is an argument. Call site: the target's
- context/mixed forward over an NVFP4 pool, once per attention layer and group of the
- step's context requests (`_nvfp4_context_plan`: slots from the context rows' request
- indices, counts `ceil(kv_len / tpb)` from the device KV lengths, the groups and the
- scratch's pages from the host ones); the row-tile kernel then reads the scratch
- through `table`.
-
- One Triton kernel, a program per (request, table page), both K/V heads and both
- roles: a head page is 32 tokens x 256 dims, 8 warps. Programs past a request's
- count store nothing.
-
- Preconditions: uint8 planes `[pages, Hkv, tpb, D/2]` (data) and `[pages, Hkv, tpb,
- D/16]` (scales), K's and V's at equal strides, unit last-dim stride and contiguous
- head pages (the QSA manager's role planes); `kv_scale_quant_orig` an fp32 [3];
- `table` an int32 tensor of `block_table`'s shape. Bitwise deterministic.
- """
-
- def __call__(
- self,
- 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,
- slots: torch.Tensor,
- counts: torch.Tensor,
- offsets: torch.Tensor,
- table: torch.Tensor,
- k_out: torch.Tensor,
- v_out: torch.Tensor,
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
- requests, listed = slots.shape[0], k_out.shape[0]
- if requests and listed:
- _nvfp4_pages_to_bf16_kernel[(requests, block_table.shape[1])](
- k_cache,
- v_cache,
- k_scale_cache,
- v_scale_cache,
- kv_scale_quant_orig,
- block_table,
- slots,
- counts,
- offsets,
- table,
- k_out,
- v_out,
- k_cache.shape[0],
- listed,
- k_cache.stride(0),
- k_cache.stride(1),
- k_cache.stride(2),
- k_scale_cache.stride(0),
- k_scale_cache.stride(1),
- block_table.stride(0),
- table.stride(0),
- k_out.stride(0),
- k_out.stride(1),
- k_out.stride(2),
- NUM_KV_HEADS=k_cache.shape[1],
- TOKENS_PER_BLOCK=k_cache.shape[2],
- HEAD_DIM=k_cache.shape[3] * _PER_BYTE,
- num_warps=8,
- )
- return k_out, v_out, table
-
- def reference(
- self,
- 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,
- slots: torch.Tensor,
- counts: torch.Tensor,
- offsets: torch.Tensor,
- table: torch.Tensor,
- k_out: torch.Tensor,
- v_out: torch.Tensor,
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
- k_ref, v_ref, table_ref = k_out.clone(), v_out.clone(), table.clone()
- listed = k_out.shape[0]
- for r in range(slots.shape[0]):
- slot, base = int(slots[r]), int(offsets[r])
- for j in range(min(int(counts[r]), block_table.shape[1])):
- dst = base + j
- if dst >= listed:
- break
- src = int(block_table[slot, j])
- resident = 0 <= src < k_cache.shape[0]
- table_ref[slot, j] = dst if resident else -1
- page = torch.tensor([src if resident else 0], device=k_cache.device)
- for out, data, sc, g, grouped in (
- (k_ref, k_cache, k_scale_cache, kv_scale_quant_orig[1], False),
- (v_ref, v_cache, v_scale_cache, kv_scale_quant_orig[2], True),
- ):
- vals = _dequantize_pages(data, sc, page, g, grouped)[0]
- out[dst] = (vals if resident else torch.zeros_like(vals)).to(out.dtype)
- return k_ref, v_ref, table_ref
-
- def is_valid(
- self,
- 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,
- slots: torch.Tensor,
- counts: torch.Tensor,
- offsets: torch.Tensor,
- table: torch.Tensor,
- k_out: torch.Tensor,
- v_out: torch.Tensor,
- ) -> None:
- pages, heads, tokens_per_block, packed = k_cache.shape
- head_dim = packed * _PER_BYTE
- assert tokens_per_block % _V_SF_QUAD == 0, "tokens_per_block: a multiple of 4"
- assert head_dim % _SF_VEC == 0 and math.log2(head_dim).is_integer(), "head_dim"
- for name, plane, width in (
- ("k_cache", k_cache, packed),
- ("v_cache", v_cache, packed),
- ("k_scale_cache", k_scale_cache, head_dim // _SF_VEC),
- ("v_scale_cache", v_scale_cache, head_dim // _SF_VEC),
- ):
- assert plane.dtype == torch.uint8, f"{name}: uint8"
- assert tuple(plane.shape) == (pages, heads, tokens_per_block, width), f"{name}: shape"
- assert plane.stride(3) == 1, f"{name}: unit last-dim stride"
- assert v_cache.stride() == k_cache.stride(), "K and V data planes at equal strides"
- assert v_scale_cache.stride() == k_scale_cache.stride(), "K and V scales at equal strides"
- assert k_scale_cache.stride(2) == head_dim // _SF_VEC, "scales: contiguous head pages"
- assert kv_scale_quant_orig.dtype == torch.float32 and kv_scale_quant_orig.shape == (3,)
- assert block_table.dim() == 2 and block_table.dtype == torch.int32
- assert block_table.stride(1) == 1, "block_table: row-contiguous"
- assert table.shape == block_table.shape and table.dtype == torch.int32
- assert table.stride(1) == 1, "table: row-contiguous"
- for name, t in (("slots", slots), ("counts", counts), ("offsets", offsets)):
- assert t.dim() == 1 and t.shape == slots.shape, f"{name}: one per request"
- assert t.dtype in (torch.int32, torch.int64), f"{name}: integer"
- for out in (k_out, v_out):
- assert out.dtype == torch.bfloat16, "out: bf16"
- assert out.shape[1:] == (heads, tokens_per_block, head_dim), (
- "out: [listed, Hkv, tpb, D]"
- )
- assert out.is_contiguous(), "out: contiguous"
-
-
-qsa_nvfp4_pages_to_bf16 = QsaNvfp4PagesToBf16()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_nvfp4_verify_ptx.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_nvfp4_verify_ptx.py
deleted file mode 100644
index d6b416ac80b1..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_nvfp4_verify_ptx.py
+++ /dev/null
@@ -1,248 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Inline-PTX pieces of `qsa_paged_sparse_gqa_nvfp4_verify`'s kernel beyond the ones the
-per-row kernel (`qsa_paged_sparse_gqa_nvfp4_decode`) already has: the union scan's bit
-operations, vector global loads and stores, the tcgen05 thread-sync fences, and tcgen05.mma issue with
-precomputed descriptors.
-
-`gemm_ts` issues every K step of a GEMM whose A operand sits in tensor memory from one
-inline-PTX block, the shared-memory descriptor of B advanced by immediates; the encodings are
-those of CUTLASS's `cute/arch/mma_sm100_desc.hpp`. `cute.gemm` rebuilds the descriptors per K
-step on the uniform datapath (about 18 instructions a UMMA), which the verify kernel's single
-MMA warp cannot hide.
-"""
-
-from __future__ import annotations
-
-import cutlass
-import cutlass.cute as cute
-from cutlass import Int32
-from cutlass._mlir.dialects import llvm
-from cutlass._mlir.dialects import nvvm as _nvvm
-from cutlass.cute.nvgpu import tcgen05
-from cutlass.cutlass_dsl import T, dsl_user_op
-
-from .qsa_paged_sparse_gqa_nvfp4_decode import _asm, _ir
-
-# ---------------------------------------------------------------------------------------
-# scalar and memory helpers
-# ---------------------------------------------------------------------------------------
-
-
-@dsl_user_op
-def popc(x, *, loc=None, ip=None):
- return cutlass.Int32(
- _asm(T.i32(), [_ir(x, cutlass.Int32, loc, ip)], "popc.b32 $0, $1;", "=r,r", loc, ip)
- )
-
-
-@dsl_user_op
-def ctz(x, *, loc=None, ip=None):
- """Trailing zeros of a nonzero word."""
- return cutlass.Int32(
- _asm(
- T.i32(),
- [_ir(x, cutlass.Int32, loc, ip)],
- "{ .reg .b32 t; brev.b32 t, $1; clz.b32 $0, t; }",
- "=r,r",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def red_or_shared(saddr, v, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [_ir(saddr, cutlass.Int32, loc, ip), _ir(v, cutlass.Int32, loc, ip)],
- "{ red.shared.or.b32 [$1], $2; mov.u32 $0, 0; }",
- "=r,r,r",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def stg4(gaddr, a, b, c, d, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [_ir(gaddr, cutlass.Int64, loc, ip)]
- + [_ir(v, cutlass.Int32, loc, ip) 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=True,
- )
-
-
-@dsl_user_op
-def stg4f(gaddr, a, b, c, d, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [_ir(gaddr, cutlass.Int64, loc, ip)]
- + [_ir(v, cutlass.Float32, loc, ip) for v in (a, b, c, d)],
- "{ st.global.v4.f32 [$1], {$2, $3, $4, $5}; mov.u32 $0, 0; }",
- "=r,l,f,f,f,f",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def ldg2(gaddr, *, loc=None, ip=None):
- res = _asm(
- llvm.StructType.get_literal([T.i32()] * 2),
- [_ir(gaddr, cutlass.Int64, loc, ip)],
- "ld.global.v2.b32 {$0, $1}, [$2];",
- "=r,=r,l",
- loc,
- ip,
- side=True,
- )
- return [cutlass.Int32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(2)]
-
-
-@dsl_user_op
-def stg2(gaddr, a, b, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [
- _ir(gaddr, cutlass.Int64, loc, ip),
- _ir(a, cutlass.Int32, loc, ip),
- _ir(b, cutlass.Int32, loc, ip),
- ],
- "{ st.global.v2.b32 [$1], {$2, $3}; mov.u32 $0, 0; }",
- "=r,l,r,r",
- loc,
- ip,
- side=True,
- )
-
-
-def groups8(x):
- """Bit i of the result: any of bits 4i .. 4i + 3 of `x` (the eight four-token groups of a
- 32-token bitmap word)."""
- y = x | (x >> 1)
- y = y | (y >> 2)
- z = y & 1
- for i in range(1, 8):
- z = z | (((y >> (4 * i)) & 1) << i)
- return z
-
-
-def spread4(nib):
- """Bits 0..3 of `nib` to bits 0, 8, 16, 24."""
- return (nib & 1) | ((nib & 2) << 7) | ((nib & 4) << 14) | ((nib & 8) << 21)
-
-
-@dsl_user_op
-def tcgen05_fence_before(*, loc=None, ip=None):
- _nvvm.tcgen05_fence(_nvvm.Tcgen05FenceKind.BEFORE_THREAD_SYNC, loc=loc, ip=ip)
-
-
-@dsl_user_op
-def tcgen05_fence_after(*, loc=None, ip=None):
- _nvvm.tcgen05_fence(_nvvm.Tcgen05FenceKind.AFTER_THREAD_SYNC, loc=loc, ip=ip)
-
-
-# ---------------------------------------------------------------------------------------
-# tcgen05.mma issue with precomputed descriptors
-# ---------------------------------------------------------------------------------------
-
-#: Swizzle: B -> the descriptor's layout type (none, 32B, 64B, 128B).
-_SW_TYPES = {0: 0, 1: 6, 2: 4, 3: 2}
-#: Layout type -> swizzle atom size along MN (16-byte units).
-_SW_ATOM_MN = {0: 1, 6: 2, 4: 4, 2: 8}
-
-
-def instr_desc(op) -> int:
- """The 32-bit instruction descriptor of an f16/bf16 x f16/bf16 -> f32 tcgen05 MMA op."""
- fmt = {cutlass.Float16: 0, cutlass.BFloat16: 1}
- a_major = 0 if op.a_major_mode == tcgen05.OperandMajorMode.K else 1
- b_major = 0 if op.b_major_mode == tcgen05.OperandMajorMode.K else 1
- m, n = op.shape_mnk[0], op.shape_mnk[1]
- assert op.acc_dtype is cutlass.Float32 and m in (64, 128) and n % 8 == 0 and 8 <= n <= 256
- d = 0
- d |= 1 << 4 # c_format f32
- d |= fmt[op.a_dtype] << 7
- d |= fmt[op.b_dtype] << 10
- d |= a_major << 15
- d |= b_major << 16
- d |= (n >> 3) << 17
- d |= (m >> 4) << 24
- return d
-
-
-def smem_desc_base(layout_u128, swizzle, k_major: bool) -> int:
- """The 64-bit shared-memory descriptor without its start address, for an operand whose
- (MN, K) layout `layout_u128` is in 16-byte units."""
- assert swizzle.num_base == 4 and swizzle.num_shift == 3
- ltype = _SW_TYPES[swizzle.num_bits]
- atom_mn = _SW_ATOM_MN[ltype]
- if not k_major:
- canon = cute.logical_divide(layout_u128, (atom_mn, 8))
- s00 = canon.stride[0][0]
- s10 = canon.stride[1][0]
- assert (ltype == 0 or s00 == 1) and s10 == atom_mn, "not a canonical UMMA MN layout"
- s01, s11 = canon.stride[0][1], canon.stride[1][1]
- sbo, lbo = (s01, s11) if ltype == 0 else (s11, s01)
- else:
- canon = cute.logical_divide(layout_u128, (8, 2))
- s00 = canon.stride[0][0]
- s10 = canon.stride[1][0]
- assert s00 == atom_mn and (ltype == 0 or s10 == 1), "not a canonical UMMA K layout"
- sbo, lbo = canon.stride[0][1], s10
- d = 0
- d |= (lbo & 0x3FFF) << 16
- d |= (sbo & 0x3FFF) << 32
- d |= 1 << 46 # descriptor version
- d |= (ltype & 0x7) << 61
- return d
-
-
-def gemm_ts(op, acc_tmem, a_tmem, tCrA, tCrB, sB, zero_init: bool):
- """acc (+)= A B over every K step: A in tensor memory at `a_tmem` (fragment `tCrA`), B the
- shared-memory operand `sB` (fragment `tCrB`, one stage). Issued by one elected thread of
- the calling warp. Plain Python run while the kernel is traced: the K loop unrolls into one
- PTX block."""
- idesc = instr_desc(op)
- k_major_b = op.b_major_mode == tcgen05.OperandMajorMode.K
- base_b = smem_desc_base(
- cute.recast_layout(128, op.b_dtype.width, sB.layout[0]),
- sB.iterator.type.swizzle_type,
- k_major_b,
- )
- b_lo = base_b & 0xFFFFFFFF
- b_hi = (base_b >> 32) & 0xFFFFFFFF
- a_layout = cute.recast_layout(32, tCrA.element_type.width, tCrA.layout)
- n_k = cute.size(tCrA.shape[2])
- off_a = [cute.crd2idx((0, 0, k), a_layout) for k in range(n_k)]
- off_b = [cute.crd2idx((0, 0, k), tCrB.layout) for k in range(n_k)]
- start_b = Int32(b_lo) | Int32((sB[None, None, 0].iterator.toint() & 0x3FFFF) >> 4)
- body = "{\n\t.reg .pred leader;\n\t.reg .b32 db_lo;\n\t.reg .b64 db;\n\telect.sync _|leader, -1;\n\t"
- for k in range(n_k):
- acc = "0" if (zero_init and k == 0) else "1"
- body += (
- f"add.u32 db_lo, $2, {hex(off_b[k])};\n\t"
- f"mov.b64 db, {{db_lo, {hex(b_hi)}}};\n\t"
- f"@leader tcgen05.mma.cta_group::1.kind::f16 [$0], [$1 + {hex(off_a[k])}], db, "
- f"{hex(idesc)}, {acc};\n\t"
- )
- body += "}\n"
- llvm.inline_asm(
- None,
- [
- Int32(cute.arch.make_warp_uniform(acc_tmem)).ir_value(),
- Int32(cute.arch.make_warp_uniform(a_tmem)).ir_value(),
- Int32(cute.arch.make_warp_uniform(start_b)).ir_value(),
- ],
- body,
- "r,r,r",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- )
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py
deleted file mode 100644
index fff162545bc7..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py
+++ /dev/null
@@ -1,245 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""QSA compressed-block index scoring, read straight out of the paged index-K side cache."""
-
-from __future__ import annotations
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-# Approved exception: QSA's selection path is Triton called as Python, with no
-# torch.ops boundary. The engine's own sparse hook (`select_qsa_paged_tokens`
-# in `_torch/attention/backends/sparse/qsa/module.py`) calls this same callable.
-from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import (
- triton_qsa_paged_index_scores as _triton_qsa_paged_index_scores,
-)
-
-_CHECKPOINT = dict(
- heads=4,
- head_dim=128,
- tokens_per_block=32,
- compress_ratio=4,
- dtype=torch.bfloat16,
- # `qsa_logical_positions` is int64, `qsa_req_idx_per_token` int32.
- positions_dtype=torch.int64,
- request_dtype=torch.int32,
- # The target's q is a head slice of its packed [T, 4 + 1 + 1, 128] index
- # scratch buffer, so its row stride is 768 rather than 512.
- q_layout="scratch_slice",
-)
-
-
-def _spec(**overrides) -> dict:
- """A cell spec: this checkpoint's geometry unless the cell says otherwise."""
- return {**_CHECKPOINT, **overrides}
-
-
-#: The serving runtime dumps: an [8, 260] table on the warm-up capture path and
-#: a [32, 260] one under the mmmu gate's mixed steps.
-_WIDTH = 260
-
-
-#: The same plane under an NVFP4 KV pool: the QSA position page joins each physical
-#: page after the index buffers (`page_pad` elements), so the trunk's view is
-#: `[pages, 32, 1, 128]` at page stride 12 x 4096 + 192 = 49344.
-_TRUNK_POOL_NVFP4 = dict(pool_buffers=12, page_pad=192)
-
-
-class QsaPagedIndexScores(OpWrapper):
- """Score every causally visible compressed block of each query row.
-
- Wraps `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels.
- triton_qsa_paged_index_scores` (one Triton launch after allocating the
- output). For row `r` and compressed block `b` of request
- `request_indices[r]`::
-
- anchor = b * compress_ratio + compress_ratio - 1
- key_b = index_cache[block_table[req, anchor // tpb], anchor % tpb, 0]
- out[r, b] = sum_h relu(q[r, h] . key_b) * index_head_dim ** -0.5
-
- The block's key is whatever the compression step left at the group's LAST
- position; nothing is pooled here. Row `r` sees `(query_positions[r] + 1) //
- compress_ratio` blocks (complete groups only). Columns at or past that are
- `-inf` when `only_visible_blocks` is False and **not written** -- the fresh
- allocation's garbage -- when it is True. The column count is the block
- table's width, `block_table.shape[1] * tpb // compress_ratio`.
-
- `context_rows` is a layout declaration ("consecutive rows mostly share a
- request"), not a phase one: with more than 256 rows and an index-head count
- that divides 128 into more than one row (1, 2, 4, ..., 64) it selects a
- query-tiled kernel sharing one key gather across `128 // heads` rows;
- otherwise the row-wise kernel runs. Each cell names the kernel it was measured on. The
- two are not bitwise equal.
-
- Fusion boundary: the table walk, the gather, the dots, relu, head sum,
- scale and causal bound are inside. Producing `q` (norm, RoPE), writing and
- compressing the index cache, and everything after the scores are not.
- Returns a fresh fp32 `[rows, columns]`; nothing is written in place.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="NVFP4 trunk plane, page stride 49344: decode MTP 4 requests x 2 rows",
- spec=_spec(
- layout="decode",
- gen_positions=(7000, 5000, 4200, 2100),
- rows_per_request=2,
- slots=32,
- pages=_WIDTH,
- only_visible_blocks=True,
- context_rows=False,
- branch="row_wise",
- **_TRUNK_POOL_NVFP4,
- ),
- ),
- Cell(
- why="decode MTP7: 32 slots x 8 rows = 256 rows, the last decode-sized count",
- spec=_spec(
- layout="decode",
- gen_positions=tuple(100 + (499 * g) % 8000 for g in range(32)),
- rows_per_request=8,
- slots=32,
- pages=_WIDTH,
- only_visible_blocks=True,
- context_rows=False,
- branch="row_wise",
- ),
- ),
- Cell(
- why="capture: 8191 rows over 8 ragged requests, visible-only",
- spec=_spec(
- layout="capture",
- lengths=(1024, 512, 2048, 256, 1500, 700, 1151, 1000),
- slots=8,
- pages=_WIDTH,
- only_visible_blocks=True,
- context_rows=True,
- branch="query_tile",
- ),
- ),
- )
-
- note = """
- Metadata consumed: none -- every input is an argument. What the target
- passes: `block_table` is `QSAAttentionMetadata.qsa_block_table` (raw
- pool-page ids, int32, -1 past each sequence); `query_positions` /
- `request_indices` are `qsa_logical_positions` / `qsa_req_idx_per_token`;
- `index_cache` is the layer's bf16 index-K plane (`get_index_k_buffer`).
-
- Preconditions (not checked): every block-table entry the causal bound
- reaches is a valid page, and `request_indices[r]` is a row of the table --
- the loads are masked by the bound, not by the pool size, so a violation is
- an out-of-bounds read (not probed). The compression step must already have
- run: an uncompressed anchor scores as a raw index key. `rows == 0` returns
- an empty `[0, columns]` tensor without launching.
-
- The op validates its own metadata loudly (`ValueError`s; the test pins the
- one-entry-per-row metadata refusal), so `is_valid` adds nothing.
-
- `only_visible_blocks=True` really leaves garbage (measured, not pinned by a
- test): with a poisoned allocation freed just before the call, the whole
- unscored tail came back poisoned. At 8191 rows from position 0 over 2080 columns that is a full
- 32-column band (2048..2079) for every row, 8.65M cells. A consumer that
- scans `[row_start, row_end)` -- `qsa_topk_select` -- is unaffected.
-
- fp32 operands are accepted but `tl.dot` rounds them to a tensor-core
- format: 8.1e-3 max abs against an fp32 reference on scores to 7.8 (measured),
- outside the test's band. No cell covers fp32; the side cache is bf16.
-
- CUDA graphs: the test captures the decode launch (`use_pdl`, `prefetch`) at
- 4 rows and replays it after `q` and the anchor keys were rewritten in
- place; every replay matches an eager call bitwise.
-
- `prefetch` (with `use_pdl` and `only_visible_blocks`, at generation-sized
- batches) reads the pages and compressed keys before the PDL wait and q after
- it. Sound only behind the decode step's own prologue
- (`qwen4_exp_attn_decode_prologue`), which writes no compressed key but the last
- visible block's, at a row whose position ends a group; that key is read again
- after the wait.
-
- Every cell runs `only_visible_blocks=True` (the target never passes False),
- 4 index heads x 128, compress ratio 4, 32-token pages, bf16 operands, int64
- positions and int32 request indices. The NVFP4-pool page stride is
- certified on the row-wise kernel only.
-
- Not certified: more than 8191 rows, tables wider than 260 pages, more than
- 32 table rows, any other head geometry, compress ratio, page size or
- operand dtype, a context step's mixed request rows, and the `-inf` padded
- mode. Hard limits: not probed; the largest certified extent is 8191 x 2080
- output elements.
- """
-
- def __call__(
- self,
- q: torch.Tensor,
- index_cache: torch.Tensor,
- block_table: torch.Tensor,
- query_positions: torch.Tensor,
- request_indices: torch.Tensor,
- tokens_per_block: int,
- compress_ratio: int,
- only_visible_blocks: bool,
- context_rows: bool,
- use_pdl: bool = False,
- prefetch: bool = False,
- ) -> torch.Tensor:
- return _triton_qsa_paged_index_scores(
- q=q,
- index_cache=index_cache,
- block_table=block_table,
- query_positions=query_positions,
- request_indices=request_indices,
- tokens_per_block=tokens_per_block,
- compress_ratio=compress_ratio,
- only_visible_blocks=only_visible_blocks,
- context_rows=context_rows,
- use_pdl=use_pdl,
- prefetch=prefetch,
- )
-
- def reference(
- self,
- q: torch.Tensor,
- index_cache: torch.Tensor,
- block_table: torch.Tensor,
- query_positions: torch.Tensor,
- request_indices: torch.Tensor,
- tokens_per_block: int,
- compress_ratio: int,
- only_visible_blocks: bool,
- context_rows: bool,
- use_pdl: bool = False,
- prefetch: bool = False,
- ) -> torch.Tensor:
- """fp32 scores, `-inf` past each row's bound -- or NaN there, meaning
- "unspecified", when `only_visible_blocks` is True. `context_rows` is
- ignored: it picks a kernel, not a result.
-
- Walked in row chunks only to bound the gathered keys, which reach
- 8.7 GiB in fp32 at 8191 rows x 2080 columns.
- """
- rows, _, head_dim = q.shape
- num_columns = block_table.shape[1] * tokens_per_block // compress_ratio
- blocks = torch.arange(num_columns, device=q.device)
- anchors = blocks * compress_ratio + compress_ratio - 1
- logical_pages, in_page = anchors // tokens_per_block, anchors % tokens_per_block
- table = block_table.long()
- requests = request_indices.long()
- scores = torch.empty(rows, num_columns, dtype=torch.float32, device=q.device)
- chunk = max(1, 2**26 // max(num_columns * head_dim, 1))
- for lo in range(0, rows, chunk):
- hi = min(lo + chunk, rows)
- # Pages past a sequence are -1; their columns are masked below.
- physical = table[requests[lo:hi]][:, logical_pages].clamp_min(0)
- keys = index_cache[physical, in_page.expand(hi - lo, num_columns)][:, :, 0, :]
- per_head = torch.einsum("rhd,rbd->rbh", q[lo:hi].float(), keys.float())
- scores[lo:hi] = torch.relu(per_head).sum(-1) * (head_dim**-0.5)
- visible = blocks[None, :] < ((query_positions.long() + 1) // compress_ratio)[:, None]
- tail = float("nan") if only_visible_blocks else float("-inf")
- return scores.masked_fill(~visible, tail)
-
-
-qsa_paged_index_scores = QsaPagedIndexScores()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_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
deleted file mode 100644
index 0100f49724d9..000000000000
--- 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
+++ /dev/null
@@ -1,400 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""QSA index side-cache store: raw index keys and position coordinates into their pages."""
-
-from __future__ import annotations
-
-from typing import Tuple
-
-import torch
-import triton
-import triton.language as tl
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-_INTEGER_DTYPES = (torch.int32, torch.int64)
-
-
-# The row count and the pool/table extents change from step to step; letting
-# Triton specialize on them (== 1, divisible by 16) only multiplies compiles.
-@triton.jit(do_not_specialize=["num_rows", "num_requests", "table_width", "num_pages"])
-def _qsa_paged_index_store_kernel(
- token_k,
- coords,
- index_cache,
- position_cache,
- request_indices,
- logical_positions,
- block_table,
- num_rows,
- num_requests,
- table_width,
- num_pages,
- stride_k_row,
- stride_k_head,
- stride_k_dim,
- stride_c_row,
- stride_c_axis,
- stride_i_page,
- stride_i_slot,
- stride_i_head,
- stride_i_dim,
- stride_p_page,
- stride_p_slot,
- stride_p_axis,
- stride_request,
- stride_position,
- stride_t_request,
- stride_t_column,
- HKV: tl.constexpr,
- HEAD_DIM: tl.constexpr,
- BLOCK_D: tl.constexpr,
- AXES: tl.constexpr,
- BLOCK_A: tl.constexpr,
- TOKENS_PER_BLOCK: tl.constexpr,
- BLOCK_ROWS: tl.constexpr,
-):
- # One program stores BLOCK_ROWS rows: the page lookup, then every head of
- # the key and the coordinate row, all through the operands' own strides.
- rows = (tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)).to(tl.int64)
- live = rows < num_rows
- request = tl.load(request_indices + rows * stride_request, mask=live, other=-1).to(tl.int64)
- position = tl.load(logical_positions + rows * stride_position, mask=live, other=-1).to(tl.int64)
- column = position // TOKENS_PER_BLOCK
- slot = position % TOKENS_PER_BLOCK
- live &= (request >= 0) & (request < num_requests) & (position >= 0) & (column < table_width)
- page = tl.load(
- block_table + request * stride_t_request + column * stride_t_column, mask=live, other=-1
- ).to(tl.int64)
- live &= (page >= 0) & (page < num_pages)
-
- dims = tl.arange(0, BLOCK_D)
- key_mask = live[:, None] & (dims < HEAD_DIM)[None, :]
- for head in tl.static_range(HKV):
- key = tl.load(
- token_k
- + rows[:, None] * stride_k_row
- + head * stride_k_head
- + dims[None, :] * stride_k_dim,
- mask=key_mask,
- other=0.0,
- )
- tl.store(
- index_cache
- + page[:, None] * stride_i_page
- + slot[:, None] * stride_i_slot
- + head * stride_i_head
- + dims[None, :] * stride_i_dim,
- key.to(index_cache.dtype.element_ty),
- mask=key_mask,
- )
-
- axes = tl.arange(0, BLOCK_A)
- axis_mask = live[:, None] & (axes < AXES)[None, :]
- coordinate = tl.load(
- coords + rows[:, None] * stride_c_row + axes[None, :] * stride_c_axis,
- mask=axis_mask,
- other=0,
- )
- tl.store(
- position_cache
- + page[:, None] * stride_p_page
- + slot[:, None] * stride_p_slot
- + axes[None, :] * stride_p_axis,
- coordinate.to(position_cache.dtype.element_ty),
- mask=axis_mask,
- )
-
-
-def _launch_rows(rows: int) -> Tuple[int, int]:
- """Rows per program and warps: one row per program while the grid is small,
- a tile of rows once there are enough rows to fill the machine. Measured on
- GB300 against 1/4/8/16/32 rows per program."""
- if rows <= 256:
- return 1, 1
- return 16, 4
-
-
-#: This checkpoint's index geometry and the engine's QSA metadata. `token_k` is
-#: the k head of the packed `[T, 4 + 1, 128]` index projection; the coordinates
-#: are the `[T, 3]` transpose of the engine's `[3, max_num_tokens]` position
-#: buffer; `slots` x `width` is `qsa_block_table`. `pages` is the pool.
-_CKPT = dict(
- hkv=1,
- head_dim=128,
- index_heads=4,
- tokens_per_block=32,
- axes=3,
- slots=32,
- width=260,
- key_dtype=torch.bfloat16,
- request_dtype=torch.int32,
- position_dtype=torch.int64,
- key_layout="packed",
- cache_layout="plain",
- hole=None,
-)
-
-#: The target's serving envelope: ISL 8192 + OSL 1024 puts decode rows up to position
-#: 9215 (page column 287), 128 concurrent requests, and `max_num_tokens` 16384
-#: rows in one prefill or mixed step.
-_WIDE = dict(slots=160, width=300)
-_PREFILL_16K = ((0, 8192), (0, 8192))
-
-#: 32 requests of one MTP7 verify block each (8 rows), at starts that put some blocks across a
-#: page edge and some inside one.
-_MTP7_32 = tuple((29 + 61 * i, 8) for i in range(32))
-
-
-def _spec(segments, pages, **overrides) -> dict:
- """`segments` is `(first position, rows)` per request, in row order."""
- return {**_CKPT, "segments": segments, "pages": pages, **overrides}
-
-
-class QsaPagedIndexStore(OpWrapper):
- """Write each row's raw index key and position coordinates into its paged slot.
-
- For row t, with `page = block_table[request_indices[t], logical_positions[t]
- // tokens_per_block]` and `slot = logical_positions[t] % tokens_per_block`::
-
- index_cache[page, slot] = token_k[t] # cast to the cache dtype
- position_cache[page, slot, :A] = position_coordinates[t] # A = coordinates' width
-
- 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,
- selection, sparse attention and the main K/V pool are the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(
- 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 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 "
- "per-slot buffer holding 12 index-K pages and the position page; one row "
- "on an unallocated (-1) table entry is skipped",
- spec=_spec(
- ((0, 1), (31, 1), (32, 1), (33, 40)),
- pages=8,
- cache_layout="shared_page",
- layers=12,
- page_stride=49344,
- position_strides=(24672, 3, 1),
- hole="unallocated_page",
- ),
- ),
- )
-
- note = """
- Metadata consumed: none -- every input is an argument. The target passes
- `qsa_req_idx_per_token` (int32) and `qsa_logical_positions` (int64)
- sliced to the step's rows, `qsa_block_table` unsliced, the layer's
- `get_index_k_buffer(i)` (`[pages, tpb, 1, 128]` bf16) and the shared
- `get_qsa_position_buffer()` (`[pages, tpb, 3]` int32, rewritten
- idempotently by every sparse layer). Every operand is read through its
- own strides, all of them runtime arguments; the cells drive the target's:
- `token_k` off the packed index projection (row stride 640) and off the fused
- input projection (14080), the `[T, 3]` transpose of the position buffer as
- the coordinates, contiguous metadata and table. Under an
- NVFP4 KV pool the QSA cache manager packs each page's index-K roles and
- then its position page into one per-slot buffer, so both planes are strided
- views of one allocation: the trunk's index plane is `[pages, 32, 1, 128]`
- at page stride 49344 (12 x 4096 + 192) with the position plane
- `[pages, 32, 3]` at (24672, 3, 1); the cell checks that the other index-K
- roles of each page survive the call.
-
- Preconditions:
- - destination slots are unique across the batch: two rows naming one
- (page, slot) both store, unordered. The engine's mapping cannot
- collide.
- Loud rejections (`ValueError`, before launching; the test pins the
- unequal-length metadata one): metadata that is not
- one-dimensional, of unequal length, or not int32/int64; `token_k` or the
- coordinates with a row count other than the metadata's; an index plane
- whose page size, head count or head width disagrees with
- `tokens_per_block` and `token_k`; a position plane whose page size
- disagrees or that is narrower than the coordinates; a table that is not a
- 2-D integer tensor. `rows == 0` returns before launching.
-
- Values the caller owns -- data, not metadata (measured, not pinned by a
- test): coordinates wider than the plane's dtype wrap (int64 `2^31 + 5`
- stores as -2147483643); floats truncate toward zero.
-
- One launch per call; HEAD_DIM, the head count, the page size and the
- coordinate width are constexprs, everything else (strides, pool and table
- extents, rows) is a runtime argument and the row count and extents are
- not specialized on, so a new batch or pool size does not recompile
- (strides still specialize on 1 and 16-divisibility). Up to 256 rows run
- one row per program (1 warp), above that 16 rows per program (4 warps).
- CUDA-graph capturable (captured at 256 rows, 32 MTP7 verify blocks); a
- replay follows positions rewritten in place.
-
- Hard limits: every index is widened to int64 inside the kernel before
- any arithmetic; an int64 position at `2^32 + 5` lands past the table and
- is skipped rather than truncated (read from the kernel, not pinned by a
- test). Rows are a grid axis: 2^31 programs.
-
- Skip certified by a cell: an unallocated (`-1`) table entry, one row of the
- NVFP4-plane cell. The other skip conditions (a negative request index among
- them) share that mask and have no cell.
-
- Not certified: an fp8 index plane, keys other than bf16, any index geometry
- other than 1 head x 128 at 32-token pages, decode batches other than 32 MTP7
- verify blocks, strided metadata or tables,
- stride-0 coordinates, multiple devices, concurrent streams onto one pool,
- duplicate destinations.
- """
-
- def __call__(
- self,
- token_k: torch.Tensor,
- position_coordinates: torch.Tensor,
- index_cache: torch.Tensor,
- position_cache: torch.Tensor,
- request_indices: torch.Tensor,
- logical_positions: torch.Tensor,
- block_table: torch.Tensor,
- tokens_per_block: int,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
- 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
-
- def reference(
- self,
- token_k: torch.Tensor,
- position_coordinates: torch.Tensor,
- index_cache: torch.Tensor,
- position_cache: torch.Tensor,
- request_indices: torch.Tensor,
- logical_positions: torch.Tensor,
- block_table: torch.Tensor,
- tokens_per_block: int,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
- """Both planes after the store, as fresh tensors; the inputs are not touched.
-
- 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()
- pool = index_cache.shape[0]
- kept, pages, slots = [], [], []
- rows = zip(request_indices.tolist(), logical_positions.tolist())
- for row, (request, position) in enumerate(rows):
- 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)
- )
- 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
deleted file mode 100644
index be0791c3ccd6..000000000000
--- 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
+++ /dev/null
@@ -1,219 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""QSA sparse path's paged K/V append: this step's rows into the HND main pool."""
-
-from __future__ import annotations
-
-from typing import Tuple
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-# Tier-2 sourcing: a vendored Triton callable reached as ordinary Python. QSA's
-# kernel plane has no torch.ops surface.
-from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels as _qsa_kernels
-
-#: This checkpoint's main K/V geometry and the engine's QSA metadata. `heads_q`
-#: places k and v as head slices of the packed `[T, 24 + 2 + 2, 256]` QKV
-#: buffer; `slots` x `width` is `qsa_block_table` (max_num_sequences x
-#: max_blocks_per_seq). `pages` is the pool, independent of the batch.
-_CKPT = dict(
- hkv=2,
- head_dim=256,
- heads_q=24,
- tokens_per_block=32,
- slots=32,
- width=260,
- row_dtype=torch.bfloat16,
- cache_dtype=torch.bfloat16,
- request_dtype=torch.int32,
- position_dtype=torch.int64,
- hole=None,
- # Heads of further columns after v in each row: 27 is the target's fused input
- # projection row [q | k | v | gate | index | pad] (55 heads of 256).
- row_pad_heads=0,
-)
-
-_F8 = torch.float8_e4m3fn
-
-_DECODE = ((0, 1), (31, 1), (32, 1), (33, 1), (100, 1), (255, 1), (7, 1), (64, 1))
-
-#: 32 requests of one MTP7 verify block each (8 rows), at starts that put some blocks across a
-#: page edge and some inside one.
-_MTP7_32 = tuple((29 + 61 * i, 8) for i in range(32))
-
-
-def _spec(segments, pages, **overrides) -> dict:
- """`segments` is `(first position, rows)` per request, in row order."""
- return {**_CKPT, "segments": segments, "pages": pages, **overrides}
-
-
-def _bits(tensor: torch.Tensor) -> torch.Tensor:
- """The same storage as integers: fp8 has no comparison ops, and -0.0 != 0.0."""
- as_int = {1: torch.uint8, 2: torch.int16, 4: torch.int32}[tensor.element_size()]
- return tensor.view(as_int)
-
-
-class QsaPagedKvStore(OpWrapper):
- """Scatter every row's K and V into its request's page of the HND pool.
-
- Wraps `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels.
- triton_qsa_paged_kv_store`: one Triton launch over `(rows, Hkv)`. For row t::
-
- page = block_table[request_indices[t], logical_positions[t] // tokens_per_block]
- slot = logical_positions[t] % tokens_per_block
- k_cache[page, :, slot, :] = k[t] v_cache[page, :, slot, :] = v[t]
-
- `logical_positions` is absolute, so prefill, decode and MTP rows are the
- same call. The store is a move, converted once to the pool's dtype
- (bitwise `Tensor.to`; an fp8-e4m3 pool saturates at 448 and takes no
- scale, 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
- index plane is never touched.
- """
-
- 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 off the fused projection row (row stride 14080), "
- "at aligned, mid and page-crossing positions, plus one row on an "
- "unallocated (-1) table entry that is skipped",
- spec=_spec(_DECODE, pages=64, row_pad_heads=27, hole="unallocated_page"),
- ),
- Cell(
- why="prefill 1024 rows off the fused projection row (row stride 14080)",
- spec=_spec(((0, 256),) * 4, pages=160, row_pad_heads=27),
- ),
- # MTP7 verify blocks: 8 rows per request, so one block spans up to two pages.
- # The non-default pool the forward reaches: fp8-e4m3.
- Cell(
- why="fp8 decode MTP7: 32 requests x 8 rows (256), blocks straddling pages",
- spec=_spec(_MTP7_32, pages=1200, slots=40, cache_dtype=_F8),
- ),
- )
-
- note = """
- Metadata consumed: none -- every input is an argument. The target passes
- `qsa_req_idx_per_token` (int32, the row's batch index) and
- `qsa_logical_positions` (int64) sliced to the step's rows,
- `qsa_block_table` unsliced (raw pool page ids, `-1` unallocated), and the
- layer pool's K and V role views of `get_buffers(i, kv_layout="HND")`.
-
- Preconditions:
- - destination slots are unique across the batch: two rows naming one
- (page, slot) both store, unordered. The engine's mapping cannot
- collide.
- - `request_indices` and `logical_positions` are contiguous: the kernel
- gets no stride for them. A strided view is read through its base
- pointer and rows land at other rows' slots, silently (`is_valid`).
- Every other operand is read through its strides; the cells drive k/v as
- head slices of the packed QKV row and of the fused projection row, the
- pool's K and V role views, and a contiguous table. The callable rejects
- loudly: mismatched or non-3-D/4-D K/V shapes, head geometry,
- `k_cache.shape[2] != tokens_per_block`, row dtypes outside bf16/fp16,
- pool dtypes outside bf16/fp16/fp8-e4m3, a non-int32 or non-2-D table,
- non-integer or wrong-length metadata, CPU operands, an empty pool or
- table (not pinned by a test). `rows == 0` returns before launching.
-
- Strides and the pool/table extents are Triton constexprs: each distinct
- layout or pool size compiles its own kernel. CUDA-graph capturable
- (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 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), fp8 conversion of values past 448, the
- packed-QKV-row (MTP) layout over a bf16 pool, fp16 rows, head geometry other than
- this checkpoint's (D 256, Hkv 2, 32-token pages), int32 positions or
- int64 request indices, head-major rows, sliced or transposed tables,
- `rows == 0`, and of the rows that resolve to no slot any but an
- unallocated (`-1`) table entry (the reference skips each).
- """
-
- def __call__(
- self,
- k: torch.Tensor,
- v: torch.Tensor,
- k_cache: torch.Tensor,
- v_cache: torch.Tensor,
- request_indices: torch.Tensor,
- logical_positions: torch.Tensor,
- block_table: torch.Tensor,
- tokens_per_block: int,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
- _qsa_kernels.triton_qsa_paged_kv_store(
- k=k,
- v=v,
- k_cache=k_cache,
- v_cache=v_cache,
- request_indices=request_indices,
- logical_positions=logical_positions,
- block_table=block_table,
- tokens_per_block=tokens_per_block,
- )
- return k_cache, v_cache
-
- def reference(
- self,
- k: torch.Tensor,
- v: torch.Tensor,
- k_cache: torch.Tensor,
- v_cache: torch.Tensor,
- request_indices: torch.Tensor,
- logical_positions: torch.Tensor,
- block_table: torch.Tensor,
- tokens_per_block: int,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
- """Both caches after the scatter, as fresh tensors; the inputs are not touched."""
- expected_k, expected_v = k_cache.clone(), v_cache.clone()
- requests = request_indices.long()
- positions = logical_positions.long()
- columns = positions.div(tokens_per_block, rounding_mode="floor")
- table_rows, table_width = block_table.shape
- addressable = (
- (requests >= 0) & (requests < table_rows) & (positions >= 0) & (columns < table_width)
- )
- pages = block_table[
- requests.clamp(0, table_rows - 1), columns.clamp(0, table_width - 1)
- ].long()
- addressable &= (pages >= 0) & (pages < k_cache.shape[0])
- page, slot = pages[addressable], (positions % tokens_per_block)[addressable]
- expected_k[page, :, slot, :] = k[addressable].to(expected_k.dtype)
- expected_v[page, :, slot, :] = v[addressable].to(expected_v.dtype)
- return expected_k, expected_v
-
- def is_valid(
- self,
- k: torch.Tensor,
- v: torch.Tensor,
- k_cache: torch.Tensor,
- v_cache: torch.Tensor,
- request_indices: torch.Tensor,
- logical_positions: torch.Tensor,
- block_table: torch.Tensor,
- tokens_per_block: int,
- ) -> None:
- # The kernel reads these two as flat unit-stride buffers. Measured: a
- # column of a 2-wide buffer is read through its base pointer and the
- # rows land at other rows' slots, with nothing raised.
- assert request_indices.is_contiguous(), "request_indices must be a contiguous 1-D tensor"
- assert logical_positions.is_contiguous(), (
- "logical_positions must be a contiguous 1-D tensor"
- )
-
-
-qsa_paged_kv_store = QsaPagedKvStore()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_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
deleted file mode 100644
index 73e6ec1bebf7..000000000000
--- 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
+++ /dev/null
@@ -1,533 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""QSA sparse path's paged K/V append into an NVFP4 pool: quantize and scatter in one launch."""
-
-from __future__ import annotations
-
-from typing import Tuple
-
-import torch
-import triton
-import triton.language as tl
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: Values per e4m3 block scale, and e2m1 values per stored byte.
-_SF_VEC = 16
-_PER_BYTE = 2
-#: Tokens of a V scale group: V's scales are stored `[tpb / 4, D / 16, 4]` per head page.
-_V_SF_QUAD = 4
-#: e2m1's largest magnitude: a block's scale maps its amax onto it.
-_E2M1_MAX = 6.0
-#: What `cvt.rn.satfinite.e2m1x2.f32` makes of a NaN of either sign: +6. A NaN reaches
-#: it as 0 x inf, from a zero in a block whose nonzero amax rounds its scale to 0.
-_E2M1_NAN_CODE = 0x7
-#: The round-to-nearest-even decision points between neighbouring e2m1 magnitudes:
-#: `(threshold, inclusive)` steps the code up when `|v| > t`, or `|v| >= t` where the
-#: upper neighbour is the even code.
-_E2M1_STEPS = (
- (0.25, False),
- (0.75, True),
- (1.25, False),
- (1.75, True),
- (2.5, False),
- (3.5, True),
- (5.0, False),
-)
-
-#: This checkpoint's K/V geometry and the engine's QSA table. `heads_q` places k and v as
-#: head slices of the packed `[T, 24 + 2 + 2, 256]` QKV row; `global_k` / `global_v` are
-#: `kv_scale_orig_quant[1]` / `[2]`; `slot_pages` role pages of one lifecycle slot hold
-#: the call's planes, layer `layer`'s K and V among them (the trunk pool holds 24).
-_CKPT = dict(
- hkv=2,
- head_dim=256,
- heads_q=24,
- tokens_per_block=32,
- slots=32,
- width=260,
- row_dtype=torch.bfloat16,
- request_dtype=torch.int32,
- position_dtype=torch.int64,
- global_k=1.0,
- global_v=1.0,
- values="randn",
- hole=None,
- thop_append=False,
- slot_pages=2,
- layer=0,
-)
-
-_DECODE = ((0, 1), (31, 1), (32, 1), (33, 1), (100, 1), (255, 1), (7, 1), (64, 1))
-
-
-def _spec(segments, pages, **overrides) -> dict:
- """`segments` is `(first position, rows)` per request, in row order."""
- return {**_CKPT, "segments": segments, "pages": pages, **overrides}
-
-
-@triton.jit
-def _rcp_approx(x):
- """`rcp.approx.ftz.f32`, the reciprocal the engine's FP4 quantizer uses."""
- return tl.inline_asm_elementwise(
- "rcp.approx.ftz.f32 $0, $1;", "=f,f", [x], dtype=tl.float32, is_pure=True, pack=1
- )
-
-
-@triton.jit
-def _e2m1x2(lo, hi):
- """One byte of two e2m1 codes, `lo` in the low nibble (`cvt.rn.satfinite`)."""
- return tl.inline_asm_elementwise(
- "{ .reg .b8 t; .reg .b16 s; cvt.rn.satfinite.e2m1x2.f32 t, $2, $1; "
- "mov.b16 s, {t, t}; cvt.u32.u16 $0, s; }",
- "=r,f,f",
- [lo, hi],
- dtype=tl.int32,
- is_pure=True,
- pack=1,
- )
-
-
-@triton.jit
-def _quantize_row(
- src,
- stride_dim,
- global_scale,
- HEAD_DIM: tl.constexpr,
-):
- """One head row: (packed e2m1 bytes [D/2], e4m3 scale bytes [D/16])."""
- pairs = tl.arange(0, HEAD_DIM // 2)
- lo = tl.load(src + (2 * pairs) * stride_dim).to(tl.float32)
- hi = tl.load(src + (2 * pairs + 1) * stride_dim).to(tl.float32)
- pair_max = tl.maximum(tl.abs(lo), tl.abs(hi))
- block_max = tl.max(tl.reshape(pair_max, [HEAD_DIM // 16, 8]), axis=1)
- sf_value = global_scale * (block_max * _rcp_approx(tl.full([HEAD_DIM // 16], 6.0, tl.float32)))
- sf_code = sf_value.to(tl.float8e4nv)
- sf_rounded = sf_code.to(tl.float32)
- out_scale = tl.where(
- block_max != 0.0,
- _rcp_approx(sf_rounded * _rcp_approx(tl.zeros_like(sf_rounded) + global_scale)),
- 0.0,
- )
- lo_scaled = tl.reshape(
- tl.reshape(lo, [HEAD_DIM // 16, 8]) * out_scale[:, None], [HEAD_DIM // 2]
- )
- hi_scaled = tl.reshape(
- tl.reshape(hi, [HEAD_DIM // 16, 8]) * out_scale[:, None], [HEAD_DIM // 2]
- )
- packed = (_e2m1x2(lo_scaled, hi_scaled) & 0xFF).to(tl.uint8)
- return packed, sf_code.to(tl.uint8, bitcast=True)
-
-
-@triton.jit
-def _qsa_paged_kv_store_nvfp4_kernel(
- k,
- v,
- k_cache,
- v_cache,
- k_scale_cache,
- v_scale_cache,
- request_indices,
- logical_positions,
- block_table,
- kv_scale_orig_quant,
- num_pages,
- num_requests,
- table_width,
- k_stride_row,
- k_stride_head,
- k_stride_dim,
- v_stride_row,
- v_stride_head,
- v_stride_dim,
- kc_stride_page,
- kc_stride_head,
- kc_stride_token,
- vc_stride_page,
- vc_stride_head,
- vc_stride_token,
- ks_stride_page,
- ks_stride_head,
- vs_stride_page,
- vs_stride_head,
- bt_stride_request,
- bt_stride_page,
- HEAD_DIM: tl.constexpr,
- TOKENS_PER_BLOCK: tl.constexpr,
-):
- row = tl.program_id(0).to(tl.int64)
- head = tl.program_id(1).to(tl.int64)
- request = tl.load(request_indices + row).to(tl.int64)
- logical = tl.load(logical_positions + row).to(tl.int64)
- column = logical // TOKENS_PER_BLOCK
- token = logical % TOKENS_PER_BLOCK
- lookup = (request >= 0) & (request < num_requests) & (logical >= 0) & (column < table_width)
- page = tl.load(
- block_table + request * bt_stride_request + column * bt_stride_page, mask=lookup, other=-1
- ).to(tl.int64)
- valid = (page >= 0) & (page < num_pages)
- page = tl.maximum(page, 0)
-
- k_data, k_sf = _quantize_row(
- k + row * k_stride_row + head * k_stride_head,
- k_stride_dim,
- tl.load(kv_scale_orig_quant + 1),
- HEAD_DIM,
- )
- v_data, v_sf = _quantize_row(
- v + row * v_stride_row + head * v_stride_head,
- v_stride_dim,
- tl.load(kv_scale_orig_quant + 2),
- HEAD_DIM,
- )
- pairs = tl.arange(0, HEAD_DIM // 2)
- blocks = tl.arange(0, HEAD_DIM // 16)
- tl.store(
- k_cache + page * kc_stride_page + head * kc_stride_head + token * kc_stride_token + pairs,
- k_data,
- mask=valid,
- )
- tl.store(
- v_cache + page * vc_stride_page + head * vc_stride_head + token * vc_stride_token + pairs,
- v_data,
- mask=valid,
- )
- # K scales token-major within the head page; V scales in four-token groups.
- tl.store(
- k_scale_cache
- + page * ks_stride_page
- + head * ks_stride_head
- + token * (HEAD_DIM // 16)
- + blocks,
- k_sf,
- mask=valid,
- )
- tl.store(
- v_scale_cache
- + page * vs_stride_page
- + head * vs_stride_head
- + (token // 4) * (4 * (HEAD_DIM // 16))
- + blocks * 4
- + token % 4,
- v_sf,
- mask=valid,
- )
-
-
-def _store_nvfp4(
- *,
- k: torch.Tensor,
- v: torch.Tensor,
- k_cache: torch.Tensor,
- v_cache: torch.Tensor,
- k_scale_cache: torch.Tensor,
- v_scale_cache: torch.Tensor,
- request_indices: torch.Tensor,
- logical_positions: torch.Tensor,
- block_table: torch.Tensor,
- tokens_per_block: int,
- kv_scale_orig_quant: torch.Tensor,
-) -> None:
- """Quantize rows of K and V to NVFP4 and store them into their pages (see the entry)."""
- if k.ndim != 3 or v.shape != k.shape:
- raise ValueError("NVFP4 K/V store: k and v must be matching [rows, heads, dim]")
- rows, heads, head_dim = k.shape
- if k.dtype not in (torch.bfloat16, torch.float16) or v.dtype != k.dtype:
- raise ValueError("NVFP4 K/V store: k and v must both be BF16 or both FP16")
- if head_dim < 32 or head_dim & (head_dim - 1):
- raise ValueError("NVFP4 K/V store: head_dim must be a power of two of at least 32")
- if tokens_per_block <= 0 or tokens_per_block % _V_SF_QUAD:
- raise ValueError("NVFP4 K/V store: tokens_per_block must be a positive multiple of 4")
- data_shape = (k_cache.shape[0], heads, tokens_per_block, head_dim // _PER_BYTE)
- scale_shape = (k_cache.shape[0], heads, tokens_per_block, head_dim // _SF_VEC)
- for name, cache, shape in (
- ("k_cache", k_cache, data_shape),
- ("v_cache", v_cache, data_shape),
- ("k_scale_cache", k_scale_cache, scale_shape),
- ("v_scale_cache", v_scale_cache, scale_shape),
- ):
- if cache.dtype != torch.uint8 or tuple(cache.shape) != shape:
- raise ValueError(
- f"NVFP4 K/V store: {name} must be uint8 {shape}, got {cache.dtype} "
- f"{tuple(cache.shape)}"
- )
- if cache.stride(3) != 1:
- raise ValueError(f"NVFP4 K/V store: {name} must be contiguous along its last dim")
- for name, cache in (("k_scale_cache", k_scale_cache), ("v_scale_cache", v_scale_cache)):
- if cache.stride(2) != head_dim // _SF_VEC:
- raise ValueError(f"NVFP4 K/V store: each head page of {name} must be contiguous")
- if block_table.ndim != 2 or block_table.dtype != torch.int32:
- raise ValueError("NVFP4 K/V store: block_table must be a 2-D int32 tensor")
- for name, meta in (
- ("request_indices", request_indices),
- ("logical_positions", logical_positions),
- ):
- if meta.shape != (rows,) or meta.dtype not in (torch.int32, torch.int64):
- raise ValueError(f"NVFP4 K/V store: {name} must be one int32/int64 entry per row")
- if not meta.is_contiguous():
- raise ValueError(f"NVFP4 K/V store: {name} must be contiguous")
- if (
- kv_scale_orig_quant.dtype != torch.float32
- or kv_scale_orig_quant.shape != (3,)
- or not kv_scale_orig_quant.is_contiguous()
- ):
- raise ValueError("NVFP4 K/V store: kv_scale_orig_quant must be a contiguous fp32 [3]")
- operands = (k, v, k_cache, v_cache, k_scale_cache, v_scale_cache, request_indices)
- operands += (logical_positions, block_table, kv_scale_orig_quant)
- if not all(t.is_cuda and t.device == k.device for t in operands):
- raise ValueError("NVFP4 K/V store: every operand must be on k's CUDA device")
- if rows == 0:
- return
- _qsa_paged_kv_store_nvfp4_kernel[(rows, heads)](
- k,
- v,
- k_cache,
- v_cache,
- k_scale_cache,
- v_scale_cache,
- request_indices,
- logical_positions,
- block_table,
- kv_scale_orig_quant,
- k_cache.shape[0],
- block_table.shape[0],
- block_table.shape[1],
- k.stride(0),
- k.stride(1),
- k.stride(2),
- v.stride(0),
- v.stride(1),
- v.stride(2),
- k_cache.stride(0),
- k_cache.stride(1),
- k_cache.stride(2),
- v_cache.stride(0),
- v_cache.stride(1),
- v_cache.stride(2),
- k_scale_cache.stride(0),
- k_scale_cache.stride(1),
- v_scale_cache.stride(0),
- v_scale_cache.stride(1),
- block_table.stride(0),
- block_table.stride(1),
- HEAD_DIM=head_dim,
- TOKENS_PER_BLOCK=tokens_per_block,
- num_warps=4,
- )
-
-
-def _e2m1_codes(values: torch.Tensor) -> torch.Tensor:
- """fp32 -> e2m1 codes 0..15, round to nearest even, saturating at 6; a NaN
- (0 x inf, in a block whose scale underflows to zero) converts to +6."""
- magnitude = values.abs()
- code = torch.zeros(values.shape, dtype=torch.int32, device=values.device)
- for threshold, inclusive in _E2M1_STEPS:
- code += (magnitude >= threshold if inclusive else magnitude > threshold).to(torch.int32)
- code |= torch.signbit(values).to(torch.int32) << 3
- return torch.where(torch.isnan(values), torch.full_like(code, _E2M1_NAN_CODE), code)
-
-
-def _quantize_reference(rows: torch.Tensor, global_scale: torch.Tensor):
- """[n, D] rows -> (packed bytes [n, D/2], e4m3 scale bytes [n, D/16]), in plain torch."""
- n, head_dim = rows.shape
- x = rows.float().reshape(n, head_dim // _SF_VEC, _SF_VEC)
- block_max = x.abs().amax(dim=-1)
- sixth = torch.tensor(_E2M1_MAX, dtype=torch.float32, device=rows.device).reciprocal()
- scale_code = (global_scale * (block_max * sixth)).to(torch.float8_e4m3fn)
- rounded = scale_code.float()
- out_scale = torch.where(
- block_max != 0,
- (rounded * global_scale.reciprocal()).reciprocal(),
- torch.zeros_like(rounded),
- )
- codes = _e2m1_codes(x * out_scale[..., None]).reshape(n, head_dim // _PER_BYTE, _PER_BYTE)
- packed = (codes[..., 0] | (codes[..., 1] << 4)).to(torch.uint8)
- return packed, scale_code.view(torch.uint8)
-
-
-class QsaPagedKvStoreNvfp4(OpWrapper):
- """Quantize every row's K and V to NVFP4 and store them into the request's page.
-
- A Triton kernel defined in this file, one launch over `(rows, Hkv)`. For row t
- and KV head h, with `page = block_table[request_indices[t], logical_positions[t]
- // tpb]` and `s = logical_positions[t] % tpb`, each 16-value block `b` of
- `k[t, h]` (and alike for v, with `kv_scale_orig_quant[2]`) is stored as::
-
- g = kv_scale_orig_quant[1]
- sf = e4m3(g * (amax_b * rcp(6))) # one byte
- code_i = e2m1(k[t, h, i] * rcp(float(sf) * rcp(g))) # 0 when amax_b == 0
- k_cache[page, h, s, i // 2] nibble i % 2 (low nibble = even i)
- k_scale_cache[page, h, s, b] = sf # token-major
- v_scale_cache[page, h, (s // 4) * 4 * (D / 16) + 4 * b + s % 4] # four-token groups
-
- `rcp` is `rcp.approx.ftz.f32`, `e4m3` and `e2m1` round to nearest even and
- saturate. This is the attention op's own NVFP4 append (`cvt_warp_fp16_to_fp4`
- behind `quantizeAndWriteFP4KVCache`), the same bits for the same row: a
- token's stored value does not depend on which path wrote it. A row that
- resolves to no slot -- negative request, request past the table, negative
- position, page column past the table, a table entry that is negative or past
- the pool -- is skipped. Returns the four caches, written in place.
-
- Fusion boundary: quantization and the scatter. The projection, norm and RoPE
- of K are the caller's; the QSA index plane is not touched.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="decode: 8 rows at aligned, mid-page and page-crossing positions, with zeros, "
- "-0.0 and a 900.0 outlier inside blocks, and a row on an unallocated (-1) page: "
- "zero scales, saturation and the skip",
- spec=_spec(_DECODE, pages=69, slots=8, values="edge", hole="unallocated_page"),
- ),
- Cell(
- why="prefill: 1024 rows, 4 requests x 256 walking 8 pages each, in the engine's "
- "trunk layout: layer 11's planes among the 24 role pages of a slot",
- spec=_spec(((0, 256),) * 4, pages=160, slot_pages=24, layer=11),
- ),
- Cell(
- why="the attention op's own append over the same rows (context, chunk, decode "
- "P=1 and P=3) at unit second-level scales: data and both scale layouts bitwise equal",
- spec=_spec(((0, 115), (0, 101), (0, 36)), pages=96, thop_append=True),
- ),
- )
-
- note = """
- Metadata consumed: none -- every input is an argument. The target passes the
- step's `qsa_req_idx_per_token` (int32) and `qsa_logical_positions` (int64),
- `qsa_block_table` unsliced (raw pool page ids, `-1` unallocated), the layer's
- K and V planes of the QSA manager's `get_buffers(i, "HND")` (uint8 `[slots,
- Hkv, tpb, D/2]`), the two planes of `get_block_scale_buffers(i, "HND")`
- (uint8 `[slots, Hkv, tpb, D/16]`), and its own `kv_scale_orig_quant` -- the
- same fp32 `[3]` tensor its dense attention call binds, so both writers apply
- one second-level scale.
-
- Preconditions:
- - destination slots are unique across the batch: two rows naming one
- (page, slot) both store, unordered.
- - the dense path's reader assumes this layout: e2m1 nibbles HND per page,
- K scales linear, V scales in four-token groups. The attention op's
- append, run over the same rows, leaves bitwise the same pool (a cell),
- at the unit second-level scales the target binds.
- The callable rejects loudly (ValueError; the test pins the non-contiguous
- metadata refusal): non-matching or non-3-D k/v, a row
- dtype outside bf16/fp16, a head_dim that is not a power of two >= 32, a
- tokens_per_block that is not a multiple of 4, caches that are not uint8 of
- the stated shapes or whose last dim or scale head page is strided, a non-2-D
- or non-int32 table, metadata that is not one contiguous int32/int64 entry per
- row, a `kv_scale_orig_quant` that is not a contiguous fp32 `[3]`, operands on
- another device. `rows == 0` returns before launching (not certified).
-
- Every offset is formed in 64 bits: rows, heads, requests, positions and pages
- are widened before any product. CUDA-graph capturable; a replay follows rows,
- pages and positions rewritten in place (the test replays 8 MTP verify blocks of
- 8 rows bitwise equal to eager).
-
- Compiled variants: the page count (`num_pages`, the pool's slot count), the
- table's height (`num_requests`) and width and every stride are plain-int
- arguments, and Triton compiles the kernel once per class of each -- equal to
- 1, a multiple of 16, anything else -- besides once per head_dim and
- tokens_per_block. A served call puts the width and each stride in one fixed
- class (the k/v head slices of the packed QKV row, the planes' slot, head and
- token strides, the contiguous table's) and varies the page count and the
- table height. The cells run (other, other) and (multiple of 16, multiple of
- 16); the other combinations, a table height of 1 among them, are not
- certified.
-
- Not finite values: a NaN or +-inf in a block is not certified (the e2m1 and
- e4m3 conversions saturate, the block maximum is the kernel's `max`).
-
- Not certified either: second-level scales other than 1.0, int32 positions
- or int64 request indices, and of the rows that resolve to no slot any but
- an unallocated (`-1`) table entry (the reference skips each).
- """
-
- def __call__(
- self,
- k: torch.Tensor,
- v: torch.Tensor,
- k_cache: torch.Tensor,
- v_cache: torch.Tensor,
- k_scale_cache: torch.Tensor,
- v_scale_cache: torch.Tensor,
- request_indices: torch.Tensor,
- logical_positions: torch.Tensor,
- block_table: torch.Tensor,
- tokens_per_block: int,
- kv_scale_orig_quant: torch.Tensor,
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
- _store_nvfp4(
- k=k,
- v=v,
- k_cache=k_cache,
- v_cache=v_cache,
- k_scale_cache=k_scale_cache,
- v_scale_cache=v_scale_cache,
- request_indices=request_indices,
- logical_positions=logical_positions,
- block_table=block_table,
- tokens_per_block=tokens_per_block,
- kv_scale_orig_quant=kv_scale_orig_quant,
- )
- return k_cache, v_cache, k_scale_cache, v_scale_cache
-
- def reference(
- self,
- k: torch.Tensor,
- v: torch.Tensor,
- k_cache: torch.Tensor,
- v_cache: torch.Tensor,
- k_scale_cache: torch.Tensor,
- v_scale_cache: torch.Tensor,
- request_indices: torch.Tensor,
- logical_positions: torch.Tensor,
- block_table: torch.Tensor,
- tokens_per_block: int,
- kv_scale_orig_quant: torch.Tensor,
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
- """The four caches after the store, as fresh tensors; the inputs are untouched.
-
- Quantization in fp32 torch with IEEE reciprocals where the kernel uses
- `rcp.approx.ftz`; the two agree on every operand the cells drive (6.0, the
- unit second-level scale, and the e4m3 scale values their rows produce), which
- the cells' bitwise gate is what pins.
- """
- out = [
- t.clone(memory_format=torch.contiguous_format)
- for t in (k_cache, v_cache, k_scale_cache, v_scale_cache)
- ]
- rows, heads, head_dim = k.shape
- requests = request_indices.long()
- positions = logical_positions.long()
- columns = positions.div(tokens_per_block, rounding_mode="floor")
- table_rows, table_width = block_table.shape
- ok = (requests >= 0) & (requests < table_rows) & (positions >= 0) & (columns < table_width)
- pages = block_table[
- requests.clamp(0, table_rows - 1), columns.clamp(0, table_width - 1)
- ].long()
- ok &= (pages >= 0) & (pages < k_cache.shape[0])
- page, slot = pages[ok], (positions % tokens_per_block)[ok]
- n = int(page.numel())
- if n == 0:
- return tuple(out)
- groups = head_dim // _SF_VEC
- quad_base = (slot // _V_SF_QUAD) * (_V_SF_QUAD * groups) + slot % _V_SF_QUAD
- for src, data, scales, global_scale, linear in (
- (k, out[0], out[2], kv_scale_orig_quant[1], True),
- (v, out[1], out[3], kv_scale_orig_quant[2], False),
- ):
- packed, sf = _quantize_reference(src[ok].reshape(n * heads, head_dim), global_scale)
- data[page, :, slot, :] = packed.reshape(n, heads, head_dim // _PER_BYTE)
- sf = sf.reshape(n, heads, groups)
- if linear:
- scales[page, :, slot, :] = sf
- else:
- flat = scales.view(scales.shape[0], heads, tokens_per_block * groups)
- offsets = quad_base[:, None] + _V_SF_QUAD * torch.arange(groups, device=sf.device)
- for h in range(heads):
- flat[page[:, None], h, offsets] = sf[:, h]
- return tuple(out)
-
-
-qsa_paged_kv_store_nvfp4 = QsaPagedKvStoreNvfp4()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py
deleted file mode 100644
index 21609fd37846..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py
+++ /dev/null
@@ -1,283 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""QSA sparse GQA: attention over an explicitly selected set of paged K/V tokens."""
-
-from __future__ import annotations
-
-from typing import Optional
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-# Approved exception: QSA's sparse path never reaches the thop attention op's
-# sparse argument group, so there is no torch.ops surface for this computation.
-from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import (
- can_fuse_qsa_splitk_output_gate as _can_fuse_qsa_splitk_output_gate,
-)
-from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import (
- triton_qsa_paged_sparse_gqa as _triton_qsa_paged_sparse_gqa,
-)
-
-_CHECKPOINT = dict(
- q_heads=24,
- kv_heads=2,
- head_dim=256,
- tokens_per_block=32,
- compress_ratio=4,
- token_topk=2048,
- cache_dtype=torch.bfloat16,
- q_dtype=torch.bfloat16,
- # q is the first 24 heads of the packed [T, 24 + 2 + 2, 256] QKV buffer;
- # "fused_row" is the target's fused input projection row, 55 heads of 256.
- q_layout="qkv_slice",
- request_dtype=torch.int32,
- # The engine's table: 260 pages per slot, -1 past each sequence.
- pages=260,
-)
-
-_SPLIT_K = ("splitk", "merge")
-_PREFILL = ("prefill",)
-_FP8 = dict(cache_dtype=torch.float8_e4m3fn)
-
-
-def _spec(**overrides) -> dict:
- """A cell spec: this checkpoint's geometry unless the cell says otherwise.
-
- `segments` is `(first position, rows)` per request, in row order;
- `prefill_pair` says whether `query_positions` / `compress_ratio` are
- passed, which the target does exactly when the step has a context request.
- """
- return {**_CHECKPOINT, **overrides}
-
-
-def _decode(segments, launches=_SPLIT_K, **overrides) -> dict:
- base = dict(segments=segments, slots=32, prefill_pair=False, launches=launches)
- return _spec(**{**base, **overrides})
-
-
-class QsaPagedSparseGqa(OpWrapper):
- """Sparse GQA over the selected tokens of an HND paged K/V pool.
-
- Wraps `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels.
- triton_qsa_paged_sparse_gqa`. For row `r`, query head `h`, KV head
- `h // (Hq // Hkv)` (contiguous groups, not interleaved), over the columns
- `j` of `selected_tokens[r]` that name a resident token::
-
- t = selected_tokens[r, j] # logical position
- page = block_table[request_indices[r], t // tokens_per_block]
- s_j = (q[r, h] . k_cache[page, kv, t % tpb]) * softmax_scale
- out[r, h] = sum_j softmax(s)_j * v_cache[page, kv, t % tpb]
-
- A negative index, a page past the table's width, a negative or
- out-of-pool table entry, and a request index outside the table's rows all
- contribute nothing; a row left with nothing returns zeros. Indices are not
- deduplicated: a token named twice is attended twice. Returns a fresh,
- contiguous `[rows, Hq, D]` tensor of `q.dtype`, even for a strided `q`.
-
- Two kernels, chosen inside the call: with `query_positions` and
- `compress_ratio` given and `rows * Hkv > 512`, a causal-prefix prefill
- kernel; otherwise split-K, plus a merge kernel when it used more than one
- split. Each cell names the launches it was measured on.
-
- **The prefill pair is a bound, not a mask.** The prefill kernel reads only
- `min(visible // cr, block_topk) * cr + visible % cr` leading columns of a
- row (`visible = query_positions[r] + 1`, rounded up to a 64-wide tile) and
- nothing after them. That is exactly what QSA's own top-k + expansion emit,
- packed at the front; a selection with valid entries past the bound is
- silently truncated (measured: a planted index moving the right answer by
- 0.77 moved the op's by nothing). The split-K kernel ignores the pair.
-
- **The output gate rides the merge when `output_gate` is passed.** Each
- merged head row is rounded to `q.dtype`, multiplied in fp32 by
- `sigmoid(output_gate)` and rounded again: the arithmetic of the ungated
- call followed by `fused_sigmoid_mul_inplace`. `output_gate` is `[rows,
- Hq * D]` or `[rows, Hq, D]` in `q.dtype`, unit stride in its last dim, any
- row stride (the gate columns of the fused projection row). Only a split-K
- launch with more than one split has a merge to carry it; `fuses_output_gate`
- says whether a call does, and the op raises `ValueError` if handed a gate
- it cannot carry.
-
- Fusion boundary: the block-table gather, scale, online softmax, P.V, the
- split-K merge and, when passed, the output gate are inside. Deciding what
- to attend to, appending K/V to the pool, and q's projection, norm and RoPE
- are not.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- # ---- generation: no prefill pair, split-K ---------------------------
- Cell(
- why="decode MTP7: 32 slots x 8 rows = 256 rows, 512 programs, 4-way split-K, "
- "q and the gate off the fused row, the gate fused into the merge",
- spec=_decode(
- tuple((100 + (499 * g) % 8000, 8) for g in range(32)),
- q_layout="fused_row",
- gate="fused_row",
- ),
- ),
- # ---- context: the prefill pair is passed ----------------------------
- Cell(
- why="fp8 context 272 rows from 2032: causal-prefix prefill kernel",
- spec=_decode(((2032, 272),), launches=_PREFILL, prefill_pair=True, **_FP8),
- ),
- Cell(
- why="q off the fused projection row (55 x 256), 8191-row prefill",
- spec=_decode(((0, 8191),), launches=_PREFILL, prefill_pair=True, q_layout="fused_row"),
- ),
- )
-
- note = """
- Metadata consumed: none -- every input is an argument. What the target
- passes: `block_table` is `qsa_block_table`, physical page ids into dim 0
- of the caches (the *unscaled* table; a reference assuming logical ==
- physical lands 306x outside tolerance); `selected_tokens` is
- `qsa_expand_block_indices`' output, logical positions, `-1`-padded;
- `k_cache` / `v_cache` are the layer pool's role planes.
-
- Preconditions:
- - Whenever the prefill kernel can run, each selection row is packed at
- the front with exactly the bound's count of valid entries (see the
- class docstring). QSA's producer -- `qsa_topk_select` with zero
- `row_starts` and `row_ends = (pos + 1) // cr`, then the expansion --
- was measured to satisfy it at 300, 8191 and 8192 rows, including with
- NaN/inf planted in the scores. Checking it needs the contents, so it
- is not an `is_valid` guard.
- - The prefill pair is a property of the launch: passed because some
- rows are context, it puts the generation rows on the prefill kernel
- too, and the bound applies to them (measured with a planted index; not
- pinned by a test).
- - There is no KV scale argument: an fp8 pool is read at face value (each code
- widens exactly to bf16 under a bf16 q), so it must hold K/V written at unit scale.
- - `D` must be a power of two (192 is refused); `Hq % Hkv == 0`.
-
- The op validates its own inputs loudly (`ValueError`s; the test pins the
- refusal of positions passed without a compress ratio; a selection width
- the compress ratio cannot produce is refused on the prefill kernel only),
- so `is_valid` adds nothing. `rows == 0` returns an empty `[0, Hq, D]`
- tensor (not certified).
-
- PDL: `TRTLLM_ENABLE_PDL` (default on) enables programmatic dependent
- launch for multi-split launches of 4 rows or fewer; no cell runs one (the
- target sends gated decode batches that small to
- `qsa_paged_sparse_gqa_decode`). Calls are bitwise deterministic (pinned
- on a 32-split launch and its merge).
-
- The decode site carries up to 8 rows per request (max_draft_len 7): the
- decode cell runs 32 requests of 8 rows, 512 `rows * Hkv` programs and a
- 4-way split-K launch with the output gate in its merge.
-
- Not certified: capture in a CUDA graph (the split-K path allocates its
- partials inside the call); PDL launches; head geometries other than 24/2
- x 256; `tokens_per_block` other than 32; more than 8191 rows; the split-K
- kernel above 256 rows, with one split (no merge) or without the gate; an
- fp8 pool on the split-K kernel; rows whose selection resolves no token
- (all padding, a foreign request, unmapped pages); an fp16 pool or q;
- standalone contiguous K/V buffers; int64 request indices.
- Hard limits: not probed; the largest certified extent is 8191 rows x 2051
- selected tokens through a [32, 260] table.
- """
-
- def __call__(
- self,
- *,
- q: torch.Tensor,
- k_cache: torch.Tensor,
- v_cache: torch.Tensor,
- block_table: torch.Tensor,
- selected_tokens: torch.Tensor,
- request_indices: torch.Tensor,
- tokens_per_block: int,
- softmax_scale: float,
- query_positions: Optional[torch.Tensor] = None,
- compress_ratio: Optional[int] = None,
- output_gate: Optional[torch.Tensor] = None,
- ) -> torch.Tensor:
- return _triton_qsa_paged_sparse_gqa(
- q=q,
- k_cache=k_cache,
- v_cache=v_cache,
- block_table=block_table,
- selected_tokens=selected_tokens,
- request_indices=request_indices,
- tokens_per_block=tokens_per_block,
- softmax_scale=softmax_scale,
- query_positions=query_positions,
- compress_ratio=compress_ratio,
- output_gate=output_gate,
- )
-
- @staticmethod
- def fuses_output_gate(
- q: torch.Tensor,
- k_cache: torch.Tensor,
- selected_tokens: torch.Tensor,
- output_gate: torch.Tensor,
- ) -> bool:
- """Whether a call with these operands and no prefill pair carries `output_gate`."""
- return _can_fuse_qsa_splitk_output_gate(q, k_cache, selected_tokens, output_gate)
-
- def reference(
- self,
- *,
- q: torch.Tensor,
- k_cache: torch.Tensor,
- v_cache: torch.Tensor,
- block_table: torch.Tensor,
- selected_tokens: torch.Tensor,
- request_indices: torch.Tensor,
- tokens_per_block: int,
- softmax_scale: float,
- query_positions: Optional[torch.Tensor] = None,
- compress_ratio: Optional[int] = None,
- output_gate: Optional[torch.Tensor] = None,
- ) -> torch.Tensor:
- """fp32 sparse GQA over every resident selected token, rounded once;
- with `output_gate`, then gated in fp32 and rounded again.
-
- `query_positions` and `compress_ratio` are ignored: on a prefix-packed
- selection the prefill kernel's bound reads exactly the valid entries.
- Walked 32 rows at a time to bound the gathered K/V.
- """
- rows, num_q_heads, head_dim = q.shape
- num_pages, num_kv_heads = k_cache.shape[0], k_cache.shape[1]
- groups = num_q_heads // num_kv_heads
- num_requests, table_width = block_table.shape
- requests = request_indices.long()
- out = torch.empty(rows, num_q_heads, head_dim, dtype=torch.float32, device=q.device)
- for start in range(0, rows, 32):
- stop = min(start + 32, rows)
- chunk = selected_tokens[start:stop].long()
- logical_page = chunk.clamp_min(0) // tokens_per_block
- in_page = chunk.clamp_min(0) % tokens_per_block
- request = requests[start:stop]
- owned = (request >= 0) & (request < num_requests)
- physical = block_table[
- request.clamp(0, num_requests - 1)[:, None].expand_as(logical_page),
- logical_page.clamp(max=table_width - 1),
- ].long()
- resident = (
- (chunk >= 0)
- & owned[:, None]
- & (logical_page < table_width)
- & (physical >= 0)
- & (physical < num_pages)
- )
- physical = physical.clamp(0, num_pages - 1)
- keys = k_cache[physical, :, in_page, :].float()
- values = v_cache[physical, :, in_page, :].float()
- grouped = q[start:stop].reshape(-1, num_kv_heads, groups, head_dim).float()
- scores = torch.einsum("bhgd,bkhd->bhgk", grouped, keys) * softmax_scale
- scores = scores.masked_fill(~resident[:, None, None, :], float("-inf"))
- # A row with nothing resident softmaxes over all -inf: NaN, i.e. zeros.
- weights = torch.nan_to_num(torch.softmax(scores, dim=-1), nan=0.0)
- pooled = torch.einsum("bhgk,bkhd->bhgd", weights, values)
- out[start:stop] = pooled.reshape(stop - start, num_q_heads, head_dim)
- if output_gate is None:
- return out.to(q.dtype)
- gate = output_gate.reshape(rows, num_q_heads, head_dim).float()
- return (out.to(q.dtype).float() * torch.sigmoid(gate)).to(q.dtype)
-
-
-qsa_paged_sparse_gqa = QsaPagedSparseGqa()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_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
deleted file mode 100644
index b6180bcb7566..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_decode.py
+++ /dev/null
@@ -1,819 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""QSA sparse GQA for decode rows: a split-K pass that loads its tokens in one round trip,
-then a merge spread over the SMs.
-
-Tier 5 (CuTe DSL written for this target), both launches.
-
-Split-K pass. One CTA (eight warps) covers `CAP` consecutive selected tokens of one KV
-head of one row, so a decode row's 2051 selected tokens take 33 CTAs a KV head:
-
-1. before its PDL wait the CTA stages the KV head's query heads in shared memory and from
- there in every warp's MMA A fragments, stages the request's block-table row (up to
- `BT_STAGE` pages) in shared memory, and releases its dependents;
-2. after it, its tokens' pages come from shared memory and every token's K and V row
- (16 bytes a lane, a warp eight tokens) is requested at once with cp.async, so the CTA
- pays one DRAM round trip -- QSA's split-K kernel walks two or three 16-token tiles,
- each a chain of dependent loads (selection, page, K/V);
-3. scores = q K^T on warp MMAs (bf16 in, fp32 out; a warp eight tokens), scaled into the
- log2 domain and masked; one softmax over the CTA's tokens (a warp two heads); P V on
- warp MMAs with P in bf16 (a warp 32 head dims); each head's partial times its
- reciprocal sum is staged in shared memory and stored in 16-byte pieces with its
- log2-sum-exp, as QSA's split-K pass stores them.
-
-Merge. A CTA covers 128 elements of a head (two CTAs a head) with up to eight warps: warp
-`w` reads its splits' partials of four adjacent elements a lane, one 16-byte load a split,
-all in flight; every warp forms the split weights itself from two LSEs a lane; the warps'
-shares are added in shared memory.
-
-Both launches release their dependents at entry, before their own waits: the merge and,
-through it, the output projection that follows launch while the top-k ahead of the
-split-K pass still runs, and the projection streams its weight through that latency-bound
-stretch. Before its wait the split-K pass reads q and the merge the output gate: the
-caller's contract is that both are complete when the grid ahead of the split-K pass
-releases it (the target's top-k does so only after its own wait; a launch without PDL
-releases at its end).
-"""
-
-from __future__ import annotations
-
-import functools
-import math
-from typing import Optional
-
-import cuda.bindings.driver as cuda_drv
-import cutlass
-import cutlass.cute as cute
-import torch
-from cutlass import const_expr
-from cutlass.cute import experimental as cute_ext
-from cutlass.cute.nvgpu import cpasync, warp
-from cutlass.cute.runtime import from_dlpack, make_ptr
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-# Approved exception (as for `qsa_paged_sparse_gqa`): QSA's sparse path has no torch.ops
-# surface; its row-count PDL policy is reached as Python.
-from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import _qsa_pdl_enabled
-
-from .qsa_paged_sparse_gqa import qsa_paged_sparse_gqa as _composition
-
-#: Output elements a merge CTA covers (four adjacent ones a lane of each warp).
-CHUNK = 128
-#: Splits the merge reads: two LSEs a lane.
-MAX_SPLITS = 64
-#: Rows the target calls this entry with: the launches that run with PDL (QSA enables it up
-#: to four rows); larger batches keep the composition.
-MAX_ROWS = 4
-#: Selected tokens a split-K CTA covers.
-CAP = 64
-#: MMA rows: the query heads of a KV head, padded.
-M = 16
-#: bf16 elements of padding after each shared K, V, q and P row (keeps the MMA fragment
-#: reads free of bank conflicts).
-PAD = 8
-WARPS = 8
-THREADS = WARPS * 32
-#: Block-table entries staged in shared memory before the wait; a wider table is read
-#: from global memory after it.
-BT_STAGE = 4096
-HEAD_DIM = 256
-LOG2E = 1.4426950408889634
-
-
-class _SplitKKernel:
- def __init__(
- self,
- *,
- group,
- width,
- splits,
- tpb,
- scale,
- bt_width,
- num_pages,
- k_strides,
- v_strides,
- num_requests,
- pdl,
- ):
- assert group <= M and CAP == 2 * 32 and CAP % WARPS == 0
- self.group, self.width, self.splits, self.tpb = group, width, splits, tpb
- self.scale = scale * LOG2E
- self.bt_width, self.num_pages, self.num_requests = bt_width, num_pages, num_requests
- self.stage_bt = bt_width <= BT_STAGE
- self.k_strides, self.v_strides = k_strides, v_strides
- self.ld = HEAD_DIM + PAD
- self.pdl = pdl
-
- def smem_bytes(self) -> int:
- def r(n):
- return -(-n // 128) * 128
-
- b = (
- r(M * self.ld * 2)
- + 2 * r(CAP * self.ld * 2)
- + r(M * (CAP + PAD) * 2)
- + r(M * (CAP + 4) * 4)
- )
- b += r(M * 2 * 4) + 2 * r(CAP * 8)
- if self.stage_bt:
- b += r(self.bt_width * 4)
- return b + 256
-
- @cute.jit
- def __call__(self, mQ, pK, pV, mBT, mSel, mReq, pPart, mLse, stream: cuda_drv.CUstream):
- rows = cute.size(mSel, mode=[0])
- kv_heads = cute.size(mQ, mode=[1]) // self.group
- mma = cute.make_tiled_mma(
- warp.MmaF16BF16Op(cutlass.BFloat16, cutlass.Float32, (16, 8, 16)),
- cute.make_layout((1, WARPS, 1)),
- )
- g2s = cute.make_copy_atom(cpasync.CopyG2SOp(), cutlass.BFloat16, num_bits_per_copy=128)
- st = cute.make_copy_atom(
- cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=128
- )
- self.kernel(mQ, pK, pV, mBT, mSel, mReq, pPart, mLse, mma, g2s, st).launch(
- grid=[rows * kv_heads * self.splits, 1, 1],
- block=[THREADS, 1, 1],
- smem=self.smem_bytes(),
- stream=stream,
- use_pdl=self.pdl,
- )
-
- @cute.kernel
- def kernel(
- self,
- mQ,
- pK,
- pV,
- mBT,
- mSel,
- mReq,
- pPart,
- mLse,
- mma: cute.TiledMma,
- g2s: cute.CopyAtom,
- st: cute.CopyAtom,
- ):
- tidx, _, _ = cute.arch.thread_idx()
- bidx, _, _ = cute.arch.block_idx()
- warp_i = cute.arch.warp_idx()
- lane = cute.arch.lane_idx()
- group: cutlass.Constexpr = self.group
- splits: cutlass.Constexpr = self.splits
- width: cutlass.Constexpr = self.width
- tpb: cutlass.Constexpr = self.tpb
- bt_width: cutlass.Constexpr = self.bt_width
- ld: cutlass.Constexpr = self.ld
- dim: cutlass.Constexpr = HEAD_DIM
- heads = cute.size(mQ, mode=[1])
- kv_heads = heads // group
- rows = cute.size(mSel, mode=[0])
- row = bidx // (kv_heads * splits)
- kvh = (bidx // splits) % kv_heads
- split = bidx % splits
- c0 = split * CAP
-
- smem = cutlass.utils.SmemAllocator()
- sQ = smem.allocate_tensor(
- cutlass.BFloat16, cute.make_layout((M, dim), stride=(ld, 1)), byte_alignment=128
- )
- sK = smem.allocate_tensor(
- cutlass.BFloat16, cute.make_layout((CAP, dim), stride=(ld, 1)), byte_alignment=128
- )
- sV = smem.allocate_tensor(
- cutlass.BFloat16, cute.make_layout((CAP, dim), stride=(ld, 1)), byte_alignment=128
- )
- sP = smem.allocate_tensor(
- cutlass.BFloat16, cute.make_layout((M, CAP), stride=(CAP + PAD, 1)), byte_alignment=128
- )
- sS = smem.allocate_tensor(
- cutlass.Float32, cute.make_layout((M, CAP), stride=(CAP + 4, 1)), byte_alignment=128
- )
- sStat = smem.allocate_tensor(
- cutlass.Float32, cute.make_layout((M, 2), stride=(2, 1)), byte_alignment=128
- )
- sKo = smem.allocate_tensor(
- cutlass.Int64, cute.make_layout((CAP,), stride=(1,)), byte_alignment=128
- )
- sVo = smem.allocate_tensor(
- cutlass.Int64, cute.make_layout((CAP,), stride=(1,)), byte_alignment=128
- )
-
- # 1. Before the wait: the KV head's query heads (complete before the grid ahead of
- # this one released it) into shared memory and the A fragments, the block-table row
- # (host-filled), the release.
- cpr: cutlass.Constexpr = dim // 8 # 16-byte pieces a row: one a lane
- q_vec = cute.logical_divide(mQ, (None, None, 8))
- sq_vec = cute.logical_divide(sQ, (None, 8))
- for i in cutlass.range_constexpr(-(-(group * cpr) // THREADS)):
- c = tidx + i * THREADS
- cc = cutlass.min(c, group * cpr - 1)
- h = cc // cpr
- ch = cc % cpr
- if c < group * cpr:
- cute.copy(g2s, q_vec[row, kvh * group + h, (None, ch)], sq_vec[h, (None, ch)])
- cute.arch.cp_async_commit_group()
- for i in cutlass.range_constexpr((M - group) * dim // THREADS):
- e = tidx + i * THREADS
- sQ[group + e // dim, e % dim] = cutlass.BFloat16(0.0)
- req = mReq[row]
- req_ok = (req >= 0) & (req < self.num_requests)
- safe_req = cutlass.min(cutlass.max(req, 0), self.num_requests - 1)
- if const_expr(self.stage_bt):
- sBT = smem.allocate_tensor(
- cutlass.Int32, cute.make_layout((bt_width,), stride=(1,)), byte_alignment=128
- )
- for i in cutlass.range_constexpr(-(-bt_width // THREADS)):
- e = tidx + i * THREADS
- ec = cutlass.min(e, bt_width - 1)
- v = mBT[safe_req, ec]
- if e < bt_width:
- sBT[e] = v
- cute.arch.cp_async_wait_group(0)
- cute.arch.sync_threads()
- thr = mma.get_slice(tidx)
- tCsQ = thr.partition_A(sQ)
- tCrQ = thr.make_fragment_A(tCsQ)
- cute.autovec_copy(tCsQ, tCrQ)
- if const_expr(self.pdl):
- # The merge waits before reading the partials and reads nothing before its wait
- # that the grid ahead of this one writes.
- cute.arch.griddepcontrol_launch_dependents()
- cute.arch.griddepcontrol_wait()
-
- # 2. The CTA's tokens: each one's K and V row offsets in elements, or -1.
- ks0, ks1, ks2, _ = self.k_strides
- vs0, vs1, vs2, _ = self.v_strides
- if tidx < CAP:
- col = c0 + tidx
- colc = cutlass.min(col, width - 1)
- tok = mSel[row, colc]
- ok = req_ok & (col < width) & (tok >= 0)
- safe_tok = cutlass.max(tok, 0)
- lp = safe_tok // tpb
- ok = ok & (lp < bt_width)
- lpc = cutlass.min(lp, bt_width - 1)
- page = cutlass.Int32(0)
- if const_expr(self.stage_bt):
- page = sBT[lpc]
- else:
- page = mBT[safe_req, lpc]
- ok = ok & (page >= 0) & (page < self.num_pages)
- pg = cutlass.Int64(cutlass.max(page, 0))
- tip = cutlass.Int64(safe_tok % tpb)
- hk = cutlass.Int64(kvh)
- sKo[tidx] = cutlass.Int64(
- cutlass.select_(ok, pg * ks0 + hk * ks1 + tip * ks2, cutlass.Int64(-1))
- )
- sVo[tidx] = cutlass.Int64(
- cutlass.select_(ok, pg * vs0 + hk * vs1 + tip * vs2, cutlass.Int64(-1))
- )
- cute.arch.sync_threads()
-
- # Every token's K and V row: warp w the tokens 8w .. 8w + 7, a lane 16 bytes of each;
- # an invalid token's rows are zeros (its score is masked; V must stay finite).
- sk_vec = cute.logical_divide(sK, (None, 8))
- sv_vec = cute.logical_divide(sV, (None, 8))
- k_base = pK.toint() + cutlass.Int64(lane * 16)
- v_base = pV.toint() + cutlass.Int64(lane * 16)
- tpw: cutlass.Constexpr = CAP // WARPS
- for i in cutlass.range_constexpr(tpw):
- t = warp_i * tpw + i
- ko = sKo[t]
- vo = sVo[t]
- ksrc = cute.make_tensor(
- cute.make_ptr(
- cutlass.BFloat16,
- k_base + cutlass.max(ko, 0) * 2,
- cute.AddressSpace.gmem,
- assumed_align=16,
- ),
- cute.make_layout((8,), stride=(1,)),
- )
- vsrc = cute.make_tensor(
- cute.make_ptr(
- cutlass.BFloat16,
- v_base + cutlass.max(vo, 0) * 2,
- cute.AddressSpace.gmem,
- assumed_align=16,
- ),
- cute.make_layout((8,), stride=(1,)),
- )
- if ko >= 0:
- cute.copy(g2s, ksrc, sk_vec[t, (None, lane)])
- cute.copy(g2s, vsrc, sv_vec[t, (None, lane)])
- else:
- for e in cutlass.range_constexpr(8):
- sK[t, lane * 8 + e] = cutlass.BFloat16(0.0)
- sV[t, lane * 8 + e] = cutlass.BFloat16(0.0)
- cute.arch.cp_async_commit_group()
- cute.arch.cp_async_wait_group(0)
- cute.arch.sync_threads()
-
- # 3. Scores: a warp eight tokens over the whole head dim, masked into shared memory.
- tCsK = thr.partition_B(sK)
- tCrK = thr.make_fragment_B(tCsK)
- cute.autovec_copy(tCsK, tCrK)
- scores = cute.make_rmem_tensor(thr.partition_shape_C((M, CAP)), cutlass.Float32)
- scores.fill(0.0)
- cute.gemm(mma, scores, tCrQ, tCrK, scores)
- cs = thr.partition_C(cute.make_identity_tensor((M, CAP)))
- neg_inf = cutlass.Float32(-math.inf)
- for j in cutlass.range_constexpr(cute.size(scores)):
- h, t = cs[j]
- v = neg_inf
- if sKo[t] >= 0:
- v = scores[j] * self.scale
- sS[h, t] = v
- cute.arch.sync_threads()
- # The softmax: warp w the heads w and w + WARPS, a lane two tokens; padding heads
- # get P = 0.
- for j in cutlass.range_constexpr(M // WARPS):
- h = warp_i + WARPS * j
- a = sS[h, lane]
- b = sS[h, lane + 32]
- mx = cute.arch.warp_reduction_max(cute.arch.fmax(a, b))
- safe_mx = cutlass.Float32(cutlass.select_(mx == neg_inf, cutlass.Float32(0.0), mx))
- keep = h < group
- pa = cutlass.Float32(
- cutlass.select_(keep, cute.math.exp2(a - safe_mx), cutlass.Float32(0.0))
- )
- pb = cutlass.Float32(
- cutlass.select_(keep, cute.math.exp2(b - safe_mx), cutlass.Float32(0.0))
- )
- den = cute.arch.warp_reduction_sum(pa + pb)
- sP[h, lane] = pa.to(cutlass.BFloat16)
- sP[h, lane + 32] = pb.to(cutlass.BFloat16)
- if lane == 0:
- sStat[h, 0] = mx
- sStat[h, 1] = den
- cute.arch.sync_threads()
-
- # P V: a warp 32 of the head dims (four 8-wide tiles).
- tCsP = thr.partition_A(sP)
- tCrP = thr.make_fragment_A(tCsP)
- cute.autovec_copy(tCsP, tCrP)
- sVt = cute.make_tensor(sV.iterator, cute.make_layout((dim, CAP), stride=(1, ld)))
- tCsV = thr.partition_B(sVt)
- tCrV = thr.make_fragment_B(tCsV)
- cute.autovec_copy(tCsV, tCrV)
- acc = cute.make_rmem_tensor(thr.partition_shape_C((M, dim)), cutlass.Float32)
- acc.fill(0.0)
- cute.gemm(mma, acc, tCrP, tCrV, acc)
-
- # Each head's partial times its reciprocal sum, staged over sK (free since the
- # softmax barrier) as [group, dim] fp32, then stored in 16-byte pieces: the CTA's
- # heads are contiguous in the partial tensor.
- sOut = cute.make_tensor(
- cute.recast_ptr(sK.iterator, dtype=cutlass.Float32),
- cute.make_layout((group, dim), stride=(dim, 1)),
- )
- co = thr.partition_C(cute.make_identity_tensor((M, dim)))
- h0 = lane // 4 # an m16n8 C fragment's rows: h0 and h0 + 8
- d0 = sStat[h0, 1]
- d1 = sStat[h0 + 8, 1]
- r0 = cutlass.Float32(0.0)
- r1 = cutlass.Float32(0.0)
- if d0 > 0.0:
- r0 = 1.0 / d0
- if d1 > 0.0:
- r1 = 1.0 / d1
- for j in cutlass.range_constexpr(cute.size(acc)):
- h, d = co[j]
- r = cutlass.Float32(cutlass.select_(h < 8, r0, r1))
- if h < group:
- sOut[h, d] = acc[j] * r
- cute.arch.sync_threads()
- base = ((cutlass.Int64(split) * rows + row) * heads + kvh * group) * dim
- out_ptr = pPart.toint() + base * 4
- so_vec = cute.make_tensor(
- sOut.iterator, cute.make_layout((group * dim // 4, 4), stride=(4, 1))
- )
- reg = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), cutlass.Float32)
- for i in cutlass.range_constexpr(-(-(group * dim // 4) // THREADS)):
- c = tidx + i * THREADS
- if c < group * dim // 4:
- cute.autovec_copy(so_vec[c, None], reg)
- dst = cute.make_tensor(
- cute.make_ptr(
- cutlass.Float32,
- out_ptr + cutlass.Int64(c * 16),
- cute.AddressSpace.gmem,
- assumed_align=16,
- ),
- cute.make_layout((4,), stride=(1,)),
- )
- cute.copy(st, reg, dst)
- if tidx < group:
- mx = sStat[tidx, 0]
- den = sStat[tidx, 1]
- lse = neg_inf
- if den > 0.0:
- lse = mx + cute.math.log2(den)
- mLse[split, row, kvh * group + tidx] = lse
-
-
-class _MergeKernel:
- def __init__(self, *, heads, dim, splits, gated, pdl):
- assert dim % CHUNK == 0 and 1 < splits <= MAX_SPLITS
- self.heads, self.dim, self.splits, self.gated, self.pdl = heads, dim, splits, gated, pdl
- self.chunks = dim // CHUNK
- self.warps = min(8, splits)
- self.spw = -(-splits // self.warps) # splits a warp, rounded up
-
- @cute.jit
- def __call__(self, mPart, mLse, mGate, mOut, stream: cuda_drv.CUstream):
- rows = cute.size(mOut, mode=[0])
- ld = cute.make_copy_atom(cute.nvgpu.CopyG2ROp(), cutlass.Float32, num_bits_per_copy=128)
- self.kernel(mPart, mLse, mGate, mOut, ld).launch(
- grid=[rows * self.heads * self.chunks, 1, 1],
- block=[self.warps * 32, 1, 1],
- smem=self.warps * CHUNK * 4 + 64,
- stream=stream,
- use_pdl=self.pdl,
- )
-
- @cute.kernel
- def kernel(self, mPart, mLse, mGate, mOut, ld: cute.CopyAtom):
- tidx, _, _ = cute.arch.thread_idx()
- bidx, _, _ = cute.arch.block_idx()
- warp_i = cute.arch.warp_idx()
- lane = cute.arch.lane_idx()
- per_row: cutlass.Constexpr = self.heads * self.chunks
- spw: cutlass.Constexpr = self.spw
- row = bidx // per_row
- head = (bidx % per_row) // self.chunks
- c0 = (bidx % self.chunks) * CHUNK
- red = cutlass.utils.SmemAllocator().allocate_tensor(
- cutlass.Float32,
- cute.make_layout((self.warps, CHUNK), stride=(CHUNK, 1)),
- byte_alignment=16,
- )
- # The gate, before the wait: see the module notes.
- g = cutlass.Float32(0.0)
- if const_expr(self.gated):
- if tidx < CHUNK:
- g = mGate[row, head * self.dim + c0 + tidx].to(cutlass.Float32)
- if const_expr(self.pdl):
- # The output projection after this reads only its weight before its own wait.
- cute.arch.griddepcontrol_launch_dependents()
- cute.arch.griddepcontrol_wait()
-
- # 1. The split weights, in every warp: exp2(lse - max), and their sum.
- neg_inf = cutlass.Float32(-math.inf)
- l0 = neg_inf
- l1 = neg_inf
- if lane < self.splits:
- l0 = mLse[lane, row, head]
- if lane + 32 < self.splits:
- l1 = mLse[lane + 32, row, head]
- # 2. Warp w's splits w * spw .. w * spw + spw - 1, four adjacent elements a lane:
- # one 16-byte load a split, all in flight. A slot past the split count reads the
- # last split again and weighs it 0.
- part_vec = cute.logical_divide(mPart, (None, None, None, 4))
- p = cute.make_rmem_tensor(cute.make_layout((spw, 4), stride=(4, 1)), cutlass.Float32)
- for j in cutlass.range_constexpr(spw):
- sidx = cutlass.min(warp_i * spw + j, self.splits - 1)
- cute.copy(ld, part_vec[sidx, row, head, (None, (c0 + 4 * lane) // 4)], p[j, None])
- mx = cute.arch.warp_reduction_max(cute.arch.fmax(l0, l1))
- w0 = cutlass.Float32(0.0)
- w1 = cutlass.Float32(0.0)
- if mx > neg_inf:
- w0 = cute.math.exp2(l0 - mx)
- w1 = cute.math.exp2(l1 - mx)
- den = cute.arch.warp_reduction_sum(w0 + w1)
- acc = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), cutlass.Float32)
- acc.fill(0.0)
- for j in cutlass.range_constexpr(spw):
- sidx = warp_i * spw + j
- wa = cute.arch.shuffle_sync(w0, sidx % 32)
- wb = cute.arch.shuffle_sync(w1, sidx % 32)
- w = cutlass.Float32(cutlass.select_(sidx < 32, wa, wb))
- w = cutlass.Float32(cutlass.select_(sidx < self.splits, w, cutlass.Float32(0.0)))
- for e in cutlass.range_constexpr(4):
- acc[e] = w * p[j, e] + acc[e]
- for e in cutlass.range_constexpr(4):
- red[warp_i, 4 * lane + e] = acc[e]
- cute.arch.sync_threads()
-
- # 3. The warps' shares summed in warp order; the output gate on the rounded
- # attention value, as the composition applies it.
- if tidx < CHUNK:
- tot = cutlass.Float32(0.0)
- for w_ in cutlass.range_constexpr(self.warps):
- tot = tot + red[w_, tidx]
- merged = cutlass.Float32(0.0)
- if den > 0.0:
- merged = tot / den
- if const_expr(self.gated):
- att = merged.to(cutlass.BFloat16).to(cutlass.Float32)
- merged = att * (1.0 / (1.0 + cute.math.exp(-g)))
- mOut[row, head, c0 + tidx] = merged.to(mOut.element_type)
-
-
-def _dl(t: torch.Tensor, align: int = 4):
- return from_dlpack(t, assumed_align=align)
-
-
-def _ptr(t: torch.Tensor, dtype):
- return make_ptr(dtype, t.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
-
-
-def splits(q, k_cache, selected_tokens) -> int:
- """The split-K CTAs a row and KV head: one per `CAP` selected tokens."""
- return -(-selected_tokens.shape[1] // CAP)
-
-
-@functools.lru_cache(maxsize=None)
-def _compiled_splitk(
- rows,
- q_shape,
- q_stride,
- group,
- width,
- tpb,
- scale,
- bt_shape,
- num_pages,
- k_strides,
- v_strides,
- pdl,
-):
- heads = q_shape[1]
- n_splits = -(-width // CAP)
- kern = _SplitKKernel(
- group=group,
- width=width,
- splits=n_splits,
- tpb=tpb,
- scale=scale,
- bt_width=bt_shape[1],
- num_pages=num_pages,
- k_strides=k_strides,
- v_strides=v_strides,
- num_requests=bt_shape[0],
- pdl=pdl,
- )
- dev = "cuda"
- q = torch.empty_strided(q_shape, q_stride, dtype=torch.bfloat16, device=dev)
- kv = torch.empty(64, dtype=torch.bfloat16, device=dev)
- bt = torch.empty(bt_shape, dtype=torch.int32, device=dev)
- sel = torch.empty(rows, width, dtype=torch.int32, device=dev)
- req = torch.empty(rows, dtype=torch.int32, device=dev)
- part = torch.empty(64, dtype=torch.float32, device=dev)
- lse = torch.empty(n_splits, rows, heads, dtype=torch.float32, device=dev)
- stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream)
- return cute_ext.compile(
- kern,
- _dl(q, 16),
- _ptr(kv, cutlass.BFloat16),
- _ptr(kv, cutlass.BFloat16),
- _dl(bt),
- _dl(sel),
- _dl(req),
- _ptr(part, cutlass.Float32),
- _dl(lse),
- stream,
- )
-
-
-@functools.lru_cache(maxsize=None)
-def _compiled_merge(rows, heads, dim, n_splits, gate_stride, dtype, pdl):
- kern = _MergeKernel(
- heads=heads, dim=dim, splits=n_splits, gated=gate_stride is not None, pdl=pdl
- )
- dev = "cuda"
- part = torch.empty(n_splits, rows, heads, dim, dtype=torch.float32, device=dev)
- lse = torch.empty(n_splits, rows, heads, dtype=torch.float32, device=dev)
- gate = torch.empty(rows, gate_stride or heads * dim, dtype=dtype, device=dev)[:, : heads * dim]
- out = torch.empty(rows, heads, dim, dtype=dtype, device=dev)
- stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream)
- return cute_ext.compile(kern, _dl(part, 16), _dl(lse), _dl(gate, 2), _dl(out, 2), stream)
-
-
-def launch(
- q,
- k_cache,
- v_cache,
- block_table,
- selected_tokens,
- request_indices,
- tokens_per_block,
- softmax_scale,
- output_gate=None,
- pdl=True,
-):
- """The launch behind `QsaPagedSparseGqaDecode.__call__`; returns `[rows, H, D]` in q's dtype."""
- rows, num_q_heads, head_dim = q.shape
- num_kv_heads = k_cache.shape[1]
- group_size = num_q_heads // num_kv_heads
- width = selected_tokens.shape[1]
- n_splits = -(-width // CAP)
- output = torch.empty((rows, num_q_heads, head_dim), dtype=q.dtype, device=q.device)
- if rows == 0:
- return output
- dev = q.device
- use_pdl = bool(pdl and _qsa_pdl_enabled(rows))
- partial_output = torch.empty(
- (n_splits, rows, num_q_heads, head_dim), dtype=torch.float32, device=dev
- )
- partial_lse = torch.empty((n_splits, rows, num_q_heads), dtype=torch.float32, device=dev)
- fn = _compiled_splitk(
- rows,
- tuple(q.shape),
- tuple(q.stride()),
- group_size,
- width,
- tokens_per_block,
- float(softmax_scale),
- tuple(block_table.shape),
- k_cache.shape[0],
- tuple(k_cache.stride()),
- tuple(v_cache.stride()),
- use_pdl,
- )
- stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream)
- fn(
- _dl(q, 16),
- _ptr(k_cache, cutlass.BFloat16),
- _ptr(v_cache, cutlass.BFloat16),
- _dl(block_table),
- _dl(selected_tokens),
- _dl(request_indices),
- _ptr(partial_output, cutlass.Float32),
- _dl(partial_lse),
- stream,
- )
- gate2d = None if output_gate is None else output_gate.reshape(rows, num_q_heads * head_dim)
- mfn = _compiled_merge(
- rows,
- num_q_heads,
- head_dim,
- n_splits,
- None if gate2d is None else gate2d.stride(0),
- q.dtype,
- use_pdl,
- )
- mfn(
- _dl(partial_output, 16),
- _dl(partial_lse),
- _dl(output.view(rows, num_q_heads * head_dim) if gate2d is None else gate2d, 2),
- _dl(output, 2),
- stream,
- )
- return output
-
-
-def supports(q, k_cache, selected_tokens) -> bool:
- """Whether these operands are ones this entry serves (see `is_valid`)."""
- return (
- q.dim() == 3
- and 0 < q.shape[0] <= MAX_ROWS
- and q.shape[2] == HEAD_DIM
- and q.dtype == torch.bfloat16
- and k_cache.dtype == torch.bfloat16
- and q.shape[1] % k_cache.shape[1] == 0
- and q.shape[1] // k_cache.shape[1] <= M
- and 1 < splits(q, k_cache, selected_tokens) <= MAX_SPLITS
- and all(t.data_ptr() % 16 == 0 for t in (q, k_cache))
- and all(s % 8 == 0 for s in (*q.stride()[:2], *k_cache.stride()[:3]))
- )
-
-
-class QsaPagedSparseGqaDecode(OpWrapper):
- """`qsa_paged_sparse_gqa` without a prefill pair, for decode rows: a split-K pass whose
- CTAs load their tokens' K/V in one round trip, then the merge -- and, with
- `output_gate`, the gate -- as one launch. Returns a new `[rows, H, D]` tensor in q's
- dtype.
-
- Fusion boundary: the split-K attention, the merge, the optional sigmoid output gate.
- The selection, the K/V store and the output projection are the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(why="decode, 1 row at position 8700, gated", spec=dict(pos=(8700,), gated=True)),
- Cell(
- why="decode, 4 rows, one a short sequence, gated",
- spec=dict(pos=(8700, 5000, 77, 3000), gated=True),
- ),
- Cell(
- why="decode, 1 row with half the selection -1",
- spec=dict(pos=(8700,), gated=True, holes=True),
- ),
- )
-
- note = """
- Metadata consumed: as `qsa_paged_sparse_gqa` (block table, request indices). Preconditions:
- as that entry's, plus: no prefill pair; at most `MAX_ROWS` rows; bf16 q and K/V with a head
- dim of 256 and at most 16 query heads a KV head; more than one and at most 64 split-K CTAs
- a row and KV head (`splits(...)`: 64 selected tokens a CTA); `output_gate` (when given)
- `[rows, H * D]` or `[rows, H, D]` with unit inner stride in q's dtype. `supports(...)` is
- the predicate a caller can check.
-
- PDL: the split-K pass reads q and the block table before its wait and the merge reads the
- output gate before its own; the selection and the K/V rows are read after the split-K
- pass's wait. So q and the gate must be complete when the grid ahead of this entry releases
- it: a launch without PDL releases at its end; a PDL producer of the selection (the
- target's top-k) must release only after its own wait.
-
- Numerics: one softmax over a CTA's 64 tokens (QSA's kernel runs an online softmax over
- 16-token tiles) and the merge sums the splits in split order (QSA's merge reduces them as
- a tree), so the fp32 values can differ in the last bits and the bf16 output by an ulp on
- some elements.
-
- Not certified: a call without `output_gate` (the target always passes one).
- """
-
- def __call__(
- self,
- q: torch.Tensor,
- k_cache: torch.Tensor,
- v_cache: torch.Tensor,
- block_table: torch.Tensor,
- selected_tokens: torch.Tensor,
- request_indices: torch.Tensor,
- tokens_per_block: int,
- softmax_scale: float,
- output_gate: Optional[torch.Tensor] = None,
- ) -> torch.Tensor:
- return launch(
- q,
- k_cache,
- v_cache,
- block_table,
- selected_tokens,
- request_indices,
- tokens_per_block,
- softmax_scale,
- output_gate,
- )
-
- def reference(
- self,
- q: torch.Tensor,
- k_cache: torch.Tensor,
- v_cache: torch.Tensor,
- block_table: torch.Tensor,
- selected_tokens: torch.Tensor,
- request_indices: torch.Tensor,
- tokens_per_block: int,
- softmax_scale: float,
- output_gate: Optional[torch.Tensor] = None,
- ) -> torch.Tensor:
- """The composition entry's reference with the same operands."""
- return _composition.reference(
- q=q,
- k_cache=k_cache,
- v_cache=v_cache,
- block_table=block_table,
- selected_tokens=selected_tokens,
- request_indices=request_indices,
- tokens_per_block=tokens_per_block,
- softmax_scale=softmax_scale,
- query_positions=None,
- compress_ratio=None,
- output_gate=output_gate,
- )
-
- def is_valid(
- self,
- q: torch.Tensor,
- k_cache: torch.Tensor,
- v_cache: torch.Tensor,
- block_table: torch.Tensor,
- selected_tokens: torch.Tensor,
- request_indices: torch.Tensor,
- tokens_per_block: int,
- softmax_scale: float,
- output_gate: Optional[torch.Tensor] = None,
- ) -> None:
- assert q.ndim == 3 and q.shape[2] == HEAD_DIM, "q [rows, H, 256]"
- assert 0 < q.shape[0] <= MAX_ROWS, f"1 .. {MAX_ROWS} rows"
- assert q.dtype == torch.bfloat16 and k_cache.dtype == v_cache.dtype == torch.bfloat16, (
- "bf16 q and K/V"
- )
- assert q.shape[1] % k_cache.shape[1] == 0 and q.shape[1] // k_cache.shape[1] <= M, (
- "<= 16 q heads a KV head"
- )
- assert 1 < splits(q, k_cache, selected_tokens) <= MAX_SPLITS, (
- "65 .. 4096 selected tokens a row"
- )
- assert q.stride(2) == 1 and k_cache.stride(3) == 1 and v_cache.stride(3) == 1, (
- "unit inner strides"
- )
- assert all(t.data_ptr() % 16 == 0 for t in (q, k_cache, v_cache)), (
- "16-byte aligned q and K/V"
- )
- assert all(
- s % 8 == 0 for s in (*q.stride()[:2], *k_cache.stride()[:3], *v_cache.stride()[:3])
- ), "q and K/V rows 16-byte aligned"
- if output_gate is not None:
- assert output_gate.dtype == q.dtype and output_gate.stride(-1) == 1, "gate in q's dtype"
- assert output_gate.numel() == q.numel(), "gate covers q"
-
-
-qsa_paged_sparse_gqa_decode = QsaPagedSparseGqaDecode()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py
deleted file mode 100644
index 75e9c55e17a8..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py
+++ /dev/null
@@ -1,1106 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""QSA sparse GQA over an NVFP4 paged pool: attention over explicitly selected tokens."""
-
-from __future__ import annotations
-
-from typing import Optional
-
-import torch
-import triton
-import triton.language as tl
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-# Approved exception (as for `qsa_paged_sparse_gqa`): QSA's sparse path has no torch.ops
-# surface; its row-count PDL policy is reached as Python. The split-K buckets are this
-# entry's own: an NVFP4 tile needs 16 tokens where the bf16 one takes 64.
-from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import _qsa_pdl_enabled
-
-#: 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,
-)
-
-#: The prefill kernel's (BLOCK_N, warps): on GB300 8192 context rows 6.5 ms, 4096 rows 2.9 ms.
-_PREFILL_TILE = (32, 2)
-
-_SPLIT_K = ("splitk", "merge")
-_PREFILL = ("prefill",)
-
-
-def _spec(**overrides) -> dict:
- """A cell spec: this checkpoint's geometry unless the cell says otherwise.
-
- `segments` is `(first position, rows)` per request, in row order; `prefill_pair` says whether `query_positions`
- / `compress_ratio` are passed, which the target does exactly when the step has a
- context request; `pool_pages` fixes the pool's page count (the requests then sit
- on its last pages); `capture` sweeps the decode site under a CUDA graph.
- """
- return {**_CHECKPOINT, **overrides}
-
-
-def _decode(segments, launches=_SPLIT_K, **overrides) -> dict:
- base = dict(segments=segments, slots=32, prefill_pair=False, launches=launches)
- return _spec(**{**base, **overrides})
-
-
-@triton.jit
-def _e2m1_to_f32(code):
- """e2m1 code (0..15, bit 3 the sign) -> its value, exactly."""
- code = code.to(tl.int32)
- magnitude_code = code & 7
- exponent = magnitude_code >> 1
- mantissa = (magnitude_code & 1).to(tl.float32)
- magnitude = tl.where(
- exponent == 0,
- 0.5 * mantissa,
- (2.0 + mantissa) * (1 << exponent).to(tl.float32) * 0.25,
- )
- return tl.where((code & 8) != 0, -magnitude, magnitude)
-
-
-@triton.jit
-def _nvfp4_halves(packed, scales):
- """`[N, D/2]` packed e2m1 bytes (low nibble the even element) and one e4m3 scale per byte
- -> the even and the odd elements as two `[N, D/2]` f16 tensors, each `e2m1 * scale`:
- `cvt.rn.f16x2.e2m1x2` and `cvt.rn.f16x2.e4m3x2`, then one f16x2 multiply, exact (two by
- four significant bits, inside f16's range)."""
- return tl.inline_asm_elementwise(
- asm="""
- {
- .reg .b8 b0, b1, b2, b3, s0, s1, s2, s3;
- .reg .b16 t0, t1, t2, t3;
- .reg .b32 h0, h1, h2, h3, c0, c1, c2, c3;
- mov.b32 {b0, b1, b2, b3}, $4;
- mov.b32 {s0, s1, s2, s3}, $5;
- mov.b16 t0, {s0, s0};
- mov.b16 t1, {s1, s1};
- mov.b16 t2, {s2, s2};
- mov.b16 t3, {s3, s3};
- cvt.rn.f16x2.e4m3x2 c0, t0;
- cvt.rn.f16x2.e4m3x2 c1, t1;
- cvt.rn.f16x2.e4m3x2 c2, t2;
- cvt.rn.f16x2.e4m3x2 c3, t3;
- cvt.rn.f16x2.e2m1x2 h0, b0;
- cvt.rn.f16x2.e2m1x2 h1, b1;
- cvt.rn.f16x2.e2m1x2 h2, b2;
- cvt.rn.f16x2.e2m1x2 h3, b3;
- mul.rn.f16x2 h0, h0, c0;
- mul.rn.f16x2 h1, h1, c1;
- mul.rn.f16x2 h2, h2, c2;
- mul.rn.f16x2 h3, h3, c3;
- prmt.b32 $0, h0, h1, 0x5410;
- prmt.b32 $1, h2, h3, 0x5410;
- prmt.b32 $2, h0, h1, 0x7632;
- prmt.b32 $3, h2, h3, 0x7632;
- }
- """,
- constraints="=r,=r,=r,=r,r,r",
- args=[packed, scales],
- dtype=(tl.float16, tl.float16),
- is_pure=True,
- pack=4,
- )
-
-
-@triton.jit
-def _load_halves(
- cache,
- scale_rows,
- pages,
- kv_head,
- tokens,
- valid,
- stride_page,
- stride_head,
- stride_token,
- SCALE_STEP: tl.constexpr,
- BLOCK_N: tl.constexpr,
- HEAD_DIM: tl.constexpr,
-):
- """A tile's even and odd elements of one role plane. `scale_rows` points at each token's
- first scale; its D/16 scales lie SCALE_STEP apart (K's token-major rows: 1; V's
- four-token groups `[tpb/4, D/16, 4]`: 4). Each byte and each scale is read once."""
- HALF: tl.constexpr = HEAD_DIM // 2
- GROUPS: tl.constexpr = HEAD_DIM // 16
- packed = tl.load(
- cache
- + pages[:, None] * stride_page
- + kv_head * stride_head
- + tokens[:, None] * stride_token
- + tl.arange(0, HALF)[None, :],
- mask=valid[:, None],
- other=0,
- )
- scales = tl.load(
- scale_rows[:, None] + tl.arange(0, GROUPS)[None, :] * SCALE_STEP,
- mask=valid[:, None],
- other=0,
- )
- scales = tl.reshape(tl.broadcast_to(scales[:, :, None], [BLOCK_N, GROUPS, 8]), [BLOCK_N, HALF])
- return _nvfp4_halves(packed, scales)
-
-
-@triton.jit
-def _attend_tile(
- q_even,
- q_odd,
- running_max,
- running_sum,
- acc_even,
- acc_odd,
- columns,
- row,
- safe_request,
- valid_request,
- kv_head,
- score_scale,
- selected_tokens,
- selected_stride_row,
- block_table,
- block_table_stride_request,
- page_table_width,
- num_cache_pages,
- k_cache,
- v_cache,
- k_scale_cache,
- v_scale_cache,
- 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,
- TOKENS_PER_BLOCK: tl.constexpr,
- TOPK: tl.constexpr,
- BLOCK_N: tl.constexpr,
- HEAD_DIM: tl.constexpr,
-):
- """One online-softmax step over the selected columns `columns` of `row`: the even and odd
- halves of K against q's, scores in the log2 domain, P rounded to f16 against V's halves."""
- GROUPS: tl.constexpr = HEAD_DIM // 16
- logical_tokens = tl.load(
- selected_tokens + row * selected_stride_row + columns, mask=columns < TOPK, other=-1
- )
- safe_tokens = tl.maximum(logical_tokens, 0)
- logical_pages = safe_tokens // TOKENS_PER_BLOCK
- tokens = safe_tokens % TOKENS_PER_BLOCK
- valid = (
- valid_request
- & (columns < TOPK)
- & (logical_tokens >= 0)
- & (logical_pages < page_table_width)
- )
- 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 &= (pages >= 0) & (pages < num_cache_pages)
- pages = tl.maximum(pages, 0)
-
- k_even, k_odd = _load_halves(
- k_cache,
- k_scale_cache + pages * ks_stride_page + kv_head * ks_stride_head + tokens * GROUPS,
- pages,
- kv_head,
- tokens,
- valid,
- k_stride_page,
- k_stride_head,
- k_stride_token,
- 1,
- BLOCK_N,
- HEAD_DIM,
- )
- scores = tl.dot(q_even, tl.trans(k_even))
- scores = tl.dot(q_odd, tl.trans(k_odd), scores)
- scores = tl.where(valid[None, :], scores * score_scale, -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)
-
- v_even, v_odd = _load_halves(
- v_cache,
- v_scale_cache
- + pages * vs_stride_page
- + kv_head * vs_stride_head
- + (tokens // 4) * (4 * GROUPS)
- + tokens % 4,
- pages,
- kv_head,
- tokens,
- valid,
- v_stride_page,
- v_stride_head,
- v_stride_token,
- 4,
- BLOCK_N,
- HEAD_DIM,
- )
- weights = probabilities.to(tl.float16)
- acc_even = tl.dot(weights, v_even, acc_even * correction[:, None])
- acc_odd = tl.dot(weights, v_odd, acc_odd * correction[:, None])
- running_sum = running_sum * correction + tl.sum(probabilities, axis=1)
- return next_max, running_sum, acc_even, acc_odd
-
-
-@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)
- # The global scales scale the scores and the normalized output, not each element.
- score_scale = softmax_scale * _LOG2_E * tl.load(kv_scale_quant_orig + 1)
- v_global = tl.load(kv_scale_quant_orig + 2)
-
- HALF: tl.constexpr = HEAD_DIM // 2
- head_offsets = tl.arange(0, BLOCK_M)
- half_offsets = tl.arange(0, HALF)
- token_offsets = tl.arange(0, BLOCK_N)
- query_heads = kv_head * GROUP_SIZE + head_offsets
- head_mask = (head_offsets < GROUP_SIZE)[:, None]
- # q's even and odd elements as f16, so both MMA operands are f16 (q is bf16 within
- # f16's range after its norm; an NVFP4 value times its scale is exact in f16).
- q_even_ptr = (
- q + row * q_stride_row + query_heads[:, None] * q_stride_head + 2 * half_offsets[None, :]
- )
- q_even = tl.load(q_even_ptr, mask=head_mask, other=0.0).to(tl.float16)
- q_odd = tl.load(q_even_ptr + 1, mask=head_mask, other=0.0).to(tl.float16)
-
- running_max = tl.full([BLOCK_M], -float("inf"), tl.float32)
- running_sum = tl.zeros([BLOCK_M], tl.float32)
- acc_even = tl.zeros([BLOCK_M, HALF], tl.float32)
- acc_odd = tl.zeros([BLOCK_M, HALF], tl.float32)
- 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):
- columns = tile * BLOCK_N + token_offsets
- running_max, running_sum, acc_even, acc_odd = _attend_tile(
- q_even,
- q_odd,
- running_max,
- running_sum,
- acc_even,
- acc_odd,
- columns,
- row,
- safe_request,
- valid_request,
- kv_head,
- score_scale,
- selected_tokens,
- selected_stride_row,
- block_table,
- block_table_stride_request,
- page_table_width,
- num_cache_pages,
- k_cache,
- v_cache,
- k_scale_cache,
- v_scale_cache,
- 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,
- TOKENS_PER_BLOCK,
- TOPK,
- BLOCK_N,
- HEAD_DIM,
- )
-
- denominator = tl.maximum(running_sum[:, None], 1.0e-20)
- has_values = running_sum[:, None] > 0
- out_ptr = (
- output
- + row * output_stride_row
- + query_heads[:, None] * output_stride_head
- + 2 * half_offsets[None, :]
- )
- tl.store(out_ptr, tl.where(has_values, acc_even / denominator, 0.0) * v_global, mask=head_mask)
- tl.store(
- out_ptr + 1, tl.where(has_values, acc_odd / denominator, 0.0) * v_global, mask=head_mask
- )
-
-
-@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)
- # The global scales scale the scores and the normalized output, not each element.
- score_scale = softmax_scale * _LOG2_E * tl.load(kv_scale_quant_orig + 1)
- v_global = tl.load(kv_scale_quant_orig + 2)
-
- HALF: tl.constexpr = HEAD_DIM // 2
- head_offsets = tl.arange(0, BLOCK_M)
- half_offsets = tl.arange(0, HALF)
- token_offsets = tl.arange(0, BLOCK_N)
- query_heads = kv_head * GROUP_SIZE + head_offsets
- head_mask = (head_offsets < GROUP_SIZE)[:, None]
- # q's even and odd elements as f16, so both MMA operands are f16 (q is bf16 within
- # f16's range after its norm; an NVFP4 value times its scale is exact in f16).
- q_even_ptr = (
- q + row * q_stride_row + query_heads[:, None] * q_stride_head + 2 * half_offsets[None, :]
- )
- q_even = tl.load(q_even_ptr, mask=head_mask, other=0.0).to(tl.float16)
- q_odd = tl.load(q_even_ptr + 1, mask=head_mask, other=0.0).to(tl.float16)
-
- running_max = tl.full([BLOCK_M], -float("inf"), tl.float32)
- running_sum = tl.zeros([BLOCK_M], tl.float32)
- acc_even = tl.zeros([BLOCK_M, HALF], tl.float32)
- acc_odd = tl.zeros([BLOCK_M, HALF], 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):
- columns = tile * BLOCK_N + token_offsets
- running_max, running_sum, acc_even, acc_odd = _attend_tile(
- q_even,
- q_odd,
- running_max,
- running_sum,
- acc_even,
- acc_odd,
- columns,
- row,
- safe_request,
- valid_request,
- kv_head,
- score_scale,
- selected_tokens,
- selected_stride_row,
- block_table,
- block_table_stride_request,
- page_table_width,
- num_cache_pages,
- k_cache,
- v_cache,
- k_scale_cache,
- v_scale_cache,
- 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,
- TOKENS_PER_BLOCK,
- TOPK,
- BLOCK_N,
- HEAD_DIM,
- )
-
- has_values = running_sum > 0
- denominator = tl.maximum(running_sum[:, None], 1.0e-20)
- even = tl.where(has_values[:, None], acc_even / denominator, 0.0) * v_global
- odd = tl.where(has_values[:, None], acc_odd / denominator, 0.0) * v_global
- if NUM_SPLITS == 1:
- out_ptr = (
- output
- + row * output_stride_row
- + query_heads[:, None] * output_stride_head
- + 2 * half_offsets[None, :]
- )
- else:
- out_ptr = (
- partial_output
- + ((split * num_rows + row) * NUM_QUERY_HEADS + query_heads[:, None]) * HEAD_DIM
- + 2 * half_offsets[None, :]
- )
- partial_lse_values = tl.where(
- has_values,
- running_max + tl.math.log2(tl.maximum(running_sum, 1.0e-20)),
- -float("inf"),
- )
- tl.store(
- partial_lse + (split * num_rows + row) * NUM_QUERY_HEADS + query_heads,
- partial_lse_values,
- mask=head_offsets < GROUP_SIZE,
- )
- tl.store(out_ptr, even, mask=head_mask)
- tl.store(out_ptr + 1, odd, mask=head_mask)
- 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) by programs = rows x Hkv, the fastest of a sweep on GB300 at
- 2051 selected tokens a row (1 row 12.6 us, 128 rows 160 us, 512 rows 560 us)."""
- base_programs = rows * num_kv_heads
- if base_programs <= 4:
- block_n, target_splits, num_warps = 32, 64, 4
- elif base_programs <= 16:
- block_n, target_splits, num_warps = 32, 16, 4
- elif base_programs <= 32:
- block_n, target_splits, num_warps = 32, 8, 4
- elif base_programs <= 64:
- block_n, target_splits, num_warps = 32, 4, 4
- elif base_programs <= 128:
- block_n, target_splits, num_warps = 64, 4, 2
- elif base_programs <= 256:
- block_n, target_splits, num_warps = 32, 8, 2
- elif base_programs <= 512:
- block_n, target_splits, num_warps = 32, 4, 2
- elif base_programs <= 1024:
- block_n, target_splits, num_warps = 32, 2, 2
- else:
- block_n, target_splits, num_warps = 32, 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 _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=_PREFILL_TILE[0],
- HEAD_DIM=head_dim,
- num_warps=_PREFILL_TILE[1],
- 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 _qsa_pdl_enabled(rows)
- if num_splits == 1:
- partial_output = partial_lse = output
- else:
- partial_output = torch.empty((num_splits, *q.shape), dtype=torch.float32, device=q.device)
- partial_lse = torch.empty(
- (num_splits, rows, num_q_heads), dtype=torch.float32, device=q.device
- )
- _nvfp4_sparse_gqa_splitk_kernel[(rows, num_kv_heads, num_splits)](
- q,
- k_cache,
- v_cache,
- k_scale_cache,
- v_scale_cache,
- kv_scale_quant_orig,
- block_table,
- selected_tokens,
- request_indices,
- partial_output,
- partial_lse,
- output,
- softmax_scale,
- rows,
- num_pages,
- block_table.shape[0],
- block_table.shape[1],
- q.stride(0),
- q.stride(1),
- *plane_strides,
- block_table.stride(0),
- selected_tokens.stride(0),
- output.stride(0),
- output.stride(1),
- GROUP_SIZE=group_size,
- NUM_QUERY_HEADS=num_q_heads,
- TOKENS_PER_BLOCK=tokens_per_block,
- TOPK=selected_tokens.shape[1],
- NUM_SPLITS=num_splits,
- NUM_TILES=num_tiles,
- BLOCK_M=block_m,
- BLOCK_N=block_n,
- HEAD_DIM=head_dim,
- USE_PDL=use_pdl,
- num_warps=num_warps,
- num_stages=2,
- launch_pdl=use_pdl,
- )
- if num_splits == 1:
- return output
- _nvfp4_merge_splitk_kernel[(rows, num_q_heads)](
- partial_output,
- partial_lse,
- output,
- rows,
- output.stride(0),
- output.stride(1),
- NUM_QUERY_HEADS=num_q_heads,
- NUM_SPLITS=num_splits,
- BLOCK_SPLITS=triton.next_power_of_2(num_splits),
- HEAD_DIM=head_dim,
- USE_PDL=use_pdl,
- num_warps=2,
- num_stages=1,
- launch_pdl=use_pdl,
- )
- return output
-
-
-def _dequantize_tokens(data, scales, pages, slots, global_scale, grouped: bool) -> torch.Tensor:
- """[..., Hkv, D] fp32 values of the tokens at (pages, slots) of one NVFP4 role plane."""
- heads, tokens_per_block, packed = data.shape[1], data.shape[2], data.shape[3]
- head_dim = packed * _PER_BYTE
- groups = head_dim // _SF_VEC
- raw = data[pages, :, slots, :].long() # [..., H, D/2]
- lut = torch.tensor(_E2M1_VALUES, dtype=torch.float32, device=data.device)
- lut = torch.cat([lut, -lut])
- values = torch.stack([lut[raw & 0xF], lut[raw >> 4]], dim=-1).flatten(-2)
- if grouped:
- flat = scales.reshape(scales.shape[0], heads, tokens_per_block * groups)
- base = (slots // _V_SF_QUAD) * (_V_SF_QUAD * groups) + slots % _V_SF_QUAD
- columns = base[..., None] + _V_SF_QUAD * torch.arange(groups, device=data.device)
- sf = torch.stack([flat[pages[..., None], h, columns] for h in range(heads)], dim=-2)
- else:
- sf = scales[pages, :, slots, :]
- sf = sf.view(torch.float8_e4m3fn).float().repeat_interleave(_SF_VEC, dim=-1)
- return values * sf * global_scale
-
-
-class QsaPagedSparseGqaNvfp4(OpWrapper):
- """Sparse GQA over the selected tokens of an NVFP4 HND paged pool.
-
- Triton kernels defined in this file. For row `r`, query head `h`, KV head
- `h // (Hq // Hkv)` (contiguous groups), over the columns `j` of
- `selected_tokens[r]` that name a resident token::
-
- t = selected_tokens[r, j] # logical position
- page = block_table[request_indices[r], t // tokens_per_block]
- K[t] = e2m1(k_cache) * e4m3(k_scale_cache) # exact in f16
- V[t] = e2m1(v_cache) * e4m3(v_scale_cache)
- s_j = (q[r, h] . K[t]) * softmax_scale * kv_scale_quant_orig[1]
- out[r, h] = kv_scale_quant_orig[2] * 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]`.
- The kernels decode each byte with `cvt.rn.f16x2.e2m1x2` and scale it in f16,
- which is exact, and run the scores and P.V on f16 MMAs with fp32
- accumulation (q converted to f16, P rounded to f16 before P.V), with the
- online softmax and split-K merge of QSA's bf16 kernels. The second-level
- scales scale the scores and the normalized output, not each element.
-
- A negative index, a page past the table's width, a negative or out-of-pool
- table entry, and a request index outside the table's rows contribute
- nothing; a row left with nothing returns zeros. Returns a fresh, contiguous
- `[rows, Hq, D]` tensor of `q.dtype`.
-
- Two kernels, chosen inside the call: with `query_positions` and
- `compress_ratio` given and `rows * Hkv > 512`, a causal-prefix prefill
- kernel; otherwise split-K, plus a merge kernel when it used more than one
- split. **The prefill pair is a bound, not a mask**, exactly as on the bf16
- entry: the prefill kernel reads only `min(visible // cr, block_topk) * cr +
- visible % cr` leading columns of a row.
-
- Fusion boundary: the block-table gather, the NVFP4 dequantize, scale,
- online softmax, P.V and the split-K merge are inside. Selection, appending
- K/V to the pool, and q's projection, norm, RoPE and output gate are not.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="decode: 4 requests x 4 rows, 8-way split-K + merge, in the engine's trunk "
- "layout: layer 11's planes among the 24 role pages of a slot",
- spec=_decode(
- tuple((4000 + 37 * i, 4) for i in range(4)), slot_pages=24, layer=11, pool_pages=600
- ),
- ),
- Cell(
- why="context 4096 rows from 0: rows see 1..4096 tokens",
- spec=_decode(((0, 4096),), launches=_PREFILL, prefill_pair=True),
- ),
- Cell(
- why="captured decode: G 1, 8 and 32 requests x 1 and 8 rows, replay == eager",
- spec=_decode((), slots=32, capture=dict(groups=(1, 8, 32), rows_per_request=(1, 8))),
- ),
- )
-
- note = """
- Metadata consumed: none -- every input is an argument. What the target
- passes: `block_table` is `qsa_block_table` (physical page ids, the unscaled
- table); `selected_tokens` is `qsa_expand_block_indices`' output, logical
- positions, `-1`-padded; `k_cache` / `v_cache` and `k_scale_cache` /
- `v_scale_cache` are the layer's planes of the QSA manager's `get_buffers(i,
- "HND")` and `get_block_scale_buffers(i, "HND")`; `kv_scale_quant_orig` is the
- fp32 [3] tensor its dense attention call binds (index 1 for K, 2 for V).
-
- Call sites (the target's `modeling.py`): the context/mixed forward passes the
- prefill pair on every step, the decode forward never does; both pass `q` as
- the 24-head slice of the packed QKV row (row stride 7168).
-
- Launch shapes at Hkv 2 (programs = rows x 2): split-K at 32-token tiles
- 64-way up to 2 rows, 16-way up to 8, 8-way up to 16 and 4-way up to 32
- (4 warps); 64-token tiles 4-way up to 64 rows; 32-token tiles 8-way up to
- 128, 4-way up to 256 and 2-way up to 512 rows (2 warps); each followed by
- the merge. Past 256 rows the prefill kernel (32-token tiles, 2 warps) when the pair
- is passed; past 512 rows otherwise one split and no merge. The decode site
- reaches 256 rows at max_draft_len 7 (32 requests x 8).
-
- Each byte is decoded once, into the even and the odd halves of the head
- dimension, which the kernels keep apart: q's halves against K's on two MMAs,
- P against V's halves into two accumulators, the output stored back
- interleaved. On GB300 at 2051 selected tokens a row: 1 row 12.6 us, 128 rows
- 160 us, against 10.4 / 108 us for QSA's bf16 kernel over the same values.
-
- 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: P is rounded to f16 here and to
- bf16 there, and the second-level scales act on the scores and the output.
-
- CUDA graphs: the decode site captures at G requests x 1..8 rows over the
- [32, 260] table; a replay after new operands are written in place is
- bitwise the eager call on them.
-
- Compiled variants: the kernels take the pool's page count (`num_pages`), the
- table's row count (`num_requests`) and the split-K kernel the step's row
- count as plain ints, and Triton compiles each apart at 1, at multiples of 16
- and otherwise. Every cell runs a 32-row table (a multiple of 16); the other
- table-height classes are not certified. Every other plain int -- the strides, the table width 260 and
- the selection's row stride 2051 -- is fixed at the served value the cells use.
-
- Preconditions:
- - Whenever the prefill kernel can run, each selection row is packed at the
- front with exactly the bound's count of valid entries (QSA's producer
- does this; see the bf16 entry's note).
- - The pool holds what the attention op's NVFP4 append writes; the
- `qsa_paged_kv_store_nvfp4` entry writes the same bits.
- - `D` a power of two >= 32; `Hq % Hkv == 0`; `tokens_per_block % 4 == 0`.
- - `q` within f16's range: the kernels convert it (the target's q follows its
- RMS norm).
- - An SM 100-family GPU: the byte decode is `cvt.rn.f16x2.e2m1x2`.
-
- The callable rejects loudly (ValueError) every operand whose dtype, shape or
- layout it cannot read (the test pins the half-passed prefill pair); `rows ==
- 0` returns an empty `[0, Hq, D]` tensor (not certified). Calls are bitwise
- deterministic.
-
- Every row, request, page and position index is widened to 64 bits before any
- product. Not certified: head geometries other than 24/2 x 256;
- `tokens_per_block` other than 32; more than 8192 rows; capture of the
- context/mixed site; a pool of one slot; second-level scales other than 1.0;
- the prefill pair on the split-K kernel (256 context rows or fewer); one split
- without a merge; the trunk layout on the prefill kernel; rows whose selection
- resolves no token (all padding, a foreign request, unmapped pages); an fp16
- or unsliced q; int64 request indices.
- """
-
- def __call__(
- self,
- q: torch.Tensor,
- k_cache: torch.Tensor,
- v_cache: torch.Tensor,
- k_scale_cache: torch.Tensor,
- v_scale_cache: torch.Tensor,
- kv_scale_quant_orig: torch.Tensor,
- block_table: torch.Tensor,
- selected_tokens: torch.Tensor,
- request_indices: torch.Tensor,
- tokens_per_block: int,
- softmax_scale: float,
- query_positions: Optional[torch.Tensor] = None,
- compress_ratio: Optional[int] = None,
- ) -> torch.Tensor:
- return _sparse_gqa_nvfp4(
- q=q,
- k_cache=k_cache,
- v_cache=v_cache,
- k_scale_cache=k_scale_cache,
- v_scale_cache=v_scale_cache,
- kv_scale_quant_orig=kv_scale_quant_orig,
- block_table=block_table,
- selected_tokens=selected_tokens,
- request_indices=request_indices,
- tokens_per_block=tokens_per_block,
- softmax_scale=softmax_scale,
- query_positions=query_positions,
- compress_ratio=compress_ratio,
- )
-
- def reference(
- self,
- q: torch.Tensor,
- k_cache: torch.Tensor,
- v_cache: torch.Tensor,
- k_scale_cache: torch.Tensor,
- v_scale_cache: torch.Tensor,
- kv_scale_quant_orig: torch.Tensor,
- block_table: torch.Tensor,
- selected_tokens: torch.Tensor,
- request_indices: torch.Tensor,
- tokens_per_block: int,
- softmax_scale: float,
- query_positions: Optional[torch.Tensor] = None,
- compress_ratio: Optional[int] = None,
- ) -> torch.Tensor:
- """fp32 sparse GQA over every resident selected token, rounded once.
-
- K and V are dequantized in fp32 and rounded to `q`'s dtype, as the kernel
- does. `query_positions` and `compress_ratio` are ignored: on a
- prefix-packed selection the prefill kernel's bound reads exactly the valid
- entries. Walked 32 rows at a time to bound the gathered K/V.
- """
- rows, num_q_heads, head_dim = q.shape
- num_pages, num_kv_heads = k_cache.shape[0], k_cache.shape[1]
- groups = num_q_heads // num_kv_heads
- num_requests, table_width = block_table.shape
- requests = request_indices.long()
- out = torch.empty(rows, num_q_heads, head_dim, dtype=torch.float32, device=q.device)
- for start in range(0, rows, 32):
- stop = min(start + 32, rows)
- chunk = selected_tokens[start:stop].long()
- logical_page = chunk.clamp_min(0) // tokens_per_block
- in_page = chunk.clamp_min(0) % tokens_per_block
- request = requests[start:stop]
- owned = (request >= 0) & (request < num_requests)
- physical = block_table[
- request.clamp(0, num_requests - 1)[:, None].expand_as(logical_page),
- logical_page.clamp(max=table_width - 1),
- ].long()
- resident = (
- (chunk >= 0)
- & owned[:, None]
- & (logical_page < table_width)
- & (physical >= 0)
- & (physical < num_pages)
- )
- physical = physical.clamp(0, num_pages - 1)
- keys = _dequantize_tokens(
- k_cache, k_scale_cache, physical, in_page, kv_scale_quant_orig[1], grouped=False
- )
- values = _dequantize_tokens(
- v_cache, v_scale_cache, physical, in_page, kv_scale_quant_orig[2], grouped=True
- )
- keys = keys.to(q.dtype).float()
- values = values.to(q.dtype).float()
- grouped = q[start:stop].reshape(-1, num_kv_heads, groups, head_dim).float()
- scores = torch.einsum("bhgd,bkhd->bhgk", grouped, keys) * softmax_scale
- scores = scores.masked_fill(~resident[:, None, None, :], float("-inf"))
- weights = torch.nan_to_num(torch.softmax(scores, dim=-1), nan=0.0)
- pooled = torch.einsum("bhgk,bkhd->bhgd", weights, values)
- out[start:stop] = pooled.reshape(stop - start, num_q_heads, head_dim)
- return out.to(q.dtype)
-
-
-qsa_paged_sparse_gqa_nvfp4 = QsaPagedSparseGqaNvfp4()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4_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_nvfp4_decode.py
deleted file mode 100644
index b25c1c1b63a6..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4_decode.py
+++ /dev/null
@@ -1,1547 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""QSA sparse GQA over an NVFP4 paged K/V pool, decode rows: CuTe DSL split-K + merge.
-
-Tier 5 (CuTe DSL written for this target), both launches. One CTA (eight warps) serves one
-(row, KV head, split) and walks the split's tiles of 64 selected tokens with an online
-softmax; the split count is chosen by row count so the grid stays about two waves deep
-(33 splits of one tile at 1 row, one split of 33 tiles from 128 rows up).
-
-Per tile:
-
-1. raw NVFP4 bytes land in shared memory with cp.async, one tile ahead (two buffers):
- each token's 128 K bytes, its 16 K scale bytes (token-major), its 128 V bytes and the
- 64-byte four-token V-scale block it sits in; a token that does not resolve (padding,
- foreign request, unmapped page) gets zeros instead;
-2. the V tile is dequantized into an f16 tile (`cvt.rn.f16x2.e2m1x2` x `cvt.rn.f16x2.e4m3x2`,
- exact) for `ldmatrix.trans`;
-3. S = q K^T on `mma.sync.m16n8k16` (f16 in, f32 out), a warp eight tokens over the whole
- head: K's B fragments are decoded straight from the raw bytes in registers. The head
- dimension is permuted identically in q and K (lane quad c of k-step pair m reads the
- bytes 16m + 4c .. 16m + 4c + 3, dims 32m + 8c .. 32m + 8c + 7, inside one scale block),
- which leaves the dot product unchanged;
-4. online softmax in the log2 domain across the warps (two shared-memory exchanges), P
- rounded to f16 as the Triton kernel does;
-5. O += P V on `mma.sync` (a warp 32 head dims), P from shared memory with `ldmatrix`,
- V with `ldmatrix.trans`.
-
-The second-level scales act on the scores (K) and the normalized output (V), as in the
-Triton op. With one split the CTA stores the output itself (optionally times the
-sigmoid of the output gate, applied to the bf16-rounded attention value); otherwise it
-stores an fp32 partial and its log2-sum-exp for the merge kernel.
-"""
-
-from __future__ import annotations
-
-import math
-from typing import Optional
-
-import cutlass
-import cutlass.cute as cute
-import torch
-from cutlass import const_expr
-from cutlass._mlir.dialects import llvm
-from cutlass.cute.runtime import from_dlpack
-from cutlass.cutlass_dsl import T, dsl_user_op
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
-from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
-
-from .qsa_paged_sparse_gqa_nvfp4 import qsa_paged_sparse_gqa_nvfp4 as _composition
-
-TILE = 64
-WARPS = 8
-THREADS = WARPS * 32
-HEAD_DIM = 256
-#: Bytes of a raw K row in shared memory (128 + 16 of padding: conflict-free fragment reads).
-RAW_LD = 144
-#: Bytes of a raw V row (unpadded: the dequant pass reads 32 bytes a thread).
-RAW_LDV = 128
-#: Raw tiles in flight (cp.async ring).
-STAGES = 3
-#: Bytes of an f16 V row in shared memory (256 + 8 elements: conflict-free ldmatrix).
-SV_LD = 528
-#: Bytes of an f16 P row (64 + 8 elements).
-SP_LD = 144
-LOG2E = 1.4426950408889634
-#: Selected tokens a merge CTA covers per head (four adjacent ones a lane of each warp).
-CHUNK = 128
-MAX_SPLITS = 64
-
-
-# ---------------------------------------------------------------------------------------
-# inline PTX helpers
-# ---------------------------------------------------------------------------------------
-
-
-def _ir(x, cls, loc, ip):
- return cls(x).ir_value(loc=loc, ip=ip)
-
-
-def _asm(ret, args, text, cons, loc, ip, side=False):
- return llvm.inline_asm(
- ret,
- args,
- text,
- cons,
- has_side_effects=side,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
-
-
-def _unpack(res, n, ty, loc, ip):
- return [
- ty(llvm.extractvalue(T.i32() if ty is cutlass.Int32 else T.f32(), res, [i], loc=loc, ip=ip))
- for i in range(n)
- ]
-
-
-@dsl_user_op
-def _cp16(saddr, gaddr, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [_ir(saddr, cutlass.Int32, loc, ip), _ir(gaddr, cutlass.Int64, loc, ip)],
- "{ cp.async.cg.shared.global [$1], [$2], 16; mov.u32 $0, 0; }",
- "=r,r,l",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _cp16z(saddr, gaddr, nbytes, *, loc=None, ip=None):
- """cp.async of 16 bytes, zero-filled past `nbytes` (0 or 16): no branch for a dead token."""
- _asm(
- T.i32(),
- [
- _ir(saddr, cutlass.Int32, loc, ip),
- _ir(gaddr, cutlass.Int64, loc, ip),
- _ir(nbytes, cutlass.Int32, loc, ip),
- ],
- "{ cp.async.cg.shared.global [$1], [$2], 16, $3; mov.u32 $0, 0; }",
- "=r,r,l,r",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _sts4(saddr, a, b, c, d, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [_ir(saddr, cutlass.Int32, loc, ip)]
- + [_ir(v, cutlass.Int32, loc, ip) for v in (a, b, c, d)],
- "{ st.shared.v4.b32 [$1], {$2, $3, $4, $5}; mov.u32 $0, 0; }",
- "=r,r,r,r,r,r",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _sts1(saddr, a, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [_ir(saddr, cutlass.Int32, loc, ip), _ir(a, cutlass.Int32, loc, ip)],
- "{ st.shared.b32 [$1], $2; mov.u32 $0, 0; }",
- "=r,r,r",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _lds1(saddr, *, loc=None, ip=None):
- return cutlass.Int32(
- _asm(
- T.i32(),
- [_ir(saddr, cutlass.Int32, loc, ip)],
- "ld.shared.b32 $0, [$1];",
- "=r,r",
- loc,
- ip,
- side=True,
- )
- )
-
-
-@dsl_user_op
-def _lds4(saddr, *, loc=None, ip=None):
- res = _asm(
- llvm.StructType.get_literal([T.i32()] * 4),
- [_ir(saddr, cutlass.Int32, loc, ip)],
- "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];",
- "=r,=r,=r,=r,r",
- loc,
- ip,
- side=True,
- )
- return _unpack(res, 4, cutlass.Int32, loc, ip)
-
-
-@dsl_user_op
-def _ldg4(gaddr, *, loc=None, ip=None):
- res = _asm(
- llvm.StructType.get_literal([T.i32()] * 4),
- [_ir(gaddr, cutlass.Int64, loc, ip)],
- "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];",
- "=r,=r,=r,=r,l",
- loc,
- ip,
- side=True,
- )
- return _unpack(res, 4, cutlass.Int32, loc, ip)
-
-
-@dsl_user_op
-def _ldsm4(saddr, *, loc=None, ip=None):
- res = _asm(
- llvm.StructType.get_literal([T.i32()] * 4),
- [_ir(saddr, cutlass.Int32, loc, ip)],
- "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {$0, $1, $2, $3}, [$4];",
- "=r,=r,=r,=r,r",
- loc,
- ip,
- side=True,
- )
- return _unpack(res, 4, cutlass.Int32, loc, ip)
-
-
-@dsl_user_op
-def _ldsm4t(saddr, *, loc=None, ip=None):
- res = _asm(
- llvm.StructType.get_literal([T.i32()] * 4),
- [_ir(saddr, cutlass.Int32, loc, ip)],
- "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {$0, $1, $2, $3}, [$4];",
- "=r,=r,=r,=r,r",
- loc,
- ip,
- side=True,
- )
- return _unpack(res, 4, cutlass.Int32, loc, ip)
-
-
-@dsl_user_op
-def _mma(a, b, c, *, loc=None, ip=None):
- """D = A B + C on m16n8k16 (f16 A/B, f32 C/D); a: 4 regs, b: 2 regs, c: 4 f32."""
- res = _asm(
- llvm.StructType.get_literal([T.f32()] * 4),
- [_ir(v, cutlass.Int32, loc, ip) for v in a]
- + [_ir(v, cutlass.Int32, loc, ip) for v in b]
- + [_ir(v, cutlass.Float32, loc, ip) for v in c],
- "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
- "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};",
- "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
- loc,
- ip,
- )
- return _unpack(res, 4, cutlass.Float32, loc, ip)
-
-
-@dsl_user_op
-def _dq_word(word, scale2, *, loc=None, ip=None):
- """Four e2m1 bytes (8 values, low nibble the even one) times an f16x2 scale -> 4 f16x2."""
- res = _asm(
- llvm.StructType.get_literal([T.i32()] * 4),
- [_ir(word, cutlass.Int32, loc, ip), _ir(scale2, cutlass.Int32, loc, ip)],
- "{ .reg .b8 b0, b1, b2, b3; .reg .b32 h0, h1, h2, h3; "
- "mov.b32 {b0, b1, b2, b3}, $4; "
- "cvt.rn.f16x2.e2m1x2 h0, b0; cvt.rn.f16x2.e2m1x2 h1, b1; "
- "cvt.rn.f16x2.e2m1x2 h2, b2; cvt.rn.f16x2.e2m1x2 h3, b3; "
- "mul.rn.f16x2 $0, h0, $5; mul.rn.f16x2 $1, h1, $5; "
- "mul.rn.f16x2 $2, h2, $5; mul.rn.f16x2 $3, h3, $5; }",
- "=r,=r,=r,=r,r,r",
- loc,
- ip,
- )
- return _unpack(res, 4, cutlass.Int32, loc, ip)
-
-
-@dsl_user_op
-def _e2m1_half(half, *, loc=None, ip=None):
- """Two e2m1 bytes (the low 16 bits: 4 values, low nibble the even one) -> 2 f16x2
- (values 0-1, values 2-3), unscaled."""
- res = _asm(
- llvm.StructType.get_literal([T.i32()] * 2),
- [_ir(half, cutlass.Int32, loc, ip)],
- "{ .reg .b8 b0, b1, b2, b3; mov.b32 {b0, b1, b2, b3}, $2; "
- "cvt.rn.f16x2.e2m1x2 $0, b0; cvt.rn.f16x2.e2m1x2 $1, b1; }",
- "=r,=r,r",
- loc,
- ip,
- )
- return _unpack(res, 2, cutlass.Int32, loc, ip)
-
-
-@dsl_user_op
-def _e4m3x2(half, *, loc=None, ip=None):
- """Two e4m3 bytes (the low 16 bits) -> f16x2, the low byte in the low half."""
- return cutlass.Int32(
- _asm(
- T.i32(),
- [_ir(half, cutlass.Int32, loc, ip)],
- "{ .reg .b16 h; cvt.u16.u32 h, $1; cvt.rn.f16x2.e4m3x2 $0, h; }",
- "=r,r",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _hmul2(a, b, *, loc=None, ip=None):
- return cutlass.Int32(
- _asm(
- T.i32(),
- [_ir(a, cutlass.Int32, loc, ip), _ir(b, cutlass.Int32, loc, ip)],
- "mul.rn.f16x2 $0, $1, $2;",
- "=r,r,r",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _prmt(a, b, sel, *, loc=None, ip=None):
- return cutlass.Int32(
- _asm(
- T.i32(),
- [
- _ir(a, cutlass.Int32, loc, ip),
- _ir(b, cutlass.Int32, loc, ip),
- _ir(sel, cutlass.Int32, loc, ip),
- ], # fmt: skip
- "prmt.b32 $0, $1, $2, $3;",
- "=r,r,r,r",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _lds_u16(saddr, *, loc=None, ip=None):
- return cutlass.Int32(
- _asm(
- T.i32(),
- [_ir(saddr, cutlass.Int32, loc, ip)],
- "ld.shared.u16 $0, [$1];",
- "=r,r",
- loc,
- ip,
- side=True,
- )
- )
-
-
-@dsl_user_op
-def _stg4(gaddr, a, b, c, d, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [_ir(gaddr, cutlass.Int64, loc, ip)]
- + [_ir(v, cutlass.Int32, loc, ip) 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=True,
- )
-
-
-@dsl_user_op
-def _scale2(sbyte, *, loc=None, ip=None):
- """An e4m3 byte (in the low 8 bits) -> f16x2 {s, s}."""
- return cutlass.Int32(
- _asm(
- T.i32(),
- [_ir(sbyte, cutlass.Int32, loc, ip)],
- "{ .reg .b32 t; .reg .b16 h; and.b32 t, $1, 255; mul.lo.u32 t, t, 257; "
- "cvt.u16.u32 h, t; cvt.rn.f16x2.e4m3x2 $0, h; }",
- "=r,r",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _scale2s(word, sel, *, loc=None, ip=None):
- """The e4m3 byte `word.byte[sel & 3]` (sel = k * 0x11) -> f16x2 {s, s}: one prmt + one cvt."""
- return cutlass.Int32(
- _asm(
- T.i32(),
- [_ir(word, cutlass.Int32, loc, ip), _ir(sel, cutlass.Int32, loc, ip)],
- "{ .reg .b32 t; .reg .b16 h; prmt.b32 t, $1, 0, $2; cvt.u16.u32 h, t; "
- "cvt.rn.f16x2.e4m3x2 $0, h; }",
- "=r,r,r",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _ex2(x, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(
- T.f32(),
- [_ir(x, cutlass.Float32, loc, ip)],
- "ex2.approx.ftz.f32 $0, $1;",
- "=f,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _bf2_to_f2(w, *, loc=None, ip=None):
- """bf16x2 -> f16x2 (through f32, round to nearest)."""
- return cutlass.Int32(
- _asm(
- T.i32(),
- [_ir(w, cutlass.Int32, loc, ip)],
- "{ .reg .b16 lo, hi; .reg .f32 fl, fh; mov.b32 {lo, hi}, $1; "
- "cvt.f32.bf16 fl, lo; cvt.f32.bf16 fh, hi; cvt.rn.f16x2.f32 $0, fh, fl; }",
- "=r,r",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _f2_pack(lo, hi, *, loc=None, ip=None):
- """Two f32 -> f16x2 (lo in the low half)."""
- return cutlass.Int32(
- _asm(
- T.i32(),
- [_ir(lo, cutlass.Float32, loc, ip), _ir(hi, cutlass.Float32, loc, ip)],
- "cvt.rn.f16x2.f32 $0, $2, $1;",
- "=r,f,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _bf2_pack(lo, hi, *, loc=None, ip=None):
- """Two f32 -> bf16x2 (lo in the low half)."""
- return cutlass.Int32(
- _asm(
- T.i32(),
- [_ir(lo, cutlass.Float32, loc, ip), _ir(hi, cutlass.Float32, loc, ip)],
- "cvt.rn.bf16x2.f32 $0, $2, $1;",
- "=r,f,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _stg1(gaddr, a, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [_ir(gaddr, cutlass.Int64, loc, ip), _ir(a, cutlass.Int32, loc, ip)],
- "{ st.global.b32 [$1], $2; mov.u32 $0, 0; }",
- "=r,l,r",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _stg2f(gaddr, a, b, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [
- _ir(gaddr, cutlass.Int64, loc, ip),
- _ir(a, cutlass.Float32, loc, ip),
- _ir(b, cutlass.Float32, loc, ip),
- ],
- "{ st.global.v2.f32 [$1], {$2, $3}; mov.u32 $0, 0; }",
- "=r,l,f,f",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _ldg1(gaddr, *, loc=None, ip=None):
- return cutlass.Int32(
- _asm(
- T.i32(),
- [_ir(gaddr, cutlass.Int64, loc, ip)],
- "ld.global.b32 $0, [$1];",
- "=r,l",
- loc,
- ip,
- side=True,
- )
- )
-
-
-@dsl_user_op
-def _bf2_lo(w, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(
- T.f32(),
- [_ir(w, cutlass.Int32, loc, ip)],
- "{ .reg .b16 lo, hi; mov.b32 {lo, hi}, $1; cvt.f32.bf16 $0, lo; }",
- "=f,r",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _bf2_hi(w, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(
- T.f32(),
- [_ir(w, cutlass.Int32, loc, ip)],
- "{ .reg .b16 lo, hi; mov.b32 {lo, hi}, $1; cvt.f32.bf16 $0, hi; }",
- "=f,r",
- loc,
- ip,
- )
- )
-
-
-# ---------------------------------------------------------------------------------------
-# split kernel
-# ---------------------------------------------------------------------------------------
-
-
-class _SplitKernel:
- def __init__(
- self,
- *,
- group,
- kv_heads,
- width,
- tpb,
- splits,
- tps,
- scale,
- num_pages,
- bt_width,
- num_requests,
- k_strides,
- v_strides,
- ks_strides,
- vs_strides,
- q_strides,
- req64,
- gated,
- pdl,
- ):
- assert group <= 16 and HEAD_DIM == 256 and tpb % 4 == 0
- self.group, self.kvh, self.width, self.tpb = group, kv_heads, width, tpb
- self.splits, self.tps = splits, tps
- self.n_tiles = -(-width // TILE)
- self.scale = scale * LOG2E
- self.num_pages, self.bt_width, self.num_requests = num_pages, bt_width, num_requests
- self.ks, self.vs, self.kss, self.vss = k_strides, v_strides, ks_strides, vs_strides
- self.qs = q_strides
- self.req64, self.gated, self.pdl = req64, gated, pdl
- self.heads = group * kv_heads
- # shared memory carve-out (bytes)
- self.o_tok = 0
- self.tok_bytes = -(-tps * TILE * 4 // 128) * 128
- self.raw_ldv = RAW_LDV
- self.raw_vsld = 64
- self.raw_k = TILE * RAW_LD
- self.raw_v = TILE * self.raw_ldv
- self.raw_ks = TILE * 16
- self.raw_vs = TILE * self.raw_vsld
- self.raw_buf = self.raw_k + self.raw_v + self.raw_ks + self.raw_vs
- self.stages = min(STAGES, tps)
- self.o_raw = self.o_tok + self.tok_bytes
- self.o_sv = self.o_raw + self.stages * self.raw_buf
- self.o_sp = self.o_sv + TILE * SV_LD
- self.o_red = self.o_sp + 16 * SP_LD
- self.smem = self.o_red + 2 * WARPS * 16 * 4 + 128
- assert k_strides == v_strides and ks_strides == vs_strides, (
- "K and V planes share one layout"
- )
- assert tpb & (tpb - 1) == 0
-
- @cute.jit
- def __call__(
- self,
- sel_t: cute.Tensor,
- q_ptr: cute.Pointer,
- k_ptr: cute.Pointer,
- v_ptr: cute.Pointer,
- ks_ptr: cute.Pointer,
- vs_ptr: cute.Pointer,
- gs_ptr: cute.Pointer,
- bt_ptr: cute.Pointer,
- req_ptr: cute.Pointer,
- out_ptr: cute.Pointer,
- lse_ptr: cute.Pointer,
- gate_ptr: cute.Pointer,
- rows: cutlass.Int32,
- gate_stride: cutlass.Int32,
- stream,
- ):
- sel = cute.make_tensor(
- sel_t.iterator, cute.make_layout((rows, self.width), stride=(self.width, 1))
- )
- bt = cute.make_tensor(
- bt_ptr, cute.make_layout((self.num_requests, self.bt_width), stride=(self.bt_width, 1))
- )
- req = cute.make_tensor(req_ptr, cute.make_layout((rows,), stride=(1,)))
- gs = cute.make_tensor(gs_ptr, cute.make_layout((3,), stride=(1,)))
- self.kernel(
- sel,
- q_ptr.toint(),
- k_ptr.toint(),
- v_ptr.toint(),
- ks_ptr.toint(),
- vs_ptr.toint(),
- gs,
- bt,
- req,
- out_ptr.toint(),
- lse_ptr.toint(),
- gate_ptr.toint(),
- rows,
- gate_stride,
- ).launch(
- grid=[rows * self.kvh * self.splits, 1, 1],
- block=[THREADS, 1, 1],
- smem=self.smem,
- stream=stream,
- use_pdl=self.pdl,
- min_blocks_per_mp=2,
- )
-
- @cute.jit
- def _issue(self, j, buf, tidx, kvh, s_tok, s_raw, k_base, v_base, ks_base, vs_base):
- """cp.async of tile `j`'s raw K, K-scale, V and V-scale-block rows into ring slot `buf`,
- committed as one group: thread t serves token t // 4 (a quarter of its rows); a token
- that does not resolve gets zeros (src-size 0)."""
- tpb: cutlass.Constexpr = self.tpb
- lt: cutlass.Constexpr = tpb.bit_length() - 1
- ks0, ks1, ks2 = self.ks
- kss0, kss1 = self.kss
- rb = s_raw + buf * self.raw_buf
- rk = rb
- rv = rb + self.raw_k
- rks = rv + self.raw_v
- rvs = rks + self.raw_ks
- tt = tidx // 4
- q4 = tidx % 4
- ent = _lds1(s_tok + (j * TILE + tt) * 4)
- ok = ent >= 0
- e = cutlass.max(ent, 0)
- pg = cutlass.Int64(e >> lt)
- sl = e & (tpb - 1)
- n = cutlass.Int32(cutlass.select_(ok, cutlass.Int32(16), cutlass.Int32(0)))
- off = pg * ks0 + cutlass.Int64(kvh * ks1 + sl * ks2 + 32 * q4)
- offs = pg * kss0 + cutlass.Int64(kvh * kss1)
- _cp16z(rk + tt * RAW_LD + 32 * q4, k_base + off, n)
- _cp16z(rk + tt * RAW_LD + 32 * q4 + 16, k_base + off + 16, n)
- _cp16z(rv + tt * self.raw_ldv + 32 * q4, v_base + off, n)
- _cp16z(rv + tt * self.raw_ldv + 32 * q4 + 16, v_base + off + 16, n)
- _cp16z(
- rvs + tt * self.raw_vsld + 16 * q4,
- vs_base + offs + cutlass.Int64((sl >> 2) * 64 + 16 * q4),
- n,
- )
- if q4 == 0:
- _cp16z(rks + tt * 16, ks_base + offs + cutlass.Int64(sl * 16), n)
- cute.arch.cp_async_commit_group()
-
- @cute.kernel
- def kernel(
- self,
- mSel: cute.Tensor,
- q_base: cutlass.Int64,
- k_base: cutlass.Int64,
- v_base: cutlass.Int64,
- ks_base: cutlass.Int64,
- vs_base: cutlass.Int64,
- mGs: cute.Tensor,
- mBT: cute.Tensor,
- mReq: cute.Tensor,
- out_base: cutlass.Int64,
- lse_base: cutlass.Int64,
- gate_base: cutlass.Int64,
- rows: cutlass.Int32,
- gate_stride: cutlass.Int32,
- ):
- tidx, _, _ = cute.arch.thread_idx()
- bidx, _, _ = cute.arch.block_idx()
- warp = cute.arch.warp_idx()
- lane = cute.arch.lane_idx()
- g = lane // 4
- c = lane % 4
- group: cutlass.Constexpr = self.group
- kvh_n: cutlass.Constexpr = self.kvh
- splits: cutlass.Constexpr = self.splits
- tps: cutlass.Constexpr = self.tps
- n_tiles_all: cutlass.Constexpr = self.n_tiles
- width: cutlass.Constexpr = self.width
- tpb: cutlass.Constexpr = self.tpb
- heads: cutlass.Constexpr = self.heads
- row = bidx // (kvh_n * splits)
- kvh = (bidx // splits) % kvh_n
- split = bidx % splits
- t0 = split * tps
- n_t = cutlass.min(tps, n_tiles_all - t0)
-
- smem = cutlass.utils.SmemAllocator()
- base_t = smem.allocate_tensor(
- cutlass.Int8, cute.make_layout((self.smem - 128,), stride=(1,)), byte_alignment=128
- )
- sbase = cutlass.Int32(base_t.iterator.toint())
- s_tok = sbase + self.o_tok
- s_raw = sbase + self.o_raw
- s_v = sbase + self.o_sv
- s_p = sbase + self.o_sp
- s_red = sbase + self.o_red
-
- if const_expr(self.pdl):
- cute.arch.griddepcontrol_launch_dependents()
- cute.arch.griddepcontrol_wait()
-
- # ---- the token table of this split: page * tpb + slot, or -1 ---------------------
- rq = cutlass.Int64(mReq[row])
- req_ok = (rq >= 0) & (rq < self.num_requests)
- safe_req = cutlass.Int32(
- cutlass.min(cutlass.max(rq, cutlass.Int64(0)), cutlass.Int64(self.num_requests - 1))
- )
- for i in cutlass.range_constexpr(-(-(tps * TILE) // THREADS)):
- cl = tidx + i * THREADS
- col = t0 * TILE + cl
- 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 < self.bt_width)
- page = mBT[safe_req, cutlass.min(lp, self.bt_width - 1)]
- ok = ok & (page >= 0) & (page < self.num_pages)
- ent = cutlass.Int32(cutlass.select_(ok, page * tpb + safe_tok % tpb, cutlass.Int32(-1)))
- if cl < tps * TILE:
- _sts1(s_tok + cl * 4, ent)
- cute.arch.sync_threads()
-
- # ---- q of this KV head: rows g and g + 8 of the m16 tile, as permuted A fragments --
- # thread (g, c) holds dims 32m + 8c .. 32m + 8c + 7 of heads g and g + 8, m = 0..7
- qa = cute.make_rmem_tensor(cute.make_layout((8, 4), stride=(4, 1)), cutlass.Int32)
- qb = cute.make_rmem_tensor(cute.make_layout((8, 4), stride=(4, 1)), cutlass.Int32)
- qs0, qs1 = self.qs
- q_row = q_base + (cutlass.Int64(row) * qs0 + cutlass.Int64(kvh * group) * qs1) * 2
- for m in cutlass.range_constexpr(8):
- off = cutlass.Int64((32 * m + 8 * c) * 2)
- w4 = _ldg4(q_row + cutlass.Int64(g) * (qs1 * 2) + off)
- for e in cutlass.range_constexpr(4):
- qa[m, e] = _bf2_to_f2(w4[e])
- hb = g + 8
- hbc = cutlass.min(hb, group - 1)
- w4b = _ldg4(q_row + cutlass.Int64(hbc) * (qs1 * 2) + off)
- for e in cutlass.range_constexpr(4):
- qb[m, e] = cutlass.Int32(
- cutlass.select_(hb < group, _bf2_to_f2(w4b[e]), cutlass.Int32(0))
- )
-
- stages: cutlass.Constexpr = self.stages
- self._issue(0, 0, tidx, kvh, s_tok, s_raw, k_base, v_base, ks_base, vs_base)
- if const_expr(stages > 2):
- if n_t > 1:
- self._issue(1, 1, tidx, kvh, s_tok, s_raw, k_base, v_base, ks_base, vs_base)
-
- neg_inf = cutlass.Float32(-math.inf)
- st = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), cutlass.Float32)
- st[0] = neg_inf # running max, head g
- st[1] = neg_inf # running max, head g + 8
- st[2] = cutlass.Float32(0.0) # this thread's share of the running sum, head g
- st[3] = cutlass.Float32(0.0) # ... head g + 8 (reduced over the CTA at the end)
- acc = cute.make_rmem_tensor(cute.make_layout((4, 4), stride=(4, 1)), cutlass.Float32)
- acc.fill(0.0)
- score_scale = cutlass.Float32(self.scale) * mGs[1]
-
- # loop-invariant per-thread offsets
- vt = tidx // 4
- vq = tidx % 4
- v_src = vt * RAW_LDV + 32 * vq
- vs_src = vt * 64 + 16 * vq
- v_dst = s_v + vt * SV_LD + 128 * vq
- sel_v = (vt % 4) * 0x11
- ktok = warp * 8 + g
- k_src = ktok * RAW_LD + 4 * c
- ks_src = ktok * 16
- sel_k0 = (c // 2) * 0x11
- sel_k1 = (2 + c // 2) * 0x11
- ta = warp * 8 + 2 * c
- p_st_g = s_p + g * SP_LD + ta * 2
- p_st_h = p_st_g + 8 * SP_LD
- p_ld = s_p + ((lane % 8) + 8 * ((lane // 8) % 2)) * SP_LD + 16 * (lane // 16)
- v_ld_row = s_v + ((lane % 8) + 8 * ((lane // 8) % 2)) * SV_LD
- ci0 = 4 * warp + lane // 16
- ci1 = ci0 + 2
- v_ld0 = v_ld_row + 16 * ((ci0 & ~7) | ((ci0 & 7) ^ (2 * ((ci0 >> 3) & 3))))
- v_ld1 = v_ld_row + 16 * ((ci1 & ~7) | ((ci1 & 7) ^ (2 * ((ci1 >> 3) & 3))))
- red_g = s_red + g * WARPS * 4
- red_h = s_red + (g + 8) * WARPS * 4
-
- for j in cutlass.range(n_t):
- buf = j % stages
- if const_expr(stages == 1):
- cute.arch.cp_async_wait_group(0)
- else:
- ahead: cutlass.Constexpr = stages - 1
- if j + ahead < n_t:
- self._issue(
- j + ahead,
- (j + ahead) % stages,
- tidx,
- kvh,
- s_tok,
- s_raw,
- k_base,
- v_base,
- ks_base,
- vs_base,
- )
- cute.arch.cp_async_wait_group(ahead)
- else:
- if const_expr(stages > 2):
- if j + 1 < n_t:
- cute.arch.cp_async_wait_group(1)
- else:
- cute.arch.cp_async_wait_group(0)
- else:
- cute.arch.cp_async_wait_group(0)
- cute.arch.sync_threads() # [A] tile j landed; previous tile's PV done
- rb = s_raw + buf * self.raw_buf
- rk = rb
- rv = rb + self.raw_k
- rks = rv + self.raw_v
- rvs = rks + self.raw_ks
-
- # -- V tile -> f16 (thread: token tidx // 4, quarter tidx % 4 = dims 64q .. 64q+63)
- vw0 = _lds4(rv + v_src)
- vw1 = _lds4(rv + v_src + 16)
- vsw = _lds4(rvs + vs_src)
- for gg in cutlass.range_constexpr(4):
- sc = _scale2s(vsw[gg], sel_v)
- wa = vw0[2 * gg] if gg < 2 else vw1[2 * gg - 4]
- wb = vw0[2 * gg + 1] if gg < 2 else vw1[2 * gg - 3]
- da = _dq_word(wa, sc)
- db = _dq_word(wb, sc)
- # logical 16-byte chunk 8 vq + k of the row sits at 8 vq + (k ^ 2 vq): the four
- # threads of a token store to different banks; ldmatrix applies the same map
- _sts4(v_dst + 16 * ((2 * gg) ^ (2 * vq)), da[0], da[1], da[2], da[3])
- _sts4(v_dst + 16 * ((2 * gg + 1) ^ (2 * vq)), db[0], db[1], db[2], db[3])
-
- # -- S = q K^T for this warp's eight tokens (two accumulator chains)
- ksw = _lds4(rks + ks_src)
- kws = [_lds1(rk + k_src + 16 * m) for m in range(8)]
- zero = cutlass.Float32(0.0)
- sa = [zero, zero, zero, zero]
- sb = [zero, zero, zero, zero]
- for m in cutlass.range_constexpr(8):
- # scale block 2m + c // 2: word m // 2, byte 2 (m % 2) + c // 2
- sc = _scale2s(ksw[m // 2], sel_k0 if m % 2 == 0 else sel_k1)
- b = _dq_word(kws[m], sc)
- if const_expr(m % 2 == 0):
- sa = _mma((qa[m, 0], qb[m, 0], qa[m, 1], qb[m, 1]), (b[0], b[1]), sa)
- sa = _mma((qa[m, 2], qb[m, 2], qa[m, 3], qb[m, 3]), (b[2], b[3]), sa)
- else:
- sb = _mma((qa[m, 0], qb[m, 0], qa[m, 1], qb[m, 1]), (b[0], b[1]), sb)
- sb = _mma((qa[m, 2], qb[m, 2], qa[m, 3], qb[m, 3]), (b[2], b[3]), sb)
-
- # mask + scale; tokens 8w + 2c, 8w + 2c + 1
- ea = _lds1(s_tok + (j * TILE + ta) * 4)
- eb = _lds1(s_tok + (j * TILE + ta + 1) * 4)
- s0 = cutlass.Float32(cutlass.select_(ea >= 0, (sa[0] + sb[0]) * score_scale, neg_inf))
- s1 = cutlass.Float32(cutlass.select_(eb >= 0, (sa[1] + sb[1]) * score_scale, neg_inf))
- s2 = cutlass.Float32(cutlass.select_(ea >= 0, (sa[2] + sb[2]) * score_scale, neg_inf))
- s3 = cutlass.Float32(cutlass.select_(eb >= 0, (sa[3] + sb[3]) * score_scale, neg_inf))
- mg = cute.arch.fmax(s0, s1)
- mh = cute.arch.fmax(s2, s3)
- mg = cute.arch.fmax(mg, cute.arch.shuffle_sync_bfly(mg, offset=1))
- mh = cute.arch.fmax(mh, cute.arch.shuffle_sync_bfly(mh, offset=1))
- mg = cute.arch.fmax(mg, cute.arch.shuffle_sync_bfly(mg, offset=2))
- mh = cute.arch.fmax(mh, cute.arch.shuffle_sync_bfly(mh, offset=2))
- if c == 0:
- _sts1(red_g + warp * 4, _f32_bits(mg))
- _sts1(red_h + warp * 4, _f32_bits(mh))
- cute.arch.sync_threads() # [B] maxes; V tile in shared memory
- rg0 = _lds4(red_g)
- rg1 = _lds4(red_g + 16)
- rh0 = _lds4(red_h)
- rh1 = _lds4(red_h + 16)
- tmg = st[0]
- tmh = st[1]
- for e in cutlass.range_constexpr(4):
- tmg = cute.arch.fmax(tmg, cute.arch.fmax(_bits_f32(rg0[e]), _bits_f32(rg1[e])))
- tmh = cute.arch.fmax(tmh, cute.arch.fmax(_bits_f32(rh0[e]), _bits_f32(rh1[e])))
- safe_g = cutlass.Float32(cutlass.select_(tmg == neg_inf, cutlass.Float32(0.0), tmg))
- safe_h = cutlass.Float32(cutlass.select_(tmh == neg_inf, cutlass.Float32(0.0), tmh))
- corr_g = _ex2(st[0] - safe_g)
- corr_h = _ex2(st[1] - safe_h)
- p0 = _ex2(s0 - safe_g)
- p1 = _ex2(s1 - safe_g)
- p2 = _ex2(s2 - safe_h)
- p3 = _ex2(s3 - safe_h)
- _sts1(p_st_g, _f2_pack(p0, p1))
- _sts1(p_st_h, _f2_pack(p2, p3))
- st[0] = tmg
- st[1] = tmh
- st[2] = st[2] * corr_g + (p0 + p1)
- st[3] = st[3] * corr_h + (p2 + p3)
- cute.arch.sync_threads() # [C] P
-
- # -- O += P V: this warp's dims 32w .. 32w + 31 (n-tiles of 8)
- if cute.arch.vote_any_sync((corr_g != 1.0) | (corr_h != 1.0)):
- for nt in cutlass.range_constexpr(4):
- acc[nt, 0] = acc[nt, 0] * corr_g
- acc[nt, 1] = acc[nt, 1] * corr_g
- acc[nt, 2] = acc[nt, 2] * corr_h
- acc[nt, 3] = acc[nt, 3] * corr_h
- for kk in cutlass.range_constexpr(4):
- a = _ldsm4(p_ld + 32 * kk)
- b0 = _ldsm4t(v_ld0 + 16 * kk * SV_LD)
- b1 = _ldsm4t(v_ld1 + 16 * kk * SV_LD)
- r0 = _mma(a, (b0[0], b0[1]), (acc[0, 0], acc[0, 1], acc[0, 2], acc[0, 3]))
- r1 = _mma(a, (b0[2], b0[3]), (acc[1, 0], acc[1, 1], acc[1, 2], acc[1, 3]))
- r2 = _mma(a, (b1[0], b1[1]), (acc[2, 0], acc[2, 1], acc[2, 2], acc[2, 3]))
- r3 = _mma(a, (b1[2], b1[3]), (acc[3, 0], acc[3, 1], acc[3, 2], acc[3, 3]))
- for e in cutlass.range_constexpr(4):
- acc[0, e] = r0[e]
- acc[1, e] = r1[e]
- acc[2, e] = r2[e]
- acc[3, e] = r3[e]
-
- # ---- the running sums over the quad and the warps -----------------------------------
- lsg = st[2]
- lsh = st[3]
- lsg = lsg + cute.arch.shuffle_sync_bfly(lsg, offset=1)
- lsh = lsh + cute.arch.shuffle_sync_bfly(lsh, offset=1)
- lsg = lsg + cute.arch.shuffle_sync_bfly(lsg, offset=2)
- lsh = lsh + cute.arch.shuffle_sync_bfly(lsh, offset=2)
- cute.arch.sync_threads()
- if c == 0:
- _sts1(red_g + warp * 4, _f32_bits(lsg))
- _sts1(red_h + warp * 4, _f32_bits(lsh))
- cute.arch.sync_threads()
- ug0 = _lds4(red_g)
- ug1 = _lds4(red_g + 16)
- uh0 = _lds4(red_h)
- uh1 = _lds4(red_h + 16)
- tsg = cutlass.Float32(0.0)
- tsh = cutlass.Float32(0.0)
- for e in cutlass.range_constexpr(4):
- tsg = tsg + _bits_f32(ug0[e])
- tsh = tsh + _bits_f32(uh0[e])
- for e in cutlass.range_constexpr(4):
- tsg = tsg + _bits_f32(ug1[e])
- tsh = tsh + _bits_f32(uh1[e])
- st[2] = tsg
- st[3] = tsh
-
- # ---- epilogue -----------------------------------------------------------------------
- gv = mGs[2]
- lg = st[2]
- lh = st[3]
- ig = cutlass.Float32(0.0)
- ih = cutlass.Float32(0.0)
- if lg > 0.0:
- ig = gv / lg
- if lh > 0.0:
- ih = gv / lh
- hg = kvh * group + g
- hh = kvh * group + g + 8
- okh = (g + 8) < group
- for nt in cutlass.range_constexpr(4):
- d = 32 * warp + 8 * nt + 2 * c
- og0 = acc[nt, 0] * ig
- og1 = acc[nt, 1] * ig
- oh0 = acc[nt, 2] * ih
- oh1 = acc[nt, 3] * ih
- if const_expr(splits == 1):
- ag = out_base + (cutlass.Int64(row) * heads + hg) * (HEAD_DIM * 2) + d * 2
- ah = out_base + (cutlass.Int64(row) * heads + hh) * (HEAD_DIM * 2) + d * 2
- if const_expr(self.gated):
- gw = _ldg1(
- gate_base + (cutlass.Int64(row) * gate_stride + hg * HEAD_DIM + d) * 2
- )
- og0 = _round_bf16(og0) * _sigmoid(_bf2_lo(gw))
- og1 = _round_bf16(og1) * _sigmoid(_bf2_hi(gw))
- _stg1(ag, _bf2_pack(og0, og1))
- if okh:
- if const_expr(self.gated):
- gw2 = _ldg1(
- gate_base + (cutlass.Int64(row) * gate_stride + hh * HEAD_DIM + d) * 2
- )
- oh0 = _round_bf16(oh0) * _sigmoid(_bf2_lo(gw2))
- oh1 = _round_bf16(oh1) * _sigmoid(_bf2_hi(gw2))
- _stg1(ah, _bf2_pack(oh0, oh1))
- else:
- pg_ = (
- out_base
- + ((cutlass.Int64(split) * rows + row) * heads + hg) * (HEAD_DIM * 4)
- + d * 4
- )
- ph_ = (
- out_base
- + ((cutlass.Int64(split) * rows + row) * heads + hh) * (HEAD_DIM * 4)
- + d * 4
- )
- _stg2f(pg_, og0, og1)
- if okh:
- _stg2f(ph_, oh0, oh1)
- if const_expr(splits > 1):
- if warp == 0:
- if c == 0:
- lse_g = neg_inf
- if lg > 0.0:
- lse_g = st[0] + cute.math.log2(lg)
- _stg1(
- lse_base + ((cutlass.Int64(split) * rows + row) * heads + hg) * 4,
- _f32_bits(lse_g),
- )
- if okh:
- lse_h = neg_inf
- if lh > 0.0:
- lse_h = st[1] + cute.math.log2(lh)
- _stg1(
- lse_base + ((cutlass.Int64(split) * rows + row) * heads + hh) * 4,
- _f32_bits(lse_h),
- )
-
-
-@dsl_user_op
-def _f32_bits(x, *, loc=None, ip=None):
- return cutlass.Int32(
- _asm(T.i32(), [_ir(x, cutlass.Float32, loc, ip)], "mov.b32 $0, $1;", "=r,f", loc, ip)
- )
-
-
-@dsl_user_op
-def _bits_f32(x, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(T.f32(), [_ir(x, cutlass.Int32, loc, ip)], "mov.b32 $0, $1;", "=f,r", loc, ip)
- )
-
-
-@dsl_user_op
-def _round_bf16(x, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(
- T.f32(),
- [_ir(x, cutlass.Float32, loc, ip)],
- "{ .reg .b16 h; cvt.rn.bf16.f32 h, $1; cvt.f32.bf16 $0, h; }",
- "=f,f",
- loc,
- ip,
- )
- )
-
-
-def _sigmoid(x):
- """1 / (1 + exp(-x)) as the bf16 merge computes it."""
- return 1.0 / (1.0 + cute.math.exp(-cutlass.Float32(x)))
-
-
-# ---------------------------------------------------------------------------------------
-# merge kernel (splits > 1)
-# ---------------------------------------------------------------------------------------
-
-
-class _MergeKernel:
- def __init__(self, *, heads, splits, gated, pdl):
- assert HEAD_DIM % CHUNK == 0 and 1 < splits <= MAX_SPLITS
- self.heads, self.splits, self.gated, self.pdl = heads, splits, gated, pdl
- self.chunks = HEAD_DIM // CHUNK
- # at least four warps: the final pass needs a thread per element of the chunk
- self.warps = max(4, min(8, splits))
- self.spw = -(-splits // self.warps)
-
- @cute.jit
- def __call__(
- self,
- gate_t: cute.Tensor,
- part_ptr: cute.Pointer,
- lse_ptr: cute.Pointer,
- out_ptr: cute.Pointer,
- rows: cutlass.Int32,
- gate_stride: cutlass.Int32,
- stream,
- ):
- lse = cute.make_tensor(
- lse_ptr,
- cute.make_layout(
- (self.splits, rows, self.heads), stride=(rows * self.heads, self.heads, 1)
- ),
- )
- out = cute.make_tensor(
- out_ptr,
- cute.make_layout(
- (rows, self.heads, HEAD_DIM), stride=(self.heads * HEAD_DIM, HEAD_DIM, 1)
- ),
- )
- self.kernel(gate_t.iterator.toint(), part_ptr.toint(), lse, out, rows, gate_stride).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,
- gate_base: cutlass.Int64,
- part_base: cutlass.Int64,
- mLse,
- mOut,
- rows: cutlass.Int32,
- gate_stride: cutlass.Int32,
- ):
- 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,
- )
- gval = cutlass.Float32(0.0)
- if const_expr(self.gated):
- if tidx < CHUNK:
- gw = _ldg1(
- gate_base
- + (cutlass.Int64(row) * gate_stride + head * HEAD_DIM + c0 + (tidx & ~1)) * 2
- )
- gval = cutlass.Float32(cutlass.select_((tidx & 1) == 0, _bf2_lo(gw), _bf2_hi(gw)))
- if const_expr(self.pdl):
- cute.arch.griddepcontrol_launch_dependents()
- cute.arch.griddepcontrol_wait()
- 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]
- p = cute.make_rmem_tensor(cute.make_layout((spw, 4), stride=(4, 1)), cutlass.Float32)
- for jj in cutlass.range_constexpr(spw):
- sidx = cutlass.min(warp_i * spw + jj, self.splits - 1)
- w4 = _ldg4(
- part_base
- + ((cutlass.Int64(sidx) * rows + row) * self.heads + head) * (HEAD_DIM * 4)
- + (c0 + 4 * lane) * 4
- )
- for e in cutlass.range_constexpr(4):
- p[jj, e] = _bits_f32(w4[e])
- 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 jj in cutlass.range_constexpr(spw):
- sidx = warp_i * spw + jj
- 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[jj, e] + acc[e]
- for e in cutlass.range_constexpr(4):
- red[warp_i, 4 * lane + e] = acc[e]
- cute.arch.sync_threads()
- 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(-gval)))
- mOut[row, head, c0 + tidx] = merged.to(cutlass.BFloat16)
-
-
-# ---------------------------------------------------------------------------------------
-# launch
-# ---------------------------------------------------------------------------------------
-
-_COMPILED: dict = {}
-
-
-#: Splits by (rows x KV heads), measured on GB300 at 2051 selected tokens (sweep_splits.py):
-#: about 250-290 CTAs where the merge is cheap, a single split from 128 rows up.
-_SPLIT_TABLE = ((4, 33), (8, 11), (16, 17), (32, 9), (64, 4), (128, 2))
-
-
-def plan_splits(rows: int, kv_heads: int, width: int) -> tuple[int, int]:
- """(splits, tiles per split) for this many (row, KV head) pairs."""
- n_tiles = -(-width // TILE)
- pairs = rows * kv_heads
- want = 1
- for bound, sp in _SPLIT_TABLE:
- if pairs <= bound:
- want = sp
- break
- want = max(1, min(n_tiles, want, MAX_SPLITS))
- tps = -(-n_tiles // want)
- return -(-n_tiles // tps), tps
-
-
-def _ptr(t: torch.Tensor, dtype, align=16):
- return make_ptr(dtype, t.data_ptr(), cute.AddressSpace.gmem, assumed_align=align)
-
-
-def nvfp4_sparse_gqa_decode(
- 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,
- output_gate: Optional[torch.Tensor] = None,
- pdl: bool = True,
- splits: Optional[int] = None,
-) -> torch.Tensor:
- """Decode-path NVFP4 sparse GQA. Returns `[rows, Hq, D]` bf16 (contiguous); with
- `output_gate` (`[rows, Hq * D]` bf16, unit column stride) the result is
- `bf16(attn) * sigmoid(gate)` rounded to bf16, as `qsa_paged_sparse_gqa_decode` fuses it."""
- rows, num_q_heads, head_dim = q.shape
- num_pages, num_kv_heads, tpb, packed = k_cache.shape
- assert head_dim == HEAD_DIM and packed == HEAD_DIM // 2 and tpb == tokens_per_block
- assert q.dtype == torch.bfloat16 and q.stride(2) == 1
- assert k_cache.dtype == torch.uint8 and v_cache.shape == k_cache.shape
- assert k_cache.stride(3) == 1 and v_cache.stride(3) == 1
- assert k_scale_cache.stride(2) == HEAD_DIM // 16 and v_scale_cache.stride(2) == HEAD_DIM // 16
- assert block_table.dtype == torch.int32 and block_table.stride(1) == 1
- assert block_table.stride(0) == block_table.shape[1], "row-contiguous block table"
- assert selected_tokens.dtype == torch.int32 and selected_tokens.is_contiguous()
- assert request_indices.is_contiguous() and request_indices.dtype in (torch.int32, torch.int64)
- if selected_tokens.data_ptr() % 16:
- # The kernel reads the selection with 16-byte loads. A row slice of a 2051-wide
- # selection (a mixed step's generation rows after its context rows) starts 16-byte
- # aligned only at every fourth row, so an unaligned one is copied (rows x 8 KB).
- selected_tokens = selected_tokens.clone()
- group = num_q_heads // num_kv_heads
- width = selected_tokens.shape[1]
- out = torch.empty((rows, num_q_heads, head_dim), dtype=torch.bfloat16, device=q.device)
- if rows == 0:
- return out
- if splits is None:
- n_splits, tps = plan_splits(rows, num_kv_heads, width)
- else:
- n_tiles = -(-width // TILE)
- tps = -(-n_tiles // splits)
- n_splits = -(-n_tiles // tps)
- use_pdl = bool(pdl and get_env_enable_pdl())
- gated = output_gate is not None
- if gated:
- assert output_gate.shape == (rows, num_q_heads * head_dim) and output_gate.stride(1) == 1
- assert output_gate.dtype == torch.bfloat16
- req64 = request_indices.dtype == torch.int64
- key = (
- group,
- num_kv_heads,
- width,
- tpb,
- n_splits,
- tps,
- float(softmax_scale),
- num_pages,
- tuple(block_table.shape),
- tuple(k_cache.stride()[:3]),
- tuple(v_cache.stride()[:3]),
- tuple(k_scale_cache.stride()[:2]),
- tuple(v_scale_cache.stride()[:2]),
- tuple(q.stride()[:2]),
- req64,
- gated and n_splits == 1,
- use_pdl,
- )
- if n_splits > 1:
- part = torch.empty(
- (n_splits, rows, num_q_heads, head_dim), dtype=torch.float32, device=q.device
- )
- lse = torch.empty((n_splits, rows, num_q_heads), dtype=torch.float32, device=q.device)
- o_t, l_t = part, lse
- else:
- o_t, l_t = out, out
- gate_t = output_gate if gated else q
- gate_stride = output_gate.stride(0) if gated else 0
- fn = _COMPILED.get(key)
- if fn is None:
- kern = _SplitKernel(
- group=group,
- kv_heads=num_kv_heads,
- width=width,
- tpb=tpb,
- splits=n_splits,
- tps=tps,
- scale=float(softmax_scale),
- num_pages=num_pages,
- bt_width=block_table.shape[1],
- num_requests=block_table.shape[0],
- k_strides=tuple(k_cache.stride()[:3]),
- v_strides=tuple(v_cache.stride()[:3]),
- ks_strides=tuple(k_scale_cache.stride()[:2]),
- vs_strides=tuple(v_scale_cache.stride()[:2]),
- q_strides=tuple(q.stride()[:2]),
- req64=req64,
- gated=gated and n_splits == 1,
- pdl=use_pdl,
- )
- fn = cute.compile(
- kern,
- from_dlpack(selected_tokens, assumed_align=16).mark_layout_dynamic(leading_dim=1),
- _ptr(q, cutlass.BFloat16),
- _ptr(k_cache, cutlass.Uint8),
- _ptr(v_cache, cutlass.Uint8),
- _ptr(k_scale_cache, cutlass.Uint8),
- _ptr(v_scale_cache, cutlass.Uint8),
- _ptr(kv_scale_quant_orig, cutlass.Float32, 4),
- _ptr(block_table, cutlass.Int32, 4),
- _ptr(request_indices, cutlass.Int64 if req64 else cutlass.Int32, 8 if req64 else 4),
- _ptr(o_t, cutlass.Float32),
- _ptr(l_t, cutlass.Float32),
- _ptr(gate_t, cutlass.BFloat16),
- cutlass.Int32(rows),
- cutlass.Int32(gate_stride),
- cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
- options="--opt-level 2 --enable-tvm-ffi",
- )
- _COMPILED[key] = fn
- fn(
- selected_tokens,
- q.data_ptr(),
- k_cache.data_ptr(),
- v_cache.data_ptr(),
- k_scale_cache.data_ptr(),
- v_scale_cache.data_ptr(),
- kv_scale_quant_orig.data_ptr(),
- block_table.data_ptr(),
- request_indices.data_ptr(),
- o_t.data_ptr(),
- l_t.data_ptr(),
- gate_t.data_ptr(),
- rows,
- gate_stride,
- )
- if n_splits == 1:
- return out
- mkey = ("merge", num_q_heads, n_splits, gated, use_pdl)
- mfn = _COMPILED.get(mkey)
- if mfn is None:
- mfn = cute.compile(
- _MergeKernel(heads=num_q_heads, splits=n_splits, gated=gated, pdl=use_pdl),
- from_dlpack(gate_t, assumed_align=16).mark_layout_dynamic(leading_dim=gate_t.dim() - 1),
- _ptr(part, cutlass.Float32),
- _ptr(lse, cutlass.Float32, 4),
- _ptr(out, cutlass.BFloat16),
- cutlass.Int32(rows),
- cutlass.Int32(gate_stride),
- cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
- options="--opt-level 2 --enable-tvm-ffi",
- )
- _COMPILED[mkey] = mfn
- mfn(gate_t, part.data_ptr(), lse.data_ptr(), out.data_ptr(), rows, gate_stride)
- return out
-
-
-def _decode(segments, **overrides) -> dict:
- """A cell spec: `(first position, rows)` per request in row order, this checkpoint's
- geometry (24 q heads over 2 KV heads of 256, 32 tokens a page, 2051 selected tokens)."""
- return dict(segments=segments, **overrides)
-
-
-class QsaPagedSparseGqaNvfp4Decode(OpWrapper):
- """`qsa_paged_sparse_gqa_nvfp4` for decode rows (no prefill pair): sparse GQA over the
- explicitly selected tokens of an NVFP4 paged pool, `[rows, Hq, D]` bf16. With
- `output_gate` (`[rows, Hq * D]` bf16, unit column stride) the result is
- `bf16(attn) * sigmoid(gate)` rounded once more to bf16, the composition
- `fused_sigmoid_mul_inplace` makes.
-
- Fusion boundary: the block-table gather, the NVFP4 dequantize, scale, online softmax,
- P.V, the split-K merge and (optionally) the output gate. Selection, the K/V append and
- q's projection, norm and RoPE are not inside.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(why="decode: 1 row, 33-way split + merge", spec=_decode(((4000, 1),))),
- Cell(
- why="MTP3 verify: 4 requests x 4 rows, gate fused into the merge",
- spec=_decode(tuple((4000 + 37 * i, 4) for i in range(4)), gated=True),
- ),
- Cell(
- why="128 rows (16 x 8), one split: gate fused into the split kernel, trunk layout",
- spec=_decode(
- tuple((100 + (499 * g) % 8000, 8) for g in range(16)),
- gated=True,
- slot_pages=24,
- layer=11,
- pool_pages=2048,
- ),
- ),
- )
-
- note = """
- Metadata consumed: none -- every input is an argument (as `qsa_paged_sparse_gqa_nvfp4`).
- Call sites (the target's `modeling.py`): `_sparse_gqa` over an NVFP4 pool for rows without
- the prefill pair -- the decode step's attention layers (with the output gate), the MTP
- draft layer and a mixed step's generation rows.
-
- Launch shapes: splits by (rows x KV heads) from `_SPLIT_TABLE` (33 splits of one 64-token
- tile at one row, a single split of 33 tiles from 128 rows up), each split followed by the
- merge. GB300, 2051 selected tokens a row, one layer: 1 row 6.7 us, 4 rows 8.5, 32 rows 22.3,
- 128 rows 67.8, 256 rows 133, 512 rows 231, 1024 rows 449 -- against 9.0 / 18.8 / 49.5 /
- 158 / 292 / 557 / 1036 us for the Triton op. A one-split launch (65 rows and up) runs the
- kernel without the f16 V tile: 72 rows 65.1 -> 61.7 us, 128 rows 71.6 -> 66.6, 192 rows
- 111.4 -> 101.6, 256 rows 137.5 -> 128.3 (12 chained calls, gated), bitwise the standard
- kernel.
-
- Numerics: dequantized K/V values are exact in f16; scores on f16 MMAs with fp32
- accumulation; P rounded to f16 (as the Triton op); the second-level scales act on the
- scores and the normalized output. Within the Triton op's band (4 ulp element, 2 ulp RMS)
- of the float64 reference; not bitwise the Triton op. Deterministic; graph-capturable
- (row count is a runtime value; one compile per split variant).
-
- Preconditions: bf16 q with head_dim 256; K and V planes (and their scale planes) share
- strides; power-of-two tokens per block; int32 row-contiguous block table; contiguous
- int32 selections.
- """
-
- def __call__(
- self,
- q,
- k_cache,
- v_cache,
- k_scale_cache,
- v_scale_cache,
- kv_scale_quant_orig,
- block_table,
- selected_tokens,
- request_indices,
- tokens_per_block: int,
- softmax_scale: float,
- output_gate=None,
- ) -> torch.Tensor:
- return nvfp4_sparse_gqa_decode(
- q,
- k_cache,
- v_cache,
- k_scale_cache,
- v_scale_cache,
- kv_scale_quant_orig,
- block_table,
- selected_tokens,
- request_indices,
- tokens_per_block,
- softmax_scale,
- output_gate=output_gate,
- )
-
- def reference(
- self,
- q,
- k_cache,
- v_cache,
- k_scale_cache,
- v_scale_cache,
- kv_scale_quant_orig,
- block_table,
- selected_tokens,
- request_indices,
- tokens_per_block: int,
- softmax_scale: float,
- output_gate=None,
- ) -> torch.Tensor:
- """`qsa_paged_sparse_gqa_nvfp4`'s float64 reference, then the gate in fp32 over its
- bf16 rounding, rounded once."""
- out = _composition.reference(
- q,
- k_cache,
- v_cache,
- k_scale_cache,
- v_scale_cache,
- kv_scale_quant_orig,
- block_table,
- selected_tokens,
- request_indices,
- tokens_per_block,
- softmax_scale,
- )
- if output_gate is None:
- return out
- rows, heads, dim = out.shape
- gate = torch.sigmoid(output_gate.float().view(rows, heads, dim))
- return (out.to(torch.bfloat16).float() * gate).to(torch.bfloat16)
-
- def is_valid(
- self,
- q,
- k_cache,
- v_cache,
- k_scale_cache,
- v_scale_cache,
- kv_scale_quant_orig,
- block_table,
- selected_tokens,
- request_indices,
- tokens_per_block: int,
- softmax_scale: float,
- output_gate=None,
- ) -> None:
- assert q.dtype == torch.bfloat16 and q.dim() == 3 and q.shape[2] == HEAD_DIM
- assert q.stride(2) == 1, "q unit last stride"
- assert k_cache.dtype == torch.uint8 and v_cache.shape == k_cache.shape
- assert k_cache.stride() == v_cache.stride(), "K and V planes share strides"
- assert k_scale_cache.stride() == v_scale_cache.stride(), "scale planes share strides"
- assert tokens_per_block & (tokens_per_block - 1) == 0, "power-of-two tokens per block"
- assert block_table.dtype == torch.int32 and block_table.is_contiguous()
- assert selected_tokens.dtype == torch.int32 and selected_tokens.is_contiguous()
- assert output_gate is None or (
- output_gate.dtype == torch.bfloat16
- and output_gate.shape == (q.shape[0], q.shape[1] * q.shape[2])
- and output_gate.stride(1) == 1
- )
-
-
-qsa_paged_sparse_gqa_nvfp4_decode = QsaPagedSparseGqaNvfp4Decode()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4_verify.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4_verify.py
deleted file mode 100644
index a71036db0466..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4_verify.py
+++ /dev/null
@@ -1,2147 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""QSA sparse GQA over an NVFP4 paged K/V pool for MTP verify rows, request-grouped: CuTe DSL
-(tcgen05) split-K, then a row merge of the split partials.
-
-Tier 5 (CuTe DSL written for this target). A verify step's rows come `T`
-(`rows_per_request`) consecutive rows a request, and the rows of one request select mostly
-the same tokens. One CTA serves one (group of T rows, KV head, split): 128 MMA rows, row
-m = t * 12 + h (verify row t, q head h of the KV head; rows past T * 12 idle), over the
-group's *union* selection, so each selected K/V token is loaded and dequantized once for the
-group's rows instead of once a row.
-
-Union (every thread, one code copy per pass): per-row token bitmaps in shared memory (the
-request's block-table row fetched in the same round trip), the union of four-token groups in
-ascending order (an expanded selection keeps column % 4 == token % 4), each group's K/V row
-offset and scale offset resolved; per tile a 64-bit selection word for each verify row. A
-(row, position) pair the row did not select is masked to -inf before its max, so the softmax
-offset is the row's own max and causality is the selection's. Splits take interleaved tiles of
-the union. Rows of a group that belong to different requests (a partial or misaligned group)
-are served in separate passes; rows of an invalid request give zeros.
-
-Launch (`plan_launch`): within one wave every (group, KV head) unit takes the same number of
-splits; past one wave the units of the full waves take one CTA each (bf16 output, the gate
-fused) and the units of the last, partial wave -- whole groups -- take two to four CTAs each,
-whose rows alone go through the merge.
-
-Roles (11 warps):
-
-* warps 5-6 `cp.async` each tile's raw NVFP4 K / V rows (eight threads a 128-byte row), K scales
- and V-scale groups into a ring of `RAW_SLOTS` raw slots, completing on the slot's mbarrier;
-* warps 7-10 (one thread a position, K then V) load a raw row and its scales into registers,
- release the raw slot, and dequantize into f16 SW128 K and V slots (`cvt.rn.f16x2.e2m1x2` x
- `cvt.rn.f16x2.e4m3x2`, exact, as the per-row kernel);
-* warp 4 issues the tcgen05 MMAs with precomputed descriptors (`qsa_nvfp4_verify_ptx`):
- S[j % 2] = Q K(j)^T (Q from tensor memory), then O += P(j - 1) V(j - 1) (P from tensor
- memory, over the S it came from);
-* warps 0-3 (one thread an MMA row) copy their q row into a raw slot at once and into tensor
- memory (f16) after the first resolve; per tile load S, take the max over the row's own
- entries, move the row's reference max only when it grew past `RESCALE_TH` (log2), rescaling
- O in tensor memory then, and write P = 2^(s - max) as f16 (0 off the row's selection);
- finally normalize O (times the V global scale) and store it as bf16 (or an fp32 partial and
- its log2-sum-exp for the merge). With the output gate the four warps read the CTA's gate
- rows (coalesced) before the last P V completes, the live warps stage bf16(O / l) in the free
- K slots, and all four apply the gate and store coalesced rows.
-
-Tensor memory: O in columns 0-255, S/P buffers at 256 and 320, Q at 384. Shared memory: two f16
-K and two f16 V slots (32 KB each), three raw slots, the union tables (up to ~225 KB): one CTA
-an SM.
-"""
-
-from __future__ import annotations
-
-import functools
-from typing import NamedTuple, Optional
-
-import cutlass
-import cutlass.cute as cute
-import cutlass.cute.nvgpu.tcgen05 as tcgen05
-import cutlass.utils.blackwell_helpers as sm100_utils
-import torch
-from cutlass import const_expr
-from cutlass.cute.nvgpu import OperandMajorMode
-from cutlass.cute.runtime import from_dlpack
-from cutlass.cute.typing import Float32, Int32, Int64
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
-
-from .qsa_nvfp4_verify_ptx import (
- ctz,
- gemm_ts,
- groups8,
- ldg2,
- popc,
- red_or_shared,
- spread4,
- stg2,
- stg4,
- stg4f,
- tcgen05_fence_after,
- tcgen05_fence_before,
-)
-from .qsa_paged_sparse_gqa_nvfp4_decode import (
- MAX_SPLITS,
- _bf2_hi,
- _bf2_lo,
- _bf2_pack,
- _bf2_to_f2,
- _bits_f32,
- _cp16z,
- _dq_word,
- _f32_bits,
- _ldg1,
- _ldg4,
- _lds1,
- _lds4,
- _ptr,
- _round_bf16,
- _scale2s,
- _stg1,
- _sts1,
- _sts4,
- nvfp4_sparse_gqa_decode,
- qsa_paged_sparse_gqa_nvfp4_decode,
-)
-
-try:
- from cutlass.memory import SmemAllocator
-except ImportError: # DSL releases before the memory namespace split
- from cutlass.utils import SmemAllocator
-
-F16 = cutlass.Float16
-LOG2E = 1.4426950408889634
-HEAD_DIM = 256
-#: MMA rows a CTA (T verify rows x the KV head's q heads) and positions a tile.
-BM = 128
-BN = 64
-#: Four-token groups a tile.
-GPT = BN // 4
-#: Bytes of a raw K or V row in shared memory (128 + 16 of padding: conflict-free reads).
-RAW_LD = 144
-RAW_SLOTS = 3
-RAW_K = BN * RAW_LD
-RAW_KS = BN * 16
-RAW_VS = GPT * 64
-#: One raw slot: K rows, V rows, K-scale rows, V-scale groups.
-RAW_BUF = 2 * RAW_K + RAW_KS + RAW_VS
-#: Bytes of the four f16 K / V slots (two each, 64 positions x 256 dims).
-KV_SLOTS_BYTES = 4 * BN * HEAD_DIM * 2
-#: Shared memory a CTA may claim (227 KB), less barriers, scan scratch and alignment.
-SMEM_BUDGET = 232448 - 4096
-#: A row's reference max moves only past this many log2 units (P <= 2^8, exact in f16 range).
-RESCALE_TH = 8.0
-SOFT_WARPS = 4
-MMA_WARP = 4
-LOAD_WARP0 = 5
-LOAD_WARPS = 2
-DQ_WARP0 = 7
-DQ_WARPS = 4
-THREADS = 32 * (DQ_WARP0 + DQ_WARPS)
-TMEM_O = 0
-TMEM_S0 = 256
-TMEM_Q = 384
-TMEM_COLS = 512
-CHUNK_COLS = 32
-#: Bytes of a staged q row (bf16, 512 + 16 of padding).
-Q_LD = 528
-#: Most splits of one (group, KV head) union.
-MAX_PLAN_SPLITS = 32
-#: Most splits of a last-wave unit (tail mode).
-MAX_TAIL_SPLITS = 4
-#: Fewer request groups than this take the per-row kernel (GB300: it is as fast or faster
-#: below 32 rows at four rows a request, the grouped kernel from 32 rows up).
-MIN_GROUPS = 8
-
-
-class _Layout(NamedTuple):
- """The shared-memory plan of one compile: bitmap words (`ngw` groups of four), the union's
- group capacity, local tiles a CTA may own in one pass, union-scratch offsets in raw slot 0
- and the table size in int32 words.
- """
-
- ngw: int
- nw: int
- ug_max: int
- tmax: int
- u_ug: int
- u_bt: int
- tab_words: int
-
-
-def _layout(rpr: int, group: int, width: int, bt_width: int, tpb: int, splits: int):
- """The kernel's shared-memory plan for `rpr` rows a group over a `width`-column expanded
- selection and a `bt_width`-page block table, or None where it does not fit: the union
- scratch (each row's token bitmap, the union's group list and the request's block-table
- row) shares raw slot 0, the staged q rows raw slots 1-2.
- """
- if tpb % 4 or tpb & (tpb - 1) or rpr < 1 or rpr * group > BM:
- return None
- ngw = -(-(bt_width * tpb) // 128)
- nw = 4 * ngw
- # an expanded selection names at most ceil(width / 4) four-token groups a row
- ug_max = min(rpr * -(-width // 4), 32 * ngw)
- tmax = -(-(-(-ug_max // GPT)) // splits)
- u_ug = -(-(rpr * nw * 4) // 16) * 16
- u_bt = -(-(u_ug + ug_max * 4) // 16) * 16
- tab_words = 3 * tmax * GPT + 2 * tmax * rpr
- fits = (
- ngw <= THREADS
- and u_bt + bt_width * 4 <= RAW_BUF
- and RAW_BUF + rpr * group * Q_LD <= RAW_SLOTS * RAW_BUF
- and 4 * tab_words + RAW_SLOTS * RAW_BUF + KV_SLOTS_BYTES <= SMEM_BUDGET
- )
- return _Layout(ngw, nw, ug_max, tmax, u_ug, u_bt, tab_words) if fits else None
-
-
-def _sigmoid_fast(x):
- """1 / (1 + 2^(-x log2 e)) on ex2.approx and rcp.approx (relative error ~1e-6)."""
- return cute.arch.rcp_approx(1.0 + cute.math.exp2(-x * LOG2E, fastmath=True))
-
-
-def _tree_fmax(vals):
- """Max of a short list, pairwise (traced as Python)."""
- while len(vals) > 1:
- vals = [
- cute.arch.fmax(vals[j], vals[j + 1]) if j + 1 < len(vals) else vals[j]
- for j in range(0, len(vals), 2)
- ]
- return vals[0]
-
-
-def _tree_sum(vals):
- """Sum of a short list, pairwise (traced as Python)."""
- while len(vals) > 1:
- vals = [
- vals[j] + vals[j + 1] if j + 1 < len(vals) else vals[j] for j in range(0, len(vals), 2)
- ]
- return vals[0]
-
-
-def _phase(g, n):
- """Parity of use `g` of a ring of `n` slots."""
- return (g // n) % 2
-
-
-class _VerifyKernel:
- def __init__(
- self,
- *,
- rpr,
- group,
- kv_heads,
- width,
- tpb,
- splits,
- tail,
- scale,
- num_pages,
- bt_width,
- num_requests,
- k_strides,
- ks_strides,
- q_strides,
- req64,
- gated,
- pdl,
- ):
- lay = _layout(rpr, group, width, bt_width, tpb, 1 if tail else splits)
- assert lay is not None, "no shared-memory plan for this geometry (see `_layout`)"
- self.T = rpr
- self.group, self.kvh, self.width, self.tpb = group, kv_heads, width, tpb
- self.splits = splits
- #: tail mode: CTAs below `n_whole` serve whole units, the rest `splits`-way splits
- self.tail = tail
- self.scale_log2 = scale * LOG2E
- self.num_pages, self.bt_width, self.num_requests = num_pages, bt_width, num_requests
- self.ks, self.kss, self.qs = k_strides, ks_strides, q_strides
- self.req64, self.gated, self.pdl = req64, gated, pdl
- self.heads = group * kv_heads
- self.rows_m = rpr * group
- #: softmax warps holding live MMA rows (the rest only keep the barriers' counts)
- self.live_warps = -(-self.rows_m // 32)
- self.ngw, self.nw, self.ug_max, self.tmax = lay.ngw, lay.nw, lay.ug_max, lay.tmax
- self.raw_k = RAW_K
- self.raw_v = RAW_K
- self.raw_ks = RAW_KS
- self.raw_vs = RAW_VS
- self.raw_buf = RAW_BUF
- self.raw_bytes = RAW_SLOTS * RAW_BUF
- # union scratch (raw slot 0, free between passes): bitmaps, group list, block-table row
- self.u_bm = 0
- self.u_ug = lay.u_ug
- self.u_bt = lay.u_bt
- # q rows (bf16, 528-byte rows: conflict-free) in raw slots 1-2 until Q is in tensor memory
- self.u_q = RAW_BUF
- # tables (int32 words). Per group: K/V row offset and scale offset of its first token
- # (16-byte units, -1 when the group does not resolve) and its (token, row) selection
- # bytes; per (tile, row): the 64-bit selection word.
- self.t_roff = 0
- self.t_soff = self.tmax * GPT
- self.t_gmask = 2 * self.tmax * GPT
- self.t_bits = 3 * self.tmax * GPT
- self.tab_words = lay.tab_words
-
- @cute.jit
- def __call__(
- self,
- sel_t: cute.Tensor,
- q_ptr: cute.Pointer,
- k_ptr: cute.Pointer,
- v_ptr: cute.Pointer,
- ks_ptr: cute.Pointer,
- vs_ptr: cute.Pointer,
- gs_ptr: cute.Pointer,
- bt_ptr: cute.Pointer,
- req_ptr: cute.Pointer,
- out_ptr: cute.Pointer,
- part_ptr: cute.Pointer,
- lse_ptr: cute.Pointer,
- gate_ptr: cute.Pointer,
- rows: cutlass.Int32,
- gate_stride: cutlass.Int32,
- n_whole: cutlass.Int32,
- tail_row0: cutlass.Int32,
- n_ctas: cutlass.Int32,
- stream,
- ):
- def qk():
- return sm100_utils.make_trivial_tiled_mma(
- F16,
- F16,
- OperandMajorMode.K,
- OperandMajorMode.K,
- Float32,
- tcgen05.CtaGroup.ONE,
- (BM, BN),
- tcgen05.OperandSource.TMEM,
- )
-
- def pv():
- return sm100_utils.make_trivial_tiled_mma(
- F16,
- F16,
- OperandMajorMode.K,
- OperandMajorMode.MN,
- Float32,
- tcgen05.CtaGroup.ONE,
- (BM, HEAD_DIM),
- tcgen05.OperandSource.TMEM,
- )
-
- qk_mma, qk_mma_acc, pv_mma_acc = qk(), qk(), pv()
- q_layout = sm100_utils.make_smem_layout_a(qk_mma, (BM, BN, HEAD_DIM), F16, 1)
- k_layout = sm100_utils.make_smem_layout_b(qk_mma, (BM, BN, HEAD_DIM), F16, 2)
- v_layout = sm100_utils.make_smem_layout_b(pv_mma_acc, (BM, HEAD_DIM, BN), F16, 2)
- p_layout = sm100_utils.make_smem_layout_a(pv_mma_acc, (BM, HEAD_DIM, BN), F16, 1)
-
- @cute.struct
- class SharedStorage:
- q_full: cute.struct.MemRange[Int64, 1]
- raw_full: cute.struct.MemRange[Int64, RAW_SLOTS]
- raw_empty: cute.struct.MemRange[Int64, RAW_SLOTS]
- k_full: cute.struct.MemRange[Int64, 2]
- k_empty: cute.struct.MemRange[Int64, 2]
- v_full: cute.struct.MemRange[Int64, 2]
- v_empty: cute.struct.MemRange[Int64, 2]
- 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
- scan: cute.struct.MemRange[Int32, 32]
- tab: cute.struct.Align[cute.struct.MemRange[Int32, self.tab_words], 16]
- raw: cute.struct.Align[cute.struct.MemRange[cutlass.Int8, self.raw_bytes], 128]
- sK: cute.struct.Align[cute.struct.MemRange[F16, cute.cosize(k_layout)], 1024]
- sV: cute.struct.Align[cute.struct.MemRange[F16, cute.cosize(k_layout)], 1024]
-
- self.shared_storage = SharedStorage
- bt = cute.make_tensor(
- bt_ptr, cute.make_layout((self.num_requests, self.bt_width), stride=(self.bt_width, 1))
- )
- req = cute.make_tensor(req_ptr, cute.make_layout((rows,), stride=(1,)))
- gs = cute.make_tensor(gs_ptr, cute.make_layout((3,), stride=(1,)))
- self.kernel(
- qk_mma,
- qk_mma_acc,
- pv_mma_acc,
- q_layout,
- k_layout,
- v_layout,
- p_layout,
- sel_t.iterator.toint(),
- q_ptr.toint(),
- k_ptr.toint(),
- v_ptr.toint(),
- ks_ptr.toint(),
- vs_ptr.toint(),
- gs,
- bt,
- req,
- out_ptr.toint(),
- part_ptr.toint(),
- lse_ptr.toint(),
- gate_ptr.toint(),
- rows,
- gate_stride,
- n_whole,
- tail_row0,
- ).launch(
- grid=[n_ctas, 1, 1],
- block=[THREADS, 1, 1],
- stream=stream,
- min_blocks_per_mp=1,
- use_pdl=self.pdl,
- )
-
- # ---- pieces every role runs (every thread of the CTA, barrier 0) -----------------------
-
- @cute.jit
- def _pass_rows(self, rq: cute.Tensor, handled):
- lead = cutlass.Int32(-1)
- for t in tuple(range(self.T - 1, -1, -1)):
- lead = cutlass.Int32(cutlass.select_(((handled >> t) & 1) == 0, rq[t], lead))
- act = cutlass.Int32(0)
- for t in cutlass.range_constexpr(self.T):
- act = act | cutlass.Int32(
- cutlass.select_(
- (((handled >> t) & 1) == 0) & (rq[t] == lead),
- cutlass.Int32(1 << t),
- cutlass.Int32(0),
- )
- )
- return lead, act
-
- @cute.jit
- def _union(
- self, act, lead, r0, rows, sel_base, tidx, warp, lane, s_bm, s_scan, s_ug, s_bt, mBT
- ):
- rpr: cutlass.Constexpr = self.T
- nw: cutlass.Constexpr = self.nw
- ngw: cutlass.Constexpr = self.ngw
- width: cutlass.Constexpr = self.width
- ntok: cutlass.Constexpr = 32 * self.nw
- n_q: cutlass.Constexpr = -(-width // 4) # four-column quads a row
- n_i: cutlass.Constexpr = -(-n_q // THREADS)
- cute.arch.sync_threads()
- for i in cutlass.range_constexpr(-(-(rpr * nw) // THREADS)):
- idx = tidx + i * THREADS
- if idx < rpr * nw:
- _sts1(s_bm + idx * 4, cutlass.Int32(0))
- # thread: columns 4q .. 4q + 3 of each row (one selected group in an expanded selection);
- # every load first, then one bit set a distinct word
- toks = []
- for t in cutlass.range_constexpr(rpr):
- row_sel = sel_base + cutlass.Int64(cutlass.min(r0 + t, rows - 1)) * (width * 4)
- for i in cutlass.range_constexpr(n_i):
- for u in cutlass.range_constexpr(4):
- colc = cutlass.min(4 * (tidx + i * THREADS) + u, width - 1)
- toks.append(_ldg1(row_sel + cutlass.Int64(colc * 4)))
- # the request's block-table row, in the same round trip
- n_bt: cutlass.Constexpr = -(-self.bt_width // THREADS)
- btv = [mBT[lead, cutlass.min(tidx + i * THREADS, self.bt_width - 1)] for i in range(n_bt)]
- cute.arch.sync_threads()
- for t in cutlass.range_constexpr(rpr):
- for i in cutlass.range_constexpr(n_i):
- q4 = tidx + i * THREADS
- live_t = ((act >> t) & 1) != 0
- okv = []
- wds = []
- bv = []
- for u in cutlass.range_constexpr(4):
- tok = toks[(t * n_i + i) * 4 + u]
- okv.append(live_t & (4 * q4 + u < width) & (tok >= 0) & (tok < ntok))
- wds.append(tok >> 5)
- bv.append(
- cutlass.Int32(
- cutlass.select_(
- okv[u], cutlass.Int32(1) << (tok & 31), cutlass.Int32(0)
- )
- )
- )
- # bits sharing token 0's word go in one atomic; any other token on its own
- m0 = bv[0]
- for u in cutlass.range_constexpr(1, 4):
- m0 = m0 | cutlass.Int32(
- cutlass.select_(wds[u] == wds[0], bv[u], cutlass.Int32(0))
- )
- if okv[0] | (m0 != 0):
- if m0 != 0:
- red_or_shared(s_bm + (t * nw + cutlass.max(wds[0], 0)) * 4, m0)
- for u in cutlass.range_constexpr(1, 4):
- if okv[u] & (wds[u] != wds[0]):
- red_or_shared(s_bm + (t * nw + wds[u]) * 4, bv[u])
- for i in cutlass.range_constexpr(n_bt):
- if tidx + i * THREADS < self.bt_width:
- _sts1(s_bt + (tidx + i * THREADS) * 4, btv[i])
- cute.arch.sync_threads()
- gmv = cutlass.Int32(0)
- if tidx < ngw:
- for k in cutlass.range_constexpr(4):
- xo = cutlass.Int32(0)
- for t in cutlass.range_constexpr(rpr):
- xo = xo | _lds1(s_bm + (t * nw + 4 * tidx + k) * 4)
- gmv = gmv | (groups8(xo) << (8 * k))
- cnt = popc(gmv)
- incl = cnt
- for o in (1, 2, 4, 8, 16):
- up = cute.arch.shuffle_sync_up(incl, offset=o)
- if lane >= o:
- incl = incl + up
- if lane == 31:
- _sts1(s_scan + warp * 4, incl)
- cute.arch.sync_threads()
- wsum = cutlass.Int32(0)
- tot = cutlass.Int32(0)
- if lane < THREADS // 32:
- wv = _lds1(s_scan + lane * 4)
- tot = wv
- if lane < warp:
- wsum = wv
- before = cute.arch.warp_reduction_sum(wsum)
- n_ug = cute.arch.warp_reduction_sum(tot)
- excl = before + incl - cnt
- if tidx < ngw:
- xm = gmv
- for j in cutlass.range(cnt):
- b = ctz(xm)
- if excl + j < self.ug_max:
- _sts1(s_ug + (excl + j) * 4, 32 * tidx + b)
- xm = xm & (xm - 1)
- cute.arch.sync_threads()
- return cutlass.min(n_ug, self.ug_max)
-
- @cute.jit
- def _resolve(self, lead, kvh, split, nsplit, n_loc, n_ug, tidx, s_bm, s_ug, s_tab, s_bt):
- """Every local tile's groups (entry of the group's first token, (token, row) selection
- bytes), then each (tile, verify row)'s 64-bit selection word.
- """
- rpr: cutlass.Constexpr = self.T
- nw: cutlass.Constexpr = self.nw
- tmax: cutlass.Constexpr = self.tmax
- n_rs: cutlass.Constexpr = -(-(tmax * GPT) // THREADS)
- s_roff = s_tab + self.t_roff * 4
- s_soff = s_tab + self.t_soff * 4
- s_gmask = s_tab + self.t_gmask * 4
- s_bits = s_tab + self.t_bits * 4
- ks0, ks1, ks2 = self.ks
- kss0, kss1 = self.kss
- pages = []
- gids = []
- for i in cutlass.range_constexpr(n_rs):
- slot = tidx + i * THREADS
- kl = slot // GPT
- ugi = (split + kl * nsplit) * GPT + slot % GPT
- gid = _lds1(s_ug + cutlass.min(ugi, self.ug_max - 1) * 4)
- gid = cutlass.min(cutlass.max(gid, 0), 8 * nw - 1)
- gids.append(gid)
- pages.append(_lds1(s_bt + cutlass.min((gid * 4) // self.tpb, self.bt_width - 1) * 4))
- for i in cutlass.range_constexpr(n_rs):
- slot = tidx + i * THREADS
- kl = slot // GPT
- ugi = (split + kl * nsplit) * GPT + slot % GPT
- gid = gids[i]
- page = pages[i]
- wsel = gid >> 3
- sh = (gid & 7) * 4
- mw = cutlass.Int32(0)
- for t in cutlass.range_constexpr(rpr):
- nib = (_lds1(s_bm + (t * nw + wsel) * 4) >> sh) & 0xF
- mw = mw | (spread4(nib) << t)
- okp = (kl < n_loc) & (ugi < n_ug) & (page >= 0) & (page < self.num_pages)
- pg64 = cutlass.Int64(cutlass.max(page, 0))
- sl0 = (gid * 4) % self.tpb
- roff = (pg64 * ks0 + cutlass.Int64(kvh * ks1 + sl0 * ks2)) >> 4
- soff = (pg64 * kss0 + cutlass.Int64(kvh * kss1 + sl0 * 16)) >> 4
- roff = cutlass.Int32(cutlass.select_(okp, cutlass.Int32(roff), cutlass.Int32(-1)))
- soff = cutlass.Int32(cutlass.select_(okp, cutlass.Int32(soff), cutlass.Int32(-1)))
- mw = cutlass.Int32(cutlass.select_(okp, mw, cutlass.Int32(0)))
- if slot < tmax * GPT:
- _sts1(s_roff + slot * 4, roff)
- _sts1(s_soff + slot * 4, soff)
- _sts1(s_gmask + slot * 4, mw)
- cute.arch.sync_threads()
- for i in cutlass.range_constexpr(-(-(tmax * rpr) // THREADS)):
- w = tidx + i * THREADS
- if w < tmax * rpr:
- kl = w // rpr
- t = w % rpr
- lo = cutlass.Int32(0)
- hi = cutlass.Int32(0)
- for gq in cutlass.range_constexpr(GPT):
- nb = (_lds1(s_gmask + (kl * GPT + gq) * 4) >> t) & 0x01010101
- nb = (nb | (nb >> 7) | (nb >> 14) | (nb >> 21)) & 0xF
- if const_expr(gq < 8):
- lo = lo | (nb << (4 * gq))
- else:
- hi = hi | (nb << (4 * (gq - 8)))
- _sts1(s_bits + w * 8, lo)
- _sts1(s_bits + w * 8 + 4, hi)
- cute.arch.sync_threads()
-
- @cute.jit
- def _pass_tiles(
- self,
- act,
- lead,
- kvh,
- r0,
- rows,
- sel_base,
- split,
- nsplit,
- tidx,
- warp,
- lane,
- s_raw,
- s_scan,
- s_tab,
- mBT,
- ):
- """Union + resolve of one pass; returns the CTA's local tile count."""
- s_bm = s_raw + self.u_bm
- s_ug = s_raw + self.u_ug
- s_bt = s_raw + self.u_bt
- n_ug = self._union(
- act, lead, r0, rows, sel_base, tidx, warp, lane, s_bm, s_scan, s_ug, s_bt, mBT
- )
- n_tiles = (n_ug + GPT - 1) // GPT
- n_loc = cutlass.Int32(
- cutlass.select_(
- split < n_tiles,
- (n_tiles - split + nsplit - 1) // nsplit,
- cutlass.Int32(0),
- )
- )
- self._resolve(lead, kvh, split, nsplit, n_loc, n_ug, tidx, s_bm, s_ug, s_tab, s_bt)
- return n_loc
-
- @cute.kernel
- def kernel(
- self,
- qk_mma: cute.TiledMma,
- qk_mma_acc: cute.TiledMma,
- pv_mma_acc: cute.TiledMma,
- q_layout: cute.ComposedLayout,
- k_layout: cute.ComposedLayout,
- v_layout: cute.ComposedLayout,
- p_layout: cute.ComposedLayout,
- sel_base: cutlass.Int64,
- q_base: cutlass.Int64,
- k_base: cutlass.Int64,
- v_base: cutlass.Int64,
- ks_base: cutlass.Int64,
- vs_base: cutlass.Int64,
- mGs: cute.Tensor,
- mBT: cute.Tensor,
- mReq: cute.Tensor,
- out_base: cutlass.Int64,
- part_base: cutlass.Int64,
- lse_base: cutlass.Int64,
- gate_base: cutlass.Int64,
- rows: cutlass.Int32,
- gate_stride: cutlass.Int32,
- n_whole: cutlass.Int32,
- tail_row0: cutlass.Int32,
- ):
- warp = cute.arch.make_warp_uniform(cute.arch.warp_idx())
- tidx, _, _ = cute.arch.thread_idx()
- bidx, _, _ = cute.arch.block_idx()
- lane = cute.arch.lane_idx()
- qk_mma.set(tcgen05.Field.ACCUMULATE, False)
- qk_mma_acc.set(tcgen05.Field.ACCUMULATE, True)
- pv_mma_acc.set(tcgen05.Field.ACCUMULATE, True)
- rpr: cutlass.Constexpr = self.T
- group: cutlass.Constexpr = self.group
- kvh_n: cutlass.Constexpr = self.kvh
- splits: cutlass.Constexpr = self.splits
- if const_expr(self.tail):
- # units (group, KV head) below n_whole: one CTA each; the rest: `splits` CTAs each
- is_whole = bidx < n_whole
- t_cta = bidx - n_whole
- unit = cutlass.Int32(cutlass.select_(is_whole, bidx, n_whole + t_cta // splits))
- split = cutlass.Int32(cutlass.select_(is_whole, cutlass.Int32(0), t_cta % splits))
- nsplit = cutlass.Int32(
- cutlass.select_(is_whole, cutlass.Int32(1), cutlass.Int32(splits))
- )
- else:
- unit = bidx // splits
- split = bidx % splits
- nsplit = cutlass.Int32(splits)
- gi = unit // kvh_n
- kvh = unit % kvh_n
- r0 = gi * rpr
-
- smem = SmemAllocator()
- storage = smem.allocate(self.shared_storage)
- q_full = storage.q_full.data_ptr()
- raw_full = storage.raw_full.data_ptr()
- raw_empty = storage.raw_empty.data_ptr()
- k_full = storage.k_full.data_ptr()
- k_empty = storage.k_empty.data_ptr()
- v_full = storage.v_full.data_ptr()
- v_empty = storage.v_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()
- s_scan = cutlass.Int32(storage.scan.data_ptr().toint())
- s_tab = cutlass.Int32(storage.tab.data_ptr().toint())
- s_raw = cutlass.Int32(storage.raw.data_ptr().toint())
- if tidx == 0:
- cute.arch.mbarrier_init(q_full, 32 * SOFT_WARPS)
- for i in cutlass.range_constexpr(RAW_SLOTS):
- cute.arch.mbarrier_init(raw_full + i, 32 * LOAD_WARPS)
- cute.arch.mbarrier_init(raw_empty + i, 32 * DQ_WARPS)
- for i in cutlass.range_constexpr(2):
- cute.arch.mbarrier_init(k_full + i, BN)
- cute.arch.mbarrier_init(k_empty + i, 1)
- cute.arch.mbarrier_init(v_full + i, BN)
- cute.arch.mbarrier_init(v_empty + i, 1)
- cute.arch.mbarrier_init(s_full + i, 1)
- cute.arch.mbarrier_init(p_full + i, 32 * SOFT_WARPS)
- cute.arch.mbarrier_init(o_full, 1)
- cute.arch.mbarrier_init(o_final, 1)
- cute.arch.mbarrier_init_fence()
- if warp == MMA_WARP:
- cute.arch.alloc_tmem(TMEM_COLS, storage.tmem_holding.ptr)
- cute.arch.relinquish_tmem_alloc_permit()
- cute.arch.sync_threads()
- tmem_ptr = cute.arch.retrieve_tmem_ptr(Float32, 16, storage.tmem_holding.ptr)
-
- qk_thr = qk_mma.get_slice(0)
- pv_thr = pv_mma_acc.get_slice(0)
- sK = storage.sK.get_tensor(k_layout.outer, swizzle=k_layout.inner)
- sVr = storage.sV.get_tensor(k_layout.outer, swizzle=k_layout.inner)
- sV = cute.make_tensor(cute.recast_ptr(sVr.iterator, v_layout.inner), v_layout.outer)
- s_k16 = cutlass.Int32(storage.sK.data_ptr().toint())
- s_v16 = cutlass.Int32(storage.sV.data_ptr().toint())
-
- # ---- softmax-side tensor-memory views (warps 0-3: one thread an MMA row) -------------
- st = tidx
- tStS = qk_thr.make_fragment_C(qk_thr.partition_shape_C((BM, BN)))
- tStS0 = cute.make_tensor(tmem_ptr + TMEM_S0, tStS.layout)
- tOtO_f = pv_thr.make_fragment_C(pv_thr.partition_shape_C((BM, HEAD_DIM)))
- tOtO = cute.make_tensor(tmem_ptr + TMEM_O, tOtO_f.layout)
- tScS = qk_thr.partition_C(cute.make_identity_tensor((BM, BN)))
- load_s = tcgen05.make_tmem_copy(
- cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), 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 = BN // 2
- tStP0 = cute.make_tensor(
- tStS0.iterator, cute.composition(tStS0.layout, cute.make_layout((BM, p_cols)))
- )
- tScP = cute.make_tensor(
- tScS.iterator, cute.composition(tScS.layout, cute.make_layout((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)
- tOcO = pv_thr.partition_C(cute.make_identity_tensor((BM, HEAD_DIM)))
- tOtO_c = cute.make_tensor(
- tOtO.iterator, cute.composition(tOtO.layout, cute.make_layout((BM, CHUNK_COLS)))
- )
- tOcO_c = cute.make_tensor(
- tOcO.iterator, cute.composition(tOcO.layout, cute.make_layout((BM, CHUNK_COLS)))
- )
- load_o = tcgen05.make_tmem_copy(
- cute.make_copy_atom(
- tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(CHUNK_COLS)), Float32
- ),
- tOtO_c,
- )
- store_o = tcgen05.make_tmem_copy(
- cute.make_copy_atom(
- tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(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)
- tview = (load_s, tLtS0, tLcS, store_p, tPtP0, tPcP, load_o, store_o, tLtO, tLcO, tStO)
-
- m = tidx
- t_row = m // group
- h = m % group
- live = (m < self.rows_m) & (r0 + t_row < rows)
- safe_row = cutlass.min(r0 + t_row, rows - 1)
- qs0, qs1 = self.qs
- q_row = q_base + (cutlass.Int64(safe_row) * qs0 + cutlass.Int64(kvh * group + h) * qs1) * 2
- q_s = s_raw + self.u_q + cutlass.min(m, self.rows_m - 1) * Q_LD
-
- if const_expr(self.pdl):
- cute.arch.griddepcontrol_launch_dependents()
- cute.arch.griddepcontrol_wait()
-
- if warp < SOFT_WARPS:
- # O = 0; this row's q copied into raw slots 1-2 now, into tensor memory (f16) after the
- # first pass's resolve. Dead rows stage zeros.
- rZ = cute.make_rmem_tensor(tLcO.shape, Float32)
- rZ.fill(0.0)
- for c in cutlass.range_constexpr(HEAD_DIM // CHUNK_COLS):
- cute.copy(
- store_o, rZ, cute.make_tensor(tStO.iterator + c * CHUNK_COLS, tStO.layout)
- )
- cute.arch.fence_view_async_tmem_store()
- if m < self.rows_m:
- qn = cutlass.Int32(cutlass.select_(live, cutlass.Int32(16), cutlass.Int32(0)))
- for i in cutlass.range_constexpr(HEAD_DIM // 8):
- _cp16z(q_s + 16 * i, q_row + cutlass.Int64(16 * i), qn)
- cute.arch.cp_async_commit_group()
-
- rq = cute.make_rmem_tensor(cute.make_layout((rpr,), stride=(1,)), cutlass.Int32)
- handled = cutlass.Int32(0)
- for t in cutlass.range_constexpr(rpr):
- row = r0 + t
- xr = cutlass.Int64(mReq[cutlass.min(row, rows - 1)])
- okr = (row < rows) & (xr >= 0) & (xr < self.num_requests)
- rq[t] = cutlass.Int32(cutlass.select_(okr, cutlass.Int32(xr), cutlass.Int32(-1)))
- handled = handled | cutlass.Int32(
- cutlass.select_(okr, cutlass.Int32(0), cutlass.Int32(1 << t))
- )
- full = (1 << rpr) - 1
- scale_log2 = Float32(self.scale_log2) * mGs[1]
- neg_inf = -Float32.inf
- row_max = neg_inf
- row_sum = Float32(0.0)
- q_pending = cutlass.Int32(1)
- g = cutlass.Int32(0)
- for _pass in cutlass.range(rpr):
- if handled != full:
- lead, act = self._pass_rows(rq, handled)
- n_loc = self._pass_tiles(
- act,
- lead,
- kvh,
- r0,
- rows,
- sel_base,
- split,
- nsplit,
- tidx,
- warp,
- lane,
- s_raw,
- s_scan,
- s_tab,
- mBT,
- )
- if warp >= DQ_WARP0:
- self.dequant_tiles(
- tidx - 32 * DQ_WARP0,
- g,
- n_loc,
- s_raw,
- s_k16,
- s_v16,
- raw_full,
- raw_empty,
- k_full,
- k_empty,
- v_full,
- v_empty,
- )
- elif warp >= LOAD_WARP0:
- self.load_tiles(
- tidx - 32 * LOAD_WARP0,
- g,
- n_loc,
- s_raw,
- s_tab,
- k_base,
- v_base,
- ks_base,
- vs_base,
- raw_full,
- raw_empty,
- q_full,
- )
- elif warp == MMA_WARP:
- self.mma_tiles(
- g,
- n_loc,
- qk_mma,
- pv_mma_acc,
- qk_thr,
- pv_thr,
- q_layout,
- p_layout,
- sK,
- sV,
- tmem_ptr,
- q_full,
- k_full,
- k_empty,
- v_full,
- v_empty,
- s_full,
- p_full,
- o_full,
- )
- else:
- if q_pending != 0:
- self.stage_q(q_s, live, tview, q_full)
- q_pending = cutlass.Int32(0)
- row_max, row_sum = self.softmax_tiles(
- g,
- n_loc,
- row_max,
- row_sum,
- live,
- t_row,
- warp,
- s_tab,
- scale_log2,
- tview,
- s_full,
- p_full,
- o_full,
- )
- g = g + n_loc
- cute.arch.sync_threads() # [P] end of the pass
- handled = handled | act
-
- if warp == MMA_WARP:
- with cute.arch.elect_one():
- tcgen05.commit(o_final)
- if warp < SOFT_WARPS:
- cute.arch.cp_async_wait_group(0)
- self.epilogue(
- row_max,
- row_sum,
- live,
- warp,
- m,
- r0,
- rows,
- kvh,
- split,
- nsplit,
- tail_row0,
- mGs,
- scale_log2,
- out_base,
- part_base,
- lse_base,
- gate_base,
- gate_stride,
- tview,
- o_final,
- tidx,
- s_k16,
- )
- # every role done with tensor memory (the softmax warps read O last)
- cute.arch.sync_threads()
- if warp == MMA_WARP:
- cute.arch.dealloc_tmem(tmem_ptr, TMEM_COLS)
-
- # ------------------------------------------------------------------ loaders (warps 5-6)
- @cute.jit
- def load_tiles(
- self,
- lt,
- g,
- n_loc,
- s_raw,
- s_tab,
- k_base,
- v_base,
- ks_base,
- vs_base,
- raw_full,
- raw_empty,
- q_full,
- ):
- """Raw tile copies: K and V rows eight threads a row (16-byte chunk lt % 8 of rows
- lt // 8 + 8 i: whole 128-byte lines a warp instruction); the K-scale row lt and chunk
- lt % 4 of the V-scale group lt // 4 (both at the group's scale offset + lt % 4). Raw slots
- 1-2 hold the q rows until Q is in tensor memory.
- """
- s_roff = s_tab + self.t_roff * 4
- s_soff = s_tab + self.t_soff * 4
- s_gmask = s_tab + self.t_gmask * 4
- tok16: cutlass.Constexpr = self.ks[2] // 16
- q8 = lt % 8
- for kl in cutlass.range(n_loc):
- gg = g + kl
- r = gg % RAW_SLOTS
- cute.arch.mbarrier_wait(raw_empty + r, _phase(gg, RAW_SLOTS) ^ 1)
- if gg >= 1:
- cute.arch.mbarrier_wait(q_full, 0)
- rb = s_raw + r * self.raw_buf
- rk = rb
- rv = rb + self.raw_k
- rks = rv + self.raw_v
- rvs = rks + self.raw_ks
- for i in cutlass.range_constexpr(BN // 8):
- n = lt // 8 + 8 * i
- slot = kl * GPT + n // 4
- roff = _lds1(s_roff + slot * 4)
- mw = _lds1(s_gmask + slot * 4)
- ok = (roff >= 0) & (((mw >> (8 * (n % 4))) & 0xFF) != 0)
- nb = cutlass.Int32(cutlass.select_(ok, cutlass.Int32(16), cutlass.Int32(0)))
- off = cutlass.Int64(cutlass.max(roff, 0) + (n % 4) * tok16 + q8) * 16
- _cp16z(rk + n * RAW_LD + 16 * q8, k_base + off, nb)
- _cp16z(rv + n * RAW_LD + 16 * q8, v_base + off, nb)
- slot = kl * GPT + lt // 4
- soff = _lds1(s_soff + slot * 4)
- mw = _lds1(s_gmask + slot * 4)
- ok = (soff >= 0) & (((mw >> (8 * (lt % 4))) & 0xFF) != 0)
- gok = (soff >= 0) & (mw != 0)
- nb = cutlass.Int32(cutlass.select_(ok, cutlass.Int32(16), cutlass.Int32(0)))
- ng = cutlass.Int32(cutlass.select_(gok, cutlass.Int32(16), cutlass.Int32(0)))
- soff64 = cutlass.Int64(cutlass.max(soff, 0) + lt % 4) * 16
- _cp16z(rks + lt * 16, ks_base + soff64, nb)
- _cp16z(rvs + (lt // 4) * 64 + 16 * (lt % 4), vs_base + soff64, ng)
- cute.arch.cp_async_mbarrier_arrive_noinc(raw_full + r)
-
- # ------------------------------------------------------------------ dequant (warps 7-10)
- @cute.jit
- def dequant_tiles(
- self,
- dt,
- g,
- n_loc,
- s_raw,
- s_k16,
- s_v16,
- raw_full,
- raw_empty,
- k_full,
- k_empty,
- v_full,
- v_empty,
- ):
- """Thread `dt` < 64 dequantizes K row `dt` of each tile, the others V row `dt - 64`: the raw
- row and its scales into registers (the raw slot is then free), then into the f16 SW128
- slot once the MMA has released it.
- """
- is_v = dt >= BN
- n = dt % BN
- for kl in cutlass.range(n_loc):
- gg = g + kl
- r = gg % RAW_SLOTS
- b = gg % 2
- cute.arch.mbarrier_wait(raw_full + r, _phase(gg, RAW_SLOTS))
- rb = s_raw + r * self.raw_buf
- row_off = (
- rb
- + cutlass.Int32(cutlass.select_(is_v, cutlass.Int32(self.raw_k), cutlass.Int32(0)))
- + n * RAW_LD
- )
- w = [_lds4(row_off + 16 * c) for c in range(8)]
- sc_off = rb + self.raw_k + self.raw_v
- sc_off = cutlass.Int32(
- cutlass.select_(is_v, sc_off + self.raw_ks + (n // 4) * 64, sc_off + n * 16)
- )
- sw = [_lds4(sc_off + 16 * i) for i in range(4)]
- cute.arch.mbarrier_arrive(raw_empty + r)
- if is_v:
- cute.arch.mbarrier_wait(v_empty + b, _phase(gg, 2) ^ 1)
- self._decode_row(w, sw, s_v16 + b * (BN * HEAD_DIM * 2), n, True)
- cute.arch.fence_view_async_shared()
- cute.arch.mbarrier_arrive(v_full + b)
- else:
- cute.arch.mbarrier_wait(k_empty + b, _phase(gg, 2) ^ 1)
- self._decode_row(w, sw, s_k16 + b * (BN * HEAD_DIM * 2), n, False)
- cute.arch.fence_view_async_shared()
- cute.arch.mbarrier_arrive(k_full + b)
-
- @cute.jit
- def _decode_row(self, w, sw, dst, n, is_v: cutlass.Constexpr):
- """Raw row words `w` (8 x 4) and scale words `sw` -> f16 SW128 row `n` at `dst`: 16-byte
- chunk ci of the row's 64-dim block ci // 8 at 8192 (ci // 8) + 128 n + 16 ((ci % 8) ^ (n % 8)).
- """
- sel_v = (n % 4) * 0x11
- for c in cutlass.range_constexpr(8):
- # blocks 2c, 2c + 1. K: bytes 2c, 2c + 1 of the token's 16 (word c // 2);
- # V: byte n % 4 of the group's words 2c, 2c + 1
- if const_expr(is_v):
- sc0 = _scale2s(sw[(2 * c) // 4][(2 * c) % 4], sel_v)
- sc1 = _scale2s(sw[(2 * c + 1) // 4][(2 * c + 1) % 4], sel_v)
- else:
- sc0 = _scale2s(sw[0][c // 2], (2 * (c % 2)) * 0x11)
- sc1 = _scale2s(sw[0][c // 2], (2 * (c % 2) + 1) * 0x11)
- for e in cutlass.range_constexpr(4):
- d4 = _dq_word(w[c][e], sc0 if e < 2 else sc1)
- ci = 4 * c + e
- _sts4(
- dst + 8192 * (ci // 8) + 128 * n + 16 * ((ci % 8) ^ (n % 8)),
- d4[0],
- d4[1],
- d4[2],
- d4[3],
- )
-
- # ------------------------------------------------------------------ MMA (warp 4)
- @cute.jit
- def mma_tiles(
- self,
- g,
- n_loc,
- qk_mma,
- pv_mma_acc,
- qk_thr,
- pv_thr,
- q_layout,
- p_layout,
- sK,
- sV,
- tmem_ptr,
- q_full,
- k_full,
- k_empty,
- v_full,
- v_empty,
- s_full,
- p_full,
- o_full,
- ):
- """QK(g), then QK(j) / PV(j - 1) pairs, then the last PV: S/P buffer j % 2 at TMEM column
- TMEM_S0 + 64 (j % 2); QK(j) overwrites it only after PV(j - 2) was issued (tcgen05 MMAs of
- one thread execute in order).
- """
- tQ = cute.make_tensor(tmem_ptr + 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 + TMEM_Q, dtype=F16), tSrQ.layout)
- tSrK = qk_thr.make_fragment_B(sK)
- tOrV = pv_thr.make_fragment_B(sV)
- tP = cute.make_tensor(tmem_ptr + TMEM_S0, p_layout.outer)
- tOrP = pv_thr.make_fragment_A(tP)[None, None, None, 0]
- tOrP = cute.make_tensor(cute.recast_ptr(tmem_ptr + TMEM_S0, dtype=F16), tOrP.layout)
- tbase = cutlass.Int32(tmem_ptr.toint())
- qk_op = qk_mma.op
- pv_op = pv_mma_acc.op
- if n_loc > 0:
- cute.arch.mbarrier_wait(q_full, 0)
- cute.arch.mbarrier_wait(k_full + g % 2, _phase(g, 2))
- tcgen05_fence_after()
- self.issue_qk(qk_op, tbase, tSrQ, tSrK, sK, g % 2)
- with cute.arch.elect_one():
- tcgen05.commit(s_full + g % 2)
- tcgen05.commit(k_empty + g % 2)
- for j in cutlass.range(1, n_loc + 1, 1, unroll=1):
- gj = g + j
- if j < n_loc:
- cute.arch.mbarrier_wait(k_full + gj % 2, _phase(gj, 2))
- tcgen05_fence_after()
- self.issue_qk(qk_op, tbase, tSrQ, tSrK, sK, gj % 2)
- with cute.arch.elect_one():
- tcgen05.commit(s_full + gj % 2)
- tcgen05.commit(k_empty + gj % 2)
- gp = gj - 1
- cute.arch.mbarrier_wait(v_full + gp % 2, _phase(gp, 2))
- cute.arch.mbarrier_wait(p_full + gp % 2, _phase(gp, 2))
- tcgen05_fence_after()
- gemm_ts(
- pv_op,
- tbase + TMEM_O,
- tbase + TMEM_S0 + (gp % 2) * BN,
- tOrP,
- tOrV[None, None, None, gp % 2],
- sV[None, None, None, gp % 2],
- False,
- )
- with cute.arch.elect_one():
- tcgen05.commit(v_empty + gp % 2)
- tcgen05.commit(o_full)
-
- @cute.jit
- def issue_qk(self, qk_op, tbase, tSrQ, tSrK, sK, slot):
- """S[slot] = Q K(slot)^T: Q from tensor memory, K from SW128 slot `slot`."""
- gemm_ts(
- qk_op,
- tbase + TMEM_S0 + slot * BN,
- tbase + TMEM_Q,
- tSrQ,
- tSrK[None, None, None, slot],
- sK[None, None, None, slot],
- True,
- )
-
- # ------------------------------------------------------------------ softmax + epilogue (warps 0-3)
- @cute.jit
- def stage_q(self, q_s, live, tview, q_full):
- """This MMA row's q (bf16 in raw slots 1-2) -> f16 into tensor memory (the QK A operand)."""
- load_s, tLtS0, tLcS, store_p, tPtP0, tPcP, load_o, store_o, tLtO, tLcO, tStO = tview
- p_cols: cutlass.Constexpr = BN // 2
- cute.arch.cp_async_wait_group(0)
- for c in cutlass.range_constexpr(HEAD_DIM // (2 * p_cols)):
- rQ = cute.make_rmem_tensor(tPcP.shape, Float32)
- for i in cutlass.range_constexpr(p_cols // 4):
- w4 = _lds4(q_s + 16 * (c * (p_cols // 4) + i))
- for e in cutlass.range_constexpr(4):
- rQ[4 * i + e] = _bits_f32(
- cutlass.Int32(cutlass.select_(live, _bf2_to_f2(w4[e]), cutlass.Int32(0)))
- )
- cute.copy(
- store_p,
- rQ,
- cute.make_tensor(tPtP0.iterator + (TMEM_Q - TMEM_S0 + c * p_cols), tPtP0.layout),
- )
- cute.arch.fence_view_async_tmem_store()
- tcgen05_fence_before()
- cute.arch.mbarrier_arrive(q_full)
-
- @cute.jit
- def softmax_tiles(
- self,
- g,
- n_loc,
- row_max,
- row_sum,
- live,
- t_row,
- warp,
- s_tab,
- scale_log2,
- tview,
- s_full,
- p_full,
- o_full,
- ):
- """Per tile: S from tensor memory, the max over the row's own entries, the reference max
- moved (and O rescaled) only past RESCALE_TH, P = 2^(s - max) as f16 over the S.
- """
- load_s, tLtS0, tLcS, store_p, tPtP0, tPcP, load_o, store_o, tLtO, tLcO, tStO = tview
- rpr: cutlass.Constexpr = self.T
- s_bits = s_tab + self.t_bits * 4
- tau = Float32(RESCALE_TH)
- neg_inf = -Float32.inf
- live_warp = warp < self.live_warps
- for kl in cutlass.range(n_loc, unroll=1):
- gg = g + kl
- buf = gg % 2
- cute.arch.mbarrier_wait(s_full + buf, _phase(gg, 2))
- tcgen05_fence_after()
- if live_warp:
- lo = cutlass.Int32(0)
- hi = cutlass.Int32(0)
- if live:
- lo = _lds1(s_bits + (kl * rpr + t_row) * 8)
- hi = _lds1(s_bits + (kl * rpr + t_row) * 8 + 4)
- rS = cute.make_rmem_tensor(tLcS.shape, Float32)
- tLtS_b = cute.make_tensor(tLtS0.iterator + buf * BN, tLtS0.layout)
- for c in cutlass.range_constexpr(cute.size(rS, mode=[1])):
- cute.copy(load_s, tLtS_b[None, c, 0, 0], rS[None, c, 0, 0])
- cute.arch.fence_view_async_tmem_load()
- tile_max = neg_inf
- for i in cutlass.range_constexpr(BN):
- word = lo if i < 32 else hi
- x = rS[i]
- x = x if (word & (1 << (i % 32))) != 0 else neg_inf
- rS[i] = x
- tile_max = cute.arch.fmax(tile_max, x)
- new_max = cute.arch.fmax(row_max, tile_max)
- use_max = row_max
- factor = Float32(1.0)
- if row_max == neg_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(..gg - 1): wait for it, then rescale this warp's rows
- cute.arch.mbarrier_wait(o_full, (gg - 1) % 2)
- tcgen05_fence_after()
- for c in cutlass.range_constexpr(HEAD_DIM // CHUNK_COLS):
- rO = cute.make_rmem_tensor(tLcO.shape, Float32)
- cute.copy(
- load_o,
- cute.make_tensor(tLtO.iterator + c * 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 * CHUNK_COLS, tStO.layout),
- )
- cute.arch.fence_view_async_tmem_store()
- row_max = use_max
- safe = use_max
- if use_max == neg_inf:
- safe = Float32(0.0)
- neg = -safe * scale_log2
- rP = cute.make_rmem_tensor(tPcP.shape, Float32)
- rP16 = cute.make_tensor(cute.recast_ptr(rP.iterator, dtype=F16), rS.layout)
- acc0 = (Float32(0.0), Float32(0.0))
- acc1 = (Float32(0.0), Float32(0.0))
- for i in cutlass.range_constexpr(0, BN, 4):
- 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):
- rS[i + u] = cute.math.exp2(e[u], fastmath=True)
- 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(F16))
- cute.copy(store_p, rP, cute.make_tensor(tPtP0.iterator + buf * BN, tPtP0.layout))
- cute.arch.fence_view_async_tmem_store()
- tcgen05_fence_before()
- cute.arch.mbarrier_arrive(p_full + buf)
- return row_max, row_sum
-
- @cute.jit
- def epilogue(
- self,
- row_max,
- row_sum,
- live,
- warp,
- m,
- r0,
- rows,
- kvh,
- split,
- nsplit,
- tail_row0,
- mGs,
- scale_log2,
- out_base,
- part_base,
- lse_base,
- gate_base,
- gate_stride,
- tview,
- o_final,
- st,
- s_stage,
- ):
- """O / l (times the V global scale): a whole unit's rows as bf16, a split's as an fp32
- partial and its log2-sum-exp for the merge (rows from `tail_row0` on).
-
- Gated whole units: the four softmax warps read the gate (coalesced, 16 bytes a chunk of
- eight values) before the last P V completes; the live warps stage bf16(O / l) in the
- free K slots; then all four apply the gate and store coalesced rows."""
- load_s, tLtS0, tLcS, store_p, tPtP0, tPcP, load_o, store_o, tLtO, tLcO, tStO = tview
- group: cutlass.Constexpr = self.group
- heads: cutlass.Constexpr = self.heads
- splits: cutlass.Constexpr = self.splits
- tailm: cutlass.Constexpr = self.tail
- has_whole: cutlass.Constexpr = tailm or splits == 1
- has_split: cutlass.Constexpr = tailm or splits > 1
- gated_whole: cutlass.Constexpr = has_whole and self.gated
- rows_m: cutlass.Constexpr = self.rows_m
- cpr: cutlass.Constexpr = HEAD_DIM // 8 # 16-byte chunks a row
- n_gc: cutlass.Constexpr = -(-(rows_m * cpr) // (32 * SOFT_WARPS))
- n_chunks: cutlass.Constexpr = HEAD_DIM // CHUNK_COLS
- neg_inf = -Float32.inf
- row = r0 + m // group
- h = m % group
- head = kvh * group + h
- prow = cutlass.Int64(row - tail_row0)
- prows = cutlass.Int64(rows - tail_row0)
- whole = nsplit == 1
- gregs = []
- if const_expr(gated_whole):
- for j in cutlass.range_constexpr(n_gc):
- k = st + 32 * SOFT_WARPS * j
- gm = cutlass.min(k // cpr, rows_m - 1)
- grow = cutlass.min(r0 + gm // group, rows - 1)
- ghead = kvh * group + gm % group
- gregs.append(
- _ldg4(
- gate_base
- + (cutlass.Int64(grow) * gate_stride + ghead * HEAD_DIM + (k % cpr) * 8) * 2
- )
- )
- cute.arch.mbarrier_wait(o_final, 0)
- tcgen05_fence_after()
- if warp < self.live_warps:
- gv = mGs[2]
- inv = Float32(0.0)
- if row_sum > Float32(0.0):
- inv = gv / row_sum
- 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 * CHUNK_COLS, tLtO.layout), rO)
- for v8 in cutlass.range_constexpr(CHUNK_COLS // 8):
- vals = [rO[8 * v8 + i] * inv for i in range(8)]
- d0 = c * CHUNK_COLS + 8 * v8
- oaddr = out_base + ((cutlass.Int64(row) * heads + head) * HEAD_DIM + d0) * 2
- paddr = (
- part_base
- + (((cutlass.Int64(split) * prows + prow) * heads + head) * HEAD_DIM + d0)
- * 4
- )
- packed = [_bf2_pack(vals[2 * e], vals[2 * e + 1]) for e in range(4)]
- saddr = s_stage + cutlass.min(m, rows_m - 1) * Q_LD + d0 * 2
- if const_expr(has_whole):
- if whole:
- if const_expr(gated_whole):
- if m < rows_m:
- _sts4(saddr, packed[0], packed[1], packed[2], packed[3])
- else:
- if live:
- stg4(oaddr, packed[0], packed[1], packed[2], packed[3])
- if const_expr(has_split):
- if live & (nsplit > 1):
- stg4f(paddr, vals[0], vals[1], vals[2], vals[3])
- stg4f(paddr + 16, vals[4], vals[5], vals[6], vals[7])
- if const_expr(has_split):
- if live & (nsplit > 1):
- lse = neg_inf
- if row_sum > Float32(0.0):
- lse = row_max * scale_log2 + cute.math.log2(row_sum)
- _stg1(
- lse_base + ((cutlass.Int64(split) * prows + prow) * heads + head) * 4,
- _f32_bits(lse),
- )
- if const_expr(gated_whole):
- if whole:
- cute.arch.barrier(barrier_id=1, number_of_threads=32 * SOFT_WARPS)
- for j in cutlass.range_constexpr(n_gc):
- k = st + 32 * SOFT_WARPS * j
- gm = cutlass.min(k // cpr, rows_m - 1)
- grow = r0 + gm // group
- ghead = kvh * group + gm % group
- a4 = _lds4(s_stage + gm * Q_LD + (k % cpr) * 16)
- g4 = gregs[j]
- res = [
- _bf2_pack(
- _bf2_lo(a4[e]) * _sigmoid_fast(_bf2_lo(g4[e])),
- _bf2_hi(a4[e]) * _sigmoid_fast(_bf2_hi(g4[e])),
- )
- for e in range(4)
- ]
- if (k < rows_m * cpr) & (grow < rows):
- stg4(
- out_base
- + ((cutlass.Int64(grow) * heads + ghead) * HEAD_DIM + (k % cpr) * 8)
- * 2,
- res[0],
- res[1],
- res[2],
- res[3],
- )
-
-
-#: Threads of a merge CTA (16 bytes, four values, a thread at a time).
-MERGE_THREADS = 256
-#: Merge CTAs a row, each a contiguous share of the row's heads x head_dim values (measured on
-#: GB300: 8 is faster than 1-4 at 64-128 rows, 16 no faster).
-MERGE_PARTS = 8
-
-
-class _RowMergeKernel:
- """The split partials of a row range merged as the decode op's merge kernel does
- (log2-sum-exp weights, fp32), the gate applied to the bf16-rounded result: `parts` CTAs a
- row, every load of a thread issued before the first use."""
-
- def __init__(self, *, heads, splits, gated, pdl, parts):
- self.heads, self.splits, self.gated, self.pdl = heads, splits, gated, pdl
- self.parts = parts
-
- @cute.jit
- def __call__(
- self,
- gate_t: cute.Tensor,
- part_ptr: cute.Pointer,
- lse_ptr: cute.Pointer,
- out_ptr: cute.Pointer,
- rows: cutlass.Int32,
- gate_stride: cutlass.Int32,
- stream,
- ):
- self.kernel(
- part_ptr.toint(),
- lse_ptr.toint(),
- out_ptr.toint(),
- gate_t.iterator.toint(),
- rows,
- gate_stride,
- ).launch(
- grid=[rows * self.parts, 1, 1],
- block=[MERGE_THREADS, 1, 1],
- stream=stream,
- use_pdl=self.pdl,
- )
-
- @cute.kernel
- def kernel(
- self,
- part_base: cutlass.Int64,
- lse_base: cutlass.Int64,
- out_base: cutlass.Int64,
- gate_base: cutlass.Int64,
- rows: cutlass.Int32,
- gate_stride: cutlass.Int32,
- ):
- tidx, _, _ = cute.arch.thread_idx()
- bidx, _, _ = cute.arch.block_idx()
- heads: cutlass.Constexpr = self.heads
- sp: cutlass.Constexpr = self.splits
- parts: cutlass.Constexpr = self.parts
- n_all: cutlass.Constexpr = heads * HEAD_DIM // 4
- n_c: cutlass.Constexpr = -(-n_all // parts)
- per: cutlass.Constexpr = -(-n_c // MERGE_THREADS)
- neg_inf = -Float32.inf
- row = bidx // parts
- c0 = (bidx % parts) * n_c
- r64 = cutlass.Int64(row)
- cs = [
- cutlass.min(c0 + cutlass.min(tidx + MERGE_THREADS * j, n_c - 1), n_all - 1)
- for j in range(per)
- ]
- gws = []
- if const_expr(self.gated):
- # the gate is an input of the attention call, ready before the split kernel ran
- gws = [ldg2(gate_base + (r64 * gate_stride + 4 * c) * 2) for c in cs]
- if const_expr(self.pdl):
- cute.arch.griddepcontrol_launch_dependents()
- cute.arch.griddepcontrol_wait()
- lss = []
- pss = []
- for j in cutlass.range_constexpr(per):
- head = cs[j] // (HEAD_DIM // 4)
- lss.append(
- [
- _bits_f32(
- _ldg1(lse_base + ((cutlass.Int64(s_) * rows + r64) * heads + head) * 4)
- )
- for s_ in range(sp)
- ]
- )
- pss.append(
- [
- _ldg4(
- part_base
- + ((cutlass.Int64(s_) * rows + r64) * (heads * HEAD_DIM) + 4 * cs[j]) * 4
- )
- for s_ in range(sp)
- ]
- )
- for j in cutlass.range_constexpr(per):
- ls = lss[j]
- mx = _tree_fmax(ls)
- ws = [
- Float32(cutlass.select_(mx > neg_inf, cute.math.exp2(lv - mx), Float32(0.0)))
- for lv in ls
- ]
- den = _tree_sum(ws)
- res = []
- for e in cutlass.range_constexpr(4):
- acc = _tree_sum([ws[s_] * _bits_f32(pss[j][s_][e]) for s_ in range(sp)])
- merged = Float32(0.0)
- if den > Float32(0.0):
- merged = acc / den
- if const_expr(self.gated):
- if const_expr(e % 2 == 0):
- gx = _bf2_lo(gws[j][e // 2])
- else:
- gx = _bf2_hi(gws[j][e // 2])
- merged = _round_bf16(merged) * _sigmoid_fast(gx)
- res.append(merged)
- if (tidx + MERGE_THREADS * j < n_c) & (c0 + tidx + MERGE_THREADS * j < n_all):
- stg2(
- out_base + (r64 * (heads * HEAD_DIM) + 4 * cs[j]) * 2,
- _bf2_pack(res[0], res[1]),
- _bf2_pack(res[2], res[3]),
- )
-
-
-# ---------------------------------------------------------------------------------------
-# launch
-# ---------------------------------------------------------------------------------------
-
-_COMPILED: dict = {}
-#: Geometries whose every reachable launch plan is compiled (`_compile_plans`).
-_PLANNED: set = set()
-
-
-@functools.lru_cache(maxsize=None)
-def _num_sms(device_index: int) -> int:
- return torch.cuda.get_device_properties(device_index).multi_processor_count
-
-
-def plan_splits(n_groups: int, kv_heads: int, width: int, num_sms: int) -> int:
- """Splits so the grid fills the GPU once (one CTA an SM): measured on GB300, a second
- partial wave costs more than the extra splits save.
- """
- est_tiles = -(-int(width * 1.25) // BN)
- return max(1, min(MAX_PLAN_SPLITS, est_tiles, num_sms // (n_groups * kv_heads)))
-
-
-def plan_launch(n_groups: int, kv_heads: int, width: int, num_sms: int):
- """(splits, n_whole) of a launch: n_whole None = every unit (group, KV head) takes `splits`
- CTAs (one wave: `plan_splits`); else units below n_whole take one CTA and the units of the
- last, partial wave (whole groups) `splits` CTAs each, merged afterwards."""
- units = n_groups * kv_heads
- if units <= num_sms:
- return plan_splits(n_groups, kv_heads, width, num_sms), None
- waves = -(-units // num_sms)
- tail = units - (waves - 1) * num_sms
- tail = -(-tail // kv_heads) * kv_heads
- tail_splits = min(MAX_TAIL_SPLITS, num_sms // tail)
- if tail_splits < 2:
- return 1, None
- return tail_splits, units - tail
-
-
-def reachable_plans(max_groups: int, kv_heads: int, width: int, num_sms: int) -> list:
- """Every `(splits, tail)` pair `plan_launch` returns for `MIN_GROUPS..max_groups` groups,
- `tail` whether the last wave is split: the kernel variants a geometry can need."""
- plans = set()
- for n_groups in range(MIN_GROUPS, max(MIN_GROUPS, max_groups) + 1):
- splits, n_whole = plan_launch(n_groups, kv_heads, width, num_sms)
- plans.add((max(1, min(MAX_SPLITS, int(splits))), n_whole is not None))
- return sorted(plans)
-
-
-def _plan(q, k_cache, block_table, selected_tokens, tokens_per_block, rows_per_request):
- """(splits, layout) of the grouped kernel for this call, or None where the per-row kernel
- takes it: fewer than `MIN_GROUPS` groups, one row a request, not SM103, or no shared-memory
- plan (more than six rows a request, or a block table wider than the union scratch holds:
- 612 pages of 32 tokens at four rows a request).
- """
- rows, num_q_heads, head_dim = q.shape
- rpr = int(rows_per_request)
- if rpr < 2 or -(-rows // rpr) < MIN_GROUPS or head_dim != HEAD_DIM:
- return None
- if torch.cuda.get_device_capability(q.device) != (10, 3):
- return None
- kv_heads = k_cache.shape[1]
- width = selected_tokens.shape[1]
- splits = plan_splits(-(-rows // rpr), kv_heads, width, _num_sms(q.device.index or 0))
- layout = _layout(
- rpr, num_q_heads // kv_heads, width, block_table.shape[1], tokens_per_block, splits
- )
- return None if layout is None else (splits, layout)
-
-
-def takes_grouped(q, k_cache, block_table, selected_tokens, tokens_per_block, rows_per_request):
- """Whether `nvfp4_sparse_gqa_verify` runs this call on the grouped kernel (else on the
- per-row kernel). A host-side decision on shapes only: a per-capture constant.
- """
- plan = _plan(q, k_cache, block_table, selected_tokens, tokens_per_block, rows_per_request)
- return plan is not None
-
-
-def _compile_verify(key: tuple, kern_kw: dict, compile_args: tuple):
- """The grouped kernel of `key` (its `splits` and `tail` at positions 5 and 6), compiled
- once."""
- fn = _COMPILED.get(key)
- if fn is None:
- kern = _VerifyKernel(splits=key[5], tail=key[6], **kern_kw)
- fn = cute.compile(kern, *compile_args, options="--opt-level 2 --enable-tvm-ffi")
- _COMPILED[key] = fn
- return fn
-
-
-def _compile_merge(heads: int, splits: int, gated: bool, pdl: bool, merge_args: tuple):
- """The row merge of `splits` partials, compiled once."""
- mkey = ("rowmerge", heads, splits, gated, pdl, MERGE_PARTS)
- mfn = _COMPILED.get(mkey)
- if mfn is None:
- mfn = cute.compile(
- _RowMergeKernel(heads=heads, splits=splits, gated=gated, pdl=pdl, parts=MERGE_PARTS),
- *merge_args,
- options="--opt-level 2 --enable-tvm-ffi",
- )
- _COMPILED[mkey] = mfn
- return mfn
-
-
-def nvfp4_sparse_gqa_verify_grouped(
- 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,
- rows_per_request: int,
- output_gate: Optional[torch.Tensor] = None,
- pdl: bool = True,
- splits: Optional[int] = None,
-) -> torch.Tensor:
- """The grouped kernel (and the merge with more than one split), whatever the row count;
- `splits` overrides the plan. `nvfp4_sparse_gqa_verify` is the routed entry.
- """
- rows, num_q_heads, head_dim = q.shape
- num_pages, num_kv_heads, tpb, packed = k_cache.shape
- rpr = int(rows_per_request)
- assert head_dim == HEAD_DIM and packed == HEAD_DIM // 2 and tpb == tokens_per_block
- assert q.dtype == torch.bfloat16 and q.stride(2) == 1 and q.stride(1) == HEAD_DIM
- assert (q.stride(0) * 2) % 16 == 0 and q.data_ptr() % 16 == 0, "16-byte aligned q rows"
- assert k_cache.dtype == torch.uint8 and k_cache.stride() == v_cache.stride()
- assert k_cache.stride(3) == 1 and v_cache.shape == k_cache.shape
- assert k_scale_cache.stride() == v_scale_cache.stride()
- assert k_scale_cache.stride(2) == HEAD_DIM // 16
- assert block_table.dtype == torch.int32 and block_table.stride(1) == 1
- assert block_table.stride(0) == block_table.shape[1], "row-contiguous block table"
- assert selected_tokens.dtype == torch.int32 and selected_tokens.is_contiguous()
- assert request_indices.is_contiguous() and request_indices.dtype in (torch.int32, torch.int64)
- group = num_q_heads // num_kv_heads
- width = selected_tokens.shape[1]
- out = torch.empty((rows, num_q_heads, head_dim), dtype=torch.bfloat16, device=q.device)
- if rows == 0:
- return out
- n_groups = -(-rows // rpr)
- units = n_groups * num_kv_heads
- n_whole = None
- splits_given = splits
- if splits is None:
- splits, n_whole = plan_launch(n_groups, num_kv_heads, width, _num_sms(q.device.index or 0))
- n_splits = max(1, min(MAX_SPLITS, int(splits)))
- tail = n_whole is not None
- assert (
- _layout(rpr, group, width, block_table.shape[1], tpb, 1 if tail else n_splits) is not None
- ), "no shared-memory plan for this geometry: route through nvfp4_sparse_gqa_verify"
- use_pdl = bool(pdl and get_env_enable_pdl())
- gated = output_gate is not None
- if gated:
- assert output_gate.shape == (rows, num_q_heads * head_dim) and output_gate.stride(1) == 1
- assert output_gate.dtype == torch.bfloat16
- assert (output_gate.stride(0) * 2) % 16 == 0 and output_gate.data_ptr() % 16 == 0
- req64 = request_indices.dtype == torch.int64
- if tail:
- tail_row0 = (n_whole // num_kv_heads) * rpr
- n_ctas = n_whole + (units - n_whole) * n_splits
- elif n_splits > 1:
- n_whole, tail_row0, n_ctas = 0, 0, units * n_splits
- else:
- n_whole, tail_row0, n_ctas = units, rows, units
- merged = rows - tail_row0
- key = (
- rpr,
- group,
- num_kv_heads,
- width,
- tpb,
- n_splits,
- tail,
- float(softmax_scale),
- num_pages,
- tuple(block_table.shape),
- tuple(k_cache.stride()[:3]),
- tuple(k_scale_cache.stride()[:2]),
- tuple(q.stride()[:2]),
- req64,
- gated,
- use_pdl,
- )
- if merged > 0:
- part = torch.empty(
- (n_splits, merged, num_q_heads, head_dim), dtype=torch.float32, device=q.device
- )
- lse = torch.empty((n_splits, merged, num_q_heads), dtype=torch.float32, device=q.device)
- else:
- part, lse = out, out
- gate_t = output_gate if gated else q
- gate_stride = output_gate.stride(0) if gated else 0
- compile_args = (
- # read with scalar loads: a row slice of the selection needs no copy
- from_dlpack(selected_tokens, assumed_align=4).mark_layout_dynamic(leading_dim=1),
- _ptr(q, cutlass.BFloat16),
- _ptr(k_cache, cutlass.Uint8),
- _ptr(v_cache, cutlass.Uint8),
- _ptr(k_scale_cache, cutlass.Uint8),
- _ptr(v_scale_cache, cutlass.Uint8),
- _ptr(kv_scale_quant_orig, cutlass.Float32, 4),
- _ptr(block_table, cutlass.Int32, 4),
- _ptr(request_indices, cutlass.Int64 if req64 else cutlass.Int32, 8 if req64 else 4),
- _ptr(out, cutlass.BFloat16),
- _ptr(part, cutlass.Float32),
- _ptr(lse, cutlass.Float32),
- _ptr(gate_t, cutlass.BFloat16),
- cutlass.Int32(rows),
- cutlass.Int32(gate_stride),
- cutlass.Int32(n_whole),
- cutlass.Int32(tail_row0),
- cutlass.Int32(n_ctas),
- cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
- )
- merge_args = (
- from_dlpack(gate_t, assumed_align=16).mark_layout_dynamic(leading_dim=gate_t.dim() - 1),
- _ptr(part, cutlass.Float32),
- _ptr(lse, cutlass.Float32, 4),
- _ptr(out, cutlass.BFloat16, 8),
- cutlass.Int32(merged),
- cutlass.Int32(gate_stride),
- cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
- )
- kern_kw = dict(
- rpr=rpr,
- group=group,
- kv_heads=num_kv_heads,
- width=width,
- tpb=tpb,
- scale=float(softmax_scale),
- num_pages=num_pages,
- bt_width=block_table.shape[1],
- num_requests=block_table.shape[0],
- k_strides=tuple(k_cache.stride()[:3]),
- ks_strides=tuple(k_scale_cache.stride()[:2]),
- q_strides=tuple(q.stride()[:2]),
- req64=req64,
- gated=gated,
- pdl=use_pdl,
- )
- # A launch plan depends on the call's row count: an eager step (a mixed step's
- # generation rows) meeting a new one would compile mid-serving, for seconds. The first
- # call of a geometry compiles every plan its row counts can reach instead.
- if splits_given is None and key[:5] + key[7:] not in _PLANNED:
- _PLANNED.add(key[:5] + key[7:])
- # A group is one request's rows: at most a block-table row's worth of them.
- max_groups = block_table.shape[0]
- for plan_splits_, plan_tail in reachable_plans(
- max_groups, num_kv_heads, width, _num_sms(q.device.index or 0)
- ):
- if (
- _layout(
- rpr, group, width, block_table.shape[1], tpb, 1 if plan_tail else plan_splits_
- )
- is None
- ):
- continue
- _compile_verify(key[:5] + (plan_splits_, plan_tail) + key[7:], kern_kw, compile_args)
- if plan_tail or plan_splits_ > 1:
- _compile_merge(num_q_heads, plan_splits_, gated, use_pdl, merge_args)
- fn = _compile_verify(key, kern_kw, compile_args)
- fn(
- selected_tokens,
- q.data_ptr(),
- k_cache.data_ptr(),
- v_cache.data_ptr(),
- k_scale_cache.data_ptr(),
- v_scale_cache.data_ptr(),
- kv_scale_quant_orig.data_ptr(),
- block_table.data_ptr(),
- request_indices.data_ptr(),
- out.data_ptr(),
- part.data_ptr(),
- lse.data_ptr(),
- gate_t.data_ptr(),
- rows,
- gate_stride,
- n_whole,
- tail_row0,
- n_ctas,
- )
- if merged <= 0:
- return out
- mfn = _compile_merge(num_q_heads, n_splits, gated, use_pdl, merge_args)
- mfn(
- gate_t[tail_row0:],
- part.data_ptr(),
- lse.data_ptr(),
- out[tail_row0:].data_ptr(),
- merged,
- gate_stride,
- )
- return out
-
-
-def nvfp4_sparse_gqa_verify(
- 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,
- rows_per_request: int,
- output_gate: Optional[torch.Tensor] = None,
- pdl: bool = True,
-) -> torch.Tensor:
- """`nvfp4_sparse_gqa_decode` for rows that come `rows_per_request` consecutive rows a
- request: the grouped kernel where `takes_grouped`, the per-row kernel otherwise. Returns
- `[rows, Hq, D]` bf16 (contiguous); with `output_gate` (`[rows, Hq * D]` bf16, unit column
- stride) `bf16(attn) * sigmoid(gate)` rounded to bf16.
- """
- plan = _plan(q, k_cache, block_table, selected_tokens, tokens_per_block, rows_per_request)
- if plan is None:
- return nvfp4_sparse_gqa_decode(
- q,
- k_cache,
- v_cache,
- k_scale_cache,
- v_scale_cache,
- kv_scale_quant_orig,
- block_table,
- selected_tokens,
- request_indices,
- tokens_per_block,
- softmax_scale,
- output_gate=output_gate,
- pdl=pdl,
- )
- return nvfp4_sparse_gqa_verify_grouped(
- q,
- k_cache,
- v_cache,
- k_scale_cache,
- v_scale_cache,
- kv_scale_quant_orig,
- block_table,
- selected_tokens,
- request_indices,
- tokens_per_block,
- softmax_scale,
- rows_per_request,
- output_gate=output_gate,
- pdl=pdl,
- splits=None,
- )
-
-
-def _verify(segments, overlap=0.05, **overrides) -> dict:
- """A cell spec: `(first position, rows)` per request in row order, four rows a request
- unless `rows_per_request` overrides it, this checkpoint's geometry (24 q heads over 2 KV
- heads of 256, 32 tokens a page, 2051 selected tokens). `overlap` is the spread of a
- request's rows' block scores around the request's (each row then takes its top 512 blocks
- in ascending order, as `qsa_topk_expand_decode` writes them); None selects every row
- independently at random, as the decode op's cells do.
- """
- return dict(segments=segments, rows_per_request=4, overlap=overlap, **overrides)
-
-
-class QsaPagedSparseGqaNvfp4Verify(OpWrapper):
- """`qsa_paged_sparse_gqa_nvfp4_decode` for rows that come `rows_per_request` consecutive
- rows a request (an MTP verify step: T = draft length + 1 rows a generation request):
- sparse GQA over the explicitly selected tokens of an NVFP4 paged pool, `[rows, Hq, D]`
- bf16, each selected K/V token loaded and dequantized once for a group of T rows. With
- `output_gate` (`[rows, Hq * D]` bf16, unit column stride) the result is
- `bf16(attn) * sigmoid(gate)` rounded once more to bf16.
-
- Fusion boundary: the block-table gather, the union of the group's selections, the NVFP4
- dequantize, scale, online softmax, P.V, the split-K merge and (optionally) the output gate.
- Selection, the K/V append and q's projection, norm and RoPE are not inside.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="MTP3 verify, 8 requests x 4 rows: the fewest the grouped kernel takes, "
- "9 splits, gate fused into the merge",
- spec=_verify(tuple((4000 + 37 * i, 4) for i in range(8)), gated=True),
- ),
- Cell(
- why="32 rows whose selections share little (independent random rows, union "
- "~2.7x a row's)",
- spec=_verify(tuple((3000 + 211 * i, 4) for i in range(8)), overlap=None),
- ),
- Cell(
- why="MTP3 verify at 128 rows (32 requests x 4, the mid-concurrency step): two "
- "splits a unit, merged by the row merge with the gate",
- spec=_verify(tuple((2100 + (53 * i) % 1900, 4) for i in range(32)), gated=True),
- ),
- Cell(
- why="MTP3 verify at 1024 rows (256 requests x 4): one CTA a unit with the gate fused "
- "into its store, the last wave's units split two ways and merged",
- spec=_verify(tuple((2100 + (37 * i) % 1900, 4) for i in range(256)), gated=True),
- ),
- Cell(
- why="groups spanning requests (3, 4, 4, 1, 4 x 5, 2 rows: misaligned and partial "
- "groups), early positions, int64 request indices, non-unit global scales, trunk "
- "layout",
- spec=_verify(
- (
- (0, 3),
- (5, 4),
- (37, 4),
- (4000, 1),
- (4100, 4),
- (4200, 4),
- (4300, 4),
- (4400, 4),
- (4500, 4),
- (4600, 2),
- ),
- request_dtype=torch.int64,
- global_k=0.37,
- global_v=2.5,
- slot_pages=24,
- layer=11,
- pool_pages=2048,
- ),
- ),
- )
-
- note = """
- Metadata consumed: none -- every input is an argument (as `qsa_paged_sparse_gqa_nvfp4`).
- Call sites (the target's `modeling.py`): `_sparse_gqa` over an NVFP4 pool for rows without
- the prefill pair that come `rows_per_request` > 1 a request -- the decode target's
- attention layers on a verify step (with the output gate) and a mixed step's generation
- rows. Draft-layer and one-row-a-request calls stay on `qsa_paged_sparse_gqa_nvfp4_decode`.
-
- Routing (`takes_grouped`, host-side, on shapes only): the grouped kernel from
- `MIN_GROUPS` (8) groups of T >= 2 rows up, on SM103, where its shared-memory plan fits
- (T <= 6; a block table of at most 612 pages of 32 tokens at T = 4, 892 at T = 3, 1364 at
- T = 2); the per-row kernel otherwise. Launch (`plan_launch`): within one wave of
- (group, KV head) units every unit takes the same splits, filling the GPU once (one CTA an
- SM), followed by the row merge; past one wave the units of the full waves take one
- CTA each (bf16 output, the gate fused) and the units of the partial last wave two to four,
- only their rows merged.
-
- GB300, one layer, cold L2, T = 4, the decode step's call (output gate on) on MTP3 verify
- selections dumped from a served run (a request's rows' union 1.4-1.7x a row's selection,
- about 8.4K tokens of context, the server's pool geometry): 64 rows 29.7 us, 128 rows 46.1,
- 256 rows 75.3, 512 rows 137, 1024 rows 233 -- against 39.5 / 69.4 / 135 / 237 / 461 us for
- the per-row kernel. Below 32 rows the call is the per-row kernel's.
-
- Numerics: dequantized K/V values are exact in f16; scores on f16 MMAs with fp32
- accumulation; P rounded to f16 (as the per-row kernel); a row's softmax offset moves only
- when its max grew by more than 2^8, the sum stays fp32 unrounded; the second-level scales
- act on the scores and the normalized output; the output gate's sigmoid on `ex2.approx` and
- `rcp.approx` (relative error ~1e-6) where the kernel applies it, as the merge's elsewhere.
- Within the decode op's band (4 ulp element,
- 2 ulp RMS) of the float64 reference and of the per-row kernel; not bitwise the per-row
- kernel (tiles hold the group's union). Deterministic; graph-capturable (row count is a
- runtime value; one compile per (T, splits, launch mode) variant).
-
- Preconditions: as `qsa_paged_sparse_gqa_nvfp4_decode` (bf16 q with head_dim 256 and
- 16-byte aligned rows; K and V planes, and their scale planes, share strides; power-of-two
- tokens per block; int32 row-contiguous block table; contiguous int32 selections), and
- expanded selections: each four-column quad of a row names one four-token group at its own
- offsets (column % 4 == token % 4), no group twice in a row -- what
- `qsa_topk_expand_decode` and `qsa_expand_block_indices` write. The rows of a group share
- the K/V tiles, so a NaN in a token one row selects reaches the group's other rows of that
- request.
- """
-
- def __call__(
- self,
- q,
- k_cache,
- v_cache,
- k_scale_cache,
- v_scale_cache,
- kv_scale_quant_orig,
- block_table,
- selected_tokens,
- request_indices,
- tokens_per_block: int,
- softmax_scale: float,
- rows_per_request: int,
- output_gate=None,
- ) -> torch.Tensor:
- return nvfp4_sparse_gqa_verify(
- q,
- k_cache,
- v_cache,
- k_scale_cache,
- v_scale_cache,
- kv_scale_quant_orig,
- block_table,
- selected_tokens,
- request_indices,
- tokens_per_block,
- softmax_scale,
- rows_per_request,
- output_gate=output_gate,
- )
-
- def reference(
- self,
- q,
- k_cache,
- v_cache,
- k_scale_cache,
- v_scale_cache,
- kv_scale_quant_orig,
- block_table,
- selected_tokens,
- request_indices,
- tokens_per_block: int,
- softmax_scale: float,
- rows_per_request: int,
- output_gate=None,
- ) -> torch.Tensor:
- """`qsa_paged_sparse_gqa_nvfp4_decode`'s reference (`rows_per_request` only groups the
- work, it does not change the result).
- """
- return qsa_paged_sparse_gqa_nvfp4_decode.reference(
- q,
- k_cache,
- v_cache,
- k_scale_cache,
- v_scale_cache,
- kv_scale_quant_orig,
- block_table,
- selected_tokens,
- request_indices,
- tokens_per_block,
- softmax_scale,
- output_gate=output_gate,
- )
-
- def is_valid(
- self,
- q,
- k_cache,
- v_cache,
- k_scale_cache,
- v_scale_cache,
- kv_scale_quant_orig,
- block_table,
- selected_tokens,
- request_indices,
- tokens_per_block: int,
- softmax_scale: float,
- rows_per_request: int,
- output_gate=None,
- ) -> None:
- qsa_paged_sparse_gqa_nvfp4_decode.is_valid(
- q,
- k_cache,
- v_cache,
- k_scale_cache,
- v_scale_cache,
- kv_scale_quant_orig,
- block_table,
- selected_tokens,
- request_indices,
- tokens_per_block,
- softmax_scale,
- output_gate=output_gate,
- )
- assert int(rows_per_request) >= 1, "rows_per_request >= 1"
- # The union is sized for expanded selections: a quad of columns names one four-token
- # group at its own offsets, and no group twice in a row. Anything else can overflow
- # the union, or count a token once where the reference counts it twice.
- rows, width = selected_tokens.shape
- quads = -(-width // 4)
- padded = torch.full((rows, 4 * quads), -1, dtype=torch.int64, device=selected_tokens.device)
- padded[:, :width] = selected_tokens.long()
- sel = padded.view(rows, quads, 4)
- valid = sel >= 0
- column = torch.arange(4, device=sel.device)
- assert bool(((sel % 4 == column) | ~valid).all()), "column % 4 == token % 4"
- group = torch.where(valid, sel // 4, torch.full_like(sel, -1)).amax(dim=2)
- assert bool(((sel // 4 == group[..., None]) | ~valid).all()), "one group a quad"
- listed = torch.where(group >= 0, group, torch.full_like(group, -1)).sort(dim=1).values
- repeat = (listed[:, 1:] == listed[:, :-1]) & (listed[:, 1:] >= 0)
- assert not bool(repeat.any()), "no four-token group twice in a row"
-
-
-qsa_paged_sparse_gqa_nvfp4_verify = QsaPagedSparseGqaNvfp4Verify()
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
deleted file mode 100644
index 415956a05579..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_prefill_compress.py
+++ /dev/null
@@ -1,276 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""QSA block compression of the paged index-key side cache: completed groups to anchors."""
-
-from __future__ import annotations
-
-from typing import Optional, Tuple
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-# Tier-2 sourcing: a vendored Triton callable reached as ordinary Python. QSA's
-# kernel plane has no torch.ops surface.
-from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels as _qsa_kernels
-
-#: The mRoPE sections the checkpoint declares; the target passes `None` today
-#: (`_QSA_MROPE_SECTION`) and these are certified, 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
-#: 128, rotary width 64 (32 cos/sin pairs), compress ratio 4, rms eps 1e-6,
-#: theta 1e7; `slots` x `width` is `qsa_block_table`. `pages` is the pool.
-_CKPT = dict(
- head_dim=128,
- tokens_per_block=32,
- slots=32,
- width=260,
- compress_ratio=4,
- eps=1e-6,
- mrope_section=None,
- logical_dtype=torch.int64,
- request_dtype=torch.int32,
- layout="plain",
-)
-
-#: 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."""
- return {**_CKPT, "segments": segments, "pages": pages, **overrides}
-
-
-def _mrope(segments, pages, **overrides) -> dict:
- """Prefill from 0 under mRoPE (11, 11, 10)."""
- return _spec(segments, pages, mrope_section=MROPE_SECTION, **overrides)
-
-
-class QsaPrefillCompress(OpWrapper):
- """Replace each completed group's anchor slot with its pooled, normed, rotated key.
-
- Wraps `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels.
- triton_qsa_prefill_compress`: one Triton launch, grid `(rows,)`. A row
- with `(L + 1) % cr == 0` (L = its logical position) is the anchor of the
- group `L - cr + 1 .. L`; every other row does nothing. Each member is
- resolved through `block_table` on its own, so groups may straddle pages::
-
- pooled = bf16(sum_f32(raw keys of the group) / cr)
- y = f32(pooled) * rsqrt(mean(f32(pooled)^2) + eps) * (1 + k_norm_weight)
- y = cos_sin.dtype(y)
- (t, h, w) = position_cache at the group's FIRST member
- pair q takes h if q % 3 == 1 and q < 3 * s1, w if q % 3 == 2 and
- q < 3 * s2, else t (t for every pair when mrope_section is None)
- NeoX-rotate y[:D/2] by cos_sin[pos[q], q], cos_sin[pos[q], R + q];
- y[D/2:] passes through (R = cos_sin.shape[1] // 2 = D / 4)
- index_cache[anchor] = bf16(y)
-
- Returns `index_cache`, written in place at anchor slots only; everything
- else, including the non-anchor members' raw keys, is untouched.
-
- Not idempotent: the anchor slot is the group's last input and its only
- output, so a second call over the same rows re-pools a compressed anchor
- (measured: all 48 of 48 anchors moved, up to 2.94).
-
- Fusion boundary: pool, Gemma norm, position gather, rotation, store. The
- raw-key and position stores (`qsa_paged_index_store`, which must run
- first), the table and metadata, and the main K/V pool are the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- 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 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, position strides (24672, 3, 1).
- Cell(
- why="NVFP4 trunk planes: 12 index roles and the position page share each page",
- spec=_mrope(((0, 64),) * 2, pages=4, layout="shared_page", roles=12),
- ),
- )
-
- note = """
- Metadata consumed: none -- every input is an argument. The target passes
- `qsa_logical_positions` (int64) and `qsa_req_idx_per_token` (int32)
- sliced to the step's rows, `qsa_block_table` unsliced, the layer's
- `get_index_k_buffer(i)` and the shared `get_qsa_position_buffer()`, the
- indexer's `k_layernorm` weight (bf16 [128]) and eps, a float32
- `[max_positions, 64]` cos|sin table built at load, and `mrope_section`
- None. It runs on every step, decode and prefill rows alike.
-
- Preconditions:
- - every member's raw key and the first member's coordinates are in the
- caches before the call; this call stores neither.
- - every logical position is non-negative. `L = -1` satisfies the anchor
- test and addresses positions -4..-1: measured, an illegal memory
- access that poisons the process. A value, so not guarded; pad with 0
- (measured an exact no-op).
- - `logical_positions`, `request_indices`, `k_norm_weight` and the last
- dim of `cos_sin` are read with an assumed unit stride while the
- callable checks only shape: each was measured accepted and wrong
- (max abs 4.72, 6.20, 0.80, 7.16). `is_valid` refuses them.
- The callable rejects loudly (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`
- or `compress_ratio`, a norm weight of the wrong length, CPU operands, and
- `mrope_section` that is not 3 entries summing to `R`. `rows == 0` returns
- before launching.
-
- `cos_sin`'s dtype is the rotation's precision: bf16 is accepted and was
- measured 1.4% of elements off by up to 1.56e-2 -- not certified. Text
- rows reduce mRoPE to plain RoPE bit for bit (measured, not pinned by a
- test). Plain `@triton.jit`, no
- autotune; CUDA-graph capturable at a fixed row count (captured at 256
- rows, 32 MTP7 verify blocks).
-
- Under an NVFP4 KV pool both caches are strided views of one per-slot
- allocation (each physical page holds the index roles, then the position
- page); the trunk-plane cell checks the other index roles and the position
- page untouched.
-
- The kernel's `tl.rsqrt` is the hardware approximation, an fp32 ulp from the
- reference's: every cell is bitwise the reference but the 32-block MTP7 one, where
- 1 of 8192 anchor elements rounds one bf16 step away.
-
- Hard limits (read from the kernel, not probed): positions, requests and
- page ids are widened to int64 before any address arithmetic, so no int32
- boundary inside the call. Certified extent: positions to 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 1928 or the table's rows, mRoPE outside
- the NVFP4 trunk-plane cell, the NVFP4 draft planes (one index role a page), multiple devices,
- `tokens_per_block` other than 32 (no group straddles a page), eps other
- than 1e-6, int32 positions or int64 request indices, a strided view of the
- index cache, position plane or table other than the NVFP4 pool's shared
- page, two rows naming one anchor in a launch, and `rows == 0`.
- """
-
- def __call__(
- self,
- *,
- logical_positions: torch.Tensor,
- request_indices: torch.Tensor,
- block_table: torch.Tensor,
- index_cache: torch.Tensor,
- position_cache: torch.Tensor,
- k_norm_weight: torch.Tensor,
- cos_sin: torch.Tensor,
- eps: float,
- tokens_per_block: int,
- compress_ratio: int,
- mrope_section: Optional[Tuple[int, int, int]],
- ) -> torch.Tensor:
- _qsa_kernels.triton_qsa_prefill_compress(
- logical_positions=logical_positions,
- request_indices=request_indices,
- block_table=block_table,
- index_cache=index_cache,
- position_cache=position_cache,
- k_norm_weight=k_norm_weight,
- cos_sin=cos_sin,
- eps=eps,
- tokens_per_block=tokens_per_block,
- compress_ratio=compress_ratio,
- mrope_section=mrope_section,
- )
- return index_cache
-
- def reference(
- self,
- *,
- logical_positions: torch.Tensor,
- request_indices: torch.Tensor,
- block_table: torch.Tensor,
- index_cache: torch.Tensor,
- position_cache: torch.Tensor,
- k_norm_weight: torch.Tensor,
- cos_sin: torch.Tensor,
- eps: float,
- tokens_per_block: int,
- compress_ratio: int,
- mrope_section: Optional[Tuple[int, int, int]],
- ) -> torch.Tensor:
- """The index cache after the call, as a fresh tensor, from the pre-call state."""
- out = index_cache.clone()
- logical = logical_positions.long()
- anchored = (logical + 1) % compress_ratio == 0
- anchor = logical[anchored]
- if anchor.numel() == 0:
- return out
- device = index_cache.device
- request = request_indices.long()[anchored]
- members = (
- anchor[:, None] - (compress_ratio - 1) + torch.arange(compress_ratio, device=device)
- )
- pages = block_table.long()[request[:, None], members // tokens_per_block]
- slots = members % tokens_per_block
-
- head_dim = index_cache.shape[3]
- raw = index_cache[pages, slots, 0].float() # [anchors, cr, D]
- pooled = (raw.sum(dim=1) / compress_ratio).to(torch.bfloat16).float()
- scale = torch.rsqrt(pooled.mul(pooled).sum(dim=-1, keepdim=True) / head_dim + eps)
- normed = (pooled * scale * (k_norm_weight.float() + 1.0)).to(cos_sin.dtype)
-
- t, h, w = position_cache[pages[:, 0], slots[:, 0], :3].long().unbind(-1)
- pairs_n = cos_sin.shape[1] // 2
- pairs = torch.arange(pairs_n, device=device)
- if mrope_section is None:
- pos = t[:, None].expand(-1, pairs_n)
- else:
- pos = torch.where(
- (pairs % 3 == 1) & (pairs < 3 * mrope_section[1]),
- h[:, None],
- torch.where(
- (pairs % 3 == 2) & (pairs < 3 * mrope_section[2]), w[:, None], t[:, None]
- ),
- )
- cos, sin = cos_sin[pos, pairs], cos_sin[pos, pairs_n + pairs]
- first, second = normed[:, :pairs_n], normed[:, pairs_n : 2 * pairs_n]
- rotated = torch.cat(
- [first * cos - second * sin, second * cos + first * sin, normed[:, 2 * pairs_n :]],
- dim=-1,
- )
- out[pages[:, -1], slots[:, -1], 0] = rotated.to(torch.bfloat16)
- return out
-
- def is_valid(
- self,
- *,
- logical_positions: torch.Tensor,
- request_indices: torch.Tensor,
- block_table: torch.Tensor,
- index_cache: torch.Tensor,
- position_cache: torch.Tensor,
- k_norm_weight: torch.Tensor,
- cos_sin: torch.Tensor,
- eps: float,
- tokens_per_block: int,
- compress_ratio: int,
- mrope_section: Optional[Tuple[int, int, int]],
- ) -> None:
- # Read by the kernel with an assumed unit stride; the callable checks
- # only their shapes. Each measured accepted and silently wrong.
- assert logical_positions.is_contiguous(), "logical_positions must be contiguous"
- assert request_indices.is_contiguous(), "request_indices must be contiguous"
- assert k_norm_weight.is_contiguous(), "k_norm_weight must be contiguous"
- # `ndim != 2` is the callable's own loud error; leave it to the callable.
- assert cos_sin.ndim != 2 or cos_sin.stride(1) == 1, (
- "cos_sin must have unit stride along its last dimension"
- )
-
-
-qsa_prefill_compress = QsaPrefillCompress()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rollback_planes.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rollback_planes.py
deleted file mode 100644
index 144df3bd24fc..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rollback_planes.py
+++ /dev/null
@@ -1,144 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""The speculative rollback of every attention layer's QSA index-K plane in one launch each way:
-gather the verify rows' slots of all planes before the trunk writes them, and after the
-draft loop write the gathered values back into the slots of the rows the sampler rejected.
-
-Tier 5 (Triton written for this target). The planes are separate buffers of the cache
-manager (one per attention layer, each `[pages, tokens_per_block, heads, dim]` with a free
-page stride and a contiguous tail); the kernels reach them through a table of their base
-addresses and page strides, one program per (plane, row). The composition this replaces ran
-one gather per layer during the trunk and a gather, a `torch.where` and an indexed write per
-layer after the draft loop.
-"""
-
-from __future__ import annotations
-
-import torch
-import triton
-import triton.language as tl
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: This checkpoint's index plane row: one index KV head of 128 bf16.
-ROW = 128
-
-
-@triton.jit
-def _gather_kernel(ptrs, page_strides, pages, within, out, rows, ROW: tl.constexpr):
- plane = tl.program_id(0)
- r = tl.program_id(1)
- base = tl.load(ptrs + plane).to(tl.pointer_type(tl.bfloat16))
- stride = tl.load(page_strides + plane)
- page = tl.load(pages + r)
- tok = tl.load(within + r)
- offs = tl.arange(0, ROW)
- vals = tl.load(base + page * stride + tok * ROW + offs)
- tl.store(out + (plane.to(tl.int64) * rows + r) * ROW + offs, vals)
-
-
-@triton.jit
-def _restore_kernel(ptrs, page_strides, pages, within, rejected, old, rows, ROW: tl.constexpr):
- plane = tl.program_id(0)
- r = tl.program_id(1)
- if tl.load(rejected + r) != 0:
- base = tl.load(ptrs + plane).to(tl.pointer_type(tl.bfloat16))
- stride = tl.load(page_strides + plane)
- page = tl.load(pages + r)
- tok = tl.load(within + r)
- offs = tl.arange(0, ROW)
- vals = tl.load(old + (plane.to(tl.int64) * rows + r) * ROW + offs)
- tl.store(base + page * stride + tok * ROW + offs, vals)
-
-
-def plane_table(planes: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
- """`(base addresses, page strides)` of the planes, int64 on their device (built once per
- cache-manager bind: the planes do not move)."""
- dev = planes[0].device
- for t in planes:
- assert t.dtype == torch.bfloat16 and t.dim() == 4, "index planes [pages, tpb, H, D] bf16"
- assert t.shape[2] * t.shape[3] == ROW, "one 128-wide row a token"
- assert t.stride(1) == ROW and t.stride(2) == t.shape[3] and t.stride(3) == 1, (
- "contiguous [tpb, H, D] tail"
- )
- ptrs = torch.tensor([t.data_ptr() for t in planes], dtype=torch.int64, device=dev)
- strides = torch.tensor([t.stride(0) for t in planes], dtype=torch.int64, device=dev)
- return ptrs, strides
-
-
-class QsaRollbackGather(OpWrapper):
- """`torch.stack([plane[pages, within] for plane in planes])` as `[L, R, 128]` bf16, with
- the planes given as `plane_table(planes)`: every row's slot of every plane, in plane
- order. Nothing is written but the output.
-
- Fusion boundary: the gather of all planes; the slots (`pages`, `within`) are the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
- CELLS: tuple[Cell, ...] = (
- Cell(why="12 planes, 4 verify rows (MTP3, one request)", spec=dict(planes=12, rows=4)),
- Cell(why="12 planes, 512 verify rows (MTP3, 128 requests)", spec=dict(planes=12, rows=512)),
- )
- note = """
- Metadata consumed: none. Call site: `SpeculativeRollback.begin`, before the verify step's
- trunk writes any plane. One program per (plane, row), one 256-byte row each.
- Preconditions: `pages` / `within` int64 `[R]`; the planes as `plane_table` accepts them.
- """
-
- def __call__(self, ptrs, page_strides, pages, within, num_planes: int) -> torch.Tensor:
- rows = pages.shape[0]
- out = torch.empty([num_planes, rows, ROW], dtype=torch.bfloat16, device=pages.device)
- if rows:
- _gather_kernel[(num_planes, rows)](
- ptrs, page_strides, pages, within, out, rows, ROW=ROW, num_warps=1
- )
- return out
-
- def reference(self, ptrs, page_strides, pages, within, num_planes: int) -> torch.Tensor:
- raise NotImplementedError("the cell harness builds planes and compares against torch")
-
- def is_valid(self, ptrs, page_strides, pages, within, num_planes: int) -> None:
- assert ptrs.dtype == torch.int64 and ptrs.numel() == num_planes
- assert page_strides.dtype == torch.int64 and page_strides.numel() == num_planes
- assert pages.dtype == torch.int64 and within.dtype == torch.int64
- assert pages.shape == within.shape and pages.dim() == 1
-
-
-class QsaRollbackRestore(OpWrapper):
- """For every row `r` with `rejected[r]` and every plane `l`: `plane_l[pages[r], within[r]]
- = old[l, r]` -- the inverse of `qsa_rollback_gather` over the rejected rows. Rows not
- rejected are not touched.
-
- Fusion boundary: the restore of all planes; `rejected` is the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
- CELLS: tuple[Cell, ...] = (
- Cell(why="12 planes, 4 verify rows, two rejected", spec=dict(planes=12, rows=4)),
- Cell(why="12 planes, 512 verify rows", spec=dict(planes=12, rows=512)),
- )
- note = """
- Metadata consumed: none. Call site: `SpeculativeRollback.commit_speculative_states`,
- after the draft loop. `rejected` int32/int8/bool `[R]` (nonzero = restore).
- """
-
- def __call__(self, ptrs, page_strides, pages, within, rejected, old) -> None:
- num_planes, rows, _ = old.shape
- if rows:
- _restore_kernel[(num_planes, rows)](
- ptrs, page_strides, pages, within, rejected, old, rows, ROW=ROW, num_warps=1
- )
-
- def reference(self, ptrs, page_strides, pages, within, rejected, old) -> None:
- raise NotImplementedError("the cell harness builds planes and compares against torch")
-
- def is_valid(self, ptrs, page_strides, pages, within, rejected, old) -> None:
- assert old.dtype == torch.bfloat16 and old.dim() == 3 and old.shape[2] == ROW
- assert old.is_contiguous()
- assert ptrs.numel() == old.shape[0] and page_strides.numel() == old.shape[0]
- assert pages.shape == (old.shape[1],) and within.shape == (old.shape[1],)
- assert rejected.shape == (old.shape[1],)
-
-
-qsa_rollback_gather = QsaRollbackGather()
-qsa_rollback_restore = QsaRollbackRestore()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py
deleted file mode 100644
index ae062832dcb6..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py
+++ /dev/null
@@ -1,651 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""QSA's block top-k and its expansion to token indices, one launch per batch of rows.
-
-Tier 5 (CuTe DSL written for this target). The decode path ran `qsa_topk_select`
-(the CUDA indexer top-k, one CTA a row) and then `qsa_expand_block_indices` (one
-more launch) per attention layer. Here one CTA per row holds four consecutive scores a
-thread and narrows to the k-th largest by radix passes over the scores' bit patterns
-(shared-memory histograms of 11 bits), starting below the bits every visible score
-shares. It stops as soon as the bin holding the threshold has at most `CAP` keys and
-ranks those exactly by pairwise comparison (ties to the higher block, the canonical
-selection's rule). Every selected block is then placed in block order by one block-wide
-scan, and expanded to token indices with the causal tail, as the expansion entry does.
-
-Scores are a sum of ReLUs times a positive scale, so non-negative: their fp32 bit
-patterns, read as int32, order like their values (a negative zero is taken as zero, a
-NaN ranks above everything).
-"""
-
-from __future__ import annotations
-
-import cutlass
-import cutlass.cute as cute
-import torch
-from cutlass import const_expr
-from cutlass.cute.runtime import from_dlpack
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
-from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
-
-from .qsa_expand_block_indices import qsa_expand_block_indices as _expand
-
-#: Consecutive scores a thread holds (one 16-byte load).
-KPT = 4
-#: Score columns a row may have: `KPT` a thread, at most 1024 threads a CTA.
-MAX_WIDTH = KPT * 1024
-#: CTAs an SM the register allocation is held to. At the 2592-wide score table a row is 672
-#: threads, which the per-subpartition register file holds once an SM at the compiler's own
-#: 48 registers a thread; at three (32 a thread) rows overlap each other's barriers and
-#: shared-memory round trips (GB300, a mixed step's 16256 context rows: 426 -> 236 us on
-#: half-zero rows, 486 -> 284 on relu-sparse ones; 128 decode rows 8.2 us both ways).
-MIN_BLOCKS = 3
-#: Radix passes at most (11 bits each; 31 magnitude bits).
-_PASSES = 3
-_HIST = 2048
-_I32_MAX = 0x7FFFFFFF
-#: At most this many keys in the threshold's bin are ranked pairwise instead of
-#: narrowed by another pass.
-CAP = 128
-#: Threads that own the 2048 bins of a scan, four each.
-_SCAN_THREADS = _HIST // 4
-#: Consecutive threads that rank one candidate of the threshold's bin, each over every
-#: `_RANK_LANES`-th candidate (`_SCAN_THREADS` threads cover `CAP` candidates).
-_RANK_LANES = _SCAN_THREADS // CAP
-
-
-def _threads(width: int) -> int:
- return max(_SCAN_THREADS, -(-width // (KPT * 32)) * 32)
-
-
-class _TopkExpandKernel:
- def __init__(self, *, width, key_stride, block_topk, compress, token_topk, pdl):
- self.width, self.key_stride, self.k, self.cr, self.token_topk, self.pdl = (
- width,
- key_stride,
- block_topk,
- compress,
- token_topk,
- pdl,
- )
- self.final = token_topk + compress - 1
- self.threads = _threads(width)
- self.warps = self.threads // 32
- self.flag_words = -(-width // 32)
-
- @cute.jit
- def __call__(
- self,
- keys_t: cute.Tensor,
- end_ptr: cute.Pointer,
- pos_ptr: cute.Pointer,
- len_ptr: cute.Pointer,
- topk_ptr: cute.Pointer,
- sel_ptr: cute.Pointer,
- rows: cutlass.Int32,
- stream,
- ):
- # The keys arrive as a tensor, as the TVM-FFI entry takes its stream from a GPU
- # tensor; their layout is rebuilt with the static width and row stride the 16-byte
- # loads need. The row count is the call's.
- mKeys = cute.make_tensor(
- keys_t.iterator, cute.make_layout((rows, self.width), stride=(self.key_stride, 1))
- )
- per_row = cute.make_layout((rows,), stride=(1,))
- mEnd = cute.make_tensor(end_ptr, per_row)
- mPos = cute.make_tensor(pos_ptr, per_row)
- mLen = cute.make_tensor(len_ptr, per_row)
- mTopk = cute.make_tensor(topk_ptr, cute.make_layout((rows, self.k), stride=(self.k, 1)))
- mSel = cute.make_tensor(
- sel_ptr, cute.make_layout((rows, self.final), stride=(self.final, 1))
- )
- ld_keys = cute.make_copy_atom(
- cute.nvgpu.CopyG2ROp(), cutlass.Int32, num_bits_per_copy=32 * KPT
- )
- self.kernel(mKeys, mEnd, mPos, mLen, mTopk, mSel, ld_keys).launch(
- grid=[rows, 1, 1],
- block=[self.threads, 1, 1],
- smem=(_HIST + self.flag_words + 2 * CAP + 64 + 16) * 4 + 128,
- stream=stream,
- use_pdl=self.pdl,
- min_blocks_per_mp=MIN_BLOCKS,
- )
-
- @cute.kernel
- def kernel(self, mKeys, mEnd, mPos, mLen, mTopk, mSel, ld_keys: cute.CopyAtom):
- tidx, _, _ = cute.arch.thread_idx()
- row, _, _ = cute.arch.block_idx()
- warp = cute.arch.warp_idx()
- lane = cute.arch.lane_idx()
- nt: cutlass.Constexpr = self.threads
- nw: cutlass.Constexpr = self.warps
- k: cutlass.Constexpr = self.k
- cr: cutlass.Constexpr = self.cr
- smem = cutlass.utils.SmemAllocator()
- hist = smem.allocate_tensor(
- cutlass.Int32, cute.make_layout((_HIST,), stride=(1,)), byte_alignment=16
- )
- flags = smem.allocate_tensor(
- cutlass.Int32, cute.make_layout((self.flag_words,), stride=(1,)), byte_alignment=16
- )
- cand_key = smem.allocate_tensor(
- cutlass.Int32, cute.make_layout((CAP,), stride=(1,)), byte_alignment=16
- )
- cand_idx = smem.allocate_tensor(
- cutlass.Int32, cute.make_layout((CAP,), stride=(1,)), byte_alignment=16
- )
- wa = smem.allocate_tensor(
- cutlass.Int32, cute.make_layout((32,), stride=(1,)), byte_alignment=16
- )
- wb = smem.allocate_tensor(
- cutlass.Int32, cute.make_layout((32,), stride=(1,)), byte_alignment=16
- )
- state = smem.allocate_tensor(
- cutlass.Int32, cute.make_layout((16,), stride=(1,)), byte_alignment=16
- )
-
- # 0. Shared state, before the wait (nothing global is read here).
- for b in cutlass.range_constexpr(-(-_HIST // nt)):
- if tidx + b * nt < _HIST:
- hist[tidx + b * nt] = cutlass.Int32(0)
- for b in cutlass.range_constexpr(-(-self.flag_words // nt)):
- if tidx + b * nt < self.flag_words:
- flags[tidx + b * nt] = cutlass.Int32(0)
- if tidx == 0:
- state[5] = cutlass.Int32(0) # candidates compacted
- # The row's metadata, host-filled: read before the wait too.
- n = mEnd[row]
- seqlen = mLen[row]
- visible = cutlass.Int32(mPos[row]) + 1
- if const_expr(self.pdl):
- cute.arch.griddepcontrol_wait()
- cute.arch.griddepcontrol_launch_dependents()
-
- # 1. Four consecutive keys a thread (a column past the window is no candidate:
- # -1); the largest and smallest key in the window.
- e0 = tidx * KPT
- keys = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32)
- keys.fill(-1)
- if e0 < self.width:
- cute.copy(ld_keys, cute.logical_divide(mKeys, (None, KPT))[row, (None, tidx)], keys)
- kmax = cutlass.Int32(-1)
- kmin = cutlass.Int32(_I32_MAX)
- for c in cutlass.range_constexpr(KPT):
- b = keys[c]
- if b < 0: # -0.0 (or a negative): rank as zero
- b = cutlass.Int32(0)
- if e0 + c >= n:
- b = cutlass.Int32(-1)
- keys[c] = b
- kmax = cutlass.max(kmax, b)
- if b >= 0:
- kmin = cutlass.min(kmin, b)
- kmax = cute.arch.warp_reduction_max(kmax)
- kmin = cute.arch.warp_reduction(kmin, lambda x, y: cutlass.min(x, y))
- if lane == 0:
- wa[warp] = kmax
- wb[warp] = kmin
- cute.arch.sync_threads()
- gmax = cutlass.Int32(-1)
- gmin = cutlass.Int32(_I32_MAX)
- if lane < nw:
- gmax = wa[lane]
- gmin = wb[lane]
- gmax = cute.arch.warp_reduction_max(gmax)
- gmin = cute.arch.warp_reduction(gmin, lambda x, y: cutlass.min(x, y))
-
- # 2. Every key in the window shares the bits above the highest bit where the
- # largest and the smallest differ (`hb`); radix passes walk down from there, 11 bits
- # at a time, until the threshold's bin holds at most CAP keys or no bits are left.
- diff = gmax ^ gmin
- hb = cutlass.Int32(-1)
- if diff != 0:
- hb = 31 - cutlass.Int32(cute.arch.clz(diff))
- low = (cutlass.Int64(1) << cutlass.Int64(hb + 1)) - 1
- mask0 = cutlass.Int32(cutlass.Int64(_I32_MAX) & (cutlass.Int64(-1) ^ low))
- if tidx == 0:
- state[0] = gmax & mask0 # threshold prefix
- state[1] = mask0 # its mask
- state[2] = cutlass.Int32(k) # still needed at the threshold
- state[3] = n # keys in the threshold's bin
- state[4] = cutlass.Int32(1) # narrowing done
- state[6] = cutlass.max(hb - 10, 0) # the next pass's shift
- if n > k:
- if hb >= 0:
- state[4] = cutlass.Int32(0)
- cute.arch.sync_threads()
- for _p in cutlass.range_constexpr(_PASSES):
- if state[4] == 0:
- prefix = state[0]
- mask = state[1]
- need = state[2]
- shift = state[6]
- for c in cutlass.range_constexpr(KPT):
- if keys[c] >= 0:
- if (keys[c] & mask) == prefix:
- digit = (keys[c] >> shift) & (_HIST - 1)
- cute.arch.atomic_add(
- hist.iterator + digit, cutlass.Int32(1), scope="cta"
- )
- cute.arch.sync_threads()
- # Thread t < 512 owns bins 2047-4t .. 2044-4t, counted from the top; each is
- # cleared for the next pass once read.
- h = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), cutlass.Int32)
- h.fill(0)
- own = cutlass.Int32(0)
- if tidx < _SCAN_THREADS:
- for j in cutlass.range_constexpr(4):
- h[j] = hist[_HIST - 1 - 4 * tidx - j]
- hist[_HIST - 1 - 4 * tidx - j] = cutlass.Int32(0)
- own = own + h[j]
- incl = own
- for off in (1, 2, 4, 8, 16):
- up = cute.arch.shuffle_sync_up(incl, offset=off)
- if lane >= off:
- incl = incl + up
- if lane == 31:
- wa[warp] = incl
- cute.arch.sync_threads()
- above = cutlass.Int32(0)
- if lane < warp:
- above = wa[lane]
- before = cute.arch.warp_reduction_sum(above) + incl - own # keys in the bins above
- if tidx < _SCAN_THREADS:
- if before < need:
- if need <= before + own:
- run = before
- for j in cutlass.range_constexpr(4):
- if run < need:
- if need <= run + h[j]:
- state[0] = prefix | ((_HIST - 1 - 4 * tidx - j) << shift)
- state[1] = mask | ((_HIST - 1) << shift)
- state[2] = need - run
- state[3] = h[j]
- state[4] = cutlass.Int32(
- cutlass.select_((h[j] <= CAP) | (shift == 0), 1, 0)
- )
- state[6] = cutlass.max(shift - 11, 0)
- run = run + h[j]
- cute.arch.sync_threads()
-
- prefix = state[0]
- mask = state[1]
- need = state[2]
- at_thr = state[3]
- # 3. At most CAP keys in the threshold's bin: rank them exactly, ties to the
- # higher block, and flag the `need` best.
- if n > k:
- if at_thr <= CAP:
- for c in cutlass.range_constexpr(KPT):
- if keys[c] >= 0:
- if (keys[c] & mask) == prefix:
- slot = cute.arch.atomic_add(
- state.iterator + 5, cutlass.Int32(1), scope="cta"
- )
- cand_key[slot] = keys[c]
- cand_idx[slot] = e0 + c
- cute.arch.sync_threads()
- # Candidate `ci` is ranked by `_RANK_LANES` consecutive lanes, lane `sub`
- # counting the better ones among candidates sub, sub + _RANK_LANES, ...;
- # butterfly shuffles add the counts up.
- ci = tidx // _RANK_LANES
- sub = tidx % _RANK_LANES
- rank = cutlass.Int32(0)
- it = cutlass.Int32(0)
- if ci < at_thr:
- kt = cand_key[ci]
- it = cand_idx[ci]
- for j in cutlass.range(sub, at_thr, _RANK_LANES, unroll=4):
- kj = cand_key[j]
- ij = cand_idx[j]
- better = (kj > kt) | ((kj == kt) & (ij > it))
- rank = rank + cutlass.Int32(cutlass.select_(better, 1, 0))
- for off in cutlass.range_constexpr(_RANK_LANES.bit_length() - 1):
- rank = rank + cute.arch.shuffle_sync_bfly(rank, offset=1 << off)
- if ci < at_thr:
- if sub == 0:
- if rank < need:
- cute.arch.atomic_or(
- flags.iterator + (it >> 5),
- cutlass.Int32(1) << (it & 31),
- scope="cta",
- )
- cute.arch.sync_threads()
-
- # 4. Block order: one scan of (selected, tied at the threshold) counts. A tied key
- # is left only when the passes ran out of bits with more than CAP keys tied; the
- # last `need` of those in block order are taken.
- take = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32)
- tied = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32)
- own = cutlass.Int32(0)
- for c in cutlass.range_constexpr(KPT):
- t_c = cutlass.Int32(0)
- q_c = cutlass.Int32(0)
- if keys[c] >= 0:
- if n <= k:
- t_c = cutlass.Int32(1)
- else:
- mk = keys[c] & mask
- if mk > prefix:
- t_c = cutlass.Int32(1)
- if mk == prefix:
- if at_thr > CAP:
- q_c = cutlass.Int32(1)
- else:
- e = e0 + c
- t_c = (flags[e >> 5] >> (e & 31)) & 1
- take[c] = t_c
- tied[c] = q_c
- own = own + t_c + (q_c << 16)
- incl = own
- for off in (1, 2, 4, 8, 16):
- up = cute.arch.shuffle_sync_up(incl, offset=off)
- if lane >= off:
- incl = incl + up
- if lane == 31:
- wb[warp] = incl
- cute.arch.sync_threads()
- above = cutlass.Int32(0)
- if lane < warp:
- above = wb[lane]
- pre = cute.arch.warp_reduction_sum(above) + incl - own
- taken_before = pre & 0xFFFF
- tied_before = pre >> 16
- # The tied keys left out, the lowest-index ones (none unless the threshold's bin
- # holds more keys than are still needed).
- skip = cutlass.max(at_thr - need, 0)
- for c in cutlass.range_constexpr(KPT):
- e = e0 + c
- sel = cutlass.Int32(0)
- pos = cutlass.Int32(0)
- if tied[c] == 1:
- if tied_before >= skip:
- sel = cutlass.Int32(1)
- pos = taken_before + tied_before - skip
- tied_before = tied_before + 1
- if take[c] == 1:
- sel = cutlass.Int32(1)
- pos = taken_before + cutlass.max(tied_before - skip, 0)
- taken_before = taken_before + 1
- if sel == 1:
- mTopk[row, pos] = e
- for j in cutlass.range_constexpr(cr):
- tok = e * cr + j
- val = cutlass.Int32(-1)
- if tok < seqlen:
- val = tok
- mSel[row, pos * cr + j] = val
-
- # 5. Padding of the block list; the causal tail after the selected tokens, then -1.
- total = n
- if n > k:
- total = cutlass.Int32(k)
- for c in cutlass.range_constexpr(-(-k // nt)):
- s = tidx + c * nt
- if s < k:
- if s >= total:
- mTopk[row, s] = cutlass.Int32(-1)
- tail_start = (visible // cr) * cr
- tail_count = visible - tail_start
- body = total * cr
- if body > self.token_topk:
- body = cutlass.Int32(self.token_topk)
- for c in cutlass.range_constexpr(-(-self.final // nt)):
- col = tidx + c * nt
- if col < self.final:
- if col >= body:
- off = col - body
- val = cutlass.Int32(-1)
- if off < cr - 1:
- if off < tail_count:
- if tail_start + off < seqlen:
- val = tail_start + off
- mSel[row, col] = val
-
-
-#: One compiled launch per score-table geometry, k, compress ratio, token budget, index
-#: dtypes and PDL; the row count is a runtime argument.
-_COMPILED: dict = {}
-_INDEX_TYPES = {torch.int32: cutlass.Int32, torch.int64: cutlass.Int64}
-
-
-def launch(
- scores,
- row_ends,
- query_positions,
- sequence_lengths,
- topk_indices,
- compress_ratio,
- token_topk,
- pdl=True,
- out=None,
-):
- """The launch behind `QsaTopkExpandDecode.__call__`: writes `topk_indices` and returns the
- token selection `[rows, token_topk + compress_ratio - 1]` int32 (`out` when given)."""
- rows, width = scores.shape
- block_topk = topk_indices.shape[1]
- sel = out
- if sel is None:
- sel = torch.empty(
- (rows, token_topk + compress_ratio - 1), dtype=torch.int32, device=scores.device
- )
- if rows == 0:
- return sel
- pdl = bool(pdl and get_env_enable_pdl())
- keys = scores.view(torch.int32)
- key = (
- width,
- scores.stride(0),
- block_topk,
- compress_ratio,
- token_topk,
- query_positions.dtype,
- sequence_lengths.dtype,
- pdl,
- )
- ptrs = (
- row_ends.data_ptr(),
- query_positions.data_ptr(),
- sequence_lengths.data_ptr(),
- topk_indices.data_ptr(),
- sel.data_ptr(),
- )
- fn = _COMPILED.get(key)
- if fn is None:
- gm = cute.AddressSpace.gmem
- fn = _COMPILED[key] = cute.compile(
- _TopkExpandKernel(
- width=width,
- key_stride=scores.stride(0),
- block_topk=block_topk,
- compress=compress_ratio,
- token_topk=token_topk,
- pdl=pdl,
- ),
- from_dlpack(keys, assumed_align=16).mark_layout_dynamic(leading_dim=1),
- make_ptr(cutlass.Int32, ptrs[0], gm, assumed_align=4),
- make_ptr(
- _INDEX_TYPES[query_positions.dtype],
- ptrs[1],
- gm,
- assumed_align=query_positions.element_size(),
- ),
- make_ptr(
- _INDEX_TYPES[sequence_lengths.dtype],
- ptrs[2],
- gm,
- assumed_align=sequence_lengths.element_size(),
- ),
- make_ptr(cutlass.Int32, ptrs[3], gm, assumed_align=4),
- make_ptr(cutlass.Int32, ptrs[4], gm, assumed_align=4),
- cutlass.Int32(rows),
- cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
- options="--enable-tvm-ffi",
- )
- fn(keys, *ptrs, rows)
- return sel
-
-
-def supports(scores) -> bool:
- """Whether a score table this wide is one this entry serves (see `is_valid`)."""
- return scores.shape[1] % KPT == 0 and scores.shape[1] <= MAX_WIDTH
-
-
-class QsaTopkExpandDecode(OpWrapper):
- """`qsa_topk_select` (zero row starts, `row_ends` the visible block counts),
- `qsa_canonical_selection` and `qsa_expand_block_indices`, in one launch: one CTA a row, for
- decode rows and a context chunk's rows alike.
-
- Per row: the `block_topk` largest of `scores[row, :row_ends[row]]` (non-negative fp32),
- ties at the threshold to the higher block, written to `topk_indices[row]` in block order
- and padded with -1 -- the composition's selection after `qsa_canonical_selection`; the
- returned token selection is the expansion entry's: each selected group's members below
- `sequence_lengths[row]`, packed at the front, then the incomplete group at the causal
- frontier (`query_positions[row] + 1` visible tokens), then -1. `topk_indices` is written
- in place, and the selection into `out` when given (a contiguous int32
- `[rows, token_topk + compress_ratio - 1]`, e.g. a row range of the step's whole
- selection); nothing else is written.
-
- Fusion boundary: the selection and the expansion. Scoring is the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="decode, 4 rows across the budget, one below it",
- spec=dict(rows=4, pos=(8700, 4095, 1500, 9900)),
- ),
- # The target writes a context chunk's rows through `out` into its rows of the
- # step's selection; relu-sparse scores put the zero class (negative zeros among it)
- # across the budget, the ties the canonical rule decides.
- Cell(
- why="context chunk: 2048 rows of one request (positions 6144..8191), relu-sparse "
- "scores, written through `out` into rows of a step-sized selection",
- spec=dict(rows=2048, chunk=(6144, 2048), relu_sparse=True, out_margin=64),
- ),
- Cell(
- why="decode, 1 row of scores within 5% of each other",
- spec=dict(rows=1, pos=(8700,), narrow=True),
- ),
- # Serving's decode rows at 8K: positive, log-normally spread over ~2 octaves, a
- # zero or two; the first pass leaves tens to ~CAP candidates at the threshold.
- Cell(
- why="decode, 16 rows of log-normal scores with a few zeros: threshold bins of "
- "up to CAP candidates ranked pairwise, ties included",
- spec=dict(rows=16, pos=tuple(8192 + 37 * r for r in range(16)), lognormal=True),
- ),
- )
-
- note = """
- Metadata consumed: `qsa_visible_blocks` (`row_ends`), the rows' logical positions
- and sequence lengths. Preconditions: zero row starts; `scores` fp32 `[rows, W]`, unit
- column stride, 16-byte aligned rows, `W % 4 == 0` and `W <= 4096` (four keys a thread,
- at most 1024 threads); `topk_indices` int32 `[rows, K]` contiguous; `row_ends` int32,
- the positions and lengths int32 or int64, each a contiguous `[rows]`. Ties: the higher
- block at the threshold, as `qsa_canonical_selection` rewrites the CUDA top-k's pick.
-
- Compiled once per score-table width and row stride (the block table's, fixed for an
- engine), k, compress ratio, token budget, index dtypes and PDL: the row count is a
- runtime argument, so an eager step's context rows, any count, reuse the kernel (read
- from the launch, not pinned by a test).
-
- Certified: decode batches of 1 and 4 rows over a 2560-wide score table, a 2048-row
- context chunk written through `out`, and a replay under a CUDA graph that reads the
- scores after its wait.
- """
-
- def __call__(
- self,
- scores,
- row_ends,
- query_positions,
- sequence_lengths,
- topk_indices,
- compress_ratio,
- token_topk,
- out=None,
- ):
- return launch(
- scores,
- row_ends,
- query_positions,
- sequence_lengths,
- topk_indices,
- compress_ratio,
- token_topk,
- out=out,
- )
-
- def reference(
- self,
- scores,
- row_ends,
- query_positions,
- sequence_lengths,
- topk_indices,
- compress_ratio,
- token_topk,
- out=None,
- ):
- """`(topk_indices_after, selection)`: a stable descending sort per row (ties to the
- higher block), the k first in block order, then the expansion entry's reference."""
- rows, _ = scores.shape
- k = topk_indices.shape[1]
- out = torch.full_like(topk_indices, -1)
- for r in range(rows):
- n = int(row_ends[r])
- s = scores[r, :n].float().clamp(min=0.0)
- # A stable sort of the reversed row puts the higher block first among equals.
- order = (n - 1) - torch.sort(torch.flip(s, [0]), descending=True, stable=True).indices[
- :k
- ]
- chosen = torch.sort(order).values
- out[r, : chosen.numel()] = chosen.to(torch.int32)
- sel = _expand.reference(
- out,
- query_positions,
- sequence_lengths,
- compress_ratio=compress_ratio,
- token_topk=token_topk,
- )
- return out, sel
-
- def is_valid(
- self,
- scores,
- row_ends,
- query_positions,
- sequence_lengths,
- topk_indices,
- compress_ratio,
- token_topk,
- out=None,
- ):
- assert scores.dtype == torch.float32 and scores.stride(1) == 1, (
- "fp32 scores, unit column stride"
- )
- assert scores.shape[1] % KPT == 0 and scores.shape[1] <= MAX_WIDTH, "W % 4 == 0, W <= 4096"
- assert scores.stride(0) % KPT == 0 and scores.data_ptr() % 16 == 0, "16-byte aligned rows"
- assert topk_indices.dtype == torch.int32 and topk_indices.is_contiguous(), (
- "int32 topk_indices"
- )
- rows = scores.shape[0]
- assert row_ends.dtype == torch.int32 and row_ends.shape == (rows,), "int32 row_ends [rows]"
- for t in (query_positions, sequence_lengths):
- assert t.dtype in _INDEX_TYPES and t.shape == (rows,), "int32/int64 [rows]"
- assert all(t.is_contiguous() for t in (row_ends, query_positions, sequence_lengths)), (
- "contiguous row_ends, positions and lengths"
- )
- assert token_topk == topk_indices.shape[1] * compress_ratio, (
- "token_topk = k * compress_ratio"
- )
- assert out is None or (
- out.dtype == torch.int32
- and out.is_contiguous()
- and out.shape == (scores.shape[0], token_topk + compress_ratio - 1)
- ), "out: contiguous int32 [rows, token_topk + compress_ratio - 1]"
-
-
-qsa_topk_expand_decode = QsaTopkExpandDecode()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/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
deleted file mode 100644
index fa7e00d6e456..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py
+++ /dev/null
@@ -1,213 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""QSA block selection: an exact per-row top-k over a caller-supplied column window."""
-
-from __future__ import annotations
-
-from typing import NamedTuple
-
-import torch
-
-import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: This checkpoint's block budget: `indexer_budget` 2048 tokens over
-#: `indexer_compress_ratio` 4. The op's own default, 2048, is the *token* budget.
-_BLOCK_TOPK = 512
-#: Score-row width the target hands over: a 260-page block table at 32 tokens
-#: per block, compressed 4:1.
-_SCORE_COLUMNS = 260 * 32 // 4
-#: 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):
- """What a correct selection must satisfy.
-
- The op's output order is unspecified, so `reference` cannot return the
- indices themselves; it returns each row's expected top values plus the
- frame the test needs to read the op's offsets back as values.
- """
-
- #: `[rows, index_topk]` fp32: the window's top values, descending, `-inf`-padded.
- values: torch.Tensor
- #: The logits as fp32 with every out-of-window column at `-inf`.
- logits: torch.Tensor
- #: int64 `row_starts`, to lift window-local offsets to columns.
- row_starts: torch.Tensor
- #: int64 window length per row, `max(row_ends - row_starts, 0)`.
- window: torch.Tensor
-
-
-class QsaTopkSelect(OpWrapper):
- """Each row's top-`index_topk` window-local column offsets, written into `indices`.
-
- Wraps `torch.ops.trtllm.indexer_topk_prefill` (CUDA, compiled into the
- wheel) 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
- is written, so `indices` needs no pre-fill. Returns `indices`.
-
- Exact, not approximate: the kernel resolves the threshold over the full 32
- bits of the fp32 pattern. The **order of the offsets is unspecified** and
- not stable between identical launches (candidates are appended through a
- shared-memory atomic; measured at `[8, 2048]`, `k=512`: four launches, four
- permutations of one selected set). Ties are broken arbitrarily within the
- tied value.
-
- Fusion boundary: selection only. Scoring, the causal bound the caller
- encodes in `row_ends`, and the block-to-token expansion are outside.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # The target calls this with `index_topk = 512` over the `[rows, 2080]`
- # score matrix `qsa_paged_index_scores` returns with
- # `only_visible_blocks=True` -- so every column at or past a row's
- # `row_end` is allocator garbage. Those cells poison it with NaN: a NaN
- # inside the window would be selected (see `note`), so they prove the
- # window is honoured, not merely that the tail happened to be small.
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="decode 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), the CUDA op",
- spec=dict(
- rows=8191,
- cols=_SCORE_COLUMNS,
- index_topk=_BLOCK_TOPK,
- windows=("chunk", 0),
- logits="relu",
- tail="nan",
- ),
- ),
- Cell(
- 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",
- ),
- ),
- )
-
- note = """
- Metadata consumed: none. Stateless; every input is an argument.
-
- Preconditions (the op checks none of these):
- - `0 <= row_starts[r] <= row_ends[r] <= logits.shape[1]`. An empty window
- is legal and yields an all-`-1` row; a `row_end` past the last column
- reads out of bounds (not probed).
- - Logits inside the window are finite. `-inf` is never selected, but a
- NaN ranks above every finite value and IS selected: NaN in columns 0-3
- of a 32-wide window with `k=8` returned all four (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, not pinned by a test): bf16/fp16 `logits`
- ("expected scalar type Float"), int64 `indices` ("expected scalar type
- Int"), `row_starts` / `row_ends` shorter than the row count ("one entry per
- row").
-
- The CUDA op dispatches by row index, not column count: rows below 12288 run
- an insertion-sort variant, rows at or above a radix variant, in one launch.
- A launch of `_CUTE_DSL_MIN_ROWS` (12288) rows or more with contiguous
- logits runs the CuTe DSL sibling instead -- QSA's stock indexer's choice on
- this GPU, measured to select identical sets (never identical permutations;
- ties, as in the CUDA op, broken arbitrarily). The tuned config's
- `max_num_tokens` 16384 reaches it at prefill-heavy mixed steps; smaller
- launches, every decode one among them, keep the compiled CUDA op, so a
- first captured decode JIT-compiles nothing. `torch.ops.trtllm.indexer_topk_decode` is the request-grouped
- sibling; QSA does not use it.
-
- CUDA graphs: captured at `[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 16384 x 512.
- """
-
- def __call__(
- self,
- logits: torch.Tensor,
- row_starts: torch.Tensor,
- row_ends: torch.Tensor,
- indices: torch.Tensor,
- index_topk: int = 2048,
- ) -> torch.Tensor:
- 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(
- self,
- logits: torch.Tensor,
- row_starts: torch.Tensor,
- row_ends: torch.Tensor,
- indices: torch.Tensor,
- index_topk: int = 2048,
- ) -> TopkExpectation:
- """Each window's top values, sorted. `indices` (the output buffer) is ignored."""
- num_rows, num_cols = logits.shape
- columns = torch.arange(num_cols, device=logits.device)[None, :]
- starts, ends = row_starts.long()[:, None], row_ends.long()[:, None]
- in_window = (columns >= starts) & (columns < ends)
- masked = logits.float().masked_fill(~in_window, float("-inf"))
- values = torch.sort(masked, dim=1, descending=True).values
- if num_cols < index_topk:
- values = torch.nn.functional.pad(
- values, (0, index_topk - num_cols), value=float("-inf")
- )
- window = (row_ends.long() - row_starts.long()).clamp_min(0)
- return TopkExpectation(values[:, :index_topk], masked, row_starts.long(), window)
-
- def is_valid(
- self,
- logits: torch.Tensor,
- row_starts: torch.Tensor,
- row_ends: torch.Tensor,
- indices: torch.Tensor,
- index_topk: int = 2048,
- ) -> None:
- # The op checks only `indices.size(1) >= index_topk`, then writes a flat
- # run of `index_topk`-wide rows and ignores the real row stride. Measured
- # with a [4, 12] buffer at k=8: row 1's answer landed across row 0's
- # columns 8-11 and its own 0-3, and the last row was never written.
- assert indices.dim() == 2 and indices.shape[1] == index_topk, (
- f"indices must be [num_rows, index_topk]; got {tuple(indices.shape)} "
- f"for index_topk={index_topk}"
- )
- assert indices.is_contiguous(), "indices must be contiguous"
-
-
-qsa_topk_select = QsaTopkSelect()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/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
deleted file mode 100644
index 4155e8a99b40..000000000000
--- 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
+++ /dev/null
@@ -1,1009 +0,0 @@
-# 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 program per row and v head stores v; one
-more program per row stores 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.
-
-Over an NVFP4 K/V pool the K/V append is `qsa_paged_kv_store_nvfp4`'s quantizing store:
-the programs holding k heads quantize their rotated, bf16-rounded rows from registers and
-the v programs quantize v from the projection row, with the store's arithmetic, so the
-pool's e2m1 bytes and e4m3 block scales are bitwise the composition's.
-
-A verify step (`rows_per_request` > 1, each request's rows consecutive at consecutive
-positions) writes several raw keys of one group in the launch: at a group's last token
-the earlier members held by the same request's preceding rows are taken from their
-projection rows, the others from the index plane.
-
-At up to `SCORE_ROWS` decode rows the index scores can ride the launch as well
-(`index_scores`): further programs run `qsa_paged_index_scores`' decode tiles, each
-rotating the row's index query itself and, at a group's last token, compressing the last
-visible block's key itself -- the same operations on the same values as the programs that
-store them, so the scores are bitwise the separate launch's.
-
-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, the K/V second-level scales, the
-compressed keys a score tile reads, and at a group end the group's earlier keys from the
-index plane with their rotation -- is read before `griddepcontrol.wait`; after it a
-program reads only projection rows (a score tile of a row without a page re-reads its
-last visible key, which the launch does not write).
-"""
-
-from __future__ import annotations
-
-from typing import Optional, Tuple
-
-import torch
-import triton
-import triton.language as tl
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
-
-from .qsa_paged_index_scores import qsa_paged_index_scores as _index_scores
-from .qsa_paged_index_store import qsa_paged_index_store as _index_store
-from .qsa_paged_kv_store import qsa_paged_kv_store as _kv_store
-from .qsa_paged_kv_store_nvfp4 import _e2m1x2, _quantize_row, _rcp_approx
-from .qsa_paged_kv_store_nvfp4 import qsa_paged_kv_store_nvfp4 as _kv_store_nvfp4
-from .qsa_prefill_compress import qsa_prefill_compress as _compress
-from .qwen4_exp_qk_norm_rope import _launch as _qk_launch
-from .qwen4_exp_qk_norm_rope import qwen4_exp_qk_norm_rope as _qk_norm_rope
-
-#: Rows up to which the index scores can ride the launch, as `qsa_paged_index_scores`' decode
-#: tiles (SCORE_BN blocks a program, two warps; the norm + rope plan is two warps there too).
-#: GB300, 12 chained layers at ~8K tokens, us a layer: 1 row 15.0 -> 13.7, 2 rows 15.8 -> 14.3;
-#: at 4 rows the 324 extra programs a row make it slower than the separate launch (17.1 -> 21.5).
-SCORE_ROWS = 2
-SCORE_BN = 8
-
-
-@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_tile(src, heads, head_mask, w2, wp2, c, s, eps, D: tl.constexpr, ROT: tl.constexpr,
- out_dtype: tl.constexpr): # 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 to `out_dtype`."""
- 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
- return (y * c[None, :] + yp * s[None, :]).to(out_dtype)
-
-
-@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
- """`_norm_rope_tile` stored at `dst`. Returns the stored (rounded) tile."""
- out = _norm_rope_tile(src, heads, head_mask, w2, wp2, c, s, eps, D, ROT, dst.dtype.element_ty)
- d = tl.arange(0, D)
- tl.store(dst + heads[:, None] * dst_head_stride + d[None, :], out, mask=head_mask[:, None])
- return out
-
-
-@triton.jit
-def _quantize_tile(x, global_scale, ROWS: tl.constexpr, D: tl.constexpr):
- """`qsa_paged_kv_store_nvfp4`'s `_quantize_row` over a `[ROWS, D]` register tile:
- (packed e2m1 bytes `[ROWS, D/2]`, e4m3 scale bytes `[ROWS, D/16]`), the same operations
- on the same values."""
- lo, hi = tl.split(tl.reshape(x.to(tl.float32), [ROWS, D // 2, 2]))
- pair_max = tl.maximum(tl.abs(lo), tl.abs(hi))
- block_max = tl.max(tl.reshape(pair_max, [ROWS, D // 16, 8]), axis=2)
- sf_value = global_scale * (block_max * _rcp_approx(tl.full([ROWS, D // 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, [ROWS, D // 16, 8]) * out_scale[:, :, None], [ROWS, D // 2]
- )
- hi_scaled = tl.reshape(
- tl.reshape(hi, [ROWS, D // 16, 8]) * out_scale[:, :, None], [ROWS, D // 2]
- )
- packed = (_e2m1x2(lo_scaled, hi_scaled) & 0xFF).to(tl.uint8)
- return packed, sf_code.to(tl.uint8, bitcast=True)
-
-
-@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,
- k_scale_cache,
- v_scale_cache,
- s_ks_page,
- s_ks_head,
- s_vs_page,
- s_vs_head,
- kv_scale_orig_quant,
- scores,
- s_sc_row,
- score_scale,
- 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,
- NVFP4: tl.constexpr,
- IN_LAUNCH_GROUPS: tl.constexpr,
- SCORE_TILES: tl.constexpr,
- SCORE_BLOCKS: tl.constexpr,
- SBN: tl.constexpr,
- SBH: tl.constexpr,
- USE_PDL: tl.constexpr,
-):
- token = tl.program_id(0).to(tl.int64)
- program = tl.program_id(1)
- if USE_PDL:
- # The consumer (the index scores, or with the scores here the top-k) reads
- # nothing of this launch before its 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 NVFP4:
- k_global = tl.load(kv_scale_orig_quant + 1)
- 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:
- is_k = (main_heads >= HQ) & (main_heads < HQ + HK) & live
- k_heads = (main_heads - HQ).to(tl.int64)[:, None]
- if NVFP4:
- # The rounded rows the composition's store reads back, quantized here.
- k_data, k_sf = _quantize_tile(out, k_global, HPB, D)
- pairs = tl.arange(0, D // 2)
- blocks = tl.arange(0, D // 16)
- tl.store(
- k_cache + page * s_kc_page + k_heads * s_kc_head + slot * s_kc_tok
- + pairs[None, :],
- k_data,
- mask=is_k[:, None],
- ) # fmt: skip
- # K scales token-major within the head page.
- tl.store(
- k_scale_cache + page * s_ks_page + k_heads * s_ks_head
- + slot * (D // 16) + blocks[None, :],
- k_sf,
- mask=is_k[:, None],
- ) # fmt: skip
- else:
- d = tl.arange(0, D)
- tl.store(
- k_cache + page * s_kc_page + k_heads * s_kc_head + slot * s_kc_tok
- + d[None, :],
- out.to(k_cache.dtype.element_ty),
- mask=is_k[:, None],
- ) # fmt: skip
- 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
- elif program >= MAIN_PROGRAMS + 2 + HK:
- if SCORE_TILES > 0:
- # `qsa_paged_index_scores`' decode tile (its prefetching kernel at up to four
- # rows: SBN blocks, the heads padded to SBH, two warps) with the row's index
- # query rotated here as the index program does, and at a group's last token the
- # last visible block's key compressed here as the index-key program does.
- tile = program - (MAIN_PROGRAMS + 2 + HK)
- visible_blocks = (logical + 1) // CR
- if tile * SBN >= visible_blocks:
- if USE_PDL:
- tl.extra.cuda.gdc_wait()
- return
- dims = tl.arange(0, ID)
- columns = tile * SBN + tl.arange(0, SBN)
- valid = (columns < SCORE_BLOCKS) & (columns < visible_blocks)
- anchors = columns * CR + CR - 1
- key_pages = tl.load(
- block_table + request * s_bt_req + (anchors // TPB) * s_bt_col, mask=valid, other=0
- ).to(tl.int64)
- key_ptrs = (
- index_cache
- + key_pages[:, None] * s_ic_page
- + (anchors % TPB)[:, None] * s_ic_slot
- + dims[None, :]
- )
- key_mask = valid[:, None] & (dims[None, :] < ID)
- keys = tl.load(key_ptrs, mask=key_mask, other=0.0, eviction_policy="evict_first")
- 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)
- # The block this launch compresses (the last visible one, at a group's last
- # token): the group's earlier raw keys, the first token's coordinates, their
- # rotation and the key norm weight, as the index-key program reads them.
- fresh_tile = ((logical + 1) % CR == 0) & ((visible_blocks - 1) // SBN == tile)
- fresh = (columns == visible_blocks - 1) & fresh_tile
- compress_here = fresh_tile & live
- 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=compress_here,
- 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=(compress_here & (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=compress_here,
- 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=compress_here, other=0)
- pos_h = tl.load(first_base + s_pc_axis, mask=compress_here, other=0)
- pos_w = tl.load(first_base + 2 * s_pc_axis, mask=compress_here, other=0)
- cosine, sine = _compress_factors(
- pos_t,
- pos_h,
- pos_w,
- cos_sin,
- stride_cs,
- ROTARY_PAIRS,
- IS_MROPE,
- MROPE_H,
- MROPE_W,
- compress_here,
- )
- key_weight = tl.load(index_k_weight + dims).to(tl.float32) + 1.0
- if USE_PDL:
- tl.extra.cuda.gdc_wait()
- query = _norm_rope_tile(
- row + IQ_COL, index_heads, index_heads < IH, w2, wp2, c, s, eps, ID, ROT,
- index_q.dtype.element_ty,
- ) # fmt: skip
- if fresh_tile:
- if compress_here:
- current_k = tl.load(row + IQ_COL + IH * ID + dims)
- group = tl.where(
- offsets[:, None] == CR - 1,
- current_k.to(index_cache.dtype.element_ty)[None, :],
- group,
- )
- pooled = (tl.sum(group.to(tl.float32), axis=0) / CR).to(tl.bfloat16)
- fresh_key = tl.reshape(
- _compress_norm_rope(
- pooled[None, :], key_weight, cosine, sine, eps, ID, ROTARY_PAIRS
- ),
- (ID,),
- )
- keys = tl.where(fresh[:, None], fresh_key[None, :].to(keys.dtype), keys)
- else:
- # No page: nothing is written; the key is the plane's.
- keys = tl.where(
- fresh[:, None],
- tl.load(key_ptrs, mask=key_mask & fresh[:, None], other=0.0),
- keys,
- )
- # The query as `[ID, SBH]`, heads past IH zero (the scores kernel's masked load).
- padded = tl.reshape(tl.broadcast_to(query[None, :, :], (SBH // IH, IH, ID)), (SBH, ID))
- head_rows = tl.arange(0, SBH)
- padded = tl.where((head_rows < IH)[:, None], padded, tl.zeros_like(padded))
- per_head = tl.dot(keys, tl.trans(padded), out_dtype=tl.float32)
- block_scores = tl.sum(
- tl.where(head_rows[None, :] < IH, tl.maximum(per_head, 0.0), 0.0), axis=1
- )
- tl.store(scores + token * s_sc_row + columns, block_scores * score_scale, mask=valid)
- elif program >= MAIN_PROGRAMS + 2:
- # One v head a program, as the K/V append stores it.
- h = program - MAIN_PROGRAMS - 2
- v_row = row + (HQ + HK + h) * D
- if NVFP4:
- v_global = tl.load(kv_scale_orig_quant + 2)
- if USE_PDL:
- tl.extra.cuda.gdc_wait()
- if live:
- if NVFP4:
- v_data, v_sf = _quantize_row(v_row, 1, v_global, D)
- pairs = tl.arange(0, D // 2)
- blocks = tl.arange(0, D // 16)
- tl.store(
- v_cache + page * s_vc_page + h * s_vc_head + slot * s_vc_tok + pairs, v_data
- )
- # V scales in four-token groups within the head page.
- tl.store(
- v_scale_cache + page * s_vs_page + h * s_vs_head
- + (slot // 4) * (4 * (D // 16)) + blocks * 4 + slot % 4,
- v_sf,
- ) # fmt: skip
- else:
- d = tl.arange(0, D)
- tl.store(
- v_cache + page * s_vc_page + h * s_vc_head + slot * s_vc_tok + d,
- tl.load(v_row + d).to(v_cache.dtype.element_ty),
- )
- else:
- # Before the wait: the coordinates; at a group's last token, the group's earlier
- # raw keys (written by earlier steps), the first token's 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)
- earlier = group_end & (offsets < CR - 1)
- if IN_LAUNCH_GROUPS:
- # A request's rows come consecutive at consecutive positions: the earlier
- # member at row `token - (CR - 1) + o`, when that row is this request's at
- # that position, is written by this launch -- its raw key is taken from its
- # projection row after the wait, not from the index plane, and as the first
- # member its coordinates from its row of `coords`.
- member_rows = token - (CR - 1) + offsets
- maybe = earlier & (member_rows >= 0)
- member_requests = tl.load(request_indices + member_rows, mask=maybe, other=-1)
- member_logical = tl.load(logical_positions + member_rows, mask=maybe, other=-1)
- in_launch = (
- maybe
- & (member_requests.to(tl.int64) == request)
- & (member_logical.to(tl.int64) == members)
- )
- else:
- in_launch = offsets < 0
- group = tl.load(
- index_cache
- + member_pages[:, None] * s_ic_page
- + (tl.maximum(members, 0) % TPB)[:, None] * s_ic_slot
- + dims[None, :],
- mask=(earlier & ~in_launch)[:, None],
- other=0.0,
- )
- first = tl.maximum(logical - (CR - 1), 0)
- first_page = tl.load(
- block_table + request * s_bt_req + (first // TPB) * s_bt_col, mask=group_end, other=0
- ).to(tl.int64)
- first_base = position_cache + first_page * s_pc_page + (first % TPB) * s_pc_slot
- if IN_LAUNCH_GROUPS:
- # A first member written by this launch has its coordinates in its row.
- first_in_launch = tl.sum(tl.where(offsets == 0, in_launch, False).to(tl.int32)) > 0
- cached = group_end & ~first_in_launch
- first_row = coords + (token - (CR - 1)) * s_c_row
- pos_t = tl.load(first_row, mask=first_in_launch, other=0)
- pos_h = tl.load(first_row + s_c_axis, mask=first_in_launch, other=0)
- pos_w = tl.load(first_row + 2 * s_c_axis, mask=first_in_launch, other=0)
- pos_t = tl.where(cached, tl.load(first_base, mask=cached, other=0), pos_t)
- pos_h = tl.where(cached, tl.load(first_base + s_pc_axis, mask=cached, other=0), pos_h)
- pos_w = tl.where(
- cached, tl.load(first_base + 2 * s_pc_axis, mask=cached, other=0), pos_w
- )
- else:
- 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:
- # The coordinates, as the index store writes them.
- tl.store(
- position_cache + page * s_pc_page + slot * s_pc_slot + axes * s_pc_axis,
- coordinate.to(position_cache.dtype.element_ty),
- mask=axes < 3,
- )
- # The raw index key -- or, at a group's last token, the group compressed in
- # its place (the compression kernel overwrites the same slot).
- current_k = tl.load(row + IQ_COL + IH * ID + dims)
- stored_k = current_k.to(index_cache.dtype.element_ty)
- if (logical + 1) % CR == 0:
- if IN_LAUNCH_GROUPS:
- launch_k = tl.load(
- proj
- + member_rows[:, None] * stride_proj
- + IQ_COL
- + IH * ID
- + dims[None, :],
- mask=in_launch[:, None],
- other=0.0,
- )
- group = tl.where(
- in_launch[:, None], launch_k.to(index_cache.dtype.element_ty), group
- )
- group = tl.where(offsets[:, None] == CR - 1, stored_k[None, :], group)
- pooled = (tl.sum(group.to(tl.float32), axis=0) / CR).to(tl.bfloat16)
- stored_k = tl.reshape(
- _compress_norm_rope(
- pooled[None, :], key_weight, cosine, sine, eps, ID, ROTARY_PAIRS
- ),
- (ID,),
- )
- tl.store(index_cache + page * s_ic_page + slot * s_ic_slot + dims, stored_k)
-
-
-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)`.
-
- Over an NVFP4 pool (`k_scale_cache`, `v_scale_cache` and `kv_scale_orig_quant` given,
- `k_cache` / `v_cache` its uint8 e2m1 planes) the K/V append is
- `qsa_paged_kv_store_nvfp4`'s: e2m1 bytes and e4m3 block scales, K's token-major and
- V's in four-token groups, at the second-level scales `kv_scale_orig_quant[1]` / `[2]`.
- `rows_per_request` > 1 declares a verify step's layout: each request's rows
- consecutive, at consecutive positions.
-
- `index_scores` (fp32 `[T, table width * tpb / compress ratio]`, up to `SCORE_ROWS`
- decode rows) takes `qsa_paged_index_scores`' output over the index plane and query
- this launch leaves (`only_visible_blocks`: each row's visible columns written, the
- rest untouched): more programs score `SCORE_BN` blocks each, rotating the row's index
- query themselves and, at a group's last token, compressing the last visible block's
- key themselves, as the programs writing them do.
-
- Fusion boundary: as the four entries (with `index_scores`, and the index scores); the
- projection before and the selection after are the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(why="decode, 1 row at position 8699 (a group's last token)", spec=dict(pos=(8699,))),
- Cell(
- why="decode, 4 rows: three group ends (two on a page's last slot), one with no compression",
- spec=dict(pos=(8699, 8700, 4095, 31)),
- ),
- Cell(
- why="decode, 2 rows, one unmapped (request out of range)",
- spec=dict(pos=(8699, 1201), unmapped=1),
- ),
- Cell(
- why="NVFP4 pool, 1 row at position 8699 (a group's last token), K/V the role "
- "planes of the engine's [pages, 2, Hkv, tpb, D/2] pool",
- spec=dict(pos=(8699,), nvfp4=True),
- ),
- Cell(
- why="NVFP4 pool, 4 rows: three group ends (two on a page's last slot), one with "
- "no compression, with zero, -0.0 and outlier blocks in k and v",
- spec=dict(pos=(8699, 8700, 4095, 31), nvfp4=True, edge=True),
- ),
- Cell(
- why="NVFP4 pool, 2 rows, one unmapped (request out of range)",
- spec=dict(pos=(8699, 1201), unmapped=1, nvfp4=True),
- ),
- Cell(
- why="NVFP4 pool, 40 rows (two heads a program): positions across groups and pages",
- spec=dict(pos=tuple(2000 + 151 * r for r in range(40)), nvfp4=True),
- ),
- Cell(
- why="NVFP4 pool, 72 rows (eight heads a program, k heads beside q heads)",
- spec=dict(pos=tuple(1000 + 113 * r for r in range(72)), nvfp4=True),
- ),
- Cell(
- why="NVFP4 pool, 1 row at a group's last token, the index scores in the launch",
- spec=dict(pos=(8699,), nvfp4=True, scores=True),
- ),
- Cell(
- why="NVFP4 pool, 2 rows (a group end on a page's last slot, one row without), the "
- "index scores in the launch",
- spec=dict(pos=(4095, 8700), nvfp4=True, scores=True, edge=True),
- ),
- Cell(
- why="bf16 pool, 2 rows (a group end, and one at 8 tokens), the index scores in the "
- "launch",
- spec=dict(pos=(6143, 7), scores=True),
- ),
- Cell(
- why="verify, NVFP4 pool, 3 requests x 4 rows: group ends at a request's third "
- "row (two earlier members in the launch), fourth (three) and first (none)",
- spec=dict(pos=(8697, 8696, 8699), rows_per_request=4, nvfp4=True, edge=True),
- ),
- Cell(
- why="verify, bf16 pool, 2 requests x 4 rows, one group end crossing a page",
- spec=dict(pos=(4093, 8698), rows_per_request=4),
- ),
- Cell(
- why="verify, NVFP4 pool, 32 requests x 4 rows (eight heads a program)",
- spec=dict(pos=tuple(3000 + 211 * r for r in range(32)), rows_per_request=4, nvfp4=True),
- ),
- )
-
- 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]`; generation
- rows: one per request (`rows_per_request` 1: a group's earlier tokens are already in
- the index plane), or a verify step's `rows_per_request` consecutive rows a request at
- consecutive positions (a group's earlier tokens in the index plane or in the
- request's preceding rows). A row whose page lookup fails writes no cache, as the
- stores and compression skip it.
- Over an NVFP4 pool: the preconditions of `qsa_paged_kv_store_nvfp4` (uint8
- `[pages, Hkv, tpb, D/2]` data and `[pages, Hkv, tpb, D/16]` scale planes, each scale
- head page contiguous, a contiguous fp32 `[3]` `kv_scale_orig_quant`); the page
- bound is the data plane's page count, as there.
-
- Numerics: the same arithmetic as the four entries; outputs bitwise theirs (the
- NVFP4 cells compare every byte of the four K/V planes).
-
- Compiled variants: one per row-count band of the norm + rope launch plan (up to 16,
- up to 64, more rows), pool kind (bf16 / NVFP4) and step kind (one row a request /
- verify); the row count, table height and width and the page count are not
- specialized on.
- """
-
- 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,
- k_scale_cache: Optional[torch.Tensor] = None,
- v_scale_cache: Optional[torch.Tensor] = None,
- kv_scale_orig_quant: Optional[torch.Tensor] = None,
- rows_per_request: int = 1,
- index_scores: Optional[torch.Tensor] = 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)
- nvfp4 = k_scale_cache is not None
- if not nvfp4:
- # Placeholders: the bf16 kernel does not read them.
- k_scale_cache = v_scale_cache = kv_scale_orig_quant = k_cache
- score_tiles = score_blocks = 0
- if index_scores is not None:
- assert num_tokens <= SCORE_ROWS and rows_per_request == 1 and warps == 2, (
- "the index scores ride decode launches of up to SCORE_ROWS rows"
- )
- score_blocks = block_table.shape[1] * tokens_per_block // compress_ratio
- score_tiles = triton.cdiv(score_blocks, SCORE_BN)
- pdl = get_env_enable_pdl()
- _attn_decode_prologue_kernel[(num_tokens, main_programs + 2 + num_heads_k + score_tiles)](
- 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),
- k_scale_cache,
- v_scale_cache,
- k_scale_cache.stride(0),
- k_scale_cache.stride(1),
- v_scale_cache.stride(0),
- v_scale_cache.stride(1),
- kv_scale_orig_quant,
- proj if index_scores is None else index_scores,
- 0 if index_scores is None else index_scores.stride(0),
- index_head_dim**-0.5,
- 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],
- NVFP4=nvfp4,
- IN_LAUNCH_GROUPS=rows_per_request > 1,
- SCORE_TILES=score_tiles,
- SCORE_BLOCKS=score_blocks,
- SBN=SCORE_BN,
- SBH=max(16, triton.next_power_of_2(index_heads)),
- 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,
- k_scale_cache: Optional[torch.Tensor] = None,
- v_scale_cache: Optional[torch.Tensor] = None,
- kv_scale_orig_quant: Optional[torch.Tensor] = None,
- rows_per_request: int = 1,
- index_scores: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
- """The four entries in their composition order, on the given buffers; with
- `index_scores`, then `qsa_paged_index_scores`' visible columns written into it."""
- 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]
- if k_scale_cache is None:
- _kv_store(
- k,
- v,
- k_cache,
- v_cache,
- request_indices,
- logical_positions,
- block_table,
- tokens_per_block,
- )
- else:
- _kv_store_nvfp4(
- k,
- v,
- k_cache,
- v_cache,
- k_scale_cache,
- v_scale_cache,
- request_indices,
- logical_positions,
- block_table,
- tokens_per_block,
- kv_scale_orig_quant,
- )
- 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
- if index_scores is not None:
- scores = _index_scores(
- q=index_q, index_cache=index_cache, block_table=block_table,
- query_positions=logical_positions, request_indices=request_indices,
- tokens_per_block=tokens_per_block, compress_ratio=compress_ratio,
- only_visible_blocks=True, context_rows=False,
- ) # fmt: skip
- for r in range(rows):
- visible = (int(logical_positions[r]) + 1) // compress_ratio
- index_scores[r, :visible] = scores[r, :visible]
- 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)"
- )
- k_scale_cache = kwargs.get("k_scale_cache")
- v_scale_cache = kwargs.get("v_scale_cache")
- global_scales = kwargs.get("kv_scale_orig_quant")
- nvfp4 = (k_scale_cache, v_scale_cache, global_scales)
- assert all(t is None for t in nvfp4) or all(t is not None for t in nvfp4), (
- "an NVFP4 pool takes both scale planes and kv_scale_orig_quant"
- )
- scores = kwargs.get("index_scores")
- if scores is not None:
- assert scores.dtype == torch.float32 and scores.stride(1) == 1, "fp32 index scores"
- assert proj.shape[0] <= SCORE_ROWS and kwargs.get("rows_per_request", 1) == 1, (
- "the index scores ride decode launches of up to SCORE_ROWS rows"
- )
- if k_scale_cache is not None:
- for scales in (k_scale_cache, v_scale_cache):
- assert scales.dtype == torch.uint8 and scales.stride(3) == 1, "uint8 scale planes"
- assert scales.stride(2) == scales.shape[3], "each scale head page contiguous"
- assert (
- global_scales.dtype == torch.float32
- and global_scales.shape == (3,)
- and global_scales.is_contiguous()
- ), "a contiguous fp32 [3] kv_scale_orig_quant"
-
-
-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
deleted file mode 100644
index 27a17ed46ace..000000000000
--- 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
+++ /dev/null
@@ -1,1193 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Gemma QK norm + partial NeoX rope for both of an attention layer's planes, one launch.
-
-The main plane (q and k heads of the fused input projection) is rewritten in
-place; the QSA index plane's query heads are read from the same projection row
-and written, normalized and rotated, into a separate buffer. The rotation reads
-the target's own fp32 `[positions, rotary_dim]` cos|sin table.
-
-Two kernels, one result: a Triton kernel, and from `_CUTE_MIN_TOKENS` tokens (a mixed
-step's prefill chunk) a persistent CuTe DSL kernel that is bitwise the Triton kernel's
-eight-head launch and streams the projection at about twice its rate.
-"""
-
-from __future__ import annotations
-
-import struct
-from typing import Dict, Tuple
-
-import cutlass
-import cutlass.cute as cute
-import torch
-import triton
-import triton.language as tl
-from cutlass._mlir.dialects import llvm
-from cutlass.cute.runtime import from_dlpack
-from cutlass.cutlass_dsl import T, dsl_user_op
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
-
-
-@triton.jit
-def _norm_rope_tile(
- src,
- dst,
- heads,
- head_mask,
- weight,
- positions,
- num_tokens,
- token,
- cos_sin,
- stride_cs,
- eps,
- D: tl.constexpr,
- ROT: tl.constexpr,
- S1: tl.constexpr,
- S2: tl.constexpr,
- WEIGHT_SPLIT: tl.constexpr,
- weight_k,
- dst_head_stride,
-):
- """One `[len(heads), D]` tile: rms over each head, `* (1 + w)`, rotate the
- first ROT dims in NeoX halves. fp32 throughout, one rounding at the store."""
- HALF: tl.constexpr = ROT // 2
- d = tl.arange(0, D)
- rot = d < ROT
- partner = tl.where(d < HALF, d + HALF, tl.where(rot, d - HALF, d))
- hm = head_mask[:, None]
- x = tl.load(src + heads[:, None] * D + d[None, :], mask=hm, other=0.0).to(tl.float32)
- xp = tl.load(src + heads[:, None] * D + partner[None, :], mask=hm & rot[None, :], other=0.0).to(
- tl.float32
- )
- w = tl.load(weight + d).to(tl.float32)
- wp = tl.load(weight + partner).to(tl.float32)
- if WEIGHT_SPLIT > 0:
- # Heads at or past WEIGHT_SPLIT take the second weight (the k heads).
- is_k = (heads >= WEIGHT_SPLIT)[:, None]
- wk = tl.load(weight_k + d).to(tl.float32)
- wkp = tl.load(weight_k + partner).to(tl.float32)
- w2 = tl.where(is_k, wk[None, :], w[None, :]) + 1.0
- wp2 = tl.where(is_k, wkp[None, :], wp[None, :]) + 1.0
- else:
- w2 = w[None, :] + 1.0
- wp2 = wp[None, :] + 1.0
- rstd = tl.rsqrt(tl.sum(x * x, axis=1) / D + eps)
- y = x * rstd[:, None] * w2
- yp = xp * rstd[:, None] * wp2
- # Interleaved mRoPE: pair j reads coordinate 1 when j % 3 == 1 and
- # j < 3 * S1, coordinate 2 when j % 3 == 2 and j < 3 * S2, else 0.
- j = d % HALF
- p0 = tl.load(positions + token)
- p1 = tl.load(positions + num_tokens + token)
- p2 = tl.load(positions + 2 * num_tokens + token)
- pos = tl.where((j % 3 == 1) & (j < 3 * S1), p1, tl.where((j % 3 == 2) & (j < 3 * S2), p2, p0))
- row = cos_sin + pos.to(tl.int64) * stride_cs
- c = tl.load(row + j, mask=rot, other=1.0)
- s = tl.load(row + HALF + j, mask=rot, other=0.0)
- s = tl.where(d < HALF, -s, s)
- out = y * c[None, :] + yp * s[None, :]
- tl.store(
- dst + heads[:, None] * dst_head_stride + d[None, :],
- out.to(dst.dtype.element_ty),
- mask=hm,
- )
-
-
-@triton.jit(do_not_specialize=["num_tokens"])
-def _qwen4_exp_qk_norm_rope_kernel(
- proj,
- stride_proj,
- index_q,
- stride_iq_token,
- stride_iq_head,
- positions,
- num_tokens,
- cos_sin,
- stride_cs,
- q_weight,
- k_weight,
- index_q_weight,
- eps,
- HQ: tl.constexpr,
- HK: tl.constexpr,
- D: tl.constexpr,
- IH: tl.constexpr,
- ID: tl.constexpr,
- IQ_COL: tl.constexpr,
- ROT: tl.constexpr,
- S1: tl.constexpr,
- S2: tl.constexpr,
- IS1: tl.constexpr,
- IS2: tl.constexpr,
- HPB: tl.constexpr,
- MAIN_PROGRAMS: tl.constexpr,
-):
- token = tl.program_id(0).to(tl.int64)
- program = tl.program_id(1)
- row = proj + token * stride_proj
- if program < MAIN_PROGRAMS:
- main_heads = program * HPB + tl.arange(0, HPB)
- _norm_rope_tile(
- row, row, main_heads, main_heads < HQ + HK, q_weight, positions, num_tokens, token,
- cos_sin, stride_cs, eps, D, ROT, S1, S2, HQ, k_weight, D,
- ) # fmt: skip
- else:
- index_heads = tl.arange(0, IH)
- _norm_rope_tile(
- row + IQ_COL, index_q + token * stride_iq_token, index_heads, index_heads < IH,
- index_q_weight, positions, num_tokens, token, cos_sin, stride_cs, eps,
- ID, ROT, IS1, IS2, 0, index_q_weight, stride_iq_head,
- ) # fmt: skip
-
-
-def _launch(num_tokens: int) -> Tuple[int, int]:
- """Main-plane heads per program and warps per program, by token count.
-
- Measured on GB300 (CUDA-graph replay, 12 rotating layers, us per call;
- HPB/warps):
-
- tokens 1/2 2/2 2/4 4/4 8/8 32/8
- 1 1.76 1.84 1.87 1.99 1.92 2.82
- 8 1.83 1.97 1.97 2.02 2.05 2.91
- 32 2.11 2.14 2.18 2.14 2.16 2.99
- 64 2.48 2.39 2.47 2.44 2.48 3.03
- 128 3.24 3.19 3.46 3.12 2.96 3.12
- 2048 29.7 25.4 30.8 24.0 22.6 25.8
- 8192 115 98.5 121 95.7 89.5 122
- """
- if num_tokens <= 16:
- return 1, 2
- if num_tokens <= 64:
- return 2, 2
- return 8, 8
-
-
-#: This checkpoint's geometry: 24 q + 2 k heads of 256 at the head of the
-#: 14080-wide fused projection row, 4 index q heads of 128 at column 13312, a
-#: 64-dim rotary prefix, `mrope_section` [11, 11, 10] on the main plane and
-#: coordinate 0 alone on the index plane.
-CKPT = dict(
- num_heads_q=24,
- num_heads_k=2,
- head_dim=256,
- index_heads=4,
- index_head_dim=128,
- index_q_col=13312,
- proj_width=14080,
- index_scratch_heads=6,
- rotary_dim=64,
- mrope_section1=11,
- mrope_section2=10,
- index_mrope_section1=0,
- index_mrope_section2=0,
-)
-
-
-# ---- CuTe DSL path: prefill-size token counts -------------------------------------------
-
-#: Token count from which the op runs the CuTe kernel. GB300, graph replay, us per call (Triton
-#: -> CuTe): 512 8.4 -> 10.3, 768 12.3 -> 8.2, 1024 14.5 -> 8.3, 4096 45.2 -> 18.5, 16384
-#: 182 -> 82.7. The Triton kernel runs its eight-head launch from 65 tokens on, and the CuTe
-#: kernel is bitwise that launch.
-_CUTE_MIN_TOKENS = 1024
-#: A token's stage in a warp's ring: the lanes' four chunks (32 x 16 bytes each), the cos and
-#: sin of the 32 rotary pairs, and three positions.
-_STAGE = 4 * 32 * 16 + 2 * 32 * 4 + 16
-_STAGE_CS = 4 * 32 * 16
-_STAGE_POS = _STAGE_CS + 2 * 32 * 4
-#: Stages in a warp's ring (the next token's loads in flight while one is formed).
-_RING = 2
-_CTAS_PER_SM = 3
-
-
-def _iv(x, loc, ip):
- return x.ir_value(loc=loc, ip=ip)
-
-
-def _asm(ret, args, text, constraints, loc, ip, side_effects=False):
- return llvm.inline_asm(
- ret,
- [_iv(a, loc, ip) for a in args],
- text,
- constraints,
- has_side_effects=side_effects,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
-
-
-def _asm4(args, text, constraints, loc, ip, side_effects=False):
- ty = T.i32()
- res = _asm(
- llvm.StructType.get_literal([ty] * 4), args, text, constraints, loc, ip, side_effects
- )
- return [cutlass.Int32(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(4)]
-
-
-@dsl_user_op
-def _ldg_v4(addr, *, loc=None, ip=None):
- return _asm4(
- [cutlass.Int64(addr)],
- "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];",
- "=r,=r,=r,=r,l",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _ldg_b32(addr, *, loc=None, ip=None):
- return cutlass.Int32(
- _asm(T.i32(), [cutlass.Int64(addr)], "ld.global.b32 $0, [$1];", "=r,l", loc, ip, True)
- )
-
-
-@dsl_user_op
-def _stg_v4(addr, off: int, a, b, c, d, *, loc=None, ip=None):
- """16 bytes at `addr + off`."""
- _asm(
- T.i32(),
- [cutlass.Int64(addr)] + [cutlass.Int32(v) for v in (a, b, c, d)],
- "{ st.global.v4.b32 [$1+%d], {$2, $3, $4, $5}; mov.u32 $0, 0; }" % off,
- "=r,l,r,r,r,r",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _lds_v4(addr, *, loc=None, ip=None):
- return _asm4(
- [cutlass.Int32(addr)],
- "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];",
- "=r,=r,=r,=r,r",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _lds_b32(addr, *, loc=None, ip=None):
- return cutlass.Int32(
- _asm(T.i32(), [cutlass.Int32(addr)], "ld.shared.b32 $0, [$1];", "=r,r", loc, ip, True)
- )
-
-
-@dsl_user_op
-def _cp16(dst, src, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [cutlass.Int32(dst), cutlass.Int64(src)],
- "{ cp.async.cg.shared.global [$1], [$2], 16; mov.u32 $0, 0; }",
- "=r,r,l",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _cp4(dst, src, pred, *, loc=None, ip=None):
- """4 bytes where `pred` is not 0."""
- _asm(
- T.i32(),
- [cutlass.Int32(dst), cutlass.Int64(src), cutlass.Int32(pred)],
- "{ .reg .pred p; setp.ne.b32 p, $3, 0; @p cp.async.ca.shared.global [$1], [$2], 4; "
- "mov.u32 $0, 0; }",
- "=r,r,l,r",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _cp_commit(*, loc=None, ip=None):
- _asm(T.i32(), [], "{ cp.async.commit_group; mov.u32 $0, 0; }", "=r", loc, ip, True)
-
-
-@dsl_user_op
-def _cp_wait(n: int, *, loc=None, ip=None):
- _asm(T.i32(), [], "{ cp.async.wait_group %d; mov.u32 $0, 0; }" % n, "=r", loc, ip, True)
-
-
-@dsl_user_op
-def _bf16_lo(w, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Int32(w)],
- "{ .reg .b32 t; shl.b32 t, $1, 16; mov.b32 $0, t; }",
- "=f,r",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _bf16_hi(w, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Int32(w)],
- "{ .reg .b32 t; and.b32 t, $1, 0xFFFF0000; mov.b32 $0, t; }",
- "=f,r",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _unpack2(w, *, loc=None, ip=None):
- """A bf16x2 word -> its two values as an fp32x2 pair, the low element first."""
- return cutlass.Int64(
- _asm(
- T.i64(),
- [cutlass.Int32(w)],
- "{ .reg .b32 a, b; shl.b32 a, $1, 16; and.b32 b, $1, 0xFFFF0000; mov.b64 $0, {a, b}; }",
- "=l,r",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _pair(lo, hi, *, loc=None, ip=None):
- """Two fp32 bit patterns -> one fp32x2 pair."""
- return cutlass.Int64(
- _asm(
- T.i64(),
- [cutlass.Int32(lo), cutlass.Int32(hi)],
- "mov.b64 $0, {$1, $2};",
- "=l,r,r",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _splat2(x, *, loc=None, ip=None):
- return cutlass.Int64(
- _asm(T.i64(), [cutlass.Float32(x)], "mov.b64 $0, {$1, $1};", "=l,f", loc, ip)
- )
-
-
-@dsl_user_op
-def _mul2(a, b, *, loc=None, ip=None):
- return cutlass.Int64(
- _asm(
- T.i64(),
- [cutlass.Int64(a), cutlass.Int64(b)],
- "mul.rn.f32x2 $0, $1, $2;",
- "=l,l,l",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _fma2(a, b, c, *, loc=None, ip=None):
- return cutlass.Int64(
- _asm(
- T.i64(),
- [cutlass.Int64(a), cutlass.Int64(b), cutlass.Int64(c)],
- "fma.rn.f32x2 $0, $1, $2, $3;",
- "=l,l,l,l",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _add2(a, b, *, loc=None, ip=None):
- return cutlass.Int64(
- _asm(
- T.i64(),
- [cutlass.Int64(a), cutlass.Int64(b)],
- "add.rn.f32x2 $0, $1, $2;",
- "=l,l,l",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _hsum2(a, *, loc=None, ip=None):
- """The low plus the high element of a pair."""
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Int64(a)],
- "{ .reg .f32 x, y; mov.b64 {x, y}, $1; add.rn.f32 $0, x, y; }",
- "=f,l",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _bf16x2_of(a, *, loc=None, ip=None):
- """A pair rounded (RN) into one bf16x2 word, the low element low."""
- return cutlass.Int32(
- _asm(
- T.i32(),
- [cutlass.Int64(a)],
- "{ .reg .f32 x, y; mov.b64 {x, y}, $1; cvt.rn.bf16x2.f32 $0, y, x; }",
- "=r,l",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _neg_bits(w, *, loc=None, ip=None):
- """An fp32 bit pattern with its sign flipped (the Triton kernel's `xor`)."""
- return cutlass.Int32(
- _asm(T.i32(), [cutlass.Int32(w)], "xor.b32 $0, $1, 0x80000000;", "=r,r", loc, ip)
- )
-
-
-@dsl_user_op
-def _fma(a, b, c, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)],
- "fma.rn.f32 $0, $1, $2, $3;",
- "=f,f,f,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _mul(a, b, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Float32(a), cutlass.Float32(b)],
- "mul.rn.f32 $0, $1, $2;",
- "=f,f,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _add(a, b, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Float32(a), cutlass.Float32(b)],
- "add.rn.f32 $0, $1, $2;",
- "=f,f,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _rstd(ss, eps, width_hex: str, *, loc=None, ip=None):
- """`rsqrt(ss / width + eps)` as the Triton kernel's `div.full`, `add`, `rsqrt.approx.ftz`:
- the `add` has no rounding mode there, so ptxas fuses it with the division (a multiply by
- the reciprocal) into one FFMA -- here the same."""
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Float32(ss), cutlass.Float32(eps)],
- "{ .reg .f32 t; div.full.f32 t, $1, " + width_hex + "; add.f32 t, t, $2; "
- "rsqrt.approx.ftz.f32 $0, t; }",
- "=f,f,f",
- loc,
- ip,
- )
- )
-
-
-def _hex(x: float) -> str:
- return "0f%08X" % struct.unpack(" bool:
- """The geometry and layouts the CuTe kernel covers; anything else runs the Triton kernel."""
- bf = torch.bfloat16
- return (
- num_heads_q % 4 == 0
- and num_heads_k == 2
- and head_dim == 256
- and index_heads == 4
- and index_head_dim == 128
- and rotary_dim == 64
- and index_q_col % 8 == 0
- and proj.dtype == index_q.dtype == bf
- and all(w.dtype == bf and w.is_contiguous() and w.data_ptr() % 16 == 0 for w in weights)
- and cos_sin.dtype == torch.float32
- and cos_sin.stride(1) == 1
- and proj.data_ptr() % 16 == 0
- and proj.stride(0) % 8 == 0
- and index_q.data_ptr() % 16 == 0
- and index_q.stride(0) % 8 == 0
- and index_q.stride(1) == index_head_dim
- )
-
-
-def _launch_cute(
- proj: torch.Tensor,
- index_q: torch.Tensor,
- position_ids: torch.Tensor,
- cos_sin: torch.Tensor,
- q_weight: torch.Tensor,
- k_weight: torch.Tensor,
- index_q_weight: torch.Tensor,
- num_heads_q: int,
- index_q_col: int,
- eps: float,
- sections: Tuple[int, int, int, int],
-) -> None:
- key = (num_heads_q, index_q_col) + tuple(sections)
- fn = _RING_KERNELS.get(key)
- if fn is None:
- gm = cute.AddressSpace.gmem
- bf = cutlass.BFloat16
- fn = cute.compile(
- _NormRopeRing(num_heads_q, index_q_col, *sections),
- from_dlpack(proj, assumed_align=16).mark_layout_dynamic(leading_dim=1),
- make_ptr(bf, index_q.data_ptr(), gm, assumed_align=16),
- make_ptr(cutlass.Int32, position_ids.data_ptr(), gm, assumed_align=4),
- make_ptr(cutlass.Float32, cos_sin.data_ptr(), gm, assumed_align=4),
- make_ptr(bf, q_weight.data_ptr(), gm, assumed_align=16),
- make_ptr(bf, k_weight.data_ptr(), gm, assumed_align=16),
- make_ptr(bf, index_q_weight.data_ptr(), gm, assumed_align=16),
- *(cutlass.Int32(1),) * 4,
- cutlass.Float32(1e-6),
- cutlass.Int32(1),
- cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
- options="--opt-level 2 --enable-tvm-ffi",
- )
- _RING_KERNELS[key] = fn
- dev = proj.device.index
- if dev not in _SM_COUNT:
- _SM_COUNT[dev] = torch.cuda.get_device_properties(proj.device).multi_processor_count
- tokens = proj.shape[0]
- fn(
- proj,
- index_q.data_ptr(),
- position_ids.data_ptr(),
- cos_sin.data_ptr(),
- q_weight.data_ptr(),
- k_weight.data_ptr(),
- index_q_weight.data_ptr(),
- tokens,
- proj.stride(0),
- index_q.stride(0),
- cos_sin.stride(0),
- eps,
- min(tokens, _SM_COUNT[dev] * _CTAS_PER_SM),
- )
-
-
-def _launch_triton(
- proj: torch.Tensor,
- index_q: torch.Tensor,
- position_ids: torch.Tensor,
- cos_sin: torch.Tensor,
- q_weight: torch.Tensor,
- k_weight: torch.Tensor,
- index_q_weight: torch.Tensor,
- num_heads_q: int,
- num_heads_k: int,
- head_dim: int,
- index_heads: int,
- index_head_dim: int,
- index_q_col: int,
- rotary_dim: int,
- eps: float,
- sections: Tuple[int, int, int, int],
-) -> None:
- num_tokens = proj.shape[0]
- hpb, warps = _launch(num_tokens)
- main_programs = triton.cdiv(num_heads_q + num_heads_k, hpb)
- _qwen4_exp_qk_norm_rope_kernel[(num_tokens, main_programs + 1)](
- proj,
- proj.stride(0),
- index_q,
- index_q.stride(0),
- index_q.stride(1),
- position_ids,
- num_tokens,
- cos_sin,
- cos_sin.stride(0),
- q_weight,
- k_weight,
- index_q_weight,
- eps,
- HQ=num_heads_q,
- HK=num_heads_k,
- D=head_dim,
- IH=index_heads,
- ID=index_head_dim,
- IQ_COL=index_q_col,
- ROT=rotary_dim,
- S1=sections[0],
- S2=sections[1],
- IS1=sections[2],
- IS2=sections[3],
- HPB=hpb,
- MAIN_PROGRAMS=main_programs,
- num_warps=warps,
- )
-
-
-class Qwen4ExpQkNormRope(OpWrapper):
- """Both planes of one attention layer's QK norm + rope, in one launch.
-
- Main plane, in place on `proj`: heads `[0, Hq)` (weight `q_weight`) and
- `[Hq, Hq + Hk)` (weight `k_weight`), each `head_dim` wide at the start of
- every row, become `rope(rmsnorm(x) * (1 + w))`. Index plane: the
- `index_heads` heads of `index_head_dim` starting at column `index_q_col` of
- `proj` are read, normalized with `index_q_weight` and rotated, and written
- to `index_q` (`[T, index_heads, index_head_dim]`, any token/head stride,
- unit dim stride); `proj`'s index columns are not written. Every other
- column of `proj` is left alone. Returns `(proj, index_q)`.
-
- The rotation is NeoX over the first `rotary_dim` dims of a head: pair `j`
- is `(j, j + rotary_dim / 2)`, rotated by `cos_sin[p, j]` and
- `cos_sin[p, rotary_dim / 2 + j]` (`cos_sin` is the fp32 `[P, rotary_dim]`
- table, cos block then sin block). `position_ids` is the contiguous int32
- `[3, T]`; pair `j` reads coordinate 1 when `j % 3 == 1 and j < 3 * s1`,
- coordinate 2 when `j % 3 == 2 and j < 3 * s2`, coordinate 0 otherwise --
- interleaved mRoPE with the plane's own `(s1, s2)`; `(0, 0)` is the plain
- rope on coordinate 0. fp32 throughout, one rounding at the store.
-
- Fusion boundary: norm and rotation of both planes. The projection GEMM
- before it, and the KV/index-cache writes and attention after it, are the
- caller's. The index plane's raw key is not touched, so the index-cache
- store may run before or after this call.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # Four cells, one per compiled launch the forward reaches: the Triton kernel's 1-, 2-
- # and 8-head launches (by decode batch size) and the CuTe kernel (prefill).
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="decode M=1: one token on the 1-head/2-warp launch, three distinct mRoPE rows",
- spec=dict(tokens=1, **CKPT),
- ),
- Cell(
- why="decode M=17: first on the 2-head launch (13 main programs)",
- spec=dict(tokens=17, **CKPT),
- ),
- Cell(
- why="decode M=128: max concurrency, 8-head launch (last main program half masked), "
- "positions up to the table end (262144 rows)",
- spec=dict(tokens=128, max_position=262144, **CKPT),
- ),
- Cell(why="prefill 16384 tokens: max_num_tokens", spec=dict(tokens=16384, **CKPT)),
- )
-
- note = """
- Metadata consumed: none (positions are an operand). Below `_CUTE_MIN_TOKENS`
- tokens one Triton kernel, no autotuner: grid `(T, cdiv(Hq + Hk, HPB) + 1)`;
- the last program of a token row runs the whole index plane as one
- `[index_heads, index_head_dim]` tile. HPB and warps by token count from
- `_launch` (its docstring has the table). Replaces, per attention layer, the
- CUDA `fused_qk_norm_rope` main call, the index-scratch copy and the second
- `fused_qk_norm_rope` call: 6.67 -> 1.76 us at M=1, 9.15 -> 2.96 at M=128, 185
- -> 89.5 at 8192 tokens (graph replay).
-
- From `_CUTE_MIN_TOKENS` tokens, where the geometry is this checkpoint's (`_cute_ok`:
- head dims 256 / 128, 2 k heads, 4 index heads, a 64-dim rotary prefix, 16-byte
- aligned rows, contiguous index heads), the CuTe DSL `_NormRopeRing`: three
- persistent CTAs an SM, each forming every G-th token, a warp's next token
- prefetched by cp.async (its class docstring has the mapping). 16384 tokens: 182 -> 82.7 us in the
- microbenchmark (the Triton kernel was L1-bound; the ring runs at ~65% of DRAM peak,
- bounded by load latency at three CTAs an SM).
-
- Numerics: fp32 throughout, one rounding at the store; the CuTe kernel is bitwise
- the Triton kernel's eight-head launch (measured at 1024-16384 tokens, signed zeros
- included; not pinned by a test). Against the fp64
- reference over the same table every cell is within one output rounding
- step (measured max 1.48 ulp of the row scale, relative RMS <= 0.011 ulp).
- Against the CUDA pair it replaces (measured, not pinned by a test), <= 1.41
- ulp of the row scale and 0.09 ulp RMS (that op derives cos/sin with `__sincosf`, whose phase error grows
- with position). Three agreeing position rows select the same table row for
- every pair, so a text batch is exactly the plain rope on row 0.
-
- Preconditions in `is_valid`: the main q/k columns end at or before
- `index_q_col` (the planes do not overlap, so reading the index plane while
- the main plane is rewritten in place is race-free); `index_q` does not
- alias `proj`; positions below the table length (an out-of-range position
- reads past the table).
-
- Hard limits: token offsets in int64; positions index the table in int64.
-
- CUDA graphs: the Triton kernel captured at M=4 replays bitwise an eager call on
- inputs rewritten in place.
-
- Not certified: non-zero index-plane mRoPE sections (the target passes (0, 0)),
- capture of the CuTe kernel, `T == 0`, fp16/fp32 activations, `rotary_dim` > either head width,
- non-power-of-two head widths or head counts (tl.arange), YaRN, non-Gemma
- norm, a k plane on the index side (its scratch slot is not written).
- """
-
- def __call__(
- self,
- proj: torch.Tensor,
- index_q: torch.Tensor,
- position_ids: torch.Tensor,
- cos_sin: torch.Tensor,
- q_weight: torch.Tensor,
- k_weight: torch.Tensor,
- index_q_weight: torch.Tensor,
- num_heads_q: int,
- num_heads_k: int,
- head_dim: int,
- index_heads: int,
- index_head_dim: int,
- index_q_col: int,
- rotary_dim: int,
- eps: float,
- mrope_section1: int,
- mrope_section2: int,
- index_mrope_section1: int = 0,
- index_mrope_section2: int = 0,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
- num_tokens = proj.shape[0]
- if num_tokens == 0:
- return proj, index_q
- sections = (mrope_section1, mrope_section2, index_mrope_section1, index_mrope_section2)
- weights = (q_weight, k_weight, index_q_weight)
- if num_tokens >= _CUTE_MIN_TOKENS and _cute_ok(
- proj, index_q, cos_sin, weights, num_heads_q, num_heads_k, head_dim, index_heads,
- index_head_dim, index_q_col, rotary_dim,
- ): # fmt: skip
- _launch_cute(
- proj, index_q, position_ids, cos_sin, *weights, num_heads_q, index_q_col, eps,
- sections,
- ) # fmt: skip
- else:
- _launch_triton(
- proj, index_q, position_ids, cos_sin, *weights, num_heads_q, num_heads_k,
- head_dim, index_heads, index_head_dim, index_q_col, rotary_dim, eps, sections,
- ) # fmt: skip
- return proj, index_q
-
- def reference(
- self,
- proj: torch.Tensor,
- index_q: torch.Tensor,
- position_ids: torch.Tensor,
- cos_sin: torch.Tensor,
- q_weight: torch.Tensor,
- k_weight: torch.Tensor,
- index_q_weight: torch.Tensor,
- num_heads_q: int,
- num_heads_k: int,
- head_dim: int,
- index_heads: int,
- index_head_dim: int,
- index_q_col: int,
- rotary_dim: int,
- eps: float,
- mrope_section1: int,
- mrope_section2: int,
- index_mrope_section1: int = 0,
- index_mrope_section2: int = 0,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
- """`(proj after, index_q's written values)` in fp64 over the same table,
- rounded once; neither input is mutated."""
- num_tokens = proj.shape[0]
- half = rotary_dim // 2
- table = cos_sin.double()
- pos = position_ids.long()
- j = torch.arange(half, device=proj.device)
-
- def angles(s1: int, s2: int):
- coord = torch.zeros(half, dtype=torch.long, device=proj.device)
- coord[(j % 3 == 1) & (j < 3 * s1)] = 1
- coord[(j % 3 == 2) & (j < 3 * s2)] = 2
- p = pos[coord, :].t() # [T, half]
- return table[p, j][:, None, :], table[p, half + j][:, None, :]
-
- def plane(x: torch.Tensor, w: torch.Tensor, cos, sin) -> torch.Tensor:
- x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * (1.0 + w)
- x1, x2 = x[..., :half], x[..., half:rotary_dim]
- rotated = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], -1)
- return torch.cat([rotated, x[..., rotary_dim:]], -1)
-
- out = proj.double().clone()
- main_heads = num_heads_q + num_heads_k
- main_cols = main_heads * head_dim
- w = torch.cat(
- [
- q_weight.double().expand(num_heads_q, head_dim),
- k_weight.double().expand(num_heads_k, head_dim),
- ]
- )
- main = out[:, :main_cols].reshape(num_tokens, main_heads, head_dim)
- out[:, :main_cols] = plane(main, w, *angles(mrope_section1, mrope_section2)).reshape(
- num_tokens, main_cols
- )
- iq = proj[:, index_q_col : index_q_col + index_heads * index_head_dim].double()
- iq = plane(
- iq.reshape(num_tokens, index_heads, index_head_dim),
- index_q_weight.double(),
- *angles(index_mrope_section1, index_mrope_section2),
- )
- return out.to(proj.dtype), iq.to(index_q.dtype)
-
- def is_valid(
- self,
- proj: torch.Tensor,
- index_q: torch.Tensor,
- position_ids: torch.Tensor,
- cos_sin: torch.Tensor,
- q_weight: torch.Tensor,
- k_weight: torch.Tensor,
- index_q_weight: torch.Tensor,
- num_heads_q: int,
- num_heads_k: int,
- head_dim: int,
- index_heads: int,
- index_head_dim: int,
- index_q_col: int,
- rotary_dim: int,
- *unused_args: object,
- **unused_kwargs: object,
- ) -> None:
- num_tokens = proj.shape[0]
- assert proj.stride(-1) == 1, "proj must have a unit column stride"
- assert (num_heads_q + num_heads_k) * head_dim <= index_q_col, (
- "the index q columns overlap the main q/k heads: the in-place main write races "
- "the index read"
- )
- assert index_q_col + index_heads * index_head_dim <= proj.shape[1], (
- "the index q columns run past the projection row"
- )
- assert index_q.shape == (num_tokens, index_heads, index_head_dim), index_q.shape
- assert index_q.stride(-1) == 1, "index_q must have a unit dim stride"
- lo, hi = index_q.data_ptr(), proj.data_ptr()
- assert lo >= hi + proj.numel() * proj.element_size() or hi >= lo + (
- index_q.shape[0] * index_q.stride(0) * index_q.element_size()
- ), "index_q aliases proj"
- assert position_ids.shape == (3, num_tokens) and position_ids.is_contiguous(), (
- f"position_ids must be a contiguous [3, {num_tokens}]; got {tuple(position_ids.shape)}"
- )
- assert position_ids.dtype == torch.int32, position_ids.dtype
- assert cos_sin.dtype == torch.float32 and cos_sin.shape[1] == rotary_dim, cos_sin.shape
- assert rotary_dim <= min(head_dim, index_head_dim) and rotary_dim % 2 == 0, rotary_dim
- if num_tokens:
- assert int(position_ids.min()) >= 0, "negative position"
- assert int(position_ids.max()) < cos_sin.shape[0], "position past the rope table"
-
-
-qwen4_exp_qk_norm_rope = Qwen4ExpQkNormRope()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/spec_rollback_commit.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/spec_rollback_commit.py
deleted file mode 100644
index 098506078e6f..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/spec_rollback_commit.py
+++ /dev/null
@@ -1,368 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""`SpeculativeRollback`'s staging and commit of a verify step in one launch each.
-
-Tier 5 (Triton written for this target). A verify step runs `p` rows per generation request
-through the trunk; the PLE pools and the QSA planes are written for every row. The staging
-launch copies what the commit rewrites from: each request's pre-step n-gram id history
-followed by its rows' ids (the `stream` whose length-`context` windows are the histories
-after each row), and its pre-step PLE conv state. After the draft loop the commit launch
-rewrites, for each generation request that accepted `a` rows:
-
-* its n-gram pool row to window `a` of its stream (the history after row `a - 1`);
-* its PLE conv pool row to columns `[a, a + state_len)` of `[pre-step state | the conv input
- of its p rows]`;
-* the QSA position plane and every attention layer's index-K plane at each of its rows of
- ordinal `>= a` to the values gathered before the trunk wrote them.
-
-The commit runs one program per (request, block): block `b` writes the request's conv
-channels `[b * 256, (b + 1) * 256)` and restores its rejected rows in index-K plane `b`; block 0
-also writes its n-gram row and position-plane slots.
-"""
-
-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
-
-from .qsa_rollback_planes import ROW
-
-#: Conv channels a program covers.
-_BLOCK_C = 256
-
-
-@triton.jit
-def _stage_kernel(
- ngram_pool,
- ngram_stride,
- conv_pool,
- conv_s0,
- conv_s1,
- conv_s2,
- slots,
- ids,
- stream,
- conv_prev,
- channels,
- P: tl.constexpr,
- CTX: tl.constexpr,
- W_BLOCK: tl.constexpr,
- STATE: tl.constexpr,
- STATE_BLOCK: tl.constexpr,
- BLOCK_C: tl.constexpr,
-):
- i = tl.program_id(0)
- cb = tl.program_id(1)
- slot = tl.load(slots + i).to(tl.int64)
- c = cb * BLOCK_C + tl.arange(0, BLOCK_C)
- s = tl.arange(0, STATE_BLOCK)
- m = (c < channels)[:, None] & (s < STATE)[None, :]
- vals = tl.load(conv_pool + slot * conv_s0 + c[:, None] * conv_s1 + s[None, :] * conv_s2, mask=m)
- tl.store(
- conv_prev + (i.to(tl.int64) * channels + c[:, None]) * STATE + s[None, :], vals, mask=m
- )
- if cb == 0:
- w = tl.arange(0, W_BLOCK)
- hist = tl.load(ngram_pool + slot * ngram_stride + w, mask=w < CTX)
- row = tl.load(ids + i * P + (w - CTX), mask=(w >= CTX) & (w < CTX + P)).to(tl.int64)
- tl.store(stream + i * (CTX + P) + w, tl.where(w < CTX, hist, row), mask=w < CTX + P)
-
-
-@triton.jit
-def _commit_kernel(
- accepted,
- slots,
- ngram_pool,
- ngram_stride,
- stream,
- conv_pool,
- conv_s0,
- conv_s1,
- conv_s2,
- conv_prev,
- conv_x,
- x_s0,
- x_s1,
- x_s2,
- position,
- pos_s0,
- pos_s1,
- pos_s2,
- position_old,
- pages,
- within,
- plane_ptrs,
- plane_strides,
- index_old,
- num_planes,
- rows,
- channels,
- P: tl.constexpr,
- P_BLOCK: tl.constexpr,
- CTX: tl.constexpr,
- CTX_BLOCK: tl.constexpr,
- STATE: tl.constexpr,
- STATE_BLOCK: tl.constexpr,
- BLOCK_C: tl.constexpr,
- ROW: tl.constexpr,
-):
- i = tl.program_id(0)
- cb = tl.program_id(1)
- a = tl.load(accepted + i).to(tl.int32)
- slot = tl.load(slots + i).to(tl.int64)
- i64 = i.to(tl.int64)
-
- # PLE conv row: column `a + s` of [pre-step state | conv input of rows 0..P-1].
- c = cb * BLOCK_C + tl.arange(0, BLOCK_C)
- s = tl.arange(0, STATE_BLOCK)
- j = a + s
- m = (c < channels)[:, None] & (s < STATE)[None, :]
- from_prev = (j < STATE)[None, :]
- prev = tl.load(
- conv_prev + (i64 * channels + c[:, None]) * STATE + j[None, :], mask=m & from_prev
- )
- xin = tl.load(
- conv_x + i64 * x_s0 + (j - STATE).to(tl.int64)[None, :] * x_s1 + c[:, None] * x_s2,
- mask=m & ~from_prev,
- )
- tl.store(
- conv_pool + slot * conv_s0 + c[:, None] * conv_s1 + s[None, :] * conv_s2,
- tl.where(from_prev, prev, xin),
- mask=m,
- )
-
- # The rows of ordinal >= a get back the plane values gathered before the trunk wrote
- # them: index-K plane `cb` here, the position plane in block 0.
- o = tl.arange(0, P_BLOCK)
- r = i64 * P + o
- rejected = (o < P) & (o >= a)
- if cb < num_planes:
- page = tl.load(pages + r, mask=rejected, other=0)
- tok = tl.load(within + r, mask=rejected, other=0)
- e = tl.arange(0, ROW)
- base = tl.load(plane_ptrs + cb).to(tl.pointer_type(tl.bfloat16))
- stride = tl.load(plane_strides + cb)
- vals = tl.load(
- index_old + (cb * rows + r[:, None]) * ROW + e[None, :], mask=rejected[:, None]
- )
- tl.store(
- base + page[:, None] * stride + tok[:, None] * ROW + e[None, :],
- vals,
- mask=rejected[:, None],
- )
-
- if cb == 0:
- # n-gram row: the id window after row `a - 1` (a negative row wraps, as torch
- # indexing does).
- t = a - 1
- t = tl.where(t < 0, t + P, t)
- k = tl.arange(0, CTX_BLOCK)
- hist = tl.load(stream + i64 * (CTX + P) + t + 1 + k, mask=k < CTX)
- tl.store(ngram_pool + slot * ngram_stride + k, hist, mask=k < CTX)
-
- page = tl.load(pages + r, mask=rejected, other=0)
- tok = tl.load(within + r, mask=rejected, other=0)
- ax = tl.arange(0, 4)
- pm = rejected[:, None] & (ax < 3)[None, :]
- old_pos = tl.load(position_old + r[:, None] * 3 + ax[None, :], mask=pm)
- tl.store(
- position + page[:, None] * pos_s0 + tok[:, None] * pos_s1 + ax[None, :] * pos_s2,
- old_pos,
- mask=pm,
- )
-
-
-def _pow2(n: int) -> int:
- return triton.next_power_of_2(max(int(n), 1))
-
-
-class SpecRollbackStage(OpWrapper):
- """`(torch.cat([ngram_pool[slots], ids.reshape(n, p).to(int64)], 1),
- conv_pool[slots].contiguous())`: per generation request, its n-gram id history followed
- by its `p` rows' ids, and its PLE conv state, read before the trunk writes either pool.
-
- Fusion boundary: the two reads of the staging; the slots and ids are the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
- CELLS: tuple[Cell, ...] = (
- Cell(why="MTP3, one request", spec=dict(n=1, p=4)),
- Cell(why="MTP3, 128 requests", spec=dict(n=128, p=4)),
- )
- note = """
- Metadata consumed: none. Call site: `SpeculativeRollback.stage_ple_verify`, before
- `ple_ngram_embed` publishes the id history. Grid (requests, conv channel blocks).
- """
-
- def __call__(self, ngram_pool, conv_pool, slots, ids, p: int):
- n = slots.shape[0]
- ctx = ngram_pool.shape[1]
- channels, state = conv_pool.shape[1], conv_pool.shape[2]
- stream = torch.empty([n, ctx + p], dtype=ngram_pool.dtype, device=ngram_pool.device)
- conv_prev = torch.empty(
- [n, channels, state], dtype=conv_pool.dtype, device=conv_pool.device
- )
- if n:
- _stage_kernel[(n, triton.cdiv(channels, _BLOCK_C))](
- ngram_pool,
- ngram_pool.stride(0),
- conv_pool,
- conv_pool.stride(0),
- conv_pool.stride(1),
- conv_pool.stride(2),
- slots,
- ids,
- stream,
- conv_prev,
- channels,
- P=p,
- CTX=ctx,
- W_BLOCK=_pow2(ctx + p),
- STATE=state,
- STATE_BLOCK=_pow2(state),
- BLOCK_C=_BLOCK_C,
- num_warps=4,
- )
- return stream, conv_prev
-
- def reference(self, ngram_pool, conv_pool, slots, ids, p: int):
- n = slots.shape[0]
- slots = slots.to(torch.long)
- rows = torch.reshape(ids[: n * p], [n, p]).to(ngram_pool.dtype)
- return torch.cat([ngram_pool[slots], rows], 1), conv_pool[slots]
-
- def is_valid(self, ngram_pool, conv_pool, slots, ids, p: int) -> None:
- assert ngram_pool.dim() == 2 and ngram_pool.stride(1) == 1
- assert ngram_pool.dtype == torch.int64
- assert conv_pool.dim() == 3
- assert slots.dim() == 1 and ids.dim() == 1 and ids.is_contiguous()
- assert ids.shape[0] >= slots.shape[0] * p
-
-
-class SpecRollbackCommit(OpWrapper):
- """For each generation request `i` with `a = accepted[i]` accepted rows (golden token
- included, `1 <= a <= p`) and slot `slots[i]`:
-
- * `ngram_pool[slot] = stream[i, a : a + context]`;
- * `conv_pool[slot] = cat([conv_prev[i], conv_x[i].T], 1)[:, a : a + state_len]`;
- * for each row `r = i * p + t` with `t >= a`: `position[pages[r], within[r]] =
- position_old[r]` and `plane_l[pages[r], within[r]] = index_old[l, r]` for every
- index-K plane `l` of `plane_table`.
-
- Fusion boundary: the whole commit of `SpeculativeRollback`; the staged tensors are what
- `SpecRollbackStage`, `qsa_rollback_gather` and the position gather produced.
- """
-
- ARCHS = frozenset({Arch.SM_103})
- CELLS: tuple[Cell, ...] = (
- Cell(why="MTP3, one request, 12 planes", spec=dict(n=1, p=4, planes=12)),
- Cell(why="MTP3, 128 requests, 12 planes", spec=dict(n=128, p=4, planes=12)),
- )
- note = """
- Metadata consumed: none. Call site: `SpeculativeRollback.commit_speculative_states`,
- after the draft loop. Grid (requests, max(conv channel blocks, planes)). Slots or plane
- slots shared by two requests (CUDA-graph padding dummies) are written by both; which one
- lands is unordered, as with the indexed writes this replaces.
- """
-
- def __call__(
- self,
- accepted,
- slots,
- ngram_pool,
- stream,
- conv_pool,
- conv_prev,
- conv_x,
- position,
- position_old,
- pages,
- within,
- plane_ptrs,
- plane_strides,
- index_old,
- ) -> None:
- n, p, channels = conv_x.shape
- state = conv_pool.shape[2]
- ctx = ngram_pool.shape[1]
- if not n:
- return
- blocks = max(triton.cdiv(channels, _BLOCK_C), index_old.shape[0])
- _commit_kernel[(n, blocks)](
- accepted,
- slots,
- ngram_pool,
- ngram_pool.stride(0),
- stream,
- conv_pool,
- conv_pool.stride(0),
- conv_pool.stride(1),
- conv_pool.stride(2),
- conv_prev,
- conv_x,
- conv_x.stride(0),
- conv_x.stride(1),
- conv_x.stride(2),
- position,
- position.stride(0),
- position.stride(1),
- position.stride(2),
- position_old,
- pages,
- within,
- plane_ptrs,
- plane_strides,
- index_old,
- index_old.shape[0],
- index_old.shape[1],
- channels,
- P=p,
- P_BLOCK=_pow2(p),
- CTX=ctx,
- CTX_BLOCK=_pow2(ctx),
- STATE=state,
- STATE_BLOCK=_pow2(state),
- BLOCK_C=_BLOCK_C,
- ROW=ROW,
- num_warps=4,
- )
-
- def reference(self, *args) -> None:
- raise NotImplementedError("test_modeling_v2_spec_rollback drives the handler end to end")
-
- def is_valid(
- self,
- accepted,
- slots,
- ngram_pool,
- stream,
- conv_pool,
- conv_prev,
- conv_x,
- position,
- position_old,
- pages,
- within,
- plane_ptrs,
- plane_strides,
- index_old,
- ) -> None:
- n, p, channels = conv_x.shape
- assert accepted.shape == (n,) and slots.shape == (n,)
- assert ngram_pool.dtype == stream.dtype == torch.int64 and ngram_pool.stride(1) == 1
- assert stream.shape == (n, ngram_pool.shape[1] + p) and stream.is_contiguous()
- assert conv_pool.shape[1] == channels and conv_prev.is_contiguous()
- assert conv_prev.shape == (n, channels, conv_pool.shape[2])
- assert position.dim() == 3 and position.shape[2] == 3
- assert position_old.shape == (n * p, 3) and position_old.is_contiguous()
- assert pages.shape == within.shape == (n * p,)
- assert pages.dtype == within.dtype == torch.int64
- assert index_old.is_contiguous() and index_old.shape[1:] == (n * p, ROW)
- assert plane_ptrs.numel() == plane_strides.numel() == index_old.shape[0]
-
-
-spec_rollback_stage = SpecRollbackStage()
-spec_rollback_commit = SpecRollbackCommit()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/verify_step_setup.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/verify_step_setup.py
deleted file mode 100644
index bdd150c32d71..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/verify_step_setup.py
+++ /dev/null
@@ -1,215 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""The decode target's per-step setup of a verify step in one launch.
-
-Tier 5 (Triton written for this target). Before the trunk runs, a verify step needs:
-
-* `pos3`: a contiguous int32 `[3, T]` copy of the step's mRoPE positions;
-* the PLE segmentation vector `cu_seqlens` (`[batch + 1]` int64, leading zero) of the
- per-request row counts;
-* the speculative rollback's staging: each row's QSA slot (`pages`, `within`, int64; page
- `block_table[request, position // tokens_per_block]` and offset `position % tokens_per_block`,
- floor semantics, a negative column wrapping as torch indexing does), the requests' recurrent
- slots as int64, and every attention layer's index-K plane at those slots (`[L, T, 128]` bf16,
- read before the trunk writes any plane).
-
-Program `(0, b)` writes block `b` of the rows' positions and slots (block 0 also the
-segmentation vector and the int64 slots); program `(1 + l, b)` copies plane `l` at the same
-rows, resolving each row's slot as program `(0, b)` does.
-"""
-
-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
-
-from .qsa_rollback_planes import ROW
-
-#: Rows a program covers.
-_BLOCK_T = 16
-#: Most requests the segmentation vector covers (`batch + 1` entries).
-_BLOCK_S = 512
-
-
-@triton.jit
-def _slot(req_idx, logical, block_table, bt_stride, bt_cols, rows, tpb):
- req = tl.load(req_idx + rows).to(tl.int64)
- pos = tl.load(logical + rows).to(tl.int64)
- q = pos // tpb
- r = pos % tpb
- neg = (r != 0) & (pos < 0)
- q = tl.where(neg, q - 1, q)
- r = tl.where(neg, r + tpb, r)
- q = tl.where(q < 0, q + bt_cols, q)
- page = tl.load(block_table + req * bt_stride + q).to(tl.int64)
- return page, r
-
-
-@triton.jit
-def _setup_kernel(
- positions,
- pos_stride,
- seq_lens,
- slots,
- req_idx,
- logical,
- block_table,
- bt_stride,
- bt_cols,
- plane_ptrs,
- plane_strides,
- pos3,
- cu_seqlens,
- slots_out,
- pages_out,
- within_out,
- index_out,
- num_rows,
- batch,
- tpb,
- BLOCK_T: tl.constexpr,
- BLOCK_S: tl.constexpr,
- ROW: tl.constexpr,
-):
- who = tl.program_id(0)
- tb = tl.program_id(1)
- t = tb * BLOCK_T + tl.arange(0, BLOCK_T)
- tm = t < num_rows
- tt = tl.where(tm, t, 0)
- page, within = _slot(req_idx, logical, block_table, bt_stride, bt_cols, tt, tpb)
- if who == 0:
- tl.store(pages_out + t, page, mask=tm)
- tl.store(within_out + t, within, mask=tm)
- for a in tl.static_range(3):
- p = tl.load(positions + a * pos_stride + t, mask=tm)
- tl.store(pos3 + a * num_rows + t, p.to(tl.int32), mask=tm)
- if tb == 0:
- s = tl.arange(0, BLOCK_S)
- sm = s < batch
- lengths = tl.load(seq_lens + s, mask=sm, other=0).to(tl.int64)
- tl.store(cu_seqlens + 1 + s, tl.cumsum(lengths, axis=0), mask=sm)
- tl.store(cu_seqlens + s, tl.zeros([BLOCK_S], dtype=tl.int64), mask=s == 0)
- tl.store(slots_out + s, tl.load(slots + s, mask=sm).to(tl.int64), mask=sm)
- else:
- plane = who - 1
- base = tl.load(plane_ptrs + plane).to(tl.pointer_type(tl.bfloat16))
- stride = tl.load(plane_strides + plane)
- e = tl.arange(0, ROW)
- vals = tl.load(
- base + page[:, None] * stride + within[:, None] * ROW + e[None, :], mask=tm[:, None]
- )
- tl.store(
- index_out + (plane.to(tl.int64) * num_rows + t[:, None]) * ROW + e[None, :],
- vals,
- mask=tm[:, None],
- )
-
-
-class VerifyStepSetup(OpWrapper):
- """`(pos3, cu_seqlens, slots.long(), pages, within, index_old)`:
-
- * `pos3 = position_ids.reshape(3, -1)` as a fresh contiguous int32 tensor;
- * `cu_seqlens = [0, cumsum(seq_lens)]` int64;
- * `pages = block_table[request_indices.long(), logical_positions // tokens_per_block]`
- and `within = logical_positions % tokens_per_block`, int64;
- * `index_old = qsa_rollback_gather(plane_ptrs, plane_strides, pages, within, L)`.
-
- Fusion boundary: the decode target's step setup before the trunk; which rows, slots and
- planes are the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
- CELLS: tuple[Cell, ...] = (
- Cell(why="MTP3, one request, 12 planes", spec=dict(batch=1, p=4, planes=12)),
- Cell(why="MTP3, 32 requests, 12 planes", spec=dict(batch=32, p=4, planes=12)),
- )
- note = """
- Metadata consumed: none (the caller passes `seq_lens_cuda[:batch]`, the QSA row mapping
- and block table). Call site: `DecodeTarget.forward` of a captured verify step, after the
- captured preprocessing corrected positions and lengths and before the trunk.
- Preconditions: `position_ids` reshapes to `[3, T]` with unit last stride; at most
- `_BLOCK_S - 1` requests.
- """
-
- def __call__(
- self,
- position_ids,
- seq_lens,
- slots,
- request_indices,
- logical_positions,
- block_table,
- tokens_per_block: int,
- plane_ptrs,
- plane_strides,
- num_planes: int,
- ):
- pos2 = torch.reshape(position_ids, [3, -1])
- rows = pos2.shape[1]
- batch = seq_lens.shape[0]
- dev = pos2.device
- pos3 = torch.empty([3, rows], dtype=torch.int32, device=dev)
- cu_seqlens = torch.empty(batch + 1, dtype=torch.int64, device=dev)
- slots_out = torch.empty(batch, dtype=torch.int64, device=dev)
- pages = torch.empty(rows, dtype=torch.int64, device=dev)
- within = torch.empty(rows, dtype=torch.int64, device=dev)
- index_old = torch.empty([num_planes, rows, ROW], dtype=torch.bfloat16, device=dev)
- if rows:
- _setup_kernel[(1 + num_planes, triton.cdiv(rows, _BLOCK_T))](
- pos2,
- pos2.stride(0),
- seq_lens,
- slots,
- request_indices,
- logical_positions,
- block_table,
- block_table.stride(0),
- block_table.shape[1],
- plane_ptrs,
- plane_strides,
- pos3,
- cu_seqlens,
- slots_out,
- pages,
- within,
- index_old,
- rows,
- batch,
- tokens_per_block,
- BLOCK_T=_BLOCK_T,
- BLOCK_S=_BLOCK_S,
- ROW=ROW,
- num_warps=4,
- )
- return pos3, cu_seqlens, slots_out, pages, within, index_old
-
- def reference(self, *args):
- raise NotImplementedError("test_modeling_v2_verify_step_setup compares the composed ops")
-
- def is_valid(
- self,
- position_ids,
- seq_lens,
- slots,
- request_indices,
- logical_positions,
- block_table,
- tokens_per_block: int,
- plane_ptrs,
- plane_strides,
- num_planes: int,
- ) -> None:
- pos2 = torch.reshape(position_ids, [3, -1])
- assert pos2.stride(1) == 1
- assert 0 < seq_lens.shape[0] < _BLOCK_S and seq_lens.is_contiguous()
- assert slots.shape == seq_lens.shape
- assert request_indices.shape[0] >= pos2.shape[1]
- assert logical_positions.shape[0] >= pos2.shape[1]
- assert block_table.stride(1) == 1 and tokens_per_block > 0
- assert plane_ptrs.numel() == plane_strides.numel() == num_planes
-
-
-verify_step_setup = VerifyStepSetup()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/vision_rope_2d.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/vision_rope_2d.py
deleted file mode 100644
index 6e5d284f768b..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/vision_rope_2d.py
+++ /dev/null
@@ -1,262 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Vision-tower 2-D axial rotary embedding (patch row / patch column), no KV
-cache.
-
-Not the text tower's mRoPE: that rotates 64 of a 256-wide head against three
-interleaved position rows (`mrope_section`, `rope_theta` 1e7); this rotates the
-whole 72-wide vision head against two spatial rows from the image grid, 18 row
-angles then 18 column angles, `theta` 1e4. Neither reads the other's config.
-
-Tier-3 composed from native torch (angle table, index gather, cos/sin, two
-multiply-adds). Every registered rope op takes 1-D positions;
-`flashinfer_apply_rope_with_cos_sin_cache_inplace` with an explicit per-token
-table would sidestep that, but at `head_size=72` it raises "misaligned
-address" and poisons the CUDA context.
-"""
-
-from __future__ import annotations
-
-from typing import List, Sequence, Tuple
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: This checkpoint: 1152 / 16 heads = a 72-wide head, spatial_merge_size 2,
-#: theta 1e4 (`model.visual.blocks.N.attn.qkv.weight` is [3456, 1152]).
-NUM_HEADS = 16
-HEAD_DIM = 72
-MERGE = 2
-THETA = 10000.0
-
-#: (frames, height, width) in patch units, from the runtime-value dump of the
-#: mmmu gate (58.33): 170 segment sizes from 256 to 12544 per frame, calls of up
-#: to 22144 rows. 62x90 is a grid this processor emits.
-SERVING_FRAME = ((1, 62, 90),)
-
-
-def _rotate_half(x: torch.Tensor, half: int) -> torch.Tensor:
- """NeoX-style half rotation: `[x0, x1] -> [-x1, x0]` over the last dim."""
- return torch.cat((-x[..., half:], x[..., :half]), dim=-1)
-
-
-def _spatial_merge_positions(grid: Sequence[Sequence[int]], merge: int) -> List[Tuple[int, int]]:
- """Per-token `(row, col)` by four nested loops, in spatial-merge order."""
- positions: List[Tuple[int, int]] = []
- for frames, height, width in grid:
- frame = [
- (block_row * merge + r, block_col * merge + c)
- for block_row in range(height // merge)
- for block_col in range(width // merge)
- for r in range(merge)
- for c in range(merge)
- ]
- positions.extend(frame * frames)
- return positions
-
-
-class VisionRope2d(OpWrapper):
- """Rotate every head of `q` and `k` by its patch's 2-D axial angles; return
- fresh `(q_out, k_out)` in the inputs' shapes and dtypes.
-
- rope_dim = head_dim // 2; axis_dim = rope_dim // 2
- inv_freq[j] = theta ** (-2j / rope_dim) j < axis_dim
- angle[t] = [row(t) * inv_freq, col(t) * inv_freq] height half first
- out[.., i] = x[i] cos a_i - x[i + rope_dim] sin a_i
- out[.., i + rope_dim] = x[i + rope_dim] cos a_i + x[i] sin a_i
-
- NeoX half-split pairing `(i, i + rope_dim)`, not GPT-J `(2j, 2j+1)`. Tokens
- are in spatial-merge order (whole `merge x merge` blocks row-major,
- row-major inside a block) -- what the patch merger consumes; a raster walk
- is silently wrong. A video's frames repeat the spatial grid: no temporal
- axis. The angle table is built in fp64 and rounded once to fp32; the
- rotation runs in fp32 and rounds once on the store.
-
- Fusion boundary: the angle table and the rotation. The qkv projection, the
- unrotated `v`, the per-frame attention segmentation and the attention are
- the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: grid, dtype; `heads` / `heads_k` / `head_dim` (default 16 / heads /
- # 72); `theta` (default 1e4); `grid_dtype` / `grid_device` for `grid_thw`
- # (default int32 on the host); `layout` "separate" (default) or
- # "packed_qkv" -- q and k as the `split` views of one [T, 3*heads, D]
- # projection, which is how the tower calls it.
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="call site: q/k split views of packed qkv, 62x90",
- spec=dict(grid=SERVING_FRAME, dtype=torch.bfloat16, layout="packed_qkv"),
- ),
- Cell(
- why="call site: packed qkv views, video + image (two grid rows)",
- spec=dict(
- grid=((3, 8, 8), (1, 16, 24)),
- dtype=torch.bfloat16,
- layout="packed_qkv",
- grid_dtype=torch.int64,
- ),
- ),
- Cell(
- why="call site: packed qkv views, video at the 32768-row ceiling",
- spec=dict(
- grid=((64, 16, 32),),
- dtype=torch.bfloat16,
- layout="packed_qkv",
- grid_dtype=torch.int64,
- ),
- ),
- )
-
- note = """
- Metadata consumed: none. `grid_thw` is read with `.tolist()`: a device
- tensor forces a sync (int32/int64, host/device all bitwise equal). `q` and
- `k` need last-dim stride 1 only; they are not mutated.
-
- `head_dim` must be divisible by 4; otherwise the angle table's width
- mismatches the head and the broadcast raises (loud, so not guarded; not
- pinned by a test).
-
- Accuracy vs fp64 angles over the cells: bf16 <= 1.56e-2 (one ulp at
- unit-normal operands). Measured, not certified: fp32 <= 4.77e-7, flat in the
- grid side up to 6272 -- the fp64 table is why (an fp32 `inv_freq` times fp32
- position moved an fp32 output 1.475e-5 at a 64 side, outside the default
- gate); wrong operators at 1024 bf16 tokens land 8.5-9.5 away: GPT-J pairing,
- axes swapped, raster order, `inv_freq` over `head_dim`.
-
- Host-bound: ~353 us/call at 1024 and 4096 tokens (`.tolist()` + small
- launches), x27 blocks per image. The table depends only on the grid, and the
- result does not depend on the grid's placement (host or device) or integer
- dtype (int32 or int64).
-
- Not certified: float16 or float32; mixed q/k dtypes or head counts; head
- geometries other than 16 x 72; float8; `spatial_merge_size` != 2; `theta`
- other than 1e4; grid sides past 90 or calls past 32768 rows; CUDA-graph
- capture (it would bake in the grid).
- """
-
- def __call__(
- self,
- q: torch.Tensor,
- k: torch.Tensor,
- grid_thw: torch.Tensor,
- spatial_merge_size: int,
- theta: float = 10000.0,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
- num_tokens, _, head_dim = q.shape
- # Angles per token: half the head (the other half is each pair's
- # partner); each axis owns half of those, height first.
- rope_dim = head_dim // 2
-
- grid: List[List[int]] = grid_thw.tolist()
-
- device = q.device
- # fp64 then one rounding to fp32: the largest angle is (side - 1)
- # radians, and fp32 throughout rounds twice there.
- inv_freq = 1.0 / (
- theta ** (torch.arange(0, rope_dim, 2, dtype=torch.float64, device=device) / rope_dim)
- )
- max_side = max(max(height, width) for _, height, width in grid)
- freq_table = torch.outer(
- torch.arange(max_side, dtype=torch.float64, device=device), inv_freq
- ).to(torch.float32)
-
- merge = spatial_merge_size
- intra = torch.arange(merge, device=device)
- pos_ids = torch.empty((num_tokens, 2), dtype=torch.long, device=device)
- offset = 0
- for frames, height, width in grid:
- blocks_h, blocks_w = height // merge, width // merge
- rows = (
- (
- torch.arange(blocks_h, device=device)[:, None, None, None] * merge
- + intra[None, None, :, None]
- )
- .expand(blocks_h, blocks_w, merge, merge)
- .reshape(-1)
- )
- cols = (
- (
- torch.arange(blocks_w, device=device)[None, :, None, None] * merge
- + intra[None, None, None, :]
- )
- .expand(blocks_h, blocks_w, merge, merge)
- .reshape(-1)
- )
- coords = torch.stack((rows, cols), dim=-1)
- if frames > 1:
- coords = coords.repeat(frames, 1)
- pos_ids[offset : offset + coords.shape[0]] = coords
- offset += coords.shape[0]
-
- # [tokens, 2, axis_dim] -> [tokens, rope_dim]: height angles, then width.
- angles = freq_table[pos_ids].flatten(1)
- duplicated = torch.cat((angles, angles), dim=-1)
- cos = duplicated.cos().unsqueeze(-2)
- sin = duplicated.sin().unsqueeze(-2)
-
- q_fp32, k_fp32 = q.float(), k.float()
- q_out = (q_fp32 * cos + _rotate_half(q_fp32, rope_dim) * sin).to(q.dtype)
- k_out = (k_fp32 * cos + _rotate_half(k_fp32, rope_dim) * sin).to(k.dtype)
- return q_out, k_out
-
- def reference(
- self,
- q: torch.Tensor,
- k: torch.Tensor,
- grid_thw: torch.Tensor,
- spatial_merge_size: int,
- theta: float = 10000.0,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
- """Positions from nested loops, angles from Python-float `inv_freq` times
- integer positions in fp64, rotation as a complex multiply in fp32."""
- rope_dim = q.shape[-1] // 2
- axis_dim = rope_dim // 2
- positions = _spatial_merge_positions(grid_thw.tolist(), spatial_merge_size)
- inverse = torch.tensor(
- [theta ** (-(2.0 * j) / rope_dim) for j in range(axis_dim)], dtype=torch.float64
- )
- pos = torch.tensor(positions, dtype=torch.float64).reshape(-1, 2)
- angles = (pos[:, :, None] * inverse).flatten(1).to(device=q.device, dtype=torch.float32)
- rotor = torch.polar(torch.ones_like(angles), angles).unsqueeze(-2)
-
- def rotate(x: torch.Tensor) -> torch.Tensor:
- fp32 = x.float()
- rotated = torch.complex(fp32[..., :rope_dim], fp32[..., rope_dim:]) * rotor
- return torch.cat((rotated.real, rotated.imag), dim=-1).to(x.dtype)
-
- return rotate(q), rotate(k)
-
- def is_valid(
- self,
- q: torch.Tensor,
- k: torch.Tensor,
- grid_thw: torch.Tensor,
- spatial_merge_size: int,
- theta: float = 10000.0,
- ) -> None:
- # Each of these returns a full-size result without an error: a side
- # that does not divide fills only the whole blocks' rows of the
- # position table and leaves the rest uninitialized (four tokens of a
- # 4x5 grid read back as unrotated); a grid describing fewer tokens than
- # q carries leaves the tail uninitialized the same way; a one-token k
- # broadcasts against every token's angles.
- assert q.dim() == 3 and k.dim() == 3, (
- f"q and k must be [tokens, heads, head_dim]; got q {tuple(q.shape)}, k {tuple(k.shape)}"
- )
- assert q.shape[0] == k.shape[0] and q.shape[-1] == k.shape[-1], (
- f"q and k must agree on tokens and head_dim; got q {tuple(q.shape)}, k {tuple(k.shape)}"
- )
- total = 0
- for frames, height, width in grid_thw.tolist():
- assert height % spatial_merge_size == 0 and width % spatial_merge_size == 0, (
- f"grid sides must be divisible by spatial_merge_size={spatial_merge_size}; "
- f"got height={height}, width={width}"
- )
- total += frames * height * width
- assert total == q.shape[0], f"grid_thw describes {total} tokens but q/k carry {q.shape[0]}"
-
-
-vision_rope_2d = VisionRope2d()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/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
deleted file mode 100644
index e3841c25f537..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/ngram_hash.py
+++ /dev/null
@@ -1,252 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""PLE n-gram hash: a token-id history becomes one global embedding row per head.
-
-Tier-3 composed from native torch integer ops (~30 small launches); exact, so
-gated bitwise. Nothing registered under ``torch.ops.trtllm`` hashes n-grams.
-``EngramHashProvider.compute_hashes`` in
-``tensorrt_llm/_torch/modules/engram/engram.py`` is the same family and a
-different function on three axes: it hashes *compressed* ids (a tokenizer
-table this checkpoint does not ship), derives its multipliers from a seed
-(six ``(seed, layer)`` combinations reproduced none of this checkpoint's
-stored ``layer_multipliers``), and pads instead of resetting at every EOS. Its
-moduli rule does agree. The built-in Qwen4-Exp implementation has its own,
-not wrapped here.
-"""
-
-from __future__ import annotations
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: This checkpoint's `ple_embedding` buffers: 3-gram, 8 heads per order.
-CKPT_MULTIPLIERS = (23703573157769, 20109073645365, 8052911324071)
-CKPT_HEAD_VOCAB_SIZES = (
- 20000003, 20000023, 20000033, 20000047, 20000059, 20000063, 20000069, 20000077,
- 20000081, 20000093, 20000107, 20000147, 20000153, 20000159, 20000161, 20000171,
-) # fmt: skip
-CKPT_HEAD_OFFSETS = (
- 0, 20000003, 40000026, 60000059, 80000106, 100000165, 120000228, 140000297,
- 160000374, 180000455, 200000548, 220000655, 240000802, 260000955, 280001114, 300001275,
-) # fmt: skip
-CKPT_EOS_TOKEN_ID = 248044
-CKPT_VOCAB_SIZE = 248320
-#: `ngram_size - 1`: the cached ids ahead of the new tokens.
-CKPT_CONTEXT_LEN = len(CKPT_MULTIPLIERS) - 1
-
-
-def _rows(batch: int, length: int, query: int, why: str, **spec) -> Cell:
- return Cell(why=why, spec=dict(batch=batch, length=length, num_query_tokens=query, **spec))
-
-
-class NgramHash(OpWrapper):
- """For a query token at history position `t` and head `h` of the `n`-gram
- block (`m = layer_multipliers`, `^` int64 XOR, `%` non-negative)::
-
- mixed[t] = (c[t] * m[0]) ^ (c[t - 1] * m[1]) ^ ... ^ (c[t - n + 1] * m[n - 1])
- row[t, h] = mixed[t] % head_vocab_sizes[h] + head_offsets[h]
-
- `c[t-k]` is `token_history[t-k]` when `t - k` lies in `t`'s segment and
- `eos_token_id` otherwise; a segment starts one past the most recent EOS
- **strictly before** `t` (an EOS closes its segment, so the token after it
- hashes as a sequence's first). Shifts before position 0 read EOS too.
- Heads run `[2-gram block | ... | ngram_size-gram block]`, `heads_per_ngram
- = head_vocab_sizes.numel() // (ngram_size - 1)` each, `ngram_size =
- layer_multipliers.numel()`.
-
- Returns a fresh contiguous int64 `[batch, num_query_tokens, ngram_heads]`
- for the **last** `num_query_tokens` positions; the leading ones only give
- the first query token its context. Indices are **global** -- the offset is
- applied, one concatenated table is addressed. No input is modified.
-
- Fusion boundary: segment scan, shifted streams, mix, modulo and offset.
- The id history (here built by `ngram_token_history`) and the table gather
- are the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: batch, length, num_query_tokens; history is uniform random 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.
- #
- # 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(4, 3, 1, "decode row whose cached context is EOS", eos_positions=(0, 1)),
- _rows(
- 32, 3, 1, "call site: 32 windows, decode / MTP verify", lead_eos=False, eos_density=0.2
- ),
- _rows(
- 8192, 3, 1, "call site: 8192 windows, largest prefill", lead_eos=False, eos_density=0.05
- ),
- )
-
- note = """
- Metadata consumed: none; rows are independent, and within a row the only
- structure seen is `eos_token_id` -- packing two sequences into one row is
- safe only when the first ends in EOS. Call site: `ngram_hash(window,
- multipliers, head_vocab, head_offsets, eos_token_id, 1)` on the
- `[T, 3]` int64 window, reshaped to `[T, 16]` by the caller.
-
- Preconditions: `head_vocab_sizes`, `head_offsets` on `token_history`'s
- device (raises otherwise; a cpu `layer_multipliers` is accepted and
- gives the identical result). Certified on an int64 `token_history`, the
- window `ngram_token_history` returns. Negative ids are accepted and hashed
- as ordinary ids, not padding. Longer histories and `num_query_tokens > 1`
- (a chunk's tail) are not certified.
- A single multiplier raises `ZeroDivisionError`. The four silent mistakes
- are `is_valid`'s.
-
- A wrong hash is a valid row: an extra shift, no EOS reset, an inclusive
- segment start, reversed multipliers, no offsets, swapped blocks all stay
- in range and inside the table. The tail slice is taken before the mix
- (bitwise equal to slicing after), and chunked prefill with the
- `ngram_size - 1` carried ids reproduces one-shot bitwise.
-
- CUDA-graph capture certified, bitwise equal to eager. Cost, not gated:
- 0.387 ms for 4096 x 16, 0.370 ms for an 8-row decode -- launch-bound.
-
- Hard limits
- token id max_id * max(layer_multipliers) < 2^63 boundary 389114 for
- this checkpoint's multipliers probed above it: accepted, no error,
- indices still inside the table (int64 wraps silently). The largest
- legal id 248319 gives 5.886e18. Not guarded: data, would sync.
- """
-
- def __call__(
- self,
- token_history: torch.Tensor,
- layer_multipliers: torch.Tensor,
- head_vocab_sizes: torch.Tensor,
- head_offsets: torch.Tensor,
- eos_token_id: int,
- num_query_tokens: int,
- ) -> torch.Tensor:
- ids = token_history.long()
- batch_size, history_len = ids.shape
- ngram_size = layer_multipliers.numel()
- total_heads = head_vocab_sizes.numel()
- heads_per_ngram = total_heads // (ngram_size - 1)
-
- # Segment bookkeeping over the FULL history: the most recent EOS
- # strictly before each position starts that position's segment.
- positions = torch.arange(history_len, device=ids.device, dtype=torch.long)
- minus_one = positions.new_full((), -1)
- eos_positions = torch.where(ids == eos_token_id, positions, minus_one)
- previous_eos_inclusive = torch.cummax(eos_positions, dim=1).values
- previous_eos = torch.cat(
- [eos_positions.new_full((batch_size, 1), -1), previous_eos_inclusive[:, :-1]],
- dim=1,
- )
- position_in_segment = positions.unsqueeze(0) - (previous_eos + 1)
-
- # Elementwise along the token axis from here on, so slicing to the
- # query tail now is bitwise equal to slicing after the modulo.
- keep_from = history_len - num_query_tokens
- eos_fill = ids.new_full((), eos_token_id)
- shifted = [ids[:, keep_from:]]
- for shift in range(1, ngram_size):
- source_positions = positions - shift
- gather_positions = source_positions.clamp_min(0).unsqueeze(0).expand(batch_size, -1)
- gathered = ids.gather(1, gather_positions)
- valid = (position_in_segment >= shift) & (source_positions.unsqueeze(0) >= 0)
- shifted.append(torch.where(valid, gathered, eos_fill)[:, keep_from:])
-
- blocks = []
- for ngram in range(2, ngram_size + 1):
- low = (ngram - 2) * heads_per_ngram
- high = low + heads_per_ngram
- mixed = shifted[0] * layer_multipliers[0]
- for position in range(1, ngram):
- mixed = torch.bitwise_xor(mixed, shifted[position] * layer_multipliers[position])
- blocks.append(
- torch.remainder(mixed.unsqueeze(-1), head_vocab_sizes[low:high].view(1, 1, -1))
- + head_offsets[low:high].view(1, 1, -1)
- )
- return torch.cat(blocks, dim=-1)
-
- def reference(
- self,
- token_history: torch.Tensor,
- layer_multipliers: torch.Tensor,
- head_vocab_sizes: torch.Tensor,
- head_offsets: torch.Tensor,
- eos_token_id: int,
- num_query_tokens: int,
- ) -> torch.Tensor:
- """Scalar Python-int arithmetic, one token and one head at a time, the
- segment start tracked by walking the row; no torch op computes it."""
- multipliers = layer_multipliers.tolist()
- sizes = head_vocab_sizes.tolist()
- offsets = head_offsets.tolist()
- ngram_size = len(multipliers)
- heads_per_ngram = len(sizes) // (ngram_size - 1)
- out = []
- for row in token_history.tolist():
- segment_start, last_eos = [], -1
- for position, token in enumerate(row):
- segment_start.append(last_eos + 1)
- if token == eos_token_id:
- last_eos = position
- row_out = []
- for position in range(len(row) - num_query_tokens, len(row)):
- context = [
- row[position - shift]
- if position - segment_start[position] >= shift and position - shift >= 0
- else eos_token_id
- for shift in range(ngram_size)
- ]
- heads = []
- for ngram in range(2, ngram_size + 1):
- mixed = context[0] * multipliers[0]
- for index in range(1, ngram):
- mixed ^= context[index] * multipliers[index]
- low = (ngram - 2) * heads_per_ngram
- heads.extend(
- mixed % sizes[low + h] + offsets[low + h] for h in range(heads_per_ngram)
- )
- row_out.append(heads)
- out.append(row_out)
- return torch.tensor(out, dtype=torch.long, device=token_history.device).reshape(
- token_history.shape[0], num_query_tokens, len(sizes)
- )
-
- def is_valid(
- self,
- token_history: torch.Tensor,
- layer_multipliers: torch.Tensor,
- head_vocab_sizes: torch.Tensor,
- head_offsets: torch.Tensor,
- eos_token_id: int,
- num_query_tokens: int,
- ) -> None:
- # Each measured to be accepted and to return plausible in-range rows.
- history_len = token_history.shape[-1]
- ngram_size = layer_multipliers.numel()
- total_heads = head_vocab_sizes.numel()
- assert 0 <= num_query_tokens <= history_len, (
- f"num_query_tokens must be in [0, history_len]; got {num_query_tokens} "
- f"against a history of {history_len} -- a value outside that range is "
- f"silently truncated to a shorter output rather than rejected"
- )
- assert layer_multipliers.dtype == torch.int64, (
- f"layer_multipliers must be int64; got {layer_multipliers.dtype} -- a "
- f"narrower dtype is accepted and silently hashes with truncated "
- f"multipliers"
- )
- assert ngram_size < 2 or total_heads % (ngram_size - 1) == 0, (
- f"head_vocab_sizes.numel()={total_heads} must be a multiple of "
- f"ngram_size-1={ngram_size - 1} -- the blocks are equal width, and a "
- f"remainder is silently dropped rather than rejected"
- )
- assert head_offsets.numel() == total_heads, (
- f"head_offsets and head_vocab_sizes must have the same length; got "
- f"{head_offsets.numel()} and {total_heads} -- an over-long head_offsets "
- f"is silently read for its first {total_heads} entries only"
- )
-
-
-ngram_hash = NgramHash()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/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
deleted file mode 100644
index c4b93e549701..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py
+++ /dev/null
@@ -1,384 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Decode-row bf16 projection launched with programmatic dependent launch (PDL),
-its weight streaming in while the kernel before it finishes.
-
-Tier 5 (CuTe DSL written for this target). At decode row counts a projection
-reads its weight once and its activation is a few KB, so the GEMV's time is the
-weight stream plus its launch and ramp. A normal launch cannot read a byte of
-that stream until the producer of the activation has exited. Launched with PDL,
-every CTA loads its weight rows -- static, so safe to read early -- into
-registers before `griddepcontrol.wait` and triggers its own dependents right
-before it; the activation is read, and the output written, only after the wait.
-
-One CTA computes `OUTS` output columns over the whole K for every row; element
-`k` of a row is held by thread `(k // VEC) % THREADS` as its vector
-`(k // VEC) // THREADS`, so each k-tile of loads is contiguous across a warp.
-Products and sums in fp32: a warp reduction, then a cross-warp one in shared
-memory; one rounding at the store.
-
-A K with a shared-memory tier (`SOUTS`) gives each CTA `SOUTS` more outputs whose
-weight rows it stages in shared memory with cp.async before the wait, each thread
-exactly the vectors it later reads (so `cp.async.wait_group` alone orders them). The
-2560 x 6144 projection then runs as one wave of 640 CTAs whose whole weight is on
-chip by the wait, instead of 1280 CTAs of which the last third start loading theirs
-only as the first ones exit.
-"""
-
-from __future__ import annotations
-
-import functools
-
-import cuda.bindings.driver as cuda_drv
-import cutlass
-import cutlass.cute as cute
-import torch
-from cutlass import const_expr
-from cutlass._mlir.dialects import llvm
-from cutlass.cute import experimental as cute_ext
-from cutlass.cute.nvgpu import cpasync
-from cutlass.cute.runtime import from_dlpack
-from cutlass.cutlass_dsl import dsl_user_op
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import (
- Arch,
- Cell,
- OpWrapper,
- true_fp32_matmul,
-)
-from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
-
-#: Largest row count one launch serves; every row re-reads the weight from registers.
-MAX_ROWS = 4
-#: (threads, outputs per CTA, elements per vector load) per K. Swept on GB300 at one row,
-#: graph replay over per-layer weights: K=2560 64 threads x 2 outputs; K=640 (80 16-byte
-#: vectors, which no warp multiple divides) 32 threads x 4 outputs on 8-byte vectors.
-TILING = {2560: (64, 2, 8), 6144: (128, 2, 8), 640: (32, 4, 4)}
-#: Outputs per CTA whose weight rows sit in shared memory (beside the `TILING` ones in
-#: registers). K=6144 swept behind the GDN recurrence and the sparse-attention merge
-#: (one row, graph replay): 2 + 2 a CTA beat 2 + 0, 2 + 6, 4 + 4 and 1 + 3.
-SOUTS = {6144: 2}
-_COMPILE_OPTIONS = "--ptxas-options -maxrregcount=64"
-_COMPILE_OPTIONS_SMEM = "--ptxas-options -maxrregcount=80"
-#: CTAs a multiprocessor must hold for a grid with a shared-memory tier: four, the
-#: 640-CTA grid's share of 160 SMs. Behind the sparse-attention merge (one row, graph
-#: replay) 4 runs the attention tail 13.5 us a layer against 14.6 at 1 and 15.6 at 2-3;
-#: the GDN tail is unchanged.
-MIN_BLOCKS_SMEM = 4
-
-
-def tiling(n: int, k: int) -> tuple[int, int, int, int]:
- """`(threads, outs, souts, vec)` for a `[*, K] x [K, N]` product: register and shared-memory
- outputs a CTA. Without a shared-memory tier that divides N the outputs shrink to divide N."""
- threads, outs, vec = TILING[k]
- souts = SOUTS.get(k, 0)
- if souts and n % (outs + souts) == 0:
- return threads, outs, souts, vec
- while n % outs:
- outs //= 2
- return threads, outs, 0, vec
-
-
-#: Before its PDL wait a CTA also has the weight rows of CTA `bidx + dist` brought into
-#: L2: behind a latency-bound kernel only part of a large grid's CTAs fit, and the later
-#: waves' weights would otherwise start streaming only after it. `dist` is about the CTAs
-#: resident before the wait: 1216 behind the hyper-connection kernels (measured), at
-#: least 912 (a 1280-CTA out projection's first wave) otherwise. Grids smaller than `dist`
-#: prefetch nothing.
-L2_DIST = 1216
-L2_DIST_MIN = 912
-
-
-def l2_dist(grid: int) -> int:
- return min(L2_DIST, max(grid // 2, L2_DIST_MIN))
-
-
-@dsl_user_op
-def _bulk_prefetch_l2(src, nbytes, *, loc=None, ip=None):
- llvm.inline_asm(
- None,
- [
- cutlass.Int64(src.toint(loc=loc, ip=ip)).ir_value(loc=loc, ip=ip),
- cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip),
- ],
- "cp.async.bulk.prefetch.L2.global [$0], $1;",
- "l,r",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
-
-
-class _GemvKernel:
- def __init__(self, *, rows, k, threads, outs, vec, pdl, dist=L2_DIST, souts=0):
- assert k % (threads * vec) == 0 and threads % 32 == 0
- self.rows, self.threads, self.outs, self.souts, self.vec = rows, threads, outs, souts, vec
- self.dist = dist
- self.k_tiles = k // (threads * vec)
- self.warps = threads // 32
- self.pdl = pdl
-
- def smem_bytes(self) -> int:
- red = self.rows * (self.outs + self.souts) * self.warps * 4
- if not self.souts:
- return red
- w = self.souts * self.k_tiles * self.threads * self.vec * 2
- return -(-w // 128) * 128 + -(-red // 128) * 128 + 128
-
- @cute.jit
- def __call__(self, mA, mW, mOut, stream: cuda_drv.CUstream):
- n = cute.size(mW, mode=[0])
- load_a = cute.make_copy_atom(
- cute.nvgpu.CopyG2ROp(),
- cutlass.BFloat16,
- num_bits_per_copy=16 * self.vec,
- load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS,
- )
- load_w = cute.make_copy_atom(
- cute.nvgpu.CopyG2ROp(),
- cutlass.BFloat16,
- num_bits_per_copy=16 * self.vec,
- load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING,
- )
- g2s = cute.make_copy_atom(
- cpasync.CopyG2SOp(), cutlass.BFloat16, num_bits_per_copy=16 * self.vec
- )
- self.kernel(mA, mW, mOut, load_a, load_w, g2s).launch(
- grid=[cute.ceil_div(n, self.outs + self.souts), 1, 1],
- block=[self.threads, 1, 1],
- smem=self.smem_bytes(),
- stream=stream,
- use_pdl=self.pdl,
- min_blocks_per_mp=MIN_BLOCKS_SMEM if self.souts else 1,
- )
-
- @cute.kernel
- def kernel(
- self, mA, mW, mOut, load_a: cute.CopyAtom, load_w: cute.CopyAtom, g2s: cute.CopyAtom
- ):
- tidx, _, _ = cute.arch.thread_idx()
- bidx, _, _ = cute.arch.block_idx()
- warp = cute.arch.warp_idx()
- rows: cutlass.Constexpr = self.rows
- outs: cutlass.Constexpr = self.outs
- threads: cutlass.Constexpr = self.threads
- vec: cutlass.Constexpr = self.vec
- k_tiles: cutlass.Constexpr = self.k_tiles
- warps: cutlass.Constexpr = self.warps
- souts: cutlass.Constexpr = self.souts
- tot: cutlass.Constexpr = outs + souts
- n0 = bidx * tot
-
- # (row, (vec, (threads, k_tiles))): thread t's vectors of a row, k-tile by k-tile.
- w_vec = cute.logical_divide(cute.logical_divide(mW, (None, vec)), (None, (None, threads)))
- a_vec = cute.logical_divide(cute.logical_divide(mA, (None, vec)), (None, (None, threads)))
- smem = cutlass.utils.SmemAllocator()
-
- # The weight is static: it streams before the wait, while the producer runs. The
- # shared-memory tier first ([output, k-tile, thread, vec]: a warp's vectors are
- # consecutive), so it lands while the register tier streams.
- if const_expr(souts > 0):
- s_w = smem.allocate_tensor(
- cutlass.BFloat16,
- cute.make_layout(
- (souts, k_tiles, threads, vec),
- stride=(k_tiles * threads * vec, threads * vec, vec, 1),
- ),
- byte_alignment=128,
- )
- for j in cutlass.range_constexpr(souts):
- s_row = w_vec[n0 + outs + j, (None, (tidx, None))]
- for kt in cutlass.range_constexpr(k_tiles):
- cute.copy(g2s, s_row[None, kt], s_w[j, kt, tidx, None])
- cute.arch.cp_async_commit_group()
- w_reg = cute.make_rmem_tensor(
- cute.make_layout((outs, k_tiles, vec), stride=(k_tiles * vec, vec, 1)), cutlass.BFloat16
- )
- for j in cutlass.range_constexpr(outs):
- w_row = w_vec[n0 + j, (None, (tidx, None))]
- for kt in cutlass.range_constexpr(k_tiles):
- cute.copy(load_w, w_row[None, kt], w_reg[j, kt, None])
- if const_expr(self.pdl):
- n_rows = cute.size(mW, mode=[0])
- k_len = cute.size(mW, mode=[1])
- if tidx == 0:
- if (bidx + self.dist + 1) * tot <= n_rows:
- _bulk_prefetch_l2(
- mW.iterator + (bidx + self.dist) * tot * k_len, tot * k_len * 2
- )
- cute.arch.griddepcontrol_launch_dependents()
- cute.arch.griddepcontrol_wait()
-
- acc = cute.make_rmem_tensor(cute.make_layout((rows, tot), stride=(tot, 1)), cutlass.Float32)
- acc.fill(0.0)
- a_reg = cute.make_rmem_tensor(
- cute.make_layout((k_tiles, vec), stride=(vec, 1)), cutlass.BFloat16
- )
- sw_reg = cute.make_rmem_tensor(cute.make_layout((vec,), stride=(1,)), cutlass.BFloat16)
- for r in cutlass.range_constexpr(rows):
- a_row = a_vec[r, (None, (tidx, None))]
- for kt in cutlass.range_constexpr(k_tiles):
- cute.copy(load_a, a_row[None, kt], a_reg[kt, None])
- for kt in cutlass.range_constexpr(k_tiles):
- for v in cutlass.range_constexpr(vec):
- av = a_reg[kt, v].to(cutlass.Float32)
- for j in cutlass.range_constexpr(outs):
- acc[r, j] = acc[r, j] + av * w_reg[j, kt, v].to(cutlass.Float32)
- if const_expr(souts > 0):
- if const_expr(r == 0):
- cute.arch.cp_async_wait_group(0)
- for j in cutlass.range_constexpr(souts):
- for kt in cutlass.range_constexpr(k_tiles):
- cute.autovec_copy(s_w[j, kt, tidx, None], sw_reg)
- for v in cutlass.range_constexpr(vec):
- acc[r, outs + j] = acc[r, outs + j] + a_reg[kt, v].to(
- cutlass.Float32
- ) * sw_reg[v].to(cutlass.Float32)
- for r in cutlass.range_constexpr(rows):
- for j in cutlass.range_constexpr(tot):
- acc[r, j] = cute.arch.warp_reduction_sum(acc[r, j])
- part = smem.allocate_tensor(
- cutlass.Float32,
- cute.make_layout((rows, tot, warps), stride=(tot * warps, warps, 1)),
- byte_alignment=16,
- )
- with cute.arch.elect_one():
- for r in cutlass.range_constexpr(rows):
- for j in cutlass.range_constexpr(tot):
- part[r, j, warp] = acc[r, j]
- cute.arch.sync_threads()
- if tidx < tot:
- for r in cutlass.range_constexpr(rows):
- total = cutlass.Float32(0.0)
- for w in cutlass.range_constexpr(warps):
- total = total + part[r, tidx, w]
- mOut[r, n0 + tidx] = total.to(cutlass.BFloat16)
-
-
-@functools.lru_cache(maxsize=None)
-def _compiled(rows, n, k, pdl):
- threads, outs, souts, vec = tiling(n, k)
- kern = _GemvKernel(
- rows=rows,
- k=k,
- threads=threads,
- outs=outs,
- vec=vec,
- pdl=pdl,
- dist=l2_dist(-(-n // (outs + souts))),
- souts=souts,
- )
- a = torch.empty(rows, k, dtype=torch.bfloat16, device="cuda")
- w = torch.empty(n, k, dtype=torch.bfloat16, device="cuda")
- out = torch.empty(rows, n, dtype=torch.bfloat16, device="cuda")
- stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream)
- return cute_ext.compile(
- kern,
- from_dlpack(a, assumed_align=32),
- from_dlpack(w, assumed_align=32),
- from_dlpack(out, assumed_align=32),
- stream,
- options=_COMPILE_OPTIONS_SMEM if souts else _COMPILE_OPTIONS,
- )
-
-
-def launch(mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool) -> torch.Tensor:
- """The launch behind `Bf16GemvPdl.__call__`. `pdl` is honoured unless
- `TRTLLM_ENABLE_PDL=0`."""
- m, k = mat_a.shape
- n = mat_b.shape[1]
- dev = mat_a.device
- pdl = pdl and get_env_enable_pdl()
- out = torch.empty((m, n), dtype=mat_a.dtype, device=dev)
- if m > 0:
- fn = _compiled(m, n, k, bool(pdl))
- stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream)
- fn(
- from_dlpack(mat_a, assumed_align=32),
- from_dlpack(mat_b.t(), assumed_align=32),
- from_dlpack(out, assumed_align=32),
- stream,
- )
- return out
-
-
-def _cell(why: str, m: int, n: int, k: int, pdl: bool = True) -> Cell:
- return Cell(why=why, spec=dict(m=m, n=n, k=k, pdl=pdl))
-
-
-class Bf16GemvPdl(OpWrapper):
- """`mat_a @ mat_b` in bf16 with fp32 accumulation, one PDL launch per call.
-
- A drop-in for `ops/gemm/low_m_bf16_mm` at decode row counts: `mat_a` is
- dense row-major `[M, K]` with `M <= MAX_ROWS`, `mat_b` the dense column-major
- `[K, N]` view of a stored `[N, K]` weight (its `.t()`), the result a fresh
- row-major `[M, N]` bf16 tensor. With `pdl` (and PDL available) the kernel is
- launched as a programmatic dependent: it may start before the kernel ahead of
- it on the stream exits, reads only its weight until `griddepcontrol.wait`,
- and lets the kernel after it launch early as well.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # Every shape is a qwen3.8-flash-next-nvfp4 tp1 decode projection; the target calls
- # it at one row (`_decode_projection`). Four cells: each `TILING` key is its own
- # compiled kernel the decode forward reaches, plus the N=1 edge.
- CELLS: tuple[Cell, ...] = (
- _cell("qwen3.8 decode M=1 GDN out / attention o, K=6144 shared-memory tier", 1, 2560, 6144),
- _cell("qwen3.8 decode M=1 GDN in-proj concat N=16480, L2 prefetch grid", 1, 16480, 2560),
- _cell("qwen3.8 decode M=1 shared down K=640", 1, 2560, 640),
- _cell("qwen3.8 decode M=1 shared gate N=1, outputs per CTA halved to 1", 1, 1, 2560),
- )
-
- note = """
- Metadata consumed: none. Preconditions: bf16 operands; `mat_a` dense
- row-major `[M, K]`, `1 <= M <= MAX_ROWS`; `mat_b` the dense column-major
- `[K, N]` view of a `[N, K]` weight (stride `(1, K)`); K a `TILING` key.
-
- PDL: with `pdl` the launch carries the programmatic-dependent attribute.
- Every read of `mat_a` and the store of the output come after
- `griddepcontrol.wait`, so a producer launched before this kernel is
- complete when they happen, whatever that producer is. Under a CUDA graph a
- cross-stream dependency is an ordinary edge and stays a full dependency.
-
- Compilation: CuTe DSL, JIT on the first call per (M, N, K, pdl); that call
- must happen before CUDA-graph capture -- the engine's warmup forward at each
- captured batch size does it.
-
- Numerics: fp32 products and sums, one rounding to bf16; only the summation
- order differs from the fp32 reference.
-
- Hard limits
- M rows per launch boundary MAX_ROWS=4 driven at 1 (cells)
- K a TILING key 2560, 6144, 640 all driven
- N any; the outputs per CTA halve until they divide N N=1 driven
-
- The test also replays a PDL chain of these GEMVs under a CUDA graph, each
- reading the one before it.
-
- Not certified: M = 2..4; PDL off (`pdl=False`, or PDL disabled in the
- environment), the same kernel without the attribute.
- """
-
- def __call__(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = True) -> torch.Tensor:
- return launch(mat_a, mat_b, pdl)
-
- def reference(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = True) -> torch.Tensor:
- """fp32-accumulated, TF32 off; `pdl` is a launch attribute and is ignored."""
- with true_fp32_matmul():
- return (mat_a.float() @ mat_b.float()).to(mat_a.dtype)
-
- def is_valid(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = True) -> None:
- assert mat_a.dtype == torch.bfloat16 and mat_b.dtype == torch.bfloat16, "bf16 operands only"
- assert mat_a.dim() == 2 and mat_a.is_contiguous(), "mat_a must be dense row-major [M, K]"
- assert 1 <= mat_a.shape[0] <= MAX_ROWS, f"M={mat_a.shape[0]} outside 1..{MAX_ROWS}"
- assert mat_b.shape[0] == mat_a.shape[1], "inner dimensions differ"
- assert mat_b.stride(0) == 1 and mat_b.stride(1) == mat_b.shape[0], (
- "mat_b must be dense column-major [K, N] (e.g. weight.t())"
- )
- assert mat_a.shape[1] in TILING, f"K={mat_a.shape[1]} has no tiling"
-
-
-bf16_gemv_pdl = Bf16GemvPdl()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/fp8_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/fp8_gemv.py
deleted file mode 100644
index ff6fa1208e36..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/fp8_gemv.py
+++ /dev/null
@@ -1,2016 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Decode-row W8A16 projection: `y[M, N] = x[M, K] @ dequant(W)^T` with FP8 e4m3 weights
-(128 x 128 block scales, fp32) and bf16 activations that are never quantized.
-
-CuTe DSL (inline PTX), one PDL launch, CUDA-graph capturable, M = 1 .. 32.
-
-Math: swap-AB warp MMA (`mma.sync.m16n8k16`, bf16 in, fp32 accumulate). The A operand is
-16 weight rows dequantized in registers with `cvt.rn.bf16x2.e4m3x2` (exact: every e4m3
-value is a bf16); B is up to 8 activation rows per n8 tile, read as raw bf16. Each 128-wide
-k-block is accumulated on its own (four independent MMA chains) and folded into the running
-fp32 sum times its block scale, so the only roundings are fp32 sums and the final bf16 store.
-
-Weight layout (`pack_weight`): 2 KB units of 16 rows x 128 k in exactly the order the MMA
-fragments read them (one conflict-free `ld.shared.v4` per lane per two k16 steps), stored
-k-block-major `[K/128][ceil(N/16)][2048]` or, for small N, tile-major
-`[ceil(N/16)][K/128][2048]`. K is permuted identically for both operands inside every
-128-block (lane quad t of k16 step j holds k = 32 (j // 2) + 8 t + 4 (j % 2) + 0..3), which
-leaves the dot product unchanged and makes every activation read a 16-byte vector.
-
-Work split (`Plan`): a CTA owns a contiguous range of 16-row tiles over its K range (all of K,
-or 1/ks of it in a `ks`-CTA cluster reduced over DSMEM). A weight-producer warp streams the
-CTA's units with `cp.async.bulk` into a ring of shared-memory slots (a slot = `kw` k-blocks of
-every tile of the CTA) completed on mbarriers; the static weight needs no PDL wait, so the ring
-fills while the kernel ahead finishes. Consumer warp (kw_i, tw_i) takes k-block kw_i of a slot
-for tiles tw_i, tw_i + tw, .... Activations are read only after `griddepcontrol.wait`, either
-straight into registers a few stages ahead (few stages per warp) or by an activation-producer
-warp with 16-byte `cp.async` into the same ring slots (long K). The `kw` k-block partials are
-summed in shared memory in a fixed order and stored as bf16: results are deterministic.
-"""
-
-from __future__ import annotations
-
-import functools
-from dataclasses import dataclass
-
-import cutlass
-import cutlass.cute as cute
-import torch
-from cutlass import const_expr
-from cutlass._mlir.dialects import llvm
-from cutlass.cute.runtime import from_dlpack, make_ptr
-from cutlass.cutlass_dsl import T, dsl_user_op
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import (
- Arch,
- Cell,
- OpWrapper,
- true_fp32_matmul,
-)
-
-TILE_N = 16
-BLOCK_K = 128
-UNIT = TILE_N * BLOCK_K # bytes of one (16-row, 128-k) unit
-#: Bytes of one NVFP4 (16-row, 128-k) unit: 1024 of e2m1 codes, then 128 of e4m3 scales.
-UNIT_W4 = TILE_N * BLOCK_K // 2 + TILE_N * BLOCK_K // 16
-FP8_MAX = 448.0
-MAX_ROWS = 32
-
-
-# ------------------------------------------------------------------------------------------
-# inline PTX helpers
-# ------------------------------------------------------------------------------------------
-
-
-def _ir(x, cls, loc, ip):
- return cls(x).ir_value(loc=loc, ip=ip)
-
-
-def _asm(ret, args, text, cons, loc, ip, side=False):
- return llvm.inline_asm(
- ret,
- args,
- text,
- cons,
- has_side_effects=side,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
-
-
-def _unpack(res, n, ty, loc, ip):
- return [
- ty(llvm.extractvalue(T.i32() if ty is cutlass.Int32 else T.f32(), res, [i], loc=loc, ip=ip))
- for i in range(n)
- ]
-
-
-@dsl_user_op
-def _lds4(saddr, *, loc=None, ip=None):
- res = _asm(
- llvm.StructType.get_literal([T.i32()] * 4),
- [_ir(saddr, cutlass.Int32, loc, ip)],
- "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];",
- "=r,=r,=r,=r,r",
- loc,
- ip,
- side=True,
- )
- return _unpack(res, 4, cutlass.Int32, loc, ip)
-
-
-@dsl_user_op
-def _lds2(saddr, *, loc=None, ip=None):
- res = _asm(
- llvm.StructType.get_literal([T.i32()] * 2),
- [_ir(saddr, cutlass.Int32, loc, ip)],
- "ld.shared.v2.b32 {$0, $1}, [$2];",
- "=r,=r,r",
- loc,
- ip,
- side=True,
- )
- return _unpack(res, 2, cutlass.Int32, loc, ip)
-
-
-@dsl_user_op
-def _lds_f32(saddr, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(
- T.f32(),
- [_ir(saddr, cutlass.Int32, loc, ip)],
- "ld.shared.f32 $0, [$1];",
- "=f,r",
- loc,
- ip,
- side=True,
- )
- )
-
-
-@dsl_user_op
-def _sts_f32(saddr, v, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [_ir(saddr, cutlass.Int32, loc, ip), _ir(v, cutlass.Float32, loc, ip)],
- "{ st.shared.f32 [$1], $2; mov.u32 $0, 0; }",
- "=r,r,f",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _ldg4p(gaddr, pred, *, loc=None, ip=None):
- """16 bytes from global (L1-cached, read-only) if `pred` != 0, else zeros."""
- res = _asm(
- llvm.StructType.get_literal([T.i32()] * 4),
- [_ir(gaddr, cutlass.Int64, loc, ip), _ir(pred, cutlass.Int32, loc, ip)],
- "{ .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.nc.v4.b32 {$0, $1, $2, $3}, [$4]; }",
- "=r,=r,=r,=r,l,r",
- loc,
- ip,
- side=True,
- )
- return _unpack(res, 4, cutlass.Int32, loc, ip)
-
-
-@dsl_user_op
-def _ldg_f32(gaddr, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(
- T.f32(),
- [_ir(gaddr, cutlass.Int64, loc, ip)],
- "ld.global.nc.f32 $0, [$1];",
- "=f,l",
- loc,
- ip,
- side=True,
- )
- )
-
-
-@dsl_user_op
-def _stg_bf16(gaddr, v, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [_ir(gaddr, cutlass.Int64, loc, ip), _ir(v, cutlass.Float32, loc, ip)],
- "{ .reg .b16 h; cvt.rn.bf16.f32 h, $2; st.global.b16 [$1], h; mov.u32 $0, 0; }",
- "=r,l,f",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _ldg_bf16p(gaddr, pred, *, loc=None, ip=None):
- """One bf16 from global as fp32 if `pred` != 0, else 0."""
- return cutlass.Float32(
- _asm(
- T.f32(),
- [_ir(gaddr, cutlass.Int64, loc, ip), _ir(pred, cutlass.Int32, loc, ip)],
- "{ .reg .pred p; .reg .b16 h; setp.ne.b32 p, $2, 0; mov.b16 h, 0; @p ld.global.b16 h, [$1]; "
- "cvt.f32.bf16 $0, h; }",
- "=f,l,r",
- loc,
- ip,
- side=True,
- )
- )
-
-
-@dsl_user_op
-def _ldg_s32p(gaddr, pred, dflt, *, loc=None, ip=None):
- """One int32 from global if `pred` != 0, else `dflt`."""
- return cutlass.Int32(
- _asm(
- T.i32(),
- [
- _ir(gaddr, cutlass.Int64, loc, ip),
- _ir(pred, cutlass.Int32, loc, ip),
- _ir(dflt, cutlass.Int32, loc, ip),
- ],
- "{ .reg .pred p; setp.ne.b32 p, $2, 0; mov.b32 $0, $3; @p ld.global.b32 $0, [$1]; }",
- "=r,l,r,r",
- loc,
- ip,
- side=True,
- )
- )
-
-
-@dsl_user_op
-def _ldg_u16p(gaddr, pred, *, loc=None, ip=None):
- """One 16-bit value from global, zero-extended, if `pred` != 0, else 0. Nothing in the asm
- reads it, so the thread runs on until the value's first use."""
- return cutlass.Int32(
- _asm(
- T.i32(),
- [_ir(gaddr, cutlass.Int64, loc, ip), _ir(pred, cutlass.Int32, loc, ip)],
- "{ .reg .pred p; setp.ne.b32 p, $2, 0; mov.b32 $0, 0; @p ld.global.u16 $0, [$1]; }",
- "=r,l,r",
- loc,
- ip,
- side=True,
- )
- )
-
-
-@dsl_user_op
-def _ldg_v2p(gaddr, pred, *, loc=None, ip=None):
- """8 bytes from global (read-only) if `pred` != 0, else zeros, as two words."""
- res = _asm(
- llvm.StructType.get_literal([T.i32()] * 2),
- [_ir(gaddr, cutlass.Int64, loc, ip), _ir(pred, cutlass.Int32, loc, ip)],
- "{ .reg .pred p; setp.ne.b32 p, $3, 0; mov.b32 $0, 0; mov.b32 $1, 0; "
- "@p ld.global.nc.v2.b32 {$0, $1}, [$2]; }",
- "=r,=r,l,r",
- loc,
- ip,
- side=True,
- )
- return _unpack(res, 2, cutlass.Int32, loc, ip)
-
-
-@dsl_user_op
-def _bf16_lo(v, *, loc=None, ip=None):
- """The bf16 in the low half of `v` as fp32."""
- return cutlass.Float32(
- _asm(T.f32(), [_ir(v, cutlass.Int32, loc, ip)], "shl.b32 $0, $1, 16;", "=f,r", loc, ip)
- )
-
-
-@dsl_user_op
-def _round_bf16(x, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(
- T.f32(),
- [_ir(x, cutlass.Float32, loc, ip)],
- "{ .reg .b16 h; cvt.rn.bf16.f32 h, $1; cvt.f32.bf16 $0, h; }",
- "=f,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _conv4_silu_bf16(c0, c1, c2, x, w01, w23, *, loc=None, ip=None):
- """`causal_conv1d_triton_update`'s arithmetic for one output of a width-4 window: the
- window `(c0, c1, c2, x)` (bf16 bits in the low half of each word) times the taps (bf16x2
- words `(w0 | w1 << 16, w2 | w3 << 16)`), each product rounded to bf16 (`mul.rn.bf16x2`),
- the four summed in fp32 in tap order from +0, then SiLU as
- `acc / (1 + ex2.approx(acc * -log2 e))` with `div.full.f32` -- the instruction sequence
- the Triton kernel compiles to. Returns the fp32 result (rounded to bf16 by the store)."""
- return cutlass.Float32(
- _asm(
- T.f32(),
- [_ir(v, cutlass.Int32, loc, ip) for v in (c0, c1, c2, x, w01, w23)],
- "{ .reg .b32 u, v; .reg .b16 p0, p1, p2, p3; .reg .f32 s, q, e; "
- "prmt.b32 u, $1, $2, 0x5410; prmt.b32 v, $3, $4, 0x5410; "
- "mul.rn.bf16x2 u, u, $5; mul.rn.bf16x2 v, v, $6; "
- "mov.b32 {p0, p1}, u; mov.b32 {p2, p3}, v; "
- "cvt.f32.bf16 s, p0; add.rn.f32 s, s, 0f00000000; "
- "cvt.f32.bf16 q, p1; add.rn.f32 s, q, s; "
- "cvt.f32.bf16 q, p2; add.rn.f32 s, q, s; "
- "cvt.f32.bf16 q, p3; add.rn.f32 s, q, s; "
- "mul.rn.f32 e, s, 0fBFB8AA3B; ex2.approx.f32 e, e; add.rn.f32 e, e, 0f3F800000; "
- "div.full.f32 $0, s, e; }",
- "=f,r,r,r,r,r,r",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _bf16_bits(x, *, loc=None, ip=None):
- """`x` rounded to bf16, as its bits in the low half of a word."""
- return cutlass.Int32(
- _asm(
- T.i32(),
- [_ir(x, cutlass.Float32, loc, ip)],
- "{ .reg .b16 h; cvt.rn.bf16.f32 h, $1; cvt.u32.u16 $0, h; }",
- "=r,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _stg_b16(gaddr, v, *, loc=None, ip=None):
- """The low half of `v` to global."""
- _asm(
- T.i32(),
- [_ir(gaddr, cutlass.Int64, loc, ip), _ir(v, cutlass.Int32, loc, ip)],
- "{ st.global.b16 [$1], $2; mov.u32 $0, 0; }",
- "=r,l,r",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _dq_chunk(w0, w1, w2, w3, *, loc=None, ip=None):
- """Four words of e4m3 (rows g / g+8 of k16 step 2c, then of step 2c+1) -> the eight bf16x2
- A registers of the two steps: (R0, R1, R2, R3) = (g lo, g+8 lo, g hi, g+8 hi)."""
- res = _asm(
- llvm.StructType.get_literal([T.i32()] * 8),
- [_ir(v, cutlass.Int32, loc, ip) for v in (w0, w1, w2, w3)],
- "{ .reg .b16 a0, a1, b0, b1, c0, c1, d0, d1; "
- "mov.b32 {a0, a1}, $8; mov.b32 {b0, b1}, $9; mov.b32 {c0, c1}, $10; mov.b32 {d0, d1}, $11; "
- "cvt.rn.bf16x2.e4m3x2 $0, a0; cvt.rn.bf16x2.e4m3x2 $1, b0; "
- "cvt.rn.bf16x2.e4m3x2 $2, a1; cvt.rn.bf16x2.e4m3x2 $3, b1; "
- "cvt.rn.bf16x2.e4m3x2 $4, c0; cvt.rn.bf16x2.e4m3x2 $5, d0; "
- "cvt.rn.bf16x2.e4m3x2 $6, c1; cvt.rn.bf16x2.e4m3x2 $7, d1; }",
- "=r,=r,=r,=r,=r,=r,=r,=r,r,r,r,r",
- loc,
- ip,
- )
- return _unpack(res, 8, cutlass.Int32, loc, ip)
-
-
-@dsl_user_op
-def _mma(a, b, c, *, loc=None, ip=None):
- """D = A B + C on m16n8k16, bf16 A/B, fp32 C/D; a: 4 regs, b: 2 regs, c: 4 f32."""
- res = _asm(
- llvm.StructType.get_literal([T.f32()] * 4),
- [_ir(v, cutlass.Int32, loc, ip) for v in a]
- + [_ir(v, cutlass.Int32, loc, ip) for v in b]
- + [_ir(v, cutlass.Float32, loc, ip) for v in c],
- "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
- "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};",
- "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
- loc,
- ip,
- )
- return _unpack(res, 4, cutlass.Float32, loc, ip)
-
-
-@dsl_user_op
-def _scales4(sw, *, loc=None, ip=None):
- """A unit's scale word `(s_g(2t), s_g8(2t), s_g(2t+1), s_g8(2t+1))` (e4m3) -> four f16x2
- registers with each scale in both halves: (row g, steps 0..3), (row g + 8, steps 0..3),
- (row g, steps 4..7), (row g + 8, steps 4..7)."""
- res = _asm(
- llvm.StructType.get_literal([T.i32()] * 4),
- [_ir(sw, cutlass.Int32, loc, ip)],
- "{ .reg .b16 lo, hi, p0, p1, q0, q1; .reg .b32 p, q; "
- "mov.b32 {lo, hi}, $4; "
- "cvt.rn.f16x2.e4m3x2 p, lo; cvt.rn.f16x2.e4m3x2 q, hi; "
- "mov.b32 {p0, p1}, p; mov.b32 {q0, q1}, q; "
- "mov.b32 $0, {p0, p0}; mov.b32 $1, {p1, p1}; mov.b32 $2, {q0, q0}; "
- "mov.b32 $3, {q1, q1}; }",
- "=r,=r,=r,=r,r",
- loc,
- ip,
- )
- return _unpack(res, 4, cutlass.Int32, loc, ip)
-
-
-@dsl_user_op
-def _dq_step(w, sg, sg8, *, loc=None, ip=None):
- """One step's code word (bytes: row g k0|k1, row g+8 k0|k1, row g k2|k3, row g+8 k2|k3)
- -> the four f16x2 A registers, each times its row's block scale (exact)."""
- res = _asm(
- llvm.StructType.get_literal([T.i32()] * 4),
- [_ir(v, cutlass.Int32, loc, ip) for v in (w, sg, sg8)],
- "{ .reg .b8 b0, b1, b2, b3; .reg .b32 h0, h1, h2, h3; "
- "mov.b32 {b0, b1, b2, b3}, $4; "
- "cvt.rn.f16x2.e2m1x2 h0, b0; cvt.rn.f16x2.e2m1x2 h1, b1; "
- "cvt.rn.f16x2.e2m1x2 h2, b2; cvt.rn.f16x2.e2m1x2 h3, b3; "
- "mul.rn.f16x2 $0, h0, $5; mul.rn.f16x2 $1, h1, $6; "
- "mul.rn.f16x2 $2, h2, $5; mul.rn.f16x2 $3, h3, $6; }",
- "=r,=r,=r,=r,r,r,r",
- loc,
- ip,
- )
- return _unpack(res, 4, cutlass.Int32, loc, ip)
-
-
-@dsl_user_op
-def _bf16x2_to_f16x2(v, *, loc=None, ip=None):
- """bf16x2 -> f16x2 (round to nearest, saturating at +-65504)."""
- return cutlass.Int32(
- _asm(
- T.i32(),
- [_ir(v, cutlass.Int32, loc, ip)],
- "{ .reg .b32 lo, hi; .reg .f32 flo, fhi; "
- "shl.b32 lo, $1, 16; and.b32 hi, $1, 0xFFFF0000; "
- "mov.b32 flo, lo; mov.b32 fhi, hi; "
- "cvt.rn.satfinite.f16x2.f32 $0, fhi, flo; }",
- "=r,r",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _mma_f16(a, b, c, *, loc=None, ip=None):
- """D = A B + C on m16n8k16, f16 A/B, fp32 C/D; a: 4 regs, b: 2 regs, c: 4 f32."""
- res = _asm(
- llvm.StructType.get_literal([T.f32()] * 4),
- [_ir(v, cutlass.Int32, loc, ip) for v in a]
- + [_ir(v, cutlass.Int32, loc, ip) for v in b]
- + [_ir(v, cutlass.Float32, loc, ip) for v in c],
- "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
- "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};",
- "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
- loc,
- ip,
- )
- return _unpack(res, 4, cutlass.Float32, loc, ip)
-
-
-@dsl_user_op
-def _mbar_init(addr, count, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [_ir(addr, cutlass.Int32, loc, ip), _ir(count, cutlass.Int32, loc, ip)],
- "{ mbarrier.init.shared::cta.b64 [$1], $2; mov.u32 $0, 0; }",
- "=r,r,r",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _mbar_init_fence(*, loc=None, ip=None):
- _asm(
- T.i32(),
- [],
- "{ fence.mbarrier_init.release.cluster; mov.u32 $0, 0; }",
- "=r",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _mbar_expect_tx(addr, nbytes, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [_ir(addr, cutlass.Int32, loc, ip), _ir(nbytes, cutlass.Int32, loc, ip)],
- "{ mbarrier.arrive.expect_tx.shared::cta.b64 _, [$1], $2; mov.u32 $0, 0; }",
- "=r,r,r",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _mbar_arrive(addr, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [_ir(addr, cutlass.Int32, loc, ip)],
- "{ mbarrier.arrive.shared::cta.b64 _, [$1]; mov.u32 $0, 0; }",
- "=r,r",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _mbar_wait(addr, parity, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [_ir(addr, cutlass.Int32, loc, ip), _ir(parity, cutlass.Int32, loc, ip)],
- "{ .reg .pred p; WAIT: mbarrier.try_wait.parity.shared::cta.b64 p, [$1], $2; @!p bra WAIT; "
- "mov.u32 $0, 0; }",
- "=r,r,r",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _bulk_g2s(saddr, gaddr, nbytes, mbar, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [
- _ir(saddr, cutlass.Int32, loc, ip),
- _ir(gaddr, cutlass.Int64, loc, ip),
- _ir(nbytes, cutlass.Int32, loc, ip),
- _ir(mbar, cutlass.Int32, loc, ip),
- ],
- "{ cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [$1], [$2], $3, [$4]; "
- "mov.u32 $0, 0; }",
- "=r,r,l,r,r",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _cp16(saddr, gaddr, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [_ir(saddr, cutlass.Int32, loc, ip), _ir(gaddr, cutlass.Int64, loc, ip)],
- "{ cp.async.cg.shared.global [$1], [$2], 16; mov.u32 $0, 0; }",
- "=r,r,l",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _cp_async_mbar_arrive(mbar, *, loc=None, ip=None):
- """Arrive on `mbar` once this thread's earlier cp.async copies have landed."""
- _asm(
- T.i32(),
- [_ir(mbar, cutlass.Int32, loc, ip)],
- "{ cp.async.mbarrier.arrive.noinc.shared::cta.b64 [$1]; mov.u32 $0, 0; }",
- "=r,r",
- loc,
- ip,
- side=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(),
- [_ir(addr, cutlass.Int32, loc, ip), _ir(rank, cutlass.Int32, loc, ip)],
- "mapa.shared::cluster.u32 $0, $1, $2;",
- "=r,r,r",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _st_dsmem_f32(addr, v, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [_ir(addr, cutlass.Int32, loc, ip), _ir(v, cutlass.Float32, loc, ip)],
- "{ st.shared::cluster.f32 [$1], $2; mov.u32 $0, 0; }",
- "=r,r,f",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _cluster_arrive_release(*, loc=None, ip=None):
- _asm(
- T.i32(),
- [],
- "{ barrier.cluster.arrive.release.aligned; mov.u32 $0, 0; }",
- "=r",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _cluster_arrive_relaxed(*, loc=None, ip=None):
- _asm(
- T.i32(),
- [],
- "{ barrier.cluster.arrive.relaxed.aligned; mov.u32 $0, 0; }",
- "=r",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _cluster_wait(*, loc=None, ip=None):
- _asm(T.i32(), [], "{ barrier.cluster.wait.aligned; mov.u32 $0, 0; }", "=r", loc, ip, side=True)
-
-
-@dsl_user_op
-def _cluster_wait_acquire(*, loc=None, ip=None):
- _asm(
- T.i32(),
- [],
- "{ barrier.cluster.wait.acquire.aligned; mov.u32 $0, 0; }",
- "=r",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _bulk_prefetch_l2(gaddr, nbytes, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [_ir(gaddr, cutlass.Int64, loc, ip), _ir(nbytes, cutlass.Int32, loc, ip)],
- "{ cp.async.bulk.prefetch.L2.global [$1], $2; mov.u32 $0, 0; }",
- "=r,l,r",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _bar_sync(bar_id, nthreads, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [_ir(bar_id, cutlass.Int32, loc, ip), _ir(nthreads, cutlass.Int32, loc, ip)],
- "{ bar.sync $1, $2; mov.u32 $0, 0; }",
- "=r,r,r",
- loc,
- ip,
- side=True,
- )
-
-
-# ------------------------------------------------------------------------------------------
-# kernel
-# ------------------------------------------------------------------------------------------
-
-
-@dataclass(frozen=True)
-class Plan:
- """Launch shape. `grid` CTAs, each a contiguous range of <= ceil(tiles / grid) 16-row
- tiles; `kw` x `tw` consumer warps (k-blocks per ring slot x tile groups); `slots` ring
- slots; `l2_ahead` slots past the ring prefetched into L2 at the start."""
-
- grid: int
- kw: int
- tw: int
- slots: int
- l2_ahead: int = 0 # rolling window of slots past the ring prefetched into L2
- xreg: int = 0 # activations: 0 auto, 1 registers, 2 through the ring
- xdepth: int = 0 # register mode: stages of activations loaded ahead (0: auto)
- ks: int = 1 # cluster K split: `ks` CTAs share a tile range, each 1/ks of K, reduced over DSMEM
-
-
-#: Most activation-fragment registers a consumer may hold to take the register path.
-XREG_MAX = 64
-XRS = 320 # bytes of one activation row slice (128 bf16 + 64 B pad: conflict-free fragment reads)
-
-
-class _Fp8Gemv:
- def __init__(
- self, *, n, k, m, plan: Plan, pdl, tile_major, conv_dim=0, pad_slot=-1, conv_p=0, w4=False
- ):
- assert k % BLOCK_K == 0 and 1 <= m <= MAX_ROWS
- self.n, self.k, self.m, self.pdl = n, k, m, pdl
- # NVFP4 units (`w4`, see `pack_weight_w4`) in place of FP8 ones: f16 MMA over the
- # activation rounded to f16, the global scale applied to the fp32 partials.
- self.w4 = w4
- self.unit = UNIT_W4 if w4 else UNIT
- self.tile_major = tile_major
- # conv1d epilogue (`conv_dim` > 0): columns below conv_dim of a non-pad row get the
- # causal conv step (window of 3, 4 taps) + SiLU; see `fp8_gemv_conv1d`. With
- # `conv_p` > 1 the rows are m / conv_p requests of conv_p consecutive tokens each, and
- # the epilogue runs the multi-token step with per-token window snapshots; see
- # `fp8_gemv_conv1d_verify`.
- self.conv_dim, self.pad_slot, self.conv_p = conv_dim, pad_slot, conv_p
- assert conv_dim == 0 or plan.ks == 1
- assert conv_p == 0 or (conv_dim > 0 and conv_p >= 2 and m % conv_p == 0), (conv_p, m)
- self.nreq = m // conv_p if conv_p else 0
- self.nt8 = -(-m // 8)
- self.n_tiles = -(-n // TILE_N)
- self.nkb = k // BLOCK_K
- self.grid = plan.grid
- self.ks = plan.ks
- assert plan.grid % self.ks == 0 and self.nkb % self.ks == 0, (plan.grid, self.ks, self.nkb)
- self.groups = plan.grid // self.ks
- self.nkb_c = self.nkb // self.ks # k-blocks per CTA
- self.rt = -(-self.n_tiles // self.groups)
- self.kw, self.tw, self.slots, self.l2_ahead = plan.kw, plan.tw, plan.slots, plan.l2_ahead
- assert self.nkb_c % self.kw == 0, (self.nkb_c, self.kw)
- self.n_stages = self.nkb_c // self.kw
- self.slots = min(self.slots, self.n_stages)
- self.l2_ahead = max(0, min(self.l2_ahead, self.n_stages - self.slots))
- self.rtw = -(-self.rt // self.tw)
- self.cw = self.kw * self.tw
- # Activations: with few stages a warp loads all of its fragments into registers right
- # after the PDL wait (lowest latency); otherwise they stream through the ring slots.
- self.xreg = self.n_stages * self.nt8 * 16 <= XREG_MAX if plan.xreg == 0 else plan.xreg == 1
- # register mode: stages of activation fragments in flight ahead of the one in use
- self.xdepth = plan.xdepth or max(1, XREG_MAX // (self.nt8 * 16) - 1)
- self.xdepth = min(self.xdepth, self.n_stages)
- # + weight producer warp (+ activation producer warp)
- self.threads = (self.cw + (1 if self.xreg else 2)) * 32
- self.w_bytes = self.kw * self.rt * self.unit
- self.x_bytes = 0 if self.xreg else self.kw * m * XRS
- self.slot_bytes = -(-(self.w_bytes + self.x_bytes) // 128) * 128
- self.np = self.nt8 * 8
- self.o_red = self.slots * self.slot_bytes
- red_bytes = self.kw * self.rt * self.np * TILE_N * 4
- # block scales of the (at most nsb) 128-row blocks the CTA's tiles touch
- self.nsb = 0 if w4 else self.rt // 8 + 2
- self.o_sc = self.o_red + -(-red_bytes // 128) * 128
- self.o_bar = self.o_sc + -(-(self.nsb * self.nkb * 4) // 128) * 128
- self.o_cred = self.o_bar + 3 * self.slots * 8
- cred_bytes = self.ks * self.rt * m * TILE_N * 4 if self.ks > 1 else 0
- self.smem = self.o_cred + cred_bytes + 128
-
- @cute.jit
- def __call__(
- self,
- out_t: cute.Tensor,
- x_ptr: cute.Pointer,
- w_ptr: cute.Pointer,
- s_ptr: cute.Pointer,
- x_stride: cutlass.Int32,
- cs_ptr: cute.Pointer,
- cw_ptr: cute.Pointer,
- idx_ptr: cute.Pointer,
- cs_s0: cutlass.Int32,
- cs_s1: cutlass.Int32,
- cs_s2: cutlass.Int32,
- iw_ptr: cute.Pointer,
- iidx_ptr: cute.Pointer,
- iw_s0: cutlass.Int32,
- iw_s1: cutlass.Int32,
- iw_s2: cutlass.Int32,
- iw_s3: cutlass.Int32,
- stream,
- ):
- args = (
- out_t.iterator.toint(),
- x_ptr.toint(),
- w_ptr.toint(),
- s_ptr.toint(),
- x_stride,
- cs_ptr.toint(),
- cw_ptr.toint(),
- idx_ptr.toint(),
- cs_s0,
- cs_s1,
- cs_s2,
- iw_ptr.toint(),
- iidx_ptr.toint(),
- iw_s0,
- iw_s1,
- iw_s2,
- iw_s3,
- )
- if const_expr(self.ks > 1):
- self.kernel(*args).launch(
- grid=[self.grid, 1, 1],
- block=[self.threads, 1, 1],
- smem=self.smem,
- stream=stream,
- cluster=[self.ks, 1, 1],
- use_pdl=self.pdl,
- )
- else:
- self.kernel(*args).launch(
- grid=[self.grid, 1, 1],
- block=[self.threads, 1, 1],
- smem=self.smem,
- stream=stream,
- use_pdl=self.pdl,
- )
-
- @cute.jit
- def _issue_w(self, dst, w_base, kbs, t0, nt, mbar):
- """Bulk copies of a stage's units (k-blocks kbs .. kbs + kw, tiles t0 .. t0 + nt)."""
- kw: cutlass.Constexpr = self.kw
- rt: cutlass.Constexpr = self.rt
- unit: cutlass.Constexpr = self.unit
- if const_expr(self.tile_major):
- for tl in cutlass.range_constexpr(rt):
- if tl < nt:
- _bulk_g2s(
- dst + tl * (kw * unit),
- w_base + (cutlass.Int64(t0 + tl) * self.nkb + kbs) * unit,
- kw * unit,
- mbar,
- )
- else:
- for j in cutlass.range_constexpr(kw):
- _bulk_g2s(
- dst + j * (rt * unit),
- w_base + (cutlass.Int64(kbs + j) * self.n_tiles + t0) * unit,
- nt * unit,
- mbar,
- )
-
- @cute.jit
- def _prefetch_w(self, w_base, kbs, t0, nt):
- kw: cutlass.Constexpr = self.kw
- rt: cutlass.Constexpr = self.rt
- unit: cutlass.Constexpr = self.unit
- if const_expr(self.tile_major):
- for tl in cutlass.range_constexpr(rt):
- if tl < nt:
- _bulk_prefetch_l2(
- w_base + (cutlass.Int64(t0 + tl) * self.nkb + kbs) * unit, kw * unit
- )
- else:
- for j in cutlass.range_constexpr(kw):
- _bulk_prefetch_l2(
- w_base + (cutlass.Int64(kbs + j) * self.n_tiles + t0) * unit, nt * unit
- )
-
- @cute.jit
- def _compute(self, acc, cur, slot, kb, kwi, twi, t0, nt, b0, scs, lane):
- """This warp's tiles of k-block `kb` (in ring slot `slot`) times the activation fragments
- `cur`, folded into `acc` with the block scale."""
- kw: cutlass.Constexpr = self.kw
- tw: cutlass.Constexpr = self.tw
- rt: cutlass.Constexpr = self.rt
- nt8: cutlass.Constexpr = self.nt8
- nkb: cutlass.Constexpr = self.nkb
- unit: cutlass.Constexpr = self.unit
- for i in cutlass.range_constexpr(self.rtw):
- tl = twi + i * tw
- if const_expr(self.w4):
- if tl < nt:
- self._compute_w4(acc, cur, slot, i, tl, kwi, lane)
- elif tl < nt:
- sc = _lds_f32(scs + (((t0 + tl) // 8 - b0) * nkb + kb) * 4)
- if const_expr(self.tile_major):
- ua = slot + (tl * kw + kwi) * unit + lane * 16
- else:
- ua = slot + (kwi * rt + tl) * unit + lane * 16
- wv = [_lds4(ua + c * 512) for c in range(4)]
- # one accumulator chain per 16-byte chunk (two k16 steps): 4-way ILP
- zero = cutlass.Float32(0.0)
- ch = [[[zero] * 4 for _ in range(nt8)] for _ in range(4)]
- for c in cutlass.range_constexpr(4):
- a = _dq_chunk(wv[c][0], wv[c][1], wv[c][2], wv[c][3])
- for st in cutlass.range_constexpr(2):
- for j in cutlass.range_constexpr(nt8):
- ch[c][j] = _mma(
- a[4 * st : 4 * st + 4],
- [cur[j][4 * c + 2 * st], cur[j][4 * c + 2 * st + 1]],
- ch[c][j],
- )
- for j in cutlass.range_constexpr(nt8):
- for e in cutlass.range_constexpr(4):
- part = (ch[0][j][e] + ch[1][j][e]) + (ch[2][j][e] + ch[3][j][e])
- acc[i, j, e] = acc[i, j, e] + part * sc
-
- @cute.jit
- def _compute_w4(self, acc, cur, slot, i, tl, kwi, lane):
- """`_compute` of one tile over an NVFP4 unit (`pack_weight_w4`): each step's e2m1 codes
- times their 16-k block scales in f16 (exact), four f16 MMA chains, `cur` f16."""
- kw: cutlass.Constexpr = self.kw
- rt: cutlass.Constexpr = self.rt
- nt8: cutlass.Constexpr = self.nt8
- unit: cutlass.Constexpr = self.unit
- if const_expr(self.tile_major):
- ua = slot + (tl * kw + kwi) * unit
- else:
- ua = slot + (kwi * rt + tl) * unit
- wv = [_lds4(ua + h * 512 + lane * 16) for h in range(2)]
- sw = _lds2(ua + 1024 + ((lane // 4) * 2 + (lane % 4) // 2) * 8)
- # (row g, row g + 8) block scales of chunks 0..3, as f16x2 pairs
- sc = _scales4(sw[0]) + _scales4(sw[1])
- zero = cutlass.Float32(0.0)
- ch = [[[zero] * 4 for _ in range(nt8)] for _ in range(4)]
- for c in cutlass.range_constexpr(4):
- for st in cutlass.range_constexpr(2):
- a = _dq_step(wv[c // 2][2 * (c % 2) + st], sc[2 * c], sc[2 * c + 1])
- for j in cutlass.range_constexpr(nt8):
- ch[c][j] = _mma_f16(
- a, [cur[j][4 * c + 2 * st], cur[j][4 * c + 2 * st + 1]], ch[c][j]
- )
- for j in cutlass.range_constexpr(nt8):
- for e in cutlass.range_constexpr(4):
- part = (ch[0][j][e] + ch[1][j][e]) + (ch[2][j][e] + ch[3][j][e])
- acc[i, j, e] = acc[i, j, e] + part
-
- @cute.jit
- def _issue_x(self, xs, src, xh, xc, x_stride):
- """cp.async of one stage's activations into the slot's activation area `xs`."""
- m: cutlass.Constexpr = self.m
- n_rc: cutlass.Constexpr = self.kw * self.m
- for i in cutlass.range_constexpr(-(-n_rc // 2)):
- rc = xh + 2 * i
- xj = rc // m
- xr = rc - xj * m
- ok = rc < n_rc
- if const_expr(n_rc % 2 == 0):
- ok = rc >= 0
- if ok:
- _cp16(
- xs + rc * XRS + xc * 16,
- src + cutlass.Int64(xr) * cutlass.Int64(x_stride) * 2 + cutlass.Int64(xj * 256),
- )
-
- @cute.kernel
- def kernel(
- self,
- out_base: cutlass.Int64,
- x_base: cutlass.Int64,
- w_base: cutlass.Int64,
- s_base: cutlass.Int64,
- x_stride: cutlass.Int32,
- cs_base: cutlass.Int64,
- cw_base: cutlass.Int64,
- idx_base: cutlass.Int64,
- cs_s0: cutlass.Int32,
- cs_s1: cutlass.Int32,
- cs_s2: cutlass.Int32,
- iw_base: cutlass.Int64,
- iidx_base: cutlass.Int64,
- iw_s0: cutlass.Int32,
- iw_s1: cutlass.Int32,
- iw_s2: cutlass.Int32,
- iw_s3: cutlass.Int32,
- ):
- tidx, _, _ = cute.arch.thread_idx()
- bidx, _, _ = cute.arch.block_idx()
- warp = cute.arch.make_warp_uniform(cute.arch.warp_idx())
- lane = cute.arch.lane_idx()
- g = lane // 4
- t = lane % 4
- m: cutlass.Constexpr = self.m
- n_tiles: cutlass.Constexpr = self.n_tiles
- nkb: cutlass.Constexpr = self.nkb
- rt: cutlass.Constexpr = self.rt
- kw: cutlass.Constexpr = self.kw
- tw: cutlass.Constexpr = self.tw
- cw: cutlass.Constexpr = self.cw
- slots: cutlass.Constexpr = self.slots
- n_stages: cutlass.Constexpr = self.n_stages
- slot_bytes: cutlass.Constexpr = self.slot_bytes
- w_bytes: cutlass.Constexpr = self.w_bytes
- nt8: cutlass.Constexpr = self.nt8
- np_: cutlass.Constexpr = self.np
- rtw: cutlass.Constexpr = self.rtw
- n_out: cutlass.Constexpr = self.n
- l2a: cutlass.Constexpr = self.l2_ahead
-
- smem = cutlass.utils.SmemAllocator()
- base_t = smem.allocate_tensor(
- cutlass.Int8, cute.make_layout((self.smem - 128,), stride=(1,)), byte_alignment=128
- )
- sbase = cutlass.Int32(base_t.iterator.toint())
- ring = sbase
- red = sbase + self.o_red
- scs = sbase + self.o_sc
- bars = sbase + self.o_bar # wfull[slots], xfull[slots], empty[slots]
-
- ks: cutlass.Constexpr = self.ks
- groups: cutlass.Constexpr = self.groups
- grp = bidx // ks
- rank = bidx % ks
- kb0 = rank * self.nkb_c
- t0 = (grp * n_tiles) // groups
- nt = ((grp + 1) * n_tiles) // groups - t0
- b0 = t0 // 8
- n_blocks: cutlass.Constexpr = -(-self.n // 128)
- n_sc: cutlass.Constexpr = self.nsb * nkb
- if const_expr(self.w4):
- # the NVFP4 global scale: `s_base` points at it
- gsc = _ldg_f32(s_base)
- for q in cutlass.range_constexpr(-(-n_sc // self.threads)):
- qi = tidx + q * self.threads
- if qi < n_sc:
- blk = b0 + qi // nkb
- if blk < n_blocks:
- _sts_f32(
- scs + qi * 4, _ldg_f32(s_base + (cutlass.Int64(blk) * nkb + qi % nkb) * 4)
- )
-
- if warp == cw:
- if lane == 0:
- for sl in cutlass.range_constexpr(slots):
- _mbar_init(bars + 8 * sl, 1)
- _mbar_init(bars + 8 * (slots + sl), 32)
- _mbar_init(bars + 8 * (2 * slots + sl), cw)
- _mbar_init_fence()
- cute.arch.sync_threads()
- if const_expr(self.ks > 1):
- _cluster_arrive_relaxed() # matched by the wait before the first DSMEM store
- if const_expr(self.pdl):
- cute.arch.griddepcontrol_launch_dependents()
-
- if warp == cw:
- # ---- weight producer (static data: no PDL wait). Slot layout: unit (j, tile) at
- # (j * rt + tile) * unit for a k-block-major weight, (tile * kw + j) * unit for a
- # tile-major one; one bulk copy per k-block row resp. per tile. ----
- if lane == 0:
- for ps in cutlass.range(n_stages):
- psl = ps % slots
- if ps >= slots:
- _mbar_wait(bars + 8 * (2 * slots + psl), ((ps // slots) - 1) & 1)
- if const_expr(l2a > 0):
- # rolling window: stages slots .. slots + l2a - 1 at the start, then one
- # stage l2a past each refill
- if ps == 0:
- for la in cutlass.range_constexpr(slots, slots + l2a):
- self._prefetch_w(w_base, kb0 + la * kw, t0, nt)
- if ps >= slots:
- if ps + l2a < n_stages:
- self._prefetch_w(w_base, kb0 + (ps + l2a) * kw, t0, nt)
- _mbar_expect_tx(bars + 8 * psl, kw * nt * self.unit)
- self._issue_w(
- ring + psl * slot_bytes, w_base, kb0 + ps * kw, t0, nt, bars + 8 * psl
- )
- elif warp == cw + 1:
- # ---- activation producer (ring mode): after the PDL wait, a stage's rows into its slot ----
- if const_expr(self.pdl):
- cute.arch.griddepcontrol_wait()
- xh = lane // 16
- xc = lane % 16
- x_lane = x_base + cutlass.Int64(xc * 16)
- for xs_ in cutlass.range(n_stages):
- xsl = xs_ % slots
- if xs_ >= slots:
- _mbar_wait(bars + 8 * (2 * slots + xsl), ((xs_ // slots) - 1) & 1)
- self._issue_x(
- ring + xsl * slot_bytes + w_bytes,
- x_lane + cutlass.Int64((kb0 + xs_ * kw) * 256),
- xh,
- xc,
- x_stride,
- )
- _cp_async_mbar_arrive(bars + 8 * (slots + xsl))
- else:
- # ---- consumers ----
- kwi = warp % kw
- twi = warp // kw
- n_ep: cutlass.Constexpr = -(-(rt * m * TILE_N) // (cw * 32))
- # verify epilogue items: (tile, request, row of the tile), each all conv_p tokens
- n_ev: cutlass.Constexpr = -(-(rt * max(self.nreq, 1) * TILE_N) // (cw * 32))
- if const_expr(self.conv_p > 0):
- # Before the PDL wait, as below: each item's slot, snapshot row, 4 taps (two
- # words) and the 3-value window it starts from, kept as raw bits: only the
- # window's address waits on a load (the slot) before the main loop.
- nreq: cutlass.Constexpr = self.nreq
- cvs = cute.make_rmem_tensor(
- cute.make_layout((n_ev, 8), stride=(8, 1)), cutlass.Int32
- )
- cslots = cute.make_rmem_tensor(
- cute.make_layout((n_ev,), stride=(1,)), cutlass.Int32
- )
- crows = cute.make_rmem_tensor(cute.make_layout((n_ev,), stride=(1,)), cutlass.Int32)
- for ic in cutlass.range_constexpr(n_ev):
- ec = tidx + ic * (cw * 32)
- rc_ = ec % TILE_N
- sc_ = (ec // TILE_N) % nreq
- tc_ = ec // (TILE_N * nreq)
- oc_ = (t0 + tc_) * TILE_N + rc_
- okc = (tc_ < nt) & (oc_ < self.conv_dim)
- sl_c = _ldg_s32p(
- idx_base + cutlass.Int64(sc_) * 4, cutlass.Int32(okc), self.pad_slot
- )
- cslots[ic] = sl_c
- okw = cutlass.Int32(okc & (sl_c != self.pad_slot))
- crows[ic] = _ldg_s32p(iidx_base + cutlass.Int64(sc_) * 4, okw, 0)
- taps = _ldg_v2p(cw_base + cutlass.Int64(oc_) * 8, cutlass.Int32(okc))
- cvs[ic, 0] = taps[0]
- cvs[ic, 1] = taps[1]
- wbase = cs_base + (cutlass.Int64(sl_c) * cs_s0 + cutlass.Int64(oc_) * cs_s1) * 2
- for q in cutlass.range_constexpr(3):
- cvs[ic, 4 + q] = _ldg_u16p(wbase + cutlass.Int64(q) * cs_s2 * 2, okw)
- elif const_expr(self.conv_dim > 0):
- # The conv operands of this thread's epilogue outputs -- slot, 4 taps and the 3-value
- # window -- before the PDL wait: the slots are host-filled, and the window is this
- # layer's, last written by this kernel a step ago.
- cvs = cute.make_rmem_tensor(
- cute.make_layout((n_ep, 8), stride=(8, 1)), cutlass.Float32
- )
- cslots = cute.make_rmem_tensor(
- cute.make_layout((n_ep,), stride=(1,)), cutlass.Int32
- )
- for ic in cutlass.range_constexpr(n_ep):
- ec = tidx + ic * (cw * 32)
- rc_ = ec % TILE_N
- nc_ = (ec // TILE_N) % m
- tc_ = ec // (TILE_N * m)
- oc_ = (t0 + tc_) * TILE_N + rc_
- okc = (tc_ < nt) & (oc_ < self.conv_dim)
- sl_c = _ldg_s32p(
- idx_base + cutlass.Int64(nc_) * 4, cutlass.Int32(okc), self.pad_slot
- )
- cslots[ic] = sl_c
- okw = cutlass.Int32(okc & (sl_c != self.pad_slot))
- for q in cutlass.range_constexpr(4):
- cvs[ic, q] = _ldg_bf16p(cw_base + (cutlass.Int64(oc_) * 4 + q) * 2, okw)
- wbase = cs_base + (cutlass.Int64(sl_c) * cs_s0 + cutlass.Int64(oc_) * cs_s1) * 2
- for q in cutlass.range_constexpr(3):
- cvs[ic, 4 + q] = _ldg_bf16p(wbase + cutlass.Int64(q) * cs_s2 * 2, okw)
- if const_expr(self.pdl):
- cute.arch.griddepcontrol_wait()
- acc = cute.make_rmem_tensor(
- cute.make_layout((rtw, nt8, 4), stride=(nt8 * 4, 4, 1)), cutlass.Float32
- )
- acc.fill(0.0)
- if const_expr(self.xreg):
- # fragments of this warp's k-block of a stage: row 8j + g, lane quad t reads
- # elements 32c + 8t .. + 7; loaded `xdepth` stages ahead (the stage loop is
- # unrolled, so every buffer is its own registers)
- xd: cutlass.Constexpr = self.xdepth
- xsrc = [
- x_base
- + cutlass.Int64(8 * j + g) * cutlass.Int64(x_stride) * 2
- + cutlass.Int64((kb0 + kwi) * (BLOCK_K * 2) + t * 16)
- for j in range(nt8)
- ]
- xok = [cutlass.Int32(8 * j + g < m) for j in range(nt8)]
- xbuf = {}
- for s0 in cutlass.range_constexpr(xd):
- xbuf[s0] = [
- [
- v
- for c in range(4)
- for v in _ldg4p(xsrc[j] + (s0 * kw * (BLOCK_K * 2) + 64 * c), xok[j])
- ]
- for j in range(nt8)
- ]
- for s in cutlass.range_constexpr(n_stages):
- if const_expr(s + xd < n_stages):
- xbuf[s + xd] = [
- [
- v
- for c in range(4)
- for v in _ldg4p(
- xsrc[j] + ((s + xd) * kw * (BLOCK_K * 2) + 64 * c), xok[j]
- )
- ]
- for j in range(nt8)
- ]
- sl = s % slots
- slot = ring + sl * slot_bytes
- _mbar_wait(bars + 8 * sl, (s // slots) & 1)
- cur = xbuf.pop(s)
- if const_expr(self.w4):
- cur = [[_bf16x2_to_f16x2(v) for v in r] for r in cur]
- self._compute(
- acc, cur, slot, kb0 + s * kw + kwi, kwi, twi, t0, nt, b0, scs, lane
- )
- cute.arch.sync_warp()
- if lane == 0:
- _mbar_arrive(bars + 8 * (2 * slots + sl))
- else:
- # activation row 8j + g (clamped to m - 1: columns past m are discarded), k-block of
- # this warp: lane quad t reads elements 32c + 8t .. + 7
- x_rows = [cutlass.min(8 * j + g, m - 1) * XRS + t * 16 for j in range(nt8)]
- for s in cutlass.range(n_stages):
- sl = s % slots
- par = (s // slots) & 1
- slot = ring + sl * slot_bytes
- _mbar_wait(bars + 8 * (slots + sl), par)
- xa = slot + w_bytes + kwi * (m * XRS)
- cur = [
- [v for c in range(4) for v in _lds4(xa + x_rows[j] + 64 * c)]
- for j in range(nt8)
- ]
- if const_expr(self.w4):
- cur = [[_bf16x2_to_f16x2(v) for v in r] for r in cur]
- _mbar_wait(bars + 8 * sl, par)
- self._compute(
- acc, cur, slot, kb0 + s * kw + kwi, kwi, twi, t0, nt, b0, scs, lane
- )
- cute.arch.sync_warp()
- if lane == 0:
- _mbar_arrive(bars + 8 * (2 * slots + sl))
-
- # ---- k-block partials -> shared memory [kw][rt][np][16] fp32 ----
- for i in cutlass.range_constexpr(rtw):
- wtl = twi + i * tw
- if wtl < nt:
- for j in cutlass.range_constexpr(nt8):
- n0 = 8 * j + 2 * t
- p = red + (((kwi * rt + wtl) * np_ + n0) * TILE_N + g) * 4
- v = [acc[i, j, e] for e in range(4)]
- if const_expr(self.w4):
- v = [ve * gsc for ve in v]
- _sts_f32(p, v[0])
- _sts_f32(p + TILE_N * 4, v[1])
- _sts_f32(p + 8 * 4, v[2])
- _sts_f32(p + TILE_N * 4 + 8 * 4, v[3])
- _bar_sync(1, cw * 32)
- if const_expr(self.conv_p > 0):
- # Item (tile, request s, row): column `orow` of the request's rows s p .. s p + p - 1.
- # A conv column runs the p tokens through the window in order -- output, snapshot
- # t = the window after token t, then the advanced window into the pool; the other
- # columns are the plain projection. Out: [m, conv_dim] conv output, then
- # [m, n - conv_dim] projection, both row-major in the one buffer.
- cp: cutlass.Constexpr = self.conv_p
- cdim: cutlass.Constexpr = self.conv_dim
- n_rest: cutlass.Constexpr = n_out - self.conv_dim
- rest_base = out_base + cutlass.Int64(m * cdim * 2)
- for it in cutlass.range_constexpr(n_ev):
- e_ = tidx + it * (cw * 32)
- er = e_ % TILE_N
- es = (e_ // TILE_N) % nreq
- etl = e_ // (TILE_N * nreq)
- orow = (t0 + etl) * TILE_N + er
- if (etl < nt) & (orow < n_out):
- if orow < cdim:
- live = cslots[it] != self.pad_slot
- # bf16 bits throughout; taps as two bf16x2 words
- c0 = cvs[it, 4]
- c1 = cvs[it, 5]
- c2 = cvs[it, 6]
- iwb = (
- iw_base
- + (cutlass.Int64(crows[it]) * iw_s0 + cutlass.Int64(orow) * iw_s2)
- * 2
- )
- for tk in cutlass.range_constexpr(cp):
- en = es * cp + tk
- tot = cutlass.Float32(0.0)
- for q in cutlass.range_constexpr(kw):
- tot = tot + _lds_f32(
- red + (((q * rt + etl) * np_ + en) * TILE_N + er) * 4
- )
- xv = _bf16_bits(tot)
- res = _bf16_lo(xv)
- if live:
- res = _conv4_silu_bf16(c0, c1, c2, xv, cvs[it, 0], cvs[it, 1])
- sb = iwb + cutlass.Int64(tk) * cutlass.Int64(iw_s1) * 2
- _stg_b16(sb, c1)
- _stg_b16(sb + cutlass.Int64(iw_s3) * 2, c2)
- _stg_b16(sb + cutlass.Int64(iw_s3) * 4, xv)
- _stg_bf16(out_base + (cutlass.Int64(en) * cdim + orow) * 2, res)
- c0 = c1
- c1 = c2
- c2 = xv
- if live:
- wb = (
- cs_base
- + (
- cutlass.Int64(cslots[it]) * cs_s0
- + cutlass.Int64(orow) * cs_s1
- )
- * 2
- )
- _stg_b16(wb, c0)
- _stg_b16(wb + cutlass.Int64(cs_s2) * 2, c1)
- _stg_b16(wb + cutlass.Int64(cs_s2) * 4, c2)
- else:
- for tk in cutlass.range_constexpr(cp):
- en = es * cp + tk
- tot = cutlass.Float32(0.0)
- for q in cutlass.range_constexpr(kw):
- tot = tot + _lds_f32(
- red + (((q * rt + etl) * np_ + en) * TILE_N + er) * 4
- )
- _stg_bf16(
- rest_base + (cutlass.Int64(en) * n_rest + orow - cdim) * 2,
- tot,
- )
- elif const_expr(ks == 1):
- for it in cutlass.range_constexpr(n_ep):
- e_ = tidx + it * (cw * 32)
- er = e_ % TILE_N
- en = (e_ // TILE_N) % m
- etl = e_ // (TILE_N * m)
- orow = (t0 + etl) * TILE_N + er
- if (etl < nt) & (orow < n_out):
- tot = cutlass.Float32(0.0)
- for q in cutlass.range_constexpr(kw):
- tot = tot + _lds_f32(
- red + (((q * rt + etl) * np_ + en) * TILE_N + er) * 4
- )
- if const_expr(self.conv_dim > 0):
- p_ = _round_bf16(tot)
- res = p_
- if (orow < self.conv_dim) & (cslots[it] != self.pad_slot):
- cv = (
- cvs[it, 0] * cvs[it, 4]
- + cvs[it, 1] * cvs[it, 5]
- + cvs[it, 2] * cvs[it, 6]
- + cvs[it, 3] * p_
- )
- res = cv / (1.0 + cute.math.exp(-cv))
- wb = (
- cs_base
- + (
- cutlass.Int64(cslots[it]) * cs_s0
- + cutlass.Int64(orow) * cs_s1
- )
- * 2
- )
- _stg_bf16(wb, cvs[it, 5])
- _stg_bf16(wb + cutlass.Int64(cs_s2) * 2, cvs[it, 6])
- _stg_bf16(wb + cutlass.Int64(cs_s2) * 4, p_)
- _stg_bf16(out_base + (cutlass.Int64(en) * n_out + orow) * 2, res)
- else:
- _stg_bf16(out_base + (cutlass.Int64(en) * n_out + orow) * 2, tot)
- if const_expr(ks > 1):
- # The CTA's K-share of each output goes to cluster rank (tile % ks), slot
- # [sender][tile][row of M][16]; after the cluster barrier the owner sums the shares in
- # rank order and stores. Every thread of the cluster takes part in the barriers.
- _cluster_wait()
- if warp < cw:
- items1: cutlass.Constexpr = rt * m * TILE_N
- for it1 in cutlass.range_constexpr(-(-items1 // (cw * 32))):
- e1 = tidx + it1 * (cw * 32)
- er1 = e1 % TILE_N
- en1 = (e1 // TILE_N) % m
- etl1 = e1 // (TILE_N * m)
- if etl1 < nt:
- tot1 = cutlass.Float32(0.0)
- for q1 in cutlass.range_constexpr(kw):
- tot1 = tot1 + _lds_f32(
- red + (((q1 * rt + etl1) * np_ + en1) * TILE_N + er1) * 4
- )
- cslot = (
- sbase
- + self.o_cred
- + (((rank * rt + etl1) * m + en1) * TILE_N + er1) * 4
- )
- _st_dsmem_f32(_mapa(cslot, etl1 % ks), tot1)
- _cluster_arrive_release()
- _cluster_wait_acquire()
- if warp < cw:
- items2: cutlass.Constexpr = rt * m * TILE_N
- for it2 in cutlass.range_constexpr(-(-items2 // (cw * 32))):
- f_ = tidx + it2 * (cw * 32)
- fr = f_ % TILE_N
- fn_ = (f_ // TILE_N) % m
- ftl = f_ // (TILE_N * m)
- frow = (t0 + ftl) * TILE_N + fr
- if (ftl < nt) & (frow < n_out) & (ftl % ks == rank):
- tot2 = cutlass.Float32(0.0)
- for q2 in cutlass.range_constexpr(ks):
- tot2 = tot2 + _lds_f32(
- sbase
- + self.o_cred
- + (((q2 * rt + ftl) * m + fn_) * TILE_N + fr) * 4
- )
- _stg_bf16(out_base + (cutlass.Int64(fn_) * n_out + frow) * 2, tot2)
-
-
-# ------------------------------------------------------------------------------------------
-# host side
-# ------------------------------------------------------------------------------------------
-
-
-@dataclass
-class PackedFp8Weight:
- """`data`: uint8 `[K/128, ceil(N/16), 2048]` (see `pack_weight`); `scale`: fp32
- `[ceil(N/128), K/128]` contiguous; logical shape `[n, k]`."""
-
- data: torch.Tensor
- scale: torch.Tensor
- n: int
- k: int
- tile_major: bool = False
-
-
-def quantize_fp8_block(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
- """bf16/fp32 `[N, K]` -> (e4m3 codes `[N, K]`, fp32 scales `[ceil(N/128), K/128]`),
- scale = amax(128x128 block) / 448, code = round_to_e4m3(w / scale)."""
- n, k = w.shape
- assert k % BLOCK_K == 0
- nb = -(-n // 128)
- wp = torch.zeros(nb * 128, k, dtype=torch.float32, device=w.device)
- wp[:n] = w.float()
- blocks = wp.reshape(nb, 128, k // BLOCK_K, BLOCK_K)
- amax = blocks.abs().amax((1, 3))
- scale = (amax / FP8_MAX).clamp_min(1e-30)
- codes = (blocks / scale[:, None, :, None]).clamp(-FP8_MAX, FP8_MAX).reshape(nb * 128, k)[:n]
- return codes.to(torch.float8_e4m3fn), scale.contiguous()
-
-
-def dequantize_fp8_block(codes: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
- """fp32 `[N, K]` = code * scale of its 128x128 block."""
- n, k = codes.shape
- s = scale.repeat_interleave(128, 0)[:n].repeat_interleave(BLOCK_K, 1)
- return codes.float() * s
-
-
-#: Weights of at most this many 16-row tiles are stored tile-major (a CTA then holds one or two
-#: tiles, and one bulk copy brings a tile's k-blocks of a stage); larger ones k-block-major.
-TILE_MAJOR_MAX_TILES = 320
-
-
-def pack_weight(
- codes: torch.Tensor, scale: torch.Tensor, tile_major: bool | None = None
-) -> PackedFp8Weight:
- """e4m3 `[N, K]` (+ its block scales `[ceil(N/128), K/128]` fp32) -> the kernel's layout.
-
- A unit (k-block kb, 16-row tile T) is 2048 bytes: rows 16T .. 16T + 15, k = 128 kb .. + 127, as
- `[c (4)][lane (32)][16 B]`, lane = 4 g + t, the 16 bytes being
- `W[16T + g][32c + 8t + 0..3], W[16T + g + 8][same], W[16T + g][32c + 8t + 4..7],
- W[16T + g + 8][same]` (rows past N are zero). Units are stored `[K/128][ceil(N/16)]`
- (k-block-major) or, with `tile_major`, `[ceil(N/16)][K/128]` (default: tile-major up to
- TILE_MAJOR_MAX_TILES tiles)."""
- n, k = codes.shape
- assert k % BLOCK_K == 0 and codes.dtype == torch.float8_e4m3fn
- n_tiles = -(-n // TILE_N)
- if tile_major is None:
- tile_major = n_tiles <= TILE_MAJOR_MAX_TILES
- b = torch.zeros(n_tiles * TILE_N, k, dtype=torch.uint8, device=codes.device)
- b[:n] = codes.view(torch.uint8)
- # source (tile, h, g, kb, c, t, jj, e); unit order (c, g, t, jj, h, e)
- v = b.reshape(n_tiles, 2, 8, k // BLOCK_K, 4, 4, 2, 4)
- if tile_major:
- data = v.permute(0, 3, 4, 2, 5, 6, 1, 7).contiguous().reshape(n_tiles, k // BLOCK_K, UNIT)
- else:
- data = v.permute(3, 0, 4, 2, 5, 6, 1, 7).contiguous().reshape(k // BLOCK_K, n_tiles, UNIT)
- nb = -(-n // 128)
- assert scale.shape == (nb, k // BLOCK_K), (scale.shape, nb)
- return PackedFp8Weight(data, scale.float().contiguous(), n, k, tile_major)
-
-
-def pack_from_bf16(w: torch.Tensor) -> PackedFp8Weight:
- """Converter from a bf16 `[N, K]` weight."""
- codes, scale = quantize_fp8_block(w)
- return pack_weight(codes, scale)
-
-
-#: Shared memory one CTA may take so that a CTA of the next launch fits beside it (PDL overlap).
-SMEM_BUDGET = 100 << 10
-
-
-#: Measured plans (GB300, 152 SMs, CUDA-graph chain with PDL, weights rotated over > L2) for the
-#: qwen3.8-flash-next-nvfp4 tp1 decode projections: (N, K) -> ((largest M, Plan fields), ...).
-#: `grid` 0 is "auto" (one CTA per 16-row tile up to 9/8 x SMs, else one per SM); -t is t tiles a
-#: CTA (-2: half the tiles).
-_TABLE: dict[tuple[int, int], tuple[tuple[int, dict], ...]] = {
- # GDN in_proj (q|k|v|z|a|b)
- (16480, 2560): (
- (8, dict(grid=0, kw=2, tw=4, slots=3, l2_ahead=1, xreg=2)),
- (16, dict(grid=0, kw=2, tw=4, slots=3, xreg=2)),
- (32, dict(grid=0, kw=2, tw=4, slots=2, xreg=2)),
- ),
- # attention in-proj (q|gate|k|v|index)
- (14080, 2560): (
- (8, dict(grid=0, kw=2, tw=6, slots=3, l2_ahead=1, xreg=2)),
- (16, dict(grid=0, kw=2, tw=6, slots=3, xreg=2)),
- (32, dict(grid=0, kw=2, tw=6, slots=2, xreg=2)),
- ),
- # GDN out / attention o
- (2560, 6144): (
- (1, dict(grid=0, kw=12, tw=1, slots=4, xreg=1, xdepth=2)),
- (8, dict(grid=0, kw=12, tw=1, slots=4, xreg=1, xdepth=1)),
- (16, dict(grid=-2, kw=6, tw=2, slots=3, xreg=2)),
- (32, dict(grid=-2, kw=6, tw=2, slots=2, xreg=2)),
- ),
- # shared expert gate_up with its gate row, and without it
- (1281, 2560): ((32, dict(grid=0, kw=20, tw=1, slots=1, xreg=1)),),
- (1280, 2560): ((32, dict(grid=0, kw=20, tw=1, slots=1, xreg=1)),),
- # shared expert down
- (2560, 640): ((32, dict(grid=0, kw=5, tw=1, slots=1, xreg=1)),),
- # vocab projection (MTP draft head / verify head): 15520 tiles, 7 a CTA -> 2218 CTAs (about
- # 15 a multiprocessor over the launch, ~66 KB shared memory each)
- (248320, 2560): ((32, dict(grid=-7, kw=2, tw=4, slots=2, l2_ahead=1, xreg=2)),),
-}
-
-#: Measured plans of the NVFP4 units (`w4`), as `_TABLE`. GDN in_proj: four k-blocks a stage
-#: (5 stages), no L2 look-ahead: 6.1 / 6.6 / 7.1 us at 1 / 4 / 8 rows alone, 6.4 / 6.1 behind a
-#: PDL predecessor (1 / 4 rows); from 9 rows the heuristic plan, behind fp8_gemv there.
-_TABLE_W4: dict[tuple[int, int], tuple[tuple[int, dict], ...]] = {
- (16480, 2560): ((8, dict(grid=0, kw=4, tw=4, slots=3, xreg=2)),),
-}
-
-#: Most 16-row tiles a CTA takes in the heuristic plan: a ring slot holds `kw` k-blocks of each of
-#: them, so a larger N gets more CTAs (several per multiprocessor over the launch) instead.
-MAX_TILES_PER_CTA = 7
-
-
-def default_plan(n: int, k: int, m: int, sms: int, w4: bool = False) -> Plan:
- """The plan `fp8_gemv` launches with: the measured table entry for (N, K, M), else a heuristic
- (`w4`: over NVFP4 units)."""
- n_tiles = -(-n // TILE_N)
- nkb = k // BLOCK_K
- auto_grid = n_tiles if n_tiles <= sms + sms // 8 else sms
- for bound, fields in (_TABLE_W4 if w4 else _TABLE).get((n, k), ()):
- if m <= bound:
- f = dict(fields)
- g = f.pop("grid")
- grid = auto_grid if g == 0 else -(-n_tiles // -g)
- return Plan(grid=grid, **f)
- grid = auto_grid
- rt = -(-n_tiles // grid)
- if rt > MAX_TILES_PER_CTA:
- grid = -(-n_tiles // MAX_TILES_PER_CTA)
- rt = -(-n_tiles // grid)
- divs = [d for d in range(1, nkb + 1) if nkb % d == 0]
- if rt >= 4:
- kw = max(d for d in divs if d <= 2)
- tw = min(rt, 4)
- else:
- kw = max(d for d in divs if d <= 12)
- tw = rt
- nt8 = -(-m // 8)
- xreg = (nkb // kw) * nt8 * 16 <= XREG_MAX
- slot = kw * rt * (UNIT_W4 if w4 else UNIT) + (0 if xreg else kw * m * XRS)
- slots = max(1, min(nkb // kw, SMEM_BUDGET // slot))
- return Plan(grid=grid, kw=kw, tw=tw, slots=slots, xreg=1 if xreg else 2)
-
-
-@functools.lru_cache(maxsize=None)
-def _compiled(
- n,
- k,
- m,
- plan: Plan,
- pdl: bool,
- tile_major: bool,
- conv_dim: int = 0,
- pad_slot: int = -1,
- conv_p: int = 0,
- w4: bool = False,
-):
- kern = _Fp8Gemv(
- n=n,
- k=k,
- m=m,
- plan=plan,
- pdl=pdl,
- tile_major=tile_major,
- conv_dim=conv_dim,
- pad_slot=pad_slot,
- conv_p=conv_p,
- w4=w4,
- )
- dev = "cuda"
- out = torch.empty(m, n, dtype=torch.bfloat16, device=dev)
- x = torch.empty(m, k, dtype=torch.bfloat16, device=dev)
- w = torch.empty(16, dtype=torch.uint8, device=dev)
- s = torch.empty(4, dtype=torch.float32, device=dev)
- return cute.compile(
- kern,
- from_dlpack(out, assumed_align=16).mark_layout_dynamic(leading_dim=1),
- make_ptr(cutlass.BFloat16, x.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
- make_ptr(cutlass.Uint8, w.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
- make_ptr(cutlass.Float32, s.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
- cutlass.Int32(k),
- make_ptr(cutlass.BFloat16, x.data_ptr(), cute.AddressSpace.gmem, assumed_align=2),
- make_ptr(cutlass.BFloat16, x.data_ptr(), cute.AddressSpace.gmem, assumed_align=2),
- make_ptr(cutlass.Int32, s.data_ptr(), cute.AddressSpace.gmem, assumed_align=4),
- cutlass.Int32(0),
- cutlass.Int32(0),
- cutlass.Int32(0),
- make_ptr(cutlass.BFloat16, x.data_ptr(), cute.AddressSpace.gmem, assumed_align=2),
- make_ptr(cutlass.Int32, s.data_ptr(), cute.AddressSpace.gmem, assumed_align=4),
- cutlass.Int32(0),
- cutlass.Int32(0),
- cutlass.Int32(0),
- cutlass.Int32(0),
- cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
- options="--opt-level 2 --enable-tvm-ffi",
- )
-
-
-@functools.lru_cache(maxsize=None)
-def _sms() -> int:
- return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count
-
-
-#: The launch arguments past `x_stride` of a call without a conv epilogue.
-_NO_CONV = (0,) * 12
-
-#: Most tokens per request of `fp8_gemv_conv1d_verify` (`max_draft_len + 1` of the widest
-#: verify step the target runs).
-VERIFY_MAX_P = 8
-
-
-def _env_pdl() -> bool:
- try:
- from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
- except ImportError:
- return True
- return bool(get_env_enable_pdl())
-
-
-def fp8_gemv(
- x: torch.Tensor, w: PackedFp8Weight, pdl: bool = True, plan: Plan | None = None
-) -> torch.Tensor:
- """`x[M, K]` bf16 (row stride >= K, 16-byte aligned rows) @ dequant(w)^T -> fresh
- `[M, N]` bf16. M = 1 .. MAX_ROWS. `w` a `PackedFp8Weight`, or NVFP4 units in this kernel's
- order (`nvfp4_gemv.SplitNvfp4Weight`, `w.w4`)."""
- m, k = x.shape
- assert k == w.k and x.dtype == torch.bfloat16 and x.stride(1) == 1
- assert 1 <= m <= MAX_ROWS
- assert x.stride(0) % 8 == 0 and x.data_ptr() % 16 == 0
- w4 = getattr(w, "w4", False)
- out = torch.empty((m, w.n), dtype=torch.bfloat16, device=x.device)
- if plan is None:
- plan = default_plan(w.n, k, m, _sms(), w4)
- fn = _compiled(w.n, k, m, plan, bool(pdl and _env_pdl()), w.tile_major, w4=w4)
- fn(
- out,
- x.data_ptr(),
- w.data.data_ptr(),
- w.scale.data_ptr(),
- x.stride(0),
- *_NO_CONV,
- )
- return out
-
-
-def fp8_gemv_conv1d(
- x: torch.Tensor,
- w: PackedFp8Weight,
- conv_state: torch.Tensor,
- conv_weight: torch.Tensor,
- conv_state_indices: torch.Tensor,
- conv_dim: int,
- pad_slot_id: int,
- pdl: bool = True,
- plan: Plan | None = None,
-) -> torch.Tensor:
- """FP8-weight `gdn_in_proj_conv1d_pdl`: `p = round_bf16(x @ dequant(w)^T)`; for a column
- c < conv_dim of a row r whose slot s = conv_state_indices[r] is not pad_slot_id,
- `out[r, c] = round_bf16(silu(sum_j conv_weight[c, j] * [conv_state[s, c, :], p[r, c]][j]))` (fp32)
- and `conv_state[s, c, :] = [conv_state[s, c, 1:], p[r, c]]`; every other output is `p`.
- `conv_state` `[slots, conv_dim, 3]` bf16 (any strides), `conv_weight` `[conv_dim, 4]` bf16
- contiguous, `conv_state_indices` int32 `[M]`. Returns the fresh `[M, N]` output. `w` as
- `fp8_gemv`'s."""
- m, k = x.shape
- assert k == w.k and x.dtype == torch.bfloat16 and x.stride(1) == 1 and 1 <= m <= MAX_ROWS
- assert x.stride(0) % 8 == 0 and x.data_ptr() % 16 == 0
- assert conv_state.dtype == torch.bfloat16 and tuple(conv_state.shape[1:]) == (conv_dim, 3)
- assert conv_weight.dtype == torch.bfloat16 and tuple(conv_weight.shape) == (conv_dim, 4)
- assert conv_weight.is_contiguous() and conv_dim <= w.n
- assert conv_state_indices.dtype == torch.int32 and conv_state_indices.is_contiguous()
- assert conv_state_indices.numel() == m
- w4 = getattr(w, "w4", False)
- out = torch.empty((m, w.n), dtype=torch.bfloat16, device=x.device)
- if plan is None:
- plan = default_plan(w.n, k, m, _sms(), w4)
- fn = _compiled(
- w.n, k, m, plan, bool(pdl and _env_pdl()), w.tile_major, conv_dim, int(pad_slot_id), w4=w4
- )
- fn(
- out,
- x.data_ptr(),
- w.data.data_ptr(),
- w.scale.data_ptr(),
- x.stride(0),
- conv_state.data_ptr(),
- conv_weight.data_ptr(),
- conv_state_indices.data_ptr(),
- *conv_state.stride(),
- *_NO_CONV[6:],
- )
- return out
-
-
-def fp8_gemv_conv1d_verify(
- x: torch.Tensor,
- w: PackedFp8Weight,
- conv_state: torch.Tensor,
- conv_weight: torch.Tensor,
- conv_state_indices: torch.Tensor,
- intermediate_conv_window: torch.Tensor,
- intermediate_state_indices: torch.Tensor,
- conv_dim: int,
- rows_per_seq: int,
- pad_slot_id: int,
- pdl: bool = True,
- plan: Plan | None = None,
-) -> tuple[torch.Tensor, torch.Tensor]:
- """`fp8_gemv` + `causal_conv1d_triton_update` of a speculative verify step in one launch.
-
- `x` is `n` requests of `p = rows_per_seq` consecutive token rows (row `s p + t` is token
- `t` of request `s`); `proj = round_bf16(x @ dequant(w)^T)`. For request `s` with slot
- `c = conv_state_indices[s]` != pad_slot_id and snapshot row `r =
- intermediate_state_indices[s]`, over the window sequence `win = [conv_state[c, :, 0..2],
- proj[s p .. s p + p - 1, :conv_dim]]` (per channel): `conv[s p + t] = silu(sum_j
- bf16(conv_weight[:, j] * win[t + j]))` (bf16 products, fp32 sum, the Triton kernel's
- SiLU), `intermediate_conv_window[r, t, :conv_dim, :] = win[t + 1 .. t + 3]` for t < p, and
- `conv_state[c] = win[p .. p + 2]`; a pad request's rows get the plain projection and
- write no state. Returns `(conv [M, conv_dim], proj[:, conv_dim:] [M, N - conv_dim])`,
- both contiguous views of one fresh buffer.
-
- `conv_state` `[slots, conv_dim, 3]` and `intermediate_conv_window` `[rows, >= p,
- conv_dim, 3]` bf16 (any strides), `conv_weight` `[conv_dim, 4]` bf16 contiguous, both
- index vectors int32 `[n]` contiguous. The slots, snapshot rows, taps and the windows are
- read before the PDL wait (as `fp8_gemv_conv1d` does). `w` as `fp8_gemv`'s."""
- m, k = x.shape
- p = int(rows_per_seq)
- assert k == w.k and x.dtype == torch.bfloat16 and x.stride(1) == 1 and 2 <= m <= MAX_ROWS
- assert x.stride(0) % 8 == 0 and x.data_ptr() % 16 == 0
- assert 2 <= p <= VERIFY_MAX_P and m % p == 0, (m, p)
- nreq = m // p
- assert conv_state.dtype == torch.bfloat16 and tuple(conv_state.shape[1:]) == (conv_dim, 3)
- assert conv_weight.dtype == torch.bfloat16 and tuple(conv_weight.shape) == (conv_dim, 4)
- assert conv_weight.is_contiguous() and conv_weight.data_ptr() % 8 == 0 and 0 < conv_dim < w.n
- iw, iidx = intermediate_conv_window, intermediate_state_indices
- assert iw.dtype == torch.bfloat16 and iw.dim() == 4 and iw.shape[1] >= p
- assert tuple(iw.shape[2:]) == (conv_dim, 3)
- for v in (conv_state_indices, iidx):
- assert v.dtype == torch.int32 and v.is_contiguous() and v.numel() == nreq
- w4 = getattr(w, "w4", False)
- buf = torch.empty(m * w.n, dtype=torch.bfloat16, device=x.device)
- if plan is None:
- plan = default_plan(w.n, k, m, _sms(), w4)
- fn = _compiled(
- w.n,
- k,
- m,
- plan,
- bool(pdl and _env_pdl()),
- w.tile_major,
- conv_dim,
- int(pad_slot_id),
- p,
- w4=w4,
- )
- fn(
- buf.view(m, w.n),
- x.data_ptr(),
- w.data.data_ptr(),
- w.scale.data_ptr(),
- x.stride(0),
- conv_state.data_ptr(),
- conv_weight.data_ptr(),
- conv_state_indices.data_ptr(),
- *conv_state.stride(),
- iw.data_ptr(),
- iidx.data_ptr(),
- *iw.stride(),
- )
- return buf[: m * conv_dim].view(m, conv_dim), buf[m * conv_dim :].view(m, w.n - conv_dim)
-
-
-def _gemv_cell(why: str, n: int, k: int, m: int) -> Cell:
- return Cell(why=why, spec=dict(n=n, k=k, m=m))
-
-
-class Fp8Gemv(OpWrapper):
- """`fp8_gemv(x, w)`: `y[M, N] = x[M, K] @ dequant(w)^T` in bf16 for 1..MAX_ROWS rows, the
- weight a `PackedFp8Weight` (e4m3 codes with 128x128 fp32 block scales, packed by
- `pack_weight` / `pack_from_bf16`), the bf16 activation never quantized. fp32 accumulation per
- 128-wide k-block, block scale applied in fp32, one bf16 rounding at the store.
-
- Fusion boundary: the projection; `fp8_gemv_conv1d` also fuses the GDN causal-conv step of
- decode rows, `fp8_gemv_conv1d_verify` the multi-token step of a verify step's rows."""
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- _gemv_cell("GDN in_proj, one decode row", 16480, 2560, 1),
- _gemv_cell("attention in-proj, MTP3 verify at con1 (4 rows)", 14080, 2560, 4),
- _gemv_cell("GDN out / attention o, 4 rows", 2560, 6144, 4),
- _gemv_cell("shared expert gate|up, 16 rows", 1280, 2560, 16),
- _gemv_cell("shared expert down, 8 rows", 2560, 640, 8),
- _gemv_cell("vocab projection (draft / verify head), 4 rows", 248320, 2560, 4),
- )
-
- note = """
- Metadata consumed: none. Call sites: the decode target's dense projections under the dense
- FP8 mode (`TRTLLM_MODELING_V2_DENSE_FP8`) at up to `_FP8_GEMV_MAX_ROWS` rows.
- Preconditions: x bf16 `[M, K]`, unit column stride, row stride a multiple of 8 elements and a
- 16-byte aligned base; K a multiple of 128. One PDL launch (weights streamed before the
- wait); CUDA-graph capturable; compiled per (N, K, M, plan) on first use.
- """
-
- def __call__(self, x: torch.Tensor, w: PackedFp8Weight) -> torch.Tensor:
- return fp8_gemv(x, w)
-
- def reference(self, x: torch.Tensor, codes: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
- """fp32 `[M, N]` over the dequantized weight (a true fp32 matmul)."""
- with true_fp32_matmul():
- return x.float() @ dequantize_fp8_block(codes, scale).t()
-
-
-class Fp8GemvConv1d(OpWrapper):
- """`fp8_gemv_conv1d`: the FP8-weight `gdn_in_proj_conv1d_pdl` -- the GDN input projection
- of decode rows, then for the conv columns of a live row the causal conv step against its
- paged state (SiLU; the state advanced in place). Returns the `[M, N]` projection row with
- the conv columns replaced by the conv output."""
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(why="STP decode: one row, live slot", spec=dict(slots=(37,))),
- Cell(why="two rows, one padding slot", spec=dict(slots=(17, -1))),
- )
-
- note = """
- Call site: the decode target's one-row GDN in_proj + conv step under the dense FP8 mode.
- `conv_state` `[slots, conv_dim, 3]` bf16 (any strides), `conv_weight` `[conv_dim, 4]` bf16
- contiguous, int32 slots; a `pad_slot_id` row gets the plain projection.
- """
-
- def __call__(self, x, w, conv_state, conv_weight, conv_state_indices, conv_dim, pad_slot_id):
- return fp8_gemv_conv1d(
- x, w, conv_state, conv_weight, conv_state_indices, conv_dim, pad_slot_id
- )
-
- def reference(self, *args, **kwargs):
- raise NotImplementedError("the test builds the fp32 reference from the dequantized weight")
-
-
-#: The checkpoint's GDN in_proj (q|k|v|z|a|b) and its q|k|v conv width.
-GDN_IN_N, GDN_IN_K, GDN_CONV_DIM = 16480, 2560, 10240
-
-
-def _verify_cell(why: str, p: int, slots: tuple, steps: int | None = None) -> Cell:
- return Cell(
- why=why,
- spec=dict(
- n=GDN_IN_N,
- k=GDN_IN_K,
- conv_dim=GDN_CONV_DIM,
- p=p,
- slots=slots,
- snapshot_steps=steps or p,
- pool_slots=461,
- snapshot_rows=32,
- ),
- )
-
-
-def verify_reference(
- x,
- deq,
- conv_state,
- conv_weight,
- conv_state_indices,
- intermediate_conv_window,
- intermediate_state_indices,
- conv_dim,
- rows_per_seq,
- pad_slot_id,
-):
- """`fp8_gemv_conv1d_verify`'s reference over a dequantized fp32 weight `deq` `[N, K]`; see
- `Fp8GemvConv1dVerify.reference`."""
- with true_fp32_matmul():
- proj = (x.float() @ deq.t()).to(torch.bfloat16)
- m, p = x.shape[0], rows_per_seq
- conv = proj[:, :conv_dim].clone()
- state, window = conv_state.clone(), intermediate_conv_window.clone()
- slots = conv_state_indices.tolist()
- rows = intermediate_state_indices.tolist()
- for s in range(m // p):
- if slots[s] == pad_slot_id:
- continue
- seq = torch.cat([conv_state[slots[s]], proj[s * p : (s + 1) * p, :conv_dim].t()], 1)
- for t in range(p):
- acc = torch.zeros(conv_dim, dtype=torch.float32, device=x.device)
- for j in range(4):
- acc = acc + (conv_weight[:, j] * seq[:, t + j]).float()
- conv[s * p + t] = (acc / (1.0 + torch.exp(-acc))).to(torch.bfloat16)
- window[rows[s], t] = seq[:, t + 1 : t + 4]
- state[slots[s]] = seq[:, p : p + 3]
- return conv, proj[:, conv_dim:].contiguous(), state, window
-
-
-class Fp8GemvConv1dVerify(OpWrapper):
- """`fp8_gemv_conv1d_verify`: the GDN input projection of a speculative verify step (n
- requests x p tokens) with `causal_conv1d_triton_update`'s multi-token conv step as its
- epilogue -- per conv column and request, the p tokens through the cached window (bf16
- products, fp32 sum, SiLU), the window after each token into the per-step snapshot buffer,
- and the pool slot advanced by all p tokens. Returns `(conv [M, conv_dim], projection
- columns conv_dim.. [M, N - conv_dim])`.
-
- Fusion boundary: projection, conv, SiLU, snapshots and the pool advance. The q/k/v and
- z/a/b splits, the recurrence and the post-acceptance promotion of a snapshot are the
- caller's."""
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec keys: n, k (the in_proj), conv_dim, p (tokens per request), slots (one per
- # request; -1 the pad sentinel), snapshot_steps (the buffer's step depth), pool_slots,
- # snapshot_rows. Every cell: bf16, width 4, no bias, SiLU, pad_slot_id -1,
- # intermediate_state_indices = arange(n), the engine's 461-slot pool and 32-row buffer.
- CELLS: tuple[Cell, ...] = (
- _verify_cell("MTP3 verify, 1 request (4 rows)", 4, (37,)),
- _verify_cell("MTP3 verify, 2 requests (8 rows)", 4, (0, 460)),
- _verify_cell("MTP3 verify, 4 requests (16 rows)", 4, (3, 17, 233, 459)),
- _verify_cell("MTP1 verify, 8 requests x 2 rows (16 rows)", 2, tuple(range(5, 461, 57))),
- _verify_cell("MTP7 verify, 2 requests x 8 rows", 8, (100, 7)),
- _verify_cell("3 rows a request into an 8-step buffer: steps 3..7 untouched", 3, (9, 41), 8),
- _verify_cell("MTP3, a pad request beside a live one", 4, (-1, 21)),
- )
-
- note = """
- Call site: the decode target's GDN in_proj of an MTP verify step under the dense FP8 mode
- in a captured graph, 2..16 rows (p = runtime_draft_len + 1 = 2..8 tokens per request,
- a request's rows consecutive); the conv output feeds the cached replay as its contiguous
- packed q|k|v, z/a/b are column views of the second output.
- `conv_state` `[slots, conv_dim, 3]` and `intermediate_conv_window` `[rows, >= p, conv_dim,
- 3]` bf16 (any strides), `conv_weight` `[conv_dim, 4]` bf16 contiguous, int32 slot and
- snapshot-row vectors of n entries. Same arithmetic as `causal_conv1d_triton_update`
- (bitwise the two-launch fp8_gemv + Triton path under the same plan); a pad request gets
- the plain projection and writes neither pool nor snapshots (the Triton op leaves its rows
- unwritten). Slots, snapshot rows, taps and windows are read before the PDL wait: nothing
- in a captured step writes this layer's pool ahead of it.
- Not guarded (content): duplicate slots or snapshot rows race; slot / row values out of
- range are illegal accesses.
- """
-
- def __call__(
- self,
- x,
- w,
- conv_state,
- conv_weight,
- conv_state_indices,
- intermediate_conv_window,
- intermediate_state_indices,
- conv_dim,
- rows_per_seq,
- pad_slot_id,
- ):
- return fp8_gemv_conv1d_verify(
- x,
- w,
- conv_state,
- conv_weight,
- conv_state_indices,
- intermediate_conv_window,
- intermediate_state_indices,
- conv_dim,
- rows_per_seq,
- pad_slot_id,
- )
-
- def reference(
- self,
- x,
- codes,
- scale,
- conv_state,
- conv_weight,
- conv_state_indices,
- intermediate_conv_window,
- intermediate_state_indices,
- conv_dim,
- rows_per_seq,
- pad_slot_id,
- ):
- """`(conv, rest, conv_state, intermediate_conv_window)` as the call leaves them, from
- an fp32 matmul over the dequantized weight rounded once to bf16, then the Triton
- op's reference conv (products in bf16, fp32 sum, `acc / (1 + exp(-acc))`); mutates
- nothing. A pad request's conv rows are its projection."""
- return verify_reference(
- x,
- dequantize_fp8_block(codes, scale),
- conv_state,
- conv_weight,
- conv_state_indices,
- intermediate_conv_window,
- intermediate_state_indices,
- conv_dim,
- rows_per_seq,
- pad_slot_id,
- )
-
- def is_valid(
- self,
- x,
- w,
- conv_state,
- conv_weight,
- conv_state_indices,
- intermediate_conv_window,
- intermediate_state_indices,
- conv_dim,
- rows_per_seq,
- pad_slot_id,
- **kwargs,
- ) -> None:
- m = x.shape[0]
- assert 2 <= rows_per_seq <= VERIFY_MAX_P and m % rows_per_seq == 0, (
- f"{m} rows are not whole requests of {rows_per_seq} tokens (2..{VERIFY_MAX_P})"
- )
- assert m <= MAX_ROWS, f"{m} rows; the kernel holds at most {MAX_ROWS}"
- assert intermediate_conv_window.shape[1] >= rows_per_seq, (
- f"a snapshot buffer of {intermediate_conv_window.shape[1]} steps holds fewer than "
- f"the {rows_per_seq} tokens a request writes"
- )
- assert conv_state.shape[2] == 3 and conv_weight.shape[1] == 4, "width-4 conv only"
- assert conv_state.shape[1] == conv_dim and intermediate_conv_window.shape[2] == conv_dim
-
-
-fp8_gemv_op = Fp8Gemv()
-fp8_gemv_conv1d_op = Fp8GemvConv1d()
-fp8_gemv_conv1d_verify_op = Fp8GemvConv1dVerify()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/fp8_per_tensor_mm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/fp8_per_tensor_mm.py
deleted file mode 100644
index 0ec92dde4f25..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/fp8_per_tensor_mm.py
+++ /dev/null
@@ -1,192 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""A dense projection in FP8: the bf16 activation quantized to e4m3 with one dynamic
-per-tensor scale, the weight stored in e4m3 with one per-tensor scale, cuBLASLt's FP8 GEMM
-with both scales in its epilogue, bf16 out.
-
-Tier 5 (Triton written for this target) for the activation quantization: two launches --
-per-block partial maxima of `|x|`, then the quantization, each of whose programs reduces the
-partials itself (no memset, no host sync) -- and `torch._scaled_mm` (cuBLASLt, tensor-wise
-scales) for the GEMM. Tensor-wise scales keep cuBLASLt on its fast FP8 kernels: on GB300 the
-checkpoint's GDN input projection runs ~1.9x bf16 at 16K rows and ~1.7x at 1-64 rows, where
-row-wise activation scales fall back to kernels only ~1.4x bf16.
-"""
-
-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
-
-E4M3_MAX = 448.0
-#: Floor of the activation maximum: an all-zero activation quantizes to zeros at this scale.
-AMAX_FLOOR = 1e-12
-#: Most programs of either launch; rows are grid-strided above it, and the quantization
-#: reduces at most this many partial maxima.
-_MAX_PROGRAMS = 1024
-_ROWS = 4
-
-
-def fp8_per_tensor_weight(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
- """`(codes, scale)` of a bf16 `[N, K]` weight: e4m3 codes `[N, K]` (contiguous) and the
- fp32 `[1]` dequantization scale `amax / 448`."""
- scale = (weight.abs().amax().float().clamp_min(AMAX_FLOOR) / E4M3_MAX).reshape(1)
- codes = (weight.float() / scale).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn)
- return codes.contiguous(), scale
-
-
-@triton.jit
-def _amax_kernel(
- x_ptr, part_ptr, M, K, stride_m, PROGRAMS, ROWS: tl.constexpr, BLOCK_K: tl.constexpr
-):
- pid = tl.program_id(0)
- acc = tl.zeros([ROWS, BLOCK_K], dtype=tl.float32)
- for r0 in range(pid * ROWS, M, PROGRAMS * ROWS):
- rows = r0 + tl.arange(0, ROWS)
- for k0 in range(0, K, BLOCK_K):
- cols = k0 + tl.arange(0, BLOCK_K)
- mask = (rows[:, None] < M) & (cols[None, :] < K)
- v = tl.load(
- x_ptr + rows[:, None].to(tl.int64) * stride_m + cols[None, :], mask=mask, other=0.0
- )
- acc = tl.maximum(acc, tl.abs(v.to(tl.float32)))
- tl.store(part_ptr + pid, tl.max(tl.max(acc, axis=1), axis=0))
-
-
-@triton.jit
-def _quant_kernel(
- x_ptr,
- part_ptr,
- q_ptr,
- scale_ptr,
- M,
- K,
- stride_m,
- PROGRAMS,
- NUM_PARTS,
- amax_floor,
- ROWS: tl.constexpr,
- BLOCK_K: tl.constexpr,
- BLOCK_P: tl.constexpr,
-):
- pid = tl.program_id(0)
- p = tl.arange(0, BLOCK_P)
- amax = tl.max(tl.load(part_ptr + p, mask=p < NUM_PARTS, other=0.0), axis=0)
- amax = tl.maximum(amax, amax_floor)
- scale = amax / 448.0
- inv = 448.0 / amax
- if pid == 0:
- tl.store(scale_ptr, scale)
- for r0 in range(pid * ROWS, M, PROGRAMS * ROWS):
- rows = r0 + tl.arange(0, ROWS)
- for k0 in range(0, K, BLOCK_K):
- cols = k0 + tl.arange(0, BLOCK_K)
- mask = (rows[:, None] < M) & (cols[None, :] < K)
- v = tl.load(
- x_ptr + rows[:, None].to(tl.int64) * stride_m + cols[None, :], mask=mask, other=0.0
- )
- q = tl.minimum(tl.maximum(v.to(tl.float32) * inv, -448.0), 448.0)
- tl.store(
- q_ptr + rows[:, None].to(tl.int64) * K + cols[None, :],
- q.to(tl.float8e4nv),
- mask=mask,
- )
-
-
-def fp8_quantize_per_tensor(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
- """`(codes [M, K] e4m3 contiguous, scale [1] fp32)` of a bf16 `[M, K]` activation with
- unit column stride: scale = max(|x|) / 448 (floored), codes = round(x / scale)."""
- m, k = x.shape
- q = torch.empty((m, k), dtype=torch.float8_e4m3fn, device=x.device)
- scale = torch.empty(1, dtype=torch.float32, device=x.device)
- programs = max(1, min(triton.cdiv(m, _ROWS), _MAX_PROGRAMS))
- block_k = min(1024, triton.next_power_of_2(k))
- part = torch.empty(programs, dtype=torch.float32, device=x.device)
- _amax_kernel[(programs,)](
- x, part, m, k, x.stride(0), programs, ROWS=_ROWS, BLOCK_K=block_k, num_warps=4
- )
- _quant_kernel[(programs,)](
- x,
- part,
- q,
- scale,
- m,
- k,
- x.stride(0),
- programs,
- part.shape[0],
- AMAX_FLOOR,
- ROWS=_ROWS,
- BLOCK_K=block_k,
- BLOCK_P=triton.next_power_of_2(part.shape[0]),
- num_warps=4,
- )
- return q, scale
-
-
-def _ckpt(m: int, k: int, n: int) -> dict:
- return dict(m=m, k=k, n=n)
-
-
-class Fp8PerTensorMm(OpWrapper):
- """`y = (q(x) * sx) @ (codes * sw)^T` in bf16, accumulated in fp32, where `sx =
- max(|x|) / 448` over the whole activation and `q(x) = round_e4m3(x / sx)`; `weight_t` is
- the `[K, N]` transpose view of the `[N, K]` e4m3 codes `fp8_per_tensor_weight` returns and
- `weight_scale` its `[1]` scale.
-
- Lossy against the bf16 projection: ~3.7% relative error per output on random operands
- (e4m3 keeps 3 mantissa bits on both operands). The activation scale is shared by every
- row of the call, so a row's rounding depends on the largest element of its batch.
-
- Fusion boundary: activation quantization + GEMM. The weight conversion
- (`fp8_per_tensor_weight`) is done once at load by the caller.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(why="GDN input projection, one decode row", spec=_ckpt(1, 2560, 16480)),
- Cell(why="GDN / attention output projection, 128 rows", spec=_ckpt(128, 6144, 2560)),
- Cell(
- why="attention input projection, a 16384-row prefill step",
- spec=_ckpt(16384, 2560, 14080),
- ),
- )
-
- note = """
- Metadata consumed: none. Call sites: the targets' GDN in/out and attention in/o
- projections when the dense FP8 mode is on (`TRTLLM_MODELING_V2_DENSE_FP8`).
- Preconditions: `x` bf16 `[M, K]` with unit column stride, K a multiple of 16; the weight
- `[K, N]` column-major view of contiguous `[N, K]` e4m3 codes with N a multiple of 16
- (cuBLASLt FP8). Three launches (two Triton, one cuBLASLt); CUDA-graph capturable: no host
- sync, the activation scale stays on the device.
- """
-
- def __call__(
- self, x: torch.Tensor, weight_t: torch.Tensor, weight_scale: torch.Tensor
- ) -> torch.Tensor:
- q, scale = fp8_quantize_per_tensor(x)
- return torch._scaled_mm(
- q, weight_t, scale_a=scale, scale_b=weight_scale, out_dtype=torch.bfloat16
- )
-
- def reference(
- self, x: torch.Tensor, weight_t: torch.Tensor, weight_scale: torch.Tensor
- ) -> torch.Tensor:
- """fp32 `[M, N]`: the activation quantized as the op does, both operands dequantized,
- an fp32 matmul."""
- scale = x.abs().amax().float().clamp_min(AMAX_FLOOR) / E4M3_MAX
- xq = (x.float() / scale).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn).float() * scale
- return xq @ (weight_t.float() * weight_scale.float())
-
- def is_valid(self, x: torch.Tensor, weight_t: torch.Tensor, weight_scale: torch.Tensor) -> None:
- assert x.dtype == torch.bfloat16 and x.dim() == 2 and x.stride(1) == 1
- assert weight_t.dtype == torch.float8_e4m3fn and weight_t.stride(0) == 1
- assert weight_t.shape[0] == x.shape[1], "K mismatch"
- assert x.shape[1] % 16 == 0 and weight_t.shape[1] % 16 == 0, "cuBLASLt FP8 needs K, N % 16"
- assert weight_scale.numel() == 1 and weight_scale.dtype == torch.float32
-
-
-fp8_per_tensor_mm = Fp8PerTensorMm()
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
deleted file mode 100644
index 967a23116a87..000000000000
--- 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
+++ /dev/null
@@ -1,231 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Decode-shaped bf16 projection: one kernel per call, its source picked per (M, N, K).
-
-At decode M a dense bf16 projection streams its weight once, and cuBLAS
-reaches only 0.55-0.65 of HBM bandwidth on 0.25-63 MB weights (and ~4.5 us
-flat on the small ones). FlashInfer's Blackwell low-M CuTe-DSL kernels -- the
-register-prefetch "direct" kernel and the tcgen05 split-K kernel, the two
-`flashinfer.mm_bf16(backend="cute-dsl")` dispatches to and
-`tensorrt_llm/_torch/modules/low_m_gemm.py` races -- are faster at M <= 32 on
-most such shapes. Which one, with which tactic, is a per-(N, K) table measured
-per M bucket on GB300 (cold L2, CUDA-graph replay); every shape the table does
-not list, and every M above it, stays on the trtllm cuBLASLt op. The rows
-today are qwen3.8-flash-next-nvfp4's projections (hidden 2560, tp1).
-
-FlashInfer API: ``flashinfer.gemm.kernels.dense_bf16_gemm_direct.run_direct_dense`` and
-``flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk.run_splitk_dense``.
-"""
-
-from __future__ import annotations
-
-import functools
-from typing import NamedTuple
-
-import torch
-from flashinfer.gemm.kernels import dense_bf16_gemm_direct as _fi_direct
-from flashinfer.gemm.kernels import dense_bf16_gemm_sm100_splitk as _fi_splitk
-
-import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import (
- Arch,
- Cell,
- OpWrapper,
- true_fp32_matmul,
-)
-
-
-class _Plan(NamedTuple):
- """`kind` is "direct", "splitk" or "cublas"; `tactic` is the kernel's own tactic tuple."""
-
- kind: str
- tactic: tuple = ()
-
-
-def _d(block_size: int, outputs_per_block: int) -> tuple:
- """A direct-kernel choice; rows_per_block is the kernel's default for the runtime M."""
- return ("direct", block_size, outputs_per_block)
-
-
-def _s(mma_m: int, mma_n: int, split_k: int, ab_stages: int) -> tuple:
- return ("splitk", mma_m, mma_n, split_k, ab_stages)
-
-
-#: (N, K) -> ((M upper bound, choice), ...) in ascending M. An M above the last
-#: bound, and any (N, K) not listed, runs on cuBLAS. Every row below is a
-#: qwen3.8-flash-next-nvfp4 tp1 projection; a target with other shapes adds its
-#: own rows after measuring them. Each choice beat cublas_mm by >= 5% at the
-#: bound's M in a sweep of these shapes (GB300, cold L2, graph replay, pdl off).
-#: M=3 uses the M=4 choice, and so on.
-_TABLE: dict[tuple[int, int], tuple[tuple[int, tuple], ...]] = {
- # GDN in_qkv: 16 q + 16 k heads x 128, 48 v heads x 128
- (10240, 2560): ((1, _d(64, 1)), (2, _d(64, 2)), (8, _s(128, 8, 1, 6))),
- # GDN in_z
- (6144, 2560): ((1, _d(64, 1)), (2, _d(64, 2))),
- # GDN in_a / in_b (48 v heads)
- (48, 2560): ((32, _d(64, 1)),),
- # GDN out and attention o: [6144 -> 2560]
- (2560, 6144): (
- (2, _d(256, 2)),
- (8, _s(64, 8, 3, 10)),
- (16, _s(64, 16, 3, 8)),
- (32, _s(64, 32, 3, 8)),
- ),
- # attention q (q and its output gate interleaved per head)
- (12288, 2560): ((2, _d(64, 2)), (4, _s(128, 8, 1, 6)), (8, _s(128, 16, 1, 6))),
- # attention k / v (2 kv heads x 256) and the MoE router (512 experts)
- (512, 2560): ((2, _d(64, 2)), (8, _d(64, 1)), (16, _d(32, 1)), (32, _s(64, 8, 4, 10))),
- # QSA index projection: 4 index heads x 128 + 1 key x 128
- (640, 2560): ((4, _d(64, 1)), (8, _d(32, 1)), (16, _s(64, 8, 4, 10)), (32, _s(64, 16, 4, 8))),
- # shared expert gate_up, and the same with its gate row appended (padded to 1296)
- (1280, 2560): ((2, _d(64, 1)), (4, _d(64, 2)), (8, _s(64, 8, 4, 10))),
- (1296, 2560): ((2, _d(64, 1)), (4, _d(64, 2)), (8, _s(64, 8, 4, 10))),
- # shared expert gate [2560 -> 1]
- (1, 2560): ((32, _d(64, 1)),),
- # the four GDN input projections concatenated (10240 + 6144 + 48 + 48)
- (16480, 2560): ((1, _d(64, 2)),),
- # the attention input projections concatenated: q|gate (12288), k, v (512 each),
- # index (640), zero-padded to 55 heads of 256 (14080)
- (14080, 2560): ((2, _d(64, 2)), (8, _s(128, 8, 1, 6))),
- # PLE value projection (decoder layer 1 only); the PLE key projection is [10240, 2560]
- (2560, 2560): ((1, _d(64, 1)), (2, _d(64, 2)), (4, _d(32, 2)), (8, _s(64, 8, 2, 10))),
-}
-# Measured and left on cuBLAS: (2560, 640) shared-expert down (the direct kernel
-# cannot tile K=640; split-K and TGV were within 4% of cuBLAS), the GDN concat
-# (16480, 2560) at M >= 2 (direct 3.6-5.3% ahead at M=2 across two sweeps, behind
-# from M=4), the attention concat (14080, 2560) at M >= 16, and M >= 64
-# everywhere (outside both kernels' M <= 32 policy). Not servable at all: the
-# hyper-connection mixer's first projection (fp32 output; both kernels write
-# bf16 only) and its up projection [10240, 320] (K=320 tiles neither kernel).
-
-
-@functools.lru_cache(maxsize=None)
-def plan(m: int, n: int, k: int) -> _Plan:
- """The kernel and tactic `LowMBf16Mm` launches for an [M, K] x [K, N] product."""
- for bound, choice in _TABLE.get((n, k), ()):
- if m <= bound:
- if choice[0] == "direct":
- t = _fi_direct.DirectTactic(
- choice[1], choice[2], _fi_direct.default_tactic(m, n, k).rows_per_block
- )
- _fi_direct.validate_tactic(t, m, n, k)
- return _Plan("direct", tuple(choice[1:]) + (t.rows_per_block,))
- _fi_splitk.validate_tactic(_fi_splitk.SplitKTactic(*choice[1:]), m, n, k)
- return _Plan("splitk", tuple(choice[1:]))
- return _Plan("cublas")
-
-
-def _cell(why: str, m: int, n: int, k: int, kind: str, pdl: bool = False) -> Cell:
- """`kind` is the branch `plan` must pick for this shape; the test holds it to that."""
- return Cell(why=why, spec=dict(m=m, n=n, k=k, kind=kind, pdl=pdl))
-
-
-class LowMBf16Mm(OpWrapper):
- """`mat_a @ mat_b` in bf16 with fp32 accumulation: one kernel launch per call.
-
- A drop-in for `catalog/gemm/cublas_mm` without bias or `out_dtype`: `mat_a`
- is dense row-major `[M, K]`, `mat_b` dense column-major `[K, N]` (a stored
- `[N, K]` weight's `.t()`), the result a fresh row-major `[M, N]` bf16 tensor.
- `plan(M, N, K)` picks the kernel on the host from three sources:
-
- - `flashinfer.gemm.kernels.dense_bf16_gemm_direct.run_direct_dense` -- SIMT,
- the weight row's whole K prefetched into registers; wins at M <= 2 and on
- narrow N;
- - `flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk.run_splitk_dense` --
- tcgen05, swap-AB, split-K reduced across a cluster in shared memory (no
- extra launch); wins at M 4-32 on the wide shapes;
- - `torch.ops.trtllm.cublas_mm` -- everything else, including every M > 32.
-
- `pdl` launches the FlashInfer kernels with programmatic dependent launch:
- they wait on the previous grid (`griddepcontrol.wait`) before reading
- `mat_a`, so it is safe after any kernel; it changes nothing on the cuBLAS
- branch.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # Every cell is a qwen3.8-flash-next-nvfp4 tp1 decode projection at M >= 2 (one row
- # takes `bf16_gemv_pdl` in the target), one per branch; `kind` pins the branch `plan`
- # takes for it. The test's host-only plan pins hold the table's other shapes.
- CELLS: tuple[Cell, ...] = (
- _cell(
- "qwen3.8 decode M=2 GDN in_z: direct, 2-row tile, pdl on",
- 2,
- 6144,
- 2560,
- "direct",
- pdl=True,
- ),
- _cell(
- "qwen3.8 decode M=32 o_proj: split-K, mma_n 32, pdl on (M boundary)",
- 32,
- 2560,
- 6144,
- "splitk",
- pdl=True,
- ),
- _cell("qwen3.8 decode M=128: cuBLAS branch above M=32", 128, 10240, 2560, "cublas"),
- )
-
- note = """
- Metadata consumed: none. Preconditions: bf16 operands, `mat_a` dense
- row-major, `mat_b` dense column-major, both 32-byte aligned (a freshly
- allocated activation and a `.t()` weight view are); the FlashInfer kernels
- raise ValueError on a layout or alignment they cannot read, the cuBLAS
- branch does not (see `is_valid`).
-
- Compilation: the FlashInfer kernels are CuTe-DSL, JIT-compiled in-process
- on the first call per (direct: M, N, K, tactic, pdl; split-K: tactic, pdl)
- and cached for the process. That first call must happen before CUDA-graph
- capture -- the engine's warmup forward at each captured batch size does it.
-
- Host cost: the FlashInfer branches spend tens of microseconds of Python per
- call (argument validation, DLPack conversion), which a CUDA graph removes
- and an eager step does not. Route only graph-captured decode calls here.
-
- Numerics: every branch accumulates in fp32 and rounds once; the branches
- differ from each other only in summation order.
-
- Hard limits
- M FlashInfer branches only at M <= 32 (both kernels' policy) boundary 32 driven at 32
- (cells) and 128 (cuBLAS cell)
- N*K largest FlashInfer-branch weight 16480*2560 = 4.2e7 elements boundary 2^31 not probed (48x below it)
- """
-
- def __call__(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = False) -> torch.Tensor:
- m, k = mat_a.shape
- n = mat_b.shape[1]
- p = plan(m, n, k)
- if p.kind == "direct":
- out = torch.empty((m, n), dtype=mat_a.dtype, device=mat_a.device)
- return _fi_direct.run_direct_dense(
- mat_a, mat_b, out, pdl, _fi_direct.DirectTactic(*p.tactic)
- )
- if p.kind == "splitk":
- out = torch.empty((m, n), dtype=mat_a.dtype, device=mat_a.device)
- return _fi_splitk.run_splitk_dense(
- mat_a, mat_b, None, out, pdl, _fi_splitk.SplitKTactic(*p.tactic)
- )
- return torch.ops.trtllm.cublas_mm(mat_a, mat_b, None, None, 0, None)
-
- def reference(
- self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = False
- ) -> torch.Tensor:
- """fp32-accumulated, TF32 off; `pdl` is a launch attribute and is ignored."""
- with true_fp32_matmul():
- return (mat_a.float() @ mat_b.float()).to(mat_a.dtype)
-
- def is_valid(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = False) -> None:
- # The cuBLAS branch ignores both leading dimensions: a row-padded
- # mat_a or a column-padded mat_b (a slice of a wider tensor) is read as
- # if dense -- silently a different product. The FlashInfer branches
- # raise on those. A dense row-major mat_b is outside the contract too
- # (cuBLAS and the direct kernel raise on it; split-K computes it).
- assert mat_a.dtype == torch.bfloat16 and mat_b.dtype == torch.bfloat16, "bf16 operands only"
- assert mat_a.is_contiguous(), "mat_a must be dense row-major [M, K]"
- assert mat_b.stride(0) == 1 and mat_b.stride(1) == mat_b.shape[0], (
- "mat_b must be dense column-major [K, N] (e.g. weight.t())"
- )
-
-
-low_m_bf16_mm = LowMBf16Mm()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/mid_rows_mm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/mid_rows_mm.py
deleted file mode 100644
index cd5c8a6eb386..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/mid_rows_mm.py
+++ /dev/null
@@ -1,1380 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Dense projections of 17..256 decode rows on tcgen05: `y[M, N] = x[M, K] @ W^T`, bf16 out.
-
-Two operand modes, one kernel:
-
-- FP8: `W` is e4m3 codes `[N, K]` with one fp32 scale (or one per weight row); the bf16
- activation is quantized first by `fp8_row_quant` -- one launch, one e4m3 scale per row
- (`max |x[m, :]| / 448`) -- and the GEMM applies `scale_x[m] * scale_w` in its epilogue.
-- bf16: `W` is the bf16 weight and `x` is read as it is.
-
-CuTe DSL (tcgen05), swap-AB: the MMA's A operand is 128 weight rows (`BM`), its B the
-activation rows (N = the call's rows rounded up to 32 / 64 / 128 / 256, rows past M
-zero-filled by TMA), so the ring holds mostly weight bytes. One CTA a 128-row weight tile and
-`K / KS` of K; the `KS` CTAs of one tile are a thread-block cluster. A TMA warp streams the
-weight tile and the activation rows into a ring of shared-memory stages (a rolling window of
-weight k-tiles past the ring is prefetched to L2), one MMA warp accumulates in tensor memory,
-eight epilogue warps (two a TMEM lane quadrant, alternate column tiles) read the accumulator.
-With `KS = 1` they scale it, round it to bf16 into shared memory transposed (over the drained
-ring), and every warp stores the tile in 16-byte rows. With `KS > 1` the tile is cut into
-units of 32 weight rows x 4 activation rows owned round-robin by the cluster's ranks: each
-CTA pushes its fp32 partial of a unit into the owner's mailbox (`st.async`, counted on the
-owner's mbarrier), and the owner sums the `KS` partials in rank order -- deterministic -- and
-stores them. The mailbox is shared memory of its own, or the drained ring behind a cluster
-barrier when it is large (256 rows).
-
-Programmatic dependent launch: the weight stages that fit the ring are issued before
-`griddepcontrol.wait` (the weight is static), the activation after it; dependents are released
-at kernel start, and the CTA gives up its tensor-memory allocation permit as soon as it has
-allocated, so a CTA of the next launch is scheduled beside it (two CTAs an SM) and fills its
-own ring meanwhile. `fp8_row_quant` releases its dependents at its start, so the GEMM after it
-launches and streams its weight while the quantization runs.
-"""
-
-from __future__ import annotations
-
-import functools
-from dataclasses import dataclass
-
-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
-from cutlass import const_expr
-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,
- true_fp32_matmul,
-)
-from tensorrt_llm._torch.kernels.blackwell.custom_pipeline import PipelineTmaUmma, PipelineUmmaAsync
-from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
-
-#: Rows one launch takes: the largest MMA N.
-MAX_ROWS = 256
-E4M3_MAX = 448.0
-#: Floor of a row's maximum: an all-zero row quantizes to zeros at this scale.
-AMAX_FLOOR = 1e-12
-_SMEM_CAP = 232448 - 1024
-
-
-# ------------------------------------------------------------------------------------------
-# inline PTX
-# ------------------------------------------------------------------------------------------
-
-
-def _asm(ret, args, text, cons, loc, ip, side=False):
- return llvm.inline_asm(
- ret,
- [a.ir_value(loc=loc, ip=ip) for a in args],
- text,
- cons,
- has_side_effects=side,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
-
-
-def _unpack(res, n, ty, loc, ip):
- mlir_ty = T.i32() if ty is cutlass.Int32 else T.f32()
- return [ty(llvm.extractvalue(mlir_ty, res, [i], loc=loc, ip=ip)) for i in range(n)]
-
-
-@dsl_user_op
-def _ldg_v4(addr, *, loc=None, ip=None):
- """16 bytes from global as four words."""
- res = _asm(
- llvm.StructType.get_literal([T.i32()] * 4),
- [cutlass.Int64(addr)],
- "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];",
- "=r,=r,=r,=r,l",
- loc,
- ip,
- side=True,
- )
- return _unpack(res, 4, cutlass.Int32, loc, ip)
-
-
-@dsl_user_op
-def _ldg_f32(addr, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(T.f32(), [cutlass.Int64(addr)], "ld.global.nc.f32 $0, [$1];", "=f,l", loc, ip, True)
- )
-
-
-@dsl_user_op
-def _stg_f32(addr, v, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [cutlass.Int64(addr), cutlass.Float32(v)],
- "{ st.global.f32 [$1], $2; mov.u32 $0, 0; }",
- "=r,l,f",
- loc,
- ip,
- True,
- )
-
-
-@dsl_user_op
-def _stg_bf16(addr, v, *, loc=None, ip=None):
- """`v` rounded to bf16, to global."""
- _asm(
- T.i32(),
- [cutlass.Int64(addr), cutlass.Float32(v)],
- "{ .reg .b16 h; cvt.rn.bf16.f32 h, $2; st.global.b16 [$1], h; mov.u32 $0, 0; }",
- "=r,l,f",
- loc,
- ip,
- True,
- )
-
-
-@dsl_user_op
-def _stg_v2(addr, a, b, *, loc=None, ip=None):
- _asm(
- 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,
- True,
- )
-
-
-@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,
- True,
- )
-
-
-@dsl_user_op
-def _bf16x2(lo, hi, *, loc=None, ip=None):
- """Two fp32 rounded to bf16, `lo` in the low half."""
- return cutlass.Int32(
- _asm(
- T.i32(),
- [cutlass.Float32(lo), cutlass.Float32(hi)],
- "cvt.rn.bf16x2.f32 $0, $2, $1;",
- "=r,f,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _sts_v4f(saddr, a, b, c, d, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [cutlass.Int32(saddr)] + [cutlass.Float32(v) for v in (a, b, c, d)],
- "{ st.shared.v4.f32 [$1], {$2, $3, $4, $5}; mov.u32 $0, 0; }",
- "=r,r,f,f,f,f",
- loc,
- ip,
- True,
- )
-
-
-@dsl_user_op
-def _sts_f32(saddr, v, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [cutlass.Int32(saddr), cutlass.Float32(v)],
- "{ st.shared.f32 [$1], $2; mov.u32 $0, 0; }",
- "=r,r,f",
- loc,
- ip,
- True,
- )
-
-
-@dsl_user_op
-def _sts_bf16(saddr, v, *, loc=None, ip=None):
- """`v` rounded to bf16, to shared memory."""
- _asm(
- T.i32(),
- [cutlass.Int32(saddr), cutlass.Float32(v)],
- "{ .reg .b16 h; cvt.rn.bf16.f32 h, $2; st.shared.b16 [$1], h; mov.u32 $0, 0; }",
- "=r,r,f",
- loc,
- ip,
- True,
- )
-
-
-@dsl_user_op
-def _lds_v4f(saddr, *, loc=None, ip=None):
- res = _asm(
- llvm.StructType.get_literal([T.f32()] * 4),
- [cutlass.Int32(saddr)],
- "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];",
- "=f,=f,=f,=f,r",
- loc,
- ip,
- side=True,
- )
- return _unpack(res, 4, cutlass.Float32, loc, ip)
-
-
-@dsl_user_op
-def _lds_v4(saddr, *, loc=None, ip=None):
- res = _asm(
- llvm.StructType.get_literal([T.i32()] * 4),
- [cutlass.Int32(saddr)],
- "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];",
- "=r,=r,=r,=r,r",
- loc,
- ip,
- side=True,
- )
- return _unpack(res, 4, cutlass.Int32, loc, ip)
-
-
-@dsl_user_op
-def _st_async_v4f(raddr, rbar, a, b, c, d, *, loc=None, ip=None):
- """Four fp32 into a cluster peer's shared memory (`raddr`, from `mapa`), crediting 16
- bytes to the peer's mbarrier `rbar`."""
- _asm(
- T.i32(),
- [cutlass.Int32(raddr), cutlass.Int32(rbar)] + [cutlass.Float32(v) for v in (a, b, c, d)],
- "{ .reg .b32 x, y, z, w; mov.b32 x, $3; mov.b32 y, $4; mov.b32 z, $5; mov.b32 w, $6; "
- "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [$1], {x, y, z, w}, [$2]; "
- "mov.u32 $0, 0; }",
- "=r,r,r,f,f,f,f",
- loc,
- ip,
- True,
- )
-
-
-@dsl_user_op
-def _mbar_expect_tx(addr, nbytes, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [cutlass.Int32(addr), cutlass.Int32(nbytes)],
- "{ mbarrier.arrive.expect_tx.shared::cta.b64 _, [$1], $2; mov.u32 $0, 0; }",
- "=r,r,r",
- loc,
- ip,
- True,
- )
-
-
-@dsl_user_op
-def _mbar_wait(addr, parity, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [cutlass.Int32(addr), cutlass.Int32(parity)],
- "{ .reg .pred p; WAIT: mbarrier.try_wait.parity.shared::cta.b64 p, [$1], $2; "
- "@!p bra WAIT; mov.u32 $0, 0; }",
- "=r,r,r",
- loc,
- ip,
- 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 _absmax_bf16x2(acc, w, *, loc=None, ip=None):
- """Per half, max(acc, |w|) over bf16x2 words (non-negative halves)."""
- return cutlass.Int32(
- _asm(
- T.i32(),
- [cutlass.Int32(acc), cutlass.Int32(w)],
- "{ .reg .b32 a; and.b32 a, $2, 0x7fff7fff; max.bf16x2 $0, $1, a; }",
- "=r,r,r",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _row_amax(v, *, loc=None, ip=None):
- """The warp's largest |x| from per-lane bf16x2 maxima, as fp32: the two halves' larger
- one (as unsigned bits: the halves are non-negative), then `redux.sync.max.u32`."""
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Int32(v)],
- "{ .reg .b32 l, h; and.b32 l, $1, 0xffff; shr.b32 h, $1, 16; max.u32 l, l, h; "
- "redux.sync.max.u32 l, l, 0xffffffff; shl.b32 l, l, 16; mov.b32 $0, l; }",
- "=f,r",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _e4m3x4(w0, w1, inv, *, loc=None, ip=None):
- """Four bf16 (two words) times `inv`, each rounded to e4m3 (saturating), packed in order."""
- return cutlass.Int32(
- _asm(
- T.i32(),
- [cutlass.Int32(w0), cutlass.Int32(w1), cutlass.Float32(inv)],
- "{ .reg .b32 t0, t1, t2, t3; .reg .f32 f0, f1, f2, f3; .reg .b16 p, q; "
- "shl.b32 t0, $1, 16; and.b32 t1, $1, 0xffff0000; "
- "shl.b32 t2, $2, 16; and.b32 t3, $2, 0xffff0000; "
- "mov.b32 f0, t0; mov.b32 f1, t1; mov.b32 f2, t2; mov.b32 f3, t3; "
- "mul.rn.f32 f0, f0, $3; mul.rn.f32 f1, f1, $3; "
- "mul.rn.f32 f2, f2, $3; mul.rn.f32 f3, f3, $3; "
- "cvt.rn.satfinite.e4m3x2.f32 p, f1, f0; cvt.rn.satfinite.e4m3x2.f32 q, f3, f2; "
- "mov.b32 $0, {p, q}; }",
- "=r,r,r,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _div_rn(a, b, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Float32(a), cutlass.Float32(b)],
- "div.rn.f32 $0, $1, $2;",
- "=f,f,f",
- loc,
- ip,
- )
- )
-
-
-# ------------------------------------------------------------------------------------------
-# per-row activation quantization
-# ------------------------------------------------------------------------------------------
-
-
-class _RowQuant:
- """One warp a row: `scale[m] = max(max |x[m, :]|, AMAX_FLOOR) / 448`,
- `q[m, k] = e4m3_satfinite(x[m, k] * (448 / max))`. K a multiple of 256."""
-
- WARPS = 4
-
- def __init__(self, k: int, pdl: bool):
- assert k % 256 == 0
- self.k, self.pdl = k, pdl
- self.nv = k // 256 # 16-byte vectors a lane
-
- @cute.jit
- def __call__(
- self,
- m: cutlass.Int32,
- x_ptr: cute.Pointer,
- x_stride: cutlass.Int32,
- q_t: cute.Tensor,
- s_ptr: cute.Pointer,
- stream: cuda.CUstream,
- ):
- self.kernel(x_ptr.toint(), x_stride, q_t.iterator.toint(), s_ptr.toint(), m).launch(
- grid=[(m + self.WARPS - 1) // self.WARPS, 1, 1],
- block=[32 * self.WARPS, 1, 1],
- stream=stream,
- use_pdl=self.pdl,
- )
-
- @cute.kernel
- def kernel(
- self,
- x_base: cutlass.Int64,
- x_stride: cutlass.Int32,
- q_base: cutlass.Int64,
- s_base: cutlass.Int64,
- m: cutlass.Int32,
- ):
- bidx, _, _ = cute.arch.block_idx()
- warp = cute.arch.make_warp_uniform(cute.arch.warp_idx())
- lane = cute.arch.lane_idx()
- if const_expr(self.pdl):
- cute.arch.griddepcontrol_launch_dependents()
- cute.arch.griddepcontrol_wait()
- row = bidx * self.WARPS + warp
- if row < m:
- xr = (
- x_base + cutlass.Int64(row) * cutlass.Int64(x_stride) * 2 + cutlass.Int64(lane * 16)
- )
- vals = [_ldg_v4(xr + v * 512) for v in range(self.nv)]
- mx = cutlass.Int32(0)
- for v in cutlass.range_constexpr(self.nv):
- for j in cutlass.range_constexpr(4):
- mx = _absmax_bf16x2(mx, vals[v][j])
- amax = cutlass.max(_row_amax(mx), cutlass.Float32(AMAX_FLOOR))
- inv = _div_rn(cutlass.Float32(E4M3_MAX), amax)
- qr = q_base + cutlass.Int64(row) * self.k + cutlass.Int64(lane * 8)
- for v in cutlass.range_constexpr(self.nv):
- _stg_v2(
- qr + v * 256,
- _e4m3x4(vals[v][0], vals[v][1], inv),
- _e4m3x4(vals[v][2], vals[v][3], inv),
- )
- if lane == 0:
- _stg_f32(s_base + cutlass.Int64(row) * 4, _div_rn(amax, cutlass.Float32(E4M3_MAX)))
-
-
-# ------------------------------------------------------------------------------------------
-# GEMM
-# ------------------------------------------------------------------------------------------
-
-#: Weight rows one CTA computes: the MMA's M.
-BM = 128
-
-
-@dataclass(frozen=True)
-class Plan:
- """`ks` K split (cluster size); `kpb` MMA k-blocks a ring stage (4: 128 bytes of K a row);
- `stages` ring depth; `l2` 1 to prefetch the CTA's weight slice past the ring into L2 at
- the start. A CTA computes `BM` weight rows for every activation row."""
-
- ks: int = 1
- kpb: int = 4
- stages: int = 4
- l2: int = 1
- epi: int = 16
-
-
-def mma_rows(m: int) -> int:
- """The MMA N (activation rows, zero-padded by TMA) a call of `m` rows runs at."""
- for b in (32, 64, 128, 256):
- if m <= b:
- return b
- raise ValueError(f"{m} rows: at most {MAX_ROWS}")
-
-
-class _MidMm:
- """See the module docstring. Swap-AB: the MMA's A operand is `BM` weight rows, its B the
- activation rows (N = `mpad`). Compile-time: N, K, the plan, `mpad`, the operand mode and
- the weight-scale mode (0 none / 1 one fp32 / 2 one a weight row). Runtime: the row count
- (<= mpad) and the pointers."""
-
- def __init__(self, *, n, k, plan: Plan, mpad, fp8, w_mode, pdl):
- epi_cols = min(plan.epi, mpad)
- assert mpad % 16 == 0 and 16 <= mpad <= 256 and mpad % epi_cols == 0
- self.n, self.k, self.plan, self.mpad, self.fp8, self.w_mode, self.pdl = (
- n,
- k,
- plan,
- mpad,
- fp8,
- w_mode,
- pdl,
- )
- self.ks = plan.ks
- self.io_dtype = cutlass.Float8E4M3FN if fp8 else cutlass.BFloat16
- self.acc_dtype = cutlass.Float32
- self.epi_cols = epi_cols
- self.cta_group = tcgen05.CtaGroup.ONE
- self.cluster_shape_mn = (1, 1)
- # Two warps a TMEM lane quadrant: warpgroup h reads the column tiles `sub % 2 == h`.
- self.epilogue_warp_id = tuple(range(8))
- self.mma_warp_id = 8
- self.tma_warp_id = 9
- self.threads = 32 * 10
- self.epilog_sync_bar_id = 1
- self.tmem_alloc_sync_bar_id = 2
- self.n_tiles = -(-n // BM)
- # The tile staged for the stores, `[mpad][BM]` bf16 (KS = 1, over the drained ring).
- self.stage_bytes = mpad * BM * 2 if plan.ks == 1 else 0
- # The split-K mailbox (KS > 1): the most units a rank owns, KS slots of 512 bytes each;
- # a large one reuses the drained ring (after a cluster barrier: every peer's mainloop is
- # done before any partial lands).
- mbox = -(-mpad // plan.ks) * plan.ks * 512 if plan.ks > 1 else 0
- self.alias_mbox = mbox > (64 << 10)
- self.mbox_bytes = 0 if self.alias_mbox else mbox
- self.mbox_need = mbox
-
- def _setup(self):
- k_major = utils.LayoutEnum.ROW_MAJOR.mma_major_mode()
- tiled_mma = utils.sm100.make_trivial_tiled_mma(
- self.io_dtype,
- self.io_dtype,
- k_major,
- k_major,
- self.acc_dtype,
- self.cta_group,
- (BM, self.mpad),
- )
- inst_k = cute.size(tiled_mma.shape_mnk, mode=[2])
- self.mma_tiler = (BM, self.mpad, inst_k * self.plan.kpb)
- 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 = (BM, self.epi_cols)
- assert self.k % (self.mma_tiler[2] * self.ks) == 0, (self.k, self.mma_tiler, self.ks)
- self.k_tiles = self.k // self.mma_tiler[2] // self.ks
- self.stages = self._stages()
- self.n_pre = self.stages
- self.a_smem_layout_staged = utils.sm100.make_smem_layout_a(
- tiled_mma, self.mma_tiler, self.io_dtype, self.stages
- )
- self.b_smem_layout_staged = utils.sm100.make_smem_layout_b(
- tiled_mma, self.mma_tiler, self.io_dtype, self.stages
- )
- acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
- acc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, 1))
- self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(acc_fake, arch="sm_100")
- return tiled_mma
-
- def _stages(self) -> int:
- """The plan's ring depth, at most the CTA's K tiles and what fits beside the mailbox."""
- bpe = 1 if self.fp8 else 2
- inst_k = 32 if self.fp8 else 16
- kb = inst_k * self.plan.kpb * bpe
- k_tiles = self.k // (inst_k * self.plan.kpb) // self.ks
- fit = (_SMEM_CAP - self.mbox_bytes - 4096) // ((BM + self.mpad) * kb)
- return max(1, min(self.plan.stages, k_tiles, fit))
-
- def smem_bytes(self) -> int:
- """Host-side estimate of the dynamic shared memory (ring or staged tile, mailbox)."""
- bpe = 1 if self.fp8 else 2
- inst_k = 32 if self.fp8 else 16
- ring = self._stages() * (BM + self.mpad) * inst_k * self.plan.kpb * bpe
- return max(ring, self.stage_bytes, self.mbox_need if self.alias_mbox else 0) + (
- self.mbox_bytes + 1024 + 2048
- )
-
- @cute.jit
- def __call__(
- self,
- m: cutlass.Int32,
- x_ptr: cute.Pointer,
- w_ptr: cute.Pointer,
- out_t: cute.Tensor,
- sx_ptr: cute.Pointer,
- sw_ptr: cute.Pointer,
- stream: cuda.CUstream,
- ):
- k = self.k
- a = cute.make_tensor(w_ptr, cute.make_layout((self.n, k, 1), stride=(k, 1, 0)))
- b = cute.make_tensor(x_ptr, cute.make_layout((m, k, 1), stride=(k, 1, 0)))
- tiled_mma = self._setup()
- 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
- )
- kern = self.kernel(
- tiled_mma,
- tma_atom_a,
- tma_a,
- tma_atom_b,
- tma_b,
- out_t.iterator.toint(),
- w_ptr.toint(),
- sx_ptr.toint(),
- sw_ptr.toint(),
- m,
- self.cluster_layout_vmnk,
- self.a_smem_layout_staged,
- self.b_smem_layout_staged,
- self.epi_tile,
- )
- # All of the SM's L1 as shared memory: a CTA of the next launch then fits beside this
- # one (programmatic dependent launch) and streams its weight stages early.
- if const_expr(self.ks > 1):
- kern.launch(
- grid=[self.ks, self.n_tiles, 1],
- block=[self.threads, 1, 1],
- cluster=[self.ks, 1, 1],
- stream=stream,
- use_pdl=self.pdl,
- preferred_smem_carveout=100,
- )
- else:
- kern.launch(
- grid=[1, self.n_tiles, 1],
- block=[self.threads, 1, 1],
- stream=stream,
- use_pdl=self.pdl,
- preferred_smem_carveout=100,
- )
-
- @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,
- out_base: cutlass.Int64,
- w_base: cutlass.Int64,
- sx_base: cutlass.Int64,
- sw_base: cutlass.Int64,
- 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())
- if const_expr(self.pdl):
- 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 # weight-row tile
- ks: cutlass.Constexpr = self.ks
- mpad: cutlass.Constexpr = self.mpad
- n0 = nt * BM
-
- @cute.struct
- class SharedStorage:
- ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.stages * 2]
- acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2]
- tmem_dealloc_mbar: cutlass.Int64
- red_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.stages,
- 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=1,
- 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=1024,
- swizzle=a_smem_layout_staged.inner,
- )
- sB = smem.allocate_tensor(
- element_type=self.io_dtype,
- layout=b_smem_layout_staged.outer,
- byte_alignment=1024,
- swizzle=b_smem_layout_staged.inner,
- )
- ring_bytes = self.stages * (BM + mpad) * self.mma_tiler[2] * (1 if self.fp8 else 2)
- if const_expr(self.stage_bytes > ring_bytes):
- smem.allocate_tensor(
- element_type=cutlass.Float32,
- layout=cute.make_layout(((self.stage_bytes - ring_bytes) // 4 + 4,)),
- byte_alignment=16,
- )
- sXt = smem.allocate_tensor(
- element_type=cutlass.Float32, layout=cute.make_layout((256,)), byte_alignment=16
- )
- sx_s = cutlass.Int32(sXt.iterator.toint())
- red_bar = cutlass.Int32(storage.red_mbar.ptr.toint())
- mbox = cutlass.Int32(0)
- if const_expr(ks > 1):
- # The split-K mailbox: units of 32 weight rows x 4 activation rows (fp32, 512
- # bytes), unit u owned by rank u % KS; slot (u // KS, sender) of the owner.
- if const_expr(self.alias_mbox):
- ring_b = self.stages * (BM + mpad) * self.mma_tiler[2] * (1 if self.fp8 else 2)
- if const_expr(self.mbox_need > ring_b):
- smem.allocate_tensor(
- element_type=cutlass.Float32,
- layout=cute.make_layout(((self.mbox_need - ring_b) // 4 + 4,)),
- byte_alignment=16,
- )
- mbox = cutlass.Int32(sA.iterator.toint())
- else:
- sMb = smem.allocate_tensor(
- element_type=cutlass.Float32,
- layout=cute.make_layout((self.mbox_bytes // 4,)),
- byte_alignment=16,
- )
- mbox = cutlass.Int32(sMb.iterator.toint())
- if tidx == 0:
- cute.arch.mbarrier_init(storage.red_mbar.ptr, 1)
- cute.arch.mbarrier_init_fence()
- # Peers st.async into this mailbox only after this CTA's barrier is initialized.
- cute.arch.cluster_arrive_relaxed()
- # The output tile (bf16, KS = 1) or the fp32 partial (KS > 1) reuses the ring (drained
- # once the accumulator is complete): it starts at the A stages, the B stages follow.
- p_base = cutlass.Int32(sA.iterator.toint())
-
- 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, 1))
-
- pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
- if const_expr(ks > 1):
- cute.arch.cluster_wait()
-
- if warp_idx == self.tma_warp_id:
- tAgA_slice = tAgA[(None, nt, None, 0)]
- tBgB_slice = tBgB[(None, 0, None, 0)]
- ab_producer.reset()
- # The weight does not depend on the previous kernel: the ring's stages of it
- # stream before the grid-dependency wait, the activation rows after it.
- pre = []
- for kk in cutlass.range_constexpr(self.n_pre):
- handle = ab_producer.acquire_and_advance()
- cute.copy(
- tma_atom_a,
- tAgA_slice[(None, k0 + kk)],
- tAsA[(None, handle.index)],
- tma_bar_ptr=handle.barrier,
- )
- pre.append(handle)
- l2a: cutlass.Constexpr = min(self.plan.l2, self.k_tiles - self.n_pre)
- if const_expr(l2a > 0):
- # A rolling window of `l2` weight k-tiles past the ring prefetched towards L2
- # (TMA prefetch): the HBM stream does not wait for ring slots.
- for kk in cutlass.range_constexpr(self.n_pre, self.n_pre + l2a):
- cute.prefetch(tma_atom_a, tAgA_slice[(None, k0 + kk)])
- if const_expr(self.pdl):
- cute.arch.griddepcontrol_wait()
- for kk in cutlass.range_constexpr(self.n_pre):
- handle = pre[kk]
- cute.copy(
- tma_atom_b,
- tBgB_slice[(None, k0 + kk)],
- tBsB[(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()
- if const_expr(l2a > 0):
- if k_tile + l2a < self.k_tiles:
- cute.prefetch(tma_atom_a, tAgA_slice[(None, k0 + k_tile + l2a)])
- 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 const_expr(ks > 1 and self.alias_mbox):
- cute.arch.cluster_arrive_relaxed()
- cute.arch.cluster_wait()
- 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)
- tCtAcc = tCtAcc_base[(None, None, None, 0)]
- acc_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1)
- ab_consumer.reset()
- peek = ab_consumer.try_wait()
- acc_pipeline.producer_acquire(acc_producer_state)
- num_kblocks = cute.size(tCrA, mode=[2])
- for k_tile in cutlass.range(0, self.k_tiles, 1, unroll=1):
- handle = ab_consumer.wait_and_advance(peek)
- for kb in cutlass.range_constexpr(num_kblocks):
- if const_expr(kb == 0):
- tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0)
- else:
- tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
- cute.gemm(
- tiled_mma,
- tCtAcc,
- tCrA[(None, None, kb, handle.index)],
- tCrB[(None, None, kb, handle.index)],
- tCtAcc,
- )
- 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()
- if const_expr(ks > 1 and self.alias_mbox):
- cute.arch.cluster_arrive_relaxed()
- acc_pipeline.producer_tail(acc_producer_state)
- if const_expr(ks > 1 and self.alias_mbox):
- cute.arch.cluster_wait()
-
- if warp_idx < self.mma_warp_id:
- tmem.allocate(self.num_tmem_alloc_cols)
- # This CTA allocates no more tensor memory: a CTA of the next launch may be
- # scheduled on this SM now (programmatic dependent launch).
- tmem.relinquish_alloc_permit()
- if const_expr(self.fp8):
- # The activation scales (the previous launch's output) into shared memory.
- if const_expr(self.pdl):
- cute.arch.griddepcontrol_wait()
- if tidx < m:
- _sts_f32(sx_s + tidx * 4, _ldg_f32(sx_base + cutlass.Int64(tidx) * 4))
- 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, 1)
- 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 % BM)
- half = tidx // BM
- tTR_tAcc_base = thr_copy_t2r.partition_S(tAcc_epi)
- tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, 0)]
- tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
- # This thread's accumulator row (a weight row) and first column (an activation
- # row): the 32dp32b load gives a thread one row's `epi_cols` consecutive columns.
- tCcC = thr_mma.partition_C(cute.make_identity_tensor((BM, mpad)))
- tTR_cC = thr_copy_t2r.partition_D(
- cute.flat_divide(tCcC[((None, None), 0, 0)], epi_tile)
- )
- frag_c = tTR_cC[(None, None, None, 0, 0)]
- crd0 = frag_c[0]
- r_loc = crd0[0]
- c_loc = crd0[1]
- tTR_rAcc = cute.make_rmem_tensor(frag_c.shape, self.acc_dtype)
- nfrag: cutlass.Constexpr = self.epi_cols
- n = n0 + r_loc
- sw = cutlass.Float32(1.0)
- if const_expr(ks == 1):
- if const_expr(self.w_mode == 1):
- sw = _ldg_f32(sw_base)
- if const_expr(self.w_mode == 2):
- if n < self.n:
- sw = _ldg_f32(sw_base + cutlass.Int64(n) * 4)
- cute.arch.barrier(
- barrier_id=self.epilog_sync_bar_id,
- number_of_threads=32 * len(self.epilogue_warp_id),
- )
- vc = (m + 3) // 4 # 4-row chunks of activation rows that hold real rows
- n_own = (4 * vc - kr + ks - 1) // ks # units this rank owns
- if const_expr(ks > 1):
- if tidx == 0:
- _mbar_expect_tx(red_bar, n_own * (ks - 1) * 512)
- acc_pipeline.consumer_wait(acc_consumer_state)
- if const_expr(ks > 1 and self.alias_mbox):
- cute.arch.cluster_arrive_relaxed()
- cute.arch.cluster_wait()
- # The tile goes to shared memory transposed, `[activation row][weight row]` (over
- # the drained ring): bf16 and scaled (KS = 1) or the fp32 partial (KS > 1).
- for sub in cutlass.range_constexpr(mpad // self.epi_cols):
- c0 = sub * self.epi_cols + c_loc
- if (c0 < m) & (half == sub % 2):
- cute.copy(tiled_copy_t2r, tTR_tAcc[(None, None, None, sub)], tTR_rAcc)
- for q in cutlass.range_constexpr(nfrag // 4):
- if const_expr(ks == 1):
- if const_expr(self.fp8):
- s4 = _lds_v4f(sx_s + (c0 + 4 * q) * 4)
- for e in cutlass.range_constexpr(4):
- i = 4 * q + e
- v = tTR_rAcc[i] * sw
- if const_expr(self.fp8):
- v = v * s4[e]
- _sts_bf16(p_base + ((c0 + i) * BM + r_loc) * 2, v)
- else:
- cq = (c0 // 4) + q
- if cq * 4 < m:
- u = (r_loc // 32) * vc + cq
- owner = u % ks
- slot = mbox + (((u // ks) * ks + kr) * 32 + r_loc % 32) * 16
- if owner == kr:
- _sts_v4f(
- slot,
- tTR_rAcc[4 * q],
- tTR_rAcc[4 * q + 1],
- tTR_rAcc[4 * q + 2],
- tTR_rAcc[4 * q + 3],
- )
- else:
- _st_async_v4f(
- _mapa(slot, owner),
- _mapa(red_bar, owner),
- tTR_rAcc[4 * q],
- tTR_rAcc[4 * q + 1],
- tTR_rAcc[4 * q + 2],
- tTR_rAcc[4 * q + 3],
- )
- 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.free(tmem_ptr)
- # The tile's stores (KS = 1): an item is 8 weight rows of one activation row (16 bytes
- # out), consecutive threads take consecutive items (coalesced).
- g8: cutlass.Constexpr = BM // 8
- if const_expr(ks == 1):
- cute.arch.barrier()
- for f in cutlass.range(tidx, m * g8, self.threads, unroll=1):
- c = f // g8
- g = f - c * g8
- n2 = n0 + 8 * g
- if n2 < self.n:
- w4 = _lds_v4(p_base + (c * BM + 8 * g) * 2)
- _stg_v4(
- out_base + (cutlass.Int64(c) * self.n + cutlass.Int64(n2)) * 2,
- w4[0],
- w4[1],
- w4[2],
- w4[3],
- )
- else:
- # The owner sums its units' KS partials (its own and the peers', which arrive on
- # the mailbox barrier) in rank order and stores them: an item is one unit's weight
- # row (lane) and its 4 activation rows; consecutive threads take consecutive rows.
- if warp_idx < self.mma_warp_id:
- _mbar_wait(red_bar, 0)
- sw2 = cutlass.Float32(1.0)
- if const_expr(self.w_mode == 1):
- sw2 = _ldg_f32(sw_base)
- vc2 = (m + 3) // 4
- n_own2 = (4 * vc2 - kr + ks - 1) // ks
- for f in cutlass.range(
- tidx, n_own2 * 32, 32 * len(self.epilogue_warp_id), unroll=1
- ):
- j = f // 32
- ln = f - j * 32
- u = kr + j * ks
- w_ = u // vc2
- cq = u - w_ * vc2
- n2 = n0 + w_ * 32 + ln
- if n2 < self.n:
- base = mbox + ((j * ks) * 32 + ln) * 16
- parts = [_lds_v4f(base + s_ * 512) for s_ in range(ks)]
- acc = parts[0]
- for s_ in cutlass.range_constexpr(1, ks):
- acc = [acc[e] + parts[s_][e] for e in range(4)]
- sc = sw2
- if const_expr(self.w_mode == 2):
- sc = _ldg_f32(sw_base + cutlass.Int64(n2) * 4)
- if const_expr(self.fp8):
- s4 = _lds_v4f(sx_s + cq * 16)
- for e in cutlass.range_constexpr(4):
- v = acc[e] * sc
- if const_expr(self.fp8):
- v = v * s4[e]
- if cq * 4 + e < m:
- _stg_bf16(
- out_base
- + (cutlass.Int64(cq * 4 + e) * self.n + cutlass.Int64(n2)) * 2,
- v,
- )
-
-
-# ------------------------------------------------------------------------------------------
-# host side
-# ------------------------------------------------------------------------------------------
-
-
-def _env_pdl() -> bool:
- try:
- from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
- except ImportError:
- return True
- return bool(get_env_enable_pdl())
-
-
-def _stream():
- return cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
-
-
-@functools.lru_cache(maxsize=None)
-def _compiled_quant(k: int, pdl: bool):
- kern = _RowQuant(k, pdl)
- dev = "cuda"
- x = torch.empty(4, k, dtype=torch.bfloat16, device=dev)
- q = torch.empty(4, k, dtype=torch.uint8, device=dev)
- s = torch.empty(4, dtype=torch.float32, device=dev)
- return cute.compile(
- kern,
- cutlass.Int32(4),
- make_ptr(cutlass.BFloat16, x.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
- cutlass.Int32(k),
- from_dlpack(q, assumed_align=16).mark_layout_dynamic(leading_dim=1),
- make_ptr(cutlass.Float32, s.data_ptr(), cute.AddressSpace.gmem, assumed_align=4),
- _stream(),
- options="--opt-level 2 --enable-tvm-ffi",
- )
-
-
-@functools.lru_cache(maxsize=None)
-def _compiled_mm(n: int, k: int, plan: Plan, mpad: int, fp8: bool, w_mode: int, pdl: bool):
- kern = _MidMm(n=n, k=k, plan=plan, mpad=mpad, fp8=fp8, w_mode=w_mode, pdl=pdl)
- assert kern.smem_bytes() <= _SMEM_CAP, (plan, kern.smem_bytes())
- dev = "cuda"
- dt = cutlass.Float8E4M3FN if fp8 else cutlass.BFloat16
- x = torch.empty(16, dtype=torch.uint8, device=dev)
- out = torch.empty(4, n, dtype=torch.bfloat16, device=dev)
- s = torch.empty(4, dtype=torch.float32, device=dev)
- return cute.compile(
- kern,
- cutlass.Int32(4),
- make_ptr(dt, x.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
- make_ptr(dt, x.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
- from_dlpack(out, assumed_align=16).mark_layout_dynamic(leading_dim=1),
- make_ptr(cutlass.Float32, s.data_ptr(), cute.AddressSpace.gmem, assumed_align=4),
- make_ptr(cutlass.Float32, s.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
- _stream(),
- options="--opt-level 2 --enable-tvm-ffi",
- )
-
-
-def fp8_row_quant(x: torch.Tensor, pdl: bool = True) -> tuple[torch.Tensor, torch.Tensor]:
- """`(q [M, K] e4m3 contiguous, scale [M] fp32)` of a bf16 `[M, K]` activation (unit column
- stride, 16-byte aligned rows): `scale = max(max |x[m]|, AMAX_FLOOR) / 448`,
- `q = e4m3_satfinite(x * (448 / max))`."""
- m, k = x.shape
- assert x.dtype == torch.bfloat16 and x.stride(1) == 1 and k % 256 == 0
- assert x.stride(0) % 8 == 0 and x.data_ptr() % 16 == 0
- q = torch.empty((m, k), dtype=torch.float8_e4m3fn, device=x.device)
- # The GEMM reads the scales of its whole MMA tile of rows (16-byte vectors): padded.
- s = torch.empty(mma_rows(m) if m <= MAX_ROWS else m, dtype=torch.float32, device=x.device)
- if m:
- _compiled_quant(k, bool(pdl and _env_pdl()))(
- m, x.data_ptr(), x.stride(0), q.view(torch.uint8), s.data_ptr()
- )
- return q, s[:m]
-
-
-def mid_mm_raw(
- x: torch.Tensor,
- w: torch.Tensor,
- plan: Plan,
- sx: torch.Tensor | None = None,
- sw: torch.Tensor | None = None,
- pdl: bool = True,
-) -> torch.Tensor:
- """The GEMM alone: `x [M, K]` (e4m3 with `sx [M]`, or bf16) contiguous, `w [N, K]` of the
- same dtype contiguous, `sw` None / fp32 [1] / fp32 [N]. Returns a fresh bf16 `[M, N]`."""
- m, k = x.shape
- n = w.shape[0]
- fp8 = x.dtype == torch.float8_e4m3fn
- assert w.dtype == x.dtype and w.shape[1] == k and x.is_contiguous() and w.is_contiguous()
- assert 1 <= m <= MAX_ROWS and n % 16 == 0
- assert fp8 == (sx is not None)
- w_mode = 0 if sw is None else (1 if sw.numel() == 1 else 2)
- out = torch.empty((m, n), dtype=torch.bfloat16, device=x.device)
- fn = _compiled_mm(n, k, plan, mma_rows(m), fp8, w_mode, bool(pdl and _env_pdl()))
- dummy = out.data_ptr()
- fn(
- m,
- x.data_ptr(),
- w.data_ptr(),
- out,
- dummy if sx is None else sx.data_ptr(),
- dummy if sw is None else sw.data_ptr(),
- )
- return out
-
-
-# ------------------------------------------------------------------------------------------
-# plans, weights, entry
-# ------------------------------------------------------------------------------------------
-
-
-@dataclass
-class MidWeight:
- """A projection weight as `mid_rows_mm` reads it: `w` `[N, K]` contiguous, e4m3 codes with
- `scale` (fp32 `[1]`, or `[N]` one a weight row) or bf16 with `scale` None."""
-
- w: torch.Tensor
- scale: torch.Tensor | None = None
-
- @property
- def n(self) -> int:
- return self.w.shape[0]
-
- @property
- def k(self) -> int:
- return self.w.shape[1]
-
- @property
- def fp8(self) -> bool:
- return self.w.dtype == torch.float8_e4m3fn
-
-
-#: Measured plans (GB300, 152 SMs, CUDA-graph chain with PDL over HBM-cold weights; the
-#: campaign's bench/gemm_mid) of the qwen3.8-flash-next-nvfp4 tp1 projections, by MMA rows:
-#: (N, K, fp8) -> ((largest rows, Plan), ...). The ring depth is clamped to what fits.
-_TABLE: dict[tuple[int, int, bool], tuple[tuple[int, Plan], ...]] = {
- # GDN in_proj (q|k|v|z|a|b)
- (16480, 2560, True): ((256, Plan(1, 4, 8, 2)),),
- # attention in-proj (q|gate|k|v|index)
- (14080, 2560, True): ((256, Plan(1, 4, 5, 2)),),
- # GDN out / attention o
- (2560, 6144, True): ((256, Plan(6, 4, 6, 2)),),
- # shared expert gate|up, down; the router (bf16)
- (1280, 2560, False): ((256, Plan(10, 4, 4, 0)),),
- (2560, 640, False): ((256, Plan(5, 4, 2, 0)),),
- (512, 2560, False): ((256, Plan(10, 4, 4, 0)),),
- # vocab head (per-row e4m3 weight)
- (248320, 2560, True): ((64, Plan(1, 4, 4, 0)), (256, Plan(1, 4, 6, 2))),
-}
-
-
-def default_plan(n: int, k: int, m: int, fp8: bool) -> Plan:
- """The table's plan for (N, K, operand mode) at `m` rows, else a heuristic: split over K
- until about one CTA an SM, as deep a ring as fits."""
- for bound, plan in _TABLE.get((n, k, fp8), ()):
- if m <= bound:
- return plan
- kpb = 4
- kt = k // ((32 if fp8 else 16) * kpb)
- tiles = -(-n // BM)
- ks = 1
- for d in range(1, 17):
- if kt % d == 0 and tiles * d <= 152:
- ks = d
- stage = (BM + mma_rows(m)) * 128
- return Plan(ks, kpb, max(2, min(10, (200 << 10) // stage)))
-
-
-def mid_rows_mm(x: torch.Tensor, mw: MidWeight, plan: Plan | None = None, pdl: bool = True):
- """`x [M, K]` bf16 (unit column stride, 16-byte aligned rows) @ `mw`^T -> fresh bf16
- `[M, N]`, M = 1 .. MAX_ROWS. An e4m3 weight quantizes `x` per row first
- (`fp8_row_quant`)."""
- m = x.shape[0]
- if plan is None:
- plan = default_plan(mw.n, mw.k, m, mw.fp8)
- if mw.fp8:
- q, sx = fp8_row_quant(x, pdl)
- return mid_mm_raw(q, mw.w, plan, sx, mw.scale, pdl)
- if not x.is_contiguous():
- x = x.contiguous()
- return mid_mm_raw(x, mw.w, plan, None, None, pdl)
-
-
-def precompile(mw: MidWeight, lo: int = 1, hi: int = MAX_ROWS, pdl: bool = True) -> int:
- """Compile every (plan, MMA rows) `mid_rows_mm` reaches over `lo`..`hi` rows of this
- weight's geometry now rather than at a first serving call; returns how many."""
- pdl = bool(pdl and _env_pdl())
- w_mode = 0 if mw.scale is None else (1 if mw.scale.numel() == 1 else 2)
- keys = {(default_plan(mw.n, mw.k, m, mw.fp8), mma_rows(m)) for m in range(lo, hi + 1)}
- for plan, mpad in sorted(keys, key=str):
- _compiled_mm(mw.n, mw.k, plan, mpad, mw.fp8, w_mode, pdl)
- if mw.fp8:
- _compiled_quant(mw.k, pdl)
- return len(keys)
-
-
-# ------------------------------------------------------------------------------------------
-# op wrappers
-# ------------------------------------------------------------------------------------------
-
-
-def _mm_cell(why: str, n: int, k: int, m: int, mode: str) -> Cell:
- """`mode`: "tensor" (e4m3, one weight scale), "row" (e4m3, one scale a weight row), or
- "bf16"."""
- return Cell(why=why, spec=dict(n=n, k=k, m=m, mode=mode))
-
-
-class MidRowsMm(OpWrapper):
- """`mid_rows_mm(x, mw)`: `y[M, N] = x[M, K] @ W^T` in bf16 for 1..MAX_ROWS rows on tcgen05.
- With an e4m3 weight (`MidWeight(codes [N, K], scale)`, one fp32 scale or one a weight row)
- the activation is quantized to e4m3 with one scale a row (`fp8_row_quant`, its own launch)
- and the product is `(q(x) * scale_x[m]) @ (codes * scale_w)^T`; with a bf16 weight it is
- the plain product. fp32 accumulation (tensor memory; split-K partials summed in rank
- order), the scales applied in fp32, one bf16 rounding.
-
- Lossy (e4m3 operands) against the bf16 projection in the FP8 mode: ~3% relative error per
- output on random operands; a row's activation rounding depends only on its own maximum.
-
- Fusion boundary: activation quantization + GEMM (two launches) / the GEMM (one)."""
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- _mm_cell("GDN in_proj, STP con32 (32 rows), e4m3", 16480, 2560, 32, "tensor"),
- _mm_cell(
- "attention in-proj, MTP3 con32 verify (128 rows), e4m3", 14080, 2560, 128, "tensor"
- ),
- _mm_cell("GDN in_proj, 256 rows (MMA N 256), e4m3", 16480, 2560, 256, "tensor"),
- _mm_cell("GDN out / attention o, 17 rows, split-K, e4m3", 2560, 6144, 17, "tensor"),
- _mm_cell(
- "GDN out / attention o, 200 rows, split-K, mailbox over the ring, e4m3",
- 2560,
- 6144,
- 200,
- "tensor",
- ),
- _mm_cell("router, 128 rows, bf16", 512, 2560, 128, "bf16"),
- _mm_cell("shared expert gate|up, 64 rows, bf16", 1280, 2560, 64, "bf16"),
- _mm_cell("shared expert down, 256 rows, bf16", 2560, 640, 256, "bf16"),
- _mm_cell("vocab head, per-row e4m3 weight, 32 rows", 248320, 2560, 32, "row"),
- )
-
- note = """
- Tier 6 (CuTe DSL written for this target): at 17-256 decode rows cuBLAS(Lt) reaches
- ~0.4-0.6 of HBM bandwidth on these weights, and the per-tensor FP8 entry adds two
- quantization launches.
- Metadata consumed: none. Call sites: the decode target's dense projections in their
- `_MID_MM_ROWS` row ranges (`_dense_mm`), the router and shared expert, the FP8 vocab
- heads above `fp8_gemv`'s rows.
- Preconditions: `x` bf16 `[M, K]`, 1 <= M <= 256, unit column stride; the FP8 mode reads
- any row stride that is a multiple of 8 elements (16-byte aligned rows), the bf16 mode a
- contiguous `x` (others are copied). `mw.w` `[N, K]` contiguous, N a multiple of 16, K a
- multiple of 128 (e4m3) / 64 (bf16) times the plan's K split.
- Launch: grid (KS, ceil(N / 128)), cluster (KS, 1, 1) when KS > 1, 10 warps (8 epilogue,
- MMA, TMA), shared memory up to ~110 KB at 32 rows so two CTAs fit an SM.
- Programmatic dependent launch (TRTLLM_ENABLE_PDL=0 turns it off): the weight stages that
- fit the ring stream before `griddepcontrol.wait`, the activation after it; dependents are
- released at kernel start (the quantization launch releases the GEMM at its start).
- Graph-capturable (current stream, no host sync). Deterministic. Compiled per (N, K, plan,
- MMA rows (32 / 64 / 128 / 256), mode); `precompile` compiles every variant a weight's
- dispatch range reaches.
- """
-
- def __call__(self, x: torch.Tensor, mw: MidWeight) -> torch.Tensor:
- return mid_rows_mm(x, mw)
-
- def reference(self, x: torch.Tensor, mw: MidWeight) -> torch.Tensor:
- """fp32 `[M, N]`: the activation quantized as the op does (FP8 mode), both operands
- dequantized, a true fp32 matmul."""
- if mw.fp8:
- q, s = fp8_row_quant(x)
- xd = q.float() * s[:, None]
- wd = mw.w.float() * (mw.scale[:, None] if mw.scale.numel() > 1 else mw.scale)
- else:
- xd, wd = x.float(), mw.w.float()
- with true_fp32_matmul():
- return xd @ wd.t()
-
- def is_valid(self, x: torch.Tensor, mw: MidWeight) -> None:
- m, k = x.shape
- assert x.dtype == torch.bfloat16 and x.stride(1) == 1, "bf16 x with unit column stride"
- assert 1 <= m <= MAX_ROWS, f"{m} rows: 1..{MAX_ROWS}"
- assert mw.w.is_contiguous() and mw.k == k and mw.n % 16 == 0
- if mw.fp8:
- assert mw.scale is not None and mw.scale.dtype == torch.float32
- assert mw.scale.numel() in (1, mw.n), "one weight scale, or one a weight row"
- assert x.stride(0) % 8 == 0 and x.data_ptr() % 16 == 0, "16-byte aligned rows"
- assert k % 256 == 0, "K a multiple of 256 (the quantization's row vectors)"
- else:
- assert mw.w.dtype == torch.bfloat16 and mw.scale is None
-
-
-class Fp8RowQuant(OpWrapper):
- """`fp8_row_quant(x)`: `(q [M, K] e4m3, scale [M] fp32)` with `scale = max(max |x[m]|,
- AMAX_FLOOR) / 448` and `q = e4m3_satfinite(x * (448 / max))` -- bitwise torch's cast of
- `x * (448 / max)`. One warp a row, one launch; releases its dependents at its start."""
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(why="in-projection input, 37 rows", spec=dict(m=37, k=2560)),
- Cell(
- why="output-projection input, 256 rows, row stride K + 64",
- spec=dict(m=256, k=6144, pad=64),
- ),
- Cell(why="one row", spec=dict(m=1, k=2560)),
- )
-
- note = """
- Call site: `mid_rows_mm`'s FP8 mode. Preconditions: `x` bf16 `[M, K]`, unit column
- stride, row stride a multiple of 8 elements, 16-byte aligned, K a multiple of 256.
- """
-
- def __call__(self, x: torch.Tensor):
- return fp8_row_quant(x)
-
- def reference(self, x: torch.Tensor):
- """Both quotients correctly rounded (float64, rounded once), as the kernel's
- `div.rn.f32`."""
- amax = x.float().abs().amax(1).clamp_min(AMAX_FLOOR).double()
- inv = (E4M3_MAX / amax).float()
- return (x.float() * inv[:, None]).to(torch.float8_e4m3fn), (amax / E4M3_MAX).float()
-
-
-mid_rows_mm_op = MidRowsMm()
-fp8_row_quant_op = Fp8RowQuant()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/nvfp4_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/nvfp4_gemv.py
deleted file mode 100644
index bbe39917e7a4..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/nvfp4_gemv.py
+++ /dev/null
@@ -1,992 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Decode-row W4A16 projection: `y[M, N] = x[M, K] @ dequant(W)^T` with NVFP4 weights (e2m1
-codes, one e4m3 scale per 16 consecutive k of a row, one fp32 global scale) and bf16
-activations, M = 1 .. 16.
-
-CuTe DSL (inline PTX), one PDL launch, CUDA-graph capturable; compiled once per (N, K, plan)
-and row octet (1..8 or 9..16 rows): the row count is a launch argument.
-
-Math: swap-AB warp MMA (`mma.sync.m16n8k16`, f16 in, fp32 accumulate). A is 16 weight rows,
-each e2m1 pair widened by `cvt.rn.f16x2.e2m1x2` and multiplied by its block scale in f16
-(`cvt.rn.f16x2.e4m3x2`, `mul.rn.f16x2`): an e2m1 value times an e4m3 scale has at most five
-significant bits and lies in [2^-10, 2688], so the dequantized weight is exact. B is up to 8
-activation rows per n8 tile, rounded once from bf16 to f16 (exact for magnitudes in
-[2^-14, 65504]; smaller ones lose low bits, larger ones saturate). The MMA sums in fp32 over
-all of K; the sum times the global scale is rounded once to bf16.
-
-Weight layout (`pack_weight`): 1152-byte units of 16 rows x 128 k -- 1024 bytes of codes then
-128 bytes of scales -- stored k-block-major `[K/128][ceil(N/16)][1152]`. K is permuted
-identically for both operands inside every 128-block: lane quad t of k16 step s holds
-k = 32 t + 4 s + 0..3, so a thread's activations of a k-block are 32 consecutive values and
-its four k of a step share one 16-element scale block (2 t for steps 0..3, 2 t + 1 for 4..7).
-Codes: `[half h (2)][lane (32)][4 words]`, word i of half h the step 4 h + i, its bytes
-`(row g, k0 | k1 << 4), (row g + 8, same k), (row g, k2 | k3 << 4), (row g + 8, same k)` --
-one byte per A register. Scales: `[lane (32)][4 bytes]`: `(row g, block 2t), (row g + 8, 2t),
-(row g, 2t + 1), (row g + 8, 2t + 1)`.
-
-Work split (`Plan`): persistent, one or two resident CTAs per multiprocessor, each a contiguous
-range of 16-row tiles taken `rt` tiles (a chunk) at a time over all of K. A producer warp streams the units
-with `cp.async.bulk` into a ring of shared-memory slots (one k-block of the chunk's tiles a
-slot) completed on mbarriers, continuously across chunks and before the PDL wait (the weight
-is static). After `griddepcontrol.wait` the consumer warps convert the activations once to
-f16 into shared memory in MMA-fragment order; consumer warp w then takes tiles w, w + warps,
-... of every slot, its accumulators in registers for the chunk, and stores its tiles through a
-per-warp staging buffer.
-
-Split entries (`nvfp4_split_gemv*`, `SplitNvfp4Weight`): the same NVFP4 weight packed in
-`fp8_gemv`'s fragment order (`pack_split`) and streamed by that kernel, which splits K across
-warps and carries the GDN conv epilogues; for projections with few 16-row tiles a
-multiprocessor (the GDN in_proj), where one warp per tile over all of K is latency-bound.
-"""
-
-from __future__ import annotations
-
-import functools
-from dataclasses import dataclass
-
-import cutlass
-import cutlass.cute as cute
-import torch
-from cutlass import const_expr
-from cutlass.cute.runtime import from_dlpack, make_ptr
-from cutlass.cutlass_dsl import T, dsl_user_op
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import (
- Arch,
- Cell,
- OpWrapper,
- true_fp32_matmul,
-)
-
-from .fp8_gemv import (
- TILE_MAJOR_MAX_TILES,
- _asm,
- _bar_sync,
- _bf16x2_to_f16x2,
- _bulk_g2s,
- _bulk_prefetch_l2,
- _dq_step,
- _env_pdl,
- _ir,
- _ldg4p,
- _lds4,
- _mbar_arrive,
- _mbar_expect_tx,
- _mbar_init,
- _mbar_init_fence,
- _mbar_wait,
- _mma_f16,
- _scales4,
- _sms,
- fp8_gemv,
- fp8_gemv_conv1d,
- fp8_gemv_conv1d_verify,
- fp8_gemv_conv1d_verify_op,
- verify_reference,
-)
-from .fp8_gemv import _compiled as _split_compiled
-from .fp8_gemv import default_plan as _split_plan
-
-TILE_N = 16
-BLOCK_K = 128
-SF_VEC = 16 # k per e4m3 block scale
-DATA = TILE_N * BLOCK_K // 2 # bytes of e2m1 codes of one (16-row, 128-k) unit
-SCALES = TILE_N * BLOCK_K // SF_VEC # bytes of its e4m3 block scales
-UNIT = DATA + SCALES
-E2M1_MAX = 6.0
-E4M3_MAX = 448.0
-MAX_ROWS = 16
-#: The e2m1 magnitudes, by the three low bits of a code (bit 3 is the sign).
-E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
-#: Rounding thresholds between consecutive e2m1 magnitudes (a tie rounds toward zero).
-_E2M1_BOUNDS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0)
-#: Shared memory one CTA may take (the SM100 opt-in maximum less a margin) and a
-#: multiprocessor holds (less a margin per resident CTA).
-SMEM_MAX = 225 << 10
-SMEM_SM = 228 << 10
-#: Activation loads a thread keeps in flight in the prologue (16 bytes each).
-_X_BATCH = 10
-
-
-# ------------------------------------------------------------------------------------------
-# inline PTX helpers
-# ------------------------------------------------------------------------------------------
-
-
-@dsl_user_op
-def _lds32(saddr, *, loc=None, ip=None):
- return cutlass.Int32(
- _asm(
- T.i32(),
- [_ir(saddr, cutlass.Int32, loc, ip)],
- "ld.shared.b32 $0, [$1];",
- "=r,r",
- loc,
- ip,
- side=True,
- )
- )
-
-
-@dsl_user_op
-def _sts4(saddr, v0, v1, v2, v3, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [_ir(saddr, cutlass.Int32, loc, ip)]
- + [_ir(v, cutlass.Int32, loc, ip) for v in (v0, v1, v2, v3)],
- "{ st.shared.v4.b32 [$1], {$2, $3, $4, $5}; mov.u32 $0, 0; }",
- "=r,r,r,r,r,r",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _stg4(gaddr, v0, v1, v2, v3, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [_ir(gaddr, cutlass.Int64, loc, ip)]
- + [_ir(v, cutlass.Int32, loc, ip) for v in (v0, v1, v2, v3)],
- "{ st.global.v4.b32 [$1], {$2, $3, $4, $5}; mov.u32 $0, 0; }",
- "=r,l,r,r,r,r",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _lds_u16(saddr, *, loc=None, ip=None):
- return cutlass.Int32(
- _asm(
- T.i32(),
- [_ir(saddr, cutlass.Int32, loc, ip)],
- "ld.shared.u16 $0, [$1];",
- "=r,r",
- loc,
- ip,
- side=True,
- )
- )
-
-
-@dsl_user_op
-def _stg_b16(gaddr, v, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [_ir(gaddr, cutlass.Int64, loc, ip), _ir(v, cutlass.Int32, loc, ip)],
- "{ .reg .b16 h; cvt.u16.u32 h, $2; st.global.b16 [$1], h; mov.u32 $0, 0; }",
- "=r,l,r",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _sts_bf16_scaled(saddr, v, s, *, loc=None, ip=None):
- """bf16(v * s) to shared memory."""
- _asm(
- T.i32(),
- [
- _ir(saddr, cutlass.Int32, loc, ip),
- _ir(v, cutlass.Float32, loc, ip),
- _ir(s, cutlass.Float32, loc, ip),
- ],
- "{ .reg .f32 f; .reg .b16 h; mul.rn.f32 f, $2, $3; cvt.rn.bf16.f32 h, f; "
- "st.shared.b16 [$1], h; mov.u32 $0, 0; }",
- "=r,r,f,f",
- loc,
- ip,
- side=True,
- )
-
-
-# ------------------------------------------------------------------------------------------
-# kernel
-# ------------------------------------------------------------------------------------------
-
-
-@dataclass(frozen=True)
-class Plan:
- """Launch shape: `grid` persistent CTAs (all resident: one or two per multiprocessor), each
- taking its tiles `rt` at a time; `warps` consumer warps; `slots` ring slots (one k-block of
- a chunk each); with `l2_ahead` > 0, issuing stage s also prefetches stage
- s + slots + l2_ahead - 1 into L2."""
-
- grid: int
- rt: int
- warps: int
- slots: int
- l2_ahead: int = 0
-
-
-class _Nvfp4Gemv:
- def __init__(self, *, n, k, nt8, plan: Plan, pdl):
- assert k % BLOCK_K == 0 and nt8 in (1, 2)
- self.n, self.k, self.nt8, self.pdl = n, k, nt8, pdl
- self.n_tiles = -(-n // TILE_N)
- self.nkb = k // BLOCK_K
- self.grid = plan.grid
- self.rt, self.cw, self.slots, self.l2_ahead = plan.rt, plan.warps, plan.slots, plan.l2_ahead
- self.rtw = -(-self.rt // self.cw)
- self.threads = (self.cw + 1) * 32 # + the weight producer warp
- self.np = nt8 * 8
- self.slot_bytes = -(-(self.rt * UNIT) // 128) * 128
- self.o_ring = -(-(self.nkb * nt8 * 4 * 512) // 128) * 128 # activation fragments first
- self.o_stg = self.o_ring + self.slots * self.slot_bytes
- self.o_bar = self.o_stg + self.cw * self.np * TILE_N * 2
- self.smem = self.o_bar + 2 * self.slots * 8 + 128
- assert self.smem <= SMEM_MAX, (self.smem, plan)
-
- @cute.jit
- def __call__(
- self,
- out_t: cute.Tensor,
- x_ptr: cute.Pointer,
- w_ptr: cute.Pointer,
- x_stride: cutlass.Int32,
- m: cutlass.Int32,
- gscale: cutlass.Float32,
- stream,
- ):
- self.kernel(
- out_t.iterator.toint(), x_ptr.toint(), w_ptr.toint(), x_stride, m, gscale
- ).launch(
- grid=[self.grid, 1, 1],
- block=[self.threads, 1, 1],
- smem=self.smem,
- stream=stream,
- use_pdl=self.pdl,
- )
-
- @cute.jit
- def _compute(self, acc, xf, slot, wi, ntc, lane):
- """This warp's tiles of the k-block in ring slot `slot` times the activation fragments
- `xf` ([nt8][16] f16x2), accumulated into `acc`."""
- cw: cutlass.Constexpr = self.cw
- nt8: cutlass.Constexpr = self.nt8
- for i in cutlass.range_constexpr(self.rtw):
- tl = wi + i * cw
- if tl < ntc:
- ua = slot + tl * UNIT
- sc = _scales4(_lds32(ua + DATA + lane * 4))
- wv = [_lds4(ua + h * 512 + lane * 16) for h in range(2)]
- for s in cutlass.range_constexpr(8):
- h = s // 4
- a = _dq_step(wv[h][s % 4], sc[2 * h], sc[2 * h + 1])
- for j in cutlass.range_constexpr(nt8):
- r = _mma_f16(
- a,
- [xf[j][2 * s], xf[j][2 * s + 1]],
- [acc[i, j, 0], acc[i, j, 1], acc[i, j, 2], acc[i, j, 3]],
- )
- for e in cutlass.range_constexpr(4):
- acc[i, j, e] = r[e]
-
- @cute.kernel
- def kernel(
- self,
- out_base: cutlass.Int64,
- x_base: cutlass.Int64,
- w_base: cutlass.Int64,
- x_stride: cutlass.Int32,
- m: cutlass.Int32,
- gscale: cutlass.Float32,
- ):
- tidx, _, _ = cute.arch.thread_idx()
- bidx, _, _ = cute.arch.block_idx()
- warp = cute.arch.make_warp_uniform(cute.arch.warp_idx())
- lane = cute.arch.lane_idx()
- g = lane // 4
- t = lane % 4
- n_tiles: cutlass.Constexpr = self.n_tiles
- nkb: cutlass.Constexpr = self.nkb
- rt: cutlass.Constexpr = self.rt
- cw: cutlass.Constexpr = self.cw
- slots: cutlass.Constexpr = self.slots
- slot_bytes: cutlass.Constexpr = self.slot_bytes
- nt8: cutlass.Constexpr = self.nt8
- np_: cutlass.Constexpr = self.np
- rtw: cutlass.Constexpr = self.rtw
- n_out: cutlass.Constexpr = self.n
- l2a: cutlass.Constexpr = self.l2_ahead
- grid: cutlass.Constexpr = self.grid
-
- smem = cutlass.utils.SmemAllocator()
- base_t = smem.allocate_tensor(
- cutlass.Int8, cute.make_layout((self.smem - 128,), stride=(1,)), byte_alignment=128
- )
- sbase = cutlass.Int32(base_t.iterator.toint())
- xs = sbase # activation fragments [kb][j][c][lane][16 B], f16
- ring = sbase + self.o_ring
- bars = sbase + self.o_bar # full[slots], empty[slots]
-
- tile0 = (bidx * n_tiles) // grid
- ntot = ((bidx + 1) * n_tiles) // grid - tile0
- nch = (ntot + rt - 1) // rt
-
- if warp == cw:
- if lane == 0:
- for sl in cutlass.range_constexpr(slots):
- _mbar_init(bars + 8 * sl, 1)
- _mbar_init(bars + 8 * (slots + sl), cw)
- _mbar_init_fence()
- cute.arch.sync_threads()
- if const_expr(self.pdl):
- cute.arch.griddepcontrol_launch_dependents()
-
- if warp == cw:
- # ---- weight producer (static data: no PDL wait); stage = (chunk, k-block) ----
- if lane == 0:
- for st in cutlass.range(nch * nkb):
- sl = st % slots
- if st >= slots:
- _mbar_wait(bars + 8 * (slots + sl), ((st // slots) - 1) & 1)
- c = st // nkb
- kb = st - c * nkb
- tc0 = tile0 + c * rt
- ntc = cutlass.min(rt, ntot - c * rt)
- if const_expr(l2a > 0):
- pst = st + slots + l2a - 1
- if pst < nch * nkb:
- pc = pst // nkb
- _bulk_prefetch_l2(
- w_base
- + (cutlass.Int64(pst - pc * nkb) * n_tiles + tile0 + pc * rt)
- * UNIT,
- cutlass.min(rt, ntot - pc * rt) * UNIT,
- )
- _mbar_expect_tx(bars + 8 * sl, ntc * UNIT)
- _bulk_g2s(
- ring + sl * slot_bytes,
- w_base + (cutlass.Int64(kb) * n_tiles + tc0) * UNIT,
- ntc * UNIT,
- bars + 8 * sl,
- )
- else:
- # ---- consumers ----
- if const_expr(self.pdl):
- cute.arch.griddepcontrol_wait()
- # activations -> f16 fragments: item q = ((kb nt8 + j) 4 + c) 32 + lane holds row
- # 8 j + g, k = 128 kb + 32 t + 8 c .. + 7 (rows past m zero)
- n_items: cutlass.Constexpr = nkb * nt8 * 4 * 32
- n_it: cutlass.Constexpr = -(-n_items // (cw * 32))
- # loads issued a batch at a time ahead of their conversions (each load asm is
- # ordered, so a load placed after a conversion would wait out its predecessor)
- for b0 in cutlass.range_constexpr(0, n_it, _X_BATCH):
- vals = []
- for it in cutlass.range_constexpr(b0, min(b0 + _X_BATCH, n_it)):
- q = tidx + it * (cw * 32)
- ql = q % 32
- row = 8 * ((q // 128) % nt8) + ql // 4
- kx = (q // (128 * nt8)) * BLOCK_K + 32 * (ql % 4) + 8 * ((q // 32) % 4)
- vals.append(
- _ldg4p(
- x_base
- + cutlass.Int64(row) * cutlass.Int64(x_stride) * 2
- + cutlass.Int64(kx * 2),
- cutlass.Int32((row < m) & (q < n_items)),
- )
- )
- for it in cutlass.range_constexpr(b0, min(b0 + _X_BATCH, n_it)):
- q = tidx + it * (cw * 32)
- if q < n_items:
- v = vals[it - b0]
- _sts4(
- xs + q * 16,
- _bf16x2_to_f16x2(v[0]),
- _bf16x2_to_f16x2(v[1]),
- _bf16x2_to_f16x2(v[2]),
- _bf16x2_to_f16x2(v[3]),
- )
- _bar_sync(1, cw * 32)
-
- stg = sbase + self.o_stg + warp * (np_ * TILE_N * 2) # [np][16] bf16 per warp
- acc = cute.make_rmem_tensor(
- cute.make_layout((rtw, nt8, 4), stride=(nt8 * 4, 4, 1)), cutlass.Float32
- )
- for c in cutlass.range(nch):
- ntc = cutlass.min(rt, ntot - c * rt)
- acc.fill(0.0)
- for kb in cutlass.range_constexpr(nkb):
- st = c * nkb + kb
- sl = st % slots
- xf = [
- [
- v
- for cc in range(4)
- for v in _lds4(xs + (((kb * nt8 + j) * 4 + cc) * 32 + lane) * 16)
- ]
- for j in range(nt8)
- ]
- _mbar_wait(bars + 8 * sl, (st // slots) & 1)
- self._compute(acc, xf, ring + sl * slot_bytes, warp, ntc, lane)
- cute.arch.sync_warp()
- if lane == 0:
- _mbar_arrive(bars + 8 * (slots + sl))
- # ---- store: through the warp's staging buffer, a row's 16 outputs per 2 lanes ----
- for i in cutlass.range_constexpr(rtw):
- tl = warp + i * cw
- if tl < ntc:
- for j in cutlass.range_constexpr(nt8):
- n0 = 8 * j + 2 * t
- p = stg + (n0 * TILE_N + g) * 2
- _sts_bf16_scaled(p, acc[i, j, 0], gscale)
- _sts_bf16_scaled(p + TILE_N * 2, acc[i, j, 1], gscale)
- _sts_bf16_scaled(p + 8 * 2, acc[i, j, 2], gscale)
- _sts_bf16_scaled(p + TILE_N * 2 + 8 * 2, acc[i, j, 3], gscale)
- cute.arch.sync_warp()
- orow0 = (tile0 + c * rt + tl) * TILE_N
- for rr in cutlass.range_constexpr(-(-(np_ * 2) // 32)):
- li = lane + rr * 32
- en = li // 2
- hh = li % 2
- if (li < np_ * 2) & (en < m):
- if const_expr(n_out % 8 == 0):
- if orow0 + 8 * hh < n_out:
- v = _lds4(stg + (en * TILE_N + 8 * hh) * 2)
- _stg4(
- out_base
- + (cutlass.Int64(en) * n_out + orow0 + 8 * hh) * 2,
- v[0],
- v[1],
- v[2],
- v[3],
- )
- else:
- for e in cutlass.range_constexpr(8):
- if orow0 + 8 * hh + e < n_out:
- _stg_b16(
- out_base
- + (cutlass.Int64(en) * n_out + orow0 + 8 * hh + e)
- * 2,
- _lds_u16(stg + (en * TILE_N + 8 * hh + e) * 2),
- )
- cute.arch.sync_warp()
-
-
-# ------------------------------------------------------------------------------------------
-# host side
-# ------------------------------------------------------------------------------------------
-
-
-@dataclass
-class PackedNvfp4Weight:
- """`data`: uint8 `[K/128, ceil(N/16), 1152]` (see `pack_weight`); `global_scale`: the fp32
- dequantization multiplier (`w = e2m1 code x e4m3 block scale x global_scale`); logical
- shape `[n, k]`."""
-
- data: torch.Tensor
- global_scale: float
- n: int
- k: int
-
-
-def _global_scale(w: torch.Tensor) -> float:
- """amax(|w|) / (448 x 6): the largest block's scale lands on e4m3's largest value."""
- return max(w.abs().amax().float().item(), 1e-30) / (E4M3_MAX * E2M1_MAX)
-
-
-def _quantize_rows(w: torch.Tensor, gscale: float):
- """(codes uint8 `[n, K]`, e4m3 block scales `[n, K/16]`) of `w` under `gscale`."""
- n, k = w.shape
- blk = w.float().reshape(n, k // SF_VEC, SF_VEC)
- sf = (blk.abs().amax(2) / (E2M1_MAX * gscale)).clamp(max=E4M3_MAX).to(torch.float8_e4m3fn)
- step = sf.float()[..., None] * gscale
- v = torch.where(step > 0, blk / step.clamp_min(1e-38), 0.0)
- bounds = torch.tensor(_E2M1_BOUNDS, device=w.device)
- idx = torch.bucketize(v.abs(), bounds).to(torch.uint8)
- codes = idx | (((v < 0) & (idx > 0)).to(torch.uint8) << 3)
- return codes.reshape(n, k), sf
-
-
-def quantize_nvfp4(w: torch.Tensor, rows_per_chunk: int = 16384):
- """bf16/fp32 `[N, K]` -> (codes uint8 `[N, K]` -- bit 3 the sign, bits 0..2 the index into
- `E2M1_VALUES` --, e4m3 block scales `[N, K/16]`, global scale): global = amax(|w|) /
- (448 x 6), a block's scale = e4m3(amax(block) / 6 / global), a code the nearest e2m1 value
- of w / (scale x global) (ties toward zero), saturated at 6."""
- n, k = w.shape
- assert k % SF_VEC == 0
- gscale = _global_scale(w)
- codes = torch.empty(n, k, dtype=torch.uint8, device=w.device)
- scales = torch.empty(n, k // SF_VEC, dtype=torch.float8_e4m3fn, device=w.device)
- for lo in range(0, n, rows_per_chunk):
- hi = min(lo + rows_per_chunk, n)
- codes[lo:hi], scales[lo:hi] = _quantize_rows(w[lo:hi], gscale)
- return codes, scales, gscale
-
-
-def dequantize_nvfp4(codes: torch.Tensor, scales: torch.Tensor, gscale: float) -> torch.Tensor:
- """fp32 `[N, K]` = e2m1(code) x its block scale x the global scale."""
- lut = torch.tensor(E2M1_VALUES, dtype=torch.float32, device=codes.device)
- val = lut[(codes & 7).long()] * (1.0 - 2.0 * (codes >> 3).float())
- return val * scales.float().repeat_interleave(SF_VEC, 1) * gscale
-
-
-def _pack_into(out: torch.Tensor, tile0: int, codes: torch.Tensor, scales: torch.Tensor) -> None:
- """Rows `16 tile0 ..` (codes `[n, K]`, scales `[n, K/16]`) into the packed buffer `out`
- `[K/128, tiles, 1152]`; rows past n up to the tile boundary are zero."""
- n, k = codes.shape
- nkb = k // BLOCK_K
- tiles = -(-n // TILE_N)
- c = torch.zeros(tiles * TILE_N, k, dtype=torch.uint8, device=codes.device)
- c[:n] = codes
- # (tile, h', g, kb, t, hh, s4, e2, e1): row 8 h' + g, k = 128 kb + 32 t + 16 hh + 4 s4 + 2 e2 + e1
- c = c.reshape(tiles, 2, 8, nkb, 4, 2, 4, 2, 2)
- byte = c[..., 0] | (c[..., 1] << 4)
- # unit byte order [hh][g][t][s4][e2][h']
- out[:, tile0 : tile0 + tiles, :DATA] = byte.permute(3, 0, 5, 2, 4, 6, 7, 1).reshape(
- nkb, tiles, DATA
- )
- s = torch.zeros(tiles * TILE_N, k // SF_VEC, dtype=torch.uint8, device=codes.device)
- s[:n] = scales.view(torch.uint8)
- # (tile, h', g, kb, t, hb): block 8 kb + 2 t + hb; unit order [g][t][hb][h']
- out[:, tile0 : tile0 + tiles, DATA:] = (
- s.reshape(tiles, 2, 8, nkb, 4, 2).permute(3, 0, 2, 4, 5, 1).reshape(nkb, tiles, SCALES)
- )
-
-
-def pack_weight(codes: torch.Tensor, scales: torch.Tensor, gscale: float) -> PackedNvfp4Weight:
- """(codes `[N, K]`, e4m3 scales `[N, K/16]`, global scale) -> the kernel's layout (module
- docstring)."""
- n, k = codes.shape
- assert k % BLOCK_K == 0 and scales.shape == (n, k // SF_VEC)
- out = torch.empty(k // BLOCK_K, -(-n // TILE_N), UNIT, dtype=torch.uint8, device=codes.device)
- _pack_into(out, 0, codes, scales)
- return PackedNvfp4Weight(out, float(gscale), n, k)
-
-
-def pack_from_bf16(w: torch.Tensor, rows_per_chunk: int = 16384) -> PackedNvfp4Weight:
- """Converter from a bf16 `[N, K]` weight, quantized and packed in row chunks (the transient
- memory stays a few hundred MB for the 248320 x 2560 vocabulary)."""
- n, k = w.shape
- assert k % BLOCK_K == 0 and rows_per_chunk % TILE_N == 0
- gscale = _global_scale(w)
- out = torch.empty(k // BLOCK_K, -(-n // TILE_N), UNIT, dtype=torch.uint8, device=w.device)
- for lo in range(0, n, rows_per_chunk):
- hi = min(lo + rows_per_chunk, n)
- _pack_into(out, lo // TILE_N, *_quantize_rows(w[lo:hi], gscale))
- return PackedNvfp4Weight(out, float(gscale), n, k)
-
-
-@dataclass
-class SplitNvfp4Weight:
- """NVFP4 units in `fp8_gemv`'s fragment order, which that kernel (k split across warps, the
- GDN conv epilogues) streams in place of FP8 ones: `data` uint8 `[K/128, ceil(N/16), 1152]`
- (k-block-major) or `[ceil(N/16), K/128, 1152]` (`tile_major`), see `pack_split`; `scale` the
- fp32 global scale as a one-element device tensor (the kernel reads it), also
- `global_scale`; logical shape `[n, k]`."""
-
- data: torch.Tensor
- scale: torch.Tensor
- global_scale: float
- n: int
- k: int
- tile_major: bool = False
- w4: bool = True
-
-
-def pack_split(
- codes: torch.Tensor, scales: torch.Tensor, gscale: float, tile_major: bool | None = None
-) -> SplitNvfp4Weight:
- """(codes `[N, K]`, e4m3 scales `[N, K/16]`, global scale) -> `fp8_gemv`'s NVFP4 units.
-
- A unit (k-block kb, 16-row tile T) is 1152 bytes. Codes, 1024 bytes `[h (2)][lane (32)][4
- words]`, lane = 4 g + t: word 2 cc + st of half h is k16 step st of chunk c = 2 h + cc,
- covering k = 128 kb + 32 c + 8 t + 4 st + 0..3 (`fp8_gemv`'s k order), its bytes `(row g,
- k0 | k1 << 4), (row g + 8, same k), (row g, k2 | k3 << 4), (row g + 8, same k)`. Scales,
- 128 bytes `[g (8)][t // 2 (2)][c (4)][row g, row g + 8]`: the block 2 c + t // 2 of the
- k-block. Rows past N are zero. Default `tile_major` as `fp8_gemv.pack_weight`'s."""
- n, k = codes.shape
- assert k % BLOCK_K == 0 and scales.shape == (n, k // SF_VEC)
- n_tiles = -(-n // TILE_N)
- nkb = k // BLOCK_K
- if tile_major is None:
- tile_major = n_tiles <= TILE_MAJOR_MAX_TILES
- c = torch.zeros(n_tiles * TILE_N, k, dtype=torch.uint8, device=codes.device)
- c[:n] = codes
- # (tile, rh, g, kb, h, cc, t, st, e2, e1): row 8 rh + g, k = 128 kb + 32 (2 h + cc) + 8 t
- # + 4 st + 2 e2 + e1
- c = c.reshape(n_tiles, 2, 8, nkb, 2, 2, 4, 2, 2, 2)
- byte = c[..., 0] | (c[..., 1] << 4)
- # unit byte order [h][g][t][cc][st][e2][rh]
- perm = (4, 2, 6, 5, 7, 8, 1)
- body = (0, 3) if tile_major else (3, 0)
- data = torch.empty(*(byte.shape[i] for i in body), UNIT, dtype=torch.uint8, device=codes.device)
- data[..., :DATA] = byte.permute(*body, *perm).reshape(*data.shape[:2], DATA)
- s = torch.zeros(n_tiles * TILE_N, k // SF_VEC, dtype=torch.uint8, device=codes.device)
- s[:n] = scales.view(torch.uint8)
- # (tile, rh, g, kb, c, th): block 8 kb + 2 c + th; unit order [g][th][c][rh]
- s = s.reshape(n_tiles, 2, 8, nkb, 4, 2)
- data[..., DATA:] = s.permute(*body, 2, 5, 4, 1).reshape(*data.shape[:2], SCALES)
- scale = torch.tensor([gscale], dtype=torch.float32, device=codes.device)
- return SplitNvfp4Weight(data, scale, float(gscale), n, k, tile_major)
-
-
-def pack_split_from_bf16(w: torch.Tensor, rows_per_chunk: int = 16384) -> SplitNvfp4Weight:
- """Converter from a bf16 `[N, K]` weight (`quantize_nvfp4`, then `pack_split`)."""
- return pack_split(*quantize_nvfp4(w, rows_per_chunk))
-
-
-#: Measured plans (GB300, 152 SMs, CUDA graph of back-to-back calls over weights rotated past
-#: L2), keyed (N, K, row octets): `ctas` resident CTAs per multiprocessor (two small ones let
-#: a launch's ramp overlap its neighbour's tail), tiles a chunk, consumer warps, ring slots,
-#: L2 look-ahead. The 16-row variant's activations (80 KB in shared memory) leave room for
-#: one CTA.
-_TABLE: dict[tuple[int, int, int], dict] = {
- (248320, 2560, 1): dict(ctas=2, rt=8, warps=8, slots=5, l2_ahead=1),
- (248320, 2560, 2): dict(ctas=1, rt=32, warps=16, slots=3, l2_ahead=1),
-}
-_DEFAULT = {
- 1: dict(ctas=2, rt=8, warps=8, slots=6, l2_ahead=1),
- 2: dict(ctas=1, rt=16, warps=16, slots=6, l2_ahead=1),
-}
-
-
-def default_plan(n: int, k: int, nt8: int, sms: int) -> Plan:
- """The plan `nvfp4_gemv` launches with: the measured table entry for (N, K, nt8), else the
- octet's default, its ring trimmed to the shared memory its CTAs per multiprocessor share."""
- n_tiles = -(-n // TILE_N)
- f = dict(_TABLE.get((n, k, nt8), _DEFAULT[nt8]))
- ctas = f.pop("ctas")
- f["rt"] = min(f["rt"], n_tiles)
- slot = -(-(f["rt"] * UNIT) // 128) * 128
- fixed = -(-((k // BLOCK_K) * nt8 * 2048) // 128) * 128 + f["warps"] * nt8 * 256 + 128
- budget = min(SMEM_MAX, SMEM_SM // ctas - (1 << 10)) - fixed - 2 * 8 * f["slots"]
- f["slots"] = max(2, min(f["slots"], budget // slot))
- return Plan(grid=min(ctas * sms, n_tiles), **f)
-
-
-@functools.lru_cache(maxsize=None)
-def _compiled(n, k, nt8, plan: Plan, pdl: bool):
- kern = _Nvfp4Gemv(n=n, k=k, nt8=nt8, plan=plan, pdl=pdl)
- buf = torch.empty(16, dtype=torch.int32, device="cuda")
- out = torch.empty(nt8 * 8, n, dtype=torch.bfloat16, device="cuda")
- return cute.compile(
- kern,
- from_dlpack(out, assumed_align=16).mark_layout_dynamic(leading_dim=1),
- make_ptr(cutlass.BFloat16, buf.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
- make_ptr(cutlass.Uint8, buf.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
- cutlass.Int32(k),
- cutlass.Int32(1),
- cutlass.Float32(1.0),
- cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
- options="--opt-level 2 --enable-tvm-ffi",
- )
-
-
-def _octet(m: int) -> int:
- return 1 if m <= 8 else 2
-
-
-def precompile(w: PackedNvfp4Weight, pdl: bool = True) -> None:
- """Compile both row octets of `w`'s shape now: a first call mid-serving would otherwise
- stall on the compile."""
- for nt8 in (1, 2):
- _compiled(w.n, w.k, nt8, default_plan(w.n, w.k, nt8, _sms()), bool(pdl and _env_pdl()))
-
-
-def nvfp4_gemv(
- x: torch.Tensor, w: PackedNvfp4Weight, pdl: bool = True, plan: Plan | None = None
-) -> torch.Tensor:
- """`x[M, K]` bf16 (row stride a multiple of 8, 16-byte aligned base) @ dequant(w)^T ->
- fresh `[M, N]` bf16. M = 1 .. MAX_ROWS."""
- m, k = x.shape
- assert k == w.k and x.dtype == torch.bfloat16 and x.stride(1) == 1
- assert 1 <= m <= MAX_ROWS
- assert x.stride(0) % 8 == 0 and x.data_ptr() % 16 == 0
- out = torch.empty((m, w.n), dtype=torch.bfloat16, device=x.device)
- nt8 = _octet(m)
- if plan is None:
- plan = default_plan(w.n, k, nt8, _sms())
- fn = _compiled(w.n, k, nt8, plan, bool(pdl and _env_pdl()))
- fn(out, x.data_ptr(), w.data.data_ptr(), x.stride(0), m, w.global_scale)
- return out
-
-
-class Nvfp4Gemv(OpWrapper):
- """`nvfp4_gemv(x, w)`: `y[M, N] = x[M, K] @ dequant(w)^T` in bf16 for 1..MAX_ROWS rows, the
- weight a `PackedNvfp4Weight` (e2m1 codes, one e4m3 scale per 16 k of a row, an fp32 global
- scale; `pack_from_bf16` / `pack_weight`). The dequantized weight is exact in f16; the
- activation is rounded once to f16 on chip (exact for |x| in [2^-14, 65504]); fp32
- accumulation over all of K, the global scale applied in fp32, one bf16 rounding.
-
- Fusion boundary: the projection."""
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec keys: n, k, m (rows). Every cell: x a column slice of a wider bf16 row (row stride
- # past K); the weight `quantize_nvfp4` of a random bf16 matrix with an outlier column and
- # a block of near-zero values.
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="MTP draft head, one request's draft step (1 row)", spec=dict(n=248320, k=2560, m=1)
- ),
- Cell(why="draft head, 3 rows: a partial row octet", spec=dict(n=248320, k=2560, m=3)),
- Cell(why="draft head, 8 rows: a full row octet", spec=dict(n=248320, k=2560, m=8)),
- Cell(
- why="draft head, 9 rows: the two-octet variant with one row in the second",
- spec=dict(n=248320, k=2560, m=9),
- ),
- Cell(why="draft head, 16 rows: the largest dispatch", spec=dict(n=248320, k=2560, m=16)),
- Cell(
- why="N not a multiple of 8 (scalar stores, a partial last tile), 3 k-blocks, "
- "fewer tiles than multiprocessors",
- spec=dict(n=1001, k=384, m=5),
- ),
- )
-
- note = """
- Metadata consumed: none. Call site: the MTP draft head's vocabulary projection at up to
- `_FP8_GEMV_HEAD_ROWS` rows under `TRTLLM_MODELING_V2_MTP_HEAD_NVFP4=1` (the verify head
- is untouched). Preconditions: x bf16 `[M, K]`, unit column stride, row stride a multiple
- of 8 elements and a 16-byte aligned base; K a multiple of 128. One PDL launch (the weight
- streamed before the wait); CUDA-graph capturable; compiled per (N, K, plan) for rows
- 1..8 and 9..16 (`precompile` builds both), the row count a launch argument.
- """
-
- def __call__(self, x: torch.Tensor, w: PackedNvfp4Weight) -> torch.Tensor:
- return nvfp4_gemv(x, w)
-
- def reference(
- self, x: torch.Tensor, codes: torch.Tensor, scales: torch.Tensor, global_scale: float
- ) -> torch.Tensor:
- """fp32 `[M, N]` over the dequantized weight (a true fp32 matmul); takes the unpacked
- NVFP4 weight (`quantize_nvfp4`'s outputs) in place of the packed `w`."""
- with true_fp32_matmul():
- return x.float() @ dequantize_nvfp4(codes, scales, global_scale).t()
-
- def is_valid(self, x: torch.Tensor, w: PackedNvfp4Weight, **kwargs) -> None:
- assert x.float().abs().amax().item() <= 65504.0, (
- "an activation beyond f16's range saturates at 65504 in the kernel"
- )
-
-
-def precompile_split(
- w: SplitNvfp4Weight,
- max_rows: int,
- conv_dim: int,
- pad_slot_id: int,
- verify_p: tuple[int, ...] = (),
- pdl: bool = True,
-) -> None:
- """Compile now every `fp8_gemv` variant over `w`'s shape a decode step can reach: the
- projection and the single-token conv entry at 1..max_rows rows, and the verify entry at
- each p of `verify_p` for every whole number of requests within max_rows. The kernel is
- compiled per row count; a first call mid-serving would stall on the compile."""
- pdl = bool(pdl and _env_pdl())
- sms = _sms()
- for m in range(1, max_rows + 1):
- plan = _split_plan(w.n, w.k, m, sms, True)
- _split_compiled(w.n, w.k, m, plan, pdl, w.tile_major, w4=True)
- _split_compiled(w.n, w.k, m, plan, pdl, w.tile_major, conv_dim, pad_slot_id, w4=True)
- for p in verify_p:
- if m % p == 0 and m >= 2:
- _split_compiled(
- w.n, w.k, m, plan, pdl, w.tile_major, conv_dim, pad_slot_id, p, w4=True
- )
-
-
-def _split_cell(why: str, n: int, k: int, m: int) -> Cell:
- return Cell(why=why, spec=dict(n=n, k=k, m=m))
-
-
-class Nvfp4SplitGemv(OpWrapper):
- """`nvfp4_split_gemv(x, w)` (`fp8_gemv` over a `SplitNvfp4Weight`): `y[M, N] = x[M, K] @
- dequant(w)^T` in bf16 for 1..32 rows, the weight NVFP4 (e2m1 codes, one e4m3 scale per 16 k
- of a row, an fp32 global scale; `pack_split_from_bf16` / `pack_split`). The dequantized
- weight is exact in f16; the activation is rounded once to f16 on chip (exact for |x| in
- [2^-14, 65504]); fp32 accumulation, the global scale applied in fp32, one bf16 rounding.
-
- Fusion boundary: the projection; `Nvfp4SplitGemvConv1d` / `Nvfp4SplitGemvConv1dVerify`
- also fuse the GDN conv step."""
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec keys: n, k, m (rows). Every cell: x a column slice of a wider bf16 row; the weight
- # `quantize_nvfp4` of a random bf16 matrix with an outlier column and near-zero blocks.
- CELLS: tuple[Cell, ...] = (
- _split_cell("GDN in_proj, one decode row", 16480, 2560, 1),
- _split_cell("GDN in_proj, 3 rows", 16480, 2560, 3),
- _split_cell("GDN in_proj, 9 rows (two n8 tiles)", 16480, 2560, 9),
- _split_cell("GDN in_proj, 16 rows", 16480, 2560, 16),
- _split_cell("small N, tile-major units, 3 k-blocks", 1001, 384, 5),
- )
-
- note = """
- Metadata consumed: none. Call site: the decode target's GDN in_proj at up to
- `_INPROJ_NVFP4_MAX_ROWS` rows of an eager step under `TRTLLM_MODELING_V2_GDN_INPROJ_NVFP4=1`.
- Preconditions: x bf16 `[M, K]`, unit column stride, row stride a multiple of 8 elements and
- a 16-byte aligned base; K a multiple of 128. One PDL launch (the weight streamed before the
- wait); CUDA-graph capturable; compiled per (N, K, M, plan) (`precompile_split`).
- """
-
- def __call__(self, x: torch.Tensor, w: SplitNvfp4Weight) -> torch.Tensor:
- return fp8_gemv(x, w)
-
- def reference(
- self, x: torch.Tensor, codes: torch.Tensor, scales: torch.Tensor, global_scale: float
- ) -> torch.Tensor:
- """fp32 `[M, N]` over the dequantized weight (a true fp32 matmul); takes the unpacked
- NVFP4 weight (`quantize_nvfp4`'s outputs) in place of the packed `w`."""
- with true_fp32_matmul():
- return x.float() @ dequantize_nvfp4(codes, scales, global_scale).t()
-
- def is_valid(self, x: torch.Tensor, w: SplitNvfp4Weight, **kwargs) -> None:
- assert x.float().abs().amax().item() <= 65504.0, (
- "an activation beyond f16's range saturates at 65504 in the kernel"
- )
-
-
-class Nvfp4SplitGemvConv1d(OpWrapper):
- """`nvfp4_split_gemv_conv1d` (`fp8_gemv_conv1d` over a `SplitNvfp4Weight`): the GDN input
- projection of decode rows over the NVFP4 copy, then for the conv columns of a live row the
- causal conv step against its paged state (SiLU; the state advanced in place). Returns the
- `[M, N]` projection with the conv columns replaced by the conv output."""
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec keys: slots (one per row; -1 the pad sentinel). Every cell: the GDN in_proj (16480 x
- # 2560, conv_dim 10240), a strided 64-slot pool, width 4, SiLU.
- CELLS: tuple[Cell, ...] = (
- Cell(why="STP decode: one row, live slot", spec=dict(slots=(37,))),
- Cell(why="two rows, one padding slot", spec=dict(slots=(17, -1))),
- Cell(why="four rows", spec=dict(slots=(0, 5, 63, 9))),
- Cell(why="sixteen rows (two n8 tiles)", spec=dict(slots=tuple(range(2, 50, 3)))),
- )
-
- note = """
- Call site: the decode target's GDN in_proj + conv step of a captured single-token step at
- up to `_INPROJ_NVFP4_MAX_ROWS` rows under `TRTLLM_MODELING_V2_GDN_INPROJ_NVFP4=1`.
- `conv_state` `[slots, conv_dim, 3]` bf16 (any strides), `conv_weight` `[conv_dim, 4]` bf16
- contiguous, int32 slots; a `pad_slot_id` row gets the plain projection. Slots, taps and the
- window are read before the PDL wait (nothing in a captured step writes this layer's pool
- ahead of it).
- """
-
- def __call__(self, x, w, conv_state, conv_weight, conv_state_indices, conv_dim, pad_slot_id):
- return fp8_gemv_conv1d(
- x, w, conv_state, conv_weight, conv_state_indices, conv_dim, pad_slot_id
- )
-
- def reference(self, *args, **kwargs):
- raise NotImplementedError("the test builds the fp32 reference from the dequantized weight")
-
- def is_valid(self, x, w, **kwargs) -> None:
- assert x.float().abs().amax().item() <= 65504.0, (
- "an activation beyond f16's range saturates at 65504 in the kernel"
- )
-
-
-class Nvfp4SplitGemvConv1dVerify(OpWrapper):
- """`nvfp4_split_gemv_conv1d_verify` (`fp8_gemv_conv1d_verify` over a `SplitNvfp4Weight`):
- the GDN input projection of a speculative verify step over the NVFP4 copy with
- `causal_conv1d_triton_update`'s multi-token conv step as its epilogue. Returns `(conv [M,
- conv_dim], projection columns conv_dim.. [M, N - conv_dim])`.
-
- Fusion boundary: projection, conv, SiLU, snapshots and the pool advance."""
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec keys as `Fp8GemvConv1dVerify`'s.
- CELLS: tuple[Cell, ...] = fp8_gemv_conv1d_verify_op.CELLS
-
- note = """
- Call site: the decode target's GDN in_proj of an MTP verify step in a captured graph, 2..16
- rows, under `TRTLLM_MODELING_V2_GDN_INPROJ_NVFP4=1`; otherwise as `Fp8GemvConv1dVerify`
- (bitwise the projection + `causal_conv1d_triton_update` over the same NVFP4 projection).
- """
-
- def __call__(
- self,
- x,
- w,
- conv_state,
- conv_weight,
- conv_state_indices,
- intermediate_conv_window,
- intermediate_state_indices,
- conv_dim,
- rows_per_seq,
- pad_slot_id,
- ):
- return fp8_gemv_conv1d_verify(
- x,
- w,
- conv_state,
- conv_weight,
- conv_state_indices,
- intermediate_conv_window,
- intermediate_state_indices,
- conv_dim,
- rows_per_seq,
- pad_slot_id,
- )
-
- def reference(
- self,
- x,
- codes,
- scales,
- global_scale,
- conv_state,
- conv_weight,
- conv_state_indices,
- intermediate_conv_window,
- intermediate_state_indices,
- conv_dim,
- rows_per_seq,
- pad_slot_id,
- ):
- """As `Fp8GemvConv1dVerify.reference`, over the dequantized NVFP4 weight."""
- return verify_reference(
- x,
- dequantize_nvfp4(codes, scales, global_scale),
- conv_state,
- conv_weight,
- conv_state_indices,
- intermediate_conv_window,
- intermediate_state_indices,
- conv_dim,
- rows_per_seq,
- pad_slot_id,
- )
-
- def is_valid(
- self,
- x,
- w,
- conv_state,
- conv_weight,
- conv_state_indices,
- intermediate_conv_window,
- intermediate_state_indices,
- conv_dim,
- rows_per_seq,
- pad_slot_id,
- **kwargs,
- ) -> None:
- assert x.float().abs().amax().item() <= 65504.0, (
- "an activation beyond f16's range saturates at 65504 in the kernel"
- )
- fp8_gemv_conv1d_verify_op.is_valid(
- x,
- w,
- conv_state,
- conv_weight,
- conv_state_indices,
- intermediate_conv_window,
- intermediate_state_indices,
- conv_dim,
- rows_per_seq,
- pad_slot_id,
- )
-
-
-nvfp4_gemv_op = Nvfp4Gemv()
-nvfp4_split_gemv_op = Nvfp4SplitGemv()
-nvfp4_split_gemv_conv1d_op = Nvfp4SplitGemvConv1d()
-nvfp4_split_gemv_conv1d_verify_op = Nvfp4SplitGemvConv1dVerify()
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
deleted file mode 100644
index 9983a66b1c55..000000000000
--- 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
+++ /dev/null
@@ -1,179 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Qwen4-Exp hyper-connection gating: low-rank stream mixer plus block-injection
-weights.
-
-Tier-3 composed from native torch (three ``F.linear`` GEMMs and four elementwise
-passes); it launches no TensorRT-LLM kernel. The ``mhc_*`` op family
-(``mhc_gemm_sqrsum_fma``, ``mhc_big_fuse``, ``mhc_post_mapping``) computes a
-structurally different hyper-connection -- a per-stream *scalar* gate from one
-linear rather than a rank-320 SiLU bottleneck, a weighted *sum* rather than a
-mean, the *raw* stream gated rather than the normalized one -- and its fused
-members do not instantiate hidden 2560. The built-in Qwen4-Exp implementation
-has fused ops for this (``trtllm::qwen4_exp_hc_*``), not wrapped here; the
-semantics follow the HuggingFace ``Qwen4ExpTextGatedResidual`` minus its norm.
-"""
-
-from __future__ import annotations
-
-from typing import Optional, Tuple
-
-import torch
-import torch.nn.functional as F
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: This checkpoint's geometry: 4 streams of 2560, rank-320 mixer bottleneck.
-CKPT_HC_COUNT = 4
-CKPT_HIDDEN = 2560
-CKPT_LOWRANK = 320
-
-
-def gate_ratio(out: torch.Tensor, ref: torch.Tensor) -> float:
- """Worst element's distance from `ref` as a fraction of the test's band:
- `rtol = eps`, `atol = eps * max|ref|`, both of `ref`'s dtype. Below 1 is
- inside."""
- eps = torch.finfo(ref.dtype).eps
- ref64 = ref.double()
- atol = eps * ref64.abs().max().item()
- return ((out.double() - ref64).abs() / (atol + eps * ref64.abs())).max().item()
-
-
-class Qwen4ExpGatedResidual(OpWrapper):
- """Collapse the `hc_count`-stream residual to one stream, and weight the
- block's injection. With `C = hc_count`, `H = width // C` and `normed` the
- `[..., C*H]` stream **after** its grouped norm, viewed stream-major as
- `[..., C, H]`::
-
- low = silu(normed @ down^T / C) # [..., lowrank]
- gate = sigmoid(low @ up^T) # [..., C*H]
- mixed = mean over streams of (gate * normed) # [..., H]
- inject = 2 * sigmoid(normed @ block_inject^T / C) # [..., C]
-
- Returns fresh contiguous `(mixed, inject)` in the activation dtype;
- `inject` is `None` exactly when `block_inject_weight` is `None` (the
- stack-final mixer), and `mixed` is then bitwise what it is with one.
-
- Both divisions by `C` sit on the pre-activation logit; `mixed` gates the
- **normed** stream; `inject` carries a factor of 2 (a zero logit is an
- identity injection). GEMMs run in the activation dtype (fp32 accumulate),
- the pointwise chain in fp32, each result rounded once.
-
- Fusion boundary: both projections, both activations, the gating, the
- stream mean and the injection projection are inside. The grouped gemma
- norm producing `normed`, and the accumulate `qwen4_exp_post_inject` (into
- the **pre-norm** stream the caller still holds), are not. Weights are used
- as stored (`nn.Linear` layout), nothing is cached.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: leading (token dims), inject (True). Every cell is this
- # checkpoint's 4x2560 bf16 geometry with synthetic weights of its measured
- # element statistics, the normed stream reaching ~+-30 as the real one does.
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="checkpoint 4x2560 bf16, decode batch 8 / MTP draft rows",
- spec=dict(leading=(8,)),
- ),
- Cell(
- why="checkpoint 4x2560 bf16, prefill 8192 tokens, largest measured",
- spec=dict(leading=(8192,)),
- ),
- Cell(
- why="stack-final mixer: block_inject_weight None",
- spec=dict(leading=(256,), inject=False),
- ),
- )
-
- note = """
- Metadata consumed: none; every row independent. Call sites: twice per
- decoder layer (attention and MoE hyper-connections, 96 calls) plus the
- stack-final mixer with `block_inject_weight=None` -- whose `mixed` *is*
- the checkpoint's final norm -- and three more in the MTP layer (two
- blocks and `shared_head`'s own mixer).
-
- Preconditions, all loud (each measured to raise `RuntimeError` before
- producing anything; the width 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`).
-
- float32 activations are accepted and not certified: `F.linear` on fp32
- follows torch's TF32 default here, ~10 mantissa bits. A literal bf16
- transcription of the pointwise chain is ~2x this entry's error (0.0449 vs
- 0.0227 on `mixed` with the real layer-0 weights).
-
- Operand range: with the real layer-0 weights the normed stream reaches
- |30.1|, the logits 47.8 / 67.5, `inject` spans 0.430 .. 1.578. At up to
- 4096x that stream every output stays finite, `|mixed| <= max|normed|`,
- `inject` in `[0, 2]`; agreement with float64 degrades there because the
- sigmoid saturates (a property of the function, outside the checkpoint's
- range).
-
- Certified: bf16 2-D `[rows, 10240]` at 8 and 8192 rows with the
- injection, 256 rows without it (the stack-final mixer); CUDA-graph replay
- at 8 rows, bitwise equal to eager. Cost, not gated, shared device: ~235 us
- at 1..512 tokens (overhead-bound), 475 us at 4096, 902 us at 8192.
-
- Not certified: other row counts, more than one leading dim, hc_count other
- than 4, hidden other than 2560, lowrank other than 320, float16 and
- float32, weight/activation dtype mixing; the operand-range bounds above
- are measured. Hard limits: not probed (largest operand 8192 x 10240 =
- 2^26.3 elements).
- """
-
- def __call__(
- self,
- hyper_input_normed: torch.Tensor,
- input_mix_weight_down: torch.Tensor,
- input_mix_weight_up: torch.Tensor,
- block_inject_weight: Optional[torch.Tensor],
- hc_count: int,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
- activation_dtype = hyper_input_normed.dtype
- hidden_size = hyper_input_normed.shape[-1] // hc_count
-
- # GEMMs in the activation dtype: a bf16 product accumulates in fp32 with
- # no ambient state, an fp32 one would follow torch's TF32 default.
- low_rank = F.linear(hyper_input_normed, input_mix_weight_down)
- low_rank = F.silu(low_rank.float() / hc_count).to(activation_dtype)
-
- mix_weight = F.linear(low_rank, input_mix_weight_up)
- mix_weight = torch.sigmoid(mix_weight.float()).unflatten(-1, (hc_count, hidden_size))
- streams = hyper_input_normed.float().unflatten(-1, (hc_count, hidden_size))
- mixed_input = (mix_weight * streams).mean(dim=-2).to(activation_dtype)
-
- injection_weights = None
- if block_inject_weight is not None:
- inject_logits = F.linear(hyper_input_normed, block_inject_weight)
- injection_weights = 2.0 * torch.sigmoid(inject_logits.float() / hc_count)
- injection_weights = injection_weights.to(activation_dtype)
-
- return mixed_input, injection_weights
-
- def reference(
- self,
- hyper_input_normed: torch.Tensor,
- input_mix_weight_down: torch.Tensor,
- input_mix_weight_up: torch.Tensor,
- block_inject_weight: Optional[torch.Tensor],
- hc_count: int,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
- """float64 throughout -- an fp32 GEMM here would be TF32 and the less
- accurate side -- rounded once to the activation dtype."""
- dtype = hyper_input_normed.dtype
- hidden = hyper_input_normed.shape[-1] // hc_count
- x = hyper_input_normed.double()
- low = F.silu((x @ input_mix_weight_down.double().t()) / hc_count)
- gate = torch.sigmoid(low @ input_mix_weight_up.double().t()).unflatten(
- -1, (hc_count, hidden)
- )
- mixed = (gate * x.unflatten(-1, (hc_count, hidden))).mean(dim=-2).to(dtype)
- inject = None
- if block_inject_weight is not None:
- inject = 2.0 * torch.sigmoid((x @ block_inject_weight.double().t()) / hc_count)
- inject = inject.to(dtype)
- return mixed, inject
-
-
-qwen4_exp_gated_residual = Qwen4ExpGatedResidual()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_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
deleted file mode 100644
index 8e58e9df35cc..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_cute.py
+++ /dev/null
@@ -1,605 +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, for one-row batches: the previous
-block's injection (with the MoE block's combine and the routed experts' finalize
-folded in, when given) and the first projection per stream, one PDL launch.
-
-Tier 5 (CuTe DSL written for this target). A stream is exactly `THREADS x VEC`
-elements, one 16-byte vector per thread (a power-of-two tile would pad a 2560-wide
-stream to 8 x 512, three eighths of it masked lanes holding registers), and the grid
-is sized so that every CTA is resident at once (one wave), so all of them stream their
-weight before the producer finishes.
-
-Order of work inside a CTA (`OUTS` weight rows x one stream):
-
-1. its weight rows and its slice of the norm weight into registers -- static, so
- before `griddepcontrol.wait`, while the kernel ahead of this one runs;
-2. `griddepcontrol.wait`; then the stream's vector of the residual and of the
- block output -- with the fold, the slot map and the weights first, then the
- slot rows -- and the injection weight;
-3. the injected residual rounded once (stored by the stream's first CTA), its sum
- of squares and its `(1 + w)`-scaled products with the weight rows, in fp32;
-4. a warp then a cross-warp reduction; thread `j < OUTS` stores partial `j`, thread
- `OUTS` of the stream's first CTA the sum of squares.
-"""
-
-from __future__ import annotations
-
-import functools
-from typing import Optional, Tuple
-
-import cuda.bindings.driver as cuda_drv
-import cutlass
-import cutlass.cute as cute
-import torch
-from cutlass import const_expr
-from cutlass.cute import experimental as cute_ext
-from cutlass.cute.runtime import from_dlpack
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
-
-#: One decode row per launch: the forward runs only 1-row decode batches through the
-#: site's GEMV pair (`_HC_DECODE_GEMV_MAX_ROWS`).
-MAX_ROWS = 1
-VEC = 8
-#: Threads per CTA: one 16-byte vector of a 2560-wide stream each.
-THREADS = 320
-#: Streaming multiprocessors the one-wave grid must fit (GB300).
-_SMS = 152
-_COMPILE_OPTIONS = "--ptxas-options -maxrregcount=128"
-
-
-def outs_for(n_cols: int, hc_count: int) -> int:
- """Weight rows per CTA: the fewest that divide `n_cols` and keep the grid in one wave."""
- for outs in range(1, n_cols + 1):
- if n_cols % outs == 0 and (n_cols // outs) * hc_count <= _SMS:
- return outs
- raise AssertionError(f"no one-wave tiling for {n_cols} columns")
-
-
-class _FirstKernel:
- def __init__(self, *, hidden, outs, inject, combine, topk, pdl):
- assert hidden == THREADS * VEC, "one vector per thread"
- self.hidden, self.outs = hidden, outs
- self.inject, self.combine, self.topk, self.pdl = inject, combine, topk, pdl
- self.warps = THREADS // 32
-
- @cute.jit
- def __call__(
- self,
- mHc,
- mBlock,
- mInj,
- mShared,
- mGate,
- mTopW,
- mTopRows,
- mNorm,
- mW,
- mHcOut,
- mPart,
- mSs,
- stream: cuda_drv.CUstream,
- ):
- n_tiles = cute.size(mPart, mode=[1]) // self.outs
- streams = cute.size(mPart, mode=[0])
- load_static = cute.make_copy_atom(
- cute.nvgpu.CopyG2ROp(),
- cutlass.BFloat16,
- num_bits_per_copy=16 * VEC,
- load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING,
- )
- load_act = cute.make_copy_atom(
- cute.nvgpu.CopyG2ROp(),
- cutlass.BFloat16,
- num_bits_per_copy=16 * VEC,
- load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS,
- )
- store = cute.make_copy_atom(
- cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=16 * VEC
- )
- self.kernel(
- mHc,
- mBlock,
- mInj,
- mShared,
- mGate,
- mTopW,
- mTopRows,
- mNorm,
- mW,
- mHcOut,
- mPart,
- mSs,
- load_static,
- load_act,
- store,
- ).launch(
- grid=[n_tiles, streams, 1],
- block=[THREADS, 1, 1],
- smem=(self.outs + 1) * self.warps * 4,
- stream=stream,
- use_pdl=self.pdl,
- # At most 96 registers a thread: the mix kernel that follows (512 threads at
- # 56) then fits on the SM beside this CTA and streams its weights before
- # this kernel ends.
- min_blocks_per_mp=2,
- )
-
- @cute.kernel
- def kernel(
- self,
- mHc,
- mBlock,
- mInj,
- mShared,
- mGate,
- mTopW,
- mTopRows,
- mNorm,
- mW,
- mHcOut,
- mPart,
- mSs,
- load_static: cute.CopyAtom,
- load_act: cute.CopyAtom,
- store: cute.CopyAtom,
- ):
- tidx, _, _ = cute.arch.thread_idx()
- bidx, s, _ = cute.arch.block_idx()
- warp = cute.arch.warp_idx()
- outs: cutlass.Constexpr = self.outs
- warps: cutlass.Constexpr = self.warps
- topk: cutlass.Constexpr = self.topk
- n0 = bidx * outs
- vi = s * THREADS + tidx # this thread's vector of the width; `tidx` of the stream
-
- w_vec = cute.logical_divide(mW, (None, VEC))
- hc_vec = cute.logical_divide(mHc, (None, VEC))
- nw_vec = cute.logical_divide(mNorm, (None, VEC))
-
- # 1. Static operands, before the wait.
- w_reg = cute.make_rmem_tensor(
- cute.make_layout((outs, VEC), stride=(VEC, 1)), cutlass.BFloat16
- )
- for j in cutlass.range_constexpr(outs):
- cute.copy(load_static, w_vec[n0 + j, (None, vi)], w_reg[j, None])
- nw_reg = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16)
- cute.copy(load_static, nw_vec[0, (None, vi)], nw_reg)
- if const_expr(self.pdl):
- cute.arch.griddepcontrol_launch_dependents()
- cute.arch.griddepcontrol_wait()
-
- # 2. The producers' outputs.
- x_reg = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16)
- cute.copy(load_act, hc_vec[0, (None, vi)], x_reg)
- xr = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.Float32)
- for v in cutlass.range_constexpr(VEC):
- xr[v] = x_reg[v].to(cutlass.Float32)
- if const_expr(self.inject):
- b_vec = cute.logical_divide(mBlock, (None, VEC))
- g = mInj[0, s].to(cutlass.Float32)
- b = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.Float32)
- if const_expr(topk > 0):
- # The runner's finalize: fp32 FMAs in slot order from zero, a slot mapped
- # to -1 skipped, one rounding. Every slot row is requested before any FMA.
- rows = cute.make_rmem_tensor(cute.make_layout((topk,), stride=(1,)), cutlass.Int32)
- wts = cute.make_rmem_tensor(cute.make_layout((topk,), stride=(1,)), cutlass.Float32)
- for k in cutlass.range_constexpr(topk):
- rows[k] = mTopRows[0, k]
- wts[k] = mTopW[0, k].to(cutlass.Float32)
- fc2 = cute.make_rmem_tensor(
- cute.make_layout((topk, VEC), stride=(VEC, 1)), cutlass.BFloat16
- )
- for k in cutlass.range_constexpr(topk):
- cute.copy(
- load_act,
- b_vec[cutlass.max(rows[k], cutlass.Int32(0)), (None, tidx)],
- fc2[k, None],
- )
- for v in cutlass.range_constexpr(VEC):
- acc = cutlass.Float32(0.0)
- for k in cutlass.range_constexpr(topk):
- if rows[k] >= 0:
- acc = cute.math.fma(fc2[k, v].to(cutlass.Float32), wts[k], acc)
- b[v] = acc.to(cutlass.BFloat16).to(cutlass.Float32)
- else:
- blk = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16)
- cute.copy(load_act, b_vec[0, (None, tidx)], blk)
- for v in cutlass.range_constexpr(VEC):
- b[v] = blk[v].to(cutlass.Float32)
- if const_expr(self.combine):
- # The MoE block output, routed + sigmoid(gate) * shared in fp32, rounded once.
- sh_vec = cute.logical_divide(mShared, (None, VEC))
- sh = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16)
- cute.copy(load_act, sh_vec[0, (None, tidx)], sh)
- z = mGate[0].to(cutlass.Float32)
- sg = 1.0 / (1.0 + cute.math.exp(-z))
- for v in cutlass.range_constexpr(VEC):
- b[v] = (
- cute.math.fma(sg, sh[v].to(cutlass.Float32), b[v])
- .to(cutlass.BFloat16)
- .to(cutlass.Float32)
- )
- out = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16)
- for v in cutlass.range_constexpr(VEC):
- # Rounded once, like `addcmul`: the norm and every later site read this value.
- out[v] = cute.math.fma(b[v], g, xr[v]).to(cutlass.BFloat16)
- xr[v] = out[v].to(cutlass.Float32)
- if bidx == 0:
- o_vec = cute.logical_divide(mHcOut, (None, VEC))
- cute.copy(store, out, o_vec[0, (None, vi)])
-
- # 3. Sum of squares and the scaled products, fp32.
- acc = cute.make_rmem_tensor(cute.make_layout((outs + 1,), stride=(1,)), cutlass.Float32)
- acc.fill(0.0)
- for v in cutlass.range_constexpr(VEC):
- y = xr[v] * (1.0 + nw_reg[v].to(cutlass.Float32))
- for j in cutlass.range_constexpr(outs):
- acc[j] = y * w_reg[j, v].to(cutlass.Float32) + acc[j]
- acc[outs] = xr[v] * xr[v] + acc[outs]
-
- # 4. Reductions; thread j < OUTS stores partial j, thread OUTS the sum of squares.
- for j in cutlass.range_constexpr(outs + 1):
- acc[j] = cute.arch.warp_reduction_sum(acc[j])
- smem = cutlass.utils.SmemAllocator()
- red = smem.allocate_tensor(
- cutlass.Float32,
- cute.make_layout((outs + 1, warps), stride=(warps, 1)),
- byte_alignment=16,
- )
- with cute.arch.elect_one():
- for j in cutlass.range_constexpr(outs + 1):
- red[j, warp] = acc[j]
- cute.arch.sync_threads()
- if tidx < outs:
- total = cutlass.Float32(0.0)
- for w in cutlass.range_constexpr(warps):
- total = total + red[tidx, w]
- mPart[s, n0 + tidx] = total
- if tidx == outs:
- if bidx == 0:
- total = cutlass.Float32(0.0)
- for w in cutlass.range_constexpr(warps):
- total = total + red[outs, w]
- mSs[s] = total
-
-
-def _dl(t: torch.Tensor, align: int = 16):
- return from_dlpack(t, assumed_align=align)
-
-
-@functools.lru_cache(maxsize=None)
-def _compiled(
- hidden, hc_count, n_cols, n_rows, inject, combine, topk, n_perm, pdl, hc_stride, block_stride
-):
- """One compiled launch per variant and geometry (strides are part of a CuTe layout)."""
- outs = outs_for(n_cols, hc_count)
- kern = _FirstKernel(
- hidden=hidden, outs=outs, inject=inject, combine=combine, topk=topk, pdl=pdl
- )
- dev, bf = "cuda", torch.bfloat16
- width = hidden * hc_count
- hc = torch.empty(1, hc_stride, dtype=bf, device=dev)[:, :width]
- block = torch.empty(max(n_perm, 1), block_stride, dtype=bf, device=dev)[:, :hidden]
- inj = torch.empty(1, hc_count, dtype=bf, device=dev)
- shared = torch.empty(1, hidden, dtype=bf, device=dev)
- gate = torch.empty(1, dtype=bf, device=dev)
- topw = torch.empty(1, max(topk, 1), dtype=bf, device=dev)
- toprows = torch.empty(1, max(topk, 1), dtype=torch.int32, device=dev)
- norm = torch.empty(1, width, dtype=bf, device=dev)
- w = torch.empty(n_rows, width, dtype=bf, device=dev)
- hc_out = torch.empty(1, width, dtype=bf, device=dev)
- part = torch.empty(hc_count, n_cols, dtype=torch.float32, device=dev)
- ss = torch.empty(hc_count, dtype=torch.float32, device=dev)
- stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream)
- return cute_ext.compile(
- kern,
- _dl(hc),
- _dl(block),
- _dl(inj, 2),
- _dl(shared),
- _dl(gate, 2),
- _dl(topw, 2),
- _dl(toprows, 4),
- _dl(norm),
- _dl(w),
- _dl(hc_out),
- _dl(part, 4),
- _dl(ss, 4),
- stream,
- options=_COMPILE_OPTIONS,
- )
-
-
-def launch(
- hyper_input,
- block_out,
- inject,
- norm_weight,
- first,
- hc_count,
- n_cols,
- shared=None,
- gate_logits=None,
- topk_weights=None,
- topk_rows=None,
-):
- """The launch behind `Qwen4ExpHcDecodeFirstCute.__call__`. Returns `(hc_out, part, ss)`,
- `part` fp32 `[C, 1, n_cols]` and `ss` fp32 `[C, 1]`."""
- rows, width = hyper_input.shape
- hidden = width // hc_count
- dev = hyper_input.device
- has_inject = block_out is not None
- has_shared = shared is not None
- top_k = topk_rows.shape[1] if topk_rows is not None else 0
- hc_out = (
- torch.empty((rows, width), dtype=hyper_input.dtype, device=dev)
- if has_inject
- else hyper_input
- )
- part = torch.empty((hc_count, rows, n_cols), dtype=torch.float32, device=dev)
- ss = torch.empty((hc_count, rows), dtype=torch.float32, device=dev)
- if rows == 0:
- return hc_out, part, ss
- pdl = get_env_enable_pdl()
- weight = first.t() # the row-major [n_rows, width] weight
- gate = gate_logits.reshape(1) if has_shared else hyper_input.new_empty(1)
- n_perm = block_out.shape[0] if (has_inject and top_k) else 1
- fn = _compiled(
- hidden,
- hc_count,
- n_cols,
- weight.shape[0],
- has_inject,
- has_shared,
- top_k,
- n_perm,
- bool(pdl),
- hyper_input.stride(0),
- block_out.stride(0) if (has_inject and top_k) else hidden,
- )
- blk = block_out if has_inject else hyper_input.new_empty(1, hidden)
- inj = inject if has_inject else hyper_input.new_empty(1, hc_count)
- sh = shared if has_shared else hyper_input.new_empty(1, hidden)
- topw = topk_weights if top_k else hyper_input.new_empty(1, 1)
- toprows = topk_rows if top_k else torch.empty(1, 1, dtype=torch.int32, device=dev)
- stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream)
- fn(
- _dl(hyper_input),
- _dl(blk),
- _dl(inj, 2),
- _dl(sh),
- _dl(gate, 2),
- _dl(topw, 2),
- _dl(toprows, 4),
- _dl(norm_weight.reshape(1, -1)),
- _dl(weight),
- _dl(hc_out),
- _dl(part[:, 0, :], 4),
- _dl(ss[:, 0], 4),
- stream,
- )
- return hc_out, part, ss
-
-
-class Qwen4ExpHcDecodeFirstCute(OpWrapper):
- """The first half of a one-row decode hyper-connection site: the previous block's
- injection and the mixer's first projection (`[down; block_inject]`) per stream. With
- `C = hc_count`, `H = width // C`, `N = n_cols`, `first` the `[width, >= N]` view of
- the load-time `[down; block_inject; 0]` weight::
-
- hc_out[0, s*H+i] = hyper_input[0, s*H+i] + block_out[0, i] * inject[0, s] (rounded once)
- part[s, 0, n] = sum_i hc_out[0, s*H+i] * (1 + norm_weight[s*H+i]) * first[s*H+i, n]
- ss[s, 0] = sum_i hc_out[0, s*H+i]^2
-
- Returns `(hc_out, part, ss)`: `hc_out` fresh in the activation dtype (it *is*
- `hyper_input` when `block_out` / `inject` are None -- the norm-only site), `part`
- fp32 `[C, 1, N]`, `ss` fp32 `[C, 1]`. The grouped RMS `rstd[s] = (ss[s, 0] / H +
- eps)^-1/2` scales stream `s`'s partial: `normed @ first[:, :N] = sum_s rstd[s] *
- part[s]`. `qwen4_exp_hc_decode_mix_cute` makes that combination.
-
- With `shared` and `gate_logits` (the site after a MoE block), `block_out` is the
- routed experts' output and the block output injected above is `block_out[0, i] +
- sigmoid(gate_logits[0]) * shared[0, i]` in fp32, rounded once to `block_out`'s
- dtype -- bitwise what `fused_sigmoid_gate_mul_add` stores. `gate_logits` is one
- pre-sigmoid scalar per row, `[T, 1]` or `[T]`.
-
- With `topk_weights` / `topk_rows` as well (the MoE runner left unfinalized),
- `block_out` is its permuted FC2 output `[P, H]` and the routed row is the runner's
- finalize, formed here: `routed[0, i] = sum_k topk_weights[0, k] *
- block_out[topk_rows[0, k], i]` as fp32 FMAs in slot order from zero (a slot whose
- row is -1 is skipped), rounded once to `block_out`'s dtype -- bitwise what the
- runner's finalize stores. `topk_rows` is the runner's expanded->permuted index map
- `[T, k]` int32, `topk_weights` its `[T, k]` routing weights. The fold needs the
- combine (`shared`, `gate_logits`).
-
- Fusion boundary: the routed experts' finalize (optional), the shared-expert combine
- (optional), the previous block's injection, the `(1 + w)` scale and the first
- projection per stream, plus the per-stream sum of squares. The routed and shared
- experts and the gate projection, the RMS itself, the activations and everything
- after are not inside.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: rows (1), inject (True), moe (False: shared + gate_logits [T, 1] too), fold
- # (0: top-k of the finalize folded in, block_out then the permuted rows), drop
- # (False: some slots -1), n_cols (324). The weight is the target's padded
- # [width, 384] (columns past n_cols never read).
- # Four cells: the forward reaches four compiled variants (inject at 324 cols, the
- # fold at 324, the fold at 320 into the stack-final mixer, norm-only), none of which
- # another covers; the -1 slots ride the 324-col fold cell.
- CELLS: tuple[Cell, ...] = (
- Cell(why="checkpoint 4x2560, 324 cols of 384, bf16, inject, 1 row", spec=dict(rows=1)),
- Cell(
- why="finalize fold, top-10, last 3 slots mapped to -1, 1 row",
- spec=dict(rows=1, moe=True, fold=10, drop=True),
- ),
- Cell(
- why="finalize fold into the stack-final mixer: 320 cols",
- spec=dict(rows=1, moe=True, fold=10, n_cols=320),
- ),
- Cell(why="norm-only site (stack start / after PLE)", spec=dict(rows=1, inject=False)),
- )
-
- note = """
- Metadata consumed: none. Call sites: `DecodeTarget`'s trunk, every
- hyper-connection site of a one-row decode step (97 per forward), followed by
- `qwen4_exp_hc_decode_mix_cute` (or its NVFP4 variant at an MoE site); batches of
- more rows take `qwen4_exp_hc_inject_norm` and the mixer. The 47 sites after a MoE
- block other than the one before the PLE add pass `shared` and `gate_logits`, and
- with them the unfinalized runner's `topk_weights` / `topk_rows`.
-
- Preconditions: 2-D operands with unit column stride (row strides free); `first`
- `[width, >= n_cols]` with unit stride along `width` (the `.t()` view of a row-major
- `[N, width]` weight; columns past `n_cols` are never read); one row; `hidden` 2560
- (320 threads x 8); bf16 operands; `n_cols` with a divisor that keeps `n_cols / outs x C <= 152`
- CTAs (324 -> 9 rows x 36 tiles x 4 streams = 144; 320 -> 10 rows, 128); the
- permuted rows' count is part of the compiled layout (one compile per value).
-
- Numerics: `hc_out` as `qwen4_exp_hc_inject_norm` (fp32 FMAs, one rounding; the
- finalize and the combine likewise); `part` and `ss` are fp32 reductions in a fixed
- order, inside the float64 reference's band (the combine and the finalize can each
- land one ulp of the block output from it, scaled into `hc_out` by `inject`).
- Deterministic (measured). Graph-capturable: compiled on the first call per
- variant, which must precede capture (the engine's warmup forward does it).
-
- PDL: with `TRTLLM_ENABLE_PDL` (default on, sm_90+) the weight rows and the norm
- slice stream before `griddepcontrol.wait`, every activation read and every
- write after it, and dependents are released before the wait.
-
- Certified: one bf16 row, the injection at 324 columns, the fold (top-10, three
- slots -1) at 324 columns, the fold into the stack-final mixer at 320 columns,
- the norm-only site; graph replay under PDL with the fold.
-
- Not certified: the combine without the fold, a flat `[T]` gate, a row-slice
- `hyper_input`, a weight narrower than the padded 384 rows (accepted; no cell
- drives them).
- """
-
- def __call__(
- self,
- hyper_input: torch.Tensor,
- block_out: Optional[torch.Tensor],
- inject: Optional[torch.Tensor],
- norm_weight: torch.Tensor,
- first: torch.Tensor,
- hc_count: int,
- n_cols: int,
- shared: Optional[torch.Tensor] = None,
- gate_logits: Optional[torch.Tensor] = None,
- topk_weights: Optional[torch.Tensor] = None,
- topk_rows: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
- return launch(
- hyper_input,
- block_out,
- inject,
- norm_weight,
- first,
- hc_count,
- n_cols,
- shared,
- gate_logits,
- topk_weights,
- topk_rows,
- )
-
- def reference(
- self,
- hyper_input: torch.Tensor,
- block_out: Optional[torch.Tensor],
- inject: Optional[torch.Tensor],
- norm_weight: torch.Tensor,
- first: torch.Tensor,
- hc_count: int,
- n_cols: int,
- shared: Optional[torch.Tensor] = None,
- gate_logits: Optional[torch.Tensor] = None,
- topk_weights: Optional[torch.Tensor] = None,
- topk_rows: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
- """float64; the finalize (with `topk_rows`) rounded once to
- `block_out`'s dtype, the combine (with `shared`) rounded once to `block_out`'s
- dtype, `hc_out` rounded once to the activation dtype, `part` and `ss`
- rounded once to fp32."""
- dtype = hyper_input.dtype
- rows, width = hyper_input.shape
- hidden = width // hc_count
- if block_out is None:
- hc_out = hyper_input
- else:
- if topk_rows is not None:
- valid = (topk_rows >= 0).unsqueeze(-1)
- rows_ = block_out.double()[topk_rows.clamp(min=0).long()]
- w = topk_weights.double().unsqueeze(-1)
- block_out = torch.where(valid, rows_ * w, 0.0).sum(1).to(block_out.dtype)
- if shared is not None:
- sg = torch.sigmoid(gate_logits.double().reshape(rows, 1))
- block_out = (block_out.double() + sg * shared.double()).to(block_out.dtype)
- streams = hyper_input.double().unflatten(-1, (hc_count, hidden))
- gated = block_out.double().unsqueeze(-2) * inject.double().unsqueeze(-1)
- hc_out = (streams + gated).flatten(-2).to(dtype)
- x = hc_out.double().unflatten(-1, (hc_count, hidden))
- y = x * (1.0 + norm_weight.double()).unflatten(-1, (hc_count, hidden))
- w = first[:, :n_cols].double().unflatten(0, (hc_count, hidden))
- part = torch.einsum("tsi,sin->stn", y, w).float()
- ss = x.pow(2).sum(dim=-1).t().float()
- return hc_out, part, ss
-
- def is_valid(
- self,
- hyper_input: torch.Tensor,
- block_out: Optional[torch.Tensor],
- inject: Optional[torch.Tensor],
- norm_weight: torch.Tensor,
- first: torch.Tensor,
- hc_count: int,
- n_cols: int,
- shared: Optional[torch.Tensor] = None,
- gate_logits: Optional[torch.Tensor] = None,
- topk_weights: Optional[torch.Tensor] = None,
- topk_rows: Optional[torch.Tensor] = None,
- ) -> None:
- # Raw row strides with unit column stride everywhere; a contiguous
- # [width, N] weight would be read with its rows and columns swapped.
- assert hyper_input.dim() == 2 and hyper_input.stride(1) == 1, "hyper_input [T, width]"
- rows, width = hyper_input.shape
- assert width % hc_count == 0 and (width // hc_count) % 16 == 0, "hidden multiple of 16"
- assert (block_out is None) == (inject is None), "block_out and inject come together"
- assert (shared is None) == (gate_logits is None), "shared and gate_logits come together"
- assert shared is None or block_out is not None, "the combine needs block_out"
- assert (topk_weights is None) == (topk_rows is None), (
- "topk_weights and topk_rows come together"
- )
- assert topk_rows is None or shared is not None, "the finalize fold needs the combine"
- if block_out is not None:
- if topk_rows is None:
- assert block_out.shape == (rows, width // hc_count) and block_out.stride(1) == 1
- else:
- # The permuted FC2 rows: any row count, read through topk_rows.
- assert block_out.dim() == 2 and block_out.shape[1] == width // hc_count
- assert block_out.stride(1) == 1, "permuted rows [P, hidden]"
- assert topk_rows.dtype == torch.int32, "topk_rows int32 (the runner's index map)"
- assert topk_rows.dim() == 2 and topk_rows.shape[0] == rows, "topk_rows [T, k]"
- assert topk_rows.stride(1) == 1 and topk_weights.stride(1) == 1
- assert topk_weights.shape == topk_rows.shape, "topk_weights [T, k]"
- assert inject.shape == (rows, hc_count) and inject.stride(1) == 1
- if shared is not None:
- assert shared.shape == (rows, width // hc_count), "shared [T, hidden]"
- assert shared.stride(1) == 1, "shared [T, hidden]"
- assert shared.dtype == block_out.dtype, "shared in block_out's dtype"
- assert gate_logits.shape in ((rows,), (rows, 1)), "gate_logits [T, 1] or [T]"
- assert rows <= 1 or gate_logits.stride(0) != 0, "one gate logit per row"
- assert norm_weight.shape == (width,) and norm_weight.is_contiguous(), "norm_weight [width]"
- assert first.shape[0] == width and first.shape[1] >= n_cols, "first [width, >= n_cols]"
- assert first.stride(0) == 1, "first must be the .t() view of a row-major [N, width] weight"
- assert hyper_input.shape[0] <= MAX_ROWS, "one row"
- assert hyper_input.shape[1] // hc_count == THREADS * VEC, "hidden 2560"
- assert hyper_input.dtype == torch.bfloat16, "bf16"
- outs_for(n_cols, hc_count)
-
-
-qwen4_exp_hc_decode_first_cute = Qwen4ExpHcDecodeFirstCute()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py
deleted file mode 100644
index df11a7567828..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_cute.py
+++ /dev/null
@@ -1,534 +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, for one-row batches:
-the grouped RMS, the SiLU bottleneck and the injection weights from the first half's
-partials, the up projection, the sigmoid gate, the stream mean and the NVFP4 of the
-mean for the routed experts -- one PDL launch.
-
-Tier 5 (CuTe DSL written for this target). A CTA owns one 16-column NVFP4 block of the
-hidden width across the 4 streams (64 rows of the up weight), and 8 threads share a
-row, 40 weights each; 160 CTAs of 512 threads, all resident at once.
-
-Order of work inside a CTA:
-
-1. its 64 weight rows and its norm-weight elements into registers (static), before
- `griddepcontrol.wait`;
-2. after the wait: the four sums of squares (every thread), one partial column per
- thread for the first 324 threads -> the shared `low` vector and, in CTA 0, the
- injection weights; each thread's residual element;
-3. each row's dot with `low` over its 8 threads, the sigmoid gate on the normed
- element, the mean over streams (16 threads), the store; with NVFP4, warp 0 also
- quantizes the block as `torch.ops.trtllm.fp4_quantize` does (linear scales,
- `rcp.approx.ftz` reciprocals).
-"""
-
-from __future__ import annotations
-
-import functools
-
-import cuda.bindings.driver as cuda_drv
-import cutlass
-import cutlass.cute as cute
-import torch
-import torch.nn.functional as F
-from cutlass import const_expr
-from cutlass._mlir.dialects import llvm
-from cutlass.cute import experimental as cute_ext
-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.quantization.fp4_quantize import (
- fp4_quantize,
-)
-from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
-
-MAX_ROWS = 1
-VEC = 8
-#: Hidden columns per CTA: one NVFP4 scale block.
-BLOCK_H = 16
-#: Threads sharing one up-weight row.
-GROUP = 8
-_COMPILE_OPTIONS = "--ptxas-options -maxrregcount=64"
-
-
-@dsl_user_op
-def _lds_v4f(addr, *, loc=None, ip=None):
- """16 bytes of this CTA's shared memory as 4 fp32 (one `LDS.128`)."""
- ty = T.f32()
- res = llvm.inline_asm(
- llvm.StructType.get_literal([ty] * 4),
- [cutlass.Int32(addr).ir_value(loc=loc, ip=ip)],
- "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];",
- "=f,=f,=f,=f,r",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- return [cutlass.Float32(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(4)]
-
-
-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: the logit is
- # formed on every thread (only its store is predicated), so the partials' loads have
- # a use outside any branch and are not sunk behind the sums of squares' divisions.
- 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)
- # This row's stream's rstd by selects (a dynamic register index would go to local memory).
- rstd_s = rstd[0]
- for t in cutlass.range_constexpr(1, streams):
- if s == t:
- rstd_s = rstd[t]
- smem = cutlass.utils.SmemAllocator()
- # `low` flat with 16 bytes of pad every 32 values: the 8 lanes of a row read 16-byte
- # vectors of 8 different 8-value rows, in 8 different bank groups.
- low = smem.allocate_tensor(
- cutlass.Float32, cute.make_layout((rank + rank // 8,), stride=(1,)), byte_alignment=16
- )
- low_a = cutlass.Int32(low.iterator.toint())
- vals = smem.allocate_tensor(
- cutlass.Float32,
- cute.make_layout((streams, BLOCK_H), stride=(BLOCK_H, 1)),
- byte_alignment=16,
- )
- 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 + (tidx // 32) * 4] = z * sig
- else:
- if const_expr(self.inject_out):
- if tidx < cols:
- 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):
- lrow = kt * GROUP + q
- lv = _lds_v4f(low_a + (lrow * VEC + (lrow // 4) * 4) * 4) + _lds_v4f(
- low_a + (lrow * VEC + (lrow // 4) * 4) * 4 + 16
- )
- for e in cutlass.range_constexpr(VEC):
- sums[e % 2] = u_reg[kt, e].to(cutlass.Float32) * lv[e] + sums[e % 2]
- dot = sums[0] + sums[1]
- dot = cute.arch.warp_reduction(dot, lambda a, b: a + b, threads_in_group=GROUP)
- if q == 0:
- gate = 1.0 / (1.0 + cute.math.exp(-dot))
- vals[s, i] = gate * (x * rstd_s * gam)
- cute.arch.sync_threads()
- if warp == 0:
- m = cutlass.Float32(0.0)
- if lane < BLOCK_H:
- tot = cutlass.Float32(0.0)
- for t in cutlass.range_constexpr(streams):
- tot = tot + vals[t, lane]
- mb = (tot * (1.0 / streams)).to(cutlass.BFloat16)
- mMixed[0, i0 + lane] = mb
- m = mb.to(cutlass.Float32)
- if const_expr(self.fp4):
- # NVFP4 of the 16 rounded values in lanes 0..15 (lanes 16..31 a second,
- # discarded group), as `fp4_quantize`: e4m3(gs * amax / 6), then
- # x * rcp(sf * rcp(gs)) rounded to e2m1, nearest, ties to even.
- amax = cute.arch.warp_reduction_max(cute.math.abs(m), threads_in_group=BLOCK_H)
- sf8 = (gs * (amax * cute.arch.rcp_approx(6.0))).to(cutlass.Float8E4M3FN)
- osc = cute.arch.rcp_approx(sf8.to(cutlass.Float32) * cute.arch.rcp_approx(gs))
- if amax == 0.0:
- osc = cutlass.Float32(0.0)
- sv = m * osc
- a = cute.math.abs(sv)
- code = cutlass.Int32(0)
- if a > 0.25:
- code = code + 1
- if a >= 0.75:
- code = code + 1
- if a > 1.25:
- code = code + 1
- if a >= 1.75:
- code = code + 1
- if a > 2.5:
- code = code + 1
- if a >= 3.5:
- code = code + 1
- if a > 5.0:
- code = code + 1
- if cute.math.copysign(1.0, sv) < 0.0: # the sign bit, -0.0 included
- code = code + 8
- hi = cute.arch.shuffle_sync_down(code, offset=1)
- if lane < BLOCK_H:
- if lane % 2 == 0:
- mXq[0, (i0 + lane) // 2] = (code + hi * 16).to(cutlass.Uint8)
- if lane == 0:
- mXsf[bidx] = sf8
-
-
-def _dl(t: torch.Tensor, align: int = 16):
- return from_dlpack(t, assumed_align=align)
-
-
-@functools.lru_cache(maxsize=None)
-def _compiled(
- hidden, streams, rank, part_cols, part_stride, ss_stride, hc_stride, inject_out, fp4, pdl
-):
- kern = _MixKernel(
- hidden=hidden, streams=streams, rank=rank, inject_out=inject_out, fp4=fp4, pdl=pdl
- )
- dev, bf = "cuda", torch.bfloat16
- width = hidden * streams
- part = torch.empty(streams, part_stride, dtype=torch.float32, device=dev)[:, :part_cols]
- ss = torch.empty(streams * ss_stride, dtype=torch.float32, device=dev)[::ss_stride]
- up = torch.empty(width, rank, dtype=bf, device=dev)
- hc = torch.empty(1, hc_stride, dtype=bf, device=dev)[:, :width]
- norm = torch.empty(1, width, dtype=bf, device=dev)
- gs = torch.empty(1, dtype=torch.float32, device=dev)
- mixed = torch.empty(1, hidden, dtype=bf, device=dev)
- inj = torch.empty(1, streams, dtype=bf, device=dev)
- xq = torch.empty(1, hidden // 2, dtype=torch.uint8, device=dev)
- xsf = torch.empty(hidden // 16, dtype=torch.float8_e4m3fn, device=dev)
- stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream)
- return cute_ext.compile(
- kern,
- _dl(part, 4),
- _dl(ss, 4),
- _dl(up),
- _dl(hc, 2),
- _dl(norm, 2),
- _dl(gs, 4),
- _dl(mixed, 2),
- _dl(inj, 2),
- _dl(xq, 1),
- _dl(xsf, 1),
- cutlass.Float32(1e-6),
- stream,
- options=_COMPILE_OPTIONS,
- )
-
-
-def launch(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale=None):
- """The launch behind `Qwen4ExpHcDecodeMixFp4Cute.__call__`. Returns `(mixed,
- inject_out, xq, xsf)`; without `global_scale` (the kernel's plain build, kept for
- measurement) `xq` / `xsf` are None."""
- rows, width = hc_out.shape
- hidden = width // hc_count
- rank = up.shape[0]
- dev, dt = hc_out.device, hc_out.dtype
- fp4 = global_scale is not None
- mixed = torch.empty((rows, hidden), dtype=dt, device=dev)
- inject_out = torch.empty((rows, hc_count), dtype=dt, device=dev) if with_inject else None
- xq = torch.empty((rows, hidden // 2), dtype=torch.uint8, device=dev) if fp4 else None
- xsf = torch.empty((rows * (hidden // 16),), dtype=torch.uint8, device=dev) if fp4 else None
- if rows == 0:
- return mixed, inject_out, xq, xsf
- pdl = get_env_enable_pdl()
- part2 = part[:, 0, :]
- fn = _compiled(
- hidden,
- hc_count,
- rank,
- part2.shape[1],
- part2.stride(0),
- ss.stride(0),
- hc_out.stride(0),
- bool(with_inject),
- fp4,
- bool(pdl),
- )
- stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream)
- gs = global_scale.reshape(1) if fp4 else mixed.new_empty(1, dtype=torch.float32)
- inj = inject_out if with_inject else mixed.new_empty(1, hc_count)
- q = xq if fp4 else mixed.new_empty(1, hidden // 2, dtype=torch.uint8)
- sf = (
- xsf.view(torch.float8_e4m3fn)
- if fp4
- else mixed.new_empty(hidden // 16, dtype=torch.float8_e4m3fn)
- )
- fn(
- _dl(part2, 4),
- _dl(ss[:, 0], 4),
- _dl(up.t()),
- _dl(hc_out, 2),
- _dl(norm_weight.reshape(1, -1), 2),
- _dl(gs, 4),
- _dl(mixed, 2),
- _dl(inj, 2),
- _dl(q, 1),
- _dl(sf, 1),
- cutlass.Float32(eps),
- stream,
- )
- return mixed, inject_out, xq, xsf
-
-
-def _check(part, ss, up, hc_out, norm_weight, hc_count, with_inject) -> None:
- rows, width = hc_out.shape
- rank = up.shape[0]
- assert hc_out.stride(1) == 1, "hc_out unit column stride"
- assert part.dtype == torch.float32 and part.stride(2) == 1, "part fp32, unit last stride"
- assert part.shape[:2] == (hc_count, rows), "part [C, T, N]"
- assert part.shape[2] >= rank + (hc_count if with_inject else 0), "part too narrow"
- assert ss.dtype == torch.float32 and ss.shape == (hc_count, rows), "ss fp32 [C, T]"
- assert up.shape == (rank, width), "up [R, width]"
- assert norm_weight.shape == (width,) and norm_weight.is_contiguous()
- assert rows <= MAX_ROWS, "one row"
- assert width // hc_count == 2560 and rank == 320 and hc_count == 4, (
- "the checkpoint's 4 x 2560, rank 320"
- )
- assert hc_out.dtype == torch.bfloat16 and up.dtype == torch.bfloat16, "bf16"
- assert up.t().is_contiguous(), "up the [R, width] view of a row-major [width, R] weight"
-
-
-def _reference(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject):
- """`(mixed, inject_out)` in float64, one rounding per output."""
- dtype = hc_out.dtype
- rank = up.shape[0]
- hidden = hc_out.shape[1] // hc_count
- rstd = torch.rsqrt(ss.double().t() / hidden + eps)
- logits = torch.einsum("ts,stn->tn", rstd, part.double())
- low = F.silu(logits[:, :rank] / hc_count)
- gate = torch.sigmoid(low @ up.double())
- x = hc_out.double().unflatten(-1, (hc_count, hidden)) * rstd.unsqueeze(-1)
- normed = x.flatten(-2) * (1.0 + norm_weight.double())
- mixed = (gate * normed).unflatten(-1, (hc_count, hidden)).mean(dim=-2).to(dtype)
- inject_out = None
- if with_inject:
- inject_out = (2.0 * torch.sigmoid(logits[:, rank : rank + hc_count] / hc_count)).to(dtype)
- return mixed, inject_out
-
-
-_INJECT = Cell(why="checkpoint 4x2560 r320 bf16, inject, 1 row (con 1)", spec=dict(rows=1))
-_ROW_SLICE = Cell(why="hc_out a row slice buf[::2], 1 row", spec=dict(rows=1, layout="row_slice"))
-
-
-class Qwen4ExpHcDecodeMixFp4Cute(OpWrapper):
- """`qwen4_exp_hc_decode_mix_cute` plus the NVFP4 of its `mixed` for the routed
- experts. With `C = hc_count`, `H = width // C`, `R = up.shape[0]`, `part` / `ss` the
- per-stream partial logits `[C, T, >= R (+ C)]` and sums of squares `[C, T]` from
- `qwen4_exp_hc_decode_first_cute`, `up` the `[R, width]` view of the up weight::
-
- rstd[t, s] = 1 / sqrt(ss[s, t] / H + eps)
- logits[t, n] = sum_s rstd[t, s] * part[s, t, n]
- low = silu(logits[:, :R] / C)
- gate = sigmoid(low @ up) # [T, C*H]
- normed[t, s*H+i] = hc_out[t, s*H+i] * rstd[t, s] * (1 + norm_weight[s*H+i])
- mixed[t, i] = mean over s of gate[t, s*H+i] * normed[t, s*H+i]
- inject_out = 2 * sigmoid(logits[:, R:R+C] / C) # with_inject only
- xq, xsf = NVFP4(mixed) # e2m1 pairs, e4m3 scale per 16, linear layout
-
- fp32 inside, one rounding per output; the quantization reads the rounded `mixed`.
- Returns fresh `(mixed, inject_out, xq [T, H/2] uint8, xsf [T*H/16] uint8)`;
- `inject_out` is None without `with_inject`. `global_scale` is the 1-element fp32
- NVFP4 global scale.
-
- Fusion boundary: everything `qwen4_exp_hc_decode_mix_cute` fuses, plus the NVFP4
- quantization of `mixed`. The partials are the producer's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
- # Its call sites always inject, so no stack-final (no-inject) cell.
- CELLS: tuple[Cell, ...] = (_INJECT, _ROW_SLICE)
-
- note = """
- Metadata consumed: none. Call sites: `DecodeTarget`'s MoE sites at a one-row
- decode step, after `qwen4_exp_hc_decode_first_cute`; `xq`/`xsf` feed
- `fp4_block_scale_moe_renorm_routed` (linear scales). Batches of more rows take
- `qwen4_exp_hc_up_gate_mean_fp4`.
-
- Preconditions: those of `qwen4_exp_hc_decode_mix_cute`, and `global_scale` fp32 with
- one element.
-
- Numerics: `mixed` and `inject_out` as `qwen4_exp_hc_decode_mix_cute`; the NVFP4 is
- `fp4_quantize`'s arithmetic on the rounded `mixed` (`rcp.approx.ftz` reciprocals,
- e4m3 scale rounded to nearest, e2m1 nearest-even with saturation at 6).
- Deterministic (measured); graph-capturable once compiled (the first call per
- variant). PDL: the weight rows and norm elements stream before
- `griddepcontrol.wait`.
-
- Certified: one bf16 row with the injection, `hc_out` compact and a row slice;
- graph replay under PDL. Not certified: `with_inject=False`.
- """
-
- def __call__(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale):
- return launch(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale)
-
- def reference(
- self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale
- ):
- """float64 with one rounding per output, then `fp4_quantize`'s native-torch NVFP4
- of the rounded `mixed` (linear scales)."""
- mixed, inject_out = _reference(
- part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject
- )
- return (mixed, inject_out) + fp4_quantize.reference(
- mixed, global_scale, BLOCK_H, False, False
- )
-
- def is_valid(
- self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale
- ) -> None:
- _check(part, ss, up, hc_out, norm_weight, hc_count, with_inject)
- # The kernel loads element 0 only; a per-row scale would be applied as row 0's.
- assert global_scale.numel() == 1 and global_scale.dtype == torch.float32, (
- "global_scale fp32 [1]"
- )
-
-
-qwen4_exp_hc_decode_mix_fp4_cute = Qwen4ExpHcDecodeMixFp4Cute()
-
-
-class Qwen4ExpHcDecodeMixCute(OpWrapper):
- """The second half of a one-row decode hyper-connection site: the grouped RMS, the
- SiLU bottleneck and the injection weights from `qwen4_exp_hc_decode_first_cute`'s
- per-stream partials, then the up projection, the sigmoid gate on the normed stream
- and the mean over streams -- `qwen4_exp_hc_decode_mix_fp4_cute`'s formulas without
- the NVFP4 output. Returns fresh `(mixed, inject_out)` in `hc_out`'s dtype;
- `inject_out` is None without `with_inject`.
-
- Fusion boundary: the grouped RMS, the bottleneck SiLU, the injection weights, the
- up projection, the sigmoid gate, the normed stream and the stream mean. The
- partials are the producer's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
- CELLS: tuple[Cell, ...] = (
- _INJECT,
- Cell(why="stack-final mixer (no inject), 1 row", spec=dict(rows=1, with_inject=False)),
- _ROW_SLICE,
- )
-
- note = """
- Metadata consumed: none. Call sites: `DecodeTarget`'s non-MoE sites at a one-row
- decode step, after `qwen4_exp_hc_decode_first_cute`; batches of more rows take
- `qwen4_exp_hc_up_gate_mean`. Beside a first kernel capped at 96 registers, its 160
- CTAs (512 threads at 49 registers) are all resident while that kernel runs, so the
- up weight is in registers when the wait returns.
-
- Preconditions: `part` fp32 `[C, T, >= R (+ C)]` and `ss` fp32 `[C, T]`, `part` with
- unit last stride; `hc_out` `[T, width]` with unit column stride; `norm_weight`
- contiguous `[width]`; one row, the checkpoint's 4 x 2560 width, rank 320, bf16;
- `up` the `.t()` view of a contiguous `[width, R]` weight.
-
- Numerics: fp32 inside (`low`, the gate logits and `normed` are never rounded), one
- rounding per output, inside the float64 reference's band. Deterministic
- (measured); graph-capturable once compiled (the first call per variant).
-
- Certified: one bf16 row with and without the injection, `hc_out` compact and a
- row slice. Not certified: graph replay of this build (the NVFP4 build's
- replay test covers the shared kernel, not this compiled variant).
- """
-
- def __call__(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject):
- mixed, inject_out, _, _ = launch(
- part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject
- )
- return mixed, inject_out
-
- def reference(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject):
- """float64, one rounding per output."""
- return _reference(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject)
-
- def is_valid(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject) -> None:
- _check(part, ss, up, hc_out, norm_weight, hc_count, with_inject)
-
-
-qwen4_exp_hc_decode_mix_cute = Qwen4ExpHcDecodeMixCute()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_rows_tc.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_rows_tc.py
deleted file mode 100644
index d3756c320875..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_rows_tc.py
+++ /dev/null
@@ -1,1338 +0,0 @@
-"""Decode hyper-connection site for 2..16 rows on warp tensor cores: two PDL launches, the
-first emitting the bottleneck `low`, the second the up projection, gate, stream mean and NVFP4.
-
-Tier 5 (CuTe DSL written for this target). A cluster spans the 4 streams x 2 K halves, so the
-first launch finishes the site's first half itself:
-
-first (`_First2`): grid (2 K halves, 4 streams, 14 N chunks of 24 columns), cluster (2, 4, 1),
-8 warps. A CTA forms its rows' injection over its 1280-element K half (rounded once; stored by the
-chunk-0 CTAs), the slice's sums of squares and `y = x * (1 + w)` in bf16; each warp multiplies a
-160-wide K slice of `y` by the CTA's 24 weight rows (staged before the PDL wait) on mma.sync;
-the warps' partials are summed in shared memory. Every CTA pushes rows (2o, 2o+1) of its partial
-tile and of its sums of squares to cluster rank o; rank o then forms, per row,
-`rstd_s = (ss_s / H + eps)^-1/2`, the logits `z = sum_s rstd_s (P_s,0 + P_s,1) / C`, and writes
-`low = bf16(silu(z))` (columns < 320), the injection weights `2 sigmoid(z)` (columns 320..323)
-and the rows' rstd.
-
-mix (`_Mix2`): 160 CTAs of 8 warps (16 hidden columns x 4 streams). Before the wait: the 64
-up-weight rows. After: `low` (640 bytes a row), rstd and the CTA's residual elements; the
-[16 x 64] up projection; gate x normed; the stream mean and its NVFP4.
-
-Accuracy as `hc_tc.py` (bf16 operands into both projections; `hc_out` bitwise).
-"""
-
-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.nvgpu import cpasync, warp
-from cutlass.cute.runtime import from_dlpack
-from cutlass.cutlass_dsl import T, dsl_user_op
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
-
-from .qwen4_exp_hc_inject_norm import _asm, _mapa
-
-#: This checkpoint's hyper-connection bottleneck width.
-RANK = 320
-#: Rows a launch serves (one m16 MMA tile); the forward sends 2..MAX_ROWS-row sites here.
-MAX_ROWS = 16
-M = 16
-VEC = 8
-PAD = 8
-KH = 2 # K halves
-KHW = 1280
-NC = 24 # columns a CTA
-NCH = 14 # 14 x 24 = 336 >= 324
-WARPS = 8
-THREADS = WARPS * 32
-KW = KHW // WARPS # K a warp
-CL = KH * 4 # cluster size
-RPO = M // CL
-#: Item iterations a thread whose activation loads issue together (all of them at <= 16 rows).
-GROUP_IT = 10
-#: Rows from which the finalize fold is split over a K half's 4 CTAs (each forms a quarter of
-#: the half's rows and pushes it to all four); below it each CTA folds its whole K half itself
-#: in its items (4x the gathers, no exchange: 1.4-1.7 us a layer faster at 2 and 3 rows behind
-#: the runner, 1.7-2.3 us slower from 4, in the chained [MoE + site + MoE site] bench).
-FOLD_SPLIT_ROWS = 4
-#: Fold vectors a thread whose gathers issue together, split (all of them at <= 16 rows) and not
-#: (ten 16-byte gathers a vector: register budget).
-GROUP_FOLD = 3
-GROUP_FOLD_ITEMS = 2
-#: Rows of First2's A stage allocated at least: its shared-memory footprint (~90 KB at 2..8
-#: rows, ~118 KB at 16) leaves room beside it for the mix kernel's CTAs of the same site.
-A_MIN_ROWS = 8
-BLOCK_H = 16
-
-
-@dsl_user_op
-def _st_dsmem_f32(addr, val, *, loc=None, ip=None):
- """4 bytes into a cluster CTA's shared memory."""
- _asm(
- T.i32(),
- [cutlass.Int32(addr), cutlass.Float32(val)],
- "{ st.shared::cluster.f32 [$1], $2; mov.u32 $0, 0; }",
- "=r,r,f",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _st_async_f32(addr, val, mbar, *, loc=None, ip=None):
- """4 bytes into a cluster CTA's shared memory, completing on that CTA's mbarrier."""
- _asm(
- T.i32(),
- [cutlass.Int32(addr), cutlass.Float32(val), cutlass.Int32(mbar)],
- "{ st.async.shared::cluster.mbarrier::complete_tx::bytes.f32 [$1], $2, [$3]; "
- "mov.u32 $0, 0; }",
- "=r,r,f,r",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _mbar_init_tx(mbar, nbytes, *, loc=None, ip=None):
- """Init an mbarrier (one arrival), make it visible to the cluster, arrive expecting
- `nbytes` of transactions."""
- _asm(
- T.i32(),
- [cutlass.Int32(mbar), cutlass.Int32(nbytes)],
- "{ mbarrier.init.shared::cta.b64 [$1], 1; fence.mbarrier_init.release.cluster; "
- "mbarrier.arrive.expect_tx.shared::cta.b64 _, [$1], $2; mov.u32 $0, 0; }",
- "=r,r,r",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _mbar_wait0(mbar, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [cutlass.Int32(mbar)],
- "{ .reg .pred p; WAIT: mbarrier.try_wait.parity.shared::cta.b64 p, [$1], 0; "
- "@!p bra WAIT; mov.u32 $0, 0; }",
- "=r,r",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _st_async_v4(addr, a, b, c, d, mbar, *, loc=None, ip=None):
- """16 bytes into a cluster CTA's shared memory, completing on that CTA's mbarrier."""
- _asm(
- T.i32(),
- [cutlass.Int32(addr)] + [cutlass.Int32(v) for v in (a, b, c, d)] + [cutlass.Int32(mbar)],
- "{ st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [$1], {$2, $3, $4, $5}, "
- "[$6]; mov.u32 $0, 0; }",
- "=r,r,r,r,r,r,r",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _bf16x2(lo, hi, *, loc=None, ip=None):
- """Two fp32 values rounded (RN) into one bf16x2 word, `lo` low (exact for bf16 values)."""
- return cutlass.Int32(
- _asm(
- T.i32(),
- [cutlass.Float32(hi), cutlass.Float32(lo)],
- "cvt.rn.bf16x2.f32 $0, $1, $2;",
- "=r,f,f",
- loc,
- ip,
- )
- )
-
-
-class _First2:
- def __init__(self, *, hidden, rank, n_cols, rows, inject, combine, pdl, topk=0):
- assert hidden == KH * KHW and n_cols <= NC * NCH and rows <= M
- assert topk == 0 or combine, "the finalize fold needs the combine"
- self.hidden, self.rank, self.n_cols, self.rows = hidden, rank, n_cols, rows
- self.inject, self.combine, self.pdl, self.topk = inject, combine, pdl, topk
- self.ld = KHW + PAD
- self.rld = NC + 2 # pushed slot: NC partials + the sum of squares (+pad)
-
- def smem_bytes(self):
- def r(n):
- return -(-n // 128) * 128
-
- assert WARPS * M * NC * 4 + M * NC * 4 <= NC * self.ld * 2
- b = r(max(self.rows, A_MIN_ROWS) * self.ld * 2) + r(NC * self.ld * 2)
- b += r(self.rows * (KHW // VEC) * 4)
- b += r(CL * RPO * self.rld * 4) + r(M * 4)
- b += r(KHW * 2) # the norm-weight slice
- if self.topk:
- b += 2 * r(self.rows * self.topk * 4) # the fold's slot rows and weights
- if self.topk and self.rows >= FOLD_SPLIT_ROWS:
- b += r(self.rows * KHW * 2) + 128 # the block rows of the K half, and an mbarrier
- return b + 256 + 128
-
- @cute.jit
- def __call__(
- self,
- mHc,
- mBlock,
- mInj,
- mShared,
- mGate,
- mTopW,
- mTopRows,
- mNorm,
- mW,
- mHcOut,
- mLow,
- mInjOut,
- mRstd,
- eps: cutlass.Float32,
- stream: cuda_drv.CUstream,
- ):
- mma = cute.make_tiled_mma(
- warp.MmaF16BF16Op(cutlass.BFloat16, cutlass.Float32, (16, 8, 16)),
- cute.make_layout((1, 1, 1)),
- )
- g2s = cute.make_copy_atom(cpasync.CopyG2SOp(), cutlass.BFloat16, num_bits_per_copy=128)
- self.kernel(
- mHc,
- mBlock,
- mInj,
- mShared,
- mGate,
- mTopW,
- mTopRows,
- mNorm,
- mW,
- mHcOut,
- mLow,
- mInjOut,
- mRstd,
- eps,
- mma,
- g2s,
- ).launch(
- grid=[KH, 4, NCH],
- block=[THREADS, 1, 1],
- cluster=[KH, 4, 1],
- smem=self.smem_bytes(),
- stream=stream,
- use_pdl=self.pdl,
- min_blocks_per_mp=1,
- )
-
- @cute.kernel
- def kernel(
- self,
- mHc,
- mBlock,
- mInj,
- mShared,
- mGate,
- mTopW,
- mTopRows,
- mNorm,
- mW,
- mHcOut,
- mLow,
- mInjOut,
- mRstd,
- eps: cutlass.Float32,
- mma: cute.TiledMma,
- g2s: cute.CopyAtom,
- ):
- tidx, _, _ = cute.arch.thread_idx()
- kh, s, nc = cute.arch.block_idx()
- warp_i = cute.arch.warp_idx()
- lane = cute.arch.lane_idx()
- rows: cutlass.Constexpr = self.rows
- hidden: cutlass.Constexpr = self.hidden
- rank: cutlass.Constexpr = self.rank
- n_cols: cutlass.Constexpr = self.n_cols
- ld: cutlass.Constexpr = self.ld
- rld: cutlass.Constexpr = self.rld
- nvec: cutlass.Constexpr = KHW // VEC
- k0 = kh * KHW
- n0 = nc * NC
- me = kh + KH * s # cluster rank
-
- smem = cutlass.utils.SmemAllocator()
- # sA holds the R live rows only; the MMA's 16-row A view runs on into sW (rows past R
- # are garbage, their outputs discarded: the MMA's rows are independent).
- sA = smem.allocate_tensor(
- cutlass.BFloat16,
- cute.make_layout((max(rows, A_MIN_ROWS), KHW), stride=(ld, 1)),
- byte_alignment=128,
- )
- sW = smem.allocate_tensor(
- cutlass.BFloat16, cute.make_layout((NC, KHW), stride=(ld, 1)), byte_alignment=128
- )
- sSs = smem.allocate_tensor(
- cutlass.Float32, cute.make_layout((rows, nvec), stride=(nvec, 1)), byte_alignment=128
- )
- # The warps' partials and their sum reuse sW once every warp holds its B fragment.
- wp_ptr = cute.recast_ptr(sW.iterator, dtype=cutlass.Float32)
- sWp = cute.make_tensor(wp_ptr, cute.make_layout((WARPS, M, NC), stride=(M * NC, NC, 1)))
- sAcc = cute.make_tensor(wp_ptr + WARPS * M * NC, cute.make_layout((M, NC), stride=(NC, 1)))
- sRed = smem.allocate_tensor(
- cutlass.Float32,
- cute.make_layout((CL, RPO, rld), stride=(RPO * rld, rld, 1)),
- byte_alignment=128,
- )
- sSsr = smem.allocate_tensor(
- cutlass.Float32, cute.make_layout((M,), stride=(1,)), byte_alignment=128
- )
- sNw = smem.allocate_tensor(
- cutlass.BFloat16, cute.make_layout((1, KHW), stride=(KHW, 1)), byte_alignment=128
- )
- topk: cutlass.Constexpr = self.topk
- split: cutlass.Constexpr = topk > 0 and rows >= FOLD_SPLIT_ROWS
- if const_expr(topk > 0):
- sTopR = smem.allocate_tensor(
- cutlass.Int32, cute.make_layout((rows * topk,), stride=(1,)), byte_alignment=128
- )
- sTopW = smem.allocate_tensor(
- cutlass.Float32, cute.make_layout((rows * topk,), stride=(1,)), byte_alignment=128
- )
- sBar = smem.allocate_tensor(
- cutlass.Int64, cute.make_layout((2 if split else 1,)), byte_alignment=8
- )
- rbar = cutlass.Int32(sBar.iterator.toint())
- bbar = rbar + 8
- if const_expr(split):
- # The K half's block rows (finalize + combine), a quarter formed by each of the 4
- # stream CTAs of this K half and pushed (st.async) into all four.
- sBlk = smem.allocate_tensor(
- cutlass.BFloat16, cute.make_layout((rows, KHW), stride=(KHW, 1)), byte_alignment=128
- )
- blk_a = cutlass.Int32(sBlk.iterator.toint())
- sblk_vec = cute.logical_divide(sBlk, (None, VEC))
- # This CTA receives rows [me * RPO, me * RPO + RPO) of the cluster's partials from all CL
- # CTAs, by st.async on `rbar` (no cluster barrier, no fence, after the pushes).
- rows_me = cutlass.max(cutlass.min(cutlass.Int32(RPO), rows - me * RPO), cutlass.Int32(0))
- if tidx == 0:
- _mbar_init_tx(rbar, rows_me * (CL * (NC + 1) * 4))
- if const_expr(split):
- _mbar_init_tx(bbar, rows * KHW * 2)
- cute.arch.cluster_arrive()
-
- # 1. Before the wait: the norm weight's slice (its own cp.async group, waited on right
- # after the PDL wait), then the CTA's 24 weight rows over its K half.
- nw_vec = cute.logical_divide(mNorm, (None, VEC))
- snw_vec = cute.logical_divide(sNw, (None, VEC))
- if tidx < KHW // VEC:
- cute.copy(
- g2s,
- nw_vec[0, (None, (s * hidden + k0) // VEC + tidx)],
- snw_vec[0, (None, tidx)],
- )
- cute.arch.cp_async_commit_group()
- w_vec = cute.logical_divide(mW, (None, VEC))
- sw_vec = cute.logical_divide(sW, (None, VEC))
- pieces: cutlass.Constexpr = NC * nvec
- col0 = (s * hidden + k0) // VEC
- for i in cutlass.range_constexpr(-(-pieces // THREADS)):
- c = tidx + i * THREADS
- if c < pieces:
- n = c // nvec
- p = c % nvec
- cute.copy(g2s, w_vec[n0 + n, (None, col0 + p)], sw_vec[n, (None, p)])
- cute.arch.cp_async_commit_group()
- if const_expr(self.pdl):
- cute.arch.griddepcontrol_launch_dependents()
- cute.arch.griddepcontrol_wait()
-
- # 2. The K half's injection, sums of squares, y (bf16) into sA. Every activation load of
- # a group of up to GROUP_IT items a thread (the row's injection weight and gate logit
- # included) issues before the first is used: one L2 round trip a group.
- hc_vec = cute.logical_divide(mHc, (None, VEC))
- b_vec = cute.logical_divide(mBlock, (None, VEC))
- sh_vec = cute.logical_divide(mShared, (None, VEC))
- o_vec = cute.logical_divide(mHcOut, (None, VEC))
- sa_vec = cute.logical_divide(sA, (None, VEC))
- items: cutlass.Constexpr = rows * nvec
- n_it: cutlass.Constexpr = -(-items // THREADS)
- if const_expr(split):
- # With the split fold, every load that does not need the slot map issues with it: the
- # items' hc rows and injection weights (one group: <= 10 items a thread at <= 16
- # rows), and the fold quarter's shared-expert rows and gate logits. The gathers
- # are then the fold's one added L2 round trip.
- assert n_it <= GROUP_IT
- fvq: cutlass.Constexpr = KHW // 4 // VEC
- f_items: cutlass.Constexpr = rows * fvq
- n_f: cutlass.Constexpr = -(-f_items // THREADS)
- assert n_f <= GROUP_FOLD
- fcol0 = (kh * KHW + s * (KHW // 4)) // VEC
- xg0 = cute.make_rmem_tensor(
- cute.make_layout((n_it, VEC), stride=(VEC, 1)), cutlass.BFloat16
- )
- ig0 = cute.make_rmem_tensor(cute.make_layout((n_it,), stride=(1,)), cutlass.Float32)
- for k in cutlass.range_constexpr(n_it):
- c = cutlass.min(tidx + k * THREADS, cutlass.Int32(items - 1))
- r = c // nvec
- p = c % nvec
- cute.autovec_copy(hc_vec[r, (None, col0 + p)], xg0[k, None])
- ig0[k] = mInj[r, s].to(cutlass.Float32)
- shq = cute.make_rmem_tensor(
- cute.make_layout((n_f, VEC), stride=(VEC, 1)), cutlass.BFloat16
- )
- zq = cute.make_rmem_tensor(cute.make_layout((n_f,), stride=(1,)), cutlass.Float32)
- for k in cutlass.range_constexpr(n_f):
- c = cutlass.min(tidx + k * THREADS, cutlass.Int32(f_items - 1))
- cute.autovec_copy(sh_vec[c // fvq, (None, fcol0 + c % fvq)], shq[k, None])
- zq[k] = mGate[c // fvq].to(cutlass.Float32)
- if const_expr(topk > 0):
- # The rows' expanded->permuted slot map and routing weights (the gathers below
- # depend on it).
- for i in cutlass.range_constexpr(-(-(rows * topk) // THREADS)):
- cs = tidx + i * THREADS
- if cs < rows * topk:
- sTopR[cs] = mTopRows[cs // topk, cs % topk]
- sTopW[cs] = mTopW[cs // topk, cs % topk].to(cutlass.Float32)
- cute.arch.cp_async_wait_group(1) # the norm slice (the weights may still be landing)
- cute.arch.sync_threads()
- if const_expr(split):
- # The finalize and the combine of this CTA's quarter of its K half (40 vectors a
- # row): the runner's fp32 FMAs in slot order from zero (a -1 slot skipped, a -1 slot
- # reads row 0, unused), rounded once; then fma(sigmoid(gate), shared, routed) rounded
- # once; pushed to the K half's 4 CTAs.
- cute.arch.cluster_wait() # every peer has started (its mbarriers initialized)
- for f0 in cutlass.range_constexpr(0, n_f, GROUP_FOLD):
- f_n: cutlass.Constexpr = min(GROUP_FOLD, n_f - f0)
- fgq = cute.make_rmem_tensor(
- cute.make_layout((f_n, topk, VEC), stride=(topk * VEC, VEC, 1)),
- cutlass.BFloat16,
- )
- for k in cutlass.range_constexpr(f_n):
- c = cutlass.min(tidx + (f0 + k) * THREADS, cutlass.Int32(f_items - 1))
- r = c // fvq
- q = c % fvq
- for j in cutlass.range_constexpr(topk):
- cute.autovec_copy(
- b_vec[
- cutlass.max(sTopR[r * topk + j], cutlass.Int32(0)),
- (None, fcol0 + q),
- ],
- fgq[k, j, None],
- )
- for k in cutlass.range_constexpr(f_n):
- c = tidx + (f0 + k) * THREADS
- if c < f_items:
- r = c // fvq
- q = c % fvq
- sg = 1.0 / (1.0 + cute.math.exp(-zq[f0 + k]))
- bq = [cutlass.Float32(0.0)] * VEC
- for v in cutlass.range_constexpr(VEC):
- acc = cutlass.Float32(0.0)
- for j in cutlass.range_constexpr(topk):
- if sTopR[r * topk + j] >= 0:
- acc = cute.math.fma(
- fgq[k, j, v].to(cutlass.Float32), sTopW[r * topk + j], acc
- )
- bv = acc.to(cutlass.BFloat16).to(cutlass.Float32)
- bq[v] = (
- cute.math.fma(sg, shq[f0 + k, v].to(cutlass.Float32), bv)
- .to(cutlass.BFloat16)
- .to(cutlass.Float32)
- )
- words = [_bf16x2(bq[2 * w], bq[2 * w + 1]) for w in range(4)]
- dst = blk_a + (r * KHW + s * (KHW // 4) + q * VEC) * 2
- for sp in cutlass.range_constexpr(4):
- peer = kh + KH * sp
- _st_async_v4(_mapa(dst, peer), *words, _mapa(bbar, peer))
- nwr = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16)
- outr = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16)
- yr = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16)
- group: cutlass.Constexpr = GROUP_FOLD_ITEMS if topk > 0 and not split else GROUP_IT
- for g0 in cutlass.range_constexpr(0, n_it, group):
- g_n: cutlass.Constexpr = min(group, n_it - g0)
- if const_expr(split):
- xg = xg0
- else:
- xg = cute.make_rmem_tensor(
- cute.make_layout((g_n, VEC), stride=(VEC, 1)), cutlass.BFloat16
- )
- bg = cute.make_rmem_tensor(
- cute.make_layout((g_n, VEC), stride=(VEC, 1)), cutlass.BFloat16
- )
- shg = cute.make_rmem_tensor(
- cute.make_layout((g_n, VEC), stride=(VEC, 1)), cutlass.BFloat16
- )
- if const_expr(split):
- ig = ig0
- else:
- ig = cute.make_rmem_tensor(cute.make_layout((g_n,), stride=(1,)), cutlass.Float32)
- zg = cute.make_rmem_tensor(cute.make_layout((g_n,), stride=(1,)), cutlass.Float32)
- if const_expr(topk > 0 and not split):
- fg = cute.make_rmem_tensor(
- cute.make_layout((g_n, topk, VEC), stride=(topk * VEC, VEC, 1)),
- cutlass.BFloat16,
- )
- for k in cutlass.range_constexpr(0 if split else g_n):
- # Past the last item a thread reads the last item's operands (unused).
- c = cutlass.min(tidx + (g0 + k) * THREADS, cutlass.Int32(items - 1))
- r = c // nvec
- p = c % nvec
- cute.autovec_copy(hc_vec[r, (None, col0 + p)], xg[k, None])
- if const_expr(self.inject):
- if const_expr(topk > 0):
- # The finalize's slot rows of the permuted FC2 output (a -1 slot reads
- # row 0, unused).
- for j in cutlass.range_constexpr(topk):
- cute.autovec_copy(
- b_vec[
- cutlass.max(sTopR[r * topk + j], cutlass.Int32(0)),
- (None, k0 // VEC + p),
- ],
- fg[k, j, None],
- )
- else:
- cute.autovec_copy(b_vec[r, (None, k0 // VEC + p)], bg[k, None])
- ig[k] = mInj[r, s].to(cutlass.Float32)
- if const_expr(self.combine):
- cute.autovec_copy(sh_vec[r, (None, k0 // VEC + p)], shg[k, None])
- zg[k] = mGate[r].to(cutlass.Float32)
- if const_expr(split and g0 == 0):
- _mbar_wait0(bbar) # the K half's block rows have landed
- for k in cutlass.range_constexpr(g_n):
- c = tidx + (g0 + k) * THREADS
- if c < items:
- r = c // nvec
- p = c % nvec
- cute.autovec_copy(snw_vec[0, (None, p)], nwr)
- ssum = cutlass.Float32(0.0)
- if const_expr(self.inject):
- g = ig[k]
- if const_expr(split):
- cute.autovec_copy(sblk_vec[r, (None, p)], bg[k, None])
- if const_expr(self.combine and not split):
- sg = 1.0 / (1.0 + cute.math.exp(-zg[k]))
- for v in cutlass.range_constexpr(VEC):
- if const_expr(topk > 0 and not split):
- # The runner's finalize: fp32 FMAs in slot order from zero, a
- # -1 slot skipped, one rounding.
- acc = cutlass.Float32(0.0)
- for j in cutlass.range_constexpr(topk):
- if sTopR[r * topk + j] >= 0:
- acc = cute.math.fma(
- fg[k, j, v].to(cutlass.Float32),
- sTopW[r * topk + j],
- acc,
- )
- bv = acc.to(cutlass.BFloat16).to(cutlass.Float32)
- else:
- bv = bg[k, v].to(cutlass.Float32)
- if const_expr(self.combine and not split):
- bv = (
- cute.math.fma(sg, shg[k, v].to(cutlass.Float32), bv)
- .to(cutlass.BFloat16)
- .to(cutlass.Float32)
- )
- outr[v] = cute.math.fma(bv, g, xg[k, v].to(cutlass.Float32)).to(
- cutlass.BFloat16
- )
- if nc == 0:
- cute.autovec_copy(outr, o_vec[r, (None, col0 + p)])
- else:
- for v in cutlass.range_constexpr(VEC):
- outr[v] = xg[k, v]
- for v in cutlass.range_constexpr(VEC):
- xv = outr[v].to(cutlass.Float32)
- ssum = xv * xv + ssum
- yr[v] = (xv * (1.0 + nwr[v].to(cutlass.Float32))).to(cutlass.BFloat16)
- cute.autovec_copy(yr, sa_vec[r, (None, p)])
- sSs[r, p] = ssum
- cute.arch.cp_async_wait_group(0)
- cute.arch.sync_threads()
-
- # 3. Warp w: rows x its 160-wide K slice x the 24 columns; the warps' partials summed in
- # warp order; the K half's sums of squares per row.
- # The warp's K slice: 320-byte offsets, so 16-byte aligned (ldmatrix rows).
- a_w = cute.make_tensor(
- cute.make_ptr(
- cutlass.BFloat16,
- sA.iterator.toint() + warp_i * (KW * 2),
- cute.AddressSpace.smem,
- assumed_align=16,
- ),
- cute.make_layout((M, KW), stride=(ld, 1)),
- )
- b_w = cute.make_tensor(
- cute.make_ptr(
- cutlass.BFloat16,
- sW.iterator.toint() + warp_i * (KW * 2),
- cute.AddressSpace.smem,
- assumed_align=16,
- ),
- cute.make_layout((NC, KW), stride=(ld, 1)),
- )
- thr = mma.get_slice(lane)
- tCsA = thr.partition_A(a_w)
- tCrA = thr.make_fragment_A(tCsA)
- tCsB = thr.partition_B(b_w)
- tCrB = thr.make_fragment_B(tCsB)
- # Fragments by ldmatrix (4 / 2 8x8 matrices an instruction).
- ldsm_a = cute.make_tiled_copy_A(
- cute.make_copy_atom(warp.LdMatrix8x8x16bOp(False, 4), cutlass.BFloat16), mma
- )
- ldsm_b = cute.make_tiled_copy_B(
- cute.make_copy_atom(warp.LdMatrix8x8x16bOp(False, 2), cutlass.BFloat16), mma
- )
- cpa = ldsm_a.get_slice(lane)
- cpb = ldsm_b.get_slice(lane)
- cute.copy(ldsm_a, cpa.partition_S(a_w), cpa.retile(tCrA))
- cute.copy(ldsm_b, cpb.partition_S(b_w), cpb.retile(tCrB))
- acc = cute.make_rmem_tensor(thr.partition_shape_C((M, NC)), cutlass.Float32)
- acc.fill(0.0)
- cute.gemm(mma, acc, tCrA, tCrB, acc)
- cute.arch.sync_threads() # every warp's fragments are loaded: sW may be reused
- co = thr.partition_C(cute.make_identity_tensor((M, NC)))
- for j in cutlass.range_constexpr(cute.size(acc)):
- rr, nn = co[j]
- sWp[warp_i, rr, nn] = acc[j]
- # Sums of squares: 16 threads a row (10 vectors each), then a 16-lane butterfly.
- segs: cutlass.Constexpr = THREADS // M
- per_seg: cutlass.Constexpr = nvec // segs
- rq = tidx // segs
- sq = tidx % segs
- t = cutlass.Float32(0.0)
- if rq < rows:
- for p in cutlass.range_constexpr(per_seg):
- t = t + sSs[rq, sq * per_seg + p]
- t = cute.arch.warp_reduction_sum(t, threads_in_group=segs)
- if sq == 0:
- if rq < rows:
- sSsr[rq] = t
- cute.arch.sync_threads()
- for i in cutlass.range_constexpr(-(-(M * NC) // THREADS)):
- e = tidx + i * THREADS
- if e < M * NC:
- rr = e // NC
- nn = e % NC
- t = cutlass.Float32(0.0)
- for w in cutlass.range_constexpr(WARPS):
- t = t + sWp[w, rr, nn]
- sAcc[rr, nn] = t
- cute.arch.sync_threads()
-
- # 4. Push rows (2o, 2o+1) (and their sums of squares) to cluster rank o, slot `me`.
- if const_expr(not split):
- cute.arch.cluster_wait()
- red_addr = cutlass.Int32(sRed.iterator.toint())
- per: cutlass.Constexpr = RPO * (NC + 1)
- for it in cutlass.range_constexpr(-(-(CL * per) // THREADS)):
- c = tidx + it * THREADS
- if c < CL * per:
- o = c // per
- rem = c % per
- rl = rem // (NC + 1)
- nn = rem % (NC + 1)
- rr = o * RPO + rl
- if rr < rows:
- v = cutlass.Float32(0.0)
- if nn < NC:
- v = sAcc[rr, nn]
- else:
- v = sSsr[rr]
- dst = red_addr + ((me * RPO + rl) * rld + nn) * 4
- _st_async_f32(_mapa(dst, o), v, _mapa(rbar, o))
- _mbar_wait0(rbar)
-
- # 5. Rank `me` finishes rows (2 me, 2 me + 1): rstd per stream, the logits, low / injection.
- for it in cutlass.range_constexpr(-(-(RPO * NC) // THREADS)):
- c = tidx + it * THREADS
- if c < RPO * NC:
- rl = c // NC
- nn = c % NC
- rr = me * RPO + rl
- col = n0 + nn
- if rr < rows:
- if col < n_cols:
- z = cutlass.Float32(0.0)
- for st in cutlass.range_constexpr(4):
- ssv = sRed[KH * st, rl, NC] + sRed[KH * st + 1, rl, NC]
- rs = cute.math.rsqrt(ssv / hidden + eps)
- z = (sRed[KH * st, rl, nn] + sRed[KH * st + 1, rl, nn]) * rs + z
- z = z * 0.25
- sig = 1.0 / (1.0 + cute.math.exp(-z))
- if col < rank:
- mLow[rr, col] = (z * sig).to(cutlass.BFloat16)
- else:
- mInjOut[rr, col - rank] = (2.0 * sig).to(cutlass.BFloat16)
- if nc == 0:
- if tidx < RPO * 4:
- rl = tidx // 4
- st = tidx % 4
- rr = me * RPO + rl
- if rr < rows:
- ssv = sRed[KH * st, rl, NC] + sRed[KH * st + 1, rl, NC]
- mRstd[rr, st] = cute.math.rsqrt(ssv / hidden + eps)
-
-
-def _dl(t: torch.Tensor, align: int = 16):
- return from_dlpack(t, assumed_align=align)
-
-
-@functools.lru_cache(maxsize=None)
-def _compiled_first(
- hidden,
- hc_count,
- rank,
- n_cols,
- n_wrows,
- rows,
- inject,
- combine,
- pdl,
- hc_stride,
- block_stride,
- topk=0,
- n_perm=0,
-):
- kern = _First2(
- hidden=hidden,
- rank=rank,
- n_cols=n_cols,
- rows=rows,
- inject=inject,
- combine=combine,
- pdl=pdl,
- topk=topk,
- )
- dev, bf = "cuda", torch.bfloat16
- width = hidden * hc_count
- hc = torch.empty(rows, hc_stride, dtype=bf, device=dev)[:, :width]
- block = torch.empty(n_perm if topk else rows, block_stride, dtype=bf, device=dev)[:, :hidden]
- topw = torch.empty(rows, max(topk, 1), dtype=bf, device=dev)
- toprows = torch.empty(rows, max(topk, 1), dtype=torch.int32, device=dev)
- inj = torch.empty(rows, hc_count, dtype=bf, device=dev)
- shared = torch.empty(rows, hidden, dtype=bf, device=dev)
- gate = torch.empty(rows, dtype=bf, device=dev)
- norm = torch.empty(1, width, dtype=bf, device=dev)
- w = torch.empty(n_wrows, width, dtype=bf, device=dev)
- hc_out = torch.empty(rows, width, dtype=bf, device=dev)
- low = torch.empty(rows, rank, dtype=bf, device=dev)
- inj_out = torch.empty(rows, hc_count, dtype=bf, device=dev)
- rstd = torch.empty(rows, 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(low),
- _dl(inj_out, 2),
- _dl(rstd, 4),
- cutlass.Float32(1e-6),
- stream,
- )
-
-
-def first_rows(
- hyper_input,
- block_out,
- inject,
- norm_weight,
- first,
- hc_count,
- n_cols,
- eps,
- shared=None,
- gate_logits=None,
- topk_weights=None,
- topk_rows=None,
- rank=320,
-):
- """Returns `(hc_out, low [R, rank] bf16, inject_out [R, C] bf16 or None, rstd [R, C] fp32)`;
- `n_cols` = rank + C with the injection, rank without (the stack-final mixer). With
- `topk_weights` / `topk_rows` (the runner left unfinalized), `block_out` is its permuted FC2
- output `[P, H]` and the routed rows are its finalize, formed here."""
- 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
- hc_out = (
- torch.empty((rows, width), dtype=hyper_input.dtype, device=dev)
- if has_inject
- else hyper_input
- )
- low = torch.empty((rows, rank), dtype=hyper_input.dtype, device=dev)
- inj_out = torch.empty((rows, hc_count), dtype=hyper_input.dtype, device=dev)
- rstd = torch.empty((rows, hc_count), dtype=torch.float32, device=dev)
- pdl = get_env_enable_pdl()
- weight = first.t()
- assert weight.shape[0] >= NC * NCH
- top_k = topk_rows.shape[1] if topk_rows is not None else 0
- fn = _compiled_first(
- hidden,
- hc_count,
- rank,
- n_cols,
- weight.shape[0],
- rows,
- has_inject,
- has_shared,
- bool(pdl),
- hyper_input.stride(0),
- block_out.stride(0) if has_inject else hidden,
- top_k,
- block_out.shape[0] if top_k else 0,
- )
- blk = block_out if has_inject else hyper_input.new_empty(rows, hidden)
- inj = inject if has_inject else hyper_input.new_empty(rows, hc_count)
- sh = shared if has_shared else hyper_input.new_empty(rows, hidden)
- gate = gate_logits.reshape(rows) if has_shared else hyper_input.new_empty(rows)
- topw = topk_weights if top_k else hyper_input.new_empty(rows, 1)
- toprows = topk_rows if top_k else torch.empty(rows, 1, dtype=torch.int32, device=dev)
- stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream)
- 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(low),
- _dl(inj_out, 2),
- _dl(rstd, 4),
- cutlass.Float32(eps),
- stream,
- )
- return hc_out, low, (inj_out if n_cols > rank else None), rstd
-
-
-class _Mix2:
- def __init__(self, *, hidden, streams, rank, rows, fp4, pdl):
- assert streams * BLOCK_H == WARPS * 8 and rows <= M and rank % 16 == 0
- self.hidden, self.streams, self.rank, self.rows, self.fp4, self.pdl = (
- hidden,
- streams,
- rank,
- rows,
- fp4,
- pdl,
- )
- self.ld = rank + PAD
-
- def smem_bytes(self):
- def r(n):
- return -(-n // 128) * 128
-
- nrow = self.streams * BLOCK_H
- b = r(nrow * self.ld * 2) + r(M * self.ld * 2)
- b += r(M * nrow * 4) + r(M * nrow * 2) + r(M * self.streams * 4) + r(nrow * 4)
- return b + 256
-
- @cute.jit
- def __call__(
- self, mLow, mRstd, mUp, mHc, mNorm, mGs, mMixed, mXq, mXsf, stream: cuda_drv.CUstream
- ):
- 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)
- self.kernel(mLow, mRstd, mUp, mHc, mNorm, mGs, mMixed, mXq, mXsf, mma, g2s).launch(
- grid=[self.hidden // BLOCK_H, 1, 1],
- block=[THREADS, 1, 1],
- smem=self.smem_bytes(),
- stream=stream,
- use_pdl=self.pdl,
- min_blocks_per_mp=2,
- )
-
- @cute.kernel
- def kernel(
- self,
- mLow,
- mRstd,
- mUp,
- mHc,
- mNorm,
- mGs,
- mMixed,
- mXq,
- mXsf,
- mma: cute.TiledMma,
- g2s: cute.CopyAtom,
- ):
- tidx, _, _ = cute.arch.thread_idx()
- bidx, _, _ = cute.arch.block_idx()
- warp_i = cute.arch.warp_idx()
- lane = cute.arch.lane_idx()
- hidden: cutlass.Constexpr = self.hidden
- streams: cutlass.Constexpr = self.streams
- rank: cutlass.Constexpr = self.rank
- rows: cutlass.Constexpr = self.rows
- ld: cutlass.Constexpr = self.ld
- nrow: cutlass.Constexpr = streams * BLOCK_H
- i0 = bidx * BLOCK_H
-
- smem = cutlass.utils.SmemAllocator()
- sU = smem.allocate_tensor(
- cutlass.BFloat16, cute.make_layout((nrow, rank), stride=(ld, 1)), byte_alignment=128
- )
- sL = smem.allocate_tensor(
- cutlass.BFloat16, cute.make_layout((M, rank), stride=(ld, 1)), byte_alignment=128
- )
- sV = smem.allocate_tensor(
- cutlass.Float32, cute.make_layout((M, nrow), stride=(nrow, 1)), byte_alignment=128
- )
- sX = smem.allocate_tensor(
- cutlass.BFloat16, cute.make_layout((M, nrow), stride=(nrow, 1)), byte_alignment=128
- )
- sR = smem.allocate_tensor(
- cutlass.Float32, cute.make_layout((M, streams), stride=(streams, 1)), byte_alignment=128
- )
- sG = smem.allocate_tensor(
- cutlass.Float32, cute.make_layout((nrow,), stride=(1,)), byte_alignment=128
- )
-
- u_vec = cute.logical_divide(mUp, (None, VEC))
- su_vec = cute.logical_divide(sU, (None, VEC))
- pieces: cutlass.Constexpr = nrow * (rank // VEC)
- for i in cutlass.range_constexpr(-(-pieces // THREADS)):
- c = tidx + i * THREADS
- if c < pieces:
- n = c // (rank // VEC)
- p = c % (rank // VEC)
- jrow = (n // BLOCK_H) * hidden + i0 + n % BLOCK_H
- cute.copy(g2s, u_vec[jrow, (None, p)], su_vec[n, (None, p)])
- cute.arch.cp_async_commit_group()
- if tidx < nrow:
- sG[tidx] = 1.0 + mNorm[0, (tidx // BLOCK_H) * hidden + i0 + tidx % BLOCK_H].to(
- cutlass.Float32
- )
- # Rows of sL past R are never written (MMA rows are independent; discarded).
- gs = cutlass.Float32(0.0)
- if const_expr(self.fp4):
- gs = mGs[0]
- # Before the wait as well: the up rows land, and every thread's B fragments go to
- # registers (the post-wait MMA then reads only its A fragments from shared memory).
- if const_expr(self.pdl):
- cute.arch.griddepcontrol_launch_dependents()
- cute.arch.cp_async_wait_group(0)
- cute.arch.sync_threads()
- thr = mma.get_slice(tidx)
- tCsB = thr.partition_B(sU)
- tCrB = thr.make_fragment_B(tCsB)
- cute.autovec_copy(tCsB, tCrB)
- if const_expr(self.pdl):
- cute.arch.griddepcontrol_wait()
-
- # low rows (16-byte pieces), rstd, this CTA's residual elements.
- l_vec = cute.logical_divide(mLow, (None, VEC))
- sl_vec = cute.logical_divide(sL, (None, VEC))
- lp: cutlass.Constexpr = rows * (rank // VEC)
- for i in cutlass.range_constexpr(-(-lp // THREADS)):
- c = tidx + i * THREADS
- if c < lp:
- rr = c // (rank // VEC)
- p = c % (rank // VEC)
- cute.copy(g2s, l_vec[rr, (None, p)], sl_vec[rr, (None, p)])
- cute.arch.cp_async_commit_group()
- if tidx < rows * streams:
- sR[tidx // streams, tidx % streams] = mRstd[tidx // streams, tidx % streams]
- h_vec = cute.logical_divide(mHc, (None, VEC))
- sx_vec = cute.logical_divide(sX, (None, VEC))
- xp: cutlass.Constexpr = rows * (nrow // VEC)
- if tidx < xp:
- rr = tidx // (nrow // VEC)
- q = tidx % (nrow // VEC)
- cute.copy(
- g2s,
- h_vec[
- rr,
- (
- None,
- ((q * VEC) // BLOCK_H) * (hidden // VEC)
- + (i0 + (q * VEC) % BLOCK_H) // VEC,
- ),
- ],
- sx_vec[rr, (None, q)],
- )
- cute.arch.cp_async_commit_group()
- cute.arch.cp_async_wait_group(0)
- cute.arch.sync_threads()
-
- tCsA = thr.partition_A(sL)
- tCrA = thr.make_fragment_A(tCsA)
- ldsm_a = cute.make_tiled_copy_A(
- cute.make_copy_atom(warp.LdMatrix8x8x16bOp(False, 4), cutlass.BFloat16), mma
- )
- cpa = ldsm_a.get_slice(tidx)
- cute.copy(ldsm_a, cpa.partition_S(sL), cpa.retile(tCrA))
- acc = cute.make_rmem_tensor(thr.partition_shape_C((M, nrow)), cutlass.Float32)
- acc.fill(0.0)
- cute.gemm(mma, acc, tCrA, tCrB, acc)
- co = thr.partition_C(cute.make_identity_tensor((M, nrow)))
- for j in cutlass.range_constexpr(cute.size(acc)):
- rr, n = co[j]
- if rr < rows:
- gate = 1.0 / (1.0 + cute.math.exp(-acc[j]))
- x = sX[rr, n].to(cutlass.Float32)
- sV[rr, n] = gate * (x * sR[rr, n // BLOCK_H] * sG[n])
- cute.arch.sync_threads()
-
- for rb in cutlass.range_constexpr(0, rows, WARPS):
- r = rb + warp_i
- if r < rows:
- m = cutlass.Float32(0.0)
- if lane < BLOCK_H:
- tot = cutlass.Float32(0.0)
- for t in cutlass.range_constexpr(streams):
- tot = tot + sV[r, t * BLOCK_H + lane]
- mb = (tot * (1.0 / streams)).to(cutlass.BFloat16)
- mMixed[r, i0 + lane] = mb
- m = mb.to(cutlass.Float32)
- if const_expr(self.fp4):
- 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:
- code = code + 8
- hi = cute.arch.shuffle_sync_down(code, offset=1)
- if lane < BLOCK_H:
- if lane % 2 == 0:
- mXq[r, (i0 + lane) // 2] = (code + hi * 16).to(cutlass.Uint8)
- if lane == 0:
- mXsf[r * (hidden // BLOCK_H) + bidx] = sf8
-
-
-@functools.lru_cache(maxsize=None)
-def _compiled_mix(hidden, streams, rank, rows, hc_stride, fp4, pdl):
- kern = _Mix2(hidden=hidden, streams=streams, rank=rank, rows=rows, fp4=fp4, pdl=pdl)
- dev, bf = "cuda", torch.bfloat16
- width = hidden * streams
- low = torch.empty(rows, rank, dtype=bf, device=dev)
- rstd = torch.empty(rows, streams, dtype=torch.float32, device=dev)
- up = torch.empty(width, rank, dtype=bf, device=dev)
- hc = torch.empty(rows, 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(rows, hidden, dtype=bf, device=dev)
- xq = torch.empty(rows, hidden // 2, dtype=torch.uint8, device=dev)
- xsf = torch.empty(rows * (hidden // 16), dtype=torch.float8_e4m3fn, device=dev)
- stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream)
- return cute_ext.compile(
- kern,
- _dl(low),
- _dl(rstd, 4),
- _dl(up),
- _dl(hc, 16),
- _dl(norm, 2),
- _dl(gs, 4),
- _dl(mixed, 2),
- _dl(xq, 1),
- _dl(xsf, 1),
- stream,
- )
-
-
-def mix_rows(low, rstd, up, hc_out, norm_weight, hc_count, global_scale=None):
- """Returns `(mixed, xq, xsf)` (xq / xsf None without `global_scale`)."""
- 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)
- 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
- pdl = get_env_enable_pdl()
- fn = _compiled_mix(hidden, hc_count, rank, rows, hc_out.stride(0), 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)
- q = xq if fp4 else mixed.new_empty(rows, hidden // 2, dtype=torch.uint8)
- sf = (
- xsf.view(torch.float8_e4m3fn)
- if fp4
- else mixed.new_empty(rows * (hidden // 16), dtype=torch.float8_e4m3fn)
- )
- fn(
- _dl(low),
- _dl(rstd, 4),
- _dl(up.t()),
- _dl(hc_out, 16),
- _dl(norm_weight.reshape(1, -1), 2),
- _dl(gs, 4),
- _dl(mixed, 2),
- _dl(q, 1),
- _dl(sf, 1),
- stream,
- )
- return mixed, xq, xsf
-
-
-def _reference_first(
- hyper_input,
- block_out,
- inject,
- norm_weight,
- first,
- hc_count,
- n_cols,
- eps,
- shared=None,
- gate_logits=None,
- topk_weights=None,
- topk_rows=None,
-):
- """float64: the finalize (with `topk_rows`), the combine and the injection each rounded once
- as the 1-row kernel does, the grouped RMS, the first projection over `normed`, `low =
- silu(z)` and `2 sigmoid(z)` (z = logits / C) each rounded once, rstd in fp32."""
- dt = hyper_input.dtype
- rows, width = hyper_input.shape
- hidden = width // hc_count
- rank = RANK
- if block_out is None:
- hc_out = hyper_input
- else:
- if topk_rows is not None:
- valid = (topk_rows >= 0).unsqueeze(-1)
- gathered = block_out.double()[topk_rows.clamp(min=0).long()]
- w = topk_weights.double().unsqueeze(-1)
- block_out = torch.where(valid, gathered * w, 0.0).sum(1).to(block_out.dtype)
- blk = block_out.double()
- if shared is not None:
- sg = torch.sigmoid(gate_logits.double().reshape(rows, 1))
- blk = (blk + sg * shared.double()).to(dt).double()
- streams = hyper_input.double().unflatten(-1, (hc_count, hidden))
- hc_out = (streams + blk.unsqueeze(-2) * inject.double().unsqueeze(-1)).flatten(-2).to(dt)
- x = hc_out.double().unflatten(-1, (hc_count, hidden))
- rstd = torch.rsqrt(x.pow(2).mean(-1) + eps)
- normed = (x * rstd.unsqueeze(-1)).flatten(-2) * (1.0 + norm_weight.double())
- z = normed @ first[:, :n_cols].double() / hc_count
- low = torch.nn.functional.silu(z[:, :rank]).to(dt)
- inj = (2.0 * torch.sigmoid(z[:, rank:])).to(dt) if n_cols > rank else None
- return hc_out, low, inj, rstd.float()
-
-
-class Qwen4ExpHcDecodeRowsFirst(OpWrapper):
- """First launch of a 2..16-row decode hyper-connection site: the previous block's injection
- (and, with `shared` / `gate_logits`, the MoE block's combine `routed + sigmoid(gate) * shared`
- rounded once), the grouped RMS, the first projection over `[down; block_inject]`, the SiLU
- bottleneck and the injection weights. Returns `(hc_out, low [R, 320] bf16, inject_out [R, C]
- or None, rstd [R, C] fp32)`; `hc_out` *is* `hyper_input` without `block_out`.
-
- `hc_out` is bitwise the 1-row kernels'; `low` and the injection weights come from bf16
- operands on tensor cores (`y = bf16(x * (1 + w))` against the bf16 weight) with fp32
- accumulation, inside 1-2 bf16 ulp of the float64 composition.
-
- 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 -1 slot skipped), rounded once -- bitwise what the runner's finalize
- stores (as `qwen4_exp_hc_decode_first_cute`'s fold). The fold needs the combine.
-
- Fusion boundary: the routed experts' finalize (optional), the combine, injection, grouped
- RMS statistics, first projection and activations.
- """
-
- ARCHS = frozenset({Arch.SM_103})
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="site after a MoE block: combine + injection, 2 rows",
- spec=dict(rows=2, site="combine"),
- ),
- Cell(
- why="site after a MoE block, runner finalize folded in, 4 rows (MTP3 verify)",
- spec=dict(rows=4, site="fold"),
- ),
- Cell(why="MoE site: injection, 8 rows", spec=dict(rows=8, site="inject")),
- Cell(
- why="norm-only site (stack start, MTP head), 16 rows, 320 cols",
- spec=dict(rows=16, site="norm"),
- ),
- )
- note = """
- Metadata consumed: none. Call sites: `DecodeTarget._hc_site` / `_hc_moe_site` at 2..16 rows
- (decode batches, MTP verify and draft-extend rows), followed by `qwen4_exp_hc_decode_rows_mix`;
- after a TRTLLM-backend MoE block (`DecodeTarget._moe_block` on the permuted runner) with the
- fold. Grid (2 K halves, 4 streams, 14 N chunks of 24 columns), cluster (2, 4, 1), 8 warps;
- weights staged before the PDL wait. With the fold from 4 rows each of a K half's 4 CTAs
- forms a quarter of the half's routed + shared rows and pushes it to all four (DSMEM); at 2
- and 3 rows each CTA folds its K half in its items. One compile per row count and variant.
- """
-
- def __call__(
- self,
- hyper_input,
- block_out,
- inject,
- norm_weight,
- first,
- hc_count,
- n_cols,
- eps,
- shared=None,
- gate_logits=None,
- topk_weights=None,
- topk_rows=None,
- ):
- return first_rows(
- hyper_input,
- block_out,
- inject,
- norm_weight,
- first,
- hc_count,
- n_cols,
- eps,
- shared,
- gate_logits,
- topk_weights,
- topk_rows,
- )
-
- def reference(
- self,
- hyper_input,
- block_out,
- inject,
- norm_weight,
- first,
- hc_count,
- n_cols,
- eps,
- shared=None,
- gate_logits=None,
- topk_weights=None,
- topk_rows=None,
- ):
- return _reference_first(
- hyper_input,
- block_out,
- inject,
- norm_weight,
- first,
- hc_count,
- n_cols,
- eps,
- shared,
- gate_logits,
- topk_weights,
- topk_rows,
- )
-
- def is_valid(
- self,
- hyper_input,
- block_out,
- inject,
- norm_weight,
- first,
- hc_count,
- n_cols,
- eps,
- shared=None,
- gate_logits=None,
- topk_weights=None,
- topk_rows=None,
- ) -> None:
- rows, width = hyper_input.shape
- assert 1 <= rows <= MAX_ROWS, f"1..{MAX_ROWS} rows"
- assert hyper_input.dtype == torch.bfloat16 and hyper_input.stride(1) == 1
- assert width // hc_count == KH * KHW, "hidden 2560"
- assert first.stride(0) == 1 and first.shape[1] >= NC * NCH, (
- "the padded [width, 384] .t() view"
- )
- assert (block_out is None) == (inject is None)
- assert (shared is None) == (gate_logits is None)
- assert (topk_weights is None) == (topk_rows is None), "topk_weights and topk_rows together"
- if topk_rows is not None:
- assert shared is not None, "the finalize fold needs the combine"
- assert topk_rows.dtype == torch.int32 and topk_rows.shape[0] == rows
- assert topk_rows.is_contiguous() and topk_weights.is_contiguous()
- assert topk_weights.shape == topk_rows.shape and topk_weights.dtype == torch.bfloat16
- assert block_out.dim() == 2 and block_out.shape[1] == width // hc_count
- assert block_out.stride(1) == 1, "permuted rows [P, hidden]"
-
-
-class Qwen4ExpHcDecodeRowsMix(OpWrapper):
- """Second launch of a 2..16-row decode hyper-connection site: the up projection of `low`, the
- sigmoid gate on the normed stream (`hc_out * rstd * (1 + w)`), the stream mean, and with
- `global_scale` the NVFP4 of the mean for the routed experts. Returns `(mixed, xq, xsf)`
- (`xq` / `xsf` None without `global_scale`); one rounding per output.
- """
-
- ARCHS = frozenset({Arch.SM_103})
- CELLS: tuple[Cell, ...] = (
- Cell(why="MoE site with NVFP4 output, 8 rows", spec=dict(rows=8, site="inject")),
- Cell(why="plain mix, 2 rows", spec=dict(rows=2, site="combine")),
- Cell(why="plain mix, 16 rows", spec=dict(rows=16, site="norm")),
- )
- note = """
- Metadata consumed: none. Call sites: after `qwen4_exp_hc_decode_rows_first`. 160 CTAs of 8
- warps (16 hidden columns x 4 streams); the up weight's 64 rows staged before the PDL wait.
- """
-
- def __call__(self, low, rstd, up, hc_out, norm_weight, hc_count, global_scale=None):
- return mix_rows(low, rstd, up, hc_out, norm_weight, hc_count, global_scale)
-
- def reference(self, low, rstd, up, hc_out, norm_weight, hc_count, global_scale=None):
- """float64 mixed, rounded once (the NVFP4 is checked against the kernel's own mixed)."""
- rows, width = hc_out.shape
- hidden = width // hc_count
- gate = torch.sigmoid(low.double() @ up.double())
- x = hc_out.double().unflatten(-1, (hc_count, hidden)) * rstd.double().unsqueeze(-1)
- normed = x.flatten(-2) * (1.0 + norm_weight.double())
- mixed = (gate * normed).unflatten(-1, (hc_count, hidden)).mean(-2).to(hc_out.dtype)
- return mixed, None, None
-
- def is_valid(self, low, rstd, up, hc_out, norm_weight, hc_count, global_scale=None) -> None:
- rows = hc_out.shape[0]
- assert 1 <= rows <= MAX_ROWS and low.shape[0] == rows and rstd.shape == (rows, hc_count)
- assert up.t().is_contiguous(), "up the [R, width] view of a row-major [width, R] weight"
- assert global_scale is None or global_scale.numel() == 1
-
-
-qwen4_exp_hc_decode_rows_first = Qwen4ExpHcDecodeRowsFirst()
-qwen4_exp_hc_decode_rows_mix = Qwen4ExpHcDecodeRowsMix()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_first_rescale_tc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_first_rescale_tc.py
deleted file mode 100644
index 9c4a5dbeadcc..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_first_rescale_tc.py
+++ /dev/null
@@ -1,969 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Hyper-connection first projection on the un-normalized residual, one tcgen05
-launch: the A operand is the site's normed stream formed on chip from `hc_out`
-and its per-(row, stream) rstd, so `[T, C*H]` normed never exists in memory.
-
-Tier 6 (CuTe DSL written for this target). A 2-CTA tcgen05 GEMM (M = 256 rows a
-CTA pair, K = 64 a stage, N = 352 as two MMAs of 192 + 160 columns). Per CTA: a
-TMA warp lands the raw 128 x 64 `hc_out` tile (128B-swizzled) in the operand
-ring; eight transform warps scale it in place by `rstd * (1 + w)`, round to bf16
-and commit the stage to the pair's leader; a TMA warp streams both halves of the
-first weight (2-CTA multicast); the leader's MMA warp accumulates in tensor
-memory. The same eight warps then run the epilogue (TMEM -> registers -> SiLU /
-2 sigmoid -> HBM): `qwen4_exp_hc_lowrank_act`'s activations without the fp32
-logits round trip.
-
-With few row tiles (`_pick_split`) each CTA pair covers `1 / split` of K and
-writes fp32 partial logits `[split, T, 352]`; a Triton combine sums them in split
-order and applies the activations.
-"""
-
-from __future__ import annotations
-
-from typing import Optional, Tuple
-
-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
-import triton
-import triton.language as tl
-from cutlass._mlir.dialects import llvm
-from cutlass.cute.nvgpu import cpasync, tcgen05
-from cutlass.cute.runtime import from_dlpack
-from cutlass.cutlass_dsl import T, dsl_user_op
-from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
-from tensorrt_llm._torch.kernels.blackwell.custom_pipeline import PipelineTmaUmma, PipelineUmmaAsync
-from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
-
-from .qwen4_exp_hc_inject_norm import _bf16_hi, _bf16_lo, _bf16x2, _ldg_v4
-from .qwen4_exp_hc_up_gate_mean import _asm, _st_global_v4
-
-#: Programmatic dependent launch (sm_90+), off with TRTLLM_ENABLE_PDL=0.
-_PDL = get_env_enable_pdl()
-
-#: K a stage (one 128B-swizzled bf16 row of the operand tile).
-BK = 64
-#: Columns of the two MMAs a K step: 192 + 160 = 352 >= rank + C, the padded first
-#: weight's columns the cuBLAS path multiplies (`hc_first_cols`). The 2-CTA atom needs
-#: N / 2 a multiple of 16 per CTA, so two 176-column halves do not build.
-NH0 = 192
-NH1 = 160
-NCOL = NH0 + NH1
-#: Operand (A) ring depth: A is the DRAM stream (`hc_out`), B (the weight) comes from L2.
-_A_STAGES = 6
-_B_STAGES = 5
-#: Transform warp groups: group g transforms k blocks j = g (mod groups), so that many
-#: k blocks are in flight in the transform at once (its per-k-block latency is what
-#: bounds a single group).
-_GROUPS = 2
-#: Largest K split (`_pick_split`).
-_MAX_SPLIT = 8
-
-
-@dsl_user_op
-def _lds_v4(addr, *, loc=None, ip=None):
- """16 bytes of shared memory. Side-effecting, so never moved across the barrier
- wait that makes them valid.
- """
- ty = T.i32()
- res = _asm(
- llvm.StructType.get_literal([ty] * 4),
- [cutlass.Int32(addr)],
- "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];",
- "=r,=r,=r,=r,r",
- loc,
- ip,
- side_effects=True,
- )
- return [cutlass.Int32(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(4)]
-
-
-@dsl_user_op
-def _sts_v4(addr, a, b, c, d, *, loc=None, ip=None):
- """16 bytes to shared memory."""
- _asm(
- T.i32(),
- [cutlass.Int32(addr)] + [cutlass.Int32(v) for v in (a, b, c, d)],
- "{ st.shared.v4.b32 [$1], {$2, $3, $4, $5}; mov.u32 $0, 0; }",
- "=r,r,r,r,r,r",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _ldg_f32(addr, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Int64(addr)],
- "ld.global.f32 $0, [$1];",
- "=f,l",
- loc,
- ip,
- side_effects=True,
- )
- )
-
-
-@dsl_user_op
-def _bits(x, *, loc=None, ip=None):
- return cutlass.Int32(_asm(T.i32(), [cutlass.Float32(x)], "mov.b32 $0, $1;", "=r,f", loc, ip))
-
-
-@dsl_user_op
-def _mul(a, b, *, loc=None, ip=None):
- """`mul.rn.f32`: never contracted into an FMA, so the rounding is the norm's."""
- 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 _add1(a, *, loc=None, ip=None):
- """`1 + a` in fp32 (`add.rn`): the grouped norm's `(1 + w)`."""
- return cutlass.Float32(
- _asm(T.f32(), [cutlass.Float32(a)], "add.rn.f32 $0, $1, 0f3F800000;", "=f,f", loc, ip)
- )
-
-
-@dsl_user_op
-def _sigmoid_fast(x, *, loc=None, ip=None):
- """1 / (1 + 2^(-x log2 e)) with ex2.approx and rcp.approx (two MUFU ops, no
- division).
- """
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Float32(x)],
- "{ .reg .f32 t; mul.f32 t, $1, 0fBFB8AA3B; ex2.approx.ftz.f32 t, t; "
- "add.f32 t, t, 0f3F800000; rcp.approx.ftz.f32 $0, t; }",
- "=f,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _st_global_v2(addr, a, b, *, loc=None, ip=None):
- """8 bytes to global memory (a row's four bf16 injection weights)."""
- _asm(
- 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,
- )
-
-
-class HcFirstRescaleKernel:
- """`low = silu(A @ first[:rank]^T / C)`, `inject = 2 sigmoid(A @ first[rank:rank+C]^T / C)`
- with `A = bf16(hc_out * rstd * (1 + w))` formed in the operand ring (`split == 1`), or the
- fp32 partial logits of K slice `sp` (`split > 1`).
- """
-
- def __init__(self, hc: int, hidden: int, rank: int, split: int):
- assert (hc * hidden) % BK == 0 and hidden % BK == 0 and (hc * hidden // BK) % split == 0
- assert rank + hc <= NCOL and rank % 16 == 0 and hc == 4
- self.hc, self.hidden, self.rank, self.split = hc, hidden, rank, split
- self.kt = hc * hidden // BK // split # k blocks a CTA pair
- self.num_stages = _A_STAGES
- self.num_b_stages = _B_STAGES
- self.groups = _GROUPS
- self.acc_dtype = cutlass.Float32
- self.io_dtype = cutlass.BFloat16
- self.cta_group = tcgen05.CtaGroup.TWO
- self.cluster_shape_mn = (2, 1)
- self.mma_tiler = (256, NH0, BK)
- self.mma_tiler1 = (256, NH1, BK)
- self.tw = 8 # transform / epilogue warps
- self.mma_warp_id = 8
- self.tma_warp_id = 9
- self.tmb_warp_id = 10
- self.threads_per_cta = 32 * 11
- self.epi_bar_id = 1
- self.tmem_bar_id = 2
-
- def _setup(self):
- k_major = utils.LayoutEnum.ROW_MAJOR.mma_major_mode()
- tiled_mma = utils.sm100.make_trivial_tiled_mma(
- self.io_dtype,
- self.io_dtype,
- k_major,
- k_major,
- self.acc_dtype,
- self.cta_group,
- self.mma_tiler[:2],
- )
- tiled_mma1 = utils.sm100.make_trivial_tiled_mma(
- self.io_dtype,
- self.io_dtype,
- k_major,
- k_major,
- self.acc_dtype,
- self.cta_group,
- self.mma_tiler1[:2],
- )
- self.cluster_layout_vmnk = cute.tiled_divide(
- cute.make_layout((*self.cluster_shape_mn, 1)), (tiled_mma.thr_id.shape,)
- )
- self.a_smem_staged = utils.sm100.make_smem_layout_a(
- tiled_mma, self.mma_tiler, self.io_dtype, self.num_stages
- )
- self.b_smem_staged = utils.sm100.make_smem_layout_b(
- tiled_mma, self.mma_tiler, self.io_dtype, self.num_b_stages
- )
- self.b1_smem_staged = utils.sm100.make_smem_layout_b(
- tiled_mma1, self.mma_tiler1, self.io_dtype, self.num_b_stages
- )
- self.num_tmem_cols = 512
- return tiled_mma, tiled_mma1
-
- @cute.jit
- def __call__(
- self,
- m: cutlass.Int32,
- hc_ptr: cute.Pointer,
- w_ptr: cute.Pointer,
- rstd: cute.Tensor,
- b_ptr: cute.Pointer,
- out_ptr: cute.Pointer,
- inj_ptr: cute.Pointer,
- stream: cuda.CUstream,
- ):
- width = self.hc * self.hidden
- tiled_mma, tiled_mma1 = self._setup()
- a = cute.make_tensor(hc_ptr, cute.make_layout((m, width, 1), stride=(width, 1, 0)))
- b = cute.make_tensor(b_ptr, cute.make_layout((NH0, width, 1), stride=(width, 1, 0)))
- b1 = cute.make_tensor(
- b_ptr + NH0 * width, cute.make_layout((NH1, width, 1), stride=(width, 1, 0))
- )
- a_one = cute.slice_(self.a_smem_staged, (None, None, None, 0))
- # The raw A tile of a CTA: a plain 1-CTA TMA into the operand stage, transformed
- # in place.
- tma_atom_a, tma_a = cute.nvgpu.make_tiled_tma_atom_A(
- cpasync.CopyBulkTensorTileG2SOp(),
- a,
- a_one,
- 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_one = cute.slice_(self.b_smem_staged, (None, None, None, 0))
- tma_atom_b, tma_b = cute.nvgpu.make_tiled_tma_atom_B(
- b_op, b, b_one, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape
- )
- b1_one = cute.slice_(self.b1_smem_staged, (None, None, None, 0))
- tma_atom_b1, tma_b1 = cute.nvgpu.make_tiled_tma_atom_B(
- b_op, b1, b1_one, self.mma_tiler1, tiled_mma1, self.cluster_layout_vmnk.shape
- )
- self.a_bytes = cute.size_in_bytes(self.io_dtype, a_one)
- # Both MMAs' B, both CTAs of the pair.
- self.b_bytes = (
- cute.size_in_bytes(self.io_dtype, b_one) + cute.size_in_bytes(self.io_dtype, b1_one)
- ) * 2
- pairs = (m + 255) // 256
- self.kernel(
- tiled_mma,
- tiled_mma1,
- tma_atom_a,
- tma_a,
- tma_atom_b,
- tma_b,
- tma_atom_b1,
- tma_b1,
- w_ptr.toint(),
- rstd.iterator.toint(),
- out_ptr.toint(),
- inj_ptr.toint(),
- m,
- self.cluster_layout_vmnk,
- self.a_smem_staged,
- self.b_smem_staged,
- self.b1_smem_staged,
- ).launch(
- grid=[2 * pairs * self.split, 1, 1],
- 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,
- tiled_mma1: cute.TiledMma,
- tma_atom_a: cute.CopyAtom,
- mA_mkl: cute.Tensor,
- tma_atom_b: cute.CopyAtom,
- mB_nkl: cute.Tensor,
- tma_atom_b1: cute.CopyAtom,
- mB1_nkl: cute.Tensor,
- w_base: cutlass.Int64,
- rstd_base: cutlass.Int64,
- out_base: cutlass.Int64,
- inj_base: cutlass.Int64,
- m: cutlass.Int32,
- cluster_layout_vmnk: cute.Layout,
- a_smem_staged: cute.ComposedLayout,
- b_smem_staged: cute.ComposedLayout,
- b1_smem_staged: cute.ComposedLayout,
- ):
- hc, hid, rank, split = self.hc, self.hidden, self.rank, self.split
- warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
- tidx, _, _ = cute.arch.thread_idx()
- bid, _, _ = cute.arch.block_idx()
- v = bid % 2
- is_leader = v == 0
- pair = bid // 2
- m_pair = pair // split
- sp = pair % split
- kt0 = sp * self.kt
- cta_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
- cluster_coord = cluster_layout_vmnk.get_flat_coord(cta_rank)
- if warp_idx == self.tma_warp_id:
- cpasync.prefetch_descriptor(tma_atom_a)
- if warp_idx == self.tmb_warp_id:
- cpasync.prefetch_descriptor(tma_atom_b)
- cpasync.prefetch_descriptor(tma_atom_b1)
-
- @cute.struct
- class SharedStorage:
- raw_full: cute.struct.MemRange[cutlass.Int64, self.num_stages]
- b_mbar: cute.struct.MemRange[cutlass.Int64, self.num_b_stages * 2]
- a_mbar: cute.struct.MemRange[cutlass.Int64, self.num_stages * 2]
- acc_mbar: cute.struct.MemRange[cutlass.Int64, 2]
- tmem_dealloc_mbar: cutlass.Int64
- tmem_holding_buf: cutlass.Int32
-
- smem = utils.SmemAllocator()
- storage = smem.allocate(SharedStorage)
- # Raw tile landed (TMA transaction count), per CTA: the transform warps wait on it.
- raw_full = storage.raw_full.data_ptr()
- if warp_idx == self.tma_warp_id:
- with cute.arch.elect_one():
- for i in cutlass.range_constexpr(self.num_stages):
- cute.arch.mbarrier_init(raw_full + i, 1)
- b_producer, b_consumer = PipelineTmaUmma.create(
- barrier_storage=storage.b_mbar.data_ptr(),
- num_stages=self.num_b_stages,
- producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
- consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1),
- tx_count=self.b_bytes,
- cta_layout_vmnk=cluster_layout_vmnk,
- ).make_participants()
- # A operand stages: the TMA warp acquires a stage (the MMA has released it), the
- # transform warps of both CTAs commit it to the leader, the leader's MMA consumes
- # and releases it.
- a_pipe = pipeline.PipelineAsyncUmma.create(
- barrier_storage=storage.a_mbar.data_ptr(),
- num_stages=self.num_stages,
- producer_group=pipeline.CooperativeGroup(
- pipeline.Agent.Thread, 2 * (self.tw // self.groups)
- ),
- consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1),
- cta_layout_vmnk=cluster_layout_vmnk,
- )
- a_consumer = a_pipe.make_participants()[1]
- acc_pipeline = PipelineUmmaAsync.create(
- barrier_storage=storage.acc_mbar.data_ptr(),
- num_stages=1,
- producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
- consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 2 * self.tw),
- cta_layout_vmnk=cluster_layout_vmnk,
- )
- tmem_alloc_barrier = pipeline.NamedBarrier(
- barrier_id=self.tmem_bar_id, num_threads=32 * (1 + self.tw)
- )
- tmem = utils.TmemAllocator(
- storage.tmem_holding_buf.ptr,
- barrier_for_retrieve=tmem_alloc_barrier,
- allocator_warp_id=0,
- is_two_cta=True,
- 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_staged.outer,
- byte_alignment=1024,
- swizzle=a_smem_staged.inner,
- )
- sB = smem.allocate_tensor(
- element_type=self.io_dtype,
- layout=b_smem_staged.outer,
- byte_alignment=1024,
- swizzle=b_smem_staged.inner,
- )
- sB1 = smem.allocate_tensor(
- element_type=self.io_dtype,
- layout=b1_smem_staged.outer,
- byte_alignment=1024,
- swizzle=b1_smem_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)
- )
- thr_mma = tiled_mma.get_slice(v)
- 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),
- )
- b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape)
- tBsB, tBgB = cpasync.tma_partition(
- tma_atom_b,
- cluster_coord[1],
- b_cta_layout,
- cute.group_modes(sB, 0, 3),
- cute.group_modes(tCgB, 0, 3),
- )
- gB1_nkl = cute.local_tile(
- mB1_nkl, cute.slice_(self.mma_tiler1, (0, None, None)), (None, None, None)
- )
- tCgB1 = tiled_mma1.get_slice(v).partition_B(gB1_nkl)
- tB1sB1, tB1gB1 = cpasync.tma_partition(
- tma_atom_b1,
- cluster_coord[1],
- b_cta_layout,
- cute.group_modes(sB1, 0, 3),
- cute.group_modes(tCgB1, 0, 3),
- )
- b_mcast = cpasync.create_tma_multicast_mask(
- cluster_layout_vmnk, cluster_coord, mcast_mode=1
- )
- tCrA = tiled_mma.make_fragment_A(sA)
- tCrB = tiled_mma.make_fragment_B(sB)
- tCrA1 = tiled_mma1.make_fragment_A(sA)
- tCrB1 = tiled_mma1.make_fragment_B(sB1)
- acc0_fake = tiled_mma.make_fragment_C(tiled_mma.partition_shape_C(self.mma_tiler[:2]))
- acc1_fake = tiled_mma1.make_fragment_C(tiled_mma1.partition_shape_C(self.mma_tiler1[:2]))
-
- pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
-
- # ------------------------------------------------------------ TMA warp (A, raw)
- if warp_idx == self.tma_warp_id:
- cute.arch.griddepcontrol_wait()
- tAgA_t = tAgA[(None, m_pair, None, 0)]
- a_state = pipeline.make_pipeline_state(
- pipeline.PipelineUserType.Producer, self.num_stages
- )
- for j in cutlass.range(self.kt, unroll=1):
- ai = a_state.index
- a_pipe.producer_acquire(a_state)
- with cute.arch.elect_one():
- cute.arch.mbarrier_arrive_and_expect_tx(raw_full + ai, self.a_bytes)
- cute.copy(
- tma_atom_a, tAgA_t[(None, kt0 + j)], tAsA[(None, ai)], tma_bar_ptr=raw_full + ai
- )
- a_state.advance()
-
- # ------------------------------------------------------------ TMA warp (B)
- if warp_idx == self.tmb_warp_id:
- for j in cutlass.range(self.kt, unroll=1):
- handle = b_producer.acquire_and_advance()
- cute.copy(
- tma_atom_b,
- tBgB[(None, 0, kt0 + j, 0)],
- tBsB[(None, handle.index)],
- tma_bar_ptr=handle.barrier,
- mcast_mask=b_mcast,
- )
- cute.copy(
- tma_atom_b1,
- tB1gB1[(None, 0, kt0 + j, 0)],
- tB1sB1[(None, handle.index)],
- tma_bar_ptr=handle.barrier,
- mcast_mask=b_mcast,
- )
- b_producer.tail()
-
- # ------------------------------------------------------------ MMA warp
- if warp_idx == self.mma_warp_id:
- tmem.wait_for_alloc()
- tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
- tCtAcc0 = cute.make_tensor(tmem_ptr, acc0_fake.layout)
- tCtAcc1 = cute.make_tensor(tmem_ptr + NH0, acc1_fake.layout)
- acc_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1)
- if is_leader:
- acc_pipeline.producer_acquire(acc_state)
- tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
- tiled_mma1.set(tcgen05.Field.ACCUMULATE, False)
- for j in cutlass.range(self.kt, unroll=1):
- b_handle = b_consumer.wait_and_advance()
- a_handle = a_consumer.wait_and_advance()
- num_kblocks = cute.size(tCrA, mode=[2])
- for kk in cutlass.range(num_kblocks, unroll_full=True):
- cute.gemm(
- tiled_mma,
- tCtAcc0,
- tCrA[(None, None, kk, a_handle.index)],
- tCrB[(None, None, kk, b_handle.index)],
- tCtAcc0,
- )
- cute.gemm(
- tiled_mma1,
- tCtAcc1,
- tCrA1[(None, None, kk, a_handle.index)],
- tCrB1[(None, None, kk, b_handle.index)],
- tCtAcc1,
- )
- tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
- tiled_mma1.set(tcgen05.Field.ACCUMULATE, True)
- b_handle.release()
- a_handle.release()
- acc_pipeline.producer_commit(acc_state)
- acc_state.advance()
- acc_pipeline.producer_tail(acc_state)
-
- # ------------------------------------------------------------ transform + epilogue
- if warp_idx < self.tw:
- cute.arch.griddepcontrol_wait()
- tmem.allocate(self.num_tmem_cols)
- # Before the transform loop: the MMA warp waits for this allocation, and the
- # MMA must run for the operand ring to drain.
- tmem.wait_for_alloc()
- tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
- # Transform mapping (128B-swizzled K-major 128 x 64 bf16 tile): within its group
- # of `gt` threads, thread t owns the 16-byte chunks at byte offsets 16 (t + gt j),
- # i.e. rows (t >> 3) + (gt / 8) j and the logical 8-column chunk
- # q = (t & 7) ^ ((t >> 3) & 7) of every row it owns. The tile is transformed in
- # place: same offsets in, same out.
- G = self.groups
- gt = 32 * self.tw // G
- nch = 1024 // gt # chunks a thread
- grp = tidx // gt
- tg = tidx % gt
- q = (tg & 7) ^ ((tg >> 3) & 7)
- row0 = m_pair * 256 + v * 128 + (tg >> 3)
- op_addr = cutlass.Int32(sA.iterator.toint()) + tg * 16
- op_stage = cute.size_in_bytes(
- self.io_dtype, cute.slice_(a_smem_staged, (None, None, None, 0))
- )
- rows64 = [cutlass.Int64(cutlass.min(row0 + (gt // 8) * jj, m - 1)) for jj in range(nch)]
- kpb = hid // BK # k blocks a stream
- S = self.num_stages
- # The group's next k block's scales are loaded one iteration ahead: rstd of the
- # thread's rows, 8 bf16 of the norm weight.
- rbuf = cute.make_rmem_tensor((nch,), cutlass.Float32)
- wbuf = cute.make_rmem_tensor((4,), cutlass.Int32)
- kfirst = kt0 + grp
- for jj in cutlass.range_constexpr(nch):
- rbuf[jj] = _ldg_f32(
- rstd_base + (rows64[jj] * hc + cutlass.Int64(kfirst // kpb)) * 4
- )
- w0 = _ldg_v4(w_base + cutlass.Int64(kfirst * BK + q * 8) * 2, cutlass.Int32(1))
- for e in cutlass.range_constexpr(4):
- wbuf[e] = w0[e]
- for j in cutlass.range(grp, self.kt, G, unroll=1):
- r = [rbuf[jj] for jj in range(nch)]
- ww = [wbuf[e] for e in range(4)]
- if j + G < self.kt:
- kn = kt0 + j + G
- s1 = kn // kpb
- for jj in cutlass.range_constexpr(nch):
- rbuf[jj] = _ldg_f32(rstd_base + (rows64[jj] * hc + cutlass.Int64(s1)) * 4)
- wn = _ldg_v4(w_base + cutlass.Int64(kn * BK + q * 8) * 2, cutlass.Int32(1))
- for e in cutlass.range_constexpr(4):
- wbuf[e] = wn[e]
- w1 = []
- for word in ww:
- w1 += [_add1(_bf16_lo(word)), _add1(_bf16_hi(word))]
- ai = j % S
- cute.arch.mbarrier_wait(raw_full + ai, (j // S) % 2)
- base = op_addr + ai * op_stage
- for jj in cutlass.range_constexpr(nch):
- words = _lds_v4(base + jj * (gt * 16))
- out = []
- for e in cutlass.range_constexpr(4):
- # (x * rstd) * (1 + w), each product rounded in fp32, then bf16:
- # `qwen4_exp_hc_inject_norm`'s normed, bitwise.
- lo = _mul(_mul(_bf16_lo(words[e]), r[jj]), w1[2 * e])
- hi = _mul(_mul(_bf16_hi(words[e]), r[jj]), w1[2 * e + 1])
- out.append(_bf16x2(lo, hi))
- _sts_v4(base + jj * (gt * 16), *out)
- cute.arch.fence_proxy("async.shared", space="cta")
- cute.arch.sync_warp()
- with cute.arch.elect_one():
- a_pipe.sync_object_full.arrive(ai, a_pipe.producer_mask)
-
- # Epilogue: warp w reads TMEM lanes 32 (w % 4).. of accumulator w // 4
- # (columns [0, NH0) and [NH0, NH0 + NH1)).
- lt = tidx % 128
- half = tidx // 128
- row = m_pair * 256 + v * 128 + lt
- row64 = cutlass.Int64(row)
- acc_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1)
- acc_pipeline.consumer_wait(acc_state)
- inv = cutlass.Float32(1.0 / hc)
- epi_tile = (128, 16)
- for hh in cutlass.range_constexpr(2):
- if half == hh:
- nhh = NH0 if hh == 0 else NH1
- acc_t = cute.make_tensor(
- tmem_ptr + hh * NH0, (acc0_fake if hh == 0 else acc1_fake).layout
- )
- copy_atom_t2r = utils.sm100.get_tmem_load_op(
- (128, nhh, BK),
- utils.LayoutEnum.ROW_MAJOR,
- self.io_dtype,
- self.acc_dtype,
- epi_tile,
- True,
- )
- tAcc_epi = cute.flat_divide(acc_t[((None, None), 0, 0)], epi_tile)
- tiled_copy_t2r = tcgen05.make_tmem_copy(
- copy_atom_t2r, tAcc_epi[(None, None, 0, 0)]
- )
- thr_copy_t2r = tiled_copy_t2r.get_slice(lt)
- tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi)
- tTR_cE = thr_copy_t2r.partition_D(cute.make_identity_tensor(epi_tile))
- frag = cute.make_rmem_tensor(tTR_cE.shape, self.acc_dtype)
- for c in cutlass.range_constexpr(nhh // 16):
- col = hh * NH0 + c * 16
- if cutlass.const_expr(col < rank + hc):
- cute.copy(tiled_copy_t2r, tTR_tAcc[(None, None, None, 0, c)], frag)
- if row < m:
- if cutlass.const_expr(split > 1):
- pbase = (
- out_base
- + ((cutlass.Int64(sp) * m + row64) * NCOL + col) * 4
- )
- for qq in cutlass.range_constexpr(4):
- _st_global_v4(
- pbase + qq * 16,
- *[_bits(frag[4 * qq + e]) for e in range(4)],
- )
- elif cutlass.const_expr(col + 16 <= rank):
- vals = []
- for e in cutlass.range_constexpr(16):
- z = frag[e] * inv
- vals.append(z * _sigmoid_fast(z))
- words = [
- _bf16x2(vals[2 * e], vals[2 * e + 1]) for e in range(8)
- ]
- obase = out_base + (row64 * rank + col) * 2
- _st_global_v4(obase, *words[0:4])
- _st_global_v4(obase + 16, *words[4:8])
- else:
- # The chunk holding the injection logits (rank % 16 == 0).
- vals = []
- for e in cutlass.range_constexpr(hc):
- z = frag[e] * inv
- vals.append(cutlass.Float32(2.0) * _sigmoid_fast(z))
- _st_global_v2(
- inj_base + row64 * (hc * 2),
- _bf16x2(vals[0], vals[1]),
- _bf16x2(vals[2], vals[3]),
- )
- cute.arch.sync_warp()
- with cute.arch.elect_one():
- acc_pipeline.consumer_release(acc_state)
- cute.arch.barrier(barrier_id=self.epi_bar_id, number_of_threads=32 * self.tw)
- tmem.relinquish_alloc_permit()
- tmem.free(tmem_ptr)
-
-
-@triton.jit
-def _combine_act_kernel(
- part_ptr,
- low_ptr,
- inj_ptr,
- T,
- SPLIT: tl.constexpr,
- RANK: tl.constexpr,
- HC: tl.constexpr,
- NCOL: tl.constexpr,
- BLOCK: tl.constexpr,
-):
- # The K-split partial logits summed in split order (fp32), then the activations of
- # the split-1 epilogue (`tl.sigmoid` here, the ex2 / rcp form there).
- row = tl.program_id(0).to(tl.int64)
- k = tl.arange(0, BLOCK)
- mask = k < RANK + HC
- z = tl.load(part_ptr + row * NCOL + k, mask=mask, other=0.0)
- for sidx in tl.static_range(1, SPLIT):
- z += tl.load(part_ptr + (sidx * T + row) * NCOL + k, mask=mask, other=0.0)
- z = z * (1.0 / HC)
- tl.store(
- low_ptr + row * RANK + k, (z * tl.sigmoid(z)).to(low_ptr.dtype.element_ty), mask=k < RANK
- )
- ki = k - RANK
- tl.store(
- inj_ptr + row * HC + ki,
- (2.0 * tl.sigmoid(z)).to(inj_ptr.dtype.element_ty),
- mask=(ki >= 0) & (ki < HC),
- )
-
-
-_COMPILED: dict = {}
-_SMS: dict = {}
-
-
-def _launcher(hc: int, hidden: int, rank: int, split: int, ptrs):
- """Compile once per static configuration; later calls reuse the TVM-FFI entry,
- which launches on the caller's current stream (graph-capturable). `ptrs` is
- `(hc_out, norm_weight, rstd tensor, first, out, inject)`.
- """
- key = (hc, hidden, rank, split)
- fn = _COMPILED.get(key)
- if fn is None:
- kern = HcFirstRescaleKernel(hc, hidden, rank, split)
- gm = cute.AddressSpace.gmem
- cptrs = [
- make_ptr(cutlass.BFloat16, ptrs[0], gm, assumed_align=16),
- make_ptr(cutlass.BFloat16, ptrs[1], gm, assumed_align=16),
- from_dlpack(ptrs[2], assumed_align=16).mark_layout_dynamic(leading_dim=1),
- make_ptr(cutlass.BFloat16, ptrs[3], gm, assumed_align=16),
- make_ptr(cutlass.BFloat16, ptrs[4], gm, assumed_align=16),
- make_ptr(cutlass.BFloat16, ptrs[5], gm, assumed_align=8),
- ]
- fn = cute.compile(
- kern,
- 128,
- *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 _pick_split(rows: int, k_blocks: int, device: torch.device) -> int:
- """CTA pairs = ceil(rows / 256) * split: split K in two while the grid stays within
- one wave of SMs (and K divides), up to `_MAX_SPLIT`.
- """
- sms = _SMS.get(device.index)
- if sms is None:
- sms = _SMS[device.index] = torch.cuda.get_device_properties(device).multi_processor_count
- pairs = -(-rows // 256)
- split = 1
- while split < _MAX_SPLIT and 2 * pairs * split * 2 <= sms and k_blocks % (2 * split) == 0:
- split *= 2
- return split
-
-
-def hc_first_rescale(
- hc_out: torch.Tensor,
- rstd: torch.Tensor,
- norm_weight: torch.Tensor,
- first_t: torch.Tensor,
- hc_count: int,
- rank: int,
- with_inject: bool,
-):
- """`(low [T, rank], inject [T, C] or None)`; `first_t` the contiguous padded
- `[N, C*H]` first weight (N >= 352).
- """
- rows, width = hc_out.shape
- hidden = width // hc_count
- dev = hc_out.device
- low = torch.empty((rows, rank), dtype=hc_out.dtype, device=dev)
- # Written in every case (the kernel stores the injection chunk); dropped without
- # `with_inject`.
- inject = torch.empty((rows, hc_count), dtype=hc_out.dtype, device=dev)
- if rows > 0:
- split = _pick_split(rows, width // BK, dev)
- if split == 1:
- out = low
- else:
- out = torch.empty((split, rows, NCOL), dtype=torch.float32, device=dev)
- ptrs = (
- hc_out.data_ptr(),
- norm_weight.data_ptr(),
- rstd,
- first_t.data_ptr(),
- out.data_ptr(),
- inject.data_ptr(),
- )
- _launcher(hc_count, hidden, rank, split, ptrs)(rows, *ptrs)
- if split > 1:
- _combine_act_kernel[(rows,)](
- out,
- low,
- inject,
- rows,
- SPLIT=split,
- RANK=rank,
- HC=hc_count,
- NCOL=NCOL,
- BLOCK=triton.next_power_of_2(rank + hc_count),
- num_warps=4,
- )
- return low, (inject if with_inject else None)
-
-
-class Qwen4ExpHcFirstRescaleTc(OpWrapper):
- """With `C = hc_count`, `H = width // C`, `first` the `[C*H, N]` transposed view
- of the padded `[N, C*H]` `[down; block_inject; 0]` weight (N >= 352)::
-
- normed = bf16(hc_out[t, s * H + i] * rstd[t, s] * (1 + norm_weight[s * H + i]))
- logits = normed @ first[:, : rank + C] # fp32
- low = silu(logits[:, :rank] / C) # [T, rank]
- inject = 2 * sigmoid(logits[:, rank : rank + C] / C) # [T, C]
-
- `normed` is bitwise `qwen4_exp_hc_inject_norm`'s (the same fp32 products, one
- rounding), formed in shared memory and never written. Returns fresh contiguous
- `(low, inject)` in `hc_out`'s dtype; `inject` is `None` exactly when
- `with_inject` is false (the stack-final mixer; `low` is unchanged).
-
- Fusion boundary: the grouped norm's scaling, the first projection and both
- activations -- `cublas_mm(normed, first[:, :352])` plus
- `qwen4_exp_hc_lowrank_act` fed by a site that stored `normed`.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: rows, inject (True). The checkpoint's geometry throughout: hc_count 4,
- # hidden 2560, rank 320, the padded 384-row first weight, bf16.
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="prefill chunk 16384, one K slice a CTA pair",
- spec=dict(rows=16384),
- ),
- Cell(why="prefill 8192 rows, K split in two", spec=dict(rows=8192)),
- Cell(why="prefill 2048 rows, K split in eight", spec=dict(rows=2048)),
- Cell(
- why="stack-final mixer (no injection), prefill chunk 16384",
- spec=dict(rows=16384, inject=False),
- ),
- Cell(why="ragged tail, 16320 + 70 rows (mixed step)", spec=dict(rows=16390)),
- Cell(why="five rows, less than one tile", spec=dict(rows=5)),
- )
-
- note = """
- Tier 6 (CuTe DSL written for this target): a tcgen05 GEMM whose A operand is
- produced by warps of the same CTA from a second tensor; no lower tier has an
- operand prologue (cuBLAS needs `normed` materialized). Metadata consumed:
- none. Call sites: `PrefillTarget`'s trunk mixers at or above
- `_HC_RSTD_MIN_ROWS` rows (every attention / GDN, MoE and stack-final site),
- fed by `qwen4_exp_hc_inject_rstd`.
-
- Preconditions: bf16; `hc_out` `[T, C*H]` contiguous; rstd `[T, C]` fp32
- contiguous; `norm_weight` `[C*H]` contiguous; `first` the transposed view of
- a contiguous `[N, C*H]` weight, N >= 352; 16-byte aligned pointers;
- `hc_count == 4`, `H % 64 == 0`, `rank % 16 == 0`, `rank + C <= 352`.
- Compiles once per `(C, H, rank, split)` on first call; the row count is a
- runtime argument. Graph-capturable (current stream, no host sync; the split
- is a host function of the row count).
-
- Launch: one CTA pair a 256-row tile, the whole K (split 1), or, when that
- leaves more than half the SMs idle, `split` pairs a tile each over `1/split`
- of K (2048 rows: split 8, 8192: split 2, 16384: split 1 on 152 SMs) with a
- Triton combine. Launched with programmatic dependent launch
- (TRTLLM_ENABLE_PDL=0 turns it off): the weight stream starts before
- `griddepcontrol.wait`. GB300, 16384 rows: 82 us, against cuBLAS 77.6 +
- lowrank_act 9.3 on a stored `normed` -- the site's saving is the `normed`
- write and re-read this removes upstream (`qwen4_exp_hc_inject_rstd`).
-
- Numerics: `low` / `inject` are one fp32 activation of fp32 logits rounded
- once (split 1: `ex2.approx` / `rcp.approx` sigmoid; split > 1: the partials
- summed in split order, `tl.sigmoid`). Measured against the cuBLAS +
- `qwen4_exp_hc_lowrank_act` path on the same `normed`: split 1 bitwise at
- 2048-16390 rows (where cuBLAS does not split K itself); otherwise the last
- bit differs on about 0.5% of `low`, inside the same band.
-
- Certified: bf16 at the checkpoint's 4 x 2560, rank 320: 16384 (split 1),
- 8192 (split 2), 2048 (split 8), 5 rows, a partial last tile (16390), the
- stack-final form; graph replay. Against the float64 product of the
- once-rounded `normed`, gate band (`rtol = atol / max|ref| = eps`).
-
- Not certified: zero rows, hc_count other than 4, hidden other than 2560,
- rank other than 320.
- """
-
- def __call__(
- self,
- hc_out: torch.Tensor,
- rstd: torch.Tensor,
- norm_weight: torch.Tensor,
- first: torch.Tensor,
- hc_count: int,
- rank: int,
- with_inject: bool,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
- return hc_first_rescale(hc_out, rstd, norm_weight, first.t(), hc_count, rank, with_inject)
-
- def reference(
- self,
- hc_out: torch.Tensor,
- rstd: torch.Tensor,
- norm_weight: torch.Tensor,
- first: torch.Tensor,
- hc_count: int,
- rank: int,
- with_inject: bool,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
- """float64: `normed` rounded once to `hc_out`'s dtype (the op's operand),
- then the projection and both activations, each output rounded once.
- """
- rows, width = hc_out.shape
- hidden = width // hc_count
- x = hc_out.double().unflatten(-1, (hc_count, hidden)) * rstd.double().unsqueeze(-1)
- normed = (x.flatten(-2) * (1.0 + norm_weight.double())).to(hc_out.dtype)
- logits = F.linear(normed.double(), first[:, : rank + hc_count].t().double())
- low = F.silu(logits[:, :rank] / hc_count).to(hc_out.dtype)
- inject = (2.0 * torch.sigmoid(logits[:, rank:] / hc_count)).to(hc_out.dtype)
- return low, (inject if with_inject else None)
-
- def is_valid(
- self,
- hc_out: torch.Tensor,
- rstd: torch.Tensor,
- norm_weight: torch.Tensor,
- first: torch.Tensor,
- hc_count: int,
- rank: int,
- with_inject: bool,
- ) -> None:
- rows, width = hc_out.shape
- assert hc_out.dtype == norm_weight.dtype == first.dtype == torch.bfloat16, "bf16"
- assert hc_out.stride() == (width, 1), "hc_out: contiguous rows"
- assert rstd.dtype == torch.float32 and rstd.shape == (rows, hc_count), "rstd [T, C] fp32"
- assert rstd.is_contiguous(), "rstd contiguous"
- assert norm_weight.shape == (width,) and norm_weight.is_contiguous(), "norm_weight [C*H]"
- assert first.dim() == 2 and first.shape[0] == width and first.shape[1] >= NCOL, (
- f"first [C*H, N >= {NCOL}]"
- )
- assert first.stride() == (1, width), "first must be the transposed view of [N, C*H]"
- # The epilogue stores the injection chunk as two bf16x2 words (C == 4) and
- # reads its 16-column chunks wholly as `low` or as injection.
- assert hc_count == 4, "hc_count 4"
- assert (width // hc_count) % BK == 0, f"hidden must be a multiple of {BK}"
- assert rank % 16 == 0 and rank + hc_count <= NCOL, "rank % 16 == 0, rank + C <= 352"
- assert isinstance(with_inject, bool)
- for t in (hc_out, norm_weight, first, rstd):
- assert t.data_ptr() % 16 == 0, "16-byte aligned operands"
-
-
-qwen4_exp_hc_first_rescale_tc = Qwen4ExpHcFirstRescaleTc()
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
deleted file mode 100644
index 9ddaa39c7e7f..000000000000
--- 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
+++ /dev/null
@@ -1,1058 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Hyper-connection post-block injection fused with the next site's grouped
-gemma RMS norm: one Triton row kernel writes the updated multi-stream residual
-and its normalized copy in the same pass. After a MoE block it also forms the
-block output itself, from the routed experts' output, the shared expert's
-output and its gate logit (the shared-expert combine).
-
-Tier 5 (Triton written for this target). No ``torch.ops.trtllm.*`` op or
-FlashInfer callable accumulates a block output into ``hc_count`` streams with
-per-stream gates and normalizes each stream with its own RMS in one launch; the
-two-launch composition (``qwen4_exp_post_inject`` + four
-``flashinfer_gemma_rmsnorm`` calls and a concat) is what this replaces, and
-with ``shared`` / ``gate_logits`` also the ``fused_sigmoid_gate_mul_add``
-launch before it -- and with ``topk_weights`` / ``topk_rows`` the MoE runner's
-finalize launch as well.
-
-The decode MoE sites (the finalize fold below `_FEW_WARPS_MIN_ROWS` rows) run a
-CuTe DSL kernel instead (Tier 6, written for this target), bitwise the Triton
-kernel's outputs there.
-"""
-
-from __future__ import annotations
-
-import struct
-from typing import Optional, Tuple
-
-import cutlass
-import cutlass.cute as cute
-import torch
-import triton
-import triton.language as tl
-from cutlass._mlir.dialects import llvm
-from cutlass.cute.runtime import from_dlpack
-from cutlass.cutlass_dsl import T, dsl_user_op
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
-from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
-
-#: Trigger dependents at kernel start (`griddepcontrol.launch_dependents`); a no-op
-#: unless the next kernel is launched with programmatic dependent launch.
-_PDL_TRIGGER = get_env_enable_pdl()
-
-#: Rows from which a wide (`hidden >= 2048`) launch runs two warps a program
-#: instead of eight: prefill-sized grids are bandwidth-bound and run faster
-#: with fewer, longer threads; decode-sized ones keep eight.
-_FEW_WARPS_MIN_ROWS = 512
-
-
-@triton.jit
-def _hc_inject_norm_kernel(
- hc_ptr,
- block_ptr,
- inject_ptr,
- shared_ptr,
- gate_ptr,
- topk_w_ptr,
- topk_row_ptr,
- weight_ptr,
- hc_out_ptr,
- normed_ptr,
- stride_hc,
- stride_block,
- stride_inject,
- stride_shared,
- stride_gate,
- stride_topk_w,
- stride_topk_row,
- stride_out,
- stride_normed,
- eps,
- HC_COUNT: tl.constexpr,
- HIDDEN: tl.constexpr,
- BLOCK: tl.constexpr,
- HAS_INJECT: tl.constexpr,
- HAS_SHARED: tl.constexpr,
- TOPK: tl.constexpr,
- STREAMS_PER_PROGRAM: tl.constexpr,
- PDL_TRIGGER: tl.constexpr,
-):
- if PDL_TRIGGER:
- # Release a programmatic-dependent-launch consumer (the decode mixer's
- # first projection) at once: it streams its static weights while this
- # kernel runs and waits for the grid to finish before reading `normed`.
- tl.extra.cuda.gdc_launch_dependents()
- # One program per `STREAMS_PER_PROGRAM` streams of a row; a stream's `HIDDEN`
- # columns are one RMS group. A row's programs are adjacent in launch order, so
- # the row's block output (and the combine's operands) come from HBM once and
- # from L2 for the others; with several streams a program forms it once for them.
- pid = tl.program_id(0).to(tl.int64)
- row = pid // (HC_COUNT // STREAMS_PER_PROGRAM)
- first_stream = ((pid % (HC_COUNT // STREAMS_PER_PROGRAM)) * STREAMS_PER_PROGRAM).to(tl.int32)
- cols = tl.arange(0, BLOCK)
- mask = cols < HIDDEN
- if HAS_INJECT:
- if TOPK > 0:
- # The routed experts' finalize: this row's TOPK permuted FC2 rows weighted
- # by their routing weights, fp32 FMAs in slot order from zero (a slot mapped
- # to -1 is skipped), rounded once to the rows' dtype.
- b = tl.zeros([BLOCK], dtype=tl.float32)
- for k in tl.static_range(TOPK):
- prow = tl.load(topk_row_ptr + row * stride_topk_row + k)
- wk = tl.load(topk_w_ptr + row * stride_topk_w + k).to(tl.float32)
- r = tl.load(
- block_ptr + tl.maximum(prow, 0).to(tl.int64) * stride_block + cols,
- mask=mask,
- other=0.0,
- ).to(tl.float32)
- b = tl.where(prow >= 0, tl.fma(r, wk, b), b)
- b = b.to(block_ptr.dtype.element_ty).to(tl.float32)
- else:
- b = tl.load(block_ptr + row * stride_block + cols, mask=mask, other=0.0).to(tl.float32)
- if HAS_SHARED:
- # The MoE block output, `routed + sigmoid(gate) * shared` in fp32, rounded
- # once to the routed output's dtype -- the value that block used to store.
- sg = tl.sigmoid(tl.load(gate_ptr + row * stride_gate).to(tl.float32))
- sh = tl.load(shared_ptr + row * stride_shared + cols, mask=mask, other=0.0)
- b = (b + sg * sh.to(tl.float32)).to(block_ptr.dtype.element_ty).to(tl.float32)
- for j in tl.static_range(STREAMS_PER_PROGRAM):
- stream = first_stream + j
- offs = stream * HIDDEN + cols
- x = tl.load(hc_ptr + row * stride_hc + offs, mask=mask, other=0.0).to(tl.float32)
- if HAS_INJECT:
- g = tl.load(inject_ptr + row * stride_inject + stream).to(tl.float32)
- # Rounded once on the store, like `addcmul`; the norm reads the rounded
- # residual, which is the value every later site sees.
- xr = (x + b * g).to(hc_out_ptr.dtype.element_ty)
- tl.store(hc_out_ptr + row * stride_out + offs, xr, mask=mask)
- x = xr.to(tl.float32)
- rstd = tl.rsqrt(tl.sum(x * x, axis=0) / HIDDEN + eps)
- w = tl.load(weight_ptr + offs, mask=mask, other=0.0).to(tl.float32)
- y = x * rstd * (1.0 + w)
- tl.store(
- normed_ptr + row * stride_normed + offs, y.to(normed_ptr.dtype.element_ty), mask=mask
- )
-
-
-# --------------------------------------------------------------------------------------
-# The decode MoE sites in CuTe DSL. There the Triton kernel runs two programs a row, each
-# gathering all `k` permuted FC2 rows and normalizing its two streams one after the other
-# (6.5 us at 128 rows in the con128 decode step on GB300). Here the row's two CTAs are a
-# cluster: CTA j forms the block output's columns [j * H / 2, (j + 1) * H / 2) -- the
-# finalize's slot-order FMA chain and the combine, the Triton kernel's instructions -- and
-# stores them into both CTAs' shared memory; after a cluster barrier each CTA reads the
-# whole block row locally and accumulates and normalizes its two streams together. A CTA
-# is 256 threads in the Triton kernel's layout at a 4096 block and eight warps (thread t:
-# columns [8t, 8t + 8) and, below the hidden, [2048 + 8t, 2048 + 8t + 8)), and the sum of
-# squares is its reduction tree: `hc_out` and `normed` are bitwise the Triton kernel's.
-# --------------------------------------------------------------------------------------
-
-_CUTE_THREADS = 256
-_CUTE_WARPS = _CUTE_THREADS // 32
-_CUTE_VEC = 8
-#: The Triton kernel's block at these hiddens is twice this: a thread's second chunk starts here.
-_CUTE_SPLIT = _CUTE_THREADS * _CUTE_VEC
-#: Streams a CTA, as the Triton fold's program; a row's `hc_count // 2` CTAs are one cluster.
-_CUTE_SPP = 2
-
-
-def _iv(x, loc, ip):
- return x.ir_value(loc=loc, ip=ip)
-
-
-def _asm(ret, args, text, constraints, loc, ip, side_effects=False):
- return llvm.inline_asm(
- ret,
- [_iv(a, loc, ip) for a in args],
- text,
- constraints,
- has_side_effects=side_effects,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
-
-
-def _asm4(args, text, constraints, loc, ip, side_effects=False):
- ty = T.i32()
- res = _asm(
- llvm.StructType.get_literal([ty] * 4), args, text, constraints, loc, ip, side_effects
- )
- return [cutlass.Int32(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(4)]
-
-
-@dsl_user_op
-def _ldg_v4(addr, pred, *, loc=None, ip=None):
- """16 bytes, or zeros where `pred` is 0 (nothing read)."""
- return _asm4(
- [cutlass.Int64(addr), cutlass.Int32(pred)],
- "{ .reg .pred p; setp.ne.b32 p, $5, 0; mov.b32 $0, 0; mov.b32 $1, 0; mov.b32 $2, 0; "
- "mov.b32 $3, 0; @p ld.global.v4.b32 {$0, $1, $2, $3}, [$4]; }",
- "=r,=r,=r,=r,l,r",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _stg_v4(addr, pred, a, b, c, d, *, loc=None, ip=None):
- """16 bytes where `pred` is not 0."""
- _asm(
- T.i32(),
- [cutlass.Int64(addr), cutlass.Int32(pred)] + [cutlass.Int32(v) for v in (a, b, c, d)],
- "{ .reg .pred p; setp.ne.b32 p, $2, 0; @p st.global.v4.b32 [$1], {$3, $4, $5, $6}; "
- "mov.u32 $0, 0; }",
- "=r,l,r,r,r,r,r",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _mapa(addr, rank, *, loc=None, ip=None):
- """This CTA's shared-memory address `addr` in cluster CTA `rank`'s window."""
- return cutlass.Int32(
- _asm(
- T.i32(),
- [cutlass.Int32(addr), cutlass.Int32(rank)],
- "mapa.shared::cluster.u32 $0, $1, $2;",
- "=r,r,r",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _st_dsmem_v4(addr, pred, a, b, c, d, *, loc=None, ip=None):
- """16 bytes into a cluster CTA's shared memory where `pred` is not 0."""
- _asm(
- T.i32(),
- [cutlass.Int32(addr), cutlass.Int32(pred)] + [cutlass.Int32(v) for v in (a, b, c, d)],
- "{ .reg .pred p; setp.ne.b32 p, $2, 0; @p st.shared::cluster.v4.b32 [$1], "
- "{$3, $4, $5, $6}; mov.u32 $0, 0; }",
- "=r,r,r,r,r,r,r",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _lds_v4(addr, pred, *, loc=None, ip=None):
- """16 bytes of this CTA's shared memory, or zeros where `pred` is 0."""
- return _asm4(
- [cutlass.Int32(addr), cutlass.Int32(pred)],
- "{ .reg .pred p; setp.ne.b32 p, $5, 0; mov.b32 $0, 0; mov.b32 $1, 0; mov.b32 $2, 0; "
- "mov.b32 $3, 0; @p ld.shared.v4.b32 {$0, $1, $2, $3}, [$4]; }",
- "=r,=r,=r,=r,r,r",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _bf16_lo(w, *, loc=None, ip=None):
- """The low bf16 of a word as fp32 (exact)."""
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Int32(w)],
- "{ .reg .b32 t; shl.b32 t, $1, 16; mov.b32 $0, t; }",
- "=f,r",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _bf16_hi(w, *, loc=None, ip=None):
- """The high bf16 of a word as fp32 (exact)."""
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Int32(w)],
- "{ .reg .b32 t; and.b32 t, $1, 0xFFFF0000; mov.b32 $0, t; }",
- "=f,r",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _bf16x2(lo, hi, *, loc=None, ip=None):
- """Two fp32 values rounded (RN) into one bf16x2 word, `lo` low."""
- return cutlass.Int32(
- _asm(
- T.i32(),
- [cutlass.Float32(hi), cutlass.Float32(lo)],
- "cvt.rn.bf16x2.f32 $0, $1, $2;",
- "=r,f,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _round_bf16(x, *, loc=None, ip=None):
- """fp32 -> bf16 (RN) -> fp32."""
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Float32(x)],
- "{ .reg .b16 h; cvt.rn.bf16.f32 h, $1; cvt.f32.bf16 $0, h; }",
- "=f,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _fma(a, b, c, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)],
- "fma.rn.f32 $0, $1, $2, $3;",
- "=f,f,f,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _fma_if(slot, a, b, c, *, loc=None, ip=None):
- """`slot >= 0 ? fma(a, b, c) : c` (the finalize skips a slot mapped to -1)."""
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Int32(slot), cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)],
- "{ .reg .pred p; .reg .f32 t; fma.rn.f32 t, $2, $3, $4; setp.ge.s32 p, $1, 0; "
- "selp.f32 $0, t, $4, p; }",
- "=f,r,f,f,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _mul(a, b, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Float32(a), cutlass.Float32(b)],
- "mul.rn.f32 $0, $1, $2;",
- "=f,f,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _add(a, b, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Float32(a), cutlass.Float32(b)],
- "add.rn.f32 $0, $1, $2;",
- "=f,f,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _sigmoid(x, *, loc=None, ip=None):
- """`tl.sigmoid` as the Triton kernel compiles it: `mul` by -log2(e), `ex2.approx`, `add`
- 1, `div.full`."""
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Float32(x)],
- "{ .reg .f32 t; mul.f32 t, $1, 0fBFB8AA3B; ex2.approx.f32 t, t; "
- "add.f32 t, t, 0f3F800000; div.full.f32 $0, 0f3F800000, t; }",
- "=f,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _rstd(ss, eps, hidden_hex: str, *, loc=None, ip=None):
- """`rsqrt(ss / hidden + eps)` as the Triton kernel's `div.full`, `add`, `rsqrt.approx.ftz`:
- the `add` has no rounding mode there, so ptxas fuses it with the division (a multiply by
- the reciprocal) into one FFMA -- here the same."""
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Float32(ss), cutlass.Float32(eps)],
- "{ .reg .f32 t; div.full.f32 t, $1, " + hidden_hex + "; add.f32 t, t, $2; "
- "rsqrt.approx.ftz.f32 $0, t; }",
- "=f,f,f",
- loc,
- ip,
- )
- )
-
-
-def _vals(words):
- """Four bf16x2 words -> their eight values as fp32, in column order."""
- out = []
- for w in words:
- out += [_bf16_lo(w), _bf16_hi(w)]
- return out
-
-
-class _FoldNormKernel:
- """CTA (j, row): the row's streams 2j and 2j + 1; the row's CTAs one cluster."""
-
- def __init__(self, hc_count: int, hidden: int, top_k: int, pdl_trigger: bool):
- self.hc, self.hidden, self.topk, self.pdl = hc_count, hidden, top_k, pdl_trigger
- self.ctas = hc_count // _CUTE_SPP
- # The block-output columns a CTA forms, one 16-byte vector a thread.
- self.per = hidden // self.ctas
- self.n_fold = self.per // _CUTE_VEC
- # Threads with a second chunk (columns [2048, hidden)).
- self.n_b = (hidden - _CUTE_SPLIT) // _CUTE_VEC
- self.hidden_hex = "0f%08X" % struct.unpack("> off))
- if tidx % 32 == 0:
- red[s, warp] = v
- cute.arch.sync_threads()
- for s in cutlass.range_constexpr(_CUTE_SPP):
- stream = rank * _CUTE_SPP + s
- v = red[s, tidx % _CUTE_WARPS]
- for off in cutlass.range_constexpr(3):
- v = _add(v, cute.arch.shuffle_sync_bfly(v, offset=4 >> off))
- rstd = _rstd(v, eps, self.hidden_hex)
- ra, rb = xrs[s]
- wa_v, wb_v = _vals(wa[s]), _vals(wb[s])
- ya = [
- _mul(_mul(rstd, ra[e]), _add(wa_v[e], cutlass.Float32(1.0)))
- for e in range(_CUTE_VEC)
- ]
- yb = [
- _mul(_mul(rstd, rb[e]), _add(wb_v[e], cutlass.Float32(1.0)))
- for e in range(_CUTE_VEC)
- ]
- nrow = nrm_base + (row64 * width + stream * self.hidden) * 2
- _stg_v4(nrow + c_a * 2, one, *[_bf16x2(ya[2 * i], ya[2 * i + 1]) for i in range(4)])
- _stg_v4(nrow + c_b * 2, has_b, *[_bf16x2(yb[2 * i], yb[2 * i + 1]) for i in range(4)])
-
-
-_FOLD_NORM: dict = {}
-
-
-def _fold_cute_ok(
- hyper_input,
- block_out,
- inject,
- norm_weight,
- hc_count,
- shared,
- gate_logits,
- topk_weights,
- topk_rows,
-) -> bool:
- """The decode MoE site this kernel certifies (the Triton kernel's eight-warp form with
- 16-byte vectors); anything else runs the Triton kernel."""
- rows, width = hyper_input.shape
- hidden = width // hc_count
- bf = torch.bfloat16
- return (
- 0 < rows < _FEW_WARPS_MIN_ROWS
- and hc_count == 4
- and _CUTE_SPLIT < hidden <= 2 * _CUTE_SPLIT
- and hidden % (_CUTE_SPP * _CUTE_VEC) == 0
- and hyper_input.dtype == block_out.dtype == shared.dtype == bf
- and inject.dtype == gate_logits.dtype == topk_weights.dtype == norm_weight.dtype == bf
- and topk_rows.dtype == torch.int32
- and all(
- t.data_ptr() % 16 == 0 and t.stride(0) % 16 == 0
- for t in (hyper_input, block_out, shared)
- )
- and norm_weight.data_ptr() % 16 == 0
- )
-
-
-def _fold_cute(
- hyper_input,
- block_out,
- inject,
- norm_weight,
- eps,
- hc_count,
- shared,
- gate_logits,
- topk_weights,
- topk_rows,
- hc_out,
- normed,
-) -> None:
- rows, width = hyper_input.shape
- hidden = width // hc_count
- top_k = topk_rows.shape[1]
- ptrs = (
- hyper_input,
- block_out.data_ptr(),
- inject.data_ptr(),
- shared.data_ptr(),
- gate_logits.data_ptr(),
- topk_weights.data_ptr(),
- topk_rows.data_ptr(),
- norm_weight.data_ptr(),
- hc_out.data_ptr(),
- normed.data_ptr(),
- )
- strides = (
- hyper_input.stride(0),
- block_out.stride(0),
- inject.stride(0),
- shared.stride(0),
- gate_logits.stride(0),
- topk_weights.stride(0),
- topk_rows.stride(0),
- )
- key = (hc_count, hidden, top_k, _PDL_TRIGGER)
- fn = _FOLD_NORM.get(key)
- if fn is None:
- gm = cute.AddressSpace.gmem
- bf = cutlass.BFloat16
- cargs = (
- from_dlpack(hyper_input, assumed_align=16).mark_layout_dynamic(leading_dim=1),
- make_ptr(bf, ptrs[1], gm, assumed_align=16),
- make_ptr(bf, ptrs[2], gm, assumed_align=2),
- make_ptr(bf, ptrs[3], gm, assumed_align=16),
- make_ptr(bf, ptrs[4], gm, assumed_align=2),
- make_ptr(bf, ptrs[5], gm, assumed_align=2),
- make_ptr(cutlass.Int32, ptrs[6], gm, assumed_align=4),
- make_ptr(bf, ptrs[7], gm, assumed_align=16),
- make_ptr(bf, ptrs[8], gm, assumed_align=16),
- make_ptr(bf, ptrs[9], gm, assumed_align=16),
- )
- fn = cute.compile(
- _FoldNormKernel(hc_count, hidden, top_k, _PDL_TRIGGER),
- *cargs,
- *(cutlass.Int32(1),) * 8,
- cutlass.Float32(1e-6),
- cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
- options="--opt-level 2 --enable-tvm-ffi",
- )
- _FOLD_NORM[key] = fn
- fn(*ptrs, rows, *strides, eps)
-
-
-class Qwen4ExpHcInjectNorm(OpWrapper):
- """With `C = hc_count`, `H = width // C`, every tensor 2-D `[T, ...]`:
-
- hc_out[t, s*H + i] = hyper_input[t, s*H + i] + block_out[t, i] * inject[t, s]
- normed[t, s*H + i] = hc_out[t, s*H + i] * rstd(t, s) * (1 + norm_weight[s*H + i])
- rstd(t, s) = 1 / sqrt(mean_i(hc_out[t, s*H + i]^2) + eps)
-
- `hc_out` is rounded once to the activation dtype and `normed` is computed
- from that rounded value (fp32 inside, one rounding on the store). Returns
- `(hc_out, normed)`, both fresh `[T, C*H]` tensors in `hyper_input`'s dtype.
-
- `block_out=None, inject=None` is the norm alone: `hc_out` *is*
- `hyper_input` (not a copy) and only `normed` is written -- the stack's first
- site and the site after the PLE add.
-
- With `shared` and `gate_logits` (after a MoE block), `block_out` is the
- routed experts' output and the block output injected above is
-
- block[t, i] = block_out[t, i] + sigmoid(gate_logits[t]) * shared[t, i]
-
- in fp32, rounded once to `block_out`'s dtype before the injection reads it:
- bitwise what `fused_sigmoid_gate_mul_add` stores. `gate_logits` is one
- pre-sigmoid scalar per row, `[T, 1]` or `[T]`; none of the inputs is written.
-
- With `topk_weights` / `topk_rows` as well (the MoE runner left unfinalized),
- `block_out` is its permuted FC2 output `[P, H]` and the routed row is the
- runner's finalize, formed here: `routed[t, i] = sum_k topk_weights[t, k] *
- block_out[topk_rows[t, k], i]` as fp32 FMAs in slot order from zero (a slot
- whose row is -1 is skipped), rounded once to `block_out`'s dtype -- bitwise
- what the runner's finalize stores. `topk_rows` is the runner's
- expanded->permuted index map `[T, k]` int32, `topk_weights` its `[T, k]`
- routing weights. The fold needs the combine (`shared`, `gate_logits`).
-
- Fusion boundary: the routed experts' finalize (optional), the
- shared-expert combine (optional), the stream-broadcast
- gated accumulate (the old `qwen4_exp_post_inject`) and the grouped gemma
- RMS norm of the next site (the old `_grouped_norm`). The routed and shared
- experts, the gate projection and the mixer that consumes `normed` are not
- inside.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: rows, inject (True), moe (False: shared + gate_logits [T, 1] too), fold
- # (0: top-k of the finalize folded in), drop (False: some slots -1). The
- # checkpoint's geometry throughout: hc_count 4, hidden 2560, bf16, contiguous rows.
- # Four cells: the forward reaches four distinct launches -- the Triton kernel's
- # eight-warp form (decode), its two-warp form (prefill, here with the combine),
- # the CuTe fold kernel (decode MoE sites; the -1 slots ride it) and the norm-only
- # variant -- none of which another covers.
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="checkpoint 4x2560 bf16, inject, decode batch 128 (eight-warp form)",
- spec=dict(rows=128),
- ),
- Cell(
- why="MoE combine, [T, 1] gate, prefill chunk 16384 (the ladder's max_num_tokens)",
- spec=dict(rows=16384, moe=True),
- ),
- Cell(
- why="MoE combine + finalize fold (top-10, CuTe kernel), three slots -1, decode 128",
- spec=dict(rows=128, moe=True, fold=10, drop=True),
- ),
- Cell(
- why="norm only (no inject: stack start / after PLE), decode 128",
- spec=dict(rows=128, inject=False),
- ),
- )
-
- note = """
- Metadata consumed: none; every (row, stream) independent. Call sites: in
- the trunk forward of both targets, once after each block (attention / GDN
- and MoE) of each layer, normalizing for the next site -- the layer's MoE
- mixer, the next layer's attention mixer, or the stack-final mixer -- plus
- norm-only calls at the stack's first site and after the PLE add. 97 calls
- per forward on the 48-layer trunk; the 48 after a MoE block pass `shared`
- and `gate_logits` (in `DecodeTarget` all but the one before the PLE add,
- which keeps the unfused combine, and those also `topk_weights` /
- `topk_rows`: the runner's finalize folded in). The MTP layer and the PLE
- block keep the unfused ops.
-
- Preconditions: 2-D operands, last-dim stride 1 (row strides free); the
- width divisible by `hc_count`; `inject` `[T, hc_count]`, `block_out` and
- `shared` `[T, hidden]` of one dtype (with the fold `block_out` is `[P,
- hidden]`, any `P`, read through int32 `topk_rows` `[T, k]`); `gate_logits` `T` elements along its
- first axis (`[T, 1]` or `[T]`, any row stride); `norm_weight` contiguous
- `[width]`. All on one device. `shared` and `gate_logits` come together and
- only with `block_out`.
-
- Numerics: the combine is fp32 with one rounding to `block_out`'s dtype
- (bitwise `fused_sigmoid_gate_mul_add`: the same `tl.sigmoid` and fma); the
- accumulate is fp32 with one rounding (bitwise what `qwen4_exp_post_inject`
- stores for 16-bit operands up to fma contraction); the norm is fp32 with
- `rsqrt`, one rounding. Against the float64 reference the combine can land
- one ulp of the block output away (its rounding midpoint), which the
- accumulate carries scaled by `inject`; the test's band states that.
- Graph-capturable by construction (no host sync, outputs allocated per
- call); not certified.
-
- Launch: one program per (row, stream), stream-minor in a 1-D grid, so a
- row's `block_out` / `shared` are read from HBM once and from L2 by the
- other `hc_count - 1` programs; with the fold, one program per two streams
- of a row (the gathered rows are read and summed half as often). Every program releases programmatic
- dependents at its start (`griddepcontrol.launch_dependents`; off with
- TRTLLM_ENABLE_PDL=0): a consumer launched with PDL (the decode mixer's
- `qwen4_exp_hc_decode_tc_first`) streams its weights meanwhile and waits for
- this grid before reading `normed`. A no-op for any other consumer. At a
- hidden of 2048 or more a program runs eight warps below
- `_FEW_WARPS_MIN_ROWS` rows and two from it; the RMS sum's fp32 reduction
- order follows the warp count.
-
- The decode MoE sites -- the fold below `_FEW_WARPS_MIN_ROWS` rows, bf16
- operands, `hc_count` 4, a hidden in (2048, 4096], 16-byte aligned rows with
- row strides a multiple of 16 -- run `_FoldNormKernel` (CuTe DSL) instead:
- a row's two CTAs one cluster, CTA j forming half the block output (the
- finalize and the combine) into both CTAs' shared memory, then each
- accumulating and normalizing two streams, in the Triton kernel's layout and
- reduction tree -- bitwise its `hc_out` and `normed` (tested at 511 rows, the
- last before the switch). GB300, graph replay,
- 16 rotating sites, us: 128 rows 7.00 -> 5.97, 64 rows 6.63 -> 5.25, 8 rows
- 3.59 -> 3.09; it releases dependents at its start as the Triton kernel does.
-
- Certified: bf16, the checkpoint's 4 x 2560: the injection at 128 rows
- (eight-warp form), the combine at 16384 rows (two-warp form), the combine
- with the fold (three slots -1) at 128 rows (CuTe kernel), norm-only at 128
- rows; the CuTe fold bitwise the Triton kernel at 511 rows.
-
- Not certified: hc_count other than 4, hidden other than 2560, float16
- operands, a flat `[T]` gate, row strides other than the width, zero rows,
- the injection without the combine at prefill sizes, the combine without the
- fold at decode sizes, CUDA-graph replay (the kernels take them; no cell
- drives them).
-
- Hard limits: row offsets are int64 (`program_id(0)` widened); the grid is
- `rows * hc_count` programs along x (2^31 - 1); the column offset
- `stream * hidden + i` is int32, bounded by the width.
- """
-
- def __call__(
- self,
- hyper_input: torch.Tensor,
- block_out: Optional[torch.Tensor],
- inject: Optional[torch.Tensor],
- norm_weight: torch.Tensor,
- eps: float,
- hc_count: int,
- shared: Optional[torch.Tensor] = None,
- gate_logits: Optional[torch.Tensor] = None,
- topk_weights: Optional[torch.Tensor] = None,
- topk_rows: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
- rows, width = hyper_input.shape
- hidden = width // hc_count
- normed = torch.empty((rows, width), dtype=hyper_input.dtype, device=hyper_input.device)
- has_inject = block_out is not None
- has_shared = shared is not None
- top_k = topk_rows.shape[1] if topk_rows is not None else 0
- hc_out = (
- torch.empty((rows, width), dtype=hyper_input.dtype, device=hyper_input.device)
- if has_inject
- else hyper_input
- )
- if (
- rows > 0
- and top_k
- and _fold_cute_ok(
- hyper_input,
- block_out,
- inject,
- norm_weight,
- hc_count,
- shared,
- gate_logits,
- topk_weights,
- topk_rows,
- )
- ):
- _fold_cute(
- hyper_input,
- block_out,
- inject,
- norm_weight,
- eps,
- hc_count,
- shared,
- gate_logits,
- topk_weights,
- topk_rows,
- hc_out,
- normed,
- )
- elif rows > 0:
- block = triton.next_power_of_2(hidden)
- # The fold gathers TOPK rows to form a row's block output: two streams a
- # program form it half as often (when hc_count is even).
- spp = 2 if top_k and hc_count % 2 == 0 else 1
- _hc_inject_norm_kernel[(rows * hc_count // spp,)](
- hyper_input,
- block_out if has_inject else hyper_input,
- inject if has_inject else hyper_input,
- shared if has_shared else hyper_input,
- gate_logits if has_shared else hyper_input,
- topk_weights if top_k else hyper_input,
- topk_rows if top_k else hyper_input,
- norm_weight,
- hc_out,
- normed,
- hyper_input.stride(0),
- block_out.stride(0) if has_inject else 0,
- inject.stride(0) if has_inject else 0,
- shared.stride(0) if has_shared else 0,
- gate_logits.stride(0) if has_shared else 0,
- topk_weights.stride(0) if top_k else 0,
- topk_rows.stride(0) if top_k else 0,
- hc_out.stride(0),
- normed.stride(0),
- eps,
- HC_COUNT=hc_count,
- HIDDEN=hidden,
- BLOCK=block,
- HAS_INJECT=has_inject,
- HAS_SHARED=has_shared,
- TOPK=top_k,
- STREAMS_PER_PROGRAM=spp,
- PDL_TRIGGER=_PDL_TRIGGER,
- num_warps=(2 if rows >= _FEW_WARPS_MIN_ROWS else 8) if block >= 2048 else 4,
- )
- return hc_out, normed
-
- def reference(
- self,
- hyper_input: torch.Tensor,
- block_out: Optional[torch.Tensor],
- inject: Optional[torch.Tensor],
- norm_weight: torch.Tensor,
- eps: float,
- hc_count: int,
- shared: Optional[torch.Tensor] = None,
- gate_logits: Optional[torch.Tensor] = None,
- topk_weights: Optional[torch.Tensor] = None,
- topk_rows: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
- """float64 throughout. The finalize (with `topk_rows`) and the combine
- (with `shared`) are each rounded once to `block_out`'s dtype, as the op
- states; `hc_out` is the float64
- accumulate rounded once; `normed` is the float64 norm of that rounded
- residual, rounded once -- the same dataflow the op states, with no
- other intermediate rounding."""
- dtype = hyper_input.dtype
- rows, width = hyper_input.shape
- hidden = width // hc_count
- if block_out is None:
- hc_out = hyper_input
- else:
- if topk_rows is not None:
- valid = (topk_rows >= 0).unsqueeze(-1)
- picked = block_out.double()[topk_rows.clamp(min=0).long()]
- w = topk_weights.double().unsqueeze(-1)
- block_out = torch.where(valid, picked * w, 0.0).sum(1).to(block_out.dtype)
- if shared is not None:
- s = torch.sigmoid(gate_logits.double().reshape(rows, 1))
- block_out = (block_out.double() + s * shared.double()).to(block_out.dtype)
- streams = hyper_input.double().unflatten(-1, (hc_count, hidden))
- gated = block_out.double().unsqueeze(-2) * inject.double().unsqueeze(-1)
- hc_out = (streams + gated).flatten(-2).to(dtype)
- x = hc_out.double().unflatten(-1, (hc_count, hidden))
- x = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps)
- normed = (x.flatten(-2) * (1.0 + norm_weight.double())).to(dtype)
- return hc_out, normed
-
- def is_valid(
- self,
- hyper_input: torch.Tensor,
- block_out: Optional[torch.Tensor],
- inject: Optional[torch.Tensor],
- norm_weight: torch.Tensor,
- eps: float,
- hc_count: int,
- shared: Optional[torch.Tensor] = None,
- gate_logits: Optional[torch.Tensor] = None,
- topk_weights: Optional[torch.Tensor] = None,
- topk_rows: Optional[torch.Tensor] = None,
- ) -> None:
- # The kernel indexes with raw row strides and unit column stride, and
- # reads `inject` / `block_out` / `shared` / the gate per row: a broadcast
- # (stride-0) or a wrong-width operand would be read out of bounds or
- # silently reused.
- assert hyper_input.dim() == 2 and hyper_input.stride(1) == 1, "hyper_input [T, width]"
- assert (block_out is None) == (inject is None), "block_out and inject come together"
- assert (shared is None) == (gate_logits is None), "shared and gate_logits come together"
- assert shared is None or block_out is not None, "the combine needs block_out"
- assert (topk_weights is None) == (topk_rows is None), (
- "topk_weights and topk_rows come together"
- )
- assert topk_rows is None or shared is not None, "the finalize fold needs the combine"
- rows, width = hyper_input.shape
- assert width % hc_count == 0, f"width {width} not divisible by hc_count {hc_count}"
- assert norm_weight.shape == (width,) and norm_weight.is_contiguous(), "norm_weight [width]"
- if block_out is not None:
- if topk_rows is None:
- assert block_out.shape == (rows, width // hc_count) and block_out.stride(1) == 1
- else:
- # The permuted FC2 rows: any row count, read through topk_rows.
- assert block_out.dim() == 2 and block_out.shape[1] == width // hc_count
- assert block_out.stride(1) == 1, "permuted rows [P, hidden]"
- assert topk_rows.dtype == torch.int32, "topk_rows int32 (the runner's index map)"
- assert topk_rows.dim() == 2 and topk_rows.shape[0] == rows, "topk_rows [T, k]"
- assert topk_rows.stride(1) == 1 and topk_weights.stride(1) == 1
- assert topk_weights.shape == topk_rows.shape, "topk_weights [T, k]"
- assert inject.shape == (rows, hc_count) and inject.stride(1) == 1
- if shared is not None:
- assert shared.shape == (rows, width // hc_count), "shared [T, hidden]"
- assert shared.stride(1) == 1, "shared [T, hidden]"
- assert shared.dtype == block_out.dtype, "shared in block_out's dtype"
- assert gate_logits.shape in ((rows,), (rows, 1)), "gate_logits [T, 1] or [T]"
- assert rows <= 1 or gate_logits.stride(0) != 0, "one gate logit per row"
-
-
-qwen4_exp_hc_inject_norm = Qwen4ExpHcInjectNorm()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_rstd.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_rstd.py
deleted file mode 100644
index 597216df40fc..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_rstd.py
+++ /dev/null
@@ -1,258 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Hyper-connection post-block injection plus the next site's per-(row, stream) RMS
-statistic, without the normalized copy: `qwen4_exp_hc_inject_norm` with its
-`normed` store replaced by one fp32 rstd per (row, stream).
-
-Tier 5 (Triton written for this target). At prefill sizes the grouped norm's output
-(`[T, C*H]`, 335 MB at 16384 rows) is a full extra write whose only readers are the
-site's two projections; those form `bf16(hc_out * rstd * (1 + w))` themselves
-(`qwen4_exp_hc_first_rescale_tc` on its A tiles, `qwen4_exp_hc_up_gate_mean_rstd` in
-its epilogue), so the residual pass moves the residual in and out once and nothing
-else.
-"""
-
-from __future__ import annotations
-
-from typing import Optional, Tuple
-
-import torch
-import triton
-import triton.language as tl
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: Rows from which a wide (`hidden >= 2048`) launch runs two warps a program instead
-#: of eight (`qwen4_exp_hc_inject_norm`'s split; the reduction tree follows it).
-_FEW_WARPS_MIN_ROWS = 512
-
-
-@triton.jit
-def _hc_inject_rstd_kernel(
- hc_ptr,
- block_ptr,
- inject_ptr,
- shared_ptr,
- gate_ptr,
- hc_out_ptr,
- rstd_ptr,
- stride_hc,
- stride_block,
- stride_inject,
- stride_shared,
- stride_gate,
- stride_out,
- eps,
- HC_COUNT: tl.constexpr,
- HIDDEN: tl.constexpr,
- BLOCK: tl.constexpr,
- HAS_INJECT: tl.constexpr,
- HAS_SHARED: tl.constexpr,
-):
- # One program per (row, stream), stream-minor: a row's block output (and the
- # combine's operands) come from HBM once and from L2 for the other programs. The
- # loads, the fp32 arithmetic and the reduction are `qwen4_exp_hc_inject_norm`'s
- # (without the fold), so `hc_out` and rstd are bitwise what that kernel stores and
- # multiplies by at the same warp count.
- pid = tl.program_id(0).to(tl.int64)
- row = pid // HC_COUNT
- stream = (pid % HC_COUNT).to(tl.int32)
- cols = tl.arange(0, BLOCK)
- mask = cols < HIDDEN
- offs = stream * HIDDEN + cols
- x = tl.load(hc_ptr + row * stride_hc + offs, mask=mask, other=0.0).to(tl.float32)
- if HAS_INJECT:
- 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.
- 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)
- g = tl.load(inject_ptr + row * stride_inject + stream).to(tl.float32)
- 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)
- tl.store(rstd_ptr + row * HC_COUNT + stream, rstd)
-
-
-class Qwen4ExpHcInjectRstd(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]
- 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 rstd is computed from that
- rounded value in fp32. Returns `(hc_out, rstd)`: `hc_out` a fresh `[T, C*H]` in
- `hyper_input`'s dtype, rstd a fresh contiguous `[T, C]` fp32. A site's normalized
- stream is `bf16(hc_out * rstd * (1 + norm_weight))`, bitwise
- `qwen4_exp_hc_inject_norm`'s `normed` (same fp32 operations and rounding).
-
- `block_out=None, inject=None` is the statistic alone: `hc_out` *is* `hyper_input`
- (not a copy) and only rstd 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 injected block output is
- `block_out + sigmoid(gate_logits) * shared` in fp32, rounded once to `block_out`'s
- dtype (`qwen4_exp_hc_inject_norm`'s combine). None of the inputs is written.
-
- Fusion boundary: the shared-expert combine (optional), the stream-broadcast gated
- accumulate and the per-(row, stream) RMS statistic. The norm weight is not an
- argument: the consumers apply it.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: rows, inject (True), moe (False: shared + gate_logits [T, 1] too). The
- # checkpoint's geometry throughout: hc_count 4, hidden 2560, bf16, contiguous rows.
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="MoE combine, prefill chunk 16384 (the mixed step's size)",
- spec=dict(rows=16384, moe=True),
- ),
- Cell(why="injection after attention / GDN, prefill chunk 16384", spec=dict(rows=16384)),
- Cell(why="MoE combine, prefill 2048", spec=dict(rows=2048, moe=True)),
- Cell(why="injection after attention / GDN, prefill 2048", spec=dict(rows=2048)),
- Cell(
- why="statistic only (stack start / after PLE), ragged 16390 rows",
- spec=dict(rows=16390, inject=False),
- ),
- )
-
- note = """
- Metadata consumed: none; every (row, stream) independent. Call sites: the trunk
- sites of `PrefillTarget` at or above `_HC_RSTD_MIN_ROWS` rows, in place of
- `qwen4_exp_hc_inject_norm`: once after each block (attention / GDN and MoE) of
- each layer, plus the statistic-only calls at the stack's first site and after
- the PLE add. Their mixers are `qwen4_exp_hc_first_rescale_tc` and
- `qwen4_exp_hc_up_gate_mean_rstd(_fp4)`, which take the site's norm weight.
-
- 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; `gate_logits` `T` elements along its first axis.
- `shared` and `gate_logits` come together and only with `block_out`.
-
- Numerics: `hc_out` and rstd are bitwise `qwen4_exp_hc_inject_norm`'s at the same
- row count (that kernel's two-warp form from 512 rows, eight below); rstd is the
- fp32 `rsqrt` of the fp32 mean of squares. Graph-capturable (no host sync, outputs
- allocated per call).
-
- Launch: one program per (row, stream), stream-minor in a 1-D grid. GB300, 16384
- rows: 0.75 GB (0.84 GB with the combine) at about 7 TB/s, against
- `qwen4_exp_hc_inject_norm`'s 1.09 / 1.17 GB.
-
- Certified: bf16, the checkpoint's 4 x 2560: the combine and the injection at 2048
- and 16384 rows, the statistic alone at 16390 rows; `hc_out` and the rebuilt
- normed stream bitwise `qwen4_exp_hc_inject_norm`'s; graph replay.
-
- Not certified: hc_count other than 4, hidden other than 2560, float16, the
- finalize fold (`topk_rows`; the prefill sites do not fold), zero rows.
-
- Hard limits: row offsets are int64; the grid is `rows * hc_count` programs along
- x (2^31 - 1).
- """
-
- def __call__(
- self,
- hyper_input: torch.Tensor,
- block_out: Optional[torch.Tensor],
- inject: Optional[torch.Tensor],
- eps: float,
- hc_count: int,
- shared: Optional[torch.Tensor] = None,
- gate_logits: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
- rows, width = hyper_input.shape
- hidden = width // hc_count
- rstd = torch.empty((rows, hc_count), dtype=torch.float32, device=hyper_input.device)
- has_inject = block_out is not None
- has_shared = shared is not None
- hc_out = torch.empty_like(hyper_input) if has_inject else hyper_input
- if rows > 0:
- block = triton.next_power_of_2(hidden)
- _hc_inject_rstd_kernel[(rows * 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,
- hc_out,
- rstd,
- 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,
- hc_out.stride(0),
- eps,
- HC_COUNT=hc_count,
- HIDDEN=hidden,
- BLOCK=block,
- HAS_INJECT=has_inject,
- HAS_SHARED=has_shared,
- num_warps=(2 if rows >= _FEW_WARPS_MIN_ROWS else 8) if block >= 2048 else 4,
- )
- return hc_out, rstd
-
- def reference(
- self,
- hyper_input: torch.Tensor,
- block_out: Optional[torch.Tensor],
- inject: Optional[torch.Tensor],
- eps: float,
- hc_count: int,
- shared: Optional[torch.Tensor] = None,
- gate_logits: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
- """float64: the combine (with `shared`) rounded once to `block_out`'s dtype,
- `hc_out` the accumulate rounded once, rstd the float64 statistic of that
- rounded residual returned in fp32 -- the dataflow the op states.
- """
- dtype = hyper_input.dtype
- rows, width = hyper_input.shape
- hidden = width // hc_count
- if block_out is None:
- hc_out = hyper_input
- else:
- 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))
- rstd = torch.rsqrt(x.pow(2).mean(dim=-1) + eps).float()
- return hc_out, rstd
-
- def is_valid(
- self,
- hyper_input: torch.Tensor,
- block_out: Optional[torch.Tensor],
- inject: Optional[torch.Tensor],
- eps: float,
- hc_count: int,
- shared: Optional[torch.Tensor] = None,
- gate_logits: 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"
- rows, width = hyper_input.shape
- assert width % hc_count == 0, f"width {width} not divisible by hc_count {hc_count}"
- if block_out is not None:
- assert block_out.shape == (rows, width // hc_count) and block_out.stride(1) == 1
- assert inject.shape == (rows, hc_count) and inject.stride(1) == 1
- if shared is not None:
- assert shared.shape == (rows, width // hc_count) and shared.stride(1) == 1
- 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_rstd = Qwen4ExpHcInjectRstd()
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
deleted file mode 100644
index 493ca58301c4..000000000000
--- 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
+++ /dev/null
@@ -1,150 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Hyper-connection mixer activations between its two projections: the SiLU of
-the rank-`R` bottleneck and the block-injection weights, from the first
-projection's fp32 logits, in one Triton row kernel.
-
-Tier 5 (Triton written for this target). Used where the first projection runs on
-cuBLAS (prefill, and decode batches past `qwen4_exp_hc_decode_tc_first`'s rows),
-between `cublas_mm` (first projection) and the fused up projection
-`qwen4_exp_hc_up_gate_mean` (or its NVFP4 variant).
-"""
-
-from __future__ import annotations
-
-from typing import Optional, Tuple
-
-import torch
-import torch.nn.functional as F
-import triton
-import triton.language as tl
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-
-@triton.jit
-def _hc_lowrank_act_kernel(
- logits_ptr,
- low_ptr,
- inject_ptr,
- stride_logits,
- stride_low,
- stride_inject,
- RANK: tl.constexpr,
- HC: tl.constexpr,
- WITH_INJECT: tl.constexpr,
- BLOCK: tl.constexpr,
-):
- row = tl.program_id(0).to(tl.int64)
- k = tl.arange(0, BLOCK)
- z = tl.load(logits_ptr + row * stride_logits + k, mask=k < RANK, other=0.0).to(tl.float32)
- z = z * (1.0 / HC)
- tl.store(
- low_ptr + row * stride_low + k,
- (z * tl.sigmoid(z)).to(low_ptr.dtype.element_ty),
- mask=k < RANK,
- )
- if WITH_INJECT:
- i = tl.arange(0, HC)
- zi = tl.load(logits_ptr + row * stride_logits + RANK + i).to(tl.float32)
- tl.store(
- inject_ptr + row * stride_inject + i,
- (2.0 * tl.sigmoid(zi * (1.0 / HC))).to(inject_ptr.dtype.element_ty),
- )
-
-
-class Qwen4ExpHcLowrankAct(OpWrapper):
- """With `R = rank`, `C = hc_count` and `logits` the fp32 `[T, N]` first
- projection `normed @ [down; block_inject; 0]^T` (columns `[0, R)` down,
- `[R, R + C)` injection when `with_inject`, the rest padding)::
-
- low = silu(logits[:, :R] / C) # [T, R], `dtype`
- inject = 2 * sigmoid(logits[:, R : R + C] / C) # [T, C], `dtype`, only with_inject
-
- fp32 inside, one rounding per output. Returns fresh `(low, inject)`;
- `inject` is `None` exactly when `with_inject` is False.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: rows, inject (True), hc_count (4), rank (320), logits_width (384), dtype (bf16).
- CELLS: tuple[Cell, ...] = (
- Cell(why="checkpoint rank 320 + 4, padded to 384, 128 rows", spec=dict(rows=128)),
- Cell(why="checkpoint rank 320 + 4, padded to 384, prefill 8192", spec=dict(rows=8192)),
- Cell(why="stack-final mixer: no inject, 8192 rows", spec=dict(rows=8192, inject=False)),
- )
-
- note = """
- Metadata consumed: none. Call sites: the trunk's hyper-connection sites in
- `PrefillTarget` (all 97), between the first projection and the up
- projection. Preconditions: `logits` fp32, unit column stride, at least
- `rank (+ hc_count)` columns; `hc_count` a power of two. Graph-capturable.
-
- Certified: bf16 outputs from 384-wide fp32 logits, with the injection at 128
- and 8192 rows, without it at 8192 rows; graph replay.
-
- Not certified: zero rows, hc_count other than 4, rank other than 320,
- float16 outputs.
- """
-
- def __call__(
- self,
- logits: torch.Tensor,
- rank: int,
- hc_count: int,
- with_inject: bool,
- dtype: torch.dtype,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
- rows = logits.shape[0]
- low = torch.empty((rows, rank), dtype=dtype, device=logits.device)
- inject = (
- torch.empty((rows, hc_count), dtype=dtype, device=logits.device)
- if with_inject
- else None
- )
- if rows > 0:
- _hc_lowrank_act_kernel[(rows,)](
- logits,
- low,
- inject if with_inject else low,
- logits.stride(0),
- low.stride(0),
- inject.stride(0) if with_inject else 0,
- RANK=rank,
- HC=hc_count,
- WITH_INJECT=with_inject,
- BLOCK=triton.next_power_of_2(rank),
- num_warps=4,
- )
- return low, inject
-
- def reference(
- self,
- logits: torch.Tensor,
- rank: int,
- hc_count: int,
- with_inject: bool,
- dtype: torch.dtype,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
- """float64, one rounding per output."""
- lg = logits.double()
- low = F.silu(lg[:, :rank] / hc_count).to(dtype)
- inject = None
- if with_inject:
- inject = (2.0 * torch.sigmoid(lg[:, rank : rank + hc_count] / hc_count)).to(dtype)
- return low, inject
-
- def is_valid(
- self,
- logits: torch.Tensor,
- rank: int,
- hc_count: int,
- with_inject: bool,
- dtype: torch.dtype,
- ) -> None:
- assert logits.dtype == torch.float32 and logits.stride(1) == 1, "fp32 logits, unit stride"
- assert logits.shape[1] >= rank + (hc_count if with_inject else 0), "logits too narrow"
- assert hc_count & (hc_count - 1) == 0, "hc_count must be a power of two"
-
-
-qwen4_exp_hc_lowrank_act = Qwen4ExpHcLowrankAct()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py
deleted file mode 100644
index 755fc8535e08..000000000000
--- 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
+++ /dev/null
@@ -1,899 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Hyper-connection mixer second half as one tcgen05 GEMM with a custom
-epilogue: the up projection `low @ up^T` accumulates in tensor memory and the
-epilogue applies the sigmoid gate to the normed stream and averages the
-streams, so the `[T, C*H]` gate logits never reach global memory.
-
-Tier 6 (CuTe DSL written for this target). A persistent warp-specialized
-kernel: one TMA warp loads `low` and the up weight, one MMA warp issues
-`tcgen05.mma` into a double-buffered TMEM accumulator, four epilogue warps
-read the accumulator, load the matching normed elements straight from global
-memory and write the mean. An MMA tile is 128 rows by `C * BH` up-projection
-outputs: the `BH`-wide hidden slice `[h0, h0 + BH)` of every stream, read
-through a strided view of the unmodified `[C*H, R]` weight (no copy, no
-re-layout), so the stream reduction stays inside one thread.
-"""
-
-from __future__ import annotations
-
-import cuda.bindings.driver as cuda
-import cutlass
-import cutlass.cute as cute
-import cutlass.pipeline as pipeline
-import cutlass.utils as utils
-import torch
-import torch.nn.functional as F
-from cutlass._mlir.dialects import llvm
-from cutlass.cute.nvgpu import cpasync, tcgen05
-from cutlass.cute.runtime import from_dlpack
-from cutlass.cutlass_dsl import T, dsl_user_op
-from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
-from tensorrt_llm._torch.kernels.blackwell.custom_pipeline import PipelineTmaUmma, PipelineUmmaAsync
-from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
-
-#: Programmatic dependent launch (sm_90+), off with TRTLLM_ENABLE_PDL=0.
-_PDL = get_env_enable_pdl()
-
-
-@dsl_user_op
-def _sigmoid(x, *, loc=None, ip=None):
- """sigmoid(x) = 0.5 * tanh(0.5 * x) + 0.5: one MUFU op (tanh.approx.f32,
- max relative error 2^-10.99) instead of ex2 + rcp."""
- return cutlass.Float32(
- llvm.inline_asm(
- T.f32(),
- [cutlass.Float32(x).ir_value(loc=loc, ip=ip)],
- "{ .reg .f32 t; mul.f32 t, $1, 0f3F000000; tanh.approx.f32 t, t; "
- "fma.rn.f32 $0, t, 0f3F000000, 0f3F000000; }",
- "=f,f",
- has_side_effects=False,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- )
-
-
-# The NVFP4 epilogue (`fp4=True`) is the `torch.ops.trtllm.fp4_quantize` recipe
-# (`cvt_warp_fp16_to_fp4`, e4m3 scales) instruction for instruction, every
-# multiply an explicit `mul.rn.f32` so that no compiler contraction or
-# reassociation can move a rounding.
-
-
-def _asm(ret, args, text, constraints, loc, ip, side_effects=False):
- return llvm.inline_asm(
- ret,
- [a.ir_value(loc=loc, ip=ip) for a in args],
- text,
- constraints,
- has_side_effects=side_effects,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
-
-
-@dsl_user_op
-def _rcp_approx_ftz(x, *, loc=None, ip=None):
- """`rcp.approx.ftz.f32`, the quantizer's `reciprocal_approximate_ftz`."""
- return cutlass.Float32(
- _asm(T.f32(), [cutlass.Float32(x)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", loc, ip)
- )
-
-
-@dsl_user_op
-def _absmax16(v, *, loc=None, ip=None):
- """max |v_i| over 16 fp32 values (exact; NaN-quiet like the bf16 `__hmax2`)."""
- body = "abs.f32 m, $1; " + "".join(f"abs.f32 t, ${i}; max.f32 m, m, t; " for i in range(2, 17))
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Float32(x) for x in v],
- "{ .reg .f32 m, t; " + body + "mov.f32 $0, m; }",
- "=f" + ",f" * 16,
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _e4m3_scale(amax, gs, rcp6, *, loc=None, ip=None):
- """The block's scale byte: e4m3(gs * (amax * rcp(6))), `cvt.rn.satfinite`."""
- return cutlass.Int32(
- _asm(
- T.i32(),
- [cutlass.Float32(amax), cutlass.Float32(gs), cutlass.Float32(rcp6)],
- "{ .reg .f32 t; .reg .b16 h; mul.rn.f32 t, $1, $3; mul.rn.f32 t, $2, t; "
- "cvt.rn.satfinite.e4m3x2.f32 h, 0f00000000, t; cvt.u32.u16 $0, h; }",
- "=r,f,f,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _e2m1_out_scale(sf, rcp_gs, amax, *, loc=None, ip=None):
- """amax != 0 ? rcp(float(e4m3 sf) * rcp(gs)) : 0, the data multiplier."""
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Int32(sf), cutlass.Float32(rcp_gs), cutlass.Float32(amax)],
- "{ .reg .b16 h, lo, hi; .reg .b32 f2; .reg .f32 s, t; .reg .pred p; "
- "cvt.u16.u32 h, $1; cvt.rn.f16x2.e4m3x2 f2, h; mov.b32 {lo, hi}, f2; "
- "cvt.f32.f16 s, lo; mul.rn.f32 t, s, $2; rcp.approx.ftz.f32 t, t; "
- "setp.neu.f32 p, $3, 0f00000000; selp.f32 $0, t, 0f00000000, p; }",
- "=f,r,f,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _e2m1x8(v, scale, *, loc=None, ip=None):
- """Eight values times `scale`, each to e2m1 (`cvt.rn.satfinite`), packed
- into 4 bytes; element `2k` in byte k's low nibble."""
- muls = "".join(f"mul.rn.f32 a{i}, ${i + 1}, $9; " for i in range(8))
- cvts = "".join(f"cvt.rn.satfinite.e2m1x2.f32 b{k}, a{2 * k + 1}, a{2 * k}; " for k in range(4))
- return cutlass.Int32(
- _asm(
- T.i32(),
- [cutlass.Float32(x) for x in v] + [cutlass.Float32(scale)],
- "{ .reg .f32 a<8>; .reg .b8 b<4>; " + muls + cvts + "mov.b32 $0, {b0, b1, b2, b3}; }",
- "=r" + ",f" * 9,
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _st_global_v4(addr, a, b, c, d, *, loc=None, ip=None):
- """One 16-byte `st.global.v4.b32` (a row's packed e2m1 for 32 columns): a
- row-per-thread epilogue pays per store instruction, not per byte."""
- _asm(
- T.i32(),
- [cutlass.Int64(addr)] + [cutlass.Int32(x) for x in (a, b, c, d)],
- "{ st.global.v4.b32 [$1], {$2, $3, $4, $5}; mov.u32 $0, 0; }",
- "=r,l,r,r,r,r",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _st_global_b16(addr, x, *, loc=None, ip=None):
- """One 2-byte store of the low half of `x` (two e4m3 scale bytes)."""
- _asm(
- T.i32(),
- [cutlass.Int64(addr), cutlass.Int32(x)],
- "{ .reg .b16 t; cvt.u16.u32 t, $2; st.global.b16 [$1], t; mov.u32 $0, 0; }",
- "=r,l,r",
- loc,
- ip,
- side_effects=True,
- )
-
-
-class HcUpGateMeanKernel:
- """`mixed[t, h] = 1/C * sum_s sigmoid((low @ up^T)[t, s*H + h]) * normed[t, s*H + h]`.
-
- Compile-time: `block_h` (hidden columns per tile, 32 or 64), `hc` (C),
- `hidden` (H), `rank` (R), `fp4` (also write the NVFP4 of `mixed`: packed
- e2m1 `[T, H/2]` and linear e4m3 scales `[T, H/16]` under the global scale
- `gs`). Runtime: the row count and the pointers; without `fp4` the `xq`,
- `sf` and `gs` pointers are never read.
- """
-
- def __init__(
- self, block_h: int, hc: int, hidden: int, rank: int, num_n_stage: int = 2, fp4: bool = False
- ):
- assert block_h in (32, 64) and hidden % block_h == 0 and rank % 64 == 0
- self.bh, self.hc, self.hidden, self.rank = block_h, hc, hidden, rank
- self.fp4 = fp4
- self.acc_dtype = cutlass.Float32
- self.io_dtype = cutlass.BFloat16
- self.mma_tiler = (128, hc * block_h, 1)
- self.epi_cols = 32
- self.cta_group = tcgen05.CtaGroup.ONE
- self.cluster_shape_mn = (1, 1)
- self.epilogue_warp_id = (0, 1, 2, 3)
- self.mma_warp_id = 4
- self.tma_warp_id = 5
- self.nrm_warp_id = 6
- self.threads_per_cta = 32 * 7
- self.num_n_stage = num_n_stage
- self.epilog_sync_bar_id = 1
- self.tmem_alloc_sync_bar_id = 2
- self.num_acc_stage = 2
-
- def _setup(self, a_major, b_major):
- tiled_mma = utils.sm100.make_trivial_tiled_mma(
- self.io_dtype, a_major, b_major, self.acc_dtype, self.cta_group, self.mma_tiler[:2]
- )
- inst_k = cute.size(tiled_mma.shape_mnk, mode=[2])
- self.mma_tiler = (self.mma_tiler[0], self.mma_tiler[1], inst_k * 4)
- self.cta_tile_shape_mnk = self.mma_tiler
- self.cluster_layout_vmnk = cute.tiled_divide(
- cute.make_layout((*self.cluster_shape_mn, 1)), (tiled_mma.thr_id.shape,)
- )
- self.epi_tile = (self.mma_tiler[0], self.epi_cols)
- a_one = utils.sm100.make_smem_layout_a(tiled_mma, self.mma_tiler, self.io_dtype, 1)
- b_one = utils.sm100.make_smem_layout_b(tiled_mma, self.mma_tiler, self.io_dtype, 1)
- ab_bytes = cute.size_in_bytes(self.io_dtype, a_one) + cute.size_in_bytes(
- self.io_dtype, b_one
- )
- smem_capacity = utils.get_smem_capacity_in_bytes()
- # One normed stage is the 32-column slice of every stream for one tile half.
- self.n_box = (self.mma_tiler[0], self.epi_cols)
- self.num_n_bytes = self.hc * self.n_box[0] * self.n_box[1] * 2
- self.num_ab_stage = min(
- 8, (smem_capacity - 2048 - self.num_n_stage * self.num_n_bytes) // ab_bytes
- )
- self.n_pre = min(self.num_ab_stage, self.rank // self.mma_tiler[2])
- self.n_smem_layout_staged = utils.sm100.make_smem_layout_epi(
- self.io_dtype, utils.LayoutEnum.ROW_MAJOR, self.n_box, self.num_n_stage * self.hc
- )
- self.a_smem_layout_staged = utils.sm100.make_smem_layout_a(
- tiled_mma, self.mma_tiler, self.io_dtype, self.num_ab_stage
- )
- self.b_smem_layout_staged = utils.sm100.make_smem_layout_b(
- tiled_mma, self.mma_tiler, self.io_dtype, self.num_ab_stage
- )
- acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
- acc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage))
- self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(acc_fake, arch="sm_100")
- return tiled_mma
-
- @cute.jit
- def __call__(
- self,
- m: cutlass.Int32,
- low_ptr: cute.Pointer,
- up_ptr: cute.Pointer,
- normed_ptr: cute.Pointer,
- mixed: cute.Tensor,
- xq_ptr: cute.Pointer,
- sf_ptr: cute.Pointer,
- gs_ptr: cute.Pointer,
- max_active_clusters: cutlass.Constexpr,
- stream: cuda.CUstream,
- ):
- bh, hc, hid, rank = self.bh, self.hc, self.hidden, self.rank
- nt = hid // bh
- width = hc * hid
- # low [m, R] (K-major A); the up weight [C*H, R] seen as N = (BH, C, H/BH) so
- # that one N tile holds hidden slice [h0, h0+BH) of every stream.
- a = cute.make_tensor(low_ptr, cute.make_layout((m, rank, 1), stride=(rank, 1, 0)))
- b = cute.make_tensor(
- up_ptr,
- cute.make_layout(((bh, hc, nt), rank, 1), stride=((rank, hid * rank, bh * rank), 1, 0)),
- )
- # normed [m, C*H] is staged through shared memory by TMA in [128, 32] boxes;
- # mixed [m, H] is written in the accumulator's N order, its stream mode
- # with stride 0 (every stream maps onto the one output column).
- nrm = cute.make_tensor(normed_ptr, cute.make_layout((m, width), stride=(width, 1)))
- mix = cute.make_tensor(
- mixed.iterator, cute.make_layout((m, (bh, hc, nt), 1), stride=(hid, (1, 0, bh), 0))
- )
- # NVFP4 outputs (`fp4` only): 8 e2m1 codes per Int32, one e4m3 byte per 16.
- xq = cute.make_tensor(xq_ptr, cute.make_layout((m, hid // 8), stride=(hid // 8, 1)))
- xsf = cute.make_tensor(sf_ptr, cute.make_layout((m, hid // 16), stride=(hid // 16, 1)))
- gsc = cute.make_tensor(gs_ptr, cute.make_layout(1))
- k_major = utils.LayoutEnum.ROW_MAJOR.mma_major_mode()
- tiled_mma = self._setup(k_major, k_major)
-
- a_op = utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id)
- a_smem = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
- tma_atom_a, tma_a = cute.nvgpu.make_tiled_tma_atom_A(
- a_op, a, a_smem, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape
- )
- b_op = utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id)
- b_smem = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
- tma_atom_b, tma_b = cute.nvgpu.make_tiled_tma_atom_B(
- b_op, b, b_smem, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape
- )
- tma_atom_n, tma_n = cpasync.make_tiled_tma_atom(
- cpasync.CopyBulkTensorTileG2SOp(),
- nrm,
- cute.select(self.n_smem_layout_staged, mode=[0, 1]),
- self.n_box,
- )
- self.num_tma_load_bytes = cute.size_in_bytes(self.io_dtype, a_smem) + cute.size_in_bytes(
- self.io_dtype, b_smem
- )
- # Rastered along N: consecutive CTAs take consecutive hidden slices of the
- # same rows, so the normed reads in flight cover whole rows rather than
- # one 64-byte column block of many rows (~2% faster at 8K-16K rows).
- sched = utils.PersistentTileSchedulerParams(
- ((m + 127) // 128, nt, 1), (*self.cluster_shape_mn, 1), 1, False
- )
- grid = utils.StaticPersistentTileScheduler.get_grid_shape(sched, max_active_clusters)
- self.kernel(
- tiled_mma,
- tma_atom_a,
- tma_a,
- tma_atom_b,
- tma_b,
- tma_atom_n,
- tma_n,
- mix,
- xq,
- xsf,
- gsc,
- m,
- self.cluster_layout_vmnk,
- self.a_smem_layout_staged,
- self.b_smem_layout_staged,
- self.n_smem_layout_staged,
- self.epi_tile,
- sched,
- ).launch(
- grid=grid,
- block=[self.threads_per_cta, 1, 1],
- cluster=(*self.cluster_shape_mn, 1),
- stream=stream,
- use_pdl=_PDL,
- )
-
- @cute.kernel
- def kernel(
- self,
- tiled_mma: cute.TiledMma,
- tma_atom_a: cute.CopyAtom,
- mA_mkl: cute.Tensor,
- tma_atom_b: cute.CopyAtom,
- mB_nkl: cute.Tensor,
- tma_atom_n: cute.CopyAtom,
- mN_mn: cute.Tensor,
- mX_mnl: cute.Tensor,
- mQ: cute.Tensor,
- mSF: cute.Tensor,
- mGS: cute.Tensor,
- m: cutlass.Int32,
- cluster_layout_vmnk: cute.Layout,
- a_smem_layout_staged: cute.ComposedLayout,
- b_smem_layout_staged: cute.ComposedLayout,
- n_smem_layout_staged: cute.ComposedLayout,
- epi_tile: cute.Tile,
- sched_params: utils.PersistentTileSchedulerParams,
- ):
- warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
- if warp_idx == self.tma_warp_id:
- cpasync.prefetch_descriptor(tma_atom_a)
- cpasync.prefetch_descriptor(tma_atom_b)
- if warp_idx == self.nrm_warp_id:
- cpasync.prefetch_descriptor(tma_atom_n)
- tidx, _, _ = cute.arch.thread_idx()
-
- @cute.struct
- class SharedStorage:
- ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]
- acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]
- n_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_n_stage * 2]
- tmem_dealloc_mbar: cutlass.Int64
- tmem_holding_buf: cutlass.Int32
-
- smem = utils.SmemAllocator()
- storage = smem.allocate(SharedStorage)
- ab_producer, ab_consumer = PipelineTmaUmma.create(
- barrier_storage=storage.ab_full_mbar_ptr.data_ptr(),
- num_stages=self.num_ab_stage,
- producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
- consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1),
- tx_count=self.num_tma_load_bytes,
- cta_layout_vmnk=cluster_layout_vmnk,
- ).make_participants()
- acc_pipeline = PipelineUmmaAsync.create(
- barrier_storage=storage.acc_full_mbar_ptr.data_ptr(),
- num_stages=self.num_acc_stage,
- producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
- consumer_group=pipeline.CooperativeGroup(
- pipeline.Agent.Thread, len(self.epilogue_warp_id)
- ),
- cta_layout_vmnk=cluster_layout_vmnk,
- )
- n_producer, n_consumer = pipeline.PipelineTmaAsync.create(
- barrier_storage=storage.n_full_mbar_ptr.data_ptr(),
- num_stages=self.num_n_stage,
- producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
- consumer_group=pipeline.CooperativeGroup(
- pipeline.Agent.Thread, len(self.epilogue_warp_id)
- ),
- tx_count=self.num_n_bytes,
- cta_layout_vmnk=cluster_layout_vmnk,
- ).make_participants()
- tmem_alloc_barrier = pipeline.NamedBarrier(
- barrier_id=self.tmem_alloc_sync_bar_id,
- num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)),
- )
- tmem = utils.TmemAllocator(
- storage.tmem_holding_buf.ptr,
- barrier_for_retrieve=tmem_alloc_barrier,
- allocator_warp_id=self.epilogue_warp_id[0],
- is_two_cta=False,
- two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
- )
- pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
-
- sA = smem.allocate_tensor(
- element_type=self.io_dtype,
- layout=a_smem_layout_staged.outer,
- byte_alignment=128,
- swizzle=a_smem_layout_staged.inner,
- )
- sB = smem.allocate_tensor(
- element_type=self.io_dtype,
- layout=b_smem_layout_staged.outer,
- byte_alignment=128,
- swizzle=b_smem_layout_staged.inner,
- )
- sN = smem.allocate_tensor(
- element_type=self.io_dtype,
- layout=n_smem_layout_staged.outer,
- byte_alignment=128,
- swizzle=n_smem_layout_staged.inner,
- )
-
- gA_mkl = cute.local_tile(
- mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
- )
- gB_nkl = cute.local_tile(
- mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
- )
- tile_mn = cute.slice_(self.mma_tiler, (None, None, 0))
- gX_mnl = cute.local_tile(mX_mnl, tile_mn, (None, None, None))
- cX_mnl = cute.local_tile(
- cute.make_identity_tensor(mX_mnl.shape), tile_mn, (None, None, None)
- )
- # normed boxes: [128 rows, 32 columns] at (row tile, column block)
- gN_box = cute.local_tile(mN_mn, self.n_box, (None, None))
- k_tile_cnt = cute.size(gA_mkl, mode=[3])
-
- thr_mma = tiled_mma.get_slice(0)
- tCgA = thr_mma.partition_A(gA_mkl)
- tCgB = thr_mma.partition_B(gB_nkl)
- tCgX = thr_mma.partition_C(gX_mnl)
- tCcX = thr_mma.partition_C(cX_mnl)
-
- tAsA, tAgA = cpasync.tma_partition(
- tma_atom_a,
- 0,
- cute.make_layout(1),
- cute.group_modes(sA, 0, 3),
- cute.group_modes(tCgA, 0, 3),
- )
- tBsB, tBgB = cpasync.tma_partition(
- tma_atom_b,
- 0,
- cute.make_layout(1),
- cute.group_modes(sB, 0, 3),
- cute.group_modes(tCgB, 0, 3),
- )
- tNsN, tNgN = cpasync.tma_partition(
- tma_atom_n,
- 0,
- cute.make_layout(1),
- cute.group_modes(sN, 0, 2),
- cute.group_modes(gN_box, 0, 2),
- )
- tCrA = tiled_mma.make_fragment_A(sA)
- tCrB = tiled_mma.make_fragment_B(sB)
- acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
- tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage))
-
- pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
- tile_sched = utils.StaticPersistentTileScheduler.create(
- sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
- )
- work_tile = tile_sched.initial_work_tile_info()
-
- if warp_idx == self.tma_warp_id:
- # Programmatic dependent launch: the first tile's up-weight stages (static)
- # stream before the grid-dependency wait, `low` -- the previous kernel's
- # output -- after it. Every CTA owns at least one tile (grid <= tiles).
- crd = work_tile.tile_idx
- tAgA_slice = tAgA[(None, crd[0], None, crd[2])]
- tBgB_slice = tBgB[(None, crd[1], None, crd[2])]
- ab_producer.reset()
- pre = []
- for kk in cutlass.range_constexpr(self.n_pre):
- handle = ab_producer.acquire_and_advance()
- cute.copy(
- tma_atom_b,
- tBgB_slice[(None, kk)],
- tBsB[(None, handle.index)],
- tma_bar_ptr=handle.barrier,
- )
- pre.append(handle)
- cute.arch.griddepcontrol_wait()
- for kk in cutlass.range_constexpr(self.n_pre):
- handle = pre[kk]
- cute.copy(
- tma_atom_a,
- tAgA_slice[(None, kk)],
- tAsA[(None, handle.index)],
- tma_bar_ptr=handle.barrier,
- )
- for k_tile in cutlass.range(self.n_pre, k_tile_cnt, 1, unroll=1):
- handle = ab_producer.acquire_and_advance()
- cute.copy(
- tma_atom_a,
- tAgA_slice[(None, k_tile)],
- tAsA[(None, handle.index)],
- tma_bar_ptr=handle.barrier,
- )
- cute.copy(
- tma_atom_b,
- tBgB_slice[(None, k_tile)],
- tBsB[(None, handle.index)],
- tma_bar_ptr=handle.barrier,
- )
- tile_sched.advance_to_next_work()
- work_tile = tile_sched.get_current_work()
- while work_tile.is_valid_tile:
- crd = work_tile.tile_idx
- tAgA_slice = tAgA[(None, crd[0], None, crd[2])]
- tBgB_slice = tBgB[(None, crd[1], None, crd[2])]
- ab_producer.reset()
- peek = ab_producer.try_acquire()
- for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):
- handle = ab_producer.acquire_and_advance(peek)
- cute.copy(
- tma_atom_a,
- tAgA_slice[(None, handle.count)],
- tAsA[(None, handle.index)],
- tma_bar_ptr=handle.barrier,
- )
- cute.copy(
- tma_atom_b,
- tBgB_slice[(None, handle.count)],
- tBsB[(None, handle.index)],
- tma_bar_ptr=handle.barrier,
- )
- peek = cutlass.Boolean(1)
- if handle.count + 1 < k_tile_cnt:
- peek = ab_producer.try_acquire()
- tile_sched.advance_to_next_work()
- work_tile = tile_sched.get_current_work()
- ab_producer.tail()
-
- if warp_idx == self.nrm_warp_id:
- cute.arch.griddepcontrol_wait()
- halves = self.bh // self.epi_cols
- col_blocks = self.hidden // self.epi_cols
- while work_tile.is_valid_tile:
- crd = work_tile.tile_idx
- for half in cutlass.range_constexpr(halves):
- handle = n_producer.acquire_and_advance()
- for s in cutlass.range_constexpr(self.hc):
- cute.copy(
- tma_atom_n,
- tNgN[(None, crd[0], s * col_blocks + crd[1] * halves + half)],
- tNsN[(None, handle.index * self.hc + s)],
- tma_bar_ptr=handle.barrier,
- )
- tile_sched.advance_to_next_work()
- work_tile = tile_sched.get_current_work()
- n_producer.tail()
-
- if warp_idx == self.mma_warp_id:
- tmem.wait_for_alloc()
- tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
- tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
- acc_producer_state = pipeline.make_pipeline_state(
- pipeline.PipelineUserType.Producer, self.num_acc_stage
- )
- while work_tile.is_valid_tile:
- tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)]
- ab_consumer.reset()
- peek = ab_consumer.try_wait()
- acc_pipeline.producer_acquire(acc_producer_state)
- tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
- for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):
- handle = ab_consumer.wait_and_advance(peek)
- num_kblocks = cute.size(tCrA, mode=[2])
- for kb in cutlass.range(num_kblocks, unroll_full=True):
- kcrd = (None, None, kb, handle.index)
- cute.gemm(tiled_mma, tCtAcc, tCrA[kcrd], tCrB[kcrd], tCtAcc)
- tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
- handle.release()
- peek = cutlass.Boolean(1)
- if handle.count + 1 < k_tile_cnt:
- peek = ab_consumer.try_wait()
- acc_pipeline.producer_commit(acc_producer_state)
- acc_producer_state.advance()
- tile_sched.advance_to_next_work()
- work_tile = tile_sched.get_current_work()
- acc_pipeline.producer_tail(acc_producer_state)
-
- if warp_idx < self.mma_warp_id:
- tmem.allocate(self.num_tmem_alloc_cols)
- tmem.wait_for_alloc()
- tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
- tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
- acc_consumer_state = pipeline.make_pipeline_state(
- pipeline.PipelineUserType.Consumer, self.num_acc_stage
- )
- copy_atom_t2r = utils.sm100.get_tmem_load_op(
- self.cta_tile_shape_mnk,
- utils.LayoutEnum.ROW_MAJOR,
- self.io_dtype,
- self.acc_dtype,
- epi_tile,
- False,
- )
- tAcc_epi = cute.flat_divide(tCtAcc_base[((None, None), 0, 0, None)], epi_tile)
- tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)])
- thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
- tTR_tAcc_base = thr_copy_t2r.partition_S(tAcc_epi)
- tTR_sN = thr_copy_t2r.partition_D(sN)
- tTR_gX = thr_copy_t2r.partition_D(
- cute.flat_divide(tCgX[((None, None), 0, 0, None, None, None)], epi_tile)
- )
- tTR_cX = thr_copy_t2r.partition_D(
- cute.flat_divide(tCcX[((None, None), 0, 0, None, None, None)], epi_tile)
- )
- frag_shape = tTR_gX[(None, None, None, 0, 0, 0, 0, 0)].shape
- tTR_rAcc = cute.make_rmem_tensor(frag_shape, self.acc_dtype)
- tTR_rX = cute.make_rmem_tensor(frag_shape, self.io_dtype)
- tTR_rSum = cute.make_rmem_tensor(frag_shape, self.acc_dtype)
- nfrag = cute.size(tTR_rAcc)
- halves = self.bh // self.epi_cols
- rN = cute.make_rmem_tensor(frag_shape, self.io_dtype)
- inv_hc = cutlass.Float32(1.0 / self.hc)
- if cutlass.const_expr(self.fp4):
- # A thread owns one row's `epi_cols` consecutive columns (the
- # 32dp32b TMEM load), i.e. whole 16-element scale blocks.
- gs = mGS[0]
- rcp_gs = _rcp_approx_ftz(gs)
- rcp6 = _rcp_approx_ftz(cutlass.Float32(6.0))
- rV = cute.make_rmem_tensor((nfrag,), self.acc_dtype)
- assert nfrag == 32, "the NVFP4 epilogue stores two scale blocks per row"
- q_base = mQ.iterator.toint()
- sf_base = mSF.iterator.toint()
-
- while work_tile.is_valid_tile:
- crd = work_tile.tile_idx
- tile_sched.advance_to_next_work()
- work_tile = tile_sched.get_current_work()
-
- tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_consumer_state.index)]
- tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
- gX_t = tTR_gX[(None, None, None, None, None, *crd)]
- gX_t = cute.group_modes(gX_t, 3, cute.rank(gX_t))
- cX_t = tTR_cX[(None, None, None, None, None, *crd)]
- cX_t = cute.group_modes(cX_t, 3, cute.rank(cX_t))
- row_ok = cX_t[(0, 0, 0, 0)][0] < m
-
- acc_pipeline.consumer_wait(acc_consumer_state)
- for half in cutlass.range_constexpr(halves):
- n_handle = n_consumer.wait_and_advance()
- for s in cutlass.range_constexpr(self.hc):
- cute.autovec_copy(
- tTR_sN[(None, None, None, n_handle.index * self.hc + s)], rN
- )
- cute.copy(
- tiled_copy_t2r,
- tTR_tAcc[(None, None, None, s * halves + half)],
- tTR_rAcc,
- )
- for i in cutlass.range_constexpr(nfrag):
- gate = _sigmoid(tTR_rAcc[i])
- if cutlass.const_expr(s == 0):
- tTR_rSum[i] = gate * rN[i].to(self.acc_dtype)
- else:
- tTR_rSum[i] = tTR_rSum[i] + gate * rN[i].to(self.acc_dtype)
- cute.arch.sync_warp()
- n_handle.release()
- for i in cutlass.range_constexpr(nfrag):
- tTR_rX[i] = (tTR_rSum[i] * inv_hc).to(self.io_dtype)
- if row_ok:
- cute.autovec_copy(tTR_rX, gX_t[(None, None, None, half)])
- if cutlass.const_expr(self.fp4):
- row = cX_t[(0, 0, 0, 0)][0]
- blk0 = (crd[1] * self.bh + half * self.epi_cols) // 16
- for i in cutlass.range_constexpr(nfrag):
- rV[i] = tTR_rX[i].to(self.acc_dtype)
- words, sfs = [], []
- for b in cutlass.range_constexpr(2):
- v = [rV[16 * b + j] for j in range(16)]
- amax = _absmax16(v)
- sf = _e4m3_scale(amax, gs, rcp6)
- osc = _e2m1_out_scale(sf, rcp_gs, amax)
- words += [_e2m1x8(v[:8], osc), _e2m1x8(v[8:], osc)]
- sfs.append(sf)
- row64 = cutlass.Int64(row)
- _st_global_v4(
- q_base + row64 * (self.hidden // 2) + cutlass.Int64(blk0) * 8,
- *words,
- )
- _st_global_b16(
- sf_base + row64 * (self.hidden // 16) + cutlass.Int64(blk0),
- sfs[0] | (sfs[1] << 8),
- )
- with cute.arch.elect_one():
- acc_pipeline.consumer_release(acc_consumer_state)
- acc_consumer_state.advance()
-
- cute.arch.barrier(
- barrier_id=self.epilog_sync_bar_id,
- number_of_threads=32 * len(self.epilogue_warp_id),
- )
- tmem.relinquish_alloc_permit()
- tmem.free(tmem_ptr)
-
-
-_COMPILED: dict = {}
-
-
-def _launcher(block_h: int, hc: int, hidden: int, rank: int, ptrs, fp4: bool = False):
- """Compile once per static configuration; later calls reuse the TVM-FFI entry,
- which launches on the caller's current stream (graph-capturable). `ptrs` is
- `(low, up, normed, mixed tensor, xq, sf, gs)`."""
- key = (block_h, hc, hidden, rank, fp4)
- fn = _COMPILED.get(key)
- if fn is None:
- kern = HcUpGateMeanKernel(block_h, hc, hidden, rank, fp4=fp4)
- cptrs = [
- make_ptr(cutlass.BFloat16, p, cute.AddressSpace.gmem, assumed_align=16)
- for p in ptrs[:3]
- ]
- out = from_dlpack(ptrs[3], assumed_align=16).mark_layout_dynamic(leading_dim=1)
- qptrs = [
- make_ptr(cutlass.Int32, ptrs[4], cute.AddressSpace.gmem, assumed_align=16),
- make_ptr(cutlass.Uint8, ptrs[5], cute.AddressSpace.gmem, assumed_align=2),
- make_ptr(cutlass.Float32, ptrs[6], cute.AddressSpace.gmem, assumed_align=4),
- ]
- fn = cute.compile(
- kern,
- 128,
- *cptrs,
- out,
- *qptrs,
- max_active_clusters=utils.HardwareInfo().get_max_active_clusters(1),
- stream=cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
- options="--opt-level 2 --enable-tvm-ffi",
- )
- _COMPILED[key] = fn
- return fn
-
-
-def hc_up_gate_mean(
- low: torch.Tensor, up: torch.Tensor, normed: torch.Tensor, hc_count: int, block_h: int
-) -> torch.Tensor:
- """Launch the fused kernel; `up` is the `[R, C*H]` transposed view of the
- `[C*H, R]` up weight, as `cublas_mm(low, up)` takes it."""
- rows, rank = low.shape
- hidden = normed.shape[1] // hc_count
- mixed = torch.empty((rows, hidden), dtype=normed.dtype, device=normed.device)
- if rows > 0:
- ptrs = (low.data_ptr(), up.data_ptr(), normed.data_ptr())
- # The NVFP4 pointers are never read without `fp4`.
- unused = (mixed.data_ptr(),) * 3
- _launcher(block_h, hc_count, hidden, rank, (*ptrs, mixed, *unused))(
- rows, *ptrs, mixed, *unused
- )
- return mixed
-
-
-def hc_up_gate_mean_fp4(
- low: torch.Tensor,
- up: torch.Tensor,
- normed: torch.Tensor,
- hc_count: int,
- global_scale: torch.Tensor,
- block_h: int,
-):
- """`hc_up_gate_mean` plus the NVFP4 of `mixed` in the same epilogue.
- Returns `(mixed [T, H], xq [T, H/2] uint8, xsf [T*H/16] uint8)`."""
- rows, rank = low.shape
- hidden = normed.shape[1] // hc_count
- dev = normed.device
- mixed = torch.empty((rows, hidden), dtype=normed.dtype, device=dev)
- xq = torch.empty((rows, hidden // 2), dtype=torch.uint8, device=dev)
- xsf = torch.empty((rows * (hidden // 16),), dtype=torch.uint8, device=dev)
- if rows > 0:
- ptrs = (low.data_ptr(), up.data_ptr(), normed.data_ptr())
- qptrs = (xq.data_ptr(), xsf.data_ptr(), global_scale.data_ptr())
- _launcher(block_h, hc_count, hidden, rank, (*ptrs, mixed, *qptrs), fp4=True)(
- rows, *ptrs, mixed, *qptrs
- )
- return mixed, xq, xsf
-
-
-# Hidden columns per tile. 32 (MMA N = 4 * 32 = 128, five A/B stages) measured
-# equal to 64 (N = 256, three stages) at 8K-16K rows and faster at 32..2048 rows.
-_BLOCK_H = 32
-
-
-class Qwen4ExpHcUpGateMean(OpWrapper):
- """With `C = hc_count`, `H = width // C`, `low` the `[T, R]` SiLU bottleneck
- and `up` the `[R, C*H]` transposed view of the `[C*H, R]` up weight::
-
- mixed[t, i] = mean over s of sigmoid((low @ up)[t, s*H + i]) * normed[t, s*H + i]
-
- The GEMM accumulates in fp32 in tensor memory and is never rounded or
- written; sigmoid, product and stream mean are fp32 with one rounding to
- `normed`'s dtype. Returns a fresh `[T, H]` tensor. Fusion boundary: the up
- projection, the sigmoid gate and the stream mean (two launches as `cublas_mm(low,
- up)` followed by a separate gate-and-mean pass).
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: rows, hc_count (4), hidden (2560), rank (320); bf16 only.
- CELLS: tuple[Cell, ...] = tuple(
- Cell(why=f"checkpoint 4x2560, rank 320, {why}", spec=dict(rows=t))
- for t, why in (
- (128, "decode 128 rows, one full tile"),
- (16384, "prefill chunk 16384 (max_num_tokens)"),
- )
- ) + (Cell(why="ragged tail, 16320 + 70 rows (mixed step)", spec=dict(rows=16390)),)
-
- note = """
- Tier 6 (CuTe DSL written for this target): a persistent tcgen05 GEMM whose
- epilogue applies the gate, so the `[T, C*H]` gate logits never reach HBM;
- tiers 1-5 have no GEMM with a custom multi-input epilogue (cuBLAS writes
- the logits). Metadata consumed: none. Call sites:
- the trunk's hyper-connection mixers (`_hc_mixer`). Preconditions (bf16):
- `low` `[T, R]` contiguous, `up` the transposed view of a contiguous
- `[C*H, R]` weight, `normed` `[T, C*H]` contiguous, 16-byte aligned
- pointers, `R % 64 == 0`, `H % 32 == 0`, `C * 32 <= 256`. Compiles once per
- `(block_h, C, H, R)` on first call (~1 s); the row count is a runtime
- argument. Graph-capturable (launches on the current stream, no host sync).
- Launched with programmatic dependent launch (TRTLLM_ENABLE_PDL=0 turns it
- off): each CTA's first tile streams its up-weight stages before
- `griddepcontrol.wait` and `low` / `normed` after it, so behind a producer
- that releases dependents early (`qwen4_exp_hc_decode_tc_first`) the weight
- read overlaps the producer.
-
- Certified: bf16 at the checkpoint's 4 x 2560, rank 320: 128 rows (one full
- tile), 16384 rows, 16390 rows (a partial last tile); graph replay.
-
- Not certified: zero rows, fewer rows than a tile, hc_count other than 4,
- hidden other than 2560, rank other than 320 (the kernel compiles for them;
- no cell drives them).
- """
-
- def __call__(
- self, low: torch.Tensor, up: torch.Tensor, normed: torch.Tensor, hc_count: int
- ) -> torch.Tensor:
- return hc_up_gate_mean(low, up, normed, hc_count, _BLOCK_H)
-
- def reference(
- self, low: torch.Tensor, up: torch.Tensor, normed: torch.Tensor, hc_count: int
- ) -> torch.Tensor:
- """float64, one rounding."""
- hidden = normed.shape[1] // hc_count
- gate = torch.sigmoid(F.linear(low.double(), up.double().t())).unflatten(
- -1, (hc_count, hidden)
- )
- x = normed.double().unflatten(-1, (hc_count, hidden))
- return (gate * x).mean(dim=-2).to(normed.dtype)
-
- def is_valid(
- self, low: torch.Tensor, up: torch.Tensor, normed: torch.Tensor, hc_count: int
- ) -> None:
- rows, rank = low.shape
- width = normed.shape[1]
- assert low.dtype == up.dtype == normed.dtype == torch.bfloat16, "bf16 operands"
- assert normed.shape[0] == rows and up.shape == (rank, width), "shapes"
- assert low.stride() == (rank, 1) and normed.stride() == (width, 1), "contiguous rows"
- assert up.stride() == (1, rank), "up must be the transposed view of a [C*H, R] weight"
- assert width % hc_count == 0 and (width // hc_count) % 32 == 0 and hc_count * 32 <= 256
- assert rank % 64 == 0, "rank must be a multiple of 64"
- for t in (low, up, normed):
- assert t.data_ptr() % 16 == 0, "16-byte aligned operands"
-
-
-qwen4_exp_hc_up_gate_mean = Qwen4ExpHcUpGateMean()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_rstd.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_rstd.py
deleted file mode 100644
index 69fcca405ae8..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_rstd.py
+++ /dev/null
@@ -1,1119 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Hyper-connection mixer second half on the un-normalized residual: the up
-projection as a 2-CTA tcgen05 GEMM whose epilogue gates `hc_out * rstd * (1 + w)`
--- the site's normed stream, formed in registers -- and averages the streams
-(optionally also writing the routed experts' NVFP4 input).
-
-Tier 6 (CuTe DSL written for this target). `qwen4_exp_hc_up_gate_mean`'s epilogue
-on a CTA pair: `tcgen05.mma.cta_group::2` with M = 256 (each CTA 128 rows of
-`low`, resident in shared memory for its row tile) and N = 128 (C streams x 32
-hidden columns), each CTA loading half of the up weight's N per k block. Tiles
-are walked row-tile-major over a contiguous range per pair (one `low` load per
-row tile); the MMA runs up to four tiles ahead of the epilogue (four accumulator
-stages).
-
-The epilogue is two sets of eight warps that take alternate tiles. A load warp
-brings each (tile, stream)'s `hc_out` box (TMA) and `(1 + w)` slice (bulk copy) to
-its own shared-memory stage; a set hands a stream's stage back as soon as it has
-consumed it, so the next tile's boxes load under the current tile's math. Within a
-set, warp w and w + 4 read the same 128 accumulator rows and split the tile's 32
-columns; the gate and the stream mean run on packed fp32 pairs (`f32x2`), and each
-thread stores its 16 outputs as one 32-byte sector.
-"""
-
-from __future__ import annotations
-
-from typing import Optional, Tuple
-
-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._experimental.modeling_v2.catalog.quantization.fp4_quantize import (
- fp4_quantize,
-)
-from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
-from tensorrt_llm._torch.kernels.blackwell.custom_pipeline import PipelineTmaUmma, PipelineUmmaAsync
-from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
-
-from .qwen4_exp_hc_first_rescale_tc import _lds_v4, _sts_v4
-from .qwen4_exp_hc_inject_norm import _bf16x2
-from .qwen4_exp_hc_up_gate_mean import _asm, _rcp_approx_ftz, _st_global_b16, _st_global_v4
-
-#: Programmatic dependent launch (sm_90+), off with TRTLLM_ENABLE_PDL=0.
-_PDL = get_env_enable_pdl()
-
-#: Hidden columns a tile (MMA N = C * 32 = 128) and a thread's epilogue columns.
-_BLOCK_H = 32
-_SUB = 16
-#: Epilogue warp sets (alternate tiles) and warps per set.
-_SETS = 2
-_SET_WARPS = 8
-#: Up-weight ring, accumulator stages, and the hc_out / (1 + w) ring (one stage per
-#: (tile, stream)).
-_B_STAGES = 8
-_ACC_STAGES = 4
-_N_STAGES = 8
-#: NVFP4 block: one e4m3 scale per this many contiguous elements.
-SF_VEC = 16
-
-
-@dsl_user_op
-def _ldg_v4_f32(addr, *, loc=None, ip=None):
- """Four fp32 values at a 16-byte aligned global address."""
- ty = T.f32()
- res = _asm(
- llvm.StructType.get_literal([ty] * 4),
- [cutlass.Int64(addr)],
- "ld.global.v4.f32 {$0, $1, $2, $3}, [$4];",
- "=f,=f,=f,=f,l",
- loc,
- ip,
- )
- return [cutlass.Float32(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(4)]
-
-
-@dsl_user_op
-def _lds_v4_f32(addr, *, loc=None, ip=None):
- """Four fp32 values at a 16-byte aligned shared address (side-effecting: ordered
- after the stage wait).
- """
- ty = T.f32()
- res = _asm(
- llvm.StructType.get_literal([ty] * 4),
- [cutlass.Int32(addr)],
- "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];",
- "=f,=f,=f,=f,r",
- loc,
- ip,
- side_effects=True,
- )
- return [cutlass.Float32(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(4)]
-
-
-@dsl_user_op
-def _bulk_g2s(dst, src, nbytes: int, mbar, *, loc=None, ip=None):
- """One non-tensor bulk copy global -> this CTA's shared memory, completing on
- `mbar`.
- """
- _asm(
- T.i32(),
- [cutlass.Int32(dst), cutlass.Int64(src), cutlass.Int32(mbar)],
- "{ cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [$1], [$2], "
- + str(nbytes)
- + ", [$3]; mov.u32 $0, 0; }",
- "=r,r,l,r",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _st_global_v8(addr, w, *, loc=None, ip=None):
- """One 32-byte `st.global.v8.b32` (sm_100+): a thread's 16 bf16 outputs, one
- full sector.
- """
- _asm(
- T.i32(),
- [cutlass.Int64(addr)] + [cutlass.Int32(x) for x in w],
- "{ st.global.v8.b32 [$1], {$2, $3, $4, $5, $6, $7, $8, $9}; mov.u32 $0, 0; }",
- "=r,l" + ",r" * 8,
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _gate2(z0, z1, xw, wa, wb, hr, a0, a1, first: bool, *, loc=None, ip=None):
- """One column pair of one stream on the packed fp32 pipe (`f32x2`: each lane is
- the scalar `.rn` operation): `t = tanh.approx(z * 0.5)`, `g = fma(t, hr, hr)`,
- `xw = x * w1`, `acc = fma(g, xw, acc)` (`g * xw` for the first stream). `xw` is
- the stream's bf16x2 word (element 2j in the low half), widened exactly by bit
- shifts.
- """
- tail = "mul.rn.f32x2 a, t, x; " if first else "mov.b64 a, {$8, $9}; fma.rn.f32x2 a, t, x, a; "
- text = (
- "{ .reg .b64 z, h, t, x, w, a; .reg .f32 c, t0, t1; .reg .b32 xl, xh; "
- "mov.f32 c, 0f3F000000; mov.b64 h, {c, c}; "
- "mov.b64 z, {$2, $3}; mul.rn.f32x2 z, z, h; "
- "mov.b64 {t0, t1}, z; tanh.approx.f32 t0, t0; tanh.approx.f32 t1, t1; "
- "mov.b64 t, {t0, t1}; mov.b64 h, {$7, $7}; fma.rn.f32x2 t, t, h, h; "
- "shl.b32 xl, $4, 16; and.b32 xh, $4, 0xFFFF0000; mov.b64 x, {xl, xh}; "
- "mov.b64 w, {$5, $6}; mul.rn.f32x2 x, x, w; " + tail + "mov.b64 {$0, $1}, a; }"
- )
- ty = T.f32()
- res = _asm(
- llvm.StructType.get_literal([ty, ty]),
- [
- cutlass.Float32(z0),
- cutlass.Float32(z1),
- cutlass.Int32(xw),
- cutlass.Float32(wa),
- cutlass.Float32(wb),
- cutlass.Float32(hr),
- cutlass.Float32(a0),
- cutlass.Float32(a1),
- ],
- text,
- "=f,=f,f,f,r,f,f,f,f,f",
- loc,
- ip,
- )
- return [cutlass.Float32(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(2)]
-
-
-@dsl_user_op
-def _nvfp4_16(words, gs, rcp_gs, rcp6, *, loc=None, ip=None):
- """NVFP4 of one 16-element block given as 8 bf16x2 words (element 2k in word k's
- low half): `(e2m1 word for elements 0..7, e2m1 word for 8..15, e4m3 scale byte)`.
- `qwen4_exp_hc_up_gate_mean`'s recipe (`_absmax16`, `_e4m3_scale`,
- `_e2m1_out_scale`, `_e2m1x8`) on the same bf16-rounded values, with the max taken
- on the packed bf16 pairs, the bf16 -> fp32 unpack done by bit shifts (both exact)
- and the scaling multiplies on `f32x2` pairs (each lane `mul.rn.f32`): bitwise its
- outputs.
- """
- amax = "abs.bf16x2 m, $3; " + "".join(
- f"abs.bf16x2 t, ${4 + k}; max.bf16x2 m, m, t; " for k in range(7)
- )
- unpack = "".join(
- f"shl.b32 a{2 * k}, ${3 + k}, 16; and.b32 a{2 * k + 1}, ${3 + k}, 0xFFFF0000; "
- for k in range(8)
- )
- muls = "mov.b64 oo, {o, o}; " + "".join(
- f"mov.b64 pp, {{a{2 * k}, a{2 * k + 1}}}; mul.rn.f32x2 pp, pp, oo; "
- f"mov.b64 {{a{2 * k}, a{2 * k + 1}}}, pp; "
- for k in range(8)
- )
- cvts = "".join(f"cvt.rn.satfinite.e2m1x2.f32 b{k}, a{2 * k + 1}, a{2 * k}; " for k in range(8))
- text = (
- "{ .reg .b32 m, t, f2; .reg .b32 a<16>; .reg .f32 x, y, mx, o; .reg .b16 h, l16, h16; "
- ".reg .pred p; .reg .b8 b<8>; .reg .b64 oo, pp; "
- + amax
- + "mov.b32 {l16, h16}, m; cvt.f32.bf16 x, l16; cvt.f32.bf16 y, h16; max.f32 mx, x, y; "
- "mul.rn.f32 x, mx, $13; mul.rn.f32 x, $11, x; "
- "cvt.rn.satfinite.e4m3x2.f32 h, 0f00000000, x; cvt.u32.u16 $2, h; "
- "cvt.rn.f16x2.e4m3x2 f2, h; mov.b32 {l16, h16}, f2; cvt.f32.f16 y, l16; "
- "mul.rn.f32 y, y, $12; rcp.approx.ftz.f32 y, y; setp.neu.f32 p, mx, 0f00000000; "
- "selp.f32 o, y, 0f00000000, p; "
- + unpack
- + muls
- + cvts
- + "mov.b32 $0, {b0, b1, b2, b3}; mov.b32 $1, {b4, b5, b6, b7}; }"
- )
- ty = T.i32()
- res = _asm(
- llvm.StructType.get_literal([ty, ty, ty]),
- [cutlass.Int32(w) for w in words]
- + [cutlass.Float32(gs), cutlass.Float32(rcp_gs), cutlass.Float32(rcp6)],
- text,
- "=r,=r,=r" + ",r" * 8 + ",f,f,f",
- loc,
- ip,
- )
- return [cutlass.Int32(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(3)]
-
-
-class HcUpGateMeanRstdKernel:
- """`mixed[t, i] = 1/C sum_s sigmoid((low @ up)[t, s*H + i]) * hc_out[t, s*H + i]
- * rstd[t, s] * w1[s*H + i]` (+ its NVFP4 with `fp4`).
- """
-
- def __init__(self, hc: int, hidden: int, rank: int, fp4: bool):
- assert rank % 64 == 0 and hidden % _BLOCK_H == 0 and hc == 4
- self.bh = _BLOCK_H
- self.sub = _SUB
- self.hc, self.hidden, self.rank, self.fp4 = hc, hidden, rank, fp4
- self.kb = rank // 64
- self.acc_dtype = cutlass.Float32
- self.io_dtype = cutlass.BFloat16
- self.mma_tiler = (256, hc * self.bh, 64)
- self.cta_group = tcgen05.CtaGroup.TWO
- self.cluster_shape_mn = (2, 1)
- self.epi_sets = _SETS
- self.set_warps = _SET_WARPS
- self.epi_warps = _SETS * _SET_WARPS
- self.mma_warp_id = self.epi_warps
- self.tma_warp_id = self.epi_warps + 1
- self.nrm_warp_id = self.epi_warps + 2
- self.threads_per_cta = 32 * (self.epi_warps + 3)
- self.num_b_stage = _B_STAGES
- self.num_n_stage = _N_STAGES
- self.num_acc_stage = _ACC_STAGES
- # The pair's j-th tile uses accumulator stage j % ACC and, for stream s, the
- # hc_out stage (C * j + s) % NS.
- assert _ACC_STAGES & (_ACC_STAGES - 1) == 0 and _SETS & (_SETS - 1) == 0
- assert _N_STAGES >= hc * _SETS
- self.acc_shift = _ACC_STAGES.bit_length() - 1
- self.set_shift = _SETS.bit_length() - 1
- self.epilog_sync_bar_id = 1
- self.tmem_alloc_sync_bar_id = 2
- # Named barriers 3..3 + 4 * SETS - 1: a set's epilogue warps w and w + 4 (the
- # NVFP4 hand-off).
- self.pair_bar_id0 = 3
-
- def _setup(self):
- k_major = utils.LayoutEnum.ROW_MAJOR.mma_major_mode()
- tiled_mma = utils.sm100.make_trivial_tiled_mma(
- self.io_dtype,
- self.io_dtype,
- k_major,
- k_major,
- self.acc_dtype,
- self.cta_group,
- self.mma_tiler[:2],
- )
- self.cluster_layout_vmnk = cute.tiled_divide(
- cute.make_layout((*self.cluster_shape_mn, 1)), (tiled_mma.thr_id.shape,)
- )
- self.n_box = (128, self.bh)
- # A stage: one stream of this CTA's tile, its hc_out box ([128, 32] bf16) and
- # (1 + w) slice ([32] fp32).
- self.num_n_bytes = 128 * self.bh * 2 + self.bh * 4
- self.a_smem_layout_staged = utils.sm100.make_smem_layout_a(
- tiled_mma, self.mma_tiler, self.io_dtype, self.kb
- )
- self.b_smem_layout_staged = utils.sm100.make_smem_layout_b(
- tiled_mma, self.mma_tiler, self.io_dtype, self.num_b_stage
- )
- 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
- )
- 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,
- hc_ptr: cute.Pointer,
- mixed: cute.Tensor,
- xq_ptr: cute.Pointer,
- sf_ptr: cute.Pointer,
- gs_ptr: cute.Pointer,
- rstd_ptr: cute.Pointer,
- w1_ptr: cute.Pointer,
- max_clusters: cutlass.Constexpr,
- stream: cuda.CUstream,
- ):
- bh, hc, hid, rank = self.bh, self.hc, self.hidden, self.rank
- nt = hid // bh
- width = hc * hid
- a = cute.make_tensor(low_ptr, cute.make_layout((m, rank, 1), stride=(rank, 1, 0)))
- # The `[C*H, R]` up weight viewed as N = (32 hidden, C streams, H/32 slices): an
- # MMA tile's N is one hidden slice of every stream.
- b = cute.make_tensor(
- up_ptr,
- cute.make_layout(((bh, hc, nt), rank, 1), stride=((rank, hid * rank, bh * rank), 1, 0)),
- )
- nrm = cute.make_tensor(hc_ptr, cute.make_layout((m, width), stride=(width, 1)))
- gsc = cute.make_tensor(gs_ptr, cute.make_layout(1))
- tiled_mma = self._setup()
- 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,
- )
- # Both CTAs' A (all k blocks) / B stage land on the leader's barrier.
- self.a_bytes = cute.size_in_bytes(self.io_dtype, a_smem) * self.kb * 2
- self.b_bytes = cute.size_in_bytes(self.io_dtype, b_smem) * 2
- total = ((m + 255) // 256) * nt
- ncl = cutlass.min(total, max_clusters)
- self.kernel(
- tiled_mma,
- tma_atom_a,
- tma_a,
- tma_atom_b,
- tma_b,
- tma_atom_n,
- tma_n,
- gsc,
- mixed.iterator.toint(),
- xq_ptr.toint(),
- sf_ptr.toint(),
- rstd_ptr.toint(),
- w1_ptr.toint(),
- m,
- total,
- self.cluster_layout_vmnk,
- self.a_smem_layout_staged,
- self.b_smem_layout_staged,
- self.n_smem_layout_staged,
- ).launch(
- grid=[2 * ncl, 1, 1],
- 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,
- mGS: cute.Tensor,
- mixed_base: cutlass.Int64,
- q_base: cutlass.Int64,
- sf_base: cutlass.Int64,
- rstd_base: cutlass.Int64,
- w1_base: cutlass.Int64,
- m: cutlass.Int32,
- total: 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,
- ):
- hc, hid, bh, sub = self.hc, self.hidden, self.bh, self.sub
- nt = hid // bh
- 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()
- bid, _, _ = cute.arch.block_idx()
- nblk, _, _ = cute.arch.grid_dim()
- v = bid % 2
- is_leader = v == 0
- cl = bid // 2
- ncl = nblk // 2
- cta_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
- cluster_coord = cluster_layout_vmnk.get_flat_coord(cta_rank)
- # The pair's contiguous range [t0, t1) of (256-row tile, hidden slice) tiles,
- # row-tile-major.
- t0 = (cl * total) // ncl
- t1 = ((cl + 1) * total) // ncl
- m_first = t0 // nt
- m_last = (t1 - 1) // nt
-
- @cute.struct
- class SharedStorage:
- a_mbar: cute.struct.MemRange[cutlass.Int64, 2]
- b_mbar: cute.struct.MemRange[cutlass.Int64, self.num_b_stage * 2]
- acc_mbar: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]
- n_mbar: 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)
- a_producer, a_consumer = PipelineTmaUmma.create(
- barrier_storage=storage.a_mbar.data_ptr(),
- num_stages=1,
- producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
- consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1),
- tx_count=self.a_bytes,
- cta_layout_vmnk=cluster_layout_vmnk,
- ).make_participants()
- b_producer, b_consumer = PipelineTmaUmma.create(
- barrier_storage=storage.b_mbar.data_ptr(),
- num_stages=self.num_b_stage,
- producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
- consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1),
- tx_count=self.b_bytes,
- cta_layout_vmnk=cluster_layout_vmnk,
- ).make_participants()
- # An accumulator stage is consumed by one set's warps in each CTA.
- acc_pipeline = PipelineUmmaAsync.create(
- barrier_storage=storage.acc_mbar.data_ptr(),
- num_stages=self.num_acc_stage,
- producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
- consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 2 * self.set_warps),
- cta_layout_vmnk=cluster_layout_vmnk,
- )
- # hc_out / (1 + w) stages are per CTA (each CTA its own rows), one set's warps
- # consuming each.
- n_pipe = pipeline.PipelineTmaAsync.create(
- barrier_storage=storage.n_mbar.data_ptr(),
- num_stages=self.num_n_stage,
- producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
- consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, self.set_warps),
- tx_count=self.num_n_bytes,
- cta_layout_vmnk=cute.make_layout((1, 1, 1, 1)),
- )
- n_producer = n_pipe.make_participants()[0]
- tmem_alloc_barrier = pipeline.NamedBarrier(
- barrier_id=self.tmem_alloc_sync_bar_id, num_threads=32 * (1 + self.epi_warps)
- )
- tmem = utils.TmemAllocator(
- storage.tmem_holding_buf.ptr,
- barrier_for_retrieve=tmem_alloc_barrier,
- allocator_warp_id=0,
- is_two_cta=True,
- 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=1024,
- swizzle=a_smem_layout_staged.inner,
- )
- sB = smem.allocate_tensor(
- element_type=self.io_dtype,
- layout=b_smem_layout_staged.outer,
- byte_alignment=1024,
- swizzle=b_smem_layout_staged.inner,
- )
- sN = smem.allocate_tensor(
- element_type=self.io_dtype,
- layout=n_smem_layout_staged.outer,
- byte_alignment=1024,
- swizzle=n_smem_layout_staged.inner,
- )
- # NVFP4: a set's upper warp group hands its half of a row's e2m1 / scale bytes to
- # the lower group through shared memory (two tile slots per set), which stores
- # the whole row piece (one 16-byte e2m1 store, one 2-byte scale store).
- sX = smem.allocate_tensor(
- element_type=cutlass.Int32,
- layout=cute.make_layout((2 * self.epi_sets, 128, 4), stride=(512, 4, 1)),
- byte_alignment=128,
- )
- sW = smem.allocate_tensor(
- element_type=cutlass.Float32,
- layout=cute.make_layout((bh, self.num_n_stage), stride=(1, bh)),
- byte_alignment=128,
- )
-
- 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)
- )
- gN_box = cute.local_tile(mN_mn, self.n_box, (None, None))
- thr_mma = tiled_mma.get_slice(v)
- tCgA = thr_mma.partition_A(gA_mkl)
- tCgB = thr_mma.partition_B(gB_nkl)
- a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape)
- tAsA, tAgA = cpasync.tma_partition(
- tma_atom_a,
- cluster_coord[2],
- a_cta_layout,
- cute.group_modes(sA, 0, 3),
- cute.group_modes(tCgA, 0, 3),
- )
- b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape)
- tBsB, tBgB = cpasync.tma_partition(
- tma_atom_b,
- cluster_coord[1],
- b_cta_layout,
- 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),
- )
- a_mcast = cpasync.create_tma_multicast_mask(
- cluster_layout_vmnk, cluster_coord, mcast_mode=2
- )
- b_mcast = cpasync.create_tma_multicast_mask(
- cluster_layout_vmnk, cluster_coord, mcast_mode=1
- )
- 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)
-
- # ---------------------------------------------------------- TMA (A, B): both CTAs
- if warp_idx == self.tma_warp_id:
- cute.arch.griddepcontrol_wait()
- for mi in cutlass.range(m_first, m_last + 1, 1, unroll=1):
- a_handle = a_producer.acquire_and_advance()
- for kb in cutlass.range_constexpr(self.kb):
- cute.copy(
- tma_atom_a,
- tAgA[(None, mi, kb, 0)],
- tAsA[(None, kb)],
- tma_bar_ptr=a_handle.barrier,
- mcast_mask=a_mcast,
- )
- n_lo = cutlass.max(t0 - mi * nt, 0)
- n_hi = cutlass.min(t1 - mi * nt, nt)
- for n in cutlass.range(n_lo, n_hi, 1, unroll=1):
- for kb in cutlass.range_constexpr(self.kb):
- handle = b_producer.acquire_and_advance()
- cute.copy(
- tma_atom_b,
- tBgB[(None, n, kb, 0)],
- tBsB[(None, handle.index)],
- tma_bar_ptr=handle.barrier,
- mcast_mask=b_mcast,
- )
- a_producer.tail()
- b_producer.tail()
-
- # ---------------------------------------------------------- hc_out boxes, (1 + w)
- if warp_idx == self.nrm_warp_id:
- cute.arch.griddepcontrol_wait()
- sw_base = cutlass.Int32(sW.iterator.toint())
- for mi in cutlass.range(m_first, m_last + 1, 1, unroll=1):
- n_lo = cutlass.max(t0 - mi * nt, 0)
- n_hi = cutlass.min(t1 - mi * nt, nt)
- for n in cutlass.range(n_lo, n_hi, 1, unroll=1):
- for s in cutlass.range_constexpr(hc):
- handle = n_producer.acquire_and_advance()
- with cute.arch.elect_one():
- _bulk_g2s(
- sw_base + handle.index * (bh * 4),
- w1_base + cutlass.Int64(s * hid + n * bh) * 4,
- bh * 4,
- cutlass.Int32(handle.barrier.toint()),
- )
- cute.copy(
- tma_atom_n,
- tNgN[(None, 2 * mi + v, s * nt + n)],
- tNsN[(None, handle.index)],
- tma_bar_ptr=handle.barrier,
- )
- n_producer.tail()
-
- # ---------------------------------------------------------- MMA (the leader issues)
- 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_state = pipeline.make_pipeline_state(
- pipeline.PipelineUserType.Producer, self.num_acc_stage
- )
- # Both CTAs walk the tiles (the accumulator state advances in both); the
- # leader issues.
- for mi in cutlass.range(m_first, m_last + 1, 1, unroll=1):
- n_lo = cutlass.max(t0 - mi * nt, 0)
- n_hi = cutlass.min(t1 - mi * nt, nt)
- if is_leader:
- a_handle = a_consumer.wait_and_advance()
- for n in cutlass.range(n_lo, n_hi, 1, unroll=1):
- tCtAcc = tCtAcc_base[(None, None, None, acc_state.index)]
- acc_pipeline.producer_acquire(acc_state)
- tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
- for kb in cutlass.range_constexpr(self.kb):
- handle = b_consumer.wait_and_advance()
- num_kblocks = cute.size(tCrA, mode=[2])
- for kk in cutlass.range(num_kblocks, unroll_full=True):
- cute.gemm(
- tiled_mma,
- tCtAcc,
- tCrA[(None, None, kk, kb)],
- tCrB[(None, None, kk, handle.index)],
- tCtAcc,
- )
- tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
- handle.release()
- acc_pipeline.producer_commit(acc_state)
- acc_state.advance()
- a_handle.release()
- else:
- for n in cutlass.range(n_lo, n_hi, 1, unroll=1):
- acc_state.advance()
- acc_pipeline.producer_tail(acc_state)
-
- # ---------------------------------------------------------- epilogue (each CTA its rows)
- if warp_idx < self.epi_warps:
- cute.arch.griddepcontrol_wait()
- 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)
- # Set `eset` takes the pair's tiles t0 + eset, t0 + eset + SETS, ...; thread
- # (grp, lt) of a set owns row lt of the CTA's 128 and columns
- # [grp * 16, grp * 16 + 16) of the tile's 32.
- eset = warp_idx // self.set_warps
- lt = tidx % 128
- grp = (tidx // 128) % 2
- epi_tile = (128, sub)
- copy_atom_t2r = utils.sm100.get_tmem_load_op(
- (128, self.mma_tiler[1], self.mma_tiler[2]),
- utils.LayoutEnum.ROW_MAJOR,
- self.io_dtype,
- self.acc_dtype,
- epi_tile,
- True,
- )
- 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(lt)
- tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi)
- tTR_sN = thr_copy_t2r.partition_D(cute.flat_divide(sN, epi_tile))
- sw_addr = cutlass.Int32(sW.iterator.toint())
- sx_addr = cutlass.Int32(sX.iterator.toint()) + eset * (2 * 128 * 16)
- frag_shape = tTR_sN[(None, None, None, 0, 0, 0)].shape
- rAcc = cute.make_rmem_tensor(frag_shape, self.acc_dtype)
- rN = cute.make_rmem_tensor(frag_shape, self.io_dtype)
- # sigmoid(z) * r / C = fma(tanh(z / 2), r / 2C, r / 2C): the stream mean and the
- # norm's rstd folded into the gate.
- scale = cutlass.Float32(0.5 / hc)
- if cutlass.const_expr(self.fp4):
- gs = mGS[0]
- rcp_gs = _rcp_approx_ftz(gs)
- rcp6 = _rcp_approx_ftz(cutlass.Float32(6.0))
-
- # The rows' rstd, reloaded when the row tile changes.
- mi_cur = cutlass.Int32(-1)
- rs0 = cutlass.Float32(0.0)
- rs1 = cutlass.Float32(0.0)
- rs2 = cutlass.Float32(0.0)
- rs3 = cutlass.Float32(0.0)
- for t in cutlass.range(t0 + eset, t1, self.epi_sets, unroll=1):
- j = t - t0
- acc_state = pipeline.PipelineState(
- self.num_acc_stage,
- j,
- j & (self.num_acc_stage - 1),
- (j >> self.acc_shift) & 1,
- )
- mi = t // nt
- n = t - mi * nt
- row = mi * 256 + v * 128 + lt
- row_ok = row < m
- if mi != mi_cur:
- rs0, rs1, rs2, rs3 = _ldg_v4_f32(
- rstd_base + cutlass.Int64(cutlass.min(row, m - 1)) * (hc * 4)
- )
- mi_cur = mi
- rs = [rs0, rs1, rs2, rs3]
- col0 = n * bh + grp * sub
- acc_pipeline.consumer_wait(acc_state)
- acc = [cutlass.Float32(0.0)] * sub
- for s in cutlass.range_constexpr(hc):
- cute.copy(
- tiled_copy_t2r,
- tTR_tAcc[(None, None, None, 0, 2 * s + grp, acc_state.index)],
- rAcc,
- )
- # tcgen05.wait::ld: the stream's accumulator columns are in registers;
- # after the last stream the stage goes back before the math.
- cute.arch.fence_view_async_tmem_load()
- if cutlass.const_expr(s == hc - 1):
- cute.arch.sync_warp()
- with cute.arch.elect_one():
- acc_pipeline.consumer_release(acc_state)
- hrs = rs[s] * scale
- q = j * hc + s
- if cutlass.const_expr(self.num_n_stage == hc * self.epi_sets):
- # The set's stream-s stage is fixed; its phase flips every own tile.
- n_state = pipeline.PipelineState(
- self.num_n_stage, q, eset * hc + s, (j >> self.set_shift) & 1
- )
- else:
- n_state = pipeline.PipelineState(
- self.num_n_stage, q, q % self.num_n_stage, (q // self.num_n_stage) & 1
- )
- n_pipe.consumer_wait(n_state)
- cute.autovec_copy(tTR_sN[(None, None, None, 0, grp, n_state.index)], rN)
- waddr = sw_addr + (n_state.index * bh + grp * sub) * 4
- w1 = []
- for c4 in cutlass.range_constexpr(sub // 4):
- w1 += _lds_v4_f32(waddr + c4 * 16)
- xw = cute.recast_tensor(rN, cutlass.Int32)
- for k in cutlass.range_constexpr(sub // 2):
- acc[2 * k], acc[2 * k + 1] = _gate2(
- rAcc[2 * k],
- rAcc[2 * k + 1],
- xw[k],
- w1[2 * k],
- w1[2 * k + 1],
- hrs,
- acc[2 * k],
- acc[2 * k + 1],
- s == 0,
- )
- # Every read of the stage (hc_out, (1 + w)) is performed before it is
- # handed back.
- cute.arch.fence_acq_rel_cta()
- cute.arch.sync_warp()
- n_pipe.consumer_release(n_state) # lane 0 of each warp arrives
- words = [_bf16x2(acc[2 * k], acc[2 * k + 1]) for k in range(sub // 2)]
- if row_ok:
- _st_global_v8(mixed_base + (cutlass.Int64(row) * hid + col0) * 2, words)
- if cutlass.const_expr(self.fp4):
- q0, q1, sf = _nvfp4_16(words, gs, rcp_gs, rcp6)
- slot = sx_addr + (((t // self.epi_sets) % 2) * 128 + lt) * 16
- if grp == 1:
- _sts_v4(slot, q0, q1, sf, sf)
- cute.arch.barrier(
- barrier_id=self.pair_bar_id0 + warp_idx % 4 + 4 * eset,
- number_of_threads=64,
- )
- if grp == 0:
- pq = _lds_v4(slot)
- if row_ok:
- row64 = cutlass.Int64(row)
- _st_global_v4(
- q_base + row64 * (hid // 2) + cutlass.Int64(n * bh // 2),
- q0,
- q1,
- pq[0],
- pq[1],
- )
- _st_global_b16(
- sf_base + row64 * (hid // 16) + cutlass.Int64(n * bh // 16),
- sf | (pq[2] << 8),
- )
-
- cute.arch.barrier(
- barrier_id=self.epilog_sync_bar_id, number_of_threads=32 * self.epi_warps
- )
- tmem.relinquish_alloc_permit()
- tmem.free(tmem_ptr)
-
-
-_COMPILED: dict = {}
-
-
-def _launcher(hc: int, hidden: int, rank: int, fp4: bool, ptrs):
- """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, hc_out, mixed tensor, xq, sf, gs, rstd, norm_w1)`.
- """
- key = (hc, hidden, rank, fp4)
- fn = _COMPILED.get(key)
- if fn is None:
- kern = HcUpGateMeanRstdKernel(hc, hidden, rank, fp4)
- gm = cute.AddressSpace.gmem
- cptrs = [
- make_ptr(cutlass.BFloat16, ptrs[0], gm, assumed_align=16),
- make_ptr(cutlass.BFloat16, ptrs[1], gm, assumed_align=16),
- make_ptr(cutlass.BFloat16, ptrs[2], gm, assumed_align=16),
- from_dlpack(ptrs[3], assumed_align=16).mark_layout_dynamic(leading_dim=1),
- make_ptr(cutlass.Int32, ptrs[4], gm, assumed_align=16),
- make_ptr(cutlass.Uint8, ptrs[5], gm, assumed_align=2),
- make_ptr(cutlass.Float32, ptrs[6], gm, assumed_align=4),
- make_ptr(cutlass.Float32, ptrs[7], gm, assumed_align=16),
- make_ptr(cutlass.Float32, ptrs[8], gm, assumed_align=16),
- ]
- fn = cute.compile(
- kern,
- 128,
- *cptrs,
- max_clusters=utils.HardwareInfo().get_max_active_clusters(2),
- 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_rstd(
- low: torch.Tensor,
- up: torch.Tensor,
- hc_out: torch.Tensor,
- rstd: torch.Tensor,
- norm_w1: torch.Tensor,
- hc_count: int,
- global_scale: Optional[torch.Tensor] = None,
-):
- """`mixed [T, H]`, and with `global_scale` also `xq [T, H/2]` and `xsf [T*H/16]`
- (linear scales).
- """
- rows, rank = low.shape
- hidden = hc_out.shape[1] // hc_count
- dev = hc_out.device
- fp4 = global_scale is not None
- mixed = torch.empty((rows, hidden), dtype=hc_out.dtype, device=dev)
- xq = xsf = None
- if fp4:
- xq = torch.empty((rows, hidden // 2), dtype=torch.uint8, device=dev)
- xsf = torch.empty((rows * (hidden // SF_VEC),), dtype=torch.uint8, device=dev)
- if rows > 0:
- # The NVFP4 pointers are never read without `fp4`.
- qp = (
- (xq.data_ptr(), xsf.data_ptr(), global_scale.data_ptr())
- if fp4
- else (mixed.data_ptr(),) * 3
- )
- ptrs = (
- low.data_ptr(),
- up.data_ptr(),
- hc_out.data_ptr(),
- mixed,
- *qp,
- rstd.data_ptr(),
- norm_w1.data_ptr(),
- )
- _launcher(hc_count, hidden, rank, fp4, ptrs)(rows, *ptrs)
- return (mixed, xq, xsf) if fp4 else mixed
-
-
-def _mixed_reference(low, up, hc_out, rstd, norm_w1, hc_count) -> torch.Tensor:
- """float64: the gate times the unrounded normed stream, stream mean, one rounding."""
- hidden = hc_out.shape[1] // hc_count
- gate = torch.sigmoid(F.linear(low.double(), up.double().t())).unflatten(-1, (hc_count, hidden))
- x = hc_out.double().unflatten(-1, (hc_count, hidden)) * rstd.double().unsqueeze(-1)
- x = x * norm_w1.double().unflatten(-1, (hc_count, hidden))
- return (gate * x).mean(dim=-2).to(hc_out.dtype)
-
-
-def _check(low, up, hc_out, rstd, norm_w1, hc_count) -> None:
- rows, rank = low.shape
- width = hc_out.shape[1]
- assert low.dtype == up.dtype == hc_out.dtype == torch.bfloat16, "bf16 operands"
- assert hc_out.shape[0] == rows and up.shape == (rank, width), "shapes"
- assert low.stride() == (rank, 1) and hc_out.stride() == (width, 1), "contiguous rows"
- assert up.stride() == (1, rank), "up must be the transposed view of a [C*H, R] weight"
- assert rstd.dtype == torch.float32 and rstd.shape == (rows, hc_count), "rstd [T, C] fp32"
- assert rstd.is_contiguous(), "rstd contiguous"
- assert norm_w1.dtype == torch.float32 and norm_w1.shape == (width,), "norm_w1 [C*H] fp32"
- assert norm_w1.is_contiguous(), "norm_w1 contiguous"
- # A row's rstd is one 16-byte load, and the MMA's N is C * 32 = 128.
- assert hc_count == 4, "hc_count 4"
- assert width % hc_count == 0 and (width // hc_count) % _BLOCK_H == 0, "H % 32 == 0"
- assert rank % 64 == 0, "rank must be a multiple of 64"
- for t in (low, up, hc_out, rstd, norm_w1):
- assert t.data_ptr() % 16 == 0, "16-byte aligned operands"
-
-
-class Qwen4ExpHcUpGateMeanRstd(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, rstd the site's
- `[T, C]` fp32 statistic and `norm_w1 = 1 + norm_weight` in fp32 `[C*H]`::
-
- normed[t, s*H + i] = hc_out[t, s*H + i] * rstd[t, s] * norm_w1[s*H + i]
- 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;
- `normed` is fp32 in registers (not rounded to bf16 as a stored normed stream
- is); gate, product and stream mean are fp32 with one rounding to `hc_out`'s
- dtype. Per element and stream, in stream order: `t = tanh.approx(z * 0.5)`,
- `g = fma(t, r, r)` with `r = rstd / 2C`, `xw = hc_out * norm_w1`,
- `acc = fma(g, xw, acc)` (`g * xw` for stream 0), each a round-to-nearest fp32
- operation. Returns a fresh `[T, H]` tensor. Fusion boundary: the up projection,
- the grouped norm's scaling, the sigmoid gate and the stream mean.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: rows. The checkpoint's geometry: hc_count 4, hidden 2560, rank 320, bf16.
- CELLS: tuple[Cell, ...] = tuple(
- Cell(why=f"checkpoint 4x2560, rank 320, {why}", spec=dict(rows=t))
- for t, why in (
- (16384, "prefill chunk 16384 (max_num_tokens)"),
- (2048, "prefill 2048 rows"),
- (16390, "ragged tail, 16320 + 70 rows (mixed step)"),
- (5, "five rows, less than one tile"),
- )
- )
-
- note = """
- Tier 6 (CuTe DSL written for this target): a 2-CTA tcgen05 GEMM with a
- multi-input epilogue (`qwen4_exp_hc_up_gate_mean`'s, plus the norm's
- scaling); no lower tier has one. Metadata consumed: none. Call sites:
- `PrefillTarget`'s attention / GDN and stack-final mixers at or above
- `_HC_RSTD_MIN_ROWS` rows, after `qwen4_exp_hc_first_rescale_tc`, on
- `qwen4_exp_hc_inject_rstd`'s `hc_out` and rstd; `norm_w1` precomputed once
- per site at load.
-
- Preconditions (bf16 activations): `low` `[T, R]` contiguous, `up` the
- transposed view of a contiguous `[C*H, R]` weight, `hc_out` `[T, C*H]`
- contiguous, rstd `[T, C]` and `norm_w1` `[C*H]` fp32 contiguous, 16-byte
- aligned pointers, `hc_count == 4`, `R % 64 == 0`, `H % 32 == 0`. Compiles
- once per `(C, H, R)` on first call; the row count is a runtime argument.
- Graph-capturable (current stream, no host sync). Launched with
- programmatic dependent launch (TRTLLM_ENABLE_PDL=0 turns it off).
-
- Launch: persistent, one cluster (CTA pair) per two SMs, each walking a
- contiguous row-tile-major range of (256-row tile, 32-column slice) tiles;
- 608 threads a CTA (two sets of eight epilogue warps on alternate tiles, a
- weight-load, an hc_out-load and an MMA warp), about 217 KB of shared memory
- (low tile, 8-stage weight ring, 8 per-stream hc_out stages) and all 512
- tensor-memory columns (4 accumulator stages). GB300, 16 calls chained over
- distinct weights: 4096 / 8192 / 16384 rows about 24 / 45 / 89 us.
-
- Numerics: the gate is `tanh.approx`-based (as `qwen4_exp_hc_up_gate_mean`'s,
- max relative error 2^-10.99) with rstd / C folded into it; against the float64
- value of the unrounded normed stream the result sits in the gate band at
- about 0.1-0.25 of it (the stored-normed op: 0.3), and is not bitwise
- `qwen4_exp_hc_up_gate_mean`'s (that op gates the bf16-rounded normed). The
- fp32 sequence above is bitwise what the GPU test emulates (zero logits).
-
- Certified: bf16 at the checkpoint's 4 x 2560, rank 320: 16384, 2048, 5
- rows and a partial last tile (16390); gate band (`rtol = atol / max|ref| =
- eps`) against the float64 reference; the epilogue's fp32 sequence bitwise
- (5, 300, 4100 rows); graph replay.
-
- Not certified: zero rows, hc_count other than 4, hidden other than 2560,
- rank other than 320.
- """
-
- def __call__(
- self,
- low: torch.Tensor,
- up: torch.Tensor,
- hc_out: torch.Tensor,
- rstd: torch.Tensor,
- norm_w1: torch.Tensor,
- hc_count: int,
- ) -> torch.Tensor:
- return hc_up_gate_mean_rstd(low, up, hc_out, rstd, norm_w1, hc_count)
-
- def reference(
- self,
- low: torch.Tensor,
- up: torch.Tensor,
- hc_out: torch.Tensor,
- rstd: torch.Tensor,
- norm_w1: torch.Tensor,
- hc_count: int,
- ) -> torch.Tensor:
- """float64, one rounding."""
- return _mixed_reference(low, up, hc_out, rstd, norm_w1, hc_count)
-
- def is_valid(
- self,
- low: torch.Tensor,
- up: torch.Tensor,
- hc_out: torch.Tensor,
- rstd: torch.Tensor,
- norm_w1: torch.Tensor,
- hc_count: int,
- ) -> None:
- _check(low, up, hc_out, rstd, norm_w1, hc_count)
-
-
-class Qwen4ExpHcUpGateMeanRstdFp4(OpWrapper):
- """`qwen4_exp_hc_up_gate_mean_rstd` at an MoE site, whose epilogue also writes
- the routed experts' NVFP4 input::
-
- mixed = qwen4_exp_hc_up_gate_mean_rstd(low, up, hc_out, rstd, norm_w1, C)
- xq, xsf = NVFP4(mixed) # e2m1 pairs, e4m3 scale per 16, linear layout
-
- `mixed` is bitwise the bf16 entry's (same kernel, same epilogue arithmetic);
- the quantization reads that bf16-rounded value with
- `qwen4_exp_hc_up_gate_mean_fp4`'s recipe. 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. Fusion boundary: the up projection, the norm's scaling,
- the sigmoid gate, the stream mean and the NVFP4 quantization of the mean.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: rows. The checkpoint's geometry: hc_count 4, hidden 2560, rank 320, bf16.
- CELLS: tuple[Cell, ...] = tuple(
- Cell(why=f"checkpoint 4x2560, rank 320, NVFP4, {why}", spec=dict(rows=t))
- for t, why in (
- (16384, "prefill chunk 16384 (max_num_tokens)"),
- (2048, "prefill 2048 rows"),
- (16390, "ragged tail, 16320 + 70 rows (mixed step)"),
- (5, "five rows, less than one tile"),
- )
- )
-
- note = """
- Tier 6 (CuTe DSL written for this target): `qwen4_exp_hc_up_gate_mean_rstd`'s
- kernel compiled with its NVFP4 epilogue. Metadata consumed: none. Call
- sites: `PrefillTarget`'s 48 MoE-site mixers at or above `_HC_RSTD_MIN_ROWS`
- rows; `xq` / `xsf` feed the routed experts (linear scales), `mixed` the
- router and the shared expert. Preconditions: those of
- `qwen4_exp_hc_up_gate_mean_rstd` and `global_scale` one fp32 element on the
- device. Compiles once per `(C, H, R)`; graph-capturable.
-
- Each thread holds 16 of a row's 32 tile columns (one scale block); within an
- epilogue set the upper warp group hands its e2m1 word pair and scale byte to
- the lower one through shared memory, which stores the row's 16 e2m1 bytes and
- 2 scale bytes in one store each. GB300, 16 calls chained over distinct
- weights: 4096 / 8192 / 16384 rows about 26 / 53 / 100 us.
-
- Numerics: `mixed` as the bf16 entry. The quantization is the trtllm NVFP4
- recipe instruction for instruction (`rcp.approx.ftz.f32` reciprocals,
- `cvt.rn.satfinite` to e4m3 and to e2m1, every multiply `mul.rn.f32`); the
- GPU test checks it against `fp4_quantize`'s native-torch reference on the
- returned `mixed` (scales bitwise, e2m1 codes bitwise outside the reciprocal
- tie window).
-
- Certified: bf16 at the checkpoint's 4 x 2560, rank 320, a calibrated
- global scale: 16384, 2048, 5 rows and a partial last tile (16390); graph
- replay.
-
- Not certified: zero rows, hc_count other than 4, hidden other than 2560,
- rank other than 320.
- """
-
- def __call__(
- self,
- low: torch.Tensor,
- up: torch.Tensor,
- hc_out: torch.Tensor,
- rstd: torch.Tensor,
- norm_w1: torch.Tensor,
- hc_count: int,
- global_scale: torch.Tensor,
- ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
- return hc_up_gate_mean_rstd(low, up, hc_out, rstd, norm_w1, hc_count, global_scale)
-
- def reference(
- self,
- low: torch.Tensor,
- up: torch.Tensor,
- hc_out: torch.Tensor,
- rstd: torch.Tensor,
- norm_w1: 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 = _mixed_reference(low, up, hc_out, rstd, norm_w1, hc_count)
- return (mixed,) + fp4_quantize.reference(mixed, global_scale, SF_VEC, False, False)
-
- def is_valid(
- self,
- low: torch.Tensor,
- up: torch.Tensor,
- hc_out: torch.Tensor,
- rstd: torch.Tensor,
- norm_w1: torch.Tensor,
- hc_count: int,
- global_scale: torch.Tensor,
- ) -> None:
- _check(low, up, hc_out, rstd, norm_w1, hc_count)
- assert (hc_out.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_rstd = Qwen4ExpHcUpGateMeanRstd()
-qwen4_exp_hc_up_gate_mean_rstd_fp4 = Qwen4ExpHcUpGateMeanRstdFp4()
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
deleted file mode 100644
index d51ddde1e610..000000000000
--- 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
+++ /dev/null
@@ -1,138 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Hyper-connection post-block injection: a per-stream gated accumulate of one
-block output into the multi-stream residual.
-
-Tier-3 composed from native torch: one kernel-launching call
-(``torch.addcmul``) plus view glue. No ``torch.ops.trtllm.*`` op computes this
-accumulate with an *identity* cross-stream mix: ``mhc_post_mapping`` applies a
-Sinkhorn doubly-stochastic mix across streams, fixes the stream count at 4 and
-instantiates its MMA paths for hidden sizes this checkpoint does not use. The
-built-in Qwen4-Exp implementation has its own fused op for it, not wrapped here.
-"""
-
-from __future__ import annotations
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: This checkpoint's hyper-connection geometry: 4 streams of 2560, packed
-#: stream-major into a 10240-wide residual.
-CKPT_HC_COUNT = 4
-CKPT_HIDDEN = 2560
-
-
-class Qwen4ExpPostInject(OpWrapper):
- """`out[..., s * H + i] = hyper_input[..., s * H + i] + block_out[..., i] *
- inject[..., s]`, with `H = block_out.shape[-1]` and `hc_count =
- inject.shape[-1]`: one block output accumulated into **every** stream, each
- scaled by its own gate. Returns a fresh contiguous tensor of `hyper_input`'s
- shape; no input is mutated.
-
- `hyper_input` is the **pre-norm** residual, packed **stream-major** (last
- dim is `(hc_count, H)`, not `(H, hc_count)` -- both have the same shape, so
- a stream-minor caller is silently wrong). The cross-stream mix is the
- identity: no stream's value enters another.
-
- Fusion boundary: the accumulate only. The grouped norm, the input mixer
- that produced the block's input and `inject` (`qwen4_exp_gated_residual`),
- and the block itself are the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: leading (token dims), hc_count, hidden, dtype; layout
- # "contiguous" (default) | "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 (
- (64, "decode batch / MTP rows, 64 tokens"),
- (8192, "prefill 8192 tokens, largest measured"),
- )
- ) + (
- Cell(
- why="hyper_input a column slice buf[:, :width]",
- spec=dict(leading=(8,), hc_count=CKPT_HC_COUNT, hidden=CKPT_HIDDEN, layout="col_slice"),
- ),
- )
-
- note = """
- Metadata consumed: none. Call sites: twice per decoder layer (after the
- attention/GDN block and after the MoE) and twice in the MTP layer, with
- `inject` the `[T, 4]` weights `qwen4_exp_gated_residual` returned for that
- block -- in `[0, 2]`, because they are `2 * sigmoid(.)`.
-
- Preconditions, all loud except the guarded one: the packed width must equal
- `inject.shape[-1] * block_out.shape[-1]` (the unflatten raises); all three
- on one CUDA device (raises). The leading dims must be **equal**, not merely
- broadcastable -- a 1-D `inject` or `block_out` was measured to be accepted
- and to broadcast one row over every token (`is_valid`).
-
- Mixed dtypes are not rejected, they promote: an fp32 `inject` beside a
- bf16 stream returns an fp32 result -- correct values, twice the bytes.
- Recorded, not guarded.
-
- Rounding: `addcmul` forms product and sum in fp32 opmath and rounds once
- on the store. Measured bit-exact against an fp32 reference rounded once at
- 1..8192 tokens; the bf16-multiply-then-bf16-add form differs by 0.0625.
- Wrong operators at `[64, 10240]` bf16, max deviation from the bit-exact
- result: stream-minor packing 10.34, no gate 3.77, gate on the residual 5.89.
-
- CUDA-graph capture certified; the output is a fresh allocation, so under
- capture it comes from the graph's private pool. Cost, not gated: 17 us at
- 64 tokens, 53 us at 2048, 213 us at 8192 (~1.78 TB/s) -- the stride-0
- broadcast operands take a non-vectorized path. An in-place form was not
- measured.
-
- Certified: bf16 2-D operands at the checkpoint's 4 x 2560, 64 and 8192
- rows contiguous, 8 rows with `hyper_input` a column slice; graph replay.
-
- Not certified: other row counts, a row-slice `hyper_input`, hc_count other
- than 4, hidden other than 2560, float16, float32 and float8 operands,
- leading dims beyond one. Hard limits: not
- probed (the largest measured operand is 8192 x 10240 = 2^26.3 elements).
- """
-
- def __call__(
- self,
- hyper_input: torch.Tensor,
- block_out: torch.Tensor,
- inject: torch.Tensor,
- ) -> torch.Tensor:
- streams = hyper_input.unflatten(-1, (inject.shape[-1], block_out.shape[-1]))
- return torch.addcmul(streams, block_out.unsqueeze(-2), inject.unsqueeze(-1)).flatten(-2)
-
- def reference(
- self,
- hyper_input: torch.Tensor,
- block_out: torch.Tensor,
- inject: torch.Tensor,
- ) -> torch.Tensor:
- """fp32 product and sum, rounded once to the promoted dtype."""
- hc_count, hidden = inject.shape[-1], block_out.shape[-1]
- streams = hyper_input.float().unflatten(-1, (hc_count, hidden))
- gated = block_out.float().unsqueeze(-2) * inject.float().unsqueeze(-1)
- out_dtype = torch.promote_types(
- torch.promote_types(hyper_input.dtype, block_out.dtype), inject.dtype
- )
- return (streams + gated).flatten(-2).to(out_dtype)
-
- def is_valid(
- self,
- hyper_input: torch.Tensor,
- block_out: torch.Tensor,
- inject: torch.Tensor,
- ) -> None:
- # Measured: a 1-D `inject` or `block_out` is accepted and returns a full
- # `[8, 10240]` result computed from one row shared by every token.
- assert hyper_input.shape[:-1] == block_out.shape[:-1] == inject.shape[:-1], (
- f"leading dims must match exactly; got hyper_input {tuple(hyper_input.shape)}, "
- f"block_out {tuple(block_out.shape)}, inject {tuple(inject.shape)}"
- )
-
-
-qwen4_exp_post_inject = Qwen4ExpPostInject()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/__init__.py
deleted file mode 100644
index 52a7a9daf028..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/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
deleted file mode 100644
index 9ba3b8122735..000000000000
--- 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
+++ /dev/null
@@ -1,484 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""CuteDSL NVFP4 FC1 of the routed experts: gather, grouped GEMM, SwiGLU and NVFP4
-requantization of the intermediate, in one trtllm call."""
-
-from __future__ import annotations
-
-from typing import NamedTuple, Optional
-
-import torch
-
-import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import SWIGLU_LIMIT_SCALAR_DISABLED
-from tensorrt_llm._torch.utils import ActivationType
-
-_SWIGLU = int(ActivationType.Swiglu)
-
-#: e2m1 magnitudes by 3-bit code, and the midpoints between neighbours. A value
-#: exactly on a midpoint rounds to the even code (`_TIE_UP[i]`: the even side is up).
-E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
-E2M1_MAX = 6.0
-E4M3_MAX = 448.0
-_E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0)
-_TIE_UP = (False, True, False, True, False, True, False)
-
-#: The pairing the fused SwiGLU epilogue reads: the weight's rows `[128 g, 128 g + 64)`
-#: are up rows, the next 64 their gate rows, and output column `64 g + r` pairs row
-#: `128 g + r` with row `128 g + 64 + r`. Fixed by the epilogue's 64-column subtile,
-#: whatever MMA N tile the autotuner picks.
-GATE_UP_GROUP = 64
-
-
-def _ckpt(
- tokens: int,
- tile_size: int,
- mma_n: int = 128,
- dist: str = "router",
- gsf: str = "calibrated",
- zero_tokens: int = 0,
-) -> dict:
- """The checkpoint's routed experts: 512 experts, top-10, hidden 2560, I 640, all
- local (tp1), SwiGLU, no clamp, 16-element e4m3 block scales.
-
- `mma_n` is the autotuner tactic's MMA N tile (128 is also the fallback when the
- tuner has no entry). `dist` names how the test draws the expert ids: `router` =
- each row's top-10 of uniform scores. `gsf` sets `global_sf` against the
- activations' magnitude: `calibrated` = the activation maximum the test
- calibrates to maps to the e4m3 maximum, as the checkpoint's per-layer `1 /
- down_proj.input_scale` does; `subnormal` = 2^-12 of it, so block scales fall
- into e4m3 subnormals and zero. `zero_tokens` makes that many leading tokens
- all-zero rows (zero block scales).
- """
- return dict(
- tokens=tokens,
- tile_size=tile_size,
- mma_n=mma_n,
- dist=dist,
- gsf=gsf,
- zero_tokens=zero_tokens,
- num_experts=512,
- top_k=10,
- hidden=2560,
- intermediate=640,
- )
-
-
-class Fc1Expected(NamedTuple):
- """What `reference` returns: the op's two outputs, plus what the test's compare reads.
-
- `c` / `c_sf` are laid out exactly as the op returns them; rows the op does not
- write hold zeros. `rows` marks the permuted rows the op writes and the compare
- gates. `scaled` is `y * global_sf` on those rows, before rounding -- the value
- the block scale and the codes are derived from, which rounding loses.
- """
-
- c: torch.Tensor
- c_sf: torch.Tensor
- rows: torch.Tensor
- scaled: torch.Tensor
-
-
-class CuteDslNvfp4GatherGroupedGemmActFusion(OpWrapper):
- """`torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell`:
- FC1 of the routed experts on the CuteDSL MoE path, one kernel launch.
-
- For every permuted row `p` the sort filled -- tile `i = p // tile_size` below
- `num_non_exiting_tiles` and `p < tile_idx_to_mn_limit[i]`:
-
- t = permuted_idx_to_expanded_idx[p] // top_k the token
- e = tile_idx_to_group_idx[i] the local expert
- acc = alpha[e] * deq(input[t]) @ deq(weight[e]).T [2 I], fp32
- y[p, j] = up * silu(gate), up = acc[128 g + r], gate = acc[128 g + 64 + r],
- j = 64 g + r
- out[p], out_sf[p] = NVFP4(y[p] * global_sf) 16-element blocks
-
- where `deq` is e2m1 times its e4m3 block scale, and NVFP4 takes per block
- `sf = e4m3(max|v| / 6)` (saturating at 448) and codes `e2m1(v / sf)` (saturating
- at 6). So `out * out_sf / global_sf` approximates `y`.
-
- Arguments, as the CuteDSL MoE passes them (`run_moe_nvfp4_impl`,
- `_torch/moe/fused_moe/fused_moe_cute_dsl.py`):
-
- - `input` `[T, H/2]` float4_e2m1fn_x2 and `input_scale` `[T, H/16]` uint8: the
- un-permuted activation as `fp4_quantize(x, 1 / input_scale_1, 16, False,
- False)` returns it -- LINEAR (row-major) e4m3 scales, viewed 2-D. Gathered
- inside; the caller never permutes it.
- - `weight` `[E, 2I, H/2]` float4_e2m1fn_x2: per expert `cat([up, gate])`
- interleaved in `GATE_UP_GROUP`-row groups (up rows 0:64, gate rows 0:64, up
- rows 64:128, ...). No 32-row shuffle (that is trtllm-gen's layout).
- - `weight_scale` `[E, 2I, H/16]` uint8: the e4m3 block scales with the same row
- interleave, then 128x4-swizzled per expert (`block_scale_interleave`).
- - `alpha` `[E]` fp32: per expert, `input_scale_1 * weight_scale_2[e]` -- one
- factor for both halves, applied to the accumulator before SwiGLU.
- - the four sort tensors, exactly as `moe_sort` returns them for the same
- `top_k` and `tile_tokens_dim == tile_size`.
- - `global_sf` fp32, one element: the FC2 input's global scale, `1 /
- max(down_proj.input_scale)` of the layer.
- - `num_experts`, `num_local_experts`, `local_expert_offset`: read only by the
- autotuner's shape buckets and cache key; the kernel takes the expert count
- from `weight.size(0)` and the expert from the tile.
-
- Returns `(c, c_sf)`: `c` `[M, I/2]` float4_e2m1fn_x2 and `c_sf` `[M * I/16]`
- uint8, `M = permuted_idx_to_expanded_idx.numel()`, one row per PERMUTED row and
- the block scales 128x4-swizzled over `[M, I/16]` -- exactly the `input` /
- `input_scale` the plain NVFP4 grouped GEMM (FC2) takes, with the same four sort
- tensors. Rows the sort did not fill are not part of the result (see `note`).
-
- Fusion boundary: the gather of the activation rows, the grouped GEMM, the
- per-expert alpha, SwiGLU and the requantization happen inside. The caller owns
- the activation quantization, the sort, FC2 and the routing-weighted combine;
- routing weights are not read here.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="tile 128: decode, 1 request without MTP",
- spec=_ckpt(1, 128),
- ),
- Cell(
- why="tile 128: 8192 rows, max_num_tokens prefill",
- spec=_ckpt(8192, 128),
- ),
- Cell(
- why="tile 128: tiny global_sf and 7 all-zero token rows; subnormal and zero "
- "block scales",
- spec=_ckpt(2048, 128, gsf="subnormal", zero_tokens=7),
- ),
- )
-
- note = """
- Path certified: the CuTe DSL kernel `BlockScaledContiguousGatherGroupedGemmKernel`
- (`_torch/moe/kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py`),
- built by `Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner`
- (`_torch/custom_ops/cute_dsl_custom_ops.py`), which JIT-compiles it with
- `cute.compile` on the first call per (tile, tactic, top_k, activation) and
- keeps it for the process. The op is registered only when the CuTe DSL imports
- (`IS_CUTLASS_DSL_AVAILABLE`), and the runner accepts SM 100 and 103 only.
-
- Metadata consumed: none beyond the arguments. Stateless: allocates `c` and
- `c_sf` on every call, on the current stream, with no host synchronization, so
- it captures into a CUDA graph together with moe_sort.
-
- Call sites: the engine's CuteDSL MoE (`run_moe_nvfp4_impl`) and the target's
- routed-expert step, eager for prefill and mixed batches (T up to
- max_num_tokens 8192) and captured for decode (T = requests x MTP verify width,
- up to 32 x 4 = 128). Tile 128 is the target's and the cells' tile; 256 is
- certified through the per-pair invariance below.
-
- Tactics. The op asks the AutoTuner for an MMA tile of (tile_size, 128) or
- (tile_size, 256) and with no tuning entry falls back to (tile_size, 128);
- serving warmup tunes per row bucket and may pick either. Every pair's codes
- and block scales are bitwise identical across tile 128 / 256 x MMA N 128 / 256
- x the fallback (the test checks it at 8192 rows), so cells driven at one
- tactic certify the others.
-
- Determinism. A (token, expert) pair's output row depends on that pair alone:
- it is bitwise unchanged by reordering an expert's rows (moe_sort orders them by
- atomics from T = 5), by repeating the call, and by cutting a batch to a prefix.
-
- What is written: the rows the sort filled hold the result; the padding rows of
- a written tile -- from its tile_idx_to_mn_limit to its end -- get zero block
- scales and zero codes, so a plain grouped GEMM over whole tiles (FC2) reads
- zeros there; the tiles at and past num_non_exiting_tiles are not written at
- all, and under a graph keep whatever the buffer last held. FC2 must be tiled by
- the same tile arrays, and the combine reads only
- expanded_idx_to_permuted_idx's rows.
-
- Scale range: with a tiny global_sf the block scales fall to e4m3 subnormals
- and to zero (a cell). An all-zero token gives scale 0 and codes of magnitude
- 0 (the same cell). With global_sf 25x the calibrated value the block scales
- clip at 448 and the codes at 6 (no NaN) -- measured, not certified. A block
- whose scale is 0 dequantizes to zero whatever its codes hold, and the test's
- compare does not read them.
-
- Preconditions the caller owns (a wrong preparation of each operand -- the
- weight's rows, its scales, the activation's scales, alpha -- runs without
- error, with finite scales, and puts most block scales and codes outside the
- test's gate):
- - `weight` / `weight_scale` in the layout the class docstring states: not
- trtllm-gen's row interleave, up rows first, scales swizzled after the same
- row interleave as the weight, alpha not shifted by an expert.
- - One `alpha` per expert scales BOTH halves, so up_proj and gate_proj must
- share `weight_scale_2`, and `input_scale` must be one value per layer; this
- checkpoint satisfies both. `global_sf` inverts the largest
- `down_proj.input_scale`; the spread of the rest is FC2's `alpha`.
- - `input_scale` LINEAR. Swizzled scales have the same size when T % 128 == 0
- and are read as garbage without error.
- - The sort from moe_sort with the same `top_k` and tile_tokens_dim ==
- `tile_size`, over the tokens `input` holds. Nothing checks `top_k` against it.
-
- Refused by `is_valid`, each silently wrong unguarded: a non-contiguous input,
- input_scale, weight, weight_scale or alpha (data_ptr() read as dense); an
- alpha shorter than weight.size(0); a weight_scale with fewer bytes than the
- weight needs (read past its end: zeros of the same allocation come back as zero
- experts, a fresh allocation faults).
-
- Loud (AssertionError from the runner before the launch): tile_size 256 with a
- tile-128 sort and 128 with a tile-256 sort; input_scale not H/16 wide;
- global_sf of 2 elements; alpha not fp32; weight_scale typed float8_e4m3fn;
- input typed uint8. tile_size 64: ValueError.
-
- Measured at bring-up and not re-driven by the test: the CUDA-graph capture
- with moe_sort, the determinism and write statements, the wrong preparations'
- distance from the gate, the short-alpha and short-weight_scale refusals and
- the loud rejections. The test re-drives the cells, the tile / tactic
- invariance and the `input` contiguity refusal.
-
- Uncertified: Relu2 and SiTU; a SwiGLU clamp (swiglu_limit_scalar >= 0);
- expert_counts (DeepEP count-native metadata); expert parallelism (the kernel
- reads neither num_experts nor local_expert_offset, so EP changes only the sort
- and the weight slice); other hidden / intermediate sizes and top_k;
- scaling_vector_size other than 16; other architectures.
-
- Hard limits (sm_103)
- written rows x I/16 B of c_sf boundary 2^31 B: an illegal memory access
- once rows written reach 53687092 at I = 640 (T ~ 5.37M at
- top-10; T = 8192 writes at most 212480 rows)
- input, input_scale, weight past 2^31 elements / bytes: exact (probed
- up to 1536 experts)
- permuted index int32, within moe_sort's 2^31-pair limit; no index passes
- through fp32
- """
-
- def __call__(
- self,
- input: torch.Tensor,
- weight: torch.Tensor,
- input_scale: torch.Tensor,
- weight_scale: torch.Tensor,
- alpha: torch.Tensor,
- tile_idx_to_group_idx: torch.Tensor,
- tile_idx_to_mn_limit: torch.Tensor,
- permuted_idx_to_expanded_idx: torch.Tensor,
- num_non_exiting_tiles: torch.Tensor,
- global_sf: torch.Tensor,
- num_experts: int,
- top_k: int,
- num_local_experts: int,
- local_expert_offset: int,
- tile_size: int,
- scaling_vector_size: int = 16,
- activation_type: int = _SWIGLU,
- swiglu_limit_scalar: float = SWIGLU_LIMIT_SCALAR_DISABLED,
- expert_counts: Optional[torch.Tensor] = None,
- expert_capacity: int = 0,
- situ_beta: Optional[float] = None,
- situ_linear_beta: Optional[float] = None,
- ) -> tuple[torch.Tensor, torch.Tensor]:
- return torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell(
- input=input,
- weight=weight,
- input_scale=input_scale,
- weight_scale=weight_scale,
- alpha=alpha,
- tile_idx_to_group_idx=tile_idx_to_group_idx,
- tile_idx_to_mn_limit=tile_idx_to_mn_limit,
- permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx,
- num_non_exiting_tiles=num_non_exiting_tiles,
- global_sf=global_sf,
- num_experts=num_experts,
- top_k=top_k,
- num_local_experts=num_local_experts,
- local_expert_offset=local_expert_offset,
- tile_size=tile_size,
- scaling_vector_size=scaling_vector_size,
- activation_type=activation_type,
- swiglu_limit_scalar=swiglu_limit_scalar,
- expert_counts=expert_counts,
- expert_capacity=expert_capacity,
- situ_beta=situ_beta,
- situ_linear_beta=situ_linear_beta,
- )
-
- def reference(
- self,
- input: torch.Tensor,
- weight: torch.Tensor,
- input_scale: torch.Tensor,
- weight_scale: torch.Tensor,
- alpha: torch.Tensor,
- tile_idx_to_group_idx: torch.Tensor,
- tile_idx_to_mn_limit: torch.Tensor,
- permuted_idx_to_expanded_idx: torch.Tensor,
- num_non_exiting_tiles: torch.Tensor,
- global_sf: torch.Tensor,
- num_experts: int,
- top_k: int,
- num_local_experts: int,
- local_expert_offset: int,
- tile_size: int,
- scaling_vector_size: int = 16,
- activation_type: int = _SWIGLU,
- swiglu_limit_scalar: float = SWIGLU_LIMIT_SCALAR_DISABLED,
- expert_counts: Optional[torch.Tensor] = None,
- expert_capacity: int = 0,
- situ_beta: Optional[float] = None,
- situ_linear_beta: Optional[float] = None,
- ) -> Fc1Expected:
- """Gather, dequantize, GEMM in float64, SwiGLU, exact NVFP4 rounding.
-
- Products of e2m1 and e4m3 values are exact, and float64 sums of 2560 of
- them are exact to ~1e-16 relative, so this side carries no accumulation
- error of its own; the block scale and codes are the exact round-to-nearest-
- even of the exact `y * global_sf`. Where a block scale rounds to zero the
- codes here are zero: they dequantize to zero whatever they hold, and the
- test's compare does not read them.
-
- Certified surface only: SwiGLU without a clamp, the sort-metadata path
- (`expert_counts` None), no SiTU. Anything else raises NotImplementedError.
- `num_experts`, `num_local_experts`, `local_expert_offset` and
- `expert_capacity` do not enter the computation.
- """
- if (
- int(activation_type) != _SWIGLU
- or swiglu_limit_scalar >= 0
- or expert_counts is not None
- or situ_beta is not None
- or situ_linear_beta is not None
- ):
- raise NotImplementedError(
- "reference covers SwiGLU without a clamp on sort metadata only"
- )
- vec = scaling_vector_size
- dev = input.device
- m = permuted_idx_to_expanded_idx.numel()
- n2 = weight.shape[1]
- inter = n2 // 2
- hidden = input.shape[1] * 2
- n_tiles = int(num_non_exiting_tiles.reshape(-1)[0])
- span = n_tiles * tile_size
-
- rows = torch.zeros(m, dtype=torch.bool, device=dev)
- limit = tile_idx_to_mn_limit[:n_tiles].long().repeat_interleave(tile_size)
- rows[:span] = torch.arange(span, device=dev) < limit
- real = torch.nonzero(rows).squeeze(1)
- token = permuted_idx_to_expanded_idx.long()[real] // top_k
- expert = tile_idx_to_group_idx[:n_tiles].long().repeat_interleave(tile_size)[real]
-
- a = _dequant(input.view(torch.uint8)[token], input_scale.view(torch.uint8)[token], vec)
- y = torch.empty(real.numel(), inter, dtype=torch.float64, device=dev)
- w_sf = weight_scale.view(torch.uint8).reshape(weight.shape[0], -1)
- group = GATE_UP_GROUP
- for e in torch.unique(expert).tolist():
- sel = torch.nonzero(expert == e).squeeze(1)
- w = _dequant(weight[e].view(torch.uint8), unswizzle(w_sf[e], n2, hidden // vec), vec)
- acc = (a[sel] @ w.T) * float(alpha[e])
- acc = acc.view(-1, n2 // (2 * group), 2, group)
- up = acc[:, :, 0].reshape(-1, inter)
- gate = acc[:, :, 1].reshape(-1, inter)
- y[sel] = up * gate * torch.sigmoid(gate)
- scaled = y * float(global_sf.reshape(-1)[0])
-
- blocks = scaled.view(-1, inter // vec, vec)
- pre = blocks.abs().amax(-1) / E2M1_MAX
- sf = pre.clamp(max=E4M3_MAX).float().to(torch.float8_e4m3fn)
- sf_val = sf.double().unsqueeze(-1)
- q = torch.where(sf_val > 0, blocks / sf_val, torch.zeros_like(blocks))
- codes = e2m1_codes(q.reshape(-1, inter))
- c = torch.zeros(m, inter // 2, dtype=torch.uint8, device=dev)
- c[real] = codes[:, 0::2] | (codes[:, 1::2] << 4)
- sf_rows = torch.zeros(m, inter // vec, dtype=torch.uint8, device=dev)
- sf_rows[real] = sf.view(torch.uint8)
- return Fc1Expected(
- c=c.view(torch.float4_e2m1fn_x2), c_sf=swizzle(sf_rows), rows=rows, scaled=scaled
- )
-
- def is_valid(
- self,
- input: torch.Tensor,
- weight: torch.Tensor,
- input_scale: torch.Tensor,
- weight_scale: torch.Tensor,
- alpha: torch.Tensor,
- tile_idx_to_group_idx: torch.Tensor,
- tile_idx_to_mn_limit: torch.Tensor,
- permuted_idx_to_expanded_idx: torch.Tensor,
- num_non_exiting_tiles: torch.Tensor,
- global_sf: torch.Tensor,
- num_experts: int,
- top_k: int,
- num_local_experts: int,
- local_expert_offset: int,
- tile_size: int,
- scaling_vector_size: int = 16,
- activation_type: int = _SWIGLU,
- swiglu_limit_scalar: float = SWIGLU_LIMIT_SCALAR_DISABLED,
- expert_counts: Optional[torch.Tensor] = None,
- expert_capacity: int = 0,
- situ_beta: Optional[float] = None,
- situ_linear_beta: Optional[float] = None,
- ) -> None:
- # The runner hands data_ptr() to the kernel with dense row-major layouts
- # assumed; it checks dims and dtypes, never strides.
- for name, t in (
- ("input", input),
- ("input_scale", input_scale),
- ("weight", weight),
- ("weight_scale", weight_scale),
- ("alpha", alpha),
- ):
- assert t.is_contiguous(), (
- f"{name} must be contiguous; the kernel reads it from data_ptr() as a "
- "dense row-major buffer and silently reads the wrong elements"
- )
- # One alpha per expert of `weight`; the runner checks only alpha.dim() == 1.
- assert alpha.numel() == weight.size(0), (
- f"alpha has {alpha.numel()} elements for {weight.size(0)} experts; the "
- "kernel reads alpha[expert] past the end without raising"
- )
- # The runner checks weight_scale.dim() == 3 only, never its size.
- expected = weight.numel() * 2 // scaling_vector_size
- assert weight_scale.numel() == expected, (
- f"weight_scale holds {weight_scale.numel()} bytes, the weight needs "
- f"{expected}; the kernel reads past its end -- silently within the "
- "allocation, with an illegal memory access beyond it"
- )
-
-
-# ── native-torch NVFP4 helpers for reference() and the test's compare ────────
-
-
-def unswizzle(buf: torch.Tensor, rows: int, cols: int) -> torch.Tensor:
- """128x4-swizzled scale bytes -> `[rows, cols]` (rows % 128 == 0, cols % 4 == 0).
-
- Byte `(r, c)` sits at `(r // 128) * 128 * cols + (c // 4) * 512 + (r % 32) * 16
- + ((r % 128) // 32) * 4 + c % 4`.
- """
- return buf.reshape(rows // 128, cols // 4, 32, 4, 4).permute(0, 3, 2, 1, 4).reshape(rows, cols)
-
-
-def swizzle(sf: torch.Tensor) -> torch.Tensor:
- """`[rows, cols]` scale bytes -> the flat 128x4-swizzled buffer; inverse of `unswizzle`."""
- rows, cols = sf.shape
- return sf.reshape(rows // 128, 4, 32, cols // 4, 4).permute(0, 3, 2, 1, 4).reshape(-1)
-
-
-def e4m3_value(byte: torch.Tensor) -> torch.Tensor:
- return byte.to(torch.uint8).view(torch.float8_e4m3fn).double()
-
-
-def _dequant(packed: torch.Tensor, sf: torch.Tensor, vec: int) -> torch.Tensor:
- """`[R, K/2]` packed e2m1 (element 2j in the low nibble) and `[R, K/vec]` e4m3
- bytes -> float64 `[R, K]`. Exact."""
- lut = torch.tensor(E2M1, dtype=torch.float64, device=packed.device)
- lo, hi = (packed & 0xF).long(), (packed >> 4).long()
- codes = torch.stack([lo, hi], dim=-1).reshape(packed.shape[0], -1)
- vals = lut[codes & 7] * torch.where(codes >= 8, -1.0, 1.0).double()
- return vals * e4m3_value(sf).repeat_interleave(vec, dim=1)
-
-
-def e2m1_codes(v: torch.Tensor) -> torch.Tensor:
- """float -> 4-bit e2m1 codes (bit 3 = sign), round-to-nearest-even, saturating at 6."""
- mid = torch.tensor(_E2M1_MIDPOINTS, device=v.device, dtype=v.dtype)
- tie_up = torch.tensor(_TIE_UP, device=v.device)
- a = v.abs().unsqueeze(-1)
- code = ((a > mid).sum(-1) + ((a == mid) & tie_up).any(-1)).to(torch.uint8)
- negative = (v < 0) | ((v == 0) & torch.signbit(v))
- return code | (negative.to(torch.uint8) << 3)
-
-
-cute_dsl_nvfp4_gather_grouped_gemm_act_fusion = CuteDslNvfp4GatherGroupedGemmActFusion()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py
deleted file mode 100644
index 3d7512323e9c..000000000000
--- 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
+++ /dev/null
@@ -1,367 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""CuteDSL NVFP4 FC2 of the routed experts: the plain grouped GEMM over the permuted
-FC1 output, one bf16 row per permuted row, in one trtllm call."""
-
-from __future__ import annotations
-
-from typing import NamedTuple
-
-import torch
-
-import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-from .cute_dsl_nvfp4_gather_grouped_gemm_act_fusion import _dequant, unswizzle
-
-
-def _ckpt(
- tokens: int,
- tile_size: int,
- mma_n: int = 128,
- cluster_n: int = 1,
- dist: str = "router",
- gsf: str = "calibrated",
- zero_tokens: int = 0,
-) -> dict:
- """The checkpoint's routed experts: 512 experts, top-10, hidden 2560, I 640, all
- local (tp1), bf16 out, 16-element e4m3 block scales.
-
- `mma_n` / `cluster_n` select the autotuner tactic: MMA tile `(tile_size,
- mma_n)`, cluster `(tile_size // 128, cluster_n)`; (128, 1) is also the
- fallback when the tuner has no entry. `dist` names how the test draws the
- expert ids: `router` = each row's top-10 of uniform scores. `gsf` is the FC1
- `global_sf` that quantized this op's input, against the SwiGLU output's
- magnitude: `calibrated` = the output maximum the test calibrates to maps to
- the e4m3 maximum, as the checkpoint's per-layer `1 / down_proj.input_scale`
- does; `subnormal` = 2^-12 of it, so input block scales fall into e4m3
- subnormals and zero. `zero_tokens` makes that many leading tokens' rows all
- zero, as FC1 returns them for an all-zero token.
- """
- return dict(
- tokens=tokens,
- tile_size=tile_size,
- mma_n=mma_n,
- cluster_n=cluster_n,
- dist=dist,
- gsf=gsf,
- zero_tokens=zero_tokens,
- num_experts=512,
- top_k=10,
- hidden=2560,
- intermediate=640,
- )
-
-
-class Fc2Expected(NamedTuple):
- """What `reference` returns: the op's output, plus what the test's compare reads.
-
- `out` is laid out exactly as the op returns it; rows the op does not write
- hold zeros. `rows` marks the rows the op writes -- always the prefix of
- `num_non_exiting_tiles * tile_size` rows. `exact` is `alpha[e] * sum_k a_k
- w_k` on those rows, before any rounding; `magnitude` is `|alpha[e]| * sum_k
- |a_k w_k|`, the scale every fp32 rounding of the accumulation is relative to.
- """
-
- out: torch.Tensor
- rows: torch.Tensor
- exact: torch.Tensor
- magnitude: torch.Tensor
-
-
-class CuteDslNvfp4GroupedGemm(OpWrapper):
- """`torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_blackwell`: FC2 of the routed
- experts on the CuteDSL MoE path, one kernel launch.
-
- For every row `p` of the first `num_non_exiting_tiles` tiles (tile `i = p //
- tile_size`):
-
- e = tile_idx_to_group_idx[i] the local expert
- acc[p] = deq(input[p]) @ deq(weight[e]).T [H], fp32
- out[p] = bf16(alpha[e] * acc[p]) fp32 product
-
- where `deq` is e2m1 times its e4m3 block scale. Every row of a live tile is
- computed, padding rows included; no row is gathered, weighted or summed.
-
- Arguments, as the CuteDSL MoE passes them (`run_moe_nvfp4_impl`,
- `_torch/moe/fused_moe/fused_moe_cute_dsl.py`, its non-fused-finalize branch):
-
- - `input` `[M, I/2]` float4_e2m1fn_x2 and `input_scale` `[M * I/16]` uint8:
- FC1's output exactly as the CuteDSL FC1
- (`cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell`) returns it --
- one row per PERMUTED row, the block scales 128x4-swizzled over `[M,
- I/16]`. `M = permuted_idx_to_expanded_idx.numel()` of the sort.
- - `weight` `[E, H, I/2]` float4_e2m1fn_x2: `down_proj` per expert in the
- checkpoint's row order. No 32-row shuffle (that is trtllm-gen's layout).
- - `weight_scale` `[E, H, I/16]` uint8: its e4m3 block scales, 128x4-swizzled
- per expert (`block_scale_interleave`).
- - `alpha` `[E]` fp32: per expert, `max(down_proj.input_scale) * weight_scale_2[e]`
- -- the inverse of FC1's `global_sf` times the expert's weight scale.
- - `tile_idx_to_group_idx` `[M / tile_size]` int32 and `num_non_exiting_tiles`
- `[1]` int32: moe_sort's `tile_idx_to_expert_idx` and
- `num_non_exiting_tiles` from the sort FC1 ran on, with tile_tokens_dim ==
- `tile_size`.
- - `num_experts`, `top_k`, `num_local_experts`, `local_expert_offset`: read only
- by the autotuner's shape buckets and cache key; the kernel takes the
- expert count from `weight.size(0)` and the expert from the tile.
- - `output_dtype`: bf16, the only one the runner builds.
-
- Returns `[M, H]` bfloat16, one row per permuted row. The pair `(t, k)` is row
- `expanded_idx_to_permuted_idx[t, k]`; rows outside the live tiles are not
- written (see `note`).
-
- Fusion boundary: the grouped GEMM, the per-expert alpha and the bf16
- conversion happen inside. The rows arrive already permuted (FC1 gathered
- them); the routing-weighted top-k combine is the caller's
- (`moe_unpermute_inplace`). Routing weights are not read here.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="tile 128: decode, 1 request without MTP; the fallback tactic",
- spec=_ckpt(1, 128),
- ),
- Cell(
- why="tile 128: MMA N 256 and cluster N 2 tactic, 8192-row max_num_tokens prefill",
- spec=_ckpt(8192, 128, mma_n=256, cluster_n=2),
- ),
- Cell(
- why="tile 128: tiny FC1 global_sf and 7 all-zero token rows; subnormal and zero "
- "input scales",
- spec=_ckpt(2048, 128, gsf="subnormal", zero_tokens=7),
- ),
- )
-
- note = """
- Path certified: the CuTe DSL kernel `Sm100BlockScaledContiguousGroupedGemmKernel`
- (`_torch/moe/kernels/blackwell/blockscaled_contiguous_grouped_gemm.py`),
- built by `Sm100BlockScaledContiguousGroupedGemmRunner`
- (`_torch/custom_ops/cute_dsl_custom_ops.py`), which JIT-compiles it with
- `cute.compile` on the first call per (tile, tactic) and keeps it for the
- process. The op is registered only when the CuTe DSL imports
- (`IS_CUTLASS_DSL_AVAILABLE`), and the runner accepts SM 100 and 103 only.
-
- Metadata consumed: none beyond the arguments. Stateless: allocates its output
- with torch.empty on every call, on the current stream, with no host
- synchronization; `num_non_exiting_tiles` is read on the device, so it
- captures into a CUDA graph together with moe_sort and FC1.
-
- Call sites: the engine's CuteDSL MoE (`run_moe_nvfp4_impl`, when the fused
- finalize is off) and the target's routed-expert step, both followed by
- `moe_unpermute_inplace`; eager for prefill and mixed batches (T up to
- max_num_tokens 8192) and captured for decode (T = requests x MTP verify
- width, up to 32 x 4 = 128). Tile 128 is the target's and the cells' tile;
- 256 is certified through the per-pair invariance below.
-
- Numerics: each written element equals the reference's `out` bit for bit on
- every cell (and, at bring-up, on the real FC1 output at 8192 tokens and on the
- checkpoint's own down_proj experts). Operands whose block scales span a wider range can round a
- partial sum; the test's compare admits that at bf16 ties only.
-
- Tactics. The op asks the AutoTuner for an MMA tile of (tile_size, 128) or
- (tile_size, 256) and a cluster of (tile_size // 128, 1) or (tile_size //
- 128, 2), and with no tuning entry falls back to ((tile_size, 128),
- (tile_size // 128, 1)); serving warmup tunes per row bucket and may pick any
- of the four. Every pair's row is bitwise identical across tile 128 / 256 x
- the four tactics x the fallback (the test checks it at 8192 rows), so cells
- driven at one tactic certify the others.
-
- Determinism. A (token, expert) pair's output row depends on that pair's input
- row and its expert alone: it is bitwise unchanged by moving it to another row
- of its expert's run (moe_sort orders them by atomics from T = 5), by repeating
- the call, and by cutting a batch to a prefix. This op does not make the MoE
- run-to-run nondeterministic; the fused-finalize variant
- (`cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell`) reduces across a
- token's experts in launch order and does.
-
- What is read and written: every row of the first num_non_exiting_tiles tiles
- is written; the pair rows hold their results and the padding rows of a live
- tile are computed from whatever the input holds there -- zero for FC1's
- output, which zeroes them. Rows at and past num_non_exiting_tiles * tile_size
- are never written: eager they hold whatever torch.empty returned, under a
- graph whatever the pool last held. Input rows past that point reach no written
- row. The combine must read only expanded_idx_to_permuted_idx's rows.
-
- Preconditions the caller owns:
- - `weight` / `weight_scale` in the layout the class docstring states, and
- `input_scale` swizzled. A wrong preparation of each runs without error,
- with finite output, and puts most of the pair rows' elements outside the
- test's gate: trtllm-gen's 32-row shuffle of the weight, input scales
- linear, one alpha for the whole layer.
- - `alpha` per expert, `max(down_proj.input_scale) * weight_scale_2[e]`, with
- the maximum the value FC1's `global_sf` inverts.
- - `tile_idx_to_group_idx` / `num_non_exiting_tiles` from the same sort as
- FC1's input, tile_tokens_dim == `tile_size`, and every expert index below
- `weight.size(0)`: nothing checks the values (a tile naming expert 512 of
- 512 ran without error and came back all zeros).
- - Every tensor on the device. Nothing checks it: a host-resident `alpha` or
- tile table was read correctly on a coherent-memory Grace host, unmeasured
- where the GPU cannot dereference a host pointer.
-
- Refused by `is_valid`, each silently wrong unguarded: a non-contiguous input,
- input_scale, weight, weight_scale, alpha or tile_idx_to_group_idx (data_ptr()
- read as dense; finite, wrong rows).
-
- Loud (AssertionError from the runner before the launch): a tile_size other
- than the sort's, either way round; output_dtype float16; alpha bf16 or short
- of weight.size(0); weight_scale with fewer experts or typed float8_e4m3fn;
- input typed uint8; input_scale 2-D or short; scaling_vector_size 32;
- num_non_exiting_tiles of 2 elements; tile_idx_to_group_idx int64 or one tile
- short. tile_size 64: ValueError.
-
- Measured at bring-up and not re-driven by the test: the CUDA-graph capture
- with moe_sort and FC1, the determinism and read / write statements, the wrong
- preparations' distance from the gate, and the loud rejections. The test
- re-drives the cells, the tile / tactic invariance and the `input` contiguity
- refusal.
-
- Uncertified: output dtypes other than bf16 (the runner builds none); expert
- parallelism (the kernel reads neither local_expert_offset nor num_experts, so
- EP changes only the sort and the weight slice); other hidden / intermediate
- sizes; scaling_vector_size other than 16; other architectures.
-
- Hard limits (sm_103)
- input, input_scale, output, weight, weight_scale past 2^31 elements /
- bytes: exact (probed up to M = 53687296 rows and 20973
- experts)
- tile index int32 tables, within moe_sort's 2^31-pair limit; no index
- passes through fp32
- """
-
- def __call__(
- self,
- input: torch.Tensor,
- weight: torch.Tensor,
- input_scale: torch.Tensor,
- weight_scale: torch.Tensor,
- alpha: torch.Tensor,
- tile_idx_to_group_idx: torch.Tensor,
- num_non_exiting_tiles: torch.Tensor,
- num_experts: int,
- top_k: int,
- num_local_experts: int,
- local_expert_offset: int,
- tile_size: int,
- output_dtype: torch.dtype,
- scaling_vector_size: int = 16,
- ) -> torch.Tensor:
- return torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_blackwell(
- input=input,
- weight=weight,
- input_scale=input_scale,
- weight_scale=weight_scale,
- alpha=alpha,
- tile_idx_to_group_idx=tile_idx_to_group_idx,
- num_non_exiting_tiles=num_non_exiting_tiles,
- num_experts=num_experts,
- top_k=top_k,
- num_local_experts=num_local_experts,
- local_expert_offset=local_expert_offset,
- tile_size=tile_size,
- output_dtype=output_dtype,
- scaling_vector_size=scaling_vector_size,
- )
-
- def reference(
- self,
- input: torch.Tensor,
- weight: torch.Tensor,
- input_scale: torch.Tensor,
- weight_scale: torch.Tensor,
- alpha: torch.Tensor,
- tile_idx_to_group_idx: torch.Tensor,
- num_non_exiting_tiles: torch.Tensor,
- num_experts: int,
- top_k: int,
- num_local_experts: int,
- local_expert_offset: int,
- tile_size: int,
- output_dtype: torch.dtype,
- scaling_vector_size: int = 16,
- ) -> Fc2Expected:
- """Dequantize, GEMM in float64, scale by alpha in fp32, round once to bf16.
-
- Products of e2m1 and e4m3 values are exact, and float64 sums of 640 of
- them are exact to ~1e-16 relative, so `exact` carries no accumulation
- error of its own; `out` is that sum rounded to fp32, multiplied by the
- fp32 alpha, rounded to bf16 -- what an fp32-accumulating GEMM with this
- epilogue returns whenever its fp32 partial sums are exact. Rows whose
- input dequantizes to all zeros are zero without a GEMM.
-
- bf16 output only (anything else raises NotImplementedError).
- `num_experts`, `top_k`, `num_local_experts` and `local_expert_offset` do
- not enter the computation.
- """
- if output_dtype != torch.bfloat16:
- raise NotImplementedError("reference covers bf16 output only")
- vec = scaling_vector_size
- dev = input.device
- m = input.shape[0]
- hidden = weight.shape[1]
- k_blocks = input.shape[1] * 2 // vec
- n_tiles = int(num_non_exiting_tiles.reshape(-1)[0])
- span = n_tiles * tile_size
-
- rows = torch.zeros(m, dtype=torch.bool, device=dev)
- rows[:span] = True
- a = _dequant(
- input.view(torch.uint8)[:span],
- unswizzle(input_scale.view(torch.uint8), m, k_blocks)[:span],
- vec,
- )
- live = a.abs().amax(dim=1) > 0
- expert = tile_idx_to_group_idx[:n_tiles].long().repeat_interleave(tile_size)
- acc = torch.zeros(span, hidden, dtype=torch.float64, device=dev)
- mag = torch.zeros(span, hidden, dtype=torch.float64, device=dev)
- w_sf = weight_scale.view(torch.uint8).reshape(weight.shape[0], -1)
- for e in torch.unique(expert[live]).tolist():
- sel = torch.nonzero(live & (expert == e)).squeeze(1)
- w = _dequant(weight[e].view(torch.uint8), unswizzle(w_sf[e], hidden, k_blocks), vec)
- acc[sel] = a[sel] @ w.T
- mag[sel] = a[sel].abs() @ w.abs().T
- scale = alpha.float()[expert].unsqueeze(1)
- out = torch.zeros(m, hidden, dtype=torch.bfloat16, device=dev)
- out[:span] = (acc.float() * scale).to(torch.bfloat16)
- return Fc2Expected(
- out=out,
- rows=rows,
- exact=acc * scale.double(),
- magnitude=mag * scale.double().abs(),
- )
-
- def is_valid(
- self,
- input: torch.Tensor,
- weight: torch.Tensor,
- input_scale: torch.Tensor,
- weight_scale: torch.Tensor,
- alpha: torch.Tensor,
- tile_idx_to_group_idx: torch.Tensor,
- num_non_exiting_tiles: torch.Tensor,
- num_experts: int,
- top_k: int,
- num_local_experts: int,
- local_expert_offset: int,
- tile_size: int,
- output_dtype: torch.dtype,
- scaling_vector_size: int = 16,
- ) -> None:
- # The runner hands data_ptr() to the kernel with dense row-major layouts
- # assumed; it checks dims, sizes and dtypes, never strides.
- for name, t in (
- ("input", input),
- ("input_scale", input_scale),
- ("weight", weight),
- ("weight_scale", weight_scale),
- ("alpha", alpha),
- ("tile_idx_to_group_idx", tile_idx_to_group_idx),
- ):
- assert t.is_contiguous(), (
- f"{name} must be contiguous; the kernel reads it from data_ptr() as a "
- "dense row-major buffer and silently reads the wrong elements"
- )
-
-
-cute_dsl_nvfp4_grouped_gemm = CuteDslNvfp4GroupedGemm()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_permuted.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_permuted.py
deleted file mode 100644
index 90e2623cf65d..000000000000
--- 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
+++ /dev/null
@@ -1,139 +0,0 @@
-# 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
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-from .fp4_block_scale_moe_renorm_routed import _ckpt, run_renorm_moe
-from .fp4_block_scale_moe_renorm_routed import fp4_block_scale_moe_renorm_routed as _finalized
-
-
-def finalize(
- rows: torch.Tensor, topk_weights: torch.Tensor, topk_rows: torch.Tensor
-) -> torch.Tensor:
- """The runner's finalize, emulated: per token `sum_k w[t, k] * rows[topk_rows[t, k]]`
- as fp32 FMAs in slot order from zero (a -1 slot skipped), rounded once to the
- rows' dtype. Each fp32 FMA is formed exactly in float64 (a product of two
- fp32 values is exact there) and rounded to fp32. Returns `[T, H]`."""
- acc = torch.zeros(topk_rows.shape[0], rows.shape[1], dtype=torch.float32, device=rows.device)
- for k in range(topk_rows.shape[1]):
- idx = topk_rows[:, k].long()
- r = rows[idx.clamp(min=0)].float().double()
- w = topk_weights[:, k].float().double().unsqueeze(-1)
- step = (r * w + acc.double()).float()
- acc = torch.where((idx >= 0).unsqueeze(-1), step, acc)
- return acc.to(rows.dtype)
-
-
-class Fp4BlockScaleMoeRenormPermuted(OpWrapper):
- """`fp4_block_scale_moe_renorm_routed` with `do_finalize=False`.
-
- Same operands, same routing, same GEMMs; the runner skips its finalize
- launch and returns three tensors:
-
- - `rows` `[P, H]` bf16: the FC2 output in the runner's permuted, padded row
- order (`P` is the runner's padded maximum; rows no slot maps to hold
- garbage);
- - `topk_weights` `[T, k]` bf16: the routing weights, slot order;
- - `topk_rows` `[T, k]` int32: slot `(t, k)`'s row in `rows`, or -1.
-
- `finalize(rows, topk_weights, topk_rows)` (this module) is bitwise the
- finalized entry's output on the same operands.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="_moe_block decode, finalize folded into the next site: one token; the block "
- "routing kernel",
- spec=_ckpt(1),
- ),
- Cell(
- why="4 rows (MTP3 verify; 2..16-row batches): finalize folded into "
- "qwen4_exp_hc_decode_rows_first",
- spec=_ckpt(4),
- ),
- Cell(
- why="128 rows: the con=128 decode batch; the cluster routing kernel",
- spec=_ckpt(128),
- ),
- Cell(why="coarse 3-level logits: ties decide most slots", spec=_ckpt(64, dist="coarse")),
- )
-
- note = """
- Measured against the finalized entry on identical operands: the emulated
- finalize of the three outputs is bitwise its output at every cell, so the
- routing semantics, the GEMMs and the weight layout are the finalized
- entry's (see its note). Which finalize kernel the runner would have
- launched (scalar below ~1184 blocks, vectorized above) does not change the
- arithmetic: both accumulate fp32 FMAs in slot order and round once.
-
- `rows` is the runner's own `[P, H]` buffer, allocated per call at the
- padded maximum the runner computes from `T`, `k` and the expert count;
- only rows some slot maps to hold data. `topk_weights` carries the output
- dtype (bf16), as the finalize reads it.
-
- Silent, refused by `is_valid`: the finalized entry's refusals (a
- non-contiguous operand, an uncertified (num_experts, top_k) pair, hidden
- not a multiple of 256).
-
- Hard limits
- T certified at 1, 4, 64 and 128 (decode); the decode forward folds at
- every site that takes the combine
- num_experts 512, top_k 10
- """
-
- def __call__(
- self,
- routing_logits: torch.Tensor,
- hidden_states: torch.Tensor,
- hidden_states_scale: torch.Tensor,
- gemm1_weights: torch.Tensor,
- gemm1_weights_scale: torch.Tensor,
- gemm2_weights: torch.Tensor,
- gemm2_weights_scale: torch.Tensor,
- output1_scale_scalar: torch.Tensor,
- output1_scale_gate_scalar: torch.Tensor,
- output2_scale_scalar: torch.Tensor,
- num_experts: int,
- top_k: int,
- intermediate_size: int,
- ) -> tuple:
- rows, topk_weights, topk_rows = run_renorm_moe(
- routing_logits,
- hidden_states,
- hidden_states_scale,
- gemm1_weights,
- gemm1_weights_scale,
- gemm2_weights,
- gemm2_weights_scale,
- output1_scale_scalar,
- output1_scale_gate_scalar,
- output2_scale_scalar,
- num_experts,
- top_k,
- intermediate_size,
- False,
- )
- 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
deleted file mode 100644
index 0377db559291..000000000000
--- 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
+++ /dev/null
@@ -1,388 +0,0 @@
-# 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
-
-from functools import cache
-
-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 tensorrt_llm._torch.autotuner import AutoTuner
-from tensorrt_llm._torch.custom_ops.trtllm_gen_custom_ops import FP4BlockScaleMoERunner
-from tensorrt_llm._torch.utils import last_positive_power_of_2
-from tensorrt_llm.logger import logger
-
-from .renorm_moe_routing_op import renorm_moe_routing_op as _renorm
-
-#: `RoutingMethodType::Renormalize`: NoOp preprocess, top-k, softmax over the k.
-ROUTING_RENORMALIZE = 1
-#: `act_type` SwiGlu.
-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)})
-
-#: The geometry `DECODE_TACTICS` was measured on: (num_experts, top_k, hidden, intermediate).
-DECODE_TACTICS_GEOMETRY = (512, 10, 2560, 640)
-
-#: trtllm-gen tactic `(tileN, FC1 x FC2 config index)` per AutoTuner bucket (the row count
-#: rounded down to a power of two) for the decode row counts, at `DECODE_TACTICS_GEOMETRY`.
-#: Each is the fastest valid tactic of its bucket on GB300 over a 48-layer chain of distinct
-#: expert stacks (HBM-cold weights, as a decode step streams them), with router logits whose
-#: distinct experts per layer match serving's routing skew, in one CUDA graph with PDL. The
-#: tactics compared per bucket (the fastest few, the runner's default, the AutoTuner's picks)
-#: all gave bitwise the same output, so the choice moves time only; the entry's test holds
-#: each pinned tactic bitwise to the AutoTuner path.
-#:
-#: A call outside AutoTuner tuning mode whose bucket is listed runs its tactic directly; any
-#: other call goes through `torch.ops.trtllm.fp4_block_scale_moe_runner` (the AutoTuner), so
-#: the larger buckets are still profiled at warmup. The config index addresses the runner's
-#: compiled cubin list: a tactic the runner does not list as valid is not run; its bucket goes
-#: through the AutoTuner instead, with one warning.
-DECODE_TACTICS: dict[int, tuple[int, int]] = {
- 1: (8, 55),
- 2: (8, 53),
- 4: (8, 44),
- 8: (8, 0),
- 16: (8, 44),
- 32: (8, 44),
- 64: (16, 38),
- 128: (16, 36),
- 256: (32, 32),
- 512: (32, 32),
- 1024: (64, 5),
-}
-
-
-@cache
-def _runner():
- """The SwiGlu trtllm-gen runner object the AutoTuner path also uses (one per act type;
- the other constructor arguments only shape the AutoTuner's tuning config)."""
- num_experts, top_k, _, intermediate = DECODE_TACTICS_GEOMETRY
- return FP4BlockScaleMoERunner(
- num_experts=num_experts,
- top_k=top_k,
- n_group=None,
- topk_group=None,
- intermediate_size=intermediate,
- local_expert_offset=0,
- local_num_experts=num_experts,
- routed_scaling_factor=None,
- routing_method_type=ROUTING_RENORMALIZE,
- do_finalize=True,
- act_type=ACT_SWIGLU,
- ).get_runner()
-
-
-@cache
-def _checked_tactic(bucket: int) -> tuple[int, int] | None:
- """`DECODE_TACTICS[bucket]` if it is a valid config for the bucket's row count, else None
- (the AutoTuner path), warned once per bucket."""
- tactic = DECODE_TACTICS[bucket]
- num_experts, top_k, hidden, intermediate = DECODE_TACTICS_GEOMETRY
- valid = _runner().get_valid_configs(top_k, hidden, intermediate, num_experts, bucket)
- if list(tactic) not in [list(t) for t in valid]:
- logger.warning(
- f"pinned tactic {tactic} for {bucket} rows is not a valid config of this build's "
- "trtllm-gen FP4 MoE runner; using the AutoTuner for this bucket. Re-measure "
- "DECODE_TACTICS."
- )
- return None
- return tactic
-
-
-def _pinned_tactic(
- num_tokens: int, num_experts: int, top_k: int, hidden: int, intermediate: int
-) -> tuple[int, int] | None:
- """The pinned tactic for this call, or None for the AutoTuner path."""
- if (num_experts, top_k, hidden, intermediate) != DECODE_TACTICS_GEOMETRY:
- return None
- bucket = last_positive_power_of_2(num_tokens)
- if bucket not in DECODE_TACTICS or AutoTuner.get().is_tuning_mode:
- return None
- return _checked_tactic(bucket)
-
-
-def run_renorm_moe(
- 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,
- do_finalize: bool,
-) -> list[torch.Tensor]:
- """One logits-routed trtllm-gen FP4 MoE call (Renormalize, SwiGlu, no expert
- parallelism): `[out]` with `do_finalize`, else `[rows, topk_weights, topk_rows]`. A
- `DECODE_TACTICS` bucket runs its pinned tactic; any other call is the AutoTuner's."""
- tactic = _pinned_tactic(
- hidden_states.shape[0], num_experts, top_k, hidden_states.shape[1] * 2, intermediate_size
- )
- if tactic is None:
- 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,
- do_finalize,
- ACT_SWIGLU,
- )
- return _runner().run_moe(
- 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,
- do_finalize,
- list(tactic),
- None,
- None,
- None,
- )
-
-
-def _ckpt(tokens: int, dist: str = "randn") -> dict:
- """qwen4_exp's routed geometry: 512 experts, top-10, H 2560 / I 640."""
- return dict(
- tokens=tokens,
- hidden=2560,
- intermediate=640,
- num_experts=512,
- top_k=10,
- logits_dtype=torch.bfloat16,
- dist=dist,
- )
-
-
-class Fp4BlockScaleMoeRenormRouted(OpWrapper):
- """One NVFP4 MoE layer whose routing happens inside the call.
-
- `routing_logits` is the router GEMM's `[T, num_experts]` output. Per token
- the kernel picks the `top_k` largest logits (ties to the LOWER expert index)
- and weights them by an fp32 softmax over just those `top_k`, rounded to
- bf16 -- the same function as `ops/moe/renorm_moe_routing_op` followed by the
- pre-routed `catalog/moe/fp4_block_scale_moe_runner`, and bitwise equal to
- that composition at every cell. Returns the combined `[T, H]` bf16 rows.
-
- Fixed, not arguments: no routing bias, Renormalize routing, SwiGlu, no
- expert parallelism (`local_num_experts == num_experts`), `do_finalize=True`.
- The weight layout (row shuffle + 128x4 scale swizzle) and the three
- per-expert scalars are exactly the pre-routed entry's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # One cell per routing kernel the call picks by row count: block, single-cluster
- # and the multi-kernel histogram path.
- CELLS: tuple[Cell, ...] = (
- Cell(why="_moe_block decode: one token; the block routing kernel", spec=_ckpt(1)),
- Cell(why="128 rows: the con=128 decode batch; the cluster routing kernel", spec=_ckpt(128)),
- Cell(
- why="16384 rows: max_num_tokens of the tuned config; past the 8192 tuning bucket",
- spec=_ckpt(16384),
- ),
- )
-
- note = """
- Routing semantics, measured against the composition it replaces
- (`renorm_moe_routing_op` -> pre-routed runner): bitwise-equal finalized
- output at every cell (the test re-drives this), and at bring-up on randn /
- coarse / large-magnitude / all-equal logits, bf16 and fp32 logits,
- 1..16384 rows. Both kernels rank by a packed
- (twiddled value, 65535 - index) key (ties to the lower expert index), take
- an fp32 `expf` softmax over the k with the same 32-lane warp reduction, and
- round the weight to bf16 once.
-
- Kernel choice inside the call is by row count (block kernel up to a small
- count, single-cluster up to the cluster cap, then the multi-kernel
- histogram path); every one of them shares the top-k and softmax code.
- `routing_bias` is None (the thop defaults the bias dtype to bf16 and never
- reads it on this policy). PDL (`TRTLLM_ENABLE_PDL`) is read by the routing
- launch.
-
- Silent, refused by `is_valid`: a non-contiguous operand (raw pointers,
- dense layout assumed) and a (num_experts, top_k) pair outside
- `CERTIFIED_PAIRS` (a tier-list miss logs and launches nothing; the test
- holds this refusal).
-
- Not certified: fp32 logits (the runner's other accepted logits dtype).
-
- Tactics: a call of 1..2047 rows outside AutoTuner tuning mode runs the
- `DECODE_TACTICS` tactic of its power-of-two bucket (measured HBM-cold with
- serving's routing skew); the AutoTuner, which profiles one layer's experts
- warm in L2 under uniform routing, takes the rest and the warmup's tuning
- pass. The test holds each pinned tactic bitwise to the AutoTuner path.
-
- Hard limits
- T certified at 1, 128 and 16384 (max_num_tokens of the tuned config).
- Autotuner buckets stop at 8192, the op's
- `tune_max_num_tokens` default, which this entry does not
- pass: a call past 8192 rows runs the 8192 bucket's tactic
- num_experts 512, top_k 10 (CERTIFIED_PAIRS)
- """
-
- def __call__(
- self,
- routing_logits: torch.Tensor,
- hidden_states: torch.Tensor,
- hidden_states_scale: torch.Tensor,
- gemm1_weights: torch.Tensor,
- gemm1_weights_scale: torch.Tensor,
- gemm2_weights: torch.Tensor,
- gemm2_weights_scale: torch.Tensor,
- output1_scale_scalar: torch.Tensor,
- output1_scale_gate_scalar: torch.Tensor,
- output2_scale_scalar: torch.Tensor,
- num_experts: int,
- top_k: int,
- intermediate_size: int,
- ) -> torch.Tensor:
- return run_renorm_moe(
- routing_logits,
- hidden_states,
- hidden_states_scale,
- gemm1_weights,
- gemm1_weights_scale,
- gemm2_weights,
- gemm2_weights_scale,
- output1_scale_scalar,
- output1_scale_gate_scalar,
- output2_scale_scalar,
- num_experts,
- top_k,
- intermediate_size,
- True,
- )[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
deleted file mode 100644
index f7bc335a4e99..000000000000
--- 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
+++ /dev/null
@@ -1,309 +0,0 @@
-# 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.*
-import tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops as _cdo
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.utils import ActivationType
-
-from ...kernels.moe.blockscaled_contiguous_grouped_gemm_finalize_fusion import (
- Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel,
-)
-from .fp4_block_scale_moe_renorm_routed import CERTIFIED_PAIRS, _ckpt
-from .fp4_block_scale_moe_renorm_routed import fp4_block_scale_moe_renorm_routed as _routed
-
-#: Token tile of the grouped GEMMs (the kernels' M tile).
-TILE_TOKENS = 128
-#: The GEMM tactics, pinned: `((mma M, mma N), cluster (M, N), raster along M)`. FC1's is what
-#: the engine's autotune warmup picks at every prefill bucket; FC2's likewise, on the finalize
-#: kernel run with a K-tile of 128 (`_FC2_MMA_INST_TILE_K`).
-FC1_TACTIC = ((128, 256), (1, 1), False)
-FC2_TACTIC = ((128, 256), (1, 2), False)
-#: MMA instructions per FC2 A/B pipeline stage (K 128): five stages beside the C staging
-#: buffer instead of two at the kernel's default K 256, so FC2's short K (640) pipelines.
-_FC2_MMA_INST_TILE_K = 2
-#: Launches of at least this many rows run FC2 with L2 cache hints (operands EVICT_FIRST,
-#: scatter-add EVICT_LAST): a gain at ~16K rows, a small loss at 8K (GB300 micro-bench).
-L2_HINT_MIN_ROWS = 12288 # capacity: fixed: dispatch threshold
-_FP4 = torch.float4_e2m1fn_x2
-_SWIGLU_LIMIT = _cdo._canonicalize_swiglu_limit_scalar(_cdo.SWIGLU_LIMIT_SCALAR_DISABLED)
-
-
-class _Fc2Kernel(Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel):
- def __init__(self, *args, **kwargs):
- super().__init__(
- *args, mma_inst_tile_k=_FC2_MMA_INST_TILE_K, row_owned_epilogue=True, **kwargs
- )
-
-
-class _Fc2HintKernel(Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel):
- def __init__(self, *args, **kwargs):
- super().__init__(
- *args,
- mma_inst_tile_k=_FC2_MMA_INST_TILE_K,
- row_owned_epilogue=True,
- l2_hints=True,
- **kwargs,
- )
-
-
-class _Fc2Runner(_cdo.Sm100BlockScaledContiguousGroupedGemmFinalizeFusionRunner):
- kernel_class = _Fc2Kernel
- kernel_cache = dict()
- tuning_config_cache = dict()
-
-
-class _Fc2HintRunner(_cdo.Sm100BlockScaledContiguousGroupedGemmFinalizeFusionRunner):
- kernel_class = _Fc2HintKernel
- kernel_cache = dict()
- tuning_config_cache = dict()
-
-
-_RUNNERS: dict = {}
-
-
-def _runners(num_experts: int, top_k: int) -> tuple:
- """`(fc1, fc2, fc2 with L2 hints)` for a geometry, built once (they hold compiled kernels)."""
- key = (num_experts, top_k)
- if key not in _RUNNERS:
- _RUNNERS[key] = (
- _cdo.Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner(
- num_experts,
- top_k,
- num_experts,
- 0,
- TILE_TOKENS,
- 16,
- activation_type=ActivationType.Swiglu,
- swiglu_limit_scalar=_SWIGLU_LIMIT,
- ),
- _Fc2Runner(num_experts, top_k, num_experts, 0, TILE_TOKENS, torch.bfloat16, 16),
- _Fc2HintRunner(num_experts, top_k, num_experts, 0, TILE_TOKENS, torch.bfloat16, 16),
- )
- return _RUNNERS[key]
-
-
-class Fp4MoeRenormCuteDsl(OpWrapper):
- """One NVFP4 MoE layer whose routing happens inside the call, on the CuTe
- DSL MoE kernels. Same function as `fp4_block_scale_moe_renorm_routed`:
- per token the `top_k` largest logits (ties to the LOWER expert index),
- weighted by an fp32 softmax over just those `top_k` rounded to bf16; the
- experts' SwiGLU FC1, the NVFP4 requantization of its output, FC2, and the
- weighted sum. Returns the combined `[T, H]` bf16 rows.
-
- Operands, in this kernel family's layout (see `weights.py`): `fc1_w`
- `[E, 2I, H/2]` uint8, `[up ; gate]` rows interleaved in 64-row groups;
- `fc1_s` `[E, 2I, H/16]` e4m3 bytes in that row order, 128x4-swizzled;
- `fc2_w` `[E, H, I/2]` as stored; `fc2_s` `[E, H, I/16]` swizzled.
- `hidden_states` / `hidden_states_scale` are the NVFP4 activation with
- **linear** block scales, quantized with the layer's FC1 activation global
- scale. Scalars: `fc1_alpha` `[E]` dequantizes the FC1 accumulator (the
- routed entry's `output1_scale_gate_scalar`); `fc2_input_scale` `[1]` is the
- FC2 activation's quantization multiplier (the reciprocal of its dequant
- scale), one number for every expert; `fc2_alpha` `[E]` dequantizes the FC2
- accumulator (the routed entry's `output2_scale_scalar`).
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="4096 rows: where the target switches to this entry; the sort's cluster kernel",
- spec=_ckpt(4096),
- ),
- Cell(
- why="16377 rows: a con=128 mixed step (16256 context + 121 generation)",
- spec=_ckpt(16377),
- ),
- Cell(
- why="16384 rows: max_num_tokens of the tuned config; the sort's cooperative kernel",
- spec=_ckpt(16384),
- ),
- )
-
- note = """
- Semantics against `fp4_block_scale_moe_renorm_routed` (trtllm-gen) on the
- same experts: the same routing (both are the `renorm_moe_routing_op`
- function, weights rounded to bf16 once) and the same NVFP4 arithmetic per
- expert (e2m1 x e4m3-block x fp32-alpha products accumulated in fp32, SwiGLU
- on the FC1 accumulator, the FC2 input requantized to NVFP4 with 16-element
- blocks), but not bitwise: the GEMMs' accumulation order differs, and the
- weighted sum over a token's experts is a scatter-add with fp32 atomics, so
- its order -- and the output's last bit -- varies between launches. Both sit
- inside the routed entry's band against its fp32 reference.
-
- One FC2 activation scale for all experts: the routed entry takes it per
- expert (folded into `output1_scale_scalar`); a checkpoint whose experts
- disagree on `down_proj.input_scale` cannot use this entry as is.
-
- Kernels (all TRT-LLM's): `trtllm::moe_topk_sort` (routing + tile metadata),
- `trtllm::moe_output_memset_inplace` (zeroes the rows the scatter-add
- touches; on `memset_stream` beside FC1 when given), the gathered FC1 runner
- and the finalize-fused FC2 runner, both with pinned tactics (no AutoTuner
- lookup); FC2's kernel runs with a K-tile of 128, a row-owned epilogue and,
- from `L2_HINT_MIN_ROWS` rows, L2 cache hints. Not CUDA-graph certified
- (prefill steps run eagerly or under the breakable prefill graph).
-
- Hard limits
- T certified 4096..16384 (the target calls it from
- `_CUTE_DSL_MOE_MIN_ROWS` = 4096 rows)
- num_experts 512, top_k 10 (the routed entry's CERTIFIED_PAIRS)
- """
-
- def __call__(
- self,
- routing_logits: torch.Tensor,
- hidden_states: torch.Tensor,
- hidden_states_scale: torch.Tensor,
- fc1_w: torch.Tensor,
- fc1_s: torch.Tensor,
- fc1_alpha: torch.Tensor,
- fc2_w: torch.Tensor,
- fc2_s: torch.Tensor,
- fc2_alpha: torch.Tensor,
- fc2_input_scale: torch.Tensor,
- num_experts: int,
- top_k: int,
- memset_stream: torch.cuda.Stream | None = None,
- ) -> torch.Tensor:
- tokens = hidden_states.shape[0]
- hidden = hidden_states.shape[1] * 2
- fc1, fc2, fc2_hint = _runners(num_experts, top_k)
- if tokens >= L2_HINT_MIN_ROWS:
- fc2 = fc2_hint
- # Renormalize routing and the tile sort in one launch: the same (expert, weight) pairs
- # (weights rounded to bf16 once) and slot order as `renorm_moe_routing_op` + `moe_sort`.
- t2e, t2lim, e2p, p2e, _, live, weights = torch.ops.trtllm.moe_topk_sort(
- routing_logits, None, num_experts, top_k, None, None, 0, num_experts, None,
- TILE_TOKENS, 1,
- ) # fmt: skip
- scales = weights.float()
- out = torch.empty((tokens, hidden), dtype=torch.bfloat16, device=hidden_states.device)
-
- def memset() -> None:
- 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,
- )
-
- # The memset of the rows FC2 scatter-adds into runs beside FC1 on `memset_stream` (a
- # stream nothing else is queued on: the join before FC2 waits for all of it). `out` is
- # a main-stream allocation the side stream is done with before main's next use of it
- # (the join), so its block needs no `record_stream` -- which would make the caching
- # allocator poll an event on every later allocation (~2000 cudaEventQuery a mixed step).
- main = torch.cuda.current_stream()
- if memset_stream is not None:
- memset_stream.wait_stream(main)
- with torch.cuda.stream(memset_stream):
- memset()
- h, h_sf = fc1.forward(
- [
- hidden_states.view(_FP4),
- fc1_w.view(_FP4),
- hidden_states_scale.view(torch.uint8).view(tokens, hidden // 16),
- fc1_s.view(torch.uint8),
- fc1_alpha,
- t2e,
- t2lim,
- p2e,
- live,
- fc2_input_scale,
- ],
- tactic=FC1_TACTIC,
- )
- if memset_stream is not None:
- main.wait_stream(memset_stream)
- else:
- memset()
- fc2.forward(
- [
- h.view(_FP4),
- fc2_w.view(_FP4),
- h_sf.view(torch.uint8),
- fc2_s.view(torch.uint8),
- fc2_alpha,
- out,
- t2e,
- t2lim,
- p2e,
- live,
- scales,
- ],
- tactic=FC2_TACTIC,
- )
- 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,
- memset_stream: torch.cuda.Stream | None = None,
- ) -> 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
deleted file mode 100644
index 2bb80ff63ecb..000000000000
--- 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
+++ /dev/null
@@ -1,365 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""DeepSeek-style FP8 block-scale mixture-of-experts layer: given top-k +
-grouped FC1 GEMM + SwiGLU + FP8 requantization + grouped FC2 GEMM +
-routing-weighted combine, in one trtllm-gen call."""
-
-from __future__ import annotations
-
-from typing import Optional
-
-import torch
-
-import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import (
- Arch,
- Cell,
- OpWrapper,
- true_fp32_matmul,
-)
-
-#: Block edge on every scale axis: 1x128 activations, 128x128 weights.
-BLOCK = 128
-
-#: Largest finite e4m3 magnitude.
-E4M3_MAX = 448.0
-
-#: The activation kernel's floor on a block amax (`kDsActAmaxEpsilon` in
-#: trtllmGenKernels/blockScaleMoe/DevKernel.cu). Only an all-zero block reaches it.
-ACT_AMAX_FLOOR = 1.0e-10
-
-
-def dequantize_block_weight(codes: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
- """One expert's `[M, K]` fp32 weight from e4m3 codes and `[M/128, K/128]` scales."""
- m, k = codes.shape
- blocks = codes.float().reshape(m // BLOCK, BLOCK, k // BLOCK, BLOCK)
- return (blocks * scales.float().reshape(m // BLOCK, 1, k // BLOCK, 1)).reshape(m, k)
-
-
-def requantize_1x128(x: torch.Tensor) -> torch.Tensor:
- """The e4m3 round trip the kernel applies to both intermediates.
-
- Per row and per 128 consecutive columns: `scale = max(amax, 1e-10) / 448`,
- then round-to-nearest-even into e4m3 saturating at +-448. The plain
- `amax / 448` scale, not the round-up power of two the MX members use.
- """
- shape = x.shape
- blocks = x.reshape(*shape[:-1], shape[-1] // BLOCK, BLOCK)
- scale = blocks.abs().amax(dim=-1, keepdim=True).clamp(min=ACT_AMAX_FLOOR) / E4M3_MAX
- codes = (blocks / scale).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn)
- return (codes.float() * scale).reshape(shape)
-
-
-def _geom(hidden, inter, num_experts, top_k, offset=0, local=None):
- """A cell's MoE geometry; the window defaults to every expert local."""
- return dict(
- hidden=hidden,
- inter=inter,
- num_experts=num_experts,
- top_k=top_k,
- local_expert_offset=offset,
- local_num_experts=num_experts if local is None else local,
- )
-
-
-#: The MTP draft layer's routed experts: H 2560, I 640, 512 experts, top-10,
-#: every expert local (tp1).
-_CKPT = _geom(2560, 640, 512, 10)
-
-
-class Fp8BlockScaleMoeRunner(OpWrapper):
- """One FP8 block-scale MoE layer over e4m3 activations, entered pre-routed.
-
- Per token `t` and slot `j`, with `A[t]` the fp32 activations
- (`hidden_states[t, k] * hidden_states_scale[k // 128, t]`) and `W_*` the
- dequantized 128x128-block weights of local expert `e = topk_ids[t, j] -
- local_expert_offset` (slots outside `[0, local_num_experts)` contribute
- nothing):
-
- up, gate = requant(A[t] @ [W_up; W_gate].T).split(I)
- gate, up = min(gate, lim), clamp(up, -lim, lim) # gemm1_clamp_limit
- out[t] += topk_weights[t, j] * (requant(silu(gate) * up) @ W_down.T)
-
- accumulated in fp32 and stored bf16, `requant` being `requantize_1x128`.
- FC1 rows are `[up ; gate]` -- up first, the engine loader's `[w3; w1]`.
-
- Fusion boundary: permutation, both GEMMs with on-the-fly dequantization,
- both intermediate requantizations, the clamped SwiGLU and the weighted
- combine are inside. The router GEMM and top-k, quantizing the bf16 hidden
- states (`fp8_quantize_1x128(x, use_ue8m0=False)`), the shared expert, any
- all-reduce and the residual are the caller's.
-
- `hidden_states_scale` is `[H/128, T]` -- transposed, exactly what
- `fp8_quantize_1x128` returns on this architecture. Pass it straight through:
- at `T == H/128` (20 tokens at H 2560) its transpose has the same shape and
- passes every check, and is 41-43 ulp RMS wrong.
-
- Returns a fresh `[T, H]` bf16 tensor, or an empty `[0]` bf16 tensor when
- `output` is given (the result is then in `output`).
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- # The token column at the checkpoint geometry, the 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="20 tokens == H/128: the activation scale buffer is square",
- spec=dict(tokens=20, **_CKPT),
- ),
- Cell(
- why="16384 rows: max_num_tokens of the tuned config; past the last tuning bucket",
- spec=dict(tokens=16384, **_CKPT),
- ),
- )
-
- note = """
- Metadata consumed: none. Two process-global caches sit behind the call --
- the autotuner's profiling cache and the C++ runner cache -- and a warm
- tuned tactic 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
- per-expert tensors in plain row-major layout, no shuffle and no swizzle:
- `gemm1_weights` `[E, 2I, H]` e4m3 rows `[up ; gate]`, `gemm1_weights_scale`
- `[E, 2I/128, H/128]` fp32, `gemm2_weights` `[E, H, I]`, `gemm2_weights_scale`
- `[E, H/128, I/128]`; `topk_ids` `[T, top_k]` int32 and `topk_weights`
- `[T, top_k]` bf16, used as given (nothing is renormalized). H and I are
- multiples of 128, `intermediate_size == gemm1_weights.shape[1] // 2`,
- `0 < top_k < num_experts`, `top_k <= 32`. The op raises on each of those.
- Supply distinct ids per row: a repeated id counted once at 16 tokens and
- twice at 256 (the token count picks the routing kernel).
-
- The window: local slot i answers for global id `local_expert_offset + i`.
- Nothing checks the window against `num_experts`; 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).
-
- Inert on the certified path, bitwise: `routing_method_type` 0/1/2/4 (the
- pre-routed path forces the routing postprocess), `n_group`/`topk_group`
- (1, 1), `routed_scaling_factor` None/1.0/2.5, `tune_max_num_tokens`,
- `use_dp`. `gemm1_clamp_limit` is live and one scalar for every expert,
- applied after the FC1 requantization; None is bitwise equal to omitting it.
-
- Silent on the raw op, refused by `is_valid` (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,
- which is the only path the C++ runner's small-tile tactic exclusion guards.
-
- The registered fake returns `[T, 2H]` (it assumes packed fp4 hidden
- states): an upstream defect, eager is unaffected.
-
- Hard limits
- T certified at 1, 20 and 16384 (max_num_tokens of the tuned config);
- chunk above it, not probed. Autotuner buckets stop at
- `tune_max_num_tokens` (default 8192, which the target
- leaves): past it a call runs the 8192 bucket's tactic; a
- larger value makes the op append a bucket at that size
- top_k <= 32 (raises past it); certified 10, the checkpoint's
- E, H, I certified at the checkpoint's 512, 2560, 640
-
- Not certified (the statements above were measured at bring-up and are not
- re-driven by the test): expert windows (`local_expert_offset` /
- `local_num_experts` other than the full stack), repeated or out-of-range
- ids, `gemm1_clamp_limit`, the inert arguments, `output=`, the registered
- fake, the op's loud rejections. Nor: the in-kernel routing path
- (`routing_logits` given), multi-rank execution, geometries other than the
- checkpoint's, sm_100.
- """
-
- def __call__(
- self,
- routing_logits: Optional[torch.Tensor],
- routing_bias: Optional[torch.Tensor],
- hidden_states: torch.Tensor,
- hidden_states_scale: torch.Tensor,
- gemm1_weights: torch.Tensor,
- gemm1_weights_scale: torch.Tensor,
- gemm2_weights: torch.Tensor,
- gemm2_weights_scale: torch.Tensor,
- num_experts: int,
- top_k: int,
- num_fused_shared_experts: Optional[int],
- n_group: Optional[int],
- topk_group: Optional[int],
- intermediate_size: int,
- local_expert_offset: int,
- local_num_experts: int,
- routed_scaling_factor: Optional[float],
- routing_method_type: int,
- topk_weights: Optional[torch.Tensor] = None,
- topk_ids: Optional[torch.Tensor] = None,
- act_type: int = 0,
- gemm1_clamp_limit: Optional[float] = None,
- output: Optional[torch.Tensor] = None,
- tune_max_num_tokens: int = 8192,
- use_dp: bool = False,
- ) -> torch.Tensor:
- return torch.ops.trtllm.fp8_block_scale_moe_runner(
- routing_logits,
- routing_bias,
- hidden_states,
- hidden_states_scale,
- gemm1_weights,
- gemm1_weights_scale,
- gemm2_weights,
- gemm2_weights_scale,
- num_experts,
- top_k,
- num_fused_shared_experts,
- n_group,
- topk_group,
- intermediate_size,
- local_expert_offset,
- local_num_experts,
- routed_scaling_factor,
- routing_method_type,
- topk_weights=topk_weights,
- topk_ids=topk_ids,
- act_type=act_type,
- gemm1_clamp_limit=gemm1_clamp_limit,
- output=output,
- tune_max_num_tokens=tune_max_num_tokens,
- use_dp=use_dp,
- )
-
- def reference(
- self,
- routing_logits: Optional[torch.Tensor],
- routing_bias: Optional[torch.Tensor],
- hidden_states: torch.Tensor,
- hidden_states_scale: torch.Tensor,
- gemm1_weights: torch.Tensor,
- gemm1_weights_scale: torch.Tensor,
- gemm2_weights: torch.Tensor,
- gemm2_weights_scale: torch.Tensor,
- num_experts: int,
- top_k: int,
- num_fused_shared_experts: Optional[int],
- n_group: Optional[int],
- topk_group: Optional[int],
- intermediate_size: int,
- local_expert_offset: int,
- local_num_experts: int,
- routed_scaling_factor: Optional[float],
- routing_method_type: int,
- topk_weights: Optional[torch.Tensor] = None,
- topk_ids: Optional[torch.Tensor] = None,
- act_type: int = 0,
- gemm1_clamp_limit: Optional[float] = None,
- output: Optional[torch.Tensor] = None,
- tune_max_num_tokens: int = 8192,
- use_dp: bool = False,
- ) -> torch.Tensor:
- """The layer over bit-identical operands, fp32 GEMMs, bf16 store.
-
- Pre-routed only: `routing_logits` must be None. Ignored: `routing_bias`,
- `num_experts`, `top_k`, `num_fused_shared_experts`, `n_group`,
- `topk_group`, `routed_scaling_factor`, `routing_method_type`,
- `act_type`, `output`, `tune_max_num_tokens`, `use_dp` -- each inert on
- that path. Returns the `[T, H]` result even when `output` is given.
- """
- assert routing_logits is None and topk_ids is not None and topk_weights is not None, (
- "the reference models the pre-routed path only"
- )
- num_tokens, hidden = hidden_states.shape
- nb = hidden // BLOCK
- act_in = (
- hidden_states.float().reshape(num_tokens, nb, BLOCK)
- * hidden_states_scale.t().reshape(num_tokens, nb, 1).float()
- ).reshape(num_tokens, hidden)
- out = torch.zeros(num_tokens, hidden, dtype=torch.float32, device=hidden_states.device)
- local = topk_ids.long() - local_expert_offset
- live = local[(local >= 0) & (local < local_num_experts)]
- with true_fp32_matmul():
- for e in torch.unique(live).tolist():
- tok, slot = (local == e).nonzero(as_tuple=True)
- w1 = dequantize_block_weight(gemm1_weights[e], gemm1_weights_scale[e])
- y1 = requantize_1x128(act_in[tok] @ w1.t())
- up, gate = y1[:, :intermediate_size], y1[:, intermediate_size:]
- if gemm1_clamp_limit is not None:
- gate = gate.clamp(max=gemm1_clamp_limit)
- up = up.clamp(-gemm1_clamp_limit, gemm1_clamp_limit)
- act = requantize_1x128(torch.nn.functional.silu(gate) * up)
- w2 = dequantize_block_weight(gemm2_weights[e], gemm2_weights_scale[e])
- y = act @ w2.t()
- out.index_add_(0, tok, y * topk_weights[tok, slot].float().unsqueeze(1))
- return out.to(torch.bfloat16)
-
- def is_valid(
- self,
- routing_logits: Optional[torch.Tensor],
- routing_bias: Optional[torch.Tensor],
- hidden_states: torch.Tensor,
- hidden_states_scale: torch.Tensor,
- gemm1_weights: torch.Tensor,
- gemm1_weights_scale: torch.Tensor,
- gemm2_weights: torch.Tensor,
- gemm2_weights_scale: torch.Tensor,
- num_experts: int,
- top_k: int,
- num_fused_shared_experts: Optional[int],
- n_group: Optional[int],
- topk_group: Optional[int],
- intermediate_size: int,
- local_expert_offset: int,
- local_num_experts: int,
- routed_scaling_factor: Optional[float],
- routing_method_type: int,
- topk_weights: Optional[torch.Tensor] = None,
- topk_ids: Optional[torch.Tensor] = None,
- act_type: int = 0,
- gemm1_clamp_limit: Optional[float] = None,
- output: Optional[torch.Tensor] = None,
- tune_max_num_tokens: int = 8192,
- use_dp: bool = False,
- ) -> None:
- """Three inputs the op accepts and answers wrongly, each measured.
-
- The transposed activation-scale buffer at `T == H/128` is the fourth,
- and no metadata check can see it.
- """
- # The kernel takes raw data pointers: a strided view is read as dense.
- for name, tensor in (
- ("routing_logits", routing_logits),
- ("routing_bias", routing_bias),
- ("hidden_states", hidden_states),
- ("hidden_states_scale", hidden_states_scale),
- ("gemm1_weights", gemm1_weights),
- ("gemm1_weights_scale", gemm1_weights_scale),
- ("gemm2_weights", gemm2_weights),
- ("gemm2_weights_scale", gemm2_weights_scale),
- ("topk_weights", topk_weights),
- ("topk_ids", topk_ids),
- ("output", output),
- ):
- if tensor is not None:
- assert tensor.is_contiguous(), (
- f"{name} must be contiguous; a strided view is read as if dense "
- "and silently produces wrong results"
- )
- # The Python runner overwrites the activation with SwiGlu.
- assert act_type == 0, (
- f"act_type={act_type} is silently ignored -- this op always computes "
- "SwiGLU. Use a different runner for a non-gated activation"
- )
- # Once topk_ids carry the routing, nothing reads the bias.
- if routing_bias is not None:
- assert topk_ids is None, (
- "routing_bias is ignored when topk_ids/topk_weights are given "
- "(routing has already happened); fold it into the router logits"
- )
-
-
-fp8_block_scale_moe_runner = Fp8BlockScaleMoeRunner()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fused_sigmoid_gate_mul_add.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fused_sigmoid_gate_mul_add.py
deleted file mode 100644
index 367a6139f8a8..000000000000
--- 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
+++ /dev/null
@@ -1,132 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Shared-expert combine: final_hidden_states + sigmoid(gate) * shared_expert_output."""
-
-from __future__ import annotations
-
-from typing import Optional
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-# No torch.ops boundary exists around this kernel: the module callable is the
-# stable surface, and its body is one Triton launch after argument checks.
-from tensorrt_llm._torch.moe.fused_shared_expert import (
- fused_sigmoid_gate_mul_add as _fused_sigmoid_gate_mul_add,
-)
-
-
-def _ckpt(tokens: int, **extra) -> dict:
- """The checkpoint's MoE combine: hidden 2560 bf16, one gate logit per token, in place."""
- spec = dict(
- tokens=tokens,
- hidden=2560,
- dtype=torch.bfloat16,
- gate_dtype=torch.bfloat16,
- gate_layout="column",
- dest="inplace",
- )
- spec.update(extra)
- return spec
-
-
-class FusedSigmoidGateMulAdd(OpWrapper):
- """`out[t, i] = final[t, i] + sigmoid(gate[t]) * shared[t, i]`, fp32 inside, rounded once.
-
- Wraps `tensorrt_llm._torch.moe.fused_shared_expert.fused_sigmoid_gate_mul_add`
- (one Triton kernel, `_sigmoid_gate_mul_add_kernel`).
-
- `gate_logits` is one pre-sigmoid scalar per token read as a flat vector of
- `num_tokens`: `[T, 1]`, `[T]`, `[1, T]` and a strided column are the same
- call, bit for bit, and its dtype is independent of the activations'.
-
- Destination: with `output=None` the result overwrites `final_hidden_states`
- and that tensor is returned; otherwise `output` is written and returned and
- both inputs are untouched. Never a fresh tensor.
-
- Fusion boundary: sigmoid, broadcast, multiply and add. The routed MoE, the
- shared expert and the `[hidden, 1]` gate projection are the caller's. The
- single rounding is the point: a sigmoid-mul then add rounds the product to
- bf16 first and moves the output by up to 0.0312 at [17, 2560] (51.6x the
- bf16 tolerance).
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="_moe_block / MTPLayer._moe decode: one token, in place; 8 warps",
- spec=_ckpt(1),
- ),
- Cell(why="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 "
- "GEMM leaves it",
- spec=_ckpt(33, gate_layout="strided"),
- ),
- )
-
- note = """
- Metadata consumed: none; stateless, no autotuner, no architecture branch.
- The Triton kernel JIT-compiles on first call.
-
- Preconditions, all enforced loudly by the callable (AssertionError before
- any launch; the test holds the wrong-length gate): `final_hidden_states`
- 2-D; `shared_expert_output` and `output` match it in shape and dtype;
- `gate_logits.numel() == num_tokens`; the destination has `stride(-1) == 1`
- (its row stride is free). A column-strided `final_hidden_states` or
- `shared_expert_output` is densified inside the call (an extra copy kernel),
- as is a non-contiguous gate. `num_tokens == 0` returns the destination
- without launching (certified).
-
- Numerics: the default `assert_close` band holds for bf16, the dtype the target
- calls it with; fp16 and fp32 run the same kernel and are not certified.
- bf16 at hidden 2560 is bitwise equal at 1-2 tokens and >= 99.993% bitwise
- at 17-8192 tokens, one ulp otherwise -- `tl.sigmoid`'s fast exponential.
- Wrong variants land 29,000-152,000x outside the bf16 band: no sigmoid,
- multiply instead of add, gate on the routed half, gate dropped, one gate
- for the whole batch.
-
- Hard limits
- num_tokens row offset widened to int64 in the kernel; largest run 8192
-
- Not certified: hidden sizes other than 2560 (the column block and warp count
- step with it), gate layouts other than `[T, 1]` and the strided column, an
- `output` buffer (including one aliasing `shared_expert_output`), row-padded
- operands, host operands, CUDA-graph capture, torch.compile, operands on two
- devices, float8 activations.
- """
-
- def __call__(
- self,
- final_hidden_states: torch.Tensor,
- gate_logits: torch.Tensor,
- shared_expert_output: torch.Tensor,
- output: Optional[torch.Tensor] = None,
- ) -> torch.Tensor:
- return _fused_sigmoid_gate_mul_add(
- final_hidden_states,
- gate_logits,
- shared_expert_output,
- output,
- )
-
- def reference(
- self,
- final_hidden_states: torch.Tensor,
- gate_logits: torch.Tensor,
- shared_expert_output: torch.Tensor,
- output: Optional[torch.Tensor] = None,
- ) -> torch.Tensor:
- """A fresh tensor in `final_hidden_states.dtype`; `output` is ignored.
-
- Takes the operands as they were before the call: in place, the op
- destroys `final_hidden_states`.
- """
- s = torch.sigmoid(gate_logits.reshape(-1, 1).float())
- out = final_hidden_states.float() + s * shared_expert_output.float()
- return out.to(final_hidden_states.dtype)
-
-
-fused_sigmoid_gate_mul_add = FusedSigmoidGateMulAdd()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/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
deleted file mode 100644
index a22c0614c1d0..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_sort.py
+++ /dev/null
@@ -1,272 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""MoE sort: bucket the routed (token, expert) pairs into per-expert row tiles
-for the CuteDSL grouped GEMMs, in one trtllm call."""
-
-from __future__ import annotations
-
-import torch
-
-import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-
-def _ckpt(tokens: int, tile_tokens_dim: int, dist: str = "router") -> dict:
- """The checkpoint's trunk MoE: 512 experts, top-10, all local (tp1), fp32 scales.
-
- `dist` names how the test draws the expert ids: `router` = each row's top-10 of
- uniform scores over the 512 experts; `tile_edges` = three experts hold exactly
- tile-1, tile and tile+1 rows, the rest `router`. Both put 10 distinct experts on
- every row, as the router does.
- """
- return dict(
- tokens=tokens,
- dist=dist,
- num_experts=512,
- top_k=10,
- local_expert_offset=0,
- local_num_experts=512,
- tile_tokens_dim=tile_tokens_dim,
- scales_dtype=torch.float32,
- )
-
-
-class MoeSort(OpWrapper):
- """`torch.ops.trtllm.moe_sort`: routed (token, slot) pairs -> per-expert row tiles.
-
- Input: `token_selected_experts` `[T, top_k]` int32 global expert ids (the
- router's output), `token_final_scales` `[T, top_k]` (its weights). Call it
- with the op's own keyword names, as the CuteDSL MoE does.
-
- Numbering: an *expanded* index is `t * top_k + k` (row-major over
- `token_selected_experts`); a *permuted* index is a row of the grouped GEMM's
- M dimension. Each local expert `e = id - local_expert_offset` owns a
- contiguous run of `count_e` permuted rows starting on a tile boundary,
- experts in increasing order, the run padded up to whole tiles of
- `tile_tokens_dim` rows. Returns a list of six int32 tensors, in this order:
-
- - `tile_idx_to_expert_idx` `[max_tiles]`: the local expert of tile `i`;
- - `tile_idx_to_mn_limit` `[max_tiles]`: one past the last real row of tile
- `i`, as an absolute permuted index: `min((i + 1) * tile, start_e + count_e)`;
- - `expanded_idx_to_permuted_idx` `[T, top_k]`: the permuted row of each
- pair, `-1` for a pair whose expert is not local;
- - `permuted_idx_to_expanded_idx` `[max_tiles * tile]`: the inverse on real rows;
- - `total_num_padded_tokens` `[1]` = `num_non_exiting_tiles * tile`;
- - `num_non_exiting_tiles` `[1]`: tiles actually used.
-
- `max_tiles = m + (T * top_k - m) // tile` with `m = min(local_num_experts,
- T * top_k)` is computed on the host from the shapes and scalars alone, so
- every output shape is fixed by `T`, which is what lets a CUDA graph capture
- it. Only the first `num_non_exiting_tiles` entries of the tile arrays and
- the rows below each tile's `mn_limit` are written; everything else is left
- as allocated (see `note`).
-
- Fusion boundary: sorting only. Routing (the top-k) is the caller's, and so
- are the permutation of the activations (the FC1 gather does it) and the
- routing-weighted unpermute; this op does not read `token_final_scales` into
- any output.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # One cell per kernel path the target reaches at its tile, 128: block (T <= 4),
- # cluster (T 5..8192) and cooperative (T >= 8193).
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="tile 128: decode, 1 request without MTP; block kernel",
- spec=_ckpt(1, 128),
- ),
- Cell(
- why="tile 128: expert counts tile-1, tile, tile+1; the mn_limit clamp; cluster kernel",
- spec=_ckpt(300, 128, "tile_edges"),
- ),
- Cell(
- why="tile 128: 8193 rows; first cooperative-kernel T (prefill)",
- spec=_ckpt(8193, 128),
- ),
- )
-
- note = """
- Metadata consumed: none. Stateless; allocates its six outputs plus scratch on
- every call and runs on the current stream, with no host synchronization, so
- it captures into a CUDA graph.
-
- Call it as the engine's CuteDSL MoE does (`run_moe_nvfp4_impl`,
- `_torch/moe/fused_moe/fused_moe_cute_dsl.py`): `num_experts` = the slot
- count, `local_expert_offset` / `local_num_experts` = this rank's slots, and
- `tile_tokens_dim` = the grouped GEMMs' tile, 128 or 256 on sm_103. The
- downstream FC1 gather GEMM, output memset and FC2 take these six tensors as
- they are.
-
- Kernel path, by `T`:
- T 1..4 routingCustom::routingIndicesBlockKernel
- T 5..8192 routingCustom::routingIndicesClusterKernel
- T 8193..471859 routingInitExpertCounts + routingIndicesCoopKernel
- T >= 471860 routingInitExpertCounts + routingIndicesHistogramKernel
- + routingIndicesOffsetsKernel
- The cooperative ceiling is (SMs - 8) * 512 * 64 / top_k rows with 512
- experts: 471859 on a 152-SM GB300, a different T on another SM count.
-
- Order within an expert's run is unspecified. The block kernel (T <= 4) lays
- each expert's rows out in token order and is bitwise repeatable; every other
- path assigns them by atomics, so `expanded_idx_to_permuted_idx` differs
- between identical calls. A test therefore compares each expert's rows as a
- set and the two mappings as exact inverses; the tile arrays and counts are
- deterministic.
-
- Not written: the tile arrays past `num_non_exiting_tiles`, and the rows of
- `permuted_idx_to_expanded_idx` from each tile's `mn_limit` to the tile's
- end. Eager, those hold whatever the allocator handed back; under a CUDA
- graph they hold the previous replay's indices, which look valid. Every
- consumer must bound its reads by `num_non_exiting_tiles` and
- `tile_idx_to_mn_limit`.
-
- `token_final_scales`: shape-checked and never written; its values never
- reach an output. Its pointer must be non-null, which is why T = 0 raises.
-
- Data hazards `is_valid` cannot see (values): a row naming the same expert
- twice keeps both pairs on the cluster kernel but loses one on the block
- kernel (that pair gets -1). An id outside [local_expert_offset,
- local_expert_offset + local_num_experts) gets -1, except that the cluster
- kernel reads ids through int16, so an id of 65536 + e sorts as expert e. The
- router produces neither.
-
- Loud: int64 ids, 1-D ids, a `[T, top_k]` size mismatch on either input,
- CPU tensors, top_k above 32, num_experts above 2048, T = 0.
-
- Silent, refused by `is_valid`: a non-contiguous `token_selected_experts`
- (read from `data_ptr()` as dense `[T, top_k]`); 2^31 or more elements (the
- host counts pairs in int32: it raises a negative-dimension error from 2^31
- and silently sorts only the wrapped remainder from 2^32); `tile_tokens_dim`
- <= 0 (the host divides by it).
-
- Certified only at the checkpoint's values: num_experts 512, top_k 10, all
- experts local (offset 0), tile 128 (the target's), fp32 scales, and T 1,
- 300 and 8193 (block, cluster and the first cooperative T); CUDA-graph replay
- on new ids at T = 128; the contiguity guard. The order, unwritten-region,
- data-hazard and loud-rejection statements above were measured at bring-up
- and are not re-driven by the test. Expert parallelism, tile 256 and other
- tiles, other expert counts and the histogram path are uncertified.
-
- Hard limits (sm_103)
- T * top_k boundary 2^31 pairs (int32 product)
- max_tiles * tile boundary 2^31 padded rows (raises past it)
- expert id int16 in the cluster kernel; ids < 2048 never reach it
- """
-
- def __call__(
- self,
- token_selected_experts: torch.Tensor,
- token_final_scales: torch.Tensor,
- num_experts: int,
- top_k: int,
- local_expert_offset: int,
- local_num_experts: int,
- tile_tokens_dim: int,
- ) -> list[torch.Tensor]:
- return torch.ops.trtllm.moe_sort(
- token_selected_experts=token_selected_experts,
- token_final_scales=token_final_scales,
- num_experts=num_experts,
- top_k=top_k,
- local_expert_offset=local_expert_offset,
- local_num_experts=local_num_experts,
- tile_tokens_dim=tile_tokens_dim,
- )
-
- def reference(
- self,
- token_selected_experts: torch.Tensor,
- token_final_scales: torch.Tensor,
- num_experts: int,
- top_k: int,
- local_expert_offset: int,
- local_num_experts: int,
- tile_tokens_dim: int,
- ) -> list[torch.Tensor]:
- """Counting sort by local expert, stable in the expanded index.
-
- Within an expert, rows are in increasing expanded index; the test's compare
- accepts any order there. What the op leaves unwritten is `-1` here.
- `token_final_scales`, `num_experts` and `top_k` are not read: the scales
- never reach an output, and the ids' shape carries the rest.
- """
- ids = token_selected_experts
- dev = ids.device
- tile = tile_tokens_dim
- num_tokens, width = ids.shape
- num_expanded = num_tokens * width
- local = ids.reshape(-1).long() - local_expert_offset
- is_local = (local >= 0) & (local < local_num_experts)
- counts = torch.bincount(local[is_local], minlength=local_num_experts)
- tiles_per_expert = (counts + tile - 1) // tile
- num_tiles = int(tiles_per_expert.sum())
- first = min(local_num_experts, num_expanded)
- max_tiles = first + (num_expanded - first) // tile
-
- run_start = (torch.cumsum(tiles_per_expert, 0) - tiles_per_expert) * tile
- tile_expert = torch.repeat_interleave(
- torch.arange(local_num_experts, device=dev), tiles_per_expert
- )
- tile_idx = torch.arange(num_tiles, device=dev)
- mn_limit = torch.minimum(
- (tile_idx + 1) * tile, run_start[tile_expert] + counts[tile_expert]
- )
-
- expanded = torch.arange(num_expanded, device=dev)[is_local]
- expert = local[is_local]
- order = torch.sort(expert, stable=True).indices
- expanded, expert = expanded[order], expert[order]
- rank = torch.arange(expert.numel(), device=dev) - (torch.cumsum(counts, 0) - counts)[expert]
- permuted = run_start[expert] + rank
-
- def unwritten(n: int) -> torch.Tensor:
- return torch.full((n,), -1, dtype=torch.long, device=dev)
-
- tile_idx_to_expert_idx = unwritten(max_tiles)
- tile_idx_to_expert_idx[:num_tiles] = tile_expert
- tile_idx_to_mn_limit = unwritten(max_tiles)
- tile_idx_to_mn_limit[:num_tiles] = mn_limit
- permuted_idx_to_expanded_idx = unwritten(max_tiles * tile)
- permuted_idx_to_expanded_idx[permuted] = expanded
- expanded_idx_to_permuted_idx = unwritten(num_expanded)
- expanded_idx_to_permuted_idx[expanded] = permuted
- outs = (
- tile_idx_to_expert_idx,
- tile_idx_to_mn_limit,
- expanded_idx_to_permuted_idx.view(num_tokens, width),
- permuted_idx_to_expanded_idx,
- torch.tensor([num_tiles * tile], device=dev),
- torch.tensor([num_tiles], device=dev),
- )
- return [t.to(torch.int32) for t in outs]
-
- def is_valid(
- self,
- token_selected_experts: torch.Tensor,
- token_final_scales: torch.Tensor,
- num_experts: int,
- top_k: int,
- local_expert_offset: int,
- local_num_experts: int,
- tile_tokens_dim: int,
- ) -> None:
- # The host computes T * top_k in int32: from 2^31 the op raises a
- # negative-dimension error, from 2^32 it sorts only the wrapped remainder.
- numel = token_selected_experts.numel()
- assert numel < 2**31, (
- f"token_selected_experts has {numel} >= 2^31 elements; moe_sort counts "
- "pairs in int32 and silently sorts only the wrapped remainder past 2^32"
- )
- # The kernels read the ids from data_ptr() as a dense [T, top_k] buffer.
- assert token_selected_experts.is_contiguous(), (
- "token_selected_experts must be contiguous; a strided view is read as a "
- "dense [T, top_k] buffer and silently sorts the wrong ids"
- )
- assert tile_tokens_dim > 0, (
- f"tile_tokens_dim must be positive, got {tile_tokens_dim}; 0 returns a "
- "negative num_non_exiting_tiles without raising"
- )
-
-
-moe_sort = MoeSort()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_unpermute_inplace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_unpermute_inplace.py
deleted file mode 100644
index e82eaa5dfdc4..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_unpermute_inplace.py
+++ /dev/null
@@ -1,246 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""MoE unpermute: each token's routing-weighted sum of its top-k expert rows, taken
-from the grouped GEMM's permuted output in a fixed slot order, in one trtllm call.
-"""
-
-from __future__ import annotations
-
-import torch
-
-import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-
-def _ckpt(tokens: int, values: str = "normal") -> dict:
- """The checkpoint's routed-expert combine: top-10 of 512 experts, hidden 2560,
- bf16 FC2 rows, fp32 routing weights.
-
- The test lays `permuted_input` out exactly as the CuteDSL FC2 receives it: it
- draws each row's top-10 of uniform scores over the 512 experts, sorts the ids
- with `moe_sort` at `tile_tokens_dim` and writes each (token, slot) row at the
- permuted index the sort gave it. `values` is how the FC2 rows are drawn:
- `normal` = standard normal; `wide` = standard normal times 2^u, u uniform in
- [-30, 30] per element, so one sum adds terms up to 2^60 apart in magnitude.
- The scales are a softmax over each row's ten slots, as the renormalizing
- router emits them.
- """
- return dict(
- tokens=tokens,
- tile_tokens_dim=128,
- values=values,
- num_experts=512,
- top_k=10,
- hidden_size=2560,
- dtype=torch.bfloat16,
- scales_dtype=torch.float32,
- )
-
-
-def _fma_f32(acc: torch.Tensor, prod: torch.Tensor) -> torch.Tensor:
- """`acc + prod` rounded once to fp32, as a fused multiply-add rounds it.
-
- `acc` is fp32; `prod` is the fp64 product of a bf16 and an fp32 value, which is
- exact (at most 8 + 24 significant bits). The fp64 sum is rounded to odd -- its
- exact rounding error from TwoSum, then the last bit forced to 1 if the sum was
- inexact -- so that the conversion to fp32 is one correct rounding of the exact
- sum rather than two roundings in a row (53 >= 24 + 2 bits).
- """
- a = acc.double()
- s = a + prod
- bb = s - a
- err = (a - (s - bb)) + (prod - bb)
- inexact = (err != 0) & torch.isfinite(s)
- # A sum rounded away from zero truncates to its neighbour one step toward zero.
- toward_zero = inexact & ((err < 0) != (s < 0))
- bits = s.view(torch.int64)
- bits = torch.where(inexact, (bits - toward_zero.long()) | 1, bits)
- return bits.view(torch.float64).float()
-
-
-class MoeUnpermuteInplace(OpWrapper):
- """`torch.ops.trtllm.moe_unpermute_inplace`: the routing-weighted top-k combine.
-
- For each token `t`, writes row `t` of `output`:
-
- output[t] = bf16( sum over k = 0 .. top_k - 1, in that order, of
- permuted_input[expanded_idx_to_permuted_idx[t, k]]
- * topk_scales[t, k] )
-
- accumulated in fp32 with one rounding per slot (each product fused into the
- running sum) and rounded once to the output dtype. A slot whose permuted index
- is negative is skipped; a token with every slot skipped gets a zero row.
-
- Inputs, as the CuteDSL MoE hands them over: `permuted_input` `[M, H]` bf16 is
- the FC2 grouped GEMM's output, one row per (token, expert) pair in permuted
- order, `M >= T * top_k` (the sort's padded row count); `output` `[T, H]` bf16
- is written in full; `expanded_idx_to_permuted_idx` `[T, top_k]` int32 is the
- third output of `moe_sort`; `topk_scales` `[T, top_k]` fp32 are the routing
- weights. Call it with these keyword names, the op's own. Returns `output`.
-
- Fusion boundary: the gather, the weighting and the sum. The FC2 GEMM and the
- sort are the caller's, and so is anything added to the routed output (the
- shared expert). The op is deterministic: one thread block per token sums its
- slots in slot order, so the result does not depend on where the sort placed
- each row, on which other tokens share the batch, or on the call. That is what
- separates it from the FC2 GEMM's fused finalize, whose reduction order varies
- between launches.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="decode, 1 request without MTP; rows from the sort's block kernel",
- spec=_ckpt(1),
- ),
- Cell(why="8192 rows, max_num_tokens prefill", spec=_ckpt(8192)),
- Cell(
- why="FC2 rows spanning 2^-30..2^30; terms 2^60 apart in one sum",
- spec=_ckpt(8192, values="wide"),
- ),
- )
-
- note = """
- Metadata consumed: none. Stateless; allocates nothing, writes `output` in place
- on the current stream with one kernel, `moeUnpermuteKernel<__nv_bfloat16,
- float, 256>` (`cpp/tensorrt_llm/kernels/moe/cuteDsl/moeUtils.cu`, registered
- in `cpp/tensorrt_llm/thop/moe/cuteDsl/cuteDslMoeUtilsOp.cpp`), one thread block
- per token. No host synchronization: it captures into a CUDA graph, and replay
- on new rows, scales and ids (moe_sort captured with it) is bitwise equal to
- eager and to `reference`.
-
- Call sites: the engine's CuteDSL MoE calls it when the FC2 finalize is not
- fused (`run_moe_nvfp4_impl`, `_torch/moe/fused_moe/fused_moe_cute_dsl.py`), and
- this target's `_moe_block` calls it in the prefill and the decode forward with
- the same four tensors: the FC2 output, the forward's dense `[num_tokens,
- hidden]` routed-output buffer, the sort's third output and the router's fp32
- weights. Prefill runs eagerly at up to max_num_tokens (8192) rows; decode runs
- captured at requests x MTP verify width, up to 128.
-
- What it reads: only the `T * top_k` rows `expanded_idx_to_permuted_idx` names,
- and the scales. The padding rows of each tile and the tiles past the sort's
- `num_non_exiting_tiles` are never read (every cell fills every other row of
- `permuted_input` with NaN, and under a graph those rows hold the previous
- replay's rows), so the FC2 output's unwritten regions need no memset before
- this op. What it writes: every element of `output`, whatever it held, and none
- of its inputs.
-
- Determinism (bitwise, measured at bring-up; the test re-drives only graph
- replay against eager): repeated calls; a second moe_sort of the same ids, which
- places an expert's rows in another order; a random permutation of all permuted
- rows; a prefix of a larger batch, with or without a fresh sort of the prefix
- alone; a shuffle of the tokens. The FC2 GEMM's fused finalize
- (`cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell`) differs from its
- own first call in about half the output elements at decode, while the plain FC2
- GEMM followed by this op does not, and is bitwise equal to the slot-order fp32
- sum of the real FC2 rows.
-
- Accuracy against the exact sum: the fp32 sum's error is bounded by its terms'
- magnitudes, so an element whose terms nearly cancel can miss the float64 sum
- rounded once to bf16 by many ulps of the element itself (827 of 20971520
- elements at T = 8192, standard-normal rows and softmax weights, up to 137.5).
-
- Silent, refused by `is_valid`: a non-contiguous `permuted_input`, `output`,
- `expanded_idx_to_permuted_idx` or `topk_scales` (all four are read or written
- from `data_ptr()` as dense row-major); an `output` whose dtype is not
- `permuted_input`'s (an fp32 output receives the bf16 bit patterns in the first
- half of its bytes); 2^31 or more tokens (nothing is written and nothing
- raises). The test holds the `permuted_input` refusal.
-
- Data hazards `is_valid` cannot see (values): a permuted index at or past
- `permuted_input.size(0)` is read without a bounds check; moe_sort never emits
- one.
-
- Loud (measured, not re-driven): int64 indices; hidden size not a multiple of 8; fp32 rows, or fp16 scales
- with bf16 rows; fewer than `T * top_k` permuted rows; a scales or index shape
- that disagrees with `output`.
-
- Accepted, not certified: T = 0 returns and writes nothing; bf16 scales; fp16
- rows; a CPU `output`.
-
- Certified only at the checkpoint's values: top_k 10, hidden 2560, bf16 rows and
- output, fp32 scales, the sort's layout at tile 128, every slot local; T 1 and
- 8192 eager, 128 under CUDA-graph replay. Slots marked -1 (another rank's
- expert) are not certified.
-
- Hard limits (sm_103)
- tokens one thread block per token, the count narrowed to int32:
- boundary 2^31, `is_valid` refuses it
- output elements T * H, offset in int64: exact past 2^31 and 2^32
- permuted rows int32 index, offset in int64; no index passes through fp32
- """
-
- def __call__(
- self,
- permuted_input: torch.Tensor,
- output: torch.Tensor,
- expanded_idx_to_permuted_idx: torch.Tensor,
- topk_scales: torch.Tensor,
- ) -> torch.Tensor:
- torch.ops.trtllm.moe_unpermute_inplace(
- permuted_input=permuted_input,
- output=output,
- expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx,
- topk_scales=topk_scales,
- )
- return output
-
- def reference(
- self,
- permuted_input: torch.Tensor,
- output: torch.Tensor,
- expanded_idx_to_permuted_idx: torch.Tensor,
- topk_scales: torch.Tensor,
- ) -> torch.Tensor:
- """The slot-order fp32 sum, one rounding per slot, rounded once at the end.
-
- Each product of a row and its weight is exact in float64; `_fma_f32` adds
- it to the fp32 running sum with one rounding, as a fused multiply-add
- does. A slot with a negative index leaves the sum untouched. `output` is
- ignored: the result has `permuted_input`'s dtype, which is what the op
- writes.
- """
- idx = expanded_idx_to_permuted_idx.long()
- acc = torch.zeros(
- idx.shape[0], permuted_input.shape[1], dtype=torch.float32, device=idx.device
- )
- for k in range(idx.shape[1]):
- slot = idx[:, k]
- rows = permuted_input.index_select(0, slot.clamp_min(0)).double()
- prod = rows * topk_scales[:, k : k + 1].double()
- acc = torch.where((slot >= 0).unsqueeze(1), _fma_f32(acc, prod), acc)
- return acc.to(permuted_input.dtype)
-
- def is_valid(
- self,
- permuted_input: torch.Tensor,
- output: torch.Tensor,
- expanded_idx_to_permuted_idx: torch.Tensor,
- topk_scales: torch.Tensor,
- ) -> None:
- # One thread block per token, the count narrowed to int32: from 2^31
- # tokens nothing is written and no error surfaces.
- assert output.size(0) < 2**31, (
- f"output has {output.size(0)} >= 2^31 rows; moe_unpermute_inplace launches one "
- "block per token from an int32 count and silently writes nothing"
- )
- # The kernel reads and writes all four through data_ptr() as dense
- # row-major buffers, and writes `output` with permuted_input's dtype.
- for name, t in (
- ("permuted_input", permuted_input),
- ("output", output),
- ("expanded_idx_to_permuted_idx", expanded_idx_to_permuted_idx),
- ("topk_scales", topk_scales),
- ):
- assert t.is_contiguous(), (
- f"{name} must be contiguous; moe_unpermute_inplace addresses it as a dense "
- "row-major buffer and silently combines or writes the wrong elements"
- )
- assert output.dtype == permuted_input.dtype, (
- f"output dtype {output.dtype} must be permuted_input's ({permuted_input.dtype}); "
- "the op writes permuted_input's element type into it without checking"
- )
-
-
-moe_unpermute_inplace = MoeUnpermuteInplace()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/renorm_moe_routing_op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/renorm_moe_routing_op.py
deleted file mode 100644
index c97fd51ca4e9..000000000000
--- 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
+++ /dev/null
@@ -1,144 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Renormalized MoE routing: top-k over the raw router logits, then a softmax
-over just those k, in one trtllm call."""
-
-from __future__ import annotations
-
-from typing import Optional
-
-import torch
-
-import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-
-def _ckpt(tokens: int, output_dtype: torch.dtype = torch.bfloat16) -> dict:
- """The checkpoint's router: 512 experts, top-10, bf16 logits."""
- return dict(
- tokens=tokens,
- num_experts=512,
- topk=10,
- logits_dtype=torch.bfloat16,
- output_dtype=output_dtype,
- dist="randn",
- )
-
-
-class RenormMoeRoutingOp(OpWrapper):
- """`norm_topk_prob=True` routing. Returns `(topk_ids, topk_weights)` -- IDS FIRST.
-
- Per token: `ids` = the `topk` largest `router_logits`, ranked in the
- logits' own dtype, by decreasing value, ties to the LOWER expert index
- (a stable descending sort -- `torch.topk` breaks ties the other way, and at
- 512 bf16 experts ties are common: it disagrees on ~30% of rows);
- `weights` = an fp32 softmax over just those `topk` logits, cast to
- `output_dtype` afterwards. Equal to HF's softmax-then-top-k-then-divide up
- to rounding (5.96e-08 at 512 / top-10).
-
- `ids` is `[T, topk]` int32; `weights` `[T, topk]` in `output_dtype`, which
- defaults to fp32 -- not the logits' dtype. Sibling routing ops return
- `(weights, ids)`; both are `[T, topk]`, so a swapped unpack is no shape
- error.
-
- Fusion boundary: routing only. The router GEMM is the caller's, and so is
- expert parallelism -- every column of the row is ranked.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(why="MTPLayer._moe decode: one token, bf16 weights", spec=_ckpt(1)),
- Cell(
- why="_moe_block at 16384 rows (max_num_tokens of the tuned config): fp32 weights, "
- "as the CuteDSL experts take them; thousands of rows carry an exact bf16 tie",
- spec=_ckpt(16384, torch.float32),
- ),
- Cell(
- why="coarse logits (3 values): ties decide most of the ranking",
- spec=dict(_ckpt(64), output_dtype=torch.float32, dist="coarse"),
- ),
- )
-
- note = """
- Metadata consumed: none; stateless, deterministic, runs on the ambient
- stream. The launch picks its kernel instance by
- `nextPowerOfTwo(num_experts)` and `nextPowerOfTwo(topk)` -- 512 / top-10 is
- the (512, 16) instance with a 256-thread block. PDL (`TRTLLM_ENABLE_PDL`)
- was measured bitwise inert.
-
- Preconditions, enforced loudly by the op: `router_logits` 2-D CUDA bf16 /
- fp16 / fp32; `output_dtype` fp32, bf16 or None (fp16 raises even for fp16
- logits); `topk <= 16`; `num_experts <= 512` -- this checkpoint sits on
- that ceiling.
-
- Silent, refused by `is_valid`: a non-contiguous `router_logits` (read as a
- 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
- 0.09-0.27 here). `RenormalizeMoeRoutingMethod`'s PyTorch fallback (used past
- 512 / 16) is not bit-compatible: it uses `torch.topk` and casts before the
- softmax.
-
- Bit-exact against `reference` over 2040 configurations 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 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__(
- self,
- router_logits: torch.Tensor,
- topk: int,
- output_dtype: Optional[torch.dtype] = None,
- ) -> tuple[torch.Tensor, torch.Tensor]:
- return torch.ops.trtllm.renorm_moe_routing_op(router_logits, topk, output_dtype)
-
- def reference(
- self,
- router_logits: torch.Tensor,
- topk: int,
- output_dtype: Optional[torch.dtype] = None,
- ) -> tuple[torch.Tensor, torch.Tensor]:
- """Stable descending sort in the logits' dtype, fp32 softmax over the
- selected k, then the cast."""
- order = torch.sort(router_logits, dim=-1, descending=True, stable=True).indices[:, :topk]
- selected = torch.gather(router_logits, 1, order).float()
- e = torch.exp(selected - selected.max(-1, keepdim=True).values)
- weights = e / e.sum(-1, keepdim=True)
- return order.to(torch.int32), weights.to(output_dtype or torch.float32)
-
- def is_valid(
- self,
- router_logits: torch.Tensor,
- topk: int,
- output_dtype: Optional[torch.dtype] = None,
- ) -> None:
- # The kernel addresses router_logits from data_ptr() as dense
- # [T, num_experts]: a column slice routes on the wider buffer's bytes.
- assert router_logits.is_contiguous(), (
- "router_logits must be contiguous; a strided view is read as a dense "
- "[num_tokens, num_experts] buffer and silently gives wrong routing"
- )
- # Surplus slots come back as duplicated expert 0 at weight 0.
- assert topk <= router_logits.shape[-1], (
- f"topk ({topk}) exceeds num_experts ({router_logits.shape[-1]}); the "
- "surplus slots come back as duplicates of expert 0 at weight 0"
- )
-
-
-renorm_moe_routing_op = RenormMoeRoutingOp()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/__init__.py
deleted file mode 100644
index 52a7a9daf028..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/flashinfer_gemma_rmsnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/flashinfer_gemma_rmsnorm.py
deleted file mode 100644
index 4ef212a17f9b..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/flashinfer_gemma_rmsnorm.py
+++ /dev/null
@@ -1,97 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Gemma-style RMS normalization (weight applied as ``1 + w``) via the flashinfer kernel."""
-
-from __future__ import annotations
-
-import torch
-
-import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: The hyper-connection residual: hc_count 4 streams of hidden 2560, row stride 10240.
-_STREAM = dict(row_width=10240, dtype=torch.bfloat16, eps=1e-6)
-
-
-class FlashinferGemmaRmsnorm(OpWrapper):
- """`x / sqrt(mean(x^2, -1) + eps) * (1 + weight)`, one op call, new tensor out.
-
- The reduction, the `1 + w` and the product are all fp32 inside the kernel;
- only the store rounds to `x.dtype` -- so `(1 + w)` is not first rounded to
- bf16. `x` is 2-D `[T, H]` or 3-D `[B, heads, D]` (a separate per-head
- kernel), normalized over the last dim; `x` is not mutated.
-
- Fusion boundary: normalization and the biased weight only. No residual add
- (`flashinfer_gemma_fused_add_rmsnorm` exists for that), no quantization; the
- plain-`w` form is `flashinfer_rmsnorm`.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # The target's grouped norm is one call per 2560-wide stream of the
- # 10240-wide residual, each a column slice at row stride 10240 -- that, not
- # a contiguous row, is what the strided-kernel cells below certify.
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="decode batch 33: stream 1, row stride 10240 > width 2560",
- spec=dict(shape=(33, 2560), col_offset=2560, **_STREAM),
- ),
- Cell(
- why="prefill 2048 tokens: stream 2 at row stride 10240",
- spec=dict(shape=(2048, 2560), col_offset=5120, **_STREAM),
- ),
- Cell(
- why="MTP pre_e_norm: the contiguous [T, 2560] embedding",
- spec=dict(shape=(4, 2560), dtype=torch.bfloat16, eps=1e-6),
- ),
- )
-
- 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`) 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 (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 --
- `is_valid` refuses them.
-
- Row stride is free; the strided kernel declares it divisible by the vector
- width (8 elements for bf16/fp16, 4 for fp32) and assumes 16-byte aligned
- pointers, so odd row strides and misaligned slices were not measured.
-
- Numerics: default tolerance holds for bf16, nothing loosened; the
- largest deviation measured is one bf16 ulp (0.03125) at 2048 and 8192 rows
- of 5120. The `+ 1` is measured: against a plain-`w` reference the output is
- ~4 away, so the cells tell the two apart.
-
- Hard limits
- T * H contiguous inputs above 2^31 - 1 elements route to the strided
- (int64-offset) kernel boundary 2^31 read from the dispatch, not probed
- """
-
- def __call__(self, x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
- return torch.ops.trtllm.flashinfer_gemma_rmsnorm(x, weight, eps)
-
- def reference(self, x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
- """fp32 throughout, `(1 + w)` formed in fp32, one rounding at the end."""
- xf = x.float()
- normed = xf * torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps)
- return (normed * (1.0 + weight.float())).to(x.dtype)
-
- def is_valid(self, x: torch.Tensor, weight: torch.Tensor, eps: float) -> None:
- assert x.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2), (
- f"{x.dtype} is accepted by the kernel's dtype table and returns a "
- "plausible fp8 result; dequantize before normalizing"
- )
-
-
-flashinfer_gemma_rmsnorm = FlashinferGemmaRmsnorm()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py
deleted file mode 100644
index 35201417be8a..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py
+++ /dev/null
@@ -1,116 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""LayerNorm over the last dim -- mean subtraction, biased variance, weight and
-bias -- in one Triton kernel, through the FLA module callable
-`tensorrt_llm._torch.modules.fla.layernorm_gated.layernorm_fn`.
-
-Tier-2 module callable, not a `torch.ops.trtllm.*` op: the registered norm
-family is RMS-norm throughout, and `fused_adaptive_layernorm` refuses every
-width but 5120. One call reshapes to 2-D, allocates the output and launches
-`_layer_norm_fwd_1pass_kernel` once.
-"""
-
-from __future__ import annotations
-
-from typing import Optional
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.modules.fla.layernorm_gated import layernorm_fn as _layernorm_fn
-
-#: The vision tower's norm1 / norm2 / merger.norm: width 1152 with a bias, eps
-#: 1e-6. The merger norm runs before the 4-to-1 merge, so it too sees 1152.
-WIDTH = 1152
-EPS = 1e-6
-
-
-class LayerNorm(OpWrapper):
- """`(x - mean) / sqrt(var + eps) * weight + bias` per row of the last dim,
- variance over `N` (not `N - 1`), new contiguous tensor in `x`'s dtype.
-
- Every operand is widened to fp32 on load, the reduction and the affine run
- in fp32, and only the store rounds. `bias=None` drops the additive term (a
- compile-time branch of the kernel, not a zero bias).
-
- Fusion boundary: the norm and its affine only -- no residual, no gate, no
- output quantization. The callable's `z`, `group_size`, `norm_before_gate`
- and `is_rms_norm` are fixed at "no gate, one group, LayerNorm"; the gated
- and RMS forms are different computations, not options of this entry.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: shape, dtype (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 = """
- Metadata consumed: none. Inputs are not mutated; two scratch fp32 [M]
- tensors (mean, rstd) are allocated per call and discarded.
-
- Layout: `x.stride(-1) == 1` is read in place through its row stride; any
- other `x`, and non-contiguous `weight`/`bias`, are silently `.contiguous()`d
- first (one extra copy). `weight`/`bias` dtype is free.
-
- Certified: contiguous bf16 `[rows, 1152]` `x` at 1, 1024 and 32768 rows
- with a bf16 weight and bias, as the vision tower calls it. Not certified:
- strided or 3-D `x`, `bias=None`, fp16/fp32 activations, weight/bias in
- another dtype than `x`.
-
- Loud failures (left to the callable): weight or bias length != N raises a
- bare AssertionError with no message; `next_pow2(N) * x.element_size() >
- 64KB` raises "This layer norm doesn't support feature dim >= 64KB" (N =
- 65536 in bf16); CPU tensors fail before or at the Triton launch.
-
- Accuracy, 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. 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__(
- self,
- x: torch.Tensor,
- weight: torch.Tensor,
- bias: Optional[torch.Tensor],
- eps: float,
- ) -> torch.Tensor:
- return _layernorm_fn(x, weight, bias, z=None, eps=eps)
-
- def reference(
- self,
- x: torch.Tensor,
- weight: torch.Tensor,
- bias: Optional[torch.Tensor],
- eps: float,
- ) -> torch.Tensor:
- """fp64 from primitive ops (not `F.layer_norm`), rounded once to `x.dtype`."""
- xd = x.double()
- mean = xd.mean(dim=-1, keepdim=True)
- var = (xd - mean).square().mean(dim=-1, keepdim=True)
- out = (xd - mean) / torch.sqrt(var + eps) * weight.double()
- if bias is not None:
- out = out + bias.double()
- return out.to(x.dtype)
-
-
-layer_norm = LayerNorm()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py
deleted file mode 100644
index fb79cf780dfd..000000000000
--- 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
+++ /dev/null
@@ -1,199 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Gated RMS normalization with the gate read from a token-major 3-D view,
-with optional fused static-scale FP8 output quantization."""
-
-from __future__ import annotations
-
-from typing import Optional
-
-import torch
-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 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")
-#: Prefill-sized launches (at least this many rows, i.e. 683 tokens of 48 heads) run
-#: the multi-row kernel at `_PREFILL_ROWS` rows a program: at 16384 tokens 94 vs 143
-#: us on GB300 (traffic roofline ~81 us), where 4 rows a program leave 196608 small
-#: programs.
-_PREFILL_MIN_ROWS = 32768 # capacity: fixed: a dispatch threshold, not a buffer size
-_PREFILL_ROWS = 8
-
-
-def _prefill_multirow(x: torch.Tensor, z: torch.Tensor, gate_activation: str) -> bool:
- # A `z` that does not cover `x` goes to the registered op, which rejects it.
- n = x.shape[-1]
- return (
- x.shape[0] >= _PREFILL_MIN_ROWS
- and x.shape[0] == z.shape[0] * z.shape[1]
- and z.shape[2] == n
- and x.stride(-1) == 1
- and z.stride(2) == 1
- and z.stride(1) == n
- and n <= 256
- and n & (n - 1) == 0
- and gate_activation in ("silu", "sigmoid")
- )
-
-
-def _prefill_launch(x, z, weight, eps, gate_activation):
- """The multi-row kernel at `_PREFILL_ROWS` rows a program (no PDL: the grid covers every SM)."""
- rows, n = x.shape
- out = torch.empty_like(x)
- _layernorm_gated._rms_norm_gated_fwd_multirow_kernel[(triton.cdiv(rows, _PREFILL_ROWS),)](
- x, out, weight, z, None, None, x.stride(0), out.stride(0), z.stride(0), rows, eps,
- N=n, ROWS=_PREFILL_ROWS, HEADS_PER_TOK=z.shape[1], GATE_SIGMOID=gate_activation == "sigmoid",
- WEIGHT_IS_DELTA=False, LAUNCH_WITH_PDL=False, num_warps=4,
- ) # fmt: skip
- return out
-
-
-class RmsNormGatedTokenMajor(OpWrapper):
- """`rmsnorm(x) * weight * act(z)`, new tensor out, row `r` of `x` gated by `z[r // heads, r % heads]`.
-
- `x` is `[T * heads, N]`; `z` is `[T, heads, N]` (heads = `z.shape[1]`),
- its `(heads, N)` block contiguous per token and its token stride free, so a
- column slice of a wider projection is read in place. `act(v)` is
- `v * sigmoid(v)` for `"silu"` and `sigmoid(v)` for `"sigmoid"`. All fp32,
- rounded once. Norm before gate, plain `w` (not `1 + w`).
-
- With `fp8_scale` (one fp32 element, the downstream linear's static input
- scale) the output is `float8_e4m3fn`: the gated value is rounded to
- `x.dtype`, multiplied by `1 / fp8_scale[0]` and saturated at +-448 --
- byte-identical to the unfused norm -> static-quantize pair.
-
- 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.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- 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),
- ),
- )
-
- note = """
- Metadata consumed: none. Two Triton kernels: a multi-row kernel (4 rows per
- program; 8 from 32768 rows, without fp8 output) when `x.stride(-1) == 1`,
- `z.stride(2) == 1`, `z.stride(1) == N` and N is a power of two <= 256 --
- the checkpoint's N = 128 -- else the generic single-row kernel on a packed
- copy of `z`. PDL on the 4-row path only when the grid is smaller than the
- SM count and `TRTLLM_ENABLE_PDL != 0`; results are bitwise identical with
- it off. The 8-row programs reduce a row's mean square in another order
- than the 4-row ones (one-ulp differences on some outputs).
-
- Loud (measured), so not guarded: 3-D `x` (ValueError), a `z` not covering `x`
- (AssertionError "does not match x shape"; 2-D `z` ValueError), an unknown
- `gate_activation` (ValueError before launch), and on the fallback a
- non-unit `x.stride(-1)` (bare AssertionError). `is_valid` covers what the
- multi-row kernel accepts silently: a weight longer than N (truncated to
- `weight[:N]`) or shorter (read out of bounds), and an `fp8_scale` with
- extra elements (only `[0]` read) or not fp32 (widened silently).
-
- `weight` and `z` are widened to fp32 and need not match `x.dtype`. The
- output is freshly allocated and contiguous; no input is mutated.
-
- Certified: bf16 `x` with the sigmoid gate over the in_proj column view of
- `z` (48 heads of N = 128): 128 tokens on the 4-row kernel, 683 tokens (the
- first 8-row launch) and 16384 tokens on the 8-row kernel. Not certified:
- fp8 output, the silu gate, the generic fallback kernel, fp16/fp32
- operands, a strided `x`.
-
- Numerics: bf16 inside the default tolerance against the fp32 reference,
- with one-ulp differences on ~1e-5 of elements at prefill row counts. fp8
- output byte-identical (measured). Nothing loosened.
-
- Hard limits
- N * itemsize fallback holds a row in 64 KB boundary 65536 B
- probed at bf16 N = 32768 (returns) and 40960 (RuntimeError)
- """
-
- def __call__(
- self,
- x: torch.Tensor,
- z: torch.Tensor,
- weight: torch.Tensor,
- eps: float,
- fp8_scale: Optional[torch.Tensor] = None,
- gate_activation: str = "silu",
- ) -> torch.Tensor:
- 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
- )
-
- def reference(
- self,
- x: torch.Tensor,
- z: torch.Tensor,
- weight: torch.Tensor,
- eps: float,
- fp8_scale: Optional[torch.Tensor] = None,
- gate_activation: str = "silu",
- ) -> torch.Tensor:
- """fp32 throughout; row `t * heads + h` pairs with `z[t, h]`.
-
- With `fp8_scale` the gated value is rounded to `x.dtype` first, then
- multiplied by the fp32 reciprocal of the scale, as the kernel does.
- """
- num_rows, width = x.shape
- xf = x.float()
- y = xf * torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) * weight.float()
- zf = z.reshape(num_rows, width).float()
- gate = torch.sigmoid(zf)
- if gate_activation != "sigmoid":
- gate = gate * zf
- y = y * gate
- if fp8_scale is None:
- return y.to(x.dtype)
- return (y.to(x.dtype).float() * (1.0 / fp8_scale.float())).to(torch.float8_e4m3fn)
-
- def is_valid(
- self,
- x: torch.Tensor,
- z: torch.Tensor,
- weight: torch.Tensor,
- eps: float,
- fp8_scale: Optional[torch.Tensor] = None,
- gate_activation: str = "silu",
- ) -> None:
- assert weight.shape == (x.shape[-1],), (
- f"weight must be 1-D of length x.shape[-1]={x.shape[-1]}; got "
- f"{tuple(weight.shape)} -- the multi-row kernel reads weight[0:N] "
- "unmasked, so a longer weight is silently truncated and a shorter "
- "one is read out of bounds"
- )
- if fp8_scale is not None:
- assert fp8_scale.numel() == 1, (
- f"fp8_scale must hold exactly one element; got {fp8_scale.numel()} "
- "-- the multi-row kernel loads fp8_scale[0] and ignores the rest"
- )
- assert fp8_scale.dtype == torch.float32, (
- f"fp8_scale must be float32; got {fp8_scale.dtype} -- the "
- "multi-row kernel widens it and quantizes against the "
- "lower-precision scale"
- )
-
-
-rms_norm_gated_token_major = RmsNormGatedTokenMajor()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/__init__.py
deleted file mode 100644
index 52a7a9daf028..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/gated_ngram_injection.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/gated_ngram_injection.py
deleted file mode 100644
index af6161baa6ee..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/gated_ngram_injection.py
+++ /dev/null
@@ -1,187 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""PLE gating: one shared n-gram value injected into every hyper-connection
-stream, each scaled by its own key-query gate.
-
-Tier-3 composed from native torch (two upcasts, one reduction, the signed
-square root, one broadcast multiply); it launches no TensorRT-LLM kernel.
-Nothing registered under ``torch.ops.trtllm`` computes a dot-product gate or a
-signed square root (``fused_sigmoid_mul_inplace`` takes a per-element gate and
-could at most stand in for the last multiply). ``Engram.forward`` in
-``tensorrt_llm/_torch/modules/engram/engram.py`` computes this gate as four
-statements inside a forward that also runs the projections, both norms (plain
-``x * w``, not this checkpoint's grouped gemma ``(1 + w)``) and a stateless
-short conv -- not an atomic callable, and a different composite. The built-in
-Qwen4-Exp implementation has its own, not wrapped here; the semantics follow
-the HuggingFace ``Qwen4ExpTextPLELayer``.
-"""
-
-from __future__ import annotations
-
-import math
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: This checkpoint's PLE geometry and its gate floor (`_PLE_CLAMP_MIN`).
-CKPT_HC_COUNT = 4
-CKPT_HIDDEN = 2560
-CKPT_CLAMP_MIN = 1e-6
-
-
-def gate_ratio(out: torch.Tensor, ref: torch.Tensor) -> float:
- """Worst element's distance from `ref` as a fraction of the test's 16-bit
- band: `rtol = eps`, `atol = tiny * eps` (one subnormal spacing). Below 1 is
- inside."""
- info = torch.finfo(ref.dtype)
- ref64 = ref.double()
- return (
- ((out.double() - ref64).abs() / (info.tiny * info.eps + info.eps * ref64.abs()))
- .max()
- .item()
- )
-
-
-class GatedNgramInjection(OpWrapper):
- """For token `t`, stream `s`, `D = value.shape[-1]`, with `key_normed` and
- `query_normed` packed stream-major as `hc_count` blocks of `D`::
-
- dot[t, s] = sum_i key_normed[t, s*D + i] * query_normed[t, s*D + i]
- g = dot / sqrt(D)
- signed = sign(g) * sqrt(max(|g|, clamp_min))
- out[t, s*D + i] = sigmoid(signed[t, s]) * value[t, i]
-
- Returns a fresh contiguous `[..., hc_count * D]` in **`value`'s dtype**
- (a bf16 key/query with an fp32 value returns fp32). The value is shared:
- every stream carries a differently scaled copy of the same vector.
- `sign(0) = 0`, so an exactly zero dot gives `sigmoid(0) = 0.5` in every
- dtype; `clamp_min` only replaces a small nonzero magnitude. The whole
- chain runs in fp32 and rounds once.
-
- Fusion boundary: the per-stream dot, scale, floor, signed root, sigmoid
- and broadcast multiply. The `key_proj` / `value_proj` GEMMs, both grouped
- gemma norms producing the operands, the short conv over the result and
- the add into the stream are the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: kind "ckpt" (default) -- operands with this checkpoint's norm-gamma
- # and value statistics: leading, value_dtype (bf16), clamp_min. kind
- # "single_channel" -- one nonzero channel per stream so each dot is an exact
- # product under the test's control: dots, value_dtype, clamp_min.
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="checkpoint 4x2560 bf16, 33 tokens / MTP verify 8 x 4 + 1",
- spec=dict(leading=(33,)),
- ),
- Cell(
- why="checkpoint 4x2560 bf16, prefill 8192 tokens, largest measured",
- spec=dict(leading=(8192,)),
- ),
- Cell(
- why="dots 0, +-5e-8 (inside the floor), 0.05, fp32 value",
- spec=dict(
- kind="single_channel",
- dots=(0.0, 5e-8, -5e-8, 0.05),
- value_dtype=torch.float32,
- clamp_min=CKPT_CLAMP_MIN,
- ),
- ),
- )
-
- note = """
- Metadata consumed: none; every token row and stream independent. Call
- site: once per forward, on the PLE layer, positional `clamp_min` =
- `_PLE_CLAMP_MIN` (1e-6), `key_normed` / `query_normed` `[T, 10240]` and
- `value` `[T, 2560]`, all bf16.
-
- Preconditions: `key_normed.shape[-1] == hc_count * value.shape[-1]` (the
- unflatten raises otherwise, which also catches a wrong `hc_count`). Leading
- dims must match exactly (`is_valid`); measured, a single query or value
- row broadcasts silently over every key row.
-
- fp32 rather than the HF reference's bf16 chain: against the unrounded
- float64 the fp32 chain sat at 0.498 of a one-epsilon bf16 band, the
- literal bf16 chain at 4.693 (4.81 against `reference`), and 31.5% of the
- output's bf16 elements differ between the two -- a bitwise comparison
- against a bf16 implementation will not match.
-
- The floor is invisible in bf16: a sub-floor dot moves the sigmoid by ~5e-4
- of its value, measurable in fp32 and below half a bf16 ulp
- (`clamp_min` 1e-6 and 0.0 bitwise identical with a bf16 `value`).
-
- Every near-miss tried (no root, unsigned root, no / misplaced / `1/D`
- scale, sigmoid before the root, tanh, stream-minor packing) stays bounded
- by `|value|`; they land 105x..1115x outside the band.
-
- CUDA-graph capture certified, bitwise equal to eager; deterministic.
- Cost, not gated: 0.114 ms at 8 tokens, 0.128 at 512, 0.495 at 4096,
- 0.956 at 8192 (launch-bound below ~512); transient allocation 3.56x the
- operand bytes (the fp32 upcasts).
-
- Hard limits: not probed (largest operand 8192 x 10240 = 2^26.3 elements).
- """
-
- def __call__(
- self,
- key_normed: torch.Tensor,
- query_normed: torch.Tensor,
- value: torch.Tensor,
- hc_count: int,
- clamp_min: float = 1e-6,
- ) -> torch.Tensor:
- activation_dtype = value.dtype
- hidden_size = value.shape[-1]
-
- keys = key_normed.unflatten(-1, (hc_count, hidden_size)).float()
- queries = query_normed.unflatten(-1, (hc_count, hidden_size)).float()
-
- gate = (keys * queries).sum(dim=-1, keepdim=True) / math.sqrt(hidden_size)
- gate = gate.abs().clamp_min(clamp_min).sqrt() * gate.sign()
-
- gated = torch.sigmoid(gate) * value.float().unsqueeze(-2)
- return gated.flatten(-2).to(activation_dtype)
-
- def reference(
- self,
- key_normed: torch.Tensor,
- query_normed: torch.Tensor,
- value: torch.Tensor,
- hc_count: int,
- clamp_min: float = 1e-6,
- ) -> torch.Tensor:
- """float64 -- the entry already computes in fp32, so an fp32 reference
- would be the same arithmetic -- rounded once to `value`'s dtype."""
- hidden = value.shape[-1]
- keys = key_normed.double().unflatten(-1, (hc_count, hidden))
- queries = query_normed.double().unflatten(-1, (hc_count, hidden))
- gate = (keys * queries).sum(dim=-1, keepdim=True) / math.sqrt(hidden)
- gate = gate.abs().clamp_min(clamp_min).sqrt() * gate.sign()
- out = torch.sigmoid(gate) * value.double().unsqueeze(-2)
- return out.flatten(-2).to(value.dtype)
-
- def is_valid(
- self,
- key_normed: torch.Tensor,
- query_normed: torch.Tensor,
- value: torch.Tensor,
- hc_count: int,
- clamp_min: float = 1e-6,
- ) -> None:
- # Both measured: a [1, width] query beside [64, width] keys, or a
- # [1, D] value, returns 64 rows computed from the wrong row.
- assert key_normed.shape == query_normed.shape, (
- f"key_normed {tuple(key_normed.shape)} and query_normed "
- f"{tuple(query_normed.shape)} must have the same shape; a mismatch "
- f"broadcasts silently instead of raising"
- )
- assert value.shape[:-1] == key_normed.shape[:-1], (
- f"value {tuple(value.shape)} must carry the same leading dims as "
- f"key_normed {tuple(key_normed.shape)}; a mismatch broadcasts silently "
- f"instead of raising"
- )
-
-
-gated_ngram_injection = GatedNgramInjection()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/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
deleted file mode 100644
index a6f8c97dbd5e..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ngram_token_history.py
+++ /dev/null
@@ -1,290 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""PLE n-gram id history: a packed token stream becomes one n-gram id window
-per token, its out-of-slice predecessors read from -- and its tail published
-back to -- the per-slot n-gram context pool.
-
-Tier-3 composed from native torch integer ops (~20 small launches); exact, so
-gated bitwise. Nothing registered under ``torch.ops.trtllm`` assembles an
-n-gram id window, and the one kernel-plane callable in the family
-(``EngramHashProvider.compute_hashes`` / ``NgramHashMapping._get_ngram_hashes``
-in ``tensorrt_llm/_torch/modules/engram/engram.py``) has no cross-step state:
-it left-pads every chunk from a pad id, so a decode step there loses the
-sequence's whole context. The built-in Qwen4-Exp implementation keeps its own,
-not wrapped here.
-"""
-
-from __future__ import annotations
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: This checkpoint's PLE constants: `ngram_size` 3, so a 2-id context row
-#: (`Qwen4ExpPLECacheParams.ngram_context_len`).
-CKPT_NGRAM_SIZE = 3
-CKPT_EOS_TOKEN_ID = 248044
-CKPT_VOCAB_SIZE = 248320
-#: 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
-
-
-def _scattered(num_seqs: int) -> tuple[int, ...]:
- """`num_seqs` distinct slots over the whole engine pool, both ends in."""
- slots = ((0, ENGINE_SLOTS - 1) + tuple(i * 13 + 4 for i in range(num_seqs)))[:num_seqs]
- assert len(set(slots)) == num_seqs and max(slots) < ENGINE_SLOTS
- return slots
-
-
-def _cell(why: str, lengths, slots, carries, pool_slots: int, **spec) -> Cell:
- return Cell(
- why=why,
- spec=dict(
- lengths=tuple(lengths),
- slots=tuple(slots),
- carries=tuple(carries),
- pool_slots=pool_slots,
- **spec,
- ),
- )
-
-
-#: The target's own argument dtypes: int32 ids and slots, int64
-#: `query_start_loc_long` (or the int64 `cu_seqlens_from_seq_lens` rebuild).
-_CALL_SITE = dict(ids_dtype=torch.int32, index_dtype=torch.int32, qsl_dtype=torch.int64)
-
-
-class NgramTokenHistory(OpWrapper):
- """Row `t` of the returned window is `[c[t-n+1], ..., c[t-1], c[t]]` --
- **newest id last**, so column `ngram_size - 1` is the token's own id --
- where `c` is the token's **sequence's** id stream: for sequence `s` over
- packed tokens `[begin, end)` with `slot = state_indices[s]`, a predecessor
- at shift `k` is `input_ids[t - k]` if `t - begin >= k`, else column
- `ngram_size - 1 + (t - begin) - k` of the carry row. The carry row is
- `context_pool[slot]` (oldest id first) when `has_initial_states[s]`, and
- all `eos_token_id` when not.
-
- Then, **in place**, each named slot's row becomes the last `ngram_size - 1`
- ids of the sequence's history including the new tokens (`window[end - 1,
- 1:]`); a zero-token sequence republishes its carry (so a fresh one is
- reseeded to EOS). Unnamed slots are bitwise untouched.
-
- The ids are **raw**: the only EOS written is the fresh-sequence seed, and
- an in-stream EOS passes through -- the hash downstream owns the
- segmentation, and `[eos, .., eos, id]` is what it treats as a sequence
- start. Seeding with 0 would be silently wrong: 0 is a real token.
-
- Returns a fresh contiguous int64 `[num_tokens, ngram_size]` window;
- `context_pool` is the only argument modified.
-
- Fusion boundary: token-to-sequence attribution, carry read, seed, the
- per-shift select and the publish. The hash, the table gather and slot
- assignment are the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: lengths, slots, carries, pool_slots (the cache manager's width, not
- # the batch's), 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, ...] = (
- _cell(
- "MTP verify p=8: 32 decodes x 8 rows, qsl rebuilt",
- (8,) * 32,
- _scattered(32),
- (True,) * 32,
- ENGINE_SLOTS,
- qsl_from="seq_lens",
- **_CALL_SITE,
- ),
- _cell(
- "MTP verify mixed: 37-token context + 8 x 3 rows",
- (37,) + (3,) * 8,
- _scattered(9),
- (False,) + (True,) * 8,
- ENGINE_SLOTS,
- qsl_from="seq_lens",
- eos_density=0.05,
- **_CALL_SITE,
- ),
- _cell(
- "zero-token rows: carry kept, fresh reseeded",
- (2, 0, 0),
- (0, 1, 2),
- (True, True, False),
- 3,
- ),
- )
-
- note = """
- Metadata consumed: none -- batch layout, pool, slot map and carry flags
- are explicit. Call site, once per forward on the PLE layer:
- `context_pool` the whole `MambaRole.PLE_NGRAM_CONTEXT` buffer
- (`[slots, ngram_size - 1]` int64, allocated by the V2 hybrid cache
- manager), `state_indices` the int32 slot vector the GDN mixer uses,
- `has_initial_states` the context prefix padded with True for the
- generation rows (the engine's own flags cover the context prefix only),
- `query_start_loc` int64 -- `query_start_loc_long`, or under MTP verify the
- `cu_seqlens_from_seq_lens` rebuild, because the metadata's vector has one
- row per generation sequence and would mis-segment silently.
-
- Preconditions:
- - Every slot a real one, `0 <= slot < slots`. The carry read
- bounds-checks, so -1 or an index past the pool is a device-side assert
- that poisons the CUDA context (observed in its own process). **A -1
- pad sentinel is not supported**; give padded rows real slots.
- - Slots distinct: two sequences on one slot is accepted and the publish
- is unordered (data; not guarded).
- - `query_start_loc[-1] == input_ids.shape[0]`: trailing tokens are
- attributed to the last sequence (data; not guarded).
- - `context_pool` int64 and `ngram_size - 1` wide, `has_initial_states`
- bool, `state_indices` and `input_ids` 1-D, `ngram_size >= 2`,
- `num_tokens >= 1`, `num_seqs >= 1`: each refused loudly.
- - The two length checks in `is_valid` are the silent ones.
- - Token ids are not range-checked (2^62 and -5 pass through verbatim).
-
- A wrong window is a window of valid ids: newest-first order, a zero
- seed, ignored carry flags, packing across sequence boundaries, a rotated
- carry read and publishing the head all hash to valid rows. Chunked
- prefill (96 tokens split at 41) is bitwise equal to one-shot, window and
- pool. At 461 slots the result is bitwise identical to an 11-slot pool
- with the same slot content.
-
- 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 4096 in timing).
- """
-
- def __call__(
- self,
- input_ids: torch.Tensor,
- context_pool: torch.Tensor,
- state_indices: torch.Tensor,
- has_initial_states: torch.Tensor,
- query_start_loc: torch.Tensor,
- eos_token_id: int,
- ngram_size: int,
- ) -> torch.Tensor:
- ids = input_ids.long()
- num_tokens = ids.shape[0]
- num_seqs = state_indices.shape[0]
- context_len = ngram_size - 1
-
- device = ids.device
- starts = query_start_loc.long()
- slots = state_indices.long()
-
- # A copy, so the in-place publish cannot disturb rows still being read
- # and a zero-token sequence has something to write back.
- carried = context_pool.index_select(0, slots)
- carried = torch.where(
- has_initial_states.reshape(-1, 1),
- carried,
- carried.new_full((), eos_token_id),
- )
-
- # Device-side and static-shape, so capturable; `repeat_interleave` on
- # device lengths would make the shape data-dependent and sync.
- positions = torch.arange(num_tokens, device=device)
- seq_id = torch.searchsorted(starts[1:].contiguous(), positions, right=True).clamp_max(
- num_seqs - 1
- )
- offset = positions - starts[seq_id]
-
- # One column per shift, oldest first: the packed vector once the token
- # is `shift` into its slice, the carry row (virtual positions
- # -context_len .. -1) before that.
- columns = []
- for shift in range(context_len, 0, -1):
- from_packed = ids.index_select(0, (positions - shift).clamp_min(0))
- pool_column = (context_len + offset - shift).clamp(0, context_len - 1)
- from_pool = carried[seq_id, pool_column]
- columns.append(torch.where(offset >= shift, from_packed, from_pool))
- columns.append(ids)
- window = torch.stack(columns, dim=1)
-
- # The last token's window minus its oldest id, or the carry when empty.
- ends = starts[1:]
- last_token = (ends - 1).clamp_min(0)
- published = torch.where(
- (ends > starts[:-1]).reshape(-1, 1),
- window.index_select(0, last_token)[:, 1:],
- carried,
- )
- context_pool[slots] = published
- return window
-
- def reference(
- self,
- input_ids: torch.Tensor,
- context_pool: torch.Tensor,
- state_indices: torch.Tensor,
- has_initial_states: torch.Tensor,
- query_start_loc: torch.Tensor,
- eos_token_id: int,
- ngram_size: int,
- ) -> tuple[torch.Tensor, torch.Tensor]:
- """`(window, context_pool)` as the call leaves them, from each
- sequence's own id stream materialized as a Python list (carry, then
- its slice) and read straight out; mutates nothing."""
- context_len = ngram_size - 1
- ids = input_ids.tolist()
- pool = context_pool.tolist()
- starts = query_start_loc.tolist()
- carries = has_initial_states.tolist()
- window = [[0] * ngram_size for _ in ids]
- published = [list(row) for row in pool]
- for s, slot in enumerate(state_indices.tolist()):
- carried = list(pool[slot]) if carries[s] else [eos_token_id] * context_len
- begin, end = starts[s], starts[s + 1]
- stream = carried + ids[begin:end]
- for step in range(end - begin):
- window[begin + step] = stream[step : step + ngram_size]
- published[slot] = stream[len(stream) - context_len :]
- device = input_ids.device
- return (
- torch.tensor(window, dtype=torch.long, device=device).reshape(len(ids), ngram_size),
- torch.tensor(published, dtype=torch.long, device=device).reshape(context_pool.shape),
- )
-
- def is_valid(
- self,
- input_ids: torch.Tensor,
- context_pool: torch.Tensor,
- state_indices: torch.Tensor,
- has_initial_states: torch.Tensor,
- query_start_loc: torch.Tensor,
- eos_token_id: int,
- ngram_size: int,
- ) -> None:
- # Both measured silent; the over-long direction of each raises instead.
- num_seqs = state_indices.shape[0]
- assert has_initial_states.shape[0] == num_seqs, (
- f"has_initial_states must be one flag per sequence; got "
- f"{tuple(has_initial_states.shape)} against {num_seqs} sequences -- a "
- f"length-1 flag tensor is broadcast over the whole batch rather than "
- f"rejected"
- )
- assert query_start_loc.shape[0] == num_seqs + 1, (
- f"query_start_loc must be [num_seqs + 1] = [{num_seqs + 1}] cumulative "
- f"token counts; got {tuple(query_start_loc.shape)} -- a short one "
- f"misattributes tokens to the wrong sequence rather than raising"
- )
-
-
-ngram_token_history = NgramTokenHistory()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/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
deleted file mode 100644
index bb2814cac29b..000000000000
--- 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
+++ /dev/null
@@ -1,683 +0,0 @@
-# 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)`, `[T, C*D]` in `hc`'s dtype: fresh, or the
- two row-strided buffers of `out` (neither overlapping `hc`) written in place.
- `conv_states` is advanced in place. When `conv_input` is given (a
- contiguous `[T, C*D]` buffer in `hc`'s dtype) the conv input `xc` is
- written there, for a caller that has to rebuild the conv state at another
- row of the sequence later (the speculative rollback of a verify step).
-
- Fusion boundary: everything from the key/query norms to the next site's
- norm. The n-gram embedding and the key / value projections are the
- caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: lengths (rows per sequence), pool_slots (default num_seqs + 2),
- # max_rows (max(lengths)), carries (None -> the argument is None). The
- # checkpoint's 4 x 2560 bf16 geometry throughout.
- CELLS: tuple[Cell, ...] = (
- _cell("decode 128 x 1 row at the 461-slot pool", (1,) * 128, pool_slots=ENGINE_SLOTS),
- _cell("MTP3 verify: 64 sequences x 4 rows", (4,) * 64, pool_slots=ENGINE_SLOTS),
- # Context-shaped (prefill / mixed) steps: the packed-row conv.
- _cell("prefill 8192 rows, fresh", (8192,), carries=(False,), pool_slots=ENGINE_SLOTS),
- _cell(
- "mixed step: 2 contexts + 64 decode rows",
- (3001, 5000) + (1,) * 64,
- carries=(False, True) + (True,) * 64,
- pool_slots=ENGINE_SLOTS,
- ),
- )
-
- note = """
- Metadata consumed: none. Call site: the decode target's PLE block, once
- per forward: `hc` the `[T, 10240]` stream, `key` the key projection
- `[T, 10240]` and `value` the value projection `[T, 2560]` (both cuBLAS,
- bf16), the three PLE norm weights and the next site's attention-side norm
- weight as contiguous `[10240]`, the `[10240, 1, 4]` conv weight, the whole
- `[slots, 10240, 9]` conv pool, int32 slots, int64 `query_start_loc_long`
- (the rebuild under MTP verify), `max_rows = gen_p`.
-
- Preconditions (data; not guarded): at most `max_rows` rows per sequence;
- distinct live slots; `query_start_loc[-1] == T`. Operands 2-D with unit
- column stride; weights contiguous. Certified at the checkpoint's 4 x 2560
- bf16 geometry with every slot live (the target never passes `pad_slot_id`),
- at a 128 x 1-row decode, a 64 x 4-row MTP3 verify, an 8192-row fresh
- prefill and a mixed 2-context + 64-decode step, all on the 461-slot pool.
- Not certified: sequences shorter than `max_rows` on the per-sequence path
- (a verify step gives every sequence `max_rows` rows).
-
- 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,
- conv_input: Optional[torch.Tensor] = None,
- out: Optional[Tuple[torch.Tensor, 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]
- if out is None:
- hc_out = torch.empty((rows, width), dtype=hc.dtype, device=hc.device)
- normed = torch.empty((rows, width), dtype=hc.dtype, device=hc.device)
- else:
- hc_out, normed = out
- 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 = (
- conv_input
- if conv_input is not None
- else 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,
- conv_input: Optional[torch.Tensor] = None,
- out: Optional[Tuple[torch.Tensor, 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,
- conv_input: Optional[torch.Tensor] = None,
- out: Optional[Tuple[torch.Tensor, 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"
- )
- assert conv_input is None or (
- conv_input.shape == hc.shape
- and conv_input.dtype == hc.dtype
- and conv_input.is_contiguous()
- ), "conv_input must be a contiguous buffer shaped and typed as hc"
- for t in out or ():
- assert t.shape == hc.shape and t.dtype == hc.dtype and t.stride(1) == 1, (
- "out buffers must be shaped and typed as hc, unit column stride"
- )
-
-
-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
deleted file mode 100644
index f8c2236854cf..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_ngram_embed.py
+++ /dev/null
@@ -1,560 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""PLE n-gram embedding for decode-shaped batches: id history, n-gram hash,
-table gather and per-tensor fp8 dequantization in one Triton launch, with the
-per-slot n-gram context published back in the same pass.
-
-Tier 5 (Triton written for this target). It replaces, on the decode path,
-the composition ``ngram_token_history`` (~20 launches) + ``ngram_hash`` (~30
-launches) + ``nn.functional.embedding`` + ``dequantize_e4m3_per_tensor`` --
-nothing registered under ``torch.ops.trtllm`` or in the installed FlashInfer hashes
-n-grams or gathers-and-dequantizes, and each of those entries is launch-bound
-at decode (0.3-0.4 ms eager each).
-
-One program per **sequence**, which loops over that sequence's rows: the carry
-read and the publish of one slot happen in one program, so no other program can
-observe a half-published row.
-
-Context-shaped batches (more than ``PER_SEQ_MAX_ROWS`` rows per sequence) take
-a second form: one program per 32 packed rows, each row binary-searching
-``query_start_loc`` for its sequence, then a one-program-per-sequence publish
-launched after it -- two launches, bitwise equal to the first form.
-"""
-
-from __future__ import annotations
-
-from typing import Optional
-
-import torch
-import triton
-import triton.language as tl
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-from ..embedding.ngram_hash import ngram_hash
-from .ngram_token_history import ngram_token_history
-
-#: This checkpoint's PLE geometry: 3-gram, 16 heads of 160 fp8 columns.
-CKPT_NGRAM_SIZE = 3
-CKPT_HEAD_DIM = 160
-CKPT_TABLE_ROWS = 320001536
-CKPT_EOS_TOKEN_ID = 248044
-CKPT_SCALE = 0.00019931793212890625
-#: Largest `ple_ngram` pool seen for this checkpoint (see `ngram_token_history`).
-ENGINE_SLOTS = 461
-#: Above this many rows per sequence (MTP verify carries at most 4) the call
-#: takes the context-shaped path: packed row tiles plus a separate publish.
-PER_SEQ_MAX_ROWS = 4
-_ROWS_BLOCK_T = 32
-#: The row-to-sequence search covers at least 2^8 sequences, so one compiled
-#: variant serves every batch up to that size.
-_MIN_LOG_S = 8
-
-
-@triton.jit
-def _ple_ngram_embed_kernel(
- ids_ptr,
- pool_ptr,
- slots_ptr,
- carry_ptr,
- qsl_ptr,
- table_ptr,
- scale_ptr,
- mult_ptr,
- vocab_ptr,
- offs_ptr,
- out_ptr,
- stride_pool,
- stride_table,
- stride_out,
- eos,
- NGRAM: tl.constexpr,
- HEADS_PER: tl.constexpr,
- HEADS: tl.constexpr,
- BLOCK_H: tl.constexpr,
- HEAD_DIM: tl.constexpr,
- BLOCK_D: tl.constexpr,
- BLOCK_P: tl.constexpr,
- BLOCK_C: tl.constexpr,
- HAS_CARRY: tl.constexpr,
-):
- CTX: tl.constexpr = NGRAM - 1
- s = tl.program_id(0)
- begin = tl.load(qsl_ptr + s).to(tl.int64)
- length = tl.load(qsl_ptr + s + 1).to(tl.int64) - begin
- slot = tl.load(slots_ptr + s).to(tl.int64)
- pool_row = pool_ptr + slot * stride_pool
- eos64 = tl.full([BLOCK_P], eos, tl.int64)
- if HAS_CARRY:
- fresh = tl.load(carry_ptr + s) == 0
-
- t = tl.arange(0, BLOCK_P)
- tmask = t < length
- heads = tl.arange(0, BLOCK_H)
- hmask = heads < HEADS
- order = heads // HEADS_PER + 2
- vocab = tl.load(vocab_ptr + heads, mask=hmask, other=1)
- offs = tl.load(offs_ptr + heads, mask=hmask, other=0)
-
- # Shift 0 is the token itself; shift k reads `k` back in the sequence's own
- # stream, i.e. the packed ids or, before the slice starts, the carry row.
- # Once a shift lands on an EOS, every deeper shift reads EOS: the hash's
- # segment starts one past the most recent EOS strictly before the token.
- mix = tl.load(ids_ptr + begin + t, mask=tmask, other=0).to(tl.int64) * tl.load(mult_ptr)
- mixed = tl.zeros([BLOCK_P, BLOCK_H], dtype=tl.int64)
- hit = t < 0
- for k in tl.static_range(1, NGRAM):
- pos = t - k
- from_ids = tl.load(ids_ptr + begin + pos, mask=tmask & (pos >= 0), other=0).to(tl.int64)
- from_pool = tl.load(pool_row + (CTX + pos), mask=tmask & (pos < 0), other=0)
- if HAS_CARRY:
- from_pool = tl.where(fresh, eos64, from_pool)
- ck = tl.where(pos >= 0, from_ids, from_pool)
- sk = tl.where(hit, eos64, ck)
- hit = hit | (ck == eos64)
- mix = mix ^ (sk * tl.load(mult_ptr + k))
- mixed = tl.where((order == k + 1)[None, :], mix[:, None], mixed)
-
- rem = mixed % vocab[None, :]
- rows = tl.where(rem < 0, rem + vocab[None, :], rem) + offs[None, :]
-
- d = tl.arange(0, BLOCK_D)
- m3 = tmask[:, None, None] & hmask[None, :, None] & (d < HEAD_DIM)[None, None, :]
- q = tl.load(table_ptr + rows[:, :, None] * stride_table + d[None, None, :], mask=m3, other=0.0)
- val = q.to(tl.float32) * tl.load(scale_ptr).to(tl.float32)
- cols = heads[None, :, None] * HEAD_DIM + d[None, None, :]
- tl.store(
- out_ptr + (begin + t)[:, None, None] * stride_out + cols,
- val.to(out_ptr.dtype.element_ty),
- mask=m3,
- )
-
- # Publish the last CTX ids of `carry + slice` -- after every thread's carry
- # read, since the row being overwritten is the one just read.
- c = tl.arange(0, BLOCK_C)
- cmask = c < CTX
- src = length + c
- from_slice = tl.load(ids_ptr + begin + src - CTX, mask=cmask & (src >= CTX), other=0).to(
- tl.int64
- )
- old = tl.load(pool_row + src, mask=cmask & (src < CTX), other=0)
- if HAS_CARRY:
- old = tl.where(fresh, tl.full([BLOCK_C], eos, tl.int64), old)
- published = tl.where(src >= CTX, from_slice, old)
- tl.debug_barrier()
- tl.store(pool_row + c, published, mask=cmask)
-
-
-@triton.jit
-def _seq_of_rows(qsl_ptr, t, num_seqs, LOG_S: tl.constexpr):
- """Per packed row `t`: the largest `s < num_seqs` with `qsl[s] <= t`, i.e.
- the sequence that owns the row (zero-row sequences are skipped over)."""
- s = tl.zeros_like(t)
- for i in tl.static_range(LOG_S):
- cand = s + (1 << (LOG_S - 1 - i))
- ok = cand < num_seqs
- start = tl.load(qsl_ptr + cand, mask=ok, other=0)
- s = tl.where(ok & (start <= t), cand, s)
- return s
-
-
-@triton.jit
-def _ple_ngram_embed_rows_kernel(
- ids_ptr,
- pool_ptr,
- slots_ptr,
- carry_ptr,
- qsl_ptr,
- table_ptr,
- scale_ptr,
- mult_ptr,
- vocab_ptr,
- offs_ptr,
- out_ptr,
- num_rows,
- num_seqs,
- stride_pool,
- stride_table,
- stride_out,
- eos,
- NGRAM: tl.constexpr,
- HEADS_PER: tl.constexpr,
- HEADS: tl.constexpr,
- BLOCK_H: tl.constexpr,
- HEAD_DIM: tl.constexpr,
- BLOCK_D: tl.constexpr,
- BLOCK_T: tl.constexpr,
- LOG_S: tl.constexpr,
- HAS_CARRY: tl.constexpr,
-):
- """Context-shaped batches: one program per BLOCK_T packed rows, each row
- finding its sequence in `query_start_loc`. Window, hash, gather and
- dequantize only; the publish is `_ple_ngram_publish_kernel`, launched
- after, because the rows that read a slot's carry may span two programs."""
- CTX: tl.constexpr = NGRAM - 1
- t = tl.program_id(0).to(tl.int64) * BLOCK_T + tl.arange(0, BLOCK_T)
- tmask = t < num_rows
- s = _seq_of_rows(qsl_ptr, t, num_seqs, LOG_S)
- begin = tl.load(qsl_ptr + s, mask=tmask, other=0).to(tl.int64)
- tpos = t - begin
- slot = tl.load(slots_ptr + s, mask=tmask, other=0).to(tl.int64)
- eos64 = tl.full([BLOCK_T], eos, tl.int64)
- if HAS_CARRY:
- fresh = tl.load(carry_ptr + s, mask=tmask, other=1) == 0
-
- heads = tl.arange(0, BLOCK_H)
- hmask = heads < HEADS
- order = heads // HEADS_PER + 2
- vocab = tl.load(vocab_ptr + heads, mask=hmask, other=1)
- offs = tl.load(offs_ptr + heads, mask=hmask, other=0)
-
- mix = tl.load(ids_ptr + t, mask=tmask, other=0).to(tl.int64) * tl.load(mult_ptr)
- mixed = tl.zeros([BLOCK_T, BLOCK_H], dtype=tl.int64)
- hit = tmask & (t < 0)
- for k in tl.static_range(1, NGRAM):
- pos = tpos - k
- from_ids = tl.load(ids_ptr + t - k, mask=tmask & (pos >= 0), other=0).to(tl.int64)
- from_pool = tl.load(
- pool_ptr + slot * stride_pool + (CTX + pos), mask=tmask & (pos < 0), other=0
- )
- if HAS_CARRY:
- from_pool = tl.where(fresh, eos64, from_pool)
- ck = tl.where(pos >= 0, from_ids, from_pool)
- sk = tl.where(hit, eos64, ck)
- hit = hit | (ck == eos64)
- mix = mix ^ (sk * tl.load(mult_ptr + k))
- mixed = tl.where((order == k + 1)[None, :], mix[:, None], mixed)
-
- rem = mixed % vocab[None, :]
- rows = tl.where(rem < 0, rem + vocab[None, :], rem) + offs[None, :]
-
- d = tl.arange(0, BLOCK_D)
- dmask = d < HEAD_DIM
- m2 = tmask[:, None] & dmask[None, :]
- scale = tl.load(scale_ptr).to(tl.float32)
- for h in tl.static_range(HEADS):
- row_h = tl.sum(tl.where((heads == h)[None, :], rows, 0), axis=1)
- q = tl.load(table_ptr + row_h[:, None] * stride_table + d[None, :], mask=m2, other=0.0)
- val = q.to(tl.float32) * scale
- tl.store(
- out_ptr + t[:, None] * stride_out + (h * HEAD_DIM + d)[None, :],
- val.to(out_ptr.dtype.element_ty),
- mask=m2,
- )
-
-
-@triton.jit
-def _ple_ngram_publish_kernel(
- ids_ptr,
- pool_ptr,
- slots_ptr,
- carry_ptr,
- qsl_ptr,
- stride_pool,
- eos,
- NGRAM: tl.constexpr,
- BLOCK_C: tl.constexpr,
- HAS_CARRY: tl.constexpr,
-):
- """One program per sequence: the last `NGRAM - 1` ids of `carry + slice`
- published to its slot, after every row's carry read (stream order)."""
- CTX: tl.constexpr = NGRAM - 1
- s = tl.program_id(0)
- begin = tl.load(qsl_ptr + s).to(tl.int64)
- length = tl.load(qsl_ptr + s + 1).to(tl.int64) - begin
- slot = tl.load(slots_ptr + s).to(tl.int64)
- pool_row = pool_ptr + slot * stride_pool
- c = tl.arange(0, BLOCK_C)
- cmask = c < CTX
- src = length + c
- from_slice = tl.load(ids_ptr + begin + src - CTX, mask=cmask & (src >= CTX), other=0).to(
- tl.int64
- )
- old = tl.load(pool_row + src, mask=cmask & (src < CTX), other=0)
- if HAS_CARRY:
- fresh = tl.load(carry_ptr + s) == 0
- old = tl.where(fresh, tl.full([BLOCK_C], eos, tl.int64), old)
- published = tl.where(src >= CTX, from_slice, old)
- tl.debug_barrier()
- tl.store(pool_row + c, published, mask=cmask)
-
-
-def _cell(why: str, lengths, **spec) -> Cell:
- return Cell(why=why, spec=dict(lengths=tuple(lengths), **spec))
-
-
-class PleNgramEmbed(OpWrapper):
- """For sequence `s` over packed rows `[b, e)` with `slot = state_indices[s]`,
- its id stream is `carry + input_ids[b:e]`, where `carry` is
- `context_pool[slot]` (oldest first) when `has_initial_states[s]` -- or
- always, when `has_initial_states` is None -- and `ngram_size - 1` copies of
- `eos_token_id` when not. Row `t` of that slice gets the window
- `stream[t : t + ngram_size]` (newest last), exactly `ngram_token_history`'s;
- each window hashes to one global row per head exactly as `ngram_hash` with
- `num_query_tokens=1`; and::
-
- out[t, h * D : (h + 1) * D] = to(scale.dtype)(float(table[row[t, h]]) * float(scale))
-
- with `D = table.shape[1]` -- `dequantize_e4m3_per_tensor`'s one rounding.
- Then, **in place**, `context_pool[slot]` becomes the last `ngram_size - 1`
- ids of the stream (a zero-row sequence republishes its carry). Unnamed slots
- are untouched.
-
- Returns a fresh contiguous `[num_tokens, heads * D]` in `scale`'s dtype.
- `context_pool` is the only argument modified.
-
- Fusion boundary: carry read and seed, window, hash, gather, dequantize,
- publish. The key/value projections that consume the result are the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: lengths (rows per sequence), pool_slots (default num_seqs + 2),
- # carries (None -> the argument is None), max_rows (default max(lengths)),
- # eos_density (0.0), table "small" | "host" | "ckpt" (the full
- # 320001536-row device table, rows filled where touched). The checkpoint's
- # 3-gram of 8 heads per order, int32 ids.
- CELLS: tuple[Cell, ...] = (
- _cell(
- "decode 128 x 1 row, host-resident table, 461 slots",
- (1,) * 128,
- pool_slots=ENGINE_SLOTS,
- table="host",
- ),
- _cell(
- "decode 128 x 1 row, the full checkpoint-height table",
- (1,) * 128,
- pool_slots=ENGINE_SLOTS,
- table="ckpt",
- eos_density=0.05,
- ),
- _cell(
- "MTP3 verify: 64 sequences x 4 rows, 461 slots",
- (4,) * 64,
- pool_slots=ENGINE_SLOTS,
- table="host",
- eos_density=0.05,
- ),
- # Context-shaped (prefill / mixed) steps: the packed-row form.
- _cell(
- "mixed step: 2 contexts + 64 decode rows",
- (3001, 5000) + (1,) * 64,
- carries=(False, True) + (True,) * 64,
- pool_slots=ENGINE_SLOTS,
- eos_density=0.05,
- ),
- )
-
- note = """
- Metadata consumed: none -- batch layout, pool, slot map are explicit.
- Call site: the decode target's PLE block, once per forward:
- `input_ids` int32 `[T]`, the whole `[slots, 2]` int64 n-gram pool,
- int32 `state_indices[:batch]`, `has_initial_states=None` (every decode
- sequence carries), int64 `query_start_loc_long` (or the
- `cu_seqlens_from_seq_lens` rebuild under MTP verify), the
- `[320001536, 160]` fp8 table (host-resident by default), the `[1, 1, 1]`
- bf16 scale, `max_rows = gen_p`.
-
- Preconditions (data; not guarded -- each would need a sync):
- - every sequence has at most `max_rows` rows; rows past it are neither
- written nor published;
- - slots real (`0 <= slot < slots`) and distinct; -1 is not supported;
- - `query_start_loc[-1] == num_tokens`;
- - token id * max multiplier < 2^63 (int64 wraps, as in `ngram_hash`).
-
- Certified at a 128 x 1-row decode on a host-resident table and on the
- full checkpoint-height device table, an MTP3 verify step of 64 x 4 rows
- on the per-sequence path, and a mixed step of 2 contexts (one fresh, one
- carrying) + 64 decode rows on the packed-row form. Not certified:
- zero-row sequences.
-
- 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/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/__init__.py
deleted file mode 100644
index 52a7a9daf028..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/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
deleted file mode 100644
index 61edc7eb2a1d..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/bf16_to_fp32.py
+++ /dev/null
@@ -1,80 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Widening bf16 -> fp32 of a contiguous tensor in one vectorized kernel: the logits
-processor's cast of the step's `[rows, vocab]` logits.
-
-Tier 5 (Triton written for this target). torch's `.float()` runs this as an unrolled,
-unvectorized `direct_copy` at under half the bandwidth; the values are the same bit for
-bit.
-"""
-
-from __future__ import annotations
-
-import torch
-import triton
-import triton.language as tl
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: This checkpoint's vocabulary: the width of the logits rows.
-VOCAB = 248320
-#: Elements a program (8 warps: 16 a thread, two 16-byte loads).
-_BLOCK = 4096
-_WARPS = 8
-
-
-@triton.jit
-def _bf16_to_fp32_kernel(x_ptr, y_ptr, n, BLOCK: tl.constexpr):
- offs = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
- mask = offs < n
- tl.store(y_ptr + offs, tl.load(x_ptr + offs, mask=mask).to(tl.float32), mask=mask)
-
-
-class Bf16ToFp32(OpWrapper):
- """`x.float()` for a contiguous bf16 `x`: a fresh contiguous fp32 tensor of `x`'s
- shape. Widening is exact, so every element is `x`'s value; NaN payloads widen as
- torch widens them (the 16 bits shifted up).
-
- Fusion boundary: the cast only; the head's GEMM before it and the sampler after it
- are the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(why="decode logits: 128 rows of the 248320 vocabulary", spec=dict(shape=(128, VOCAB))),
- Cell(why="one row; a ragged last program", spec=dict(shape=(1, VOCAB))),
- Cell(
- why="every bf16 bit pattern: NaNs, infinities, subnormals, signed zeros",
- spec=dict(shape=(1, 65536), source="all_codes"),
- ),
- )
-
- note = """
- Metadata consumed: none; stateless. One Triton kernel, `cdiv(numel, 4096)` programs
- of 8 warps over the flat tensor. Call site: the target's logits processor
- (`_WideningLogitsProcessor`), on the head's bf16 logits; the shell's `.float()`
- after it is then a no-op.
-
- GB300, graph replay: 128 x 248320 31.3 us against 82.1 for `.float()`, 124 rows
- 30.7 against 79.9, 33 rows 8.3 against 16.5, one row 4.1 against 4.2.
-
- Preconditions: bf16, contiguous (the kernel reads the flat storage).
- """
-
- def __call__(self, x: torch.Tensor) -> torch.Tensor:
- y = torch.empty(x.shape, dtype=torch.float32, device=x.device)
- n = x.numel()
- if n:
- _bf16_to_fp32_kernel[(triton.cdiv(n, _BLOCK),)](x, y, n, BLOCK=_BLOCK, num_warps=_WARPS)
- return y
-
- def reference(self, x: torch.Tensor) -> torch.Tensor:
- return x.float()
-
- def is_valid(self, x: torch.Tensor) -> None:
- assert x.dtype == torch.bfloat16, "bf16 input"
- assert x.is_contiguous(), "contiguous input: the kernel reads the flat storage"
-
-
-bf16_to_fp32 = Bf16ToFp32()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/dequantize_e4m3_per_tensor.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/dequantize_e4m3_per_tensor.py
deleted file mode 100644
index 7f3922dea06a..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/dequantize_e4m3_per_tensor.py
+++ /dev/null
@@ -1,94 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Per-tensor FP8 dequantization: float8_e4m3fn -> bf16/fp16/fp32, multiplied by one scale."""
-
-from __future__ import annotations
-
-import torch
-
-import tensorrt_llm._torch.custom_ops # noqa: F401 — loads the extension registering torch.ops.tensorrt_llm.*
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: The PLE n-gram table's per-tensor dequantization scale in this checkpoint.
-PLE_SCALE = 0.00019931793212890625
-
-#: Scale values the whole-domain cells sweep: the checkpoint's, unity, a power
-#: of two, an arbitrary one, and the smallest e4m3 subnormal.
-_DOMAIN_SCALES = (PLE_SCALE, 1.0, 0.5, 0.037, 2.0**-9)
-
-
-def _ple(tokens: int) -> dict:
- """`_ple_block`'s call: 16 gathered 160-wide rows per token, one bf16 scale."""
- return dict(shape=(tokens, 16, 160), dtype=torch.bfloat16, scales=(PLE_SCALE,), source="gather")
-
-
-class DequantizeE4m3PerTensor(OpWrapper):
- """`out = to(scales.dtype)(float(input) * float(scales[0]))`, one rounding.
-
- Wraps `torch.ops.tensorrt_llm.dequantize_e4m3_per_tensor` -- the
- `tensorrt_llm` namespace; there is no `trtllm::` alias.
-
- It MULTIPLIES: pass a checkpoint's `weight_scale` straight, never
- reciprocated (the inverse, `static_quantize_e4m3_per_tensor`, divides). The
- output dtype follows `scales`, not `input`: a bf16 scale yields bf16 with no
- cast. Shape is preserved exactly. `scales` has `input`'s rank with every dim
- 1 (`[1, 1, 1]` for the PLE's rank-3 gather).
-
- Fusion boundary: scale-and-widen only. The table gather before it and the
- flatten to `[T, 2560]` after it (a view) are the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(why="_ple_block decode: [1, 16, 160] gathered rows, bf16 scale", spec=_ple(1)),
- Cell(
- why="_ple_block prefill at max_num_tokens: [16384, 16, 160] through the grid-stride "
- "loop",
- spec=_ple(16384),
- ),
- Cell(
- why="every e4m3 byte into bf16, five scale values, NaN codes included",
- spec=dict(
- shape=(1, 256), dtype=torch.bfloat16, scales=_DOMAIN_SCALES, source="all_codes"
- ),
- ),
- )
-
- note = """
- Metadata consumed: none; stateless. One kernel per call,
- `scaleMatrix` at a fixed 1024-block
- grid with a grid-stride loop; there is no vectorized variant, so alignment
- is not an axis.
-
- Preconditions, enforced by the op (`symmetric_dequantize_per_tensor` in
- thop/fp8Op.cpp): `input` float8_e4m3fn, CUDA, contiguous, rank >= 2,
- non-empty (zero rows raise -- an empty batch must branch first; the test
- holds this rejection); `scales` CUDA, bf16 / fp16 / fp32, of `input`'s rank
- with every dim 1.
-
- NaN codes stay NaN; `448 * PLE_SCALE` reads back as bf16 0.08935546875 (the
- all-codes cell).
-
- CUDA-graph capture is certified: replay picks up a new `input` and a new
- `scales` value written in place (the scale is read from memory at replay).
-
- Hard limits
- numel certified to 41,943,040 (16384 tokens); probed to 167,772,160; the
- int32 boundary 2^31 not probed
-
- Not certified: fp16 and fp32 outputs (the target's scale is bf16), ranks
- other than 2 and 3, degenerate scale values (0, negative, inf, NaN) and
- overflow, the op's other loud rejections (dtype, layout, rank, device),
- `TRTLLM_ENABLE_PDL=1`, `scales` on a different device than `input`.
- """
-
- def __call__(self, input: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
- return torch.ops.tensorrt_llm.dequantize_e4m3_per_tensor(input, scales)
-
- def reference(self, input: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
- """One fp32 multiply, one conversion to the scales' dtype."""
- return (input.float() * scales.float().flatten()[0]).to(scales.dtype)
-
-
-dequantize_e4m3_per_tensor = DequantizeE4m3PerTensor()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/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
deleted file mode 100644
index 51e97e8e5444..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/fp8_quantize_1x128.py
+++ /dev/null
@@ -1,131 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Dynamic 1x128 block-wise FP8 quantization: bf16 -> e4m3 data + transposed fp32 block scales."""
-
-from __future__ import annotations
-
-import torch
-
-import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: Largest finite e4m3 magnitude.
-E4M3_MAX = 448.0
-
-#: The kernel's amax floor, as the bf16 constant it compares against:
-#: bfloat16(1e-10) == 1.000444171950221e-10.
-AMAX_FLOOR = float(torch.tensor(1e-10, dtype=torch.bfloat16))
-
-#: The two divisions in the kernel are approximate (`--use_fast_math`), each
-#: within CUDA's documented 2 ulp; the returned scale is two deep.
-SCALE_ULP_GATE = 4
-
-
-def block_scales(input: torch.Tensor, use_ue8m0: bool) -> tuple[torch.Tensor, torch.Tensor]:
- """`(quant_scale, dequant_scale)`, each `[T, ceil(H/128)]` fp32, exactly rounded.
-
- The amax runs in bf16, as the kernel's running maximum does (exact: |x| of
- a bf16 is a bf16), floored at `AMAX_FLOOR`; a short tail block takes the
- amax of its own elements. The kernel multiplies by the first before
- rounding to e4m3 and returns the second.
- """
- t, n = input.shape
- nb = (n + 127) // 128
- mag = torch.nn.functional.pad(input.abs(), (0, nb * 128 - n))
- amax = mag.reshape(t, nb, 128).amax(dim=2).float().clamp(min=AMAX_FLOOR)
- qs = E4M3_MAX / amax
- if not use_ue8m0:
- return qs, 1.0 / qs
- raw = 1.0 / qs
- mant, exp = torch.frexp(raw)
- ds = torch.ldexp(torch.ones_like(raw), torch.where(mant == 0.5, exp - 1, exp))
- return 1.0 / ds, ds
-
-
-#: The MTP draft layer's hidden width, quantized with the plain scale the
-#: block-scale MoE runner wants.
-_CKPT = dict(hidden=2560, use_ue8m0=False)
-
-
-class Fp8Quantize1x128(OpWrapper):
- """Quantize a bf16 `[T, H]` matrix to e4m3 with one fp32 scale per 1x128 row block.
-
- Per token `t` and block `b` of 128 consecutive columns (the last one short
- when `H % 128 != 0`): `amax = max(|x[t, block]|, bf16(1e-10))`,
- `data = e4m3(x * (448 / amax))` saturating at +-448, `scale[b, t] =
- 1 / (448 / amax)`. Returns `(data [T, H] float8_e4m3fn, scales [ceil(H/128),
- T] fp32)`.
-
- Three readings are each a silent wrong answer: `scales` is TRANSPOSED
- (`[NB, T]`, the layout `fp8_block_scale_moe_runner` consumes); it
- MULTIPLIES the code back (dividing is 6780 relative RMS off); and with
- `use_ue8m0=False` it is the plain `amax / 448`, while `use_ue8m0=True`
- rounds every scale up to a power of two -- a different quantization of the
- same input (99% of scales, 96% of codes change), not a layout choice.
-
- Fusion boundary: the amax reduction, scaling, rounding and (on the sm_100
- family) the compaction into `[NB, T]` are inside; no GEMM is fused.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- CELLS: tuple[Cell, ...] = (
- Cell(why="MTPLayer._moe decode: one token at hidden 2560", spec=dict(tokens=1, **_CKPT)),
- Cell(
- why="20 tokens == H/128: the scale buffer is square, its transpose undetectable",
- spec=dict(tokens=20, **_CKPT),
- ),
- Cell(
- why="16384 rows: max_num_tokens of the tuned config",
- spec=dict(tokens=16384, **_CKPT),
- ),
- )
-
- note = """
- Metadata consumed: none. The returned scale layout is chosen by an
- architecture query (`isSM100Family()` in thop/fp8Quantize.cpp): off the
- sm_100 family the op returns a 1-D padded buffer instead. That branch was
- read, not run.
-
- Preconditions, 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 at 1, 20 and 16384 at H 2560 (16384: max_num_tokens of the
- tuned config); above that not probed
- H multiple of 16 (raises otherwise); certified at 2560, the draft layer's
- width
-
- Not certified: `T == 0` (the op accepts it and returns `[0, H]` /
- `[NB, 0]`); a short tail block (`H % 128 != 0`); all-zero blocks (the
- `AMAX_FLOOR` scale), `inf` / NaN inputs and saturation; `use_ue8m0=True`;
- the registered fake.
- """
-
- def __call__(
- self,
- input: torch.Tensor,
- use_ue8m0: bool = False,
- ) -> tuple[torch.Tensor, torch.Tensor]:
- return torch.ops.trtllm.fp8_quantize_1x128(input, use_ue8m0)
-
- def reference(
- self,
- input: torch.Tensor,
- use_ue8m0: bool = False,
- ) -> tuple[torch.Tensor, torch.Tensor]:
- """`(data, scales[NB, T])` from exactly-rounded fp32 divisions."""
- qs, ds = block_scales(input, use_ue8m0)
- cols = qs.repeat_interleave(128, dim=1)[:, : input.shape[1]]
- data = (input.float() * cols).to(torch.float8_e4m3fn)
- return data, ds.t().contiguous()
-
-
-fp8_quantize_1x128 = Fp8Quantize1x128()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/__init__.py
deleted file mode 100644
index 52a7a9daf028..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_fwd.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_fwd.py
deleted file mode 100644
index af83f235c2f7..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_fwd.py
+++ /dev/null
@@ -1,250 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Causal depthwise conv1d over a whole (ragged) batch of sequences, with each
-sequence's convolution state read from and published to a paged state pool in
-the same call."""
-
-from __future__ import annotations
-
-from typing import Optional
-
-import torch
-
-import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: This checkpoint's GDN mixer: `[q|k|v]` is 10240 wide, the conv is width 4,
-#: so each pool slot holds a 3-column window.
-CKPT_DIM = 10240
-#: The largest conv-pool slot count the cache manager was seen to allocate for
-#: this checkpoint (10, 35, 366, 368, 396, 399, 432 and 461 across twelve
-#: runtime-value dumps, while `max_num_sequences` never exceeds 32). The target
-#: hands the pool over unsliced, so 461 tops the certified interval.
-ENGINE_SLOTS = 461
-
-
-class CausalConv1dFwd(OpWrapper):
- """Depthwise causal conv1d (+ optional SiLU) of every sequence in a batch,
- against its carried window, publishing the new window -- one launch.
-
- For sequence `s` over tokens `[b, e)` and `slot = cache_indices[s]` (or `s`):
- `xpad = cat([conv_states[slot] if has_initial_state[s] else 0, x[:, b:e]])`,
- `out[c, b + t] = act(sum_k weight[c, k] * xpad[c, t + k] + bias[c])` with the
- taps accumulated in fp32 and rounded once, and
- `conv_states[slot] = xpad[:, L : L + W - 1]` -- the raw input tail, not the
- activation, exact in every dtype.
-
- Fusion boundary: the state read, the convolution, the bias, the SiLU and the
- state publish are inside. The projection that produces `x`, and slicing the
- result into q/k/v, are the caller's.
-
- Writes the activation into `out`, or into `x` in place when `out is None`,
- and returns that tensor. `conv_states` is advanced in place for every
- non-padded sequence and is not returned.
-
- Two kernels sit behind the op, selected from `x`'s layout: unit stride on
- the token axis is channel-major (in place allowed; `out`, if given, must
- have `x`'s exact strides); unit stride on the channel axis with a token
- stride > 1, `dim > 1` and at least one token is channel-last (`out`
- required, channel-last, disjoint from `x`). The two agree bitwise.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec keys: dim, lengths (varlen), pool_slots, slots (cache_indices),
- # initial (has_initial_state); bf16, W 4, no bias, SiLU, pad_slot_id -1,
- # channel-major in place throughout.
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="checkpoint prefill at the engine's 461-slot pool: slots 460, 0, 230",
- spec=dict(
- dim=CKPT_DIM,
- lengths=(37, 1, 128),
- pool_slots=ENGINE_SLOTS,
- slots=(ENGINE_SLOTS - 1, 0, ENGINE_SLOTS // 2),
- initial=(True, False, True),
- ),
- ),
- Cell(
- why="pad_slot_id -1 (the engine's sentinel): sequence and slot skipped",
- spec=dict(
- dim=1024,
- lengths=(16, 8, 16),
- pool_slots=4,
- slots=(0, -1, 2),
- initial=(True, True, True),
- ),
- ),
- Cell(
- why="one 2048-token sequence: the vectorised load, 128-thread block",
- spec=dict(dim=CKPT_DIM, lengths=(2048,), pool_slots=1, slots=(0,), initial=(True,)),
- ),
- )
-
- note = """
- Metadata consumed: none. Batch layout (`query_start_loc`, int32), pool
- (`conv_states`, `[slots, dim, W - 1]`), slot map (`cache_indices`, int32)
- and carry flags (`has_initial_state`, bool) are explicit arguments.
-
- Call site: the GDN prefill half -- `x` the `[10240, num_prefill_tokens]`
- channel-major output of `extract_transpose_prefill_slice`, W 4, no bias,
- SiLU on, the layer's whole conv pool unsliced, `pad_slot_id` -1, `out` None.
- The cells are that call (bf16, W 4, no bias, SiLU, channel-major in
- place): a three-sequence prefill at the engine's 461-slot pool, which
- tops every width the cache manager was seen to allocate (the 64-thread
- block with scalar loads); one 2048-token sequence (the 128-thread block
- with vectorised loads); and the pad sentinel on a small pool. The
- 128-thread block's scalar-load form (seen at a single 2047-token
- sequence) is not certified.
-
- Preconditions the op checks itself (all raise): dtypes in {bf16, fp16,
- fp32}, `conv_states`/`bias` dtype, int32 `query_start_loc`/`cache_indices`,
- bool `has_initial_state` of length batch, `out` dtype/shape/layout for the
- selected kernel, channel-last in place, `x` rank.
-
- Not checked, not guarded (content, not metadata, or an illegal access
- rather than a wrong number): a `cache_indices` value outside `[0, slots)`;
- `has_initial_state[s]` True with `conv_states=None` (reads through null);
- `conv_states`' shape (only its dtype is validated). Sequences shorter
- than `W - 1` are legal; one (a single token without carried state) is
- certified, one with carried state is not. A zero-length sequence, which
- the target never schedules, is not certified. Unnamed slots are bitwise
- untouched.
-
- Only the channel-major kernel (the target's layout) is certified; the
- channel-last kernel and `TRTLLM_CAUSAL_CONV1D_SILU_MODE`, which selects
- faster SiLU forms on it, are not.
-
- Hard limits
- x token/channel strides uint32 stride fields boundary 2^32 elements
- not probed.
- """
-
- def __call__(
- self,
- x: torch.Tensor,
- weight: torch.Tensor,
- bias: Optional[torch.Tensor],
- conv_states: Optional[torch.Tensor],
- query_start_loc: Optional[torch.Tensor],
- cache_indices: Optional[torch.Tensor],
- has_initial_state: Optional[torch.Tensor],
- silu_activation: bool,
- pad_slot_id: int,
- out: Optional[torch.Tensor] = None,
- ) -> torch.Tensor:
- torch.ops.trtllm.causal_conv1d_fwd(
- x,
- weight,
- bias,
- conv_states,
- query_start_loc,
- cache_indices,
- has_initial_state,
- silu_activation,
- pad_slot_id,
- out,
- )
- return x if out is None else out
-
- def reference(
- self,
- x: torch.Tensor,
- weight: torch.Tensor,
- bias: Optional[torch.Tensor],
- conv_states: Optional[torch.Tensor],
- query_start_loc: Optional[torch.Tensor],
- cache_indices: Optional[torch.Tensor],
- has_initial_state: Optional[torch.Tensor],
- silu_activation: bool,
- pad_slot_id: int,
- out: Optional[torch.Tensor] = None,
- ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
- """`(activation, conv_states)` as the call leaves them; mutates nothing.
-
- The activation starts from `out` (or `x`) so a padded sequence's
- columns keep whatever was there, as the op leaves them.
- """
- width = weight.shape[-1]
- act = (x if out is None else out).clone()
- states = None if conv_states is None else conv_states.clone()
- if query_start_loc is not None:
- starts = query_start_loc.tolist()
- spans = [(slice(None), slice(starts[s], starts[s + 1])) for s in range(len(starts) - 1)]
- else:
- spans = [(s,) for s in range(x.shape[0])]
- slots = None if cache_indices is None else cache_indices.tolist()
- carry = None if has_initial_state is None else has_initial_state.tolist()
- for s, where in enumerate(spans):
- slot = s if slots is None else slots[s]
- if slots is not None and slot == pad_slot_id:
- continue
- tokens = x[where].float() # [dim, L]
- length = tokens.shape[1]
- if conv_states is not None and carry is not None and carry[s]:
- prefix = conv_states[slot].float()
- else:
- prefix = tokens.new_zeros(tokens.shape[0], width - 1)
- padded = torch.cat([prefix, tokens], dim=1)
- acc = torch.zeros_like(tokens)
- for k in range(width):
- acc = acc + weight[:, k].float().unsqueeze(1) * padded[:, k : k + length]
- if bias is not None:
- acc = acc + bias.float().unsqueeze(1)
- if silu_activation:
- acc = acc * torch.sigmoid(acc)
- if length:
- act[where] = acc.to(act.dtype)
- if states is not None:
- states[slot] = padded[:, length : length + width - 1].to(states.dtype)
- return act, states
-
- def is_valid(
- self,
- x: torch.Tensor,
- weight: torch.Tensor,
- bias: Optional[torch.Tensor],
- conv_states: Optional[torch.Tensor],
- query_start_loc: Optional[torch.Tensor],
- cache_indices: Optional[torch.Tensor],
- has_initial_state: Optional[torch.Tensor],
- silu_activation: bool,
- pad_slot_id: int,
- out: Optional[torch.Tensor] = None,
- ) -> None:
- # Widths outside 2..4 fall off the dispatcher with no launch at all:
- # measured at W 1 and W 5, destination and pool bitwise unchanged.
- assert weight.shape[-1] in (2, 3, 4), (
- f"width {weight.shape[-1]} has no kernel; only 2, 3 and 4 are dispatched, "
- "and any other width is a silent no-op"
- )
- # The dispatch reads weight (and bias) as x's type without converting:
- # an fp32 weight with bf16 x moved the result by 11.5.
- assert weight.dtype == x.dtype, (
- "weight must have x's dtype; the kernel reads it as x's type without converting"
- )
- varlen = query_start_loc is not None
- c_axis, l_axis = (0, 1) if varlen else (1, 2)
- if x.dim() != (2 if varlen else 3):
- return # the op rejects a wrong rank itself
- channel_last = (
- x.stride(c_axis) == 1
- and x.stride(l_axis) > 1
- and x.shape[c_axis] > 1
- and x.shape[l_axis] > 0
- )
- if not channel_last:
- # The channel-major kernel walks the token axis and the state's last
- # axis contiguously, ignoring both strides: token stride 2 landed
- # 16.3 off; a state last-dim stride 2 landed 7.9 off (5.4 in the
- # published state). The channel-last kernel honours every stride.
- assert x.shape[l_axis] == 0 or x.stride(l_axis) == 1, (
- "the channel-major kernel loads the token axis contiguously; "
- "x must have stride 1 along it"
- )
- assert conv_states is None or conv_states.stride(-1) == 1, (
- "the channel-major kernel ignores conv_states' last-dim stride; it must be 1"
- )
-
-
-causal_conv1d_fwd = CausalConv1dFwd()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/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
deleted file mode 100644
index cbcd77719624..000000000000
--- 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
+++ /dev/null
@@ -1,283 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Multi-token causal depthwise conv1d step (Triton), with optional per-step
-window snapshots for speculative verification.
-
-Wraps ``tensorrt_llm._torch.modules.mamba.causal_conv1d_triton.causal_conv1d_update``
--- one Triton launch (``_causal_conv1d_update_kernel``) after allocating the
-output. There is no torch-op boundary around this kernel. It is *not*
-``torch.ops.trtllm.causal_conv1d_update``: that op writes `x` in place, has a
-ring mode, takes `state_len` from the pool, has no snapshots, and multiplies in
-fp32 -- the two are not substitutable.
-"""
-
-from __future__ import annotations
-
-from typing import Optional, Union
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.modules.mamba.causal_conv1d_triton import (
- causal_conv1d_update as _causal_conv1d_triton_update,
-)
-
-#: This checkpoint's GDN `[q|k|v]` width; the conv is width 4 over a 3-column pool.
-CKPT_DIM = 10240
-#: 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
-
-
-def _pinned(batch: int) -> tuple[int, ...]:
- """`batch` distinct slots of the engine pool, slot 0 and the last pinned in."""
- return (0,) + tuple(i * 14 + 3 for i in range(1, batch - 1)) + (ENGINE_SLOTS - 1,)
-
-
-class CausalConv1dTritonUpdate(OpWrapper):
- """Convolve each row's new tokens against its cached window, out of place,
- and advance the window; optionally snapshot the window after every token.
-
- Row `b`, slot `s = conv_state_indices[b]` (or `b`), `off = 0` (or
- `num_accepted_tokens[b] - 1`): `window = [conv_state[s, :, off : off+W-1],
- x[b]]`, `out[b, c, i] = act(bias[c] + sum_j fp32(weight[c, j] * window[i+j]))`
- -- each product rounded to the operand dtype, accumulated in fp32, SiLU as
- `acc / (1 + exp(-acc))`, rounded once. The state is then rewritten from
- column 0: the last `W - 1` window columns (or, with `num_accepted_tokens`,
- `W - 1 + seqlen - 1` columns rebased past the offset). Columns beyond that
- are untouched -- `state_len` comes from `W`, never from the pool.
-
- With `intermediate_conv_window`, snapshot `i` of row `b` --
- `[intermediate_state_indices[b], i, :, :W-1]` -- is the window after token
- `i`; after acceptance of `k` tokens the cache is snapshot `k - 1`.
-
- Fusion boundary: window read, conv, bias, SiLU, state write and snapshots
- are inside. The projection, the q/k/v split and the post-acceptance
- promotion of a snapshot into the pool are the caller's.
-
- Returns a fresh tensor with `x`'s shape (2-D in, 2-D out); `x` is read-only.
- A row whose slot equals `pad_slot_id` is not written at all -- its result
- rows are allocator garbage -- and its pool slot and snapshot rows stay
- bitwise unchanged. `conv_state` and `intermediate_conv_window` are written
- in place and not returned.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec keys: dim, 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 = """
- 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, 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 the 461-slot pool
- with a 32-row snapshot buffer and `x` as the target's q|k|v column view
- of the projection: 32 rows at seqlen 3 and at seqlen 8 (the widest verify
- step), and 8 rows at seqlen 5 into an 8-step buffer (a buffer deeper than
- seqlen keeps its later steps untouched); bf16, W 4, no bias, SiLU only.
- Not certified: a contiguous `x`, other seqlens, and pad-sentinel rows
- (the target never sends one).
-
- Raise loudly, so not guarded: width 1 (triton CompilationError); a missing
- `intermediate_state_indices` with snapshots and slot indices given.
-
- Not guarded (content): duplicate slots or duplicate snapshot rows race;
- `num_accepted_tokens[b]` must lie in `[1, seqlen]` (the accepted-token
- roll is uncertified; the target passes None). Out-of-range slot values:
- illegal access, not probed.
-
- Accepted, uncertified: `launch_dependent_kernels` (the target passes
- False). Accepted but not exposed: `cache_seqlens` (ignored by the kernel --
- there is no ring mode), the eagle-tree mask trio, `validate_data`,
- `metadata`, the `_block_n`/`_num_warps`/`_num_stages` overrides. A
- `conv_state` dtype different from `x` is accepted and uncertified.
-
- Hard limits
- none probed.
- """
-
- def __call__(
- self,
- x: torch.Tensor,
- conv_state: torch.Tensor,
- weight: torch.Tensor,
- bias: Optional[torch.Tensor],
- activation: Union[bool, str, None],
- conv_state_indices: Optional[torch.Tensor],
- pad_slot_id: int,
- num_accepted_tokens: Optional[torch.Tensor] = None,
- intermediate_conv_window: Optional[torch.Tensor] = None,
- intermediate_state_indices: Optional[torch.Tensor] = None,
- launch_dependent_kernels: bool = False,
- ) -> torch.Tensor:
- return _causal_conv1d_triton_update(
- x,
- conv_state,
- weight,
- bias=bias,
- activation=activation,
- conv_state_indices=conv_state_indices,
- num_accepted_tokens=num_accepted_tokens,
- intermediate_conv_window=intermediate_conv_window,
- intermediate_state_indices=intermediate_state_indices,
- pad_slot_id=pad_slot_id,
- launch_dependent_kernels=launch_dependent_kernels,
- )
-
- def reference(
- self,
- x: torch.Tensor,
- conv_state: torch.Tensor,
- weight: torch.Tensor,
- bias: Optional[torch.Tensor],
- activation: Union[bool, str, None],
- conv_state_indices: Optional[torch.Tensor],
- pad_slot_id: int,
- num_accepted_tokens: Optional[torch.Tensor] = None,
- intermediate_conv_window: Optional[torch.Tensor] = None,
- intermediate_state_indices: Optional[torch.Tensor] = None,
- launch_dependent_kernels: bool = False,
- ) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
- """`(out, conv_state, intermediate_conv_window)` as the call leaves
- them; mutates nothing. `launch_dependent_kernels` is ignored.
-
- Rows the op leaves unwritten (padded) are NaN in `out`; the test
- skips them. Products are formed in the operand dtype, as the kernel
- forms them -- multiplying in fp32 instead misses the default band on
- about a third of the elements.
- """
- x3 = x.unsqueeze(-1) if x.dim() == 2 else x
- batch, dim, seqlen = x3.shape
- width = weight.shape[1]
- silu = activation in (True, "silu", "swish")
- state_len = width - 1 + (seqlen - 1 if num_accepted_tokens is not None else 0)
- out = torch.full(x3.shape, float("nan"), dtype=x.dtype, device=x.device)
- new_state = conv_state.clone()
- inter = None if intermediate_conv_window is None else intermediate_conv_window.clone()
- slots = None if conv_state_indices is None else conv_state_indices.tolist()
- accepted = None if num_accepted_tokens is None else num_accepted_tokens.tolist()
- rows = None if intermediate_state_indices is None else intermediate_state_indices.tolist()
- for b in range(batch):
- slot = b if slots is None else slots[b]
- if slots is not None and slot == pad_slot_id:
- continue
- off = 0 if accepted is None else accepted[b] - 1
- state = conv_state[slot]
- full = torch.cat([state[:, off : off + width - 1], x3[b]], dim=1)
- acc = torch.zeros(dim, seqlen, dtype=torch.float32, device=x.device)
- for j in range(width):
- acc = acc + (weight[:, j].unsqueeze(1) * full[:, j : j + seqlen]).float()
- if bias is not None:
- acc = acc + bias.float().unsqueeze(1)
- if silu:
- acc = acc / (1.0 + torch.exp(-acc))
- out[b] = acc.to(x.dtype)
- if inter is not None:
- row = b if rows is None else rows[b]
- for i in range(seqlen):
- inter[row, i, :dim, : width - 1] = full[:, i + 1 : i + width]
- if accepted is None:
- new_state[slot, :, : width - 1] = full[:, seqlen : seqlen + width - 1]
- else:
- new_state[slot, :, :state_len] = torch.cat(
- [state[:, off + 1 : off + width - 1], x3[b]], dim=1
- )
- return (out.squeeze(-1) if x.dim() == 2 else out), new_state, inter
-
- def is_valid(
- self,
- x: torch.Tensor,
- conv_state: torch.Tensor,
- weight: torch.Tensor,
- bias: Optional[torch.Tensor],
- activation: Union[bool, str, None],
- conv_state_indices: Optional[torch.Tensor],
- pad_slot_id: int,
- num_accepted_tokens: Optional[torch.Tensor] = None,
- intermediate_conv_window: Optional[torch.Tensor] = None,
- intermediate_state_indices: Optional[torch.Tensor] = None,
- launch_dependent_kernels: bool = False,
- ) -> None:
- dim = x.shape[1]
- seqlen = 1 if x.dim() == 2 else x.shape[2]
- width = weight.shape[1]
- state_len = width - 1 + (seqlen - 1 if num_accepted_tokens is not None else 0)
- # The tap loop is specialised for W 2..4: W 5 runs and returned an
- # answer up to 47.2 off (mean 1.89). W 1 fails to compile.
- assert 2 <= width <= 4, (
- f"weight width must be 2, 3 or 4; got {width} -- width 5 is accepted "
- "by the kernel and silently computes a wrong answer"
- )
- # Reads and writes `state_len` columns unmasked against the pool: a
- # 2-column pool at W 4 wrote one element into the next slot.
- assert conv_state.shape[2] >= state_len, (
- f"conv_state must hold at least {state_len} columns (width={width}, "
- f"seqlen={seqlen}, num_accepted_tokens="
- f"{'given' if num_accepted_tokens is not None else 'None'}); got "
- f"{conv_state.shape[2]} -- a shorter pool is accepted and writes past "
- "the slot into the next one"
- )
- # The channel axis is masked against x's dim only: a half-width pool
- # changed 384 elements of unnamed slots; a half-height weight was read
- # past its end.
- assert conv_state.shape[1] >= dim and weight.shape[0] >= dim, (
- f"conv_state and weight must cover all {dim} channels of x; got "
- f"{conv_state.shape[1]} and {weight.shape[0]} -- a narrower pool or "
- "weight is accepted and reads/writes out of bounds"
- )
- if intermediate_conv_window is not None:
- # No bound on the step axis: a 2-step buffer at seqlen 4 wrote
- # into the next row.
- assert (
- intermediate_conv_window.shape[1] >= seqlen
- and intermediate_conv_window.shape[2] >= dim
- and intermediate_conv_window.shape[3] >= width - 1
- ), (
- f"intermediate_conv_window must be at least [rows, {seqlen}, {dim}, "
- f"{width - 1}]; got {tuple(intermediate_conv_window.shape)} -- an "
- "undersized buffer is accepted and writes past the row it was given"
- )
-
-
-causal_conv1d_triton_update = CausalConv1dTritonUpdate()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py
deleted file mode 100644
index 14413d38be2b..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py
+++ /dev/null
@@ -1,481 +0,0 @@
-# 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.
-
-The decode step's shape (one token a row, SiLU, no bias, a width-4 conv over a
-3-column bf16 pool, slot indices) from `_CUTE_MIN_ROWS` rows runs a CuTe DSL
-kernel (Tier 6, written for this target) that is bitwise the TRT-LLM kernel it
-replaces there; every other call is `torch.ops.trtllm.causal_conv1d_update`.
-"""
-
-from __future__ import annotations
-
-from typing import Optional
-
-import cutlass
-import cutlass.cute as cute
-import torch
-from cutlass._mlir.dialects import llvm
-from cutlass.cute.runtime import from_dlpack
-from cutlass.cutlass_dsl import T, dsl_user_op
-
-import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.*
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
-
-#: This checkpoint's GDN `[q|k|v]` width; the conv is width 4 over a 3-column pool.
-CKPT_DIM = 10240
-#: Largest conv-pool slot count the cache manager was seen to allocate for this
-#: checkpoint (10 .. 461 over twelve runtime-value dumps, batch never above 32).
-#: The decode step hands the whole pool over unsliced.
-ENGINE_SLOTS = 461
-
-
-# --------------------------------------------------------------------------------------
-# 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."""
- return (0,) + tuple(i * 14 + 3 for i in range(1, batch - 1)) + (ENGINE_SLOTS - 1,)
-
-
-class CausalConv1dUpdate(OpWrapper):
- """One causal depthwise conv step per batch row against that row's cached
- window, then the window's advance -- `x` and the pool both written in place.
-
- Row `b`, slot `s = conv_state_indices[b]` (or `b`), new tokens `i < seqlen`:
- `history = [conv_state[s, c, :], x[b, c, :]]`, the window is its last
- `W - 1 + seqlen` columns, `x[b, c, i] = act(bias[c] + sum_j weight[c, j] *
- window[i + j])` in fp32 rounded once, and `conv_state[s] = history[-state_len:]`
- with `state_len = conv_state.shape[2]`. With `cache_seqlens` the pool is a
- ring instead: position `p` lives at `p % state_len`, `cache_seqlens[b]` is
- the first new token's absolute position, and nothing shifts.
-
- Fusion boundary: window read, convolution, bias, SiLU and state write are
- inside. The projection producing `x`, slot allocation and the q/k/v split
- are the caller's.
-
- Returns `x` -- the same object, written in place. A 2-D `[batch, dim]` `x`
- is handed to the op as an `unsqueeze(-1)` view (metadata only; the op
- writes through it). `conv_state` is advanced in place and not returned.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec keys: dim, 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="x the q|k|v column view of the [rows, 16480] in_proj, batch 32",
- spec=dict(dim=CKPT_DIM, batch=32, pool_slots=ENGINE_SLOTS, slots=_pinned(32)),
- ),
- Cell(
- why="x the q|k|v column view of the [rows, 16480] in_proj, batch 128",
- spec=dict(dim=CKPT_DIM, batch=128, pool_slots=256, slots="random"),
- ),
- Cell(
- why="64 rows (the CuTe kernel's first) over a layer view, pad -1 rows untouched",
- spec=dict(
- dim=CKPT_DIM,
- batch=64,
- pool_slots=ENGINE_SLOTS,
- slots=tuple(-1 if i % 21 == 5 else i * 7 for i in range(64)),
- pool_layers=2,
- ),
- ),
- )
-
- note = """
- Metadata consumed: none.
-
- Call site: the GDN decode step -- 2-D `[num_decodes, 10240]` x, the q|k|v
- column view (row stride 16480) of the decode rows of the layer's one
- `[T, 16480]` input-projection output (the op writes those columns only);
- the layer's whole conv pool `[slots, 10240, 3]`, W 4, no bias, SiLU on,
- `cache_seqlens` None, int32 slots, `pad_slot_id` -1.
- That exact shape (silu, no ring, `state_len == W - 1`) runs TRT-LLM's
- dedicated seqlen-1 kernel; any other shape runs its general one, which the
- target never reaches and which is not certified. From `_CUTE_MIN_ROWS`
- (64) rows, with bf16 operands, a contiguous weight, a pool whose slot is
- `[dim, 3]` contiguous and 16-byte aligned rows, it runs
- `_DecodeConvKernel` (CuTe DSL) instead: 8 channels a thread by 16-byte
- loads, the TRT-LLM kernel's instructions, bitwise its result (x and pool,
- pad rows untouched); 128 rows 7.46 -> 5.41 us a layer (GB300, graph
- replay over 8 rotating layers). Certified at that call: 32 rows (the
- TRT-LLM kernel) over the 461-slot pool, 128 rows (the CuTe kernel) over a
- 256-slot pool, and 64 rows (the CuTe kernel's first row count) with pad
- rows over the 461-slot pool seen as one layer of a two-layer buffer.
-
- The op checks metadata itself, loudly: width 2..4, `state_len >= W - 1`,
- one dtype across x/conv_state/weight (bias matches weight), weight shape,
- int32 contiguous length-batch `conv_state_indices`, int32
- `cache_seqlens`, a batch-sized pool when indices are absent. So no guard.
-
- Not guarded because it is content, not metadata: two rows naming one slot
- race -- the slot ended up matching neither row's expected state in 10 of
- 10 runs while both outputs matched. Slot uniqueness is the caller's.
- Out-of-range slot values are an illegal access, not probed.
-
- A row whose index equals `pad_slot_id` (-1 at the call site) is skipped:
- its x rows and the pool stay bitwise unchanged (certified on the CuTe
- kernel only).
-
- Hard limits
- none probed.
- """
-
- def __call__(
- self,
- x: torch.Tensor,
- conv_state: torch.Tensor,
- weight: torch.Tensor,
- bias: Optional[torch.Tensor],
- silu_activation: bool,
- cache_seqlens: Optional[torch.Tensor],
- conv_state_indices: Optional[torch.Tensor],
- pad_slot_id: int,
- ) -> torch.Tensor:
- 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,
- weight,
- bias,
- silu_activation,
- cache_seqlens,
- conv_state_indices,
- pad_slot_id,
- )
- return x
-
- def reference(
- self,
- x: torch.Tensor,
- conv_state: torch.Tensor,
- weight: torch.Tensor,
- bias: Optional[torch.Tensor],
- silu_activation: bool,
- cache_seqlens: Optional[torch.Tensor],
- conv_state_indices: Optional[torch.Tensor],
- pad_slot_id: int,
- ) -> tuple[torch.Tensor, torch.Tensor]:
- """`(x, conv_state)` as the call leaves them; mutates nothing."""
- x3 = x.unsqueeze(-1) if x.dim() == 2 else x
- batch, dim, seqlen = x3.shape
- width = weight.shape[1]
- state_len = conv_state.shape[2]
- out = x3.clone()
- new_state = conv_state.clone()
- slots = None if conv_state_indices is None else conv_state_indices.tolist()
- starts = None if cache_seqlens is None else cache_seqlens.tolist()
- for b in range(batch):
- slot = b if slots is None else slots[b]
- if slots is not None and slot == pad_slot_id:
- continue
- state = conv_state[slot].float()
- tokens = x3[b].float()
- if starts is None:
- history = torch.cat([state, tokens], dim=1)
- window = history[:, state_len - (width - 1) :]
- new_state[slot] = history[:, -state_len:].to(conv_state.dtype)
- else:
- start = starts[b] % state_len
- read = [(start - (width - 1) + i) % state_len for i in range(width - 1)]
- window = torch.cat([state[:, read], tokens], dim=1)
- advanced = state.clone()
- for i in range(seqlen):
- advanced[:, (start + i) % state_len] = tokens[:, i]
- new_state[slot] = advanced.to(conv_state.dtype)
- acc = torch.zeros(dim, seqlen, dtype=torch.float32, device=x.device)
- for j in range(width):
- acc = acc + weight[:, j].float().unsqueeze(1) * window[:, j : j + seqlen]
- if bias is not None:
- acc = acc + bias.float().unsqueeze(1)
- if silu_activation:
- acc = acc * torch.sigmoid(acc)
- out[b] = acc.to(x.dtype)
- return (out.squeeze(-1) if x.dim() == 2 else out), new_state
-
-
-causal_conv1d_update = CausalConv1dUpdate()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/chunk_gated_delta_rule.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/chunk_gated_delta_rule.py
deleted file mode 100644
index 701fa42c82ed..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/chunk_gated_delta_rule.py
+++ /dev/null
@@ -1,268 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Chunked gated delta rule: the Gated DeltaNet prefill core.
-
-Wraps the in-tree module callable
-``tensorrt_llm._torch.modules.fla.chunk.chunk_gated_delta_rule`` -- the
-**vendored Triton** implementation (tier-2 sourcing: the gated delta rule has
-no ``torch.ops.trtllm`` surface). Its interior is a fixed sequence of Triton
-steps (gate cumsum, scaled K K^T, triangular solve, WY recompute, chunked state
-recurrence, output epilogue) no caller invokes on its own: one fused op.
-
-The GDN mixer resolves between this and a FlashInfer-backed sibling of the same
-name (``tensorrt_llm._torch.modules.fla.flashinfer_chunk``) through
-``gdn_mixer._resolve_chunk_gated_delta_rule``, which picks FlashInfer on sm_103
-unless ``TLLM_USE_FLASHINFER_GDN_PREFILL=0``. This entry names the Triton one
-directly, so neither the resolver nor the variable changes what it runs.
-"""
-
-from __future__ import annotations
-
-from typing import Optional, Tuple
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.modules.fla.chunk import chunk_gated_delta_rule as _chunk_gated_delta_rule
-
-#: This checkpoint's Gated DeltaNet geometry at tp1.
-NUM_K_HEADS, NUM_V_HEADS, HEAD_K_DIM, HEAD_V_DIM = 16, 48, 128, 128
-#: A recurrent-pool slot count of a full-size engine run. The kernels address a slot by
-#: its index alone, so a pool of another size (a bf16 pool holds more slots in the same
-#: memory) takes the same path. The target passes the layer's whole pool unsliced.
-ENGINE_SLOTS = 461
-
-
-def _engine_slots(num_seqs: int) -> tuple[int, ...]:
- """`num_seqs` distinct slots of the engine pool, slot 0 and the last live."""
- return ((0, ENGINE_SLOTS - 1) + tuple(i * 13 + 4 for i in range(num_seqs)))[:num_seqs]
-
-
-class ChunkGatedDeltaRule(OpWrapper):
- """Run the gated delta rule over whole sequences, returning `(o, final_state)`.
-
- Per token `t` and value head, `h` the `[V, K]` state, `q_h`/`k_h` the key
- heads broadcast by `repeat_interleave` (value head `i` reads key head
- `i // (H // Hg)`), fp32 accumulation:
-
- h *= exp(g[t]); u = (v[t] - h k_h[t]) * beta[t]; h += u k_h[t]^T
- o[t] = h (scale * q_h[t])
-
- computed 64 tokens at a time over a WY representation rounded to
- `q.dtype` -- the same recurrence, a different sum. `scale=None` is
- `K ** -0.5`; `g` is in log space. `use_qk_l2norm_in_kernel` normalizes q/k
- as `x / sqrt(sum(x^2) + 1e-6)` first.
-
- With `cu_seqlens` the batch is one packed row (`q.shape[0] == 1`) split at
- those offsets; without it each row is a sequence. Initial state: zeros
- (`initial_state=None`); `initial_state[n]`, read-only; or, with
- `initial_state_indices` and `inplace_indexed_state_update=True`,
- `initial_state[indices[n]]` with the final state **written back into that
- slot**. `output_final_state` only decides whether a fresh fp32
- `[N, H, V, K]` final state is also returned (on the indexed path, a copy
- gathered after the write-back).
-
- Fusion boundary: gate cumsum, optional q/k norm, scale, recurrence, state
- read and write, output epilogue. The conv, q/k/v/g/beta preparation, slot
- allocation and the gated norm on `o` are the caller's.
-
- Returns `(o, final_state)` unchanged: `o` `[B, T, H, V]` in `q.dtype`
- (`output` itself when given), `final_state` or `None`. The pool write is in
- place and not returned.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: seqlens (packed with cu_seqlens); the initial state is the
- # `ENGINE_SLOTS`-row pool at `indices`, written in place; pool_dtype (fp32 |
- # bf16); output: None | "buffer" (a fresh `output=` tensor).
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="fp32 pool (a verify step's context rows), 461 slots, output= buffer",
- spec=dict(seqlens=(200, 37, 91), indices=_engine_slots(3), output="buffer"),
- ),
- Cell(
- why="bf16 pool, 461-slot pool, seqlens [1024]",
- spec=dict(seqlens=(1024,), indices=_engine_slots(1), pool_dtype=torch.bfloat16),
- ),
- Cell(
- why="bf16 pool, mixed step: a context and four 1-token rows [200, 1, 1, 1, 1]",
- spec=dict(
- seqlens=(200, 1, 1, 1, 1), indices=_engine_slots(5), pool_dtype=torch.bfloat16
- ),
- ),
- )
-
- note = """
- Metadata consumed: none.
-
- Call site (`_gdn_block`): `(q, k, v, g, beta, None, ssm_states, si_all,
- True, False, qsl_long, False, None)` after `fused_gdn_post_conv` (q/k
- already normalized) -- the layer's whole pool `[slots, 48, 128, 128]`, fp32
- or bf16 as `mamba_ssm_cache_dtype` allocated it, int32 slots, int64
- `cu_seqlens`; the MTP-verify context half passes the `si_p` / `qsl_long_p`
- prefixes, and a mixed batch's one-token decode rows ride in the same call.
- Certified at 461 slots: three ragged context rows over an fp32 pool into
- an `output=` buffer, one 1024-token sequence over a bf16 pool, and a
- mixed step (a context and four one-token rows) over a bf16 pool.
-
- Over a bf16 pool the initial state is read as fp32 and the final state is stored
- rounded once to bf16, once per call and sequence; the accumulation is unchanged.
-
- The first indexed in-place call for an autotune key in a process is as correct as a
- warm one: the state kernel's autotuner snapshots and restores the indexed state slots
- its benchmark launches write (`chunk_delta_h.py`, `_save_autotune_state`). The key
- includes the pool dtype, so a bf16 pool has its own. Certified in a fresh interpreter
- at this geometry over a bf16 pool (the fp32 pool's key takes the same snapshot and
- restore). The positional path never aliased its state.
-
- Loud, from the callable: fp32 or mixed-dtype q/k/v, a 2-D `beta`,
- `q.shape[0] != 1` with `cu_seqlens`, a sequence-count mismatch of indices
- or initial states, indices without `inplace_indexed_state_update`.
-
- Not checked: slots must not alias (not probed); q/k must be unit-norm,
- by the caller or the in-kernel flag (raw q/k with the flag off return
- finite values ~1e25). `head_first` is pinned False.
-
- Sequence lengths are unconstrained: certified at 1, 37, 91, 200 and 1024,
- off and on the 64 boundary, mixed ragged. Memory: a `[B, NT, H, K, V]`
- per-chunk state scratch in `q.dtype`, 24 MiB per 1024 tokens here.
- Deterministic.
-
- Not certified: the FlashInfer sibling, fp16 (accepted), bf16 `g`/`beta`, an
- fp16 pool, `K = 256` (accepted, unmeasured), empty `T` or zero-length spans,
- and what the target never passes: `initial_state=None` or a positional
- initial state, `output_final_state`, `use_qk_l2norm_in_kernel`, an explicit
- `scale`, no `cu_seqlens`.
-
- Hard limits
- K state kernel asserts K <= 256 (read, not probed)
- """
-
- def __call__(
- self,
- q: torch.Tensor,
- k: torch.Tensor,
- v: torch.Tensor,
- g: torch.Tensor,
- beta: torch.Tensor,
- scale: Optional[float] = None,
- initial_state: Optional[torch.Tensor] = None,
- initial_state_indices: Optional[torch.Tensor] = None,
- inplace_indexed_state_update: bool = False,
- output_final_state: bool = False,
- cu_seqlens: Optional[torch.Tensor] = None,
- use_qk_l2norm_in_kernel: bool = False,
- output: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
- return _chunk_gated_delta_rule(
- q=q,
- k=k,
- v=v,
- g=g,
- beta=beta,
- scale=scale,
- initial_state=initial_state,
- initial_state_indices=initial_state_indices,
- inplace_indexed_state_update=inplace_indexed_state_update,
- output_final_state=output_final_state,
- cu_seqlens=cu_seqlens,
- head_first=False,
- use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
- output=output,
- )
-
- def reference(
- self,
- q: torch.Tensor,
- k: torch.Tensor,
- v: torch.Tensor,
- g: torch.Tensor,
- beta: torch.Tensor,
- scale: Optional[float] = None,
- initial_state: Optional[torch.Tensor] = None,
- initial_state_indices: Optional[torch.Tensor] = None,
- inplace_indexed_state_update: bool = False,
- output_final_state: bool = False,
- cu_seqlens: Optional[torch.Tensor] = None,
- use_qk_l2norm_in_kernel: bool = False,
- output: Optional[torch.Tensor] = None,
- ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
- """`(o, final_state, initial_state as the call leaves it)`.
-
- fp64 token-at-a-time recurrence, rounded once: `o` to `q.dtype`, the
- states to fp32 / the pool's dtype. Mutates nothing. `output` is ignored
- (it only chooses where `o` lands).
- """
- dt = torch.float64
- batch, num_tokens, num_k_heads, head_k = q.shape
- num_v_heads, head_v = v.shape[2], v.shape[3]
- repeats = num_v_heads // num_k_heads
- scale = head_k**-0.5 if scale is None else scale
- if cu_seqlens is None:
- spans = [(b, 0, num_tokens) for b in range(batch)]
- else:
- bounds = cu_seqlens.tolist()
- spans = [(0, bounds[n], bounds[n + 1]) for n in range(len(bounds) - 1)]
- slots = None if initial_state_indices is None else initial_state_indices.tolist()
- out = torch.zeros(batch, num_tokens, num_v_heads, head_v, dtype=dt, device=q.device)
- finals = []
- for n, (row, bos, eos) in enumerate(spans):
- if initial_state is None:
- state = torch.zeros(num_v_heads, head_v, head_k, dtype=dt, device=q.device)
- else:
- state = initial_state[n if slots is None else slots[n]].to(dt)
- for t in range(bos, eos):
- q_t, k_t = q[row, t].to(dt), k[row, t].to(dt)
- if use_qk_l2norm_in_kernel:
- q_t = q_t / torch.sqrt(q_t.pow(2).sum(-1, keepdim=True) + 1e-6)
- k_t = k_t / torch.sqrt(k_t.pow(2).sum(-1, keepdim=True) + 1e-6)
- q_t = q_t.repeat_interleave(repeats, dim=0) * scale
- k_t = k_t.repeat_interleave(repeats, dim=0)
- state = state * g[row, t].to(dt).exp()[:, None, None]
- u = (v[row, t].to(dt) - (state * k_t[:, None, :]).sum(-1)) * beta[row, t].to(dt)[
- :, None
- ]
- state = state + u[:, :, None] * k_t[:, None, :]
- out[row, t] = (state * q_t[:, None, :]).sum(-1)
- finals.append(state)
- final_state = torch.stack(finals).float() if output_final_state else None
- pool = None if initial_state is None else initial_state.clone()
- if slots is not None and inplace_indexed_state_update:
- for slot, state in zip(slots, finals):
- pool[slot] = state.to(pool.dtype)
- return out.to(q.dtype), final_state, pool
-
- def is_valid(
- self,
- q: torch.Tensor,
- k: torch.Tensor,
- v: torch.Tensor,
- g: torch.Tensor,
- beta: torch.Tensor,
- scale: Optional[float] = None,
- initial_state: Optional[torch.Tensor] = None,
- initial_state_indices: Optional[torch.Tensor] = None,
- inplace_indexed_state_update: bool = False,
- output_final_state: bool = False,
- cu_seqlens: Optional[torch.Tensor] = None,
- use_qk_l2norm_in_kernel: bool = False,
- output: Optional[torch.Tensor] = None,
- ) -> None:
- if output is None:
- return
- # Measured: a head-slice view (strides (1179648, 12288, 128, 1)) was
- # accepted and came back 9.9e-2 off where correct is ~4e-4; an fp32
- # buffer was filled correctly but the returned bf16 tensor no longer
- # aliased it. A token-window slice that reports contiguous is fine.
- assert output.is_contiguous(), (
- f"output must be contiguous; got strides {tuple(output.stride())} -- the "
- "epilogue addresses it through strides derived from its shape"
- )
- assert output.dtype == q.dtype, (
- f"output must have q's dtype {q.dtype}; got {output.dtype} -- the result "
- "is cast to q.dtype and stops aliasing the buffer"
- )
-
-
-chunk_gated_delta_rule = ChunkGatedDeltaRule()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/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
deleted file mode 100644
index 904246678bc0..000000000000
--- 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
+++ /dev/null
@@ -1,131 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Prepend-zero cumulative sum of a per-request query-row count: the
-``cu_seqlens`` segmentation vector every varlen recurrent op takes.
-
-Tier-3 composed from native torch (``torch.zeros`` + ``torch.cumsum``, three
-launches). No registered op or ``modules/`` callable produces this vector: the
-recurrent metadata's ``query_start_loc*`` is built inline by
-``Mamba2Metadata.prepare`` and segments the generation tail one row per
-sequence, which is short under speculative decoding.
-"""
-
-from __future__ import annotations
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: Integer dtypes accepted on either side. ``torch.cumsum`` silently accepts
-#: bool and floating dtypes and answers wrongly with them.
-_INTEGER_DTYPES = (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64)
-
-#: Batch shapes this checkpoint's serving path produces 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 = {
- "mixed_p2": (1024, 37, 2, 2, 2, 2),
- "padded_decode_p2": (2, 2, 2, 0, 0, 0),
-}
-
-
-class CuSeqlensFromSeqLens(OpWrapper):
- """`out[0] = 0`, `out[s] = seq_lens[0] + ... + seq_lens[s - 1]`: a fresh
- contiguous `[num_seqs + 1]` tensor of `dtype` on `seq_lens.device`.
-
- The scan accumulates in `dtype`, not in `seq_lens`'s dtype. Fusion boundary:
- the leading zero, the scan and the dtype. Choosing which length vector to
- scan, and keeping the packed activation exactly `out[-1]` rows long, are the
- caller's. `seq_lens` is read-only.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: lengths, in_dtype, out_dtype; `layout` "contiguous" (default) or
- # "prefix" (the first len(lengths) of a longer buffer -- the call site's
- # `seq_lens_cuda[:batch]`).
- CELLS: tuple[Cell, ...] = (
- 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="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="padded_decode_p2: zero-length padded rows, int32->int64",
- spec=dict(
- lengths=BATCHES["padded_decode_p2"], in_dtype=torch.int32, out_dtype=torch.int64
- ),
- ),
- )
-
- note = """
- Metadata consumed: none. The call site reads `attn_metadata.seq_lens_cuda`
- prefix, which still holds the true per-request query-row count during the
- trunk's forward (the spec worker rewrites it only afterwards). Not the
- recurrent metadata's `query_start_loc_long`: that one has one row per
- generation sequence and is short under MTP; the two agree exactly when
- every generation request contributes one row.
-
- Preconditions, none checked (data, not metadata -- a check would sync):
- entries non-negative (a negative length was silently wrong in three
- consumer probes and a device-side assert in the fourth); a padded row is
- length 0, never -1; `out[-1]` equals the packed activation's row count.
-
- `dtype` has no default: `causal_conv1d_fwd` requires int32 and
- `chunk_gated_delta_rule` keys its autotuner on `cu_seqlens`' dtype.
-
- Certified bitwise at the call site's prefix view of `seq_lens_cuda` (a
- mixed step at p 2, and the p 8 verify of 32 requests) and at a decode
- batch padded with zero-length rows, int32 in and int64 out. CUDA-graph
- capture is certified (only `seq_lens.shape[0]` reaches Python; a captured
- graph follows lengths rewritten in place). Three launches, about 0.03 ms
- at batch 8..256 -- launch bound.
-
- Not certified: a `seq_lens` dtype other than int32 (the guard admits any
- integer dtype), a non-contiguous `seq_lens` view.
-
- Hard limits
- out[-1] accumulated in `dtype` boundary 2^31 - 1 for int32
- probed at [2e9, 2e9]: int64 gives [0, 2e9, 4e9]; int32 wraps silently
- to [0, 2000000000, -294967296]. Not guarded (data).
- """
-
- def __call__(self, seq_lens: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
- out = torch.zeros(seq_lens.shape[0] + 1, dtype=dtype, device=seq_lens.device)
- torch.cumsum(seq_lens, dim=0, dtype=dtype, out=out[1:])
- return out
-
- def reference(self, seq_lens: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
- """A scalar Python prefix sum; no torch op computes it."""
- total, out = 0, [0]
- for length in seq_lens.tolist():
- total += int(length)
- out.append(total)
- return torch.tensor(out, dtype=dtype, device=seq_lens.device)
-
- def is_valid(self, seq_lens: torch.Tensor, dtype: torch.dtype) -> None:
- # Measured: a [1, n] row vector returned a length-2 [0, n_0]; an [n, 2]
- # input a length-(n + 1) vector of nonsense (only a resize warning); a
- # float32 [7.9, 1.2, 13.6] was truncated to [0, 7, 8, 21]; a bool input
- # was counted as nonzero rows; a float `dtype` returned a float vector.
- assert seq_lens.dim() == 1, (
- f"seq_lens must be a 1-D per-request row count, got shape {tuple(seq_lens.shape)}"
- )
- assert seq_lens.dtype in _INTEGER_DTYPES, (
- f"seq_lens must have an integer dtype, got {seq_lens.dtype}"
- )
- assert dtype in _INTEGER_DTYPES, f"dtype must be an integer dtype, got {dtype}"
-
-
-cu_seqlens_from_seq_lens = CuSeqlensFromSeqLens()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/dilated_short_conv_paged.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/dilated_short_conv_paged.py
deleted file mode 100644
index 8ff846801270..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/dilated_short_conv_paged.py
+++ /dev/null
@@ -1,292 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Dilated depthwise short convolution over a packed, ragged serving batch, with
-each sequence's convolution state read from and published to a paged pool.
-
-Composed from native torch ops (sourcing tier 3): no TensorRT-LLM kernel takes
-a dilation. The ladder, as measured on this machine in this paged form:
-
-* tier 1 -- no registered op schema mentions `dilation`; the three convolutions
- (`causal_conv1d_fwd`, `causal_conv1d_update`, `mamba_conv1d`) are the
- non-dilated family. `causal_conv1d_fwd` fits the batch layout exactly, and a
- zero-interleaved width-10 weight (taps at 0, 3, 6, 9) over the very
- `[slots, channels, 9]` pool **returns normally and computes nothing** --
- activation and pool bitwise unchanged -- because its dispatcher has no branch
- past width 4 and the width check sits on the update entry point only.
- `causal_conv1d_update` raises on the same weight.
-* tier 2 -- the vendored Triton `causal_conv1d_fn` does not compile at width
- 10; engram's `ShortConv` is dilated but rectangular and stateless.
-* tier 3 -- `F.conv1d(dilation=, groups=)` over a packed history plus index
- arithmetic. This entry.
-
-Not taken: `d` phase-separated non-dilated convolutions through
-`causal_conv1d_fwd` (three engine calls plus gathers -- not one entry), and a
-rectangular `[num_seqs, channels, state_len + max_tokens]` buffer (5.4 GB at 32
-x 8192 tokens). A fused dilation-aware paged kernel is the faster path.
-"""
-
-from __future__ import annotations
-
-from typing import Optional
-
-import torch
-import torch.nn.functional as F
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: `hc_count * hidden_size = 4 * 2560` channels, kernel 4, dilation `ngram_size`
-#: 3, hence the 9-column state the mamba cache manager allocates per slot.
-CKPT_CHANNELS = 10240
-CKPT_DILATION = 3
-#: A PLE-conv-pool slot count of a full-size engine run. The kernel addresses a slot by its
-#: index alone, so a pool of another size (the pool shares the recurrent pool's slot space,
-#: and a bf16 recurrent pool holds more slots in the same memory) takes the same path. The
-#: target hands the whole pool over unsliced.
-ENGINE_SLOTS = 461
-
-
-def _scattered(num_seqs: int) -> tuple[int, ...]:
- """`num_seqs` distinct slots over the whole engine pool, both ends in."""
- slots = ((0, ENGINE_SLOTS - 1) + tuple(i * 13 + 4 for i in range(num_seqs)))[:num_seqs]
- assert len(set(slots)) == num_seqs and max(slots) < ENGINE_SLOTS
- return slots
-
-
-def _ckpt(why: str, lengths: tuple[int, ...], **spec) -> Cell:
- return Cell(
- why=why,
- spec=dict(channels=CKPT_CHANNELS, lengths=lengths, pool_slots=ENGINE_SLOTS, **spec),
- )
-
-
-class DilatedShortConvPaged(OpWrapper):
- """For sequence `s` over columns `[b, e)` of the packed channel-major `x`,
- `slot = cache_indices[s]` (or `s`), `state_len = (kernel_size - 1) *
- dilation`: `window = [conv_states[slot] if has_initial_state[s] else 0,
- x[:, b:e]]`, `out[c, b + t] = silu(round(sum_k weight[c, 0, k] *
- window[c, t + k * dilation] + bias[c]))` -- taps in fp32, rounded by the
- convolution, then the SiLU reads that back and rounds again -- and
- `conv_states[slot] = window[:, L : L + state_len]`, the raw input tail.
-
- The highest tap index multiplies the token being emitted. A padded
- sequence (`cache_indices[s] == pad_slot_id`) publishes nothing and still
- gets output columns -- its tokens against a zero carry, defined and
- meaningless. Two sequences naming one slot: the last live one's state
- survives (defined by construction, not a race; uncertified, as the
- scheduler never names one slot twice).
-
- Fusion boundary: state read, zero fill for a fresh sequence, the dilated
- depthwise conv, bias, SiLU and the state publish are inside. The
- pre-conv norm producing `x`, slot allocation and reset, and the residual
- add are the caller's.
-
- Returns a fresh contiguous `[channels, num_tokens]` activation; `x`,
- `weight` and `bias` are read-only. `conv_states` is advanced in place and
- not returned.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec keys: lengths, slots, initial (default all True). Every cell is the
- # call site's own form: a bare .t() of the token-major [T, 10240]
- # activation, kernel 4, dilation 3, bf16, no bias, SiLU,
- # query_start_loc_long (int64), int32 slots, the whole 461-slot pool.
- CELLS: tuple[Cell, ...] = (
- _ckpt("decode: 32 one-token rows", (1,) * 32, slots=_scattered(32)),
- _ckpt(
- "mixed: a 2047-token context + 3 one-token rows", (2047, 1, 1, 1), slots=_scattered(4)
- ),
- _ckpt(
- "MTP verify: a fresh 37-token context + 8 rows of 3 drafts",
- (37,) + (3,) * 8,
- slots=_scattered(9),
- initial=(False,) + (True,) * 8,
- ),
- )
-
- note = """
- Metadata consumed: none -- the same explicit argument set as
- `causal_conv1d_fwd` (packed batch, pool, int32/int64 slot map, bool carry).
-
- Call site: the PLE short conv on one layer -- `x` a bare `.t()` of the
- `[T, 10240]` normed activation, weight `[10240, 1, 4]` bf16, no bias,
- dilation 3 (state 9), SiLU on, the whole `[slots, 10240, 9]` pool, int32
- slots, `query_start_loc_long` (int64, rebuilt from `seq_lens` under MTP),
- carry flags covering the whole batch, `pad_slot_id` -1. Certified in that
- form at the 461-slot pool: 32 one-token decode rows, a 2047-token context
- mixed with three one-token rows, and an MTP verify step (a fresh 37-token
- context and 8 requests of 3 tokens). A padded row is not certified.
-
- Preconditions: `query_start_loc[-1] == x.shape[1]`; a short one indexes
- past the internal history and fires a device-side assert that kills the
- CUDA context (loud, not guarded). An out-of-range slot reaches the same
- bounds check (reasoned, not probed). Pool, weight and bias share `x`'s
- dtype (each mismatch raises); `weight` is 3-D depthwise (2-D raises).
- Only the checkpoint's kernel size 4 and dilation 3 are certified; there is
- no width dispatch.
-
- CUDA-graph capture is certified (at 32 requests of 8 tokens): no host
- sync, all batch reads on device, a captured graph stays correct when the
- batch tensors' values change.
- Launch-bound on decode (8 rows: 0.461 ms eager vs 0.029 ms of conv; 0.094
- ms replayed from a graph).
-
- Not measured: pools above 461 slots (a bf16 recurrent state gives the engine
- larger ones; a slot is addressed by its index alone), token counts above 8192,
- `torch.compile`, a pool dtype unlike the model's.
- """
-
- def __call__(
- self,
- x: torch.Tensor,
- weight: torch.Tensor,
- bias: Optional[torch.Tensor],
- conv_states: torch.Tensor,
- query_start_loc: torch.Tensor,
- cache_indices: Optional[torch.Tensor],
- has_initial_state: Optional[torch.Tensor],
- dilation: int,
- silu_activation: bool,
- pad_slot_id: int,
- ) -> torch.Tensor:
- channels, num_tokens = x.shape
- kernel_size = weight.shape[-1]
- state_len = (kernel_size - 1) * dilation
- num_slots = conv_states.shape[0]
- num_seqs = query_start_loc.numel() - 1
- device = x.device
-
- # Device-resident arithmetic only: no `.item()`, no boolean indexing,
- # no data-dependent shape -- which is what makes the call capturable.
- starts = query_start_loc[:-1].to(torch.int64)
- ends = query_start_loc[1:].to(torch.int64)
- rows = torch.arange(num_seqs, device=device)
- columns = torch.arange(num_tokens, device=device)
- taps = torch.arange(state_len, device=device)
-
- # The first sequence whose end is strictly past the column; skips
- # zero-token sequences by construction.
- seq_of_column = torch.bucketize(columns, ends.contiguous(), right=True)
-
- if cache_indices is None:
- slots = rows
- live = torch.ones(num_seqs, dtype=torch.bool, device=device)
- else:
- live = cache_indices != pad_slot_id
- slots = torch.where(live, cache_indices.to(torch.int64), torch.zeros_like(rows))
- carries = (
- torch.zeros_like(live) if has_initial_state is None else (has_initial_state & live)
- )
-
- # Each sequence's carry placed right ahead of its own tokens, so one
- # convolution over the buffer never reads across a sequence boundary.
- history = x.new_zeros((channels, num_tokens + num_seqs * state_len))
- history.index_copy_(1, columns + state_len * (seq_of_column + 1), x)
-
- previous = conv_states.index_select(0, slots)
- history.index_copy_(
- 1,
- ((starts + rows * state_len).unsqueeze(1) + taps).reshape(-1),
- torch.where(carries.reshape(num_seqs, 1, 1), previous, torch.zeros_like(previous))
- .permute(1, 0, 2)
- .reshape(channels, -1),
- )
-
- convolved = F.conv1d(
- history,
- weight,
- bias=bias,
- stride=1,
- padding=0,
- dilation=dilation,
- groups=channels,
- )
- if silu_activation:
- convolved = F.silu(convolved)
- out = convolved.index_select(1, columns + state_len * seq_of_column)
-
- published = (
- history.index_select(1, ((ends + rows * state_len).unsqueeze(1) + taps).reshape(-1))
- .view(channels, num_seqs, state_len)
- .permute(1, 0, 2)
- )
-
- # Skip padded rows without a host sync: every row writes, but writes
- # the bytes of the highest-numbered live row owning that slot (or the
- # slot's own contents when no live row does), so duplicate
- # destinations carry identical bytes and landing order cannot matter.
- writer = torch.full((num_slots,), -1, dtype=torch.int64, device=device)
- writer.scatter_reduce_(
- 0,
- slots,
- torch.where(live, rows, torch.full_like(rows, -1)),
- reduce="amax",
- include_self=True,
- )
- owner = writer.index_select(0, slots)
- conv_states.index_copy_(
- 0,
- slots,
- torch.where(
- (owner >= 0).reshape(num_seqs, 1, 1),
- published.index_select(0, owner.clamp_min(0)),
- previous,
- ),
- )
- return out
-
- def reference(
- self,
- x: torch.Tensor,
- weight: torch.Tensor,
- bias: Optional[torch.Tensor],
- conv_states: torch.Tensor,
- query_start_loc: torch.Tensor,
- cache_indices: Optional[torch.Tensor],
- has_initial_state: Optional[torch.Tensor],
- dilation: int,
- silu_activation: bool,
- pad_slot_id: int,
- ) -> tuple[torch.Tensor, torch.Tensor]:
- """`(out, conv_states)` as the call leaves them, one sequence at a time
- as an explicit fp32 tap-sum (no convolution op); mutates nothing.
-
- Rounds where the entry's two calls round: once after the taps, again
- after the SiLU. Publishes in sequence order, so a duplicated slot ends
- with the last live sequence's state.
- """
- channels = x.shape[0]
- kernel_size = weight.shape[-1]
- state_len = (kernel_size - 1) * dilation
- starts = query_start_loc.tolist()
- slots = None if cache_indices is None else cache_indices.tolist()
- carry = None if has_initial_state is None else has_initial_state.tolist()
- w = weight.float()
- out = torch.zeros(channels, x.shape[1], dtype=x.dtype, device=x.device)
- pool = conv_states.clone()
- for s in range(len(starts) - 1):
- begin, end = starts[s], starts[s + 1]
- length = end - begin
- slot = s if slots is None else slots[s]
- live = slots is None or slot != pad_slot_id
- if live and carry is not None and carry[s]:
- prefix = conv_states[slot]
- else:
- prefix = x.new_zeros(channels, state_len)
- window = torch.cat([prefix, x[:, begin:end]], dim=1)
- xf = window.float()
- acc = torch.zeros(channels, length, dtype=torch.float32, device=x.device)
- for k in range(kernel_size):
- acc = acc + w[:, 0, k].reshape(-1, 1) * xf[:, k * dilation : k * dilation + length]
- if bias is not None:
- acc = acc + bias.float().reshape(-1, 1)
- y = acc.to(x.dtype)
- if silu_activation:
- y = F.silu(y.float()).to(x.dtype)
- out[:, begin:end] = y
- if live:
- pool[slot] = window[:, length : length + state_len]
- return out, pool
-
-
-dilated_short_conv_paged = DilatedShortConvPaged()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/extract_transpose_prefill_slice.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/extract_transpose_prefill_slice.py
deleted file mode 100644
index 434e48ce4b55..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/extract_transpose_prefill_slice.py
+++ /dev/null
@@ -1,129 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Channel-major extraction of a prefill row/column slice, the operand a causal
-conv1d over a token-major projection needs.
-
-Wraps the in-tree Triton callable
-``tensorrt_llm._torch.modules.mamba.fuse_elementwise_ops.extract_transpose_prefill_slice``
-(one kernel launch, ``_extract_transpose_prefill_kernel``). A module-plane
-callable rather than a ``torch.ops.trtllm.*`` op: there is no op boundary
-around this kernel.
-"""
-
-from __future__ import annotations
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.modules.mamba.fuse_elementwise_ops import (
- extract_transpose_prefill_slice as _extract_transpose_prefill_slice,
-)
-
-#: The GDN mixer's packed `[q|k|v]` width for this checkpoint; the target hands
-#: this op the `[q|k|v]` column view of its one input-projection output.
-GDN_MIXED_QKV_WIDTH = 10240
-#: That output's width: `[q|k|v | z | a | b]` = 10240 + 6144 + 48 + 48.
-GDN_IN_PROJ_WIDTH = 16480
-
-
-class ExtractTransposePrefillSlice(OpWrapper):
- """`out[c, t] = src[t, start_col + c]` for `t < num_prefill_tokens`,
- `c < width`: `src[:n, a:a+w].transpose(0, 1).contiguous()` in one pass.
-
- Fusion boundary: the row crop, the column crop and the transpose. No
- convolution, no conv state, no channel split -- the caller hands the result
- to `causal_conv1d_fwd`. Returns a fresh contiguous `[width,
- num_prefill_tokens]` tensor in `src`'s dtype; `src` is read-only.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # `num_rows`/`num_cols` build the [T, 16480] projection; `row_stop` and
- # `col_slice` turn it into the view the call site passes (its q|k|v
- # columns, and in a mixed step the prefill rows of those).
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="mixed batch: prefill rows of the q|k|v view of a [512, 16480] in_proj",
- spec=dict(
- num_rows=512,
- num_cols=GDN_IN_PROJ_WIDTH,
- row_stop=400,
- col_slice=(0, GDN_MIXED_QKV_WIDTH),
- n=400,
- start_col=0,
- width=GDN_MIXED_QKV_WIDTH,
- dtype=torch.bfloat16,
- ),
- ),
- *(
- Cell(
- why=f"GDN prefill, the q|k|v column view of the in_proj, {n} tokens",
- spec=dict(
- num_rows=n,
- num_cols=GDN_IN_PROJ_WIDTH,
- col_slice=(0, GDN_MIXED_QKV_WIDTH),
- n=n,
- start_col=0,
- width=GDN_MIXED_QKV_WIDTH,
- dtype=torch.bfloat16,
- ),
- )
- for n in (33, 8192)
- ),
- )
-
- note = """
- Metadata consumed: none. Stateless.
-
- Certified bitwise in bf16 at the call site's operand, the q|k|v column
- view of the `[T, 16480]` input projection: 33 and 8192 tokens, and a mixed
- step's 400 prefill rows of a 512-row projection.
-
- Preconditions: `src` 2-D on CUDA with `src.stride(1) == 1`; the row stride
- is free. A non-unit column stride (`src.t()`, `src[:, ::2]`) is refused
- loudly by the callable's own bare `assert`, before any launch.
-
- The kernel masks its loads only against the requested `num_prefill_tokens`
- and `width`, never against `src`'s extent: 128 rows of a 64-row `src`, or
- columns `[200, 328)` of a 256-column one, were accepted and returned a
- full tensor read from out-of-bounds memory. `is_valid` bounds both.
-
- Hard limits
- width * num_prefill_tokens int32 destination offset boundary 2^31
- probed at width 16384: n 131000 (2.146e9) bitwise exact; n 140000
- (2.294e9) raises `CUDA error: an illegal memory access` (loud, so not
- a guard). The source offset is widened to int64 in-kernel: a source
- offset of 2.62e9 was bitwise exact. At width 10240 the destination
- boundary is n > 209715, within reach of a large `max_num_tokens`.
-
- Not certified: dtypes other than bf16, `start_col` other than 0, a row crop
- shorter than `src`, empty extents.
- """
-
- def __call__(
- self, src: torch.Tensor, num_prefill_tokens: int, start_col: int, width: int
- ) -> torch.Tensor:
- return _extract_transpose_prefill_slice(src, num_prefill_tokens, start_col, width)
-
- def reference(
- self, src: torch.Tensor, num_prefill_tokens: int, start_col: int, width: int
- ) -> torch.Tensor:
- return src[:num_prefill_tokens, start_col : start_col + width].transpose(0, 1).contiguous()
-
- def is_valid(
- self, src: torch.Tensor, num_prefill_tokens: int, start_col: int, width: int
- ) -> None:
- # A 3-D src with a unit dim-1 stride passes the callable's own assert
- # and is read as 2-D through stride(0).
- assert src.dim() == 2, f"src must be 2D, got {src.dim()}D"
- assert 0 <= num_prefill_tokens <= src.shape[0], (
- f"num_prefill_tokens {num_prefill_tokens} outside src rows [0, {src.shape[0]}]; "
- "the kernel reads past src rather than failing"
- )
- assert 0 <= start_col and 0 <= width and start_col + width <= src.shape[1], (
- f"column range [{start_col}, {start_col + width}) outside src columns "
- f"[0, {src.shape[1]}); the kernel reads past src rather than failing"
- )
-
-
-extract_transpose_prefill_slice = ExtractTransposePrefillSlice()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/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
deleted file mode 100644
index c0f7de9908f2..000000000000
--- 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
+++ /dev/null
@@ -1,414 +0,0 @@
-# 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 ``[N, HV, V, K]`` buffer in the pool's dtype (the sm_100/103 kernel keeps
-the state in fp32 inside and reads and writes it in that dtype), 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. With ``has_initial_state`` the call is the same
-sequence in this module, its gather a Triton kernel that writes zeros for a sequence
-carrying nothing instead of reading its slot, so the caller need not zero the slot first.
-"""
-
-from __future__ import annotations
-
-from typing import Optional, Tuple
-
-import torch
-import triton
-import triton.language as tl
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.modules.fla.flashinfer_chunk import (
- chunk_gated_delta_rule as _fi_chunk_gated_delta_rule,
-)
-from tensorrt_llm._torch.modules.fla.fused_state_io import cast_scatter_fp32_vk_to_vk
-
-#: 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]
-
-
-#: Elements of a `[V, K]` state plane a gather program copies.
-_GATHER_BLOCK = 4096
-
-
-@triton.jit
-def _gather_carried_states_kernel(
- pool_ptr,
- dst_ptr,
- slot_ptr,
- init_ptr,
- stride_ps,
- stride_ph,
- stride_pv,
- stride_pk,
- H: tl.constexpr,
- K: tl.constexpr,
- PLANE: tl.constexpr,
- BLOCK: tl.constexpr,
-):
- """One program: `BLOCK` elements of sequence `n`'s head `h` plane, into the packed
- `[N, H, V, K]` buffer: the pool row `slot_ptr[n]` when `init_ptr[n]`, else +0.0 (the
- pool row is not read)."""
- n = tl.program_id(0)
- h = tl.program_id(1)
- i = tl.program_id(2) * BLOCK + tl.arange(0, BLOCK)
- mask = i < PLANE
- dst = dst_ptr + (n.to(tl.int64) * H + h) * PLANE + i
- if tl.load(init_ptr + n) != 0:
- slot = tl.load(slot_ptr + n).to(tl.int64)
- src = (
- pool_ptr + slot * stride_ps + h * stride_ph + (i // K) * stride_pv + (i % K) * stride_pk
- )
- tl.store(dst, tl.load(src, mask=mask), mask=mask)
- else:
- tl.store(dst, tl.zeros([BLOCK], dtype=dst_ptr.dtype.element_ty), mask=mask)
-
-
-def _gather_carried_states(
- pool: torch.Tensor, indices: torch.Tensor, has_initial_state: torch.Tensor
-) -> torch.Tensor:
- """`[N, H, V, K]` in the pool's dtype: `pool[indices[n]]` for a carried sequence, zeros
- for a fresh one."""
- num_seqs = indices.shape[0]
- _, h, v, k = pool.shape
- out = torch.empty(num_seqs, h, v, k, dtype=pool.dtype, device=pool.device)
- _gather_carried_states_kernel[(num_seqs, h, triton.cdiv(v * k, _GATHER_BLOCK))](
- pool, out, indices, has_initial_state, *pool.stride(),
- H=h, K=k, PLANE=v * k, BLOCK=_GATHER_BLOCK, num_warps=4,
- ) # fmt: skip
- return out
-
-
-def _chunk_carried(
- q, k, v, g, beta, scale, pool, indices, has_initial_state, cu_seqlens, output, linear_gate
-) -> torch.Tensor:
- """The in-place indexed call of the adapter, its gather `_gather_carried_states`: the
- packed states in the pool's dtype (the sm_100/103 kernel keeps fp32 inside), FlashInfer's
- kernel (`g` exponentiated first unless `linear_gate`), and the fused scatter of the final
- states into the same slots."""
- import flashinfer
-
- q3, k3, v3 = q.squeeze(0).contiguous(), k.squeeze(0).contiguous(), v.squeeze(0).contiguous()
- packed = _gather_carried_states(pool, indices, has_initial_state)
- out = (
- output.squeeze(0)
- if output is not None
- else q3.new_empty(q3.shape[0], v3.shape[1], v3.shape[2])
- )
- final = torch.empty_like(packed)
- flashinfer.chunk_gated_delta_rule(
- q=q3,
- k=k3,
- v=v3,
- g=g.squeeze(0) if linear_gate else torch.exp(g.squeeze(0)),
- beta=beta.squeeze(0).to(torch.float32),
- scale=scale,
- initial_state=packed,
- output_final_state=True,
- cu_seqlens=cu_seqlens,
- use_qk_l2norm_in_kernel=False,
- output=out,
- output_state=final,
- )
- cast_scatter_fp32_vk_to_vk(final, pool, indices)
- return out.unsqueeze(0)
-
-
-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), or with `linear_gate` (the
- `has_initial_state` call only) the factor `exp(g)` itself, as
- `gdn_prefill_conv_post(linear_gate=True)` returns 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]]`, or, when `has_initial_state`
- (bool `[N]`) is given and `has_initial_state[n]` is false, from zeros, its
- slot not read (the in-place indexed call only). 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 (zeros for a fresh sequence with `has_initial_state`), 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); pool_dtype (fp32); output (None | "buffer"). Every cell is
- # qwen3.8-flash-next-nvfp4's geometry, in place.
- CELLS: tuple[Cell, ...] = tuple(
- Cell(
- why=f"qwen3.8 prefill into output=, 461-slot pool, seqlens {list(s)}",
- spec=dict(
- geometry=QWEN38_GEOMETRY,
- seqlens=s,
- slots=ENGINE_SLOTS,
- indices=_engine_slots(len(s)),
- output="buffer",
- ),
- )
- for s in ((200, 37, 91), (8192,))
- ) + (
- Cell(
- why="qwen3.8 context chunk + 1-token rows [8, 1, 1, 1], 461-slot pool",
- spec=dict(
- geometry=QWEN38_GEOMETRY,
- seqlens=(8, 1, 1, 1),
- slots=ENGINE_SLOTS,
- indices=_engine_slots(4),
- ),
- ),
- Cell(
- why="bf16 pool (mamba_ssm_cache_dtype), 461 slots, seqlens [200, 37, 91]",
- spec=dict(
- geometry=QWEN38_GEOMETRY,
- seqlens=(200, 37, 91),
- slots=ENGINE_SLOTS,
- indices=_engine_slots(3),
- pool_dtype=torch.bfloat16,
- ),
- ),
- Cell(
- why="target call: bf16 pool, has_initial_state [1, 0, 1], seqlens [300, 1024, 500]",
- spec=dict(
- geometry=QWEN38_GEOMETRY,
- seqlens=(300, 1024, 500),
- slots=ENGINE_SLOTS,
- indices=_engine_slots(3),
- pool_dtype=torch.bfloat16,
- output="buffer",
- initial=(True, False, True),
- linear_gate=True,
- ),
- ),
- )
-
- note = """
- Metadata consumed: none.
-
- Call site as qwen3.8-flash-next-nvfp4's `_gdn_block` makes it (in place of
- the Triton `chunk_gated_delta_rule`):
- the same positional arguments, `(q, k, v, g, beta, None, ssm_states,
- si, True, False, qsl_long, False, False, None)` with `head_first=False`
- in the 12th slot, after `fused_gdn_post_conv` (q/k already normalized,
- contiguous `[1, T, H, 128]`; `g` fp32 log-space; `beta` fp32). Give it
- context rows only: single-token rows are slow here. Measured per layer
- at 8065 + 127 x 1 tokens: 1078.8 us with every row in this op, 399.4 us
- with the 127 decode rows sent to FlashInfer's decode kernel
- (`gdn_decode.gated_delta_rule_decode_pretranspose`) instead (both halves
- plus the concatenation).
-
- The target (`PrefillTarget._gdn_core`, `_gdn_core_verify`) passes the
- context rows by keyword after `gdn_prefill_conv_post(linear_gate=True)`,
- with `has_initial_state` and `linear_gate=True`.
-
- Call site with `has_initial_state` (`hi_p` / `hi_all`, the flags
- `reset_recurrent_state_rows` would zero the fresh slots by): a fresh
- sequence's slot is never read, and the scatter writes every named slot,
- so the reset of the recurrent pool before this call is not needed (the
- conv pass reads no carry for such a sequence and its publish writes the
- whole conv window). Bitwise the reset followed by the plain call; the
- gather moves only the carried sequences' rows.
-
- 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; and at the 461-slot bf16 pool: three
- ragged context rows, and the target's call (300 + 1024 + 500 tokens,
- the middle one fresh, `has_initial_state`, `linear_gate`, `output=`).
-
- Memory per call: a packed copy of the named slots in and out in the pool's
- dtype (2 x 3 MiB per sequence fp32, half that bf16), `exp(g)` (not with
- `linear_gate`), and the
- output. q/k/v/beta are not copied
- when already contiguous / fp32.
-
- JIT and streams: the first call per process compiles the CuTe-DSL
- kernel in-process (seconds), and the installed FlashInfer keeps one TMA
- workspace per device, so two streams must not run this op concurrently.
- Needs CUDA >= 13 torch (FlashInfer raises otherwise).
-
- Numerics (measured vs an fp64 recurrence, 1 x 8192 tokens, qwen3.8
- geometry): o rel-L2 4.06e-3, max abs 1.26e-3; state rel-L2 3.31e-3 --
- the Triton op's are
- 4.00e-3 / 1.13e-3 / 3.26e-3. Not bitwise to it. A bf16 pool rounds each
- sequence's state to bf16 once per call, on the scatter. 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=`, an fp16 pool, `scale` other than the
- default, fp16 operands, K or V other than 128, other head ratios, and what the
- target never passes: `inplace_indexed_state_update=False` (final states
- returned) and `use_qk_l2norm_in_kernel`.
-
- 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,
- has_initial_state: Optional[torch.Tensor] = None,
- linear_gate: bool = False,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
- assert not linear_gate or has_initial_state is not None, (
- "linear_gate is taken with has_initial_state only"
- )
- if has_initial_state is not None:
- assert (
- inplace_indexed_state_update
- and initial_state_indices is not None
- and not output_final_state
- and not head_first
- and not use_qk_l2norm_in_kernel
- and state_workspace is None
- and cu_seqlens is not None
- and g.dtype == torch.float32
- ), "has_initial_state: the in-place indexed call only, q/k normalized by the caller"
- o = _chunk_carried(
- q, k, v, g, beta, scale, initial_state, initial_state_indices,
- has_initial_state, cu_seqlens, output, linear_gate,
- ) # fmt: skip
- return o, None
- 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,
- has_initial_state: Optional[torch.Tensor] = None,
- linear_gate: bool = False,
- ) -> 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()
- carried = [True] * len(slots) if has_initial_state is None else has_initial_state.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)
- if not carried[n]:
- state = torch.zeros_like(state)
- 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)
- decay = g[0, t].to(dt) if linear_gate else g[0, t].to(dt).exp()
- state = state * decay[:, None, None]
- u = (v[0, t].to(dt) - (state * k_t[:, None, :]).sum(-1)) * beta[0, t].to(dt)[
- :, None
- ]
- state = state + u[:, :, None] * k_t[:, None, :]
- out[0, t] = (state * q_t[:, None, :]).sum(-1)
- finals.append(state)
- pool = initial_state.clone()
- final_state = None
- if inplace_indexed_state_update:
- for slot, state in zip(slots, finals):
- pool[slot] = state.to(pool.dtype)
- elif output_final_state:
- final_state = torch.stack(finals).to(initial_state.dtype)
- return out.to(q.dtype), final_state, pool
-
-
-flashinfer_chunk_gated_delta_rule = FlashinferChunkGatedDeltaRule()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_gating.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_gating.py
deleted file mode 100644
index bc996ad6dd0d..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_gating.py
+++ /dev/null
@@ -1,153 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Gated-DeltaNet decay gate: g = -exp(A_log) * softplus(a + dt_bias), per token
-and per value head, in one Triton kernel.
-
-Wraps the in-tree module callable
-``tensorrt_llm._torch.modules.mamba.gdn_mixer.fused_gdn_gating`` (tier-2
-sourcing: no op boundary exists). Its body allocates the fp32 output and
-launches ``fused_gdn_gating_kernel`` once.
-"""
-
-from __future__ import annotations
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.modules.mamba.gdn_mixer import fused_gdn_gating as _fused_gdn_gating
-
-#: This checkpoint's GDN value-head count at tp1, and the `[b | a]` width the
-#: stock mixer slices `a` out of.
-NUM_V_HEADS = 48
-BA_WIDTH = 2 * NUM_V_HEADS
-
-#: Read off this checkpoint's 48 layers of `linear_attn.A_log` /
-#: `linear_attn.dt_bias` (bf16, 48 each): exp(A_log) spans [3.839e-03, 1.394e+02],
-#: and dt_bias reaching 19.25 against the default threshold of 20.0 makes the
-#: softplus's linear branch a live path, not a corner case.
-A_LOG_RANGE = (-5.5625, 4.9375)
-DT_BIAS_RANGE = (-8.9375, 19.25)
-
-#: Relative error of the SASS approximate transcendentals (`ex2.approx.f32`,
-#: `lg2.approx.f32`) Triton lowers `tl.exp` / `tl.log` to.
-APPROX_ULP = 2.0**-22
-
-
-class FusedGdnGating(OpWrapper):
- """`g[t, h] = -exp(A_log[h]) * softplus(a[t, h] + dt_bias[h])`, computed in
- fp32 whatever the operand dtypes, into a fresh contiguous fp32
- `[batch, num_heads]` tensor. `batch, num_heads = a.shape`.
-
- `softplus(x) = log(1 + exp(beta * x)) / beta` while `beta * x <= threshold`,
- else `x` -- the unstabilised form: below `x = -24 ln 2 = -16.6355` the fp32
- sum drops `exp(x)` and `g` is exactly 0 (absolute error <= 6e-8 times
- exp(A_log), relative error 1.0). Harmless for a log-domain decay.
-
- Fusion boundary: bias add, softplus with its threshold branch, `exp(A_log)`
- and the product are inside. The projection producing `a`, the
- recurrence's other gate `beta = sigmoid(b)`, and the recurrence are not.
- No input is modified.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: rows; layout "contiguous" | "target_rows" (the decode rows of a
- # contiguous [T, 48] `a` projection, past T - rows prefill rows); fill
- # (constant a). `a` bf16, softplus beta 1.0, threshold 20 throughout.
- CELLS: tuple[Cell, ...] = (
- Cell(why="target verify call: 32 seqs x 2 tokens, contiguous a", spec=dict(rows=64)),
- Cell(
- why="target mixed verify: decode rows of a [T, 48] projection",
- spec=dict(rows=16, layout="target_rows"),
- ),
- Cell(
- why="a = 200: linear branch; exp(x) is evaluated, inf never leaks",
- spec=dict(rows=4, fill=200.0),
- ),
- )
-
- note = """
- Metadata consumed: none. Stateless and per (token, head).
-
- Call site (`_gdn_verify`): `(A_log, a_d, dt_bias, 1.0, 20.0)` with `a_d`
- the `[n * p, 48]` bf16 decode rows of the `a` projection (all of it on a
- pure verify step) and fp32 `A_log` / `dt_bias`; `g` is then reshaped to
- `[n, p, 48]` for `fused_recurrent_gated_delta_rule_update`. Certified at
- that call: 32 sequences of 2 tokens, the 16 decode rows of a mixed step's
- `[T, 48]` projection, and `a = 200` (the softplus's linear branch, where
- `exp` overflows).
-
- Raise loudly, so not guarded: a non-2-D `a` (`ValueError` unpacking
- `a.shape`), CPU tensors (`ValueError: Pointer argument cannot be accessed
- from Triton`).
-
- `a`'s row stride is honoured (a row is read at `row * a.stride(0)`; the
- target passes contiguous rows); everything else is addressed without
- strides. `is_valid` guards the three violations measured to return a finite wrong `g`: a stride-2 A_log
- view moved g by 6.031e+01 and a stride-2 dt_bias by 3.832e-01; 40-element
- A_log / dt_bias against 48 heads moved it by 1.009e+02 / 1.010e+02 (the
- per-head load is masked against `a.shape[1]` only); an `a` with head
- stride 2 by 1.153e+00.
-
- Deterministic: two launches agree bitwise.
-
- Hard limits
- none probed.
- """
-
- def __call__(
- self,
- A_log: torch.Tensor,
- a: torch.Tensor,
- dt_bias: torch.Tensor,
- beta: float = 1.0,
- threshold: float = 20.0,
- ) -> torch.Tensor:
- return _fused_gdn_gating(A_log, a, dt_bias, beta, threshold)
-
- def reference(
- self,
- A_log: torch.Tensor,
- a: torch.Tensor,
- dt_bias: torch.Tensor,
- beta: float = 1.0,
- threshold: float = 20.0,
- ) -> torch.Tensor:
- """fp64 with the stabilised `log1p(exp(x))`, rounded once to fp32.
-
- The stabilised form makes the reference the accurate side; the
- kernel's unstabilised tail differs inside the test's absolute floor.
- """
- num_heads = a.shape[-1]
- x = a.double() + dt_bias[:num_heads].double()
- softplus = torch.where(beta * x <= threshold, torch.log1p(torch.exp(beta * x)) / beta, x)
- return (-torch.exp(A_log[:num_heads].double()) * softplus).float()
-
- def is_valid(
- self,
- A_log: torch.Tensor,
- a: torch.Tensor,
- dt_bias: torch.Tensor,
- beta: float = 1.0,
- threshold: float = 20.0,
- ) -> None:
- # `shape[-1]` / `stride(-1)` so a non-2-D `a` still reaches the op's own
- # ValueError rather than being preempted here.
- num_heads = a.shape[-1]
- assert a.stride(-1) == 1, (
- f"`a`'s head axis must be unit-stride; got strides {a.stride()} -- the kernel "
- "reads a row as `row * a.stride(0) + head`, so another column stride returns "
- "a finite wrong g"
- )
- assert A_log.stride(0) == 1 and dt_bias.stride(0) == 1, (
- f"A_log and dt_bias must be contiguous 1-D; got strides {A_log.stride()} and "
- f"{dt_bias.stride()} -- both are read as `ptr + head`"
- )
- assert A_log.numel() >= num_heads and dt_bias.numel() >= num_heads, (
- f"A_log and dt_bias must hold at least a.shape[-1]={num_heads} elements; got "
- f"{A_log.numel()} and {dt_bias.numel()} -- the per-head load is masked against "
- "a.shape[-1] only, so a short tensor is read out of bounds"
- )
-
-
-fused_gdn_gating = FusedGdnGating()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_post_conv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_post_conv.py
deleted file mode 100644
index e34c3fdb77f1..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_post_conv.py
+++ /dev/null
@@ -1,232 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Gated-DeltaNet post-convolution preparation: the q/k/v split of the causal
-conv output with Q/K L2 normalization, plus the decay gate and beta.
-
-Wraps the in-tree module callable
-``tensorrt_llm._torch.modules.mamba.fuse_elementwise_ops.fused_gdn_post_conv``
-(tier-2 sourcing: no op boundary exists). One Triton launch,
-``_fused_gdn_post_conv_kernel``, after allocating its five outputs.
-"""
-
-from __future__ import annotations
-
-from typing import Optional
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.modules.mamba.fuse_elementwise_ops import (
- fused_gdn_post_conv as _fused_gdn_post_conv,
-)
-
-#: This checkpoint's linear-attention geometry at tp1.
-NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS, HEAD_V_DIM = 16, 128, 48, 128
-QKV_DIM = 2 * NUM_K_HEADS * HEAD_K_DIM + NUM_V_HEADS * HEAD_V_DIM # 10240
-
-#: Read off this checkpoint's `linear_attn.A_log` / `linear_attn.dt_bias`
-#: (bf16, 48 per layer): A_log spans [-5.5625, 3.5625] and dt_bias
-#: [-8.9375, 19.25]. dt_bias reaching 19.25 makes the softplus threshold branch
-#: a live path, not a corner case.
-A_LOG_RANGE = (-5.5625, 3.5625)
-DT_BIAS_RANGE = (-8.9375, 19.25)
-SOFTPLUS_THRESHOLD = 20.0
-
-
-class FusedGdnPostConv(OpWrapper):
- """Split a causal-conv output into q/k/v and build the GDN gates.
-
- The token axis is the `num_prefill_tokens` columns of the dim-major
- `prefill` `[>= qkv_dim, n_p]`, then the rows of the token-major `decode`
- `[n_d, >= qkv_dim]` (`None` for a prefill-only batch), with
- `qkv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim` packed
- `[q | k | v]`. Per token and head, in fp32:
-
- q, k = x / sqrt(sum(x^2) + l2_norm_eps) (per head)
- v = copied bit-exact
- g = -exp(A_log) * softplus(a + dt_bias) (fp32 always)
- beta = sigmoid(b) (fp32, rounded once to beta_dtype)
-
- `softplus(x)` is the unstabilised `log(1 + exp(x))` below
- `softplus_threshold` and `x` above it.
-
- Fusion boundary: split, Q/K norm, `g` and `beta` are inside. The conv is
- not (`prefill`/`decode` are already its output), and the recurrence that
- consumes these must not normalize q/k again (`use_qk_l2norm_in_kernel=False`).
-
- Returns `(q, k, v, g, beta)`, five fresh contiguous tensors:
- `[1, T, num_k_heads, head_k_dim]` x2 and `[1, T, num_v_heads, head_v_dim]`
- in `prefill.dtype`, then `[1, T, num_v_heads]` fp32 and `beta_dtype`.
- No input is modified.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: num_prefill, num_decode; dt_bias_fill. Every cell is the target's
- # one [T, 16480] input projection: `decode` the q|k|v column view of its
- # decode rows, `a`/`b` its column views, `prefill` the transposing
- # extraction of its prefill rows.
- CELLS: tuple[Cell, ...] = (
- *(
- Cell(
- why=f"target in_proj views, {p} prefill + {d} decode tokens",
- spec=dict(num_prefill=p, num_decode=d),
- )
- for p, d in ((200, 32), (2048, 0))
- ),
- Cell(
- why="every head above softplus_threshold: the linear branch",
- spec=dict(num_prefill=64, num_decode=0, dt_bias_fill=SOFTPLUS_THRESHOLD + 8.0),
- ),
- )
-
- note = """
- Metadata consumed: none. Stateless and per token; it never needs sequence
- boundaries.
-
- Call sites (`_gdn_block`): prefill-only `(mq_t, None, a, b, ...)`, mixed
- `(mq_p_t, mq_d, a, b, ...)` and the MTP-verify context half
- `(mq_p_t, None, a_p, b_p, ...)`, all with `l2_norm_eps=1e-6`,
- `softplus_threshold=20.0` and the fp32 `beta_dtype` default; `prefill` is
- `extract_transpose_prefill_slice`'s contiguous output; `decode` and `a`/`b`
- column views (row stride 16480) of the layer's one `[T, 16480]`
- input-projection output. The stock mixer's verify call passes
- `beta_dtype=b.dtype` (not certified here). Certified in that form at a
- mixed step of 200 prefill and 32 decode tokens, a 2048-token prefill-only
- step, and a 64-token step with every head past the softplus threshold.
-
- Preconditions: `prefill.shape[0]` and `decode.shape[1]` >= qkv_dim (extra
- trailing features ignored; a short one is read out of bounds); `decode`
- has `prefill`'s dtype. `a`/`b` may have extra columns (the in_proj
- case) and are read through their own strides. `head_k_dim`/`head_v_dim`
- are rounded up to a power of two and masked.
-
- The op raises on none of the metadata violations probed; `is_valid`
- guards the four the kernel reads out of bounds or through the wrong
- addresses, measured: stride-2 A_log moved g by 2.8e+02, stride-2 dt_bias by
- 3.2e+02; A_log, dt_bias, a, b short by rows/elements moved g by 4.9e+12,
- 2.9e+02, 3.2e+37 and made beta NaN.
-
- The softplus is unstabilised: for very negative `x` the fp32 `1 + exp(x)`
- loses `exp(x)` and `g` rounds to exactly 0. Absolute error stays ~1 fp32
- ulp of the log argument times `exp(A_log)` (4.8e-07 measured over
- x in [-9, 5]); relative error in that tail reaches 1.0. Inconsequential
- for a log-domain decay. `exp(x)` is evaluated on the linear side too
- (`tl.where` is a select); harmless below fp32 overflow at x ~ 88, which no
- checkpoint `a + dt_bias` reaches.
-
- Not certified: fp16/fp32 activations, other head geometries, a bf16
- `beta_dtype`, bf16 `A_log`/`dt_bias`, a non-contiguous `prefill`, contiguous
- `decode`/`a`/`b` (not views of the projection), single-token halves, an
- empty batch.
-
- Hard limits
- none probed.
- """
-
- def __call__(
- self,
- prefill: torch.Tensor,
- decode: Optional[torch.Tensor],
- a: torch.Tensor,
- b: torch.Tensor,
- A_log: torch.Tensor,
- dt_bias: torch.Tensor,
- num_k_heads: int,
- head_k_dim: int,
- num_v_heads: int,
- head_v_dim: int,
- l2_norm_eps: float = 1e-6,
- softplus_threshold: float = 20.0,
- beta_dtype: torch.dtype = torch.float32,
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
- return _fused_gdn_post_conv(
- prefill,
- decode,
- a,
- b,
- A_log,
- dt_bias,
- num_k_heads,
- head_k_dim,
- num_v_heads,
- head_v_dim,
- l2_norm_eps,
- softplus_threshold,
- beta_dtype,
- )
-
- def reference(
- self,
- prefill: torch.Tensor,
- decode: Optional[torch.Tensor],
- a: torch.Tensor,
- b: torch.Tensor,
- A_log: torch.Tensor,
- dt_bias: torch.Tensor,
- num_k_heads: int,
- head_k_dim: int,
- num_v_heads: int,
- head_v_dim: int,
- l2_norm_eps: float = 1e-6,
- softplus_threshold: float = 20.0,
- beta_dtype: torch.dtype = torch.float32,
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
- """q/k normalized in fp32 and rounded once; `g` in fp64 with the
- stabilised `log1p(exp(x))` (the kernel's unstabilised form differs
- inside the test's floor, not around it), rounded to fp32; `beta` is
- the fp32 sigmoid rounded once to `beta_dtype`."""
- blocks = [prefill.t()[:, : 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim]]
- if decode is not None:
- blocks.append(decode[:, : blocks[0].shape[1]])
- mixed = torch.cat(blocks, dim=0)
- n = mixed.shape[0]
- kd = num_k_heads * head_k_dim
-
- def l2norm(x):
- x = x.reshape(n, num_k_heads, head_k_dim).float()
- return (x / torch.sqrt(x.pow(2).sum(-1, keepdim=True) + l2_norm_eps)).to(prefill.dtype)
-
- q, k = l2norm(mixed[:, :kd]), l2norm(mixed[:, kd : 2 * kd])
- v = mixed[:, 2 * kd :].reshape(n, num_v_heads, head_v_dim).contiguous()
- x = a[:n, :num_v_heads].double() + dt_bias[:num_v_heads].double()
- softplus = torch.where(x <= softplus_threshold, torch.log1p(torch.exp(x)), x)
- g = (-torch.exp(A_log[:num_v_heads].double()) * softplus).float()
- beta = torch.sigmoid(b[:n, :num_v_heads].float()).to(beta_dtype)
- return tuple(t.unsqueeze(0) for t in (q, k, v, g, beta))
-
- def is_valid(
- self,
- prefill: torch.Tensor,
- decode: Optional[torch.Tensor],
- a: torch.Tensor,
- b: torch.Tensor,
- A_log: torch.Tensor,
- dt_bias: torch.Tensor,
- num_k_heads: int,
- head_k_dim: int,
- num_v_heads: int,
- head_v_dim: int,
- l2_norm_eps: float = 1e-6,
- softplus_threshold: float = 20.0,
- beta_dtype: torch.dtype = torch.float32,
- ) -> None:
- # The kernel indexes A_log/dt_bias as `ptr + head` with no stride and
- # no mask, and does not mask a/b's token axis against their extent.
- num_tokens = prefill.shape[1] + (0 if decode is None else decode.shape[0])
- assert A_log.stride(0) == 1 and dt_bias.stride(0) == 1, (
- f"A_log and dt_bias must be contiguous 1-D; got strides {A_log.stride()} "
- f"and {dt_bias.stride()} -- read as `ptr + head`, a strided view gives a "
- "finite wrong g"
- )
- assert A_log.numel() >= num_v_heads and dt_bias.numel() >= num_v_heads, (
- f"A_log and dt_bias must hold at least num_v_heads={num_v_heads} elements; "
- f"got {A_log.numel()} and {dt_bias.numel()} -- the load is unmasked"
- )
- assert a.shape[0] >= num_tokens and b.shape[0] >= num_tokens, (
- f"a and b must have at least num_tokens={num_tokens} rows (prefill + decode); "
- f"got {a.shape[0]} and {b.shape[0]} -- the token axis is unmasked"
- )
-
-
-fused_gdn_post_conv = FusedGdnPostConv()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_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
deleted file mode 100644
index 7ed03ff9eb23..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py
+++ /dev/null
@@ -1,512 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Gated DeltaNet speculative-verify recurrence under the engine's cached-replay
-state update: a checkpoint state per slot plus a compact history of the tokens
-accepted since it, instead of a per-step state buffer.
-
-Wraps the in-tree module callable
-``tensorrt_llm._torch.modules.fla.cached_replay.fused_recurrent_gated_delta_rule_cached_replay_update``
-(tier-2 sourcing: the recurrence has no op boundary). One launch of the vendored
-Triton kernel ``_cached_replay_kernel`` -- plain ``triton.jit``, no autotuner;
-the Python body picks the launch configuration from the shapes.
-"""
-
-from __future__ import annotations
-
-from typing import Optional
-
-import torch
-import torch.nn.functional as F
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.modules.fla.cached_replay import (
- fused_recurrent_gated_delta_rule_cached_replay_update as _cached_replay_update,
-)
-
-#: This checkpoint's GDN geometry at tp1.
-NUM_K_HEADS, NUM_V_HEADS, HEAD_K_DIM, HEAD_V_DIM = 16, 48, 128, 128
-#: `[q | k | v]` row width of the post-conv activations a verify step packs.
-PACKED_QKV_WIDTH = 2 * NUM_K_HEADS * HEAD_K_DIM + NUM_V_HEADS * HEAD_V_DIM # 10240
-#: The cache manager's `replay_history_size`: max(MIN_REPLAY_HISTORY_SIZE=16, T).
-HISTORY_SIZE = 16
-#: From this batch size on the engine runs the partitioned replay (work items, no
-#: in-kernel commit, one all-layer commit after every layer has run):
-#: `CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE` in `modules/fla/cached_replay.py`.
-ALL_LAYER_COMMIT_MIN_BATCH = 16
-#: The recurrent pool's slot count of a full-size engine run.
-ENGINE_SLOTS = 461
-#: Layers per slot at which the last of `ENGINE_SLOTS` slots starts past 2^31 elements of
-#: the pool (`460 * L * 48 * 128 * 128 > 2**31` from L = 6).
-PAST_INT32_LAYERS = 8
-
-
-def _cell_spec(path, n, t, **spec) -> dict:
- return dict(path=path, n=n, t=t, **spec)
-
-
-def _bf16_cell(why, path, n, t, **spec) -> Cell:
- return Cell(
- why=f"bf16 state {why}", spec=_cell_spec(path, n, t, state_dtype=torch.bfloat16, **spec)
- )
-
-
-class FusedRecurrentGatedDeltaRuleCachedReplayUpdate(OpWrapper):
- """One verify step of `N` sequences x `T` tokens through one GDN layer, for
- the cached-replay state layout. Per live row (`slot = state_indices[n] >= 0`),
- with `p = prev_num_accepted_tokens[slot]`, `b = cache_buf_idx[slot]`:
-
- 1. the state the step starts from is rebuilt from the checkpoint
- `h0 = ssm_states[slot]` (`[HV, V, K]`, K innermost) and the `p` pending
- history entries in half `b` -- update vectors `u_j = old_u[slot, b, j]`,
- normalized keys `k_j = old_k[slot, b, j]`, cumulative log-decays
- `G_j = old_G[slot, b, :, j]` measured from the checkpoint:
- `h = exp(G_{p-1}) h0 + sum_j exp(G_{p-1} - G_j) u_j k_j^T`;
- 2. the `T` new tokens run the gated delta rule in order (q/k head
- `hv // (HV // H)`, so a v-head group shares one q/k head):
- `h *= exp(g_t); u_t = beta_t (v_t - h k_t); h += u_t k_t^T; o_t = h (scale q_t)`,
- q and k L2-normalized first (`x / sqrt(sum(x^2) + 1e-6)`) when
- `use_qk_l2norm_in_kernel`, `scale=None` meaning `K ** -0.5`; with
- `A_log` given, `g = -exp(A_log) * softplus(g + dt_bias)` and
- `beta = sigmoid(beta)` are applied to the raw projections inside;
- 3. the new tokens' `(u_t, k_t, G_t)` are appended to the history. When
- `p + T <= history_size` they go to half `b` at `p .. p+T-1` with `G`
- still cumulative from the checkpoint. Otherwise (a *wrap*) they go to
- half `1 - b` at `0 .. T-1` with `G` restarted from zero, and the
- checkpoint must advance to the state after the `p` pending tokens:
- without work items the kernel writes that state into `ssm_states[slot]`
- itself; with `replay_work_items` (`use_all_layer_commit=True`) it does
- not, and the caller owes one `commit_gdn_cached_replay_history_layers`
- over every layer before the next step.
-
- Nothing about acceptance happens here: the caller advances `p`/`b` after
- sampling (`p <- accepted` and `b <- 1 - b` on a wrap, else `p += accepted`,
- `accepted` counting the bonus token), so rejected tokens' history entries
- are simply never read.
-
- Fusion boundary: L2 norms, scale, the gating of the raw `a`/`b`
- projections, the history rebuild, the recurrence, the history append and
- (small batches) the checkpoint commit are inside. The conv, the projections,
- the all-layer commit and the `p`/`b` bookkeeping are not.
-
- Returns `[N, T, HV, V]` in `output`'s dtype, else `q`'s; it *is* `output`
- when one is given. `ssm_states`, `old_u`, `old_k`, `old_G` are written in
- place and not returned. Rows with a negative slot are skipped entirely:
- nothing is read or written for them, their output rows included.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: path ("small": no work items, the kernel commits a wrap itself |
- # "all_layer": work items + use_all_layer_commit, the test then runs the
- # manager's all-layer commit), n sequences x t tokens per verify step,
- # steps (consecutive verify steps, each followed by an acceptance per row),
- # accept ("random": 1..t per row, the default | "full": t every step, so from
- # a fresh slot at t=8 the second step appends at `pnat + t == 16` without
- # wrapping and every second step after wraps),
- # state_dtype (the checkpoint pool's dtype: fp32 | bf16, as
- # `mamba_ssm_cache_dtype` allocates it; fp32 by default),
- # layers (each with its own activations, gates
- # and history; one V2-style pool interleaving the layers per slot), slots
- # (the pool's slot count), start ("warm": random pending histories in both
- # halves, a wrap and an append forced at step 0 | "fresh": pnat 0 and
- # zeroed history, as the manager leaves a slot after prefill), edge_slots
- # (rows 0 and n-1 on the pool's last and first slot), dup (rows given the
- # SAME slot, as CUDA-graph padding repeats one dummy slot; only the other
- # rows and slots are held to the reference). Every cell passes `packed_qkv`,
- # separate raw `a`/`b` projections, no `output` and `launch_with_pdl=False`.
- #
- # Four cells, one over the three-cell budget: the forward reaches two commit
- # paths (in-kernel below 16 requests, all-layer from 16) at two checkpoint
- # dtypes, and two edges no main-shape cell can carry -- CUDA-graph padding
- # rows sharing one slot, and slot offsets past 2^31 elements at the engine's
- # pool.
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="all_layer n=32 T=8 full accept: commit every second step",
- spec=_cell_spec("all_layer", 32, 8, start="fresh", accept="full", steps=8),
- ),
- _bf16_cell(
- "small n=4 T=8 full accept: pnat+T==16 appends, then wraps",
- "small",
- 4,
- 8,
- start="fresh",
- accept="full",
- steps=8,
- ),
- Cell(
- why="small n=8 T=3 dup: rows 5..7 share one padding slot",
- spec=_cell_spec("small", 8, 3, dup=(5, 6, 7), steps=6),
- ),
- Cell(
- why="all_layer n=16 T=3 L=8 S=461: slot 460's offset is past 2^31 elements",
- spec=_cell_spec(
- "all_layer",
- 16,
- 3,
- layers=PAST_INT32_LAYERS,
- slots=ENGINE_SLOTS,
- edge_slots=True,
- steps=3,
- ),
- ),
- )
-
- note = """
- Call sites. The engine's GDN mixer (`modules/mamba/gdn_mixer.py`) is the
- reference caller: the target-verify step of a generation-only batch, and the
- decode rows of a mixed context + verify batch -- `N = num_decodes`, `T =
- runtime_draft_len + 1` (2..8 here: max_draft_len 1..7). A plain decode step never reaches it. q/k/v
- are the `[N, T, 16|48, 128]` views of the post-conv `[N*T, 10240]` rows, plus
- `packed_qkv` = those rows for a generation-only batch (a mixed batch passes the
- views alone and the op copies them contiguous); `g=a`, `beta=b` are the raw
- bf16 `[N, T, 48]` projections, with the layer's fp32 `[48]` `A_log` and
- `dt_bias`; `ssm_states` is the layer's pool view `[S, 48, 128, 128]`, fp32 or bf16 as
- `mamba_ssm_cache_dtype` allocated it (V2 interleaves the layers inside each slot, so
- the slot stride is a multiple of one state); `state_indices` the decode rows' int32 slots; `old_u, old_k, old_G,
- old_beta` the layer's `old_x, old_B, old_dt, old_dA_cumsum` from
- `mamba_layer_cache(layer)`; `cache_buf_idx`, `prev_num_accepted_tokens`,
- `history_size` from `get_replay_state_update_metadata()` (16 here);
- `use_qk_l2norm_in_kernel=True`; and exactly when `N >= 16` and the manager's
- `use_gdn_cached_replay_all_layer_commit` holds, `replay_work_items =
- mamba_metadata.replay_work_items[:N]`, `n_writes = mamba_metadata.replay_n_writes`,
- `use_all_layer_commit=True`. That condition has to be the manager's own:
- `update_mamba_states` runs the all-layer commit iff it holds, so work items
- passed when it does not leave every wrap uncommitted, and work items withheld
- when it does commit twice -- both silent at the op and far off the exact
- recurrence one step later.
-
- Metadata consumed: `state_indices` (one slot per row), and per slot
- `prev_num_accepted_tokens` and `cache_buf_idx`. `replay_work_items` /
- `n_writes` are only checked (shape, dtype); the kernel never reads them, their
- presence switches the in-kernel commit off. `old_beta` is never read or written.
-
- Preconditions not checkable from metadata:
- * rows on distinct slots, except padding: two rows on one slot both append to the
- same history positions and the appended u is a per-element mixture of the two
- rows, so that slot's state and those rows' outputs are unspecified. The rows of
- CUDA-graph padding repeat one shared dummy slot (the manager's `_padding_slot`);
- the rows on unique slots are exact whatever the padding does, and the dummy slot
- is nobody's state;
- * `prev_num_accepted_tokens` in `[0, history_size]`, `cache_buf_idx` in `{0, 1}`,
- and the pending entries written by earlier calls of this op. A slot fresh from
- prefill starts at 0 pending (the manager zeroes the count, the half and the
- history);
- * after every step, before the next: with work items,
- `commit_gdn_cached_replay_history_layers` over every local GDN layer; then the
- advance (`_advance_replay_state` in `modules/fla/cache_manager.py`)
- with the accepted counts including the bonus token. Advancing as if every
- draft were accepted is silent at the op and wrong from the next step on;
- * the conv state is not this op's: under replay the conv still writes its
- per-step window and the manager promotes it.
-
- The launch is one V tile, `num_warps=4`, `num_stages=1`, grid `(1, N * 48)`: the
- tuned mappings need HV = 2H or 4H and never engage at HV/H = 3. The checkpoint's dtype
- is the kernel's `STATE_FP32` specialization and nothing else below N = 128: an fp32
- checkpoint enters the history rebuild's tensor-core dots as a bf16 hi/lo pair, a bf16
- one as stored, and every commit is computed in fp32 and rounded once to the pool's
- dtype. From N = 128 a bf16 state alone also turns on the L2-streaming mapping
- (`main_num_stages=3`, a 5-stage commit), which no cell drives. `FLA_USE_FAST_OPS=1` at
- import switches the kernel's `exp` to the approximate one; the cells run with it unset.
-
- CUDA-graph replay on new data is bit for bit the eager call (the test captures
- N = 32 at T = 8 over an fp32 checkpoint, with the all-layer commit). The two commit
- paths, `packed_qkv` against separate q/k/v, and `launch_with_pdl` on and off are not
- certified against each other; each form the target calls is certified against the
- reference by the cells.
-
- Certified: the all-layer commit over an fp32 checkpoint at N = 32, T = 8 and at
- N = 16, T = 3 with 8 interleaved layers over the engine's 461-slot pool (the last
- slot's offset past 2^31 elements); the in-kernel commit over a bf16 checkpoint at
- N = 4, T = 8 and over an fp32 one at N = 8, T = 3 with three CUDA-graph padding rows
- on one shared slot. T = 5..8 share one compiled token tile (`BT = next_pow2(T) = 8`);
- the cells drive the BT = 4 and BT = 8 tiles, not BT = 2 (T = 2). At T = 8 and history
- 16 a slot whose every draft is accepted appends rows 8..15 at `pnat + T == 16` without
- wrapping and wraps on the next step, so it commits every second step; the
- `accept="full"` cells drive that boundary (in-kernel and all-layer commit).
-
- Loud refusals: work items without `use_all_layer_commit` (ValueError) and the
- reverse (ValueError); one of `replay_work_items` / `n_writes` alone
- (ValueError); a strided `output`; a strided head axis of `g`; `history_size <
- T` (a bare assert); and `history_size <= 8`, which does not compile -- the
- history tile is `next_pow2(history_size)` wide and feeds a `tl.dot` whose
- reduction must be 16, so only 9..16 work (the engine's is 16).
-
- `is_valid` guards what is accepted and silently wrong: `state_indices` shorter
- than N (the trailing rows read past it and append to a slot nobody named);
- history buffers whose shape disagrees with `history_size` or the head geometry
- (the kernel addresses them by `history_size`, so 16-deep buffers called at 12
- read other slots' entries); history buffers, `prev_num_accepted_tokens` or
- `cache_buf_idx` covering fewer slots than `ssm_states` (a slot of the pool past
- their end is read and written out of bounds); a strided `prev_num_accepted_tokens`
- or `cache_buf_idx` (read at element `slot`). A strided `state_indices` is
- harmless: the entry's input guard makes it contiguous.
-
- Not certified: pre-gated `g`/`beta` (`A_log=None`); `use_qk_l2norm_in_kernel=
- False`; an `ssm_states` other than fp32 or bf16, or a bf16 one at N >= 128; HV/H other
- than 3, K or V other than 128; T = 2, or T outside 2..8 (T = 9 raises: the kernel's
- bare `assert T <= 8`), at either `ssm_states` dtype; `history_size` other than 16; `old_u` /
- `old_k` other than bf16, `old_G` other than fp32; a `scale`; any launch-knob
- override (`block_v`, `num_warps`, `use_l2_*`, `main_num_stages`); an output
- dtype other than bf16; int64 `state_indices`; N above 32.
- """
-
- def __call__(
- self,
- q: torch.Tensor,
- k: torch.Tensor,
- v: torch.Tensor,
- g: torch.Tensor,
- beta: torch.Tensor,
- ssm_states: torch.Tensor,
- state_indices: torch.Tensor,
- old_u: torch.Tensor,
- old_k: torch.Tensor,
- old_G: torch.Tensor,
- old_beta: torch.Tensor,
- cache_buf_idx: torch.Tensor,
- prev_num_accepted_tokens: torch.Tensor,
- history_size: int,
- scale: Optional[float] = None,
- use_qk_l2norm_in_kernel: bool = False,
- A_log: Optional[torch.Tensor] = None,
- dt_bias: Optional[torch.Tensor] = None,
- launch_with_pdl: bool = False,
- replay_work_items: Optional[torch.Tensor] = None,
- n_writes: Optional[torch.Tensor] = None,
- block_v: Optional[int] = None,
- num_warps: Optional[int] = None,
- use_l2_state_cache: Optional[bool] = None,
- use_l2_streaming_inputs: Optional[bool] = None,
- use_l2_shared_inputs: Optional[bool] = None,
- main_num_stages: Optional[int] = None,
- packed_qkv: Optional[torch.Tensor] = None,
- use_all_layer_commit: bool = False,
- output: Optional[torch.Tensor] = None,
- ) -> torch.Tensor:
- return _cached_replay_update(
- q=q,
- k=k,
- v=v,
- g=g,
- beta=beta,
- ssm_states=ssm_states,
- state_indices=state_indices,
- old_u=old_u,
- old_k=old_k,
- old_G=old_G,
- old_beta=old_beta,
- cache_buf_idx=cache_buf_idx,
- prev_num_accepted_tokens=prev_num_accepted_tokens,
- history_size=history_size,
- scale=scale,
- use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
- A_log=A_log,
- dt_bias=dt_bias,
- launch_with_pdl=launch_with_pdl,
- replay_work_items=replay_work_items,
- n_writes=n_writes,
- block_v=block_v,
- num_warps=num_warps,
- use_l2_state_cache=use_l2_state_cache,
- use_l2_streaming_inputs=use_l2_streaming_inputs,
- use_l2_shared_inputs=use_l2_shared_inputs,
- main_num_stages=main_num_stages,
- packed_qkv=packed_qkv,
- use_all_layer_commit=use_all_layer_commit,
- output=output,
- )
-
- def reference(
- self,
- q: torch.Tensor,
- k: torch.Tensor,
- v: torch.Tensor,
- g: torch.Tensor,
- beta: torch.Tensor,
- ssm_states: torch.Tensor,
- state_indices: torch.Tensor,
- old_u: torch.Tensor,
- old_k: torch.Tensor,
- old_G: torch.Tensor,
- old_beta: torch.Tensor,
- cache_buf_idx: torch.Tensor,
- prev_num_accepted_tokens: torch.Tensor,
- history_size: int,
- scale: Optional[float] = None,
- use_qk_l2norm_in_kernel: bool = False,
- A_log: Optional[torch.Tensor] = None,
- dt_bias: Optional[torch.Tensor] = None,
- launch_with_pdl: bool = False,
- replay_work_items: Optional[torch.Tensor] = None,
- n_writes: Optional[torch.Tensor] = None,
- block_v: Optional[int] = None,
- num_warps: Optional[int] = None,
- use_l2_state_cache: Optional[bool] = None,
- use_l2_streaming_inputs: Optional[bool] = None,
- use_l2_shared_inputs: Optional[bool] = None,
- main_num_stages: Optional[int] = None,
- packed_qkv: Optional[torch.Tensor] = None,
- use_all_layer_commit: bool = False,
- output: Optional[torch.Tensor] = None,
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
- """`(out, ssm_states, old_u, old_k, old_G)` as the call leaves them;
- mutates nothing. The state is rebuilt from the checkpoint and history
- in fp32 and the new tokens are stepped one at a time in fp32 -- no
- matmul, so no TF32 -- and every store is rounded once to its buffer's
- dtype. The history is read at whatever dtype and depth its buffers
- have, so the same function driven on fp32 history deeper than any run
- ever fills is the exact token-by-token recurrence.
-
- Rows with a negative slot keep `output`'s prior content, or are NaN
- (unspecified) when no `output` is given. `old_beta`, `launch_with_pdl`,
- `n_writes` and the launch-tuning knobs are ignored; `replay_work_items`
- only decides whether a wrap commits the checkpoint here."""
- N, T, H, K = k.shape
- HV, V = v.shape[2], v.shape[3]
- rep = HV // H
- dev = v.device
- scale = K**-0.5 if scale is None else scale
- if packed_qkv is not None:
- rows = packed_qkv.view(N, T, -1)
- q = rows[..., : H * K].view(N, T, H, K)
- k = rows[..., H * K : 2 * H * K].view(N, T, H, K)
- v = rows[..., 2 * H * K :].view(N, T, HV, V)
- qf, kf, vf = q.float(), k.float(), v.float()
- if use_qk_l2norm_in_kernel:
- qf = qf / torch.sqrt(qf.square().sum(-1, keepdim=True) + 1e-6)
- kf = kf / torch.sqrt(kf.square().sum(-1, keepdim=True) + 1e-6)
- qf = (qf * scale).repeat_interleave(rep, dim=2) # [N, T, HV, K]
- kv = kf.repeat_interleave(rep, dim=2) # [N, T, HV, K]
- if A_log is not None:
- gf = -torch.exp(A_log.float()) * F.softplus(g.float() + dt_bias.float())
- bf = torch.sigmoid(beta.float())
- else:
- gf, bf = g.float(), beta.float()
-
- out_dtype = q.dtype if output is None else output.dtype
- out = (
- torch.full((N, T, HV, V), float("nan"), device=dev)
- if output is None
- else output.float().clone()
- )
- pool, u_buf, k_buf, G_buf = (x.clone() for x in (ssm_states, old_u, old_k, old_G))
- live = (state_indices[:N] >= 0).nonzero().flatten()
- if live.numel() == 0:
- return out.to(out_dtype), pool, u_buf, k_buf, G_buf
- slot = state_indices[:N][live].long()
- p = prev_num_accepted_tokens[slot].long()
- b = cache_buf_idx[slot].long()
- M, depth = live.numel(), old_u.shape[2]
-
- # 1. the pending state: checkpoint plus the p history entries of half b
- pending = torch.arange(depth, device=dev)[None, :] < p[:, None] # [M, depth]
- Gh = torch.where(pending[:, None, :], old_G[slot, b].float(), 0.0) # [M, HV, depth]
- Uh = torch.where(pending[:, :, None, None], old_u[slot, b].float(), 0.0)
- Kh = torch.where(pending[:, :, None, None], old_k[slot, b].float(), 0.0)
- Kh = Kh.repeat_interleave(rep, dim=2) # [M, depth, HV, K]
- last = (p - 1).clamp(min=0)[:, None, None].expand(M, HV, 1)
- g_start = torch.where(p[:, None] > 0, Gh.gather(2, last).squeeze(2), 0.0) # [M, HV]
- weight = torch.where(pending[:, None, :], torch.exp(g_start[:, :, None] - Gh), 0.0)
- h = torch.exp(g_start)[:, :, None, None] * ssm_states[slot].float() # [M, HV, V, K]
- for j in range(depth):
- h = h + weight[:, :, j, None, None] * Uh[:, j, :, :, None] * Kh[:, j, :, None, :]
- checkpoint = h
-
- # 2. the new tokens, in order
- gl, bl = gf[live], bf[live] # [M, T, HV]
- ql, kl, vl = qf[live], kv[live], vf[live]
- G_local = torch.zeros(M, HV, device=dev)
- u_new, G_new = [], []
- for t in range(T):
- h = h * torch.exp(gl[:, t])[:, :, None, None]
- u = bl[:, t, :, None] * (vl[:, t] - (h * kl[:, t, :, None, :]).sum(-1)) # [M, HV, V]
- h = h + u[:, :, :, None] * kl[:, t, :, None, :]
- out[live, t] = (h * ql[:, t, :, None, :]).sum(-1)
- G_local = G_local + gl[:, t]
- u_new.append(u)
- G_new.append(G_local.clone())
-
- # 3. the history append, and the checkpoint commit on a wrap
- wrap = p + T > history_size
- w_buf = torch.where(wrap, 1 - b, b)
- w_off = torch.where(wrap, 0, p)
- for t in range(T):
- pos = w_off + t
- u_buf[slot, w_buf, pos] = u_new[t].to(u_buf.dtype)
- k_buf[slot, w_buf, pos] = kf[live, t].to(k_buf.dtype)
- G_store = torch.where(wrap[:, None], G_new[t], g_start + G_new[t])
- G_buf[slot, w_buf, :, pos] = G_store.to(G_buf.dtype)
- if replay_work_items is None and bool(wrap.any()):
- pool[slot[wrap]] = checkpoint[wrap].to(pool.dtype)
- return out.to(out_dtype), pool, u_buf, k_buf, G_buf
-
- def is_valid(
- self,
- q: torch.Tensor,
- k: torch.Tensor,
- v: torch.Tensor,
- g: torch.Tensor,
- beta: torch.Tensor,
- ssm_states: torch.Tensor,
- state_indices: torch.Tensor,
- old_u: torch.Tensor,
- old_k: torch.Tensor,
- old_G: torch.Tensor,
- old_beta: torch.Tensor,
- cache_buf_idx: torch.Tensor,
- prev_num_accepted_tokens: torch.Tensor,
- history_size: int,
- scale: Optional[float] = None,
- use_qk_l2norm_in_kernel: bool = False,
- A_log: Optional[torch.Tensor] = None,
- dt_bias: Optional[torch.Tensor] = None,
- launch_with_pdl: bool = False,
- replay_work_items: Optional[torch.Tensor] = None,
- n_writes: Optional[torch.Tensor] = None,
- block_v: Optional[int] = None,
- num_warps: Optional[int] = None,
- use_l2_state_cache: Optional[bool] = None,
- use_l2_streaming_inputs: Optional[bool] = None,
- use_l2_shared_inputs: Optional[bool] = None,
- main_num_stages: Optional[int] = None,
- packed_qkv: Optional[torch.Tensor] = None,
- use_all_layer_commit: bool = False,
- output: Optional[torch.Tensor] = None,
- ) -> None:
- N, _, H, K = k.shape
- HV, V = v.shape[2], v.shape[3]
- assert state_indices.dim() == 1 and state_indices.shape[0] >= N, (
- f"state_indices has shape {tuple(state_indices.shape)} for {N} sequences -- the "
- "kernel reads one slot per sequence with no bound, so a shorter tensor hands the "
- "trailing rows whatever follows it in memory, and they write that slot's history"
- )
- pool_slots, slots = ssm_states.shape[0], old_u.shape[0]
- assert slots >= pool_slots, (
- f"the history buffers have {slots} slots for a {pool_slots}-slot pool -- a slot "
- "past the history's end is read and written out of bounds"
- )
- for name, tensor in (
- ("prev_num_accepted_tokens", prev_num_accepted_tokens),
- ("cache_buf_idx", cache_buf_idx),
- ):
- assert tensor.dim() == 1 and tensor.is_contiguous() and tensor.shape[0] >= pool_slots, (
- f"{name} has shape {tuple(tensor.shape)} and stride {tensor.stride()}; the kernel "
- f"reads one element per slot of the {pool_slots}-slot pool, at the slot's index"
- )
- expected = {
- "old_u": (old_u, (slots, 2, history_size, HV, V)),
- "old_k": (old_k, (slots, 2, history_size, H, K)),
- "old_G": (old_G, (slots, 2, HV, history_size)),
- }
- for name, (tensor, shape) in expected.items():
- assert tuple(tensor.shape) == shape, (
- f"{name} has shape {tuple(tensor.shape)}; history_size={history_size} and this "
- f"geometry need {shape} -- the kernel addresses the history by history_size and "
- "the head geometry, never by the buffer's own shape"
- )
-
-
-fused_recurrent_gated_delta_rule_cached_replay_update = (
- FusedRecurrentGatedDeltaRuleCachedReplayUpdate()
-)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_update.py
deleted file mode 100644
index 072f874f6d4e..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_update.py
+++ /dev/null
@@ -1,298 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Gated DeltaNet recurrent core over a slot pool, several tokens per sequence:
-the speculative verify step's recurrence.
-
-Wraps the in-tree module callable
-``tensorrt_llm._torch.modules.fla.fused_recurrent.fused_recurrent_gated_delta_rule_update``
-(tier-2 sourcing: the gated delta rule has no op boundary). One launch of the
-vendored Triton kernel ``fused_recurrent_gated_delta_rule_update_fwd_kernel`` --
-``triton.heuristics`` + ``triton.jit``, no autotuner.
-"""
-
-from __future__ import annotations
-
-from typing import Optional
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.modules.fla.fused_recurrent import (
- fused_recurrent_gated_delta_rule_update as _fused_recurrent_gated_delta_rule_update,
-)
-
-#: This checkpoint's GDN geometry at tp1.
-NUM_K_HEADS, NUM_V_HEADS, HEAD_K_DIM, HEAD_V_DIM = 16, 48, 128, 128
-
-#: Rows of the engine's per-step buffer (`max_num_sequences`), as the
-#: runtime-value dumps record `intermediate_ssm_states`: [32, p, 48, 128, 128].
-ENGINE_VERIFY_ROWS = 32
-#: Default pool and buffer rows of the test's builds.
-SLOTS = 8
-
-
-def _verify(n, p, why, **spec):
- return Cell(why=why, spec=dict(mode="verify", n=n, p=p, **spec))
-
-
-class FusedRecurrentGatedDeltaRuleUpdate(OpWrapper):
- """Advance each sequence's `[K, V]` recurrent state by *all* its tokens, in
- order. Per token and v-head `hv` (q/k head `hv // (HV // H)`):
-
- q, k = q / sqrt(q.q + 1e-6), k / sqrt(k.k + 1e-6) # if use_qk_l2norm_in_kernel
- h *= exp(g); u = (v - h^T k) * beta; h += k u^T; o = h^T (scale * q)
-
- in fp32 throughout. `h` starts at `initial_state_source[idx]` (`[S, HV, V,
- K]`, K innermost), or zero when `idx < 0` or there is no pool. `scale=None`
- is `K ** -0.5`. `beta` with `v`'s rank is read per channel.
-
- Two independent sinks for `h`:
- `disable_state_update=False` stores the final state back into slot `idx`;
- `intermediate_states_buffer` (`[S', cache_steps, HV, V, K]`) receives the
- state after every token at `[idx, step]`. The verify call is the second
- alone: the pool must not move until the sampler has accepted a prefix.
-
- Fusion boundary: L2 norm, scale, the indexed gather and both scatters are
- inside. The gates are *not*: `g = -exp(A_log) * softplus(a + dt_bias)` and
- `beta = sigmoid(b)` arrive evaluated. The conv and q/k/v split are the
- caller's too.
-
- Returns the `[N, P, HV, V]` output (`[1, T, HV, V]` with `cu_seqlens`) in
- `output`'s dtype, else `q`'s; it *is* `output` when one is given. The pool
- and the per-step buffer are written in place and not returned.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: mode ("verify": disable_state_update + buffer with cache_steps = p
- # | "write_back"), n sequences x p tokens, pool_slots, cache_rows, slots
- # (None: arange(n)), pool_dtype / cache_dtype, beta ("fp32" | "bf16":
- # sigmoid of a bf16 projection, as the target passes it), qkv ("views"
- # over a packed [T, 10240] row block | "contiguous", as the target passes
- # them).
- CELLS: tuple[Cell, ...] = tuple(
- _verify(
- n,
- p,
- f"{label}target verify: n={n} p={p}, n-row gather, 32-row buffer",
- pool_slots=n,
- cache_rows=ENGINE_VERIFY_ROWS,
- beta="bf16",
- qkv="contiguous",
- pool_dtype=dtype,
- cache_dtype=dtype,
- )
- for dtype, label, n, p in (
- (torch.float32, "", 32, 2),
- (torch.float32, "", 8, 4),
- (torch.bfloat16, "bf16 state: ", 32, 2),
- )
- )
-
- note = """
- Metadata consumed: none; every piece of state is an argument.
-
- Call site (`_gdn_verify`, the legacy verify path of `TRTLLM_USE_GDN_REPLAY=0`, `p` up
- to 4), all positional: `(q, k, v, g, beta, src, int_idx, None, True, None, True, False,
- int_ssm, p, None)` -- q/k/v the contiguous `[n, p, 16|48, 128]` reshapes of
- `pack_gdn_decode_qkv`'s outputs (so the in-kernel L2 norm is required), `g` fp32 from
- `fused_gdn_gating`, `beta = sigmoid(b)` bf16, `src` an `[n, 48, 128, 128]` gather of
- the live slots and `int_ssm` the layer's `[32, p, 48, 128, 128]` per-step buffer, both
- at the pool's dtype (fp32 or bf16, as `mamba_ssm_cache_dtype` allocated it; the buffer's
- states are rounded once to it), `int_idx = intermediate_state_indices[:n]` (int32
- arange). The buffer's rows are the batch rows `update_mamba_states` reads, not slot
- ids, and its step axis is `tokens_per_gen_step`: `p` (as `cache_steps`) must equal it.
-
- Preconditions not checkable from metadata: `HV % H == 0` (q/k heads are
- repeat_interleaved, not tiled); indices in `[0, S)` and `[0, S')` or
- negative (the device check is `idx >= 0` only; out of range is an
- out-of-bounds access, not probed); distinct indices when writing back --
- measured at two sequences on one slot: outputs exact, the slot a
- per-(v-tile, head) mixture 0.83 .. 0.89 from a sequential application; in the
- varlen form `cache_steps >= max(seq_len)`; q/k unit-norm-conditioned (pass
- `use_qk_l2norm_in_kernel=True` or normalize upstream). `g` is log-space.
-
- Raise loudly, so not guarded: `cu_seqlens` with `q.shape[0] > 1` and an
- index count other than `len(cu_seqlens) - 1` (ValueError); a `None` pool
- with the write-back on (triton CompilationError).
-
- `is_valid` guards what is accepted and silently wrong, each measured:
- a strided `output` (the one tensor not made contiguous) is written with
- packed arithmetic -- 7.04 off and the neighbouring half overwritten;
- `output` with `disable_output_calculation=True` is never written and the
- return is a [1, 1, 1, 1] scratch; a strided pool or per-step buffer is
- silently replaced by a contiguous copy, so the writes land in the copy and
- the caller's tensor comes back bitwise unchanged; `cache_steps` below the
- token count (or None, i.e. 0) overwrites the next row's steps -- 0.850 off
- at cache_steps 2, P 4; `cache_steps` other than the buffer's step axis
- misplaces every row past 0 against the buffer's own indexing.
-
- The first call in a fresh process is bitwise the second (measured on the
- write-back sink; the kernel has no autotuner).
-
- Certified: the call site's form alone -- the per-step buffer sink with the pool
- untouched, at 32 sequences of P = 2 and 8 of P = 4 over an fp32 pool and buffer, and
- 32 of P = 2 over a bf16 one. Not certified: the write-back sink
- (`disable_state_update=False`), alone or with the buffer; the varlen form
- (`cu_seqlens`); negative, scattered or int64 indices; no pool; fp32 or per-channel
- `beta`; q/k/v as views of the packed rows; an `output` buffer;
- `disable_output_calculation=True`; `K` or `V` other than 128; fp16/fp32 q/k/v; P = 3
- and P outside 2..4; N above 32.
-
- Hard limits
- N * HV gridDim.z boundary 65535 probed at N=1300 (62400): runs;
- N=1400 (67200): RuntimeError "invalid argument", context survives
- """
-
- def __call__(
- self,
- q: torch.Tensor,
- k: torch.Tensor,
- v: torch.Tensor,
- g: torch.Tensor,
- beta: torch.Tensor,
- initial_state_source: Optional[torch.Tensor],
- initial_state_indices: torch.Tensor,
- scale: Optional[float] = None,
- use_qk_l2norm_in_kernel: bool = False,
- cu_seqlens: Optional[torch.Tensor] = None,
- disable_state_update: bool = False,
- disable_output_calculation: bool = False,
- intermediate_states_buffer: Optional[torch.Tensor] = None,
- cache_steps: Optional[int] = None,
- output: Optional[torch.Tensor] = None,
- ) -> torch.Tensor:
- return _fused_recurrent_gated_delta_rule_update(
- q=q,
- k=k,
- v=v,
- g=g,
- beta=beta,
- scale=scale,
- initial_state_source=initial_state_source,
- initial_state_indices=initial_state_indices,
- cu_seqlens=cu_seqlens,
- use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
- disable_state_update=disable_state_update,
- disable_output_calculation=disable_output_calculation,
- intermediate_states_buffer=intermediate_states_buffer,
- cache_steps=cache_steps,
- output=output,
- )
-
- def reference(
- self,
- q: torch.Tensor,
- k: torch.Tensor,
- v: torch.Tensor,
- g: torch.Tensor,
- beta: torch.Tensor,
- initial_state_source: Optional[torch.Tensor],
- initial_state_indices: torch.Tensor,
- scale: Optional[float] = None,
- use_qk_l2norm_in_kernel: bool = False,
- cu_seqlens: Optional[torch.Tensor] = None,
- disable_state_update: bool = False,
- disable_output_calculation: bool = False,
- intermediate_states_buffer: Optional[torch.Tensor] = None,
- cache_steps: Optional[int] = None,
- output: Optional[torch.Tensor] = None,
- ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
- """`(out, initial_state_source, intermediate_states_buffer)` as the call
- leaves them; mutates nothing. Token-sequential, fp32 throughout, each
- store rounded once. `disable_output_calculation` and `cache_steps` are
- ignored (the output is always computed; the buffer is indexed by its
- own step axis)."""
- H, K = k.shape[2], k.shape[3]
- HV, V = v.shape[2], v.shape[3]
- rep = HV // H
- scale = K**-0.5 if scale is None else scale
- if cu_seqlens is None:
- spans = [(b, 0, q.shape[1]) for b in range(q.shape[0])]
- else:
- cu = cu_seqlens.tolist()
- spans = [(0, cu[i], cu[i + 1]) for i in range(len(cu) - 1)]
- out = torch.zeros(v.shape, dtype=torch.float32, device=v.device)
- pool = None if initial_state_source is None else initial_state_source.clone()
- cache = None if intermediate_states_buffer is None else intermediate_states_buffer.clone()
- per_channel = beta.dim() == v.dim()
- for i, (b, start, end) in enumerate(spans):
- slot = int(initial_state_indices[i])
- if slot >= 0 and initial_state_source is not None:
- h = initial_state_source[slot].float().transpose(1, 2) # [HV, K, V]
- else:
- h = torch.zeros(HV, K, V, dtype=torch.float32, device=v.device)
- for step, t in enumerate(range(start, end)):
- q_t = q[b, t].float().repeat_interleave(rep, dim=0) # [HV, K]
- k_t = k[b, t].float().repeat_interleave(rep, dim=0)
- if use_qk_l2norm_in_kernel:
- q_t = q_t / torch.sqrt(q_t.square().sum(-1, keepdim=True) + 1e-6)
- k_t = k_t / torch.sqrt(k_t.square().sum(-1, keepdim=True) + 1e-6)
- q_t = q_t * scale
- h = h * torch.exp(g[b, t].float())[:, None, None]
- u = v[b, t].float() - (h * k_t[:, :, None]).sum(1) # [HV, V]
- u = u * (beta[b, t].float() if per_channel else beta[b, t].float()[:, None])
- h = h + k_t[:, :, None] * u[:, None, :]
- out[b, t] = (h * q_t[:, :, None]).sum(1)
- if cache is not None and slot >= 0:
- cache[slot, step] = h.transpose(1, 2).to(cache.dtype)
- if pool is not None and not disable_state_update and slot >= 0:
- pool[slot] = h.transpose(1, 2).to(pool.dtype)
- return out.to(q.dtype if output is None else output.dtype), pool, cache
-
- def is_valid(
- self,
- q: torch.Tensor,
- k: torch.Tensor,
- v: torch.Tensor,
- g: torch.Tensor,
- beta: torch.Tensor,
- initial_state_source: Optional[torch.Tensor],
- initial_state_indices: torch.Tensor,
- scale: Optional[float] = None,
- use_qk_l2norm_in_kernel: bool = False,
- cu_seqlens: Optional[torch.Tensor] = None,
- disable_state_update: bool = False,
- disable_output_calculation: bool = False,
- intermediate_states_buffer: Optional[torch.Tensor] = None,
- cache_steps: Optional[int] = None,
- output: Optional[torch.Tensor] = None,
- ) -> None:
- assert not (output is not None and disable_output_calculation), (
- "output and disable_output_calculation=True are mutually exclusive: with the "
- "output calculation off the buffer is never written and the return value is a "
- "[1, 1, 1, 1] scratch tensor"
- )
- if output is not None:
- assert output.is_contiguous(), (
- f"output must be contiguous; got strides {output.stride()} for shape "
- f"{tuple(output.shape)} -- the kernel writes it with packed [B, T, HV, V] "
- "arithmetic, misplacing values and clobbering whatever lies between"
- )
- if initial_state_source is not None and not disable_state_update:
- assert initial_state_source.is_contiguous(), (
- f"initial_state_source must be contiguous when disable_state_update is False; "
- f"got strides {initial_state_source.stride()} -- a strided pool is copied "
- "before the call and the final state is written into the copy"
- )
- if intermediate_states_buffer is not None:
- assert intermediate_states_buffer.is_contiguous(), (
- f"intermediate_states_buffer must be contiguous; got strides "
- f"{intermediate_states_buffer.stride()} -- a strided buffer is copied before "
- "the call and every per-step state is written into the copy"
- )
- if cu_seqlens is None:
- assert cache_steps is not None and cache_steps >= q.shape[1], (
- f"cache_steps must be at least the per-sequence token count when an "
- f"intermediate_states_buffer is given; got cache_steps={cache_steps} for "
- f"{q.shape[1]} tokens per sequence -- a smaller stride makes each "
- "sequence's tail steps overwrite the next row's"
- )
- assert cache_steps == intermediate_states_buffer.shape[1], (
- f"cache_steps={cache_steps} must equal intermediate_states_buffer's step axis "
- f"({intermediate_states_buffer.shape[1]}) -- the kernel strides rows by "
- "cache_steps, so any other value misplaces every row past 0"
- )
-
-
-fused_recurrent_gated_delta_rule_update = FusedRecurrentGatedDeltaRuleUpdate()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py
deleted file mode 100644
index 19e5cfdc4d9a..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py
+++ /dev/null
@@ -1,297 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Gated DeltaNet recurrent core with the gating fused in, over a slot pool:
-the GDN decode step.
-
-Wraps the in-tree module callable
-``tensorrt_llm._torch.modules.fla.fused_sigmoid_gating_recurrent.fused_sigmoid_gating_delta_rule_update``
-(tier-2 sourcing: the gated delta rule has no ``torch.ops.trtllm`` surface).
-"""
-
-from __future__ import annotations
-
-from typing import Optional
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.modules.fla.fused_sigmoid_gating_recurrent import (
- fused_sigmoid_gating_delta_rule_update as _fused_sigmoid_gating_delta_rule_update,
-)
-
-#: A recurrent-pool slot count of a full-size engine run. The kernels address a slot by
-#: its index alone, so a pool of another size (a bf16 pool holds more slots in the same
-#: memory) takes the same path. The target passes the layer's whole pool unsliced.
-ENGINE_SLOTS = 461
-
-
-def _scattered(n_seqs: int) -> tuple[int, ...]:
- """`n_seqs` distinct engine-pool slots, both ends first, then a stride-13 walk over the
- interior (13 is coprime to its `ENGINE_SLOTS - 2` slots)."""
- rest = tuple((i * 13 + 4) % (ENGINE_SLOTS - 2) + 1 for i in range(max(n_seqs - 2, 0)))
- return ((0, ENGINE_SLOTS - 1) + rest)[:n_seqs]
-
-
-class FusedSigmoidGatingDeltaRuleUpdate(OpWrapper):
- """Advance the gated delta rule over every sequence, gates computed in-kernel.
-
- Sequence `n` (tokens `cu_seqlens[n]:cu_seqlens[n+1]`, or row `n` without
- `cu_seqlens`), v-head `hv` reading k/q head `hv // (HV // H)`, state `h`
- `[K, V]` loaded from `initial_state_source[initial_state_indices[n]]`
- (stored `[V, K]`, K innermost), per token, all in fp32:
-
- g = -exp(A_log) * softplus_b(a + dt_bias) beta = sigmoid(b)
- q, k = q / sqrt(q.q + 1e-6), k / sqrt(k.k + 1e-6) (only with use_qk_l2norm_in_kernel)
- h *= exp(g); u = (v - h^T k) * beta; h += k u^T; o = h^T (scale * q)
-
- `softplus_b(x) = log(1 + exp(beta_s x)) / beta_s` below `softplus_threshold`
- (in `beta_s x`), `x` above; `scale=None` is `K ** -0.5`, bitwise. The final
- `h` is written back to the same slot. A negative index means no slot: zero
- initial state, nothing written. `initial_state_source=None`: zero state,
- nothing written.
-
- Fusion boundary: both gates, the optional q/k norm, scale, recurrence,
- indexed gather and scatter are inside. The conv, the q/k/v split (plain
- views) and the gated RMS norm on `o` are the caller's.
-
- Returns `o` `[B, T, HV, V]` in `output`'s dtype (it *is* `output` when
- given), else `q`'s. The pool is advanced in place and not returned.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: seq_lens, slots (pool rows), indices (None -> arange), pool_dtype
- # (fp32), ab ("separate": two contiguous [T, 48] | "in_proj": q/k/v/a/b all
- # column views of one [T, 16480] projection, as the target's one input GEMM
- # makes them), layout ("varlen": [1, T, ...] with int64 cu_seqlens |
- # "rows": cu_seqlens None, one token per row, q/k/v [N, 1, heads, 128]
- # views of the packed [N, 10240] block whose token stride is the row
- # stride, and a/b [N, 1, 48] -- the target's bf16 call). In-kernel q/k
- # norm on, no output buffer.
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="target bf16 decode call, rows form N=128, 461-slot pool",
- spec=dict(
- seq_lens=(1,) * 128,
- slots=ENGINE_SLOTS,
- indices=_scattered(128),
- ab="separate",
- layout="rows",
- pool_dtype=torch.bfloat16,
- ),
- ),
- Cell(
- why="bf16 rows form, pad rows (-1) between slot 0 and 460",
- spec=dict(
- seq_lens=(1, 1, 1, 1),
- slots=ENGINE_SLOTS,
- indices=(0, -1, ENGINE_SLOTS - 1, -1),
- ab="separate",
- layout="rows",
- pool_dtype=torch.bfloat16,
- ),
- ),
- Cell(
- why="fp32 pool, in_proj column views, N=128: the block test's reference",
- spec=dict(seq_lens=(1,) * 128, slots=128, ab="in_proj"),
- ),
- )
-
- note = """
- Metadata consumed: none.
-
- Call site (`_gdn_block` decode over a bf16 pool), one token per sequence,
- int32 slots, the layer's whole pool `[slots, 48, 128, 128]`: the rows form,
- `cu_seqlens` None -- q/k/v `[T, 1, heads, 128]` split from the whole
- `[T, 10240]` block reshaped to `[T, 1, 80, 128]`, so the token stride is the
- row stride (`is_valid` refuses other strides), `a`/`b` `[T, 1, 48]`. The
- cells drive it at 128 rows of a 461-slot pool, and with CUDA-graph pad rows
- (-1) at 4 rows. One more cell certifies the packed form the GDN block test
- uses as its fp32-pool reference: `(a_log, a, dt_bias, 1.0, 20.0, q, k, v,
- b, ssm_states, si_all, None, True, qsl_long, None)`, q/k/v `[1, T, heads,
- 128]` and `a`/`b` column views of one `[T, 16480]` input-projection output
- (row stride 16480), int64 `cu_seqlens`, at 128 rows.
-
- Branch: the callable takes a FlashInfer CuTe-DSL kernel
- (`gdn_decode_bf16_state.gated_delta_rule`) when the pool is bf16,
- `cu_seqlens` is given and `T == N` (head dims 128, supported arch);
- everything else runs the vendored Triton kernel, and only the Triton
- branch is certified. The target's rows form (no `cu_seqlens`) and the reference's
- packed form over an fp32 pool both stay on it. The bf16 call takes
- the rows form because FlashInfer is not batch-invariant: at 48 value
- heads it runs one kernel up to N = 10 and another from N = 11, and one
- fixed row's single-step output and state differ between the two, so a
- row's tokens would depend on its batch. The Triton rows are one program
- per (row, head); the test shows a fixed row's output and state bitwise
- equal across batch sizes 1..32. FlashInfer also differs from `reference`
- in reading and writing a negative slot as slot 0 (Triton gives a zero
- state and no write) and in its L2-norm epsilon.
-
- Preconditions not checked by anything:
- - `HV % H == 0`.
- - Pool trailing axes `(V, K)`, K innermost. The callable asserts the
- strides `(V*K, K, 1)`, but at K == V a transposed `[S, HV, K, V]` pool
- has the same strides: accepted, measured 2.15e-2 off (vs 3.8e-6
- correct). The slot stride is free.
- - Indices distinct (a shared slot races on write-back: outputs fine, the
- slot holds one writer's state) and in `[0, S)` or negative
- (out-of-range is an unchecked device access -- `tl.device_assert` is
- compiled out -- not probed).
- - q/k unit-norm conditioned: in-kernel norm on, or normalized by the
- caller. Raw multi-token q/k with the norm off drifted the fp32 state
- 1.7e-2 relative from an fp32 sequential reference: domain, not defect.
- - `output` contiguous, `[B, T, HV, V]`: guarded by `is_valid`. Its dtype
- wins (an fp32 `output` with bf16 q returns fp32).
-
- Deterministic: repeated calls are bitwise identical. Grid
- `(1, ceil(V/32), min(N*HV, 65535))`, grid-striding past the cap (read,
- not measured: N > 1365 at this geometry).
-
- Not certified: the FlashInfer branch, K or V other than 128, fp16/fp32
- operands, an fp16 pool, `softplus_beta`/`softplus_threshold` other than
- 1.0/20.0, `B > 1` with `cu_seqlens`; the rows form over an fp32 pool or above
- 128 rows; CUDA-graph replay against the eager call; the packed form over a
- bf16 pool, with more than one token per sequence, with an `output` buffer,
- with `l2norm` off, without a pool or on a layer view of one.
-
- Hard limits
- none probed.
- """
-
- def __call__(
- self,
- A_log: torch.Tensor,
- a: torch.Tensor,
- dt_bias: torch.Tensor,
- softplus_beta: float,
- softplus_threshold: float,
- q: torch.Tensor,
- k: torch.Tensor,
- v: torch.Tensor,
- b: torch.Tensor,
- initial_state_source: Optional[torch.Tensor],
- initial_state_indices: torch.Tensor,
- scale: Optional[float] = None,
- use_qk_l2norm_in_kernel: bool = False,
- cu_seqlens: Optional[torch.Tensor] = None,
- output: Optional[torch.Tensor] = None,
- ) -> torch.Tensor:
- return _fused_sigmoid_gating_delta_rule_update(
- A_log=A_log,
- a=a,
- dt_bias=dt_bias,
- softplus_beta=softplus_beta,
- softplus_threshold=softplus_threshold,
- q=q,
- k=k,
- v=v,
- b=b,
- initial_state_source=initial_state_source,
- initial_state_indices=initial_state_indices,
- scale=scale,
- use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
- cu_seqlens=cu_seqlens,
- output=output,
- )
-
- def reference(
- self,
- A_log: torch.Tensor,
- a: torch.Tensor,
- dt_bias: torch.Tensor,
- softplus_beta: float,
- softplus_threshold: float,
- q: torch.Tensor,
- k: torch.Tensor,
- v: torch.Tensor,
- b: torch.Tensor,
- initial_state_source: Optional[torch.Tensor],
- initial_state_indices: torch.Tensor,
- scale: Optional[float] = None,
- use_qk_l2norm_in_kernel: bool = False,
- cu_seqlens: Optional[torch.Tensor] = None,
- output: Optional[torch.Tensor] = None,
- ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
- """`(o, pool_after)`: sequential over tokens, fp32 throughout like the
- kernel, `o` rounded once to the result dtype and the state once to the
- pool's. `output` contributes only its dtype; nothing is mutated."""
- batch, total, num_k_heads, k_dim = q.shape
- hv, v_dim = v.shape[2], v.shape[3]
- repeat = hv // num_k_heads
- scale = k_dim**-0.5 if scale is None else scale
- if cu_seqlens is None:
- spans = [(n, 0, total) for n in range(batch)]
- else:
- bounds = cu_seqlens.tolist()
- spans = [(0, bounds[n], bounds[n + 1]) for n in range(len(bounds) - 1)]
- a3, b3 = (a, b) if a.dim() == 3 else (a.unsqueeze(0), b.unsqueeze(0))
- out = torch.zeros(batch, total, hv, v_dim, device=q.device, dtype=torch.float32)
- pool = None if initial_state_source is None else initial_state_source.clone()
- neg_exp_a = -torch.exp(A_log.float())
- slots = initial_state_indices.tolist()
- for n, (row, bos, eos) in enumerate(spans):
- slot = slots[n]
- if pool is not None and slot >= 0:
- state = initial_state_source[slot].float().transpose(1, 2)
- else:
- state = torch.zeros(hv, k_dim, v_dim, device=q.device, dtype=torch.float32)
- for t in range(bos, eos):
- q_t = q[row, t].float().repeat_interleave(repeat, dim=0)
- k_t = k[row, t].float().repeat_interleave(repeat, dim=0)
- x = a3[row, t].float() + dt_bias.float()
- bx = softplus_beta * x
- softplus = torch.where(
- bx <= softplus_threshold, torch.log(1.0 + torch.exp(bx)) / softplus_beta, x
- )
- beta = torch.sigmoid(b3[row, t].float())
- if use_qk_l2norm_in_kernel:
- q_t = q_t / torch.sqrt(q_t.pow(2).sum(-1, keepdim=True) + 1e-6)
- k_t = k_t / torch.sqrt(k_t.pow(2).sum(-1, keepdim=True) + 1e-6)
- state = state * torch.exp(neg_exp_a * softplus)[:, None, None]
- u = (v[row, t].float() - (state * k_t[:, :, None]).sum(1)) * beta[:, None]
- state = state + k_t[:, :, None] * u[:, None, :]
- out[row, t] = (state * (q_t * scale)[:, :, None]).sum(1)
- if pool is not None and slot >= 0:
- pool[slot] = state.transpose(1, 2).to(pool.dtype)
- return out.to(q.dtype if output is None else output.dtype), pool
-
- def is_valid(
- self,
- A_log: torch.Tensor,
- a: torch.Tensor,
- dt_bias: torch.Tensor,
- softplus_beta: float,
- softplus_threshold: float,
- q: torch.Tensor,
- k: torch.Tensor,
- v: torch.Tensor,
- b: torch.Tensor,
- initial_state_source: Optional[torch.Tensor],
- initial_state_indices: torch.Tensor,
- scale: Optional[float] = None,
- use_qk_l2norm_in_kernel: bool = False,
- cu_seqlens: Optional[torch.Tensor] = None,
- output: Optional[torch.Tensor] = None,
- ) -> None:
- # Measured on the Triton branch: a last-dim slice of a doubly wide
- # buffer and a token-strided view were both accepted, came back 5.3e-2
- # and 4.8e-2 off at an output absmax of 4.1e-2, and overwrote elements
- # outside the view.
- assert output is None or output.is_contiguous(), (
- f"output must be contiguous; got strides {output.stride()} for shape "
- f"{tuple(output.shape)} -- the kernel writes it with packed "
- "[NK, B, T, HV, V] arithmetic, misplacing values and clobbering the gaps"
- )
- if cu_seqlens is None:
- # q/k/v reshaped to [N, 1, heads, 128] from column slices of the packed
- # [N, 10240] block carry stride(1) = heads * 128, not the row stride: they are
- # accepted and read other rows' columns.
- for name, t in (("q", q), ("k", k), ("v", v), ("a", a), ("b", b)):
- assert t.dim() < 3 or t.stride(0) == t.shape[1] * t.stride(1), (
- f"{name} {tuple(t.shape)} has strides {t.stride()}: without "
- "cu_seqlens the kernel finds row n's tokens at n * T * stride(1), "
- "so stride(0) must equal T * stride(1)"
- )
-
-
-fused_sigmoid_gating_delta_rule_update = FusedSigmoidGatingDeltaRuleUpdate()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_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
deleted file mode 100644
index e9179d141741..000000000000
--- 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
+++ /dev/null
@@ -1,509 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""One Gated-DeltaNet decode step per row, and the gated RMS norm on its output, in
-one PDL launch.
-
-Tier 5 (CuTe DSL written for this target; the recurrence is the gated delta rule of
-FlashInfer's `gdn_decode.gated_delta_rule_decode_pretranspose`, reimplemented here). At one
-decode row the step reads and writes each v-head's 64 KB fp32 state once and
-otherwise does a few hundred flops per state row; the rest of its time is latency.
-The state is not produced by anything in the step ahead of this kernel -- the
-pool row is this layer's own and was last written by this kernel one step ago, and
-the slot index is filled before the step's graph is launched -- so it is read
-before `griddepcontrol.wait`, while the input projection ahead of this kernel
-streams its weight.
-
-A v-head is split over a cluster of `CLUSTER` CTAs of `WARPS` warps; each warp owns
-`V / (CLUSTER * WARPS)` value rows of the head's `[V, K]` state and lane `l` the
-`K / 32` key columns from `l * K / 32` of each. Small CTAs fit beside the input
-projection's last wave, so their state reads start before it ends. The gated RMS
-norm (per v-head, 128 wide) needs the head's whole output: every warp writes the
-sum of squares of its rows into each CTA of the cluster through distributed shared
-memory, and every CTA then normalizes and gates its own rows.
-"""
-
-from __future__ import annotations
-
-import cuda.bindings.driver as cuda_drv
-import cutlass
-import cutlass.cute as cute
-import torch
-from cutlass import const_expr
-from cutlass._mlir.dialects import llvm
-from cutlass.cute import experimental as cute_ext
-from cutlass.cute.runtime import from_dlpack
-from cutlass.cutlass_dsl import dsl_user_op
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
-
-from ..norm.rms_norm_gated_token_major import rms_norm_gated_token_major as _gated_norm
-
-#: Decode rows per launch this entry is certified for.
-MAX_ROWS = 4
-#: CTAs per v-head (one cluster) and warps per CTA; their product must be <= 32.
-CLUSTER = 8
-WARPS = 4
-
-
-@dsl_user_op
-def _st_shared_cluster_f32(ptr, val, *, loc=None, ip=None):
- """`st.shared::cluster.f32` to a pointer from `cute.arch.map_dsmem_ptr`."""
- llvm.inline_asm(
- None,
- [
- cutlass.Int32(ptr.toint(loc=loc, ip=ip)).ir_value(loc=loc, ip=ip),
- cutlass.Float32(val).ir_value(loc=loc, ip=ip),
- ],
- "st.shared::cluster.f32 [$0], $1;",
- "r,f",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
-
-
-class _GdnKernel:
- def __init__(self, *, k_heads, v_heads, kdim, vdim, pdl):
- assert vdim % (CLUSTER * WARPS) == 0 and CLUSTER * WARPS <= 32
- assert kdim % 32 == 0 and v_heads % k_heads == 0
- self.k_heads, self.v_heads, self.kdim, self.vdim, self.pdl = (
- k_heads,
- v_heads,
- kdim,
- vdim,
- pdl,
- )
- self.rows_per_cta = vdim // CLUSTER
- self.rows_per_warp = vdim // (CLUSTER * WARPS)
- self.per_lane = kdim // 32
- self.k_off = k_heads * kdim
- self.v_off = 2 * k_heads * kdim
- self.z_off = self.v_off + v_heads * vdim
- self.a_off = self.z_off + v_heads * vdim
- self.b_off = self.a_off + v_heads
-
- @cute.jit
- def __call__(
- self,
- mProj,
- mPool,
- mSlot,
- mALog,
- mDt,
- mW,
- mOut,
- eps: cutlass.Float32,
- stream: cuda_drv.CUstream,
- ):
- rows = cute.size(mProj, mode=[0])
- ld_state = cute.make_copy_atom(
- cute.nvgpu.CopyG2ROp(),
- cutlass.Float32,
- num_bits_per_copy=32 * self.per_lane,
- load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING,
- )
- st_state = cute.make_copy_atom(
- cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=32 * self.per_lane
- )
- ld_act = cute.make_copy_atom(
- cute.nvgpu.CopyG2ROp(),
- cutlass.BFloat16,
- num_bits_per_copy=16 * self.per_lane,
- load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS,
- )
- self.kernel(
- mProj, mPool, mSlot, mALog, mDt, mW, mOut, eps, ld_state, st_state, ld_act
- ).launch(
- grid=[self.v_heads * CLUSTER, rows, 1],
- block=[WARPS * 32, 1, 1],
- cluster=[CLUSTER, 1, 1],
- smem=32 * 4 + 16,
- stream=stream,
- use_pdl=self.pdl,
- # <= 64 registers a thread: the CTAs fit beside the producer's last wave.
- min_blocks_per_mp=max(1, 8 * 4 // WARPS),
- )
-
- @cute.kernel
- def kernel(
- self,
- mProj,
- mPool,
- mSlot,
- mALog,
- mDt,
- mW,
- mOut,
- eps: cutlass.Float32,
- ld_state: cute.CopyAtom,
- st_state: cute.CopyAtom,
- ld_act: cute.CopyAtom,
- ):
- bidx, n, _ = cute.arch.block_idx()
- warp = cute.arch.warp_idx()
- lane = cute.arch.lane_idx()
- rpw: cutlass.Constexpr = self.rows_per_warp
- pl: cutlass.Constexpr = self.per_lane
- kdim: cutlass.Constexpr = self.kdim
- vdim: cutlass.Constexpr = self.vdim
- parts: cutlass.Constexpr = CLUSTER * WARPS
- hv = bidx // CLUSTER
- part = bidx % CLUSTER # the CTA's rank in its cluster: clusters tile the grid's x
- kh = hv // (self.v_heads // self.k_heads)
- v0 = part * self.rows_per_cta + warp * rpw
- k0 = lane * pl
- ss_parts = cutlass.utils.SmemAllocator().allocate_tensor(
- cutlass.Float32, cute.make_layout((32,), stride=(1,)), byte_alignment=16
- )
- if const_expr(CLUSTER > 1):
- # Every CTA of the cluster must have started before another writes into its
- # shared memory; this arrival is waited on just before those writes.
- cute.arch.cluster_arrive_relaxed()
-
- # 1. Static and step-invariant operands, before the wait: the gates' constants,
- # the norm weight of the warp's rows, the slot and the warp's state rows.
- a_log_h = mALog[hv]
- dtb = mDt[hv]
- wn = cute.make_rmem_tensor(cute.make_layout((rpw,), stride=(1,)), cutlass.Float32)
- for r in cutlass.range_constexpr(rpw):
- wn[r] = mW[v0 + r].to(cutlass.Float32)
- slot = mSlot[n]
- # 64-bit: the cache manager coalesces every GDN layer's state of a slot into one
- # page, so a layer's slot stride is ~28M elements and slot * stride overflows int32.
- slot64 = cutlass.Int64(slot)
- # (slot, hv, v, (pl, lane)): lane l's 16-byte vector of a state row.
- pool_vec = cute.logical_divide(mPool, (None, None, None, pl))
- h = cute.make_rmem_tensor(cute.make_layout((rpw, pl), stride=(pl, 1)), cutlass.Float32)
- h.fill(0.0)
- if slot >= 0:
- for r in cutlass.range_constexpr(rpw):
- cute.copy(ld_state, pool_vec[slot64, hv, v0 + r, (None, lane)], h[r, None])
- if const_expr(self.pdl):
- cute.arch.griddepcontrol_launch_dependents()
- cute.arch.griddepcontrol_wait()
-
- # 2. The projection's q, k, v, z, a, b for this row and head.
- # (n, (pl, cols / pl)): 8-byte vectors of the row; every offset below is a multiple of pl.
- proj_vec = cute.logical_divide(mProj, (None, pl))
- raw = cute.make_rmem_tensor(cute.make_layout((2, pl), stride=(pl, 1)), cutlass.BFloat16)
- cute.copy(ld_act, proj_vec[n, (None, (kh * kdim + k0) // pl)], raw[0, None])
- cute.copy(ld_act, proj_vec[n, (None, (self.k_off + kh * kdim + k0) // pl)], raw[1, None])
- qn = cute.make_rmem_tensor(cute.make_layout((pl,), stride=(1,)), cutlass.Float32)
- kn = cute.make_rmem_tensor(cute.make_layout((pl,), stride=(1,)), cutlass.Float32)
- zv = cute.make_rmem_tensor(cute.make_layout((rpw,), stride=(1,)), cutlass.Float32)
- vv = cute.make_rmem_tensor(cute.make_layout((rpw,), stride=(1,)), cutlass.Float32)
- for e in cutlass.range_constexpr(pl):
- qn[e] = raw[0, e].to(cutlass.Float32)
- kn[e] = raw[1, e].to(cutlass.Float32)
- for r in cutlass.range_constexpr(rpw):
- vv[r] = mProj[n, self.v_off + hv * vdim + v0 + r].to(cutlass.Float32)
- zv[r] = mProj[n, self.z_off + hv * vdim + v0 + r].to(cutlass.Float32)
- a = mProj[n, self.a_off + hv].to(cutlass.Float32)
- b = mProj[n, self.b_off + hv].to(cutlass.Float32)
-
- # q, k L2 norm (eps 1e-6 inside the root), the 1/sqrt(K) scale on q, the gates.
- sq = cutlass.Float32(0.0)
- sk = cutlass.Float32(0.0)
- for e in cutlass.range_constexpr(pl):
- sq = qn[e] * qn[e] + sq
- sk = kn[e] * kn[e] + sk
- sq = cute.arch.warp_reduction_sum(sq)
- sk = cute.arch.warp_reduction_sum(sk)
- rq = cute.math.rsqrt(sq + 1e-6) * (1.0 / cute.math.sqrt(cutlass.Float32(kdim)))
- rk = cute.math.rsqrt(sk + 1e-6)
- for e in cutlass.range_constexpr(pl):
- qn[e] = qn[e] * rq
- kn[e] = kn[e] * rk
- x = a + dtb
- sp = x
- if x <= 20.0:
- sp = cute.math.log1p(cute.math.exp(x))
- decay = cute.math.exp(-cute.math.exp(a_log_h) * sp)
- beta = 1.0 / (1.0 + cute.math.exp(-b))
-
- # 3. The delta rule, one value row at a time; the new rows go straight back.
- ov = cute.make_rmem_tensor(cute.make_layout((rpw,), stride=(1,)), cutlass.Float32)
- ssw = cutlass.Float32(0.0)
- for r in cutlass.range_constexpr(rpw):
- hk = cutlass.Float32(0.0)
- for e in cutlass.range_constexpr(pl):
- hk = h[r, e] * kn[e] + hk
- hk = cute.arch.warp_reduction_sum(hk)
- u = beta * (vv[r] - decay * hk)
- o = cutlass.Float32(0.0)
- for e in cutlass.range_constexpr(pl):
- h[r, e] = decay * h[r, e] + u * kn[e]
- o = h[r, e] * qn[e] + o
- o = cute.arch.warp_reduction_sum(o)
- if slot >= 0:
- cute.copy(st_state, h[r, None], pool_vec[slot64, hv, v0 + r, (None, lane)])
- else:
- o = cutlass.Float32(0.0)
- # The recurrence's output is a bf16 tensor in the unfused composition.
- ov[r] = o.to(cutlass.BFloat16).to(cutlass.Float32)
- ssw = ov[r] * ov[r] + ssw
-
- # 4. The head's sum of squares: warp `w` of CTA `c` writes its rows' share into
- # slot `c * WARPS + w` of every CTA of the cluster (lane `i` into CTA `i`).
- if const_expr(CLUSTER > 1):
- cute.arch.cluster_wait()
- if lane < CLUSTER:
- _st_shared_cluster_f32(
- cute.arch.map_dsmem_ptr(ss_parts.iterator + (part * WARPS + warp), lane), ssw
- )
- cute.arch.cluster_arrive()
- cute.arch.cluster_wait()
- else:
- if lane == 0:
- ss_parts[warp] = ssw
- cute.arch.sync_threads()
- ss = cutlass.Float32(0.0)
- if lane < parts:
- ss = ss_parts[lane]
- ss = cute.arch.warp_reduction_sum(ss)
- rstd = cute.math.rsqrt(ss / vdim + eps)
-
- # 5. Gated RMS norm of the warp's rows: lane `r` writes row `r`.
- o_r = ov[0]
- z_r = zv[0]
- w_r = wn[0]
- for r in cutlass.range_constexpr(1, rpw):
- o_r = cutlass.Float32(cutlass.select_(lane == r, ov[r], o_r))
- z_r = cutlass.Float32(cutlass.select_(lane == r, zv[r], z_r))
- w_r = cutlass.Float32(cutlass.select_(lane == r, wn[r], w_r))
- y = o_r * rstd * w_r * (1.0 / (1.0 + cute.math.exp(-z_r)))
- if lane < rpw:
- mOut[n, hv * vdim + v0 + lane] = y.to(cutlass.BFloat16)
-
-
-def _dl(t: torch.Tensor, align: int = 4):
- return from_dlpack(t, assumed_align=align)
-
-
-_COMPILED: dict = {}
-
-
-def _operands(proj, pool, slots, A_log, dt_bias, norm_weight, out, eps, stream):
- return (
- _dl(proj, 8),
- _dl(pool, 16),
- _dl(slots),
- _dl(A_log),
- _dl(dt_bias),
- _dl(norm_weight, 2),
- _dl(out, 2),
- cutlass.Float32(eps),
- stream,
- )
-
-
-def launch(proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, pdl=True):
- """The launch behind `GdnDecodeGatedNormPdl.__call__`; updates `pool` in place and
- returns the gated-normed output `[rows, HV * V]` bf16."""
- rows = proj.shape[0]
- v_heads, vdim = pool.shape[1], pool.shape[2]
- out = torch.empty((rows, v_heads * vdim), dtype=torch.bfloat16, device=proj.device)
- if rows == 0:
- return out
- pdl = bool(pdl and get_env_enable_pdl())
- stream = cuda_drv.CUstream(torch.cuda.current_stream(proj.device).cuda_stream)
- operands = (proj, pool, slots, A_log, dt_bias, norm_weight, out, eps, stream)
- key = (
- rows,
- k_heads,
- v_heads,
- kdim,
- vdim,
- proj.stride(0),
- tuple(pool.shape),
- pool.stride(),
- pdl,
- )
- fn = _COMPILED.get(key)
- if fn is None:
- # Compiled from the call's own operands (a stand-in for the pool would span every
- # layer of the coalesced state cache).
- kern = _GdnKernel(k_heads=k_heads, v_heads=v_heads, kdim=kdim, vdim=vdim, pdl=pdl)
- fn = _COMPILED[key] = cute_ext.compile(kern, *_operands(*operands))
- fn(*_operands(*operands))
- return out
-
-
-def _recurrence_reference(q, k, v, A_log, a, dt_bias, b, pool, slots):
- """`(o, pool_after)` of one gated-delta-rule decode step per row in fp64, rounded once:
- `o` to bf16, the state to the pool's dtype. Row `n` reads and writes `pool[slots[n]]`
- (a negative slot skips the row: zero output, no state access); q and k are
- L2-normalized, the decay is `exp(-exp(A_log) * softplus(a + dt_bias))`, `beta =
- sigmoid(b)`, the scale `K^-0.5`. `q` / `k` `[n, 1, H, K]`, `v` `[n, 1, HV, V]`, `a` /
- `b` `[n, 1, HV]`, `pool` `[slots, HV, V, K]`. Mutates nothing."""
- dt = torch.float64
- batch, _, num_k_heads, k_dim = q.shape
- hv = v.shape[2]
- rep = hv // num_k_heads
- scale = k_dim**-0.5
- reads = slots.tolist()
- pool_after = pool.clone()
- o = torch.zeros(batch, 1, hv, v.shape[3], dtype=dt, device=q.device)
- neg_a = -torch.exp(A_log.to(dt))
- for n in range(batch):
- if reads[n] < 0:
- continue
- q_n = q[n, 0].to(dt).repeat_interleave(rep, dim=0)
- k_n = k[n, 0].to(dt).repeat_interleave(rep, dim=0)
- q_n = q_n * torch.rsqrt(q_n.pow(2).sum(-1, keepdim=True) + 1e-6)
- k_n = k_n * torch.rsqrt(k_n.pow(2).sum(-1, keepdim=True) + 1e-6)
- x = a[n, 0].to(dt) + dt_bias.to(dt)
- softplus = torch.where(x <= 20.0, torch.log1p(torch.exp(x)), x)
- beta = torch.sigmoid(b[n, 0].to(dt))
- h = pool[reads[n]].to(dt) * torch.exp(neg_a * softplus)[:, None, None]
- u = (v[n, 0].to(dt) - torch.einsum("hvk,hk->hv", h, k_n)) * beta[:, None]
- h = h + u[:, :, None] * k_n[:, None, :]
- o[n, 0] = torch.einsum("hvk,hk->hv", h, q_n * scale)
- pool_after[reads[n]] = h.to(pool.dtype)
- return o.to(torch.bfloat16), pool_after
-
-
-class GdnDecodeGatedNormPdl(OpWrapper):
- """A single-token Gated-DeltaNet step per row over the fp32 slot pool, then the
- sigmoid-gated RMS norm of its output per v-head. With `proj` the decode input
- projection's output row `[q (H*K) | k (H*K) | v (HV*V) | z (HV*V) | a (HV) | b (HV)]`
- (after the conv step: q, k, v are the conv's SiLU outputs)::
-
- o, pool = gated_delta_rule_decode_pretranspose(q, k, v, A_log, a, dt_bias, b,
- use_qk_l2norm=True, pool, slots)
- out = rms_norm_gated_token_major(o (bf16), z, norm_weight, eps, "sigmoid")
-
- Returns `out` `[rows, HV * V]` bf16 and updates `pool[slots[n]]` in place; a
- negative slot skips the row's state (no read or write) and its output is zero.
- `o` is rounded to bf16 between the two, as the composition stores it.
-
- Fusion boundary: both gates, the q/k L2 norm, the recurrence, the indexed state
- read and write, the gated RMS norm. The projection and the conv are the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: rows; pool (default a contiguous 461-slot pool | "coalesced": the
- # cache manager's pool, slot stride 36 layers' states, slots 100 and 101);
- # pad (row 1 given slot -1).
- CELLS: tuple[Cell, ...] = (
- Cell(why="qwen3.8 decode, 1 row (con 1), 461-slot pool", spec=dict(rows=1)),
- Cell(why="qwen3.8 decode, 4 rows (MAX_ROWS), scattered slots", spec=dict(rows=4)),
- Cell(
- why="2 rows, coalesced pool (slot stride 36 layers, slot 100), row 1 a pad slot (-1)",
- spec=dict(rows=2, pool="coalesced", pad=True),
- ),
- )
-
- note = """
- Metadata consumed: `slots` (the cache manager's state indices for the step).
-
- Preconditions: bf16 `proj` `[rows, 2*H*K + 2*HV*V + 2*HV]` with unit column stride;
- the pool fp32 `[slots, HV, V, K]` with K innermost (any outer strides); `A_log`,
- `dt_bias` fp32 `[HV]`; `norm_weight` bf16 `[V]`; `V % 32 == 0` (rows a warp:
- `V / 32`), `K % 32 == 0`; int32 `slots` `[rows]`. **The pool rows and `slots` must not be
- written by any kernel of the same step queued ahead of this one**: both are read
- before `griddepcontrol.wait`. That holds for this target's decode steps (each
- layer's pool is written only by its own recurrence, `slots` is filled before the
- step's graph launches) and is why the entry is not for verify or prefill steps.
-
- Launch: `HV * 8` CTAs a row, clusters of 8 (one v-head), 128 threads, <= 64
- registers. Numerics: fp32 inside; `o` rounded to bf16 before the norm, the output
- rounded once. Inside torch's default band for bf16 / fp32 of the fp64 composition.
- Deterministic. Graph-capturable after the first call per geometry (compiled then).
-
- Certified: 1 and 4 rows over a contiguous 461-slot pool, and 2 rows over the cache
- manager's coalesced pool (slot stride 36 layers' states, slot 100 ~2.8G elements past
- the base) with the second row a pad slot (-1). Under graph replay the op reads the
- projection after `griddepcontrol.wait` (1 row). Not certified: 3 rows, `pdl=False`.
- """
-
- def __call__(
- self,
- proj: torch.Tensor,
- pool: torch.Tensor,
- slots: torch.Tensor,
- A_log: torch.Tensor,
- dt_bias: torch.Tensor,
- norm_weight: torch.Tensor,
- eps: float,
- k_heads: int,
- kdim: int,
- pdl: bool = True,
- ) -> torch.Tensor:
- return launch(proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, pdl)
-
- def reference(
- self,
- proj: torch.Tensor,
- pool: torch.Tensor,
- slots: torch.Tensor,
- A_log: torch.Tensor,
- dt_bias: torch.Tensor,
- norm_weight: torch.Tensor,
- eps: float,
- k_heads: int,
- kdim: int,
- pdl: bool = True,
- ) -> tuple[torch.Tensor, torch.Tensor]:
- """`(out, pool_after)`: the fp64 recurrence (`_recurrence_reference`; o rounded to
- bf16, the pool to fp32), then the gated norm entry's fp32 reference."""
- rows = proj.shape[0]
- v_heads, vdim = pool.shape[1], pool.shape[2]
- q, k, v, z, a, b = torch.split(
- proj,
- [k_heads * kdim, k_heads * kdim, v_heads * vdim, v_heads * vdim, v_heads, v_heads],
- 1,
- )
- o, pool_after = _recurrence_reference(
- q.reshape(rows, 1, k_heads, kdim),
- k.reshape(rows, 1, k_heads, kdim),
- v.reshape(rows, 1, v_heads, vdim),
- A_log,
- a.reshape(rows, 1, v_heads),
- dt_bias,
- b.reshape(rows, 1, v_heads),
- pool,
- slots,
- )
- out = _gated_norm.reference(
- o.reshape(-1, vdim), z.reshape(rows, v_heads, vdim), norm_weight, eps, None, "sigmoid"
- )
- return out.reshape(rows, v_heads * vdim), pool_after
-
- def is_valid(
- self,
- proj: torch.Tensor,
- pool: torch.Tensor,
- slots: torch.Tensor,
- A_log: torch.Tensor,
- dt_bias: torch.Tensor,
- norm_weight: torch.Tensor,
- eps: float,
- k_heads: int,
- kdim: int,
- pdl: bool = True,
- ) -> None:
- rows = proj.shape[0]
- v_heads, vdim, kd = pool.shape[1], pool.shape[2], pool.shape[3]
- assert rows <= MAX_ROWS, f"at most {MAX_ROWS} rows"
- assert proj.dtype == torch.bfloat16 and proj.stride(1) == 1, "bf16 proj, unit column stride"
- assert proj.shape[1] == 2 * k_heads * kdim + 2 * v_heads * vdim + 2 * v_heads, "proj width"
- assert pool.dtype == torch.float32 and pool.stride(3) == 1 and kd == kdim, (
- "fp32 pool [S, HV, V, K]"
- )
- assert vdim % (CLUSTER * WARPS) == 0 and kdim % 32 == 0, "V % (CLUSTER * WARPS), K % 32"
- assert v_heads % k_heads == 0, "HV a multiple of H"
- assert slots.dtype == torch.int32 and slots.shape == (rows,), "int32 slots [rows]"
- assert A_log.dtype == torch.float32 and dt_bias.dtype == torch.float32, (
- "fp32 gate constants"
- )
- assert norm_weight.dtype == torch.bfloat16 and norm_weight.shape == (vdim,), (
- "bf16 norm weight [V]"
- )
-
-
-gdn_decode_gated_norm_pdl = GdnDecodeGatedNormPdl()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py
deleted file mode 100644
index 7fb0f8144a6b..000000000000
--- 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
+++ /dev/null
@@ -1,1035 +0,0 @@
-# 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 or bf16 state pool: the decode step at batch sizes past the cluster kernel's
-over an fp32 pool, and at every batch size over a bf16 pool.
-
-Tier 6 (CuTe DSL written for this target; the recurrence is the gated delta rule of
-FlashInfer's `gdn_decode.gated_delta_rule_decode_pretranspose`, reimplemented here). At a
-decode batch the step is the pool's traffic -- each (row, v-head) state, 64 KB fp32 or
-32 KB bf16, read and written once -- so a CTA owns one whole state and streams it: warp w
-the value rows [w * RPW, (w + 1) * RPW), lane l the key columns [4l, 4l + 4) of each, so
-a warp's access is one contiguous state row. A bf16 state is widened to fp32 on the load
-and rounded back (nearest even) on the store. The state is loaded before
-`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.
-
-Two variants of that CTA, the same arithmetic instruction for instruction (bitwise the same
-output and state): the state rows loaded into registers (128 registers a thread, two CTAs
-an SM), or -- from the batch where the register variant needs a second wave -- each warp's
-rows brought into shared memory by one bulk copy (TMA) issued as soon as the slot is known,
-which holds no registers while in flight (64 a thread, three to four CTAs an SM): the pool
-streams at 6.8 TB/s read + write instead of 5.3-5.7.
-"""
-
-from __future__ import annotations
-
-import math
-
-import cutlass
-import cutlass.cute as cute
-import torch
-from cutlass._mlir.dialects import llvm
-from cutlass.cute.runtime import from_dlpack
-from cutlass.cutlass_dsl import T, dsl_user_op
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
-from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
-
-from .gdn_decode_gated_norm_pdl import _recurrence_reference
-from .gdn_decode_gated_norm_pdl import gdn_decode_gated_norm_pdl as _one_launch
-
-#: Key columns a lane (one 16-byte fp32 or 8-byte bf16 vector of a state row).
-_PL = 4
-#: Warps a CTA; a warp owns `V // _WARPS` value rows of the head's state.
-_WARPS = 8
-#: Largest recurrent-pool slot count the cache manager was seen to allocate for this
-#: checkpoint; the kernel addresses a slot by its index alone, any pool size alike.
-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_state_bf16(addr, pred, *, loc=None, ip=None):
- """Four bf16 of a state row (streaming) as fp32, or zeros where `pred` is 0."""
- r = _asm_n(
- 4,
- T.f32(),
- [cutlass.Int64(addr), cutlass.Int32(pred)],
- "{ .reg .pred p; .reg .b32 w0, w1; setp.ne.b32 p, $5, 0; mov.b32 w0, 0; "
- "mov.b32 w1, 0; @p ld.global.cs.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,r",
- loc,
- ip,
- side_effects=True,
- )
- return [cutlass.Float32(v) for v in r]
-
-
-@dsl_user_op
-def _st_state_bf16(addr, pred, a, b, c, d, *, loc=None, ip=None):
- """Four fp32 rounded to bf16 into 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; .reg .b32 r0, r1; setp.ne.b32 p, $2, 0; "
- "cvt.rn.bf16x2.f32 r0, $4, $3; cvt.rn.bf16x2.f32 r1, $6, $5; "
- "@p st.global.cs.v2.b32 [$1], {r0, r1}; 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,
- )
- )
-
-
-@dsl_user_op
-def _lds_state(saddr, pred, *, loc=None, ip=None):
- """Four fp32 of a staged state row (shared memory), or zeros where `pred` is 0."""
- r = _asm_n(
- 4,
- T.f32(),
- [cutlass.Int32(saddr), 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.f32 {$0, $1, $2, $3}, [$4]; }",
- "=f,=f,=f,=f,r,r",
- loc,
- ip,
- side_effects=True,
- )
- return [cutlass.Float32(v) for v in r]
-
-
-@dsl_user_op
-def _lds_state_bf16(saddr, pred, *, loc=None, ip=None):
- """Four bf16 of a staged state row (shared memory) as fp32, or zeros where `pred` is 0."""
- r = _asm_n(
- 4,
- T.f32(),
- [cutlass.Int32(saddr), cutlass.Int32(pred)],
- "{ .reg .pred p; .reg .b32 w0, w1; setp.ne.b32 p, $5, 0; mov.b32 w0, 0; "
- "mov.b32 w1, 0; @p ld.shared.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,r,r",
- loc,
- ip,
- side_effects=True,
- )
- return [cutlass.Float32(v) for v in r]
-
-
-@dsl_user_op
-def _fma(a, b, c, *, loc=None, ip=None):
- """`a * b + c` rounded once (pinned: not left to the backend's contraction)."""
- 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):
- """`a * b` rounded (pinned: not contracted into a neighbouring add)."""
- 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 _sts_f32(saddr, x, *, loc=None, ip=None):
- _asm(
- T.i32(),
- [cutlass.Int32(saddr), cutlass.Float32(x)],
- "{ st.shared.f32 [$1], $2; mov.u32 $0, 0; }",
- "=r,r,f",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _lds_f32(saddr, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(
- T.f32(),
- [cutlass.Int32(saddr)],
- "ld.shared.f32 $0, [$1];",
- "=f,r",
- loc,
- ip,
- side_effects=True,
- )
- )
-
-
-@dsl_user_op
-def _mbar_init(saddr, *, loc=None, ip=None):
- """An mbarrier of arrival count 1, made visible to the async proxy."""
- _asm(
- T.i32(),
- [cutlass.Int32(saddr)],
- "{ mbarrier.init.shared::cta.b64 [$1], 1; fence.mbarrier_init.release.cluster; "
- "mov.u32 $0, 0; }",
- "=r,r",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _bulk_g2s(saddr, gaddr, nbytes, mbar, *, loc=None, ip=None):
- """One 1-D bulk copy global -> shared of `nbytes`, completing on the mbarrier.
-
- With `mbarrier.arrive.expect_tx`; the L2 policy is evict-first (the state is read once a
- step).
- """
- _asm(
- T.i32(),
- [cutlass.Int32(saddr), cutlass.Int64(gaddr), cutlass.Int32(nbytes), cutlass.Int32(mbar)],
- "{ .reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol, 1.0; "
- "mbarrier.arrive.expect_tx.shared::cta.b64 _, [$4], $3; "
- "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint "
- "[$1], [$2], $3, [$4], pol; mov.u32 $0, 0; }",
- "=r,r,l,r,r",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _mbar_wait(mbar, pred, *, loc=None, ip=None):
- """Wait for phase 0 of the mbarrier where `pred` is not 0."""
- _asm(
- T.i32(),
- [cutlass.Int32(mbar), cutlass.Int32(pred)],
- "{ .reg .pred P1, P2; setp.ne.b32 P2, $2, 0; @!P2 bra DONE; LAB_WAIT: "
- "mbarrier.try_wait.parity.shared::cta.b64 P1, [$1], 0; @P1 bra DONE; bra LAB_WAIT; "
- "DONE: mov.u32 $0, 0; }",
- "=r,r,r",
- loc,
- ip,
- side_effects=True,
- )
-
-
-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.
-
- `bulk`: each warp's `RPW` state rows (contiguous: `s_row == K`) come into shared memory
- by one 1-D bulk copy (TMA, an mbarrier a warp) issued as soon as the slot is known, and
- both passes over the rows read them from there; the state then holds no registers while
- it is in flight (64 registers a thread, three to four CTAs an SM, against two). Else the
- rows are loaded into registers (any 4-element-aligned layout). The arithmetic is the same
- instruction for instruction: the bulk variant pins the state update to the
- `fma(decay, h, u * k)` the backend contracts the register variant's to."""
-
- def __init__(self, k_heads, v_heads, kdim, vdim, pdl, norm=True, state_bf16=False, bulk=False):
- 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
- # The pool's element size and the lane's 4-column load and store of a state row.
- self.sb = 2 if state_bf16 else 4
- self.ld_state, self.st_state = (
- (_ld_state_bf16, _st_state_bf16) if state_bf16 else (_ld_state, _st_state)
- )
- self.lds_state = _lds_state_bf16 if state_bf16 else _lds_state
- self.bulk = bulk
- # The register cap that lets the bulk variant hold four CTAs an SM (three with fp32).
- self.minb = (4 if state_bf16 else 3) if bulk else 0
- 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
- # Shared memory: the norm's per-warp partials, the mbarriers, the staged rows.
- self.wbytes = self.rpw * kdim * self.sb
- self.o_mb = 64
- self.o_st = 128
- self.smem = self.o_st + (_WARPS * self.wbytes if bulk else 0)
-
- @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=self.smem + 128 if self.bulk else _WARPS * 4 + 64,
- stream=stream,
- use_pdl=self.pdl,
- min_blocks_per_mp=self.minb,
- )
-
- @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
- if cutlass.const_expr(self.bulk):
- sm = cutlass.utils.SmemAllocator().allocate_tensor(
- cutlass.Int8, cute.make_layout((self.smem,), stride=(1,)), byte_alignment=128
- )
- sbase = cutlass.Int32(sm.iterator.toint())
- s_red = sbase
- s_mb = sbase + self.o_mb + warp * 8
- s_st = sbase + self.o_st + warp * self.wbytes
- else:
- 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.
- h = []
- if cutlass.const_expr(self.bulk):
- wbase = (
- pool_base
- + (
- cutlass.Int64(cutlass.max(slot, cutlass.Int32(0))) * s_slot
- + cutlass.Int64(hv * s_head + v0 * s_row)
- )
- * self.sb
- )
- base = wbase + cutlass.Int64(lane * _PL * self.sb)
- if lane == 0:
- _mbar_init(s_mb)
- if slot >= 0:
- _bulk_g2s(s_st, wbase, cutlass.Int32(self.wbytes), s_mb)
- else:
- base = (
- pool_base
- + (
- cutlass.Int64(cutlass.max(slot, cutlass.Int32(0))) * s_slot
- + cutlass.Int64(hv * s_head + v0 * s_row + lane * _PL)
- )
- * self.sb
- )
- h = [
- self.ld_state(base + cutlass.Int64(j * s_row * self.sb), 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.
- if cutlass.const_expr(self.bulk):
- srow = s_st + lane * (_PL * self.sb)
- cute.arch.sync_warp()
- _mbar_wait(s_mb, live)
- part = []
- for j in cutlass.range_constexpr(self.rpw):
- if cutlass.const_expr(self.bulk):
- hj = self.lds_state(srow + j * (self.k * self.sb), live)
- else:
- hj = h[j]
- acc = cutlass.Float32(0.0)
- for e in cutlass.range_constexpr(_PL):
- acc = hj[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)
- if cutlass.const_expr(self.bulk):
- hj = self.lds_state(srow + j * (self.k * self.sb), live)
- else:
- hj = h[j]
- acc = cutlass.Float32(0.0)
- for e in cutlass.range_constexpr(_PL):
- if cutlass.const_expr(self.bulk):
- # the register variant's contraction, pinned (rows from shared memory
- # would let the backend contract the other product)
- hj[e] = _fma(decay, hj[e], _mul(u, kn[e]))
- else:
- hj[e] = u * kn[e] + decay * hj[e]
- acc = hj[e] * qn[e] + acc
- part.append(acc)
- self.st_state(base + cutlass.Int64(j * s_row * self.sb), live, *hj)
- 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 cutlass.const_expr(self.bulk):
- if lane == 0:
- _sts_f32(s_red + warp * 4, ss)
- cute.arch.sync_threads()
- tot = _lds_f32(s_red)
- for i in cutlass.range_constexpr(1, _WARPS):
- tot = tot + _lds_f32(s_red + i * 4)
- else:
- 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 _bulk_eligible(pool, kdim) -> bool:
- """The bulk copy's layout: a warp's rows contiguous, every copy 16-byte aligned."""
- es = pool.element_size()
- return (
- pool.stride(2) == kdim
- and pool.data_ptr() % 16 == 0
- and all(s_ * es % 16 == 0 for s_ in pool.stride()[:2])
- )
-
-
-def _use_bulk(rows, v_heads, device, state_bf16=True) -> bool:
- """Whether a batch takes the bulk-copy variant.
-
- Over a bf16 pool, from the batch the register variant cannot hold in one wave (two of its
- CTAs an SM): below that the register variant's lead (its rows are in registers when the
- wait returns) is worth more than the occupancy. An fp32 pool keeps the register variant:
- its 64 state registers a thread already hold the CTA's 64 KB in flight (6.9 TB/s; the
- bulk copy 6.85).
- """
- sms = torch.cuda.get_device_properties(device).multi_processor_count
- return state_bf16 and rows * v_heads > 2 * sms
-
-
-def _compile(key, k_heads, v_heads, kdim, vdim, pdl, norm, state_bf16, bulk, proj, ptrs):
- gm = cute.AddressSpace.gmem
- cargs = (
- from_dlpack(proj, assumed_align=8).mark_layout_dynamic(leading_dim=1),
- make_ptr(
- cutlass.BFloat16 if state_bf16 else cutlass.Float32,
- ptrs[1],
- gm,
- assumed_align=16 if (bulk or not state_bf16) else 8,
- ),
- 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, state_bf16, bulk),
- *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
- return fn
-
-
-def launch(
- proj,
- pool,
- slots,
- A_log,
- dt_bias,
- norm_weight,
- eps,
- k_heads,
- kdim,
- pdl=True,
- out=None,
- bulk=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.
-
- `bulk`: None picks the variant by pool dtype and batch (`_use_bulk`), True / False force
- it (True falls back to the register variant where the pool's layout does not admit the
- bulk copy). Both variants are compiled together on the first call of a configuration, so
- a batch size seen first under serving does not compile."""
- rows = proj.shape[0]
- v_heads, vdim = pool.shape[1], pool.shape[2]
- norm = norm_weight is not None
- if out is None:
- out = torch.empty((rows, v_heads * vdim), dtype=torch.bfloat16, device=proj.device)
- if rows == 0:
- return out
- pdl = bool(pdl and get_env_enable_pdl())
- state_bf16 = pool.dtype == torch.bfloat16
- eligible = _bulk_eligible(pool, kdim)
- if bulk is None:
- bulk = _use_bulk(rows, v_heads, proj.device, state_bf16)
- bulk = bool(bulk) and eligible
- 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(),
- )
- base_key = (k_heads, v_heads, kdim, vdim, pdl, norm, state_bf16)
- fn = _COMPILED.get(base_key + (bulk,))
- if fn is None:
- for variant in (False, True) if eligible else (False,):
- if base_key + (variant,) not in _COMPILED:
- _compile(base_key + (variant,), *base_key, variant, proj, ptrs)
- fn = _COMPILED[base_key + (bulk,)]
- 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 or bf16 slot pool, then
- the sigmoid-gated RMS norm of its output per v-head. With `proj` the decode input
- projection's output row `[q (H*K) | k (H*K) | v (HV*V) | z (HV*V) | a (HV) | b (HV)]`
- (after the conv step: q, k, v are the conv's SiLU outputs)::
-
- o, pool = gated_delta_rule_decode_pretranspose(q, k, v, A_log, a, dt_bias, b,
- use_qk_l2norm=True, pool, slots)
- out = rms_norm_gated_token_major(o (bf16), z, norm_weight, eps, "sigmoid")
-
- Returns `out` `[rows, HV * V]` bf16 and updates `pool[slots[n]]` in place; a
- negative slot skips the row's state (no read or write) and its output is zero.
- `o` is rounded to bf16 between the two, as the composition stores it.
-
- Fusion boundary: both gates, the q/k L2 norm, the recurrence, the indexed state
- read and write, the gated RMS norm. The projection and the conv are the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: rows; pool (default a contiguous 461-slot pool | "coalesced": the
- # cache manager's pool, slot stride 36 layers' states, slots 100 and 101);
- # pad (row 1 given slot -1).
- CELLS: tuple[Cell, ...] = (
- Cell(why="qwen3.8 decode, 128 rows (con 128), 461-slot pool", spec=dict(rows=128)),
- Cell(why="qwen3.8 decode, 3 rows (past the conv-fused path)", spec=dict(rows=3)),
- Cell(
- why="2 rows, coalesced pool (slot stride 36 layers, slot 100), row 1 a pad slot (-1)",
- spec=dict(rows=2, pool="coalesced", pad=True),
- ),
- Cell(why="bf16 pool (mamba_ssm_cache_dtype), 128 rows", spec=dict(rows=128, state="bf16")),
- Cell(why="bf16 pool, 1 row (con 1)", spec=dict(rows=1, state="bf16")),
- Cell(
- why="bf16 pool, 33 rows, one pad slot (-1)", spec=dict(rows=33, pad=True, state="bf16")
- ),
- Cell(
- why="bf16 pool, 2 rows, coalesced: slot stride 36 layers, slot 100",
- spec=dict(rows=2, pool="coalesced", state="bf16"),
- ),
- )
-
- 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 or bf16 `[slots, HV, V, K]` with K innermost
- (any outer strides, multiples of 4) and aligned to 4 elements; `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); a bf16 pool's state by bulk copy into shared memory from
- `rows * HV > 2 * SMs` (7 rows here; an fp32 pool, or a layout the copy cannot take,
- stays on registers), both variants compiled on the first call. GB300, bf16 pool, a CUDA
- graph of 36 layers each streaming its own pool from HBM, PDL on / off, us a layer
- (register-only kernel -> this): 256 rows 152.4 / 141.6 -> 118.8 / 119.4 (6.8 TB/s read +
- write), 128 rows 77.7 / 72.6 -> 61.1 / 61.7, 32 rows 21.7 / 20.7 -> 17.9 / 18.5, 8 rows
- 7.2 / 7.5 -> 6.5 / 6.9, 1 row 2.6 / 3.5 (the register variant).
-
- 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. A bf16 pool's state is rounded once per step, at the store, and `o`
- comes from the state before that rounding. Inside torch's default band for bf16 and
- the pool's dtype of the composition's float64 reference, not bitwise FlashInfer's
- kernel. Batch-invariant: a row's output and state do not depend on the other rows.
-
- Certified: 128 and 3 rows (the fewest a decode step sends here once the conv-fused
- path stops) over a contiguous 461-slot pool, and 2 rows over the cache manager's
- coalesced pool (slot stride 36 layers' states, slot 100 ~2.8G elements past the
- base) with the second row a pad slot (-1); the same over a bf16 pool. Under graph
- replay the op reads the projection after `griddepcontrol.wait` (16 rows). The bulk-copy
- and register variants bitwise equal (outputs and state) at 1-33 rows over contiguous bf16
- and fp32 pools and the coalesced pool (slot 100 by the bulk copy), pad rows included. Not
- certified: `pdl=False`.
- """
-
- def __call__(
- self,
- proj: torch.Tensor,
- pool: torch.Tensor,
- slots: torch.Tensor,
- A_log: torch.Tensor,
- dt_bias: torch.Tensor,
- norm_weight: torch.Tensor,
- eps: float,
- k_heads: int,
- kdim: int,
- pdl: bool = True,
- ) -> torch.Tensor:
- return launch(proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, pdl)
-
- def reference(
- self,
- proj: torch.Tensor,
- pool: torch.Tensor,
- slots: torch.Tensor,
- A_log: torch.Tensor,
- dt_bias: torch.Tensor,
- norm_weight: torch.Tensor,
- eps: float,
- k_heads: int,
- kdim: int,
- pdl: bool = True,
- ) -> tuple[torch.Tensor, torch.Tensor]:
- """`(out, pool_after)`: `gdn_decode_gated_norm_pdl`'s reference, the same contract --
- the fp64 recurrence (o rounded to bf16, the pool to its dtype), then the gated norm
- entry's fp32 reference."""
- return _one_launch.reference(
- proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, pdl
- )
-
- def is_valid(
- self,
- proj: torch.Tensor,
- pool: torch.Tensor,
- slots: torch.Tensor,
- A_log: torch.Tensor,
- dt_bias: torch.Tensor,
- norm_weight: torch.Tensor,
- eps: float,
- k_heads: int,
- kdim: int,
- pdl: bool = True,
- ) -> None:
- rows = proj.shape[0]
- v_heads, vdim, kd = pool.shape[1], pool.shape[2], pool.shape[3]
- assert proj.dtype == torch.bfloat16 and proj.stride(1) == 1, "bf16 proj, unit column stride"
- assert proj.shape[1] == 2 * k_heads * kdim + 2 * v_heads * vdim + 2 * v_heads, "proj width"
- assert proj.data_ptr() % 8 == 0 and proj.stride(0) % 4 == 0, "8-byte aligned proj rows"
- assert pool.dtype in (torch.float32, torch.bfloat16), "fp32 or bf16 pool"
- assert pool.stride(3) == 1 and kd == kdim, "pool [S, HV, V, K]"
- assert pool.data_ptr() % (_PL * pool.element_size()) == 0, "pool aligned to 4 elements"
- assert all(s % _PL == 0 for s in pool.stride()[:3]), "pool strides multiples of 4"
- 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"),
- ),
- Cell(
- why="bf16 pool: 121 generation rows into a row range of the step output",
- spec=dict(rows=121, out="row_range", state="bf16"),
- ),
- Cell(
- why="bf16 pool: 3 rows, one pad slot (-1)",
- spec=dict(rows=3, pad=True, out="row_range", state="bf16"),
- ),
- )
-
- note = """
- Metadata consumed: `slots`. Preconditions and launch: `GdnDecodeRowsGatedNorm`'s, plus
- a bf16 `out` `[rows, HV * V]` with unit column stride. The pool rows and `slots` are
- read before `griddepcontrol.wait`: in a mixed step nothing queued ahead of the call
- writes the generation rows' slots (the context rows' reset and chunked recurrence touch
- only theirs). Numerics: the recurrence of `GdnDecodeRowsGatedNorm` (the same kernel).
-
- Certified: a mixed step's 121 generation rows, 3 rows with a pad slot (-1), and 2 rows
- over the cache manager's coalesced pool, each into a row range of the step output.
- """
-
- def __call__(
- self,
- proj: torch.Tensor,
- pool: torch.Tensor,
- slots: torch.Tensor,
- A_log: torch.Tensor,
- dt_bias: torch.Tensor,
- k_heads: int,
- kdim: int,
- out: torch.Tensor,
- pdl: bool = True,
- ) -> torch.Tensor:
- return launch(proj, pool, slots, A_log, dt_bias, None, 0.0, k_heads, kdim, pdl, out=out)
-
- def reference(
- self,
- proj: torch.Tensor,
- pool: torch.Tensor,
- slots: torch.Tensor,
- A_log: torch.Tensor,
- dt_bias: torch.Tensor,
- k_heads: int,
- kdim: int,
- out: torch.Tensor,
- pdl: bool = True,
- ) -> tuple[torch.Tensor, torch.Tensor]:
- """`(o, pool_after)`: the fp64 recurrence, o rounded to bf16 (the pool to its dtype)."""
- rows = proj.shape[0]
- v_heads, vdim = pool.shape[1], pool.shape[2]
- q, k, v, _, a, b = torch.split(
- proj,
- [k_heads * kdim, k_heads * kdim, v_heads * vdim, v_heads * vdim, v_heads, v_heads],
- 1,
- )
- o, pool_after = _recurrence_reference(
- q.reshape(rows, 1, k_heads, kdim),
- k.reshape(rows, 1, k_heads, kdim),
- v.reshape(rows, 1, v_heads, vdim),
- A_log,
- a.reshape(rows, 1, v_heads),
- dt_bias,
- b.reshape(rows, 1, v_heads),
- pool,
- slots,
- )
- return o.reshape(rows, v_heads * vdim), pool_after
-
- def is_valid(
- self,
- proj: torch.Tensor,
- pool: torch.Tensor,
- slots: torch.Tensor,
- A_log: torch.Tensor,
- dt_bias: torch.Tensor,
- k_heads: int,
- kdim: int,
- out: torch.Tensor,
- pdl: bool = True,
- ) -> None:
- v_heads, vdim = pool.shape[1], pool.shape[2]
- gdn_decode_rows_gated_norm.is_valid(
- proj,
- pool,
- slots,
- A_log,
- dt_bias,
- torch.empty(vdim, dtype=torch.bfloat16),
- 0.0,
- k_heads,
- kdim,
- )
- assert out.dtype == torch.bfloat16 and out.dim() == 2, "bf16 out [rows, HV * V]"
- assert out.shape == (proj.shape[0], v_heads * vdim) and out.stride(1) == 1, (
- "out [rows, HV * V], unit column stride"
- )
-
-
-gdn_decode_rows_gated_norm = GdnDecodeRowsGatedNorm()
-gdn_decode_rows = GdnDecodeRows()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py
deleted file mode 100644
index 4110469820d5..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py
+++ /dev/null
@@ -1,432 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""The Gated-DeltaNet input projection at decode rows with the causal conv1d step
-as its epilogue, one PDL launch.
-
-Tier 5 (CuTe DSL written for this target). The decode GDN block projects into one
-`[T, q|k|v | z | a | b]` buffer and then runs a separate depthwise conv kernel
-over the q|k|v columns. That conv touches 3 cached values and 4 taps per column:
-far too little work for a launch of its own. Here every CTA computes `OUTS`
-output columns over the whole K, and for a column below `conv_dim` it finishes
-the conv step on the projected value before storing it.
-
-Order of work inside a CTA:
-
-1. its weight rows into registers, the conv taps of its columns, the row's slot
- and the cached window of each of its conv columns -- static, host-filled, or
- last written by this kernel one step ago -- so before `griddepcontrol.wait`
- (PDL), while the kernel ahead of this one is still running;
-2. `griddepcontrol.wait`;
-3. the activation, the products, a warp then a cross-warp reduction;
-4. per column: the projection rounded to bf16 (the value the unfused path
- stores), and for a conv column of a non-pad row the conv in fp32 over
- `[window, projected]`, SiLU, rounded once; the window advances to
- `[window[1:], projected]`.
-
-Element `k` of a row is held by thread `(k // VEC) % THREADS` as its vector
-`(k // VEC) // THREADS`, so each k-tile of loads is contiguous across a warp.
-"""
-
-from __future__ import annotations
-
-import cuda.bindings.driver as cuda_drv
-import cutlass
-import cutlass.cute as cute
-import torch
-from cutlass import const_expr
-from cutlass._mlir.dialects import llvm
-from cutlass.cute import experimental as cute_ext
-from cutlass.cute.runtime import from_dlpack
-from cutlass.cutlass_dsl import dsl_user_op
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import (
- Arch,
- Cell,
- OpWrapper,
- true_fp32_matmul,
-)
-from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
-
-VEC = 8
-#: Before its PDL wait a CTA also has the weight rows of CTAs `bidx + j * L2_DIST`
-#: (j = 1 .. L2_AHEAD) brought into L2: the kernel ahead of this one is latency-bound and
-#: only a few of these CTAs fit beside it, so the later waves' weights would otherwise
-#: start streaming only after it.
-L2_AHEAD = 2
-L2_DIST = 1216
-
-
-@dsl_user_op
-def _bulk_prefetch_l2(src, nbytes, *, loc=None, ip=None):
- llvm.inline_asm(
- None,
- [
- cutlass.Int64(src.toint(loc=loc, ip=ip)).ir_value(loc=loc, ip=ip),
- cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip),
- ],
- "cp.async.bulk.prefetch.L2.global [$0], $1;",
- "l,r",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
-
-
-#: Largest decode row count this entry serves (every row re-reads the weight from registers).
-MAX_ROWS = 2
-#: (threads, outputs per CTA) per K; swept on GB300 at the checkpoint's 16480 x 2560.
-TILING = {2560: (64, 2)}
-WIDTH = 4 # conv taps
-_COMPILE_OPTIONS = "--ptxas-options -maxrregcount=64"
-
-
-class _InProjConvKernel:
- def __init__(self, *, rows, k, threads, outs, conv_dim, pad_slot, pdl):
- assert k % (threads * VEC) == 0 and threads % 32 == 0
- self.rows, self.threads, self.outs = rows, threads, outs
- self.k_tiles = k // (threads * VEC)
- self.warps = threads // 32
- self.conv_dim, self.pad_slot, self.pdl = conv_dim, pad_slot, pdl
-
- @cute.jit
- def __call__(self, mA, mW, mOut, mState, mConvW, mSlot, stream: cuda_drv.CUstream):
- n = cute.size(mW, mode=[0])
- load_a = cute.make_copy_atom(
- cute.nvgpu.CopyG2ROp(),
- cutlass.BFloat16,
- num_bits_per_copy=16 * VEC,
- load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS,
- )
- load_w = cute.make_copy_atom(
- cute.nvgpu.CopyG2ROp(),
- cutlass.BFloat16,
- num_bits_per_copy=16 * VEC,
- load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING,
- )
- self.kernel(mA, mW, mOut, mState, mConvW, mSlot, load_a, load_w).launch(
- grid=[cute.ceil_div(n, self.outs), 1, 1],
- block=[self.threads, 1, 1],
- smem=self.rows * self.outs * self.warps * 4,
- stream=stream,
- use_pdl=self.pdl,
- min_blocks_per_mp=1,
- )
-
- @cute.kernel
- def kernel(
- self,
- mA,
- mW,
- mOut,
- mState,
- mConvW,
- mSlot,
- load_a: cute.CopyAtom,
- load_w: cute.CopyAtom,
- ):
- tidx, _, _ = cute.arch.thread_idx()
- bidx, _, _ = cute.arch.block_idx()
- warp = cute.arch.warp_idx()
- rows: cutlass.Constexpr = self.rows
- outs: cutlass.Constexpr = self.outs
- threads: cutlass.Constexpr = self.threads
- k_tiles: cutlass.Constexpr = self.k_tiles
- warps: cutlass.Constexpr = self.warps
- n0 = bidx * outs
- my_col = n0 + tidx
-
- # (row, (VEC, (threads, k_tiles))): thread t's vectors of a row, k-tile by k-tile.
- w_vec = cute.logical_divide(cute.logical_divide(mW, (None, VEC)), (None, (None, threads)))
- a_vec = cute.logical_divide(cute.logical_divide(mA, (None, VEC)), (None, (None, threads)))
-
- # 1. Static operands, before the wait.
- w_reg = cute.make_rmem_tensor(
- cute.make_layout((outs, k_tiles, VEC), stride=(k_tiles * VEC, VEC, 1)), cutlass.BFloat16
- )
- for j in cutlass.range_constexpr(outs):
- w_row = w_vec[n0 + j, (None, (tidx, None))]
- for kt in cutlass.range_constexpr(k_tiles):
- cute.copy(load_w, w_row[None, kt], w_reg[j, kt, None])
- tap = cute.make_rmem_tensor(cute.make_layout((WIDTH,), stride=(1,)), cutlass.BFloat16)
- win = cute.make_rmem_tensor(
- cute.make_layout((rows, WIDTH - 1), stride=(WIDTH - 1, 1)), cutlass.BFloat16
- )
- if tidx < outs:
- if my_col < self.conv_dim:
- for j4 in cutlass.range_constexpr(WIDTH):
- tap[j4] = mConvW[my_col, j4]
- for r in cutlass.range_constexpr(rows):
- slot = mSlot[r]
- if slot != self.pad_slot:
- for j3 in cutlass.range_constexpr(WIDTH - 1):
- win[r, j3] = mState[slot, my_col, j3]
-
- if const_expr(self.pdl):
- n_rows = cute.size(mW, mode=[0])
- k_len = cute.size(mW, mode=[1])
- if tidx < L2_AHEAD:
- tgt = bidx + (tidx + 1) * L2_DIST
- if (tgt + 1) * outs <= n_rows:
- _bulk_prefetch_l2(mW.iterator + tgt * outs * k_len, outs * k_len * 2)
-
- # 2. The producer's output only after the wait.
- if const_expr(self.pdl):
- cute.arch.griddepcontrol_launch_dependents()
- cute.arch.griddepcontrol_wait()
-
- # 3. Products and reductions.
- acc = cute.make_rmem_tensor(
- cute.make_layout((rows, outs), stride=(outs, 1)), cutlass.Float32
- )
- acc.fill(0.0)
- a_reg = cute.make_rmem_tensor(
- cute.make_layout((k_tiles, VEC), stride=(VEC, 1)), cutlass.BFloat16
- )
- for r in cutlass.range_constexpr(rows):
- a_row = a_vec[r, (None, (tidx, None))]
- for kt in cutlass.range_constexpr(k_tiles):
- cute.copy(load_a, a_row[None, kt], a_reg[kt, None])
- for kt in cutlass.range_constexpr(k_tiles):
- for v in cutlass.range_constexpr(VEC):
- av = a_reg[kt, v].to(cutlass.Float32)
- for j in cutlass.range_constexpr(outs):
- acc[r, j] = acc[r, j] + av * w_reg[j, kt, v].to(cutlass.Float32)
- for r in cutlass.range_constexpr(rows):
- for j in cutlass.range_constexpr(outs):
- acc[r, j] = cute.arch.warp_reduction_sum(acc[r, j])
- smem = cutlass.utils.SmemAllocator()
- part = smem.allocate_tensor(
- cutlass.Float32,
- cute.make_layout((rows, outs, warps), stride=(outs * warps, warps, 1)),
- byte_alignment=16,
- )
- with cute.arch.elect_one():
- for r in cutlass.range_constexpr(rows):
- for j in cutlass.range_constexpr(outs):
- part[r, j, warp] = acc[r, j]
- cute.arch.sync_threads()
-
- # 4. Thread j < outs finishes column n0 + j.
- if tidx < outs:
- for r in cutlass.range_constexpr(rows):
- total = cutlass.Float32(0.0)
- for w in cutlass.range_constexpr(warps):
- total = total + part[r, tidx, w]
- proj = total.to(cutlass.BFloat16)
- res = proj
- if my_col < self.conv_dim:
- slot = mSlot[r]
- if slot != self.pad_slot:
- c = (
- tap[0].to(cutlass.Float32) * win[r, 0].to(cutlass.Float32)
- + tap[1].to(cutlass.Float32) * win[r, 1].to(cutlass.Float32)
- + tap[2].to(cutlass.Float32) * win[r, 2].to(cutlass.Float32)
- + tap[3].to(cutlass.Float32) * proj.to(cutlass.Float32)
- )
- res = (c / (1.0 + cute.math.exp(-c))).to(cutlass.BFloat16)
- mState[slot, my_col, 0] = win[r, 1]
- mState[slot, my_col, 1] = win[r, 2]
- mState[slot, my_col, 2] = proj
- mOut[r, my_col] = res
-
-
-def _dl(t: torch.Tensor, align: int):
- return from_dlpack(t, assumed_align=align)
-
-
-_COMPILED: dict = {}
-
-
-def _operands(x, mat_b, out, conv_state, conv_weight, conv_state_indices, stream):
- return (
- _dl(x, 32),
- _dl(mat_b.t(), 32),
- _dl(out, 32),
- _dl(conv_state, 2),
- _dl(conv_weight, 8),
- _dl(conv_state_indices, 4),
- stream,
- )
-
-
-def launch(x, mat_b, conv_state, conv_weight, conv_state_indices, conv_dim, pad_slot_id, pdl):
- """The launch behind `GdnInProjConv1dPdl.__call__`. Returns the fresh `[M, N]` output."""
- rows, k = x.shape
- n = mat_b.shape[1]
- dev = x.device
- pdl = bool(pdl and get_env_enable_pdl())
- out = torch.empty(rows, n, dtype=x.dtype, device=dev)
- stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream)
- operands = (x, mat_b, out, conv_state, conv_weight, conv_state_indices, stream)
- key = (
- rows,
- n,
- k,
- conv_dim,
- int(pad_slot_id),
- pdl,
- tuple(conv_state.shape),
- conv_state.stride(),
- )
- fn = _COMPILED.get(key)
- if fn is None:
- # Compiled from the call's own operands (a stand-in for the pool would span every
- # layer of the coalesced conv cache).
- threads, outs = TILING[k]
- kern = _InProjConvKernel(
- rows=rows,
- k=k,
- threads=threads,
- outs=outs,
- conv_dim=conv_dim,
- pad_slot=int(pad_slot_id),
- pdl=pdl,
- )
- fn = _COMPILED[key] = cute_ext.compile(
- kern, *_operands(*operands), options=_COMPILE_OPTIONS
- )
- fn(*_operands(*operands))
- return out
-
-
-def _cell(why: str, rows: int, slots) -> Cell:
- return Cell(why=why, spec=dict(rows=rows, slots=slots, pdl=True, pool=461))
-
-
-class GdnInProjConv1dPdl(OpWrapper):
- """`out = x @ mat_b`, then for `c < conv_dim` and a row whose slot is not
- `pad_slot_id` the causal conv step on the projected value, one launch.
-
- `x` is the dense row-major `[M, K]` activation (`M <= MAX_ROWS`), `mat_b` the
- dense column-major `[K, N]` view of the stored `[N, K]` projection (its
- `.t()`). Row `r`, slot `s = conv_state_indices[r]`, column `c < conv_dim`:
- `p = round_bf16(sum_k x[r, k] * W[c, k])` (fp32 products and sums),
- `out[r, c] = round_bf16(silu(sum_j conv_weight[c, j] * [conv_state[s, c, :], p][j]))`
- in fp32, and `conv_state[s, c, :] = [conv_state[s, c, 1:], p]`. Every other
- column, and every column of a pad row, is `p`; a pad row's slot is not read or
- written. Returns the fresh `[M, N]` output; `conv_state` is advanced in place.
-
- Equivalent to `low_m_bf16_mm` followed by `causal_conv1d_update(x=out[:, :conv_dim],
- conv_state, conv_weight, bias=None, silu_activation=True, cache_seqlens=None,
- conv_state_indices, pad_slot_id)` up to the GEMV's summation order.
-
- Fusion boundary: the projection, the conv step, SiLU and the window's advance.
- The split into q|k|v, z, a, b and everything after are the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: rows, slots (tuple per row; -1 is the pad slot), pdl, pool (slots in the pool).
- CELLS: tuple[Cell, ...] = (
- _cell("decode 1 row, slot 37 of the 461-slot pool", 1, (37,)),
- _cell("decode 2 rows, two slots", 2, (5, 211)),
- _cell("decode 2 rows, one a pad row (-1)", 2, (17, -1)),
- )
-
- note = """
- Metadata consumed: the decode rows' conv slots (`conv_state_indices`, int32
- `[M]`). Preconditions: bf16 `x` dense `[M, K]`, `1 <= M <= MAX_ROWS`; `mat_b`
- the dense column-major `[K, N]` view of a `[N, K]` weight with K in `TILING`
- and `N` a multiple of the tile's outputs; `conv_state` `[slots, conv_dim, 3]`
- bf16 (any strides; they are compiled in), `conv_weight` `[conv_dim, 4]` bf16
- contiguous, `conv_dim <= N`, no bias, SiLU on, `cache_seqlens` None (the
- window shifts; not a ring).
-
- PDL: the weight rows, the conv taps, the slots and the pool's windows are read
- before `griddepcontrol.wait` (the pool row is this layer's own, written by this
- kernel one step ago; nothing between writes it); the activation after it. The
- launch triggers its dependents right before its own wait.
-
- Compilation: CuTe DSL, JIT on the first call per (M, N, K, conv_dim,
- pad_slot_id, pdl, pool shape and strides); that call must come before
- CUDA-graph capture -- the engine's warmup forward does it.
-
- Numerics: the projection is fp32 products and sums rounded once to bf16 (the
- value stored for non-conv columns and pushed into the window); the conv is
- fp32 over the bf16 window, taps and projection, rounded once. Only the
- GEMV's summation order differs from `low_m_bf16_mm` + `causal_conv1d_update`.
-
- Certified with PDL on, over the 461-slot pool: 1 and 2 rows, and 2 rows with a pad
- row; two consecutive steps under CUDA-graph replay (1 row). Not certified: `pdl=False`.
-
- Hard limits
- M rows per launch boundary MAX_ROWS=2 driven at 1, 2 (cells)
- K a TILING key 2560 the checkpoint's hidden size
- pool any slot count and strides 461 driven (64 under graph replay)
- """
-
- def __call__(
- self,
- x: torch.Tensor,
- mat_b: torch.Tensor,
- conv_state: torch.Tensor,
- conv_weight: torch.Tensor,
- conv_state_indices: torch.Tensor,
- conv_dim: int,
- pad_slot_id: int,
- pdl: bool = True,
- ) -> torch.Tensor:
- return launch(
- x, mat_b, conv_state, conv_weight, conv_state_indices, conv_dim, pad_slot_id, pdl
- )
-
- def reference(
- self,
- x: torch.Tensor,
- mat_b: torch.Tensor,
- conv_state: torch.Tensor,
- conv_weight: torch.Tensor,
- conv_state_indices: torch.Tensor,
- conv_dim: int,
- pad_slot_id: int,
- pdl: bool = True,
- ) -> tuple[torch.Tensor, torch.Tensor]:
- """`(out, conv_state)` as the call leaves them, in fp32 then rounded; mutates nothing."""
- with true_fp32_matmul():
- proj = (x.float() @ mat_b.float()).to(x.dtype)
- out = proj.clone()
- state = conv_state.clone()
- for r, s in enumerate(conv_state_indices.tolist()):
- if s == pad_slot_id:
- continue
- p = proj[r, :conv_dim].float()
- window = torch.cat([conv_state[s].float(), p[:, None]], dim=1)
- c = (window * conv_weight.float()).sum(1)
- out[r, :conv_dim] = (c * torch.sigmoid(c)).to(x.dtype)
- state[s] = window[:, 1:].to(conv_state.dtype)
- return out, state
-
- def is_valid(
- self,
- x: torch.Tensor,
- mat_b: torch.Tensor,
- conv_state: torch.Tensor,
- conv_weight: torch.Tensor,
- conv_state_indices: torch.Tensor,
- conv_dim: int,
- pad_slot_id: int,
- pdl: bool = True,
- ) -> None:
- assert x.dtype == mat_b.dtype == conv_state.dtype == conv_weight.dtype == torch.bfloat16
- assert x.dim() == 2 and x.is_contiguous() and 1 <= x.shape[0] <= MAX_ROWS, tuple(x.shape)
- assert mat_b.shape[0] == x.shape[1] and x.shape[1] in TILING, (
- tuple(mat_b.shape),
- tuple(x.shape),
- )
- assert mat_b.stride(0) == 1 and mat_b.stride(1) == mat_b.shape[0], (
- "mat_b must be dense column-major [K, N] (e.g. weight.t())"
- )
- assert mat_b.shape[1] % TILING[x.shape[1]][1] == 0, (
- "N must be a multiple of the tile's outputs"
- )
- assert conv_state.dim() == 3 and tuple(conv_state.shape[1:]) == (conv_dim, WIDTH - 1)
- assert tuple(conv_weight.shape) == (conv_dim, WIDTH) and conv_weight.is_contiguous()
- assert conv_dim <= mat_b.shape[1]
- assert conv_state_indices.dtype == torch.int32 and tuple(conv_state_indices.shape) == (
- x.shape[0],
- )
-
-
-gdn_in_proj_conv1d_pdl = GdnInProjConv1dPdl()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/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
deleted file mode 100644
index 99ea41dfd83d..000000000000
--- 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
+++ /dev/null
@@ -1,1074 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Gated-DeltaNet prefill preparation in one pass: the causal depthwise conv1d
-of the context rows, read token-major straight off the input projection, with
-the SiLU, the q/k split and L2 norm, the v copy and the decay/beta gates fused
-in, and each sequence's conv window published to the paged pool.
-
-Tier 6 (CuTe DSL written for this target) for the conv pass, Tier 5 (Triton)
-for the gates and the window publish. The three-call
-composition it replaces -- `extract_transpose_prefill_slice` (token-major ->
-channel-major copy), `causal_conv1d_fwd` (channel-major conv, in place) and
-`fused_gdn_post_conv` (split + norm + gates) -- moves the `[T, 10240]`
-activation through HBM three times; no ``torch.ops.trtllm.*`` op, in-tree
-callable or installed FlashInfer API fuses the conv with the norm, and the
-channel-last `causal_conv1d_fwd` kernel alone is slower than this whole pass.
-"""
-
-from __future__ import annotations
-
-import cuda.bindings.driver as cuda
-import cutlass
-import cutlass.cute as cute
-import torch
-import triton
-import triton.language as tl
-from cutlass._mlir.dialects import llvm
-from cutlass.cute.runtime import from_dlpack
-from cutlass.cutlass_dsl import T, dsl_user_op
-from triton.language.extra.cuda import libdevice
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
-
-#: This checkpoint's linear-attention geometry at tp1.
-NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS, HEAD_V_DIM = 16, 128, 48, 128
-QKV_DIM = 2 * NUM_K_HEADS * HEAD_K_DIM + NUM_V_HEADS * HEAD_V_DIM # 10240
-#: The target's one GDN input-projection output: [q|k|v | z | a | b].
-IN_PROJ_WIDTH = QKV_DIM + NUM_V_HEADS * HEAD_V_DIM + 2 * NUM_V_HEADS # 16480
-CONV_WIDTH = 4
-#: Largest conv-pool slot count the cache manager was seen to allocate for this
-#: checkpoint (the `causal_conv1d_fwd` entry's ENGINE_SLOTS).
-ENGINE_SLOTS = 461
-A_LOG_RANGE = (-5.5625, 3.5625)
-DT_BIAS_RANGE = (-8.9375, 19.25)
-SOFTPLUS_THRESHOLD = 20.0
-
-#: Channels per program of the window publish.
-_PUBLISH_BLOCK_C = 512
-#: Tokens per program of the gates: at 16384 tokens on GB300 4.4 us (6.1 at 64, 5.3 at
-#: 16 or 32).
-_GATES_BLOCK_T = 8
-
-
-@triton.jit
-def _gdn_prefill_conv_publish_kernel(
- x_ptr,
- stride_xt,
- state_ptr,
- stride_ss,
- stride_sc,
- stride_sw,
- qsl_ptr,
- slot_ptr,
- init_ptr,
- conv_dim,
- pad_slot_id,
- BLOCK_C: tl.constexpr,
-):
- """One program: one sequence x BLOCK_C channels of its new 3-column window.
-
- New column j holds position `e - 3 + j`: the raw input when it lies inside
- the sequence, else the carried column `L + j` (or 0 without a carry). All
- three are loaded before any is stored, and each channel's three columns
- live in one thread, so a short sequence's shift reads the old window."""
- n = tl.program_id(0)
- c = tl.program_id(1) * BLOCK_C + tl.arange(0, BLOCK_C)
- c_mask = c < conv_dim
- slot = tl.load(slot_ptr + n)
- if slot != pad_slot_id:
- s0 = tl.load(qsl_ptr + n)
- e = tl.load(qsl_ptr + n + 1)
- carry = tl.load(init_ptr + n) != 0
- base = state_ptr + slot.to(tl.int64) * stride_ss + c.to(tl.int64) * stride_sc
- p0 = e - 3
- p1 = e - 2
- p2 = e - 1
- x0 = tl.load(x_ptr + p0.to(tl.int64) * stride_xt + c, mask=c_mask & (p0 >= s0), other=0.0)
- x1 = tl.load(x_ptr + p1.to(tl.int64) * stride_xt + c, mask=c_mask & (p1 >= s0), other=0.0)
- x2 = tl.load(x_ptr + p2.to(tl.int64) * stride_xt + c, mask=c_mask & (p2 >= s0), other=0.0)
- s_0 = tl.load(base + (3 + p0 - s0) * stride_sw, mask=c_mask & (p0 < s0) & carry, other=0.0)
- s_1 = tl.load(base + (3 + p1 - s0) * stride_sw, mask=c_mask & (p1 < s0) & carry, other=0.0)
- s_2 = tl.load(base + (3 + p2 - s0) * stride_sw, mask=c_mask & (p2 < s0) & carry, other=0.0)
- n0 = tl.where(p0 >= s0, x0, s_0)
- n1 = tl.where(p1 >= s0, x1, s_1)
- n2 = tl.where(p2 >= s0, x2, s_2)
- tl.store(base, n0, mask=c_mask)
- tl.store(base + stride_sw, n1, mask=c_mask)
- tl.store(base + 2 * stride_sw, n2, mask=c_mask)
-
-
-# --------------------------------------------------------------------------------------
-# The conv pass in CuTe DSL. A CTA is one 128-channel head x `_TT` tokens, 128 threads:
-# it first copies its `_TT + 3`-row tile (3 history rows) into shared memory with
-# 16-byte `cp.async` (rows outside `[0, T)` zero-filled), so the bytes in flight cost no
-# registers; then a thread (4 channels x `_TPT` consecutive tokens) slides the width-4
-# window over shared-memory rows, every element converted to fp32 once. The conv runs
-# as `fma.rn.f32x2` on channel pairs (first tap from zero, then in tap order: per element
-# an `fma.rn.f32` chain); the SiLU is `acc * rcp(1 + 2^(-acc * log2e))`
-# with `ex2.approx.ftz` and `rcp.approx.ftz`, rounded in pairs by `cvt.rn.bf16x2.f32`.
-# --------------------------------------------------------------------------------------
-
-_VEC = 4
-_TPT = 16
-_RG = 4
-_TT = _TPT * _RG
-_TILE_ROWS = _TT + 3
-_THREADS = 128
-_CHUNKS = _TILE_ROWS * 16
-
-
-def _iv(x, loc, ip):
- return x.ir_value(loc=loc, ip=ip)
-
-
-def _asm1(ret, args, text, constraints, loc, ip, side_effects=False):
- return llvm.inline_asm(
- ret,
- [_iv(a, loc, ip) for a in args],
- text,
- constraints,
- has_side_effects=side_effects,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
-
-
-def _asm2(ty, args, text, constraints, loc, ip, side_effects=False):
- res = _asm1(
- llvm.StructType.get_literal([ty, ty]), args, text, constraints, loc, ip, side_effects
- )
- return (
- llvm.extractvalue(ty, res, [0], loc=loc, ip=ip),
- llvm.extractvalue(ty, res, [1], loc=loc, ip=ip),
- )
-
-
-@dsl_user_op
-def _ffma2(a0, a1, b0, b1, c0, c1, *, loc=None, ip=None):
- """`fma.rn.f32x2`: two fp32 FMAs, each rounded once."""
- r = _asm2(
- T.f32(),
- [cutlass.Float32(v) for v in (a0, a1, b0, b1, c0, c1)],
- "{ .reg .b64 a, b, c, d; mov.b64 a, {$2, $3}; mov.b64 b, {$4, $5}; "
- "mov.b64 c, {$6, $7}; fma.rn.f32x2 d, a, b, c; mov.b64 {$0, $1}, d; }",
- "=f,=f,f,f,f,f,f,f",
- loc,
- ip,
- )
- return cutlass.Float32(r[0]), cutlass.Float32(r[1])
-
-
-@dsl_user_op
-def _fmul2(a0, a1, b0, b1, *, loc=None, ip=None):
- """`mul.rn.f32x2`."""
- r = _asm2(
- T.f32(),
- [cutlass.Float32(v) for v in (a0, a1, b0, b1)],
- "{ .reg .b64 a, b, d; mov.b64 a, {$2, $3}; mov.b64 b, {$4, $5}; "
- "mul.rn.f32x2 d, a, b; mov.b64 {$0, $1}, d; }",
- "=f,=f,f,f,f,f",
- loc,
- ip,
- )
- return cutlass.Float32(r[0]), cutlass.Float32(r[1])
-
-
-@dsl_user_op
-def _silu2(a0, a1, *, loc=None, ip=None):
- """`acc * rcp(1 + 2^(acc * -log2e))` for two values: `mul.rn.f32x2`, `ex2.approx.ftz`,
- `add.rn.f32x2`, `rcp.approx.ftz`, `mul.rn`. For `1 + e` in [1, 2^126] this equals a
- Triton `acc / (1 + exp(-acc))` (`div.full.f32`); past it (acc < -87.3) it gives -0
- where that gives a value below 1e-36."""
- r = _asm2(
- T.f32(),
- [cutlass.Float32(a0), cutlass.Float32(a1)],
- "{ .reg .b64 a, k, t, o, d; .reg .f32 c, t0, t1, d0, d1; "
- "mov.f32 c, 0fBFB8AA3B; mov.b64 k, {c, c}; mov.f32 c, 0f3F800000; mov.b64 o, {c, c}; "
- "mov.b64 a, {$2, $3}; mul.rn.f32x2 t, a, k; mov.b64 {t0, t1}, t; "
- "ex2.approx.ftz.f32 t0, t0; ex2.approx.ftz.f32 t1, t1; mov.b64 t, {t0, t1}; "
- "add.rn.f32x2 d, t, o; mov.b64 {d0, d1}, d; "
- "rcp.approx.ftz.f32 d0, d0; rcp.approx.ftz.f32 d1, d1; "
- "mul.rn.f32 $0, $2, d0; mul.rn.f32 $1, $3, d1; }",
- "=f,=f,f,f",
- loc,
- ip,
- )
- return cutlass.Float32(r[0]), cutlass.Float32(r[1])
-
-
-@dsl_user_op
-def _fma_rn(a, b, c, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm1(
- T.f32(),
- [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)],
- "fma.rn.f32 $0, $1, $2, $3;",
- "=f,f,f,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _inv_norm(ss, eps, *, loc=None, ip=None):
- """1 / sqrt(ss + eps): `add.f32`, `sqrt.approx.ftz.f32`, `div.full.f32` (a Triton
- `1.0 / tl.sqrt(...)`)."""
- return cutlass.Float32(
- _asm1(
- T.f32(),
- [cutlass.Float32(ss), cutlass.Float32(eps)],
- "{ .reg .f32 t; add.f32 t, $1, $2; sqrt.approx.ftz.f32 t, t; "
- "div.full.f32 $0, 0f3F800000, t; }",
- "=f,f,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _bf16x2(lo, hi, *, loc=None, ip=None):
- """`cvt.rn.bf16x2.f32`: `lo` in the low half (the lower address)."""
- return cutlass.Int32(
- _asm1(
- T.i32(),
- [cutlass.Float32(hi), cutlass.Float32(lo)],
- "cvt.rn.bf16x2.f32 $0, $1, $2;",
- "=r,f,f",
- loc,
- ip,
- )
- )
-
-
-@dsl_user_op
-def _unpack2(w, *, loc=None, ip=None):
- """A bf16x2 word as `(low, high)` fp32."""
- r = _asm2(
- T.f32(),
- [cutlass.Int32(w)],
- "{ .reg .b32 l, h; shl.b32 l, $2, 16; and.b32 h, $2, 0xFFFF0000; "
- "mov.b32 $0, l; mov.b32 $1, h; }",
- "=f,=f,r",
- loc,
- ip,
- )
- return cutlass.Float32(r[0]), cutlass.Float32(r[1])
-
-
-@dsl_user_op
-def _st_v2(addr, a, b, *, loc=None, ip=None):
- _asm1(
- T.i32(),
- [cutlass.Int64(addr), cutlass.Int32(a), cutlass.Int32(b)],
- "{ st.global.v2.b32 [$1], {$2, $3}; mov.u32 $0, 0; }",
- "=r,l,r,r",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _cp_async16(smem_addr, gmem_addr, src_bytes, *, loc=None, ip=None):
- """16 bytes global -> shared, the rest of the 16 zero-filled past `src_bytes`."""
- _asm1(
- T.i32(),
- [cutlass.Int32(smem_addr), cutlass.Int64(gmem_addr), cutlass.Int32(src_bytes)],
- "{ cp.async.cg.shared.global [$1], [$2], 16, $3; mov.u32 $0, 0; }",
- "=r,r,l,r",
- loc,
- ip,
- side_effects=True,
- )
-
-
-@dsl_user_op
-def _cp_async_wait_all(*, loc=None, ip=None):
- _asm1(T.i32(), [], "{ cp.async.wait_all; mov.u32 $0, 0; }", "=r", loc, ip, side_effects=True)
-
-
-@dsl_user_op
-def _lds_v2(addr, *, loc=None, ip=None):
- r = _asm2(
- T.i32(),
- [cutlass.Int32(addr)],
- "ld.shared.v2.b32 {$0, $1}, [$2];",
- "=r,=r,r",
- loc,
- ip,
- side_effects=True,
- )
- return cutlass.Int32(r[0]), cutlass.Int32(r[1])
-
-
-class _ConvPassKernel:
- """q / k / v of the conv pass; the gates and the publish are separate launches."""
-
- def __init__(self, num_k_heads: int, num_v_heads: int, head_dim: int, x_stride: int):
- assert head_dim == _VEC * 32
- self.hk, self.hv, self.d = num_k_heads, num_v_heads, head_dim
- self.heads = 2 * num_k_heads + num_v_heads
- self.qkv = self.heads * head_dim
- self.sx = x_stride
-
- @cute.jit
- def __call__(
- self,
- x_ptr: cute.Pointer,
- w_ptr: cute.Pointer,
- st_ptr: cute.Pointer,
- qsl_ptr: cute.Pointer,
- slot_ptr: cute.Pointer,
- init_ptr: cute.Pointer,
- q_t: cute.Tensor,
- k_ptr: cute.Pointer,
- v_ptr: cute.Pointer,
- num_tokens: cutlass.Int32,
- num_seqs: cutlass.Int32,
- search_steps: cutlass.Int32,
- pad_slot_id: cutlass.Int32,
- eps: cutlass.Float32,
- w_s0: cutlass.Int32,
- w_s1: cutlass.Int32,
- st_s0: cutlass.Int32,
- st_s1: cutlass.Int32,
- st_s2: cutlass.Int32,
- stream: cuda.CUstream,
- ):
- weight = cute.make_tensor(w_ptr, cute.make_layout((self.qkv, 4), stride=(w_s0, w_s1)))
- # The pool's slot extent is never read: only `cache_indices` slots are addressed.
- pool = cute.make_tensor(
- st_ptr, cute.make_layout((1 << 20, self.qkv, 3), stride=(st_s0, st_s1, st_s2))
- )
- qsl = cute.make_tensor(qsl_ptr, cute.make_layout((num_seqs + 1,), stride=(1,)))
- slots = cute.make_tensor(slot_ptr, cute.make_layout((num_seqs,), stride=(1,)))
- init = cute.make_tensor(init_ptr, cute.make_layout((num_seqs,), stride=(1,)))
- blocks = (num_tokens + _TT - 1) // _TT
- # `q` arrives as a tensor: the TVM-FFI entry takes its stream from a GPU tensor.
- self.kernel(
- x_ptr.toint(),
- weight,
- pool,
- qsl,
- slots,
- init,
- q_t.iterator.toint(),
- k_ptr.toint(),
- v_ptr.toint(),
- num_tokens,
- num_seqs,
- search_steps,
- pad_slot_id,
- eps,
- ).launch(
- grid=[self.heads, blocks, 1],
- block=[_THREADS, 1, 1],
- smem=_TILE_ROWS * 128 * 2 + 128,
- stream=stream,
- )
-
- @cute.kernel
- def kernel(
- self,
- x_base: cutlass.Int64,
- mW: cute.Tensor,
- mState: cute.Tensor,
- mQsl: cute.Tensor,
- mSlot: cute.Tensor,
- mInit: cute.Tensor,
- q_base: cutlass.Int64,
- k_base: cutlass.Int64,
- v_base: cutlass.Int64,
- num_tokens: cutlass.Int32,
- num_seqs: cutlass.Int32,
- search_steps: cutlass.Int32,
- pad_slot_id: cutlass.Int32,
- eps: cutlass.Float32,
- ):
- tidx, _, _ = cute.arch.thread_idx()
- head, tb, _ = cute.arch.block_idx()
- lane = tidx % 32
- rg = tidx // 32
- t0 = tb * _TT
-
- smem = cutlass.utils.SmemAllocator()
- tile = smem.allocate_tensor(
- cutlass.BFloat16, cute.make_layout((_TILE_ROWS * 128,), stride=(1,)), byte_alignment=16
- )
- tile_base = cutlass.Int32(tile.iterator.toint())
-
- # 1. The tile: rows t0-3 .. t0+TT-1 of this head's 128 channels, zero-filled outside [0, T).
- for it in cutlass.range_constexpr((_CHUNKS + _THREADS - 1) // _THREADS):
- chunk = tidx + it * _THREADS
- if chunk < _CHUNKS:
- r = chunk // 16
- cc = chunk - r * 16
- row = t0 - 3 + r
- valid = cutlass.Int32(0)
- if row >= 0:
- if row < num_tokens:
- valid = cutlass.Int32(16)
- rowc = cutlass.min(cutlass.max(row, cutlass.Int32(0)), num_tokens - 1)
- _cp_async16(
- tile_base + (r * 128 + cc * 8) * 2,
- x_base + (cutlass.Int64(rowc) * self.sx + head * 128 + cc * 8) * 2,
- valid,
- )
-
- # 2. While the copies land: the weights, the first token's sequence, the output row.
- c0 = head * 128 + lane * _VEC
- w = cute.make_rmem_tensor(cute.make_layout((_VEC, 4), stride=(4, 1)), cutlass.Float32)
- for i in cutlass.range_constexpr(_VEC * 4):
- w[i // 4, i % 4] = mW[c0 + i // 4, i % 4].to(cutlass.Float32)
- tf = t0 + rg * _TPT
- lo = cutlass.Int32(0)
- hi = num_seqs
- for _ in cutlass.range(search_steps, unroll=1):
- mid = (lo + hi) // 2
- if mQsl[mid] <= tf:
- lo = mid
- else:
- hi = mid
- s0f = mQsl[lo]
- ef = mQsl[lo + 1]
- # One uniform branch a thread: its tokens and their windows inside one sequence.
- fast = cutlass.Int32(0)
- if tf - s0f >= 3:
- if tf + _TPT <= ef:
- if tf + _TPT <= num_tokens:
- fast = cutlass.Int32(1)
- is_qk = head < 2 * self.hk
- o_base = v_base + cutlass.Int64((head - 2 * self.hk) * 128 + lane * _VEC) * 2
- o_stride = cutlass.Int64(self.hv * 128 * 2)
- if head < self.hk:
- o_base = q_base + cutlass.Int64(head * 128 + lane * _VEC) * 2
- o_stride = cutlass.Int64(self.hk * 128 * 2)
- else:
- if is_qk:
- o_base = k_base + cutlass.Int64((head - self.hk) * 128 + lane * _VEC) * 2
- o_stride = cutlass.Int64(self.hk * 128 * 2)
-
- _cp_async_wait_all()
- cute.arch.sync_threads()
-
- # 3. The window over tile rows rg*TPT .. (tile row r is token t0 - 3 + r).
- win = cute.make_rmem_tensor(cute.make_layout((4, _VEC), stride=(_VEC, 1)), cutlass.Float32)
- for i in cutlass.range_constexpr(3):
- lw, hw = _lds_v2(tile_base + ((rg * _TPT + i) * 128 + lane * _VEC) * 2)
- win[i, 0], win[i, 1] = _unpack2(lw)
- win[i, 2], win[i, 3] = _unpack2(hw)
- acc = cute.make_rmem_tensor(cute.make_layout((_VEC,), stride=(1,)), cutlass.Float32)
- for j in cutlass.range_constexpr(_TPT):
- t = tf + j
- lw, hw = _lds_v2(tile_base + ((rg * _TPT + j + 3) * 128 + lane * _VEC) * 2)
- win[(j + 3) % 4, 0], win[(j + 3) % 4, 1] = _unpack2(lw)
- win[(j + 3) % 4, 2], win[(j + 3) % 4, 3] = _unpack2(hw)
- if fast == 1:
- for p in cutlass.range_constexpr(_VEC // 2):
- a0, a1 = _ffma2(
- win[j % 4, 2 * p],
- win[j % 4, 2 * p + 1],
- w[2 * p, 0],
- w[2 * p + 1, 0],
- cutlass.Float32(0.0),
- cutlass.Float32(0.0),
- )
- for kk in cutlass.range_constexpr(1, 4):
- a0, a1 = _ffma2(
- win[(j + kk) % 4, 2 * p],
- win[(j + kk) % 4, 2 * p + 1],
- w[2 * p, kk],
- w[2 * p + 1, kk],
- a0,
- a1,
- )
- acc[2 * p] = a0
- acc[2 * p + 1] = a1
- else:
- # The token's sequence; a tap before its start reads the carried
- # column W-1-(s0-pos), or 0 without a carry.
- sq = lo
- s0 = s0f
- if t >= ef:
- lo2 = cutlass.Int32(0)
- hi2 = num_seqs
- for _ in cutlass.range(search_steps, unroll=1):
- mid = (lo2 + hi2) // 2
- if mQsl[mid] <= t:
- lo2 = mid
- else:
- hi2 = mid
- sq = lo2
- s0 = mQsl[lo2]
- slot = mSlot[sq]
- init = cutlass.Int32(mInit[sq])
- for v in cutlass.range_constexpr(_VEC):
- a = cutlass.Float32(0.0)
- for kk in cutlass.range_constexpr(4):
- pos = t - 3 + kk
- xv = win[(j + kk) % 4, v]
- if pos < s0:
- xv = cutlass.Float32(0.0)
- if init != 0:
- if slot != pad_slot_id:
- xv = mState[slot, c0 + v, 3 + pos - s0].to(cutlass.Float32)
- a = _fma_rn(xv, w[v, kk], a)
- acc[v] = a
- y0, y1 = _silu2(acc[0], acc[1])
- y2, y3 = _silu2(acc[2], acc[3])
- p01 = _bf16x2(y0, y1)
- p23 = _bf16x2(y2, y3)
- if is_qk:
- # L2 norm of the rounded head row: fp32 squares, a 32-lane butterfly.
- f0, f1 = _unpack2(p01)
- f2, f3 = _unpack2(p23)
- s_a, s_b = _ffma2(f0, f1, f0, f1, cutlass.Float32(0.0), cutlass.Float32(0.0))
- s_a, s_b = _ffma2(f2, f3, f2, f3, s_a, s_b)
- ss = s_a + s_b
- for off in cutlass.range_constexpr(5):
- ss = ss + cute.arch.shuffle_sync_bfly(ss, offset=1 << off)
- inv = _inv_norm(ss, eps)
- n0, n1 = _fmul2(f0, f1, inv, inv)
- n2, n3 = _fmul2(f2, f3, inv, inv)
- p01 = _bf16x2(n0, n1)
- p23 = _bf16x2(n2, n3)
- if t < num_tokens:
- _st_v2(o_base + cutlass.Int64(t) * o_stride, p01, p23)
-
-
-@triton.jit
-def _gdn_gates_kernel(
- a_ptr,
- stride_at,
- stride_ah,
- b_ptr,
- stride_bt,
- stride_bh,
- A_log_ptr,
- dt_bias_ptr,
- g_ptr,
- beta_ptr,
- num_tokens,
- softplus_threshold,
- NUM_V_HEADS: tl.constexpr,
- HP: tl.constexpr,
- BLOCK_T: tl.constexpr,
- LINEAR_GATE: tl.constexpr,
-):
- """The decay/beta gates over `[T, NUM_V_HEADS]`, one row-coalesced tile a program;
- with `LINEAR_GATE` the decay stored as `exp(g)` (libdevice `expf`, as `torch.exp`)."""
- 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
- g = -tl.exp(A_log) * softplus
- if LINEAR_GATE:
- g = libdevice.exp(g)
- tl.store(g_ptr + g_off, g, 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,
- linear_gate=False,
-):
- assert head_k_dim == head_v_dim and weight.shape[-1] == CONV_WIDTH, (
- "one program per 128-channel head: head_k_dim must equal head_v_dim, and the "
- f"conv width must be {CONV_WIDTH}"
- )
- num_tokens = x.shape[0]
- num_seqs = query_start_loc.shape[0] - 1
- conv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim
- dev, dt = x.device, x.dtype
- q = torch.empty((1, num_tokens, num_k_heads, head_k_dim), dtype=dt, device=dev)
- k = torch.empty_like(q)
- v = torch.empty((1, num_tokens, num_v_heads, head_v_dim), dtype=dt, device=dev)
- g = torch.empty((1, num_tokens, num_v_heads), dtype=torch.float32, device=dev)
- beta = torch.empty((1, num_tokens, num_v_heads), dtype=torch.float32, device=dev)
- if num_seqs == 0:
- return q, k, v, g, beta
- if num_tokens > 0:
- assert _cute_conv_pass_ok(x, weight, conv_states, num_k_heads, head_k_dim, num_v_heads), (
- "the conv pass takes bf16 x with 128-channel heads, 16-byte aligned rows and "
- f"int32 weight/pool strides; got {x.dtype}, head_k_dim {head_k_dim}, x at "
- f"{x.data_ptr() % 16} mod 16 with row stride {x.stride(0)}"
- )
- _cute_conv_pass(
- x, weight, conv_states, query_start_loc, cache_indices, has_initial_state, q, k, v,
- num_k_heads, num_v_heads, head_k_dim, pad_slot_id, l2_norm_eps,
- ) # fmt: skip
- _gdn_gates_kernel[(triton.cdiv(num_tokens, _GATES_BLOCK_T),)](
- 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=_GATES_BLOCK_T,
- LINEAR_GATE=linear_gate,
- num_warps=4,
- )
- _gdn_prefill_conv_publish_kernel[(num_seqs, triton.cdiv(conv_dim, _PUBLISH_BLOCK_C))](
- x,
- x.stride(0),
- conv_states,
- *conv_states.stride(),
- query_start_loc,
- cache_indices,
- has_initial_state,
- conv_dim,
- pad_slot_id,
- BLOCK_C=_PUBLISH_BLOCK_C,
- num_warps=4,
- )
- return q, k, v, g, beta
-
-
-def _engine_slots(num_seqs: int) -> tuple[int, ...]:
- """`num_seqs` distinct slots of the engine pool, the last and slot 0 first."""
- rest = tuple((i * 37 + 4) % (ENGINE_SLOTS - 2) + 1 for i in range(max(num_seqs - 2, 0)))
- slots = ((ENGINE_SLOTS - 1, 0) + rest)[:num_seqs]
- assert len(set(slots)) == num_seqs
- return slots
-
-
-class GdnPrefillConvPost(OpWrapper):
- """Causal conv + SiLU of the context rows, then the GDN q/k/v/g/beta.
-
- `x` is `[T, >= qkv_dim]` token-major with unit channel stride (any row
- stride: the target passes the q|k|v column view of its `[T, 16480]` input
- projection), packed `[q | k | v]`, `qkv_dim = 2 * num_k_heads * head_k_dim
- + num_v_heads * head_v_dim`. For sequence `s` over rows `[b, e)` of
- `query_start_loc` and `slot = cache_indices[s]`:
-
- xpad = cat([conv_states[slot] if has_initial_state[s] else 0, x[b:e]])
- y[t] = bf16(silu(sum_k weight[:, k] * xpad[t + k])) (fp32 taps, in order)
- q, k = y / sqrt(sum(y^2) + l2_norm_eps) (per head, fp32, rounded once)
- v = y (bit-exact)
- g = -exp(A_log) * softplus(a + dt_bias) (fp32; exp(g) with linear_gate)
- beta = sigmoid(b) (fp32)
- conv_states[slot] = xpad[e - b : e - b + W - 1] (raw input, exact)
-
- -- the composition `extract_transpose_prefill_slice` -> `causal_conv1d_fwd`
- (SiLU, no bias) -> `fused_gdn_post_conv(prefill, None, ...)`, in one pass
- over the activation plus a per-sequence window publish.
-
- Fusion boundary: the conv-state read, conv, SiLU, split, q/k norm, gates
- and the state publish are inside. The projection, the generation rows'
- conv and the recurrence are the caller's (`use_qk_l2norm_in_kernel=False`
- downstream).
-
- Returns `(q, k, v, g, beta)`: fresh contiguous `[1, T, num_k_heads,
- head_k_dim]` x2, `[1, T, num_v_heads, head_v_dim]` in `x.dtype`, and
- `[1, T, num_v_heads]` fp32 twice. `conv_states` is advanced in place for
- every sequence whose slot is not `pad_slot_id`; nothing else is written.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: lengths; slots (pool rows used, None -> _engine_slots); pool_slots;
- # initial; layout "in_proj" (x and a/b column views of one [T, 16480]
- # buffer, the target's call) | "rows" (x/a/b the first rows of a longer
- # projection, as a mixed step's context half), extra_rows.
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="target call, engine 461-slot pool: 37+1+128 tokens at slots 460, 0, 230",
- spec=dict(
- lengths=(37, 1, 128),
- pool_slots=ENGINE_SLOTS,
- slots=(ENGINE_SLOTS - 1, 0, ENGINE_SLOTS // 2),
- initial=(True, False, True),
- ),
- ),
- Cell(
- why="mixed step's context half: first 8065 rows of a 8192-row projection",
- spec=dict(
- lengths=(8065,),
- pool_slots=ENGINE_SLOTS,
- initial=(False,),
- layout="rows",
- extra_rows=127,
- ),
- ),
- Cell(
- why="target call with linear_gate: 3000+8192+5000 tokens, the middle one fresh",
- spec=dict(
- lengths=(3000, 8192, 5000),
- pool_slots=ENGINE_SLOTS,
- initial=(True, False, True),
- linear_gate=True,
- ),
- ),
- Cell(
- why="128 short contexts (the batch ceiling), some shorter than the window",
- spec=dict(
- lengths=tuple(1 + (i * 7) % 64 for i in range(128)),
- pool_slots=ENGINE_SLOTS,
- initial=tuple(i % 3 != 0 for i in range(128)),
- ),
- ),
- )
-
- note = """
- Metadata consumed: none. Batch layout (`query_start_loc`, int32 `[N + 1]`,
- `query_start_loc[0] == 0`, `query_start_loc[N] == x.shape[0]`), slot map
- (`cache_indices`, int32 `[N]`) and carry flags (`has_initial_state`, bool
- `[N]`) are explicit.
-
- Call sites (`PrefillTarget._gdn_prefill_prep`, from `_gdn_block`): the
- context rows of every prefill-carrying step -- prefill-only `(mixed_qkv,
- a, b, qsl_all, si_all, hi_all)`, mixed and MTP-verify `(mq_p, a_p, b_p,
- qsl_p, si_p, hi_p)` -- with the layer's `[10240, 4]` conv weight, its
- whole conv pool unsliced, `pad_slot_id` -1, `l2_norm_eps` 1e-6 and
- `softplus_threshold` 20. `x`, `a`, `b` are column (and row) views, row
- stride 16480, of the one input-projection output; the outputs feed
- `flashinfer_chunk_gated_delta_rule` with `use_qk_l2norm_in_kernel=False`,
- and `linear_gate=True` on both: the decay leaves this pass as `exp(g)`,
- the factor FlashInfer's kernel takes, bitwise the `torch.exp` the chunk
- op would otherwise launch.
-
- 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 `[8, 48]` tile a
- program (`exp(g)` with `linear_gate`, libdevice `expf`). Then a
- per-sequence window publish (Triton), after both, so no
- program reads a window another has already advanced.
-
- Speed (one layer, 461-slot pool, in_proj views, us, graph replay over 4
- layers' inputs): 16384 tokens / 3 contexts 136 (the three-call chain it
- replaced 528).
-
- Numerics: the conv result is rounded to bf16 before the norm reads it, as
- the unfused path stores it. Conv taps: fp32 FMAs from zero in tap order
- (`fma.rn.f32x2` on channel pairs). SiLU: `acc * rcp(1 + 2^(-acc *
- log2e))` with `ex2.approx.ftz` and `rcp.approx.ftz` -- a Triton
- `acc / (1 + exp(-acc))`'s `ex2.approx` / `div.full` values wherever
- `1 + e <= 2^126`; below acc = -87.3 it returns -0 where those return a
- value under 1e-36. q/k norm: the 128-channel fp32 sum of squares as each
- thread's pairs plus a 32-lane butterfly, `sqrt.approx.ftz`, `div.full`.
- Against the unfused path (measured at the target-call cell): q/k/v within
- one bf16 ulp on 1e-6..1e-4 of elements, g/beta within the fp32 band, conv
- pool bitwise.
-
- Certified over the 461-slot pool, seen as one layer of a two-layer buffer:
- three contexts of 37, 1 and 128 tokens; a mixed step's 8065-token context
- half, the first rows of an 8192-row projection; 128 short contexts (1 to
- 64 tokens, some shorter than the window); and the target's call with
- `linear_gate` (3000 + 8192 + 5000 tokens, the middle one fresh).
-
- A `pad_slot_id` context reads no carry and publishes nothing; its outputs
- are computed from a zero window (the unfused conv leaves those columns
- unconvolved instead -- nothing consumes either). The engine gives every
- context row a real slot; a pad context is not certified.
-
- Preconditions (`is_valid`): `x` 2-D, unit channel stride, at least
- `qkv_dim` wide; `a`/`b` at least `T` rows (read through both strides);
- A_log/dt_bias contiguous with `num_v_heads` entries; `weight` `[qkv_dim,
- 4]` in `x`'s dtype (any strides); `conv_states` `[slots, qkv_dim, 3]` in
- `x`'s dtype (any strides); int32 `query_start_loc`/`cache_indices`, bool
- `has_initial_state`. Loud (`__call__`): `head_k_dim != head_v_dim`, conv
- width other than 4, and what the conv pass cannot take: an `x` not bf16,
- heads other than 128 channels, a row of `x` not 16-byte aligned, a weight
- or pool stride from 2^31 elements.
-
- Not checked: `cache_indices` within the pool and distinct;
- `query_start_loc` non-decreasing and ending at `T`.
-
- Not certified: fp16/fp32 activations, other head geometries, bias, a row
- stride of `x` other than the projection's 16480 (the conv pass compiles one
- kernel per row stride), zero-token contexts, `a + dt_bias` past the softplus
- threshold on every head.
-
- Hard limits
- element offsets 64-bit for the rows of x and the outputs; weight and
- pool strides int32 (refused from 2^31) not probed beyond 8192 tokens.
- """
-
- def __call__(
- self,
- x: torch.Tensor,
- weight: torch.Tensor,
- conv_states: torch.Tensor,
- query_start_loc: torch.Tensor,
- cache_indices: torch.Tensor,
- has_initial_state: torch.Tensor,
- a: torch.Tensor,
- b: torch.Tensor,
- A_log: torch.Tensor,
- dt_bias: torch.Tensor,
- num_k_heads: int,
- head_k_dim: int,
- num_v_heads: int,
- head_v_dim: int,
- pad_slot_id: int = -1,
- l2_norm_eps: float = 1e-6,
- softplus_threshold: float = 20.0,
- linear_gate: bool = False,
- ) -> 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,
- linear_gate,
- )
-
- 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,
- linear_gate: bool = False,
- ) -> 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 (with `linear_gate`, `exp` of that fp32 value in fp64, 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()
- if linear_gate:
- g = torch.exp(g.double()).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,
- linear_gate: bool = False,
- ) -> None:
- # Every address below is computed from the arguments' own extents and
- # strides only where stated; the rest is `ptr + index`, unmasked.
- qkv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim
- n_tok = x.shape[0]
- assert x.dim() == 2 and x.stride(1) == 1 and x.shape[1] >= qkv_dim, (
- f"x must be [T, >= {qkv_dim}] with unit channel stride; got {tuple(x.shape)} "
- f"strides {x.stride()} -- the channel axis is read as `ptr + c`"
- )
- assert weight.dtype == x.dtype and weight.shape[0] == qkv_dim, (
- f"weight must be [{qkv_dim}, 4] in x's dtype; got {tuple(weight.shape)} {weight.dtype}"
- )
- assert conv_states.dtype == x.dtype and tuple(conv_states.shape[1:]) == (qkv_dim, 3), (
- f"conv_states must be [slots, {qkv_dim}, 3] in x's dtype; got "
- f"{tuple(conv_states.shape)} {conv_states.dtype}"
- )
- assert a.shape[0] >= n_tok and b.shape[0] >= n_tok, (
- f"a and b must have at least T={n_tok} rows; the token axis is unmasked"
- )
- assert A_log.stride(0) == 1 and dt_bias.stride(0) == 1, (
- "A_log and dt_bias must be contiguous 1-D; they are read as `ptr + head`"
- )
- assert A_log.numel() >= num_v_heads and dt_bias.numel() >= num_v_heads, (
- f"A_log and dt_bias must hold num_v_heads={num_v_heads} entries; the load is unmasked"
- )
- assert query_start_loc.dtype == torch.int32 and cache_indices.dtype == torch.int32, (
- "query_start_loc and cache_indices must be int32"
- )
- assert has_initial_state.dtype == torch.bool, "has_initial_state must be bool"
- n_seq = query_start_loc.shape[0] - 1
- assert cache_indices.shape[0] >= n_seq and has_initial_state.shape[0] >= n_seq, (
- f"cache_indices and has_initial_state must cover the {n_seq} sequences"
- )
-
-
-gdn_prefill_conv_post = GdnPrefillConvPost()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/pack_gdn_decode_qkv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/pack_gdn_decode_qkv.py
deleted file mode 100644
index 8fe4119dbaaa..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/pack_gdn_decode_qkv.py
+++ /dev/null
@@ -1,135 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Gated-DeltaNet decode-phase q/k/v packing: the plain split of the causal
-conv-update output into per-head q, k and v, with no normalization and no gates.
-
-Wraps the in-tree module callable
-``tensorrt_llm._torch.modules.mamba.fuse_elementwise_ops.pack_gdn_decode_qkv``
-(tier-2 sourcing: no op boundary exists). One Triton launch,
-``_pack_gdn_decode_qkv_kernel``, after allocating the three outputs.
-"""
-
-from __future__ import annotations
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.modules.mamba.fuse_elementwise_ops import (
- pack_gdn_decode_qkv as _pack_gdn_decode_qkv,
-)
-
-#: This checkpoint's linear-attention geometry at tp1.
-NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS, HEAD_V_DIM = 16, 128, 48, 128
-CKPT_HEADS = (NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS, HEAD_V_DIM)
-QKV_DIM = 2 * NUM_K_HEADS * HEAD_K_DIM + NUM_V_HEADS * HEAD_V_DIM # 10240
-
-
-class PackGdnDecodeQkv(OpWrapper):
- """Cut a token-major `[num_tokens, >= qkv_dim]` conv output, packed
- `[q | k | v]` along features with
- `qkv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim`, into
-
- q, k [1, num_tokens, num_k_heads, head_k_dim]
- v [1, num_tokens, num_v_heads, head_v_dim]
-
- Every element is a bit-exact copy: no Q/K L2 norm, no gate, no dtype
- change. Columns at or past `qkv_dim` are dropped.
-
- Fusion boundary: only the split and per-head reshape are inside. The conv
- (and its activation) is the caller's, and so is the Q/K normalization --
- the recurrence consuming these must run `use_qk_l2norm_in_kernel=True`.
- That is the one way to misuse this op and still get plausible output.
-
- Row-wise, with no notion of a sequence: a verify step of `B` sequences x
- `P` tokens is passed as `B * P` rows, the token axis keeps the input row
- order, and `q.reshape(B, P, ...)` is a view, not a re-gather.
-
- Returns three fresh contiguous allocations in `mixed_qkv.dtype`;
- `mixed_qkv` is read through both of its strides and is not modified.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: num_tokens, or (num_decodes, draft) for a verify batch, whose rows
- # are the token-major `[n * p, qkv_dim]` view of the conv's `[n, C, p]` output.
- CELLS: tuple[Cell, ...] = tuple(
- Cell(
- why=f"target verify call: {b} seqs x {p} tokens, token-major rows",
- spec=dict(num_decodes=b, draft=p),
- )
- for b, p in ((32, 2), (8, 4))
- ) + (Cell(why="127 tokens, contiguous: a partial final block", spec=dict(num_tokens=127)),)
-
- note = """
- Metadata consumed: none. Stateless; sequence boundaries and draft grouping
- are carried by the caller's row order alone.
-
- Call site (`_gdn_verify`): `mixed_qkv` the `[n * p, 10240]` token-major
- conv output (contiguous: the conv result carries its input's strides, so
- transposing back is a view), heads `16, 128, 48, 128`, bf16. Certified bitwise
- at that call for 32 sequences of 2 tokens and 8 of 4, and at 127 contiguous
- rows (a partial final block).
-
- Raise loudly, so not guarded: `head_k_dim` and `head_v_dim` rounding up to
- different powers of two (`128 / 64`: triton `CompilationError`
- "Mismatched type for dim_mask", before any launch); a 1-D input
- (`IndexError` reading `stride(1)`).
-
- The feature axis is never masked against the input's own extent; the
- token axis is. `is_valid` guards that: `[32, 9216]` at qkv_dim 10240 was
- accepted and returned finite values (83.3% of `v` bit-equal to the block
- it should have read, the rest up to 6.25 off); over-declaring
- `num_v_heads` likewise; a 3-D `[4, 2, 10240]` input faulted with an
- illegal memory access that poisons the process.
-
- Hard limits
- none probed.
- """
-
- def __call__(
- self,
- mixed_qkv: torch.Tensor,
- num_k_heads: int,
- head_k_dim: int,
- num_v_heads: int,
- head_v_dim: int,
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
- return _pack_gdn_decode_qkv(mixed_qkv, num_k_heads, head_k_dim, num_v_heads, head_v_dim)
-
- def reference(
- self,
- mixed_qkv: torch.Tensor,
- num_k_heads: int,
- head_k_dim: int,
- num_v_heads: int,
- head_v_dim: int,
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
- """`[q | k | v]` cut with `narrow` and reshaped per head; views where torch allows."""
- n = mixed_qkv.shape[0]
- qk_width, v_width = num_k_heads * head_k_dim, num_v_heads * head_v_dim
- q = mixed_qkv.narrow(1, 0, qk_width).reshape(n, num_k_heads, head_k_dim)
- k = mixed_qkv.narrow(1, qk_width, qk_width).reshape(n, num_k_heads, head_k_dim)
- v = mixed_qkv.narrow(1, 2 * qk_width, v_width).reshape(n, num_v_heads, head_v_dim)
- return q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0)
-
- def is_valid(
- self,
- mixed_qkv: torch.Tensor,
- num_k_heads: int,
- head_k_dim: int,
- num_v_heads: int,
- head_v_dim: int,
- ) -> None:
- # `shape[1]` rather than a rank check: a 1-D input still reaches the
- # op's own IndexError, and a 3-D one (whose shape[1] is unrelated) is
- # refused here instead of faulting.
- qkv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim
- assert mixed_qkv.shape[1] >= qkv_dim, (
- f"mixed_qkv must be at least qkv_dim={qkv_dim} wide "
- f"(2 * {num_k_heads} * {head_k_dim} + {num_v_heads} * {head_v_dim}); got "
- f"shape {tuple(mixed_qkv.shape)} -- the feature axis is not masked against "
- "the input's extent, so a short row is read out of bounds"
- )
-
-
-pack_gdn_decode_qkv = PackGdnDecodeQkv()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/reset_recurrent_state_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/reset_recurrent_state_rows.py
deleted file mode 100644
index b61b97d3bb67..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/reset_recurrent_state_rows.py
+++ /dev/null
@@ -1,163 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Zero the recurrent (and optional convolution) pool rows of fresh requests.
-
-Wraps the in-tree module callable
-``tensorrt_llm._torch.modules.mamba.recurrent_state_cache.reset_recurrent_state_rows``
-(tier-2 sourcing: no ``torch.ops.trtllm`` op exists). One Triton launch,
-``_reset_recurrent_state_rows_kernel``; the rest of the callable is shape and
-stride validation.
-"""
-
-from __future__ import annotations
-
-from typing import Optional
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-from tensorrt_llm._torch.modules.mamba import recurrent_state_cache as _recurrent_state_cache
-
-#: This checkpoint's per-layer pool rows: a [48, 128, 128] recurrent state and a
-#: [conv_dim, d_conv - 1] = [10240, 3] convolution window.
-REC_SHAPE = (48, 128, 128)
-CONV_SHAPE = (10240, 3)
-#: A slot count of a full-size engine run, shared by the recurrent and conv pools. The
-#: kernel addresses a slot by its index alone, so a pool of another size (a bf16 recurrent
-#: pool holds more slots in the same memory) takes the same path. The target hands both
-#: pools over unsliced.
-ENGINE_SLOTS = 461
-
-
-def _engine_case(num_seqs: int, flip: bool) -> dict:
- """`num_seqs` distinct slots over the engine pool, slot 0 and 460 live;
- `flip` moves which of the two is the one reset."""
- slots = (0, ENGINE_SLOTS - 1) + tuple(i * 13 + 4 for i in range(num_seqs - 2))
- flags = tuple((i % 2 == 1) != flip for i in range(num_seqs))
- return dict(slots=ENGINE_SLOTS, indices=slots, flags=flags)
-
-
-class ResetRecurrentStateRows(OpWrapper):
- """For each request `i` with `slot = state_indices[i]`: when
- `has_initial_states[i] == 0` and `0 <= slot < recurrent_states.shape[0]`,
- fill row `slot` of `recurrent_states` -- and of `conv_states`, when given --
- with +0.0. Nothing is read from either pool.
-
- Untouched, bitwise: every unnamed slot, every named slot whose flag is
- nonzero (the row a resumed or chunked request keeps), every negative or
- out-of-range slot (skipped, not clamped), padding between rows, the whole
- conv pool when it is `None`, both index tensors. A slot named twice is
- zeroed if any of its entries asks.
-
- Fusion boundary: the call decides nothing. Computing `has_initial_states`,
- assigning slots, and slicing both vectors to the context prefix are the
- caller's.
-
- Returns `(recurrent_states, conv_states)` -- the pools it was handed,
- written in place (`conv_states` is `None` when not passed).
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: slots (pool rows), indices (int32), flags (bool), rec_dtype (fp32),
- # layout: "plain" | "layer_view" (pool[1] of a [3, slots, ...] allocation).
- # Both pools are passed throughout.
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="engine pool 461, 32 requests, slot 460 reset",
- spec=_engine_case(32, flip=True),
- ),
- Cell(
- why="bf16 recurrent pool 461, 32 requests, slot 0 reset",
- spec=dict(_engine_case(32, flip=False), rec_dtype=torch.bfloat16),
- ),
- Cell(
- why="pool[layer] of a [layers, slots, ...] allocation",
- spec=dict(slots=4, indices=(2, 0), flags=(False, True), layout="layer_view"),
- ),
- )
-
- note = """
- Metadata consumed: none. At the call site `state_indices` /
- `has_initial_states` are the context prefix of `Mamba2Metadata`'s int32 /
- bool buffers, and the pools are the layer's views of the cache manager's
- `[layers, slots, ...]` allocations, passed unsliced. Certified at 461 slots
- over an fp32 and a bf16 recurrent pool, and at 4 slots over one layer's view
- of a 3-layer allocation. Not certified: negative or out-of-range slots,
- rows with padding between them, `conv_states=None`.
-
- Loud, from the callable itself (ValueError): an empty recurrent pool, conv
- and recurrent slot counts differing (the in-range test uses the recurrent
- pool's), flag/index shape mismatch, non-contiguous or overlapping rows in
- either pool, any CPU tensor.
-
- Silent, so `is_valid` guards it: both index tensors are read as flat
- contiguous buffers of `state_indices.shape[0]` elements.
-
- Not guarded (data, not metadata): an index dtype too narrow for its slot
- wraps -- an int8 130 arrives as -126 and is skipped as invalid. int32 is
- what the runtime uses.
-
- Launch: grid `(num_requests, cdiv(max(rec_row, conv_row), 256))`, the two
- row widths masked independently; ordered on the current stream.
-
- Hard limits
- slot * stride(0) int64 row base probed at 460 * 786432 = 3.6e8 exact
- within-row offset int32 boundary 2^31 elements per row not probed
- (this checkpoint's largest row is 786432)
- """
-
- def __call__(
- self,
- recurrent_states: torch.Tensor,
- state_indices: torch.Tensor,
- has_initial_states: torch.Tensor,
- conv_states: Optional[torch.Tensor] = None,
- ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
- _recurrent_state_cache.reset_recurrent_state_rows(
- recurrent_states,
- state_indices,
- has_initial_states,
- conv_states,
- )
- return recurrent_states, conv_states
-
- def reference(
- self,
- recurrent_states: torch.Tensor,
- state_indices: torch.Tensor,
- has_initial_states: torch.Tensor,
- conv_states: Optional[torch.Tensor] = None,
- ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
- """Both pools as the call leaves them: clone, zero the rows the rule names."""
- slots = recurrent_states.shape[0]
- rec = recurrent_states.clone()
- conv = None if conv_states is None else conv_states.clone()
- for slot, flag in zip(state_indices.tolist(), has_initial_states.tolist()):
- if flag or not 0 <= slot < slots:
- continue
- rec[slot] = 0
- if conv is not None:
- conv[slot] = 0
- return rec, conv
-
- def is_valid(
- self,
- recurrent_states: torch.Tensor,
- state_indices: torch.Tensor,
- has_initial_states: torch.Tensor,
- conv_states: Optional[torch.Tensor] = None,
- ) -> None:
- # Measured: a stride-2 view of a 6-element index buffer reset rows
- # [0, 1] instead of [0, 1, 2] (read through the base pointer); a
- # stride-2 flag view reset row 1 instead of nothing; a [2, 3]
- # index/flag pair considered only its first 2 of 6 requests.
- assert state_indices.dim() == 1 and state_indices.is_contiguous(), (
- "state_indices must be a contiguous 1-D tensor"
- )
- assert has_initial_states.dim() == 1 and has_initial_states.is_contiguous(), (
- "has_initial_states must be a contiguous 1-D tensor"
- )
-
-
-reset_recurrent_state_rows = ResetRecurrentStateRows()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/__init__.py
deleted file mode 100644
index 52a7a9daf028..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/patch_embed_conv3d.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/patch_embed_conv3d.py
deleted file mode 100644
index 16b691b4008a..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/patch_embed_conv3d.py
+++ /dev/null
@@ -1,141 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Vision patch embedding: one 3-D convolution per pixel patch.
-
-Tier-3 composed: one `torch.nn.functional.conv3d` call between two views that
-launch nothing. No registered op is a 3-D convolution, and the CuTe DSL
-implicit-GEMM conv3d kernels in `visual_gen/kernels/blackwell/conv/` refuse this
-checkpoint's three input channels (bf16 needs `in_channels % 8 == 0`;
-`can_implement` swallows the refusal and returns False).
-"""
-
-from __future__ import annotations
-
-from typing import Optional
-
-import torch
-import torch.nn.functional as F
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import (
- Arch,
- Cell,
- OpWrapper,
- true_fp32_matmul,
-)
-
-#: This checkpoint: in_channels 3, kernel (temporal_patch_size 2, patch_size
-#: 16, 16), out_channels 1152 -- weight [1152, 3, 2, 16, 16], a 1536-pixel row.
-IN_CHANNELS = 3
-KERNEL = (2, 16, 16)
-EMBED_DIM = 1152
-
-
-class PatchEmbedConv3d(OpWrapper):
- """`out[n, k] = bias[k] + sum_{c,t,i,j} pixel[n, c, t, i, j] * weight[k, c, t, i, j]`.
-
- One whole `[in_channels, kt, kh, kw]` patch per row of `pixel_values`, flat
- `[N, C*kt*kh*kw]` (what the image processor emits, channel-slowest) or
- explicit `[N, C, kt, kh, kw]`; the kernel extent is read from `weight`.
- Cross-correlation (no kernel flip), like `nn.Conv3d`. Returns a fresh
- contiguous `[N, out_channels]` in the input dtype.
-
- Fusion boundary: the patch projection and its bias. Patchification, the
- position embedding added after it, and the row order (rows pass through in
- the order they arrive -- the merge-block order the merger needs) are the
- caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: num_patches, dtype. The checkpoint's geometry, flat contiguous
- # rows and a bias: what the vision tower passes.
- CELLS: tuple[Cell, ...] = tuple(
- Cell(
- why=f"checkpoint geometry, {n} patches bf16",
- spec=dict(num_patches=n, dtype=torch.bfloat16),
- )
- for n in (64, 4096)
- ) + (
- Cell(
- why="tower call ceiling: 8192 tokens x merge 2x2 = 32768 patches bf16",
- spec=dict(num_patches=32768, dtype=torch.bfloat16),
- ),
- )
-
- note = """
- Metadata consumed: none. `pixel_values` is not mutated. All three operands
- share one dtype and device -- torch raises "... should be the same"
- otherwise. Several images, and a video's frame groups, are just more rows;
- there is no batch axis.
-
- fp32 runs in TF32 unless the caller sets `torch.backends.cudnn.allow_tf32 =
- False` (global state the caller owns): measured 5.84e-4 from an exact fp32
- matmul with it on (99.99% of outputs differ), 3.58e-6 with it off. bf16 and
- fp16 are unaffected.
-
- At this shape the op is a GEMM: one output position per row, and the
- `[1152, 1536]` reshape of the filter is a view of the same bytes. Not timed.
- CUDA-graph capturable; the output is a fresh allocation.
-
- Not batch-invariant: cuDNN picks its kernel by row count, so a row's last bits
- depend on how many rows share the call. Check rows against `reference`, never
- against a call of a different size.
-
- Not certified: dtypes other than bf16 (the fp32 TF32 behavior above is
- measured, not certified); other geometries, `bias=None` and non-contiguous
- pixel rows; patch counts below 64, between the cells, and past 32768
- (cuDNN picks its kernel by row count).
- """
-
- def __call__(
- self,
- pixel_values: torch.Tensor,
- weight: torch.Tensor,
- bias: Optional[torch.Tensor] = None,
- ) -> torch.Tensor:
- out_channels, in_channels, kt, kh, kw = weight.shape
- patches = pixel_values.view(-1, in_channels, kt, kh, kw)
- return F.conv3d(patches, weight, bias).view(-1, out_channels)
-
- def reference(
- self,
- pixel_values: torch.Tensor,
- weight: torch.Tensor,
- bias: Optional[torch.Tensor] = None,
- ) -> torch.Tensor:
- """True-fp32 matmul of the flattened rows against the flattened filter."""
- rows = pixel_values.reshape(pixel_values.shape[0], -1).float()
- filters = weight.reshape(weight.shape[0], -1).float()
- with true_fp32_matmul():
- out = rows @ filters.t()
- if bias is not None:
- out = out + bias.float()
- return out.to(pixel_values.dtype)
-
- def is_valid(
- self,
- pixel_values: torch.Tensor,
- weight: torch.Tensor,
- bias: Optional[torch.Tensor] = None,
- ) -> None:
- # The view reinterprets anything whose element count divides: a whole
- # [1, 3, 2, 64, 64] volume becomes 16 "patches" -- a [16, 1152] result
- # of the right shape from the wrong pixels, 2.52 from the strided conv
- # the caller meant -- and [N, 3072] rows become 2N, without an error.
- _, in_channels, kt, kh, kw = weight.shape
- patch_numel = in_channels * kt * kh * kw
- if pixel_values.dim() == 2:
- one_patch_per_row = pixel_values.shape[1] == patch_numel
- elif pixel_values.dim() == 5:
- one_patch_per_row = tuple(pixel_values.shape[1:]) == (in_channels, kt, kh, kw)
- else:
- one_patch_per_row = False
- assert one_patch_per_row, (
- f"each row must be exactly one patch: pixel_values must be "
- f"[num_patches, {patch_numel}] or "
- f"[num_patches, {in_channels}, {kt}, {kh}, {kw}] for this weight; got "
- f"{tuple(pixel_values.shape)}"
- )
-
-
-patch_embed_conv3d = PatchEmbedConv3d()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/pos_embed_interpolate.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/pos_embed_interpolate.py
deleted file mode 100644
index b616a48a86de..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/pos_embed_interpolate.py
+++ /dev/null
@@ -1,243 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Vision position-embedding bilinear interpolation onto each image's patch
-grid, emitted in spatial-merge order.
-
-Tier-3 composed: index and corner-weight arithmetic on `[h]`, `[w]` and
-`[h, w]` tensors, then one `torch.nn.functional.embedding_bag` for the whole
-batch. No registered op interpolates anything, and `modules/` has no vision
-position-embedding callable. The exact-math sibling,
-`_triton_pos_embed_interpolate` in `models/modeling_qwen3vl.py`, is a private
-symbol of another model's modeling file; the other in-tree resamplers are
-bicubic and/or `align_corners=False`, a different function.
-"""
-
-from __future__ import annotations
-
-import math
-from typing import Sequence
-
-import torch
-import torch.nn.functional as F
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper
-
-#: This checkpoint: num_position_embeddings 2304 (a 48 x 48 grid), hidden 1152,
-#: spatial_merge_size 2.
-NUM_GRID = 48
-HIDDEN = 1152
-MERGE = 2
-
-
-def _merge_order(x: torch.Tensor, height: int, width: int, merge: int) -> torch.Tensor:
- """`[h, w, ...]` raster -> `[h*w, ...]` in (h_block, w_block, row, col) order."""
- tail = x.shape[2:]
- return (
- x.reshape(height // merge, merge, width // merge, merge, *tail)
- .permute(0, 2, 1, 3, *range(4, 4 + len(tail)))
- .reshape(height * width, *tail)
- )
-
-
-class PosEmbedInterpolate(OpWrapper):
- """Resample a square row-major `[N*N, hidden]` table onto each `(t, h, w)`
- patch grid; return `[sum(t*h*w), hidden]` in the table's dtype.
-
- Corner-aligned bilinear: `pos = linspace(0, N - 1, n)` per axis, `lo =
- floor(pos)`, `hi = min(lo + 1, N - 1)`; `h` indexes table rows (stride `N`),
- `w` columns. Rows come out in spatial-merge order -- whole `merge x merge`
- blocks row-major, row-major inside a block -- which is what the patch
- merger's 4-to-1 reshape consumes; raster order is silently wrong. `t`
- repeats the frame's rows (no temporal axis); images concatenate in order.
- The four-term sum accumulates in fp32 and rounds once.
-
- Fusion boundary: the position embeddings alone. Adding them to the patch
- embeddings, the 2-D rope and the attention segmentation are the caller's.
- """
-
- ARCHS = frozenset({Arch.SM_103})
-
- # spec: grid (tuple of (t, h, w)), merge, dtype, rows (table rows), hidden.
- CELLS: tuple[Cell, ...] = (
- Cell(
- why="a 4-image batch incl. the 62x90 serving grid",
- spec=dict(
- grid=((1, 4, 8), (2, 6, 4), (1, 48, 48), (1, 62, 90)),
- merge=MERGE,
- dtype=torch.bfloat16,
- rows=NUM_GRID**2,
- hidden=HIDDEN,
- ),
- ),
- Cell(
- why="video: 3 frames repeat one 6x10 frame",
- spec=dict(
- grid=((3, 6, 10),),
- merge=MERGE,
- dtype=torch.bfloat16,
- rows=NUM_GRID**2,
- hidden=HIDDEN,
- ),
- ),
- # The largest grid one call carries (32768 patches).
- Cell(
- why="checkpoint table bf16, grid [[1, 128, 256]]",
- spec=dict(
- grid=((1, 128, 256),),
- merge=MERGE,
- dtype=torch.bfloat16,
- rows=NUM_GRID**2,
- hidden=HIDDEN,
- ),
- ),
- )
-
- note = """
- Metadata consumed: none. `grid_thw` is host ints (a device tensor works but
- `int()` syncs per element). The table is not mutated; the result is fresh.
-
- Loud failures, left to torch: an `h` or `w` not divisible by
- `spatial_merge_size` (or a merge larger than the grid) raises "shape ... is
- invalid for input of size ..."; a non-2-D table raises "weight has to be a
- 2D Tensor". An empty grid side returns `[0, hidden]` rather than raising.
-
- At `h = w = N` the samples are exactly `arange(N)` and the result is
- bitwise the table in merge order. The bf16 result is within one bf16
- half-ulp of the fp64 value (0.0147 at 62x90); a bf16 multiply-add chain
- (a literal HF transcription) lands up to one bf16 ulp (0.031) from this.
-
- CUDA graphs: capturable, but the grid is a host argument, so it is frozen
- at capture. Cost is ~225 us of launch-bound index construction per image,
- flat in token count; the gather is the only part that scales.
-
- Certified at the checkpoint's bf16 table on a 4-image batch (4x8, a 2x6x4
- video, the 48x48 exact fit and the 62x90 serving grid), a 3-frame 6x10
- video and the 128x256 grid one call carries at most. Not certified: a CPU
- table (it runs and returns on CPU); fp16 and fp32 tables; a
- `spatial_merge_size` other than the checkpoint's 2.
- """
-
- def __call__(
- self,
- pos_embed_weight: torch.Tensor,
- grid_thw: Sequence[Sequence[int]],
- spatial_merge_size: int,
- ) -> torch.Tensor:
- num_grid_per_side = math.isqrt(pos_embed_weight.shape[0])
- device = pos_embed_weight.device
- merge = spatial_merge_size
- idx_parts = []
- weight_parts = []
-
- for frames, height, width in grid_thw:
- frames, height, width = int(frames), int(height), int(width)
-
- # Corner-aligned; non-negative positions, so the int64 cast floors.
- h_pos = torch.linspace(
- 0, num_grid_per_side - 1, height, device=device, dtype=torch.float32
- )
- w_pos = torch.linspace(
- 0, num_grid_per_side - 1, width, device=device, dtype=torch.float32
- )
- h_lo = h_pos.to(torch.int64)
- w_lo = w_pos.to(torch.int64)
- h_hi = (h_lo + 1).clamp(max=num_grid_per_side - 1)
- w_hi = (w_lo + 1).clamp(max=num_grid_per_side - 1)
- dh = (h_pos - h_lo).unsqueeze(1)
- dw = (w_pos - w_lo).unsqueeze(0)
-
- row_lo = (h_lo * num_grid_per_side).unsqueeze(1)
- row_hi = (h_hi * num_grid_per_side).unsqueeze(1)
- col_lo = w_lo.unsqueeze(0)
- col_hi = w_hi.unsqueeze(0)
-
- # [h, w, 4]: corners in (lo,lo) (lo,hi) (hi,lo) (hi,hi) order.
- idx = torch.stack(
- (row_lo + col_lo, row_lo + col_hi, row_hi + col_lo, row_hi + col_hi), dim=-1
- )
- weights = torch.stack(
- ((1 - dh) * (1 - dw), (1 - dh) * dw, dh * (1 - dw), dh * dw), dim=-1
- )
-
- # Spatial-merge reorder; `reshape` is loud when a side does not divide.
- merged_h, merged_w = height // merge, width // merge
- idx = (
- idx.reshape(merged_h, merge, merged_w, merge, 4)
- .permute(0, 2, 1, 3, 4)
- .reshape(height * width, 4)
- )
- weights = (
- weights.reshape(merged_h, merge, merged_w, merge, 4)
- .permute(0, 2, 1, 3, 4)
- .reshape(height * width, 4)
- )
-
- if frames > 1:
- idx = idx.repeat(frames, 1)
- weights = weights.repeat(frames, 1)
- idx_parts.append(idx)
- weight_parts.append(weights)
-
- idx = idx_parts[0] if len(idx_parts) == 1 else torch.cat(idx_parts)
- weights = weight_parts[0] if len(weight_parts) == 1 else torch.cat(weight_parts)
-
- # One launch: each row of `idx` is a 4-element bag scaled by `weights`
- # and summed, fp32 table and fp32 weights.
- table = pos_embed_weight.float()
- out = F.embedding_bag(idx, table, mode="sum", per_sample_weights=weights)
- return out.to(pos_embed_weight.dtype)
-
- def reference(
- self,
- pos_embed_weight: torch.Tensor,
- grid_thw: Sequence[Sequence[int]],
- spatial_merge_size: int,
- ) -> torch.Tensor:
- """The explicit four-corner weighted sum in fp64, rounded once.
-
- Shares the entry's sample positions (fp32 `linspace`) but not its
- gather or accumulation; the test cross-checks the structure separately
- against `F.interpolate`.
- """
- num_grid = math.isqrt(pos_embed_weight.shape[0])
- hidden = pos_embed_weight.shape[1]
- grid = pos_embed_weight.double().view(num_grid, num_grid, hidden)
- device = pos_embed_weight.device
- outs = []
- for frames, height, width in grid_thw:
- frames, height, width = int(frames), int(height), int(width)
- h_pos = torch.linspace(0, num_grid - 1, height, device=device)
- w_pos = torch.linspace(0, num_grid - 1, width, device=device)
- h_lo, w_lo = h_pos.to(torch.int64), w_pos.to(torch.int64)
- h_hi = (h_lo + 1).clamp(max=num_grid - 1)
- w_hi = (w_lo + 1).clamp(max=num_grid - 1)
- d_h = (h_pos - h_lo).double().unsqueeze(1)
- d_w = (w_pos - w_lo).double().unsqueeze(0)
- raster = (
- grid[h_lo][:, w_lo] * ((1 - d_h) * (1 - d_w))[..., None]
- + grid[h_lo][:, w_hi] * ((1 - d_h) * d_w)[..., None]
- + grid[h_hi][:, w_lo] * (d_h * (1 - d_w))[..., None]
- + grid[h_hi][:, w_hi] * (d_h * d_w)[..., None]
- )
- rows = _merge_order(raster, height, width, spatial_merge_size)
- outs.append(rows.repeat(frames, 1) if frames > 1 else rows)
- out = torch.cat(outs) if len(outs) > 1 else outs[0]
- return out.to(pos_embed_weight.dtype)
-
- def is_valid(
- self,
- pos_embed_weight: torch.Tensor,
- grid_thw: Sequence[Sequence[int]],
- spatial_merge_size: int,
- ) -> None:
- # With 2300 rows isqrt truncates to 47, every gather index stays in
- # range, and the result is a correctly shaped resample of the wrong grid.
- num_rows = pos_embed_weight.shape[0]
- side = math.isqrt(num_rows)
- assert side * side == num_rows, (
- f"pos_embed_weight must hold a square grid; got {num_rows} rows, "
- f"which is not a perfect square"
- )
-
-
-pos_embed_interpolate = PosEmbedInterpolate()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py
deleted file mode 100644
index 56753f9c6a1b..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py
+++ /dev/null
@@ -1,1438 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Weight table and loader: Qwen3.8-Flash-Next-NVFP4 / sm_103 / tp1.
-
-One `W` per weight role (shape, checkpoint source, dtype, load-time transform);
-`Qwen4ExpTp1Weights` unrolls the table into declaration, manifest and load. tp1
-means no per-rank slicing. Storage is HF `[out, in]`; GEMM views, grouped-norm
-slices, fp32 gating params and NVFP4 scalars are derived in `modeling.py` after load.
-
-Transforms:
-
-* trunk attention input projection: `q_proj`, `k_proj`, `v_proj` and the indexer's
- `index_qk_proj` read the same activation, so they are one
- `[q | k | v | gate | index | pad, hidden]` operand and one GEMM (`_attn_in_w`). `q_proj`
- interleaves q and the output gate per head; its q halves go first so the GEMM output's
- leading columns are the packed `[q | k | v]` the fused norm+rope rewrites in place, then
- the gate halves, the index projection and `index_head_dim` zero rows (a row is a whole
- number of `head_dim` heads). The draft (MTP) layer keeps its separate projections;
-* fused shared-expert gate|up: `[gate ; up]` (the half order `flashinfer_silu_and_mul` reads);
-* NVFP4 routed-expert stacks, laid out for the trunk's routed backend
- (`core.trunk_moe_backend`): `[up ; gate]`, then 128x4 scale swizzle. trtllm-gen also
- interleaves FC1 row by row and 32-row block shuffles weight and scale bytes (FC2:
- shuffle only); CuteDSL interleaves FC1 in 64-row up/gate groups and leaves FC2 in
- checkpoint row order. Under trtllm-gen a second copy in the CuteDSL layout (`fc*_cd`)
- serves the prefill-sized launches (`modeling._CUTE_DSL_MOE_MIN_ROWS`) when the core keeps
- it (`core.moe_prefill_copy`);
-* n-gram table: 128 row chunks concatenated into one `[rows, 160]` buffer (the hash
- emits global row indices).
-
-The n-gram table is not in the table: its height is a checkpoint fact, so `load`
-registers it with `torch.empty` and adds its rows to the manifest;
-`post_load_weights` reconciles the height. Residency (device or host) is set by
-`TRTLLM_MODELING_V2_PLE_HOST_OFFLOAD` (`_host_resident_empty`); the manifest is the
-same either way. Offloading frees ~47.7 GiB, which the KV cache manager turns into
-more slots; several entries certify at most 461 slots, so raising
-`kv_cache_fraction` with the offload on needs those cells re-checked.
-
-Not loaded: `mtp.*` when there is no `speculative_config`. With one, the whole MTP module
-is loaded, including the draft layer's own QSA indexer (`_index_branch`). MTP experts are
-plain row-major FP8 128x128: `[up ; gate]` concat, bf16
-`weight_scale_inv` cast to fp32. `lm_head` is the shell's parameter.
-
-`load(model, weights)` fills every declared parameter once and asserts coverage
-both ways: every key is consumed or in the non-load set.
-"""
-
-import mmap
-import os
-from dataclasses import dataclass
-from types import SimpleNamespace
-from typing import Any, Callable
-
-import torch
-from cuda.bindings import runtime as cudart
-from torch import nn
-
-from tensorrt_llm._utils import CUASSERT
-from tensorrt_llm.llmapi.utils import get_numa_aware_cpu_affinity
-from tensorrt_llm.logger import logger
-
-#: `layers=ALL_LAYERS` is the common case and the default. A sentinel rather
-#: than a lambda so the table reads as data, and rather than `True` so the
-#: field has one type of answer: which layers.
-ALL_LAYERS = object()
-
-
-@dataclass(frozen=True)
-class W:
- """One weight role: what it is, where its bytes come from, what happens on the way.
-
- `shape` takes the dimension bundle and returns the declared shape; every
- shape in these models is derived from configuration rather than constant.
-
- `src` says where the bytes come from, in one of three forms:
-
- * a checkpoint key template, with `{p}` standing for the layer prefix;
- * a sequence of `(template, slice_fn)` pairs, when one parameter is
- assembled from several checkpoint tensors side by side -- `slice_fn`
- takes the dimension bundle and returns an index into the destination;
- * a callable `(d, i) -> [(key, index), ...]`, taking the dimension bundle
- and the layer index. The first two forms are special cases of it; this
- one exists for sources that are not a fixed template -- a per-expert
- row, or a prefix that differs between layer families. A `key` may itself
- be a tuple, meaning the transform is handed those tensors together.
-
- `layers` says which layers carry this weight: `ALL_LAYERS` for every one,
- `None` for a weight that is not per-layer at all, or a callable
- `d -> iterable[int]` when only some layers have it.
-
- `transform` runs on the source tensor before the copy, on the
- destination's device. `None` means the checkpoint tensor is already the
- declared shape.
- """
-
- name: str
- shape: Callable[[Any], tuple[int, ...]]
- src: str | tuple[tuple[str, Callable], ...] | Callable
- dtype: torch.dtype | None = None
- transform: Callable | None = None
- layers: object | None = ALL_LAYERS
-
- def layer_indices(self, d: SimpleNamespace) -> list[int] | None:
- """The layers this weight exists on, or None if it is not per-layer."""
- if self.layers is None:
- return None
- if self.layers is ALL_LAYERS:
- return list(range(d.num_layers))
- return list(self.layers(d))
-
- def sources(self, d: SimpleNamespace, i: int | None, prefix: str) -> list[tuple[Any, Any]]:
- """`[(checkpoint key or key tuple, index into the param | None), ...]`."""
- if callable(self.src):
- return list(self.src(d, i))
- if isinstance(self.src, str):
- return [(self.src.format(p=prefix), None)]
- return [(template.format(p=prefix), slice_fn(d)) for template, slice_fn in self.src]
-
-
-# ----------------------------------------------------------- host residency --
-#
-# Allocates the one PLE n-gram table. The table stays a `cuda` tensor of the same
-# shape/dtype. Registered, device-mapped pages (not a host gather) because the gather's
-# row indices are produced on device inside the captured decode graph.
-
-#: `cudaMemoryType` values; the core asserts on them after load.
-CUDA_MEMORY_TYPE_HOST = 1
-CUDA_MEMORY_TYPE_DEVICE = 2
-
-
-def memory_kind(ptr: int) -> int:
- """The driver's own answer to "what kind of memory is this pointer".
-
- Returns a `cudaMemoryType`. This is the residency question asked of the
- process rather than of `nvidia-smi`, and it is the only one a
- working-but-resident implementation cannot answer the way an offloaded one
- does.
- """
- (attributes,) = CUASSERT(cudart.cudaPointerGetAttributes(ptr))
- return int(attributes.type)
-
-
-#: `cudaHostRegister` flags (`driver_types.h`): mapped, so a kernel may dereference the
-#: pages, and portable, for every context in the process.
-_CUDA_HOST_REGISTER_PORTABLE = 0x01
-_CUDA_HOST_REGISTER_MAPPED = 0x02
-
-
-def _local_cpus(device) -> set | None:
- """The CPUs of `device`'s NUMA affinity, as NVML reports it (`get_numa_aware_cpu_affinity`),
- that this process may run on, or None when NVML cannot tell."""
- import pynvml
-
- index = torch.device(device).index
- try:
- local = get_numa_aware_cpu_affinity(torch.cuda.current_device() if index is None else index)
- except pynvml.NVMLError:
- return None
- return (set(local) & os.sched_getaffinity(0)) or None
-
-
-class _RegisteredHostBuffer:
- """A pageable CPU mapping of its own, registered (pinned and mapped) for the
- device, and unregistered before its pages are unmapped: freeing registered pages
- would leave the driver mapping memory the OS may hand out again.
-
- The mapping is advised `MADV_HUGEPAGE` before the registration faults it in, so
- the table sits on transparent huge pages whatever torch's CPU allocator does
- (TensorRT LLM imports torch with mimalloc's THP advice off). On GB300 the
- device's random row gathers from 64 KiB pages run 2.5x (64 tokens) to 23x
- (1024 tokens) slower.
-
- Not torch's pinned allocator: it rounds the 47.7 GiB table up to 64 GiB, and on
- GB300 the device's random row gathers from its pages run an order of magnitude
- slower than from registered pageable memory."""
-
- def __init__(self, nbytes: int, device=None) -> None:
- # The registration faults the pages in on this thread, so under the default local
- # allocation policy they land on the NUMA node of the CPU it runs on. Pinned to the
- # CPUs local to `device` meanwhile, the table sits next to the GPU that gathers from
- # it: on GB300, random row gathers from the other socket's memory run ~1.7x slower,
- # and with a process per GPU half of them would otherwise land there.
- local = _local_cpus(device) if device is not None else None
- previous = os.sched_getaffinity(0) if local else None
- if local:
- os.sched_setaffinity(0, local)
- try:
- self._mapping = mmap.mmap(-1, nbytes, flags=mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS)
- try:
- self._mapping.madvise(mmap.MADV_HUGEPAGE)
- except OSError as error: # a kernel without transparent huge pages
- logger.warning(
- f"madvise(MADV_HUGEPAGE) failed for the host-resident PLE table ({error}); "
- "the device's row gathers from its base pages run several times slower"
- )
- # Shares the mapping's pages, and keeps the mapping alive while it lives.
- self._storage = torch.frombuffer(self._mapping, dtype=torch.uint8)
- self.pointer = self._storage.data_ptr()
- CUASSERT(
- cudart.cudaHostRegister(
- self.pointer, nbytes, _CUDA_HOST_REGISTER_PORTABLE | _CUDA_HOST_REGISTER_MAPPED
- )
- )
- finally:
- if previous:
- os.sched_setaffinity(0, previous)
- # Bound now: at interpreter shutdown the module may already be torn down.
- self._unregister = cudart.cudaHostUnregister
-
- def __del__(self) -> None:
- pointer, self.pointer = self.pointer, 0
- if pointer:
- self._unregister(pointer)
-
-
-def _host_resident_empty(
- rows: int, cols: int, dtype, device
-) -> tuple[torch.Tensor, _RegisteredHostBuffer]:
- """`[rows, cols]` of `dtype` on `device`, whose bytes live in host memory.
-
- Returns the tensor and the registered CPU allocation that owns the pages. The
- caller **must** keep the second alive for as long as the first: the tensor
- is a view through a device address and owns nothing.
- """
- itemsize = torch.empty(0, dtype=dtype).element_size()
- assert itemsize == 1, (
- f"the host-resident allocation reinterprets a `|u1` buffer, so its "
- f"element must be one byte; {dtype} is {itemsize}"
- )
- assert device.type == "cuda", f"host offload needs a cuda device to map into; got {device}"
-
- host = _RegisteredHostBuffer(rows * cols, device)
- (device_pointer,) = CUASSERT(cudart.cudaHostGetDevicePointer(host.pointer, 0))
-
- # No public torch API builds a tensor over a pointer the caching allocator
- # did not hand out; `__cuda_array_interface__` is the numba/cupy path.
- # `|u1` because the protocol has no float8 typestr; reinterpreted below.
- interface = SimpleNamespace(
- __cuda_array_interface__={
- "shape": (rows, cols),
- "typestr": "|u1",
- "data": (int(device_pointer), False),
- "version": 3,
- "strides": None,
- }
- )
- table = torch.as_tensor(interface, device=device).view(dtype)
- # `torch.as_tensor` copies when it decides the buffer is not already on the
- # requested device; a copy here would be a silent 48 GiB device allocation.
- assert table.data_ptr() == int(device_pointer), (
- "torch.as_tensor copied the host-mapped buffer instead of viewing it, "
- "so the table would be device-resident after all"
- )
- assert table.shape == (rows, cols) and table.dtype == dtype, (table.shape, table.dtype)
- return table, host
-
-
-# ------------------------------------------------------------------ the table --
-
-_LM = "model.language_model"
-_VISUAL = "model.visual"
-_MTP = "mtp"
-_MTP_LAYER = "mtp.layers.0"
-
-# NVFP4 block size: one e4m3 scale per 16 contiguous elements along K.
-SF_VEC = 16
-
-# The FP8 block-scale family's block edge, on both axes of the MTP experts'
-# weight scales and along the activation row.
-FP8_BLOCK = 128
-
-# Parameter-key suffixes after which to release transform scratch. The expert
-# relayouts allocate per-expert scratch 512 times per parameter and the table
-# fill stages one shard at a time; without this, peak load memory is the whole
-# model rather than one parameter.
-_RELEASE_AFTER = (
- "_fc2_s",
- "_fc1_s",
- "_fc2_s_cd",
- "_fc1_s_cd",
- "ple_table",
- "mtp_fc1_w",
- "mtp_fc2_w",
-)
-
-
-def _block32_perm(rows: int, device) -> torch.Tensor:
- """Gather index of the 32-row block shuffle: inside each aligned block of
- 32 rows, source row `4u + v` (0<=u<8, 0<=v<4) lands at destination row
- `8v + u`."""
- assert rows % 32 == 0, rows
- src = torch.arange(32)
- dst = (src % 4) * 8 + src // 4
- idx = torch.empty(32, dtype=torch.long)
- idx[dst] = src
- blocks = rows // 32
- return (idx.repeat(blocks) + torch.arange(blocks).repeat_interleave(32) * 32).to(device)
-
-
-def _interleave_perm(rows: int, device) -> torch.Tensor:
- """Gather index that interleaves the `[up | gate]` halves of a `2*I` row
- stack into up0, gate0, up1, gate1, ..."""
- p = torch.empty(rows, dtype=torch.long)
- p[0::2] = torch.arange(0, rows // 2)
- p[1::2] = torch.arange(rows // 2, rows)
- return p.to(device)
-
-
-#: The group the CuteDSL FC1 epilogue pairs up and gate rows in: rows `[128k, 128k + 64)`
-#: of the fused operand are up rows and the next 64 their gate rows.
-_CUTEDSL_GATE_UP_GROUP = 64
-
-
-def _group_interleave_perm(rows: int, group: int, device) -> torch.Tensor:
- """Gather index that interleaves the `[up | gate]` halves of a `2*I` row stack in
- `group`-row blocks: up[0:g], gate[0:g], up[g:2g], gate[g:2g], ..."""
- assert rows % (2 * group) == 0, (rows, group)
- idx = torch.arange(rows, device=device).view(2, rows // (2 * group), group)
- return idx.transpose(0, 1).reshape(-1)
-
-
-def _fc1_perm(rows: int, device, core) -> torch.Tensor:
- """The FC1 row gather of the trunk's routed backend. trtllm-gen: interleave then
- 32-row block shuffle, composed into one gather. CuteDSL: the 64-row group
- interleave alone."""
- if core.trunk_moe_backend == "CUTEDSL":
- return _group_interleave_perm(rows, _CUTEDSL_GATE_UP_GROUP, device)
- return _interleave_perm(rows, device)[_block32_perm(rows, device)]
-
-
-def _swizzle(x: torch.Tensor) -> torch.Tensor:
- """128x4 block-scale swizzle of an `[M, C]` uint8 scale matrix, returned
- flat (`M % 128 == 0` and `C % 4 == 0` hold at this geometry, so the result
- is exactly `M * C` bytes)."""
- return torch.ops.trtllm.block_scale_interleave(x.contiguous())
-
-
-#: Rows after the shared expert's gate|up rows in its fused operand: the per-token gate's
-#: one row, then zeros, so the GEMM's output row stays a whole number of 16-byte vectors.
-SH_GATE_ROWS = 8
-
-
-def _cat_rows(t: tuple[torch.Tensor, ...], core) -> torch.Tensor:
- """Fused shared-expert weight bytes: **gate** rows first, the half order
- `flashinfer_silu_and_mul` reads, then up, then the per-token gate's row and
- `SH_GATE_ROWS - 1` zero rows -- one GEMM forms gate|up and the gate logit."""
- gate, up, sh_gate = t
- pad = sh_gate.new_zeros(SH_GATE_ROWS - 1, sh_gate.shape[-1])
- return torch.cat((gate, up, sh_gate, pad), dim=0).contiguous()
-
-
-def _attn_in_w(t: tuple[torch.Tensor, ...], core) -> torch.Tensor:
- """(q_proj, k_proj, v_proj, index_qk_proj) -> the fused `[q | k | v | gate | index |
- pad]` input projection of one trunk attention layer (see the module docstring).
-
- `q_proj`'s rows are `[heads, 2, head_dim]`: head `h`'s q rows, then its gate rows.
- """
- q, k, v, idx = t
- heads, head_dim, hidden = core.heads_q, core.head_dim, q.shape[-1]
- qg = q.reshape(heads, 2, head_dim, hidden)
- pad = q.new_zeros(core.index_head_dim, hidden)
- rows = (qg[:, 0].reshape(heads * head_dim, hidden), k, v)
- rows += (qg[:, 1].reshape(heads * head_dim, hidden), idx, pad)
- return torch.cat(rows, dim=0).contiguous()
-
-
-def _expert_fc1_w(t: tuple[torch.Tensor, ...], core) -> torch.Tensor:
- """(up, gate) packed nibbles -> one expert's FC1 operand."""
- up, gate = t
- w = torch.cat([up, gate], dim=0)
- return torch.index_select(w, 0, _fc1_perm(w.shape[0], w.device, core)).contiguous()
-
-
-def _expert_fc1_s(t: tuple[torch.Tensor, ...], core) -> torch.Tensor:
- """(up, gate) e4m3 scales -> one expert's FC1 scale operand: the weight's own row
- permutation, then the swizzle."""
- up, gate = t
- s = torch.cat([up.view(torch.uint8), gate.view(torch.uint8)], dim=0)
- rows, cols = s.shape
- s = torch.index_select(s, 0, _fc1_perm(rows, s.device, core))
- return _swizzle(s).reshape(rows, cols)
-
-
-def _expert_fc2_w(t: torch.Tensor, core) -> torch.Tensor:
- """`down_proj` packed nibbles -> one expert's FC2 operand (no interleave): the
- 32-row block shuffle on trtllm-gen, checkpoint row order on CuteDSL."""
- if core.trunk_moe_backend == "CUTEDSL":
- return t.contiguous()
- return torch.index_select(t, 0, _block32_perm(t.shape[0], t.device)).contiguous()
-
-
-def _expert_fc2_s(t: torch.Tensor, core) -> torch.Tensor:
- """`down_proj` e4m3 scales -> one expert's FC2 scale operand: the weight's own row
- order, then the swizzle."""
- s = t.view(torch.uint8)
- rows, cols = s.shape
- if core.trunk_moe_backend != "CUTEDSL":
- s = torch.index_select(s, 0, _block32_perm(rows, s.device))
- return _swizzle(s).reshape(rows, cols)
-
-
-#: The layout argument of the second, CuteDSL-layout expert copy (`fc*_cd`).
-_CUTEDSL_LAYOUT = SimpleNamespace(trunk_moe_backend="CUTEDSL")
-
-
-def _expert_fc1_w_cd(t: tuple[torch.Tensor, ...], core) -> torch.Tensor:
- """(up, gate) packed nibbles -> one expert's FC1 operand in the CuteDSL layout."""
- return _expert_fc1_w(t, _CUTEDSL_LAYOUT)
-
-
-def _expert_fc1_s_cd(t: tuple[torch.Tensor, ...], core) -> torch.Tensor:
- """(up, gate) e4m3 scales -> one expert's FC1 scale operand in the CuteDSL layout."""
- return _expert_fc1_s(t, _CUTEDSL_LAYOUT)
-
-
-def _expert_fc2_s_cd(t: torch.Tensor, core) -> torch.Tensor:
- """`down_proj` e4m3 scales -> one expert's FC2 scale operand in the CuteDSL layout."""
- return _expert_fc2_s(t, _CUTEDSL_LAYOUT)
-
-
-def _mtp_fc1_w(t: tuple[torch.Tensor, ...], core) -> torch.Tensor:
- """(up, gate) e4m3 codes -> one draft expert's FC1 operand, `[up ; gate]` on
- the row axis and nothing else. The FP8 block-scale runner reads the
- checkpoint's plain row-major layout."""
- up, gate = t
- return torch.cat([up, gate], dim=0).contiguous()
-
-
-def _mtp_fc1_s(t: tuple[torch.Tensor, ...], core) -> torch.Tensor:
- """(up, gate) 128x128 block scales -> one draft expert's FC1 scale operand:
- the same `[up ; gate]` concatenation, cast bf16 -> fp32."""
- up, gate = t
- return torch.cat([up, gate], dim=0).to(torch.float32).contiguous()
-
-
-def _mtp_fc2_s(t: torch.Tensor, core) -> torch.Tensor:
- """`down_proj.weight_scale_inv` -> the FC2 scale operand: a bf16 -> fp32
- cast, no relayout."""
- return t.to(torch.float32).contiguous()
-
-
-#: The layout argument of the draft experts' NVFP4 copy: the draft MoE runs it on trtllm-gen
-#: whatever the trunk's routed backend is.
-_TRTLLM_GEN_LAYOUT = SimpleNamespace(trunk_moe_backend="TRTLLM")
-
-#: `448 * 6`: an NVFP4 global scale is this over the tensor's amax (modelopt's recipe).
-NVFP4_AMAX_TO_GLOBAL = 448.0 * 6.0
-
-
-def _fp8_block_dequant(codes: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
- """`[E, N, K]` e4m3 codes times their `[E, N/128, K/128]` fp32 block scales, in fp32."""
- s = scales.repeat_interleave(FP8_BLOCK, 1).repeat_interleave(FP8_BLOCK, 2)
- return codes.float() * s
-
-
-def mtp_experts_to_nvfp4(fc1_w, fc1_s, fc2_w, fc2_s, chunk: int = 32) -> tuple:
- """Requantize the draft layer's FP8 128x128 block-scaled experts to NVFP4 in the
- trunk's trtllm-gen layout (`_expert_fc1_w/_s`, `_expert_fc2_w/_s`).
-
- Per expert and projection (FC1's up|gate share one, as the fused FC1's alpha
- requires): dequantize with the block scales, take the global scale
- `448 * 6 / amax`, then `fp4_quantize` (16-element blocks, e4m3 block scales,
- linear order), which then gets the trunk's row permutation and 128x4 swizzle.
- Returns `(fc1_w [E, 2I, H/2], fc1_s [E, 2I, H/16], fc2_w [E, H, I/2],
- fc2_s [E, H, I/16])` uint8 and the fp32 `[E]` weight global scales `(g1, g2)`;
- a weight's real value is its NVFP4 reconstruction over its global scale.
- `fc1_w` / `fc1_s` are the `[up ; gate]` stacks `_mtp_fc1_w` / `_mtp_fc1_s` build."""
- experts, two_i, hidden = fc1_w.shape
- inter = two_i // 2
- dev = fc1_w.device
- u8 = torch.uint8
- out = (
- torch.empty(experts, two_i, hidden // 2, dtype=u8, device=dev),
- torch.empty(experts, two_i, hidden // SF_VEC, dtype=u8, device=dev),
- torch.empty(experts, hidden, inter // 2, dtype=u8, device=dev),
- torch.empty(experts, hidden, inter // SF_VEC, dtype=u8, device=dev),
- )
- g1 = torch.empty(experts, dtype=torch.float32, device=dev)
- g2 = torch.empty(experts, dtype=torch.float32, device=dev)
- layout = _TRTLLM_GEN_LAYOUT
- for e0 in range(0, experts, chunk):
- e1 = min(e0 + chunk, experts)
- w1 = _fp8_block_dequant(fc1_w[e0:e1], fc1_s[e0:e1])
- w2 = _fp8_block_dequant(fc2_w[e0:e1], fc2_s[e0:e1])
- g1[e0:e1] = NVFP4_AMAX_TO_GLOBAL / w1.abs().amax(dim=(1, 2)).clamp_min(1e-30)
- g2[e0:e1] = NVFP4_AMAX_TO_GLOBAL / w2.abs().amax(dim=(1, 2)).clamp_min(1e-30)
- w1, w2 = w1.to(torch.bfloat16), w2.to(torch.bfloat16)
- for j, e in enumerate(range(e0, e1)):
- q, s = torch.ops.trtllm.fp4_quantize(w1[j], g1[e : e + 1], SF_VEC, False, False)
- s = s.view(two_i, hidden // SF_VEC)
- out[0][e] = _expert_fc1_w((q[:inter], q[inter:]), layout)
- out[1][e] = _expert_fc1_s((s[:inter], s[inter:]), layout)
- q, s = torch.ops.trtllm.fp4_quantize(w2[j], g2[e : e + 1], SF_VEC, False, False)
- out[2][e] = _expert_fc2_w(q, layout)
- out[3][e] = _expert_fc2_s(s.view(hidden, inter // SF_VEC), layout)
- del w1, w2
- return (*out, g1, g2)
-
-
-def _materialize(entry) -> torch.Tensor:
- """Realize one checkpoint entry: `[:]` on a lazy safetensors slice, but a
- 0-dim tensor (every NVFP4 per-tensor scalar on this checkpoint) rejects that
- index."""
- return entry if getattr(entry, "ndim", 1) == 0 else entry[:]
-
-
-# --- which layers carry a weight ---
-
-
-def _attn_layers(d):
- """The full-attention layers (`i % 4 == 3` on this checkpoint)."""
- return [i for i, a in enumerate(d.is_attn) if a]
-
-
-def _gdn_layers(d):
- """The Gated-DeltaNet (linear-attention) layers."""
- return [i for i, a in enumerate(d.is_attn) if not a]
-
-
-#: The Gated-DeltaNet input projections, in the row order of the concatenated
-#: `in_proj_w` (and so the column order of its GEMM output).
-_IN_PROJ_PARTS = ("in_proj_qkv", "in_proj_z", "in_proj_a", "in_proj_b")
-
-
-def _in_proj_rows(d, part: int) -> tuple[slice]:
- """The row range of `in_proj_w` that input projection `part` fills."""
- start = sum(d.in_proj_split[:part])
- return (slice(start, start + d.in_proj_split[part]),)
-
-
-def _cute_dsl_copy_layers(d):
- """Every layer when the core keeps the second, CuteDSL-layout copy that prefill-sized
- launches read under trtllm-gen (`core.moe_prefill_copy`); none otherwise, CuteDSL among
- them, whose one copy is it."""
- return range(d.num_layers) if d.moe_prefill_copy else ()
-
-
-def _vision_blocks(d):
- """The vision tower's blocks. Keyed `l{b}_vis_*`: block `b`, not decoder layer `b`."""
- return range(d.v_depth)
-
-
-def _mtp_only(d):
- """The MTP module, at its own layer index `num_hidden_layers`, only when it is live."""
- return [d.num_layers] if d.mtp_enabled else []
-
-
-# --- where the bytes come from ---
-
-
-def _ple(suffix):
- """A PLE weight. The PLE layer's index is config, so the key is computed."""
- return lambda d, i: [(f"{_LM}.layers.{d.ple_layer}.ple.{suffix}", None)]
-
-
-def _ngram_prefix(ple_layer: int) -> str:
- """The quantized n-gram table submodule: the 128 shards and the one
- per-tensor scale that dequantizes them both hang off this, one level below
- `ple_embedding` itself."""
- return f"{_LM}.layers.{ple_layer}.ple.ple_embedding.ngram_embedding"
-
-
-def _visual(suffix):
- """A vision-tower weight that is not per-block."""
- return lambda d, i: [(f"{_VISUAL}.{suffix}", None)]
-
-
-def _vision_block(suffix):
- """One vision block's weight."""
- return lambda d, i: [(f"{_VISUAL}.blocks.{i}.{suffix}", None)]
-
-
-def _mtp_key(key):
- """One MTP-module tensor, at its own checkpoint key (the module's prefix is not the
- trunk's layer prefix)."""
- return lambda d, i: [(key, None)]
-
-
-def _experts(prefix, *suffixes):
- """One source per routed expert, written at index `e`; a pair of suffixes hands the
- transform both tensors together."""
-
- def build(d, i):
- base = prefix(i)
- out = []
- for e in range(d.num_experts):
- keys = tuple(f"{base}.experts.{e}.{s}" for s in suffixes)
- out.append((keys if len(keys) > 1 else keys[0], (e,)))
- return out
-
- return build
-
-
-def _trunk_mlp(i):
- return f"{_LM}.layers.{i}.mlp"
-
-
-def _mtp_mlp(i):
- return f"{_MTP_LAYER}.mlp"
-
-
-#: The six per-tensor NVFP4 scalars every routed expert carries, as
-#: (parameter suffix, checkpoint suffix). The checkpoint stores modelopt
-#: reciprocals; the core folds them into the runner's scalars after load.
-_EXPERT_SCALARS = (
- ("e_isc1", "gate_proj.input_scale"),
- ("e_isc1_up", "up_proj.input_scale"),
- ("e_ws2_1", "gate_proj.weight_scale_2"),
- ("e_ws2_1_up", "up_proj.weight_scale_2"),
- ("e_isc2", "down_proj.input_scale"),
- ("e_ws2_2", "down_proj.weight_scale_2"),
-)
-
-
-#: (parameter tag, checkpoint module) of the two per-layer hyper-connection gates.
-_HC_MODULES = (("ahc", "attn_hyper_connection"), ("mhc", "mlp_hyper_connection"))
-
-
-def _hyper_connection(tag, src, layers):
- """The four weights of one gated-residual hyper-connection; `src(name)` names the source."""
- return (
- W(f"{tag}_norm", shape=lambda d: (d.stream,), src=src("hc_norm"), layers=layers),
- W(
- f"{tag}_down",
- shape=lambda d: (d.hc_lowrank, d.stream),
- src=src("input_mix_weight_down"),
- layers=layers,
- ),
- W(
- f"{tag}_up",
- shape=lambda d: (d.stream, d.hc_lowrank),
- src=src("input_mix_weight_up"),
- layers=layers,
- ),
- W(
- f"{tag}_inject",
- shape=lambda d: (d.hc_count, d.stream),
- src=src("block_inject_weight"),
- layers=layers,
- ),
- )
-
-
-def _mixer(tag, prefix, layers):
- """A stack-final hyper-connection mixer: a layer's gate minus the injection weight."""
- src = _mtp_key if layers is not None else (lambda key: key)
- return (
- W(
- f"{tag}_norm",
- shape=lambda d: (d.stream,),
- src=src(f"{prefix}.hc_norm.weight"),
- layers=layers,
- ),
- W(
- f"{tag}_down",
- shape=lambda d: (d.hc_lowrank, d.stream),
- src=src(f"{prefix}.input_mix_weight_down.weight"),
- layers=layers,
- ),
- W(
- f"{tag}_up",
- shape=lambda d: (d.stream, d.hc_lowrank),
- src=src(f"{prefix}.input_mix_weight_up.weight"),
- layers=layers,
- ),
- )
-
-
-def _attn_in_fused(tag, src, layers):
- """The trunk attention layers' one input projection: q_proj, k_proj, v_proj and the
- QSA indexer's index_qk_proj, re-laid by `_attn_in_w`."""
- keys = ("q_proj.weight", "k_proj.weight", "v_proj.weight", "indexer.index_qk_proj.weight")
- return (
- W(
- f"{tag}attn_in_w",
- shape=lambda d: (d.attn_in_rows, d.hidden),
- src=lambda d, i: [(tuple(src(k).format(p=f"{_LM}.layers.{i}") for k in keys), None)],
- transform=_attn_in_w,
- layers=layers,
- ),
- )
-
-
-def _gated_attention(tag, src, layers, fused_in=False):
- """The gated GQA block's six weights. `q_proj` is twice as wide as q: q and the output
- gate interleave per head. With `fused_in`, q/k/v are not declared here: they are part
- of `_attn_in_fused`'s one operand."""
- return (
- *_attn_in_separate(tag, src, layers, fused_in),
- *_attn_out_and_norms(tag, src, layers),
- )
-
-
-def _attn_in_separate(tag, src, layers, fused_in):
- if fused_in:
- return ()
- return (
- W(
- f"{tag}q_w",
- shape=lambda d: (2 * d.q_width, d.hidden),
- src=src("q_proj.weight"),
- layers=layers,
- ),
- W(
- f"{tag}k_w",
- shape=lambda d: (d.kv_width, d.hidden),
- src=src("k_proj.weight"),
- layers=layers,
- ),
- W(
- f"{tag}v_w",
- shape=lambda d: (d.kv_width, d.hidden),
- src=src("v_proj.weight"),
- layers=layers,
- ),
- )
-
-
-def _attn_out_and_norms(tag, src, layers):
- return (
- W(
- f"{tag}o_w",
- shape=lambda d: (d.hidden, d.q_width),
- src=src("o_proj.weight"),
- layers=layers,
- ),
- W(f"{tag}q_norm", shape=lambda d: (d.head_dim,), src=src("q_norm.weight"), layers=layers),
- W(f"{tag}k_norm", shape=lambda d: (d.head_dim,), src=src("k_norm.weight"), layers=layers),
- )
-
-
-def _index_branch(tag, src, layers, fused_in=False):
- """The QSA index branch's three weights: one replicated, unquantized projection
- producing `[index_heads q | 1 k]`, and two per-head gemma norms. With `fused_in`, the
- projection is part of `_attn_in_fused`'s one operand and only the norms are declared."""
- projection = W(
- f"{tag}idx_w",
- shape=lambda d: (d.index_width, d.hidden),
- src=src("index_qk_proj.weight"),
- layers=layers,
- )
- return (
- *(() if fused_in else (projection,)),
- W(
- f"{tag}idx_qn",
- shape=lambda d: (d.index_head_dim,),
- src=src("q_layernorm.weight"),
- layers=layers,
- ),
- W(
- f"{tag}idx_kn",
- shape=lambda d: (d.index_head_dim,),
- src=src("k_layernorm.weight"),
- layers=layers,
- ),
- )
-
-
-def _shared_expert(tag, prefix, layers):
- """The bf16 shared expert beside the routed ones, its per-token gate in the same
- operand as its gate|up (`_cat_rows`)."""
- return (
- W(
- f"{tag}sh_gu",
- shape=lambda d: (2 * d.shared_inter + SH_GATE_ROWS, d.hidden),
- src=lambda d, i: [
- (
- (
- f"{prefix(i)}.shared_expert.gate_proj.weight",
- f"{prefix(i)}.shared_expert.up_proj.weight",
- f"{prefix(i)}.shared_expert_gate.weight",
- ),
- None,
- )
- ],
- transform=_cat_rows,
- layers=layers,
- ),
- W(
- f"{tag}sh_dn",
- shape=lambda d: (d.hidden, d.shared_inter),
- src=lambda d, i: [(f"{prefix(i)}.shared_expert.down_proj.weight", None)],
- layers=layers,
- ),
- )
-
-
-u8, f8, f32, i64 = torch.uint8, torch.float8_e4m3fn, torch.float32, torch.int64
-
-
-class Qwen4ExpTp1Weights:
- """Qwen3.8-Flash-Next-NVFP4 / sm_103 / tp1.
-
- The table, and the passes over it: `declare` allocates every weight it names,
- `manifest` maps checkpoint keys onto them, `load` fills them. `declare` and
- `manifest` share one "which parameters exist" decision (`_targets`), so they
- cannot disagree about a weight.
- """
-
- WEIGHTS: tuple[W, ...] = (
- # --- the two per-layer hyper-connection gates: every trunk layer ---
- *(
- w
- for tag, module in _HC_MODULES
- for w in _hyper_connection(
- tag, lambda name, m=module: f"{{p}}.{m}.{name}.weight", ALL_LAYERS
- )
- ),
- # --- the gated GQA layers and their QSA index branch ---
- # q/k/v and the indexer's replicated, unquantized [index_heads q | 1 k]
- # projection as one operand, then o_proj, the q/k norms and the indexer's two
- # per-head gemma norms.
- *_attn_in_fused("", lambda s: f"{{p}}.self_attn.{s}", _attn_layers),
- *_gated_attention("", lambda s: f"{{p}}.self_attn.{s}", _attn_layers, fused_in=True),
- *_index_branch("", lambda s: f"{{p}}.self_attn.indexer.{s}", _attn_layers, fused_in=True),
- # --- the Gated-DeltaNet layers ---
- # The four input projections read the same activation, so they are one
- # weight and one GEMM: rows `[q|k|v (conv_dim) | z (value_dim) | a | b]`
- # (`d.in_proj_split`), each checkpoint tensor copied into its own row
- # range. The forward splits the output into column views.
- W(
- "in_proj_w",
- shape=lambda d: (sum(d.in_proj_split), d.hidden),
- src=tuple(
- (f"{{p}}.linear_attn.{name}.weight", lambda d, j=j: _in_proj_rows(d, j))
- for j, name in enumerate(_IN_PROJ_PARTS)
- ),
- layers=_gdn_layers,
- ),
- W(
- "A_log",
- shape=lambda d: (d.num_v_heads,),
- src="{p}.linear_attn.A_log",
- layers=_gdn_layers,
- ),
- W(
- "dt_bias",
- shape=lambda d: (d.num_v_heads,),
- src="{p}.linear_attn.dt_bias",
- layers=_gdn_layers,
- ),
- # Stored [conv_dim, 1, width]; the conv entries want [conv_dim, width], which is
- # a zero-copy view the targets take once, after load.
- W(
- "conv_w",
- shape=lambda d: (d.conv_dim, 1, d.conv_kernel),
- src="{p}.linear_attn.conv1d.weight",
- layers=_gdn_layers,
- ),
- W(
- "gnorm",
- shape=lambda d: (d.head_v_dim,),
- src="{p}.linear_attn.norm.weight",
- layers=_gdn_layers,
- ),
- W(
- "out_w",
- shape=lambda d: (d.hidden, d.value_dim),
- src="{p}.linear_attn.out_proj.weight",
- layers=_gdn_layers,
- ),
- # --- MoE on every layer: NVFP4 routed experts, a bf16 router and shared expert ---
- W("router", shape=lambda d: (d.num_experts, d.hidden), src="{p}.mlp.gate.weight"),
- W(
- "fc1_w",
- shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // 2),
- dtype=u8,
- src=_experts(_trunk_mlp, "up_proj.weight", "gate_proj.weight"),
- transform=_expert_fc1_w,
- ),
- W(
- "fc1_s",
- shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // SF_VEC),
- dtype=u8,
- src=_experts(_trunk_mlp, "up_proj.weight_scale", "gate_proj.weight_scale"),
- transform=_expert_fc1_s,
- ),
- W(
- "fc2_w",
- shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // 2),
- dtype=u8,
- src=_experts(_trunk_mlp, "down_proj.weight"),
- transform=_expert_fc2_w,
- ),
- W(
- "fc2_s",
- shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // SF_VEC),
- dtype=u8,
- src=_experts(_trunk_mlp, "down_proj.weight_scale"),
- transform=_expert_fc2_s,
- ),
- # Under trtllm-gen, the same experts again in the CuteDSL layout: prefill-sized
- # launches run on them (`modeling._CUTE_DSL_MOE_MIN_ROWS`).
- W(
- "fc1_w_cd",
- shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // 2),
- dtype=u8,
- src=_experts(_trunk_mlp, "up_proj.weight", "gate_proj.weight"),
- transform=_expert_fc1_w_cd,
- layers=_cute_dsl_copy_layers,
- ),
- W(
- "fc1_s_cd",
- shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // SF_VEC),
- dtype=u8,
- src=_experts(_trunk_mlp, "up_proj.weight_scale", "gate_proj.weight_scale"),
- transform=_expert_fc1_s_cd,
- layers=_cute_dsl_copy_layers,
- ),
- W(
- "fc2_w_cd",
- shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // 2),
- dtype=u8,
- src=_experts(_trunk_mlp, "down_proj.weight"),
- layers=_cute_dsl_copy_layers,
- ),
- W(
- "fc2_s_cd",
- shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // SF_VEC),
- dtype=u8,
- src=_experts(_trunk_mlp, "down_proj.weight_scale"),
- transform=_expert_fc2_s_cd,
- layers=_cute_dsl_copy_layers,
- ),
- *(
- W(name, shape=lambda d: (d.num_experts,), dtype=f32, src=_experts(_trunk_mlp, suffix))
- for name, suffix in _EXPERT_SCALARS
- ),
- *_shared_expert("", _trunk_mlp, ALL_LAYERS),
- # --- not per-layer: the stack-final mixer, which is also the final norm ---
- *_mixer("fhc", f"{_LM}.hyper_connection_mixer", None),
- W(
- "embed",
- shape=lambda d: (d.vocab, d.hidden),
- src=f"{_LM}.embed_tokens.weight",
- layers=None,
- ),
- # --- PLE, on its one decoder layer (the n-gram table itself is declared at load) ---
- W(
- "ple_key_w",
- shape=lambda d: (d.stream, d.hidden),
- src=_ple("key_proj.weight"),
- layers=None,
- ),
- W(
- "ple_val_w",
- shape=lambda d: (d.ple_embed_dim, d.hidden),
- src=_ple("value_proj.weight"),
- layers=None,
- ),
- W("ple_nk", shape=lambda d: (d.stream,), src=_ple("norm_key.weight"), layers=None),
- W("ple_nq", shape=lambda d: (d.stream,), src=_ple("norm_query.weight"), layers=None),
- W("ple_nc", shape=lambda d: (d.stream,), src=_ple("norm_conv.weight"), layers=None),
- W(
- "ple_conv_w",
- shape=lambda d: (d.stream, 1, d.ple_conv_kernel),
- src=_ple("conv1d.weight"),
- layers=None,
- ),
- W(
- "ple_mult",
- shape=lambda d: (d.ngram_size,),
- dtype=i64,
- src=_ple("ple_embedding.layer_multipliers"),
- layers=None,
- ),
- W(
- "ple_vocab",
- shape=lambda d: (d.ngram_heads,),
- dtype=i64,
- src=_ple("ple_embedding.ngram_heads_vocab_sizes"),
- layers=None,
- ),
- W(
- "ple_offs",
- shape=lambda d: (d.ngram_heads,),
- dtype=i64,
- src=_ple("ple_embedding.ngram_heads_offsets"),
- layers=None,
- ),
- # The table's per-tensor scale sits with the **table**, one level below the three
- # metadata tensors above: the shards and this scale are the quantized
- # `ngram_embedding` submodule, while the multipliers, offsets and vocabulary sizes
- # describe the id space and stay on `ple_embedding` itself. The two groups also
- # live in different safetensors files, which is the other way to notice the split.
- W(
- "ple_tscale",
- shape=lambda d: (1,),
- src=lambda d, i: [(f"{_ngram_prefix(d.ple_layer)}.weight_scale", None)],
- layers=None,
- ),
- # --- the vision tower: all bf16, one-to-one, no transform ---
- W(
- "vis_patch_w",
- shape=lambda d: (d.v_hidden, d.v_channels, d.v_tpatch, d.v_patch, d.v_patch),
- src=_visual("patch_embed.proj.weight"),
- layers=None,
- ),
- W(
- "vis_patch_b",
- shape=lambda d: (d.v_hidden,),
- src=_visual("patch_embed.proj.bias"),
- layers=None,
- ),
- W(
- "vis_pos",
- shape=lambda d: (d.v_pos_rows, d.v_hidden),
- src=_visual("pos_embed.weight"),
- layers=None,
- ),
- # LayerNorm **with a bias**, unlike every norm in the language tower.
- W(
- "vis_n1_w",
- shape=lambda d: (d.v_hidden,),
- src=_vision_block("norm1.weight"),
- layers=_vision_blocks,
- ),
- W(
- "vis_n1_b",
- shape=lambda d: (d.v_hidden,),
- src=_vision_block("norm1.bias"),
- layers=_vision_blocks,
- ),
- # One fused `[q | k | v]` projection, MHA, with a bias.
- W(
- "vis_qkv_w",
- shape=lambda d: (3 * d.v_hidden, d.v_hidden),
- src=_vision_block("attn.qkv.weight"),
- layers=_vision_blocks,
- ),
- W(
- "vis_qkv_b",
- shape=lambda d: (3 * d.v_hidden,),
- src=_vision_block("attn.qkv.bias"),
- layers=_vision_blocks,
- ),
- W(
- "vis_proj_w",
- shape=lambda d: (d.v_hidden, d.v_hidden),
- src=_vision_block("attn.proj.weight"),
- layers=_vision_blocks,
- ),
- W(
- "vis_proj_b",
- shape=lambda d: (d.v_hidden,),
- src=_vision_block("attn.proj.bias"),
- layers=_vision_blocks,
- ),
- W(
- "vis_n2_w",
- shape=lambda d: (d.v_hidden,),
- src=_vision_block("norm2.weight"),
- layers=_vision_blocks,
- ),
- W(
- "vis_n2_b",
- shape=lambda d: (d.v_hidden,),
- src=_vision_block("norm2.bias"),
- layers=_vision_blocks,
- ),
- # A **non-gated** MLP: nothing here is split in half, unlike every SwiGLU in the
- # language tower.
- W(
- "vis_fc1_w",
- shape=lambda d: (d.v_inter, d.v_hidden),
- src=_vision_block("mlp.linear_fc1.weight"),
- layers=_vision_blocks,
- ),
- W(
- "vis_fc1_b",
- shape=lambda d: (d.v_inter,),
- src=_vision_block("mlp.linear_fc1.bias"),
- layers=_vision_blocks,
- ),
- W(
- "vis_fc2_w",
- shape=lambda d: (d.v_hidden, d.v_inter),
- src=_vision_block("mlp.linear_fc2.weight"),
- layers=_vision_blocks,
- ),
- W(
- "vis_fc2_b",
- shape=lambda d: (d.v_hidden,),
- src=_vision_block("mlp.linear_fc2.bias"),
- layers=_vision_blocks,
- ),
- # The patch merger. Its norm is applied **before** the 4-to-1 shuffle, so it is
- # `hidden` wide and not `merge_dim` wide.
- W(
- "vis_m_n_w",
- shape=lambda d: (d.v_hidden,),
- src=_visual("merger.norm.weight"),
- layers=None,
- ),
- W("vis_m_n_b", shape=lambda d: (d.v_hidden,), src=_visual("merger.norm.bias"), layers=None),
- W(
- "vis_m_fc1_w",
- shape=lambda d: (d.v_merge_dim, d.v_merge_dim),
- src=_visual("merger.linear_fc1.weight"),
- layers=None,
- ),
- W(
- "vis_m_fc1_b",
- shape=lambda d: (d.v_merge_dim,),
- src=_visual("merger.linear_fc1.bias"),
- layers=None,
- ),
- W(
- "vis_m_fc2_w",
- shape=lambda d: (d.hidden, d.v_merge_dim),
- src=_visual("merger.linear_fc2.weight"),
- layers=None,
- ),
- W(
- "vis_m_fc2_b",
- shape=lambda d: (d.hidden,),
- src=_visual("merger.linear_fc2.bias"),
- layers=None,
- ),
- # --- the MTP draft module, at layer index num_hidden_layers, only when live ---
- #
- # Structurally one more decoder layer of this architecture with a two-branch
- # front end; every shape equals a trunk full-attention layer's except the front
- # end and the expert quantization (FP8 128x128 block scales, not NVFP4).
- W(
- "mtp_pre_e_norm",
- shape=lambda d: (d.hidden,),
- src=_mtp_key(f"{_MTP}.pre_fc_norm_embedding.weight"),
- layers=_mtp_only,
- ),
- # `[hc_count * hidden]`: the module's input is the trunk's four-stream residual.
- W(
- "mtp_pre_h_norm",
- shape=lambda d: (d.stream,),
- src=_mtp_key(f"{_MTP}.pre_fc_norm_hidden.weight"),
- layers=_mtp_only,
- ),
- W(
- "mtp_fc_e",
- shape=lambda d: (d.hidden, d.hidden),
- src=_mtp_key(f"{_MTP}.fc_embedding.weight"),
- layers=_mtp_only,
- ),
- W(
- "mtp_fc_h",
- shape=lambda d: (d.hidden, d.hidden),
- src=_mtp_key(f"{_MTP}.fc_hidden.weight"),
- layers=_mtp_only,
- ),
- *(
- w
- for tag, module in _HC_MODULES
- for w in _hyper_connection(
- f"mtp_{tag}",
- lambda name, m=module: _mtp_key(f"{_MTP_LAYER}.{m}.{name}.weight"),
- _mtp_only,
- )
- ),
- *_gated_attention("mtp_", lambda s: _mtp_key(f"{_MTP_LAYER}.self_attn.{s}"), _mtp_only),
- *_index_branch(
- "mtp_", lambda s: _mtp_key(f"{_MTP_LAYER}.self_attn.indexer.{s}"), _mtp_only
- ),
- W(
- "mtp_router",
- shape=lambda d: (d.num_experts, d.hidden),
- src=_mtp_key(f"{_MTP_LAYER}.mlp.gate.weight"),
- layers=_mtp_only,
- ),
- # e4m3 codes stored row-major with one fp32 scale per 128x128 weight block, the
- # FC1 halves concatenated `[up ; gate]` on the row axis.
- W(
- "mtp_fc1_w",
- shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden),
- dtype=f8,
- src=_experts(_mtp_mlp, "up_proj.weight", "gate_proj.weight"),
- transform=_mtp_fc1_w,
- layers=_mtp_only,
- ),
- W(
- "mtp_fc1_s",
- shape=lambda d: (
- d.num_experts,
- (2 * d.moe_inter) // FP8_BLOCK,
- d.hidden // FP8_BLOCK,
- ),
- dtype=f32,
- src=_experts(_mtp_mlp, "up_proj.weight_scale_inv", "gate_proj.weight_scale_inv"),
- transform=_mtp_fc1_s,
- layers=_mtp_only,
- ),
- # FC2 needs no relayout at all: the checkpoint's `[hidden, I]` e4m3 block is
- # exactly the operand.
- W(
- "mtp_fc2_w",
- shape=lambda d: (d.num_experts, d.hidden, d.moe_inter),
- dtype=f8,
- src=_experts(_mtp_mlp, "down_proj.weight"),
- layers=_mtp_only,
- ),
- W(
- "mtp_fc2_s",
- shape=lambda d: (d.num_experts, d.hidden // FP8_BLOCK, d.moe_inter // FP8_BLOCK),
- dtype=f32,
- src=_experts(_mtp_mlp, "down_proj.weight_scale_inv"),
- transform=_mtp_fc2_s,
- layers=_mtp_only,
- ),
- *_shared_expert("mtp_", _mtp_mlp, _mtp_only),
- # The module's own stack-final mixer, which is its `shared_head`.
- *_mixer("mtp_fhc", f"{_MTP}.hyper_connection_mixer", _mtp_only),
- )
-
- def dims(self, core) -> SimpleNamespace:
- """Shapes are read off the normalized text config and the composite's vision
- config, which the core holds as `vision_config` before it declares anything."""
- cfg = core.model_config.pretrained_config
- vcfg = core.vision_config
- hidden = cfg.hidden_size
- index_heads, index_kv_heads = cfg.indexer_n_heads, cfg.indexer_kv_heads
- num_k_heads, num_v_heads = cfg.linear_num_key_heads, cfg.linear_num_value_heads
- key_dim = num_k_heads * cfg.linear_key_head_dim
- value_dim = num_v_heads * cfg.linear_value_head_dim
- ngram_heads = cfg.heads_per_ngram * (cfg.ngram_size - 1)
- v_hidden = vcfg.hidden_size
- return SimpleNamespace(
- num_layers=cfg.num_hidden_layers,
- is_attn=[t == "full_attention" for t in cfg.layer_types],
- hidden=hidden,
- hc_count=cfg.hc_count,
- hc_lowrank=cfg.hc_lowrank,
- stream=cfg.hc_count * hidden,
- head_dim=cfg.head_dim,
- q_width=cfg.num_attention_heads * cfg.head_dim,
- kv_width=cfg.num_key_value_heads * cfg.head_dim,
- index_head_dim=cfg.indexer_head_dim,
- index_width=(index_heads + index_kv_heads) * cfg.indexer_head_dim,
- # [q | k | v | gate | index q, k | pad of one index head]
- attn_in_rows=2 * cfg.num_attention_heads * cfg.head_dim
- + 2 * cfg.num_key_value_heads * cfg.head_dim
- + (index_heads + index_kv_heads + 1) * cfg.indexer_head_dim,
- num_v_heads=num_v_heads,
- head_v_dim=cfg.linear_value_head_dim,
- conv_dim=2 * key_dim + value_dim,
- value_dim=value_dim,
- in_proj_split=(2 * key_dim + value_dim, value_dim, num_v_heads, num_v_heads),
- conv_kernel=cfg.linear_conv_kernel_dim,
- num_experts=cfg.num_experts,
- moe_inter=cfg.moe_intermediate_size,
- shared_inter=cfg.shared_expert_intermediate_size,
- vocab=cfg.vocab_size,
- ple_layer=list(cfg.ple_layer_ids)[0] - 1, # one-based in the config
- ple_embed_dim=cfg.ple_embed_dim,
- ple_conv_kernel=cfg.ple_conv_kernel_size,
- ngram_size=cfg.ngram_size,
- ngram_heads=ngram_heads,
- v_depth=vcfg.depth,
- v_hidden=v_hidden,
- v_inter=vcfg.intermediate_size,
- v_channels=vcfg.in_channels,
- v_tpatch=vcfg.temporal_patch_size,
- v_patch=vcfg.patch_size,
- v_pos_rows=vcfg.num_position_embeddings,
- v_merge_dim=v_hidden * vcfg.spatial_merge_size * vcfg.spatial_merge_size,
- mtp_enabled=getattr(core.model_config, "spec_config", None) is not None,
- trunk_moe_backend=core.trunk_moe_backend,
- moe_prefill_copy=core.moe_prefill_copy,
- dtype=core.model_config.torch_dtype,
- )
-
- def expected_unconsumed(self, core, weights) -> set:
- """With MTP off, the whole `mtp.` subtree; with it on, nothing, so that any
- unloaded `mtp.*` key fails the coverage assert instead of hiding."""
- if core.mtp_enabled:
- return set()
- return {k for k in weights if k.startswith(f"{_MTP}.")}
-
- def declare(self, d: SimpleNamespace) -> nn.ParameterDict:
- """Allocate every weight the table declares.
-
- Meta-init intercepts `torch.empty` here -- real CUDA storage arrives
- when the engine materializes the registry.
- """
- w = nn.ParameterDict()
- for entry in self.WEIGHTS:
- dtype = entry.dtype or d.dtype
- for key, _, _ in self._targets(entry, d):
- w[key] = nn.Parameter(
- torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False
- )
- return w
-
- def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None, str]]:
- """`[(param key, layer index | None, checkpoint layer prefix), ...]`.
-
- One place decides which parameters an entry declares, so `declare` and
- `manifest` cannot disagree about whether a weight exists.
- """
- layers = entry.layer_indices(d)
- if layers is None:
- return [(entry.name, None, "")]
- return [(f"l{i}_{entry.name}", i, f"{_LM}.layers.{i}") for i in layers]
-
- def manifest(self, core) -> dict[str, list[tuple[Any, Any, Callable | None]]]:
- """param key -> [(checkpoint key, index into the param | None, transform | None)]."""
- d = self.dims(core)
- rows: dict[str, list[tuple[Any, Any, Callable | None]]] = {}
- for entry in self.WEIGHTS:
- for key, i, prefix in self._targets(entry, d):
- rows[key] = [
- (ckpt, index, entry.transform) for ckpt, index in entry.sources(d, i, prefix)
- ]
- return rows
-
- def load(self, model, weights) -> None:
- """Fill every declared parameter. The n-gram table is declared here first,
- multi-key rows feed their transform several tensors at once, the per-tensor
- scalars are 0-dim and reject `[:]`, and coverage is asserted in both
- directions."""
- core = model.model
- # The one load-time declaration, before the manifest is built: the table's
- # height comes from the shards, and the manifest addresses it by row range.
- shard_keys = _declare_ngram_table(core, weights)
- manifest = self.manifest(core)
- manifest["ple_table"] = [
- (key, (slice(s * core.ple_rows_per_shard, (s + 1) * core.ple_rows_per_shard),), None)
- for s, key in enumerate(shard_keys)
- ]
- consumed: set = set()
-
- def fill(param: nn.Parameter, ckpt_key, index, transform) -> None:
- keys = ckpt_key if isinstance(ckpt_key, tuple) else (ckpt_key,)
- for key in keys:
- assert key in weights, f"checkpoint key missing: {key}"
- dst = param.data if index is None else param.data[index]
- # Materialize where the destination lives: the expert relayouts are row
- # gathers and the scale swizzle is a kernel, both on the device.
- src_all = tuple(
- _materialize(weights[key]).to(dst.device, non_blocking=True) for key in keys
- )
- if transform is not None:
- src = transform(src_all if isinstance(ckpt_key, tuple) else src_all[0], core)
- else:
- assert len(src_all) == 1, "a multi-key row needs a transform"
- src = src_all[0]
- assert dst.shape == src.shape, (ckpt_key, tuple(dst.shape), tuple(src.shape))
- assert src.dtype == dst.dtype, (ckpt_key, src.dtype, dst.dtype)
- dst.copy_(src, non_blocking=True)
- consumed.update(keys)
-
- assert set(manifest.keys()) == set(core.w.keys()), (
- "manifest/parameter drift",
- set(manifest.keys()) ^ set(core.w.keys()),
- )
- for param_key, sources in manifest.items():
- for ckpt_key, index, transform in sources:
- fill(core.w[param_key], ckpt_key, index, transform)
- if param_key.endswith(_RELEASE_AFTER):
- torch.cuda.empty_cache()
-
- # Shell-registered exception: the base class owns lm_head (untied), and this
- # checkpoint leaves it bf16, so it is a single dense weight. Asserted rather than
- # adapted — a shape change is a different logits path.
- assert tuple(model.lm_head.weight.shape) == (core.vocab, core.hidden), (
- f"the shell built lm_head {tuple(model.lm_head.weight.shape)} instead "
- f"of {(core.vocab, core.hidden)}"
- )
- fill(model.lm_head.weight, "lm_head.weight", None, None)
- # Parameter-side coverage, the other direction: name every registered parameter
- # outside the target's own ParameterDict instead of assuming lm_head is the only
- # one.
- shell = {n for n, _ in model.named_parameters() if not n.startswith("model.w.")}
- assert shell == {"lm_head.weight"}, f"unfed shell parameters: {sorted(shell)}"
-
- torch.cuda.synchronize()
- leftover = set(weights.keys()) - consumed
- expected = self.expected_unconsumed(core, weights)
- what = "empty" if core.mtp_enabled else "the MTP module"
- assert leftover == expected, (
- f"checkpoint coverage: leftover keys are not exactly {what}; "
- f"unexpected {sorted(leftover - expected)[:8]}, "
- f"missed {sorted(expected - leftover)[:8]}"
- )
-
-
-MODEL_WEIGHTS = Qwen4ExpTp1Weights()
-
-
-def _declare_ngram_table(core, weights) -> list[str]:
- """Register the n-gram table with the height its shards concatenate to, and return
- the shard keys in concatenation order.
-
- The shards are **contiguous row chunks**, not an interleaving: global row `r` lives in
- shard `r // rows_per_shard` at local row `r % rows_per_shard`. That is what makes the
- flat buffer indexable by the hash's global rows.
-
- **This is the whole of PLE host offload.** `core.ple_host_offload`, resolved at
- construction from `TRTLLM_MODELING_V2_PLE_HOST_OFFLOAD`, picks between an ordinary
- device allocation and host memory the device addresses. Nothing else changes: same
- shape, same dtype, same `ParameterDict` key, same rows filling it shard by shard, same
- gather in the forward. An offloaded tensor is **loaded, not skipped** — it stays
- fully covered in both directions. What the manifest cannot record is *where* it
- landed, so the core asserts that separately against the driver.
- """
- prefix = _ngram_prefix(core.ple_layer)
- keys = [f"{prefix}.shard_{s}.weight" for s in range(core.ple_shards)]
- for key in keys:
- assert key in weights, f"checkpoint key missing: {key}"
- shapes = {tuple(weights[key].shape) for key in keys}
- assert len(shapes) == 1, (
- f"the n-gram table's {len(keys)} shards do not share one shape: "
- f"{sorted(shapes)}; they are concatenated as equal contiguous row "
- f"chunks and an uneven one would shift every row above it"
- )
- rows_per_shard, head_dim = shapes.pop()
- assert head_dim == core.ple_head_dim, (
- f"the n-gram table is {head_dim} wide but {core.ngram_heads} heads must "
- f"tile ple_embed_dim {core.ple_embed_dim}, i.e. {core.ple_head_dim}"
- )
- dtype = weights[keys[0]].dtype
- assert dtype == core.ple_table_dtype, (
- f"the n-gram table is stored {dtype}; this target declares "
- f"{core.ple_table_dtype} and does not convert it"
- )
- device = core.w["embed"].device
- rows = core.ple_shards * rows_per_shard
- if core.ple_host_offload:
- table, core._ple_table_host = _host_resident_empty(rows, head_dim, dtype, device)
- else:
- table = torch.empty(rows, head_dim, dtype=dtype, device=device)
- core._ple_table_host = None
- core.w["ple_table"] = nn.Parameter(table, requires_grad=False)
- # `nn.Parameter` is `_make_subclass`, which shares storage rather than copying — but
- # "rather than copying" is the entire feature here, so it is asserted instead of
- # assumed.
- assert core.w["ple_table"].data_ptr() == table.data_ptr(), (
- "wrapping the n-gram table in nn.Parameter moved its storage"
- )
- core.ple_rows_per_shard = rows_per_shard
- return keys
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py
index af8a50119783..cc8135690654 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py
@@ -2,12 +2,16 @@
# SPDX-License-Identifier: Apache-2.0
"""Routing for Qwen4ExpForConditionalGeneration.
-`route` compares the facts that fix which code runs against the one deployment the
-target was written for, and raises `NoTarget` stating them otherwise (`auto` falls
+`route` compares the facts that fix which code runs against the deployments the
+targets were written for, and raises `NoTarget` stating them otherwise (`auto` falls
back to the built-in implementation; `require` raises). The checkpoint's published
architecture name is kept by `load_pretrained_config`, so `route` registers under it.
All criteria are read off `text_config` / `vision_config`; the composite parent's
mirrored fields are unreliable.
+
+Qwen3.8-Flash-Next-NVFP4 on sm103 lands on `qwen3_8_flash_next_nvfp4__sm_103__default`
+at one of its parallel layouts (`layout`): tp1 on one rank, and tp4, tep4, dep4 or dep2 in
+one node.
"""
from __future__ import annotations
@@ -25,31 +29,54 @@
from tensorrt_llm._torch.model_config import ModelConfig
from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
+# The GPU the target is certified on.
+_SM = (10, 3)
+
# (layers, hidden, experts, attention heads, linear value heads, vocab). Shape-sniffing
# also matches fine-tunes; the accuracy gate pins identity.
_TEXT_SHAPE = (48, 2560, 512, 24, 48, 248320)
# (depth, hidden, heads, merge, out_hidden, num_pos, deepstack_indexes, act).
# `deepstack_visual_indexes == ()` is the key member: a deepstack sibling has a merger
-# list this target lacks.
+# list these targets lack.
_VISION_SHAPE = (27, 1152, 16, 2, 2560, 2304, (), "gelu_pytorch_tanh")
-# `moe_config.backend` requests the target implements (`_TRUNK_MOE_BACKENDS` in modeling.py).
-_MOE_BACKENDS = ("AUTO", "TRTLLM", "CUTEDSL")
-
-# The longest MTP draft the target runs (`_MAX_VERIFY_WIDTH - 1` in modeling.py); the legacy
-# per-step recurrent states (`_MAX_LEGACY_VERIFY_WIDTH - 1`) and the CUTEDSL experts
-# (`_CUTEDSL_MAX_DRAFT_LEN`) are certified at shorter ones.
+# The longest MTP draft the target runs (`_MAX_VERIFY_WIDTH - 1` in modeling.py).
_MAX_DRAFT_LEN = 7
-_MAX_LEGACY_DRAFT_LEN = 3
+
+# The engine choices the target implements at one rank (tp1), all in its modeling.py: the
+# `moe_config.backend` requests (`_TRUNK_MOE_BACKENDS`), the longest draft of the CUTEDSL
+# experts (`_CUTEDSL_MAX_DRAFT_LEN`) and the recurrent-state cache dtypes
+# (`_ONE_RANK_SSM_STATE_DTYPES`).
+_ONE_RANK_MOE_BACKENDS = ("AUTO", "TRTLLM", "CUTEDSL")
_MAX_CUTEDSL_DRAFT_LEN = 3
+_ONE_RANK_SSM_STATE_DTYPES = (torch.float32, torch.bfloat16)
+
+# Their counterparts at every other layout: no CUTEDSL experts (`_TRUNK_MOE_BACKENDS` less
+# CUTEDSL) and a bf16 recurrent state (`_SSM_STATE_DTYPES`).
+_MOE_BACKENDS = ("AUTO", "TRTLLM")
+_SSM_STATE_DTYPES = (torch.bfloat16,)
+
+# (world size, attention DP, MoE TP, MoE EP) -> the target's layout. Every layout spans its
+# whole world in one node, attention TP over all of it unless attention DP replicates
+# attention. One rank (tp1) holds everything. At four ranks the routed experts are split by
+# intermediate (tp4) or by expert (tep4, dep4). At two ranks only attention DP with two
+# 256-expert windows (dep2), where each rank's attention, Gated-DeltaNet and PLE geometry is
+# dep4's.
+_LAYOUTS = {
+ (1, False, 1, 1): "tp1",
+ (4, False, 4, 1): "tp4",
+ (4, False, 1, 4): "tep4",
+ (4, True, 1, 4): "dep4",
+ (2, True, 1, 2): "dep2",
+}
-# Recurrent-state cache dtypes the target serves (`_SSM_STATE_DTYPES` in modeling.py).
-_SSM_STATE_DTYPES = (torch.float32, torch.bfloat16)
+# The attention-DP layouts MTP is admitted at: those its draft acceptance was measured at.
+_ATTENTION_DP_MTP_LAYOUTS = ("dep4",)
class EngineChoices(NamedTuple):
- """Engine choices that fix which code the target runs but are not in the checkpoint's
+ """Engine choices that fix which code a target runs but are not in the checkpoint's
config, as the engine will resolve them."""
torch_dtype: torch.dtype # the dtype the engine runs the model in
@@ -60,7 +87,7 @@ class EngineChoices(NamedTuple):
def engine_choices(model_config: ModelConfig, llm_args: TorchLlmArgs) -> EngineChoices:
- # Imported on the claimed branch only, as the target is.
+ # Imported on a claimed branch only, as the targets are.
from tensorrt_llm._torch.pyexecutor.model_loader import effective_mamba_ssm_cache_dtype
from tensorrt_llm._torch.utils import is_gdn_replay_enabled
@@ -76,14 +103,57 @@ def engine_choices(model_config: ModelConfig, llm_args: TorchLlmArgs) -> EngineC
)
+def layout(mapping) -> Optional[str]:
+ """The target's layout ``mapping`` describes (``tp1``, ``tp4``, ``tep4``, ``dep4``,
+ ``dep2``), or None.
+
+ Pipeline and context parallelism, an attention TP narrower than the TP group, MoE
+ clusters, lm_head TP under attention DP and a group spanning nodes are other
+ forwards.
+ """
+ if mapping.tp_size != mapping.world_size or (mapping.pp_size, mapping.cp_size) != (1, 1):
+ return None
+ if mapping.is_multi_node() or mapping.moe_cluster_size != 1:
+ return None
+ if mapping.enable_attention_dp:
+ if mapping.enable_lm_head_tp_in_adp:
+ return None
+ elif mapping.attn_tp_size != mapping.tp_size:
+ return None
+ return _LAYOUTS.get(
+ (
+ mapping.world_size,
+ mapping.enable_attention_dp,
+ mapping.moe_tp_size,
+ mapping.moe_ep_size,
+ )
+ )
+
+
+def _describe(mapping) -> str:
+ return (
+ f"world_size={mapping.world_size} (tp={mapping.tp_size}, pp={mapping.pp_size}, "
+ f"cp={mapping.cp_size}, attention_dp={mapping.enable_attention_dp}, "
+ f"moe_tp={mapping.moe_tp_size}, moe_ep={mapping.moe_ep_size})"
+ )
+
+
def _check_engine(
- spec_config: Optional[Any], kv_cache_quant_algo: Optional[Any], engine: EngineChoices
+ spec_config: Optional[Any],
+ kv_cache_quant_algo: Optional[Any],
+ engine: EngineChoices,
+ member: str,
) -> None:
- """Decline what the target's constructor or first forward would refuse.
+ """Decline what the target's constructor or first forward would refuse at layout
+ `member`: at tp1 the CuteDSL experts and the fp32 recurrent state are served too.
Without this ``auto`` claims the deployment and then fails mid-build instead of
falling back to the built-in implementation.
"""
+ if member == "tp1":
+ moe_backends, ssm_state_dtypes = _ONE_RANK_MOE_BACKENDS, _ONE_RANK_SSM_STATE_DTYPES
+ else:
+ moe_backends, ssm_state_dtypes = _MOE_BACKENDS, _SSM_STATE_DTYPES
if spec_config is not None and (
getattr(spec_config, "draft_len_schedule", None) is not None
or getattr(spec_config, "acceptance_rate_threshold", None) is not None
@@ -100,15 +170,16 @@ def _check_engine(
)
if engine.torch_dtype != torch.bfloat16:
raise NoTarget(f"the engine runs the model in {engine.torch_dtype}; the target is bf16")
- if str(engine.moe_backend).upper() not in _MOE_BACKENDS:
+ if str(engine.moe_backend).upper() not in moe_backends:
raise NoTarget(
- f"moe_config.backend={engine.moe_backend!r}; the target runs the trunk's routed "
- f"experts on {', '.join(_MOE_BACKENDS)}"
+ f"moe_config.backend={engine.moe_backend!r}; at {member} the target runs the "
+ f"trunk's routed experts on {', '.join(moe_backends)}"
)
- if engine.mamba_ssm_cache_dtype not in _SSM_STATE_DTYPES:
+ if engine.mamba_ssm_cache_dtype not in ssm_state_dtypes:
+ served = " or ".join(str(dtype).removeprefix("torch.") for dtype in ssm_state_dtypes)
raise NoTarget(
- f"a {engine.mamba_ssm_cache_dtype} recurrent-state cache; the target serves "
- "float32 or bfloat16 (kv_cache_config.mamba_ssm_cache_dtype)"
+ f"a {engine.mamba_ssm_cache_dtype} recurrent-state cache; at {member} the target "
+ f"serves {served} (kv_cache_config.mamba_ssm_cache_dtype)"
)
if engine.prefill_cuda_graph_backend == "piecewise":
raise NoTarget(
@@ -119,11 +190,11 @@ def _check_engine(
if spec_config is None:
return
draft = spec_config.max_draft_len
- if not engine.gdn_replay and draft > _MAX_LEGACY_DRAFT_LEN:
+ if not engine.gdn_replay:
raise NoTarget(
- f"max_draft_len={draft} with TRTLLM_USE_GDN_REPLAY=0: the legacy per-step recurrent "
- f"states are certified at max_draft_len up to {_MAX_LEGACY_DRAFT_LEN}; leave "
- "TRTLLM_USE_GDN_REPLAY unset for a longer draft"
+ "MTP with TRTLLM_USE_GDN_REPLAY=0: the target verifies its Gated-DeltaNet rows on "
+ "the cached replay only, not on the legacy per-step recurrent states; leave "
+ "TRTLLM_USE_GDN_REPLAY unset"
)
if str(engine.moe_backend).upper() == "CUTEDSL" and draft > _MAX_CUTEDSL_DRAFT_LEN:
raise NoTarget(
@@ -132,6 +203,28 @@ def _check_engine(
)
+def _check_attention_dp(member: str, spec_config: Optional[Any], engine: EngineChoices) -> None:
+ """Decline what the target does not run under attention DP (dep4, dep2), where
+ each rank serves its own requests and only the routed experts cross ranks.
+
+ MTP is admitted at dep4, where each rank drafts for its own requests and the draft layer's
+ FP8 expert window runs the trunk's attention-DP round trip. Routing admits both MTP and
+ the breakable prefill graph at the attention-TP layouts (tp4, tep4).
+ """
+ if spec_config is not None and member not in _ATTENTION_DP_MTP_LAYOUTS:
+ raise NoTarget(
+ f"speculative decoding with attention DP ({member}): no draft acceptance was "
+ "measured at this layout; MTP under attention DP is admitted at "
+ f"{', '.join(_ATTENTION_DP_MTP_LAYOUTS)}"
+ )
+ if engine.prefill_cuda_graph_backend == "breakable":
+ raise NoTarget(
+ f"prefill_cuda_graph_backend='breakable' with attention DP ({member}): the engine "
+ "pads every rank to one prefill bucket, so a rank holding only generation requests "
+ "is handed more rows than its batch, which the target's decode step refuses"
+ )
+
+
# Qwen3.8-Flash-Next publishes ``Qwen4ExpForConditionalGeneration``, and the
# config loader's composite branch writes the same name back.
@routes("Qwen4ExpForConditionalGeneration")
@@ -140,8 +233,8 @@ def route(
) -> type:
"""Return the target class for this deployment, or raise ``NoTarget`` saying why none fits.
- The target is imported where it is returned, so declining costs no import
- of modeling code. It also declines the engine choices (`EngineChoices`) the
+ A target is imported where it is returned, so declining costs no import of
+ modeling code. Routing also declines the engine choices (`EngineChoices`) the
target's constructor would refuse.
"""
pretrained_config = model_config.pretrained_config
@@ -193,31 +286,34 @@ def route(
and spec_config.max_total_draft_tokens == spec_config.max_draft_len
)
- # Qwen3.8-Flash-Next-NVFP4 on one GB300 (sm103), MTP-Eagle or no speculation.
- if (
- compute_capability == (10, 3)
- and shape == _TEXT_SHAPE
- and vshape == _VISION_SHAPE
- and mapping.world_size == 1
- and mtp_eagle_or_none
- ):
- # One line on purpose: the claims suite reads the routed targets off
- # `from ..modeling import ` lines.
+ checkpoint = compute_capability == _SM and shape == _TEXT_SHAPE and vshape == _VISION_SHAPE
+ # Qwen3.8-Flash-Next-NVFP4 on GB300s (sm103): on one at tp1, on four at tp4, tep4 or dep4,
+ # MTP-Eagle or no speculation; on two at dep2, without speculation.
+ member = layout(mapping)
+ if checkpoint and member is not None and mtp_eagle_or_none:
+ engine = engine_choices(model_config, llm_args)
+ if mapping.enable_attention_dp:
+ _check_attention_dp(member, spec_config, engine)
+ _check_engine(spec_config, kv_cache_quant_algo, engine, member)
+ # One line on purpose: the claims suite reads the routed target off the
+ # `from ..modeling import ` line.
# fmt: off
# isort: off
- from .qwen3_8_flash_next_nvfp4__sm_103__tp1.modeling import ModelingV2Qwen38FlashNextNvfp4Sm103Tp1
+ from .qwen3_8_flash_next_nvfp4__sm_103__default.modeling import ModelingV2Qwen38FlashNextNvfp4Sm103Default
# isort: on
# fmt: on
- _check_engine(spec_config, kv_cache_quant_algo, engine_choices(model_config, llm_args))
- return ModelingV2Qwen38FlashNextNvfp4Sm103Tp1
+ return ModelingV2Qwen38FlashNextNvfp4Sm103Default
raise NoTarget(
f"compute capability {compute_capability}, text shape {shape}, vision shape "
- f"{vshape}, {mapping.world_size} ranks, speculative decoding "
+ f"{vshape}, {_describe(mapping)}, speculative decoding "
f"{spec_config.spec_dec_mode.name if spec_config is not None else 'off'}"
f"{f' max_draft_len={spec_config.max_draft_len}' if spec_config is not None else ''}, "
- f"KV cache {kv_cache_quant_algo or 'bf16'}; "
- "the one qwen4_exp target is Qwen3.8-Flash-Next-NVFP4 on sm103, tp1, with "
- f"one-model MTP-Eagle (max_draft_len 1..{_MAX_DRAFT_LEN}) or no speculative decoding"
+ f"KV cache {kv_cache_quant_algo or 'bf16'}; the qwen4_exp target is "
+ "Qwen3.8-Flash-Next-NVFP4 on sm103 in one node: at 1 rank (tp1), and at 4 ranks tp4 "
+ "(attention TP4, experts split by intermediate), tep4 (attention TP4, experts EP4) or "
+ f"dep4 (attention DP, experts EP4), with one-model MTP-Eagle (max_draft_len "
+ f"1..{_MAX_DRAFT_LEN}) or no speculative decoding; at 2 ranks dep2 (attention DP, "
+ "experts EP2) without speculative decoding"
)
diff --git a/tests/integration/defs/.test_durations b/tests/integration/defs/.test_durations
index 7b596be823f1..d86b7b84f838 100644
--- a/tests/integration/defs/.test_durations
+++ b/tests/integration/defs/.test_durations
@@ -20,12 +20,12 @@
"accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp4::test_gsm8k": 1247.1575,
"accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp8::test_eagle3_gsm8k_and_acceptance": 998.465,
"accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp8::test_gsm8k": 317.73875,
- "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_breakable_prefill_matches_eager": 730.6153846153846,
- "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_gsm8k": 643.2887001455605,
- "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_gsm8k_mmmu_and_video": 969.7372977443608,
- "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_cutedsl_fp8_kv_bf16_state": 658.6539230769231,
- "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_long_context_passkey": 492.2441946022727,
- "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp7_long_context_passkey": 537.4055384615385,
+ "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_breakable_prefill_matches_eager": 730.6153846153846,
+ "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_gsm8k": 643.2887001455605,
+ "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp1_gsm8k_mmmu_and_video": 969.7372977443608,
+ "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp3_cutedsl_fp8_kv_bf16_state": 658.6539230769231,
+ "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp3_gsm8k_and_long_context": 492.2441946022727,
+ "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp7_long_context_passkey": 537.4055384615385,
"accuracy/test_disaggregated_serving.py::TestDeepSeekR1::test_kv_cache_v2_nixl_python[cache_mgr_v2]": 606.7294444444445,
"accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=0-overlap_scheduler=False]": 348.56913636363635,
"accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=0-overlap_scheduler=True]": 207.74877272727272,
@@ -1598,20 +1598,6 @@
"test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_rank_agreement.py]": 23.462214285714285,
"unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py": 23.585939655172414,
"test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py]": 23.585939655172414,
- "unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_kv_quant.py": 29.04516883116883,
- "test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_kv_quant.py]": 29.04516883116883,
- "unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_moe_scalars.py": 25.506426666666666,
- "test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_moe_scalars.py]": 25.506426666666666,
- "unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_multimodal_embeds.py": 25.56098734177215,
- "test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_multimodal_embeds.py]": 25.56098734177215,
- "unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_pool_contract.py": 25.537,
- "test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_pool_contract.py]": 25.537,
- "unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_thop_attention_calls.py": 23.835025,
- "test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_thop_attention_calls.py]": 23.835025,
- "unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_validate_request.py": 25.437897435897437,
- "test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_validate_request.py]": 25.437897435897437,
- "unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_weight_table.py": 32.944,
- "test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_weight_table.py]": 32.944,
"unittest/_torch/modeling_v2/moe/test_modeling_v2_mxe4m3_mxe2m1_block_scale_moe_runner.py": 24.025819444444444,
"test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/moe/test_modeling_v2_mxe4m3_mxe2m1_block_scale_moe_runner.py]": 24.025819444444444,
"unittest/_torch/modeling_v2/moe/test_modeling_v2_noaux_tc_op.py::test_kimi_k3_config": 10.210739463601533,
diff --git a/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py b/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py
index 11761258ad08..a8edd1cdd25a 100644
--- a/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py
+++ b/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py
@@ -14,9 +14,16 @@
# limitations under the License.
"""Whole-model gates for the qwen4_exp modeling_v2 target.
-Every test passes ``modeling_v2="require"``: under ``"auto"`` a configuration
-that missed the target's routing would fall back to the built-in model, pass,
-and report the built-in's numbers as the target's.
+``TestModelingV2Qwen38FlashNextNvfp4Sm103Default``: Qwen3.8-Flash-Next-NVFP4 on sm_103 at every
+layout the target serves: one rank (tp1, the ``test_tp1_*`` cases), four (tp4, tep4, dep4) and
+two (dep2).
+
+Each test runs one engine (two for the breakable-graph comparison) and checks every feature
+that engine carries, so a feature is covered without a test of its own. Every test passes
+``modeling_v2="require"``: under ``"auto"`` a configuration that missed the target's routing
+would fall back to the built-in model, pass, and report the built-in's numbers as the
+target's. The checkpoint is also served disaggregated, the target on both sides of the hop:
+one rank per engine with the AgentX engines, and dep2 on both.
"""
import random
@@ -29,7 +36,7 @@
from tensorrt_llm import LLM
from tensorrt_llm._torch._experimental.modeling_v2 import MODELING_V2_VALIDATE_ENV
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501
modeling as target_modeling,
)
from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph_runner import BreakableCUDAGraphRunner
@@ -50,6 +57,7 @@
QSASparseAttentionConfig,
SamplingParams,
)
+from tensorrt_llm.quantization import QuantAlgo
from ...conftest import llm_models_root
from ..accuracy_core import (
@@ -59,6 +67,7 @@
assert_acceptance_length,
compute_acceptance_length,
)
+from ..test_disaggregated_serving import launch_disaggregated_llm, run_accuracy_test
from ._env import scoped_env_overrides
# The target is certified on sm_103 only.
@@ -70,6 +79,32 @@
# engine prepared; the flag has to reach the worker rank.
_VALIDATE = {MODELING_V2_VALIDATE_ENV: "1"}
+# Every worker of a disaggregated test: the step-contract check, and what a worker spawning
+# its second rank over MPI needs in a container running as root (docker/Dockerfile.multi).
+_DISAGG_WORKER_ENV = {
+ **_VALIDATE,
+ "PRTE_ALLOW_RUN_AS_ROOT": "1",
+ "PRTE_ALLOW_RUN_AS_ROOT_CONFIRM": "1",
+}
+
+
+def _disagg_server_config(ctx_router=None, gen_router=None) -> dict:
+ """One context and one generation server on this host, with optional router types."""
+
+ def servers(port, router):
+ config = {"num_instances": 1, "urls": [f"localhost:{port}"]}
+ if router is not None:
+ config["router"] = {"type": router}
+ return config
+
+ return {
+ "hostname": "localhost",
+ "port": 8000,
+ "backend": "pytorch",
+ "context_servers": servers(8001, ctx_router),
+ "generation_servers": servers(8002, gen_router),
+ }
+
def _target_core(llm: LLM):
"""The target's core module.
@@ -306,6 +341,21 @@ def _synthetic_image(color: tuple, size: int = 256) -> Image.Image:
return image
+def _mmmu(llm, model_name: str) -> None:
+ """MMMU on ``llm`` against the reference row the engine resolves."""
+ # MMMU adds <|endoftext|> to the stop token, as the built-in model's MMMU test does.
+ sampling_params = SamplingParams(
+ max_tokens=MMMU.MAX_OUTPUT_LEN,
+ truncate_prompt_tokens=MMMU.MAX_INPUT_LEN,
+ stop="<|endoftext|>",
+ )
+ MMMU(model_name).evaluate(
+ llm,
+ sampling_params=sampling_params,
+ extra_evaluator_kwargs=dict(chat_template_kwargs=dict(enable_thinking=False)),
+ )
+
+
def _media_request(tokenizer, processor, text: str, items: list[tuple[str, object]]) -> dict:
"""One request carrying ``items``, ``(modality, data)`` pairs in prompt order.
@@ -331,11 +381,37 @@ def _media_request(tokenizer, processor, text: str, items: list[tuple[str, objec
return dict(prompt=prompt, multi_modal_data=media, mm_item_order=tracker.item_order())
-class TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1(LlmapiAccuracyTestHarness):
- """qwen3.8-flash-next-nvfp4 / sm_103 / tp1: text, MTP and multimodal."""
+# The layouts routing admits, as LLM API parallel arguments. Each cuts the weights
+# differently and crosses ranks differently, and a wrong cut still produces fluent text, so
+# every layout runs one gsm8k against an anchor: tep4, tp4 and dep4 below, dep2 in the
+# disaggregated case.
+# Attention split by head over four ranks, the routed experts by expert.
+_TEP4 = dict(tensor_parallel_size=4, moe_expert_parallel_size=4, enable_attention_dp=False)
+# Attention split by head, every expert split by intermediate (160 of 640 rows a rank).
+_TP4 = dict(
+ tensor_parallel_size=4,
+ moe_tensor_parallel_size=4,
+ moe_expert_parallel_size=1,
+ enable_attention_dp=False,
+)
+# Attention replicated, each rank serving its own requests; the routed experts by expert.
+_DEP4 = dict(tensor_parallel_size=4, moe_expert_parallel_size=4, enable_attention_dp=True)
+
+# The quantized KV pools, one layout each: NVFP4 at tp4 with the breakable prefill graph
+# (routed at attention TP only), FP8 at dep4 (MTP under attention DP, each rank drafting for
+# its own requests).
+_QUANTIZED_KV = [
+ pytest.param(_TP4, "nvfp4", True, id="tp4_nvfp4_kv_breakable"),
+ pytest.param(_DEP4, "fp8", False, id="dep4_fp8_kv"),
+]
+
+
+class TestModelingV2Qwen38FlashNextNvfp4Sm103Default(LlmapiAccuracyTestHarness):
+ """qwen3.8-flash-next-nvfp4 / sm_103 / default: tp1, tp4, tep4, dep4 and dep2."""
# References are keyed by model name plus the quantization the engine
- # resolves for this checkpoint (MIXED_PRECISION).
+ # resolves for this checkpoint (MIXED_PRECISION), and the MTP row by the
+ # speculative mode as well.
MODEL_NAME = "Qwen/Qwen3.8-Flash-Next"
MODEL_PATH = f"{llm_models_root()}/Qwen3.8-Flash-Next-NVFP4"
@@ -351,17 +427,60 @@ class TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1(LlmapiAccuracyTestHarness):
# The anchor's `max_tokens`; the class default (256) truncates this checkpoint's answers.
GSM8K_MAX_OUTPUT_LEN = 512
- # `mamba_ssm_cache_dtype` is float32 because the checkpoint declares
+ # The checkpoint's bf16 K/V, a bf16 recurrent state (the one the target serves), no
+ # block reuse.
+ KV_CACHE = KvCacheConfig(
+ free_gpu_memory_fraction=0.6,
+ enable_block_reuse=False,
+ mamba_ssm_cache_dtype="bfloat16",
+ )
+
+ # The routed experts run on trtllm-gen, the MoE backend the target implements at every
+ # layout (CuteDSL is opt-in at tp1). It is also what the draft layer's FP8 block-scale
+ # experts need: the default CUTLASS MoE does not support them on sm_103.
+ MOE = MoeConfig(backend="TRTLLM")
+
+ ENGINE = dict(
+ max_seq_len=9216,
+ max_batch_size=128,
+ max_num_tokens=8192,
+ enable_chunked_prefill=False,
+ kv_cache_config=KV_CACHE,
+ cuda_graph_config=CudaGraphConfig(enable_padding=True, max_batch_size=128),
+ moe_config=MOE,
+ )
+
+ # One-model MTP-Eagle with the checkpoint's single MTP module replayed three times, on
+ # the engine the MTP configuration is served with (a 16384-token sequence bound).
+ MTP3 = MTPDecodingConfig(max_draft_len=3)
+ ENGINE_MTP = dict(ENGINE, max_seq_len=16384)
+
+ # The breakable prefill CUDA graph, up to max_num_tokens.
+ BREAKABLE_PREFILL = dict(
+ prefill_cuda_graph_backend="breakable",
+ prefill_capture_num_tokens=[512, 1024, 2048, 4096, 8192],
+ )
+
+ # A collapsed or miscomputed draft path accepts nothing and reads 1.0; every draft
+ # accepted reads 1 + max_draft_len. The midpoint of that range is the gsm8k MTP
+ # cases' tripwire: at or above it, more than half the drafts are accepted.
+ MTP3_COLLAPSE_FLOOR = 1.0 + MTP3.max_draft_len / 2
+
+ # The pass keys' acceptance-length floor in references/acceptance_length.yaml.
+ ACCEPTANCE_KEY_LONG_CONTEXT = "ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3_long_context"
+
+ # ---- one rank (tp1) ----------------------------------------------------------------
+ # At tp1 `mamba_ssm_cache_dtype` is float32 because the checkpoint declares
# `mamba_ssm_dtype: float32` (the engine default is bfloat16); the NVFP4- and fp8-KV cases
# set it to bfloat16, what `auto` resolves to for this checkpoint. `dtype` is left at the
# checkpoint's bf16 KV, which is what the anchors were measured at; only the fp8 and NVFP4
# KV cases set it.
- KV_CACHE = KvCacheConfig(
+ TP1_KV_CACHE = KvCacheConfig(
free_gpu_memory_fraction=0.6,
enable_block_reuse=True,
mamba_ssm_cache_dtype="float32",
)
- KV_CACHE_FP8_BF16_STATE = KvCacheConfig(
+ TP1_KV_CACHE_FP8_BF16_STATE = KvCacheConfig(
dtype="fp8",
free_gpu_memory_fraction=0.6,
enable_block_reuse=True,
@@ -369,113 +488,82 @@ class TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1(LlmapiAccuracyTestHarness):
)
# Block reuse is on, which (with chunked prefill and the MTP drafts) is what lets the
# attention op read an NVFP4 pool: its kernels exist for the paged-context path only.
- KV_CACHE_NVFP4_BF16_STATE = KvCacheConfig(
+ TP1_KV_CACHE_NVFP4_BF16_STATE = KvCacheConfig(
dtype="nvfp4",
free_gpu_memory_fraction=0.6,
enable_block_reuse=True,
mamba_ssm_cache_dtype="bfloat16",
)
# Block reuse is off for the vision path, which requires it.
- KV_CACHE_MULTIMODAL = KvCacheConfig(
+ TP1_KV_CACHE_MULTIMODAL = KvCacheConfig(
free_gpu_memory_fraction=0.6,
enable_block_reuse=False,
mamba_ssm_cache_dtype="float32",
)
- # Required, not a tuning choice: the MTP draft layer's FP8 block-scale
- # experts are not supported by the default CUTLASS MoE on sm_103. Set for
- # every case so they share one engine configuration.
- MOE = MoeConfig(backend="TRTLLM")
-
- ENGINE = dict(
+ # One rank: chunked prefill over an 8192-token sequence bound, 32 requests.
+ TP1_ENGINE = dict(
max_seq_len=8192,
max_batch_size=32,
enable_chunked_prefill=True,
- kv_cache_config=KV_CACHE,
+ kv_cache_config=TP1_KV_CACHE,
cuda_graph_config=CudaGraphConfig(),
moe_config=MOE,
)
- # The breakable prefill CUDA graph is opt-in; these buckets cover the prompts of
- # `test_breakable_prefill_matches_eager`.
- BREAKABLE_PREFILL = dict(
+ # These buckets cover the prompts of `test_tp1_breakable_prefill_matches_eager`.
+ TP1_BREAKABLE_PREFILL = dict(
prefill_cuda_graph_backend="breakable",
prefill_capture_num_tokens=[128, 256, 512, 1024, 2048, 3072],
)
- ENGINE_BREAKABLE = dict(ENGINE, **BREAKABLE_PREFILL)
- # One-model MTP-Eagle with the checkpoint's single MTP module replayed
- # `max_draft_len` times.
+ # One draft, and the longest the target runs.
MTP1 = MTPDecodingConfig(max_draft_len=1)
- MTP3 = MTPDecodingConfig(max_draft_len=3)
MTP7 = MTPDecodingConfig(max_draft_len=7)
- # Acceptance-length anchors in references/acceptance_length.yaml, keyed by
- # target and variant.
- ACCEPTANCE_KEY = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1"
- ACCEPTANCE_KEY_LONG_CONTEXT = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_long_context"
- ACCEPTANCE_KEY_OPT_IN = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl_fp8_kv_bf16_state"
- ACCEPTANCE_KEY_OPT_IN_LONG_CONTEXT = (
- "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl_fp8_kv_bf16_state_long_context"
+ # The tp1 cases' acceptance-length anchors in references/acceptance_length.yaml.
+ TP1_ACCEPTANCE_KEY = "ModelingV2Qwen38FlashNextNvfp4Sm103Default::tp1_mtp1"
+ TP1_ACCEPTANCE_KEY_LONG_CONTEXT = (
+ "ModelingV2Qwen38FlashNextNvfp4Sm103Default::tp1_mtp3_long_context"
)
- ACCEPTANCE_KEY_MTP7_LONG_CONTEXT = (
- "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp7_nvfp4_kv_bf16_state_long_context"
+ TP1_ACCEPTANCE_KEY_OPT_IN = (
+ "ModelingV2Qwen38FlashNextNvfp4Sm103Default::tp1_mtp3_cutedsl_fp8_kv_bf16_state"
+ )
+ TP1_ACCEPTANCE_KEY_OPT_IN_LONG_CONTEXT = (
+ "ModelingV2Qwen38FlashNextNvfp4Sm103Default::"
+ "tp1_mtp3_cutedsl_fp8_kv_bf16_state_long_context"
+ )
+ TP1_ACCEPTANCE_KEY_MTP7_LONG_CONTEXT = (
+ "ModelingV2Qwen38FlashNextNvfp4Sm103Default::tp1_mtp7_nvfp4_kv_bf16_state_long_context"
)
- def _gsm8k_and_acceptance(self, mocker, llm, key, label):
- """gsm8k on ``llm``, then the same run's acceptance length against ``key``.
+ def _gsm8k(self, mocker, llm, label: str, key: str | None = None) -> None:
+ """gsm8k against the engine's reference row, then its acceptance length.
- gsm8k is read against the reference row the engine resolves; the acceptance length comes
- from the engine's iteration statistics, so the engine needs ``enable_iter_perf_stats``.
+ The acceptance length is held to the floor under `key` in
+ references/acceptance_length.yaml, or without one above the three-draft collapse
+ tripwire. Rejection sampling holds the emitted distribution to the target model's, so
+ a miscomputed draft layer leaves the score in place and only slows generation; the
+ acceptance length is what sees it. It comes from the engine's iteration statistics,
+ so the engine needs ``enable_iter_perf_stats``.
"""
mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN)
GSM8K(self.MODEL_NAME).evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS)
acceptance_length = compute_acceptance_length(llm)
print(f"[AL] {label} acceptance_length = {acceptance_length:.3f}")
- assert_acceptance_length(key, acceptance_length)
-
- def _passkeys(self, llm, lengths, draft_len, key, label):
- """Pass key retrieval on ``llm`` over prompts past the indexer's budget.
+ if key is not None:
+ assert_acceptance_length(key, acceptance_length)
+ return
+ assert acceptance_length >= self.MTP3_COLLAPSE_FLOOR, (
+ f"{label}: acceptance length {acceptance_length:.3f} is below "
+ f"{self.MTP3_COLLAPSE_FLOOR}: the draft path is rejected more often than accepted"
+ )
- Every key must be stated back, and two identical requests far apart in one batch
- (different slots) must give the same tokens, which fails if a request's selection leaks
- into another slot. The acceptance length, from the requests' own metrics, is held to the
- floor under ``key``.
- """
+ def _passkeys(self, llm, lengths, draft_len, key, label) -> None:
+ """Pass keys over prompts of `lengths` tokens, the acceptance length held to `key`."""
tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH)
prompts = _passkey_prompts(tokenizer, lengths)
- texts = [prompt for prompt, _ in prompts] + [prompts[len(prompts) // 2][0]]
- outputs = llm.generate(
- texts,
- SamplingParams(max_tokens=_PASSKEY_MAX_TOKENS, temperature=0, return_perf_metrics=True),
- use_tqdm=False,
- )
- missed = [
- k
- for (_, k), output in zip(prompts, outputs)
- if k not in output.outputs[0].text[:_PASSKEY_MAX_TOKENS]
- ]
- assert not missed, f"pass keys not stated back: {missed} of {[k for _, k in prompts]}"
- twins = outputs[len(prompts) // 2], outputs[-1]
- assert twins[0].outputs[0].token_ids == twins[1].outputs[0].token_ids
- acceptance_length = _acceptance_length_from_requests(outputs, draft_len)
- print(f"[AL] {label} acceptance_length = {acceptance_length:.3f}")
- assert_acceptance_length(key, acceptance_length)
-
- def _mmmu(self, llm):
- """MMMU on ``llm`` against the reference row the engine resolves."""
- # MMMU adds <|endoftext|> to the stop token, as the built-in model's
- # MMMU test does.
- sampling_params = SamplingParams(
- max_tokens=MMMU.MAX_OUTPUT_LEN,
- truncate_prompt_tokens=MMMU.MAX_INPUT_LEN,
- stop="<|endoftext|>",
- )
- MMMU(self.MODEL_NAME).evaluate(
- llm,
- sampling_params=sampling_params,
- extra_evaluator_kwargs=dict(chat_template_kwargs=dict(enable_thinking=False)),
- )
+ assert_acceptance_length(key, self._passkey_acceptance(llm, prompts, label, draft_len))
def _video_checks(self, llm):
"""Synthetic videos through ``llm``.
@@ -488,7 +576,7 @@ def _video_checks(self, llm):
bound (a 3072x3072 image is 192x192 = 36864 patches) must fail as that request's own
error, and the request after it must still complete; ``max_seq_len`` is raised so that
image reaches the admission check rather than the sequence length check. The order of
- the splice itself is pinned by ``test_modeling_v2_qwen4_exp_multimodal_embeds``.
+ the splice itself is pinned by the target's ``test_vision``.
"""
tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH)
processor = AutoProcessor.from_pretrained(
@@ -522,8 +610,234 @@ def ask(text, items):
answer = ask(_VIDEO_QUESTION, [("video", _synthetic_video("red_rightward"))]).lower()
assert "red" in answer and "right" in answer, answer
+ def _long_context_passkey(
+ self, monkeypatch, speculative_config, engine, index_share, lengths, key, label, mocker=None
+ ):
+ """`_passkeys` with the executor in this process, checking the op calls too.
+
+ With ``mocker``, gsm8k runs first on the same engine. The draft layer's IndexShare
+ reuse op must run exactly when IndexShare is on, and no prefill graph may replay.
+ """
+ monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1")
+ reuse_calls = []
+ reuse = target_modeling.qsa_mtp_reuse_topk
+
+ def counting_reuse(*args, **kwargs):
+ reuse_calls.append(1)
+ return reuse(*args, **kwargs)
+
+ monkeypatch.setattr(target_modeling, "qsa_mtp_reuse_topk", counting_reuse)
+ replays = _count_prefill_graph_replays(monkeypatch)
+ engine = dict(
+ engine,
+ sparse_attention_config=QSASparseAttentionConfig(
+ index_share_for_mtp_iteration=index_share
+ ),
+ )
+ with (
+ scoped_env_overrides(_VALIDATE) as rank_env,
+ LLM(
+ self.MODEL_PATH,
+ speculative_config=speculative_config,
+ **engine,
+ modeling_v2="require",
+ env_overrides=rank_env,
+ ) as llm,
+ ):
+ assert _target_core(llm).mtp_index_share == index_share
+ if mocker is not None:
+ mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN)
+ GSM8K(self.MODEL_NAME).evaluate(
+ llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS
+ )
+ self._passkeys(llm, lengths, speculative_config.max_draft_len, key, label)
+ assert not replays, replays
+ if index_share:
+ assert reuse_calls, "IndexShare is on and the reuse op never ran"
+ else:
+ assert not reuse_calls, "IndexShare is off and the reuse op ran"
+
+ def _passkey_acceptance(self, llm, prompts, label, draft_len: int = 3) -> float:
+ """Pass-key retrieval on ``llm`` and the requests' mean acceptance length.
+
+ Every key must be stated back, and two identical requests far apart in one batch
+ (different slots) must give the same tokens, which fails if a request's selection
+ leaks into another slot.
+ """
+ texts = [prompt for prompt, _ in prompts] + [prompts[len(prompts) // 2][0]]
+ outputs = llm.generate(
+ texts,
+ SamplingParams(max_tokens=_PASSKEY_MAX_TOKENS, temperature=0, return_perf_metrics=True),
+ use_tqdm=False,
+ )
+ missed = [
+ key
+ for (_, key), output in zip(prompts, outputs)
+ if key not in output.outputs[0].text[:_PASSKEY_MAX_TOKENS]
+ ]
+ assert not missed, (
+ f"{label}: pass keys not stated back: {missed} of {[key for _, key in prompts]}"
+ )
+ twins = outputs[len(prompts) // 2], outputs[-1]
+ assert twins[0].outputs[0].token_ids == twins[1].outputs[0].token_ids, (
+ f"{label}: the same prompt in two slots gave different tokens"
+ )
+ acceptance_length = _acceptance_length_from_requests(outputs, draft_len)
+ print(f"[AL] {label} acceptance_length = {acceptance_length:.3f}")
+ return acceptance_length
+
+ @skip_not_sm103
+ @pytest.mark.skip_less_device(4)
+ def test_mtp3_gsm8k_and_long_context(self, mocker):
+ """tep4 with three drafts and IndexShare: gsm8k, then pass keys over 3k-6k tokens.
+
+ gsm8k reads the MTP row of references/gsm8k.yaml, its acceptance length held above
+ the collapse tripwire. The pass keys are past the QSA indexer's 2048-token budget, so
+ the target's decode steps and its draft layer's take the sparse attention path, which
+ gsm8k never leaves, and IndexShare makes the draft loop reuse the selection of its
+ first step. A miscomputed draft path keeps the text correct and only lowers the
+ acceptance length, which is held to its floor in references/acceptance_length.yaml.
+ """
+ tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH)
+ prompts = _passkey_prompts(tokenizer, _PASSKEY_LENGTHS)
+ with (
+ scoped_env_overrides(_VALIDATE) as rank_env,
+ LLM(
+ self.MODEL_PATH,
+ speculative_config=self.MTP3,
+ enable_iter_perf_stats=True,
+ sparse_attention_config=QSASparseAttentionConfig(
+ index_share_for_mtp_iteration=True
+ ),
+ **self.ENGINE_MTP,
+ **_TEP4,
+ modeling_v2="require",
+ env_overrides=rank_env,
+ ) as llm,
+ ):
+ self._gsm8k(mocker, llm, "tep4 mtp3 gsm8k")
+ acceptance_length = self._passkey_acceptance(llm, prompts, "tep4 mtp3 long context")
+ assert_acceptance_length(self.ACCEPTANCE_KEY_LONG_CONTEXT, acceptance_length)
+
@skip_not_sm103
- def test_gsm8k(self, mocker):
+ @pytest.mark.skip_less_device(4)
+ @pytest.mark.parametrize("parallel, kv_dtype, breakable", _QUANTIZED_KV)
+ def test_mtp3_gsm8k_quantized_kv(self, mocker, parallel, kv_dtype, breakable):
+ """gsm8k with three drafts over a quantized KV pool, block reuse and chunked prefill.
+
+ gsm8k reads the MTP row of references/gsm8k.yaml for that KV dtype, its acceptance
+ length held above the collapse tripwire.
+ """
+ kv_cache_config = KvCacheConfig(
+ dtype=kv_dtype,
+ free_gpu_memory_fraction=0.6,
+ enable_block_reuse=True,
+ mamba_ssm_cache_dtype="bfloat16",
+ )
+ engine = dict(self.ENGINE_MTP, enable_chunked_prefill=True, kv_cache_config=kv_cache_config)
+ if breakable:
+ engine.update(self.BREAKABLE_PREFILL)
+ with (
+ scoped_env_overrides(_VALIDATE) as rank_env,
+ LLM(
+ self.MODEL_PATH,
+ speculative_config=self.MTP3,
+ enable_iter_perf_stats=True,
+ **engine,
+ **parallel,
+ modeling_v2="require",
+ env_overrides=rank_env,
+ ) as llm,
+ ):
+ self._gsm8k(mocker, llm, f"mtp3 {kv_dtype} kv")
+
+ @skip_not_sm103
+ @pytest.mark.skip_less_device(4)
+ def test_mtp3_mmmu(self):
+ """MMMU at tep4 with three drafts, against the MTP row of references/mmmu.yaml.
+
+ It runs the vision tower, replicated on every rank, and the mrope path.
+ """
+ with (
+ scoped_env_overrides(_VALIDATE) as rank_env,
+ LLM(
+ self.MODEL_PATH,
+ speculative_config=self.MTP3,
+ **self.ENGINE_MTP,
+ **_TEP4,
+ modeling_v2="require",
+ env_overrides=rank_env,
+ ) as llm,
+ ):
+ _mmmu(llm, self.MODEL_NAME)
+
+ @skip_not_sm103
+ @pytest.mark.timeout(5400)
+ @pytest.mark.skip_less_device(4)
+ @pytest.mark.skip_less_device_memory(145000)
+ @pytest.mark.skip_less_host_memory(204800)
+ def test_disaggregated_dep2(self, mocker):
+ """dep2 on both engines, one context and one generation server of two GPUs each.
+
+ Attention DP over two ranks, the routed experts in two 256-expert windows; both
+ engines run one layout, so each rank's K/V, QSA index, Gated-DeltaNet and PLE pools
+ have one layout on both sides of the Python NIXL transceiver.
+ """
+ # The core configuration on both engines: the checkpoint's bf16 K/V, a bf16
+ # recurrent state, no block reuse, no chunked prefill, the trtllm-gen experts.
+ common_config = {
+ "modeling_v2": "require",
+ "tensor_parallel_size": 2,
+ "moe_expert_parallel_size": 2,
+ "enable_attention_dp": True,
+ "max_seq_len": 9216,
+ "max_batch_size": 128,
+ "max_num_tokens": 8192,
+ "enable_chunked_prefill": False,
+ "moe_config": {"backend": "TRTLLM"},
+ "kv_cache_config": {
+ "enable_block_reuse": False,
+ "mamba_ssm_cache_dtype": "bfloat16",
+ "free_gpu_memory_fraction": 0.6,
+ },
+ # The V2 hybrid manager's pools travel over the Python transceiver.
+ "cache_transceiver_config": {"backend": "NIXL", "transceiver_runtime": "PYTHON"},
+ }
+ ctx_server_config = {
+ **common_config,
+ "disable_overlap_scheduler": True,
+ "cuda_graph_config": None,
+ }
+ gen_server_config = {
+ **common_config,
+ "disable_overlap_scheduler": False,
+ "cuda_graph_config": {"enable_padding": True, "max_batch_size": 128},
+ }
+
+ mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN)
+ with launch_disaggregated_llm(
+ _disagg_server_config(),
+ ctx_server_config,
+ gen_server_config,
+ self.MODEL_PATH,
+ max_workers=128,
+ extra_env=_DISAGG_WORKER_ENV,
+ # What every rank of both engines prints once it has bound the manager's pools.
+ assert_worker_log_contains="modeling_v2 qwen4_exp default target: rank",
+ ) as llm:
+ # The served quantization the reference is keyed by; the launcher reads NVFP4
+ # off the checkpoint's name.
+ llm.args.quant_config.quant_algo = QuantAlgo.MIXED_PRECISION
+ run_accuracy_test(
+ llm,
+ self.MODEL_NAME,
+ ["GSM8K"],
+ extra_evaluator_kwargs={GSM8K: self.EXTRA_EVALUATOR_KWARGS},
+ timeout=3000,
+ )
+
+ @skip_not_sm103
+ def test_tp1_gsm8k(self, mocker):
"""The text model without speculation, against the gsm8k anchor."""
mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN)
@@ -531,7 +845,7 @@ def test_gsm8k(self, mocker):
scoped_env_overrides(_VALIDATE) as rank_env,
LLM(
self.MODEL_PATH,
- **self.ENGINE,
+ **self.TP1_ENGINE,
modeling_v2="require",
env_overrides=rank_env,
) as llm,
@@ -540,7 +854,7 @@ def test_gsm8k(self, mocker):
task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS)
@skip_not_sm103
- def test_mtp1_gsm8k_mmmu_and_video(self, mocker):
+ def test_tp1_mtp1_gsm8k_mmmu_and_video(self, mocker):
"""gsm8k with its acceptance-length floor, MMMU and synthetic videos on one MTP engine.
- gsm8k against the MTP row of references/gsm8k.yaml. A miscomputed draft layer only
@@ -553,7 +867,9 @@ def test_mtp1_gsm8k_mmmu_and_video(self, mocker):
Block reuse is off, as the vision path requires; it does not change the text numerics.
``max_seq_len`` is raised for the video checks' oversized image.
"""
- engine = dict(self.ENGINE, kv_cache_config=self.KV_CACHE_MULTIMODAL, max_seq_len=16384)
+ engine = dict(
+ self.TP1_ENGINE, kv_cache_config=self.TP1_KV_CACHE_MULTIMODAL, max_seq_len=16384
+ )
with (
scoped_env_overrides(_VALIDATE) as rank_env,
LLM(
@@ -565,12 +881,12 @@ def test_mtp1_gsm8k_mmmu_and_video(self, mocker):
env_overrides=rank_env,
) as llm,
):
- self._gsm8k_and_acceptance(mocker, llm, self.ACCEPTANCE_KEY, "mtp1")
- self._mmmu(llm)
+ self._gsm8k(mocker, llm, "tp1 mtp1", self.TP1_ACCEPTANCE_KEY)
+ _mmmu(llm, self.MODEL_NAME)
self._video_checks(llm)
@skip_not_sm103
- def test_mtp3_cutedsl_fp8_kv_bf16_state(self, mocker):
+ def test_tp1_mtp3_cutedsl_fp8_kv_bf16_state(self, mocker):
"""gsm8k and 3k-6k-token pass keys with three drafts over the opt-in choices.
One engine runs gsm8k with an acceptance-length floor, then the pass keys. The trunk's
@@ -585,9 +901,9 @@ def test_mtp3_cutedsl_fp8_kv_bf16_state(self, mocker):
cannot use, which accuracy does not.
"""
engine = dict(
- self.ENGINE,
+ self.TP1_ENGINE,
moe_config=MoeConfig(backend="CUTEDSL"),
- kv_cache_config=self.KV_CACHE_FP8_BF16_STATE,
+ kv_cache_config=self.TP1_KV_CACHE_FP8_BF16_STATE,
sparse_attention_config=QSASparseAttentionConfig(index_share_for_mtp_iteration=False),
)
with (
@@ -601,81 +917,41 @@ def test_mtp3_cutedsl_fp8_kv_bf16_state(self, mocker):
env_overrides=rank_env,
) as llm,
):
- self._gsm8k_and_acceptance(mocker, llm, self.ACCEPTANCE_KEY_OPT_IN, "mtp3 opt-in")
+ self._gsm8k(mocker, llm, "tp1 mtp3 opt-in", self.TP1_ACCEPTANCE_KEY_OPT_IN)
self._passkeys(
llm,
_PASSKEY_LENGTHS,
self.MTP3.max_draft_len,
- self.ACCEPTANCE_KEY_OPT_IN_LONG_CONTEXT,
- "mtp3 opt-in long context",
+ self.TP1_ACCEPTANCE_KEY_OPT_IN_LONG_CONTEXT,
+ "tp1 mtp3 opt-in long context",
)
- def _long_context_passkey(
- self, monkeypatch, speculative_config, engine, index_share, lengths, key, label
- ):
- """`_passkeys` with the executor in this process, checking the op calls too.
-
- The draft layer's IndexShare reuse op must run exactly when IndexShare is on, and no
- prefill graph may replay.
- """
- monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1")
- reuse_calls = []
- reuse = target_modeling.qsa_mtp_reuse_topk
-
- def counting_reuse(*args, **kwargs):
- reuse_calls.append(1)
- return reuse(*args, **kwargs)
-
- monkeypatch.setattr(target_modeling, "qsa_mtp_reuse_topk", counting_reuse)
- replays = _count_prefill_graph_replays(monkeypatch)
- engine = dict(
- engine,
- sparse_attention_config=QSASparseAttentionConfig(
- index_share_for_mtp_iteration=index_share
- ),
- )
- with (
- scoped_env_overrides(_VALIDATE) as rank_env,
- LLM(
- self.MODEL_PATH,
- speculative_config=speculative_config,
- **engine,
- modeling_v2="require",
- env_overrides=rank_env,
- ) as llm,
- ):
- assert _target_core(llm).mtp_index_share == index_share
- self._passkeys(llm, lengths, speculative_config.max_draft_len, key, label)
- assert not replays, replays
- if index_share:
- assert reuse_calls, "IndexShare is on and the reuse op never ran"
- else:
- assert not reuse_calls, "IndexShare is off and the reuse op ran"
-
@skip_not_sm103
- def test_mtp3_long_context_passkey(self, monkeypatch):
- """Pass key retrieval over 3k-6k-token prompts with three drafts and IndexShare.
-
- The KV cache is bf16; the prompts are long enough to run the sparse attention path. The
- gsm8k prompts never leave the dense path (they are under the indexer's 2048-token
- budget), so the pass-key cases are the ones that run the target's QSA selection and the
- draft layer's end to end; a pass key stated back correctly needs the right blocks
- selected at every depth. IndexShare makes the draft loop reuse the selection of its
- first step: the acceptance length must not fall below the floor taken from the same
- prompts with the draft layer selecting on every step, and the reuse op must have run.
+ def test_tp1_mtp3_gsm8k_and_long_context(self, mocker, monkeypatch):
+ """gsm8k, then 3k-6k-token pass keys, with three drafts and IndexShare on one engine.
+
+ The KV cache is bf16, the recurrent state fp32, with block reuse and chunked prefill.
+ gsm8k reads the MTP row of references/gsm8k.yaml; its prompts stay on the dense
+ attention path (they are under the indexer's 2048-token budget). The pass keys are
+ past it, so they run the target's QSA selection and the draft layer's end to end; a
+ pass key stated back correctly needs the right blocks selected at every depth.
+ IndexShare makes the draft loop reuse the selection of its first step: the pass keys'
+ acceptance length must not fall below the floor taken from the same prompts with the
+ draft layer selecting on every step, and the reuse op must have run.
"""
self._long_context_passkey(
monkeypatch,
self.MTP3,
- self.ENGINE,
+ self.TP1_ENGINE,
True,
_PASSKEY_LENGTHS,
- self.ACCEPTANCE_KEY_LONG_CONTEXT,
- "mtp3 long context",
+ self.TP1_ACCEPTANCE_KEY_LONG_CONTEXT,
+ "tp1 mtp3 long context",
+ mocker=mocker,
)
@skip_not_sm103
- def test_breakable_prefill_matches_eager(self, monkeypatch):
+ def test_tp1_breakable_prefill_matches_eager(self, monkeypatch):
"""Greedy continuations of the same prompts, with the breakable prefill graph and without.
MTP with three drafts, a 3072-token chunk budget, and prompts of a natural document
@@ -748,13 +1024,15 @@ def run(engine):
long_tokens = long_future.result().outputs[0].token_ids
return single, again, answers, long_tokens
- single, again, answers, long_tokens = run(dict(self.ENGINE_BREAKABLE, max_num_tokens=3072))
+ single, again, answers, long_tokens = run(
+ dict(dict(self.TP1_ENGINE, **self.TP1_BREAKABLE_PREFILL), max_num_tokens=3072)
+ )
assert {128, 256, 1024, 2048, 3072} <= set(replays), replays
assert single == again, "the same engine repeated itself with a different answer"
replays.clear()
eager_single, _, eager_answers, eager_long_tokens = run(
- dict(self.ENGINE, max_num_tokens=3072)
+ dict(self.TP1_ENGINE, max_num_tokens=3072)
)
assert not replays, replays
@@ -771,7 +1049,7 @@ def run(engine):
assert long_tokens[:4] == eager_long_tokens[:4], (long_tokens, eager_long_tokens)
@skip_not_sm103
- def test_mtp7_long_context_passkey(self, monkeypatch):
+ def test_tp1_mtp7_long_context_passkey(self, monkeypatch):
"""Pass key retrieval over 3k-12k-token prompts with seven drafts and IndexShare.
The KV cache is NVFP4 and the recurrent state bf16. Seven drafts run the sparse path
@@ -785,7 +1063,7 @@ def test_mtp7_long_context_passkey(self, monkeypatch):
tokens bound that drift. 12k tokens past the 8k default needs ``max_seq_len`` 16384.
"""
engine = dict(
- self.ENGINE, max_seq_len=16384, kv_cache_config=self.KV_CACHE_NVFP4_BF16_STATE
+ self.TP1_ENGINE, max_seq_len=16384, kv_cache_config=self.TP1_KV_CACHE_NVFP4_BF16_STATE
)
self._long_context_passkey(
monkeypatch,
@@ -793,6 +1071,112 @@ def test_mtp7_long_context_passkey(self, monkeypatch):
engine,
True,
_PASSKEY_LENGTHS_LONG,
- self.ACCEPTANCE_KEY_MTP7_LONG_CONTEXT,
- "mtp7 nvfp4 kv bf16 state long context",
+ self.TP1_ACCEPTANCE_KEY_MTP7_LONG_CONTEXT,
+ "tp1 mtp7 nvfp4 kv bf16 state long context",
)
+
+ # The target's AgentX deployment environment at tp1 on both engines: PDL, eight K/V
+ # transfer threads and every opt-in projection of the target.
+ AGENTX_ENV = {
+ "TRTLLM_ENABLE_PDL": "1",
+ "TRTLLM_KV_TRANSFER_NUM_THREADS": "8",
+ "TRTLLM_QWEN4_EXP_PLE_HOST_OFFLOAD": "1",
+ "TRTLLM_MODELING_V2_DENSE_FP8": "1",
+ "TRTLLM_MODELING_V2_GDN_INPROJ_NVFP4": "1",
+ "TRTLLM_MODELING_V2_MTP_MOE_NVFP4": "1",
+ "TRTLLM_MODELING_V2_MTP_HEAD_NVFP4": "1",
+ "TRTLLM_MODELING_V2_TARGET_HEAD_NVFP4": "1",
+ "TRTLLM_MODELING_V2_MTP_DENSE_FP8": "1",
+ }
+
+ @skip_not_sm103
+ @pytest.mark.timeout(5400)
+ @pytest.mark.skip_less_device(2)
+ @pytest.mark.skip_less_device_memory(145000)
+ @pytest.mark.skip_less_host_memory(204800)
+ def test_tp1_disaggregated_mtp2(self, mocker):
+ """One context and one generation GPU, the engines of the target's AgentX deployment.
+
+ One rank each, so each runs the target at tp1.
+
+ NVFP4 K/V, a bf16 recurrent state, MTP with two drafts, per-conversation block reuse
+ on the context engine and rejection sampling on the generation engine, over the
+ Python NIXL transceiver. Acceptance is real.
+ """
+ common_config = {
+ "modeling_v2": "require",
+ "max_seq_len": 262144,
+ "max_num_tokens": 8192,
+ "enable_chunked_prefill": True,
+ "disable_mm_encoder": True,
+ "moe_config": {"backend": "TRTLLM"},
+ "scheduler_config": {"capacity_scheduler_policy": "GUARANTEED_NO_EVICT"},
+ "cache_transceiver_config": {"backend": "NIXL", "transceiver_runtime": "PYTHON"},
+ }
+ kv_cache_config = {
+ "dtype": "nvfp4",
+ "free_gpu_memory_fraction": 0.85,
+ "mamba_ssm_cache_dtype": "bfloat16",
+ "use_kv_cache_manager_v2": True,
+ "tokens_per_block": 32,
+ }
+ # The context engine's breakable prefill graphs, up to max_num_tokens.
+ prefill_capture_num_tokens = (
+ [2**i for i in range(10)] + list(range(768, 3328, 256)) + [4096, 6144, 8192]
+ )
+ ctx_server_config = {
+ **common_config,
+ "max_batch_size": 32,
+ "kv_cache_config": {
+ **kv_cache_config,
+ "enable_block_reuse": True,
+ "block_reuse_config": {"policy": "per_conversation", "max_num_turns": 3},
+ # One recurrent-state snapshot two tokens before each prompt's end.
+ "mamba_state_config": {
+ "periodic_snapshot_interval": 0,
+ "additional_snapshot_offsets_from_end": [2],
+ },
+ },
+ "cuda_graph_config": {"enable_padding": True, "max_batch_size": 32},
+ "prefill_cuda_graph_backend": "breakable",
+ "prefill_capture_num_tokens": prefill_capture_num_tokens,
+ "speculative_config": {"decoding_type": "MTP", "max_draft_len": 2},
+ }
+ gen_server_config = {
+ **common_config,
+ "max_batch_size": 128,
+ "kv_cache_config": {**kv_cache_config, "enable_block_reuse": False},
+ "cuda_graph_config": {
+ "enable_padding": True,
+ "batch_sizes": list(range(1, 17)) + list(range(24, 129, 8)),
+ },
+ "speculative_config": {
+ "decoding_type": "MTP",
+ "max_draft_len": 2,
+ "use_rejection_sampling": True,
+ },
+ }
+
+ mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN)
+ with launch_disaggregated_llm(
+ _disagg_server_config(ctx_router="conversation", gen_router="load_balancing"),
+ ctx_server_config,
+ gen_server_config,
+ self.MODEL_PATH,
+ max_workers=128,
+ extra_env={**_DISAGG_WORKER_ENV, **self.AGENTX_ENV},
+ # What the target prints once it has packed the target head for decode rows.
+ assert_worker_log_contains="target head NVFP4",
+ ) as llm:
+ # The served quantization and draft the reference is keyed by; the launcher
+ # reads NVFP4 off the checkpoint's name.
+ llm.args.quant_config.quant_algo = QuantAlgo.MIXED_PRECISION
+ llm.args.quant_config.kv_cache_quant_algo = QuantAlgo.NVFP4
+ llm.args.speculative_config = MTPDecodingConfig(max_draft_len=2)
+ run_accuracy_test(
+ llm,
+ self.MODEL_NAME,
+ ["GSM8K"],
+ extra_evaluator_kwargs={GSM8K: self.EXTRA_EVALUATOR_KWARGS},
+ timeout=3000,
+ )
diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml
index fc3eaad3ac9b..308420d9a5aa 100644
--- a/tests/integration/defs/accuracy/references/acceptance_length.yaml
+++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml
@@ -153,8 +153,9 @@ ModelingV2Glm52Nvfp4Sm103Tep4::mtp1:
# self-referential; min_al is the automatic 95% of it.
ref_al: 1.942036
min_al: 1.844935
-# Read by the acceptance-length cases of TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1
-# (test_qwen4_exp.py). Keyed by target and variant for the same reason as the
+# Read by the tp1 acceptance-length cases of TestModelingV2Qwen38FlashNextNvfp4Sm103Default
+# (test_qwen4_exp.py, the `test_tp1_*` cases: the target at one rank). Keyed by target, layout
+# and variant for the same reason as the
# ModelingV2DeepseekR10528Nvfp4Sm10xDefault::mtp3 entry above, and one key per variant:
# TRTLLM_POPULATE_ACCEPTANCE_LENGTH=1 on one case must not overwrite another's floor.
#
@@ -165,8 +166,8 @@ ModelingV2Glm52Nvfp4Sm103Tep4::mtp1:
# every draft accepted. Each min_al below is a hand-set collapse tripwire unless its
# comment says it is 95% of ref_al.
-# test_mtp1_gsm8k_mmmu_and_video, over gsm8k
-ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1:
+# test_tp1_mtp1_gsm8k_mmmu_and_video, over gsm8k
+ModelingV2Qwen38FlashNextNvfp4Sm103Default::tp1_mtp1:
# Smoke-scale, and labelled as such: measured on GB300 over the boot prompt
# pair, STOCK 1.7403 against modeling_v2's 1.8063 on the same prompts. The
# number below is the STOCK reading, so the gate is not self-referential.
@@ -184,9 +185,9 @@ ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1:
# measured, and tighten min_al then.
min_al: 1.5
-# test_mtp3_cutedsl_fp8_kv_bf16_state, over gsm8k: three drafts with the trunk's experts on
+# test_tp1_mtp3_cutedsl_fp8_kv_bf16_state, over gsm8k: three drafts with the trunk's experts on
# CuteDSL over an fp8 KV cache and a bf16 recurrent state.
-ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl_fp8_kv_bf16_state:
+ModelingV2Qwen38FlashNextNvfp4Sm103Default::tp1_mtp3_cutedsl_fp8_kv_bf16_state:
# modeling_v2's own full-gsm8k reading on GB300 (1319 samples); there is no built-in reading
# for this combination (the built-in model has no CuteDSL trunk). 3.705 over bf16 KV.
ref_al: 3.509
@@ -194,11 +195,11 @@ ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl_fp8_kv_bf16_state:
# accepted fails it.
min_al: 3.334
-# test_mtp3_long_context_passkey: three drafts with IndexShare on, bf16 KV, over
-# 13 prompts of 3k-6k tokens (the sparse attention path). Most of the output is a copy of the
-# document's section headings, which the draft predicts almost always, so the readings sit
-# at the range's ceiling of 4.0.
-ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_long_context:
+# test_tp1_mtp3_gsm8k_and_long_context, over its pass keys: three drafts with IndexShare on, bf16
+# KV, over 13 prompts of 3k-6k tokens (the sparse attention path). Most of the output is a copy
+# of the document's section headings, which the draft predicts almost always, so the readings
+# sit at the range's ceiling of 4.0.
+ModelingV2Qwen38FlashNextNvfp4Sm103Default::tp1_mtp3_long_context:
# The same prompts with the draft layer selecting on every step (IndexShare off): 4.000,
# equal to the IndexShare reading.
ref_al: 4.000
@@ -207,20 +208,32 @@ ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_long_context:
# 3.760, a 6% drop the automatic floor would pass.
min_al: 3.9
-# test_mtp3_cutedsl_fp8_kv_bf16_state, over the pass keys of mtp3_long_context with IndexShare
-# off.
-ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl_fp8_kv_bf16_state_long_context:
+# test_tp1_mtp3_cutedsl_fp8_kv_bf16_state, over the pass keys of tp1_mtp3_long_context with
+# IndexShare off.
+ModelingV2Qwen38FlashNextNvfp4Sm103Default::tp1_mtp3_cutedsl_fp8_kv_bf16_state_long_context:
# 3.962 here, against 4.000 over bf16 KV.
ref_al: 3.962
- # 2.5% below ref_al, as for mtp3_long_context.
+ # 2.5% below ref_al, as for tp1_mtp3_long_context.
min_al: 3.86
-# test_mtp7_long_context_passkey: seven drafts with IndexShare on, over an NVFP4 KV cache and a
+# test_tp1_mtp7_long_context_passkey: seven drafts with IndexShare on, over an NVFP4 KV cache and a
# bf16 recurrent state, over 12 prompts of 3k-12k tokens (the sparse attention path), whose
# output mostly copies the document's section headings. The KV dtype sets the reading, and seven
# chained drafts amplify it: 7.857 over bf16 KV, 7.500 over fp8, 6.867 over NVFP4 (6.961 with
# IndexShare off, 7.003 over an fp32 state).
-ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp7_nvfp4_kv_bf16_state_long_context:
+ModelingV2Qwen38FlashNextNvfp4Sm103Default::tp1_mtp7_nvfp4_kv_bf16_state_long_context:
ref_al: 6.867
- # 2.5% below ref_al, as for mtp3_long_context.
+ # 2.5% below ref_al, as for tp1_mtp3_long_context.
min_al: 6.69
+
+# Read by TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_and_long_context
+# (test_qwen4_exp.py), over the pass keys of tp1_mtp3_long_context at tep4: three drafts with
+# IndexShare on, bf16 KV.
+ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3_long_context:
+ # modeling_v2's own reading on GB300, at the range's ceiling as at tp1. The
+ # built-in implementation does not start this engine (its decode front's router GEMM
+ # asserts a power-of-two K), so there is no stock reading.
+ ref_al: 4.000
+ # 2.5% below ref_al, as for tp1_mtp3_long_context, whose deliberately wrong
+ # reused selection scores 3.760.
+ min_al: 3.9
diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt
index 2ef5f1fd47cf..edd4420f9e91 100644
--- a/tests/integration/test_lists/qa/llm_function_core.txt
+++ b/tests/integration/test_lists/qa/llm_function_core.txt
@@ -784,6 +784,22 @@ unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latenc
unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_draft_sample_glue.py
unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/test_modeling_v2_intra_node_reduce_op_matrix.py
+# Qwen3.8-Flash-Next-NVFP4, modeling_v2 qwen4_exp target (sm_103 only): every accuracy test of
+# the target, the CI ones included, and those that run only here (tp1 without speculation,
+# tp1's breakable prefill graph against eager, mmmu at tep4).
+accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_gsm8k
+accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp3_gsm8k_and_long_context
+accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp1_gsm8k_mmmu_and_video
+accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp3_cutedsl_fp8_kv_bf16_state
+accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp7_long_context_passkey
+accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_breakable_prefill_matches_eager
+accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_disaggregated_mtp2 TIMEOUT (90)
+accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_and_long_context TIMEOUT (60)
+accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[tp4_nvfp4_kv_breakable] TIMEOUT (60)
+accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[dep4_fp8_kv] TIMEOUT (60)
+accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_mmmu TIMEOUT (90)
+accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_disaggregated_dep2 TIMEOUT (90)
+
# Dense FP4 MLA accuracy (SM107 only)
accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_mla
accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_nvfp4_kv_cache_nixl
diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml
index 58fb98d6d86d..5e8f195c6f93 100644
--- a/tests/integration/test_lists/test-db/l0_b300.yml
+++ b/tests/integration/test_lists/test-db/l0_b300.yml
@@ -205,14 +205,12 @@ l0_b300:
- unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_rms_norm_gated_token_major.py
- unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_situ_and_mul.py
- accuracy/modeling_v2/test_gpt_oss.py::TestModelingV2GptOss120bSm10xDefault::test_gsm8k[tp1]
- # qwen4_exp target, pre-merge: gsm8k; one MTP engine for gsm8k with the acceptance-length
- # floor (the only gate that sees a miscomputed draft layer), mmmu (the vision tower) and
- # synthetic videos; and 3k-6k-token pass keys with three drafts and IndexShare over bf16 KV,
- # the only prompts long enough to run the sparse attention path. The opt-in variants run
- # post-merge (below).
- - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_gsm8k
- - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_gsm8k_mmmu_and_video
- - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_long_context_passkey
+ # qwen4_exp target at tp1, pre-merge: one engine with three drafts and IndexShare over bf16 KV,
+ # gsm8k and then 3k-6k-token pass keys, the only prompts long enough to run the sparse
+ # attention path, with their acceptance-length floor (what sees a miscomputed draft layer).
+ # The other engines run post-merge (below); no speculation and the breakable prefill graph's
+ # comparison with eager run in QA (llm_function_core).
+ - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp3_gsm8k_and_long_context
- condition:
ranges:
system_gpu_count:
@@ -233,11 +231,11 @@ l0_b300:
- unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "TRTLLM"
- unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "CUTEDSL"
- unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "DEEPGEMM"
- # qwen4_exp target, post-merge: the opt-in variants -- three drafts over CuteDSL experts, fp8
- # KV and a bf16 recurrent state in one engine (gsm8k and the pass keys); seven drafts over
- # 3k-12k-token pass keys on NVFP4 KV and a bf16 state (the verify step's eight rows per
- # request); and the breakable prefill CUDA graph replayed against eager greedy continuations
- # at the bucket edges and with context and generation rows in one batch.
- - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_cutedsl_fp8_kv_bf16_state
- - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp7_long_context_passkey
- - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_breakable_prefill_matches_eager
+ # qwen4_exp target at tp1, post-merge: one draft with mmmu (the vision tower) and synthetic
+ # videos; three drafts over CuteDSL experts, fp8 KV and a bf16 recurrent state (gsm8k and
+ # the pass keys); seven drafts over 3k-12k-token pass keys on NVFP4 KV and a bf16 state (the
+ # verify step's eight rows per request). Its disaggregated engines run in
+ # l0_gb300_multi_gpus.
+ - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp1_gsm8k_mmmu_and_video
+ - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp3_cutedsl_fp8_kv_bf16_state
+ - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp7_long_context_passkey
diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml
index 09cbf54d5a74..2e816594cd1f 100644
--- a/tests/integration/test_lists/test-db/l0_cpu.yml
+++ b/tests/integration/test_lists/test-db/l0_cpu.yml
@@ -66,6 +66,9 @@ l0_cpu:
- unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_draft_exchange.py
- unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_small_batch_mtp.py
- unittest/_torch/modeling_v2/models/glm_moe_dsa/test_routing.py
+ # The qwen4_exp default target's cpu_only tests (its routing among them); this stage runs
+ # -m cpu_only, so the directory's GPU tests are left to l0_b300's modeling_v2 entry.
+ - unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default
- unittest/_torch/pyexecutor
- unittest/_torch/test_connector.py::test_a_partial_grouped_scheduler_is_rejected_at_bring_up
- unittest/_torch/test_connector.py::test_a_grouped_scheduler_and_a_single_group_cache_are_accepted
@@ -77,11 +80,8 @@ l0_cpu:
- unittest/_torch/modeling
# modeling_v2 routing and routing-table claims: config shapes and a monkeypatched
# compute capability, no device. The default targets' shard tests: small CPU
- # tensors and the CPU block_scale_interleave. The qwen4_exp entries are the target's
- # tests that run on the source text, stubbed weights and meta tensors: its multimodal
- # splice and admission check on a stub tower, the KV pool contract, its KV storage
- # derivation and attention-call arguments, its routed-expert scalars and its weight table.
- # The rest of unittest/_torch/modeling_v2 needs a GPU.
+ # tensors and the CPU block_scale_interleave. The rest of unittest/_torch/modeling_v2
+ # needs a GPU.
- unittest/_torch/modeling_v2/test_modeling_v2_claims.py
- unittest/_torch/modeling_v2/test_modeling_v2_native_claims.py
- unittest/_torch/modeling_v2/test_modeling_v2_native_loader.py
@@ -110,13 +110,6 @@ l0_cpu:
- unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_draft_attention_bcg_route.py
- unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py
- unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_dense_backend_library_order
- - unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_multimodal_embeds.py
- - unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_moe_scalars.py
- - unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_validate_request.py
- - unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_pool_contract.py
- - unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_kv_quant.py
- - unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_thop_attention_calls.py
- - unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_weight_table.py
- unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_fp8_indexer_rejects_different_qk_norm_epsilons
- unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_moe_reduces_only_local_terms
- unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_decoder_layer_sets_post_fusion_from_moe_scheduler
diff --git a/tests/integration/test_lists/test-db/l0_gb300.yml b/tests/integration/test_lists/test-db/l0_gb300.yml
index d860454553c8..3db15194ff62 100644
--- a/tests/integration/test_lists/test-db/l0_gb300.yml
+++ b/tests/integration/test_lists/test-db/l0_gb300.yml
@@ -123,6 +123,10 @@ l0_gb300:
- accuracy/test_llm_api_pytorch.py::TestDeepSeekV41::test_gsm8k_dspark[nvfp4-tp4] TIMEOUT (120)
- accuracy/test_llm_api_pytorch.py::TestDeepSeekV41::test_gsm8k_dspark[nvfp4-dp4] TIMEOUT (120)
- accuracy/test_glm53_flash.py::TestGLM53FlashFP8::test_mtp[tp_size=4-ep_size=4] TIMEOUT (60)
+ # qwen4_exp target, pre-merge: tep4 with three drafts and IndexShare, gsm8k and then
+ # 3k-6k-token pass keys with their acceptance-length floor. The other layouts and KV dtypes
+ # run post-merge (l0_gb300_multi_gpus).
+ - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_and_long_context TIMEOUT (60)
# K3 collectives and decode-path exchanges, four ranks in one NVLink domain.
- condition:
diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
index 95691e11722e..8707a434dabd 100644
--- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
+++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
@@ -33,6 +33,17 @@ l0_gb300_multi_gpus:
- accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_gsm8k_chunked_prefill[dep4]
- accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp4::test_gsm8k
- accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp4::test_eagle3_gsm8k_and_acceptance
+ # qwen4_exp target, post-merge: three drafts over NVFP4 KV at tp4 with the
+ # breakable prefill graph and over FP8 KV at dep4 (MTP under attention DP), both with block
+ # reuse and chunked prefill; and dep2 context and generation engines on one node. tep4 runs
+ # pre-merge (l0_gb300); mmmu runs in QA (llm_function_core).
+ - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[tp4_nvfp4_kv_breakable] TIMEOUT (60)
+ - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[dep4_fp8_kv] TIMEOUT (60)
+ - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_disaggregated_dep2 TIMEOUT (90)
+ # qwen4_exp target at tp1, post-merge: its AgentX engines (NVFP4 KV, MTP2 with rejection
+ # sampling, the breakable prefill graph, every opt-in projection) on one context and one
+ # generation GPU.
+ - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_disaggregated_mtp2 TIMEOUT (90)
# Covers the GPU cases of tests/unittest/_torch/attention/ on GB300, kernels/ and
# multi_gpu/ included; its cpu_only cases run on CPU (l0_cpu). The GPU cases of
# kernels/parallel_hw_agnostic/ run here too, although l0_b200 and l0_b300 ignore that
diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py
index 4bc00766af36..0610b489bc1a 100644
--- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py
+++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py
@@ -180,7 +180,7 @@ def _deviation(out: torch.Tensor, ref: torch.Tensor) -> str:
),
),
),
- "qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1": (
+ "qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default": (
Cell(
why="qwen4_exp main rope, decode: one token, Gemma norm and interleaved mRoPE "
"with three distinct position rows",
diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py
index ccbd73d80daa..b534b1a08b47 100644
--- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py
+++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py
@@ -7550,7 +7550,7 @@ def _spread(batch: int, longest: int) -> tuple[int, ...]:
),
),
),
- "qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1": (
+ "qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default": (
Cell(
why="qwen4_exp prefill: GQA 24q/2kv d256, bf16 pool, paged context; "
"a fresh 70-row sequence beside 33 rows over a 48-token cached prefix",
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_index_store.py
similarity index 52%
rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py
rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_index_store.py
index d6d7c5b03b47..8d840ae68ecb 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_index_store.py
@@ -8,18 +8,27 @@
its allocation, rows past the batch holding a previous step's plausible page
ids. Both planes are pre-filled with a sentinel so a stray write shows.
A `hole` cell appends one row that resolves to no slot.
+
+The default target's cells store into the stock manager's bf16 layout -- one
+layer's plane of the 12 sparse layers' coalesced index pool, the other 11
+checked untouched -- at the measured tep4 pool, with the batch's last row on the
+pool's last page. `test_the_stock_managers_planes_are_the_cells_layout` builds
+the stock manager itself (through the target's pool-layout test) and stores
+through its own views.
"""
+import importlib.util
import zlib
+from pathlib import Path
import pytest
import torch
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_index_store import ( # noqa: E501
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_index_store import ( # noqa: E501
qsa_paged_index_store as op,
)
-__extra_import_path__ = ["../../../.."]
+__extra_import_path__ = ["../../.."]
from _validating import certifying, validating # noqa: E402,F401
assert torch.cuda.is_available(), "qsa_paged_index_store requires a CUDA device"
@@ -80,8 +89,43 @@ def _shared_page_planes(spec, extras: dict):
return index_cache, position_cache
-def _build(spec, seed: int):
- """Real planes, table and rows for one cell; returns the op's kwargs and extras."""
+def _coalesced_planes(spec, extras: dict):
+ """The bf16 index pool of the default target: the `layers` sparse layers' INDEX_KEY buffers
+ coalesced into one allocation, so the cell's plane (layer `layer`) is a strided view, and the
+ position plane its own contiguous pool. The whole allocation holds `KEY_FILL`; the pool is
+ recorded for the untouched check."""
+ layers, layer, tpb, hkv, head_dim, axes, pages = (
+ spec[k] for k in ("layers", "layer", "tokens_per_block", "hkv", "head_dim", "axes", "pages")
+ )
+ pool = torch.full(
+ (pages, layers, tpb, hkv, head_dim), KEY_FILL, dtype=torch.bfloat16, device=DEV
+ )
+ index_cache = pool[:, layer]
+ assert index_cache.stride() == (spec["page_stride"], hkv * head_dim, head_dim, 1)
+ position_cache = torch.full((pages, tpb, axes), POS_FILL, dtype=torch.int32, device=DEV)
+ extras["coalesced"] = (pool, layer)
+ return index_cache, position_cache
+
+
+def _other_layers_untouched(pool: torch.Tensor, layer: int) -> None:
+ """Every other layer's plane in the coalesced pool still holds `KEY_FILL`, checked a band of
+ pages at a time so the check stays a fraction of the pool's size."""
+ fill = int(_bits(torch.tensor([KEY_FILL], dtype=torch.bfloat16))[0])
+ bits = pool.view(torch.int16)
+ for start in range(0, pool.shape[0], 16384):
+ band = bits[start : start + 16384]
+ for other in range(pool.shape[1]):
+ if other != layer:
+ assert bool((band[:, other] == fill).all()), (
+ f"layer {other}'s index plane moved in pages {start}..{start + band.shape[0]}"
+ )
+
+
+def _build(spec, seed: int, planes=None):
+ """Real planes, table and rows for one cell; returns the op's kwargs and extras.
+
+ `planes`, if given, is an `(index_cache, position_cache)` pair to store into instead of the
+ cell's own (a stock manager's views)."""
gen = torch.Generator(device="cpu").manual_seed(seed)
tpb, hkv, head_dim, axes = (spec[k] for k in ("tokens_per_block", "hkv", "head_dim", "axes"))
slots, width, pages = spec["slots"], spec["width"], spec["pages"]
@@ -99,7 +143,17 @@ def _build(spec, seed: int):
table[r, :cols] = order[used : used + cols]
used += cols
allocated_cols.append(cols)
- assert used <= pages and len(segments) < slots, "pool or table too small for the batch"
+ assert used <= pages and len(segments) <= slots, "pool or table too small for the batch"
+ if spec.get("straddle"):
+ # The batch's last row stores into the pool's last page; every other row keeps its page
+ # from the permutation, so the rows land on both sides of 2^31 elements of the plane.
+ last_start, last_rows = segments[-1]
+ column = (last_start + last_rows - 1) // tpb
+ live = table[: len(segments)]
+ holder = (live == pages - 1).nonzero()
+ if len(holder):
+ live[holder[0, 0], holder[0, 1]] = live[len(segments) - 1, column]
+ live[len(segments) - 1, column] = pages - 1
if spec["hole"] is not None:
request, position = _hole(spec, table, allocated_cols)
requests.append(request)
@@ -113,9 +167,12 @@ def _build(spec, seed: int):
token_k = projection[:, spec["index_heads"] :]
if spec["key_layout"] == "fused_row":
# The target's fused input projection row, [q | k | v | gate | index q, k |
- # pad] = 110 heads of 128: the raw index key is head 108.
- row = torch.randn(rows, 110, head_dim, generator=gen, dtype=torch.float32)
- token_k = row.to(device=DEV, dtype=spec["key_dtype"])[:, 108 : 108 + hkv]
+ # pad of one index head], in heads of 128: the raw index key is the second to
+ # last (108 of 110 at 14080 wide; 32 of 34 at the default target's 4352).
+ heads = spec.get("row_width", 14080) // head_dim
+ row = torch.randn(rows, heads, head_dim, generator=gen, dtype=torch.float32)
+ token_k = row.to(device=DEV, dtype=spec["key_dtype"])[:, heads - 2 : heads - 2 + hkv]
+ assert token_k.stride(0) == heads * head_dim
request_indices = torch.tensor(requests, dtype=spec["request_dtype"], device=DEV)
logical_positions = torch.tensor(positions, dtype=spec["position_dtype"], device=DEV)
@@ -131,8 +188,12 @@ def _build(spec, seed: int):
extras = {}
position_cache = None
- if spec["cache_layout"] == "shared_page":
+ if planes is not None:
+ index_cache, position_cache = planes
+ elif spec["cache_layout"] == "shared_page":
index_cache, position_cache = _shared_page_planes(spec, extras)
+ elif spec["cache_layout"] == "coalesced":
+ index_cache, position_cache = _coalesced_planes(spec, extras)
else:
index_cache = torch.full((pages, tpb, hkv, head_dim), KEY_FILL, device=DEV).to(
torch.bfloat16
@@ -156,6 +217,9 @@ def _build(spec, seed: int):
def test_certified_cells(cell, certifying): # noqa: F811
certifying(op)
args, extras = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
+ if cell.spec.get("straddle"):
+ # The last row's page starts past 2^31 elements of the plane it is stored through.
+ assert (cell.spec["pages"] - 1) * args["index_cache"].stride(0) >= 2**31
ref = op.reference(**args)
out = op(**args)
assert out[0] is args["index_cache"] and out[1] is args["position_cache"]
@@ -165,6 +229,8 @@ def test_certified_cells(cell, certifying): # noqa: F811
assert bool((extras["roles"].view(torch.int16) == fill).all()), (
"another index-K role of the shared page moved"
)
+ if "coalesced" in extras:
+ _other_layers_untouched(*extras["coalesced"])
#: The cell of 32 MTP7 verify blocks (256 rows): the segments, pool and table it runs over.
@@ -178,39 +244,57 @@ def _small(seed: int, **overrides):
return _build({**spec, **overrides}, seed=seed)[0]
-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_()
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(warmup):
- for _ in range(3):
- op(**args)
- torch.cuda.current_stream().wait_stream(warmup)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- op(**args)
-
- positions.copy_(replayed)
- args["index_cache"].fill_(KEY_FILL)
- args["position_cache"].fill_(POS_FILL)
- ref = op.reference(**args)
- graph.replay()
- torch.cuda.synchronize()
- _compare((args["index_cache"], args["position_cache"]), ref)
-
-
-def test_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]
- before = (args["index_cache"].clone(), args["position_cache"].clone())
- with pytest.raises(ValueError, match="equally long"):
+def _pool_layout():
+ """The target's pool-layout test, loaded by path: its `_manager` builds the stock manager the
+ engine builds on a rank of a member (core cache config, bf16 KV), and `ATTN_LAYERS` names the
+ sparse layers. Loaded rather than restated so the two cannot disagree on what the stock pool
+ is."""
+ path = (
+ Path(__file__).resolve().parents[1]
+ / "qwen3_8_flash_next_nvfp4__sm_103__default"
+ / "test_pool_layout.py"
+ )
+ spec = importlib.util.spec_from_file_location("_qsa_index_store_pool_layout", path)
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+@pytest.mark.parametrize("member", ("tp4", "tep4", "dep4"))
+def test_the_stock_managers_planes_are_the_cells_layout(member) -> None:
+ """The stock manager's own index and position planes at rank 0 of `member` have the layout
+ the member's cells model, and the op stores through them with the member's row width: four
+ rows, the last on the pool's last page, compared against `reference` over the whole plane,
+ every other sparse layer's plane untouched."""
+ layout = _pool_layout()
+ cell = next(c.spec for c in op.CELLS if c.spec.get("member") == member)
+ tpb, hkv, head_dim, axes = (cell[k] for k in ("tokens_per_block", "hkv", "head_dim", "axes"))
+ with layout._manager(member, 0) as cm:
+ planes = [cm.get_index_k_buffer(i) for i in layout.ATTN_LAYERS]
+ index_cache, position_cache = planes[cell["layer"]], cm.get_qsa_position_buffer()
+ assert len(planes) == cell["layers"]
+ assert index_cache.stride() == (cell["page_stride"], hkv * head_dim, head_dim, 1)
+ assert position_cache.stride() == (tpb * axes, axes, 1)
+ assert position_cache.shape[0] == index_cache.shape[0]
+ plane_bytes = tpb * hkv * head_dim * index_cache.element_size()
+ offsets = [p.data_ptr() - planes[0].data_ptr() for p in planes]
+ assert offsets == [j * plane_bytes for j in range(len(planes))]
+
+ for plane in planes:
+ plane.fill_(KEY_FILL)
+ position_cache.fill_(POS_FILL)
+ spec = dict(
+ cell,
+ segments=((0, 1), (31, 1), (33, 2)),
+ pages=index_cache.shape[0],
+ slots=8,
+ )
+ args, _ = _build(spec, seed=71, planes=(index_cache, position_cache))
+ assert int(args["block_table"][2, 1]) == index_cache.shape[0] - 1
+ ref = op.reference(**args)
op(**args)
- assert torch.equal(_bits(args["index_cache"]), _bits(before[0]))
- assert torch.equal(args["position_cache"], before[1])
+ _compare((index_cache, position_cache), ref)
+ fill = int(_bits(torch.tensor([KEY_FILL], dtype=torch.bfloat16))[0])
+ for j, plane in enumerate(planes):
+ if j != cell["layer"]:
+ assert bool((_bits(plane) == fill).all()), f"sparse layer {j}'s plane moved"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py
similarity index 72%
rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py
rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py
index 4278c0822fcc..6c5e478ff738 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py
@@ -4,10 +4,9 @@
The experts use the shared test helpers' drawn NVFP4 codes and
power-of-two scales, so every weight is exact; this entry's layout is built
-from their raw per-expert tensors by the target's own `weights.py` transforms,
+from their raw per-expert tensors by the default target's `weights.py` transforms,
so the test also pins those. Per cell: inside the routed entry's band against
-the fp32 reference. Beyond the cells: the guard against a per-expert FC2 input
-scale.
+the fp32 reference.
"""
import zlib
@@ -16,17 +15,17 @@
import pytest
import torch
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501
- weights as W,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_moe_renorm_cute_dsl import ( # noqa: E501
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_moe_renorm_cute_dsl import ( # noqa: E501
fp4_moe_renorm_cute_dsl as op,
)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.renorm_moe_routing_op import ( # noqa: E501
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.renorm_moe_routing_op import ( # noqa: E501
renorm_moe_routing_op,
)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501
+ weights as W,
+)
-__extra_import_path__ = ["../../../.."]
+__extra_import_path__ = ["../../.."]
import _fp4_moe_test_utils as H # noqa: E402
from _fp4_moe_test_utils import many_expert_layer # noqa: E402,F401 — fixture by name
from _precision import true_fp32_matmul # noqa: E402
@@ -84,27 +83,3 @@ def test_certified_cells(cell, certifying, many_expert_layer, cute_dsl_experts):
with true_fp32_matmul():
exp = op.reference(logits, x, spec["top_k"], ref["up"], ref["gate"], ref["down"], g2)
H._compare(y, exp.to(torch.bfloat16))
-
-
-def test_guard_refuses_a_per_expert_fc2_scale(many_expert_layer, cute_dsl_experts): # noqa: F811
- spec = {**op.CELLS[0].spec, "tokens": 1024}
- args, ref, data, sf, x, logits, _, _, g2, _ = H.renorm_operands(
- many_expert_layer(), spec, 9, renorm_moe_routing_op
- )
- fc1_w, fc1_s, fc2_w, fc2_s = cute_dsl_experts
- s1, sg, s2 = H._scalars(spec["num_experts"], G1, g2)
- with validating(op), pytest.raises(AssertionError):
- op(
- logits,
- data,
- sf,
- fc1_w,
- fc1_s,
- sg,
- fc2_w,
- fc2_s,
- s2,
- s2,
- spec["num_experts"],
- spec["top_k"],
- )
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/__init__.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/__init__.py
new file mode 100644
index 000000000000..113c2e6ba7f9
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/__init__.py
@@ -0,0 +1,7 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Package marker.
+
+`test_routing.py` repeats the glm_moe_dsa family's basename. Without a package, pytest's
+default import mode loads both under one module name and refuses the second.
+"""
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_nvfp4_pages_to_bf16.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_qsa_fp8_pages_to_bf16__prefill.py
similarity index 52%
rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_nvfp4_pages_to_bf16.py
rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_qsa_fp8_pages_to_bf16__prefill.py
index 99c939511024..74f52efb90b5 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_nvfp4_pages_to_bf16.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_qsa_fp8_pages_to_bf16__prefill.py
@@ -1,14 +1,14 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the qsa_nvfp4_pages_to_bf16 entry.
-
-The paged state is real: one NVFP4 pool as the QSA cache manager serves it -- e2m1 data
-`[pages, role pages, Hkv, tpb, D/2]` and e4m3 block scales `[pages, role pages, Hkv, tpb,
-D/16]`, the layer's K and V two of the role planes -- and the engine's `[slots, width]`
-table, the listed requests on seeded distinct slots, their pages a seeded permutation,
-`-1` past each sequence. Scratch and table are pre-filled with sentinels, so a missing or
-a stray write shows. Every cell is bitwise `reference`, and the scratch is what the NVFP4
-sparse GQA kernels dequantize.
+"""GPU test for this target's qsa_fp8_pages_to_bf16 entry (`__prefill`).
+
+The paged state is the engine's kv_fp8 trunk: e4m3 `[pages, 24, Hkv, tpb, D]`, layer 11's K
+and V two of the role planes (a straddle cell sized so the requests' pages lie on both sides
+of 2^31 elements of them), and the `[slots, 8196]` table at the tune's max_seq_len, the
+listed requests on seeded distinct slots, their pages a seeded permutation, `-1` past each
+sequence. Scratch and table are pre-filled with sentinels, so a missing or a stray write
+shows. Every cell is bitwise `reference` (a widening copy), and the scratch is what the
+e4m3 sparse GQA kernel reads: e4m3 codes at unit scale.
"""
import math
@@ -17,21 +17,29 @@
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 _gqa,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_nvfp4_pages_to_bf16 import ( # noqa: E501
- qsa_nvfp4_pages_to_bf16 as op,
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.attention.qsa_fp8_pages_to_bf16__prefill import ( # noqa: E501
+ qsa_fp8_pages_to_bf16 as op,
)
__extra_import_path__ = ["../../../.."]
from _validating import certifying, validating # noqa: E402,F401
-assert torch.cuda.is_available(), "qsa_nvfp4_pages_to_bf16 requires a CUDA device"
+assert torch.cuda.is_available(), "qsa_fp8_pages_to_bf16 requires a CUDA device"
DEV = "cuda"
SCRATCH_FILL = 0x7FC1 # a bf16 NaN
TABLE_FILL = -7
+_EDGE = 2**31
+
+
+def _pool_pages(spec, needed: int) -> int:
+ if not spec["straddle"]:
+ return needed + 5
+ role_page = spec["kv_heads"] * spec["tokens_per_block"] * spec["head_dim"]
+ edge_page = math.ceil(
+ (_EDGE - 2 * spec["layer"] * role_page) / (spec["slot_pages"] * role_page)
+ )
+ return max(edge_page + needed // 2, needed) + 5
def _build(spec, seed):
@@ -39,38 +47,36 @@ def _build(spec, seed):
tpb, kv, d = spec["tokens_per_block"], spec["kv_heads"], spec["head_dim"]
counts = [math.ceil(n / tpb) for n in spec["kv_lens"]]
needed = sum(counts) + 7
- pool_pages = needed + 5
+ pool_pages = _pool_pages(spec, needed)
role_pages, layer = spec["slot_pages"], spec["layer"]
- data = torch.randint(
- 0, 256, (pool_pages, role_pages, kv, tpb, d // 2), generator=gen, device=DEV
- ).to(torch.uint8)
- scales = torch.randint(
- 0x08, 0x31, (pool_pages, role_pages, kv, tpb, d // 16), generator=gen, device=DEV
- ).to(torch.uint8)
+ # Random e4m3 codes (NaN codes 0x7F / 0xFF excluded), in the layer's two planes only.
+ data = torch.zeros((pool_pages, role_pages, kv, tpb, d), dtype=torch.uint8, device=DEV)
+ first = pool_pages - needed
+ for p0 in range(first, pool_pages, 1024):
+ p1 = min(p0 + 1024, pool_pages)
+ codes = torch.randint(0, 254, (p1 - p0, 2, kv, tpb, d), generator=gen, device=DEV)
+ codes = torch.where(codes >= 0x7F, codes + 1, codes)
+ data[p0:p1, 2 * layer : 2 * layer + 2] = codes.to(torch.uint8)
+ pool = data.view(torch.float8_e4m3fn)
width, slots_total = spec["width"], spec["slots"]
table = torch.full((slots_total, width), -1, dtype=torch.int32, device=DEV)
slots = torch.randperm(slots_total, generator=torch.Generator().manual_seed(seed))
slots = slots[: len(counts)]
- page_ids = torch.randperm(pool_pages, generator=gen, device=DEV).to(torch.int32)
+ page_ids = first + torch.randperm(needed, generator=gen, device=DEV).to(torch.int32)
cursor = 0
for slot, count in zip(slots.tolist(), counts):
table[slot, :count] = page_ids[cursor : cursor + count]
cursor += count
if spec["holes"]:
- first = slots[0].item()
- table[first, 1] = -1 # unmapped inside the sequence
- table[first, 2] = pool_pages + 3 # past the pool
+ head = slots[0].item()
+ table[head, 1] = -1 # unmapped inside the sequence
+ table[head, 2] = pool_pages + 3 # past the pool
counts_t = torch.tensor(counts, dtype=torch.int64, device=DEV)
listed = int(counts_t.sum())
out_shape = (listed, kv, tpb, d)
- args = dict(
- 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
- ),
+ return dict(
+ k_cache=pool[:, 2 * layer],
+ v_cache=pool[:, 2 * layer + 1],
block_table=table,
slots=slots.to(DEV, torch.int32),
counts=counts_t,
@@ -83,7 +89,6 @@ def _build(spec, seed):
torch.bfloat16
),
)
- return args
def _bitwise(a: torch.Tensor, b: torch.Tensor) -> bool:
@@ -92,10 +97,21 @@ def _bitwise(a: torch.Tensor, b: torch.Tensor) -> bool:
return a.shape == b.shape and torch.equal(a, b)
+def _assert_straddles(args) -> None:
+ table = args["block_table"]
+ pages = table[(table >= 0) & (table < args["k_cache"].shape[0])].long()
+ for name in ("k_cache", "v_cache"):
+ plane = args[name]
+ ends = plane.storage_offset() + pages * plane.stride(0) + plane[0].numel() - 1
+ assert bool((ends >= _EDGE).any()) and bool((ends < _EDGE).any()), name
+
+
@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
def test_certified_cells(cell, certifying) -> None: # noqa: F811 -- fixture by name
certifying(op)
args = _build(cell.spec, zlib.crc32(cell.why.encode()))
+ if cell.spec["straddle"]:
+ _assert_straddles(args)
ref = op.reference(**{k: v.clone() if torch.is_tensor(v) else v for k, v in args.items()})
out = op(**args)
for name, got, want in zip(("k_out", "v_out", "table"), out, ref):
@@ -108,32 +124,19 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 -- fixture by
assert bool((out[2][~listed_rows] == TABLE_FILL).all())
-def test_scratch_is_what_the_sparse_gqa_kernels_read() -> None:
- """The scratch holds `_dequantize_tokens`' values (the NVFP4 sparse GQA entry's own
- dequantize, fp32 rounded once to bf16) for every listed page."""
- spec = op.CELLS[4].spec
- args = _build({**spec, "global_k": 0.37, "global_v": 2.5}, 11)
+def test_scratch_is_the_e4m3_values_at_unit_scale() -> None:
+ """Every listed page's scratch holds the pool's e4m3 values exactly (e4m3 widens to bf16
+ without rounding), and the table addresses it."""
+ spec = op.CELLS[3].spec
+ args = _build(spec, 11)
k_out, v_out, table = op(**args)
- tpb = spec["tokens_per_block"]
for slot, count, base in zip(
args["slots"].tolist(), args["counts"].tolist(), args["offsets"].tolist()
):
pages = args["block_table"][slot, :count].long()
- tokens = torch.arange(tpb, device=DEV)[None, :].expand(count, -1)
- for got, data, scales, index, grouped in (
- (k_out, args["k_cache"], args["k_scale_cache"], 1, False),
- (v_out, args["v_cache"], args["v_scale_cache"], 2, True),
- ):
- want = _gqa._dequantize_tokens(
- data,
- scales,
- pages[:, None].expand(-1, tpb),
- tokens,
- args["kv_scale_quant_orig"][index],
- grouped,
- ) # [count, tpb, H, D]
- want = want.transpose(1, 2).to(torch.bfloat16)
- assert _bitwise(got[base : base + count], want)
+ for got, plane in ((k_out, args["k_cache"]), (v_out, args["v_cache"])):
+ want = plane[pages].float()
+ assert torch.equal(got[base : base + count].float(), want)
assert torch.equal(
table[slot, :count], torch.arange(base, base + count, device=DEV, dtype=torch.int32)
)
@@ -141,21 +144,10 @@ def test_scratch_is_what_the_sparse_gqa_kernels_read() -> None:
def test_listing_past_the_scratch_is_dropped() -> None:
"""A request whose pages run past `listed` writes only those below it."""
- spec = op.CELLS[2].spec
- args = _build(spec, 5)
+ args = _build(op.CELLS[3].spec, 5)
short = args["k_out"].shape[0] - 3
args["k_out"], args["v_out"] = args["k_out"][:short], args["v_out"][:short]
ref = op.reference(**{k: v.clone() if torch.is_tensor(v) else v for k, v in args.items()})
out = op(**args)
for got, want in zip(out, ref):
assert _bitwise(got, want)
-
-
-def test_guards() -> None:
- args = _build(op.CELLS[0].spec, 3)
- with validating(op), pytest.raises(AssertionError):
- op(**{**args, "k_out": args["k_out"].float()})
- with validating(op), pytest.raises(AssertionError):
- op(**{**args, "block_table": args["block_table"].long()})
- with validating(op), pytest.raises(AssertionError):
- op(**{**args, "table": args["table"][:, :-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_sparse_gqa_rowtile.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_qsa_sparse_gqa_rowtile__prefill.py
similarity index 58%
rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py
rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_qsa_sparse_gqa_rowtile__prefill.py
index a79e34fa0142..62429f6004ac 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__default/attention/test_modeling_v2_qsa_sparse_gqa_rowtile__prefill.py
@@ -1,27 +1,32 @@
# 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.
+"""GPU test for this target's qsa_sparse_gqa_rowtile variant (`__prefill`).
+
+The paged state is the target's: a trunk cell's K and V are layer 11's planes of a `[pages,
+24, kv, tpb, D]` HND pool (24 role pages a slot), sized so the requests' pages straddle 2^31
+elements, every other page and plane NaN; a scratch cell's are a contiguous `[pages, kv, tpb,
+D]` bf16 pool, as `qsa_fp8_pages_to_bf16` widens a kv_fp8 pool into. The slot table is the
+engine's at the tune's max_seq_len (8196 columns, a row per slot, a seeded permutation, -1
+past each sequence). Selection rows are laid out the way QSA's top-k + expansion emit them
+(complete groups expanded to tokens, the trailing partial group, then -1, packed at the
+front). The band and the reference are `qsa_paged_sparse_gqa`'s: this op replaces it on
+context launches.
"""
import math
-import 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__default.ops.attention import ( # noqa: E501
+ qsa_sparse_gqa_rowtile__prefill as _module,
)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_sparse_gqa_rowtile import ( # noqa: E501
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.attention.qsa_sparse_gqa_rowtile__prefill import ( # noqa: E501
qsa_sparse_gqa_rowtile as op,
)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.attention.qsa_sparse_gqa_rowtile__prefill import ( # noqa: E501
+ rows_per_group,
+)
__extra_import_path__ = ["../../../.."]
from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401
@@ -32,8 +37,10 @@
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."""
+ """`qsa_paged_sparse_gqa`'s band: 4 ulp of each (row, head)'s magnitude, 2 ulp relative
+ RMS, in q's dtype, against the fp32 reference rounded once. The row-tile pass rounds P to
+ bf16 for the PV MMA and sums it in fp32, as the per-row kernel rounds its probabilities to
+ the cache dtype before P.V."""
assert out.dtype == ref.dtype, (out.dtype, ref.dtype)
assert out.shape == ref.shape, (out.shape, ref.shape)
if out.numel() == 0:
@@ -42,8 +49,8 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
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."""
+ """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
@@ -75,11 +82,10 @@ def launches(monkeypatch):
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)."""
+ """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)
@@ -90,8 +96,8 @@ def _selected_tokens(positions, seed, compress_ratio, token_topk, shared_scores=
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)
+ for start in range(0, rows, 1024):
+ stop = min(start + 1024, rows)
avail, take, vis = available[start:stop], taken[start:stop], visible[start:stop]
if shared_scores:
scores = torch.rand(1, width, generator=gen, device=DEV).expand(stop - start, -1)
@@ -112,6 +118,20 @@ def _selected_tokens(positions, seed, compress_ratio, token_topk, shared_scores=
return out
+def _nan_pool(shape, dtype):
+ assert dtype == torch.bfloat16
+ return torch.full(shape, 0x7FC0, dtype=torch.int16, device=DEV).view(dtype)
+
+
+def _straddling(spec, needed: int) -> int:
+ """The pool page count that puts the requests -- the pool's last `needed` pages -- on
+ both sides of `offset_edge` elements of layer `layer`'s K plane."""
+ role_page = spec["kv_heads"] * spec["tokens_per_block"] * spec["head_dim"]
+ edge, slot_pages, layer = spec["offset_edge"], spec["slot_pages"], spec["layer"]
+ edge_page = math.ceil((edge - 2 * layer * role_page) / (slot_pages * role_page))
+ return max(edge_page + needed // 2, needed)
+
+
def _build(spec, seed):
"""Real paged state for one spec, keyed by `__call__`'s parameter names."""
gen = torch.Generator(device=DEV).manual_seed(seed)
@@ -129,17 +149,34 @@ def _build(spec, seed):
slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(segments)]
mapped = [math.ceil((s + n) / tpb) for s, n in segments]
assert max(mapped) <= spec["pages"]
- pool_pages = sum(mapped) + 7
- page_ids = torch.randperm(pool_pages, generator=gen, device=DEV).to(torch.int32)
+ needed = sum(mapped) + 7
+ trunk = spec["cache_layout"] == "trunk"
+ pool_pages = _straddling(spec, needed) if trunk else needed
+ first = pool_pages - needed
+ page_ids = first + torch.randperm(needed, generator=gen, device=DEV).to(torch.int32)
block_table = torch.full((slots, spec["pages"]), -1, dtype=torch.int32, device=DEV)
cursor = 0
for slot, count in zip(slot_of.tolist(), mapped):
block_table[slot, :count] = page_ids[cursor : cursor + count]
cursor += count
- pool = torch.randn(pool_pages, 2, kv, tpb, d, generator=gen, device=DEV).to(spec["cache_dtype"])
- k_cache, v_cache = pool[:, 0], pool[:, 1]
- assert spec["cache_layout"] == "pool_planes" and spec["q_layout"] == "qkv_slice"
+ dtype = spec["cache_dtype"]
+ if trunk:
+ # Layer `layer`'s K and V planes of the trunk; every other page and plane is NaN.
+ layer = spec["layer"]
+ pool = _nan_pool((pool_pages, spec["slot_pages"], kv, tpb, d), dtype)
+ for p0 in range(first, pool_pages, 1024):
+ p1 = min(p0 + 1024, pool_pages)
+ pool[p0:p1, 2 * layer : 2 * layer + 2] = torch.randn(
+ (p1 - p0, 2, kv, tpb, d), generator=gen, device=DEV
+ ).to(dtype)
+ k_cache, v_cache = pool[:, 2 * layer], pool[:, 2 * layer + 1]
+ else:
+ # The widened scratch: two contiguous bf16 pools.
+ assert spec["cache_layout"] == "scratch"
+ k_cache = torch.randn((pool_pages, kv, tpb, d), generator=gen, device=DEV).to(dtype)
+ v_cache = torch.randn((pool_pages, kv, tpb, d), generator=gen, device=DEV).to(dtype)
+ assert spec["q_layout"] == "qkv_slice"
qkv = torch.randn(rows, q_heads + 2 * kv, d, generator=gen, device=DEV).to(spec["q_dtype"])
q = qkv[:, :q_heads]
@@ -172,11 +209,24 @@ def _is_prefix_packed(table):
return (valid == (columns < valid.sum(1)[:, None])).all(1)
+def _assert_straddles(args, edge: int) -> None:
+ """The mapped pages put the K and V planes' element offsets on both sides of `edge`."""
+ table = args["block_table"]
+ pages = table[table >= 0].long()
+ for name in ("k_cache", "v_cache"):
+ plane = args[name]
+ ends = plane.storage_offset() + pages * plane.stride(0) + plane[0].numel() - 1
+ assert bool((ends >= edge).any()), f"{name}: no mapped page lies past {edge}"
+ assert bool((ends < edge).any()), f"{name}: every mapped page lies past {edge}"
+
+
@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811
certifying(op)
spec = cell.spec
args = _build(spec, seed=sum(map(ord, cell.why)))
+ if spec["cache_layout"] == "trunk":
+ _assert_straddles(args, spec["offset_edge"])
rows = args["q"].shape[0]
# The per-row kernel's precondition, which this op shares.
selected = args["selected_tokens"]
@@ -184,14 +234,15 @@ def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811
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)
+ group_rows = rows_per_group(spec["q_heads"] // spec["kv_heads"])
+ assert group_rows == {12: 10, 6: 20}[spec["q_heads"] // spec["kv_heads"]]
+ groups = math.ceil(rows / group_rows)
assert launches[0][1] == (rows,) and launches[1][1] == (groups,)
assert launches[2][1] == (min(rows, 256), spec["kv_heads"])
# The row-tile pass is persistent: one CTA per SM, never more than (group, KV head) items.
items = groups * spec["kv_heads"]
- assert launches[3][1] == (_module._rowtile_ctas(items, args["q"].device),)
- assert launches[3][1][0] == min(
- items, torch.cuda.get_device_properties(0).multi_processor_count
+ assert launches[3][1] == (
+ min(items, torch.cuda.get_device_properties(0).multi_processor_count),
)
assert out.shape == (rows, spec["q_heads"], spec["head_dim"])
assert out.dtype == args["q"].dtype and out.is_contiguous()
@@ -199,39 +250,26 @@ def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811
_compare(out, op.reference(**args))
-CKPT = op.CELLS[0].spec
-assert (CKPT["q_heads"], CKPT["kv_heads"], CKPT["head_dim"]) == (24, 2, 256)
+TP4 = op.CELLS[0].spec
+assert (TP4["q_heads"], TP4["kv_heads"], TP4["head_dim"]) == (6, 1, 256)
def _with(segments, **overrides):
- return {**CKPT, "segments": tuple(segments), **overrides}
+ return {**TP4, "segments": tuple(segments), **overrides}
@pytest.mark.parametrize(
"segments",
[
- # Fewer items (2 groups x 2 KV heads) than SMs: a 4-CTA grid, mostly one item a CTA.
- [(5000, 17)],
+ # Fewer items (2 groups x 1 KV head) than SMs: a 2-CTA grid, one item a CTA.
+ [(5000, 37)],
# A lone group (rows of two requests) between row-tile groups of each request.
- [(3000, 25), (7000, 25)],
+ [(3000, 50), (7000, 50)],
],
ids=["fewer items than SMs", "a lone group between row-tile groups"],
)
def test_small_launches_match_the_reference(segments) -> None:
- """The persistent pass's edges: CTAs left without a second item, and an item the pass
- skips (a lone group) between items it serves."""
+ """The persistent pass's edges at 20-row groups: CTAs left without a second item, and
+ an item the pass skips (a lone group) between items it serves."""
args = _build(_with(segments, slots=8), seed=7700 + len(segments))
_compare(op(**args), op.reference(**args))
-
-
-def test_calls_are_bitwise_deterministic() -> None:
- args = _build(_with([(0, 2000), (500, 300), (7000, 1)], slots=8), seed=8800)
- assert torch.equal(op(**args), op(**args))
-
-
-def test_rejects_a_call_without_the_prefill_pair() -> None:
- """The op's own loud refusal: the row-tile pass is bounded by the prefill pair."""
- args = _build(_with([(600, 4)], slots=8), seed=9100)
- args.update(query_positions=None, compress_ratio=None)
- with pytest.raises(ValueError, match=re.escape("the row-tile kernel needs the prefill pair")):
- op(**args)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue__decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue__decode.py
new file mode 100644
index 000000000000..db073f810319
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue__decode.py
@@ -0,0 +1,261 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the default target's qwen4_exp_attn_decode_prologue entry: every cell against
+the entry's fp32 `reference` on every allocation either writes -- the packed QKV, the index
+query, the K/V pool, the index plane, the position cache -- and a stale-read check under
+graph replay with PDL.
+
+The rotated values (the packed q and k, the index query, the K pool's k and the compressed
+keys) are judged in units of their row's scale: the kernel and the reference evaluate the
+same fp32 formulas and differ only in the rms's reduction order, which moves a value across
+a bf16 rounding boundary at most once -- one bf16 spacing of an element no larger than the
+row's largest, at most 2^-7 of the row scale (2 ulp in `assert_within_ulp`'s 2^-8 units).
+Everything else -- v in the packed buffer and the V pool, the raw index keys, the
+coordinates and every byte neither side writes -- is compared bitwise, and the K pool's
+rows are bitwise the kernel's own packed k converted to the pool's dtype.
+
+The pools are laid out as the engine serves them: K/V the two role planes of a
+`[pages, 2, Hkv, tpb, D]` pool, requests on random pages."""
+
+import zlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.attention.qwen4_exp_attn_decode_prologue__decode import ( # noqa: E501
+ qwen4_exp_attn_decode_prologue as op,
+)
+
+__extra_import_path__ = ["../../../.."]
+from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401
+
+assert torch.cuda.is_available(), "qwen4_exp_attn_decode_prologue requires a CUDA device"
+
+DEV = "cuda"
+D, IH, ID, TPB, CR, ROT, WIDTH, SLOTS = 256, 4, 128, 32, 4, 64, 8196, 8
+#: Per-rank geometry: (q heads, KV heads); the fused row is
+#: `[q | k | v | gate | index q 4 x 128, k 128 | v slot 128]`.
+GEOMS = {"tp4": (6, 1), "dep4": (24, 2)}
+#: One bf16 spacing of an element no larger than its row's largest, in 2^-8 units of the
+#: row scale; the relative RMS bound is measured, see the module docstring.
+ELEMENT_ULP, RMS_ULP = 2.0, 0.05
+
+
+def _geometry(geom):
+ hq, hk = GEOMS[geom]
+ qkv = (hq + 2 * hk) * D
+ iq_col = qkv + hq * D
+ return hq, hk, qkv, iq_col, iq_col + (IH + 2) * ID
+
+
+def _edge_rows(proj, hq, hk) -> None:
+ """Zero, -0.0 and outlier blocks in the k head (past the rotary prefix, where the norm
+ keeps zeros zero) and in the v head, over the first rows."""
+ k0, v0 = hq * D, (hq + hk) * D
+ edges = (
+ (0, slice(k0 + 128, k0 + 144), 0.0),
+ (1, slice(k0 + 160, k0 + 176), -0.0),
+ (2, slice(k0 + 200, k0 + 201), 900.0),
+ (0, slice(v0 + 16, v0 + 32), 0.0),
+ (1, slice(v0 + 48, v0 + 64), torch.tensor([-0.0] * 8 + [0.0] * 8)),
+ (3, slice(v0 + 40, v0 + 41), -900.0),
+ (2, slice(v0 + 64, v0 + 80), 1e-30),
+ )
+ for r, cols, value in edges:
+ if r < proj.shape[0]:
+ proj[r, cols] = value
+
+
+def _rows(spec):
+ """(request, position) per row: one row a request, or `rows_per_request` consecutive
+ positions of each request from its first."""
+ per = spec.get("rows_per_request", 1)
+ return [(r, p + i) for r, p in enumerate(spec["pos"]) for i in range(per)]
+
+
+def _build(spec, seed):
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ hq, hk, qkv_width, iq_col, row_width = _geometry(spec["geom"])
+ layout = _rows(spec)
+ pos = [p for _, p in layout]
+ rows = len(pos)
+ last = {}
+ for req, p in layout:
+ last[req] = max(last.get(req, 0), p)
+ slots = max(SLOTS, len(last) + 3)
+ pages = sum(p // TPB + 1 for p in last.values()) + 64
+ perm = torch.randperm(pages, generator=g, device=DEV).to(torch.int32)
+ block_table = torch.full((slots, WIDTH), -1, dtype=torch.int32, device=DEV)
+ used = 0
+ for req, p in last.items():
+ n = p // TPB + 1
+ block_table[req, :n] = perm[used : used + n]
+ used += n
+ request = torch.tensor([req for req, _ in layout], dtype=torch.int32, device=DEV)
+ if spec.get("unmapped") is not None:
+ request[spec["unmapped"]] = slots + 5
+ logical = torch.tensor(pos, dtype=torch.int64, device=DEV)
+ ids = torch.stack([logical, logical + 3, logical + 7]).to(torch.int32).contiguous()
+ buf = torch.empty(3, 1, rows + 5, dtype=torch.int32, device=DEV)
+ buf[:, 0, :rows] = ids
+ coords = torch.transpose(buf[:, 0, :rows], 0, 1)
+ table_rows = max(pos) + 8
+ angles = torch.rand(table_rows, ROT // 2, generator=g, device=DEV) * 6.28
+ cos_sin = torch.cat([torch.cos(angles), torch.sin(angles)], dim=1).contiguous()
+ proj = (torch.randn(rows, row_width, generator=g, device=DEV) * 2.0).to(torch.bfloat16)
+ if spec.get("edge"):
+ _edge_rows(proj, hq, hk)
+ out = torch.randn(rows, qkv_width + 64, generator=g, device=DEV).to(torch.bfloat16)
+ index_scratch = torch.randn(rows, IH + 2, ID, generator=g, device=DEV).to(torch.bfloat16)
+ weights = [
+ (0.3 * torch.randn(n, generator=g, device=DEV)).to(torch.bfloat16) for n in (D, D, ID, ID)
+ ]
+ kv_dtype = torch.float8_e4m3fn if spec.get("fp8") else torch.bfloat16
+ pool = torch.randn(pages, 2, hk, TPB, D, generator=g, device=DEV).to(kv_dtype)
+ index_cache = torch.randn(pages, TPB, 1, ID, generator=g, device=DEV).to(torch.bfloat16)
+ position_cache = torch.randint(
+ 0, table_rows, (pages, TPB, 3), generator=g, device=DEV, dtype=torch.int32
+ )
+ return dict(
+ proj=proj, out=out[:, :qkv_width], index_q=index_scratch[:, :IH], position_ids=ids,
+ coords=coords, cos_sin=cos_sin, q_weight=weights[0], k_weight=weights[1],
+ index_q_weight=weights[2], index_k_weight=weights[3], k_cache=pool[:, 0],
+ v_cache=pool[:, 1], index_cache=index_cache, position_cache=position_cache,
+ request_indices=request, logical_positions=logical, block_table=block_table,
+ tokens_per_block=TPB, compress_ratio=CR, num_heads_q=hq, num_heads_k=hk, head_dim=D,
+ index_heads=IH, index_head_dim=ID, index_q_col=iq_col, rotary_dim=ROT, eps=1e-6,
+ mrope_section1=11, mrope_section2=10, rows_per_request=spec.get("rows_per_request", 1),
+ ) # fmt: skip
+
+
+def _bytes(storage, device) -> torch.Tensor:
+ return torch.empty(0, dtype=torch.uint8, device=device).set_(storage)
+
+
+def _clone(args):
+ """The arguments over fresh copies of every allocation they view, each view rebuilt at
+ its offset and strides; returns them and the copies, keyed by the first argument
+ viewing each allocation."""
+ out, copies, keys = dict(args), {}, {}
+ for name, t in args.items():
+ if not isinstance(t, torch.Tensor):
+ continue
+ key = keys.setdefault(t.untyped_storage().data_ptr(), name)
+ if key not in copies:
+ copies[key] = _bytes(t.untyped_storage(), t.device).clone()
+ view = torch.empty(0, dtype=t.dtype, device=t.device)
+ view.set_(copies[key].untyped_storage(), t.storage_offset(), t.shape, t.stride())
+ out[name] = view
+ return out, copies
+
+
+def _written_slots(args):
+ """(row, page, slot) of every row whose page lookup holds, from the host's view."""
+ table = args["block_table"].cpu()
+ pages = args["k_cache"].shape[0]
+ found = []
+ for t, (req, p) in enumerate(
+ zip(args["request_indices"].tolist(), args["logical_positions"].tolist())
+ ):
+ column = p // TPB
+ if 0 <= req < table.shape[0] and column < table.shape[1]:
+ page = int(table[req, column])
+ if 0 <= page < pages:
+ found.append((t, page, p % TPB))
+ return found
+
+
+def _check(got, want, where="") -> dict:
+ """Judge the rotated values within the band, then neutralize them in `got` and require
+ every byte of every allocation bitwise. Returns the measured worst errors."""
+ g, w = got[0], want[0]
+ hq, hk = g["num_heads_q"], g["num_heads_k"]
+ width = (hq + hk) * D
+ measured = {}
+
+ def banded(name, a, b):
+ a, b = a.float().reshape(-1, a.shape[-1]), b.float().reshape(-1, b.shape[-1])
+ scale = b.abs().amax(dim=1, keepdim=True).clamp_min(1e-9)
+ measured[name] = ((a - b).abs() / scale).max().item() / 2**-8
+ assert_within_ulp(a.to(torch.bfloat16), b.to(torch.bfloat16), ELEMENT_ULP, RMS_ULP)
+
+ banded("packed q/k", g["out"][:, :width], w["out"][:, :width])
+ banded("index query", g["index_q"], w["index_q"])
+ slots = _written_slots(g)
+ k_got = g["out"][:, hq * D : width].reshape(-1, hk, D)
+ for t, page, slot in slots:
+ own = k_got[t].to(g["k_cache"].dtype)
+ assert torch.equal(g["k_cache"][page, :, slot].view(torch.uint8), own.view(torch.uint8)), (
+ f"{where}row {t}: the K pool's row is not the packed k converted"
+ )
+ g["k_cache"][page, :, slot] = w["k_cache"][page, :, slot]
+ ends = [(t, page, slot) for t, page, slot in slots if (slot + 1) % CR == 0]
+ if ends:
+ a = torch.stack([g["index_cache"][page, slot, 0] for _, page, slot in ends])
+ b = torch.stack([w["index_cache"][page, slot, 0] for _, page, slot in ends])
+ banded("compressed keys", a, b)
+ for _, page, slot in ends:
+ g["index_cache"][page, slot, 0] = w["index_cache"][page, slot, 0]
+ g["out"][:, :width] = w["out"][:, :width]
+ g["index_q"].copy_(w["index_q"])
+ assert got[1].keys() == want[1].keys()
+ for key, copy in got[1].items():
+ differ = (copy != want[1][key]).nonzero()
+ assert differ.numel() == 0, f"{where}{key}: {differ.shape[0]} bytes differ"
+ return measured
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
+ certifying(op)
+ args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31))
+ want = _clone(args)
+ op.reference(**want[0])
+ got = _clone(args)
+ op(**got[0])
+ torch.cuda.synchronize()
+ measured = _check(got, want)
+ print(f"{cell.why[:48]}: worst element / row scale in 2^-8 units {measured}")
+
+
+@pytest.mark.parametrize(
+ "spec",
+ [
+ dict(geom="tp4", pos=(8699, 8700)),
+ dict(geom="tp4", pos=(84000,), rows_per_request=4),
+ dict(geom="dep4", pos=(8697, 8700), rows_per_request=4, fp8=True),
+ ],
+ ids=["tp4_decode", "tp4_verify", "dep4_fp8_verify"],
+)
+def test_reads_after_the_wait_under_graph_replay(spec) -> None:
+ """The projection is rewritten inside the graph right before the op; with PDL a read of
+ it before the wait would see the previous replay's row (in a verify step, also the raw
+ index keys of a group's earlier rows)."""
+ args = _build(spec, seed=23)
+ src = args["proj"].clone()
+ work, work_copies = _clone(args)
+
+ def step():
+ work["proj"].copy_(src * 1.0)
+ op(**work)
+
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ step()
+ torch.cuda.current_stream().wait_stream(side)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ step()
+ # The replays' projections from a seeded generator: the global RNG's state depends on the
+ # tests run before this one in the process.
+ g = torch.Generator(device=DEV).manual_seed(29)
+ for trial in range(3):
+ src.copy_((torch.randn(src.shape, generator=g, device=DEV) * 2.0).to(src.dtype))
+ before = _clone(work)
+ graph.replay()
+ torch.cuda.synchronize()
+ want = _clone(before[0])
+ want[0]["proj"].copy_(src)
+ op.reference(**want[0])
+ _check((work, work_copies), want, where=f"replay {trial}: ")
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_vision_rope_2d_packed_qk.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_vision_rope_2d_packed_qk.py
new file mode 100644
index 000000000000..e4e3949fd5b4
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_vision_rope_2d_packed_qk.py
@@ -0,0 +1,174 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the vision_rope_2d_packed_qk variant: the vision tower's 2-D rotary,
+one Triton launch in place on the q|k prefix of the packed qkv projection.
+
+Each cell builds a packed `[T, 3 * 16, 72]` bf16 qkv projection and the caller's angle
+table from the cell's grid (spatial-merge positions, fp64 angles rounded once to fp32,
+fp32 cos / sin), rotates `qkv[:, :32]` in place, and gates:
+ 1. bitwise against the entry's `reference` on the same table;
+ 2. bitwise against the whole-op reference of `vision_rope_2d__r52` (copied below:
+ positions by nested loops, fp64 angles, `polar`, complex multiply), q and k
+ separately, with that entry's distribution gates beside it;
+ 3. the v heads of every row untouched.
+"""
+
+import zlib
+from typing import List, Sequence, Tuple
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.attention.vision_rope_2d_packed_qk import ( # noqa: E501
+ HEAD_DIM,
+ NUM_HEADS,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.attention.vision_rope_2d_packed_qk import ( # noqa: E501
+ vision_rope_2d_packed_qk as op,
+)
+
+__extra_import_path__ = ["../../../.."]
+from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
+
+assert torch.cuda.is_available(), "vision_rope_2d_packed_qk requires a CUDA device"
+
+MERGE = 2
+THETA = 10000.0
+HALF = HEAD_DIM // 2
+
+
+def _num_tokens(grid) -> int:
+ return sum(t * h * w for t, h, w in grid)
+
+
+def _spatial_merge_positions(grid: Sequence[Sequence[int]], merge: int) -> List[Tuple[int, int]]:
+ """Per-token `(row, col)` by four nested loops, in spatial-merge order."""
+ positions: List[Tuple[int, int]] = []
+ for frames, height, width in grid:
+ frame = [
+ (block_row * merge + r, block_col * merge + c)
+ for block_row in range(height // merge)
+ for block_col in range(width // merge)
+ for r in range(merge)
+ for c in range(merge)
+ ]
+ positions.extend(frame * frames)
+ return positions
+
+
+def _angle_table(grid_thw: torch.Tensor, merge: int, theta: float, head_dim: int):
+ """The caller's glue: per-token fp32 `cos` / `sin` `[T, head_dim // 2]`, height angles
+ then width, from fp64 angles rounded once to fp32."""
+ rope_dim = head_dim // 2
+ inv_freq = 1.0 / (theta ** (torch.arange(0, rope_dim, 2, dtype=torch.float64) / rope_dim))
+ pos = torch.tensor(_spatial_merge_positions(grid_thw.tolist(), merge), dtype=torch.float64)
+ angles = (pos[:, :, None] * inv_freq).flatten(1).to(device="cuda", dtype=torch.float32)
+ return angles.cos().contiguous(), angles.sin().contiguous()
+
+
+def _r52_reference(q, k, grid_thw, spatial_merge_size, theta=THETA):
+ """`vision_rope_2d__r52.reference`, verbatim: the whole op from the grid, as a fresh
+ `(q_out, k_out)`."""
+ rope_dim = q.shape[-1] // 2
+ axis_dim = rope_dim // 2
+ positions = _spatial_merge_positions(grid_thw.tolist(), spatial_merge_size)
+ inverse = torch.tensor(
+ [theta ** (-(2.0 * j) / rope_dim) for j in range(axis_dim)], dtype=torch.float64
+ )
+ pos = torch.tensor(positions, dtype=torch.float64).reshape(-1, 2)
+ angles = (pos[:, :, None] * inverse).flatten(1).to(device=q.device, dtype=torch.float32)
+ rotor = torch.polar(torch.ones_like(angles), angles).unsqueeze(-2)
+
+ def rotate(x: torch.Tensor) -> torch.Tensor:
+ fp32 = x.float()
+ rotated = torch.complex(fp32[..., :rope_dim], fp32[..., rope_dim:]) * rotor
+ return torch.cat((rotated.real, rotated.imag), dim=-1).to(x.dtype)
+
+ return rotate(q), rotate(k)
+
+
+def _packed_qkv(tokens: int, dtype: torch.dtype, seed: int) -> torch.Tensor:
+ g = torch.Generator(device="cuda").manual_seed(seed)
+ return torch.randn(tokens, 3 * NUM_HEADS, HEAD_DIM, generator=g, device="cuda").to(dtype)
+
+
+def _qk(qkv: torch.Tensor) -> torch.Tensor:
+ qk = qkv[:, : 2 * NUM_HEADS]
+ assert qk.stride() == (3 * NUM_HEADS * HEAD_DIM, HEAD_DIM, 1) and not qk.is_contiguous()
+ return qk
+
+
+def _distribution(got: torch.Tensor, want: torch.Tensor) -> Tuple[float, float, float]:
+ """(cosine, relative L2, p99.9 absolute error) of `got` against `want`, flattened."""
+ g, w = got.float().reshape(-1), want.float().reshape(-1)
+ cos = torch.nn.functional.cosine_similarity(g, w, dim=0).item()
+ rel_l2 = ((g - w).norm() / w.norm().clamp_min(1e-30)).item()
+ err = (g - w).abs()
+ # torch.quantile caps its input size; a strided sample of 2^24 is plenty for p99.9.
+ if err.numel() > 2**24:
+ err = err[:: -(-err.numel() // 2**24)]
+ return cos, rel_l2, torch.quantile(err, 0.999).item()
+
+
+#: `vision_rope_2d__r52`'s distribution gates: both sides round the same fp32 rotation to bf16
+#: once, so they differ by at most one bf16 ulp per element, at most 2^-7 of its value: the
+#: relative L2 is bounded by that, and 1 - cosine by half its square.
+_REL_L2_BOUND = 2.0**-7
+_COS_BOUND = 1.0 - 2.0**-15
+
+
+def _mismatch(got: torch.Tensor, want: torch.Tensor) -> Tuple[int, float]:
+ """(elements differing in their bits, largest absolute difference)."""
+ differing = (got.reshape(-1).view(torch.int16) != want.reshape(-1).view(torch.int16)).sum()
+ return int(differing.item()), (got.float() - want.float()).abs().max().item()
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name
+ certifying(op)
+ spec = cell.spec
+ grid_thw = torch.tensor([list(row) for row in spec["grid"]], dtype=spec["grid_dtype"])
+ tokens = _num_tokens(spec["grid"])
+ qkv = _packed_qkv(tokens, spec["dtype"], seed=zlib.crc32(cell.why.encode()) % 2**31)
+ before = qkv.clone()
+ cos, sin = _angle_table(grid_thw, MERGE, THETA, HEAD_DIM)
+ qk = _qk(qkv)
+
+ want = op.reference(qk, cos, sin)
+ assert torch.equal(qkv, before), "reference must not mutate its input"
+ out = op(qk, cos, sin)
+ torch.cuda.synchronize()
+
+ assert out is qk, "the op returns the view it rotated"
+ assert torch.equal(qkv[:, 2 * NUM_HEADS :], before[:, 2 * NUM_HEADS :]), "v was written"
+ assert not torch.equal(qk, before[:, : 2 * NUM_HEADS]), "q|k was not rotated"
+
+ # 1. The entry's own reference on the same table: bitwise.
+ differing, max_dev = _mismatch(qk, want)
+ print(f"[{cell.why}] tokens={tokens} vs reference: {differing} differing, max {max_dev:.3e}")
+ assert differing == 0, f"{differing} elements differ from reference (max {max_dev:.3e})"
+
+ # 2. vision_rope_2d__r52's whole-op reference from the grid: bitwise, plus its
+ # distribution gates.
+ q_r52, k_r52 = _r52_reference(
+ before[:, :NUM_HEADS], before[:, NUM_HEADS : 2 * NUM_HEADS], grid_thw, MERGE, THETA
+ )
+ for name, got, ref in (
+ ("q", qkv[:, :NUM_HEADS], q_r52),
+ ("k", qkv[:, NUM_HEADS : 2 * NUM_HEADS], k_r52),
+ ):
+ differing, max_dev = _mismatch(got, ref)
+ cos_sim, rel_l2, p999 = _distribution(got, ref)
+ print(
+ f"[{cell.why}] {name} vs r52: {differing} differing, max {max_dev:.3e}, "
+ f"cos={cos_sim:.9f} rel_l2={rel_l2:.3e} p99.9={p999:.3e}"
+ )
+ assert cos_sim >= _COS_BOUND and rel_l2 <= _REL_L2_BOUND, (name, cos_sim, rel_l2)
+ assert differing == 0, f"{name}: {differing} elements differ from r52 (max {max_dev:.3e})"
+
+
+def _small():
+ grid_thw = torch.tensor([(1, 8, 12)], dtype=torch.int64)
+ qkv = _packed_qkv(96, torch.bfloat16, seed=17)
+ cos, sin = _angle_table(grid_thw, MERGE, THETA, HEAD_DIM)
+ return qkv, cos, sin
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first__decode.py
similarity index 70%
rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first.py
rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first__decode.py
index bf2d62d45c13..4550b39ed729 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first__decode.py
@@ -1,14 +1,15 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the qwen4_exp_hc_decode_tc_first entry: cells in the band against
-the float64 reference, the row-count guard, graph capture."""
+"""GPU test for this target's qwen4_exp_hc_decode_tc_first variant
+(ops/hc/qwen4_exp_hc_decode_tc_first__decode): cells in the band against the float64
+reference at the decode rows the members run."""
import zlib
import pytest
import torch
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_tc_first import ( # noqa: E501
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.hc.qwen4_exp_hc_decode_tc_first__decode import ( # noqa: E501
qwen4_exp_hc_decode_tc_first as op,
)
@@ -66,28 +67,3 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by
low, inject = op(*args)
assert low.shape == (cell.spec["rows"], RANK)
_compare((low, inject), op.reference(*args))
-
-
-def test_guard_refuses_more_than_two_tiles_of_rows() -> None:
- normed, first, rank, c, wi = _build(dict(rows=257), seed=4)
- with validating(op), pytest.raises(AssertionError):
- op(normed, first, rank, c, wi)
-
-
-@pytest.mark.parametrize("rows", [8, 256])
-def test_cuda_graph_capture_and_replay(rows) -> None:
- normed, first, rank, c, wi = _build(dict(rows=rows), seed=5)
- side = torch.cuda.Stream()
- side.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(side):
- op(normed, first, rank, c, wi)
- torch.cuda.current_stream().wait_stream(side)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- g_low, g_inject = op(normed, first, rank, c, wi)
- fresh = _build(dict(rows=rows), seed=6)[0]
- normed.copy_(fresh)
- graph.replay()
- torch.cuda.synchronize()
- low, inject = op(fresh, first, rank, c, wi)
- assert torch.equal(g_low, low) and torch.equal(g_inject, inject)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4__moe.py
similarity index 51%
rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py
rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4__moe.py
index eecc9d836bef..b25c216ac330 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4__moe.py
@@ -1,9 +1,9 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the qwen4_exp_hc_up_gate_mean_fp4 entry: `mixed` in the band
-`qwen4_exp_hc_up_gate_mean`'s test uses, the NVFP4 outputs against the
-native-torch quantization of the returned `mixed` outside the reciprocal tie
-window, the global-scale guard, graph capture."""
+"""GPU test for the qwen4 default target's qwen4_exp_hc_up_gate_mean_fp4__moe
+entry: `mixed` in the band `qwen4_exp_hc_up_gate_mean`'s test uses and bitwise
+that entry's own output, the NVFP4 outputs against the native-torch
+quantization of the returned `mixed` outside the reciprocal tie window."""
import zlib
@@ -14,10 +14,10 @@
Fp4Quantize,
fp4_quantize,
)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501
- gate_ratio,
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_up_gate_mean import (
+ qwen4_exp_hc_up_gate_mean,
)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_up_gate_mean_fp4 import ( # noqa: E501
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.hc.qwen4_exp_hc_up_gate_mean_fp4__moe import ( # noqa: E501
qwen4_exp_hc_up_gate_mean_fp4 as op,
)
@@ -29,16 +29,14 @@
DEV = "cuda"
-def nvfp4_reference(x, gs):
- """`fp4_quantize`'s native-torch NVFP4 of `x` `[T, K]` with linear scales:
- `(packed [T, K/2] uint8, sf [T*K/16] uint8)`."""
- return fp4_quantize.reference(x, gs, 16, False, False)
-
-
-def nvfp4_near_tie(x, gs):
- """`[T, K]` bool: elements within `fp4_quantize`'s tie window of an e2m1 midpoint,
- where the kernels' approximate reciprocals may round either way."""
- return Fp4Quantize.near_tie(fp4_quantize.scaled(x, gs, 16))
+def gate_ratio(out: torch.Tensor, ref: torch.Tensor) -> float:
+ """Worst element's distance from `ref` as a fraction of the band `rtol = eps`,
+ `atol = eps * max|ref|` (both of `ref`'s dtype), the band
+ `qwen4_exp_hc_up_gate_mean`'s test certifies `mixed` in. Below 1 is inside."""
+ eps = torch.finfo(ref.dtype).eps
+ ref64 = ref.double()
+ atol = eps * ref64.abs().max().item()
+ return ((out.double() - ref64).abs() / (atol + eps * ref64.abs())).max().item()
def _unpack(xq: torch.Tensor) -> torch.Tensor:
@@ -49,25 +47,16 @@ def _check_quant(mixed, xq, xsf, gs) -> None:
"""NVFP4 outputs against the native-torch quantization of the `mixed` this
call returned: scale bytes equal, e2m1 codes equal except inside the
reciprocal tie window."""
- want_q, want_sf = nvfp4_reference(mixed, gs)
+ want_q, want_sf = fp4_quantize.reference(mixed, gs, 16, False, False)
assert xq.dtype == torch.uint8 and xq.shape == want_q.shape and xq.is_contiguous()
assert xsf.dtype == torch.uint8 and xsf.shape == want_sf.shape
assert torch.equal(xsf, want_sf), f"{(xsf != want_sf).sum().item()} scale bytes differ"
diff = _unpack(xq) != _unpack(want_q)
- stray = diff & ~nvfp4_near_tie(mixed, gs)
+ near = Fp4Quantize.near_tie(fp4_quantize.scaled(mixed, gs, 16))
+ stray = diff & ~near
assert not stray.any(), f"{stray.sum().item()} e2m1 codes differ outside the tie window"
-def _compare(out, ref, gs) -> None:
- """`mixed`: gate_ratio <= 1, the band `qwen4_exp_hc_up_gate_mean` is
- certified in. Quantization: exact on the returned `mixed`."""
- (mixed, xq, xsf), (ref_mixed, _, _) = out, ref
- assert mixed.dtype == ref_mixed.dtype and mixed.shape == ref_mixed.shape
- assert mixed.is_contiguous()
- assert gate_ratio(mixed, ref_mixed) <= 1.0, f"{gate_ratio(mixed, ref_mixed):.4g}x the band"
- _check_quant(mixed, xq, xsf, gs)
-
-
def _build(spec, seed):
hc, hidden, rank = spec.get("hc_count", 4), spec.get("hidden", 2560), spec.get("rank", 320)
g = torch.Generator(device=DEV).manual_seed(seed)
@@ -92,34 +81,11 @@ def _seed(cell) -> int:
def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
certifying(op)
args = _build(cell.spec, seed=_seed(cell))
- out = op(*args)
- assert out[0].shape == (cell.spec["rows"], cell.spec.get("hidden", 2560))
- _compare(out, op.reference(*args), args[4])
-
-
-def test_guards() -> None:
- """The kernel loads `global_scale[0]` only; a per-row scale is refused."""
- low, up, normed, c, gs = _build(dict(rows=4), seed=4)
- with validating(op), pytest.raises(AssertionError):
- op(low, up, normed, c, torch.cat([gs, gs]))
-
-
-def test_cuda_graph_capture_and_replay() -> None:
- low, up, normed, c, gs = _build(dict(rows=8), seed=5)
- side = torch.cuda.Stream()
- side.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(side):
- op(low, up, normed, c, gs)
- torch.cuda.current_stream().wait_stream(side)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- g_out = op(low, up, normed, c, gs)
- fresh = _build(dict(rows=8), seed=6)
- low.copy_(fresh[0])
- normed.copy_(fresh[2])
- gs.mul_(0.5)
- graph.replay()
- torch.cuda.synchronize()
- want = op(fresh[0], up, fresh[2], c, gs)
- for got, exp in zip(g_out, want):
- assert torch.equal(got, exp)
+ mixed, xq, xsf = op(*args)
+ assert mixed.shape == (cell.spec["rows"], cell.spec.get("hidden", 2560))
+ ref_mixed = op.reference(*args)[0]
+ assert mixed.dtype == ref_mixed.dtype and mixed.is_contiguous()
+ assert gate_ratio(mixed, ref_mixed) <= 1.0, f"{gate_ratio(mixed, ref_mixed):.4g}x the band"
+ # The same kernel and epilogue arithmetic as the bf16-only build the target runs elsewhere.
+ assert torch.equal(mixed, qwen4_exp_hc_up_gate_mean(*args[:4]))
+ _check_quant(mixed, xq, xsf, args[4])
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ssm/test_modeling_v2_gdn_cached_replay_cute__verify.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ssm/test_modeling_v2_gdn_cached_replay_cute__verify.py
new file mode 100644
index 000000000000..fabebf6190ed
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ssm/test_modeling_v2_gdn_cached_replay_cute__verify.py
@@ -0,0 +1,276 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the CuTe DSL cached-replay entry (`gdn_cached_replay_cute__verify`) at a
+tp4/tep4 rank's 4 key / 12 value heads.
+
+Each cell drives consecutive verify steps the way the engine does, through the Triton
+cached-replay entry's test harness (`_Bench`: the stock manager's bf16 state page at the tep4 engine's slot
+count, work items, the all-layer commit, `_advance_replay_state` with random acceptance), and
+holds every call to this entry's `reference` on that call's own inputs (output and all four
+state buffers, the Triton entry's bands) and every output to the exact twin (the fp32 token-by-token
+recurrence over every accepted token). With `strided_gates` the raw a / b reach the kernel as
+column views of an input-projection-shaped buffer, as the target hands them over.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+
+import tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.ssm.gdn_cached_replay_cute__verify as cr # noqa: E501
+
+_MV2_TESTS = Path(__file__).resolve().parents[4]
+if str(_MV2_TESTS) not in sys.path:
+ sys.path.insert(0, str(_MV2_TESTS))
+from _validating import certifying # noqa: E402,F401 — fixture by name
+
+_spec = importlib.util.spec_from_file_location(
+ "cached_replay_harness",
+ # The Triton entry's test, in the family tree (tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm).
+ _MV2_TESTS
+ / "models/qwen4_exp/ssm"
+ / "test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py",
+)
+trep = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(trep)
+
+assert torch.cuda.is_available(), "gdn_cached_replay_cute requires a CUDA device"
+
+op = cr.gdn_cached_replay_cute
+DEV = "cuda"
+K, V = cr.K, cr.V
+#: The input projection's column split at a tp4/tep4 rank: conv (q|k|v), z, a, b.
+_CONV, _Z = 2 * cr.RANK_K_HEADS * K + cr.RANK_V_HEADS * V, cr.RANK_V_HEADS * V
+
+
+def _gates(bench, act, strided):
+ """The raw a / b `[N, T, HV]`: the harness's contiguous tensors, or the same values as
+ column views of a `[N*T, conv + z + 2 HV]` projection (row stride = its width)."""
+ if not strided:
+ return act["a"], act["b"]
+ n, t, hv = bench.n, bench.t, bench.hv
+ proj = torch.randn(n * t, _CONV + _Z + 2 * hv, device=DEV, dtype=torch.bfloat16)
+ a = proj[:, _CONV + _Z : _CONV + _Z + hv]
+ b = proj[:, _CONV + _Z + hv :]
+ a.copy_(act["a"].reshape(n * t, hv))
+ b.copy_(act["b"].reshape(n * t, hv))
+ a, b = a.reshape(n, t, hv), b.reshape(n, t, hv)
+ assert a.stride() == (t * proj.shape[1], proj.shape[1], 1), a.stride()
+ return a, b
+
+
+def _reference(bench, kw):
+ """This entry's `reference` on a call's arguments, every per-slot buffer through the
+ harness's `compact` (the touched slots' rows of an engine page)."""
+ if bench.idx is None:
+ return op.reference(**kw)
+ per_slot = (
+ "ssm_states",
+ "old_u",
+ "old_k",
+ "old_G",
+ "old_beta",
+ "cache_buf_idx",
+ "prev_num_accepted_tokens",
+ )
+ compact = {name: bench.compact(kw[name]) for name in per_slot}
+ return op.reference(
+ **dict(kw, **compact, state_indices=bench.compact_indices(kw["state_indices"]))
+ )
+
+
+def _drive(cell, bench) -> None:
+ strided = cell.spec.get("strided_gates", False)
+ wraps = appends = 0
+ for _ in range(bench.steps):
+ mask, wrapped, _ = bench.writes()
+ wraps += wrapped.numel()
+ appends += bench.live.numel() - wrapped.numel()
+ bench.prepare()
+ after_commit = []
+ for layer in range(bench.layers):
+ act = bench.inputs()
+ kw = bench.kwargs(layer, act)
+ kw["g"], kw["beta"] = _gates(bench, act, strided)
+ _, u, k, G = bench.state(layer)
+ before = dict(u=u.clone(), k=k.clone(), G=G.clone(), dA=bench.old_dA[layer].clone())
+ ref = _reference(bench, kw)
+ exact = op.reference(**bench.exact_kwargs(layer, kw))
+ committed = None
+ if bench.path == "all_layer":
+ committed = _reference(
+ bench,
+ dict(kw, replay_work_items=None, n_writes=None, use_all_layer_commit=False),
+ )
+ got = op(**kw)
+ trep._compare((got, *bench.observed(layer)), ref)
+ _, u, k, G = bench.state(layer)
+ assert torch.equal(bench.old_dA[layer], before["dA"]), "old_beta is never written"
+ assert torch.equal(u[~mask], before["u"][~mask]), "old_u written outside the append"
+ assert torch.equal(k[~mask], before["k"][~mask]), "old_k written outside the append"
+ Gt, Gt0 = G.transpose(-1, -2), before["G"].transpose(-1, -2)
+ assert torch.equal(Gt[~mask], Gt0[~mask]), "old_G written outside the append"
+ bench.adopt_exact(layer, exact)
+ after_commit.append((got, exact, committed))
+ bench.commit()
+ for layer, (got, exact, committed) in enumerate(after_commit):
+ if committed is not None:
+ # The partitioned call plus the manager's commit is the call that commits
+ # in-kernel.
+ trep._compare((got, *bench.observed(layer)), committed)
+ trep._compare((got, None, None, None, None), (exact[0], None, None, None, None))
+ if bench.page is not None:
+ bench.page.assert_rest_unchanged()
+ bench.accept(bench.draw_acceptance())
+ assert wraps > 0 and appends > 0, f"the cell never wrapped or appended ({wraps}, {appends})"
+
+
+def _drive_padding(cell) -> None:
+ """Trailing rows sharing one slot, as CUDA-graph padding repeats the manager's dummy
+ slot: the rows on unique slots, and every slot but the shared one, are held to the
+ reference run on the unique rows alone (a dense pool: the shared slot's rows are
+ indexed out)."""
+ spec = {k_: v_ for k_, v_ in cell.spec.items() if k_ not in ("page_layers",)}
+ spec["slots"] = 24
+ bench = trep._Bench(spec, seed=trep._seed(cell.why), exact=False)
+ dup = spec["dup"]
+ unique = bench.n - len(dup)
+ shared = int(bench.state_indices[unique])
+ others = torch.ones(bench.slots, dtype=torch.bool, device=DEV)
+ others[shared] = False
+ for _ in range(bench.steps):
+ for layer in range(bench.layers):
+ act = bench.inputs()
+ kw = bench.kwargs(layer, act)
+ kw["g"], kw["beta"] = _gates(bench, act, spec.get("strided_gates", False))
+ alone = dict(
+ kw,
+ q=kw["q"][:unique],
+ k=kw["k"][:unique],
+ v=kw["v"][:unique],
+ g=kw["g"][:unique],
+ beta=kw["beta"][:unique],
+ state_indices=kw["state_indices"][:unique],
+ packed_qkv=act["packed"][: unique * bench.t],
+ )
+ ref = op.reference(**alone)
+ got = op(**kw)
+ pool, u, k, G = bench.state(layer)
+ trep._compare(
+ (got[:unique], pool[others], u[others], k[others], G[others]),
+ (ref[0], *(x[others] for x in ref[1:])),
+ )
+ bench.accept(bench.draw_acceptance())
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS])
+def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
+ certifying(op)
+ trep._MEASURED.clear()
+ if "dup" in cell.spec:
+ _drive_padding(cell)
+ else:
+ bench = trep._Bench(cell.spec, seed=trep._seed(cell.why))
+ try:
+ _drive(cell, bench)
+ finally:
+ bench.close()
+ print(
+ f"\n[{cell.why}] "
+ + ", ".join(f"{k_} {v_:.3f}" for k_, v_ in sorted(trep._MEASURED.items()))
+ )
+
+
+def _rank_spec(**spec):
+ return dict(state_dtype=torch.bfloat16, H=cr.RANK_K_HEADS, HV=cr.RANK_V_HEADS, **spec)
+
+
+@pytest.mark.parametrize("pdl", [False, True])
+@pytest.mark.parametrize("n,path", [(1, "small"), (4, "small"), (16, "all_layer")])
+def test_cuda_graph_replay_on_new_data_is_bitwise_eager(n, path, pdl) -> None:
+ """Two layers captured once (strided a / b, the all-layer commit after them when
+ partitioned), replayed three times on new activations with acceptance in between,
+ bitwise against the same steps run eagerly on a twin state; with `pdl` the replay
+ launches with programmatic dependent launch, as the target's verify step does."""
+ t = 4
+ spec = _rank_spec(path=path, n=n, t=t, layers=2, slots=40)
+ seed = 9100 + n
+ eager, graph, warm = (trep._Bench(spec, seed, exact=False) for _ in range(3))
+ acts = [warm.inputs() for _ in range(2)]
+ gates = [_gates(warm, act, True) for act in acts]
+
+ def step(b):
+ b.prepare()
+ outs = []
+ for layer in range(2):
+ kw = b.kwargs(layer, acts[layer])
+ kw["g"], kw["beta"] = gates[layer]
+ kw["launch_with_pdl"] = pdl
+ outs.append(op(**kw))
+ b.commit()
+ return outs
+
+ step(warm)
+ torch.cuda.synchronize()
+ captured = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(captured):
+ outs_graph = step(graph)
+ refill = torch.Generator(device=DEV).manual_seed(seed + 1)
+ for _ in range(3):
+ for act, (a, b) in zip(acts, gates):
+ for x in (act["packed"], a, b):
+ x.copy_(torch.randn(x.shape, generator=refill, device=DEV, dtype=x.dtype))
+ outs_eager = step(eager)
+ captured.replay()
+ torch.cuda.synchronize()
+ for e, g in zip(outs_eager, outs_graph):
+ assert torch.equal(e, g), "replayed output differs from eager"
+ for name in ("backing", "old_x", "old_B", "old_dt", "pnat", "buf"):
+ assert torch.equal(getattr(eager, name), getattr(graph, name)), f"{name} differs"
+ accepted = eager.draw_acceptance()
+ eager.accept(accepted)
+ graph.accept(accepted)
+
+
+#: Launch mappings the 4 / 12 geometry can take, each checked against the tp1 default
+#: small-batch mapping: every one computes the same instructions per state row.
+_MAPPINGS = (
+ dict(tpc=1, warps=1, pf=0, minb=4, early_acc=True),
+ dict(tpc=2, warps=2, pf=0, minb=4, early_acc=True),
+ dict(tpc=4, warps=4, pf=0, minb=4, early_acc=True),
+ dict(tpc=8, warps=4, pf=1, minb=4),
+ dict(tpc=24, warps=8, pf=1, minb=2),
+)
+
+
+@pytest.mark.parametrize("pdl", [False, True])
+@pytest.mark.parametrize("n,t", [(1, 4), (4, 4), (2, 2), (1, 8)])
+def test_launch_mappings_agree(n, t, pdl) -> None:
+ """One step bitwise equal across launch mappings (output, checkpoint with in-kernel
+ commits, u / k / G history), with wrapping and appending rows and, from 4 rows, a
+ negative slot."""
+ spec = _rank_spec(path="small", n=n, t=t, steps=1, layers=1, slots=24)
+ bench = trep._Bench(spec, 4242 + n * 10 + t, exact=False)
+ sidx = bench.state_indices.clone()
+ if n > 2:
+ sidx[n // 2] = -1
+ act = bench.inputs()
+ a, b = _gates(bench, act, True)
+ snap = [x.clone() for x in bench.state(0)]
+ results = []
+ for cfg in _MAPPINGS:
+ for x, y in zip(bench.state(0), snap):
+ x.copy_(y)
+ kw = bench.kwargs(0, act)
+ sentinel = torch.full((n, t, bench.hv, V), 3.0, dtype=torch.bfloat16, device=DEV)
+ kw.update(g=a, beta=b, state_indices=sidx, output=sentinel, config=cfg, launch_with_pdl=pdl)
+ out = op(**kw)
+ torch.cuda.synchronize()
+ results.append([out.clone()] + [x.clone() for x in bench.state(0)])
+ for cfg, got in zip(_MAPPINGS[1:], results[1:]):
+ for name, x, y in zip(("out", "state", "u", "k", "G"), got, results[0]):
+ assert torch.equal(x, y), f"{name}: {cfg} vs {_MAPPINGS[0]}"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_attention_plumbing.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_attention_plumbing.py
new file mode 100644
index 000000000000..2daa23f1e54d
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_attention_plumbing.py
@@ -0,0 +1,386 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The target's attention plumbing around its kernels.
+
+* **The FP8 KV range probe** (`TRTLLM_MODELING_V2_VALIDATE` over an fp8 pool). The pool stores
+ K/V at the unit scale, so a magnitude past the e4m3 maximum saturates. The probe reports the
+ largest magnitude of the first eager forward's attention layers, once, and never reads a
+ tensor under CUDA graph capture.
+* **The context pass over an NVFP4 pool, a group of requests at a time** (one rank). The
+ context rows read a bf16 scratch of their requests' pages (`_nvfp4_context_plan`), filled and
+ attended one group of consecutive requests at a time, so it holds the largest group's pages
+ rather than the step's. Under a budget of a few pages, with one request longer than the
+ budget, a prefill step's rows match one group over the whole step bitwise: its requests'
+ rows come in multiples of the row-tile pass's 10-row groups, so a launch per group splits
+ none of them and every row takes the same pass as in the one launch.
+* **MTP IndexShare.** The draft selection (`_selection`) in its fused and its two-launch form,
+ the capture of each request's draft-extend row into the shared buffers (decode and
+ mixed-batch forms), and the reuse entry's order over canonical captures. The targets are
+ built without weights: the capture methods read only their arguments, the metadata and the
+ core's draft binding.
+"""
+
+import math
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_mtp_reuse_topk import ( # noqa: E501
+ qsa_mtp_reuse_topk as reuse,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501
+ modeling as M,
+)
+
+_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="runs on a CUDA device")
+DEV = "cuda"
+
+
+# ---- the FP8 KV range probe ---------------------------------------------------------------
+
+LAST = 11
+
+
+def _kv(peak: float = 1.0):
+ """K and V of one layer, 5 rows x 2 heads x 8 dims, with `peak` as V's largest value."""
+ k = torch.linspace(-1.0, 1.0, 80, device=DEV).reshape(5, 2, 8).to(torch.bfloat16)
+ v = torch.linspace(-0.5, 0.5, 80, device=DEV).reshape(5, 2, 8).to(torch.bfloat16)
+ v[3, 1, 4] = peak
+ return k, v
+
+
+def _run_forward(probe, peaks):
+ """One forward's attention layers 3, 7 and LAST; returns each `observe` result."""
+ return [probe.observe(layer, *_kv(peak)) for layer, peak in zip((3, 7, LAST), peaks)]
+
+
+@_cuda
+def test_a_forward_within_the_e4m3_range_is_silent(monkeypatch):
+ logger = Mock()
+ monkeypatch.setattr(M, "logger", logger)
+
+ assert _run_forward(M._KvRangeProbe(LAST), (1.0, 300.0, 448.0)) == [True, True, False]
+
+ logger.warning.assert_not_called()
+
+
+@_cuda
+@pytest.mark.parametrize("peak", [-600.0, float("nan")])
+def test_a_magnitude_past_the_e4m3_maximum_is_reported_once(monkeypatch, peak):
+ """A negative value counts by its magnitude; a NaN counts as out of range."""
+ logger = Mock()
+ monkeypatch.setattr(M, "logger", logger)
+
+ # The oversized value is in a middle layer: the report waits for the forward's last one.
+ assert _run_forward(M._KvRangeProbe(LAST), (1.0, peak, 1.0)) == [True, True, False]
+
+ logger.warning.assert_called_once()
+ assert "e4m3 maximum 448" in logger.warning.call_args.args[0]
+
+
+@pytest.mark.cpu_only
+def test_the_probe_runs_only_when_asked_over_an_fp8_pool():
+ core = SimpleNamespace(quant_mode=M.QM_FP8_KV_CACHE, attn_layers=[3, 7, LAST])
+ assert M._KvRangeProbe.for_target(core, True)._last == LAST
+ assert M._KvRangeProbe.for_target(core, False) is None
+ for quant_mode in (0, M.QM_NVFP4_KV_CACHE):
+ other = SimpleNamespace(quant_mode=quant_mode, attn_layers=core.attn_layers)
+ assert M._KvRangeProbe.for_target(other, True) is None
+
+
+@_cuda
+def test_observing_under_capture_reads_nothing(monkeypatch):
+ logger = Mock()
+ monkeypatch.setattr(M, "logger", logger)
+ probe = M._KvRangeProbe(LAST)
+ k, v = _kv(1000.0)
+ scratch = torch.zeros(1, device=DEV)
+ stream = torch.cuda.Stream()
+ stream.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(stream):
+ scratch.add_(1) # warm the allocator off the capture stream
+ torch.cuda.current_stream().wait_stream(stream)
+ graph = torch.cuda.CUDAGraph()
+ kept = []
+ with torch.cuda.graph(graph):
+ scratch.add_(1)
+ # The last layer, with an out-of-range value: a read here would fail the capture
+ # (a host sync) and a report would fire.
+ kept.append(probe.observe(LAST, k, v))
+
+ assert kept == [True], "a captured forward does not end the probe"
+ logger.warning.assert_not_called()
+ # The first eager forward after capture is what reports.
+ assert _run_forward(probe, (1.0, 1.0, 1000.0)) == [True, True, False]
+ logger.warning.assert_called_once()
+
+
+# ---- the context pass over an NVFP4 pool ---------------------------------------------------
+
+TPB, HKV, D, QH, RATIO, TOPK = 32, 2, 256, 24, 4, 2048
+# (cached, new) tokens a context request: 10, 30, 1, 47 and 4 pages.
+SEGMENTS = ((0, 300), (900, 40), (20, 10), (1400, 100), (64, 60))
+BUDGET_PAGES = 32 # request 3's 47 pages form a group of their own
+
+
+def _context_step(seed: int = 5):
+ gen = torch.Generator(device=DEV).manual_seed(seed)
+ kv_lens = [s + n for s, n in SEGMENTS]
+ pages = [math.ceil(k / TPB) for k in kv_lens]
+ pool_pages = sum(pages) + 3
+ shape = (pool_pages, 2, HKV, TPB)
+ data = torch.randint(0, 256, (*shape, D // 2), generator=gen, device=DEV).to(torch.uint8)
+ scales = torch.randint(0x08, 0x31, (*shape, D // 16), generator=gen, device=DEV)
+ scales = scales.to(torch.uint8)
+ slots = 2 * len(SEGMENTS)
+ slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))
+ slot_of = slot_of[: len(SEGMENTS)].to(DEV, torch.int32)
+ page_ids = torch.randperm(pool_pages, generator=gen, device=DEV).to(torch.int32)
+ block_table = torch.full((slots, max(pages)), -1, dtype=torch.int32, device=DEV)
+ cursor = 0
+ for slot, count in zip(slot_of.tolist(), pages):
+ block_table[slot, :count] = page_ids[cursor : cursor + count]
+ cursor += count
+ new = torch.tensor([n for _, n in SEGMENTS], dtype=torch.int32)
+ md = SimpleNamespace(
+ num_contexts=len(SEGMENTS),
+ kv_lens_runtime=torch.tensor(kv_lens, dtype=torch.int32),
+ kv_lens_cuda_runtime=torch.tensor(kv_lens, dtype=torch.int32, device=DEV),
+ seq_lens=new,
+ seq_lens_cuda=new.to(DEV),
+ )
+ logical = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in SEGMENTS])
+ rows = logical.numel()
+ # Every visible token, as the expansion lays out a selection under its top-k.
+ column = torch.arange(TOPK + RATIO - 1, device=DEV)[None, :]
+ selected = torch.where(column <= logical[:, None], column, -1).to(torch.int32)
+ core = SimpleNamespace(
+ heads_kv=HKV,
+ head_dim=D,
+ q_width=QH * D,
+ kv_scale_quant_orig=torch.ones(3, dtype=torch.float32, device=DEV),
+ softmax_scale=D**-0.5,
+ compress_ratio=RATIO,
+ )
+ kv = SimpleNamespace(k=data[:, 0], v=data[:, 1], k_scale=scales[:, 0], v_scale=scales[:, 1])
+ return dict(
+ core=core,
+ kv=kv,
+ md=md,
+ block_table=block_table,
+ req_idx=slot_of.repeat_interleave(new.to(DEV).long()),
+ q=torch.randn(rows, QH, D, generator=gen, device=DEV).to(torch.bfloat16),
+ selected=selected,
+ logical=logical.to(torch.int32),
+ rows=rows,
+ )
+
+
+def _context_pass(step):
+ plan = M._nvfp4_context_plan(
+ step["core"], step["md"], step["req_idx"], step["block_table"], TPB
+ )
+ out = torch.empty(step["rows"], QH * D, dtype=torch.bfloat16, device=DEV)
+ M._sparse_gqa_nvfp4_context(
+ step["core"],
+ step["kv"],
+ plan,
+ step["q"],
+ step["block_table"],
+ step["selected"],
+ step["req_idx"],
+ step["logical"],
+ TPB,
+ step["rows"],
+ out,
+ 1,
+ )
+ return plan, out
+
+
+@_cuda
+def test_grouped_context_pass_equals_one_group(monkeypatch) -> None:
+ step = _context_step()
+ whole, want = _context_pass(step)
+ assert len(whole[0]) == 1 and whole[2].shape[0] == 92
+
+ page_bytes = 2 * HKV * TPB * D * 2
+ monkeypatch.setattr(M, "_NVFP4_CONTEXT_SCRATCH_BYTES", BUDGET_PAGES * page_bytes)
+ (groups, _, k_scratch, _), got = _context_pass(step)
+ # Requests [0], [1, 2], [3] (over the budget, alone), [4].
+ assert [(g[0], g[1]) for g in groups] == [(0, 300), (300, 350), (350, 450), (450, 510)]
+ assert k_scratch.shape[0] == 47
+ assert torch.equal(got.view(torch.int16), want.view(torch.int16))
+
+
+# ---- MTP IndexShare -----------------------------------------------------------------------
+
+K = 512
+SLOT = 1
+COLS = 2080
+
+
+def _target(cls):
+ target = object.__new__(cls)
+ target.core = SimpleNamespace(_draft_bound=SimpleNamespace(share_slot=SLOT))
+ return target
+
+
+def _metadata(batch: int, accepted, seq_lens=None, num_contexts: int = 0):
+ return SimpleNamespace(
+ num_seqs=batch,
+ num_contexts=num_contexts,
+ seq_lens_cuda=None
+ if seq_lens is None
+ else torch.tensor(seq_lens, dtype=torch.int32, device=DEV),
+ qsa_mtp_num_accepted=torch.tensor(accepted, dtype=torch.int32, device=DEV),
+ qsa_shared_topk_indices=torch.full((3, 8, K), -7, dtype=torch.int32, device=DEV),
+ qsa_shared_topk_visible_blocks=torch.full((3, 8), -7, dtype=torch.int32, device=DEV),
+ qsa_shared_topk_captured_slots=set(),
+ )
+
+
+def _scores_and_windows(rows: int, gen: torch.Generator):
+ """Relu-shaped scores (half zeros: ties at the budget's last score whenever fewer than
+ K are positive) over windows of 0 to 3K groups; the columns past a window are NaN."""
+ ends = torch.randint(0, COLS + 1, (rows,), generator=gen).to(torch.int32).to(DEV)
+ scores = torch.relu(torch.randn(rows, COLS, generator=gen)).to(DEV)
+ scores = scores * (torch.rand(rows, COLS, generator=gen).to(DEV) < 0.3)
+ outside = torch.arange(COLS, device=DEV)[None, :] >= ends[:, None]
+ scores.masked_fill_(outside, float("nan"))
+ return scores, torch.zeros(rows, dtype=torch.int32, device=DEV), ends
+
+
+def _select(scores, starts, ends) -> torch.Tensor:
+ out = torch.empty((scores.shape[0], K), dtype=torch.int32, device=DEV)
+ M._select(SimpleNamespace(block_topk=K), scores, starts, ends, out)
+ return out
+
+
+@_cuda
+@pytest.mark.parametrize("width", [2080, 65600])
+@pytest.mark.parametrize("fused", [False, True])
+def test_selection_matches_the_canonical_composition(width, fused, monkeypatch):
+ """The draft selection keeps the canonical tie order and expansion in the fused launch
+ and in the selection and the expansion apart."""
+ positions = torch.tensor([2043, 2047, 84000, 99999], dtype=torch.int64, device=DEV)
+ visible = ((positions + 1) // 4).to(torch.int32)
+ visible.clamp_(max=width)
+ positions = torch.minimum(positions, (visible.to(torch.int64) * 4 + 2))
+ lengths = (positions + 3).to(torch.int32)
+ starts = torch.zeros_like(visible)
+ generator = torch.Generator(device=DEV).manual_seed(85)
+ scores = torch.relu(torch.randn((4, width), generator=generator, device=DEV) - 1.8)
+ cols = torch.arange(width, device=DEV)
+ scores.masked_fill_(cols[None, :] >= visible[:, None], float("nan"))
+ expected = torch.empty((4, K), dtype=torch.int32, device=DEV)
+ core = SimpleNamespace(block_topk=K, compress_ratio=4, token_topk=2048)
+ M._select(core, scores, starts, visible, expected)
+ expected_tokens = M.qsa_expand_block_indices(
+ expected, positions, lengths, compress_ratio=4, token_topk=2048
+ )
+ if not fused:
+ monkeypatch.setattr(M, "_topk_expand_supports", lambda scores: False)
+ result = torch.full_like(expected, -7)
+ selected = torch.full_like(expected_tokens, -7)
+ M._selection(core, scores, visible, positions, lengths, result, starts, selected)
+ assert torch.equal(result, expected)
+ assert torch.equal(selected, expected_tokens)
+
+
+@_cuda
+def test_decode_capture_at_draft_seven_takes_row_s_times_8_plus_a_minus_1():
+ """max_draft_len 7: eight rows per request, accepted 1..8, so the captured row is
+ `s * 8 + a - 1` -- the first and the last row of a block included."""
+ p, batch = 8, 4
+ accepted = [1, 8, 5, 3]
+ topk = torch.arange(batch * p * K, dtype=torch.int32, device=DEV).reshape(batch * p, K)
+ visible = torch.arange(300, 300 + batch * p, dtype=torch.int32, device=DEV)
+ md = _metadata(batch, accepted)
+
+ _target(M.DecodeTarget)._capture_selection(md, topk, visible, batch, p)
+
+ rows = [s * p + a - 1 for s, a in enumerate(accepted)]
+ assert rows == [0, 15, 20, 26]
+ torch.testing.assert_close(md.qsa_shared_topk_indices[SLOT, :batch], topk[rows])
+ torch.testing.assert_close(md.qsa_shared_topk_visible_blocks[SLOT, :batch], visible[rows])
+ # Other slots and rows past the batch are untouched.
+ assert torch.all(md.qsa_shared_topk_indices[SLOT, batch:] == -7)
+ assert torch.all(md.qsa_shared_topk_indices[0] == -7)
+ assert md.qsa_shared_topk_captured_slots == {SLOT}
+
+
+@_cuda
+@pytest.mark.parametrize("compact", [False, True])
+def test_prefill_capture_at_draft_seven(compact):
+ """A mixed batch at max_draft_len 7: two context chunks of 5 and 3 rows, then three
+ generation requests of 8 rows with 8, 1 and 6 accepted. A context request captures its
+ chunk's last row whatever accepted count the worker left there. A `compact` selection
+ (the draft-extend step's kept rows under attention TP) holds each context request's last
+ row, then every generation row."""
+ seq_lens = [5, 3, 8, 8, 8]
+ accepted = [2, 3, 8, 1, 6]
+ rows = [4, 7, 8 + 7, 16 + 0, 24 + 5]
+ if compact:
+ seq_lens_rows = [1, 1, 8, 8, 8]
+ rows = [0, 1, 2 + 7, 10 + 0, 18 + 5]
+ else:
+ seq_lens_rows = seq_lens
+ rows_total = sum(seq_lens_rows)
+ topk = torch.arange(rows_total * K, dtype=torch.int32, device=DEV).reshape(rows_total, K)
+ visible = torch.arange(400, 400 + rows_total, dtype=torch.int32, device=DEV)
+ md = _metadata(5, accepted, seq_lens=seq_lens, num_contexts=2)
+ md.num_ctx_tokens = 8
+
+ _target(M.PrefillTarget)._capture_selection(md, topk, visible, compact=compact)
+
+ torch.testing.assert_close(md.qsa_shared_topk_indices[SLOT, :5], topk[rows])
+ torch.testing.assert_close(md.qsa_shared_topk_visible_blocks[SLOT, :5], visible[rows])
+ assert md.qsa_shared_topk_captured_slots == {SLOT}
+
+
+def _reuse_case(fresh: int):
+ """A canonical capture of 16 requests and the reuse call's arguments for it."""
+ gen = torch.Generator().manual_seed(11 + fresh)
+ batch = 16
+ scores, starts, ends = _scores_and_windows(batch, gen)
+ captured_visible = ends.cpu()
+ captured = _select(scores, starts, ends)
+ shared = torch.full((8, 32, K), -1, dtype=torch.int32, device=DEV)
+ shared[SLOT, :batch] = captured
+ shared_visible = torch.zeros((8, 32), dtype=torch.int32, device=DEV)
+ shared_visible[SLOT, :batch] = captured_visible.to(DEV)
+ args = dict(
+ shared_indices=shared[SLOT],
+ shared_visible_blocks=shared_visible[SLOT],
+ request_indices=torch.arange(batch, dtype=torch.int32, device=DEV),
+ visible_blocks=(captured_visible + fresh).to(DEV),
+ output=torch.empty((batch, K), dtype=torch.int32, device=DEV),
+ )
+ return args, captured_visible
+
+
+@_cuda
+def test_reuse_of_a_canonical_capture_stays_ascending():
+ """A canonical capture plus the groups completed since is still ascending: appended ids
+ are newer than every captured one, and a full row gives up its oldest ids to them, so
+ the newest captured groups -- the ones next to the new ones -- are kept."""
+ for fresh in (0, 2):
+ args, captured_visible = _reuse_case(fresh)
+
+ out = reuse(**args)
+
+ assert torch.equal(out, reuse.reference(**args))
+ for row, shared_row, vis in zip(
+ out.cpu().tolist(), args["shared_indices"].cpu().tolist(), captured_visible.tolist()
+ ):
+ ids = [i for i in row if i >= 0]
+ assert row == ids + [-1] * (K - len(ids)), "padding must stay last"
+ assert ids == sorted(ids) and len(set(ids)) == len(ids)
+ assert len(ids) == min(vis + fresh, K)
+ assert all(i < vis + fresh for i in ids)
+ captured = [i for i in shared_row if i >= 0]
+ kept = captured[max(0, len(captured) + fresh - K) :]
+ assert ids == kept + list(range(vis, vis + fresh))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_core_layout.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_core_layout.py
new file mode 100644
index 000000000000..57b3e5f537e5
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_core_layout.py
@@ -0,0 +1,769 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The default target's core at tp4, tep4, dep4 and dep2: each rank's forward widths and its
+experts.
+
+The core is built on the meta device from the checkpoint's text and vision config values at
+every rank of a tp4, tep4, dep4 or dep2 mapping, so nothing here touches a device. What is pinned is
+what the forward addresses: the fused attention row and the split the norm + rope reads it
+through, the Gated-DeltaNet projection's column split, the routed experts a rank holds (an
+expert window, or tp4's padded intermediate slice of every expert) and the per-expert scalars
+cut to them, and the attention-DP round trip's padding and slicing. Under attention DP the
+draft layer holds every attention head and runs its FP8 expert window over the rows the group
+gathers, padded by the row counts the spec worker hands each draft step.
+"""
+
+from __future__ import annotations
+
+import ast
+import inspect
+from types import SimpleNamespace
+from unittest import mock
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501
+ modeling as M,
+)
+from tensorrt_llm._torch.configs import Qwen4ExpTextConfig, Qwen4ExpVisionConfig
+from tensorrt_llm._torch.model_config import ModelConfig
+from tensorrt_llm.llmapi import QSASparseAttentionConfig
+from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig
+from tensorrt_llm.mapping import Mapping
+from tensorrt_llm.models.modeling_utils import QuantConfig
+from tensorrt_llm.quantization.mode import QuantAlgo
+
+pytestmark = pytest.mark.cpu_only
+
+NUM_LAYERS = 48
+WORLD = 4
+
+# Qwen3.8-Flash-Next's text config.
+_TEXT = dict(
+ architectures=["Qwen4ExpForCausalLM"],
+ attention_bias=False,
+ attention_dropout=0.0,
+ bos_token_id=248044,
+ dtype="bfloat16",
+ eos_token_id=248044,
+ full_attention_interval=4,
+ hc_count=4,
+ hc_lowrank=320,
+ head_dim=256,
+ heads_per_ngram=8,
+ hidden_act="silu",
+ hidden_size=2560,
+ indexer_budget=2048,
+ indexer_compress_ratio=4,
+ indexer_head_dim=128,
+ indexer_kv_heads=1,
+ indexer_n_heads=4,
+ layer_types=["full_attention" if i % 4 == 3 else "linear_attention" for i in range(NUM_LAYERS)],
+ linear_conv_kernel_dim=4,
+ linear_key_head_dim=128,
+ linear_num_key_heads=16,
+ linear_num_value_heads=48,
+ linear_value_head_dim=128,
+ make_ngram_vocab_size_divisible_by=128,
+ mamba_ssm_dtype="float32",
+ max_position_embeddings=262144,
+ moe_intermediate_size=640,
+ mtp=dict(
+ hybrid=True,
+ layer_types=["full_attention"],
+ mtp_use_hidden_state_from_layer=None,
+ num_hidden_layers=1,
+ rope_theta=10000000,
+ ),
+ mtp_num_hidden_layers=1,
+ mtp_use_dedicated_embeddings=False,
+ ngram_size=3,
+ ngram_vocab_size_base=20000000,
+ norm_topk_prob=True,
+ num_attention_heads=24,
+ num_experts=512,
+ num_experts_per_tok=10,
+ num_hidden_layers=NUM_LAYERS,
+ num_key_value_heads=2,
+ number_of_conv_states=3,
+ output_gate_type="sigmoid",
+ partial_rotary_factor=0.25,
+ ple_conv_kernel_size=4,
+ ple_embed_dim=2560,
+ ple_layer_ids=[2],
+ rms_norm_eps=1e-06,
+ rope_parameters=dict(
+ mrope_interleaved=True,
+ mrope_section=[11, 11, 10],
+ partial_rotary_factor=0.25,
+ rope_theta=10000000,
+ rope_type="default",
+ ),
+ shared_expert_intermediate_size=640,
+ split_ngram_parts=128,
+ tie_word_embeddings=False,
+ vocab_size=248320,
+)
+_VISION = dict(
+ deepstack_visual_indexes=[],
+ depth=27,
+ dtype="bfloat16",
+ hidden_act="gelu_pytorch_tanh",
+ hidden_size=1152,
+ in_channels=3,
+ intermediate_size=4304,
+ num_heads=16,
+ num_position_embeddings=2304,
+ out_hidden_size=2560,
+ patch_size=16,
+ spatial_merge_size=2,
+ temporal_patch_size=2,
+)
+
+# (attention DP, MoE TP, MoE EP) of the three members.
+_TEP4 = dict(enable_attention_dp=False, moe_tp_size=1, moe_ep_size=4)
+_DEP4 = dict(enable_attention_dp=True, moe_tp_size=1, moe_ep_size=4)
+_DEP2 = dict(enable_attention_dp=True, moe_tp_size=1, moe_ep_size=2)
+_TP4 = dict(enable_attention_dp=False, moe_tp_size=4, moe_ep_size=1)
+
+
+def _core(rank: int, layout: dict, spec_config=None, world: int = WORLD):
+ """The core of rank `rank` of a `world`-rank engine at `layout`, every parameter on meta."""
+ model_config = ModelConfig(
+ pretrained_config=Qwen4ExpTextConfig(**_TEXT),
+ mapping=Mapping(world_size=world, tp_size=world, rank=rank, **layout),
+ quant_config=QuantConfig(quant_algo=QuantAlgo.MIXED_PRECISION),
+ moe_backend="TRTLLM",
+ moe_backend_requested="TRTLLM",
+ sparse_attention_config=QSASparseAttentionConfig(),
+ spec_config=spec_config,
+ )
+ composite = SimpleNamespace(
+ vision_config=Qwen4ExpVisionConfig(**_VISION), image_token_id=248056, video_token_id=248057
+ )
+ # The core asserts the device it runs on; the weights it declares stay on meta.
+ with (
+ mock.patch.object(torch.cuda, "get_device_capability", return_value=M._SM),
+ torch.device("meta"),
+ ):
+ return M.Qwen4ExpDefaultModelingV2(model_config, composite)
+
+
+@pytest.mark.parametrize("rank", range(WORLD))
+def test_a_tep4_rank_runs_its_share_of_the_heads(rank) -> None:
+ """Six q heads and one of the two KV heads (repeated on rank pairs) over a 4352-wide
+ fused row: `[q 1536 | k 256 | v 256 | gate 1536 | index q 512, k 128 | v slot 128]`.
+ The main norm + rope reads it as 6 q, 1 k and 10 "v" heads of 256 (the row stride
+ alone), which is the cell `fused_qk_norm_rope` certifies."""
+ core = _core(rank, _TEP4)
+ assert (core.heads_q, core.heads_kv, core.kv_head_start) == (6, 1, rank * 2 // WORLD)
+ assert (core.q_width, core.kv_width, core.qkv_width) == (1536, 256, 2048)
+ assert (core.index_width, core.index_scratch_width) == (640, 768)
+ assert core.attn_in_width == 4352
+ assert core.attn_in_tail_heads == 10
+ assert core.index_col == 3584
+ assert tuple(core.w["l3_attn_in_w"].shape) == (core.attn_in_width, 2560)
+ assert tuple(core.w["l3_o_w"].shape) == (2560, core.q_width)
+
+
+@pytest.mark.parametrize("rank", range(WORLD))
+def test_a_tep4_rank_runs_its_share_of_the_recurrent_heads(rank) -> None:
+ """Four key and twelve value heads: the conv runs over `[q | k | v]` = 512 + 512 + 1536
+ channels, and the one input projection's output splits `[q|k|v 2560 | z 1536 | a 12 |
+ b 12]`."""
+ core = _core(rank, _TEP4)
+ assert (core.num_k_heads, core.num_v_heads) == (4, 12)
+ assert (core.key_dim, core.value_dim, core.conv_dim) == (512, 1536, 2560)
+ assert tuple(core.in_proj_split) == (2560, 1536, 12, 12)
+ assert tuple(core.w["l0_in_proj_w"].shape) == (sum(core.in_proj_split), 2560)
+ assert tuple(core.w["l0_out_w"].shape) == (2560, core.value_dim)
+ assert tuple(core.w["l0_conv_w"].shape) == (core.conv_dim, 1, 4)
+
+
+@pytest.mark.parametrize("rank", range(WORLD))
+def test_a_tep4_rank_holds_one_128_expert_window(rank) -> None:
+ """The routed experts' window is `[128r, 128r + 128)`; everything the router, the shared
+ expert and the hyper-connections read is whole on every rank."""
+ core = _core(rank, _TEP4)
+ assert (core.local_experts, core.expert_offset) == (128, 128 * rank)
+ assert core.w["l0_fc1_w"].shape[0] == 128
+ assert core.w["l0_fc2_w"].shape[0] == 128
+ assert tuple(core.w["l0_router"].shape) == (512, 2560)
+ assert tuple(core.w["embed"].shape) == (248320, 2560)
+
+
+@pytest.mark.parametrize("rank", range(WORLD))
+def test_a_tp4_rank_holds_every_expert_at_its_padded_intermediate_slice(rank) -> None:
+ """tp4 cuts attention and Gated-DeltaNet as tep4 does, and holds all 512 experts at a
+ 160-row intermediate slice stored 256 wide: the runner is handed the whole routing space
+ (`local_expert_offset` 0, `local_num_experts` 512) at `intermediate_size` 256, the
+ stack `fp4_block_scale_moe_runner` certifies, and the scalars span every expert."""
+ core = _core(rank, _TP4)
+ assert (core.heads_q, core.heads_kv, core.kv_head_start) == (6, 1, rank * 2 // WORLD)
+ assert (core.num_k_heads, core.num_v_heads, core.attn_in_width) == (4, 12, 4352)
+ assert (core.dp_size, core.attn_tp_size) == (1, WORLD)
+ assert (core.local_experts, core.expert_offset, core.expert_inter) == (512, 0, 256)
+ assert core.moe_inter == 640
+ assert tuple(core.w["l0_fc1_w"].shape) == (512, 2 * 256, 2560 // 2)
+ assert tuple(core.w["l0_fc1_s"].shape) == (512, 2 * 256, 2560 // 16)
+ assert tuple(core.w["l0_fc2_w"].shape) == (512, 2560, 256 // 2)
+ assert tuple(core.w["l0_fc2_s"].shape) == (512, 2560, 256 // 16)
+ assert tuple(core.w["l0_e_isc2"].shape) == (512,)
+
+
+@pytest.mark.parametrize("layout", [_TEP4, _DEP4], ids=["tep4", "dep4"])
+def test_an_expert_window_runs_at_the_whole_intermediate(layout) -> None:
+ core = _core(1, layout)
+ assert (core.local_experts, core.expert_offset, core.expert_inter) == (128, 128, 640)
+
+
+def test_the_expert_scalars_are_cut_to_the_window() -> None:
+ """The runner takes one scalar per expert of its window: each is the whole routing
+ space's product, sliced, and the FC1 activation scale is the layer's one number."""
+ num_experts, start, count = 512, 256, 128
+ gen = torch.Generator().manual_seed(0)
+ isc1 = torch.full((num_experts,), 0.5)
+ ws2_1 = torch.rand(num_experts, generator=gen) + 0.5
+ isc2 = torch.rand(num_experts, generator=gen) + 0.5
+ ws2_2 = torch.rand(num_experts, generator=gen) + 0.5
+ w = {
+ "l0_e_isc1": isc1,
+ "l0_e_isc1_up": isc1.clone(),
+ "l0_e_ws2_1": ws2_1,
+ "l0_e_ws2_1_up": ws2_1.clone(),
+ "l0_e_isc2": isc2,
+ "l0_e_ws2_2": ws2_2,
+ }
+ s = M._moe_layer_scalars(w, 0, start, count)
+ window = slice(start, start + count)
+ gate1 = isc1 * ws2_1
+ torch.testing.assert_close(s["global_scale"], torch.tensor([2.0]), rtol=0, atol=0)
+ torch.testing.assert_close(s["output1_scale_gate_scalar"], gate1[window], rtol=0, atol=0)
+ torch.testing.assert_close(s["output1_scale_scalar"], (gate1 / isc2)[window], rtol=0, atol=0)
+ torch.testing.assert_close(s["output2_scale_scalar"], (isc2 * ws2_2)[window], rtol=0, atol=0)
+ assert all(t.shape == (count,) for k, t in s.items() if k != "global_scale")
+
+
+def test_a_spread_fc1_input_scale_is_refused() -> None:
+ """One activation quantization feeds every expert, so the layer must have one
+ `gate_proj.input_scale`."""
+ isc1 = torch.full((512,), 0.5)
+ isc1[7] = 0.25
+ ones = torch.ones(512)
+ w = {
+ "l0_e_isc1": isc1,
+ "l0_e_isc1_up": isc1.clone(),
+ "l0_e_ws2_1": ones,
+ "l0_e_ws2_1_up": ones.clone(),
+ "l0_e_isc2": ones,
+ "l0_e_ws2_2": ones,
+ }
+ with pytest.raises(AssertionError, match="one gate_proj.input_scale"):
+ M._moe_layer_scalars(w, 0, 0, 128)
+
+
+@pytest.mark.parametrize("rank", range(WORLD))
+def test_a_dep4_rank_runs_every_head(rank) -> None:
+ """Attention DP keeps the single-rank geometry: 24 q and both KV heads over the 14080-wide
+ fused row (55 heads of 256: 24 q, 2 k, 29 "v"), and all 16 / 48 GDN heads over 10240 conv
+ channels. Nothing is summed after the mixers."""
+ core = _core(rank, _DEP4)
+ assert (core.heads_q, core.heads_kv, core.kv_head_start) == (24, 2, 0)
+ assert (core.q_width, core.kv_width, core.qkv_width) == (6144, 512, 7168)
+ assert core.attn_in_width == 14080
+ assert core.attn_in_tail_heads == 29
+ assert core.index_col == 13312
+ assert tuple(core.w["l3_attn_in_w"].shape) == (14080, 2560)
+ assert tuple(core.w["l3_o_w"].shape) == (2560, 6144)
+ assert (core.num_k_heads, core.num_v_heads, core.conv_dim) == (16, 48, 10240)
+ assert tuple(core.in_proj_split) == (10240, 6144, 48, 48)
+ assert core.attn_tp_size == 1
+
+
+@pytest.mark.parametrize("rank", range(WORLD))
+def test_a_dep4_rank_runs_one_128_expert_window_for_the_group(rank) -> None:
+ """The MoE group is the TP group, which spans the DP ranks: rank `r` is DP rank `r` and
+ holds the window `[128r, 128r + 128)`, over the rows the whole group gathers."""
+ core = _core(rank, _DEP4)
+ assert (core.dp_size, core.dp_rank, core.dp_group) == (WORLD, rank, list(range(WORLD)))
+ assert (core.local_experts, core.expert_offset) == (128, 128 * rank)
+ assert core.w["l0_fc1_w"].shape[0] == 128
+ assert _core(rank, _TEP4).dp_size == 1
+
+
+@pytest.mark.parametrize("rank", range(2))
+def test_a_dep2_rank_runs_every_head_and_one_256_expert_window(rank) -> None:
+ """dep2, the disaggregated-serving engines' layout, is dep4 at half the group: the same
+ single-rank attention and Gated-DeltaNet geometry, and rank `r` holds the 256-expert
+ window `[256r, 256r + 256)` over the rows both ranks gather."""
+ core, dep4 = _core(rank, _DEP2, world=2), _core(rank, _DEP4)
+ for name in (
+ "heads_q",
+ "heads_kv",
+ "kv_head_start",
+ "qkv_width",
+ "attn_in_width",
+ "attn_in_tail_heads",
+ "index_col",
+ "num_k_heads",
+ "num_v_heads",
+ "conv_dim",
+ "attn_tp_size",
+ ):
+ assert getattr(core, name) == getattr(dep4, name), name
+ assert tuple(core.in_proj_split) == tuple(dep4.in_proj_split)
+ for key in ("l3_attn_in_w", "l3_o_w", "l0_in_proj_w", "l0_router", "l0_sh_gu"):
+ assert core.w[key].shape == dep4.w[key].shape, key
+ assert (core.dp_size, core.dp_rank, core.dp_group) == (2, rank, [0, 1])
+ assert (core.local_experts, core.expert_offset) == (256, 256 * rank)
+ assert core.w["l0_fc1_w"].shape[0] == 256
+ assert core.w["l0_fc1_w"].shape[1:] == dep4.w["l0_fc1_w"].shape[1:]
+
+
+def test_attention_dp_pads_to_the_groups_largest_row_count() -> None:
+ """Each rank pads to the largest of `all_rank_num_tokens`; a count that disagrees with
+ the rows this rank holds would gather the wrong split, and is refused."""
+ core = SimpleNamespace(dp_size=WORLD, dp_rank=2)
+ md = SimpleNamespace(all_rank_num_tokens=[5, 17, 9, 1])
+ assert M.Qwen4ExpDefaultModelingV2._dp_pad_rows(core, md, 9) == 8
+ with pytest.raises(AssertionError, match="wrong split"):
+ M.Qwen4ExpDefaultModelingV2._dp_pad_rows(core, md, 10)
+ # Without attention DP the metadata is not read.
+ assert M.Qwen4ExpDefaultModelingV2._dp_pad_rows(SimpleNamespace(dp_size=1), None, 3) == 0
+
+
+@pytest.mark.parametrize("rank", range(WORLD))
+def test_the_attention_dp_round_trip_returns_this_ranks_rows_through_every_window(
+ rank, monkeypatch
+) -> None:
+ """The round trip pads this rank's rows, gathers the group's in rank order, runs this
+ rank's window over all of them, reduce-scatters and drops the padding: with each window
+ a function of the gathered rows, this rank gets back the sum over windows of its own rows.
+ The collectives are stood in for by the group's arithmetic."""
+ counts = [3, 7, 1, 4]
+ hidden = 8
+ gen = torch.Generator().manual_seed(rank)
+ xs = [torch.randn(n, hidden, generator=gen) for n in counts]
+ rows = max(counts)
+ gathered = torch.cat([nn_pad(x, rows) for x in xs], 0)
+
+ def window(w):
+ return lambda t: t * (w + 1) + w
+
+ full = sum(window(w)(gathered) for w in range(WORLD))
+ seen = {}
+
+ def fake_allgather(t, sizes, group, process_group):
+ assert sizes is None and group == list(range(WORLD)) and process_group is None
+ assert tuple(t.shape) == (rows, hidden)
+ torch.testing.assert_close(t, nn_pad(xs[rank], rows), rtol=0, atol=0)
+ return gathered
+
+ def fake_reducescatter(t, sizes, group, process_group):
+ assert sizes is None and group == list(range(WORLD)) and process_group is None
+ seen["partial"] = t
+ return full[rank * rows : (rank + 1) * rows]
+
+ monkeypatch.setattr(M, "allgather", fake_allgather)
+ monkeypatch.setattr(M, "reducescatter", fake_reducescatter)
+ monkeypatch.setattr(M, "mpi_disabled", lambda: False)
+ core = SimpleNamespace(dp_size=WORLD, dp_group=list(range(WORLD)), mapping=None)
+ expert_sum = M.Qwen4ExpDefaultModelingV2._expert_sum_op(core)
+ out = expert_sum(window(rank), xs[rank], rows - counts[rank])
+ torch.testing.assert_close(seen["partial"], window(rank)(gathered), rtol=0, atol=0)
+ torch.testing.assert_close(out, sum(window(w)(xs[rank]) for w in range(WORLD)))
+ assert tuple(out.shape) == (counts[rank], hidden)
+
+
+def nn_pad(x: torch.Tensor, rows: int) -> torch.Tensor:
+ """`x` with zero rows appended up to `rows`."""
+ return torch.cat([x, x.new_zeros(rows - x.shape[0], x.shape[1])], 0)
+
+
+@pytest.mark.parametrize("max_draft_len", [1, 3, 7])
+def test_mtp_widens_the_verify_step_and_keys_the_draft_layer_after_the_trunk(
+ max_draft_len,
+) -> None:
+ """A verify step carries `max_draft_len + 1` rows a generation request; the draft layer is
+ global layer 48, the index every KV manager maps it at and the weight table keys it at.
+ Without speculation the MTP module is not declared."""
+ core = _core(0, _TEP4, MTPDecodingConfig(max_draft_len=max_draft_len))
+ assert core.mtp_enabled and core.max_verify_rows == max_draft_len + 1
+ assert core.mtp_layer_idx == NUM_LAYERS
+ assert core.mtp_theta == core.theta == 1e7
+ assert core.spec_rollback is not None
+ assert any(k.startswith(f"l{NUM_LAYERS}_mtp_") for k in core.w)
+ plain = _core(0, _TEP4)
+ assert not plain.mtp_enabled and plain.max_verify_rows == 1 and plain.spec_rollback is None
+ assert not any("_mtp_" in k for k in plain.w)
+
+
+@pytest.mark.parametrize("rank", range(WORLD))
+def test_a_tep4_rank_holds_its_share_of_the_draft_layer(rank) -> None:
+ """The draft layer's attention is cut like a trunk attention layer's: six q heads (their
+ q and interleaved gate rows), one KV head and the o_proj columns of the six heads; its
+ QSA indexer, front end, router, shared expert and mixers are whole. Its FP8 experts are
+ the rank's window `[128r, 128r + 128)` over the TP group, FP8 128x128 block scales."""
+ core = _core(rank, _TEP4, MTPDecodingConfig(max_draft_len=3))
+ p = f"l{NUM_LAYERS}_mtp_"
+ w = core.w
+ assert tuple(w[f"{p}q_w"].shape) == (2 * 1536, 2560)
+ assert tuple(w[f"{p}k_w"].shape) == tuple(w[f"{p}v_w"].shape) == (256, 2560)
+ assert tuple(w[f"{p}o_w"].shape) == (2560, 1536)
+ assert tuple(w[f"{p}idx_w"].shape) == (640, 2560)
+ assert tuple(w[f"{p}fc_e"].shape) == tuple(w[f"{p}fc_h"].shape) == (2560, 2560)
+ assert tuple(w[f"{p}pre_h_norm"].shape) == (4 * 2560,)
+ assert tuple(w[f"{p}router"].shape) == (512, 2560)
+ assert tuple(w[f"{p}sh_gu"].shape) == (2 * 640 + 8, 2560)
+ assert (core.mtp_local_experts, core.mtp_expert_offset) == (128, 128 * rank)
+ assert tuple(w[f"{p}fc1_w"].shape) == (128, 2 * 640, 2560)
+ assert w[f"{p}fc1_w"].dtype == torch.float8_e4m3fn
+ assert tuple(w[f"{p}fc1_s"].shape) == (128, 2 * 640 // 128, 2560 // 128)
+ assert tuple(w[f"{p}fc2_w"].shape) == (128, 2560, 640)
+ assert tuple(w[f"{p}fc2_s"].shape) == (128, 2560 // 128, 640 // 128)
+
+
+@pytest.mark.parametrize("rank", range(WORLD))
+def test_a_tp4_rank_holds_whole_expert_draft_windows_beside_its_intermediate_slice(rank) -> None:
+ """At tp4 the trunk's routed experts are every expert at a 160-row intermediate slice
+ stored 256 wide, but the draft layer's FP8 experts cannot be cut that way (their 128x128
+ weight blocks would split mid-block). They are the rank's whole-expert window
+ `[128r, 128r + 128)` over the TP group, exactly tep4's, and the four windows partition
+ the routing space. The draft attention is cut as at tep4."""
+ max_draft_len = 3
+ core = _core(rank, _TP4, MTPDecodingConfig(max_draft_len=max_draft_len))
+ assert core.mtp_enabled and core.max_verify_rows == max_draft_len + 1
+ assert (core.local_experts, core.expert_offset, core.expert_inter) == (512, 0, 256)
+ assert (core.mtp_local_experts, core.mtp_expert_offset, core.moe_inter) == (
+ 128,
+ 128 * rank,
+ 640,
+ )
+ p = f"l{NUM_LAYERS}_mtp_"
+ w = core.w
+ assert tuple(w[f"{p}fc1_w"].shape) == (128, 2 * 640, 2560)
+ assert w[f"{p}fc1_w"].dtype == torch.float8_e4m3fn
+ assert tuple(w[f"{p}fc1_s"].shape) == (128, 2 * 640 // 128, 2560 // 128)
+ assert tuple(w[f"{p}fc2_w"].shape) == (128, 2560, 640)
+ assert tuple(w[f"{p}fc2_s"].shape) == (128, 2560 // 128, 640 // 128)
+ assert tuple(w[f"{p}q_w"].shape) == (2 * 1536, 2560)
+ assert tuple(w[f"{p}k_w"].shape) == tuple(w[f"{p}v_w"].shape) == (256, 2560)
+ assert tuple(w[f"{p}o_w"].shape) == (2560, 1536)
+ assert tuple(w[f"{p}router"].shape) == (512, 2560)
+ windows = [_core(r, _TP4, MTPDecodingConfig(max_draft_len=1)) for r in range(WORLD)]
+ covered = sorted(
+ e
+ for c in windows
+ for e in range(c.mtp_expert_offset, c.mtp_expert_offset + c.mtp_local_experts)
+ )
+ assert covered == list(range(512))
+
+
+def test_the_tp4_draft_windows_sum_to_the_whole_draft_moe(monkeypatch) -> None:
+ """`_mtp_moe_block` at tp4 hands the FP8 runner this rank's whole-expert window -- the
+ draft layer's `local_expert_offset` / `local_num_experts` and the whole 640-row
+ intermediate, never the trunk's 256-wide padded slice -- and completes the window's
+ partial with the TP sum once. The entries are stood in for by arithmetic that names the
+ expert a slot routes to, so the four ranks' partials sum to the all-expert output."""
+ tokens, hidden, num_experts, topk = 5, 2560, 512, 10
+ gen = torch.Generator().manual_seed(0)
+ x = torch.randn(tokens, hidden, generator=gen)
+ router = torch.randn(hidden, num_experts, generator=gen)
+ logits = x @ router
+ top_logits, top_ids = torch.topk(logits, topk, dim=-1)
+ top_w = torch.softmax(top_logits, dim=-1)
+
+ def expert_out(ids, wts, start, count):
+ keep = ((ids >= start) & (ids < start + count)).to(x.dtype)
+ return (wts * keep * (ids + 1).to(x.dtype)).sum(-1, keepdim=True) * x
+
+ full = expert_out(top_ids, top_w, 0, num_experts)
+
+ def fake_cublas_mm(a, b, *args, **kwargs):
+ return a @ b
+
+ def fake_routing(lg, topk, output_dtype):
+ assert topk == 10 and lg.shape == (tokens, num_experts)
+ return top_ids, top_w
+
+ def fake_runner(**kw):
+ core = seen["core"]
+ assert kw["intermediate_size"] == core.moe_inter == 640
+ assert kw["local_expert_offset"] == core.mtp_expert_offset
+ assert kw["local_num_experts"] == core.mtp_local_experts == 128
+ assert tuple(kw["gemm1_weights"].shape) == (128, 2 * 640, hidden)
+ assert tuple(kw["gemm2_weights"].shape) == (128, hidden, 640)
+ assert kw["num_experts"] == num_experts and kw["top_k"] == topk
+ return expert_out(
+ kw["topk_ids"], kw["topk_weights"], kw["local_expert_offset"], kw["local_num_experts"]
+ )
+
+ monkeypatch.setattr(M, "cublas_mm", fake_cublas_mm)
+ monkeypatch.setattr(M, "renorm_moe_routing_op", fake_routing)
+ monkeypatch.setattr(M, "fp8_quantize_1x128", lambda t, use_ue8m0: (t, None))
+ monkeypatch.setattr(M, "fp8_block_scale_moe_runner", fake_runner)
+ monkeypatch.setattr(M, "silu_and_mul_rows", lambda gug, inter: gug[:, :inter])
+ # The shared expert and its gate are replicated and checked elsewhere: keep the routed sum.
+ monkeypatch.setattr(M, "fused_sigmoid_gate_mul_add", lambda routed, gate, shared: routed)
+
+ seen = {}
+ partials = []
+ for rank in range(WORLD):
+ core = _core(rank, _TP4, MTPDecodingConfig(max_draft_len=3))
+ seen["core"] = core
+ summed = []
+ core.tp_sum = lambda t: summed.append(t) or t
+ p = f"l{NUM_LAYERS}_mtp_"
+ views = SimpleNamespace(
+ router=router,
+ fc1_w=core.w[f"{p}fc1_w"],
+ fc1_s=core.w[f"{p}fc1_s"],
+ fc2_w=core.w[f"{p}fc2_w"],
+ fc2_s=core.w[f"{p}fc2_s"],
+ sh_gug=torch.zeros(hidden, 2 * 640 + 8),
+ sh_dn=torch.zeros(640, hidden),
+ # The FP8 experts (no NVFP4 copy away from one rank).
+ moe_gscale=None,
+ )
+ out = M._mtp_moe_block(core, views, x)
+ assert len(summed) == 1, "the window's partial is summed over the TP group once"
+ partials.append(out)
+ torch.testing.assert_close(sum(partials), full)
+
+
+def test_the_draft_layer_views_are_the_trunks_gemm_operands() -> None:
+ """`_mtp_views` hands the GEMMs column-major views: the shared expert's one
+ `[gate | up | gate logit | pad]` operand (the trunk layers' 2560 -> 1288 shape), the q
+ projection with its interleaved gate, and each hyper-connection site in the trunk's
+ form -- its whole grouped-norm weight, its `[down; block_inject; 0]` first projection
+ padded to 384 rows (the checkpoint's parameters rebound as its row views) and its up
+ projection -- plus the unit injection weight of the front end's site."""
+ core = _core(0, _TEP4, MTPDecodingConfig(max_draft_len=3))
+ with torch.device("cpu"):
+ for key, t in core.w.items():
+ core.w[key] = torch.nn.Parameter(
+ torch.empty(t.shape, dtype=t.dtype, device="cpu"), requires_grad=False
+ )
+ core._fuse_hc_first_projections()
+ M_ = M._mtp_views(core)
+ assert tuple(M_.sh_gug.shape) == (2560, 1288) and M_.sh_gug.stride() == (1, 2560)
+ assert tuple(M_.sh_dn.shape) == (640, 2560) and M_.sh_dn.stride() == (1, 640)
+ assert tuple(M_.o.shape) == (1536, 2560) and M_.o.stride() == (1, 1536)
+ assert tuple(M_.q.shape) == (2560, 3072) and M_.q.stride() == (1, 2560)
+ for norm in (M_.pre_h_norm, M_.ahc_norm, M_.mhc_norm, M_.fhc_norm):
+ assert tuple(norm.shape) == (4 * 2560,) and norm.is_contiguous()
+ p = f"l{NUM_LAYERS}_mtp_"
+ for tag, first, up, inject in (
+ ("ahc", M_.ahc_first, M_.ahc_up, True),
+ ("mhc", M_.mhc_first, M_.mhc_up, True),
+ ("fhc", M_.fhc_first, M_.fhc_up, False),
+ ):
+ assert tuple(first.shape) == (4 * 2560, 384) and first.stride() == (1, 4 * 2560)
+ assert tuple(up.shape) == (320, 4 * 2560) and up.stride() == (1, 320)
+ assert core.w[f"{p}{tag}_down"].data_ptr() == first.data_ptr()
+ if inject:
+ assert core.w[f"{p}{tag}_inject"].data_ptr() == first[:, 320:].data_ptr()
+ assert tuple(M_.unit_inject.shape) == (core.max_num_tokens, 4)
+ assert bool((M_.unit_inject == 1).all())
+
+
+# ---- MTP under attention DP ---------------------------------------------------------------
+
+HIDDEN, NUM_EXPERTS, TOPK = 2560, 512, 10
+
+
+def _mtp(max_draft_len: int = 3) -> MTPDecodingConfig:
+ return MTPDecodingConfig(max_draft_len=max_draft_len)
+
+
+@pytest.mark.parametrize("rank", range(WORLD))
+def test_a_dep4_rank_holds_the_whole_draft_attention_and_one_fp8_window(rank) -> None:
+ """The constructor builds MTP under attention DP. The draft layer's attention is unsharded,
+ as the trunk's dep4 attention is: 24 gated q heads, both KV heads and the whole o_proj;
+ nothing sums its output. Its FP8 experts are this rank's 128-expert window at the whole
+ 640-row intermediate, which the group's four windows complete."""
+ core = _core(rank, _DEP4, _mtp())
+ assert core.mtp_enabled and core.max_verify_rows == 4
+ assert (core.dp_size, core.dp_rank, core.attn_tp_size) == (WORLD, rank, 1)
+ assert core.spec_rollback is not None
+ p = f"l{NUM_LAYERS}_mtp_"
+ w = core.w
+ assert tuple(w[f"{p}q_w"].shape) == (2 * 6144, HIDDEN)
+ assert tuple(w[f"{p}k_w"].shape) == tuple(w[f"{p}v_w"].shape) == (512, HIDDEN)
+ assert tuple(w[f"{p}o_w"].shape) == (HIDDEN, 6144)
+ assert tuple(w[f"{p}idx_w"].shape) == (640, HIDDEN)
+ assert tuple(w[f"{p}router"].shape) == (NUM_EXPERTS, HIDDEN)
+ assert (core.mtp_local_experts, core.mtp_expert_offset, core.moe_inter) == (
+ 128,
+ 128 * rank,
+ 640,
+ )
+ assert tuple(w[f"{p}fc1_w"].shape) == (128, 2 * 640, HIDDEN)
+ assert w[f"{p}fc1_w"].dtype == torch.float8_e4m3fn
+ assert tuple(w[f"{p}fc2_w"].shape) == (128, HIDDEN, 640)
+ windows = [_core(r, _DEP4, _mtp(1)) for r in range(WORLD)]
+ covered = sorted(
+ e
+ for c in windows
+ for e in range(c.mtp_expert_offset, c.mtp_expert_offset + c.mtp_local_experts)
+ )
+ assert covered == list(range(NUM_EXPERTS))
+
+
+def test_a_draft_step_pads_by_the_counts_the_worker_hands_it() -> None:
+ """The spec worker hands each draft step the group's row counts -- every rank's trunk
+ tokens on the loop's first step, its request count on the later ones -- and does not write
+ them into the metadata. A count that disagrees with this rank's rows would gather the wrong
+ split, and a missing one cannot pad at all; without attention DP nothing is padded."""
+ core = _core(2, _DEP4, _mtp())
+ # Draft-extend: rank 2 holds 9 trunk rows, the group's largest is 17.
+ assert M._draft_pad_rows(core, [5, 17, 9, 1], 9) == 8
+ # A later draft step: one row a request, attention-DP dummy requests included.
+ assert M._draft_pad_rows(core, [2, 4, 3, 1], 3) == 1
+ # A captured decode step pads every rank to one batch.
+ assert M._draft_pad_rows(core, [4, 4, 4, 4], 4) == 0
+ with pytest.raises(AssertionError, match="wrong split"):
+ M._draft_pad_rows(core, [5, 17, 9, 1], 10)
+ with pytest.raises(AssertionError, match="without the group's row counts"):
+ M._draft_pad_rows(core, None, 9)
+ assert M._draft_pad_rows(_core(0, _TEP4, _mtp()), None, 7) == 0
+
+
+def _expert_out(x, ids, wts, start, count):
+ """A stand-in MoE whose output names the experts a row routes to inside `[start, start +
+ count)`: windows of whole experts sum to the all-expert output."""
+ keep = ((ids >= start) & (ids < start + count)).to(x.dtype)
+ return (wts * keep * (ids + 1).to(x.dtype)).sum(-1, keepdim=True) * x
+
+
+@pytest.mark.parametrize(
+ "counts",
+ [[3, 7, 1, 4], [1, 2, 1, 3], [4, 4, 4, 4]],
+ ids=["draft-extend", "draft-decode-with-a-dummy-rank", "captured-batch"],
+)
+def test_every_rank_completes_its_draft_rows_through_the_round_trip(counts, monkeypatch) -> None:
+ """At every draft step every rank pads its rows to the group's largest count, gathers the
+ group's rows (one allgather), routes, quantizes and runs its window over all of them in
+ chunks of at most `_MTP_MOE_MAX_ROWS`, and reduce-scatters (one reduce-scatter): the same
+ two collectives at the same shape on every rank, a rank holding one dummy request included.
+ Each rank gets back the all-expert MoE of its own rows; nothing goes through the TP sum."""
+ chunk_rows = 5
+ monkeypatch.setattr(M, "_MTP_MOE_MAX_ROWS", chunk_rows)
+ gen = torch.Generator().manual_seed(0)
+ dt = torch.float64
+ xs = [torch.randn(n, HIDDEN, generator=gen, dtype=dt) for n in counts]
+ router = torch.randn(HIDDEN, NUM_EXPERTS, generator=gen, dtype=dt)
+ rows = max(counts)
+ gathered = torch.cat([nn_pad(x, rows) for x in xs], 0)
+
+ def route(t):
+ top_logits, top_ids = torch.topk(t @ router, TOPK, dim=-1)
+ return top_ids, torch.softmax(top_logits, dim=-1)
+
+ ids, wts = route(gathered)
+ full = _expert_out(gathered, ids, wts, 0, NUM_EXPERTS)
+
+ state = {}
+
+ def fake_cublas_mm(a, b, *args, **kwargs):
+ return a @ b
+
+ def fake_routing(lg, topk, output_dtype):
+ assert topk == TOPK and lg.shape[1] == NUM_EXPERTS and output_dtype == dt
+ top_logits, top_ids = torch.topk(lg, topk, dim=-1)
+ return top_ids, torch.softmax(top_logits, dim=-1)
+
+ def fake_runner(**kw):
+ core = state["core"]
+ assert kw["intermediate_size"] == core.moe_inter == 640
+ assert kw["local_expert_offset"] == 128 * core.dp_rank
+ assert kw["local_num_experts"] == core.mtp_local_experts == 128
+ x = kw["hidden_states"]
+ assert 0 < x.shape[0] <= chunk_rows
+ state["runner_rows"] += x.shape[0]
+ return _expert_out(x, kw["topk_ids"], kw["topk_weights"], kw["local_expert_offset"], 128)
+
+ def fake_allgather(t, sizes, group, process_group):
+ assert sizes is None and group == list(range(WORLD)) and process_group is None
+ state["collectives"].append(("allgather", tuple(t.shape)))
+ torch.testing.assert_close(t, nn_pad(xs[state["rank"]], rows), rtol=0, atol=0)
+ return gathered
+
+ def fake_reducescatter(t, sizes, group, process_group):
+ assert sizes is None and group == list(range(WORLD)) and process_group is None
+ state["collectives"].append(("reducescatter", tuple(t.shape)))
+ state["partial"] = t
+ r = state["rank"]
+ return full[r * rows : (r + 1) * rows]
+
+ monkeypatch.setattr(M, "cublas_mm", fake_cublas_mm)
+ monkeypatch.setattr(M, "renorm_moe_routing_op", fake_routing)
+ monkeypatch.setattr(M, "fp8_quantize_1x128", lambda t, use_ue8m0: (t, None))
+ monkeypatch.setattr(M, "fp8_block_scale_moe_runner", fake_runner)
+ monkeypatch.setattr(M, "silu_and_mul_rows", lambda gug, inter: gug[:, :inter])
+ # The shared expert and its gate are replicated and checked elsewhere: keep the routed sum.
+ monkeypatch.setattr(M, "fused_sigmoid_gate_mul_add", lambda routed, gate, shared: routed)
+ monkeypatch.setattr(M, "allgather", fake_allgather)
+ monkeypatch.setattr(M, "reducescatter", fake_reducescatter)
+ monkeypatch.setattr(M, "mpi_disabled", lambda: False)
+
+ per_rank = []
+ for rank in range(WORLD):
+ core = _core(rank, _DEP4, _mtp())
+ core.expert_sum = core._expert_sum_op()
+
+ def no_tp_sum(t):
+ raise AssertionError("the draft MoE under attention DP went through the TP sum")
+
+ core.tp_sum = no_tp_sum
+ state.update(core=core, rank=rank, collectives=[], runner_rows=0)
+ p = f"l{NUM_LAYERS}_mtp_"
+ views = SimpleNamespace(
+ router=router,
+ fc1_w=core.w[f"{p}fc1_w"],
+ fc1_s=core.w[f"{p}fc1_s"],
+ fc2_w=core.w[f"{p}fc2_w"],
+ fc2_s=core.w[f"{p}fc2_s"],
+ sh_gug=torch.zeros(HIDDEN, 2 * 640 + 8, dtype=dt),
+ sh_dn=torch.zeros(640, HIDDEN, dtype=dt),
+ # The FP8 experts (no NVFP4 copy away from one rank).
+ moe_gscale=None,
+ )
+ pad_rows = M._draft_pad_rows(core, counts, counts[rank])
+ out = M._mtp_moe_block(core, views, xs[rank], pad_rows)
+ # The window ran over every gathered row once, chunk by chunk.
+ assert state["runner_rows"] == WORLD * rows
+ torch.testing.assert_close(
+ state["partial"], _expert_out(gathered, ids, wts, 128 * rank, 128)
+ )
+ assert tuple(out.shape) == (counts[rank], HIDDEN)
+ torch.testing.assert_close(out, _expert_out(xs[rank], *route(xs[rank]), 0, NUM_EXPERTS))
+ per_rank.append(state["collectives"])
+ assert all(c == per_rank[0] for c in per_rank), per_rank
+ assert per_rank[0] == [
+ ("allgather", (rows, HIDDEN)),
+ ("reducescatter", (WORLD * rows, HIDDEN)),
+ ]
+
+
+def _calls(fn: ast.AST, name: str) -> list[ast.Call]:
+ return [
+ node
+ for node in ast.walk(fn)
+ if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == name
+ ]
+
+
+@pytest.mark.parametrize("target", [M.PrefillTarget, M.DecodeTarget])
+def test_both_draft_targets_pad_the_draft_moe_by_the_workers_counts(target) -> None:
+ """A draft step runs on the prefill target while the metadata still carries context rows
+ and on the decode target otherwise; ranks of one attention-DP group can be on either at the
+ same draft step. Both pad the draft MoE by the `all_rank_num_tokens` the worker passes them
+ for that step, never by the metadata's, which still holds the trunk step's counts."""
+ fn = ast.parse(inspect.getsource(target.mtp_forward).lstrip())
+ (pad,) = _calls(fn, "_draft_pad_rows")
+ assert [ast.unparse(a) for a in pad.args] == ["core", "all_rank_num_tokens", "rows"]
+ (tail,) = _calls(fn, "_mtp_tail")
+ assert ast.unparse(tail.args[-1]) == "pad_rows"
+ assert "_dp_pad_rows" not in ast.unparse(fn)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_decode_paths.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_decode_paths.py
new file mode 100644
index 000000000000..6c88e8525cb2
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_decode_paths.py
@@ -0,0 +1,295 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The target's decode-side paths at one rank (tp1): the draft layer's front end, the heads
+and their logits, and the Gated-DeltaNet verify step's fused conv.
+
+* **The draft front end**: the token embedding under rejection sampling, and the dense
+ projections' dispatch between `fp8_gemv` and the bf16 projection.
+* **The logits**: the target's logits processor returns the shell's logits bit for bit (the
+ same row gather, the same fp32 values, its head's bf16 logits widened by one vectorized
+ kernel); a consumer of the row argmax alone may take the bf16 logits; the dense FP8 verify
+ head reads the NVFP4 copy at decode rows.
+* **The verify step's conv** (`DecodeTarget._gdn_block`) under the dense FP8 mode and over the
+ in_proj's NVFP4 copy (`TRTLLM_MODELING_V2_GDN_INPROJ_NVFP4`, up to
+ `_INPROJ_NVFP4_MAX_ROWS`). A captured verify step runs the in_proj with the multi-token conv
+ step as its epilogue (`fp8_gemv_conv1d_verify` / `nvfp4_split_gemv_conv1d_verify`); an eager
+ one runs the projection and then `causal_conv1d_triton_update`. Both must hand the cached
+ replay the same thing -- the conv output (contiguous, as the replay reads it), `a`, `b` and
+ `z` bitwise -- and leave the conv pool and the per-step snapshot buffer bitwise the same. The
+ replay is stubbed to record its inputs (it is the same call on both sides), the epilogue to
+ record `z`. Over a bf16 pool (the CuTe replay) `a` and `b` reach it as views of the
+ projection, not copies.
+"""
+
+import types
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.fp8_gemv import ( # noqa: E501
+ pack_from_bf16,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.nvfp4_gemv import ( # noqa: E501
+ pack_split_from_bf16,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501
+ modeling as M,
+)
+from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor
+from tensorrt_llm._torch.speculative.interface import rejection_sampling_one_model
+
+_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="runs on a CUDA device")
+DEV, BF16 = "cuda", torch.bfloat16
+
+
+# ---- the draft front end ------------------------------------------------------------------
+
+
+@_cuda
+def test_draft_embedding_reads_row_0_for_rejected_slots() -> None:
+ """Rejection sampling pads the slots past a request's accepted count with -1 (here every
+ draft is rejected at slot 0); the front end embeds those as row 0 and every valid id as
+ the plain lookup does."""
+ batch, draft, vocab, hidden = 8, 3, 4096, 256
+ draft_tokens = torch.arange(batch * draft, dtype=torch.int32, device=DEV).reshape(batch, draft)
+ draft_probs = torch.zeros(batch, draft, vocab, device=DEV)
+ draft_probs.scatter_(2, draft_tokens.long().unsqueeze(-1), 1.0)
+ target_probs = torch.zeros(batch, draft + 1, vocab, device=DEV)
+ target_probs[..., vocab - 1] = 1.0
+ seed = torch.zeros(batch, dtype=torch.int64, device=DEV)
+ offset = torch.arange(batch, dtype=torch.int64, device=DEV)
+ accepted, _ = rejection_sampling_one_model(
+ draft_probs, draft_tokens, target_probs, seed, offset
+ )
+ assert bool((accepted[:, 1:] == -1).all()) and bool((accepted[:, 0] == vocab - 1).all())
+
+ g = torch.Generator(device=DEV).manual_seed(0)
+ table = torch.randn(vocab, hidden, generator=g, device=DEV).to(BF16)
+ valid_ids = torch.randint(0, vocab, (257,), generator=g, device=DEV, dtype=torch.int32)
+ for ids in (accepted.flatten(), valid_ids, valid_ids.long()):
+ e = M._draft_embedding(ids, table)
+ valid = ids >= 0
+ assert torch.equal(e[valid], table[ids[valid].long()])
+ assert torch.equal(e[~valid], table[0].expand(int((~valid).sum()), -1))
+
+
+@_cuda
+@pytest.mark.parametrize("rows", [1, 16, 17])
+def test_draft_projection_dispatch(rows: int) -> None:
+ """`_draft_mm`: `fp8_gemv` over the packed copy up to `_MTP_DENSE_FP8_MAX_ROWS` rows, the
+ bf16 projection past them and without a copy."""
+ g = torch.Generator(device=DEV).manual_seed(rows)
+ w = (torch.randn(2560, 2560, generator=g, device=DEV) * 0.02).to(BF16)
+ x = torch.randn(rows, 2560, generator=g, device=DEV).to(BF16)
+ packed = pack_from_bf16(w)
+ got = M._draft_mm(x, w.t(), packed, cublas_mm)
+ if rows <= M._MTP_DENSE_FP8_MAX_ROWS:
+ assert torch.equal(got, M.fp8_gemv_op(x, packed))
+ ref = x.float() @ w.float().t()
+ assert ((got.float() - ref).norm() / ref.norm()).item() < 0.05
+ else:
+ assert torch.equal(got, cublas_mm(x, w.t()))
+ assert torch.equal(M._draft_mm(x, w.t(), None, cublas_mm), cublas_mm(x, w.t()))
+
+
+# ---- the logits ---------------------------------------------------------------------------
+
+VOCAB, HIDDEN = 248320, 2560
+
+
+@pytest.fixture(scope="module")
+def head():
+ g = torch.Generator(device=DEV).manual_seed(0)
+ lin = torch.nn.Linear(HIDDEN, VOCAB, bias=False, device=DEV, dtype=BF16)
+ with torch.no_grad():
+ lin.weight.copy_(torch.randn(VOCAB, HIDDEN, generator=g, device=DEV) * 0.02)
+ return lin
+
+
+@_cuda
+@pytest.mark.parametrize("seq_lens", [[1] * 128, [5000, 3000, 1] + [1] * 121])
+def test_logits_match_the_shell_processor_bitwise(head, seq_lens) -> None:
+ g = torch.Generator(device=DEV).manual_seed(len(seq_lens))
+ lens = torch.tensor(seq_lens, dtype=torch.int32, device=DEV)
+ hidden = torch.randn(int(lens.sum()), HIDDEN, generator=g, device=DEV).to(BF16)
+ md = SimpleNamespace(seq_lens_cuda=lens)
+ with torch.no_grad():
+ want = LogitsProcessor().forward(hidden, head, md)
+ got = M._WideningLogitsProcessor().forward(hidden, head, md)
+ assert got.dtype == torch.float32 and got.shape == (len(seq_lens), VOCAB)
+ assert torch.equal(got.view(torch.int32), want.view(torch.int32))
+
+
+@_cuda
+def test_context_logits_and_no_metadata(head) -> None:
+ g = torch.Generator(device=DEV).manual_seed(9)
+ hidden = torch.randn(33, HIDDEN, generator=g, device=DEV).to(BF16)
+ with torch.no_grad():
+ for md, ctx in ((None, False), (SimpleNamespace(seq_lens_cuda=None), True)):
+ want = LogitsProcessor().forward(hidden, head, md, ctx)
+ got = M._WideningLogitsProcessor().forward(hidden, head, md, ctx)
+ assert torch.equal(got.view(torch.int32), want.view(torch.int32))
+
+
+@_cuda
+@pytest.mark.parametrize("rows", [1, 4, 256])
+def test_unwidened_logits_give_the_same_greedy_tokens(head, rows) -> None:
+ """`widen=False` returns the head's own bf16 logits, whose greedy samplers (the verify
+ sampler's `greedy_argmax`, the draft sampler's int32 one) pick the widened logits' tokens."""
+ from tensorrt_llm._torch.speculative.greedy_argmax import greedy_argmax
+
+ g = torch.Generator(device=DEV).manual_seed(rows)
+ hidden = torch.randn(rows, HIDDEN, generator=g, device=DEV).to(BF16)
+ proc = M._WideningLogitsProcessor()
+ with torch.no_grad():
+ wide = proc.forward(hidden, head, None, True)
+ narrow = proc.forward(hidden, head, None, True, widen=False)
+ # Without return_context_logits the shell's gather and fp32 contract stay.
+ gathered = proc.forward(hidden, head, None, False, widen=False)
+ assert narrow.dtype == BF16 and wide.dtype == torch.float32
+ assert gathered.dtype == torch.float32
+ assert torch.equal(narrow.float().view(torch.int32), wide.view(torch.int32))
+ narrow[0, 17] = narrow[0].max() # an exact tie, broken to the lower index either way
+ wide = narrow.float()
+ assert torch.equal(greedy_argmax(narrow), greedy_argmax(wide))
+ assert torch.equal(greedy_argmax(narrow, torch.int32), greedy_argmax(wide, torch.int32))
+ assert torch.equal(greedy_argmax(narrow), torch.argmax(wide, dim=-1))
+
+
+@_cuda
+@pytest.mark.parametrize("rows", [1, 3, 16, 17])
+def test_fp8_verify_head_reads_the_nvfp4_copy_at_decode_rows(head, rows) -> None:
+ """With an NVFP4 copy (`_TARGET_HEAD_NVFP4_ENV`) the dense FP8 verify head's decode rows
+ are `nvfp4_gemv` over it; past `_FP8_GEMV_HEAD_ROWS` they take the per-row FP8 copy as
+ the head without it does."""
+ rowwise = M._fp8_rowwise(head.weight.data)
+ nv = M.nvfp4_pack_from_bf16(head.weight.data)
+ g = torch.Generator(device=DEV).manual_seed(100 + rows)
+ hidden = torch.randn(rows, HIDDEN, generator=g, device=DEV).to(BF16)
+ with torch.no_grad():
+ got = M._Fp8Head(*rowwise, None, nv)(hidden)
+ if rows <= M._FP8_GEMV_HEAD_ROWS:
+ want = M.nvfp4_gemv_op(hidden, nv)
+ else:
+ want = M._Fp8Head(*rowwise, None, None)(hidden)
+ assert got.shape == (rows, VOCAB) and torch.equal(got, want)
+
+
+def _gate(guided=None, greedy=True, rejection=False, penalty=False, relaxed=False):
+ worker = SimpleNamespace(
+ guided_decoder=guided,
+ spec_config=SimpleNamespace(use_relaxed_acceptance_for_thinking=relaxed),
+ )
+ md = SimpleNamespace(
+ is_all_greedy_sample=greedy,
+ wants_advanced_draft_sampling=rejection and not greedy,
+ enable_penalty=penalty,
+ )
+ return M._spec_argmax_only(worker, md)
+
+
+@pytest.mark.cpu_only
+def test_argmax_only_gate() -> None:
+ assert _gate() == (True, True)
+ # A non-greedy batch: the verify sampler samples; the draft stays argmax unless
+ # rejection sampling needs its distribution.
+ assert _gate(greedy=False) == (False, True)
+ assert _gate(greedy=False, rejection=True) == (False, False)
+ # The occurrence penalties and relaxed acceptance read the verify logits only.
+ assert _gate(penalty=True) == (False, True)
+ assert _gate(relaxed=True) == (False, True)
+ # A guided decoder edits both.
+ assert _gate(guided=object()) == (False, False)
+
+
+# ---- the Gated-DeltaNet verify step's conv ------------------------------------------------
+
+KEY, HV, V, CONV, VD, SLOTS, ROWS = 16, 48, 128, 10240, 6144, 461, 32
+
+
+@pytest.fixture(scope="module", params=["fp8", "nvfp4"])
+def gdn_layer(request):
+ g = torch.Generator(device=DEV).manual_seed(0)
+ w = (torch.randn(CONV + VD + 2 * HV, HIDDEN, generator=g, device=DEV) * HIDDEN**-0.5).to(BF16)
+ nv = request.param == "nvfp4"
+ return types.SimpleNamespace(
+ in_proj=w.t(),
+ in_proj8=None,
+ in_proj_m=None,
+ in_proj_g=None if nv else pack_from_bf16(w),
+ in_proj_nv=pack_split_from_bf16(w) if nv else None,
+ conv_w=(torch.randn(CONV, 4, generator=g, device=DEV) * 0.5).to(BF16),
+ a_log=None,
+ dt_bias=None,
+ )
+
+
+def _gdn_block(layer, x, si, int_idx, conv0, win0, p, captured, monkeypatch):
+ """One verify `_gdn_block` on copies of the pools; returns what the cached replay and the
+ epilogue were handed, and the pools."""
+ seen = {}
+
+ def replay(**kw):
+ mq_out = kw["packed_qkv"]
+ n, p_ = kw["g"].shape[:2]
+ seen.update(
+ mq_out=mq_out.clone(),
+ a=kw["g"].clone(),
+ b=kw["beta"].clone(),
+ pdl=kw["launch_with_pdl"],
+ )
+ assert mq_out.is_contiguous() and mq_out.shape == (n * p_, CONV)
+ # The CuTe replay reads a / b through their row stride: views, not copies.
+ assert kw["g"].stride(1) > HV and kw["beta"].stride(1) > HV
+ return torch.zeros(n, p_, HV, V, dtype=BF16, device=DEV)
+
+ monkeypatch.setattr(M, "gdn_cached_replay_cute", replay)
+ history = torch.zeros(1, device=DEV)
+ core = types.SimpleNamespace(
+ conv_dim=CONV,
+ key_dim=KEY * 128,
+ value_dim=VD,
+ num_k_heads=KEY,
+ head_k_dim=128,
+ num_v_heads=HV,
+ head_v_dim=V,
+ in_proj_split=(CONV, VD, HV, HV),
+ decode_in_mm=cublas_mm,
+ _bound=types.SimpleNamespace(
+ pools=types.SimpleNamespace(
+ conv={0: conv0.clone()}, ssm={0: torch.zeros(1, dtype=BF16, device=DEV)}
+ ),
+ int_conv={0: win0.clone()},
+ replay=types.SimpleNamespace(
+ old={0: (history,) * 4}, cache_buf_idx=None, prev_num_accepted=None, history=16
+ ),
+ ),
+ )
+ target = object.__new__(M.DecodeTarget)
+ target.core = core
+ target._gdn_epilogue = lambda out, z, L, num_tokens: z.clone()
+ verify = (int_idx, None, None, False)
+ z = target._gdn_block(x, 0, layer, si, verify, p, captured=captured)
+ pools = core._bound.pools
+ return dict(seen, z=z, conv=pools.conv[0], win=core._bound.int_conv[0])
+
+
+@_cuda
+@pytest.mark.parametrize("n,p", [(1, 4), (2, 4), (8, 2), (2, 8)])
+def test_fused_verify_conv_matches_the_two_launches(gdn_layer, n, p, monkeypatch) -> None:
+ if gdn_layer.in_proj_nv is not None and n * p > M._INPROJ_NVFP4_MAX_ROWS:
+ pytest.skip("past the NVFP4 copy's rows")
+ g = torch.Generator(device=DEV).manual_seed(10 * n + p)
+ conv0 = torch.randn(SLOTS, CONV, 3, generator=g, device=DEV).to(BF16)
+ win0 = torch.randn(ROWS, p, CONV, 3, generator=g, device=DEV).to(BF16)
+ x = torch.randn(n * p, HIDDEN, generator=g, device=DEV).to(BF16)
+ si = torch.randperm(SLOTS, generator=g, device=DEV)[:n].to(torch.int32)
+ int_idx = torch.arange(ROWS, dtype=torch.int32, device=DEV)[:n]
+ eager = _gdn_block(gdn_layer, x, si, int_idx, conv0, win0, p, False, monkeypatch)
+ fused = _gdn_block(gdn_layer, x, si, int_idx, conv0, win0, p, True, monkeypatch)
+ for key in ("mq_out", "a", "b", "z", "conv", "win"):
+ assert torch.equal(eager[key].view(torch.int16), fused[key].view(torch.int16)), key
+ assert fused["pdl"] == eager["pdl"] == M._GDN_REPLAY_PDL
+ assert not torch.equal(fused["conv"], conv0) and not torch.equal(fused["win"], win0)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_modeling_v2_spec_rollback.py
similarity index 72%
rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py
rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_modeling_v2_spec_rollback.py
index fc1099ea6537..1936894dbafa 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_modeling_v2_spec_rollback.py
@@ -1,35 +1,37 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the target's speculative rollback handler (`SpeculativeRollback`).
-
-A verify step is simulated with the target's own state-writing ops -- the n-gram id
-history, the dilated short conv and the QSA index store, plus an anchor-slot overwrite
-standing in for the compression -- run over every draft row the way the forward runs
-them, with the handler staged at the forward's call sites and committed afterwards.
-
-The reference shares no code with the commit: it runs the same ops from the same
-pre-step state over each generation request's accepted prefix only, which is by
-definition the state a verify step must leave behind. Every pool and plane is compared
-bitwise in full, so a write to a slot the step does not own fails too. The step is also
-run sync-free and under CUDA-graph capture and replay with the accepted counts changed
-in place. Each check runs over the three handler forms: per-layer plane staging, the
-batched plane gather / restore, and the fused one-launch staging and commit.
+"""GPU test for the default target's speculative rollback handler (`SpeculativeRollback`).
+
+A verify step is simulated with the state writes the forward makes over every draft row:
+the n-gram id history (`ngram_token_history`), the PLE dilated-conv state (a plain-torch
+stand-in for the conv inside `ple_gate_conv_inject_norm`: the state after a run of rows is
+the last `state_len` columns of `[carried state | the rows' conv input]`), the QSA index
+store, and an anchor-slot overwrite standing in for the compression. The handler is staged
+at the forward's call sites and committed afterwards, as the spec worker commits it.
+
+The reference shares no code with the commit: it runs the same writes from the same
+pre-step state over each generation request's accepted prefix only, which is by definition
+the state a verify step must leave behind. Every pool and plane is compared bitwise in full,
+so a write to a slot the step does not own fails too. The step also runs sync-free and under
+CUDA-graph capture and replay with the accepted counts changed in place. Each check runs over
+both handler forms: the batched plane gather / restore with torch PLE staging (the prefill
+target's and an eager decode step's), and the one-launch staging and commit (a captured
+decode step's).
"""
+from __future__ import annotations
+
import pytest
import torch
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.modeling import ( # noqa: E501
- SpeculativeRollback,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_index_store import ( # noqa: E501
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_index_store import ( # noqa: E501
qsa_paged_index_store,
)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ngram_token_history import ( # noqa: E501
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ngram_token_history import ( # noqa: E501
ngram_token_history,
)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.dilated_short_conv_paged import ( # noqa: E501
- dilated_short_conv_paged,
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.modeling import ( # noqa: E501
+ SpeculativeRollback,
)
assert torch.cuda.is_available(), "the rollback handler test requires a CUDA device"
@@ -52,12 +54,12 @@
class _Batch:
- """One verify step: `ctx_lens` context requests (fresh) followed by
- generation requests of `p` rows each, at `starts` cached tokens, in the
- recurrent `slots` given (contexts first). Generation request `s` reuses the
- block-table row of request `dup_of[s]` when given -- the CUDA-graph padding
- dummy, repeated. With `shared_page` the live index and position planes are strided
- views of one per-slot buffer, as an NVFP4 KV cache's manager serves them."""
+ """One verify step: `ctx_lens` context requests (fresh) followed by generation requests
+ of `p` rows each, at `starts` cached tokens, in the recurrent `slots` given (contexts
+ first). Generation request `s` reuses the block-table row of request `dup_of[s]` when
+ given -- the CUDA-graph padding dummy, repeated. With `shared_page` the live index and
+ position planes are strided views of one per-slot buffer, as an NVFP4 KV cache's manager
+ serves them."""
def __init__(self, seed, ctx_lens, gen_starts, p, slots, dup_of=None, shared_page=False):
g = torch.Generator().manual_seed(seed)
@@ -99,7 +101,6 @@ def __init__(self, seed, ctx_lens, gen_starts, p, slots, dup_of=None, shared_pag
self.position0 = torch.randint(0, 1 << 20, (PAGES, TPB, 3), generator=g).to(
DEV, torch.int32
)
- self.weight = (0.3 * torch.randn(CHANNELS, 1, KERNEL, generator=g)).to(DEV, torch.bfloat16)
# Step inputs.
self.ids = torch.randint(0, VOCAB, (self.rows,), generator=g).to(DEV, torch.int32)
@@ -166,9 +167,22 @@ def _qsl(lengths):
return torch.tensor(starts, dtype=torch.int64, device=DEV)
+def _conv_state(conv, x, lengths, slots, carry):
+ """The PLE conv pool after each sequence's rows: the last `STATE_LEN` columns of its
+ carried state (zeros for a fresh one) followed by its rows' conv input, channel-major.
+ Host-sized slices only, so it captures."""
+ start = 0
+ for s, length in enumerate(lengths):
+ slot = slots[s : s + 1].long()
+ prev = conv[slot][0] if carry[s] else torch.zeros_like(conv[0])
+ window = torch.cat([prev, x[start : start + length].t()], 1)
+ conv[slot] = window[:, -STATE_LEN:].unsqueeze(0)
+ start += length
+
+
def _write_planes(b, layer_index, position, req_idx, logical, keys, anchor_keys, coords):
- """The index store, then a compression stand-in that overwrites each anchor
- row's own slot -- the two writers of an attention layer's planes."""
+ """The index store, then a compression stand-in that overwrites each anchor row's own
+ slot -- the two writers of an attention layer's planes."""
qsa_paged_index_store(keys, coords, layer_index, position, req_idx, logical, b.block_table, TPB)
pages = b.block_table[req_idx.long(), logical // TPB].long()
within = logical % TPB
@@ -176,13 +190,12 @@ def _write_planes(b, layer_index, position, req_idx, logical, keys, anchor_keys,
layer_index[pages, within] = torch.where(anchor, anchor_keys, layer_index[pages, within])
-def _verify_step(b, rollback, fused=False):
- """The forward's write sequence over all `p` rows of every generation
- request, staged at the forward's call sites, then the worker's commit. The
- PLE pools stage as the fused PLE launches do: `stage_ple_verify` before the
- id history publishes, `stage_ple_conv_input` from the conv input."""
+def _verify_step(b, rollback, fused):
+ """The forward's write sequence over all `p` rows of every generation request, staged
+ at the forward's call sites, then the worker's commit: `stage_ple_verify` before the id
+ history publishes, `stage_ple_conv_input` from the conv input, `stage_index` before
+ each attention layer's store."""
rollback.discard()
- qsl = b.qsl
rollback.begin(
num_contexts=b.num_contexts,
prefix_tokens=b.prefix,
@@ -196,11 +209,9 @@ def _verify_step(b, rollback, fused=False):
fused=fused,
)
rollback.stage_ple_verify(b.ngram, b.conv, b.ids)
- ngram_token_history(b.ids, b.ngram, b.slots, b.carry, qsl, EOS, NGRAM)
+ ngram_token_history(b.ids, b.ngram, b.slots, b.carry, b.qsl, EOS, NGRAM)
rollback.stage_ple_conv_input(b.x)
- dilated_short_conv_paged(
- b.x.t(), b.weight, None, b.conv, qsl, b.slots, b.carry, DILATION, True, -1
- )
+ _conv_state(b.conv, b.x, b.lengths, b.slots, [False] * b.num_contexts + [True] * b.n)
for layer in range(LAYERS):
rollback.stage_index(b.index[layer], b.position)
_write_planes(
@@ -217,8 +228,8 @@ def _verify_step(b, rollback, fused=False):
def _reference(b):
- """The state after running the same ops from the pre-step state over the
- context rows and each generation request's accepted prefix only."""
+ """The state after running the same writes from the pre-step state over the context
+ rows and each generation request's accepted prefix only."""
keep = list(range(b.prefix))
lengths = list(b.ctx_lens)
for s, a in enumerate(b.gen_accepted):
@@ -226,14 +237,11 @@ def _reference(b):
keep += range(first, first + a)
lengths.append(a)
keep = torch.tensor(keep, dtype=torch.long, device=DEV)
- qsl = _qsl(lengths)
ngram, conv = b.ngram0.clone(), b.conv0.clone()
index = [t.clone() for t in b.index0]
position = b.position0.clone()
- ngram_token_history(b.ids[keep], ngram, b.slots, b.carry, qsl, EOS, NGRAM)
- dilated_short_conv_paged(
- b.x[keep].t(), b.weight, None, conv, qsl, b.slots, b.carry, DILATION, True, -1
- )
+ ngram_token_history(b.ids[keep], ngram, b.slots, b.carry, _qsl(lengths), EOS, NGRAM)
+ _conv_state(conv, b.x[keep], lengths, b.slots, [False] * b.num_contexts + [True] * b.n)
for layer in range(LAYERS):
_write_planes(
b,
@@ -261,12 +269,13 @@ def _assert_state(b, expected, what):
assert torch.equal(b.position, position), f"{what}: position plane"
-# (ctx_lens, gen_starts, p, gen_accepted). Positions straddle pages and land on
-# compression anchors; one request accepts everything. All three are the accepted-prefix check;
-# `decode_p8` also drives the shared-page and capture tests.
+# (ctx_lens, gen_starts, p, gen_accepted). Positions straddle pages and land on compression
+# anchors; one request accepts everything. `decode_p8` also drives the shared-page test.
CASES = {
"decode": ((), (5, 30, 61, 100), 4, (1, 2, 4, 3)),
"mixed": ((7, 3), (12, 29, 95), 4, (2, 1, 3)),
+ # max_draft_len 1: the shortest verify step.
+ "decode_p2": ((), (3, 31, 64), 2, (1, 2, 1)),
# max_draft_len 7: every acceptance 1..8 once; blocks at 27, 60 and 90 straddle a page,
# and an 8-row block completes two compression groups.
"decode_p8": ((), (0, 27, 60, 90, 101, 5, 44, 119), 8, (1, 2, 3, 4, 5, 6, 7, 8)),
@@ -277,26 +286,25 @@ def _slots(num):
return list(range(3, 3 + 2 * num, 2))
-#: Handler forms: per-layer plane staging; the planes bound (`begin` gathers them all in one
-#: launch and the commit restores them in one, `qsa_rollback_gather` / `qsa_rollback_restore`);
-#: the planes bound and the step opened `fused` (`spec_rollback_stage` / `spec_rollback_commit`).
-MODES = ["per_layer", "batched", "fused"]
+#: Handler forms: torch PLE staging with the planes gathered / restored in one launch each
+#: (`qsa_rollback_gather` / `qsa_rollback_restore`), and the step opened `fused`
+#: (`spec_rollback_stage` / `spec_rollback_commit`).
+MODES = ["batched", "fused"]
-def _rollback(b, mode):
+def _rollback(b):
rollback = SpeculativeRollback()
- if mode != "per_layer":
- rollback.set_index_planes(b.index)
+ rollback.set_index_planes(b.index)
return rollback
@pytest.mark.parametrize("mode", MODES)
-@pytest.mark.parametrize("case", ["decode", "mixed", "decode_p8"])
+@pytest.mark.parametrize("case", sorted(CASES))
def test_commit_matches_the_accepted_prefix(case, mode):
ctx_lens, gen_starts, p, gen_accepted = CASES[case]
b = _Batch(11, ctx_lens, gen_starts, p, _slots(len(ctx_lens) + len(gen_starts)))
b.set_accepted(gen_accepted)
- _verify_step(b, _rollback(b, mode), mode == "fused")
+ _verify_step(b, _rollback(b), mode == "fused")
_assert_state(b, _reference(b), case)
@@ -312,7 +320,7 @@ def test_commit_over_the_strided_planes_of_a_shared_page(mode):
13, ctx_lens, gen_starts, p, _slots(len(ctx_lens) + len(gen_starts)), shared_page=True
)
b.set_accepted(gen_accepted)
- _verify_step(b, _rollback(b, mode), mode == "fused")
+ _verify_step(b, _rollback(b), mode == "fused")
_assert_state(b, _reference(b), case)
@@ -320,8 +328,7 @@ def test_commit_without_staging_is_a_no_op():
ctx_lens, gen_starts, p, gen_accepted = CASES["decode"]
b = _Batch(14, ctx_lens, gen_starts, p, _slots(len(gen_starts)))
b.set_accepted(gen_accepted)
- rollback = SpeculativeRollback()
- rollback.commit_speculative_states(b.accepted, b.slots, 0)
+ _rollback(b).commit_speculative_states(b.accepted, b.slots, 0)
_assert_state(b, (b.ngram0, b.conv0, b.index0, b.position0), "unstaged")
@@ -330,7 +337,7 @@ def test_discard_drops_a_stale_staging():
ctx_lens, gen_starts, p, gen_accepted = CASES["decode"]
b = _Batch(15, ctx_lens, gen_starts, p, _slots(len(gen_starts)))
b.set_accepted(gen_accepted)
- rollback = SpeculativeRollback()
+ rollback = _rollback(b)
rollback.begin(
num_contexts=0,
prefix_tokens=0,
@@ -347,16 +354,34 @@ def test_discard_drops_a_stale_staging():
_assert_state(b, (b.ngram0, b.conv0, b.index0, b.position0), "discarded")
+def test_a_step_without_bound_planes_is_refused():
+ """The trunk's bind gives the handler every attention layer's index plane before any
+ verify step; a step opened without them would leave the planes unrolled."""
+ ctx_lens, gen_starts, p, _ = CASES["decode"]
+ b = _Batch(19, ctx_lens, gen_starts, p, _slots(len(gen_starts)))
+ with pytest.raises(AssertionError, match="index planes are not bound"):
+ SpeculativeRollback().begin(
+ num_contexts=0,
+ prefix_tokens=0,
+ num_gens=b.n,
+ rows_per_gen=p,
+ slots=b.slots,
+ request_indices=b.req_idx,
+ logical_positions=b.logical,
+ block_table=b.block_table,
+ tokens_per_block=TPB,
+ )
+
+
@pytest.mark.parametrize("mode", MODES)
def test_commit_is_sync_free(mode):
ctx_lens, gen_starts, p, gen_accepted = CASES["mixed"]
b = _Batch(16, ctx_lens, gen_starts, p, _slots(len(ctx_lens) + len(gen_starts)))
b.set_accepted(gen_accepted)
- rollback = _rollback(b, mode)
- if mode == "fused":
- # Triton compiles the launches on their first call, outside the checked region.
- _verify_step(b, rollback, True)
- b.reset()
+ rollback = _rollback(b)
+ # Triton compiles the launches on their first call, outside the checked region.
+ _verify_step(b, rollback, mode == "fused")
+ b.reset()
torch.cuda.synchronize()
torch.cuda.set_sync_debug_mode("error")
try:
@@ -368,15 +393,14 @@ def test_commit_is_sync_free(mode):
@pytest.mark.parametrize("mode", MODES)
def test_padding_duplicates_leave_the_unique_request_exact(mode):
- # Requests 2 and 3 are the same padding dummy as request 1: one slot, one
- # block-table row, written concurrently with no defined winner, so only the
- # unique request is compared.
+ # Requests 2 and 3 are the same padding dummy as request 1: one slot, one block-table
+ # row, written concurrently with no defined winner, so only the unique request is
+ # compared.
gen_starts, p = (5, 40, 40, 40), 4
slots = [3, 7, 7, 7]
b = _Batch(17, (), gen_starts, p, slots, dup_of={2: 1, 3: 1})
b.set_accepted((2, 3, 3, 3))
- _verify_step(b, _rollback(b, mode), mode == "fused")
- # The unique request's slot is exactly its accepted prefix.
+ _verify_step(b, _rollback(b), mode == "fused")
pages = b.block_table[0].long()
expected = _reference(_single(b, 0))
assert torch.equal(b.ngram[3], expected[0][3])
@@ -409,8 +433,8 @@ def _single(b, s):
return one
-#: Per case, the acceptance counts one captured graph is replayed with (the accepted count
-#: of every generation request, bonus token included).
+#: Per case, the acceptance counts one captured graph is replayed with (the accepted count of
+#: every generation request, bonus token included).
_REPLAYS = {
"decode_p8": [
(1,) * 8,
@@ -424,11 +448,11 @@ def _single(b, s):
@pytest.mark.parametrize("mode", MODES)
@pytest.mark.parametrize("case", sorted(_REPLAYS))
def test_capture_and_replay_follow_the_accepted_counts(case, mode):
- """The whole staged step and commit captured once, replayed with the
- accepted counts and the pre-step state rewritten in place."""
+ """The whole staged step and commit captured once, replayed with the accepted counts and
+ the pre-step state rewritten in place."""
ctx_lens, gen_starts, p, _ = CASES[case]
b = _Batch(18, ctx_lens, gen_starts, p, _slots(len(gen_starts)))
- rollback = _rollback(b, mode)
+ rollback = _rollback(b)
fused = mode == "fused"
b.set_accepted((p,) * len(gen_starts))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_mtp_draft_hc.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_mtp_draft_hc.py
new file mode 100644
index 000000000000..79aa96a99447
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_mtp_draft_hc.py
@@ -0,0 +1,382 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The MTP draft layer: its hyper-connection sites on the trunk's fused entries, against the
+HuggingFace math in float64 and against the composition they replace, and its token embedding
+under rejection sampling.
+
+`_mtp_front`, `_mtp_tail` and `MTPLayer.shared_head` run over a stand-in core that holds the
+draft layer's hyper-connection weights (synthetic, at the checkpoint's element statistics): the
+attention-side site with the front end's `x_h + x_e` as its injection at unit weight, the
+MoE-side site after the output projection, and the head's stack-final site. `_mtp_tail`'s MoE
+is the identity on its mixed rows and its TP sum the identity: the sites around them are under
+test, at the rows a draft step runs -- 1 (the later draft steps and every shared head), 4 (the
+first draft step of an MTP-3 verify of one request), 32 and 1024 (larger batches -- an
+attention-DP rank's decode batch -- and a prefill).
+
+The float64 reference rounds to bf16 where the runtime's activations are bf16 between modules
+(the norm outputs, the projections, the residual) and once at each mixer output; the mixers of
+the fused path keep their logits in fp32 (the trunk's numerics). Each output must be within
+the band of that reference that the replaced composition (per-stream `flashinfer_gemma_rmsnorm`
+and `qwen4_exp_gated_residual`) reaches, widened by half.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import (
+ cublas_mm as mtp_cublas_mm,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_gated_residual import (
+ qwen4_exp_gated_residual,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_inject_norm import (
+ qwen4_exp_hc_inject_norm,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_post_inject import (
+ qwen4_exp_post_inject,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.norm.flashinfer_gemma_rmsnorm import (
+ flashinfer_gemma_rmsnorm,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501
+ modeling as MOD,
+)
+from tensorrt_llm._torch.speculative.interface import rejection_sampling_one_model
+
+pytestmark = pytest.mark.skipif(
+ not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3),
+ reason="the fused hyper-connection entries are certified on sm_103",
+)
+
+DEV = "cuda"
+BF16 = torch.bfloat16
+HC = 4
+HIDDEN = 2560
+WIDTH = HC * HIDDEN
+RANK = 320
+EPS = 1e-6
+#: A TP / TEP rank's gated attention output: 6 of the 24 query heads of 256.
+Q_WIDTH = 1536
+VOCAB = 4096
+P = "l0_mtp_"
+
+# Element statistics of the checkpoint's hyper-connection tensors (the gated_residual
+# entry's test): the norm weight's tail puts the normed stream at ~+-30.
+DOWN_RMS, UP_RMS, INJECT_RMS, GAMMA_RMS, GAMMA_ABSMAX = 0.0296, 0.0681, 0.0107, 0.477, 6.625
+
+
+class _Weights(dict):
+ """The draft layer's weights; the operands of its attention and MoE, which no site here
+ reads, are placeholders."""
+
+ def __missing__(self, key):
+ return torch.zeros(1, 1, dtype=BF16, device=DEV)
+
+
+def _raw(seed: int) -> dict:
+ g = torch.Generator(device=DEV).manual_seed(seed)
+
+ def normal(*shape, scale):
+ return (torch.randn(*shape, generator=g, device=DEV) * scale).to(BF16)
+
+ def gamma(n):
+ t = normal(n, scale=GAMMA_RMS)
+ t[0], t[n // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX
+ return t
+
+ w = {
+ "embed": normal(VOCAB, HIDDEN, scale=1.0),
+ "fhc_down": normal(RANK, WIDTH, scale=DOWN_RMS),
+ f"{P}pre_e_norm": gamma(HIDDEN),
+ f"{P}pre_h_norm": gamma(WIDTH),
+ f"{P}fc_e": normal(HIDDEN, HIDDEN, scale=HIDDEN**-0.5),
+ f"{P}fc_h": normal(HIDDEN, HIDDEN, scale=HIDDEN**-0.5),
+ f"{P}o_w": normal(HIDDEN, Q_WIDTH, scale=Q_WIDTH**-0.5),
+ }
+ for tag in ("ahc", "mhc", "fhc"):
+ w[f"{P}{tag}_norm"] = gamma(WIDTH)
+ w[f"{P}{tag}_down"] = normal(RANK, WIDTH, scale=DOWN_RMS)
+ w[f"{P}{tag}_up"] = normal(WIDTH, RANK, scale=UP_RMS)
+ if tag != "fhc":
+ w[f"{P}{tag}_inject"] = normal(HC, WIDTH, scale=INJECT_RMS)
+ return w
+
+
+def _core(raw: dict):
+ """A stand-in core over copies of `raw`, loaded as `post_load_weights` loads the draft
+ layer: the first projections fused, then the views."""
+ w = _Weights({k: v.clone() for k, v in raw.items()})
+ core = SimpleNamespace(
+ w=w,
+ num_layers=0,
+ mtp_layer_idx=0,
+ mtp_enabled=True,
+ hidden=HIDDEN,
+ hc_count=HC,
+ stream=WIDTH,
+ eps=EPS,
+ hc_lowrank=RANK,
+ hc_first_cols=-(-(RANK + HC) // 32) * 32,
+ max_num_tokens=8192,
+ # The shared expert's operands are placeholders (`_Weights`).
+ shared_inter=1,
+ tp_sum=lambda t: t,
+ _hc_first=None,
+ )
+ MOD.Qwen4ExpDefaultModelingV2._fuse_hc_first_projections(core)
+ core._mtp = MOD._mtp_views(core)
+ return core, core._mtp
+
+
+@pytest.fixture
+def moe_identity(monkeypatch):
+ monkeypatch.setattr(MOD, "_mtp_moe_block", lambda core, M, x, pad_rows=0: x)
+
+
+# ------------------------------------------------------------------- references ---
+
+
+def _bf(x: torch.Tensor) -> torch.Tensor:
+ return x.to(BF16)
+
+
+def _norm64(x: torch.Tensor, gamma: torch.Tensor, groups: int) -> torch.Tensor:
+ """Gemma RMS norm of each of `groups` equal column groups, float64, rounded once."""
+ xs = x.double().unflatten(-1, (groups, -1))
+ xs = xs * torch.rsqrt(xs.pow(2).mean(-1, keepdim=True) + EPS)
+ return _bf(xs.flatten(-2) * (1.0 + gamma.double()))
+
+
+def _mm64(a: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
+ """`a @ w^T` for an `nn.Linear` weight, float64, rounded once."""
+ return _bf(a.double() @ w.double().t())
+
+
+def _mixer64(raw, tag, normed, inject=True):
+ return qwen4_exp_gated_residual.reference(
+ normed,
+ raw[f"{P}{tag}_down"],
+ raw[f"{P}{tag}_up"],
+ raw[f"{P}{tag}_inject"] if inject else None,
+ HC,
+ )
+
+
+def _front64(raw, ids, hs):
+ rows = hs.shape[0]
+ x_e = _mm64(_norm64(raw["embed"][ids], raw[f"{P}pre_e_norm"], 1), raw[f"{P}fc_e"])
+ h_n = _norm64(hs, raw[f"{P}pre_h_norm"], HC)
+ x_h = _mm64(h_n.reshape(rows * HC, HIDDEN), raw[f"{P}fc_h"]).reshape(rows, WIDTH)
+ x = _bf(x_h.double() + x_e.double().repeat(1, HC))
+ return (x, *_mixer64(raw, "ahc", _norm64(x, raw[f"{P}ahc_norm"], HC)))
+
+
+def _inject64(x, block, inject):
+ gated = block.double().unsqueeze(-2) * inject.double().unsqueeze(-1)
+ return _bf((x.double().unflatten(-1, (HC, HIDDEN)) + gated).flatten(-2))
+
+
+def _tail64(raw, x, attn, inject):
+ x = _inject64(x, _mm64(attn, raw[f"{P}o_w"]), inject)
+ mixed, inject = _mixer64(raw, "mhc", _norm64(x, raw[f"{P}mhc_norm"], HC))
+ return _inject64(x, mixed, inject)
+
+
+def _head64(raw, hs):
+ return _mixer64(raw, "fhc", _norm64(hs, raw[f"{P}fhc_norm"], HC), inject=False)[0]
+
+
+# The composition the fused sites replace: one `flashinfer_gemma_rmsnorm` a stream, then
+# `qwen4_exp_gated_residual` (bf16 GEMM outputs), and the separate accumulate.
+
+
+def _grouped(x, gamma):
+ parts = torch.split(x, [HIDDEN] * HC, 1)
+ return torch.cat(
+ [flashinfer_gemma_rmsnorm(p, g, EPS) for g, p in zip(torch.split(gamma, HIDDEN), parts)],
+ 1,
+ )
+
+
+def _mixer_old(raw, tag, normed, inject=True):
+ return qwen4_exp_gated_residual(
+ normed,
+ raw[f"{P}{tag}_down"],
+ raw[f"{P}{tag}_up"],
+ raw[f"{P}{tag}_inject"] if inject else None,
+ HC,
+ )
+
+
+def _front_old(raw, ids, hs):
+ rows = hs.shape[0]
+ e = torch.nn.functional.embedding(ids, raw["embed"])
+ x_e = mtp_cublas_mm(
+ flashinfer_gemma_rmsnorm(e, raw[f"{P}pre_e_norm"], EPS), raw[f"{P}fc_e"].t()
+ )
+ h_n = _grouped(hs, raw[f"{P}pre_h_norm"])
+ x_h = mtp_cublas_mm(h_n.reshape(rows * HC, HIDDEN), raw[f"{P}fc_h"].t()).reshape(rows, WIDTH)
+ x = torch.add(x_h, torch.cat([x_e] * HC, -1))
+ return (x, *_mixer_old(raw, "ahc", _grouped(x, raw[f"{P}ahc_norm"])))
+
+
+def _tail_old(raw, x, attn, inject):
+ x = qwen4_exp_post_inject(x, mtp_cublas_mm(attn, raw[f"{P}o_w"].t()), inject)
+ mixed, inject = _mixer_old(raw, "mhc", _grouped(x, raw[f"{P}mhc_norm"]))
+ return qwen4_exp_post_inject(x, mixed, inject)
+
+
+def _head_old(raw, hs):
+ return _mixer_old(raw, "fhc", _grouped(hs, raw[f"{P}fhc_norm"]), inject=False)[0]
+
+
+def _err(got: torch.Tensor, ref: torch.Tensor) -> float:
+ """Worst element's distance in units of `eps * max|ref|` (bf16 eps)."""
+ ref64 = ref.double()
+ unit = torch.finfo(BF16).eps * ref64.abs().max().item()
+ return (got.double() - ref64).abs().max().item() / unit
+
+
+def _rel_l2(got: torch.Tensor, ref: torch.Tensor) -> float:
+ ref64 = ref.double()
+ return ((got.double() - ref64).norm() / ref64.norm()).item()
+
+
+def _within(name, got, ref, old) -> None:
+ """`got` as close to `ref` as the replaced composition `old` is, widened by half (and
+ never tighter than one eps of max|ref| / 1e-3 relative L2, both below the noise of one
+ bf16 rounding of a mixer output)."""
+ assert got.shape == ref.shape and got.dtype == ref.dtype, (name, got.shape, ref.shape)
+ assert bool(torch.isfinite(got).all()), name
+ e_new, e_old = _err(got, ref), _err(old, ref)
+ l_new, l_old = _rel_l2(got, ref), _rel_l2(old, ref)
+ print(
+ f"{name}: max {e_new:.3f} (old {e_old:.3f}) eps*max|ref|, relL2 {l_new:.3g} (old {l_old:.3g})"
+ )
+ assert e_new <= max(1.5 * e_old, 1.0), (name, e_new, e_old)
+ assert l_new <= max(1.5 * l_old, 1e-3), (name, l_new, l_old)
+
+
+# ------------------------------------------------------------------------ tests ---
+
+
+def _inputs(rows: int, seed: int, strided: bool = False):
+ """Token ids, the trunk's `[rows, stream]` residual (with `strided`, the column view of
+ the trunk's packed `[final | hc]` output the first draft step reads), the attention
+ output and an attention-site injection in the checkpoint's range."""
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ ids = torch.randint(0, VOCAB, (rows,), generator=g, device=DEV)
+ if strided:
+ packed = (torch.randn(rows, HIDDEN + WIDTH, generator=g, device=DEV) * 3).to(BF16)
+ hs = packed[:, HIDDEN:]
+ else:
+ hs = (torch.randn(rows, WIDTH, generator=g, device=DEV) * 3).to(BF16)
+ attn = torch.randn(rows, Q_WIDTH, generator=g, device=DEV).to(BF16)
+ inject = (0.43 + 1.15 * torch.rand(rows, HC, generator=g, device=DEV)).to(BF16)
+ return ids, hs, attn, inject
+
+
+ROWS = (1, 4, 32, 1024)
+
+
+@pytest.mark.parametrize(
+ "rows,strided", [(r, False) for r in ROWS] + [(4, True)], ids=lambda v: str(v)
+)
+def test_the_front_site(rows, strided) -> None:
+ raw = _raw(11)
+ core, M = _core(raw)
+ ids, hs, _, _ = _inputs(rows, 101 + rows, strided)
+ x, mixed, inject = MOD._mtp_front(core, M, raw["embed"], ids, hs)
+ ref, old = _front64(raw, ids, hs), _front_old(raw, ids, hs)
+ for name, got, r, o in zip(("x", "mixed", "inject"), (x, mixed, inject), ref, old):
+ _within(f"front[{rows}].{name}", got, r, o)
+ # The unit injection is the bf16 add: the residual is bitwise the replaced one's, whose
+ # projections are the same calls.
+ _, h_n = qwen4_exp_hc_inject_norm(hs, None, None, M.pre_h_norm, EPS, HC)
+ x_h = mtp_cublas_mm(h_n.reshape(rows * HC, HIDDEN), M.fc_h).reshape(rows, WIDTH)
+ e = torch.nn.functional.embedding(ids, raw["embed"])
+ x_e = mtp_cublas_mm(flashinfer_gemma_rmsnorm(e, M.pre_e_norm, EPS), M.fc_e)
+ assert torch.equal(x, torch.add(x_h, torch.cat([x_e] * HC, -1)))
+
+
+@pytest.mark.parametrize("rows", ROWS)
+def test_the_tail_sites(rows, moe_identity) -> None:
+ raw = _raw(12)
+ core, M = _core(raw)
+ _, x, attn, inject = _inputs(rows, 201 + rows)
+ out = MOD._mtp_tail(core, M, x, attn, inject, 0)
+ _within(f"tail[{rows}]", out, _tail64(raw, x, attn, inject), _tail_old(raw, x, attn, inject))
+
+
+# 32: a shared head over one attention-DP rank's decode batch (ht_dep4 drafts under MTP).
+@pytest.mark.parametrize("rows", (1, 4, 32))
+def test_the_head_site(rows) -> None:
+ raw = _raw(13)
+ core, _ = _core(raw)
+ _, hs, _, _ = _inputs(rows, 301 + rows)
+ head = SimpleNamespace(
+ core=core, logits_processor=SimpleNamespace(forward=lambda final, *args: final)
+ )
+ final = MOD.MTPLayer.shared_head(head, hs, None, None, False)
+ _within(f"head[{rows}]", final, _head64(raw, hs), _head_old(raw, hs))
+
+
+def test_a_captured_draft_step_replays_the_eager_bits(moe_identity) -> None:
+ """Front, tail and head of a 4-row draft step captured in one CUDA graph: the replay is
+ bitwise the eager result."""
+ raw = _raw(14)
+ core, M = _core(raw)
+ ids, hs, attn, _ = _inputs(4, 401, strided=True)
+ head = SimpleNamespace(
+ core=core, logits_processor=SimpleNamespace(forward=lambda final, *args: final)
+ )
+
+ def step():
+ x, _, inject = MOD._mtp_front(core, M, raw["embed"], ids, hs)
+ out = MOD._mtp_tail(core, M, x, attn, inject, 0)
+ return out, MOD.MTPLayer.shared_head(head, out, None, None, False)
+
+ eager = step()
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ step()
+ torch.cuda.current_stream().wait_stream(side)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ captured = step()
+ graph.replay()
+ torch.cuda.synchronize()
+ for got, want in zip(captured, eager):
+ assert torch.equal(got, want)
+
+
+def test_draft_embedding_reads_row_0_for_rejected_slots() -> None:
+ """Rejection sampling pads the slots past a request's accepted count with -1 (here every
+ draft is rejected at slot 0); the front end embeds those as row 0 and every valid id as
+ the plain lookup does."""
+ batch, draft, vocab, hidden = 8, 3, 4096, 256
+ draft_tokens = torch.arange(batch * draft, dtype=torch.int32, device="cuda").reshape(
+ batch, draft
+ )
+ draft_probs = torch.zeros(batch, draft, vocab, device="cuda")
+ draft_probs.scatter_(2, draft_tokens.long().unsqueeze(-1), 1.0)
+ target_probs = torch.zeros(batch, draft + 1, vocab, device="cuda")
+ target_probs[..., vocab - 1] = 1.0
+ seed = torch.zeros(batch, dtype=torch.int64, device="cuda")
+ offset = torch.arange(batch, dtype=torch.int64, device="cuda")
+ accepted, _ = rejection_sampling_one_model(
+ draft_probs, draft_tokens, target_probs, seed, offset
+ )
+ assert bool((accepted[:, 1:] == -1).all()) and bool((accepted[:, 0] == vocab - 1).all())
+
+ g = torch.Generator(device="cuda").manual_seed(0)
+ table = torch.randn(vocab, hidden, generator=g, device="cuda").to(torch.bfloat16)
+ valid_ids = torch.randint(0, vocab, (257,), generator=g, device="cuda", dtype=torch.int32)
+ for ids in (accepted.flatten(), valid_ids, valid_ids.long()):
+ e = MOD._draft_embedding(ids, table)
+ valid = ids >= 0
+ assert torch.equal(e[valid], table[ids[valid].long()])
+ assert torch.equal(e[~valid], table[0].expand(int((~valid).sum()), -1))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_pool_contract.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_pool_contract.py
similarity index 55%
rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_pool_contract.py
rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_pool_contract.py
index 2cbe7c9046bf..edf3dcd623a8 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_pool_contract.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_pool_contract.py
@@ -1,10 +1,11 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
-"""qwen4_exp's KV cache pool contract rejects the mismatches the kernels would not.
+"""The qwen4_exp target's KV cache pool contract (`pool.py`) rejects the mismatches the
+kernels would not.
The checks read only `.shape` / `.dtype`, so these run on meta tensors with no GPU and
no engine. Each case is one wrong pool that the kernel consuming it would accept and
-mis-address.
+mis-address. The geometry is one rank's at tp1, which holds every head.
"""
from __future__ import annotations
@@ -14,27 +15,25 @@
import pytest
import torch
-__extra_import_path__ = ["../../.."] # noqa: F841 -- repo's file-scoped import hook
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp import routing
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501
- modeling as M,
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501
+ pool as P,
)
-from tensorrt_llm._torch.modules.fla.cache_manager import GDNIntermediateState, GDNReplayState
-from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import Qwen4ExpHybridCacheManager
pytestmark = pytest.mark.cpu_only
-G = M._PoolGeometry(
+G = P.PoolGeometry(
tokens_per_block=32,
heads_kv=2,
+ kv_head_start=0,
head_dim=256,
index_kv_heads=1,
index_head_dim=128,
conv_dim=8192,
conv_kernel=4,
num_k_heads=16,
+ k_head_start=0,
num_v_heads=48,
+ v_head_start=0,
head_v_dim=128,
head_k_dim=128,
stream=10240,
@@ -74,103 +73,42 @@ def _history(slots=SLOTS, history=HISTORY):
def test_valid_pools_pass() -> None:
- M.check_position_plane(G, _t(9, 32, 3, dtype=torch.int32))
- M.check_attention_pools(
+ P.check_position_plane(G, _t(9, 32, 3, dtype=torch.int32))
+ P.check_attention_pools(
G, 3, _t(9, 32, 1, 128), _t(9, 2, G.heads_kv, G.tokens_per_block, G.head_dim)
)
- M.check_attention_pools(
+ P.check_attention_pools(
G_FP8,
3,
_t(9, 32, 1, 128),
_t(9, 2, G.heads_kv, G.tokens_per_block, G.head_dim, dtype=torch.float8_e4m3fn),
)
nvfp4_kv = _t(9, 2, G.heads_kv, G.tokens_per_block, 128, dtype=torch.uint8)
- M.check_attention_pools(G_NVFP4, 3, _t(9, 32, 1, 128), nvfp4_kv)
+ P.check_attention_pools(G_NVFP4, 3, _t(9, 32, 1, 128), nvfp4_kv)
for dtype in (torch.float32, torch.bfloat16):
- M.check_recurrent_pools(G, _conv(), _ssm(dtype))
- M.check_ple_pools(
+ P.check_recurrent_pools(G, _conv(), _ssm(dtype))
+ P.check_ple_pools(
G,
_t(SLOTS, G.stream, G.ple_state_len),
_t(SLOTS, G.ngram_context_len, dtype=torch.int64),
_conv(),
)
- M.check_step_buffers(
+ P.check_step_buffers(
G, _t(SLOTS, 2, G.conv_dim, G.conv_kernel - 1), _conv(), _t(SLOTS, dtype=torch.int32)
)
for dtype in (torch.float32, torch.bfloat16):
- M.check_legacy_state(
- G, _t(SLOTS, 2, G.num_v_heads, G.head_v_dim, G.head_k_dim, dtype=dtype), _ssm(dtype)
- )
- for dtype in (torch.float32, torch.bfloat16):
- for width in (2, 5, M._MAX_VERIFY_WIDTH):
- M.check_replay_state(
+ for width in (2, 5, P.MAX_VERIFY_WIDTH):
+ P.check_replay_state(
HISTORY,
width,
_t(SLOTS, dtype=torch.int32),
_t(SLOTS, dtype=torch.int32),
_ssm(dtype),
)
- M.check_replay_history(G, 3, HISTORY, SLOTS, *_history())
- for rows in range(1, M._MAX_VERIFY_WIDTH + 1):
- M.check_verify_width(rows)
- M.check_shared_selection(_t(2, 8, 512, dtype=torch.int32), {3: 1}, 3, 512)
-
-
-@pytest.mark.parametrize("replay", [False, True])
-def test_bind_step_buffers_uses_manager_layer_views(replay: bool) -> None:
- state = GDNReplayState(tokens_per_gen_step=2) if replay else GDNIntermediateState()
- state.intermediate_conv = _t(2, SLOTS, 2, G.conv_dim, G.conv_kernel - 1)
- if replay:
- state.prev_num_accepted_tokens = _t(SLOTS, dtype=torch.int32)
- state.cache_buf_idx = _t(SLOTS, dtype=torch.int32)
- state.old_x = _t(2, SLOTS, 2, HISTORY, G.num_v_heads, G.head_v_dim)
- state.old_B = _t(2, SLOTS, 2, HISTORY, G.num_k_heads, G.head_k_dim)
- state.old_dt = _t(2, SLOTS, 2, G.num_v_heads, HISTORY, dtype=torch.float32)
- state.old_dA_cumsum = torch.empty_like(state.old_dt)
- state._state_strides = (2, 0, 1)
- else:
- state.intermediate_ssm = _t(
- 2, SLOTS, 2, G.num_v_heads, G.head_v_dim, G.head_k_dim, dtype=torch.float32
- )
- state.intermediate_indices = _t(SLOTS, dtype=torch.int32)
- manager = Qwen4ExpHybridCacheManager.__new__(Qwen4ExpHybridCacheManager)
- manager.spec_config = SimpleNamespace(tokens_per_gen_step=2)
- manager._speculative_state = state
- manager.mamba_layer_offsets = {3: 1, 7: 0}
- manager.all_conv_states = [_conv(), _conv()]
- manager.all_ssm_states = [_ssm(), _ssm()]
-
- core = M.Qwen4ExpTp1ModelingV2.__new__(M.Qwen4ExpTp1ModelingV2)
- torch.nn.Module.__init__(core)
- core.gdn_layers = [3, 7]
- metadata = SimpleNamespace(
- replay_work_items=_t(SLOTS, 4, dtype=torch.int32),
- replay_n_writes=_t(1, dtype=torch.int32),
- )
- bound = core._bind_step_buffers(manager, G, metadata, _conv(), _ssm())
-
- assert bound["int_idx"] is state.intermediate_indices
- for layer, offset in manager.mamba_layer_offsets.items():
- views = [(bound["int_conv"][layer], state.intermediate_conv)]
- if replay:
- assert "int_ssm" not in bound
- replay_bound = bound["replay"]
- assert replay_bound.all_layer_commit
- assert replay_bound.cache_buf_idx is state.cache_buf_idx
- assert replay_bound.prev_num_accepted is state.prev_num_accepted_tokens
- views.extend(
- zip(
- replay_bound.old[layer],
- (state.old_x, state.old_B, state.old_dt, state.old_dA_cumsum),
- )
- )
- else:
- assert "replay" not in bound
- views.append((bound["int_ssm"][layer], state.intermediate_ssm))
- for view, storage in views:
- assert view._base is storage
- assert view.storage_offset() == offset * storage.stride(0)
- assert view.shape == storage.shape[1:]
+ P.check_replay_history(G, 3, HISTORY, SLOTS, *_history())
+ for rows in range(1, P.MAX_VERIFY_WIDTH + 1):
+ P.check_verify_width(rows)
+ P.check_shared_selection(_t(2, 8, 512, dtype=torch.int32), {3: 1}, 3, 512)
@pytest.mark.parametrize(
@@ -184,23 +122,23 @@ def test_bind_step_buffers_uses_manager_layer_views(replay: bool) -> None:
)
def test_position_plane_rejects(plane, match) -> None:
with pytest.raises(AssertionError, match=match):
- M.check_position_plane(G, plane)
+ P.check_position_plane(G, plane)
def test_attention_pools_reject_wrong_index_plane_and_kv_layout() -> None:
kv = _t(9, 2, G.heads_kv, G.tokens_per_block, G.head_dim)
with pytest.raises(AssertionError, match="QSA index plane"):
- M.check_attention_pools(G, 3, _t(9, 32, 1, 64), kv)
+ P.check_attention_pools(G, 3, _t(9, 32, 1, 64), kv)
with pytest.raises(AssertionError, match="QSA index plane"):
- M.check_attention_pools(G, 3, _t(9, 32, 1, 128, dtype=torch.float16), kv)
+ P.check_attention_pools(G, 3, _t(9, 32, 1, 128, dtype=torch.float16), kv)
with pytest.raises(AssertionError, match="paged K/V"):
- M.check_attention_pools(G, 3, _t(9, 32, 1, 128), _t(9, 2, 4, 32, 256))
+ P.check_attention_pools(G, 3, _t(9, 32, 1, 128), _t(9, 2, 4, 32, 256))
# NHD instead of HND: heads and block swapped.
with pytest.raises(AssertionError, match="paged K/V"):
- M.check_attention_pools(G, 3, _t(9, 32, 1, 128), _t(9, 2, 32, 2, 256))
+ P.check_attention_pools(G, 3, _t(9, 32, 1, 128), _t(9, 2, 32, 2, 256))
# K and V not adjacent role planes.
with pytest.raises(AssertionError, match="role planes"):
- M.check_attention_pools(G, 3, _t(9, 32, 1, 128), _t(9, 3, 2, 32, 256))
+ P.check_attention_pools(G, 3, _t(9, 32, 1, 128), _t(9, 3, 2, 32, 256))
def test_attention_pools_reject_the_other_storage_dtype() -> None:
@@ -210,21 +148,21 @@ def test_attention_pools_reject_the_other_storage_dtype() -> None:
bf16_kv = _t(9, 2, G.heads_kv, G.tokens_per_block, G.head_dim)
fp8_kv = bf16_kv.to(torch.float8_e4m3fn)
with pytest.raises(AssertionError, match="paged K/V"):
- M.check_attention_pools(G, 3, index, fp8_kv)
+ P.check_attention_pools(G, 3, index, fp8_kv)
with pytest.raises(AssertionError, match="paged K/V"):
- M.check_attention_pools(G, 3, index, bf16_kv.to(torch.float16))
+ P.check_attention_pools(G, 3, index, bf16_kv.to(torch.float16))
with pytest.raises(AssertionError, match="paged K/V"):
- M.check_attention_pools(G_FP8, 3, index, bf16_kv)
+ P.check_attention_pools(G_FP8, 3, index, bf16_kv)
# The index plane is a kernel contract, not KV storage: bf16 under either algorithm.
with pytest.raises(AssertionError, match="QSA index plane"):
- M.check_attention_pools(G_FP8, 3, index.to(torch.float8_e4m3fn), fp8_kv)
+ P.check_attention_pools(G_FP8, 3, index.to(torch.float8_e4m3fn), fp8_kv)
# An NVFP4 pool is half a head wide: a full-width uint8 pool is the other storage.
with pytest.raises(AssertionError, match="paged K/V"):
- M.check_attention_pools(
+ P.check_attention_pools(
G_NVFP4, 3, index, _t(9, 2, G.heads_kv, G.tokens_per_block, 256, dtype=torch.uint8)
)
with pytest.raises(AssertionError, match="paged K/V"):
- M.check_attention_pools(
+ P.check_attention_pools(
G, 3, index, _t(9, 2, G.heads_kv, G.tokens_per_block, 128, dtype=torch.uint8)
)
@@ -244,128 +182,111 @@ def test_kv_scale_planes_must_hold_the_data_pools_pages_per_slot() -> None:
against 24576 of scales."""
data = _strided((9, 2, G.heads_kv, G.tokens_per_block, 128), 24 * 8192)
scales = _strided((9, 2, G.heads_kv, G.tokens_per_block, 16), 24 * 1024)
- M.check_kv_scale_planes(G_NVFP4, 3, data, scales)
+ P.check_kv_scale_planes(G_NVFP4, 3, data, scales)
# A side page coalesced into the data pool: 36 pages per slot against 24.
coalesced = _strided((9, 2, G.heads_kv, G.tokens_per_block, 128), 36 * 8192)
with pytest.raises(AssertionError, match="36 pages per slot.*24"):
- M.check_kv_scale_planes(G_NVFP4, 3, coalesced, scales)
+ P.check_kv_scale_planes(G_NVFP4, 3, coalesced, scales)
with pytest.raises(AssertionError, match="whole number"):
- M.check_kv_scale_planes(G_NVFP4, 3, data, _strided((9, 2, 2, 32, 16), 24 * 1024 + 1))
+ P.check_kv_scale_planes(G_NVFP4, 3, data, _strided((9, 2, 2, 32, 16), 24 * 1024 + 1))
def test_kv_scale_planes_reject_the_wrong_block_scale_width() -> None:
data = _t(9, 2, G.heads_kv, G.tokens_per_block, 128, dtype=torch.uint8)
with pytest.raises(AssertionError, match="NVFP4 block scales"):
- M.check_kv_scale_planes(G_NVFP4, 3, data, _t(9, 2, 2, 32, 8, dtype=torch.uint8))
+ P.check_kv_scale_planes(G_NVFP4, 3, data, _t(9, 2, 2, 32, 8, dtype=torch.uint8))
def test_recurrent_pools_reject_wrong_key_dim_order_and_dtype() -> None:
with pytest.raises(AssertionError, match="recurrent state pool"):
- M.check_recurrent_pools(
+ P.check_recurrent_pools(
G, _conv(), _t(SLOTS, G.num_v_heads, G.head_k_dim, G.head_v_dim + 1)
)
with pytest.raises(AssertionError, match="conv state pool"):
- M.check_recurrent_pools(G, _conv(dtype=torch.float32), _ssm())
+ P.check_recurrent_pools(G, _conv(dtype=torch.float32), _ssm())
# The engine accepts an fp16 pool; no recurrence entry the target binds is certified at it.
with pytest.raises(AssertionError, match="recurrent state pool"):
- M.check_recurrent_pools(G, _conv(), _ssm(torch.float16))
+ P.check_recurrent_pools(G, _conv(), _ssm(torch.float16))
def test_ple_pools_reject_slot_mismatch_and_int32_history() -> None:
ple_conv = _t(SLOTS, G.stream, G.ple_state_len)
ngram = _t(SLOTS, G.ngram_context_len, dtype=torch.int64)
with pytest.raises(AssertionError, match="one slot space"):
- M.check_ple_pools(
+ P.check_ple_pools(
G, ple_conv, _t(SLOTS + 1, G.ngram_context_len, dtype=torch.int64), _conv()
)
with pytest.raises(AssertionError, match="n-gram context pool"):
- M.check_ple_pools(G, ple_conv, _t(SLOTS, G.ngram_context_len, dtype=torch.int32), _conv())
+ P.check_ple_pools(G, ple_conv, _t(SLOTS, G.ngram_context_len, dtype=torch.int32), _conv())
with pytest.raises(AssertionError, match="did not allocate"):
- M.check_ple_pools(G, None, ngram, _conv())
+ P.check_ple_pools(G, None, ngram, _conv())
def test_step_buffers_reject_missing_conv_rank_dtype_and_index() -> None:
int_conv = _t(SLOTS, 2, G.conv_dim, G.conv_kernel - 1)
idx = _t(SLOTS, dtype=torch.int32)
with pytest.raises(AssertionError, match="no per-step conv"):
- M.check_step_buffers(G, None, _conv(), idx)
+ P.check_step_buffers(G, None, _conv(), idx)
# Rank without the steps dim.
with pytest.raises(AssertionError, match="per-step conv buffer"):
- M.check_step_buffers(G, _conv(), _conv(), idx)
+ P.check_step_buffers(G, _conv(), _conv(), idx)
# The promotion kernel copies without converting, so the dtype must match the pool.
with pytest.raises(AssertionError, match="per-step conv buffer"):
- M.check_step_buffers(G, int_conv.to(torch.float32), _conv(), idx)
+ P.check_step_buffers(G, int_conv.to(torch.float32), _conv(), idx)
with pytest.raises(AssertionError, match="intermediate_state_indices"):
- M.check_step_buffers(G, int_conv, _conv(), idx.to(torch.int64))
+ P.check_step_buffers(G, int_conv, _conv(), idx.to(torch.int64))
with pytest.raises(AssertionError, match="intermediate_state_indices"):
- M.check_step_buffers(G, int_conv, _conv(), None)
-
-
-def test_legacy_state_rejects_missing_rank_dtype_drift_and_too_many_rows() -> None:
- def per_step(steps):
- return _t(SLOTS, steps, G.num_v_heads, G.head_v_dim, G.head_k_dim, dtype=torch.float32)
-
- int_ssm = per_step(2)
- with pytest.raises(AssertionError, match="no per-step recurrent"):
- M.check_legacy_state(G, None, _ssm())
- with pytest.raises(AssertionError, match="per-step recurrent buffer"):
- M.check_legacy_state(G, _ssm(), _ssm())
- with pytest.raises(AssertionError, match="per-step recurrent buffer"):
- M.check_legacy_state(G, int_ssm.to(torch.bfloat16), _ssm())
- # More rows per request than the legacy recurrence is certified at.
- M.check_legacy_state(G, per_step(M._MAX_LEGACY_VERIFY_WIDTH), _ssm())
- with pytest.raises(AssertionError, match="TRTLLM_USE_GDN_REPLAY"):
- M.check_legacy_state(G, per_step(M._MAX_LEGACY_VERIFY_WIDTH + 1), _ssm())
+ P.check_step_buffers(G, int_conv, _conv(), None)
def test_replay_state_rejects_pool_dtype_geometry_and_bookkeeping() -> None:
pnat = _t(SLOTS, dtype=torch.int32)
buf = _t(SLOTS, dtype=torch.int32)
with pytest.raises(AssertionError, match="fp32 or a bf16 checkpoint"):
- M.check_replay_state(HISTORY, 3, pnat, buf, _ssm(torch.float16))
+ P.check_replay_state(HISTORY, 3, pnat, buf, _ssm(torch.float16))
with pytest.raises(AssertionError, match="replay_history_size"):
- M.check_replay_state(12, 3, pnat, buf, _ssm())
- for width in (1, M._MAX_VERIFY_WIDTH + 1):
+ P.check_replay_state(12, 3, pnat, buf, _ssm())
+ for width in (1, P.MAX_VERIFY_WIDTH + 1):
with pytest.raises(AssertionError, match="replay_step_width"):
- M.check_replay_state(HISTORY, width, pnat, buf, _ssm())
+ P.check_replay_state(HISTORY, width, pnat, buf, _ssm())
# The bookkeeping is one int32 per pool slot, so it must cover the whole pool.
with pytest.raises(AssertionError, match="prev_num_accepted_tokens"):
- M.check_replay_state(HISTORY, 3, _t(SLOTS - 1, dtype=torch.int32), buf, _ssm())
+ P.check_replay_state(HISTORY, 3, _t(SLOTS - 1, dtype=torch.int32), buf, _ssm())
with pytest.raises(AssertionError, match="cache_buf_idx"):
- M.check_replay_state(HISTORY, 3, pnat, buf.to(torch.int64), _ssm())
+ P.check_replay_state(HISTORY, 3, pnat, buf.to(torch.int64), _ssm())
def test_verify_width_admits_one_to_eight_rows_per_request() -> None:
- for rows in (0, M._MAX_VERIFY_WIDTH + 1):
+ for rows in (0, P.MAX_VERIFY_WIDTH + 1):
with pytest.raises(AssertionError, match="predicted_tokens_per_seq"):
- M.check_verify_width(rows)
+ P.check_verify_width(rows)
def test_replay_width_admits_only_the_managers_step_width() -> None:
replay = SimpleNamespace(width=3)
- M.check_replay_width(None, 1) # the legacy recurrence takes any width
- M.check_replay_width(replay, 3)
+ P.check_replay_width(None, 1) # without the cached replay any width
+ P.check_replay_width(replay, 3)
for rows in (1, 2, 4):
with pytest.raises(AssertionError, match="step width 3"):
- M.check_replay_width(replay, rows)
+ P.check_replay_width(replay, rows)
def test_replay_history_rejects_missing_short_pool_and_wrong_geometry() -> None:
old_x, old_B, old_dt = _history()
with pytest.raises(AssertionError, match="no old_x"):
- M.check_replay_history(G, 3, HISTORY, SLOTS, None, old_B, old_dt)
+ P.check_replay_history(G, 3, HISTORY, SLOTS, None, old_B, old_dt)
with pytest.raises(AssertionError, match="slots for a"):
- M.check_replay_history(G, 3, HISTORY, SLOTS, *_history(slots=SLOTS - 1))
+ P.check_replay_history(G, 3, HISTORY, SLOTS, *_history(slots=SLOTS - 1))
# The kernel addresses the history by `history_size`, so a buffer of another depth
# would be read at the wrong offsets.
with pytest.raises(AssertionError, match="old_x"):
- M.check_replay_history(G, 3, HISTORY, SLOTS, *_history(history=12))
+ P.check_replay_history(G, 3, HISTORY, SLOTS, *_history(history=12))
with pytest.raises(AssertionError, match="old_B"):
- M.check_replay_history(G, 3, HISTORY, SLOTS, old_x, old_B.to(torch.float32), old_dt)
+ P.check_replay_history(G, 3, HISTORY, SLOTS, old_x, old_B.to(torch.float32), old_dt)
with pytest.raises(AssertionError, match="old_dt"):
- M.check_replay_history(G, 3, HISTORY, SLOTS, old_x, old_B, old_dt.to(torch.bfloat16))
+ P.check_replay_history(G, 3, HISTORY, SLOTS, old_x, old_B, old_dt.to(torch.bfloat16))
with pytest.raises(AssertionError, match="double-buffered"):
- M.check_replay_history(
+ P.check_replay_history(
G,
3,
HISTORY,
@@ -379,22 +300,10 @@ def test_replay_history_rejects_missing_short_pool_and_wrong_geometry() -> None:
def test_shared_selection_rejects_missing_row_and_width() -> None:
shared = _t(2, 8, 512, dtype=torch.int32)
with pytest.raises(AssertionError, match="no shared buffers"):
- M.check_shared_selection(None, {3: 0}, 3, 512)
+ P.check_shared_selection(None, {3: 0}, 3, 512)
with pytest.raises(AssertionError, match="no row"):
- M.check_shared_selection(shared, {4: 0}, 3, 512)
+ P.check_shared_selection(shared, {4: 0}, 3, 512)
with pytest.raises(AssertionError, match="no row"):
- M.check_shared_selection(shared, {3: 2}, 3, 512)
+ P.check_shared_selection(shared, {3: 2}, 3, 512)
with pytest.raises(AssertionError, match="blocks per row"):
- M.check_shared_selection(shared, {3: 1}, 3, 256)
-
-
-def test_routing_admits_exactly_what_the_target_implements() -> None:
- """`route` declines a MoE backend, a draft length or a recurrent-state dtype the
- constructor or the first bind would assert on. `test_modeling_v2_routing.py` pins
- routing's behavior at today's values; this keeps routing's constants and the target's
- from drifting apart."""
- assert set(routing._MOE_BACKENDS) == set(M._TRUNK_MOE_BACKENDS)
- assert routing._MAX_DRAFT_LEN == M._MAX_VERIFY_WIDTH - 1
- assert routing._MAX_LEGACY_DRAFT_LEN == M._MAX_LEGACY_VERIFY_WIDTH - 1
- assert routing._MAX_CUTEDSL_DRAFT_LEN == M._CUTEDSL_MAX_DRAFT_LEN
- assert routing._SSM_STATE_DTYPES == M._SSM_STATE_DTYPES
+ P.check_shared_selection(shared, {3: 1}, 3, 256)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_pool_layout.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_pool_layout.py
new file mode 100644
index 000000000000..e84867c58dc6
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_pool_layout.py
@@ -0,0 +1,345 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The per-rank pool layout this target binds at tp1, tp4, tep4, dep4 and dep2, against the
+stock manager.
+
+`pool.pool_geometry` derives one rank's pool shapes from the text config and the rank's `Mapping`.
+These tests hold it to `LAYOUT` (what `core` binds at each member), then build the stock
+`Qwen4ExpHybridCacheManager` the way the engine does -- `get_kv_cache_manager_cls` and
+`_create_kv_cache_manager` over a `ModelConfig`, with the core runtime's cache configuration
+(bf16 recurrent state, block reuse off, KV cache manager V2, which the engine resolves `auto` to
+for this model) -- on the first and the last rank of each member, one GPU, and take every pool
+view through `pool.take_pools`, which checks each against the geometry. Ranks 0 and 3 hold
+different KV and GDN heads at tp4 and tep4; the shapes are the same, the heads are not. dep2, the
+layout of both disaggregated-serving engines, binds dep4's geometry on both of its ranks, so the
+transceiver moves pools of one layout between identical layouts. tp1 holds every head, the
+geometry dep4 and dep2 bind on each rank.
+"""
+
+import contextlib
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501
+ pool as P,
+)
+from tensorrt_llm._torch.configs import Qwen4ExpTextConfig
+from tensorrt_llm._torch.model_config import ModelConfig
+from tensorrt_llm._torch.models.modeling_utils import duplicate_kv_weight
+from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import Qwen4ExpHybridCacheManager
+from tensorrt_llm._torch.pyexecutor._util import _create_kv_cache_manager, get_kv_cache_manager_cls
+from tensorrt_llm._torch.pyexecutor.model_loader import validate_and_set_mamba_ssm_cache_dtype
+from tensorrt_llm.bindings import DataType
+from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MTPDecodingConfig, QSASparseAttentionConfig
+from tensorrt_llm.mapping import Mapping
+from tensorrt_llm.models.modeling_utils import QuantConfig
+from tensorrt_llm.quantization.mode import QuantAlgo
+
+NUM_LAYERS = 48
+
+# Qwen3.8-Flash-Next's text config: the fields the manager and the geometry read, and what the
+# QSA lowering and the config class need beside them.
+_TEXT = dict(
+ architectures=["Qwen4ExpForCausalLM"],
+ dtype="bfloat16",
+ full_attention_interval=4,
+ hc_count=4,
+ hc_lowrank=320,
+ head_dim=256,
+ heads_per_ngram=8,
+ hidden_size=2560,
+ indexer_budget=2048,
+ indexer_compress_ratio=4,
+ indexer_head_dim=128,
+ indexer_kv_heads=1,
+ indexer_n_heads=4,
+ layer_types=["full_attention" if i % 4 == 3 else "linear_attention" for i in range(NUM_LAYERS)],
+ linear_conv_kernel_dim=4,
+ linear_key_head_dim=128,
+ linear_num_key_heads=16,
+ linear_num_value_heads=48,
+ linear_value_head_dim=128,
+ make_ngram_vocab_size_divisible_by=128,
+ mamba_ssm_dtype="float32",
+ max_position_embeddings=262144,
+ moe_intermediate_size=640,
+ mtp=dict(
+ hybrid=True,
+ layer_types=["full_attention"],
+ mtp_use_hidden_state_from_layer=None,
+ num_hidden_layers=1,
+ rope_theta=10000000,
+ ),
+ mtp_num_hidden_layers=1,
+ mtp_use_dedicated_embeddings=False,
+ ngram_size=3,
+ ngram_vocab_size_base=20000000,
+ num_attention_heads=24,
+ num_experts=512,
+ num_experts_per_tok=10,
+ num_hidden_layers=NUM_LAYERS,
+ num_key_value_heads=2,
+ partial_rotary_factor=0.25,
+ ple_conv_kernel_size=4,
+ ple_embed_dim=2560,
+ ple_layer_ids=[2],
+ rms_norm_eps=1e-06,
+ rope_parameters=dict(
+ mrope_interleaved=True,
+ mrope_section=[11, 11, 10],
+ partial_rotary_factor=0.25,
+ rope_theta=10000000,
+ rope_type="default",
+ ),
+ shared_expert_intermediate_size=640,
+ split_ngram_parts=128,
+ tie_word_embeddings=False,
+ vocab_size=248320,
+)
+
+ATTN_LAYERS = [i for i in range(NUM_LAYERS) if i % 4 == 3]
+GDN_LAYERS = [i for i in range(NUM_LAYERS) if i % 4 != 3]
+PLE_LAYER = 1 # `ple_layer_ids: [2]` is one-based
+MTP_LAYER = NUM_LAYERS
+TOKENS_PER_BLOCK = 32 # `kv_cache_config.tokens_per_block`'s default, which the runtime keeps
+
+#: Each member's mapping: one rank (tp1), attention TP4 with the experts split by intermediate
+#: (tp4) or by expert (tep4), and attention DP with expert parallelism (dep4, dep2). The expert
+#: split moves no pool.
+MEMBERS = {
+ "tp1": dict(world_size=1, tp_size=1, moe_tp_size=1, moe_ep_size=1, enable_attention_dp=False),
+ "tp4": dict(world_size=4, tp_size=4, moe_tp_size=4, moe_ep_size=1, enable_attention_dp=False),
+ "tep4": dict(world_size=4, tp_size=4, moe_tp_size=1, moe_ep_size=4, enable_attention_dp=False),
+ "dep4": dict(world_size=4, tp_size=4, moe_tp_size=1, moe_ep_size=4, enable_attention_dp=True),
+ "dep2": dict(world_size=2, tp_size=2, moe_tp_size=1, moe_ep_size=2, enable_attention_dp=True),
+}
+
+#: What `core` binds on each member and rank: KV heads and the first one's global index, GDN key
+#: heads and the first, GDN value heads and the first, and the conv channels ([q | k | v]).
+LAYOUT = {
+ ("tp1", 0): dict(heads_kv=2, kv_head_start=0, num_k_heads=16, k_head_start=0, num_v_heads=48,
+ v_head_start=0, conv_dim=10240),
+ ("tp4", 0): dict(heads_kv=1, kv_head_start=0, num_k_heads=4, k_head_start=0, num_v_heads=12,
+ v_head_start=0, conv_dim=2560),
+ ("tp4", 3): dict(heads_kv=1, kv_head_start=1, num_k_heads=4, k_head_start=12, num_v_heads=12,
+ v_head_start=36, conv_dim=2560),
+ ("tep4", 0): dict(heads_kv=1, kv_head_start=0, num_k_heads=4, k_head_start=0, num_v_heads=12,
+ v_head_start=0, conv_dim=2560),
+ ("tep4", 3): dict(heads_kv=1, kv_head_start=1, num_k_heads=4, k_head_start=12, num_v_heads=12,
+ v_head_start=36, conv_dim=2560),
+ ("dep4", 0): dict(heads_kv=2, kv_head_start=0, num_k_heads=16, k_head_start=0, num_v_heads=48,
+ v_head_start=0, conv_dim=10240),
+ ("dep4", 3): dict(heads_kv=2, kv_head_start=0, num_k_heads=16, k_head_start=0, num_v_heads=48,
+ v_head_start=0, conv_dim=10240),
+ ("dep2", 0): dict(heads_kv=2, kv_head_start=0, num_k_heads=16, k_head_start=0, num_v_heads=48,
+ v_head_start=0, conv_dim=10240),
+ ("dep2", 1): dict(heads_kv=2, kv_head_start=0, num_k_heads=16, k_head_start=0, num_v_heads=48,
+ v_head_start=0, conv_dim=10240),
+} # fmt: skip
+
+#: The same on every member and rank: the page size, the head dims, the replicated QSA index and
+#: PLE pools, and the activation dtype the conv, index and PLE pools hold.
+COMMON = dict(
+ tokens_per_block=TOKENS_PER_BLOCK,
+ head_dim=256,
+ index_kv_heads=1,
+ index_head_dim=128,
+ conv_kernel=4,
+ head_k_dim=128,
+ head_v_dim=128,
+ stream=4 * 2560,
+ ple_state_len=(4 - 1) * 3,
+ ngram_context_len=3 - 1,
+ dtype=torch.bfloat16,
+)
+
+#: The K/V pool's storage per resolved KV-cache algorithm: (pool dtype, innermost extent, the
+#: manager's binding dtype).
+KV_STORAGE = {
+ None: (torch.bfloat16, 256, DataType.BF16),
+ QuantAlgo.FP8: (torch.float8_e4m3fn, 256, DataType.FP8),
+ QuantAlgo.NVFP4: (torch.uint8, 128, DataType.NVFP4),
+}
+KV_IDS = {None: "kv_bf16", QuantAlgo.FP8: "kv_fp8", QuantAlgo.NVFP4: "kv_nvfp4"}
+
+_needs_gpu = pytest.mark.skipif(
+ not torch.cuda.is_available(), reason="the manager allocates its pools on the GPU"
+)
+
+
+def _mapping(member: str, rank: int) -> Mapping:
+ return Mapping(rank=rank, gpus_per_node=4, pp_size=1, **MEMBERS[member])
+
+
+def _geometry(member: str, rank: int, kv_cache_quant_algo=None) -> P.PoolGeometry:
+ return P.pool_geometry(
+ Qwen4ExpTextConfig(**_TEXT),
+ _mapping(member, rank),
+ tokens_per_block=TOKENS_PER_BLOCK,
+ kv_cache_quant_algo=kv_cache_quant_algo,
+ dtype=torch.bfloat16,
+ )
+
+
+# The first and the last rank of each member.
+_MEMBER_RANKS = [
+ (member, rank) for member, m in MEMBERS.items() for rank in sorted({0, m["world_size"] - 1})
+]
+_MEMBER_RANK_IDS = [f"{member}-rank{rank}" for member, rank in _MEMBER_RANKS]
+# dep2 serves without speculation (routing declines a `speculative_config` there).
+_MTP_MEMBER_RANKS = [(member, rank) for member, rank in _MEMBER_RANKS if member != "dep2"]
+_MTP_MEMBER_RANK_IDS = [f"{member}-rank{rank}" for member, rank in _MTP_MEMBER_RANKS]
+
+
+@pytest.mark.parametrize("member, rank", _MEMBER_RANKS, ids=_MEMBER_RANK_IDS)
+def test_the_geometry_is_the_layout_core_binds(member, rank):
+ g = _geometry(member, rank)
+ expected = dict(COMMON, **LAYOUT[member, rank], kv_dtype=torch.bfloat16, kv_width=256)
+ assert g._asdict() == expected
+ assert g.conv_sections == (
+ g.num_k_heads * 128,
+ g.num_k_heads * 128,
+ g.num_v_heads * 128,
+ )
+
+
+@pytest.mark.parametrize("algo", list(KV_STORAGE), ids=list(KV_IDS.values()))
+def test_the_kv_storage_follows_the_resolved_algorithm(algo):
+ g = _geometry("dep4", 0, algo)
+ assert (g.kv_dtype, g.kv_width) == KV_STORAGE[algo][:2]
+ assert g.kv_block_scales == (algo == QuantAlgo.NVFP4)
+
+
+def test_another_kv_algorithm_is_refused():
+ with pytest.raises(ValueError, match="kv_cache_quant_algo"):
+ _geometry("dep4", 0, QuantAlgo.INT8)
+
+
+@pytest.mark.parametrize("member, rank", _MEMBER_RANKS, ids=_MEMBER_RANK_IDS)
+def test_a_rank_holds_the_kv_head_the_weight_split_gives_it(member, rank):
+ """With more attention-TP ranks than KV heads the stock loader repeats each head
+ (`duplicate_kv_weight`) before the even row split, so the K/V a rank writes into its pages
+ is that of global head `kv_head_start`; under attention DP a rank holds all of them."""
+ g = _geometry(member, rank)
+ heads, rows_per_head = 2, 4
+ weight = torch.arange(heads * rows_per_head, dtype=torch.float32).reshape(-1, 1)
+ tp = P.attention_tp_size(_mapping(member, rank))
+ local = duplicate_kv_weight(weight, heads, tp).chunk(tp)[rank if tp > 1 else 0]
+ start, stop = g.kv_head_start * rows_per_head, (g.kv_head_start + g.heads_kv) * rows_per_head
+ assert torch.equal(local, weight[start:stop])
+
+
+@contextlib.contextmanager
+def _manager(member: str, rank: int, kv_cache_quant_algo=None, spec_config=None):
+ """The trunk manager the engine builds on `rank` of `member`, through the engine's own class
+ choice and factory."""
+ mapping = _mapping(member, rank)
+ model_config = ModelConfig(
+ pretrained_config=Qwen4ExpTextConfig(**_TEXT),
+ mapping=mapping,
+ quant_config=QuantConfig(
+ quant_algo=QuantAlgo.MIXED_PRECISION, kv_cache_quant_algo=kv_cache_quant_algo
+ ),
+ sparse_attention_config=QSASparseAttentionConfig(),
+ spec_config=spec_config,
+ )
+ # What the loader does with `kv_cache_config.mamba_ssm_cache_dtype: bfloat16`.
+ validate_and_set_mamba_ssm_cache_dtype(model_config, "bfloat16")
+ kv_cache_config = KvCacheConfig(
+ enable_block_reuse=False,
+ mamba_ssm_cache_dtype="bfloat16",
+ max_tokens=4096,
+ use_kv_cache_manager_v2=True,
+ )
+ cls = get_kv_cache_manager_cls(model_config, kv_cache_config)
+ cm = _create_kv_cache_manager(
+ None,
+ cls,
+ mapping,
+ kv_cache_config,
+ tokens_per_block=TOKENS_PER_BLOCK,
+ max_seq_len=2048,
+ max_batch_size=4,
+ spec_config=spec_config,
+ sparse_attention_config=model_config.sparse_attention_config,
+ max_num_tokens=2048,
+ max_beam_width=1,
+ kv_connector_manager=None,
+ model_config=model_config,
+ dtype=torch.bfloat16,
+ max_cuda_graph_batch_size=4,
+ )
+ try:
+ yield cm
+ finally:
+ cm.shutdown()
+
+
+@_needs_gpu
+@pytest.mark.parametrize("algo", list(KV_STORAGE), ids=list(KV_IDS.values()))
+@pytest.mark.parametrize("member, rank", _MEMBER_RANKS, ids=_MEMBER_RANK_IDS)
+def test_the_stock_manager_serves_the_layout(member, rank, algo):
+ """Every trunk pool view of the stock manager passes the contract at this rank's geometry,
+ and the manager's own per-rank numbers are the geometry's."""
+ g = _geometry(member, rank, algo)
+ with _manager(member, rank, algo) as cm:
+ assert type(cm) is Qwen4ExpHybridCacheManager
+ assert cm.dtype == KV_STORAGE[algo][2]
+ assert cm.ssm_state_dtype == torch.bfloat16
+ pools = P.take_pools(cm, g, ATTN_LAYERS, GDN_LAYERS, PLE_LAYER)
+
+ assert {cm.num_kv_heads_per_layer[cm.layer_offsets[i]] for i in ATTN_LAYERS} == {g.heads_kv}
+ assert (cm.qsa_index_kv_heads, cm.qsa_index_dim) == (g.index_kv_heads, g.index_head_dim)
+ assert list(cm.conv_state_shape) == [g.conv_dim, g.conv_kernel - 1]
+ assert list(cm.ssm_state_shape) == [g.num_v_heads, g.head_v_dim, g.head_k_dim]
+ # Each attention layer has a pool offset of its own.
+ assert len(set(pools.pool_idx.values())) == len(ATTN_LAYERS)
+ # One slot space for every QSA view of the layers, and one for the recurrent and PLE
+ # state.
+ slots = {view.shape[0] for view in (pools.position_cache, *pools.index_cache.values())}
+ slots |= {planes.k.shape[0] for planes in pools.kv.values()}
+ assert len(slots) == 1
+ assert {t.shape[0] for t in (*pools.conv.values(), *pools.ssm.values())} == {
+ pools.ple_state[0].shape[0]
+ }
+
+
+@_needs_gpu
+@pytest.mark.parametrize("member, rank", _MTP_MEMBER_RANKS, ids=_MTP_MEMBER_RANK_IDS)
+def test_the_mtp_verify_buffers_hold_this_ranks_heads(monkeypatch, member, rank):
+ """With MTP (three drafts, the GDN cached replay) the verify buffers -- the per-step conv
+ windows and every GDN layer's replay history -- are this rank's heads. Under attention TP the
+ draft layer has a separate draft KV cache; under attention DP it is in the trunk's manager,
+ with the trunk's per-rank pool shapes."""
+ monkeypatch.delenv("TRTLLM_USE_GDN_REPLAY", raising=False)
+ g = _geometry(member, rank)
+ spec = MTPDecodingConfig(max_draft_len=3, num_nextn_predict_layers=1)
+ with _manager(member, rank, spec_config=spec) as cm:
+ pools = P.take_pools(cm, g, ATTN_LAYERS, GDN_LAYERS, PLE_LAYER)
+ assert cm.is_speculative()
+ ssm0 = pools.ssm[GDN_LAYERS[0]]
+ caches = {i: cm.mamba_layer_cache(i) for i in GDN_LAYERS}
+ for i, cache in caches.items():
+ P.check_step_buffers(
+ g, cache.intermediate_conv_window, pools.conv[i], cm.intermediate_state_indices
+ )
+ assert cache.intermediate_ssm is None, "the cached replay keeps no per-step state"
+ meta = cm.get_replay_state_update_metadata()
+ assert meta is not None
+ P.check_replay_state(
+ meta.replay_history_size,
+ meta.replay_step_width,
+ meta.prev_num_accepted_tokens,
+ meta.cache_buf_idx,
+ ssm0,
+ )
+ assert meta.replay_step_width == spec.max_draft_len + 1
+ history, slots = meta.replay_history_size, ssm0.shape[0]
+ for i, cache in caches.items():
+ P.check_replay_history(g, i, history, slots, cache.old_x, cache.old_B, cache.old_dt)
+
+ draft_in_trunk = MTP_LAYER in cm.layer_offsets
+ assert draft_in_trunk == MEMBERS[member]["enable_attention_dp"]
+ if draft_in_trunk:
+ P.kv_planes(cm, g, MTP_LAYER, cm.get_index_k_buffer(MTP_LAYER))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routed_experts.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routed_experts.py
new file mode 100644
index 000000000000..2eb055a0bb2c
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routed_experts.py
@@ -0,0 +1,671 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The routed experts at one rank (tp1): the CuteDSL experts and the NVFP4 draft experts.
+
+**CuteDSL.** The CuteDSL op tests build their operands with layout code of their own. Here the
+checkpoint's expert tensors (512 experts: packed NVFP4 `up_proj` / `gate_proj` / `down_proj`
+codes, linear e4m3 block scales, one `weight_scale_2` per expert, per-expert `input_scale`s)
+go through the target's own load-time transforms (`weights._expert_fc1_w` / `_expert_fc1_s` /
+`_expert_fc2_w` / `_expert_fc2_s`) and scalar derivation (`modeling._moe_layer_scalars`,
+`_moe_layer_scalars_cd`), once per routed backend, then through the call sequence of the
+target's routed experts -- CuteDSL (`_cute_dsl_routed`): sort, FC1 with SwiGLU, FC2,
+unpermute; trtllm-gen: the one-call block-scale runner -- with explicit routing.
+
+Two routings hit all 512 experts. In the first every token's only non-zero routing weight is
+on its own expert, so each expert's output is checked alone: a layout error confined to one
+expert (a permutation, a scale swizzle, an alpha) cannot hide in a sum over ten. On it both
+backends are held to a float64 reference over the dequantized checkpoint tensors, with
+`down_proj.input_scale` spread across the experts as a calibrated checkpoint's is. The second
+is a weighted top-10: with a uniform `down_proj.input_scale` the two backends requantize the
+SwiGLU output identically, so on it they are held to each other. The reference models the one
+rounding step that is not exact: the SwiGLU output is requantized to NVFP4 before FC2, at the
+backend's own global scale (the layer's single `global_sf` on CuteDSL, each expert's
+`1 / down_proj.input_scale` on trtllm-gen). The scalars, the CuteDSL weight row orders and the
+prefill copy's decision are also checked on the host.
+
+**NVFP4 draft experts** (`TRTLLM_MODELING_V2_MTP_MOE_NVFP4`). `weights.mtp_experts_to_nvfp4`
+requantizes FP8 128x128 block-scaled experts to NVFP4 in the trunk's trtllm-gen layout,
+`modeling._mtp_moe_nvfp4` adds the runner's scalars and `modeling._mtp_routed_nvfp4` runs
+them. Against the FP8 block-scale path the draft layer runs by default, on the same random
+experts, activations and router, the NVFP4 output must stay inside a relative-error band of
+the FP8 one, keep its magnitude (a wrong runner scalar moves the norm, which unbiased
+quantization noise does not), and leave no FP8 copy behind.
+"""
+
+import functools
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import (
+ fp4_block_scale_moe_runner,
+ quantize_nvfp4,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import (
+ fp4_quantize,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp8_block_scale_moe_runner import ( # noqa: E501
+ fp8_block_scale_moe_runner,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.renorm_moe_routing_op import ( # noqa: E501
+ renorm_moe_routing_op,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.quantization.fp8_quantize_1x128 import ( # noqa: E501
+ fp8_quantize_1x128,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501
+ modeling as M,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501
+ weights as W,
+)
+
+_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="the experts run on a CUDA device")
+DEV = "cuda"
+
+
+# ---- the scalars and the weight layout, on the host ---------------------------------------
+
+SCALAR_EXPERTS = 8
+
+
+def _scalar_rows(isc2: torch.Tensor, seed: int = 0) -> dict:
+ gen = torch.Generator().manual_seed(seed)
+ isc1 = torch.full((SCALAR_EXPERTS,), 0.0123)
+ ws2_1 = torch.rand(SCALAR_EXPERTS, generator=gen) * 1e-3 + 1e-4
+ return {
+ "l0_e_isc1": isc1,
+ "l0_e_isc1_up": isc1.clone(),
+ "l0_e_ws2_1": ws2_1,
+ "l0_e_ws2_1_up": ws2_1.clone(),
+ "l0_e_isc2": isc2,
+ "l0_e_ws2_2": torch.rand(SCALAR_EXPERTS, generator=gen) * 1e-3 + 1e-4,
+ }
+
+
+@pytest.mark.cpu_only
+def test_uniform_isc2_keeps_the_per_expert_products_bitwise() -> None:
+ w = _scalar_rows(torch.full((SCALAR_EXPERTS,), 0.0371))
+ got = M._moe_layer_scalars_cd(w, 0)
+ assert torch.equal(got["alpha2"], w["l0_e_isc2"] * w["l0_e_ws2_2"])
+ assert torch.equal(got["global_sf"], 1.0 / w["l0_e_isc2"][:1])
+ assert got["global_sf"].shape == (1,)
+ assert torch.equal(got["alpha1"], w["l0_e_isc1"] * w["l0_e_ws2_1"])
+
+
+@pytest.mark.cpu_only
+def test_isc2_spread_is_folded_into_alpha2() -> None:
+ """CuteDSL requantizes FC1's output at ONE global scale per layer, where trtllm-gen folds
+ each expert's `down_proj.input_scale` into its own FC1 scalar. Whatever the spread (here
+ 40x), `global_sf * alpha2[e]` is expert e's own `ws2_2[e]`, the product trtllm-gen forms
+ from `(1 / isc2[e]) * (isc2[e] * ws2_2[e])`."""
+ isc2 = 0.01 * torch.linspace(1.0, 40.0, SCALAR_EXPERTS)
+ w = _scalar_rows(isc2.flip(0))
+ cute = M._moe_layer_scalars_cd(w, 0)
+ gen = M._moe_layer_scalars(w, 0, 0, SCALAR_EXPERTS)
+ isc2 = w["l0_e_isc2"]
+ assert torch.equal(cute["global_sf"], 1.0 / isc2.max().reshape(1))
+ # No expert's FC2 input is scaled above what the checkpoint calibrated for it.
+ assert (cute["global_sf"] <= 1.0 / isc2).all()
+ assert torch.allclose(
+ cute["global_sf"] * cute["alpha2"], (1.0 / isc2) * gen["output2_scale_scalar"], rtol=1e-6
+ )
+ assert torch.allclose(cute["global_sf"] * cute["alpha2"], w["l0_e_ws2_2"], rtol=1e-6)
+
+
+@pytest.mark.cpu_only
+def test_gate_up_input_scale_mismatch_is_refused() -> None:
+ w = _scalar_rows(torch.full((SCALAR_EXPERTS,), 0.0371))
+ w["l0_e_isc1_up"] = w["l0_e_isc1_up"] * 1.5
+ with pytest.raises(AssertionError, match="gate/up input_scale differ"):
+ M._moe_layer_scalars(w, 0, 0, SCALAR_EXPERTS)
+
+
+@pytest.mark.cpu_only
+def test_cutedsl_fc1_pairs_up_and_gate_rows_in_64_row_groups() -> None:
+ """Destination rows `[128k, 128k + 64)` hold up rows `[64k, 64k + 64)` and the next 64
+ hold the matching gate rows, which sit `I` rows later in the `[up ; gate]` stack."""
+ inter = 640
+ perm = W._fc1_perm(2 * inter, "cpu", SimpleNamespace(trunk_moe_backend="CUTEDSL"))
+ assert sorted(perm.tolist()) == list(range(2 * inter))
+ for k in range(inter // 64):
+ assert perm[128 * k : 128 * k + 64].tolist() == list(range(64 * k, 64 * k + 64))
+ assert perm[128 * k + 64 : 128 * k + 128].tolist() == list(
+ range(inter + 64 * k, inter + 64 * k + 64)
+ )
+
+
+@pytest.mark.cpu_only
+def test_cutedsl_fc2_keeps_the_checkpoint_row_order_and_trtllm_gen_shuffles_it() -> None:
+ """trtllm-gen's 32-row block shuffle moves source row `4u + v` to destination row
+ `8v + u`; CuteDSL reads `down_proj` as the checkpoint stores it."""
+ rows = torch.arange(64).unsqueeze(1).expand(64, 8).contiguous()
+ kept = W._expert_fc2_w(rows, SimpleNamespace(trunk_moe_backend="CUTEDSL"))
+ assert torch.equal(kept, rows)
+ shuffled = W._expert_fc2_w(rows, SimpleNamespace(trunk_moe_backend="TRTLLM"))
+ for u in range(8):
+ for v in range(4):
+ assert int(shuffled[8 * v + u, 0]) == 4 * u + v
+ assert int(shuffled[32 + 8 * v + u, 0]) == 32 + 4 * u + v
+
+
+@pytest.mark.cpu_only
+@pytest.mark.parametrize(
+ "one_rank, backend, max_num_tokens, env, kept",
+ [
+ (True, "TRTLLM", 4096, None, True),
+ (True, "TRTLLM", 4095, None, False),
+ (True, "TRTLLM", 16384, "0", False),
+ (True, "CUTEDSL", 16384, None, False),
+ (False, "TRTLLM", 16384, None, False),
+ ],
+)
+def test_second_expert_copy_kept_only_where_its_launches_can_happen(
+ monkeypatch, one_rank, backend, max_num_tokens, env, kept
+):
+ """The copy costs about 1.3 GiB a layer: it is kept only at one rank, under trtllm-gen,
+ when a step's rows can reach the CuteDSL dispatch threshold, and not when the env var is
+ off."""
+ if env is None:
+ monkeypatch.delenv(M._MOE_PREFILL_COPY_ENV, raising=False)
+ else:
+ monkeypatch.setenv(M._MOE_PREFILL_COPY_ENV, env)
+ assert M._moe_prefill_copy(one_rank, backend, max_num_tokens) is kept
+
+
+# ---- CuteDSL: the experts' geometry and gates ---------------------------------------------
+
+#: The trunk's routed experts: count, top-k, hidden, intermediate, NVFP4 block.
+E, K, H, INTER, VEC = 512, 10, 2560, 640, 16
+#: The activation's input scale, one number for all experts: |x| = 5 maps to the e4m3 maximum.
+ISC1 = 5.0 / (448 * 6)
+#: The SwiGLU output maximum `down_proj.input_scale` is calibrated to. With `spread` each
+#: expert's own value ranges over 2^-1/2 .. 2^1/2 of it, so the layer's `max` is not every
+#: expert's.
+ISC2 = 64.0 / (448 * 6)
+E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
+#: The routed runners read activation scales as a linear buffer, not the 128x4-swizzled
+#: one of a dense NVFP4 GEMM; byte counts match at multiples of 128 rows, so the wrong
+#: one passes silently.
+SF_LINEAR = False
+
+#: Tokens of the two routings: one per expert, and a weighted top-10.
+ONE_HOT_TOKENS = E
+WEIGHTED_TOKENS = 1024
+
+BACKENDS = ("CUTEDSL", "TRTLLM")
+#: The float64 reference's gate, relative per token. Both backends sit at 2.5e-3 or below:
+#: the bf16 rounding of the output is 1e-3 of it, the rest is the reference's fp32 / tie
+#: differences. A wrong layout is ~1; a wrong per-expert scalar is its own relative error.
+REFERENCE_TOLERANCE = 5e-3
+#: The same, between the backends under a uniform `down_proj.input_scale`. Each backend is
+#: within `REFERENCE_TOLERANCE` of the reference, so the two are within twice that of each
+#: other. Both requantize the SwiGLU output to NVFP4, and an fp32 difference in the last
+#: place moves an element at an e2m1 rounding boundary by a whole code: the token holding
+#: it differs by ~5e-3 while the mean stays near 5e-6. A wrong layout is still ~1.
+BACKEND_TOLERANCE = 2 * REFERENCE_TOLERANCE
+
+
+# ---- CuteDSL: the checkpoint -------------------------------------------------
+
+
+def _quantize(dense: torch.Tensor):
+ """A bf16 `[R, C]` matrix as the checkpoint stores it: packed e2m1 codes `[R, C/2]`
+ uint8, e4m3 block scales `[R, C/16]`, and the `weight_scale_2` that undoes the global
+ scale the block scales were derived under."""
+ rows, cols = dense.shape
+ gs = (448 * 6 / dense.abs().amax().float()).reshape(1)
+ codes, sf = torch.ops.trtllm.fp4_quantize(dense, gs, VEC, False, False)
+ return (
+ codes.view(torch.uint8).view(rows, cols // 2),
+ sf.view(torch.uint8).view(rows, cols // VEC).view(torch.float8_e4m3fn),
+ (1.0 / gs).squeeze(0),
+ )
+
+
+@functools.lru_cache(maxsize=None)
+def _checkpoint(spread: bool) -> SimpleNamespace:
+ """All 512 experts of one layer, with the checkpoint's tensor names and dtypes.
+
+ `up_proj` and `gate_proj` are quantized together, so they share one `weight_scale_2`, as
+ the fused FC1 requires. Expert magnitudes spread over 2^-1 .. 2^1/2 (FC1) and 2^-1 .. 2^1
+ (FC2), so a wrong per-expert scalar shows."""
+ g = torch.Generator(device=DEV).manual_seed(20261014)
+ spread1 = torch.exp2(torch.rand(E, generator=g, device=DEV) * 1.5 - 1)
+ spread2 = torch.exp2(torch.rand(E, generator=g, device=DEV) * 2 - 1)
+ fields = {k: [] for k in ("up_w", "up_s", "gate_w", "gate_s", "down_w", "down_s")}
+ ws2_gu, ws2_dn = [], []
+ for e in range(E):
+ gate_up = torch.randn(2 * INTER, H, generator=g, device=DEV) * (spread1[e] / H**0.5)
+ down = torch.randn(H, INTER, generator=g, device=DEV) * (spread2[e] / INTER**0.5)
+ codes, sf, ws2 = _quantize(gate_up.to(torch.bfloat16))
+ fields["up_w"].append(codes[:INTER])
+ fields["gate_w"].append(codes[INTER:])
+ fields["up_s"].append(sf[:INTER])
+ fields["gate_s"].append(sf[INTER:])
+ ws2_gu.append(ws2)
+ codes, sf, ws2 = _quantize(down.to(torch.bfloat16))
+ fields["down_w"].append(codes)
+ fields["down_s"].append(sf)
+ ws2_dn.append(ws2)
+ ws2_gu, ws2_dn = torch.stack(ws2_gu), torch.stack(ws2_dn)
+ isc1 = torch.full((E,), ISC1, dtype=torch.float32, device=DEV)
+ isc2 = torch.full((E,), ISC2, dtype=torch.float32, device=DEV)
+ if spread:
+ isc2 = isc2 * torch.exp2(torch.rand(E, generator=g, device=DEV) - 0.5)
+ return SimpleNamespace(
+ **fields,
+ ws2_gu=ws2_gu,
+ ws2_dn=ws2_dn,
+ isc2=isc2,
+ # The rows `modeling._moe_layer_scalars` / `_moe_layer_scalars_cd` read, under layer 0.
+ scalars={
+ "l0_e_isc1": isc1,
+ "l0_e_isc1_up": isc1.clone(),
+ "l0_e_ws2_1": ws2_gu,
+ "l0_e_ws2_1_up": ws2_gu.clone(),
+ "l0_e_isc2": isc2,
+ "l0_e_ws2_2": ws2_dn,
+ },
+ )
+
+
+# ---- CuteDSL: what the target builds from it ---------------------------------
+
+
+@functools.lru_cache(maxsize=None)
+def _layer(backend: str, spread: bool) -> SimpleNamespace:
+ """The layer's expert operands for `backend`, built as the weight table builds them:
+ one transform call per expert and tensor, stacked, then viewed as `_layer_views`
+ views them."""
+ ck = _checkpoint(spread)
+ core = SimpleNamespace(trunk_moe_backend=backend)
+
+ def stack(fn, *tensors):
+ return torch.stack([fn(tuple(t[e] for t in tensors), core) for e in range(E)])
+
+ def stack1(fn, tensor):
+ return torch.stack([fn(tensor[e], core) for e in range(E)])
+
+ fc1_w = stack(W._expert_fc1_w, ck.up_w, ck.gate_w)
+ fc1_s = stack(W._expert_fc1_s, ck.up_s, ck.gate_s)
+ fc2_w = stack1(W._expert_fc2_w, ck.down_w)
+ fc2_s = stack1(W._expert_fc2_s, ck.down_s)
+ if backend == "CUTEDSL":
+ w_dtype, s_dtype = torch.float4_e2m1fn_x2, torch.uint8
+ else:
+ w_dtype, s_dtype = torch.uint8, torch.float8_e4m3fn
+ return SimpleNamespace(
+ fc1_w=fc1_w.view(w_dtype),
+ fc1_s=fc1_s.view(s_dtype),
+ fc2_w=fc2_w.view(w_dtype),
+ fc2_s=fc2_s.view(s_dtype),
+ scalars=_expert_scalars(ck.scalars, backend),
+ )
+
+
+def _expert_scalars(rows: dict, backend: str) -> dict:
+ """Layer 0's routed-expert scalars as the target derives them for `backend`: the
+ trtllm-gen runner's (with the FC1 activation's global scale both backends quantize at),
+ plus on CuteDSL its FC1 / FC2 alphas and FC2 global scale."""
+ scalars = M._moe_layer_scalars(rows, 0, 0, E)
+ if backend == "CUTEDSL":
+ scalars.update(M._moe_layer_scalars_cd(rows, 0))
+ return scalars
+
+
+def _cutedsl_routed(layer, x, ids, weights) -> torch.Tensor:
+ """The CuteDSL branch of `_moe_block`, from the quantization to the unpermute."""
+ sc = layer.scalars
+ tokens = x.shape[0]
+ xq, xsf = fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, SF_LINEAR)
+ tile_expert, tile_limit, expanded_to_permuted, permuted_to_expanded, _, live_tiles = M.moe_sort(
+ token_selected_experts=ids,
+ token_final_scales=weights,
+ num_experts=E,
+ top_k=K,
+ local_expert_offset=0,
+ local_num_experts=E,
+ tile_tokens_dim=M._CUTEDSL_TILE_SIZE,
+ )
+ h, h_sf = M.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion(
+ input=xq.view(torch.float4_e2m1fn_x2),
+ weight=layer.fc1_w,
+ input_scale=xsf.view(tokens, H // M._SF_VEC),
+ weight_scale=layer.fc1_s,
+ alpha=sc["alpha1"],
+ tile_idx_to_group_idx=tile_expert,
+ tile_idx_to_mn_limit=tile_limit,
+ permuted_idx_to_expanded_idx=permuted_to_expanded,
+ num_non_exiting_tiles=live_tiles,
+ global_sf=sc["global_sf"],
+ num_experts=E,
+ top_k=K,
+ num_local_experts=E,
+ local_expert_offset=0,
+ tile_size=M._CUTEDSL_TILE_SIZE,
+ scaling_vector_size=M._SF_VEC,
+ activation_type=int(M.ActivationType.Swiglu),
+ )
+ y = M.cute_dsl_nvfp4_grouped_gemm(
+ input=h,
+ weight=layer.fc2_w,
+ input_scale=h_sf,
+ weight_scale=layer.fc2_s,
+ alpha=sc["alpha2"],
+ tile_idx_to_group_idx=tile_expert,
+ num_non_exiting_tiles=live_tiles,
+ num_experts=E,
+ top_k=K,
+ num_local_experts=E,
+ local_expert_offset=0,
+ tile_size=M._CUTEDSL_TILE_SIZE,
+ output_dtype=torch.bfloat16,
+ scaling_vector_size=M._SF_VEC,
+ )
+ return M.moe_unpermute_inplace(
+ permuted_input=y,
+ output=torch.empty(tokens, H, dtype=torch.bfloat16, device=x.device),
+ expanded_idx_to_permuted_idx=expanded_to_permuted,
+ topk_scales=weights,
+ )
+
+
+def _trtllm_gen_routed(layer, x, ids, weights) -> torch.Tensor:
+ """The trtllm-gen branch of `_moe_block`."""
+ sc = layer.scalars
+ xq, xsf = fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, SF_LINEAR)
+ return fp4_block_scale_moe_runner(
+ routing_logits=None,
+ routing_bias=None,
+ hidden_states=xq,
+ hidden_states_scale=xsf.view(torch.float8_e4m3fn),
+ gemm1_weights=layer.fc1_w,
+ gemm1_weights_scale=layer.fc1_s,
+ gemm1_bias=None,
+ gemm1_alpha=None,
+ gemm1_beta=None,
+ gemm1_clamp_limit=None,
+ gemm2_weights=layer.fc2_w,
+ gemm2_weights_scale=layer.fc2_s,
+ gemm2_bias=None,
+ output1_scale_scalar=sc["output1_scale_scalar"],
+ output1_scale_gate_scalar=sc["output1_scale_gate_scalar"],
+ output2_scale_scalar=sc["output2_scale_scalar"],
+ num_experts=E,
+ top_k=K,
+ n_group=None,
+ topk_group=None,
+ intermediate_size=INTER,
+ local_expert_offset=0,
+ local_num_experts=E,
+ routed_scaling_factor=None,
+ routing_method_type=M._ROUTING_METHOD_RENORM,
+ do_finalize=True,
+ act_type=M._ACT_TYPE_SWIGLU,
+ topk_weights=weights.to(torch.bfloat16),
+ topk_ids=ids,
+ )[0]
+
+
+# ---- CuteDSL: the inputs and the reference -----------------------------------
+
+
+@functools.lru_cache(maxsize=None)
+def _problem(case: str):
+ """`(x, ids, weights)`: bf16 activations, expert ids `[T, K]` int32 and routing weights
+ `[T, K]` fp32 (exact in bf16, which is the dtype trtllm-gen takes them in)."""
+ tokens = ONE_HOT_TOKENS if case == "one_hot" else WEIGHTED_TOKENS
+ g = torch.Generator(device=DEV).manual_seed(7 if case == "one_hot" else 8)
+ x = torch.randn(tokens, H, generator=g, device=DEV).to(torch.bfloat16)
+ scores = torch.rand(tokens, E, generator=g, device=DEV)
+ if case == "one_hot":
+ # The forced expert has the top score, so it sits in slot 0.
+ scores[torch.arange(tokens), torch.arange(tokens) % E] = 2.0
+ weights = torch.zeros(tokens, K, device=DEV)
+ weights[:, 0] = 1.0
+ else:
+ weights = torch.softmax(torch.randn(tokens, K, generator=g, device=DEV), dim=-1)
+ weights = weights.to(torch.bfloat16).float()
+ ids = scores.topk(K, dim=-1).indices.to(torch.int32)
+ return x, ids, weights
+
+
+def _dequant(codes: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
+ """`[R, C/2]` packed e2m1 (element 2j in the low nibble) and `[R, C/16]` e4m3 scales
+ -> float64 `[R, C]`."""
+ magnitude = torch.tensor(E2M1, dtype=torch.float64, device=codes.device)
+ table = torch.cat([magnitude, -magnitude])
+ lo, hi = table[(codes & 0xF).long()], table[(codes >> 4).long()]
+ values = torch.stack([lo, hi], dim=-1).reshape(codes.shape[0], -1)
+ return values * scales.view(torch.float8_e4m3fn).double().repeat_interleave(VEC, dim=1)
+
+
+def _reference(backend: str, case: str, spread: bool) -> torch.Tensor:
+ """The routed output in float64, `[T, H]`.
+
+ The activation is what the kernels see, the dequantized NVFP4 of the target's own
+ quantization. Each expert's weights are dequantized from the checkpoint tensors, times
+ its `weight_scale_2`. The SwiGLU output is requantized at `backend`'s global scale."""
+ ck, layer = _checkpoint(spread), _layer(backend, spread)
+ x, ids, weights = _problem(case)
+ sc = layer.scalars
+ xq, xsf = fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, SF_LINEAR)
+ x_dq = _dequant(xq.view(torch.uint8), xsf.view(x.shape[0], H // VEC)) * ISC1
+ if backend == "CUTEDSL":
+ g2 = [float(sc["global_sf"])] * E
+ else:
+ g2 = (1.0 / ck.isc2).tolist()
+ out = torch.zeros(x.shape[0], H, dtype=torch.float64, device=DEV)
+ for e in range(E):
+ token, slot = (ids == e).nonzero(as_tuple=True)
+ w = weights[token, slot].double()
+ token, w = token[w != 0], w[w != 0]
+ if token.numel() == 0:
+ continue
+ ws2 = ck.ws2_gu[e].double()
+ xe = x_dq[token]
+ up = xe @ (_dequant(ck.up_w[e], ck.up_s[e]) * ws2).T
+ gate = xe @ (_dequant(ck.gate_w[e], ck.gate_s[e]) * ws2).T
+ act = up * gate * torch.sigmoid(gate)
+ act = quantize_nvfp4(act.float(), g2[e]).double() / g2[e]
+ y = act @ (_dequant(ck.down_w[e], ck.down_s[e]) * ck.ws2_dn[e].double()).T
+ out.index_add_(0, token, y * w[:, None])
+ return out
+
+
+@functools.lru_cache(maxsize=None)
+def _routed(backend: str, case: str, spread: bool) -> torch.Tensor:
+ x, ids, weights = _problem(case)
+ run = _cutedsl_routed if backend == "CUTEDSL" else _trtllm_gen_routed
+ return run(_layer(backend, spread), x, ids, weights).double()
+
+
+def _token_error(out: torch.Tensor, ref: torch.Tensor) -> torch.Tensor:
+ """Each token's `|out - ref| / |ref|` (L2 over the hidden width)."""
+ return (out - ref).norm(dim=1) / ref.norm(dim=1).clamp_min(1e-12)
+
+
+@_cuda
+@pytest.mark.parametrize("backend", BACKENDS)
+def test_routed_output_matches_the_float64_reference(backend: str) -> None:
+ case = "one_hot"
+ _, ids, weights = _problem(case)
+ assert torch.unique(ids[weights != 0]).numel() == E, "the routing misses an expert"
+ err = _token_error(_routed(backend, case, True), _reference(backend, case, True))
+ print(f"[err] {backend} {case}: mean {err.mean():.4e} max {err.max():.4e}")
+ assert err.max() <= REFERENCE_TOLERANCE, (err.mean().item(), err.max().item())
+
+
+@_cuda
+def test_cutedsl_matches_trtllm_gen_under_a_uniform_input_scale() -> None:
+ case = "weighted"
+ _, ids, weights = _problem(case)
+ assert torch.unique(ids[weights != 0]).numel() == E, "the routing misses an expert"
+ err = _token_error(_routed("CUTEDSL", case, False), _routed("TRTLLM", case, False))
+ print(f"[err] CUTEDSL vs TRTLLM {case}: mean {err.mean():.4e} max {err.max():.4e}")
+ assert err.max() <= BACKEND_TOLERANCE, (err.mean().item(), err.max().item())
+
+
+# ---- the NVFP4 draft experts --------------------------------------------------------------
+
+#: The draft experts' FP8 weight block.
+BLK = 128
+#: A trunk calibration to derive the static draft scales from: the checkpoint's last trunk layer
+#: (FC1 input, FC2 input amax), the layer that feeds the draft.
+TRUNK_AMAX = (11.1875, 57.75)
+#: Relative RMS of NVFP4 - FP8 over the FP8 output, the floor on each row's cosine, and the
+#: band on the output norm ratio. Measured here: 0.239-0.251 RMS, row cosine >= 0.961, at every
+#: row count; on the checkpoint's real draft experts 0.222-0.231 and >= 0.964 (bench). That is
+#: the NVFP4 chain's own noise -- four e2m1 roundings (activation, FC1 weight, FC1 output, FC2
+#: weight) at ~0.1 each, compounded by the SwiGLU product -- not a bias: it is flat across
+#: activation scales 0.5-4x, and FP8 sits at 0.07 from the fp32 reference of the same weights.
+REL_RMS_BAND = 0.3
+COS_FLOOR = 0.95
+NORM_RATIO_BAND = 0.1
+
+
+def _fp8_blocks(rows: int, cols: int, gen: torch.Generator):
+ """`[E, rows, cols]` e4m3 codes and `[E, rows/128, cols/128]` fp32 scales of a random
+ expert stack whose 128x128 blocks span 16x in magnitude, as trained experts' do."""
+ w = torch.randn(E, rows, cols, generator=gen, device=DEV) * cols**-0.5
+ spread = torch.exp2(
+ torch.randint(-2, 3, (E, rows // BLK, 1, cols // BLK, 1), generator=gen, device=DEV)
+ ).float()
+ blocks = w.view(E, rows // BLK, BLK, cols // BLK, BLK) * spread
+ scale = blocks.abs().amax(dim=(2, 4), keepdim=True) / 448.0
+ codes = (blocks / scale).to(torch.float8_e4m3fn).view(E, rows, cols)
+ return codes, scale.view(E, rows // BLK, cols // BLK).contiguous()
+
+
+@pytest.fixture(scope="module")
+def draft_experts():
+ gen = torch.Generator(device=DEV).manual_seed(20261022)
+ fc1_w, fc1_s = _fp8_blocks(2 * INTER, H, gen)
+ fc2_w, fc2_s = _fp8_blocks(H, INTER, gen)
+ router = (torch.randn(E, H, generator=gen, device=DEV) * H**-0.5).to(torch.bfloat16)
+ fp8 = SimpleNamespace(fc1_w=fc1_w, fc1_s=fc1_s, fc2_w=fc2_w, fc2_s=fc2_s, router=router.t())
+ p = "l1_mtp_"
+ w = {p + k: getattr(fp8, k).clone() for k in ("fc1_w", "fc1_s", "fc2_w", "fc2_s")}
+ w["l0_e_isc1"] = torch.full((E,), TRUNK_AMAX[0] / 2688.0, device=DEV)
+ w["l0_e_isc2"] = torch.full((E,), TRUNK_AMAX[1] / 2688.0, device=DEV)
+ # No packed FP8 draft projections: the shared expert runs its bf16 GEMMs.
+ core = SimpleNamespace(
+ w=w,
+ mtp_layer_idx=1,
+ num_layers=1,
+ num_experts=E,
+ topk=K,
+ moe_inter=INTER,
+ _mtp_gemv8={},
+ )
+ fields = M._mtp_moe_nvfp4(core)
+ nvfp4 = SimpleNamespace(router=fp8.router, **fields)
+ return core, fp8, nvfp4
+
+
+def _fp8_routed(fp8, x):
+ """The FP8 block-scale draft path, call for call as `_mtp_moe_block` makes it."""
+ logits = M.cublas_mm(x, fp8.router)
+ ids, wts = renorm_moe_routing_op(logits, topk=K, output_dtype=torch.bfloat16)
+ xq, xsf = fp8_quantize_1x128(x, use_ue8m0=False)
+ return fp8_block_scale_moe_runner(
+ routing_logits=None,
+ routing_bias=None,
+ hidden_states=xq,
+ hidden_states_scale=xsf,
+ gemm1_weights=fp8.fc1_w,
+ gemm1_weights_scale=fp8.fc1_s,
+ gemm2_weights=fp8.fc2_w,
+ gemm2_weights_scale=fp8.fc2_s,
+ num_experts=E,
+ top_k=K,
+ num_fused_shared_experts=None,
+ n_group=None,
+ topk_group=None,
+ intermediate_size=INTER,
+ local_expert_offset=0,
+ local_num_experts=E,
+ routed_scaling_factor=None,
+ routing_method_type=M._ROUTING_METHOD_RENORM,
+ topk_weights=wts,
+ topk_ids=ids,
+ act_type=M._ACT_TYPE_SWIGLU,
+ gemm1_clamp_limit=None,
+ )
+
+
+@_cuda
+def test_draft_fp8_copies_freed_and_layout(draft_experts):
+ core, _, nvfp4 = draft_experts
+ for k in ("fc1_w", "fc1_s", "fc2_w", "fc2_s"):
+ assert core.w[f"l1_mtp_{k}"].numel() == 0, k
+ assert nvfp4.fc1_w.shape == (E, 2 * INTER, H // 2) and nvfp4.fc1_w.dtype == torch.uint8
+ assert nvfp4.fc1_s.shape == (E, 2 * INTER, H // 16) and nvfp4.fc1_s.dtype == torch.float8_e4m3fn
+ assert nvfp4.fc2_w.shape == (E, H, INTER // 2) and nvfp4.fc2_s.shape == (E, H, INTER // 16)
+ for s in (nvfp4.out1_scale, nvfp4.out1_gate_scale, nvfp4.out2_scale):
+ assert s.shape == (E,) and s.dtype == torch.float32 and torch.isfinite(s).all()
+ assert nvfp4.moe_gscale.shape == (1,)
+
+
+@_cuda
+@pytest.mark.parametrize("rows", [1, 4, 32, 128, 1024])
+def test_draft_nvfp4_tracks_fp8(draft_experts, rows):
+ core, fp8, nvfp4 = draft_experts
+ gen = torch.Generator(device=DEV).manual_seed(rows)
+ x = torch.randn(rows, H, generator=gen, device=DEV).to(torch.bfloat16)
+ ref = _fp8_routed(fp8, x).float()
+ out = M._mtp_routed_nvfp4(core, nvfp4, x)
+ assert out.shape == (rows, H) and out.dtype == torch.bfloat16
+ out = out.float()
+ rel = ((out - ref).norm() / ref.norm()).item()
+ cos = torch.nn.functional.cosine_similarity(out, ref, dim=1).min().item()
+ ratio = (out.norm() / ref.norm()).item()
+ print(f"rows={rows}: rel RMS {rel:.4f}, min row cosine {cos:.4f}, norm ratio {ratio:.4f}")
+ assert rel < REL_RMS_BAND, rel
+ assert cos > COS_FLOOR, cos
+ assert abs(ratio - 1.0) < NORM_RATIO_BAND, ratio
+
+
+@_cuda
+@pytest.mark.parametrize("rows", [1, 64, 512])
+def test_draft_shared_expert_overlap_preserves_serial_output(draft_experts, rows):
+ """`DecodeTarget._mtp_moe` over the NVFP4 experts runs the shared expert on the side stream
+ beside them: its output is the serial composition's bitwise, eager and replayed from a
+ CUDA graph over new inputs."""
+ core, _, nvfp4 = draft_experts
+ gen = torch.Generator(device=DEV).manual_seed(42)
+ # The checkpoint's fused shared expert: gate|up rows, then the one gate-logit row.
+ sh_gu = (torch.randn(2 * INTER + 1, H, generator=gen, device=DEV) / H**0.5).to(torch.bfloat16)
+ sh_dn = (torch.randn(H, INTER, generator=gen, device=DEV) / INTER**0.5).to(torch.bfloat16)
+ weights = SimpleNamespace(
+ **vars(nvfp4),
+ sh_gu=sh_gu[: 2 * INTER].t(),
+ sh_dn=sh_dn.t(),
+ sh_gate=sh_gu[2 * INTER :].t(),
+ )
+ target = object.__new__(M.DecodeTarget)
+ target.core = SimpleNamespace(
+ **vars(core), _mtp=weights, _side_stream=torch.cuda.Stream(), shared_inter=INTER
+ )
+ x = torch.randn(rows, H, generator=gen, device=DEV).to(torch.bfloat16)
+
+ def serial():
+ routed = M._mtp_routed_nvfp4(target.core, weights, x)
+ up = M.silu_and_mul_rows(M.cublas_mm(x, weights.sh_gu), INTER)
+ shared = M.cublas_mm(up, weights.sh_dn)
+ gate = M.cublas_mm(x, weights.sh_gate)
+ return M.fused_sigmoid_gate_mul_add(routed, gate, shared)
+
+ expected = serial().clone()
+ torch.testing.assert_close(target._mtp_moe(x), expected, rtol=0, atol=0)
+ if rows == 64:
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ captured = target._mtp_moe(x)
+ for _ in range(3):
+ x.normal_(generator=gen)
+ expected = serial().clone()
+ graph.replay()
+ torch.testing.assert_close(captured, expected, rtol=0, atol=0)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py
new file mode 100644
index 000000000000..5a6960713b27
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py
@@ -0,0 +1,396 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Qwen3.8-Flash-Next-NVFP4 routing at its layouts, driven by synthetic configs on CPU.
+
+`qwen4_exp` routes to its one target one rank (tp1), tp4 and tep4 with or without MTP, and
+the attention-DP layouts dep4 (with or without MTP) and dep2 (the disaggregated-serving
+engines' layout, without speculation), neither with a breakable prefill graph. The layouts
+the target does not serve fall back under `auto`. The tp1 engine choices are exercised in
+`tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py`.
+"""
+
+from dataclasses import replace
+
+import pytest
+import torch
+from transformers import PretrainedConfig
+
+from tensorrt_llm._torch._experimental.modeling_v2._router_index import modeling_v2_resolve
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp import routing
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501
+ modeling as M,
+)
+from tensorrt_llm._torch.model_config import ModelConfig
+from tensorrt_llm.llmapi.llm_args import (
+ KvCacheConfig,
+ MTPDecodingConfig,
+ NGramDecodingConfig,
+ TorchLlmArgs,
+)
+from tensorrt_llm.mapping import Mapping
+from tensorrt_llm.models.modeling_utils import QuantConfig
+from tensorrt_llm.quantization.mode import QuantAlgo
+
+pytestmark = pytest.mark.cpu_only
+
+_DUMMY_MODEL = "/tmp/dummy_model"
+_ARCH = "Qwen4ExpForConditionalGeneration"
+_TARGET = (
+ "models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.modeling."
+ "ModelingV2Qwen38FlashNextNvfp4Sm103Default"
+)
+
+# The four-rank layouts, as the LLM API's parallel arguments build them. tp4 leaves the
+# expert split to Mapping's default, which is the TP size.
+_LAYOUTS = {
+ "tp4": dict(world_size=4, tp_size=4),
+ "tep4": dict(world_size=4, tp_size=4, moe_tp_size=1, moe_ep_size=4),
+ "dep4": dict(world_size=4, tp_size=4, moe_tp_size=1, moe_ep_size=4, enable_attention_dp=True),
+}
+
+# The two-rank layout both engines of the disaggregated-serving test run: attention DP,
+# the routed experts in two 256-expert windows.
+_DEP2 = dict(world_size=2, tp_size=2, moe_tp_size=1, moe_ep_size=2, enable_attention_dp=True)
+
+# The layouts with attention TP, which serve MTP and the breakable prefill graph, and the
+# attention-DP ones, which serve neither.
+_ATTENTION_TP = ["tp4", "tep4"]
+_ATTENTION_DP = {"dep4": _LAYOUTS["dep4"], "dep2": _DEP2}
+
+# Every draft length the target runs.
+_DRAFT_LENS = list(range(1, 8))
+
+
+def _qualified(target) -> str:
+ """``models...modeling.`` for a resolved target class."""
+ return f"{target.__module__.split('.modeling_v2.')[1]}.{target.__name__}"
+
+
+@pytest.fixture(autouse=True)
+def _on_sm103(monkeypatch):
+ """Route as if this were a GB300, wherever the test actually runs, with the GDN
+ cached replay at its default (on)."""
+ monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+ monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (10, 3))
+ monkeypatch.delenv("TRTLLM_USE_GDN_REPLAY", raising=False)
+
+
+def _llm_args(mode="auto", **kwargs) -> TorchLlmArgs:
+ return TorchLlmArgs(model=_DUMMY_MODEL, modeling_v2=mode, **kwargs)
+
+
+def _qwen4_exp_config(text_overrides=None, **vision_overrides):
+ """The shape Qwen3.8-Flash-Next-NVFP4's normalized composite config declares."""
+ text = dict(
+ num_hidden_layers=48,
+ hidden_size=2560,
+ num_experts=512,
+ num_attention_heads=24,
+ linear_num_value_heads=48,
+ vocab_size=248320,
+ )
+ text.update(text_overrides or {})
+ vision = dict(
+ depth=27,
+ hidden_size=1152,
+ num_heads=16,
+ spatial_merge_size=2,
+ out_hidden_size=2560,
+ num_position_embeddings=2304,
+ deepstack_visual_indexes=[],
+ hidden_act="gelu_pytorch_tanh",
+ )
+ vision.update(vision_overrides)
+ return PretrainedConfig(
+ architectures=[_ARCH],
+ text_config=PretrainedConfig(**text),
+ vision_config=PretrainedConfig(**vision),
+ )
+
+
+def _model_config(spec_config=None, pretrained_config=None, **mapping_kwargs) -> ModelConfig:
+ """The loaded config; the checkpoint's own KV is bf16 (no KV-cache algorithm)."""
+ return ModelConfig(
+ pretrained_config=pretrained_config or _qwen4_exp_config(),
+ mapping=Mapping(**mapping_kwargs),
+ spec_config=spec_config,
+ quant_config=QuantConfig(kv_cache_quant_algo=None),
+ )
+
+
+def _claimed(config, **llm_args_kwargs) -> None:
+ for mode in ("auto", "require"):
+ target = modeling_v2_resolve(config, _llm_args(mode, **llm_args_kwargs))
+ assert target is not None and _qualified(target) == _TARGET, mode
+
+
+def _declined(config, match=None, **llm_args_kwargs) -> None:
+ """``auto`` falls back to the built-in model and ``require`` says why."""
+ assert modeling_v2_resolve(config, _llm_args("auto", **llm_args_kwargs)) is None
+ with pytest.raises(ValueError, match=match):
+ modeling_v2_resolve(config, _llm_args("require", **llm_args_kwargs))
+
+
+def test_routing_admits_exactly_what_the_target_implements():
+ """`route` declines a MoE backend, a draft length or a recurrent-state dtype the
+ target does not run; this keeps routing's constants and the target's together."""
+ # At tp1 the CuteDSL backend and the fp32 state too; at the other layouts neither.
+ assert set(routing._ONE_RANK_MOE_BACKENDS) == set(M._TRUNK_MOE_BACKENDS)
+ assert set(routing._MOE_BACKENDS) == set(M._TRUNK_MOE_BACKENDS) - {"CUTEDSL"}
+ assert routing._MAX_CUTEDSL_DRAFT_LEN == M._CUTEDSL_MAX_DRAFT_LEN
+ assert routing._MAX_DRAFT_LEN == M._MAX_VERIFY_WIDTH - 1
+ assert routing._SSM_STATE_DTYPES == M._SSM_STATE_DTYPES
+ assert routing._ONE_RANK_SSM_STATE_DTYPES == M._ONE_RANK_SSM_STATE_DTYPES
+ assert routing._SM == M._SM
+
+
+# --- each layout lands on the default target ------------------------------------
+
+
+@pytest.mark.parametrize("member", sorted(_LAYOUTS))
+def test_each_member_routes_to_the_default_target(member):
+ config = _model_config(**_LAYOUTS[member])
+ assert routing.layout(config.mapping) == member
+ for dtype in ("bfloat16", "auto"):
+ _claimed(config, kv_cache_config=KvCacheConfig(mamba_ssm_cache_dtype=dtype))
+
+
+@pytest.mark.parametrize("max_draft_len", _DRAFT_LENS)
+@pytest.mark.parametrize("member", _ATTENTION_TP)
+def test_mtp_at_every_draft_length_routes_under_attention_tp(member, max_draft_len):
+ """The tune runs MTP at tp4 and tep4 over draft lengths 1 to 7."""
+ config = _model_config(MTPDecodingConfig(max_draft_len=max_draft_len), **_LAYOUTS[member])
+ assert routing.layout(config.mapping) == member
+ for dtype in ("bfloat16", "auto"):
+ _claimed(config, kv_cache_config=KvCacheConfig(mamba_ssm_cache_dtype=dtype))
+
+
+@pytest.mark.parametrize("dtype", ["fp8", "nvfp4", "auto"])
+def test_dep2_routes_to_the_default_target(dtype):
+ config = _model_config(**_DEP2)
+ assert routing.layout(config.mapping) == "dep2"
+ _claimed(config, kv_cache_config=KvCacheConfig(dtype=dtype, mamba_ssm_cache_dtype="bfloat16"))
+
+
+@pytest.mark.parametrize("max_draft_len", _DRAFT_LENS)
+def test_mtp_at_every_draft_length_routes_at_dep4(max_draft_len):
+ """Each dep4 rank drafts for its own requests and the draft layer's FP8 expert window runs
+ the trunk's attention-DP round trip: the tune's ht_dep4 config may serve MTP."""
+ config = _model_config(MTPDecodingConfig(max_draft_len=max_draft_len), **_LAYOUTS["dep4"])
+ assert routing.layout(config.mapping) == "dep4"
+ for dtype in ("bfloat16", "auto"):
+ _claimed(config, kv_cache_config=KvCacheConfig(mamba_ssm_cache_dtype=dtype))
+ _claimed(config, kv_cache_config=KvCacheConfig(dtype="fp8", mamba_ssm_cache_dtype="bfloat16"))
+
+
+# --- what the target does not run under attention DP ------------------------------
+
+
+@pytest.mark.parametrize("max_draft_len", [1, 3, 7])
+def test_dep2_declines_speculation(max_draft_len):
+ """MTP under attention DP is admitted only where its draft acceptance was measured (dep4);
+ neither family claims dep2 with speculation."""
+ config = _model_config(MTPDecodingConfig(max_draft_len=max_draft_len), **_DEP2)
+ assert routing.layout(config.mapping) == "dep2"
+ _declined(config, match=r"speculative decoding with attention DP \(dep2\)")
+
+
+@pytest.mark.parametrize("member", sorted(_ATTENTION_DP))
+def test_attention_dp_declines_the_breakable_prefill_graph(member):
+ """The engine pads every attention-DP rank to one prefill bucket, which a rank holding
+ only generation requests cannot take; neither family claims it."""
+ config = _model_config(**_ATTENTION_DP[member])
+ _declined(
+ config,
+ match=rf"prefill_cuda_graph_backend='breakable' with attention DP \({member}\)",
+ prefill_cuda_graph_backend="breakable",
+ )
+ _claimed(config, prefill_cuda_graph_backend="disabled")
+
+
+def test_attention_dp_declines_speculation_before_the_engine_choices(monkeypatch):
+ """MTP without the cached replay is declined at dep2 for attention DP: that is the
+ structural reason, and no engine choice would lift it. At dep4 the same draft is declined
+ for the missing cached replay alone."""
+ monkeypatch.setenv("TRTLLM_USE_GDN_REPLAY", "0")
+ config = _model_config(MTPDecodingConfig(max_draft_len=7), **_DEP2)
+ _declined(config, match=r"speculative decoding with attention DP \(dep2\)")
+ config = _model_config(MTPDecodingConfig(max_draft_len=7), **_LAYOUTS["dep4"])
+ _declined(config, match="TRTLLM_USE_GDN_REPLAY=0")
+
+
+def test_dep4_declines_mtp_with_the_breakable_prefill_graph():
+ """Admitting MTP at dep4 keeps the breakable prefill graph's refusal under attention DP."""
+ config = _model_config(MTPDecodingConfig(max_draft_len=3), **_LAYOUTS["dep4"])
+ _declined(
+ config,
+ match=r"prefill_cuda_graph_backend='breakable' with attention DP \(dep4\)",
+ prefill_cuda_graph_backend="breakable",
+ )
+
+
+@pytest.mark.parametrize("spec_config", [None, MTPDecodingConfig(max_draft_len=3)])
+def test_the_one_rank_deployment_routes_to_the_default_target_at_tp1(spec_config):
+ config = _model_config(spec_config)
+ assert routing.layout(config.mapping) == "tp1"
+ _claimed(config)
+ # The one-rank engine choices: the CuteDSL experts and the fp32 recurrent state.
+ _claimed(
+ replace(config, moe_backend_requested="CUTEDSL"),
+ kv_cache_config=KvCacheConfig(mamba_ssm_cache_dtype="float32"),
+ )
+
+
+def test_attention_dp_at_one_rank_declines():
+ config = _model_config(world_size=1, tp_size=1, enable_attention_dp=True)
+ assert routing.layout(config.mapping) is None
+ _declined(config)
+
+
+@pytest.mark.parametrize(
+ "mapping_kwargs",
+ [
+ dict(world_size=2, tp_size=2),
+ dict(world_size=2, tp_size=2, moe_tp_size=1, moe_ep_size=2),
+ dict(_DEP2, moe_tp_size=2, moe_ep_size=1),
+ dict(_DEP2, enable_lm_head_tp_in_adp=True),
+ dict(world_size=8, tp_size=8),
+ dict(world_size=8, tp_size=8, moe_tp_size=1, moe_ep_size=8),
+ dict(world_size=8, tp_size=8, moe_tp_size=1, moe_ep_size=8, enable_attention_dp=True),
+ ],
+ ids=["tp2", "tep2", "dep2-moetp2", "dep2-lm-head-tp", "tp8", "tep8", "dep8"],
+)
+def test_other_two_or_eight_rank_layouts_decline(mapping_kwargs):
+ """The target does not claim it, and ``require`` names the layout and what the target
+ serves."""
+ config = _model_config(**mapping_kwargs)
+ assert modeling_v2_resolve(config, _llm_args("auto")) is None
+ with pytest.raises(ValueError) as raised:
+ modeling_v2_resolve(config, _llm_args("require"))
+ message = str(raised.value)
+ assert "the qwen4_exp target is" in message, message
+ assert f"world_size={mapping_kwargs['world_size']}" in message, message
+
+
+# --- near misses at four ranks ------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "mapping_kwargs",
+ [
+ # attention DP with the experts split by intermediate
+ dict(_LAYOUTS["dep4"], moe_tp_size=4, moe_ep_size=1),
+ # experts split two ways by intermediate and two by expert
+ dict(world_size=4, tp_size=4, moe_tp_size=2, moe_ep_size=2),
+ # pipeline parallelism
+ dict(world_size=4, tp_size=2, pp_size=2),
+ # lm_head TP under attention DP
+ dict(_LAYOUTS["dep4"], enable_lm_head_tp_in_adp=True),
+ # context parallelism
+ dict(world_size=4, tp_size=2, cp_size=2),
+ # four ranks over two nodes
+ dict(_LAYOUTS["tep4"], gpus_per_node=2),
+ ],
+ ids=["dep4-moetp4", "tp4-moetp2-ep2", "tp2-pp2", "dep4-lm-head-tp", "tp2-cp2", "two-nodes"],
+)
+def test_other_layouts_at_four_ranks_decline(mapping_kwargs):
+ config = _model_config(**mapping_kwargs)
+ assert routing.layout(config.mapping) is None
+ _declined(config)
+
+
+@pytest.mark.parametrize(
+ "spec_config",
+ [
+ NGramDecodingConfig(max_draft_len=3, max_matching_ngram_size=2),
+ MTPDecodingConfig(max_draft_len=8),
+ MTPDecodingConfig(max_draft_len=3, use_mtp_vanilla=True),
+ ],
+ ids=["ngram", "draft8", "vanilla-mtp"],
+)
+def test_unsupported_speculation_declines(spec_config):
+ _declined(_model_config(spec_config, **_LAYOUTS["tep4"]))
+
+
+def test_another_checkpoint_shape_declines():
+ config = _model_config(
+ pretrained_config=_qwen4_exp_config(dict(num_hidden_layers=60)), **_LAYOUTS["tep4"]
+ )
+ _declined(config, match="text shape")
+
+
+def test_a_deepstack_vision_tower_declines():
+ config = _model_config(
+ pretrained_config=_qwen4_exp_config(deepstack_visual_indexes=[5, 11, 17]),
+ **_LAYOUTS["tep4"],
+ )
+ _declined(config, match="vision shape")
+
+
+def test_sm100_declines(monkeypatch):
+ monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (10, 0))
+ _declined(_model_config(**_LAYOUTS["tep4"]), match=r"compute capability \(10, 0\)")
+
+
+# --- engine choices the target does not implement ---------------------------------
+
+
+@pytest.mark.parametrize("dtype", ["fp8", "nvfp4"])
+@pytest.mark.parametrize("member", sorted(_LAYOUTS))
+def test_quantized_kv_cache_dtypes_match(member, dtype):
+ _claimed(_model_config(**_LAYOUTS[member]), kv_cache_config=KvCacheConfig(dtype=dtype))
+
+
+def test_other_kv_cache_algorithms_decline():
+ choices = routing.engine_choices(_model_config(**_LAYOUTS["tep4"]), _llm_args())
+ for algo in (QuantAlgo.INT8, QuantAlgo.FP8_BLOCK_SCALES):
+ with pytest.raises(routing.NoTarget, match="KV cache"):
+ routing._check_engine(None, algo, choices, "tep4")
+ _declined(
+ _model_config(**_LAYOUTS["tep4"]),
+ match="KV cache",
+ kv_cache_config=KvCacheConfig(dtype="fp8_ds_mla"),
+ )
+
+
+@pytest.mark.parametrize("backend", ["AUTO", "TRTLLM"])
+def test_the_trtllm_moe_backend_matches(backend):
+ _claimed(replace(_model_config(**_LAYOUTS["tep4"]), moe_backend_requested=backend))
+
+
+@pytest.mark.parametrize("backend", ["CUTEDSL", "CUTLASS", "TRITON", "DEEPGEMM"])
+def test_other_moe_backends_decline(backend):
+ """CUTEDSL included: its expert pipeline runs every expert on one rank (tp1)."""
+ config = replace(_model_config(**_LAYOUTS["tep4"]), moe_backend_requested=backend)
+ _declined(config, match="moe_config.backend")
+
+
+def test_a_float32_state_cache_declines():
+ _declined(
+ _model_config(**_LAYOUTS["dep4"]),
+ match="recurrent-state cache",
+ kv_cache_config=KvCacheConfig(mamba_ssm_cache_dtype="float32"),
+ )
+
+
+@pytest.mark.parametrize("backend", ["disabled", "breakable"])
+@pytest.mark.parametrize("member", _ATTENTION_TP)
+def test_no_or_the_breakable_prefill_cuda_graph_matches(member, backend):
+ _claimed(_model_config(**_LAYOUTS[member]), prefill_cuda_graph_backend=backend)
+
+
+def test_the_piecewise_prefill_cuda_graph_declines():
+ _declined(
+ _model_config(**_LAYOUTS["tep4"]), match="piecewise", prefill_cuda_graph_backend="piecewise"
+ )
+
+
+@pytest.mark.parametrize("max_draft_len", [1, 3, 7])
+def test_mtp_needs_the_cached_replay(monkeypatch, max_draft_len):
+ """The target verifies on the GDN cached replay only; without speculation it is unused."""
+ config = _model_config(MTPDecodingConfig(max_draft_len=max_draft_len), **_LAYOUTS["tep4"])
+ monkeypatch.setenv("TRTLLM_USE_GDN_REPLAY", "0")
+ _declined(config, match="TRTLLM_USE_GDN_REPLAY=0")
+ _claimed(_model_config(**_LAYOUTS["tep4"]))
+ monkeypatch.delenv("TRTLLM_USE_GDN_REPLAY")
+ _claimed(config)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_select_canonical.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_select_canonical.py
new file mode 100644
index 000000000000..6ab7a2d73a3f
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_select_canonical.py
@@ -0,0 +1,69 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the default target's block selection (`_select`).
+
+`_select` picks each row's canonical top-`block_topk` groups of its window: the exact top-k
+and then the canonical rewrite, two launches, which serve a draft step's rows (the rejection
+kernel's ragged windows) at every table width the engine's graph buckets use.
+
+Both checks run over windows short of the budget, at it and past it, with a coarse score
+grid so the k-th value ties: the rows equal the canonical definition
+(`qsa_canonical_selection.reference`: LARGE tie-break, ascending, `-1` past the window), and
+bounding the rewrite to the live columns returns the same rows.
+"""
+
+from __future__ import annotations
+
+import types
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_canonical_selection import ( # noqa: E501
+ qsa_canonical_selection,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501
+ modeling as U,
+)
+
+assert torch.cuda.is_available(), "the block selection runs on a CUDA device"
+
+K = 512
+
+
+def _case(rows: int, cols: int, seed: int) -> tuple:
+ """Ragged windows over a coarse score grid: every row's window ends somewhere in
+ `[1, cols]`, so rows sit below, at and past the budget, and the k-th value ties."""
+ g = torch.Generator(device="cuda").manual_seed(seed)
+ scores = torch.randint(0, 5, (rows, cols), generator=g, device="cuda").float()
+ row_starts = torch.zeros(rows, dtype=torch.int32, device="cuda")
+ ends = torch.randint(1, cols + 1, (rows,), generator=g, device="cuda").to(torch.int32)
+ ends[0] = cols
+ ends[1] = K
+ ends[2] = 1
+ indices = torch.full((rows, K), -1, dtype=torch.int32, device="cuda")
+ return scores, row_starts, ends, indices
+
+
+@pytest.mark.parametrize("rows,cols", [(24, 2304), (24, 4096), (8, 1024)])
+def test_rows_equal_the_canonical_definition(rows: int, cols: int) -> None:
+ scores, row_starts, row_ends, indices = _case(rows, cols, seed=rows * cols)
+ U._select(types.SimpleNamespace(block_topk=K), scores, row_starts, row_ends, indices)
+
+ expected = qsa_canonical_selection.reference(scores, row_starts, row_ends, indices.clone())
+ assert torch.equal(indices, expected)
+
+
+def test_live_cols_returns_the_same_rows() -> None:
+ scores, row_starts, row_ends, indices = _case(16, 2304, seed=7)
+ bounded = indices.clone()
+ U._select(types.SimpleNamespace(block_topk=K), scores, row_starts, row_ends, indices)
+ U._select(
+ types.SimpleNamespace(block_topk=K),
+ scores,
+ row_starts,
+ row_ends,
+ bounded,
+ int(row_ends.max()),
+ )
+ assert torch.equal(indices, bounded)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_thop_attention_calls.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_thop_attention_calls.py
similarity index 75%
rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_thop_attention_calls.py
rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_thop_attention_calls.py
index 87584e317806..1b11587e562e 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_thop_attention_calls.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_thop_attention_calls.py
@@ -5,11 +5,11 @@
The attention op reads the pool through `quant_mode` and the two KV scales, and the three
have to agree with the pool's dtype. A call that leaves one at its default reads an fp8 pool
as bf16 (or the reverse) and still returns finite values, so no run notices. The check reads
-`modeling.py` as text: it needs no device and no import of the target.
+`modeling.py` and `vision.py` as text: it needs no device and no import of the target.
-The vision tower's call is the one exception: its attention is the encoder-only mode of the
-op, which has no KV cache, so it passes `kv_scale_*` None and `quant_mode` 0 through
-`_VISION_CALL_CONSTANTS`.
+The vision tower's call (`vision.py`) is the one exception: its attention is the encoder-only
+mode of the op, which has no KV cache, so it passes `kv_scale_*` None and `quant_mode` 0
+through `_VISION_CALL_CONSTANTS`.
"""
from __future__ import annotations
@@ -23,33 +23,34 @@
pytestmark = pytest.mark.cpu_only
-_MODELING = (
+_TARGET = (
Path(_modeling_v2.__file__).resolve().parent
- / "models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py"
+ / "models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default"
)
_KV_STORAGE_ARGS = ("quant_mode", "kv_scale_orig_quant", "kv_scale_quant_orig")
-_TREE = ast.parse(_MODELING.read_text())
+_TREE = ast.parse((_TARGET / "modeling.py").read_text())
+_VISION_TREE = ast.parse((_TARGET / "vision.py").read_text())
-def _thop_attention_calls() -> list[ast.Call]:
+def _thop_attention_calls(tree: ast.Module) -> list[ast.Call]:
return [
node
- for node in ast.walk(_TREE)
+ for node in ast.walk(tree)
if isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and node.func.id == "thop_attention"
]
-def _module_dict(name: str) -> ast.Call:
- """The `dict(...)` a module-level `name = dict(...)` assigns."""
- for node in _TREE.body:
+def _module_dict(tree: ast.Module, name: str) -> ast.Call:
+ """The `dict(...)` a module-level `name = dict(...)` of `tree` assigns."""
+ for node in tree.body:
if isinstance(node, ast.Assign) and any(
isinstance(t, ast.Name) and t.id == name for t in node.targets
):
assert isinstance(node.value, ast.Call), name
return node.value
- raise AssertionError(f"modeling.py has no module-level {name}")
+ raise AssertionError(f"no module-level {name}")
def _keywords(call: ast.Call) -> dict[str, ast.expr]:
@@ -67,13 +68,14 @@ def _is_vision(call: ast.Call) -> bool:
return "_VISION_CALL_CONSTANTS" in _unpacked(call)
-_KV_CACHE_CALLS = [call for call in _thop_attention_calls() if not _is_vision(call)]
+_KV_CACHE_CALLS = _thop_attention_calls(_TREE)
def test_every_kv_cache_call_passes_the_pool_storage() -> None:
# Against going vacuous: the trunk and the draft layer each call it from the prefill
# and the decode target.
assert len(_KV_CACHE_CALLS) >= 4
+ assert not any(_is_vision(call) for call in _KV_CACHE_CALLS)
for call in _KV_CACHE_CALLS:
keywords = _keywords(call)
for name in _KV_STORAGE_ARGS:
@@ -89,7 +91,7 @@ def test_the_unpacked_arguments_carry_no_pool_storage() -> None:
"""`**_CALL_CONSTANTS` and `**step` (each target's `step_args`) are unpacked into every
KV cache call; a storage argument in either would raise as a duplicate keyword or, if the
call dropped its own, silently pin the pool to one storage."""
- constants = {kw.arg for kw in _module_dict("_CALL_CONSTANTS").keywords}
+ constants = {kw.arg for kw in _module_dict(_TREE, "_CALL_CONSTANTS").keywords}
assert constants
assert not constants & set(_KV_STORAGE_ARGS)
for call in _KV_CACHE_CALLS:
@@ -110,11 +112,13 @@ def test_the_unpacked_arguments_carry_no_pool_storage() -> None:
def test_the_vision_tower_runs_without_kv_storage() -> None:
"""The encoder-only call, made once, states a bf16, unscaled, cache-less configuration."""
- constants = {kw.arg: kw.value for kw in _module_dict("_VISION_CALL_CONSTANTS").keywords}
+ constants = {
+ kw.arg: kw.value for kw in _module_dict(_VISION_TREE, "_VISION_CALL_CONSTANTS").keywords
+ }
assert ast.literal_eval(constants["quant_mode"]) == 0
assert ast.literal_eval(constants["kv_scale_orig_quant"]) is None
assert ast.literal_eval(constants["kv_scale_quant_orig"]) is None
- vision = [call for call in _thop_attention_calls() if _is_vision(call)]
- assert len(vision) == 1
+ vision = _thop_attention_calls(_VISION_TREE)
+ assert len(vision) == 1 and _is_vision(vision[0])
for call in vision:
assert not set(_keywords(call)) & set(_KV_STORAGE_ARGS)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision.py
new file mode 100644
index 000000000000..6490c6a81cba
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision.py
@@ -0,0 +1,568 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The default target's vision path on the CPU: the tower's weights, its calls and angle table,
+the splice's bookkeeping and the admission check.
+
+The tower is built over the weight table's own `vis_*` declarations on the meta device, so
+what it reads is what `weights.py` declares. Its per-call body is a stub where a test needs one
+(the vision ops need a GPU), and so is the whole tower where the shell's bookkeeping is under
+test: every row of a stub tower names the item it came from, so the modality dispatch, item
+metadata, chunk slice, frame grouping and the order of concatenated calls run with no engine.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501
+ modeling as M,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501
+ vision as V,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501
+ weights as W,
+)
+from tensorrt_llm._torch.configs import Qwen4ExpVisionConfig
+from tensorrt_llm.inputs.multimodal import (
+ MULTIMODAL_ENCODER_ITEM_METADATA_KEY,
+ MultimodalItemMetadata,
+ MultimodalParams,
+)
+
+pytestmark = pytest.mark.cpu_only
+
+Shell = M.ModelingV2Qwen38FlashNextNvfp4Sm103Default
+
+# Qwen3.8-Flash-Next's vision config.
+_VISION = dict(
+ deepstack_visual_indexes=[],
+ depth=27,
+ dtype="bfloat16",
+ hidden_act="gelu_pytorch_tanh",
+ hidden_size=1152,
+ in_channels=3,
+ intermediate_size=4304,
+ num_heads=16,
+ num_position_embeddings=2304,
+ out_hidden_size=2560,
+ patch_size=16,
+ spatial_merge_size=2,
+ temporal_patch_size=2,
+)
+HIDDEN_OUT = 2560
+
+
+def _tower_over_the_weight_table() -> V.VisionTower:
+ """The tower over every `vis_*` weight the table declares, at the checkpoint's shapes,
+ on meta; `embed` stands in for the activation dtype the tower reads off it."""
+ vcfg = Qwen4ExpVisionConfig(**_VISION)
+ d = SimpleNamespace(
+ hidden=HIDDEN_OUT,
+ dtype=torch.bfloat16,
+ v_depth=vcfg.depth,
+ v_hidden=vcfg.hidden_size,
+ v_inter=vcfg.intermediate_size,
+ v_channels=vcfg.in_channels,
+ v_tpatch=vcfg.temporal_patch_size,
+ v_patch=vcfg.patch_size,
+ v_pos_rows=vcfg.num_position_embeddings,
+ v_merge_dim=vcfg.hidden_size * vcfg.spatial_merge_size**2,
+ )
+ w = {"embed": torch.empty(8, HIDDEN_OUT, dtype=torch.bfloat16, device="meta")}
+ for entry in W.MODEL_WEIGHTS.WEIGHTS:
+ if not entry.name.startswith("vis_"):
+ continue
+ for key, _, _ in W.MODEL_WEIGHTS._targets(entry, d):
+ w[key] = torch.empty(*entry.shape(d), dtype=entry.dtype or d.dtype, device="meta")
+ return V.VisionTower(w, vcfg)
+
+
+def test_the_tower_reads_the_weight_tables_vision_rows_at_the_ops_layouts():
+ """Every view the tower builds exists in the table, and each GEMM operand is the
+ column-major `[K, N]` view `cublas_mm` reads: the packed qkv row is `3 * 1152` wide and
+ the merger reads `2 x 2` patches of 1152."""
+ tower = _tower_over_the_weight_table()
+ assert (tower.depth, tower.heads, tower.head_dim, tower.merge) == (27, 16, 72, 2)
+ assert (tower.pixel_width, tower.merge_dim, tower.grid_side) == (1536, 4608, 48)
+ assert tower.dtype == torch.bfloat16
+ views = tower._w
+ assert tuple(views.patch_w.shape) == (1152, 3, 2, 16, 16)
+ assert tuple(views.pos.shape) == (2304, 1152)
+ assert len(views.blocks) == 27
+ for block in views.blocks:
+ for gemm, (k, n) in (
+ (block.qkv, (1152, 3456)),
+ (block.proj, (1152, 1152)),
+ (block.fc1, (1152, 4304)),
+ (block.fc2, (4304, 1152)),
+ ):
+ assert tuple(gemm.shape) == (k, n) and gemm.stride() == (1, k)
+ assert tuple(block.qkv_b.shape) == (3456,) and tuple(block.fc1_b.shape) == (4304,)
+ assert tuple(views.m_fc1.shape) == (4608, 4608) and views.m_fc1.stride() == (1, 4608)
+ assert tuple(views.m_fc2.shape) == (4608, HIDDEN_OUT) and views.m_fc2.stride() == (1, 4608)
+
+
+# --------------------------------------------------------------- angle table ---
+
+
+def _nested_loop_positions(grid, merge):
+ """Per-row `(row, col)` by four nested loops, in spatial-merge order: the template the
+ rotary entry's test builds its certified tables from."""
+ positions = []
+ for frames, height, width in grid:
+ frame = [
+ (block_row * merge + r, block_col * merge + c)
+ for block_row in range(height // merge)
+ for block_col in range(width // merge)
+ for r in range(merge)
+ for c in range(merge)
+ ]
+ positions.extend(frame * frames)
+ return positions
+
+
+@pytest.mark.parametrize(
+ "grid", [[(1, 62, 90)], [(3, 8, 8), (1, 16, 24)], [(2, 4, 6)]], ids=["62x90", "mixed", "t2"]
+)
+def test_the_angle_table_is_the_rotary_entrys_certified_table(grid):
+ """`_rope_tables` builds, bit for bit, the table `vision_rope_2d_packed_qk` is
+ certified with: positions in spatial-merge order, fp64 angles rounded once to fp32,
+ height angles then width angles, `[rows, 36]` contiguous."""
+ assert V._merge_order_positions(grid, 2).tolist() == [
+ list(p) for p in _nested_loop_positions(grid, 2)
+ ]
+ cos, sin = V._rope_tables(grid, 2, V.VISION_ROPE_THETA, 72, "cpu")
+ rope_dim = 36
+ inv_freq = 1.0 / (
+ V.VISION_ROPE_THETA ** (torch.arange(0, rope_dim, 2, dtype=torch.float64) / rope_dim)
+ )
+ pos = torch.tensor(_nested_loop_positions(grid, 2), dtype=torch.float64)
+ angles = (pos[:, :, None] * inv_freq).flatten(1).to(torch.float32)
+ rows = sum(t * h * w for t, h, w in grid)
+ for got, want in ((cos, angles.cos()), (sin, angles.sin())):
+ assert got.dtype == torch.float32 and got.is_contiguous()
+ assert tuple(got.shape) == (rows, rope_dim)
+ assert torch.equal(got, want)
+
+
+# ---------------------------------------------------------------- tower calls ---
+
+HIDDEN = 3
+MERGE = 2
+PIXEL_WIDTH = 4
+
+
+def _frames(grid):
+ """The item-by-item frame list a grid describes: one `(h, w)` per frame."""
+ return [(h, w) for t, h, w in grid for _ in range(t)]
+
+
+def _random_grid(gen):
+ return [
+ (int(torch.randint(1, 6, (1,), generator=gen)), 2 * h, 2 * w)
+ for h, w in torch.randint(
+ 1, 5, (int(torch.randint(1, 6, (1,), generator=gen)), 2), generator=gen
+ ).tolist()
+ ]
+
+
+def test_frames_are_grouped_whole_and_in_order():
+ # 16-row frames, 32 rows to a call: one item's frames split across calls.
+ assert V._frame_calls([(5, 4, 4)], 32) == [[(2, 4, 4)], [(2, 4, 4)], [(1, 4, 4)]]
+ # A call takes the end of one item and the start of the next.
+ grid = [(1, 8, 4), (3, 4, 4), (1, 2, 2)]
+ assert V._frame_calls(grid, 48) == [[(1, 8, 4), (1, 4, 4)], [(2, 4, 4), (1, 2, 2)]]
+ assert V._frame_calls(grid, 1 << 20) == [grid]
+
+
+def test_a_call_is_filled_to_the_ceiling_and_never_past_it():
+ gen = torch.Generator().manual_seed(0)
+ for _ in range(200):
+ grid = _random_grid(gen)
+ cap = int(torch.randint(64, 300, (1,), generator=gen))
+ calls = V._frame_calls(grid, cap)
+ assert [f for call in calls for f in _frames(call)] == _frames(grid)
+ sizes = [sum(t * h * w for t, h, w in call) for call in calls]
+ assert all(0 < size <= cap for size in sizes)
+ # Greedy: the next call's first frame would not have fit.
+ for size, call in zip(sizes, calls[1:]):
+ h, w = _frames(call)[0]
+ assert size + h * w > cap
+
+
+def test_one_frame_past_the_ceiling_is_refused_and_one_at_it_is_not():
+ assert V._frame_calls([(2, 8, 8)], 64) == [[(1, 8, 8)], [(1, 8, 8)]]
+ with pytest.raises(AssertionError, match="one frame is 8x10 = 80 patch rows"):
+ V._frame_calls([(1, 4, 4), (3, 8, 10)], 64)
+
+
+def test_the_tower_call_ceiling_is_the_vision_entries_largest_cell():
+ """32768 patch rows (8192 merged rows) per call: the ceiling the patch embedding,
+ position interpolation, layer norm, rotary, GEMM and attention cells reach."""
+ assert V.VISION_CALL_MAX_PATCHES == 32768
+
+
+def _stub_tower(monkeypatch, cap):
+ """A `VisionTower` whose per-call body is a stub: merged row `j` of a call carries the
+ pixel value of the call's patch row `4 * j`, which `_pixels` set to the patch's index in
+ the request."""
+ monkeypatch.setattr(V, "VISION_CALL_MAX_PATCHES", cap)
+ tower = object.__new__(V.VisionTower)
+ tower.dtype = torch.float32
+ tower.pixel_width = PIXEL_WIDTH
+ tower.merge = MERGE
+ tower.calls = []
+
+ def run(pixel_values, grid):
+ tower.calls.append((pixel_values, grid))
+ return pixel_values[:: MERGE * MERGE, :1].expand(-1, HIDDEN).clone()
+
+ tower._run = run
+ return tower
+
+
+def _pixels(grid):
+ patches = sum(t * h * w for t, h, w in grid)
+ return torch.arange(patches, dtype=torch.float32)[:, None].expand(-1, PIXEL_WIDTH).clone()
+
+
+def test_the_tower_calls_whole_frames_and_concatenates_their_rows_in_order(monkeypatch):
+ grid = [(1, 8, 4), (3, 4, 4), (2, 2, 6), (1, 2, 2)]
+ tower = _stub_tower(monkeypatch, 48)
+ out = tower(_pixels(grid), torch.tensor(grid))
+ patches = sum(t * h * w for t, h, w in grid)
+ assert out[:, 0].tolist() == [float(4 * j) for j in range(patches // 4)]
+ assert [call[1] for call in tower.calls] == V._frame_calls(grid, 48)
+ start = 0
+ for pixel_values, call in tower.calls:
+ rows = sum(t * h * w for t, h, w in call)
+ assert pixel_values[:, 0].tolist() == [float(i) for i in range(start, start + rows)]
+ assert rows <= 48
+ start += rows
+ assert start == patches
+
+
+def test_a_request_past_one_calls_ceiling_is_served_in_two(monkeypatch):
+ # 33 temporal groups of 32x32 patches (66 frames of 512x512 pixels): 33792 rows, 1024
+ # rows past what a single call may carry.
+ ceiling = V.VISION_CALL_MAX_PATCHES
+ grid = [(33, 32, 32)]
+ tower = _stub_tower(monkeypatch, ceiling)
+ out = tower(_pixels(grid), torch.tensor(grid))
+ assert [call[1] for call in tower.calls] == [[(32, 32, 32)], [(1, 32, 32)]]
+ assert out.shape[0] == 33 * 32 * 32 // 4
+ assert out[:, 0].tolist() == [float(4 * j) for j in range(out.shape[0])]
+
+
+def test_the_tower_refuses_pixels_its_grid_does_not_describe(monkeypatch):
+ tower = _stub_tower(monkeypatch, 64)
+ with pytest.raises(AssertionError, match="describes 16 patch rows"):
+ tower(_pixels([(1, 4, 4)])[:12], torch.tensor([(1, 4, 4)]))
+ with pytest.raises(AssertionError, match="divide by spatial_merge_size 2"):
+ tower(_pixels([(1, 4, 3)]), torch.tensor([(1, 4, 3)]))
+ with pytest.raises(AssertionError, match="activation dtype"):
+ tower(_pixels([(1, 4, 4)]).to(torch.bfloat16), torch.tensor([(1, 4, 4)]))
+
+
+# ------------------------------------------------------------------- splice ---
+
+IMAGE, VIDEO = 1, 2
+
+
+class _Tower:
+ """Stands in for `VisionTower`: every row of item `i` of modality `m` holds `100 * m +
+ i`, so a spliced row says where it came from."""
+
+ merge = MERGE
+
+ def __init__(self):
+ self.calls = 0
+
+ def __call__(self, pixel_values, grid_thw):
+ self.calls += 1
+ modality = int(pixel_values[0, 0])
+ return torch.cat(
+ [
+ torch.full((t * h * w // (MERGE * MERGE), HIDDEN), float(100 * modality + item))
+ for item, (t, h, w) in enumerate(grid_thw.tolist())
+ ]
+ )
+
+
+def _media(modality, grids):
+ """One `multimodal_data` entry; `pixel_values` only carries the stub's modality."""
+ patches = sum(t * h * w for t, h, w in grids)
+ key = "pixel_values" if modality == IMAGE else "pixel_values_videos"
+ grid_key = "image_grid_thw" if modality == IMAGE else "video_grid_thw"
+ return {key: torch.full((patches, 1), float(modality)), grid_key: torch.tensor(grids)}
+
+
+def _params(image=None, video=None, order=None, runtime=None, mrope=True):
+ data = {}
+ if mrope:
+ data["mrope_config"] = {"mrope_position_ids": torch.zeros(3, 1, 8, dtype=torch.int32)}
+ if image is not None:
+ data["image"] = _media(IMAGE, image)
+ if video is not None:
+ data["video"] = _media(VIDEO, video)
+ if order is not None:
+ data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalItemMetadata(item_refs=order)
+ return MultimodalParams(multimodal_data=data, multimodal_runtime=runtime)
+
+
+def _splice(tower, params, num_rows, padding=0):
+ """The tower's rows spliced, as the prefill target's break does, over a prompt whose
+ tokens 1..num_rows are the placeholders and which `padding` rows follow; returns the rows
+ that landed there, one value per row, or None when the step carries no media."""
+ kwargs = {
+ "multimodal_params": [params],
+ "mm_token_indices": torch.arange(1, 1 + num_rows, dtype=torch.int64),
+ }
+ vision = V.multimodal_embeds(tower, kwargs)
+ if vision is None:
+ return None
+ total = num_rows + 2
+ h = torch.full((total + padding, HIDDEN), -1.0)
+ target = SimpleNamespace(_step=SimpleNamespace(vision=vision, num_tokens=total))
+ M.PrefillTarget._splice(target, h)
+ assert torch.count_nonzero(h[total:]) == 0, "padding rows must be zeroed"
+ assert torch.all(h[0] == -1.0) and torch.all(h[total - 1] == -1.0)
+ return h[1 : 1 + num_rows, 0].tolist()
+
+
+def _item(modality, index):
+ return modality, index
+
+
+# Two videos of 4 and 8 rows, two images of 4 and 1.
+VIDEOS = [(2, 2, 4), (2, 4, 4)]
+IMAGES = [(1, 4, 4), (1, 2, 2)]
+MIXED_ORDER = [_item("video", 0), _item("image", 0), _item("video", 1), _item("image", 1)]
+MIXED_ROWS = [200.0] * 4 + [100.0] * 4 + [201.0] * 8 + [101.0]
+
+
+def test_a_text_step_runs_no_tower():
+ assert V.multimodal_embeds(None, {}) is None
+ assert V.multimodal_embeds(None, {"multimodal_params": []}) is None
+ # Position ids alone (a text prompt through a multimodal processor) are not media.
+ assert V.multimodal_embeds(None, {"multimodal_params": [_params()]}) is None
+
+
+def test_the_splice_fills_the_placeholders_and_zeroes_the_padding():
+ """An image request's rows land on its placeholders in one tower call and the padding
+ rows after the step are zeroed (`_splice` checks); a text step writes only the padding."""
+ tower = _Tower()
+ assert _splice(tower, _params(image=IMAGES), 5, padding=3) == [100.0] * 4 + [101.0]
+ assert tower.calls == 1
+ target = SimpleNamespace(_step=SimpleNamespace(vision=None, num_tokens=2))
+ h = torch.full((4, HIDDEN), -1.0)
+ M.PrefillTarget._splice(target, h)
+ assert torch.all(h[:2] == -1.0) and torch.count_nonzero(h[2:]) == 0
+
+
+def test_video_rows_reach_the_placeholders_item_by_item():
+ tower = _Tower()
+ assert _splice(tower, _params(video=VIDEOS), 12) == [200.0] * 4 + [201.0] * 8
+ assert tower.calls == 1
+
+
+def test_a_mixed_request_is_reassembled_in_prompt_order():
+ tower = _Tower()
+ params = _params(image=IMAGES, video=VIDEOS, order=MIXED_ORDER)
+ assert _splice(tower, params, len(MIXED_ROWS)) == MIXED_ROWS
+ assert tower.calls == 2
+
+
+def test_a_mixed_request_needs_its_manifest():
+ with pytest.raises(AssertionError, match="needs item metadata"):
+ _splice(_Tower(), _params(image=IMAGES, video=VIDEOS), len(MIXED_ROWS))
+
+
+def test_a_chunk_takes_its_slice_of_the_prompt_order_rows():
+ runtime = SimpleNamespace(
+ num_cached_mm_tokens=3, num_mm_tokens_in_chunk=6, total_embeds_in_request=len(MIXED_ROWS)
+ )
+ params = _params(image=IMAGES, video=VIDEOS, order=MIXED_ORDER, runtime=runtime)
+ assert _splice(_Tower(), params, 6) == MIXED_ROWS[3:9]
+
+
+def test_a_chunk_without_placeholders_skips_the_tower():
+ runtime = SimpleNamespace(
+ num_cached_mm_tokens=len(MIXED_ROWS),
+ num_mm_tokens_in_chunk=0,
+ total_embeds_in_request=len(MIXED_ROWS),
+ )
+ tower = _Tower()
+ params = _params(image=IMAGES, video=VIDEOS, order=MIXED_ORDER, runtime=runtime)
+ assert _splice(tower, params, 0) is None
+ assert tower.calls == 0
+
+
+def test_a_media_request_needs_the_processors_position_ids():
+ with pytest.raises(AssertionError, match="mrope_position_ids"):
+ _splice(_Tower(), _params(video=VIDEOS, mrope=False), 12)
+
+
+def test_the_placeholder_rows_must_match_the_towers_rows():
+ kwargs = {
+ "multimodal_params": [_params(image=IMAGES)],
+ "mm_token_indices": torch.arange(4, dtype=torch.int64),
+ }
+ with pytest.raises(AssertionError, match="4 placeholder rows and the tower produced 5"):
+ V.multimodal_embeds(_Tower(), kwargs)
+
+
+def test_the_grids_stay_on_the_host():
+ """The executor moves only the pixels and the position ids; the grids the tower and the
+ admission check read on the host stay there."""
+ paths = Shell.multimodal_data_device_paths.fget(None)
+ assert set(paths) == {
+ "image.pixel_values",
+ "video.pixel_values_videos",
+ "mrope_config.mrope_position_ids",
+ "mrope_config.mrope_position_deltas",
+ }
+ assert not any("grid_thw" in path for path in paths)
+
+
+# ---------------------------------------------------------------- admission ---
+
+PATCH = 16
+TEXT_TOKENS = 5
+
+_GEOMETRY = SimpleNamespace(merge=MERGE, patch=PATCH, pixel_width=PIXEL_WIDTH)
+
+
+def _request_media(modality, grids):
+ """A payload whose pixel tensor has the right shape and no storage."""
+ patches = sum(t * h * w for t, h, w in grids)
+ key = "pixel_values" if modality == "image" else "pixel_values_videos"
+ grid_key = "image_grid_thw" if modality == "image" else "video_grid_thw"
+ return {
+ key: torch.empty(patches, PIXEL_WIDTH, device="meta"),
+ grid_key: torch.tensor(grids),
+ }
+
+
+def _request(image=None, video=None, order=None, **overrides):
+ """A request whose prompt is `TEXT_TOKENS` text tokens around the placeholders the media
+ produce, with the position ids and placeholder cumsum the processor emits."""
+ rows = sum(t * h * w for t, h, w in (image or []) + (video or [])) // (MERGE * MERGE)
+ prompt_len = rows + TEXT_TOKENS
+ data = {
+ "mrope_config": {"mrope_position_ids": torch.zeros(3, 1, prompt_len, dtype=torch.int32)},
+ "multimodal_embed_mask_cumsum": torch.cat(
+ [torch.arange(1, rows + 1), torch.full((TEXT_TOKENS,), rows)]
+ ),
+ }
+ if image is not None:
+ data["image"] = _request_media("image", image)
+ if video is not None:
+ data["video"] = _request_media("video", video)
+ if order is not None:
+ data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalItemMetadata(item_refs=order)
+ fields = {"py_multimodal_data": data, "prompt_len": prompt_len}
+ fields.update(overrides)
+ return SimpleNamespace(**fields)
+
+
+def _admit(request):
+ """Through the shell's own hook, as the executor calls it."""
+ Shell.validate_request(SimpleNamespace(vision_tower=_GEOMETRY), request)
+
+
+def test_clean_requests_are_admitted():
+ # A single-modality request needs no manifest.
+ _admit(_request(image=IMAGES))
+ _admit(_request(video=VIDEOS))
+ _admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER))
+
+
+def test_a_request_without_media_costs_nothing():
+ # Not even the other request attributes are read.
+ _admit(SimpleNamespace(py_multimodal_data=None))
+ _admit(SimpleNamespace(py_multimodal_data={}))
+ # A warmup request carries position ids and no media.
+ _admit(
+ SimpleNamespace(
+ py_multimodal_data={"mrope_config": {"mrope_position_ids": torch.zeros(3, 1, 8)}}
+ )
+ )
+
+
+def test_the_longest_admitted_frame_is_the_calls_ceiling(monkeypatch):
+ monkeypatch.setattr(V, "VISION_CALL_MAX_PATCHES", 64)
+ _admit(_request(image=[(1, 8, 8)]))
+ _admit(_request(video=[(3, 4, 16)]))
+ with pytest.raises(ValueError, match="frames of 8x10 = 80 patches"):
+ _admit(_request(image=[(1, 8, 8), (1, 8, 10)]))
+
+
+def test_a_frame_past_the_ceiling_fails_the_request():
+ # 3072x3072 pixels: a 192x192 grid, 36864 patches in one frame.
+ request = _request(video=[(1, 192, 192)])
+ with pytest.raises(ValueError, match=r"video 0 has frames of 192x192 = 36864 patches"):
+ _admit(request)
+ # The message names the pixels and the bound.
+ with pytest.raises(ValueError, match=r"\(3072x3072 pixels\).*at most 32768 patches"):
+ _admit(request)
+
+
+def test_the_frame_bound_is_per_frame_not_per_item():
+ # 33 temporal groups of 32x32 patches: 33792 rows in all, 1024 past one call.
+ _admit(_request(video=[(33, 32, 32)]))
+ with pytest.raises(ValueError, match="video 1 has frames of 192x192"):
+ _admit(_request(video=[(4, 128, 128), (1, 192, 192)]))
+
+
+def test_a_malformed_grid_is_refused():
+ request = _request(image=[(1, 4, 4)])
+ request.py_multimodal_data["image"]["image_grid_thw"] = torch.tensor([(1, 4, 5)])
+ with pytest.raises(ValueError, match="spatial_merge_size 2"):
+ _admit(request)
+ request.py_multimodal_data["image"]["image_grid_thw"] = torch.tensor([1, 4, 4])
+ with pytest.raises(ValueError, match="one \\(t, h, w\\) row per image"):
+ _admit(request)
+
+
+def test_pixels_that_disagree_with_the_grid_are_refused():
+ request = _request(video=VIDEOS)
+ video = request.py_multimodal_data["video"]
+ video["pixel_values_videos"] = torch.empty(5, PIXEL_WIDTH, device="meta")
+ with pytest.raises(ValueError, match="pixel_values_videos is .*describe 48 patch rows"):
+ _admit(request)
+ del video["video_grid_thw"]
+ with pytest.raises(ValueError, match="this target reads 'pixel_values_videos'"):
+ _admit(request)
+
+
+def test_missing_or_mis_sized_position_ids_fail_the_request():
+ request = _request(video=VIDEOS)
+ request.py_multimodal_data["mrope_config"]["mrope_position_ids"] = torch.zeros(3, 1, 40)
+ with pytest.raises(ValueError, match="one position per prompt token"):
+ _admit(request)
+ del request.py_multimodal_data["mrope_config"]
+ with pytest.raises(ValueError, match="mrope_position_ids.*got None"):
+ _admit(request)
+
+
+def test_a_mixed_request_lists_its_items():
+ with pytest.raises(ValueError, match="needs item metadata"):
+ _admit(_request(image=IMAGES, video=VIDEOS))
+ with pytest.raises(ValueError, match="lists image items \\[0\\]; the payload carries 2"):
+ _admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER[:3]))
+ with pytest.raises(ValueError, match="is not an item of this request"):
+ _admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER + [_item("audio", 0)]))
+
+
+def test_the_placeholder_count_is_the_sum_over_both_modalities():
+ request = _request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER)
+ request.py_multimodal_data["multimodal_embed_mask_cumsum"] = torch.arange(1, 14)
+ with pytest.raises(ValueError, match="holds 13 image/video placeholder tokens.*17"):
+ _admit(request)
+ del request.py_multimodal_data["multimodal_embed_mask_cumsum"]
+ _admit(request)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_weight_shards.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_weight_shards.py
new file mode 100644
index 000000000000..a25b790cd6f0
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_weight_shards.py
@@ -0,0 +1,799 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The default target's per-rank weight table, checked against the checkpoint on the CPU.
+
+From `model.safetensors.index.json` and the safetensors headers alone -- no tensor data is
+read -- at each member (tp4, tep4, dep4; dep2, the disaggregated engines' layout, at the end):
+
+* **coverage**: at every rank, with MTP off and on, every checkpoint tensor is either read by
+ the rank's table or predicted unread by `expected_unconsumed`, never both; across a member's
+ ranks every tensor is read by some rank, and each routed or draft expert's tensors by
+ exactly one -- except tp4's routed experts, which every rank reads at its intermediate
+ slice;
+* **shapes**: at ranks 0 and 3, every source read through its cut lies inside the stored
+ tensor, and the cut run through the table's own transform (on meta tensors) has exactly
+ the shape and dtype of the parameter it fills. A rank reads a split tensor at its share
+ only, and every other tensor whole, once;
+* **the modules' shapes**: the per-rank widths the forward is written against, and the head
+ ranges the same rank's K/V and recurrent pools hold (`pool.pool_geometry`).
+
+And, on small synthetic tensors whose values name their rows, that each cut selects the
+rank's own heads in the order the fused operands lay them out, and that tp4's routed
+experts hold the rank's intermediate slice with zero padding at the end of each half.
+"""
+
+from __future__ import annotations
+
+import collections
+import json
+import re
+import struct
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+from utils.llm_data import llm_models_root
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501
+ pool as P,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501
+ weights as WT,
+)
+from tensorrt_llm._torch.configs import Qwen4ExpTextConfig, Qwen4ExpVisionConfig
+from tensorrt_llm.mapping import Mapping
+
+pytestmark = pytest.mark.cpu_only
+
+CHECKPOINT = "Qwen3.8-Flash-Next-NVFP4"
+
+#: Each member's mapping: attention TP4 with the experts split by intermediate (tp4) or by
+#: expert (tep4), and attention DP with expert parallelism (dep4).
+MEMBERS = {
+ "tp4": dict(world_size=4, tp_size=4, moe_tp_size=4, moe_ep_size=1, enable_attention_dp=False),
+ "tep4": dict(world_size=4, tp_size=4, moe_tp_size=1, moe_ep_size=4, enable_attention_dp=False),
+ "dep4": dict(world_size=4, tp_size=4, moe_tp_size=1, moe_ep_size=4, enable_attention_dp=True),
+}
+#: Both engines of the disaggregated-serving test: attention DP over two ranks, the routed
+#: experts in two 256-expert windows. It serves without speculation.
+DEP2 = dict(world_size=2, tp_size=2, moe_tp_size=1, moe_ep_size=2, enable_attention_dp=True)
+_LAYOUTS = {**MEMBERS, "dep2": DEP2}
+RANKS = (0, 1, 2, 3)
+SHAPE_RANKS = (0, 3)
+
+_STORED_DTYPES = {
+ "BF16": torch.bfloat16,
+ "F32": torch.float32,
+ "U8": torch.uint8,
+ "F8_E4M3": torch.float8_e4m3fn,
+ "I64": torch.int64,
+}
+
+#: A routed or draft expert's tensors, as opposed to its per-tensor scalars (which every
+#: rank reads, for the whole routing space).
+_EXPERT_TENSOR = re.compile(
+ r"\.mlp\.experts\.\d+\.\w+_proj\.(weight|weight_scale|weight_scale_inv)$"
+)
+_TRUNK_EXPERT_TENSOR = re.compile(
+ r"^model\.language_model\.layers\.\d+\.mlp\.experts\.\d+\.\w+_proj\.(weight|weight_scale)$"
+)
+
+
+# ------------------------------------------------------------------ fixtures --
+
+
+@pytest.fixture(scope="module")
+def checkpoint() -> Path:
+ root = llm_models_root()
+ path = None if root is None else root / CHECKPOINT
+ if path is None or not (path / "model.safetensors.index.json").is_file():
+ pytest.skip(f"{CHECKPOINT} is not under LLM_MODELS_ROOT")
+ return path
+
+
+@pytest.fixture(scope="module")
+def stored(checkpoint) -> dict[str, tuple[tuple[int, ...], torch.dtype]]:
+ """Every checkpoint tensor's stored shape and dtype, from the index and the headers."""
+ index = json.loads((checkpoint / "model.safetensors.index.json").read_text())["weight_map"]
+ specs = {}
+ for file in sorted(set(index.values())):
+ with open(checkpoint / file, "rb") as fh:
+ (length,) = struct.unpack(" Mapping:
+ return Mapping(rank=rank, gpus_per_node=4, pp_size=1, **_LAYOUTS[member])
+
+
+def _dims(configs, member: str, rank: int, mtp: bool = True, prefill_copy: bool = False):
+ text, vision = configs
+ return WT.rank_dims(
+ text,
+ vision,
+ _mapping(member, rank),
+ mtp_enabled=mtp,
+ trunk_moe_backend="TRTLLM",
+ moe_prefill_copy=prefill_copy,
+ dtype=torch.bfloat16,
+ )
+
+
+def _attention_tp(member: str) -> int:
+ layout = _LAYOUTS[member]
+ return 1 if layout["enable_attention_dp"] else layout["tp_size"]
+
+
+def _lm_head_rows(member: str, vocab: int) -> int:
+ """The rows of the lm_head the shell builds: vocab-parallel under attention TP, whole
+ under attention DP (`DecoderModelForCausalLM`)."""
+ return vocab // _attention_tp(member)
+
+
+def _rows_per_shard(stored, d) -> int:
+ shapes = {stored[key][0] for key in WT.ngram_shard_keys(d)}
+ assert len(shapes) == 1, shapes
+ return next(iter(shapes))[0]
+
+
+def _rank_manifest(stored, d, member: str) -> dict[str, list[tuple]]:
+ """The rank's manifest as `load` builds it, with the n-gram table and the shell's lm_head."""
+ manifest = WT.MODEL_WEIGHTS.manifest(d)
+ manifest["ple_table"] = WT.ngram_table_sources(d, _rows_per_shard(stored, d))
+ cut, pad = WT.lm_head_share(_lm_head_rows(member, d.vocab), d.vocab, d.tp_rank)
+ manifest["lm_head.weight"] = [("lm_head.weight", None, cut, pad)]
+ return manifest
+
+
+def _keys_and_cuts(ckpt_key, cut):
+ if isinstance(ckpt_key, tuple):
+ return list(zip(ckpt_key, cut if cut is not None else (None,) * len(ckpt_key)))
+ return [(ckpt_key, cut)]
+
+
+def _claimed(manifest) -> set:
+ return {
+ key
+ for sources in manifest.values()
+ for ckpt_key, _, cut, _ in sources
+ for key, _ in _keys_and_cuts(ckpt_key, cut)
+ }
+
+
+def _cut_numel(shape: tuple[int, ...], cut) -> int:
+ """Elements of a `shape` tensor `cut` reads, after checking each slice lies inside it."""
+ if cut is None:
+ return int(torch.Size(shape).numel())
+ extents = list(shape)
+ for axis, s in enumerate(cut):
+ assert isinstance(s, slice) and s.step in (None, 1), (shape, cut)
+ start = 0 if s.start is None else s.start
+ stop = shape[axis] if s.stop is None else s.stop
+ assert 0 <= start < stop <= shape[axis], f"cut {cut} leaves the stored {shape}"
+ extents[axis] = stop - start
+ return int(torch.Size(extents).numel())
+
+
+class _HeaderEntry:
+ """A checkpoint tensor known by its header: indexing it returns a meta tensor of the
+ indexed shape, and nothing is ever read."""
+
+ def __init__(self, shape, dtype):
+ self.shape, self.dtype = torch.Size(shape), dtype
+ self.ndim = len(shape)
+
+ def __getitem__(self, index):
+ return torch.empty(self.shape, dtype=self.dtype, device="meta")[index]
+
+
+def _share(member: str, key: str) -> float:
+ """The fraction of `key` a rank of `member` holds, worked out here from the member alone:
+ attention q/o by q head and k/v by KV head (two heads over the attention TP ranks), the
+ Gated-DeltaNet tensors by head over the same ranks, the shell's lm_head by vocabulary,
+ and at tp4 each routed expert's tensors by intermediate. Everything else is read whole."""
+ tp = _attention_tp(member)
+ if member == "tp4" and _TRUNK_EXPERT_TENSOR.match(key):
+ return 1 / MEMBERS[member]["moe_tp_size"]
+ if re.search(r"\.self_attn\.(q_proj|o_proj)\.weight$", key) or key == "lm_head.weight":
+ return 1 / tp
+ if re.search(r"\.self_attn\.(k_proj|v_proj)\.weight$", key):
+ return 1 / min(tp, 2)
+ if re.search(
+ r"\.linear_attn\.((in_proj_qkv|in_proj_z|in_proj_a|in_proj_b|conv1d|out_proj)\.weight"
+ r"|A_log|dt_bias)$",
+ key,
+ ):
+ return 1 / tp
+ return 1.0
+
+
+# ----------------------------------------------------------------- coverage --
+
+
+@pytest.mark.parametrize("mtp", [False, True], ids=["no-mtp", "mtp"])
+@pytest.mark.parametrize("member", list(MEMBERS))
+def test_every_checkpoint_tensor_is_claimed(stored, configs, member, mtp):
+ everything = set(stored)
+ read_by = collections.Counter()
+ for rank in RANKS:
+ d = _dims(configs, member, rank, mtp)
+ claimed = _claimed(_rank_manifest(stored, d, member))
+ unread = WT.MODEL_WEIGHTS.expected_unconsumed(d, stored.keys())
+ assert claimed <= everything, f"rank {rank} reads keys the checkpoint lacks: " + str(
+ sorted(claimed - everything)[:8]
+ )
+ assert not claimed & unread, f"rank {rank}: read and predicted unread: " + str(
+ sorted(claimed & unread)[:8]
+ )
+ assert claimed | unread == everything, (
+ f"rank {rank}: neither read nor predicted {sorted(everything - claimed - unread)[:8]}"
+ )
+ read_by.update(claimed)
+
+ mtp_keys = set() if mtp else {k for k in everything if k.startswith("mtp.")}
+ assert set(read_by) == everything - mtp_keys, "tensors no rank reads: " + str(
+ sorted(everything - mtp_keys - set(read_by))[:8]
+ )
+
+ def readers(key):
+ # tp4: every rank reads every routed expert, at its slice (checked by the shape test).
+ return len(RANKS) if member == "tp4" and _TRUNK_EXPERT_TENSOR.match(key) else 1
+
+ wrong = sorted(k for k in read_by if _EXPERT_TENSOR.search(k) and read_by[k] != readers(k))
+ assert not wrong, f"expert tensors not read by their expected rank count: {wrong[:8]}"
+
+
+# -------------------------------------------------------------------- shapes --
+
+
+@pytest.mark.parametrize("prefill_copy", [False, True], ids=["trtllm-gen", "with-cutedsl-copy"])
+@pytest.mark.parametrize("rank", SHAPE_RANKS)
+@pytest.mark.parametrize("member", list(MEMBERS))
+def test_rank_operands_have_their_parameter_shapes(stored, configs, member, rank, prefill_copy):
+ """Every source, read through its cut and run through its transform on meta tensors,
+ fills its parameter exactly; the rank reads each tensor at its share, and once per
+ layout it is kept in (the routed experts twice with the CuteDSL prefill copy)."""
+ d = _dims(configs, member, rank, mtp=True, prefill_copy=prefill_copy)
+ with torch.device("meta"):
+ params = dict(WT.MODEL_WEIGHTS.declare(d))
+ # What `load` adds beside the table: the n-gram table and the shell's lm_head.
+ params["ple_table"] = torch.empty(
+ d.ple_shards * _rows_per_shard(stored, d), d.ple_head_dim, dtype=d.ple_table_dtype
+ )
+ params["lm_head.weight"] = torch.empty(
+ _lm_head_rows(member, d.vocab), d.hidden, dtype=torch.bfloat16
+ )
+ entries = {key: _HeaderEntry(*spec) for key, spec in stored.items()}
+ manifest = _rank_manifest(stored, d, member)
+ assert manifest.keys() == params.keys()
+
+ read = collections.Counter()
+ results: dict = {}
+ for param_key, sources in manifest.items():
+ param = params[param_key]
+ for ckpt_key, index, cut, transform in sources:
+ key_cuts = _keys_and_cuts(ckpt_key, cut)
+ for key, key_cut in key_cuts:
+ read[key] += _cut_numel(stored[key][0], key_cut)
+ dst = param if index is None else param[index]
+ # Sources that agree on the stored tensors, their cuts and the transform agree
+ # on the result: run each distinct one once.
+ signature = (transform, tuple((stored[k], c) for k, c in key_cuts))
+ if signature not in results:
+ src = WT.MODEL_WEIGHTS.read(entries, ckpt_key, cut, transform, d)
+ results[signature] = (tuple(src.shape), src.dtype)
+ assert results[signature] == (tuple(dst.shape), dst.dtype), (
+ param_key,
+ ckpt_key,
+ results[signature],
+ (tuple(dst.shape), dst.dtype),
+ )
+
+ def copies(key):
+ return 2 if prefill_copy and _TRUNK_EXPERT_TENSOR.match(key) else 1
+
+ wrong = {
+ key: (n, _share(member, key), copies(key))
+ for key, n in read.items()
+ if n != round(torch.Size(stored[key][0]).numel() * _share(member, key)) * copies(key)
+ }
+ assert not wrong, f"tensors not read at the rank's share, once: {sorted(wrong.items())[:8]}"
+
+
+#: Per-rank shapes of the attention and Gated-DeltaNet weights the forward is written
+#: against: (under attention TP, under attention DP).
+_SPLIT_SHAPES = {
+ # [q 6x256 | k 256 | v 256 | gate 6x256 | index 4x128 + 128 | pad 128]
+ "l3_attn_in_w": ((4352, 2560), (14080, 2560)),
+ "l3_o_w": ((2560, 1536), (2560, 6144)),
+ # [q 4x128 | k 4x128 | v 12x128 | z 12x128 | a 12 | b 12]
+ "l0_in_proj_w": ((4120, 2560), (16480, 2560)),
+ "l0_conv_w": ((2560, 1, 4), (10240, 1, 4)),
+ "l0_A_log": ((12,), (48,)),
+ "l0_dt_bias": ((12,), (48,)),
+ "l0_out_w": ((2560, 1536), (2560, 6144)),
+ "l48_mtp_q_w": ((3072, 2560), (12288, 2560)),
+ "l48_mtp_k_w": ((256, 2560), (512, 2560)),
+ "l48_mtp_v_w": ((256, 2560), (512, 2560)),
+ "l48_mtp_o_w": ((2560, 1536), (2560, 6144)),
+}
+
+#: The same at every member and rank: replicated weights and the draft layer's 128-expert
+#: FP8 window.
+_COMMON_SHAPES = {
+ "l3_idx_qn": (128,),
+ "l0_gnorm": (128,),
+ "l0_router": (512, 2560),
+ "l0_sh_gu": (2 * 640 + 8, 2560),
+ "l0_sh_dn": (2560, 640),
+ "l0_e_isc1": (512,),
+ "l0_ahc_down": (320, 10240),
+ "fhc_up": (10240, 320),
+ "embed": (248320, 2560),
+ "ple_key_w": (10240, 2560),
+ "l48_mtp_idx_w": (640, 2560),
+ "l48_mtp_fc_e": (2560, 2560),
+ "l48_mtp_fc1_w": (128, 1280, 2560),
+ "l48_mtp_fc1_s": (128, 10, 20),
+ "l48_mtp_fc2_w": (128, 2560, 640),
+ "l48_mtp_fc2_s": (128, 20, 5),
+}
+
+#: The routed experts a rank holds: a 128-expert window of whole experts (tep4, dep4), or
+#: all 512 at a 160-row intermediate slice stored 256 wide (tp4).
+_EXPERT_SHAPES = {
+ "tep4": {
+ "l0_fc1_w": (128, 1280, 1280),
+ "l0_fc1_s": (128, 1280, 160),
+ "l0_fc2_w": (128, 2560, 320),
+ "l0_fc2_s": (128, 2560, 40),
+ },
+ "tp4": {
+ "l0_fc1_w": (512, 512, 1280),
+ "l0_fc1_s": (512, 512, 160),
+ "l0_fc2_w": (512, 2560, 128),
+ "l0_fc2_s": (512, 2560, 16),
+ },
+}
+_EXPERT_SHAPES["dep4"] = _EXPERT_SHAPES["tep4"]
+
+
+@pytest.mark.parametrize("rank", SHAPE_RANKS)
+@pytest.mark.parametrize("member", list(MEMBERS))
+def test_rank_shapes_are_what_the_modules_expect(stored, configs, member, rank):
+ d = _dims(configs, member, rank)
+ with torch.device("meta"):
+ params = WT.MODEL_WEIGHTS.declare(d)
+ dp = MEMBERS[member]["enable_attention_dp"]
+ for key, shapes in _SPLIT_SHAPES.items():
+ assert tuple(params[key].shape) == shapes[dp], (key, tuple(params[key].shape))
+ for key, shape in _COMMON_SHAPES.items():
+ assert tuple(params[key].shape) == shape, (key, tuple(params[key].shape))
+ for key, shape in _EXPERT_SHAPES[member].items():
+ assert tuple(params[key].shape) == shape, (key, tuple(params[key].shape))
+ window = range(512) if member == "tp4" else range(128 * rank, 128 * (rank + 1))
+ experts = [s[0] for s in WT.MODEL_WEIGHTS.manifest(d)["l0_fc2_w"]]
+ assert experts == [
+ f"model.language_model.layers.0.mlp.experts.{e}.down_proj.weight" for e in window
+ ]
+ drafts = [s[0] for s in WT.MODEL_WEIGHTS.manifest(d)["l48_mtp_fc2_w"]]
+ assert drafts == [
+ f"mtp.layers.0.mlp.experts.{e}.down_proj.weight"
+ for e in range(128 * rank, 128 * (rank + 1))
+ ]
+
+ # The heads this rank's weights hold are the heads its pools hold.
+ g = P.pool_geometry(
+ configs[0],
+ _mapping(member, rank),
+ tokens_per_block=32,
+ kv_cache_quant_algo=None,
+ dtype=torch.bfloat16,
+ )
+ assert (d.heads_kv, d.kv_head_start) == (g.heads_kv, g.kv_head_start)
+ assert (d.num_k_heads, d.k_head_start) == (g.num_k_heads, g.k_head_start)
+ assert (d.num_v_heads, d.v_head_start) == (g.num_v_heads, g.v_head_start)
+ assert params["l0_conv_w"].shape[0] == g.conv_dim
+ assert params["l0_A_log"].shape[0] == g.num_v_heads
+
+
+def test_lm_head_reads_the_rows_the_shell_built():
+ vocab = 248320
+ for rank in RANKS:
+ cut, _ = WT.lm_head_share(vocab // 4, vocab, rank)
+ assert cut == (slice(rank * 62080, (rank + 1) * 62080),)
+ cut, _ = WT.lm_head_share(vocab, vocab, rank)
+ assert cut == (slice(0, vocab),)
+
+
+# ---------------------------------------------------------------------- load --
+
+
+def _stub_model(configs, member: str, rank: int, mtp: bool):
+ """The shell and core `load` reads, on the meta device: the core's configs, mapping and
+ engine choices, its declared parameters, and the shell's lm_head."""
+ text, vision = configs
+ mapping = _mapping(member, rank)
+ core = SimpleNamespace(
+ model_config=SimpleNamespace(
+ pretrained_config=text,
+ mapping=mapping,
+ spec_config=object() if mtp else None,
+ torch_dtype=torch.bfloat16,
+ ),
+ vision_config=vision,
+ trunk_moe_backend="TRTLLM",
+ moe_prefill_copy=False,
+ ple_host_offload=False,
+ )
+ d = WT.MODEL_WEIGHTS.dims(core)
+ with torch.device("meta"):
+ core.w = WT.MODEL_WEIGHTS.declare(d)
+ lm_head = torch.nn.Parameter(
+ torch.empty(_lm_head_rows(member, d.vocab), d.hidden, dtype=torch.bfloat16),
+ requires_grad=False,
+ )
+
+ def named_parameters():
+ yield from ((f"model.w.{k}", v) for k, v in core.w.items())
+ yield "lm_head.weight", lm_head
+
+ return SimpleNamespace(
+ model=core, lm_head=SimpleNamespace(weight=lm_head), named_parameters=named_parameters
+ )
+
+
+def _meta_checkpoint(stored) -> dict:
+ """Every checkpoint tensor as its header: 0-dim scalars as meta tensors, the rest as
+ entries that index into meta tensors."""
+ return {
+ key: torch.empty((), dtype=dtype, device="meta")
+ if not shape
+ else _HeaderEntry(shape, dtype)
+ for key, (shape, dtype) in stored.items()
+ }
+
+
+@pytest.mark.parametrize(
+ "member, rank, mtp",
+ [("tep4", 3, True), ("dep4", 0, True), ("tep4", 1, False), ("tp4", 2, True)],
+)
+def test_load_fills_every_parameter_and_covers_the_checkpoint(
+ stored, configs, member, rank, mtp, monkeypatch
+):
+ """`load` end to end at the real shapes, on meta tensors: the n-gram table declared from
+ its shards, every manifest row filled through its cut, the shell's lm_head rows, and
+ both coverage asserts."""
+ monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None)
+ monkeypatch.setattr(torch.cuda, "empty_cache", lambda *a, **k: None)
+ model = _stub_model(configs, member, rank, mtp)
+ WT.MODEL_WEIGHTS.load(model, _meta_checkpoint(stored))
+ assert tuple(model.model.w["ple_table"].shape) == (320_001_536, 160)
+
+
+# --------------------------------------------------- the cuts pick the heads --
+
+_HIDDEN = 2 # the projections' input axis, shrunk: no cut touches it
+_FAKE_OFFSETS = dict(q=1, k=2, v=3, idx=4, o=5, qkv=10, z=11, a=12, b=13, conv=14, a_log=15,
+ dtb=16, out=17) # fmt: skip
+
+
+def _ids(tag: str, rows: int, *trailing: int) -> torch.Tensor:
+ """`rows` rows whose every value is `tag`'s offset plus the row's index."""
+ ids = torch.arange(rows, dtype=torch.int64) + _FAKE_OFFSETS[tag] * 10**7
+ return ids.view(rows, *([1] * len(trailing))).expand(rows, *trailing).contiguous()
+
+
+def _col_ids(tag: str, cols: int) -> torch.Tensor:
+ return _ids(tag, cols, _HIDDEN).t().contiguous()
+
+
+def _fake_checkpoint(stored) -> dict[str, torch.Tensor]:
+ """Attention layer 3, Gated-DeltaNet layer 0 and the MTP attention, at their stored row
+ counts and a two-wide input axis."""
+ fake = {}
+
+ def rows_of(key):
+ return stored[key][0][0]
+
+ for p in ("model.language_model.layers.3.self_attn", "mtp.layers.0.self_attn"):
+ for tag, name in (("q", "q_proj"), ("k", "k_proj"), ("v", "v_proj")):
+ fake[f"{p}.{name}.weight"] = _ids(tag, rows_of(f"{p}.{name}.weight"), _HIDDEN)
+ key = f"{p}.indexer.index_qk_proj.weight"
+ fake[key] = _ids("idx", rows_of(key), _HIDDEN)
+ fake[f"{p}.o_proj.weight"] = _col_ids("o", stored[f"{p}.o_proj.weight"][0][1])
+ p = "model.language_model.layers.0.linear_attn"
+ for tag, name in (("qkv", "in_proj_qkv"), ("z", "in_proj_z"), ("a", "in_proj_a"),
+ ("b", "in_proj_b")): # fmt: skip
+ fake[f"{p}.{name}.weight"] = _ids(tag, rows_of(f"{p}.{name}.weight"), _HIDDEN)
+ fake[f"{p}.conv1d.weight"] = _ids("conv", rows_of(f"{p}.conv1d.weight"), 1, 4)
+ fake[f"{p}.A_log"] = _ids("a_log", rows_of(f"{p}.A_log"))
+ fake[f"{p}.dt_bias"] = _ids("dtb", rows_of(f"{p}.dt_bias"))
+ fake[f"{p}.out_proj.weight"] = _col_ids("out", stored[f"{p}.out_proj.weight"][0][1])
+ return fake
+
+
+def _assemble(d, fake, sources) -> torch.Tensor:
+ """A parameter filled from its sources the way `load` fills it."""
+ if len(sources) == 1 and sources[0][1] is None:
+ ckpt_key, _, cut, transform = sources[0]
+ return WT.MODEL_WEIGHTS.read(fake, ckpt_key, cut, transform, d)
+ parts = [
+ (index, WT.MODEL_WEIGHTS.read(fake, ckpt_key, cut, transform, d))
+ for ckpt_key, index, cut, transform in sources
+ ]
+ rows = max(index[0].stop for index, _ in parts)
+ out = torch.full((rows, *parts[0][1].shape[1:]), -1, dtype=parts[0][1].dtype)
+ for index, part in parts:
+ out[index] = part
+ return out
+
+
+def _expected_rows(tag: str, starts) -> torch.Tensor:
+ return torch.tensor([_FAKE_OFFSETS[tag] * 10**7 + r for r in starts], dtype=torch.int64)
+
+
+def _block(start: int, width: int) -> list[int]:
+ return list(range(start, start + width))
+
+
+@pytest.mark.parametrize("rank", RANKS)
+@pytest.mark.parametrize("member", ["tep4", "dep4"])
+def test_cuts_select_the_rank_heads(stored, configs, member, rank):
+ """tp4 cuts attention and Gated-DeltaNet exactly as tep4 does: the members differ only
+ in the routed experts."""
+ d = _dims(configs, member, rank)
+ fake = _fake_checkpoint(stored)
+ manifest = WT.MODEL_WEIGHTS.manifest(d)
+ tp = _attention_tp(member)
+ # Under attention DP each rank holds every head: the split is over one rank.
+ r = rank if tp > 1 else 0
+ q_heads = _block(r * 24 // tp, 24 // tp)
+ kv_heads = [r * 2 // tp] if tp > 1 else [0, 1]
+ k_heads = _block(r * 16 // tp, 16 // tp)
+ v_heads = _block(r * 48 // tp, 48 // tp)
+
+ def rows(heads, width, base=0):
+ return [base + h * width + j for h in heads for j in range(width)]
+
+ q_rows = [h * 512 + j for h in q_heads for j in range(256)]
+ gate_rows = [h * 512 + 256 + j for h in q_heads for j in range(256)]
+ attn_in = _assemble(d, fake, manifest["l3_attn_in_w"])
+ expected = torch.cat(
+ (
+ _expected_rows("q", q_rows),
+ _expected_rows("k", rows(kv_heads, 256)),
+ _expected_rows("v", rows(kv_heads, 256)),
+ _expected_rows("q", gate_rows),
+ _expected_rows("idx", range(640)),
+ torch.zeros(128, dtype=torch.int64),
+ )
+ )
+ assert torch.equal(attn_in[:, 0], expected) and torch.equal(attn_in[:, 1], expected)
+
+ o_cols = _expected_rows("o", rows(q_heads, 256))
+ for key in ("l3_o_w", "l48_mtp_o_w"):
+ assert torch.equal(_assemble(d, fake, manifest[key])[0], o_cols), key
+ assert torch.equal(_assemble(d, fake, manifest["l48_mtp_q_w"])[:, 0],
+ _expected_rows("q", rows(q_heads, 512))) # fmt: skip
+ assert torch.equal(_assemble(d, fake, manifest["l48_mtp_k_w"])[:, 0],
+ _expected_rows("k", rows(kv_heads, 256))) # fmt: skip
+
+ def qkv(tag):
+ return torch.cat(
+ (
+ _expected_rows(tag, rows(k_heads, 128)),
+ _expected_rows(tag, rows(k_heads, 128, 2048)),
+ _expected_rows(tag, rows(v_heads, 128, 4096)),
+ )
+ )
+
+ in_proj = torch.cat(
+ (
+ qkv("qkv"),
+ _expected_rows("z", rows(v_heads, 128)),
+ _expected_rows("a", v_heads),
+ _expected_rows("b", v_heads),
+ )
+ )
+ assert torch.equal(_assemble(d, fake, manifest["l0_in_proj_w"])[:, 0], in_proj)
+ assert torch.equal(_assemble(d, fake, manifest["l0_conv_w"])[:, 0, 0], qkv("conv"))
+ assert torch.equal(_assemble(d, fake, manifest["l0_A_log"]), _expected_rows("a_log", v_heads))
+ assert torch.equal(_assemble(d, fake, manifest["l0_dt_bias"]), _expected_rows("dtb", v_heads))
+ assert torch.equal(_assemble(d, fake, manifest["l0_out_w"])[0],
+ _expected_rows("out", rows(v_heads, 128))) # fmt: skip
+
+
+# ------------------------------------------ tp4: the padded intermediate slice --
+
+
+def _blk32_gather(rows: int) -> torch.Tensor:
+ """The 32-row block shuffle's gather index, written out here: inside each block of 32
+ rows, source row `4u + v` lands at destination row `8v + u`."""
+ j = torch.arange(32)
+ idx = torch.empty(32, dtype=torch.long)
+ idx[(j % 4) * 8 + j // 4] = j
+ return idx.repeat(rows // 32) + torch.arange(rows // 32).repeat_interleave(32) * 32
+
+
+def _interleave_gather(rows: int) -> torch.Tensor:
+ """`[up | gate]` -> up0, gate0, up1, gate1, ..."""
+ p = torch.empty(rows, dtype=torch.long)
+ p[0::2] = torch.arange(rows // 2)
+ p[1::2] = torch.arange(rows // 2, rows)
+ return p
+
+
+def _swizzle_flat(s: torch.Tensor) -> torch.Tensor:
+ """The 128x4 block-scale swizzle of one `[M, C]` matrix, flat: scale `(m, c)` lands at
+ `(m // 128) * 512 * (C // 4) + (c // 4) * 512 + (m % 32) * 16 + ((m % 128) // 32) * 4
+ + c % 4`."""
+ m, c = s.shape
+ v = s.reshape(m // 128, 4, 32, c // 4, 4).permute(0, 3, 2, 1, 4)
+ return v.reshape(-1).contiguous()
+
+
+def _zero_rows(x: torch.Tensor, rows: int) -> torch.Tensor:
+ return torch.cat((x, x.new_zeros(rows - x.shape[0], *x.shape[1:])), 0)
+
+
+def _zero_cols(x: torch.Tensor, cols: int) -> torch.Tensor:
+ return torch.cat((x, x.new_zeros(x.shape[0], cols - x.shape[1])), 1)
+
+
+@pytest.mark.parametrize("expert", [0, 511])
+@pytest.mark.parametrize("rank", RANKS)
+def test_tp4_experts_hold_the_rank_slice_zero_padded_per_half(configs, rank, expert, monkeypatch):
+ """At tp4 rank `r` reads rows `[160r, 160r + 160)` of an expert's `up_proj` and
+ `gate_proj` (codes and block scales) and the same columns of its `down_proj`; each half is
+ zero-extended to 256 rows (FC2 to 256 columns) *before* the `[up ; gate]` concatenation,
+ then interleaved, block-shuffled and swizzled as at the whole width. The bytes are
+ nonzero, so every zero in the operand is padding. The scale swizzle is the formula
+ `fp4_block_scale_moe_runner`'s test writes out, standing in for the native op on
+ the CPU; the shrunk axes (FC1's input bytes, FC2's output rows) take no cut."""
+ monkeypatch.setattr(WT, "_swizzle", _swizzle_flat)
+ d = _dims(configs, "tp4", rank, mtp=False)
+ assert (d.expert_inter_share, d.expert_inter_start, d.expert_inter) == (160, 160 * rank, 256)
+ gen = torch.Generator().manual_seed(rank * 1000 + expert)
+
+ def nonzero(*shape):
+ return torch.randint(1, 256, shape, generator=gen, dtype=torch.uint8)
+
+ p = f"model.language_model.layers.0.mlp.experts.{expert}"
+ fake = {
+ f"{p}.up_proj.weight": nonzero(640, 4),
+ f"{p}.gate_proj.weight": nonzero(640, 4),
+ f"{p}.up_proj.weight_scale": nonzero(640, 4).view(torch.float8_e4m3fn),
+ f"{p}.gate_proj.weight_scale": nonzero(640, 4).view(torch.float8_e4m3fn),
+ f"{p}.down_proj.weight": nonzero(128, 320),
+ f"{p}.down_proj.weight_scale": nonzero(128, 40).view(torch.float8_e4m3fn),
+ }
+ manifest = WT.MODEL_WEIGHTS.manifest(d)
+
+ def operand(param_key):
+ (row,) = [s for s in manifest[param_key] if s[1] == (expert,)]
+ ckpt_key, _, cut, transform = row
+ return WT.MODEL_WEIGHTS.read(fake, ckpt_key, cut, transform, d).view(torch.uint8)
+
+ rows = slice(160 * rank, 160 * rank + 160)
+
+ def fc1(kind):
+ up = fake[f"{p}.up_proj.{kind}"].view(torch.uint8)[rows]
+ gate = fake[f"{p}.gate_proj.{kind}"].view(torch.uint8)[rows]
+ both = torch.cat((_zero_rows(up, 256), _zero_rows(gate, 256)), 0)
+ return both[_interleave_gather(512)[_blk32_gather(512)]]
+
+ assert torch.equal(operand("l0_fc1_w"), fc1("weight"))
+ assert torch.equal(operand("l0_fc1_s"), _swizzle_flat(fc1("weight_scale")).view(512, 4))
+
+ down = fake[f"{p}.down_proj.weight"][:, 80 * rank : 80 * rank + 80]
+ assert torch.equal(operand("l0_fc2_w"), _zero_cols(down, 128)[_blk32_gather(128)])
+ down_s = fake[f"{p}.down_proj.weight_scale"].view(torch.uint8)[:, 10 * rank : 10 * rank + 10]
+ expected = _swizzle_flat(_zero_cols(down_s, 16)[_blk32_gather(128)]).view(128, 16)
+ assert torch.equal(operand("l0_fc2_s"), expected)
+
+ # The padding is where the runner's `intermediate_size` = 256 puts it: the last 96 rows
+ # of each half and the last 96 of FC2's 256 input columns, nothing else.
+ w1 = operand("l0_fc1_w")[torch.argsort(_interleave_gather(512)[_blk32_gather(512)])]
+ for half in (w1[:256], w1[256:]):
+ assert bool((half[:160] != 0).all()) and not bool(half[160:].any())
+
+
+@pytest.mark.parametrize("member", list(MEMBERS))
+def test_the_stored_intermediate_follows_the_engines_storage_rule(configs, member):
+ """The width a rank stores each expert's intermediate at is the one the engine's own
+ NVFP4 trtllm-gen MoE stores that shard at, and the width the runner is certified at."""
+ from tensorrt_llm._torch.moe.fused_moe.quantization import NVFP4TRTLLMGenFusedMoEMethod
+
+ d = _dims(configs, member, 1)
+ align, _ = NVFP4TRTLLMGenFusedMoEMethod.resolve_alignments(d.hidden, d.expert_inter_share)
+ assert d.expert_inter == -(-d.expert_inter_share // align) * align
+ assert (d.expert_inter_share, d.expert_inter) == ((160, 256) if member == "tp4" else (640, 640))
+
+
+# ---------------------------------------------------------------------- dep2 --
+
+
+def test_dep2_ranks_read_the_checkpoint_in_two_256_expert_windows(stored, configs):
+ """dep2 reads what dep4 reads, except the routed experts: rank `r` reads the window
+ `[256r, 256r + 256)`, so across the two ranks every tensor is read, each expert's by one
+ rank, and the `mtp.*` subtree by none (no speculation at dep2)."""
+ everything = set(stored)
+ mtp_keys = {k for k in everything if k.startswith("mtp.")}
+ read_by = collections.Counter()
+ for rank in range(DEP2["world_size"]):
+ d = _dims(configs, "dep2", rank, mtp=False)
+ claimed = _claimed(_rank_manifest(stored, d, "dep2"))
+ unread = WT.MODEL_WEIGHTS.expected_unconsumed(d, stored.keys())
+ assert claimed <= everything and not claimed & unread, rank
+ assert claimed | unread == everything, (
+ f"rank {rank}: neither read nor predicted {sorted(everything - claimed - unread)[:8]}"
+ )
+ read_by.update(claimed)
+
+ dep4 = _dims(configs, "dep4", rank, mtp=False)
+ non_expert = {k for k in claimed if not _TRUNK_EXPERT_TENSOR.match(k)}
+ assert non_expert == {
+ k
+ for k in _claimed(_rank_manifest(stored, dep4, "dep4"))
+ if not _TRUNK_EXPERT_TENSOR.match(k)
+ }, rank
+
+ with torch.device("meta"):
+ params = WT.MODEL_WEIGHTS.declare(d)
+ dep4_params = WT.MODEL_WEIGHTS.declare(dep4)
+ # Every parameter dep4 declares, at dep4's shape, but the expert window's: 256
+ # experts where dep4 holds 128.
+ assert set(params.keys()) == set(dep4_params.keys())
+ for key in dep4_params:
+ ours, theirs = tuple(params[key].shape), tuple(dep4_params[key].shape)
+ if ours != theirs:
+ assert (ours[0], theirs[0], ours[1:]) == (256, 128, theirs[1:]), (key, ours)
+ for key, shape in _EXPERT_SHAPES["dep4"].items():
+ assert tuple(params[key].shape) == (256, *shape[1:]), (key, tuple(params[key].shape))
+ experts = [s[0] for s in WT.MODEL_WEIGHTS.manifest(d)["l0_fc2_w"]]
+ assert experts == [
+ f"model.language_model.layers.0.mlp.experts.{e}.down_proj.weight"
+ for e in range(256 * rank, 256 * (rank + 1))
+ ]
+ g = P.pool_geometry(
+ configs[0],
+ _mapping("dep2", rank),
+ tokens_per_block=32,
+ kv_cache_quant_algo=None,
+ dtype=torch.bfloat16,
+ )
+ assert (d.heads_kv, d.num_k_heads, d.num_v_heads) == (
+ g.heads_kv,
+ g.num_k_heads,
+ g.num_v_heads,
+ )
+
+ assert set(read_by) == everything - mtp_keys
+ wrong = sorted(k for k in read_by if _EXPERT_TENSOR.search(k) and read_by[k] != 1)
+ assert not wrong, f"expert tensors not read by exactly one rank: {wrong[:8]}"
+
+
+def test_dep2_load_fills_every_parameter_and_covers_the_checkpoint(stored, configs, monkeypatch):
+ """`load` end to end at dep2's last rank, without speculation, on meta tensors."""
+ monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None)
+ monkeypatch.setattr(torch.cuda, "empty_cache", lambda *a, **k: None)
+ model = _stub_model(configs, "dep2", 1, False)
+ WT.MODEL_WEIGHTS.load(model, _meta_checkpoint(stored))
+ assert model.lm_head.weight.shape[0] == 248320
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision/test_modeling_v2_triton_pos_embed_interpolate.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision/test_modeling_v2_triton_pos_embed_interpolate.py
new file mode 100644
index 000000000000..6a8b63a8f50e
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision/test_modeling_v2_triton_pos_embed_interpolate.py
@@ -0,0 +1,116 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test for the default target's triton_pos_embed_interpolate entry (its vision
+tower).
+
+`op.reference` is the explicit four-corner sum in fp64, returned unrounded as
+fp32, so the band measures the op's own bf16 arithmetic. `_resize64` runs
+torch's `F.interpolate` in fp64 and shares no arithmetic with either: each cell
+holds the reference to it (the structure: which axis is which, corner
+alignment, merge row order), and its wrong variants (`align_corners=False`,
+raster rows, swapped table axes) must fail the band.
+"""
+
+import zlib
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.vision.triton_pos_embed_interpolate import ( # noqa: E501
+ triton_pos_embed_interpolate as op,
+)
+
+__extra_import_path__ = ["../../../.."]
+from _validating import ULP, assert_within_ulp, certifying, validating # noqa: E402,F401
+
+assert torch.cuda.is_available(), "triton_pos_embed_interpolate requires a CUDA device"
+
+NUM_GRID = 48
+HIDDEN = 1152
+MERGE = 2
+
+
+def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
+ """4 bf16 ulp per element against the row scale, 1 ulp relative RMS.
+
+ The op rounds its four corner weights to bf16 and multiply-adds in bf16 (the
+ HF reference's arithmetic); `ref` is the exact value. The worst element
+ reaches about 3.5 ulp and the RMS 0.75, on `_table` and on the checkpoint's
+ own table alike. The nearest wrong variant, `align_corners=False`, lands an
+ order of magnitude outside the band, raster order and swapped axes two.
+ """
+ assert out.dtype == torch.bfloat16, out.dtype
+ assert_within_ulp(out.float(), ref, element_ulp=4.0, rms_ulp=1.0, ulp=ULP[torch.bfloat16])
+
+
+def _table(seed: int, num_grid: int = NUM_GRID, hidden: int = HIDDEN) -> torch.Tensor:
+ """A smooth `[num_grid**2, hidden]` bf16 table, the way a learned 2-D
+ position table is: a 12x12 normal field upsampled bicubically, plus 0.1
+ noise. Neighbouring rows have cosine 0.911; the checkpoint table's have 0.907.
+
+ Contiguous on purpose: the upsampled field's permuted view reshapes to a
+ strided table, which the kernel reads wrongly (`is_valid` refuses it).
+ """
+ g = torch.Generator(device="cuda").manual_seed(seed)
+ coarse = torch.randn(1, hidden, 12, 12, generator=g, device="cuda")
+ up = F.interpolate(coarse, size=(num_grid, num_grid), mode="bicubic", align_corners=True)
+ up = up[0].permute(1, 2, 0) + 0.1 * torch.randn(
+ num_grid, num_grid, hidden, generator=g, device="cuda"
+ )
+ return up.reshape(num_grid * num_grid, hidden).to(torch.bfloat16).contiguous()
+
+
+def _resize64(table, t, h, w, m_size, *, align=True, raster=False, swap=False):
+ """torch's bilinear resize in fp64, in merge order unless `raster`."""
+ num_grid = int(round(table.shape[0] ** 0.5))
+ grid = table.double().view(num_grid, num_grid, -1)
+ if swap:
+ grid = grid.transpose(0, 1)
+ src = grid.permute(2, 0, 1).unsqueeze(0).contiguous()
+ out = F.interpolate(src, size=(h, w), mode="bilinear", align_corners=align)[0].permute(1, 2, 0)
+ if raster:
+ rows = out.reshape(h * w, -1)
+ else:
+ rows = (
+ out.reshape(h // m_size, m_size, w // m_size, m_size, -1)
+ .permute(0, 2, 1, 3, 4)
+ .reshape(h * w, -1)
+ )
+ return rows.repeat(t, 1) if t > 1 else rows
+
+
+def _args(spec, table):
+ return (
+ table,
+ spec["t"],
+ spec["h"],
+ spec["w"],
+ spec["num_grid_per_side"],
+ spec["m_size"],
+ spec["dtype"],
+ )
+
+
+@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
+def test_certified_cells(cell, certifying): # noqa: F811
+ certifying(op)
+ spec = cell.spec
+ t, h, w, m_size = spec["t"], spec["h"], spec["w"], spec["m_size"]
+ table = _table(zlib.crc32(cell.why.encode()) % 2**31, spec["num_grid_per_side"], spec["hidden"])
+ before = table.clone()
+
+ out = op(*_args(spec, table))
+ assert out.shape == (t * h * w, spec["hidden"]) and out.dtype == spec["dtype"]
+ assert torch.equal(table, before) and out.data_ptr() != table.data_ptr()
+
+ ref = op.reference(*_args(spec, table))
+ _compare(out, ref)
+
+ # The reference's structure, from an implementation that shares nothing with it.
+ torch.testing.assert_close(ref, _resize64(table, t, h, w, m_size).float())
+
+ # The band refuses each wrong variant of the function.
+ for wrong in (dict(align=False), dict(raster=True), dict(swap=True)):
+ with pytest.raises(AssertionError):
+ _compare(out, _resize64(table, t, h, w, m_size, **wrong).float())
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py
deleted file mode 100644
index 1cfb3de32ee3..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py
+++ /dev/null
@@ -1,72 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the fused_sigmoid_mul_inplace entry."""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.activation.fused_sigmoid_mul_inplace import ( # noqa: E501
- fused_sigmoid_mul_inplace as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "fused_sigmoid_mul_inplace requires a CUDA device"
-
-
-def _seed(text: str) -> int:
- return zlib.crc32(text.encode())
-
-
-def _build(spec, seed):
- """`(out, gate_buffer, gate)` for one spec.
-
- `gate` is a view of `gate_buffer`: the second half of each
- `[T, heads, 2 * D]` q_proj head, or the 2-D gate window of the target's
- fused projection row.
- """
- g = torch.Generator(device="cuda").manual_seed(seed)
- t, heads, d, dtype = spec["tokens"], spec["heads"], spec["head_dim"], spec["dtype"]
- hidden = heads * d
-
- def randn(*shape):
- return torch.randn(*shape, generator=g, device="cuda").to(dtype)
-
- out = randn(t, hidden)
- if spec["gate"] == "interleaved":
- gate_buffer = randn(t, heads, 2 * d)
- gate = gate_buffer[:, :, d:]
- else:
- assert spec["gate"] == "fused_row", spec["gate"]
- # The target's fused input projection row [q | k | v | gate | index | pad]:
- # 14080 columns, the gate a 2-D [T, hidden] window at column 7168.
- gate_buffer = randn(t, 14080)
- gate = gate_buffer[:, 7168 : 7168 + hidden]
- return out, gate_buffer, gate
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- out, gate_buffer, gate = _build(cell.spec, seed=_seed(cell.why))
- gate_before = gate_buffer.clone()
- expected = op.reference(out, gate)
-
- result = op(out, gate)
-
- assert result.data_ptr() == out.data_ptr() and result.shape == out.shape
- torch.testing.assert_close(result, expected)
- assert torch.equal(gate_buffer, gate_before)
-
-
-def test_metadata_misuse_raises() -> None:
- """A gate with a strided last dim is refused loudly, as is every layout
- violation; hence the entry has no `is_valid`."""
- g = torch.Generator(device="cuda").manual_seed(8)
- attention_output = torch.randn(4, 128, generator=g, device="cuda").to(torch.bfloat16)
- gate = torch.randn(4, 256, generator=g, device="cuda").to(torch.bfloat16)[:, ::2]
- with pytest.raises(AssertionError):
- op(attention_output, gate)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/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
deleted file mode 100644
index 5255d4197ee8..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_erf.py
+++ /dev/null
@@ -1,141 +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 gelu_erf entry.
-
-The entry's identity is *which* GELU it is, and the wrong answer -- the tanh
-GELU this same tower's block MLPs use -- rounds to within about one bf16 ulp of
-it. The discrimination test is the point of this file; only fp32 tells the two
-apart.
-"""
-
-import math
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.activation.gelu_erf import ( # noqa: E501
- gelu_erf as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401
-
-assert torch.cuda.is_available(), "gelu_erf requires a CUDA device"
-
-MERGER_WIDTH = 4608
-TOWER_HIDDEN = 1152
-#: model.visual.merger, measured from the checkpoint: linear_fc1.weight std
-#: 0.010927, linear_fc1.bias std 0.002407, norm.weight mean 1.174581 / std
-#: 0.140536. The real tensors give an activand at std 0.875 / absmax 4.23.
-FC1_WEIGHT_STD = 0.010927
-FC1_BIAS_STD = 0.002407
-NORM_GAMMA_MEAN = 1.174581
-NORM_GAMMA_STD = 0.140536
-
-SQRT_2 = math.sqrt(2.0)
-SQRT_2_OVER_PI = math.sqrt(2.0 / math.pi)
-
-#: torch.testing's untouched per-dtype defaults, restated only so the
-#: discrimination test can say how far outside that gate each variant lands.
-_RTOL = {torch.bfloat16: 1.6e-2, torch.float16: 1e-3, torch.float32: 1.3e-6}
-_ATOL = 1e-5
-
-
-def _seed(text: str) -> int:
- return zlib.crc32(text.encode())
-
-
-def _uniform(shape, lo, hi, dtype, seed):
- g = torch.Generator(device="cuda").manual_seed(seed)
- x = torch.empty(shape, device="cuda", dtype=torch.float32).uniform_(lo, hi, generator=g)
- return x.to(dtype)
-
-
-def _build(spec, seed):
- shape, dtype = spec["shape"], spec["dtype"]
- if spec.get("dist", "uniform") == "merger":
- g = torch.Generator(device="cuda").manual_seed(seed)
- rows = shape[0]
- gamma = NORM_GAMMA_MEAN + NORM_GAMMA_STD * torch.randn(
- TOWER_HIDDEN, generator=g, device="cuda"
- )
- hidden = torch.randn(4 * rows, TOWER_HIDDEN, generator=g, device="cuda")
- grouped = (hidden * gamma.to(dtype).float()).reshape(rows, MERGER_WIDTH)
- weight = FC1_WEIGHT_STD * torch.randn(
- MERGER_WIDTH, MERGER_WIDTH, generator=g, device="cuda"
- )
- bias = FC1_BIAS_STD * torch.randn(MERGER_WIDTH, generator=g, device="cuda")
- return (grouped @ weight.to(dtype).float().t() + bias.to(dtype).float()).to(dtype)
- lo, hi = spec["lo"], spec["hi"]
- return _uniform(shape, lo, hi, dtype, seed)
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying): # noqa: F811
- certifying(op)
- x = _build(cell.spec, seed=_seed(cell.why))
- before = x.clone()
- out = op(x)
- assert out.shape == x.shape and out.dtype == x.dtype and out.device == x.device
- 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"
- torch.testing.assert_close(out, op.reference(x))
-
-
-def _erf64(x):
- x64 = x.double()
- return 0.5 * x64 * (1.0 + torch.erf(x64 / SQRT_2))
-
-
-def _tanh64(x):
- """The *other* GELU -- `gelu_pytorch_tanh`, this tower's block-MLP activation."""
- x64 = x.double()
- return 0.5 * x64 * (1.0 + torch.tanh(SQRT_2_OVER_PI * (x64 + 0.044715 * x64 * x64 * x64)))
-
-
-def test_identity_is_the_erf_form_and_not_tanh() -> None:
- """Measured in fp32, where the question has an answer; and shown to have
- none in bf16, so a bf16 match cannot quietly be read as confirmation."""
- x32 = _uniform((512, MERGER_WIDTH), -4.0, 4.0, torch.float32, 24)
- out = op(x32).double()
- ref = _erf64(x32)
- correct = (out - ref).abs().max().item()
- gate = _ATOL + _RTOL[torch.float32] * ref.abs().max().item()
-
- # The entry's own gate refuses the tanh form in fp32.
- with pytest.raises(AssertionError):
- torch.testing.assert_close(out.float(), _tanh64(x32).float())
-
- x64 = x32.double()
- variants = {
- "tanh gelu (gelu_pytorch_tanh)": _tanh64(x32),
- "sigmoid approx x*sigmoid(1.702x)": x64 * torch.sigmoid(1.702 * x64),
- "silu": x64 * torch.sigmoid(x64),
- "relu": torch.clamp(x64, min=0.0),
- }
- print(f"\n fp32: gate {gate:.3e}; correct-path deviation {correct:.3e}")
- for name, wrong in variants.items():
- dev = (out - wrong).abs().max().item()
- print(
- f" {name:36s} dev {dev:.6e} ({dev / correct:9.1f}x correct, {dev / gate:7.1f}x gate)"
- )
- assert dev > 100.0 * correct, (name, dev, correct)
- assert dev > 10.0 * gate, (name, dev, gate)
-
- # erf's scale truncated from 1/sqrt(2) to 1/1.414 -- a plausible
- # transcription slip, weak by construction, so a lower multiple.
- perturbed = 0.5 * x64 * (1.0 + torch.erf(x64 / 1.414))
- dev = (out - perturbed).abs().max().item()
- print(f" {'erf scale 1/sqrt(2) -> 1/1.414':36s} dev {dev:.6e} ({dev / gate:7.1f}x gate)")
- assert dev > 20.0 * correct, (dev, correct)
- assert dev > 1.5 * gate, (dev, gate)
-
- # bf16 cannot see the difference.
- xb = x32.to(torch.bfloat16)
- outb = op(xb).double()
- correct_b = (outb - _erf64(xb)).abs().max().item()
- tanh_b = (outb - _tanh64(xb)).abs().max().item()
- print(f" bf16: correct {correct_b:.6e}; tanh variant {tanh_b:.6e} ({tanh_b / correct_b:.3f}x)")
- assert tanh_b < 4.0 * correct_b, (tanh_b, correct_b)
- torch.testing.assert_close(outb.to(torch.bfloat16), op.reference(xb))
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
deleted file mode 100644
index b09168720ce5..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_tanh.py
+++ /dev/null
@@ -1,133 +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 gelu_tanh entry.
-
-The entry's identity is *which* GELU it is, so the discrimination test is the
-point of this file: in bf16 the tanh and erf GELUs round to within about one
-ulp of each other, and only fp32 tells them apart.
-"""
-
-import math
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.activation.gelu_tanh import ( # noqa: E501
- gelu_tanh as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401
-
-assert torch.cuda.is_available(), "gelu_tanh requires a CUDA device"
-
-INTERMEDIATE_SIZE = 4304
-HIDDEN_SIZE = 1152
-#: model.visual.blocks.*.mlp.linear_fc1, measured from the checkpoint: weight
-#: std 0.0197 (block 0), bias std 1.676 / absmax 9.44 -- the bias dominates.
-FC1_WEIGHT_STD = 0.0197
-FC1_BIAS_STD = 1.676
-
-SQRT_2_OVER_PI = math.sqrt(2.0 / math.pi)
-
-#: torch.testing's untouched per-dtype defaults, restated only so the
-#: discrimination test can say how far outside that gate each variant lands.
-_RTOL = {torch.bfloat16: 1.6e-2, torch.float16: 1e-3, torch.float32: 1.3e-6}
-_ATOL = 1e-5
-
-
-def _seed(text: str) -> int:
- return zlib.crc32(text.encode())
-
-
-def _uniform(shape, lo, hi, dtype, seed):
- g = torch.Generator(device="cuda").manual_seed(seed)
- x = torch.empty(shape, device="cuda", dtype=torch.float32).uniform_(lo, hi, generator=g)
- return x.to(dtype)
-
-
-def _build(spec, seed):
- shape, dtype = spec["shape"], spec["dtype"]
- if spec.get("dist", "uniform") == "fc1":
- g = torch.Generator(device="cuda").manual_seed(seed)
- rows = shape[0]
- hidden = torch.randn(rows, HIDDEN_SIZE, generator=g, device="cuda").to(dtype)
- weight = (
- FC1_WEIGHT_STD * torch.randn(shape[-1], HIDDEN_SIZE, generator=g, device="cuda")
- ).to(dtype)
- bias = (FC1_BIAS_STD * torch.randn(shape[-1], generator=g, device="cuda")).to(dtype)
- return (hidden.float() @ weight.float().t() + bias.float()).to(dtype)
- lo, hi = spec["lo"], spec["hi"]
- return _uniform(shape, lo, hi, dtype, seed)
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying): # noqa: F811
- certifying(op)
- x = _build(cell.spec, seed=_seed(cell.why))
- before = x.clone()
- out = op(x)
- assert out.shape == x.shape and out.dtype == x.dtype and out.device == x.device
- 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"
- torch.testing.assert_close(out, op.reference(x))
-
-
-def _tanh64(x):
- x64 = x.double()
- return 0.5 * x64 * (1.0 + torch.tanh(SQRT_2_OVER_PI * (x64 + 0.044715 * x64 * x64 * x64)))
-
-
-def _erf64(x):
- """The *other* GELU -- exact, erf-based; `torch.nn.GELU()` and `moe_gelu`."""
- x64 = x.double()
- return 0.5 * x64 * (1.0 + torch.erf(x64 / math.sqrt(2.0)))
-
-
-def test_identity_is_the_tanh_form_and_not_erf() -> None:
- """Measured in fp32, where the question has an answer; and shown to have
- none in bf16, so a bf16 match cannot quietly be read as confirmation."""
- x32 = _uniform((512, INTERMEDIATE_SIZE), -4.0, 4.0, torch.float32, 24)
- out = op(x32).double()
- ref = _tanh64(x32)
- correct = (out - ref).abs().max().item()
- gate = _ATOL + _RTOL[torch.float32] * ref.abs().max().item()
-
- # The entry's own gate refuses the erf form in fp32.
- with pytest.raises(AssertionError):
- torch.testing.assert_close(out.float(), _erf64(x32).float())
-
- x64 = x32.double()
- variants = {
- "erf gelu (nn.GELU, and moe_gelu)": _erf64(x32),
- "sigmoid approx x*sigmoid(1.702x)": x64 * torch.sigmoid(1.702 * x64),
- "silu": x64 * torch.sigmoid(x64),
- "relu": torch.clamp(x64, min=0.0),
- }
- print(f"\n fp32: gate {gate:.3e}; correct-path deviation {correct:.3e}")
- for name, wrong in variants.items():
- dev = (out - wrong).abs().max().item()
- print(
- f" {name:36s} dev {dev:.6e} ({dev / correct:9.1f}x correct, {dev / gate:7.1f}x gate)"
- )
- assert dev > 100.0 * correct, (name, dev, correct)
- assert dev > 10.0 * gate, (name, dev, gate)
-
- # A perturbation of this same formula: weak by construction (bounded by the
- # coefficient change times the tanh derivative), so a lower multiple.
- perturbed = 0.5 * x64 * (1.0 + torch.tanh(SQRT_2_OVER_PI * (x64 + 0.045 * x64**3)))
- dev = (out - perturbed).abs().max().item()
- print(f" {'coefficient 0.044715 -> 0.045':36s} dev {dev:.6e} ({dev / gate:7.1f}x gate)")
- assert dev > 100.0 * correct, (dev, correct)
- assert dev > 5.0 * gate, (dev, gate)
-
- # bf16 cannot see the difference: the erf variant's max deviation is within
- # a small multiple of the correct path's own.
- xb = x32.to(torch.bfloat16)
- outb = op(xb).double()
- correct_b = (outb - _tanh64(xb)).abs().max().item()
- erf_b = (outb - _erf64(xb)).abs().max().item()
- print(f" bf16: correct {correct_b:.6e}; erf variant {erf_b:.6e} ({erf_b / correct_b:.3f}x)")
- assert erf_b < 4.0 * correct_b, (erf_b, correct_b)
- torch.testing.assert_close(outb.to(torch.bfloat16), op.reference(xb))
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
deleted file mode 100644
index e5291a35eb15..000000000000
--- 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
+++ /dev/null
@@ -1,66 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the silu_and_mul_rows entry: the cells bitwise `flashinfer_silu_and_mul` on a
-contiguous copy of the gate|up columns (and so inside its reference band) with the fused
-output's extra columns present, graph capture, the width guard."""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.flashinfer_silu_and_mul import ( # noqa: E501
- flashinfer_silu_and_mul,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.activation.silu_and_mul_rows import ( # noqa: E501
- silu_and_mul_rows as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "silu_and_mul_rows requires a CUDA device"
-
-INTER = 640
-
-
-def _build(spec, seed):
- g = torch.Generator(device="cuda").manual_seed(seed)
- rows, width = spec["rows"], 2 * INTER + 8
- x = (torch.randn(rows, width, generator=g, device="cuda") * 3.0).to(torch.bfloat16)
- if rows:
- x[0, :4] = torch.tensor([0.0, -0.0, 88.0, -100.0], device="cuda").to(torch.bfloat16)
- return x
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
- certifying(op)
- x = _build(cell.spec, zlib.crc32(cell.why.encode()))
- out = op(x, INTER)
- assert out.shape == (x.shape[0], INTER) and out.is_contiguous()
- torch.testing.assert_close(out, op.reference(x, INTER))
- if x.shape[0]:
- assert torch.equal(out, flashinfer_silu_and_mul(x[:, : 2 * INTER].contiguous()))
-
-
-def test_cuda_graph_capture_and_replay() -> None:
- x = _build(dict(rows=128), 8)
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(warmup):
- op(x, INTER)
- torch.cuda.current_stream().wait_stream(warmup)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- out = op(x, INTER)
- x.copy_(_build(dict(rows=128), 9))
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(out, op(x, INTER))
-
-
-def test_guards() -> None:
- x = _build(dict(rows=4), 5)
- with validating(op), pytest.raises(AssertionError):
- op(x[:, : 2 * INTER - 2], INTER)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_canonical_selection.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_canonical_selection.py
deleted file mode 100644
index 0561d585eeae..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_canonical_selection.py
+++ /dev/null
@@ -1,167 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the qsa_canonical_selection entry.
-
-The rewrite fixes every slot of every row, so `_compare` is index equality against a
-composite-key top-k. The input rows come from an adversary that keeps a random part of the
-tie class at the k-th value in a random order, or from the radix engine the target calls.
-"""
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_canonical_selection import ( # noqa: E501
- qsa_canonical_selection as op,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_topk_select import ( # noqa: E501
- qsa_topk_select,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401
-
-assert torch.cuda.is_available(), "qsa_canonical_selection requires a CUDA device"
-
-DEV = "cuda"
-COMPRESS_RATIO = 4
-BLOCK_TOPK = 512
-COLS = 2080
-
-
-def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
- """Index for index. The rule fixes every slot of every row, so nothing is tolerated: a
- selection does no arithmetic."""
- assert out.dtype == torch.int32, out.dtype
- assert out.shape == ref.shape, (out.shape, ref.shape)
- wrong = (out != ref).any(1)
- if bool(wrong.any()):
- row = int(wrong.nonzero()[0, 0])
- raise AssertionError(
- f"{int(wrong.sum())} of {out.shape[0]} rows differ from the canonical "
- f"selection; first is row {row}"
- )
-
-
-def _logits(kind: str, rows: int, cols: int, gen: torch.Generator) -> torch.Tensor:
- if kind == "relu":
- return torch.relu(torch.randn(rows, cols, generator=gen, device=DEV)) * 3.0
- if kind == "relu_sparse":
- x = torch.relu(torch.randn(rows, cols, generator=gen, device=DEV)) * 3.0
- return x * (torch.rand(rows, cols, generator=gen, device=DEV) < 0.2)
- raise ValueError(kind)
-
-
-def _windows(spec: dict, rows: int, cols: int, gen: torch.Generator):
- kind = spec["windows"][0]
- starts = torch.zeros(rows, dtype=torch.int32, device=DEV)
- if kind == "chunk":
- positions = spec["windows"][1] + torch.arange(rows, device=DEV)
- else:
- assert kind == "positions_random", kind
- lo, hi = spec["windows"][1:]
- positions = torch.randint(lo, hi + 1, (rows,), generator=gen, device=DEV)
- ends = ((positions + 1) // COMPRESS_RATIO).clamp_max(cols).to(torch.int32)
- return starts, ends
-
-
-def _poison(logits, starts, ends, kind) -> None:
- columns = torch.arange(logits.shape[1], device=DEV)[None, :]
- outside = (columns < starts[:, None]) | (columns >= ends[:, None])
- if kind == "nan":
- logits.masked_fill_(outside, float("nan"))
- return
- # Values that would win any comparison they were allowed into.
- pattern = torch.tensor([float("nan"), float("inf"), 3e38, float("nan")], device=DEV)
- logits.copy_(torch.where(outside, pattern[columns % 4].expand_as(logits), logits))
-
-
-def _order_key(logits: torch.Tensor) -> torch.Tensor:
- """int64 that orders as the fp32 values do, -0.0 folded into +0.0."""
- bits = (logits.float() + 0.0).view(torch.int32).long()
- return torch.where(bits < 0, bits ^ 0x7FFFFFFF, bits)
-
-
-def _adversary(logits, starts, ends, index_topk, gen) -> torch.Tensor:
- """An exact top-k of each window that keeps a random part of the tie class at the
- k-th value, in a random order, `-1` last."""
- rows, cols = logits.shape
- columns = torch.arange(cols, device=DEV)[None, :]
- in_window = (columns >= starts.long()[:, None]) & (columns < ends.long()[:, None])
- tiebreak = torch.rand(rows, cols, generator=gen, device=DEV).argsort(1)
- key = torch.where(
- in_window, (_order_key(logits) << 32) | tiebreak, torch.iinfo(torch.int64).min
- )
- width = min(index_topk, cols)
- top = torch.topk(key, width, dim=1).indices
- local = torch.where(torch.gather(in_window, 1, top), top - starts.long()[:, None], -1)
- shuffle = torch.where(local >= 0, torch.rand(local.shape, generator=gen, device=DEV), 2.0)
- local = torch.gather(local, 1, shuffle.argsort(1))
- out = torch.full((rows, index_topk), -1, dtype=torch.int32, device=DEV)
- out[:, :width] = local.to(torch.int32)
- return out
-
-
-def _radix(logits, starts, ends) -> torch.Tensor:
- out = torch.empty(logits.shape[0], BLOCK_TOPK, dtype=torch.int32, device=DEV)
- return qsa_topk_select(logits, starts, ends, out, index_topk=BLOCK_TOPK)
-
-
-def _build(spec: dict, rows: int, seed: int):
- gen = torch.Generator(device=DEV).manual_seed(seed)
- logits = _logits(spec["logits"], rows, spec["cols"], gen)
- starts, ends = _windows(spec, rows, spec["cols"], gen)
- _poison(logits, starts, ends, spec["tail"])
- if spec["selection"] == "random":
- indices = _adversary(logits, starts, ends, spec["index_topk"], gen)
- else:
- indices = _radix(logits, starts, ends)
- return logits, starts, ends, indices
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811
- certifying(op)
- spec = cell.spec
- seed = sum(map(ord, cell.why))
- row_counts = spec["rows"] if isinstance(spec["rows"], tuple) else (spec["rows"],)
- for rows in row_counts:
- args = _build(spec, rows, seed + rows)
- expected = op.reference(*args)
- out = op(*args)
- assert out is args[3], "the rewrite writes `indices` and returns it"
- _compare(out, expected)
-
-
-def test_narrowing_the_table_to_a_bound_on_the_windows_changes_nothing() -> None:
- """A caller that knows every window ends by column 1500 passes `logits[:, :1500]` of a
- wide table: the rows are the full table's, and a 2^24-element table is tiled."""
- gen = torch.Generator(device=DEV).manual_seed(21)
- rows, cols, bound = 4096, 8192, 1500
- logits = _logits("relu_sparse", rows, cols, gen)
- starts = torch.zeros(rows, dtype=torch.int32, device=DEV)
- ends = torch.randint(0, bound + 1, (rows,), generator=gen, device=DEV).to(torch.int32)
- _poison(logits, starts, ends, "poison")
- raw = _radix(logits, starts, ends)
- wide = op(logits, starts, ends, raw.clone())
- narrow = op(logits[:, :bound], starts, ends, raw.clone())
- _compare(wide, op.reference(logits, starts, ends, raw))
- assert torch.equal(wide, narrow)
-
-
-def test_guard_refuses_a_broadcast_row_start() -> None:
- """A one-element `row_starts` broadcasts over every row without an error and frames
- every row by row 0's start. Shown unguarded first, then refused."""
- gen = torch.Generator(device=DEV).manual_seed(9)
- rows = 4
- logits = torch.relu(torch.randn(rows, COLS, generator=gen, device=DEV))
- starts = torch.tensor([0, 10, 20, 30], dtype=torch.int32, device=DEV)
- ends = torch.full((rows,), 2000, dtype=torch.int32, device=DEV)
- raw = _adversary(logits, starts, ends, BLOCK_TOPK, gen)
- expected = op.reference(logits, starts, ends, raw)
- out = op(logits, starts[:1], ends, raw.clone())
- assert out.shape == expected.shape and not torch.equal(out, expected), (
- "the broadcast hazard stopped reproducing"
- )
- with pytest.raises(AssertionError, match="one entry per logits row"):
- with validating(op):
- op(logits, starts[:1], ends, raw.clone())
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_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
deleted file mode 100644
index 0c83dcf90793..000000000000
--- 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
+++ /dev/null
@@ -1,156 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the qsa_expand_block_indices catalog entry.
-
-Selections are built the way `qsa_topk_select` emits them: a random subset of
-the complete visible groups, prefix-packed, then `-1`. Every gate is bitwise.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_expand_block_indices import ( # noqa: E501
- qsa_expand_block_indices as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401
-
-assert torch.cuda.is_available(), "qsa_expand_block_indices requires a CUDA device"
-
-
-def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
- """Bitwise: these are KV-cache token indices, and one wrong entry reads
- a different key. There is no rounding for a tolerance to absorb."""
- assert out.dtype == ref.dtype == torch.int32, (out.dtype, ref.dtype)
- assert out.shape == ref.shape, (out.shape, ref.shape)
- torch.testing.assert_close(out, ref, rtol=0, atol=0)
-
-
-DEV = "cuda"
-CR, TOPK = 4, 2048
-ENGINE_MAX_SEQ_LEN = 8192
-
-
-def _serving_rows(requests: int, context_lens: tuple) -> tuple[list[int], list[int]]:
- """Per-row positions and lengths of one packed step.
-
- Chunk 0 starts at 0, so its rows span every visible-group count up to the
- budget; later chunks are chunked-prefill continuations. Each remaining
- request is an MTP generation request with two rows at consecutive positions.
- """
- positions: list[int] = []
- lengths: list[int] = []
- for chunk, length in enumerate(context_lens):
- start = 0 if chunk == 0 else (907 * chunk) % (ENGINE_MAX_SEQ_LEN - length + 1)
- positions.extend(range(start, start + length))
- lengths.extend([start + length] * length)
- for gen in range(requests - len(context_lens)):
- position = 100 + (499 * gen) % 8000
- positions.extend((position, position + 1))
- lengths.extend((position + 2, position + 2))
- assert max(lengths, default=0) <= ENGINE_MAX_SEQ_LEN
- return positions, lengths
-
-
-def _rows(batch) -> tuple[list[int], list[int]]:
- kind = batch[0]
- if kind == "sweep":
- positions = [4096 + 13 * i for i in range(batch[1])]
- elif kind == "serving":
- return _serving_rows(batch[1], batch[2])
- else:
- assert kind == "verify", kind
- starts, per_request = batch[1], batch[2]
- positions = [s + i for s in starts for i in range(per_request)]
- lengths = [s + per_request for s in starts for _ in range(per_request)]
- assert max(lengths) <= ENGINE_MAX_SEQ_LEN
- return positions, lengths
- return positions, [p + 1 for p in positions]
-
-
-def _packed_selection(positions: torch.Tensor, compress_ratio: int, token_topk: int, seed: int):
- """A legal Top-K result: distinct complete visible groups, then -1 padding."""
- block_topk = token_topk // compress_ratio
- rows = positions.numel()
- if rows == 0:
- return torch.empty((0, block_topk), dtype=torch.int32, device=DEV)
- available = (positions.long() + 1) // compress_ratio
- width = max(block_topk, int(available.max()))
- gen = torch.Generator(device=DEV).manual_seed(seed)
- scores = torch.rand(rows, width, generator=gen, device=DEV)
- groups = torch.arange(width, device=DEV)
- scores = torch.where(groups[None, :] < available[:, None], scores, torch.inf)
- chosen = scores.argsort(dim=1)[:, :block_topk]
- taken = available.clamp(max=block_topk)
- chosen = torch.where(torch.arange(block_topk, device=DEV)[None, :] < taken[:, None], chosen, -1)
- return chosen.to(torch.int32)
-
-
-def _build(spec, seed: int):
- positions, lengths = _rows(spec["batch"])
- cr, topk = spec["compress_ratio"], spec["token_topk"]
- 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)
- return (blocks, query_positions, sequence_lengths), dict(compress_ratio=cr, token_topk=topk)
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying): # noqa: F811
- certifying(op)
- args, geometry = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
- out = op(*args, **geometry)
- rows = args[0].shape[0]
- final = geometry["token_topk"] + geometry["compress_ratio"] - 1
- assert out.dtype == torch.int32 and out.is_cuda and out.is_contiguous()
- assert tuple(out.shape) == (rows, final)
- _compare(out, op.reference(*args, **geometry))
- if cell.spec["batch"][0] == "serving":
- positions, lengths = args[1], args[2]
- assert bool((lengths.long() > positions.long() + 1).any()), (
- "a serving cell must drive sequence_lengths past position + 1"
- )
-
-
-def 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)
-
- def call():
- return op(blocks, positions, lengths, compress_ratio=CR, token_topk=TOPK)
-
- eager = call().clone()
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(warmup):
- for _ in range(3):
- call()
- torch.cuda.current_stream().wait_stream(warmup)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- captured = call()
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(captured, eager)
-
- blocks.copy_(_packed_selection(positions, CR, TOPK, seed=32))
- graph.replay()
- torch.cuda.synchronize()
- _compare(captured, op.reference(blocks, positions, lengths, compress_ratio=CR, token_topk=TOPK))
-
-
-def test_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)
- 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
deleted file mode 100644
index 1f8e34122860..000000000000
--- 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
+++ /dev/null
@@ -1,248 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the qsa_mtp_reuse_topk catalog entry.
-
-Captured rows are built the way `qsa_topk_select` emits them: `min(K, visible)`
-distinct complete groups in an arbitrary order as a prefix, then `-1`. Inputs
-sit in the engine's buffers -- one slot of a `[slots, max_num_sequences, K]`
-shared buffer holding stale rows elsewhere, and row slices of the
-`max_num_tokens`-tall request map, visible-count and Top-K buffers. Every gate
-is bitwise.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_mtp_reuse_topk import ( # noqa: E501
- qsa_mtp_reuse_topk as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401
-
-assert torch.cuda.is_available(), "qsa_mtp_reuse_topk requires a CUDA device"
-
-
-def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
- """Bitwise: these are group ids, and one wrong entry attends to a different
- group. There is no rounding for a tolerance to absorb."""
- assert out.dtype == ref.dtype, (out.dtype, ref.dtype)
- assert out.shape == ref.shape, (out.shape, ref.shape)
- torch.testing.assert_close(out, ref, rtol=0, atol=0)
-
-
-DEV = "cuda"
-CR = 4
-#: The engine's `max_num_tokens`: height of the request-map and Top-K buffers.
-ENGINE_MAX_NUM_TOKENS = 8192
-ENGINE_MAX_SEQ_LEN = 8192
-#: Captured positions for the first requests of a "mixed" batch: empty (0
-#: groups), 1, 100, 255, 511 and 512 groups, rows that gain a group and rows
-#: that do not, up to the last position the 8192-token envelope allows.
-ANCHORS = (2, 3, 402, 2046, 2050, 8188, 2051, 2047, 1022, 4095, 5998, 7001, 37, 2044, 3001, 6143)
-
-
-def _captured_row(visible: int, block_topk: int, gen: torch.Generator) -> torch.Tensor:
- """A legal capture: `min(K, visible)` distinct groups below `visible`, then -1."""
- row = torch.full((block_topk,), -1, dtype=torch.int32)
- if visible > 0:
- take = min(visible, block_topk)
- row[:take] = torch.randperm(visible, generator=gen)[:take].to(torch.int32)
- return row
-
-
-def _requests(contexts, drafted, gen) -> list[tuple[int, int]]:
- """(captured_visible, current_visible) per request."""
- kind = contexts[0]
- if kind == "explicit":
- positions = list(contexts[1])
- else:
- count = contexts[1]
- low, high = {"mixed": (0, 8188), "full": (2047, 8188)}[kind]
- fixed = {"mixed": ANCHORS, "full": (2047, 2050, 8188, 2051)}[kind]
- positions = list(fixed[:count])
- positions += torch.randint(low, high + 1, (count - len(positions),), generator=gen).tolist()
- out = []
- for q, position in enumerate(positions):
- d = drafted[q % len(drafted)]
- position = min(position, ENGINE_MAX_SEQ_LEN - 1 - d)
- out.append(((position + 1) // CR, (position + d + 1) // CR))
- return out
-
-
-def _build(spec, seed: int, requests: int | None = None) -> dict:
- """Turn a cell spec into the op's keyword arguments, plus the buffers behind them.
-
- Returns `dict(kwargs=..., buffers=...)`; `buffers` holds the full tensors the
- arguments are views of, so a test can check that nothing outside the view moved.
- """
- gen = torch.Generator().manual_seed(seed)
- block_topk = spec["block_topk"]
- contexts = spec["contexts"]
- if requests is not None:
- contexts = (contexts[0], requests)
- pairs = _requests(contexts, spec["drafted"], gen)
- num_requests = len(pairs)
- per_request = spec["rows_per_request"]
- rows = num_requests * per_request
- num_slots, slot = spec["slots"]
- shared_rows = max(spec["shared_rows"], num_requests)
- meta_dtype = spec["metadata_dtype"]
-
- # Stale rows everywhere: other layers' slots and requests past this batch.
- shared = torch.randint(-1, 3000, (num_slots, shared_rows, block_topk), generator=gen)
- shared = shared.to(torch.int32)
- shared_vis = torch.randint(0, 3000, (num_slots, shared_rows), generator=gen).to(meta_dtype)
- for q, (captured, _) in enumerate(pairs):
- shared[slot, q] = _captured_row(captured, block_topk, gen)
- shared_vis[slot, q] = captured
-
- request, visible = [], []
- for q in range(num_requests):
- captured, current = pairs[q]
- for j in range(per_request):
- request.append(q)
- # A later row of the same request sees at least as many groups.
- visible.append(current + j // CR)
- request = torch.tensor(request, dtype=meta_dtype).reshape(rows)
- visible = torch.tensor(visible, dtype=meta_dtype).reshape(rows)
-
- height = max(ENGINE_MAX_NUM_TOKENS, rows + 37)
- tall_out = torch.randint(-1, 3000, (height, block_topk), generator=gen).to(torch.int32)
- tall_req = torch.randint(0, shared_rows, (height,), generator=gen).to(meta_dtype)
- tall_vis = torch.randint(0, 3000, (height,), generator=gen).to(meta_dtype)
- tall_req[:rows], tall_vis[:rows] = request, visible
- buffers = dict(
- shared=shared.to(DEV),
- shared_vis=shared_vis.to(DEV),
- request=tall_req.to(DEV),
- visible=tall_vis.to(DEV),
- output=tall_out.to(DEV),
- )
- kwargs = dict(
- shared_indices=buffers["shared"][slot],
- shared_visible_blocks=buffers["shared_vis"][slot],
- request_indices=buffers["request"][:rows],
- visible_blocks=buffers["visible"][:rows],
- output=buffers["output"][:rows],
- )
- return dict(kwargs=kwargs, buffers=buffers, rows=rows)
-
-
-def _assert_legal_selection(out: torch.Tensor, kwargs: dict) -> None:
- """What the producer's contract buys: no duplicate group, none past the horizon."""
- if out.shape[0] == 0:
- return
- ids = out.long()
- ordered = torch.sort(ids, dim=1).values
- assert not bool(((ordered[:, 1:] == ordered[:, :-1]) & (ordered[:, 1:] >= 0)).any()), (
- "a group id appears twice in a row"
- )
- captured = kwargs["shared_visible_blocks"].long()[kwargs["request_indices"].long()]
- horizon = torch.maximum(captured, kwargs["visible_blocks"].long())
- assert bool((ids < horizon[:, None]).all()), "a group id reaches past the row's horizon"
- assert bool((ids >= -1).all())
-
-
-def _check_cell(spec, seed: int, requests: int | None = None) -> torch.Tensor:
- built = _build(spec, seed, requests)
- kwargs, buffers, rows = built["kwargs"], built["buffers"], built["rows"]
- before = {name: tensor.clone() for name, tensor in buffers.items()}
- out = op(**kwargs)
- torch.cuda.synchronize()
- assert out.data_ptr() == kwargs["output"].data_ptr(), "the op returns the buffer it wrote"
- assert out.dtype == torch.int32 and tuple(out.shape) == (rows, spec["block_topk"])
- _compare(out, op.reference(**kwargs))
- _assert_legal_selection(out, kwargs)
- for name, tensor in buffers.items():
- if name == "output":
- assert torch.equal(tensor[rows:], before[name][rows:]), "rows past the view moved"
- else:
- assert torch.equal(tensor, before[name]), f"{name} was written"
- return out
-
-
-def _check_graph(spec, seed: int) -> None:
- """At each request count: eager, capture, then replay on new data."""
- for requests in spec["graph_batches"]:
- first = _build(spec, seed + requests, requests)
- kwargs = first["kwargs"]
- eager = op(**kwargs).clone()
- _compare(eager, op.reference(**kwargs))
-
- side = torch.cuda.Stream()
- side.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(side):
- op(**kwargs)
- torch.cuda.current_stream().wait_stream(side)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- captured = op(**kwargs)
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(captured, eager), f"{requests} requests: replay != eager"
-
- # The mixed batch's first requests sit at fixed positions, so the new
- # data comes from full rows: every request's captured row changes.
- refill = {**spec, "contexts": ("full", requests)}
- fresh = _build(refill, seed + 1000 + requests)["kwargs"]
- for name, tensor in kwargs.items():
- if name != "output":
- tensor.copy_(fresh[name])
- graph.replay()
- torch.cuda.synchronize()
- expected = op.reference(**kwargs)
- assert not torch.equal(expected, eager), f"{requests} requests: new data changes nothing"
- _compare(captured, expected)
- rerun = op(**{**kwargs, "output": torch.empty_like(kwargs["output"])})
- assert torch.equal(captured, rerun), f"{requests} requests: replay != eager on new data"
- del graph
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying): # noqa: F811
- certifying(op)
- seed = zlib.crc32(cell.why.encode())
- if cell.spec["graph_batches"]:
- _check_graph(cell.spec, seed)
- else:
- _check_cell(cell.spec, seed)
-
-
-def _spec_of(prefix: str) -> dict:
- (cell,) = [c for c in op.CELLS if c.why.startswith(prefix)]
- return cell.spec
-
-
-def _guard_case():
- """Legal arguments, then the whole `[slots, seqs]` visible-count buffer passed where
- one slot's view belongs.
-
- Returns (intended, kwargs): `intended` is the reference on the legal
- arguments -- what a caller passing the bad layout meant to get.
- """
- spec = _spec_of("decode MTP")
- kwargs = _build({**spec, "slots": (3, 2)}, seed=13, requests=6)["kwargs"]
- intended = op.reference(**kwargs)
- whole = torch.zeros((3, kwargs["shared_visible_blocks"].numel()), dtype=torch.int32)
- whole = whole.to(DEV)
- whole[2] = kwargs["shared_visible_blocks"]
- kwargs["shared_visible_blocks"] = whole
- return intended, kwargs
-
-
-def test_guard_refuses_a_whole_multi_slot_count_buffer() -> None:
- """Unguarded the op reads slot 0's counts and returns a plausible wrong answer;
- `is_valid` refuses it."""
- intended, kwargs = _guard_case()
- out = op(**kwargs)
- torch.cuda.synchronize()
- assert not torch.equal(out, intended), "the op got it right after all"
- _, kwargs = _guard_case()
- with (
- validating(op),
- pytest.raises(ValueError, match="shared_visible_blocks must be one slot's"),
- ):
- op(**kwargs)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py
deleted file mode 100644
index 8498ad6cded2..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py
+++ /dev/null
@@ -1,284 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the qsa_paged_index_scores catalog entry.
-
-The paged state is real: a `[pages, tokens_per_block, 1, dim]` index cache
-(a view of a wider page under the NVFP4 cells) sized from the pages the table maps, and an
-int32 block table whose rows are request *slots* -- more of them than the batch has requests, assigned by a
-seeded permutation -- with shuffled physical pages and `-1` past each
-sequence on the serving layouts, which is what the engine's table carries.
-
-The call has two kernels behind it. `_watch_launches` swaps the two Triton
-entry points for recording proxies that still launch the real kernel, so every
-cell asserts the branch it names.
-"""
-
-import math
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_index_scores import ( # noqa: E501
- qsa_paged_index_scores as op,
-)
-from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels as _qsa_kernels
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401
-
-assert torch.cuda.is_available(), "qsa_paged_index_scores requires a CUDA device"
-
-
-def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
- """torch's default fp32 band on every specified column; `-inf` exactly.
-
- bf16 and fp16 products are exact in the kernel's fp32 accumulator, so
- only summation order separates it from the reference: measured at most
- 2.9e-6 over every cell (2.4e-6 at 8191 rows x 2080 columns, scores to
- 11.6), inside the default rtol 1.3e-6 / atol 1e-5. Wrong variants (no
- relu, mean over heads, first group member as key, unscaled,
- 1/head_dim scale, block table ignored) were measured 3e5x to 4e6x
- outside it. Columns the reference marks NaN are unspecified and not read.
- """
- assert out.dtype == ref.dtype == torch.float32, (out.dtype, ref.dtype)
- assert out.shape == ref.shape, (out.shape, ref.shape)
- specified = ~torch.isnan(ref)
- torch.testing.assert_close(out[specified], ref[specified])
-
-
-DEV = torch.device("cuda")
-#: The callable's own threshold: at or below it every call is row-wise.
-DECODE_ROW_LIMIT = 256
-#: Index heads must divide this into more than one row for the query tile.
-DOT_TILE_WIDTH = 128
-
-
-class _LaunchRecorder:
- """Transparent proxy over a Triton entry point: records the grid, then launches."""
-
- def __init__(self, inner, name, log):
- self._inner, self._name, self._log = inner, name, log
-
- def __getitem__(self, grid):
- self._log.append((self._name, tuple(grid)))
- return self._inner[grid]
-
- def __getattr__(self, item):
- return getattr(self._inner, item)
-
-
-class _watch_launches:
- """Record which of the entry's two Triton kernels a call launches."""
-
- _NAMES = {
- "_qsa_paged_index_scores_kernel": "row_wise",
- "_qsa_paged_index_scores_qtile_kernel": "query_tile",
- }
-
- def __enter__(self):
- self.log = []
- self._saved = {name: getattr(_qsa_kernels, name) for name in self._NAMES}
- for name, tag in self._NAMES.items():
- setattr(_qsa_kernels, name, _LaunchRecorder(self._saved[name], tag, self.log))
- return self
-
- def __exit__(self, *exc):
- for name, kernel in self._saved.items():
- setattr(_qsa_kernels, name, kernel)
- return False
-
- @property
- def branch(self):
- assert len(self.log) == 1, f"expected exactly one kernel launch, saw {self.log}"
- return self.log[0][0]
-
- @property
- def grid(self):
- assert len(self.log) == 1, f"expected exactly one kernel launch, saw {self.log}"
- return self.log[0][1]
-
-
-def _case(rows, requests, **overrides):
- """A narrow-table spec for the non-cell tests, same geometry as the entry's own."""
- spec = dict(
- layout="random",
- rows=rows,
- requests=requests,
- slots=requests + 4,
- pages=8,
- heads=4,
- head_dim=128,
- tokens_per_block=32,
- compress_ratio=4,
- dtype=torch.bfloat16,
- positions_dtype=torch.int32,
- request_dtype=torch.int32,
- q_layout="packed",
- only_visible_blocks=False,
- context_rows=False,
- )
- spec.update(overrides)
- return spec
-
-
-def _rows_and_requests(spec, gen):
- """Per-row positions, per-row request ordinal, and each request's mapped pages."""
- layout, tpb, pages = spec["layout"], spec["tokens_per_block"], spec["pages"]
-
- def pages_for(seq_len):
- return math.ceil(seq_len / tpb)
-
- if layout == "random":
- rows, requests = spec["rows"], spec["requests"]
- req = torch.arange(rows, device=DEV) * requests // max(rows, 1)
- positions = torch.randint(0, pages * tpb, (rows,), generator=gen, device=DEV)
- return positions, req, [pages] * requests
-
- # (first position, row count) per request, in row order
- if layout == "capture":
- segments = [(0, n) for n in spec["lengths"]]
- elif layout == "decode":
- segments = [(p, spec["rows_per_request"]) for p in spec["gen_positions"]]
- else:
- raise AssertionError(f"unknown layout {layout!r}")
- positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments])
- req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)])
- mapped = [pages_for(s + n) for s, n in segments]
- assert max(mapped) <= pages
- return positions, req, mapped
-
-
-def _build(spec, seed):
- """Real paged state for one spec, keyed by `__call__`'s parameter names."""
- gen = torch.Generator(device=DEV).manual_seed(seed)
- heads, head_dim, dtype = spec["heads"], spec["head_dim"], spec["dtype"]
- positions, req, mapped = _rows_and_requests(spec, gen)
- rows, slots = positions.numel(), spec["slots"]
- assert len(mapped) <= slots
-
- slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(mapped)]
- pool_pages = sum(mapped) + 3
- page_ids = torch.randperm(pool_pages, generator=gen, device=DEV).to(torch.int32)
- block_table = torch.full((slots, spec["pages"]), -1, dtype=torch.int32, device=DEV)
- cursor = 0
- for slot, count in zip(slot_of.tolist(), mapped):
- block_table[slot, :count] = page_ids[cursor : cursor + count]
- cursor += count
- tpb = spec["tokens_per_block"]
- keys = torch.randn(pool_pages, tpb, 1, head_dim, generator=gen, device=DEV).to(dtype)
- if "page_pad" in spec:
- # The NVFP4 pool's page: the other index buffers, this layer's, then the position
- # page, so the plane is a view at page stride buffers * tpb * dim + page_pad.
- buffers, buffer = spec["pool_buffers"], tpb * head_dim
- page = torch.zeros(pool_pages, buffers * buffer + spec["page_pad"], dtype=dtype, device=DEV)
- index_cache = page[:, (buffers - 1) * buffer : buffers * buffer].view(
- pool_pages, tpb, 1, head_dim
- )
- index_cache.copy_(keys)
- assert index_cache.stride() == (page.stride(0), head_dim, head_dim, 1)
- else:
- index_cache = keys
-
- if spec["q_layout"] == "scratch_slice":
- # [T, index heads | 1 k | 1 v, D], q is the first `heads` slots.
- scratch = torch.randn(rows, heads + 2, head_dim, generator=gen, device=DEV).to(dtype)
- q = scratch[:, :heads]
- else:
- q = torch.randn(rows, heads, head_dim, generator=gen, device=DEV).to(dtype)
- if spec["q_layout"] != "packed" and rows > 1: # one row is trivially contiguous
- assert not q.is_contiguous()
-
- return dict(
- q=q,
- index_cache=index_cache,
- block_table=block_table,
- query_positions=positions.to(spec["positions_dtype"]),
- request_indices=slot_of.to(DEV)[req].to(spec["request_dtype"]),
- tokens_per_block=spec["tokens_per_block"],
- compress_ratio=spec["compress_ratio"],
- only_visible_blocks=spec["only_visible_blocks"],
- context_rows=spec["context_rows"],
- )
-
-
-def _branch_rule(rows, spec):
- heads = spec["heads"]
- packs = DOT_TILE_WIDTH % heads == 0 and DOT_TILE_WIDTH // heads > 1
- tiled = spec["context_rows"] and rows > DECODE_ROW_LIMIT and packs
- return "query_tile" if tiled else "row_wise"
-
-
-def _columns(args):
- return args["block_table"].shape[1] * args["tokens_per_block"] // args["compress_ratio"]
-
-
-def _visible(args):
- blocks = torch.arange(_columns(args), device=DEV)
- bound = (args["query_positions"].long() + 1) // args["compress_ratio"]
- return blocks[None, :] < bound[:, None]
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811
- certifying(op)
- args = _build(cell.spec, seed=sum(map(ord, cell.why)))
- rows = args["q"].shape[0]
- assert cell.spec["branch"] == _branch_rule(rows, cell.spec), "the cell names the wrong kernel"
- with _watch_launches() as watcher:
- out = op(**args)
- torch.cuda.synchronize()
- assert watcher.branch == cell.spec["branch"]
- assert out.shape == (rows, _columns(args)) and out.dtype == torch.float32
- _compare(out, op.reference(**args))
-
-
-def test_prefetch_reads_the_fresh_anchor_key_and_q_after_the_wait() -> None:
- """`prefetch` reads pages and keys before the PDL wait. Under graph replay a producer
- rewrites q and, at rows whose position ends a group, the last visible block's anchor
- key -- what the decode prologue writes -- right before the op; both must be seen.
- The other keys are static, as they are in the decode step."""
- args = _build(_case(4, 4, only_visible_blocks=True), seed=411)
- cr, tpb = args["compress_ratio"], args["tokens_per_block"]
- pos = args["query_positions"].clone()
- assert bool((pos >= cr).all())
- ends = torch.tensor([True, False, True, True], device=DEV)
- args["query_positions"] = torch.where(ends, pos - (pos + 1) % cr, pos).to(pos.dtype)
- pos = args["query_positions"].long()
- assert bool(((pos + 1) % cr == 0).eq(ends).all())
- pages = args["block_table"][args["request_indices"].long(), pos // tpb].long()
- anchors = (pages[ends], (pos % tpb)[ends])
- visible = _visible(args)
- plain = op(**args)
- assert torch.equal(op(**{**args, "use_pdl": True, "prefetch": True})[visible], plain[visible])
- src_q = args["q"].clone()
- src_k = args["index_cache"][anchors].clone()
-
- def step():
- args["q"].copy_(src_q * 1.0)
- args["index_cache"][anchors] = src_k * 1.0
- return op(**{**args, "use_pdl": True, "prefetch": True})
-
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(warmup):
- step()
- torch.cuda.current_stream().wait_stream(warmup)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- graph_out = step()
- for trial in range(3):
- fresh = _build(_case(4, 4, only_visible_blocks=True), seed=412 + trial)
- src_q.copy_(fresh["q"])
- src_k.copy_(torch.randn(src_k.shape, device=DEV).to(src_k.dtype))
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(graph_out[visible], op(**args)[visible]), f"replay {trial}"
-
-
-def test_rejects_metadata_shorter_than_the_rows() -> None:
- """The callable's own refusal, message as observed -- the reason this entry has no
- `is_valid`: positions for fewer rows than `q` carries."""
- args = _build(_case(16, 2), seed=1701)
- args["query_positions"] = args["query_positions"][:4]
- with pytest.raises(ValueError, match="one entry per row"):
- op(**args)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py
deleted file mode 100644
index caf83f14f1c7..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py
+++ /dev/null
@@ -1,190 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the qsa_paged_kv_store catalog entry.
-
-The paged state is real: one `[pages, 2, Hkv, tpb, D]` HND pool whose K and V
-role views are the caches, and the engine's `[slots, width]` table -- a row
-per request slot, more slots than requests, each live row's pages drawn from a
-seeded permutation of the pool, `-1` past its allocation, and the rows past
-the batch holding a previous step's plausible page ids. Pool sizes are the old
-receipts'; the pool is pre-filled with a sentinel so a stray write shows.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_kv_store import ( # noqa: E501
- qsa_paged_kv_store as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401
-
-assert torch.cuda.is_available(), "qsa_paged_kv_store requires a CUDA device"
-
-
-def _compare(out, ref) -> None:
- """Bitwise over both role planes, compared as raw bits.
-
- A store is a move with one conversion, and `reference` performs the
- same conversion, so there is no rounding to spend a tolerance on. In
- the target both planes together are the whole pool allocation, so this
- is also the stray-write check.
- """
- for name, got, want in zip(("k_cache", "v_cache"), out, ref):
- assert got.dtype == want.dtype, (name, got.dtype, want.dtype)
- assert got.shape == want.shape, (name, got.shape, want.shape)
- torch.testing.assert_close(
- _bits(got), _bits(want), rtol=0, atol=0, msg=lambda m, n=name: f"{n}: {m}"
- )
-
-
-DEV = "cuda"
-FILL = -2.0
-
-
-def _bits(tensor: torch.Tensor) -> torch.Tensor:
- as_int = {1: torch.uint8, 2: torch.int16, 4: torch.int32}[tensor.element_size()]
- return tensor.contiguous().view(as_int)
-
-
-def _hole(spec, table, allocated_cols):
- """One row that resolves to no slot: its (request, position)."""
- assert spec["hole"] == "unallocated_page", spec["hole"]
- assert int(table[0, allocated_cols[0]]) == -1
- return 0, allocated_cols[0] * spec["tokens_per_block"]
-
-
-def _build(spec, seed: int):
- """Real pool, table and rows for one cell; returns the op's kwargs and the pool."""
- gen = torch.Generator(device="cpu").manual_seed(seed)
- tpb, hkv, head_dim = spec["tokens_per_block"], spec["hkv"], spec["head_dim"]
- slots, width, pages = spec["slots"], spec["width"], spec["pages"]
- segments = spec["segments"]
- requests = [r for r, (_, n) in enumerate(segments) for _ in range(n)]
- positions = [p for start, n in segments for p in range(start, start + n)]
-
- table = torch.randint(0, pages, (slots, width), generator=gen, dtype=torch.int32)
- order = torch.randperm(pages, generator=gen).to(torch.int32)
- used, allocated_cols = 0, []
- for r, (start, n) in enumerate(segments):
- cols = (start + n - 1) // tpb + 1
- table[r] = -1
- table[r, :cols] = order[used : used + cols]
- used += cols
- allocated_cols.append(cols)
- assert used <= pages and len(segments) < slots, "pool or table too small for the batch"
- if spec["hole"] is not None:
- request, position = _hole(spec, table, allocated_cols)
- requests.append(request)
- positions.append(position)
-
- rows = len(positions)
- qkv = torch.randn(
- rows,
- spec["heads_q"] + 2 * hkv + spec.get("row_pad_heads", 0),
- head_dim,
- generator=gen,
- dtype=torch.float32,
- ).to(device=DEV, dtype=spec["row_dtype"])
- k = qkv[:, spec["heads_q"] : spec["heads_q"] + hkv]
- v = qkv[:, spec["heads_q"] + hkv : spec["heads_q"] + 2 * hkv]
- table = table.to(DEV)
-
- pool = torch.full((pages, 2, hkv, tpb, head_dim), FILL, dtype=torch.bfloat16, device=DEV).to(
- spec["cache_dtype"]
- )
- args = dict(
- k=k,
- v=v,
- k_cache=pool[:, 0],
- v_cache=pool[:, 1],
- request_indices=torch.tensor(requests, dtype=spec["request_dtype"], device=DEV),
- logical_positions=torch.tensor(positions, dtype=spec["position_dtype"], device=DEV),
- block_table=table,
- tokens_per_block=tpb,
- )
- return args, pool
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying): # noqa: F811
- certifying(op)
- args, pool = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
- before = pool.clone()
- ref = op.reference(**args)
- out = op(**args)
- assert out[0] is args["k_cache"] and out[1] is args["v_cache"]
- _compare(out, ref)
- wrote = not torch.equal(_bits(pool), _bits(before))
- assert wrote == bool(cell.spec["segments"]), (
- "a cell with rows must write, an empty one must not"
- )
-
-
-def _decode_args(seed: int, **overrides):
- """The decode cell's layout over a bf16 pool, without its skipped row."""
- spec = {**op.CELLS[0].spec, "hole": None}
- return _build({**spec, **overrides}, seed=seed)
-
-
-#: The cell of 32 MTP7 verify blocks (256 rows): the segments, pool and table it runs over.
-_MTP7_32 = next(c.spec for c in op.CELLS if c.why.startswith("fp8 decode MTP7: 32 requests"))
-_MTP7_BATCH = {key: _MTP7_32[key] for key in ("segments", "pages", "slots")}
-
-
-def test_cuda_graph_replay_follows_updated_positions() -> None:
- """Captured once at 32 MTP7 verify blocks, replayed against new positions written
- into the same buffer."""
- args, pool = _decode_args(seed=9, **_MTP7_BATCH)
- positions = args["logical_positions"]
- saved = positions.clone()
- positions.zero_()
-
- def call():
- op(**args)
-
- for _ in range(3): # compile the Triton kernel before capture
- call()
- torch.cuda.synchronize()
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- call()
-
- pool.fill_(FILL)
- positions.copy_(saved)
- ref = op.reference(**args)
- graph.replay()
- torch.cuda.synchronize()
- _compare((args["k_cache"], args["v_cache"]), ref)
-
-
-def test_guard_refuses_strided_row_metadata() -> None:
- """`is_valid`'s one guard, after the hazard it exists for.
-
- The kernel receives strides for k, v, both caches and the table, but reads
- the two metadata vectors as flat unit-stride buffers. A column of a 2-wide
- buffer is read through its base pointer: accepted, and rows land elsewhere.
- """
- args, pool = _decode_args(seed=10)
- requests = args["request_indices"]
- rows = requests.numel()
- wide_requests = torch.full((rows, 2), rows - 1, dtype=requests.dtype, device=DEV)
- wide_requests[:, 0] = requests
- ref = op.reference(**args)
- strided = {**args, "request_indices": wide_requests[:, 0]}
-
- pool.fill_(FILL)
- out = op(**strided) # unguarded: accepted
- assert not all(torch.equal(_bits(got), _bits(want)) for got, want in zip(out, ref)), (
- "strided metadata was expected to place rows elsewhere"
- )
-
- pool.fill_(FILL)
- before = pool.clone()
- with pytest.raises(AssertionError, match="must be a contiguous 1-D tensor"):
- with validating(op):
- op(**strided)
- assert torch.equal(_bits(pool), _bits(before)), "a refused call still wrote"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_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
deleted file mode 100644
index ce3ddec3389e..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py
+++ /dev/null
@@ -1,369 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the qsa_paged_kv_store_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._torch.attention.backends.fmha.fallback import FallbackFmha
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401
-
-assert torch.cuda.is_available(), "qsa_paged_kv_store_nvfp4 requires a CUDA device"
-
-DEV = "cuda"
-DATA_FILL, SCALE_FILL = 0xA5, 0x5A
-QUANT_MODE_NVFP4_KV = 8192
-
-
-def _compare(out, ref) -> None:
- """Bitwise over all four planes: the store is a pure function of the row's
- bits, and `reference` performs the same rounding steps."""
- names = ("k_cache", "v_cache", "k_scale_cache", "v_scale_cache")
- for name, got, want in zip(names, out, ref):
- assert got.dtype == want.dtype and got.shape == want.shape, (
- name,
- got.shape,
- want.shape,
- )
- differ = (got != want).nonzero()
- if differ.numel():
- first = tuple(differ[0].tolist())
- raise AssertionError(
- f"{name}: {differ.shape[0]} bytes differ from the reference, first at "
- f"{first}: {int(got[first]):#04x} vs {int(want[first]):#04x}"
- )
-
-
-def _hole(spec, table, allocated_cols):
- """One row that resolves to no slot: its (request, position)."""
- assert spec["hole"] == "unallocated_page", spec["hole"]
- assert int(table[0, allocated_cols[0]]) == -1
- return 0, allocated_cols[0] * spec["tokens_per_block"]
-
-
-def _fill_rows(spec, k, v, gen):
- """Row values: random with per-block magnitudes over 2^-10..2^5, plus edge blocks
- (zero, -0.0, an outlier) when the cell asks for them."""
- hkv, head_dim = spec["hkv"], spec["head_dim"]
- rows = k.shape[0]
- for t in (k, v):
- base = torch.randn(rows, hkv, head_dim // 16, 16, generator=gen)
- exponent = torch.randint(-10, 6, (rows, hkv, head_dim // 16, 1), generator=gen)
- t.copy_((base * torch.pow(2.0, exponent.float())).reshape(t.shape).to(DEV))
- if spec["values"] == "edge":
- k[0, 0, :16] = 0.0
- v[0, 1, 16:32] = -0.0
- k[1, 1, 32:48] = torch.tensor([-0.0] * 8 + [0.0] * 8)
- k[2, 0, 5] = 900.0
- v[3, 0, 200] = -900.0
- k[4, 1, 64:80] = 1e-30 # far below the smallest e4m3 scale
- else:
- assert spec["values"] == "randn"
-
-
-def _pool(spec):
- """The allocation the cell's planes view: data and scales of `slot_pages` role
- pages per slot, every byte its sentinel."""
- pages, slot_pages = spec["pages"], spec["slot_pages"]
- hkv, tpb, head_dim = spec["hkv"], spec["tokens_per_block"], spec["head_dim"]
- data = torch.full(
- (pages, slot_pages, hkv, tpb, head_dim // 2), DATA_FILL, dtype=torch.uint8, device=DEV
- )
- scales = torch.full(
- (pages, slot_pages, hkv, tpb, head_dim // 16), SCALE_FILL, dtype=torch.uint8, device=DEV
- )
- return data, scales
-
-
-def _planes(args):
- return (args["k_cache"], args["v_cache"], args["k_scale_cache"], args["v_scale_cache"])
-
-
-def _reset(planes) -> None:
- for plane, fill in zip(planes, (DATA_FILL, DATA_FILL, SCALE_FILL, SCALE_FILL)):
- plane.fill_(fill)
-
-
-def _assert_other_roles_untouched(spec, pool) -> None:
- """Every role page of a slot but the call's K and V still holds its sentinel."""
- data, scales = pool
- for role in range(spec["slot_pages"]):
- if role // 2 != spec["layer"]:
- assert bool((data[:, role] == DATA_FILL).all()), f"data role {role} moved"
- assert bool((scales[:, role] == SCALE_FILL).all()), f"scale role {role} moved"
-
-
-def _build(spec, seed: int, pool=None):
- """Real pool, table and rows for one cell; returns the op's kwargs, the pool,
- and the packed QKV rows the K/V slices view. `pool` reuses an allocation
- from `_pool(spec)`."""
- gen = torch.Generator(device="cpu").manual_seed(seed)
- tpb, hkv, head_dim = spec["tokens_per_block"], spec["hkv"], spec["head_dim"]
- slots, width, pages = spec["slots"], spec["width"], spec["pages"]
- segments = spec["segments"]
- requests = [r for r, (_, n) in enumerate(segments) for _ in range(n)]
- positions = [p for start, n in segments for p in range(start, start + n)]
-
- table = torch.randint(0, pages, (slots, width), generator=gen, dtype=torch.int32)
- order = torch.randperm(pages, generator=gen).to(torch.int32)
- used, allocated_cols = 0, []
- for r, (start, n) in enumerate(segments):
- cols = (start + n - 1) // tpb + 1
- table[r] = -1
- table[r, :cols] = order[used : used + cols]
- used += cols
- allocated_cols.append(cols)
- assert used <= pages and len(segments) <= slots, "pool or table too small for the batch"
- if spec["hole"] is not None:
- request, position = _hole(spec, table, allocated_cols)
- requests.append(request)
- positions.append(position)
-
- rows = len(positions)
- qkv = torch.zeros(
- rows, spec["heads_q"] + 2 * hkv, head_dim, dtype=spec["row_dtype"], device=DEV
- )
- k = qkv[:, spec["heads_q"] : spec["heads_q"] + hkv]
- v = qkv[:, spec["heads_q"] + hkv :]
- qkv[:, : spec["heads_q"]] = torch.randn(rows, spec["heads_q"], head_dim, generator=gen).to(
- device=DEV, dtype=spec["row_dtype"]
- )
- _fill_rows(spec, k, v, gen)
-
- layer = spec["layer"]
- data, scales = _pool(spec) if pool is None else pool
- args = dict(
- k=k,
- v=v,
- k_cache=data[:, 2 * layer],
- v_cache=data[:, 2 * layer + 1],
- k_scale_cache=scales[:, 2 * layer],
- v_scale_cache=scales[:, 2 * layer + 1],
- request_indices=torch.tensor(requests, dtype=spec["request_dtype"], device=DEV),
- logical_positions=torch.tensor(positions, dtype=spec["position_dtype"], device=DEV),
- block_table=table.to(DEV),
- tokens_per_block=tpb,
- kv_scale_orig_quant=torch.tensor(
- [1.0, spec["global_k"], spec["global_v"]], dtype=torch.float32, device=DEV
- ),
- )
- return args, (data, scales), qkv
-
-
-def _snapshot(pool):
- return tuple(t.clone() for t in pool)
-
-
-def _pools_equal(a, b) -> bool:
- return all(torch.equal(x, y) for x, y in zip(a, b))
-
-
-def _graph_replay(spec, seed: int, graph_batches) -> None:
- """Each G: capture one step of the first G requests, replay the next step.
-
- The next step lives in the captured buffers: new k/v rows, new pages in the
- table, and every request's positions advanced by its rows, the way a decode
- step moves. The replay must equal the reference and an eager call on the
- same inputs, bit for bit. One pool serves the whole sweep, as the engine's does.
- """
- pool = _pool(spec)
- steps = {n for _, n in spec["segments"]}
- assert len(steps) == 1, "a graph cell's requests all carry the same row count"
- step = steps.pop()
- for g in graph_batches:
- sub = {**spec, "segments": spec["segments"][:g]}
- args, _, _ = _build(sub, seed=seed + g, pool=pool)
- following, _, _ = _build(sub, seed=seed + 1000 + g, pool=pool)
- planes = _planes(args)
- positions = args["logical_positions"]
- next_positions = positions.clone()
- positions.sub_(step)
- assert int(positions.min()) >= 0
-
- def call():
- op(**args)
-
- for _ in range(3): # compile the Triton kernel before capture
- call()
- torch.cuda.synchronize()
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- call()
-
- for name in ("k", "v", "block_table"):
- args[name].copy_(following[name])
- positions.copy_(next_positions)
- _reset(planes)
- op(**args)
- eager = tuple(plane.clone() for plane in planes)
- _reset(planes)
- ref = op.reference(**args)
- graph.replay()
- torch.cuda.synchronize()
- _compare(planes, ref)
- assert all(torch.equal(p, e) for p, e in zip(planes, eager)), (
- f"G={g} x {step} rows: replay differs from eager"
- )
- del graph
- _assert_other_roles_untouched(spec, pool)
-
-
-def _thop_append(qkv, pool, table, requests, new, cached, num_contexts, p, scale_orig_quant):
- """The attention op's own append over the same rows (e4m3 output: the append
- does not depend on the output type)."""
- data, scales = pool
- tpb = 32
- hq, hkv, d = 24, 2, 256
- ns = len(requests)
- width = table.shape[1]
- offsets = torch.zeros(1, ns, 2, width, dtype=torch.int32, device=DEV)
- for row, r in enumerate(requests):
- live = table[r].clamp_min(0)
- offsets[0, row, 0] = 2 * live
- offsets[0, row, 1] = 2 * live + 1
- pointers = torch.zeros(1, 2, 2, dtype=torch.int64)
- pointers[0, 0, 0], pointers[0, 0, 1] = data.data_ptr(), scales.data_ptr()
- kv = [c + n for c, n in zip(cached, new)]
- ctx = [n if i < num_contexts else kv[i] for i, n in enumerate(new)]
- scale_quant_orig = scale_orig_quant.reciprocal()
- rows = qkv.reshape(qkv.shape[0], -1)
- FallbackFmha.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=None,
- spec_decoding_packed_mask=None, spec_decoding_bl_tree_mask_offset=None,
- spec_decoding_bl_tree_mask=None, spec_bl_tree_first_sparse_mask_offset_kv=None,
- sparse_kv_indices=None, sparse_kv_offsets=None, sparse_attn_indices=None,
- sparse_attn_offsets=None, sparse_attn_indices_block_size=0, num_contexts=num_contexts,
- num_ctx_tokens=sum(new[:num_contexts]),
- ) # fmt: skip
- torch.cuda.synchronize()
-
-
-def _thop_equality(spec, seed: int) -> None:
- """Steps of a context, a chunked continuation, decode P=1 and MTP verify P=3 over
- three requests, each written once by the attention op's append into one pool and
- once by this op into another, at the cell's second-level scales: the two pools stay
- bitwise equal, and this op's output equals its reference at every step."""
- base, ours, _ = _build(spec, seed=seed)
- theirs = _snapshot(ours)
- gen = torch.Generator(device="cpu").manual_seed(seed + 1)
- table = base["block_table"]
- cached = [0, 0, 0]
- for new, nc, p in (
- ([70, 33, 5], 3, 1),
- ([41, 64, 27], 3, 1),
- ([1, 1, 1], 0, 1),
- ([3] * 3, 0, 3),
- ):
- rows = sum(new)
- qkv = torch.randn(rows, 28, 256, generator=gen)
- qkv *= torch.pow(2.0, torch.randint(-6, 4, (rows, 28, 1), generator=gen).float())
- qkv = qkv.to(device=DEV, dtype=torch.bfloat16)
- requests = [r for r, n in enumerate(new) for _ in range(n)]
- positions = [c + i for c, n in zip(cached, new) for i in range(n)]
- _thop_append(qkv, theirs, table, [0, 1, 2], new, cached, nc, p, base["kv_scale_orig_quant"])
- args = {
- **base,
- "k": qkv[:, 24:26],
- "v": qkv[:, 26:28],
- "request_indices": torch.tensor(requests, dtype=torch.int32, device=DEV),
- "logical_positions": torch.tensor(positions, dtype=torch.int64, device=DEV),
- }
- ref = op.reference(**args)
- out = op(**args)
- _compare(out, ref)
- for plane, (a, b) in enumerate(zip(ours, theirs)):
- assert torch.equal(a, b), (
- f"step {new}: {('data', 'scales')[plane]} differ from the attention op's "
- f"append in {int((a != b).sum())} bytes"
- )
- cached = [c + n for c, n in zip(cached, new)]
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying): # noqa: F811
- certifying(op)
- seed = zlib.crc32(cell.why.encode())
- if cell.spec["thop_append"]:
- _thop_equality(cell.spec, seed=seed)
- return
- args, pool, _ = _build(cell.spec, seed=seed)
- before = _snapshot(pool)
- ref = op.reference(**args)
- out = op(**args)
- assert out[0] is args["k_cache"] and out[3] is args["v_scale_cache"]
- _compare(out, ref)
- wrote = not _pools_equal(pool, before)
- assert wrote == bool(cell.spec["segments"]), (
- "a cell with rows must write, an empty one must not"
- )
- _assert_other_roles_untouched(cell.spec, pool)
-
-
-def test_cuda_graph_replay_matches_eager() -> None:
- """Decode MTP under a graph: 8 requests x 8 rows in the prefill cell's trunk layout,
- captured once, replayed on the next step's rows, pages and positions."""
- prefill = next(c.spec for c in op.CELLS if c.why.startswith("prefill:"))
- spec = {**prefill, "segments": tuple((29 + 61 * i, 8) for i in range(8)), "pages": 256}
- _graph_replay(spec, seed=8, graph_batches=(8,))
-
-
-def _decode_args(seed: int, **overrides):
- return _build({**op.CELLS[0].spec, **overrides}, seed=seed)
-
-
-def test_op_rejects_strided_request_indices() -> None:
- """The callable's own loud error -- why the entry carries no `is_valid` -- for the
- metadata a flat-pointer kernel would otherwise read at the wrong rows."""
- args, _, _ = _decode_args(seed=5)
- strided = torch.stack([args["request_indices"]] * 2, 1)[:, 0]
- with pytest.raises(ValueError, match="request_indices must be contiguous"):
- with validating(op):
- op(**{**args, "request_indices": strided})
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py
deleted file mode 100644
index 28ed2deb6c8a..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py
+++ /dev/null
@@ -1,250 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the qsa_paged_sparse_gqa catalog entry.
-
-The paged state is real: one `[pages, 2, kv, tpb, D]` HND pool sized from the
-pages its requests map, and the engine's slot table -- a row per request slot,
-more slots than requests, assigned by a seeded permutation, `-1` past each
-sequence. Selection rows are built the way QSA's top-k + expansion lay them
-out: complete groups expanded to tokens, then the trailing partial group, then
-`-1`, all packed at the front.
-
-The call has two kernels behind it (and a merge). The `launches` fixture
-records which ones a call indexed with a grid, so every cell asserts its
-branch.
-"""
-
-import math
-import re
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_sparse_gqa import ( # noqa: E501
- qsa_paged_sparse_gqa as op,
-)
-from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels as _qsa_kernels
-
-__extra_import_path__ = ["../../../.."]
-from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401
-
-assert torch.cuda.is_available(), "qsa_paged_sparse_gqa requires a CUDA device"
-
-
-def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
- """4 ulp of each (row, head)'s magnitude, 2 ulp relative RMS, in q's dtype.
-
- The kernel rounds its online-softmax probabilities to the cache dtype
- before P.V, so the error tracks the magnitude of the values averaged,
- not of the output: against an fp32 reference it measured 1.2e-3 on rows
- averaging a thousand tokens (outputs 0.14 to 0.4) and 1.2e-2 on rows
- averaging a handful (outputs to 5). Scaled by each head row's own
- magnitude that is one band. Against this reference (fp32, rounded once)
- every cell measured at most 2.0 element ulp -- a single final-rounding
- flip, since bf16's step is 2 of these units at the bottom of a binade --
- and at most 0.62 ulp RMS: 2x and 3.2x margin. Six wrong references (no
- scale, interleaved heads, padding attended, identity pages, K/V
- swapped, page/slot transposed) were each measured more than 0.1 off.
- """
- assert out.dtype == ref.dtype, (out.dtype, ref.dtype)
- assert out.shape == ref.shape, (out.shape, ref.shape)
- if out.numel() == 0:
- return
- assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0)
-
-
-DEV = "cuda"
-
-_KERNELS = {
- "_qsa_paged_sparse_gqa_kernel": "prefill",
- "_qsa_paged_sparse_gqa_splitk_kernel": "splitk",
- "_qsa_merge_splitk_kernel": "merge",
-}
-
-
-class _LaunchRecorder:
- """Records a Triton kernel's grid indexing, then delegates to the kernel."""
-
- def __init__(self, inner, name, log):
- self._inner, self._name, self._log = inner, name, log
-
- def __getitem__(self, grid):
- self._log.append((self._name, tuple(grid)))
- return self._inner[grid]
-
-
-@pytest.fixture
-def launches(monkeypatch):
- """Which branch kernels the next call launches, in order, with grids."""
- log: list = []
- for name, tag in _KERNELS.items():
- inner = getattr(_qsa_kernels, name)
- monkeypatch.setattr(_qsa_kernels, name, _LaunchRecorder(inner, tag, log))
- return log
-
-
-# --------------------------------------------------------------------------
-# building real state from a cell spec
-# --------------------------------------------------------------------------
-
-
-def _selected_tokens(positions, seed, compress_ratio, token_topk):
- """QSA's expansion layout, built in one pass.
-
- A random `min(available, block_topk)`-subset of the complete visible
- groups, each expanded to its `compress_ratio` tokens, then the incomplete
- trailing group, then `-1` -- packed at the front, as the producer emits.
- """
- block_topk = token_topk // compress_ratio
- final_topk = token_topk + compress_ratio - 1
- gen = torch.Generator(device=DEV).manual_seed(seed)
- rows = positions.numel()
- visible = positions.long() + 1
- available = visible // compress_ratio
- taken = available.clamp(max=block_topk)
- width = max(block_topk, int(available.max().item()) if rows else 0)
-
- scores = torch.rand(rows, width, generator=gen, device=DEV)
- group_ids = torch.arange(width, device=DEV)[None, :]
- scores = torch.where(group_ids < available[:, None], scores, torch.inf)
- chosen = scores.argsort(dim=1)[:, :block_topk]
-
- column = torch.arange(final_topk, device=DEV)[None, :]
- group_of_column = (column // compress_ratio).clamp(max=block_topk - 1).expand(rows, -1)
- body = torch.gather(chosen, 1, group_of_column) * compress_ratio + column % compress_ratio
- in_body = column < (taken * compress_ratio)[:, None]
- tail_column = column - (taken * compress_ratio)[:, None]
- tail = (available * compress_ratio)[:, None] + tail_column
- in_tail = ~in_body & (tail_column >= 0) & (tail < visible[:, None])
- selected = torch.where(in_body, body, torch.where(in_tail, tail, torch.full_like(body, -1)))
- return selected.to(torch.int32)
-
-
-def _build(spec, seed):
- """Real paged state for one spec, keyed by `__call__`'s parameter names."""
- gen = torch.Generator(device=DEV).manual_seed(seed)
- tpb, kv, d, q_heads = (
- spec["tokens_per_block"],
- spec["kv_heads"],
- spec["head_dim"],
- spec["q_heads"],
- )
- segments = list(spec["segments"])
- positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments])
- req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)])
- rows, slots = positions.numel(), spec["slots"]
- assert len(segments) <= slots
-
- slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(segments)]
- mapped = [math.ceil((s + n) / tpb) for s, n in segments]
- assert max(mapped) <= spec["pages"]
- pool_pages = sum(mapped) + 7
- page_ids = torch.randperm(pool_pages, generator=gen, device=DEV).to(torch.int32)
- block_table = torch.full((slots, spec["pages"]), -1, dtype=torch.int32, device=DEV)
- cursor = 0
- for slot, count in zip(slot_of.tolist(), mapped):
- block_table[slot, :count] = page_ids[cursor : cursor + count]
- cursor += count
-
- pool = torch.randn(pool_pages, 2, kv, tpb, d, generator=gen, device=DEV).to(spec["cache_dtype"])
- k_cache, v_cache = pool[:, 0], pool[:, 1]
-
- q_dtype = spec["q_dtype"]
- row_heads = q_heads + 2 * kv if spec["q_layout"] == "qkv_slice" else 55
- qkv = torch.randn(rows, row_heads, d, generator=gen, device=DEV).to(q_dtype)
- q = qkv[:, :q_heads]
-
- selected = _selected_tokens(positions, seed + 3, spec["compress_ratio"], spec["token_topk"])
- request_indices = slot_of.to(DEV)[req].to(spec["request_dtype"])
-
- gate = None
- if spec.get("gate") == "fused_row":
- # The gate columns of the target's fused projection row (55 heads of
- # D): heads 28..51, after q, k and v.
- fused = torch.randn(rows, 55 * d, generator=gen, device=DEV).to(q_dtype)
- gate = fused[:, 28 * d : (28 + q_heads) * d]
-
- args = dict(
- q=q,
- k_cache=k_cache,
- v_cache=v_cache,
- block_table=block_table,
- selected_tokens=selected,
- request_indices=request_indices,
- tokens_per_block=tpb,
- softmax_scale=d**-0.5,
- )
- if spec["prefill_pair"]:
- args.update(query_positions=positions, compress_ratio=spec["compress_ratio"])
- if gate is not None:
- args.update(output_gate=gate)
- return args
-
-
-def _bound(positions, spec):
- """The prefill kernel's per-row scan bound, from the positions alone."""
- cr = spec["compress_ratio"]
- visible = positions.long() + 1
- return (visible // cr).clamp(max=spec["token_topk"] // cr) * cr + visible % cr
-
-
-def _is_prefix_packed(table):
- valid = table >= 0
- columns = torch.arange(table.shape[1], device=table.device)[None, :]
- return (valid == (columns < valid.sum(1)[:, None])).all(1)
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811
- certifying(op)
- spec = cell.spec
- args = _build(spec, seed=sum(map(ord, cell.why)))
- rows = args["q"].shape[0]
- if spec["prefill_pair"]:
- # The precondition the op cannot check: front-packed, and exactly as
- # many valid entries as the prefill kernel will scan.
- selected = args["selected_tokens"]
- assert bool(_is_prefix_packed(selected).all())
- valid = (selected >= 0).sum(1)
- assert torch.equal(valid, _bound(args["query_positions"], spec))
-
- out = op(**args)
- assert [name for name, _ in launches] == list(spec["launches"])
- assert launches[0][1][:2] == (rows, spec["kv_heads"])
- assert out.shape == (rows, spec["q_heads"], spec["head_dim"])
- assert out.dtype == args["q"].dtype and out.is_contiguous()
- assert not bool(out.isnan().any())
- _compare(out, op.reference(**args))
-
-
-#: The first cell is this checkpoint's own decode geometry; the tests below
-#: start from it.
-CKPT = op.CELLS[0].spec
-assert (CKPT["q_heads"], CKPT["kv_heads"], CKPT["head_dim"]) == (24, 2, 256)
-
-
-def _with(**overrides):
- return {**CKPT, **overrides}
-
-
-def test_calls_are_bitwise_deterministic() -> None:
- """Including a 32-split launch whose merge is a second kernel."""
- spec = _with(segments=tuple((4000 + 37 * i, 2) for i in range(4)), gate=None)
- args = _build(spec, seed=8800)
- assert torch.equal(op(**args), op(**args))
-
-
-# --------------------------------------------------------------------------
-# what the op refuses
-# --------------------------------------------------------------------------
-
-
-def test_rejects_positions_without_a_compress_ratio() -> None:
- """The op's own loud refusal, verbatim -- the reason `is_valid` is empty: half of
- the prefill pair."""
- args = _build(_with(segments=((600, 1),) * 4, gate=None), seed=9100)
- args["selected_tokens"] = args["selected_tokens"][:, :64].contiguous()
- args["query_positions"] = torch.zeros(4, dtype=torch.int64, device=DEV)
- message = "QSA causal-prefix bounds require positions and compression ratio"
- with pytest.raises(ValueError, match=re.escape(message)):
- op(**args)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_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
deleted file mode 100644
index 7d24c80e8471..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.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 qsa_paged_sparse_gqa_decode entry: cells against the composition
-entry's fp32 reference in that entry's band (4 ulp of each head row's magnitude, 2 ulp
-relative RMS), the selection's read-after-wait under graph replay with PDL, and the row
-guard."""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_sparse_gqa_decode import ( # noqa: E501
- qsa_paged_sparse_gqa_decode as op,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_topk_expand_decode import ( # noqa: E501
- qsa_topk_expand_decode,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import ( # noqa: E402,F401 — path above; fixture by name
- assert_within_ulp,
- certifying,
- validating,
-)
-
-assert torch.cuda.is_available(), "qsa_paged_sparse_gqa_decode requires a CUDA device"
-
-DEV = "cuda"
-Q_HEADS, KV_HEADS, D, TPB, CR, TOKEN_TOPK, PAGES, SLOTS = 24, 2, 256, 32, 4, 2048, 320, 6
-
-
-def _build(spec, seed):
- """Paged K/V for one decode row per position, the selection QSA's expansion emits
- (512 random complete groups, then the incomplete group, then -1), q and the gate as
- column views of the target's 55-head fused projection row."""
- g = torch.Generator(device=DEV).manual_seed(seed)
- pos = torch.tensor(spec["pos"], dtype=torch.int64, device=DEV)
- rows = pos.numel()
- pages_per = [int(p) // TPB + 1 for p in spec["pos"]]
- pool_pages = sum(pages_per) + 5
- page_ids = torch.randperm(pool_pages, generator=g, device=DEV).to(torch.int32)
- block_table = torch.full((SLOTS, PAGES), -1, dtype=torch.int32, device=DEV)
- slot_of = list(range(1, rows + 1))
- cursor = 0
- for slot, count in zip(slot_of, pages_per):
- block_table[slot, :count] = page_ids[cursor : cursor + count]
- cursor += count
- pool = torch.randn(pool_pages, 2, KV_HEADS, TPB, D, generator=g, device=DEV).to(torch.bfloat16)
- k_cache, v_cache = pool[:, 0], pool[:, 1]
- fused = torch.randn(rows, 55 * D, generator=g, device=DEV).to(torch.bfloat16)
- q = fused.view(rows, 55, D)[:, :Q_HEADS]
- gate = fused[:, 28 * D : (28 + Q_HEADS) * D] if spec["gated"] else None
- block_topk = TOKEN_TOPK // CR
- selected = torch.full((rows, TOKEN_TOPK + CR - 1), -1, dtype=torch.int32, device=DEV)
- for r, p in enumerate(spec["pos"]):
- visible = p + 1
- groups = visible // CR
- chosen = torch.randperm(groups, generator=g, device=DEV)[: min(groups, block_topk)]
- toks = (chosen[:, None] * CR + torch.arange(CR, device=DEV)).reshape(-1)
- tail = torch.arange(groups * CR, visible, device=DEV)
- row = torch.cat([toks, tail]).to(torch.int32)
- selected[r, : row.numel()] = row
- if spec.get("holes"):
- selected[r, 1 : row.numel() : 2] = -1
- request_indices = torch.tensor(slot_of, dtype=torch.int32, device=DEV)
- return [q, k_cache, v_cache, block_table, selected, request_indices, TPB, D**-0.5, gate]
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
- certifying(op)
- args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31))
- out = op(*args)
- ref = op.reference(*args)
- assert out.dtype == torch.bfloat16 and out.shape == ref.shape and out.is_contiguous()
- assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0)
-
-
-def _replays_like_eager(step, refresh) -> None:
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(warmup):
- step()
- torch.cuda.current_stream().wait_stream(warmup)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- graph_out = step()
- for trial in range(3):
- refresh()
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(graph_out, step()), f"replay {trial}"
-
-
-def test_reads_the_selection_after_the_wait_under_graph_replay() -> None:
- """The target's top-k runs inside the graph right before the op: a PDL producer that
- releases the op right after its own wait, before it writes the selection. A read of the
- selection before the op's wait would see the previous replay's."""
- args = _build(dict(pos=(8700,), gated=True), seed=23)
- pos = torch.tensor([8700], dtype=torch.int64, device=DEV)
- ends = ((pos + 1) // CR).to(torch.int32)
- seqlen = (pos + 1).to(torch.int32)
- scores = torch.rand(1, 2304, device=DEV)
- topk = torch.empty(1, TOKEN_TOPK // CR, dtype=torch.int32, device=DEV)
-
- def step():
- sel = qsa_topk_expand_decode(scores, ends, pos, seqlen, topk, CR, TOKEN_TOPK)
- return op(*args[:4], sel, *args[5:])
-
- _replays_like_eager(step, lambda: scores.copy_(torch.rand(scores.shape, device=DEV)))
-
-
-def test_guard_refuses_rows_past_max_rows() -> None:
- args = _build(dict(pos=(8700,), gated=True), seed=5)
- many = [a.expand(300, *a.shape[1:]) if i in (0, 4, 5) else a for i, a in enumerate(args[:6])]
- with validating(op), pytest.raises(AssertionError):
- op(*many, *args[6:8], None) # 300 rows: past MAX_ROWS
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py
deleted file mode 100644
index ab4313b61e21..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py
+++ /dev/null
@@ -1,321 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the qsa_paged_sparse_gqa_nvfp4 entry.
-
-The paged state is real: one NVFP4 pool as the QSA cache manager serves it --
-e2m1 data `[pages, 2, kv, tpb, D/2]` and e4m3 block scales `[pages, 2, kv, tpb,
-D/16]` -- sized from the pages its requests map, and the engine's slot table,
-a row per request slot, more slots than requests, assigned by a seeded
-permutation, `-1` past each sequence. Selection rows are built the way QSA's
-top-k + expansion lay them out, packed at the front.
-
-Beside `_compare` against the reference, every cell runs QSA's bf16 sparse GQA
-over a bf16 pool that holds the dequantized values and holds the two outputs to
-the same band: the NVFP4 kernels are the bf16 ones with a dequantizing load.
-"""
-
-import math
-import re
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention import ( # noqa: E501
- qsa_paged_sparse_gqa_nvfp4 as _entry,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_sparse_gqa_nvfp4 import ( # noqa: E501
- qsa_paged_sparse_gqa_nvfp4 as op,
-)
-from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import triton_qsa_paged_sparse_gqa
-
-__extra_import_path__ = ["../../../.."]
-from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401
-
-assert torch.cuda.is_available(), "qsa_paged_sparse_gqa_nvfp4 requires a CUDA device"
-
-DEV = "cuda"
-ENGINE_MAX_SEQ_LEN = 8192
-#: The e4m3fn NaN bit pattern.
-E4M3_NAN = 0x7F
-
-_KERNELS = {
- "_nvfp4_sparse_gqa_prefill_kernel": "prefill",
- "_nvfp4_sparse_gqa_splitk_kernel": "splitk",
- "_nvfp4_merge_splitk_kernel": "merge",
-}
-
-
-class _LaunchRecorder:
- """Records a Triton kernel's grid indexing, then delegates to the kernel."""
-
- def __init__(self, inner, name, log):
- self._inner, self._name, self._log = inner, name, log
-
- def __getitem__(self, grid):
- self._log.append((self._name, tuple(grid)))
- return self._inner[grid]
-
-
-@pytest.fixture
-def launches(monkeypatch):
- """Which branch kernels the next call launches, in order, with grids."""
- log: list = []
- for name, tag in _KERNELS.items():
- monkeypatch.setattr(_entry, name, _LaunchRecorder(getattr(_entry, name), tag, log))
- return log
-
-
-def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
- """4 ulp of each (row, head)'s magnitude, 2 ulp relative RMS, in q's dtype.
-
- After the dequantize the kernels are the bf16 sparse GQA kernels, whose band
- this is (they round the online-softmax probabilities to the value dtype
- before P.V); the dequantize itself is exact at unit scale and rounded once,
- as `reference` rounds it, otherwise.
- """
- assert out.dtype == ref.dtype, (out.dtype, ref.dtype)
- assert out.shape == ref.shape, (out.shape, ref.shape)
- if out.numel() == 0:
- return
- assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0)
-
-
-def _selected_tokens(positions, seed, compress_ratio, token_topk):
- """QSA's expansion layout: chosen complete groups, the trailing partial group, -1."""
- block_topk = token_topk // compress_ratio
- final_topk = token_topk + compress_ratio - 1
- gen = torch.Generator(device=DEV).manual_seed(seed)
- rows = positions.numel()
- visible = positions.long() + 1
- available = visible // compress_ratio
- taken = available.clamp(max=block_topk)
- width = max(block_topk, int(available.max().item()) if rows else 0)
- scores = torch.rand(rows, width, generator=gen, device=DEV)
- group_ids = torch.arange(width, device=DEV)[None, :]
- scores = torch.where(group_ids < available[:, None], scores, torch.inf)
- chosen = scores.argsort(dim=1)[:, :block_topk]
- column = torch.arange(final_topk, device=DEV)[None, :]
- group_of_column = (column // compress_ratio).clamp(max=block_topk - 1).expand(rows, -1)
- body = torch.gather(chosen, 1, group_of_column) * compress_ratio + column % compress_ratio
- in_body = column < (taken * compress_ratio)[:, None]
- tail_column = column - (taken * compress_ratio)[:, None]
- tail = (available * compress_ratio)[:, None] + tail_column
- in_tail = ~in_body & (tail_column >= 0) & (tail < visible[:, None])
- selected = torch.where(in_body, body, torch.where(in_tail, tail, torch.full_like(body, -1)))
- return selected.to(torch.int32)
-
-
-def _random_pool(pages, kv, tpb, d, gen):
- """Pool bytes: any e2m1 nibble, and positive e4m3 scales from 2^-6 to 0.5."""
- data = torch.randint(0, 256, (pages, 2, kv, tpb, d // 2), generator=gen, device=DEV)
- codes = torch.randint(0x08, 0x31, (pages, 2, kv, tpb, d // 16), generator=gen, device=DEV)
- return data.to(torch.uint8), codes.to(torch.uint8)
-
-
-def _build(spec, seed):
- """Real paged state for one spec, keyed by the op's parameter names."""
- gen = torch.Generator(device=DEV).manual_seed(seed)
- tpb, kv, d, q_heads = (spec[k] for k in ("tokens_per_block", "kv_heads", "head_dim", "q_heads"))
- segments = list(spec["segments"])
- positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments])
- req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)])
- rows, slots = positions.numel(), spec["slots"]
- assert len(segments) <= slots
- slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(segments)]
- mapped = [math.ceil((s + n) / tpb) for s, n in segments]
- assert not mapped or max(mapped) <= spec["pages"]
- needed = sum(mapped) + 7
- slot_pages, layer = spec["slot_pages"], spec["layer"]
- planes = slice(2 * layer, 2 * layer + 2)
- # The requests sit on the pool's last `needed` pages, and every scale outside them is
- # NaN, so a page address that strayed could not come back finite.
- pool_pages = spec.get("pool_pages", needed)
- assert pool_pages >= needed
- first = pool_pages - needed
- page_ids = first + torch.randperm(needed, generator=gen, device=DEV).to(torch.int32)
- data = torch.zeros(pool_pages, slot_pages, kv, tpb, d // 2, dtype=torch.uint8, device=DEV)
- scales = torch.full(
- (pool_pages, slot_pages, kv, tpb, d // 16), E4M3_NAN, dtype=torch.uint8, device=DEV
- )
- data[first:, planes], scales[first:, planes] = _random_pool(needed, kv, tpb, d, gen)
- block_table = torch.full((slots, spec["pages"]), -1, dtype=torch.int32, device=DEV)
- cursor = 0
- for slot, count in zip(slot_of.tolist(), mapped):
- block_table[slot, :count] = page_ids[cursor : cursor + count]
- cursor += count
-
- assert spec["q_layout"] == "qkv_slice", spec["q_layout"]
- qkv = torch.randn(rows, q_heads + 2 * kv, d, generator=gen, device=DEV).to(spec["q_dtype"])
- q = qkv[:, :q_heads]
- selected = _selected_tokens(positions, seed + 3, spec["compress_ratio"], spec["token_topk"])
- request_indices = slot_of.to(DEV)[req].to(spec["request_dtype"])
-
- args = dict(
- q=q,
- k_cache=data[:, 2 * layer],
- v_cache=data[:, 2 * layer + 1],
- k_scale_cache=scales[:, 2 * layer],
- v_scale_cache=scales[:, 2 * layer + 1],
- kv_scale_quant_orig=torch.tensor(
- [1.0, spec["global_k"], spec["global_v"]], dtype=torch.float32, device=DEV
- ),
- block_table=block_table,
- selected_tokens=selected,
- request_indices=request_indices,
- tokens_per_block=tpb,
- softmax_scale=d**-0.5,
- )
- if spec["prefill_pair"]:
- args.update(query_positions=positions, compress_ratio=spec["compress_ratio"])
- return args
-
-
-def _bf16_twin(args):
- """QSA's bf16 sparse GQA over a bf16 pool holding the dequantized values."""
- pages = args["k_cache"].shape[0]
- all_pages = torch.arange(pages, device=DEV)[:, None].expand(-1, args["tokens_per_block"])
- slots = torch.arange(args["tokens_per_block"], device=DEV)[None, :].expand(pages, -1)
- planes = []
- for data, scales, index, grouped in (
- (args["k_cache"], args["k_scale_cache"], 1, False),
- (args["v_cache"], args["v_scale_cache"], 2, True),
- ):
- values = _entry._dequantize_tokens(
- data, scales, all_pages, slots, args["kv_scale_quant_orig"][index], grouped
- ) # [pages, tpb, H, D]
- planes.append(values.to(args["q"].dtype).transpose(1, 2).contiguous())
- twin = {
- name: args[name]
- for name in ("q", "block_table", "selected_tokens", "request_indices", "tokens_per_block")
- }
- twin.update(softmax_scale=args["softmax_scale"], k_cache=planes[0], v_cache=planes[1])
- if "query_positions" in args:
- twin.update(query_positions=args["query_positions"], compress_ratio=args["compress_ratio"])
- return triton_qsa_paged_sparse_gqa(**twin)
-
-
-def _bound(positions, spec):
- cr = spec["compress_ratio"]
- visible = positions.long() + 1
- return (visible // cr).clamp(max=spec["token_topk"] // cr) * cr + visible % cr
-
-
-def _is_prefix_packed(table):
- valid = table >= 0
- columns = torch.arange(table.shape[1], device=table.device)[None, :]
- return (valid == (columns < valid.sum(1)[:, None])).all(1)
-
-
-_OPERANDS = (
- "q",
- "k_cache",
- "v_cache",
- "k_scale_cache",
- "v_scale_cache",
- "block_table",
- "selected_tokens",
- "request_indices",
-)
-
-
-def _capture_sweep(spec, seed, launches):
- """The decode site under a CUDA graph, at each G for each rows-per-request:
- capture, write a second state's operands in place, replay, and require the
- eager call on that second state bitwise."""
- sweep = spec["capture"]
- for per_request in sweep["rows_per_request"]:
- for groups in sweep["groups"]:
- _capture_once(
- _decode_step(spec, groups, per_request),
- seed + 100 * per_request + groups,
- launches,
- where=f"G={groups} x {per_request} rows",
- )
-
-
-def _start(request: int, rows: int) -> int:
- """A request's first position, across the budget."""
- return min(100 + (499 * request) % 8000, ENGINE_MAX_SEQ_LEN - rows)
-
-
-def _decode_step(spec, groups, per_request):
- """The decode site's step: `groups` requests x `per_request` rows."""
- segments = tuple((_start(g, per_request), per_request) for g in range(groups))
- return {**spec, "segments": segments}
-
-
-def _capture_once(shape, seed, launches, where) -> None:
- """One call of `shape` captured, a second state's operands written in place,
- replayed: the replay must be bitwise the eager call on the second state."""
- where = f"{where}, [{shape['slots']}, {shape['pages']}] table"
- args = _build(shape, seed)
- fresh = _build(shape, seed + 50)
- launches.clear()
- op(**args)
- assert [name for name, _ in launches] == list(shape["launches"]), where
- torch.cuda.synchronize()
- graph = torch.cuda.CUDAGraph()
- try:
- with torch.cuda.graph(graph):
- captured = op(**args)
- for name in _OPERANDS:
- args[name].copy_(fresh[name])
- graph.replay()
- torch.cuda.synchronize()
- eager = op(**fresh)
- assert torch.equal(captured, eager), f"replay differs from eager at {where}"
- _compare(captured, op.reference(**fresh))
- finally:
- del graph
- torch.cuda.synchronize()
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811
- certifying(op)
- spec = cell.spec
- if "capture" in spec:
- _capture_sweep(spec, seed=sum(map(ord, cell.why)), launches=launches)
- return
- args = _build(spec, seed=sum(map(ord, cell.why)))
- rows = args["q"].shape[0]
- if spec["prefill_pair"]:
- selected = args["selected_tokens"]
- assert bool(_is_prefix_packed(selected).all())
- valid = (selected >= 0).sum(1)
- assert torch.equal(valid, _bound(args["query_positions"], spec))
-
- out = op(**args)
- assert [name for name, _ in launches] == list(spec["launches"])
- assert launches[0][1][:2] == (rows, spec["kv_heads"])
- assert out.shape == (rows, spec["q_heads"], spec["head_dim"])
- assert out.dtype == args["q"].dtype and out.is_contiguous()
- assert not bool(out.isnan().any())
- _compare(out, op.reference(**args))
- # QSA's bf16 kernels over the same values: inside the same band, not bitwise
- # (the dequantizing load changes the compiled kernel).
- _compare(out, _bf16_twin(args))
-
-
-CKPT = op.CELLS[0].spec
-
-
-def _with(**overrides):
- return {**CKPT, **overrides}
-
-
-def test_calls_are_bitwise_deterministic() -> None:
- args = _build(_with(segments=tuple((4000 + 37 * i, 2) for i in range(4))), seed=8800)
- assert torch.equal(op(**args), op(**args))
-
-
-def test_rejects_positions_without_a_compress_ratio() -> None:
- """The op's own loud refusal -- the reason the entry carries no `is_valid`: half of
- the prefill pair."""
- args = _build(_with(segments=((600, 1),) * 4), seed=9100)
- args["selected_tokens"] = args["selected_tokens"][:, :64].contiguous()
- args["query_positions"] = torch.zeros(4, dtype=torch.int64, device=DEV)
- message = "pass both query_positions and compress_ratio, or neither"
- with pytest.raises(ValueError, match=re.escape(message)):
- with validating(op):
- op(**args)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_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_nvfp4_decode.py
deleted file mode 100644
index d8ddeb55e29d..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_decode.py
+++ /dev/null
@@ -1,209 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test: `qsa_paged_sparse_gqa_nvfp4_decode` against the float64 reference of
-`qsa_paged_sparse_gqa_nvfp4` and QSA's bf16 kernel over the dequantized pool, inside the
-NVFP4 op's band (4 ulp of each (row, head)'s max, 2 ulp relative RMS), over its cells with and
-without the output gate; plus determinism and graph replay against eager."""
-
-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.attention import ( # noqa: E501
- qsa_paged_sparse_gqa_nvfp4 as _nvfp4_module,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_sparse_gqa_nvfp4_decode import ( # noqa: E501
- qsa_paged_sparse_gqa_nvfp4_decode as op,
-)
-from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import triton_qsa_paged_sparse_gqa
-
-assert torch.cuda.is_available(), "qsa_paged_sparse_gqa_nvfp4_decode requires a CUDA device"
-
-DEV = "cuda"
-E4M3_NAN = 0x7F
-ULP = 2.0**-7 # bf16
-
-
-def within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0):
- o = out.float().reshape(-1, out.shape[-1])
- r = ref.float().reshape(-1, ref.shape[-1])
- scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9)
- el = ((o - r).abs() / scale).max().item() / ULP
- rms = (((o - r) ** 2).sum() / (r**2).sum().clamp_min(1e-30)).sqrt().item() / ULP
- return el, rms, el <= element_ulp and rms <= rms_ulp
-
-
-def _selected_tokens(positions, seed, compress_ratio, token_topk):
- block_topk = token_topk // compress_ratio
- final_topk = token_topk + compress_ratio - 1
- gen = torch.Generator(device=DEV).manual_seed(seed)
- rows = positions.numel()
- visible = positions.long() + 1
- available = visible // compress_ratio
- taken = available.clamp(max=block_topk)
- width = max(block_topk, int(available.max().item()) if rows else 0)
- scores = torch.rand(rows, width, generator=gen, device=DEV)
- group_ids = torch.arange(width, device=DEV)[None, :]
- scores = torch.where(group_ids < available[:, None], scores, torch.inf)
- chosen = scores.argsort(dim=1)[:, :block_topk]
- column = torch.arange(final_topk, device=DEV)[None, :]
- group_of_column = (column // compress_ratio).clamp(max=block_topk - 1).expand(rows, -1)
- body = torch.gather(chosen, 1, group_of_column) * compress_ratio + column % compress_ratio
- in_body = column < (taken * compress_ratio)[:, None]
- tail_column = column - (taken * compress_ratio)[:, None]
- tail = (available * compress_ratio)[:, None] + tail_column
- in_tail = ~in_body & (tail_column >= 0) & (tail < visible[:, None])
- selected = torch.where(in_body, body, torch.where(in_tail, tail, torch.full_like(body, -1)))
- return selected.to(torch.int32)
-
-
-def build(
- segments,
- seed,
- slots=32,
- pages=260,
- slot_pages=2,
- layer=0,
- global_k=1.0,
- global_v=1.0,
- request_dtype=torch.int32,
- pool_pages=None,
-):
- gen = torch.Generator(device=DEV).manual_seed(seed)
- tpb, kv, d, q_heads = 32, 2, 256, 24
- positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments])
- req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)])
- rows = positions.numel()
- slots = max(slots, len(segments))
- slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(segments)]
- mapped = [math.ceil((s + n) / tpb) for s, n in segments]
- needed = sum(mapped) + 7
- planes = slice(2 * layer, 2 * layer + 2)
- pool_pages = pool_pages or needed
- first = pool_pages - needed
- page_ids = first + torch.randperm(needed, generator=gen, device=DEV).to(torch.int32)
- data = torch.zeros(pool_pages, slot_pages, kv, tpb, d // 2, dtype=torch.uint8, device=DEV)
- scales = torch.full(
- (pool_pages, slot_pages, kv, tpb, d // 16), E4M3_NAN, dtype=torch.uint8, device=DEV
- )
- data[first:, planes] = torch.randint(
- 0, 256, (needed, 2, kv, tpb, d // 2), generator=gen, device=DEV
- ).to(torch.uint8)
- scales[first:, planes] = torch.randint(
- 0x08, 0x31, (needed, 2, kv, tpb, d // 16), generator=gen, device=DEV
- ).to(torch.uint8)
- block_table = torch.full((slots, pages), -1, dtype=torch.int32, device=DEV)
- cursor = 0
- for slot, count in zip(slot_of.tolist(), mapped):
- block_table[slot, :count] = page_ids[cursor : cursor + count]
- cursor += count
- qkv = torch.randn(rows, q_heads + 2 * kv, d, generator=gen, device=DEV).to(torch.bfloat16)
- selected = _selected_tokens(positions, seed + 3, 4, 2048)
- return dict(
- q=qkv[:, :q_heads],
- 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, global_k, global_v], dtype=torch.float32, device=DEV
- ),
- block_table=block_table,
- selected_tokens=selected,
- request_indices=slot_of.to(DEV)[req].to(request_dtype),
- tokens_per_block=tpb,
- softmax_scale=d**-0.5,
- )
-
-
-def bf16_twin(args):
- pages = args["k_cache"].shape[0]
- all_pages = torch.arange(pages, device=DEV)[:, None].expand(-1, args["tokens_per_block"])
- slots = torch.arange(args["tokens_per_block"], device=DEV)[None, :].expand(pages, -1)
- planes = []
- for data, scales, index, grouped in (
- (args["k_cache"], args["k_scale_cache"], 1, False),
- (args["v_cache"], args["v_scale_cache"], 2, True),
- ):
- values = _nvfp4_module._dequantize_tokens(
- data, scales, all_pages, slots, args["kv_scale_quant_orig"][index], grouped
- )
- planes.append(values.to(args["q"].dtype).transpose(1, 2).contiguous())
- twin = {
- n: args[n]
- for n in ("q", "block_table", "selected_tokens", "request_indices", "tokens_per_block")
- }
- twin.update(softmax_scale=args["softmax_scale"], k_cache=planes[0], v_cache=planes[1])
- return triton_qsa_paged_sparse_gqa(**twin)
-
-
-def _args(spec, seed):
- kw = {k: v for k, v in spec.items() if k not in ("segments", "gated")}
- args = build(spec["segments"], seed=seed, **kw)
- if spec.get("gated"):
- rows = args["q"].shape[0]
- g = torch.Generator(device=DEV).manual_seed(seed + 1)
- args["output_gate"] = torch.randn(rows, 24 * 256, generator=g, device=DEV).to(
- torch.bfloat16
- )
- return args
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell) -> None:
- args = _args(cell.spec, zlib.crc32(cell.why.encode()) % (2**31))
- out = op(**args)
- assert out.dtype == torch.bfloat16 and not out.isnan().any()
- el, rms, ok = within_ulp(out, op.reference(**args))
- assert ok, f"vs reference: {el:.2f} ulp element, {rms:.2f} ulp RMS"
- if "output_gate" not in args:
- el, rms, ok = within_ulp(out, bf16_twin(args))
- assert ok, f"vs QSA bf16 kernel: {el:.2f} ulp element, {rms:.2f} ulp RMS"
-
-
-def test_unaligned_row_slice_of_the_selection() -> None:
- """A mixed step hands its generation rows as a row slice of the step's 2051-wide
- selection, 16-byte aligned only at every fourth row: an unaligned slice gives the same
- rows as an aligned copy."""
- args = build(tuple((3000 + 211 * i, 1) for i in range(9)), seed=4242)
- aligned = dict(args)
- for name in ("q", "selected_tokens", "request_indices"):
- args[name] = args[name][1:]
- aligned[name] = args[name].clone()
- assert args["selected_tokens"].data_ptr() % 16
- assert torch.equal(op(**args), op(**aligned))
-
-
-def test_calls_are_bitwise_deterministic() -> None:
- args = build(tuple((4000 + 37 * i, 2) for i in range(4)), seed=8800)
- assert torch.equal(op(**args), op(**args))
-
-
-@pytest.mark.parametrize("rows", [1, 64])
-def test_replay_matches_eager(rows) -> None:
- segments = (
- ((4000, 1),) if rows == 1 else tuple((100 + (499 * g) % 8000, 1) for g in range(rows))
- )
- args = build(segments, seed=77)
- fresh = build(segments, seed=127)
- op(**args)
- torch.cuda.synchronize()
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- captured = op(**args)
- for name in (
- "q",
- "k_cache",
- "v_cache",
- "k_scale_cache",
- "v_scale_cache",
- "block_table",
- "selected_tokens",
- "request_indices",
- ):
- args[name].copy_(fresh[name])
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(captured, op(**fresh))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py
deleted file mode 100644
index 664da3193ba6..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py
+++ /dev/null
@@ -1,316 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test: `qsa_paged_sparse_gqa_nvfp4_verify` against the reference of
-`qsa_paged_sparse_gqa_nvfp4_decode` and against the per-row kernel, inside the NVFP4 op's band
-(4 ulp of each (row, head)'s max, 2 ulp relative RMS), over its cells; plus the routing to the
-per-row kernel, forced splits and other rows a request, the launch plan (a partial last wave's
-units split and merged), invalid request rows, an unaligned row slice of the selection, the
-selections `is_valid` refuses, determinism and graph replay against eager."""
-
-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.attention import ( # noqa: E501
- qsa_paged_sparse_gqa_nvfp4_verify as _verify_module,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_sparse_gqa_nvfp4_decode import ( # noqa: E501
- nvfp4_sparse_gqa_decode as per_row,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_sparse_gqa_nvfp4_verify import ( # noqa: E501
- qsa_paged_sparse_gqa_nvfp4_verify as op,
-)
-
-assert torch.cuda.is_available(), "qsa_paged_sparse_gqa_nvfp4_verify requires a CUDA device"
-
-DEV = "cuda"
-E4M3_NAN = 0x7F
-ULP = 2.0**-7 # bf16
-COMPRESS_RATIO, TOKEN_TOPK = 4, 2048
-WIDTH = TOKEN_TOPK + COMPRESS_RATIO - 1
-
-
-def within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0):
- o = out.float().reshape(-1, out.shape[-1])
- r = ref.float().reshape(-1, ref.shape[-1])
- scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9)
- el = ((o - r).abs() / scale).max().item() / ULP
- rms = (((o - r) ** 2).sum() / (r**2).sum().clamp_min(1e-30)).sqrt().item() / ULP
- return el, rms, el <= element_ulp and rms <= rms_ulp
-
-
-def _expand(blocks, taken, available, visible):
- """Block ids `[rows, 512]` (the first `taken` valid) to the 2051-wide token selection:
- each block's four tokens, then the causal tail, then -1."""
- column = torch.arange(WIDTH, device=DEV)[None, :]
- rows = blocks.shape[0]
- block_of_column = (column // COMPRESS_RATIO).clamp(max=blocks.shape[1] - 1).expand(rows, -1)
- body = torch.gather(blocks, 1, block_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 _selected_tokens(positions, requests, seed, overlap):
- """Each row's top 512 visible four-token blocks, expanded. `overlap` None: independent
- random scores a row, in random order (the decode op's test); else the request's scores plus
- `overlap` x N(0, 1) a row, blocks ascending (as `qsa_topk_expand_decode` writes them)."""
- block_topk = TOKEN_TOPK // COMPRESS_RATIO
- 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)
- block_ids = torch.arange(width, device=DEV)[None, :]
- if overlap is None:
- scores = torch.rand(rows, width, generator=gen, device=DEV)
- scores = torch.where(block_ids < available[:, None], scores, torch.inf)
- return _expand(scores.argsort(dim=1)[:, :block_topk], taken, available, visible)
- base = torch.rand(int(requests.max().item()) + 1, width, generator=gen, device=DEV)
- noise = torch.randn(rows, width, generator=gen, device=DEV)
- scores = base[requests] + overlap * noise
- scores = torch.where(block_ids < available[:, None], scores, -torch.inf)
- top = scores.topk(block_topk, dim=1).indices
- valid = torch.arange(block_topk, device=DEV)[None, :] < taken[:, None]
- blocks = torch.where(valid, top, torch.full_like(top, 1 << 30)).sort(dim=1).values
- return _expand(blocks, taken, available, visible)
-
-
-def build(
- segments,
- seed,
- overlap=0.05,
- slots=32,
- pages=260,
- slot_pages=2,
- layer=0,
- global_k=1.0,
- global_v=1.0,
- request_dtype=torch.int32,
- pool_pages=None,
-):
- """The decode op's test state for `segments` ((first position, rows) per request, in row
- order) with `_selected_tokens(overlap)`."""
- gen = torch.Generator(device=DEV).manual_seed(seed)
- tpb, kv, d, q_heads = 32, 2, 256, 24
- positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments])
- req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)])
- rows = positions.numel()
- slots = max(slots, len(segments))
- slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(segments)]
- mapped = [math.ceil((s + n) / tpb) for s, n in segments]
- needed = sum(mapped) + 7
- planes = slice(2 * layer, 2 * layer + 2)
- pool_pages = pool_pages or needed
- first = pool_pages - needed
- page_ids = first + torch.randperm(needed, generator=gen, device=DEV).to(torch.int32)
- data = torch.zeros(pool_pages, slot_pages, kv, tpb, d // 2, dtype=torch.uint8, device=DEV)
- scales = torch.full(
- (pool_pages, slot_pages, kv, tpb, d // 16), E4M3_NAN, dtype=torch.uint8, device=DEV
- )
- data[first:, planes] = torch.randint(
- 0, 256, (needed, 2, kv, tpb, d // 2), generator=gen, device=DEV
- ).to(torch.uint8)
- scales[first:, planes] = torch.randint(
- 0x08, 0x31, (needed, 2, kv, tpb, d // 16), generator=gen, device=DEV
- ).to(torch.uint8)
- block_table = torch.full((slots, pages), -1, dtype=torch.int32, device=DEV)
- cursor = 0
- for slot, count in zip(slot_of.tolist(), mapped):
- block_table[slot, :count] = page_ids[cursor : cursor + count]
- cursor += count
- qkv = torch.randn(rows, q_heads + 2 * kv, d, generator=gen, device=DEV).to(torch.bfloat16)
- selected = _selected_tokens(positions, req.long(), seed + 3, overlap)
- return dict(
- q=qkv[:, :q_heads],
- 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, global_k, global_v], dtype=torch.float32, device=DEV
- ),
- block_table=block_table,
- selected_tokens=selected,
- request_indices=slot_of.to(DEV)[req].to(request_dtype),
- tokens_per_block=tpb,
- softmax_scale=d**-0.5,
- )
-
-
-def _args(spec, seed):
- """A cell's arguments, `rows_per_request` and the gate included."""
- kw = {k: v for k, v in spec.items() if k not in ("segments", "gated", "rows_per_request")}
- args = build(spec["segments"], seed=seed, **kw)
- args["rows_per_request"] = spec["rows_per_request"]
- if spec.get("gated"):
- rows = args["q"].shape[0]
- g = torch.Generator(device=DEV).manual_seed(seed + 1)
- args["output_gate"] = torch.randn(rows, 24 * 256, generator=g, device=DEV).to(
- torch.bfloat16
- )
- return args
-
-
-def _takes_grouped(args):
- return _verify_module.takes_grouped(
- args["q"],
- args["k_cache"],
- args["block_table"],
- args["selected_tokens"],
- args["tokens_per_block"],
- args["rows_per_request"],
- )
-
-
-def _per_row(args):
- return per_row(**{k: v for k, v in args.items() if k != "rows_per_request"})
-
-
-def _segments(requests, rows_each, first=4000, step=37):
- return tuple((first + step * i, rows_each) for i in range(requests))
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell) -> None:
- args = _args(cell.spec, zlib.crc32(cell.why.encode()) % (2**31))
- op.is_valid(**args)
- assert _takes_grouped(args), "every cell runs on the grouped kernel"
- out = op(**args)
- assert out.dtype == torch.bfloat16 and not out.isnan().any()
- el, rms, ok = within_ulp(out, op.reference(**args))
- assert ok, f"vs reference: {el:.2f} ulp element, {rms:.2f} ulp RMS"
- el, rms, ok = within_ulp(out, _per_row(args))
- assert ok, f"vs the per-row kernel: {el:.2f} ulp element, {rms:.2f} ulp RMS"
-
-
-@pytest.mark.parametrize(
- "segments, rows_per_request",
- [(_segments(1, 4), 4), (_segments(7, 4), 4), (_segments(128, 1, step=29), 1)],
- ids=["4 rows", "7 groups of 4", "128 rows of 1"],
-)
-def test_small_batches_and_single_rows_take_the_per_row_kernel(segments, rows_per_request):
- args = build(segments, seed=31)
- args["rows_per_request"] = rows_per_request
- assert not _takes_grouped(args)
- assert torch.equal(op(**args), _per_row(args))
-
-
-@pytest.mark.parametrize("rows_per_request, splits", [(2, None), (4, 2)], ids=str)
-def test_other_rows_a_request_and_forced_splits(rows_per_request, splits) -> None:
- args = build(_segments(8, rows_per_request), seed=500 + rows_per_request)
- out = _verify_module.nvfp4_sparse_gqa_verify_grouped(
- **args, rows_per_request=rows_per_request, splits=splits
- )
- el, rms, ok = within_ulp(out, op.reference(**args, rows_per_request=rows_per_request))
- assert ok, f"vs reference: {el:.2f} ulp element, {rms:.2f} ulp RMS"
-
-
-def test_launch_plan_splits_only_a_partial_last_wave() -> None:
- plan = _verify_module.plan_launch
- # within one wave every unit (group, KV head) takes the same splits
- assert plan(32, 2, WIDTH, 152) == (2, None)
- # 512 units on 152 SMs: three full waves of one CTA a unit, the last 56 units split in two
- assert plan(256, 2, WIDTH, 152) == (2, 456)
- # a last wave too full to split (104 units): one CTA a unit
- assert plan(128, 2, WIDTH, 152) == (1, None)
-
-
-@pytest.mark.parametrize("gated", [False, True], ids=["plain", "gated"])
-def test_split_last_wave_matches_one_cta_a_unit(gated) -> None:
- """The units of the full waves are bitwise the one-CTA-a-unit launch; the split and
- merged last wave stays inside the band of it and of the reference."""
- args = build(tuple((2100 + (37 * i) % 1900, 4) for i in range(256)), seed=2718)
- rows = args["q"].shape[0]
- if gated:
- g = torch.Generator(device=DEV).manual_seed(2719)
- args["output_gate"] = torch.randn(rows, 24 * 256, generator=g, device=DEV).to(
- torch.bfloat16
- )
- sms = torch.cuda.get_device_properties(0).multi_processor_count
- splits, n_whole = _verify_module.plan_launch(rows // 4, 2, WIDTH, sms)
- if n_whole is None:
- pytest.skip(f"{sms} SMs: no partial last wave at {rows} rows")
- out = op(**args, rows_per_request=4)
- one = _verify_module.nvfp4_sparse_gqa_verify_grouped(**args, rows_per_request=4, splits=1)
- head = (n_whole // 2) * 4
- assert torch.equal(out[:head], one[:head])
- el, rms, ok = within_ulp(out, one)
- assert ok, f"vs one CTA a unit: {el:.2f} ulp element, {rms:.2f} ulp RMS"
- el, rms, ok = within_ulp(out, op.reference(**args, rows_per_request=4))
- assert ok, f"vs reference: {el:.2f} ulp element, {rms:.2f} ulp RMS"
-
-
-def test_rows_of_an_invalid_request_are_zero() -> None:
- args = build(_segments(8, 4), seed=99)
- args["request_indices"][5] = -1
- out = op(**args, rows_per_request=4)
- assert not out[5].any()
- el, rms, ok = within_ulp(out, op.reference(**args, rows_per_request=4))
- assert ok, f"vs reference: {el:.2f} ulp element, {rms:.2f} ulp RMS"
-
-
-def test_unaligned_row_slice_of_the_selection() -> None:
- """A mixed step hands its generation rows as a row slice of the step's 2051-wide
- selection, 16-byte aligned only at every fourth row: an unaligned slice gives the same
- rows as an aligned copy."""
- args = build(((3000, 1),) + _segments(8, 4), seed=4242)
- aligned = dict(args)
- for name in ("q", "selected_tokens", "request_indices"):
- args[name] = args[name][1:]
- aligned[name] = args[name].clone()
- assert args["selected_tokens"].data_ptr() % 16
- assert torch.equal(op(**args, rows_per_request=4), op(**aligned, rows_per_request=4))
-
-
-def test_is_valid_refuses_selections_that_are_not_expanded() -> None:
- args = build(_segments(8, 4), seed=7)
- op.is_valid(**args, rows_per_request=4)
- good = args["selected_tokens"]
- shifted = good.clone()
- shifted[0, 0] += 1 # off its column's offset within the group
- mixed = good.clone()
- mixed[0, 1] = good[0, 5] # two groups in one quad
- repeated = good.clone()
- repeated[0, 4:8] = good[0, 0:4] # one group twice in a row
- for bad in (shifted, mixed, repeated):
- args["selected_tokens"] = bad
- with pytest.raises(AssertionError):
- op.is_valid(**args, rows_per_request=4)
-
-
-def test_calls_are_bitwise_deterministic() -> None:
- args = build(_segments(8, 4), seed=8800)
- assert torch.equal(op(**args, rows_per_request=4), op(**args, rows_per_request=4))
-
-
-@pytest.mark.parametrize("rows", [32])
-def test_replay_matches_eager(rows) -> None:
- segments = tuple((100 + (499 * g) % 8000, 4) for g in range(rows // 4))
- args = build(segments, seed=77)
- fresh = build(segments, seed=127)
- op(**args, rows_per_request=4)
- torch.cuda.synchronize()
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- captured = op(**args, rows_per_request=4)
- for name in (
- "q",
- "k_cache",
- "v_cache",
- "k_scale_cache",
- "v_scale_cache",
- "block_table",
- "selected_tokens",
- "request_indices",
- ):
- args[name].copy_(fresh[name])
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(captured, op(**fresh, rows_per_request=4))
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
deleted file mode 100644
index cac393b9865c..000000000000
--- 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
+++ /dev/null
@@ -1,237 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the qsa_prefill_compress catalog entry.
-
-The paged state is real: an index-key pool of random raw keys, a position
-plane holding `(t, h, w)` for every allocated position, and the engine's
-`[slots, width]` table -- more slots than requests, each live row's pages drawn
-from a seeded permutation of the pool, `-1` past its allocation, rows past the
-batch holding a previous step's plausible page ids. The cos|sin table is built
-the way the target builds it at load.
-
-A padding row at logical position -1 is deliberately absent: it is an illegal
-memory access that poisons the process for every test after it.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_prefill_compress import ( # noqa: E501
- qsa_prefill_compress as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401
-
-assert torch.cuda.is_available(), "qsa_prefill_compress requires a CUDA device"
-
-
-def _compare(out: torch.Tensor, ref: torch.Tensor, 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)
- 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:
- """The target's `[max_positions, 2 * pairs]` float32 cos|sin table."""
- exponent = torch.arange(0, ROTARY_DIM, 2, dtype=torch.float64, device=DEV)
- inv_freq = (1.0 / (THETA ** (exponent / ROTARY_DIM))).to(torch.float32)
- positions = torch.arange(MAX_POSITIONS, dtype=torch.float32, device=DEV)
- angles = positions[:, None] * inv_freq[None, :]
- return torch.cat([torch.cos(angles), torch.sin(angles)], dim=1).contiguous()
-
-
-def _build(spec, seed: int):
- """Real caches, table and rows for one cell; returns the op's kwargs and extras."""
- gen = torch.Generator(device="cpu").manual_seed(seed)
- tpb, head_dim = spec["tokens_per_block"], spec["head_dim"]
- slots, width, pages = spec["slots"], spec["width"], spec["pages"]
- segments = spec["segments"]
- requests = [r for r, (_, n) in enumerate(segments) for _ in range(n)]
- positions = [p for start, n in segments for p in range(start, start + n)]
-
- table = torch.randint(0, pages, (slots, width), generator=gen, dtype=torch.int32)
- order = torch.randperm(pages, generator=gen).to(torch.int32)
- index_cache = torch.randn(pages, tpb, 1, head_dim, generator=gen).to(torch.bfloat16)
- position_cache = torch.zeros(pages, tpb, 3, dtype=torch.int32)
- used = 0
- for r, (start, n) in enumerate(segments):
- cols = (start + n - 1) // tpb + 1
- table[r] = -1
- table[r, :cols] = order[used : used + cols]
- used += cols
- p = torch.arange(cols * tpb).view(cols, tpb)
- 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"])
- table, index_cache, position_cache = table.to(DEV), index_cache.to(DEV), position_cache.to(DEV)
-
- extras = {}
- if spec["layout"] == "shared_page":
- # An NVFP4 pool's page: `roles` index pages, then the position page. The other
- # index roles hold a sentinel and must survive, as must the position page.
- roles = spec["roles"]
- index_bytes = tpb * head_dim * 2
- page = torch.empty(pages, roles * index_bytes + tpb * 3 * 4, dtype=torch.uint8, device=DEV)
- others = page[:, : (roles - 1) * index_bytes].view(torch.int16)
- others.fill_(SENTINEL)
- shared_index = page[:, (roles - 1) * index_bytes : roles * index_bytes]
- shared_index = shared_index.view(torch.bfloat16).view(pages, tpb, 1, head_dim)
- shared_position = page[:, roles * index_bytes :].view(torch.int32).view(pages, tpb, 3)
- shared_index.copy_(index_cache)
- shared_position.copy_(position_cache)
- assert shared_index.stride(0) == roles * tpb * head_dim + tpb * 3 * 2
- assert shared_position.stride() == (roles * tpb * head_dim // 2 + tpb * 3, 3, 1)
- index_cache, position_cache = shared_index, shared_position
- extras["neighbour"] = (others, others.clone())
- extras["position_page"] = (position_cache, position_cache.clone())
- else:
- assert spec["layout"] == "plain", spec["layout"]
-
- args = dict(
- logical_positions=logical.to(DEV),
- request_indices=request.to(DEV),
- block_table=table,
- index_cache=index_cache,
- position_cache=position_cache,
- k_norm_weight=k_norm_weight.to(DEV),
- cos_sin=_cos_sin(),
- eps=spec["eps"],
- tokens_per_block=tpb,
- compress_ratio=spec["compress_ratio"],
- mrope_section=spec["mrope_section"],
- )
- return args, extras
-
-
-def _anchors(args) -> int:
- """Distinct anchor slots the rows name."""
- logical = args["logical_positions"].long()
- keyed = args["request_indices"].long() * (1 << 32) + logical
- return int(torch.unique(keyed[(logical + 1) % args["compress_ratio"] == 0]).numel())
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying): # noqa: F811
- certifying(op)
- args, extras = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
- before = args["index_cache"].clone()
- ref = op.reference(**args)
- out = op(**args)
- assert out is args["index_cache"]
- _compare(out, ref, 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 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, 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]
-
-
-#: The cell of 32 MTP7 verify blocks (256 rows, two anchors each): the segments, pool and table.
-_MTP7_32 = next(c.spec for c in op.CELLS if c.why.startswith("decode MTP7: 32 requests"))
-_MTP7_BATCH = {key: _MTP7_32[key] for key in ("segments", "pages", "slots", "rounding_flips")}
-
-
-def test_cuda_graph_capture_and_replay() -> None:
- """The launch grid is fixed by the row count, so it captures and replays: 32 MTP7
- verify blocks, the decode target's captured batch."""
- batch = _MTP7_BATCH
- args = _scene(seed=16, **batch)
- cache = args["index_cache"]
- before = cache.clone()
- ref = op.reference(**args)
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(warmup):
- for _ in range(3):
- cache.copy_(before)
- op(**args)
- torch.cuda.current_stream().wait_stream(warmup)
- torch.cuda.synchronize()
- graph = torch.cuda.CUDAGraph()
- cache.copy_(before)
- with torch.cuda.graph(graph):
- op(**args)
- for _ in range(2):
- cache.copy_(before)
- graph.replay()
- torch.cuda.synchronize()
- _compare(cache, ref, batch["rounding_flips"])
-
-
-def _strided(tensor: torch.Tensor, dim: int) -> torch.Tensor:
- """The same values as a stride-2 view along `dim` of a zero-filled buffer."""
- shape = list(tensor.shape)
- shape[dim] *= 2
- wide = torch.zeros(shape, dtype=tensor.dtype, device=tensor.device)
- index = [slice(None)] * tensor.dim()
- index[dim] = slice(None, None, 2)
- wide[tuple(index)] = tensor
- return wide[tuple(index)]
-
-
-def test_guard_refuses_strided_logical_positions() -> None:
- """Shown unguarded first -- accepted and wrong -- then refused: a 1-D operand the
- kernel reads flat (`request_indices` and `k_norm_weight` carry the same
- `is_contiguous` guard). Old receipt's max abs error against the right answer: 4.72;
- the floor below states a bound, not one draw.
- """
- args = _scene(seed=12)
- before = args["index_cache"].clone()
- ref = op.reference(**args)
- bad = dict(args, logical_positions=_strided(args["logical_positions"], 0))
- assert not bad["logical_positions"].is_contiguous()
-
- out = op(**bad)
- error = (out.float() - ref.float()).abs().max().item()
- assert error > 0.25, (
- f"strided logical_positions was expected to be read wrongly; off by {error}"
- )
-
- args["index_cache"].copy_(before)
- with pytest.raises(AssertionError, match="logical_positions must be contiguous"):
- with validating(op):
- op(**bad)
- assert torch.equal(args["index_cache"], before), "a refused call still wrote"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_rollback_planes.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_rollback_planes.py
deleted file mode 100644
index 59c29b3fccec..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_rollback_planes.py
+++ /dev/null
@@ -1,84 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test: `qsa_rollback_gather` / `qsa_rollback_restore` against torch's advanced indexing,
-bitwise, over planes that are strided views of one per-page buffer (as an NVFP4 KV cache's
-manager packs the side planes) and over contiguous planes."""
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_rollback_planes import ( # noqa: E501
- ROW,
- plane_table,
- qsa_rollback_gather,
- qsa_rollback_restore,
-)
-
-assert torch.cuda.is_available(), "qsa_rollback_planes requires a CUDA device"
-
-DEV = "cuda"
-TPB = 32
-PAGES = 600
-
-
-def _planes(count, strided, g):
- if strided:
- page = torch.zeros(PAGES, count * TPB * ROW * 2 + 384, dtype=torch.uint8, device=DEV)
- planes = [
- page[:, i * TPB * ROW * 2 : (i + 1) * TPB * ROW * 2]
- .view(torch.bfloat16)
- .view(PAGES, TPB, 1, ROW)
- for i in range(count)
- ]
- else:
- planes = [
- torch.empty(PAGES, TPB, 1, ROW, dtype=torch.bfloat16, device=DEV) for _ in range(count)
- ]
- for t in planes:
- t.copy_(torch.randn(t.shape, generator=g, device=DEV).to(torch.bfloat16))
- return planes
-
-
-def _slots(rows, g):
- pages = torch.randperm(PAGES, generator=torch.Generator().manual_seed(rows))[:rows].to(DEV)
- within = torch.randint(0, TPB, (rows,), generator=g, device=DEV)
- return pages.to(torch.int64), within.to(torch.int64)
-
-
-@pytest.mark.parametrize("strided", [False, True], ids=["contiguous", "shared_page"])
-@pytest.mark.parametrize(
- "cell", qsa_rollback_gather.CELLS, ids=[c.why[:40] for c in qsa_rollback_gather.CELLS]
-)
-def test_gather_cells(cell, strided):
- g = torch.Generator(device=DEV).manual_seed(7)
- planes = _planes(cell.spec["planes"], strided, g)
- pages, within = _slots(cell.spec["rows"], g)
- ptrs, strides = plane_table(planes)
- got = qsa_rollback_gather(ptrs, strides, pages, within, len(planes))
- want = torch.stack([t[pages, within].reshape(-1, ROW) for t in planes])
- assert torch.equal(got.view(torch.int16), want.view(torch.int16))
-
-
-@pytest.mark.parametrize("strided", [False, True], ids=["contiguous", "shared_page"])
-@pytest.mark.parametrize(
- "cell", qsa_rollback_restore.CELLS, ids=[c.why[:40] for c in qsa_rollback_restore.CELLS]
-)
-def test_restore_cells(cell, strided):
- g = torch.Generator(device=DEV).manual_seed(9)
- planes = _planes(cell.spec["planes"], strided, g)
- pages, within = _slots(cell.spec["rows"], g)
- ptrs, strides = plane_table(planes)
- old = qsa_rollback_gather(ptrs, strides, pages, within, len(planes))
- # The step overwrites every slot; the restore puts back the rejected ones only.
- for t in planes:
- t[pages, within] = torch.randn(t[pages, within].shape, generator=g, device=DEV).to(t.dtype)
- after_step = [t.clone() for t in planes]
- # Every other row rejected: each cell restores some slots and keeps the rest.
- rejected = (torch.arange(pages.shape[0], device=DEV) % 2).to(torch.int32)
- qsa_rollback_restore(ptrs, strides, pages, within, rejected, old)
- for i, (t, step) in enumerate(zip(planes, after_step)):
- want = step.clone()
- want[pages, within] = torch.where(
- rejected.bool()[:, None, None], old[i].view(-1, 1, ROW), step[pages, within]
- )
- assert torch.equal(t.view(torch.int16), want.view(torch.int16)), f"plane {i}"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py
deleted file mode 100644
index 0013c41db1bb..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py
+++ /dev/null
@@ -1,108 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the qsa_topk_expand_decode entry.
-
-Cells against the entry's reference (a stable descending sort, ties to the higher block,
-then the expansion entry's reference): the block list and the token selection exactly.
-Columns past each row's visible count carry garbage (huge values, NaN) that must never
-be read. Beyond the cells: a stale-read check under graph replay.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_topk_expand_decode import ( # noqa: E501
- qsa_topk_expand_decode as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "qsa_topk_expand_decode requires a CUDA device"
-
-DEV = "cuda"
-W, K, CR, TOKEN_TOPK = 2560, 512, 4, 2048
-
-
-def _build(spec, seed):
- g = torch.Generator(device=DEV).manual_seed(seed)
- if "chunk" in spec:
- # One request's consecutive context rows: every row's KV length is the chunk's end.
- start, count = spec["chunk"]
- pos = torch.arange(start, start + count, dtype=torch.int64, device=DEV)
- seqlen = torch.full((count,), start + count, dtype=torch.int32, device=DEV)
- else:
- pos = torch.tensor(spec["pos"], dtype=torch.int64, device=DEV)
- seqlen = (pos + 1).to(torch.int32)
- rows = pos.numel()
- scores = torch.randn(rows, W, generator=g, device=DEV).abs() * 3.0
- if spec.get("relu_sparse"):
- # Fewer positive scores than the budget, so the zero class straddles it.
- scores = torch.relu(torch.randn(rows, W, generator=g, device=DEV) - 2.2)
- scores[:, ::7] = -0.0 # negative zero ties with zero
- if spec.get("narrow"):
- scores = 1.0 + 0.05 * torch.rand(rows, W, generator=g, device=DEV)
- if spec.get("lognormal"):
- scores = torch.exp(1.74 + 0.2 * torch.randn(rows, W, generator=g, device=DEV))
- # bf16-like resolution, so equal scores tie; a zero in every other row.
- scores = scores.to(torch.bfloat16).float()
- scores[::2, 5] = 0.0
- ends = ((pos + 1) // CR).to(torch.int32)
- for r in range(rows): # garbage past the visible blocks
- n = int(ends[r])
- scores[r, n:] = 1.0e30
- if n + 3 < W:
- scores[r, n + 3] = float("nan")
- topk = torch.full((rows, K), -7, dtype=torch.int32, device=DEV)
- return [scores, ends, pos, seqlen, topk, CR, TOKEN_TOPK]
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
- certifying(op)
- args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31)
- want_topk, want_sel = op.reference(*args)
- margin = cell.spec.get("out_margin")
- if margin is None:
- sel = op(*args)
- else:
- # Rows [margin, margin + rows) of a step-sized selection; the rows around stay put.
- rows = args[0].shape[0]
- whole = torch.full(
- (rows + 2 * margin, TOKEN_TOPK + CR - 1), -9, dtype=torch.int32, device=DEV
- )
- sel = op(*args, out=whole[margin : margin + rows])
- assert sel.data_ptr() == whole[margin].data_ptr()
- assert (whole[:margin] == -9).all() and (whole[margin + rows :] == -9).all(), (
- "rows outside out"
- )
- assert sel.dtype == torch.int32 and sel.shape == want_sel.shape
- assert torch.equal(args[4], want_topk), "block list"
- assert torch.equal(sel, want_sel.to(torch.int32)), "token selection"
-
-
-def test_reads_after_the_wait_under_graph_replay() -> None:
- args = _build(dict(pos=(8700,)), seed=21)
- src = args[0].clone()
-
- def step():
- args[0].copy_(src * 1.0)
- return op(*args)
-
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(warmup):
- step()
- torch.cuda.current_stream().wait_stream(warmup)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- graph_sel = step()
- for trial in range(3):
- src.copy_(_build(dict(pos=(8700,)), seed=100 + trial)[0])
- graph.replay()
- torch.cuda.synchronize()
- got_topk = args[4].clone()
- sel = op(src.clone(), *args[1:])
- assert torch.equal(graph_sel, sel) and torch.equal(got_topk, args[4]), f"replay {trial}"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_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
deleted file mode 100644
index e0ee54787a3a..000000000000
--- 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
+++ /dev/null
@@ -1,147 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the qsa_topk_select catalog entry.
-
-The op selects; it computes nothing. `_compare` therefore checks the output's
-structure exactly and the *multiset of selected values* bit for bit, never the
-permutation, which the op leaves unspecified.
-"""
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_topk_select import ( # noqa: E501
- TopkExpectation,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_topk_select import ( # noqa: E501
- qsa_topk_select as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401
-
-assert torch.cuda.is_available(), "qsa_topk_select requires a CUDA device"
-
-
-def _compare(out: torch.Tensor, ref: TopkExpectation) -> None:
- """Structure exactly, then the selected values bit for bit.
-
- A selection does no arithmetic, so nothing is tolerated. Values rather
- than offsets are compared so that a tied k-th score -- common here: the
- relu cells are half zeros -- is a pass rather than a coin flip.
- """
- assert out.dtype == torch.int32, out.dtype
- assert out.shape == ref.values.shape, (out.shape, ref.values.shape)
- index_topk = out.shape[1]
- valid = out >= 0
- assert not bool((valid[:, 1:] & ~valid[:, :-1]).any()), "a valid offset follows a -1 pad"
- assert bool((out >= -1).all()), "padding must be exactly -1"
- assert torch.equal(valid.sum(1), ref.window.clamp(max=index_topk)), (
- "each row must select exactly min(index_topk, window) offsets"
- )
- assert bool((out < ref.window[:, None]).all()), "an offset reaches past row_end"
- descending = torch.sort(out, dim=1, descending=True).values
- assert not bool(((descending[:, :-1] == descending[:, 1:]) & (descending[:, 1:] >= 0)).any()), (
- "duplicate offset within a row"
- )
- absolute = (out.long() + ref.row_starts[:, None]).clamp(0, ref.logits.shape[1] - 1)
- picked = torch.gather(ref.logits, 1, absolute).masked_fill(~valid, float("-inf"))
- got = torch.sort(picked, dim=1, descending=True).values
- assert torch.equal(got, ref.values), "selected values are not the window's top-k"
-
-
-DEV = "cuda"
-COMPRESS_RATIO = 4
-BLOCK_TOPK = 512
-#: Pre-fill for every output buffer, so a slot the op failed to write shows up
-#: as an offset past `row_end` in `_compare`.
-SENTINEL = 12345
-
-
-def _logits(kind: str, rows: int, cols: int, gen: torch.Generator) -> torch.Tensor:
- x = torch.randn(rows, cols, generator=gen, device=DEV)
- # QSA's scores are relu(q.k) summed over heads: non-negative with ties at 0.
- return torch.relu(x) * 3.0 if kind == "relu" else x
-
-
-def _build(spec: dict, seed: int):
- gen = torch.Generator(device=DEV).manual_seed(seed)
- rows, cols, index_topk = spec["rows"], spec["cols"], spec["index_topk"]
- logits = _logits(spec["logits"], rows, cols, gen)
-
- kind, start = spec["windows"]
- if kind == "positions":
- positions = torch.tensor(start, device=DEV)
- else:
- 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, :]
- outside = (columns < row_starts[:, None]) | (columns >= row_ends[:, None])
- logits.masked_fill_(outside, float("nan"))
- indices = torch.full((rows, index_topk), SENTINEL, dtype=torch.int32, device=DEV)
- return logits, row_starts, row_ends, indices, index_topk
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811
- certifying(op)
- args = _build(cell.spec, seed=sum(map(ord, cell.why)))
- out = op(*args)
- assert out is args[3], "the op writes `indices` and returns it"
- _compare(out, op.reference(*args))
-
-
-def test_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)
- row_starts = torch.zeros(rows, dtype=torch.int32, device=DEV)
- row_ends = torch.full((rows,), cols, dtype=torch.int32, device=DEV)
- indices = torch.empty(rows, BLOCK_TOPK, dtype=torch.int32, device=DEV)
-
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(warmup):
- op(logits, row_starts, row_ends, indices, BLOCK_TOPK)
- torch.cuda.current_stream().wait_stream(warmup)
- torch.cuda.synchronize()
-
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- op(logits, row_starts, row_ends, indices, BLOCK_TOPK)
-
- logits.copy_(_logits("relu", rows, cols, gen))
- row_ends.fill_(1500)
- indices.fill_(SENTINEL)
- graph.replay()
- torch.cuda.synchronize()
- _compare(indices, op.reference(logits, row_starts, row_ends, indices, BLOCK_TOPK))
-
-
-def _layout_case():
- torch.manual_seed(10)
- rows, cols, top_k = 4, 128, 8
- logits = torch.randn(rows, cols, device=DEV)
- row_starts = torch.zeros(rows, dtype=torch.int32, device=DEV)
- row_ends = torch.full((rows,), cols, dtype=torch.int32, device=DEV)
- return logits, row_starts, row_ends, top_k
-
-
-def test_guard_refuses_the_token_budget_default() -> None:
- """The default `index_topk` is 2048, the token budget; the target's buffer
- is 512 wide. That is the wider-buffer fault the other way round, and the
- guard catches it at the call that forgot to pass the block budget."""
- logits, row_starts, row_ends, _ = _layout_case()
- indices = torch.empty(logits.shape[0], BLOCK_TOPK, dtype=torch.int32, device=DEV)
- with pytest.raises(AssertionError, match="index_topk=2048"):
- with validating(op):
- op(logits, row_starts, row_ends, indices)
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
deleted file mode 100644
index cb7a12755720..000000000000
--- 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
+++ /dev/null
@@ -1,219 +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_attn_decode_prologue entry: every cell bitwise the
-composition it replaces (`qwen4_exp_qk_norm_rope`, `qsa_paged_kv_store` or
-`qsa_paged_kv_store_nvfp4`, `qsa_paged_index_store`, `qsa_prefill_compress`) on every
-allocation either writes -- the projection, the index query, the K/V pool (with its block
-scales over NVFP4), the index plane, the position cache; a stale-read check under graph
-replay with PDL over both pool kinds.
-
-The NVFP4 cells lay the pools out as the engine serves them: K/V the two role planes of a
-`[pages, 2, Hkv, tpb, D/2]` e2m1 pool and its `[pages, 2, Hkv, tpb, D/16]` scale pool,
-and the index plane and the position cache two views of one padded page (page stride
-12 x 4096 + 192 bf16 elements, the position page at its tail)."""
-
-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
-#: The NVFP4 engine's index page: 12 index buffers of tpb x 128, then the position page.
-INDEX_PAGE, POSITION_PAD = 12 * TPB * ID + 192, 192
-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 _edge_rows(proj) -> None:
- """Zero, -0.0 and outlier blocks in the k heads (past the rotary prefix, where the
- norm keeps zeros zero) and in the v heads, over the first rows (up to four)."""
- k0, v0 = HQ * D, (HQ + HK) * D
- edges = (
- (0, slice(k0 + 128, k0 + 144), 0.0),
- (1, slice(k0 + D + 160, k0 + D + 176), -0.0),
- (2, slice(k0 + 200, k0 + 201), 900.0),
- (0, slice(v0 + 16, v0 + 32), 0.0),
- (1, slice(v0 + D, v0 + D + 16), torch.tensor([-0.0] * 8 + [0.0] * 8)),
- (3, slice(v0 + 40, v0 + 41), -900.0),
- (2, slice(v0 + D + 64, v0 + D + 80), 1e-30),
- )
- for r, cols, value in edges:
- if r < proj.shape[0]:
- proj[r, cols] = value
-
-
-def _rows(spec):
- """(request slot, position) per row: one row a request, or under `verify` the
- `rows_per_request` consecutive positions of each request from its first."""
- per = spec.get("rows_per_request", 1)
- return [(1 + r, p + i) for r, p in enumerate(spec["pos"]) for i in range(per)]
-
-
-def _build(spec, seed):
- g = torch.Generator(device=DEV).manual_seed(seed)
- layout = _rows(spec)
- pos = [p for _, p in layout]
- rows = len(pos)
- last = {}
- for req, p in layout:
- last[req] = max(last.get(req, 0), p)
- slots = max(SLOTS, max(last) + 2)
- pages = max(SLOTS * WIDTH // 2, sum(p // TPB + 1 for p in last.values()) + 64)
- perm = torch.randperm(pages, generator=g, device=DEV).to(torch.int32)
- block_table = torch.full((slots, WIDTH), -1, dtype=torch.int32, device=DEV)
- used = 0
- for req, p in last.items():
- n = p // TPB + 1
- block_table[req, :n] = perm[used : used + n]
- used += n
- request = torch.tensor([req for req, _ in layout], dtype=torch.int32, device=DEV)
- if spec.get("unmapped") is not None:
- request[spec["unmapped"]] = slots + 5
- logical = torch.tensor(pos, dtype=torch.int64, device=DEV)
- ids = torch.stack([logical, logical + 3, logical + 7]).to(torch.int32).contiguous()
- buf = torch.empty(3, 1, rows + 5, dtype=torch.int32, device=DEV)
- buf[:, 0, :rows] = ids
- coords = torch.transpose(buf[:, 0, :rows], 0, 1)
- 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)
- if spec.get("edge"):
- _edge_rows(proj)
- 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)
- ]
- if spec.get("nvfp4"):
- data = torch.randint(
- 0, 256, (pages, 2, HK, TPB, D // 2), generator=g, device=DEV, dtype=torch.uint8
- )
- scales = torch.randint(
- 0, 256, (pages, 2, HK, TPB, D // 16), generator=g, device=DEV, dtype=torch.uint8
- )
- plane = torch.randn(pages, INDEX_PAGE, generator=g, device=DEV).to(torch.bfloat16)
- index_cache = plane[:, : TPB * ID].view(pages, TPB, 1, ID)
- position_cache = plane[:, INDEX_PAGE - POSITION_PAD :].view(torch.int32)
- position_cache = position_cache.view(pages, TPB, 3)
- position_cache.copy_(
- torch.randint(0, TABLE, (pages, TPB, 3), generator=g, device=DEV, dtype=torch.int32)
- )
- global_scales = spec.get("global_scales", (1.0, 1.0, 1.0))
- kv = dict(
- k_cache=data[:, 0], v_cache=data[:, 1], k_scale_cache=scales[:, 0],
- v_scale_cache=scales[:, 1],
- kv_scale_orig_quant=torch.tensor(global_scales, dtype=torch.float32, device=DEV),
- ) # fmt: skip
- else:
- pool = 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
- )
- kv = dict(k_cache=pool[:, 0], v_cache=pool[:, 1])
- if spec.get("scores"):
- # Columns past a row's visible blocks are not written: a sentinel shows a stray one.
- kv["index_scores"] = torch.full(
- (rows, WIDTH * TPB // CR), -7.0, dtype=torch.float32, device=DEV
- )
- 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],
- index_cache=index_cache, position_cache=position_cache, request_indices=request,
- logical_positions=logical, block_table=block_table,
- rows_per_request=spec.get("rows_per_request", 1), **kv, **GEOM,
- ) # fmt: skip
-
-
-def _bytes(storage, device) -> torch.Tensor:
- return torch.empty(0, dtype=torch.uint8, device=device).set_(storage)
-
-
-def _clone(args):
- """The arguments over fresh copies of every allocation they view, each view rebuilt at
- its offset and strides; returns them and the copies, keyed by the first argument
- viewing each allocation."""
- out, copies, keys = dict(args), {}, {}
- for name, t in args.items():
- if not isinstance(t, torch.Tensor):
- continue
- key = keys.setdefault(t.untyped_storage().data_ptr(), name)
- if key not in copies:
- copies[key] = _bytes(t.untyped_storage(), t.device).clone()
- view = torch.empty(0, dtype=t.dtype, device=t.device)
- view.set_(copies[key].untyped_storage(), t.storage_offset(), t.shape, t.stride())
- out[name] = view
- return out, copies
-
-
-def _assert_same(got, want, where="") -> None:
- """Every allocation bitwise equal: what the op writes and everything it must not."""
- assert got[1].keys() == want[1].keys()
- for key, copy in got[1].items():
- differ = (copy != want[1][key]).nonzero()
- assert differ.numel() == 0, f"{where}{key}: {differ.shape[0]} bytes differ"
-
-
-@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[0])
- got = _clone(args)
- op(**got[0])
- torch.cuda.synchronize()
- _assert_same(got, want)
-
-
-@pytest.mark.parametrize(
- "spec",
- [
- dict(pos=(8699, 8700)),
- dict(pos=(8697, 8700), rows_per_request=4, nvfp4=True),
- dict(pos=(8699, 8700), nvfp4=True, scores=True),
- ],
- ids=["bf16_pool", "nvfp4_pool_verify", "nvfp4_pool_scores"],
-)
-def test_reads_after_the_wait_under_graph_replay(spec) -> None:
- """The projection is rewritten inside the graph right before the op; with PDL a read of
- it before the wait would see the previous replay's row (in a verify step, also the
- raw index keys of a group's earlier rows)."""
- args = _build(spec, seed=23)
- src = args["proj"].clone()
- work, work_copies = _clone(args)
-
- def step():
- work["proj"].copy_(src * 1.0)
- op(**work)
-
- side = torch.cuda.Stream()
- side.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(side):
- step()
- torch.cuda.current_stream().wait_stream(side)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- step()
- 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[0])
- want[0]["proj"].copy_(src)
- op.reference(**want[0])
- _assert_same((work, work_copies), want, where=f"replay {trial}: ")
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__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
deleted file mode 100644
index d5404666e2c6..000000000000
--- 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
+++ /dev/null
@@ -1,162 +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_qk_norm_rope entry."""
-
-import functools
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qwen4_exp_qk_norm_rope import ( # noqa: E501
- CKPT,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qwen4_exp_qk_norm_rope import ( # noqa: E501
- qwen4_exp_qk_norm_rope as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import ( # noqa: E402,F401 — path above; fixture by name
- assert_within_ulp,
- certifying,
- validating,
-)
-
-assert torch.cuda.is_available(), "qwen4_exp_qk_norm_rope requires a CUDA device"
-
-THETA = 1e7
-EPS = 1e-6
-TABLE_ROWS = 262144 # the checkpoint's max_position_embeddings
-#: Spec keys that shape the test's inputs rather than the op call.
-_NOT_OP_ARGS = ("tokens", "max_position", "proj_width", "index_scratch_heads")
-
-
-def _seed(text: str) -> int:
- return zlib.crc32(text.encode())
-
-
-@functools.lru_cache(maxsize=1)
-def _table() -> torch.Tensor:
- """The target's rope table, built the way `post_load_weights` builds it."""
- rot = CKPT["rotary_dim"]
- exponent = torch.arange(0, rot, 2, dtype=torch.float64, device="cuda")
- inv_freq = (1.0 / (THETA ** (exponent / rot))).to(torch.float32)
- positions = torch.arange(TABLE_ROWS, dtype=torch.float32, device="cuda")
- angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0)
- return torch.cat([torch.cos(angles), torch.sin(angles)], dim=1).contiguous()
-
-
-def _distinct_rows(tokens: int, max_position: int, g: torch.Generator) -> torch.Tensor:
- """A contiguous int32 `[3, T]` whose three rows differ at every token."""
- row0 = torch.randint(0, max_position, (tokens,), generator=g, device="cuda")
- d1 = torch.randint(1, max_position, (tokens,), generator=g, device="cuda")
- d2 = torch.randint(1, max_position - 1, (tokens,), generator=g, device="cuda")
- d2 = d2 + (d2 >= d1).long()
- rows = torch.stack([row0, (row0 + d1) % max_position, (row0 + d2) % max_position])
- return rows.to(torch.int32).contiguous()
-
-
-def _build(spec, seed):
- """`(proj, scratch, kwargs)`: the fused projection row, the index scratch
- whose first `index_heads` heads are the `index_q` view, and the op's
- keyword arguments. Norm weights near 1 put q/k at their largest magnitude
- under the Gemma `1 + w`."""
- g = torch.Generator(device="cuda").manual_seed(seed)
- t = spec["tokens"]
- proj = torch.randn(t, spec["proj_width"], generator=g, device="cuda").to(torch.bfloat16)
- scratch = torch.randn(
- t, spec["index_scratch_heads"], spec["index_head_dim"], generator=g, device="cuda"
- ).to(torch.bfloat16)
-
- def weight(n):
- return (torch.randn(n, generator=g, device="cuda") * 0.1 + 1.0).to(torch.bfloat16)
-
- positions = _distinct_rows(t, spec.get("max_position", 32768), g)
- kwargs = {k: v for k, v in spec.items() if k not in _NOT_OP_ARGS}
- kwargs.update(
- index_q=scratch[:, : spec["index_heads"]],
- position_ids=positions,
- cos_sin=_table(),
- q_weight=weight(spec["head_dim"]),
- k_weight=weight(spec["head_dim"]),
- index_q_weight=weight(spec["index_head_dim"]),
- eps=EPS,
- )
- return proj, scratch, kwargs
-
-
-def _deviation(out: torch.Tensor, ref: torch.Tensor) -> str:
- ulp = 2.0**-8
- o, r = out.float(), ref.float()
- row = r.abs().amax(dim=-1, keepdim=True).clamp_min(1e-9)
- elt = ((o - r).abs() / row).max().item() / ulp
- rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ulp
- return f"max abs {(o - r).abs().max().item():.6f}, elt {elt:.2f} ulp, rms {rms:.3f} ulp"
-
-
-def _main_cols(spec) -> int:
- return (spec["num_heads_q"] + spec["num_heads_k"]) * spec["head_dim"]
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811
- certifying(op)
- spec = cell.spec
- proj, scratch, kwargs = _build(spec, seed=_seed(cell.why))
- before, scratch_before = proj.clone(), scratch.clone()
- ref_proj, ref_iq = op.reference(proj, **kwargs)
- assert torch.equal(proj, before), "the reference mutated its input"
- out_proj, out_iq = op(proj, **kwargs)
- assert out_proj is proj and out_iq is kwargs["index_q"]
- main = _main_cols(spec)
- t = spec["tokens"]
- flat_ref = ref_iq.reshape(t, -1)
- flat_out = out_iq.reshape(t, -1)
- print(f"\n {cell.why[:50]}: main {_deviation(proj[:, :main], ref_proj[:, :main])}")
- print(f" {cell.why[:50]}: index {_deviation(flat_out, flat_ref)}")
- # fp32 throughout and one rounding, against an fp64 reference over the same
- # table: the two can differ by one output rounding step, which is up to 2
- # ulp of the row scale (an element just above a power of two in a row whose
- # largest magnitude sits just below the next one). Measured: max 1.48 ulp
- # element-wise, relative RMS <= 0.011 ulp, so the RMS gate is tight.
- assert_within_ulp(proj[:, :main], ref_proj[:, :main], element_ulp=2.0, rms_ulp=0.25)
- assert_within_ulp(flat_out, flat_ref, element_ulp=2.0, rms_ulp=0.25)
- assert torch.equal(proj[:, main:], before[:, main:]), "a column past the main k heads moved"
- heads = spec["index_heads"]
- assert torch.equal(scratch[:, heads:], scratch_before[:, heads:]), (
- "index_q wrote past its heads"
- )
-
-
-def test_cuda_graph_replay_reads_fresh_inputs() -> None:
- """Captured once at M=4, replayed after the inputs are rewritten in place:
- the replay equals an eager call on the new inputs, bit for bit."""
- spec = dict(tokens=4, **CKPT)
- proj, scratch, kwargs = _build(spec, seed=_seed("graph"))
- proj_new, _, kwargs_new = _build(spec, seed=_seed("graph-new"))
- eager_proj, eager_scratch = proj_new.clone(), scratch.clone()
- eager_kwargs = dict(kwargs_new, index_q=eager_scratch[:, : CKPT["index_heads"]])
- op(eager_proj, **eager_kwargs)
- stream = torch.cuda.Stream()
- stream.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(stream):
- op(proj.clone(), **dict(kwargs, index_q=scratch.clone()[:, : CKPT["index_heads"]]))
- torch.cuda.current_stream().wait_stream(stream)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- op(proj, **kwargs)
- proj.copy_(proj_new)
- for name in ("position_ids", "q_weight", "k_weight", "index_q_weight"):
- kwargs[name].copy_(kwargs_new[name])
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(proj, eager_proj)
- assert torch.equal(kwargs["index_q"], eager_kwargs["index_q"])
-
-
-def test_guard_refuses_a_position_past_the_table() -> None:
- """An out-of-range position would read past the rope table."""
- proj, _, kwargs = _build(dict(tokens=3, **CKPT), seed=_seed("guard"))
- kwargs["position_ids"][0, 0] = TABLE_ROWS
- with validating(op), pytest.raises(AssertionError, match="past the rope table"):
- op(proj, **kwargs)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_spec_rollback_commit.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_spec_rollback_commit.py
deleted file mode 100644
index 9d8456e59f62..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_spec_rollback_commit.py
+++ /dev/null
@@ -1,101 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU tests for the spec_rollback_commit entries, driven by their CELLS.
-
-`spec_rollback_stage` is compared bitwise with its reference. `spec_rollback_commit` has no
-reference of its own (test_modeling_v2_spec_rollback drives it through the handler); here each
-cell is checked bitwise against the docstring's per-request rewrite, written out in torch.
-"""
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.spec_rollback_commit import ( # noqa: E501
- ROW,
- spec_rollback_commit,
- spec_rollback_stage,
-)
-
-assert torch.cuda.is_available(), "spec_rollback_commit requires a CUDA device"
-
-DEV = "cuda"
-CTX = 2
-CHANNELS = 640
-STATE = 3
-TPB = 32
-
-
-def _slots(n: int, g: torch.Generator) -> torch.Tensor:
- return (torch.randperm(2 * n + 3, generator=g)[:n] + 1).to(DEV, torch.int32)
-
-
-@pytest.mark.parametrize(
- "cell", spec_rollback_stage.CELLS, ids=[c.why for c in spec_rollback_stage.CELLS]
-)
-def test_stage_cells(cell) -> None:
- n, p = cell.spec["n"], cell.spec["p"]
- g = torch.Generator().manual_seed(n * 31 + p)
- slots = _slots(n, g)
- pool_slots = 2 * n + 4
- ngram_pool = torch.randint(0, 1 << 40, (pool_slots, CTX), generator=g).to(DEV)
- conv_pool = torch.randn(pool_slots, CHANNELS, STATE, generator=g).to(DEV, torch.bfloat16)
- ids = torch.randint(0, 1 << 20, (n * p,), generator=g).to(DEV, torch.int32)
- spec_rollback_stage.is_valid(ngram_pool, conv_pool, slots, ids, p)
- stream, conv_prev = spec_rollback_stage(ngram_pool, conv_pool, slots, ids, p)
- want_stream, want_prev = spec_rollback_stage.reference(ngram_pool, conv_pool, slots, ids, p)
- assert torch.equal(stream, want_stream)
- assert torch.equal(conv_prev, want_prev)
-
-
-@pytest.mark.parametrize(
- "cell", spec_rollback_commit.CELLS, ids=[c.why for c in spec_rollback_commit.CELLS]
-)
-def test_commit_cells(cell) -> None:
- n, p, planes = cell.spec["n"], cell.spec["p"], cell.spec["planes"]
- g = torch.Generator().manual_seed(n * 17 + planes)
- rows = n * p
- slots = _slots(n, g)
- pool_slots = 2 * n + 4
- accepted = torch.randint(1, p + 1, (n,), generator=g).to(DEV, torch.int32)
- ngram_pool = torch.randint(0, 1 << 40, (pool_slots, CTX), generator=g).to(DEV)
- stream = torch.randint(0, 1 << 40, (n, CTX + p), generator=g).to(DEV)
- conv_pool = torch.randn(pool_slots, CHANNELS, STATE, generator=g).to(DEV, torch.bfloat16)
- conv_prev = torch.randn(n, CHANNELS, STATE, generator=g).to(DEV, torch.bfloat16)
- conv_x = torch.randn(n, p, CHANNELS, generator=g).to(DEV, torch.bfloat16)
- # Each row its own (page, slot within the page).
- flat = torch.randperm(rows * 2, generator=g)[:rows]
- pages, within = (flat // TPB).to(DEV), (flat % TPB).to(DEV)
- num_pages = int(pages.max()) + 1
- position = torch.randint(0, 1 << 20, (num_pages, TPB, 3), generator=g).to(DEV, torch.int32)
- position_old = torch.randint(0, 1 << 20, (rows, 3), generator=g).to(DEV, torch.int32)
- plane_list = [
- torch.randn(num_pages, TPB, ROW, generator=g).to(DEV, torch.bfloat16) for _ in range(planes)
- ]
- plane_ptrs = torch.tensor([t.data_ptr() for t in plane_list], dtype=torch.int64, device=DEV)
- plane_strides = torch.tensor([t.stride(0) for t in plane_list], dtype=torch.int64, device=DEV)
- index_old = torch.randn(planes, rows, ROW, generator=g).to(DEV, torch.bfloat16)
-
- want_ngram, want_conv = ngram_pool.clone(), conv_pool.clone()
- want_position, want_planes = position.clone(), [t.clone() for t in plane_list]
- for i in range(n):
- a, slot = int(accepted[i]), int(slots[i])
- want_ngram[slot] = stream[i, a : a + CTX]
- cat = torch.cat([conv_prev[i], conv_x[i].T], 1)
- want_conv[slot] = cat[:, a : a + STATE]
- for t in range(a, p):
- r = i * p + t
- pg, tok = int(pages[r]), int(within[r])
- want_position[pg, tok] = position_old[r]
- for plane, want in zip(range(planes), want_planes):
- want[pg, tok] = index_old[plane, r]
-
- args = (accepted, slots, ngram_pool, stream, conv_pool, conv_prev, conv_x, position,
- position_old, pages, within, plane_ptrs, plane_strides, index_old) # fmt: skip
- spec_rollback_commit.is_valid(*args)
- spec_rollback_commit(*args)
- torch.cuda.synchronize()
- assert torch.equal(ngram_pool, want_ngram)
- assert torch.equal(conv_pool, want_conv)
- assert torch.equal(position, want_position)
- for got, want in zip(plane_list, want_planes):
- assert torch.equal(got, want)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_verify_step_setup.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_verify_step_setup.py
deleted file mode 100644
index d7824aaa72e2..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_verify_step_setup.py
+++ /dev/null
@@ -1,139 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test: `verify_step_setup` against the composed ops of the decode target's step setup
-(`pos3` copy, `cu_seqlens_from_seq_lens`, `SpeculativeRollback.begin`'s slot arithmetic and
-`qsa_rollback_gather`), bitwise, eagerly and replayed from a CUDA graph."""
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_rollback_planes import ( # noqa: E501
- plane_table,
- qsa_rollback_gather,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.verify_step_setup import ( # noqa: E501
- verify_step_setup,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.cu_seqlens_from_seq_lens import ( # noqa: E501
- cu_seqlens_from_seq_lens,
-)
-
-assert torch.cuda.is_available(), "the verify step setup test requires a CUDA device"
-
-DEV = "cuda"
-TPB = 32
-PAGES = 96
-LAYERS = 12
-MAX_BLOCKS = 6
-
-
-def _case(batch, p, seed, negative=False):
- g = torch.Generator().manual_seed(seed)
- rows = batch * p
- full = torch.randint(0, 1 << 20, (3, 1, rows + 5), generator=g).to(DEV, torch.int32)
- position_ids = full[..., :rows]
- seq_lens = torch.full((batch,), p, dtype=torch.int32, device=DEV)
- slots = torch.randint(0, 64, (batch,), generator=g).to(DEV, torch.int32)
- req = torch.arange(batch).repeat_interleave(p).to(DEV, torch.int32)
- starts = torch.randint(0, MAX_BLOCKS * TPB - p - 20, (batch,), generator=g)
- logical = (starts.repeat_interleave(p) + torch.arange(p).repeat(batch)).to(DEV, torch.int64)
- if negative:
- logical[:p] -= starts[0] + 40 # a padded request: negative positions, one page back
- block_table = torch.randint(0, PAGES, (batch + 3, MAX_BLOCKS + 2), generator=g)
- block_table = block_table.to(DEV, torch.int32)[:, :MAX_BLOCKS]
- planes = [
- torch.randn(PAGES, TPB, 1, 128, generator=g).to(DEV, torch.bfloat16) for _ in range(LAYERS)
- ]
- return position_ids, seq_lens, slots, req, logical, block_table, planes
-
-
-def _reference(position_ids, seq_lens, slots, req, logical, block_table, ptrs, strides):
- pos3 = torch.empty([3, position_ids.shape[-1]], dtype=torch.int32, device=DEV)
- pos3.copy_(torch.reshape(position_ids, [3, -1]))
- qsl = cu_seqlens_from_seq_lens(seq_lens, torch.int64)
- positions = logical.to(torch.long)
- pages = block_table[req.to(torch.long), positions // TPB].to(torch.long)
- within = positions % TPB
- index = qsa_rollback_gather(ptrs, strides, pages, within, LAYERS)
- return pos3, qsl, slots.to(torch.long), pages, within, index
-
-
-def _bits(t):
- return t.view(torch.int16) if t.dtype == torch.bfloat16 else t
-
-
-@pytest.mark.parametrize("negative", [False, True], ids=["positions", "negative_positions"])
-@pytest.mark.parametrize("batch,p", [(1, 4), (17, 8), (255, 2)])
-def test_setup_matches_composed_ops(batch, p, negative):
- position_ids, seq_lens, slots, req, logical, block_table, planes = _case(
- batch, p, batch * 7 + p, negative
- )
- ptrs, strides = plane_table(planes)
- ref = _reference(position_ids, seq_lens, slots, req, logical, block_table, ptrs, strides)
- got = verify_step_setup(
- position_ids, seq_lens, slots, req, logical, block_table, TPB, ptrs, strides, LAYERS
- )
- for a, b in zip(got, ref):
- assert a.dtype == b.dtype and a.shape == b.shape
- assert torch.equal(_bits(a), _bits(b))
-
-
-def test_setup_zero_length_rows_in_segmentation():
- position_ids, seq_lens, slots, req, logical, block_table, planes = _case(6, 2, 5)
- seq_lens.copy_(torch.tensor([2, 2, 2, 2, 2, 2], dtype=torch.int32))
- seq_lens[3:] = 0
- ptrs, strides = plane_table(planes)
- got = verify_step_setup(
- position_ids, seq_lens, slots, req, logical, block_table, TPB, ptrs, strides, LAYERS
- )
- assert torch.equal(got[1], cu_seqlens_from_seq_lens(seq_lens, torch.int64))
-
-
-def test_setup_cuda_graph_follows_inputs():
- position_ids, seq_lens, slots, req, logical, block_table, planes = _case(8, 4, 11)
- ptrs, strides = plane_table(planes)
- args = (position_ids, seq_lens, slots, req, logical, block_table, TPB, ptrs, strides, LAYERS)
- verify_step_setup(*args)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- out = verify_step_setup(*args)
- g = torch.Generator().manual_seed(12)
- for _ in range(3):
- logical.add_(torch.randint(0, 3, logical.shape, generator=g).to(DEV))
- position_ids.add_(1)
- for plane in planes:
- plane.copy_(torch.randn(plane.shape, generator=g).to(DEV, torch.bfloat16))
- graph.replay()
- torch.cuda.synchronize()
- ref = _reference(position_ids, seq_lens, slots, req, logical, block_table, ptrs, strides)
- for a, b in zip(out, ref):
- assert torch.equal(_bits(a), _bits(b))
-
-
-@pytest.mark.parametrize(
- "cell", verify_step_setup.CELLS, ids=[c.why for c in verify_step_setup.CELLS]
-)
-def test_setup_cells(cell) -> None:
- """Each cell against the docstring's outputs written out in torch indexing."""
- batch, p, num_planes = cell.spec["batch"], cell.spec["p"], cell.spec["planes"]
- position_ids, seq_lens, slots, req, logical, block_table, planes = _case(batch, p, 21 * batch)
- planes = planes[:num_planes]
- ptrs, strides = plane_table(planes)
- args = (position_ids, seq_lens, slots, req, logical, block_table, TPB, ptrs, strides,
- num_planes) # fmt: skip
- verify_step_setup.is_valid(*args)
- pos3, qsl, slots_long, pages, within, index = verify_step_setup(*args)
- rows = batch * p
- assert torch.equal(pos3, position_ids.reshape(3, rows).to(torch.int32))
- want_qsl = torch.zeros(batch + 1, dtype=torch.int64, device=DEV)
- want_qsl[1:] = torch.cumsum(seq_lens.to(torch.int64), 0)
- assert torch.equal(qsl, want_qsl)
- assert torch.equal(slots_long, slots.to(torch.int64))
- want_pages = block_table[req.long(), torch.div(logical, TPB, rounding_mode="floor")].long()
- want_within = torch.remainder(logical, TPB)
- assert torch.equal(pages, want_pages)
- assert torch.equal(within, want_within)
- want_index = torch.stack(
- [plane[want_pages, want_within].reshape(rows, 128) for plane in planes]
- )
- assert torch.equal(_bits(index), _bits(want_index))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_vision_rope_2d.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_vision_rope_2d.py
deleted file mode 100644
index 2a83d03b3fc4..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_vision_rope_2d.py
+++ /dev/null
@@ -1,120 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the qwen4_exp vision_rope_2d entry.
-
-`op.reference` shares no code with the entry: positions from nested loops,
-angles in fp64 from `theta ** (-2j / rope_dim)`, rotation as a complex
-multiply.
-"""
-
-import zlib
-from typing import Tuple
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.vision_rope_2d import ( # noqa: E501
- vision_rope_2d as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401
-
-assert torch.cuda.is_available(), "vision_rope_2d requires a CUDA device"
-
-
-def _compare(
- out: Tuple[torch.Tensor, torch.Tensor],
- ref: Tuple[torch.Tensor, torch.Tensor],
-) -> None:
- """Default band on each of `(q_out, k_out)`; nothing is loosened."""
- for got, want in zip(out, ref, strict=True):
- torch.testing.assert_close(got, want)
-
-
-HEADS = 16
-HEAD_DIM = 72
-MERGE = 2
-THETA = 10000.0
-
-
-def _num_tokens(grid) -> int:
- return sum(t * h * w for t, h, w in grid)
-
-
-def _operands(tokens, heads, heads_k, head_dim, dtype, seed, layout="separate"):
- g = torch.Generator(device="cuda").manual_seed(seed)
- if layout == "packed_qkv":
- assert heads == heads_k
- qkv = torch.randn(tokens, 3 * heads, head_dim, generator=g, device="cuda").to(dtype)
- q, k, _ = torch.split(qkv, [heads] * 3, 1)
- assert not q.is_contiguous() and q.stride(-1) == 1
- return q, k
- q = torch.randn(tokens, heads, head_dim, generator=g, device="cuda").to(dtype)
- k = torch.randn(tokens, heads_k, head_dim, generator=g, device="cuda").to(dtype)
- return q, k
-
-
-def _build(spec, seed):
- grid = spec["grid"]
- heads = spec.get("heads", HEADS)
- q, k = _operands(
- _num_tokens(grid),
- heads,
- spec.get("heads_k", heads),
- spec.get("head_dim", HEAD_DIM),
- spec["dtype"],
- seed,
- spec.get("layout", "separate"),
- )
- grid_thw = torch.tensor(
- [list(row) for row in grid],
- dtype=spec.get("grid_dtype", torch.int32),
- device=spec.get("grid_device", "cpu"),
- )
- return q, k, grid_thw, MERGE, spec.get("theta", THETA)
-
-
-def _max_deviation(lhs, rhs) -> float:
- return (lhs.float() - rhs.float()).abs().max().item()
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying): # noqa: F811
- certifying(op)
- q, k, grid_thw, merge, theta = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
- q_before, k_before = q.clone(), k.clone()
- q_out, k_out = op(q, k, grid_thw, merge, theta)
- assert q_out.shape == q.shape and q_out.dtype == q.dtype
- assert k_out.shape == k.shape and k_out.dtype == k.dtype
- assert torch.equal(q, q_before) and torch.equal(k, k_before), "inputs must not be mutated"
- assert q_out.data_ptr() != q.data_ptr() and k_out.data_ptr() != k.data_ptr()
- q_ref, k_ref = op.reference(q, k, grid_thw, merge, theta)
- print(
- f"[{cell.why}] q_dev={_max_deviation(q_out, q_ref):.3e} "
- f"k_dev={_max_deviation(k_out, k_ref):.3e}"
- )
- _compare((q_out, k_out), (q_ref, k_ref))
-
-
-def test_grid_placement_and_integer_dtype_do_not_change_the_result() -> None:
- """The target hands over an int64 grid on the device; the cells pass host grids."""
- grid = [(1, 16, 24), (1, 8, 12)]
- q, k = _operands(_num_tokens(grid), HEADS, HEADS, HEAD_DIM, torch.bfloat16, seed=29)
- host = torch.tensor(grid, dtype=torch.int32)
- q_host, k_host = op(q, k, host, MERGE, THETA)
- for other in (host.to("cuda"), host.to(torch.int64), host.to(torch.int64).to("cuda")):
- q_other, k_other = op(q, k, other, MERGE, THETA)
- assert torch.equal(q_host, q_other) and torch.equal(k_host, k_other)
-
-
-def test_guard_refuses_a_one_token_k() -> None:
- """The hazard, unguarded: a [1, H, D] k broadcasts to [T, H, D]."""
- grid_thw = torch.tensor([(1, 4, 4)], dtype=torch.int32)
- q, k = _operands(16, HEADS, HEADS, HEAD_DIM, torch.bfloat16, seed=5)
- _, k_out = op(q, k[:1], grid_thw, MERGE, THETA)
- assert k_out.shape == (16, HEADS, HEAD_DIM)
-
- with pytest.raises(AssertionError, match="must agree on tokens and head_dim"):
- with validating(op):
- op(q, k[:1], grid_thw, MERGE, THETA)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/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
deleted file mode 100644
index a402400fc12d..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/embedding/test_modeling_v2_ngram_hash.py
+++ /dev/null
@@ -1,132 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the ngram_hash entry.
-
-Exact integer arithmetic, gated bitwise against the entry's scalar reference.
-Beyond the cells: the EOS reset that `ngram_token_history`'s EOS seed relies
-on, the narrow-multiplier guard and graph capture.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.embedding.ngram_hash import ( # noqa: E501
- CKPT_CONTEXT_LEN,
- CKPT_EOS_TOKEN_ID,
- CKPT_HEAD_OFFSETS,
- CKPT_HEAD_VOCAB_SIZES,
- CKPT_MULTIPLIERS,
- CKPT_VOCAB_SIZE,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.embedding.ngram_hash import ( # noqa: E501
- ngram_hash as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "ngram_hash requires a CUDA device"
-
-
-def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
- """Bitwise: exact integer arithmetic, no tolerance to loosen."""
- assert out.dtype == ref.dtype == torch.int64, (out.dtype, ref.dtype)
- assert out.shape == ref.shape, (out.shape, ref.shape)
- assert out.is_contiguous()
- assert torch.equal(out, ref), f"{int((out != ref).sum())} indices differ"
-
-
-DEV = "cuda"
-EOS = CKPT_EOS_TOKEN_ID
-NGRAM_HEADS = len(CKPT_HEAD_VOCAB_SIZES)
-PADDED_TABLE_ROWS = 320001536 # sum(head_vocab_sizes) rounded up to a multiple of 128
-
-
-def _constants():
- return (
- torch.tensor(CKPT_MULTIPLIERS, dtype=torch.long, device=DEV),
- torch.tensor(CKPT_HEAD_VOCAB_SIZES, dtype=torch.long, device=DEV),
- torch.tensor(CKPT_HEAD_OFFSETS, dtype=torch.long, device=DEV),
- )
-
-
-def _tokens(batch, length, seed, eos_positions=(), lead_eos=True, eos_density=0.0):
- g = torch.Generator().manual_seed(seed)
- ids = torch.randint(0, CKPT_VOCAB_SIZE, (batch, length), generator=g, dtype=torch.long)
- if eos_density > 0.0:
- ids[torch.rand(batch, length, generator=g) < eos_density] = EOS
- if lead_eos:
- ids[:, : min(CKPT_CONTEXT_LEN, length)] = EOS
- for position in eos_positions:
- ids[:, position] = EOS
- return ids.to(DEV)
-
-
-def _hash(ids, num_query_tokens):
- return op(ids, *_constants(), EOS, num_query_tokens)
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- spec = cell.spec
- ids = _tokens(
- spec["batch"],
- spec["length"],
- seed=zlib.crc32(cell.why.encode()),
- eos_positions=spec.get("eos_positions", ()),
- lead_eos=spec.get("lead_eos", True),
- eos_density=spec.get("eos_density", 0.0),
- )
- before = ids.clone()
- args = (ids, *_constants(), EOS, spec["num_query_tokens"])
- out = op(*args)
- assert out.shape == (spec["batch"], spec["num_query_tokens"], NGRAM_HEADS)
- assert out.data_ptr() != ids.data_ptr()
- assert torch.equal(ids, before), "the history is read-only"
- _compare(out, op.reference(*args))
-
-
-def test_eos_restarts_the_window() -> None:
- """One past an EOS hashes as a sequence's first token, and only there."""
- token = 4242
- fresh = torch.tensor([[EOS] * CKPT_CONTEXT_LEN + [token]], device=DEV)
- after_eos = torch.tensor([[11, 22, 33, EOS, token]], device=DEV)
- no_eos = torch.tensor([[11, 22, 33, 44, token]], device=DEV)
- assert torch.equal(_hash(fresh, 1), _hash(after_eos, 1))
- assert not torch.equal(_hash(fresh, 1), _hash(no_eos, 1))
-
-
-def test_guard_refuses_narrow_multipliers() -> None:
- ids = _tokens(1, CKPT_CONTEXT_LEN + 4, seed=47)
- multipliers, sizes, offsets = _constants()
- narrow = multipliers.to(torch.int32)
- unguarded = op(ids, narrow, sizes, offsets, EOS, 4)
- assert not torch.equal(unguarded, _hash(ids, 4)), "the hazard this guard covers is gone"
- assert int(unguarded.min()) >= 0 and int(unguarded.max()) < PADDED_TABLE_ROWS
- with pytest.raises(AssertionError, match="int64"):
- with validating(op):
- op(ids, narrow, sizes, offsets, EOS, 4)
-
-
-def test_cuda_graph_capture_and_replay() -> None:
- ids = _tokens(8, CKPT_CONTEXT_LEN + 1, seed=53, lead_eos=False, eos_density=0.2)
- constants = _constants()
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(warmup):
- for _ in range(3):
- op(ids, *constants, EOS, 1)
- torch.cuda.current_stream().wait_stream(warmup)
-
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- captured = op(ids, *constants, EOS, 1)
-
- fresh = _tokens(8, CKPT_CONTEXT_LEN + 1, seed=54, lead_eos=False, eos_density=0.2)
- ids.copy_(fresh)
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(captured, op.reference(fresh, *constants, EOS, 1))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_bf16_gemv_pdl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_bf16_gemv_pdl.py
deleted file mode 100644
index ba0d3a05c9df..000000000000
--- 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
+++ /dev/null
@@ -1,98 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the bf16_gemv_pdl entry.
-
-Cells against the fp32 reference within the band `low_m_bf16_mm` uses (the
-same products, summed in another order). Beyond the cells: a PDL chain under
-CUDA-graph capture and replay, where each GEMV consumes the previous one's
-output and so would read a stale activation if it read before its wait; and
-the layout guard.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.bf16_gemv_pdl import ( # noqa: E501
- bf16_gemv_pdl as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import ( # noqa: E402,F401 — path above; fixture by name
- ULP,
- assert_within_ulp,
- certifying,
- validating,
-)
-
-assert torch.cuda.is_available(), "bf16_gemv_pdl requires a CUDA device"
-
-
-def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
- """Accumulation order only: 4 bf16 ulp element-wise against the row scale,
- 2 ulp of relative RMS, as for `low_m_bf16_mm` and `cublas_mm`."""
- assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0, ulp=ULP[torch.bfloat16])
-
-
-def _build(m, n, k, seed):
- """`mat_b` is a `[N, K]` weight's `.t()`, as the target stores its projections."""
- g = torch.Generator(device="cuda").manual_seed(seed)
- mat_a = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16)
- weight = (torch.randn(n, k, generator=g, device="cuda") * k**-0.5).to(torch.bfloat16)
- return mat_a, weight.t()
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811
- certifying(op)
- s = cell.spec
- mat_a, mat_b = _build(s["m"], s["n"], s["k"], seed=zlib.crc32(cell.why.encode()) % 2**31)
- out = op(mat_a, mat_b, s["pdl"])
- assert out.shape == (s["m"], s["n"]) and out.dtype == torch.bfloat16 and out.is_contiguous()
- _compare(out, op.reference(mat_a, mat_b, s["pdl"]))
-
-
-def test_dependent_chain_under_graph_replay() -> None:
- """Decode runs this under CUDA graphs, back to back with other kernels. Each
- link of the chain reads the output of the one before it: a GEMV that read
- its activation before `griddepcontrol.wait` would see the previous replay's
- values, and the comparison below would fail on the second replay."""
- pdl = True
- dims = [2560, 6144, 2560, 640, 2560] # K of each link; N is the next link's K
- weights = [_build(1, n, k, seed=i)[1] for i, (k, n) in enumerate(zip(dims, dims[1:] + [16480]))]
- x = torch.randn(1, dims[0], device="cuda").to(torch.bfloat16)
-
- def chain():
- h = x
- outs = []
- for w in weights:
- h = op(h, w, pdl)
- outs.append(h)
- return outs
-
- stream = torch.cuda.Stream()
- stream.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(stream):
- for _ in range(2):
- chain()
- torch.cuda.current_stream().wait_stream(stream)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- outs = chain()
- for trial in range(3):
- g = torch.Generator(device="cuda").manual_seed(100 + trial)
- x.copy_(torch.randn(x.shape, generator=g, device="cuda").to(x.dtype))
- graph.replay()
- torch.cuda.synchronize()
- h = x
- for w, out in zip(weights, outs):
- # Each link against the reference of what the previous link actually stored.
- _compare(out, op.reference(h, w))
- h = out
-
-
-def test_guard_refuses_a_row_major_mat_b() -> None:
- mat_a, mat_b = _build(1, 512, 2560, seed=3)
- with validating(op), pytest.raises(AssertionError):
- op(mat_a, mat_b.contiguous())
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_fp8_gemv.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_fp8_gemv.py
deleted file mode 100644
index 2500b3fa4722..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_fp8_gemv.py
+++ /dev/null
@@ -1,221 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test: `fp8_gemv` / `fp8_gemv_conv1d` / `fp8_gemv_conv1d_verify` against an fp32
-reference over the dequantized weight (128x128-block e4m3), per cell; the verify entry also
-bitwise against the two-launch path it replaces (`fp8_gemv` + `causal_conv1d_triton_update`);
-strided activations; CUDA-graph replay."""
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.fp8_gemv import ( # noqa: E501
- dequantize_fp8_block,
- fp8_gemv_conv1d_op,
- fp8_gemv_conv1d_verify_op,
- fp8_gemv_op,
- pack_weight,
- quantize_fp8_block,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_triton_update import ( # noqa: E501
- causal_conv1d_triton_update,
-)
-
-assert torch.cuda.is_available(), "fp8_gemv requires a CUDA device"
-
-DEV = "cuda"
-
-
-def _weight(n, k, seed):
- g = torch.Generator(device=DEV).manual_seed(seed)
- w = torch.randn(n, k, generator=g, device=DEV) * k**-0.5
- w[:, 7] *= 30.0 # outlier column
- w[: min(n, 128), 128:256] *= 1e-4 # a tiny block: subnormal codes
- codes, scale = quantize_fp8_block(w.to(torch.bfloat16))
- return codes, scale, pack_weight(codes, scale)
-
-
-def _close(y, ref, x, deq) -> None:
- """Within one bf16 rounding of the fp32 reference (plus fp32 summation noise)."""
- with true_fp32_matmul():
- absref = x.float().abs() @ deq.abs().t()
- ulp = torch.exp2(torch.floor(torch.log2(ref.abs().clamp_min(1e-30))) - 7)
- excess = ((y.float() - ref).abs() - 0.5 * ulp).clamp_min(0) / absref.clamp_min(1e-30)
- assert excess.max().item() < 1e-5, f"error beyond one rounding: {excess.max().item():.2e}"
-
-
-@pytest.mark.parametrize("cell", fp8_gemv_op.CELLS, ids=[c.why[:40] for c in fp8_gemv_op.CELLS])
-def test_gemv_cells(cell) -> None:
- n, k, m = cell.spec["n"], cell.spec["k"], cell.spec["m"]
- codes, scale, packed = _weight(n, k, n + m)
- g = torch.Generator(device=DEV).manual_seed(m)
- # A column slice of a wider row, as the target's split projections hand over.
- wide = torch.randn(m, k + 256, generator=g, device=DEV).to(torch.bfloat16)
- x = wide[:, :k]
- y = fp8_gemv_op(x, packed)
- assert y.shape == (m, n) and y.dtype == torch.bfloat16
- _close(y, fp8_gemv_op.reference(x, codes, scale), x, dequantize_fp8_block(codes, scale))
-
-
-@pytest.mark.parametrize(
- "cell", fp8_gemv_conv1d_op.CELLS, ids=[c.why[:40] for c in fp8_gemv_conv1d_op.CELLS]
-)
-def test_conv1d_cells(cell) -> None:
- n, k, conv_dim, pool, pad = 16480, 2560, 10240, 64, -1
- codes, scale, packed = _weight(n, k, 3)
- deq = dequantize_fp8_block(codes, scale)
- slots = list(cell.spec["slots"])
- idx = torch.tensor(slots, dtype=torch.int32, device=DEV)
- g = torch.Generator(device=DEV).manual_seed(11)
- backing = torch.randn(pool, conv_dim, 4, generator=g, device=DEV).to(torch.bfloat16)
- conv_state = backing[:, :, 1:] # a strided pool view
- conv_weight = torch.randn(conv_dim, 4, generator=g, device=DEV).to(torch.bfloat16)
- x = torch.randn(len(slots), k, generator=g, device=DEV).to(torch.bfloat16)
- before = conv_state.clone()
- with true_fp32_matmul():
- proj = (x.float() @ deq.t()).to(torch.bfloat16)
- want, want_state = proj.clone(), conv_state.clone()
- for r, s in enumerate(slots):
- if s == pad:
- continue
- window = torch.cat([conv_state[s].float(), proj[r, :conv_dim].float()[:, None]], 1)
- c = (window * conv_weight.float()).sum(1)
- want[r, :conv_dim] = (c * torch.sigmoid(c)).to(torch.bfloat16)
- want_state[s] = window[:, 1:].to(torch.bfloat16)
- out = fp8_gemv_conv1d_op(x, packed, conv_state, conv_weight, idx, conv_dim, pad)
- # The projection may round one ulp differently than the reference's fp32 matmul; a flip
- # carries through the conv. Nearly every output is bitwise the reference's.
- assert (out == want).float().mean().item() > 0.999
- assert ((out.float() - want.float()).abs() / want.float().abs().clamp_min(1e-2)).max() < 0.02
- assert ((conv_state.float() - want_state.float()).abs()).max().item() < 0.02
- others = [s for s in range(pool) if s not in slots]
- assert torch.equal(conv_state[others], before[others]), "other slots untouched"
-
-
-def test_graph_replay_matches_eager() -> None:
- _, _, packed = _weight(2560, 6144, 5)
- x = torch.randn(4, 6144, device=DEV).to(torch.bfloat16)
- fp8_gemv_op(x, packed)
- torch.cuda.synchronize()
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- out = fp8_gemv_op(x, packed)
- x.copy_(torch.randn(4, 6144, device=DEV).to(torch.bfloat16))
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(out, fp8_gemv_op(x, packed))
-
-
-def _verify_args(spec, seed):
- """`fp8_gemv_conv1d_verify`'s arguments for a cell: the activation a column slice of a wider
- row, the pool a strided view (every 4th column skipped), slot 0's window zero (a fresh
- sequence), the snapshot buffer pre-filled so untouched steps and rows are visible."""
- k, conv_dim, p = spec["k"], spec["conv_dim"], spec["p"]
- slots = list(spec["slots"])
- m = len(slots) * p
- g = torch.Generator(device=DEV).manual_seed(seed)
- wide = torch.randn(m, k + 128, generator=g, device=DEV).to(torch.bfloat16)
- backing = torch.randn(spec["pool_slots"], conv_dim, 4, generator=g, device=DEV)
- conv_state = backing.to(torch.bfloat16)[:, :, 1:]
- conv_state[0] = 0
- conv_weight = torch.randn(conv_dim, 4, generator=g, device=DEV).to(torch.bfloat16)
- window = torch.randn(
- spec["snapshot_rows"], spec["snapshot_steps"], conv_dim, 3, generator=g, device=DEV
- ).to(torch.bfloat16)
- idx = torch.tensor(slots, dtype=torch.int32, device=DEV)
- rows = torch.arange(len(slots), dtype=torch.int32, device=DEV)
- return [wide[:, :k], conv_state, conv_weight, idx, window, rows, conv_dim, p, -1]
-
-
-def _two_launch(x, packed, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad):
- """The target's verify path before the fusion: the projection, then the Triton conv over
- the transposed q|k|v view; returns `(conv, rest)` in the fused entry's layout."""
- proj = fp8_gemv_op(x, packed)
- n = x.shape[0] // p
- x3 = torch.transpose(torch.reshape(proj[:, :conv_dim], [n, p, conv_dim]), 1, 2)
- out = causal_conv1d_triton_update(
- x3, conv_state, conv_weight, None, "silu", idx, pad, None, window, rows, False
- )
- return torch.reshape(torch.transpose(out, 1, 2), [n * p, conv_dim]), proj[:, conv_dim:]
-
-
-@pytest.mark.parametrize(
- "cell",
- fp8_gemv_conv1d_verify_op.CELLS,
- ids=[c.why[:40] for c in fp8_gemv_conv1d_verify_op.CELLS],
-)
-def test_conv1d_verify_cells(cell) -> None:
- spec = cell.spec
- codes, scale, packed = _weight(spec["n"], spec["k"], 3)
- x, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad = _verify_args(spec, 7)
- live = [r for r in range(x.shape[0]) if spec["slots"][r // p] != pad]
-
- state_a, window_a = conv_state.clone(), window.clone()
- conv, rest = fp8_gemv_conv1d_verify_op(
- x, packed, state_a, conv_weight, idx, window_a, rows, conv_dim, p, pad
- )
- assert conv.shape == (x.shape[0], conv_dim) and conv.is_contiguous()
- assert rest.shape == (x.shape[0], spec["n"] - conv_dim)
-
- # Bitwise the two-launch path (same projection plan, the Triton conv's arithmetic).
- state_b, window_b = conv_state.clone(), window.clone()
- conv_b, rest_b = _two_launch(
- x, packed, state_b, conv_weight, idx, window_b, rows, conv_dim, p, pad
- )
- assert torch.equal(rest.view(torch.int16), rest_b.view(torch.int16))
- assert torch.equal(conv[live].view(torch.int16), conv_b[live].view(torch.int16))
- assert torch.equal(state_a.view(torch.int16), state_b.view(torch.int16)), "whole pool"
- assert torch.equal(window_a.view(torch.int16), window_b.view(torch.int16)), "whole buffer"
-
- # Against the fp32 reference: the projection may round one ulp differently (fp32 summation
- # order), and a flip carries through the conv (one tap times one ulp of the input: within
- # 2^-5 of the row's peak), the snapshots and the pool (one ulp: within 2^-7 of the peak).
- want, want_rest, want_state, want_window = fp8_gemv_conv1d_verify_op.reference(
- x, codes, scale, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad
- )
- deq = dequantize_fp8_block(codes, scale)
- with true_fp32_matmul():
- ref = x.float() @ deq.t()
- _close(rest, ref[:, conv_dim:], x, deq[conv_dim:])
- assert (conv == want).float().mean().item() > 0.999
- peak = want.float().abs().amax(1, keepdim=True)
- assert ((conv.float() - want.float()).abs() <= 2.0**-5 * peak).all()
- for got, ref_ in ((state_a, want_state), (window_a, want_window)):
- assert (got == ref_).float().mean().item() > 0.9999
- assert (got.float() - ref_.float()).abs().max() <= 2.0**-7 * ref_.float().abs().max()
- # Non-vacuity: every live slot and snapshot row was rewritten.
- for r, slot in enumerate(spec["slots"]):
- if slot != pad:
- assert not torch.equal(state_a[slot], conv_state[slot]), slot
- assert not torch.equal(window_a[r, :p], window[r, :p]), r
- if spec["snapshot_steps"] > p:
- assert torch.equal(window_a[:, p:], window[:, p:]), "steps past p untouched"
-
-
-def test_conv1d_verify_graph_replay_matches_eager() -> None:
- spec = fp8_gemv_conv1d_verify_op.CELLS[2].spec
- _, _, packed = _weight(spec["n"], spec["k"], 5)
- x, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad = _verify_args(spec, 9)
- state0, window0 = conv_state.clone(), window.clone()
- fp8_gemv_conv1d_verify_op(
- x, packed, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad
- )
- torch.cuda.synchronize()
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- conv, rest = fp8_gemv_conv1d_verify_op(
- x, packed, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad
- )
- x.copy_(torch.randn(x.shape, device=DEV).to(torch.bfloat16))
- conv_state.copy_(state0)
- window.copy_(window0)
- graph.replay()
- torch.cuda.synchronize()
- got = (conv.clone(), rest.clone(), conv_state.clone(), window.clone())
- conv_state.copy_(state0)
- window.copy_(window0)
- want = fp8_gemv_conv1d_verify_op(
- x, packed, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad
- )
- for a, b in zip(got, (*want, conv_state, window)):
- assert torch.equal(a.view(torch.int16), b.view(torch.int16))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_fp8_per_tensor_mm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_fp8_per_tensor_mm.py
deleted file mode 100644
index dce1286b3460..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_fp8_per_tensor_mm.py
+++ /dev/null
@@ -1,83 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test: `fp8_per_tensor_mm` against its fp32 reference (the activation quantized as the
-op does, both operands dequantized) per cell, the activation quantization bitwise against
-torch's e4m3 cast, and CUDA-graph replay with a new activation."""
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.fp8_per_tensor_mm import ( # noqa: E501
- E4M3_MAX,
- fp8_per_tensor_mm,
- fp8_per_tensor_weight,
- fp8_quantize_per_tensor,
-)
-
-assert torch.cuda.is_available(), "fp8_per_tensor_mm requires a CUDA device"
-
-DEV = "cuda"
-
-
-def _operands(m, k, n, seed):
- g = torch.Generator(device=DEV).manual_seed(seed)
- x = (torch.randn(m, k, generator=g, device=DEV) * 0.7).to(torch.bfloat16)
- w = (torch.randn(n, k, generator=g, device=DEV) * k**-0.5).to(torch.bfloat16)
- codes, scale = fp8_per_tensor_weight(w)
- return x, w, codes.t(), scale
-
-
-@pytest.mark.parametrize(
- "cell", fp8_per_tensor_mm.CELLS, ids=[c.why[:40] for c in fp8_per_tensor_mm.CELLS]
-)
-def test_cells(cell) -> None:
- m, k, n = cell.spec["m"], cell.spec["k"], cell.spec["n"]
- x, w, codes_t, scale = _operands(m, k, n, 1000 + m)
- got = fp8_per_tensor_mm(x, codes_t, scale)
- want = fp8_per_tensor_mm.reference(x, codes_t, scale)
- assert got.dtype == torch.bfloat16 and got.shape == (m, n)
- rel = (got.float() - want).norm() / want.norm()
- # The bf16 output rounding and cuBLASLt's accumulation separate them: its small-M FP8
- # kernels accumulate below fp32 precision (measured 8.0e-3 at one row, under 4e-3 from
- # 128 rows) -- a fifth of the e4m3 quantization error the op is gated on below.
- assert rel < 1.2e-2, f"relative error {rel:.2e} against the dequantized reference"
- full = x.float() @ w.float().t()
- lossy = (got.float() - full).norm() / full.norm()
- assert lossy < 0.06, f"FP8 error against the bf16 projection {lossy:.3f}"
-
-
-def test_quantization_is_torchs_cast() -> None:
- g = torch.Generator(device=DEV).manual_seed(5)
- x = (torch.randn(37, 2560, generator=g, device=DEV) * 3).to(torch.bfloat16)
- q, scale = fp8_quantize_per_tensor(x)
- want_scale = x.abs().amax().float() / E4M3_MAX
- assert torch.equal(scale, want_scale.reshape(1))
- want = (x.float() * (E4M3_MAX / x.abs().amax().float())).to(torch.float8_e4m3fn)
- assert (q.view(torch.uint8) != want.view(torch.uint8)).float().mean() < 1e-4
-
-
-def test_zero_activation_and_strided_rows() -> None:
- x, _, codes_t, scale = _operands(8, 2560, 2560, 3)
- zeros = torch.zeros_like(x)
- assert torch.equal(
- fp8_per_tensor_mm(zeros, codes_t, scale),
- torch.zeros(8, 2560, dtype=torch.bfloat16, device=DEV),
- )
- wide = torch.zeros(8, 4096, dtype=torch.bfloat16, device=DEV)
- wide[:, :2560] = x
- assert torch.equal(
- fp8_per_tensor_mm(wide[:, :2560], codes_t, scale), fp8_per_tensor_mm(x, codes_t, scale)
- )
-
-
-def test_graph_replay_matches_eager() -> None:
- x, _, codes_t, scale = _operands(64, 2560, 16480, 7)
- fp8_per_tensor_mm(x, codes_t, scale)
- torch.cuda.synchronize()
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- out = fp8_per_tensor_mm(x, codes_t, scale)
- x.copy_(_operands(64, 2560, 16480, 8)[0] * 4)
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(out, fp8_per_tensor_mm(x, codes_t, scale))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_low_m_bf16_mm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_low_m_bf16_mm.py
deleted file mode 100644
index d3b3a17e22a0..000000000000
--- 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
+++ /dev/null
@@ -1,148 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the low_m_bf16_mm entry.
-
-Every cell names the branch `plan` must pick, and a host-only check pins the
-branch of the table's other shapes and M buckets, so a table edit that moves a
-shape to another kernel fails here before it can change a number silently.
-Beyond the cells: CUDA-graph capture and replay after warmup (the way the
-decode forward runs it), and the guard against a padded `mat_b`, which the
-cuBLAS branch would read as dense.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.low_m_bf16_mm import ( # noqa: E501
- low_m_bf16_mm as op,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.low_m_bf16_mm import ( # noqa: E501
- plan,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import ( # noqa: E402,F401 — path above; fixture by name
- ULP,
- assert_within_ulp,
- certifying,
- validating,
-)
-
-assert torch.cuda.is_available(), "low_m_bf16_mm requires a CUDA device"
-
-
-def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
- """Accumulation order only: 4 bf16 ulp element-wise against the row scale,
- 2 ulp of relative RMS -- the band `cublas_mm` uses, for the same reason.
- Every branch sums the same bf16 products in fp32 and rounds once; what is
- left between kernel and fp32 reference is the final bf16 rounding (half an
- ulp) plus the summation order."""
- assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0, ulp=ULP[torch.bfloat16])
-
-
-def _build(m, n, k, seed):
- """`mat_b` is a `[N, K]` weight's `.t()`, as a target stores its projections."""
- g = torch.Generator(device="cuda").manual_seed(seed)
- mat_a = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16)
- weight = (torch.randn(n, k, generator=g, device="cuda") * k**-0.5).to(torch.bfloat16)
- return mat_a, weight.t()
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811
- certifying(op)
- s = cell.spec
- assert plan(s["m"], s["n"], s["k"]).kind == s["kind"], plan(s["m"], s["n"], s["k"])
- mat_a, mat_b = _build(s["m"], s["n"], s["k"], seed=zlib.crc32(cell.why.encode()) % 2**31)
- out = op(mat_a, mat_b, s["pdl"])
- assert out.shape == (s["m"], s["n"]) and out.dtype == torch.bfloat16 and out.is_contiguous()
- _compare(out, op.reference(mat_a, mat_b, s["pdl"]))
-
-
-#: `(M, N, K, branch)` of the table's shapes and M buckets the cells do not drive:
-#: what `plan` picks for each, held without a launch.
-_PLAN_PINS = (
- (1, 10240, 2560, "direct"),
- (1, 48, 2560, "direct"),
- (32, 48, 2560, "direct"),
- (8, 2560, 6144, "splitk"),
- (2, 16480, 2560, "cublas"),
- (1, 2560, 6144, "direct"),
- (3, 12288, 2560, "splitk"),
- (8, 12288, 2560, "splitk"),
- (16, 12288, 2560, "cublas"),
- (16, 512, 2560, "direct"),
- (24, 512, 2560, "splitk"),
- (16, 640, 2560, "splitk"),
- (4, 1280, 2560, "direct"),
- (8, 1296, 2560, "splitk"),
- (1, 1, 2560, "direct"),
- (1, 16480, 2560, "direct"),
- (2, 14080, 2560, "direct"),
- (3, 14080, 2560, "splitk"),
- (8, 14080, 2560, "splitk"),
- (16, 14080, 2560, "cublas"),
- (1, 2560, 2560, "direct"),
- (4, 2560, 2560, "direct"),
- (8, 2560, 2560, "splitk"),
- (1, 2560, 640, "cublas"),
- (8192, 2560, 6144, "cublas"),
- (1, 1024, 2048, "cublas"),
-)
-
-
-def test_plan_pins_the_branch_of_the_other_table_shapes() -> None:
- picked = [(m, n, k, want, plan(m, n, k).kind) for m, n, k, want in _PLAN_PINS]
- wrong = [p for p in picked if p[3] != p[4]]
- assert not wrong, f"(M, N, K, pinned, picked): {wrong}"
-
-
-def test_graph_replay_after_warmup() -> None:
- """Decode runs this under CUDA graphs, with PDL: compile in warmup, capture,
- then replay with new activations copied into the captured input. One call per
- branch, in a chain, so a FlashInfer kernel follows cuBLAS and vice versa."""
- pdl = True
- shapes = [(1, 10240, 2560), (8, 2560, 6144), (1, 2560, 640), (8, 12288, 2560)]
- assert [plan(*s).kind for s in shapes] == ["direct", "splitk", "cublas", "splitk"]
- ins = [_build(m, n, k, seed=i) for i, (m, n, k) in enumerate(shapes)]
- static = [a.clone() for a, _ in ins]
- stream = torch.cuda.Stream()
- stream.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(stream):
- for _ in range(2):
- [op(a, b, pdl) for a, (_, b) in zip(static, ins)]
- torch.cuda.current_stream().wait_stream(stream)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- outs = [op(a, b, pdl) for a, (_, b) in zip(static, ins)]
- for trial in range(3):
- g = torch.Generator(device="cuda").manual_seed(100 + trial)
- for a in static:
- a.copy_(torch.randn(a.shape, generator=g, device="cuda").to(a.dtype))
- graph.replay()
- torch.cuda.synchronize()
- for a, (_, b), out in zip(static, ins, outs):
- _compare(out, op.reference(a, b))
-
-
-def _padded(m, n, k):
- """A row-padded mat_a and a column-padded mat_b: dense along K, leading
- dimension 2K. Both are views a caller gets by slicing a wider tensor."""
- a = torch.randn(m, 2 * k, dtype=torch.bfloat16, device="cuda")[:, :k]
- w = (torch.randn(n, 2 * k, device="cuda") * k**-0.5).to(torch.bfloat16)[:, :k]
- return a, w.t()
-
-
-def test_guard_refuses_a_padded_mat_b() -> None:
- """Why the guard exists: the cuBLAS branch ignores mat_b's leading dimension
- and silently computes a different product."""
- a, b = _padded(4, 640, 640)
- a = a.contiguous()
- assert plan(4, 640, 640).kind == "cublas" and b.stride() == (1, 1280)
- with pytest.raises(AssertionError):
- _compare(op(a, b), op.reference(a, b))
- with pytest.raises(AssertionError, match="column-major"):
- with validating(op):
- op(a, b)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_mid_rows_mm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_mid_rows_mm.py
deleted file mode 100644
index 959822b4b891..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_mid_rows_mm.py
+++ /dev/null
@@ -1,119 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test: `mid_rows_mm` (FP8 with the per-row activation quantization, and bf16) against an
-fp32 reference over the operands it multiplies, per cell; `fp8_row_quant` bitwise against
-torch's e4m3 cast; strided activation rows; CUDA-graph replay with new activations."""
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.mid_rows_mm import ( # noqa: E501
- AMAX_FLOOR,
- E4M3_MAX,
- MidWeight,
- fp8_row_quant,
- fp8_row_quant_op,
- mid_rows_mm_op,
- precompile,
-)
-
-assert torch.cuda.is_available(), "mid_rows_mm requires a CUDA device"
-
-DEV = "cuda"
-
-
-def _weight(n, k, mode, seed):
- g = torch.Generator(device=DEV).manual_seed(seed)
- w = torch.randn(n, k, generator=g, device=DEV) * k**-0.5
- w[:, 7] *= 20.0 # outlier column
- w = w.to(torch.bfloat16)
- if mode == "bf16":
- return w, MidWeight(w.contiguous())
- if mode == "row":
- scale = (w.float().abs().amax(1) / E4M3_MAX).clamp_min(1e-12)
- return w, MidWeight((w.float() / scale[:, None]).to(torch.float8_e4m3fn), scale)
- scale = (w.float().abs().amax() / E4M3_MAX).clamp_min(1e-12).reshape(1)
- return w, MidWeight((w.float() / scale).to(torch.float8_e4m3fn), scale)
-
-
-def _activation(m, k, seed, pad=0):
- g = torch.Generator(device=DEV).manual_seed(seed)
- wide = torch.randn(m, k + pad, generator=g, device=DEV) * 0.7
- wide[:, 3] *= 12.0 # outlier column
- if m > 2:
- wide[1] = 0.0 # an all-zero row
- return wide.to(torch.bfloat16)[:, :k]
-
-
-def _dequantized(x, mw):
- """The fp32 operands the call multiplies."""
- if not mw.fp8:
- return x.float(), mw.w.float()
- q, s = fp8_row_quant(x)
- wd = mw.w.float() * (mw.scale[:, None] if mw.scale.numel() > 1 else mw.scale)
- return q.float() * s[:, None], wd
-
-
-def _close(y, x, mw) -> None:
- """Within one bf16 rounding of the fp32 product of the operands (plus fp32 summation
- noise relative to sum |x w|)."""
- xd, wd = _dequantized(x, mw)
- with true_fp32_matmul():
- ref = xd @ wd.t()
- absref = xd.abs() @ wd.abs().t()
- ulp = torch.exp2(torch.floor(torch.log2(ref.abs().clamp_min(1e-30))) - 7)
- excess = ((y.float() - ref).abs() - 0.5 * ulp).clamp_min(0) / absref.clamp_min(1e-30)
- assert torch.isfinite(y.float()).all()
- assert excess.max().item() < 1e-5, f"error beyond one rounding: {excess.max().item():.2e}"
-
-
-@pytest.mark.parametrize(
- "cell", mid_rows_mm_op.CELLS, ids=[c.why[:48] for c in mid_rows_mm_op.CELLS]
-)
-def test_cells(cell) -> None:
- n, k, m, mode = cell.spec["n"], cell.spec["k"], cell.spec["m"], cell.spec["mode"]
- w, mw = _weight(n, k, mode, n + m)
- x = _activation(m, k, m, pad=64 if mode != "bf16" else 0)
- mid_rows_mm_op.is_valid(x, mw)
- y = mid_rows_mm_op(x, mw)
- assert y.shape == (m, n) and y.dtype == torch.bfloat16 and y.is_contiguous()
- _close(y, x, mw)
- if mw.fp8:
- # The FP8 product against the bf16 projection: e4m3 keeps 3 mantissa bits on both
- # operands; per-row activation scales are at least as accurate as one per tensor.
- full = x.float() @ w.float().t()
- assert ((y.float() - full).norm() / full.norm()).item() < 0.06
-
-
-@pytest.mark.parametrize(
- "cell", fp8_row_quant_op.CELLS, ids=[c.why[:48] for c in fp8_row_quant_op.CELLS]
-)
-def test_row_quant_is_torchs_cast(cell) -> None:
- m, k, pad = cell.spec["m"], cell.spec["k"], cell.spec.get("pad", 0)
- x = _activation(m, k, 5 + m, pad=pad) * 3
- if m > 4:
- x[4, 9] = 900.0 # a row whose maximum is far above the rest
- q, s = fp8_row_quant_op(x)
- # The kernel divides correctly rounded (torch divides by a scalar as a product with
- # its reciprocal): the quotients in float64, rounded once.
- amax = x.float().abs().amax(1).clamp_min(AMAX_FLOOR).double()
- assert torch.equal(s, (amax / E4M3_MAX).float())
- want = (x.float() * (E4M3_MAX / amax).float()[:, None]).to(torch.float8_e4m3fn)
- assert torch.equal(q.view(torch.uint8), want.view(torch.uint8))
-
-
-def test_graph_replay_matches_eager() -> None:
- for mode, n, k, m in (("tensor", 16480, 2560, 64), ("bf16", 512, 2560, 200)):
- _, mw = _weight(n, k, mode, 3)
- precompile(mw)
- x = _activation(m, k, 1).contiguous()
- mid_rows_mm_op(x, mw)
- torch.cuda.synchronize()
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- out = mid_rows_mm_op(x, mw)
- x.copy_(_activation(m, k, 2) * 4)
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(out, mid_rows_mm_op(x, mw))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_nvfp4_gemv.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_nvfp4_gemv.py
deleted file mode 100644
index 7f5a273112d2..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_nvfp4_gemv.py
+++ /dev/null
@@ -1,326 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test: `nvfp4_gemv` against an fp32 reference over the dequantized NVFP4 weight, per
-cell; the quantizer against its definition; CUDA-graph replay; one compiled kernel per row
-octet; the f16-range guard. The split entries (NVFP4 units on `fp8_gemv`'s kernel) per cell:
-the projection against the fp32 reference, the conv entries also against the two-launch
-paths they replace (the split projection + `causal_conv1d_update` /
-`causal_conv1d_triton_update`)."""
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm import ( # noqa: E501
- nvfp4_gemv as nv,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_triton_update import ( # noqa: E501
- causal_conv1d_triton_update,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_update import ( # noqa: E501
- causal_conv1d_update,
-)
-
-assert torch.cuda.is_available(), "nvfp4_gemv requires a CUDA device"
-
-DEV = "cuda"
-nvfp4_gemv_op = nv.nvfp4_gemv_op
-
-
-def _weight(n, k, seed):
- g = torch.Generator(device=DEV).manual_seed(seed)
- w = torch.randn(n, k, generator=g, device=DEV) * k**-0.5
- w[:, 7] *= 30.0 # outlier column: its blocks' scales run up to e4m3's top
- w[: min(n, 64), 128:256] *= 1e-4 # near-zero blocks: small (subnormal) e4m3 scales
- w = w.to(torch.bfloat16)
- codes, scales, gscale = nv.quantize_nvfp4(w)
- return w, codes, scales, gscale, nv.pack_weight(codes, scales, gscale)
-
-
-def _x(m, k, seed):
- g = torch.Generator(device=DEV).manual_seed(seed)
- # A column slice of a wider row, as the target's split activations hand over.
- return torch.randn(m, k + 256, generator=g, device=DEV).to(torch.bfloat16)[:, :k]
-
-
-def _close(y, ref, x, deq) -> None:
- """Within one bf16 rounding of the fp32 reference, plus fp32 summation noise."""
- with true_fp32_matmul():
- absref = x.float().abs() @ deq.abs().t()
- ulp = torch.exp2(torch.floor(torch.log2(ref.abs().clamp_min(1e-30))) - 7)
- excess = ((y.float() - ref).abs() - 0.5 * ulp).clamp_min(0) / absref.clamp_min(1e-30)
- assert excess.max().item() < 1e-5, f"error beyond one rounding: {excess.max().item():.2e}"
-
-
-@pytest.mark.parametrize("cell", nvfp4_gemv_op.CELLS, ids=[c.why[:40] for c in nvfp4_gemv_op.CELLS])
-def test_gemv_cells(cell) -> None:
- n, k, m = cell.spec["n"], cell.spec["k"], cell.spec["m"]
- _, codes, scales, gscale, packed = _weight(n, k, n + m)
- x = _x(m, k, m)
- nvfp4_gemv_op.is_valid(x, packed)
- y = nvfp4_gemv_op(x, packed)
- assert y.shape == (m, n) and y.dtype == torch.bfloat16
- ref = nvfp4_gemv_op.reference(x, codes, scales, gscale)
- _close(y, ref, x, nv.dequantize_nvfp4(codes, scales, gscale))
-
-
-def test_quantizer_matches_its_definition() -> None:
- """Codes are the nearest e2m1 value of w / (block scale x global scale); the largest block
- scale is e4m3's 448; the relative error is NVFP4's (well under 2^-2 a row)."""
- w, codes, scales, gscale, _ = _weight(512, 1024, 3)
- assert codes.max().item() <= 15 and torch.isfinite(scales.float()).all()
- assert scales.float().max().item() == 448.0
- step = scales.float().repeat_interleave(nv.SF_VEC, 1) * gscale
- v = torch.where(step > 0, w.float() / step.clamp_min(1e-38), 0.0)
- grid = torch.tensor(nv.E2M1_VALUES, device=DEV)
- best = (v.abs()[..., None] - grid).abs().min(-1).values
- got = nv.dequantize_nvfp4(codes, scales, gscale) / step.clamp_min(1e-38)
- err = (v - torch.where(step > 0, got, 0.0)).abs()
- assert (err <= best + 1e-6).all(), "a code is not the nearest e2m1 value"
- deq = nv.dequantize_nvfp4(codes, scales, gscale)
- rel = (deq - w.float()).norm(dim=1) / w.float().norm(dim=1)
- assert rel.max().item() < 0.2, rel.max().item()
-
-
-def test_chunked_pack_matches_whole() -> None:
- """`pack_from_bf16` (row chunks, a partial last tile) is `pack_weight` of the whole."""
- w = torch.randn(1000, 256, device=DEV).to(torch.bfloat16)
- a = nv.pack_from_bf16(w, rows_per_chunk=192)
- b = nv.pack_weight(*nv.quantize_nvfp4(w))
- assert a.global_scale == b.global_scale and torch.equal(a.data, b.data)
-
-
-def test_one_kernel_per_row_octet() -> None:
- """The row count is a launch argument: rows 1..16 build two kernels per shape."""
- _, codes, scales, gscale, packed = _weight(2048, 512, 5)
- before = nv._compiled.cache_info().currsize
- for m in range(1, nv.MAX_ROWS + 1):
- x = _x(m, 512, m)
- _close(
- nvfp4_gemv_op(x, packed),
- nvfp4_gemv_op.reference(x, codes, scales, gscale),
- x,
- nv.dequantize_nvfp4(codes, scales, gscale),
- )
- assert nv._compiled.cache_info().currsize - before <= 2
-
-
-def test_graph_replay_matches_eager() -> None:
- _, _, _, _, packed = _weight(4096, 2560, 9)
- x = torch.randn(4, 2560, device=DEV).to(torch.bfloat16)
- nvfp4_gemv_op(x, packed)
- torch.cuda.synchronize()
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- out = nvfp4_gemv_op(x, packed)
- x.copy_(torch.randn(4, 2560, device=DEV).to(torch.bfloat16))
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(out, nvfp4_gemv_op(x, packed))
-
-
-def test_is_valid_rejects_activation_beyond_f16() -> None:
- _, _, _, _, packed = _weight(256, 256, 1)
- x = torch.randn(2, 256, device=DEV).to(torch.bfloat16)
- nvfp4_gemv_op.is_valid(x, packed)
- x[1, 3] = 1e5
- with pytest.raises(AssertionError):
- nvfp4_gemv_op.is_valid(x, packed)
-
-
-# ------------------------------------------------------------------------------------------
-# split entries: NVFP4 units on fp8_gemv's kernel
-# ------------------------------------------------------------------------------------------
-
-split_op = nv.nvfp4_split_gemv_op
-split_conv_op = nv.nvfp4_split_gemv_conv1d_op
-split_verify_op = nv.nvfp4_split_gemv_conv1d_verify_op
-GDN_N, GDN_K, GDN_CONV = 16480, 2560, 10240
-
-
-@pytest.mark.parametrize("cell", split_op.CELLS, ids=[c.why[:40] for c in split_op.CELLS])
-def test_split_gemv_cells(cell) -> None:
- n, k, m = cell.spec["n"], cell.spec["k"], cell.spec["m"]
- _, codes, scales, gscale, _ = _weight(n, k, n + m)
- packed = nv.pack_split(codes, scales, gscale)
- x = _x(m, k, m)
- split_op.is_valid(x, packed)
- y = split_op(x, packed)
- assert y.shape == (m, n) and y.dtype == torch.bfloat16
- _close(
- y,
- split_op.reference(x, codes, scales, gscale),
- x,
- nv.dequantize_nvfp4(codes, scales, gscale),
- )
-
-
-def test_split_layouts_agree() -> None:
- """Tile-major and k-block-major units give the same projection (same plan)."""
- _, codes, scales, gscale, _ = _weight(1001, 384, 4)
- x = _x(5, 384, 5)
- a = split_op(x, nv.pack_split(codes, scales, gscale, tile_major=True))
- b = split_op(x, nv.pack_split(codes, scales, gscale, tile_major=False))
- assert torch.equal(a, b)
-
-
-@pytest.mark.parametrize("cell", split_conv_op.CELLS, ids=[c.why[:40] for c in split_conv_op.CELLS])
-def test_split_conv1d_cells(cell) -> None:
- pool, pad = 64, -1
- _, codes, scales, gscale, _ = _weight(GDN_N, GDN_K, 3)
- packed = nv.pack_split(codes, scales, gscale)
- deq = nv.dequantize_nvfp4(codes, scales, gscale)
- slots = list(cell.spec["slots"])
- idx = torch.tensor(slots, dtype=torch.int32, device=DEV)
- g = torch.Generator(device=DEV).manual_seed(11)
- backing = torch.randn(pool, GDN_CONV, 4, generator=g, device=DEV).to(torch.bfloat16)
- conv_state = backing[:, :, 1:] # a strided pool view
- conv_weight = torch.randn(GDN_CONV, 4, generator=g, device=DEV).to(torch.bfloat16)
- x = torch.randn(len(slots), GDN_K, generator=g, device=DEV).to(torch.bfloat16)
- before = conv_state.clone()
- split_conv_op.is_valid(x, packed)
- out = split_conv_op(x, packed, conv_state, conv_weight, idx, GDN_CONV, pad)
-
- # Against the two-launch path over the same projection.
- proj = split_op(x, packed)
- state_b = before.clone()
- two = proj.clone()
- causal_conv1d_update(
- x=two[:, :GDN_CONV],
- conv_state=state_b,
- weight=conv_weight,
- bias=None,
- silu_activation=True,
- cache_seqlens=None,
- conv_state_indices=idx,
- pad_slot_id=pad,
- )
- assert torch.equal(out[:, GDN_CONV:], proj[:, GDN_CONV:])
- assert (out == two).float().mean().item() > 0.999
- assert ((out.float() - two.float()).abs() / two.float().abs().clamp_min(1e-2)).max() < 0.02
- assert torch.equal(conv_state, state_b), "the pool as the conv kernel leaves it"
-
- # Against the fp32 reference: the projection may round one ulp differently (fp32 summation
- # order, the activation rounded to f16), and a flip carries through the conv (one tap times
- # one ulp of the input: within 2^-5 of the row's peak).
- with true_fp32_matmul():
- ref = (x.float() @ deq.t()).to(torch.bfloat16)
- want = ref.clone()
- for r, s in enumerate(slots):
- if s == pad:
- continue
- window = torch.cat([before[s].float(), ref[r, :GDN_CONV].float()[:, None]], 1)
- c = (window * conv_weight.float()).sum(1)
- want[r, :GDN_CONV] = (c * torch.sigmoid(c)).to(torch.bfloat16)
- assert (out == want).float().mean().item() > 0.999
- peak = want.float().abs().amax(1, keepdim=True)
- assert ((out.float() - want.float()).abs() <= 2.0**-5 * peak).all()
- others = [s for s in range(pool) if s not in slots]
- assert torch.equal(conv_state[others], before[others]), "other slots untouched"
-
-
-def _verify_args(spec, seed):
- """As the fp8 verify test: a column slice of a wider row, a strided pool (slot 0 zero), a
- pre-filled snapshot buffer."""
- k, conv_dim, p = spec["k"], spec["conv_dim"], spec["p"]
- slots = list(spec["slots"])
- m = len(slots) * p
- g = torch.Generator(device=DEV).manual_seed(seed)
- wide = torch.randn(m, k + 128, generator=g, device=DEV).to(torch.bfloat16)
- backing = torch.randn(spec["pool_slots"], conv_dim, 4, generator=g, device=DEV)
- conv_state = backing.to(torch.bfloat16)[:, :, 1:]
- conv_state[0] = 0
- conv_weight = torch.randn(conv_dim, 4, generator=g, device=DEV).to(torch.bfloat16)
- window = torch.randn(
- spec["snapshot_rows"], spec["snapshot_steps"], conv_dim, 3, generator=g, device=DEV
- ).to(torch.bfloat16)
- idx = torch.tensor(slots, dtype=torch.int32, device=DEV)
- rows = torch.arange(len(slots), dtype=torch.int32, device=DEV)
- return [wide[:, :k], conv_state, conv_weight, idx, window, rows, conv_dim, p, -1]
-
-
-@pytest.mark.parametrize(
- "cell", split_verify_op.CELLS, ids=[c.why[:40] for c in split_verify_op.CELLS]
-)
-def test_split_conv1d_verify_cells(cell) -> None:
- spec = cell.spec
- _, codes, scales, gscale, _ = _weight(spec["n"], spec["k"], 3)
- packed = nv.pack_split(codes, scales, gscale)
- x, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad = _verify_args(spec, 7)
- live = [r for r in range(x.shape[0]) if spec["slots"][r // p] != pad]
- split_verify_op.is_valid(
- x, packed, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad
- )
- state_a, window_a = conv_state.clone(), window.clone()
- conv, rest = split_verify_op(
- x, packed, state_a, conv_weight, idx, window_a, rows, conv_dim, p, pad
- )
- # Bitwise the two-launch path over the same projection.
- proj = split_op(x, packed)
- n = x.shape[0] // p
- state_b, window_b = conv_state.clone(), window.clone()
- x3 = torch.transpose(torch.reshape(proj[:, :conv_dim], [n, p, conv_dim]), 1, 2)
- two = causal_conv1d_triton_update(
- x3, state_b, conv_weight, None, "silu", idx, pad, None, window_b, rows, False
- )
- two = torch.reshape(torch.transpose(two, 1, 2), [n * p, conv_dim])
- assert torch.equal(rest, proj[:, conv_dim:])
- assert torch.equal(conv[live], two[live])
- assert torch.equal(state_a, state_b) and torch.equal(window_a, window_b)
- # Against the fp32 reference.
- want, want_rest, want_state, want_window = split_verify_op.reference(
- x, codes, scales, gscale, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad
- )
- deq = nv.dequantize_nvfp4(codes, scales, gscale)
- with true_fp32_matmul():
- ref = x.float() @ deq.t()
- _close(rest, ref[:, conv_dim:], x, deq[conv_dim:])
- assert (conv == want).float().mean().item() > 0.999
- peak = want.float().abs().amax(1, keepdim=True)
- assert ((conv.float() - want.float()).abs() <= 2.0**-5 * peak).all()
- for got, ref_ in ((state_a, want_state), (window_a, want_window)):
- assert (got == ref_).float().mean().item() > 0.9999
- assert (got.float() - ref_.float()).abs().max() <= 2.0**-7 * ref_.float().abs().max()
-
-
-def test_split_graph_replay_matches_eager() -> None:
- _, codes, scales, gscale, _ = _weight(GDN_N, GDN_K, 9)
- packed = nv.pack_split(codes, scales, gscale)
- x = torch.randn(4, GDN_K, device=DEV).to(torch.bfloat16)
- split_op(x, packed)
- torch.cuda.synchronize()
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- out = split_op(x, packed)
- x.copy_(torch.randn(4, GDN_K, device=DEV).to(torch.bfloat16))
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(out, split_op(x, packed))
-
-
-def test_split_precompile_covers_the_decode_variants() -> None:
- """After `precompile_split`, no projection / conv / verify call within its rows compiles."""
- from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm import ( # noqa: E501
- fp8_gemv as f8,
- )
-
- _, codes, scales, gscale, _ = _weight(GDN_N, GDN_K, 2)
- packed = nv.pack_split(codes, scales, gscale)
- nv.precompile_split(packed, 4, GDN_CONV, -1, (2, 4))
- before = f8._compiled.cache_info().misses
- conv_state = torch.zeros(8, GDN_CONV, 3, device=DEV).to(torch.bfloat16)
- conv_weight = torch.randn(GDN_CONV, 4, device=DEV).to(torch.bfloat16)
- window = torch.zeros(4, 4, GDN_CONV, 3, device=DEV).to(torch.bfloat16)
- for m in range(1, 5):
- x = torch.randn(m, GDN_K, device=DEV).to(torch.bfloat16)
- idx = torch.arange(m, dtype=torch.int32, device=DEV)
- split_op(x, packed)
- split_conv_op(x, packed, conv_state, conv_weight, idx, GDN_CONV, -1)
- for p in (2, 4):
- if m % p == 0:
- n = m // p
- ids = torch.arange(n, dtype=torch.int32, device=DEV)
- split_verify_op(
- x, packed, conv_state, conv_weight, ids, window, ids, GDN_CONV, p, -1
- )
- assert f8._compiled.cache_info().misses == before
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
deleted file mode 100644
index 9d2a95df0a97..000000000000
--- 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
+++ /dev/null
@@ -1,156 +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_gated_residual entry.
-
-The thing under test is a formula -- which operand is gated, which axis is
-reduced, where each division by `hc_count` sits, the factor on the injection
--- and every mis-reading of it keeps the shapes; the cells' band sits 32x..192x
-inside the nearest wrong formula (measured). Beyond the cells: graph replay and
-one loud misuse.
-"""
-
-import zlib
-from typing import Optional, Tuple
-
-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_gated_residual import ( # noqa: E501
- qwen4_exp_gated_residual as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "qwen4_exp_gated_residual requires a CUDA device"
-
-
-def _compare(
- out: Tuple[torch.Tensor, Optional[torch.Tensor]],
- ref: Tuple[torch.Tensor, Optional[torch.Tensor]],
-) -> None:
- """Each result within `rtol = eps`, `atol = eps * max|ref|` of the
- activation dtype (`gate_ratio <= 1`).
-
- The absolute term is not torch's 1e-5: `mixed` is a mean of signed
- products, so its near-zero elements are cancellations whose error is
- set by the whole stream's rounding, and a relative gate on them gates
- nothing. The relative term is half torch's bf16 default. Measured on
- sm_103 at 512 tokens: the correct result at 0.341 (`mixed`) and 0.404
- (`inject`) of this band (0.198 / 0.210 against the unrounded float64,
- the rest is the reference's own rounding), ten wrong formulas at
- 32x..192x.
- """
- (mixed, inject), (ref_mixed, ref_inject) = out, ref
- assert (inject is None) == (ref_inject is None), "inject must be None exactly when ref's is"
- pairs = [(mixed, ref_mixed)] + ([] if inject is None else [(inject, ref_inject)])
- for got, want in pairs:
- assert got.dtype == want.dtype, (got.dtype, want.dtype)
- assert got.shape == want.shape, (got.shape, want.shape)
- assert got.is_contiguous()
- ratio = gate_ratio(got, want)
- assert ratio <= 1.0, f"{ratio:.4g}x the band"
-
-
-DEV = "cuda"
-HC_COUNT = 4
-HIDDEN = 2560
-WIDTH = HC_COUNT * HIDDEN
-LOWRANK = 320
-RMS_NORM_EPS = 1e-6
-
-# Element statistics of this checkpoint's layer-0 attn_hyper_connection tensors.
-# The normed stream is scaled by (1 + gamma), so gamma's tail -- |6.625| -- is
-# what puts it at ~+-30 rather than ~+-4, and the error floor scales with it.
-CKPT_DOWN_RMS = 0.0296
-CKPT_UP_RMS = 0.0681
-CKPT_INJECT_RMS = 0.0107
-CKPT_GAMMA_RMS = 0.477
-CKPT_GAMMA_ABSMAX = 6.625
-
-
-def _make_weights(seed, dtype=torch.bfloat16, hc_count=HC_COUNT):
- g = torch.Generator(device=DEV).manual_seed(seed)
- width = hc_count * HIDDEN
-
- def normal(*shape, scale):
- return (torch.randn(*shape, generator=g, device=DEV) * scale).to(dtype)
-
- down = normal(LOWRANK, width, scale=CKPT_DOWN_RMS)
- up = normal(width, LOWRANK, scale=CKPT_UP_RMS)
- inject = normal(hc_count, width, scale=CKPT_INJECT_RMS)
- gamma = normal(width, scale=CKPT_GAMMA_RMS)
- # A Gaussian does not reproduce the real gamma's ~14-sigma tail; plant it.
- gamma[0] = CKPT_GAMMA_ABSMAX
- gamma[width // 2] = -CKPT_GAMMA_ABSMAX
- return down, up, inject, gamma
-
-
-def _grouped_gemma_rmsnorm(x, gamma, hc_count):
- """The hc_norm the entry's input comes from -- here only to give the
- operand its real distribution, not under test."""
- hidden = x.shape[-1] // hc_count
- xf = x.float().reshape(*x.shape[:-1], hc_count, hidden)
- xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + RMS_NORM_EPS)
- return (xf.flatten(-2) * (1.0 + gamma.float())).to(x.dtype)
-
-
-def _build(leading, seed, hc_count=HC_COUNT, dtype=torch.bfloat16, inject=True):
- down, up, inject_w, gamma = _make_weights(seed, dtype, hc_count)
- g = torch.Generator(device=DEV).manual_seed(seed + 1)
- raw = torch.randn(*leading, hc_count * HIDDEN, generator=g, device=DEV).to(dtype)
- normed = _grouped_gemma_rmsnorm(raw, gamma, hc_count)
- return raw, (normed, down, up, inject_w if inject else None, hc_count)
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- spec = cell.spec
- hc_count = spec.get("hc_count", HC_COUNT)
- _, args = _build(
- spec["leading"],
- seed=zlib.crc32(cell.why.encode()),
- hc_count=hc_count,
- dtype=spec.get("dtype", torch.bfloat16),
- inject=spec.get("inject", True),
- )
- before = args[0].clone()
- mixed, inject = op(*args)
- assert mixed.shape == (*spec["leading"], HIDDEN)
- if spec.get("inject", True):
- assert inject.shape == (*spec["leading"], hc_count)
- assert torch.equal(args[0], before), "the normed stream is read-only"
- _compare((mixed, inject), op.reference(*args))
-
-
-def test_malformed_operands_raise_rather_than_return() -> None:
- """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)
-
-
-def test_cuda_graph_capture_and_replay() -> None:
- _, (static_in, down, up, inject_w, c) = _build((8,), seed=8)
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(warmup):
- for _ in range(3):
- op(static_in, down, up, inject_w, c)
- torch.cuda.current_stream().wait_stream(warmup)
-
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- graph_mixed, graph_inject = op(static_in, down, up, inject_w, c)
-
- _, (fresh, *_) = _build((8,), seed=18)
- static_in.copy_(fresh)
- graph.replay()
- torch.cuda.synchronize()
- eager_mixed, eager_inject = op(fresh, down, up, inject_w, c)
- assert torch.equal(graph_mixed, eager_mixed)
- assert torch.equal(graph_inject, eager_inject)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_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
deleted file mode 100644
index 2cd74f1c0b3d..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_cute.py
+++ /dev/null
@@ -1,196 +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_cute entry.
-
-Cells against the entry's float64 reference: `hc_out` within one ulp (plus the
-combine's and the finalize's own rounding at the MoE sites), `part` within `2^-16` of
-the reference relative to its scale, `ss` to fp32 accuracy. Beyond the cells: a
-stale-read check under CUDA-graph replay with PDL, where a kernel that read its
-activations before `griddepcontrol.wait` would see the previous replay's; the one-row
-guard.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_first_cute import ( # noqa: E501
- qwen4_exp_hc_decode_first_cute as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "qwen4_exp_hc_decode_first_cute requires a CUDA device"
-
-DEV = "cuda"
-HC_COUNT, HIDDEN, RANK = 4, 2560, 320
-DOWN_RMS, INJECT_RMS, GAMMA_RMS, GAMMA_ABSMAX = 0.0296, 0.0107, 0.477, 6.625
-PART_EPS = 2.0**-16
-
-
-def _band(got, want, eps) -> float:
- w = want.double()
- return ((got.double() - w).abs() / (eps * w.abs().max().item() + eps * w.abs())).max().item()
-
-
-def _ulp(got, want) -> float:
- eps, tiny = torch.finfo(want.dtype).eps, torch.finfo(want.dtype).tiny
- w = want.double()
- return ((got.double() - w).abs() / (eps * w.abs() + tiny)).max().item()
-
-
-def _compare(out, ref) -> None:
- (hc, part, ss), (ref_hc, ref_part, ref_ss) = out, ref
- assert hc.dtype == ref_hc.dtype and hc.shape == ref_hc.shape
- assert part.dtype == torch.float32 and part.shape == ref_part.shape and part.stride(2) == 1
- assert ss.dtype == torch.float32 and ss.shape == ref_ss.shape
- assert _ulp(hc, ref_hc) <= 1.0, f"hc_out {_ulp(hc, ref_hc):.4g} ulp"
- assert _band(part, ref_part, PART_EPS) <= 1.0, f"part {_band(part, ref_part, PART_EPS):.4g}x"
- assert torch.allclose(ss, ref_ss, rtol=1e-5, atol=0), "ss"
-
-
-def _routed(block_out, topk_weights, topk_rows):
- if topk_rows is None:
- return block_out
- valid = (topk_rows >= 0).unsqueeze(-1)
- rows = block_out.double()[topk_rows.clamp(min=0).long()]
- return (
- torch.where(valid, rows * topk_weights.double().unsqueeze(-1), 0.0)
- .sum(1)
- .to(block_out.dtype)
- )
-
-
-def _hc_band_ratio(
- hc, ref_hc, block_out, inject, shared, gate, topk_weights=None, topk_rows=None
-) -> float:
- """The `hc_out` band at the MoE sites: one ulp of `hc_out`, plus one
- ulp of the combine (and with the fold, of the routed row) times `|inject|`."""
- feps, tiny = torch.finfo(ref_hc.dtype).eps, torch.finfo(ref_hc.dtype).tiny
- routed = _routed(block_out, topk_weights, topk_rows)
- s = torch.sigmoid(gate.double().reshape(-1, 1))
- b = (routed.double() + s * shared.double()).to(routed.dtype).double().abs()
- if topk_rows is not None:
- b = b + routed.double().abs()
- band = feps * ref_hc.double().abs() + (
- feps * b.unsqueeze(-2) * inject.double().abs().unsqueeze(-1)
- ).flatten(-2)
- return ((hc.double() - ref_hc.double()).abs() / (band + tiny)).max().item()
-
-
-def _compare_moe(args, out, ref) -> None:
- hyper_input, block_out, inject, gamma, first, c, n, shared, gate = args[:9]
- topk_weights, topk_rows = args[9:] if len(args) > 9 else (None, None)
- (hc, part, ss), (ref_hc, _, _) = out, ref
- ratio = _hc_band_ratio(hc, ref_hc, block_out, inject, shared, gate, topk_weights, topk_rows)
- assert ratio <= 1.0, f"hc_out {ratio:.4g}x the band"
- _compare((hc, part, ss), op.reference(hc, None, None, gamma, first, c, n))
-
-
-def _build(spec, seed):
- """The test operands at one row."""
- rows, hc_count, hidden, dtype = 1, HC_COUNT, HIDDEN, torch.bfloat16
- width = hc_count * hidden
- n_cols = spec.get("n_cols", RANK + hc_count)
- g = torch.Generator(device=DEV).manual_seed(seed)
-
- def normal(*shape, scale):
- return torch.randn(*shape, generator=g, device=DEV) * scale
-
- hyper_input = normal(rows, width, scale=3.0).to(dtype)
- gamma = normal(width, scale=GAMMA_RMS)
- gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX
- gamma = gamma.to(dtype)
- n_rows = -(-n_cols // 128) * 128
- weight = torch.zeros(n_rows, width, device=DEV)
- weight[:n_cols] = normal(n_cols, width, scale=DOWN_RMS)
- weight[n_cols - hc_count : n_cols] *= INJECT_RMS / DOWN_RMS
- first = weight.to(dtype).t()
- block_out = inject = shared = gate_logits = None
- if spec.get("inject", True):
- block_out = normal(rows, hidden, scale=2.0).to(dtype)
- inject = (2.0 * torch.sigmoid(normal(rows, hc_count, scale=1.0))).to(dtype)
- if spec.get("moe", False):
- shared = normal(rows, hidden, scale=2.0).to(dtype)
- gate_logits = normal(rows, 1, scale=2.0).to(dtype)
- top_k = spec.get("fold", 0)
- if not top_k:
- return hyper_input, block_out, inject, gamma, first, hc_count, n_cols, shared, gate_logits
- n_perm = rows * top_k + 37
- block_out = normal(n_perm, hidden, scale=2.0).to(dtype)
- order = torch.randperm(n_perm, generator=g, device=DEV)[: rows * top_k]
- topk_rows = order.reshape(rows, top_k).to(torch.int32)
- if spec.get("drop"):
- topk_rows[:, -3:] = -1
- topk_weights = torch.softmax(normal(rows, top_k, scale=1.0), -1).to(dtype)
- return (
- hyper_input,
- block_out,
- inject,
- gamma,
- first,
- hc_count,
- n_cols,
- shared,
- gate_logits,
- topk_weights,
- topk_rows,
- )
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
- certifying(op)
- args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31))
- before = [a.clone() if isinstance(a, torch.Tensor) else a for a in args]
- out = op(*args)
- for a, b in zip(args, before):
- assert not isinstance(a, torch.Tensor) or torch.equal(a, b), "every input is read-only"
- if args[1] is None:
- assert out[0] is args[0], "norm-only returns hyper_input itself"
- if args[7] is None:
- _compare(out, op.reference(*args))
- else:
- _compare_moe(args, out, op.reference(*args))
-
-
-def test_reads_after_the_wait_under_graph_replay() -> None:
- """A producer kernel rewrites every activation inside the graph right before this op.
- With PDL the op may start before the producer finishes; reading any activation before
- `griddepcontrol.wait` would return the previous replay's value."""
- args = list(_build(dict(moe=True, fold=10), seed=8))
- sources = [a.clone() for a in (args[0], args[1], args[2], args[7], args[8], args[9])]
- targets = (args[0], args[1], args[2], args[7], args[8], args[9])
-
- def step():
- for dst, src in zip(targets, sources):
- dst.copy_(src * 1.0) # an elementwise producer kernel per operand
- return op(*args)
-
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(warmup):
- for _ in range(2):
- step()
- torch.cuda.current_stream().wait_stream(warmup)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- graph_out = step()
- for trial in range(3):
- fresh = _build(dict(moe=True, fold=10), seed=100 + trial)
- for src, new in zip(sources, (fresh[0], fresh[1], fresh[2], fresh[7], fresh[8], fresh[9])):
- src.copy_(new)
- graph.replay()
- torch.cuda.synchronize()
- want = op(*[*args[:10], args[10]])
- for a, b in zip(graph_out, want):
- assert torch.equal(a, b), f"replay {trial}"
-
-
-def test_guards() -> None:
- args = list(_build(dict(), seed=6))
- two = args[0].expand(2, -1).contiguous()
- with validating(op), pytest.raises(AssertionError, match="one row"):
- op(two, args[1].expand(2, -1).contiguous(), args[2].expand(2, -1).contiguous(), *args[3:])
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py
deleted file mode 100644
index 5958e069eb6a..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_cute.py
+++ /dev/null
@@ -1,162 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the two entries of qwen4_exp_hc_decode_mix_fp4_cute: the NVFP4 build
-(`qwen4_exp_hc_decode_mix_fp4_cute`) and the plain build (`qwen4_exp_hc_decode_mix_cute`)
-of the CuTe mix kernel. Each entry's cells against its float64 reference in its band
-(`mixed` and `inject_out` in `rtol = eps`, `atol = eps * max|ref|`; for the NVFP4 build,
-the NVFP4 exact on the returned `mixed` outside the reciprocal tie window), a
-stale-read check under CUDA-graph replay with PDL for the NVFP4 build, the one-row
-guard."""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import (
- Fp4Quantize,
- fp4_quantize,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_mix_fp4_cute import ( # noqa: E501
- qwen4_exp_hc_decode_mix_cute,
- qwen4_exp_hc_decode_mix_fp4_cute,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "qwen4_exp_hc_decode_mix_fp4_cute requires a CUDA device"
-
-DEV = "cuda"
-HC_COUNT, HIDDEN, RANK, EPS = 4, 2560, 320, 1e-6
-UP_RMS, GAMMA_RMS, GAMMA_ABSMAX = 0.0681, 0.477, 6.625
-
-ENTRIES = {"fp4": qwen4_exp_hc_decode_mix_fp4_cute, "plain": qwen4_exp_hc_decode_mix_cute}
-CASES = [(name, cell) for name, entry in ENTRIES.items() for cell in entry.CELLS]
-
-
-def nvfp4_reference(x, gs):
- """`fp4_quantize`'s native-torch NVFP4 of `x` `[T, K]` with linear scales:
- `(packed [T, K/2] uint8, sf [T*K/16] uint8)`."""
- return fp4_quantize.reference(x, gs, 16, False, False)
-
-
-def nvfp4_near_tie(x, gs):
- """`[T, K]` bool: elements within `fp4_quantize`'s tie window of an e2m1 midpoint,
- where the kernels' approximate reciprocals may round either way."""
- return Fp4Quantize.near_tie(fp4_quantize.scaled(x, gs, 16))
-
-
-def _band(got, want) -> float:
- eps = torch.finfo(want.dtype).eps
- w = want.double()
- return ((got.double() - w).abs() / (eps * w.abs().max().item() + eps * w.abs())).max().item()
-
-
-def _compare(out, ref) -> None:
- (mixed, inj), (ref_mixed, ref_inj) = out[:2], ref[:2]
- assert (inj is None) == (ref_inj is None)
- for got, want in [(mixed, ref_mixed)] + ([] if inj is None else [(inj, ref_inj)]):
- assert got.dtype == want.dtype and got.shape == want.shape and got.is_contiguous()
- assert _band(got, want) <= 1.0, f"{_band(got, want):.4g}x the band"
-
-
-def _unpack(xq):
- return torch.stack([xq & 0xF, xq >> 4], -1).flatten(-2)
-
-
-def _check_quant(mixed, xq, xsf, gs) -> None:
- """The NVFP4 outputs exact on the returned `mixed`: scales bitwise, e2m1 codes
- bitwise outside the window where the approximate reciprocals may pick either side."""
- want_q, want_sf = nvfp4_reference(mixed, gs)
- assert xq.dtype == torch.uint8 and xq.shape == want_q.shape and xq.is_contiguous()
- assert xsf.dtype == torch.uint8 and xsf.shape == want_sf.shape
- assert torch.equal(xsf, want_sf), f"{(xsf != want_sf).sum().item()} scale bytes differ"
- stray = (_unpack(xq) != _unpack(want_q)) & ~nvfp4_near_tie(mixed, gs)
- assert not stray.any(), f"{stray.sum().item()} e2m1 codes differ outside the tie window"
-
-
-def _build(spec, seed, fp4=True):
- """The entry's arguments; `fp4` appends the NVFP4 build's global scale."""
- rows, hc_count, hidden, rank, dtype = 1, HC_COUNT, HIDDEN, RANK, torch.bfloat16
- with_inject = spec.get("with_inject", True)
- width = hc_count * hidden
- g = torch.Generator(device=DEV).manual_seed(seed)
-
- def normal(*shape, scale):
- return torch.randn(*shape, generator=g, device=DEV) * scale
-
- hc_out = normal(rows, width, scale=3.0).to(dtype)
- if spec.get("layout") == "row_slice":
- buf = normal(2 * rows, width, scale=1.0).to(dtype)
- buf[::2] = hc_out
- hc_out = buf[::2]
- ss = hc_out.double().unflatten(-1, (hc_count, hidden)).pow(2).sum(-1).t().float().contiguous()
- gamma = normal(width, scale=GAMMA_RMS)
- gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX
- gamma = gamma.to(dtype)
- n_pad = -(-(rank + hc_count) // 4) * 4
- part = normal(hc_count, rows, n_pad + 4, scale=60.0)[:, :, :n_pad]
- up = normal(width, rank, scale=UP_RMS).to(dtype).t()
- args = [part, ss, up, hc_out, gamma, EPS, hc_count, with_inject]
- if fp4:
- # A calibrated activation global scale: 448 * 6 / amax of the mix (~|8|).
- args.append(torch.tensor([448.0 * 6.0 / 8.0], device=DEV, dtype=torch.float32))
- return args
-
-
-@pytest.mark.parametrize(
- "name,cell", CASES, ids=[f"{name}-{cell.why[:44]}" for name, cell in CASES]
-)
-def test_certified_cells(name, cell, certifying) -> None: # noqa: F811 — fixture by name
- op = ENTRIES[name]
- certifying(op)
- fp4 = name == "fp4"
- args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31), fp4=fp4)
- out = op(*args)
- assert out[0].shape == (1, HIDDEN)
- _compare(out, op.reference(*args))
- if fp4:
- _check_quant(out[0], out[2], out[3], args[8])
-
-
-def test_reads_after_the_wait_under_graph_replay() -> None:
- """Producer kernels rewrite `part`, `ss` and `hc_out` inside the graph right before
- the op; with PDL a read before `griddepcontrol.wait` would see the previous replay's.
- The NVFP4 build: the MoE sites' variant, with the most outputs."""
- op = qwen4_exp_hc_decode_mix_fp4_cute
- args = _build(dict(), seed=8)
- targets = (args[0], args[1], args[3])
- sources = [t.clone() for t in targets]
-
- def step():
- for dst, src in zip(targets, sources):
- dst.copy_(src * 1.0)
- return op(*args)
-
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(warmup):
- for _ in range(2):
- step()
- torch.cuda.current_stream().wait_stream(warmup)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- graph_out = step()
- for trial in range(3):
- fresh = _build(dict(), seed=100 + trial)
- for src, new in zip(sources, (fresh[0], fresh[1], fresh[3])):
- src.copy_(new)
- graph.replay()
- torch.cuda.synchronize()
- for a, b in zip(graph_out, op(*args)):
- assert torch.equal(a, b), f"replay {trial}"
-
-
-def test_guards() -> None:
- """Both entries share the shape checks; the NVFP4 build's `is_valid` refuses two rows."""
- op = qwen4_exp_hc_decode_mix_fp4_cute
- args = _build(dict(), seed=6)
- two = args[3].expand(2, -1).contiguous()
- with validating(op), pytest.raises(AssertionError):
- op(args[0].expand(-1, 2, -1), args[1].expand(-1, 2), args[2], two, *args[4:])
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_rows_tc.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_rows_tc.py
deleted file mode 100644
index 9f4f27c1fe15..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_rows_tc.py
+++ /dev/null
@@ -1,173 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test: the 2..16-row hyper-connection site (`qwen4_exp_hc_decode_rows_first` +
-`qwen4_exp_hc_decode_rows_mix`) against the 1-row GEMV kernels applied row by row, over the
-site kinds the forward sends (after a MoE block: combine + injection, with or without the
-runner's finalize folded in; MoE site: injection + NVFP4; norm-only); and the folded site
-bitwise against the runner's finalize (`fp4_block_scale_moe_renorm_permuted.finalize`, the
-runner's arithmetic) followed by the unfolded site. `hc_out` must be bitwise; `mixed` within
-1e-2 of the reference's largest magnitude (1-2 bf16 ulp: the projections take bf16 operands on
-tensor cores); the injection weights within one bf16 ulp; the NVFP4 codes may differ where
-`mixed` does."""
-
-import importlib
-
-import pytest
-import torch
-
-P = "tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc."
-rows_mod = importlib.import_module(P + "qwen4_exp_hc_decode_rows_tc")
-first_op = rows_mod.qwen4_exp_hc_decode_rows_first
-mix_op = rows_mod.qwen4_exp_hc_decode_rows_mix
-first1 = importlib.import_module(
- P + "qwen4_exp_hc_decode_first_cute"
-).qwen4_exp_hc_decode_first_cute
-mix1 = importlib.import_module(P + "qwen4_exp_hc_decode_mix_fp4_cute")
-finalize = importlib.import_module(
- P.replace(".ops.hc.", ".ops.moe.") + "fp4_block_scale_moe_renorm_permuted"
-).finalize
-
-assert torch.cuda.is_available(), "the multi-row hyper-connection site requires a CUDA device"
-
-DEV = "cuda"
-HIDDEN, C, RANK = 2560, 4, 320
-WIDTH = HIDDEN * C
-EPS = 1e-6
-TOPK = 10
-
-
-def _operands(rows, seed):
- g = torch.Generator(device=DEV).manual_seed(seed)
- bf = torch.bfloat16
- rnd = lambda *s, scale=1.0: (torch.randn(*s, generator=g, device=DEV) * scale).to(bf) # noqa: E731
- return dict(
- first=rnd(384, WIDTH, scale=0.01).t(),
- up=rnd(WIDTH, RANK, scale=0.05).t(),
- norm=rnd(WIDTH, scale=0.1),
- hc=rnd(rows, WIDTH, scale=0.5),
- block=rnd(rows, HIDDEN, scale=0.5),
- shared=rnd(rows, HIDDEN, scale=0.5),
- gate=rnd(rows, 1),
- inject=(torch.rand(rows, C, generator=g, device=DEV) * 2).to(bf),
- gscale=torch.tensor([448.0 * 6 / 4.0], device=DEV),
- # The unfinalized runner's permuted FC2 rows (37 rows no slot maps), slot map
- # (the last 3 slots of the odd rows -1) and routing weights.
- perm=rnd(rows * TOPK + 37, HIDDEN, scale=0.5),
- topk_rows=_slot_map(rows, g),
- topk_w=torch.softmax(torch.randn(rows, TOPK, generator=g, device=DEV), -1).to(bf),
- )
-
-
-def _slot_map(rows, g):
- order = torch.randperm(rows * TOPK + 37, generator=g, device=DEV)[: rows * TOPK]
- m = order.reshape(rows, TOPK).to(torch.int32)
- m[1::2, -3:] = -1
- return m.contiguous()
-
-
-def _site(o, site, rows_path: bool):
- """The site through either path: `(hc_out, mixed, inject_out, xq)`."""
- blk, inj = (None, None) if site == "norm" else (o["block"], o["inject"])
- kw = dict(shared=o["shared"], gate_logits=o["gate"]) if site in ("combine", "fold") else {}
- if site == "fold":
- blk = o["perm"]
- kw.update(topk_weights=o["topk_w"], topk_rows=o["topk_rows"])
- n_cols = RANK if site == "norm" else RANK + C
- gs = o["gscale"] if site == "inject" else None
- if rows_path:
- hc, low, inj_out, rstd = first_op(
- o["hc"], blk, inj, o["norm"], o["first"], C, n_cols, EPS, **kw
- )
- mixed, xq, _ = mix_op(low, rstd, o["up"], hc, o["norm"], C, gs)
- return hc, mixed, inj_out, xq
- out = []
- for r in range(o["hc"].shape[0]):
- one = lambda t: None if t is None else t[r : r + 1] # noqa: E731
- h, part, ss = first1(
- one(o["hc"]),
- blk if site == "fold" else one(blk),
- one(inj),
- o["norm"],
- o["first"],
- C,
- n_cols,
- **{k: one(v) for k, v in kw.items()},
- )
- if gs is not None:
- m, i, q, _ = mix1.qwen4_exp_hc_decode_mix_fp4_cute(
- part, ss, o["up"], h, o["norm"], EPS, C, True, gs
- )
- else:
- m, i = mix1.qwen4_exp_hc_decode_mix_cute(
- part, ss, o["up"], h, o["norm"], EPS, C, n_cols > RANK
- )
- q = None
- out.append((h, m, i, q))
- cat = lambda k: None if out[0][k] is None else torch.cat([x[k] for x in out]) # noqa: E731
- return cat(0), cat(1), cat(2), cat(3)
-
-
-def _check(cell):
- rows, site = cell.spec["rows"], cell.spec["site"]
- o = _operands(rows, 1000 + rows)
- got = _site(o, site, True)
- want = _site(o, site, False)
- assert torch.equal(got[0], want[0]), "hc_out bitwise"
- rel = (got[1].float() - want[1].float()).abs().max() / want[1].float().abs().max()
- assert rel <= 1e-2, f"mixed {rel:.2e} of the reference's max"
- if want[2] is not None:
- ulp = (got[2].float() - want[2].float()).abs() / want[2].float().abs().clamp_min(1e-6)
- assert ulp.max() <= 2.0**-7 * 1.01, "injection weights within one bf16 ulp"
- if want[3] is not None:
- assert (got[3] != want[3]).float().mean() < 0.01, "NVFP4 codes"
-
-
-@pytest.mark.parametrize("cell", first_op.CELLS, ids=[c.why[:44] for c in first_op.CELLS])
-def test_first_cells(cell) -> None:
- _check(cell)
-
-
-@pytest.mark.parametrize("cell", mix_op.CELLS, ids=[c.why[:44] for c in mix_op.CELLS])
-def test_mix_cells(cell) -> None:
- _check(cell)
-
-
-@pytest.mark.parametrize("rows", [2, 16])
-@pytest.mark.parametrize("n_cols", [RANK + C, RANK])
-def test_fold_bitwise_vs_runner_finalize(rows, n_cols) -> None:
- """Every output of the folded site equals the unfolded site fed the runner's finalize."""
- o = _operands(rows, 300 + rows)
- kw = dict(shared=o["shared"], gate_logits=o["gate"])
- got = first_op(
- o["hc"],
- o["perm"],
- o["inject"],
- o["norm"],
- o["first"],
- C,
- n_cols,
- EPS,
- topk_weights=o["topk_w"],
- topk_rows=o["topk_rows"],
- **kw,
- )
- routed = finalize(o["perm"], o["topk_w"], o["topk_rows"])
- want = first_op(o["hc"], routed, o["inject"], o["norm"], o["first"], C, n_cols, EPS, **kw)
- for a, b in zip(got, want):
- assert (a is None and b is None) or torch.equal(a, b)
-
-
-def test_replay_matches_eager() -> None:
- o = _operands(8, 77)
- _site(o, "inject", True)
- torch.cuda.synchronize()
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- captured = _site(o, "inject", True)
- fresh = _operands(8, 78)
- o["hc"].copy_(fresh["hc"])
- o["block"].copy_(fresh["block"])
- graph.replay()
- torch.cuda.synchronize()
- eager = _site(o, "inject", True)
- assert torch.equal(captured[1], eager[1]) and torch.equal(captured[0], eager[0])
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_first_rescale_tc.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_first_rescale_tc.py
deleted file mode 100644
index 1645a331fa44..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_first_rescale_tc.py
+++ /dev/null
@@ -1,122 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the qwen4_exp_hc_first_rescale_tc entry: `low` and `inject` in
-the gate band (`qwen4_exp_gated_residual`'s) of the float64 projection of the
-once-rounded normed stream, the stack-final form, the weight-layout guard,
-graph capture.
-"""
-
-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_first_rescale_tc import ( # noqa: E501
- qwen4_exp_hc_first_rescale_tc 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_first_rescale_tc requires a CUDA device"
-
-DEV = "cuda"
-HC_COUNT = 4
-HIDDEN = 2560
-RANK = 320
-EPS = 1e-6
-#: The padded first weight's rows (`_fuse_hc_first_projections`: 324 rounded up to 128).
-FIRST_ROWS = 384
-# This checkpoint's grouped-norm weight statistics (layer-0 hyper-connection norm).
-GAMMA_RMS = 0.477
-GAMMA_ABSMAX = 6.625
-
-
-def _compare(out, ref) -> None:
- """gate_ratio <= 1 for `low` and `inject`: the band the mixer's outputs are
- certified in (`rtol = eps`, `atol = eps * max|ref|`); the fp32 logits are a
- 10240-term tensor-core sum against the float64 one.
- """
- 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 gate_ratio(got, want) <= 1.0, f"{gate_ratio(got, want):.4g}x the band"
-
-
-def _build(spec, seed):
- rows, width, dtype = spec["rows"], HC_COUNT * HIDDEN, torch.bfloat16
- g = torch.Generator(device=DEV).manual_seed(seed)
- hc_out = (torch.randn(rows, width, generator=g, device=DEV) * 3.0).to(dtype)
- x = hc_out.double().unflatten(-1, (HC_COUNT, HIDDEN))
- rstd = torch.rsqrt(x.pow(2).mean(-1) + EPS).float().contiguous()
- gamma = torch.randn(width, generator=g, device=DEV) * GAMMA_RMS
- gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX
- # `[down; block_inject; 0]`, logits up to about +-60 (this checkpoint's reach |67|).
- first = torch.zeros(FIRST_ROWS, width, device=DEV)
- first[: RANK + HC_COUNT] = torch.randn(RANK + HC_COUNT, width, generator=g, device=DEV) * 0.15
- return (
- hc_out,
- rstd,
- gamma.to(dtype),
- first.to(dtype).t(),
- HC_COUNT,
- RANK,
- bool(spec.get("inject", True)),
- )
-
-
-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))
- low, inject = op(*args)
- assert low.shape == (cell.spec["rows"], RANK)
- assert (inject is None) == (not args[6])
- _compare((low, inject), op.reference(*args))
-
-
-def test_stack_final_low_is_the_injecting_form_bitwise() -> None:
- args = _build(dict(rows=300), seed=7)
- low, inject = op(*args[:6], True)
- low_f, none = op(*args[:6], False)
- assert none is None and inject.shape == (300, HC_COUNT)
- assert torch.equal(low, low_f)
-
-
-def test_guard_refuses_a_contiguous_first() -> None:
- """The kernel streams the `[N, C*H]` weight row-major: `first` must be its
- transposed view, not a `[C*H, N]` copy.
- """
- args = list(_build(dict(rows=4), seed=4))
- args[3] = args[3].contiguous()
- with validating(op), pytest.raises(AssertionError, match="transposed"):
- op(*args)
-
-
-def test_cuda_graph_capture_and_replay() -> None:
- args = _build(dict(rows=8), seed=5)
- side = torch.cuda.Stream()
- side.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(side):
- op(*args)
- torch.cuda.current_stream().wait_stream(side)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- g_low, g_inject = op(*args)
- fresh = _build(dict(rows=8), seed=6)
- args[0].copy_(fresh[0])
- args[1].copy_(fresh[1])
- graph.replay()
- torch.cuda.synchronize()
- e_low, e_inject = op(*args)
- assert torch.equal(g_low, e_low) and torch.equal(g_inject, e_inject)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py
deleted file mode 100644
index b69e82f5c120..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_inject_norm.py
+++ /dev/null
@@ -1,202 +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_inject_norm entry.
-
-Cells are gated at one ulp of the activation dtype against the float64
-reference; the MoE-combine cells widen the residual's band by the combine's own
-rounding (see `_compare_moe`). Beyond them: the CuTe fold kernel bitwise the
-Triton kernel at the last row count before the switch, and the slot-map guard.
-"""
-
-import sys
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_inject_norm import ( # noqa: E501
- qwen4_exp_hc_inject_norm as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "qwen4_exp_hc_inject_norm requires a CUDA device"
-
-#: The op's module (its kernels and dispatch guards).
-M = sys.modules[type(op).__module__]
-
-DEV = "cuda"
-HC_COUNT = 4
-HIDDEN = 2560
-EPS = 1e-6
-# This checkpoint's grouped-norm weight statistics (layer-0 hyper-connection norm).
-GAMMA_RMS = 0.477
-GAMMA_ABSMAX = 6.625
-
-
-def _ulp_ratio(out: torch.Tensor, ref: torch.Tensor) -> float:
- """Worst element's distance from `ref` in units of `eps(dtype) * |ref|`
- (one ulp of the reference value, or a little more). 1 means one ulp."""
- eps = torch.finfo(ref.dtype).eps
- ref64 = ref.double()
- tiny = torch.finfo(ref.dtype).tiny
- return ((out.double() - ref64).abs() / (eps * ref64.abs() + tiny)).max().item()
-
-
-def _compare(out, ref) -> None:
- """Both outputs within one ulp of the float64 result rounded once.
-
- The op rounds each output once from an fp32 value whose own error is
- ~1e-6 relative (a 16-bit product is exact in fp32; the RMS is a 2560-term
- fp32 sum), so it can differ from the once-rounded float64 value only where
- the two straddle a rounding midpoint: by exactly one ulp."""
- (hc, normed), (ref_hc, ref_normed) = out, ref
- for got, want in ((hc, ref_hc), (normed, ref_normed)):
- assert got.dtype == want.dtype and got.shape == want.shape
- if got.numel():
- ratio = _ulp_ratio(got, want)
- assert ratio <= 1.0, f"{ratio:.4g} ulp"
-
-
-def _combine(block_out, shared, gate_logits) -> torch.Tensor:
- """The MoE block output in float64, rounded once to `block_out`'s dtype."""
- s = torch.sigmoid(gate_logits.double().reshape(-1, 1))
- return (block_out.double() + s * shared.double()).to(block_out.dtype)
-
-
-def _routed(block_out, topk_weights, topk_rows):
- """The routed rows: `block_out` itself, or with the fold the float64 finalize of
- the permuted rows rounded once."""
- if topk_rows is None:
- return block_out
- valid = (topk_rows >= 0).unsqueeze(-1)
- picked = block_out.double()[topk_rows.clamp(min=0).long()]
- w = topk_weights.double().unsqueeze(-1)
- return torch.where(valid, picked * w, 0.0).sum(1).to(block_out.dtype)
-
-
-def _compare_moe(args, out, ref) -> None:
- """MoE-combine cells.
-
- `hc_out`: the op's combine is an fp32 value rounded once, the reference's a
- float64 value rounded once; where the two straddle a rounding midpoint they
- differ by one ulp of the block output, and the accumulate carries that
- scaled by `inject`. So the band is one ulp of `hc_out` plus one ulp of the
- combine times `|inject|`, element by element -- still well under one ulp
- of `hc_out` wherever the block term is small against the residual. With the
- finalize fold the routed row has its own rounding (fp32 FMAs against the
- float64 sum), one more ulp of the routed row times `|inject|`.
-
- `normed`: one ulp of the float64 norm of the residual the op stored,
- `reference(hc_out, None, None, ...)` -- the dataflow the op states."""
- hyper_input, block_out, inject, gamma, eps, c, shared, gate = args[:8]
- topk_weights, topk_rows = args[8:] if len(args) > 8 else (None, None)
- (hc, normed), (ref_hc, _) = out, ref
- assert hc.dtype == ref_hc.dtype and hc.shape == ref_hc.shape
- if not hc.numel():
- return
- feps, tiny = torch.finfo(ref_hc.dtype).eps, torch.finfo(ref_hc.dtype).tiny
- routed = _routed(block_out, topk_weights, topk_rows)
- b = _combine(routed, shared, gate).double().abs()
- if topk_rows is not None:
- b = b + routed.double().abs()
- band = feps * ref_hc.double().abs() + (
- feps * b.unsqueeze(-2) * inject.double().abs().unsqueeze(-1)
- ).flatten(-2)
- ratio = ((hc.double() - ref_hc.double()).abs() / (band + tiny)).max().item()
- assert ratio <= 1.0, f"hc_out {ratio:.4g}x the band"
- own = op.reference(hc, None, None, gamma, eps, c)[1]
- assert _ulp_ratio(normed, own) <= 1.0, f"normed {_ulp_ratio(normed, own):.4g} ulp"
-
-
-def _build(spec, seed):
- rows, hc_count, hidden, dtype = spec["rows"], HC_COUNT, HIDDEN, torch.bfloat16
- width = hc_count * hidden
- g = torch.Generator(device=DEV).manual_seed(seed)
- hyper_input = (torch.randn(rows, width, generator=g, device=DEV) * 3.0).to(dtype)
- gamma = torch.randn(width, generator=g, device=DEV) * GAMMA_RMS
- gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX
- gamma = gamma.to(dtype)
- block_out = inject = shared = gate_logits = None
- if spec.get("inject", True):
- block_out = (torch.randn(rows, hidden, generator=g, device=DEV) * 2.0).to(dtype)
- inject = (2.0 * torch.sigmoid(torch.randn(rows, hc_count, generator=g, device=DEV))).to(
- dtype
- )
- if spec.get("moe", False):
- shared = (torch.randn(rows, hidden, generator=g, device=DEV) * 2.0).to(dtype)
- gate_logits = (torch.randn(rows, 1, generator=g, device=DEV) * 2.0).to(dtype)
- top_k = spec.get("fold", 0)
- if not top_k:
- return hyper_input, block_out, inject, gamma, EPS, hc_count, shared, gate_logits
- # The runner's unfinalized output: permuted FC2 rows (more rows than slots, as
- # its padded buffer has), each slot's row and its routing weight.
- n_perm = rows * top_k + 37
- block_out = (torch.randn(n_perm, hidden, generator=g, device=DEV) * 2.0).to(dtype)
- order = torch.randperm(n_perm, generator=g, device=DEV)[: rows * top_k]
- topk_rows = order.reshape(rows, top_k).to(torch.int32)
- if spec.get("drop"):
- topk_rows[:, -3:] = -1
- topk_weights = torch.softmax(torch.randn(rows, top_k, generator=g, device=DEV), -1).to(dtype)
- return (
- hyper_input,
- block_out,
- inject,
- gamma,
- EPS,
- hc_count,
- shared,
- gate_logits,
- topk_weights,
- topk_rows,
- )
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
- certifying(op)
- args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
- before = [a.clone() if isinstance(a, torch.Tensor) else a for a in args]
- hc, normed = op(*args)
- for a, b in zip(args, before):
- assert not isinstance(a, torch.Tensor) or torch.equal(a, b), "every input is read-only"
- if args[1] is None:
- assert hc is args[0], "norm-only returns hyper_input itself"
- else:
- assert hc.is_contiguous()
- assert args[0].numel() == 0 or hc.data_ptr() != args[0].data_ptr(), "hc_out is fresh"
- assert normed.is_contiguous()
- if args[6] is None:
- _compare((hc, normed), op.reference(*args))
- else:
- _compare_moe(args, (hc, normed), op.reference(*args))
-
-
-def test_guard_refuses_an_int64_slot_map() -> None:
- """The fold's slot map is the runner's int32 index map; `is_valid` refuses an int64 one."""
- fold = _build(dict(rows=4, moe=True, fold=10), seed=9)
- with validating(op), pytest.raises(AssertionError, match="int32"):
- op(*fold[:9], fold[9].long())
-
-
-def test_fold_cute_is_the_triton_kernel_bitwise(monkeypatch) -> None:
- """Below `_FEW_WARPS_MIN_ROWS` rows the fold runs the CuTe DSL kernel; its
- `hc_out` and `normed` are bitwise the Triton kernel's (run by refusing the
- CuTe guard) at the last row count before the switch."""
- rows = M._FEW_WARPS_MIN_ROWS - 1
- args = _build(dict(rows=rows, moe=True, fold=10), seed=500 + rows)
- calls = []
- cute_fold = M._fold_cute
-
- def recording(*a):
- calls.append(a[0].shape[0])
- cute_fold(*a)
-
- monkeypatch.setattr(M, "_fold_cute", recording)
- hc, normed = op(*args)
- assert calls == [rows], "the decode MoE site runs the CuTe kernel"
- monkeypatch.setattr(M, "_fold_cute_ok", lambda *a: False)
- want_hc, want_normed = op(*args)
- assert torch.equal(hc, want_hc)
- assert torch.equal(normed, want_normed)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_rstd.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_rstd.py
deleted file mode 100644
index 20558035d0b0..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_inject_rstd.py
+++ /dev/null
@@ -1,165 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the qwen4_exp_hc_inject_rstd entry.
-
-Cells: `hc_out` within one ulp of the float64 reference (the MoE-combine cells
-widened by the combine's own rounding, as `qwen4_exp_hc_inject_norm`'s test),
-rstd within 2^-16 relative of the float64 statistic of the stored residual.
-Beyond them: `hc_out` and the normed stream rebuilt from rstd bitwise
-`qwen4_exp_hc_inject_norm`'s at the cells' sizes, the operand guard, graph
-capture.
-"""
-
-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.hc.qwen4_exp_hc_inject_rstd import ( # noqa: E501
- qwen4_exp_hc_inject_rstd as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "qwen4_exp_hc_inject_rstd requires a CUDA device"
-
-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
-#: rstd band, relative: an fp32 2560-term sum of squares and an fp32 rsqrt.
-RSTD_RTOL = 2.0**-16
-
-
-def _ulp_ratio(out: torch.Tensor, ref: torch.Tensor) -> float:
- """Worst element's distance from `ref` in units of `eps(dtype) * |ref|`."""
- 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 _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 _compare(args, out, ref) -> None:
- """`hc_out`: one ulp of the once-rounded float64 accumulate; with the MoE combine,
- plus one ulp of the combine times `|inject|` (the combine is an fp32 value
- rounded once, the reference's a float64 one). rstd: `RSTD_RTOL` of the float64
- statistic of the residual the op stored -- the dataflow the op states.
- """
- hyper_input, block_out, inject, eps, c, shared, gate = args
- (hc, rstd), (ref_hc, ref_rstd) = out, ref
- assert hc.dtype == ref_hc.dtype and hc.shape == ref_hc.shape
- assert rstd.dtype == torch.float32 and rstd.shape == ref_rstd.shape and rstd.is_contiguous()
- if not hc.numel():
- return
- if shared is None:
- assert _ulp_ratio(hc, ref_hc) <= 1.0, f"hc_out {_ulp_ratio(hc, ref_hc):.4g} ulp"
- else:
- feps, tiny = torch.finfo(ref_hc.dtype).eps, torch.finfo(ref_hc.dtype).tiny
- b = _combine(block_out, shared, gate).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, eps, c)[1].double()
- rel = ((rstd.double() - own).abs() / own).max().item()
- assert rel <= RSTD_RTOL, f"rstd {rel:.3g} relative"
-
-
-def _gamma(seed):
- g = torch.Generator(device=DEV).manual_seed(seed)
- width = HC_COUNT * HIDDEN
- gamma = torch.randn(width, generator=g, device=DEV) * GAMMA_RMS
- gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX
- return gamma.to(torch.bfloat16)
-
-
-def _build(spec, seed):
- rows, dtype = spec["rows"], 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)
- 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)
- return hyper_input, block_out, inject, EPS, HC_COUNT, shared, gate_logits
-
-
-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))
- before = [a.clone() if isinstance(a, torch.Tensor) else a for a in args]
- hc, rstd = 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], "statistic-only returns hyper_input itself"
- else:
- assert hc.is_contiguous() and hc.data_ptr() != args[0].data_ptr(), "hc_out is fresh"
- _compare(args, (hc, rstd), op.reference(*args))
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_normed_rebuilt_from_rstd_is_inject_norm_bitwise(cell) -> None:
- """The consumers' operand `bf16((hc_out * rstd) * (1 + w))` (fp32 products) is
- bitwise `qwen4_exp_hc_inject_norm`'s `normed`, and `hc_out` its `hc_out`.
- """
- hyper_input, block_out, inject, eps, c, shared, gate = _build(cell.spec, seed=_seed(cell))
- gamma = _gamma(_seed(cell) + 1)
- hc, rstd = op(hyper_input, block_out, inject, eps, c, shared, gate)
- want_hc, want_normed = qwen4_exp_hc_inject_norm(
- hyper_input, block_out, inject, gamma, eps, c, shared, gate
- )
- assert torch.equal(hc, want_hc), f"{(hc != want_hc).sum().item()} hc_out elements differ"
- x = hc.float().unflatten(-1, (c, -1)) * rstd.unsqueeze(-1)
- normed = (x.flatten(-2) * (1.0 + gamma.float())).to(hc.dtype)
- diff = (normed != want_normed).sum().item()
- assert diff == 0, f"{diff} normed elements differ"
-
-
-def test_guard_refuses_block_out_without_inject() -> None:
- hyper_input, block_out, _, eps, c, _, _ = _build(dict(rows=4), seed=4)
- with validating(op), pytest.raises(AssertionError, match="together"):
- op(hyper_input, block_out, None, eps, c)
-
-
-def test_cuda_graph_capture_and_replay() -> None:
- args = list(_build(dict(rows=8, moe=True), seed=5))
- side = torch.cuda.Stream()
- side.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(side):
- op(*args)
- torch.cuda.current_stream().wait_stream(side)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- g_hc, g_rstd = op(*args)
- fresh = _build(dict(rows=8, moe=True), seed=6)
- for i in (0, 1, 2, 5, 6):
- args[i].copy_(fresh[i])
- graph.replay()
- torch.cuda.synchronize()
- e_hc, e_rstd = op(*fresh)
- assert torch.equal(g_hc, e_hc) and torch.equal(g_rstd, e_rstd)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py
deleted file mode 100644
index 8641b2b8b7b5..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_lowrank_act.py
+++ /dev/null
@@ -1,79 +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_lowrank_act entry: cells within one ulp of the
-float64 reference, the fp32-logits guard, graph capture."""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_lowrank_act import ( # noqa: E501
- qwen4_exp_hc_lowrank_act as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "qwen4_exp_hc_lowrank_act requires a CUDA device"
-
-DEV = "cuda"
-RANK = 320
-
-
-def _ulp_ratio(out, ref) -> float:
- eps = torch.finfo(ref.dtype).eps
- tiny = torch.finfo(ref.dtype).tiny
- return ((out.double() - ref.double()).abs() / (eps * ref.double().abs() + tiny)).max().item()
-
-
-def _compare(out, ref) -> None:
- """One ulp: each output is one fp32 activation of an fp32 input, rounded
- once, so it can differ from the once-rounded float64 value only at a
- rounding midpoint."""
- for got, want in zip(out, ref):
- assert (got is None) == (want is None)
- if got is None:
- continue
- assert got.dtype == want.dtype and got.shape == want.shape and got.is_contiguous()
- assert _ulp_ratio(got, want) <= 1.0, f"{_ulp_ratio(got, want):.4g} ulp"
-
-
-def _build(spec, seed):
- hc_count = spec.get("hc_count", 4)
- g = torch.Generator(device=DEV).manual_seed(seed)
- # The first projection's logits on this checkpoint reach |67| (down) / |48| (inject).
- logits = torch.randn(spec["rows"], spec.get("logits_width", 384), generator=g, device=DEV) * 20
- return logits, RANK, hc_count, spec.get("inject", True), spec.get("dtype", torch.bfloat16)
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
- certifying(op)
- args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31))
- low, inject = op(*args)
- assert low.shape == (cell.spec["rows"], RANK)
- _compare((low, inject), op.reference(*args))
-
-
-def test_guard_refuses_bf16_logits() -> None:
- logits, rank, c, wi, dt = _build(dict(rows=4), seed=4)
- with validating(op), pytest.raises(AssertionError):
- op(logits.to(torch.bfloat16), rank, c, wi, dt)
-
-
-def test_cuda_graph_capture_and_replay() -> None:
- logits, rank, c, wi, dt = _build(dict(rows=8), seed=5)
- side = torch.cuda.Stream()
- side.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(side):
- op(logits, rank, c, wi, dt)
- torch.cuda.current_stream().wait_stream(side)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- g_low, g_inject = op(logits, rank, c, wi, dt)
- logits.copy_(_build(dict(rows=8), seed=6)[0])
- graph.replay()
- torch.cuda.synchronize()
- e_low, e_inject = op(logits, rank, c, wi, dt)
- assert torch.equal(g_low, e_low) and torch.equal(g_inject, e_inject)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py
deleted file mode 100644
index 3e392c8e6f75..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_up_gate_mean.py
+++ /dev/null
@@ -1,78 +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_up_gate_mean entry: cells in the band
-`qwen4_exp_gated_residual`'s test derives for `mixed`, the guard, graph
-capture."""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501
- gate_ratio,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_up_gate_mean import ( # noqa: E501
- qwen4_exp_hc_up_gate_mean as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "qwen4_exp_hc_up_gate_mean requires a CUDA device"
-
-DEV = "cuda"
-
-
-def _compare(out, ref) -> None:
- """`rtol = eps`, `atol = eps * max|ref|` (`gate_ratio <= 1`), the band
- `qwen4_exp_gated_residual` derives for `mixed`."""
- assert out.dtype == ref.dtype and out.shape == ref.shape and out.is_contiguous()
- assert gate_ratio(out, ref) <= 1.0, f"{gate_ratio(out, ref):.4g}x the band"
-
-
-def _build(spec, seed):
- hc, hidden, rank = spec.get("hc_count", 4), spec.get("hidden", 2560), spec.get("rank", 320)
- g = torch.Generator(device=DEV).manual_seed(seed)
- rows, width = spec["rows"], hc * hidden
- # This checkpoint's statistics: low = silu(.) of the bottleneck (|x| < ~3), the up
- # weight ~N(0, 0.068) (gate logits within about +-8), the normed stream ~+-30.
- low = torch.nn.functional.silu(torch.randn(rows, rank, generator=g, device=DEV) * 1.5).to(
- torch.bfloat16
- )
- up_w = (torch.randn(width, rank, generator=g, device=DEV) * 0.068).to(torch.bfloat16)
- normed = (torch.randn(rows, width, generator=g, device=DEV) * 6).to(torch.bfloat16)
- return low, up_w.t(), normed, hc
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
- certifying(op)
- args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31))
- mixed = op(*args)
- assert mixed.shape == (cell.spec["rows"], cell.spec.get("hidden", 2560))
- _compare(mixed, op.reference(*args))
-
-
-def test_guard_refuses_a_contiguous_up() -> None:
- low, up, normed, c = _build(dict(rows=4), seed=4)
- with validating(op), pytest.raises(AssertionError):
- op(low, up.contiguous(), normed, c)
-
-
-def test_cuda_graph_capture_and_replay() -> None:
- low, up, normed, c = _build(dict(rows=8), seed=5)
- side = torch.cuda.Stream()
- side.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(side):
- op(low, up, normed, c)
- torch.cuda.current_stream().wait_stream(side)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- g_mixed = op(low, up, normed, c)
- fresh = _build(dict(rows=8), seed=6)
- low.copy_(fresh[0])
- normed.copy_(fresh[2])
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(g_mixed, op(fresh[0], up, fresh[2], c))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_rstd.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_rstd.py
deleted file mode 100644
index 3f4c2f41f174..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_up_gate_mean_rstd.py
+++ /dev/null
@@ -1,175 +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_up_gate_mean_rstd entries (bf16 and NVFP4):
-`mixed` in the gate band of the float64 reference, the epilogue's fp32 operation
-sequence bitwise (gate at zero logits), the NVFP4 outputs against the native-torch
-quantization of the returned `mixed` outside the reciprocal tie window, the NVFP4
-build's `mixed` bitwise the bf16 build's, the global-scale guard, graph capture.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import (
- Fp4Quantize,
- fp4_quantize,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501
- gate_ratio,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_up_gate_mean_rstd import ( # noqa: E501
- qwen4_exp_hc_up_gate_mean_rstd as op,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_up_gate_mean_rstd import ( # noqa: E501
- qwen4_exp_hc_up_gate_mean_rstd_fp4 as op4,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "qwen4_exp_hc_up_gate_mean_rstd requires a CUDA device"
-
-DEV = "cuda"
-HC_COUNT = 4
-HIDDEN = 2560
-RANK = 320
-EPS = 1e-6
-GAMMA_RMS = 0.477
-GAMMA_ABSMAX = 6.625
-
-
-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 = fp4_quantize.reference(mixed, gs, 16, False, False)
- assert xq.dtype == torch.uint8 and xq.shape == want_q.shape and xq.is_contiguous()
- assert xsf.dtype == torch.uint8 and xsf.shape == want_sf.shape
- assert torch.equal(xsf, want_sf), f"{(xsf != want_sf).sum().item()} scale bytes differ"
- diff = _unpack(xq) != _unpack(want_q)
- stray = diff & ~Fp4Quantize.near_tie(fp4_quantize.scaled(mixed, gs, 16))
- assert not stray.any(), f"{stray.sum().item()} e2m1 codes differ outside the tie window"
-
-
-def _check_mixed(mixed, ref_mixed) -> None:
- """gate_ratio <= 1, the band `qwen4_exp_hc_up_gate_mean` is certified in."""
- assert mixed.dtype == ref_mixed.dtype and mixed.shape == ref_mixed.shape
- assert mixed.is_contiguous()
- assert gate_ratio(mixed, ref_mixed) <= 1.0, f"{gate_ratio(mixed, ref_mixed):.4g}x the band"
-
-
-def _build(spec, seed):
- """`(low, up, hc_out, rstd, norm_w1, hc_count, global_scale)`."""
- rows, width = spec["rows"], HC_COUNT * HIDDEN
- g = torch.Generator(device=DEV).manual_seed(seed)
- # This checkpoint's statistics: low = silu(.) of the bottleneck (|x| < ~3), the up
- # weight ~N(0, 0.068) (gate logits within about +-8).
- 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)
- hc_out = (torch.randn(rows, width, generator=g, device=DEV) * 3.0).to(torch.bfloat16)
- x = hc_out.double().unflatten(-1, (HC_COUNT, HIDDEN))
- rstd = torch.rsqrt(x.pow(2).mean(-1) + EPS).float().contiguous()
- gamma = torch.randn(width, generator=g, device=DEV) * GAMMA_RMS
- gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX
- norm_w1 = (1.0 + gamma.to(torch.bfloat16).float()).contiguous()
- # A calibrated activation global scale: 448 * 6 / amax over a typical batch.
- gs = torch.tensor([448.0 * 6.0 / 4.0], device=DEV, dtype=torch.float32)
- return low, up_w.t(), hc_out, rstd, norm_w1, HC_COUNT, 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))[:6]
- mixed = op(*args)
- assert mixed.shape == (cell.spec["rows"], HIDDEN)
- _check_mixed(mixed, op.reference(*args))
-
-
-@pytest.mark.parametrize("cell", op4.CELLS, ids=[c.why[:44] for c in op4.CELLS])
-def test_certified_cells_fp4(cell, certifying) -> None: # noqa: F811 — fixture by name
- certifying(op4)
- args = _build(cell.spec, seed=_seed(cell))
- mixed, xq, xsf = op4(*args)
- assert mixed.shape == (cell.spec["rows"], HIDDEN)
- ref_mixed = op4.reference(*args)[0]
- _check_mixed(mixed, ref_mixed)
- _check_quant(mixed, xq, xsf, args[6])
-
-
-@pytest.mark.parametrize("rows", [5, 300, 4100])
-def test_epilogue_operation_sequence_bitwise(rows) -> None:
- """With `low = 0` every logit is 0 and `tanh.approx(0) = 0`, so the gate is the
- folded `rstd / 2C` and `mixed` is the documented fp32 sequence alone:
- `acc = (rstd_0 / 8) * (x_0 * w1_0)`, then `acc = fma(rstd_s / 8, x_s * w1_s, acc)`
- in stream order, one bf16 rounding. Each step is emulated in float64 (products
- of the fp32 operands are exact there) and rounded to fp32.
- """
- low, up, hc_out, rstd, w1, c, _ = _build(dict(rows=rows), seed=200 + rows)
- low = torch.zeros_like(low)
- mixed = op(low, up, hc_out, rstd, w1, c)
- x = hc_out.double().unflatten(-1, (c, HIDDEN))
- xw = (x * w1.double().unflatten(-1, (c, HIDDEN))).float().double()
- hr = (rstd * (0.5 / c)).double()
- acc = (hr[:, :1] * xw[:, 0]).float().double()
- for s in range(1, c):
- acc = (hr[:, s : s + 1] * xw[:, s] + acc).float().double()
- want = acc.float().to(torch.bfloat16)
- assert torch.equal(mixed, want), f"{(mixed != want).sum().item()} elements differ"
-
-
-@pytest.mark.parametrize("rows", [5, 2048, 4100])
-def test_fp4_mixed_is_the_bf16_entry_bitwise(rows) -> None:
- args = _build(dict(rows=rows), seed=100 + rows)
- assert torch.equal(op4(*args)[0], op(*args[:6]))
-
-
-def test_guards() -> None:
- """The kernel loads `global_scale[0]` only; a per-row scale is refused. rstd must be
- fp32.
- """
- low, up, hc_out, rstd, w1, c, gs = _build(dict(rows=4), seed=4)
- with validating(op4), pytest.raises(AssertionError):
- op4(low, up, hc_out, rstd, w1, c, torch.cat([gs, gs]))
- with validating(op), pytest.raises(AssertionError, match="rstd"):
- op(low, up, hc_out, rstd.to(torch.bfloat16), w1, c)
-
-
-@pytest.mark.parametrize("fp4", [False, True])
-def test_cuda_graph_capture_and_replay(fp4) -> None:
- args = list(_build(dict(rows=8), seed=5))
- if not fp4:
- args = args[:6]
- entry = op4 if fp4 else op
- side = torch.cuda.Stream()
- side.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(side):
- entry(*args)
- torch.cuda.current_stream().wait_stream(side)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- g_out = entry(*args)
- fresh = _build(dict(rows=8), seed=6)
- for i in (0, 2, 3):
- args[i].copy_(fresh[i])
- if fp4:
- args[6].mul_(0.5)
- graph.replay()
- torch.cuda.synchronize()
- want = entry(*args)
- g_out, want = (g_out, want) if fp4 else ((g_out,), (want,))
- for got, exp in zip(g_out, want):
- assert torch.equal(got, exp)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_post_inject.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_post_inject.py
deleted file mode 100644
index e609e8632440..000000000000
--- 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
+++ /dev/null
@@ -1,106 +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_post_inject entry.
-
-The cells are gated bitwise (16-bit) by `_compare`. Beyond them: the
-leading-dim guard and graph capture.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_post_inject import ( # noqa: E501
- qwen4_exp_post_inject as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "qwen4_exp_post_inject requires a CUDA device"
-
-
-def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
- """Bitwise for 16-bit operands: the product of two bf16 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()
- assert torch.equal(out, ref), (out.float() - ref.float()).abs().max().item()
-
-
-DEV = "cuda"
-HC_COUNT = 4
-HIDDEN = 2560
-
-
-def _operands(leading, hc_count=HC_COUNT, hidden=HIDDEN, dtype=torch.bfloat16, seed=0):
- """Operands in the ranges the model produces: `inject` is 2*sigmoid(.), so [0, 2]."""
- g = torch.Generator(device=DEV).manual_seed(seed)
- hyper_input = torch.randn(*leading, hc_count * hidden, generator=g, device=DEV).to(dtype)
- block_out = torch.randn(*leading, hidden, generator=g, device=DEV).to(dtype)
- inject = (2.0 * torch.sigmoid(torch.randn(*leading, hc_count, generator=g, device=DEV))).to(
- dtype
- )
- return hyper_input, block_out, inject
-
-
-def _build(spec, seed):
- hc_count, hidden = spec["hc_count"], spec["hidden"]
- dtype = spec.get("dtype", torch.bfloat16)
- hyper_input, block_out, inject = _operands(spec["leading"], hc_count, hidden, dtype, seed)
- layout = spec.get("layout", "contiguous")
- width = hc_count * hidden
- if layout == "col_slice":
- buf = torch.randn(hyper_input.shape[0], 2 * width, device=DEV).to(dtype)
- buf[:, :width] = hyper_input
- hyper_input = buf[:, :width]
- if layout != "contiguous":
- assert not hyper_input.is_contiguous() and hyper_input.stride(-1) == 1
- return hyper_input, block_out, inject
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- hyper_input, block_out, inject = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
- before = hyper_input.clone()
- out = op(hyper_input, block_out, inject)
- assert out.shape == hyper_input.shape and out.dtype == hyper_input.dtype
- assert out.data_ptr() != hyper_input.data_ptr(), "the result is a fresh tensor"
- assert torch.equal(hyper_input, before), "the residual it was given is read-only"
- _compare(out, op.reference(hyper_input, block_out, inject))
-
-
-def test_guard_refuses_unequal_leading_dims() -> None:
- """Unguarded, a 1-D `inject` broadcasts one row over every token."""
- hyper_input, block_out, inject = _operands((8,), seed=10)
- unguarded = op(hyper_input, block_out, inject[0])
- assert unguarded.shape == hyper_input.shape, "the hazard this guard covers is gone"
- assert not torch.equal(unguarded, op.reference(hyper_input, block_out, inject))
-
- with pytest.raises(AssertionError, match="leading dims must match"):
- with validating(op):
- op(hyper_input, block_out, inject[0])
-
-
-def test_cuda_graph_capture_and_replay() -> None:
- hyper_input, block_out, inject = _operands((64,), seed=14)
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(warmup):
- for _ in range(3):
- op(hyper_input, block_out, inject)
- torch.cuda.current_stream().wait_stream(warmup)
-
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- captured = op(hyper_input, block_out, inject)
-
- fresh, _, _ = _operands((64,), seed=15)
- hyper_input.copy_(fresh)
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(captured, op.reference(hyper_input, block_out, inject))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py
deleted file mode 100644
index f23d2789ee81..000000000000
--- 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
+++ /dev/null
@@ -1,399 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the CuteDSL NVFP4 FC1 entry of the Qwen3.8-Flash-Next target.
-
-The entry owns the contract. This file builds the checkpoint-shaped operands --
-512 experts in the CuteDSL weight layout, the activation as the target quantizes
-it, the sort as moe_sort returns it -- and holds the claims no cell carries:
-per-pair invariance across tiles and tactics, and the contiguity guard.
-"""
-
-import functools
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion import ( # noqa: E501
- E2M1,
- E2M1_MAX,
- GATE_UP_GROUP,
- Fc1Expected,
- e2m1_codes,
- e4m3_value,
- unswizzle,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion import ( # noqa: E501
- cute_dsl_nvfp4_gather_grouped_gemm_act_fusion as op,
-)
-from tensorrt_llm._torch.autotuner import AutoTuner
-from tensorrt_llm._torch.utils import ActivationType
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "the CuteDSL FC1 requires a CUDA device"
-
-DEV = "cuda"
-E, K, H, INTER, VEC = 512, 10, 2560, 640, 16
-SWIGLU = int(ActivationType.Swiglu)
-#: The activation's input scale: |x| = 5 maps to the e4m3 maximum, as the
-#: checkpoint's per-layer `input_scale` does for its own calibration maximum.
-ISC1 = 5.0 / (448 * 6)
-#: The SwiGLU output maximum `global_sf` is calibrated to (see `_ckpt` in the entry).
-ACT_AMAX = 48.0
-GSF = {
- "calibrated": 448 * 6 / ACT_AMAX,
- "subnormal": 2.0**-12 * 448 * 6 / ACT_AMAX,
-}
-
-#: Tie windows of `_compare`, derived in its docstring.
-SF_TIE = 2.0**-20
-CODE_TIE = 2.0**-17
-#: At most this fraction of block scales / codes may take the other side of a tie
-#: (4 of them when that is more).
-FLIP_BUDGET = 1e-5
-
-
-# ── the gate ──────────────────────────────────────────────────────────────
-
-
-def _e2m1_rank(packed: torch.Tensor) -> torch.Tensor:
- """Packed codes -> signed magnitude rank in -7..7 per element; -0 and +0 are both 0."""
- lo, hi = (packed & 0xF).long(), (packed >> 4).long()
- codes = torch.stack([lo, hi], dim=-1).reshape(packed.shape[0], -1)
- return torch.where(codes >= 8, -(codes & 7), codes & 7)
-
-
-def _e2m1_rank_of(v: torch.Tensor) -> torch.Tensor:
- codes = e2m1_codes(v).long()
- return torch.where(codes >= 8, -(codes & 7), codes & 7)
-
-
-def _rank_value(rank: torch.Tensor) -> torch.Tensor:
- lut = torch.tensor(E2M1, dtype=torch.float64, device=rank.device)
- return lut[rank.abs()] * torch.sign(rank).double()
-
-
-def _deviation(out: tuple[torch.Tensor, torch.Tensor], ref: Fc1Expected) -> dict:
- """Every disagreement between the op and the reference on the rows it writes.
-
- Block scales: equal, or the two neighbours of an e4m3 midpoint that the exact
- pre-rounding scale lies within `SF_TIE` (relative) of -- a `sf_flip`; anything
- else is `sf_outside`. Codes are re-derived under the op's OWN block scale (so
- a flipped scale does not cascade into its 16 codes): equal, or the two
- neighbours of an e2m1 midpoint that the exact `scaled / sf` lies within
- `CODE_TIE` of -- a `code_flip`; anything else is `code_outside`. Blocks whose
- scale is zero dequantize to zero whatever their codes hold and are not
- code-checked. `*_far` counts the disagreements that are not between
- neighbours at all.
- """
- c, c_sf = out
- m, half = c.shape
- inter = half * 2
- nb = c_sf.numel() // m
- vec = inter // nb
- rows = ref.rows
- s_op = unswizzle(c_sf.view(torch.uint8), m, nb)[rows].long()
- s_ref = unswizzle(ref.c_sf.view(torch.uint8), m, nb)[rows].long()
- pre = ref.scaled.view(-1, nb, vec).abs().amax(-1) / E2M1_MAX
-
- v_op, v_ref = e4m3_value(s_op), e4m3_value(s_ref)
- sf_diff = s_op != s_ref
- finite = (s_op < 0x7F) & (s_ref < 0x7F)
- sf_adj = sf_diff & finite & ((s_op - s_ref).abs() == 1)
- sf_dist = ((pre - (v_op + v_ref) / 2).abs() / pre.clamp_min(1e-300)).masked_fill(~sf_adj, 0.0)
- sf_flip = sf_adj & (sf_dist <= SF_TIE)
- sf_outside = sf_diff & ~sf_flip
-
- rank_op = _e2m1_rank(c.view(torch.uint8)[rows])
- scale = v_op.repeat_interleave(vec, dim=1)
- live = scale > 0
- q = torch.where(live, ref.scaled / scale, torch.zeros_like(scale))
- rank_exp = _e2m1_rank_of(q)
- code_diff = live & (rank_op != rank_exp)
- code_adj = code_diff & ((rank_op - rank_exp).abs() == 1)
- mid = (_rank_value(rank_op) + _rank_value(rank_exp)) / 2
- code_dist = (q - mid).abs().masked_fill(~code_adj, 0.0)
- code_flip = code_adj & (code_dist <= CODE_TIE)
- code_outside = code_diff & ~code_flip
-
- return dict(
- rows=int(rows.sum()),
- blocks=s_op.numel(),
- sf_flips=int(sf_flip.sum()),
- sf_outside=int(sf_outside.sum()),
- sf_far=int((sf_diff & ~sf_adj).sum()),
- elements=int(live.sum()),
- code_flips=int(code_flip.sum()),
- code_outside=int(code_outside.sum()),
- code_far=int((code_diff & ~code_adj).sum()),
- )
-
-
-def _compare(out: tuple[torch.Tensor, torch.Tensor], ref: Fc1Expected) -> None:
- """Bitwise on every written row, except at rounding ties of bounded width.
-
- The op accumulates in fp32 on the tensor cores, evaluates SwiGLU with
- `ex2.approx` / `rcp.approx`, and divides by the block scale through
- `rcp.approx`; the reference is exact. The op's pre-rounding value is
- therefore off by a handful of fp32 roundings (2^-24 relative each), and
- only where the exact value lies that close to a rounding midpoint may the
- op land on its other side. Everywhere else the two agree bit for bit.
-
- - Block scale: a flip only within `SF_TIE` = 2^-20 (16 fp32 roundings)
- relative of an e4m3 midpoint.
- - Code: a flip only within `CODE_TIE` = 2^-17 of an e2m1 midpoint, in units
- of the block scale -- 16 fp32 roundings of a block maximum of at most 6,
- rounded up to a power of two.
- - Flips: at most `FLIP_BUDGET` = 1e-5 of the scales and of the codes (4 when
- that is fewer), so the op may not take the far side of every near-tie.
-
- Every wrongly prepared operand puts most of the block scales outside the gate.
- """
- c, c_sf = out
- assert c.dtype == torch.float4_e2m1fn_x2, c.dtype
- assert c_sf.dtype == torch.uint8, c_sf.dtype
- assert c.shape == ref.c.shape, (tuple(c.shape), tuple(ref.c.shape))
- assert c_sf.shape == ref.c_sf.shape, (tuple(c_sf.shape), tuple(ref.c_sf.shape))
- d = _deviation(out, ref)
- assert d["rows"] > 0, d
- assert d["sf_outside"] == 0, f"block scales off by more than a tie: {d}"
- assert d["code_outside"] == 0, f"codes off by more than a tie: {d}"
- assert d["sf_flips"] <= max(4, FLIP_BUDGET * d["blocks"]), f"too many scale ties flipped: {d}"
- assert d["code_flips"] <= max(4, FLIP_BUDGET * d["elements"]), (
- f"too many code ties flipped: {d}"
- )
-
-
-# ── operands ──────────────────────────────────────────────────────────────
-
-
-@functools.lru_cache(maxsize=None)
-def _experts(num_experts: int = E):
- """Checkpoint-format routed experts: per expert `[up; gate]` NVFP4 rows with
- LINEAR e4m3 block scales and one `weight_scale_2` shared by both halves, as the
- checkpoint stores them. Expert magnitudes spread over 2^-1..2^1 so a wrong
- per-expert alpha shows."""
- g = torch.Generator(device=DEV).manual_seed(20261013)
- spread = torch.exp2(torch.rand(num_experts, generator=g, device=DEV) * 2 - 1)
- w = torch.empty(num_experts, 2 * INTER, H // 2, dtype=torch.uint8, device=DEV)
- s = torch.empty(num_experts, 2 * INTER, H // VEC, dtype=torch.uint8, device=DEV)
- ws2 = torch.empty(num_experts, dtype=torch.float32, device=DEV)
- for e in range(num_experts):
- dense = torch.randn(2 * INTER, H, generator=g, device=DEV) * (float(spread[e]) / H**0.5)
- dense = dense.to(torch.bfloat16)
- gs = (448 * 6 / dense.abs().amax().float()).reshape(1)
- q, sf = torch.ops.trtllm.fp4_quantize(dense, gs, VEC, False, False)
- w[e] = q.view(2 * INTER, H // 2)
- s[e] = sf.view(2 * INTER, H // VEC)
- ws2[e] = 1.0 / gs
- return w, s, ws2
-
-
-def _row_perm(n: int) -> torch.Tensor:
- """Destination row -> source row of the `[up; gate]` stack, in `GATE_UP_GROUP`-row
- groups."""
- idx = torch.arange(n, device=DEV)
- return idx.view(2, n // (2 * GATE_UP_GROUP), GATE_UP_GROUP).transpose(0, 1).reshape(-1)
-
-
-@functools.lru_cache(maxsize=None)
-def _weights():
- """`(weight, weight_scale, alpha)` in the layout CuteDSL reads -- what the
- target's load-time transform produces."""
- w, s, ws2 = _experts()
- perm = _row_perm(2 * INTER)
- w_k = w[:, perm].contiguous()
- s_k = torch.ops.trtllm.block_scale_interleave(s[:, perm].contiguous()).view(s.shape)
- alpha = (ISC1 * ws2).contiguous()
- return w_k.view(torch.float4_e2m1fn_x2), s_k, alpha
-
-
-def _activation(tokens: int, seed: int, zero_tokens: int = 0):
- """The target's FC1 input: bf16 hidden rows through `fp4_quantize(x, 1/isc1,
- 16, False, False)`, scales viewed `[T, H/16]`."""
- g = torch.Generator(device=DEV).manual_seed(seed)
- x = torch.randn(tokens, H, generator=g, device=DEV).to(torch.bfloat16)
- x[:zero_tokens] = 0
- gs = torch.tensor([1.0 / ISC1], dtype=torch.float32, device=DEV)
- q, sf = torch.ops.trtllm.fp4_quantize(x, gs, VEC, False, False)
- return q.view(torch.float4_e2m1fn_x2), sf.view(tokens, H // VEC)
-
-
-def _ids(dist: str, tokens: int, seed: int) -> torch.Tensor:
- """Expert ids `[tokens, K]` int32: each row's top-K of uniform scores."""
- assert dist == "router", dist
- g = torch.Generator(device=DEV).manual_seed(seed)
- scores = torch.rand(tokens, E, device=DEV, generator=g)
- return scores.topk(K, dim=-1).indices.to(torch.int32)
-
-
-def _sort(ids: torch.Tensor, tile: int):
- """moe_sort as the CuteDSL MoE calls it; the routing weights do not reach FC1."""
- scales = torch.full(ids.shape, 0.1, dtype=torch.float32, device=DEV)
- return torch.ops.trtllm.moe_sort(
- token_selected_experts=ids,
- token_final_scales=scales,
- num_experts=E,
- top_k=K,
- local_expert_offset=0,
- local_num_experts=E,
- tile_tokens_dim=tile,
- )
-
-
-def _args(x, x_sf, sort, tile, gsf):
- """The keyword arguments the target passes, by the op's own names."""
- weight, weight_scale, alpha = _weights()
- t2e, t2mn, _e2p, p2e, _total, n_tiles = sort
- return dict(
- input=x,
- weight=weight,
- input_scale=x_sf,
- weight_scale=weight_scale,
- alpha=alpha,
- tile_idx_to_group_idx=t2e,
- tile_idx_to_mn_limit=t2mn,
- permuted_idx_to_expanded_idx=p2e,
- num_non_exiting_tiles=n_tiles,
- global_sf=torch.tensor([gsf], dtype=torch.float32, device=DEV),
- num_experts=E,
- top_k=K,
- num_local_experts=E,
- local_expert_offset=0,
- tile_size=tile,
- scaling_vector_size=VEC,
- activation_type=SWIGLU,
- )
-
-
-def _problem(tokens, tile, seed, dist="router", gsf="calibrated", zero_tokens=0):
- x, x_sf = _activation(tokens, seed, zero_tokens)
- sort = _sort(_ids(dist, tokens, seed + 1), tile)
- return _args(x, x_sf, sort, tile, GSF[gsf]), sort
-
-
-def _call(args, mma_n):
- """One op call with the autotuner tactic whose MMA tile is `(tile_size, mma_n)`,
- selected through the tuner's own capture/replay so the call still goes through
- the op."""
- tuner = AutoTuner.get()
- with tuner.capture() as cap:
- op(**args)
- for config in cap:
- ((_runner, tactic),) = config
- if isinstance(tactic, tuple) and tactic[0] == (args["tile_size"], mma_n):
- with tuner.replay(config):
- return op(**args)
- raise AssertionError(f"no tactic with MMA tile ({args['tile_size']}, {mma_n})")
-
-
-def _pairs(out, sort, tokens):
- """Each (token, slot) pair's output row and its block-scale row, in expanded
- order -- what FC2 and the combine see of it, whatever row the sort gave it."""
- c, c_sf = out
- e2p = sort[2].reshape(-1).long()[: tokens * K]
- sf = unswizzle(c_sf, c.shape[0], c_sf.numel() // c.shape[0])
- return c.view(torch.uint8)[e2p], sf[e2p]
-
-
-def _clone(v):
- """A copy that does not depend on copy kernels for the fp4 dtype."""
- if not isinstance(v, torch.Tensor):
- return v
- return v.view(torch.uint8).clone().view(v.dtype) if v.element_size() == 1 else v.clone()
-
-
-def _assert_pairs_equal(a, b, what):
- assert torch.equal(a[0], b[0]), f"{what}: codes differ"
- assert torch.equal(a[1], b[1]), f"{what}: block scales differ"
-
-
-def _cell_id(cell):
- s = cell.spec
- return (
- f"t{s['tile_size']}-n{s['mma_n']}-T{s['tokens']}-{s['dist']}-{s['gsf']}-z{s['zero_tokens']}"
- )
-
-
-# ── the certified cells ───────────────────────────────────────────────────
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[_cell_id(c) for c in op.CELLS])
-def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- s = cell.spec
- assert (s["num_experts"], s["top_k"], s["hidden"], s["intermediate"]) == (E, K, H, INTER)
- seed = zlib.crc32(cell.why.encode())
- args, sort = _problem(s["tokens"], s["tile_size"], seed, s["dist"], s["gsf"], s["zero_tokens"])
- before = {k: _clone(v) for k, v in args.items() if isinstance(v, torch.Tensor)}
- out = _call(args, s["mma_n"])
- _compare(out, op.reference(**args))
- for k, v in before.items():
- assert torch.equal(args[k].view(torch.uint8), v.view(torch.uint8)), f"{k} was written"
- if s["zero_tokens"]:
- zero = sort[2][: s["zero_tokens"]].reshape(-1).long()
- c, c_sf = out
- sf = unswizzle(c_sf, c.shape[0], INTER // VEC)
- assert bool((sf[zero] == 0).all()), "an all-zero token's block scales are not zero"
- assert bool((c.view(torch.uint8)[zero] & 0x77 == 0).all()), "a zero token's codes are not 0"
-
-
-# ── invariance: a pair's output depends on the pair alone ─────────────────
-
-
-def test_tiles_and_tactics_agree_per_pair(certifying): # noqa: F811
- """Tile 128 / 256, MMA N 128 / 256 and the tuner's fallback all produce every
- pair's codes and block scales bit for bit the same at the 8192-row prefill, so a
- cell certifies the tactic the autotuner picks in serving as well as the one it
- ran."""
- certifying(op)
- tokens = 8192
- x, x_sf = _activation(tokens, seed=tokens)
- ids = _ids("router", tokens, seed=tokens + 1)
- seen = {}
- for tile in (128, 256):
- sort = _sort(ids, tile)
- args = _args(x, x_sf, sort, tile, GSF["calibrated"])
- for mma_n in (128, 256):
- seen[(tile, mma_n)] = _pairs(_call(args, mma_n), sort, tokens)
- seen[(tile, "fallback")] = _pairs(op(**args), sort, tokens)
- if tile == 128:
- _compare(op(**args), op.reference(**args))
- first = seen[(128, 128)]
- for key, got in seen.items():
- _assert_pairs_equal(got, first, f"T={tokens} {key} vs (128, 128)")
-
-
-# ── the guard ─────────────────────────────────────────────────────────────
-
-
-def _strided(t: torch.Tensor) -> torch.Tensor:
- """The same values as `t`, as a non-contiguous view into a wider buffer."""
- u = t.view(torch.uint8) if t.element_size() == 1 else t
- wide = torch.zeros(*u.shape[:-1], 2 * u.shape[-1], dtype=u.dtype, device=DEV)
- wide[..., : u.shape[-1]] = u
- view = wide[..., : u.shape[-1]]
- return view.view(t.dtype) if t.element_size() == 1 else view
-
-
-def _u8(t: torch.Tensor) -> torch.Tensor:
- return t.view(torch.uint8) if t.element_size() == 1 else t
-
-
-def test_noncontiguous_operand_answers_wrongly_and_is_refused():
- """The runner hands data_ptr() to the kernel and never reads a stride: the
- same values through a strided view come back finite and wrong. The activation;
- the guard checks input_scale, weight, weight_scale and alpha the same way."""
- name = "input"
- tokens, tile = 64, 128
- args, _ = _problem(tokens, tile, seed=21)
- ref = op.reference(**args)
- bad = dict(args, **{name: _strided(args[name])})
- assert not bad[name].is_contiguous()
- assert torch.equal(_u8(bad[name]), _u8(args[name]))
- with pytest.raises(AssertionError):
- _compare(op(**bad), ref)
- with validating(op), pytest.raises(AssertionError, match=f"^{name} must be contiguous"):
- op(**bad)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py
deleted file mode 100644
index f744f13a4402..000000000000
--- 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
+++ /dev/null
@@ -1,375 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the CuteDSL NVFP4 FC2 entry of the Qwen3.8-Flash-Next target.
-
-The entry owns the contract. This file builds the checkpoint-shaped operands --
-512 experts' down_proj in the CuteDSL weight layout, the input as the CuteDSL
-FC1 returns it (permuted NVFP4 rows, swizzled block scales, zero padding rows,
-unwritten bytes past the live tiles), the sort as moe_sort returns it -- and
-holds the claims no cell carries: per-pair invariance across tiles and tactics,
-and the contiguity guard.
-"""
-
-import functools
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.cute_dsl_nvfp4_grouped_gemm import ( # noqa: E501
- Fc2Expected,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.cute_dsl_nvfp4_grouped_gemm import ( # noqa: E501
- cute_dsl_nvfp4_grouped_gemm as op,
-)
-from tensorrt_llm._torch.autotuner import AutoTuner
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "the CuteDSL FC2 requires a CUDA device"
-
-DEV = "cuda"
-E, K, H, INTER, VEC = 512, 10, 2560, 640, 16
-NB = INTER // VEC
-#: The SwiGLU output maximum FC1's `global_sf` is calibrated to (see `_ckpt` in the entry).
-ACT_AMAX = 48.0
-GSF = {
- "calibrated": 448 * 6 / ACT_AMAX,
- "subnormal": 2.0**-12 * 448 * 6 / ACT_AMAX,
-}
-
-#: Tie window of `_compare`, in units of the element's accumulation magnitude
-#: (derived in its docstring).
-TIE = 2.0**-20
-#: At most this fraction of the gated elements with a nonzero magnitude may take
-#: the other side of a tie (4 of them when that is more).
-FLIP_BUDGET = 1e-5
-
-
-# ── the gate ──────────────────────────────────────────────────────────────
-
-
-def _bf16_order(x: torch.Tensor) -> torch.Tensor:
- """bf16 -> an integer that is monotonic in its value, consecutive between
- neighbouring bf16 values, and equal for +0 and -0."""
- bits = x.view(torch.int16).int() & 0xFFFF
- mag = bits & 0x7FFF
- return torch.where(bits >= 0x8000, -mag, mag)
-
-
-def _deviation(out: torch.Tensor, ref: Fc2Expected) -> dict:
- """Every disagreement between the op and the reference on the rows it writes.
-
- An element is `equal` when its value matches the reference's `out` (+0 and
- -0 alike). Otherwise it is a `flip` when the two are neighbouring bf16
- values and the exact value lies within `TIE * magnitude` of the midpoint
- between them, and `outside` when it is anything else. `far` counts the
- disagreements that are not between neighbours at all.
- """
- span = ref.exact.shape[0]
- idx = torch.nonzero(ref.rows[:span]).squeeze(1)
- o = out[idx]
- r = ref.out[idx]
- v = ref.exact[idx]
- s = ref.magnitude[idx]
- of, rf = o.double(), r.double()
- finite = torch.isfinite(of)
- diff = (of != rf) | ~finite
- adj = diff & finite & ((_bf16_order(o) - _bf16_order(r)).abs() == 1)
- dist = ((v - (of + rf) / 2).abs() / s.clamp_min(1e-300)).masked_fill(~adj, 0.0)
- flip = adj & (dist <= TIE)
- outside = diff & ~flip
- return dict(
- rows=int(idx.numel()),
- elements=o.numel(),
- nonzero=int((s > 0).sum()),
- nonfinite=int((~finite).sum()),
- flips=int(flip.sum()),
- outside=int(outside.sum()),
- far=int((diff & ~adj).sum()),
- )
-
-
-def _compare(out: torch.Tensor, ref: Fc2Expected) -> None:
- """Bitwise on every written row, except at bf16 rounding ties of derived width.
-
- The op accumulates the 640 products of a row in fp32 on the tensor cores --
- ten MMA k-steps of 64 -- multiplies by alpha in fp32 and rounds to bf16. Each
- step rounds a partial sum no larger than the element's `magnitude` (`|alpha| *
- sum |a_k w_k|`), so the op's fp32 value is within 11 roundings of `2^-24 *
- magnitude` of the exact one; only where the exact value lies that close to a
- bf16 midpoint may the op land on the other neighbour. Everywhere else the two
- agree bit for bit.
-
- - Flip: only to the neighbouring bf16 value, and only within `TIE` = 2^-20
- (16 fp32 roundings) of `magnitude` from the midpoint.
- - Flips: at most `FLIP_BUDGET` = 1e-5 of the gated elements whose magnitude is
- nonzero (4 when that is fewer) -- padding rows do not widen the budget.
- - No NaN or infinity.
-
- Every wrongly prepared operand puts most of the pair rows' elements outside the
- gate, most of them not even a bf16 neighbour away.
- """
- assert out.dtype == torch.bfloat16, out.dtype
- assert out.shape == ref.out.shape, (tuple(out.shape), tuple(ref.out.shape))
- d = _deviation(out, ref)
- assert d["rows"] > 0, d
- assert d["nonfinite"] == 0, f"non-finite output on written rows: {d}"
- assert d["outside"] == 0, f"elements off by more than a tie: {d}"
- assert d["flips"] <= max(4, FLIP_BUDGET * d["nonzero"]), f"too many ties flipped: {d}"
-
-
-# ── operands ──────────────────────────────────────────────────────────────
-
-
-def _swizzle(sf: torch.Tensor) -> torch.Tensor:
- """`[rows, cols]` scale bytes -> the flat 128x4-swizzled buffer; inverse of `unswizzle`."""
- rows, cols = sf.shape
- return sf.reshape(rows // 128, 4, 32, cols // 4, 4).permute(0, 3, 2, 1, 4).reshape(-1)
-
-
-def _quantized_experts(num_experts: int, rows: int, cols: int, seed: int):
- """Per expert, a Gaussian `[rows, cols]` matrix quantized the way the checkpoint
- stores it: NVFP4 codes, LINEAR e4m3 block scales, one `weight_scale_2`. Expert
- magnitudes spread over 2^-1..2^1 so a wrong per-expert alpha shows."""
- g = torch.Generator(device=DEV).manual_seed(seed)
- spread = torch.exp2(torch.rand(num_experts, generator=g, device=DEV) * 2 - 1)
- w = torch.empty(num_experts, rows, cols // 2, dtype=torch.uint8, device=DEV)
- s = torch.empty(num_experts, rows, cols // VEC, dtype=torch.uint8, device=DEV)
- ws2 = torch.empty(num_experts, dtype=torch.float32, device=DEV)
- for e in range(num_experts):
- dense = torch.randn(rows, cols, generator=g, device=DEV) * (float(spread[e]) / cols**0.5)
- dense = dense.to(torch.bfloat16)
- gs = (448 * 6 / dense.abs().amax().float()).reshape(1)
- q, sf = torch.ops.trtllm.fp4_quantize(dense, gs, VEC, False, False)
- w[e] = q.view(rows, cols // 2)
- s[e] = sf.view(rows, cols // VEC)
- ws2[e] = 1.0 / gs
- return w, s, ws2
-
-
-@functools.lru_cache(maxsize=None)
-def _down_proj():
- """The checkpoint's `down_proj` experts: `[E, H, I/2]` codes, `[E, H, I/16]`
- linear scales, `weight_scale_2` per expert."""
- return _quantized_experts(E, H, INTER, seed=20261013)
-
-
-@functools.lru_cache(maxsize=None)
-def _weights(gsf: str = "calibrated"):
- """`(weight, weight_scale, alpha)` in the layout CuteDSL reads -- what the target's
- load-time transform produces. `alpha = down_proj.input_scale * weight_scale_2`,
- with the input scale the inverse of FC1's `global_sf`."""
- w, s, ws2 = _down_proj()
- s_k = torch.ops.trtllm.block_scale_interleave(s.contiguous()).view(s.shape)
- alpha = (ws2 / GSF[gsf]).contiguous()
- return w.contiguous().view(torch.float4_e2m1fn_x2), s_k, alpha
-
-
-def _pair_rows(pairs: int, seed: int, zero_pairs: int = 0) -> torch.Tensor:
- """SwiGLU outputs of FC1, one bf16 `[I]` row per (token, slot) pair, in expanded
- order: `up * silu(gate)` of Gaussian halves."""
- g = torch.Generator(device=DEV).manual_seed(seed)
- up = torch.randn(pairs, INTER, generator=g, device=DEV) * 2
- gate = torch.randn(pairs, INTER, generator=g, device=DEV) * 2
- y = (up * gate * torch.sigmoid(gate)).to(torch.bfloat16)
- y[:zero_pairs] = 0
- return y
-
-
-def _fc2_input(y: torch.Tensor, sort, gsf: str, garbage_seed: int = 0):
- """This op's input as the CuteDSL FC1 returns it: each pair's row NVFP4-quantized
- with `global_sf` at its permuted row, block scales 128x4-swizzled, the padding
- rows of live tiles zero, and the rows past the live tiles holding stale bytes."""
- t2e, _t2mn, e2p, p2e, _total, n_tiles = sort
- m = p2e.numel()
- tile = m // t2e.numel()
- span = int(n_tiles) * tile
- x = torch.zeros(m, INTER, dtype=torch.bfloat16, device=DEV)
- x[e2p.reshape(-1).long()[: y.shape[0]]] = y
- gs = torch.tensor([GSF[gsf]], dtype=torch.float32, device=DEV)
- q, sf = torch.ops.trtllm.fp4_quantize(x, gs, VEC, False, False)
- q = q.view(torch.uint8).reshape(m, INTER // 2)
- sf = sf.view(torch.uint8).reshape(m, NB)
- g = torch.Generator(device=DEV).manual_seed(garbage_seed)
- q[span:] = torch.randint(0, 256, q[span:].shape, generator=g, device=DEV, dtype=torch.uint8)
- sf[span:] = torch.randint(0, 256, sf[span:].shape, generator=g, device=DEV, dtype=torch.uint8)
- return q.view(torch.float4_e2m1fn_x2), _swizzle(sf).contiguous()
-
-
-def _ids(dist: str, tokens: int, seed: int) -> torch.Tensor:
- """Expert ids `[tokens, K]` int32: each row's top-K of uniform scores."""
- assert dist == "router", dist
- g = torch.Generator(device=DEV).manual_seed(seed)
- scores = torch.rand(tokens, E, device=DEV, generator=g)
- return scores.topk(K, dim=-1).indices.to(torch.int32)
-
-
-def _sort(ids: torch.Tensor, tile: int):
- """moe_sort as the CuteDSL MoE calls it; the routing weights do not reach FC2."""
- scales = torch.full(ids.shape, 0.1, dtype=torch.float32, device=DEV)
- return torch.ops.trtllm.moe_sort(
- token_selected_experts=ids,
- token_final_scales=scales,
- num_experts=E,
- top_k=K,
- local_expert_offset=0,
- local_num_experts=E,
- tile_tokens_dim=tile,
- )
-
-
-def _args(x, x_sf, sort, tile, gsf="calibrated"):
- """The keyword arguments the target passes, by the op's own names."""
- weight, weight_scale, alpha = _weights(gsf)
- return dict(
- input=x,
- weight=weight,
- input_scale=x_sf,
- weight_scale=weight_scale,
- alpha=alpha,
- tile_idx_to_group_idx=sort[0],
- num_non_exiting_tiles=sort[5],
- num_experts=E,
- top_k=K,
- num_local_experts=E,
- local_expert_offset=0,
- tile_size=tile,
- output_dtype=torch.bfloat16,
- scaling_vector_size=VEC,
- )
-
-
-def _problem(tokens, tile, seed, dist="router", gsf="calibrated", zero_tokens=0):
- sort = _sort(_ids(dist, tokens, seed + 1), tile)
- y = _pair_rows(tokens * K, seed, zero_pairs=zero_tokens * K)
- x, x_sf = _fc2_input(y, sort, gsf, garbage_seed=seed + 2)
- return _args(x, x_sf, sort, tile, gsf), sort
-
-
-def _call(args, mma_n, cluster_n):
- """One op call with the autotuner tactic whose MMA tile is `(tile_size, mma_n)`
- and cluster `(tile_size // 128, cluster_n)`, selected through the tuner's own
- capture/replay so the call still goes through the op."""
- tuner = AutoTuner.get()
- with tuner.capture() as cap:
- op(**args)
- tile = args["tile_size"]
- want = ((tile, mma_n), (tile // 128, cluster_n))
- for config in cap:
- ((_runner, tactic),) = config
- if tactic == want:
- with tuner.replay(config):
- return op(**args)
- raise AssertionError(f"no tactic {want}")
-
-
-def _pairs(out, sort, tokens):
- """Each (token, slot) pair's output row, in expanded order -- what the combine
- reads, whatever row the sort gave it."""
- return out[sort[2].reshape(-1).long()[: tokens * K]]
-
-
-def _clone(v):
- """A copy that does not depend on copy kernels for the fp4 dtype."""
- if not isinstance(v, torch.Tensor):
- return v
- return v.view(torch.uint8).clone().view(v.dtype) if v.element_size() == 1 else v.clone()
-
-
-def _bits(t: torch.Tensor) -> torch.Tensor:
- return t.view(torch.int16)
-
-
-def _assert_pairs_equal(a, b, what):
- assert torch.equal(_bits(a), _bits(b)), f"{what}: rows differ"
-
-
-def _cell_id(cell):
- s = cell.spec
- return (
- f"t{s['tile_size']}-n{s['mma_n']}-c{s['cluster_n']}-T{s['tokens']}-{s['dist']}"
- f"-{s['gsf']}-z{s['zero_tokens']}"
- )
-
-
-# ── the certified cells ───────────────────────────────────────────────────
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[_cell_id(c) for c in op.CELLS])
-def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- s = cell.spec
- assert (s["num_experts"], s["top_k"], s["hidden"], s["intermediate"]) == (E, K, H, INTER)
- seed = zlib.crc32(cell.why.encode())
- args, sort = _problem(s["tokens"], s["tile_size"], seed, s["dist"], s["gsf"], s["zero_tokens"])
- before = {k: _clone(v) for k, v in args.items() if isinstance(v, torch.Tensor)}
- out = _call(args, s["mma_n"], s["cluster_n"])
- _compare(out, op.reference(**args))
- for k, v in before.items():
- assert torch.equal(_u8(args[k]), _u8(v)), f"{k} was written"
- if s["zero_tokens"]:
- zero = sort[2][: s["zero_tokens"]].reshape(-1).long()
- assert bool((out[zero] == 0).all()), "an all-zero token's rows are not zero"
-
-
-# ── invariance: a pair's row depends on the pair alone ────────────────────
-
-
-def test_tiles_and_tactics_agree_per_pair(certifying): # noqa: F811
- """Tile 128 / 256, MMA N 128 / 256, cluster N 1 / 2 and the tuner's fallback
- all produce every pair's row bit for bit the same at the 8192-row prefill, so a
- cell certifies the tactic the autotuner picks in serving as well as the one it
- ran."""
- certifying(op)
- tokens = 8192
- ids = _ids("router", tokens, seed=tokens + 1)
- y = _pair_rows(tokens * K, seed=tokens)
- seen = {}
- for tile in (128, 256):
- sort = _sort(ids, tile)
- x, x_sf = _fc2_input(y, sort, "calibrated", garbage_seed=tile)
- args = _args(x, x_sf, sort, tile)
- for mma_n in (128, 256):
- for cluster_n in (1, 2):
- seen[(tile, mma_n, cluster_n)] = _pairs(_call(args, mma_n, cluster_n), sort, tokens)
- fallback = op(**args)
- seen[(tile, "fallback")] = _pairs(fallback, sort, tokens)
- _compare(fallback, op.reference(**args))
- first = seen[(128, 128, 1)]
- for key, got in seen.items():
- _assert_pairs_equal(got, first, f"T={tokens} {key} vs (128, 128, 1)")
-
-
-# ── the guard ─────────────────────────────────────────────────────────────
-
-
-def _u8(t: torch.Tensor) -> torch.Tensor:
- return t.view(torch.uint8) if t.element_size() == 1 else t
-
-
-def _strided(t: torch.Tensor) -> torch.Tensor:
- """The same values as `t`, as a non-contiguous view into a wider buffer."""
- u = _u8(t)
- wide = torch.zeros(*u.shape[:-1], 2 * u.shape[-1], dtype=u.dtype, device=DEV)
- wide[..., : u.shape[-1]] = u
- view = wide[..., : u.shape[-1]]
- return view.view(t.dtype) if t.element_size() == 1 else view
-
-
-def test_noncontiguous_operand_answers_wrongly_and_is_refused():
- """The runner hands data_ptr() to the kernel and never reads a stride: the
- same values through a strided view come back finite and wrong. The activation;
- the guard checks the other five operands the same way."""
- name = "input"
- tokens, tile = 64, 128
- args, _ = _problem(tokens, tile, seed=21)
- ref = op.reference(**args)
- bad = dict(args, **{name: _strided(args[name])})
- assert not bad[name].is_contiguous()
- assert torch.equal(_u8(bad[name]), _u8(args[name]))
- with pytest.raises(AssertionError):
- _compare(op(**bad), ref)
- with validating(op), pytest.raises(AssertionError, match=f"^{name} must be contiguous"):
- op(**bad)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py
deleted file mode 100644
index 6d7963ed32db..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py
+++ /dev/null
@@ -1,371 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""The routed experts of the Qwen3.8-Flash-Next target, from checkpoint-format tensors to output.
-
-The CuteDSL op tests build their operands with layout code of their own. This file builds
-the checkpoint's expert tensors (512 experts: packed NVFP4 `up_proj` / `gate_proj` /
-`down_proj` codes, linear e4m3 block scales, one `weight_scale_2` per expert, per-expert
-`input_scale`s) and runs them through the target's own load-time transforms
-(`weights._expert_fc1_w` / `_expert_fc1_s` / `_expert_fc2_w` / `_expert_fc2_s`) and scalar
-derivation (`modeling._moe_layer_scalars`), once per routed backend. The operands then go
-through the call sequence of the target's `_moe_block` -- CuteDSL: sort, FC1 with SwiGLU,
-FC2, unpermute; trtllm-gen: the one-call block-scale runner -- with explicit routing.
-
-Two routings hit all 512 experts. In the first every token's only non-zero routing weight is
-on its own expert, so each expert's output is checked alone: a layout error confined to one
-expert (a permutation, a scale swizzle, an alpha) cannot hide in a sum over ten. On it both
-backends are held to a float64 reference over the dequantized checkpoint tensors, with
-`down_proj.input_scale` spread across the experts as a calibrated checkpoint's is. The second
-is a weighted top-10: with a uniform `down_proj.input_scale` the two backends requantize the
-SwiGLU output identically, so on it they are held to each other -- the area's one
-cross-backend check.
-
-The reference does not model the kernels' accumulation. It models the one rounding step
-that is not exact: the SwiGLU output is requantized to NVFP4 before FC2, at the backend's
-own global scale (the layer's single `global_sf` on CuteDSL, each expert's
-`1 / down_proj.input_scale` on trtllm-gen).
-"""
-
-import functools
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import (
- fp4_block_scale_moe_runner,
- quantize_nvfp4,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import (
- fp4_quantize,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501
- modeling as M,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501
- weights as W,
-)
-
-assert torch.cuda.is_available(), "the CuteDSL experts require a CUDA device"
-
-DEV = "cuda"
-E, K, H, INTER, VEC = 512, 10, 2560, 640, 16
-#: The activation's input scale, one number for all experts: |x| = 5 maps to the e4m3 maximum.
-ISC1 = 5.0 / (448 * 6)
-#: The SwiGLU output maximum `down_proj.input_scale` is calibrated to. With `spread` each
-#: expert's own value ranges over 2^-1/2 .. 2^1/2 of it, so the layer's `max` is not every
-#: expert's.
-ISC2 = 64.0 / (448 * 6)
-E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
-#: The routed runners read activation scales as a linear buffer, not the 128x4-swizzled
-#: one of a dense NVFP4 GEMM; byte counts match at multiples of 128 rows, so the wrong
-#: one passes silently.
-SF_LINEAR = False
-
-#: Tokens of the two routings: one per expert, and a weighted top-10.
-ONE_HOT_TOKENS = E
-WEIGHTED_TOKENS = 1024
-
-BACKENDS = ("CUTEDSL", "TRTLLM")
-#: The float64 reference's gate, relative per token. Both backends sit at 2.5e-3 or below:
-#: the bf16 rounding of the output is 1e-3 of it, the rest is the reference's fp32 / tie
-#: differences. A wrong layout is ~1; a wrong per-expert scalar is its own relative error.
-REFERENCE_TOLERANCE = 5e-3
-#: The same, between the backends under a uniform `down_proj.input_scale`. Each backend is
-#: within `REFERENCE_TOLERANCE` of the reference, so the two are within twice that of each
-#: other. Both requantize the SwiGLU output to NVFP4, and an fp32 difference in the last
-#: place moves an element at an e2m1 rounding boundary by a whole code: the token holding
-#: it differs by ~5e-3 while the mean stays near 5e-6. A wrong layout is still ~1.
-BACKEND_TOLERANCE = 2 * REFERENCE_TOLERANCE
-
-
-# ── the checkpoint ────────────────────────────────────────────────────────
-
-
-def _quantize(dense: torch.Tensor):
- """A bf16 `[R, C]` matrix as the checkpoint stores it: packed e2m1 codes `[R, C/2]`
- uint8, e4m3 block scales `[R, C/16]`, and the `weight_scale_2` that undoes the global
- scale the block scales were derived under."""
- rows, cols = dense.shape
- gs = (448 * 6 / dense.abs().amax().float()).reshape(1)
- codes, sf = torch.ops.trtllm.fp4_quantize(dense, gs, VEC, False, False)
- return (
- codes.view(torch.uint8).view(rows, cols // 2),
- sf.view(torch.uint8).view(rows, cols // VEC).view(torch.float8_e4m3fn),
- (1.0 / gs).squeeze(0),
- )
-
-
-@functools.lru_cache(maxsize=None)
-def _checkpoint(spread: bool) -> SimpleNamespace:
- """All 512 experts of one layer, with the checkpoint's tensor names and dtypes.
-
- `up_proj` and `gate_proj` are quantized together, so they share one `weight_scale_2`, as
- the fused FC1 requires. Expert magnitudes spread over 2^-1 .. 2^1/2 (FC1) and 2^-1 .. 2^1
- (FC2), so a wrong per-expert scalar shows."""
- g = torch.Generator(device=DEV).manual_seed(20261014)
- spread1 = torch.exp2(torch.rand(E, generator=g, device=DEV) * 1.5 - 1)
- spread2 = torch.exp2(torch.rand(E, generator=g, device=DEV) * 2 - 1)
- fields = {k: [] for k in ("up_w", "up_s", "gate_w", "gate_s", "down_w", "down_s")}
- ws2_gu, ws2_dn = [], []
- for e in range(E):
- gate_up = torch.randn(2 * INTER, H, generator=g, device=DEV) * (spread1[e] / H**0.5)
- down = torch.randn(H, INTER, generator=g, device=DEV) * (spread2[e] / INTER**0.5)
- codes, sf, ws2 = _quantize(gate_up.to(torch.bfloat16))
- fields["up_w"].append(codes[:INTER])
- fields["gate_w"].append(codes[INTER:])
- fields["up_s"].append(sf[:INTER])
- fields["gate_s"].append(sf[INTER:])
- ws2_gu.append(ws2)
- codes, sf, ws2 = _quantize(down.to(torch.bfloat16))
- fields["down_w"].append(codes)
- fields["down_s"].append(sf)
- ws2_dn.append(ws2)
- ws2_gu, ws2_dn = torch.stack(ws2_gu), torch.stack(ws2_dn)
- isc1 = torch.full((E,), ISC1, dtype=torch.float32, device=DEV)
- isc2 = torch.full((E,), ISC2, dtype=torch.float32, device=DEV)
- if spread:
- isc2 = isc2 * torch.exp2(torch.rand(E, generator=g, device=DEV) - 0.5)
- return SimpleNamespace(
- **fields,
- ws2_gu=ws2_gu,
- ws2_dn=ws2_dn,
- isc2=isc2,
- # The rows `modeling._moe_layer_scalars` reads, under layer 0.
- scalars={
- "l0_e_isc1": isc1,
- "l0_e_isc1_up": isc1.clone(),
- "l0_e_ws2_1": ws2_gu,
- "l0_e_ws2_1_up": ws2_gu.clone(),
- "l0_e_isc2": isc2,
- "l0_e_ws2_2": ws2_dn,
- },
- )
-
-
-# ── what the target builds from it ────────────────────────────────────────
-
-
-@functools.lru_cache(maxsize=None)
-def _layer(backend: str, spread: bool) -> SimpleNamespace:
- """The layer's expert operands for `backend`, built as the weight table builds them:
- one transform call per expert and tensor, stacked, then viewed as `_layer_views`
- views them."""
- ck = _checkpoint(spread)
- core = SimpleNamespace(trunk_moe_backend=backend)
-
- def stack(fn, *tensors):
- return torch.stack([fn(tuple(t[e] for t in tensors), core) for e in range(E)])
-
- def stack1(fn, tensor):
- return torch.stack([fn(tensor[e], core) for e in range(E)])
-
- fc1_w = stack(W._expert_fc1_w, ck.up_w, ck.gate_w)
- fc1_s = stack(W._expert_fc1_s, ck.up_s, ck.gate_s)
- fc2_w = stack1(W._expert_fc2_w, ck.down_w)
- fc2_s = stack1(W._expert_fc2_s, ck.down_s)
- if backend == "CUTEDSL":
- w_dtype, s_dtype = torch.float4_e2m1fn_x2, torch.uint8
- else:
- w_dtype, s_dtype = torch.uint8, torch.float8_e4m3fn
- return SimpleNamespace(
- fc1_w=fc1_w.view(w_dtype),
- fc1_s=fc1_s.view(s_dtype),
- fc2_w=fc2_w.view(w_dtype),
- fc2_s=fc2_s.view(s_dtype),
- scalars=M._moe_layer_scalars(ck.scalars, 0, backend),
- )
-
-
-def _cutedsl_routed(layer, x, ids, weights) -> torch.Tensor:
- """The CuteDSL branch of `_moe_block`, from the quantization to the unpermute."""
- sc = layer.scalars
- tokens = x.shape[0]
- xq, xsf = fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, SF_LINEAR)
- tile_expert, tile_limit, expanded_to_permuted, permuted_to_expanded, _, live_tiles = M.moe_sort(
- token_selected_experts=ids,
- token_final_scales=weights,
- num_experts=E,
- top_k=K,
- local_expert_offset=0,
- local_num_experts=E,
- tile_tokens_dim=M._CUTEDSL_TILE_SIZE,
- )
- h, h_sf = M.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion(
- input=xq.view(torch.float4_e2m1fn_x2),
- weight=layer.fc1_w,
- input_scale=xsf.view(tokens, H // M._SF_VEC),
- weight_scale=layer.fc1_s,
- alpha=sc["alpha1"],
- tile_idx_to_group_idx=tile_expert,
- tile_idx_to_mn_limit=tile_limit,
- permuted_idx_to_expanded_idx=permuted_to_expanded,
- num_non_exiting_tiles=live_tiles,
- global_sf=sc["global_sf"],
- num_experts=E,
- top_k=K,
- num_local_experts=E,
- local_expert_offset=0,
- tile_size=M._CUTEDSL_TILE_SIZE,
- scaling_vector_size=M._SF_VEC,
- activation_type=int(M.ActivationType.Swiglu),
- )
- y = M.cute_dsl_nvfp4_grouped_gemm(
- input=h,
- weight=layer.fc2_w,
- input_scale=h_sf,
- weight_scale=layer.fc2_s,
- alpha=sc["alpha2"],
- tile_idx_to_group_idx=tile_expert,
- num_non_exiting_tiles=live_tiles,
- num_experts=E,
- top_k=K,
- num_local_experts=E,
- local_expert_offset=0,
- tile_size=M._CUTEDSL_TILE_SIZE,
- output_dtype=torch.bfloat16,
- scaling_vector_size=M._SF_VEC,
- )
- return M.moe_unpermute_inplace(
- permuted_input=y,
- output=torch.empty(tokens, H, dtype=torch.bfloat16, device=x.device),
- expanded_idx_to_permuted_idx=expanded_to_permuted,
- topk_scales=weights,
- )
-
-
-def _trtllm_gen_routed(layer, x, ids, weights) -> torch.Tensor:
- """The trtllm-gen branch of `_moe_block`."""
- sc = layer.scalars
- xq, xsf = fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, SF_LINEAR)
- return fp4_block_scale_moe_runner(
- routing_logits=None,
- routing_bias=None,
- hidden_states=xq,
- hidden_states_scale=xsf.view(torch.float8_e4m3fn),
- gemm1_weights=layer.fc1_w,
- gemm1_weights_scale=layer.fc1_s,
- gemm1_bias=None,
- gemm1_alpha=None,
- gemm1_beta=None,
- gemm1_clamp_limit=None,
- gemm2_weights=layer.fc2_w,
- gemm2_weights_scale=layer.fc2_s,
- gemm2_bias=None,
- output1_scale_scalar=sc["output1_scale_scalar"],
- output1_scale_gate_scalar=sc["output1_scale_gate_scalar"],
- output2_scale_scalar=sc["output2_scale_scalar"],
- num_experts=E,
- top_k=K,
- n_group=None,
- topk_group=None,
- intermediate_size=INTER,
- local_expert_offset=0,
- local_num_experts=E,
- routed_scaling_factor=None,
- routing_method_type=M._ROUTING_METHOD_RENORM,
- do_finalize=True,
- act_type=M._ACT_TYPE_SWIGLU,
- topk_weights=weights.to(torch.bfloat16),
- topk_ids=ids,
- )[0]
-
-
-# ── the inputs and the reference ──────────────────────────────────────────
-
-
-@functools.lru_cache(maxsize=None)
-def _problem(case: str):
- """`(x, ids, weights)`: bf16 activations, expert ids `[T, K]` int32 and routing weights
- `[T, K]` fp32 (exact in bf16, which is the dtype trtllm-gen takes them in)."""
- tokens = ONE_HOT_TOKENS if case == "one_hot" else WEIGHTED_TOKENS
- g = torch.Generator(device=DEV).manual_seed(7 if case == "one_hot" else 8)
- x = torch.randn(tokens, H, generator=g, device=DEV).to(torch.bfloat16)
- scores = torch.rand(tokens, E, generator=g, device=DEV)
- if case == "one_hot":
- # The forced expert has the top score, so it sits in slot 0.
- scores[torch.arange(tokens), torch.arange(tokens) % E] = 2.0
- weights = torch.zeros(tokens, K, device=DEV)
- weights[:, 0] = 1.0
- else:
- weights = torch.softmax(torch.randn(tokens, K, generator=g, device=DEV), dim=-1)
- weights = weights.to(torch.bfloat16).float()
- ids = scores.topk(K, dim=-1).indices.to(torch.int32)
- return x, ids, weights
-
-
-def _dequant(codes: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
- """`[R, C/2]` packed e2m1 (element 2j in the low nibble) and `[R, C/16]` e4m3 scales
- -> float64 `[R, C]`."""
- magnitude = torch.tensor(E2M1, dtype=torch.float64, device=codes.device)
- table = torch.cat([magnitude, -magnitude])
- lo, hi = table[(codes & 0xF).long()], table[(codes >> 4).long()]
- values = torch.stack([lo, hi], dim=-1).reshape(codes.shape[0], -1)
- return values * scales.view(torch.float8_e4m3fn).double().repeat_interleave(VEC, dim=1)
-
-
-def _reference(backend: str, case: str, spread: bool) -> torch.Tensor:
- """The routed output in float64, `[T, H]`.
-
- The activation is what the kernels see, the dequantized NVFP4 of the target's own
- quantization. Each expert's weights are dequantized from the checkpoint tensors, times
- its `weight_scale_2`. The SwiGLU output is requantized at `backend`'s global scale."""
- ck, layer = _checkpoint(spread), _layer(backend, spread)
- x, ids, weights = _problem(case)
- sc = layer.scalars
- xq, xsf = fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, SF_LINEAR)
- x_dq = _dequant(xq.view(torch.uint8), xsf.view(x.shape[0], H // VEC)) * ISC1
- if backend == "CUTEDSL":
- g2 = [float(sc["global_sf"])] * E
- else:
- g2 = (1.0 / ck.isc2).tolist()
- out = torch.zeros(x.shape[0], H, dtype=torch.float64, device=DEV)
- for e in range(E):
- token, slot = (ids == e).nonzero(as_tuple=True)
- w = weights[token, slot].double()
- token, w = token[w != 0], w[w != 0]
- if token.numel() == 0:
- continue
- ws2 = ck.ws2_gu[e].double()
- xe = x_dq[token]
- up = xe @ (_dequant(ck.up_w[e], ck.up_s[e]) * ws2).T
- gate = xe @ (_dequant(ck.gate_w[e], ck.gate_s[e]) * ws2).T
- act = up * gate * torch.sigmoid(gate)
- act = quantize_nvfp4(act.float(), g2[e]).double() / g2[e]
- y = act @ (_dequant(ck.down_w[e], ck.down_s[e]) * ck.ws2_dn[e].double()).T
- out.index_add_(0, token, y * w[:, None])
- return out
-
-
-@functools.lru_cache(maxsize=None)
-def _routed(backend: str, case: str, spread: bool) -> torch.Tensor:
- x, ids, weights = _problem(case)
- run = _cutedsl_routed if backend == "CUTEDSL" else _trtllm_gen_routed
- return run(_layer(backend, spread), x, ids, weights).double()
-
-
-def _token_error(out: torch.Tensor, ref: torch.Tensor) -> torch.Tensor:
- """Each token's `|out - ref| / |ref|` (L2 over the hidden width)."""
- return (out - ref).norm(dim=1) / ref.norm(dim=1).clamp_min(1e-12)
-
-
-# ── tests ─────────────────────────────────────────────────────────────────
-
-
-@pytest.mark.parametrize("backend", BACKENDS)
-def test_routed_output_matches_the_float64_reference(backend: str) -> None:
- case = "one_hot"
- _, ids, weights = _problem(case)
- assert torch.unique(ids[weights != 0]).numel() == E, "the routing misses an expert"
- err = _token_error(_routed(backend, case, True), _reference(backend, case, True))
- print(f"[err] {backend} {case}: mean {err.mean():.4e} max {err.max():.4e}")
- assert err.max() <= REFERENCE_TOLERANCE, (err.mean().item(), err.max().item())
-
-
-def test_cutedsl_matches_trtllm_gen_under_a_uniform_input_scale() -> None:
- case = "weighted"
- _, ids, weights = _problem(case)
- assert torch.unique(ids[weights != 0]).numel() == E, "the routing misses an expert"
- err = _token_error(_routed("CUTEDSL", case, False), _routed("TRTLLM", case, False))
- print(f"[err] CUTEDSL vs TRTLLM {case}: mean {err.mean():.4e} max {err.max():.4e}")
- assert err.max() <= BACKEND_TOLERANCE, (err.mean().item(), err.max().item())
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_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
deleted file mode 100644
index 6d6cecf25444..000000000000
--- 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
+++ /dev/null
@@ -1,117 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for 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 uses the shared NVFP4 test helpers.
-"""
-
-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.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,
-)
-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__ = ["../../../.."]
-import _fp4_moe_test_utils as H # noqa: E402
-from _fp4_moe_test_utils import many_expert_layer # noqa: E402,F401 — fixture by name
-from _precision import true_fp32_matmul # noqa: E402
-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 _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, many_expert_layer): # 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 = H.renorm_operands(
- many_expert_layer(), spec, seed, renorm_moe_routing_op
- )
- 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.
- with true_fp32_matmul():
- exp = op.reference(logits, x, k, ref["up"], ref["gate"], ref["down"], g2)
- H._compare(y, exp.to(torch.bfloat16))
-
-
-def test_uncertified_pair_refused(many_expert_layer): # noqa: F811
- spec = op.CELLS[0].spec
- args, _, data, sf, _, logits, _, _, _, scalars = H.renorm_operands(
- many_expert_layer(), spec, 13, renorm_moe_routing_op
- )
- with validating(op), pytest.raises(AssertionError, match="certified Renormalize"):
- op(*_args(args, data, sf, logits, scalars, spec, top_k=9))
-
-
-def test_cuda_graph_capture_and_replay(many_expert_layer): # noqa: F811
- spec = next(c.spec for c in op.CELLS if c.spec["tokens"] == 1)
- args, _, data, sf, _, logits, _, _, _, scalars = H.renorm_operands(
- many_expert_layer(), spec, 43, renorm_moe_routing_op
- )
- call = list(_args(args, data, sf, logits, scalars, spec))
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(warmup):
- for _ in range(3):
- op(*call)
- torch.cuda.current_stream().wait_stream(warmup)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- g_rows, g_w, g_idx = op(*call)
- logits.copy_(torch.randn_like(logits.float()).to(logits.dtype))
- graph.replay()
- torch.cuda.synchronize()
- want = finalized(*call)
- assert torch.equal(finalize(g_rows, g_w, g_idx), want)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py
deleted file mode 100644
index 249a87534aa9..000000000000
--- 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
+++ /dev/null
@@ -1,155 +0,0 @@
-# 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) uses ordinary shared test helpers. Beyond the cells: the refusal
-of an uncertified (num_experts, top_k) pair, and per `DECODE_TACTICS` bucket
-that the pinned tactic is a valid runner config and answers bitwise what the
-AutoTuner path's untuned (runner-default) tactic answers.
-"""
-
-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.fp4_block_scale_moe_renorm_routed import ( # noqa: E501
- ACT_SWIGLU,
- CERTIFIED_PAIRS,
- DECODE_TACTICS,
- DECODE_TACTICS_GEOMETRY,
- ROUTING_RENORMALIZE,
- _ckpt,
- _pinned_tactic,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_block_scale_moe_renorm_routed import ( # noqa: E501
- 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__ = ["../../../.."]
-import _fp4_moe_test_utils as H # noqa: E402
-from _fp4_moe_test_utils import many_expert_layer # noqa: E402,F401 — fixture by name
-from _precision import true_fp32_matmul # noqa: E402
-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 _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, many_expert_layer): # 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) = H.renorm_operands(
- many_expert_layer(), spec, seed, renorm_moe_routing_op
- )
-
- 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.
- with true_fp32_matmul():
- exp = op.reference(logits, x, spec["top_k"], ref["up"], ref["gate"], ref["down"], g2)
- H._compare(y, exp.to(torch.bfloat16))
-
-
-def test_uncertified_pair_refused(many_expert_layer): # noqa: F811
- spec = op.CELLS[0].spec
- assert (spec["num_experts"], 9) not in CERTIFIED_PAIRS
- args, _, data, sf, _, logits, _, _, _, scalars = H.renorm_operands(
- many_expert_layer(), spec, 13, renorm_moe_routing_op
- )
- with validating(op), pytest.raises(AssertionError, match="certified Renormalize"):
- _routed(args, data, sf, logits, scalars, spec, top_k=9)
-
-
-@pytest.mark.parametrize("bucket", sorted(DECODE_TACTICS))
-def test_decode_tactic_valid_and_bitwise(bucket, many_expert_layer): # noqa: F811
- """The pinned tactic of each decode bucket is what a call there runs, is a valid config of
- the runner, and gives bitwise the output of the AutoTuner path outside tuning mode (the
- runner's default tactic for the shape)."""
- spec = _ckpt(bucket)
- num_experts, top_k, hidden, inter = DECODE_TACTICS_GEOMETRY
- assert (spec["num_experts"], spec["top_k"], spec["hidden"], spec["intermediate"]) == (
- num_experts,
- top_k,
- hidden,
- inter,
- )
- # None (the AutoTuner path, with a warning) if the runner does not list it as valid here.
- assert _pinned_tactic(bucket, num_experts, top_k, hidden, inter) == DECODE_TACTICS[bucket]
- args, _, data, sf, _, logits, _, _, _, scalars = H.renorm_operands(
- many_expert_layer(), spec, 7000 + bucket, renorm_moe_routing_op
- )
- y = _routed(args, data, sf, logits, scalars, spec)
- y_tuner = torch.ops.trtllm.fp4_block_scale_moe_runner(
- logits,
- None,
- data,
- sf.view(torch.float8_e4m3fn),
- args["gemm1_weights"],
- args["gemm1_weights_scale"],
- None,
- None,
- None,
- None,
- args["gemm2_weights"],
- args["gemm2_weights_scale"],
- None,
- *scalars,
- num_experts,
- top_k,
- None,
- None,
- inter,
- 0,
- num_experts,
- None,
- ROUTING_RENORMALIZE,
- True,
- ACT_SWIGLU,
- )[0]
- differing = (y.view(torch.int16) != y_tuner.view(torch.int16)).any(dim=1).sum().item()
- assert differing == 0, f"{differing} of {bucket} rows differ from the AutoTuner path"
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
deleted file mode 100644
index b2d77dcf4336..000000000000
--- 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
+++ /dev/null
@@ -1,195 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the fp8_block_scale_moe_runner catalog entry.
-
-The entry owns the contract: `CELLS`, `reference`, `is_valid`. This
-file turns a cell into operands and holds the claims no cell carries: a warm
-autotuned tactic is bitwise the cold fallback the cells run, and the
-contiguity guard.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp8_block_scale_moe_runner import ( # noqa: E501
- BLOCK,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp8_block_scale_moe_runner import ( # noqa: E501
- fp8_block_scale_moe_runner as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "fp8_block_scale_moe_runner requires a CUDA device"
-
-
-def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
- """16 ulp element-wise, 0.75 ulp RMS, both in bf16 ulp of each row's max.
-
- Both statistics normalize each element by its reference row's largest
- magnitude: `d = (out - ref) / (rowmax * 2^-8)`, gated on `max|d|` and
- on `sqrt(mean(d^2))` -- the RMS of the row-normalized deviation, not
- `assert_within_ulp`'s ratio of norms, which reads higher on the same
- data (a row's max sits several sigma above its RMS). Kernel and
- reference consume bit-identical operands and model both
- requantizations, so they differ by accumulation order and by FC1 sums
- that land on an e4m3 rounding boundary.
-
- Measured worst over the cells: 6.65 element-wise (T=4096; the max over
- T*H samples grows with T -- 0.91 at T=1) and 0.306 RMS (H=I=512), the
- RMS sitting at 0.17-0.31 everywhere; earlier draws reached 10.03 and
- 0.342. Wrong references at the checkpoint geometry, RMS: gate/up swap
- 65, transposed activation scales 43, strided hidden_states 94, no
- FC1-output requant 3.5, neither requant 3.1, no activation requant 1.58
- -- the closest.
- """
- assert out.dtype == ref.dtype == torch.bfloat16, (out.dtype, ref.dtype)
- assert out.shape == ref.shape, (out.shape, ref.shape)
- o, r = out.float(), ref.float()
- rowmax = r.abs().amax(dim=-1, keepdim=True).clamp(min=1e-6)
- d = (o - r) / (rowmax * 2.0**-8)
- elt, rms = d.abs().max().item(), d.pow(2).mean().sqrt().item()
- assert elt <= 16.0, f"{elt:.2f} ulp element-wise > 16"
- assert rms <= 0.75, f"{rms:.3f} ulp RMS > 0.75"
-
-
-DEV = "cuda"
-ULP = 2.0**-8
-CKPT = dict(hidden=2560, inter=640, num_experts=512, top_k=10)
-
-# ── operands ──────────────────────────────────────────────────────────────
-
-_WEIGHTS: dict = {}
-
-
-def _weights(hidden: int, inter: int, experts: int, seed: int):
- """`(w1, w1s, w2, w2s)` drawn directly in e4m3 / fp32, so every value is exact.
-
- Cached: the checkpoint stack is 2.5 GB and every checkpoint-geometry cell
- reads the same one. Drawn 32 experts at a time to cap the fp32 transient.
- """
- key = (hidden, inter, experts, seed)
- if key not in _WEIGHTS:
- g = torch.Generator(device=DEV).manual_seed(seed)
- w1 = torch.empty(experts, 2 * inter, hidden, dtype=torch.float8_e4m3fn, device=DEV)
- w2 = torch.empty(experts, hidden, inter, dtype=torch.float8_e4m3fn, device=DEV)
- for lo in range(0, experts, 32):
- hi = min(lo + 32, experts)
- w1[lo:hi] = torch.randn(hi - lo, 2 * inter, hidden, device=DEV, generator=g)
- w2[lo:hi] = torch.randn(hi - lo, hidden, inter, device=DEV, generator=g)
- w1s = 0.01 + 0.02 * torch.rand(
- experts, 2 * inter // BLOCK, hidden // BLOCK, device=DEV, generator=g
- )
- w2s = 0.01 + 0.02 * torch.rand(
- experts, hidden // BLOCK, inter // BLOCK, device=DEV, generator=g
- )
- _WEIGHTS[key] = (w1, w1s, w2, w2s)
- return _WEIGHTS[key]
-
-
-def _args(
- tokens, hidden, inter, num_experts, top_k, local_expert_offset=0, local_num_experts=None, seed=0
-):
- """The op's arguments by name, pre-routed, every expert local."""
- assert local_expert_offset == 0 and local_num_experts in (None, num_experts)
- 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.
- hs_scale = 0.5 + torch.rand(hidden // BLOCK, tokens, device=DEV, generator=g)
- logits = torch.randn(tokens, num_experts, device=DEV, generator=g)
- weights, ids = torch.topk(torch.softmax(logits, dim=-1), top_k, dim=-1)
- ids = ids.to(torch.int32)
- return dict(
- routing_logits=None,
- routing_bias=None,
- hidden_states=hs,
- hidden_states_scale=hs_scale,
- 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=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=None,
- )
-
-
-def _ulp(out: torch.Tensor, ref: torch.Tensor):
- """(max, RMS) deviation in bf16 ulp of the row max -- `_compare`'s statistics."""
- o, r = out.float(), ref.float()
- d = (o - r) / (r.abs().amax(dim=-1, keepdim=True).clamp(min=1e-6) * ULP)
- return d.abs().max().item(), d.pow(2).mean().sqrt().item()
-
-
-def _seed(cell) -> int:
- return zlib.crc32(cell.why.encode())
-
-
-# ── the certified cells ───────────────────────────────────────────────────
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- args = _args(**cell.spec, seed=_seed(cell))
- out = op(**args)
- assert out.shape == args["hidden_states"].shape and out.dtype == torch.bfloat16
- ref = op.reference(**args)
- elt, rms = _ulp(out, ref)
- print(f" {cell.why[:44]}: {elt:.2f} ulp element-wise, {rms:.3f} ulp RMS")
- _compare(out, ref)
-
-
-# ── what a cell cannot carry ──────────────────────────────────────────────
-
-
-def test_autotuned_tactic_matches_the_fallback_tactic():
- """A warm profiling cache changes the tactic, not the result."""
- from tensorrt_llm._torch.autotuner import AutoTuner, autotune
-
- a = _args(256, **CKPT, seed=42)
- AutoTuner.get().clear_cache()
- try:
- cold = op(**a)
- with autotune():
- op(**a)
- entries = len(AutoTuner.get().profiling_cache)
- warm = op(**a)
- print(f" warming wrote {entries} cache entries")
- assert entries > 0, "warming profiled nothing -- the comparison would be vacuous"
- assert torch.equal(warm, cold), "a tuned tactic changed the result"
- _compare(warm, op.reference(**a))
- finally:
- AutoTuner.get().clear_cache()
-
-
-# ── the guards ────────────────────────────────────────────────────────────
-
-
-def _strided(t: torch.Tensor) -> torch.Tensor:
- wide = torch.zeros(*t.shape[:-1], 2 * t.shape[-1], device=DEV, dtype=t.dtype)
- wide[..., ::2] = t
- view = wide[..., ::2]
- assert not view.is_contiguous()
- return view
-
-
-def test_guard_refuses_a_strided_hidden_states():
- """What `is_valid` exists for: the raw op reads a strided operand as dense."""
- a = _args(16, 256, 128, 8, 2, seed=11)
- with validating(op), pytest.raises(AssertionError, match="must be contiguous"):
- op(**dict(a, hidden_states=_strided(a["hidden_states"])))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py
deleted file mode 100644
index 035ca9c348ce..000000000000
--- 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
+++ /dev/null
@@ -1,83 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the fused_sigmoid_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: the empty
-batch and the loud rejection of a wrong-length gate.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fused_sigmoid_gate_mul_add import ( # noqa: E501
- fused_sigmoid_gate_mul_add as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "fused_sigmoid_gate_mul_add requires a CUDA device"
-
-DEV = "cuda"
-HIDDEN = 2560
-
-
-def _operands(spec, seed):
- """`(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)
-
- final, shared = rand(t, h), rand(t, h)
- gate = rand(t, 1, dtype=spec["gate_dtype"])
- if spec["gate_layout"] == "strided":
- wide = rand(t, 4, dtype=spec["gate_dtype"])
- wide[:, 2:3] = gate
- gate = wide[:, 2:3]
- assert not gate.is_contiguous()
- else:
- assert spec["gate_layout"] == "column", spec["gate_layout"]
- return final, gate, shared
-
-
-# ── the certified cells ───────────────────────────────────────────────────
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- final, gate, shared = _operands(cell.spec, zlib.crc32(cell.why.encode()))
- before = {k: v.clone() for k, v in dict(final=final, gate=gate, shared=shared).items()}
- ref = op.reference(before["final"], gate, shared)
-
- out = op(final, gate, shared)
-
- 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"])
-
-
-# ── what a cell cannot carry ──────────────────────────────────────────────
-
-
-def test_zero_tokens_is_a_no_op():
- final = torch.empty(0, HIDDEN, dtype=torch.bfloat16, device=DEV)
- shared = torch.empty(0, HIDDEN, dtype=torch.bfloat16, device=DEV)
- gate = torch.empty(0, 1, dtype=torch.bfloat16, device=DEV)
- returned = op(final, gate, shared)
- assert returned.shape == (0, HIDDEN) and returned.data_ptr() == final.data_ptr()
-
-
-def test_a_gate_of_the_wrong_length_raises() -> None:
- """Metadata violations are loud, which is why the entry has no guard."""
- bf = dict(dtype=torch.bfloat16, device=DEV)
- final, shared = torch.randn(4, 128, **bf), torch.randn(4, 128, **bf)
- with pytest.raises(AssertionError):
- op(final, torch.randn(3, 1, **bf), shared)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_sort.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_sort.py
deleted file mode 100644
index 8d97d4b0dd4f..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_sort.py
+++ /dev/null
@@ -1,215 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the moe_sort entry of the Qwen3.8-Flash-Next target.
-
-The entry owns the contract. This file turns a cell into expert ids and holds
-the claims no cell carries: CUDA-graph replay at the largest captured row
-count, and the contiguity guard.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.moe_sort import ( # noqa: E501
- moe_sort as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "moe_sort requires a CUDA device"
-
-DEV = "cuda"
-E, K = 512, 10
-
-_NAMES = (
- "tile_idx_to_expert_idx",
- "tile_idx_to_mn_limit",
- "expanded_idx_to_permuted_idx",
- "permuted_idx_to_expanded_idx",
- "total_num_padded_tokens",
- "num_non_exiting_tiles",
-)
-
-
-def _compare(out: list[torch.Tensor], ref: list[torch.Tensor]) -> None:
- """Integer outputs, so exact -- up to the one freedom the op has.
-
- Counts and the tile arrays' written prefix are bitwise. The order of an
- expert's rows is the kernel's atomic order (it varies between identical calls
- from T = 5), so `permuted_idx_to_expanded_idx` is gated per expert run as a
- set: each real row's value sorted within its run must equal the reference's
- run. `expanded_idx_to_permuted_idx` must then be the exact inverse: -1
- exactly where the expert is not local, and otherwise a real row of the right
- run that maps back. Padding rows and tiles past `num_non_exiting_tiles` are
- not written by the op and are not compared.
- """
- assert len(out) == len(ref) == len(_NAMES), (len(out), len(ref))
- for name, o, r in zip(_NAMES, out, ref):
- assert o.dtype == r.dtype == torch.int32, (name, o.dtype, r.dtype)
- assert o.shape == r.shape, (name, tuple(o.shape), tuple(r.shape))
- t2e, t2mn, e2p, p2e, total, n_tiles = out
- r_t2e, r_t2mn, r_e2p, r_p2e, r_total, r_n_tiles = ref
- assert torch.equal(n_tiles.cpu(), r_n_tiles.cpu()), (n_tiles, r_n_tiles)
- assert torch.equal(total.cpu(), r_total.cpu()), (total, r_total)
- n = int(r_n_tiles)
- assert torch.equal(t2e[:n], r_t2e[:n]), "tile_idx_to_expert_idx differs"
- assert torch.equal(t2mn[:n], r_t2mn[:n]), "tile_idx_to_mn_limit differs"
-
- tile = p2e.numel() // t2e.numel()
- num_expanded = e2p.numel()
- real = r_p2e >= 0
- rows = torch.nonzero(real).squeeze(1)
- run = r_t2e[rows // tile].long()
- got = p2e[real].long()
- assert bool(((got >= 0) & (got < num_expanded)).all()), (
- "permuted_idx_to_expanded_idx holds an out-of-range value on a real row"
- )
- assert torch.equal(
- torch.sort(run * num_expanded + got).values, run * num_expanded + r_p2e[real].long()
- ), "an expert's run of permuted_idx_to_expanded_idx holds the wrong pairs"
-
- flat, r_flat = e2p.reshape(-1).long(), r_e2p.reshape(-1).long()
- not_local = r_flat < 0
- assert bool((flat[not_local] == -1).all()), "a non-local pair is not -1"
- p = flat[~not_local]
- assert bool(((p >= 0) & (p < p2e.numel())).all()), "a local pair points outside the rows"
- assert bool(real[p].all()), "a local pair points at a padding row"
- assert torch.equal(p2e[p].long(), torch.arange(num_expanded, device=p.device)[~not_local]), (
- "expanded_idx_to_permuted_idx is not the inverse of permuted_idx_to_expanded_idx"
- )
-
-
-def _ids(dist, tokens, tile, seed, num_experts=E, top_k=K):
- """Expert ids, `[tokens, top_k]` int32, every row naming distinct experts."""
- g = torch.Generator(device=DEV).manual_seed(seed)
- if dist == "router":
- scores = torch.rand(tokens, num_experts, device=DEV, generator=g)
- return scores.topk(top_k, dim=-1).indices.to(torch.int32)
- if dist == "tile_edges":
- edge = (7, 300, 511)
- scores = torch.rand(tokens, num_experts, device=DEV, generator=g)
- scores[:, list(edge)] = -1.0
- ids = scores.topk(top_k, dim=-1).indices.to(torch.int32)
- for slot, (expert, rows) in enumerate(zip(edge, (tile - 1, tile, tile + 1))):
- ids[:rows, slot] = expert
- return ids
- raise ValueError(dist)
-
-
-def _args(ids, tile, seed=0, scales_dtype=torch.float32):
- """The CuteDSL call site's keyword arguments at the checkpoint's values."""
- g = torch.Generator(device=DEV).manual_seed(seed + 1)
- scales = torch.rand(ids.shape, device=DEV, generator=g).to(scales_dtype)
- return dict(
- token_selected_experts=ids,
- token_final_scales=scales,
- num_experts=E,
- top_k=K,
- local_expert_offset=0,
- local_num_experts=E,
- tile_tokens_dim=tile,
- )
-
-
-def _check(args):
- """Drive the op and gate it through the entry's own reference and `_compare`."""
- out = op(**args)
- ref = op.reference(**args)
- _compare(out, ref)
- return out, ref
-
-
-def _cell_id(cell):
- s = cell.spec
- return f"t{s['tile_tokens_dim']}-T{s['tokens']}-{s['dist']}"
-
-
-# ── the certified cells ───────────────────────────────────────────────────
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[_cell_id(c) for c in op.CELLS])
-def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- s = cell.spec
- seed = zlib.crc32(cell.why.encode())
- ids = _ids(s["dist"], s["tokens"], s["tile_tokens_dim"], seed)
- args = _args(ids, s["tile_tokens_dim"], seed, s["scales_dtype"])
- assert (s["num_experts"], s["top_k"]) == (args["num_experts"], args["top_k"])
- assert (s["local_expert_offset"], s["local_num_experts"]) == (0, E)
- ids_before = ids.clone()
- scales_before = args["token_final_scales"].clone()
- out, ref = _check(args)
- assert torch.equal(ids, ids_before), "token_selected_experts was written"
- assert torch.equal(args["token_final_scales"], scales_before), "token_final_scales was written"
- if s["dist"] == "tile_edges":
- tile = s["tile_tokens_dim"]
- n = int(ref[5])
- limits = out[1][:n].long()
- starts = torch.arange(n, device=DEV) * tile
- fill = (limits - starts).tolist()
- experts = out[0][:n].tolist()
- per_expert = {e: [] for e in (7, 300, 511)}
- for e, f in zip(experts, fill):
- if e in per_expert:
- per_expert[e].append(f)
- assert per_expert == {7: [tile - 1], 300: [tile], 511: [tile, 1]}, per_expert
-
-
-# ── CUDA graphs: captured row counts, on new ids ──────────────────────────
-
-#: The largest decode capture: 32 requests at MTP verify width 4, cluster kernel.
-_CAPTURED_ROWS = 128
-
-
-def _capture(args):
- side = torch.cuda.Stream()
- side.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(side):
- for _ in range(2):
- op(**args)
- torch.cuda.current_stream().wait_stream(side)
- torch.cuda.synchronize()
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- static_out = op(**args)
- return graph, static_out
-
-
-def test_captured_rows_replay_on_new_ids(certifying): # noqa: F811
- """The graph is replayed on two fresh id sets and gated against the reference,
- and against an eager call on the same ids: counts and tile arrays bitwise. At
- the target's tile, 128."""
- certifying(op)
- tile, rows = 128, _CAPTURED_ROWS
- args = _args(_ids("router", rows, tile, seed=rows), tile, seed=rows)
- graph, static_out = _capture(args)
- for trial in (1, 2):
- fresh = _ids("router", rows, tile, seed=1000 * trial + rows)
- args["token_selected_experts"].copy_(fresh)
- graph.replay()
- eager = op(**dict(args, token_selected_experts=fresh.clone()))
- ref = op.reference(**args)
- _compare(static_out, ref)
- _compare(eager, ref)
- n = int(ref[5])
- for i in (0, 1):
- assert torch.equal(static_out[i][:n], eager[i][:n]), (trial, i)
- for i in (4, 5):
- assert torch.equal(static_out[i], eager[i]), (trial, i)
- del graph, static_out
- torch.cuda.synchronize()
-
-
-# ── the guard ─────────────────────────────────────────────────────────────
-
-
-def test_guard_refuses_noncontiguous_ids():
- """What `is_valid` exists for: the op reads a strided view as the dense buffer
- that starts at its data pointer, and answers without raising."""
- wide = _ids("router", 300, 128, seed=51, top_k=16)
- sc = torch.rand(300, K, device=DEV)
- with validating(op), pytest.raises(AssertionError, match="contiguous"):
- op(wide[:, :K], sc, E, K, 0, E, 128)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_unpermute_inplace.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_unpermute_inplace.py
deleted file mode 100644
index 1ff8171c9b9e..000000000000
--- 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
+++ /dev/null
@@ -1,268 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the moe_unpermute_inplace entry of the Qwen3.8-Flash-Next target.
-
-The entry owns the contract. This file lays a cell's FC2 rows out the way the
-CuteDSL FC2 hands them over -- placed by moe_sort, every row the op must not read
-filled with NaN, the output pre-filled with NaN -- and holds the claims no cell
-carries: CUDA-graph replay at the largest captured row count, bitwise equal to
-eager, and the contiguity guard.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.moe_unpermute_inplace import ( # noqa: E501
- moe_unpermute_inplace as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "moe_unpermute_inplace requires a CUDA device"
-
-DEV = "cuda"
-E, K, H = 512, 10, 2560
-NAN = float("nan")
-
-#: The largest decode capture: 32 requests at MTP verify width 4.
-_CAPTURED_ROWS = 128
-
-
-def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
- """Bitwise.
-
- `reference` is the computation the op performs, in its order and with its
- roundings, so no band is needed and none is given. The gate sees the
- accumulation order: the same sum taken in reverse slot order, with each
- product rounded before its add, as the float64 sum rounded once, or
- accumulated in bf16, differs in many elements and is rejected. A failure in a
- few isolated one-ulp elements points at the kernel's rounding (a lost
- multiply-add contraction, say); a wrong row or scale moves whole rows.
- """
- assert out.dtype == ref.dtype, (out.dtype, ref.dtype)
- assert out.shape == ref.shape, (tuple(out.shape), tuple(ref.shape))
- bits = {2: torch.int16, 4: torch.int32}[out.element_size()]
- differ = out.view(bits) != ref.view(bits)
- n = int(differ.sum())
- if n:
- first = torch.nonzero(differ)[0].tolist()
- raise AssertionError(
- f"{n} of {out.numel()} elements differ bitwise from the slot-order fp32 "
- f"sum; first at {first}: {out[tuple(first)].item()} vs "
- f"{ref[tuple(first)].item()}"
- )
-
-
-def _spec(tokens, tile_tokens_dim=128, values="normal"):
- """The fields `_problem` reads; the same names as a cell's spec."""
- return dict(tokens=tokens, tile_tokens_dim=tile_tokens_dim, values=values)
-
-
-def _ids(tokens: int, seed: int) -> torch.Tensor:
- """Expert ids, `[tokens, K]` int32: each row's top-10 of uniform scores."""
- g = torch.Generator(device=DEV).manual_seed(seed)
- scores = torch.rand(tokens, E, device=DEV, generator=g)
- return scores.topk(K, dim=-1).indices.to(torch.int32)
-
-
-def _scales(tokens: int, seed: int) -> torch.Tensor:
- """fp32 routing weights, a softmax over each row's ten slots."""
- g = torch.Generator(device=DEV).manual_seed(seed + 1)
- return torch.softmax(torch.randn(tokens, K, device=DEV, generator=g), dim=-1)
-
-
-def _rows(tokens: int, seed: int, values: str = "normal") -> torch.Tensor:
- """The FC2 output row of each (token, slot) pair, `[tokens, K, H]` bf16."""
- g = torch.Generator(device=DEV).manual_seed(seed + 2)
- v = torch.randn(tokens, K, H, device=DEV, generator=g)
- if values == "wide":
- v = v * torch.exp2(torch.randint(-30, 31, v.shape, device=DEV, generator=g).float())
- elif values != "normal":
- raise ValueError(values)
- return v.to(torch.bfloat16)
-
-
-def _sort(ids: torch.Tensor, scales: torch.Tensor, tile: int):
- """moe_sort as the CuteDSL MoE calls it: (expanded_idx_to_permuted_idx, padded rows)."""
- out = torch.ops.trtllm.moe_sort(
- token_selected_experts=ids,
- token_final_scales=scales,
- num_experts=E,
- top_k=K,
- local_expert_offset=0,
- local_num_experts=E,
- tile_tokens_dim=tile,
- )
- return out[2], out[3].numel()
-
-
-def _place(rows: torch.Tensor, e2p: torch.Tensor, num_rows: int) -> torch.Tensor:
- """The FC2 output as the op receives it: each live pair's row at its permuted
- index, NaN in every row the op has no business reading."""
- perm = torch.full((num_rows, H), NAN, dtype=rows.dtype, device=DEV)
- flat = e2p.reshape(-1)
- live = flat >= 0
- perm[flat[live].long()] = rows.reshape(-1, H)[live]
- return perm
-
-
-def _output(tokens: int) -> torch.Tensor:
- return torch.full((tokens, H), NAN, dtype=torch.bfloat16, device=DEV)
-
-
-def _problem(spec: dict, seed: int):
- """A cell's arguments, plus the per-pair rows and ids they were built from."""
- tokens = spec["tokens"]
- ids = _ids(tokens, seed)
- scales = _scales(tokens, seed)
- e2p, num_rows = _sort(ids, scales, spec["tile_tokens_dim"])
- rows = _rows(tokens, seed, spec["values"])
- args = dict(
- permuted_input=_place(rows, e2p, num_rows),
- output=_output(tokens),
- expanded_idx_to_permuted_idx=e2p,
- topk_scales=scales,
- )
- return args
-
-
-def _bits(t: torch.Tensor) -> torch.Tensor:
- return t.contiguous().view(torch.uint8)
-
-
-def _same_bits(a: torch.Tensor, b: torch.Tensor) -> bool:
- return a.shape == b.shape and a.dtype == b.dtype and torch.equal(_bits(a), _bits(b))
-
-
-def _cell_id(cell):
- s = cell.spec
- return f"t{s['tile_tokens_dim']}-T{s['tokens']}-{s['values']}"
-
-
-# ── the certified cells ───────────────────────────────────────────────────
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[_cell_id(c) for c in op.CELLS])
-def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- s = cell.spec
- assert (s["num_experts"], s["top_k"], s["hidden_size"]) == (E, K, H)
- assert (s["dtype"], s["scales_dtype"]) == (torch.bfloat16, torch.float32)
- args = _problem(s, zlib.crc32(cell.why.encode()))
- before = {
- n: args[n].clone()
- for n in ("permuted_input", "expanded_idx_to_permuted_idx", "topk_scales")
- }
- out = op(**args)
- assert out is args["output"], "the op must return the tensor it wrote"
- _compare(out, op.reference(**args))
- for name, snapshot in before.items():
- assert _same_bits(args[name], snapshot), f"{name} was written"
-
-
-# ── CUDA graphs: captured row counts, on new data ─────────────────────────
-
-
-def _capture(fn):
- side = torch.cuda.Stream()
- side.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(side):
- for _ in range(2):
- fn()
- torch.cuda.current_stream().wait_stream(side)
- torch.cuda.synchronize()
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- static = fn()
- return graph, static
-
-
-def test_captured_rows_replay_on_new_data(certifying): # noqa: F811
- """The sort, the FC2 rows' placement and this op are captured together, as the
- target runs them; the graph is replayed on two fresh sets of ids, scales and
- rows. The permuted buffer is not cleared between replays, so the rows the new
- sort does not use hold the previous placement's rows. Each replay is gated
- against the reference, and is bitwise equal to an eager call on the same data
- through its own sort. At the target's tile, 128.
- """
- certifying(op)
- tile = 128
- stale_rows = 0
- tokens = _CAPTURED_ROWS
- ids = _ids(tokens, seed=tokens)
- scales = _scales(tokens, seed=tokens)
- rows = _rows(tokens, seed=tokens)
- num_rows = _sort(ids, scales, tile)[1]
- perm = torch.full((num_rows, H), NAN, dtype=torch.bfloat16, device=DEV)
- out = _output(tokens)
-
- def step(ids=ids, scales=scales, rows=rows, perm=perm, out=out, tile=tile):
- e2p = _sort(ids, scales, tile)[0]
- perm.index_copy_(0, e2p.reshape(-1).long(), rows.reshape(-1, H))
- op(
- permuted_input=perm,
- output=out,
- expanded_idx_to_permuted_idx=e2p,
- topk_scales=scales,
- )
- return e2p
-
- graph, e2p = _capture(step)
- for trial in (1, 2):
- seed = 1000 * trial + tokens
- ids.copy_(_ids(tokens, seed))
- scales.copy_(_scales(tokens, seed))
- rows.copy_(_rows(tokens, seed))
- out.fill_(NAN)
- graph.replay()
- args = dict(
- permuted_input=perm,
- output=out,
- expanded_idx_to_permuted_idx=e2p,
- topk_scales=scales,
- )
- _compare(out, op.reference(**args))
- used = torch.zeros(num_rows, dtype=torch.bool, device=DEV)
- used[e2p.reshape(-1).long()] = True
- stale_rows += int((~used & ~perm.isnan().any(1)).sum())
-
- e2p_eager, n = _sort(ids.clone(), scales.clone(), tile)
- eager = op(
- permuted_input=_place(rows.clone(), e2p_eager, n),
- output=_output(tokens),
- expanded_idx_to_permuted_idx=e2p_eager,
- topk_scales=scales.clone(),
- )
- assert _same_bits(out, eager), trial
- del graph, e2p
- torch.cuda.synchronize()
- assert stale_rows > 0, "no replay ran over a stale row: the stale-row claim is untested"
-
-
-# ── the guards ────────────────────────────────────────────────────────────
-
-
-def _strided(t: torch.Tensor, pad: int) -> torch.Tensor:
- """The same values in a view whose rows are `pad` elements wider apart."""
- wide = torch.zeros(t.shape[0], t.shape[1] + pad, dtype=t.dtype, device=t.device)
- wide[:, : t.shape[1]] = t
- return wide[:, : t.shape[1]]
-
-
-def test_noncontiguous_operand_answers_wrongly_and_is_refused():
- """The FC2 rows the op reads; the guard checks all four operands the same way."""
- name = "permuted_input"
- args = _problem(_spec(300), seed=61)
- ref = op.reference(**args)
- bad = dict(args, **{name: _strided(args[name], 64)})
- out = op(**bad)
- with pytest.raises(AssertionError, match="differ bitwise"):
- _compare(out.contiguous(), ref)
- with validating(op):
- with pytest.raises(AssertionError, match=f"{name} must be contiguous"):
- op(**bad)
- out = op(**dict(args, output=_output(300)))
- _compare(out, ref)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_mtp_moe_nvfp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_mtp_moe_nvfp4.py
deleted file mode 100644
index 6e47da43f2fe..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_mtp_moe_nvfp4.py
+++ /dev/null
@@ -1,142 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the NVFP4 copy of the draft layer's routed experts
-(`TRTLLM_MODELING_V2_MTP_MOE_NVFP4`).
-
-`weights.mtp_experts_to_nvfp4` requantizes FP8 128x128 block-scaled experts to NVFP4 in the
-trunk's trtllm-gen layout, `modeling._mtp_moe_nvfp4` adds the runner's scalars and
-`modeling._mtp_routed_nvfp4` runs them. Checked against the FP8 block-scale path the draft
-layer runs by default, on the same random experts, activations and router: the NVFP4 output
-must stay inside a relative-error band of the FP8 one, keep its magnitude (a wrong runner
-scalar moves the norm, which unbiased quantization noise does not), and leave no FP8 copy
-behind. A layout error is O(1) relative, far outside the band.
-"""
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501
- modeling as m,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp8_block_scale_moe_runner import ( # noqa: E501
- fp8_block_scale_moe_runner,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.renorm_moe_routing_op import ( # noqa: E501
- renorm_moe_routing_op,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.quantization.fp8_quantize_1x128 import ( # noqa: E501
- fp8_quantize_1x128,
-)
-
-assert torch.cuda.is_available(), "the NVFP4 draft MoE requires a CUDA device"
-
-DEV = "cuda"
-E, H, INTER, K = 512, 2560, 640, 10
-BLK = 128
-#: A trunk calibration to derive the static draft scales from: the checkpoint's last trunk layer
-#: (FC1 input, FC2 input amax), the layer that feeds the draft.
-TRUNK_AMAX = (11.1875, 57.75)
-#: Relative RMS of NVFP4 - FP8 over the FP8 output, the floor on each row's cosine, and the
-#: band on the output norm ratio. Measured here: 0.239-0.251 RMS, row cosine >= 0.961, at every
-#: row count; on the checkpoint's real draft experts 0.222-0.231 and >= 0.964 (bench). That is
-#: the NVFP4 chain's own noise -- four e2m1 roundings (activation, FC1 weight, FC1 output, FC2
-#: weight) at ~0.1 each, compounded by the SwiGLU product -- not a bias: it is flat across
-#: activation scales 0.5-4x, and FP8 sits at 0.07 from the fp32 reference of the same weights.
-REL_RMS_BAND = 0.3
-COS_FLOOR = 0.95
-NORM_RATIO_BAND = 0.1
-
-
-def _fp8_blocks(rows: int, cols: int, gen: torch.Generator):
- """`[E, rows, cols]` e4m3 codes and `[E, rows/128, cols/128]` fp32 scales of a random
- expert stack whose 128x128 blocks span 16x in magnitude, as trained experts' do."""
- w = torch.randn(E, rows, cols, generator=gen, device=DEV) * cols**-0.5
- spread = torch.exp2(
- torch.randint(-2, 3, (E, rows // BLK, 1, cols // BLK, 1), generator=gen, device=DEV)
- ).float()
- blocks = w.view(E, rows // BLK, BLK, cols // BLK, BLK) * spread
- scale = blocks.abs().amax(dim=(2, 4), keepdim=True) / 448.0
- codes = (blocks / scale).to(torch.float8_e4m3fn).view(E, rows, cols)
- return codes, scale.view(E, rows // BLK, cols // BLK).contiguous()
-
-
-@pytest.fixture(scope="module")
-def experts():
- gen = torch.Generator(device=DEV).manual_seed(20261022)
- fc1_w, fc1_s = _fp8_blocks(2 * INTER, H, gen)
- fc2_w, fc2_s = _fp8_blocks(H, INTER, gen)
- router = (torch.randn(E, H, generator=gen, device=DEV) * H**-0.5).to(torch.bfloat16)
- fp8 = SimpleNamespace(fc1_w=fc1_w, fc1_s=fc1_s, fc2_w=fc2_w, fc2_s=fc2_s, router=router.t())
- p = "l1_mtp_"
- w = {p + k: getattr(fp8, k).clone() for k in ("fc1_w", "fc1_s", "fc2_w", "fc2_s")}
- w["l0_e_isc1"] = torch.full((E,), TRUNK_AMAX[0] / 2688.0, device=DEV)
- w["l0_e_isc2"] = torch.full((E,), TRUNK_AMAX[1] / 2688.0, device=DEV)
- core = SimpleNamespace(
- w=w, mtp_layer_idx=1, num_layers=1, num_experts=E, topk=K, moe_inter=INTER
- )
- fields = m._mtp_moe_nvfp4(core)
- nvfp4 = SimpleNamespace(router=fp8.router, **fields)
- return core, fp8, nvfp4
-
-
-def _fp8_routed(fp8, x):
- """The FP8 block-scale draft path, call for call as `_mtp_moe_block` makes it."""
- logits = m.cublas_mm(x, fp8.router)
- ids, wts = renorm_moe_routing_op(logits, topk=K, output_dtype=torch.bfloat16)
- xq, xsf = fp8_quantize_1x128(x, use_ue8m0=False)
- return fp8_block_scale_moe_runner(
- routing_logits=None,
- routing_bias=None,
- hidden_states=xq,
- hidden_states_scale=xsf,
- gemm1_weights=fp8.fc1_w,
- gemm1_weights_scale=fp8.fc1_s,
- gemm2_weights=fp8.fc2_w,
- gemm2_weights_scale=fp8.fc2_s,
- num_experts=E,
- top_k=K,
- num_fused_shared_experts=None,
- n_group=None,
- topk_group=None,
- intermediate_size=INTER,
- local_expert_offset=0,
- local_num_experts=E,
- routed_scaling_factor=None,
- routing_method_type=m._ROUTING_METHOD_RENORM,
- topk_weights=wts,
- topk_ids=ids,
- act_type=m._ACT_TYPE_SWIGLU,
- gemm1_clamp_limit=None,
- )
-
-
-def test_fp8_copies_freed_and_layout(experts):
- core, _, nvfp4 = experts
- for k in ("fc1_w", "fc1_s", "fc2_w", "fc2_s"):
- assert core.w[f"l1_mtp_{k}"].numel() == 0, k
- assert nvfp4.fc1_w.shape == (E, 2 * INTER, H // 2) and nvfp4.fc1_w.dtype == torch.uint8
- assert nvfp4.fc1_s.shape == (E, 2 * INTER, H // 16) and nvfp4.fc1_s.dtype == torch.float8_e4m3fn
- assert nvfp4.fc2_w.shape == (E, H, INTER // 2) and nvfp4.fc2_s.shape == (E, H, INTER // 16)
- for s in (nvfp4.out1_scale, nvfp4.out1_gate_scale, nvfp4.out2_scale):
- assert s.shape == (E,) and s.dtype == torch.float32 and torch.isfinite(s).all()
- assert nvfp4.moe_gscale.shape == (1,)
-
-
-@pytest.mark.parametrize("rows", [1, 4, 32, 128, 1024])
-def test_nvfp4_tracks_fp8(experts, rows):
- core, fp8, nvfp4 = experts
- gen = torch.Generator(device=DEV).manual_seed(rows)
- x = torch.randn(rows, H, generator=gen, device=DEV).to(torch.bfloat16)
- ref = _fp8_routed(fp8, x).float()
- out = m._mtp_routed_nvfp4(core, nvfp4, x)
- assert out.shape == (rows, H) and out.dtype == torch.bfloat16
- out = out.float()
- rel = ((out - ref).norm() / ref.norm()).item()
- cos = torch.nn.functional.cosine_similarity(out, ref, dim=1).min().item()
- ratio = (out.norm() / ref.norm()).item()
- print(f"rows={rows}: rel RMS {rel:.4f}, min row cosine {cos:.4f}, norm ratio {ratio:.4f}")
- assert rel < REL_RMS_BAND, rel
- assert cos > COS_FLOOR, cos
- assert abs(ratio - 1.0) < NORM_RATIO_BAND, ratio
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/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
deleted file mode 100644
index 5c7fe656dd28..000000000000
--- 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
+++ /dev/null
@@ -1,125 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the renorm_moe_routing_op catalog entry.
-
-The entry owns the contract. This file turns a cell into logits and holds the
-claims no cell carries: equivalence to HF's `norm_topk_prob` routing, and the
-contiguity guard.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.renorm_moe_routing_op import ( # noqa: E501
- renorm_moe_routing_op as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "renorm_moe_routing_op requires a CUDA device"
-
-
-def _compare(
- out: tuple[torch.Tensor, torch.Tensor],
- ref: tuple[torch.Tensor, torch.Tensor],
-) -> None:
- """Ids equal and weights bit-exact.
-
- The reference is not an approximation of the kernel: over 2040
- configurations and 16384-token runs at 512 / top-10, not one id and not
- one weight bit differed. A one-ulp `assert_close` runs first only so a
- drift reports a readable diff. The nearest wrong answers sit 0.4476
- away (un-renormalized softmax, `1/topk`), and a `torch.topk` reference
- differs on a third of the rows at 512 bf16 experts.
- """
- ids, weights = out
- ref_ids, ref_weights = ref
- assert ids.dtype == ref_ids.dtype == torch.int32, (ids.dtype, ref_ids.dtype)
- assert weights.dtype == ref_weights.dtype, (weights.dtype, ref_weights.dtype)
- assert ids.shape == ref_ids.shape == weights.shape, (ids.shape, ref_ids.shape)
- torch.testing.assert_close(ids, ref_ids, rtol=0, atol=0)
- eps = torch.finfo(weights.dtype).eps
- torch.testing.assert_close(weights, ref_weights, rtol=eps, atol=0.0)
- mismatched = int((weights != ref_weights).sum())
- assert mismatched == 0, f"{mismatched} of {weights.numel()} weights are not bit-exact"
-
-
-DEV = "cuda"
-
-
-def _logits(tokens, num_experts, logits_dtype, dist="randn", magnitude=1.0, seed=0):
- g = torch.Generator(device=DEV).manual_seed(seed)
- if dist == "coarse":
- x = torch.randint(0, 3, (tokens, num_experts), device=DEV, generator=g).float()
- else:
- x = torch.randn(tokens, num_experts, device=DEV, generator=g) * magnitude
- return x.to(logits_dtype)
-
-
-def _check(logits, topk, output_dtype=None):
- """Drive the op and gate it through the entry's own reference and `_compare`."""
- ids, weights = op(logits, topk, output_dtype)
- assert ids.is_contiguous() and weights.is_contiguous()
- assert ids.device == weights.device == logits.device
- _compare((ids, weights), op.reference(logits, topk, output_dtype))
- return ids, weights
-
-
-# ── the certified cells ───────────────────────────────────────────────────
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- s = cell.spec
- logits = _logits(
- s["tokens"],
- s["num_experts"],
- s["logits_dtype"],
- s["dist"],
- s.get("magnitude", 1.0),
- seed=zlib.crc32(cell.why.encode()),
- )
- ids, weights = _check(logits, s["topk"], s["output_dtype"])
- assert weights.dtype == s["output_dtype"]
- # renormalized: every row sums to 1
- torch.testing.assert_close(
- weights.float().sum(-1),
- torch.ones(s["tokens"], device=DEV),
- rtol=8e-3 if s["output_dtype"] is torch.bfloat16 else 1e-6,
- atol=0.0,
- )
-
-
-# ── 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."""
- 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(
- 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
-
-
-# ── the guard ─────────────────────────────────────────────────────────────
-
-
-def test_guard_refuses_a_column_slice():
- """What `is_valid` exists for: the raw op reads a slice as the dense buffer."""
- wide = _logits(8, 1024, torch.bfloat16, seed=12)
- with validating(op), pytest.raises(AssertionError, match="contiguous"):
- op(wide[:, :512], 10)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py
deleted file mode 100644
index 3d86b01b2f55..000000000000
--- 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
+++ /dev/null
@@ -1,70 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the flashinfer_gemma_rmsnorm entry."""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.norm.flashinfer_gemma_rmsnorm import ( # noqa: E501
- flashinfer_gemma_rmsnorm as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "flashinfer_gemma_rmsnorm requires a CUDA device"
-
-
-def _seed(text: str) -> int:
- return zlib.crc32(text.encode())
-
-
-def _build(spec, seed):
- """`(buffer, x, weight, eps)`: `x` is a view of `buffer` laid out as the spec says."""
- g = torch.Generator(device="cuda").manual_seed(seed)
- shape, dtype = spec["shape"], spec["dtype"]
- width = shape[-1]
-
- def randn(*s):
- return torch.randn(*s, generator=g, device="cuda").to(dtype)
-
- if "row_width" in spec:
- rows, off = shape[0], spec["col_offset"]
- buffer = randn(rows, spec["row_width"])
- x = buffer[:, off : off + width]
- else:
- buffer = x = randn(*shape)
- weight = torch.randn(width, generator=g, device="cuda").to(dtype)
- return buffer, x, weight, spec["eps"]
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- buffer, x, weight, eps = _build(cell.spec, seed=_seed(cell.why))
- assert tuple(x.shape) == cell.spec["shape"] and x.stride(-1) == 1
- if "row_width" in cell.spec and x.shape[0] > 1:
- assert x.stride(0) == cell.spec["row_width"] and not x.is_contiguous()
- buffer_before = buffer.clone()
-
- out = op(x, weight, eps)
-
- assert out.shape == x.shape and out.dtype == x.dtype
- assert out.data_ptr() != x.data_ptr(), "the op must return a new tensor"
- assert torch.equal(buffer, buffer_before), "x (and the stream it slices) is read-only"
- torch.testing.assert_close(out, op.reference(x, weight, eps))
-
-
-def test_guard_refuses_fp8_before_the_kernel_answers() -> None:
- """fp8 is in the DSL dtype table: accepted, and a tensor comes back."""
- x = torch.randn(4, 2560, device="cuda").to(torch.float8_e4m3fn)
- w = torch.randn(2560, device="cuda").to(torch.float8_e4m3fn)
- unguarded = torch.ops.trtllm.flashinfer_gemma_rmsnorm(x, w, 1e-6)
- assert unguarded.dtype is x.dtype and unguarded.shape == x.shape, (
- "the hazard this guard covers is gone"
- )
- with pytest.raises(AssertionError, match="dequantize"):
- with validating(op):
- op(x, w, 1e-6)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_layer_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_layer_norm.py
deleted file mode 100644
index 0e07d7975518..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_layer_norm.py
+++ /dev/null
@@ -1,71 +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 layer_norm entry.
-
-Weight and bias are drawn over this checkpoint's own 110 vision LayerNorm
-tensors' ranges: the affine is wide enough (weights to 17.5, biases to -25.6)
-that a wrong normalization is not masked by a near-unit scale.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.norm.layer_norm import ( # noqa: E501
- layer_norm as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401
-
-assert torch.cuda.is_available(), "layer_norm requires a CUDA device"
-
-DEVICE = torch.device("cuda")
-WIDTH = 1152
-EPS = 1e-6
-
-#: Over all of model.visual.*.norm{1,2} and model.visual.merger.norm (bf16).
-WEIGHT_LO, WEIGHT_HI = -1.210938, 17.5
-BIAS_LO, BIAS_HI = -25.625, 14.875
-
-
-def _gen(seed: int) -> torch.Generator:
- return torch.Generator(device=DEVICE).manual_seed(seed)
-
-
-def _affine(gen, width=WIDTH, dtype=torch.bfloat16):
- weight = torch.rand(width, generator=gen, device=DEVICE) * (WEIGHT_HI - WEIGHT_LO) + WEIGHT_LO
- bias = torch.rand(width, generator=gen, device=DEVICE) * (BIAS_HI - BIAS_LO) + BIAS_LO
- return weight.to(dtype), bias.to(dtype)
-
-
-def _build(spec, seed):
- gen = _gen(seed)
- shape, dtype = spec["shape"], spec["dtype"]
- 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)]
- 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 = [x, weight, bias]
- assert all(torch.equal(a, b) for a, b in zip(before, after)), "inputs must not be mutated"
- assert torch.equal(out, op(x, weight, bias, eps)), "two launches must agree bitwise"
- torch.testing.assert_close(out, op.reference(x, weight, bias, eps))
-
-
-def test_rejects_a_weight_whose_length_is_not_the_normalized_width():
- """Loud: `_layer_norm_fwd` asserts the shapes, with no message."""
- gen = _gen(88)
- x = torch.randn(8, WIDTH, generator=gen, device=DEVICE).to(torch.bfloat16)
- weight, bias = _affine(gen)
- with pytest.raises(AssertionError):
- op(x, weight[:-1].contiguous(), bias, EPS)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py
deleted file mode 100644
index 21569aeef4e0..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py
+++ /dev/null
@@ -1,80 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the rms_norm_gated_token_major entry."""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.norm.rms_norm_gated_token_major import ( # noqa: E501
- rms_norm_gated_token_major as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "rms_norm_gated_token_major requires a CUDA device"
-
-
-#: 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 guard test.
-_PACKED = dict(op.CELLS[0].spec, tokens=4, z="packed")
-
-
-def _seed(text: str) -> int:
- return zlib.crc32(text.encode())
-
-
-def _build(spec, seed):
- """`__call__`'s positional arguments for one spec."""
- g = torch.Generator(device="cuda").manual_seed(seed)
- t, heads, n, dtype = spec["tokens"], spec["heads"], spec["width"], spec["dtype"]
-
- def randn(*shape):
- return torch.randn(*shape, generator=g, device="cuda").to(dtype)
-
- x = randn(t * heads, n)
-
- layout = spec.get("z", "packed")
- if layout == "in_proj":
- z = randn(t, _IN_PROJ_WIDTH)[:, _Z_OFFSET : _Z_OFFSET + heads * n].view(t, heads, n)
- else:
- assert layout == "packed", layout
- z = randn(t, heads, n)
-
- weight = randn(n)
- 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 = 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() and out.dtype == x.dtype
- for t, b in zip((x, z, weight), before):
- assert torch.equal(t, b), "inputs are read-only"
- torch.testing.assert_close(out, op.reference(*args))
-
-
-def test_guard_refuses_a_weight_of_the_wrong_length() -> None:
- """Unguarded, a length-192 weight is silently truncated to weight[:128]."""
- x, z, weight, eps, _, act = _build(_PACKED, seed=12)
- long_weight = torch.cat([weight, torch.randn(64, dtype=weight.dtype, device="cuda")])
- assert torch.equal(
- torch.ops.trtllm.rms_norm_gated_token_major(x, z, long_weight, eps, None, act),
- torch.ops.trtllm.rms_norm_gated_token_major(x, z, weight, eps, None, act),
- ), "the hazard this guard covers is gone"
-
- with pytest.raises(AssertionError, match="weight must be 1-D of length"):
- with validating(op):
- op(x, z, long_weight, eps, None, act)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_gated_ngram_injection.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_gated_ngram_injection.py
deleted file mode 100644
index f93fdf147cd6..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_gated_ngram_injection.py
+++ /dev/null
@@ -1,164 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the gated_ngram_injection entry: the cells (the sub-floor and
-zero dots are a cell), the leading-dim guard and CUDA-graph replay."""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.gated_ngram_injection import ( # noqa: E501
- gate_ratio,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.gated_ngram_injection import ( # noqa: E501
- gated_ngram_injection as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "gated_ngram_injection requires a CUDA device"
-
-
-def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
- """16-bit: `rtol = eps`, `atol` one subnormal spacing (`gate_ratio <= 1`)
- -- half torch's bf16 `rtol`. Nothing cancels (each element is a sigmoid
- times one `value` element), so the budget is the final rounding; with
- the reference also rounded once, op and reference agree or sit one ulp
- apart. One ulp is at most `eps * |ref|`, so this band always admits it
- and never two. Measured at 2048 tokens: 0.992 (a one-ulp element just
- above a binade boundary), eight wrong readings at 105x..1115x. The
- `atol` is for Gaussian `value` elements that land subnormal (measured
- on fp16).
-
- fp32: torch's default band. The 2560-term signed dot is accumulated at
- the width it is reported in, so its conditioning, not the cast, sets
- the error; a custom pair would be fitted to a result.
- """
- assert out.dtype == ref.dtype, (out.dtype, ref.dtype)
- assert out.shape == ref.shape, (out.shape, ref.shape)
- assert out.is_contiguous()
- if ref.dtype == torch.float32:
- torch.testing.assert_close(out, ref)
- return
- ratio = gate_ratio(out, ref)
- assert ratio <= 1.0, f"{ratio:.4g}x the band"
-
-
-DEV = "cuda"
-HC_COUNT = 4
-HIDDEN = 2560
-WIDTH = HC_COUNT * HIDDEN
-RMS_NORM_EPS = 1e-6
-CLAMP_MIN = 1e-6
-
-# This checkpoint's PLE tensors: norm_key / norm_query gammas (element RMS
-# 0.1358, extremes 0.863 and -0.326), and the value's scale derived as
-# sqrt(2560) * 0.015295 (value_proj RMS) * 0.00747 (dequantized table RMS).
-# The entry is exactly linear in `value`, so its absolute scale decides nothing.
-CKPT_NORM_GAMMA_RMS = 0.1358
-CKPT_NORM_GAMMA_MAX = 0.863281
-CKPT_NORM_GAMMA_MIN = -0.326172
-CKPT_VALUE_RMS = 0.0058
-
-
-def _gemma_grouped_rmsnorm(x, gamma, hc_count):
- """Where key/query come from -- here only for their distribution."""
- hidden = x.shape[-1] // hc_count
- xf = x.float().reshape(*x.shape[:-1], hc_count, hidden)
- xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + RMS_NORM_EPS)
- return (xf.flatten(-2) * (1.0 + gamma.float())).to(x.dtype)
-
-
-def _make_operands(seed, dtype, leading, hc_count=HC_COUNT, value_dtype=None):
- g = torch.Generator(device=DEV).manual_seed(seed)
- width = hc_count * HIDDEN
-
- def normal(*dims, scale, dt=dtype):
- return (torch.randn(*dims, generator=g, device=DEV) * scale).to(dt)
-
- gammas = []
- for _ in range(2):
- gamma = normal(width, scale=CKPT_NORM_GAMMA_RMS)
- gamma[0] = CKPT_NORM_GAMMA_MAX
- gamma[width // 2] = CKPT_NORM_GAMMA_MIN
- gammas.append(gamma)
- key = _gemma_grouped_rmsnorm(normal(*leading, width, scale=1.0), gammas[0], hc_count)
- query = _gemma_grouped_rmsnorm(normal(*leading, width, scale=1.0), gammas[1], hc_count)
- value = normal(*leading, HIDDEN, scale=CKPT_VALUE_RMS, dt=value_dtype or dtype)
- return key, query, value
-
-
-def _single_channel_operands(dots, value_dtype):
- """One nonzero channel per stream: each dot is exactly `dots[s]`."""
- hc_count = len(dots)
- key = torch.zeros(8, hc_count * HIDDEN, dtype=torch.bfloat16, device=DEV)
- query = torch.zeros_like(key)
- for stream, dot in enumerate(dots):
- key[:, stream * HIDDEN] = 1.0
- query[:, stream * HIDDEN] = dot
- value = torch.full((8, HIDDEN), 0.5, dtype=value_dtype, device=DEV)
- return key, query, value
-
-
-def _build(spec, seed):
- if spec.get("kind") == "single_channel":
- dots = spec["dots"]
- key, query, value = _single_channel_operands(dots, spec["value_dtype"])
- return key, query, value, len(dots), spec["clamp_min"]
- key, query, value = _make_operands(
- seed, torch.bfloat16, spec["leading"], value_dtype=spec.get("value_dtype")
- )
- return key, query, value, HC_COUNT, spec.get("clamp_min", CLAMP_MIN)
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
- before = [a.clone() for a in args[:3]]
- out = op(*args)
- key, _, value, hc_count, _ = args
- assert out.shape == (*key.shape[:-1], hc_count * value.shape[-1])
- assert out.dtype == value.dtype
- for a, b in zip(args[:3], before):
- assert torch.equal(a, b), "the operands are read-only"
- _compare(out, op.reference(*args))
-
-
-def test_guard_refuses_unequal_leading_dims() -> None:
- """Unguarded, one query (or value) row is broadcast over every key row."""
- key, query, value = _make_operands(60, torch.bfloat16, (64,))
- unguarded = op(key, query[:1], value, HC_COUNT)
- assert unguarded.shape == (64, WIDTH), "the hazard this guard covers is gone"
- assert not torch.equal(unguarded, op.reference(key, query, value, HC_COUNT))
- assert op(key, query, value[:1], HC_COUNT).shape == (64, WIDTH)
-
- with pytest.raises(AssertionError, match="broadcasts silently"):
- with validating(op):
- op(key, query[:1], value, HC_COUNT)
- with pytest.raises(AssertionError, match="broadcasts silently"):
- with validating(op):
- op(key, query, value[:1], HC_COUNT)
-
-
-def test_cuda_graph_capture_and_replay() -> None:
- key, query, value = _make_operands(100, torch.bfloat16, (8,))
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(warmup):
- for _ in range(3):
- op(key, query, value, HC_COUNT, CLAMP_MIN)
- torch.cuda.current_stream().wait_stream(warmup)
-
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- captured = op(key, query, value, HC_COUNT, CLAMP_MIN)
-
- fresh = _make_operands(101, torch.bfloat16, (8,))
- for static, new in zip((key, query, value), fresh):
- static.copy_(new)
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(captured, op(*fresh, HC_COUNT, CLAMP_MIN))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/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
deleted file mode 100644
index 0e16ce5bdedf..000000000000
--- 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
+++ /dev/null
@@ -1,212 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the ngram_token_history entry.
-
-Integer bookkeeping, gated bitwise -- window and the whole published pool --
-against the entry's scalar reference, which builds each sequence's stream as
-a list and shares no structure with the masked per-shift select under test.
-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
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ngram_token_history import ( # noqa: E501
- CKPT_EOS_TOKEN_ID,
- CKPT_NGRAM_SIZE,
- CKPT_VOCAB_SIZE,
- ENGINE_SLOTS,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ngram_token_history import ( # noqa: E501
- ngram_token_history as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "ngram_token_history requires a CUDA device"
-
-
-def _compare(
- out: tuple[torch.Tensor, torch.Tensor], ref: tuple[torch.Tensor, torch.Tensor]
-) -> None:
- """`(window, context_pool)` -- the returned window with the pool as the
- call left it -- against `reference`'s pair, both bitwise: integer
- bookkeeping, no tolerance. The whole pool is compared, which is the
- lock on every slot the batch did not name."""
- (window, pool), (ref_window, ref_pool) = out, ref
- assert window.dtype == ref_window.dtype == torch.int64, (window.dtype, ref_window.dtype)
- assert window.shape == ref_window.shape, (window.shape, ref_window.shape)
- assert window.is_contiguous()
- assert torch.equal(window, ref_window), "window disagrees with the reference"
- assert torch.equal(pool, ref_pool), "published context pool disagrees with the reference"
-
-
-DEV = "cuda"
-EOS = CKPT_EOS_TOKEN_ID
-NGRAM = CKPT_NGRAM_SIZE
-CONTEXT_LEN = NGRAM - 1
-
-
-def _qsl(lengths, dtype, qsl_from="literal"):
- if qsl_from == "seq_lens":
- # The call site under MTP verify: an int32 prefix of the metadata's
- # seq_lens buffer, scanned into an int64 vector with a leading zero.
- buf = torch.full((len(lengths) + 9,), 99, dtype=torch.int32, device=DEV)
- buf[: len(lengths)] = torch.tensor(lengths, dtype=torch.int32, device=DEV)
- seq_lens = buf[: len(lengths)]
- out = torch.zeros(len(lengths) + 1, dtype=torch.int64, device=DEV)
- torch.cumsum(seq_lens, dim=0, dtype=torch.int64, out=out[1:])
- return out
- starts = [0]
- for length in lengths:
- starts.append(starts[-1] + length)
- return torch.tensor(starts, dtype=dtype, device=DEV)
-
-
-def _build(spec, seed):
- """The op's positional arguments for one cell."""
- g = torch.Generator().manual_seed(seed)
- 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
- pool = torch.randint(
- 0, CKPT_VOCAB_SIZE, (spec["pool_slots"], CONTEXT_LEN), generator=g, dtype=torch.long
- ).to(DEV)
- index_dtype = spec.get("index_dtype", torch.int32)
- 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,
- )
-
-
-def _expect(args):
- """The reference, driven from a snapshot of the pool before the call."""
- return op.reference(args[0], args[1].clone(), *args[2:])
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
- read_only = [a.clone() for a in (args[0], args[2], args[3], args[4])]
- pool_before = args[1].clone()
- expected = _expect(args)
- window = op(*args)
- _compare((window, args[1]), expected)
- for a, b in zip((args[0], args[2], args[3], args[4]), read_only):
- assert torch.equal(a, b), "only context_pool may be written"
- # The whole-pool comparison locks the unnamed slots; this keeps it from
- # being vacuous -- every cell has a sequence that publishes new ids.
- assert not torch.equal(args[1], pool_before), "no slot moved"
-
-
-# ----------------------------------------------------------- multi-step
-
-
-def test_chunked_prefill_matches_one_shot() -> None:
- """96 tokens split at 41 over two calls: window and pool bitwise one-shot."""
- g = torch.Generator().manual_seed(5001)
- ids = torch.randint(0, CKPT_VOCAB_SIZE, (96,), generator=g).to(DEV)
- start = torch.randint(0, CKPT_VOCAB_SIZE, (3, CONTEXT_LEN), generator=g).to(DEV)
- slot = torch.tensor([1], dtype=torch.int32, device=DEV)
-
- def call(pool, chunk, carry):
- flags = torch.tensor([carry], device=DEV)
- qsl = torch.tensor([0, chunk.shape[0]], dtype=torch.int32, device=DEV)
- return op(chunk, pool, slot, flags, qsl, EOS, NGRAM)
-
- one_shot_pool = start.clone()
- one_shot = call(one_shot_pool, ids, False)
- pool = start.clone()
- split = torch.cat([call(pool, ids[:41], False), call(pool, ids[41:], True)])
- assert torch.equal(split, one_shot)
- assert torch.equal(pool, one_shot_pool)
-
-
-# ----------------------------------------------------------- guards
-
-
-def _two_sequences():
- pool = torch.randint(0, CKPT_VOCAB_SIZE, (4, CONTEXT_LEN), device=DEV)
- return (
- torch.tensor([11, 12, 13, 21], device=DEV),
- pool,
- torch.tensor([0, 1], dtype=torch.int32, device=DEV),
- torch.tensor([True, False], device=DEV),
- torch.tensor([0, 3, 4], dtype=torch.int32, device=DEV),
- EOS,
- NGRAM,
- )
-
-
-def _misses(args, bad_index, bad_value):
- """Run the op unguarded with one argument swapped; True if it returned
- and its answer differs from the reference on the intended arguments."""
- expected = _expect(args)
- bad = list(args)
- bad[1] = args[1].clone()
- bad[bad_index] = bad_value
- window = op(*bad)
- return not (torch.equal(window, expected[0]) and torch.equal(bad[1], expected[1]))
-
-
-def test_guard_refuses_a_short_query_start_loc() -> None:
- """Unguarded, a 2-entry vector for 2 sequences misattributes the tail."""
- args = _two_sequences()
- short = torch.tensor([0, 3], dtype=torch.int32, device=DEV)
- assert _misses(args, 4, short), "the hazard this guard covers is gone"
- bad = list(args)
- bad[4] = short
- with pytest.raises(AssertionError, match=r"query_start_loc must be \[num_seqs"):
- with validating(op):
- op(*bad)
-
-
-# ----------------------------------------------------------- capture
-
-
-def test_cuda_graph_capture_and_replay() -> None:
- """The captured decode shape, 8 sequences at the engine pool width."""
- spec = dict(
- lengths=(1,) * 8,
- slots=(0, ENGINE_SLOTS - 1, 4, 17, 30, 43, 56, 69),
- carries=(True,) * 8,
- pool_slots=ENGINE_SLOTS,
- ids_dtype=torch.int32,
- qsl_dtype=torch.int64,
- )
- args = _build(spec, seed=11001)
- start = args[1].clone()
-
- side = torch.cuda.Stream()
- side.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(side):
- for _ in range(3):
- op(*args)
- torch.cuda.current_stream().wait_stream(side)
-
- args[1].copy_(start)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- captured = op(*args)
-
- replay_ids = torch.randint(0, CKPT_VOCAB_SIZE, (8,), dtype=torch.int32, device=DEV)
- args[0].copy_(replay_ids)
- args[1].copy_(start)
- expected = op.reference(replay_ids, start.clone(), *args[2:])
- graph.replay()
- torch.cuda.synchronize()
- _compare((captured, args[1]), expected)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py
deleted file mode 100644
index de9c02457d60..000000000000
--- 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
+++ /dev/null
@@ -1,212 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the ple_gate_conv_inject_norm entry.
-
-Cells are gated against the float64 reference that rounds where the op rounds:
-the stream (`hc_out`) and the published conv state within one ulp of each
-element, the normed stream within two (see `_compare`). Beyond the cells:
-the operand-width guard and graph capture.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ple_gate_conv_inject_norm import ( # noqa: E501
- CKPT_CLAMP_MIN,
- CKPT_DILATION,
- CKPT_HC_COUNT,
- CKPT_HIDDEN,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ple_gate_conv_inject_norm import ( # noqa: E501
- ple_gate_conv_inject_norm as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import ( # noqa: E402,F401 — path above; fixture by name
- assert_within_ulp,
- certifying,
- validating,
-)
-
-assert torch.cuda.is_available(), "ple_gate_conv_inject_norm requires a CUDA device"
-
-DEV = "cuda"
-EPS = 1e-6
-PAD = -1
-GAMMA_RMS = 0.477
-
-
-def _ulp_ratio(out: torch.Tensor, ref: torch.Tensor) -> torch.Tensor:
- """Per-element distance from `ref` in units of one ulp of `|ref|`."""
- eps = torch.finfo(ref.dtype).eps
- tiny = torch.finfo(ref.dtype).tiny
- ref64 = ref.double()
- return (out.double() - ref64).abs() / (eps * ref64.abs() + tiny)
-
-
-def _stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]:
- ratio = _ulp_ratio(out, ref)
- return ratio.max().item(), (out != ref).double().mean().item()
-
-
-#: The band. Every rounding point of the composition is kept, so the op and the
-#: float64 reference differ only where an fp32 reduction (three sums of squares
-#: and the gate's dot per stream) or fp32 `exp` lands a value across a bf16
-#: midpoint -- measured on 2e-5..2.3e-4 of the elements. A flipped upstream
-#: element then propagates, and through `hc + ple` it meets cancellation
-#: (`hc ~ -ple`), so a per-element ulp is unbounded there: 73 ulp of a
-#: 0.003-sized element in a row whose largest is ~10. The stream and its norm
-#: are therefore gated in units of their row's scale (`assert_within_ulp`):
-#: 2 ulp of the row maximum per element and 0.1 ulp relative RMS. The
-#: composition this replaces sits at 0.87 / 0.33 of a row-max ulp against the
-#: same reference (measured: normed / hc_out at 128 rows), the op at 0.72 /
-#: 0.38. The published conv state is one rounding of the conv-input norm,
-#: whose only perturbation is a one-ulp flip of a gated element: 2 ulp per
-#: element (measured worst 1.78, the old composition 1.56).
-ROW_ULP, RMS_ULP, STATE_ULP = 2.0, 0.1, 2.0
-#: The per-element row-max gate is a tail statistic, and a context-shaped cell
-#: draws ~100x more elements than a 128-row one. Against the same float64
-#: reference over three seeds at 16384 rows, the composition this replaces
-#: reaches 2.91 row-max ulp (normed) and the op 2.94; at 8192 rows 1.89 / 1.89.
-#: Cells above 1024 rows are therefore gated at 4; the RMS gate is unchanged.
-CONTEXT_ROWS, CONTEXT_ROW_ULP = 1024, 4.0
-
-
-def _compare(out, ref) -> None:
- row_ulp = CONTEXT_ROW_ULP if out[0].shape[0] > CONTEXT_ROWS else ROW_ULP
- for name, got, want in zip(("hc_out", "normed", "conv_states"), out, ref):
- assert got.dtype == want.dtype and got.shape == want.shape, name
- if not got.numel():
- continue
- if name == "conv_states":
- worst = _ulp_ratio(got, want).max().item()
- assert worst <= STATE_ULP, f"{name}: {worst:.4g} ulp"
- else:
- assert_within_ulp(got, want, row_ulp, RMS_ULP)
-
-
-def _build(spec, seed):
- g = torch.Generator(device=DEV).manual_seed(seed)
- gc = torch.Generator().manual_seed(seed)
- lengths = spec["lengths"]
- c, d, dtype = CKPT_HC_COUNT, CKPT_HIDDEN, torch.bfloat16
- w = c * d
- rows = sum(lengths)
- num_seqs = len(lengths)
- pool_slots = spec.get("pool_slots", num_seqs + 2)
-
- def rand(*shape, scale=1.0):
- return (torch.randn(*shape, generator=g, device=DEV) * scale).to(dtype)
-
- # Streams at different scales, as the trained residual's are.
- scales = torch.tensor([0.5 * 2.0 ** (i % 4) for i in range(c)], device=DEV)
- hc = (rand(rows, w, scale=3.0).float() * scales.repeat_interleave(d)).to(dtype)
- key = rand(rows, w)
- value = rand(rows, d, scale=0.5)
- norms = [rand(w, scale=GAMMA_RMS) for _ in range(4)]
- conv_weight = rand(w, 1, 4, scale=0.3)
- conv_states = rand(pool_slots, w, 3 * CKPT_DILATION)
- order = (torch.randperm(pool_slots - 2, generator=gc) + 1).tolist()
- slots = ([0, pool_slots - 1] + order)[:num_seqs] if pool_slots >= 64 else list(range(num_seqs))
- slots = torch.tensor(slots, dtype=torch.int32, device=DEV)
- starts = [0]
- for length in lengths:
- starts.append(starts[-1] + length)
- qsl = torch.tensor(starts, dtype=torch.int64, device=DEV)
- max_rows = spec.get("max_rows", max(lengths))
- carries = spec.get("carries")
- carries = None if carries is None else torch.tensor(carries, dtype=torch.bool, device=DEV)
- return [
- hc, key, value, norms[0], norms[1], norms[2], conv_weight, conv_states, slots, qsl,
- norms[3], EPS, c, CKPT_CLAMP_MIN, CKPT_DILATION, PAD, max_rows, carries,
- ] # fmt: skip
-
-
-def _run(args):
- """`(hc_out, normed, conv_states)` from the op, on a copy of the pool."""
- args = list(args)
- args[7] = args[7].clone()
- hc_out, normed = op(*args)
- return hc_out, normed, args[7]
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
- certifying(op)
- args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
- hc_before = args[0].clone()
- ref = op.reference(*args)
- out = _run(args)
- assert torch.equal(args[0], hc_before), "hc is read-only"
- assert out[0].is_contiguous() and out[1].is_contiguous()
- _compare(out, ref)
- for name, got, want in zip(("hc_out", "normed", "conv_states"), out, ref):
- worst, frac = _stats(got, want)
- print(
- f"\n {cell.why[:40]:<40} {name:<12} worst {worst:.3g} ulp, differ {frac:.2e}", end=""
- )
-
-
-def test_conv_input_is_what_the_conv_consumed() -> None:
- """`conv_input` receives the conv input of every row, and changes nothing else: with
- four rows a sequence, the published state's last four columns are those rows."""
- args = _build(dict(lengths=(4,) * 16, pool_slots=461), seed=11)
- plain = _run(args)
- xc = torch.empty_like(args[0])
- args[7] = args[7].clone()
- hc_out, normed = op(*args, conv_input=xc)
- assert torch.equal(hc_out, plain[0]) and torch.equal(normed, plain[1])
- assert torch.equal(args[7], plain[2])
- tail = args[7][args[8].long()][:, :, -4:]
- assert torch.equal(torch.reshape(torch.transpose(tail, 1, 2), xc.shape), xc)
-
-
-def test_out_writes_the_rows_of_larger_buffers() -> None:
- """`out` takes the leading rows of two padded `[T + 8, C*D]` buffers: the same bits as
- the fresh outputs, the padding rows untouched -- the prefill target's in-place write."""
- args = _build(dict(lengths=(3001, 5000) + (1,) * 64, carries=(False, True) + (True,) * 64,
- pool_slots=461), seed=12) # fmt: skip
- plain = _run(args)
- rows = args[0].shape[0]
- bufs = [torch.full((rows + 8, args[0].shape[1]), 7.0, dtype=args[0].dtype, device=DEV)
- for _ in range(2)] # fmt: skip
- args[7] = args[7].clone()
- hc_out, normed = op(*args, out=(bufs[0][:rows], bufs[1][:rows]))
- assert hc_out.data_ptr() == bufs[0].data_ptr() and normed.data_ptr() == bufs[1].data_ptr()
- assert torch.equal(bufs[0][:rows], plain[0]) and torch.equal(bufs[1][:rows], plain[1])
- assert torch.equal(args[7], plain[2])
- assert bool((bufs[0][rows:] == 7.0).all()) and bool((bufs[1][rows:] == 7.0).all())
-
-
-def test_guard_refuses_a_short_value() -> None:
- args = _build(dict(lengths=(1, 1)), seed=4)
- args[2] = args[2][:, :-1]
- with validating(op), pytest.raises(AssertionError):
- op(*args)
-
-
-def test_cuda_graph_capture_and_replay() -> None:
- args = _build(dict(lengths=(1,) * 8, pool_slots=461), seed=8)
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- pool = args[7].clone()
- with torch.cuda.stream(warmup):
- for _ in range(2):
- op(*args)
- torch.cuda.current_stream().wait_stream(warmup)
- args[7].copy_(pool)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- graph_hc, graph_normed = op(*args)
- fresh = _build(dict(lengths=(1,) * 8, pool_slots=461), seed=18)
- args[0].copy_(fresh[0])
- args[1].copy_(fresh[1])
- args[7].copy_(pool)
- eager = _run(args)
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(graph_hc, eager[0])
- assert torch.equal(graph_normed, eager[1])
- assert torch.equal(args[7], eager[2])
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_ngram_embed.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_ngram_embed.py
deleted file mode 100644
index 1051e07320fd..000000000000
--- 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
+++ /dev/null
@@ -1,230 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the ple_ngram_embed entry.
-
-Integer bookkeeping plus a one-rounding dequantize, so gated bitwise -- the
-embedding and the whole published pool -- against the entry's reference, which
-is the scalar references of `ngram_token_history` and `ngram_hash` and a
-float64 gather-and-scale. Beyond the cells: the `query_start_loc` length guard,
-graph capture, and the host table's huge-page advice.
-"""
-
-import errno
-import mmap
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.embedding.ngram_hash import ( # noqa: E501
- CKPT_HEAD_OFFSETS,
- CKPT_HEAD_VOCAB_SIZES,
- CKPT_MULTIPLIERS,
- ngram_hash,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ngram_token_history import ( # noqa: E501
- ngram_token_history,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ple_ngram_embed import ( # noqa: E501
- CKPT_EOS_TOKEN_ID,
- CKPT_HEAD_DIM,
- CKPT_SCALE,
- CKPT_TABLE_ROWS,
- ENGINE_SLOTS,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ple_ngram_embed import ( # noqa: E501
- ple_ngram_embed as op,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.weights import ( # noqa: E501
- _host_resident_empty,
-)
-from tensorrt_llm.logger import logger
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "ple_ngram_embed requires a CUDA device"
-
-DEV = "cuda"
-EOS = CKPT_EOS_TOKEN_ID
-VOCAB = 248320
-FP8 = torch.float8_e4m3fn
-#: Keeps the hosts alive for as long as the tables they back.
-_HOSTS = []
-
-
-def _scattered(n: int, pool_slots: int, g) -> list:
- """`n` distinct slots over the whole pool, both ends included."""
- middle = (torch.randperm(pool_slots - 2, generator=g) + 1).tolist()
- slots = ([0, pool_slots - 1] + middle)[:n]
- assert len(set(slots)) == n and max(slots) < pool_slots
- return slots
-
-
-def _geometry(spec, g):
- """`(multipliers, vocab, offsets)` for the cell: the checkpoint's own over
- the full-height table, the checkpoint's multipliers over small stand-in head
- vocabularies otherwise."""
- if spec.get("table") == "ckpt":
- return CKPT_MULTIPLIERS, CKPT_HEAD_VOCAB_SIZES, CKPT_HEAD_OFFSETS
- mult = CKPT_MULTIPLIERS
- vocab = tuple(
- 1000 + 7 * h + int(x)
- for h, x in enumerate(torch.randint(0, 50, (len(CKPT_HEAD_VOCAB_SIZES),), generator=g))
- )
- offsets, acc = [], 0
- for v in vocab:
- offsets.append(acc)
- acc += v
- return mult, vocab, tuple(offsets)
-
-
-def _table(kind: str, rows: int, touched: torch.Tensor, g) -> torch.Tensor:
- if kind == "ckpt":
- table = torch.empty(CKPT_TABLE_ROWS, CKPT_HEAD_DIM, dtype=FP8, device=DEV)
- table[touched] = torch.randn(touched.numel(), CKPT_HEAD_DIM, generator=g).to(FP8).to(DEV)
- return table
- src = torch.randn(rows, CKPT_HEAD_DIM, generator=g).to(FP8)
- if kind == "host":
- table, host = _host_resident_empty(rows, CKPT_HEAD_DIM, FP8, torch.device(DEV))
- _HOSTS.append(host)
- table.copy_(src.to(DEV))
- return table
- return src.to(DEV)
-
-
-def _with_eos(x: torch.Tensor, density: float, g) -> torch.Tensor:
- if density > 0.0:
- x = torch.where(torch.rand(x.shape, generator=g) < density, torch.full_like(x, EOS), x)
- return x
-
-
-def _build(spec, seed):
- g = torch.Generator().manual_seed(seed)
- lengths = spec["lengths"]
- n = len(CKPT_MULTIPLIERS)
- num_seqs = len(lengths)
- pool_slots = spec.get("pool_slots", num_seqs + 2)
- density = spec.get("eos_density", 0.0)
- ids = _with_eos(torch.randint(0, VOCAB, (sum(lengths),), generator=g), density, g)
- ids = ids.to(torch.int32).to(DEV)
- pool = _with_eos(torch.randint(0, VOCAB, (pool_slots, n - 1), generator=g), density, g).to(DEV)
- slots = _scattered(num_seqs, pool_slots, g) if pool_slots >= 64 else list(range(num_seqs))
- slots = torch.tensor(slots, dtype=torch.int32, device=DEV)
- carries = spec.get("carries")
- carries = None if carries is None else torch.tensor(carries, dtype=torch.bool, device=DEV)
- starts = [0]
- for length in lengths:
- starts.append(starts[-1] + length)
- qsl = torch.tensor(starts, dtype=torch.int64, device=DEV)
- mult, vocab, offsets = _geometry(spec, g)
- mult = torch.tensor(mult, dtype=torch.int64, device=DEV)
- vocab = torch.tensor(vocab, dtype=torch.int64, device=DEV)
- offsets = torch.tensor(offsets, dtype=torch.int64, device=DEV)
- # The rows this batch touches, from the two scalar references, so that the
- # full-height table needs bytes only where they are read.
- seeded = torch.ones(num_seqs, dtype=torch.bool, device=DEV) if carries is None else carries
- window, _ = ngram_token_history.reference(ids, pool, slots, seeded, qsl, EOS, n)
- touched = ngram_hash.reference(window, mult, vocab, offsets, EOS, 1).reshape(-1)
- rows = int(offsets[-1] + vocab[-1])
- table = _table(spec.get("table", "small"), rows, touched, g)
- scale = torch.tensor([CKPT_SCALE], dtype=torch.bfloat16, device=DEV).reshape(1, 1, 1)
- max_rows = spec.get("max_rows", max(lengths))
- return [ids, pool, slots, carries, qsl, table, scale, mult, vocab, offsets, EOS, max_rows]
-
-
-def _compare(out, ref) -> None:
- (emb, pool), (ref_emb, ref_pool) = out, ref
- assert emb.dtype == ref_emb.dtype and emb.shape == ref_emb.shape
- assert emb.is_contiguous()
- assert torch.equal(emb, ref_emb), (
- f"embedding differs at {(emb != ref_emb).any(-1).nonzero().flatten()[:8].tolist()}"
- )
- assert torch.equal(pool, ref_pool), "published pool differs"
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
- certifying(op)
- args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
- ref = op.reference(*args)
- pool_before = args[1].clone()
- out = op(*args)
- _compare((out, args[1]), ref)
- if cell.spec.get("table") == "host":
- assert args[1].data_ptr() != 0 and not torch.equal(args[1], pool_before)
-
-
-def test_guard_refuses_short_query_start_loc() -> None:
- args = _build(dict(lengths=(1, 1, 1)), seed=3)
- args[4] = args[4][:-1]
- with validating(op), pytest.raises(AssertionError):
- op(*args)
-
-
-def test_cuda_graph_capture_and_replay() -> None:
- args = _build(dict(lengths=(1,) * 16, pool_slots=ENGINE_SLOTS, eos_density=0.1), seed=11)
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(warmup):
- op(*args)
- torch.cuda.current_stream().wait_stream(warmup)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- graph_out = op(*args)
- fresh = _build(dict(lengths=(1,) * 16, pool_slots=ENGINE_SLOTS, eos_density=0.1), seed=12)
- args[0].copy_(fresh[0])
- pool_before = args[1].clone()
- want, want_pool = op.reference(*args)
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(graph_out, want)
- assert torch.equal(args[1], want_pool)
- assert not torch.equal(args[1], pool_before)
-
-
-def _huge_page_advice_accepted() -> bool:
- """Whether this kernel takes MADV_HUGEPAGE (it has transparent huge pages)."""
- probe = mmap.mmap(-1, mmap.PAGESIZE)
- try:
- probe.madvise(mmap.MADV_HUGEPAGE)
- except OSError:
- return False
- finally:
- probe.close()
- return True
-
-
-@pytest.mark.skipif(not _huge_page_advice_accepted(), reason="the kernel refuses MADV_HUGEPAGE")
-def test_host_table_is_advised_transparent_huge_pages() -> None:
- # The host-resident table is a mapping of its own, advised MADV_HUGEPAGE
- # before it is faulted in, whatever torch's CPU allocator does with huge pages.
- table, host = _host_resident_empty(1 << 16, CKPT_HEAD_DIM, FP8, torch.device(DEV))
- flags, inside = None, False
- with open("/proc/self/smaps") as smaps:
- for line in smaps:
- fields = line.split()
- if line[:1] in "0123456789abcdef" and "-" in fields[0]:
- lo, hi = (int(x, 16) for x in fields[0].split("-"))
- inside = lo <= host.pointer < hi
- elif inside and line.startswith("VmFlags:"):
- flags = fields[1:]
- break
- assert flags is not None and "hg" in flags, flags
- assert table.shape == (1 << 16, CKPT_HEAD_DIM)
-
-
-def test_host_table_loads_when_huge_page_advice_is_refused(monkeypatch) -> None:
- # A kernel without transparent huge pages refuses the advice: the table is
- # still allocated and registered, on base pages, with a warning.
- class _Refused(mmap.mmap):
- def madvise(self, *args):
- raise OSError(errno.EINVAL, "Invalid argument")
-
- warnings = []
- monkeypatch.setattr(mmap, "mmap", _Refused)
- monkeypatch.setattr(logger, "warning", warnings.append)
- table, host = _host_resident_empty(1 << 16, CKPT_HEAD_DIM, FP8, torch.device(DEV))
- assert isinstance(host._mapping, _Refused) and host.pointer
- assert table.shape == (1 << 16, CKPT_HEAD_DIM)
- assert len(warnings) == 1 and "MADV_HUGEPAGE" in warnings[0], warnings
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
deleted file mode 100644
index 3bf21aa13b26..000000000000
--- 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
+++ /dev/null
@@ -1,67 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the bf16_to_fp32 entry: the cells bit for bit `x.float()`, graph capture, the
-contiguity guard."""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.quantization.bf16_to_fp32 import ( # noqa: E501
- bf16_to_fp32 as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "bf16_to_fp32 requires a CUDA device"
-
-
-def _build(spec, seed):
- shape = spec["shape"]
- if spec.get("source") == "all_codes":
- return (
- torch.arange(-32768, 32768, dtype=torch.int32, device="cuda")
- .to(torch.int16)
- .view(torch.bfloat16)
- .reshape(shape)
- )
- g = torch.Generator(device="cuda").manual_seed(seed)
- return (torch.randn(shape, generator=g, device="cuda") * 8).to(torch.bfloat16)
-
-
-def _bitwise(a: torch.Tensor, b: torch.Tensor) -> bool:
- return a.shape == b.shape and torch.equal(a.view(torch.int32), b.view(torch.int32))
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
- certifying(op)
- x = _build(cell.spec, zlib.crc32(cell.why.encode()))
- out = op(x)
- ref = op.reference(x)
- assert out.dtype == torch.float32 and out.is_contiguous()
- assert _bitwise(out, ref), "widening is exact; a bit difference is a defect"
-
-
-def test_cuda_graph_capture_and_replay() -> None:
- x = _build(dict(shape=(128, 4099)), 1)
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(warmup):
- op(x)
- torch.cuda.current_stream().wait_stream(warmup)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- out = op(x)
- x.copy_(_build(dict(shape=(128, 4099)), 2))
- graph.replay()
- torch.cuda.synchronize()
- assert _bitwise(out, x.float())
-
-
-def test_guard_refuses_a_strided_input() -> None:
- x = _build(dict(shape=(8, 64)), 3)
- with validating(op), pytest.raises(AssertionError):
- op(x[:, ::2])
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_dequantize_e4m3_per_tensor.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_dequantize_e4m3_per_tensor.py
deleted file mode 100644
index ef52eab0be66..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_dequantize_e4m3_per_tensor.py
+++ /dev/null
@@ -1,140 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the dequantize_e4m3_per_tensor catalog entry.
-
-The entry owns the contract. This file turns a cell into operands and holds
-the claims no cell carries: graph capture with the scale read at replay, and
-the loud rejection of an empty batch.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.quantization.dequantize_e4m3_per_tensor import ( # noqa: E501
- PLE_SCALE,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.quantization.dequantize_e4m3_per_tensor import ( # noqa: E501
- dequantize_e4m3_per_tensor as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "dequantize_e4m3_per_tensor requires a CUDA device"
-
-
-_BIT_VIEW = {
- torch.bfloat16: torch.int16,
- torch.float16: torch.int16,
- torch.float32: torch.int32,
-}
-
-
-def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
- """Bit-exact, NaN for NaN.
-
- Both sides form the same exact fp32 product and round it once, so any
- difference is a defect. Measured bit-exact on every case, including all
- 256 codes x three output dtypes x five scales, with -0.0 and the
- subnormals compared on bit patterns. NaN payloads are not compared.
- """
- assert out.dtype == ref.dtype, (out.dtype, ref.dtype)
- assert out.shape == ref.shape, (out.shape, ref.shape)
- nan = torch.isnan(ref)
- assert torch.equal(torch.isnan(out), nan), "NaN lanes differ"
- view = _BIT_VIEW[ref.dtype]
- assert torch.equal(out.contiguous().view(view)[~nan], ref.contiguous().view(view)[~nan]), (
- "dequantization is exact; a bit difference is a defect"
- )
-
-
-DEV = "cuda"
-
-_TABLE = {}
-
-
-def _fp8(x: torch.Tensor) -> torch.Tensor:
- return x.to(torch.float8_e4m3fn)
-
-
-def _scale(v: float, rank: int, dtype=torch.bfloat16) -> torch.Tensor:
- """The only shape the op accepts: `input`'s rank, every dimension 1."""
- return torch.full((1,) * rank, v, dtype=dtype, device=DEV)
-
-
-def _input(spec, seed: int) -> torch.Tensor:
- g = torch.Generator(device=DEV).manual_seed(seed)
- source, shape = spec["source"], spec["shape"]
- if source == "gather":
- if "table" not in _TABLE:
- tg = torch.Generator(device=DEV).manual_seed(11)
- _TABLE["table"] = _fp8(torch.randn(1 << 20, 160, device=DEV, generator=tg) * 40.0)
- idx = torch.randint(0, 1 << 20, shape[:2], device=DEV, generator=g)
- x = _TABLE["table"][idx]
- else:
- assert source == "all_codes", source
- x = torch.arange(256, dtype=torch.uint8, device=DEV).view(shape).view(torch.float8_e4m3fn)
- assert x.is_contiguous() and x.shape == shape
- return x
-
-
-def _bits(t: torch.Tensor) -> torch.Tensor:
- view = {torch.bfloat16: torch.int16, torch.float16: torch.int16, torch.float32: torch.int32}
- return t.view(view[t.dtype])
-
-
-# ── the certified cells ───────────────────────────────────────────────────
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- x = _input(cell.spec, zlib.crc32(cell.why.encode()))
- for value in cell.spec["scales"]:
- scales = _scale(value, x.dim(), cell.spec["dtype"])
- out = op(x, scales)
- assert out.shape == x.shape and out.dtype == scales.dtype and out.is_contiguous()
- _compare(out, op.reference(x, scales))
- if cell.spec["source"] == "gather":
- flat = out.reshape(x.shape[0], -1)
- assert flat.shape[1] == 2560 and flat.data_ptr() == out.data_ptr(), "flatten is a view"
-
-
-# ── what a cell cannot carry ──────────────────────────────────────────────
-
-
-def test_cuda_graph_capture_and_replay():
- """Both operands stay live inputs: the scale is read at replay."""
- g = torch.Generator(device=DEV).manual_seed(16)
- static_in = _fp8(torch.randn(256, 16, 160, device=DEV, generator=g) * 8.0)
- static_scales = _scale(PLE_SCALE, 3)
- eager = op(static_in, static_scales).clone()
- warm = torch.cuda.Stream()
- warm.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(warm):
- for _ in range(3):
- op(static_in, static_scales)
- torch.cuda.current_stream().wait_stream(warm)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- static_out = op(static_in, static_scales)
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(_bits(static_out), _bits(eager))
- fresh = _fp8(torch.randn(256, 16, 160, device=DEV, generator=g) * 8.0)
- static_in.copy_(fresh)
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(_bits(static_out), _bits(op(fresh, static_scales)))
- static_scales.copy_(_scale(0.5, 3))
- graph.replay()
- torch.cuda.synchronize()
- _compare(static_out, op.reference(fresh, _scale(0.5, 3)))
-
-
-def test_rejects_an_empty_batch():
- """The op raises on zero rows rather than launching; the caller branches first."""
- with pytest.raises((RuntimeError, NotImplementedError)):
- op(_fp8(torch.randn(0, 128, device=DEV)), _scale(PLE_SCALE, 2))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/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
deleted file mode 100644
index caa087518151..000000000000
--- 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
+++ /dev/null
@@ -1,161 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the fp8_quantize_1x128 catalog entry.
-
-The entry owns the contract. This file turns a cell into an input, 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
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.quantization.fp8_quantize_1x128 import ( # noqa: E501
- SCALE_ULP_GATE,
- block_scales,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.quantization.fp8_quantize_1x128 import ( # noqa: E501
- fp8_quantize_1x128 as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "fp8_quantize_1x128 requires a CUDA device"
-
-
-def _ordinal(codes: torch.Tensor) -> torch.Tensor:
- """Sign-magnitude ordinal of an e4m3 byte: adjacent codes differ by 1."""
- b = codes.view(torch.uint8).int()
- return torch.where(b >= 128, -(b & 0x7F), b & 0x7F)
-
-
-def _compare(
- out: tuple[torch.Tensor, torch.Tensor],
- ref: tuple[torch.Tensor, torch.Tensor],
-) -> None:
- """Scales within 4 fp32 ulp; codes at most one e4m3 step apart, >=99.5% bit-exact.
-
- This build compiles the kernel's two fp32 divisions approximately
- (`--use_fast_math`), each within CUDA's documented 2 ulp, so neither
- output is reproducible to the last bit with `use_ue8m0=False`. Measured
- over the certified column: scales within 2 ulp of both `amax / 448` and
- `1 / (448 / amax)` (~80% bit-exact at T=1024); codes at worst
- 99.6875% bit-exact at H 2560 (99.8566% elsewhere), every miss exactly
- one code away with the exact product within 1.122e-07 relative of the
- midpoint between the two -- a rounding tie the scale's last ulp decides.
- `use_ue8m0=True` is bit-exact (the power-of-two scale makes the
- multiply exact), and passes this gate trivially.
-
- Wrong recipes at T=256, H=2560 (bytes agreeing / worst relative scale
- error): power-of-two 0.043 / 0.99, per-tensor 0.0002 / 2.0, per-row
- 0.083 / 1.6, blocks down the token axis 0.19 / 1.0, blocks offset by 64
- 0.48 / 0.93, e4m3 max as 240 0 / 0.87, signed max 0.58 / 0.58; the
- correct one 0.9993 / 1.4e-07.
- """
- data, scales = out
- ref_data, ref_scales = ref
- assert data.dtype == ref_data.dtype == torch.float8_e4m3fn, (data.dtype, ref_data.dtype)
- assert scales.dtype == ref_scales.dtype == torch.float32, (scales.dtype, ref_scales.dtype)
- assert data.shape == ref_data.shape, (data.shape, ref_data.shape)
- assert scales.shape == ref_scales.shape, (scales.shape, ref_scales.shape)
- if scales.numel():
- ulp = (
- scales.contiguous().view(torch.int32).long()
- - ref_scales.contiguous().view(torch.int32).long()
- )
- worst = int(ulp.abs().max())
- assert worst <= SCALE_ULP_GATE, f"scale {worst} fp32 ulp off > {SCALE_ULP_GATE}"
- if data.numel():
- step = (_ordinal(data) - _ordinal(ref_data)).abs()
- assert int(step.max()) <= 1, f"a code is {int(step.max())} e4m3 steps off"
- exact = 1.0 - (step != 0).float().mean().item()
- assert exact >= 0.995, f"only {100 * exact:.4f}% of codes bit-exact"
-
-
-DEV = "cuda"
-# A code can only disagree when the exact product sits within the first
-# division's 2 ulp of an e4m3 rounding midpoint.
-TIE_REL_GATE = 2.0 * 2.0**-23
-
-
-def _expand(v: torch.Tensor, n: int) -> torch.Tensor:
- """`[T, nb]` block values -> `[T, n]` per-column values."""
- return v.repeat_interleave(128, dim=1)[:, :n].contiguous()
-
-
-def _input(tokens: int, hidden: int, seed: int, magnitude: float = 3.0) -> torch.Tensor:
- g = torch.Generator(device=DEV).manual_seed(seed)
- return (torch.randn(tokens, hidden, device=DEV, generator=g) * magnitude).to(torch.bfloat16)
-
-
-def _assert_misses_are_ties(x: torch.Tensor, data: torch.Tensor, ref_data: torch.Tensor) -> float:
- """Every code disagreement is a rounding tie: the exact product sits within
- the scale's 2-ulp band of the midpoint between the two codes."""
- qs, _ = block_scales(x, False)
- prod = x.float() * _expand(qs, x.shape[1])
- diff = data.float() != ref_data.float() # -0 and +0 are the same code here
- if not bool(diff.any()):
- return 0.0
- mid = ((data.float() + ref_data.float()) / 2).double()
- rel = ((prod.double()[diff] - mid[diff]).abs() / mid[diff].abs()).max().item()
- assert rel <= TIE_REL_GATE, f"a miss {rel:.3e} from its midpoint is not a tie"
- return rel
-
-
-# ── the certified cells ───────────────────────────────────────────────────
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- spec = cell.spec
- x = _input(spec["tokens"], spec["hidden"], seed=zlib.crc32(cell.why.encode()))
- data, scales = op(x, spec["use_ue8m0"])
- nb = (spec["hidden"] + 127) // 128
- assert data.shape == x.shape and data.is_contiguous() and data.device == x.device
- assert scales.shape == (nb, spec["tokens"]) and scales.is_contiguous()
- ref = op.reference(x, spec["use_ue8m0"])
- _compare((data, scales), ref)
- rel = _assert_misses_are_ties(x, data, ref[0]) if not spec["use_ue8m0"] else 0.0
- exact = (data.view(torch.uint8) == ref[0].view(torch.uint8)).float().mean().item()
- print(f" {cell.why[:44]}: {100 * exact:.4f}% bit-exact, worst tie {rel:.3e}")
-
-
-# ── what a cell cannot carry ──────────────────────────────────────────────
-
-
-def test_cuda_graph_capture_and_replay():
- """A replay reproduces a fresh eager call bitwise; so does a side stream."""
- x = _input(64, 2560, seed=17, magnitude=1.0)
- side = torch.cuda.Stream()
- side.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(side):
- for _ in range(3):
- op(x)
- torch.cuda.current_stream().wait_stream(side)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- g_data, g_sf = op(x)
- x.copy_(_input(64, 2560, seed=18, magnitude=1.0))
- e_data, e_sf = op(x)
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(g_data.view(torch.uint8), e_data.view(torch.uint8))
- assert torch.equal(g_sf.view(torch.int32), e_sf.view(torch.int32))
- alt = torch.cuda.Stream()
- with torch.cuda.stream(alt):
- a_data, a_sf = op(x)
- alt.synchronize()
- assert torch.equal(a_data.view(torch.uint8), e_data.view(torch.uint8))
- assert torch.equal(a_sf.view(torch.int32), e_sf.view(torch.int32))
-
-
-def test_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_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
deleted file mode 100644
index 7f8a993b6785..000000000000
--- 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
+++ /dev/null
@@ -1,116 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the causal_conv1d_triton_update entry."""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_triton_update import ( # noqa: E501
- ENGINE_SLOTS,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_triton_update import ( # noqa: E501
- causal_conv1d_triton_update as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "causal_conv1d_triton_update requires a CUDA device"
-
-
-def _compare(out: tuple, ref: tuple) -> None:
- """`(out, conv_state, intermediate_conv_window)` triples: the returned
- tensor with the two buffers as the call left them, against
- `reference`'s triple.
-
- The result at torch's default band over the rows the op writes (NaN in
- the reference marks the rest): against the operand-dtype-product
- reference bf16/fp16 were bit-exact at dim 512 and at dim 10240 off by
- one ulp in at most 9 of 2.6M elements (the fp32 `exp`); fp32 is off by
- 1.4e-6. The pool and the snapshots are copies, gated bitwise over the
- whole buffer -- the lock on every unnamed slot and snapshot row.
- """
- (act, state, inter), (ref_act, ref_state, ref_inter) = out, ref
- assert act.dtype == ref_act.dtype and act.shape == ref_act.shape
- written = ~torch.isnan(ref_act)
- torch.testing.assert_close(act[written], ref_act[written])
- assert torch.equal(state, ref_state), "the advanced state must be bitwise exact"
- assert (inter is None) == (ref_inter is None)
- if inter is not None:
- assert torch.equal(inter, ref_inter), "the snapshots must be bitwise exact"
-
-
-PAD_SLOT_ID = -1
-
-
-def _build(spec, seed):
- """`__call__`'s positional arguments for one spec, freshly allocated."""
- g = torch.Generator(device="cuda").manual_seed(seed)
- dim, batch, seqlen, width = spec["dim"], spec["batch"], spec["seqlen"], 4
-
- def randn(*shape):
- return torch.randn(*shape, generator=g, device="cuda").to(torch.bfloat16)
-
- # The call site under the one input GEMM: the q|k|v columns of an
- # [n * p, C + 6240] projection, reshaped [n, p, C] and seen as [n, C, p].
- wide = dim + 6144 + 2 * 48
- x = randn(batch * seqlen, wide)[:, :dim].reshape(batch, seqlen, dim).transpose(1, 2)
- assert x.stride() == (seqlen * wide, 1, wide)
- conv_state = randn(spec["pool_slots"], dim, width - 1)
- weight = randn(dim, width)
- indices = torch.tensor(spec["slots"], dtype=torch.int32, device="cuda")
- steps = spec.get("snapshot_steps", seqlen)
- window = torch.zeros(
- spec["snapshots"], steps, dim, width - 1, dtype=torch.bfloat16, device="cuda"
- )
- rows_idx = torch.arange(batch, dtype=torch.int32, device="cuda")
- return [
- x,
- conv_state,
- weight,
- None,
- "silu",
- indices,
- PAD_SLOT_ID,
- None,
- window,
- rows_idx,
- False,
- ]
-
-
-def _run(args):
- """Drive the op on copies of the two buffers; return the triple compare takes."""
- got = list(args)
- got[1], got[8] = args[1].clone(), args[8].clone()
- returned = op(*got)
- return returned, got[1], got[8]
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
- x, weight = args[0].clone(), args[2].clone()
- ref = op.reference(*args)
- got = _run(args)
- assert got[0] is not args[0] and got[0].shape == args[0].shape
- assert torch.equal(args[0], x) and torch.equal(args[2], weight), "x and weight are read-only"
- _compare(got, ref)
- if cell.spec["pool_slots"] == ENGINE_SLOTS:
- for slot in cell.spec["slots"]: # non-vacuity of the whole-pool gate
- assert not torch.equal(got[1][slot], args[1][slot]), slot
-
-
-def test_guard_refuses_undersized_buffers() -> None:
- """A snapshot buffer shallower than seqlen is accepted by the kernel and
- written past the row it was given (measured); the guard refuses it
- before the launch."""
- spec = dict(dim=512, batch=2, seqlen=2, pool_slots=8, slots=(0, 2), snapshots=2)
- args = _build(spec, seed=15)
- args[8] = torch.zeros(2, 1, 512, 3, dtype=torch.bfloat16, device="cuda")
- with pytest.raises(AssertionError, match="intermediate_conv_window must be at least"):
- with validating(op):
- op(*args)
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
deleted file mode 100644
index bbd3b86e9575..000000000000
--- 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
+++ /dev/null
@@ -1,111 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the causal_conv1d_update entry."""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_update import ( # noqa: E501
- ENGINE_SLOTS,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_update import ( # noqa: E501
- causal_conv1d_update as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "causal_conv1d_update requires a CUDA device"
-
-
-def _compare(
- out: tuple[torch.Tensor, torch.Tensor], ref: tuple[torch.Tensor, torch.Tensor]
-) -> None:
- """`(x, conv_state)` pairs: the returned `x` with the pool as the call
- left it, against `reference`'s pair.
-
- Activation at torch's default band: bf16/fp16 were measured bitwise
- equal to the fp32 reference in every cell; fp32 differs by tap order
- (max 1.43e-6, inside the default fp32 atol). The pool is a copy of
- input values, so it is gated bitwise -- over the whole pool, which is
- the lock on every slot the batch did not name.
- """
- (act, state), (ref_act, ref_state) = out, ref
- assert act.dtype == ref_act.dtype and act.shape == ref_act.shape
- torch.testing.assert_close(act, ref_act)
- assert torch.equal(state, ref_state), "the advanced state must be bitwise exact"
-
-
-PAD_SLOT_ID = -1
-#: Rows of prefill ahead of the decode rows in the in_proj output.
-_PREFILL_ROWS = 37
-#: Columns past q|k|v in the target's one GDN input-projection output: z | a | b.
-_IN_PROJ_TAIL = 6144 + 2 * 48
-
-
-def _build(spec, seed):
- """`__call__`'s arguments for one spec, plus the buffers x and the pool view."""
- g = torch.Generator(device="cuda").manual_seed(seed)
- dim, batch, layers = spec["dim"], spec["batch"], spec.get("pool_layers", 1)
-
- def randn(*shape):
- 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)
- else:
- indices = torch.tensor(slots, dtype=torch.int32, device="cuda")
- 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)
- 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"
- 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]
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- args, extras = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
- weight = args[2].clone()
- ref = op.reference(*args)
- returned, state = _run(args, extras)
- assert torch.equal(args[2], weight), "weight is read-only"
- _compare((returned, state), ref)
- if cell.spec["pool_slots"] == ENGINE_SLOTS:
- # Non-vacuity for the large pool: every named slot moved.
- for slot in cell.spec["slots"]:
- if slot != PAD_SLOT_ID:
- assert not torch.equal(state[slot], args[1][slot]), slot
-
-
-def test_op_rejects_metadata_violations_loudly() -> None:
- """Why the entry has no `is_valid`: the op raises on metadata violations
- itself (probed here at int64 slot indices)."""
- spec = dict(dim=512, batch=4, pool_slots=16, slots=(0, 1, 2, 3))
- args = _build(spec, seed=14)[0]
- args[6] = args[6].long()
- with pytest.raises(RuntimeError, match="conv_state_indices.scalar_type"):
- with validating(op):
- op(*args)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_chunk_gated_delta_rule.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_chunk_gated_delta_rule.py
deleted file mode 100644
index 6a8cf2fa2286..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_chunk_gated_delta_rule.py
+++ /dev/null
@@ -1,242 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the chunk_gated_delta_rule entry.
-
-The first indexed in-place call per autotune key in a process must be as correct as a
-warm one: the state kernel's autotuner restores the indexed state its benchmark launches
-write, and the pool's dtype is part of the key. "First call in the process" is a property
-of the process, so the cold-call test re-runs this file as a child:
-`python `.
-"""
-
-import os
-import subprocess
-import sys
-import zlib
-from pathlib import Path
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.chunk_gated_delta_rule import ( # noqa: E501
- ENGINE_SLOTS,
- HEAD_K_DIM,
- HEAD_V_DIM,
- NUM_K_HEADS,
- NUM_V_HEADS,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.chunk_gated_delta_rule import ( # noqa: E501
- chunk_gated_delta_rule as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "chunk_gated_delta_rule requires a CUDA device"
-
-
-#: `_compare`'s band; see its docstring.
-RTOL, ATOL_OUTPUT, ATOL_STATE = 1.6e-2, 2e-3, 2e-2
-
-
-def _compare(out: tuple, ref: tuple) -> None:
- """`(o, final_state, pool)`: the call's two returns plus the
- `initial_state` it was handed, against `reference`'s triple; `None`
- components must match.
-
- Wider than bf16's default because the op is a chunked reformulation
- whose WY operands are rounded to bf16 -- the same recurrence, not the
- same sum. rtol stays at the bf16 default; atol is 2e-3 on `o` and 2e-2
- on states, about 4x and 6x the largest deviations measured warm over
- every cell against this reference (`o` 4.88e-4, one bf16 ulp, on
- magnitudes up to 7.9e-2; states 3.57e-3 on up to 8.3e-1).
- Discriminated when measured: seven wrong references (no gating, beta = 1, q/k swapped,
- initial state dropped, scale = 1, key heads reversed, `[K, V]`
- transposed state) all fail the `o` gate, the closest 11x outside it.
- """
- for got, want, atol in zip(out, ref, (ATOL_OUTPUT, ATOL_STATE, ATOL_STATE)):
- assert (got is None) == (want is None), (got is None, want is None)
- if got is None:
- continue
- assert got.dtype == want.dtype and got.shape == want.shape, (
- (got.dtype, got.shape),
- (want.dtype, want.shape),
- )
- torch.testing.assert_close(got.float(), want.float(), rtol=RTOL, atol=atol)
-
-
-DEV = "cuda"
-
-
-def _l2norm(x: torch.Tensor) -> torch.Tensor:
- xf = x.float()
- return xf / torch.sqrt(xf.pow(2).sum(-1, keepdim=True) + 1e-6)
-
-
-def _operands(seqlens: tuple, seed: int):
- """Operands in the domain the GDN prefill path produces: unit-norm q/k,
- `g = -exp(A_log) * softplus(a + dt_bias) < 0`, `beta` a sigmoid. The delta
- rule is only contractive there."""
- gen = torch.Generator(device=DEV).manual_seed(seed)
- total = sum(seqlens)
- cu_seqlens = torch.tensor(
- [0] + torch.tensor(seqlens).cumsum(0).tolist(), dtype=torch.long, device=DEV
- )
- shape_qk = (1, total, NUM_K_HEADS, HEAD_K_DIM)
- q = torch.randn(shape_qk, generator=gen, device=DEV)
- k = torch.randn(shape_qk, generator=gen, device=DEV)
- q, k = _l2norm(q).to(torch.bfloat16), _l2norm(k).to(torch.bfloat16)
- v = torch.randn(
- 1, total, NUM_V_HEADS, HEAD_V_DIM, generator=gen, device=DEV, dtype=torch.bfloat16
- )
- a_log = torch.log(torch.rand(NUM_V_HEADS, generator=gen, device=DEV) * 15 + 1)
- dt_bias = torch.randn(NUM_V_HEADS, generator=gen, device=DEV)
- a = torch.randn(1, total, NUM_V_HEADS, generator=gen, device=DEV)
- g = (-torch.exp(a_log) * torch.nn.functional.softplus(a + dt_bias)).float()
- beta = torch.rand(1, total, NUM_V_HEADS, generator=gen, device=DEV).sigmoid().float()
- return q, k, v, g, beta, cu_seqlens
-
-
-def _pool(num_slots: int, seed: int, dtype: torch.dtype = torch.float32) -> torch.Tensor:
- """A recurrent-state pool, `[slots, H, V, K]` (K innermost), at `dtype`."""
- gen = torch.Generator(device=DEV).manual_seed(seed)
- shape = (num_slots, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM)
- return (torch.randn(shape, generator=gen, device=DEV) * 0.1).to(dtype)
-
-
-def _gate_output(o: torch.Tensor, ref: tuple) -> None:
- """`_compare`, on `o` alone."""
- _compare((o, None, None), (ref[0], None, None))
-
-
-def _gate_pool(pool: torch.Tensor, ref: tuple) -> None:
- """`_compare`, on the pool alone."""
- _compare((ref[0], None, pool), (ref[0], None, ref[2]))
-
-
-def _build(spec, seed):
- """`__call__`'s 13 arguments for one spec."""
- q, k, v, g, beta, cu_seqlens = _operands(spec["seqlens"], seed)
- pool = _pool(ENGINE_SLOTS, seed=seed + 1, dtype=spec.get("pool_dtype", torch.float32))
- indices = torch.tensor(spec["indices"], dtype=torch.int32, device=DEV)
- output = None
- if spec.get("output") == "buffer":
- shape_o = (q.shape[0], q.shape[1], NUM_V_HEADS, HEAD_V_DIM)
- output = torch.empty(shape_o, device=DEV, dtype=torch.bfloat16)
- return [q, k, v, g, beta, None, pool, indices, True, False, cu_seqlens, False, output]
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31)
- pool = args[6]
- pool_before = pool.clone()
- ref = op.reference(*args)
-
- o, final_state = op(*args)
-
- _compare((o, final_state, pool), ref)
- if args[12] is not None:
- assert o.data_ptr() == args[12].data_ptr(), "output= must be written in place"
- named = args[7].tolist()
- unnamed = [s for s in range(pool.shape[0]) if s not in set(named)]
- assert torch.equal(pool[unnamed], pool_before[unnamed]), "unnamed slots written"
- assert not torch.equal(pool[named], pool_before[named]), "named slots did not move"
-
-
-def test_guard_refuses_a_noncontiguous_output() -> None:
- """Unguarded, a head-slice view of a wider buffer is written through
- shape-derived strides and the result is wrong."""
- q, k, v, g, beta, cu = _operands((96,), seed=61)
- ref = op.reference(q, k, v, g, beta, cu_seqlens=cu)
- head_slice = torch.zeros(1, 96, 2 * NUM_V_HEADS, HEAD_V_DIM, device=DEV, dtype=torch.bfloat16)[
- :, :, :NUM_V_HEADS
- ]
- assert not head_slice.is_contiguous()
- o, _ = op(q, k, v, g, beta, cu_seqlens=cu, output=head_slice)
- with pytest.raises(AssertionError):
- _gate_output(o, ref)
- with pytest.raises(AssertionError, match="output must be contiguous"):
- with validating(op):
- op(q, k, v, g, beta, cu_seqlens=cu, output=head_slice)
-
-
-# ---------------------------------------------------------------------------
-# The cold call. The cells above have already benchmarked this autotune key by
-# the time pytest reaches it, so it runs in a fresh interpreter.
-
-CHILD_TIMEOUT_S = 420
-#: The GDN prefill call shape, packed varlen, state by slot index.
-CHILD_SEQLENS = (256, 96)
-
-
-def _state_kernel_autotuner():
- """The `triton.autotune` wrapper around `chunk_gated_delta_rule_fwd_kernel_h_blockdim64`;
- a key in its `cache` is a key whose benchmark launches already happened."""
- from tensorrt_llm._torch.modules.fla.chunk_delta_h import (
- chunk_gated_delta_rule_fwd_kernel_h_blockdim64 as node,
- )
-
- while node is not None:
- if isinstance(getattr(node, "cache", None), dict):
- return node
- node = getattr(node, "fn", None)
- raise AssertionError("no triton autotuner found around the chunked state kernel")
-
-
-def _checkpoint_case(seed: int = 101, pool_dtype: torch.dtype = torch.float32):
- q, k, v, g, beta, cu = _operands(CHILD_SEQLENS, seed=seed)
- pool = _pool(4, seed=seed + 1, dtype=pool_dtype)
- indices = torch.tensor([2, 0], dtype=torch.int32, device=DEV)
- buf = torch.empty(
- 1, sum(CHILD_SEQLENS), NUM_V_HEADS, HEAD_V_DIM, device=DEV, dtype=torch.bfloat16
- )
- return (q, k, v, g, beta, None, pool, indices, True, False, cu, False, buf)
-
-
-def _child_cold_call(pool_dtype: torch.dtype = torch.float32) -> None:
- autotuner = _state_kernel_autotuner()
- assert not autotuner.cache, "the child must be the first indexed call in its process"
- args = _checkpoint_case(pool_dtype=pool_dtype)
- ref = op.reference(*args)
- o, _ = op(*args)
- assert len(autotuner.cache) == 1, "the call must have benchmarked exactly one key"
- _gate_output(o, ref)
- _gate_pool(args[6], ref)
-
-
-_CHILD_MODES = {
- "cold-call-bf16-pool": lambda: _child_cold_call(torch.bfloat16),
-}
-
-
-def _run_child(mode: str) -> None:
- # `_validating` lives four directories up; the child has no pytest to
- # resolve `__extra_import_path__` for it.
- root = str(Path(__file__).resolve().parents[4])
- existing = os.environ.get("PYTHONPATH")
- env = dict(os.environ, PYTHONPATH=root if not existing else os.pathsep.join([root, existing]))
- proc = subprocess.run(
- [sys.executable, __file__, mode],
- capture_output=True,
- text=True,
- timeout=CHILD_TIMEOUT_S,
- env=env,
- )
- print(proc.stdout)
- assert proc.returncode == 0, (
- f"child `{mode}` exited {proc.returncode}\n--- stdout ---\n{proc.stdout}\n"
- f"--- stderr ---\n{proc.stderr}"
- )
-
-
-def test_cold_indexed_call_over_a_bf16_pool() -> None:
- """The first indexed call of a process at 16 / 48 heads, K = V = 128, over a bf16
- pool, in the mixer's shape (256 + 96 packed, by slot, into `output=`), matches the
- reference."""
- _run_child("cold-call-bf16-pool")
-
-
-if __name__ == "__main__":
- _CHILD_MODES[sys.argv[1]]()
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_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
deleted file mode 100644
index 64926845e39c..000000000000
--- 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
+++ /dev/null
@@ -1,105 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the cu_seqlens_from_seq_lens entry.
-
-Integer arithmetic, gated bitwise. A near-miss vector is monotonic and
-plausible, so each cell also proves the vector actually splits a packed
-buffer.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.cu_seqlens_from_seq_lens import ( # noqa: E501
- cu_seqlens_from_seq_lens as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "cu_seqlens_from_seq_lens requires a CUDA device"
-
-
-def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
- """Bitwise: integer arithmetic, no tolerance to loosen."""
- assert out.dtype == ref.dtype, (out.dtype, ref.dtype)
- assert out.shape == ref.shape, (out.shape, ref.shape)
- assert out.is_contiguous()
- assert torch.equal(out, ref), f"{out.tolist()} != {ref.tolist()}"
-
-
-DEV = "cuda"
-
-
-def _seq_lens(lengths, dtype=torch.int32):
- return torch.tensor(list(lengths), dtype=dtype, device=DEV)
-
-
-def _build(spec):
- lengths, dtype = spec["lengths"], spec["in_dtype"]
- layout = spec.get("layout", "contiguous")
- if layout == "prefix":
- # The metadata buffer is sized for max_num_requests; the call site
- # passes its first `batch` entries.
- buf = torch.full((len(lengths) + 9,), 99, dtype=dtype, device=DEV)
- buf[: len(lengths)] = _seq_lens(lengths, dtype)
- return buf[: len(lengths)]
- return _seq_lens(lengths, dtype)
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- spec = cell.spec
- seq_lens = _build(spec)
- before = seq_lens.clone()
- out = op(seq_lens, spec["out_dtype"])
- assert torch.equal(seq_lens, before), "seq_lens is read-only"
- assert out.device.type == "cuda"
- _compare(out, op.reference(seq_lens, spec["out_dtype"]))
-
- # The entry's purpose: slicing a packed buffer with the vector recovers
- # exactly the sequences packed into it.
- lengths = spec["lengths"]
- g = torch.Generator(device=DEV).manual_seed(zlib.crc32(cell.why.encode()))
- pieces = [torch.randint(0, 10_000, (n,), generator=g, device=DEV) for n in lengths]
- packed = torch.cat(pieces) if pieces else torch.empty(0, dtype=torch.int64, device=DEV)
- bounds = out.tolist()
- assert bounds[-1] == packed.shape[0]
- for s, piece in enumerate(pieces):
- assert torch.equal(packed[bounds[s] : bounds[s + 1]], piece), f"sequence {s} mis-sliced"
-
-
-def test_cuda_graph_capture_and_replay() -> None:
- """A captured graph follows lengths rewritten in place: only the shape
- reaches Python, the values are read on device and never synchronized on.
- 32 requests of 8 rows (the MTP-7 verify's 256 rows)."""
- batch, p = 32, 8
- static_lens = _seq_lens([p] * batch)
- side = torch.cuda.Stream()
- side.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(side):
- for _ in range(3):
- op(static_lens, torch.int32)
- torch.cuda.current_stream().wait_stream(side)
- torch.cuda.synchronize()
-
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- static_out = op(static_lens, torch.int32)
- for lengths in ([1] * batch, [p] * batch, [p + 1] + [1] * (batch - 1)):
- static_lens.copy_(_seq_lens(lengths))
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(static_out, op.reference(static_lens, torch.int32))
-
-
-def test_guard_refuses_non_integer_seq_lens() -> None:
- """Unguarded, a float32 seq_lens is truncated toward zero before the scan."""
- floats = torch.tensor([7.9, 1.2, 13.6], dtype=torch.float32, device=DEV)
- assert op(floats, torch.int32).tolist() == [0, 7, 8, 21], "hazard gone"
- with pytest.raises(AssertionError, match="seq_lens must have an integer dtype"):
- with validating(op):
- op(floats, torch.int32)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_dilated_short_conv_paged.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_dilated_short_conv_paged.py
deleted file mode 100644
index 8908d558b5ad..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_dilated_short_conv_paged.py
+++ /dev/null
@@ -1,147 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the dilated_short_conv_paged entry."""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.dilated_short_conv_paged import ( # noqa: E501
- ENGINE_SLOTS,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.dilated_short_conv_paged import ( # noqa: E501
- dilated_short_conv_paged as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "dilated_short_conv_paged requires a CUDA device"
-
-
-def _compare(
- out: tuple[torch.Tensor, torch.Tensor], ref: tuple[torch.Tensor, torch.Tensor]
-) -> None:
- """`(out, conv_states)` pairs: the returned activation with the pool as
- the call left it, against `reference`'s pair.
-
- Activation at torch's default band: bf16 and fp16 were measured bitwise
- equal to the reference (at 10240 channels and 461 slots too); fp32
- differs only by tap order. The pool is a copy, gated bitwise over the
- whole pool -- the lock on every slot the batch did not name.
- """
- (act, pool), (ref_act, ref_pool) = out, ref
- assert act.dtype == ref_act.dtype and act.shape == ref_act.shape
- assert act.is_contiguous()
- torch.testing.assert_close(act, ref_act)
- assert torch.equal(pool, ref_pool), "the published state must be bitwise exact"
-
-
-KERNEL_SIZE = 4
-DILATION = 3
-STATE_LEN = (KERNEL_SIZE - 1) * DILATION
-PAD_SLOT_ID = -1
-
-
-def _build(spec, seed):
- """`__call__`'s positional arguments for one spec, freshly allocated, in the
- call site's form: `x` a bare `.t()` of a token-major activation, int64
- `query_start_loc`, int32 slots.
-
- The checkpoint's taps are small (its four positions average 1e-4 .. 0.0185),
- so the weight is scaled down rather than left at unit variance.
- """
- g = torch.Generator(device="cuda").manual_seed(seed)
- channels, lengths = spec["channels"], spec["lengths"]
- num_seqs = len(lengths)
-
- def randn(*shape, scale=1.0):
- return (scale * torch.randn(*shape, generator=g, device="cuda")).to(torch.bfloat16)
-
- starts = [0]
- for length in lengths:
- starts.append(starts[-1] + length)
- qsl = torch.tensor(starts, dtype=torch.int64, device="cuda")
- x = randn(starts[-1], channels).t()
- weight = randn(channels, 1, KERNEL_SIZE, scale=0.3)
- pool = randn(spec.get("pool_slots", 8), channels, STATE_LEN)
- slots = spec.get("slots") or tuple(range(num_seqs))
- cache_indices = torch.tensor(slots, dtype=torch.int32, device="cuda")
- initial = spec.get("initial", (True,) * num_seqs)
- his = torch.tensor(initial, dtype=torch.bool, device="cuda")
- return [x, weight, None, pool, qsl, cache_indices, his, DILATION, True, PAD_SLOT_ID]
-
-
-def _run(args):
- """Drive the op on a copy of the pool; return (out, pool after)."""
- got = list(args)
- got[3] = args[3].clone()
- return op(*got), got[3]
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
- x, weight = args[0].clone(), args[1].clone()
- ref = op.reference(*args)
- out, pool = _run(args)
- assert out.data_ptr() != args[0].data_ptr()
- assert torch.equal(args[0], x) and torch.equal(args[1], weight), "x and weight are read-only"
- _compare((out, pool), ref)
- if cell.spec.get("pool_slots") == ENGINE_SLOTS:
- for slot in cell.spec["slots"]: # non-vacuity of the whole-pool gate
- assert not torch.equal(pool[slot], args[3][slot]), slot
-
-
-def _small(lengths, slots, seed, **spec):
- return _build(dict(channels=256, lengths=lengths, slots=slots, **spec), seed=seed)
-
-
-def test_cuda_graph_capture_and_replay() -> None:
- """Captured once at the MTP-7 verify's 32 requests of 8 tokens, replayed after
- the batch tensors' values change in place."""
- batch, p = 32, 8
- args = _small((p,) * batch, tuple(range(batch)), seed=25, pool_slots=batch + 2)
- x, weight, _, pool, qsl, cidx, his = args[:7]
- live = pool.clone()
- stream = torch.cuda.Stream()
- stream.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(stream):
- for _ in range(3):
- op(x, weight, None, live, qsl, cidx, his, DILATION, True, PAD_SLOT_ID)
- torch.cuda.current_stream().wait_stream(stream)
-
- graph = torch.cuda.CUDAGraph()
- live.copy_(pool)
- with torch.cuda.graph(graph):
- captured = op(x, weight, None, live, qsl, cidx, his, DILATION, True, PAD_SLOT_ID)
-
- g = torch.Generator(device="cuda").manual_seed(26)
- # Every row but the last shrinks to one token; the slot map reverses and shifts by one.
- qsl.copy_(torch.tensor([*range(batch), batch * p], dtype=torch.int32, device="cuda"))
- cidx.copy_(torch.tensor(list(range(batch, 0, -1)), dtype=torch.int32, device="cuda"))
- his.copy_(torch.tensor([i % 2 == 0 for i in range(batch)], dtype=torch.bool, device="cuda"))
- x.copy_(torch.randn(x.shape, generator=g, device="cuda").to(x.dtype))
- live.copy_(pool)
- graph.replay()
- torch.cuda.synchronize()
-
- ref_out, ref_pool = op.reference(
- x, weight, None, pool, qsl, cidx, his, DILATION, True, PAD_SLOT_ID
- )
- torch.testing.assert_close(captured, ref_out)
- assert torch.equal(live, ref_pool)
-
-
-def test_mismatched_operands_raise() -> None:
- """Refused loudly by the torch calls the entry is composed from -- why it
- has no `is_valid`. Probed at a pool of another dtype (the index copy)."""
- args = _small((4, 4), None, seed=27)
- args[3] = args[3].float()
- with pytest.raises(RuntimeError) as excinfo:
- with validating(op):
- op(*args)
- fragment = "index_copy_(): self and source expected to have the same dtype"
- assert fragment in str(excinfo.value), str(excinfo.value)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_extract_transpose_prefill_slice.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_extract_transpose_prefill_slice.py
deleted file mode 100644
index d0e4fff2e5ad..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_extract_transpose_prefill_slice.py
+++ /dev/null
@@ -1,65 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the extract_transpose_prefill_slice entry."""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.extract_transpose_prefill_slice import ( # noqa: E501
- extract_transpose_prefill_slice as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "extract_transpose_prefill_slice requires a CUDA device"
-
-
-def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
- """Bitwise: the kernel moves elements with no cast and no arithmetic,
- so any tolerance would only hide a wrong index. Compared as raw bytes so
- dtypes without an equality kernel (`float8_e4m3fn`) are gated the same."""
- assert out.dtype == ref.dtype, (out.dtype, ref.dtype)
- assert out.shape == ref.shape, (out.shape, ref.shape)
- assert out.is_contiguous()
- if out.numel():
- assert torch.equal(out.reshape(-1).view(torch.uint8), ref.reshape(-1).view(torch.uint8))
-
-
-def _src(num_rows: int, num_cols: int, dtype: torch.dtype, seed: int) -> torch.Tensor:
- g = torch.Generator(device="cuda").manual_seed(seed)
- return torch.randn(num_rows, num_cols, generator=g, device="cuda").to(dtype)
-
-
-def _build(spec, seed):
- src = _src(spec["num_rows"], spec["num_cols"], spec["dtype"], seed)
- if "row_stop" in spec:
- src = src[: spec["row_stop"]]
- if "col_slice" in spec:
- lo, hi = spec["col_slice"]
- src = src[:, lo:hi]
- assert not src.is_contiguous() and src.stride(1) == 1
- return src, spec["n"], spec["start_col"], spec["width"]
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- src, n, start_col, width = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
- before = src.contiguous().view(torch.uint8).clone()
- out = op(src, n, start_col, width)
- assert out.data_ptr() != src.data_ptr() or out.numel() == 0
- assert torch.equal(src.contiguous().view(torch.uint8), before), "src is read-only"
- _compare(out, op.reference(src, n, start_col, width))
-
-
-def test_guard_refuses_a_range_outside_src() -> None:
- """The kernel masks only against the requested extents and reads past
- `src` rather than failing; the guard bounds both axes (probed here at 128
- rows of a 64-row `src`)."""
- src = torch.randn(64, 256, dtype=torch.bfloat16, device="cuda")
- with pytest.raises(AssertionError, match="outside src"):
- with validating(op):
- op(src, 128, 0, 256)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py
deleted file mode 100644
index c02361fd0d2e..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py
+++ /dev/null
@@ -1,143 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the flashinfer_chunk_gated_delta_rule entry (FlashInfer GDN prefill)."""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.flashinfer_chunk_gated_delta_rule import ( # noqa: E501
- flashinfer_chunk_gated_delta_rule as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "flashinfer_chunk_gated_delta_rule requires a CUDA device"
-
-DEV = "cuda"
-#: The band the Triton chunk op is certified at (rtol bf16 default; atol 2e-3
-#: on `o`, 2e-2 on states), derived for a chunked WY reformulation with bf16
-#: operands. This op is the same class of sum: measured against fp64 at
-#: 1 x 8192 tokens, o max abs 1.26e-3 and state max abs 8.6e-3 (Triton op:
-#: 1.13e-3 / 7.7e-3), inside it.
-RTOL, ATOL_OUTPUT, ATOL_STATE = 1.6e-2, 2e-3, 2e-2
-
-
-def _compare(out: tuple, ref: tuple) -> None:
- """`(o, final_state, pool)` against `reference`'s triple; `None` must match."""
- for got, want, atol in zip(out, ref, (ATOL_OUTPUT, ATOL_STATE, ATOL_STATE)):
- assert (got is None) == (want is None), (got is None, want is None)
- if got is None:
- continue
- assert got.dtype == want.dtype and got.shape == want.shape, (
- (got.dtype, got.shape),
- (want.dtype, want.shape),
- )
- assert torch.isfinite(got).all()
- torch.testing.assert_close(got.float(), want.float(), rtol=RTOL, atol=atol)
-
-
-def _build(spec, seed):
- """Keyword arguments in the domain a GDN post-conv step produces:
- unit-norm q/k, g < 0 in log space, beta a sigmoid."""
- gen = torch.Generator(device=DEV).manual_seed(seed)
- num_k_heads, num_v_heads, head_k_dim, head_v_dim = spec["geometry"]
- seqlens = spec["seqlens"]
- total = sum(seqlens)
- q = torch.randn(1, total, num_k_heads, head_k_dim, generator=gen, device=DEV)
- k = torch.randn(1, total, num_k_heads, head_k_dim, generator=gen, device=DEV)
- q = q / torch.sqrt(q.pow(2).sum(-1, keepdim=True) + 1e-6)
- k = k / torch.sqrt(k.pow(2).sum(-1, keepdim=True) + 1e-6)
- q, k = q.bfloat16(), k.bfloat16()
- v = torch.randn(
- 1, total, num_v_heads, head_v_dim, generator=gen, device=DEV, dtype=torch.bfloat16
- )
- a_log = torch.log(torch.rand(num_v_heads, generator=gen, device=DEV) * 15 + 1)
- x = torch.randn(1, total, num_v_heads, generator=gen, device=DEV) + torch.randn(
- num_v_heads, generator=gen, device=DEV
- )
- g = (-torch.exp(a_log) * torch.nn.functional.softplus(x)).float()
- beta = torch.rand(1, total, num_v_heads, generator=gen, device=DEV).sigmoid().float()
- pool = (
- torch.randn(spec["slots"], num_v_heads, head_v_dim, head_k_dim, generator=gen, device=DEV)
- * 0.1
- ).to(spec.get("pool_dtype", torch.float32))
- cu = torch.tensor([0] + torch.tensor(seqlens).cumsum(0).tolist(), dtype=torch.long, device=DEV)
- indices = torch.tensor(spec["indices"], dtype=torch.int32, device=DEV)
- output = None
- if spec.get("output") == "buffer":
- output = torch.full(
- (1, total, num_v_heads, head_v_dim), -7.0, device=DEV, dtype=torch.bfloat16
- )
- extra = {}
- if "initial" in spec:
- extra["has_initial_state"] = torch.tensor(spec["initial"], dtype=torch.bool, device=DEV)
- if spec.get("linear_gate"):
- extra["linear_gate"] = True
- g = torch.exp(g)
- return dict(
- **extra,
- q=q,
- k=k,
- v=v,
- g=g,
- beta=beta,
- scale=None,
- initial_state=pool,
- initial_state_indices=indices,
- inplace_indexed_state_update=True,
- output_final_state=False,
- cu_seqlens=cu,
- head_first=False,
- use_qk_l2norm_in_kernel=False,
- output=output,
- state_workspace=None,
- )
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
- pool = args["initial_state"]
- pool_before = pool.clone()
- ref = op.reference(**args)
-
- o, final_state = op(**args)
-
- _compare((o, final_state, pool), ref)
- if args["output"] is not None:
- assert o.data_ptr() == args["output"].data_ptr(), "output= must be written in place"
- named = set(args["initial_state_indices"].tolist())
- unnamed = [s for s in range(pool.shape[0]) if s not in named]
- assert torch.equal(pool[unnamed], pool_before[unnamed]), "unnamed slots written"
-
-
-_CARRIED = next(c for c in op.CELLS if "initial" in c.spec)
-
-
-def test_has_initial_state_is_the_reset_then_the_plain_call() -> None:
- """The flags replace `reset_recurrent_state_rows` before the call: zeroing the fresh
- slots and running the plain in-place call gives the same output and pool, bitwise; a
- linear gate that is `torch.exp` of the plain call's log gate changes nothing."""
- args = _build(_CARRIED.spec, seed=5)
- flags = args.pop("has_initial_state")
- linear_gate = args.pop("linear_gate", False)
- pool = args["initial_state"]
- fresh = args["initial_state_indices"][~flags].long()
- pool[fresh] = torch.randn_like(pool[fresh]) # stale rows the flags must not read
- plain = dict(args, initial_state=pool.clone(), output=args["output"].clone())
- plain["initial_state"][fresh] = 0.0
- if linear_gate:
- plain["g"] = torch.log(args["g"])
- # The linear gate of the target call is torch.exp of the log gate, bitwise.
- args["g"] = torch.exp(plain["g"])
-
- o_plain, _ = op(**plain)
- o, final_state = op(**args, has_initial_state=flags, linear_gate=linear_gate)
-
- assert final_state is None
- assert torch.equal(o, o_plain)
- assert torch.equal(pool, plain["initial_state"])
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_gating.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_gating.py
deleted file mode 100644
index 4022900963ef..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_gating.py
+++ /dev/null
@@ -1,115 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the fused_gdn_gating entry.
-
-Geometry is a Qwen3.8-Flash-Next linear-attention layer: 48 value heads, `a`
-bf16, `A_log` / `dt_bias` fp32 drawn across the checkpoint's own ranges.
-"""
-
-import math
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_gdn_gating import ( # noqa: E501
- A_LOG_RANGE,
- APPROX_ULP,
- BA_WIDTH,
- DT_BIAS_RANGE,
- NUM_V_HEADS,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_gdn_gating import ( # noqa: E501
- fused_gdn_gating as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "fused_gdn_gating requires a CUDA device"
-
-
-#: The softplus `beta` the absolute floor is sized for; the tail error scales as
-#: 1 / beta. The target calls 1.0, so 0.5 leaves a factor 2 of headroom.
-_MIN_CERTIFIED_BETA = 0.5
-
-
-def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
- """Both fp32. `g` multiplies two approximate transcendentals, so the
- relative term is `4 * 2^-22`. The absolute floor is the unstabilised
- softplus's tail error -- about one approximate ulp of the log's
- argument -- scaled by `1 / beta` and by the `exp(A_log)` it multiplies;
- `_compare` cannot see either, so it takes exp(A_log) <= 139.42 (the
- checkpoint's extreme) and beta >= 0.5 (`_MIN_CERTIFIED_BETA`): atol =
- 4 * 2^-22 * 139.42 / 0.5 = 2.66e-04.
-
- Measured against this gate on 256 checkpoint-range rows: worst error /
- allowance 0.215. Five wrong formulas (dt_bias dropped, sign dropped,
- A_log without exp, exp(-A_log), sigmoid for softplus) land 8.7e+05 to
- 1.5e+07 allowances away.
- """
- assert out.dtype == ref.dtype == torch.float32, (out.dtype, ref.dtype)
- assert out.shape == ref.shape, (out.shape, ref.shape)
- atol = 4.0 * APPROX_ULP * math.exp(A_LOG_RANGE[1]) / _MIN_CERTIFIED_BETA
- torch.testing.assert_close(out, ref, rtol=4.0 * APPROX_ULP, atol=atol)
-
-
-DEV = "cuda"
-SOFTPLUS_THRESHOLD = 20.0
-
-
-def _params(g, num_heads=NUM_V_HEADS, dtype=torch.float32):
- lo, hi = A_LOG_RANGE
- A_log = torch.rand(num_heads, generator=g, device=DEV) * (hi - lo) + lo
- lo, hi = DT_BIAS_RANGE
- dt_bias = torch.rand(num_heads, generator=g, device=DEV) * (hi - lo) + lo
- return A_log.to(dtype), dt_bias.to(dtype)
-
-
-def _build(spec, seed):
- """`__call__`'s positional arguments for one spec."""
- g = torch.Generator(device=DEV).manual_seed(seed)
- rows, heads = spec["rows"], NUM_V_HEADS
-
- def randn(*shape):
- return torch.randn(*shape, generator=g, device=DEV).to(torch.bfloat16)
-
- if "fill" in spec:
- a = torch.full((rows, heads), spec["fill"], dtype=torch.bfloat16, device=DEV)
- elif spec.get("layout", "contiguous") == "target_rows":
- # `_gdn_block` splits the [T, 48] projection into prefill and decode rows.
- a = randn(rows + 23, heads)[23:]
- assert a.is_contiguous() and a.storage_offset() == 23 * heads
- else:
- a = randn(rows, heads)
- A_log, dt_bias = _params(g, heads)
- return A_log, a, dt_bias, 1.0, SOFTPLUS_THRESHOLD
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
- g = op(*args)
- _compare(g, op.reference(*args))
- assert g.is_contiguous() and torch.isfinite(g).all()
-
-
-def test_guard_refuses_what_the_kernel_reads_at_the_wrong_addresses() -> None:
- """Unguarded, an `a` with head stride 2 returns a finite `g` -- exactly the
- one for the memory the kernel actually addressed (`row * stride(0) +
- head`: the first 48 columns of each row), not the tensor it was handed;
- the guard refuses it."""
- g = torch.Generator(device=DEV).manual_seed(111)
- A_log, dt_bias = _params(g)
- wide = torch.randn(32, BA_WIDTH, generator=g, device=DEV).to(torch.bfloat16)
- bad = (A_log, wide[:, ::2], dt_bias)
- addressed = (A_log, wide[:, :NUM_V_HEADS], dt_bias)
- unguarded = op(*bad)
- assert torch.isfinite(unguarded).all()
- assert torch.equal(unguarded, op(*addressed)), "the hazard this guard covers is gone"
- with pytest.raises(AssertionError):
- _compare(unguarded, op.reference(*bad))
- with pytest.raises(AssertionError, match="head axis must be unit-stride"):
- with validating(op):
- op(*bad)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_post_conv.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_post_conv.py
deleted file mode 100644
index 01b35ad0d378..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_post_conv.py
+++ /dev/null
@@ -1,143 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the fused_gdn_post_conv entry.
-
-Geometry is a Qwen3.8-Flash-Next linear-attention layer: 16 x 128 key heads,
-48 x 128 value heads, qkv_dim 10240, bf16 activations.
-"""
-
-import math
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_gdn_post_conv import ( # noqa: E501
- A_LOG_RANGE,
- DT_BIAS_RANGE,
- HEAD_K_DIM,
- HEAD_V_DIM,
- NUM_K_HEADS,
- NUM_V_HEADS,
- QKV_DIM,
- SOFTPLUS_THRESHOLD,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_gdn_post_conv import ( # noqa: E501
- fused_gdn_post_conv as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "fused_gdn_post_conv requires a CUDA device"
-
-
-#: Relative error of the SASS approximate transcendentals (`ex2.approx.f32`,
-#: `lg2.approx.f32`) Triton lowers `tl.exp` / `tl.log` to.
-_APPROX_ULP = 2.0**-22
-
-
-def _compare(out: tuple, ref: tuple) -> None:
- """Per output, each gate derived rather than tuned.
-
- q, k -- torch's bf16 band: both sides round one fp32 value, so they
- differ by at most one bf16 ulp. v -- bitwise: the kernel moves bits.
- g -- a product of two SASS approximate transcendentals (relative error
- ~2^-22 each) with the softplus magnitude bounded by the threshold and
- every term scaled by `exp(A_log)`:
- `atol = 4 * 2^-22 * softplus_threshold * exp(max A_log)`, `rtol = 8 * 2^-22`.
- At this checkpoint's A_log ceiling that is 6.7e-04 against a measured
- worst case of 7.1e-05; five wrong gate formulas land at 1.2e+02 and up.
- beta -- fp32: torch's fp32 band (one well-conditioned transcendental);
- narrower: bitwise, the fp32 sigmoid rounded exactly once.
- """
- assert len(out) == len(ref) == 5
- for got, want in zip(out, ref):
- assert got.dtype == want.dtype and got.shape == want.shape, (
- (got.dtype, got.shape),
- (want.dtype, want.shape),
- )
- assert got.is_contiguous()
- q, k, v, g, beta = out
- q_ref, k_ref, v_ref, g_ref, beta_ref = ref
- torch.testing.assert_close(q, q_ref)
- torch.testing.assert_close(k, k_ref)
- assert torch.equal(v, v_ref), "v must be a bit-exact copy"
- g_atol = 4 * _APPROX_ULP * SOFTPLUS_THRESHOLD * math.exp(A_LOG_RANGE[1])
- torch.testing.assert_close(g, g_ref, rtol=8 * _APPROX_ULP, atol=g_atol)
- if beta.dtype == torch.float32:
- torch.testing.assert_close(beta, beta_ref)
- else:
- assert torch.equal(beta, beta_ref), "beta must be the fp32 sigmoid rounded once"
-
-
-DEV = "cuda"
-DTYPE = torch.bfloat16
-HEADS = (NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS, HEAD_V_DIM)
-L2_NORM_EPS = 1e-6
-#: The `z` columns of the input projection.
-Z_WIDTH = NUM_V_HEADS * HEAD_V_DIM
-#: The target's one GDN input-projection output: [q|k|v | z | a | b].
-IN_PROJ_WIDTH = QKV_DIM + Z_WIDTH + 2 * NUM_V_HEADS # 16480
-
-
-def _build(spec, seed):
- """`__call__`'s positional arguments (its keyword `beta_dtype` left at fp32)."""
- g = torch.Generator(device=DEV).manual_seed(seed)
- n_p, n_d = spec["num_prefill"], spec["num_decode"]
- n = n_p + n_d
-
- def randn(*shape):
- return torch.randn(*shape, generator=g, device=DEV).to(DTYPE)
-
- # `_gdn_block`: one [T, 16480] GEMM output; q|k|v, a and b are its column
- # views, the decode half the row split of q|k|v, the prefill half the
- # transposing extraction of the rest.
- proj = randn(n, IN_PROJ_WIDTH)
- mixed = proj[:, :QKV_DIM]
- prefill = mixed[:n_p].t().contiguous()
- decode = mixed[n_p:] if n_d else None
- a = proj[:, QKV_DIM + Z_WIDTH : QKV_DIM + Z_WIDTH + NUM_V_HEADS]
- b = proj[:, QKV_DIM + Z_WIDTH + NUM_V_HEADS :]
- assert a.stride(0) == b.stride(0) == IN_PROJ_WIDTH
- assert decode is None or decode.stride() == (IN_PROJ_WIDTH, 1)
-
- lo, hi = A_LOG_RANGE
- A_log = torch.rand(NUM_V_HEADS, generator=g, device=DEV) * (hi - lo) + lo
- lo, hi = DT_BIAS_RANGE
- dt_bias = torch.rand(NUM_V_HEADS, generator=g, device=DEV) * (hi - lo) + lo
- if "dt_bias_fill" in spec:
- dt_bias = torch.full((NUM_V_HEADS,), spec["dt_bias_fill"], device=DEV)
- return [prefill, decode, a, b, A_log, dt_bias, *HEADS, L2_NORM_EPS, SOFTPLUS_THRESHOLD]
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
- if "dt_bias_fill" in cell.spec:
- # Every head strictly on the side of the threshold the cell names.
- x = args[2].float() + args[5].float()
- above = cell.spec["dt_bias_fill"] > SOFTPLUS_THRESHOLD
- assert bool((x > SOFTPLUS_THRESHOLD).all() if above else (x <= SOFTPLUS_THRESHOLD).all())
- before = [t.clone() if isinstance(t, torch.Tensor) else t for t in args]
- out = op(*args)
- for arg, old in zip(args, before):
- if isinstance(arg, torch.Tensor):
- assert torch.equal(arg, old), "inputs are read-only"
- _compare(out, op.reference(*args))
-
-
-def test_guard_refuses_strided_gate_parameters() -> None:
- """Unguarded, a stride-2 `A_log` view is read as `ptr + head` -- the
- interleaved zeros -- and returns a finite wrong g."""
- args = _build(dict(num_prefill=32, num_decode=0), seed=108)
- strided = torch.stack([args[4], torch.zeros_like(args[4])], dim=1)[:, 0]
- assert strided.stride() == (2,)
- bad = list(args)
- bad[4] = strided
- with pytest.raises(AssertionError):
- _compare(op(*bad), op.reference(*bad))
- with pytest.raises(AssertionError, match="must be contiguous 1-D"):
- with validating(op):
- op(*bad)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_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
deleted file mode 100644
index 5ebd22d61d5e..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py
+++ /dev/null
@@ -1,547 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the fused_recurrent_gated_delta_rule_cached_replay_update entry.
-
-The op is stateful across verify steps: what one call leaves in the checkpoint
-pool and the history double buffer is what the next call reads, after the
-engine's own bookkeeping ran in between. So every cell drives several
-consecutive verify steps the way the engine does -- work items before the
-forward (partitioned batches), the op once per layer, then the all-layer
-commit and `_advance_replay_state` with a random partial acceptance -- and
-checks two things at every step:
-
-* each call against the entry's `reference` on that call's own inputs (the
- output and all four state buffers), plus bitwise locks on everything the
- call must not touch;
-* every output against an *exact twin*: the entry's `reference` driven on the
- same checkpoint and history held in fp32 and deep enough never to wrap,
- advanced by `p += accepted` alone. The twin's state is the plain
- token-by-token recurrence over every accepted token, independent of the
- wrap, the double buffer and the commit, so a wrong commit or wrong
- bookkeeping surfaces as an output that drifts off it.
-"""
-
-from __future__ import annotations
-
-import zlib
-
-import pytest
-import torch
-import torch.nn.functional as F
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_recurrent_gated_delta_rule_cached_replay_update import ( # noqa: E501
- ALL_LAYER_COMMIT_MIN_BATCH,
- HEAD_K_DIM,
- HEAD_V_DIM,
- HISTORY_SIZE,
- NUM_K_HEADS,
- NUM_V_HEADS,
- PACKED_QKV_WIDTH,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_recurrent_gated_delta_rule_cached_replay_update import ( # noqa: E501
- fused_recurrent_gated_delta_rule_cached_replay_update as op,
-)
-from tensorrt_llm._torch.modules.fla.cache_manager import (
- ReplayStateUpdateMetadata,
- _advance_replay_state,
-)
-from tensorrt_llm._torch.modules.fla.cached_replay import (
- CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE,
- commit_gdn_cached_replay_history_layers,
-)
-from tensorrt_llm._torch.modules.mamba.mamba2_metadata import _build_replay_work_items_triton
-
-__extra_import_path__ = ["../../../.."]
-from _validating import ULP, certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), (
- "fused_recurrent_gated_delta_rule_cached_replay_update requires a CUDA device"
-)
-
-DEV = "cuda"
-H, HV, K, V = NUM_K_HEADS, NUM_V_HEADS, HEAD_K_DIM, HEAD_V_DIM
-W = PACKED_QKV_WIDTH
-HIST = HISTORY_SIZE
-BF16_ULP = ULP[torch.bfloat16]
-
-#: This checkpoint's 36 GDN layers of `linear_attn.A_log` / `linear_attn.dt_bias`
-#: (bf16, 48 each): mean, standard deviation, range.
-A_LOG_STATS = (1.006, 1.487, (-7.125, 5.0625))
-DT_BIAS_STATS = (-4.411, 1.732, (-11.4375, 2.53125))
-
-#: Element gate of the bf16 members, in bf16 ulps of the block they belong to.
-_ELEMENT_ULPS = 4.0
-#: Relative-RMS gate of the output, in bf16 ulps.
-_OUTPUT_RMS_ULPS = 1.5
-#: Element gate of an fp32 checkpoint state, relative to its (slot, head) scale.
-_STATE_TOL = 2.0**-13
-#: Element gate of a bf16 checkpoint state, in bf16 ulps of its (slot, head) scale.
-_BF16_STATE_ULPS = 2.0
-#: Gate of the stored cumulative log-decays: |dG| <= atol + rtol * |G|.
-_G_ATOL, _G_RTOL = 2.0**-11, 2.0**-18
-
-
-def _gate(name, got, want, element, rms=None):
- """`got` against `want` per block of the last two dims (a token's
- `[heads, dim]`, or one head's `[V, K]` state): every element within
- `element` of its block's absmax, NaN blocks of `want` excluded."""
- block = got.shape[-2] * got.shape[-1]
- g, w = got.float().reshape(-1, block), want.float().reshape(-1, block)
- keep = ~torch.isnan(w).any(1)
- g, w = g[keep], w[keep]
- scale = w.abs().amax(1, keepdim=True).clamp_min(torch.finfo(torch.float32).tiny)
- ratio = ((g - w).abs() / scale).nan_to_num(nan=float("inf"))
- worst = ratio.max().item() if ratio.numel() else 0.0
- assert worst <= element, (
- f"{name}: an element is {worst:.3e} of its block's absmax from the reference "
- f"(gate {element:.3e})"
- )
- if rms is not None and w.numel():
- rel = ((g - w).pow(2).mean().sqrt() / w.pow(2).mean().sqrt().clamp_min(1e-30)).item()
- assert rel <= rms, f"{name}: relative RMS {rel:.3e} exceeds {rms:.3e}"
-
-
-def _compare(out: tuple, ref: tuple) -> None:
- """`(returned, ssm_states, old_u, old_k, old_G)` as the call left them,
- against `reference`'s five; a member `None` on both sides is skipped (the
- exact twin is compared on the output alone).
-
- The kernel is not fp32 throughout, and the bands say where. It rounds the
- L2-normalized q and k to bf16 before every product (2^-9 relative each),
- rounds the history-mixing coefficients to bf16 for its tensor-core dots, and
- stores the output and the history's u and k in bf16. So the output, `old_u`
- and `old_k` are gated per block -- one token's `[heads, dim]` -- at 4 bf16
- ulps of the block's absmax, which keeps a near-silent head from turning
- rounding into relative noise, and the output's relative RMS at 1.5 ulps. An fp32
- checkpoint state is fp32 by construction: the commit splits the decayed
- history update into a bf16 hi and lo part, ~16 significant bits per term, and
- up to 16 terms each within 2^-17 give the 2^-13 gate, relative to one (slot,
- head) state's absmax. A bf16 checkpoint is that fp32 commit rounded once, as the
- reference's is, so the two stored values are equal or adjacent, and adjacent bf16
- values are at most 2 bf16 ulps of their head's absmax apart: the bf16 gate is one
- representable step. `G` only ever enters as exp(G_t - G_j), so an absolute
- error is a relative decay error: the kernel's softplus runs on approximate
- transcendentals whose fp32 sum drops exp(x)'s low bits; atol 2^-11 is an
- eighth of a bf16 ulp of the u it scales, and rtol 2^-18 covers the cumulative
- sum's rounding where |G| reaches thousands on fast-forgetting heads.
- """
- names = ("output", "ssm_states", "old_u", "old_k", "old_G")
- assert len(out) == len(ref) == len(names)
- for name, got, want in zip(names, out, ref):
- assert (got is None) == (want is None), f"{name}: present on one side only"
- if got is None:
- continue
- assert got.dtype == want.dtype and got.shape == want.shape, (
- name,
- (got.dtype, tuple(got.shape)),
- (want.dtype, tuple(want.shape)),
- )
- if name == "old_G":
- torch.testing.assert_close(got, want, rtol=_G_RTOL, atol=_G_ATOL, msg=name)
- elif name == "ssm_states":
- bf16 = got.dtype == torch.bfloat16
- _gate(name, got, want, _BF16_STATE_ULPS * BF16_ULP if bf16 else _STATE_TOL)
- else:
- rms = _OUTPUT_RMS_ULPS * BF16_ULP if name == "output" else None
- _gate(name, got, want, _ELEMENT_ULPS * BF16_ULP, rms)
-
-
-def _seed(why):
- return zlib.crc32(why.encode())
-
-
-def _gating(gen, layers):
- """Per-layer A_log / dt_bias drawn around the checkpoint's own statistics;
- head 0 pinned to the corner that forgets everything every token
- (g ~ -400), head 1 to the one that hardly decays (g ~ -1e-8)."""
- (a_mean, a_std, a_range), (d_mean, d_std, d_range) = A_LOG_STATS, DT_BIAS_STATS
- a_log = torch.randn(layers, HV, generator=gen, device=DEV) * a_std + a_mean
- dt_bias = torch.randn(layers, HV, generator=gen, device=DEV) * d_std + d_mean
- a_log, dt_bias = a_log.clamp(*a_range), dt_bias.clamp(*d_range)
- a_log[:, 0], dt_bias[:, 0] = a_range[1], d_range[1]
- a_log[:, 1], dt_bias[:, 1] = a_range[0], d_range[0]
- return a_log.contiguous(), dt_bias.contiguous()
-
-
-def _unit(x):
- return (x / x.norm(dim=-1, keepdim=True)).to(torch.bfloat16)
-
-
-class _Bench:
- """The state the cache manager keeps for `layers` GDN layers, laid out as
- `MambaHybridCacheManagerV2` lays it out, driven the way the engine drives
- it: work items before the forward (`_build_replay_work_items_triton`, what
- mamba metadata's prepare runs for a partitioned batch), the op per layer,
- then `update_mamba_states`'s order -- the all-layer commit, then
- `_advance_replay_state`.
-
- With `exact=True` it also carries the exact twin (see the module docstring).
- """
-
- def __init__(self, spec, seed, exact=True):
- self.spec = spec
- self.path, self.n, self.t = spec["path"], spec["n"], spec["t"]
- self.layers = spec.get("layers", 2)
- self.slots = spec.get("slots", 64)
- self.steps = spec.get("steps", 8)
- self.gen = torch.Generator(device=DEV).manual_seed(seed)
- self.cpu = torch.Generator().manual_seed(seed)
- n, t, L, S = self.n, self.t, self.layers, self.slots
-
- if spec.get("edge_slots"):
- inner = torch.randperm(S - 2, generator=self.cpu)[: n - 2] + 1
- chosen = torch.cat([torch.tensor([S - 1]), inner, torch.tensor([0])])
- else:
- chosen = torch.randperm(S, generator=self.cpu)[:n]
- chosen = chosen.to(torch.int32)
- live_rows = torch.ones(n, dtype=torch.bool)
- dup = tuple(spec.get("dup", ()))
- if dup:
- chosen[list(dup)] = int(chosen[dup[0]])
- # The padding rows' shared slot is nobody's state: the bookkeeping never names it.
- live_rows[list(dup)] = False
- self.state_indices = chosen.to(DEV)
- self.live = self.state_indices[live_rows.to(DEV)].clone()
-
- # [slots, layers, HV, V, K]: each layer's state is a slot-strided view,
- # the affine layout whose strides the all-layer commit is handed.
- self.backing = (self._randn(S, L, HV, V, K) * 0.5).to(
- spec.get("state_dtype", torch.float32)
- )
- self.pool = [self.backing[:, i] for i in range(L)]
- self.old_x = torch.zeros(L, S, 2, HIST, HV, V, device=DEV, dtype=torch.bfloat16)
- self.old_B = torch.zeros(L, S, 2, HIST, H, K, device=DEV, dtype=torch.bfloat16)
- self.old_dt = torch.zeros(L, S, 2, HV, HIST, device=DEV)
- self.old_dA = self._randn(L, S, 2, HV, HIST) # handed over as old_beta; never read
- self.pnat = torch.zeros(S, dtype=torch.int32, device=DEV)
- self.buf = torch.zeros(S, dtype=torch.int32, device=DEV)
- self.a_log, self.dt_bias = _gating(self.gen, L)
- if spec.get("start", "warm") == "warm":
- self.old_x.copy_(self._randn(L, S, 2, HIST, HV, V) * 0.5)
- self.old_B.copy_(_unit(self._randn(L, S, 2, HIST, H, K)))
- raw_a = self._randn(L, S, 2, HV, HIST)
- g = -torch.exp(self.a_log)[:, None, None, :, None] * F.softplus(
- raw_a + self.dt_bias[:, None, None, :, None]
- )
- self.old_dt.copy_(g.cumsum(-1))
- pnat = torch.randint(0, HIST + 1, (S,), generator=self.cpu, dtype=torch.int32)
- live = self.live.long().cpu()
- pnat[live[0]] = HIST # a wrap at step 0
- if live.numel() > 1:
- pnat[live[1]] = 0 # an append at step 0
- self.pnat.copy_(pnat)
- self.buf.copy_(torch.randint(0, 2, (S,), generator=self.cpu, dtype=torch.int32))
- self.meta = ReplayStateUpdateMetadata(self.pnat, self.buf, t, HIST)
- self.work_items = self.n_writes = None
- if self.path == "all_layer":
- self.work_items = torch.zeros(n, 4, dtype=torch.int32, device=DEV)
- self.n_writes = torch.zeros(1, dtype=torch.int32, device=DEV)
-
- self.exact = exact
- if exact:
- self.depth = HIST + (self.steps + 1) * t
- self.x_pool = [p.clone() for p in self.pool]
- self.x_u = torch.zeros(L, S, 1, self.depth, HV, V, device=DEV)
- self.x_k = torch.zeros(L, S, 1, self.depth, H, K, device=DEV)
- self.x_G = torch.zeros(L, S, 1, HV, self.depth, device=DEV)
- self.x_pnat = self.pnat.clone()
- self.x_buf = torch.zeros_like(self.buf)
- for s in self.live.tolist():
- p, b = int(self.pnat[s]), int(self.buf[s])
- self.x_u[:, s, 0, :p] = self.old_x[:, s, b, :p].float()
- self.x_k[:, s, 0, :p] = self.old_B[:, s, b, :p].float()
- self.x_G[:, s, 0, :, :p] = self.old_dt[:, s, b, :, :p]
-
- def _randn(self, *shape, dtype=torch.float32):
- return torch.randn(*shape, generator=self.gen, device=DEV, dtype=dtype)
-
- def inputs(self):
- """One layer's fresh activations: post-conv `[q|k|v]` rows and the raw
- `a`/`b` projections, as the verify step hands them over."""
- n, t = self.n, self.t
- packed = self._randn(n * t, W, dtype=torch.bfloat16)
- rows = packed.view(n, t, W)
- q = rows[..., : H * K].view(n, t, H, K)
- k = rows[..., H * K : 2 * H * K].view(n, t, H, K)
- v = rows[..., 2 * H * K :].view(n, t, HV, V)
- a = self._randn(n, t, HV, dtype=torch.bfloat16)
- b = self._randn(n, t, HV, dtype=torch.bfloat16)
- return dict(q=q, k=k, v=v, a=a, b=b, packed=packed)
-
- def kwargs(self, layer, act):
- return dict(
- q=act["q"],
- k=act["k"],
- v=act["v"],
- g=act["a"],
- beta=act["b"],
- ssm_states=self.pool[layer],
- state_indices=self.state_indices,
- old_u=self.old_x[layer],
- old_k=self.old_B[layer],
- old_G=self.old_dt[layer],
- old_beta=self.old_dA[layer],
- cache_buf_idx=self.buf,
- prev_num_accepted_tokens=self.pnat,
- history_size=HIST,
- scale=None,
- use_qk_l2norm_in_kernel=True,
- A_log=self.a_log[layer],
- dt_bias=self.dt_bias[layer],
- launch_with_pdl=False,
- replay_work_items=self.work_items,
- n_writes=self.n_writes,
- packed_qkv=act["packed"],
- use_all_layer_commit=self.path == "all_layer",
- output=None,
- )
-
- def exact_kwargs(self, layer, kw):
- return dict(
- kw,
- ssm_states=self.x_pool[layer],
- old_u=self.x_u[layer],
- old_k=self.x_k[layer],
- old_G=self.x_G[layer],
- cache_buf_idx=self.x_buf,
- prev_num_accepted_tokens=self.x_pnat,
- history_size=self.depth,
- replay_work_items=None,
- n_writes=None,
- use_all_layer_commit=False,
- output=None,
- )
-
- def adopt_exact(self, layer, ref):
- assert torch.equal(ref[1], self.x_pool[layer]), "the twin must never wrap"
- self.x_u[layer].copy_(ref[2])
- self.x_k[layer].copy_(ref[3])
- self.x_G[layer].copy_(ref[4])
-
- def state(self, layer):
- return self.pool[layer], self.old_x[layer], self.old_B[layer], self.old_dt[layer]
-
- def prepare(self):
- if self.path == "all_layer":
- self.work_items.zero_()
- _build_replay_work_items_triton(
- self.live, self.pnat, self.buf, self.work_items, self.n_writes, self.t, HIST
- )
-
- def commit(self):
- if self.path == "all_layer":
- commit_gdn_cached_replay_history_layers(
- ssm_states=self.pool[0],
- ssm_state_num_layers=self.layers,
- ssm_state_layer_stride=self.backing.stride(1),
- ssm_state_slot_stride=self.backing.stride(0),
- old_u=self.old_x,
- old_k=self.old_B,
- old_G=self.old_dt,
- replay_work_items=self.work_items,
- n_writes=self.n_writes,
- history_size=HIST,
- )
-
- def draw_acceptance(self):
- """Accepted tokens per live row, bonus token included, by `spec["accept"]`."""
- mode = self.spec.get("accept", "random")
- rows = self.live.numel()
- if mode == "full":
- return torch.full((rows,), self.t, dtype=torch.int32, device=DEV)
- assert mode == "random", mode
- return torch.randint(1, self.t + 1, (rows,), generator=self.cpu, dtype=torch.int32).to(DEV)
-
- def accept(self, accepted):
- _advance_replay_state(self.meta, self.live, accepted)
- if self.exact:
- self.x_pnat[self.live.long()] += accepted
-
- def writes(self):
- """Where this step's calls append, by the documented rule, which slots wrap, and how
- many append exactly up to the history's end: `[slots, 2, HIST]` mask, wrapping slot
- ids, the `pnat + t == HIST` count."""
- s = self.live.long()
- p, b = self.pnat[s].long(), self.buf[s].long()
- wrap = p + self.t > HIST
- half, start = torch.where(wrap, 1 - b, b), torch.where(wrap, 0, p)
- mask = torch.zeros(self.slots, 2, HIST, dtype=torch.bool, device=DEV)
- for t in range(self.t):
- mask[s, half, start + t] = True
- return mask, s[wrap], int((p + self.t == HIST).sum())
-
-
-def _check_duplicate_padding_slot(cell) -> None:
- """The last rows of a batch share one slot, as CUDA-graph padding repeats the manager's
- dummy slot: the rows on unique slots, and every slot but the shared one, are held to
- the reference run on those rows alone -- the padding rows change nothing about them.
- The shared slot's state and the padding rows' outputs are unspecified."""
- bench = _Bench(cell.spec, seed=_seed(cell.why), exact=False)
- dup = cell.spec["dup"]
- unique = bench.n - len(dup)
- assert dup == tuple(range(unique, bench.n)), "the padding rows trail the batch"
- shared = int(bench.state_indices[unique])
- assert (bench.state_indices[unique:] == shared).all()
- assert shared not in bench.live.tolist()
- others = torch.ones(bench.slots, dtype=torch.bool, device=DEV)
- others[shared] = False
- for _ in range(bench.steps):
- for layer in range(bench.layers):
- act = bench.inputs()
- kw = bench.kwargs(layer, act)
- alone = dict(
- kw,
- q=kw["q"][:unique],
- k=kw["k"][:unique],
- v=kw["v"][:unique],
- g=kw["g"][:unique],
- beta=kw["beta"][:unique],
- state_indices=kw["state_indices"][:unique],
- packed_qkv=act["packed"][: unique * bench.t],
- )
- ref = op.reference(**alone)
- got = op(**kw)
- pool, u, k, G = bench.state(layer)
- _compare(
- (got[:unique], pool[others], u[others], k[others], G[others]),
- (ref[0], *(x[others] for x in ref[1:])),
- )
- bench.accept(bench.draw_acceptance())
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- if "dup" in cell.spec:
- return _check_duplicate_padding_slot(cell)
- bench = _Bench(cell.spec, seed=_seed(cell.why))
- wraps = appends = to_the_end = 0
- for _ in range(bench.steps):
- mask, wrapped, at_end = bench.writes()
- wraps += wrapped.numel()
- to_the_end += at_end
- appends += bench.live.numel() - wrapped.numel()
- bench.prepare()
- after_commit = []
- for layer in range(bench.layers):
- act = bench.inputs()
- kw = bench.kwargs(layer, act)
- pool, u, k, G = bench.state(layer)
- before = dict(
- pool=pool.clone(),
- u=u.clone(),
- k=k.clone(),
- G=G.clone(),
- old_beta=bench.old_dA[layer].clone(),
- )
- ref = op.reference(**kw)
- exact = op.reference(**bench.exact_kwargs(layer, kw))
- committed = None
- if bench.path == "all_layer":
- committed = op.reference(
- **dict(kw, replay_work_items=None, n_writes=None, use_all_layer_commit=False)
- )
- got = op(**kw)
- _compare((got, *bench.state(layer)), ref)
- # The locks: in place, and nothing outside the documented writes.
- pool, u, k, G = bench.state(layer)
- assert torch.equal(bench.old_dA[layer], before["old_beta"]), "old_beta is never written"
- assert torch.equal(u[~mask], before["u"][~mask]), "old_u written outside the append"
- assert torch.equal(k[~mask], before["k"][~mask]), "old_k written outside the append"
- Gt, Gt0 = G.transpose(-1, -2), before["G"].transpose(-1, -2)
- assert torch.equal(Gt[~mask], Gt0[~mask]), "old_G written outside the append"
- untouched = torch.ones(bench.slots, dtype=torch.bool, device=DEV)
- if bench.path == "small":
- untouched[wrapped] = False
- assert torch.equal(pool[untouched], before["pool"][untouched]), (
- "the pool moved where the call commits nothing (a partitioned call commits none)"
- )
- bench.adopt_exact(layer, exact)
- after_commit.append((got, exact, committed))
- bench.commit()
- for layer, (got, exact, committed) in enumerate(after_commit):
- if committed is not None:
- # The partitioned call plus the manager's commit is the call
- # that commits in-kernel.
- _compare((got, *bench.state(layer)), committed)
- _compare((got, None, None, None, None), (exact[0], None, None, None, None))
- bench.accept(bench.draw_acceptance())
- assert wraps > 0 and appends > 0, (
- f"the cell never wrapped or never appended ({wraps}, {appends})"
- )
- if cell.spec.get("accept") == "full":
- live = bench.live.numel()
- assert wraps >= 3 * live and to_the_end >= 3 * live, (
- f"full acceptance should wrap every row at least 3 times and append up to the "
- f"history's end as often ({wraps}, {to_the_end} over {live} rows)"
- )
-
-
-# ── CUDA graphs ───────────────────────────────────────────────────────────
-
-
-def test_cuda_graph_replay_on_new_data_is_bitwise_eager() -> None:
- """Two layers of 32 sequences x 8 tokens over an fp32 checkpoint captured with the
- all-layer commit after them; replayed three times on new activations, with acceptance
- in between, against the same steps run eagerly on a twin state."""
- n, t = 32, 8
- assert ALL_LAYER_COMMIT_MIN_BATCH == CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE
- assert n >= ALL_LAYER_COMMIT_MIN_BATCH
- spec = dict(path="all_layer", n=n, t=t, layers=2, slots=40, state_dtype=torch.float32)
- seed = 7000 + 10 * n + t
- eager, graph, warm = (_Bench(spec, seed, exact=False) for _ in range(3))
- acts = [warm.inputs() for _ in range(2)]
- warm.prepare()
- for layer in range(2):
- op(**warm.kwargs(layer, acts[layer]))
- warm.commit()
- torch.cuda.synchronize()
-
- graph.prepare()
- captured = torch.cuda.CUDAGraph()
- with torch.cuda.graph(captured):
- outs_graph = [op(**graph.kwargs(layer, acts[layer])) for layer in range(2)]
- graph.commit()
- assert torch.equal(graph.backing, eager.backing) and torch.equal(graph.old_x, eager.old_x)
-
- refill = torch.Generator(device=DEV).manual_seed(seed + 1)
- for _ in range(3):
- for act in acts:
- for key in ("packed", "a", "b"):
- act[key].copy_(
- torch.randn(act[key].shape, generator=refill, device=DEV, dtype=act[key].dtype)
- )
- eager.prepare()
- outs_eager = [op(**eager.kwargs(layer, acts[layer])) for layer in range(2)]
- eager.commit()
- graph.prepare()
- captured.replay()
- torch.cuda.synchronize()
- for e, g in zip(outs_eager, outs_graph):
- assert torch.equal(e, g), "replayed output differs from eager"
- for name in ("backing", "old_x", "old_B", "old_dt", "pnat", "buf"):
- assert torch.equal(getattr(eager, name), getattr(graph, name)), f"{name} differs"
- accepted = eager.draw_acceptance()
- eager.accept(accepted)
- graph.accept(accepted)
-
-
-# ── guards: accepted and silently wrong ──────────────────────────────────
-
-
-def test_a_short_state_indices_reads_past_its_end_and_is_refused() -> None:
- """Unguarded, three indices for four rows: the fourth row reads whatever
- follows the view -- here a slot the call never named -- and writes its
- history."""
- bench = _Bench(dict(path="small", n=4, t=2, layers=1), seed=101, exact=False)
- stranger = next(s for s in range(bench.slots) if s not in bench.live.tolist())
- parent = torch.cat(
- [bench.state_indices[:3], torch.tensor([stranger], dtype=torch.int32, device=DEV)]
- )
- kw = dict(bench.kwargs(0, bench.inputs()), state_indices=parent[:3])
- before = bench.old_x[0, stranger].clone()
- op(**kw)
- assert not torch.equal(bench.old_x[0, stranger], before), "the hazard this guard covers is gone"
- with pytest.raises(AssertionError, match="state_indices has shape"):
- with validating(op):
- op(**kw)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_update.py
deleted file mode 100644
index 3b2041b63919..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_update.py
+++ /dev/null
@@ -1,250 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the fused_recurrent_gated_delta_rule_update entry.
-
-The axis this entry exists for is more than one token per sequence per step --
-the speculative verify call -- so the cells are that call as the target makes it
-(per-step states into the rollback buffer, the pool untouched) at both pool
-dtypes. The write-back sink is driven only by the cold-call test below.
-
-The neighbouring `chunk_gated_delta_rule` was once wrong on the first indexed
-in-place call of a process (an autotuner trial launch consumes the state), so
-`test_cold_call_in_a_fresh_process_is_correct` measures the first call of a
-process that has made none.
-"""
-
-from __future__ import annotations
-
-import json
-import subprocess
-import sys
-import zlib
-from pathlib import Path
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_recurrent_gated_delta_rule_update import ( # noqa: E501
- HEAD_K_DIM,
- HEAD_V_DIM,
- NUM_K_HEADS,
- NUM_V_HEADS,
- SLOTS,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_recurrent_gated_delta_rule_update import ( # noqa: E501
- fused_recurrent_gated_delta_rule_update as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "fused_recurrent_gated_delta_rule_update requires a CUDA device"
-
-
-def _compare(out: tuple, ref: tuple) -> None:
- """`(returned, initial_state_source, intermediate_states_buffer)` as the
- call left them, against `reference`'s triple, each at torch's default
- dtype-aware band -- nothing is loosened.
-
- The kernel and the reference both accumulate in fp32 and differ only in
- reduction order. Measured over every cell shape: output <= 2.27e-4 at a
- bf16 absmax of 3.9e-2 .. 8.2e-2 (under one bf16 ulp there), written-back
- state <= 2.98e-7 at absmax ~1, per-step states <= 1.8e-7; a bf16
- per-step buffer costs one bf16 ulp (1.94e-3 at absmax 0.89). Eight wrong
- recurrences (q/k heads tiled not interleaved, q/k swapped, beta 1, no
- gating, reversed tokens, every token from the initial state, a
- transposed or absent initial state) land 1.56e-2 .. 8.17e-2 away on the
- output -- 128x the correct result's own deviation at the closest.
- """
- (o, pool, cache), (o_ref, pool_ref, cache_ref) = out, ref
- assert o.dtype == o_ref.dtype and o.shape == o_ref.shape, (
- (o.dtype, o.shape),
- (o_ref.dtype, o_ref.shape),
- )
- torch.testing.assert_close(o, o_ref)
- assert (pool is None) == (pool_ref is None)
- if pool is not None:
- torch.testing.assert_close(pool, pool_ref)
- assert (cache is None) == (cache_ref is None)
- if cache is not None:
- torch.testing.assert_close(cache, cache_ref)
-
-
-DEV = "cuda"
-MIXED_QKV_WIDTH = 2 * NUM_K_HEADS * HEAD_K_DIM + NUM_V_HEADS * HEAD_V_DIM # 10240
-CACHE_SENTINEL = -3.0
-
-
-def _build(spec, seed):
- """`__call__`'s arguments by name for one spec, plus the live slots.
-
- q/k/v are views over one packed `[T, 10240]` row block (or contiguous
- copies, as the target passes them); `g` is built the way the gating step
- builds it -- negative, so `exp(g) < 1` -- and `beta` is a sigmoid.
- """
- gen = torch.Generator(device=DEV).manual_seed(seed)
- mode, n, tokens = spec["mode"], spec["n"], spec["p"]
- total = n * tokens
-
- def randn(*shape, dtype=torch.float32):
- return torch.randn(*shape, generator=gen, device=DEV, dtype=dtype)
-
- mixed = randn(total, MIXED_QKV_WIDTH, dtype=torch.bfloat16)
- kd = NUM_K_HEADS * HEAD_K_DIM
- q = mixed[:, :kd].view(n, tokens, NUM_K_HEADS, HEAD_K_DIM)
- k = mixed[:, kd : 2 * kd].view(n, tokens, NUM_K_HEADS, HEAD_K_DIM)
- v = mixed[:, 2 * kd :].view(n, tokens, NUM_V_HEADS, HEAD_V_DIM)
- if spec.get("qkv") == "contiguous":
- q, k, v = q.contiguous(), k.contiguous(), v.contiguous()
-
- a, a_log, dt_bias = randn(total, NUM_V_HEADS), randn(NUM_V_HEADS), randn(NUM_V_HEADS)
- g = -torch.exp(a_log) * torch.nn.functional.softplus(a + dt_bias)
- g = g.view(n, tokens, NUM_V_HEADS)
- b = randn(total, NUM_V_HEADS, dtype=torch.bfloat16)
- beta = b.sigmoid() if spec.get("beta", "fp32") == "bf16" else b.float().sigmoid()
- beta = beta.view(n, tokens, NUM_V_HEADS)
-
- pool_slots = spec.get("pool_slots", SLOTS)
- pool = randn(pool_slots, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM) * 0.05
- pool = pool.to(spec.get("pool_dtype", torch.float32))
-
- slots = spec.get("slots") or tuple(range(n))
- indices = torch.tensor(slots, dtype=torch.int32, device=DEV)
-
- buffer = cache_steps = None
- if mode == "verify":
- rows = spec.get("cache_rows", pool_slots)
- cache_steps = tokens
- buffer = torch.full(
- (rows, cache_steps, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM),
- CACHE_SENTINEL,
- dtype=spec.get("cache_dtype", torch.float32),
- device=DEV,
- )
-
- kwargs = dict(
- q=q,
- k=k,
- v=v,
- g=g,
- beta=beta,
- initial_state_source=pool,
- initial_state_indices=indices,
- scale=None,
- use_qk_l2norm_in_kernel=True,
- cu_seqlens=None,
- disable_state_update=mode == "verify",
- disable_output_calculation=False,
- intermediate_states_buffer=buffer,
- cache_steps=cache_steps,
- output=None,
- )
- return kwargs, list(slots)
-
-
-def _seed(why):
- return zlib.crc32(why.encode())
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- kw, live = _build(cell.spec, seed=_seed(cell.why))
- pool, buffer = kw["initial_state_source"], kw["intermediate_states_buffer"]
- pool_before, buffer_before = pool.clone(), buffer.clone()
-
- ref = op.reference(**kw)
- got = op(**kw)
- _compare((got, pool, buffer), ref)
-
- # The in-place locks, bitwise, independent of the reference.
- assert torch.equal(pool, pool_before), "the pool must not move on a verify call"
- for s in live:
- assert not (buffer[s] == CACHE_SENTINEL).any(), f"row {s} not fully cached"
- buffer_before[live] = buffer[live]
- assert torch.equal(buffer, buffer_before), "a buffer row outside the indices was written"
-
-
-# ── the mechanism the neighbouring entry failed on ────────────────────────
-
-
-_COLD_SPEC = dict(mode="write_back", n=2, p=3, slots=(5, 2))
-
-
-def cold_call_measurement() -> dict:
- """The first indexed in-place call of the process, then the second.
-
- Public because the child process below calls it by name.
- """
- kw, _ = _build(_COLD_SPEC, seed=7)
- pool = kw["initial_state_source"]
- pool_before = pool.clone()
- ref = op.reference(**kw)
- cold_out = op(**kw).clone()
- cold_pool = pool.clone()
- try:
- _compare((cold_out, cold_pool, None), ref)
- cold_ok = True
- except AssertionError:
- cold_ok = False
- pool.copy_(pool_before)
- warm_out = op(**kw)
- return {
- "cold_compare_ok": cold_ok,
- "cold_out_err": (cold_out.float() - ref[0].float()).abs().max().item(),
- "cold_pool_err": (cold_pool - ref[1]).abs().max().item(),
- "out_bitwise_equal": bool(torch.equal(cold_out, warm_out)),
- "pool_bitwise_equal": bool(torch.equal(cold_pool, pool)),
- }
-
-
-_COLD_CHILD = """
-import importlib.util, json, sys
-sys.path.insert(0, {root!r})
-spec = importlib.util.spec_from_file_location("_cold_entry_test", {path!r})
-mod = importlib.util.module_from_spec(spec)
-sys.modules["_cold_entry_test"] = mod
-spec.loader.exec_module(mod)
-print("COLD_RESULT " + json.dumps(mod.cold_call_measurement()))
-"""
-
-
-def test_cold_call_in_a_fresh_process_is_correct() -> None:
- """Measured in a child that has launched nothing: a warm-up here would
- hide exactly the defect the sibling prefill core has. Cold must pass the
- entry's own gate *and* be bitwise the warm call."""
- root = str(Path(__file__).resolve().parents[4]) # where `_validating` lives
- proc = subprocess.run(
- [sys.executable, "-c", _COLD_CHILD.format(root=root, path=__file__)],
- capture_output=True,
- text=True,
- timeout=480,
- check=False,
- )
- assert proc.returncode == 0, (
- f"cold-call child exited {proc.returncode}\n{proc.stdout[-4000:]}\n{proc.stderr[-4000:]}"
- )
- line = [ln for ln in proc.stdout.splitlines() if ln.startswith("COLD_RESULT ")]
- assert line, f"child produced no result\n{proc.stdout[-4000:]}"
- result = json.loads(line[0][len("COLD_RESULT ") :])
- print(f"\ncold call: {result}")
- assert result["cold_compare_ok"], result
- assert result["out_bitwise_equal"] and result["pool_bitwise_equal"], result
-
-
-# ── what the guard refuses ─────────────────────────────────────────────
-
-
-def test_cache_steps_off_the_buffer_step_axis_misplaces_rows_and_is_refused() -> None:
- """Unguarded, a 4-step buffer driven at cache_steps 2: row 1's states land in
- row 0's steps 2..3 and row 1 is never written -- the rows the promotion
- reads are not the ones the kernel filled."""
- kw = _build(dict(mode="verify", n=2, p=2), seed=75)[0]
- deep = torch.full((SLOTS, 4, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM), CACHE_SENTINEL, device=DEV)
- op(**dict(kw, intermediate_states_buffer=deep, cache_steps=2))
- assert (deep[1] == CACHE_SENTINEL).all() and not (deep[0, 2:] == CACHE_SENTINEL).any(), (
- "the hazard this guard covers is gone"
- )
- with pytest.raises(AssertionError, match="must equal intermediate_states_buffer's step axis"):
- with validating(op):
- op(**dict(kw, intermediate_states_buffer=deep.clone(), cache_steps=2))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py
deleted file mode 100644
index b1039ed71fd6..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py
+++ /dev/null
@@ -1,231 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the fused_sigmoid_gating_delta_rule_update entry.
-
-Everything here certifies the vendored Triton branch: the target's decode call -- the
-rows form (no `cu_seqlens`) over a bf16 pool -- and the packed call over an fp32 pool that
-the GDN block test uses as its reference. `test_rows_form_is_batch_invariant` shows the
-property the bf16 call takes the rows form for.
-"""
-
-import zlib
-from typing import Optional
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_sigmoid_gating_delta_rule_update import ( # noqa: E501
- ENGINE_SLOTS,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_sigmoid_gating_delta_rule_update import ( # noqa: E501
- fused_sigmoid_gating_delta_rule_update as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "fused_sigmoid_gating_delta_rule_update requires a CUDA device"
-
-
-def _compare(
- out: tuple[torch.Tensor, Optional[torch.Tensor]],
- ref: tuple[torch.Tensor, Optional[torch.Tensor]],
-) -> None:
- """`(o, pool)`: the returned output with the pool as the call left it.
-
- Both at torch's default dtype-aware band: the kernel and the reference
- accumulate the same fp32 sequence and round once (bf16 `o`: measured
- 3.05e-05 max abs at N=32, 461 slots; fp32 pool 2.4e-07). Over a bf16 pool the
- state is rounded once on store, as the reference rounds it, so a stored value
- can sit one bf16 step (2^-8 relative) from the reference's, inside the band. Five wrong
- variants (a/b swapped, norm dropped, pool `(V, K)` transposed, initial
- state ignored, q-head order reversed) land 1.7e-2 to 4.5 away, three
- orders of magnitude outside. The pool is compared whole, so the
- unnamed slots are gated too (the reference leaves them bit-identical).
- """
- (o, pool), (ref_o, ref_pool) = out, ref
- assert o.dtype == ref_o.dtype and o.shape == ref_o.shape, (
- (o.dtype, o.shape),
- (ref_o.dtype, ref_o.shape),
- )
- assert torch.isfinite(o).all()
- torch.testing.assert_close(o, ref_o)
- assert (pool is None) == (ref_pool is None)
- if pool is not None:
- assert pool.dtype == ref_pool.dtype and pool.shape == ref_pool.shape
- torch.testing.assert_close(pool, ref_pool)
-
-
-DEV = "cuda"
-NUM_K_HEADS, NUM_V_HEADS, HEAD_K_DIM, HEAD_V_DIM = 16, 48, 128, 128
-MIXED_QKV_WIDTH = 2 * NUM_K_HEADS * HEAD_K_DIM + NUM_V_HEADS * HEAD_V_DIM # 10240
-#: The target's one GDN input-projection output: [q|k|v | z | a | b], 16480 wide.
-IN_PROJ_A = MIXED_QKV_WIDTH + NUM_V_HEADS * HEAD_V_DIM # 16384
-IN_PROJ_WIDTH = IN_PROJ_A + 2 * NUM_V_HEADS
-SOFTPLUS_BETA, SOFTPLUS_THRESHOLD = 1.0, 20.0
-
-
-def _rows_form(mixed, a, b):
- """The target's bf16 decode operands: q/k/v `[N, 1, heads, 128]` views of the packed
- `[N, 10240]` block, reshaped whole so the size-1 token axis carries the row stride, and
- a/b `[N, 1, 48]`."""
- n = mixed.shape[0]
- qkv = mixed.view(n, 1, MIXED_QKV_WIDTH // HEAD_K_DIM, HEAD_K_DIM)
- q, k, v = torch.split(qkv, [NUM_K_HEADS, NUM_K_HEADS, NUM_V_HEADS], 2)
- return q, k, v, a.reshape(n, 1, NUM_V_HEADS), b.reshape(n, 1, NUM_V_HEADS)
-
-
-def _build(spec, seed):
- """`__call__`'s 15 positional arguments, plus the packed block the rows form
- views. q/k/v are views over one packed `[T, 10240]` conv output, exactly as
- the decode call reshapes its column splits."""
- g = torch.Generator(device=DEV).manual_seed(seed)
- seq_lens = spec["seq_lens"]
- total = sum(seq_lens)
- if spec["ab"] == "in_proj":
- proj = torch.randn(total, IN_PROJ_WIDTH, generator=g, device=DEV, dtype=torch.bfloat16)
- mixed = proj[:, :MIXED_QKV_WIDTH]
- else:
- mixed = torch.randn(total, MIXED_QKV_WIDTH, generator=g, device=DEV, dtype=torch.bfloat16)
- kd = NUM_K_HEADS * HEAD_K_DIM
- q = mixed[:, :kd].view(1, total, NUM_K_HEADS, HEAD_K_DIM)
- k = mixed[:, kd : 2 * kd].view(1, total, NUM_K_HEADS, HEAD_K_DIM)
- v = mixed[:, 2 * kd :].view(1, total, NUM_V_HEADS, HEAD_V_DIM)
- if spec["ab"] == "in_proj":
- a = proj[:, IN_PROJ_A : IN_PROJ_A + NUM_V_HEADS]
- b = proj[:, IN_PROJ_A + NUM_V_HEADS :]
- assert mixed.stride(0) == a.stride(0) == b.stride(0) == IN_PROJ_WIDTH
- else:
- a = torch.randn(total, NUM_V_HEADS, generator=g, device=DEV, dtype=torch.bfloat16)
- b = torch.randn(total, NUM_V_HEADS, generator=g, device=DEV, dtype=torch.bfloat16)
- A_log = torch.randn(NUM_V_HEADS, generator=g, device=DEV)
- dt_bias = torch.randn(NUM_V_HEADS, generator=g, device=DEV)
-
- state_shape = (spec["slots"], NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM)
- pool = torch.randn(state_shape, generator=g, device=DEV) * 0.05
- pool = pool.to(spec.get("pool_dtype", torch.float32))
- indices = spec.get("indices") or tuple(range(len(seq_lens)))
- indices = torch.tensor(indices, dtype=torch.int32, device=DEV)
- cu_seqlens = torch.tensor(
- [0] + torch.tensor(seq_lens).cumsum(0).tolist(), dtype=torch.int64, device=DEV
- )
- if spec.get("layout", "varlen") == "rows":
- q, k, v, a, b = _rows_form(mixed, a, b)
- cu_seqlens = None
- args = [
- A_log,
- a,
- dt_bias,
- SOFTPLUS_BETA,
- SOFTPLUS_THRESHOLD,
- q,
- k,
- v,
- b,
- pool,
- indices,
- None,
- True,
- cu_seqlens,
- None,
- ]
- return args, {"mixed": mixed}
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- spec = cell.spec
- args, _ = _build(spec, seed=zlib.crc32(cell.why.encode()))
- pool = args[9]
- pool_before = pool.clone()
- ref = op.reference(*args)
-
- out = op(*args)
-
- _compare((out, pool), ref)
- # The whole-pool compare is the lock on unnamed slots; make it bitwise.
- named = {s for s in args[10].tolist() if s >= 0}
- unnamed = [s for s in range(pool.shape[0]) if s not in named]
- assert torch.equal(pool[unnamed], pool_before[unnamed]), "unnamed slots written"
- if spec["slots"] == ENGINE_SLOTS:
- for slot in named:
- assert not torch.equal(pool[slot], pool_before[slot]), slot
-
-
-#: Batch sizes either side of every FlashInfer dispatch boundary at 48 value heads (its
-#: kernel and tile change at 2|3, 6|7, 10|11 and 21|22) and the engine's CUDA-graph batch
-#: sizes up to 32.
-_BATCHES = (1, 2, 3, 10, 11, 16, 21, 22, 32)
-_FIXED_SLOT = 301
-_FIXED_ROW_SLOTS = 512
-
-
-def test_rows_form_is_batch_invariant() -> None:
- """The property the target's bf16 decode call relies on: a row's output and state are
- bitwise independent of the batch it rides in, at batch sizes either side of every
- FlashInfer dispatch boundary and at both ends of the batch. One fixed row (activations
- and slot 301 of a bf16 pool) sits at position 0 and N-1 of random N-row batches."""
- gen = torch.Generator(device=DEV).manual_seed(3)
- cpu = torch.Generator().manual_seed(3)
- pool0 = (
- torch.randn(
- _FIXED_ROW_SLOTS, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM, generator=gen, device=DEV
- )
- * 0.05
- ).to(torch.bfloat16)
- row = torch.randn(MIXED_QKV_WIDTH, generator=gen, device=DEV, dtype=torch.bfloat16)
- row_a, row_b = (
- torch.randn(NUM_V_HEADS, generator=gen, device=DEV, dtype=torch.bfloat16) for _ in range(2)
- )
- A_log = torch.randn(NUM_V_HEADS, generator=gen, device=DEV)
- dt_bias = torch.randn(NUM_V_HEADS, generator=gen, device=DEV)
- first = None
- for n in _BATCHES:
- for pos in sorted({0, n - 1}):
- g = torch.Generator(device=DEV).manual_seed(1000 * n + pos)
- mixed = torch.randn(n, MIXED_QKV_WIDTH, generator=g, device=DEV, dtype=torch.bfloat16)
- a = torch.randn(n, NUM_V_HEADS, generator=g, device=DEV, dtype=torch.bfloat16)
- b = torch.randn(n, NUM_V_HEADS, generator=g, device=DEV, dtype=torch.bfloat16)
- mixed[pos], a[pos], b[pos] = row, row_a, row_b
- others = [s for s in torch.randperm(_FIXED_ROW_SLOTS, generator=cpu).tolist()]
- others = [s for s in others if s != _FIXED_SLOT][:n]
- others[pos] = _FIXED_SLOT
- slots = torch.tensor(others, dtype=torch.int32, device=DEV)
- pool = pool0.clone()
- q, k, v, a_in, b_in = _rows_form(mixed, a, b)
- out = op(A_log, a_in, dt_bias, 1.0, 20.0, q, k, v, b_in, pool, slots, None, True, None)
- got = (out.reshape(n, NUM_V_HEADS, HEAD_V_DIM)[pos].clone(), pool[_FIXED_SLOT].clone())
- first = first or got
- assert torch.equal(got[0], first[0]) and torch.equal(got[1], first[1]), (
- f"the fixed row differs at batch {n}, position {pos}"
- )
-
-
-def test_guard_refuses_rows_views_whose_token_stride_is_not_the_row_stride() -> None:
- """Unguarded, q/k/v reshaped to `[N, 1, heads, 128]` from column slices of the packed
- block carry `stride(1) = heads * 128`: without `cu_seqlens` the kernel reads row n at
- `n * stride(1)`, from other rows' columns."""
- n, kd = 8, NUM_K_HEADS * HEAD_K_DIM
- spec = dict(
- seq_lens=(1,) * n, slots=64, ab="separate", layout="rows", pool_dtype=torch.bfloat16
- )
- args, extras = _build(spec, seed=600)
- mixed = extras["mixed"]
- pool_before = args[9].clone()
- ref = op.reference(*args)
- bad = list(args)
- bad[5] = mixed[:, :kd].reshape(n, 1, NUM_K_HEADS, HEAD_K_DIM)
- bad[6] = mixed[:, kd : 2 * kd].reshape(n, 1, NUM_K_HEADS, HEAD_K_DIM)
- bad[7] = mixed[:, 2 * kd :].reshape(n, 1, NUM_V_HEADS, HEAD_V_DIM)
- assert all(bad[i].stride(1) != bad[i].stride(0) for i in (5, 6, 7))
- for i in (5, 6, 7):
- assert torch.equal(bad[i], args[i]), "the same values, other strides"
- got = op(*bad)
- with pytest.raises(AssertionError):
- _compare((got, bad[9]), ref)
- guarded = list(bad)
- guarded[9] = pool_before
- with pytest.raises(AssertionError, match=r"stride\(0\) must equal"):
- with validating(op):
- op(*guarded)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py
deleted file mode 100644
index 6447b4cf57fd..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py
+++ /dev/null
@@ -1,416 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test: the target's whole Gated-DeltaNet block against the all-Triton path.
-
-Drives `DecodeTarget._gdn_block` and `PrefillTarget._gdn_block` at the
-checkpoint's geometry (hidden 2560, 16 k-heads, 48 v-heads, head dim 128,
-fp32 or bf16 recurrent pool of 461 slots) on a decode and a mixed (context + one-token
-generation rows) batch, and compares the block output
-and both state pools with a reference that composes the Triton recurrences
-(`fused_sigmoid_gating_delta_rule_update`'s rows form for decode, bitwise its
-packed call at either pool dtype; `chunk_gated_delta_rule` over every row of a
-prefill step) around the same projections, convolutions and epilogue.
-
-The reference's context rows run the unfused `extract_transpose_prefill_slice`
--> `causal_conv1d_fwd` -> `fused_gdn_post_conv` chain the target's
-`gdn_prefill_conv_post` replaces.
-
-Bands: the conv pools are bitwise (both sides publish the raw projection). The
-recurrent pool is held to the looser of the two recurrence entries' state
-bands (`flashinfer_chunk_gated_delta_rule`: rtol 1.6e-2, atol 2e-2), and the
-block output, after the gated norm and the output projection, to a relative
-L2 distance of 1e-2 and a max-abs distance of 3e-2 of its own peak.
-"""
-
-import types
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501
- modeling as M,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_fwd import ( # noqa: E501
- causal_conv1d_fwd,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.chunk_gated_delta_rule import ( # noqa: E501
- chunk_gated_delta_rule as triton_chunk,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.extract_transpose_prefill_slice import ( # noqa: E501
- extract_transpose_prefill_slice,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_gdn_post_conv import ( # noqa: E501
- fused_gdn_post_conv,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_sigmoid_gating_delta_rule_update import ( # noqa: E501
- fused_sigmoid_gating_delta_rule_update as triton_decode,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.reset_recurrent_state_rows import ( # noqa: E501
- reset_recurrent_state_rows,
-)
-
-assert torch.cuda.is_available(), "the Gated-DeltaNet block requires a CUDA device"
-
-DEV = "cuda"
-BF16 = torch.bfloat16
-HIDDEN, H, HV, K, V, WIDTH = 2560, 16, 48, 128, 128, 4
-KD, VD = H * K, HV * V
-CONV = 2 * KD + VD
-SLOTS = 461
-MAX_BATCH = 256
-RTOL_STATE, ATOL_STATE = 1.6e-2, 2e-2
-OUT_REL_L2, OUT_MAX_FRAC = 1e-2, 3e-2
-
-
-def _slots(n: int) -> list[int]:
- """`n` distinct slots of the pool, both ends first."""
- rest = [(i * 37 + 4) % (SLOTS - 2) + 1 for i in range(max(n - 2, 0))]
- return ([SLOTS - 1, 0] + rest)[:n]
-
-
-def _core(seed: int, state: torch.dtype = torch.float32) -> types.SimpleNamespace:
- """The attributes `_gdn_block` reads off the core, with layer 0's pools (the recurrent
- one in `state`)."""
- g = torch.Generator(device=DEV).manual_seed(seed)
- conv = torch.randn(SLOTS, CONV, WIDTH - 1, generator=g, device=DEV, dtype=BF16)
- ssm = (torch.randn(SLOTS, HV, V, K, generator=g, device=DEV) * 0.05).to(state)
- return types.SimpleNamespace(
- conv_dim=CONV,
- key_dim=KD,
- value_dim=VD,
- num_k_heads=H,
- head_k_dim=K,
- num_v_heads=HV,
- head_v_dim=V,
- eps=1e-6,
- conv_kernel=WIDTH,
- in_proj_split=(CONV, VD, HV, HV),
- _bound=types.SimpleNamespace(conv={0: conv}, ssm={0: ssm}),
- )
-
-
-def _clone_core(core: types.SimpleNamespace) -> types.SimpleNamespace:
- out = types.SimpleNamespace(**vars(core))
- out._bound = types.SimpleNamespace(
- conv={0: core._bound.conv[0].clone()}, ssm={0: core._bound.ssm[0].clone()}
- )
- return out
-
-
-def _layer(seed: int) -> types.SimpleNamespace:
- """Layer weights as the target views them: HF `[out, in]` storage read
- through `.t()`, fp32 gating vectors over the checkpoint's ranges."""
- g = torch.Generator(device=DEV).manual_seed(seed)
-
- def lin(n_out, n_in):
- return (torch.randn(n_out, n_in, generator=g, device=DEV) * n_in**-0.5).to(BF16).t()
-
- # The target's one [q|k|v | z | a | b] input projection, read through `.t()`.
- # The reference blocks run the same GEMM, so the conv pools stay bitwise.
- return types.SimpleNamespace(
- in_proj=lin(CONV + VD + 2 * HV, HIDDEN),
- out=lin(HIDDEN, VD),
- conv_w=(torch.randn(CONV, WIDTH, generator=g, device=DEV) * 0.5).to(BF16),
- a_log=torch.empty(HV, device=DEV).uniform_(-5.5, 3.5, generator=g),
- dt_bias=torch.empty(HV, device=DEV).uniform_(-8.9, 19.2, generator=g),
- gnorm=(1.0 + 0.1 * torch.randn(VD // HV, generator=g, device=DEV)).to(BF16),
- # No FP8 copies: the projections stay on the bf16 entries (dense FP8 mode off).
- in_proj8=None,
- out8=None,
- in_proj_g=None,
- out_g=None,
- in_proj_m=None,
- out_m=None,
- in_proj_nv=None,
- )
-
-
-def _x(tokens: int, seed: int) -> torch.Tensor:
- g = torch.Generator(device=DEV).manual_seed(seed)
- return torch.randn(tokens, HIDDEN, generator=g, device=DEV, dtype=BF16)
-
-
-def _target(cls, core, L=None, step=None):
- """A target over `core` alone; a prefill target's `_gdn_core` break reads layer 0's
- weights and the step's bookkeeping off it."""
- target = object.__new__(cls)
- target.core = core
- target._layers = [L]
- target._step = step
- return target
-
-
-def _prefill_block(core, x, L, st):
- """`PrefillTarget._gdn_block` over the step `st`, the break's output buffer allocated
- as the forward allocates it."""
- tokens = x.shape[0]
- step = types.SimpleNamespace(gdn=st, num_tokens=tokens)
- out = torch.empty([1, tokens, HV, V], dtype=BF16, device=DEV)
- return _target(M.PrefillTarget, core, L, step)._gdn_block(x, 0, L, out)
-
-
-def _decode_state(m: int):
- """`DecodeTarget.forward`'s `gdn_state` for `m` generation sequences: the
- slot vector is the head of an engine-sized buffer."""
- buf = torch.zeros(MAX_BATCH, dtype=torch.int32, device=DEV)
- buf[:m] = torch.tensor(_slots(m), dtype=torch.int32, device=DEV)
- return (buf[:m], False, m, 1, None, None)
-
-
-def _prefill_state(ctx_lens: tuple, n_dec: int, has_init: tuple):
- """`PrefillTarget.forward`'s `gdn_state` for a context prefix plus `n_dec`
- one-token generation rows, built the way the forward builds it from the
- engine's `Mamba2Metadata` buffers."""
- n_ctx = len(ctx_lens)
- batch = n_ctx + n_dec
- n_p = sum(ctx_lens)
- si_buf = torch.zeros(MAX_BATCH, dtype=torch.int32, device=DEV)
- si_buf[:batch] = torch.tensor(_slots(batch), dtype=torch.int32, device=DEV)
- hi_buf = torch.zeros(MAX_BATCH, dtype=torch.bool, device=DEV)
- hi_buf[:n_ctx] = torch.tensor(has_init, dtype=torch.bool, device=DEV)
- si_all, hi_all = si_buf[:batch], hi_buf[:batch]
- si_p, si_d = torch.split(si_all, [n_ctx, n_dec], 0)
- si_d = torch.clone(si_d)
- hi_p = hi_all[:n_ctx]
- bounds = [0]
- for n in ctx_lens:
- bounds.append(bounds[-1] + n)
- bounds += [n_p + j for j in range(1, n_dec + 1)]
- cu_buf = torch.zeros(MAX_BATCH + 1, dtype=torch.int32, device=DEV)
- cu_buf[: batch + 1] = torch.tensor(bounds, dtype=torch.int32, device=DEV)
- cu_long_buf = cu_buf.to(torch.int64)
- qsl_all = cu_buf[: batch + 1]
- qsl_p = qsl_all[: n_ctx + 1]
- qsl_long = cu_long_buf[: batch + 1]
- qsl_long_p = qsl_long[: n_ctx + 1]
- return (
- n_p,
- n_dec,
- si_p,
- si_all,
- si_d,
- hi_p,
- hi_all,
- qsl_p,
- qsl_all,
- qsl_long,
- False,
- n_dec,
- 1,
- None,
- qsl_long_p,
- None,
- )
-
-
-def _post_conv(prefill, decode, a, b, L):
- return fused_gdn_post_conv(
- prefill=prefill,
- decode=decode,
- a=a,
- b=b,
- A_log=L.a_log,
- dt_bias=L.dt_bias,
- num_k_heads=H,
- head_k_dim=K,
- num_v_heads=HV,
- head_v_dim=V,
- l2_norm_eps=M._L2_NORM_EPS,
- softplus_threshold=M._SOFTPLUS_THRESHOLD,
- )
-
-
-def _triton_decode_block(core, x, L, si_all):
- """The decode block with the Triton single-token recurrence in its rows form (q/k/v
- `[m, 1, heads, 128]` from a contiguous copy, `cu_seqlens` None), which stays on the
- Triton kernel at either pool dtype."""
- m = x.shape[0]
- mixed, z, a, b = torch.split(M.cublas_mm(x, L.in_proj), [CONV, VD, HV, HV], 1)
- M.causal_conv1d_update(
- x=mixed,
- conv_state=core._bound.conv[0],
- weight=L.conv_w,
- bias=None,
- silu_activation=True,
- cache_seqlens=None,
- conv_state_indices=si_all,
- pad_slot_id=M._PAD_SLOT_ID,
- )
- qv, kv, vv = torch.split(torch.reshape(mixed.contiguous(), [m, 1, -1, K]), [H, H, HV], 2)
- out = triton_decode(
- A_log=L.a_log,
- a=torch.reshape(a.contiguous(), [m, 1, HV]),
- dt_bias=L.dt_bias,
- softplus_beta=M._SOFTPLUS_BETA,
- softplus_threshold=M._SOFTPLUS_THRESHOLD,
- q=qv,
- k=kv,
- v=vv,
- b=torch.reshape(b.contiguous(), [m, 1, HV]),
- initial_state_source=core._bound.ssm[0],
- initial_state_indices=si_all,
- scale=None,
- use_qk_l2norm_in_kernel=True,
- cu_seqlens=None,
- output=None,
- )
- return M.DecodeTarget._gdn_epilogue(_target(M.DecodeTarget, core), out, z, L, m)
-
-
-def _triton_prefill_block(core, x, L, st):
- """The prefill block with every row -- context and generation -- in one
- Triton chunk call."""
- n_p, n_d, si_p, si_all, si_d, hi_p, hi_all, qsl_p, qsl_all, qsl_long = st[:10]
- T = x.shape[0]
- conv, ssm = core._bound.conv[0], core._bound.ssm[0]
- mixed, z, a, b = torch.split(M.cublas_mm(x, L.in_proj), [CONV, VD, HV, HV], 1)
- reset_recurrent_state_rows(ssm, si_p, hi_p, conv)
- conv_args = dict(
- weight=L.conv_w,
- bias=None,
- conv_states=conv,
- silu_activation=True,
- pad_slot_id=M._PAD_SLOT_ID,
- out=None,
- )
- if n_d > 0:
- mq_p, mq_d = torch.split(mixed, [n_p, n_d], 0)
- mq_p_t = extract_transpose_prefill_slice(mq_p, n_p, 0, CONV)
- causal_conv1d_fwd(
- x=mq_p_t, query_start_loc=qsl_p, cache_indices=si_p, has_initial_state=hi_p, **conv_args
- )
- M.causal_conv1d_update(
- x=mq_d,
- conv_state=conv,
- weight=L.conv_w,
- bias=None,
- silu_activation=True,
- cache_seqlens=None,
- conv_state_indices=si_d,
- pad_slot_id=M._PAD_SLOT_ID,
- )
- q, k, v, g, beta = _post_conv(mq_p_t, mq_d, a, b, L)
- else:
- mq_t = extract_transpose_prefill_slice(mixed, T, 0, CONV)
- causal_conv1d_fwd(
- x=mq_t,
- query_start_loc=qsl_all,
- cache_indices=si_all,
- has_initial_state=hi_all,
- **conv_args,
- )
- q, k, v, g, beta = _post_conv(mq_t, None, a, b, L)
- out, _ = triton_chunk(
- q,
- k,
- v,
- g,
- beta,
- scale=None,
- initial_state=ssm,
- initial_state_indices=si_all,
- inplace_indexed_state_update=True,
- output_final_state=False,
- cu_seqlens=qsl_long,
- use_qk_l2norm_in_kernel=False,
- output=None,
- )
- return M.PrefillTarget._gdn_epilogue(_target(M.PrefillTarget, core), out, z, L, T)
-
-
-def _check(y, y_ref, core, core_ref, what):
- """Output and both pools against the reference block's, at the module bands."""
- assert y.shape == y_ref.shape and y.dtype == y_ref.dtype
- assert torch.isfinite(y).all()
- d = y.float() - y_ref.float()
- rel = (d.norm() / y_ref.float().norm()).item()
- peak = (d.abs().max() / y_ref.float().abs().max()).item()
- s_d = core._bound.ssm[0] - core_ref._bound.ssm[0]
- s_rel = (s_d.norm() / core_ref._bound.ssm[0].norm()).item()
- print(
- f"{what}: out rel-L2 {rel:.3e}, out max-abs/peak {peak:.3e}, "
- f"state rel-L2 {s_rel:.3e}, state max-abs {s_d.abs().max().item():.3e}"
- )
- assert rel <= OUT_REL_L2 and peak <= OUT_MAX_FRAC, (rel, peak)
- assert torch.equal(core._bound.conv[0], core_ref._bound.conv[0]), "conv pools differ"
- torch.testing.assert_close(
- core._bound.ssm[0], core_ref._bound.ssm[0], rtol=RTOL_STATE, atol=ATOL_STATE
- )
-
-
-@pytest.fixture(scope="module")
-def warm_triton_chunk():
- """The Triton chunk op's first indexed in-place call per autotune key
- returns a wrong answer (its entry's note); spend it on scratch."""
- core = _core(0)
- L = _layer(0)
- st = _prefill_state((64,), 0, (False,))
- _triton_prefill_block(core, _x(64, 0), L, st)
- torch.cuda.synchronize()
-
-
-#: The recurrent pool dtypes the target serves (`mamba_ssm_cache_dtype`).
-STATES = {"fp32": torch.float32, "bf16": torch.bfloat16}
-
-
-@pytest.mark.parametrize("state", STATES)
-def test_decode_block(state) -> None:
- """A con-128 decode step; the graph-replay test below runs 32 rows."""
- m = 128
- core, L, x = _core(100 + m, STATES[state]), _layer(200 + m), _x(m, 300 + m)
- core_ref = _clone_core(core)
- st = _decode_state(m)
- y = _target(M.DecodeTarget, core)._gdn_block(x, 0, L, st)
- y_ref = _triton_decode_block(core_ref, x, L, st[0])
- _check(y, y_ref, core, core_ref, f"decode m={m}, {state} pool")
- named = set(_slots(m))
- unnamed = [s for s in range(SLOTS) if s not in named]
- original = _core(100 + m, STATES[state])
- assert torch.equal(core._bound.ssm[0][unnamed], original._bound.ssm[0][unnamed]), (
- "unnamed slots written"
- )
-
-
-@pytest.mark.parametrize("state", STATES)
-def test_prefill_block(state, warm_triton_chunk) -> None: # noqa: F811 — fixture by name
- """A mixed step: three variable-length contexts (two carrying state) plus five
- one-token generation rows."""
- ctx_lens, n_dec, has_init = (200, 37, 91), 5, (True, False, True)
- tokens = sum(ctx_lens) + n_dec
- core, L, x = _core(400 + tokens, STATES[state]), _layer(500 + tokens), _x(tokens, 600 + tokens)
- core_ref = _clone_core(core)
- st = _prefill_state(ctx_lens, n_dec, has_init)
- y = _prefill_block(core, x, L, st)
- y_ref = _triton_prefill_block(core_ref, x, L, st)
- _check(y, y_ref, core, core_ref, f"prefill {ctx_lens}+{n_dec}x1, {state} pool")
-
-
-@pytest.mark.parametrize("state", STATES)
-def test_decode_block_replays_under_cuda_graph(state) -> None:
- """The decode block captured into a CUDA graph after one eager call (the
- eager call compiles the decode kernel and sizes its per-batch scratch, as
- the engine's warmup does) replays to exactly the eager result."""
- m = 32
- core0, L, x = _core(700, STATES[state]), _layer(701), _x(m, 702)
- st = _decode_state(m)
- eager_core = _clone_core(core0)
- y_eager = _target(M.DecodeTarget, eager_core)._gdn_block(x, 0, L, st).clone()
- graph_core = _clone_core(core0)
- target = _target(M.DecodeTarget, graph_core)
- stream = torch.cuda.Stream()
- stream.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(stream):
- target._gdn_block(x, 0, L, st)
- torch.cuda.current_stream().wait_stream(stream)
- torch.cuda.synchronize()
- graph_core._bound.conv[0].copy_(core0._bound.conv[0])
- graph_core._bound.ssm[0].copy_(core0._bound.ssm[0])
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- y_graph = target._gdn_block(x, 0, L, st)
- assert torch.equal(graph_core._bound.ssm[0], core0._bound.ssm[0]), "capture must not execute"
- graph.replay()
- torch.cuda.synchronize()
- assert torch.equal(y_graph, y_eager)
- assert torch.equal(graph_core._bound.ssm[0], eager_core._bound.ssm[0])
- assert torch.equal(graph_core._bound.conv[0], eager_core._bound.conv[0])
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_cached_replay_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_cached_replay_cute.py
deleted file mode 100644
index 78902d1b7580..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_cached_replay_cute.py
+++ /dev/null
@@ -1,364 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test: the CuTe DSL cached-replay kernel (`gdn_cached_replay_cute`) over consecutive
-verify steps against the Triton entry's fp32 `reference` and the Triton kernel itself, driven by
-that entry's test harness (`_Bench`: work items, all-layer commit, `_advance_replay_state` with
-random partial acceptance), in both commit modes: (a) no in-kernel commit + the manager's
-all-layer commit when partitioned, (b) every wrapping row committed in-kernel. Plus rows with a
-negative slot and CUDA-graph replay (PDL on and off) bitwise against eager."""
-
-import importlib.util
-import os
-
-import pytest
-import torch
-
-import tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_cached_replay_cute as cr # noqa: E501
-
-_spec = importlib.util.spec_from_file_location(
- "trep",
- os.path.join(
- os.path.dirname(os.path.abspath(__file__)),
- "test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py",
- ),
-)
-trep = importlib.util.module_from_spec(_spec)
-_spec.loader.exec_module(trep)
-
-assert torch.cuda.is_available(), "gdn_cached_replay_cute requires a CUDA device"
-
-BF = torch.bfloat16
-ULP = trep.BF16_ULP
-
-
-def blk_err(got, want):
- """max |got - want| / absmax of its block (last two dims), in bf16 ulps; relative RMS."""
- block = got.shape[-2] * got.shape[-1]
- g, w = got.float().reshape(-1, block), want.float().reshape(-1, block)
- keep = ~torch.isnan(w).any(1)
- g, w = g[keep], w[keep]
- if g.numel() == 0:
- return 0.0, 0.0
- scale = w.abs().amax(1, keepdim=True).clamp_min(1e-30)
- worst = ((g - w).abs() / scale).nan_to_num(nan=float("inf")).max().item() / ULP
- rms = ((g - w).pow(2).mean().sqrt() / w.pow(2).mean().sqrt().clamp_min(1e-30)).item() / ULP
- return worst, rms
-
-
-class Stat:
- def __init__(self):
- self.d = {}
-
- def add(self, name, val):
- a, b = val if isinstance(val, tuple) else (val, 0.0)
- o = self.d.get(name, (0.0, 0.0))
- self.d[name] = (max(o[0], a), max(o[1], b))
-
- def fmt(self):
- return " ".join(f"{k}={v[0]:.2f}/{v[1]:.2f}" for k, v in self.d.items())
-
-
-def run_cell(name, spec, mode, config=None, seed=1234):
- """mode 'a': mine without commit (+ the manager's all-layer commit when partitioned);
- mode 'b': mine commits wrapping rows in-kernel (no all-layer commit for mine)."""
- spec = dict(spec, state_dtype=BF)
- bt = trep._Bench(spec, seed, exact=False)
- bm = trep._Bench(spec, seed, exact=False)
- dup = tuple(spec.get("dup", ()))
- n = bt.n
- unique = n - len(dup)
- others = torch.ones(bm.slots, dtype=torch.bool, device="cuda")
- if dup:
- others[int(bm.state_indices[unique])] = False
- vs_ref, vs_tri, tri_ref = Stat(), Stat(), Stat()
- fails = []
- wraps = 0
- for step in range(bt.steps):
- _, wrapped, _ = bm.writes()
- wraps += wrapped.numel()
- bt.prepare()
- bm.prepare()
- wi_all = nw_all = None
- if bm.path == "all_layer":
- # the engine's work items cover every decode row (padding included)
- wi_all = torch.zeros(bm.n, 4, dtype=torch.int32, device="cuda")
- nw_all = torch.zeros(1, dtype=torch.int32, device="cuda")
- trep._build_replay_work_items_triton(
- bm.state_indices, bm.pnat, bm.buf, wi_all, nw_all, bm.t, trep.HIST
- )
- for layer in range(bt.layers):
- at, am = bt.inputs(), bm.inputs()
- kt, km = bt.kwargs(layer, at), bm.kwargs(layer, am)
- rkw = dict(km)
- if mode == "b":
- rkw.update(replay_work_items=None, n_writes=None, use_all_layer_commit=False)
- if dup:
- rkw.update(
- q=km["q"][:unique],
- k=km["k"][:unique],
- v=km["v"][:unique],
- g=km["g"][:unique],
- beta=km["beta"][:unique],
- state_indices=km["state_indices"][:unique],
- packed_qkv=am["packed"][: unique * bm.t],
- )
- ref = trep.op.reference(**rkw)
- tkw = dict(kt)
- if mode == "b":
- tkw.update(replay_work_items=None, n_writes=None, use_all_layer_commit=False)
- rt = None if dup else trep.op.reference(**tkw)
- got_t = trep.op(**tkw)
- if rt is not None:
- tri_ref.add("out", blk_err(got_t, rt[0]))
- tri_ref.add("u", blk_err(bt.state(layer)[1], rt[2])[0])
- tri_ref.add("state", blk_err(bt.state(layer)[0], rt[1])[0])
- kmm = dict(km)
- if wi_all is not None:
- kmm.update(replay_work_items=wi_all, n_writes=nw_all)
- got_m = cr.cached_replay_update(**kmm, in_kernel_commit=(mode == "b"), config=config)
- pool, u, k, G = bm.state(layer)
- gm = (got_m[:unique], pool[others], u[others], k[others], G[others])
- rf = (ref[0], *(x[others] for x in ref[1:]))
- try:
- trep._compare(gm, rf)
- except AssertionError as e: # record, keep going
- fails.append(f"step {step} layer {layer}: {str(e)[:200]}")
- vs_ref.add("out", blk_err(gm[0], rf[0]))
- vs_ref.add("state", blk_err(gm[1], rf[1])[0])
- vs_ref.add("u", blk_err(gm[2], rf[2])[0])
- vs_ref.add("k", blk_err(gm[3], rf[3])[0])
- vs_ref.add("G", (gm[4] - rf[4]).abs().max().item())
- tp, tu, tk, tG = bt.state(layer)
- vs_tri.add("out", blk_err(got_m[:unique], got_t[:unique]))
- vs_tri.add("u", blk_err(u[others], tu[others])[0])
- vs_tri.add("k", blk_err(k[others], tk[others])[0])
- vs_tri.add("G", (G[others] - tG[others]).abs().max().item())
- if mode == "a":
- bt.commit()
- bm.commit()
- for layer in range(bt.layers):
- vs_tri.add("state", blk_err(bm.state(layer)[0][others], bt.state(layer)[0][others])[0])
- acc = bt.draw_acceptance()
- bm.draw_acceptance()
- bt.accept(acc)
- bm.accept(acc)
- torch.cuda.synchronize()
- ok = not fails
- print(
- f"[{'PASS' if ok else 'FAIL'}] {name} mode={mode} cfg={config} wraps={wraps}\n"
- f" vs reference (max ulp of block absmax / rel-RMS ulp; G abs): {vs_ref.fmt()}\n"
- f" vs Triton: {vs_tri.fmt()}\n"
- f" (Triton vs reference: {tri_ref.fmt()})",
- flush=True,
- )
- for f in fails[:4]:
- print(" ", f)
- return ok
-
-
-def run_negative(config=None, n=8, t=4, seed=99):
- """Rows with a negative slot: nothing read or written for them, output rows untouched."""
- spec = dict(path="small", n=n, t=t, steps=1, layers=1, state_dtype=BF)
- bm = trep._Bench(spec, seed, exact=False)
- sidx = bm.state_indices.clone()
- dead = [1, n - 2]
- sidx[dead] = -1
- ok = True
- for mode in ("b", "a"):
- act = bm.inputs()
- kw = bm.kwargs(0, act)
- sentinel = torch.full((n, t, trep.HV, trep.V), 3.0, dtype=BF, device="cuda")
- kw.update(state_indices=sidx, output=sentinel.clone())
- if mode == "a":
- # every row listed (row, slot, pending count, half), dead rows with their negative slot
- sl = sidx.long()
- safe = sl.clamp(min=0)
- wi = (
- torch.stack(
- [torch.arange(n, device="cuda"), sl, bm.pnat[safe].long(), bm.buf[safe].long()],
- 1,
- )
- .to(torch.int32)
- .flip(0)
- .contiguous()
- )
- nw = torch.zeros(1, dtype=torch.int32, device="cuda")
- kw.update(replay_work_items=wi, n_writes=nw, use_all_layer_commit=True)
- ref = trep.op.reference(**kw)
- got = cr.cached_replay_update(**kw, config=config)
- try:
- trep._compare((got, *bm.state(0)), ref)
- assert torch.equal(got[dead], sentinel[dead]), "dead rows written"
- except AssertionError as e:
- ok = False
- print(" negative-slot FAIL", mode, str(e)[:200])
- print(f"[{'PASS' if ok else 'FAIL'}] negative slots cfg={config}", flush=True)
- return ok
-
-
-def run_graph(config=None, n=32, t=4, mode="b", pdl=False):
- """Captured once, replayed on new activations with acceptance in between: bitwise the eager
- calls on a twin state (the op test's CUDA-graph check, bf16 state, both commit modes)."""
- spec = dict(path="all_layer", n=n, t=t, layers=2, slots=40, state_dtype=BF)
- seed = 7000 + 10 * n + t
- eager, graph, warm = (trep._Bench(spec, seed, exact=False) for _ in range(3))
- acts = [warm.inputs() for _ in range(2)]
-
- def step(b):
- b.prepare()
- outs = []
- for layer in range(2):
- kw = b.kwargs(layer, acts[layer])
- kw.update(launch_with_pdl=pdl)
- outs.append(
- cr.cached_replay_update(**kw, in_kernel_commit=(mode == "b"), config=config)
- )
- if mode == "a":
- b.commit()
- return outs
-
- step(warm)
- torch.cuda.synchronize()
- captured = torch.cuda.CUDAGraph()
- s_ = torch.cuda.Stream()
- s_.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(s_):
- step(warm)
- torch.cuda.current_stream().wait_stream(s_)
- torch.cuda.synchronize()
- with torch.cuda.graph(captured):
- outs_graph = step(graph)
- refill = torch.Generator(device="cuda").manual_seed(seed + 1)
- ok = True
- for _ in range(3):
- for act in acts:
- for key in ("packed", "a", "b"):
- act[key].copy_(
- torch.randn(
- act[key].shape, generator=refill, device="cuda", dtype=act[key].dtype
- )
- )
- outs_eager = step(eager)
- captured.replay()
- torch.cuda.synchronize()
- for e, g in zip(outs_eager, outs_graph):
- ok &= torch.equal(e, g)
- for name in ("backing", "old_x", "old_B", "old_dt", "pnat", "buf"):
- ok &= torch.equal(getattr(eager, name), getattr(graph, name))
- accepted = eager.draw_acceptance()
- eager.accept(accepted)
- graph.accept(accepted)
- print(
- f"[{'PASS' if ok else 'FAIL'}] CUDA graph replay bitwise eager mode={mode} pdl={pdl} cfg={config}",
- flush=True,
- )
- return ok
-
-
-CELLS = {
- "small_n4_t8_full": dict(path="small", n=4, t=8, start="fresh", accept="full", steps=8),
- "small_n8_t3_dup": dict(path="small", n=8, t=3, dup=(5, 6, 7), steps=6),
- "all_n32_t8_full": dict(path="all_layer", n=32, t=8, start="fresh", accept="full", steps=8),
- "all_n16_t3_L8_S461": dict(
- path="all_layer", n=16, t=3, layers=8, slots=461, edge_slots=True, steps=3
- ),
- "t4_n1": dict(path="small", n=1, t=4, steps=6),
- "t4_n4": dict(path="small", n=4, t=4, steps=6),
- "t4_n16": dict(path="all_layer", n=16, t=4, steps=5),
- "t4_n64": dict(path="all_layer", n=64, t=4, steps=4, slots=96),
- "t4_n256": dict(path="all_layer", n=256, t=4, steps=3, slots=300, layers=1),
- "t2_n8": dict(path="small", n=8, t=2, steps=6),
- "t5_n8": dict(path="small", n=8, t=5, steps=6),
- "t4_n32_dup": dict(path="all_layer", n=32, t=4, dup=(28, 29, 30, 31), steps=4),
-}
-
-
-_RUNS = [
- (name, mode)
- for name, spec in CELLS.items()
- for mode in ("a", "b")
- if not (mode == "a" and spec["path"] == "small")
-]
-
-
-@pytest.mark.parametrize("name,mode", _RUNS, ids=[f"{n}-{m}" for n, m in _RUNS])
-def test_steps_match_the_reference(name, mode) -> None:
- assert run_cell(name, CELLS[name], mode)
-
-
-@pytest.mark.parametrize(
- "cell",
- cr.gdn_cached_replay_cute.CELLS,
- ids=[c.why[:40] for c in cr.gdn_cached_replay_cute.CELLS],
-)
-def test_certified_cells(cell) -> None:
- # The mode the target uses at the cell's size: in-kernel commit on the small path, the
- # manager's all-layer commit on the partitioned one.
- assert run_cell(cell.why, cell.spec, "b" if cell.spec["path"] == "small" else "a")
-
-
-def test_negative_slots() -> None:
- assert run_negative()
-
-
-@pytest.mark.parametrize("pdl", [False, True])
-@pytest.mark.parametrize("mode", ["a", "b"])
-def test_graph_replay_matches_eager(mode, pdl) -> None:
- assert run_graph(mode=mode, pdl=pdl)
-
-
-#: Launch mappings a batch can take, and the mapping each is checked against: the state tiles
-#: all by bulk copy with the state dots before the first barrier (the small-batch default),
-#: against the first tile in registers and the dots after it.
-_MAPPINGS = (
- (dict(tpc=2, warps=2, pf=0, minb=4, early_acc=True), dict(tpc=2, warps=2, pf=1, minb=4)),
- (dict(tpc=4, warps=4, pf=0, minb=4, early_acc=True), dict(tpc=4, warps=4, pf=1, minb=4)),
- (dict(tpc=8, warps=4, pf=1, minb=4, early_acc=True), dict(tpc=8, warps=4, pf=1, minb=4)),
- (dict(tpc=8, warps=4, pf=0, minb=4), dict(tpc=24, warps=8, pf=1, minb=2)),
-)
-
-
-@pytest.mark.parametrize("n,t,pdl", [(1, 4, True), (5, 3, False), (32, 4, True), (8, 8, True)])
-def test_launch_mappings_agree(n, t, pdl) -> None:
- """Every launch mapping computes the same instructions per state row.
-
- One step bitwise equal across mappings (output, checkpoint with in-kernel commits, u / k /
- G history), with wrapping and appending rows and a negative slot, with and without the
- engine's work items.
- """
- spec = dict(path="small", n=n, t=t, steps=1, layers=1, state_dtype=BF)
- bm = trep._Bench(spec, 4242 + n, exact=False)
- sidx = bm.state_indices.clone()
- if n > 2:
- sidx[n // 2] = -1
- act = bm.inputs()
- snap = [x.clone() for x in bm.state(0)]
- # every row listed (row, slot, pending count, half), reversed, the dead row with its -1
- sl = sidx.long()
- safe = sl.clamp(min=0)
- wi = (
- torch.stack(
- [torch.arange(n, device="cuda"), sl, bm.pnat[safe].long(), bm.buf[safe].long()], 1
- )
- .to(torch.int32)
- .flip(0)
- .contiguous()
- )
- nw = torch.zeros(1, dtype=torch.int32, device="cuda")
- for new_cfg, old_cfg in _MAPPINGS:
- for with_wi in (False, True):
- got = []
- for cfg in (new_cfg, old_cfg):
- for x, y in zip(bm.state(0), snap):
- x.copy_(y)
- kw = bm.kwargs(0, act)
- # a dead row's output is never written: a sentinel, the same for every call
- sentinel = torch.full((n, t, trep.HV, trep.V), 3.0, dtype=BF, device="cuda")
- kw.update(state_indices=sidx, launch_with_pdl=pdl, output=sentinel)
- if with_wi:
- kw.update(replay_work_items=wi, n_writes=nw, use_all_layer_commit=False)
- out = cr.cached_replay_update(**kw, in_kernel_commit=True, config=cfg)
- torch.cuda.synchronize()
- got.append([out.clone()] + [x.clone() for x in bm.state(0)])
- for name, x, y in zip(("out", "state", "u", "k", "G"), *got):
- assert torch.equal(x, y), f"{name}: {new_cfg} vs {old_cfg} (work items {with_wi})"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py
deleted file mode 100644
index 24ebe8c685a8..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py
+++ /dev/null
@@ -1,122 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the gdn_decode_gated_norm_pdl entry.
-
-Cells against the composition's fp64 reference (the gated delta rule's recurrence,
-then the gated norm entry's) in torch's default band for the bf16 output and the fp32
-pool, the pool compared whole so slots the call did not name are
-gated too. Beyond the cells: a stale-read check under graph replay (a producer
-rewrites the projection inside the graph right before the op), and the row guard.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_decode_gated_norm_pdl import ( # noqa: E501
- gdn_decode_gated_norm_pdl as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "gdn_decode_gated_norm_pdl requires a CUDA device"
-
-DEV = "cuda"
-H, HV, K, V = 16, 48, 128, 128
-SLOTS, EPS = 461, 1e-6
-WIDTH = 2 * H * K + 2 * HV * V + 2 * HV
-
-
-def _scattered(n: int, slots: int) -> tuple[int, ...]:
- """`n` distinct slots of a `slots`-row pool, both ends first, then a stride-37 walk
- over the interior (37 is coprime to every pool size here)."""
- rest = tuple((i * 37 + 4) % (slots - 2) + 1 for i in range(max(n - 2, 0)))
- return ((slots - 1, 0) + rest)[:n]
-
-
-def _build(spec, seed):
- g = torch.Generator(device=DEV).manual_seed(seed)
- rows = spec["rows"]
- proj = torch.randn(rows, WIDTH, generator=g, device=DEV).to(torch.bfloat16)
- if spec.get("pool") == "coalesced":
- # The KV-cache-v2 manager's view: every GDN layer's state of a slot in one page,
- # this layer's state one sub-page of it, so the slot stride is 36 layers' states
- # and slot 100 sits ~2.8G elements past the base (past int32).
- layers, slots = 36, 102
- page = torch.empty(slots * layers * HV * V * K, device=DEV)
- page[: HV * V * K * 4].normal_(0.0, 0.05, generator=g)
- pool = page.as_strided((slots, HV, V, K), (layers * HV * V * K, V * K, K, 1))
- for s in (100, 101):
- pool[s].copy_(torch.randn(HV, V, K, generator=g, device=DEV) * 0.05)
- else:
- pool = torch.randn(SLOTS, HV, V, K, generator=g, device=DEV) * 0.05
- if spec.get("pool") == "coalesced":
- slots = torch.tensor([100, 101][:rows], dtype=torch.int32, device=DEV)
- else:
- slots = torch.tensor(_scattered(rows, SLOTS), dtype=torch.int32, device=DEV)
- if spec.get("pad"):
- slots[1] = -1
- A_log = torch.empty(HV, device=DEV).uniform_(-5.5, 1.0, generator=g)
- dt_bias = torch.randn(HV, generator=g, device=DEV) * 4.0
- w = (1.0 + 0.2 * torch.randn(V, generator=g, device=DEV)).to(torch.bfloat16)
- return [proj, pool, slots, A_log, dt_bias, w, EPS, H, K]
-
-
-def _compare(out, pool, ref_out, ref_pool) -> None:
- assert out.dtype == torch.bfloat16 and out.shape == ref_out.shape and out.is_contiguous()
- assert torch.isfinite(out).all()
- torch.testing.assert_close(out, ref_out)
- torch.testing.assert_close(pool, ref_pool)
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
- certifying(op)
- args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31))
- ref_out, ref_pool = op.reference(*args)
- out = op(*args)
- _compare(out, args[1], ref_out, ref_pool)
- if cell.spec.get("pad"):
- assert (out[1] == 0).all(), "a pad row's output is zero"
-
-
-def test_reads_after_the_wait_under_graph_replay() -> None:
- """A producer rewrites the projection inside the graph right before the op; with PDL
- a read of it before `griddepcontrol.wait` would see the previous replay's row. The
- pool advances one step per replay, as in decode."""
- args = _build(dict(rows=1), seed=41)
- src = args[0].clone()
-
- def step():
- args[0].copy_(src * 1.0)
- return op(*args)
-
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- pool0 = args[1].clone()
- with torch.cuda.stream(warmup):
- step()
- torch.cuda.current_stream().wait_stream(warmup)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- graph_out = step()
- for trial in range(3):
- src.copy_(_build(dict(rows=1), seed=100 + trial)[0])
- before = args[1].clone()
- graph.replay()
- torch.cuda.synchronize()
- after = args[1].clone()
- args[1].copy_(before)
- want = op(src.clone(), *args[1:])
- assert torch.equal(graph_out, want), f"replay {trial} output"
- assert torch.equal(after, args[1]), f"replay {trial} pool"
- del pool0
-
-
-def test_guards() -> None:
- """More rows than `MAX_ROWS` are refused."""
- args = _build(dict(rows=1), seed=5)
- with validating(op), pytest.raises(AssertionError):
- op(args[0].expand(5, -1).contiguous(), args[1], args[2].expand(5).contiguous(), *args[3:])
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py
deleted file mode 100644
index 206ca93d1d3a..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py
+++ /dev/null
@@ -1,252 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the gdn_decode_rows_gated_norm entry.
-
-Cells against the composition's fp64 reference (the gated delta rule's recurrence,
-then the gated norm entry's) in torch's default band for the bf16 output and the pool's
-dtype (fp32 or bf16), the pool compared whole so slots the call did not name are
-gated too. Beyond the cells: a fixed row bitwise across batch sizes over a bf16 pool, a
-stale-read check under graph replay (a producer rewrites the projection inside the graph
-right before the op), the two launch variants (state rows by bulk copy into shared memory,
-or loaded into registers) bitwise each other, and the guards. The `gdn_decode_rows` entry
-(the recurrence alone) has its own cells below.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm import ( # noqa: E501
- gdn_decode_rows_gated_norm as rows_mod,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_decode_rows_gated_norm import ( # noqa: E501
- gdn_decode_rows as rows_op,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_decode_rows_gated_norm import ( # noqa: E501
- gdn_decode_rows_gated_norm as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "gdn_decode_rows_gated_norm requires a CUDA device"
-
-DEV = "cuda"
-H, HV, K, V = 16, 48, 128, 128
-SLOTS, EPS = 461, 1e-6
-WIDTH = 2 * H * K + 2 * HV * V + 2 * HV
-
-
-def _scattered(n: int, slots: int) -> tuple[int, ...]:
- """`n` distinct slots of a `slots`-row pool, both ends first, then a stride-37 walk
- over the interior (37 is coprime to every pool size here)."""
- rest = tuple((i * 37 + 4) % (slots - 2) + 1 for i in range(max(n - 2, 0)))
- return ((slots - 1, 0) + rest)[:n]
-
-
-def _build(spec, seed):
- g = torch.Generator(device=DEV).manual_seed(seed)
- rows = spec["rows"]
- state = torch.bfloat16 if spec.get("state") == "bf16" else torch.float32
- proj = torch.randn(rows, WIDTH, generator=g, device=DEV).to(torch.bfloat16)
- if spec.get("pool") == "coalesced":
- # The KV-cache-v2 manager's view: every GDN layer's state of a slot in one page,
- # this layer's state one sub-page of it, so the slot stride is 36 layers' states
- # and slot 100 sits ~2.8G elements past the base (past int32).
- layers, slots = 36, 102
- page = torch.empty(slots * layers * HV * V * K, device=DEV, dtype=state)
- page[: HV * V * K * 4].normal_(0.0, 0.05, generator=g)
- pool = page.as_strided((slots, HV, V, K), (layers * HV * V * K, V * K, K, 1))
- for s in (100, 101):
- pool[s].copy_(torch.randn(HV, V, K, generator=g, device=DEV) * 0.05)
- else:
- pool = (torch.randn(SLOTS, HV, V, K, generator=g, device=DEV) * 0.05).to(state)
- if spec.get("pool") == "coalesced":
- slots = torch.tensor([100, 101][:rows], dtype=torch.int32, device=DEV)
- else:
- slots = torch.tensor(_scattered(rows, SLOTS), dtype=torch.int32, device=DEV)
- if spec.get("pad"):
- slots[1] = -1
- A_log = torch.empty(HV, device=DEV).uniform_(-5.5, 1.0, generator=g)
- dt_bias = torch.randn(HV, generator=g, device=DEV) * 4.0
- w = (1.0 + 0.2 * torch.randn(V, generator=g, device=DEV)).to(torch.bfloat16)
- return [proj, pool, slots, A_log, dt_bias, w, EPS, H, K]
-
-
-def _compare(out, pool, ref_out, ref_pool) -> None:
- assert out.dtype == torch.bfloat16 and out.shape == ref_out.shape and out.is_contiguous()
- assert torch.isfinite(out).all()
- torch.testing.assert_close(out, ref_out)
- torch.testing.assert_close(pool, ref_pool)
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
- certifying(op)
- args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31))
- ref_out, ref_pool = op.reference(*args)
- out = op(*args)
- _compare(out, args[1], ref_out, ref_pool)
- if cell.spec.get("pad"):
- assert (out[1] == 0).all(), "a pad row's output is zero"
-
-
-#: Decode batch sizes up to the engine's largest CUDA-graph batch.
-_BATCHES = (1, 2, 4, 5, 16, 33, 128, 256)
-_FIXED_SLOT = 301
-
-
-def test_bf16_pool_is_batch_invariant() -> None:
- """The property the bf16 decode step relies on: a row's output and state are bitwise
- independent of the batch it rides in. One fixed row (its projection and slot 301 of a
- bf16 pool) sits first and last in random batches of every size."""
- gen = torch.Generator(device=DEV).manual_seed(3)
- cpu = torch.Generator().manual_seed(3)
- pool0 = (torch.randn(SLOTS, HV, V, K, generator=gen, device=DEV) * 0.05).to(torch.bfloat16)
- row = torch.randn(WIDTH, generator=gen, device=DEV).to(torch.bfloat16)
- rest = _build(dict(rows=1), seed=3)[3:]
- first = None
- for n in _BATCHES:
- for pos in sorted({0, n - 1}):
- g = torch.Generator(device=DEV).manual_seed(1000 * n + pos)
- proj = torch.randn(n, WIDTH, generator=g, device=DEV).to(torch.bfloat16)
- proj[pos] = row
- others = [s for s in torch.randperm(SLOTS, generator=cpu).tolist() if s != _FIXED_SLOT]
- others = others[:n]
- others[pos] = _FIXED_SLOT
- slots = torch.tensor(others, dtype=torch.int32, device=DEV)
- pool = pool0.clone()
- out = op(proj, pool, slots, *rest)
- got = (out[pos].clone(), pool[_FIXED_SLOT].clone())
- first = first or got
- assert torch.equal(got[0], first[0]) and torch.equal(got[1], first[1]), (
- f"the fixed row differs at batch {n}, position {pos}"
- )
-
-
-def test_reads_after_the_wait_under_graph_replay() -> None:
- """A producer rewrites the projection inside the graph right before the op; with PDL
- a read of it before `griddepcontrol.wait` would see the previous replay's row. The
- pool advances one step per replay, as in decode."""
- args = _build(dict(rows=16), seed=41)
- src = args[0].clone()
-
- def step():
- args[0].copy_(src * 1.0)
- return op(*args)
-
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(warmup):
- step()
- torch.cuda.current_stream().wait_stream(warmup)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- graph_out = step()
- for trial in range(3):
- src.copy_(_build(dict(rows=16), seed=100 + trial)[0])
- before = args[1].clone()
- graph.replay()
- torch.cuda.synchronize()
- after = args[1].clone()
- args[1].copy_(before)
- want = op(src.clone(), *args[1:])
- assert torch.equal(graph_out, want), f"replay {trial} output"
- assert torch.equal(after, args[1]), f"replay {trial} pool"
-
-
-_VARIANT_SPECS = (
- dict(rows=1, state="bf16"),
- dict(rows=33, state="bf16"),
- dict(rows=33, pad=True),
- dict(rows=2, pool="coalesced", state="bf16", pad=True),
-)
-
-
-@pytest.mark.parametrize("norm", [True, False], ids=["norm", "rows"])
-@pytest.mark.parametrize("spec", _VARIANT_SPECS, ids=lambda s: "-".join(map(str, s.values())))
-def test_bulk_and_register_variants_agree(spec, norm) -> None:
- """The bulk-copy and the register variants: outputs and state bitwise equal.
-
- The bulk-copy variant stages the state rows in shared memory; the register variant holds
- them in registers; both compute the same instructions. Over contiguous bf16 and fp32 pools
- and the cache manager's coalesced pool (slot 100, ~2.8G elements past the base, by the
- bulk copy), with pad rows.
- """
- args = _build(spec, seed=zlib.crc32(repr(spec).encode()) % (2**31))
- proj, pool, slots = args[0], args[1], args[2]
- assert rows_mod._bulk_eligible(pool, K)
- live = slots[slots >= 0].long()
- before = pool[live].clone()
- got = []
- for bulk in (False, True):
- pool[live] = before
- w = args[5] if norm else None
- out = rows_mod.launch(proj, pool, slots, args[3], args[4], w, EPS, H, K, bulk=bulk)
- got.append((out.clone(), pool[live].clone()))
- assert torch.equal(got[0][0], got[1][0]), "outputs differ between the variants"
- assert torch.equal(got[0][1], got[1][1]), "states differ between the variants"
-
-
-def test_variant_choice_and_unaligned_pool() -> None:
- """The variant choice, and the fallback for a pool the bulk copy cannot take.
-
- The bulk variant from the batch the register variant needs a second wave for; a pool 8-
- but not 16-byte aligned stays on the register variant whatever is asked, and matches the
- aligned pool's result.
- """
- sms = torch.cuda.get_device_properties(DEV).multi_processor_count
- edge = (2 * sms) // HV
- assert not rows_mod._use_bulk(edge, HV, DEV) and rows_mod._use_bulk(edge + 1, HV, DEV)
- args = _build(dict(rows=33, state="bf16"), seed=19)
- pool = args[1]
- flat = torch.empty(pool.numel() + 4, dtype=pool.dtype, device=DEV)
- shifted = flat[4:].view(pool.shape)
- shifted.copy_(pool)
- assert not rows_mod._bulk_eligible(shifted, K)
- want = rows_mod.launch(args[0], pool, *args[2:], bulk=True)
- got = rows_mod.launch(args[0], shifted, *args[2:], bulk=True)
- assert torch.equal(got, want) and torch.equal(shifted, pool)
-
-
-def test_guards() -> None:
- """An fp16 pool, int64 slot indices and a projection of the wrong width are refused."""
- args = _build(dict(rows=4), seed=5)
- with validating(op), pytest.raises(AssertionError):
- op(args[0], args[1].to(torch.float16), *args[2:])
- with validating(op), pytest.raises(AssertionError):
- op(args[0], args[1], args[2].long(), *args[3:])
- with validating(op), pytest.raises(AssertionError):
- op(args[0][:, :-2], *args[1:])
-
-
-# ---- gdn_decode_rows: the recurrence alone, into a row range of a step's output ----------
-
-
-def _rows_args(spec, seed):
- """`(args, buf)`: the entry's arguments, `out` rows 5.. of a `[1, rows + 9, HV, V]` buffer
- (the step's output, the call's rows a range of it), and that buffer."""
- proj, pool, slots, A_log, dt_bias, _, _, h, k = _build(spec, seed)
- rows = proj.shape[0]
- g = torch.Generator(device=DEV).manual_seed(seed + 1)
- buf = torch.randn(1, rows + 9, HV, V, generator=g, device=DEV).to(torch.bfloat16)
- out = buf[0, 5 : 5 + rows].view(rows, HV * V)
- return [proj, pool, slots, A_log, dt_bias, h, k, out], buf
-
-
-@pytest.mark.parametrize("cell", rows_op.CELLS, ids=[c.why[:44] for c in rows_op.CELLS])
-def test_rows_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name
- certifying(rows_op)
- args, buf = _rows_args(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31))
- outside = torch.cat([buf[0, :5], buf[0, 5 + args[0].shape[0] :]]).clone()
- ref_o, ref_pool = rows_op.reference(*args)
- got = rows_op(*args)
- assert got.data_ptr() == args[-1].data_ptr(), "written into out"
- torch.testing.assert_close(got, ref_o)
- torch.testing.assert_close(args[1], ref_pool)
- assert torch.equal(torch.cat([buf[0, :5], buf[0, 5 + args[0].shape[0] :]]), outside), (
- "rows outside the range untouched"
- )
- if cell.spec.get("pad"):
- assert (got[1] == 0).all(), "a pad row's output is zero"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py
deleted file mode 100644
index e6a5f55bd730..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py
+++ /dev/null
@@ -1,103 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the gdn_in_proj_conv1d_pdl entry.
-
-Cells against the entry's fp32 reference: the output within the band the
-decode projections use (4 bf16 ulp element-wise against the row scale, 2 ulp
-relative RMS -- the GEMV's summation order is the only difference), the pool
-the same, a pad row's slot bitwise untouched. Beyond the cells: two decode
-steps under CUDA-graph replay, the second reading the window the first one
-wrote; and the row guard.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_in_proj_conv1d_pdl import ( # noqa: E501
- MAX_ROWS,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_in_proj_conv1d_pdl import ( # noqa: E501
- gdn_in_proj_conv1d_pdl as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import ( # noqa: E402,F401 — path above; fixture by name
- ULP,
- assert_within_ulp,
- certifying,
- validating,
-)
-
-assert torch.cuda.is_available(), "gdn_in_proj_conv1d_pdl requires a CUDA device"
-
-N, K, CONV_DIM, PAD = 16480, 2560, 10240, -1
-
-
-def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
- assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0, ulp=ULP[torch.bfloat16])
-
-
-def _build(rows, pool, seed):
- g = torch.Generator(device="cuda").manual_seed(seed)
- x = torch.randn(rows, K, generator=g, device="cuda").to(torch.bfloat16)
- weight = (torch.randn(N, K, generator=g, device="cuda") * K**-0.5).to(torch.bfloat16)
- state = torch.randn(pool, CONV_DIM, 3, generator=g, device="cuda").to(torch.bfloat16)
- taps = (torch.randn(CONV_DIM, 4, generator=g, device="cuda") * 0.3).to(torch.bfloat16)
- return x, weight.t(), state, taps
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811
- certifying(op)
- s = cell.spec
- x, mat_b, state, taps = _build(s["rows"], s["pool"], seed=zlib.crc32(cell.why.encode()) % 2**31)
- slots = torch.tensor(s["slots"], dtype=torch.int32, device="cuda")
- pool = state.clone()
- out = op(x, mat_b, pool, taps, slots, CONV_DIM, PAD, s["pdl"])
- ref_out, ref_state = op.reference(x, mat_b, state, taps, slots, CONV_DIM, PAD, s["pdl"])
- assert out.shape == (s["rows"], N) and out.dtype == torch.bfloat16 and out.is_contiguous()
- _compare(out, ref_out)
- _compare(pool, ref_state)
- touched = {v for v in s["slots"] if v != PAD}
- for slot in range(s["pool"]):
- if slot not in touched:
- assert torch.equal(pool[slot], state[slot]), f"slot {slot} was written"
-
-
-def test_two_steps_under_graph_replay() -> None:
- """Decode runs this under CUDA graphs with PDL, layer after layer: capture one
- step, replay it twice with new activations. The second replay's conv reads the
- window the first replay advanced, so a stale or misplaced window write fails
- the second comparison."""
- pdl = True
- x, mat_b, state, taps = _build(1, 64, seed=3)
- slots = torch.tensor([9], dtype=torch.int32, device="cuda")
- pool = state.clone()
- stream = torch.cuda.Stream()
- stream.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(stream):
- for _ in range(2):
- op(x, mat_b, pool.clone(), taps, slots, CONV_DIM, PAD, pdl)
- torch.cuda.current_stream().wait_stream(stream)
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- out = op(x, mat_b, pool, taps, slots, CONV_DIM, PAD, pdl)
- want_state = state.clone()
- for trial in range(2):
- g = torch.Generator(device="cuda").manual_seed(50 + trial)
- x.copy_(torch.randn(x.shape, generator=g, device="cuda").to(x.dtype))
- graph.replay()
- torch.cuda.synchronize()
- ref_out, want_state = op.reference(x, mat_b, want_state, taps, slots, CONV_DIM, PAD)
- _compare(out, ref_out)
- _compare(pool, want_state)
-
-
-def test_guards() -> None:
- """More rows than `MAX_ROWS` are refused."""
- x, mat_b, state, taps = _build(MAX_ROWS + 1, 16, seed=21)
- slots = torch.zeros(MAX_ROWS + 1, dtype=torch.int32, device="cuda")
- with validating(op), pytest.raises(AssertionError):
- op(x, mat_b, state, taps, slots, CONV_DIM, PAD)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_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
deleted file mode 100644
index 9cfca29e1712..000000000000
--- 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
+++ /dev/null
@@ -1,156 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the gdn_prefill_conv_post entry.
-
-Every cell is checked against the entry's fp32/fp64 reference, the conv pool
-bitwise. Geometry is a Qwen3.8-Flash-Next linear-attention layer: 16 x 128 key heads,
-48 x 128 value heads, qkv_dim 10240, conv width 4, bf16 activations.
-"""
-
-import math
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_prefill_conv_post import ( # noqa: E501
- A_LOG_RANGE,
- CONV_WIDTH,
- DT_BIAS_RANGE,
- HEAD_K_DIM,
- HEAD_V_DIM,
- IN_PROJ_WIDTH,
- NUM_K_HEADS,
- NUM_V_HEADS,
- QKV_DIM,
- SOFTPLUS_THRESHOLD,
- _engine_slots,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_prefill_conv_post import ( # noqa: E501
- gdn_prefill_conv_post as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "gdn_prefill_conv_post requires a CUDA device"
-
-DEV = "cuda"
-DTYPE = torch.bfloat16
-HEADS = (NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS, HEAD_V_DIM)
-PAD = -1
-L2_NORM_EPS = 1e-6
-#: `g`'s band, derived as in the fused_gdn_post_conv entry test: two SASS
-#: approximate transcendentals (~2^-22 relative each), the softplus bounded by
-#: the threshold, every term scaled by exp(A_log).
-_APPROX_ULP = 2.0**-22
-G_ATOL = 4 * _APPROX_ULP * SOFTPLUS_THRESHOLD * math.exp(A_LOG_RANGE[1])
-A_COL = QKV_DIM + NUM_V_HEADS * HEAD_V_DIM
-
-
-def _build(spec, seed):
- """`__call__`'s positional arguments for one cell."""
- g = torch.Generator(device=DEV).manual_seed(seed)
- lengths = spec["lengths"]
- n_tok = sum(lengths)
-
- def randn(*shape):
- return torch.randn(*shape, generator=g, device=DEV).to(DTYPE)
-
- # "in_proj" and "rows" both carve x, a and b out of one projection; "rows"
- # leaves `extra_rows` trailing rows the call must not read.
- proj = randn(n_tok + spec.get("extra_rows", 0), IN_PROJ_WIDTH)
- x = proj[:n_tok, :QKV_DIM]
- a = proj[:n_tok, A_COL : A_COL + NUM_V_HEADS]
- b = proj[:n_tok, A_COL + NUM_V_HEADS :]
- assert x.stride(0) == a.stride(0) == IN_PROJ_WIDTH
- weight = (torch.randn(QKV_DIM, CONV_WIDTH, generator=g, device=DEV) * 0.5).to(DTYPE)
- pool_slots = spec["pool_slots"]
- # The pool is a per-layer view of a two-layer buffer, as a cache manager
- # hands it out: slot stride twice the slot's own size.
- conv_states = randn(pool_slots, 2, QKV_DIM, CONV_WIDTH - 1)[:, 1]
- slots = spec.get("slots") or _engine_slots(len(lengths))
- bounds = [0]
- for n in lengths:
- bounds.append(bounds[-1] + n)
- qsl = torch.tensor(bounds, dtype=torch.int32, device=DEV)
- si = torch.tensor(slots, dtype=torch.int32, device=DEV)
- hi = torch.tensor(spec["initial"], dtype=torch.bool, device=DEV)
- lo, hi_ = A_LOG_RANGE
- A_log = torch.rand(NUM_V_HEADS, generator=g, device=DEV) * (hi_ - lo) + lo
- lo, hi_ = DT_BIAS_RANGE
- dt_bias = torch.rand(NUM_V_HEADS, generator=g, device=DEV) * (hi_ - lo) + lo
- return [
- x,
- weight,
- conv_states,
- qsl,
- si,
- hi,
- a,
- b,
- A_log,
- dt_bias,
- *HEADS,
- PAD,
- L2_NORM_EPS,
- SOFTPLUS_THRESHOLD,
- spec.get("linear_gate", False),
- ]
-
-
-def _compare(out, pool, want, what=""):
- q, k, v, g, beta = out
- q_r, k_r, v_r, g_r, beta_r, pool_r = want
- for got, ref in zip(out, want[:5]):
- assert got.dtype == ref.dtype and got.shape == ref.shape and got.is_contiguous()
- for name, got, ref in (("q", q, q_r), ("k", k, k_r), ("v", v, v_r)):
- d = (got.float() - ref.float()).abs()
- flips = (d > 0).float().mean().item() if d.numel() else 0.0
- print(
- f"{what} {name}: max abs {d.max().item() if d.numel() else 0:.3e}, "
- f"differing fraction {flips:.2e}"
- )
- torch.testing.assert_close(got, ref)
- torch.testing.assert_close(g, g_r, rtol=8 * _APPROX_ULP, atol=G_ATOL)
- torch.testing.assert_close(beta, beta_r)
- assert torch.equal(pool, pool_r), "conv pool differs"
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
- want = op.reference(*args)
- inputs = [t.clone() if isinstance(t, torch.Tensor) else t for t in args]
- out = op(*args)
- for i, (arg, before) in enumerate(zip(args, inputs)):
- if isinstance(arg, torch.Tensor) and i != 2:
- assert torch.equal(arg, before), f"input {i} must be read-only"
- _compare(out, args[2], want, what=f"{cell.why[:30]} vs reference")
-
-
-def test_guard_refuses_channel_strided_x() -> None:
- args = _build(dict(lengths=(16,), pool_slots=2, slots=(0,), initial=(True,)), seed=302)
- bad = list(args)
- bad[0] = torch.empty(QKV_DIM, 16, dtype=DTYPE, device=DEV).t()
- with pytest.raises(AssertionError, match="unit channel stride"):
- with validating(op):
- op(*bad)
-
-
-def test_linear_gate_is_torch_exp_of_the_log_gate() -> None:
- """`linear_gate=True` stores `torch.exp` of the log-space decay bitwise, and leaves
- every other output and the conv pool as without it."""
- spec = next(c for c in op.CELLS if c.spec.get("linear_gate")).spec
- args = _build(spec, seed=77)
- log_args = list(args)
- log_args[-1] = False
- log_args[2] = args[2].clone()
- pool_lin = args[2]
- log_out = op(*log_args)
- lin_out = op(*args)
- assert torch.equal(lin_out[3], torch.exp(log_out[3]))
- for got, want in zip(lin_out[:3] + lin_out[4:], log_out[:3] + log_out[4:]):
- assert torch.equal(got, want)
- assert torch.equal(pool_lin, log_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_replay_verify_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_replay_verify_cute.py
deleted file mode 100644
index 322d994709bd..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_replay_verify_cute.py
+++ /dev/null
@@ -1,128 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the target's cached-replay verify dispatch over a bf16 pool (the CuTe kernel).
-
-`_replay_verify` as the decode target calls it -- `a`/`b` as column views of the input
-projection, PDL after the fused in_proj GEMV -- against the call it replaces (contiguous
-`a`/`b`, no PDL), each on its own copy of the pool and the history, over consecutive verify
-steps with random partial acceptance and `_advance_replay_state` in between: outputs, the
-checkpoint (committed in-kernel on every wrap, as without the manager's all-layer commit) and
-the history bitwise equal at every step. The partitioned replay (work items and the manager's
-all-layer commit) is the CuTe entry's own test's.
-"""
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501
- modeling as M,
-)
-from tensorrt_llm._torch.modules.fla.cache_manager import (
- ReplayStateUpdateMetadata,
- _advance_replay_state,
-)
-
-assert torch.cuda.is_available(), "the GDN verify dispatch test requires a CUDA device"
-
-DEV = "cuda"
-BF = torch.bfloat16
-H, HV, K, V = 16, 48, 128, 128
-KEY_DIM, VALUE_DIM = H * K, HV * V
-CONV_DIM = 2 * KEY_DIM + VALUE_DIM
-HIST = 16
-SLOTS = 40
-
-
-class _Side:
- """One copy of a layer's replay state and the `core` stand-in `_replay_verify` reads.
-
- The state: the pool, the history and the bookkeeping (pending count, half).
- """
-
- def __init__(self, pool, old, pnat, buf, t, gating):
- self.pool = pool.clone()
- self.old = [x.clone() for x in old]
- self.pnat, self.buf = pnat.clone(), buf.clone()
- self.meta = ReplayStateUpdateMetadata(self.pnat, self.buf, t, HIST)
- replay = M._Replay(
- old={0: tuple(self.old)},
- cache_buf_idx=self.buf,
- prev_num_accepted=self.pnat,
- history=HIST,
- width=t,
- all_layer_commit=False,
- cute=True,
- )
- self.core = SimpleNamespace(
- conv_dim=CONV_DIM,
- key_dim=KEY_DIM,
- value_dim=VALUE_DIM,
- num_k_heads=H,
- head_k_dim=K,
- num_v_heads=HV,
- head_v_dim=V,
- _bound=SimpleNamespace(replay=replay),
- )
- self.layer = SimpleNamespace(a_log=gating[0], dt_bias=gating[1])
-
- def state(self):
- return [self.pool, *self.old[:3]]
-
-
-@pytest.mark.parametrize("n,t", [(1, 4), (4, 4), (32, 4), (5, 3)])
-def test_cute_dispatch_matches_the_plain_call(n, t):
- gen = torch.Generator(device=DEV).manual_seed(77 + n)
- cpu = torch.Generator().manual_seed(77 + n)
- slots = torch.randperm(SLOTS, generator=cpu)[:n].to(DEV, torch.int32)
- pool = (0.3 * torch.randn(SLOTS, HV, V, K, generator=gen, device=DEV)).to(BF)
- old = (
- (0.5 * torch.randn(SLOTS, 2, HIST, HV, V, generator=gen, device=DEV)).to(BF),
- torch.nn.functional.normalize(
- torch.randn(SLOTS, 2, HIST, H, K, generator=gen, device=DEV), dim=-1
- ).to(BF),
- -(0.2 * torch.rand(SLOTS, 2, HV, HIST, generator=gen, device=DEV)).cumsum(-1),
- torch.zeros(SLOTS, 2, HV, HIST, device=DEV),
- )
- pnat = torch.randint(0, HIST + 1, (SLOTS,), generator=cpu, dtype=torch.int32).to(DEV)
- buf = torch.randint(0, 2, (SLOTS,), generator=cpu, dtype=torch.int32).to(DEV)
- gating = (
- (torch.randn(HV, generator=gen, device=DEV) * 1.5 + 1.0).clamp(-7.0, 5.0),
- (torch.randn(HV, generator=gen, device=DEV) * 1.7 - 4.4).clamp(-11.0, 2.5),
- )
- plain = _Side(pool, old, pnat, buf, t, gating)
- disp = _Side(pool, old, pnat, buf, t, gating)
- wraps = 0
- for step in range(HIST // t + 3):
- wraps += int((disp.pnat[slots.long()] + t > HIST).sum())
- replay = M._replay_work(disp.core, SimpleNamespace(), n)
- assert replay == (None, None, False), "no all-layer commit: the kernel commits itself"
- mq_out = (0.5 * torch.randn(n * t, CONV_DIM, generator=gen, device=DEV)).to(BF)
- # z | a | b: column views of one projection output, as the decode target holds them
- zab = torch.randn(n * t, VALUE_DIM + 2 * HV, generator=gen, device=DEV).to(BF)
- _, a, b = torch.split(zab, [VALUE_DIM, HV, HV], 1)
- want = M._replay_verify(
- plain.core,
- mq_out,
- a.contiguous(),
- b.contiguous(),
- plain.pool,
- 0,
- plain.layer,
- slots,
- n,
- t,
- (None, None, False),
- )
- got = M._replay_verify(
- disp.core, mq_out, a, b, disp.pool, 0, disp.layer, slots, n, t, replay, pdl=True
- )
- torch.cuda.synchronize()
- assert torch.equal(got, want), f"step {step}: outputs differ"
- for name, x, y in zip(("pool", "u", "k", "G"), disp.state(), plain.state()):
- assert torch.equal(x, y), f"step {step}: {name} differs"
- accepted = torch.randint(1, t + 1, (n,), generator=cpu, dtype=torch.int32).to(DEV)
- for side in (plain, disp):
- _advance_replay_state(side.meta, slots, accepted)
- assert wraps > 0, "no history wrapped: the in-kernel commit was never exercised"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_verify_fused_conv.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_verify_fused_conv.py
deleted file mode 100644
index 8476e49feada..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_verify_fused_conv.py
+++ /dev/null
@@ -1,100 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test: `DecodeTarget._gdn_block`'s verify dispatch under the dense FP8 mode, and over the
-in_proj's NVFP4 copy (`TRTLLM_MODELING_V2_GDN_INPROJ_NVFP4`, up to `_INPROJ_NVFP4_MAX_ROWS`).
-
-A captured verify step (`fused_conv=True`) runs the in_proj with the multi-token conv step as its
-epilogue (`fp8_gemv_conv1d_verify` / `nvfp4_split_gemv_conv1d_verify`); an eager one runs the
-projection and then `causal_conv1d_triton_update`. Both must hand the recurrence the same thing: the conv output
-(contiguous, as the cached replay reads it), `a`, `b` and `z` bitwise, and leave the conv pool
-and the per-step snapshot buffer bitwise the same. The recurrence is stubbed to record its
-inputs (it is the same call on both sides), the epilogue to record `z`. Over a bf16 pool (the
-CuTe replay) `a` and `b` reach the recurrence as views of the projection (no copies), and only
-the fused step asks for PDL (its GEMV signals its dependents; the Triton conv does not).
-"""
-
-import types
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501
- modeling as M,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.fp8_gemv import ( # noqa: E501
- pack_from_bf16,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.nvfp4_gemv import ( # noqa: E501
- pack_split_from_bf16,
-)
-
-assert torch.cuda.is_available(), "the Gated-DeltaNet block requires a CUDA device"
-
-DEV, BF16 = "cuda", torch.bfloat16
-HIDDEN, HV, V, CONV, VD, SLOTS, ROWS = 2560, 48, 128, 10240, 6144, 461, 32
-
-
-@pytest.fixture(scope="module", params=["fp8", "nvfp4"])
-def layer(request):
- g = torch.Generator(device=DEV).manual_seed(0)
- w = (torch.randn(CONV + VD + 2 * HV, HIDDEN, generator=g, device=DEV) * HIDDEN**-0.5).to(BF16)
- nv = request.param == "nvfp4"
- return types.SimpleNamespace(
- in_proj=w.t(),
- in_proj8=None,
- in_proj_m=None,
- in_proj_g=None if nv else pack_from_bf16(w),
- in_proj_nv=pack_split_from_bf16(w) if nv else None,
- conv_w=(torch.randn(CONV, 4, generator=g, device=DEV) * 0.5).to(BF16),
- )
-
-
-def _block(layer, x, si, int_idx, conv0, win0, p, fused, monkeypatch):
- """One verify `_gdn_block` on copies of the pools; returns what the recurrence and the
- epilogue were handed, and the pools."""
- seen = {}
-
- def replay(core, mq_out, a_d, b_d, ssm_states, i, L, si_d, n, p_, work, pdl=False):
- seen.update(mq_out=mq_out.clone(), a=a_d.clone(), b=b_d.clone(), pdl=pdl)
- assert mq_out.is_contiguous() and mq_out.shape == (n * p_, CONV)
- # the CuTe replay reads a / b through their row stride: views, not copies
- assert a_d.stride(0) > HV and b_d.stride(0) > HV
- return torch.zeros(1, n * p_, HV, V, dtype=BF16, device=DEV)
-
- monkeypatch.setattr(M, "_replay_verify", replay)
- core = types.SimpleNamespace(
- conv_dim=CONV,
- in_proj_split=(CONV, VD, HV, HV),
- _bound=types.SimpleNamespace(
- conv={0: conv0.clone()},
- ssm={0: torch.zeros(1, device=DEV)},
- int_conv={0: win0.clone()},
- replay=types.SimpleNamespace(cute=True),
- ),
- )
- target = object.__new__(M.DecodeTarget)
- target.core = core
- target._gdn_epilogue = lambda out, z, L, num_tokens, gemm=None: z.clone()
- n = si.shape[0]
- gdn_state = (si, True, n, p, int_idx, (None, None, False))
- z = target._gdn_block(x, 0, layer, gdn_state, cublas_mm, fused)
- return dict(seen, z=z, conv=core._bound.conv[0], win=core._bound.int_conv[0])
-
-
-@pytest.mark.parametrize("n,p", [(1, 4), (2, 4), (8, 2), (2, 8)])
-def test_fused_verify_matches_two_launch(layer, n, p, monkeypatch) -> None:
- if layer.in_proj_nv is not None and n * p > M._INPROJ_NVFP4_MAX_ROWS:
- pytest.skip("past the NVFP4 copy's rows")
- g = torch.Generator(device=DEV).manual_seed(10 * n + p)
- conv0 = torch.randn(SLOTS, CONV, 3, generator=g, device=DEV).to(BF16)
- win0 = torch.randn(ROWS, p, CONV, 3, generator=g, device=DEV).to(BF16)
- x = torch.randn(n * p, HIDDEN, generator=g, device=DEV).to(BF16)
- si = torch.randperm(SLOTS, generator=g, device=DEV)[:n].to(torch.int32)
- int_idx = torch.arange(ROWS, dtype=torch.int32, device=DEV)[:n]
- eager = _block(layer, x, si, int_idx, conv0, win0, p, False, monkeypatch)
- fused = _block(layer, x, si, int_idx, conv0, win0, p, True, monkeypatch)
- for key in ("mq_out", "a", "b", "z", "conv", "win"):
- assert torch.equal(eager[key].view(torch.int16), fused[key].view(torch.int16)), key
- assert fused["pdl"] and not eager["pdl"], "PDL exactly after the fused GEMV"
- assert not torch.equal(fused["conv"], conv0) and not torch.equal(fused["win"], win0)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_pack_gdn_decode_qkv.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_pack_gdn_decode_qkv.py
deleted file mode 100644
index 4adc8c137f6f..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_pack_gdn_decode_qkv.py
+++ /dev/null
@@ -1,117 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the pack_gdn_decode_qkv entry.
-
-The op is a pure copy, so every comparison is bitwise.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.pack_gdn_decode_qkv import ( # noqa: E501
- CKPT_HEADS,
- HEAD_K_DIM,
- HEAD_V_DIM,
- NUM_K_HEADS,
- NUM_V_HEADS,
- QKV_DIM,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.pack_gdn_decode_qkv import ( # noqa: E501
- pack_gdn_decode_qkv as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "pack_gdn_decode_qkv requires a CUDA device"
-
-
-def _compare(out: tuple, ref: tuple) -> None:
- """Bitwise, all three: the op moves bits and changes nothing."""
- assert len(out) == len(ref) == 3
- for name, got, want in zip("qkv", out, ref):
- assert got.dtype == want.dtype and got.shape == want.shape, (
- name,
- (got.dtype, got.shape),
- (want.dtype, want.shape),
- )
- assert torch.equal(got, want), f"{name} must be a bit-exact copy"
-
-
-DEV = "cuda"
-Q_WIDTH = NUM_K_HEADS * HEAD_K_DIM # 2048
-V_WIDTH = NUM_V_HEADS * HEAD_V_DIM # 6144
-
-
-def _build(spec, seed):
- """`(mixed_qkv, conv_out or None)`: the bf16 input, and for a verify batch the
- `[B, qkv_dim, P]` conv output it was flattened from."""
- g = torch.Generator(device=DEV).manual_seed(seed)
-
- def randn(*shape):
- return torch.randn(*shape, generator=g, device=DEV).to(torch.bfloat16)
-
- if "draft" in spec:
- b, p = spec["num_decodes"], spec["draft"]
- # The target: the conv result carries its `[n, C, p]` input view's
- # strides, so transposing back and flattening is a contiguous view.
- conv_out = randn(b, p, QKV_DIM).transpose(1, 2)
- mixed = conv_out.transpose(1, 2).reshape(b * p, QKV_DIM)
- assert mixed.is_contiguous() and mixed.data_ptr() == conv_out.data_ptr()
- return mixed, conv_out
- return randn(spec["num_tokens"], QKV_DIM), None
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- spec = cell.spec
- mixed, conv_out = _build(spec, seed=zlib.crc32(cell.why.encode()))
- snapshot = mixed.clone()
- out = op(mixed, *CKPT_HEADS)
- _compare(out, op.reference(mixed, *CKPT_HEADS))
-
- assert all(t.is_contiguous() for t in out)
- assert torch.equal(mixed, snapshot), "mixed_qkv is read-only"
- storages = {t.untyped_storage().data_ptr() for t in out}
- assert len(storages) == 3 and mixed.untyped_storage().data_ptr() not in storages
-
- if "draft" in spec:
- # The verify caller regroups each return as [B, P, heads, dim]: that is
- # a view of the same storage, and it lands on (sequence, draft) exactly.
- b, p = spec["num_decodes"], spec["draft"]
- grouped = (
- out[0].reshape(b, p, NUM_K_HEADS, HEAD_K_DIM),
- out[1].reshape(b, p, NUM_K_HEADS, HEAD_K_DIM),
- out[2].reshape(b, p, NUM_V_HEADS, HEAD_V_DIM),
- )
- assert all(gt.data_ptr() == t.data_ptr() for gt, t in zip(grouped, out))
- per_token = conv_out.permute(0, 2, 1) # [B, P, qkv_dim], straight off the conv
- want = (
- per_token.narrow(2, 0, Q_WIDTH).reshape(b, p, NUM_K_HEADS, HEAD_K_DIM),
- per_token.narrow(2, Q_WIDTH, Q_WIDTH).reshape(b, p, NUM_K_HEADS, HEAD_K_DIM),
- per_token.narrow(2, 2 * Q_WIDTH, V_WIDTH).reshape(b, p, NUM_V_HEADS, HEAD_V_DIM),
- )
- assert all(torch.equal(gt, w) for gt, w in zip(grouped, want))
-
-
-def test_a_short_row_reads_the_next_row_and_is_refused() -> None:
- """Unguarded, a `[32, 9216]` input at qkv_dim 10240 runs and silently reads
- the next row's first 1024 columns into the tail of `v`; the guard refuses it.
-
- The input is carved from a longer flat buffer so the over-read lands in
- known memory rather than in allocator garbage.
- """
- rows, width = 32, QKV_DIM - 1024
- flat = torch.randn(rows * width + 1024, device=DEV).to(torch.bfloat16)
- mixed = flat[: rows * width].view(rows, width)
- _, _, v = op(mixed, *CKPT_HEADS)
- torch.cuda.synchronize()
- tail = v[0].reshape(rows, V_WIDTH)[:, V_WIDTH - 1024 :]
- assert torch.equal(tail[:-1], mixed[1:, :1024]), "the hazard this guard covers is gone"
-
- with pytest.raises(AssertionError, match="must be at least qkv_dim=10240 wide"):
- with validating(op):
- op(mixed, *CKPT_HEADS)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_qwen4_exp_causal_conv1d_fwd.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_qwen4_exp_causal_conv1d_fwd.py
deleted file mode 100644
index 5ca8d32dfed5..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_qwen4_exp_causal_conv1d_fwd.py
+++ /dev/null
@@ -1,131 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the causal_conv1d_fwd entry."""
-
-import zlib
-from typing import Optional
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_fwd import ( # noqa: E501
- ENGINE_SLOTS,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_fwd import ( # noqa: E501
- causal_conv1d_fwd as op,
-)
-
-# Cells sized at the engine's pool get a non-vacuity check below.
-assert ENGINE_SLOTS == 461
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "causal_conv1d_fwd requires a CUDA device"
-
-
-def _compare(
- out: tuple[torch.Tensor, Optional[torch.Tensor]],
- ref: tuple[torch.Tensor, Optional[torch.Tensor]],
-) -> None:
- """`(activation, conv_states)` pairs: the returned activation with the
- pool as the call left it, against `reference`'s pair.
-
- Activation at torch's default band: the kernel accumulates the taps in
- fp32 and rounds once, so it sits within one ulp of the fp32 reference
- (checkpoint geometry: max abs 0.015625, bitwise fraction 0.999996; fp16
- one ulp; fp32 1.4e-6). The state is a copy of the input tail, so it is
- gated bitwise over the whole pool -- which is also the lock on every
- slot the batch did not name.
- """
- (act, states), (ref_act, ref_states) = out, ref
- assert act.dtype == ref_act.dtype and act.shape == ref_act.shape
- torch.testing.assert_close(act, ref_act)
- assert (states is None) == (ref_states is None)
- if states is not None:
- assert torch.equal(states, ref_states), "the published state must be bitwise exact"
-
-
-PAD_SLOT_ID = -1
-
-
-def _seed(text: str) -> int:
- return zlib.crc32(text.encode())
-
-
-def _build(spec, seed):
- """`__call__`'s positional arguments for one spec, freshly allocated."""
- g = torch.Generator(device="cuda").manual_seed(seed)
- dim, width, dtype = spec["dim"], 4, torch.bfloat16
-
- def randn(*shape):
- return torch.randn(*shape, generator=g, device="cuda").to(dtype)
-
- lengths = spec["lengths"]
- starts = [0]
- for length in lengths:
- starts.append(starts[-1] + length)
- qsl = torch.tensor(starts, dtype=torch.int32, device="cuda")
- x = randn(sum(lengths), dim).t().contiguous()
- weight = randn(dim, width)
- pool = randn(spec["pool_slots"], dim, width - 1)
- cache_indices = torch.tensor(spec["slots"], dtype=torch.int32, device="cuda")
- his = torch.tensor(spec["initial"], dtype=torch.bool, device="cuda")
- return [x, weight, None, pool, qsl, cache_indices, his, True, PAD_SLOT_ID, None]
-
-
-def _clone(args):
- """The mutable operands cloned (x, conv_states)."""
- got = list(args)
- for i in (0, 3):
- got[i] = got[i].clone()
- return got
-
-
-def _run(args):
- """Drive the op on clones; return (returned activation, pool after, clones)."""
- got = _clone(args)
- returned = op(*got)
- return returned, got[3], got
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- args = _build(cell.spec, seed=_seed(cell.why))
- ref = op.reference(*args)
- returned, pool, got = _run(args)
- assert returned is got[0], "returns the tensor it wrote"
- _compare((returned, pool), ref)
- if cell.spec["pool_slots"] == ENGINE_SLOTS:
- # The whole-pool bitwise gate above already locks the unnamed slots;
- # this keeps it from being vacuous on the large pool.
- for slot in cell.spec["slots"]:
- assert not torch.equal(pool[slot], args[3][slot]), slot
-
-
-def _token_strided_args():
- """A fresh argument list whose `x` has token stride 2: the channel-major
- kernel ignores that stride and answers silently wrong."""
- dim, tokens = 512, 64
- spec = dict(dim=dim, lengths=(tokens,), pool_slots=1, slots=(0,), initial=(True,))
- args = _build(spec, seed=15)
- g = torch.Generator(device="cuda").manual_seed(16)
- wide = torch.randn(2 * tokens, dim, generator=g, device="cuda").to(torch.bfloat16)
- args[0] = wide.t().contiguous()[:, ::2]
- assert args[0].stride(1) == 2
- return args
-
-
-def test_guard_refuses_what_the_op_answers_silently_wrong() -> None:
- """First the hazard, unguarded: the op returns without raising and its
- answer misses the reference. Then the guard refuses the same call."""
- ref = op.reference(*_token_strided_args())
- unguarded = _token_strided_args() # strided views must not be cloned dense
- returned = op(*unguarded)
- with pytest.raises(AssertionError):
- _compare((returned, unguarded[3]), ref)
-
- with pytest.raises(AssertionError):
- with validating(op):
- op(*_token_strided_args())
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_reset_recurrent_state_rows.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_reset_recurrent_state_rows.py
deleted file mode 100644
index 13a81044a508..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_reset_recurrent_state_rows.py
+++ /dev/null
@@ -1,117 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the reset_recurrent_state_rows entry."""
-
-import zlib
-from typing import Optional
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.reset_recurrent_state_rows import ( # noqa: E501
- CONV_SHAPE,
- REC_SHAPE,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.reset_recurrent_state_rows import ( # noqa: E501
- reset_recurrent_state_rows as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name
-
-assert torch.cuda.is_available(), "reset_recurrent_state_rows requires a CUDA device"
-
-
-def _compare(
- out: tuple[torch.Tensor, Optional[torch.Tensor]],
- ref: tuple[torch.Tensor, Optional[torch.Tensor]],
-) -> None:
- """Bitwise, both pools whole: every element is its old value or +0.0,
- so a tolerance could only hide a wrong row -- and comparing the whole
- pool is the lock on every slot the batch did not name. A zeroed row
- must be +0.0 by bit pattern; `torch.equal` alone would accept -0.0."""
- for got, want in zip(out, ref):
- assert (got is None) == (want is None)
- if got is None:
- continue
- assert got.dtype == want.dtype and got.shape == want.shape
- assert torch.equal(got, want), "pool differs from the reference"
- zero = want.reshape(want.shape[0], -1).eq(0).all(dim=1)
- if zero.any():
- bits = {1: torch.uint8, 2: torch.int16, 4: torch.int32}[got.element_size()]
- assert (got[zero].contiguous().view(bits) == 0).all(), "a reset row is not +0.0"
-
-
-DEV = "cuda"
-_LAYERS = 3
-
-
-def _build(spec, seed):
- """`__call__`'s arguments, plus the backing allocations a layer view views."""
- g = torch.Generator(device=DEV).manual_seed(seed)
- slots = spec["slots"]
- rec_dtype = spec.get("rec_dtype", torch.float32)
- layer_view = spec.get("layout", "plain") == "layer_view"
-
- def randn(*shape, dtype):
- return torch.randn(*shape, generator=g, device=DEV).to(dtype)
-
- backing = {}
- if layer_view:
- backing["rec"] = randn(_LAYERS, slots, *REC_SHAPE, dtype=rec_dtype)
- backing["conv"] = randn(_LAYERS, slots, *CONV_SHAPE, dtype=torch.bfloat16)
- rec, conv = backing["rec"][1], backing["conv"][1]
- else:
- rec = randn(slots, *REC_SHAPE, dtype=rec_dtype)
- conv = randn(slots, *CONV_SHAPE, dtype=torch.bfloat16)
-
- indices = torch.tensor(spec["indices"], dtype=torch.int32, device=DEV)
- flags = torch.tensor(spec["flags"], dtype=torch.bool, device=DEV)
- return [rec, indices, flags, conv], backing
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name
- certifying(op)
- spec = cell.spec
- args, backing = _build(spec, seed=zlib.crc32(cell.why.encode()))
- rec, indices, flags, conv = args
- before = {name: t.clone() for name, t in backing.items()}
- indices_before, flags_before = indices.clone(), flags.clone()
- ref = op.reference(*args)
-
- out = op(*args)
-
- assert out[0] is rec and out[1] is conv, "returns the pools it was handed"
- _compare(out, ref)
- assert torch.equal(indices, indices_before) and torch.equal(flags, flags_before)
-
- named = {s for s in indices.tolist() if 0 <= s < spec["slots"]}
- reset = {s for s, f in zip(indices.tolist(), flags.tolist()) if not f and s in named}
- if spec["slots"] > 64:
- # Non-vacuity: the engine-width cells reset at least one row, so the
- # whole-pool compare is also the lock on the ~430 unnamed slots.
- assert reset
- if spec.get("layout") == "layer_view":
- for name in backing:
- for other in (0, 2):
- assert torch.equal(backing[name][other], before[name][other]), (name, other)
-
-
-def _pools():
- torch.manual_seed(7)
- rec = torch.randn(6, *REC_SHAPE[1:], dtype=torch.float32, device=DEV)
- conv = torch.randn(6, *CONV_SHAPE, dtype=torch.bfloat16, device=DEV)
- return rec, conv
-
-
-def test_guard_refuses_index_tensors_the_op_misreads() -> None:
- """Unguarded, a stride-2 `state_indices` view is read through its base pointer
- as a flat buffer: indices [0, 1, 2] arrive as [0, 9, 1], so rows 0 and 1 are reset
- and row 2 is not (measured). The guard refuses it."""
- rec, conv = _pools()
- indices = torch.tensor([0, 9, 1, 9, 2, 9], dtype=torch.int32, device=DEV)[::2]
- flags = torch.zeros(3, dtype=torch.bool, device=DEV)
- with pytest.raises(AssertionError, match="state_indices must be a contiguous 1-D tensor"):
- with validating(op):
- op(rec, indices, flags, conv)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py
deleted file mode 100644
index b67a01545790..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py
+++ /dev/null
@@ -1,357 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the target's Gated-DeltaNet verify step in both recurrent modes.
-
-Drives the two targets' own `_gdn_verify` -- the prefill target's (generation rows
-behind a context prefix) and the decode target's -- over consecutive verify steps of
-this checkpoint's GDN geometry, with the cache manager's bookkeeping run in between as
-`update_mamba_states` runs it: the conv window promotion in both modes, then either the
-legacy per-step state promotion, or the cached replay's all-layer commit (partitioned
-batches only) and `_advance_replay_state`.
-
-Three things are checked at every step, with a random partial acceptance (or, in the
-full-acceptance cases, every draft accepted):
-
-* the replay mode's output against the legacy mode's on identical inputs and identical
- accepted histories: the two modes compute one recurrence with different kernels and
- state disciplines, so a mis-bound history buffer, a wrong partition condition or a
- skipped commit drifts off within a step;
-* the prefill and the decode target's replay calls bitwise equal to each other (each on
- its own copy of the state);
-* the conv pool identical across the modes (it is promoted the same way).
-"""
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501
- modeling as M,
-)
-from tensorrt_llm._torch.modules.fla.cache_manager import (
- ReplayStateUpdateMetadata,
- _advance_replay_state,
-)
-from tensorrt_llm._torch.modules.fla.cached_replay import (
- CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE,
- commit_gdn_cached_replay_history_layers,
-)
-from tensorrt_llm._torch.modules.mamba.mamba2_metadata import _build_replay_work_items_torch
-
-assert torch.cuda.is_available(), "the GDN verify test requires a CUDA device"
-
-DEV = "cuda"
-H, HV, K, V = 16, 48, 128, 128
-KEY_DIM, VALUE_DIM = H * K, HV * V
-CONV_DIM = 2 * KEY_DIM + VALUE_DIM
-CONV_KERNEL = 4
-HIST = 16
-LAYERS = 2
-SLOTS = 40
-STEPS = 7
-
-
-def _steps(t):
- """Enough steps that every slot fills its history at least once: each step accepts
- at least one token, and a slot wraps once `pnat + t > HIST`."""
- return HIST - t + 3
-
-
-BF16_ULP = 2.0**-8
-#: Replay against legacy, per token block: the replay kernel rounds the history's update
-#: vectors and keys to bf16 and rebuilds the state from them, where the legacy kernel
-#: carries an fp32 state, so the two agree to a few bf16 ulps of the block.
-ELEMENT_GATE = 6 * BF16_ULP
-RMS_GATE = 1.5 * BF16_ULP
-
-
-class _Mode:
- """One cache manager's worth of state for `LAYERS` GDN layers, in one mode, plus the
- `core` / `L` stand-ins `_gdn_verify` reads."""
-
- def __init__(self, replay, t, conv0, ssm0, weights, gating, max_batch):
- self.conv = [c.clone() for c in conv0]
- self.int_conv = [
- torch.zeros(max_batch, t, CONV_DIM, CONV_KERNEL - 1, dtype=torch.bfloat16, device=DEV)
- for _ in range(LAYERS)
- ]
- # V2 lays the layers out inside each slot; the all-layer commit is handed that
- # affine layout.
- self.backing = torch.stack(ssm0, 1).contiguous()
- self.ssm = [self.backing[:, i] for i in range(LAYERS)]
- bound = dict(
- conv=dict(enumerate(self.conv)),
- int_conv=dict(enumerate(self.int_conv)),
- int_ssm=None,
- replay=None,
- )
- if replay:
- self.old_x = torch.zeros(
- LAYERS, SLOTS, 2, HIST, HV, V, dtype=torch.bfloat16, device=DEV
- )
- self.old_B = torch.zeros(LAYERS, SLOTS, 2, HIST, H, K, dtype=torch.bfloat16, device=DEV)
- self.old_dt = torch.zeros(LAYERS, SLOTS, 2, HV, HIST, device=DEV)
- self.old_dA = torch.zeros(LAYERS, SLOTS, 2, HV, HIST, device=DEV)
- self.pnat = torch.zeros(SLOTS, dtype=torch.int32, device=DEV)
- self.buf = torch.zeros(SLOTS, dtype=torch.int32, device=DEV)
- bound["replay"] = M._Replay(
- old={
- i: (self.old_x[i], self.old_B[i], self.old_dt[i], self.old_dA[i])
- for i in range(LAYERS)
- },
- cache_buf_idx=self.buf,
- prev_num_accepted=self.pnat,
- history=HIST,
- width=t,
- all_layer_commit=True,
- # The Triton replay and the manager's all-layer commit, as an fp32 pool runs.
- cute=False,
- )
- self.meta = ReplayStateUpdateMetadata(self.pnat, self.buf, t, HIST)
- self.mm = SimpleNamespace(
- replay_work_items=torch.zeros(max_batch, 4, dtype=torch.int32, device=DEV),
- replay_n_writes=torch.zeros(1, dtype=torch.int32, device=DEV),
- )
- else:
- self.int_ssm = [torch.zeros(max_batch, t, HV, V, K, device=DEV) for _ in range(LAYERS)]
- bound["int_ssm"] = dict(enumerate(self.int_ssm))
- self.core = SimpleNamespace(
- conv_dim=CONV_DIM,
- key_dim=KEY_DIM,
- value_dim=VALUE_DIM,
- num_k_heads=H,
- head_k_dim=K,
- num_v_heads=HV,
- head_v_dim=V,
- _bound=SimpleNamespace(**bound),
- )
- self.layers = [
- SimpleNamespace(conv_w=weights[i], a_log=gating[i][0], dt_bias=gating[i][1])
- for i in range(LAYERS)
- ]
-
- def work(self, slots, n, t):
- """What the forward hands `_gdn_verify` as `replay`: the metadata's prepare builds
- the work items, then `_replay_work` applies the manager's partition condition."""
- partitioned = n >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE
- if partitioned:
- _build_replay_work_items_torch(
- slots,
- self.pnat,
- self.buf,
- self.mm.replay_work_items,
- self.mm.replay_n_writes,
- t,
- HIST,
- )
- replay = M._replay_work(self.core, self.mm, n)
- assert replay[2] == partitioned
- return replay
-
- def verify(self, target_cls, layer, mq, a, b, slots, int_idx, n, t, replay):
- return target_cls._gdn_verify(
- SimpleNamespace(core=self.core),
- mq,
- a,
- b,
- self.ssm[layer],
- layer,
- self.layers[layer],
- slots,
- int_idx,
- n,
- t,
- replay,
- )
-
- def commit(self, slots, n):
- """The manager's all-layer commit, then the replay bookkeeping advance is the
- caller's (it needs the accepted counts)."""
- commit_gdn_cached_replay_history_layers(
- ssm_states=self.ssm[0],
- ssm_state_num_layers=LAYERS,
- ssm_state_layer_stride=self.backing.stride(1),
- ssm_state_slot_stride=self.backing.stride(0),
- old_u=self.old_x,
- old_k=self.old_B,
- old_G=self.old_dt,
- replay_work_items=self.mm.replay_work_items[:n],
- n_writes=self.mm.replay_n_writes,
- history_size=HIST,
- )
-
-
-def _rows(gen, rows):
- """One layer's verify-row activations: the pre-conv `[q|k|v]` projection and the raw
- `a` and `b` projections, each its own contiguous GEMM output."""
- return tuple(
- torch.randn(rows, width, generator=gen, device=DEV).to(torch.bfloat16)
- for width in (CONV_DIM, HV, HV)
- )
-
-
-def _behind_prefix(x, prefix=5):
- """`x` as the trailing rows of a buffer that starts with `prefix` context rows."""
- buf = torch.zeros(prefix + x.shape[0], x.shape[1], dtype=x.dtype, device=DEV)
- buf[prefix:] = x
- return torch.split(buf, [prefix, x.shape[0]], 0)[1]
-
-
-def _promote(pool, inter, slots, accepted):
- """`update_mamba_states`' copy: step `accepted - 1` of each batch row's per-step
- buffer into the row's slot."""
- rows = torch.arange(slots.shape[0], device=DEV)
- pool[slots.long()] = inter[rows, accepted.long() - 1]
-
-
-def _gate(got, want, what):
- block = got.shape[-2] * got.shape[-1]
- g, w = got.float().reshape(-1, block), want.float().reshape(-1, block)
- scale = w.abs().amax(1, keepdim=True).clamp_min(1e-30)
- worst = ((g - w).abs() / scale).max().item()
- rms = ((g - w).pow(2).mean().sqrt() / w.pow(2).mean().sqrt()).item()
- assert worst <= ELEMENT_GATE and rms <= RMS_GATE, (
- f"{what}: replay vs legacy {worst / BF16_ULP:.2f} ulps of the block, "
- f"RMS {rms / BF16_ULP:.2f} ulps"
- )
-
-
-def _state(gen, cpu, n, t, gating=None):
- """Random weights, gating parameters and pre-step pools shared by the modes, and a
- random slot assignment of `n` requests."""
- slots = torch.randperm(SLOTS, generator=cpu)[:n].to(DEV, torch.int32)
- weights = [
- (0.3 * torch.randn(CONV_DIM, CONV_KERNEL, generator=gen, device=DEV)).to(torch.bfloat16)
- for _ in range(LAYERS)
- ]
- if gating is None:
- gating = [
- (
- (torch.randn(HV, generator=gen, device=DEV) * 1.5 + 1.0).clamp(-7.0, 5.0),
- (torch.randn(HV, generator=gen, device=DEV) * 1.7 - 4.4).clamp(-11.0, 2.5),
- )
- for _ in range(LAYERS)
- ]
- conv0 = [
- torch.randn(SLOTS, CONV_DIM, CONV_KERNEL - 1, generator=gen, device=DEV).to(torch.bfloat16)
- for _ in range(LAYERS)
- ]
- ssm0 = [0.3 * torch.randn(SLOTS, HV, V, K, generator=gen, device=DEV) for _ in range(LAYERS)]
- return slots, weights, gating, conv0, ssm0
-
-
-@pytest.mark.parametrize(
- "n,t,accept",
- [
- # At the partition threshold, at max_draft_len 7 (T = 8, the replay kernel's
- # ceiling): the history wraps through the all-layer commit.
- (CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, 8, "random"),
- # Below the threshold, every draft accepted: from a fresh slot at T = 8 the second
- # step appends exactly up to the history's end (`pnat + T == 16`) and every second
- # step after it wraps in-kernel.
- (4, 8, "full"),
- ],
- ids=lambda v: str(v),
-)
-def test_replay_verify_tracks_legacy_in_both_targets(n, t, accept):
- gen = torch.Generator(device=DEV).manual_seed(1000 * n + t)
- cpu = torch.Generator().manual_seed(1000 * n + t)
- slots, weights, gating, conv0, ssm0 = _state(gen, cpu, n, t)
- int_idx = torch.arange(n, dtype=torch.int32, device=DEV)
- legacy = _Mode(False, t, conv0, ssm0, weights, gating, n)
- prefill = _Mode(True, t, conv0, ssm0, weights, gating, n)
- decode = _Mode(True, t, conv0, ssm0, weights, gating, n)
- partitioned = n >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE
-
- wraps = to_the_end = 0
- for step in range(_steps(t)):
- pending = prefill.pnat[slots.long()]
- wraps += int((pending + t > HIST).sum())
- to_the_end += int((pending + t == HIST).sum())
- replays = {id(mode): mode.work(slots, n, t) for mode in (prefill, decode)}
- for layer in range(LAYERS):
- mq, a, b = _rows(gen, n * t)
- # The prefill target's generation rows sit behind a context prefix: row views
- # at an offset into the batch's projections.
- mq_p, a_p, b_p = (_behind_prefix(x) for x in (mq, a, b))
- want = legacy.verify(M.DecodeTarget, layer, mq, a, b, slots, int_idx, n, t, None)
- got_p = prefill.verify(
- M.PrefillTarget, layer, mq_p, a_p, b_p, slots, int_idx, n, t, replays[id(prefill)]
- )
- got_d = decode.verify(
- M.DecodeTarget, layer, mq, a, b, slots, int_idx, n, t, replays[id(decode)]
- )
- assert got_p.shape == got_d.shape == want.shape == (1, n * t, HV, V)
- assert torch.equal(got_p, got_d), (
- f"step {step} layer {layer}: the prefill and decode targets' replay calls disagree"
- )
- _gate(got_d, want, f"step {step} layer {layer}")
- if accept == "full":
- accepted = torch.full((n,), t, dtype=torch.int32, device=DEV)
- else:
- accepted = torch.randint(1, t + 1, (n,), generator=cpu, dtype=torch.int32).to(DEV)
- for layer in range(LAYERS):
- for mode in (legacy, prefill, decode):
- _promote(mode.conv[layer], mode.int_conv[layer], slots, accepted)
- _promote(legacy.ssm[layer], legacy.int_ssm[layer], slots, accepted)
- for mode in (prefill, decode):
- if partitioned:
- mode.commit(slots, n)
- _advance_replay_state(mode.meta, slots, accepted)
- for layer in range(LAYERS):
- assert torch.equal(prefill.conv[layer], legacy.conv[layer])
- assert torch.equal(decode.conv[layer], legacy.conv[layer])
- assert wraps > 0, "the steps never filled a history, so the checkpoint never advanced"
- if accept == "full":
- assert wraps >= 3 * n and to_the_end >= 3 * n, (
- f"full acceptance at T={t} should wrap every row at least 3 times and append up to "
- f"the history's end as often ({wraps}, {to_the_end})"
- )
-
-
-def test_a_wrong_partition_drifts_off_legacy():
- """The control: the work items withheld from a partitioned step while the manager
- still runs its all-layer commit -- the checkpoint committed twice -- must fail the
- same gate."""
- n, t = CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, 4
- gen = torch.Generator(device=DEV).manual_seed(7)
- cpu = torch.Generator().manual_seed(7)
- gating = [
- (torch.full((HV,), 0.5, device=DEV), torch.full((HV,), -4.0, device=DEV))
- for _ in range(LAYERS)
- ]
- slots, weights, gating, conv0, ssm0 = _state(gen, cpu, n, t, gating)
- int_idx = torch.arange(n, dtype=torch.int32, device=DEV)
- legacy = _Mode(False, t, conv0, ssm0, weights, gating, n)
- wrong = _Mode(True, t, conv0, ssm0, weights, gating, n)
- failed = False
- for _ in range(STEPS):
- # The manager's prepare builds the work items; the step does not receive them.
- _build_replay_work_items_torch(
- slots,
- wrong.pnat,
- wrong.buf,
- wrong.mm.replay_work_items,
- wrong.mm.replay_n_writes,
- t,
- HIST,
- )
- for layer in range(LAYERS):
- mq, a, b = _rows(gen, n * t)
- want = legacy.verify(M.DecodeTarget, layer, mq, a, b, slots, int_idx, n, t, None)
- got = wrong.verify(
- M.DecodeTarget, layer, mq, a, b, slots, int_idx, n, t, (None, None, False)
- )
- try:
- _gate(got, want, "control")
- except AssertionError:
- failed = True
- accepted = torch.randint(1, t + 1, (n,), generator=cpu, dtype=torch.int32).to(DEV)
- for layer in range(LAYERS):
- _promote(legacy.conv[layer], legacy.int_conv[layer], slots, accepted)
- _promote(wrong.conv[layer], wrong.int_conv[layer], slots, accepted)
- _promote(legacy.ssm[layer], legacy.int_ssm[layer], slots, accepted)
- wrong.commit(slots, n)
- _advance_replay_state(wrong.meta, slots, accepted)
- assert failed, "a doubled checkpoint commit passed the replay-vs-legacy gate"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py
deleted file mode 100644
index 548008dacdf4..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py
+++ /dev/null
@@ -1,168 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU tests for the target's MTP IndexShare plumbing.
-
-The target's canonical selection, the capture of each request's draft-extend row
-into the shared buffers (decode and mixed-batch forms), and the reuse entry's
-order preservation over canonical captures. The targets are built without weights:
-the capture methods read only their arguments, the metadata and the core's
-draft binding.
-"""
-
-from types import SimpleNamespace
-
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501
- modeling as m,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_canonical_selection import ( # noqa: E501
- qsa_canonical_selection as canonical,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_mtp_reuse_topk import ( # noqa: E501
- qsa_mtp_reuse_topk as reuse,
-)
-
-assert torch.cuda.is_available(), "the IndexShare plumbing runs on a CUDA device"
-
-DEV = "cuda"
-K = 512
-SLOT = 1
-
-
-def _target(cls):
- target = object.__new__(cls)
- target.core = SimpleNamespace(_draft_bound=SimpleNamespace(share_slot=SLOT))
- return target
-
-
-def _metadata(batch: int, accepted, seq_lens=None, num_contexts: int = 0):
- return SimpleNamespace(
- num_seqs=batch,
- num_contexts=num_contexts,
- seq_lens_cuda=None
- if seq_lens is None
- else torch.tensor(seq_lens, dtype=torch.int32, device=DEV),
- qsa_mtp_num_accepted=torch.tensor(accepted, dtype=torch.int32, device=DEV),
- qsa_shared_topk_indices=torch.full((3, 8, K), -7, dtype=torch.int32, device=DEV),
- qsa_shared_topk_visible_blocks=torch.full((3, 8), -7, dtype=torch.int32, device=DEV),
- qsa_shared_topk_captured_slots=set(),
- )
-
-
-COLS = 2080
-
-
-def _scores_and_windows(rows: int, gen: torch.Generator):
- """Relu-shaped scores (half zeros: ties at the budget's last score whenever fewer than
- K are positive) over windows of 0 to 3K groups; the columns past a window are NaN."""
- ends = torch.randint(0, COLS + 1, (rows,), generator=gen).to(torch.int32).to(DEV)
- scores = torch.relu(torch.randn(rows, COLS, generator=gen)).to(DEV)
- scores = scores * (torch.rand(rows, COLS, generator=gen).to(DEV) < 0.3)
- outside = torch.arange(COLS, device=DEV)[None, :] >= ends[:, None]
- scores.masked_fill_(outside, float("nan"))
- return scores, torch.zeros(rows, dtype=torch.int32, device=DEV), ends
-
-
-def _select(scores, starts, ends) -> torch.Tensor:
- out = torch.empty((scores.shape[0], K), dtype=torch.int32, device=DEV)
- m._select(SimpleNamespace(block_topk=K), scores, starts, ends, out)
- return out
-
-
-def test_select_is_the_canonical_row_ascending_with_padding_last():
- gen = torch.Generator().manual_seed(3)
- scores, starts, ends = _scores_and_windows(64, gen)
-
- out = _select(scores, starts, ends)
-
- assert out.dtype == torch.int32
- assert torch.equal(out, canonical.reference(scores, starts, ends, out))
- for row, end in zip(out.cpu().tolist(), ends.tolist()):
- ids = [i for i in row if i >= 0]
- assert ids == sorted(ids) and len(set(ids)) == len(ids) == min(end, K)
- assert row == ids + [-1] * (K - len(ids)), "padding must stay last"
-
-
-def test_decode_capture_at_draft_seven_takes_row_s_times_8_plus_a_minus_1():
- """max_draft_len 7: eight rows per request, accepted 1..8, so the captured row is
- `s * 8 + a - 1` -- the first and the last row of a block included."""
- p, batch = 8, 4
- accepted = [1, 8, 5, 3]
- topk = torch.arange(batch * p * K, dtype=torch.int32, device=DEV).reshape(batch * p, K)
- visible = torch.arange(300, 300 + batch * p, dtype=torch.int32, device=DEV)
- md = _metadata(batch, accepted)
-
- _target(m.DecodeTarget)._capture_selection(md, topk, visible, batch, p)
-
- rows = [s * p + a - 1 for s, a in enumerate(accepted)]
- assert rows == [0, 15, 20, 26]
- torch.testing.assert_close(md.qsa_shared_topk_indices[SLOT, :batch], topk[rows])
- torch.testing.assert_close(md.qsa_shared_topk_visible_blocks[SLOT, :batch], visible[rows])
- # Other slots and rows past the batch are untouched.
- assert torch.all(md.qsa_shared_topk_indices[SLOT, batch:] == -7)
- assert torch.all(md.qsa_shared_topk_indices[0] == -7)
- assert md.qsa_shared_topk_captured_slots == {SLOT}
-
-
-def test_prefill_capture_at_draft_seven():
- """A mixed batch at max_draft_len 7: two context chunks of 5 and 3 rows, then three
- generation requests of 8 rows with 8, 1 and 6 accepted. A context request captures its
- chunk's last row whatever accepted count the worker left there."""
- seq_lens = [5, 3, 8, 8, 8]
- accepted = [2, 3, 8, 1, 6]
- rows_total = sum(seq_lens)
- topk = torch.arange(rows_total * K, dtype=torch.int32, device=DEV).reshape(rows_total, K)
- visible = torch.arange(400, 400 + rows_total, dtype=torch.int32, device=DEV)
- md = _metadata(5, accepted, seq_lens=seq_lens, num_contexts=2)
-
- _target(m.PrefillTarget)._capture_selection(md, topk, visible)
-
- rows = [4, 7, 8 + 7, 16 + 0, 24 + 5]
- torch.testing.assert_close(md.qsa_shared_topk_indices[SLOT, :5], topk[rows])
- torch.testing.assert_close(md.qsa_shared_topk_visible_blocks[SLOT, :5], visible[rows])
- assert md.qsa_shared_topk_captured_slots == {SLOT}
-
-
-def _reuse_case(fresh: int):
- """A canonical capture of 16 requests and the reuse call's arguments for it."""
- gen = torch.Generator().manual_seed(11 + fresh)
- batch = 16
- scores, starts, ends = _scores_and_windows(batch, gen)
- captured_visible = ends.cpu()
- captured = _select(scores, starts, ends)
- shared = torch.full((8, 32, K), -1, dtype=torch.int32, device=DEV)
- shared[SLOT, :batch] = captured
- shared_visible = torch.zeros((8, 32), dtype=torch.int32, device=DEV)
- shared_visible[SLOT, :batch] = captured_visible.to(DEV)
- args = dict(
- shared_indices=shared[SLOT],
- shared_visible_blocks=shared_visible[SLOT],
- request_indices=torch.arange(batch, dtype=torch.int32, device=DEV),
- visible_blocks=(captured_visible + fresh).to(DEV),
- output=torch.empty((batch, K), dtype=torch.int32, device=DEV),
- )
- return args, captured_visible
-
-
-def test_reuse_of_a_canonical_capture_stays_ascending():
- """A canonical capture plus the groups completed since is still ascending: appended ids
- are newer than every captured one, and a full row gives up its oldest ids to them, so
- the newest captured groups -- the ones next to the new ones -- are kept."""
- for fresh in (0, 2):
- args, captured_visible = _reuse_case(fresh)
-
- out = reuse(**args)
-
- assert torch.equal(out, reuse.reference(**args))
- for row, shared_row, vis in zip(
- out.cpu().tolist(), args["shared_indices"].cpu().tolist(), captured_visible.tolist()
- ):
- ids = [i for i in row if i >= 0]
- assert row == ids + [-1] * (K - len(ids)), "padding must stay last"
- assert ids == sorted(ids) and len(set(ids)) == len(ids)
- assert len(ids) == min(vis + fresh, K)
- assert all(i < vis + fresh for i in ids)
- captured = [i for i in shared_row if i >= 0]
- kept = captured[max(0, len(captured) + fresh - K) :]
- assert ids == kept + list(range(vis, vis + fresh))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_kv_range_probe.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_kv_range_probe.py
deleted file mode 100644
index 2603aaf63e00..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_kv_range_probe.py
+++ /dev/null
@@ -1,82 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""The FP8 KV range probe (`TRTLLM_MODELING_V2_VALIDATE` over an fp8 pool).
-
-The pool stores K/V at the unit scale, so a magnitude past the e4m3 maximum saturates. The probe
-reports the largest magnitude of the first eager forward's attention layers, once, and never reads
-a tensor under CUDA graph capture.
-"""
-
-from unittest.mock import Mock
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501
- modeling as m,
-)
-
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-LAST = 11
-
-
-def _kv(peak: float = 1.0):
- """K and V of one layer, 5 rows x 2 heads x 8 dims, with `peak` as V's largest value."""
- k = torch.linspace(-1.0, 1.0, 80, device=DEVICE).reshape(5, 2, 8).to(torch.bfloat16)
- v = torch.linspace(-0.5, 0.5, 80, device=DEVICE).reshape(5, 2, 8).to(torch.bfloat16)
- v[3, 1, 4] = peak
- return k, v
-
-
-def _run_forward(probe, peaks):
- """One forward's attention layers 3, 7 and LAST; returns each `observe` result."""
- return [probe.observe(layer, *_kv(peak)) for layer, peak in zip((3, 7, LAST), peaks)]
-
-
-def test_a_forward_within_the_e4m3_range_is_silent(monkeypatch):
- logger = Mock()
- monkeypatch.setattr(m, "logger", logger)
-
- assert _run_forward(m._KvRangeProbe(LAST), (1.0, 300.0, 448.0)) == [True, True, False]
-
- logger.warning.assert_not_called()
-
-
-@pytest.mark.parametrize("peak", [-600.0, float("nan")])
-def test_a_magnitude_past_the_e4m3_maximum_is_reported_once(monkeypatch, peak):
- """A negative value counts by its magnitude; a NaN counts as out of range."""
- logger = Mock()
- monkeypatch.setattr(m, "logger", logger)
-
- # The oversized value is in a middle layer: the report waits for the forward's last one.
- assert _run_forward(m._KvRangeProbe(LAST), (1.0, peak, 1.0)) == [True, True, False]
-
- logger.warning.assert_called_once()
- assert "e4m3 maximum 448" in logger.warning.call_args.args[0]
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
-def test_observing_under_capture_reads_nothing(monkeypatch):
- logger = Mock()
- monkeypatch.setattr(m, "logger", logger)
- probe = m._KvRangeProbe(LAST)
- k, v = _kv(1000.0)
- scratch = torch.zeros(1, device="cuda")
- stream = torch.cuda.Stream()
- stream.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(stream):
- scratch.add_(1) # warm the allocator off the capture stream
- torch.cuda.current_stream().wait_stream(stream)
- graph = torch.cuda.CUDAGraph()
- kept = []
- with torch.cuda.graph(graph):
- scratch.add_(1)
- # The last layer, with an out-of-range value: a read here would fail the capture
- # (a host sync) and a report would fire.
- kept.append(probe.observe(LAST, k, v))
-
- assert kept == [True], "a captured forward does not end the probe"
- logger.warning.assert_not_called()
- # The first eager forward after capture is what reports.
- assert _run_forward(probe, (1.0, 1.0, 1000.0)) == [True, True, False]
- logger.warning.assert_called_once()
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_logits_widening.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_logits_widening.py
deleted file mode 100644
index 62dc201db2db..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_logits_widening.py
+++ /dev/null
@@ -1,100 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""The target's logits processor returns the shell's logits bit for bit: the same row gather,
-the same fp32 values, its head's bf16 logits widened by one vectorized kernel."""
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501
- modeling as T,
-)
-from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor
-
-assert torch.cuda.is_available(), "the logits processor test requires a CUDA device"
-
-VOCAB, HIDDEN = 248320, 2560
-
-
-@pytest.fixture(scope="module")
-def head():
- g = torch.Generator(device="cuda").manual_seed(0)
- lin = torch.nn.Linear(HIDDEN, VOCAB, bias=False, device="cuda", dtype=torch.bfloat16)
- with torch.no_grad():
- lin.weight.copy_(torch.randn(VOCAB, HIDDEN, generator=g, device="cuda") * 0.02)
- return lin
-
-
-@pytest.mark.parametrize("seq_lens", [[1] * 128, [5000, 3000, 1] + [1] * 121])
-def test_matches_the_shell_processor_bitwise(head, seq_lens) -> None:
- g = torch.Generator(device="cuda").manual_seed(len(seq_lens))
- lens = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")
- hidden = torch.randn(int(lens.sum()), HIDDEN, generator=g, device="cuda").to(torch.bfloat16)
- md = SimpleNamespace(seq_lens_cuda=lens)
- with torch.no_grad():
- want = LogitsProcessor().forward(hidden, head, md)
- got = T._WideningLogitsProcessor().forward(hidden, head, md)
- assert got.dtype == torch.float32 and got.shape == (len(seq_lens), VOCAB)
- assert torch.equal(got.view(torch.int32), want.view(torch.int32))
-
-
-def test_context_logits_and_no_metadata(head) -> None:
- g = torch.Generator(device="cuda").manual_seed(9)
- hidden = torch.randn(33, HIDDEN, generator=g, device="cuda").to(torch.bfloat16)
- with torch.no_grad():
- for md, ctx in ((None, False), (SimpleNamespace(seq_lens_cuda=None), True)):
- want = LogitsProcessor().forward(hidden, head, md, ctx)
- got = T._WideningLogitsProcessor().forward(hidden, head, md, ctx)
- assert torch.equal(got.view(torch.int32), want.view(torch.int32))
-
-
-@pytest.mark.parametrize("rows", [1, 4, 256])
-def test_unwidened_logits_give_the_same_greedy_tokens(head, rows) -> None:
- """`widen=False` returns the head's own bf16 logits, whose greedy samplers (the verify
- sampler's `greedy_argmax`, the draft sampler's int32 one) pick the widened logits' tokens."""
- from tensorrt_llm._torch.speculative.greedy_argmax import greedy_argmax
-
- g = torch.Generator(device="cuda").manual_seed(rows)
- hidden = torch.randn(rows, HIDDEN, generator=g, device="cuda").to(torch.bfloat16)
- proc = T._WideningLogitsProcessor()
- with torch.no_grad():
- wide = proc.forward(hidden, head, None, True)
- narrow = proc.forward(hidden, head, None, True, widen=False)
- # Without return_context_logits the shell's gather and fp32 contract stay.
- gathered = proc.forward(hidden, head, None, False, widen=False)
- assert narrow.dtype == torch.bfloat16 and wide.dtype == torch.float32
- assert gathered.dtype == torch.float32
- assert torch.equal(narrow.float().view(torch.int32), wide.view(torch.int32))
- narrow[0, 17] = narrow[0].max() # an exact tie, broken to the lower index either way
- wide = narrow.float()
- assert torch.equal(greedy_argmax(narrow), greedy_argmax(wide))
- assert torch.equal(greedy_argmax(narrow, torch.int32), greedy_argmax(wide, torch.int32))
- assert torch.equal(greedy_argmax(narrow), torch.argmax(wide, dim=-1))
-
-
-def _gate(guided=None, greedy=True, rejection=False, penalty=False, relaxed=False):
- worker = SimpleNamespace(
- guided_decoder=guided,
- spec_config=SimpleNamespace(use_relaxed_acceptance_for_thinking=relaxed),
- )
- md = SimpleNamespace(
- is_all_greedy_sample=greedy,
- wants_advanced_draft_sampling=rejection and not greedy,
- enable_penalty=penalty,
- )
- return T._spec_argmax_only(worker, md)
-
-
-def test_argmax_only_gate() -> None:
- assert _gate() == (True, True)
- # A non-greedy batch: the verify sampler samples; the draft stays argmax unless
- # rejection sampling needs its distribution.
- assert _gate(greedy=False) == (False, True)
- assert _gate(greedy=False, rejection=True) == (False, False)
- # The occurrence penalties and relaxed acceptance read the verify logits only.
- assert _gate(penalty=True) == (False, True)
- assert _gate(relaxed=True) == (False, True)
- # A guided decoder edits both.
- assert _gate(guided=object()) == (False, False)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_nvfp4_context_groups.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_nvfp4_context_groups.py
deleted file mode 100644
index 847d5d8abc5e..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_nvfp4_context_groups.py
+++ /dev/null
@@ -1,136 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test: the target's context pass over an NVFP4 pool, a group of requests at a time.
-
-The context rows read a bf16 scratch of their requests' pages (`_nvfp4_context_plan`), filled
-and attended one group of consecutive requests at a time, so it holds the largest group's
-pages rather than the step's. Under a budget of a few pages, with one request longer than
-the budget, the rows equal, bitwise, one group over the whole step -- for a prefill step's
-rows and for a mixed step's one kept row a request. The prefill requests' rows come in
-multiples of the row-tile pass's 10-row groups, so a launch per group splits none of them and
-every row takes the same pass as in the one launch (a split group's rows go to the per-row
-pass, which rounds differently).
-"""
-
-import math
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501
- modeling as M,
-)
-
-assert torch.cuda.is_available(), "the NVFP4 context pass requires a CUDA device"
-
-DEV = "cuda"
-TPB, HKV, D, QH, RATIO, TOPK = 32, 2, 256, 24, 4, 2048
-# (cached, new) tokens a context request: 10, 30, 1, 47 and 4 pages.
-SEGMENTS = ((0, 300), (900, 40), (20, 10), (1400, 100), (64, 60))
-BUDGET_PAGES = 32 # request 3's 47 pages form a group of their own
-
-
-def _step(one_row_per_request: bool, seed: int = 5):
- gen = torch.Generator(device=DEV).manual_seed(seed)
- kv_lens = [s + n for s, n in SEGMENTS]
- pages = [math.ceil(k / TPB) for k in kv_lens]
- pool_pages = sum(pages) + 3
- shape = (pool_pages, 2, HKV, TPB)
- data = torch.randint(0, 256, (*shape, D // 2), generator=gen, device=DEV).to(torch.uint8)
- scales = torch.randint(0x08, 0x31, (*shape, D // 16), generator=gen, device=DEV)
- scales = scales.to(torch.uint8)
- slots = 2 * len(SEGMENTS)
- slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))
- slot_of = slot_of[: len(SEGMENTS)].to(DEV, torch.int32)
- page_ids = torch.randperm(pool_pages, generator=gen, device=DEV).to(torch.int32)
- block_table = torch.full((slots, max(pages)), -1, dtype=torch.int32, device=DEV)
- cursor = 0
- for slot, count in zip(slot_of.tolist(), pages):
- block_table[slot, :count] = page_ids[cursor : cursor + count]
- cursor += count
- new = torch.tensor([n for _, n in SEGMENTS], dtype=torch.int32)
- md = SimpleNamespace(
- num_contexts=len(SEGMENTS),
- kv_lens_runtime=torch.tensor(kv_lens, dtype=torch.int32),
- kv_lens_cuda_runtime=torch.tensor(kv_lens, dtype=torch.int32, device=DEV),
- seq_lens=new,
- seq_lens_cuda=new.to(DEV),
- )
- # The step's rows: a request's new tokens, or the last of them alone.
- kept = [1 if one_row_per_request else n for _, n in SEGMENTS]
- logical = torch.cat(
- [torch.arange(s + n - k, s + n, device=DEV) for (s, n), k in zip(SEGMENTS, kept)]
- )
- rows = logical.numel()
- # Every visible token, as the expansion lays out a selection under its top-k.
- column = torch.arange(TOPK + RATIO - 1, device=DEV)[None, :]
- selected = torch.where(column <= logical[:, None], column, -1).to(torch.int32)
- core = SimpleNamespace(
- heads_kv=HKV,
- head_dim=D,
- q_width=QH * D,
- kv_scale_quant_orig=torch.ones(3, dtype=torch.float32, device=DEV),
- softmax_scale=D**-0.5,
- compress_ratio=RATIO,
- )
- kv = SimpleNamespace(k=data[:, 0], v=data[:, 1], k_scale=scales[:, 0], v_scale=scales[:, 1])
- return dict(
- core=core,
- kv=kv,
- md=md,
- block_table=block_table,
- req_all=slot_of.repeat_interleave(new.to(DEV).long()),
- req_rows=slot_of.repeat_interleave(torch.tensor(kept, device=DEV)),
- q=torch.randn(rows, QH, D, generator=gen, device=DEV).to(torch.bfloat16),
- selected=selected,
- logical=logical.to(torch.int32),
- rows=rows,
- )
-
-
-def _context_pass(step, one_row_per_request):
- plan = M._nvfp4_context_plan(
- step["core"],
- step["md"],
- step["req_all"],
- step["block_table"],
- TPB,
- one_row_per_request=one_row_per_request,
- )
- out = torch.empty(step["rows"], QH * D, dtype=torch.bfloat16, device=DEV)
- M._sparse_gqa_nvfp4_context(
- step["core"],
- step["kv"],
- plan,
- step["q"],
- step["block_table"],
- step["selected"],
- step["req_rows"],
- step["logical"],
- TPB,
- step["rows"],
- out,
- )
- return plan, out
-
-
-@pytest.mark.parametrize("one_row_per_request", [False, True], ids=["prefill", "mixed_kept"])
-def test_grouped_context_pass_equals_one_group(one_row_per_request, monkeypatch) -> None:
- step = _step(one_row_per_request)
- whole, want = _context_pass(step, one_row_per_request)
- assert len(whole[0]) == 1 and whole[2].shape[0] == 92
-
- page_bytes = 2 * HKV * TPB * D * 2
- monkeypatch.setattr(M, "_NVFP4_CONTEXT_SCRATCH_BYTES", BUDGET_PAGES * page_bytes)
- (groups, _, k_scratch, _), got = _context_pass(step, one_row_per_request)
- ends = [1, 2, 3, 4, 5] if one_row_per_request else [300, 340, 350, 450, 510]
- # Requests [0], [1, 2], [3] (over the budget, alone), [4].
- assert [(g[0], g[1]) for g in groups] == [
- (0, ends[0]),
- (ends[0], ends[2]),
- (ends[2], ends[3]),
- (ends[3], ends[4]),
- ]
- assert k_scratch.shape[0] == 47
- assert torch.equal(got.view(torch.int16), want.view(torch.int16))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_patch_embed_conv3d.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_patch_embed_conv3d.py
deleted file mode 100644
index 21bdba2b0d0c..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_patch_embed_conv3d.py
+++ /dev/null
@@ -1,118 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the qwen4_exp patch_embed_conv3d entry.
-
-The reference is a true-fp32 matmul of the flattened patch against the
-flattened filter: a different op, accumulation order and memory layout from
-the convolution under test, so it checks the contraction pairing rather than
-restating it. Beyond the cells: the one-patch-per-row guard and graph capture.
-"""
-
-import zlib
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.vision.patch_embed_conv3d import ( # noqa: E501
- patch_embed_conv3d as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import ULP, certifying, validating # noqa: E402,F401
-
-assert torch.cuda.is_available(), "patch_embed_conv3d requires a CUDA device"
-
-
-def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
- """Default rtol; atol raised to 4 ulps of the dtype at the result's RMS.
-
- Summing 1536 signed products leaves a few outputs per thousand near
- zero, where the relative gate measures nothing and torch's 1e-5 atol is
- three orders below one bf16 ulp. Measured at the checkpoint's geometry:
- max |conv - fp32 reference| was 4.31 / 4.38 / 7.81 ulps at the RMS for
- 64 / 256 / 4096 bf16 patches (1-2 ulps of the element's own magnitude,
- which the untouched rtol covers); fp16 4.36 / 4.39 / 8.15. Wrong
- variants land 325-392x outside this gate (bias dropped: 10.9x, bounded
- by `max|bias|` ~ 0.02). For fp32 the 1e-5 floor binds, so nothing is
- loosened there.
- """
- assert out.dtype == ref.dtype, (out.dtype, ref.dtype)
- assert out.shape == ref.shape, (out.shape, ref.shape)
- rms = ref.float().pow(2).mean().sqrt().item()
- rtol = {torch.bfloat16: 1.6e-2, torch.float16: 1e-3, torch.float32: 1.3e-6}[ref.dtype]
- atol = max(1e-5, 4.0 * rms * ULP[ref.dtype])
- torch.testing.assert_close(out, ref, rtol=rtol, atol=atol)
-
-
-IN_CHANNELS = 3
-KT, KH, KW = 2, 16, 16
-EMBED_DIM = 1152
-PATCH_NUMEL = IN_CHANNELS * KT * KH * KW # 1536
-
-
-def _operands(
- num_patches,
- dtype=torch.bfloat16,
- seed=0,
- in_channels=IN_CHANNELS,
- kernel=(KT, KH, KW),
- embed_dim=EMBED_DIM,
- bias=True,
-):
- """Pixels in [-1, 1] (processor mean/std 0.5), weights at the tower's
- initializer_range 0.02."""
- g = torch.Generator(device="cuda").manual_seed(seed)
- kt, kh, kw = kernel
- pixel_values = torch.empty(num_patches, in_channels * kt * kh * kw, device="cuda")
- pixel_values = pixel_values.uniform_(-1.0, 1.0, generator=g).to(dtype)
- weight = (
- 0.02 * torch.randn(embed_dim, in_channels, kt, kh, kw, generator=g, device="cuda")
- ).to(dtype)
- bias_t = (0.02 * torch.randn(embed_dim, generator=g, device="cuda")).to(dtype) if bias else None
- return pixel_values, weight, bias_t
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying): # noqa: F811
- certifying(op)
- spec = cell.spec
- pixel_values, weight, bias = _operands(
- spec["num_patches"], dtype=spec["dtype"], seed=zlib.crc32(cell.why.encode())
- )
- before = pixel_values.clone()
- out = op(pixel_values, weight, bias)
- assert out.shape == (pixel_values.shape[0], weight.shape[0])
- assert out.dtype == pixel_values.dtype and out.is_contiguous()
- assert out.data_ptr() != pixel_values.data_ptr()
- assert torch.equal(pixel_values, before), "the pixels are read-only"
- _compare(out, op.reference(pixel_values, weight, bias))
-
-
-def test_guard_refuses_rows_that_are_not_one_patch() -> None:
- """Two patches per row are split into twice the rows, without an error."""
- pixel_values, weight, bias = _operands(8, seed=8)
- doubled = pixel_values.reshape(4, 2 * PATCH_NUMEL)
- assert op(doubled, weight, bias).shape == (8, EMBED_DIM)
- with pytest.raises(AssertionError, match="each row must be exactly one patch"):
- with validating(op):
- op(doubled, weight, bias)
-
-
-def test_cuda_graph_capture_and_replay() -> None:
- pixel_values, weight, bias = _operands(64, seed=12)
-
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(warmup):
- for _ in range(3):
- op(pixel_values, weight, bias)
- torch.cuda.current_stream().wait_stream(warmup)
-
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- captured = op(pixel_values, weight, bias)
-
- pixel_values.copy_(_operands(64, seed=13)[0])
- graph.replay()
- torch.cuda.synchronize()
- _compare(captured, op.reference(pixel_values, weight, bias))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_pos_embed_interpolate.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_pos_embed_interpolate.py
deleted file mode 100644
index 3517e2f3d18e..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_pos_embed_interpolate.py
+++ /dev/null
@@ -1,110 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the qwen4_exp pos_embed_interpolate entry.
-
-`op.reference` is the explicit four-corner sum in fp64 -- exact, but sharing
-the entry's sample positions, so it gates numerics. `_ref_resize` here runs
-torch's own `F.interpolate` in fp64 and shares no arithmetic with the entry, so
-it gates the structure: which axis is which, how the table is read, and the
-spatial-merge row order. It sits up to 8.6e-6 from the exact value (its sample
-coordinates are float): far inside a bf16 cell's band, though not an fp32 one's.
-Beyond the cells: the square-table guard and graph capture.
-"""
-
-import zlib
-
-import pytest
-import torch
-import torch.nn.functional as F
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.vision.pos_embed_interpolate import ( # noqa: E501
- pos_embed_interpolate as op,
-)
-
-__extra_import_path__ = ["../../../.."]
-from _validating import certifying, validating # noqa: E402,F401
-
-assert torch.cuda.is_available(), "pos_embed_interpolate requires a CUDA device"
-
-NUM_GRID = 48
-NUM_POS = NUM_GRID * NUM_GRID
-HIDDEN = 1152
-MERGE = 2
-
-
-def _table(rows=NUM_POS, hidden=HIDDEN, dtype=torch.bfloat16, seed=0):
- g = torch.Generator(device="cuda").manual_seed(seed)
- return torch.randn(rows, hidden, generator=g, device="cuda").to(dtype)
-
-
-def _merge_order(x, height, width, merge):
- hidden = x.shape[-1]
- return (
- x.reshape(height // merge, merge, width // merge, merge, hidden)
- .permute(0, 2, 1, 3, 4)
- .reshape(height * width, hidden)
- )
-
-
-def _ref_resize(table, grid_thw, merge):
- """Structural reference via torch's resize kernel, in fp64."""
- num_grid = int(round(table.shape[0] ** 0.5))
- hidden = table.shape[1]
- src = table.double().view(num_grid, num_grid, hidden).permute(2, 0, 1).unsqueeze(0).contiguous()
- outs = []
- for frames, height, width in grid_thw:
- resized = F.interpolate(src, size=(height, width), mode="bilinear", align_corners=True)
- rows = _merge_order(resized[0].permute(1, 2, 0), height, width, merge)
- outs.append(rows.repeat(frames, 1) if frames > 1 else rows)
- return torch.cat(outs) if len(outs) > 1 else outs[0]
-
-
-@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS])
-def test_certified_cells(cell, certifying): # noqa: F811
- certifying(op)
- spec = cell.spec
- table = _table(spec["rows"], spec["hidden"], spec["dtype"], seed=zlib.crc32(cell.why.encode()))
- grid = [list(g) for g in spec["grid"]]
- before = table.clone()
- out = op(table, grid, spec["merge"])
- assert out.shape == (sum(t * h * w for t, h, w in grid), spec["hidden"])
- assert out.dtype == table.dtype
- assert torch.equal(table, before) and out.data_ptr() != table.data_ptr()
- torch.testing.assert_close(out, op.reference(table, grid, spec["merge"]))
- # The structural gate, from a reference that shares no arithmetic.
- torch.testing.assert_close(out, _ref_resize(table, grid, spec["merge"]).to(out.dtype))
-
-
-def test_guard_refuses_a_non_square_table() -> None:
- """The hazard, unguarded: 2300 rows -> isqrt 47, every index in range, a
- correctly shaped result off the wrong grid."""
- table = _table(rows=2300, hidden=64, seed=10)
- out = op(table, [(1, 4, 4)], MERGE)
- assert out.shape == (16, 64)
- assert torch.equal(out, op(table[: 47 * 47], [(1, 4, 4)], MERGE)), "read as a 47x47 grid"
-
- with pytest.raises(AssertionError, match="must hold a square grid"):
- with validating(op):
- op(table, [(1, 4, 4)], MERGE)
-
-
-def test_cuda_graph_capture_and_replay() -> None:
- """The grid is frozen at capture; the table's contents are not."""
- table = _table(seed=12)
- grid = [(1, 8, 12)]
-
- warmup = torch.cuda.Stream()
- warmup.wait_stream(torch.cuda.current_stream())
- with torch.cuda.stream(warmup):
- for _ in range(3):
- op(table, grid, MERGE)
- torch.cuda.current_stream().wait_stream(warmup)
-
- graph = torch.cuda.CUDAGraph()
- with torch.cuda.graph(graph):
- captured = op(table, grid, MERGE)
-
- table.copy_(_table(seed=13))
- graph.replay()
- torch.cuda.synchronize()
- torch.testing.assert_close(captured, op.reference(table, grid, MERGE))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_vision_tower_calls.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_vision_tower_calls.py
deleted file mode 100644
index d8a1f9396844..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_vision_tower_calls.py
+++ /dev/null
@@ -1,165 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the qwen4_exp vision tower's call grouping.
-
-A request's frames run in calls of whole frames under `_VISION_CALL_MAX_PATCHES` rows.
-The tower here has the checkpoint's geometry (1152 wide, 16 heads of 72, 2x2 merge) with
-two random blocks, so every op the real tower runs runs here; the cap is a module
-constant, lowered so a few hundred rows span several calls. The reference is the same
-tower in one call.
-
-Cutting between frames leaves attention, the position embedding, the rope and every
-row-wise op unchanged; what can differ is which kernel a call picks (a GEMM's split, the
-fused attention kernel for a call whose longest frame is at most 64 rows), within the ops'
-reference bands.
-"""
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501
- modeling as M,
-)
-
-assert torch.cuda.is_available(), "the vision tower requires a CUDA device"
-
-HIDDEN = 1152
-HEADS = 16
-HEAD_DIM = HIDDEN // HEADS
-INTER = 4304
-MERGE = 2
-MERGE_DIM = HIDDEN * MERGE * MERGE
-OUT = 2560
-POS_ROWS = 48 * 48
-PIXEL_WIDTH = 3 * 2 * 16 * 16
-DEPTH = 2
-DT = torch.bfloat16
-
-
-def _core():
- g = torch.Generator(device="cuda").manual_seed(0)
-
- def w(*shape, scale=0.02):
- return (scale * torch.randn(*shape, generator=g, device="cuda")).to(DT)
-
- def norm_w():
- return (1.0 + 0.1 * torch.randn(HIDDEN, generator=g, device="cuda")).to(DT)
-
- weights = {
- "embed": torch.zeros(1, dtype=DT, device="cuda"),
- "vis_patch_w": w(HIDDEN, 3, 2, 16, 16),
- "vis_patch_b": w(HIDDEN),
- "vis_pos": w(POS_ROWS, HIDDEN, scale=0.5),
- "vis_m_n_w": norm_w(),
- "vis_m_n_b": w(HIDDEN),
- "vis_m_fc1_w": w(MERGE_DIM, MERGE_DIM),
- "vis_m_fc1_b": w(MERGE_DIM),
- "vis_m_fc2_w": w(OUT, MERGE_DIM),
- "vis_m_fc2_b": w(OUT),
- }
- for b in range(DEPTH):
- weights.update(
- {
- f"l{b}_vis_n1_w": norm_w(),
- f"l{b}_vis_n1_b": w(HIDDEN),
- f"l{b}_vis_qkv_w": w(3 * HIDDEN, HIDDEN),
- f"l{b}_vis_qkv_b": w(3 * HIDDEN),
- f"l{b}_vis_proj_w": w(HIDDEN, HIDDEN),
- f"l{b}_vis_proj_b": w(HIDDEN),
- f"l{b}_vis_n2_w": norm_w(),
- f"l{b}_vis_n2_b": w(HIDDEN),
- f"l{b}_vis_fc1_w": w(INTER, HIDDEN),
- f"l{b}_vis_fc1_b": w(INTER),
- f"l{b}_vis_fc2_w": w(HIDDEN, INTER),
- f"l{b}_vis_fc2_b": w(HIDDEN),
- }
- )
- return SimpleNamespace(
- w=weights,
- v_depth=DEPTH,
- v_hidden=HIDDEN,
- v_heads=HEADS,
- v_head_dim=HEAD_DIM,
- v_merge=MERGE,
- v_merge_dim=MERGE_DIM,
- v_pixel_width=PIXEL_WIDTH,
- )
-
-
-def _pixels(grid, seed=1):
- patches = sum(t * h * w for t, h, w in grid)
- g = torch.Generator(device="cuda").manual_seed(seed)
- return torch.empty(patches, PIXEL_WIDTH, device="cuda").uniform_(-1, 1, generator=g).to(DT)
-
-
-def _count_calls(tower, monkeypatch):
- """Wrap `tower._run`; the returned list gets one `(rows, grid)` entry per call."""
- calls = []
- run = tower._run
-
- def counted(pixel_values, grid, grid_thw):
- calls.append((pixel_values.shape[0], list(grid)))
- return run(pixel_values, grid, grid_thw)
-
- monkeypatch.setattr(tower, "_run", counted)
- return calls
-
-
-def _assert_close(out, ref):
- """Two runs of the tower over different row counts differ by bf16 rounding: another
- call picks another GEMM kernel, and the difference compounds through the blocks and the
- merger. Measured here: 0.6% RMS relative error and at most 8 ulps (of the RMS) in one
- element, against 117% RMS between unrelated rows. The gates sit between the two."""
- assert out.dtype == ref.dtype and out.shape == ref.shape, (out.shape, ref.shape)
- error = out.float() - ref.float()
- rms = ref.float().pow(2).mean().sqrt()
- assert error.pow(2).mean().sqrt() <= 0.02 * rms, (error.pow(2).mean().sqrt(), rms)
- assert error.abs().max() <= 16 * 2**-8 * rms, (error.abs().max(), rms)
-
-
-@pytest.fixture(scope="module")
-def tower():
- return M.VisionTower(_core())
-
-
-# Frames of 256, 144 and 32 rows under a 512-row cap: three calls -- the first item, the
-# second with two of the third's six frames, the other four -- so a cut between items and one
-# inside an item; the last call's longest frame is 32 rows, where the fused attention kernel
-# is the short-frame one.
-GRID, CAP = [(2, 16, 16), (3, 12, 12), (6, 4, 8)], 512
-
-
-def test_calls_of_whole_frames_equal_one_call(tower, monkeypatch):
- grid, cap = GRID, CAP
- pixels = _pixels(grid)
- grid_thw = torch.tensor(grid)
- monkeypatch.setattr(M, "_VISION_CALL_MAX_PATCHES", 1 << 20)
- whole = tower(pixels, grid_thw)
-
- monkeypatch.setattr(M, "_VISION_CALL_MAX_PATCHES", cap)
- calls = _count_calls(tower, monkeypatch)
- chunked = tower(pixels, grid_thw)
-
- expected = M._frame_calls([tuple(row) for row in grid], cap)
- assert len(calls) == len(expected) > 1
- assert all(rows <= cap for rows, _ in calls)
- _assert_close(chunked, whole)
-
-
-def test_a_request_past_the_old_whole_request_ceiling_runs(tower):
- # 66 frames of 512x512 pixels: 33 temporal groups of 32x32 patches, 33792 rows in two
- # calls (32 groups, then 1). The last group's rows, alone in its call, and the first
- # group's, in a call with 31 others, equal each group run on its own.
- grid = [(33, 32, 32)]
- pixels = _pixels(grid)
- out = tower(pixels, torch.tensor(grid))
- assert out.shape == (33 * 32 * 32 // 4, OUT)
- assert torch.isfinite(out).all()
- group_rows, group_out = 32 * 32, 32 * 32 // 4
- for group in (0, 32):
- alone = tower(
- pixels[group * group_rows : (group + 1) * group_rows], torch.tensor([(1, 32, 32)])
- )
- _assert_close(out[group * group_out : (group + 1) * group_out], alone)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_kv_quant.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_kv_quant.py
deleted file mode 100644
index 46c8dad1a1ab..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_kv_quant.py
+++ /dev/null
@@ -1,229 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""qwen4_exp derives its KV storage from the engine's `kv_cache_quant_algo`.
-
-`None` keeps the checkpoint's bf16 K/V with the attention op's `quant_mode` 0 and no KV
-scales; `FP8` stores K/V as e4m3 under `QM_FP8_KV_CACHE` with unit fp32 `[1]` scales; `NVFP4`
-stores e2m1 pairs as uint8 half a head wide under `QM_NVFP4_KV_CACHE` with unit fp32 `[3]`
-scales; any other algorithm is refused at construction. The core is built on the meta device from the
-checkpoint's text and vision config values, so nothing here touches a device; the
-`post_load_weights` cases stub out everything that reads weight values.
-"""
-
-from __future__ import annotations
-
-from types import SimpleNamespace
-from unittest import mock
-
-import pytest
-import torch
-
-__extra_import_path__ = ["../../.."] # noqa: F841 -- repo's file-scoped import hook
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501
- modeling as M,
-)
-from tensorrt_llm._torch.configs import Qwen4ExpTextConfig, Qwen4ExpVisionConfig
-from tensorrt_llm._torch.model_config import ModelConfig
-from tensorrt_llm.llmapi import QSASparseAttentionConfig
-from tensorrt_llm.mapping import Mapping
-from tensorrt_llm.models.modeling_utils import QuantConfig
-from tensorrt_llm.quantization.mode import QuantAlgo
-
-pytestmark = pytest.mark.cpu_only
-
-NUM_LAYERS = 48
-
-# Qwen3.8-Flash-Next's text config.
-_TEXT = dict(
- architectures=["Qwen4ExpForCausalLM"],
- attention_bias=False,
- attention_dropout=0.0,
- bos_token_id=248044,
- dtype="bfloat16",
- eos_token_id=248044,
- full_attention_interval=4,
- hc_count=4,
- hc_lowrank=320,
- head_dim=256,
- heads_per_ngram=8,
- hidden_act="silu",
- hidden_size=2560,
- indexer_budget=2048,
- indexer_compress_ratio=4,
- indexer_head_dim=128,
- indexer_kv_heads=1,
- indexer_n_heads=4,
- initializer_range=0.02,
- layer_types=["full_attention" if i % 4 == 3 else "linear_attention" for i in range(NUM_LAYERS)],
- linear_conv_kernel_dim=4,
- linear_key_head_dim=128,
- linear_num_key_heads=16,
- linear_num_value_heads=48,
- linear_value_head_dim=128,
- make_ngram_vocab_size_divisible_by=128,
- mamba_ssm_dtype="float32",
- max_position_embeddings=262144,
- moe_intermediate_size=640,
- mtp=dict(
- hybrid=True,
- layer_types=["full_attention"],
- mtp_use_hidden_state_from_layer=None,
- num_hidden_layers=1,
- rope_theta=10000000,
- ),
- mtp_num_hidden_layers=1,
- mtp_use_dedicated_embeddings=False,
- ngram_size=3,
- ngram_vocab_size_base=20000000,
- norm_topk_prob=True,
- num_attention_heads=24,
- num_experts=512,
- num_experts_per_tok=10,
- num_hidden_layers=NUM_LAYERS,
- num_key_value_heads=2,
- number_of_conv_states=3,
- output_gate_type="sigmoid",
- output_router_logits=False,
- pad_token_id=None,
- partial_rotary_factor=0.25,
- ple_conv_kernel_size=4,
- ple_embed_dim=2560,
- ple_layer_ids=[2],
- rms_norm_eps=1e-06,
- rope_parameters=dict(
- mrope_interleaved=True,
- mrope_section=[11, 11, 10],
- partial_rotary_factor=0.25,
- rope_theta=10000000,
- rope_type="default",
- ),
- router_aux_loss_coef=0.001,
- seed=1234,
- shared_expert_intermediate_size=640,
- split_ngram_parts=128,
- tie_word_embeddings=False,
- use_cache=True,
- vocab_size=248320,
-)
-_VISION = dict(
- deepstack_visual_indexes=[],
- depth=27,
- dtype="bfloat16",
- hidden_act="gelu_pytorch_tanh",
- hidden_size=1152,
- in_channels=3,
- intermediate_size=4304,
- num_heads=16,
- num_position_embeddings=2304,
- out_hidden_size=2560,
- patch_size=16,
- spatial_merge_size=2,
- temporal_patch_size=2,
-)
-
-
-def _core(kv_cache_quant_algo):
- """The core of a Qwen3.8-Flash-Next engine whose resolved `kv_cache_quant_algo` is
- `kv_cache_quant_algo`, with every parameter on the meta device."""
- model_config = ModelConfig(
- pretrained_config=Qwen4ExpTextConfig(**_TEXT),
- mapping=Mapping(world_size=1, tp_size=1, rank=0),
- quant_config=QuantConfig(
- quant_algo=QuantAlgo.MIXED_PRECISION, kv_cache_quant_algo=kv_cache_quant_algo
- ),
- moe_backend="CUTLASS",
- moe_backend_requested="TRTLLM",
- sparse_attention_config=QSASparseAttentionConfig(),
- )
- composite = SimpleNamespace(
- vision_config=Qwen4ExpVisionConfig(**_VISION), image_token_id=248056, video_token_id=248057
- )
- # The core asserts the device it runs on; the weights it declares stay on meta.
- with (
- mock.patch.object(torch.cuda, "get_device_capability", return_value=(10, 3)),
- torch.device("meta"),
- ):
- return M.Qwen4ExpTp1ModelingV2(model_config, composite)
-
-
-def _load(core, monkeypatch) -> None:
- """`post_load_weights` on a core with the weights it reads replaced by CPU stubs.
-
- Everything that reads weight values (the expert scalars, the n-gram table's driver
- residency, the per-step targets) is replaced; what is left is the derivation under test
- and the shape checks around it. The derived scales land on `embed`'s device, here CPU,
- so their values can be read.
- """
- table_rows = core.ngram_vocab_divisor
- stubs = dict(
- embed=torch.empty(1),
- ple_table=torch.empty(table_rows, 1),
- ple_offs=torch.zeros(1, dtype=torch.int64),
- ple_vocab=torch.full((1,), table_rows, dtype=torch.int64),
- )
- for key, value in stubs.items():
- core.w[key] = torch.nn.Parameter(value, requires_grad=False)
- monkeypatch.setattr(M, "_moe_layer_scalars", lambda w, i, backend: {})
- for target in ("PrefillTarget", "DecodeTarget", "VisionTower"):
- monkeypatch.setattr(M, target, lambda core: None)
- monkeypatch.setattr(core, "_report_ple_residency", lambda: None)
- core.post_load_weights()
-
-
-def test_unquantized_kv_cache_keeps_the_checkpoints_bf16(monkeypatch) -> None:
- core = _core(None)
- assert core.quant_mode == 0
- assert core.kv_pool_dtype == torch.bfloat16
- geometry = core._pool_geometry(32)
- assert geometry.kv_dtype == torch.bfloat16
- assert geometry.kv_width == core.head_dim
- assert core.kv_scale_orig_quant is None and core.kv_scale_quant_orig is None
- _load(core, monkeypatch)
- assert core.kv_scale_orig_quant is None and core.kv_scale_quant_orig is None
-
-
-def test_fp8_kv_cache_stores_e4m3_with_unit_fp32_single_element_scales(monkeypatch) -> None:
- """The generation kernels read `kvCacheScale[0]` and the QSA store and sparse GQA take
- no scale, so the scale is a real tensor holding exactly 1.0: the dense and the sparse
- path then agree on what a cache slot holds."""
- core = _core(QuantAlgo.FP8)
- assert core.quant_mode == M.QM_FP8_KV_CACHE
- assert core.kv_pool_dtype == torch.float8_e4m3fn
- # The pool contract holds every manager's K/V pool to this dtype; the index and position
- # planes keep the activation dtype.
- geometry = core._pool_geometry(32)
- assert geometry.kv_dtype == torch.float8_e4m3fn
- assert geometry.kv_width == core.head_dim
- assert geometry.dtype == torch.bfloat16
- _load(core, monkeypatch)
- for scale in (core.kv_scale_orig_quant, core.kv_scale_quant_orig):
- assert scale is not None
- assert scale.dtype == torch.float32 and tuple(scale.shape) == (1,)
- assert torch.equal(scale, torch.ones(1))
-
-
-def test_nvfp4_kv_cache_stores_e2m1_pairs_with_unit_fp32_scale_triples(monkeypatch) -> None:
- """The attention op reads the q, K and V second-level scales out of one fp32 `[3]` tensor,
- and the QSA NVFP4 store and sparse GQA take the same tensor, so every writer and reader
- applies the one value, 1.0 (the checkpoint has no KV scales)."""
- core = _core(QuantAlgo.NVFP4)
- assert core.quant_mode == M.QM_NVFP4_KV_CACHE
- # An NVFP4 byte holds two values: the pool is uint8, half a head wide.
- assert core.kv_pool_dtype == torch.uint8
- geometry = core._pool_geometry(32)
- assert geometry.kv_dtype == torch.uint8
- assert geometry.kv_width == core.head_dim // 2
- assert geometry.dtype == torch.bfloat16
- _load(core, monkeypatch)
- for scale in (core.kv_scale_orig_quant, core.kv_scale_quant_orig):
- assert scale is not None
- assert scale.dtype == torch.float32 and tuple(scale.shape) == (3,)
- assert torch.equal(scale, torch.ones(3))
-
-
-def test_other_kv_cache_algorithms_are_refused() -> None:
- """The QSA backend serves its sparse path over half, bf16, fp8 and NVFP4 pools and falls
- back to dense for any other, so a pool of another storage must not be accepted."""
- with pytest.raises(AssertionError, match="kv_cache_quant_algo"):
- _core(QuantAlgo.INT8)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_moe_scalars.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_moe_scalars.py
deleted file mode 100644
index 7706883208f5..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_moe_scalars.py
+++ /dev/null
@@ -1,74 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""qwen4_exp's routed-expert scalars, derived on CPU from the checkpoint's scale rows.
-
-CuteDSL requantizes FC1's output with ONE global scale per layer, where trtllm-gen folds
-each expert's `down_proj.input_scale` into its own FC1 scalar. The CuteDSL scalars are
-exact for any spread of `down_proj.input_scale` across experts: they are the engine's
-`max(w2 input_scale)` convention.
-"""
-
-from __future__ import annotations
-
-import pytest
-import torch
-
-__extra_import_path__ = ["../../.."] # noqa: F841 -- repo's file-scoped import hook
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501
- modeling as M,
-)
-
-pytestmark = pytest.mark.cpu_only
-
-EXPERTS = 8
-
-
-def _rows(isc2: torch.Tensor, seed: int = 0) -> dict:
- gen = torch.Generator().manual_seed(seed)
- isc1 = torch.full((EXPERTS,), 0.0123)
- ws2_1 = torch.rand(EXPERTS, generator=gen) * 1e-3 + 1e-4
- return {
- "l0_e_isc1": isc1,
- "l0_e_isc1_up": isc1.clone(),
- "l0_e_ws2_1": ws2_1,
- "l0_e_ws2_1_up": ws2_1.clone(),
- "l0_e_isc2": isc2,
- "l0_e_ws2_2": torch.rand(EXPERTS, generator=gen) * 1e-3 + 1e-4,
- }
-
-
-def test_uniform_isc2_keeps_the_per_expert_products_bitwise() -> None:
- w = _rows(torch.full((EXPERTS,), 0.0371))
- got = M._moe_layer_scalars(w, 0, "CUTEDSL")
- assert torch.equal(got["alpha2"], w["l0_e_isc2"] * w["l0_e_ws2_2"])
- assert torch.equal(got["global_sf"], 1.0 / w["l0_e_isc2"][:1])
- assert got["global_sf"].shape == (1,)
- assert torch.equal(got["alpha1"], w["l0_e_isc1"] * w["l0_e_ws2_1"])
-
-
-def test_isc2_spread_is_folded_into_alpha2() -> None:
- """Whatever the spread (here 40x), `global_sf * alpha2[e]` is expert e's own `ws2_2[e]`,
- the product trtllm-gen forms from `(1 / isc2[e]) * (isc2[e] * ws2_2[e])`: the FC2 value an
- expert produces is the same whether its input was quantized at its own `1 / isc2[e]` or
- at the layer's single scale."""
- isc2 = 0.01 * torch.linspace(1.0, 40.0, EXPERTS)
- w = _rows(isc2.flip(0))
- cute = M._moe_layer_scalars(w, 0, "CUTEDSL")
- gen = M._moe_layer_scalars(w, 0, "TRTLLM")
- isc2 = w["l0_e_isc2"]
- assert torch.equal(cute["global_sf"], 1.0 / isc2.max().reshape(1))
- # The CuteDSL FC1 quantizes at the smallest per-expert global scale: no expert's FC2
- # input is scaled above what the checkpoint calibrated for it.
- assert (cute["global_sf"] <= 1.0 / isc2).all()
- assert torch.allclose(
- cute["global_sf"] * cute["alpha2"], (1.0 / isc2) * gen["output2_scale_scalar"], rtol=1e-6
- )
- assert torch.allclose(cute["global_sf"] * cute["alpha2"], w["l0_e_ws2_2"], rtol=1e-6)
-
-
-def test_gate_up_input_scale_mismatch_is_refused() -> None:
- w = _rows(torch.full((EXPERTS,), 0.0371))
- w["l0_e_isc1_up"] = w["l0_e_isc1_up"] * 1.5
- with pytest.raises(AssertionError, match="gate/up input_scale differ"):
- M._moe_layer_scalars(w, 0, "CUTEDSL")
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_multimodal_embeds.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_multimodal_embeds.py
deleted file mode 100644
index d7c71b3b43b9..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_multimodal_embeds.py
+++ /dev/null
@@ -1,274 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""qwen4_exp's multimodal splice runs the vision tower once per modality a request
-carries and puts the rows back in prompt order; the tower itself runs a request's frames
-in calls of whole frames under a row ceiling.
-
-The tower is a stub whose rows name the item they came from, so the bookkeeping --
-modality dispatch, item metadata, the chunk slice, the frame grouping and the order of
-the concatenated calls -- runs on the CPU without the vision ops or an engine.
-"""
-
-from __future__ import annotations
-
-import types
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-__extra_import_path__ = ["../../.."] # noqa: F841 -- repo's file-scoped import hook
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501
- modeling as M,
-)
-from tensorrt_llm.inputs.multimodal import (
- MULTIMODAL_ENCODER_ITEM_METADATA_KEY,
- MultimodalItemMetadata,
- MultimodalParams,
-)
-
-pytestmark = pytest.mark.cpu_only
-
-Shell = M.ModelingV2Qwen38FlashNextNvfp4Sm103Tp1
-
-HIDDEN = 3
-MERGE = 2
-IMAGE, VIDEO = 1, 2
-
-
-class _Tower:
- """Stands in for `VisionTower`: every row of item `i` of modality `m` holds
- `100 * m + i`, so a spliced row says where it came from."""
-
- def __init__(self):
- self.calls = 0
-
- def __call__(self, pixel_values, grid_thw):
- self.calls += 1
- modality = int(pixel_values[0, 0])
- return torch.cat(
- [
- torch.full((t * h * w // (MERGE * MERGE), HIDDEN), float(100 * modality + item))
- for item, (t, h, w) in enumerate(grid_thw.tolist())
- ]
- )
-
-
-def _shell(tower):
- shell = SimpleNamespace(
- model=SimpleNamespace(
- v_merge=MERGE, w={"embed": torch.zeros(16, HIDDEN)}, vision_tower=tower
- )
- )
- shell._run_tower = types.MethodType(Shell._run_tower, shell)
- shell._in_prompt_order = types.MethodType(Shell._in_prompt_order, shell)
- return shell
-
-
-def _media(modality, grids):
- """One `multimodal_data` entry; `pixel_values` only carries the stub's modality."""
- patches = sum(t * h * w for t, h, w in grids)
- key = "pixel_values" if modality == IMAGE else "pixel_values_videos"
- grid_key = "image_grid_thw" if modality == IMAGE else "video_grid_thw"
- return {key: torch.full((patches, 1), float(modality)), grid_key: torch.tensor(grids)}
-
-
-def _params(image=None, video=None, order=None, runtime=None, mrope=True):
- data = {}
- if mrope:
- data["mrope_config"] = {"mrope_position_ids": torch.zeros(3, 1, 8, dtype=torch.int32)}
- if image is not None:
- data["image"] = _media(IMAGE, image)
- if video is not None:
- data["video"] = _media(VIDEO, video)
- if order is not None:
- data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalItemMetadata(item_refs=order)
- return MultimodalParams(multimodal_data=data, multimodal_runtime=runtime)
-
-
-def _splice(tower, params, num_rows, padding=0):
- """The tower's rows spliced, as the prefill target's break does, over a prompt whose
- tokens 1..num_rows are the placeholders and which `padding` rows follow; returns the
- rows that landed there, one value per row, or None when the step carries no media."""
- kwargs = {
- "multimodal_params": [params],
- "mm_token_indices": torch.arange(1, 1 + num_rows, dtype=torch.int64),
- }
- vision = Shell._multimodal_embeds(_shell(tower), kwargs)
- if vision is None:
- return None
- total = num_rows + 2
- h = torch.full((total + padding, HIDDEN), -1.0)
- target = SimpleNamespace(_step=SimpleNamespace(vision=vision, num_tokens=total))
- M.PrefillTarget._splice(target, h)
- assert torch.count_nonzero(h[total:]) == 0, "padding rows must be zeroed"
- assert torch.all(h[0] == -1.0) and torch.all(h[total - 1] == -1.0)
- return h[1 : 1 + num_rows, 0].tolist()
-
-
-def _item(modality, index):
- return modality, index
-
-
-# Two videos of 4 and 8 rows, two images of 4 and 1.
-VIDEOS = [(2, 2, 4), (2, 4, 4)]
-IMAGES = [(1, 4, 4), (1, 2, 2)]
-MIXED_ORDER = [_item("video", 0), _item("image", 0), _item("video", 1), _item("image", 1)]
-MIXED_ROWS = [200.0] * 4 + [100.0] * 4 + [201.0] * 8 + [101.0]
-
-
-def test_the_splice_fills_the_placeholders_and_zeroes_the_padding():
- """An image request's rows land on its placeholders in one tower call and the padding
- rows after the step are zeroed (`_splice` checks); a text step writes only the padding."""
- tower = _Tower()
- assert _splice(tower, _params(image=IMAGES), 5, padding=3) == [100.0] * 4 + [101.0]
- assert tower.calls == 1
- target = SimpleNamespace(_step=SimpleNamespace(vision=None, num_tokens=2))
- h = torch.full((4, HIDDEN), -1.0)
- M.PrefillTarget._splice(target, h)
- assert torch.all(h[:2] == -1.0) and torch.count_nonzero(h[2:]) == 0
-
-
-def test_video_rows_reach_the_placeholders_item_by_item():
- tower = _Tower()
- assert _splice(tower, _params(video=VIDEOS), 12) == [200.0] * 4 + [201.0] * 8
- assert tower.calls == 1
-
-
-def test_a_mixed_request_is_reassembled_in_prompt_order():
- tower = _Tower()
- params = _params(image=IMAGES, video=VIDEOS, order=MIXED_ORDER)
- assert _splice(tower, params, len(MIXED_ROWS)) == MIXED_ROWS
- assert tower.calls == 2
-
-
-def test_a_mixed_request_needs_its_manifest():
- with pytest.raises(AssertionError, match="needs item metadata"):
- _splice(_Tower(), _params(image=IMAGES, video=VIDEOS), len(MIXED_ROWS))
-
-
-def test_a_chunk_takes_its_slice_of_the_prompt_order_rows():
- runtime = SimpleNamespace(
- num_cached_mm_tokens=3, num_mm_tokens_in_chunk=6, total_embeds_in_request=len(MIXED_ROWS)
- )
- params = _params(image=IMAGES, video=VIDEOS, order=MIXED_ORDER, runtime=runtime)
- assert _splice(_Tower(), params, 6) == MIXED_ROWS[3:9]
-
-
-def test_a_chunk_without_placeholders_skips_the_tower():
- runtime = SimpleNamespace(
- num_cached_mm_tokens=len(MIXED_ROWS),
- num_mm_tokens_in_chunk=0,
- total_embeds_in_request=len(MIXED_ROWS),
- )
- tower = _Tower()
- params = _params(image=IMAGES, video=VIDEOS, order=MIXED_ORDER, runtime=runtime)
- assert _splice(tower, params, 0) is None
- assert tower.calls == 0
-
-
-def _frames(grid):
- """The item-by-item frame list a grid describes: one `(h, w)` per frame."""
- return [(h, w) for t, h, w in grid for _ in range(t)]
-
-
-def _random_grid(gen):
- return [
- (int(torch.randint(1, 6, (1,), generator=gen)), 2 * h, 2 * w)
- for h, w in torch.randint(
- 1, 5, (int(torch.randint(1, 6, (1,), generator=gen)), 2), generator=gen
- ).tolist()
- ]
-
-
-def test_frames_are_grouped_whole_and_in_order():
- # 16-row frames, 32 rows to a call: one item's frames split across calls.
- assert M._frame_calls([(5, 4, 4)], 32) == [[(2, 4, 4)], [(2, 4, 4)], [(1, 4, 4)]]
- # A call takes the end of one item and the start of the next.
- grid = [(1, 8, 4), (3, 4, 4), (1, 2, 2)]
- assert M._frame_calls(grid, 48) == [[(1, 8, 4), (1, 4, 4)], [(2, 4, 4), (1, 2, 2)]]
- assert M._frame_calls(grid, 1 << 20) == [grid]
-
-
-def test_a_call_is_filled_to_the_ceiling_and_never_past_it():
- gen = torch.Generator().manual_seed(0)
- for _ in range(200):
- grid = _random_grid(gen)
- cap = int(torch.randint(64, 300, (1,), generator=gen))
- calls = M._frame_calls(grid, cap)
- assert [f for call in calls for f in _frames(call)] == _frames(grid)
- sizes = [sum(t * h * w for t, h, w in call) for call in calls]
- assert all(0 < size <= cap for size in sizes)
- # Greedy: the next call's first frame would not have fit.
- for size, call in zip(sizes, calls[1:]):
- h, w = _frames(call)[0]
- assert size + h * w > cap
-
-
-def test_one_frame_past_the_ceiling_is_refused_and_one_at_it_is_not():
- assert M._frame_calls([(2, 8, 8)], 64) == [[(1, 8, 8)], [(1, 8, 8)]]
- with pytest.raises(AssertionError, match="one frame is 8x10 = 80 patch rows"):
- M._frame_calls([(1, 4, 4), (3, 8, 10)], 64)
-
-
-PIXEL_WIDTH = 4
-
-
-def _tower(monkeypatch, cap):
- """A `VisionTower` whose per-call body is a stub: merged row `j` of a call carries the
- pixel value of the call's patch row `4 * j`, which `_pixels` set to the patch's
- index in the request."""
- monkeypatch.setattr(M, "_VISION_CALL_MAX_PATCHES", cap)
- tower = object.__new__(M.VisionTower)
- tower.core = SimpleNamespace(
- w={"embed": torch.zeros(1)}, v_pixel_width=PIXEL_WIDTH, v_merge=MERGE
- )
- tower.calls = []
-
- def run(pixel_values, grid, grid_thw):
- tower.calls.append((pixel_values, grid, grid_thw))
- return pixel_values[:: MERGE * MERGE, :1].expand(-1, HIDDEN).clone()
-
- tower._run = run
- return tower
-
-
-def _pixels(grid):
- patches = sum(t * h * w for t, h, w in grid)
- return torch.arange(patches, dtype=torch.float32)[:, None].expand(-1, PIXEL_WIDTH).clone()
-
-
-def test_the_tower_calls_whole_frames_and_concatenates_their_rows_in_order(monkeypatch):
- grid = [(1, 8, 4), (3, 4, 4), (2, 2, 6), (1, 2, 2)]
- tower = _tower(monkeypatch, 48)
- out = tower(_pixels(grid), torch.tensor(grid))
- patches = sum(t * h * w for t, h, w in grid)
- assert out[:, 0].tolist() == [float(4 * j) for j in range(patches // 4)]
- assert [call[1] for call in tower.calls] == M._frame_calls(grid, 48)
- start = 0
- for pixel_values, call, grid_thw in tower.calls:
- rows = sum(t * h * w for t, h, w in call)
- assert pixel_values[:, 0].tolist() == [float(i) for i in range(start, start + rows)]
- assert rows <= 48
- assert grid_thw.tolist() == [list(row) for row in call]
- assert grid_thw.dtype == torch.int64
- start += rows
- assert start == patches
-
-
-def test_a_request_past_the_old_whole_request_ceiling_is_served(monkeypatch):
- # 33 temporal groups of 32x32 patches (66 frames of 512x512 pixels): 33792 rows, 1024
- # rows past what a single call may carry.
- ceiling = M._VISION_CALL_MAX_PATCHES
- grid = [(33, 32, 32)]
- tower = _tower(monkeypatch, ceiling)
- out = tower(_pixels(grid), torch.tensor(grid))
- assert [call[1] for call in tower.calls] == [[(32, 32, 32)], [(1, 32, 32)]]
- assert out.shape[0] == 33 * 32 * 32 // 4
- assert out[:, 0].tolist() == [float(4 * j) for j in range(out.shape[0])]
-
-
-def test_a_media_request_needs_the_processors_position_ids():
- with pytest.raises(AssertionError, match="mrope_position_ids"):
- _splice(_Tower(), _params(video=VIDEOS, mrope=False), 12)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_validate_request.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_validate_request.py
deleted file mode 100644
index 80688e4b3cf9..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_validate_request.py
+++ /dev/null
@@ -1,209 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""qwen4_exp's admission check fails a bad media request before the forward would assert.
-
-The check reads only the request's host-side payload, so these run on fake request
-objects with meta pixel tensors: no GPU, vision ops or engine. Every refusal is a
-`ValueError`, which the executor turns into that one request's error.
-"""
-
-from __future__ import annotations
-
-import types
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-__extra_import_path__ = ["../../.."] # noqa: F841 -- repo's file-scoped import hook
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501
- modeling as M,
-)
-from tensorrt_llm.inputs.multimodal import (
- MULTIMODAL_ENCODER_ITEM_METADATA_KEY,
- MultimodalItemMetadata,
-)
-
-pytestmark = pytest.mark.cpu_only
-
-Shell = M.ModelingV2Qwen38FlashNextNvfp4Sm103Tp1
-
-MERGE = 2
-PATCH = 16
-PIXEL_WIDTH = 4
-TEXT_TOKENS = 5
-
-
-def _shell():
- shell = SimpleNamespace(
- model=SimpleNamespace(v_merge=MERGE, v_patch=PATCH, v_pixel_width=PIXEL_WIDTH)
- )
- shell.validate_request = types.MethodType(Shell.validate_request, shell)
- shell._media_grid = types.MethodType(Shell._media_grid, shell)
- shell._check_item_order = Shell._check_item_order
- return shell
-
-
-def _media(modality, grids):
- """A payload whose pixel tensor has the right shape and no storage."""
- patches = sum(t * h * w for t, h, w in grids)
- key = "pixel_values" if modality == "image" else "pixel_values_videos"
- grid_key = "image_grid_thw" if modality == "image" else "video_grid_thw"
- return {
- key: torch.empty(patches, PIXEL_WIDTH, device="meta"),
- grid_key: torch.tensor(grids),
- }
-
-
-def _request(image=None, video=None, order=None, **overrides):
- """A request whose prompt is `TEXT_TOKENS` text tokens around the placeholders the
- media produce, with the position ids and placeholder cumsum the processor emits."""
- rows = sum(t * h * w for t, h, w in (image or []) + (video or [])) // (MERGE * MERGE)
- prompt_len = rows + TEXT_TOKENS
- data = {
- "mrope_config": {"mrope_position_ids": torch.zeros(3, 1, prompt_len, dtype=torch.int32)},
- "multimodal_embed_mask_cumsum": torch.cat(
- [torch.arange(1, rows + 1), torch.full((TEXT_TOKENS,), rows)]
- ),
- }
- if image is not None:
- data["image"] = _media("image", image)
- if video is not None:
- data["video"] = _media("video", video)
- if order is not None:
- data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalItemMetadata(item_refs=order)
- fields = {"py_multimodal_data": data, "prompt_len": prompt_len}
- fields.update(overrides)
- return SimpleNamespace(**fields)
-
-
-def _admit(request):
- _shell().validate_request(request)
-
-
-def _item(modality, index):
- return modality, index
-
-
-IMAGES = [(1, 4, 4), (1, 2, 2)]
-VIDEOS = [(2, 2, 4), (2, 4, 4)]
-MIXED_ORDER = [_item("video", 0), _item("image", 0), _item("video", 1), _item("image", 1)]
-
-
-def test_clean_requests_are_admitted():
- # A single-modality request needs no manifest.
- _admit(_request(image=IMAGES))
- _admit(_request(video=VIDEOS))
- _admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER))
-
-
-def test_a_request_without_media_costs_nothing():
- # Not even the other request attributes are read.
- _admit(SimpleNamespace(py_multimodal_data=None))
- _admit(SimpleNamespace(py_multimodal_data={}))
- # A warmup request carries position ids and no media.
- _admit(
- SimpleNamespace(
- py_multimodal_data={"mrope_config": {"mrope_position_ids": torch.zeros(3, 1, 8)}}
- )
- )
- # Precomputed embeddings arrive without pixels; the splice skips them.
- _admit(SimpleNamespace(py_multimodal_data={"multimodal_embedding": torch.zeros(4, 3)}))
-
-
-def test_the_longest_admitted_frame_is_the_calls_ceiling(monkeypatch):
- monkeypatch.setattr(M, "_VISION_CALL_MAX_PATCHES", 64)
- _admit(_request(image=[(1, 8, 8)]))
- _admit(_request(video=[(3, 4, 16)]))
- with pytest.raises(ValueError, match="frames of 8x10 = 80 patches"):
- _admit(_request(image=[(1, 8, 8), (1, 8, 10)]))
-
-
-def test_a_frame_past_the_ceiling_fails_the_request():
- # 3072x3072 pixels: a 192x192 grid, 36864 patches in one frame.
- request = _request(video=[(1, 192, 192)])
- with pytest.raises(ValueError, match=r"video 0 has frames of 192x192 = 36864 patches"):
- _admit(request)
- # The message names the pixels and the bound.
- with pytest.raises(ValueError, match=r"\(3072x3072 pixels\).*at most 32768 patches"):
- _admit(request)
-
-
-def test_the_frame_bound_is_per_frame_not_per_item():
- # 33 temporal groups of 32x32 patches: 33792 rows in all, 1024 past one call.
- _admit(_request(video=[(33, 32, 32)]))
- # 4 frames of 128x128 patches (16384 each) are fine; 192x192 in the second item is not.
- _admit(_request(video=[(4, 128, 128)]))
- with pytest.raises(ValueError, match="video 1 has frames of 192x192"):
- _admit(_request(video=[(4, 128, 128), (1, 192, 192)]))
-
-
-def test_a_malformed_grid_is_refused():
- request = _request(image=[(1, 4, 4)])
- request.py_multimodal_data["image"]["image_grid_thw"] = torch.tensor([(1, 4, 5)])
- with pytest.raises(ValueError, match="spatial_merge_size 2"):
- _admit(request)
-
-
-def test_a_grid_that_is_not_one_row_per_item_is_refused():
- request = _request(image=IMAGES)
- request.py_multimodal_data["image"]["image_grid_thw"] = torch.tensor([1, 4, 4])
- with pytest.raises(ValueError, match="one \\(t, h, w\\) row per image"):
- _admit(request)
-
-
-def test_pixels_that_disagree_with_the_grid_are_refused():
- request = _request(video=VIDEOS)
- video = request.py_multimodal_data["video"]
- video["pixel_values_videos"] = torch.empty(5, PIXEL_WIDTH, device="meta")
- with pytest.raises(ValueError, match="pixel_values_videos is .*describe 48 patch rows"):
- _admit(request)
-
-
-def test_a_payload_without_its_pixels_or_grid_is_refused():
- request = _request(video=VIDEOS)
- del request.py_multimodal_data["video"]["video_grid_thw"]
- with pytest.raises(ValueError, match="this target reads 'pixel_values_videos'"):
- _admit(request)
-
-
-def test_missing_or_mis_sized_position_ids_fail_the_request():
- request = _request(video=VIDEOS)
- request.py_multimodal_data["mrope_config"]["mrope_position_ids"] = torch.zeros(3, 1, 40)
- with pytest.raises(ValueError, match="one position per prompt token"):
- _admit(request)
- del request.py_multimodal_data["mrope_config"]
- with pytest.raises(ValueError, match="mrope_position_ids.*got None"):
- _admit(request)
-
-
-def test_a_mixed_request_needs_its_manifest():
- with pytest.raises(ValueError, match="needs item metadata"):
- _admit(_request(image=IMAGES, video=VIDEOS))
- with pytest.raises(ValueError, match="needs item metadata"):
- _admit(_request(image=IMAGES, video=VIDEOS, order=[]))
-
-
-def test_a_manifest_that_disagrees_with_the_payload_is_refused():
- with pytest.raises(ValueError, match="lists image items \\[0\\]; the payload carries 2"):
- _admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER[:3]))
- with pytest.raises(ValueError, match="lists video items \\[0, 1, 0\\]"):
- _admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER + [_item("video", 0)]))
- with pytest.raises(ValueError, match="is not an item of this request"):
- _admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER + [_item("audio", 0)]))
- with pytest.raises(ValueError, match="is not an item of this request"):
- _admit(_request(image=IMAGES, video=VIDEOS, order=[("image",)]))
-
-
-def test_the_placeholder_count_is_the_sum_over_both_modalities():
- request = _request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER)
- request.py_multimodal_data["multimodal_embed_mask_cumsum"] = torch.arange(1, 14)
- with pytest.raises(ValueError, match="holds 13 image/video placeholder tokens.*17"):
- _admit(request)
-
-
-def test_a_request_without_a_cumsum_skips_the_count():
- request = _request(video=VIDEOS)
- del request.py_multimodal_data["multimodal_embed_mask_cumsum"]
- _admit(request)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_weight_table.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_weight_table.py
deleted file mode 100644
index a3bba2823db9..000000000000
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_weight_table.py
+++ /dev/null
@@ -1,401 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""qwen4_exp's weight table declares what the hand-written code did.
-
-The declaration reference below is a transcription of the `__init__` block the table
-replaced, keyed the way the table keys it: vision block `b`'s weights are `l{b}_vis_*`
-and the MTP module's are `l{num_hidden_layers}_mtp_*`, its layer index in the checkpoint
-and the KV pool. The n-gram table is not in either: its height is a checkpoint fact,
-declared at load.
-
-No GPU: nothing here allocates, and `dims` reads the configs alone.
-"""
-
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-__extra_import_path__ = ["../../.."] # noqa: F841 -- repo's file-scoped import hook
-
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501
- modeling as M,
-)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501
- weights as W,
-)
-
-pytestmark = pytest.mark.cpu_only
-
-# Qwen3.8-Flash-Next's own numbers, so the comparison is against the shapes this target
-# actually ships rather than a toy configuration.
-_TEXT = dict(
- num_hidden_layers=48,
- layer_types=["full_attention" if i % 4 == 3 else "linear_attention" for i in range(48)],
- hidden_size=2560,
- hc_count=4,
- hc_lowrank=320,
- head_dim=256,
- num_attention_heads=24,
- num_key_value_heads=2,
- indexer_n_heads=4,
- indexer_kv_heads=1,
- indexer_head_dim=128,
- linear_num_key_heads=16,
- linear_num_value_heads=48,
- linear_key_head_dim=128,
- linear_value_head_dim=128,
- linear_conv_kernel_dim=4,
- num_experts=512,
- moe_intermediate_size=640,
- shared_expert_intermediate_size=640,
- vocab_size=248320,
- ple_layer_ids=[2],
- ple_embed_dim=2560,
- ple_conv_kernel_size=4,
- ngram_size=3,
- heads_per_ngram=8,
-)
-_VISION = dict(
- depth=27,
- hidden_size=1152,
- intermediate_size=4304,
- in_channels=3,
- temporal_patch_size=2,
- patch_size=16,
- num_position_embeddings=2304,
- spatial_merge_size=2,
-)
-
-
-def _core(mtp: bool, backend: str = "TRTLLM", prefill_copy: bool | None = None):
- """The configuration surface the table reads; the second expert copy defaults to kept
- under trtllm-gen."""
- return SimpleNamespace(
- model_config=SimpleNamespace(
- pretrained_config=SimpleNamespace(**_TEXT),
- torch_dtype=torch.bfloat16,
- spec_config=object() if mtp else None,
- ),
- vision_config=SimpleNamespace(**_VISION),
- mtp_enabled=mtp,
- trunk_moe_backend=backend,
- moe_prefill_copy=backend == "TRTLLM" if prefill_copy is None else prefill_copy,
- )
-
-
-def test_every_declared_weight_has_a_source():
- """A declared weight with no source allocates memory nothing ever fills -- silently
- zero rather than an error, which is what `load`'s coverage assert catches from the
- other side. With MTP live: the rows without it are a subset
- (`test_mtp_rows_appear_only_when_mtp_is_live`)."""
- core = _core(True)
- manifest = W.MODEL_WEIGHTS.manifest(core)
- for key, sources in manifest.items():
- assert sources, f"{key} is declared with no checkpoint source"
- with torch.device("meta"):
- declared = W.MODEL_WEIGHTS.declare(W.MODEL_WEIGHTS.dims(core))
- assert set(manifest) == set(declared)
-
-
-def test_mtp_rows_appear_only_when_mtp_is_live():
- off = W.MODEL_WEIGHTS.manifest(_core(False))
- on = W.MODEL_WEIGHTS.manifest(_core(True))
- added = set(on) - set(off)
- assert set(off) <= set(on)
- m = f"l{_TEXT['num_hidden_layers']}_mtp_"
- assert added and all(k.startswith(m) for k in added)
- # The draft layer runs QSA over its own indexer, read from the module's keys.
- indexer = {f"{m}idx_w", f"{m}idx_qn", f"{m}idx_kn"}
- assert indexer <= added
- assert all(
- str(src).startswith("mtp.layers.0.self_attn.indexer.")
- for k in indexer
- for src, _, _ in on[k]
- )
-
-
-def _reference_declaration(mtp: bool) -> dict:
- """The declaration block from modeling.py, transcribed."""
- t, v = SimpleNamespace(**_TEXT), SimpleNamespace(**_VISION)
- bf16, u8, f8, f32, i64 = (
- torch.bfloat16,
- torch.uint8,
- torch.float8_e4m3fn,
- torch.float32,
- torch.int64,
- )
- hidden, stream, rank = t.hidden_size, t.hc_count * t.hidden_size, t.hc_lowrank
- q_width = t.num_attention_heads * t.head_dim
- kv_width = t.num_key_value_heads * t.head_dim
- index_width = (t.indexer_n_heads + t.indexer_kv_heads) * t.indexer_head_dim
- key_dim = t.linear_num_key_heads * t.linear_key_head_dim
- value_dim = t.linear_num_value_heads * t.linear_value_head_dim
- conv_dim = 2 * key_dim + value_dim
- experts, inter, shared = (
- t.num_experts,
- t.moe_intermediate_size,
- t.shared_expert_intermediate_size,
- )
- ngram_heads = t.heads_per_ngram * (t.ngram_size - 1)
- vh, vi = v.hidden_size, v.intermediate_size
- vm = vh * v.spatial_merge_size**2
- nl = t.num_hidden_layers
- w = {}
-
- def P(name, *shape, dtype=bf16):
- w[name] = (tuple(shape), dtype)
-
- def hyper_connection(prefix):
- P(f"{prefix}_norm", stream)
- P(f"{prefix}_down", rank, stream)
- P(f"{prefix}_up", stream, rank)
- P(f"{prefix}_inject", t.hc_count, stream)
-
- def gated_attention(prefix, fused_in=False):
- if fused_in:
- # [q | k | v | gate | index q, k | one index head of zero pad]
- rows = 2 * q_width + 2 * kv_width + index_width + t.indexer_head_dim
- P(f"{prefix}attn_in_w", rows, hidden)
- else:
- P(f"{prefix}q_w", 2 * q_width, hidden)
- P(f"{prefix}k_w", kv_width, hidden)
- P(f"{prefix}v_w", kv_width, hidden)
- P(f"{prefix}o_w", hidden, q_width)
- P(f"{prefix}q_norm", t.head_dim)
- P(f"{prefix}k_norm", t.head_dim)
-
- def shared_expert(prefix):
- # gate|up, then the per-token gate's row and zero rows (`weights._cat_rows`).
- P(f"{prefix}sh_gu", 2 * shared + 8, hidden)
- P(f"{prefix}sh_dn", hidden, shared)
-
- for i in range(nl):
- hyper_connection(f"l{i}_ahc")
- hyper_connection(f"l{i}_mhc")
- if t.layer_types[i] == "full_attention":
- gated_attention(f"l{i}_", fused_in=True)
- P(f"l{i}_idx_qn", t.indexer_head_dim)
- P(f"l{i}_idx_kn", t.indexer_head_dim)
- else:
- # [in_proj_qkv | in_proj_z | in_proj_a | in_proj_b], one row block each.
- P(f"l{i}_in_proj_w", conv_dim + value_dim + 2 * t.linear_num_value_heads, hidden)
- P(f"l{i}_A_log", t.linear_num_value_heads)
- P(f"l{i}_dt_bias", t.linear_num_value_heads)
- P(f"l{i}_conv_w", conv_dim, 1, t.linear_conv_kernel_dim)
- P(f"l{i}_gnorm", t.linear_value_head_dim)
- P(f"l{i}_out_w", hidden, value_dim)
- P(f"l{i}_router", experts, hidden)
- P(f"l{i}_fc1_w", experts, 2 * inter, hidden // 2, dtype=u8)
- P(f"l{i}_fc1_s", experts, 2 * inter, hidden // 16, dtype=u8)
- P(f"l{i}_fc2_w", experts, hidden, inter // 2, dtype=u8)
- P(f"l{i}_fc2_s", experts, hidden, inter // 16, dtype=u8)
- # The trtllm-gen backend's second, CuteDSL-layout copy of the experts.
- P(f"l{i}_fc1_w_cd", experts, 2 * inter, hidden // 2, dtype=u8)
- P(f"l{i}_fc1_s_cd", experts, 2 * inter, hidden // 16, dtype=u8)
- P(f"l{i}_fc2_w_cd", experts, hidden, inter // 2, dtype=u8)
- P(f"l{i}_fc2_s_cd", experts, hidden, inter // 16, dtype=u8)
- for name in ("e_isc1", "e_isc1_up", "e_ws2_1", "e_ws2_1_up", "e_isc2", "e_ws2_2"):
- P(f"l{i}_{name}", experts, dtype=f32)
- shared_expert(f"l{i}_")
- P("fhc_norm", stream)
- P("fhc_down", rank, stream)
- P("fhc_up", stream, rank)
- P("embed", t.vocab_size, hidden)
- P("ple_key_w", stream, hidden)
- P("ple_val_w", t.ple_embed_dim, hidden)
- P("ple_nk", stream)
- P("ple_nq", stream)
- P("ple_nc", stream)
- P("ple_conv_w", stream, 1, t.ple_conv_kernel_size)
- P("ple_mult", t.ngram_size, dtype=i64)
- P("ple_vocab", ngram_heads, dtype=i64)
- P("ple_offs", ngram_heads, dtype=i64)
- P("ple_tscale", 1)
- P("vis_patch_w", vh, v.in_channels, v.temporal_patch_size, v.patch_size, v.patch_size)
- P("vis_patch_b", vh)
- P("vis_pos", v.num_position_embeddings, vh)
- for b in range(v.depth):
- P(f"l{b}_vis_n1_w", vh)
- P(f"l{b}_vis_n1_b", vh)
- P(f"l{b}_vis_qkv_w", 3 * vh, vh)
- P(f"l{b}_vis_qkv_b", 3 * vh)
- P(f"l{b}_vis_proj_w", vh, vh)
- P(f"l{b}_vis_proj_b", vh)
- P(f"l{b}_vis_n2_w", vh)
- P(f"l{b}_vis_n2_b", vh)
- P(f"l{b}_vis_fc1_w", vi, vh)
- P(f"l{b}_vis_fc1_b", vi)
- P(f"l{b}_vis_fc2_w", vh, vi)
- P(f"l{b}_vis_fc2_b", vh)
- P("vis_m_n_w", vh)
- P("vis_m_n_b", vh)
- P("vis_m_fc1_w", vm, vm)
- P("vis_m_fc1_b", vm)
- P("vis_m_fc2_w", hidden, vm)
- P("vis_m_fc2_b", hidden)
- if mtp:
- m = f"l{nl}_mtp_"
- P(f"{m}pre_e_norm", hidden)
- P(f"{m}pre_h_norm", stream)
- P(f"{m}fc_e", hidden, hidden)
- P(f"{m}fc_h", hidden, hidden)
- hyper_connection(f"{m}ahc")
- hyper_connection(f"{m}mhc")
- gated_attention(m)
- P(f"{m}idx_w", index_width, hidden)
- P(f"{m}idx_qn", t.indexer_head_dim)
- P(f"{m}idx_kn", t.indexer_head_dim)
- P(f"{m}router", experts, hidden)
- P(f"{m}fc1_w", experts, 2 * inter, hidden, dtype=f8)
- P(f"{m}fc1_s", experts, (2 * inter) // 128, hidden // 128, dtype=f32)
- P(f"{m}fc2_w", experts, hidden, inter, dtype=f8)
- P(f"{m}fc2_s", experts, hidden // 128, inter // 128, dtype=f32)
- shared_expert(m)
- P(f"{m}fhc_norm", stream)
- P(f"{m}fhc_down", rank, stream)
- P(f"{m}fhc_up", stream, rank)
- return w
-
-
-def test_the_table_declares_what_the_init_block_did():
- """With MTP live; without it the table drops exactly the MTP module's rows
- (`test_mtp_rows_appear_only_when_mtp_is_live`)."""
- core = _core(True)
- with torch.device("meta"):
- declared = W.MODEL_WEIGHTS.declare(W.MODEL_WEIGHTS.dims(core))
- got = {k: (tuple(p.shape), p.dtype) for k, p in declared.items()}
- want = _reference_declaration(True)
- assert set(got) == set(want), (
- sorted(set(got) - set(want))[:8],
- sorted(set(want) - set(got))[:8],
- )
- wrong = {k: (got[k], want[k]) for k in got if got[k] != want[k]}
- assert not wrong, list(wrong.items())[:8]
-
-
-def test_gdn_input_projections_fill_adjacent_row_blocks():
- """The concatenated GDN input projection: each checkpoint tensor lands in its own
- row block, in the column order the forward splits the GEMM output by, and the
- blocks tile the declared rows exactly."""
- t = SimpleNamespace(**_TEXT)
- core = _core(False)
- manifest = W.MODEL_WEIGHTS.manifest(core)
- key_dim = t.linear_num_key_heads * t.linear_key_head_dim
- value_dim = t.linear_num_value_heads * t.linear_value_head_dim
- widths = (
- 2 * key_dim + value_dim,
- value_dim,
- t.linear_num_value_heads,
- t.linear_num_value_heads,
- )
- gdn = [i for i, kind in enumerate(t.layer_types) if kind == "linear_attention"]
- assert len(gdn) == 36
- for i in gdn:
- rows = manifest[f"l{i}_in_proj_w"]
- names = [ckpt.rsplit(".", 2)[-2] for ckpt, _, _ in rows]
- assert names == ["in_proj_qkv", "in_proj_z", "in_proj_a", "in_proj_b"], names
- assert all(ckpt.startswith(f"{W._LM}.layers.{i}.linear_attn.") for ckpt, _, _ in rows)
- assert all(transform is None for _, _, transform in rows)
- spans = [(index[0].start, index[0].stop) for _, index, _ in rows]
- starts = [0, widths[0], widths[0] + widths[1], widths[0] + widths[1] + widths[2]]
- assert spans == [(s, s + w) for s, w in zip(starts, widths)], spans
- assert spans[-1][1] == sum(widths) == 16480
-
-
-def test_attention_input_projection_relayout():
- """`_attn_in_w` puts q, k, v, the output gate and the index projection side by side
- so that one GEMM's columns are exactly what the four separate projections gave:
- q and gate de-interleaved from q_proj's per-head `[q | gate]` rows, then a zero pad
- of one index head. Checked at the checkpoint's head geometry, a narrow hidden."""
- t = SimpleNamespace(**_TEXT)
- heads, d, kvh, hidden = t.num_attention_heads, t.head_dim, t.num_key_value_heads, 64
- idx_rows = (t.indexer_n_heads + t.indexer_kv_heads) * t.indexer_head_dim
- g = torch.Generator().manual_seed(0)
- q = torch.randn(2 * heads * d, hidden, generator=g, dtype=torch.float64)
- k = torch.randn(kvh * d, hidden, generator=g, dtype=torch.float64)
- v = torch.randn(kvh * d, hidden, generator=g, dtype=torch.float64)
- idx = torch.randn(idx_rows, hidden, generator=g, dtype=torch.float64)
- core = SimpleNamespace(heads_q=heads, head_dim=d, index_head_dim=t.indexer_head_dim)
- fused = W._attn_in_w((q, k, v, idx), core)
- rows = 2 * heads * d + 2 * kvh * d + idx_rows + t.indexer_head_dim
- assert fused.shape == (rows, hidden) and fused.is_contiguous()
-
- x = torch.randn(3, hidden, generator=g, dtype=torch.float64)
- y = x @ fused.t()
- qg = (x @ q.t()).reshape(3, heads, 2 * d)
- want = torch.cat(
- [
- qg[..., :d].reshape(3, -1),
- x @ k.t(),
- x @ v.t(),
- qg[..., d:].reshape(3, -1),
- x @ idx.t(),
- torch.zeros(3, t.indexer_head_dim, dtype=torch.float64),
- ],
- 1,
- )
- torch.testing.assert_close(y, want, rtol=0, atol=1e-12)
- # The row is a whole number of head_dim heads: the main norm+rope addresses it so.
- assert rows % d == 0
-
-
-def test_cutedsl_fc1_pairs_up_and_gate_rows_in_64_row_groups():
- """Destination rows `[128k, 128k + 64)` hold up rows `[64k, 64k + 64)` and the next 64
- hold the matching gate rows, which sit `I` rows later in the `[up ; gate]` stack."""
- inter = _TEXT["moe_intermediate_size"]
- core = SimpleNamespace(trunk_moe_backend="CUTEDSL")
- perm = W._fc1_perm(2 * inter, "cpu", core)
- assert sorted(perm.tolist()) == list(range(2 * inter))
- for k in range(inter // 64):
- assert perm[128 * k : 128 * k + 64].tolist() == list(range(64 * k, 64 * k + 64))
- assert perm[128 * k + 64 : 128 * k + 128].tolist() == list(
- range(inter + 64 * k, inter + 64 * k + 64)
- )
-
-
-def test_cutedsl_fc2_keeps_the_checkpoint_row_order_and_trtllm_gen_shuffles_it():
- """trtllm-gen's 32-row block shuffle moves source row `4u + v` to destination row
- `8v + u`; CuteDSL reads `down_proj` as the checkpoint stores it."""
- rows = torch.arange(64).unsqueeze(1).expand(64, 8).contiguous()
- kept = W._expert_fc2_w(rows, SimpleNamespace(trunk_moe_backend="CUTEDSL"))
- assert torch.equal(kept, rows)
- shuffled = W._expert_fc2_w(rows, SimpleNamespace(trunk_moe_backend="TRTLLM"))
- for u in range(8):
- for v in range(4):
- assert int(shuffled[8 * v + u, 0]) == 4 * u + v
- assert int(shuffled[32 + 8 * v + u, 0]) == 32 + 4 * u + v
-
-
-@pytest.mark.parametrize("backend, prefill_copy", [("TRTLLM", True), ("CUTEDSL", False)])
-def test_cute_dsl_expert_copy_rows_only_when_kept(backend, prefill_copy):
- """With the second expert copy (trtllm-gen, a step able to reach the CuteDSL dispatch
- threshold) every layer also carries its experts in the CuteDSL layout (`fc*_cd`);
- without it, and under CuteDSL whose one copy is that layout, no `_cd` row is declared."""
- manifest = W.MODEL_WEIGHTS.manifest(_core(False, backend, prefill_copy))
- rows = [k for k in manifest if k.endswith(("_fc1_w_cd", "_fc1_s_cd", "_fc2_w_cd", "_fc2_s_cd"))]
- layers = _TEXT["num_hidden_layers"]
- assert len(rows) == (4 * layers if prefill_copy else 0), len(rows)
-
-
-@pytest.mark.parametrize(
- "backend, max_num_tokens, env, kept",
- [
- ("TRTLLM", 4096, None, True),
- ("TRTLLM", 4095, None, False),
- ("TRTLLM", 16384, "0", False),
- ("CUTEDSL", 16384, None, False),
- ],
-)
-def test_second_expert_copy_kept_only_where_its_launches_can_happen(
- monkeypatch, backend, max_num_tokens, env, kept
-):
- """The copy costs about 1.3 GiB a layer: it is kept only under trtllm-gen, only when a
- step's rows can reach the CuteDSL dispatch threshold, and not when the env var is off."""
- if env is None:
- monkeypatch.delenv(M._MOE_PREFILL_COPY_ENV, raising=False)
- else:
- monkeypatch.setenv(M._MOE_PREFILL_COPY_ENV, env)
- assert M._moe_prefill_copy(backend, max_num_tokens) is kept
diff --git a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fp4_block_scale_moe_runner.py b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fp4_block_scale_moe_runner.py
index a7a802d8c6d9..13dee2f51c6d 100644
--- a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fp4_block_scale_moe_runner.py
+++ b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fp4_block_scale_moe_runner.py
@@ -1925,7 +1925,7 @@ def _cell_layer(spec: dict, layer_factory: Callable[[], tuple[dict, dict]]) -> t
),
),
),
- "qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1": (
+ "qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default": (
Cell(
why="qwen4_exp decode: one token over 512 experts, top-10 renormalized, H 2560 / I 640",
spec=dict(tokens=1, **_QWEN4_EXP),
diff --git a/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py b/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py
index 105296323d6c..cb794c987b0a 100644
--- a/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py
+++ b/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py
@@ -50,7 +50,7 @@ def _unpack(packed):
spec=dict(shape=(16, 256), dtype=torch.bfloat16, vec=16, swizzled=True),
),
),
- "qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1": (
+ "qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default": (
Cell(
why="qwen3.8 decode: 1 token, hidden 2560, linear -- its routed-expert input",
spec=dict(shape=(1, 2560), dtype=torch.bfloat16, vec=16, swizzled=False),
diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_reference_precision.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_reference_precision.py
index f070ddd2f96e..4218ce9bb5df 100644
--- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_reference_precision.py
+++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_reference_precision.py
@@ -151,12 +151,9 @@ def test_moe_reference_does_not_override_precision(relative_path: str) -> None:
"moe/test_modeling_v2_fp4_block_scale_moe_runner.py",
"moe/test_modeling_v2_mxe4m3_mxe2m1_block_scale_moe_runner.py",
"models/deepseek_v3/r1_0528_nvfp4__sm_10x__default/moe/test_modeling_v2_fused_moe.py",
- "models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/"
- "test_modeling_v2_fp4_block_scale_moe_renorm_routed.py",
- "models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/"
- "test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py",
- "models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/"
- "test_modeling_v2_fp4_moe_renorm_cute_dsl.py",
+ "models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py",
+ "models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py",
+ "models/qwen4_exp/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py",
],
)
def test_moe_test_reference_calls_scope_precision(relative_path: str) -> None:
diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py
index 0c592a14603c..c23aa0c587e7 100644
--- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py
+++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py
@@ -1152,7 +1152,7 @@ def test_kimi_k3_on_another_sm_does_not_match(monkeypatch):
# --- qwen4_exp: one-model MTP-Eagle over a linear chain, one rank ---------------
-_QWEN4_EXP_TARGET = "models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.modeling."
+_QWEN4_EXP_TARGET = "models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.modeling."
@pytest.fixture
@@ -1270,7 +1270,7 @@ def test_qwen4_exp_other_kv_cache_algorithms_are_declined(algo):
_model_config(_qwen4_exp_config()), _llm_args("auto")
)
with pytest.raises(NoTarget, match="KV cache"):
- qwen4_exp_routing._check_engine(None, algo, choices)
+ qwen4_exp_routing._check_engine(None, algo, choices, "tp1")
@pytest.mark.usefixtures("_qwen4_exp_sm")
@@ -1319,10 +1319,10 @@ def test_qwen4_exp_runtime_draft_length_falls_back(spec_kwargs):
@pytest.mark.usefixtures("_qwen4_exp_sm")
-@pytest.mark.parametrize("max_draft_len", [4, 7])
-def test_qwen4_exp_draft_beyond_three_needs_the_cached_replay(monkeypatch, max_draft_len):
- """The legacy per-step recurrent states are certified at up to three drafted tokens and
- hold one fp32 state per verify row."""
+@pytest.mark.parametrize("max_draft_len", [1, 3, 7])
+def test_qwen4_exp_mtp_needs_the_cached_replay(monkeypatch, max_draft_len):
+ """The target verifies on the GDN cached replay only, never on the legacy per-step
+ recurrent states."""
monkeypatch.setenv("TRTLLM_USE_GDN_REPLAY", "0")
config = _model_config(
_qwen4_exp_config(), spec_config=MTPDecodingConfig(max_draft_len=max_draft_len)
@@ -1333,13 +1333,6 @@ def test_qwen4_exp_draft_beyond_three_needs_the_cached_replay(monkeypatch, max_d
_claimed(config, dict(kv_cache_config=kv_cache_config))
-@pytest.mark.usefixtures("_qwen4_exp_sm")
-def test_qwen4_exp_legacy_verify_matches_up_to_three_drafted_tokens(monkeypatch):
- monkeypatch.setenv("TRTLLM_USE_GDN_REPLAY", "0")
- config = _model_config(_qwen4_exp_config(), spec_config=MTPDecodingConfig(max_draft_len=3))
- _claimed(config, dict(kv_cache_config=KvCacheConfig(mamba_ssm_cache_dtype="float32")))
-
-
@pytest.mark.usefixtures("_qwen4_exp_sm")
@pytest.mark.parametrize("max_draft_len", [4, 7])
def test_qwen4_exp_cutedsl_declines_drafts_beyond_three(monkeypatch, max_draft_len):
@@ -1396,10 +1389,9 @@ def test_qwen4_exp_mtp_on_a_float16_state_cache_declines_with_replay_off(monkeyp
@pytest.mark.usefixtures("_qwen4_exp_sm")
-def test_qwen4_exp_mtp_on_a_bf16_state_cache_matches_with_replay_off(monkeypatch):
+def test_qwen4_exp_without_speculation_matches_with_replay_off(monkeypatch):
monkeypatch.setenv("TRTLLM_USE_GDN_REPLAY", "0")
- config = _model_config(_qwen4_exp_config(), spec_config=MTPDecodingConfig(max_draft_len=3))
- _claimed(config)
+ _claimed(_model_config(_qwen4_exp_config()))
@pytest.mark.usefixtures("_qwen4_exp_sm")
From 49b94395f7ce516b234173f6230f17336e269e43 Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Wed, 4 Nov 2026 04:27:26 -0800
Subject: [PATCH 03/15] [None][test] modeling_v2 qwen4_exp: restore op guard
and CUDA-graph replay tests
The family op tests again cover what each op's note states:
- CUDA-graph capture and replay: gated_residual, hc_up_gate_mean,
hc_up_gate_mean_fp4__moe, ple_ngram_embed (16 x 1 rows over 461 slots,
128 x 1 and 128 x 8 over 1874), fp8_quantize_1x128 (64, 128, 1024 rows),
qsa_expand_block_indices (1, 128, 256, 1024 rows), qsa_paged_index_store,
qsa_paged_kv_store_nvfp4, qsa_topk_select (every GRAPH_CASES entry) and
qsa_paged_sparse_gqa_nvfp4_verify (with its determinism check);
- guards and loud rejections: malformed operands, a contiguous up projection,
a short query_start_loc, a strided fp8 input, a mismatched selection width,
a whole multi-slot count buffer, unequal metadata lengths, strided
request_indices, the token-budget default, a position past the rope table,
pages_to_bf16's dtype and shape checks, unexpanded verify selections, a
per-row fp4 global scale, and each silent table fault of the packed vision
rope (bf16, short, narrow, non-contiguous);
- the fp32 erf/tanh GELU discrimination tests.
The large paged-KV cells carry the registered high_cuda_memory marker instead
of an unregistered one.
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
.../activation/test_modeling_v2_gelu_erf.py | 50 +++++++++++-
.../activation/test_modeling_v2_gelu_tanh.py | 53 ++++++++++++-
...st_modeling_v2_qsa_expand_block_indices.py | 41 ++++++++++
.../test_modeling_v2_qsa_mtp_reuse_topk.py | 15 ++++
...est_modeling_v2_qsa_nvfp4_pages_to_bf16.py | 10 +++
.../test_modeling_v2_qsa_paged_index_store.py | 38 +++++++++
...st_modeling_v2_qsa_paged_kv_store_nvfp4.py | 78 ++++++++++++++++++-
...ng_v2_qsa_paged_sparse_gqa_nvfp4_decode.py | 2 +-
...ng_v2_qsa_paged_sparse_gqa_nvfp4_verify.py | 50 +++++++++++-
.../test_modeling_v2_qsa_topk_select.py | 51 ++++++++++++
...test_modeling_v2_qwen4_exp_qk_norm_rope.py | 8 ++
...st_modeling_v2_qwen4_exp_gated_residual.py | 32 +++++++-
...t_modeling_v2_qwen4_exp_hc_up_gate_mean.py | 27 ++++++-
.../ple/test_modeling_v2_ple_ngram_embed.py | 40 +++++++++-
.../test_modeling_v2_fp8_quantize_1x128.py | 40 +++++++++-
...st_modeling_v2_vision_rope_2d_packed_qk.py | 22 ++++++
...g_v2_qwen4_exp_hc_up_gate_mean_fp4__moe.py | 31 +++++++-
17 files changed, 575 insertions(+), 13 deletions(-)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_erf.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_erf.py
index 8e30429fc93f..ed7a1d95f5ac 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_erf.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_erf.py
@@ -4,7 +4,8 @@
The entry's identity is *which* GELU it is, and the wrong answer -- the tanh
GELU this same tower's block MLPs use -- rounds to within about one bf16 ulp of
-it; only fp32 tells the two apart.
+it. The discrimination test is the point of this file; only fp32 tells the two
+apart.
"""
import math
@@ -91,3 +92,50 @@ def _tanh64(x):
"""The *other* GELU -- `gelu_pytorch_tanh`, this tower's block-MLP activation."""
x64 = x.double()
return 0.5 * x64 * (1.0 + torch.tanh(SQRT_2_OVER_PI * (x64 + 0.044715 * x64 * x64 * x64)))
+
+
+def test_identity_is_the_erf_form_and_not_tanh() -> None:
+ """Measured in fp32, where the question has an answer; and shown to have
+ none in bf16, so a bf16 match cannot quietly be read as confirmation."""
+ x32 = _uniform((512, MERGER_WIDTH), -4.0, 4.0, torch.float32, 24)
+ out = op(x32).double()
+ ref = _erf64(x32)
+ correct = (out - ref).abs().max().item()
+ gate = _ATOL + _RTOL[torch.float32] * ref.abs().max().item()
+
+ # The entry's own gate refuses the tanh form in fp32.
+ with pytest.raises(AssertionError):
+ torch.testing.assert_close(out.float(), _tanh64(x32).float())
+
+ x64 = x32.double()
+ variants = {
+ "tanh gelu (gelu_pytorch_tanh)": _tanh64(x32),
+ "sigmoid approx x*sigmoid(1.702x)": x64 * torch.sigmoid(1.702 * x64),
+ "silu": x64 * torch.sigmoid(x64),
+ "relu": torch.clamp(x64, min=0.0),
+ }
+ print(f"\n fp32: gate {gate:.3e}; correct-path deviation {correct:.3e}")
+ for name, wrong in variants.items():
+ dev = (out - wrong).abs().max().item()
+ print(
+ f" {name:36s} dev {dev:.6e} ({dev / correct:9.1f}x correct, {dev / gate:7.1f}x gate)"
+ )
+ assert dev > 100.0 * correct, (name, dev, correct)
+ assert dev > 10.0 * gate, (name, dev, gate)
+
+ # erf's scale truncated from 1/sqrt(2) to 1/1.414 -- a plausible
+ # transcription slip, weak by construction, so a lower multiple.
+ perturbed = 0.5 * x64 * (1.0 + torch.erf(x64 / 1.414))
+ dev = (out - perturbed).abs().max().item()
+ print(f" {'erf scale 1/sqrt(2) -> 1/1.414':36s} dev {dev:.6e} ({dev / gate:7.1f}x gate)")
+ assert dev > 20.0 * correct, (dev, correct)
+ assert dev > 1.5 * gate, (dev, gate)
+
+ # bf16 cannot see the difference.
+ xb = x32.to(torch.bfloat16)
+ outb = op(xb).double()
+ correct_b = (outb - _erf64(xb)).abs().max().item()
+ tanh_b = (outb - _tanh64(xb)).abs().max().item()
+ print(f" bf16: correct {correct_b:.6e}; tanh variant {tanh_b:.6e} ({tanh_b / correct_b:.3f}x)")
+ assert tanh_b < 4.0 * correct_b, (tanh_b, correct_b)
+ torch.testing.assert_close(outb.to(torch.bfloat16), op.reference(xb))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_tanh.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_tanh.py
index 14b83e9943ef..4f818110d2f6 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_tanh.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_tanh.py
@@ -2,8 +2,9 @@
# SPDX-License-Identifier: Apache-2.0
"""GPU test for the qwen4_exp gelu_tanh entry.
-The entry's identity is *which* GELU it is: in bf16 the tanh and erf GELUs
-round to within about one ulp of each other, and only fp32 tells them apart.
+The entry's identity is *which* GELU it is, so the discrimination test is the
+point of this file: in bf16 the tanh and erf GELUs round to within about one
+ulp of each other, and only fp32 tells them apart.
"""
import math
@@ -82,3 +83,51 @@ def _erf64(x):
"""The *other* GELU -- exact, erf-based; `torch.nn.GELU()` and `moe_gelu`."""
x64 = x.double()
return 0.5 * x64 * (1.0 + torch.erf(x64 / math.sqrt(2.0)))
+
+
+def test_identity_is_the_tanh_form_and_not_erf() -> None:
+ """Measured in fp32, where the question has an answer; and shown to have
+ none in bf16, so a bf16 match cannot quietly be read as confirmation."""
+ x32 = _uniform((512, INTERMEDIATE_SIZE), -4.0, 4.0, torch.float32, 24)
+ out = op(x32).double()
+ ref = _tanh64(x32)
+ correct = (out - ref).abs().max().item()
+ gate = _ATOL + _RTOL[torch.float32] * ref.abs().max().item()
+
+ # The entry's own gate refuses the erf form in fp32.
+ with pytest.raises(AssertionError):
+ torch.testing.assert_close(out.float(), _erf64(x32).float())
+
+ x64 = x32.double()
+ variants = {
+ "erf gelu (nn.GELU, and moe_gelu)": _erf64(x32),
+ "sigmoid approx x*sigmoid(1.702x)": x64 * torch.sigmoid(1.702 * x64),
+ "silu": x64 * torch.sigmoid(x64),
+ "relu": torch.clamp(x64, min=0.0),
+ }
+ print(f"\n fp32: gate {gate:.3e}; correct-path deviation {correct:.3e}")
+ for name, wrong in variants.items():
+ dev = (out - wrong).abs().max().item()
+ print(
+ f" {name:36s} dev {dev:.6e} ({dev / correct:9.1f}x correct, {dev / gate:7.1f}x gate)"
+ )
+ assert dev > 100.0 * correct, (name, dev, correct)
+ assert dev > 10.0 * gate, (name, dev, gate)
+
+ # A perturbation of this same formula: weak by construction (bounded by the
+ # coefficient change times the tanh derivative), so a lower multiple.
+ perturbed = 0.5 * x64 * (1.0 + torch.tanh(SQRT_2_OVER_PI * (x64 + 0.045 * x64**3)))
+ dev = (out - perturbed).abs().max().item()
+ print(f" {'coefficient 0.044715 -> 0.045':36s} dev {dev:.6e} ({dev / gate:7.1f}x gate)")
+ assert dev > 100.0 * correct, (dev, correct)
+ assert dev > 5.0 * gate, (dev, gate)
+
+ # bf16 cannot see the difference: the erf variant's max deviation is within
+ # a small multiple of the correct path's own.
+ xb = x32.to(torch.bfloat16)
+ outb = op(xb).double()
+ correct_b = (outb - _tanh64(xb)).abs().max().item()
+ erf_b = (outb - _erf64(xb)).abs().max().item()
+ print(f" bf16: correct {correct_b:.6e}; erf variant {erf_b:.6e} ({erf_b / correct_b:.3f}x)")
+ assert erf_b < 4.0 * correct_b, (erf_b, correct_b)
+ torch.testing.assert_close(outb.to(torch.bfloat16), op.reference(xb))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_expand_block_indices.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_expand_block_indices.py
index b2e01fa46b19..f45b9470a169 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_expand_block_indices.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_expand_block_indices.py
@@ -132,3 +132,44 @@ def test_certified_cells(cell, certifying): # noqa: F811
assert bool((lengths.long() > positions.long() + 1).any()), (
"a context cell must drive sequence_lengths past position + 1"
)
+
+
+@pytest.mark.parametrize("rows", [1, 128, 256, 1024])
+def test_cuda_graph_capture_and_replay(rows) -> None:
+ """Capturable at decode and verify row counts (both launch shapes), and a replay
+ recomputes from whatever the buffers now hold."""
+ positions = torch.tensor([4096 + 7 * i for i in range(rows)], device=DEV)
+ lengths = (positions + 1).to(torch.int32)
+ blocks = _packed_selection(positions, CR, TOPK, seed=31)
+
+ def call():
+ return op(blocks, positions, lengths, compress_ratio=CR, token_topk=TOPK)
+
+ eager = call().clone()
+ warmup = torch.cuda.Stream()
+ warmup.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(warmup):
+ for _ in range(3):
+ call()
+ torch.cuda.current_stream().wait_stream(warmup)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ captured = call()
+ graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(captured, eager)
+
+ blocks.copy_(_packed_selection(positions, CR, TOPK, seed=32))
+ graph.replay()
+ torch.cuda.synchronize()
+ _compare(captured, op.reference(blocks, positions, lengths, compress_ratio=CR, token_topk=TOPK))
+
+
+def test_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)
+ 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/attention/test_modeling_v2_qsa_mtp_reuse_topk.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_mtp_reuse_topk.py
index da6de8790424..d689bfd91d89 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_mtp_reuse_topk.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_mtp_reuse_topk.py
@@ -253,3 +253,18 @@ def _guard_case():
whole[2] = kwargs["shared_visible_blocks"]
kwargs["shared_visible_blocks"] = whole
return intended, kwargs
+
+
+def test_guard_refuses_a_whole_multi_slot_count_buffer() -> None:
+ """Unguarded the op reads slot 0's counts and returns a plausible wrong answer;
+ `is_valid` refuses it."""
+ intended, kwargs = _guard_case()
+ out = op(**kwargs)
+ torch.cuda.synchronize()
+ assert not torch.equal(out, intended), "the op got it right after all"
+ _, kwargs = _guard_case()
+ with (
+ validating(op),
+ pytest.raises(ValueError, match="shared_visible_blocks must be one slot's"),
+ ):
+ op(**kwargs)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_nvfp4_pages_to_bf16.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_nvfp4_pages_to_bf16.py
index 61b73119dcc7..8adaf4bdce6e 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_nvfp4_pages_to_bf16.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_nvfp4_pages_to_bf16.py
@@ -149,3 +149,13 @@ def test_listing_past_the_scratch_is_dropped() -> None:
out = op(**args)
for got, want in zip(out, ref):
assert _bitwise(got, want)
+
+
+def test_guards() -> None:
+ args = _build(op.CELLS[0].spec, 3)
+ with validating(op), pytest.raises(AssertionError):
+ op(**{**args, "k_out": args["k_out"].float()})
+ with validating(op), pytest.raises(AssertionError):
+ op(**{**args, "block_table": args["block_table"].long()})
+ with validating(op), pytest.raises(AssertionError):
+ op(**{**args, "table": args["table"][:, :-1]})
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_index_store.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_index_store.py
index 8d840ae68ecb..3852fbcc9523 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_index_store.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_index_store.py
@@ -298,3 +298,41 @@ def test_the_stock_managers_planes_are_the_cells_layout(member) -> None:
for j, plane in enumerate(planes):
if j != cell["layer"]:
assert bool((_bits(plane) == fill).all()), f"sparse layer {j}'s plane moved"
+
+
+def test_cuda_graph_capture_and_replay() -> None:
+ """Captured once at the 32 MTP7 verify blocks (256 rows), replayed against new
+ positions written into the same buffer."""
+ args = _small(seed=52, **_MTP7_BATCH)
+ positions = args["logical_positions"]
+ replayed = positions.clone()
+ positions.zero_()
+ warmup = torch.cuda.Stream()
+ warmup.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(warmup):
+ for _ in range(3):
+ op(**args)
+ torch.cuda.current_stream().wait_stream(warmup)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ op(**args)
+
+ positions.copy_(replayed)
+ args["index_cache"].fill_(KEY_FILL)
+ args["position_cache"].fill_(POS_FILL)
+ ref = op.reference(**args)
+ graph.replay()
+ torch.cuda.synchronize()
+ _compare((args["index_cache"], args["position_cache"]), ref)
+
+
+def test_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]
+ before = (args["index_cache"].clone(), args["position_cache"].clone())
+ with pytest.raises(ValueError, match="equally long"):
+ op(**args)
+ assert torch.equal(_bits(args["index_cache"]), _bits(before[0]))
+ assert torch.equal(args["position_cache"], before[1])
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py
index 65066fdaf510..a03f5821bdfd 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py
@@ -295,9 +295,10 @@ def _assert_straddles(spec, args) -> None:
assert (spec["pages"] - 1) * args[name].stride(0) >= 2**31, (name, args[name].stride())
-# A pool sized to a member's measured page count does not fit a normal cell's 40 GiB.
+# A pool sized to a member's measured page count takes 47 GB of device memory.
_CELL_PARAMS = [
- pytest.param(c, marks=pytest.mark.large) if c.spec.get("large") else c for c in op.CELLS
+ pytest.param(c, marks=pytest.mark.high_cuda_memory) if c.spec.get("large") else c
+ for c in op.CELLS
]
@@ -323,3 +324,76 @@ def test_certified_cells(cell, certifying): # noqa: F811
def _decode_args(seed: int, **overrides):
return _build({**op.CELLS[0].spec, **overrides}, seed=seed)
+
+
+def _reset(planes) -> None:
+ for plane, fill in zip(planes, (DATA_FILL, DATA_FILL, SCALE_FILL, SCALE_FILL)):
+ plane.fill_(fill)
+
+
+def _graph_replay(spec, seed: int, graph_batches) -> None:
+ """Each G: capture one step of the first G requests, replay the next step.
+
+ The next step lives in the captured buffers: new k/v rows, new pages in the
+ table, and every request's positions advanced by its rows, the way a decode
+ step moves. The replay must equal the reference and an eager call on the
+ same inputs, bit for bit. One pool serves the whole sweep, as the engine's does.
+ """
+ pool = _pool(spec)
+ steps = {n for _, n in spec["segments"]}
+ assert len(steps) == 1, "a graph cell's requests all carry the same row count"
+ step = steps.pop()
+ for g in graph_batches:
+ sub = {**spec, "segments": spec["segments"][:g]}
+ args, _, _ = _build(sub, seed=seed + g, pool=pool)
+ following, _, _ = _build(sub, seed=seed + 1000 + g, pool=pool)
+ planes = _planes(args)
+ positions = args["logical_positions"]
+ next_positions = positions.clone()
+ positions.sub_(step)
+ assert int(positions.min()) >= 0
+
+ def call():
+ op(**args)
+
+ for _ in range(3): # compile the Triton kernel before capture
+ call()
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ call()
+
+ for name in ("k", "v", "block_table"):
+ args[name].copy_(following[name])
+ positions.copy_(next_positions)
+ _reset(planes)
+ op(**args)
+ eager = tuple(plane.clone() for plane in planes)
+ _reset(planes)
+ ref = op.reference(**args)
+ graph.replay()
+ torch.cuda.synchronize()
+ _compare(planes, ref)
+ assert all(torch.equal(p, e) for p, e in zip(planes, eager)), (
+ f"G={g} x {step} rows: replay differs from eager"
+ )
+ del graph
+ _assert_other_roles_untouched(spec, pool)
+
+
+def test_cuda_graph_replay_matches_eager() -> None:
+ """Decode MTP under a graph: 8 requests x 8 rows in the prefill cell's trunk layout,
+ captured once, replayed on the next step's rows, pages and positions."""
+ 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 test_op_rejects_strided_request_indices() -> None:
+ """The callable's own loud error -- why the entry carries no `is_valid` -- for the
+ metadata a flat-pointer kernel would otherwise read at the wrong rows."""
+ args, _, _ = _decode_args(seed=5)
+ strided = torch.stack([args["request_indices"]] * 2, 1)[:, 0]
+ with pytest.raises(ValueError, match="request_indices must be contiguous"):
+ with validating(op):
+ op(**{**args, "request_indices": strided})
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_decode.py
index 93d998a93d4e..354a8469963d 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_decode.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_decode.py
@@ -227,7 +227,7 @@ def _replay_in_place(args, seed) -> None:
def _marks(cell):
- return (pytest.mark.large, pytest.mark.high_cuda_memory) if cell.spec.get("large") else ()
+ return (pytest.mark.high_cuda_memory,) if cell.spec.get("large") else ()
@pytest.mark.parametrize(
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py
index d13d08a9e60a..bbe42f210b46 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py
@@ -4,7 +4,8 @@
`qsa_paged_sparse_gqa_nvfp4_decode` and against the per-row kernel, inside the NVFP4 op's band
(4 ulp of each (row, head)'s max, 2 ulp relative RMS), over its cells; plus the routing to the
per-row kernel, forced splits and other rows a request, the launch plan (a partial last wave's
-units split and merged), invalid request rows, and an unaligned row slice of the selection."""
+units split and merged), invalid request rows, an unaligned row slice of the selection, the
+selections `is_valid` refuses, determinism and graph replay against eager."""
import math
import zlib
@@ -266,3 +267,50 @@ def test_unaligned_row_slice_of_the_selection() -> None:
aligned[name] = args[name].clone()
assert args["selected_tokens"].data_ptr() % 16
assert torch.equal(op(**args, rows_per_request=4), op(**aligned, rows_per_request=4))
+
+
+def test_is_valid_refuses_selections_that_are_not_expanded() -> None:
+ args = build(_segments(8, 4), seed=7)
+ op.is_valid(**args, rows_per_request=4)
+ good = args["selected_tokens"]
+ shifted = good.clone()
+ shifted[0, 0] += 1 # off its column's offset within the group
+ mixed = good.clone()
+ mixed[0, 1] = good[0, 5] # two groups in one quad
+ repeated = good.clone()
+ repeated[0, 4:8] = good[0, 0:4] # one group twice in a row
+ for bad in (shifted, mixed, repeated):
+ args["selected_tokens"] = bad
+ with pytest.raises(AssertionError):
+ op.is_valid(**args, rows_per_request=4)
+
+
+def test_calls_are_bitwise_deterministic() -> None:
+ args = build(_segments(8, 4), seed=8800)
+ assert torch.equal(op(**args, rows_per_request=4), op(**args, rows_per_request=4))
+
+
+@pytest.mark.parametrize("rows", [32])
+def test_replay_matches_eager(rows) -> None:
+ segments = tuple((100 + (499 * g) % 8000, 4) for g in range(rows // 4))
+ args = build(segments, seed=77)
+ fresh = build(segments, seed=127)
+ op(**args, rows_per_request=4)
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ captured = op(**args, rows_per_request=4)
+ for name in (
+ "q",
+ "k_cache",
+ "v_cache",
+ "k_scale_cache",
+ "v_scale_cache",
+ "block_table",
+ "selected_tokens",
+ "request_indices",
+ ):
+ args[name].copy_(fresh[name])
+ graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(captured, op(**fresh, rows_per_request=4))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_select.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_select.py
index b13321c3b687..ecb3dfe2c753 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_select.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_select.py
@@ -103,3 +103,54 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811
#: histogram path); the replays move the bound, 128 rows down to the budget's
#: copy-through shortcut.
GRAPH_CASES = ((256, 2048, 1500), (1, 65536, 65535), (128, 65536, 511), (1024, 65536, 30000))
+
+
+def _layout_case():
+ torch.manual_seed(10)
+ rows, cols, top_k = 4, 128, 8
+ logits = torch.randn(rows, cols, device=DEV)
+ row_starts = torch.zeros(rows, dtype=torch.int32, device=DEV)
+ row_ends = torch.full((rows,), cols, dtype=torch.int32, device=DEV)
+ return logits, row_starts, row_ends, top_k
+
+
+@pytest.mark.parametrize(
+ "rows, cols, window", GRAPH_CASES, ids=[f"{r}x{c}-to-{w}" for r, c, w in GRAPH_CASES]
+)
+def test_cuda_graph_capture_and_replay(rows, cols, window) -> None:
+ """Captured over full windows, then replayed against new scores and a new causal
+ bound through the captured addresses."""
+ gen = torch.Generator(device=DEV).manual_seed(8)
+ logits = _logits("relu", rows, cols, gen)
+ row_starts = torch.zeros(rows, dtype=torch.int32, device=DEV)
+ row_ends = torch.full((rows,), cols, dtype=torch.int32, device=DEV)
+ indices = torch.empty(rows, BLOCK_TOPK, dtype=torch.int32, device=DEV)
+
+ warmup = torch.cuda.Stream()
+ warmup.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(warmup):
+ op(logits, row_starts, row_ends, indices, BLOCK_TOPK)
+ torch.cuda.current_stream().wait_stream(warmup)
+ torch.cuda.synchronize()
+
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ op(logits, row_starts, row_ends, indices, BLOCK_TOPK)
+
+ logits.copy_(_logits("relu", rows, cols, gen))
+ row_ends.fill_(window)
+ indices.fill_(SENTINEL)
+ graph.replay()
+ torch.cuda.synchronize()
+ _compare(indices, op.reference(logits, row_starts, row_ends, indices, BLOCK_TOPK))
+
+
+def test_guard_refuses_the_token_budget_default() -> None:
+ """The default `index_topk` is 2048, the token budget; the target's buffer
+ is 512 wide. That is the wider-buffer fault the other way round, and the
+ guard catches it at the call that forgot to pass the block budget."""
+ logits, row_starts, row_ends, _ = _layout_case()
+ indices = torch.empty(logits.shape[0], BLOCK_TOPK, dtype=torch.int32, device=DEV)
+ with pytest.raises(AssertionError, match="index_topk=2048"):
+ with validating(op):
+ op(logits, row_starts, row_ends, indices)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py
index 306742125fc6..0674746fabec 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py
@@ -152,3 +152,11 @@ def test_cuda_graph_replay_reads_fresh_inputs() -> None:
torch.cuda.synchronize()
assert torch.equal(proj, eager_proj)
assert torch.equal(kwargs["index_q"], eager_kwargs["index_q"])
+
+
+def test_guard_refuses_a_position_past_the_table() -> None:
+ """An out-of-range position would read past the rope table."""
+ proj, _, kwargs = _build(dict(tokens=3, **CKPT), seed=_seed("guard"))
+ kwargs["position_ids"][0, 0] = TABLE_ROWS
+ with validating(op), pytest.raises(AssertionError, match="past the rope table"):
+ op(proj, **kwargs)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_gated_residual.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_gated_residual.py
index a5f97c64c529..155544922e6b 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_gated_residual.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_gated_residual.py
@@ -5,7 +5,8 @@
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; the cells' band sits 32x..192x
-inside the nearest wrong formula (measured).
+inside the nearest wrong formula (measured). Beyond the cells: graph replay and
+one loud misuse.
"""
import zlib
@@ -124,3 +125,32 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res
assert inject.shape == (*spec["leading"], hc_count)
assert torch.equal(args[0], before), "the normed stream is read-only"
_compare((mixed, inject), op.reference(*args))
+
+
+def test_malformed_operands_raise_rather_than_return() -> None:
+ """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)
+
+
+def test_cuda_graph_capture_and_replay() -> None:
+ _, (static_in, down, up, inject_w, c) = _build((8,), seed=8)
+ warmup = torch.cuda.Stream()
+ warmup.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(warmup):
+ for _ in range(3):
+ op(static_in, down, up, inject_w, c)
+ torch.cuda.current_stream().wait_stream(warmup)
+
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ graph_mixed, graph_inject = op(static_in, down, up, inject_w, c)
+
+ _, (fresh, *_) = _build((8,), seed=18)
+ static_in.copy_(fresh)
+ graph.replay()
+ torch.cuda.synchronize()
+ eager_mixed, eager_inject = op(fresh, down, up, inject_w, c)
+ assert torch.equal(graph_mixed, eager_mixed)
+ assert torch.equal(graph_inject, eager_inject)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py
index 55e12cf73622..458a7dbb0390 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py
@@ -1,7 +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`."""
+`qwen4_exp_gated_residual`'s test derives for `mixed`, the guard, graph
+capture."""
import zlib
@@ -65,3 +66,27 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by
# The answer came from the CuTe DSL kernel, JIT-compiled for this device.
assert (_BLOCK_H, args[3], mixed.shape[1], args[0].shape[1], False) in _COMPILED
_compare(mixed, op.reference(*args))
+
+
+def test_guard_refuses_a_contiguous_up() -> None:
+ low, up, normed, c = _build(dict(rows=4), seed=4)
+ with validating(op), pytest.raises(AssertionError):
+ op(low, up.contiguous(), normed, c)
+
+
+def test_cuda_graph_capture_and_replay() -> None:
+ low, up, normed, c = _build(dict(rows=8), seed=5)
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ op(low, up, normed, c)
+ torch.cuda.current_stream().wait_stream(side)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ g_mixed = op(low, up, normed, c)
+ fresh = _build(dict(rows=8), seed=6)
+ low.copy_(fresh[0])
+ normed.copy_(fresh[2])
+ graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(g_mixed, op(fresh[0], up, fresh[2], c))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_ngram_embed.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_ngram_embed.py
index 7eb14dc43f01..aad5d62cc216 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_ngram_embed.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_ngram_embed.py
@@ -5,7 +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.
+float64 gather-and-scale. Beyond the cells: the `query_start_loc` length guard
+and graph capture at the note's replay cells.
"""
import zlib
@@ -29,6 +30,8 @@
CKPT_HEAD_DIM,
CKPT_SCALE,
CKPT_TABLE_ROWS,
+ ENGINE_SLOTS,
+ TEP4_SLOTS,
)
from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ple_ngram_embed import ( # noqa: E501
ple_ngram_embed as op,
@@ -179,3 +182,38 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by
_compare((out, args[1]), ref)
if cell.spec.get("table") == "host":
assert args[1].data_ptr() != 0 and not torch.equal(args[1], pool_before)
+
+
+def test_guard_refuses_short_query_start_loc() -> None:
+ args = _build(dict(lengths=(1, 1, 1)), seed=3)
+ args[4] = args[4][:-1]
+ with validating(op), pytest.raises(AssertionError):
+ op(*args)
+
+
+@pytest.mark.parametrize(
+ "lengths, pool_slots",
+ [((1,) * 16, ENGINE_SLOTS), ((1,) * 128, TEP4_SLOTS), ((8,) * 128, TEP4_SLOTS)],
+ ids=["16x1-461", "128x1-1874", "128x8-1874"],
+)
+def test_cuda_graph_capture_and_replay(lengths, pool_slots) -> None:
+ """The note's replay cells: a replay reads the ids and pool written since capture."""
+ spec = dict(lengths=lengths, pool_slots=pool_slots, eos_density=0.1)
+ args = _build(spec, seed=11)
+ warmup = torch.cuda.Stream()
+ warmup.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(warmup):
+ op(*args)
+ torch.cuda.current_stream().wait_stream(warmup)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ graph_out = op(*args)
+ fresh = _build(spec, 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/quantization/test_modeling_v2_fp8_quantize_1x128.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/quantization/test_modeling_v2_fp8_quantize_1x128.py
index 0a502ed944fd..719270a41cd4 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/quantization/test_modeling_v2_fp8_quantize_1x128.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/quantization/test_modeling_v2_fp8_quantize_1x128.py
@@ -2,8 +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
-every code disagreement of a cell to a rounding tie.
+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
@@ -128,3 +130,37 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th
# ── what a cell cannot carry ──────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("rows", [64, 128, 1024])
+def test_cuda_graph_capture_and_replay(rows):
+ """A replay reproduces a fresh eager call bitwise; so does a side stream."""
+ x = _input(rows, 2560, seed=17, magnitude=1.0)
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ for _ in range(3):
+ op(x)
+ torch.cuda.current_stream().wait_stream(side)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ g_data, g_sf = op(x)
+ x.copy_(_input(rows, 2560, seed=18, magnitude=1.0))
+ e_data, e_sf = op(x)
+ graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(g_data.view(torch.uint8), e_data.view(torch.uint8))
+ assert torch.equal(g_sf.view(torch.int32), e_sf.view(torch.int32))
+ alt = torch.cuda.Stream()
+ with torch.cuda.stream(alt):
+ a_data, a_sf = op(x)
+ alt.synchronize()
+ assert torch.equal(a_data.view(torch.uint8), e_data.view(torch.uint8))
+ assert torch.equal(a_sf.view(torch.int32), e_sf.view(torch.int32))
+
+
+def test_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__default/attention/test_modeling_v2_vision_rope_2d_packed_qk.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_vision_rope_2d_packed_qk.py
index e4e3949fd5b4..f47dc0ec1c37 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_vision_rope_2d_packed_qk.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_vision_rope_2d_packed_qk.py
@@ -172,3 +172,25 @@ def _small():
qkv = _packed_qkv(96, torch.bfloat16, seed=17)
cos, sin = _angle_table(grid_thw, MERGE, THETA, HEAD_DIM)
return qkv, cos, sin
+
+
+#: The silent faults the op's note lists, each a table the kernel accepts: (cos, sin) -> table.
+_FAULTS = {
+ "bf16 table": lambda t: t.to(torch.bfloat16),
+ "short table": lambda t: t[:-8],
+ "narrow table": lambda t: t[:, : HALF - 4].contiguous(),
+ "non-contiguous table": lambda t: torch.cat((t, t), dim=-1)[:, :HALF],
+}
+
+
+@pytest.mark.parametrize("fault", list(_FAULTS))
+def test_is_valid_refuses_each_silent_table_fault(fault):
+ """Unguarded, the kernel returns a wrong rotation for each fault; `is_valid` refuses it."""
+ qkv, cos, sin = _small()
+ want = op.reference(_qk(qkv), cos, sin)
+ bad_cos, bad_sin = _FAULTS[fault](cos), _FAULTS[fault](sin)
+ got = op(_qk(qkv.clone()), bad_cos, bad_sin)
+ torch.cuda.synchronize()
+ assert not torch.equal(got, want), f"{fault}: the kernel got it right after all"
+ with validating(op), pytest.raises(AssertionError):
+ op(_qk(qkv.clone()), bad_cos, bad_sin)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4__moe.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4__moe.py
index b25c216ac330..1d9bcf987b1d 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4__moe.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4__moe.py
@@ -3,7 +3,8 @@
"""GPU test for the qwen4 default target's qwen4_exp_hc_up_gate_mean_fp4__moe
entry: `mixed` in the band `qwen4_exp_hc_up_gate_mean`'s test uses and bitwise
that entry's own output, the NVFP4 outputs against the native-torch
-quantization of the returned `mixed` outside the reciprocal tie window."""
+quantization of the returned `mixed` outside the reciprocal tie window, the
+global-scale guard, graph capture."""
import zlib
@@ -89,3 +90,31 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by
# The same kernel and epilogue arithmetic as the bf16-only build the target runs elsewhere.
assert torch.equal(mixed, qwen4_exp_hc_up_gate_mean(*args[:4]))
_check_quant(mixed, xq, xsf, args[4])
+
+
+def test_guards() -> None:
+ """The kernel loads `global_scale[0]` only; a per-row scale is refused."""
+ low, up, normed, c, gs = _build(dict(rows=4), seed=4)
+ with validating(op), pytest.raises(AssertionError):
+ op(low, up, normed, c, torch.cat([gs, gs]))
+
+
+def test_cuda_graph_capture_and_replay() -> None:
+ low, up, normed, c, gs = _build(dict(rows=8), seed=5)
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ op(low, up, normed, c, gs)
+ torch.cuda.current_stream().wait_stream(side)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ g_out = op(low, up, normed, c, gs)
+ fresh = _build(dict(rows=8), seed=6)
+ low.copy_(fresh[0])
+ normed.copy_(fresh[2])
+ gs.mul_(0.5)
+ graph.replay()
+ torch.cuda.synchronize()
+ want = op(fresh[0], up, fresh[2], c, gs)
+ for got, exp in zip(g_out, want):
+ assert torch.equal(got, exp)
From db9657de1106c678685e51ff4d6febd13544b55c Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Wed, 4 Nov 2026 18:15:48 -0800
Subject: [PATCH 04/15] [None][fix] modeling_v2 qwen4_exp: keep the tp1 draft
head; decline an int32-overflowing score table
The post-load walk calls the core's post_load_weights after the shell's own, which builds
the tp1 FP8 draft head; the core's call reset it, so MTP at tp1 drafted on the verify or
the bf16 head and TRTLLM_MODELING_V2_MTP_HEAD_NVFP4 had no effect. The core now sets
draft_head only at construction, and the spec forward asserts it exists at tp1.
A prefill step scores and selects its context rows in one call, and the score and top-k
kernels index the score table in int32. Routing declines a max_num_tokens whose rows times
the table's columns reach 2^31 (at max_seq_len 262144: more than 32752 rows).
qwen4_exp_hc_decode_rows_tc.py gets its SPDX header.
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
.../ops/hc/qwen4_exp_hc_decode_rows_tc.py | 2 ++
.../modeling.py | 9 +++--
.../modeling_v2/models/qwen4_exp/routing.py | 36 +++++++++++++++++++
.../test_routing.py | 15 ++++++++
4 files changed, 60 insertions(+), 2 deletions(-)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_rows_tc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_rows_tc.py
index 39c27a557c7b..b923f1477630 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_rows_tc.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_rows_tc.py
@@ -1,3 +1,5 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
"""Decode hyper-connection site for 2..16 rows on warp tensor cores: two PDL launches, the
first emitting the bottleneck `low`, the second the up projection, gate, stream mean and NVFP4.
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py
index bb040f05c87b..ed7631b9540e 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py
@@ -1083,6 +1083,9 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig
# The draft experts' NVFP4 copy at one rank (`_mtp_moe_nvfp4`), built once: it frees
# the FP8 experts, and the shell's post-load walk reaches `post_load_weights` again.
self._mtp_moe_nvfp4 = None
+ # The FP8 draft head at tp1, set by the shell's `post_load_weights` after this core's;
+ # the post-load walk calls this core's `post_load_weights` again, which leaves it.
+ self.draft_head = None
# ---- PLE -----------------------------------------------------------
# One-based in the config.
@@ -1308,8 +1311,6 @@ def post_load_weights(self) -> None:
)
self._prefill = PrefillTarget(self)
self._decode = DecodeTarget(self)
- # The FP8 draft head at tp1, set by the shell's `post_load_weights`.
- self.draft_head = None
def _build_draft_one_rank(self, device) -> None:
"""At one rank the draft layer's operands for `DecodeTarget._mtp_layer`: the fused
@@ -6621,6 +6622,10 @@ def forward(
**kwargs,
)
assert spec_metadata is not None, "the spec worker is built but no spec_metadata came"
+ assert not self.one_rank or self.model.draft_head is not None, (
+ "MTP at tp1 runs its draft steps on the FP8 draft head `_build_heads` makes, but "
+ "the core has none"
+ )
final, hc = self.model.split_spec_output(
self.model(
attn_metadata=attn_metadata,
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 cc8135690654..4ad745b1d0e6 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
@@ -74,6 +74,15 @@
# The attention-DP layouts MTP is admitted at: those its draft acceptance was measured at.
_ATTENTION_DP_MTP_LAYOUTS = ("dep4",)
+# A prefill step scores and selects its context rows (at most max_num_tokens) in one call,
+# and `qsa_paged_index_scores` (output rows) and `qsa_topk_expand_decode` (row layouts) index
+# the score table in int32: rows x columns must stay below 2^31. The table has
+# `indexer_compress_ratio`-compressed columns over the KV manager's `max_blocks_per_seq`
+# pages: the sequence plus its generation lead (at most twice the widest verify step),
+# padded to a multiple of 4 pages.
+_SCORE_TABLE_ELEMENTS = 2**31
+_GENERATION_LEAD_TOKENS = 2 * (_MAX_DRAFT_LEN + 1)
+
class EngineChoices(NamedTuple):
"""Engine choices that fix which code a target runs but are not in the checkpoint's
@@ -203,6 +212,27 @@ def _check_engine(
)
+def _check_score_table(
+ max_num_tokens: int,
+ max_seq_len: Optional[int],
+ tokens_per_block: int,
+ compress_ratio: Optional[int],
+) -> None:
+ """Decline a prefill step whose score table would pass 2^31 elements (see
+ `_SCORE_TABLE_ELEMENTS`); an unknown max_seq_len or compression is not checked."""
+ if max_seq_len is None or compress_ratio is None:
+ return
+ pages = -(-(max_seq_len + _GENERATION_LEAD_TOKENS) // tokens_per_block)
+ columns = -(-pages // 4) * 4 * tokens_per_block // compress_ratio
+ if max_num_tokens * columns >= _SCORE_TABLE_ELEMENTS:
+ raise NoTarget(
+ f"max_num_tokens={max_num_tokens} at max_seq_len={max_seq_len}: a prefill step's "
+ f"context rows score a {columns}-column table, and the target indexes it in int32 "
+ f"up to 2^31 elements; set max_num_tokens to at most "
+ f"{(_SCORE_TABLE_ELEMENTS - 1) // columns}"
+ )
+
+
def _check_attention_dp(member: str, spec_config: Optional[Any], engine: EngineChoices) -> None:
"""Decline what the target does not run under attention DP (dep4, dep2), where
each rank serves its own requests and only the routed experts cross ranks.
@@ -295,6 +325,12 @@ def route(
if mapping.enable_attention_dp:
_check_attention_dp(member, spec_config, engine)
_check_engine(spec_config, kv_cache_quant_algo, engine, member)
+ _check_score_table(
+ llm_args.max_num_tokens or model_config.max_num_tokens,
+ llm_args.max_seq_len or getattr(text, "max_position_embeddings", None),
+ llm_args.kv_cache_config.tokens_per_block,
+ getattr(text, "indexer_compress_ratio", None),
+ )
# One line on purpose: the claims suite reads the routed target off the
# `from ..modeling import ` line.
# fmt: off
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py
index 5a6960713b27..9bb9303d1cd1 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py
@@ -385,6 +385,21 @@ def test_the_piecewise_prefill_cuda_graph_declines():
)
+@pytest.mark.parametrize("mapping_kwargs", [{}, _LAYOUTS["dep4"]], ids=["tp1", "dep4"])
+def test_a_prefill_step_past_the_int32_score_table_declines(mapping_kwargs):
+ """At the checkpoint's max_seq_len 262144 the score table is 65568 columns (8196 pages of
+ 32 tokens at compression 4): 32752 context rows stay under 2^31 elements, 32768 do not."""
+ config = _model_config(
+ pretrained_config=_qwen4_exp_config(
+ {"max_position_embeddings": 262144, "indexer_compress_ratio": 4}
+ ),
+ **mapping_kwargs,
+ )
+ _claimed(config, max_num_tokens=32752)
+ _declined(config, match="max_num_tokens=32768 at max_seq_len=262144", max_num_tokens=32768)
+ _claimed(config, max_num_tokens=32768, max_seq_len=131072)
+
+
@pytest.mark.parametrize("max_draft_len", [1, 3, 7])
def test_mtp_needs_the_cached_replay(monkeypatch, max_draft_len):
"""The target verifies on the GDN cached replay only; without speculation it is unused."""
From 52cc48603a0208ed439fdfe918c6e29856b7d4d9 Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Wed, 4 Nov 2026 19:36:55 -0800
Subject: [PATCH 05/15] [None][fix] modeling_v2 qwen4_exp: second-review fixes,
restored op tests, Modeling V2 docs
Target:
- A prefill step's context rows are scored and selected in chunks of fewer than 2^31 score
elements (`_score_and_select`), so a large max_num_tokens no longer passes the int32
indexing of qsa_paged_index_scores and qsa_topk_expand_decode; one chunk at the defaults.
- `allreduce_strategy` reaches the TP sum; `TRTLLM_ENABLE_PDL=0` turns PDL off for the decode
projections, the GDN replay chain and low_m_bf16_mm.
- One switch parser for every TRTLLM_MODELING_V2_* variable (unknown spellings raise); the
tp1-only switches warn at the other layouts. `one_rank` is set once on the core.
- Load checks: the n-gram table's padding rows; the layout's recurrent-state dtype at bind.
- causal_conv1d_update refuses a conv-state pool whose slot offset passes 32 bits in
TRT-LLM's kernel.
- The row-tile kernel lives once, in the family ops, with the 8-row swizzle-atom check.
- Dead code removed: unused PTX helpers, always-true expert-split checks, duplicate op
aliases, stage_index's unused argument.
Routing declines EPLB at tep4, dep4 and dep2.
Tests: the family operations' guard and CUDA-graph replay tests their notes state, the
16/48-head GDN replay cells, score-table chunking, switch parsing, the TP sum's strategy,
vision admission refusals; per-case acceptance-length floors for the four-GPU gsm8k cases;
the tp1 pre-merge and post-merge tiers as on main; durations for the new cases. Notes describe
the operations as they stand.
Docs: a Modeling V2 section in the Qwen3.8-Flash-Next deployment guide.
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
...-guide-for-qwen3.8-flash-next-on-trtllm.md | 44 +
.../attention/qsa_rowtile_attention_cute.py | 3 +
.../ops/attention/qsa_sparse_gqa_rowtile.py | 3 +-
.../qwen4_exp_attn_decode_prologue.py | 2 +-
.../qwen4_exp/ops/gemm/bf16_gemv_pdl.py | 4 +-
.../qwen4_exp/ops/gemm/low_m_bf16_mm.py | 3 +
.../ops/hc/qwen4_exp_gated_residual.py | 31 +-
.../ops/hc/qwen4_exp_hc_lowrank_act.py | 24 +-
.../qwen4_exp/ops/hc/qwen4_exp_post_inject.py | 15 +-
.../moe/fp4_block_scale_moe_renorm_routed.py | 2 +-
.../qwen4_exp/ops/ple/ngram_token_history.py | 9 +-
.../ops/ple/ple_gate_conv_inject_norm.py | 23 +-
.../qwen4_exp/ops/ple/ple_ngram_embed.py | 14 +-
.../qwen4_exp/ops/ssm/causal_conv1d_update.py | 15 +-
.../ops/ssm/cu_seqlens_from_seq_lens.py | 7 +-
.../ssm/flashinfer_chunk_gated_delta_rule.py | 37 +-
.../ops/ssm/gdn_decode_gated_norm_pdl.py | 15 +-
.../ops/ssm/gdn_decode_rows_gated_norm.py | 11 +-
.../ops/ssm/gdn_prefill_conv_post.py | 12 +-
.../__init__.py | 2 +-
.../modeling.py | 312 ++--
.../qsa_rowtile_attention_cute__prefill.py | 1386 -----------------
.../qsa_sparse_gqa_rowtile__prefill.py | 18 +-
.../ops/attention/vision_rope_2d_packed_qk.py | 14 +-
.../ops/ssm/gdn_cached_replay_cute__verify.py | 152 +-
.../pool.py | 4 +-
.../weights.py | 28 +-
.../modeling_v2/models/qwen4_exp/routing.py | 63 +-
tests/integration/defs/.test_durations | 6 +
.../accuracy/modeling_v2/test_qwen4_exp.py | 46 +-
.../references/acceptance_length.yaml | 17 +
.../test_lists/test-db/l0_b300.yml | 23 +-
...t_modeling_v2_fused_sigmoid_mul_inplace.py | 91 +-
.../test_modeling_v2_silu_and_mul_rows.py | 26 +-
...est_modeling_v2_qsa_canonical_selection.py | 71 +-
...test_modeling_v2_qsa_paged_index_scores.py | 44 +-
.../test_modeling_v2_qsa_paged_kv_store.py | 67 +
.../test_modeling_v2_qsa_paged_sparse_gqa.py | 23 +-
...modeling_v2_qsa_paged_sparse_gqa_decode.py | 10 +-
..._modeling_v2_qsa_paged_sparse_gqa_nvfp4.py | 19 +
...ng_v2_qsa_paged_sparse_gqa_nvfp4_decode.py | 36 +-
.../test_modeling_v2_qsa_prefill_compress.py | 71 +
.../test_modeling_v2_qsa_rollback_planes.py | 87 ++
...test_modeling_v2_qsa_sparse_gqa_rowtile.py | 18 +-
...test_modeling_v2_qsa_topk_expand_decode.py | 23 +-
...eling_v2_qwen4_exp_attn_decode_prologue.py | 2 +-
.../embedding/test_modeling_v2_ngram_hash.py | 49 +-
.../gemm/test_modeling_v2_bf16_gemv_pdl.py | 9 +-
.../gemm/test_modeling_v2_fp8_gemv.py | 51 +-
.../test_modeling_v2_fp8_per_tensor_mm.py | 16 +-
.../gemm/test_modeling_v2_low_m_bf16_mm.py | 44 +
.../gemm/test_modeling_v2_mid_rows_mm.py | 19 +-
.../gemm/test_modeling_v2_nvfp4_gemv.py | 48 +-
...eling_v2_qwen4_exp_hc_decode_first_cute.py | 10 +-
...ing_v2_qwen4_exp_hc_decode_mix_fp4_cute.py | 14 +-
...modeling_v2_qwen4_exp_hc_decode_rows_tc.py | 18 +-
...deling_v2_qwen4_exp_hc_first_rescale_tc.py | 32 +-
...st_modeling_v2_qwen4_exp_hc_inject_norm.py | 9 +-
...st_modeling_v2_qwen4_exp_hc_inject_rstd.py | 28 +-
...st_modeling_v2_qwen4_exp_hc_lowrank_act.py | 25 +-
...eling_v2_qwen4_exp_hc_up_gate_mean_rstd.py | 40 +-
.../test_modeling_v2_qwen4_exp_post_inject.py | 35 +-
...sl_nvfp4_gather_grouped_gemm_act_fusion.py | 37 +-
...modeling_v2_cute_dsl_nvfp4_grouped_gemm.py | 31 +-
..._v2_fp4_block_scale_moe_renorm_permuted.py | 34 +-
...ng_v2_fp4_block_scale_moe_renorm_routed.py | 22 +-
...4_block_scale_moe_renorm_routed_windows.py | 4 +-
...est_modeling_v2_fp4_moe_renorm_cute_dsl.py | 27 +-
..._modeling_v2_fp8_block_scale_moe_runner.py | 57 +-
..._modeling_v2_fused_sigmoid_gate_mul_add.py | 38 +-
.../moe/test_modeling_v2_moe_sort.py | 16 +-
.../test_modeling_v2_moe_unpermute_inplace.py | 28 +-
.../test_modeling_v2_renorm_moe_routing_op.py | 37 +-
...st_modeling_v2_flashinfer_gemma_rmsnorm.py | 41 +
.../norm/test_modeling_v2_layer_norm.py | 63 +-
...v2_qwen4_exp_rms_norm_gated_token_major.py | 18 +
.../test_modeling_v2_ngram_token_history.py | 80 +-
...t_modeling_v2_ple_gate_conv_inject_norm.py | 71 +-
.../test_modeling_v2_bf16_to_fp32.py | 25 +-
...st_modeling_v2_vision_rope_2d_packed_qk.py | 27 +-
...v2_qwen4_exp_hc_decode_tc_first__decode.py | 28 +-
...eling_v2_gdn_cached_replay_cute__verify.py | 21 +-
.../test_attention_plumbing.py | 68 +
.../test_core_layout.py | 48 +-
.../test_modeling_v2_spec_rollback.py | 2 +-
.../test_mtp_draft_hc.py | 71 +-
.../test_routing.py | 28 +-
.../test_select_canonical.py | 3 +
.../test_vision.py | 6 +
...modeling_v2_causal_conv1d_triton_update.py | 15 +-
.../test_modeling_v2_causal_conv1d_update.py | 26 +
...st_modeling_v2_cu_seqlens_from_seq_lens.py | 51 +-
...ng_v2_flashinfer_chunk_gated_delta_rule.py | 37 +-
...t_gated_delta_rule_cached_replay_update.py | 57 +
..._modeling_v2_gdn_decode_rows_gated_norm.py | 18 +-
.../test_modeling_v2_gdn_prefill_conv_post.py | 13 +-
.../test_modeling_v2_patch_embed_conv3d.py | 110 +-
97 files changed, 2658 insertions(+), 1980 deletions(-)
delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_rowtile_attention_cute__prefill.py
diff --git a/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md
index 4cb49d003240..98c05da231b7 100644
--- a/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md
+++ b/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md
@@ -198,6 +198,50 @@ encoder_max_num_tokens: 65536
Image input is supported in aggregated serving with a local encoder. A separate encoder-to-prefill handoff is not part of this guide.
+### Modeling V2 (Experimental)
+
+On GB300 (sm_103), the NVFP4 checkpoint can also be served by the modeling_v2 implementation, which is off by default. Set `modeling_v2` in the YAML file:
+
+```yaml
+modeling_v2: auto # use modeling_v2 when it supports this deployment, otherwise the built-in model
+# modeling_v2: require # fail at startup, with the reason, when it does not
+```
+
+It supports these single-node layouts:
+
+| Layout | GPUs | Parallel settings | Speculative decoding |
+|---|---|---|---|
+| tp1 | 1 | `tensor_parallel_size: 1` | none, or MTP with `max_draft_len` 1-7 |
+| tp4 | 4 | `tensor_parallel_size: 4`, `moe_tensor_parallel_size: 4` | none, or MTP with `max_draft_len` 1-7 |
+| tep4 | 4 | `tensor_parallel_size: 4`, `moe_expert_parallel_size: 4` | none, or MTP with `max_draft_len` 1-7 |
+| dep4 | 4 | as tep4, plus `enable_attention_dp: true` | none, or MTP with `max_draft_len` 1-7 |
+| dep2 | 2 | `tensor_parallel_size: 2`, `moe_expert_parallel_size: 2`, `enable_attention_dp: true` | none |
+
+Under `auto`, these settings make the server fall back to the built-in model; under `require`, the server fails at startup:
+
+* Any other layout, including `enable_lm_head_tp_in_adp: true` with attention DP.
+* Speculative decoding other than one-model MTP, or a draft length that changes at runtime.
+* MTP with `TRTLLM_USE_GDN_REPLAY=0`.
+* A KV cache dtype other than the checkpoint's bf16, `fp8` or `nvfp4`.
+* A recurrent-state cache other than `bfloat16`. tp1 also accepts `float32`.
+* `moe_config.backend` other than `AUTO` or `TRTLLM`. tp1 also accepts `CUTEDSL` with `max_draft_len` up to 3.
+* `prefill_cuda_graph_backend: piecewise`; with attention DP, also `breakable`.
+* `moe_config.load_balancer` (EPLB) at tep4, dep4 or dep2.
+
+Other differences from the built-in model:
+
+* The PLE n-gram table is in pinned host memory by default, one copy per rank (about 48 GiB each, so about 190 GiB of host memory for a four-GPU node). `TRTLLM_MODELING_V2_PLE_HOST_OFFLOAD=0` keeps it on the device. `TRTLLM_QWEN4_EXP_PLE_HOST_OFFLOAD` applies only to the built-in model.
+* `disable_mm_encoder` has no effect. The vision tower is always built and image requests are accepted.
+* At tp1 these environment switches enable lossy FP8 or NVFP4 copies of projections. All are off by default, and the other layouts ignore them with a warning:
+ * `TRTLLM_MODELING_V2_DENSE_FP8`
+ * `TRTLLM_MODELING_V2_GDN_INPROJ_NVFP4`
+ * `TRTLLM_MODELING_V2_TARGET_HEAD_NVFP4` (with `DENSE_FP8`)
+ * `TRTLLM_MODELING_V2_MTP_HEAD_NVFP4`
+ * `TRTLLM_MODELING_V2_MTP_MOE_NVFP4`
+ * `TRTLLM_MODELING_V2_MTP_DENSE_FP8`
+* At tp1, `TRTLLM_MODELING_V2_MOE_PREFILL_COPY=0` drops the second copy of the routed experts that serves large prefill steps.
+* A `TRTLLM_MODELING_V2_*` switch accepts `1`/`0`, `true`/`false`, `yes`/`no` or `on`/`off`, and the server fails at startup on any other value.
+
### Key Configuration Options
These options control TensorRT LLM behavior and are set in the YAML file passed to `trtllm-serve` with the `--config` argument.
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_rowtile_attention_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_rowtile_attention_cute.py
index 1a9cac72e24e..4fa6fef27d45 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_rowtile_attention_cute.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_rowtile_attention_cute.py
@@ -101,6 +101,9 @@ def __init__(
"the TMEM plan (O 256 + 2 x S 64 + Q 128 columns) is laid out for D = 256"
)
assert group_size * rows_per_tile <= 128
+ # The O staging's four SW128 boxes sit `64 G R` elements apart: each must start on
+ # a 1024-byte swizzle atom, so the live MMA rows come in whole 8-row atoms.
+ assert group_size * rows_per_tile % 8 == 0, (group_size, rows_per_tile)
self.group_size = group_size
self.rows_per_tile = rows_per_tile
self.d = head_dim
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_sparse_gqa_rowtile.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_sparse_gqa_rowtile.py
index a7b5ad5e6a8e..328d93499bbc 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_sparse_gqa_rowtile.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_sparse_gqa_rowtile.py
@@ -853,8 +853,7 @@ class QsaSparseGqaRowtile(OpWrapper):
it is slower than the per-row kernel (generation-only launches 1.4-2x);
the target's served context launches are thousands of rows. The tensor-core
row-tile pass alone, on six dumped 16384-row context launches of the
- target (GB300, layers 0-5 of one step): 581-775 us, against 694-886 us for
- the one-program-per-item pass it replaced (the same output, bit for bit).
+ target (GB300, layers 0-5 of one step): 581-775 us.
Compiled once per (group size, rows per tile, tokens per block, PDL):
nothing in the row-tile pass's compile key depends on the row count or
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qwen4_exp_attn_decode_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qwen4_exp_attn_decode_prologue.py
index a3b705eb80d5..c9c1f0a581c3 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qwen4_exp_attn_decode_prologue.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qwen4_exp_attn_decode_prologue.py
@@ -3,7 +3,7 @@
"""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
+The composition of entries it fuses is 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
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/bf16_gemv_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/bf16_gemv_pdl.py
index 75652e898bae..7fd325fde41e 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/bf16_gemv_pdl.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/bf16_gemv_pdl.py
@@ -355,12 +355,12 @@ class Bf16GemvPdl(OpWrapper):
order differs from the fp32 reference.
Hard limits
- M rows per launch boundary MAX_ROWS=4 driven at 1 (cells)
+ M rows per launch boundary MAX_ROWS=4 driven at 1 and 4 (cells)
K a TILING key 2560, 6144, 640, 1536 all driven
N any; the outputs per CTA halve until they divide N N=1 driven
The test also replays a PDL chain of these GEMVs under a CUDA graph, each
- reading the one before it.
+ reading the one before it, and pins the guard's refusal of a row-major `mat_b`.
Not certified: M = 2..4 at K other than 1536; PDL disabled in the environment (the
PDL-off cell is the same kernel without the attribute).
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/low_m_bf16_mm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/low_m_bf16_mm.py
index 6a6722a355a6..6e2089f0928b 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/low_m_bf16_mm.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/low_m_bf16_mm.py
@@ -34,6 +34,7 @@
OpWrapper,
true_fp32_matmul,
)
+from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
class _Plan(NamedTuple):
@@ -221,6 +222,8 @@ def __call__(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = False)
m, k = mat_a.shape
n = mat_b.shape[1]
p = plan(m, n, k)
+ # `TRTLLM_ENABLE_PDL=0` turns programmatic dependent launch off, as for every PDL op.
+ pdl = bool(pdl and get_env_enable_pdl())
if p.kind == "direct":
out = torch.empty((m, n), dtype=mat_a.dtype, device=mat_a.device)
return _fi_direct.run_direct_dense(
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_gated_residual.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_gated_residual.py
index 8fe4616fcea7..7a4a5b79475f 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_gated_residual.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_gated_residual.py
@@ -9,8 +9,11 @@
the raw stream gated) and does not instantiate hidden 2560. The semantics follow
the HuggingFace ``Qwen4ExpTextGatedResidual`` minus its norm.
-Called by the qwen4_exp target (`qwen3_8_flash_next_nvfp4__sm_103__default`) at every
-layout.
+No target calls it: the qwen4_exp target (`qwen3_8_flash_next_nvfp4__sm_103__default`)
+runs every mixer, the MTP layer's included, on the fused hyper-connection entries. This
+entry is their unfused form: `gate_ratio` is the band their tests gate `low`, `inject` and
+`mixed` in, and the target's MTP draft hyper-connection test takes `reference` as the
+float64 mixer and runs the entry as the unfused composition the fused sites are held to.
"""
from __future__ import annotations
@@ -83,7 +86,7 @@ class Qwen4ExpGatedResidual(OpWrapper):
spec=dict(leading=(256,), inject=False),
),
# The MTP layer's shared_head mixer (block_inject_weight None) at the
- # rows DecodeTarget.mtp_forward runs: one request, the largest decode
+ # rows a draft step's head mixes: one request, the largest decode
# graph batch, and mtp7's first draft over 128 x (1 + 7) verify rows.
Cell(
why="1 row, no injection: one request's MTP shared_head",
@@ -104,10 +107,11 @@ class Qwen4ExpGatedResidual(OpWrapper):
)
note = """
- Metadata consumed: none; every row independent. Call sites: twice per
- decoder layer (attention and MoE hyper-connections, 96 calls) plus the
- stack-final mixer with `block_inject_weight=None` -- whose `mixed` *is*
- the checkpoint's final norm -- and three more in the MTP layer (two
+ Metadata consumed: none; every row independent. Call sites: none in a
+ target (see the module docstring). The formula is every hyper-connection
+ mixer's: twice per decoder layer (attention and MoE hyper-connections),
+ the stack-final mixer with `block_inject_weight=None` -- whose `mixed`
+ *is* the checkpoint's final norm -- and three in the MTP layer (two
blocks and `shared_head`'s own mixer).
Preconditions, all loud (`RuntimeError` before producing anything): the
@@ -130,13 +134,12 @@ class Qwen4ExpGatedResidual(OpWrapper):
`mixed` within 0.40 of the test's band, cosine >= 0.9999982 and relative
L2 <= 1.88e-3 against float64.
- Not certified: other row counts; `inject` at 128 and 1024 rows (the MTP
- layer's unfused attention and MoE sites; `mixed` there is bitwise the
- no-injection cells'); more than one leading dim, hc_count other than 4,
- hidden other than 2560, lowrank other than 320, float16 and float32,
- weight/activation dtype mixing; the operand-range bounds above are
- measured. Hard limits: not probed; a 2^31-element operand would take
- 209716 rows, past the 8192-row envelope.
+ Not certified: other row counts; `inject` at 128 and 1024 rows (`mixed`
+ there is bitwise the no-injection cells'); more than one leading dim,
+ hc_count other than 4, hidden other than 2560, lowrank other than 320,
+ float16 and float32, weight/activation dtype mixing; the operand-range
+ bounds above are measured. Hard limits: not probed; a 2^31-element
+ operand would take 209716 rows, past the 8192-row envelope.
"""
def __call__(
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_lowrank_act.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_lowrank_act.py
index bb5368c55c33..90b7dd8a2746 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_lowrank_act.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_lowrank_act.py
@@ -66,9 +66,9 @@ class Qwen4ExpHcLowrankAct(OpWrapper):
ARCHS = frozenset({Arch.SM_103})
# spec: rows, inject (True), hc_count (4), rank (320), logits_width (384), dtype (bf16).
- # The default target's HC mixers run this form at every row count, on the fresh
- # fp32 `cublas_mm` output over the first `hc_first_cols` = ceil((320 + 4) / 32) * 32
- # = 352 columns of the first projection: hence the 352-wide cells.
+ # The default target's HC mixers run this form on the fresh fp32 `cublas_mm` output
+ # over the first `hc_first_cols` = ceil((320 + 4) / 32) * 32 = 352 columns of the
+ # first projection: hence the 352-wide cells.
CELLS: tuple[Cell, ...] = (
Cell(why="checkpoint rank 320 + 4, padded to 384, 128 rows", spec=dict(rows=128)),
Cell(why="checkpoint rank 320 + 4, padded to 384, prefill 8192", spec=dict(rows=8192)),
@@ -96,14 +96,16 @@ class Qwen4ExpHcLowrankAct(OpWrapper):
)
note = """
- Metadata consumed: none. Call sites: every hyper-connection mixer, between
- the first projection (`cublas_mm`, fp32 out, 352 columns) and the up
- projection. HC is replicated, so the rows are this rank's tokens at every
- member: decode 1..128 (tp4 / tep4) or 1..32 (a dep4 rank), MTP verify up to
- 128 x 8 = 1024 rows (tep4, draft 7), prefill one chunk of up to 8192 rows
- (`max_num_tokens`). Preconditions: `logits` fp32, unit column stride, at
- least `rank (+ hc_count)` columns; `hc_count` a power of two.
- Graph-capturable.
+ Metadata consumed: none. Call sites: `_hc_first`, every hyper-connection
+ mixer of a step past `_HC_TC_FIRST_MAX_ROWS` (256) rows -- smaller steps
+ run `qwen4_exp_hc_decode_tc_first` -- between the first projection
+ (`cublas_mm`, fp32 out, 352 columns) and the up projection. HC is
+ replicated, so the rows are this rank's tokens at every member: MTP verify
+ up to 128 x 8 = 1024 rows (tep4, draft 7) and prefill chunks of up to 8192
+ rows (`max_num_tokens`); at one rank, prefill steps from
+ `_HC_RSTD_MIN_ROWS` rows take `qwen4_exp_hc_first_rescale_tc` instead.
+ Preconditions: `logits` fp32, unit column stride, at least
+ `rank (+ hc_count)` columns; `hc_count` a power of two. Graph-capturable.
Source: a Triton kernel JIT-compiled at first call for the running device;
the cells ran on sm_103. One program per row with an int64 row offset;
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_post_inject.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_post_inject.py
index eadadb7e44ba..be2bac2a555d 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_post_inject.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_post_inject.py
@@ -35,8 +35,8 @@ class Qwen4ExpPostInject(OpWrapper):
identity: no stream's value enters another.
Fusion boundary: the accumulate only. The grouped norm, the input mixer
- that produced the block's input and `inject` (`qwen4_exp_gated_residual`),
- and the block itself are the caller's.
+ that produced the block's input and `inject` (`qwen4_exp_gated_residual`'s
+ formula), and the block itself are the caller's.
"""
ARCHS = frozenset({Arch.SM_103})
@@ -77,10 +77,13 @@ class Qwen4ExpPostInject(OpWrapper):
)
note = """
- Metadata consumed: none. Call sites: twice per decoder layer (after the
- attention/GDN block and after the MoE) and twice in the MTP layer, with
- `inject` the `[T, 4]` weights `qwen4_exp_gated_residual` returned for that
- block -- in `[0, 2]`, because they are `2 * sigmoid(.)`.
+ Metadata consumed: none. Call sites: the MTP layer's MoE injection
+ (`_mtp_tail` and `DecodeTarget`'s draft tail), and in `DecodeTarget`'s
+ trunk the block output still pending before the PLE layer; every other
+ site's accumulate is fused into the next site's first kernel
+ (`qwen4_exp_hc_inject_norm` and its rstd and decode forms). `inject` is
+ the `[T, 4]` weights the site's mixer returned for that block -- in
+ `[0, 2]`, because they are `2 * sigmoid(.)`.
Preconditions, all loud except the guarded one: the packed width must equal
`inject.shape[-1] * block_out.shape[-1]` (the unflatten raises); all three
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_block_scale_moe_renorm_routed.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_block_scale_moe_renorm_routed.py
index f6257976bc31..d34bda4302d7 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_block_scale_moe_renorm_routed.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_block_scale_moe_renorm_routed.py
@@ -309,7 +309,7 @@ class Fp4BlockScaleMoeRenormRouted(OpWrapper):
)
note = """
- Routing semantics, measured against the composition it replaces
+ Routing semantics, measured against the two-entry composition
(`renorm_moe_routing_op` -> pre-routed runner): bitwise-equal finalized
output at every cell (the test re-drives this), and at bring-up on randn /
coarse / large-magnitude / all-equal logits, bf16 and fp32 logits,
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ngram_token_history.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ngram_token_history.py
index f0d557a20929..5a9c1760dbb9 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ngram_token_history.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ngram_token_history.py
@@ -5,8 +5,9 @@
back to -- the per-slot n-gram context pool.
A macro entry, composed from native torch integer ops (~20 small launches);
-exact, so gated bitwise. Called by the qwen4_exp target at every layout; its
-per-rank slot pools are among the cells. Nothing registered
+exact, so gated bitwise. The qwen4_exp target calls ``ple_ngram_embed``, whose
+``reference`` composes this entry's; the default target's per-rank slot pools
+are among the cells. Nothing registered
under ``torch.ops.trtllm`` assembles an n-gram id window, and
``EngramHashProvider.compute_hashes`` / ``NgramHashMapping._get_ngram_hashes``
(``tensorrt_llm/_torch/modules/engram/engram.py``) has no cross-step state: it
@@ -219,8 +220,8 @@ class NgramTokenHistory(OpWrapper):
The pool the manager hands over is a contiguous `[slots, 2]` int64
(stride (2, 1), offset 0) at tp4, tep4 and dep4, with the GDN pools'
slot count. The call site is `ple_ngram_embed`, whose `reference`
- composes this entry's `reference`; the tp1 forward launches the fused
- entry, not this op.
+ composes this entry's `reference`; the target launches the fused entry,
+ not this op.
Not certified: tp4's pool (pure tp4 is not served), the MTP and
multimodal members' pools, and the dep4 variants' pools (kv_fp8,
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_gate_conv_inject_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_gate_conv_inject_norm.py
index f1eebf8224e8..e976cf44e696 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_gate_conv_inject_norm.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_gate_conv_inject_norm.py
@@ -7,11 +7,11 @@
norm -- three Triton launches behind one call.
Called by the qwen4_exp target at every layout; uses the row-to-sequence
-search of ``ple_ngram_embed``. It computes the composition of three grouped
-norms, ``gated_ngram_injection``, ``dilated_short_conv_paged``, two adds and
-the norm-only ``qwen4_exp_hc_inject_norm`` after the PLE add. No
-``torch.ops.trtllm`` op or FlashInfer callable takes a dilation (see
-``dilated_short_conv_paged``).
+search of ``ple_ngram_embed``. Its ``reference`` is that chain in float64,
+rounded to the activation dtype after each step: three grouped norms, the
+sigmoid-gated value, the dilated depthwise conv over the carried window, two
+adds and the next site's grouped norm. No ``torch.ops.trtllm`` op or FlashInfer
+callable takes a dilation.
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
@@ -302,7 +302,8 @@ def _state_head(state_len: int) -> int:
CKPT_HIDDEN = 2560
CKPT_DILATION = 3
CKPT_CLAMP_MIN = 1e-6
-#: Largest PLE-conv pool seen for this checkpoint at tp1 (see `dilated_short_conv_paged`).
+#: The PLE-conv slot count of a full-size tp1 engine run (the n-gram pool's slot space). The
+#: op addresses a slot by its index alone, so a pool of another size takes the same path.
ENGINE_SLOTS = 461
#: The default target's per-rank PLE-conv pools: the conv state and the n-gram context are
#: bound over one slot space (the stock manager asserts their slot counts equal), so the
@@ -319,18 +320,18 @@ def _spec(lengths, **spec) -> dict:
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
- unfused composition rounds -- for row `t` of sequence `s`, stream `c`::
+ * (1 + w)`), and each step rounded to the activation dtype as `reference`
+ rounds it -- 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
+ gt = round(sigmoid(sg) * value[t])
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
+ ple = round(gt + round(silu(conv)))
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
+ normed[t, c] = round(n(hc_out[t, c], next_norm[c]))
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
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_ngram_embed.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_ngram_embed.py
index f727731c8386..9d9ae7917fd6 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_ngram_embed.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_ngram_embed.py
@@ -4,9 +4,10 @@
table gather and per-tensor fp8 dequantization in one Triton launch, with the
per-slot n-gram context published back in the same pass.
-Called by the qwen4_exp target at every layout. It computes the composition
-``ngram_token_history`` + ``ngram_hash`` + ``nn.functional.embedding`` +
-``dequantize_e4m3_per_tensor``, whose references it reads. Nothing registered
+Called by the qwen4_exp target at every layout. Its ``reference`` is the
+scalar references of ``ngram_token_history`` (window and published pool) and
+``ngram_hash`` (rows), then a float64 gather of the fp8 table rows times the
+per-tensor scale, rounded once. Nothing registered
under ``torch.ops.trtllm`` or in the installed FlashInfer gathers and
dequantizes n-gram rows (``torch.ops.trtllm.qwen4_exp_ple_ngram_hash`` hashes a
window only).
@@ -303,7 +304,7 @@ class PleNgramEmbed(OpWrapper):
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.
+ with `D = table.shape[1]` -- one rounding of the scaled fp8 value.
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.
@@ -432,8 +433,9 @@ class PleNgramEmbed(OpWrapper):
stock PLE host offload's placement. The checkpoint-height cells allocate
the 47.7 GiB table on the device.
- Bitwise equal to `ngram_token_history` + `ngram_hash` + embedding +
- `dequantize_e4m3_per_tensor`, window, rows, values and pool alike.
+ Bitwise equal to `reference` (`ngram_token_history`'s and `ngram_hash`'s
+ references, then the gather and one rounding of the scale), 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.
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_update.py
index f8f98d2a9807..a35358e48c8e 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_update.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_update.py
@@ -25,6 +25,9 @@
#: This checkpoint's GDN `[q|k|v]` width; the conv is width 4 over a 3-column pool.
CKPT_DIM = 10240
+#: TRT-LLM's kernel forms a slot's offset as `int slot * uint32_t conv_state_batch_stride`:
+#: a pool of this many elements or more (slots x slot stride) wraps, so the op refuses it.
+TRTLLM_SLOT_OFFSET_LIMIT = 2**32
#: Largest conv-pool slot count the cache manager allocates for this checkpoint at a batch
#: of at most 32. The decode step hands the whole pool over unsliced.
ENGINE_SLOTS = 461
@@ -504,8 +507,9 @@ class CausalConv1dUpdate(OpWrapper):
conv_state_batch_stride`, 32-bit unsigned elements: it wraps from slot
15535 at stride 276480 (tp4 / tep4) and slot 3884 at 1105920 (dep4).
The measured pools reach 518,123,519 (tep4) and 940,031,999 (dep4)
- elements, under 2^31 elements and bytes. The CuTe kernel forms slot and
- row offsets in 64 bits.
+ elements, under 2^31 elements and bytes. A call that would run it over a
+ pool of 2^32 elements or more (`TRTLLM_SLOT_OFFSET_LIMIT`) raises
+ ValueError instead. The CuTe kernel forms slot and row offsets in 64 bits.
"""
def __call__(
@@ -524,6 +528,13 @@ def __call__(
):
_decode_conv(x, conv_state, weight, conv_state_indices, pad_slot_id)
return x
+ if conv_state.shape[0] * conv_state.stride(0) >= TRTLLM_SLOT_OFFSET_LIMIT:
+ raise ValueError(
+ f"a conv-state pool of {conv_state.shape[0]} slots at stride "
+ f"{conv_state.stride(0)}: TRT-LLM's causal_conv1d_update forms a slot's offset "
+ f"in 32 bits and would wrap past slot {(TRTLLM_SLOT_OFFSET_LIMIT - 1) // conv_state.stride(0)}; "
+ "lower free_gpu_memory_fraction or the recurrent-state share of the KV cache"
+ )
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/ops/ssm/cu_seqlens_from_seq_lens.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/cu_seqlens_from_seq_lens.py
index c94f9601b4d3..2ab3f9740489 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/cu_seqlens_from_seq_lens.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/cu_seqlens_from_seq_lens.py
@@ -18,7 +18,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 at tp1 reuses the
+#: under MTP (2..8: `max_draft_len` 1..7); at ``p = 1`` the target reuses the
#: recurrent metadata's vector instead, and ``decode_p1_128`` certifies the call
#: there all the same.
#: The ``_128`` batches are the serving envelope's widest step at every member of
@@ -123,8 +123,9 @@ class CuSeqlensFromSeqLens(OpWrapper):
device-side assert in the consumers); a padded row is length 0, never -1;
`out[-1]` equals the packed activation's row count.
- `dtype` has no default: `causal_conv1d_fwd` requires int32 and
- `chunk_gated_delta_rule` keys its autotuner on `cu_seqlens`' dtype.
+ `dtype` has no default: it is the consumer's index dtype. The target asks
+ for int64, the `query_start_loc` that `ple_ngram_embed` and
+ `ple_gate_conv_inject_norm` take.
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
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/flashinfer_chunk_gated_delta_rule.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/flashinfer_chunk_gated_delta_rule.py
index 85a0d7c477f9..9a9b1a5637e0 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/flashinfer_chunk_gated_delta_rule.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/flashinfer_chunk_gated_delta_rule.py
@@ -391,26 +391,23 @@ class FlashinferChunkGatedDeltaRule(OpWrapper):
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, and belong to
- FlashInfer's decode kernel (`gdn_decode.gated_delta_rule_decode_pretranspose`).
-
- The target (`PrefillTarget._gdn_core`, `_gdn_core_verify`) passes the
- context rows by keyword after `gdn_prefill_conv_post(linear_gate=True)`,
- with `has_initial_state` and `linear_gate=True`.
-
- Call site with `has_initial_state` (`hi_p` / `hi_all`, the flags
- `reset_recurrent_state_rows` would zero the fresh slots by): a fresh
- sequence's slot is never read, and the scatter writes every named slot,
- so the reset of the recurrent pool before this call is not needed (the
- conv pass reads no carry for such a sequence and its publish writes the
- whole conv window). Bitwise the reset followed by the plain call; the
- gather moves only the carried sequences' rows.
+ Call site (`PrefillTarget._gdn_core`, every step with context rows, verify steps
+ included): the context rows by keyword after
+ `gdn_prefill_conv_post(linear_gate=True)` -- q/k already L2-normalized, contiguous
+ `[1, T, H, 128]`; `g` the linear decay `exp(g)`; `beta` fp32 -- with `scale=None`,
+ the layer's pool as `initial_state` indexed by the context rows' slots,
+ `inplace_indexed_state_update=True`, `output_final_state=False`, int64
+ `cu_seqlens`, `head_first=False`, `use_qk_l2norm_in_kernel=False`, `output=` a
+ prefix view, `has_initial_state` and `linear_gate=True`. Give it context rows
+ only: single-token rows are slow here, and belong to FlashInfer's decode kernel
+ (`gdn_decode.gated_delta_rule_decode_pretranspose`).
+
+ With `has_initial_state` (`hi_p` / `hi_all`): a fresh sequence's slot is never
+ read, and the scatter writes every named slot, so the fresh slots of the
+ recurrent pool need no zeroing before this call (the conv pass reads no carry
+ for such a sequence and its publish writes the whole conv window). Bitwise
+ zeroing the fresh slots followed by the plain call; the gather moves only the
+ carried sequences' rows.
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
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_decode_gated_norm_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_decode_gated_norm_pdl.py
index 2112962a1b7a..582e3f5007a7 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_decode_gated_norm_pdl.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_decode_gated_norm_pdl.py
@@ -438,13 +438,14 @@ class GdnDecodeGatedNormPdl(OpWrapper):
Source: the CuTe DSL kernel in this file, JIT-compiled (`cute_ext.compile`) for the
running device at the first call per geometry; the cells ran that path on sm_103.
- The default target: its routing admits a bf16 recurrent-state cache only, and this kernel takes
- an fp32 pool only -- a bf16 pool fails its compile loudly ("`copy` currently only
- supports equal source and destination element type bit width") -- so the default target's
- decode step does not launch it. Its bf16 decode recurrence is
- `gdn_decode_rows_gated_norm`, which reads this entry's `reference` and
- `_recurrence_reference` (both round the state to the pool's own dtype) and not its
- kernel.
+ The default target: this kernel takes an fp32 pool only -- a bf16 pool fails its
+ compile loudly ("`copy` currently only supports equal source and destination element
+ type bit width") -- and the routing admits an fp32 recurrent-state cache at tp1 alone.
+ There `DecodeTarget._gdn_block` launches it for a captured decode step of at most
+ `MAX_ROWS` rows over an fp32 pool; every other decode step, and every bf16 pool (the
+ only one at tp4, tep4 and dep4), runs `gdn_decode_rows_gated_norm`, which reads this
+ entry's `reference` and `_recurrence_reference` (both round the state to the pool's own
+ dtype) and not its kernel.
Certified, sm_103:
- tp1's 16 / 48 heads: 1 and 4 rows over a contiguous 461-slot pool, and 2 rows over a
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_decode_rows_gated_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_decode_rows_gated_norm.py
index c3c9d94b7133..890d41dfe4e4 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_decode_rows_gated_norm.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_decode_rows_gated_norm.py
@@ -892,9 +892,11 @@ class GdnDecodeRowsGatedNorm(OpWrapper):
Source: the CuTe DSL kernel in this file, JIT-compiled (`cute.compile`) for the running
device at the first call per geometry; the cells ran that path on sm_103 (GB300).
- The default target: its routing admits a bf16 recurrent-state cache only, so its decode step
- runs this entry over the bf16 pool at every batch size (the fp32-only one-launch kernel,
- `gdn_decode_gated_norm_pdl`, is never launched here; its references are this entry's).
+ The default target (`DecodeTarget._gdn_block`, a decode step without verify rows): at
+ tp4, tep4 and dep4 its routing admits a bf16 recurrent-state cache only, so the decode
+ step runs this entry over the bf16 pool at every batch size. At tp1 it also serves an
+ fp32 cache; there a captured step of at most `gdn_decode_gated_norm_pdl.MAX_ROWS` rows
+ over an fp32 pool runs that one-launch kernel, and every other decode step this entry.
A tp4/tep4 rank holds 4 k / 12 v heads, its projection `[rows, 4120]`; a dep4 rank all
16 / 48, `[rows, 16480]`. The pool is the stock manager's coalesced page: every GDN
layer's state of a slot in one page, slot stride 36 layers' states (7077888 elements at
@@ -933,8 +935,7 @@ class GdnDecodeRowsGatedNorm(OpWrapper):
`Int64(slot) * s_slot`: any slot count an int32 index can name.
Not certified: `pdl=False`; the slot counts of the members' variants (kv_fp8, kv_nvfp4,
- chunked, reuse at dep4; MTP at tep4, `max_seq_len` 16384) and of tp4's own engine; this
- target's own call, which its modeling does not make yet.
+ chunked, reuse at dep4; MTP at tep4, `max_seq_len` 16384) and of tp4's own engine.
"""
def __call__(
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_prefill_conv_post.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_prefill_conv_post.py
index acb71d5577f5..6bf32fb34f91 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_prefill_conv_post.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_prefill_conv_post.py
@@ -818,9 +818,8 @@ class GdnPrefillConvPost(OpWrapper):
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.
+ -- in one pass over the activation plus a per-sequence window publish;
+ `reference` evaluates each line above per sequence.
Fusion boundary: the conv-state read, conv, SiLU, split, q/k norm, gates
and the state publish are inside. The projection, the generation rows'
@@ -939,7 +938,7 @@ class GdnPrefillConvPost(OpWrapper):
(`cache_indices`, int32 `[N]`) and carry flags (`has_initial_state`, bool
`[N]`) are explicit.
- Call sites (`PrefillTarget._gdn_prefill_prep`, from `_gdn_block`): the
+ Call sites (`PrefillTarget._gdn_prefill_prep`, from `_gdn_core`): 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
@@ -960,9 +959,8 @@ class GdnPrefillConvPost(OpWrapper):
dep4): one slot holds all 36 GDN layers' `[C, 3]` states, so the view
has strides `(36 * C * 3, 3, 1)` --
(276480, 3, 1) at C 2560, (1105920, 3, 1) at C 10240 -- and layer `i`
- starts `i * C * 3` elements into the page. The target's forward does
- not call this variant yet; the rank cells mirror tp1's call at the
- rank's widths.
+ starts `i * C * 3` elements into the page. The rank cells below are that
+ call at the rank's widths.
Three launches. The conv pass (CuTe DSL, `_ConvPassKernel`): grid
`heads x cdiv(T, 64)` (80 heads at tp1 and dep4, 20 on a tp4/tep4
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/__init__.py
index eba053a6d445..579261bda67a 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/__init__.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/__init__.py
@@ -1,3 +1,3 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
-"""Qwen3.8-Flash-Next-NVFP4 / sm_103 / default: tp4, tep4, dep4 and dep2."""
+"""Qwen3.8-Flash-Next-NVFP4 / sm_103 / default: tp1, tp4, tep4, dep4 and dep2."""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py
index ed7631b9540e..fb78e4aa8757 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py
@@ -52,15 +52,15 @@
of its own, settled at construction (`one_rank`): the decode projections' one-rank dispatch
(`_one_rank_projection`); the hyper-connection sites' row-count forms (`hc_fast_sites`: the
decode GEMV and tcgen05 rows pairs, the prefill rstd sites), which the other layouts leave
-out so that a request's tokens do not depend on what its steps carry beside it; a trtllm-gen
+out; a trtllm-gen
decode MoE whose finalize the next site folds in (`moe_fold`); the CuteDSL copy of the
routed experts for prefill steps past `_CUTE_DSL_MOE_MIN_ROWS` (`moe_prefill_copy`), or the
CuteDSL pipeline for every step (`moe_config.backend: CUTEDSL`, `_cute_dsl_routed`); an fp32
recurrent state, whose captured single-token decode runs `gdn_decode_gated_norm_pdl`; the
decode draft layer and its FP8 head; the NVFP4 pool's context rows attended a group of
requests at a time (`_nvfp4_context_plan`); and the opt-in FP8 / NVFP4 copies of the dense
-projections and heads (`_DENSE_FP8_ENV`, `_GDN_INPROJ_NVFP4_ENV`, `_TARGET_HEAD_NVFP4_ENV`,
-`_MTP_HEAD_NVFP4_ENV`, `_MTP_MOE_NVFP4_ENV`).
+projections and heads (`_ONE_RANK_SWITCHES`: `_DENSE_FP8_ENV`, `_GDN_INPROJ_NVFP4_ENV`,
+`_TARGET_HEAD_NVFP4_ENV`, `_MTP_HEAD_NVFP4_ENV`, `_MTP_MOE_NVFP4_ENV`, `_MTP_DENSE_FP8_ENV`).
**QSA sparse attention.** The index plane is maintained on every step so a request can
cross `dense_seq_len_threshold` later without rebuilding its prefix; above it the sparse
@@ -150,9 +150,6 @@
from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import (
fused_qk_norm_rope,
)
-from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import (
- fused_qk_norm_rope as mtp_fused_qk_norm_rope,
-)
from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import (
QM_FP8_KV_CACHE,
QM_NVFP4_KV_CACHE,
@@ -165,9 +162,6 @@
)
from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.reducescatter import reducescatter
from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm
-from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import (
- cublas_mm as mtp_cublas_mm,
-)
from tensorrt_llm._torch._experimental.modeling_v2.catalog.graph.eager_on_graph import (
eager_on_graph,
)
@@ -180,9 +174,6 @@
from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.activation.fused_sigmoid_mul_inplace import (
fused_sigmoid_mul_inplace,
)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.activation.fused_sigmoid_mul_inplace import (
- fused_sigmoid_mul_inplace as mtp_fused_sigmoid_mul_inplace,
-)
from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.activation.silu_and_mul_rows import (
silu_and_mul_rows,
)
@@ -433,6 +424,7 @@
from tensorrt_llm._torch.attention.backends.fmha.phased import get_spec_decoding_position_offsets
from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata
from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata
+from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
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
@@ -474,7 +466,8 @@
# one without the GDN cached replay (the per-step recurrent states).
_MAX_VERIFY_WIDTH = _pool.MAX_VERIFY_WIDTH
-# Recurrent-state cache dtypes the target serves: at one rank the checkpoint's fp32 too.
+# Recurrent-state cache dtypes the target serves, checked at bind: at one rank the checkpoint's
+# fp32 too.
_SSM_STATE_DTYPES = (torch.bfloat16,)
_ONE_RANK_SSM_STATE_DTYPES = (torch.float32, torch.bfloat16)
@@ -562,8 +555,8 @@
_PAD_SLOT_ID = -1
#: Whether a captured decode step's projections on `low_m_bf16_mm` launch with programmatic
#: dependent launch (their FlashInfer kernels stream the weight before waiting on the
-#: activation's producer; the cuBLAS branch ignores it).
-_DECODE_PROJ_PDL = True
+#: activation's producer; the cuBLAS branch ignores it): `TRTLLM_ENABLE_PDL`, on by default.
+_DECODE_PROJ_PDL = get_env_enable_pdl()
#: The K at which `_decode_projection` runs `bf16_gemv_pdl` up to `_GEMV_MAX_ROWS` rows: a
#: tp4/tep4 rank's o and GDN out projections.
_GEMV_TP4_K = 1536
@@ -574,8 +567,8 @@
#: Whether a verify step's conv signals its dependents at its start and the CuTe replay after it
#: launches with programmatic dependent launch: the replay then streams the checkpoint and the
#: history (which nothing in the step writes ahead of it) while the conv runs, and reads the
-#: conv's output only after `griddepcontrol.wait`.
-_GDN_REPLAY_PDL = True
+#: conv's output only after `griddepcontrol.wait`. `TRTLLM_ENABLE_PDL`, on by default.
+_GDN_REPLAY_PDL = get_env_enable_pdl()
_SOFTPLUS_THRESHOLD = 20.0
_L2_NORM_EPS = 1e-6
_PLE_CLAMP_MIN = 1e-6
@@ -591,8 +584,8 @@
#: there up to `_HC_DECODE_ROWS_MAX` the tcgen05 pair `qwen4_exp_hc_decode_rows_first` +
#: `qwen4_exp_hc_decode_rows_mix`; larger steps, and every step at the other layouts,
#: `qwen4_exp_hc_inject_norm` and the mixer. The two fast pairs form the mixer's low-rank
-#: activations within 1-2 bf16 ulp of that path, not bitwise, so a request's tokens would depend
-#: on how many rows its step carries; the other layouts keep one path at every batch size.
+#: activations within 1-2 bf16 ulp of that path, not bitwise; the other layouts do not take
+#: them.
_HC_DECODE_GEMV_MAX_ROWS = 1 # capacity: fixed: dispatch threshold, not a buffer bound
#: Decode rows from which a folded MoE block (`DecodeTarget._moe`) launches its router before
#: the shared expert: launched after the side stream's gate|up GEMM its CTAs waited for SMs
@@ -701,10 +694,23 @@
_FP8_CONV_FUSED_MAX_ROWS = 8 # capacity: fixed: a dispatch threshold
-def _switch(env: str) -> bool:
- mode = os.environ.get(env, "0")
- assert mode in ("0", "1"), f"{env}={mode!r}: expected 0 or 1"
- return mode == "1"
+#: The spellings of a `TRTLLM_MODELING_V2_*` switch: these turn it on, these off.
+_SWITCH_ON = frozenset({"1", "true", "yes", "on"})
+_SWITCH_OFF = frozenset({"0", "false", "no", "off"})
+
+
+def _switch(env: str, default: bool = False) -> bool:
+ """A `TRTLLM_MODELING_V2_*` switch: unset or empty is `default`; any spelling outside
+ `_SWITCH_ON` / `_SWITCH_OFF` raises."""
+ raw = os.environ.get(env, "")
+ value = raw.strip().lower()
+ if not value:
+ return default
+ if value in _SWITCH_ON:
+ return True
+ if value in _SWITCH_OFF:
+ return False
+ raise ValueError(f"{env}={raw!r}: expected 1 / 0 (or true / false, yes / no, on / off)")
def _gemv_ready(x: torch.Tensor) -> bool:
@@ -802,7 +808,34 @@ def _precompile_mid(mid: dict) -> None:
_MTP_MOE_NVFP4_FC2_MARGIN = 4.0
_PLE_HOST_OFFLOAD_ENV = "TRTLLM_MODELING_V2_PLE_HOST_OFFLOAD"
-_PLE_HOST_OFFLOAD_OFF = frozenset({"0", "false", "no", "off"})
+
+#: The opt-in switches certified at tp1; every other layout runs the checkpoint's bf16 paths.
+_ONE_RANK_SWITCHES = (
+ _DENSE_FP8_ENV,
+ _GDN_INPROJ_NVFP4_ENV,
+ _TARGET_HEAD_NVFP4_ENV,
+ _MTP_HEAD_NVFP4_ENV,
+ _MTP_MOE_NVFP4_ENV,
+ _MTP_DENSE_FP8_ENV,
+)
+
+
+def _check_switches(one_rank: bool) -> None:
+ """Read every opt-in switch at construction, at any layout, so a misspelled value raises,
+ and warn about one set where it takes no effect: away from tp1, or the target head's NVFP4
+ copy without the dense FP8 mode whose head it replaces."""
+ on = [env for env in _ONE_RANK_SWITCHES if _switch(env)]
+ if on and not one_rank:
+ logger.warning(
+ f"{', '.join(on)}: certified at tp1 only; this layout runs the bf16 paths and "
+ "ignores them"
+ )
+ if one_rank and _TARGET_HEAD_NVFP4_ENV in on and _DENSE_FP8_ENV not in on:
+ logger.warning(
+ f"{_TARGET_HEAD_NVFP4_ENV}=1 takes effect only with {_DENSE_FP8_ENV}=1; the target "
+ "head stays bf16"
+ )
+
#: At one rank, context-step MoE launches of at least this many rows (and at most
#: `_CUTE_DSL_MOE_MAX_ROWS`, the op's certified top) run on the CuTe DSL MoE
@@ -813,7 +846,7 @@ def _precompile_mid(mid: dict) -> None:
_CUTE_DSL_MOE_MIN_ROWS = 4096 # capacity: fixed: a dispatch threshold, not a buffer size
_CUTE_DSL_MOE_MAX_ROWS = 16384 # capacity: fixed: the op's certified rows
# Whether one rank keeps that second copy (about 1.3 GiB a layer). No engine arg owns it, hence
-# an env var read at construction; "off" (the spellings above) drops the copy and every launch
+# an env var read at construction (`_switch`, default on); off drops the copy and every launch
# runs on trtllm-gen.
_MOE_PREFILL_COPY_ENV = "TRTLLM_MODELING_V2_MOE_PREFILL_COPY"
@@ -823,18 +856,18 @@ def _moe_prefill_copy(one_rank: bool, trunk_moe_backend: str, max_num_tokens: in
is local, the layout the CuteDSL MoE is certified at), only under trtllm-gen, only when a
step's rows (at most `max_num_tokens`) can reach the dispatch threshold, and not when the
env var turns it off."""
+ wanted = _switch(_MOE_PREFILL_COPY_ENV, default=True)
return (
one_rank
and trunk_moe_backend == "TRTLLM"
and max_num_tokens >= _CUTE_DSL_MOE_MIN_ROWS
- and os.environ.get(_MOE_PREFILL_COPY_ENV, "1").strip().lower() not in _PLE_HOST_OFFLOAD_OFF
+ and wanted
)
def _ple_host_offload_enabled() -> bool:
"""Whether the n-gram table is declared in host memory (the default)."""
- value = os.environ.get(_PLE_HOST_OFFLOAD_ENV, "1").strip().lower()
- return value not in _PLE_HOST_OFFLOAD_OFF
+ return _switch(_PLE_HOST_OFFLOAD_ENV, default=True)
def _moe_layer_scalars(w, i: int, start: int, count: int) -> dict[str, torch.Tensor]:
@@ -929,6 +962,8 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig
super().__init__(model_config)
cfg = model_config.pretrained_config
mapping = model_config.mapping
+ # One rank holds the whole model (tp1), which takes forms of its own (module docstring).
+ self.one_rank = mapping.world_size == 1
assert torch.cuda.get_device_capability() == _SM, (
f"target certified on sm_{_SM[0]}{_SM[1]}, running on "
f"sm_{''.join(map(str, torch.cuda.get_device_capability()))}"
@@ -1065,7 +1100,7 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig
# Fixes the load-time layout of the expert stacks (`weights.py`).
self.trunk_moe_backend = _TRUNK_MOE_BACKENDS[requested]
# The CuteDSL pipeline runs every expert at once: one rank only.
- assert self.trunk_moe_backend != "CUTEDSL" or mapping.world_size == 1, (
+ assert self.trunk_moe_backend != "CUTEDSL" or self.one_rank, (
"moe_config.backend=CUTEDSL is implemented at one rank (tp1) only; this layout "
f"has {mapping.world_size} ranks"
)
@@ -1078,7 +1113,7 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig
# At one rank the CuteDSL-layout copy of the experts for prefill-sized launches
# (`_moe_prefill_copy`); `weights.py` declares it from this.
self.moe_prefill_copy = _moe_prefill_copy(
- mapping.world_size == 1, self.trunk_moe_backend, model_config.max_num_tokens
+ self.one_rank, self.trunk_moe_backend, model_config.max_num_tokens
)
# The draft experts' NVFP4 copy at one rank (`_mtp_moe_nvfp4`), built once: it frees
# the FP8 experts, and the shell's post-load walk reaches `post_load_weights` again.
@@ -1097,6 +1132,7 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig
self.eos_token_id = int(eos[0] if isinstance(eos, (list, tuple)) else eos)
self.ple_table_dtype = _weights.PLE_TABLE_DTYPE
self.ple_host_offload = _ple_host_offload_enabled()
+ _check_switches(self.one_rank)
# CPU allocation behind the offloaded table; the `ple_table` parameter is a
# device-address view that owns none of those pages.
self._ple_table_host = None
@@ -1122,7 +1158,6 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig
# ---- this rank's share ----------------------------------------------
d = _weights.MODEL_WEIGHTS.dims(self)
- assert d.expert_split in _weights.LOADED_EXPERT_SPLITS, d.expert_split
self.heads_q = d.heads_q
self.heads_kv = d.heads_kv
self.kv_head_start = d.kv_head_start
@@ -1169,8 +1204,10 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig
# The shared expert's stream (see `_moe_block`); made after load, because a stream
# created inside a graph capture is not capturable.
self._side_stream: torch.cuda.Stream | None = None
- # The tensor-parallel sum of a block's partial output (`_tp_sum_op`), and the routed
- # experts' completion over this rank's rows (`_expert_sum_op`).
+ # The tensor-parallel sum of a block's partial output (`_tp_sum_op`, over the engine's
+ # `allreduce_strategy`), and the routed experts' completion over this rank's rows
+ # (`_expert_sum_op`).
+ self.allreduce_strategy = int(model_config.allreduce_strategy)
self.tp_sum = None
self.expert_sum = None
self._prefill: "PrefillTarget | None" = None
@@ -1262,6 +1299,12 @@ def post_load_weights(self) -> None:
f"sizes address {used_rows}"
)
assert table_rows % self.ngram_vocab_divisor == 0, (table_rows, self.ngram_vocab_divisor)
+ assert table_rows - used_rows < self.ngram_vocab_divisor, (
+ f"the n-gram table carries {table_rows - used_rows} padding rows, more than one "
+ f"divisor of {self.ngram_vocab_divisor}: its shards do not concatenate to the padded "
+ "id space the checkpoint declares, and every id past the first shard would read a "
+ "shifted row"
+ )
self._cos_sin = self._index_rope_table(self.theta, device)
# Before the draft layer's views, which read its fused first projections.
self._fuse_hc_first_projections()
@@ -1295,14 +1338,13 @@ def post_load_weights(self) -> None:
self.expert_sum = self._expert_sum_op()
# A decode step's bf16 projection entry, and its GDN in_proj's: at tp1 the one-rank
# dispatch (`_one_rank_projection`), at the other layouts the shapes' own.
- one_rank = self.mapping.world_size == 1
+ one_rank = self.one_rank
self.decode_mm = _one_rank_projection if one_rank else _decode_projection
# At one rank a trtllm-gen decode MoE leaves its finalize to the next site
# (`DecodeTarget._moe`).
self.moe_fold = one_rank and self.trunk_moe_backend == "TRTLLM"
# At one rank the hyper-connection sites' fast forms by row count (the decode GEMV and
- # tcgen05 rows pairs, the prefill rstd sites); at the other layouts one form serves every
- # batch size, so a request's tokens do not depend on what its steps carry beside it.
+ # tcgen05 rows pairs, the prefill rstd sites); the other layouts do not take them.
self.hc_fast_sites = one_rank
self.decode_in_mm = (
_one_rank_projection
@@ -1319,7 +1361,7 @@ def _build_draft_one_rank(self, device) -> None:
packed FP8 copies of its dense projections; None / empty at every other layout."""
self._mtp_attn_in = None
self._mtp_gemv8 = {}
- if self.mapping.world_size != 1:
+ if not self.one_rank:
return
w = self.w
p = f"l{self.mtp_layer_idx}_mtp_"
@@ -1349,7 +1391,7 @@ def _build_dense_copies(self, device) -> None:
(`_GDN_INPROJ_NVFP4_ENV`), with every decode variant compiled now; both only at tp1,
the layout they are certified at, and empty otherwise."""
w = self.w
- one_rank = self.mapping.world_size == 1
+ one_rank = self.one_rank
self._dense8, self._gemv8, self._mid, self._gemv_nv = {}, {}, {}, {}
if one_rank and _switch(_DENSE_FP8_ENV):
for name in w:
@@ -1384,11 +1426,12 @@ def _build_dense_copies(self, device) -> None:
def _tp_sum_op(self):
"""The sum that completes a block's per-rank partial output, settled once: under
- attention TP the plain all-reduce over the TP group (`comm/allreduce`, op `NONE`),
- which every rank gets the same bits back from; otherwise nothing to sum."""
+ attention TP the plain all-reduce over the TP group (`comm/allreduce`, op `NONE`, at
+ the engine's `allreduce_strategy`), which every rank gets the same bits back from;
+ otherwise nothing to sum."""
if self.attn_tp_size == 1:
return lambda x: x
- group = allreduce_group_args(self.mapping)
+ group = allreduce_group_args(self.mapping, self.allreduce_strategy)
def tp_sum(x: torch.Tensor) -> torch.Tensor:
return allreduce(
@@ -1558,6 +1601,13 @@ def _bind_cache_manager(self, md, position_ids) -> None:
mm_missing = [name for name in _MAMBA_FIELDS if not hasattr(mm, name)]
assert not mm_missing, f"mamba metadata fields missing: {mm_missing}"
rank_pools = _pool.take_pools(cm, g, self.attn_layers, self.gdn_layers, self.ple_layer)
+ # The pool serves both state dtypes; this layout serves its own (routing declines the rest).
+ served = _ONE_RANK_SSM_STATE_DTYPES if self.one_rank else _SSM_STATE_DTYPES
+ state_dtype = rank_pools.ssm[self.gdn_layers[0]].dtype
+ assert state_dtype in served, (
+ f"a {state_dtype} recurrent-state cache; this layout serves "
+ f"{' or '.join(str(d) for d in served)} (kv_cache_config.mamba_ssm_cache_dtype)"
+ )
# No sliding window: the targets pass `max_seq_len` as the window.
windows = cm.max_attention_window_vec
attn_windows = {windows[off] for off in rank_pools.pool_idx.values()}
@@ -2301,12 +2351,12 @@ def _mtp_front(core, M: _MtpWeights, embed_tokens, input_ids, hidden_states) ->
f"a draft step of {rows} rows exceeds max_num_tokens {M.unit_inject.shape[0]}"
)
e = _draft_embedding(input_ids, embed_tokens)
- x_e = mtp_cublas_mm(flashinfer_gemma_rmsnorm(e, M.pre_e_norm, core.eps), M.fc_e)
+ x_e = cublas_mm(flashinfer_gemma_rmsnorm(e, M.pre_e_norm, core.eps), M.fc_e)
_, h_n = qwen4_exp_hc_inject_norm(
hidden_states, None, None, M.pre_h_norm, core.eps, core.hc_count
)
x_h = torch.reshape(
- mtp_cublas_mm(torch.reshape(h_n, [rows * core.hc_count, core.hidden]), M.fc_h),
+ cublas_mm(torch.reshape(h_n, [rows * core.hc_count, core.hidden]), M.fc_h),
[rows, core.stream],
)
x, normed = qwen4_exp_hc_inject_norm(
@@ -2322,7 +2372,7 @@ def _mtp_qkv(core, M: _MtpWeights, mixed, pos3) -> tuple:
place at the module's rope base. Returns `(qkv, gate, k, v)`: `gate` is the
`[T, heads_q, head_dim]` gate half of the q projection, `k` and v` views of `qkv`."""
rows = mixed.shape[0]
- qg = mtp_cublas_mm(mixed, M.q)
+ qg = cublas_mm(mixed, M.q)
qkv = torch.empty([rows, core.qkv_width], dtype=mixed.dtype, device=mixed.device)
q_dst, k_dst, v_dst = torch.split(
torch.reshape(qkv, [rows, core.heads_q + 2 * core.heads_kv, core.head_dim]),
@@ -2335,9 +2385,9 @@ def _mtp_qkv(core, M: _MtpWeights, mixed, pos3) -> tuple:
2,
)
q_dst.copy_(q_src)
- k_dst.copy_(torch.reshape(mtp_cublas_mm(mixed, M.k), [rows, core.heads_kv, core.head_dim]))
- v_dst.copy_(torch.reshape(mtp_cublas_mm(mixed, M.v), [rows, core.heads_kv, core.head_dim]))
- mtp_fused_qk_norm_rope(
+ k_dst.copy_(torch.reshape(cublas_mm(mixed, M.k), [rows, core.heads_kv, core.head_dim]))
+ v_dst.copy_(torch.reshape(cublas_mm(mixed, M.v), [rows, core.heads_kv, core.head_dim]))
+ fused_qk_norm_rope(
qkv,
num_heads_q=core.heads_q,
num_heads_k=core.heads_kv,
@@ -2366,7 +2416,7 @@ def _mtp_index_branch(core, M: _MtpWeights, mixed, draft: _DraftBound, qsa) -> t
step completes. Returns the normalized, rotated index query."""
req_idx, logical, block_table, tokens_per_block, pos, coords = qsa
rows = mixed.shape[0]
- qk_index = mtp_cublas_mm(mixed, M.idx)
+ 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],
@@ -2684,7 +2734,7 @@ def _decode_router_mid(core, L: _LayerWeights):
"""A decode step's router projection for layer `L`: `mid_rows_mm` over its view in the
dense FP8 mode's rows (`_MID_MM_ROWS["router"]`), else at tp1 the step's projection entry
(`core.decode_mm`) and at the other layouts `_decode_router`."""
- fallback = core.decode_mm if core.mapping.world_size == 1 else _decode_router
+ fallback = core.decode_mm if core.one_rank else _decode_router
if L.router_m is None:
return fallback
return lambda x, router: _dense_mm(x, router, None, fallback, "router", mid=L.router_m)
@@ -3072,6 +3122,64 @@ def _selection(
)
+#: Elements of one score table: `qsa_paged_index_scores` writes its output rows and
+#: `qsa_topk_expand_decode` lays out its rows with int32 index arithmetic (row times the
+#: table's width), so a call scores fewer rows than this over the table's columns.
+_SCORE_TABLE_ELEMENTS = 2**31
+
+
+def _score_and_select(
+ core,
+ index_cache,
+ q_index,
+ block_table,
+ logical,
+ req_idx,
+ tokens_per_block: int,
+ parts,
+ visible_blocks,
+ seq_lens_q,
+ topk_indices,
+ row_starts,
+ selected,
+ live_cols,
+) -> None:
+ """Score each part's rows against the index cache and select into their rows of
+ `selected` (`_selection`). A part's rows go in chunks of fewer than
+ `_SCORE_TABLE_ELEMENTS / columns` rows: a prefill step's context rows reach
+ `max_num_tokens`, and at `max_seq_len` 262144 (65568 columns) more than 32752 of them
+ pass the int32 range. A row's scores and selection read only that row's query, position
+ and request, so the chunks select what one call would; below the bound a part is one
+ call."""
+ columns = block_table.shape[1] * tokens_per_block // core.compress_ratio
+ chunk = max(1, (_SCORE_TABLE_ELEMENTS - 1) // columns)
+ for lo, hi, context_rows in parts:
+ for a in range(lo, hi, chunk):
+ b = min(a + chunk, hi)
+ scores = qsa_paged_index_scores(
+ q=q_index[a:b],
+ index_cache=index_cache,
+ block_table=block_table,
+ query_positions=logical[a:b],
+ request_indices=req_idx[a:b],
+ tokens_per_block=tokens_per_block,
+ compress_ratio=core.compress_ratio,
+ only_visible_blocks=True,
+ context_rows=context_rows,
+ )
+ _selection(
+ core,
+ scores,
+ visible_blocks[a:b],
+ logical[a:b],
+ seq_lens_q[a:b],
+ topk_indices[a:b],
+ row_starts[a:b],
+ selected[a:b],
+ live_cols,
+ )
+
+
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
@@ -3594,9 +3702,7 @@ def _attention_core(self, proj, proj_qkv, index_scratch, attn_out, i) -> None:
qkv = proj_qkv[:num_tokens]
qkv.copy_(proj[:, : core.qkv_width])
if step.stage:
- core.spec_rollback.stage_index(
- core._bound.pools.index_cache[i], core._bound.pools.position_cache
- )
+ core.spec_rollback.stage_index(core._bound.pools.position_cache)
_index_plane(core, proj, i, L, (req_idx, logical, block_table, tokens_per_block, coords))
if sparse:
@@ -3617,32 +3723,25 @@ def _attention_core(self, proj, proj_qkv, index_scratch, attn_out, i) -> None:
dtype=torch.int32,
device=proj.device,
)
- for lo, hi, context_rows in parts:
- scores = qsa_paged_index_scores(
- q=q_index[lo:hi],
- index_cache=core._bound.pools.index_cache[i],
- block_table=block_table,
- query_positions=logical[lo:hi],
- request_indices=req_idx[lo:hi],
- tokens_per_block=tokens_per_block,
- compress_ratio=core.compress_ratio,
- only_visible_blocks=True,
- context_rows=context_rows,
- )
- _selection(
- core,
- scores,
- visible_blocks[lo:hi],
- logical[lo:hi],
- seq_lens_q[lo:hi],
- topk_indices[lo:hi],
- topk_row_starts[lo:hi],
- selected[lo:hi],
- step.live_cols,
- )
+ _score_and_select(
+ core,
+ core._bound.pools.index_cache[i],
+ q_index,
+ block_table,
+ logical,
+ req_idx,
+ tokens_per_block,
+ parts,
+ visible_blocks,
+ seq_lens_q,
+ topk_indices,
+ topk_row_starts,
+ selected,
+ step.live_cols,
+ )
# Context rows with the prefill pair (a bound on a row's scan), generation rows
# without.
- if core.mapping.world_size == 1 and kv.k_scale is not None:
+ if core.one_rank and kv.k_scale is not None:
_sparse_gqa_nvfp4_context(
core,
kv,
@@ -4113,7 +4212,7 @@ def prepare(
fp8_pages = _fp8_context_pages(md, md.tokens_per_block)
nvfp4_context = None
if (
- core.mapping.world_size == 1
+ core.one_rank
and sparse
and core.quant_mode == QM_NVFP4_KV_CACHE
and num_prefill_tokens > _NVFP4_ROWTILE_MIN_ROWS
@@ -4397,7 +4496,7 @@ def mtp_forward(
self._capture_selection(md, topk_indices, visible_blocks, compact=True)
q_att = _packed_q(core, qkv)[keep]
attn = torch.empty([q_rows, core.q_width], dtype=qkv.dtype, device=qkv.device)
- if kv.k_scale is not None and core.mapping.world_size == 1:
+ if kv.k_scale is not None and core.one_rank:
# One row per context request, then `p` per generation request: the split-K
# kernel for the first, the verify kernel for the rest where it serves.
_sparse_gqa_nvfp4_context(
@@ -4435,29 +4534,22 @@ def mtp_forward(
dtype=torch.int32,
device=qkv.device,
)
- for lo, hi, context_rows in parts:
- scores = qsa_paged_index_scores(
- q=q_index[lo:hi],
- index_cache=draft.index_cache,
- block_table=block_table,
- query_positions=logical[lo:hi],
- request_indices=req_idx[lo:hi],
- tokens_per_block=tokens_per_block,
- compress_ratio=core.compress_ratio,
- only_visible_blocks=True,
- context_rows=context_rows,
- )
- _selection(
- core,
- scores,
- visible_blocks[lo:hi],
- logical[lo:hi],
- seq_lens_q[lo:hi],
- topk_indices[lo:hi],
- row_starts[lo:hi],
- selected[lo:hi],
- _live_score_columns(core, md, batch),
- )
+ _score_and_select(
+ core,
+ draft.index_cache,
+ q_index,
+ block_table,
+ logical,
+ req_idx,
+ tokens_per_block,
+ parts,
+ visible_blocks,
+ seq_lens_q,
+ topk_indices,
+ row_starts,
+ selected,
+ _live_score_columns(core, md, batch),
+ )
if share:
self._capture_selection(md, topk_indices, visible_blocks)
attn = torch.empty([rows, core.q_width], dtype=qkv.dtype, device=qkv.device)
@@ -4487,7 +4579,7 @@ def mtp_forward(
**step,
**_CALL_CONSTANTS,
)
- mtp_fused_sigmoid_mul_inplace(attn, gate)
+ fused_sigmoid_mul_inplace(attn, gate)
out = x
if keep is not None:
# The sparse path attended the kept rows only; the dense call wrote every row.
@@ -4578,7 +4670,7 @@ def _attention_prologue(self, proj, q_index, index_view, index_cache, L, qsa, si
1,
)[1]
if stage:
- core.spec_rollback.stage_index(index_cache, pools.position_cache)
+ core.spec_rollback.stage_index(pools.position_cache)
qsa_paged_index_store(
token_k,
coords,
@@ -4639,7 +4731,7 @@ def _attention_one_rank(self, x, i, L, qkv, index_scratch, attn_out, qsa, step,
if fused_prologue and stage:
side.wait_stream(main)
with torch.cuda.stream(side):
- core.spec_rollback.stage_index(index_cache, pools.position_cache)
+ core.spec_rollback.stage_index(pools.position_cache)
proj = _dense_mm(
x, L.attn_in, L.attn_in8, core.decode_mm, "in", L.attn_in_g, mid=L.attn_in_m
)
@@ -4829,7 +4921,7 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step, sta
index slots are staged for `SpeculativeRollback` before the store. At one rank
`_attention_one_rank` runs instead."""
core = self.core
- if core.mapping.world_size == 1:
+ if core.one_rank:
return self._attention_one_rank(x, i, L, qkv, index_scratch, attn_out, qsa, step, stage)
(
req_idx,
@@ -4857,7 +4949,7 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step, sta
# index store and compression are one launch; it writes the packed buffer itself.
prologue = sparse and pools.kv[i].k_scale is None
if stage:
- core.spec_rollback.stage_index(pools.index_cache[i], pools.position_cache)
+ core.spec_rollback.stage_index(pools.position_cache)
if prologue:
kv = pools.kv[i]
qwen4_exp_attn_decode_prologue(
@@ -5540,7 +5632,7 @@ def _gdn_block(self, x, i, L, si_all, verify, gen_p, captured=False) -> torch.Te
)
if (
captured
- and core.mapping.world_size == 1
+ and core.one_rank
and pools.ssm[i].dtype == torch.float32
and num_tokens <= _GDN_DECODE_FUSED_MAX_ROWS
):
@@ -5905,7 +5997,7 @@ def mtp_forward(
block_table = md.qsa_block_table
tokens_per_block = md.tokens_per_block
- if core.mapping.world_size == 1:
+ if core.one_rank:
return self._mtp_layer(
embed_tokens,
input_ids,
@@ -6018,7 +6110,7 @@ def mtp_forward(
**step,
**_CALL_CONSTANTS,
)
- mtp_fused_sigmoid_mul_inplace(attn, gate)
+ fused_sigmoid_mul_inplace(attn, gate)
return _mtp_tail(core, M, x, attn, inject, pad_rows)
@@ -6270,7 +6362,7 @@ def stage_ple_conv_input(self, conv_input: torch.Tensor) -> None:
conv_x = torch.reshape(conv_input[prefix:], [n, p, -1])
pending["conv_x"] = conv_x.clone() if prefix else conv_x
- def stage_index(self, index_cache: torch.Tensor, position_cache: torch.Tensor) -> None:
+ def stage_index(self, position_cache: torch.Tensor) -> None:
"""Before an attention layer's index store: on the step's first attention layer (the
first to write the shared plane) the generation rows' position slots. The index-K
slots of every layer were gathered when the step was opened."""
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_rowtile_attention_cute__prefill.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_rowtile_attention_cute__prefill.py
deleted file mode 100644
index 4fa6fef27d45..000000000000
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_rowtile_attention_cute__prefill.py
+++ /dev/null
@@ -1,1386 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""The row-tile pass of `qsa_sparse_gqa_rowtile` on SM103 tensor cores (CuTe DSL, tcgen05).
-
-Inputs and output: per group of `rows_per_tile` rows of one request, the plan's packed
-slot list (`slots[group]`, `listed` entries), each row's bitmap re-indexed onto it
-(`packed_bits`, as int32 halves) and `group_plan[group] = (tiles, listed)`; groups with
-`tiles < 0` belong to the lone-row pass.
-
-A persistent kernel: one CTA per SM, each taking work items -- a (group, KV head) pair, 128
-MMA rows = `rows_per_tile` query rows x `group_size` heads -- from a global ticket counter,
-longest-first (the plan's last group first). 9 warps:
-
-* warp 5 lane 0 draws the tickets and publishes the items through a shared-memory ring that
- every role reads; warps 5-8 gather each item's Q rows (into a 60 KB Q buffer) and its list's
- K and V rows (64 a tile) with `cp.async` into one ring of `kv_slots` 32 KB SW128 slots, filled
- in the MMA warp's order K0, K1, V0, K2, V1, ...; the next item's Q is fetched while the
- current item's last tiles load;
-* warp 4 issues the tcgen05 MMAs: S[t % 2] = Q K(t)^T (Q from tensor memory), then
- O += P(t - 1) V(t - 1) (P from tensor memory, overwriting the S it came from), `t` counting
- the CTA's tiles across items;
-* warps 0-3, one thread per MMA row, per tile load S with its row max (`tcgen05.ld.red`),
- rescale O when that max grew by more than `rescale_threshold` in log2 units (lazily;
- otherwise exponentiate against the stale max), write P = 2^(s - max) on the row's own bits
- and 0 elsewhere. After an item's last tile they copy the next item's Q from the Q buffer
- into tensor memory (so its first QK MMAs overlap this epilogue), then normalize O, stage it
- in the Q buffer as four 64-dim SW128 boxes of the item's G heads x R rows and store it with
- TMA (thread 0 hands the buffer back to the loaders once the store has read it).
-
-The Q buffer is shared by Q and the staged O: a softmax thread's Q row and its O row occupy
-the same bytes, so the next item's Q is read out before this item's O is staged over it. An
-item followed by a lone (or empty) group stores O row by row from registers instead, as the
-loaders may then fetch the following item's Q into the buffer at any time.
-
-Tensor memory: O in columns 0-255, S/P buffers at 256 and 320, Q at 384 (512 columns).
-The tile's row max covers every listed entry, an upper bound of the row's own: a valid
-softmax offset. P is rounded to bf16 for the PV MMA and summed in fp32 unrounded.
-"""
-
-from __future__ import annotations
-
-import functools
-
-import cuda.bindings.driver as cuda_drv
-import cutlass
-import cutlass.cute as cute
-import cutlass.cute.nvgpu.tcgen05 as tcgen05
-import cutlass.pipeline as pipeline
-import cutlass.utils.blackwell_helpers as sm100_utils
-import torch
-from cutlass._mlir.dialects import nvvm as _nvvm
-from cutlass.cute.nvgpu import OperandMajorMode, cpasync
-from cutlass.cute.runtime import make_ptr
-from cutlass.cute.typing import Float32, Int32, Int64
-from cutlass.cutlass_dsl import dsl_user_op
-
-try:
- from cutlass.memory import SmemAllocator
-except ImportError: # DSL releases before the memory namespace split
- from cutlass.utils import SmemAllocator
-
-BF16 = cutlass.BFloat16
-_LOG2_E = 1.4426950408889634
-#: The head dimension the tensor-memory plan is laid out for.
-HEAD_DIM = 256
-#: MMA rows per CTA and tokens per tile (one 64-bit bitmap word).
-BLOCK_M = 128
-BLOCK_N = 64
-#: K/V ring slots (32 KB each); with the 60 KB Q buffer, about 220 KB of shared memory.
-KV_SLOTS = 5
-#: Work items in flight in the shared-memory schedule ring.
-SCHED_SLOTS = 8
-#: The next item's Q is fetched once this many tiles of the current item remain to load.
-Q_LEAD = 3
-
-
-@dsl_user_op
-def tcgen05_fence_before(*, loc=None, ip=None):
- """tcgen05.fence::before_thread_sync: order this thread's prior tcgen05 ops before a thread sync."""
- _nvvm.tcgen05_fence(_nvvm.Tcgen05FenceKind.BEFORE_THREAD_SYNC, loc=loc, ip=ip)
-
-
-@dsl_user_op
-def tcgen05_fence_after(*, loc=None, ip=None):
- """tcgen05.fence::after_thread_sync: order later tcgen05 ops after a thread sync."""
- _nvvm.tcgen05_fence(_nvvm.Tcgen05FenceKind.AFTER_THREAD_SYNC, loc=loc, ip=ip)
-
-
-class _RowtileAttentionKernel:
- def __init__(
- self,
- group_size: int,
- rows_per_tile: int,
- head_dim: int = 256,
- tokens_per_block: int = 32,
- kv_slots: int = KV_SLOTS,
- rescale_threshold: float = 8.0,
- use_pdl: bool = True,
- ):
- assert head_dim == 256, (
- "the TMEM plan (O 256 + 2 x S 64 + Q 128 columns) is laid out for D = 256"
- )
- assert group_size * rows_per_tile <= 128
- # The O staging's four SW128 boxes sit `64 G R` elements apart: each must start on
- # a 1024-byte swizzle atom, so the live MMA rows come in whole 8-row atoms.
- assert group_size * rows_per_tile % 8 == 0, (group_size, rows_per_tile)
- self.group_size = group_size
- self.rows_per_tile = rows_per_tile
- self.d = head_dim
- self.tpb = tokens_per_block
- self.bm = 128
- self.bn = 64
- self.kv_slots = kv_slots
- self.sched_slots = SCHED_SLOTS
- self.q_lead = Q_LEAD
- self.tau = rescale_threshold
- self.use_pdl = use_pdl
- self.mma_warp = 4
- self.load_warp0 = 5
- self.threads = 32 * 9
- self.tmem_o = 0
- self.tmem_s0 = 256
- self.tmem_q = 384 # Q as the QK MMA's A operand: 256 bf16 a row, two a column
- self.tmem_cols = 512
- self.chunk_cols = 32 # O columns a TMEM load/store moves per thread
-
- @cute.jit
- def __call__(
- self,
- q: cute.Pointer,
- k: cute.Pointer,
- v: cute.Pointer,
- slots: cute.Pointer,
- bits: cute.Pointer,
- plan: cute.Pointer,
- out: cute.Pointer,
- work: cute.Pointer,
- rows: Int32,
- num_groups: Int32,
- num_kv_heads: Int32,
- num_ctas: Int32,
- q_stride_row: Int32,
- k_stride_page: Int32,
- k_stride_head: Int32,
- k_stride_token: Int32,
- v_stride_page: Int32,
- v_stride_head: Int32,
- v_stride_token: Int32,
- slots_stride_group: Int32,
- bits_stride_row: Int32,
- out_stride_row: Int32,
- scale_log2: Float32,
- stream,
- ):
- def qk():
- return sm100_utils.make_trivial_tiled_mma(
- BF16,
- BF16,
- OperandMajorMode.K,
- OperandMajorMode.K,
- Float32,
- tcgen05.CtaGroup.ONE,
- (self.bm, self.bn),
- tcgen05.OperandSource.TMEM,
- )
-
- def pv():
- return sm100_utils.make_trivial_tiled_mma(
- BF16,
- BF16,
- OperandMajorMode.K,
- OperandMajorMode.MN,
- Float32,
- tcgen05.CtaGroup.ONE,
- (self.bm, self.d),
- tcgen05.OperandSource.TMEM,
- )
-
- # Overwrite / accumulate twins of each MMA: the ACCUMULATE field is set once at kernel
- # scope, so no loop or branch rebinds an MMA value.
- qk_mma, qk_mma_acc, pv_mma, pv_mma_acc = qk(), qk(), pv(), pv()
- q_layout = sm100_utils.make_smem_layout_a(qk_mma, (self.bm, self.bn, self.d), BF16, 1)
- k_layout = sm100_utils.make_smem_layout_b(
- qk_mma, (self.bm, self.bn, self.d), BF16, self.kv_slots
- )
- v_layout = sm100_utils.make_smem_layout_b(
- pv_mma, (self.bm, self.d, self.bn), BF16, self.kv_slots
- )
- p_layout = sm100_utils.make_smem_layout_a(pv_mma, (self.bm, self.d, self.bn), BF16, 1)
-
- @cute.struct
- class SharedStorage:
- q_full: cute.struct.MemRange[Int64, 1]
- kv_full: cute.struct.MemRange[Int64, self.kv_slots]
- kv_empty: cute.struct.MemRange[Int64, self.kv_slots]
- s_full: cute.struct.MemRange[Int64, 2]
- p_full: cute.struct.MemRange[Int64, 2]
- o_full: cute.struct.MemRange[Int64, 1]
- o_final: cute.struct.MemRange[Int64, 1]
- qs_full: cute.struct.MemRange[Int64, 1]
- qs_empty: cute.struct.MemRange[Int64, 1]
- sch_full: cute.struct.MemRange[Int64, self.sched_slots]
- sch_empty: cute.struct.MemRange[Int64, self.sched_slots]
- sch_val: cute.struct.MemRange[Int32, self.sched_slots]
- tmem_holding: Int32
- sKV: cute.struct.Align[cute.struct.MemRange[BF16, cute.cosize(k_layout)], 1024]
- # The Q buffer / O staging: four 64-dim boxes of the live G x R MMA rows.
- sQ: cute.struct.Align[
- cute.struct.MemRange[BF16, self.group_size * self.rows_per_tile * self.d], 1024
- ]
-
- self.shared_storage = SharedStorage
-
- # The epilogue's TMA store: out viewed as (D, heads, rows), one box = 64 dims of the
- # group's G heads x R rows, staged in the Q buffer as 128-byte SW128 rows (MMA row m =
- # local row x G + head), four boxes 64 G R elements apart.
- G = self.group_size
- R = self.rows_per_tile
- o_box = cute.make_composed_layout(
- cute.make_swizzle(3, 4, 3), 0, cute.make_layout((64, G, R), stride=(1, 64, 64 * G))
- )
- mO = cute.make_tensor(
- out,
- cute.make_layout((self.d, num_kv_heads * G, rows), stride=(1, self.d, out_stride_row)),
- )
- tma_atom_o, tma_tensor_o = cpasync.make_tiled_tma_atom(
- cpasync.CopyBulkTensorTileS2GOp(), mO, o_box, (64, G, R)
- )
-
- self.kernel(
- qk_mma,
- qk_mma_acc,
- pv_mma,
- pv_mma_acc,
- q_layout,
- k_layout,
- v_layout,
- p_layout,
- tma_atom_o,
- tma_tensor_o,
- q,
- k,
- v,
- slots,
- bits,
- plan,
- out,
- work,
- rows,
- num_groups,
- num_kv_heads,
- q_stride_row,
- k_stride_page,
- k_stride_head,
- k_stride_token,
- v_stride_page,
- v_stride_head,
- v_stride_token,
- slots_stride_group,
- bits_stride_row,
- out_stride_row,
- scale_log2,
- ).launch(
- grid=[num_ctas, 1, 1],
- block=[self.threads, 1, 1],
- stream=stream,
- min_blocks_per_mp=1,
- use_pdl=self.use_pdl,
- )
-
- @cute.kernel
- def kernel(
- self,
- qk_mma: cute.TiledMma,
- qk_mma_acc: cute.TiledMma,
- pv_mma: cute.TiledMma,
- pv_mma_acc: cute.TiledMma,
- q_layout: cute.ComposedLayout,
- k_layout: cute.ComposedLayout,
- v_layout: cute.ComposedLayout,
- p_layout: cute.ComposedLayout,
- tma_atom_o: cute.CopyAtom,
- tma_tensor_o: cute.Tensor,
- q: cute.Pointer,
- k: cute.Pointer,
- v: cute.Pointer,
- slots: cute.Pointer,
- bits: cute.Pointer,
- plan: cute.Pointer,
- out: cute.Pointer,
- work: cute.Pointer,
- rows: Int32,
- num_groups: Int32,
- num_kv_heads: Int32,
- q_stride_row: Int32,
- k_stride_page: Int32,
- k_stride_head: Int32,
- k_stride_token: Int32,
- v_stride_page: Int32,
- v_stride_head: Int32,
- v_stride_token: Int32,
- slots_stride_group: Int32,
- bits_stride_row: Int32,
- out_stride_row: Int32,
- scale_log2: Float32,
- ):
- warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
- tidx, _, _ = cute.arch.thread_idx()
- qk_mma.set(tcgen05.Field.ACCUMULATE, False)
- qk_mma_acc.set(tcgen05.Field.ACCUMULATE, True)
- pv_mma.set(tcgen05.Field.ACCUMULATE, False)
- pv_mma_acc.set(tcgen05.Field.ACCUMULATE, True)
-
- smem = SmemAllocator()
- storage = smem.allocate(self.shared_storage)
- q_full = storage.q_full.data_ptr()
- kv_full = storage.kv_full.data_ptr()
- kv_empty = storage.kv_empty.data_ptr()
- s_full = storage.s_full.data_ptr()
- p_full = storage.p_full.data_ptr()
- o_full = storage.o_full.data_ptr()
- o_final = storage.o_final.data_ptr()
- qs_full = storage.qs_full.data_ptr()
- qs_empty = storage.qs_empty.data_ptr()
- sch_full = storage.sch_full.data_ptr()
- sch_empty = storage.sch_empty.data_ptr()
- sch_val = storage.sch_val.data_ptr()
- if tidx == 0:
- cute.arch.mbarrier_init(q_full, 128)
- for i in cutlass.range_constexpr(self.kv_slots):
- cute.arch.mbarrier_init(kv_full + i, 128)
- cute.arch.mbarrier_init(kv_empty + i, 1)
- for i in cutlass.range_constexpr(2):
- cute.arch.mbarrier_init(s_full + i, 1)
- cute.arch.mbarrier_init(p_full + i, 128)
- cute.arch.mbarrier_init(o_full, 1)
- cute.arch.mbarrier_init(o_final, 1)
- cute.arch.mbarrier_init(qs_full, 128)
- # One arrival a hand-back: softmax thread 0, after the softmax warps are done with it.
- cute.arch.mbarrier_init(qs_empty, 1)
- for i in cutlass.range_constexpr(self.sched_slots):
- cute.arch.mbarrier_init(sch_full + i, 1)
- # One elected lane of each of the 9 warps releases an item.
- cute.arch.mbarrier_init(sch_empty + i, 9)
- cute.arch.mbarrier_init_fence()
-
- tmem_users = pipeline.NamedBarrier(barrier_id=1, num_threads=32 * 5)
- if warp_idx == self.mma_warp:
- cute.arch.alloc_tmem(self.tmem_cols, storage.tmem_holding.ptr)
- cute.arch.relinquish_tmem_alloc_permit()
- cute.arch.sync_threads()
-
- # Slices taken before the role branches: the MMA warp's `set(ACCUMULATE)` rebinds the
- # tiled MMAs inside its branch, which the other branches must not see.
- qk_thr = qk_mma.get_slice(0)
- pv_thr = pv_mma.get_slice(0)
- sK = storage.sKV.get_tensor(k_layout.outer, swizzle=k_layout.inner)
- sV = cute.make_tensor(cute.recast_ptr(sK.iterator, v_layout.inner), v_layout.outer)
- sQ = storage.sQ.data_ptr()
- sO = storage.sQ.get_tensor(
- cute.make_layout(
- (64, self.group_size, self.rows_per_tile, self.d // 64),
- stride=(1, 64, 64 * self.group_size, 64 * self.group_size * self.rows_per_tile),
- ),
- swizzle=cute.make_swizzle(3, 4, 3),
- )
-
- if warp_idx >= self.load_warp0:
- self.load_role(
- tidx - 32 * self.load_warp0,
- sK,
- sQ,
- q,
- k,
- v,
- slots,
- plan,
- work,
- rows,
- num_groups,
- num_kv_heads,
- q_stride_row,
- k_stride_page,
- k_stride_head,
- k_stride_token,
- v_stride_page,
- v_stride_head,
- v_stride_token,
- slots_stride_group,
- kv_full,
- kv_empty,
- qs_full,
- qs_empty,
- sch_full,
- sch_empty,
- sch_val,
- )
- elif warp_idx == self.mma_warp:
- tmem_ptr = cute.arch.retrieve_tmem_ptr(Float32, 16, storage.tmem_holding.ptr)
- self.mma_role(
- qk_mma,
- qk_mma_acc,
- pv_mma,
- pv_mma_acc,
- qk_thr,
- pv_thr,
- q_layout,
- p_layout,
- sK,
- sV,
- tmem_ptr,
- plan,
- num_kv_heads,
- q_full,
- kv_full,
- kv_empty,
- s_full,
- p_full,
- o_full,
- o_final,
- sch_full,
- sch_empty,
- sch_val,
- )
- tmem_users.arrive_and_wait()
- cute.arch.dealloc_tmem(tmem_ptr, self.tmem_cols)
- else:
- tmem_ptr = cute.arch.retrieve_tmem_ptr(Float32, 16, storage.tmem_holding.ptr)
- self.softmax_role(
- qk_thr,
- pv_thr,
- tmem_ptr,
- tidx,
- plan,
- num_kv_heads,
- rows,
- bits,
- bits_stride_row,
- out,
- out_stride_row,
- scale_log2,
- q_full,
- s_full,
- p_full,
- o_full,
- o_final,
- qs_full,
- qs_empty,
- sch_full,
- sch_empty,
- sch_val,
- sQ,
- sO,
- tma_atom_o,
- tma_tensor_o,
- )
- tmem_users.arrive_and_wait()
-
- if cutlass.const_expr(self.use_pdl):
- # The lone-row pass may still run: finish only after it has, so a later launch sees every row.
- cute.arch.griddepcontrol_wait()
-
- # ------------------------------------------------------------------ schedule
- @cute.jit
- def read_item(self, n: Int32, sch_full, sch_val):
- """Item `n` of this CTA (group * kv_heads + head, or -1 past the last), once published."""
- slot = n % self.sched_slots
- cute.arch.mbarrier_wait(sch_full + slot, (n // self.sched_slots) % 2)
- return cute.make_tensor(sch_val, cute.make_layout((self.sched_slots,)))[slot]
-
- @cute.jit
- def release_item(self, n: Int32, sch_empty):
- with cute.arch.elect_one():
- cute.arch.mbarrier_arrive(sch_empty + n % self.sched_slots)
-
- @cute.jit
- def item_info(self, item: Int32, plan: cute.Pointer, num_kv_heads: Int32):
- """(group, KV head, tiles, listed) of a published item; tiles = -1 past the last."""
- group = Int32(0)
- head = Int32(0)
- tiles = Int32(-1)
- listed = Int32(0)
- if item >= 0:
- group = item // num_kv_heads
- head = item - group * num_kv_heads
- mPlan = cute.make_tensor(plan, cute.make_layout((1 << 30,)))
- tiles = mPlan[group * 2]
- listed = mPlan[group * 2 + 1]
- return group, head, tiles, listed
-
- # ------------------------------------------------------------------ loaders (warps 5-8)
- @cute.jit
- def publish(self, n: Int32, ticket: Int32, num_items: Int32, num_groups: Int32,
- num_kv_heads: Int32, sch_full, sch_empty, sch_val): # fmt: skip
- """Item `n` from a ticket: the plan's groups last to first, both heads of a group in turn."""
- item = Int32(-1)
- if ticket < num_items:
- back = ticket // num_kv_heads
- item = (num_groups - 1 - back) * num_kv_heads + (ticket - back * num_kv_heads)
- slot = n % self.sched_slots
- cute.arch.mbarrier_wait(sch_empty + slot, ((n // self.sched_slots) % 2) ^ 1)
- cute.make_tensor(sch_val, cute.make_layout((self.sched_slots,)))[slot] = item
- cute.arch.mbarrier_arrive(sch_full + slot)
-
- @cute.jit
- def load_role(
- self,
- lt: Int32,
- sK,
- sQ: cute.Pointer,
- q: cute.Pointer,
- k: cute.Pointer,
- v: cute.Pointer,
- slots: cute.Pointer,
- plan: cute.Pointer,
- work: cute.Pointer,
- rows: Int32,
- num_groups: Int32,
- num_kv_heads: Int32,
- q_stride_row: Int32,
- k_stride_page: Int32,
- k_stride_head: Int32,
- k_stride_token: Int32,
- v_stride_page: Int32,
- v_stride_head: Int32,
- v_stride_token: Int32,
- slots_stride_group: Int32,
- kv_full,
- kv_empty,
- qs_full,
- qs_empty,
- sch_full,
- sch_empty,
- sch_val,
- ):
- """Per item: its Q rows into the Q buffer (unless fetched during the previous item), then
- the ring in the MMA warp's order K0, K1, V0, K2, V1, ..., V(T-1). Thread `lt` gathers
- tokens 16 (lt // 32) + 4 t + (lt % 32) // 8 (t < 4), one 16-byte chunk (lt % 8) of each of
- their four 128-byte dim blocks; the slots of tile j + 1 are fetched while tile j's copies
- issue. Lane 0 of warp 5 publishes item n + 1 when the warps start item n.
- """
- D = self.d
- atom = cute.make_copy_atom(
- cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), BF16, num_bits_per_copy=128
- )
- chunk = lt % 8
- # Element-coordinate view of the swizzled ring: (8 elements, (chunk, 64-dim block), token, slot).
- sKV_c = cute.make_tensor(
- sK.iterator,
- cute.make_layout(
- (8, (8, D // 64), self.bn, self.kv_slots),
- stride=(1, (8, self.bn * 64), 64, self.bn * D),
- ),
- )
- num_items = num_groups * num_kv_heads
- first_token = 16 * (lt // 32) + (lt % 32) // 8
- ticket = Int32(0)
- if lt == 0:
- ticket = cute.arch.atomic_add(work, Int32(1))
- self.publish(
- Int32(0), ticket, num_items, num_groups, num_kv_heads, sch_full, sch_empty, sch_val
- )
- ticket = cute.arch.atomic_add(work, Int32(1))
- n = Int32(0)
- ring = Int32(0)
- ring_phase = Int32(1)
- qs_phase = Int32(1)
- q_ready = Int32(0)
- item = self.read_item(n, sch_full, sch_val)
- while item >= 0:
- if lt == 0:
- self.publish(
- n + 1, ticket, num_items, num_groups, num_kv_heads, sch_full, sch_empty, sch_val
- )
- ticket = cute.arch.atomic_add(work, Int32(1))
- group, kv_head, tiles, listed = self.item_info(item, plan, num_kv_heads)
- nitem = self.read_item(n + 1, sch_full, sch_val)
- ngroup, nhead, ntiles, _ = self.item_info(nitem, plan, num_kv_heads)
- if tiles > 0:
- if q_ready == 0:
- cute.arch.mbarrier_wait(qs_empty, qs_phase)
- qs_phase = qs_phase ^ 1
- self.issue_q(lt, group, kv_head, rows, q, q_stride_row, sQ)
- cute.arch.cp_async_mbarrier_arrive_noinc(qs_full)
- q_ready = Int32(0)
- q_at = cutlass.max(tiles - self.q_lead, Int32(0))
- mSlots = cute.make_tensor(
- slots + Int64(group) * slots_stride_group, cute.make_layout((1 << 30,))
- )
- slot_next = cute.make_rmem_tensor((4,), Int32)
- k_cur = cute.make_rmem_tensor((4,), Int64)
- v_cur = cute.make_rmem_tensor((4,), Int64)
- v_prev = cute.make_rmem_tensor((4,), Int64)
- for t in cutlass.range_constexpr(4):
- slot_next[t] = mSlots[cutlass.min(first_token + 4 * t, listed - 1)]
- for j in cutlass.range(0, tiles, 1, unroll=1):
- for t in cutlass.range_constexpr(4):
- slot = slot_next[t]
- page = slot // self.tpb
- tok = slot % self.tpb
- k_cur[t] = (
- Int64(page) * k_stride_page
- + kv_head * k_stride_head
- + tok * k_stride_token
- )
- v_cur[t] = (
- Int64(page) * v_stride_page
- + kv_head * v_stride_head
- + tok * v_stride_token
- )
- if j + 1 < tiles:
- for t in cutlass.range_constexpr(4):
- slot_next[t] = mSlots[
- cutlass.min((j + 1) * self.bn + first_token + 4 * t, listed - 1)
- ]
- # K(j)
- cute.arch.mbarrier_wait(kv_empty + ring, ring_phase)
- self.issue_tile(atom, k, k_cur, sKV_c, ring, chunk, first_token)
- cute.arch.cp_async_mbarrier_arrive_noinc(kv_full + ring)
- ring, ring_phase = self.advance(ring, ring_phase)
- if j > 0:
- # V(j - 1)
- cute.arch.mbarrier_wait(kv_empty + ring, ring_phase)
- self.issue_tile(atom, v, v_prev, sKV_c, ring, chunk, first_token)
- cute.arch.cp_async_mbarrier_arrive_noinc(kv_full + ring)
- ring, ring_phase = self.advance(ring, ring_phase)
- for t in cutlass.range_constexpr(4):
- v_prev[t] = v_cur[t]
- if j == q_at and ntiles > 0:
- # The next item's Q, once this item's epilogue no longer needs the buffer.
- cute.arch.mbarrier_wait(qs_empty, qs_phase)
- qs_phase = qs_phase ^ 1
- self.issue_q(lt, ngroup, nhead, rows, q, q_stride_row, sQ)
- cute.arch.cp_async_mbarrier_arrive_noinc(qs_full)
- q_ready = Int32(1)
- # V(tiles - 1)
- cute.arch.mbarrier_wait(kv_empty + ring, ring_phase)
- self.issue_tile(atom, v, v_prev, sKV_c, ring, chunk, first_token)
- cute.arch.cp_async_mbarrier_arrive_noinc(kv_full + ring)
- ring, ring_phase = self.advance(ring, ring_phase)
- self.release_item(n, sch_empty)
- n = n + 1
- item = nitem
- # Tail: every slot released, so no load is in flight when the CTA retires.
- for i in cutlass.range_constexpr(self.kv_slots):
- cute.arch.mbarrier_wait(kv_empty + ring, ring_phase)
- ring, ring_phase = self.advance(ring, ring_phase)
-
- @cute.jit
- def issue_q(
- self,
- lt: Int32,
- group: Int32,
- kv_head: Int32,
- rows: Int32,
- q: cute.Pointer,
- q_stride_row: Int32,
- sQ: cute.Pointer,
- ):
- """An item's live Q rows into the Q buffer: MMA row m = local row x G + head at m x D,
- its 16-byte chunk c at chunk c ^ (m % 8). Consecutive threads take consecutive chunks of
- a row's G heads (contiguous in q).
- """
- G = self.group_size
- R = self.rows_per_tile
- D = self.d
- atom = cute.make_copy_atom(
- cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), BF16, num_bits_per_copy=128
- )
- per_row = G * (D // 8)
- total = R * per_row
- for it in cutlass.range_constexpr((total + 127) // 128):
- x = lt + 128 * it
- local_row = x // per_row
- rem = x - local_row * per_row
- head = rem // (D // 8)
- c = rem - head * (D // 8)
- row = group * R + local_row
- if x < total and row < rows:
- m = local_row * G + head
- src = cute.make_tensor(
- q
- + cute.assume(
- Int64(row) * q_stride_row + (kv_head * G + head) * D + c * 8, divby=8
- ),
- cute.make_layout((8,)),
- )
- dst = cute.make_tensor(
- sQ + cute.assume(self.box_offset(m, c), divby=8), cute.make_layout((8,))
- )
- cute.copy(atom, src, dst)
-
- @cute.jit
- def box_offset(self, m, c):
- """Element offset of row m's 16-byte chunk c (of D / 8) in the Q buffer: the layout of the
- epilogue's TMA boxes -- box c // 8 (64 dims of G heads x R rows), its 128-byte row m,
- chunk (c % 8) ^ (m % 8) (SW128) -- so a softmax thread's Q row and its O row occupy the
- same bytes.
- """
- return (
- (c // 8) * (64 * self.group_size * self.rows_per_tile)
- + m * 64
- + (((c % 8) ^ (m % 8)) * 8)
- )
-
- @cute.jit
- def issue_tile(
- self,
- atom,
- src_base: cute.Pointer,
- bases,
- sKV_c,
- ring: Int32,
- chunk: Int32,
- first_token: Int32,
- ):
- for t in cutlass.range_constexpr(4):
- n = first_token + 4 * t
- for b in cutlass.range_constexpr(self.d // 64):
- offset = cute.assume(bases[t] + b * 64 + chunk * 8, divby=8)
- src = cute.make_tensor(src_base + offset, cute.make_layout((8,)))
- dst = sKV_c[(None, (chunk, b), n, ring)]
- cute.copy(atom, src, dst)
-
- # ------------------------------------------------------------------ MMA (warp 4)
- @cute.jit
- def mma_role(
- self,
- qk_mma: cute.TiledMma,
- qk_mma_acc: cute.TiledMma,
- pv_mma: cute.TiledMma,
- pv_mma_acc: cute.TiledMma,
- qk_thr,
- pv_thr,
- q_layout,
- p_layout,
- sK,
- sV,
- tmem_ptr,
- plan: cute.Pointer,
- num_kv_heads: Int32,
- q_full,
- kv_full,
- kv_empty,
- s_full,
- p_full,
- o_full,
- o_final,
- sch_full,
- sch_empty,
- sch_val,
- ):
- """Per item: QK(0), then QK(j) / PV(j - 1) pairs, then the last PV. S/P buffer b sits at
- TMEM columns tmem_s0 + 64 b, b = the CTA's tile count mod 2; QK(t) overwrites S[t % 2] only
- after PV(t - 2) was issued (tcgen05 MMAs of one thread execute in order), and P(t - 2) was
- its last reader. An item's PV(0) overwrites O: it waits on P(0), which the softmax warps
- publish only after reading the previous item's O out.
- """
- tQ = cute.make_tensor(tmem_ptr + self.tmem_q, q_layout.outer)
- tSrQ = qk_thr.make_fragment_A(tQ)[None, None, None, 0]
- tSrQ = cute.make_tensor(cute.recast_ptr(tmem_ptr + self.tmem_q, dtype=BF16), tSrQ.layout)
- tSrK = qk_thr.make_fragment_B(sK)
- tOrV = pv_thr.make_fragment_B(sV)
- tStS = qk_thr.make_fragment_C(qk_thr.partition_shape_C((self.bm, self.bn)))
- tOtO_f = pv_thr.make_fragment_C(pv_thr.partition_shape_C((self.bm, self.d)))
- tOtO = cute.make_tensor(tmem_ptr + self.tmem_o, tOtO_f.layout)
- tP = cute.make_tensor(tmem_ptr + self.tmem_s0, p_layout.outer)
- tOrP = pv_thr.make_fragment_A(tP)[None, None, None, 0]
- n = Int32(0)
- g = Int32(0)
- nq = Int32(0)
- ring = Int32(0)
- ring_phase = Int32(0)
- item = self.read_item(n, sch_full, sch_val)
- while item >= 0:
- group, kv_head, tiles, listed = self.item_info(item, plan, num_kv_heads)
- if tiles > 0:
- cute.arch.mbarrier_wait(q_full, nq % 2)
- tcgen05_fence_after()
- # QK(0)
- buf0 = g % 2
- cute.arch.mbarrier_wait(kv_full + ring, ring_phase)
- self.gemm_qk(
- qk_mma,
- qk_mma_acc,
- tSrQ,
- tSrK[None, None, None, ring],
- cute.make_tensor(tmem_ptr + (self.tmem_s0 + buf0 * self.bn), tStS.layout),
- )
- with cute.arch.elect_one():
- tcgen05.commit(s_full + buf0)
- tcgen05.commit(kv_empty + ring)
- ring, ring_phase = self.advance(ring, ring_phase)
- for j in cutlass.range(1, tiles + 1, 1, unroll=1):
- t = g + j
- if j < tiles:
- # QK(j) into S[t % 2]
- buf = t % 2
- cute.arch.mbarrier_wait(kv_full + ring, ring_phase)
- self.gemm_qk(
- qk_mma,
- qk_mma_acc,
- tSrQ,
- tSrK[None, None, None, ring],
- cute.make_tensor(
- tmem_ptr + (self.tmem_s0 + buf * self.bn), tStS.layout
- ),
- )
- with cute.arch.elect_one():
- tcgen05.commit(s_full + buf)
- tcgen05.commit(kv_empty + ring)
- ring, ring_phase = self.advance(ring, ring_phase)
- # PV(j - 1)
- prev = t - 1
- pbuf = prev % 2
- cute.arch.mbarrier_wait(kv_full + ring, ring_phase)
- cute.arch.mbarrier_wait(p_full + pbuf, (prev // 2) % 2)
- tcgen05_fence_after()
- tOrP_b = cute.make_tensor(
- cute.recast_ptr(
- tmem_ptr + (self.tmem_s0 + pbuf * self.bn), dtype=tOrP.element_type
- ),
- tOrP.layout,
- )
- if j == 1:
- self.gemm_pv(pv_mma, pv_mma_acc, tOrP_b, tOrV[None, None, None, ring], tOtO)
- else:
- self.gemm_pv(
- pv_mma_acc, pv_mma_acc, tOrP_b, tOrV[None, None, None, ring], tOtO
- )
- with cute.arch.elect_one():
- tcgen05.commit(kv_empty + ring)
- tcgen05.commit(o_full)
- ring, ring_phase = self.advance(ring, ring_phase)
- with cute.arch.elect_one():
- tcgen05.commit(o_final)
- g = g + tiles
- nq = nq + 1
- self.release_item(n, sch_empty)
- n = n + 1
- item = self.read_item(n, sch_full, sch_val)
-
- @cute.jit
- def advance(self, ring: Int32, ring_phase: Int32):
- ring = ring + 1
- if ring == self.kv_slots:
- ring = Int32(0)
- ring_phase = ring_phase ^ 1
- return ring, ring_phase
-
- @cute.jit
- def gemm_qk(self, first: cute.TiledMma, rest: cute.TiledMma, tSrQ, tSrK, tStS):
- for kb in cutlass.range_constexpr(cute.size(tSrQ, mode=[2])):
- if cutlass.const_expr(kb == 0):
- cute.gemm(first, tStS, tSrQ[None, None, kb], tSrK[None, None, kb], tStS)
- else:
- cute.gemm(rest, tStS, tSrQ[None, None, kb], tSrK[None, None, kb], tStS)
-
- @cute.jit
- def gemm_pv(self, first: cute.TiledMma, rest: cute.TiledMma, tOrP, tOrV, tOtO):
- for kb in cutlass.range_constexpr(cute.size(tOrP, mode=[2])):
- if cutlass.const_expr(kb == 0):
- cute.gemm(first, tOtO, tOrP[None, None, kb], tOrV[None, None, kb], tOtO)
- else:
- cute.gemm(rest, tOtO, tOrP[None, None, kb], tOrV[None, None, kb], tOtO)
-
- # ------------------------------------------------------------------ softmax + epilogue (warps 0-3)
- @cute.jit
- def copy_q(self, m: Int32, live, sQ: cute.Pointer, store_p, tPtP0, tPcP, p_cols: int, q_full):
- """Row m's Q from the Q buffer into TMEM columns tmem_q.. (the QK MMA's A operand), 64
- elements a store; dead rows store zeros (their S rows are masked out).
- """
- D = self.d
- for c in cutlass.range_constexpr(D // (2 * p_cols)):
- rQ = cute.make_rmem_tensor(tPcP.shape, Float32)
- rQ16 = cute.make_tensor(
- cute.recast_ptr(rQ.iterator, dtype=BF16), cute.make_layout((2 * p_cols,))
- )
- if live:
- for i in cutlass.range_constexpr(2 * p_cols // 8):
- cidx = c * (2 * p_cols // 8) + i
- off = cute.assume(self.box_offset(m, cidx), divby=8)
- cute.autovec_copy(
- cute.make_tensor(sQ + off, cute.make_layout((8,))),
- cute.make_tensor(rQ16.iterator + i * 8, cute.make_layout((8,))),
- )
- else:
- rQ.fill(0.0)
- cute.copy(
- store_p,
- rQ,
- cute.make_tensor(
- tPtP0.iterator + (self.tmem_q - self.tmem_s0 + c * p_cols), tPtP0.layout
- ),
- )
- cute.arch.fence_view_async_tmem_store()
- tcgen05_fence_before()
- cute.arch.mbarrier_arrive(q_full)
-
- @cute.jit
- def softmax_role(
- self,
- qk_thr,
- pv_thr,
- tmem_ptr,
- st: Int32,
- plan: cute.Pointer,
- num_kv_heads: Int32,
- rows: Int32,
- bits: cute.Pointer,
- bits_stride_row: Int32,
- out: cute.Pointer,
- out_stride_row: Int32,
- scale_log2: Float32,
- q_full,
- s_full,
- p_full,
- o_full,
- o_final,
- qs_full,
- qs_empty,
- sch_full,
- sch_empty,
- sch_val,
- sQ: cute.Pointer,
- sO,
- tma_atom_o,
- tma_tensor_o,
- ):
- G = self.group_size
- R = self.rows_per_tile
- D = self.d
- tStS = qk_thr.make_fragment_C(qk_thr.partition_shape_C((self.bm, self.bn)))
- tStS0 = cute.make_tensor(tmem_ptr + self.tmem_s0, tStS.layout)
- tOtO_f = pv_thr.make_fragment_C(pv_thr.partition_shape_C((self.bm, self.d)))
- tOtO = cute.make_tensor(tmem_ptr + self.tmem_o, tOtO_f.layout)
- # The softmax warps' own barrier (hand-backs of the Q buffer, the TMA store's staging).
- softmax_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=128)
-
- # S: one 64-column load a tile; P: 32 packed-bf16 columns stored over the S it came from
- # (or two 16-column halves).
- cS = cute.make_identity_tensor((self.bm, self.bn))
- tScS = qk_thr.partition_C(cS)
- # The load also reduces the row max over the tile's 64 columns (tcgen05.ld.red, SM103): an
- # upper bound of the row's masked max, which serves as the softmax offset just as well.
- load_s = tcgen05.make_tmem_copy(
- cute.make_copy_atom(
- tcgen05.copy.LdRed32x32bOp(
- tcgen05.copy.Repetition(32), redOp=tcgen05.TmemLoadRedOp.MAX
- ),
- Float32,
- ),
- tStS0,
- )
- thr_load_s = load_s.get_slice(st)
- tLtS0 = thr_load_s.partition_S(tStS0)
- tLcS = thr_load_s.partition_D(tScS)
- p_cols = self.bn // 2
- tStP_layout = cute.composition(tStS0.layout, cute.make_layout((self.bm, p_cols)))
- tStP0 = cute.make_tensor(tStS0.iterator, tStP_layout)
- tScP = cute.make_tensor(
- tScS.iterator, cute.composition(tScS.layout, cute.make_layout((self.bm, p_cols)))
- )
- store_p = tcgen05.make_tmem_copy(
- cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(p_cols)), Float32),
- tStP0,
- )
- thr_store_p = store_p.get_slice(st)
- tPtP0 = thr_store_p.partition_D(tStP0)
- tPcP = thr_store_p.partition_S(tScP)
- # O in 32-column chunks (rescale and epilogue).
- cO = cute.make_identity_tensor((self.bm, self.d))
- tOcO = pv_thr.partition_C(cO)
- tOtO_c = cute.make_tensor(
- tOtO.iterator,
- cute.composition(tOtO.layout, cute.make_layout((self.bm, self.chunk_cols))),
- )
- tOcO_c = cute.make_tensor(
- tOcO.iterator,
- cute.composition(tOcO.layout, cute.make_layout((self.bm, self.chunk_cols))),
- )
- load_o = tcgen05.make_tmem_copy(
- cute.make_copy_atom(
- tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(self.chunk_cols)), Float32
- ),
- tOtO_c,
- )
- store_o = tcgen05.make_tmem_copy(
- cute.make_copy_atom(
- tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(self.chunk_cols)), Float32
- ),
- tOtO_c,
- )
- thr_load_o = load_o.get_slice(st)
- thr_store_o = store_o.get_slice(st)
- tLtO = thr_load_o.partition_S(tOtO_c)
- tLcO = thr_load_o.partition_D(tOcO_c)
- tStO = thr_store_o.partition_D(tOtO_c)
- n_chunks = self.d // self.chunk_cols
-
- # The TMA store's view: per item, box b of (group, KV head).
- gO = cute.flat_divide(tma_tensor_o, (64, G, R))
- tOsO, tOgO = cpasync.tma_partition(
- tma_atom_o,
- 0,
- cute.make_layout(1),
- cute.group_modes(sO, 0, 3),
- cute.group_modes(gO, 0, 3),
- )
-
- m = st
- local_row = m // G
- head = m % G
- tau = Float32(self.tau)
- n = Int32(0)
- g = Int32(0)
- nq = Int32(0)
- qs_phase = Int32(0)
- q_ready = Int32(0)
- store_pending = Int32(0)
- item = self.read_item(n, sch_full, sch_val)
- while item >= 0:
- group, kv_head, tiles, listed = self.item_info(item, plan, num_kv_heads)
- row = group * R + local_row
- live = (local_row < R) and (row < rows)
- if tiles > 0:
- if q_ready == 0:
- cute.arch.mbarrier_wait(qs_full, qs_phase)
- qs_phase = qs_phase ^ 1
- self.copy_q(m, live, sQ, store_p, tPtP0, tPcP, p_cols, q_full)
- # Every row read out: hand the buffer back.
- if st < 32:
- softmax_bar.arrive_and_wait()
- if st == 0:
- cute.arch.mbarrier_arrive(qs_empty)
- else:
- softmax_bar.arrive()
- q_ready = Int32(0)
- safe_row = group * R
- if live:
- safe_row = row
- mBits = cute.make_tensor(
- bits + Int64(safe_row) * bits_stride_row, cute.make_layout((1 << 30,))
- )
- row_max = -Float32.inf
- row_sum = Float32(0.0)
- lo_next = Int32(0)
- hi_next = Int32(0)
- if live:
- lo_next = mBits[0]
- hi_next = mBits[1]
- for j in cutlass.range(0, tiles, 1, unroll=1):
- t = g + j
- lo = lo_next
- hi = hi_next
- if live and j + 1 < tiles:
- lo_next = mBits[2 * j + 2]
- hi_next = mBits[2 * j + 3]
- buf = t % 2
- cute.arch.mbarrier_wait(s_full + buf, (t // 2) % 2)
- tcgen05_fence_after()
- rS = cute.make_rmem_tensor(tLcS.shape, Float32)
- n_ld = cute.size(rS, mode=[1])
- rMax = cute.make_rmem_tensor(cute.make_layout((1, n_ld)), Float32)
- tLtS_b = cute.make_tensor(tLtS0.iterator + buf * self.bn, tLtS0.layout)
- for c in cutlass.range_constexpr(n_ld):
- cute.copy_atom_call(
- load_s, tLtS_b[None, c, 0, 0], (rS[None, c, 0, 0], rMax[None, c])
- )
- cute.arch.fence_view_async_tmem_load()
- tile_max = rMax[0]
- for c in cutlass.range_constexpr(1, n_ld):
- tile_max = cute.arch.fmax(tile_max, rMax[c])
- new_max = cute.arch.fmax(row_max, tile_max)
- use_max = row_max
- factor = Float32(1.0)
- if row_max == -Float32.inf:
- use_max = new_max
- elif (new_max - row_max) * scale_log2 > tau:
- use_max = new_max
- factor = cute.math.exp2((row_max - new_max) * scale_log2, fastmath=True)
- need = factor != Float32(1.0)
- if cute.arch.vote_any_sync(need):
- # O holds PV(0..j-1): wait for PV(j-1), then rescale this warp's rows.
- cute.arch.mbarrier_wait(o_full, (t - 1) % 2)
- tcgen05_fence_after()
- for c in cutlass.range_constexpr(n_chunks):
- rO = cute.make_rmem_tensor(tLcO.shape, Float32)
- cute.copy(
- load_o,
- cute.make_tensor(tLtO.iterator + c * self.chunk_cols, tLtO.layout),
- rO,
- )
- for i in cutlass.range_constexpr(cute.size(rO)):
- rO[i] = rO[i] * factor
- cute.copy(
- store_o,
- rO,
- cute.make_tensor(tStO.iterator + c * self.chunk_cols, tStO.layout),
- )
- cute.arch.fence_view_async_tmem_store()
- row_max = use_max
- neg = -use_max * scale_log2
- rP = cute.make_rmem_tensor(tPcP.shape, Float32)
- rP16 = cute.make_tensor(cute.recast_ptr(rP.iterator, dtype=BF16), rS.layout)
- # p = 2^(s * scale - max * scale) on the row's own entries, 0 elsewhere; four
- # partial sums.
- acc0 = (Float32(0.0), Float32(0.0))
- acc1 = (Float32(0.0), Float32(0.0))
- for i in cutlass.range_constexpr(0, self.bn, 4):
- word = lo if i < 32 else hi
- e = cute.make_rmem_tensor((4,), Float32)
- e[0], e[1] = cute.arch.fma_packed_f32x2(
- (rS[i], rS[i + 1]), (scale_log2, scale_log2), (neg, neg)
- )
- e[2], e[3] = cute.arch.fma_packed_f32x2(
- (rS[i + 2], rS[i + 3]), (scale_log2, scale_log2), (neg, neg)
- )
- for u in cutlass.range_constexpr(4):
- x = cute.math.exp2(e[u], fastmath=True)
- rS[i + u] = x if (word & (1 << ((i + u) % 32))) != 0 else Float32(0.0)
- acc0 = cute.arch.add_packed_f32x2(acc0, (rS[i], rS[i + 1]))
- acc1 = cute.arch.add_packed_f32x2(acc1, (rS[i + 2], rS[i + 3]))
- row_sum = row_sum * factor + ((acc0[0] + acc0[1]) + (acc1[0] + acc1[1]))
- rP16.store(rS.load().to(BF16))
- cute.copy(
- store_p, rP, cute.make_tensor(tPtP0.iterator + buf * self.bn, tPtP0.layout)
- )
- cute.arch.fence_view_async_tmem_store()
- tcgen05_fence_before()
- cute.arch.mbarrier_arrive(p_full + buf)
- if j == 0 and store_pending != 0:
- # The previous item's TMA store has read the Q buffer: hand it back.
- if st == 0:
- cute.arch.cp_async_bulk_wait_group(0, read=True)
- cute.arch.mbarrier_arrive(qs_empty)
- store_pending = Int32(0)
-
- # The last QK of this item has completed (its S arrived): the next item's Q goes
- # into tensor memory now, so its first QK MMAs overlap this epilogue.
- nitem = self.read_item(n + 1, sch_full, sch_val)
- ngroup, nhead, ntiles, _ = self.item_info(nitem, plan, num_kv_heads)
- next_q = ntiles > 0
- if next_q:
- nrow = ngroup * R + local_row
- nlive = (local_row < R) and (nrow < rows)
- cute.arch.mbarrier_wait(qs_full, qs_phase)
- qs_phase = qs_phase ^ 1
- self.copy_q(m, nlive, sQ, store_p, tPtP0, tPcP, p_cols, q_full)
- q_ready = Int32(1)
- # The Q buffer is free for staging O unless a later item still has to load its Q
- # there (a following lone or empty group): then each thread stores its own row.
- staged = next_q or nitem < 0
-
- # Epilogue: O / l in bf16.
- cute.arch.mbarrier_wait(o_final, nq % 2)
- tcgen05_fence_after()
- inv = Float32(0.0)
- if row_sum > Float32(0.0):
- inv = Float32(1.0) / row_sum
- if staged:
- if nitem < 0:
- # No next Q orders this staging after the previous item's TMA store read
- # the buffer (thread 0 waited for it at this item's first tile): sync.
- softmax_bar.arrive_and_wait()
- for c in cutlass.range_constexpr(n_chunks):
- rO = cute.make_rmem_tensor(tLcO.shape, Float32)
- cute.copy(
- load_o,
- cute.make_tensor(tLtO.iterator + c * self.chunk_cols, tLtO.layout),
- rO,
- )
- for i in cutlass.range_constexpr(cute.size(rO)):
- rO[i] = rO[i] * inv
- rB = cute.make_rmem_tensor((self.chunk_cols,), BF16)
- rB.store(rO.load().to(BF16))
- if m < G * R:
- for i in cutlass.range_constexpr(self.chunk_cols // 8):
- offset = cute.assume(
- self.box_offset(m, c * (self.chunk_cols // 8) + i), divby=8
- )
- cute.autovec_copy(
- cute.make_tensor(rB.iterator + i * 8, cute.make_layout((8,))),
- cute.make_tensor(sQ + offset, cute.make_layout((8,))),
- )
- cute.arch.fence_view_async_shared()
- if st < 32:
- softmax_bar.arrive_and_wait()
- if st == 0:
- for b in cutlass.range_constexpr(self.d // 64):
- cute.copy(tma_atom_o, tOsO[None, b], tOgO[None, b, kv_head, group])
- cute.arch.cp_async_bulk_commit_group()
- else:
- softmax_bar.arrive()
- store_pending = Int32(1)
- if nitem < 0:
- if st == 0:
- cute.arch.cp_async_bulk_wait_group(0, read=True)
- store_pending = Int32(0)
- else:
- for c in cutlass.range_constexpr(n_chunks):
- rO = cute.make_rmem_tensor(tLcO.shape, Float32)
- cute.copy(
- load_o,
- cute.make_tensor(tLtO.iterator + c * self.chunk_cols, tLtO.layout),
- rO,
- )
- for i in cutlass.range_constexpr(cute.size(rO)):
- rO[i] = rO[i] * inv
- rB = cute.make_rmem_tensor((self.chunk_cols,), BF16)
- rB.store(rO.load().to(BF16))
- if live:
- for i in cutlass.range_constexpr(self.chunk_cols // 8):
- dst_off = cute.assume(
- Int64(row) * out_stride_row
- + (kv_head * G + head) * D
- + c * self.chunk_cols
- + i * 8,
- divby=8,
- )
- cute.autovec_copy(
- cute.make_tensor(rB.iterator + i * 8, cute.make_layout((8,))),
- cute.make_tensor(out + dst_off, cute.make_layout((8,))),
- )
- g = g + tiles
- nq = nq + 1
- elif tiles == 0:
- # Nothing resident for this group: zeros.
- if live:
- rZ = cute.make_rmem_tensor((8,), BF16)
- rZ.fill(0.0)
- for i in cutlass.range_constexpr(D // 8):
- dst_off = cute.assume(
- Int64(row) * out_stride_row + (kv_head * G + head) * D + i * 8,
- divby=8,
- )
- cute.autovec_copy(
- rZ, cute.make_tensor(out + dst_off, cute.make_layout((8,)))
- )
- self.release_item(n, sch_empty)
- n = n + 1
- item = self.read_item(n, sch_full, sch_val)
-
-
-@functools.lru_cache(maxsize=None)
-def _compiled(group_size: int, rows_per_tile: int, tokens_per_block: int, use_pdl: bool):
- kernel = _RowtileAttentionKernel(
- group_size, rows_per_tile, HEAD_DIM, tokens_per_block, kv_slots=KV_SLOTS, use_pdl=use_pdl
- )
- gmem = cute.AddressSpace.gmem
-
- def fake(dtype, align):
- return make_ptr(dtype, 16 * 1024, gmem, assumed_align=align)
-
- args = (
- fake(BF16, 16),
- fake(BF16, 16),
- fake(BF16, 16),
- fake(Int32, 4),
- fake(Int32, 4),
- fake(Int32, 4),
- fake(BF16, 16),
- fake(Int32, 4),
- 1,
- 1,
- 1,
- 1,
- 8,
- 8,
- 8,
- 8,
- 8,
- 8,
- 8,
- 1,
- 1,
- 8,
- 1.0,
- cuda_drv.CUstream(0),
- )
- return cute.compile(kernel, *args)
-
-
-@functools.lru_cache(maxsize=None)
-def _sm_count(device_index: int) -> int:
- return torch.cuda.get_device_properties(device_index).multi_processor_count
-
-
-def num_ctas(num_items: int, device: torch.device) -> int:
- """The persistent grid: one CTA per SM, never more than there are items."""
- return max(1, min(num_items, _sm_count(device.index if device.index is not None else 0)))
-
-
-def _aligned(t: torch.Tensor, strides: int) -> bool:
- return t.data_ptr() % 16 == 0 and all(t.stride(i) % 8 == 0 for i in range(strides))
-
-
-def eligible(q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor) -> bool:
- """Whether a launch fits this pass: SM103, bf16 q and K/V, D = 256, unit-stride head
- dimension and 16-byte-aligned rows.
- """
- if q.shape[2] != HEAD_DIM:
- return False
- if not (q.dtype == k_cache.dtype == v_cache.dtype == torch.bfloat16):
- return False
- if torch.cuda.get_device_capability(q.device) != (10, 3):
- return False
- if (
- q.stride(2) != 1
- or q.stride(1) != HEAD_DIM
- or k_cache.stride(3) != 1
- or v_cache.stride(3) != 1
- ):
- return False
- return _aligned(q, 1) and _aligned(k_cache, 3) and _aligned(v_cache, 3)
-
-
-class _Launcher:
- """`rowtile_attention_cute[grid](...)`: grid-indexed like the entry's Triton launches;
- the grid is `(ctas,)`, the persistent CTAs (`num_ctas`).
- """
-
- def __getitem__(self, grid):
- return functools.partial(self._launch, tuple(grid))
-
- @staticmethod
- def _launch(
- grid,
- q,
- k_cache,
- v_cache,
- slots,
- packed_bits,
- group_plan,
- output,
- work,
- softmax_scale,
- rows,
- *,
- group_size,
- rows_per_tile,
- tokens_per_block,
- use_pdl,
- ):
- assert len(grid) == 1, grid
- assert (
- output.is_contiguous()
- and packed_bits.dtype == torch.int64
- and slots.dtype == torch.int32
- and work.dtype == torch.int32
- )
- fn = _compiled(group_size, rows_per_tile, tokens_per_block, bool(use_pdl))
- gmem = cute.AddressSpace.gmem
-
- def ptr(t, dtype, align):
- return make_ptr(dtype, t.data_ptr(), gmem, assumed_align=align)
-
- fn(
- ptr(q, BF16, 16),
- ptr(k_cache, BF16, 16),
- ptr(v_cache, BF16, 16),
- ptr(slots, Int32, 4),
- ptr(packed_bits, Int32, 4),
- ptr(group_plan, Int32, 4),
- ptr(output, BF16, 16),
- ptr(work, Int32, 4),
- rows,
- group_plan.shape[0],
- k_cache.shape[1],
- grid[0],
- q.stride(0),
- k_cache.stride(0),
- k_cache.stride(1),
- k_cache.stride(2),
- v_cache.stride(0),
- v_cache.stride(1),
- v_cache.stride(2),
- slots.stride(0),
- 2 * packed_bits.stride(0),
- output.stride(0),
- softmax_scale * _LOG2_E,
- cuda_drv.CUstream(torch.cuda.current_stream(q.device).cuda_stream),
- )
-
-
-rowtile_attention_cute = _Launcher()
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_sparse_gqa_rowtile__prefill.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_sparse_gqa_rowtile__prefill.py
index 2d7cb622703f..4a37ca42dfd4 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_sparse_gqa_rowtile__prefill.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_sparse_gqa_rowtile__prefill.py
@@ -35,14 +35,22 @@
from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa import (
qsa_paged_sparse_gqa as _per_row_entry,
)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_rowtile_attention_cute import ( # noqa: E501
+ BLOCK_M,
+ BLOCK_N,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_rowtile_attention_cute import ( # noqa: E501
+ eligible as _cute_eligible,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_rowtile_attention_cute import ( # noqa: E501
+ num_ctas as _rowtile_ctas,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_rowtile_attention_cute import ( # noqa: E501
+ rowtile_attention_cute as _rowtile_attention_cute,
+)
from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
-from .qsa_rowtile_attention_cute__prefill import BLOCK_M, BLOCK_N
-from .qsa_rowtile_attention_cute__prefill import eligible as _cute_eligible
-from .qsa_rowtile_attention_cute__prefill import num_ctas as _rowtile_ctas
-from .qsa_rowtile_attention_cute__prefill import rowtile_attention_cute as _rowtile_attention_cute
-
_INTEGER_DTYPES = (torch.int32, torch.int64)
# The per-row kernel's scan tile: the prefill bound is rounded up to it, so the
# bitmap reads exactly the columns that kernel reads.
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/vision_rope_2d_packed_qk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/vision_rope_2d_packed_qk.py
index 6b8bdc160131..7a7b746d84a0 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/vision_rope_2d_packed_qk.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/vision_rope_2d_packed_qk.py
@@ -31,11 +31,11 @@
# `flash_attn/ops/triton/rotary.py`: `rotary_kernel` and `apply_rotary`.
"""Vision-tower 2-D rotary, in place on the q|k prefix of the packed qkv projection.
-The rotation half of the vision tower's 2-D axial rotary (`vision_rope_2d__r52` is the
-whole op, angle table included). The angle table depends only on the request's
-`grid_thw`, so the caller builds it once per tower forward -- spatial-merge `(row, col)`
-per token, fp64 angles rounded once to fp32, then `cos` / `sin` -- and every one of the 27
-blocks hands it here with that block's qkv projection output.
+The rotation half of the vision tower's 2-D axial rotary; the tower builds the angle
+table. The angle table depends only on the request's `grid_thw`, so the caller builds it
+once per tower forward -- spatial-merge `(row, col)` per token, fp64 angles rounded once
+to fp32, then `cos` / `sin` -- and every one of the 27 blocks hands it here with that
+block's qkv projection output.
Tier 5: one Triton launch, flash-attn 2's `rotary_kernel` (BSD-3-Clause, header above)
cut down to the path this call takes. flash-attn 2.x is not a pinned dependency
@@ -186,8 +186,8 @@ class VisionRope2dPackedQk(OpWrapper):
itself, and the caller's qkv buffer holds the rotated q|k on return.
Numerics: bitwise equal to `reference` (a complex multiply in fp32, rounded once) and
- to `vision_rope_2d__r52.reference` (positions by nested loops, fp64 angles, `polar`)
- at every cell, q and k alike: relative L2 0, cosine 1.
+ to the test's whole-op reference from the grid (positions by nested loops, fp64
+ angles, `polar`) at every cell, q and k alike: relative L2 0, cosine 1.
Silent, refused by `is_valid` (each demonstrated unguarded in the test):
- a bf16 (or any non-fp32) table: the kernel widens it and rotates by rounded
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/ssm/gdn_cached_replay_cute__verify.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/ssm/gdn_cached_replay_cute__verify.py
index 9c4ceeef43f3..fa0c4738b1bc 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/ssm/gdn_cached_replay_cute__verify.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/ssm/gdn_cached_replay_cute__verify.py
@@ -221,21 +221,6 @@ def _ldg_i32(gaddr, *, loc=None, ip=None):
)
-@dsl_user_op
-def _ldg_f32(gaddr, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(
- _ty.f32(),
- [_ir(gaddr, cutlass.Int64, loc, ip)],
- "ld.global.f32 $0, [$1];",
- "=f,l",
- loc,
- ip,
- side=True,
- )
- )
-
-
@dsl_user_op
def _ldg_f32p(gaddr, pred, *, loc=None, ip=None):
"""An f32, or 0 where `pred` is 0 (nothing read)."""
@@ -285,52 +270,6 @@ def _b16_f32(x, *, loc=None, ip=None):
)
-@dsl_user_op
-def _f32_bits(x, *, loc=None, ip=None):
- return cutlass.Int32(
- _asm(_ty.i32(), [_ir(x, cutlass.Float32, loc, ip)], "mov.b32 $0, $1;", "=r,f", loc, ip)
- )
-
-
-@dsl_user_op
-def _bits_f32(x, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(_ty.f32(), [_ir(x, cutlass.Int32, loc, ip)], "mov.b32 $0, $1;", "=f,r", loc, ip)
- )
-
-
-@dsl_user_op
-def _ldg_bf16p(gaddr, pred, *, loc=None, ip=None):
- """A bf16 widened to f32, or 0 where `pred` is 0 (nothing read)."""
- return cutlass.Float32(
- _asm(
- _ty.f32(),
- [_ir(gaddr, cutlass.Int64, loc, ip), _ir(pred, cutlass.Int32, loc, ip)],
- "{ .reg .pred q; .reg .b16 h; setp.ne.b32 q, $2, 0; mov.b16 h, 0; @q ld.global.b16 h, [$1]; "
- "cvt.f32.bf16 $0, h; }",
- "=f,l,r",
- loc,
- ip,
- side=True,
- )
- )
-
-
-@dsl_user_op
-def _ldg_bf16(gaddr, *, loc=None, ip=None):
- return cutlass.Float32(
- _asm(
- _ty.f32(),
- [_ir(gaddr, cutlass.Int64, loc, ip)],
- "{ .reg .b16 h; ld.global.b16 h, [$1]; cvt.f32.bf16 $0, h; }",
- "=f,l",
- loc,
- ip,
- side=True,
- )
- )
-
-
@dsl_user_op
def _cp16z(saddr, gaddr, nbytes, *, loc=None, ip=None):
"""cp.async of 16 bytes, zero-filled past `nbytes` (0 or 16)."""
@@ -740,31 +679,6 @@ def _mbar_wait(mbar, phase, *, loc=None, ip=None):
)
-@dsl_user_op
-def _mbar_arrive(mbar, *, loc=None, ip=None):
- _asm(
- _ty.i32(),
- [_ir(mbar, cutlass.Int32, loc, ip)],
- "{ mbarrier.arrive.shared::cta.b64 _, [$1]; mov.u32 $0, 0; }",
- "=r,r",
- loc,
- ip,
- side=True,
- )
-
-
-@dsl_user_op
-def _fence_proxy_async(*, loc=None, ip=None):
- _asm(
- _ty.i32(), [], "{ fence.proxy.async.shared::cta; mov.u32 $0, 0; }", "=r", loc, ip, side=True
- )
-
-
-@dsl_user_op
-def _cp_wait1(*, loc=None, ip=None):
- _asm(_ty.i32(), [], "{ cp.async.wait_group 1; mov.u32 $0, 0; }", "=r", loc, ip, side=True)
-
-
def _exp(x):
return _ex2(x * LOG2E)
@@ -1946,6 +1860,8 @@ def cached_replay_update(
#: A tp4/tep4 rank's GDN heads: 16 key / 48 value split four ways.
RANK_K_HEADS, RANK_V_HEADS = 4, 12
+#: A tp1 or dep4 rank's GDN heads: the whole layer's.
+ONE_RANK_K_HEADS, ONE_RANK_V_HEADS = 16, 48
#: The stock tep4 engine's recurrent-state slots a rank (bf16 state), and the GDN layers whose
#: states share one slot page in the stock manager: the slot stride is 36 layers' states.
TEP4_SLOTS, PAGE_LAYERS = 1874, 36
@@ -1970,6 +1886,23 @@ def _tep4_cell(why, path, n, t, **spec) -> Cell:
)
+def _whole_layer_cell(why, path, n, t, **spec) -> Cell:
+ """The whole layer's heads (tp1, dep4) over a bf16 state pool."""
+ return Cell(
+ why=f"16/48 {why}",
+ spec=dict(
+ path=path,
+ n=n,
+ t=t,
+ state_dtype=torch.bfloat16,
+ H=ONE_RANK_K_HEADS,
+ HV=ONE_RANK_V_HEADS,
+ history_size=HIST,
+ **spec,
+ ),
+ )
+
+
class GdnCachedReplayCute(OpWrapper):
"""`cached_replay_update` as a catalog entry: one verify step of `N` sequences x `T` tokens
through one GDN layer under the cached-replay state layout, bf16 checkpoint state only, with
@@ -2038,12 +1971,47 @@ class GdnCachedReplayCute(OpWrapper):
strided_gates=True,
),
_tep4_cell("all_layer n=128 T=4: a large partitioned batch", "all_layer", 128, 4, steps=4),
+ _whole_layer_cell(
+ "small n=1 T=3: MTP2 verify at con1 (the tp1 AgentX engines), in-kernel commit",
+ "small",
+ 1,
+ 3,
+ steps=6,
+ strided_gates=True,
+ ),
+ _whole_layer_cell(
+ "small n=1 T=4: MTP3 verify at con1, in-kernel commit",
+ "small",
+ 1,
+ 4,
+ steps=6,
+ strided_gates=True,
+ ),
+ _whole_layer_cell(
+ "all_layer n=32 T=4: work items + the all-layer commit",
+ "all_layer",
+ 32,
+ 4,
+ steps=4,
+ slots=72,
+ strided_gates=True,
+ ),
+ _whole_layer_cell(
+ "all_layer n=256 T=4: a large partitioned batch",
+ "all_layer",
+ 256,
+ 4,
+ steps=3,
+ slots=300,
+ layers=1,
+ ),
)
note = """
Call site: `_gdn_verify` (`DecodeTarget._gdn_block` / `PrefillTarget._gdn_block` of a verify
- step) over the bf16 recurrent pool, at a tp4/tep4 rank's 4 / 12 heads: q/k/v from the conv's
- fresh contiguous `[N*T, 2560]` rows (`packed_qkv`), `g` / `beta` the input projection's raw
+ step) over the bf16 recurrent pool, at a tp4/tep4 rank's 4 / 12 heads (tp1 and dep4: the
+ whole layer's 16 / 48): q/k/v from the conv's fresh contiguous `[N*T, 2560]` rows
+ (`packed_qkv`; 10240 at 16 / 48), `g` / `beta` the input projection's raw
a / b columns read through its row stride (no copies), `A_log` / `dt_bias` fp32, the
layer's view of the stock manager's bf16 state page (slot stride 36 layers' states), the
history buffers `[slots, 2, 16, 12|4, 128]` / `[slots, 2, 12, 16]`, and work items +
@@ -2055,7 +2023,9 @@ class GdnCachedReplayCute(OpWrapper):
T, launch mapping, commit, PDL, ...) on first use (5-10 s): the warm-up before CUDA-graph
capture covers it.
- Certified at 4 / 12 heads over a bf16 checkpoint (cells: the stock tep4 engine's 36-layer
+ Certified at the whole layer's 16 / 48 heads (tp1, dep4) over a bf16 state: the in-kernel
+ commit at N = 1 and T = 3, 4, the all-layer commit at N = 32 and 256 (T = 4). Certified at
+ 4 / 12 heads over a bf16 checkpoint (cells: the stock tep4 engine's 36-layer
state page at 1874 slots, rows on both ends of it, several consecutive verify steps from
warm histories with random acceptance): the in-kernel commit at N = 1, 2, 4 and T = 2, 4,
8 (T = 8 with full acceptance appending to the history's end, then wrapping), CUDA-graph
@@ -2065,8 +2035,8 @@ class GdnCachedReplayCute(OpWrapper):
<= 0.71 of 1.5 ulps, bf16 checkpoint <= 0.68 of 2 ulps of its head's absmax, `old_u` <=
0.50 of 4 ulps, `old_G` <= 0.06 -- the Triton entry's own margins (0.52 / 0.64 / 0.84 /
0.50 / 0.08). A reference with the key heads reversed is refused by the output gate.
- CUDA-graph replay (PDL on and off, N = 1, 4 and 16) is bitwise eager, and the launch
- mappings are bitwise each other at N = 1, 2, 4 and T = 2, 4, 8.
+ CUDA-graph replay (PDL on and off, N = 1, 4 and 16, at both geometries) is bitwise eager,
+ and the 4 / 12 launch mappings are bitwise each other at N = 1, 2, 4 and T = 2, 4, 8.
Numerics: the Triton kernel's rounding points (bf16 L2-normalized q / k and history
coefficients into the tensor-core dots, fp32 recurrence, fp32 commit rounded once to the
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/pool.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/pool.py
index 4d88a74eaeb0..90d8c3259a2a 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/pool.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/pool.py
@@ -357,8 +357,8 @@ def check_replay_width(replay, rows: int) -> None:
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"
+ "draft-length schedule, or skipped drafting) cannot run under it: this target runs "
+ "one draft length, and routing declines a variable one"
)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/weights.py
index ffa1185ca5e0..864ff5efea9e 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/weights.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/weights.py
@@ -312,11 +312,6 @@ def _host_resident_empty(
#: stored (the gather dequantizes the rows it reads).
PLE_TABLE_DTYPE = torch.float8_e4m3fn
-#: How the trunk's routed experts are split (`rank_dims`'s `expert_split`) that this table
-#: declares and `load` fills: by expert (`ep`, a window of whole experts) or by each
-#: expert's intermediate (`tp`, every expert at a zero-padded slice).
-LOADED_EXPERT_SPLITS = ("ep", "tp")
-
#: The multiple a rank's stored routed-expert intermediate is padded up to. It is the
#: engine's own storage rule for an NVFP4 trtllm-gen shard
#: (`NVFP4TRTLLMGenFusedMoEMethod.resolve_alignments`: a padded intermediate that is not a
@@ -696,9 +691,9 @@ def _gdn_layers(d):
def _trunk_expert_layers(d):
- """Every layer when the routed experts are split in a way this table declares
- (`LOADED_EXPERT_SPLITS`); none otherwise."""
- return range(d.num_layers) if d.expert_split in LOADED_EXPERT_SPLITS else ()
+ """Every layer: each holds routed experts, split by expert (`ep`, a window of whole
+ experts) or by each expert's intermediate (`tp`, every expert at a zero-padded slice)."""
+ return range(d.num_layers)
def _cute_dsl_copy_layers(d):
@@ -1486,14 +1481,11 @@ def expected_unconsumed(self, d: SimpleNamespace, keys) -> set:
whole `mtp.` subtree. The windows are enumerated from `d`, not read off the
checkpoint, so a key that belongs to a real weight can never land here."""
out = set()
- if d.expert_split in LOADED_EXPERT_SPLITS:
- start, count = _trunk_window(d)
- for i in range(d.num_layers):
- for e in range(d.num_experts):
- if not start <= e < start + count:
- out.update(
- f"{_trunk_mlp(i)}.experts.{e}.{t}" for t in _TRUNK_EXPERT_TENSORS
- )
+ start, count = _trunk_window(d)
+ for i in range(d.num_layers):
+ for e in range(d.num_experts):
+ if not start <= e < start + count:
+ out.update(f"{_trunk_mlp(i)}.experts.{e}.{t}" for t in _TRUNK_EXPERT_TENSORS)
if not d.mtp_enabled:
return out | {k for k in keys if k.startswith(f"{_MTP}.")}
start, count = _mtp_window(d)
@@ -1567,10 +1559,6 @@ def load(self, model, weights) -> None:
tensors at once, and coverage is asserted in both directions."""
core = model.model
d = self.dims(core)
- assert d.expert_split in LOADED_EXPERT_SPLITS, (
- f"routed experts split by {d.expert_split!r}: this table declares the splits "
- f"{LOADED_EXPERT_SPLITS}"
- )
# The one load-time declaration, before the manifest is built: the table's
# height comes from the shards, and the manifest addresses it by row range.
rows_per_shard = _declare_ngram_table(core, d, weights)
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 4ad745b1d0e6..1dcdda18b0ab 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
@@ -16,7 +16,7 @@
from __future__ import annotations
-from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Tuple
+from typing import TYPE_CHECKING, NamedTuple, Optional, Tuple
import torch
@@ -28,6 +28,7 @@
if TYPE_CHECKING:
from tensorrt_llm._torch.model_config import ModelConfig
from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
+ from tensorrt_llm.mapping import Mapping
# The GPU the target is certified on.
_SM = (10, 3)
@@ -74,14 +75,9 @@
# The attention-DP layouts MTP is admitted at: those its draft acceptance was measured at.
_ATTENTION_DP_MTP_LAYOUTS = ("dep4",)
-# A prefill step scores and selects its context rows (at most max_num_tokens) in one call,
-# and `qsa_paged_index_scores` (output rows) and `qsa_topk_expand_decode` (row layouts) index
-# the score table in int32: rows x columns must stay below 2^31. The table has
-# `indexer_compress_ratio`-compressed columns over the KV manager's `max_blocks_per_seq`
-# pages: the sequence plus its generation lead (at most twice the widest verify step),
-# padded to a multiple of 4 pages.
-_SCORE_TABLE_ELEMENTS = 2**31
-_GENERATION_LEAD_TOKENS = 2 * (_MAX_DRAFT_LEN + 1)
+# The layouts with expert parallelism, where the engine builds an EPLB balancer for
+# `moe_config.load_balancer`; the target cuts fixed contiguous expert windows.
+_EXPERT_PARALLEL_LAYOUTS = tuple(member for key, member in _LAYOUTS.items() if key[3] > 1)
class EngineChoices(NamedTuple):
@@ -93,6 +89,7 @@ class EngineChoices(NamedTuple):
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
+ load_balancer: bool # `moe_config.load_balancer` is set (EPLB)
def engine_choices(model_config: ModelConfig, llm_args: TorchLlmArgs) -> EngineChoices:
@@ -109,10 +106,11 @@ def engine_choices(model_config: ModelConfig, llm_args: TorchLlmArgs) -> EngineC
),
gdn_replay=is_gdn_replay_enabled(),
prefill_cuda_graph_backend=str(llm_args.prefill_cuda_graph_backend.value),
+ load_balancer=llm_args.moe_config.load_balancer is not None,
)
-def layout(mapping) -> Optional[str]:
+def layout(mapping: Mapping) -> Optional[str]:
"""The target's layout ``mapping`` describes (``tp1``, ``tp4``, ``tep4``, ``dep4``,
``dep2``), or None.
@@ -139,7 +137,7 @@ def layout(mapping) -> Optional[str]:
)
-def _describe(mapping) -> str:
+def _describe(mapping: Mapping) -> str:
return (
f"world_size={mapping.world_size} (tp={mapping.tp_size}, pp={mapping.pp_size}, "
f"cp={mapping.cp_size}, attention_dp={mapping.enable_attention_dp}, "
@@ -148,8 +146,8 @@ def _describe(mapping) -> str:
def _check_engine(
- spec_config: Optional[Any],
- kv_cache_quant_algo: Optional[Any],
+ spec_config: Optional[MTPDecodingConfig],
+ kv_cache_quant_algo: Optional[QuantAlgo],
engine: EngineChoices,
member: str,
) -> None:
@@ -190,6 +188,12 @@ def _check_engine(
f"a {engine.mamba_ssm_cache_dtype} recurrent-state cache; at {member} the target "
f"serves {served} (kv_cache_config.mamba_ssm_cache_dtype)"
)
+ if engine.load_balancer and member in _EXPERT_PARALLEL_LAYOUTS:
+ raise NoTarget(
+ f"moe_config.load_balancer (EPLB) at {member}: the target runs fixed contiguous "
+ "expert windows and registers no layer with the balancer, so the placement would "
+ "do nothing"
+ )
if engine.prefill_cuda_graph_backend == "piecewise":
raise NoTarget(
"prefill_cuda_graph_backend='piecewise' makes the engine torch.compile the model, "
@@ -212,28 +216,9 @@ def _check_engine(
)
-def _check_score_table(
- max_num_tokens: int,
- max_seq_len: Optional[int],
- tokens_per_block: int,
- compress_ratio: Optional[int],
+def _check_attention_dp(
+ member: str, spec_config: Optional[MTPDecodingConfig], engine: EngineChoices
) -> None:
- """Decline a prefill step whose score table would pass 2^31 elements (see
- `_SCORE_TABLE_ELEMENTS`); an unknown max_seq_len or compression is not checked."""
- if max_seq_len is None or compress_ratio is None:
- return
- pages = -(-(max_seq_len + _GENERATION_LEAD_TOKENS) // tokens_per_block)
- columns = -(-pages // 4) * 4 * tokens_per_block // compress_ratio
- if max_num_tokens * columns >= _SCORE_TABLE_ELEMENTS:
- raise NoTarget(
- f"max_num_tokens={max_num_tokens} at max_seq_len={max_seq_len}: a prefill step's "
- f"context rows score a {columns}-column table, and the target indexes it in int32 "
- f"up to 2^31 elements; set max_num_tokens to at most "
- f"{(_SCORE_TABLE_ELEMENTS - 1) // columns}"
- )
-
-
-def _check_attention_dp(member: str, spec_config: Optional[Any], engine: EngineChoices) -> None:
"""Decline what the target does not run under attention DP (dep4, dep2), where
each rank serves its own requests and only the routed experts cross ranks.
@@ -304,8 +289,8 @@ def route(
# 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.
+ # module. The target's constructor relies on each of these, so declining here
+ # is what lets ``auto`` fall back to the built-in model.
mtp_eagle_or_none = spec_config is None or (
isinstance(spec_config, MTPDecodingConfig)
and spec_config.spec_dec_mode.is_mtp_eagle_one_model()
@@ -325,12 +310,6 @@ def route(
if mapping.enable_attention_dp:
_check_attention_dp(member, spec_config, engine)
_check_engine(spec_config, kv_cache_quant_algo, engine, member)
- _check_score_table(
- llm_args.max_num_tokens or model_config.max_num_tokens,
- llm_args.max_seq_len or getattr(text, "max_position_embeddings", None),
- llm_args.kv_cache_config.tokens_per_block,
- getattr(text, "indexer_compress_ratio", None),
- )
# One line on purpose: the claims suite reads the routed target off the
# `from ..modeling import ` line.
# fmt: off
diff --git a/tests/integration/defs/.test_durations b/tests/integration/defs/.test_durations
index d86b7b84f838..571426fce20f 100644
--- a/tests/integration/defs/.test_durations
+++ b/tests/integration/defs/.test_durations
@@ -20,7 +20,13 @@
"accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp4::test_gsm8k": 1247.1575,
"accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp8::test_eagle3_gsm8k_and_acceptance": 998.465,
"accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp8::test_gsm8k": 317.73875,
+ "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_disaggregated_dep2": 382.8,
+ "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_and_long_context": 468.77,
+ "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[dep4_fp8_kv]": 377.08,
+ "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[tp4_nvfp4_kv_breakable]": 597.49,
+ "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_mmmu": 793.46,
"accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_breakable_prefill_matches_eager": 730.6153846153846,
+ "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_disaggregated_mtp2": 481.54,
"accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_gsm8k": 643.2887001455605,
"accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp1_gsm8k_mmmu_and_video": 969.7372977443608,
"accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp3_cutedsl_fp8_kv_bf16_state": 658.6539230769231,
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 a8edd1cdd25a..6e879e5bd793 100644
--- a/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py
+++ b/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py
@@ -461,13 +461,11 @@ class TestModelingV2Qwen38FlashNextNvfp4Sm103Default(LlmapiAccuracyTestHarness):
prefill_capture_num_tokens=[512, 1024, 2048, 4096, 8192],
)
- # A collapsed or miscomputed draft path accepts nothing and reads 1.0; every draft
- # accepted reads 1 + max_draft_len. The midpoint of that range is the gsm8k MTP
- # cases' tripwire: at or above it, more than half the drafts are accepted.
- MTP3_COLLAPSE_FLOOR = 1.0 + MTP3.max_draft_len / 2
-
# The pass keys' acceptance-length floor in references/acceptance_length.yaml.
ACCEPTANCE_KEY_LONG_CONTEXT = "ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3_long_context"
+ # The four-GPU gsm8k cases' floors there: tep4 over bf16 KV, and per quantized KV dtype
+ # (`mtp3_{kv_dtype}_kv`), each 95% of its GB300 reading.
+ ACCEPTANCE_KEY_MTP3 = "ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3"
# ---- one rank (tp1) ----------------------------------------------------------------
# At tp1 `mamba_ssm_cache_dtype` is float32 because the checkpoint declares
@@ -537,12 +535,11 @@ class TestModelingV2Qwen38FlashNextNvfp4Sm103Default(LlmapiAccuracyTestHarness):
"ModelingV2Qwen38FlashNextNvfp4Sm103Default::tp1_mtp7_nvfp4_kv_bf16_state_long_context"
)
- def _gsm8k(self, mocker, llm, label: str, key: str | None = None) -> None:
+ def _gsm8k(self, mocker, llm, label: str, key: str) -> None:
"""gsm8k against the engine's reference row, then its acceptance length.
The acceptance length is held to the floor under `key` in
- references/acceptance_length.yaml, or without one above the three-draft collapse
- tripwire. Rejection sampling holds the emitted distribution to the target model's, so
+ references/acceptance_length.yaml. Rejection sampling holds the emitted distribution to the target model's, so
a miscomputed draft layer leaves the score in place and only slows generation; the
acceptance length is what sees it. It comes from the engine's iteration statistics,
so the engine needs ``enable_iter_perf_stats``.
@@ -551,13 +548,7 @@ def _gsm8k(self, mocker, llm, label: str, key: str | None = None) -> None:
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}")
- if key is not None:
- assert_acceptance_length(key, acceptance_length)
- return
- assert acceptance_length >= self.MTP3_COLLAPSE_FLOOR, (
- f"{label}: acceptance length {acceptance_length:.3f} is below "
- f"{self.MTP3_COLLAPSE_FLOOR}: the draft path is rejected more often than accepted"
- )
+ assert_acceptance_length(key, acceptance_length)
def _passkeys(self, llm, lengths, draft_len, key, label) -> None:
"""Pass keys over prompts of `lengths` tokens, the acceptance length held to `key`."""
@@ -691,12 +682,13 @@ def _passkey_acceptance(self, llm, prompts, label, draft_len: int = 3) -> float:
def test_mtp3_gsm8k_and_long_context(self, mocker):
"""tep4 with three drafts and IndexShare: gsm8k, then pass keys over 3k-6k tokens.
- gsm8k reads the MTP row of references/gsm8k.yaml, its acceptance length held above
- the collapse tripwire. The pass keys are past the QSA indexer's 2048-token budget, so
- the target's decode steps and its draft layer's take the sparse attention path, which
- gsm8k never leaves, and IndexShare makes the draft loop reuse the selection of its
- first step. A miscomputed draft path keeps the text correct and only lowers the
- acceptance length, which is held to its floor in references/acceptance_length.yaml.
+ gsm8k reads the MTP row of references/gsm8k.yaml, its acceptance length held to the
+ `mtp3` floor in references/acceptance_length.yaml. The pass keys are past the QSA
+ indexer's 2048-token budget, so the target's decode steps and its draft layer's take the
+ sparse attention path, which gsm8k never leaves, and IndexShare makes the draft loop
+ reuse the selection of its first step. A miscomputed draft path keeps the text correct
+ and only lowers the acceptance length, which is held to its floor in
+ references/acceptance_length.yaml.
"""
tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH)
prompts = _passkey_prompts(tokenizer, _PASSKEY_LENGTHS)
@@ -715,7 +707,7 @@ def test_mtp3_gsm8k_and_long_context(self, mocker):
env_overrides=rank_env,
) as llm,
):
- self._gsm8k(mocker, llm, "tep4 mtp3 gsm8k")
+ self._gsm8k(mocker, llm, "tep4 mtp3 gsm8k", self.ACCEPTANCE_KEY_MTP3)
acceptance_length = self._passkey_acceptance(llm, prompts, "tep4 mtp3 long context")
assert_acceptance_length(self.ACCEPTANCE_KEY_LONG_CONTEXT, acceptance_length)
@@ -726,7 +718,7 @@ def test_mtp3_gsm8k_quantized_kv(self, mocker, parallel, kv_dtype, breakable):
"""gsm8k with three drafts over a quantized KV pool, block reuse and chunked prefill.
gsm8k reads the MTP row of references/gsm8k.yaml for that KV dtype, its acceptance
- length held above the collapse tripwire.
+ length held to that dtype's floor in references/acceptance_length.yaml.
"""
kv_cache_config = KvCacheConfig(
dtype=kv_dtype,
@@ -749,7 +741,9 @@ def test_mtp3_gsm8k_quantized_kv(self, mocker, parallel, kv_dtype, breakable):
env_overrides=rank_env,
) as llm,
):
- self._gsm8k(mocker, llm, f"mtp3 {kv_dtype} kv")
+ self._gsm8k(
+ mocker, llm, f"mtp3 {kv_dtype} kv", f"{self.ACCEPTANCE_KEY_MTP3}_{kv_dtype}_kv"
+ )
@skip_not_sm103
@pytest.mark.skip_less_device(4)
@@ -1080,7 +1074,6 @@ def test_tp1_mtp7_long_context_passkey(self, monkeypatch):
AGENTX_ENV = {
"TRTLLM_ENABLE_PDL": "1",
"TRTLLM_KV_TRANSFER_NUM_THREADS": "8",
- "TRTLLM_QWEN4_EXP_PLE_HOST_OFFLOAD": "1",
"TRTLLM_MODELING_V2_DENSE_FP8": "1",
"TRTLLM_MODELING_V2_GDN_INPROJ_NVFP4": "1",
"TRTLLM_MODELING_V2_MTP_MOE_NVFP4": "1",
@@ -1110,7 +1103,8 @@ def test_tp1_disaggregated_mtp2(self, mocker):
"enable_chunked_prefill": True,
"disable_mm_encoder": True,
"moe_config": {"backend": "TRTLLM"},
- "scheduler_config": {"capacity_scheduler_policy": "GUARANTEED_NO_EVICT"},
+ # What KV cache manager v2 runs; it overrides any other policy with a warning.
+ "scheduler_config": {"capacity_scheduler_policy": "MAX_UTILIZATION"},
"cache_transceiver_config": {"backend": "NIXL", "transceiver_runtime": "PYTHON"},
}
kv_cache_config = {
diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml
index 308420d9a5aa..7ee8d7875f08 100644
--- a/tests/integration/defs/accuracy/references/acceptance_length.yaml
+++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml
@@ -227,6 +227,23 @@ ModelingV2Qwen38FlashNextNvfp4Sm103Default::tp1_mtp7_nvfp4_kv_bf16_state_long_co
min_al: 6.69
# Read by TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_and_long_context
+# (test_qwen4_exp.py), over gsm8k at tep4: three drafts with IndexShare on, bf16 KV.
+# modeling_v2's own full-gsm8k reading on GB300; min_al is 95% of it, above the 2.8 that a
+# never-accepted third draft reads.
+ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3:
+ ref_al: 3.689
+ min_al: 3.50
+# Read by test_mtp3_gsm8k_quantized_kv[tp4_nvfp4_kv_breakable]: tp4 over NVFP4 KV with the
+# breakable prefill graph. Measured and floored as the mtp3 entry.
+ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3_nvfp4_kv:
+ ref_al: 3.535
+ min_al: 3.36
+# Read by test_mtp3_gsm8k_quantized_kv[dep4_fp8_kv]: dep4 (MTP under attention DP) over FP8 KV.
+# Measured and floored as the mtp3 entry.
+ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3_fp8_kv:
+ ref_al: 3.512
+ min_al: 3.34
+# Read by TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_and_long_context
# (test_qwen4_exp.py), over the pass keys of tp1_mtp3_long_context at tep4: three drafts with
# IndexShare on, bf16 KV.
ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3_long_context:
diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml
index 5e8f195c6f93..fb98d7d27cec 100644
--- a/tests/integration/test_lists/test-db/l0_b300.yml
+++ b/tests/integration/test_lists/test-db/l0_b300.yml
@@ -205,11 +205,13 @@ l0_b300:
- unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_rms_norm_gated_token_major.py
- unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_situ_and_mul.py
- accuracy/modeling_v2/test_gpt_oss.py::TestModelingV2GptOss120bSm10xDefault::test_gsm8k[tp1]
- # qwen4_exp target at tp1, pre-merge: one engine with three drafts and IndexShare over bf16 KV,
- # gsm8k and then 3k-6k-token pass keys, the only prompts long enough to run the sparse
- # attention path, with their acceptance-length floor (what sees a miscomputed draft layer).
- # The other engines run post-merge (below); no speculation and the breakable prefill graph's
- # comparison with eager run in QA (llm_function_core).
+ # qwen4_exp target at tp1, pre-merge: gsm8k without speculation; one MTP engine for gsm8k
+ # with its acceptance-length floor, mmmu (the vision tower) and synthetic videos; and one
+ # engine with three drafts and IndexShare over bf16 KV, gsm8k and then 3k-6k-token pass
+ # keys, the only prompts long enough to run the sparse attention path. The other engines
+ # run post-merge (below).
+ - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_gsm8k
+ - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp1_gsm8k_mmmu_and_video
- accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp3_gsm8k_and_long_context
- condition:
ranges:
@@ -231,11 +233,12 @@ 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 at tp1, post-merge: one draft with mmmu (the vision tower) and synthetic
- # videos; three drafts over CuteDSL experts, fp8 KV and a bf16 recurrent state (gsm8k and
- # the pass keys); seven drafts over 3k-12k-token pass keys on NVFP4 KV and a bf16 state (the
- # verify step's eight rows per request). Its disaggregated engines run in
+ # qwen4_exp target at tp1, post-merge: three drafts over CuteDSL experts, fp8 KV and a bf16
+ # recurrent state (gsm8k and the pass keys); seven drafts over 3k-12k-token pass keys on
+ # NVFP4 KV and a bf16 state (the verify step's eight rows per request); and the breakable
+ # prefill CUDA graph replayed against eager greedy continuations at the bucket edges and
+ # with context and generation rows in one batch. Its disaggregated engines run in
# l0_gb300_multi_gpus.
- - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp1_gsm8k_mmmu_and_video
- accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp3_cutedsl_fp8_kv_bf16_state
- accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp7_long_context_passkey
+ - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_breakable_prefill_matches_eager
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py
index bb60f45d720c..6e1b621026ca 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py
@@ -1,7 +1,9 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
"""GPU test for the fused_sigmoid_mul_inplace catalog entry: every cell against
-the entry's reference with nothing written past `attention_output`'s last row."""
+the entry's reference with nothing written past `attention_output`'s last row,
+graph capture of the default target's fused-row and MTP forms, the band refusing
+a wrong gate read, and the loud layout refusal."""
import zlib
@@ -58,6 +60,14 @@ def randn(*shape):
return out, out_buffer, gate_buffer, gate
+def _mtp_spec(tokens: int) -> dict:
+ """The default MTP layer's 6-head interleaved form at `tokens` rows."""
+ spec = next(
+ c.spec for c in op.CELLS if c.spec["gate"] == "interleaved" and c.spec["heads"] == 6
+ )
+ return dict(spec, tokens=tokens)
+
+
@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)
@@ -72,3 +82,82 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th
torch.testing.assert_close(result, expected)
assert torch.equal(gate_buffer, gate_before)
assert torch.equal(out_buffer[out.numel() :], tail_before), "a masked lane stored"
+
+
+def test_cuda_graph_capture_and_replay() -> None:
+ """The decode call sites run captured: the tp4 / tep4 form (6 heads off the
+ 4352-wide row, two 1024-wide blocks a row) replays on refilled inputs."""
+ spec = next(c.spec for c in op.CELLS if c.spec.get("row_stride") == 4352)
+ spec = dict(spec, tokens=8)
+ out, _, gate_buffer, gate = _build(spec, seed=5)
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ op(out, gate) # JIT compile outside capture
+ torch.cuda.current_stream().wait_stream(side)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ op(out, gate)
+
+ for seed in (11, 12):
+ fresh_out, _, fresh_gate_buffer, _ = _build(spec, seed=seed)
+ out.copy_(fresh_out)
+ gate_buffer.copy_(fresh_gate_buffer)
+ expected = op.reference(out, gate)
+ graph.replay()
+ torch.cuda.synchronize()
+ torch.testing.assert_close(out, expected)
+
+
+@pytest.mark.parametrize("tokens", [1, 128, 1024])
+def test_mtp_gate_graph_capture_and_replay(tokens) -> None:
+ """The MTP drafter and verify run captured: the 6-head `[T, 6, 512]` gate half
+ replays on refilled inputs, on both block widths (1024 below 1024 rows, 2048 from)."""
+ spec = _mtp_spec(tokens)
+ out, _, gate_buffer, gate = _build(spec, seed=5)
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ op(out, gate) # JIT compile outside capture
+ torch.cuda.current_stream().wait_stream(side)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ op(out, gate)
+
+ for seed in (11, 12):
+ fresh_out, _, fresh_gate_buffer, _ = _build(spec, seed=seed)
+ out.copy_(fresh_out)
+ gate_buffer.copy_(fresh_gate_buffer)
+ expected = op.reference(out, gate)
+ graph.replay()
+ torch.cuda.synchronize()
+ torch.testing.assert_close(out, expected)
+
+
+@pytest.mark.parametrize("tokens", [1, 1024])
+def test_the_gate_discriminates_a_wrong_gate_read(tokens) -> None:
+ """The default bf16 band refuses the two reads a wrong per-head indexing gives:
+ the q half of each head, and the first `hidden` elements of the flat
+ `[T, 6 * 512]` row (head stride ignored)."""
+ spec = _mtp_spec(tokens)
+ hidden = spec["heads"] * spec["head_dim"]
+ out, _, gate_buffer, gate = _build(spec, seed=_seed(f"wrong read {tokens}"))
+ original = out.clone()
+ op(out, gate)
+ for name, wrong in (
+ ("q half", gate_buffer[:, :, : spec["head_dim"]]),
+ ("flat row", gate_buffer.reshape(tokens, -1)[:, :hidden]),
+ ):
+ with pytest.raises(AssertionError):
+ torch.testing.assert_close(out, op.reference(original, wrong))
+ assert (out.float() - op.reference(original, wrong).float()).abs().max() > 0.5, name
+
+
+def test_metadata_misuse_raises() -> None:
+ """A gate with a strided last dim is refused loudly, as is every layout
+ violation; hence the entry has no `is_valid`."""
+ g = torch.Generator(device="cuda").manual_seed(8)
+ attention_output = torch.randn(4, 128, generator=g, device="cuda").to(torch.bfloat16)
+ gate = torch.randn(4, 256, generator=g, device="cuda").to(torch.bfloat16)[:, ::2]
+ with pytest.raises(AssertionError):
+ op(attention_output, gate)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_silu_and_mul_rows.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_silu_and_mul_rows.py
index 7a4874047e69..a5fff4fd0fa2 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_silu_and_mul_rows.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_silu_and_mul_rows.py
@@ -2,7 +2,7 @@
# SPDX-License-Identifier: Apache-2.0
"""GPU test for the silu_and_mul_rows entry: the cells bitwise `flashinfer_silu_and_mul` on a
contiguous copy of the gate|up columns (and so inside its reference band) with the fused
-output's extra columns present."""
+output's extra columns present, graph capture, the width guard."""
import zlib
@@ -42,3 +42,27 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by
torch.testing.assert_close(out, op.reference(x, INTER))
if x.shape[0]:
assert torch.equal(out, flashinfer_silu_and_mul(x[:, : 2 * INTER].contiguous()))
+
+
+@pytest.mark.parametrize("rows", [1, 32, 128])
+def test_cuda_graph_capture_and_replay(rows) -> None:
+ """The decode call sites run captured: one request, a dep4 rank's 32, a 128-request batch."""
+ x = _build(dict(rows=rows), 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=rows), 9))
+ graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(out, op(x, INTER))
+
+
+def test_guards() -> None:
+ x = _build(dict(rows=4), 5)
+ with validating(op), pytest.raises(AssertionError):
+ op(x[:, : 2 * INTER - 2], INTER)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_canonical_selection.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_canonical_selection.py
index 8701e8322c8c..37092d1c827a 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_canonical_selection.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_canonical_selection.py
@@ -3,11 +3,12 @@
"""GPU test for the qsa_canonical_selection entry.
The rewrite fixes every slot of every row, so `_compare` is index equality against a
-composite-key top-k, and against `_torch_rewrite`, the composed-torch rewrite the fused
-kernel replaced, which also pins the rows outside the contract (NaN in a window, a
-window cut by the table's width). The input rows come from an adversary that keeps a
-random part of the tie class at the k-th value in a random order, or from the radix
-engine the target calls.
+composite-key top-k, and against `_torch_rewrite`, the same rule as composed torch ops,
+which also pins the rows outside the contract (NaN in a window, a window cut by the
+table's width). The input rows come from an adversary that keeps a random part of the
+tie class at the k-th value in a random order, or from the radix engine the target
+calls. Beyond the cells: graph replay at 256 and 1024 rows, and the guard against a
+one-element `row_starts`.
"""
import pytest
@@ -49,9 +50,9 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None:
def _torch_rewrite(logits, row_starts, row_ends, indices) -> torch.Tensor:
- """The rewrite as composed torch ops, in row tiles of 2^24 elements: the
- implementation the fused kernel replaced, kept as a second oracle that is bitwise the
- old rows wherever they disagree with `reference` (outside the contract)."""
+ """The rewrite as composed torch ops, in row tiles of 2^24 elements: a second oracle,
+ and the one that fixes the rows outside the contract, where it and `reference` may
+ disagree."""
num_rows, num_cols = logits.shape
index_topk = indices.shape[1]
tile = max(1, (1 << 24) // max(num_cols, 1))
@@ -173,11 +174,11 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811
for rows in row_counts:
args = _build(spec, rows, seed + rows)
expected = op.reference(*args)
- old = _torch_rewrite(args[0], args[1], args[2], args[3].clone())
+ composed = _torch_rewrite(args[0], args[1], args[2], args[3].clone())
out = op(*args)
assert out is args[3], "the rewrite writes `indices` and returns it"
_compare(out, expected)
- _compare(out, old)
+ _compare(out, composed)
def _random_windows(rows: int, cols: int, gen: torch.Generator):
@@ -266,6 +267,37 @@ def test_rows_outside_the_contract_match_the_torch_rewrite(rows) -> None:
_compare(op(scores, starts, ends, raw.clone()), expected)
+@pytest.mark.parametrize("rows", (256, _SPLIT_ROWS))
+def test_replays_in_a_cuda_graph(rows) -> None:
+ """Captured once, replayed over new scores and new picks in the same buffers."""
+ gen = torch.Generator(device=DEV).manual_seed(77)
+ cols = 32768
+ logits = torch.empty(rows, cols, device=DEV)
+ starts = torch.empty(rows, dtype=torch.int32, device=DEV)
+ ends = torch.empty(rows, dtype=torch.int32, device=DEV)
+ indices = torch.empty(rows, BLOCK_TOPK, dtype=torch.int32, device=DEV)
+
+ def stage(kind):
+ scores = _logits(kind, rows, cols, gen)
+ lo, hi = _random_windows(rows, cols, gen)
+ _poison(scores, lo, hi, "poison")
+ logits.copy_(scores)
+ starts.copy_(lo)
+ ends.copy_(hi)
+ indices.copy_(_adversary(scores, lo, hi, BLOCK_TOPK, gen))
+ return op.reference(logits, starts, ends, indices)
+
+ stage("relu")
+ op(logits, starts, ends, indices) # compile outside the capture
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ op(logits, starts, ends, indices)
+ for kind in ("relu_rare", "quantized"):
+ expected = stage(kind)
+ graph.replay()
+ _compare(indices, 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."""
@@ -280,3 +312,22 @@ def test_narrowing_the_table_to_a_bound_on_the_windows_changes_nothing() -> None
narrow = op(logits[:, :bound], starts, ends, raw.clone())
_compare(wide, op.reference(logits, starts, ends, raw))
assert torch.equal(wide, narrow)
+
+
+def test_guard_refuses_a_broadcast_row_start() -> None:
+ """A one-element `row_starts` broadcasts over every row without an error and frames
+ every row by row 0's start. Shown unguarded first, then refused."""
+ gen = torch.Generator(device=DEV).manual_seed(9)
+ rows = 4
+ logits = torch.relu(torch.randn(rows, COLS, generator=gen, device=DEV))
+ starts = torch.tensor([0, 10, 20, 30], dtype=torch.int32, device=DEV)
+ ends = torch.full((rows,), 2000, dtype=torch.int32, device=DEV)
+ raw = _adversary(logits, starts, ends, BLOCK_TOPK, gen)
+ expected = op.reference(logits, starts, ends, raw)
+ out = op(logits, starts[:1], ends, raw.clone())
+ assert out.shape == expected.shape and not torch.equal(out, expected), (
+ "the broadcast hazard stopped reproducing"
+ )
+ with pytest.raises(AssertionError, match="one entry per logits row"):
+ with validating(op):
+ op(logits, starts[:1], ends, raw.clone())
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_index_scores.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_index_scores.py
index 84f02ccb594e..df021ce8d39f 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_index_scores.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_index_scores.py
@@ -17,7 +17,8 @@
ungrouped decode-sized visible-only calls and request-group kernel for verify-shaped
ones, and the engine's row-wise and query-tile kernels. `_watch_launches` swaps the four Triton entry points
for recording proxies that still launch the real kernel, so every cell asserts
-the branch it names.
+the branch it names. Beyond the cells: capture and replay at every measured-pool decode
+and verify cell, and the refusal of metadata shorter than the rows by both launchers.
"""
import math
@@ -330,6 +331,36 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811
_CAPTURED = tuple(c for c in op.CELLS if "pool_pages" in c.spec and not c.spec["context_rows"])
+@pytest.mark.parametrize("cell", _CAPTURED, ids=[c.why[:44] for c in _CAPTURED])
+def test_captured_call_replays_like_eager(cell) -> None:
+ """Capture the call at the cell's batch and flags over the measured pool, rewrite q and
+ every mapped key in place, replay, and match an eager call bitwise on every specified
+ column.
+ """
+ args = _build(cell.spec, seed=sum(map(ord, cell.why)) + 1)
+ visible = _visible(args)
+ table = args["block_table"]
+ mapped = torch.unique(table[table >= 0]).long()
+ warmup = torch.cuda.Stream()
+ warmup.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(warmup):
+ op(**args)
+ torch.cuda.current_stream().wait_stream(warmup)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ captured = op(**args)
+ gen = torch.Generator(device=DEV).manual_seed(7)
+ q, cache = args["q"], args["index_cache"]
+ for trial in range(2):
+ q.copy_(torch.randn(q.shape, generator=gen, device=DEV).to(q.dtype))
+ cache[mapped] = torch.randn(
+ (mapped.numel(), *cache.shape[1:]), generator=gen, device=DEV
+ ).to(cache.dtype)
+ graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(captured[visible], op(**args)[visible]), f"replay {trial}"
+
+
#: The visible-split cells on the tune's table: their grid is fixed at capture, the windows not.
_SPLIT_CAPTURED = tuple(
c for c in _CAPTURED if c.spec["branch"] == "visible_split" and c.spec["pages"] > 516
@@ -458,3 +489,14 @@ def test_request_group_is_bitwise_the_row_wise_kernel(group, rows, heads, mixed,
"the request-group kernel's bits differ from the row-wise kernel's"
)
_compare(grouped, op.reference(**args))
+
+
+@pytest.mark.parametrize("only_visible_blocks", [False, True], ids=["engine", "visible_split"])
+def test_rejects_metadata_shorter_than_the_rows(only_visible_blocks) -> None:
+ """The callable's own refusal, message as observed -- the reason this entry has no
+ `is_valid`: positions for fewer rows than `q` carries. Both kernels' launchers refuse.
+ """
+ args = _build(_case(16, 2, only_visible_blocks=only_visible_blocks), 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/attention/test_modeling_v2_qsa_paged_kv_store.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_kv_store.py
index dc4581afb8c9..4a574340111d 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_kv_store.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_kv_store.py
@@ -15,6 +15,9 @@
memory mapped under the pages the rows touch only (`_EnginePool`); the pool's last page is
always one of them. The reference runs on a dense copy of those pages, and every other mapped
byte is checked unchanged.
+
+Beyond the cells: graph replay at the 32 MTP7 verify blocks and at 128 tep4 decode rows over
+the engine pool, and the contiguity guard on the row metadata.
"""
import zlib
@@ -319,6 +322,9 @@ def _decode_args(seed: int, **overrides):
_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")}
+#: The tep4 decode cell of 128 rows, over its engine pool.
+_TEP4_128 = next(c.spec for c in op.CELLS if c.why.startswith("tep4 decode, 128 rows"))
+
def test_cuda_graph_replay_follows_updated_positions() -> None:
"""Captured once at 32 MTP7 verify blocks, replayed against new positions written
@@ -344,3 +350,64 @@ def call():
graph.replay()
torch.cuda.synchronize()
_compare((args["k_cache"], args["v_cache"]), ref)
+
+
+def test_cuda_graph_replay_on_the_tep4_engine_pool() -> None:
+ """The tep4 decode site as the engine captures it: 128 rows of one KV head into a
+ 257109-page pool. Captured with every row on its page's first slot, replayed after the
+ real positions are written into the same buffer."""
+ args, pool = _build(_TEP4_128, seed=11)
+ try:
+ positions = args["logical_positions"]
+ saved = positions.clone()
+ tpb = args["tokens_per_block"]
+ positions.copy_(saved - saved % tpb) # the same pages, so only mapped memory is written
+
+ def call():
+ op(**args)
+
+ for _ in range(3): # compile the Triton kernel before capture
+ call()
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ call()
+
+ pool.fill()
+ positions.copy_(saved)
+ ref = _engine_reference(args, pool)
+ graph.replay()
+ torch.cuda.synchronize()
+ _compare(pool.gathered(), ref)
+ pool.assert_rest_unchanged()
+ finally:
+ pool.close()
+
+
+def test_guard_refuses_strided_row_metadata() -> None:
+ """`is_valid`'s one guard, after the hazard it exists for.
+
+ The kernel receives strides for k, v, both caches and the table, but reads
+ the two metadata vectors as flat unit-stride buffers. A column of a 2-wide
+ buffer is read through its base pointer: accepted, and rows land elsewhere.
+ """
+ args, pool = _decode_args(seed=10)
+ requests = args["request_indices"]
+ rows = requests.numel()
+ wide_requests = torch.full((rows, 2), rows - 1, dtype=requests.dtype, device=DEV)
+ wide_requests[:, 0] = requests
+ ref = op.reference(**args)
+ strided = {**args, "request_indices": wide_requests[:, 0]}
+
+ pool.fill_(FILL)
+ out = op(**strided) # unguarded: accepted
+ assert not all(torch.equal(_bits(got), _bits(want)) for got, want in zip(out, ref)), (
+ "strided metadata was expected to place rows elsewhere"
+ )
+
+ pool.fill_(FILL)
+ before = pool.clone()
+ with pytest.raises(AssertionError, match="must be a contiguous 1-D tensor"):
+ with validating(op):
+ op(**strided)
+ assert torch.equal(_bits(pool), _bits(before)), "a refused call still wrote"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa.py
index 91ea1f705957..3efa3deba45d 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa.py
@@ -15,10 +15,13 @@
records which ones a call indexed with a grid, and whether each launched with
PDL, so every cell asserts its branch (and, where it names `splits`, the
split count). A cell marked `capture` is also captured in a CUDA graph,
-refilled in place and replayed, as the engine's decode step is.
+refilled in place and replayed, as the engine's decode step is. Beyond the cells:
+bitwise determinism on a 32-split launch, and the refusal of positions passed
+without a compress ratio.
"""
import math
+import re
import pytest
import torch
@@ -351,3 +354,21 @@ def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811
def _with(**overrides):
return {**CKPT, **overrides}
+
+
+def test_calls_are_bitwise_deterministic() -> None:
+ """Including a 32-split launch whose merge is a second kernel."""
+ spec = _with(segments=tuple((4000 + 37 * i, 2) for i in range(4)), gate=None)
+ args = _build(spec, seed=8800)
+ assert torch.equal(op(**args), op(**args))
+
+
+def test_rejects_positions_without_a_compress_ratio() -> None:
+ """The op's own loud refusal, verbatim -- the reason `is_valid` is empty: half of
+ the prefill pair."""
+ args = _build(_with(segments=((600, 1),) * 4, gate=None), seed=9100)
+ args["selected_tokens"] = args["selected_tokens"][:, :64].contiguous()
+ args["query_positions"] = torch.zeros(4, dtype=torch.int64, device=DEV)
+ message = "QSA causal-prefix bounds require positions and compression ratio"
+ with pytest.raises(ValueError, match=re.escape(message)):
+ op(**args)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py
index 7ec9a3535fbc..a30140a8e158 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py
@@ -2,7 +2,8 @@
# 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), and the selection's read-after-wait under graph replay with PDL."""
+relative RMS), the selection's read-after-wait under graph replay with PDL, and the row
+guard."""
import zlib
@@ -108,3 +109,10 @@ def step():
return op(*args[:4], sel, *args[5:])
_replays_like_eager(step, lambda: scores.copy_(torch.rand(scores.shape, device=DEV)))
+
+
+def test_guard_refuses_rows_past_max_rows() -> None:
+ args = _build(dict(pos=(8700,), gated=True), seed=5)
+ many = [a.expand(300, *a.shape[1:]) if i in (0, 4, 5) else a for i, a in enumerate(args[:6])]
+ with validating(op), pytest.raises(AssertionError):
+ op(*many, *args[6:8], None) # 300 rows: past MAX_ROWS
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py
index f7db6f517d43..d6f27432a1a0 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py
@@ -12,9 +12,11 @@
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.
+Beyond the cells: bitwise determinism and the half-passed prefill pair's refusal.
"""
import math
+import re
import pytest
import torch
@@ -330,3 +332,20 @@ def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811
def _with(**overrides):
return {**CKPT, **overrides}
+
+
+def test_calls_are_bitwise_deterministic() -> None:
+ args = _build(_with(segments=tuple((4000 + 37 * i, 2) for i in range(4))), seed=8800)
+ assert torch.equal(op(**args), op(**args))
+
+
+def test_rejects_positions_without_a_compress_ratio() -> None:
+ """The op's own loud refusal -- the reason the entry carries no `is_valid`: half of
+ the prefill pair."""
+ args = _build(_with(segments=((600, 1),) * 4), seed=9100)
+ args["selected_tokens"] = args["selected_tokens"][:, :64].contiguous()
+ args["query_positions"] = torch.zeros(4, dtype=torch.int64, device=DEV)
+ message = "pass both query_positions and compress_ratio, or neither"
+ with pytest.raises(ValueError, match=re.escape(message)):
+ with validating(op):
+ op(**args)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_decode.py
index 354a8469963d..bb091e0b5cbb 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_decode.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_decode.py
@@ -4,7 +4,8 @@
`qsa_paged_sparse_gqa_nvfp4`
and QSA's bf16 kernel over the dequantized pool, inside the NVFP4 op's band (4 ulp of each
(row, head)'s max, 2 ulp relative RMS), over its cells with and without the output gate; plus
-an unaligned row slice of the selection.
+an unaligned row slice of the selection, determinism, and graph replay against eager at 1 and
+64 rows after every operand is rewritten in place.
The tp4-heads cells run one tp4 / tep4 rank's 6 q heads over 1 KV head, q a view of the packed
`[q | k | v]` rows as the target hands it. The dep4 cells run at one rank's pool as the engine
@@ -266,3 +267,36 @@ def test_unaligned_row_slice_of_the_selection() -> None:
aligned[name] = args[name].clone()
assert args["selected_tokens"].data_ptr() % 16
assert torch.equal(op(**args), op(**aligned))
+
+
+def test_calls_are_bitwise_deterministic() -> None:
+ args = build(tuple((4000 + 37 * i, 2) for i in range(4)), seed=8800)
+ assert torch.equal(op(**args), op(**args))
+
+
+@pytest.mark.parametrize("rows", [1, 64])
+def test_replay_matches_eager(rows) -> None:
+ segments = (
+ ((4000, 1),) if rows == 1 else tuple((100 + (499 * g) % 8000, 1) for g in range(rows))
+ )
+ args = build(segments, seed=77)
+ fresh = build(segments, seed=127)
+ op(**args)
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ captured = op(**args)
+ for name in (
+ "q",
+ "k_cache",
+ "v_cache",
+ "k_scale_cache",
+ "v_scale_cache",
+ "block_table",
+ "selected_tokens",
+ "request_indices",
+ ):
+ args[name].copy_(fresh[name])
+ graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(captured, op(**fresh))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_prefill_compress.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_prefill_compress.py
index f2426eb8d1ec..9bc40b9cec03 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_prefill_compress.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_prefill_compress.py
@@ -15,6 +15,9 @@
them out at bf16, the other 11 layers' planes holding a sentinel that must
survive. Those pools are drawn on the device.
+Beyond the cells: graph capture and replay at the 32 MTP7 verify blocks and at 128 decode
+rows over the tep4 pool, and the contiguity guard on `logical_positions`.
+
A padding row at logical position -1 is deliberately absent: it is an illegal
memory access that poisons the process for every test after it.
"""
@@ -221,11 +224,54 @@ def test_certified_cells(cell, certifying): # noqa: F811
assert top * out.stride(0) >= 2**31, top * out.stride(0)
+def _scene(seed: int, **overrides):
+ """Two requests of 64 tokens from 0, the target's call, at the checkpoint geometry."""
+ spec = {**op.CELLS[0].spec, "segments": ((0, 64),) * 2, "pages": 4, **overrides}
+ return _build(spec, seed=seed)[0]
+
+
#: 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")}
+#: The 128-request decode over the default target's tep4 pool, at its cell's seed.
+_DECODE_128 = next(c for c in op.CELLS if c.why.startswith("tep4 decode: 128 requests"))
+
+
+@pytest.mark.parametrize("batch", ["mtp7_32", "tep4_decode_128"])
+def test_cuda_graph_capture_and_replay(batch) -> None:
+ """The launch grid is fixed by the row count, so it captures and replays: 32 MTP7
+ verify blocks, and 128 decode rows over the default target's coalesced tep4 pool, the
+ decode target's captured batches."""
+ if batch == "mtp7_32":
+ args, flips = _scene(seed=16, **_MTP7_BATCH), _MTP7_BATCH["rounding_flips"]
+ else:
+ spec = _DECODE_128.spec
+ args = _build(spec, seed=zlib.crc32(_DECODE_128.why.encode()))[0]
+ flips = spec.get("rounding_flips", 0)
+ cache = args["index_cache"]
+ before = cache.clone()
+ ref = op.reference(**args)
+ warmup = torch.cuda.Stream()
+ warmup.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(warmup):
+ for _ in range(3):
+ cache.copy_(before)
+ op(**args)
+ torch.cuda.current_stream().wait_stream(warmup)
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ cache.copy_(before)
+ with torch.cuda.graph(graph):
+ op(**args)
+ for _ in range(2):
+ cache.copy_(before)
+ graph.replay()
+ torch.cuda.synchronize()
+ _compare(cache, ref, flips)
+
+
def _strided(tensor: torch.Tensor, dim: int) -> torch.Tensor:
"""The same values as a stride-2 view along `dim` of a zero-filled buffer."""
shape = list(tensor.shape)
@@ -235,3 +281,28 @@ def _strided(tensor: torch.Tensor, dim: int) -> torch.Tensor:
index[dim] = slice(None, None, 2)
wide[tuple(index)] = tensor
return wide[tuple(index)]
+
+
+def test_guard_refuses_strided_logical_positions() -> None:
+ """Shown unguarded first -- accepted and wrong -- then refused: a 1-D operand the
+ kernel reads flat (`request_indices` and `k_norm_weight` carry the same
+ `is_contiguous` guard). One measured max abs error against the right answer: 4.72;
+ the floor below states a bound, not one draw.
+ """
+ args = _scene(seed=12)
+ before = args["index_cache"].clone()
+ ref = op.reference(**args)
+ bad = dict(args, logical_positions=_strided(args["logical_positions"], 0))
+ assert not bad["logical_positions"].is_contiguous()
+
+ out = op(**bad)
+ error = (out.float() - ref.float()).abs().max().item()
+ assert error > 0.25, (
+ f"strided logical_positions was expected to be read wrongly; off by {error}"
+ )
+
+ args["index_cache"].copy_(before)
+ with pytest.raises(AssertionError, match="logical_positions must be contiguous"):
+ with validating(op):
+ op(**bad)
+ assert torch.equal(args["index_cache"], before), "a refused call still wrote"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_rollback_planes.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_rollback_planes.py
index 1568de7ac079..54f16581903c 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_rollback_planes.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_rollback_planes.py
@@ -9,6 +9,8 @@
side-plane page of an FP8 / NVFP4 pool; the default target's cells over core's coalesced bf16 pool at a
tep4 rank's measured page count, with the requests' pages drawn across it and the pool's first
and last page among them, so the last is addressed past 2^31 elements from a plane's base.
+Beyond the cells: graph capture and replay of both calls at 1024 rows, and the guards on a
+strided `old`, a `num_planes` past the table and a short `within`.
"""
import zlib
@@ -210,6 +212,52 @@ def test_restore_cells(cell, layout, certifying): # noqa: F811
_assert_far_page(cell.spec, pool, pages[rejected.bool()])
+def test_cuda_graph_replay():
+ """The verify step is graph-captured: both calls at the envelope's 1024 rows, captured
+ once, then replayed over a new step's slots written into the captured inputs."""
+ (cell,) = [c for c in qsa_rollback_gather.CELLS if c.spec["rows"] == 1024]
+ spec = cell.spec
+ g, pool, pages, within = _setup(cell, spec["layout"])
+ ptrs, strides = plane_table(pool.planes)
+ count, rows = len(pool.planes), pages.shape[0]
+ rejected = _rejected(spec, g)
+ old = torch.empty(count, rows, ROW, dtype=torch.bfloat16, device=DEV)
+
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side): # Triton compiles on the first launch, outside the capture
+ old.copy_(qsa_rollback_gather(ptrs, strides, pages, within, count))
+ qsa_rollback_restore(ptrs, strides, pages, within, rejected, old)
+ torch.cuda.current_stream().wait_stream(side)
+ torch.cuda.synchronize()
+
+ gather_graph, restore_graph = torch.cuda.CUDAGraph(), torch.cuda.CUDAGraph()
+ with torch.cuda.graph(gather_graph):
+ gathered = qsa_rollback_gather(ptrs, strides, pages, within, count)
+ with torch.cuda.graph(restore_graph, pool=gather_graph.pool()):
+ qsa_rollback_restore(ptrs, strides, pages, within, rejected, old)
+
+ new_pages, new_within = _verify_slots(spec, 991)
+ assert not torch.equal(new_pages, pages)
+ pages.copy_(new_pages)
+ within.copy_(new_within)
+ rejected.copy_(_rejected(spec, g))
+ _fill(pool.planes, pages, within, g)
+
+ want_old = qsa_rollback_gather.reference(pool.planes, pages, within, count)
+ gather_graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(gathered.view(torch.int16), want_old.view(torch.int16))
+
+ old.copy_(gathered)
+ _fill(pool.planes, pages, within, g)
+ want = qsa_rollback_restore.reference(pool.planes, pages, within, rejected, old)
+ restore_graph.replay()
+ torch.cuda.synchronize()
+ got = _slots_now(pool.planes, pages, within)
+ assert torch.equal(got.view(torch.int16), want.view(torch.int16))
+
+
def _small():
spec = dict(planes=12, rows=16)
g = torch.Generator(device=DEV).manual_seed(5)
@@ -217,3 +265,42 @@ def _small():
pages, within = _slots(spec, 5, g)
_fill(pool.planes, pages, within, g)
return g, pool, pages, within
+
+
+def test_guard_refuses_a_strided_old():
+ """The kernel reads `old` as dense `[L, R, 128]`: a strided view restores the wrong values,
+ with no error."""
+ g, pool, pages, within = _small()
+ ptrs, strides = plane_table(pool.planes)
+ count, rows = len(pool.planes), pages.shape[0]
+ wide = torch.randn(count, rows, 2 * ROW, generator=g, device=DEV).to(torch.bfloat16)
+ old = wide[..., :ROW]
+ assert not old.is_contiguous()
+ rejected = torch.ones(rows, dtype=torch.bool, device=DEV)
+
+ want = qsa_rollback_restore.reference(pool.planes, pages, within, rejected, old)
+ qsa_rollback_restore(ptrs, strides, pages, within, rejected, old)
+ got = _slots_now(pool.planes, pages, within)
+ assert not torch.equal(got.view(torch.int16), want.view(torch.int16)), (
+ "the hazard this guard covers is gone"
+ )
+
+ with validating(qsa_rollback_restore), pytest.raises(AssertionError, match="contiguous"):
+ qsa_rollback_restore(ptrs, strides, pages, within, rejected, old)
+
+
+def test_guard_refuses_more_planes_than_the_table_holds():
+ """Not driven unguarded: the kernel would read addresses past the table's end."""
+ _, pool, pages, within = _small()
+ ptrs, strides = plane_table(pool.planes[:-1])
+ with validating(qsa_rollback_gather), pytest.raises(AssertionError, match="num_planes"):
+ qsa_rollback_gather(ptrs, strides, pages, within, len(pool.planes))
+
+
+def test_guard_refuses_a_short_within():
+ """Not driven unguarded: rows are counted off `pages`, so `within` would be read past its
+ end."""
+ _, pool, pages, within = _small()
+ ptrs, strides = plane_table(pool.planes)
+ with validating(qsa_rollback_gather), pytest.raises(AssertionError, match=r"one \[R\] pair"):
+ qsa_rollback_gather(ptrs, strides, pages, within[:-1], len(pool.planes))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py
index cc612e23c742..12bb8eada7fa 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py
@@ -7,10 +7,13 @@
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.
+reference are that entry's: the target calls this op in its place on context
+launches. Beyond the cells: bitwise determinism and the refusal of a call without
+the prefill pair.
"""
import math
+import re
import pytest
import torch
@@ -236,3 +239,16 @@ def test_long_context_disjoint_groups_match_the_reference() -> None:
pass, however long its list."""
args = _build(_with([(86000, 197)], slots=4, pages=2700), seed=9300)
_compare(op(**args), op.reference(**args))
+
+
+def test_calls_are_bitwise_deterministic() -> None:
+ args = _build(_with([(0, 2000), (500, 300), (7000, 1)], slots=8), seed=8800)
+ assert torch.equal(op(**args), op(**args))
+
+
+def test_rejects_a_call_without_the_prefill_pair() -> None:
+ """The op's own loud refusal: the row-tile pass is bounded by the prefill pair."""
+ args = _build(_with([(600, 4)], slots=8), seed=9100)
+ args.update(query_positions=None, compress_ratio=None)
+ with pytest.raises(ValueError, match=re.escape("the row-tile kernel needs the prefill pair")):
+ op(**args)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_expand_decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_expand_decode.py
index 28f84c8d508f..68c7a17d6285 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_expand_decode.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_expand_decode.py
@@ -5,8 +5,9 @@
Cells against the entry's reference (a stable descending sort, ties to the higher block,
then the expansion entry's reference): the block list and the token selection exactly.
Columns past each row's visible count carry garbage (huge values, NaN) that must never
-be read. Beyond the cells: the wide rows against the canonical composition, and a
-stale-read check under graph replay at each decode-path cell's shape.
+be read. Beyond the cells: the wide rows against the canonical composition, a
+stale-read check under graph replay at each decode-path cell's shape, and the
+`token_topk` guard.
"""
import zlib
@@ -177,3 +178,21 @@ def step():
want_topk, want_sel = op.reference(src, *args[1:])
assert torch.equal(got_topk, want_topk), f"replay {trial}: block list"
assert torch.equal(graph_sel, want_sel), f"replay {trial}: token selection"
+
+
+def test_guard_refuses_a_token_budget_that_is_not_k_blocks() -> None:
+ """`token_topk` sizes the selection's rows while `k * compress_ratio` places the selected
+ tokens: a smaller budget writes past each row, into the next row and, for the last one,
+ past `out` -- here into the margin rows of a larger buffer, where it shows."""
+ spec = dict(rows=2, pos=(9215, 9211), relu_sparse=True, width=2304)
+ args = _build(spec, seed=5)
+ budget = TOKEN_TOPK // 2 # 256 blocks' worth, while 512 are selected
+ margin = 2
+ whole = torch.full((2 + 2 * margin, budget + CR - 1), -9, dtype=torch.int32, device=DEV)
+ args[6] = budget
+ op(*args, out=whole[margin : margin + 2])
+ torch.cuda.synchronize()
+ assert (whole[margin + 2 :] != -9).any(), "the hazard this guard covers is gone"
+
+ with validating(op), pytest.raises(AssertionError, match="token_topk = k \\* compress_ratio"):
+ op(*args, out=whole[margin : margin + 2])
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue.py
index a49258c73753..aa976cf4b66f 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue.py
@@ -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 qwen4_exp_attn_decode_prologue entry: every cell bitwise the
-composition it replaces (`qwen4_exp_qk_norm_rope`, `qsa_paged_kv_store` or
+composition of the entries it fuses (`qwen4_exp_qk_norm_rope`, `qsa_paged_kv_store` or
`qsa_paged_kv_store_nvfp4`, `qsa_paged_index_store`, `qsa_prefill_compress`) on every
allocation either writes -- the projection, the index query, the K/V pool (with its block
scales over NVFP4), the index plane, the position cache; a stale-read check under graph
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/embedding/test_modeling_v2_ngram_hash.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/embedding/test_modeling_v2_ngram_hash.py
index ce804c238fc9..bc3a4db7efd2 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/embedding/test_modeling_v2_ngram_hash.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/embedding/test_modeling_v2_ngram_hash.py
@@ -4,7 +4,8 @@
Exact integer arithmetic, gated bitwise against the entry's scalar reference.
Beyond the cells: the EOS reset that `ngram_token_history`'s EOS seed relies
-on.
+on, the narrow-multiplier and query-count guards, and graph capture at the
+captured decode and MTP-7 verify widths.
"""
import zlib
@@ -97,3 +98,49 @@ def test_eos_restarts_the_window() -> None:
no_eos = torch.tensor([[11, 22, 33, 44, token]], device=DEV)
assert torch.equal(_hash(fresh, 1), _hash(after_eos, 1))
assert not torch.equal(_hash(fresh, 1), _hash(no_eos, 1))
+
+
+def test_guard_refuses_narrow_multipliers() -> None:
+ ids = _tokens(1, CKPT_CONTEXT_LEN + 4, seed=47)
+ multipliers, sizes, offsets = _constants()
+ narrow = multipliers.to(torch.int32)
+ unguarded = op(ids, narrow, sizes, offsets, EOS, 4)
+ assert not torch.equal(unguarded, _hash(ids, 4)), "the hazard this guard covers is gone"
+ assert int(unguarded.min()) >= 0 and int(unguarded.max()) < PADDED_TABLE_ROWS
+ with pytest.raises(AssertionError, match="int64"):
+ with validating(op):
+ op(ids, narrow, sizes, offsets, EOS, 4)
+
+
+def test_guard_refuses_more_queries_than_history() -> None:
+ """A history of 3 asked for 8 query tokens returns 3 of them, without an error."""
+ ids = _tokens(128, CKPT_CONTEXT_LEN + 1, seed=59, lead_eos=False, eos_density=0.05)
+ unguarded = _hash(ids, 8)
+ assert unguarded.shape == (128, CKPT_CONTEXT_LEN + 1, NGRAM_HEADS), (
+ "the hazard this guard covers is gone"
+ )
+ with pytest.raises(AssertionError, match="num_query_tokens must be in"):
+ with validating(op):
+ _hash(ids, 8)
+
+
+@pytest.mark.parametrize("windows", [128, 1024], ids=["decode_128", "mtp7_verify_1024"])
+def test_cuda_graph_capture_and_replay(windows) -> None:
+ ids = _tokens(windows, CKPT_CONTEXT_LEN + 1, seed=53, lead_eos=False, eos_density=0.2)
+ constants = _constants()
+ warmup = torch.cuda.Stream()
+ warmup.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(warmup):
+ for _ in range(3):
+ op(ids, *constants, EOS, 1)
+ torch.cuda.current_stream().wait_stream(warmup)
+
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ captured = op(ids, *constants, EOS, 1)
+
+ fresh = _tokens(windows, CKPT_CONTEXT_LEN + 1, seed=54, lead_eos=False, eos_density=0.2)
+ ids.copy_(fresh)
+ graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(captured, op.reference(fresh, *constants, EOS, 1))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_bf16_gemv_pdl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_bf16_gemv_pdl.py
index 6f2cd3278b41..0a5219b6a0f6 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_bf16_gemv_pdl.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_bf16_gemv_pdl.py
@@ -5,7 +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.
+output and so would read a stale activation if it read before its wait; and
+the layout guard.
"""
import zlib
@@ -100,3 +101,9 @@ def chain():
# Each link against the reference of what the previous link actually stored.
_compare(out, op.reference(h, w))
h = out
+
+
+def test_guard_refuses_a_row_major_mat_b() -> None:
+ mat_a, mat_b = _build(1, 512, 2560, seed=3)
+ with validating(op), pytest.raises(AssertionError):
+ op(mat_a, mat_b.contiguous())
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_gemv.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_gemv.py
index 368839e703e2..ebba4790a141 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_gemv.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_gemv.py
@@ -2,8 +2,8 @@
# SPDX-License-Identifier: Apache-2.0
"""GPU test: `fp8_gemv` / `fp8_gemv_conv1d` / `fp8_gemv_conv1d_verify` against an fp32
reference over the dequantized weight (128x128-block e4m3), per cell; the verify entry also
-bitwise against the two-launch path it replaces (`fp8_gemv` + `causal_conv1d_triton_update`);
-strided activations."""
+bitwise against the same step in two launches (`fp8_gemv` + `causal_conv1d_triton_update`);
+strided activations; CUDA-graph replay of `fp8_gemv` and the verify entry."""
import pytest
import torch
@@ -114,8 +114,8 @@ def _verify_args(spec, seed):
def _two_launch(x, packed, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad):
- """The target's verify path before the fusion: the projection, then the Triton conv over
- the transposed q|k|v view; returns `(conv, rest)` in the fused entry's layout."""
+ """The verify step in two launches: the projection, then the Triton conv over the
+ transposed q|k|v view; returns `(conv, rest)` in the fused entry's layout."""
proj = fp8_gemv_op(x, packed)
n = x.shape[0] // p
x3 = torch.transpose(torch.reshape(proj[:, :conv_dim], [n, p, conv_dim]), 1, 2)
@@ -176,3 +176,46 @@ def test_conv1d_verify_cells(cell) -> None:
assert not torch.equal(window_a[r, :p], window[r, :p]), r
if spec["snapshot_steps"] > p:
assert torch.equal(window_a[:, p:], window[:, p:]), "steps past p untouched"
+
+
+def test_graph_replay_matches_eager() -> None:
+ _, _, packed = _weight(2560, 6144, 5)
+ x = torch.randn(4, 6144, device=DEV).to(torch.bfloat16)
+ fp8_gemv_op(x, packed)
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ out = fp8_gemv_op(x, packed)
+ x.copy_(torch.randn(4, 6144, device=DEV).to(torch.bfloat16))
+ graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(out, fp8_gemv_op(x, packed))
+
+
+def test_conv1d_verify_graph_replay_matches_eager() -> None:
+ spec = fp8_gemv_conv1d_verify_op.CELLS[2].spec
+ _, _, packed = _weight(spec["n"], spec["k"], 5)
+ x, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad = _verify_args(spec, 9)
+ state0, window0 = conv_state.clone(), window.clone()
+ fp8_gemv_conv1d_verify_op(
+ x, packed, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad
+ )
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ conv, rest = fp8_gemv_conv1d_verify_op(
+ x, packed, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad
+ )
+ x.copy_(torch.randn(x.shape, device=DEV).to(torch.bfloat16))
+ conv_state.copy_(state0)
+ window.copy_(window0)
+ graph.replay()
+ torch.cuda.synchronize()
+ got = (conv.clone(), rest.clone(), conv_state.clone(), window.clone())
+ conv_state.copy_(state0)
+ window.copy_(window0)
+ want = fp8_gemv_conv1d_verify_op(
+ x, packed, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad
+ )
+ for a, b in zip(got, (*want, conv_state, window)):
+ assert torch.equal(a.view(torch.int16), b.view(torch.int16))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_per_tensor_mm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_per_tensor_mm.py
index 7329eb4f2890..366b632ae579 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_per_tensor_mm.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_per_tensor_mm.py
@@ -2,7 +2,8 @@
# SPDX-License-Identifier: Apache-2.0
"""GPU test: `fp8_per_tensor_mm` against its fp32 reference (the activation quantized as the
op does, both operands dequantized) per cell, the activation quantization bitwise against
-torch's e4m3 cast, and a zero activation and strided rows."""
+torch's e4m3 cast, a zero activation and strided rows, and CUDA-graph replay with a new
+activation."""
import pytest
import torch
@@ -68,3 +69,16 @@ def test_zero_activation_and_strided_rows() -> None:
assert torch.equal(
fp8_per_tensor_mm(wide[:, :2560], codes_t, scale), fp8_per_tensor_mm(x, codes_t, scale)
)
+
+
+def test_graph_replay_matches_eager() -> None:
+ x, _, codes_t, scale = _operands(64, 2560, 16480, 7)
+ fp8_per_tensor_mm(x, codes_t, scale)
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ out = fp8_per_tensor_mm(x, codes_t, scale)
+ x.copy_(_operands(64, 2560, 16480, 8)[0] * 4)
+ graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(out, fp8_per_tensor_mm(x, codes_t, scale))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_low_m_bf16_mm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_low_m_bf16_mm.py
index 699442c6d0f3..48475ca7df7b 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_low_m_bf16_mm.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_low_m_bf16_mm.py
@@ -5,6 +5,9 @@
Every cell names the branch `plan` must pick, and a host-only check pins the
branch of the table's other shapes and M buckets, so a table edit that moves a
shape to another kernel fails here before it can change a number silently.
+Beyond the cells: CUDA-graph capture and replay after warmup (the way the
+decode forward runs it), and the guard against a padded `mat_b`, which the
+cuBLAS branch would read as dense.
"""
import zlib
@@ -105,9 +108,50 @@ def test_plan_pins_the_branch_of_the_other_table_shapes() -> None:
assert not wrong, f"(M, N, K, pinned, picked): {wrong}"
+def test_graph_replay_after_warmup() -> None:
+ """Decode runs this under CUDA graphs, with PDL: compile in warmup, capture,
+ then replay with new activations copied into the captured input. One call per
+ branch, in a chain, so a FlashInfer kernel follows cuBLAS and vice versa."""
+ pdl = True
+ shapes = [(1, 10240, 2560), (8, 2560, 6144), (1, 2560, 640), (8, 12288, 2560)]
+ assert [plan(*s).kind for s in shapes] == ["direct", "splitk", "cublas", "splitk"]
+ ins = [_build(m, n, k, seed=i) for i, (m, n, k) in enumerate(shapes)]
+ static = [a.clone() for a, _ in ins]
+ stream = torch.cuda.Stream()
+ stream.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(stream):
+ for _ in range(2):
+ [op(a, b, pdl) for a, (_, b) in zip(static, ins)]
+ torch.cuda.current_stream().wait_stream(stream)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ outs = [op(a, b, pdl) for a, (_, b) in zip(static, ins)]
+ for trial in range(3):
+ g = torch.Generator(device="cuda").manual_seed(100 + trial)
+ for a in static:
+ a.copy_(torch.randn(a.shape, generator=g, device="cuda").to(a.dtype))
+ graph.replay()
+ torch.cuda.synchronize()
+ for a, (_, b), out in zip(static, ins, outs):
+ _compare(out, op.reference(a, b))
+
+
def _padded(m, n, k):
"""A row-padded mat_a and a column-padded mat_b: dense along K, leading
dimension 2K. Both are views a caller gets by slicing a wider tensor."""
a = torch.randn(m, 2 * k, dtype=torch.bfloat16, device="cuda")[:, :k]
w = (torch.randn(n, 2 * k, device="cuda") * k**-0.5).to(torch.bfloat16)[:, :k]
return a, w.t()
+
+
+def test_guard_refuses_a_padded_mat_b() -> None:
+ """Why the guard exists: the cuBLAS branch ignores mat_b's leading dimension
+ and silently computes a different product."""
+ a, b = _padded(4, 640, 640)
+ a = a.contiguous()
+ assert plan(4, 640, 640).kind == "cublas" and b.stride() == (1, 1280)
+ with pytest.raises(AssertionError):
+ _compare(op(a, b), op.reference(a, b))
+ with pytest.raises(AssertionError, match="column-major"):
+ with validating(op):
+ op(a, b)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_mid_rows_mm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_mid_rows_mm.py
index d6dfb627ad4f..e6351288ed57 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_mid_rows_mm.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_mid_rows_mm.py
@@ -2,7 +2,7 @@
# SPDX-License-Identifier: Apache-2.0
"""GPU test: `mid_rows_mm` (FP8 with the per-row activation quantization, and bf16) against an
fp32 reference over the operands it multiplies, per cell; `fp8_row_quant` bitwise against
-torch's e4m3 cast; strided activation rows."""
+torch's e4m3 cast; strided activation rows; CUDA-graph replay with new activations."""
import pytest
import torch
@@ -15,6 +15,7 @@
fp8_row_quant,
fp8_row_quant_op,
mid_rows_mm_op,
+ precompile,
)
assert torch.cuda.is_available(), "mid_rows_mm requires a CUDA device"
@@ -100,3 +101,19 @@ def test_row_quant_is_torchs_cast(cell) -> None:
assert torch.equal(s, (amax / E4M3_MAX).float())
want = (x.float() * (E4M3_MAX / amax).float()[:, None]).to(torch.float8_e4m3fn)
assert torch.equal(q.view(torch.uint8), want.view(torch.uint8))
+
+
+def test_graph_replay_matches_eager() -> None:
+ for mode, n, k, m in (("tensor", 16480, 2560, 64), ("bf16", 512, 2560, 200)):
+ _, mw = _weight(n, k, mode, 3)
+ precompile(mw)
+ x = _activation(m, k, 1).contiguous()
+ mid_rows_mm_op(x, mw)
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ out = mid_rows_mm_op(x, mw)
+ x.copy_(_activation(m, k, 2) * 4)
+ graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(out, mid_rows_mm_op(x, mw))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_nvfp4_gemv.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_nvfp4_gemv.py
index 6fe44edb45fa..55d1ae2a38bb 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_nvfp4_gemv.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_nvfp4_gemv.py
@@ -1,11 +1,11 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
"""GPU test: `nvfp4_gemv` against an fp32 reference over the dequantized NVFP4 weight, per
-cell; the quantizer against its definition; one compiled kernel per row
-octet. The split entries (NVFP4 units on `fp8_gemv`'s kernel) per cell:
-the projection against the fp32 reference, the conv entries also against the two-launch
-paths they replace (the split projection + `causal_conv1d_update` /
-`causal_conv1d_triton_update`)."""
+cell; the quantizer against its definition; CUDA-graph replay; one compiled kernel per row
+octet; the f16-range guard. The split entries (NVFP4 units on `fp8_gemv`'s kernel) per cell:
+the projection against the fp32 reference, the conv entries also against the same steps in
+two launches (the split projection + `causal_conv1d_update` /
+`causal_conv1d_triton_update`); the split projection under CUDA-graph replay."""
import pytest
import torch
@@ -105,6 +105,29 @@ def test_one_kernel_per_row_octet() -> None:
assert nv._compiled.cache_info().currsize - before <= 2
+def test_graph_replay_matches_eager() -> None:
+ _, _, _, _, packed = _weight(4096, 2560, 9)
+ x = torch.randn(4, 2560, device=DEV).to(torch.bfloat16)
+ nvfp4_gemv_op(x, packed)
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ out = nvfp4_gemv_op(x, packed)
+ x.copy_(torch.randn(4, 2560, device=DEV).to(torch.bfloat16))
+ graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(out, nvfp4_gemv_op(x, packed))
+
+
+def test_is_valid_rejects_activation_beyond_f16() -> None:
+ _, _, _, _, packed = _weight(256, 256, 1)
+ x = torch.randn(2, 256, device=DEV).to(torch.bfloat16)
+ nvfp4_gemv_op.is_valid(x, packed)
+ x[1, 3] = 1e5
+ with pytest.raises(AssertionError):
+ nvfp4_gemv_op.is_valid(x, packed)
+
+
# ------------------------------------------------------------------------------------------
# split entries: NVFP4 units on fp8_gemv's kernel
# ------------------------------------------------------------------------------------------
@@ -286,3 +309,18 @@ def test_split_precompile_covers_the_decode_variants() -> None:
x, packed, conv_state, conv_weight, ids, window, ids, GDN_CONV, p, -1
)
assert f8._compiled.cache_info().misses == before
+
+
+def test_split_graph_replay_matches_eager() -> None:
+ _, codes, scales, gscale, _ = _weight(GDN_N, GDN_K, 9)
+ packed = nv.pack_split(codes, scales, gscale)
+ x = torch.randn(4, GDN_K, device=DEV).to(torch.bfloat16)
+ split_op(x, packed)
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ out = split_op(x, packed)
+ x.copy_(torch.randn(4, GDN_K, device=DEV).to(torch.bfloat16))
+ graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(out, split_op(x, packed))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py
index 877e8a866659..e50207eccb6d 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py
@@ -6,7 +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.
+activations before `griddepcontrol.wait` would see the previous replay's; the one-row
+guard.
"""
import zlib
@@ -186,3 +187,10 @@ def step():
want = op(*[*args[:10], args[10]])
for a, b in zip(graph_out, want):
assert torch.equal(a, b), f"replay {trial}"
+
+
+def test_guards() -> None:
+ args = list(_build(dict(), seed=6))
+ two = args[0].expand(2, -1).contiguous()
+ with validating(op), pytest.raises(AssertionError, match="one row"):
+ op(two, args[1].expand(2, -1).contiguous(), args[2].expand(2, -1).contiguous(), *args[3:])
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py
index 831b6bfba981..fd96bfbef797 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py
@@ -4,8 +4,9 @@
(`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), and a
-stale-read check under CUDA-graph replay with PDL for the NVFP4 build."""
+the NVFP4 exact on the returned `mixed` outside the reciprocal tie window), a
+stale-read check under CUDA-graph replay with PDL for the NVFP4 build, the one-row
+guard."""
import zlib
@@ -150,3 +151,12 @@ def step():
torch.cuda.synchronize()
for a, b in zip(graph_out, op(*args)):
assert torch.equal(a, b), f"replay {trial}"
+
+
+def test_guards() -> None:
+ """Both entries share the shape checks; the NVFP4 build's `is_valid` refuses two rows."""
+ op = qwen4_exp_hc_decode_mix_fp4_cute
+ args = _build(dict(), seed=6)
+ two = args[3].expand(2, -1).contiguous()
+ with validating(op), pytest.raises(AssertionError, match="one row"):
+ op(args[0].expand(-1, 2, -1), args[1].expand(-1, 2), args[2], two, *args[4:])
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_rows_tc.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_rows_tc.py
index c9a8f96ab984..6dda064654ea 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_rows_tc.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_rows_tc.py
@@ -8,7 +8,7 @@
runner's arithmetic) followed by the unfolded site. `hc_out` must be bitwise; `mixed` within
1e-2 of the reference's largest magnitude (1-2 bf16 ulp: the projections take bf16 operands on
tensor cores); the injection weights within one bf16 ulp; the NVFP4 codes may differ where
-`mixed` does."""
+`mixed` does. Beyond the cells: CUDA-graph replay of the MoE site at 8 rows, bitwise eager."""
import importlib
@@ -170,3 +170,19 @@ def test_fold_bitwise_vs_runner_finalize(rows, n_cols) -> None:
want = first_op(o["hc"], routed, o["inject"], o["norm"], o["first"], C, n_cols, EPS, **kw)
for a, b in zip(got, want):
assert (a is None and b is None) or torch.equal(a, b)
+
+
+def test_replay_matches_eager() -> None:
+ o = _operands(8, 77)
+ _site(o, "inject", True)
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ captured = _site(o, "inject", True)
+ fresh = _operands(8, 78)
+ o["hc"].copy_(fresh["hc"])
+ o["block"].copy_(fresh["block"])
+ graph.replay()
+ torch.cuda.synchronize()
+ eager = _site(o, "inject", True)
+ assert torch.equal(captured[1], eager[1]) and torch.equal(captured[0], eager[0])
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_first_rescale_tc.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_first_rescale_tc.py
index cd294ad6c6d8..fc9cd3acf0f6 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_first_rescale_tc.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_first_rescale_tc.py
@@ -2,7 +2,8 @@
# SPDX-License-Identifier: Apache-2.0
"""GPU test for the qwen4_exp_hc_first_rescale_tc entry: `low` and `inject` in
the gate band (`qwen4_exp_gated_residual`'s) of the float64 projection of the
-once-rounded normed stream, and the stack-final form.
+once-rounded normed stream, the stack-final form, the weight-layout guard,
+graph capture.
"""
import zlib
@@ -90,3 +91,32 @@ def test_stack_final_low_is_the_injecting_form_bitwise() -> None:
low_f, none = op(*args[:6], False)
assert none is None and inject.shape == (300, HC_COUNT)
assert torch.equal(low, low_f)
+
+
+def test_guard_refuses_a_contiguous_first() -> None:
+ """The kernel streams the `[N, C*H]` weight row-major: `first` must be its
+ transposed view, not a `[C*H, N]` copy.
+ """
+ args = list(_build(dict(rows=4), seed=4))
+ args[3] = args[3].contiguous()
+ with validating(op), pytest.raises(AssertionError, match="transposed"):
+ op(*args)
+
+
+def test_cuda_graph_capture_and_replay() -> None:
+ args = _build(dict(rows=8), seed=5)
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ op(*args)
+ torch.cuda.current_stream().wait_stream(side)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ g_low, g_inject = op(*args)
+ fresh = _build(dict(rows=8), seed=6)
+ args[0].copy_(fresh[0])
+ args[1].copy_(fresh[1])
+ graph.replay()
+ torch.cuda.synchronize()
+ e_low, e_inject = op(*args)
+ assert torch.equal(g_low, e_low) and torch.equal(g_inject, e_inject)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py
index f87d5115878d..af2419ce51aa 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py
@@ -5,7 +5,7 @@
Cells are gated at one ulp of the activation dtype against the float64
reference; the MoE-combine cells widen the residual's band by the combine's own
rounding (see `_compare_moe`). Beyond them: the CuTe fold kernel bitwise the
-Triton kernel at the last row count before the switch.
+Triton kernel at the last row count before the switch, and the slot-map guard.
"""
import sys
@@ -204,3 +204,10 @@ def recording(*a):
want_hc, want_normed = op(*args)
assert torch.equal(hc, want_hc)
assert torch.equal(normed, want_normed)
+
+
+def test_guard_refuses_an_int64_slot_map() -> None:
+ """The fold's slot map is the runner's int32 index map; `is_valid` refuses an int64 one."""
+ fold = _build(dict(rows=4, moe=True, fold=10), seed=9)
+ with validating(op), pytest.raises(AssertionError, match="int32"):
+ op(*fold[:9], fold[9].long())
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_rstd.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_rstd.py
index 5159e2281416..07e4dd651117 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_rstd.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_rstd.py
@@ -6,7 +6,8 @@
widened by the combine's own rounding, as `qwen4_exp_hc_inject_norm`'s test),
rstd within 2^-16 relative of the float64 statistic of the stored residual.
Beyond them: `hc_out` and the normed stream rebuilt from rstd bitwise
-`qwen4_exp_hc_inject_norm`'s at the cells' sizes.
+`qwen4_exp_hc_inject_norm`'s at the cells' sizes, the operand guard, graph
+capture.
"""
import zlib
@@ -137,3 +138,28 @@ def test_normed_rebuilt_from_rstd_is_inject_norm_bitwise(cell) -> None:
normed = (x.flatten(-2) * (1.0 + gamma.float())).to(hc.dtype)
diff = (normed != want_normed).sum().item()
assert diff == 0, f"{diff} normed elements differ"
+
+
+def test_guard_refuses_block_out_without_inject() -> None:
+ hyper_input, block_out, _, eps, c, _, _ = _build(dict(rows=4), seed=4)
+ with validating(op), pytest.raises(AssertionError, match="together"):
+ op(hyper_input, block_out, None, eps, c)
+
+
+def test_cuda_graph_capture_and_replay() -> None:
+ args = list(_build(dict(rows=8, moe=True), seed=5))
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ op(*args)
+ torch.cuda.current_stream().wait_stream(side)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ g_hc, g_rstd = op(*args)
+ fresh = _build(dict(rows=8, moe=True), seed=6)
+ for i in (0, 1, 2, 5, 6):
+ args[i].copy_(fresh[i])
+ graph.replay()
+ torch.cuda.synchronize()
+ e_hc, e_rstd = op(*fresh)
+ assert torch.equal(g_hc, e_hc) and torch.equal(g_rstd, e_rstd)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py
index 07b816f9ae61..c44e4f046098 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py
@@ -2,7 +2,7 @@
# SPDX-License-Identifier: Apache-2.0
"""GPU test for the default target's qwen4_exp_hc_lowrank_act variant
(models/qwen4_exp/ops/hc/qwen4_exp_hc_lowrank_act): cells within one ulp of the float64
-reference."""
+reference, the fp32-logits guard, graph capture."""
import zlib
@@ -55,3 +55,26 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by
low, inject = op(*args)
assert low.shape == (cell.spec["rows"], RANK)
_compare((low, inject), op.reference(*args))
+
+
+def test_guard_refuses_bf16_logits() -> None:
+ logits, rank, c, wi, dt = _build(dict(rows=4), seed=4)
+ with validating(op), pytest.raises(AssertionError):
+ op(logits.to(torch.bfloat16), rank, c, wi, dt)
+
+
+def test_cuda_graph_capture_and_replay() -> None:
+ logits, rank, c, wi, dt = _build(dict(rows=8), seed=5)
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ op(logits, rank, c, wi, dt)
+ torch.cuda.current_stream().wait_stream(side)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ g_low, g_inject = op(logits, rank, c, wi, dt)
+ logits.copy_(_build(dict(rows=8), seed=6)[0])
+ graph.replay()
+ torch.cuda.synchronize()
+ e_low, e_inject = op(logits, rank, c, wi, dt)
+ assert torch.equal(g_low, e_low) and torch.equal(g_inject, e_inject)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_rstd.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_rstd.py
index cee012140e7b..4d1e44e53f76 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_rstd.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_rstd.py
@@ -4,7 +4,7 @@
`mixed` in the gate band of the float64 reference, the epilogue's fp32 operation
sequence bitwise (gate at zero logits), the NVFP4 outputs against the native-torch
quantization of the returned `mixed` outside the reciprocal tie window, the NVFP4
-build's `mixed` bitwise the bf16 build's.
+build's `mixed` bitwise the bf16 build's, the global-scale guard, graph capture.
"""
import zlib
@@ -135,3 +135,41 @@ def test_epilogue_operation_sequence_bitwise(rows) -> None:
def test_fp4_mixed_is_the_bf16_entry_bitwise(rows) -> None:
args = _build(dict(rows=rows), seed=100 + rows)
assert torch.equal(op4(*args)[0], op(*args[:6]))
+
+
+def test_guards() -> None:
+ """The kernel loads `global_scale[0]` only; a per-row scale is refused. rstd must be
+ fp32.
+ """
+ low, up, hc_out, rstd, w1, c, gs = _build(dict(rows=4), seed=4)
+ with validating(op4), pytest.raises(AssertionError):
+ op4(low, up, hc_out, rstd, w1, c, torch.cat([gs, gs]))
+ with validating(op), pytest.raises(AssertionError, match="rstd"):
+ op(low, up, hc_out, rstd.to(torch.bfloat16), w1, c)
+
+
+@pytest.mark.parametrize("fp4", [False, True])
+def test_cuda_graph_capture_and_replay(fp4) -> None:
+ args = list(_build(dict(rows=8), seed=5))
+ if not fp4:
+ args = args[:6]
+ entry = op4 if fp4 else op
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ entry(*args)
+ torch.cuda.current_stream().wait_stream(side)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ g_out = entry(*args)
+ fresh = _build(dict(rows=8), seed=6)
+ for i in (0, 2, 3):
+ args[i].copy_(fresh[i])
+ if fp4:
+ args[6].mul_(0.5)
+ graph.replay()
+ torch.cuda.synchronize()
+ want = entry(*args)
+ g_out, want = (g_out, want) if fp4 else ((g_out,), (want,))
+ for got, exp in zip(g_out, want):
+ assert torch.equal(got, exp)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_post_inject.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_post_inject.py
index 917c563af1b0..72965faf5107 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_post_inject.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_post_inject.py
@@ -2,7 +2,8 @@
# SPDX-License-Identifier: Apache-2.0
"""GPU test for the qwen4_exp_post_inject catalog entry.
-The cells are gated bitwise (16-bit) by `_compare`.
+The cells are gated bitwise (16-bit) by `_compare`. Beyond them: the
+leading-dim guard and graph capture.
"""
import zlib
@@ -74,3 +75,35 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res
assert out.data_ptr() != hyper_input.data_ptr(), "the result is a fresh tensor"
assert torch.equal(hyper_input, before), "the residual it was given is read-only"
_compare(out, op.reference(hyper_input, block_out, inject))
+
+
+def test_guard_refuses_unequal_leading_dims() -> None:
+ """Unguarded, a 1-D `inject` broadcasts one row over every token."""
+ hyper_input, block_out, inject = _operands((8,), seed=10)
+ unguarded = op(hyper_input, block_out, inject[0])
+ assert unguarded.shape == hyper_input.shape, "the hazard this guard covers is gone"
+ assert not torch.equal(unguarded, op.reference(hyper_input, block_out, inject))
+
+ with pytest.raises(AssertionError, match="leading dims must match"):
+ with validating(op):
+ op(hyper_input, block_out, inject[0])
+
+
+def test_cuda_graph_capture_and_replay() -> None:
+ hyper_input, block_out, inject = _operands((64,), seed=14)
+ warmup = torch.cuda.Stream()
+ warmup.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(warmup):
+ for _ in range(3):
+ op(hyper_input, block_out, inject)
+ torch.cuda.current_stream().wait_stream(warmup)
+
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ captured = op(hyper_input, block_out, inject)
+
+ fresh, _, _ = _operands((64,), seed=15)
+ hyper_input.copy_(fresh)
+ graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(captured, op.reference(hyper_input, block_out, inject))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py
index a75ead10c13d..69f7d3cb15c3 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py
@@ -4,8 +4,8 @@
The entry owns the contract. This file builds the checkpoint-shaped operands --
512 experts in the CuteDSL weight layout, the activation as the target quantizes
-it, the sort as moe_sort returns it -- and holds the claim no cell carries:
-per-pair invariance across tiles and tactics.
+it, the sort as moe_sort returns it -- and holds the claims no cell carries:
+per-pair invariance across tiles and tactics, and the contiguity guard.
"""
import functools
@@ -364,3 +364,36 @@ def test_tiles_and_tactics_agree_per_pair(certifying): # noqa: F811
first = seen[(128, 128)]
for key, got in seen.items():
_assert_pairs_equal(got, first, f"T={tokens} {key} vs (128, 128)")
+
+
+# ── the guard ─────────────────────────────────────────────────────────────
+
+
+def _strided(t: torch.Tensor) -> torch.Tensor:
+ """The same values as `t`, as a non-contiguous view into a wider buffer."""
+ u = t.view(torch.uint8) if t.element_size() == 1 else t
+ wide = torch.zeros(*u.shape[:-1], 2 * u.shape[-1], dtype=u.dtype, device=DEV)
+ wide[..., : u.shape[-1]] = u
+ view = wide[..., : u.shape[-1]]
+ return view.view(t.dtype) if t.element_size() == 1 else view
+
+
+def _u8(t: torch.Tensor) -> torch.Tensor:
+ return t.view(torch.uint8) if t.element_size() == 1 else t
+
+
+def test_noncontiguous_operand_answers_wrongly_and_is_refused():
+ """The runner hands data_ptr() to the kernel and never reads a stride: the
+ same values through a strided view come back finite and wrong. The activation;
+ the guard checks input_scale, weight, weight_scale and alpha the same way."""
+ name = "input"
+ tokens, tile = 64, 128
+ args, _ = _problem(tokens, tile, seed=21)
+ ref = op.reference(**args)
+ bad = dict(args, **{name: _strided(args[name])})
+ assert not bad[name].is_contiguous()
+ assert torch.equal(_u8(bad[name]), _u8(args[name]))
+ with pytest.raises(AssertionError):
+ _compare(op(**bad), ref)
+ with validating(op), pytest.raises(AssertionError, match=f"^{name} must be contiguous"):
+ op(**bad)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py
index b7f9d97efb50..4029939cc546 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py
@@ -6,7 +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 claim no cell carries: per-pair invariance across tiles and tactics.
+holds the claims no cell carries: per-pair invariance across tiles and tactics,
+and the contiguity guard.
"""
import functools
@@ -341,8 +342,34 @@ def test_tiles_and_tactics_agree_per_pair(certifying): # noqa: F811
_assert_pairs_equal(got, first, f"T={tokens} {key} vs (128, 128, 1)")
-# ── helpers ───────────────────────────────────────────────────────────────
+# ── the guard ─────────────────────────────────────────────────────────────
def _u8(t: torch.Tensor) -> torch.Tensor:
return t.view(torch.uint8) if t.element_size() == 1 else t
+
+
+def _strided(t: torch.Tensor) -> torch.Tensor:
+ """The same values as `t`, as a non-contiguous view into a wider buffer."""
+ u = _u8(t)
+ wide = torch.zeros(*u.shape[:-1], 2 * u.shape[-1], dtype=u.dtype, device=DEV)
+ wide[..., : u.shape[-1]] = u
+ view = wide[..., : u.shape[-1]]
+ return view.view(t.dtype) if t.element_size() == 1 else view
+
+
+def test_noncontiguous_operand_answers_wrongly_and_is_refused():
+ """The runner hands data_ptr() to the kernel and never reads a stride: the
+ same values through a strided view come back finite and wrong. The activation;
+ the guard checks the other five operands the same way."""
+ name = "input"
+ tokens, tile = 64, 128
+ args, _ = _problem(tokens, tile, seed=21)
+ ref = op.reference(**args)
+ bad = dict(args, **{name: _strided(args[name])})
+ assert not bad[name].is_contiguous()
+ assert torch.equal(_u8(bad[name]), _u8(args[name]))
+ with pytest.raises(AssertionError):
+ _compare(op(**bad), ref)
+ with validating(op), pytest.raises(AssertionError, match=f"^{name} must be contiguous"):
+ op(**bad)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py
index 34ce9252cc48..49531f0bfbea 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py
@@ -7,7 +7,8 @@
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 uses the shared NVFP4 test helpers.
+operand construction uses the shared NVFP4 test helpers. Beyond the cells: the
+uncertified-pair refusal and CUDA-graph replay at one token.
"""
import zlib
@@ -84,3 +85,34 @@ def test_certified_cells(cell, certifying, many_expert_layer): # noqa: F811
with true_fp32_matmul():
exp = op.reference(logits, x, k, ref["up"], ref["gate"], ref["down"], g2)
H._compare(y, exp.to(torch.bfloat16))
+
+
+def test_uncertified_pair_refused(many_expert_layer): # noqa: F811
+ spec = op.CELLS[0].spec
+ args, _, data, sf, _, logits, _, _, _, scalars = H.renorm_operands(
+ many_expert_layer(), spec, 13, renorm_moe_routing_op
+ )
+ with validating(op), pytest.raises(AssertionError, match="certified Renormalize"):
+ op(*_args(args, data, sf, logits, scalars, spec, top_k=9))
+
+
+def test_cuda_graph_capture_and_replay(many_expert_layer): # noqa: F811
+ spec = next(c.spec for c in op.CELLS if c.spec["tokens"] == 1)
+ args, _, data, sf, _, logits, _, _, _, scalars = H.renorm_operands(
+ many_expert_layer(), spec, 43, renorm_moe_routing_op
+ )
+ call = list(_args(args, data, sf, logits, scalars, spec))
+ warmup = torch.cuda.Stream()
+ warmup.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(warmup):
+ for _ in range(3):
+ op(*call)
+ torch.cuda.current_stream().wait_stream(warmup)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ g_rows, g_w, g_idx = op(*call)
+ logits.copy_(torch.randn_like(logits.float()).to(logits.dtype))
+ graph.replay()
+ torch.cuda.synchronize()
+ want = finalized(*call)
+ assert torch.equal(finalize(g_rows, g_w, g_idx), want)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py
index 78019fcf751d..fa7fb853458a 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py
@@ -2,12 +2,13 @@
# 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`),
+Two claims per cell: the logits-routed call is **bitwise** the two-entry composition
+(`renorm_moe_routing_op` -> pre-routed `fp4_block_scale_moe_runner`),
and it sits inside the pre-routed entry's band against this entry's own
fp32 reference. The NVFP4 operand construction (layout, activations, scalars,
-the band) uses ordinary shared test helpers. Beyond the cells: the prefill
-tactic resolution bitwise the custom op's, and per `DECODE_TACTICS` bucket
+the band) uses ordinary shared test helpers. Beyond the cells: the refusal
+of an uncertified (num_experts, top_k) pair, the prefill tactic resolution
+bitwise the custom op's, and per `DECODE_TACTICS` bucket
that the pinned tactic is a valid runner config and answers bitwise what the
AutoTuner path's untuned (runner-default) tactic answers.
"""
@@ -19,6 +20,7 @@
from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_block_scale_moe_renorm_routed import ( # noqa: E501
ACT_SWIGLU,
+ CERTIFIED_PAIRS,
DECODE_TACTICS,
DECODE_TACTICS_GEOMETRY,
ROUTING_RENORMALIZE,
@@ -76,7 +78,7 @@ def test_certified_cells(cell, certifying, many_expert_layer): # noqa: F811
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.
+ # 1. Bitwise the two-entry composition, on identical operands.
y_pre = H._call(
data,
sf,
@@ -98,6 +100,16 @@ def test_certified_cells(cell, certifying, many_expert_layer): # noqa: F811
H._compare(y, exp.to(torch.bfloat16))
+def test_uncertified_pair_refused(many_expert_layer): # noqa: F811
+ spec = op.CELLS[0].spec
+ assert (spec["num_experts"], 9) not in CERTIFIED_PAIRS
+ args, _, data, sf, _, logits, _, _, _, scalars = H.renorm_operands(
+ many_expert_layer(), spec, 13, renorm_moe_routing_op
+ )
+ with validating(op), pytest.raises(AssertionError, match="certified Renormalize"):
+ _routed(args, data, sf, logits, scalars, spec, top_k=9)
+
+
def test_prefill_tactic_resolution_matches_the_custom_op_bitwise(
monkeypatch,
many_expert_layer, # noqa: F811
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed_windows.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed_windows.py
index 80a89853c909..d7a109cb9017 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed_windows.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed_windows.py
@@ -4,7 +4,7 @@
are in test_modeling_v2_fp4_block_scale_moe_renorm_routed.py).
Two claims per cell: the logits-routed call over a rank's expert window is **bitwise**
-the composition it replaces at the call site (`renorm_moe_routing_op` -> pre-routed
+the two-entry composition (`renorm_moe_routing_op` -> pre-routed
`fp4_block_scale_moe_runner` over the same window), and it sits inside the pre-routed
entry's band against this entry's own fp32 reference. The NVFP4 operand construction
(layout, the 512-expert stack, windows, activations, scalars, the band) is the pre-routed
@@ -173,7 +173,7 @@ def test_certified_cells(cell, certifying, many_expert_layer): # noqa: F811
y = _routed(win_args, data, sf, logits, scalars, spec)
assert y.shape == (spec["tokens"], spec["hidden"]) and y.dtype == torch.bfloat16
- # 1. Bitwise the composition it replaces, on identical operands.
+ # 1. Bitwise the two-entry composition, on identical operands.
y_pre = H._call(
data,
sf,
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py
index 6c5e478ff738..ee8626ce9376 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py
@@ -6,7 +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 default target's `weights.py` transforms,
so the test also pins those. Per cell: inside the routed entry's band against
-the fp32 reference.
+the fp32 reference. Beyond the cells: the guard against a per-expert FC2 input
+scale.
"""
import zlib
@@ -83,3 +84,27 @@ def test_certified_cells(cell, certifying, many_expert_layer, cute_dsl_experts):
with true_fp32_matmul():
exp = op.reference(logits, x, spec["top_k"], ref["up"], ref["gate"], ref["down"], g2)
H._compare(y, exp.to(torch.bfloat16))
+
+
+def test_guard_refuses_a_per_expert_fc2_scale(many_expert_layer, cute_dsl_experts): # noqa: F811
+ spec = {**op.CELLS[0].spec, "tokens": 1024}
+ args, ref, data, sf, x, logits, _, _, g2, _ = H.renorm_operands(
+ many_expert_layer(), spec, 9, renorm_moe_routing_op
+ )
+ fc1_w, fc1_s, fc2_w, fc2_s = cute_dsl_experts
+ s1, sg, s2 = H._scalars(spec["num_experts"], G1, g2)
+ with validating(op), pytest.raises(AssertionError):
+ op(
+ logits,
+ data,
+ sf,
+ fc1_w,
+ fc1_s,
+ sg,
+ fc2_w,
+ fc2_s,
+ s2,
+ s2,
+ spec["num_experts"],
+ spec["top_k"],
+ )
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp8_block_scale_moe_runner.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp8_block_scale_moe_runner.py
index 7fc469633394..26a296165a7c 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp8_block_scale_moe_runner.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp8_block_scale_moe_runner.py
@@ -6,7 +6,9 @@
The entry owns the contract: `CELLS`, `reference`, `is_valid`. This
file turns a cell into operands and holds the claims no cell carries: a warm
autotuned tactic is bitwise the cold fallback the cells run (whole stack and
-window), and the default target's four windows sum to the 512-expert layer.
+window), the default target's four windows sum to the 512-expert layer, a
+window call refuses the whole stack's tensors loudly, and the two guards
+(contiguity, a window outside the routing space).
"""
import zlib
@@ -229,6 +231,41 @@ def test_windows_sum_to_the_whole_stack(tokens):
_compare(short.to(torch.bfloat16), ref)
+@pytest.mark.parametrize(
+ "name", ["gemm1_weights", "gemm1_weights_scale", "gemm2_weights", "gemm2_weights_scale"]
+)
+def test_window_call_refuses_the_whole_stack_loudly(name):
+ """A window call handed the whole stack's tensor raises: unlike the NVFP4
+ runner, this op compares every weight and scale's expert dim with
+ `local_num_experts`, so `is_valid` need not."""
+ a = _args(16, **CKPT, local_expert_offset=128, local_num_experts=128, seed=5)
+ stack = dict(
+ zip(
+ ("gemm1_weights", "gemm1_weights_scale", "gemm2_weights", "gemm2_weights_scale"),
+ _weights(2560, 640, 512, seed=2560 + 640 + 512),
+ )
+ )
+ with pytest.raises(RuntimeError, match=f"{name} has incorrect"):
+ op(**dict(a, **{name: stack[name]}))
+
+
+def test_guard_refuses_a_window_outside_the_routing_space():
+ """`num_experts` passed as the window's own count: the raw op answers all
+ zeros with no error, and `is_valid` refuses it.
+
+ Probed at one token only. At 16 rows the same call reads out of bounds (an
+ illegal memory access that would end this process), recorded in the note.
+ """
+ a = _args(1, **CKPT, local_expert_offset=128, local_num_experts=128, seed=3)
+ assert _in_window(a).any(), "the probe needs a slot in the window"
+ wrong = dict(a, num_experts=128)
+ unguarded = op(**wrong)
+ assert torch.count_nonzero(unguarded) == 0, "the hazard this guard covers is gone"
+ assert torch.count_nonzero(op.reference(**a)) > 0
+ with validating(op), pytest.raises(AssertionError, match="outside the routing space"):
+ op(**wrong)
+
+
# ── what a cell cannot carry ──────────────────────────────────────────────
@@ -255,3 +292,21 @@ def test_autotuned_tactic_matches_the_fallback_tactic(offset, local):
_compare(warm, op.reference(**a))
finally:
AutoTuner.get().clear_cache()
+
+
+# ── the guards ────────────────────────────────────────────────────────────
+
+
+def _strided(t: torch.Tensor) -> torch.Tensor:
+ wide = torch.zeros(*t.shape[:-1], 2 * t.shape[-1], device=DEV, dtype=t.dtype)
+ wide[..., ::2] = t
+ view = wide[..., ::2]
+ assert not view.is_contiguous()
+ return view
+
+
+def test_guard_refuses_a_strided_hidden_states():
+ """What `is_valid` exists for: the raw op reads a strided operand as dense."""
+ a = _args(16, 256, 128, 8, 2, seed=11)
+ with validating(op), pytest.raises(AssertionError, match="must be contiguous"):
+ op(**dict(a, hidden_states=_strided(a["hidden_states"])))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py
index 9fb348eeb829..c10b7b266307 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py
@@ -5,8 +5,9 @@
The entry owns the contract. This file turns a cell into operands (a strided gate
is the call site's `gug` column), locks both sides of the in-place write, and holds
-the claims no cell carries: the strided gate equal to the dense column, and the
-empty batch.
+the claims no cell carries: the strided gate equal to the dense column, CUDA-graph
+capture and replay at the 128-row site, the empty batch and the loud rejection of
+a wrong-length gate.
"""
import zlib
@@ -93,9 +94,42 @@ def test_strided_gate_is_the_dense_column_bit_for_bit() -> None:
assert torch.equal(final, dense_final)
+def test_cuda_graph_capture_and_replay() -> None:
+ """The decode sites run captured: one capture at the 128-row site, replayed on
+ refilled operands (the strided gate's densifying copy is captured with it)."""
+ spec = _strided_spec(128)
+ final, gate, shared = _operands(spec, seed=5)
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ op(final.clone(), gate, shared) # warm-up: the Triton JIT compiles outside capture
+ torch.cuda.current_stream().wait_stream(side)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ out = op(final, gate, shared)
+ assert out.data_ptr() == final.data_ptr()
+
+ for seed in (21, 22):
+ new_final, new_gate, new_shared = _operands(spec, seed=seed)
+ final.copy_(new_final)
+ gate.copy_(new_gate)
+ shared.copy_(new_shared)
+ graph.replay()
+ torch.cuda.synchronize()
+ torch.testing.assert_close(final, op.reference(new_final, new_gate, new_shared))
+
+
def test_zero_tokens_is_a_no_op():
final = torch.empty(0, HIDDEN, dtype=torch.bfloat16, device=DEV)
shared = torch.empty(0, HIDDEN, dtype=torch.bfloat16, device=DEV)
gate = torch.empty(0, 1, dtype=torch.bfloat16, device=DEV)
returned = op(final, gate, shared)
assert returned.shape == (0, HIDDEN) and returned.data_ptr() == final.data_ptr()
+
+
+def test_a_gate_of_the_wrong_length_raises() -> None:
+ """Metadata violations are loud, which is why the entry has no guard."""
+ bf = dict(dtype=torch.bfloat16, device=DEV)
+ final, shared = torch.randn(4, 128, **bf), torch.randn(4, 128, **bf)
+ with pytest.raises(AssertionError):
+ op(final, torch.randn(3, 1, **bf), shared)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_sort.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_sort.py
index f240c17b8240..0bef34594a86 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_sort.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_sort.py
@@ -3,8 +3,8 @@
"""GPU test for the moe_sort entry of the Qwen3.8-Flash-Next target.
The entry owns the contract. This file turns a cell into expert ids and holds
-the claim no cell carries: CUDA-graph replay on new ids at the largest captured
-row count.
+the claims no cell carries: CUDA-graph replay on new ids at the largest captured
+row count, and the contiguity guard.
"""
import zlib
@@ -201,3 +201,15 @@ def test_captured_rows_replay_on_new_ids(certifying): # noqa: F811
assert torch.equal(static_out[i], eager[i]), (trial, i)
del graph, static_out
torch.cuda.synchronize()
+
+
+# ── the guard ─────────────────────────────────────────────────────────────
+
+
+def test_guard_refuses_noncontiguous_ids():
+ """What `is_valid` exists for: the op reads a strided view as the dense buffer
+ that starts at its data pointer, and answers without raising."""
+ wide = _ids("router", 300, 128, seed=51, top_k=16)
+ sc = torch.rand(300, K, device=DEV)
+ with validating(op), pytest.raises(AssertionError, match="contiguous"):
+ op(wide[:, :K], sc, E, K, 0, E, 128)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_unpermute_inplace.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_unpermute_inplace.py
index 723fdf39ca5a..573fe7ed8f49 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_unpermute_inplace.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_unpermute_inplace.py
@@ -6,7 +6,7 @@
CuteDSL FC2 hands them over -- placed by moe_sort, every row the op must not read
filled with NaN, the output pre-filled with NaN -- and holds the claims no cell
carries: CUDA-graph replay at the largest captured row count, bitwise equal to
-eager.
+eager, and the contiguity guard.
"""
import zlib
@@ -240,3 +240,29 @@ def step(ids=ids, scales=scales, rows=rows, perm=perm, out=out, tile=tile):
del graph, e2p
torch.cuda.synchronize()
assert stale_rows > 0, "no replay ran over a stale row: the stale-row claim is untested"
+
+
+# ── the guards ────────────────────────────────────────────────────────────
+
+
+def _strided(t: torch.Tensor, pad: int) -> torch.Tensor:
+ """The same values in a view whose rows are `pad` elements wider apart."""
+ wide = torch.zeros(t.shape[0], t.shape[1] + pad, dtype=t.dtype, device=t.device)
+ wide[:, : t.shape[1]] = t
+ return wide[:, : t.shape[1]]
+
+
+def test_noncontiguous_operand_answers_wrongly_and_is_refused():
+ """The FC2 rows the op reads; the guard checks all four operands the same way."""
+ name = "permuted_input"
+ args = _problem(_spec(300), seed=61)
+ ref = op.reference(**args)
+ bad = dict(args, **{name: _strided(args[name], 64)})
+ out = op(**bad)
+ with pytest.raises(AssertionError, match="differ bitwise"):
+ _compare(out.contiguous(), ref)
+ with validating(op):
+ with pytest.raises(AssertionError, match=f"{name} must be contiguous"):
+ op(**bad)
+ out = op(**dict(args, output=_output(300)))
+ _compare(out, ref)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_renorm_moe_routing_op.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_renorm_moe_routing_op.py
index c98408673a3e..b835a2ad6b6c 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_renorm_moe_routing_op.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_renorm_moe_routing_op.py
@@ -4,7 +4,8 @@
(models/qwen4_exp/ops/moe/renorm_moe_routing_op).
The entry owns the contract. This file turns a cell into logits and holds the
-claim no cell carries: equivalence to HF's `norm_topk_prob` routing.
+claims no cell carries: equivalence to HF's `norm_topk_prob` routing, CUDA-graph
+replay at the 128-row captured batch, and the contiguity guard.
"""
import zlib
@@ -113,3 +114,37 @@ def test_equivalent_to_hf_norm_topk_prob():
atol=1e-7,
)
assert (weights - hf_values).abs().max().item() > 1e-2
+
+
+def test_cuda_graph_replay_at_the_captured_batch():
+ """The trunk's decode router runs inside the engine's CUDA graphs: captured
+ once at the largest captured batch (128 rows), each replay over refilled
+ logits answers bit-exactly what the reference does on them."""
+ tokens, num_experts, topk = 128, 512, 10
+ logits = _logits(tokens, num_experts, torch.bfloat16, seed=21)
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ op(logits, topk, torch.bfloat16)
+ torch.cuda.current_stream().wait_stream(side)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ ids, weights = op(logits, topk, torch.bfloat16)
+ for seed in (22, 23):
+ logits.copy_(_logits(tokens, num_experts, torch.bfloat16, seed=seed))
+ graph.replay()
+ torch.cuda.synchronize()
+ _compare(
+ (ids.clone(), weights.clone()),
+ op.reference(logits, topk, torch.bfloat16),
+ )
+
+
+# ── the guard ─────────────────────────────────────────────────────────────
+
+
+def test_guard_refuses_a_column_slice():
+ """What `is_valid` exists for: the raw op reads a slice as the dense buffer."""
+ wide = _logits(8, 1024, torch.bfloat16, seed=12)
+ with validating(op), pytest.raises(AssertionError, match="contiguous"):
+ op(wide[:, :512], 10)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py
index cb2e79386da1..f801516daa43 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py
@@ -58,3 +58,44 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res
assert out.data_ptr() != x.data_ptr(), "the op must return a new tensor"
assert torch.equal(buffer, buffer_before), "x (and the stream it slices) is read-only"
torch.testing.assert_close(out, op.reference(x, weight, eps))
+
+
+def test_cuda_graph_replay_at_the_decode_stream_slice() -> None:
+ """The decode site is captured: a graph over the 128-row stream slice, replayed
+ after the residual is refilled, answers bitwise what the eager call does."""
+ spec = next(
+ c.spec
+ for c in op.CELLS
+ if c.spec["shape"] == (128, 2560) and c.spec.get("row_width") == 10240
+ )
+ buffer, x, weight, eps = _build(spec, seed=_seed("graph replay"))
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ for _ in range(2):
+ op(x, weight, eps)
+ torch.cuda.current_stream().wait_stream(side)
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ out = op(x, weight, eps)
+
+ g = torch.Generator(device="cuda").manual_seed(_seed("graph refill"))
+ buffer.copy_(torch.randn(buffer.shape, generator=g, device="cuda").to(buffer.dtype))
+ graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(out, op(x, weight, eps)), "replay differs from the eager call"
+ torch.testing.assert_close(out, op.reference(x, weight, eps))
+
+
+def test_guard_refuses_fp8_before_the_kernel_answers() -> None:
+ """fp8 is in the DSL dtype table: accepted, and a tensor comes back."""
+ x = torch.randn(4, 2560, device="cuda").to(torch.float8_e4m3fn)
+ w = torch.randn(2560, device="cuda").to(torch.float8_e4m3fn)
+ unguarded = torch.ops.trtllm.flashinfer_gemma_rmsnorm(x, w, 1e-6)
+ assert unguarded.dtype is x.dtype and unguarded.shape == x.shape, (
+ "the hazard this guard covers is gone"
+ )
+ with pytest.raises(AssertionError, match="dequantize"):
+ with validating(op):
+ op(x, w, 1e-6)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_layer_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_layer_norm.py
index b64080199d2b..8e70742eb225 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_layer_norm.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_layer_norm.py
@@ -5,7 +5,10 @@
Weight and bias are drawn over this checkpoint's own 110 vision LayerNorm
tensors' ranges: the affine is wide enough (weights to 17.5, biases to -25.6)
that a wrong normalization is not masked by a near-unit scale. The default
-bf16 gate against the entry's fp64 reference is the band.
+bf16 gate against the entry's fp64 reference is the band; the negative
+controls below show it refuses the wrong normalizations a LayerNorm is
+usually confused with. Beyond that: the note's 1024-row graph replay and the
+loud weight-length refusal.
"""
import zlib
@@ -61,3 +64,61 @@ def test_certified_cells(cell, certifying): # noqa: F811
assert all(torch.equal(a, b) for a, b in zip(before, after)), "inputs must not be mutated"
assert torch.equal(out, op(x, weight, bias, eps)), "two launches must agree bitwise"
torch.testing.assert_close(out, op.reference(x, weight, bias, eps))
+
+
+def _wrong_reference(x, weight, bias, eps, *, ddof=0, center=True, with_bias=True):
+ """The entry's fp64 formula with one deliberate mistake, rounded to `x.dtype`."""
+ xd = x.double()
+ mean = xd.mean(dim=-1, keepdim=True) if center else torch.zeros_like(xd[..., :1])
+ var = (xd - mean).square().sum(dim=-1, keepdim=True) / (xd.shape[-1] - ddof)
+ out = (xd - mean) / torch.sqrt(var + eps) * weight.double()
+ if with_bias:
+ out = out + bias.double()
+ return out.to(x.dtype)
+
+
+@pytest.mark.parametrize(
+ "mistake",
+ [dict(ddof=1), dict(center=False), dict(with_bias=False)],
+ ids=["unbiased_variance", "no_mean_subtraction", "no_bias"],
+)
+def test_the_gate_discriminates_a_wrong_normalization(mistake):
+ """The cells' gate must refuse a reference that normalizes wrongly.
+
+ The op's output differs from the correctly rounded fp64 value on a tiny
+ fraction of the elements; each of these mistakes moves far more of them
+ and fails the default bf16 `assert_close`.
+ """
+ cell = next(c for c in op.CELLS if c.spec["shape"][0] == 1024)
+ x, weight, bias, eps = _build(cell.spec, seed=zlib.crc32(cell.why.encode()))
+ out = op(x, weight, bias, eps)
+ torch.testing.assert_close(out, op.reference(x, weight, bias, eps))
+ with pytest.raises(AssertionError):
+ torch.testing.assert_close(out, _wrong_reference(x, weight, bias, eps, **mistake))
+
+
+def test_capturable_in_a_cuda_graph():
+ """A 1024-row capture, replayed on refilled input, is bitwise the eager call."""
+ gen = _gen(77)
+ static_x = torch.randn(1024, WIDTH, generator=gen, device=DEVICE).to(torch.bfloat16)
+ weight, bias = _affine(gen)
+ 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_(torch.randn(1024, WIDTH, generator=gen, device=DEVICE).to(torch.bfloat16))
+ graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(static_out, op(static_x, weight, bias, EPS))
+
+
+def test_rejects_a_weight_whose_length_is_not_the_normalized_width():
+ """Loud: `_layer_norm_fwd` asserts the shapes, with no message."""
+ gen = _gen(88)
+ x = torch.randn(8, WIDTH, generator=gen, device=DEVICE).to(torch.bfloat16)
+ weight, bias = _affine(gen)
+ with pytest.raises(AssertionError):
+ op(x, weight[:-1].contiguous(), bias, EPS)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py
index 6db54e25f309..d0d157032a2c 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py
@@ -30,6 +30,10 @@ def _z_offset(row_stride: int, heads: int, n: int) -> int:
return offset
+#: 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:
return zlib.crc32(text.encode())
@@ -76,3 +80,17 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res
for t, b in zip((x, z, weight), before):
assert torch.equal(t, b), "inputs are read-only"
torch.testing.assert_close(out, op.reference(*args))
+
+
+def test_guard_refuses_a_weight_of_the_wrong_length() -> None:
+ """Unguarded, a length-192 weight is silently truncated to weight[:128]."""
+ x, z, weight, eps, _, act = _build(_PACKED, seed=12)
+ long_weight = torch.cat([weight, torch.randn(64, dtype=weight.dtype, device="cuda")])
+ assert torch.equal(
+ torch.ops.trtllm.rms_norm_gated_token_major(x, z, long_weight, eps, None, act),
+ torch.ops.trtllm.rms_norm_gated_token_major(x, z, weight, eps, None, act),
+ ), "the hazard this guard covers is gone"
+
+ with pytest.raises(AssertionError, match="weight must be 1-D of length"):
+ with validating(op):
+ op(x, z, long_weight, eps, None, act)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ngram_token_history.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ngram_token_history.py
index a5e97badb438..345be975f2ff 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ngram_token_history.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ngram_token_history.py
@@ -6,7 +6,8 @@
against the entry's scalar reference, which builds each sequence's stream as
a list and shares no structure with the masked per-shift select under test.
Beyond the cells: chunked prefill against one-shot (the cross-step carry
-the forward relies on).
+the forward relies on), the `query_start_loc` length guard and CUDA-graph
+replay.
"""
import zlib
@@ -18,6 +19,7 @@
CKPT_EOS_TOKEN_ID,
CKPT_NGRAM_SIZE,
CKPT_VOCAB_SIZE,
+ ENGINE_SLOTS,
)
from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ngram_token_history import ( # noqa: E501
ngram_token_history as op,
@@ -132,3 +134,79 @@ def call(pool, chunk, carry):
split = torch.cat([call(pool, ids[:41], False), call(pool, ids[41:], True)])
assert torch.equal(split, one_shot)
assert torch.equal(pool, one_shot_pool)
+
+
+# ----------------------------------------------------------- guards
+
+
+def _two_sequences():
+ pool = torch.randint(0, CKPT_VOCAB_SIZE, (4, CONTEXT_LEN), device=DEV)
+ return (
+ torch.tensor([11, 12, 13, 21], device=DEV),
+ pool,
+ torch.tensor([0, 1], dtype=torch.int32, device=DEV),
+ torch.tensor([True, False], device=DEV),
+ torch.tensor([0, 3, 4], dtype=torch.int32, device=DEV),
+ EOS,
+ NGRAM,
+ )
+
+
+def _misses(args, bad_index, bad_value):
+ """Run the op unguarded with one argument swapped; True if it returned
+ and its answer differs from the reference on the intended arguments."""
+ expected = _expect(args)
+ bad = list(args)
+ bad[1] = args[1].clone()
+ bad[bad_index] = bad_value
+ window = op(*bad)
+ return not (torch.equal(window, expected[0]) and torch.equal(bad[1], expected[1]))
+
+
+def test_guard_refuses_a_short_query_start_loc() -> None:
+ """Unguarded, a 2-entry vector for 2 sequences misattributes the tail."""
+ args = _two_sequences()
+ short = torch.tensor([0, 3], dtype=torch.int32, device=DEV)
+ assert _misses(args, 4, short), "the hazard this guard covers is gone"
+ bad = list(args)
+ bad[4] = short
+ with pytest.raises(AssertionError, match=r"query_start_loc must be \[num_seqs"):
+ with validating(op):
+ op(*bad)
+
+
+# ----------------------------------------------------------- capture
+
+
+def test_cuda_graph_capture_and_replay() -> None:
+ """The captured decode shape, 8 sequences at the engine pool width."""
+ spec = dict(
+ lengths=(1,) * 8,
+ slots=(0, ENGINE_SLOTS - 1, 4, 17, 30, 43, 56, 69),
+ carries=(True,) * 8,
+ pool_slots=ENGINE_SLOTS,
+ ids_dtype=torch.int32,
+ qsl_dtype=torch.int64,
+ )
+ args = _build(spec, seed=11001)
+ start = args[1].clone()
+
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ for _ in range(3):
+ op(*args)
+ torch.cuda.current_stream().wait_stream(side)
+
+ args[1].copy_(start)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ captured = op(*args)
+
+ replay_ids = torch.randint(0, CKPT_VOCAB_SIZE, (8,), dtype=torch.int32, device=DEV)
+ args[0].copy_(replay_ids)
+ args[1].copy_(start)
+ expected = op.reference(replay_ids, start.clone(), *args[2:])
+ graph.replay()
+ torch.cuda.synchronize()
+ _compare((captured, args[1]), expected)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_gate_conv_inject_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_gate_conv_inject_norm.py
index 4581ad5e52db..00ec41d65016 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_gate_conv_inject_norm.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_gate_conv_inject_norm.py
@@ -6,8 +6,9 @@
the stream (`hc_out`) and the normed stream in units of their row's scale, the
published conv state per element (see `_compare`); a verify cell's staged conv
input is checked bitwise against the state it published. Beyond the cells: the
-staged conv input is what the conv consumed, and `out` writes only its rows of a
-larger buffer.
+staged conv input is what the conv consumed, `out` writes only its rows of a
+larger buffer, the operand-width guard, and graph capture at the 461-slot pool and
+over a rank's whole pool (the tep4 decode and MTP7 verify graphs).
"""
import zlib
@@ -20,6 +21,8 @@
CKPT_DILATION,
CKPT_HC_COUNT,
CKPT_HIDDEN,
+ ENGINE_SLOTS,
+ TEP4_SLOTS,
)
from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ple_gate_conv_inject_norm import ( # noqa: E501
ple_gate_conv_inject_norm as op,
@@ -61,17 +64,17 @@ def _stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]:
#: (`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).
+#: 2 ulp of the row maximum per element and 0.1 ulp relative RMS. The same
+#: steps unfused (one bf16 kernel per step) sit 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 unfused steps 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.
+#: reference over three seeds at 16384 rows, the unfused steps reach 2.91
+#: row-max ulp (normed) and the op 2.94; at 8192 rows 1.89 / 1.89.
#: Cells above 1024 rows are therefore gated at 4; the RMS gate is unchanged.
#: The default target's per-rank pools (1874 / 850 slots) sit inside the same gates; a
#: reference reading each sequence's neighbouring slot falls far outside the RMS gate
@@ -197,3 +200,51 @@ def test_out_writes_the_rows_of_larger_buffers() -> None:
assert torch.equal(bufs[0][:rows], plain[0]) and torch.equal(bufs[1][:rows], plain[1])
assert torch.equal(args[7], plain[2])
assert bool((bufs[0][rows:] == 7.0).all()) and bool((bufs[1][rows:] == 7.0).all())
+
+
+def test_guard_refuses_a_short_value() -> None:
+ args = _build(dict(lengths=(1, 1)), seed=4)
+ args[2] = args[2][:, :-1]
+ with validating(op), pytest.raises(AssertionError):
+ op(*args)
+
+
+@pytest.mark.parametrize(
+ "spec",
+ [
+ dict(lengths=(1,) * 8, pool_slots=ENGINE_SLOTS),
+ # The tep4 decode graph and the MTP7 verify graph (packed-row conv, staged
+ # conv input) over a rank's whole pool.
+ dict(lengths=(1,) * 128, pool_slots=TEP4_SLOTS),
+ dict(lengths=(8,) * 128, pool_slots=TEP4_SLOTS, conv_input=True),
+ ],
+ ids=lambda s: f"{len(s['lengths'])}x{s['lengths'][0]}-{s['pool_slots']}slots",
+)
+def test_cuda_graph_capture_and_replay(spec) -> None:
+ args = _build(spec, seed=8)
+ xc = torch.empty_like(args[0]) if spec.get("conv_input") else None
+ 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, conv_input=xc)
+ 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, conv_input=xc)
+ fresh = _build(spec, seed=18)
+ args[0].copy_(fresh[0])
+ args[1].copy_(fresh[1])
+ args[2].copy_(fresh[2])
+ args[7].copy_(pool)
+ eager_xc = None if xc is None else torch.empty_like(xc)
+ eager = _run(args, conv_input=eager_xc)
+ 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])
+ if xc is not None:
+ assert torch.equal(xc, eager_xc)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/quantization/test_modeling_v2_bf16_to_fp32.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/quantization/test_modeling_v2_bf16_to_fp32.py
index 37cca3ab7333..9d7d7cbb9d22 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/quantization/test_modeling_v2_bf16_to_fp32.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/quantization/test_modeling_v2_bf16_to_fp32.py
@@ -1,6 +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()`."""
+"""GPU test for the bf16_to_fp32 entry: the cells bit for bit `x.float()`, graph capture, the
+contiguity guard."""
import zlib
@@ -42,3 +43,25 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by
ref = op.reference(x)
assert out.dtype == torch.float32 and out.is_contiguous()
assert _bitwise(out, ref), "widening is exact; a bit difference is a defect"
+
+
+def test_cuda_graph_capture_and_replay() -> None:
+ x = _build(dict(shape=(128, 4099)), 1)
+ warmup = torch.cuda.Stream()
+ warmup.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(warmup):
+ op(x)
+ torch.cuda.current_stream().wait_stream(warmup)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ out = op(x)
+ x.copy_(_build(dict(shape=(128, 4099)), 2))
+ graph.replay()
+ torch.cuda.synchronize()
+ assert _bitwise(out, x.float())
+
+
+def test_guard_refuses_a_strided_input() -> None:
+ x = _build(dict(shape=(8, 64)), 3)
+ with validating(op), pytest.raises(AssertionError):
+ op(x[:, ::2])
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_vision_rope_2d_packed_qk.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_vision_rope_2d_packed_qk.py
index f47dc0ec1c37..98908916f232 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_vision_rope_2d_packed_qk.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_vision_rope_2d_packed_qk.py
@@ -7,9 +7,9 @@
table from the cell's grid (spatial-merge positions, fp64 angles rounded once to fp32,
fp32 cos / sin), rotates `qkv[:, :32]` in place, and gates:
1. bitwise against the entry's `reference` on the same table;
- 2. bitwise against the whole-op reference of `vision_rope_2d__r52` (copied below:
- positions by nested loops, fp64 angles, `polar`, complex multiply), q and k
- separately, with that entry's distribution gates beside it;
+ 2. bitwise against a whole-op reference from the grid (`_grid_reference`: positions by
+ nested loops, fp64 angles, `polar`, complex multiply), q and k separately, with
+ distribution gates beside it;
3. the v heads of every row untouched.
"""
@@ -66,9 +66,8 @@ def _angle_table(grid_thw: torch.Tensor, merge: int, theta: float, head_dim: int
return angles.cos().contiguous(), angles.sin().contiguous()
-def _r52_reference(q, k, grid_thw, spatial_merge_size, theta=THETA):
- """`vision_rope_2d__r52.reference`, verbatim: the whole op from the grid, as a fresh
- `(q_out, k_out)`."""
+def _grid_reference(q, k, grid_thw, spatial_merge_size, theta=THETA):
+ """The whole op from the grid, angle table included, as a fresh `(q_out, k_out)`."""
rope_dim = q.shape[-1] // 2
axis_dim = rope_dim // 2
positions = _spatial_merge_positions(grid_thw.tolist(), spatial_merge_size)
@@ -110,7 +109,7 @@ def _distribution(got: torch.Tensor, want: torch.Tensor) -> Tuple[float, float,
return cos, rel_l2, torch.quantile(err, 0.999).item()
-#: `vision_rope_2d__r52`'s distribution gates: both sides round the same fp32 rotation to bf16
+#: The distribution gates against `_grid_reference`: both sides round the same fp32 rotation to bf16
#: once, so they differ by at most one bf16 ulp per element, at most 2^-7 of its value: the
#: relative L2 is bounded by that, and 1 - cosine by half its square.
_REL_L2_BOUND = 2.0**-7
@@ -148,23 +147,25 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th
print(f"[{cell.why}] tokens={tokens} vs reference: {differing} differing, max {max_dev:.3e}")
assert differing == 0, f"{differing} elements differ from reference (max {max_dev:.3e})"
- # 2. vision_rope_2d__r52's whole-op reference from the grid: bitwise, plus its
+ # 2. The whole-op reference from the grid: bitwise, plus its
# distribution gates.
- q_r52, k_r52 = _r52_reference(
+ q_grid, k_grid = _grid_reference(
before[:, :NUM_HEADS], before[:, NUM_HEADS : 2 * NUM_HEADS], grid_thw, MERGE, THETA
)
for name, got, ref in (
- ("q", qkv[:, :NUM_HEADS], q_r52),
- ("k", qkv[:, NUM_HEADS : 2 * NUM_HEADS], k_r52),
+ ("q", qkv[:, :NUM_HEADS], q_grid),
+ ("k", qkv[:, NUM_HEADS : 2 * NUM_HEADS], k_grid),
):
differing, max_dev = _mismatch(got, ref)
cos_sim, rel_l2, p999 = _distribution(got, ref)
print(
- f"[{cell.why}] {name} vs r52: {differing} differing, max {max_dev:.3e}, "
+ f"[{cell.why}] {name} vs the grid reference: {differing} differing, max {max_dev:.3e}, "
f"cos={cos_sim:.9f} rel_l2={rel_l2:.3e} p99.9={p999:.3e}"
)
assert cos_sim >= _COS_BOUND and rel_l2 <= _REL_L2_BOUND, (name, cos_sim, rel_l2)
- assert differing == 0, f"{name}: {differing} elements differ from r52 (max {max_dev:.3e})"
+ assert differing == 0, (
+ f"{name}: {differing} elements differ from the grid reference (max {max_dev:.3e})"
+ )
def _small():
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first__decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first__decode.py
index 4550b39ed729..466f910ea2ff 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first__decode.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first__decode.py
@@ -2,7 +2,8 @@
# SPDX-License-Identifier: Apache-2.0
"""GPU test for this target's qwen4_exp_hc_decode_tc_first variant
(ops/hc/qwen4_exp_hc_decode_tc_first__decode): cells in the band against the float64
-reference at the decode rows the members run."""
+reference at the decode rows the members run, the row-count guard, graph replay at the
+note's 1, 4, 16 and 256 rows (every tile variant)."""
import zlib
@@ -67,3 +68,28 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by
low, inject = op(*args)
assert low.shape == (cell.spec["rows"], RANK)
_compare((low, inject), op.reference(*args))
+
+
+def test_guard_refuses_more_than_two_tiles_of_rows() -> None:
+ normed, first, rank, c, wi = _build(dict(rows=257), seed=4)
+ with validating(op), pytest.raises(AssertionError, match="at most 256 rows"):
+ op(normed, first, rank, c, wi)
+
+
+@pytest.mark.parametrize("rows", [1, 4, 16, 256])
+def test_cuda_graph_capture_and_replay(rows) -> None:
+ normed, first, rank, c, wi = _build(dict(rows=rows), seed=5)
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ op(normed, first, rank, c, wi)
+ torch.cuda.current_stream().wait_stream(side)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ g_low, g_inject = op(normed, first, rank, c, wi)
+ fresh = _build(dict(rows=rows), seed=6)[0]
+ normed.copy_(fresh)
+ graph.replay()
+ torch.cuda.synchronize()
+ low, inject = op(fresh, first, rank, c, wi)
+ assert torch.equal(g_low, low) and torch.equal(g_inject, inject)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ssm/test_modeling_v2_gdn_cached_replay_cute__verify.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ssm/test_modeling_v2_gdn_cached_replay_cute__verify.py
index fabebf6190ed..98b9e62c5284 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ssm/test_modeling_v2_gdn_cached_replay_cute__verify.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ssm/test_modeling_v2_gdn_cached_replay_cute__verify.py
@@ -43,8 +43,6 @@
op = cr.gdn_cached_replay_cute
DEV = "cuda"
K, V = cr.K, cr.V
-#: The input projection's column split at a tp4/tep4 rank: conv (q|k|v), z, a, b.
-_CONV, _Z = 2 * cr.RANK_K_HEADS * K + cr.RANK_V_HEADS * V, cr.RANK_V_HEADS * V
def _gates(bench, act, strided):
@@ -52,10 +50,12 @@ def _gates(bench, act, strided):
column views of a `[N*T, conv + z + 2 HV]` projection (row stride = its width)."""
if not strided:
return act["a"], act["b"]
- n, t, hv = bench.n, bench.t, bench.hv
- proj = torch.randn(n * t, _CONV + _Z + 2 * hv, device=DEV, dtype=torch.bfloat16)
- a = proj[:, _CONV + _Z : _CONV + _Z + hv]
- b = proj[:, _CONV + _Z + hv :]
+ n, t, h, hv = bench.n, bench.t, bench.h, bench.hv
+ # The input projection's column split at the rank's heads: conv (q|k|v), z, a, b.
+ conv, z = 2 * h * K + hv * V, hv * V
+ proj = torch.randn(n * t, conv + z + 2 * hv, device=DEV, dtype=torch.bfloat16)
+ a = proj[:, conv + z : conv + z + hv]
+ b = proj[:, conv + z + hv :]
a.copy_(act["a"].reshape(n * t, hv))
b.copy_(act["b"].reshape(n * t, hv))
a, b = a.reshape(n, t, hv), b.reshape(n, t, hv)
@@ -185,19 +185,20 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by
)
-def _rank_spec(**spec):
- return dict(state_dtype=torch.bfloat16, H=cr.RANK_K_HEADS, HV=cr.RANK_V_HEADS, **spec)
+def _rank_spec(heads=(cr.RANK_K_HEADS, cr.RANK_V_HEADS), **spec):
+ return dict(state_dtype=torch.bfloat16, H=heads[0], HV=heads[1], **spec)
+@pytest.mark.parametrize("heads", cr.GEOMETRIES, ids=[f"{h}x{hv}" for h, hv in cr.GEOMETRIES])
@pytest.mark.parametrize("pdl", [False, True])
@pytest.mark.parametrize("n,path", [(1, "small"), (4, "small"), (16, "all_layer")])
-def test_cuda_graph_replay_on_new_data_is_bitwise_eager(n, path, pdl) -> None:
+def test_cuda_graph_replay_on_new_data_is_bitwise_eager(n, path, pdl, heads) -> None:
"""Two layers captured once (strided a / b, the all-layer commit after them when
partitioned), replayed three times on new activations with acceptance in between,
bitwise against the same steps run eagerly on a twin state; with `pdl` the replay
launches with programmatic dependent launch, as the target's verify step does."""
t = 4
- spec = _rank_spec(path=path, n=n, t=t, layers=2, slots=40)
+ spec = _rank_spec(heads, path=path, n=n, t=t, layers=2, slots=40)
seed = 9100 + n
eager, graph, warm = (trep._Bench(spec, seed, exact=False) for _ in range(3))
acts = [warm.inputs() for _ in range(2)]
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_attention_plumbing.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_attention_plumbing.py
index 2daa23f1e54d..5dde8f4ddd7d 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_attention_plumbing.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_attention_plumbing.py
@@ -18,6 +18,9 @@
mixed-batch forms), and the reuse entry's order over canonical captures. The targets are
built without weights: the capture methods read only their arguments, the metadata and the
core's draft binding.
+* **Score tables past the int32 range.** A part whose rows times the table's columns would
+ reach 2^31 is scored and selected in chunks under it, each row once, with its part's
+ context flag; at the checkpoint's `max_seq_len` (65568 columns) 32752 rows are one call.
"""
import math
@@ -384,3 +387,68 @@ def test_reuse_of_a_canonical_capture_stays_ascending():
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))
+
+
+# ---- score tables past the int32 range ----------------------------------------------------
+
+
+def _chunked_calls(monkeypatch, rows: int, pages: int, parts, limit=None):
+ """`_score_and_select` over `rows` rows and a `pages`-page table (32 tokens, compression
+ 4), with the scoring and the selection recorded: `(first row, rows, context_rows)` per
+ scoring call, and `selected` with each row's position in its first column."""
+ columns = pages * 32 // 4
+ calls = []
+
+ def scores_op(*, q, query_positions, context_rows, **_):
+ calls.append((int(query_positions[0]), q.shape[0], context_rows))
+ assert q.shape[0] * columns < (limit or M._SCORE_TABLE_ELEMENTS)
+ return q[:, 0, :1].expand(q.shape[0], columns)
+
+ def selection(core, scores, visible, logical, seq_lens, topk, starts, out, live_cols):
+ assert scores.shape == (logical.shape[0], columns) and live_cols == 7
+ out.copy_(logical[:, None].to(out.dtype).expand_as(out))
+
+ monkeypatch.setattr(M, "qsa_paged_index_scores", scores_op)
+ monkeypatch.setattr(M, "_selection", selection)
+ if limit is not None:
+ monkeypatch.setattr(M, "_SCORE_TABLE_ELEMENTS", limit)
+ logical = torch.arange(rows)
+ selected = torch.full((rows, 3), -1, dtype=torch.int32)
+ table = torch.zeros(2, pages, dtype=torch.int32)
+ rows_q = logical.float()[:, None, None]
+ M._score_and_select(
+ SimpleNamespace(compress_ratio=4),
+ None,
+ rows_q,
+ table,
+ logical,
+ logical,
+ 32,
+ parts,
+ logical,
+ logical,
+ logical,
+ logical,
+ selected,
+ 7,
+ )
+ assert torch.equal(selected[:, 0], logical.to(torch.int32)), "every row selected once"
+ return calls
+
+
+@pytest.mark.cpu_only
+def test_a_part_past_the_score_table_bound_goes_in_chunks(monkeypatch) -> None:
+ # 4 pages: 32 columns; a bound of 5 x 32 + 1 elements takes 5 rows a call.
+ calls = _chunked_calls(monkeypatch, 13, 4, [(0, 11, True), (11, 13, False)], 5 * 32 + 1)
+ assert calls == [(0, 5, True), (5, 5, True), (10, 1, True), (11, 2, False)]
+
+
+@pytest.mark.cpu_only
+def test_at_max_seq_len_262144_a_call_takes_32752_rows(monkeypatch) -> None:
+ # 8196 pages (262144 tokens plus the generation lead, padded to 4 pages): 65568 columns.
+ assert 32752 * 65568 < 2**31 <= 32753 * 65568
+ assert _chunked_calls(monkeypatch, 8192, 8196, [(0, 8192, True)]) == [(0, 8192, True)]
+ assert _chunked_calls(monkeypatch, 40000, 8196, [(0, 40000, True)]) == [
+ (0, 32752, True),
+ (32752, 7248, True),
+ ]
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_core_layout.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_core_layout.py
index 57b3e5f537e5..b7edbfb9c3d7 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_core_layout.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_core_layout.py
@@ -28,6 +28,7 @@
)
from tensorrt_llm._torch.configs import Qwen4ExpTextConfig, Qwen4ExpVisionConfig
from tensorrt_llm._torch.model_config import ModelConfig
+from tensorrt_llm.functional import AllReduceStrategy
from tensorrt_llm.llmapi import QSASparseAttentionConfig
from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig
from tensorrt_llm.mapping import Mapping
@@ -128,8 +129,9 @@
_TP4 = dict(enable_attention_dp=False, moe_tp_size=4, moe_ep_size=1)
-def _core(rank: int, layout: dict, spec_config=None, world: int = WORLD):
- """The core of rank `rank` of a `world`-rank engine at `layout`, every parameter on meta."""
+def _core(rank: int, layout: dict, spec_config=None, world: int = WORLD, **config):
+ """The core of rank `rank` of a `world`-rank engine at `layout`, every parameter on meta;
+ `config` sets further `ModelConfig` fields."""
model_config = ModelConfig(
pretrained_config=Qwen4ExpTextConfig(**_TEXT),
mapping=Mapping(world_size=world, tp_size=world, rank=rank, **layout),
@@ -138,6 +140,7 @@ def _core(rank: int, layout: dict, spec_config=None, world: int = WORLD):
moe_backend_requested="TRTLLM",
sparse_attention_config=QSASparseAttentionConfig(),
spec_config=spec_config,
+ **config,
)
composite = SimpleNamespace(
vision_config=Qwen4ExpVisionConfig(**_VISION), image_token_id=248056, video_token_id=248057
@@ -155,7 +158,7 @@ def test_a_tep4_rank_runs_its_share_of_the_heads(rank) -> None:
"""Six q heads and one of the two KV heads (repeated on rank pairs) over a 4352-wide
fused row: `[q 1536 | k 256 | v 256 | gate 1536 | index q 512, k 128 | v slot 128]`.
The main norm + rope reads it as 6 q, 1 k and 10 "v" heads of 256 (the row stride
- alone), which is the cell `fused_qk_norm_rope` certifies."""
+ alone)."""
core = _core(rank, _TEP4)
assert (core.heads_q, core.heads_kv, core.kv_head_start) == (6, 1, rank * 2 // WORLD)
assert (core.q_width, core.kv_width, core.qkv_width) == (1536, 256, 2048)
@@ -767,3 +770,42 @@ def test_both_draft_targets_pad_the_draft_moe_by_the_workers_counts(target) -> N
(tail,) = _calls(fn, "_mtp_tail")
assert ast.unparse(tail.args[-1]) == "pad_rows"
assert "_dp_pad_rows" not in ast.unparse(fn)
+
+
+# ---- engine knobs the target reads ---------------------------------------------------------
+
+
+@pytest.mark.parametrize("strategy", [AllReduceStrategy.AUTO, AllReduceStrategy.NCCL])
+def test_the_tp_sum_runs_at_the_engines_allreduce_strategy(strategy, monkeypatch) -> None:
+ core = _core(0, _TEP4, allreduce_strategy=strategy)
+ seen = []
+ monkeypatch.setattr(M, "allreduce_group_args", lambda mapping, s: seen.append(s) or {})
+ core._tp_sum_op()
+ assert seen == [int(strategy)]
+
+
+@pytest.mark.parametrize("value", ["1", "true", "On", "0", "no", ""])
+def test_a_switch_takes_the_usual_spellings(value, monkeypatch) -> None:
+ monkeypatch.setenv(M._DENSE_FP8_ENV, value)
+ assert M._switch(M._DENSE_FP8_ENV) is (value.lower() in ("1", "true", "on"))
+ monkeypatch.setenv(M._PLE_HOST_OFFLOAD_ENV, value)
+ assert M._ple_host_offload_enabled() is (value.lower() not in ("0", "no"))
+
+
+@pytest.mark.parametrize(
+ "env", [M._DENSE_FP8_ENV, M._PLE_HOST_OFFLOAD_ENV, M._MOE_PREFILL_COPY_ENV]
+)
+def test_a_misspelled_switch_refuses_the_build(env, monkeypatch) -> None:
+ monkeypatch.setenv(env, "disable")
+ with pytest.raises(ValueError, match=env):
+ _core(0, _TEP4)
+
+
+def test_a_tp1_switch_away_from_tp1_is_reported(monkeypatch) -> None:
+ monkeypatch.setenv(M._DENSE_FP8_ENV, "1")
+ monkeypatch.setenv(M._MTP_HEAD_NVFP4_ENV, "1")
+ warning = mock.Mock()
+ monkeypatch.setattr(M.logger, "warning", warning)
+ _core(0, _TEP4)
+ (message,) = [c.args[0] for c in warning.call_args_list if "certified at tp1" in c.args[0]]
+ assert M._DENSE_FP8_ENV in message and M._MTP_HEAD_NVFP4_ENV in message
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_modeling_v2_spec_rollback.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_modeling_v2_spec_rollback.py
index 1936894dbafa..ff2d83ad0721 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_modeling_v2_spec_rollback.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_modeling_v2_spec_rollback.py
@@ -213,7 +213,7 @@ def _verify_step(b, rollback, fused):
rollback.stage_ple_conv_input(b.x)
_conv_state(b.conv, b.x, b.lengths, b.slots, [False] * b.num_contexts + [True] * b.n)
for layer in range(LAYERS):
- rollback.stage_index(b.index[layer], b.position)
+ rollback.stage_index(b.position)
_write_planes(
b,
b.index[layer],
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_mtp_draft_hc.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_mtp_draft_hc.py
index 79aa96a99447..e85dea0869e5 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_mtp_draft_hc.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_mtp_draft_hc.py
@@ -1,8 +1,7 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
-"""The MTP draft layer: its hyper-connection sites on the trunk's fused entries, against the
-HuggingFace math in float64 and against the composition they replace, and its token embedding
-under rejection sampling.
+"""The MTP draft layer's hyper-connection sites on the trunk's fused entries, against the
+HuggingFace math in float64 and against the unfused composition of the same math.
`_mtp_front`, `_mtp_tail` and `MTPLayer.shared_head` run over a stand-in core that holds the
draft layer's hyper-connection weights (synthetic, at the checkpoint's element statistics): the
@@ -16,7 +15,7 @@
The float64 reference rounds to bf16 where the runtime's activations are bf16 between modules
(the norm outputs, the projections, the residual) and once at each mixer output; the mixers of
the fused path keep their logits in fp32 (the trunk's numerics). Each output must be within
-the band of that reference that the replaced composition (per-stream `flashinfer_gemma_rmsnorm`
+the band of that reference that the unfused composition (per-stream `flashinfer_gemma_rmsnorm`
and `qwen4_exp_gated_residual`) reaches, widened by half.
"""
@@ -43,7 +42,6 @@
from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501
modeling as MOD,
)
-from tensorrt_llm._torch.speculative.interface import rejection_sampling_one_model
pytestmark = pytest.mark.skipif(
not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3),
@@ -200,7 +198,7 @@ def _grouped(x, gamma):
)
-def _mixer_old(raw, tag, normed, inject=True):
+def _mixer_unfused(raw, tag, normed, inject=True):
return qwen4_exp_gated_residual(
normed,
raw[f"{P}{tag}_down"],
@@ -210,7 +208,7 @@ def _mixer_old(raw, tag, normed, inject=True):
)
-def _front_old(raw, ids, hs):
+def _front_unfused(raw, ids, hs):
rows = hs.shape[0]
e = torch.nn.functional.embedding(ids, raw["embed"])
x_e = mtp_cublas_mm(
@@ -219,17 +217,17 @@ def _front_old(raw, ids, hs):
h_n = _grouped(hs, raw[f"{P}pre_h_norm"])
x_h = mtp_cublas_mm(h_n.reshape(rows * HC, HIDDEN), raw[f"{P}fc_h"].t()).reshape(rows, WIDTH)
x = torch.add(x_h, torch.cat([x_e] * HC, -1))
- return (x, *_mixer_old(raw, "ahc", _grouped(x, raw[f"{P}ahc_norm"])))
+ return (x, *_mixer_unfused(raw, "ahc", _grouped(x, raw[f"{P}ahc_norm"])))
-def _tail_old(raw, x, attn, inject):
+def _tail_unfused(raw, x, attn, inject):
x = qwen4_exp_post_inject(x, mtp_cublas_mm(attn, raw[f"{P}o_w"].t()), inject)
- mixed, inject = _mixer_old(raw, "mhc", _grouped(x, raw[f"{P}mhc_norm"]))
+ mixed, inject = _mixer_unfused(raw, "mhc", _grouped(x, raw[f"{P}mhc_norm"]))
return qwen4_exp_post_inject(x, mixed, inject)
-def _head_old(raw, hs):
- return _mixer_old(raw, "fhc", _grouped(hs, raw[f"{P}fhc_norm"]), inject=False)[0]
+def _head_unfused(raw, hs):
+ return _mixer_unfused(raw, "fhc", _grouped(hs, raw[f"{P}fhc_norm"]), inject=False)[0]
def _err(got: torch.Tensor, ref: torch.Tensor) -> float:
@@ -244,16 +242,16 @@ def _rel_l2(got: torch.Tensor, ref: torch.Tensor) -> float:
return ((got.double() - ref64).norm() / ref64.norm()).item()
-def _within(name, got, ref, old) -> None:
- """`got` as close to `ref` as the replaced composition `old` is, widened by half (and
+def _within(name, got, ref, unfused) -> None:
+ """`got` as close to `ref` as the unfused composition `unfused` is, widened by half (and
never tighter than one eps of max|ref| / 1e-3 relative L2, both below the noise of one
bf16 rounding of a mixer output)."""
assert got.shape == ref.shape and got.dtype == ref.dtype, (name, got.shape, ref.shape)
assert bool(torch.isfinite(got).all()), name
- e_new, e_old = _err(got, ref), _err(old, ref)
- l_new, l_old = _rel_l2(got, ref), _rel_l2(old, ref)
+ e_new, e_old = _err(got, ref), _err(unfused, ref)
+ l_new, l_old = _rel_l2(got, ref), _rel_l2(unfused, ref)
print(
- f"{name}: max {e_new:.3f} (old {e_old:.3f}) eps*max|ref|, relL2 {l_new:.3g} (old {l_old:.3g})"
+ f"{name}: max {e_new:.3f} (unfused {e_old:.3f}) eps*max|ref|, relL2 {l_new:.3g} (unfused {l_old:.3g})"
)
assert e_new <= max(1.5 * e_old, 1.0), (name, e_new, e_old)
assert l_new <= max(1.5 * l_old, 1e-3), (name, l_new, l_old)
@@ -289,8 +287,8 @@ def test_the_front_site(rows, strided) -> None:
core, M = _core(raw)
ids, hs, _, _ = _inputs(rows, 101 + rows, strided)
x, mixed, inject = MOD._mtp_front(core, M, raw["embed"], ids, hs)
- ref, old = _front64(raw, ids, hs), _front_old(raw, ids, hs)
- for name, got, r, o in zip(("x", "mixed", "inject"), (x, mixed, inject), ref, old):
+ ref, unfused = _front64(raw, ids, hs), _front_unfused(raw, ids, hs)
+ for name, got, r, o in zip(("x", "mixed", "inject"), (x, mixed, inject), ref, unfused):
_within(f"front[{rows}].{name}", got, r, o)
# The unit injection is the bf16 add: the residual is bitwise the replaced one's, whose
# projections are the same calls.
@@ -307,7 +305,9 @@ def test_the_tail_sites(rows, moe_identity) -> None:
core, M = _core(raw)
_, x, attn, inject = _inputs(rows, 201 + rows)
out = MOD._mtp_tail(core, M, x, attn, inject, 0)
- _within(f"tail[{rows}]", out, _tail64(raw, x, attn, inject), _tail_old(raw, x, attn, inject))
+ _within(
+ f"tail[{rows}]", out, _tail64(raw, x, attn, inject), _tail_unfused(raw, x, attn, inject)
+ )
# 32: a shared head over one attention-DP rank's decode batch (ht_dep4 drafts under MTP).
@@ -320,7 +320,7 @@ def test_the_head_site(rows) -> None:
core=core, logits_processor=SimpleNamespace(forward=lambda final, *args: final)
)
final = MOD.MTPLayer.shared_head(head, hs, None, None, False)
- _within(f"head[{rows}]", final, _head64(raw, hs), _head_old(raw, hs))
+ _within(f"head[{rows}]", final, _head64(raw, hs), _head_unfused(raw, hs))
def test_a_captured_draft_step_replays_the_eager_bits(moe_identity) -> None:
@@ -351,32 +351,3 @@ def step():
torch.cuda.synchronize()
for got, want in zip(captured, eager):
assert torch.equal(got, want)
-
-
-def test_draft_embedding_reads_row_0_for_rejected_slots() -> None:
- """Rejection sampling pads the slots past a request's accepted count with -1 (here every
- draft is rejected at slot 0); the front end embeds those as row 0 and every valid id as
- the plain lookup does."""
- batch, draft, vocab, hidden = 8, 3, 4096, 256
- draft_tokens = torch.arange(batch * draft, dtype=torch.int32, device="cuda").reshape(
- batch, draft
- )
- draft_probs = torch.zeros(batch, draft, vocab, device="cuda")
- draft_probs.scatter_(2, draft_tokens.long().unsqueeze(-1), 1.0)
- target_probs = torch.zeros(batch, draft + 1, vocab, device="cuda")
- target_probs[..., vocab - 1] = 1.0
- seed = torch.zeros(batch, dtype=torch.int64, device="cuda")
- offset = torch.arange(batch, dtype=torch.int64, device="cuda")
- accepted, _ = rejection_sampling_one_model(
- draft_probs, draft_tokens, target_probs, seed, offset
- )
- assert bool((accepted[:, 1:] == -1).all()) and bool((accepted[:, 0] == vocab - 1).all())
-
- g = torch.Generator(device="cuda").manual_seed(0)
- table = torch.randn(vocab, hidden, generator=g, device="cuda").to(torch.bfloat16)
- valid_ids = torch.randint(0, vocab, (257,), generator=g, device="cuda", dtype=torch.int32)
- for ids in (accepted.flatten(), valid_ids, valid_ids.long()):
- e = MOD._draft_embedding(ids, table)
- valid = ids >= 0
- assert torch.equal(e[valid], table[ids[valid].long()])
- assert torch.equal(e[~valid], table[0].expand(int((~valid).sum()), -1))
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py
index 9bb9303d1cd1..f46db3ed01ea 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py
@@ -23,6 +23,7 @@
from tensorrt_llm._torch.model_config import ModelConfig
from tensorrt_llm.llmapi.llm_args import (
KvCacheConfig,
+ MoeConfig,
MTPDecodingConfig,
NGramDecodingConfig,
TorchLlmArgs,
@@ -385,19 +386,22 @@ def test_the_piecewise_prefill_cuda_graph_declines():
)
-@pytest.mark.parametrize("mapping_kwargs", [{}, _LAYOUTS["dep4"]], ids=["tp1", "dep4"])
-def test_a_prefill_step_past_the_int32_score_table_declines(mapping_kwargs):
- """At the checkpoint's max_seq_len 262144 the score table is 65568 columns (8196 pages of
- 32 tokens at compression 4): 32752 context rows stay under 2^31 elements, 32768 do not."""
- config = _model_config(
- pretrained_config=_qwen4_exp_config(
- {"max_position_embeddings": 262144, "indexer_compress_ratio": 4}
- ),
- **mapping_kwargs,
+@pytest.mark.parametrize("member", ["tep4", "dep4", "dep2"])
+def test_eplb_declines_at_the_expert_parallel_layouts(member):
+ """The engine builds an EPLB balancer only with expert parallelism; the target's windows
+ are fixed, so `auto` keeps the built-in model, which balances."""
+ mapping = _DEP2 if member == "dep2" else _LAYOUTS[member]
+ eplb = MoeConfig(load_balancer={"num_slots": 512})
+ _declined(_model_config(**mapping), match=r"load_balancer \(EPLB\)", moe_config=eplb)
+ _claimed(_model_config(**mapping))
+
+
+@pytest.mark.parametrize("mapping_kwargs", [{}, _LAYOUTS["tp4"]], ids=["tp1", "tp4"])
+def test_eplb_without_expert_parallelism_still_matches(mapping_kwargs):
+ """Without expert parallelism the engine builds no balancer, as for the built-in model."""
+ _claimed(
+ _model_config(**mapping_kwargs), moe_config=MoeConfig(load_balancer={"num_slots": 512})
)
- _claimed(config, max_num_tokens=32752)
- _declined(config, match="max_num_tokens=32768 at max_seq_len=262144", max_num_tokens=32768)
- _claimed(config, max_num_tokens=32768, max_seq_len=131072)
@pytest.mark.parametrize("max_draft_len", [1, 3, 7])
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_select_canonical.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_select_canonical.py
index 6ab7a2d73a3f..7d6fb0ed96cd 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_select_canonical.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_select_canonical.py
@@ -56,6 +56,9 @@ def test_rows_equal_the_canonical_definition(rows: int, cols: int) -> None:
def test_live_cols_returns_the_same_rows() -> None:
scores, row_starts, row_ends, indices = _case(16, 2304, seed=7)
+ # Every window ends short of the table, so the bound drops columns the rewrite would read.
+ row_ends.clamp_(max=1500)
+ assert int(row_ends.max()) < scores.shape[1]
bounded = indices.clone()
U._select(types.SimpleNamespace(block_topk=K), scores, row_starts, row_ends, indices)
U._select(
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision.py
index 6490c6a81cba..5c9bd6a64fc4 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision.py
@@ -553,10 +553,16 @@ def test_missing_or_mis_sized_position_ids_fail_the_request():
def test_a_mixed_request_lists_its_items():
with pytest.raises(ValueError, match="needs item metadata"):
_admit(_request(image=IMAGES, video=VIDEOS))
+ with pytest.raises(ValueError, match="needs item metadata"):
+ _admit(_request(image=IMAGES, video=VIDEOS, order=[]))
with pytest.raises(ValueError, match="lists image items \\[0\\]; the payload carries 2"):
_admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER[:3]))
+ with pytest.raises(ValueError, match="lists video items \\[0, 1, 0\\]"):
+ _admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER + [_item("video", 0)]))
with pytest.raises(ValueError, match="is not an item of this request"):
_admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER + [_item("audio", 0)]))
+ with pytest.raises(ValueError, match="is not an item of this request"):
+ _admit(_request(image=IMAGES, video=VIDEOS, order=[("image",)]))
def test_the_placeholder_count_is_the_sum_over_both_modalities():
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_triton_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_triton_update.py
index a289411fbafb..40b69748a206 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_triton_update.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_triton_update.py
@@ -1,6 +1,7 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
-"""GPU test for the causal_conv1d_triton_update entry."""
+"""GPU test for the causal_conv1d_triton_update entry. Beyond the cells: the
+snapshot-depth guard."""
import zlib
@@ -128,3 +129,15 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res
# Non-vacuity of the whole-pool gate: every slot the batch named moved.
for slot in cell.spec["slots"]:
assert not torch.equal(got[1][slot], args[1][slot]), slot
+
+
+def test_guard_refuses_undersized_buffers() -> None:
+ """A snapshot buffer shallower than seqlen is accepted by the kernel and
+ written past the row it was given (measured); the guard refuses it
+ before the launch."""
+ spec = dict(dim=512, batch=2, seqlen=2, pool_slots=8, slots=(0, 2), snapshots=2)
+ args, _ = _build(spec, seed=15)
+ args[8] = torch.zeros(2, 1, 512, 3, dtype=torch.bfloat16, device="cuda")
+ with pytest.raises(AssertionError, match="intermediate_conv_window must be at least"):
+ with validating(op):
+ op(*args)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_update.py
index d958f21cb0e3..345516d0b8dd 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_update.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_update.py
@@ -121,3 +121,29 @@ def test_which_kernel_each_member_runs() -> None:
spec = dict(dim=dim, x_row_stride=row_stride, batch=128, pool_slots=256, slots="random")
args = _build(spec, seed=dim)[0]
assert _decode_conv_ok(*args[:7]) is expect, (dim, expect)
+
+
+def test_a_pool_past_the_32_bit_slot_offset_is_refused() -> None:
+ """TRT-LLM's kernel forms a slot's offset in 32 bits: a pool of 2^32 elements or more is
+ refused before any launch (meta tensors: nothing is allocated)."""
+ stride = 1105920 # a dep4 rank's slot stride
+ slots = -(-(2**32) // stride)
+ x = torch.empty((4, 10240), dtype=torch.bfloat16, device="meta")
+ weight = torch.empty((10240, 4), dtype=torch.bfloat16, device="meta")
+ indices = torch.zeros(4, dtype=torch.int32, device="meta")
+ pool = torch.empty_strided(
+ (slots, 10240, 3), (stride, 3, 1), dtype=torch.bfloat16, device="meta"
+ )
+ with pytest.raises(ValueError, match="32 bits"):
+ op(x, pool, weight, None, True, None, indices, -1)
+
+
+def test_op_rejects_metadata_violations_loudly() -> None:
+ """Why the entry has no `is_valid`: the op raises on metadata violations
+ itself (probed here at int64 slot indices)."""
+ spec = dict(dim=512, batch=4, pool_slots=16, slots=(0, 1, 2, 3))
+ args = _build(spec, seed=14)[0]
+ args[6] = args[6].long()
+ with pytest.raises(RuntimeError, match="conv_state_indices.scalar_type"):
+ with validating(op):
+ op(*args)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py
index cc65c056494b..4dab4cb3bac0 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py
@@ -4,7 +4,8 @@
Integer arithmetic, gated bitwise. A near-miss vector is monotonic and
plausible, so each cell also proves the vector actually splits a packed
-buffer.
+buffer. Beyond the cells: CUDA-graph replay on rewritten lengths and the
+integer-dtype guard.
"""
import zlib
@@ -70,3 +71,51 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res
assert bounds[-1] == packed.shape[0]
for s, piece in enumerate(pieces):
assert torch.equal(packed[bounds[s] : bounds[s + 1]], piece), f"sequence {s} mis-sliced"
+
+
+#: Graph-captured batches: the MTP-7 verify of 32 requests, and the largest captured
+#: batch (128 requests) at the call site's int64 output.
+_GRAPH_BATCHES = ((32, torch.int32), (128, torch.int64))
+
+
+@pytest.mark.parametrize(
+ "batch, dtype", _GRAPH_BATCHES, ids=[f"{b}req-{str(d)[6:]}" for b, d in _GRAPH_BATCHES]
+)
+def test_cuda_graph_capture_and_replay(batch, dtype) -> None:
+ """A captured graph follows lengths rewritten in place: only the shape
+ reaches Python, the values are read on device and never synchronized on.
+ Replayed at one row a request, at MTP-7's 8, at one longer row, and at a
+ mixed step whose rows fill `max_num_tokens` (8192)."""
+ p, max_num_tokens = 8, 8192
+ static_lens = _seq_lens([p] * batch)
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ for _ in range(3):
+ op(static_lens, dtype)
+ torch.cuda.current_stream().wait_stream(side)
+ torch.cuda.synchronize()
+
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ static_out = op(static_lens, dtype)
+ for lengths in (
+ [1] * batch,
+ [p] * batch,
+ [p + 1] + [1] * (batch - 1),
+ [max_num_tokens - p * (batch - 1)] + [p] * (batch - 1),
+ ):
+ static_lens.copy_(_seq_lens(lengths))
+ graph.replay()
+ torch.cuda.synchronize()
+ assert static_out.dtype == dtype
+ assert torch.equal(static_out, op.reference(static_lens, dtype))
+
+
+def test_guard_refuses_non_integer_seq_lens() -> None:
+ """Unguarded, a float32 seq_lens is truncated toward zero before the scan."""
+ floats = torch.tensor([7.9, 1.2, 13.6], dtype=torch.float32, device=DEV)
+ assert op(floats, torch.int32).tolist() == [0, 7, 8, 21], "hazard gone"
+ with pytest.raises(AssertionError, match="seq_lens must have an integer dtype"):
+ with validating(op):
+ op(floats, torch.int32)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py
index da97829e7f5d..2392877d0746 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py
@@ -13,7 +13,8 @@
reference runs on a dense copy of those slots, and every other byte of their pages is
checked unchanged; a write to any other slot would fault.
-Beyond the cells: the has_initial_state equivalence.
+Beyond the cells: the has_initial_state equivalence, and a negative control that the band
+refuses a value head reading the wrong key head at 4/12.
"""
import warnings
@@ -27,6 +28,7 @@
from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.flashinfer_chunk_gated_delta_rule import ( # noqa: E501
INPLACE_MAX_TOKENS,
+ RANK4_GEOMETRY,
)
from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.flashinfer_chunk_gated_delta_rule import ( # noqa: E501
flashinfer_chunk_gated_delta_rule as op,
@@ -42,7 +44,7 @@
#: on `o`, 2e-2 on states), derived for a chunked WY reformulation with bf16
#: operands. This op is the same class of sum and sits inside it against fp64, CP
#: pipeline and chunked kernel alike; a value head reading the wrong key head at 4/12
-#: lands far outside it (measured).
+#: lands far outside it (`test_band_refuses_a_misgrouped_key_head`).
RTOL, ATOL_OUTPUT, ATOL_STATE = 1.6e-2, 2e-3, 2e-2
#: Rows past the call's own in the output buffer of a `slice` cell: the target passes
#: `out[:, :num_tokens]` of a wider buffer, and the rows past it must stay as they were.
@@ -353,9 +355,9 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res
def test_has_initial_state_is_the_reset_then_the_plain_call() -> None:
- """The flags replace `reset_recurrent_state_rows` before the call: zeroing the fresh
- slots and running the plain in-place call gives the same output and pool, bitwise; a
- linear gate that is `torch.exp` of the plain call's log gate changes nothing."""
+ """The flags are zeroing the fresh slots before the call: zeroing them and running the
+ plain in-place call gives the same output and pool, bitwise; a linear gate that is
+ `torch.exp` of the plain call's log gate changes nothing."""
args, _ = _build(_CARRIED.spec, seed=5)
flags = args.pop("has_initial_state")
linear_gate = args.pop("linear_gate", False)
@@ -375,3 +377,28 @@ def test_has_initial_state_is_the_reset_then_the_plain_call() -> None:
assert final_state is None
assert torch.equal(o, o_plain)
assert torch.equal(pool, plain["initial_state"])
+
+
+def test_band_refuses_a_misgrouped_key_head() -> None:
+ """The band is wider than torch's bf16 default, so it has to refuse the error a 4/12
+ rank can make and 16/48 shares: a value head reading the wrong key head. Rolling the
+ key heads by one makes value head i read key head (i // 3 - 1) mod 4; the op's answer
+ then has to fall outside the band around the true reference."""
+ spec = dict(
+ geometry=RANK4_GEOMETRY,
+ seqlens=(200, 37, 91),
+ slots=64,
+ indices=(63, 0, 17),
+ pool_dtype=torch.bfloat16,
+ output="buffer",
+ initial=(True, False, True),
+ linear_gate=True,
+ )
+ args, _ = _build(spec, seed=11)
+ pool = args["initial_state"]
+ ref = op.reference(**args)
+ rolled = dict(args, q=args["q"].roll(1, dims=2), k=args["k"].roll(1, dims=2))
+ o, _ = op(**rolled)
+ print(f"\n[misgrouped key head] {_stats('o', o, ref[0])}; {_stats('state', pool, ref[2])}")
+ with pytest.raises(AssertionError):
+ _compare((o, None, pool), ref)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py
index d811aa537d70..e8bedff31515 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py
@@ -25,6 +25,10 @@
reserved virtual range with memory mapped under the touched slots only (`_SparsePage`).
The reference and the twin run on a dense copy of those slots, and every byte of their
pages outside the cell's sub-pages is checked unchanged.
+
+Beyond the cells: a reference with the key heads reversed that the output gate must
+refuse, CUDA-graph replay against the same steps run eagerly, and the `state_indices`
+length guard.
"""
from __future__ import annotations
@@ -44,6 +48,7 @@
NUM_K_HEADS,
NUM_V_HEADS,
PACKED_QKV_WIDTH,
+ RANK_V_HEADS,
)
from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.fused_recurrent_gated_delta_rule_cached_replay_update import ( # noqa: E501
fused_recurrent_gated_delta_rule_cached_replay_update as op,
@@ -710,6 +715,37 @@ def _drive_cell(cell, bench) -> None:
)
+def test_the_band_refuses_a_wrong_head_pairing() -> None:
+ """The output gate is wider than torch's per-element default near a block's small
+ elements, so it has to refuse a plausible wrong answer. At a tp4/tep4 rank's 4 / 12
+ heads value head `hv` reads key head `hv // 3`; the same call's reference with the
+ key heads reversed (value heads 0..2 on key head 3) must fall outside it, by a margin
+ printed with the passing reference's."""
+ cell = next(
+ c
+ for c in op.CELLS
+ if c.spec.get("HV") == RANK_V_HEADS and c.spec["n"] == 128 and c.spec["t"] == 8
+ )
+ bench = _Bench(dict(cell.spec, start="fresh", steps=1), seed=4646, exact=False)
+ try:
+ kw = bench.kwargs(0, bench.inputs())
+ reversed_heads = dict(
+ kw, q=kw["q"].flip(2).contiguous(), k=kw["k"].flip(2).contiguous(), packed_qkv=None
+ )
+ right, wrong = bench.reference(kw)[0], bench.reference(reversed_heads)[0]
+ bench.prepare()
+ got = op(**kw)
+ _MEASURED.clear()
+ _compare((got, None, None, None, None), (right, None, None, None, None))
+ passing = dict(_MEASURED)
+ _MEASURED.clear()
+ with pytest.raises(AssertionError, match="output"):
+ _compare((got, None, None, None, None), (wrong, None, None, None, None))
+ print(f"\n[wrong head pairing] passing {passing}, wrong {dict(_MEASURED)}")
+ finally:
+ bench.close()
+
+
# ── CUDA graphs ───────────────────────────────────────────────────────────
@@ -757,3 +793,24 @@ def test_cuda_graph_replay_on_new_data_is_bitwise_eager() -> None:
accepted = eager.draw_acceptance()
eager.accept(accepted)
graph.accept(accepted)
+
+
+# ── guards: accepted and silently wrong ──────────────────────────────────
+
+
+def test_a_short_state_indices_reads_past_its_end_and_is_refused() -> None:
+ """Unguarded, three indices for four rows: the fourth row reads whatever
+ follows the view -- here a slot the call never named -- and writes its
+ history."""
+ bench = _Bench(dict(path="small", n=4, t=2, layers=1), seed=101, exact=False)
+ stranger = next(s for s in range(bench.slots) if s not in bench.live.tolist())
+ parent = torch.cat(
+ [bench.state_indices[:3], torch.tensor([stranger], dtype=torch.int32, device=DEV)]
+ )
+ kw = dict(bench.kwargs(0, bench.inputs()), state_indices=parent[:3])
+ before = bench.old_x[0, stranger].clone()
+ op(**kw)
+ assert not torch.equal(bench.old_x[0, stranger], before), "the hazard this guard covers is gone"
+ with pytest.raises(AssertionError, match="state_indices has shape"):
+ with validating(op):
+ op(**kw)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py
index d9688b4400f2..23365a9c97b3 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py
@@ -17,9 +17,9 @@
Beyond the cells: a fixed row bitwise across batch sizes over a bf16 pool, a stale-read
check under graph replay (a producer rewrites the projection inside the graph right
before the op), the two launch variants (state rows by bulk copy into shared memory, or
-loaded into registers) bitwise each other, and the variant choice over an unaligned
-pool, each at tp1's 16 / 48 heads and a tp4/tep4 rank's 4 / 12. The `gdn_decode_rows`
-entry (the recurrence alone) has its own cells below.
+loaded into registers) bitwise each other, the variant choice over an unaligned pool,
+and the guards, each at tp1's 16 / 48 heads and a tp4/tep4 rank's 4 / 12. The
+`gdn_decode_rows` entry (the recurrence alone) has its own cells below.
"""
import zlib
@@ -439,6 +439,18 @@ def test_variant_choice_and_unaligned_pool(v_heads) -> None:
assert torch.equal(got, want) and torch.equal(shifted, pool)
+@pytest.mark.parametrize("k_heads, v_heads", [(H, HV), RANK_HEADS], ids=["16x48", "4x12"])
+def test_guards(k_heads, v_heads) -> None:
+ """An fp16 pool, int64 slot indices and a projection of the wrong width are refused."""
+ args = _build(dict(rows=4, k_heads=k_heads, v_heads=v_heads), seed=5)[0]
+ with validating(op), pytest.raises(AssertionError):
+ op(args[0], args[1].to(torch.float16), *args[2:])
+ with validating(op), pytest.raises(AssertionError):
+ op(args[0], args[1], args[2].long(), *args[3:])
+ with validating(op), pytest.raises(AssertionError):
+ op(args[0][:, :-2], *args[1:])
+
+
# ---- gdn_decode_rows: the recurrence alone, into a row range of a step's output ----------
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_prefill_conv_post.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_prefill_conv_post.py
index a49ba3fde83a..57035ecb7246 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_prefill_conv_post.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_prefill_conv_post.py
@@ -10,7 +10,8 @@
bigger buffer, which the test also checks the call leaves alone outside the view: a
two-layer buffer for the tp1 cells, the stock manager's page (36 GDN layers a slot,
the view at the last) for a rank's cells: at an engine's measured slot count, or at
-slots whose offset passes 2^31 elements.
+slots whose offset passes 2^31 elements. Beyond the cells: `linear_gate` against the
+log-space gate, and the unit-channel-stride guard.
"""
import math
@@ -28,6 +29,7 @@
IN_PROJ_WIDTH,
NUM_K_HEADS,
NUM_V_HEADS,
+ QKV_DIM,
SOFTPLUS_THRESHOLD,
_engine_slots,
)
@@ -195,3 +197,12 @@ def test_linear_gate_is_torch_exp_of_the_log_gate() -> None:
for got, want in zip(lin_out[:3] + lin_out[4:], log_out[:3] + log_out[4:]):
assert torch.equal(got, want)
assert torch.equal(pool_lin, log_args[2])
+
+
+def test_guard_refuses_channel_strided_x() -> None:
+ args = _build(dict(lengths=(16,), pool_slots=2, slots=(0,), initial=(True,)), seed=302)
+ bad = list(args)
+ bad[0] = torch.empty(QKV_DIM, 16, dtype=DTYPE, device=DEV).t()
+ with pytest.raises(AssertionError, match="unit channel stride"):
+ with validating(op):
+ op(*bad)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/vision/test_modeling_v2_patch_embed_conv3d.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/vision/test_modeling_v2_patch_embed_conv3d.py
index 62f900fd6ed8..25c52c1c7752 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/vision/test_modeling_v2_patch_embed_conv3d.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/vision/test_modeling_v2_patch_embed_conv3d.py
@@ -6,7 +6,8 @@
flattened filter. The op is the same contraction made by cuBLASLt, so each
cell also holds the reference to an fp32 `F.conv3d` of the 5-D patches: that
pins the flatten order the two share to `nn.Conv3d`'s rather than letting one
-restate the other.
+restate the other. Beyond the cells: each guard against a measured silent
+answer, and graph capture.
"""
import contextlib
@@ -92,6 +93,14 @@ def _operands(
return pixel_values, weight, bias_t
+def _is_close(out, ref) -> bool:
+ try:
+ torch.testing.assert_close(out, ref)
+ except AssertionError:
+ return False
+ return True
+
+
@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)
@@ -111,3 +120,102 @@ def test_certified_cells(cell, certifying): # noqa: F811
torch.testing.assert_close(ref32, _conv3d_fp32(pixel_values, weight, bias))
_compare(out, op.reference(pixel_values, weight, bias))
+
+
+def test_guard_refuses_rows_that_are_not_one_patch() -> None:
+ """Two patches per row are split into twice the rows, without an error."""
+ pixel_values, weight, bias = _operands(8, seed=8)
+ doubled = pixel_values.reshape(4, 2 * PATCH_NUMEL)
+ assert op(doubled, weight, bias).shape == (8, EMBED_DIM)
+ with pytest.raises(AssertionError, match="each row must be exactly one patch"):
+ with validating(op):
+ op(doubled, weight, bias)
+
+
+def test_guard_refuses_strided_pixel_rows() -> None:
+ """A row-strided view (2-D, or 5-D with a batch stride) still views to
+ `[N, 1536]` and is read as if dense."""
+ pixel_values, weight, bias = _operands(256, seed=21)
+ ref = op.reference(pixel_values, weight, bias)
+ wide = torch.zeros(256, 2 * PATCH_NUMEL, dtype=pixel_values.dtype, device="cuda")
+ wide[:, :PATCH_NUMEL] = pixel_values
+ batched = torch.zeros(256, 2, IN_CHANNELS, KT, KH, KW, dtype=pixel_values.dtype, device="cuda")
+ batched[:, 0] = pixel_values.view(256, IN_CHANNELS, KT, KH, KW)
+ for strided in (wide[:, :PATCH_NUMEL], batched[:, 0]):
+ assert not strided.is_contiguous()
+ assert not _is_close(op(strided, weight, bias), ref), "the hazard this guard covers is gone"
+ with pytest.raises(AssertionError, match="pixel_values must be contiguous"):
+ with validating(op):
+ op(strided, weight, bias)
+
+
+def test_guard_refuses_a_strided_weight() -> None:
+ """A weight sliced along its output channels views to `[1152, 1536]` and is
+ read as if dense."""
+ pixel_values, weight, bias = _operands(256, seed=22)
+ ref = op.reference(pixel_values, weight, bias)
+ stacked = torch.zeros(2 * EMBED_DIM, IN_CHANNELS, KT, KH, KW, dtype=weight.dtype, device="cuda")
+ stacked[::2] = weight
+ sliced = stacked[::2]
+ assert not _is_close(op(pixel_values, sliced, bias), ref), (
+ "the hazard this guard covers is gone"
+ )
+ with pytest.raises(AssertionError, match="weight must be contiguous"):
+ with validating(op):
+ op(pixel_values, sliced, bias)
+
+
+def test_guard_refuses_a_bias_the_epilogue_misreads() -> None:
+ """A bias in another dtype, or strided, is read as dense output-dtype
+ elements: a plausible answer from the wrong bias. A short one is read past
+ its end, so it is refused without calling the op."""
+ pixel_values, weight, bias = _operands(256, seed=23)
+ ref = op.reference(pixel_values, weight, bias)
+ spread = torch.zeros(2 * EMBED_DIM, dtype=bias.dtype, device="cuda")
+ spread[::2] = bias
+ for bad, match in (
+ (bias.float(), "bias dtype"),
+ (bias.half(), "bias dtype"),
+ (spread[::2], "contiguous"),
+ ):
+ assert not _is_close(op(pixel_values, weight, bad), ref), (
+ "the hazard this guard covers is gone"
+ )
+ with pytest.raises(AssertionError, match=match):
+ with validating(op):
+ op(pixel_values, weight, bad)
+ with pytest.raises(AssertionError, match="contiguous \\[1152\\]"):
+ with validating(op):
+ op(pixel_values, weight, bias[: EMBED_DIM // 2].contiguous())
+
+
+def test_guard_refuses_fp32_with_a_bias() -> None:
+ """fp32 input drops the bias: the answer is the bias-free product."""
+ pixel_values, weight, bias = _operands(256, dtype=torch.float32, seed=24)
+ out = op(pixel_values, weight, bias)
+ torch.testing.assert_close(out, op.reference(pixel_values, weight, None))
+ assert not _is_close(out, op.reference(pixel_values, weight, bias))
+ with pytest.raises(AssertionError, match="dropped for fp32"):
+ with validating(op):
+ op(pixel_values, weight, bias)
+
+
+@pytest.mark.parametrize("num_patches", [64, 32768])
+def test_cuda_graph_capture_and_replay(num_patches) -> None:
+ pixel_values, weight, bias = _operands(num_patches, seed=12)
+
+ warmup = torch.cuda.Stream()
+ warmup.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(warmup):
+ for _ in range(3):
+ op(pixel_values, weight, bias)
+ torch.cuda.current_stream().wait_stream(warmup)
+
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ captured = op(pixel_values, weight, bias)
+
+ pixel_values.copy_(_operands(num_patches, seed=13)[0])
+ graph.replay()
+ torch.cuda.synchronize()
+ _compare(captured, op.reference(pixel_values, weight, bias))
From 8f72bda302d4956f83fd8531a88afe58bbfe3f33 Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Wed, 4 Nov 2026 22:57:08 -0800
Subject: [PATCH 06/15] [None][fix] modeling_v2 qwen4_exp: third-review
follow-ups
- The PLE host table is placed by NVML's CPU affinity
(get_numa_aware_cpu_affinity), as on main.
- No kernel compiles in a serving step: the CuTe GDN replay's launch mappings
are compiled at bind (precompile), and the vision tower warms up at load,
its rope and position-embedding kernels taking the grid unspecialized (the
position-embedding kernel now lives in the op, with no private
modeling_qwen3vl import).
- The row-tile op refuses a group size whose tile breaks the 8-row swizzle
atoms.
- causal_conv1d_update's pool-size message no longer matches bandit's SQL
pattern (B608).
- Tests: spec_rollback_commit guards, the fp4 window weight-stack refusal, tp4
head cuts, the gated-norm test's pool zeroed; the FlashInfer chunk note's
bitwise claim narrowed to the packed path.
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
.../attention/qsa_rowtile_attention_cute.py | 4 +-
.../ops/attention/qsa_sparse_gqa_rowtile.py | 8 +-
.../moe/fp4_block_scale_moe_renorm_routed.py | 5 +-
.../qwen4_exp/ops/ssm/causal_conv1d_update.py | 7 +-
.../ssm/flashinfer_chunk_gated_delta_rule.py | 9 +-
.../modeling.py | 14 ++
.../qsa_sparse_gqa_rowtile__prefill.py | 5 +-
.../ops/attention/vision_rope_2d_packed_qk.py | 15 +-
.../ops/ssm/gdn_cached_replay_cute__verify.py | 98 +++++++++---
.../vision/triton_pos_embed_interpolate.py | 140 +++++++++++++++---
.../vision.py | 15 ++
.../weights.py | 34 ++---
...test_modeling_v2_qsa_sparse_gqa_rowtile.py | 8 +
.../test_modeling_v2_spec_rollback_commit.py | 44 ++++++
...4_block_scale_moe_renorm_routed_windows.py | 23 +++
...st_modeling_v2_vision_rope_2d_packed_qk.py | 19 +++
...eling_v2_gdn_cached_replay_cute__verify.py | 25 ++++
.../test_vision_warm_up.py | 108 ++++++++++++++
.../test_weight_shards.py | 6 +-
...odeling_v2_triton_pos_embed_interpolate.py | 17 +++
...ng_v2_flashinfer_chunk_gated_delta_rule.py | 8 +-
..._modeling_v2_gdn_decode_rows_gated_norm.py | 4 +-
22 files changed, 534 insertions(+), 82 deletions(-)
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision_warm_up.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_rowtile_attention_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_rowtile_attention_cute.py
index 4fa6fef27d45..7abd032541a6 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_rowtile_attention_cute.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_rowtile_attention_cute.py
@@ -66,6 +66,8 @@
#: MMA rows per CTA and tokens per tile (one 64-bit bitmap word).
BLOCK_M = 128
BLOCK_N = 64
+#: The O staging's swizzle atom in MMA rows: a tile's rows x heads come in whole atoms.
+ROW_ATOM = 8
#: K/V ring slots (32 KB each); with the 60 KB Q buffer, about 220 KB of shared memory.
KV_SLOTS = 5
#: Work items in flight in the shared-memory schedule ring.
@@ -103,7 +105,7 @@ def __init__(
assert group_size * rows_per_tile <= 128
# The O staging's four SW128 boxes sit `64 G R` elements apart: each must start on
# a 1024-byte swizzle atom, so the live MMA rows come in whole 8-row atoms.
- assert group_size * rows_per_tile % 8 == 0, (group_size, rows_per_tile)
+ assert group_size * rows_per_tile % ROW_ATOM == 0, (group_size, rows_per_tile)
self.group_size = group_size
self.rows_per_tile = rows_per_tile
self.d = head_dim
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_sparse_gqa_rowtile.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_sparse_gqa_rowtile.py
index 328d93499bbc..dd103e6f695c 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_sparse_gqa_rowtile.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_sparse_gqa_rowtile.py
@@ -34,7 +34,7 @@
from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr
-from .qsa_rowtile_attention_cute import BLOCK_M, BLOCK_N
+from .qsa_rowtile_attention_cute import BLOCK_M, BLOCK_N, ROW_ATOM
from .qsa_rowtile_attention_cute import eligible as _cute_eligible
from .qsa_rowtile_attention_cute import num_ctas as _rowtile_ctas
from .qsa_rowtile_attention_cute import rowtile_attention_cute as _rowtile_attention_cute
@@ -613,6 +613,12 @@ def _check_inputs(q, k_cache, v_cache, block_table, selected_tokens, request_ind
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")
+ group_size = num_q_heads // num_kv_heads
+ if BLOCK_M // group_size * group_size % ROW_ATOM:
+ raise ValueError(
+ f"a row tile of {BLOCK_M // group_size} rows x {group_size} query heads per KV head "
+ f"must fill whole {ROW_ATOM}-row swizzle atoms of the output staging"
+ )
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) "
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_block_scale_moe_renorm_routed.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_block_scale_moe_renorm_routed.py
index d34bda4302d7..38bbc62e883d 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_block_scale_moe_renorm_routed.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_block_scale_moe_renorm_routed.py
@@ -340,8 +340,9 @@ class Fp4BlockScaleMoeRenormRouted(OpWrapper):
Windows: a window's weight stacks are its own experts only (the op compares the scales'
expert dimension with `local_num_experts`; the whole stack beside a window's scales would be
- read as its first experts, refused by `is_valid`). Not certified: windows of the padded
- I 256 stack (tp4 does not split experts); window rows past 8192 (the caller chunks).
+ read as its first experts, refused by `is_valid`; the windows test holds this refusal).
+ Not certified: windows of the padded I 256 stack (tp4 does not split experts); window rows
+ past 8192 (the caller chunks).
Hard limits
window the whole 512 at I 640 (tp1), 128 of 512 at I 640 (every offset), 256 of
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_update.py
index a35358e48c8e..6637f807dd9c 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_update.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_update.py
@@ -529,11 +529,12 @@ def __call__(
_decode_conv(x, conv_state, weight, conv_state_indices, pad_slot_id)
return x
if conv_state.shape[0] * conv_state.stride(0) >= TRTLLM_SLOT_OFFSET_LIMIT:
+ last = (TRTLLM_SLOT_OFFSET_LIMIT - 1) // conv_state.stride(0)
raise ValueError(
f"a conv-state pool of {conv_state.shape[0]} slots at stride "
- f"{conv_state.stride(0)}: TRT-LLM's causal_conv1d_update forms a slot's offset "
- f"in 32 bits and would wrap past slot {(TRTLLM_SLOT_OFFSET_LIMIT - 1) // conv_state.stride(0)}; "
- "lower free_gpu_memory_fraction or the recurrent-state share of the KV cache"
+ f"{conv_state.stride(0)}: TRT-LLM's C++ conv kernel addresses a slot in 32 bits "
+ f"and would wrap past slot {last}; lower free_gpu_memory_fraction or the "
+ "recurrent-state share of the KV cache"
)
torch.ops.trtllm.causal_conv1d_update(
x.unsqueeze(-1) if x.dim() == 2 else x,
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/flashinfer_chunk_gated_delta_rule.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/flashinfer_chunk_gated_delta_rule.py
index 9a9b1a5637e0..affaf202259a 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/flashinfer_chunk_gated_delta_rule.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/flashinfer_chunk_gated_delta_rule.py
@@ -405,9 +405,12 @@ class FlashinferChunkGatedDeltaRule(OpWrapper):
With `has_initial_state` (`hi_p` / `hi_all`): a fresh sequence's slot is never
read, and the scatter writes every named slot, so the fresh slots of the
recurrent pool need no zeroing before this call (the conv pass reads no carry
- for such a sequence and its publish writes the whole conv window). Bitwise
- zeroing the fresh slots followed by the plain call; the gather moves only the
- carried sequences' rows.
+ for such a sequence and its publish writes the whole conv window). Above
+ `INPLACE_MAX_TOKENS` context tokens (the packed path) it is bitwise zeroing the
+ fresh slots followed by the plain call, which the test holds; the gather moves
+ only the carried sequences' rows. A call of at most `INPLACE_MAX_TOKENS` runs in
+ place (below), within the same band of the fp64 reference but not bitwise to the
+ plain call, which may take the context-parallel pipeline instead.
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
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py
index fb78e4aa8757..d76b8a91a1cc 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py
@@ -448,6 +448,7 @@
from .ops.hc.qwen4_exp_hc_decode_tc_first__decode import qwen4_exp_hc_decode_tc_first
from .ops.hc.qwen4_exp_hc_up_gate_mean_fp4__moe import qwen4_exp_hc_up_gate_mean_fp4
from .ops.ssm.gdn_cached_replay_cute__verify import gdn_cached_replay_cute
+from .ops.ssm.gdn_cached_replay_cute__verify import precompile as gdn_replay_precompile
# The GPU architecture this target is certified for; a direct instantiation could
# bypass routing.
@@ -1608,6 +1609,18 @@ def _bind_cache_manager(self, md, position_ids) -> None:
f"a {state_dtype} recurrent-state cache; this layout serves "
f"{' or '.join(str(d) for d in served)} (kv_cache_config.mamba_ssm_cache_dtype)"
)
+ if self.mtp_enabled and state_dtype == torch.bfloat16:
+ # Every CuTe replay kernel a verify step of up to `max_num_requests` requests
+ # launches (`_gdn_verify`), compiled now: the captured decode graphs do not reach
+ # every launch mapping, and a compile in a serving step stalls it 5-10 s.
+ gdn_replay_precompile(
+ self.max_verify_rows,
+ self.num_k_heads,
+ self.num_v_heads,
+ md.max_num_requests,
+ pdl=_GDN_REPLAY_PDL,
+ partition_from=CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE,
+ )
# No sliding window: the targets pass `max_seq_len` as the window.
windows = cm.max_attention_window_vec
attn_windows = {windows[off] for off in rank_pools.pool_idx.values()}
@@ -6771,6 +6784,7 @@ def post_load_weights(self):
super().post_load_weights()
self.model.post_load_weights()
self.vision_tower = _vision.VisionTower(self.model.w, self.model.vision_config)
+ self.vision_tower.warm_up()
if self.one_rank:
self._build_heads(self.lm_head.weight.data)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_sparse_gqa_rowtile__prefill.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_sparse_gqa_rowtile__prefill.py
index 4a37ca42dfd4..62ec41278514 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_sparse_gqa_rowtile__prefill.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_sparse_gqa_rowtile__prefill.py
@@ -38,6 +38,7 @@
from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_rowtile_attention_cute import ( # noqa: E501
BLOCK_M,
BLOCK_N,
+ ROW_ATOM,
)
from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_rowtile_attention_cute import ( # noqa: E501
eligible as _cute_eligible,
@@ -64,8 +65,6 @@
# this many quarters of a per-row tile (1.5x), cost no more than its rows' own columns.
_SHARED_TILE_COST = 6
_PLAN_WARPS = 4
-#: The O staging's swizzle atom in MMA rows: a group's rows x heads come in whole atoms.
-_ROW_ATOM = 8
_LONE_WARPS = 2
_LONE_STAGES = 2
@@ -548,7 +547,7 @@ def rows_per_group(group_size: int) -> int:
"""Rows of one request a row-tile item serves: the most whose `rows x group_size` MMA
rows fill whole 8-row swizzle atoms of the 128-row tile (10 at 12 heads, 20 at 6)."""
rows = BLOCK_M // group_size
- while rows > 1 and rows * group_size % _ROW_ATOM:
+ while rows > 1 and rows * group_size % ROW_ATOM:
rows -= 1
return rows
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/vision_rope_2d_packed_qk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/vision_rope_2d_packed_qk.py
index 7a7b746d84a0..ebaa515b8e06 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/vision_rope_2d_packed_qk.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/vision_rope_2d_packed_qk.py
@@ -71,7 +71,8 @@
_BLOCK_M = 8
-@triton.jit
+# The token counts vary per image: not specialized, so one compile serves every grid.
+@triton.jit(do_not_specialize=["seqlen", "seqlen_ro"])
def _rotary_neox_inplace_kernel(
X,
COS,
@@ -175,11 +176,13 @@ class VisionRope2dPackedQk(OpWrapper):
(tp4, tep4, dep4) runs the tower at these shapes.
Source: Triton, JIT-compiled at first call for the running device; the cells ran that
- path on sm_103 (GB300). Launch: `(cdiv(heads, 4), cdiv(T, 8))` programs, one per 4
- heads x 8 tokens, each loading its 8 tokens' cos / sin once and both half-heads of
- its 4 heads with 64-lane masked loads (36 live). Token offsets are int64, so no
- address product overflows (the origin's int32 product would pass 2^31 at 621379 rows
- of the 3456-wide qkv row); the cos / sin index is `t * half`.
+ path on sm_103 (GB300). The token counts are not specialized, so one compile serves
+ every grid (the tower's `warm_up` makes it at load; the test holds one variant).
+ Launch: `(cdiv(heads, 4), cdiv(T, 8))` programs, one per 4 heads x 8 tokens, each
+ loading its 8 tokens' cos / sin once and both half-heads of its 4 heads with 64-lane
+ masked loads (36 live). Token offsets are int64, so no address product overflows (the
+ origin's int32 product would pass 2^31 at 621379 rows of the 3456-wide qkv row); the
+ cos / sin index is `t * half`.
Preconditions: `qk` 3-D, any strides (the kernel reads all three); `cos` and `sin`
fp32, contiguous, both `[T, head_dim // 2]`. In place: the return value is `qk`
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/ssm/gdn_cached_replay_cute__verify.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/ssm/gdn_cached_replay_cute__verify.py
index fa0c4738b1bc..29d8bba0b215 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/ssm/gdn_cached_replay_cute__verify.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/ssm/gdn_cached_replay_cute__verify.py
@@ -1687,6 +1687,75 @@ def default_config(n: int, t: int, commit: bool = False) -> dict:
return dict(tpc=8, warps=4, pf=1, minb=5)
+def _kernel_kw(
+ n: int,
+ t: int,
+ h: int,
+ hv: int,
+ *,
+ commit: bool,
+ pdl: bool,
+ l2norm: bool,
+ gating: bool,
+ g_f32: bool,
+ b_f32: bool,
+ work_items: bool,
+ config: Optional[dict] = None,
+) -> dict:
+ """The kernel's compile-time parameters for a call of `n` requests: its key in `_COMPILED`."""
+ cfg = dict(default_config(n, t, commit))
+ if config:
+ cfg.update(config)
+ # with PDL, issue the state / history (not written by the previous kernel) before the wait;
+ # without it, the step's own small inputs first
+ cfg.setdefault("qk_first", not pdl)
+ cfg["use_wi"] = bool(cfg.pop("use_wi", True)) and work_items
+ return dict(
+ T=t,
+ h=h,
+ hv=hv,
+ commit=commit,
+ pdl=pdl,
+ l2norm=l2norm,
+ gating=gating,
+ g_f32=g_f32,
+ b_f32=b_f32,
+ **cfg,
+ )
+
+
+def precompile(
+ t: int, h: int, hv: int, max_requests: int, *, pdl: bool, partition_from: int
+) -> int:
+ """Compile every kernel a verify step of 1..`max_requests` requests launches, the way a
+ target calls this entry: work items, and with them no in-kernel commit, from
+ `partition_from` requests on; the q/k L2 norm and the gating in-kernel; bf16 `a` / `b`.
+ A compile costs 5-10 s, so a target runs this at bind, before any step: the captured
+ decode graphs do not reach every mapping (padded graphs skip 9 to 15 requests, which an
+ eager mixed step carries). Returns the number of distinct kernels."""
+ keys = set()
+ for n in range(1, max_requests + 1):
+ partitioned = n >= partition_from
+ kw = _kernel_kw(
+ n,
+ t,
+ h,
+ hv,
+ commit=not partitioned,
+ pdl=pdl,
+ l2norm=True,
+ gating=True,
+ g_f32=False,
+ b_f32=False,
+ work_items=partitioned,
+ )
+ key = tuple(sorted(kw.items()))
+ if key not in keys:
+ keys.add(key)
+ _compile(key, _Replay(**kw))
+ return len(keys)
+
+
def _compile(key, kern):
fn = _COMPILED.get(key)
if fn is None:
@@ -1800,27 +1869,21 @@ def cached_replay_update(
ap, dp = A_log.data_ptr(), dt_bias.data_ptr()
else:
ap = dp = 0
- cfg = dict(default_config(N, T, commit))
- if config:
- cfg.update(config)
- # with PDL, issue the state / history (not written by the previous kernel) before the wait;
- # without it, the step's own small inputs first
- cfg.setdefault("qk_first", not bool(launch_with_pdl))
- use_wi = bool(cfg.pop("use_wi", True)) and replay_work_items is not None
- cfg["use_wi"] = use_wi
- pdl = bool(launch_with_pdl)
- kw = dict(
- T=T,
- h=H,
- hv=HV,
+ kw = _kernel_kw(
+ N,
+ T,
+ H,
+ HV,
commit=commit,
- pdl=pdl,
+ pdl=bool(launch_with_pdl),
l2norm=bool(use_qk_l2norm_in_kernel),
gating=gating,
g_f32=g.dtype == torch.float32,
b_f32=beta.dtype == torch.float32,
- **cfg,
+ work_items=replay_work_items is not None,
+ config=config,
)
+ use_wi = kw["use_wi"]
key = tuple(sorted(kw.items()))
fn = _COMPILED.get(key)
if fn is None:
@@ -2020,8 +2083,9 @@ class GdnCachedReplayCute(OpWrapper):
dependents at its start (`launch_dependent_kernels`), so the kernel streams the checkpoint
and the history -- which nothing in the step writes ahead of it -- while the conv runs, and
reads q/k/v, a and b only after `griddepcontrol.wait`. Compiles one kernel per (geometry,
- T, launch mapping, commit, PDL, ...) on first use (5-10 s): the warm-up before CUDA-graph
- capture covers it.
+ T, launch mapping, commit, PDL, ...) on first use (5-10 s); `precompile` builds every one a
+ target's verify steps reach, and the target runs it at bind (the captured decode graphs
+ alone skip 9 to 15 requests, which an eager mixed step carries).
Certified at the whole layer's 16 / 48 heads (tp1, dep4) over a bf16 state: the in-kernel
commit at N = 1 and T = 3, 4, the all-layer commit at N = 32 and 256 (T = 4). Certified at
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/vision/triton_pos_embed_interpolate.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/vision/triton_pos_embed_interpolate.py
index 5b395305068b..04c231b99f14 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/vision/triton_pos_embed_interpolate.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/vision/triton_pos_embed_interpolate.py
@@ -4,21 +4,22 @@
grid, emitted in spatial-merge order, by the stock Qwen3-VL tower's fused
Triton kernel.
-`__call__` is one call to `_triton_pos_embed_interpolate` in
-`tensorrt_llm/_torch/models/modeling_qwen3vl.py`, which allocates the output
-and makes one launch of the Triton JIT kernel `_bilinear_pos_embed_kernel`
-beside it: the four corner gathers, the bilinear weights and the merge-order
-row permutation in one pass. It is the kernel the stock `modeling_qwen4_exp`
-vision tower runs (`Qwen3VisionModel.fast_pos_embed_interpolate`), one call per
-image. Concatenating several images' rows is the caller's.
+`__call__` allocates the output and makes one launch of the Triton JIT kernel
+`_bilinear_pos_embed_kernel`: the four corner gathers, the bilinear weights and
+the merge-order row permutation in one pass. The kernel and its launch are the
+ones `tensorrt_llm/_torch/models/modeling_qwen3vl.py` runs for the stock vision
+towers (`Qwen3VisionModel.fast_pos_embed_interpolate`, one call per image), with
+the grid sides left unspecialized so one compile serves every image.
+Concatenating several images' rows is the caller's.
"""
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
-from tensorrt_llm._torch.models.modeling_qwen3vl import _triton_pos_embed_interpolate
#: This checkpoint: num_position_embeddings 2304 (a 48 x 48 grid), hidden 1152,
#: spatial_merge_size 2.
@@ -40,6 +41,108 @@ def _ckpt(t: int, h: int, w: int) -> dict:
)
+# The grid sides vary per image: not specialized, so one compile serves every grid.
+@triton.jit(do_not_specialize=["H", "W"])
+def _bilinear_pos_embed_kernel(
+ embed_ptr,
+ output_ptr,
+ H,
+ W,
+ h_scale,
+ w_scale,
+ NUM_GRID: tl.constexpr,
+ M_SIZE: tl.constexpr,
+ HIDDEN_DIM: tl.constexpr,
+ BLOCK_D: tl.constexpr,
+):
+ """One output row: the bilinear blend of the table's four corner rows around the
+ row's patch, the patch taken in spatial-merge order."""
+ pid = tl.program_id(0)
+ total_spatial = H * W
+ spatial_idx = pid % total_spatial
+
+ num_blocks_w = W // M_SIZE
+ block_idx = spatial_idx // (M_SIZE * M_SIZE)
+ local_idx = spatial_idx % (M_SIZE * M_SIZE)
+ br = block_idx // num_blocks_w
+ bc = block_idx % num_blocks_w
+ lr = local_idx // M_SIZE
+ lc = local_idx % M_SIZE
+ row = br * M_SIZE + lr
+ col = bc * M_SIZE + lc
+
+ h_frac = row.to(tl.float32) * h_scale
+ w_frac = col.to(tl.float32) * w_scale
+
+ hf = tl.math.floor(h_frac).to(tl.int32)
+ wf = tl.math.floor(w_frac).to(tl.int32)
+ hc = tl.minimum(hf + 1, NUM_GRID - 1)
+ wc = tl.minimum(wf + 1, NUM_GRID - 1)
+
+ dh = h_frac - hf.to(tl.float32)
+ dw = w_frac - wf.to(tl.float32)
+ w11 = dh * dw
+ w10 = dh - w11
+ w01 = dw - w11
+ w00 = 1.0 - dh - w01
+
+ off00 = (hf * NUM_GRID + wf) * HIDDEN_DIM
+ off01 = (hf * NUM_GRID + wc) * HIDDEN_DIM
+ off10 = (hc * NUM_GRID + wf) * HIDDEN_DIM
+ off11 = (hc * NUM_GRID + wc) * HIDDEN_DIM
+ out_off = pid * HIDDEN_DIM
+
+ # The corner weights rounded to the output dtype; the multiply-adds run in it.
+ out_dtype = output_ptr.dtype.element_ty
+ w00_c = w00.to(out_dtype)
+ w01_c = w01.to(out_dtype)
+ w10_c = w10.to(out_dtype)
+ w11_c = w11.to(out_dtype)
+
+ for d in tl.range(0, HIDDEN_DIM, BLOCK_D):
+ cols = d + tl.arange(0, BLOCK_D)
+ mask = cols < HIDDEN_DIM
+
+ e00 = tl.load(embed_ptr + off00 + cols, mask=mask)
+ e01 = tl.load(embed_ptr + off01 + cols, mask=mask)
+ e10 = tl.load(embed_ptr + off10 + cols, mask=mask)
+ e11 = tl.load(embed_ptr + off11 + cols, mask=mask)
+
+ val = w00_c * e00 + w01_c * e01 + w10_c * e10 + w11_c * e11
+
+ tl.store(output_ptr + out_off + cols, val, mask=mask)
+
+
+def _pos_embed_interpolate(
+ embed_weight: torch.Tensor,
+ t: int,
+ h: int,
+ w: int,
+ num_grid_per_side: int,
+ m_size: int,
+ dtype: torch.dtype,
+) -> torch.Tensor:
+ """One launch over a `(t, h, w)` grid: `[t * h * w, hidden]` in spatial-merge order."""
+ hidden_dim = embed_weight.shape[1]
+ total_out = t * h * w
+ output = torch.empty(total_out, hidden_dim, device=embed_weight.device, dtype=dtype)
+ h_scale = (num_grid_per_side - 1) / (h - 1) if h > 1 else 0.0
+ w_scale = (num_grid_per_side - 1) / (w - 1) if w > 1 else 0.0
+ _bilinear_pos_embed_kernel[(total_out,)](
+ embed_weight,
+ output,
+ h,
+ w,
+ h_scale,
+ w_scale,
+ num_grid_per_side,
+ m_size,
+ hidden_dim,
+ triton.next_power_of_2(hidden_dim),
+ )
+ return output
+
+
class TritonPosEmbedInterpolate(OpWrapper):
"""Resample a square row-major `[N*N, hidden]` table onto one `(t, h, w)`
patch grid; return a fresh `[t*h*w, hidden]` tensor in `dtype`.
@@ -75,11 +178,12 @@ class TritonPosEmbedInterpolate(OpWrapper):
note = """
Source: a Triton JIT kernel, compiled at first call for the device it runs
- on. `H` and `W` are int32, the scales fp32, and `NUM_GRID`, `M_SIZE`,
+ on: the stock Qwen3-VL towers' kernel (`modeling_qwen3vl`), kept here so its
+ grid sides can stay unspecialized. `H` and `W` are int32 and not specialized,
+ so every grid shares one compile (the tower's `warm_up` makes it at load; the
+ test holds one variant); the scales are fp32, and `NUM_GRID`, `M_SIZE`,
`HIDDEN_DIM` (1152) and `BLOCK_D` (2048) are constexprs, so another table
- geometry compiles another kernel. Tier 2: a callable already in the tree (a
- private function of the stock Qwen3-VL modeling file), not a registered
- `torch.ops.trtllm` op. Importing this entry imports `modeling_qwen3vl`.
+ geometry compiles another kernel. Not a registered `torch.ops.trtllm` op.
Accuracy, against the exact (fp64) corner-aligned bilinear value, as
row-scale element error / relative RMS in bf16 ulp (`assert_within_ulp`):
@@ -104,9 +208,9 @@ class TritonPosEmbedInterpolate(OpWrapper):
argument, not read off the table. 47 on the 2304-row table returned a
correctly shaped result 564 / 340 ulp off.
- an empty side (`t`, `h` or `w` of 0): returns `[0, hidden]`, no error.
- - `h` or `w` not divisible by `m_size`: the wrapper's own check is a bare
- `assert`. Under `python -O` a 6x9 grid returned a finite `[54, 1152]`
- tensor, its rows read from positions outside the grid and the table.
+ - `h` or `w` not divisible by `m_size`: a 6x9 grid returned a finite
+ `[54, 1152]` tensor, its rows read from positions outside the grid and the
+ table.
Hard limits: the output offset is `program_id * hidden` in int32, so
`t*h*w*hidden` must stay below 2^31 (1,864,135 rows at hidden 1152; the
@@ -129,9 +233,7 @@ def __call__(
m_size: int,
dtype: torch.dtype,
) -> torch.Tensor:
- return _triton_pos_embed_interpolate(
- embed_weight, t, h, w, num_grid_per_side, m_size, dtype
- )
+ return _pos_embed_interpolate(embed_weight, t, h, w, num_grid_per_side, m_size, dtype)
def reference(
self,
@@ -203,7 +305,7 @@ def is_valid(
)
assert h % m_size == 0 and w % m_size == 0, (
f"grid sides h={h}, w={w} are not whole m_size={m_size} merge blocks; "
- f"the kernel's bare assert is all that stops an out-of-grid read"
+ f"the kernel would read rows outside the grid and the table"
)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision.py
index cc2e590770f1..0aa802cb3803 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision.py
@@ -163,6 +163,11 @@
_POSITION_AXES = 3
+#: The grids `VisionTower.warm_up` runs: one call whose patch count is a multiple of 16 (4x4)
+#: and one whose count is not (2x2), the two variants Triton compiles a row count into.
+WARM_UP_GRIDS = (((1, 4, 4),), ((1, 2, 2),))
+
+
# --------------------------------------------------------------- weight views ---
@@ -333,6 +338,16 @@ def __init__(self, w: dict, vision_config) -> None:
# engine's own workspace.
self._ws: torch.Tensor | None = None
+ def warm_up(self) -> None:
+ """Run the tower once per `WARM_UP_GRIDS` grid on zero pixels, so every JIT kernel it
+ launches is compiled at load rather than on the first media request. The rope and
+ position-embedding kernels take the grid as unspecialized ints; the layer norm's row
+ count is specialized only on divisibility by 16, and the two grids cover both."""
+ device = self._w.pos.device
+ for grid in WARM_UP_GRIDS:
+ rows = sum(t * h * w for t, h, w in grid)
+ self._run(torch.zeros(rows, self.pixel_width, dtype=self.dtype, device=device), grid)
+
def __call__(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor:
"""`pixel_values` is `[sum(t*h*w), in_channels*kt*kh*kw]`, one flat patch per row,
channel-slowest. `grid_thw` is one `(frames, height, width)` row per image or video
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/weights.py
index 864ff5efea9e..2a405f1ef4a5 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/weights.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/weights.py
@@ -99,6 +99,7 @@
from tensorrt_llm._torch.pyexecutor.config_utils import mamba_effective_tp_size
from tensorrt_llm._utils import CUASSERT
+from tensorrt_llm.llmapi.utils import get_numa_aware_cpu_affinity
from .pool import attention_tp_size
@@ -193,23 +194,16 @@ def memory_kind(ptr: int) -> int:
def _local_cpus(device) -> set | None:
- """The CPUs sysfs lists as local to `device`'s PCI function and this process may run on,
- or None when that is unknown."""
+ """The CPUs of `device`'s NUMA affinity, as NVML reports it (`get_numa_aware_cpu_affinity`),
+ that this process may run on, or None when NVML cannot tell."""
+ import pynvml
+
+ index = torch.device(device).index
try:
- props = torch.cuda.get_device_properties(device)
- path = (
- f"/sys/bus/pci/devices/{props.pci_domain_id:04x}:{props.pci_bus_id:02x}:"
- f"{props.pci_device_id:02x}.0/local_cpulist"
- )
- with open(path) as fh:
- text = fh.read().strip()
- except (OSError, AttributeError, RuntimeError):
+ local = get_numa_aware_cpu_affinity(torch.cuda.current_device() if index is None else index)
+ except pynvml.NVMLError:
return None
- cpus = set()
- for part in filter(None, text.split(",")):
- lo, _, hi = part.partition("-")
- cpus.update(range(int(lo), int(hi or lo) + 1))
- return (cpus & os.sched_getaffinity(0)) or None
+ return (set(local) & os.sched_getaffinity(0)) or None
class _RegisteredHostBuffer:
@@ -218,16 +212,16 @@ class _RegisteredHostBuffer:
allocation with `mmap` and returns it with `munmap`, so freeing registered pages
would leave the driver mapping memory the OS may hand out again.
- Not torch's pinned allocator: it rounds the 47.7 GiB table up to 64 GiB, and the
- device's random row gathers from its pages are slower than from registered pageable
- memory."""
+ Not torch's pinned allocator: it rounds the 47.7 GiB table up to 64 GiB, and on
+ GB300 the device's random row gathers from its pages run an order of magnitude
+ slower than from registered pageable memory."""
def __init__(self, nbytes: int, device=None) -> None:
# The registration faults the pages in on this thread, so under the default local
# allocation policy they land on the NUMA node of the CPU it runs on. Pinned to the
# CPUs local to `device` meanwhile, the table sits next to the GPU that gathers from
- # it: with a process per GPU, half the ranks would otherwise gather from the other
- # socket's memory.
+ # it: on GB300, random row gathers from the other socket's memory run ~1.7x slower,
+ # and with a process per GPU half of them would otherwise land there.
local = _local_cpus(device) if device is not None else None
previous = os.sched_getaffinity(0) if local else None
if local:
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py
index 12bb8eada7fa..f3ec7837bdf6 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py
@@ -252,3 +252,11 @@ def test_rejects_a_call_without_the_prefill_pair() -> None:
args.update(query_positions=None, compress_ratio=None)
with pytest.raises(ValueError, match=re.escape("the row-tile kernel needs the prefill pair")):
op(**args)
+
+
+def test_rejects_a_group_size_whose_tile_breaks_the_row_atoms() -> None:
+ """6 query heads per KV head would take 21 rows (126 MMA rows), not whole 8-row swizzle
+ atoms: the op refuses it before the kernel is built."""
+ args = _build(_with([(600, 4)], slots=8, q_heads=12), seed=9200)
+ with pytest.raises(ValueError, match="must fill whole 8-row swizzle atoms"):
+ op(**args)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_spec_rollback_commit.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_spec_rollback_commit.py
index 2631822136f5..b5c229535841 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_spec_rollback_commit.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_spec_rollback_commit.py
@@ -317,3 +317,47 @@ def test_commit_graph_replay_rereads_its_inputs() -> None:
_GUARD_SPEC = dict(n=8, p=4, planes=12)
+
+#: Argument positions in `_Step.call_args()`.
+_STREAM, _PAGES, _PTRS, _STRIDES, _INDEX_OLD = 3, 9, 11, 12, 13
+
+
+def _strided_index_old(step: _Step, args: list) -> None:
+ wide = torch.zeros(step.count, step.n * step.p, 2 * ROW, dtype=torch.bfloat16, device=DEV)
+ args[_INDEX_OLD] = wide[..., :ROW]
+
+
+def _narrow_stream(step: _Step, args: list) -> None:
+ args[_STREAM] = args[_STREAM][:, :-1].contiguous()
+
+
+def _int32_pages(step: _Step, args: list) -> None:
+ args[_PAGES] = args[_PAGES].to(torch.int32)
+
+
+def _short_plane_table(step: _Step, args: list) -> None:
+ args[_PTRS], args[_STRIDES] = args[_PTRS][:-1], args[_STRIDES][:-1]
+
+
+@pytest.mark.parametrize(
+ "perturb, message",
+ [
+ (_strided_index_old, "index_old must be the contiguous"),
+ (_narrow_stream, "stream must be the contiguous"),
+ (_int32_pages, "pages and within must be int64"),
+ (_short_plane_table, "the plane table holds 11 planes and index_old 12"),
+ ],
+ ids=["strided_index_old", "narrow_stream", "int32_pages", "short_plane_table"],
+)
+def test_guard_refuses_a_misshapen_staged_input(perturb, message) -> None:
+ """Each staged tensor the kernel reads by raw offset, one wrong at a time: `is_valid`
+ refuses it before the launch, so no pool is written."""
+ step = _Step(_GUARD_SPEC, 31)
+ before = step.snapshot()
+ args = list(step.call_args())
+ perturb(step, args)
+ with validating(spec_rollback_commit), pytest.raises(AssertionError, match=message):
+ spec_rollback_commit(*args)
+ torch.cuda.synchronize()
+ for got, want, what in zip(step.snapshot(), before, ("n-gram", "conv", "position")):
+ _equal(got, want, f"{what} pool after a refused call")
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed_windows.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed_windows.py
index d7a109cb9017..780ed0335aff 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed_windows.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed_windows.py
@@ -205,3 +205,26 @@ def test_certified_cells(cell, certifying, many_expert_layer): # noqa: F811
H._compare(y, exp.to(torch.bfloat16))
elt, rms = H._dev(y, exp.to(torch.bfloat16))
print(f" cell {cell.why[:64]!r}: {elt:.2f} elt / {rms:.2f} rms ulp")
+
+
+def test_guard_refuses_the_whole_stack_beside_a_window(many_expert_layer): # noqa: F811
+ """The whole 512-expert weight stack passed beside a window's scales: the kernel sizes the
+ window from the scales and would read the stack's first experts as the window's;
+ `is_valid` refuses it."""
+ cell = next(
+ c
+ for c in _WINDOWS
+ if c.spec["intermediate"] == 640 and c.spec["local_num_experts"] < c.spec["num_experts"]
+ )
+ spec = cell.spec
+ win_args, _, data, sf, _, logits, _, _, _, scalars = _operands(spec, 7, many_expert_layer)
+ whole, _ = _stack(spec, many_expert_layer)
+ weights = dict(
+ win_args, gemm1_weights=whole["gemm1_weights"], gemm2_weights=whole["gemm2_weights"]
+ )
+ message = (
+ f"gemm1_weights holds {spec['num_experts']} experts but the window has "
+ f"{spec['local_num_experts']}"
+ )
+ with validating(op), pytest.raises(AssertionError, match=message):
+ _routed(win_args, data, sf, logits, scalars, spec, weights=weights)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_vision_rope_2d_packed_qk.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_vision_rope_2d_packed_qk.py
index 98908916f232..c8a43ca0ee9d 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_vision_rope_2d_packed_qk.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_vision_rope_2d_packed_qk.py
@@ -22,6 +22,7 @@
from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.attention.vision_rope_2d_packed_qk import ( # noqa: E501
HEAD_DIM,
NUM_HEADS,
+ _rotary_neox_inplace_kernel,
)
from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.attention.vision_rope_2d_packed_qk import ( # noqa: E501
vision_rope_2d_packed_qk as op,
@@ -195,3 +196,21 @@ def test_is_valid_refuses_each_silent_table_fault(fault):
assert not torch.equal(got, want), f"{fault}: the kernel got it right after all"
with validating(op), pytest.raises(AssertionError):
op(_qk(qkv.clone()), bad_cos, bad_sin)
+
+
+def test_one_compile_serves_every_grid():
+ """The token counts are not specialized: grids whose counts are and are not multiples of
+ 16 run the one compiled kernel, bitwise the reference."""
+ variants = []
+ for grid in ([(1, 2, 2)], [(1, 4, 4)], [(1, 6, 10)], [(2, 8, 12)], [(1, 62, 90)]):
+ grid_thw = torch.tensor(grid, dtype=torch.int64)
+ qkv = _packed_qkv(_num_tokens(grid), torch.bfloat16, seed=23)
+ cos, sin = _angle_table(grid_thw, MERGE, THETA, HEAD_DIM)
+ want = op.reference(_qk(qkv), cos, sin)
+ got = op(_qk(qkv), cos, sin)
+ torch.cuda.synchronize()
+ assert _mismatch(got, want)[0] == 0, grid
+ variants.append(
+ len(_rotary_neox_inplace_kernel.device_caches[torch.cuda.current_device()][0])
+ )
+ assert variants == [variants[0]] * len(variants), variants
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ssm/test_modeling_v2_gdn_cached_replay_cute__verify.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ssm/test_modeling_v2_gdn_cached_replay_cute__verify.py
index 98b9e62c5284..d7396fe75f7d 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ssm/test_modeling_v2_gdn_cached_replay_cute__verify.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ssm/test_modeling_v2_gdn_cached_replay_cute__verify.py
@@ -22,6 +22,7 @@
import torch
import tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.ssm.gdn_cached_replay_cute__verify as cr # noqa: E501
+from tensorrt_llm._torch.modules.fla.cached_replay import CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE
_MV2_TESTS = Path(__file__).resolve().parents[4]
if str(_MV2_TESTS) not in sys.path:
@@ -275,3 +276,27 @@ def test_launch_mappings_agree(n, t, pdl) -> None:
for cfg, got in zip(_MAPPINGS[1:], results[1:]):
for name, x, y in zip(("out", "state", "u", "k", "G"), got, results[0]):
assert torch.equal(x, y), f"{name}: {cfg} vs {_MAPPINGS[0]}"
+
+
+@pytest.mark.parametrize("heads", cr.GEOMETRIES, ids=[f"{h}x{hv}" for h, hv in cr.GEOMETRIES])
+def test_precompile_covers_every_step(heads) -> None:
+ """After `precompile`, as the target runs it at bind, a verify step of any request count
+ compiles nothing: steps at each launch mapping's edges, 9 to 15 requests (an eager mixed
+ step's, which no padded decode graph captures) among them, the partitioned ones with work
+ items and the all-layer commit as the target hands them over."""
+ t, partition_from = 4, CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE
+ kernels = cr.precompile(t, *heads, 40, pdl=True, partition_from=partition_from)
+ assert kernels > 0, kernels
+ compiled = len(cr._COMPILED)
+ for n in (1, 2, 3, 8, 9, 12, 15, 16, 17, 40):
+ path = "all_layer" if n >= partition_from else "small"
+ spec = _rank_spec(heads, path=path, n=n, t=t, steps=1, layers=1, slots=64)
+ bench = trep._Bench(spec, 5150 + n, exact=False)
+ bench.prepare()
+ act = bench.inputs()
+ kw = bench.kwargs(0, act)
+ kw["g"], kw["beta"] = _gates(bench, act, True)
+ kw["launch_with_pdl"] = True
+ op(**kw)
+ torch.cuda.synchronize()
+ assert len(cr._COMPILED) == compiled, f"a step of {n} requests compiled a kernel"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision_warm_up.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision_warm_up.py
new file mode 100644
index 000000000000..73a73c4d1d78
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision_warm_up.py
@@ -0,0 +1,108 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""GPU test of the default target's vision tower warm-up: after `VisionTower.warm_up`, which
+`post_load_weights` runs, a media request of any grid compiles no Triton kernel.
+
+The tower is one block deep (the block's kernels are every block's) over random weights at the
+checkpoint's shapes, built from the weight table's own `vis_*` declarations.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501
+ vision as V,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501
+ weights as W,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.attention.vision_rope_2d_packed_qk import ( # noqa: E501
+ _rotary_neox_inplace_kernel,
+)
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.vision.triton_pos_embed_interpolate import ( # noqa: E501
+ _bilinear_pos_embed_kernel,
+)
+from tensorrt_llm._torch.configs import Qwen4ExpVisionConfig
+from tensorrt_llm._torch.modules.fla.layernorm_gated import _layer_norm_fwd_1pass_kernel
+
+assert torch.cuda.is_available(), "the vision tower runs on a CUDA device"
+
+# Qwen3.8-Flash-Next's vision config, one block deep.
+_VISION = dict(
+ deepstack_visual_indexes=[],
+ depth=1,
+ dtype="bfloat16",
+ hidden_act="gelu_pytorch_tanh",
+ hidden_size=1152,
+ in_channels=3,
+ intermediate_size=4304,
+ num_heads=16,
+ num_position_embeddings=2304,
+ out_hidden_size=2560,
+ patch_size=16,
+ spatial_merge_size=2,
+ temporal_patch_size=2,
+)
+HIDDEN_OUT = 2560
+
+#: Every Triton kernel the tower launches.
+_KERNELS = {
+ "rope": _rotary_neox_inplace_kernel,
+ "pos_embed": _bilinear_pos_embed_kernel,
+ "layer_norm": _layer_norm_fwd_1pass_kernel,
+}
+
+
+def _tower() -> V.VisionTower:
+ vcfg = Qwen4ExpVisionConfig(**_VISION)
+ d = SimpleNamespace(
+ hidden=HIDDEN_OUT,
+ dtype=torch.bfloat16,
+ v_depth=vcfg.depth,
+ v_hidden=vcfg.hidden_size,
+ v_inter=vcfg.intermediate_size,
+ v_channels=vcfg.in_channels,
+ v_tpatch=vcfg.temporal_patch_size,
+ v_patch=vcfg.patch_size,
+ v_pos_rows=vcfg.num_position_embeddings,
+ v_merge_dim=vcfg.hidden_size * vcfg.spatial_merge_size**2,
+ )
+ gen = torch.Generator(device="cuda").manual_seed(7)
+ w = {"embed": torch.empty(8, HIDDEN_OUT, dtype=torch.bfloat16, device="cuda")}
+ for entry in W.MODEL_WEIGHTS.WEIGHTS:
+ if not entry.name.startswith("vis_"):
+ continue
+ for key, _, _ in W.MODEL_WEIGHTS._targets(entry, d):
+ shape, dtype = entry.shape(d), entry.dtype or d.dtype
+ w[key] = (torch.randn(*shape, generator=gen, device="cuda") * 0.02).to(dtype)
+ return V.VisionTower(w, vcfg)
+
+
+def _variants() -> dict:
+ """Compiled variants per kernel on this device (through any `heuristics` wrappers)."""
+ device = torch.cuda.current_device()
+ counts = {}
+ for name, kernel in _KERNELS.items():
+ while not hasattr(kernel, "device_caches"):
+ kernel = kernel.fn
+ counts[name] = len(kernel.device_caches[device][0])
+ return counts
+
+
+def test_after_warm_up_no_grid_compiles_a_kernel():
+ tower = _tower()
+ tower.warm_up()
+ torch.cuda.synchronize()
+ warmed = _variants()
+ assert all(warmed.values()), f"the warm-up compiled nothing: {warmed}"
+ gen = torch.Generator(device="cuda").manual_seed(11)
+ for grid in ([(1, 16, 16)], [(1, 6, 10)], [(2, 8, 12)], [(1, 62, 90)], [(1, 2, 34), (1, 4, 4)]):
+ rows = sum(t * h * w for t, h, w in grid)
+ pixels = torch.randn(rows, tower.pixel_width, generator=gen, device="cuda").to(tower.dtype)
+ out = tower(pixels, torch.tensor(grid, dtype=torch.int64))
+ torch.cuda.synchronize()
+ assert out.shape == (rows // 4, HIDDEN_OUT) and bool(torch.isfinite(out).all()), grid
+ assert _variants() == warmed, f"{grid} compiled a kernel after the warm-up"
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_weight_shards.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_weight_shards.py
index a25b790cd6f0..11a0dd88d3df 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_weight_shards.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_weight_shards.py
@@ -554,10 +554,10 @@ def _block(start: int, width: int) -> list[int]:
@pytest.mark.parametrize("rank", RANKS)
-@pytest.mark.parametrize("member", ["tep4", "dep4"])
+@pytest.mark.parametrize("member", ["tp4", "tep4", "dep4"])
def test_cuts_select_the_rank_heads(stored, configs, member, rank):
- """tp4 cuts attention and Gated-DeltaNet exactly as tep4 does: the members differ only
- in the routed experts."""
+ """Each rank's attention and Gated-DeltaNet rows: tp4 and tep4 cut them alike (the
+ members differ only in the routed experts), dep4 holds every head."""
d = _dims(configs, member, rank)
fake = _fake_checkpoint(stored)
manifest = WT.MODEL_WEIGHTS.manifest(d)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision/test_modeling_v2_triton_pos_embed_interpolate.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision/test_modeling_v2_triton_pos_embed_interpolate.py
index 6a8b63a8f50e..5c7e014517ba 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision/test_modeling_v2_triton_pos_embed_interpolate.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision/test_modeling_v2_triton_pos_embed_interpolate.py
@@ -17,6 +17,9 @@
import torch
import torch.nn.functional as F
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.vision.triton_pos_embed_interpolate import ( # noqa: E501
+ _bilinear_pos_embed_kernel,
+)
from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.vision.triton_pos_embed_interpolate import ( # noqa: E501
triton_pos_embed_interpolate as op,
)
@@ -114,3 +117,17 @@ def test_certified_cells(cell, certifying): # noqa: F811
for wrong in (dict(align=False), dict(raster=True), dict(swap=True)):
with pytest.raises(AssertionError):
_compare(out, _resize64(table, t, h, w, m_size, **wrong).float())
+
+
+def test_one_compile_serves_every_grid():
+ """The grid sides are not specialized: sides that are and are not multiples of 16 run
+ the one compiled kernel, each within the band."""
+ table = _table(41)
+ variants = []
+ for t, h, w in ((1, 2, 2), (1, 16, 16), (1, 4, 34), (2, 62, 90), (1, 32, 6)):
+ args = (table, t, h, w, NUM_GRID, MERGE, torch.bfloat16)
+ _compare(op(*args), op.reference(*args))
+ variants.append(
+ len(_bilinear_pos_embed_kernel.device_caches[torch.cuda.current_device()][0])
+ )
+ assert variants == [variants[0]] * len(variants), variants
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py
index 2392877d0746..81d48a46f2ac 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py
@@ -355,9 +355,11 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res
def test_has_initial_state_is_the_reset_then_the_plain_call() -> None:
- """The flags are zeroing the fresh slots before the call: zeroing them and running the
- plain in-place call gives the same output and pool, bitwise; a linear gate that is
- `torch.exp` of the plain call's log gate changes nothing."""
+ """On the packed path (above `INPLACE_MAX_TOKENS`), the flags are zeroing the fresh slots
+ before the call: zeroing them and running the plain in-place call gives the same output
+ and pool, bitwise; a linear gate that is `torch.exp` of the plain call's log gate changes
+ nothing."""
+ assert sum(_CARRIED.spec["seqlens"]) > INPLACE_MAX_TOKENS, "the cell must take the packed path"
args, _ = _build(_CARRIED.spec, seed=5)
flags = args.pop("has_initial_state")
linear_gate = args.pop("linear_gate", False)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py
index 23365a9c97b3..ec234f839502 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py
@@ -226,8 +226,10 @@ def _build(spec, seed):
# and slot 100 sits ~2.8G elements past the base (past int32).
layers, num = PAGE_LAYERS, 102
page = torch.empty(num * layers * v_heads * V * K, device=DEV, dtype=state)
- page[: v_heads * V * K * 4].normal_(0.0, 0.05, generator=g)
pool = page.as_strided((num, v_heads, V, K), (layers * v_heads * V * K, V * K, K, 1))
+ # The pool is compared whole: no slot of the view is left uninitialized.
+ pool.zero_()
+ page[: v_heads * V * K * 4].normal_(0.0, 0.05, generator=g)
for s in (100, 101):
pool[s].copy_(torch.randn(v_heads, V, K, generator=g, device=DEV) * 0.05)
slots = torch.tensor([100, 101][:rows], dtype=torch.int32, device=DEV)
From cc1a94681946130f9d7d523be8f02db2adfa23fc Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Thu, 5 Nov 2026 00:34:15 -0800
Subject: [PATCH 07/15] [None][feat] modeling_v2 qwen4_exp: FlashInfer ring
replay, NVFP4 KV without paged context, quantized KV at tp4/tep4
- Routing declines a checkpoint of the target's shape whose config changes
what it computes: norm_topk_prob, top-k, hidden_act, attention bias, the
output gate, the rope settings or tied embeddings.
- The verify step runs FlashInfer's ring replay when the cache manager lays
the replay history out for it (TRTLLM_USE_GDN_FLASHINFER_REPLAY).
- An NVFP4 KV cache takes the paged-context attention path without chunked
prefill or block reuse.
- Quantized KV at tp4/tep4 (6 query heads a KV head): cells for the FP8
sparse GQA, the NVFP4 prefill and verify, the e4m3 decode prologue and the
KV stores at one KV head; the quantized-KV accuracy case runs pass keys too,
with tep4 FP8 (ring replay) and tep4 NVFP4 (no paged context) members.
- A whole-layer Gated-DeltaNet test against an fp32 reference at 16/48 and
4/12 heads.
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
...-guide-for-qwen3.8-flash-next-on-trtllm.md | 1 +
.../ops/attention/qsa_paged_kv_store.py | 23 +-
.../ops/attention/qsa_paged_kv_store_nvfp4.py | 20 +
.../ops/attention/qsa_paged_sparse_gqa.py | 38 +-
.../attention/qsa_paged_sparse_gqa_nvfp4.py | 28 +-
.../qsa_paged_sparse_gqa_nvfp4_verify.py | 30 ++
.../modeling.py | 110 +++++-
.../qwen4_exp_attn_decode_prologue__decode.py | 11 +
.../modeling_v2/models/qwen4_exp/routing.py | 51 +++
tests/integration/defs/.test_durations | 2 +
.../accuracy/modeling_v2/test_qwen4_exp.py | 50 ++-
.../references/acceptance_length.yaml | 9 +
.../test_lists/qa/llm_function_core.txt | 2 +
.../test-db/l0_gb300_multi_gpus.yml | 2 +
...ng_v2_qsa_paged_sparse_gqa_nvfp4_verify.py | 9 +-
.../ssm/test_modeling_v2_gdn_block.py | 349 ++++++++++++++++++
.../test_routing.py | 87 +++++
.../modeling_v2/test_modeling_v2_routing.py | 13 +
18 files changed, 791 insertions(+), 44 deletions(-)
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ssm/test_modeling_v2_gdn_block.py
diff --git a/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md
index 98c05da231b7..d8a111121b32 100644
--- a/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md
+++ b/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md
@@ -220,6 +220,7 @@ It supports these single-node layouts:
Under `auto`, these settings make the server fall back to the built-in model; under `require`, the server fails at startup:
* Any other layout, including `enable_lm_head_tp_in_adp: true` with attention DP.
+* A checkpoint of the same shape whose config changes what it computes, for example `norm_topk_prob: false`, another `num_experts_per_tok`, or other rope settings.
* Speculative decoding other than one-model MTP, or a draft length that changes at runtime.
* MTP with `TRTLLM_USE_GDN_REPLAY=0`.
* A KV cache dtype other than the checkpoint's bf16, `fp8` or `nvfp4`.
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_kv_store.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_kv_store.py
index 5d5af38707ad..e061883c1236 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_kv_store.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_kv_store.py
@@ -54,6 +54,8 @@
#: A tp4 / tep4 rank: one KV head, k and v at heads 6 and 7 of the fused projection row
#: `[q 6 | k 1 | v 1 | gate 6 | index 3]` (17 heads of 256, row stride 4352).
_TP4 = dict(_ENGINE, hkv=1, heads_q=6, row_pad_heads=9, row_stride=4352)
+#: The same rank under `kv_cache_config.dtype: fp8`: one KV head into an e4m3 pool.
+_TP4_FP8 = dict(_TP4, cache_dtype=torch.float8_e4m3fn)
#: A dep4 rank (attention DP, every head): the tp1 row `[q 24 | k 2 | v 2 | gate 24 | index 3]`
#: (55 heads, row stride 14080) over the kv_fp8 variant's e4m3 pool.
@@ -152,6 +154,11 @@ class QsaPagedKvStore(OpWrapper):
why="fp8 decode MTP7: 32 requests x 8 rows (256), blocks straddling pages",
spec=_spec(_MTP7_32, pages=1200, slots=40, cache_dtype=_F8),
),
+ Cell(
+ why="fp8 MTP7 at a tp4 / tep4 rank's draft layer: 32 requests x 8 rows of the packed "
+ "[6 + 1 + 1] row, one KV head",
+ spec=_spec(_MTP7_32, pages=1200, slots=40, cache_dtype=_F8, hkv=1, heads_q=6),
+ ),
# The default target's members over the engine's own pool: one rank's page count, the
# 24-role-page stride, layer 47's view, a 128-row table.
Cell(
@@ -172,6 +179,20 @@ class QsaPagedKvStore(OpWrapper):
"request's first 1023",
spec=_engine(_TP4, "tp4", _PREFILL_8192, PAGES_TEP4),
),
+ Cell(
+ why="tep4 kv_fp8 decode, 128 rows: hkv 1 off the 4352-wide row into an e4m3 pool "
+ "of tep4's 257109 pages (a kv_fp8 engine holds more)",
+ spec=_engine(_TP4_FP8, "tep4", _DECODE_128, PAGES_TEP4),
+ ),
+ Cell(
+ why="tep4 kv_fp8 MTP7 verify, 128 requests x 8 rows (1024) to position 16383: hkv 1 "
+ "into an e4m3 pool, table width 516",
+ spec=_engine(_TP4_FP8, "tep4", _MTP7_128, PAGES_TEP4, width=516),
+ ),
+ Cell(
+ why="tp4 kv_fp8 context, 8192 rows (max_num_tokens): hkv 1 into an e4m3 pool",
+ spec=_engine(_TP4_FP8, "tp4", _PREFILL_8192, PAGES_TEP4),
+ ),
Cell(
why="dep4 kv_fp8 decode, 32 rows: hkv 2 off the 14080-wide row into a "
"171362-page e4m3 pool",
@@ -228,7 +249,7 @@ class QsaPagedKvStore(OpWrapper):
Not certified: duplicate destinations, multiple devices, a scaled fp8
pool (there is no scale argument), fp8 conversion of values past 448, the
- packed-QKV-row (MTP draft layer) layout over a bf16 pool or at one KV head,
+ packed-QKV-row (MTP draft layer) layout over a bf16 pool,
fp16 rows, head geometry other than D 256 with Hkv 1 or 2 and 32-token
pages, int32 positions or int64 request indices, head-major rows, sliced
or transposed tables, `rows == 0`, and of the rows that resolve to no slot
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_kv_store_nvfp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_kv_store_nvfp4.py
index 719c548b5fd2..74022d993a8f 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_kv_store_nvfp4.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_kv_store_nvfp4.py
@@ -74,6 +74,16 @@
)
+#: One tp4 / tep4 rank's kv_nvfp4 pool: one KV head (k and v at heads 6 and 7 of the packed
+#: `[q 6 | k 1 | v 1]` row), the trunk layout and a `[128, 516]` table (MTP's max_seq_len
+#: 16384). A slot's 24 role pages of one KV head are half dep4's, so the same pool bytes hold
+#: twice the slots: 426590 puts the stores past 2^31 bytes of both planes.
+_TP4_PAGES = 426590
+_TP4 = dict(
+ hkv=1, heads_q=6, slots=128, width=516, slot_pages=24, layer=11, last_page=0, large=True
+)
+
+
def _spec(segments, pages, **overrides) -> dict:
"""`segments` is `(first position, rows)` per request, in row order."""
return {**_CKPT, "segments": segments, "pages": pages, **overrides}
@@ -414,6 +424,16 @@ class QsaPagedKvStoreNvfp4(OpWrapper):
"across page edges, on the rank's 213295-slot pool",
spec=_spec(_DEP4_DECODE, pages=_DEP4_PAGES, **_DEP4),
),
+ Cell(
+ why="tep4 kv_nvfp4 decode: 32 requests at positions 0..16383, one KV head, the last "
+ "request's last page on the pool's last slot",
+ spec=_spec(((16383, 1),) + _DEP4_DECODE[1:], pages=_TP4_PAGES, **_TP4),
+ ),
+ Cell(
+ why="tp4 kv_nvfp4 context chunk: 8192 rows (max_num_tokens) at one KV head, a prompt's "
+ "chunk to position 16383 and another's first 1024",
+ spec=_spec(((9216, 7168), (0, 1024)), pages=_TP4_PAGES, **_TP4),
+ ),
Cell(
why="dep4 kv_nvfp4 context chunk: 8192 rows (max_num_tokens), one prompt's chunk at "
"2048..9215 and a new prompt's first 1024, on the rank's 213295-slot pool",
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa.py
index b65b55eac0cf..1b96989177e4 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa.py
@@ -74,6 +74,9 @@ def _decode(segments, launches=_SPLIT_K, **overrides) -> dict:
#: 516 at MTP's 16384.
_TP4 = dict(q_heads=6, kv_heads=1, slots=128, pages=292, **_TRUNK)
_TP4_MTP = dict(_TP4, pages=516)
+#: The same ranks under `kv_cache_config.dtype: fp8`: an e4m3 pool at 6 q heads a KV head.
+_TP4_FP8 = dict(_TP4, **_FP8)
+_TP4_MTP_FP8 = dict(_TP4_MTP, **_FP8)
#: One dep4 rank (attention DP): tp1's 24 / 2 heads and a 32-row table (128 requests over 4
#: ranks); a bf16 pool, or the kv_fp8 variant's e4m3 pool.
_DEP4 = dict(q_heads=24, kv_heads=2, slots=32, pages=292, **_TRUNK)
@@ -251,6 +254,35 @@ class QsaPagedSparseGqa(OpWrapper):
**_TP4_ROW,
),
),
+ Cell(
+ why="tep4 kv_fp8 gated decode 32 x 1 row: 6 q / 1 kv head over an e4m3 pool, q and "
+ "the gate off the 17-head fused row, 32-way split-K + gated merge",
+ spec=_default(_spread(32, 1, 9215), _TP4_FP8, _SPLIT_K, splits=32, **_TP4_ROW),
+ ),
+ Cell(
+ why="tep4 kv_fp8 MTP3 gated verify 128 x 4 rows = 512: 6 q / 1 kv head over an e4m3 "
+ "pool, 4-way split-K + gated merge, [128, 516] table, rows to position 16383",
+ spec=_default(
+ _spread(128, 4, 16380),
+ _TP4_MTP_FP8,
+ _SPLIT_K,
+ splits=4,
+ rows_per_request=4,
+ **_TP4_ROW,
+ ),
+ ),
+ Cell(
+ why="tep4 kv_fp8 8192-row context step with the prefill pair: 6 q / 1 kv head over "
+ "an e4m3 pool on the prefill kernel, q off the 17-head fused row",
+ spec=_default(
+ ((0, 8065),) + _spread(127, 1, 9215),
+ _TP4_FP8,
+ _PREFILL,
+ prefill_pair=True,
+ q_layout="fused_row",
+ row_heads=17,
+ ),
+ ),
Cell(
why="dep4 core gated decode 32 x 1 row: 24 q / 2 kv heads over a bf16 pool, q and "
"the gate off the 55-head (14080-wide) fused row, 8-way split-K + gated merge",
@@ -367,9 +399,9 @@ class docstring). QSA's producer -- `qsa_topk_select` with zero
one-split launch (the 1024-row MTP7 verify, which allocates only its
output), of the 512-row MTP3 verify and of the prefill kernel; dep4's
gated decode below 16 rows (its 64- and 32-way splits at 24/2); an fp8
- pool at 6/1 (no tp4 / tep4 member stores one) and on the prefill kernel
- with q off the 55-head row (the fp8 context cell reads the packed
- `[q | k | v]` slice); positions past 16383; rows whose selection
+ pool on the prefill kernel with q off the 55-head row (the dep4 fp8
+ context cell reads the packed `[q | k | v]` slice; the tep4 one, at 6/1,
+ reads q off the 17-head row); positions past 16383; rows whose selection
resolves no token (all padding, a foreign
request, unmapped pages); an fp16 pool or q; standalone contiguous K/V
buffers; int64 request indices. Hard limits: not probed; the largest
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4.py
index 734370a39e0b..d16d4658d55c 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4.py
@@ -83,6 +83,26 @@ def _decode(segments, launches=_SPLIT_K, **overrides) -> dict:
_OFFSET_EDGE = 2**31
+#: One tp4 / tep4 rank's engine under `kv_nvfp4`: 6 q heads and one of the 2 KV heads,
+#: `max_batch_size` 128 and MTP's `max_seq_len` 16384 give a `[128, 516]` table, the same trunk
+#: layout. A page of layer 11's scale plane advances 24 x 1 x 32 x 16 bytes, so 2^31 bytes lies
+#: at page 174763 (the data plane's at 21846).
+_TP4 = dict(q_heads=6, kv_heads=1, pages=516, slots=128, slot_pages=24, layer=11)
+
+
+def _tp4_context(segments, pool_pages) -> dict:
+ """A context step of a tp4 / tep4 rank's kv_nvfp4 engine: the prefill pair, the prefill
+ kernel, at 6 q heads a KV head."""
+ return _decode(
+ segments,
+ launches=_PREFILL,
+ prefill_pair=True,
+ pool_pages=pool_pages,
+ offset_edge=_OFFSET_EDGE,
+ **_TP4,
+ )
+
+
def _dep4_context(segments, pool_pages) -> dict:
"""A context step of dep4's kv_nvfp4 engine: the prefill pair, the prefill kernel."""
return _decode(
@@ -970,6 +990,12 @@ class QsaPagedSparseGqaNvfp4(OpWrapper):
"selection budget; 87401 pages (not a multiple of 16)",
spec=_dep4_context(((8960, 256), (0, 160), (2000, 64), (5000, 32)), pool_pages=87401),
),
+ Cell(
+ why="tp4/tep4 8192-row context chunk at 6 q / 1 kv head: a request at 8320..16383 "
+ "(MTP's max_seq_len, table column 511) and one from 0, [128, 516] table, trunk "
+ "layout, pages straddling 2^31 bytes of the scale plane; 174900 pages",
+ spec=_tp4_context(((8320, 8064), (0, 128)), pool_pages=174900),
+ ),
)
note = """
@@ -1048,7 +1074,7 @@ class QsaPagedSparseGqaNvfp4(OpWrapper):
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;
+ product. Not certified: head geometries other than 24/2 and 6/1 x 256;
`tokens_per_block` other than 32; more than 8192 rows; capture of the
context/mixed site; a pool of one slot; second-level scales other than 1.0
(the target binds ones); the prefill pair on the split-K kernel (256 context
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4_verify.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4_verify.py
index 4b4aec3eff8f..9e0eb4ec42db 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4_verify.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4_verify.py
@@ -1960,6 +1960,32 @@ class QsaPagedSparseGqaNvfp4Verify(OpWrapper):
"into its store, the last wave's units split two ways and merged",
spec=_verify(tuple((2100 + (37 * i) % 1900, 4) for i in range(256)), gated=True),
),
+ Cell(
+ why="tep4 MTP3 verify, 32 requests x 4 rows at a rank's 6 q / 1 kv head: 24 MMA rows "
+ "a group, gate fused, the MTP engine's [128, 516] table",
+ spec=_verify(
+ tuple((2100 + (53 * i) % 1900, 4) for i in range(32)),
+ gated=True,
+ q_heads=6,
+ kv_heads=1,
+ slots=128,
+ pages=516,
+ ),
+ ),
+ Cell(
+ why="tep4 MTP3 verify at 512 rows (128 requests x 4, the largest captured batch) at "
+ "6 q / 1 kv head, gate fused, trunk layout",
+ spec=_verify(
+ tuple((2100 + (37 * i) % 14000, 4) for i in range(128)),
+ gated=True,
+ q_heads=6,
+ kv_heads=1,
+ slots=128,
+ pages=516,
+ slot_pages=24,
+ layer=11,
+ ),
+ ),
Cell(
why="groups spanning requests (3, 4, 4, 1, 4 x 5, 2 rows: misaligned and partial "
"groups), early positions, int64 request indices, non-unit global scales, trunk "
@@ -2003,6 +2029,10 @@ class QsaPagedSparseGqaNvfp4Verify(OpWrapper):
CTA each (bf16 output, the gate fused) and the units of the partial last wave two to four,
only their rows merged.
+ Head geometries: tp1 / dep4's 24 q over 2 KV heads and a tp4 / tep4 rank's 6 over 1 (24
+ MMA rows a group at T = 4), each held by the cells against `reference` and the per-row
+ kernel.
+
GB300, one layer, cold L2, T = 4, the decode step's call (output gate on) on MTP3 verify
selections dumped from a served run (a request's rows' union 1.4-1.7x a row's selection,
about 8.4K tokens of context, the server's pool geometry): 64 rows 29.7 us, 128 rows 46.1,
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py
index d76b8a91a1cc..5388a8e9f36f 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py
@@ -428,6 +428,10 @@
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.fla.flashinfer_cached_replay import (
+ FLASHINFER_REPLAY_RING_SIZE,
+ flashinfer_cached_replay_update,
+)
from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor
from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks
from tensorrt_llm._torch.speculative import get_spec_worker, should_use_separate_draft_kv_cache
@@ -921,13 +925,33 @@ def _moe_layer_scalars_cd(w, i: int) -> dict[str, torch.Tensor]:
class _Replay(NamedTuple):
- """The GDN cached replay's history buffers and bookkeeping of one manager."""
+ """The GDN cached replay's history buffers and bookkeeping of one manager: two history
+ halves a slot (`old`, the Triton / CuTe replay), or, under FlashInfer's ring replay
+ (`TRTLLM_USE_GDN_FLASHINFER_REPLAY`), a 32-row ring a slot (`ring`, `cache_base`)."""
- old: dict # GDN layer -> (old_x, old_B, old_dt, old_dA_cumsum)
+ old: dict # GDN layer -> (old_x, old_B, old_dt, old_dA_cumsum); empty under the ring
cache_buf_idx: torch.Tensor # [slots] int32: the history half each slot writes next
prev_num_accepted: torch.Tensor # [slots] int32: tokens pending since the checkpoint
history: int # rows a history half holds
width: int # the one verify width a step may carry per generation request
+ ring: dict | None = None # GDN layer -> (k_ring, u_ring, g_ring), FlashInfer's layout
+ cache_base: torch.Tensor | None = None # [slots] int32: each slot's ring start
+
+
+def _check_replay_ring(g, layer: int, slots: int, k_ring, u_ring, g_ring) -> None:
+ """One GDN layer's FlashInfer ring buffers at this rank's heads: `[slots, H, 32, K]`,
+ `[slots, HV, 32, V]` in the state's bf16 and `[slots, HV, 32]` fp32 cumulative gates."""
+ ring = FLASHINFER_REPLAY_RING_SIZE
+ shapes = (
+ (k_ring, (slots, g.num_k_heads, ring, g.head_k_dim), torch.bfloat16),
+ (u_ring, (slots, g.num_v_heads, ring, g.head_v_dim), torch.bfloat16),
+ (g_ring, (slots, g.num_v_heads, ring), torch.float32),
+ )
+ for t, shape, dtype in shapes:
+ assert t is not None and tuple(t.shape) == shape and t.dtype == dtype, (
+ f"layer {layer}: a FlashInfer replay ring of {None if t is None else tuple(t.shape)} "
+ f"{None if t is None else t.dtype}; this rank's heads need {shape} {dtype}"
+ )
class _Bound(NamedTuple):
@@ -1567,10 +1591,6 @@ def _bind_cache_manager(self, md, position_ids) -> None:
f"KV addressing over {len(pools)} pools is not certified; layer->pool ids "
f"{sorted(pools)}"
)
- assert self.quant_mode != QM_NVFP4_KV_CACHE or md.use_paged_context_fmha, (
- "an NVFP4 KV cache is read by the paged-context attention kernels only, which the "
- "engine enables for chunked prefill, KV cache block reuse or speculative decoding"
- )
g = _pool.pool_geometry(
self.model_config.pretrained_config,
self.mapping,
@@ -1677,10 +1697,6 @@ def _bind_step_buffers(self, cm, g: _pool.PoolGeometry, mm, pools: _pool.RankPoo
"(TRTLLM_USE_GDN_REPLAY=0), which this target does not assemble; leave "
"TRTLLM_USE_GDN_REPLAY unset for the cached replay"
)
- assert not getattr(cm, "use_gdn_flashinfer_ring_replay", False), (
- "the manager laid its replay history out for the FlashInfer ring kernel; this "
- "target runs the double-buffered Triton cached replay"
- )
assert layer_caches[g0].intermediate_ssm is None, (
"the cache manager runs the cached-replay verify path and also allocated a "
"per-step recurrent buffer; the runtime allocates one or the other"
@@ -1690,6 +1706,27 @@ def _bind_step_buffers(self, cm, g: _pool.PoolGeometry, mm, pools: _pool.RankPoo
_pool.check_replay_state(
history, width, meta.prev_num_accepted_tokens, meta.cache_buf_idx, ssm0
)
+ if getattr(cm, "use_gdn_flashinfer_ring_replay", False):
+ # FlashInfer's ring replay (the manager chose it at load: MTP with 3 or 7 drafts
+ # over a bf16 state): a 32-row ring of history a slot and layer, folded in-kernel.
+ ring = {}
+ for i, cache in layer_caches.items():
+ ring[i] = (cache.k_ring, cache.u_ring, cache.g_ring)
+ _check_replay_ring(g, i, ssm0.shape[0], *ring[i])
+ cache_base = layer_caches[g0].cache_base
+ assert cache_base is not None and tuple(cache_base.shape) == (ssm0.shape[0],), (
+ "the ring replay's per-slot start is missing or not one entry per slot"
+ )
+ replay = _Replay(
+ {},
+ meta.cache_buf_idx,
+ meta.prev_num_accepted_tokens,
+ history,
+ width,
+ ring,
+ cache_base,
+ )
+ return dict(int_idx=int_idx, int_conv=int_conv, replay=replay)
old = {}
for i, cache in layer_caches.items():
old[i] = (cache.old_x, cache.old_B, cache.old_dt, cache.old_dA_cumsum)
@@ -2226,9 +2263,12 @@ def _verify_state(core, mm, num_decodes: int) -> tuple:
`update_mamba_states` promotes the conv windows from, and the replay's work items. Work
items go to the kernel from `CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE` requests on, and the
kernel leaves the history commit to the manager's all-layer commit exactly when the
- manager runs one after the step -- the stock Gated-DeltaNet layer's two conditions."""
+ manager runs one after the step -- the stock Gated-DeltaNet layer's two conditions. The
+ FlashInfer ring replay folds in-kernel and reads no work items."""
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.ring is not None:
+ return int_idx, None, None, False
partitioned = num_decodes >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE
all_layer = partitioned and bool(
getattr(core._bound.cm, "use_gdn_cached_replay_all_layer_commit", False)
@@ -2260,7 +2300,8 @@ def _gdn_verify(core, mq, a, b, i, L, si_d, n, p, verify, mq_out=None) -> torch.
replay = core._bound.replay
ssm = pools.ssm[i]
cute = ssm.dtype == torch.bfloat16
- pdl = _GDN_REPLAY_PDL and cute
+ # Only the CuTe replay is launched as the conv's programmatic dependent.
+ pdl = _GDN_REPLAY_PDL and cute and replay.ring is None
if mq_out is None:
# [n*p, C] -> [n, C, p]: the multi-token conv reads channel-major.
x3 = torch.transpose(torch.reshape(mq, [n, p, core.conv_dim]), 1, 2)
@@ -2280,8 +2321,32 @@ def _gdn_verify(core, mq, a, b, i, L, si_d, n, p, verify, mq_out=None) -> torch.
)
# Back to token-major: one contiguous copy, which the kernel reads q/k/v from directly.
mq_out = torch.reshape(torch.transpose(conv_out, 1, 2), [n * p, core.conv_dim])
- old_x, old_B, old_dt, old_dA = replay.old[i]
qv, kv, vv = torch.split(mq_out, [core.key_dim, core.key_dim, core.value_dim], 1)
+ if replay.ring is not None:
+ # FlashInfer's ring replay reads q/k/v as column slices of `mq_out` (one token stride)
+ # and `a` / `b` through their row stride, and the per-row history length and ring
+ # start gathered by slot; the manager advances both after the step.
+ k_ring, u_ring, g_ring = replay.ring[i]
+ slots = si_d.long()
+ out = flashinfer_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]),
+ a=torch.reshape(a, [n, p, core.num_v_heads]),
+ b=torch.reshape(b, [n, p, core.num_v_heads]),
+ A_log=L.a_log,
+ dt_bias=L.dt_bias,
+ ssm_states=ssm,
+ state_indices=si_d,
+ k_ring=k_ring,
+ u_ring=u_ring,
+ g_ring=g_ring,
+ hist_len=replay.prev_num_accepted[slots],
+ cache_base=replay.cache_base[slots],
+ history_size=replay.history,
+ )
+ return torch.reshape(out, [1, n * p, core.num_v_heads, core.head_v_dim])
+ old_x, old_B, old_dt, old_dA = replay.old[i]
if cute:
replay_op = gdn_cached_replay_cute
else:
@@ -3506,9 +3571,11 @@ def _copy_positions(position_ids: torch.Tensor) -> torch.Tensor:
return pos3
-def _step_args(md: TrtllmAttentionMetadata, num_contexts, num_ctx_tokens) -> dict:
+def _step_args(md: TrtllmAttentionMetadata, num_contexts, num_ctx_tokens, paged_context) -> dict:
"""thop_attention's batch state off the prepared metadata. Tensors are engine-owned
- persistent buffers refreshed in place; Python ints are per-capture constants."""
+ persistent buffers refreshed in place; Python ints are per-capture constants.
+ `paged_context` forces the paged-context path the engine enables only for chunked prefill
+ or KV cache block reuse: an NVFP4 pool's context kernels exist on it alone."""
return dict(
sequence_length=md.kv_lens_cuda_runtime,
host_past_key_value_lengths=md.kv_lens_runtime,
@@ -3529,7 +3596,7 @@ def _step_args(md: TrtllmAttentionMetadata, num_contexts, num_ctx_tokens) -> dic
num_contexts=num_contexts,
num_ctx_tokens=num_ctx_tokens,
trtllm_gen_jit_warmup=md.trtllm_gen_jit_warmup,
- use_paged_context_fmha=md.use_paged_context_fmha,
+ use_paged_context_fmha=md.use_paged_context_fmha or paged_context,
beam_width=md.effective_beam_width,
cache_indirection=md.cache_indirection,
block_ids_per_seq=md.block_ids_per_seq,
@@ -3659,7 +3726,9 @@ def _check_step_contract(self, md) -> None:
def step_args(self, md: TrtllmAttentionMetadata) -> dict:
"""Project the prepared metadata onto thop_attention's batch state."""
- return _step_args(md, md.num_contexts, md.num_ctx_tokens)
+ return _step_args(
+ md, md.num_contexts, md.num_ctx_tokens, self.core.quant_mode == QM_NVFP4_KV_CACHE
+ )
# --------------------------------------------------------------- blocks ---
@@ -4634,7 +4703,12 @@ def _check_step_contract(self, md) -> None:
def step_args(self, md: TrtllmAttentionMetadata) -> dict:
"""See PrefillTarget.step_args. Context counts are stated, not read: routing
guarantees them, and a decode-only CUDA graph would freeze whatever it read."""
- return _step_args(md, num_contexts=0, num_ctx_tokens=0)
+ return _step_args(
+ md,
+ num_contexts=0,
+ num_ctx_tokens=0,
+ paged_context=self.core.quant_mode == QM_NVFP4_KV_CACHE,
+ )
# --------------------------------------------------------------- blocks ---
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qwen4_exp_attn_decode_prologue__decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qwen4_exp_attn_decode_prologue__decode.py
index ff83f898bf92..bc5f535b8256 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qwen4_exp_attn_decode_prologue__decode.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qwen4_exp_attn_decode_prologue__decode.py
@@ -432,6 +432,17 @@ class Qwen4ExpAttnDecodePrologue(OpWrapper):
"masks its eighth head)",
spec=dict(_TP4, pos=tuple(1000 + 113 * r for r in range(72))),
),
+ Cell(
+ why="tp4/tep4 rank, decode 32 requests, e4m3 pool (a rank of a kv_fp8 engine)",
+ spec=dict(_TP4, pos=tuple(2000 + 151 * r for r in range(32)), fp8=True),
+ ),
+ Cell(
+ why="tp4/tep4 rank, e4m3 pool, MTP-3 verify 10 requests x 4 rows, a group end among "
+ "them",
+ spec=dict(
+ _TP4, pos=tuple(3000 + 211 * r for r in range(10)), rows_per_request=4, fp8=True
+ ), # fmt: skip
+ ),
Cell(
why="dep4 rank, decode 16 requests at 20k..94k (ht's step), e4m3 pool",
spec=dict(_DEP4, pos=_HT16, fp8=True),
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 1dcdda18b0ab..a7b1d7a36598 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
@@ -37,6 +37,31 @@
# also matches fine-tunes; the accuracy gate pins identity.
_TEXT_SHAPE = (48, 2560, 512, 24, 48, 248320)
+# The checkpoint's settings that change what it computes without changing a weight's shape,
+# at the values the target implements: a fine-tune of the same shape that changed one would
+# load and run with wrong logits and no error. Read off the text config as the target reads
+# them (`_text_settings`).
+_TEXT_SETTINGS = {
+ # The routed experts: Renormalize routing over the top 10 and a SiLU activation.
+ "norm_topk_prob": True,
+ "num_experts_per_tok": 10,
+ "hidden_act": "silu",
+ # Attention: two KV heads of 256, no projection bias, a sigmoid output gate, and an output
+ # head of its own.
+ "num_key_value_heads": 2,
+ "head_dim": 256,
+ "attention_bias": False,
+ "output_gate_type": "sigmoid",
+ "tie_word_embeddings": False,
+ # Rope: a quarter of the head (64 dims), theta 1e7, interleaved mRoPE over sections
+ # 11 / 11 / 10, unscaled.
+ "rope_theta": 1e7,
+ "partial_rotary_factor": 0.25,
+ "mrope_section": (11, 11, 10),
+ "mrope_interleaved": True,
+ "rope_factor": None,
+}
+
# (depth, hidden, heads, merge, out_hidden, num_pos, deepstack_indexes, act).
# `deepstack_visual_indexes == ()` is the key member: a deepstack sibling has a merger
# list these targets lack.
@@ -110,6 +135,24 @@ def engine_choices(model_config: ModelConfig, llm_args: TorchLlmArgs) -> EngineC
)
+def _text_settings(text) -> dict:
+ """`_TEXT_SETTINGS`' entries as `text` declares them. The rope fields are the target's
+ reading: the `rope_parameters` dict, or the flattened `rope_scaling` with the theta and
+ the partial factor at the top level."""
+ rope = getattr(text, "rope_parameters", None) or getattr(text, "rope_scaling", None) or {}
+ settings = {name: getattr(text, name, None) for name in _TEXT_SETTINGS}
+ settings.update(
+ rope_theta=rope.get("rope_theta", getattr(text, "rope_theta", None)),
+ partial_rotary_factor=rope.get(
+ "partial_rotary_factor", getattr(text, "partial_rotary_factor", None)
+ ),
+ mrope_section=tuple(rope.get("mrope_section") or ()),
+ mrope_interleaved=rope.get("mrope_interleaved"),
+ rope_factor=rope.get("factor"),
+ )
+ return settings
+
+
def layout(mapping: Mapping) -> Optional[str]:
"""The target's layout ``mapping`` describes (``tp1``, ``tp4``, ``tep4``, ``dep4``,
``dep2``), or None.
@@ -302,6 +345,14 @@ def route(
)
checkpoint = compute_capability == _SM and shape == _TEXT_SHAPE and vshape == _VISION_SHAPE
+ if checkpoint:
+ settings = _text_settings(text)
+ differs = {k: v for k, v in settings.items() if v != _TEXT_SETTINGS[k]}
+ if differs:
+ raise NoTarget(
+ f"the checkpoint has Qwen3.8-Flash-Next's shape but declares {differs}; the "
+ f"target implements { ({k: _TEXT_SETTINGS[k] for k in differs}) }"
+ )
# Qwen3.8-Flash-Next-NVFP4 on GB300s (sm103): on one at tp1, on four at tp4, tep4 or dep4,
# MTP-Eagle or no speculation; on two at dep2, without speculation.
member = layout(mapping)
diff --git a/tests/integration/defs/.test_durations b/tests/integration/defs/.test_durations
index 571426fce20f..a61a1390407b 100644
--- a/tests/integration/defs/.test_durations
+++ b/tests/integration/defs/.test_durations
@@ -23,6 +23,8 @@
"accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_disaggregated_dep2": 382.8,
"accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_and_long_context": 468.77,
"accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[dep4_fp8_kv]": 377.08,
+ "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[tep4_fp8_kv_flashinfer_ring]": 700.0,
+ "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[tep4_nvfp4_kv_no_reuse]": 700.0,
"accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[tp4_nvfp4_kv_breakable]": 597.49,
"accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_mmmu": 793.46,
"accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_breakable_prefill_matches_eager": 730.6153846153846,
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 6e879e5bd793..3c14de2dc351 100644
--- a/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py
+++ b/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py
@@ -397,12 +397,17 @@ def _media_request(tokenizer, processor, text: str, items: list[tuple[str, objec
# Attention replicated, each rank serving its own requests; the routed experts by expert.
_DEP4 = dict(tensor_parallel_size=4, moe_expert_parallel_size=4, enable_attention_dp=True)
-# The quantized KV pools, one layout each: NVFP4 at tp4 with the breakable prefill graph
-# (routed at attention TP only), FP8 at dep4 (MTP under attention DP, each rank drafting for
-# its own requests).
+# The quantized KV pools: NVFP4 at tp4 with the breakable prefill graph (routed at attention
+# TP only); FP8 at dep4 (MTP under attention DP, each rank drafting for its own requests); FP8
+# at tep4 over FlashInfer's ring replay (`TRTLLM_USE_GDN_FLASHINFER_REPLAY=force`, which fails
+# the build if the ring cannot run); NVFP4 at tep4 with chunked prefill and block reuse off, so
+# the engine leaves its paged-context attention off and the target takes it for the NVFP4 pool.
+# The attention-TP cases run 6 query heads a KV head on each rank.
_QUANTIZED_KV = [
- pytest.param(_TP4, "nvfp4", True, id="tp4_nvfp4_kv_breakable"),
- pytest.param(_DEP4, "fp8", False, id="dep4_fp8_kv"),
+ pytest.param(_TP4, "nvfp4", dict(breakable=True), id="tp4_nvfp4_kv_breakable"),
+ pytest.param(_DEP4, "fp8", {}, id="dep4_fp8_kv"),
+ pytest.param(_TEP4, "fp8", dict(flashinfer_ring=True), id="tep4_fp8_kv_flashinfer_ring"),
+ pytest.param(_TEP4, "nvfp4", dict(paged_context=False), id="tep4_nvfp4_kv_no_reuse"),
]
@@ -713,24 +718,34 @@ def test_mtp3_gsm8k_and_long_context(self, mocker):
@skip_not_sm103
@pytest.mark.skip_less_device(4)
- @pytest.mark.parametrize("parallel, kv_dtype, breakable", _QUANTIZED_KV)
- def test_mtp3_gsm8k_quantized_kv(self, mocker, parallel, kv_dtype, breakable):
- """gsm8k with three drafts over a quantized KV pool, block reuse and chunked prefill.
-
- gsm8k reads the MTP row of references/gsm8k.yaml for that KV dtype, its acceptance
- length held to that dtype's floor in references/acceptance_length.yaml.
+ @pytest.mark.parametrize("parallel, kv_dtype, variant", _QUANTIZED_KV)
+ def test_mtp3_gsm8k_quantized_kv(self, mocker, parallel, kv_dtype, variant):
+ """gsm8k, then pass keys over 3k-6k tokens, with three drafts over a quantized KV pool.
+
+ Block reuse and chunked prefill are on unless the case turns them off. gsm8k reads the
+ MTP row of references/gsm8k.yaml for that KV dtype; the pass keys are past the QSA
+ indexer's 2048-token budget, so they take the sparse attention path over the quantized
+ pool, which gsm8k never leaves. Both acceptance lengths are held to that dtype's floors
+ in references/acceptance_length.yaml.
"""
+ paged_context = variant.get("paged_context", True)
kv_cache_config = KvCacheConfig(
dtype=kv_dtype,
free_gpu_memory_fraction=0.6,
- enable_block_reuse=True,
+ enable_block_reuse=paged_context,
mamba_ssm_cache_dtype="bfloat16",
)
- engine = dict(self.ENGINE_MTP, enable_chunked_prefill=True, kv_cache_config=kv_cache_config)
- if breakable:
+ engine = dict(
+ self.ENGINE_MTP, enable_chunked_prefill=paged_context, kv_cache_config=kv_cache_config
+ )
+ if variant.get("breakable"):
engine.update(self.BREAKABLE_PREFILL)
+ env = dict(_VALIDATE)
+ if variant.get("flashinfer_ring"):
+ env["TRTLLM_USE_GDN_FLASHINFER_REPLAY"] = "force"
+ key = f"{self.ACCEPTANCE_KEY_MTP3}_{kv_dtype}_kv"
with (
- scoped_env_overrides(_VALIDATE) as rank_env,
+ scoped_env_overrides(env) as rank_env,
LLM(
self.MODEL_PATH,
speculative_config=self.MTP3,
@@ -741,8 +756,9 @@ def test_mtp3_gsm8k_quantized_kv(self, mocker, parallel, kv_dtype, breakable):
env_overrides=rank_env,
) as llm,
):
- self._gsm8k(
- mocker, llm, f"mtp3 {kv_dtype} kv", f"{self.ACCEPTANCE_KEY_MTP3}_{kv_dtype}_kv"
+ self._gsm8k(mocker, llm, f"mtp3 {kv_dtype} kv", key)
+ self._passkeys(
+ llm, _PASSKEY_LENGTHS, 3, f"{key}_long_context", f"mtp3 {kv_dtype} kv long context"
)
@skip_not_sm103
diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml
index 7ee8d7875f08..9e8b518d5be5 100644
--- a/tests/integration/defs/accuracy/references/acceptance_length.yaml
+++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml
@@ -243,6 +243,15 @@ ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3_nvfp4_kv:
ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3_fp8_kv:
ref_al: 3.512
min_al: 3.34
+# Read by test_mtp3_gsm8k_quantized_kv's pass keys (3k-6k tokens, the sparse attention path):
+# the NVFP4 cases at tp4 and tep4 (6 query heads a KV head) and the FP8 cases at dep4 and tep4.
+# modeling_v2's own readings on GB300; min_al 2.5% below, as for tp1_mtp3_long_context.
+ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3_nvfp4_kv_long_context:
+ ref_al: 3.962
+ min_al: 3.86
+ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3_fp8_kv_long_context:
+ ref_al: 3.962
+ min_al: 3.86
# Read by TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_and_long_context
# (test_qwen4_exp.py), over the pass keys of tp1_mtp3_long_context at tep4: three drafts with
# IndexShare on, bf16 KV.
diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt
index edd4420f9e91..721d4da56fa1 100644
--- a/tests/integration/test_lists/qa/llm_function_core.txt
+++ b/tests/integration/test_lists/qa/llm_function_core.txt
@@ -797,6 +797,8 @@ accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103D
accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_and_long_context TIMEOUT (60)
accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[tp4_nvfp4_kv_breakable] TIMEOUT (60)
accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[dep4_fp8_kv] TIMEOUT (60)
+accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[tep4_fp8_kv_flashinfer_ring] TIMEOUT (60)
+accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[tep4_nvfp4_kv_no_reuse] TIMEOUT (60)
accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_mmmu TIMEOUT (90)
accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_disaggregated_dep2 TIMEOUT (90)
diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
index 8707a434dabd..ad22d5a2e7fc 100644
--- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
+++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
@@ -39,6 +39,8 @@ l0_gb300_multi_gpus:
# pre-merge (l0_gb300); mmmu runs in QA (llm_function_core).
- accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[tp4_nvfp4_kv_breakable] TIMEOUT (60)
- accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[dep4_fp8_kv] TIMEOUT (60)
+ - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[tep4_fp8_kv_flashinfer_ring] TIMEOUT (60)
+ - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[tep4_nvfp4_kv_no_reuse] TIMEOUT (60)
- accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_disaggregated_dep2 TIMEOUT (90)
# qwen4_exp target at tp1, post-merge: its AgentX engines (NVFP4 KV, MTP2 with rejection
# sampling, the breakable prefill graph, every opt-in projection) on one context and one
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py
index bbe42f210b46..e4b2cf2113a4 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py
@@ -94,11 +94,13 @@ def build(
global_v=1.0,
request_dtype=torch.int32,
pool_pages=None,
+ q_heads=24,
+ kv_heads=2,
):
"""The decode op's test state for `segments` ((first position, rows) per request, in row
order) with `_selected_tokens(overlap)`."""
gen = torch.Generator(device=DEV).manual_seed(seed)
- tpb, kv, d, q_heads = 32, 2, 256, 24
+ tpb, kv, d = 32, kv_heads, 256
positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments])
req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)])
rows = positions.numel()
@@ -152,9 +154,8 @@ def _args(spec, seed):
if spec.get("gated"):
rows = args["q"].shape[0]
g = torch.Generator(device=DEV).manual_seed(seed + 1)
- args["output_gate"] = torch.randn(rows, 24 * 256, generator=g, device=DEV).to(
- torch.bfloat16
- )
+ width = args["q"].shape[1] * args["q"].shape[2]
+ args["output_gate"] = torch.randn(rows, width, generator=g, device=DEV).to(torch.bfloat16)
return args
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ssm/test_modeling_v2_gdn_block.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ssm/test_modeling_v2_gdn_block.py
new file mode 100644
index 000000000000..7ef5b34c4408
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ssm/test_modeling_v2_gdn_block.py
@@ -0,0 +1,349 @@
+# 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 layer against an fp32 PyTorch reference.
+
+Drives `DecodeTarget._gdn_block` and `PrefillTarget._gdn_block` at a rank's two geometries --
+the whole layer's 16 key / 48 value heads (tp1, dep4) and a tp4 / tep4 rank's 4 / 12 -- over
+the recurrent pool dtypes each serves (fp32 or bf16 at 16 / 48, bf16 at 4 / 12), on a decode
+batch, a mixed batch (contexts with and without carried state plus one-token generation rows)
+and a context-only batch, and compares the layer output and both state pools with a reference
+written from the layer's definition in PyTorch: the input projection (the same GEMM, so the
+conv pools are bitwise), the causal conv over each sequence's carried window, SiLU, the q / k
+L2 norm, the decay and beta gates, the delta-rule recurrence per token, the sigmoid-gated RMS
+norm and the output projection. The reference shares no kernel with the target past the two
+GEMMs, so a wrong head split, gate or slot at either geometry fails it.
+
+Bands: the conv pools are bitwise; the recurrent pool is held to rtol 1.6e-2 / atol 2e-2 (the
+recurrence entries' state bands), the layer output to a relative L2 distance of 1e-2 and a
+max-abs distance of 3e-2 of its own peak.
+"""
+
+from __future__ import annotations
+
+import types
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501
+ modeling as M,
+)
+
+assert torch.cuda.is_available(), "the Gated-DeltaNet layer requires a CUDA device"
+
+DEV = "cuda"
+BF16 = torch.bfloat16
+HIDDEN, K, V, WIDTH = 2560, 128, 128, 4
+SLOTS = 461
+MAX_BATCH = 256
+RTOL_STATE, ATOL_STATE = 1.6e-2, 2e-2
+OUT_REL_L2, OUT_MAX_FRAC = 1e-2, 3e-2
+
+#: (key heads, value heads) of a rank: the whole layer's, and a tp4 / tep4 rank's.
+GEOMETRIES = {"16x48": (16, 48), "4x12": (4, 12)}
+#: (geometry, pool dtype): the pool dtypes each layout serves.
+CASES = [("16x48", "fp32"), ("16x48", "bf16"), ("4x12", "bf16")]
+STATES = {"fp32": torch.float32, "bf16": BF16}
+
+
+def _dims(h: int, hv: int) -> types.SimpleNamespace:
+ kd, vd = h * K, hv * V
+ return types.SimpleNamespace(h=h, hv=hv, kd=kd, vd=vd, conv=2 * kd + vd)
+
+
+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(d, seed: int, state: torch.dtype, one_rank: bool) -> types.SimpleNamespace:
+ """The attributes `_gdn_block` reads off the core, with layer 0's pools (the recurrent
+ one in `state`); the TP sum is the identity (one rank's partial)."""
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ conv = torch.randn(SLOTS, d.conv, WIDTH - 1, generator=g, device=DEV, dtype=BF16)
+ ssm = (torch.randn(SLOTS, d.hv, V, K, generator=g, device=DEV) * 0.05).to(state)
+ return types.SimpleNamespace(
+ conv_dim=d.conv,
+ key_dim=d.kd,
+ value_dim=d.vd,
+ num_k_heads=d.h,
+ head_k_dim=K,
+ num_v_heads=d.hv,
+ head_v_dim=V,
+ eps=1e-6,
+ conv_kernel=WIDTH,
+ in_proj_split=(d.conv, d.vd, d.hv, d.hv),
+ one_rank=one_rank,
+ tp_sum=lambda x: x,
+ decode_mm=M._one_rank_projection if one_rank else M._decode_projection,
+ decode_in_mm=(
+ M._one_rank_projection
+ if one_rank
+ else (lambda mat_a, mat_b: M.low_m_bf16_mm(mat_a, mat_b, M._DECODE_PROJ_PDL))
+ ),
+ _bound=types.SimpleNamespace(pools=types.SimpleNamespace(conv={0: conv}, ssm={0: ssm})),
+ )
+
+
+def _clone_core(core: types.SimpleNamespace) -> types.SimpleNamespace:
+ out = types.SimpleNamespace(**vars(core))
+ pools = core._bound.pools
+ out._bound = types.SimpleNamespace(
+ pools=types.SimpleNamespace(conv={0: pools.conv[0].clone()}, ssm={0: pools.ssm[0].clone()})
+ )
+ return out
+
+
+def _layer(d, 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), with no FP8 or NVFP4 copies."""
+ 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()
+
+ return types.SimpleNamespace(
+ in_proj=lin(d.conv + d.vd + 2 * d.hv, HIDDEN),
+ out=lin(HIDDEN, d.vd),
+ conv_w=(torch.randn(d.conv, WIDTH, generator=g, device=DEV) * 0.5).to(BF16),
+ a_log=torch.empty(d.hv, device=DEV).uniform_(-5.5, 3.5, generator=g),
+ dt_bias=torch.empty(d.hv, device=DEV).uniform_(-8.9, 19.2, generator=g),
+ gnorm=(1.0 + 0.1 * torch.randn(V, generator=g, device=DEV)).to(BF16),
+ in_proj8=None,
+ out8=None,
+ in_proj_g=None,
+ out_g=None,
+ in_proj_m=None,
+ out_m=None,
+ in_proj_nv=None,
+ )
+
+
+def _x(tokens: int, seed: int) -> torch.Tensor:
+ g = torch.Generator(device=DEV).manual_seed(seed)
+ return torch.randn(tokens, HIDDEN, generator=g, device=DEV, dtype=BF16)
+
+
+def _target(cls, core, L=None, step=None):
+ """A target over `core` alone; a prefill target's `_gdn_core` break reads layer 0's
+ weights and the step's bookkeeping off it."""
+ target = object.__new__(cls)
+ target.core = core
+ target._layers = [L]
+ target._step = step
+ return target
+
+
+def _decode_slots(m: int) -> torch.Tensor:
+ """The decode step's slot vector: the head of an engine-sized buffer."""
+ buf = torch.zeros(MAX_BATCH, dtype=torch.int32, device=DEV)
+ buf[:m] = torch.tensor(_slots(m), dtype=torch.int32, device=DEV)
+ return buf[:m]
+
+
+def _prefill_step(ctx_lens: tuple, n_dec: int, has_init: tuple) -> tuple:
+ """`PrefillTarget`'s `step.gdn` 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)
+ 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)
+ qsl_all = cu_buf[: batch + 1]
+ qsl_long = qsl_all.to(torch.int64)
+ return (
+ n_p, # num_prefill_tokens
+ n_dec, # num_decode_tokens
+ si_p,
+ si_all,
+ torch.clone(si_d),
+ hi_all[:n_ctx], # hi_p
+ hi_all,
+ qsl_all[: n_ctx + 1], # qsl_p
+ qsl_all,
+ qsl_long,
+ qsl_long[: n_ctx + 1], # qsl_long_p
+ n_dec, # num_decodes
+ 1, # gen_p
+ None, # verify
+ )
+
+
+# ------------------------------------------------------------------ the reference ---
+
+
+def _reference(d, core, L, x, seqs) -> torch.Tensor:
+ """The layer over `x`, sequence by sequence: `seqs` is `(first row, rows, slot, carries)`
+ per sequence (a generation row carries its state). Updates `core`'s pools in place (the
+ conv window bitwise, the recurrent state rounded once to the pool's dtype) and returns
+ the layer output."""
+ conv_pool, ssm_pool = core._bound.pools.conv[0], core._bound.pools.ssm[0]
+ proj = M.cublas_mm(x, L.in_proj)
+ mixed, z, a, b = torch.split(proj, [d.conv, d.vd, d.hv, d.hv], 1)
+ rep = d.hv // d.h
+ w = L.conv_w.float()
+ neg_a = -torch.exp(L.a_log.float())
+ o = torch.empty(x.shape[0], d.hv, V, dtype=torch.float32, device=DEV)
+ for first, rows, slot, carries in seqs:
+ seq = slice(first, first + rows)
+ window = conv_pool[slot] if carries else torch.zeros_like(conv_pool[slot])
+ full = torch.cat([window, mixed[seq].t()], 1) # [C, W - 1 + rows]
+ conv_pool[slot] = full[:, -(WIDTH - 1) :]
+ conv = sum(w[:, j : j + 1] * full[:, j : j + rows].float() for j in range(WIDTH))
+ conv = F.silu(conv).to(BF16).float().t() # [rows, C], the kernels' bf16 output
+ q, k, v = torch.split(conv, [d.kd, d.kd, d.vd], 1)
+ q = q.reshape(rows, d.h, K).repeat_interleave(rep, dim=1)
+ k = k.reshape(rows, d.h, K).repeat_interleave(rep, dim=1)
+ v = v.reshape(rows, d.hv, V)
+ q = q * torch.rsqrt(q.pow(2).sum(-1, keepdim=True) + 1e-6) * K**-0.5
+ k = k * torch.rsqrt(k.pow(2).sum(-1, keepdim=True) + 1e-6)
+ decay = torch.exp(neg_a * F.softplus(a[seq].float() + L.dt_bias.float(), threshold=20.0))
+ beta = torch.sigmoid(b[seq].float())
+ h = (
+ ssm_pool[slot].float()
+ if carries
+ else torch.zeros_like(ssm_pool[slot], dtype=torch.float32)
+ )
+ for t in range(rows):
+ h = h * decay[t][:, None, None]
+ u = (v[t] - torch.einsum("hvk,hk->hv", h, k[t])) * beta[t][:, None]
+ h = h + u[:, :, None] * k[t][:, None, :]
+ o[first + t] = torch.einsum("hvk,hk->hv", h, q[t])
+ ssm_pool[slot] = h.to(ssm_pool.dtype)
+ o = o.to(BF16).float() # the recurrence entries' bf16 output
+ y = o * torch.rsqrt(o.pow(2).mean(-1, keepdim=True) + core.eps) * L.gnorm.float()
+ y = y * torch.sigmoid(z.reshape(-1, d.hv, V).float())
+ return M.cublas_mm(y.reshape(-1, d.vd).to(BF16), L.out)
+
+
+def _check(y, y_ref, core, core_ref, what) -> None:
+ """Output and both pools against the reference's, at the module bands."""
+ assert y.shape == y_ref.shape and y.dtype == y_ref.dtype
+ assert torch.isfinite(y).all()
+ diff = y.float() - y_ref.float()
+ rel = (diff.norm() / y_ref.float().norm()).item()
+ peak = (diff.abs().max() / y_ref.float().abs().max()).item()
+ ssm, ssm_ref = core._bound.pools.ssm[0], core_ref._bound.pools.ssm[0]
+ s_rel = ((ssm.float() - ssm_ref.float()).norm() / ssm_ref.float().norm()).item()
+ print(f"{what}: out rel-L2 {rel:.3e}, out max-abs/peak {peak:.3e}, state rel-L2 {s_rel:.3e}")
+ assert rel <= OUT_REL_L2 and peak <= OUT_MAX_FRAC, (rel, peak)
+ assert torch.equal(core._bound.pools.conv[0], core_ref._bound.pools.conv[0]), (
+ "conv pools differ"
+ )
+ torch.testing.assert_close(ssm.float(), ssm_ref.float(), rtol=RTOL_STATE, atol=ATOL_STATE)
+
+
+def _case(geometry: str, state: str, seed: int):
+ d = _dims(*GEOMETRIES[geometry])
+ core = _core(d, seed, STATES[state], one_rank=geometry == "16x48")
+ return d, core, _layer(d, seed + 1)
+
+
+# ---------------------------------------------------------------------- the tests ---
+
+
+@pytest.mark.parametrize("geometry, state", CASES, ids=[f"{g}-{s}" for g, s in CASES])
+def test_decode_block(geometry, state) -> None:
+ """A 128-row decode step; slots the step does not name are not written."""
+ m = 128
+ d, core, L = _case(geometry, state, 100)
+ x = _x(m, 300)
+ core_ref, before = _clone_core(core), _clone_core(core)
+ si = _decode_slots(m)
+ y = _target(M.DecodeTarget, core)._gdn_block(x, 0, L, si, None, 1)
+ y_ref = _reference(d, core_ref, L, x, [(r, 1, s, True) for r, s in enumerate(_slots(m))])
+ _check(y, y_ref, core, core_ref, f"decode m={m}, {geometry} {state} pool")
+ unnamed = [s for s in range(SLOTS) if s not in set(_slots(m))]
+ assert torch.equal(core._bound.pools.ssm[0][unnamed], before._bound.pools.ssm[0][unnamed])
+
+
+def test_captured_one_rank_fp32_decode_block() -> None:
+ """At one rank over an fp32 pool a captured step's recurrence and gated norm are one
+ launch (`gdn_decode_gated_norm_pdl`); run eagerly here against the same reference."""
+ m = 32
+ d, core, L = _case("16x48", "fp32", 110)
+ x = _x(m, 310)
+ core_ref = _clone_core(core)
+ y = _target(M.DecodeTarget, core)._gdn_block(x, 0, L, _decode_slots(m), None, 1, True)
+ y_ref = _reference(d, core_ref, L, x, [(r, 1, s, True) for r, s in enumerate(_slots(m))])
+ _check(y, y_ref, core, core_ref, f"captured-form decode m={m}, 16x48 fp32 pool")
+
+
+def _prefill(core, L, x, step):
+ tokens = x.shape[0]
+ out = torch.empty([1, tokens, core.num_v_heads, V], dtype=BF16, device=DEV)
+ target = _target(M.PrefillTarget, core, L, types.SimpleNamespace(gdn=step, num_tokens=tokens))
+ return target._gdn_block(x, 0, L, out)
+
+
+def _prefill_seqs(ctx_lens, n_dec, has_init) -> list:
+ slots = _slots(len(ctx_lens) + n_dec)
+ seqs, first = [], 0
+ for n, (rows, carries) in enumerate(zip(ctx_lens, has_init)):
+ seqs.append((first, rows, slots[n], carries))
+ first += rows
+ for j in range(n_dec):
+ seqs.append((first + j, 1, slots[len(ctx_lens) + j], True))
+ return seqs
+
+
+@pytest.mark.parametrize(
+ "ctx_lens, n_dec, has_init",
+ [((200, 37, 91), 5, (True, False, True)), ((130, 64), 0, (True, False))],
+ ids=["mixed", "context_only"],
+)
+@pytest.mark.parametrize("geometry, state", CASES, ids=[f"{g}-{s}" for g, s in CASES])
+def test_prefill_block(geometry, state, ctx_lens, n_dec, has_init) -> None:
+ """Variable-length contexts, with and without carried state, and (mixed) one-token
+ generation rows in the same step."""
+ tokens = sum(ctx_lens) + n_dec
+ d, core, L = _case(geometry, state, 400 + tokens)
+ x = _x(tokens, 600 + tokens)
+ core_ref = _clone_core(core)
+ y = _prefill(core, L, x, _prefill_step(ctx_lens, n_dec, has_init))
+ y_ref = _reference(d, core_ref, L, x, _prefill_seqs(ctx_lens, n_dec, has_init))
+ _check(y, y_ref, core, core_ref, f"prefill {ctx_lens}+{n_dec}x1, {geometry} {state} pool")
+
+
+@pytest.mark.parametrize("geometry, state", CASES, ids=[f"{g}-{s}" for g, s in CASES])
+def test_decode_block_replays_under_cuda_graph(geometry, state) -> None:
+ """The decode block captured into a CUDA graph after one eager call (as the engine's
+ warm-up runs it) replays to exactly the eager result."""
+ m = 32
+ _, core0, L = _case(geometry, state, 700)
+ x = _x(m, 702)
+ si = _decode_slots(m)
+ eager_core = _clone_core(core0)
+ y_eager = _target(M.DecodeTarget, eager_core)._gdn_block(x, 0, L, si, None, 1).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, si, None, 1)
+ torch.cuda.current_stream().wait_stream(stream)
+ torch.cuda.synchronize()
+ pools, pools0 = graph_core._bound.pools, core0._bound.pools
+ pools.conv[0].copy_(pools0.conv[0])
+ pools.ssm[0].copy_(pools0.ssm[0])
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ y_graph = target._gdn_block(x, 0, L, si, None, 1)
+ assert torch.equal(pools.ssm[0], pools0.ssm[0]), "capture must not execute"
+ graph.replay()
+ torch.cuda.synchronize()
+ assert torch.equal(y_graph, y_eager)
+ assert torch.equal(pools.ssm[0], eager_core._bound.pools.ssm[0])
+ assert torch.equal(pools.conv[0], eager_core._bound.pools.conv[0])
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py
index f46db3ed01ea..b766831b715a 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py
@@ -89,6 +89,19 @@ def _qwen4_exp_config(text_overrides=None, **vision_overrides):
num_attention_heads=24,
linear_num_value_heads=48,
vocab_size=248320,
+ # The settings the target pins (`routing._TEXT_SETTINGS`), as the normalized config
+ # carries them: rope theta and partial factor flattened to the top level.
+ norm_topk_prob=True,
+ num_experts_per_tok=10,
+ hidden_act="silu",
+ num_key_value_heads=2,
+ head_dim=256,
+ attention_bias=False,
+ output_gate_type="sigmoid",
+ tie_word_embeddings=False,
+ rope_theta=10000000,
+ partial_rotary_factor=0.25,
+ rope_scaling=dict(type="mrope", mrope_section=[11, 11, 10], mrope_interleaved=True),
)
text.update(text_overrides or {})
vision = dict(
@@ -320,6 +333,55 @@ def test_another_checkpoint_shape_declines():
_declined(config, match="text shape")
+_ROPE = dict(type="mrope", mrope_section=[11, 11, 10], mrope_interleaved=True)
+
+
+@pytest.mark.parametrize(
+ "override, field",
+ [
+ (dict(norm_topk_prob=False), "norm_topk_prob"),
+ (dict(num_experts_per_tok=8), "num_experts_per_tok"),
+ (dict(hidden_act="gelu"), "hidden_act"),
+ (dict(num_key_value_heads=4), "num_key_value_heads"),
+ (dict(attention_bias=True), "attention_bias"),
+ (dict(output_gate_type="silu"), "output_gate_type"),
+ (dict(tie_word_embeddings=True), "tie_word_embeddings"),
+ (dict(rope_theta=1e6), "rope_theta"),
+ (dict(partial_rotary_factor=0.5), "partial_rotary_factor"),
+ (dict(rope_scaling=dict(_ROPE, mrope_interleaved=False)), "mrope_interleaved"),
+ (dict(rope_scaling=dict(_ROPE, mrope_section=[16, 8, 8])), "mrope_section"),
+ (dict(rope_scaling=dict(_ROPE, factor=4.0)), "rope_factor"),
+ ],
+ ids=lambda v: v if isinstance(v, str) else None,
+)
+def test_a_same_shape_checkpoint_with_other_settings_declines(override, field):
+ """A fine-tune with the checkpoint's shape that changes what it computes, which the target
+ would run with wrong logits and no error: `auto` keeps the built-in model."""
+ config = _model_config(pretrained_config=_qwen4_exp_config(override), **_LAYOUTS["tep4"])
+ _declined(config, match=f"'{field}'")
+
+
+def test_the_settings_read_off_the_rope_parameters_dict():
+ """A config that keeps its rope fields in `rope_parameters` (theta and partial factor
+ inside) is read the same way, as the target reads it."""
+ rope = dict(
+ rope_type="default",
+ rope_theta=10000000,
+ partial_rotary_factor=0.25,
+ mrope_section=[11, 11, 10],
+ mrope_interleaved=True,
+ )
+ text = dict(
+ rope_parameters=rope, rope_scaling=None, rope_theta=None, partial_rotary_factor=None
+ )
+ _claimed(_model_config(pretrained_config=_qwen4_exp_config(text), **_LAYOUTS["tep4"]))
+ text["rope_parameters"] = dict(rope, rope_theta=5e6)
+ _declined(
+ _model_config(pretrained_config=_qwen4_exp_config(text), **_LAYOUTS["tep4"]),
+ match="'rope_theta'",
+ )
+
+
def test_a_deepstack_vision_tower_declines():
config = _model_config(
pretrained_config=_qwen4_exp_config(deepstack_visual_indexes=[5, 11, 17]),
@@ -342,6 +404,31 @@ def test_quantized_kv_cache_dtypes_match(member, dtype):
_claimed(_model_config(**_LAYOUTS[member]), kv_cache_config=KvCacheConfig(dtype=dtype))
+@pytest.mark.parametrize(
+ "chunked, reuse",
+ [(False, False), (True, False), (False, True)],
+ ids=["neither", "chunked_prefill", "block_reuse"],
+)
+@pytest.mark.parametrize("dtype", ["fp8", "nvfp4"])
+def test_a_quantized_kv_cache_matches_with_or_without_the_paged_context_path(dtype, chunked, reuse):
+ """The engine enables its paged-context attention only for chunked prefill or block reuse;
+ the target takes it for an NVFP4 pool either way, so neither setting is declined."""
+ _claimed(
+ _model_config(**_LAYOUTS["tep4"]),
+ enable_chunked_prefill=chunked,
+ kv_cache_config=KvCacheConfig(dtype=dtype, enable_block_reuse=reuse),
+ )
+
+
+@pytest.mark.parametrize("value", ["1", "force"])
+@pytest.mark.parametrize("max_draft_len", [3, 7])
+def test_mtp_matches_with_the_flashinfer_ring_replay(monkeypatch, max_draft_len, value):
+ """`TRTLLM_USE_GDN_FLASHINFER_REPLAY` lays the replay history out as FlashInfer's ring at 3
+ and 7 drafts over a bf16 state, which the target's verify step runs."""
+ monkeypatch.setenv("TRTLLM_USE_GDN_FLASHINFER_REPLAY", value)
+ _claimed(_model_config(MTPDecodingConfig(max_draft_len=max_draft_len), **_LAYOUTS["tep4"]))
+
+
def test_other_kv_cache_algorithms_decline():
choices = routing.engine_choices(_model_config(**_LAYOUTS["tep4"]), _llm_args())
for algo in (QuantAlgo.INT8, QuantAlgo.FP8_BLOCK_SCALES):
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 c23aa0c587e7..b7a579756d49 100644
--- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py
+++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py
@@ -181,6 +181,19 @@ def _qwen4_exp_config(**vision_overrides):
num_attention_heads=24,
linear_num_value_heads=48,
vocab_size=248320,
+ # The settings the target pins (`routing._TEXT_SETTINGS`), as the normalized config
+ # carries them: rope theta and partial factor flattened to the top level.
+ norm_topk_prob=True,
+ num_experts_per_tok=10,
+ hidden_act="silu",
+ num_key_value_heads=2,
+ head_dim=256,
+ attention_bias=False,
+ output_gate_type="sigmoid",
+ tie_word_embeddings=False,
+ rope_theta=10000000,
+ partial_rotary_factor=0.25,
+ rope_scaling=dict(type="mrope", mrope_section=[11, 11, 10], mrope_interleaved=True),
)
vision = dict(
depth=27,
From 8ae17252682b0ada77ad045f969ac32b9f57191a Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Thu, 5 Nov 2026 01:32:11 -0800
Subject: [PATCH 08/15] [None][chore] modeling_v2 qwen4_exp: drop the
conv-state pool guard now that #21684 is merged
TRT-LLM's causal_conv1d kernels form a conv-state slot's offset in 64 bits
since #21684, so the target's wrapper of causal_conv1d_update no longer
refuses a pool of 2^32 elements or more; its note and test follow.
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
.../qwen4_exp/ops/ssm/causal_conv1d_update.py | 24 ++++---------------
.../test_modeling_v2_causal_conv1d_update.py | 15 ------------
2 files changed, 5 insertions(+), 34 deletions(-)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_update.py
index 6637f807dd9c..da22d1fc172e 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_update.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_update.py
@@ -25,9 +25,6 @@
#: This checkpoint's GDN `[q|k|v]` width; the conv is width 4 over a 3-column pool.
CKPT_DIM = 10240
-#: TRT-LLM's kernel forms a slot's offset as `int slot * uint32_t conv_state_batch_stride`:
-#: a pool of this many elements or more (slots x slot stride) wraps, so the op refuses it.
-TRTLLM_SLOT_OFFSET_LIMIT = 2**32
#: Largest conv-pool slot count the cache manager allocates for this checkpoint at a batch
#: of at most 32. The decode step hands the whole pool over unsliced.
ENGINE_SLOTS = 461
@@ -502,14 +499,11 @@ class CausalConv1dUpdate(OpWrapper):
`cpp/tensorrt_llm/thop/causalConv1dOp.cpp`); the CuTe kernel is JIT-compiled
by `cute.compile` on its first call for the running device, sm_103 here.
- Hard limits (from source, not probed: the measured pools stay far below)
- The TRT-LLM kernel forms a slot's offset as `int slot * uint32_t
- conv_state_batch_stride`, 32-bit unsigned elements: it wraps from slot
- 15535 at stride 276480 (tp4 / tep4) and slot 3884 at 1105920 (dep4).
- The measured pools reach 518,123,519 (tep4) and 940,031,999 (dep4)
- elements, under 2^31 elements and bytes. A call that would run it over a
- pool of 2^32 elements or more (`TRTLLM_SLOT_OFFSET_LIMIT`) raises
- ValueError instead. The CuTe kernel forms slot and row offsets in 64 bits.
+ Offsets: both kernels form a slot's offset in 64 bits (the TRT-LLM kernel
+ as `int64_t slot * conv_state_batch_stride`), so a pool past 2^32 elements
+ is addressed whole; the TRT-LLM op refuses a slot stride of 2^32 elements
+ or more, which it keeps in 32 bits. The measured pools reach 518,123,519
+ (tep4) and 940,031,999 (dep4) elements.
"""
def __call__(
@@ -528,14 +522,6 @@ def __call__(
):
_decode_conv(x, conv_state, weight, conv_state_indices, pad_slot_id)
return x
- if conv_state.shape[0] * conv_state.stride(0) >= TRTLLM_SLOT_OFFSET_LIMIT:
- last = (TRTLLM_SLOT_OFFSET_LIMIT - 1) // conv_state.stride(0)
- raise ValueError(
- f"a conv-state pool of {conv_state.shape[0]} slots at stride "
- f"{conv_state.stride(0)}: TRT-LLM's C++ conv kernel addresses a slot in 32 bits "
- f"and would wrap past slot {last}; lower free_gpu_memory_fraction or the "
- "recurrent-state share of the KV cache"
- )
torch.ops.trtllm.causal_conv1d_update(
x.unsqueeze(-1) if x.dim() == 2 else x,
conv_state,
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_update.py
index 345516d0b8dd..89c89a820bc1 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_update.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_update.py
@@ -123,21 +123,6 @@ def test_which_kernel_each_member_runs() -> None:
assert _decode_conv_ok(*args[:7]) is expect, (dim, expect)
-def test_a_pool_past_the_32_bit_slot_offset_is_refused() -> None:
- """TRT-LLM's kernel forms a slot's offset in 32 bits: a pool of 2^32 elements or more is
- refused before any launch (meta tensors: nothing is allocated)."""
- stride = 1105920 # a dep4 rank's slot stride
- slots = -(-(2**32) // stride)
- x = torch.empty((4, 10240), dtype=torch.bfloat16, device="meta")
- weight = torch.empty((10240, 4), dtype=torch.bfloat16, device="meta")
- indices = torch.zeros(4, dtype=torch.int32, device="meta")
- pool = torch.empty_strided(
- (slots, 10240, 3), (stride, 3, 1), dtype=torch.bfloat16, device="meta"
- )
- with pytest.raises(ValueError, match="32 bits"):
- op(x, pool, weight, None, True, None, indices, -1)
-
-
def test_op_rejects_metadata_violations_loudly() -> None:
"""Why the entry has no `is_valid`: the op raises on metadata violations
itself (probed here at int64 slot indices)."""
From e64db8dc8841bf1d9defccfbc40beb63a83842be Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Thu, 5 Nov 2026 02:04:06 -0800
Subject: [PATCH 09/15] [None][test] modeling_v2 qwen4_exp: the replay stub's
ring fields; the tep4 fp8 decode cell's split count; measured long-context AL
references
The decode-paths test's stand-in replay names the FlashInfer ring fields the
verify step now reads, the tep4 kv_fp8 32-row gated decode cell expects the
8-way split-K the op takes at that row count, and the quantized-KV pass keys'
acceptance-length references carry the GB300 readings.
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
.../models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa.py | 4 ++--
.../defs/accuracy/references/acceptance_length.yaml | 5 +++--
.../test_decode_paths.py | 7 ++++++-
3 files changed, 11 insertions(+), 5 deletions(-)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa.py
index 1b96989177e4..f1343f321ad6 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa.py
@@ -256,8 +256,8 @@ class QsaPagedSparseGqa(OpWrapper):
),
Cell(
why="tep4 kv_fp8 gated decode 32 x 1 row: 6 q / 1 kv head over an e4m3 pool, q and "
- "the gate off the 17-head fused row, 32-way split-K + gated merge",
- spec=_default(_spread(32, 1, 9215), _TP4_FP8, _SPLIT_K, splits=32, **_TP4_ROW),
+ "the gate off the 17-head fused row, 8-way split-K + gated merge",
+ spec=_default(_spread(32, 1, 9215), _TP4_FP8, _SPLIT_K, splits=8, **_TP4_ROW),
),
Cell(
why="tep4 kv_fp8 MTP3 gated verify 128 x 4 rows = 512: 6 q / 1 kv head over an e4m3 "
diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml
index 9e8b518d5be5..440f6c75d679 100644
--- a/tests/integration/defs/accuracy/references/acceptance_length.yaml
+++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml
@@ -245,12 +245,13 @@ ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3_fp8_kv:
min_al: 3.34
# Read by test_mtp3_gsm8k_quantized_kv's pass keys (3k-6k tokens, the sparse attention path):
# the NVFP4 cases at tp4 and tep4 (6 query heads a KV head) and the FP8 cases at dep4 and tep4.
-# modeling_v2's own readings on GB300; min_al 2.5% below, as for tp1_mtp3_long_context.
+# modeling_v2's own readings on GB300 (NVFP4 3.951-3.968, FP8 3.960-3.975 over the cases and
+# repeats); min_al 2.5% below, as for tp1_mtp3_long_context.
ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3_nvfp4_kv_long_context:
ref_al: 3.962
min_al: 3.86
ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3_fp8_kv_long_context:
- ref_al: 3.962
+ ref_al: 3.965
min_al: 3.86
# Read by TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_and_long_context
# (test_qwen4_exp.py), over the pass keys of tp1_mtp3_long_context at tep4: three drafts with
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_decode_paths.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_decode_paths.py
index 6c88e8525cb2..e453ad16f25f 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_decode_paths.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_decode_paths.py
@@ -263,7 +263,12 @@ def replay(**kw):
),
int_conv={0: win0.clone()},
replay=types.SimpleNamespace(
- old={0: (history,) * 4}, cache_buf_idx=None, prev_num_accepted=None, history=16
+ old={0: (history,) * 4},
+ cache_buf_idx=None,
+ prev_num_accepted=None,
+ history=16,
+ ring=None,
+ cache_base=None,
),
),
)
From 13074572912d98cea3642d358fffdaf289c4b10a Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Thu, 5 Nov 2026 03:58:33 -0800
Subject: [PATCH 10/15] [None][fix] modeling_v2 qwen4_exp: one router kernel
across attention-DP ranks; the MTP settings routing pins; the NVFP4 context
row-tile pass at every layout
- Under attention DP every rank computes the router logits of the rows its group gathers. An
eager decode step now takes the prefill router (cublas_mm), so a decode rank and a prefill
rank of one step pick the same top-k; a captured step, a decode step on every rank, keeps
_decode_router.
- Routing pins the MTP module's settings (one full-attention layer that reads the end of the
stack, on the trunk's embedding) and attn_output_gate, reading the mtp block as a mapping or a
config object; the target reads the block's own rope theta in either form.
- Over an NVFP4 pool, a step with more than 256 context rows attends them on the row-tile pass
over the dequantized bf16 scratch at every layout: the family's pass at one rank, the
target's (6/1 and 24/2 heads) elsewhere. A 6/1 scratch cell for the target's row-tile op.
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
...-guide-for-qwen3.8-flash-next-on-trtllm.md | 2 +-
.../modeling.py | 74 ++++++++++-------
.../qsa_sparse_gqa_rowtile__prefill.py | 8 ++
.../modeling_v2/models/qwen4_exp/routing.py | 32 +++++++-
.../references/acceptance_length.yaml | 2 +-
.../test_attention_plumbing.py | 80 ++++++++++++++-----
.../test_core_layout.py | 38 ++++++++-
.../test_routing.py | 40 +++++++++-
8 files changed, 218 insertions(+), 58 deletions(-)
diff --git a/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md
index d8a111121b32..bff2472e1313 100644
--- a/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md
+++ b/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md
@@ -220,7 +220,7 @@ It supports these single-node layouts:
Under `auto`, these settings make the server fall back to the built-in model; under `require`, the server fails at startup:
* Any other layout, including `enable_lm_head_tp_in_adp: true` with attention DP.
-* A checkpoint of the same shape whose config changes what it computes, for example `norm_topk_prob: false`, another `num_experts_per_tok`, or other rope settings.
+* A checkpoint of the same shape whose config changes what it computes, for example `norm_topk_prob: false`, another `num_experts_per_tok`, other rope settings, or an MTP module that reads another layer's hidden state (`mtp_use_hidden_state_from_layer`).
* Speculative decoding other than one-model MTP, or a draft length that changes at runtime.
* MTP with `TRTLLM_USE_GDN_REPLAY=0`.
* A KV cache dtype other than the checkpoint's bf16, `fp8` or `nvfp4`.
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py
index 5388a8e9f36f..79159b457996 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py
@@ -57,8 +57,8 @@
routed experts for prefill steps past `_CUTE_DSL_MOE_MIN_ROWS` (`moe_prefill_copy`), or the
CuteDSL pipeline for every step (`moe_config.backend: CUTEDSL`, `_cute_dsl_routed`); an fp32
recurrent state, whose captured single-token decode runs `gdn_decode_gated_norm_pdl`; the
-decode draft layer and its FP8 head; the NVFP4 pool's context rows attended a group of
-requests at a time (`_nvfp4_context_plan`); and the opt-in FP8 / NVFP4 copies of the dense
+decode draft layer and its FP8 head; the family's row-tile pass for the NVFP4 pool's context
+rows (`_sparse_gqa_nvfp4_context`); and the opt-in FP8 / NVFP4 copies of the dense
projections and heads (`_ONE_RANK_SWITCHES`: `_DENSE_FP8_ENV`, `_GDN_INPROJ_NVFP4_ENV`,
`_TARGET_HEAD_NVFP4_ENV`, `_MTP_HEAD_NVFP4_ENV`, `_MTP_MOE_NVFP4_ENV`, `_MTP_DENSE_FP8_ENV`).
@@ -629,9 +629,9 @@
# Over an fp8 pool the launch reads its requests' pages widened into a bf16 scratch; a launch
# whose scratch would pass this many bytes stays on the per-row kernel.
_FP8_ROWTILE_MAX_SCRATCH = 1 << 30 # capacity: fixed: bounds a per-call scratch allocation
-# At one rank a step with more context rows than this over an NVFP4 pool attends them on the
-# row-tile pass over a bf16 scratch of their requests' pages (`_sparse_gqa_nvfp4_context`):
-# 16384 context rows 1.36 ms a layer there against 77-948 ms for the NVFP4 per-row kernel.
+# A step with more context rows than this over an NVFP4 pool attends them on the row-tile pass
+# over a bf16 scratch of their requests' pages (`_sparse_gqa_nvfp4_context`): at one rank 16384
+# context rows 1.36 ms a layer there against 77-948 ms for the NVFP4 per-row kernel.
_NVFP4_ROWTILE_MIN_ROWS = 256 # capacity: fixed: a dispatch threshold, not a buffer size
# The bf16 K/V scratch those context rows read holds every page of each context request's KV
# history. Consecutive context requests are dequantized and attended in groups whose pages fit
@@ -1098,14 +1098,16 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig
# ---- the MTP module ------------------------------------------------
# Global layer index of the draft layer (the KV managers map it there, and the
- # weight table keys it there); its rope base is its own block's, else the trunk's.
+ # weight table keys it there); its rope base is its own block's (a mapping or a config
+ # object), else the trunk's. Routing pins the block's other settings.
self.mtp_layer_idx = self.num_layers
mtp_block = getattr(cfg, "mtp", None)
- self.mtp_theta = (
- float(mtp_block.get("rope_theta", self.theta))
+ mtp_theta = (
+ mtp_block.get("rope_theta")
if isinstance(mtp_block, dict)
- else self.theta
+ else getattr(mtp_block, "rope_theta", None)
)
+ self.mtp_theta = self.theta if mtp_theta is None else float(mtp_theta)
# ---- gated deltanet ----------------------------------------------
self.head_k_dim = cfg.linear_key_head_dim
@@ -2808,10 +2810,19 @@ def _hc_moe_mixer_rstd(core, hc, rstd, L: _LayerWeights):
return mixed, inject, (xq, xsf)
-def _decode_router_mid(core, L: _LayerWeights):
+def _decode_router_mid(core, L: _LayerWeights, captured: bool):
"""A decode step's router projection for layer `L`: `mid_rows_mm` over its view in the
dense FP8 mode's rows (`_MID_MM_ROWS["router"]`), else at tp1 the step's projection entry
- (`core.decode_mm`) and at the other layouts `_decode_router`."""
+ (`core.decode_mm`) and under attention TP `_decode_router`.
+
+ Under attention DP every rank routes all the rows its group gathers, so every rank must
+ compute their logits with one kernel: two kernels can differ in the last bit and flip a
+ near-tie in the top-k, and the ranks would then disagree on which window serves a slot. A
+ `captured` step is a decode step on every rank (the engine captures only when every rank
+ can) and takes `_decode_router`; an eager one can meet another rank's prefill step and takes
+ the prefill router, `cublas_mm`."""
+ if core.dp_size > 1:
+ return _decode_router if captured else cublas_mm
fallback = core.decode_mm if core.one_rank else _decode_router
if L.router_m is None:
return fallback
@@ -3341,17 +3352,20 @@ def _nvfp4_context_plan(core, md, req_idx, block_table, tpb: int) -> tuple:
def _sparse_gqa_nvfp4_context(
core, kv, plan, q, block_table, selected, req_idx, logical, tpb, num_ctx, out, rows_per_request
) -> torch.Tensor:
- """Sparse GQA of a step with context rows over an NVFP4 pool at one rank, into `out` `[rows,
- q_width]`. With a plan (`_nvfp4_context_plan`) the first `num_ctx` rows, the context rows,
- take the row-tile pass over the plan's scratch, filled here from this layer's pool a group of
- requests at a time; the other rows, and every row without a plan, go split-K over the pool
- itself. The generation rows after the context rows come `rows_per_request` a request: more
- than one (a verify step's) takes the request-grouped entry for them on their own."""
+ """Sparse GQA of a step with context rows over an NVFP4 pool, into `out` `[rows, q_width]`.
+ With a plan (`_nvfp4_context_plan`) the first `num_ctx` rows, the context rows, take the
+ row-tile pass over the plan's scratch, filled here from this layer's pool a group of
+ requests at a time: at one rank the family's pass, elsewhere the target's, which also serves
+ a tp4 / tep4 rank's 6 / 1 heads. The other rows, and every row without a plan, go split-K
+ over the pool itself. The generation rows after the context rows come `rows_per_request` a
+ request: more than one (a verify step's) takes the request-grouped entry for them on their
+ own."""
rows = q.shape[0]
head = 0
if plan is not None:
groups, table, k_scratch, v_scratch = plan
assert groups[-1][1] == num_ctx, (groups[-1][1], num_ctx)
+ rowtile = one_rank_sparse_gqa_rowtile if core.one_rank else qsa_sparse_gqa_rowtile
for first_row, end_row, slots, counts, offsets in groups:
qsa_nvfp4_pages_to_bf16(
kv.k,
@@ -3367,7 +3381,7 @@ def _sparse_gqa_nvfp4_context(
k_scratch,
v_scratch,
)
- context = one_rank_sparse_gqa_rowtile(
+ context = rowtile(
q=q[first_row:end_row],
k_cache=k_scratch,
v_cache=v_scratch,
@@ -3682,7 +3696,7 @@ class _PrefillStep(NamedTuple):
vision: tuple | None
# A verify step: the generation rows' side-cache writes are staged for the rollback.
stage: bool = False
- # At one rank over an NVFP4 pool, `_nvfp4_context_plan`'s row-tile plan or None.
+ # Over an NVFP4 pool, `_nvfp4_context_plan`'s row-tile plan or None.
nvfp4_context: tuple | None = None
# Over an fp8 pool, the context requests' pages the row-tile pass widens
# (`_fp8_context_pages`); None otherwise.
@@ -3822,8 +3836,9 @@ def _attention_core(self, proj, proj_qkv, index_scratch, attn_out, i) -> None:
step.live_cols,
)
# Context rows with the prefill pair (a bound on a row's scan), generation rows
- # without.
- if core.one_rank and kv.k_scale is not None:
+ # without. Over an NVFP4 pool at one rank, or with the row-tile plan, the rows go
+ # through `_sparse_gqa_nvfp4_context`.
+ if kv.k_scale is not None and (core.one_rank or step.nvfp4_context is not None):
_sparse_gqa_nvfp4_context(
core,
kv,
@@ -4294,8 +4309,7 @@ def prepare(
fp8_pages = _fp8_context_pages(md, md.tokens_per_block)
nvfp4_context = None
if (
- core.one_rank
- and sparse
+ sparse
and core.quant_mode == QM_NVFP4_KV_CACHE
and num_prefill_tokens > _NVFP4_ROWTILE_MIN_ROWS
):
@@ -5560,14 +5574,15 @@ def _hc_moe_site(self, hc, block_out, inject, L: _LayerWeights):
mixed, inject, fp4 = _hc_moe_mixer(core, normed, L)
return hc, mixed, inject, fp4
- def _moe(self, x, L: _LayerWeights, pad_rows: int, fp4, fold: bool) -> tuple:
+ def _moe(self, x, L: _LayerWeights, pad_rows: int, fp4, fold: bool, captured: bool) -> tuple:
"""One decode MoE block, left uncombined for the next site: `(routed, (gate_logits,
shared, topk_weights, topk_rows))`. With `fold` (one rank, where the window is the whole
layer and needs no sum) the runner skips its finalize: `routed` is its permuted FC2 rows
and the top-k pair its routing weights and index map, which the next site's first kernel
- gathers; otherwise the pair is None and `routed` the finalized rows (`_moe_block`)."""
+ gathers; otherwise the pair is None and `routed` the finalized rows (`_moe_block`).
+ `captured` is the step's CUDA-graph state, which picks the router (`_decode_router_mid`)."""
core = self.core
- router_mm = _decode_router_mid(core, L)
+ router_mm = _decode_router_mid(core, L, captured)
if not fold:
routed, gate_logits, shared = _moe_block(
core, x, L, pad_rows, fp4, router_mm, core.decode_mm
@@ -5976,7 +5991,12 @@ def forward(
block_out = self._gdn_block(mixed, i, L, si_all, verify, gen_p, md.is_cuda_graph)
hc, mixed, inject, fp4 = self._hc_moe_site(hc, block_out, inject, L)
block_out, moe = self._moe(
- mixed, L, pad_rows, fp4, core.moe_fold and i + 1 != core.ple_layer
+ mixed,
+ L,
+ pad_rows,
+ fp4,
+ core.moe_fold and i + 1 != core.ple_layer,
+ md.is_cuda_graph,
)
# The stack-final site with no injection weight: its mix *is* the final norm.
hc, final, _ = self._hc_site(
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_sparse_gqa_rowtile__prefill.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_sparse_gqa_rowtile__prefill.py
index 62ec41278514..5c4b6b6ed4a1 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_sparse_gqa_rowtile__prefill.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_sparse_gqa_rowtile__prefill.py
@@ -663,6 +663,9 @@ def qsa_rowtile_sparse_gqa(
#: (`qsa_fp8_pages_to_bf16`) addressed through its own table.
_DEP4 = dict(q_heads=24, kv_heads=2, slots=32, **_TRUNK)
_DEP4_SCRATCH = dict(q_heads=24, kv_heads=2, slots=32, cache_layout="scratch")
+#: A tp4 / tep4 rank's quantized pool: an fp8 pool widened, or an NVFP4 pool dequantized
+#: (`qsa_nvfp4_pages_to_bf16`), into a contiguous `[pages, 1, 32, 256]` bf16 scratch.
+_TP4_SCRATCH = dict(q_heads=6, kv_heads=1, slots=128, cache_layout="scratch")
def _spec(segments, member, **overrides) -> dict:
@@ -732,6 +735,11 @@ class QsaSparseGqaRowtile(OpWrapper):
why="tp4/tep4 2049 rows from 9000: a last group of 9 rows, 6 q / 1 kv heads",
spec=_spec([(9000, 2049)], _TP4),
),
+ Cell(
+ why="tp4/tep4 kv_fp8 / kv_nvfp4 chunk: 8192 rows from 16384 over the bf16 scratch, "
+ "6 q / 1 kv heads (20-row groups), [128, 8196] table",
+ spec=_spec([(16384, 8192)], _TP4_SCRATCH),
+ ),
Cell(
why="dep4 kv_fp8 deep chunk: 8160 rows from 188416 over the widened bf16 "
"scratch, 24 q / 2 kv heads (10-row groups), [32, 8196] table",
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 a7b1d7a36598..598e03ee13bb 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
@@ -46,12 +46,13 @@
"norm_topk_prob": True,
"num_experts_per_tok": 10,
"hidden_act": "silu",
- # Attention: two KV heads of 256, no projection bias, a sigmoid output gate, and an output
- # head of its own.
+ # Attention: two KV heads of 256, no projection bias, a sigmoid output gate on every head
+ # (absent `attn_output_gate` reads as gated), and an output head of its own.
"num_key_value_heads": 2,
"head_dim": 256,
"attention_bias": False,
"output_gate_type": "sigmoid",
+ "attn_output_gate": True,
"tie_word_embeddings": False,
# Rope: a quarter of the head (64 dims), theta 1e7, interleaved mRoPE over sections
# 11 / 11 / 10, unscaled.
@@ -60,6 +61,15 @@
"mrope_section": (11, 11, 10),
"mrope_interleaved": True,
"rope_factor": None,
+ # The MTP module: one full-attention layer on the trunk's embedding that reads the end of
+ # the stack (`mtp_use_hidden_state_from_layer: null`). The `mtp.` fields are the `mtp`
+ # block's, a mapping or a config object; a field it omits reads as the value here. (Its
+ # own rope theta the target reads.)
+ "mtp_num_hidden_layers": 1,
+ "mtp_use_dedicated_embeddings": False,
+ "mtp.num_hidden_layers": 1,
+ "mtp.layer_types": ("full_attention",),
+ "mtp.mtp_use_hidden_state_from_layer": None,
}
# (depth, hidden, heads, merge, out_hidden, num_pos, deepstack_indexes, act).
@@ -135,6 +145,13 @@ def engine_choices(model_config: ModelConfig, llm_args: TorchLlmArgs) -> EngineC
)
+def _field(block, name: str, default=None):
+ """`block`'s `name` -- a mapping's entry or an object's attribute -- or `default` where it
+ is absent or null."""
+ value = block.get(name) if isinstance(block, dict) else getattr(block, name, None)
+ return default if value is None else value
+
+
def _text_settings(text) -> dict:
"""`_TEXT_SETTINGS`' entries as `text` declares them. The rope fields are the target's
reading: the `rope_parameters` dict, or the flattened `rope_scaling` with the theta and
@@ -149,6 +166,17 @@ def _text_settings(text) -> dict:
mrope_section=tuple(rope.get("mrope_section") or ()),
mrope_interleaved=rope.get("mrope_interleaved"),
rope_factor=rope.get("factor"),
+ attn_output_gate=_field(text, "attn_output_gate", True),
+ mtp_num_hidden_layers=_field(text, "mtp_num_hidden_layers", 1),
+ mtp_use_dedicated_embeddings=_field(text, "mtp_use_dedicated_embeddings", False),
+ )
+ mtp = getattr(text, "mtp", None)
+ settings.update(
+ {
+ "mtp.num_hidden_layers": _field(mtp, "num_hidden_layers", 1),
+ "mtp.layer_types": tuple(_field(mtp, "layer_types", ("full_attention",))),
+ "mtp.mtp_use_hidden_state_from_layer": _field(mtp, "mtp_use_hidden_state_from_layer"),
+ }
)
return settings
diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml
index 440f6c75d679..1cd64aada6a7 100644
--- a/tests/integration/defs/accuracy/references/acceptance_length.yaml
+++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml
@@ -245,7 +245,7 @@ ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3_fp8_kv:
min_al: 3.34
# Read by test_mtp3_gsm8k_quantized_kv's pass keys (3k-6k tokens, the sparse attention path):
# the NVFP4 cases at tp4 and tep4 (6 query heads a KV head) and the FP8 cases at dep4 and tep4.
-# modeling_v2's own readings on GB300 (NVFP4 3.951-3.968, FP8 3.960-3.975 over the cases and
+# modeling_v2's own readings on GB300 (NVFP4 3.949-3.968, FP8 3.960-3.975 over the cases and
# repeats); min_al 2.5% below, as for tp1_mtp3_long_context.
ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3_nvfp4_kv_long_context:
ref_al: 3.962
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_attention_plumbing.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_attention_plumbing.py
index 5dde8f4ddd7d..e9cd75616e6c 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_attention_plumbing.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_attention_plumbing.py
@@ -6,13 +6,15 @@
K/V at the unit scale, so a magnitude past the e4m3 maximum saturates. The probe reports the
largest magnitude of the first eager forward's attention layers, once, and never reads a
tensor under CUDA graph capture.
-* **The context pass over an NVFP4 pool, a group of requests at a time** (one rank). The
+* **The context pass over an NVFP4 pool, a group of requests at a time**, at one rank (24 / 2
+ heads, the family's row-tile pass) and at a tp4 / tep4 rank (6 / 1 heads, the target's). The
context rows read a bf16 scratch of their requests' pages (`_nvfp4_context_plan`), filled and
attended one group of consecutive requests at a time, so it holds the largest group's pages
rather than the step's. Under a budget of a few pages, with one request longer than the
budget, a prefill step's rows match one group over the whole step bitwise: its requests'
- rows come in multiples of the row-tile pass's 10-row groups, so a launch per group splits
- none of them and every row takes the same pass as in the one launch.
+ rows come in multiples of the row-tile pass's groups (10 rows at 24 / 2 heads, 20 at 6 / 1),
+ so a launch per group splits none of them and every row takes the same pass as in the one
+ launch. At 6 / 1 heads the pass agrees with the NVFP4 per-row prefill kernel.
* **MTP IndexShare.** The draft selection (`_selection`) in its fused and its two-launch form,
the capture of each request's draft-extend row into the shared buffers (decode and
mixed-batch forms), and the reuse entry's order over canonical captures. The targets are
@@ -122,47 +124,56 @@ def test_observing_under_capture_reads_nothing(monkeypatch):
# ---- the context pass over an NVFP4 pool ---------------------------------------------------
-TPB, HKV, D, QH, RATIO, TOPK = 32, 2, 256, 24, 4, 2048
-# (cached, new) tokens a context request: 10, 30, 1, 47 and 4 pages.
+TPB, D, RATIO, TOPK = 32, 256, 4, 2048
+# (cached, new) tokens a context request: 10, 30, 1 (2 at tp4), 47 and 4 pages. Each request's
+# rows are a multiple of the row-tile pass's groups: 10 rows at 24 / 2 heads, 20 at 6 / 1.
SEGMENTS = ((0, 300), (900, 40), (20, 10), (1400, 100), (64, 60))
+TP4_SEGMENTS = ((0, 300), (900, 40), (20, 20), (1400, 100), (64, 60))
+# (one rank, query heads, KV heads, segments, the whole step's pages)
+GEOMETRIES = {
+ "tp1": (True, 24, 2, SEGMENTS, 92),
+ "tp4": (False, 6, 1, TP4_SEGMENTS, 93),
+}
BUDGET_PAGES = 32 # request 3's 47 pages form a group of their own
-def _context_step(seed: int = 5):
+def _context_step(geometry: str, seed: int = 5):
+ one_rank, qh, hkv, segments, _ = GEOMETRIES[geometry]
gen = torch.Generator(device=DEV).manual_seed(seed)
- kv_lens = [s + n for s, n in SEGMENTS]
+ kv_lens = [s + n for s, n in segments]
pages = [math.ceil(k / TPB) for k in kv_lens]
pool_pages = sum(pages) + 3
- shape = (pool_pages, 2, HKV, TPB)
+ shape = (pool_pages, 2, hkv, TPB)
data = torch.randint(0, 256, (*shape, D // 2), generator=gen, device=DEV).to(torch.uint8)
scales = torch.randint(0x08, 0x31, (*shape, D // 16), generator=gen, device=DEV)
scales = scales.to(torch.uint8)
- slots = 2 * len(SEGMENTS)
+ slots = 2 * len(segments)
slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))
- slot_of = slot_of[: len(SEGMENTS)].to(DEV, torch.int32)
+ slot_of = slot_of[: len(segments)].to(DEV, torch.int32)
page_ids = torch.randperm(pool_pages, generator=gen, device=DEV).to(torch.int32)
block_table = torch.full((slots, max(pages)), -1, dtype=torch.int32, device=DEV)
cursor = 0
for slot, count in zip(slot_of.tolist(), pages):
block_table[slot, :count] = page_ids[cursor : cursor + count]
cursor += count
- new = torch.tensor([n for _, n in SEGMENTS], dtype=torch.int32)
+ new = torch.tensor([n for _, n in segments], dtype=torch.int32)
md = SimpleNamespace(
- num_contexts=len(SEGMENTS),
+ num_contexts=len(segments),
kv_lens_runtime=torch.tensor(kv_lens, dtype=torch.int32),
kv_lens_cuda_runtime=torch.tensor(kv_lens, dtype=torch.int32, device=DEV),
seq_lens=new,
seq_lens_cuda=new.to(DEV),
)
- logical = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in SEGMENTS])
+ logical = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments])
rows = logical.numel()
# Every visible token, as the expansion lays out a selection under its top-k.
column = torch.arange(TOPK + RATIO - 1, device=DEV)[None, :]
selected = torch.where(column <= logical[:, None], column, -1).to(torch.int32)
core = SimpleNamespace(
- heads_kv=HKV,
+ one_rank=one_rank,
+ heads_kv=hkv,
head_dim=D,
- q_width=QH * D,
+ q_width=qh * D,
kv_scale_quant_orig=torch.ones(3, dtype=torch.float32, device=DEV),
softmax_scale=D**-0.5,
compress_ratio=RATIO,
@@ -174,7 +185,7 @@ def _context_step(seed: int = 5):
md=md,
block_table=block_table,
req_idx=slot_of.repeat_interleave(new.to(DEV).long()),
- q=torch.randn(rows, QH, D, generator=gen, device=DEV).to(torch.bfloat16),
+ q=torch.randn(rows, qh, D, generator=gen, device=DEV).to(torch.bfloat16),
selected=selected,
logical=logical.to(torch.int32),
rows=rows,
@@ -185,7 +196,7 @@ def _context_pass(step):
plan = M._nvfp4_context_plan(
step["core"], step["md"], step["req_idx"], step["block_table"], TPB
)
- out = torch.empty(step["rows"], QH * D, dtype=torch.bfloat16, device=DEV)
+ out = torch.empty(step["rows"], step["core"].q_width, dtype=torch.bfloat16, device=DEV)
M._sparse_gqa_nvfp4_context(
step["core"],
step["kv"],
@@ -204,20 +215,45 @@ def _context_pass(step):
@_cuda
-def test_grouped_context_pass_equals_one_group(monkeypatch) -> None:
- step = _context_step()
+@pytest.mark.parametrize("geometry", sorted(GEOMETRIES))
+def test_grouped_context_pass_equals_one_group(monkeypatch, geometry) -> None:
+ step = _context_step(geometry)
whole, want = _context_pass(step)
- assert len(whole[0]) == 1 and whole[2].shape[0] == 92
+ assert len(whole[0]) == 1 and whole[2].shape[0] == GEOMETRIES[geometry][4]
- page_bytes = 2 * HKV * TPB * D * 2
+ page_bytes = 2 * step["core"].heads_kv * TPB * D * 2
monkeypatch.setattr(M, "_NVFP4_CONTEXT_SCRATCH_BYTES", BUDGET_PAGES * page_bytes)
(groups, _, k_scratch, _), got = _context_pass(step)
# Requests [0], [1, 2], [3] (over the budget, alone), [4].
- assert [(g[0], g[1]) for g in groups] == [(0, 300), (300, 350), (350, 450), (450, 510)]
+ ends = [0, 300]
+ for _, n in GEOMETRIES[geometry][3][1:]:
+ ends.append(ends[-1] + n)
+ del ends[2] # requests 1 and 2 share a group
+ assert [(g[0], g[1]) for g in groups] == list(zip(ends, ends[1:]))
assert k_scratch.shape[0] == 47
assert torch.equal(got.view(torch.int16), want.view(torch.int16))
+@_cuda
+def test_the_tp4_context_pass_matches_the_per_row_kernel() -> None:
+ """At a tp4 / tep4 rank (6 / 1 heads) the context rows take the target's row-tile pass over
+ the dequantized scratch, and agree with the NVFP4 per-row prefill kernel over the pool."""
+ step = _context_step("tp4")
+ _, got = _context_pass(step)
+ want = M._sparse_gqa(
+ step["core"],
+ step["kv"],
+ step["q"],
+ step["block_table"],
+ step["selected"],
+ step["req_idx"],
+ TPB,
+ step["logical"],
+ ).reshape(got.shape)
+ err = (got.float() - want.float()).norm() / want.float().norm()
+ assert err < 1e-2, err
+
+
# ---- MTP IndexShare -----------------------------------------------------------------------
K = 512
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_core_layout.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_core_layout.py
index b7edbfb9c3d7..f278f21a2131 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_core_layout.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_core_layout.py
@@ -8,7 +8,8 @@
what the forward addresses: the fused attention row and the split the norm + rope reads it
through, the Gated-DeltaNet projection's column split, the routed experts a rank holds (an
expert window, or tp4's padded intermediate slice of every expert) and the per-expert scalars
-cut to them, and the attention-DP round trip's padding and slicing. Under attention DP the
+cut to them, the attention-DP round trip's padding and slicing, and the one router kernel every
+attention-DP rank runs over the gathered rows whatever its own step. Under attention DP the
draft layer holds every attention head and runs its FP8 expert window over the rows the group
gathers, padded by the row counts the spec worker hands each draft step.
"""
@@ -129,11 +130,11 @@
_TP4 = dict(enable_attention_dp=False, moe_tp_size=4, moe_ep_size=1)
-def _core(rank: int, layout: dict, spec_config=None, world: int = WORLD, **config):
+def _core(rank: int, layout: dict, spec_config=None, world: int = WORLD, text=None, **config):
"""The core of rank `rank` of a `world`-rank engine at `layout`, every parameter on meta;
- `config` sets further `ModelConfig` fields."""
+ `text` overrides text config fields and `config` sets further `ModelConfig` fields."""
model_config = ModelConfig(
- pretrained_config=Qwen4ExpTextConfig(**_TEXT),
+ pretrained_config=Qwen4ExpTextConfig(**{**_TEXT, **(text or {})}),
mapping=Mapping(world_size=world, tp_size=world, rank=rank, **layout),
quant_config=QuantConfig(quant_algo=QuantAlgo.MIXED_PRECISION),
moe_backend="TRTLLM",
@@ -379,6 +380,24 @@ def fake_reducescatter(t, sizes, group, process_group):
assert tuple(out.shape) == (counts[rank], hidden)
+@pytest.mark.parametrize("captured", [False, True])
+def test_every_attention_dp_rank_routes_the_gathered_rows_with_one_kernel(captured) -> None:
+ """Each attention-DP rank computes the router logits of every row its group gathers, so the
+ kernel cannot depend on whether this rank's step is a prefill or a decode: an eager decode
+ step takes the prefill router (`_moe_block`'s default), a captured one -- a decode step on
+ every rank -- `_decode_router` on every rank. Under attention TP every rank runs the same
+ step, and a decode step keeps `_decode_router`."""
+ prefill_router = inspect.signature(M._moe_block).parameters["router_mm"].default
+ assert prefill_router is M.cublas_mm
+ layer = SimpleNamespace(router_m=None)
+ for dp_size in (4, 2):
+ core = SimpleNamespace(dp_size=dp_size, one_rank=False, decode_mm=None)
+ want = M._decode_router if captured else prefill_router
+ assert M._decode_router_mid(core, layer, captured) is want
+ core = SimpleNamespace(dp_size=1, one_rank=False, decode_mm=None)
+ assert M._decode_router_mid(core, layer, captured) is M._decode_router
+
+
def nn_pad(x: torch.Tensor, rows: int) -> torch.Tensor:
"""`x` with zero rows appended up to `rows`."""
return torch.cat([x, x.new_zeros(rows - x.shape[0], x.shape[1])], 0)
@@ -402,6 +421,17 @@ def test_mtp_widens_the_verify_step_and_keys_the_draft_layer_after_the_trunk(
assert not any("_mtp_" in k for k in plain.w)
+@pytest.mark.parametrize("form", ["mapping", "object"])
+def test_the_draft_layer_ropes_at_its_own_blocks_theta(form) -> None:
+ """The `mtp` block's own rope theta sets the draft layer's base, whether the block is a
+ mapping or a config object."""
+ block = dict(_TEXT["mtp"], rope_theta=5e6)
+ if form == "object":
+ block = SimpleNamespace(**block)
+ core = _core(0, _TEP4, MTPDecodingConfig(max_draft_len=3), text=dict(mtp=block))
+ assert (core.mtp_theta, core.theta) == (5e6, 1e7)
+
+
@pytest.mark.parametrize("rank", range(WORLD))
def test_a_tep4_rank_holds_its_share_of_the_draft_layer(rank) -> None:
"""The draft layer's attention is cut like a trunk attention layer's: six q heads (their
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py
index b766831b715a..caec32aee91d 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py
@@ -80,6 +80,16 @@ def _llm_args(mode="auto", **kwargs) -> TorchLlmArgs:
return TorchLlmArgs(model=_DUMMY_MODEL, modeling_v2=mode, **kwargs)
+# The checkpoint's `mtp` block.
+_MTP = dict(
+ hybrid=True,
+ layer_types=["full_attention"],
+ mtp_use_hidden_state_from_layer=None,
+ num_hidden_layers=1,
+ rope_theta=10000000,
+)
+
+
def _qwen4_exp_config(text_overrides=None, **vision_overrides):
"""The shape Qwen3.8-Flash-Next-NVFP4's normalized composite config declares."""
text = dict(
@@ -102,6 +112,9 @@ def _qwen4_exp_config(text_overrides=None, **vision_overrides):
rope_theta=10000000,
partial_rotary_factor=0.25,
rope_scaling=dict(type="mrope", mrope_section=[11, 11, 10], mrope_interleaved=True),
+ mtp=dict(_MTP),
+ mtp_num_hidden_layers=1,
+ mtp_use_dedicated_embeddings=False,
)
text.update(text_overrides or {})
vision = dict(
@@ -351,16 +364,41 @@ def test_another_checkpoint_shape_declines():
(dict(rope_scaling=dict(_ROPE, mrope_interleaved=False)), "mrope_interleaved"),
(dict(rope_scaling=dict(_ROPE, mrope_section=[16, 8, 8])), "mrope_section"),
(dict(rope_scaling=dict(_ROPE, factor=4.0)), "rope_factor"),
+ (dict(attn_output_gate=False), "attn_output_gate"),
+ (dict(mtp_num_hidden_layers=2), "mtp_num_hidden_layers"),
+ (dict(mtp_use_dedicated_embeddings=True), "mtp_use_dedicated_embeddings"),
+ (dict(mtp=dict(_MTP, num_hidden_layers=2)), "mtp.num_hidden_layers"),
+ (dict(mtp=dict(_MTP, layer_types=["linear_attention"])), "mtp.layer_types"),
+ (
+ dict(mtp=dict(_MTP, mtp_use_hidden_state_from_layer=40)),
+ "mtp.mtp_use_hidden_state_from_layer",
+ ),
],
ids=lambda v: v if isinstance(v, str) else None,
)
def test_a_same_shape_checkpoint_with_other_settings_declines(override, field):
"""A fine-tune with the checkpoint's shape that changes what it computes, which the target
- would run with wrong logits and no error: `auto` keeps the built-in model."""
+ would run with wrong logits (or wrong drafts) and no error: `auto` keeps the built-in
+ model."""
config = _model_config(pretrained_config=_qwen4_exp_config(override), **_LAYOUTS["tep4"])
_declined(config, match=f"'{field}'")
+def test_the_mtp_block_is_read_as_a_mapping_or_a_config_object():
+ """An `mtp` block that is a config object is read field by field as a mapping is, and
+ one that omits the block's fields (or the whole block) takes the checkpoint's values."""
+ as_object = PretrainedConfig(**_MTP)
+ tep4 = _LAYOUTS["tep4"]
+ _claimed(_model_config(pretrained_config=_qwen4_exp_config(dict(mtp=as_object)), **tep4))
+ as_object.mtp_use_hidden_state_from_layer = 40
+ _declined(
+ _model_config(pretrained_config=_qwen4_exp_config(dict(mtp=as_object)), **tep4),
+ match="'mtp.mtp_use_hidden_state_from_layer'",
+ )
+ bare = dict(mtp=None, mtp_num_hidden_layers=None, mtp_use_dedicated_embeddings=None)
+ _claimed(_model_config(pretrained_config=_qwen4_exp_config(bare), **tep4))
+
+
def test_the_settings_read_off_the_rope_parameters_dict():
"""A config that keeps its rope fields in `rope_parameters` (theta and partial factor
inside) is read the same way, as the target reads it."""
From 616104461ea43d81f15ae84c3679a1595284454e Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Thu, 5 Nov 2026 21:12:00 -0800
Subject: [PATCH 11/15] [None][chore] modeling_v2 qwen4_exp: keep route short,
move the admission checks onto the target
route now decides on the GPU, the parallel layout (one _LAYOUTS table) and the draft path.
The checkpoint shape and settings checks and the engine-choice checks move to the target's
check_deployment staticmethod, next to the constants they read; route calls it before it
returns the class, so auto still falls back on everything it declined before.
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
.../modeling.py | 235 +++++++++-
.../vision.py | 2 +-
.../modeling_v2/models/qwen4_exp/routing.py | 423 +++---------------
.../test_routing.py | 29 +-
.../modeling_v2/test_modeling_v2_routing.py | 12 +-
5 files changed, 305 insertions(+), 396 deletions(-)
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py
index 79159b457996..2807813eae62 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py
@@ -140,13 +140,17 @@
import math
import os
import re
-from typing import NamedTuple
+from typing import TYPE_CHECKING, NamedTuple
import torch
from torch import nn
from transformers import PretrainedConfig
-from tensorrt_llm._torch._experimental.modeling_v2._router_index import step_contract_enabled
+from tensorrt_llm._torch._experimental.modeling_v2._router_index import (
+ NoTarget,
+ engine_kv_cache_quant_algo,
+ step_contract_enabled,
+)
from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import (
fused_qk_norm_rope,
)
@@ -433,9 +437,10 @@
flashinfer_cached_replay_update,
)
from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor
+from tensorrt_llm._torch.pyexecutor.model_loader import effective_mamba_ssm_cache_dtype
from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks
from tensorrt_llm._torch.speculative import get_spec_worker, should_use_separate_draft_kv_cache
-from tensorrt_llm._torch.utils import ActivationType
+from tensorrt_llm._torch.utils import ActivationType, is_gdn_replay_enabled
from tensorrt_llm._utils import mpi_disabled
from tensorrt_llm.functional import AllReduceFusionOp
from tensorrt_llm.logger import logger
@@ -454,6 +459,9 @@
from .ops.ssm.gdn_cached_replay_cute__verify import gdn_cached_replay_cute
from .ops.ssm.gdn_cached_replay_cute__verify import precompile as gdn_replay_precompile
+if TYPE_CHECKING:
+ from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
+
# The GPU architecture this target is certified for; a direct instantiation could
# bypass routing.
_SM = (10, 3)
@@ -476,6 +484,48 @@
_SSM_STATE_DTYPES = (torch.bfloat16,)
_ONE_RANK_SSM_STATE_DTYPES = (torch.float32, torch.bfloat16)
+# The checkpoint this target is written for, read off the normalized composite's halves
+# (`check_deployment`): the text tower's (layers, hidden, experts, attention heads, linear value
+# heads, vocab) and the vision tower's (depth, hidden, heads, merge, out_hidden, num_pos,
+# deepstack_indexes, act). A deepstack vision sibling has a merger list this target lacks.
+_TEXT_SHAPE = (48, 2560, 512, 24, 48, 248320)
+_VISION_SHAPE = (27, 1152, 16, 2, 2560, 2304, (), "gelu_pytorch_tanh")
+
+# The checkpoint's settings that change what it computes without changing a weight's shape,
+# at the values the target implements: a fine-tune of the same shape that changed one would
+# load and run with wrong logits and no error. Read off the text config as the target reads
+# them (`_text_settings`).
+_TEXT_SETTINGS = {
+ # The routed experts: Renormalize routing over the top 10 and a SiLU activation.
+ "norm_topk_prob": True,
+ "num_experts_per_tok": 10,
+ "hidden_act": "silu",
+ # Attention: two KV heads of 256, no projection bias, a sigmoid output gate on every head
+ # (absent `attn_output_gate` reads as gated), and an output head of its own.
+ "num_key_value_heads": 2,
+ "head_dim": 256,
+ "attention_bias": False,
+ "output_gate_type": "sigmoid",
+ "attn_output_gate": True,
+ "tie_word_embeddings": False,
+ # Rope: a quarter of the head (64 dims), theta 1e7, interleaved mRoPE over sections
+ # 11 / 11 / 10, unscaled.
+ "rope_theta": 1e7,
+ "partial_rotary_factor": 0.25,
+ "mrope_section": (11, 11, 10),
+ "mrope_interleaved": True,
+ "rope_factor": None,
+ # The MTP module: one full-attention layer on the trunk's embedding that reads the end of
+ # the stack (`mtp_use_hidden_state_from_layer: null`). The `mtp.` fields are the `mtp`
+ # block's, a mapping or a config object; a field it omits reads as the value here. (Its
+ # own rope theta the target reads.)
+ "mtp_num_hidden_layers": 1,
+ "mtp_use_dedicated_embeddings": False,
+ "mtp.num_hidden_layers": 1,
+ "mtp.layer_types": ("full_attention",),
+ "mtp.mtp_use_hidden_state_from_layer": None,
+}
+
# QSA backend metadata buffers rebuilt every step. Their absence means the engine did not
# select the QSA metadata and the index plane would never be written.
_QSA_FIELDS = (
@@ -6671,6 +6721,42 @@ def _spec_argmax_only(spec_worker, spec_metadata) -> tuple[bool, bool]:
return verify, draft
+def _field(block, name: str, default=None):
+ """`block`'s `name` -- a mapping's entry or an object's attribute -- or `default` where it
+ is absent or null."""
+ value = block.get(name) if isinstance(block, dict) else getattr(block, name, None)
+ return default if value is None else value
+
+
+def _text_settings(text) -> dict:
+ """`_TEXT_SETTINGS`' entries as `text` declares them. The rope fields are the target's
+ reading: the `rope_parameters` dict, or the flattened `rope_scaling` with the theta and
+ the partial factor at the top level."""
+ rope = getattr(text, "rope_parameters", None) or getattr(text, "rope_scaling", None) or {}
+ settings = {name: getattr(text, name, None) for name in _TEXT_SETTINGS}
+ settings.update(
+ rope_theta=rope.get("rope_theta", getattr(text, "rope_theta", None)),
+ partial_rotary_factor=rope.get(
+ "partial_rotary_factor", getattr(text, "partial_rotary_factor", None)
+ ),
+ mrope_section=tuple(rope.get("mrope_section") or ()),
+ mrope_interleaved=rope.get("mrope_interleaved"),
+ rope_factor=rope.get("factor"),
+ attn_output_gate=_field(text, "attn_output_gate", True),
+ mtp_num_hidden_layers=_field(text, "mtp_num_hidden_layers", 1),
+ mtp_use_dedicated_embeddings=_field(text, "mtp_use_dedicated_embeddings", False),
+ )
+ mtp = getattr(text, "mtp", None)
+ settings.update(
+ {
+ "mtp.num_hidden_layers": _field(mtp, "num_hidden_layers", 1),
+ "mtp.layer_types": tuple(_field(mtp, "layer_types", ("full_attention",))),
+ "mtp.mtp_use_hidden_state_from_layer": _field(mtp, "mtp_use_hidden_state_from_layer"),
+ }
+ )
+ return settings
+
+
class ModelingV2Qwen38FlashNextNvfp4Sm103Default(
DecoderModelForCausalLM[Qwen4ExpDefaultModelingV2, PretrainedConfig]
):
@@ -6689,6 +6775,149 @@ def get_preferred_kv_cache_manager_version(cls, pretrained_config=None) -> str:
position planes exist only there."""
return "V2"
+ @staticmethod
+ def check_deployment(model_config: ModelConfig, llm_args: TorchLlmArgs, layout: str) -> None:
+ """Raise `NoTarget` for a checkpoint or an engine choice this target does not run.
+
+ `route` calls it once the GPU, the parallel layout (`layout`, a key of routing's
+ `_LAYOUTS`) and the draft path match, before it returns this class. Where the
+ constructor, the bind or the first forward would refuse the deployment, or a same-shape
+ checkpoint would compute something else, `auto` then falls back to the built-in
+ implementation.
+ """
+ cfg = model_config.pretrained_config
+ text = getattr(cfg, "text_config", None)
+ vision = getattr(cfg, "vision_config", None)
+ if text is None or vision is None:
+ raise NoTarget(
+ f"{type(cfg).__name__} carries no text_config/vision_config; the target reads "
+ "both halves of the normalized Qwen4-Exp composite"
+ )
+ shape = (
+ text.num_hidden_layers,
+ text.hidden_size,
+ text.num_experts,
+ text.num_attention_heads,
+ text.linear_num_value_heads,
+ text.vocab_size,
+ )
+ if shape != _TEXT_SHAPE:
+ raise NoTarget(f"text shape {shape}; the target is Qwen3.8-Flash-Next's {_TEXT_SHAPE}")
+ vshape = (
+ vision.depth,
+ vision.hidden_size,
+ vision.num_heads,
+ vision.spatial_merge_size,
+ vision.out_hidden_size,
+ vision.num_position_embeddings,
+ tuple(getattr(vision, "deepstack_visual_indexes", None) or ()),
+ vision.hidden_act,
+ )
+ if vshape != _VISION_SHAPE:
+ raise NoTarget(
+ f"vision shape {vshape}; the target is Qwen3.8-Flash-Next's {_VISION_SHAPE}"
+ )
+ differs = {k: v for k, v in _text_settings(text).items() if v != _TEXT_SETTINGS[k]}
+ if differs:
+ raise NoTarget(
+ f"the checkpoint has Qwen3.8-Flash-Next's shape but declares {differs}; the "
+ f"target implements { ({k: _TEXT_SETTINGS[k] for k in differs}) }"
+ )
+
+ mapping = model_config.mapping
+ one_rank = mapping.world_size == 1
+ spec = model_config.spec_config
+ prefill_graph = str(llm_args.prefill_cuda_graph_backend.value)
+ # The draft loop runs the checkpoint's single MTP module as one linear chain, verified
+ # at one width, `max_draft_len + 1` rows a request.
+ if spec is not None:
+ if (
+ getattr(spec, "num_nextn_predict_layers", None) not in (None, 1)
+ or getattr(spec, "use_dynamic_tree", False)
+ or getattr(spec, "eagle_choices", None) is not None
+ or spec.max_total_draft_tokens != spec.max_draft_len
+ ):
+ raise NoTarget(
+ "an MTP draft other than one linear chain over the checkpoint's single MTP "
+ "module"
+ )
+ if not 1 <= spec.max_draft_len < _MAX_VERIFY_WIDTH:
+ raise NoTarget(
+ f"max_draft_len={spec.max_draft_len}; the verify step is certified at "
+ f"max_draft_len 1..{_MAX_VERIFY_WIDTH - 1}"
+ )
+ if (
+ getattr(spec, "draft_len_schedule", None) is not None
+ or getattr(spec, "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 mapping.enable_attention_dp and prefill_graph == "breakable":
+ raise NoTarget(
+ f"prefill_cuda_graph_backend='breakable' with attention DP ({layout}): the engine "
+ "pads every rank to one prefill bucket, so a rank holding only generation "
+ "requests is handed more rows than its batch, which the target's decode step "
+ "refuses"
+ )
+ kv_cache_quant_algo = engine_kv_cache_quant_algo(model_config, llm_args)
+ 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 model_config.torch_dtype != torch.bfloat16:
+ raise NoTarget(
+ f"the engine runs the model in {model_config.torch_dtype}; the target is bf16"
+ )
+ # The CuteDSL expert pipeline runs every expert at once: one rank only.
+ requested = str(model_config.moe_backend_requested).upper()
+ moe_backends = [
+ name for name, runs in _TRUNK_MOE_BACKENDS.items() if one_rank or runs != "CUTEDSL"
+ ]
+ if requested not in moe_backends:
+ raise NoTarget(
+ f"moe_config.backend={model_config.moe_backend_requested!r}; at {layout} the "
+ f"target runs the trunk's routed experts on {', '.join(moe_backends)}"
+ )
+ state_dtype = effective_mamba_ssm_cache_dtype(
+ model_config, llm_args.kv_cache_config.mamba_ssm_cache_dtype
+ )
+ state_dtypes = _ONE_RANK_SSM_STATE_DTYPES if one_rank else _SSM_STATE_DTYPES
+ if state_dtype not in state_dtypes:
+ served = " or ".join(str(dtype).removeprefix("torch.") for dtype in state_dtypes)
+ raise NoTarget(
+ f"a {state_dtype} recurrent-state cache; at {layout} the target serves {served} "
+ "(kv_cache_config.mamba_ssm_cache_dtype)"
+ )
+ if llm_args.moe_config.load_balancer is not None and mapping.moe_ep_size > 1:
+ raise NoTarget(
+ f"moe_config.load_balancer (EPLB) at {layout}: the target runs fixed contiguous "
+ "expert windows and registers no layer with the balancer, so the placement "
+ "would do nothing"
+ )
+ if prefill_graph == "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 is None:
+ return
+ if not is_gdn_replay_enabled():
+ raise NoTarget(
+ "MTP with TRTLLM_USE_GDN_REPLAY=0: the target verifies its Gated-DeltaNet rows "
+ "on the cached replay only, not on the legacy per-step recurrent states; leave "
+ "TRTLLM_USE_GDN_REPLAY unset"
+ )
+ if requested == "CUTEDSL" and spec.max_draft_len > _CUTEDSL_MAX_DRAFT_LEN:
+ raise NoTarget(
+ f"max_draft_len={spec.max_draft_len} with moe_config.backend=CUTEDSL: the CuteDSL "
+ f"expert pipeline is certified at max_draft_len up to {_CUTEDSL_MAX_DRAFT_LEN}"
+ )
+
def __init__(self, model_config: ModelConfig):
cfg = model_config.pretrained_config
assert cfg is not None
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision.py
index 0aa802cb3803..818d919c6551 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision.py
@@ -318,7 +318,7 @@ class VisionTower:
def __init__(self, w: dict, vision_config) -> None:
"""`w` is the core's weight table after load; `vision_config` the composite
- config's vision half, which `routing.py` matched against this checkpoint's."""
+ config's vision half, which `check_deployment` matched against this checkpoint's."""
vcfg = vision_config
self.depth = vcfg.depth
self.hidden = vcfg.hidden_size
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 598e03ee13bb..1ae51c1cc9e0 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py
@@ -2,28 +2,20 @@
# SPDX-License-Identifier: Apache-2.0
"""Routing for Qwen4ExpForConditionalGeneration.
-`route` compares the facts that fix which code runs against the deployments the
-targets were written for, and raises `NoTarget` stating them otherwise (`auto` falls
-back to the built-in implementation; `require` raises). The checkpoint's published
-architecture name is kept by `load_pretrained_config`, so `route` registers under it.
-All criteria are read off `text_config` / `vision_config`; the composite parent's
-mirrored fields are unreliable.
-
-Qwen3.8-Flash-Next-NVFP4 on sm103 lands on `qwen3_8_flash_next_nvfp4__sm_103__default`
-at one of its parallel layouts (`layout`): tp1 on one rank, and tp4, tep4, dep4 or dep2 in
-one node.
+Qwen3.8-Flash-Next-NVFP4 on sm103 lands on `qwen3_8_flash_next_nvfp4__sm_103__default` at one
+of the layouts in `_LAYOUTS`, with one-model MTP-Eagle (not at dep2) or no speculative
+decoding. The target's `check_deployment` then declines the checkpoints and engine choices it
+does not run, so `auto` falls back to the built-in implementation instead of failing in the
+build.
"""
from __future__ import annotations
-from typing import TYPE_CHECKING, NamedTuple, Optional, Tuple
-
-import torch
+from typing import TYPE_CHECKING, Optional, Tuple
from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig
-from tensorrt_llm.quantization.mode import QuantAlgo
-from ..._router_index import NoTarget, engine_kv_cache_quant_algo, routes
+from ..._router_index import NoTarget, routes
if TYPE_CHECKING:
from tensorrt_llm._torch.model_config import ModelConfig
@@ -33,72 +25,10 @@
# The GPU the target is certified on.
_SM = (10, 3)
-# (layers, hidden, experts, attention heads, linear value heads, vocab). Shape-sniffing
-# also matches fine-tunes; the accuracy gate pins identity.
-_TEXT_SHAPE = (48, 2560, 512, 24, 48, 248320)
-
-# The checkpoint's settings that change what it computes without changing a weight's shape,
-# at the values the target implements: a fine-tune of the same shape that changed one would
-# load and run with wrong logits and no error. Read off the text config as the target reads
-# them (`_text_settings`).
-_TEXT_SETTINGS = {
- # The routed experts: Renormalize routing over the top 10 and a SiLU activation.
- "norm_topk_prob": True,
- "num_experts_per_tok": 10,
- "hidden_act": "silu",
- # Attention: two KV heads of 256, no projection bias, a sigmoid output gate on every head
- # (absent `attn_output_gate` reads as gated), and an output head of its own.
- "num_key_value_heads": 2,
- "head_dim": 256,
- "attention_bias": False,
- "output_gate_type": "sigmoid",
- "attn_output_gate": True,
- "tie_word_embeddings": False,
- # Rope: a quarter of the head (64 dims), theta 1e7, interleaved mRoPE over sections
- # 11 / 11 / 10, unscaled.
- "rope_theta": 1e7,
- "partial_rotary_factor": 0.25,
- "mrope_section": (11, 11, 10),
- "mrope_interleaved": True,
- "rope_factor": None,
- # The MTP module: one full-attention layer on the trunk's embedding that reads the end of
- # the stack (`mtp_use_hidden_state_from_layer: null`). The `mtp.` fields are the `mtp`
- # block's, a mapping or a config object; a field it omits reads as the value here. (Its
- # own rope theta the target reads.)
- "mtp_num_hidden_layers": 1,
- "mtp_use_dedicated_embeddings": False,
- "mtp.num_hidden_layers": 1,
- "mtp.layer_types": ("full_attention",),
- "mtp.mtp_use_hidden_state_from_layer": None,
-}
-
-# (depth, hidden, heads, merge, out_hidden, num_pos, deepstack_indexes, act).
-# `deepstack_visual_indexes == ()` is the key member: a deepstack sibling has a merger
-# list these targets lack.
-_VISION_SHAPE = (27, 1152, 16, 2, 2560, 2304, (), "gelu_pytorch_tanh")
-
-# The longest MTP draft the target runs (`_MAX_VERIFY_WIDTH - 1` in modeling.py).
-_MAX_DRAFT_LEN = 7
-
-# The engine choices the target implements at one rank (tp1), all in its modeling.py: the
-# `moe_config.backend` requests (`_TRUNK_MOE_BACKENDS`), the longest draft of the CUTEDSL
-# experts (`_CUTEDSL_MAX_DRAFT_LEN`) and the recurrent-state cache dtypes
-# (`_ONE_RANK_SSM_STATE_DTYPES`).
-_ONE_RANK_MOE_BACKENDS = ("AUTO", "TRTLLM", "CUTEDSL")
-_MAX_CUTEDSL_DRAFT_LEN = 3
-_ONE_RANK_SSM_STATE_DTYPES = (torch.float32, torch.bfloat16)
-
-# Their counterparts at every other layout: no CUTEDSL experts (`_TRUNK_MOE_BACKENDS` less
-# CUTEDSL) and a bf16 recurrent state (`_SSM_STATE_DTYPES`).
-_MOE_BACKENDS = ("AUTO", "TRTLLM")
-_SSM_STATE_DTYPES = (torch.bfloat16,)
-
-# (world size, attention DP, MoE TP, MoE EP) -> the target's layout. Every layout spans its
-# whole world in one node, attention TP over all of it unless attention DP replicates
-# attention. One rank (tp1) holds everything. At four ranks the routed experts are split by
-# intermediate (tp4) or by expert (tep4, dep4). At two ranks only attention DP with two
-# 256-expert windows (dep2), where each rank's attention, Gated-DeltaNet and PLE geometry is
-# dep4's.
+# (world size, attention DP, MoE TP, MoE EP) -> the target's layout, each in one node. tp1 holds
+# everything on one rank; tp4 splits the routed experts by intermediate, tep4 and dep4 by
+# expert; dep2 is attention DP with two 256-expert windows. MTP runs at every layout but dep2,
+# where no draft acceptance was measured.
_LAYOUTS = {
(1, False, 1, 1): "tp1",
(4, False, 4, 1): "tp4",
@@ -107,307 +37,66 @@
(2, True, 1, 2): "dep2",
}
-# The attention-DP layouts MTP is admitted at: those its draft acceptance was measured at.
-_ATTENTION_DP_MTP_LAYOUTS = ("dep4",)
-
-# The layouts with expert parallelism, where the engine builds an EPLB balancer for
-# `moe_config.load_balancer`; the target cuts fixed contiguous expert windows.
-_EXPERT_PARALLEL_LAYOUTS = tuple(member for key, member in _LAYOUTS.items() if key[3] > 1)
-
-
-class EngineChoices(NamedTuple):
- """Engine choices that fix which code a target runs but are not in the checkpoint's
- config, as the engine will resolve them."""
-
- torch_dtype: torch.dtype # the dtype the engine runs the model in
- 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
- load_balancer: bool # `moe_config.load_balancer` is set (EPLB)
-
-
-def engine_choices(model_config: ModelConfig, llm_args: TorchLlmArgs) -> EngineChoices:
- # Imported on a claimed branch only, as the targets are.
- from tensorrt_llm._torch.pyexecutor.model_loader import effective_mamba_ssm_cache_dtype
- from tensorrt_llm._torch.utils import is_gdn_replay_enabled
-
- 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),
- load_balancer=llm_args.moe_config.load_balancer is not None,
- )
-
-
-def _field(block, name: str, default=None):
- """`block`'s `name` -- a mapping's entry or an object's attribute -- or `default` where it
- is absent or null."""
- value = block.get(name) if isinstance(block, dict) else getattr(block, name, None)
- return default if value is None else value
-
-
-def _text_settings(text) -> dict:
- """`_TEXT_SETTINGS`' entries as `text` declares them. The rope fields are the target's
- reading: the `rope_parameters` dict, or the flattened `rope_scaling` with the theta and
- the partial factor at the top level."""
- rope = getattr(text, "rope_parameters", None) or getattr(text, "rope_scaling", None) or {}
- settings = {name: getattr(text, name, None) for name in _TEXT_SETTINGS}
- settings.update(
- rope_theta=rope.get("rope_theta", getattr(text, "rope_theta", None)),
- partial_rotary_factor=rope.get(
- "partial_rotary_factor", getattr(text, "partial_rotary_factor", None)
- ),
- mrope_section=tuple(rope.get("mrope_section") or ()),
- mrope_interleaved=rope.get("mrope_interleaved"),
- rope_factor=rope.get("factor"),
- attn_output_gate=_field(text, "attn_output_gate", True),
- mtp_num_hidden_layers=_field(text, "mtp_num_hidden_layers", 1),
- mtp_use_dedicated_embeddings=_field(text, "mtp_use_dedicated_embeddings", False),
- )
- mtp = getattr(text, "mtp", None)
- settings.update(
- {
- "mtp.num_hidden_layers": _field(mtp, "num_hidden_layers", 1),
- "mtp.layer_types": tuple(_field(mtp, "layer_types", ("full_attention",))),
- "mtp.mtp_use_hidden_state_from_layer": _field(mtp, "mtp_use_hidden_state_from_layer"),
- }
- )
- return settings
-
def layout(mapping: Mapping) -> Optional[str]:
- """The target's layout ``mapping`` describes (``tp1``, ``tp4``, ``tep4``, ``dep4``,
- ``dep2``), or None.
+ """The layout ``mapping`` describes, or None.
- Pipeline and context parallelism, an attention TP narrower than the TP group, MoE
- clusters, lm_head TP under attention DP and a group spanning nodes are other
- forwards.
+ Attention runs TP over the whole world, or DP with the LM head whole. Pipeline and context
+ parallelism, MoE clusters and a world spanning nodes are other forwards.
"""
- if mapping.tp_size != mapping.world_size or (mapping.pp_size, mapping.cp_size) != (1, 1):
- return None
- if mapping.is_multi_node() or mapping.moe_cluster_size != 1:
- return None
if mapping.enable_attention_dp:
- if mapping.enable_lm_head_tp_in_adp:
- return None
- elif mapping.attn_tp_size != mapping.tp_size:
+ attention_as_built = not mapping.enable_lm_head_tp_in_adp
+ else:
+ attention_as_built = mapping.attn_tp_size == mapping.tp_size
+ if (
+ not attention_as_built
+ or mapping.tp_size != mapping.world_size
+ or (mapping.pp_size, mapping.cp_size, mapping.moe_cluster_size) != (1, 1, 1)
+ or mapping.is_multi_node()
+ ):
return None
return _LAYOUTS.get(
- (
- mapping.world_size,
- mapping.enable_attention_dp,
- mapping.moe_tp_size,
- mapping.moe_ep_size,
- )
- )
-
-
-def _describe(mapping: Mapping) -> str:
- return (
- f"world_size={mapping.world_size} (tp={mapping.tp_size}, pp={mapping.pp_size}, "
- f"cp={mapping.cp_size}, attention_dp={mapping.enable_attention_dp}, "
- f"moe_tp={mapping.moe_tp_size}, moe_ep={mapping.moe_ep_size})"
+ (mapping.world_size, mapping.enable_attention_dp, mapping.moe_tp_size, mapping.moe_ep_size)
)
-def _check_engine(
- spec_config: Optional[MTPDecodingConfig],
- kv_cache_quant_algo: Optional[QuantAlgo],
- engine: EngineChoices,
- member: str,
-) -> None:
- """Decline what the target's constructor or first forward would refuse at layout
- `member`: at tp1 the CuteDSL experts and the fp32 recurrent state are served too.
-
- Without this ``auto`` claims the deployment and then fails mid-build instead of
- falling back to the built-in implementation.
- """
- if member == "tp1":
- moe_backends, ssm_state_dtypes = _ONE_RANK_MOE_BACKENDS, _ONE_RANK_SSM_STATE_DTYPES
- else:
- moe_backends, ssm_state_dtypes = _MOE_BACKENDS, _SSM_STATE_DTYPES
- if spec_config is not None and (
- getattr(spec_config, "draft_len_schedule", None) is not None
- or getattr(spec_config, "acceptance_rate_threshold", None) is not None
- ):
- 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.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}; at {member} the target runs the "
- f"trunk's routed experts on {', '.join(moe_backends)}"
- )
- if engine.mamba_ssm_cache_dtype not in ssm_state_dtypes:
- served = " or ".join(str(dtype).removeprefix("torch.") for dtype in ssm_state_dtypes)
- raise NoTarget(
- f"a {engine.mamba_ssm_cache_dtype} recurrent-state cache; at {member} the target "
- f"serves {served} (kv_cache_config.mamba_ssm_cache_dtype)"
- )
- if engine.load_balancer and member in _EXPERT_PARALLEL_LAYOUTS:
- raise NoTarget(
- f"moe_config.load_balancer (EPLB) at {member}: the target runs fixed contiguous "
- "expert windows and registers no layer with the balancer, so the placement would "
- "do nothing"
- )
- 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:
- raise NoTarget(
- "MTP with TRTLLM_USE_GDN_REPLAY=0: the target verifies its Gated-DeltaNet rows on "
- "the cached replay only, not on the legacy per-step recurrent states; leave "
- "TRTLLM_USE_GDN_REPLAY unset"
- )
- if str(engine.moe_backend).upper() == "CUTEDSL" and draft > _MAX_CUTEDSL_DRAFT_LEN:
- raise NoTarget(
- 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 _check_attention_dp(
- member: str, spec_config: Optional[MTPDecodingConfig], engine: EngineChoices
-) -> None:
- """Decline what the target does not run under attention DP (dep4, dep2), where
- each rank serves its own requests and only the routed experts cross ranks.
-
- MTP is admitted at dep4, where each rank drafts for its own requests and the draft layer's
- FP8 expert window runs the trunk's attention-DP round trip. Routing admits both MTP and
- the breakable prefill graph at the attention-TP layouts (tp4, tep4).
- """
- if spec_config is not None and member not in _ATTENTION_DP_MTP_LAYOUTS:
- raise NoTarget(
- f"speculative decoding with attention DP ({member}): no draft acceptance was "
- "measured at this layout; MTP under attention DP is admitted at "
- f"{', '.join(_ATTENTION_DP_MTP_LAYOUTS)}"
- )
- if engine.prefill_cuda_graph_backend == "breakable":
- raise NoTarget(
- f"prefill_cuda_graph_backend='breakable' with attention DP ({member}): the engine "
- "pads every rank to one prefill bucket, so a rank holding only generation requests "
- "is handed more rows than its batch, which the target's decode step refuses"
- )
-
-
# Qwen3.8-Flash-Next publishes ``Qwen4ExpForConditionalGeneration``, and the
# config loader's composite branch writes the same name back.
@routes("Qwen4ExpForConditionalGeneration")
def route(
model_config: ModelConfig, llm_args: TorchLlmArgs, compute_capability: Tuple[int, int]
) -> type:
- """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits.
-
- A target is imported where it is returned, so declining costs no import of
- modeling code. Routing also declines the engine choices (`EngineChoices`) the
- target's constructor would refuse.
- """
- pretrained_config = model_config.pretrained_config
- mapping = model_config.mapping
- spec_config = model_config.spec_config
- kv_cache_quant_algo = engine_kv_cache_quant_algo(model_config, llm_args)
- text = getattr(pretrained_config, "text_config", None)
- vision = getattr(pretrained_config, "vision_config", None)
- # Absent means this is not the normalized Qwen4-Exp composite at all, and
- # nothing below would be readable.
- if text is None or vision is None:
+ """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits."""
+ mapping, spec = model_config.mapping, model_config.spec_config
+ member = layout(mapping)
+ if compute_capability != _SM or member is None:
raise NoTarget(
- f"{type(pretrained_config).__name__} carries no text_config/vision_config; "
- "the qwen4_exp target reads both halves of the normalized composite"
+ f"compute capability {compute_capability}, world_size={mapping.world_size} "
+ f"(tp={mapping.tp_size}, pp={mapping.pp_size}, cp={mapping.cp_size}, "
+ f"attention_dp={mapping.enable_attention_dp}, moe_tp={mapping.moe_tp_size}, "
+ f"moe_ep={mapping.moe_ep_size}); the qwen4_exp target is Qwen3.8-Flash-Next-NVFP4 "
+ "on sm103 in one node: tp1, tp4 (experts split by intermediate), tep4, dep4 "
+ "(attention DP) or dep2 (attention DP, experts EP2)"
)
-
- shape = (
- text.num_hidden_layers,
- text.hidden_size,
- text.num_experts,
- text.num_attention_heads,
- text.linear_num_value_heads,
- text.vocab_size,
- )
- vshape = (
- vision.depth,
- vision.hidden_size,
- vision.num_heads,
- vision.spatial_merge_size,
- vision.out_hidden_size,
- vision.num_position_embeddings,
- tuple(getattr(vision, "deepstack_visual_indexes", None) or ()),
- vision.hidden_act,
- )
-
- # The target implements the one-model MTP-Eagle loop over the checkpoint's
- # single MTP module, and no other draft path. A draft longer than 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's constructor relies on each of these, so declining here
- # is what lets ``auto`` fall back to the built-in model.
- 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 <= _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
- )
-
- checkpoint = compute_capability == _SM and shape == _TEXT_SHAPE and vshape == _VISION_SHAPE
- if checkpoint:
- settings = _text_settings(text)
- differs = {k: v for k, v in settings.items() if v != _TEXT_SETTINGS[k]}
- if differs:
- raise NoTarget(
- f"the checkpoint has Qwen3.8-Flash-Next's shape but declares {differs}; the "
- f"target implements { ({k: _TEXT_SETTINGS[k] for k in differs}) }"
- )
- # Qwen3.8-Flash-Next-NVFP4 on GB300s (sm103): on one at tp1, on four at tp4, tep4 or dep4,
- # MTP-Eagle or no speculation; on two at dep2, without speculation.
- member = layout(mapping)
- if checkpoint and member is not None and mtp_eagle_or_none:
- engine = engine_choices(model_config, llm_args)
- if mapping.enable_attention_dp:
- _check_attention_dp(member, spec_config, engine)
- _check_engine(spec_config, kv_cache_quant_algo, engine, member)
- # One line on purpose: the claims suite reads the routed target off the
- # `from ..modeling import ` line.
- # fmt: off
- # isort: off
- from .qwen3_8_flash_next_nvfp4__sm_103__default.modeling import ModelingV2Qwen38FlashNextNvfp4Sm103Default
- # isort: on
- # fmt: on
-
- return ModelingV2Qwen38FlashNextNvfp4Sm103Default
-
- raise NoTarget(
- f"compute capability {compute_capability}, text shape {shape}, vision shape "
- f"{vshape}, {_describe(mapping)}, speculative decoding "
- f"{spec_config.spec_dec_mode.name if spec_config is not None else 'off'}"
- f"{f' max_draft_len={spec_config.max_draft_len}' if spec_config is not None else ''}, "
- f"KV cache {kv_cache_quant_algo or 'bf16'}; the qwen4_exp target is "
- "Qwen3.8-Flash-Next-NVFP4 on sm103 in one node: at 1 rank (tp1), and at 4 ranks tp4 "
- "(attention TP4, experts split by intermediate), tep4 (attention TP4, experts EP4) or "
- f"dep4 (attention DP, experts EP4), with one-model MTP-Eagle (max_draft_len "
- f"1..{_MAX_DRAFT_LEN}) or no speculative decoding; at 2 ranks dep2 (attention DP, "
- "experts EP2) without speculative decoding"
- )
+ if spec is not None and not (
+ isinstance(spec, MTPDecodingConfig) and spec.spec_dec_mode.is_mtp_eagle_one_model()
+ ):
+ raise NoTarget(
+ f"speculative decoding {spec.spec_dec_mode.name}; the target runs one-model "
+ "MTP-Eagle or no speculative decoding"
+ )
+ if spec is not None and member == "dep2":
+ raise NoTarget(
+ "speculative decoding with attention DP (dep2): no draft acceptance was measured at "
+ "this layout; MTP under attention DP is admitted at dep4"
+ )
+ # One line on purpose: the claims suite reads the routed target off the
+ # `from ..modeling import ` line.
+ # fmt: off
+ # isort: off
+ from .qwen3_8_flash_next_nvfp4__sm_103__default.modeling import ModelingV2Qwen38FlashNextNvfp4Sm103Default
+ # isort: on
+ # fmt: on
+
+ ModelingV2Qwen38FlashNextNvfp4Sm103Default.check_deployment(model_config, llm_args, member)
+ return ModelingV2Qwen38FlashNextNvfp4Sm103Default
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py
index caec32aee91d..f0b9e287d48a 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py
@@ -99,7 +99,7 @@ def _qwen4_exp_config(text_overrides=None, **vision_overrides):
num_attention_heads=24,
linear_num_value_heads=48,
vocab_size=248320,
- # The settings the target pins (`routing._TEXT_SETTINGS`), as the normalized config
+ # The settings the target pins (`modeling._TEXT_SETTINGS`), as the normalized config
# carries them: rope theta and partial factor flattened to the top level.
norm_topk_prob=True,
num_experts_per_tok=10,
@@ -158,16 +158,9 @@ def _declined(config, match=None, **llm_args_kwargs) -> None:
modeling_v2_resolve(config, _llm_args("require", **llm_args_kwargs))
-def test_routing_admits_exactly_what_the_target_implements():
- """`route` declines a MoE backend, a draft length or a recurrent-state dtype the
- target does not run; this keeps routing's constants and the target's together."""
- # At tp1 the CuteDSL backend and the fp32 state too; at the other layouts neither.
- assert set(routing._ONE_RANK_MOE_BACKENDS) == set(M._TRUNK_MOE_BACKENDS)
- assert set(routing._MOE_BACKENDS) == set(M._TRUNK_MOE_BACKENDS) - {"CUTEDSL"}
- assert routing._MAX_CUTEDSL_DRAFT_LEN == M._CUTEDSL_MAX_DRAFT_LEN
- assert routing._MAX_DRAFT_LEN == M._MAX_VERIFY_WIDTH - 1
- assert routing._SSM_STATE_DTYPES == M._SSM_STATE_DTYPES
- assert routing._ONE_RANK_SSM_STATE_DTYPES == M._ONE_RANK_SSM_STATE_DTYPES
+def test_routing_and_the_target_agree_on_the_gpu():
+ """`route` gates on the GPU before it imports the target, whose constructor asserts the
+ same one; the target's `check_deployment` holds every other admission fact."""
assert routing._SM == M._SM
@@ -467,11 +460,15 @@ def test_mtp_matches_with_the_flashinfer_ring_replay(monkeypatch, max_draft_len,
_claimed(_model_config(MTPDecodingConfig(max_draft_len=max_draft_len), **_LAYOUTS["tep4"]))
-def test_other_kv_cache_algorithms_decline():
- choices = routing.engine_choices(_model_config(**_LAYOUTS["tep4"]), _llm_args())
- for algo in (QuantAlgo.INT8, QuantAlgo.FP8_BLOCK_SCALES):
- with pytest.raises(routing.NoTarget, match="KV cache"):
- routing._check_engine(None, algo, choices, "tep4")
+@pytest.mark.parametrize("algo", [QuantAlgo.INT8, QuantAlgo.FP8_BLOCK_SCALES])
+def test_another_checkpoint_kv_cache_algorithm_declines(algo):
+ """`kv_cache_config.dtype: auto` keeps the checkpoint's own KV-cache algorithm."""
+ config = _model_config(**_LAYOUTS["tep4"])
+ config.quant_config.kv_cache_quant_algo = algo
+ _declined(config, match="KV cache")
+
+
+def test_the_fp8_ds_mla_kv_cache_declines():
_declined(
_model_config(**_LAYOUTS["tep4"]),
match="KV cache",
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 b7a579756d49..635f655dcf54 100644
--- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py
+++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py
@@ -38,9 +38,6 @@
modeling_v2_resolve,
resolved_mode,
)
-from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp import (
- routing as qwen4_exp_routing,
-)
from tensorrt_llm._torch.configs.kimi_k3 import KimiK3Config
from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce
from tensorrt_llm._torch.model_config import ModelConfig
@@ -181,7 +178,7 @@ def _qwen4_exp_config(**vision_overrides):
num_attention_heads=24,
linear_num_value_heads=48,
vocab_size=248320,
- # The settings the target pins (`routing._TEXT_SETTINGS`), as the normalized config
+ # The settings the target pins (`modeling._TEXT_SETTINGS`), as the normalized config
# carries them: rope theta and partial factor flattened to the top level.
norm_topk_prob=True,
num_experts_per_tok=10,
@@ -1279,11 +1276,8 @@ 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):
- 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, "tp1")
+ """`kv_cache_config.dtype: auto` keeps the checkpoint's own KV-cache algorithm."""
+ _declined(_model_config(_qwen4_exp_config(), kv_cache_quant_algo=algo), match="KV cache")
@pytest.mark.usefixtures("_qwen4_exp_sm")
From 598f1118e7767abd883c95fe9f2359278447804a Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Thu, 5 Nov 2026 21:12:38 -0800
Subject: [PATCH 12/15] [None][test] modeling_v2 qwen4_exp: name both readers
of the quantized-KV gsm8k AL keys
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
.../defs/accuracy/references/acceptance_length.yaml | 9 +++++----
1 file changed, 5 insertions(+), 4 deletions(-)
diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml
index 1cd64aada6a7..59c638a096e5 100644
--- a/tests/integration/defs/accuracy/references/acceptance_length.yaml
+++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml
@@ -233,13 +233,14 @@ ModelingV2Qwen38FlashNextNvfp4Sm103Default::tp1_mtp7_nvfp4_kv_bf16_state_long_co
ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3:
ref_al: 3.689
min_al: 3.50
-# Read by test_mtp3_gsm8k_quantized_kv[tp4_nvfp4_kv_breakable]: tp4 over NVFP4 KV with the
-# breakable prefill graph. Measured and floored as the mtp3 entry.
+# Read by test_mtp3_gsm8k_quantized_kv[tp4_nvfp4_kv_breakable] (tp4 over NVFP4 KV with the
+# breakable prefill graph, where it was measured) and [tep4_nvfp4_kv_no_reuse]. Floored as the
+# mtp3 entry.
ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3_nvfp4_kv:
ref_al: 3.535
min_al: 3.36
-# Read by test_mtp3_gsm8k_quantized_kv[dep4_fp8_kv]: dep4 (MTP under attention DP) over FP8 KV.
-# Measured and floored as the mtp3 entry.
+# Read by test_mtp3_gsm8k_quantized_kv[dep4_fp8_kv] (dep4, MTP under attention DP, over FP8 KV,
+# where it was measured) and [tep4_fp8_kv_flashinfer_ring]. Floored as the mtp3 entry.
ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3_fp8_kv:
ref_al: 3.512
min_al: 3.34
From bfbc840db2c37b2bd738e4aaaeb558efad2640c0 Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Fri, 6 Nov 2026 00:44:52 -0800
Subject: [PATCH 13/15] [None][doc] modeling_v2 qwen4_exp: drop the Modeling V2
section from the Qwen3.8-Flash-Next deployment guide
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
...-guide-for-qwen3.8-flash-next-on-trtllm.md | 45 -------------------
1 file changed, 45 deletions(-)
diff --git a/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md
index bff2472e1313..4cb49d003240 100644
--- a/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md
+++ b/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md
@@ -198,51 +198,6 @@ encoder_max_num_tokens: 65536
Image input is supported in aggregated serving with a local encoder. A separate encoder-to-prefill handoff is not part of this guide.
-### Modeling V2 (Experimental)
-
-On GB300 (sm_103), the NVFP4 checkpoint can also be served by the modeling_v2 implementation, which is off by default. Set `modeling_v2` in the YAML file:
-
-```yaml
-modeling_v2: auto # use modeling_v2 when it supports this deployment, otherwise the built-in model
-# modeling_v2: require # fail at startup, with the reason, when it does not
-```
-
-It supports these single-node layouts:
-
-| Layout | GPUs | Parallel settings | Speculative decoding |
-|---|---|---|---|
-| tp1 | 1 | `tensor_parallel_size: 1` | none, or MTP with `max_draft_len` 1-7 |
-| tp4 | 4 | `tensor_parallel_size: 4`, `moe_tensor_parallel_size: 4` | none, or MTP with `max_draft_len` 1-7 |
-| tep4 | 4 | `tensor_parallel_size: 4`, `moe_expert_parallel_size: 4` | none, or MTP with `max_draft_len` 1-7 |
-| dep4 | 4 | as tep4, plus `enable_attention_dp: true` | none, or MTP with `max_draft_len` 1-7 |
-| dep2 | 2 | `tensor_parallel_size: 2`, `moe_expert_parallel_size: 2`, `enable_attention_dp: true` | none |
-
-Under `auto`, these settings make the server fall back to the built-in model; under `require`, the server fails at startup:
-
-* Any other layout, including `enable_lm_head_tp_in_adp: true` with attention DP.
-* A checkpoint of the same shape whose config changes what it computes, for example `norm_topk_prob: false`, another `num_experts_per_tok`, other rope settings, or an MTP module that reads another layer's hidden state (`mtp_use_hidden_state_from_layer`).
-* Speculative decoding other than one-model MTP, or a draft length that changes at runtime.
-* MTP with `TRTLLM_USE_GDN_REPLAY=0`.
-* A KV cache dtype other than the checkpoint's bf16, `fp8` or `nvfp4`.
-* A recurrent-state cache other than `bfloat16`. tp1 also accepts `float32`.
-* `moe_config.backend` other than `AUTO` or `TRTLLM`. tp1 also accepts `CUTEDSL` with `max_draft_len` up to 3.
-* `prefill_cuda_graph_backend: piecewise`; with attention DP, also `breakable`.
-* `moe_config.load_balancer` (EPLB) at tep4, dep4 or dep2.
-
-Other differences from the built-in model:
-
-* The PLE n-gram table is in pinned host memory by default, one copy per rank (about 48 GiB each, so about 190 GiB of host memory for a four-GPU node). `TRTLLM_MODELING_V2_PLE_HOST_OFFLOAD=0` keeps it on the device. `TRTLLM_QWEN4_EXP_PLE_HOST_OFFLOAD` applies only to the built-in model.
-* `disable_mm_encoder` has no effect. The vision tower is always built and image requests are accepted.
-* At tp1 these environment switches enable lossy FP8 or NVFP4 copies of projections. All are off by default, and the other layouts ignore them with a warning:
- * `TRTLLM_MODELING_V2_DENSE_FP8`
- * `TRTLLM_MODELING_V2_GDN_INPROJ_NVFP4`
- * `TRTLLM_MODELING_V2_TARGET_HEAD_NVFP4` (with `DENSE_FP8`)
- * `TRTLLM_MODELING_V2_MTP_HEAD_NVFP4`
- * `TRTLLM_MODELING_V2_MTP_MOE_NVFP4`
- * `TRTLLM_MODELING_V2_MTP_DENSE_FP8`
-* At tp1, `TRTLLM_MODELING_V2_MOE_PREFILL_COPY=0` drops the second copy of the routed experts that serves large prefill steps.
-* A `TRTLLM_MODELING_V2_*` switch accepts `1`/`0`, `true`/`false`, `yes`/`no` or `on`/`off`, and the server fails at startup on any other value.
-
### Key Configuration Options
These options control TensorRT LLM behavior and are set in the YAML file passed to `trtllm-serve` with the `--config` argument.
From 2f9bf038fc9b14a2112b6fe1f7468bdf303c124f Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Fri, 6 Nov 2026 02:25:42 -0800
Subject: [PATCH 14/15] [None][fix] modeling_v2 qwen4_exp: put the default
target's host-resident PLE table on transparent huge pages, as #21463 does
for the tp1 target
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
.../weights.py | 24 +++++--
.../test_ple_host_table.py | 66 +++++++++++++++++++
2 files changed, 86 insertions(+), 4 deletions(-)
create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_ple_host_table.py
diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/weights.py
index 2a405f1ef4a5..f94e9a0abe81 100644
--- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/weights.py
+++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/weights.py
@@ -88,6 +88,7 @@
rank's windows, and the MTP module when there is no speculation.
"""
+import mmap
import os
from dataclasses import dataclass
from types import SimpleNamespace
@@ -100,6 +101,7 @@
from tensorrt_llm._torch.pyexecutor.config_utils import mamba_effective_tp_size
from tensorrt_llm._utils import CUASSERT
from tensorrt_llm.llmapi.utils import get_numa_aware_cpu_affinity
+from tensorrt_llm.logger import logger
from .pool import attention_tp_size
@@ -207,11 +209,16 @@ def _local_cpus(device) -> set | None:
class _RegisteredHostBuffer:
- """A pageable CPU allocation registered (pinned and mapped) for the device, and
- unregistered before its storage is released: at this size glibc serves the
- allocation with `mmap` and returns it with `munmap`, so freeing registered pages
+ """A pageable CPU mapping of its own, registered (pinned and mapped) for the
+ device, and unregistered before its pages are unmapped: freeing registered pages
would leave the driver mapping memory the OS may hand out again.
+ The mapping is advised `MADV_HUGEPAGE` before the registration faults it in, so
+ the table sits on transparent huge pages whatever torch's CPU allocator does
+ (TensorRT LLM imports torch with mimalloc's THP advice off). On GB300 the
+ device's random row gathers from 64 KiB pages run 2.5x (64 tokens) to 23x
+ (1024 tokens) slower.
+
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."""
@@ -227,7 +234,16 @@ def __init__(self, nbytes: int, device=None) -> None:
if local:
os.sched_setaffinity(0, local)
try:
- self._storage = torch.empty(nbytes, dtype=torch.uint8)
+ self._mapping = mmap.mmap(-1, nbytes, flags=mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS)
+ try:
+ self._mapping.madvise(mmap.MADV_HUGEPAGE)
+ except OSError as error: # a kernel without transparent huge pages
+ logger.warning(
+ f"madvise(MADV_HUGEPAGE) failed for the host-resident PLE table ({error}); "
+ "the device's row gathers from its base pages run several times slower"
+ )
+ # Shares the mapping's pages, and keeps the mapping alive while it lives.
+ self._storage = torch.frombuffer(self._mapping, dtype=torch.uint8)
self.pointer = self._storage.data_ptr()
CUASSERT(
cudart.cudaHostRegister(
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_ple_host_table.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_ple_host_table.py
new file mode 100644
index 000000000000..6c18e8651c44
--- /dev/null
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_ple_host_table.py
@@ -0,0 +1,66 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""The host-resident PLE n-gram table (`weights._host_resident_empty`): a mapping of its own,
+advised `MADV_HUGEPAGE` before the registration faults it in, and still allocated and
+registered, with a warning, on a kernel that refuses the advice."""
+
+import errno
+import mmap
+
+import pytest
+import torch
+
+from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.weights import ( # noqa: E501
+ _host_resident_empty,
+)
+from tensorrt_llm.logger import logger
+
+_ROWS = 1 << 16
+_HEAD_DIM = 256
+
+
+def _huge_page_advice_accepted() -> bool:
+ """Whether this kernel takes MADV_HUGEPAGE (it has transparent huge pages)."""
+ probe = mmap.mmap(-1, mmap.PAGESIZE)
+ try:
+ probe.madvise(mmap.MADV_HUGEPAGE)
+ except OSError:
+ return False
+ finally:
+ probe.close()
+ return True
+
+
+@pytest.mark.skipif(not _huge_page_advice_accepted(), reason="the kernel refuses MADV_HUGEPAGE")
+def test_host_table_is_advised_transparent_huge_pages() -> None:
+ # The host-resident table is a mapping of its own, advised MADV_HUGEPAGE
+ # before it is faulted in, whatever torch's CPU allocator does with huge pages.
+ table, host = _host_resident_empty(_ROWS, _HEAD_DIM, torch.float8_e4m3fn, torch.device("cuda"))
+ flags, inside = None, False
+ with open("/proc/self/smaps") as smaps:
+ for line in smaps:
+ fields = line.split()
+ if line[:1] in "0123456789abcdef" and "-" in fields[0]:
+ lo, hi = (int(x, 16) for x in fields[0].split("-"))
+ inside = lo <= host.pointer < hi
+ elif inside and line.startswith("VmFlags:"):
+ flags = fields[1:]
+ break
+ assert flags is not None and "hg" in flags, flags
+ assert table.shape == (_ROWS, _HEAD_DIM)
+
+
+def test_host_table_loads_when_huge_page_advice_is_refused(monkeypatch) -> None:
+ # A kernel without transparent huge pages refuses the advice: the table is
+ # still allocated and registered, on base pages, with a warning.
+ class _Refused(mmap.mmap):
+ def madvise(self, *args):
+ raise OSError(errno.EINVAL, "Invalid argument")
+
+ warnings = []
+ monkeypatch.setattr(mmap, "mmap", _Refused)
+ monkeypatch.setattr(logger, "warning", warnings.append)
+ table, host = _host_resident_empty(_ROWS, _HEAD_DIM, torch.float8_e4m3fn, torch.device("cuda"))
+ assert isinstance(host._mapping, _Refused) and host.pointer
+ assert table.shape == (_ROWS, _HEAD_DIM)
+ assert len(warnings) == 1 and "MADV_HUGEPAGE" in warnings[0], warnings
From 4c04aecd2059f6f227b0e98cb8f09f18be633485 Mon Sep 17 00:00:00 2001
From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Date: Fri, 6 Nov 2026 02:46:29 -0800
Subject: [PATCH 15/15] [None][test] modeling_v2 qwen4_exp: count only the
MADV_HUGEPAGE warning; name every post-merge quantized-KV case
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
---
.../test_lists/test-db/l0_gb300_multi_gpus.yml | 9 +++++----
.../test_ple_host_table.py | 3 ++-
2 files changed, 7 insertions(+), 5 deletions(-)
diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
index ad22d5a2e7fc..f6bf56fa630b 100644
--- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
+++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
@@ -33,10 +33,11 @@ l0_gb300_multi_gpus:
- accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_gsm8k_chunked_prefill[dep4]
- accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp4::test_gsm8k
- accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp4::test_eagle3_gsm8k_and_acceptance
- # qwen4_exp target, post-merge: three drafts over NVFP4 KV at tp4 with the
- # breakable prefill graph and over FP8 KV at dep4 (MTP under attention DP), both with block
- # reuse and chunked prefill; and dep2 context and generation engines on one node. tep4 runs
- # pre-merge (l0_gb300); mmmu runs in QA (llm_function_core).
+ # qwen4_exp target, post-merge: three drafts over quantized KV -- NVFP4 KV at tp4 with the
+ # breakable prefill graph, FP8 KV at dep4 (MTP under attention DP), FP8 KV at tep4 on
+ # FlashInfer's ring replay, NVFP4 KV at tep4 without block reuse -- and dep2 context and
+ # generation engines on one node. tep4 over bf16 KV runs pre-merge (l0_gb300); mmmu runs in
+ # QA (llm_function_core).
- accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[tp4_nvfp4_kv_breakable] TIMEOUT (60)
- accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[dep4_fp8_kv] TIMEOUT (60)
- accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[tep4_fp8_kv_flashinfer_ring] TIMEOUT (60)
diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_ple_host_table.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_ple_host_table.py
index 6c18e8651c44..e18caefb9549 100644
--- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_ple_host_table.py
+++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_ple_host_table.py
@@ -63,4 +63,5 @@ def madvise(self, *args):
table, host = _host_resident_empty(_ROWS, _HEAD_DIM, torch.float8_e4m3fn, torch.device("cuda"))
assert isinstance(host._mapping, _Refused) and host.pointer
assert table.shape == (_ROWS, _HEAD_DIM)
- assert len(warnings) == 1 and "MADV_HUGEPAGE" in warnings[0], warnings
+ # Other modules on the path (the NUMA placement) may warn too; count only this one.
+ assert len([w for w in warnings if "MADV_HUGEPAGE" in str(w)]) == 1, warnings