diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py index 2eeadb915157..b2ae46ce1798 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py @@ -1591,6 +1591,7 @@ def _prepare_next_new_tokens( batch_indices_cuda, batch_size, num_accepted_tokens, + fused: bool = False, ): """The rows the last draft-step kernel wrote when these draft tokens are that loop's.""" drafted, self._drafted_rows = self._drafted_rows, None @@ -1602,6 +1603,7 @@ def _prepare_next_new_tokens( batch_indices_cuda, batch_size, num_accepted_tokens, + fused=fused, ) def _sample_and_accept_draft_tokens_base( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py index 29ac78174102..b0e6656af651 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py @@ -80,8 +80,8 @@ is also what ``auto`` resolves to for this checkpoint (`_SSM_STATE_DTYPES`): a lossy substitution, since the state is rounded to bf16 once per decode token, once per prefill call and at every cached-replay checkpoint commit. The conv, PLE and replay-history buffers keep their own -types under both. A bf16 pool makes the decode step call its recurrence without ``cu_seqlens``; -see ``DecodeTarget._gdn_block``. +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 @@ -105,8 +105,11 @@ forward, which cannot express a host readback. """ +import functools import math import os +import re +import types from typing import NamedTuple import torch @@ -129,6 +132,9 @@ 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.interface import AttentionMetadata from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.model_config import ModelConfig @@ -149,7 +155,9 @@ 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_scores_verify import qsa_paged_index_scores_verify 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 @@ -158,29 +166,69 @@ 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.embedding.ngram_hash import ngram_hash +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 amax_partials, 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_gated_residual import qwen4_exp_gated_residual from .ops.hc.qwen4_exp_hc_decode_first_cute import qwen4_exp_hc_decode_first_cute from .ops.hc.qwen4_exp_hc_decode_mix_fp4_cute import ( qwen4_exp_hc_decode_mix_cute, qwen4_exp_hc_decode_mix_fp4_cute, ) +from .ops.hc.qwen4_exp_hc_decode_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, @@ -193,32 +241,27 @@ 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.nvfp4_moe_gemv import nvfp4_moe_gemv_op +from .ops.moe.nvfp4_moe_gemv import precompile as nvfp4_moe_gemv_precompile +from .ops.moe.nvfp4_moe_gemv import register_weights as nvfp4_moe_gemv_register from .ops.moe.renorm_moe_routing_op import renorm_moe_routing_op from .ops.norm.flashinfer_gemma_rmsnorm import flashinfer_gemma_rmsnorm from .ops.norm.layer_norm import layer_norm from .ops.norm.rms_norm_gated_token_major import rms_norm_gated_token_major -from .ops.ple.gated_ngram_injection import gated_ngram_injection -from .ops.ple.ngram_token_history import ngram_token_history from .ops.ple.ple_gate_conv_inject_norm import ple_gate_conv_inject_norm from .ops.ple.ple_ngram_embed import ple_ngram_embed from .ops.quantization.bf16_to_fp32 import bf16_to_fp32 -from .ops.quantization.dequantize_e4m3_per_tensor import dequantize_e4m3_per_tensor from .ops.quantization.fp8_quantize_1x128 import fp8_quantize_1x128 -from .ops.ssm.causal_conv1d_fwd import causal_conv1d_fwd from .ops.ssm.causal_conv1d_triton_update import causal_conv1d_triton_update from .ops.ssm.causal_conv1d_update import causal_conv1d_update -from .ops.ssm.chunk_gated_delta_rule import chunk_gated_delta_rule from .ops.ssm.cu_seqlens_from_seq_lens import cu_seqlens_from_seq_lens -from .ops.ssm.dilated_short_conv_paged import dilated_short_conv_paged -from .ops.ssm.extract_transpose_prefill_slice import extract_transpose_prefill_slice from .ops.ssm.flashinfer_chunk_gated_delta_rule import flashinfer_chunk_gated_delta_rule from .ops.ssm.fused_gdn_gating import fused_gdn_gating -from .ops.ssm.fused_gdn_post_conv import fused_gdn_post_conv from .ops.ssm.fused_recurrent_gated_delta_rule_cached_replay_update import ( fused_recurrent_gated_delta_rule_cached_replay_update, ) from .ops.ssm.fused_recurrent_gated_delta_rule_update import fused_recurrent_gated_delta_rule_update -from .ops.ssm.fused_sigmoid_gating_delta_rule_update import fused_sigmoid_gating_delta_rule_update +from .ops.ssm.gdn_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 @@ -226,7 +269,6 @@ from .ops.ssm.gdn_in_proj_conv1d_pdl import gdn_in_proj_conv1d_pdl from .ops.ssm.gdn_prefill_conv_post import gdn_prefill_conv_post from .ops.ssm.pack_gdn_decode_qkv import pack_gdn_decode_qkv -from .ops.ssm.reset_recurrent_state_rows import reset_recurrent_state_rows from .ops.vision.patch_embed_conv3d import patch_embed_conv3d from .ops.vision.pos_embed_interpolate import pos_embed_interpolate @@ -259,6 +301,10 @@ # 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 #: 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 @@ -267,18 +313,26 @@ #: 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. More rows run `low_m_bf16_mm` with its FlashInfer - branches' PDL on (its cuBLAS branch, above M=32 and off its table, ignores the flag). - Every kernel either path launches reads the activation only after - `griddepcontrol.wait`. + 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`. """ - if mat_a.shape[0] == 1 and mat_b.shape[0] in _GEMV_TILED_K: + 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) @@ -477,9 +531,8 @@ def _decode_projection(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor # process. _PAD_SLOT_ID = -1 -# Gated-DeltaNet gating form: the only pair `fused_sigmoid_gating_delta_rule_update` -# is certified at; matches `gdn_prefill_conv_post`'s default so prefill and decode -# compute the same `g`. +# 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 @@ -554,6 +607,228 @@ def _moe_prefill_copy(trunk_moe_backend: str, max_num_tokens: int) -> bool: #: draft rejected; only the acceptance gate can tell. _MTP_EMBED_BLOCK_ALL_STREAMS = True +#: Set to 1 to run the decode draft step on the composed (unfused) draft layer instead of the +#: trunk's fused site, attention and projection launches (`DecodeTarget.mtp_forward`). +_MTP_UNFUSED_ENV = "TRTLLM_MODELING_V2_MTP_UNFUSED" + +#: The draft head's lm_head: "fp8" (default: a per-vocabulary-row e4m3 copy read by cuBLASLt with +#: per-token activation scales) or "bf16" (the checkpoint's, as the verify step reads it). The +#: draft head 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 AL). +_MTP_DRAFT_HEAD_ENV = "TRTLLM_MODELING_V2_MTP_DRAFT_HEAD" + + +def _mtp_draft_head_mode() -> str: + mode = os.environ.get(_MTP_DRAFT_HEAD_ENV, "fp8").lower() + assert mode in ("bf16", "fp8"), f"{_MTP_DRAFT_HEAD_ENV}={mode!r}: expected bf16 or fp8" + return mode + + +#: Set to N > 0 to let the FP8 draft head propose only the first N token ids (the tokenizer's +#: most frequent merges come first) and the vocabulary's tail from the added tokens on +#: (end-of-text, the chat and thinking markers), rounded to whole 128-row blocks: a row +#: subset of the head's FP8 copies, its draft ids mapped back through the worker's `d2t` +#: offsets. A draft can only miss more often -- the verify step's full head decides what is +#: emitted -- so outputs are unchanged and the acceptance length is the cost to watch. +#: 0 (default): the whole vocabulary. +_MTP_DRAFT_VOCAB_ENV = "TRTLLM_MODELING_V2_MTP_DRAFT_VOCAB" + + +def _mtp_draft_vocab() -> int: + vocab = int(os.environ.get(_MTP_DRAFT_VOCAB_ENV, "0")) + assert vocab >= 0, f"{_MTP_DRAFT_VOCAB_ENV}={vocab}: expected a non-negative row count" + return vocab + + +#: 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, + amax_parts=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). `amax_parts` are the partial maxima `x`'s producer + folded (`_fp8_amax_parts`) for the FP8 GEMM.""" + 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, amax_parts) + 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) + + +def _fp8_amax_parts(rows: int, fp8, kind: str, device) -> torch.Tensor | None: + """For the producer of a prefill `_dense_mm` input (no packed decode copy) of `rows` rows: + a zeroed partial-maximum buffer when the call takes the per-tensor FP8 GEMM, so the + producer folds `max(|x|)` as it stores `x` and the GEMM skips its maximum pass; else None.""" + if fp8 is None or rows < _DENSE_FP8_MIN_ROWS[kind]: + return None + return amax_partials(device) + + #: 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` @@ -565,11 +840,22 @@ def _moe_prefill_copy(trunk_moe_backend: str, max_num_tokens: int) -> bool: #: `_HC_DECODE_GEMV_MAX_ROWS + 1` rows) runs its first half (projection and #: activations) as the one tcgen05 launch `qwen4_exp_hc_decode_tc_first` instead of #: cuBLAS split-K (+ its reduce) + `qwen4_exp_hc_lowrank_act`: the entry's MAX_ROWS -#: (one 128-row MMA tile); larger batches keep cuBLAS. +#: (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 = ( - 128 # capacity: fixed: dispatch threshold (one MMA tile), not a buffer bound + 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. @@ -579,6 +865,47 @@ def _moe_prefill_copy(trunk_moe_backend: str, max_num_tokens: int) -> bool: #: (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" + + +#: Set to 1 to run the decode step's routed experts on `nvfp4_moe_gemv` at the row counts in +#: `_MOE_GEMV_ROWS`, where the runner's finalize folds into the next hyper-connection site: +#: Renormalize routing, FC1, SwiGLU, NVFP4 requantization and FC2 in two PDL launches on the +#: tcgen05 block-scaled MMA over the trtllm-gen expert stacks (bitwise the runner's outputs on +#: the checkpoint). The TMA descriptors of every layer's stacks are built and the kernels +#: compiled at load. 0 (default): the trtllm-gen runner. +_MOE_GEMV_ENV = "TRTLLM_MODELING_V2_MOE_GEMV" +#: Decode row counts `nvfp4_moe_gemv` takes: where it measured faster than the runner (GB300, +#: 48-layer chain with the router projection). +_MOE_GEMV_ROWS = frozenset({1, 2, 3, 4}) + + +def _moe_gemv_enabled() -> bool: + mode = os.environ.get(_MOE_GEMV_ENV, "0") + assert mode in ("0", "1"), f"{_MOE_GEMV_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. @@ -669,6 +996,7 @@ class _Replay(NamedTuple): 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, committing wrapping rows in-kernel class _Bound(NamedTuple): @@ -996,6 +1324,8 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig # 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 " @@ -1280,6 +1610,9 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig self.moe_prefill_copy = _moe_prefill_copy( self.trunk_moe_backend, model_config.max_num_tokens ) + # `_MOE_GEMV_ENV`: the small-row decode MoE over the trtllm-gen stacks. + self.moe_gemv = _moe_gemv_enabled() and self.trunk_moe_backend == "TRTLLM" + self.max_num_tokens = model_config.max_num_tokens # ---- PLE ----------------------------------------------------------- ple_ids = list(getattr(cfg, "ple_layer_ids", None) or []) @@ -1509,6 +1842,8 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig 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. @@ -1517,9 +1852,22 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig ) self._prefill: "PrefillTarget | None" = None self._decode: "DecodeTarget | None" = None + self._mtp_fused: "_MtpFused | None" = None + # (e4m3 [vocab, hidden], fp32 [1, vocab] row scales) when the draft head reads FP8. + 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 when it reads FP8, and the target head's under the dense FP8 mode. + self._lm_head_g = None + # The draft head's own packed copy when it reads a vocabulary slice + # (`_MTP_DRAFT_VOCAB_ENV`); else the draft head shares `_lm_head_g`. + self._mtp_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. @@ -1570,6 +1918,11 @@ def fuse(*params): 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. @@ -1589,6 +1942,11 @@ def post_load_weights(self) -> None: ) self._moe_scalars = moe_scalars self._moe_scalars_cd = moe_scalars_cd + if self.moe_gemv and device.type == "cuda": + # Before any CUDA-graph capture: the descriptors are uploaded, the kernels built. + for i in range(self.num_layers): + nvfp4_moe_gemv_register(w[f"l{i}_fc1_w"], w[f"l{i}_fc2_w"]) + nvfp4_moe_gemv_precompile(max(_MOE_GEMV_ROWS)) # Both delta-rule entries index A_log and dt_bias as `ptr + head` and certify # fp32; the checkpoint is bf16. self._gdn_gating = { @@ -1633,11 +1991,67 @@ def post_load_weights(self) -> None: 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) @@ -1854,6 +2268,8 @@ def _bind_cache_manager(self, md, position_ids) -> None: 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 @@ -1915,6 +2331,18 @@ def _bind_replay(self, cm, g: _PoolGeometry, mm, ssm0) -> _Replay: # 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, which commits every wrapping row in-kernel at + # any batch size (the checkpoint read once and written once): the manager's all-layer + # commit would advance those checkpoints a second time, so it is switched off for this + # manager, and with it the work items the metadata builds for it. + cute = ssm0.dtype == torch.bfloat16 + if cute and all_layer: + cm.disable_gdn_cached_replay_all_layer_commit() + assert not cm.use_gdn_cached_replay_all_layer_commit, ( + f"{type(cm).__name__} kept its all-layer replay commit on: the CuTe replay " + f"would commit every wrapping checkpoint twice" + ) + all_layer = False if all_layer: missing = [n for n in ("replay_work_items", "replay_n_writes") if not hasattr(mm, n)] assert not missing, ( @@ -1922,7 +2350,7 @@ def _bind_replay(self, cm, g: _PoolGeometry, mm, ssm0) -> _Replay: f"{missing}; the partitioned verify call needs both" ) return _Replay( - old, meta.cache_buf_idx, meta.prev_num_accepted_tokens, history, width, all_layer + old, meta.cache_buf_idx, meta.prev_num_accepted_tokens, history, width, all_layer, cute ) def _bind_draft_cache_manager(self, md) -> None: @@ -2121,6 +2549,33 @@ class _LayerWeights(NamedTuple): # `[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]: @@ -2218,6 +2673,24 @@ def _layer_views(core: Qwen4ExpTp1ModelingV2) -> list[_LayerWeights]: 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 @@ -2232,9 +2705,6 @@ class _StackWeights(NamedTuple): fhc_up: torch.Tensor ple_key: torch.Tensor ple_val: torch.Tensor - ple_nk: tuple - ple_nq: tuple - ple_nc: tuple ple_nk_w: torch.Tensor ple_nq_w: torch.Tensor ple_nc_w: torch.Tensor @@ -2243,6 +2713,15 @@ class _StackWeights(NamedTuple): 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: @@ -2253,11 +2732,8 @@ def _stack_views(core: Qwen4ExpTp1ModelingV2) -> _StackWeights: fhc_up=w["fhc_up"].t(), ple_key=w["ple_key_w"].t(), ple_val=w["ple_val_w"].t(), - ple_nk=core._norm_slices("ple_nk"), - ple_nq=core._norm_slices("ple_nq"), - ple_nc=core._norm_slices("ple_nc"), - # The same three weights whole, for the decode target's fused tail, which - # normalizes each stream inside one program. + # 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"], @@ -2305,6 +2781,13 @@ class _MtpWeights(NamedTuple): 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: @@ -2346,6 +2829,124 @@ def _mtp_views(core: Qwen4ExpTp1ModelingV2) -> _MtpWeights: ) +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).""" + + 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. @@ -2408,10 +3009,16 @@ def _sparse_gqa( 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.""" + `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, @@ -2424,6 +3031,43 @@ def _sparse_gqa( 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, @@ -2515,6 +3159,123 @@ def _cute_dsl_routed(core, L, logits, xq, xsf, routed_out) -> torch.Tensor: ) +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 at scratch pages + `offsets[r]` on. Returns `(slots, counts, offsets, table, k_scratch, v_scratch)`, the + operands `qsa_nvfp4_pages_to_bf16` fills per layer: `table` (all -1 here) gets the + requests' slot rows, the scratch pools `[listed, Hkv, tpb, D]` bf16 the pages.""" + 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 + # size the scratch, the device's place the pages. + listed = int(((md.kv_lens_runtime[:num_contexts] + (tpb - 1)) // tpb).sum()) + 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" + ) + offsets = torch.cumsum(counts, 0) - counts + 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 slots, counts, offsets, 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; 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: + slots, counts, offsets, table, k_scratch, v_scratch = plan + 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[:num_ctx], + k_cache=k_scratch, + v_cache=v_scratch, + block_table=table, + selected_tokens=selected[:num_ctx], + request_indices=req_idx[:num_ctx], + tokens_per_block=tpb, + softmax_scale=core.softmax_scale, + query_positions=logical[:num_ctx], + compress_ratio=core.compress_ratio, + ) + out[:num_ctx].copy_(torch.reshape(context, [num_ctx, 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 @@ -2523,22 +3284,37 @@ def _live_score_columns(core, md, batch: int) -> int: return (longest + _MAX_VERIFY_WIDTH) // core.compress_ratio + 1 -def _replay_verify(core, mq_out, a_d, b_d, ssm_states, i, L, si_d, n, p, replay) -> torch.Tensor: +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). - `mq_out` is the conv's fresh, contiguous output, so the kernel reads q/k/v from it - directly (`packed_qkv`) and never from the views. The conv before it does not signal - dependents, so the kernel launches without PDL.""" + `replay` is `(work_items, n_writes, partitioned)`. For the Triton kernel (an fp32 pool) + 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). A bf16 pool runs the CuTe kernel (`gdn_cached_replay_cute`), + which commits every wrap here: the manager's all-layer commit is off for it + (`_bind_replay`), and the work items its metadata still builds over every decode row + hand the kernel each row's slot, pending count and history half in one load. `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) - out = fused_recurrent_gated_delta_rule_cached_replay_update( + entry = ( + functools.partial(gdn_cached_replay_cute, in_kernel_commit=True) + 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]), @@ -2557,7 +3333,7 @@ def _replay_verify(core, mq_out, a_d, b_d, ssm_states, i, L, si_d, n, p, replay) use_qk_l2norm_in_kernel=True, A_log=L.a_log, dt_bias=L.dt_bias, - launch_with_pdl=False, + launch_with_pdl=bool(pdl and state.cute), replay_work_items=work_items, n_writes=n_writes, block_v=None, @@ -2576,7 +3352,15 @@ def _replay_verify(core, mq_out, a_d, b_d, ssm_states, i, L, si_d, n, p, replay) 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.""" + all-layer commit after the step. The CuTe kernel (a bf16 pool, never partitioned) + takes the work items the metadata builds over all `num_decodes` rows whenever they + are there, with no all-layer commit.""" + state = core._bound.replay + if state.cute: + built = getattr(mm, "replay_num_decodes", None) == num_decodes + if built and getattr(mm, "replay_work_items", None) is not None: + return mm.replay_work_items[:num_decodes], mm.replay_n_writes, False + return None, None, False partitioned = ( core._bound.replay.all_layer_commit and num_decodes >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE @@ -2630,8 +3414,9 @@ class _PrefillStep(NamedTuple): `n` ids. `attn` is thop_attention's per-step arguments; `gdn`, `ple` and `qsa` are the recurrent, PLE and QSA bookkeeping tuples their breaks unpack, each cut to the batch's rows; `stage` says the speculative rollback stages this verify step; `vision` is the - shell's `(rows, mm_token_indices)` or None. A new one replaces it every step, so no - break reads a previous step's.""" + 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 @@ -2642,6 +3427,7 @@ class _PrefillStep(NamedTuple): qsa: tuple stage: bool vision: tuple | None + nvfp4_context: tuple | None = None class PrefillTarget: @@ -2660,8 +3446,15 @@ def __init__(self, core: Qwen4ExpTp1ModelingV2) -> None: if self._contract_pending and core.quant_mode == QM_FP8_KV_CACHE else None ) - self._layers = _layer_views(core) - self._stack = _stack_views(core) + # The 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 @@ -2757,6 +3550,30 @@ def _hc_moe_mixer(self, normed, first, up, gscale): 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 @@ -2817,6 +3634,7 @@ def _moe_block(self, x, xq, xsf, L: _LayerWeights, routed_out) -> tuple: L.global_sf_cd, core.num_experts, core.topk, + memset_stream=core._moe_memset_stream, ) else: routed = fp4_block_scale_moe_renorm_routed( @@ -2852,7 +3670,7 @@ def _attention_block(self, x, i, L, proj_qkv, index_scratch, attn_out, pos3): core = self.core num_tokens = x.shape[0] - proj = cublas_mm(x, L.attn_in) + 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. @@ -2896,8 +3714,9 @@ def _attention_block(self, x, i, L, proj_qkv, index_scratch, attn_out, pos3): if not self._kv_probe.observe(i, k_dst, v_dst): self._kv_probe = None eager_on_graph(proj, proj_qkv, index_scratch, attn_out, i, inner=self._attention_core) - fused_sigmoid_mul_inplace(attn_out, gate) - return cublas_mm(attn_out, L.o) + parts = _fp8_amax_parts(attn_out.shape[0], L.o8, "out", attn_out.device) + fused_sigmoid_mul_inplace(attn_out, gate, parts) + return _dense_mm(attn_out, L.o, L.o8, cublas_mm, "out", amax_parts=parts) 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 @@ -3046,11 +3865,26 @@ def _attention_core(self, proj, proj_qkv, index_scratch, attn_out, i) -> None: token_topk=core.token_topk, ) ) - # A context launch of thousands of rows over a bf16 pool takes the row-tile - # pass: consecutive rows of one request share their K/V tiles; rows of groups - # that span requests, the generation rows of a mixed step among them, are served - # row by row. Below the threshold the per-row kernel is faster. - if num_tokens >= _ROWTILE_MIN_ROWS and kv.k.dtype == q_dst.dtype: + 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, @@ -3063,6 +3897,7 @@ def _attention_core(self, proj, proj_qkv, index_scratch, attn_out, i) -> None: 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, @@ -3077,7 +3912,7 @@ def _attention_core(self, proj, proj_qkv, index_scratch, attn_out, i) -> None: query_positions=logical, compress_ratio=core.compress_ratio, ) - attn_out[:num_tokens].copy_(torch.reshape(out, [num_tokens, core.q_width])) + 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] @@ -3192,7 +4027,7 @@ def _gdn_block(self, x, i, L, core_out) -> torch.Tensor: gated-norm epilogue run over all `T` rows here, the recurrence in the `_gdn_core` break, which leaves its output in `core_out` `[1, T, num_v_heads, head_v_dim]`.""" core = self.core - proj = cublas_mm(x, L.in_proj) + 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]) @@ -3202,10 +4037,13 @@ def _gdn_core(self, proj, out, i) -> None: `[1, T, num_v_heads, head_v_dim]` in place and zeroes rows `[n:]`. Every branch is on a host int of the step, never a tensor value. - Over an fp32 recurrent pool outside a verify step the context rows take one fused + 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 FlashInfer's single-token decode kernel; - otherwise `_gdn_core_generic`.""" + 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 @@ -3228,16 +4066,13 @@ def _gdn_core(self, proj, out, i) -> None: replay, ) = step.gdn num_tokens = step.num_tokens - ssm_states = core._bound.ssm[i] - if is_verify or ssm_states.dtype != torch.float32: - self._gdn_core_generic(proj, out, i) + 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] - # A context sequence carrying nothing starts from a zeroed slot; only the - # context prefix is a candidate (a generation row's state must survive). - reset_recurrent_state_rows(ssm_states, si_p, hi_p, conv_states) if num_decode_tokens > 0: mq_p, mq_d = torch.split(mixed_qkv, [num_prefill_tokens, num_decode_tokens], 0) a_p, b_p = a[:num_prefill_tokens], b[:num_prefill_tokens] @@ -3271,12 +4106,14 @@ def _gdn_core(self, proj, out, i) -> None: 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' reset and recurrence touch only theirs). + # (the context rows' recurrence touches only theirs). gdn_decode_rows( proj[num_prefill_tokens:num_tokens], ssm_states, @@ -3308,13 +4145,17 @@ def _gdn_core(self, proj, out, i) -> None: 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_generic(self, proj, out, i) -> None: - """`_gdn_core` on a verify step or over a bf16 recurrent pool: the conv, the - post-conv gating and the chunked delta rule over the step's `n` rows, or under - speculation the verify path for the generation rows. The q|k|v columns are read in + 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.""" @@ -3343,124 +4184,17 @@ def _gdn_core_generic(self, proj, out, i) -> None: mixed_qkv, _, a, b = torch.split(proj[:num_tokens], core.in_proj_split, 1) a, b = a.contiguous(), b.contiguous() - conv_states = core._bound.conv[i] ssm_states = core._bound.ssm[i] - # A context sequence carrying nothing starts from a zeroed slot; only the - # context prefix is a candidate (a generation row's state must survive). - reset_recurrent_state_rows(ssm_states, si_p, hi_p, conv_states) - if num_decode_tokens > 0: - mq_p, mq_d = torch.split(mixed_qkv, [num_prefill_tokens, num_decode_tokens], 0) - # Prefill conv wants channel-major, decode token-major. - mq_p_t = extract_transpose_prefill_slice(mq_p, num_prefill_tokens, 0, core.conv_dim) - causal_conv1d_fwd( - x=mq_p_t, - weight=L.conv_w, - bias=None, - conv_states=conv_states, - query_start_loc=qsl_p, - cache_indices=si_p, - has_initial_state=hi_p, - silu_activation=True, - pad_slot_id=_PAD_SLOT_ID, - out=None, - ) - if is_verify: - # A context row advances its slot as on the base path; a generation - # row is a draft chain of unknown accepted length, so it goes to the - # rollback buffers. - a_p, a_d = torch.split(a, [num_prefill_tokens, num_decode_tokens], 0) - b_p, b_d = torch.split(b, [num_prefill_tokens, num_decode_tokens], 0) - q, k, v, g, beta = fused_gdn_post_conv( - prefill=mq_p_t, - decode=None, - a=a_p, - b=b_p, - A_log=L.a_log, - dt_bias=L.dt_bias, - num_k_heads=core.num_k_heads, - head_k_dim=core.head_k_dim, - num_v_heads=core.num_v_heads, - head_v_dim=core.head_v_dim, - l2_norm_eps=_L2_NORM_EPS, - softplus_threshold=_SOFTPLUS_THRESHOLD, - ) - # q and k are already L2-normalized by the post-conv call. - chunk_gated_delta_rule( - q, - k, - v, - g, - beta, - scale=None, - initial_state=ssm_states, - initial_state_indices=si_p, - inplace_indexed_state_update=True, - output_final_state=False, - cu_seqlens=qsl_long_p, - use_qk_l2norm_in_kernel=False, - output=out[:, :num_prefill_tokens], - ) - out_d = self._gdn_verify( - mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, num_decodes, gen_p, replay - ) - out[:, num_prefill_tokens:num_tokens].copy_(out_d) - out[:, num_tokens:].zero_() - return - causal_conv1d_update( - x=mq_d, - conv_state=conv_states, - weight=L.conv_w, - bias=None, - silu_activation=True, - cache_seqlens=None, - conv_state_indices=si_d, - pad_slot_id=_PAD_SLOT_ID, - ) - q, k, v, g, beta = fused_gdn_post_conv( - prefill=mq_p_t, - decode=mq_d, - a=a, - b=b, - A_log=L.a_log, - dt_bias=L.dt_bias, - num_k_heads=core.num_k_heads, - head_k_dim=core.head_k_dim, - num_v_heads=core.num_v_heads, - head_v_dim=core.head_v_dim, - l2_norm_eps=_L2_NORM_EPS, - softplus_threshold=_SOFTPLUS_THRESHOLD, - ) - else: - mq_t = extract_transpose_prefill_slice(mixed_qkv, num_tokens, 0, core.conv_dim) - causal_conv1d_fwd( - x=mq_t, - weight=L.conv_w, - bias=None, - conv_states=conv_states, - query_start_loc=qsl_all, - cache_indices=si_all, - has_initial_state=hi_all, - silu_activation=True, - pad_slot_id=_PAD_SLOT_ID, - out=None, - ) - q, k, v, g, beta = fused_gdn_post_conv( - prefill=mq_t, - decode=None, - a=a, - b=b, - A_log=L.a_log, - dt_bias=L.dt_bias, - num_k_heads=core.num_k_heads, - head_k_dim=core.head_k_dim, - num_v_heads=core.num_v_heads, - head_v_dim=core.head_v_dim, - l2_norm_eps=_L2_NORM_EPS, - softplus_threshold=_SOFTPLUS_THRESHOLD, - ) - # q and k are already L2-normalized; the recurrence must not repeat it. - chunk_gated_delta_rule( + # 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, @@ -3468,20 +4202,29 @@ def _gdn_core_generic(self, proj, out, i) -> None: beta, scale=None, initial_state=ssm_states, - initial_state_indices=si_all, + initial_state_indices=si_p, inplace_indexed_state_update=True, output_final_state=False, - cu_seqlens=qsl_long, + cu_seqlens=qsl_long_p, + head_first=False, use_qk_l2norm_in_kernel=False, - output=out[:, :num_tokens], + 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 -- one pass over `mq`, the token-major - q|k|v column view of the projection.""" + 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, @@ -3501,12 +4244,14 @@ def _gdn_prefill_prep(self, mq, a, b, L, i, qsl, si, hi): 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`). + parts = _fp8_amax_parts(num_tokens, L.out8, "out", out.device) 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]), @@ -3514,8 +4259,10 @@ def _gdn_epilogue(self, out, z, L, num_tokens) -> torch.Tensor: eps=core.eps, fp8_scale=None, gate_activation="sigmoid", + amax_parts=parts, ) - return cublas_mm(torch.reshape(o, [num_tokens, core.value_dim]), L.out) + o = torch.reshape(o, [num_tokens, core.value_dim]) + return _dense_mm(o, L.out, L.out8, cublas_mm, "out", amax_parts=parts) 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, @@ -3532,12 +4279,12 @@ def _ple_block(self, hc, hc_out, normed_out) -> None: the hash cannot distinguish. The id history and the conv state publish the state after the whole chunk. - Outside a verify step the window, hash, gather and dequantize are - `ple_ngram_embed`, and the gate, the conv, the residual add and the next site's - grouped norm `ple_gate_conv_inject_norm`. On a verify step (`stage`) the - generation rows' pre-step pool rows, id windows and conv input go to - `SpeculativeRollback`, which rewrites both pools to the accepted prefix once the - sampler has decided it, so that step takes `_ple_delta`, which exposes them.""" + 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 @@ -3545,102 +4292,50 @@ def _ple_block(self, hc, hc_out, normed_out) -> None: num_tokens = step.num_tokens next_norm = self._layers[core.ple_layer].ahc_norm hc = hc[:num_tokens] + conv_input = None if step.stage: - stream = torch.add(hc, self._ple_delta(hc)) - stream, normed = qwen4_exp_hc_inject_norm( - stream, None, None, next_norm, core.eps, core.hc_count - ) - else: - ngram = ple_ngram_embed( - step.input_ids, - ngram_pool, - si_all, - carry_all, - qsl, - core.w["ple_table"], - S.ple_tscale, - S.ple_mult, - S.ple_vocab, - S.ple_offs, - core.eos_token_id, - num_tokens, - ) - stream, normed = ple_gate_conv_inject_norm( - hc, - cublas_mm(ngram, S.ple_key), - cublas_mm(ngram, S.ple_val), - S.ple_nk_w, - S.ple_nq_w, - S.ple_nc_w, - S.ple_conv_w, - conv_states, - si_all, - qsl, - next_norm, - core.eps, - core.hc_count, - _PLE_CLAMP_MIN, - core.ple_dilation, - _PAD_SLOT_ID, - num_tokens, - carry_all, - ) - hc_out[:num_tokens].copy_(stream) - normed_out[:num_tokens].copy_(normed) - hc_out[num_tokens:].zero_() - normed_out[num_tokens:].zero_() - - def _ple_delta(self, hc) -> torch.Tensor: - """A verify step's PLE output for the batch's `n` stream rows `hc`, before it is - added into the stream, through the unfused ops whose intermediates - `SpeculativeRollback` stages: the pre-step pool rows, the id windows and the conv - input.""" - core = self.core - S = self._stack - step = self._step - conv_states, ngram_pool, si_all, carry_all, qsl = step.ple - core.spec_rollback.stage_ple_pools(ngram_pool, conv_states) - window = ngram_token_history( - step.input_ids, ngram_pool, si_all, carry_all, qsl, core.eos_token_id, core.ngram_size - ) - rows = torch.reshape( - ngram_hash( - window, - layer_multipliers=S.ple_mult, - head_vocab_sizes=S.ple_vocab, - head_offsets=S.ple_offs, - eos_token_id=core.eos_token_id, - num_query_tokens=1, - ), - [-1, core.ngram_heads], - ) - gathered = nn.functional.embedding(rows, core.w["ple_table"]) - ngram = torch.reshape( - dequantize_e4m3_per_tensor(gathered, S.ple_tscale), [-1, core.ple_embed_dim] - ) - - key_normed = self._grouped_norm(cublas_mm(ngram, S.ple_key), S.ple_nk) - query_normed = self._grouped_norm(hc, S.ple_nq) - value = cublas_mm(ngram, S.ple_val) - gated = gated_ngram_injection( - key_normed, query_normed, value, core.hc_count, _PLE_CLAMP_MIN + 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, ) - conv_in = self._grouped_norm(gated, S.ple_nc) - core.spec_rollback.stage_ple_rows(window, conv_in) - # The short conv reads channel-major; a bare transpose is the certified form. - convolved = dilated_short_conv_paged( - x=torch.transpose(conv_in, 0, 1), - weight=S.ple_conv_w, - bias=None, - conv_states=conv_states, - query_start_loc=qsl, - cache_indices=si_all, - has_initial_state=carry_all, - dilation=core.ple_dilation, - silu_activation=True, - pad_slot_id=_PAD_SLOT_ID, + 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]), ) - return torch.add(gated, torch.transpose(convolved, 0, 1)) + 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 @@ -3838,6 +4533,16 @@ def prepare( 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, @@ -3848,6 +4553,7 @@ def prepare( qsa=qsa, stage=stage, vision=vision, + nvfp4_context=nvfp4_context, ) def forward( @@ -3874,7 +4580,12 @@ def forward( 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`).""" + 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, ( @@ -3918,34 +4629,68 @@ def forward( hc_count = core.hc_count hc = torch.cat([h] * hc_count, -1) layers, S = self._layers, self._stack - _, normed = qwen4_exp_hc_inject_norm(hc, None, None, layers[0].ahc_norm, core.eps, hc_count) + # 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 - mixed, inject = self._hc_mixer(normed, L.ahc_first, L.ahc_up, True) + # 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) - 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) + 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) - hc, normed = qwen4_exp_hc_inject_norm( - hc, routed, inject, next_norm, core.eps, hc_count, shared, gate_logits - ) + 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. - final, _ = self._hc_mixer(normed, S.fhc_first, S.fhc_up, False) + 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 @@ -3956,6 +4701,23 @@ def forward( # ------------------------------------------------------------- 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. @@ -3963,9 +4725,12 @@ def _mtp_moe_block(self, x, dt) -> torch.Tensor: 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).""" + 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) @@ -3997,7 +4762,7 @@ def _mtp_moe_block(self, x, dt) -> torch.Tensor: gate_logits = cublas_mm(x, M.sh_gate) return fused_sigmoid_gate_mul_add(routed, gate_logits, shared) - def _capture_selection(self, md, topk_indices, visible_blocks) -> None: + 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. @@ -4006,7 +4771,8 @@ def _capture_selection(self, md, topk_indices, visible_blocks) -> None: 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.""" + 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 " @@ -4018,6 +4784,10 @@ def _capture_selection(self, md, topk_indices, visible_blocks) -> None: 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_( @@ -4161,31 +4931,66 @@ def mtp_forward( cublas_mm(torch.reshape(h_n, [rows * core.hc_count, core.hidden]), M.fc_h), [rows, core.stream], ) - if _MTP_EMBED_BLOCK_ALL_STREAMS: - x = torch.add(x_h, torch.cat([x_e] * core.hc_count, -1)) - else: - x = torch.add( - x_h, - nn.functional.pad(x_e, [0, core.stream - core.hidden], mode="constant", value=0.0), + # With the fused views 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 + fused = F is not None and not core._decode._mtp_unfused and _MTP_EMBED_BLOCK_ALL_STREAMS + if fused: + x, mixed, inject = self._mtp_site( + x_h, x_e, F.unit_inject[:rows], F.ahc_norm, F.ahc_first, F.ahc_up ) + else: + if _MTP_EMBED_BLOCK_ALL_STREAMS: + x = torch.add(x_h, torch.cat([x_e] * core.hc_count, -1)) + else: + x = torch.add( + x_h, + nn.functional.pad( + x_e, [0, core.stream - core.hidden], mode="constant", value=0.0 + ), + ) # ---- the decoder layer ---------------------------------------------- - mixed, inject = qwen4_exp_gated_residual( - self._grouped_norm(x, M.ahc_norm), M.ahc_down, M.ahc_up, M.ahc_inject, core.hc_count - ) + if not fused: + mixed, inject = qwen4_exp_gated_residual( + self._grouped_norm(x, M.ahc_norm), + M.ahc_down, + M.ahc_up, + M.ahc_inject, + core.hc_count, + ) + # 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([rows, core.q_width], dtype=dt, device=dev) - qg = cublas_mm(mixed, M.q) + 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, [rows, core.heads_q, 2 * core.head_dim]), + torch.reshape(qg, [q_rows, core.heads_q, 2 * core.head_dim]), [core.head_dim, core.head_dim], 2, ) - q_dst.copy_(q_src) + 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`). @@ -4234,8 +5039,12 @@ def mtp_forward( block_table, tokens_per_block, ) - index_scratch = torch.empty([rows, core.index_scratch_width], dtype=dt, device=dev) - index_scratch[:, : core.index_width].copy_(qk_index) + 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, @@ -4248,7 +5057,7 @@ def mtp_forward( k_weight=M.idx_kn, base=core.mtp_theta, is_neox=True, - position_ids=pos3 if _QSA_INDEX_USE_MROPE else pos, + 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, @@ -4258,7 +5067,7 @@ def mtp_forward( q_index = torch.split( torch.reshape( index_scratch, - [rows, core.index_heads + 2 * core.index_kv_heads, core.index_head_dim], + [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, @@ -4282,50 +5091,78 @@ def mtp_forward( # 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=logical, - request_indices=req_idx, + query_positions=q_logical, + request_indices=q_req, tokens_per_block=tokens_per_block, compress_ratio=core.compress_ratio, only_visible_blocks=True, - context_rows=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[:rows] - visible_blocks = md.qsa_visible_blocks[:rows] + topk_indices = md.qsa_topk_indices[:q_rows] _select( core, scores, - md.qsa_topk_row_starts[:rows], + row_starts, visible_blocks, topk_indices, _live_score_columns(core, md, batch), ) if share: - self._capture_selection(md, topk_indices, visible_blocks) + self._capture_selection(md, topk_indices, visible_blocks, q_keep is not None) selected = qsa_expand_block_indices( topk_indices, - logical, - md.qsa_sequence_lengths[:rows], + q_logical, + seq_lengths, compress_ratio=core.compress_ratio, token_topk=core.token_topk, ) - attn_out = torch.reshape( - _sparse_gqa( + 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) + _sparse_gqa_nvfp4_context( core, kv, - q=q_dst, - block_table=block_table, - selected_tokens=selected, - request_indices=req_idx, - tokens_per_block=tokens_per_block, - query_positions=logical, - compress_ratio=core.compress_ratio, - ), - [rows, core.q_width], - ) + 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, @@ -4342,12 +5179,43 @@ def mtp_forward( **_CALL_CONSTANTS, ) fused_sigmoid_mul_inplace(attn_out, gate) - x = qwen4_exp_post_inject(x, cublas_mm(attn_out, M.o), inject) + if keep is not None: + if q_keep is None: + attn_out = attn_out[keep] + out, inject = x, inject[keep] + x = x[keep] + if fused: + x, mixed, inject = self._mtp_site( + x, cublas_mm(attn_out, M.o), inject, F.mhc_norm, F.mhc_first, F.mhc_up + ) + else: + x = qwen4_exp_post_inject(x, cublas_mm(attn_out, M.o), inject) + mixed, inject = qwen4_exp_gated_residual( + self._grouped_norm(x, M.mhc_norm), + M.mhc_down, + M.mhc_up, + M.mhc_inject, + core.hc_count, + ) + 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) - mixed, inject = qwen4_exp_gated_residual( - self._grouped_norm(x, M.mhc_norm), M.mhc_down, M.mhc_up, M.mhc_inject, core.hc_count - ) - return qwen4_exp_post_inject(x, self._mtp_moe_block(mixed, dt), inject) + +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: @@ -4366,6 +5234,10 @@ def __init__(self, core: Qwen4ExpTp1ModelingV2) -> None: ) self._layers = _layer_views(core) self._stack = _stack_views(core) + self._mtp_unfused = os.environ.get(_MTP_UNFUSED_ENV, "0") == "1" + # `(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.""" @@ -4472,10 +5344,29 @@ def _hc_site(self, hc, block_out, inject, norm, first, up, with_inject, normed=N `(gate_logits, shared, topk_weights, topk_rows)` `_moe_block` left uncombined: the site's first kernel forms the block output itself. With the top-k pair `block_out` is the runner's permuted FC2 rows and that - kernel (on either route) also forms their finalize. A site handed + 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( @@ -4519,6 +5410,21 @@ def _hc_moe_site(self, hc, block_out, inject, norm, first, up, gscale): 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 @@ -4551,12 +5457,14 @@ def _moe_block( per-capture constant) the runner also skips its finalize: `routed` is then its permuted FC2 rows and `topk_weights` / `topk_rows` its routing weights and index map, which that site's first kernel - (`qwen4_exp_hc_decode_first_cute` or `qwen4_exp_hc_inject_norm`) gathers; + (`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. `x` is the bf16 hidden states (router and shared expert); `xq` / + 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` + @@ -4571,17 +5479,31 @@ def _moe_block( # 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): - shared = gemm(flashinfer_silu_and_mul(gemm(x, L.sh_gu)), L.sh_dn) + 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) - # The routing reads its logits as a dense row-major buffer from data_ptr() and - # ignores strides; both projection entries allocate them contiguous. - logits = gemm(x, L.router) + if 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) - moe_op = fp4_block_scale_moe_renorm_permuted if fold else fp4_block_scale_moe_renorm_routed + if fold and core.moe_gemv and x.shape[0] in _MOE_GEMV_ROWS: + moe_op = nvfp4_moe_gemv_op + elif fold: + moe_op = fp4_block_scale_moe_renorm_permuted + else: + moe_op = fp4_block_scale_moe_renorm_routed out = moe_op( logits, xq, @@ -4701,12 +5623,12 @@ def _attention_block( advanced, so a request that grows past the threshold finds its prefix already compressed; above it the dense core is replaced outright and the main K/V append becomes this layer's own call, because the op that would - have done it is not the one running. Above the threshold over a bf16 pool, - on a step with one row per request, everything up to the index scores is one - launch, `qwen4_exp_attn_decode_prologue`. A verify step stays on the composed - prologue: at a group's last token that launch reads the group's earlier keys - from the index plane, and in a verify group the rows before it are written by - the same launch.""" + 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, @@ -4723,9 +5645,19 @@ def _attention_block( ) = qsa num_tokens = x.shape[0] - proj = gemm(x, L.attn_in) 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 ) @@ -4741,13 +5673,20 @@ def _attention_block( 1, )[0] index_section = (_QSA_MROPE_SECTION or (0, 0, 0)) if _QSA_INDEX_USE_MROPE else (0, 0, 0) - index_cache = core._bound.index_cache[i] - kv = core._bound.kv[i] - fused_prologue = sparse and not stage and kv.k_scale is None and kv.k.dtype == proj.dtype + # 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, and the index branch's raw-key and coordinate store and group - # compression in one launch on this stream. + # 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, @@ -4780,6 +5719,11 @@ def _attention_block( 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( @@ -4796,24 +5740,41 @@ def _attention_block( # Append the fused-norm+rope'd K/V; the selection reads the current # token from the pool. Over an fp8 or NVFP4 pool as in the prefill target. _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block) - scores = qsa_paged_index_scores( - q=q_index, - index_cache=index_cache, - 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. - use_pdl=fused_prologue, - prefetch=fused_prologue and num_tokens <= _SPARSE_DECODE_MAX_ROWS, - ) + rows_per_request = step["predicted_tokens_per_seq"] + if scores is None and stage and rows_per_request * core.index_heads <= 16: + # A verify step: one gather of each request's compressed keys serves its + # rows (bitwise the row-wise entry; its keys are read after the wait). + scores = qsa_paged_index_scores_verify( + q_index, + index_cache, + block_table, + logical, + req_idx, + tokens_per_block, + core.compress_ratio, + rows_per_request, + use_pdl=fused_prologue, + ) + 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). @@ -4836,6 +5797,8 @@ def _attention_block( 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, @@ -4847,7 +5810,10 @@ def _attention_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. @@ -4904,34 +5870,41 @@ def _attention_block( ) if not gated: fused_sigmoid_mul_inplace(out, gate) - return gemm(out, L.o) + 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 + 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`.""" + `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 - # [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]) + # 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) + 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 ) @@ -4981,17 +5954,83 @@ def _gdn_block(self, x, i, L, gdn_state, gemm=cublas_mm, fused_conv=False) -> to rows runs the conv step as that GEMM's epilogue (`gdn_in_proj_conv1d_pdl`), and over an fp32 recurrent pool at up to `_GDN_DECODE_FUSED_MAX_ROWS` rows the recurrence and the gated norm as one launch (`gdn_decode_gated_norm_pdl`); past - them, after the conv step, as one streaming launch (`gdn_decode_rows_gated_norm`). - Over a bf16 pool the recurrence is `fused_sigmoid_gating_delta_rule_update`'s rows - form, which is batch-invariant.""" + 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] - fp32_state = ssm_states.dtype == torch.float32 conv_fused = fused_conv and not is_verify and num_tokens <= _GDN_CONV_FUSED_MAX_ROWS - if conv_fused: + # 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, @@ -5002,11 +6041,21 @@ def _gdn_block(self, x, i, L, gdn_state, gemm=cublas_mm, fused_conv=False) -> to _PAD_SLOT_ID, True, ) + elif nv is not None: + proj = nvfp4_split_gemv_op(x, nv) else: - proj = gemm(x, L.in_proj) - mixed_qkv, z, a, b = torch.split(proj, core.in_proj_split, 1) + 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 fp32_state and conv_fused and num_tokens <= _GDN_DECODE_FUSED_MAX_ROWS: + 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 @@ -5022,12 +6071,14 @@ def _gdn_block(self, x, i, L, gdn_state, gemm=cublas_mm, fused_conv=False) -> to core.num_k_heads, core.head_k_dim, ) - return gemm(o, L.out) + 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.contiguous(), - b.contiguous(), + a if strided_ab else a.contiguous(), + b if strided_ab else b.contiguous(), ssm_states, i, L, @@ -5036,9 +6087,10 @@ def _gdn_block(self, x, i, L, gdn_state, gemm=cublas_mm, fused_conv=False) -> to num_decodes, gen_p, replay, + mq_out, ) return self._gdn_epilogue(out, z, L, num_tokens, gemm) - if not conv_fused: + if not (conv_fused or fp8_conv or nv_conv): causal_conv1d_update( x=mixed_qkv, conv_state=core._bound.conv[i], @@ -5049,51 +6101,23 @@ def _gdn_block(self, x, i, L, gdn_state, gemm=cublas_mm, fused_conv=False) -> to conv_state_indices=si_all, pad_slot_id=_PAD_SLOT_ID, ) - if fp32_state: - # The recurrence and the gated norm as one launch over `proj`, whose q/k/v - # columns the conv step rewrote in place; it computes both gates and the - # q/k L2 norm itself, at the softplus pair (1, 20) and epsilon 1e-6 this - # target's constants name, and reads this layer's pool rows before its PDL - # wait (nothing in a single-token step writes them ahead of it). - o = gdn_decode_rows_gated_norm( - proj, - ssm_states, - si_all, - L.a_log, - L.dt_bias, - L.gnorm, - core.eps, - core.num_k_heads, - core.head_k_dim, - ) - return gemm(o, L.out) - else: - # The rows form, `cu_seqlens` None: it finds row n's tokens at - # `n * q.stride(1)`, so q/k/v are split from a contiguous `[T, 10240]` copy - # reshaped to `[T, 1, 80, 128]` (K and V heads are both 128 wide), where the - # size-1 token axis carries the row stride; a/b are contiguous `[T, 48]`. - qkv_rows = torch.reshape(mixed_qkv.contiguous(), [num_tokens, 1, -1, core.head_k_dim]) - q, k, v = torch.split( - qkv_rows, [core.num_k_heads, core.num_k_heads, core.num_v_heads], 2 - ) - out = fused_sigmoid_gating_delta_rule_update( - A_log=L.a_log, - a=torch.reshape(a.contiguous(), [num_tokens, 1, core.num_v_heads]), - dt_bias=L.dt_bias, - softplus_beta=_SOFTPLUS_BETA, - softplus_threshold=_SOFTPLUS_THRESHOLD, - q=q, - k=k, - v=v, - b=torch.reshape(b.contiguous(), [num_tokens, 1, core.num_v_heads]), - initial_state_source=ssm_states, - initial_state_indices=si_all, - scale=None, - use_qk_l2norm_in_kernel=True, - cu_seqlens=None, - output=None, - ) - return self._gdn_epilogue(out, z, L, num_tokens, gemm) + # 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 @@ -5107,7 +6131,8 @@ def _gdn_epilogue(self, out, z, L, num_tokens, gemm=cublas_mm) -> torch.Tensor: fp8_scale=None, gate_activation="sigmoid", ) - return gemm(torch.reshape(o, [num_tokens, core.value_dim]), L.out) + 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 @@ -5129,54 +6154,14 @@ def _ple_block(self, hc, input_ids, ple_state, next_norm, stage, gemm=cublas_mm) (sequence, stream). Every generation sequence carries its history by construction, so no carry - flags are passed. A verify step (`stage`) takes the unfused ops instead, whose - intermediates `SpeculativeRollback` stages (see `PrefillTarget._ple_delta`).""" + 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_pools(ngram_pool, conv_states) - window = ngram_token_history( - input_ids, ngram_pool, si_all, carry_all, qsl, core.eos_token_id, core.ngram_size - ) - rows = torch.reshape( - ngram_hash( - window, - layer_multipliers=S.ple_mult, - head_vocab_sizes=S.ple_vocab, - head_offsets=S.ple_offs, - eos_token_id=core.eos_token_id, - num_query_tokens=1, - ), - [-1, core.ngram_heads], - ) - gathered = nn.functional.embedding(rows, core.w["ple_table"]) - ngram = torch.reshape( - dequantize_e4m3_per_tensor(gathered, S.ple_tscale), [-1, core.ple_embed_dim] - ) - key_normed = self._grouped_norm(cublas_mm(ngram, S.ple_key), S.ple_nk) - query_normed = self._grouped_norm(hc, S.ple_nq) - value = cublas_mm(ngram, S.ple_val) - gated = gated_ngram_injection( - key_normed, query_normed, value, core.hc_count, _PLE_CLAMP_MIN - ) - conv_in = self._grouped_norm(gated, S.ple_nc) - core.spec_rollback.stage_ple_rows(window, conv_in) - # The short conv reads channel-major; a bare transpose is the certified form. - convolved = dilated_short_conv_paged( - x=torch.transpose(conv_in, 0, 1), - weight=S.ple_conv_w, - bias=None, - conv_states=conv_states, - query_start_loc=qsl, - cache_indices=si_all, - has_initial_state=carry_all, - dilation=core.ple_dilation, - silu_activation=True, - pad_slot_id=_PAD_SLOT_ID, - ) - hc = torch.add(hc, torch.add(gated, torch.transpose(convolved, 0, 1))) - return qwen4_exp_hc_inject_norm(hc, None, None, next_norm, core.eps, core.hc_count) + core.spec_rollback.stage_ple_verify(ngram_pool, conv_states, input_ids) ngram = ple_ngram_embed( input_ids, ngram_pool, @@ -5191,7 +6176,8 @@ def _ple_block(self, hc, input_ids, ple_state, next_norm, stage, gemm=cublas_mm) core.eos_token_id, gen_p, ) - return ple_gate_conv_inject_norm( + 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), @@ -5209,7 +6195,11 @@ def _ple_block(self, hc, input_ids, ple_state, next_norm, stage, gemm=cublas_mm) 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 --- @@ -5230,15 +6220,22 @@ def forward( # (`_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 = 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])) + pos3 = None if fused_setup else _copy_positions(position_ids) h = nn.functional.embedding(input_ids, core.w["embed"]) num_tokens = h.shape[0] @@ -5297,14 +6294,27 @@ def forward( replay = _replay_work(core, mm, batch) gdn_state = (si_all, is_verify, batch, gen_p, int_idx, replay) - # Every generation sequence carries its n-gram history; the engine's - # `has_initial_states` covers the context prefix only (empty here). - ple_carry = torch.ones( - [batch], dtype=mm.has_initial_states.dtype, device=mm.has_initial_states.device - ) # PLE segmentation vector: rebuilt from `seq_lens_cuda` when `gen_p > 1` (see - # PrefillTarget.prepare). - if gen_p == 1: + # 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) @@ -5312,10 +6322,11 @@ def forward( core._bound.ple_state[0], core._bound.ple_state[1], si_all, - ple_carry, + 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. @@ -5338,10 +6349,18 @@ def forward( # Staging for the speculative rollback, as on the prefill target (see # PrefillTarget.prepare); every row here is a generation row. stage = is_verify - rollback = core.spec_rollback if rollback is not None: rollback.discard() - if stage: + 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, @@ -5352,6 +6371,7 @@ def forward( 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 @@ -5431,9 +6451,12 @@ def forward( # ------------------------------------------------------------- the draft --- def _mtp_moe_block(self, x, dt) -> torch.Tensor: - """FP8 block-scale draft MoE; see `PrefillTarget._mtp_moe_block`.""" + """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) @@ -5465,6 +6488,291 @@ def _mtp_moe_block(self, x, dt) -> torch.Tensor: gate_logits = cublas_mm(x, M.sh_gate) return fused_sigmoid_gate_mul_add(routed, gate_logits, shared) + def _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 @@ -5595,17 +6903,41 @@ def mtp_forward( 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 + # 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) + if core._mtp_fused is not None and not self._mtp_unfused: + 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, + ) + # ---- 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) @@ -6105,14 +7437,31 @@ def shared_head( mixer with its own parameters, not a plain RMSNorm), then the trunk's `lm_head` through the shell's `logits_processor`.""" core = self.core - M = core._mtp - parts = torch.split(hidden_states, [core.hidden] * core.hc_count, 1) - normed = torch.cat( - [flashinfer_gemma_rmsnorm(part, w, core.eps) for w, part in zip(M.fhc_norm, parts)], - 1, + F = core._mtp_fused + if F is not None and not core._decode._mtp_unfused: + # 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 + ) + else: + M = core._mtp + parts = torch.split(hidden_states, [core.hidden] * core.hc_count, 1) + normed = torch.cat( + [flashinfer_gemma_rmsnorm(part, w, core.eps) for w, part in zip(M.fhc_norm, parts)], + 1, + ) + final, _ = qwen4_exp_gated_residual(normed, M.fhc_down, M.fhc_up, None, core.hc_count) + if core._mtp_head_fp8 is not None: + packed = core._mtp_head_g if core._mtp_head_g is not None else core._lm_head_g + lm_head = _Fp8DraftHead(*core._mtp_head_fp8, packed, 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), ) - final, _ = qwen4_exp_gated_residual(normed, M.fhc_down, M.fhc_up, None, core.hc_count) - return self.logits_processor.forward(final, lm_head, attn_metadata, return_context_logits) def forward( self, @@ -6179,17 +7528,73 @@ class SpeculativeRollback: * every index-K and position slot written by a row of ordinal `>= a` to its pre-verify contents. - The targets stage during the forward (`begin`, then `stage_ple_pools` / - `stage_ple_rows` in the PLE block and `stage_index` before each attention layer's - store) and the spec worker calls `commit_speculative_states` after its draft loop, - inside the same model forward. Context rows and non-verify steps stage nothing, and a - commit with nothing staged returns. Staging and commit use metadata-sized + 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.""" + 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 --- @@ -6209,12 +7614,14 @@ def begin( 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.""" + 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} " @@ -6236,7 +7643,18 @@ def begin( 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' @@ -6247,16 +7665,39 @@ def stage_ple_pools(self, ngram_pool: torch.Tensor, conv_pool: torch.Tensor) -> pending["conv_pool"] = conv_pool pending["conv_prev"] = conv_pool[slots] - def stage_ple_rows(self, window: torch.Tensor, conv_input: torch.Tensor) -> None: - """After the id history and before the conv: the generation rows' id windows - (`[T, ngram_size]`, newest id last) and conv input (`[T, C]`, token-major). Copied - rather than viewed, so a long context prefix in a mixed batch is not kept alive + 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"] - # Row t of the window without its oldest id is the history after token t. - pending["ngram_candidates"] = torch.reshape(window[prefix:, 1:], [n, p, -1]).clone() - pending["conv_x"] = torch.reshape(conv_input[prefix:], [n, p, -1]).clone() + 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 @@ -6266,6 +7707,9 @@ def stage_index(self, index_cache: torch.Tensor, position_cache: torch.Tensor) - 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 --- @@ -6293,6 +7737,26 @@ def commit_speculative_states( "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) @@ -6308,8 +7772,14 @@ def commit_speculative_states( torch.arange(p, device=accepted.device).unsqueeze(0) >= accepted.unsqueeze(1), [-1] ) pages, within = pending["pages"], pending["within"] - for cache, old in pending["index"]: - cache[pages, within] = torch.where(rejected[:, None, None], old, cache[pages, within]) + 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]) @@ -6338,17 +7808,124 @@ def __call__(self, hidden_states: torch.Tensor) -> torch.Tensor: 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`.""" + """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 + 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] ): @@ -6771,12 +8348,16 @@ def forward( ) ) # `gather_ids`: one row per context request, `runtime_draft_len + 1` per - # generation request. + # 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. @@ -6801,3 +8382,66 @@ def load_weights(self, weights, *args, **kwargs): def post_load_weights(self): super().post_load_weights() self.model.post_load_weights() + draft_fp8 = self.model.mtp_enabled and _mtp_draft_head_mode() == "fp8" + 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) + vocab = _mtp_draft_vocab() + if draft_fp8 and vocab > 0: + self._slice_draft_vocab(rowwise, vocab) + + 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, + ) + + def _slice_draft_vocab(self, rowwise: tuple, vocab: int) -> None: + """The draft head over token ids `[0, vocab)` and the vocabulary's tail from the + added tokens on (`_MTP_DRAFT_VOCAB_ENV`); the worker maps a draft id `i` back to + `i + d2t[i]`.""" + weight = self.lm_head.weight.data + full = weight.shape[0] + tail = full - -(-(full - self.model.eos_token_id) // 128) * 128 + assert vocab % 128 == 0 and vocab <= tail, ( + f"{_MTP_DRAFT_VOCAB_ENV}={vocab}: a multiple of 128 below the added tokens' " + f"block at {tail}" + ) + if vocab == tail: + return + rows = torch.cat([torch.arange(vocab), torch.arange(tail, full)]).to(weight.device) + codes, scale = rowwise + self.model._mtp_head_fp8 = (codes[rows], scale[:, rows].contiguous()) + mid_rows_precompile( + _head_mid(*self.model._mtp_head_fp8), _FP8_GEMV_HEAD_ROWS + 1, _MID_HEAD_ROWS + ) + if self.model._mtp_head_nv is not None: + self._build_draft_head_nvfp4(weight[rows]) + else: + self.model._mtp_head_g = pack_from_bf16(weight[rows]) + d2t = torch.zeros(rows.numel(), dtype=torch.long, device=weight.device) + d2t[vocab:] = tail - vocab + self.draft_model.model = types.SimpleNamespace(d2t=d2t) + self.spec_worker.set_draft_model(self.draft_model) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py index 7987c692ddfd..89fa2e2418fc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py @@ -4,7 +4,11 @@ from __future__ import annotations +from typing import Optional + import torch +import triton +import triton.language as tl # Registered by the module that defines it: after `import # tensorrt_llm._torch.custom_ops` alone, torch.ops.trtllm has no @@ -16,6 +20,70 @@ _CKPT = dict(heads=24, head_dim=256, dtype=torch.bfloat16, gate="interleaved") +@triton.jit +def _sigmoid_mul_amax_kernel( + out_ptr, + gate_ptr, + parts_ptr, + out_stride_token, + gate_stride_token, + gate_stride_head, + num_parts, + hidden_size: tl.constexpr, + head_dim: tl.constexpr, + block_size: tl.constexpr, +): + """The registered `_fused_sigmoid_mul_kernel`'s arithmetic in place, plus the maximum of + the stored `|out|` folded into `parts[program % num_parts]`.""" + token = tl.program_id(0).to(tl.int64) + block = tl.program_id(1) + offsets = block * block_size + tl.arange(0, block_size) + mask = offsets < hidden_size + head = offsets // head_dim + dim = offsets - head * head_dim + attention = tl.load(out_ptr + token * out_stride_token + offsets, mask=mask, other=0.0).to( + tl.float32 + ) + gate = tl.load( + gate_ptr + token * gate_stride_token + head * gate_stride_head + dim, mask=mask, other=0.0 + ).to(tl.float32) + y = (attention * tl.sigmoid(gate)).to(out_ptr.dtype.element_ty) + tl.store(out_ptr + token * out_stride_token + offsets, y, mask=mask) + peak = tl.max(tl.where(mask, tl.abs(y.to(tl.float32)), 0.0), axis=0) + program = token * tl.num_programs(1) + block + tl.atomic_max(parts_ptr + program % num_parts, peak) + + +def _sigmoid_mul_amax(attention_output, gate, amax_parts) -> None: + # The registered op's layout checks, as loud. + assert attention_output.dim() == 2 and attention_output.stride(-1) == 1 + num_tokens, hidden_size = attention_output.shape + assert gate.stride(-1) == 1 + if gate.dim() == 3: + assert gate.shape[0] == num_tokens and gate.shape[1] * gate.shape[2] == hidden_size + head_dim, stride_token, stride_head = gate.shape[2], gate.stride(0), gate.stride(1) + else: + assert gate.dim() == 2 and gate.shape == attention_output.shape + head_dim, stride_token, stride_head = hidden_size, gate.stride(0), hidden_size + assert amax_parts.dtype == torch.float32 and amax_parts.dim() == 1 + if num_tokens == 0: + return + block_size = min(triton.next_power_of_2(hidden_size), 1024 if num_tokens < 1024 else 2048) + _sigmoid_mul_amax_kernel[(num_tokens, triton.cdiv(hidden_size, block_size))]( + attention_output, + gate, + amax_parts, + attention_output.stride(0), + stride_token, + stride_head, + amax_parts.shape[0], + hidden_size=hidden_size, + head_dim=head_dim, + block_size=block_size, + num_warps=4, + ) + + class FusedSigmoidMulInplace(OpWrapper): """`attention_output[t, i] *= sigmoid(gate[t, i // D, i % D])` in place, `D = gate.shape[-1]`. @@ -26,8 +94,14 @@ class FusedSigmoidMulInplace(OpWrapper): 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. + With `amax_parts` (a zeroed fp32 buffer, `fp8_per_tensor_mm.amax_partials`) the same + arithmetic runs in this target's copy of the kernel, which also folds the maximum of the + stored `|attention_output|` into the buffer (atomic maxima, one slot per program modulo + its length) for the per-tensor FP8 o projection. + + Fusion boundary: sigmoid and multiply (and the output maximum) only. The caller owns + producing both operands and everything downstream (the o projection). `gate` is + read-only. """ ARCHS = frozenset({Arch.SM_103}) @@ -45,6 +119,17 @@ class FusedSigmoidMulInplace(OpWrapper): why="MTP decode: 1 token, 24x256 gate as the strided half of q_proj", spec=dict(tokens=1, **_CKPT), ), + Cell( + why="prefill 16384 tokens folding the output maximum for the FP8 o projection", + spec=dict( + tokens=16384, + heads=24, + head_dim=256, + dtype=torch.bfloat16, + gate="fused_row", + amax=True, + ), + ), ) note = """ @@ -82,11 +167,24 @@ class FusedSigmoidMulInplace(OpWrapper): 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) + def __call__( + self, + attention_output: torch.Tensor, + gate: torch.Tensor, + amax_parts: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if amax_parts is None: + torch.ops.trtllm.fused_sigmoid_mul_inplace(attention_output, gate) + else: + _sigmoid_mul_amax(attention_output, gate, amax_parts) return attention_output - def reference(self, attention_output: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + def reference( + self, + attention_output: torch.Tensor, + gate: torch.Tensor, + amax_parts: Optional[torch.Tensor] = None, + ) -> 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 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 new file mode 100644 index 000000000000..7e497b738c47 --- /dev/null +++ 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 @@ -0,0 +1,389 @@ +# 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, for 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, `listed` 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 new file mode 100644 index 000000000000..d6b416ac80b1 --- /dev/null +++ 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 @@ -0,0 +1,248 @@ +# 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_verify.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_verify.py new file mode 100644 index 000000000000..1284b4c14a64 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores_verify.py @@ -0,0 +1,281 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA compressed-block index scores for a verify step's rows: one key gather per request. + +`qsa_paged_index_scores` scores each row on its own, so a verify step's `P` rows of one +request gather that request's compressed keys `P` times. Here one program serves a +request's `P` rows: it gathers a tile of keys once and multiplies it by the `P` rows' +index queries packed side by side (`P` x 4 heads <= 16 columns, the width the row-wise +kernel pads its 4 heads to), then reduces each row's four head columns. + +Bitwise the row-wise kernel's scores: the same tile shapes, warps and dot (a score column +of an MMA depends on its own query column alone), each row's relu'd head sum over the +same thread's pair of columns and the same shuffles, and the same scale. +""" + +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_paged_index_scores import qsa_paged_index_scores as _row_wise + +#: The row-wise kernel's head padding: `P * heads` query columns fit in it. +_BLOCK_H = 16 + +#: This checkpoint's index geometry and the NVFP4 trunk's index page (12 index buffers of +#: 32 x 128, then the 192-element position page: page stride 49344). +_CKPT = dict(heads=4, head_dim=128, tokens_per_block=32, compress_ratio=4, width=324) + + +@triton.jit(do_not_specialize=["rows"]) +def _qsa_index_scores_verify_kernel( + q, + index_cache, + block_table, + query_positions, + request_indices, + output, + score_scale, + rows, + q_stride_row, + q_stride_head, + q_stride_dim, + cache_stride_page, + cache_stride_token, + cache_stride_dim, + bt_stride_request, + bt_stride_page, + out_stride_row, + P: tl.constexpr, + IH: tl.constexpr, + ID: tl.constexpr, + BLOCK_D: tl.constexpr, + TPB: tl.constexpr, + CR: tl.constexpr, + MAX_BLOCKS: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_N: tl.constexpr, + TILES_PER_PROGRAM: tl.constexpr, + USE_PDL: tl.constexpr, +): + if USE_PDL: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + first_row = tl.program_id(0) * P + request = tl.load(request_indices + first_row).to(tl.int64) + # The request's last row sees the most blocks (its rows sit at consecutive positions). + last_visible = (tl.load(query_positions + first_row + P - 1) + 1) // CR + first_tile = tl.program_id(1) * TILES_PER_PROGRAM + if first_tile * BLOCK_N >= last_visible: + return + last_tile = tl.minimum(first_tile + TILES_PER_PROGRAM, tl.cdiv(last_visible, BLOCK_N)) + + # Query column c is head c % IH of the group's row c // IH; columns past P rows are zero. + columns = tl.arange(0, BLOCK_H) + column_row = columns // IH + dim_offsets = tl.arange(0, BLOCK_D) + query_values = tl.load( + q + + (first_row + column_row).to(tl.int64)[None, :] * q_stride_row + + (columns % IH)[None, :] * q_stride_head + + dim_offsets[:, None] * q_stride_dim, + mask=(dim_offsets[:, None] < ID) & (column_row < P)[None, :], + other=0.0, + ) + block_offsets = tl.arange(0, BLOCK_N) + for tile in tl.range(first_tile, last_tile, num_stages=2): + block_columns = tile * BLOCK_N + block_offsets + valid_blocks = (block_columns < MAX_BLOCKS) & (block_columns < last_visible) + anchor_positions = block_columns * CR + CR - 1 + physical_pages = tl.load( + block_table + request * bt_stride_request + (anchor_positions // TPB) * bt_stride_page, + mask=valid_blocks, + other=0, + ).to(tl.int64) + keys = tl.load( + index_cache + + physical_pages[:, None] * cache_stride_page + + (anchor_positions % TPB)[:, None] * cache_stride_token + + dim_offsets[None, :] * cache_stride_dim, + mask=valid_blocks[:, None] & (dim_offsets[None, :] < ID), + other=0.0, + eviction_policy="evict_first", + ) + per_head_scores = tl.maximum(tl.dot(keys, query_values, out_dtype=tl.float32), 0.0) + for r in tl.static_range(P): + row = first_row + r + visible = (tl.load(query_positions + row) + 1) // CR + scores = tl.sum(tl.where((column_row == r)[None, :], per_head_scores, 0.0), axis=1) + tl.store( + output + row.to(tl.int64) * out_stride_row + block_columns, + scores * score_scale, + mask=valid_blocks & (block_columns < visible), + ) + + +def _plan(rows: int): + """The row-wise kernel's tile plan for this many rows: (BLOCK_N, tiles a program).""" + return (64 if rows > 4 else 8), (8 if rows > 256 else 1) + + +class QsaPagedIndexScoresVerify(OpWrapper): + """`qsa_paged_index_scores` (`only_visible_blocks=True`, `context_rows=False`) for rows that + come `rows_per_request` consecutive rows a request at consecutive positions -- an MTP + verify step's generation rows. Returns a fresh fp32 `[rows, columns]`: each row's + visible columns bitwise the row-wise entry's, the rest not written. + + One Triton launch over `(requests, column tiles)`; each program gathers a tile of the + request's compressed keys once for all its rows. + + Fusion boundary: as `qsa_paged_index_scores`. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="MTP3 verify, 1 request x 4 rows (the row-wise plan's 8-block tiles), a group " + "end mid-request, NVFP4 trunk page", + spec=dict(starts=(8697,), per=4, nvfp4=True), + ), + Cell( + why="MTP3 verify, 32 requests x 4 rows (64-block tiles), starts at every residue " + "mod 4, NVFP4 trunk page", + spec=dict(starts=tuple(2000 + 197 * r for r in range(32)), per=4, nvfp4=True), + ), + Cell( + why="MTP3 verify, 256 requests x 4 rows (eight tiles a program), short and long " + "contexts, plain index plane", + spec=dict(starts=tuple(3 + (577 * r) % 10000 for r in range(256)), per=4), + ), + Cell( + why="MTP1 verify, 8 requests x 2 rows", + spec=dict(starts=tuple(4093 + 701 * r for r in range(8)), per=2, nvfp4=True), + ), + ) + + note = """ + Metadata consumed: as `qsa_paged_index_scores` (`qsa_block_table`, + `qsa_logical_positions`, `qsa_req_idx_per_token`). + + Preconditions: `rows % rows_per_request == 0`; each group of `rows_per_request` + consecutive rows is one request's, at consecutive increasing positions (the verify + step's layout); `rows_per_request * heads <= 16`; the row-wise entry's preconditions + (valid pages up to each row's causal bound). Not checked on the device: a group that + breaks the layout is scored against its first row's request and its last row's bound. + + Numerics: bitwise the row-wise entry on every visible column (the cells compare them + exactly). The tile plan (`BLOCK_N` 8 up to 4 rows else 64, eight tiles a program + past 256 rows, two warps) is the row-wise entry's at the same row count. + + Compiled variants: one per `rows_per_request`, tile plan, table width and PDL; the row + count is not specialized on. + """ + + 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, + rows_per_request: int, + use_pdl: bool = False, + ) -> torch.Tensor: + rows, heads, head_dim = q.shape + columns = block_table.shape[1] * tokens_per_block // compress_ratio + output = torch.empty((rows, columns), dtype=torch.float32, device=q.device) + if rows == 0: + return output + block_n, tiles = _plan(rows) + _qsa_index_scores_verify_kernel[ + (rows // rows_per_request, triton.cdiv(columns, block_n * tiles)) + ]( + q, + index_cache, + block_table, + query_positions, + request_indices, + output, + 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), + P=rows_per_request, + IH=heads, + ID=head_dim, + BLOCK_D=max(16, triton.next_power_of_2(head_dim)), + TPB=tokens_per_block, + CR=compress_ratio, + MAX_BLOCKS=columns, + BLOCK_H=_BLOCK_H, + BLOCK_N=block_n, + TILES_PER_PROGRAM=tiles, + USE_PDL=use_pdl, + num_warps=2, + launch_pdl=use_pdl, + ) + return output + + 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, + rows_per_request: int, + use_pdl: bool = False, + ) -> torch.Tensor: + """The row-wise entry on the same rows (its visible columns are the contract).""" + return _row_wise( + 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=True, + context_rows=False, + ) + + def is_valid( + self, + q, + index_cache, + block_table, + query_positions, + request_indices, + tokens_per_block, + compress_ratio, + rows_per_request, + use_pdl=False, + ) -> None: + rows, heads, _ = q.shape + assert rows_per_request >= 1 and rows % rows_per_request == 0, "whole request groups" + assert rows_per_request * heads <= _BLOCK_H, "rows_per_request x heads <= 16" + assert index_cache.ndim == 4 and index_cache.shape[2] == 1, "one index KV head" + assert q.dtype == index_cache.dtype, "matching q and index dtypes" + assert block_table.dtype == torch.int32 and block_table.ndim == 2, "int32 table" + assert query_positions.shape == (rows,) and request_indices.shape == (rows,), ( + "one position and request a row" + ) + + +qsa_paged_index_scores_verify = QsaPagedIndexScoresVerify() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py index 6909d8cde45d..75e9c55e17a8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py @@ -13,11 +13,9 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper # Approved exception (as for `qsa_paged_sparse_gqa`): QSA's sparse path has no torch.ops -# surface; its split-K buckets and row-count PDL policy are reached as Python. -from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import ( - _qsa_pdl_enabled, - _qsa_splitk_launch_config, -) +# 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) @@ -50,6 +48,9 @@ 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",) @@ -86,93 +87,195 @@ def _e2m1_to_f32(code): @triton.jit -def _load_keys( - k_cache, - k_scale_cache, +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, - dims, - k_global, - k_stride_page, - k_stride_head, - k_stride_token, - ks_stride_page, - ks_stride_head, - SCALE_GROUPS: tl.constexpr, + stride_page, + stride_head, + stride_token, + SCALE_STEP: tl.constexpr, + BLOCK_N: tl.constexpr, + HEAD_DIM: tl.constexpr, ): - """K for the selected tokens as an fp32 [HEAD_DIM, BLOCK_N] tile: e2m1 x scale x global.""" + """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( - k_cache - + pages[None, :] * k_stride_page - + kv_head * k_stride_head - + tokens[None, :] * k_stride_token - + dims[:, None] // 2, - mask=valid[None, :], + cache + + pages[:, None] * stride_page + + kv_head * stride_head + + tokens[:, None] * stride_token + + tl.arange(0, HALF)[None, :], + mask=valid[:, None], other=0, - ).to(tl.int32) - codes = (packed >> ((dims[:, None] % 2) * 4)) & 0xF - scales = ( - tl.load( - k_scale_cache - + pages[None, :] * ks_stride_page - + kv_head * ks_stride_head - + tokens[None, :] * SCALE_GROUPS - + dims[:, None] // 16, - mask=valid[None, :], - other=0, - ) - .to(tl.float8e4nv, bitcast=True) - .to(tl.float32) ) - return _e2m1_to_f32(codes) * scales * k_global + 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 _load_values( +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, - pages, - kv_head, - tokens, - valid, - dims, - v_global, + 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, - SCALE_GROUPS: tl.constexpr, + TOKENS_PER_BLOCK: tl.constexpr, + TOPK: tl.constexpr, + BLOCK_N: tl.constexpr, + HEAD_DIM: tl.constexpr, ): - """V for the selected tokens as an fp32 [BLOCK_N, HEAD_DIM] tile; V's scales sit in - four-token groups, `[tpb / 4, D / 16, 4]` per head page.""" - packed = tl.load( - v_cache - + pages[:, None] * v_stride_page - + kv_head * v_stride_head - + tokens[:, None] * v_stride_token - + dims[None, :] // 2, - mask=valid[:, None], - other=0, - ).to(tl.int32) - codes = (packed >> ((dims[None, :] % 2) * 4)) & 0xF - scales = ( - tl.load( - v_scale_cache - + pages[:, None] * vs_stride_page - + kv_head * vs_stride_head - + (tokens[:, None] // 4) * (4 * SCALE_GROUPS) - + (dims[None, :] // 16) * 4 - + tokens[:, None] % 4, - mask=valid[:, None], - other=0, - ) - .to(tl.float8e4nv, bitcast=True) - .to(tl.float32) + """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, ) - return _e2m1_to_f32(codes) * scales * v_global + 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 @@ -222,114 +325,84 @@ def _nvfp4_sparse_gqa_prefill_kernel( request = tl.load(request_indices + row).to(tl.int64) valid_request = (request >= 0) & (request < num_requests) safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) - k_global = tl.load(kv_scale_quant_orig + 1) + # 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) - dim_offsets = tl.arange(0, HEAD_DIM) + half_offsets = tl.arange(0, HALF) + token_offsets = tl.arange(0, BLOCK_N) query_heads = kv_head * GROUP_SIZE + head_offsets - query_values = tl.load( - q + row * q_stride_row + query_heads[:, None] * q_stride_head + dim_offsets[None, :], - mask=(head_offsets < GROUP_SIZE)[:, None], - other=0.0, + 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) - accumulator = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32) - token_offsets = tl.arange(0, BLOCK_N) + 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): - selected_columns = tile * BLOCK_N + token_offsets - logical_tokens = tl.load( - selected_tokens + row * selected_stride_row + selected_columns, - mask=selected_columns < TOPK, - other=-1, - ) - safe_tokens = tl.maximum(logical_tokens, 0) - logical_pages = safe_tokens // TOKENS_PER_BLOCK - token_in_page = safe_tokens % TOKENS_PER_BLOCK - valid = ( - valid_request - & (selected_columns < TOPK) - & (logical_tokens >= 0) - & (logical_pages < page_table_width) - ) - physical_pages = tl.load( - block_table - + safe_request * block_table_stride_request - + tl.minimum(logical_pages, page_table_width - 1), - mask=valid, - other=-1, - ).to(tl.int64) - valid &= (physical_pages >= 0) & (physical_pages < num_cache_pages) - safe_pages = tl.maximum(physical_pages, 0) - - keys = _load_keys( + 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, - safe_pages, - kv_head, - token_in_page, - valid, - dim_offsets, - k_global, + v_scale_cache, k_stride_page, k_stride_head, k_stride_token, - ks_stride_page, - ks_stride_head, - HEAD_DIM // 16, - ).to(query_values.dtype) - values = _load_values( - v_cache, - v_scale_cache, - safe_pages, - kv_head, - token_in_page, - valid, - dim_offsets, - v_global, v_stride_page, v_stride_head, v_stride_token, + ks_stride_page, + ks_stride_head, vs_stride_page, vs_stride_head, - HEAD_DIM // 16, - ).to(query_values.dtype) - scores = tl.dot(query_values, keys) * softmax_scale * _LOG2_E - scores = tl.where(valid[None, :], scores, -float("inf")) - next_max = tl.maximum(running_max, tl.max(scores, axis=1)) - correction = tl.math.exp2(running_max - next_max) - probabilities = tl.where( - valid[None, :], - tl.math.exp2(scores - next_max[:, None]), - 0.0, + TOKENS_PER_BLOCK, + TOPK, + BLOCK_N, + HEAD_DIM, ) - accumulator = tl.dot( - probabilities.to(values.dtype), - values, - accumulator * correction[:, None], - ) - running_sum = running_sum * correction + tl.sum(probabilities, axis=1) - running_max = next_max - normalized = tl.where( - running_sum[:, None] > 0, - accumulator / tl.maximum(running_sum[:, None], 1.0e-20), - 0.0, - ) - tl.store( + 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 - + dim_offsets[None, :], - normalized, - mask=(head_offsets < GROUP_SIZE)[:, None], + + 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 ) @@ -387,135 +460,100 @@ def _nvfp4_sparse_gqa_splitk_kernel( request = tl.load(request_indices + row).to(tl.int64) valid_request = (request >= 0) & (request < num_requests) safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) - k_global = tl.load(kv_scale_quant_orig + 1) + # 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) - dim_offsets = tl.arange(0, HEAD_DIM) + half_offsets = tl.arange(0, HALF) token_offsets = tl.arange(0, BLOCK_N) query_heads = kv_head * GROUP_SIZE + head_offsets - query_values = tl.load( - q + row * q_stride_row + query_heads[:, None] * q_stride_head + dim_offsets[None, :], - mask=(head_offsets < GROUP_SIZE)[:, None], - other=0.0, + 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) - accumulator = tl.zeros([BLOCK_M, HEAD_DIM], 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): - selected_columns = tile * BLOCK_N + token_offsets - logical_tokens = tl.load( - selected_tokens + row * selected_stride_row + selected_columns, - mask=selected_columns < TOPK, - other=-1, - ) - safe_tokens = tl.maximum(logical_tokens, 0) - logical_pages = safe_tokens // TOKENS_PER_BLOCK - token_in_page = safe_tokens % TOKENS_PER_BLOCK - valid = ( - valid_request - & (selected_columns < TOPK) - & (logical_tokens >= 0) - & (logical_pages < page_table_width) - ) - physical_pages = tl.load( - block_table - + safe_request * block_table_stride_request - + tl.minimum(logical_pages, page_table_width - 1), - mask=valid, - other=-1, - ).to(tl.int64) - valid &= (physical_pages >= 0) & (physical_pages < num_cache_pages) - safe_pages = tl.maximum(physical_pages, 0) - - keys = _load_keys( + 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, - safe_pages, - kv_head, - token_in_page, - valid, - dim_offsets, - k_global, + v_scale_cache, k_stride_page, k_stride_head, k_stride_token, - ks_stride_page, - ks_stride_head, - HEAD_DIM // 16, - ).to(query_values.dtype) - values = _load_values( - v_cache, - v_scale_cache, - safe_pages, - kv_head, - token_in_page, - valid, - dim_offsets, - v_global, v_stride_page, v_stride_head, v_stride_token, + ks_stride_page, + ks_stride_head, vs_stride_page, vs_stride_head, - HEAD_DIM // 16, - ).to(query_values.dtype) - scores = tl.dot(query_values, keys) - scores *= softmax_scale * _LOG2_E - scores = tl.where(valid[None, :], scores, -float("inf")) - next_max = tl.maximum(running_max, tl.max(scores, axis=1)) - correction = tl.math.exp2(running_max - next_max) - probabilities = tl.where( - valid[None, :], - tl.math.exp2(scores - next_max[:, None]), - 0.0, - ) - accumulator = tl.dot( - probabilities.to(values.dtype), - values, - accumulator * correction[:, None], + TOKENS_PER_BLOCK, + TOPK, + BLOCK_N, + HEAD_DIM, ) - running_sum = running_sum * correction + tl.sum(probabilities, axis=1) - running_max = next_max has_values = running_sum > 0 - normalized = tl.where( - has_values[:, None], - accumulator / tl.maximum(running_sum[:, None], 1.0e-20), - 0.0, - ) - output_mask = (head_offsets < GROUP_SIZE)[:, None] + 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: - tl.store( + out_ptr = ( output + row * output_stride_row + query_heads[:, None] * output_stride_head - + dim_offsets[None, :], - normalized, - mask=output_mask, + + 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_output - + ((split * num_rows + row) * NUM_QUERY_HEADS + query_heads[:, None]) * HEAD_DIM - + dim_offsets[None, :], - normalized, - mask=output_mask, - ) tl.store( partial_lse + (split * num_rows + row) * NUM_QUERY_HEADS + query_heads, partial_lse_values, mask=head_offsets < GROUP_SIZE, ) + 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() @@ -563,6 +601,35 @@ def _nvfp4_merge_splitk_kernel( tl.store(output + row * output_stride_row + head * output_stride_head + dim_offsets, merged) +def _splitk_launch_config( + rows: int, num_kv_heads: int, selected_width: int +) -> tuple[int, int, int]: + """(BLOCK_N, splits, warps) 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, @@ -711,14 +778,14 @@ def _sparse_gqa_nvfp4( BLOCK_TOPK=token_topk // compress_ratio, COMPRESS_RATIO=compress_ratio, BLOCK_M=block_m, - BLOCK_N=64, + BLOCK_N=_PREFILL_TILE[0], HEAD_DIM=head_dim, - num_warps=2, + num_warps=_PREFILL_TILE[1], num_stages=2, ) return output - block_n, num_splits, num_warps = _qsa_splitk_launch_config( + 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) @@ -819,19 +886,19 @@ class QsaPagedSparseGqaNvfp4(OpWrapper): t = selected_tokens[r, j] # logical position page = block_table[request_indices[r], t // tokens_per_block] - K[t] = bf16(e2m1(k_cache) * e4m3(k_scale_cache) * kv_scale_quant_orig[1]) - V[t] = bf16(e2m1(v_cache) * e4m3(v_scale_cache) * kv_scale_quant_orig[2]) - s_j = (q[r, h] . K[t]) * softmax_scale - out[r, h] = sum_j softmax(s)_j * V[t] + 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]`. - Each value is dequantized in fp32 and rounded once to `q`'s dtype; at unit - second-level scales every NVFP4 value is exact in bf16, so a call is the - bf16 sparse GQA over the bf16 values the pool holds. After the dequantize the - kernels are QSA's bf16 sparse GQA kernels operation for operation: the same - tiles, online softmax and split-K merge. + 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 @@ -883,16 +950,24 @@ class QsaPagedSparseGqaNvfp4(OpWrapper): prefill pair on every step, the decode forward never does; both pass `q` as the 24-head slice of the packed QKV row (row stride 7168). - Launch shapes at Hkv 2 (programs = rows x 2), as on the bf16 entry: split-K - 64-way up to 2 rows, 32-way up to 15 (16-token tiles), 8-way up to 128 and - 4-way up to 256 (64-token tiles), each followed by the merge; past 256 rows - the prefill kernel when the pair is passed, else one split and no merge. - The decode site reaches 256 rows at max_draft_len 7 (32 requests x 8). + 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 but the captured one the two - outputs agree within the test's band. They are not bitwise equal: the - dequantizing load changes the compiled kernel around the same arithmetic. + 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 @@ -912,6 +987,9 @@ class QsaPagedSparseGqaNvfp4(OpWrapper): - 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 == 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 new file mode 100644 index 000000000000..150773d62e54 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4_decode.py @@ -0,0 +1,1878 @@ +# 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`. + +A one-split launch (65 rows and up) skips the f16 V tile: each warp builds its P V B +fragments from the raw V rows in registers (the f16 tile's stores and ldmatrix reads, 64 KB a +tile, bound the standard kernel there), bitwise the same. + +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 +#: One-split launches (from 65 rows on, `_SPLIT_TABLE`) run the lean split kernel: V's B +#: fragments dequantized from the raw tile in registers, no f16 V tile (its 64 KB a tile of +#: shared-memory stores and ldmatrix reads is what bounds the standard kernel at many rows). +#: Raw V rows and V-scale groups of the lean kernel, padded so that the four tokens a load +#: instruction touches (two apart) fall in different banks. +LEAN_RAW_LDV = 144 +LEAN_VS_LD = 80 + + +# --------------------------------------------------------------------------------------- +# 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, + lean=False, + ): + assert group <= 16 and HEAD_DIM == 256 and tpb % 4 == 0 + assert not lean or splits == 1, "the lean kernel runs one split" + self.lean = lean + 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 = LEAN_RAW_LDV if lean else RAW_LDV + self.raw_vsld = LEAN_VS_LD if lean else 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 + if lean: + # tokens | raw tiles | P | reductions + self.o_sp = self.o_raw + self.stages * self.raw_buf + 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,))) + kernel = self.kernel_lean if self.lean else self.kernel + 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), + ) + + @cute.kernel + def kernel_lean( + 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, + ): + """`kernel` for one split without the f16 V tile: the same MMAs on the same values. + + P V takes V's B fragments straight from the raw tile: lane (g, c) of warp w serves the + four head dims 32w + 4g .. 32w + 4g + 3 (one byte pair of each token, one scale block) + as n-index g of four n-tiles, n-tile j holding dim 32w + 4n + j. An output element sees + the same P row and V column in the same token order as in `kernel`, so the result is + bitwise its. The raw tile is read until the end of the tile's P V, so the next tile is + issued after barrier [A] (every warp's P V of the tile its buffer held is done).""" + 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 + 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 + stages: cutlass.Constexpr = self.stages + row = bidx // kvh_n + kvh = bidx % kvh_n + n_t = cutlass.min(tps, n_tiles_all) + + 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_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: 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 + colc = cutlass.min(cl, width - 1) + tok = mSel[row, colc] + ok = req_ok & (cl < 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 -- + 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)) + ) + + 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 + st[1] = neg_inf + st[2] = cutlass.Float32(0.0) + st[3] = cutlass.Float32(0.0) + 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] + + 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) + red_g = s_red + g * WARPS * 4 + red_h = s_red + (g + 8) * WARPS * 4 + # P V: dims 32w + 4g .. + 3 are byte pair 16w + 2g of a token's V row, scale block + # 2w + g // 4 of its four-token scale group. + v_pair = 16 * warp + 2 * g + v_blk = (2 * warp + g // 4) * 4 + lo_sel = cutlass.Int32(0x5410) + hi_sel = cutlass.Int32(0x7632) + ldv: cutlass.Constexpr = self.raw_ldv + vsld: cutlass.Constexpr = self.raw_vsld + + ahead: cutlass.Constexpr = stages - 1 + for j in cutlass.range(n_t): + buf = j % stages + 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; every warp's P V of tile j - 1 done + if j + ahead < n_t: + # into the buffer of tile j - 1 + self._issue( + j + ahead, + (j + ahead) % stages, + tidx, + kvh, + s_tok, + s_raw, + k_base, + v_base, + ks_base, + vs_base, + ) + rb = s_raw + buf * self.raw_buf + rk = rb + rv = rb + self.raw_k + rks = rv + self.raw_v + rvs = rks + self.raw_ks + + # -- S = q K^T for this warp's eight tokens (two accumulator chains), q per k-step + 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): + 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) + + 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 + 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: n-tile j of this warp holds dims 32w + 4n + j + 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) + # tokens t, t + 1 (one scale group: their two scale bytes are adjacent) and + # t + 8, t + 9; e2m1 pairs unscaled, each token's scale applied after the + # pairing (an exact product either way) + lo = [] + hi = [] + scl = [] + for dt in (0, 8): + t0 = 16 * kk + 2 * c + dt + scl.append(_e4m3x2(_lds_u16(rvs + t0 * vsld + v_blk + (t0 & 3)))) + for d1 in (0, 1): + ee = _e2m1_half(_lds_u16(rv + (t0 + d1) * ldv + v_pair)) + lo.append(ee[0]) + hi.append(ee[1]) + for nt in cutlass.range_constexpr(4): + src = lo if nt < 2 else hi + sel = lo_sel if nt % 2 == 0 else hi_sel + b0 = _hmul2(_prmt(src[0], src[1], sel), scl[0]) + b1 = _hmul2(_prmt(src[2], src[3], sel), scl[1]) + r = _mma(a, (b0, b1), (acc[nt, 0], acc[nt, 1], acc[nt, 2], acc[nt, 3])) + for e in cutlass.range_constexpr(4): + acc[nt, e] = r[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]) + + # ---- epilogue: this lane's dims 32w + 8c .. 32w + 8c + 7 of heads g and g + 8 -------- + gv = mGs[2] + ig = cutlass.Float32(0.0) + ih = cutlass.Float32(0.0) + if tsg > 0.0: + ig = gv / tsg + if tsh > 0.0: + ih = gv / tsh + d0 = 32 * warp + 8 * c + for half in cutlass.range_constexpr(2): + h = kvh * group + g + 8 * half + inv = ig if half == 0 else ih + if (g + 8 * half) < group: + # columns 2c, 2c + 1 of n-tile j: dims d0 + j and d0 + 4 + j + o = [acc[j, 2 * half + k] * inv for k in range(2) for j in range(4)] + if const_expr(self.gated): + gw = _ldg4( + gate_base + (cutlass.Int64(row) * gate_stride + h * HEAD_DIM + d0) * 2 + ) + for e in cutlass.range_constexpr(4): + o[2 * e] = _round_bf16(o[2 * e]) * _sigmoid(_bf2_lo(gw[e])) + o[2 * e + 1] = _round_bf16(o[2 * e + 1]) * _sigmoid(_bf2_hi(gw[e])) + _stg4( + out_base + (cutlass.Int64(row) * heads + h) * (HEAD_DIM * 2) + d0 * 2, + _bf2_pack(o[0], o[1]), + _bf2_pack(o[2], o[3]), + _bf2_pack(o[4], o[5]), + _bf2_pack(o[6], o[7]), + ) + + +@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, + lean: Optional[bool] = 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()) + lean = bool((True if lean is None else lean) and n_splits == 1) + 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, + lean, + ) + 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, + lean=lean, + ) + 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 new file mode 100644 index 000000000000..a71036db0466 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4_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.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_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 new file mode 100644 index 000000000000..144df3bd24fc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rollback_planes.py @@ -0,0 +1,144 @@ +# 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_rowtile_attention_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rowtile_attention_cute.py index 03e2c925518e..1a9cac72e24e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rowtile_attention_cute.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rowtile_attention_cute.py @@ -7,19 +7,30 @@ (`packed_bits`, as int32 halves) and `group_plan[group] = (tiles, listed)`; groups with `tiles < 0` belong to the lone-row pass. -One CTA per (group, KV head), 128 MMA rows = `rows_per_tile` query rows x `group_size` -heads, 9 warps: - -* warps 5-8 gather the list's K and V rows (64 a tile) with `cp.async` into one ring of - `kv_slots` 32 KB SW128 shared-memory slots, filled in the MMA warp's order K0, K1, V0, - K2, V1, ...; each slot completes on its mbarrier through `cp.async.mbarrier.arrive`; -* warp 4 issues the tcgen05 MMAs: S[j % 2] = Q K(j)^T (Q from tensor memory), then - O += P(j - 1) V(j - 1) (P from tensor memory, overwriting the S it came from); -* warps 0-3, one thread per MMA row, first copy their Q row into tensor memory, then per - tile load S with its row max (`tcgen05.ld.red`), rescale O when that max grew by more - than `rescale_threshold` in log2 units (lazily; otherwise exponentiate against the stale - max), write P = 2^(s - max) on the row's own bits and 0 elsewhere, and at the end - normalize O and store it through the idle ring as whole 512-byte rows. +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 @@ -55,6 +66,12 @@ #: 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 @@ -76,7 +93,7 @@ def __init__( rows_per_tile: int, head_dim: int = 256, tokens_per_block: int = 32, - kv_slots: int = 6, + kv_slots: int = KV_SLOTS, rescale_threshold: float = 8.0, use_pdl: bool = True, ): @@ -91,6 +108,8 @@ def __init__( 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 @@ -112,9 +131,11 @@ def __call__( 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, @@ -173,11 +194,36 @@ class SharedStorage: 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, @@ -187,6 +233,8 @@ class SharedStorage: k_layout, v_layout, p_layout, + tma_atom_o, + tma_tensor_o, q, k, v, @@ -194,7 +242,10 @@ class SharedStorage: bits, plan, out, + work, rows, + num_groups, + num_kv_heads, q_stride_row, k_stride_page, k_stride_head, @@ -207,7 +258,7 @@ class SharedStorage: out_stride_row, scale_log2, ).launch( - grid=[num_groups, num_kv_heads, 1], + grid=[num_ctas, 1, 1], block=[self.threads, 1, 1], stream=stream, min_blocks_per_mp=1, @@ -225,6 +276,8 @@ def kernel( 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, @@ -232,7 +285,10 @@ def kernel( 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, @@ -247,141 +303,218 @@ def kernel( ): warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) tidx, _, _ = cute.arch.thread_idx() - group, kv_head, _ = cute.arch.block_idx() qk_mma.set(tcgen05.Field.ACCUMULATE, False) qk_mma_acc.set(tcgen05.Field.ACCUMULATE, True) pv_mma.set(tcgen05.Field.ACCUMULATE, False) pv_mma_acc.set(tcgen05.Field.ACCUMULATE, True) - mPlan = cute.make_tensor(plan, cute.make_layout((2 * 1048576,))) - tiles = mPlan[group * 2] - listed = mPlan[group * 2 + 1] - - if tiles >= 0: - smem = SmemAllocator() - storage = smem.allocate(self.shared_storage) - q_full = storage.q_full.data_ptr() - kv_full = storage.kv_full.data_ptr() - kv_empty = storage.kv_empty.data_ptr() - s_full = storage.s_full.data_ptr() - p_full = storage.p_full.data_ptr() - o_full = storage.o_full.data_ptr() - o_final = storage.o_final.data_ptr() - if tidx == 0: - cute.arch.mbarrier_init(q_full, 128) - for i in cutlass.range_constexpr(self.kv_slots): - cute.arch.mbarrier_init(kv_full + i, 128) - cute.arch.mbarrier_init(kv_empty + i, 1) - for i in cutlass.range_constexpr(2): - cute.arch.mbarrier_init(s_full + i, 1) - cute.arch.mbarrier_init(p_full + i, 128) - cute.arch.mbarrier_init(o_full, 1) - cute.arch.mbarrier_init(o_final, 1) - cute.arch.mbarrier_init_fence() - - tmem_users = pipeline.NamedBarrier(barrier_id=1, num_threads=32 * 5) - if warp_idx == self.mma_warp: - cute.arch.alloc_tmem(self.tmem_cols, storage.tmem_holding.ptr) - cute.arch.relinquish_tmem_alloc_permit() - cute.arch.sync_threads() - - # Slices taken before the role branches: the MMA warp's `set(ACCUMULATE)` rebinds the - # tiled MMAs inside its branch, which the other branches must not see. - qk_thr = qk_mma.get_slice(0) - pv_thr = pv_mma.get_slice(0) - sK = storage.sKV.get_tensor(k_layout.outer, swizzle=k_layout.inner) - sV = cute.make_tensor(cute.recast_ptr(sK.iterator, v_layout.inner), v_layout.outer) - - if warp_idx >= self.load_warp0: - self.load_role( - tidx - 32 * self.load_warp0, - group, - kv_head, - tiles, - listed, - sK, - k, - v, - slots, - k_stride_page, - k_stride_head, - k_stride_token, - v_stride_page, - v_stride_head, - v_stride_token, - slots_stride_group, - kv_full, - kv_empty, - ) - elif warp_idx == self.mma_warp: - tmem_ptr = cute.arch.retrieve_tmem_ptr(Float32, 16, storage.tmem_holding.ptr) - self.mma_role( - qk_mma, - qk_mma_acc, - pv_mma, - pv_mma_acc, - qk_thr, - pv_thr, - q_layout, - p_layout, - sK, - sV, - tmem_ptr, - tiles, - q_full, - kv_full, - kv_empty, - s_full, - p_full, - o_full, - o_final, - ) - tmem_users.arrive_and_wait() - cute.arch.dealloc_tmem(tmem_ptr, self.tmem_cols) - else: - tmem_ptr = cute.arch.retrieve_tmem_ptr(Float32, 16, storage.tmem_holding.ptr) - self.softmax_role( - qk_thr, - pv_thr, - tmem_ptr, - tidx, - group, - kv_head, - tiles, - rows, - bits, - bits_stride_row, - out, - out_stride_row, - scale_log2, - q_full, - s_full, - p_full, - o_full, - o_final, - storage.sKV.data_ptr(), - q, - q_stride_row, - ) - tmem_users.arrive_and_wait() + 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, - group: Int32, - kv_head: Int32, - tiles: Int32, - listed: 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, @@ -391,11 +524,18 @@ def load_role( slots_stride_group: Int32, kv_full, kv_empty, + qs_full, + qs_empty, + sch_full, + sch_empty, + sch_val, ): - """Fill the K/V ring in the MMA warp's order K0, K1, V0, K2, V1, ..., V(T-1). Thread `lt` - gathers tokens 16 (lt // 32) + 4 t + (lt % 32) // 8 (t < 4), one 16-byte chunk (lt % 8) of - each of their four 128-byte dim blocks; the slots of tile j + 1 are fetched while tile j's - copies issue.""" + """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 @@ -409,58 +549,156 @@ def load_role( stride=(1, (8, self.bn * 64), 64, self.bn * D), ), ) - if tiles > 0: - mSlots = cute.make_tensor( - slots + Int64(group) * slots_stride_group, cute.make_layout((1 << 30,)) + 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 ) - first_token = 16 * (lt // 32) + (lt % 32) // 8 - slot_next = cute.make_rmem_tensor((4,), Int32) - k_cur = cute.make_rmem_tensor((4,), Int64) - v_cur = cute.make_rmem_tensor((4,), Int64) - v_prev = cute.make_rmem_tensor((4,), Int64) - for t in cutlass.range_constexpr(4): - slot_next[t] = mSlots[cutlass.min(first_token + 4 * t, listed - 1)] - ring = Int32(0) - ring_phase = Int32(1) - for j in cutlass.range(0, tiles, 1, unroll=1): + 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 = 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: + 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_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) + 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, v, v_prev, sKV_c, ring, chunk, first_token) + 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) - for t in cutlass.range_constexpr(4): - v_prev[t] = v_cur[t] - # V(tiles - 1) - cute.arch.mbarrier_wait(kv_empty + ring, ring_phase) - self.issue_tile(atom, v, v_prev, sKV_c, ring, chunk, first_token) - cute.arch.cp_async_mbarrier_arrive_noinc(kv_full + ring) - ring, ring_phase = self.advance(ring, ring_phase) - # Tail: every slot released, so no load is in flight when the CTA retires and the ring - # is free for the epilogue's staging. - for i in cutlass.range_constexpr(self.kv_slots): + 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( @@ -496,7 +734,8 @@ def mma_role( sK, sV, tmem_ptr, - tiles: Int32, + plan: cute.Pointer, + num_kv_heads: Int32, q_full, kv_full, kv_empty, @@ -504,10 +743,16 @@ def mma_role( p_full, o_full, o_final, + sch_full, + sch_empty, + sch_val, ): - """QK(0), then QK(j) / PV(j - 1) pairs, then the last PV. S/P buffer b sits at TMEM columns - tmem_s0 + 64 b and is picked by pointer offset; QK(j) overwrites S[j % 2] only after PV(j - 2) - was issued (tcgen05 MMAs of one thread execute in order), and P(j - 2) was its last reader.""" + """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) @@ -518,62 +763,79 @@ def mma_role( tOtO = cute.make_tensor(tmem_ptr + self.tmem_o, tOtO_f.layout) tP = cute.make_tensor(tmem_ptr + self.tmem_s0, p_layout.outer) tOrP = pv_thr.make_fragment_A(tP)[None, None, None, 0] - if tiles > 0: - cute.arch.mbarrier_wait(q_full, 0) - tcgen05_fence_after() - ring = Int32(0) - ring_phase = Int32(0) - # QK(0) - cute.arch.mbarrier_wait(kv_full + ring, ring_phase) - self.gemm_qk( - qk_mma, - qk_mma_acc, - tSrQ, - tSrK[None, None, None, ring], - cute.make_tensor(tmem_ptr + self.tmem_s0, tStS.layout), - ) - with cute.arch.elect_one(): - tcgen05.commit(s_full) - tcgen05.commit(kv_empty + ring) - ring, ring_phase = self.advance(ring, ring_phase) - for j in cutlass.range(1, tiles + 1, 1, unroll=1): - buf = j % 2 - if j < tiles: - # QK(j) into S[j % 2] + 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) - 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), + 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(s_full + buf) tcgen05.commit(kv_empty + ring) + tcgen05.commit(o_full) ring, ring_phase = self.advance(ring, ring_phase) - # PV(j - 1) - prev = j - 1 - pbuf = prev % 2 - cute.arch.mbarrier_wait(kv_full + ring, ring_phase) - cute.arch.mbarrier_wait(p_full + pbuf, (prev // 2) % 2) - tcgen05_fence_after() - tOrP_b = cute.make_tensor( - cute.recast_ptr( - tmem_ptr + (self.tmem_s0 + pbuf * self.bn), dtype=tOrP.element_type - ), - tOrP.layout, - ) - if prev == 0: - self.gemm_pv(pv_mma, pv_mma_acc, tOrP_b, tOrV[None, None, None, ring], tOtO) - else: - self.gemm_pv(pv_mma_acc, pv_mma_acc, tOrP_b, tOrV[None, None, None, ring], tOtO) with cute.arch.elect_one(): - tcgen05.commit(kv_empty + ring) - tcgen05.commit(o_full) - ring, ring_phase = self.advance(ring, ring_phase) - with cute.arch.elect_one(): - tcgen05.commit(o_final) + 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): @@ -600,6 +862,38 @@ def gemm_pv(self, first: cute.TiledMma, rest: cute.TiledMma, tOrP, tOrV, tOtO): 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, @@ -607,9 +901,8 @@ def softmax_role( pv_thr, tmem_ptr, st: Int32, - group: Int32, - kv_head: Int32, - tiles: Int32, + plan: cute.Pointer, + num_kv_heads: Int32, rows: Int32, bits: cute.Pointer, bits_stride_row: Int32, @@ -621,9 +914,15 @@ def softmax_role( p_full, o_full, o_final, - stage, - q: cute.Pointer, - q_stride_row: Int32, + 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 @@ -632,8 +931,11 @@ def softmax_role( 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. + # 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 @@ -663,7 +965,6 @@ def softmax_role( 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) @@ -694,185 +995,254 @@ def softmax_role( 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 - row = group * R + local_row - live = (local_row < R) and (row < rows) - safe_row = group * R - if live: - safe_row = row - mBits = cute.make_tensor( - bits + Int64(safe_row) * bits_stride_row, cute.make_layout((1 << 30,)) - ) - - # Q row -> TMEM columns tmem_q.. (the QK MMA's A operand), 64 elements a store. Dead rows - # store zeros; their S rows are masked out. - if tiles > 0: - q_row = cute.make_tensor( - q + cute.assume(Int64(safe_row) * q_stride_row + (kv_head * G + head) * D, divby=8), - cute.make_layout((D,)), - ) - for c in cutlass.range_constexpr(D // (2 * p_cols)): - rQ = cute.make_rmem_tensor(tPcP.shape, Float32) - rQ16 = cute.make_tensor( - cute.recast_ptr(rQ.iterator, dtype=BF16), cute.make_layout((2 * p_cols,)) + 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: - for i in cutlass.range_constexpr(2 * p_cols // 8): - cute.autovec_copy( - cute.make_tensor( - q_row.iterator + (c * 2 * p_cols + i * 8), cute.make_layout((8,)) - ), - cute.make_tensor(rQ16.iterator + i * 8, cute.make_layout((8,))), + 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]) ) - else: - rQ.fill(0.0) - cute.copy( - store_p, - rQ, - cute.make_tensor( - tPtP0.iterator + (self.tmem_q - self.tmem_s0 + c * p_cols), tPtP0.layout - ), - ) - cute.arch.fence_view_async_tmem_store() - tcgen05_fence_before() - cute.arch.mbarrier_arrive(q_full) - - row_max = -Float32.inf - row_sum = Float32(0.0) - lo_next = Int32(0) - hi_next = Int32(0) - if live and tiles > 0: - lo_next = mBits[0] - hi_next = mBits[1] - tau = Float32(self.tau) - for j in cutlass.range(0, tiles, 1, unroll=1): - lo = lo_next - hi = hi_next - if live and j + 1 < tiles: - lo_next = mBits[2 * j + 2] - hi_next = mBits[2 * j + 3] - buf = j % 2 - cute.arch.mbarrier_wait(s_full + buf, (j // 2) % 2) - tcgen05_fence_after() - rS = cute.make_rmem_tensor(tLcS.shape, Float32) - n_ld = cute.size(rS, mode=[1]) - rMax = cute.make_rmem_tensor(cute.make_layout((1, n_ld)), Float32) - tLtS_b = cute.make_tensor(tLtS0.iterator + buf * self.bn, tLtS0.layout) - for c in cutlass.range_constexpr(n_ld): - cute.copy_atom_call( - load_s, tLtS_b[None, c, 0, 0], (rS[None, c, 0, 0], rMax[None, c]) - ) - cute.arch.fence_view_async_tmem_load() - tile_max = rMax[0] - for c in cutlass.range_constexpr(1, n_ld): - tile_max = cute.arch.fmax(tile_max, rMax[c]) - new_max = cute.arch.fmax(row_max, tile_max) - use_max = row_max - factor = Float32(1.0) - if row_max == -Float32.inf: - use_max = new_max - elif (new_max - row_max) * scale_log2 > tau: - use_max = new_max - factor = cute.math.exp2((row_max - new_max) * scale_log2, fastmath=True) - need = factor != Float32(1.0) - if cute.arch.vote_any_sync(need): - # O holds PV(0..j-1): wait for PV(j-1), then rescale this warp's rows. - cute.arch.mbarrier_wait(o_full, (j - 1) % 2) - tcgen05_fence_after() - for c in cutlass.range_constexpr(n_chunks): - rO = cute.make_rmem_tensor(tLcO.shape, Float32) - cute.copy( - load_o, - cute.make_tensor(tLtO.iterator + c * self.chunk_cols, tLtO.layout), - rO, - ) - for i in cutlass.range_constexpr(cute.size(rO)): - rO[i] = rO[i] * factor + cute.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_o, - rO, - cute.make_tensor(tStO.iterator + c * self.chunk_cols, tStO.layout), + store_p, rP, cute.make_tensor(tPtP0.iterator + buf * self.bn, tPtP0.layout) ) - cute.arch.fence_view_async_tmem_store() - row_max = use_max - neg = -use_max * scale_log2 - rP = cute.make_rmem_tensor(tPcP.shape, Float32) - rP16 = cute.make_tensor(cute.recast_ptr(rP.iterator, dtype=BF16), rS.layout) - # p = 2^(s * scale - max * scale) on the row's own entries, 0 elsewhere; four partial sums. - acc0 = (Float32(0.0), Float32(0.0)) - acc1 = (Float32(0.0), Float32(0.0)) - for i in cutlass.range_constexpr(0, self.bn, 4): - word = lo if i < 32 else hi - e = cute.make_rmem_tensor((4,), Float32) - e[0], e[1] = cute.arch.fma_packed_f32x2( - (rS[i], rS[i + 1]), (scale_log2, scale_log2), (neg, neg) - ) - e[2], e[3] = cute.arch.fma_packed_f32x2( - (rS[i + 2], rS[i + 3]), (scale_log2, scale_log2), (neg, neg) - ) - for u in cutlass.range_constexpr(4): - x = cute.math.exp2(e[u], fastmath=True) - rS[i + u] = x if (word & (1 << ((i + u) % 32))) != 0 else Float32(0.0) - acc0 = cute.arch.add_packed_f32x2(acc0, (rS[i], rS[i + 1])) - acc1 = cute.arch.add_packed_f32x2(acc1, (rS[i + 2], rS[i + 3])) - row_sum = row_sum * factor + ((acc0[0] + acc0[1]) + (acc1[0] + acc1[1])) - rP16.store(rS.load().to(BF16)) - cute.copy(store_p, rP, cute.make_tensor(tPtP0.iterator + buf * self.bn, tPtP0.layout)) - cute.arch.fence_view_async_tmem_store() - tcgen05_fence_before() - cute.arch.mbarrier_arrive(p_full + buf) - - # Epilogue: O / l in bf16, staged through the (now idle) K/V ring -- a row's 32 16-byte - # chunks XOR-swizzled by the row -- so each warp then stores whole 512-byte rows. - if tiles > 0: - cute.arch.mbarrier_wait(o_final, 0) - tcgen05_fence_after() - inv = Float32(0.0) - if row_sum > Float32(0.0): - inv = Float32(1.0) / row_sum - swz = m % 32 - for c in cutlass.range_constexpr(n_chunks): - rO = cute.make_rmem_tensor(tLcO.shape, Float32) - if tiles > 0: - cute.copy( - load_o, cute.make_tensor(tLtO.iterator + c * self.chunk_cols, tLtO.layout), rO - ) - for i in cutlass.range_constexpr(cute.size(rO)): - rO[i] = rO[i] * inv - else: - rO.fill(0.0) - rB = cute.make_rmem_tensor((self.chunk_cols,), BF16) - rB.store(rO.load().to(BF16)) - for i in cutlass.range_constexpr(self.chunk_cols // 8): - logical = c * (self.chunk_cols // 8) + i - offset = cute.assume(m * D + ((logical ^ swz) * 8), divby=8) - dst = cute.make_tensor(stage + offset, cute.make_layout((8,))) - cute.autovec_copy( - cute.make_tensor(rB.iterator + i * 8, cute.make_layout((8,))), dst - ) - cute.arch.sync_warp() - lane = st % 32 - warp = st // 32 - for r in cutlass.range_constexpr(32): - mr = warp * 32 + r - lr = mr // G - row_r = group * R + lr - if lr < R and row_r < rows: - src_off = cute.assume(mr * D + ((lane ^ r) * 8), divby=8) - rV = cute.make_rmem_tensor((8,), BF16) - cute.autovec_copy(cute.make_tensor(stage + src_off, cute.make_layout((8,))), rV) - dst_off = cute.assume( - Int64(row_r) * out_stride_row + (kv_head * G + mr % G) * D + lane * 8, divby=8 - ) - cute.autovec_copy(rV, cute.make_tensor(out + dst_off, cute.make_layout((8,)))) + 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, use_pdl=use_pdl + group_size, rows_per_tile, HEAD_DIM, tokens_per_block, kv_slots=KV_SLOTS, use_pdl=use_pdl ) gmem = cute.AddressSpace.gmem @@ -887,6 +1257,8 @@ def fake(dtype, align): fake(Int32, 4), fake(Int32, 4), fake(BF16, 16), + fake(Int32, 4), + 1, 1, 1, 1, @@ -906,13 +1278,24 @@ def fake(dtype, align): 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.""" + 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): @@ -931,7 +1314,8 @@ def eligible(q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor) -> b class _Launcher: """`rowtile_attention_cute[grid](...)`: grid-indexed like the entry's Triton launches; - the grid is `(groups, kv_heads)`.""" + the grid is `(ctas,)`, the persistent CTAs (`num_ctas`). + """ def __getitem__(self, grid): return functools.partial(self._launch, tuple(grid)) @@ -946,6 +1330,7 @@ def _launch( packed_bits, group_plan, output, + work, softmax_scale, rows, *, @@ -954,11 +1339,12 @@ def _launch( tokens_per_block, use_pdl, ): - assert grid == (group_plan.shape[0], k_cache.shape[1]), grid + 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 @@ -974,9 +1360,11 @@ def ptr(t, dtype, align): 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], - grid[1], q.stride(0), k_cache.stride(0), k_cache.stride(1), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py index 86a4f4a92b8f..27a191801ec4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py @@ -6,15 +6,17 @@ lets it read) into a bitmap over logical token positions. A plan launch forms, per group of consecutive rows of one request, the union of its rows' resident selected tokens as a packed list of physical cache slots, and each row's -bitmap re-indexed onto that list. The row-tile launch runs one program per -(group, KV head) that walks the packed list in 64-slot tiles, each tile's K/V -loaded once for every row and head of the group, each row masked to its own -bits. A group whose rows span requests (generation rows of a mixed batch) -walks each row's own selected columns instead, as the per-row kernel does. +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 functools import math import cutlass @@ -33,6 +35,7 @@ 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 _INTEGER_DTYPES = (torch.int32, torch.int64) @@ -229,6 +232,335 @@ def _row_bits( return bits, row_info +# -------------------------------------------------------------------------------------- +# The plan pass (CuTe DSL). One warp per group of `rows_per_tile` rows, a lane per bitmap +# word: the group's union word (OR of its rows' words), its resident part (tokens on pages +# the table maps inside the pool), a warp scan of the words' token counts for their list +# offsets, the list's slots staged in shared memory and stored by the whole warp, and each +# row's word compressed onto the resident bits (Hacker's Delight `compress`, the masks computed +# once a word for every row) and OR-ed into the row's packed words in shared memory +# (`red.shared.or.b64`), stored once the list is complete. Same plan, slot list, packed words +# and lone list as the Triton kernel below, which serves int64 request indices (and tables +# too long for a CTA's shared memory). +# -------------------------------------------------------------------------------------- + +#: Groups (one a warp) a plan CTA takes. +_PLAN_GROUPS_PER_CTA = 4 +#: Shared memory the CuTe plan may take a CTA; past it (long tables, small groups) the Triton +#: kernel plans. +_PLAN_SMEM_LIMIT = 200 * 1024 + + +def _plan_smem_bytes(rows_per_tile: int, num_words: int) -> int: + """A plan CTA's shared memory: per warp, its rows' packed words and a step's staged slots.""" + return _PLAN_GROUPS_PER_CTA * (rows_per_tile * num_words * 8 + 32 * 64 * 4) + 256 + + +class _RowtilePlanKernel: + """The plan pass's slots, packed bitmaps, group plan and lone list, a warp a group.""" + + def __init__(self, rows_per_tile, tokens_per_block, table_width, num_words, shared_tile_cost): + assert 64 % tokens_per_block == 0 or tokens_per_block % 64 == 0 + self.R = rows_per_tile + self.tpb = tokens_per_block + self.width = table_width + self.nw = num_words + self.cost = shared_tile_cost + self.warps = _PLAN_GROUPS_PER_CTA + # Pages a 64-token word spans, and the tokens of the word on each. + self.ppw = max(1, 64 // tokens_per_block) + self.seg = min(64, tokens_per_block) + + @cute.jit + def __call__( + self, + req_t: cute.Tensor, + info_ptr: cute.Pointer, + bits_ptr: cute.Pointer, + table_ptr: cute.Pointer, + slots_ptr: cute.Pointer, + packed_ptr: cute.Pointer, + plan_ptr: cute.Pointer, + lone_rows_ptr: cute.Pointer, + lone_count_ptr: cute.Pointer, + rows: cutlass.Int32, + num_groups: cutlass.Int32, + num_requests: cutlass.Int32, + num_cache_pages: cutlass.Int32, + table_stride: cutlass.Int32, + slots_stride: cutlass.Int32, + bits_stride: cutlass.Int32, + stream, + ): + ctas = (num_groups + self.warps - 1) // self.warps + # `req_t` arrives as a tensor: the TVM-FFI entry takes its stream from a GPU tensor. + self.kernel( + req_t, + info_ptr, + bits_ptr, + table_ptr, + slots_ptr, + packed_ptr, + plan_ptr, + lone_rows_ptr, + lone_count_ptr, + rows, + num_groups, + num_requests, + num_cache_pages, + table_stride, + slots_stride, + bits_stride, + ).launch( + grid=[ctas, 1, 1], + block=[32 * self.warps, 1, 1], + smem=_plan_smem_bytes(self.R, self.nw), + stream=stream, + ) + + @cute.jit + def compress_masks(self, m): + """The six move masks of `compress` for mask `m` (as `_compress_bits`).""" + one = cutlass.Uint64(1) + mk = (~m) << one + mvs = [] + for i in cutlass.range_constexpr(6): + mp = mk ^ (mk << one) + mp = mp ^ (mp << cutlass.Uint64(2)) + mp = mp ^ (mp << cutlass.Uint64(4)) + mp = mp ^ (mp << cutlass.Uint64(8)) + mp = mp ^ (mp << cutlass.Uint64(16)) + mp = mp ^ (mp << cutlass.Uint64(32)) + mv = mp & m + m = (m ^ mv) | (mv >> cutlass.Uint64(1 << i)) + mk = mk & (~mp) + mvs.append(mv) + return mvs + + @cute.kernel + def kernel( + self, + mReq: cute.Tensor, + info_ptr: cute.Pointer, + bits_ptr: cute.Pointer, + table_ptr: cute.Pointer, + slots_ptr: cute.Pointer, + packed_ptr: cute.Pointer, + plan_ptr: cute.Pointer, + lone_rows_ptr: cute.Pointer, + lone_count_ptr: cute.Pointer, + rows: cutlass.Int32, + num_groups: cutlass.Int32, + num_requests: cutlass.Int32, + num_cache_pages: cutlass.Int32, + table_stride: cutlass.Int32, + slots_stride: cutlass.Int32, + bits_stride: cutlass.Int32, + ): + R = self.R + NW = self.nw + I32 = cutlass.Int32 + U64 = cutlass.Uint64 + big = cute.make_layout((1 << 30,)) + tidx, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(tidx // 32) + lane = tidx % 32 + group = bid * self.warps + warp + mInfo = cute.make_tensor(info_ptr, big) + mBits = cute.make_tensor(bits_ptr, big) + mTable = cute.make_tensor(table_ptr, big) + mSlots = cute.make_tensor(slots_ptr, big) + mPacked = cute.make_tensor(packed_ptr, big) + mPlan = cute.make_tensor(plan_ptr, big) + mLone = cute.make_tensor(lone_rows_ptr, big) + mLoneCount = cute.make_tensor(lone_count_ptr, cute.make_layout((1,))) + + smem = cutlass.utils.SmemAllocator() + buf = smem.allocate_tensor( + U64, cute.make_layout((self.warps * R * NW,), stride=(1,)), byte_alignment=16 + ) + buf_base = I32(buf.iterator.toint()) + # A step's slots (at most 32 words x 64 tokens), staged for coalesced stores. + sstage = smem.allocate_tensor( + I32, cute.make_layout((self.warps * 32 * 64,), stride=(1,)), byte_alignment=16 + ) + sbase = warp * (32 * 64) + wbase = warp * (R * NW) + for i in cutlass.range(lane, R * NW, 32, unroll=1): + buf[wbase + i] = U64(0) + cute.arch.sync_warp() + + if group < num_groups: + row = group * R + lane + live = lane < R and row < rows + req = I32(-1) + hi_tok = I32(-1) + own = I32(0) + if live: + req = I32(mReq[row]) + hi_tok = mInfo[row * 2] + own = (mInfo[row * 2 + 1] + 63) // 64 + lowest = req if live else num_requests + highest = req + for o in cutlass.range_constexpr(5): + lowest = cutlass.min(lowest, cute.arch.shuffle_sync_bfly(lowest, offset=1 << o)) + highest = cutlass.max(highest, cute.arch.shuffle_sync_bfly(highest, offset=1 << o)) + hi_tok = cutlass.max(hi_tok, cute.arch.shuffle_sync_bfly(hi_tok, offset=1 << o)) + own = own + cute.arch.shuffle_sync_bfly(own, offset=1 << o) + one_request = lowest == highest and lowest >= 0 and lowest < num_requests + words_needed = I32(0) + if one_request: + words_needed = (hi_tok + 64) // 64 + req_c = cutlass.max(cutlass.min(lowest, num_requests - 1), I32(0)) + listed = I32(0) + for base in cutlass.range(0, words_needed, 32, unroll=1): + w = base + lane + inr = w < words_needed + rw = [] + union = U64(0) + for r in cutlass.range_constexpr(R): + x = U64(0) + if inr and group * R + r < rows: + x = U64(mBits[(group * R + r) * bits_stride + w]) + rw.append(x) + union = union | x + # Resident tokens: the word's pages mapped inside the pool. + resident = U64(0) + pages = [] + for pp in cutlass.range_constexpr(self.ppw): + lp = (w * 64 + pp * self.seg) // self.tpb + pid = I32(-1) + if inr and lp < self.width: + pid = mTable[req_c * table_stride + lp] + pages.append(pid) + seg_bits = ( + U64(0xFFFFFFFFFFFFFFFF) + if self.seg == 64 + else U64(((1 << self.seg) - 1) << (pp * self.seg)) + ) + if pid >= 0 and pid < num_cache_pages: + resident = resident | seg_bits + resident = resident & union + count = I32(cute.arch.popc(resident)) + incl = count + for o in cutlass.range_constexpr(5): + up = cute.arch.shuffle_sync_up(incl, offset=1 << o) + if lane >= (1 << o): + incl = incl + up + total = cute.arch.shuffle_sync(incl, 31) + first = listed + incl - count + # The list: resident tokens in logical order, as physical cache slots, staged in + # shared memory and stored by the whole warp. + pos = incl - count + for b in cutlass.range_constexpr(64): + if (resident >> U64(b)) & U64(1) != U64(0): + if cutlass.const_expr(self.tpb <= 64): + tok = b % self.tpb + else: + tok = (w * 64) % self.tpb + b + sstage[sbase + pos] = pages[b // self.seg] * self.tpb + tok + pos = pos + 1 + cute.arch.sync_warp() + slot_row = group * slots_stride + listed + for k in cutlass.range(lane, total, 32, unroll=1): + mSlots[slot_row + k] = sstage[sbase + k] + cute.arch.sync_warp() + # Each row's word compressed onto the resident bits, at the list offset. + mvs = self.compress_masks(resident) + d0 = first // 64 + sh = U64(first % 64) + spills = (first % 64) + count > 64 + if count > 0: + for r in cutlass.range_constexpr(R): + x = rw[r] & resident + for i in cutlass.range_constexpr(6): + moved = x & mvs[i] + x = (x ^ moved) | (moved >> U64(1 << i)) + addr = buf_base + (wbase + r * NW + d0) * 8 + _red_or_shared(addr, cutlass.Int64(x << sh)) + if spills: + _red_or_shared( + addr + 8, + cutlass.Int64((x >> U64(1)) >> (U64(63) - sh)), + ) + listed = listed + total + cute.arch.sync_warp() + tiles = (listed + 63) // 64 + shared = one_request and tiles * self.cost <= own * 4 + if lane == 0: + mPlan[group * 2] = tiles if shared else I32(-1) + mPlan[group * 2 + 1] = listed + for r in cutlass.range_constexpr(R): + rr = group * R + r + if rr < rows: + for k in cutlass.range(lane, tiles, 32, unroll=1): + mPacked[rr * bits_stride + k] = buf[wbase + r * NW + k] + if not shared: + n_live = I32(1) if live else I32(0) + for o in cutlass.range_constexpr(5): + n_live = n_live + cute.arch.shuffle_sync_bfly(n_live, offset=1 << o) + start = I32(0) + if lane == 0: + start = cute.arch.atomic_add(mLoneCount.iterator, n_live) + start = cute.arch.shuffle_sync(start, 0) + if live: + mLone[start + lane] = row + + +_ROW_PLAN: dict = {} + + +def _cute_plan_supported(request_indices, tokens_per_block, rows_per_tile, num_words) -> bool: + """Whether `_RowtilePlanKernel` plans this launch: int32 request indices, pages that tile a + 64-token word (or whole words), and a CTA's packed words within `_PLAN_SMEM_LIMIT`.""" + return ( + request_indices.dtype == torch.int32 + and (64 % tokens_per_block == 0 or tokens_per_block % 64 == 0) + and _plan_smem_bytes(rows_per_tile, num_words) <= _PLAN_SMEM_LIMIT + ) + + +def _rowtile_plan_cute( + request_indices, row_info, bits, block_table, slots, packed_bits, group_plan, lone_rows, + lone_count, rows, num_groups, *, rows_per_tile, tokens_per_block, num_cache_pages, +): # fmt: skip + """The plan pass on `_RowtilePlanKernel` (int32 request indices). Writes each row's packed + words below its group's tile count only.""" + table_width = block_table.shape[1] + num_words = bits.shape[1] + key = (rows_per_tile, tokens_per_block, table_width, num_words) + fn = _ROW_PLAN.get(key) + ptrs = (row_info, bits, block_table, slots, packed_bits, group_plan, lone_rows, lone_count) + args = (request_indices, *(t.data_ptr() for t in ptrs)) + if fn is None: + gm = cute.AddressSpace.gmem + dtypes = (cutlass.Int32, cutlass.Uint64, cutlass.Int32, cutlass.Int32, cutlass.Uint64, + cutlass.Int32, cutlass.Int32, cutlass.Int32) # fmt: skip + aligns = (8, 16, 4, 4, 16, 4, 4, 4) + fn = cute.compile( + _RowtilePlanKernel( + rows_per_tile, tokens_per_block, table_width, num_words, _SHARED_TILE_COST + ), + from_dlpack(request_indices, assumed_align=4).mark_layout_dynamic(leading_dim=0), + *(make_ptr(dt, p, gm, assumed_align=al) for dt, p, al in zip(dtypes, args[1:], aligns)), + cutlass.Int32(1), + cutlass.Int32(1), + cutlass.Int32(1), + cutlass.Int32(1), + 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", + ) + _ROW_PLAN[key] = fn + fn( + *args, rows, num_groups, block_table.shape[0], num_cache_pages, block_table.stride(0), + slots.stride(0), bits.stride(0), + ) # fmt: skip + + @triton.jit def _or_combine(a, b): return a | b @@ -257,12 +589,13 @@ def _compress_bits(x, m): @triton.jit -def _rowtile_plan_kernel( +def _rowtile_plan_kernel_triton( request_indices, row_info, bits, block_table, slots, packed_bits, group_plan, lone_rows, lone_count, rows, bits_stride_row, block_table_stride_request, block_table_stride_page, slots_stride_group, - ROWS_PER_TILE: tl.constexpr, NUM_REQUESTS: tl.constexpr, ROWS_POW2: tl.constexpr, - BLOCK_N: tl.constexpr, SHARED_TILE_COST: tl.constexpr, NUM_CACHE_PAGES: tl.constexpr, + 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, ): # fmt: skip """Per group of rows of one request: list the union of its rows' resident @@ -342,6 +675,44 @@ def _rowtile_plan_kernel( tl.store(lone_rows + base + offsets, members.to(tl.int32), mask=live) +def _rowtile_plan( + grid, request_indices, row_info, bits, block_table, slots, packed_bits, group_plan, + lone_rows, lone_count, rows, *, rows_per_tile, tokens_per_block, num_cache_pages, +): # fmt: skip + """The plan pass over `grid = (groups,)`: `_RowtilePlanKernel` for int32 request indices + (packed words written below each group's tile count), else the Triton kernel (packed words + OR-ed into a zeroed `packed_bits`).""" + (num_groups,) = grid + if _cute_plan_supported(request_indices, tokens_per_block, rows_per_tile, bits.shape[1]): + _rowtile_plan_cute( + request_indices, row_info, bits, block_table, slots, packed_bits, group_plan, + lone_rows, lone_count, rows, num_groups, rows_per_tile=rows_per_tile, + tokens_per_block=tokens_per_block, num_cache_pages=num_cache_pages, + ) # fmt: skip + return + table_width = block_table.shape[1] + _rowtile_plan_kernel_triton[(num_groups,)]( + request_indices, row_info, bits, block_table, slots, packed_bits, group_plan, + lone_rows, lone_count, rows, + bits.stride(0), block_table.stride(0), block_table.stride(1), slots.stride(0), + ROWS_PER_TILE=rows_per_tile, NUM_REQUESTS=block_table.shape[0], + ROWS_POW2=triton.next_power_of_2(rows_per_tile), BLOCK_N=BLOCK_N, + SHARED_TILE_COST=_SHARED_TILE_COST, NUM_CACHE_PAGES=num_cache_pages, + PAGE_TABLE_WIDTH=table_width, TOKENS_PER_BLOCK=tokens_per_block, + PLAN_WORDS=_PLAN_WORDS, num_warps=_PLAN_WARPS, + ) # fmt: skip + + +class _PlanLauncher: + """`_rowtile_plan_kernel[(groups,)](...)`: grid-indexed like the entry's other launches.""" + + def __getitem__(self, grid): + return functools.partial(_rowtile_plan, tuple(grid)) + + +_rowtile_plan_kernel = _PlanLauncher() + + @triton.jit def _rowtile_lone_kernel( q, k_cache, v_cache, block_table, selected_tokens, request_indices, row_info, @@ -355,10 +726,11 @@ def _rowtile_lone_kernel( selected_stride_row: tl.constexpr, selected_stride_token: tl.constexpr, output_stride_row: tl.constexpr, output_stride_head: tl.constexpr, output_stride_dim: tl.constexpr, - NUM_CACHE_PAGES: tl.constexpr, NUM_REQUESTS: tl.constexpr, PAGE_TABLE_WIDTH: tl.constexpr, + 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, LONE_PROGRAMS: tl.constexpr, USE_PDL: 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.""" @@ -391,7 +763,7 @@ def _rowtile_lone_row( block_table_stride_request, block_table_stride_page, selected_stride_row, selected_stride_token, output_stride_row, output_stride_head, output_stride_dim, - NUM_CACHE_PAGES: tl.constexpr, NUM_REQUESTS: tl.constexpr, PAGE_TABLE_WIDTH: tl.constexpr, + 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 @@ -568,18 +940,19 @@ def qsa_rowtile_sparse_gqa( 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) + # The CuTe plan writes every packed word the row-tile pass reads; the Triton one ORs. + on_cute = _cute_plan_supported(request_indices, tokens_per_block, rows_per_tile, num_words) + alloc = torch.empty if on_cute else torch.zeros + packed_bits = alloc((rows, num_words), dtype=torch.int64, device=q.device) group_plan = torch.empty((num_groups, 2), dtype=torch.int32, device=q.device) lone_rows = torch.empty((rows,), dtype=torch.int32, device=q.device) - lone_count = torch.zeros((1,), dtype=torch.int32, device=q.device) + # [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] _rowtile_plan_kernel[(num_groups,)]( request_indices, row_info, bits, block_table, slots, packed_bits, group_plan, - lone_rows, lone_count, rows, - bits.stride(0), block_table.stride(0), block_table.stride(1), slots.stride(0), - ROWS_PER_TILE=rows_per_tile, NUM_REQUESTS=block_table.shape[0], ROWS_POW2=rows_pow2, - BLOCK_N=BLOCK_N, SHARED_TILE_COST=_SHARED_TILE_COST, NUM_CACHE_PAGES=k_cache.shape[0], - PAGE_TABLE_WIDTH=table_width, TOKENS_PER_BLOCK=tokens_per_block, - PLAN_WORDS=_PLAN_WORDS, num_warps=_PLAN_WARPS, + lone_rows, lone_count, rows, rows_per_tile=rows_per_tile, + tokens_per_block=tokens_per_block, num_cache_pages=k_cache.shape[0], ) # fmt: skip use_pdl = get_env_enable_pdl() _rowtile_lone_kernel[(min(rows, _LONE_PROGRAMS), num_kv_heads)]( @@ -590,8 +963,9 @@ def qsa_rowtile_sparse_gqa( LONE_PROGRAMS=min(rows, _LONE_PROGRAMS), USE_PDL=use_pdl, LONE_STAGES=_LONE_STAGES, num_warps=_LONE_WARPS, **geometry, ) # fmt: skip - _rowtile_attention_cute[(num_groups, num_kv_heads)]( - q, k_cache, v_cache, slots, packed_bits, group_plan, output, softmax_scale, rows, + _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 @@ -655,13 +1029,15 @@ class QsaSparseGqaRowtile(OpWrapper): less than the rows' own columns would; otherwise -- or when its rows span requests -- its rows join the lone list. A lone-row pass (256 programs per KV head striding the list) serves listed rows over their own columns, - as the per-row kernel does. The row-tile pass runs one program per - (group, KV head), loads each listed tile's K/V once for the 120 (row, - head) pairs and masks each row to its own re-indexed bits; it is launched + 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), so the entry - serves SM103, bf16 q and K/V and D = 256. + 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, @@ -717,8 +1093,13 @@ class QsaSparseGqaRowtile(OpWrapper): does (16292-row mixed step: 3859 vs 3836 us). Below a few hundred rows it is slower than the per-row kernel (generation-only launches 1.4-2x); the target's served context launches are thousands of rows. The tensor-core - row-tile pass alone, on two dumped 16384-row context launches of the - target (GB300): 631 and 887 us. + 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 and the re-indexed bitmap, each `rows x ceil(table_width * tpb / 64)` int64 words (21 MB at 16384 rows and 320 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py index 542d4d9a4e70..ae062832dcb6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py @@ -50,6 +50,9 @@ 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: @@ -151,15 +154,16 @@ def kernel(self, mKeys, mEnd, mPos, mLen, mTopk, mSel, ld_keys: cute.CopyAtom): 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. The row's metadata and four consecutive keys a thread (a column past the - # window is no candidate: -1); the largest and smallest key in the window. - n = mEnd[row] - seqlen = mLen[row] - visible = cutlass.Int32(mPos[row]) + 1 + # 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) @@ -282,19 +286,31 @@ def kernel(self, mKeys, mEnd, mPos, mLen, mTopk, mSel, ld_keys: cute.CopyAtom): cand_key[slot] = keys[c] cand_idx[slot] = e0 + c cute.arch.sync_threads() - if tidx < at_thr: - kt = cand_key[tidx] - it = cand_idx[tidx] - rank = cutlass.Int32(0) - for j in cutlass.range(0, at_thr, 1, unroll=8): + # 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)) - if rank < need: - cute.arch.atomic_or( - flags.iterator + (it >> 5), cutlass.Int32(1) << (it & 31), scope="cta" - ) + 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 @@ -513,6 +529,13 @@ class QsaTopkExpandDecode(OpWrapper): 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 = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_attn_decode_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_attn_decode_prologue.py index 854f7efb1d6c..4155e8a99b40 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_attn_decode_prologue.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_attn_decode_prologue.py @@ -8,15 +8,34 @@ the side stream the index plane's raw-key and coordinate store (`qsa_paged_index_store`) and the group compression (`qsa_prefill_compress`), then a cross-stream join. Here the programs of `qwen4_exp_qk_norm_rope` run as they do there; the ones holding k heads also -write their rotated rows into the K pool; one more program per row stores v, the index -key and the coordinates, and -- at a group's last token -- pools the group with this -token's key taken from the projection row (the cached copy is the one being written), -normalizes, rotates and stores it at the anchor slot, as the compression kernel does. +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, and at a group end the group's earlier -keys from the index plane with their rotation -- is read before `griddepcontrol.wait`; -after it a program reads only the projection row. +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 @@ -30,12 +49,22 @@ 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, @@ -83,11 +112,10 @@ def _norm_weights( @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 +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 at the store. Returns the stored - (rounded) tile.""" + `* (1 + w)`, rotate; fp32 throughout, one rounding to `out_dtype`.""" HALF: tl.constexpr = ROT // 2 d = tl.arange(0, D) rot = d < ROT @@ -100,11 +128,45 @@ def _norm_rope_apply(src, dst, heads, head_mask, w2, wp2, c, s, eps, D: tl.const 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 (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 @@ -188,6 +250,16 @@ def _attn_decode_prologue_kernel( 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, @@ -207,12 +279,19 @@ def _attn_decode_prologue_kernel( 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 index scores read nothing of this launch before their own wait. + # 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 @@ -232,23 +311,42 @@ def _attn_decode_prologue_kernel( 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: - d = tl.arange(0, D) is_k = (main_heads >= HQ) & (main_heads < HQ + HK) & live - tl.store( - k_cache - + page * s_kc_page - + (main_heads - HQ).to(tl.int64)[:, None] * s_kc_head - + slot * s_kc_tok - + d[None, :], - out.to(k_cache.dtype.element_ty), - mask=is_k[:, None], - ) + 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) @@ -259,6 +357,148 @@ def _attn_decode_prologue_kernel( 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 @@ -274,12 +514,30 @@ def _attn_decode_prologue_kernel( 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=group_end & (offsets < CR - 1)[:, None], + mask=(earlier & ~in_launch)[:, None], other=0.0, ) first = tl.maximum(logical - (CR - 1), 0) @@ -287,9 +545,23 @@ def _attn_decode_prologue_kernel( block_table + request * s_bt_req + (first // TPB) * s_bt_col, mask=group_end, other=0 ).to(tl.int64) first_base = position_cache + first_page * s_pc_page + (first % TPB) * s_pc_slot - pos_t = tl.load(first_base, mask=group_end, other=0) - pos_h = tl.load(first_base + s_pc_axis, mask=group_end, other=0) - pos_w = tl.load(first_base + 2 * s_pc_axis, mask=group_end, other=0) + 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, @@ -306,14 +578,6 @@ def _attn_decode_prologue_kernel( if USE_PDL: tl.extra.cuda.gdc_wait() if live: - # v, as the K/V append stores it. - d = tl.arange(0, D) - for h in tl.static_range(HK): - v = tl.load(row + (HQ + HK + h) * D + d) - tl.store( - v_cache + page * s_vc_page + h * s_vc_head + slot * s_vc_tok + d, - v.to(v_cache.dtype.element_ty), - ) # The coordinates, as the index store writes them. tl.store( position_cache + page * s_pc_page + slot * s_pc_slot + axes * s_pc_axis, @@ -325,6 +589,19 @@ def _attn_decode_prologue_kernel( 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( @@ -348,8 +625,22 @@ class Qwen4ExpAttnDecodePrologue(OpWrapper): last token, the compressed group) and the coordinates into the index plane and the position cache, at the row's page and slot. Returns `(proj, index_q)`. - Fusion boundary: as the four entries; the projection before and the index scores - after are the caller's. + 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}) @@ -364,17 +655,79 @@ class Qwen4ExpAttnDecodePrologue(OpWrapper): 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]`; decode rows, - one per request (a group's earlier tokens are already in the index plane). A row - whose page lookup fails writes no cache, as the stores and compression skip it. - - Numerics: the same arithmetic as the four entries; outputs bitwise theirs. + `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__( @@ -410,6 +763,11 @@ def __call__( 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: @@ -417,8 +775,19 @@ def __call__( hpb, warps = _qk_launch(num_tokens) main_programs = triton.cdiv(num_heads_q + num_heads_k, hpb) section = compress_mrope_section or (0, 0, 0) + 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)]( + _attn_decode_prologue_kernel[(num_tokens, main_programs + 2 + num_heads_k + score_tiles)]( proj, proj.stride(0), index_q, @@ -459,6 +828,16 @@ def __call__( 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, @@ -478,6 +857,12 @@ def __call__( 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, @@ -517,8 +902,14 @@ def reference( 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.""" + """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, @@ -531,16 +922,31 @@ def reference( heads = proj.view(rows, -1, head_dim) k = heads[:, num_heads_q : num_heads_q + num_heads_k] v = heads[:, num_heads_q + num_heads_k : num_heads_q + 2 * num_heads_k] - _kv_store( - k, - v, - k_cache, - v_cache, - request_indices, - logical_positions, - block_table, - tokens_per_block, - ) + 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 ) @@ -555,6 +961,16 @@ def reference( 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: @@ -566,6 +982,28 @@ def is_valid(self, proj, index_q, position_ids, coords, cos_sin, *args, **kwargs 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/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 new file mode 100644 index 000000000000..098506078e6f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/spec_rollback_commit.py @@ -0,0 +1,368 @@ +# 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 new file mode 100644 index 000000000000..bdd150c32d71 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/verify_step_setup.py @@ -0,0 +1,215 @@ +# 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/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 new file mode 100644 index 000000000000..ff6fa1208e36 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/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/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 new file mode 100644 index 000000000000..71a5c049e9dd --- /dev/null +++ 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 @@ -0,0 +1,226 @@ +# 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. A producer that already wrote the partial maxima of what it stored +(`amax_partials`, filled by atomic maxima) hands them over and the first launch is skipped. +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 +#: Slots of a producer-filled partial-maximum buffer (`amax_partials`): a producer program +#: folds its maximum into slot `program % AMAX_PARTS`, spreading the atomics. +AMAX_PARTS = 1024 + + +def amax_partials(device: torch.device) -> torch.Tensor: + """A zeroed fp32 `[AMAX_PARTS]` buffer for a producer to fold `max(|out|)` into.""" + return torch.zeros(AMAX_PARTS, dtype=torch.float32, device=device) + + +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, amax_parts: torch.Tensor | None = None +) -> 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). + `amax_parts`, when given, holds partial maxima whose maximum is `max(|x|)`.""" + 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 = amax_parts + if part is None: + 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. `amax_parts` (fp32, at most + `AMAX_PARTS` partial maxima whose maximum is `max(|x|)`, as a producer that stored `x` + folds them) skips the maximum pass; the result is the same bits. + """ + + 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, + amax_parts: torch.Tensor | None = None, + ) -> torch.Tensor: + q, scale = fp8_quantize_per_tensor(x, amax_parts) + 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, + amax_parts: torch.Tensor | None = None, + ) -> 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, + amax_parts: torch.Tensor | None = None, + ) -> 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 + assert amax_parts is None or ( + amax_parts.dtype == torch.float32 and 0 < amax_parts.numel() <= AMAX_PARTS + ), "amax_parts: fp32, at most AMAX_PARTS partial maxima" + + +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/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 new file mode 100644 index 000000000000..cd5c8a6eb386 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/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/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 new file mode 100644 index 000000000000..bbe39917e7a4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py index 8e58e9df35cc..8deac6c0ee69 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py @@ -26,6 +26,7 @@ from __future__ import annotations import functools +import os from typing import Optional, Tuple import cuda.bindings.driver as cuda_drv @@ -33,8 +34,10 @@ 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 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 @@ -48,6 +51,31 @@ #: Streaming multiprocessors the one-wave grid must fit (GB300). _SMS = 152 _COMPILE_OPTIONS = "--ptxas-options -maxrregcount=128" +#: ns the weight-register loads are held after a CTA starts (dependents released first). +#: In a chain of sites alone the next site's first kernel becomes resident as this site's +#: mix starts its critical path, and holding its 6.6 MB weight burst ~0.5 us keeps that +#: path's L2 round trips short (-0.4 us a site in the chained micro-bench). In the full STP +#: decode step it costs more than it saves (con1: 412.5 tok/s/user at 0, 399.0 at 500), so +#: it is off by default. +_WEIGHT_DELAY_NS = int(os.environ.get("TRTLLM_HC_FIRST_WEIGHT_DELAY_NS", "0")) + + +@dsl_user_op +def _spin_ns(ns, *, loc=None, ip=None): + """Spin (nanosleep 64) for `ns` nanoseconds of %globaltimer.""" + llvm.inline_asm( + T.i32(), + [cutlass.Int32(ns).ir_value(loc=loc, ip=ip)], + "{ .reg .pred p; .reg .u64 t0, t, e; .reg .u64 d; mov.u64 t0, %globaltimer; " + "cvt.u64.u32 d, $1; add.u64 e, t0, d; SPIN: mov.u64 t, %globaltimer; " + "setp.lt.u64 p, t, e; @!p bra DONE; nanosleep.u32 64; bra SPIN; DONE: mov.u32 $0, 0; }", + "=r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) def outs_for(n_cols: int, hc_count: int) -> int: @@ -159,6 +187,10 @@ def kernel( hc_vec = cute.logical_divide(mHc, (None, VEC)) nw_vec = cute.logical_divide(mNorm, (None, VEC)) + if const_expr(_WEIGHT_DELAY_NS > 0): + if const_expr(self.pdl): + cute.arch.griddepcontrol_launch_dependents() + _spin_ns(_WEIGHT_DELAY_NS) # 1. Static operands, before the wait. w_reg = cute.make_rmem_tensor( cute.make_layout((outs, VEC), stride=(VEC, 1)), cutlass.BFloat16 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py index 45a3c08057bd..df11a7567828 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py @@ -32,8 +32,10 @@ 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 ( @@ -50,6 +52,24 @@ _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 @@ -140,7 +160,9 @@ def kernel( # 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. + # 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)] @@ -149,29 +171,33 @@ def kernel( 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` in rows of VEC with one pad word, so the 8 threads of a row (stride VEC + 1 - # words) read 8 different banks. + # `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 // VEC, VEC), stride=(VEC + 1, 1)), - byte_alignment=16, + 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, ) - if tidx < cols: - logit = cutlass.Float32(0.0) - for t in cutlass.range_constexpr(streams): - logit = pv[t] * rstd[t] + logit - z = logit * (1.0 / streams) - sig = 1.0 / (1.0 + cute.math.exp(-z)) - if tidx < rank: - low[tidx // VEC, tidx % VEC] = z * sig - else: - if const_expr(self.inject_out): + 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() @@ -180,15 +206,17 @@ def kernel( # 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) * low[kt * GROUP + q, e] + sums[e % 2] - ) + 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) + vals[s, i] = gate * (x * rstd_s * gam) cute.arch.sync_threads() if warp == 0: m = cutlass.Float32(0.0) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_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 new file mode 100644 index 000000000000..d3756c320875 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_rows_tc.py @@ -0,0 +1,1338 @@ +"""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_decode_tc_first.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py index df942de81ed0..834196008834 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py @@ -1,13 +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..128 rows) as one +"""Hyper-connection mixer first half at decode row counts (2..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 -`[128 rows, BN columns]` partial over `K / KS` of the 10240-wide stream: a TMA -warp streams the normed rows (rows past `T` zero-filled by TMA, never read) +`[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) and the weight rows (only the `n_cols` real ones; the load-time zero padding to 384 rows is never read), one MMA warp accumulates in tensor memory, four epilogue warps move the fp32 partial into shared memory. The `KS` CTAs of one @@ -41,8 +42,8 @@ from tensorrt_llm._torch.kernels.blackwell.custom_pipeline import PipelineTmaUmma, PipelineUmmaAsync from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr -#: Largest row count this entry is certified for: one 128-row MMA tile. -MAX_ROWS = 128 +#: Largest row count this entry is certified for: two 128-row MMA tiles. +MAX_ROWS = 256 #: Programmatic dependent launch (sm_90+), off with TRTLLM_ENABLE_PDL=0. Without it #: the grid-dependency wait and the dependents trigger are no-ops. @@ -101,7 +102,7 @@ class HcDecodeTcFirstKernel: Compile-time: `bn` (MMA N, a multiple of 16), `ks` (K split = cluster size), `hc`, `rank`, `n_cols`, `k`, `with_inject`. Runtime: the row count - (<= 128) and the pointers. + (one grid-z M tile per 128 rows) and the pointers. """ def __init__(self, bn, ks, hc, rank, n_cols, k, with_inject, epi_cols=16): @@ -210,7 +211,7 @@ def __call__( self.b_smem_layout_staged, self.epi_tile, ).launch( - grid=[ks, self.n_tiles, 1], + grid=[ks, self.n_tiles, (m + 127) // 128], block=[self.threads_per_cta, 1, 1], cluster=[ks, 1, 1], stream=stream, @@ -241,9 +242,10 @@ def kernel( cpasync.prefetch_descriptor(tma_atom_a) cpasync.prefetch_descriptor(tma_atom_b) tidx, _, _ = cute.arch.thread_idx() - bidx, bidy, _ = cute.arch.block_idx() + bidx, bidy, bidz = cute.arch.block_idx() kr = bidx # K slice == rank in the (KS, 1, 1) cluster nt = bidy # column tile + mt = bidz # 128-row M tile @cute.struct class SharedStorage: @@ -335,7 +337,7 @@ class SharedStorage: pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) if warp_idx == self.tma_warp_id: - tAgA_slice = tAgA[(None, 0, None, 0)] + tAgA_slice = tAgA[(None, mt, None, 0)] tBgB_slice = tBgB[(None, nt, None, 0)] ab_producer.reset() # The weight does not depend on the previous kernel: the first stages' weight @@ -448,10 +450,12 @@ class SharedStorage: cute.arch.cluster_arrive() cute.arch.cluster_wait() - # This CTA's share of the tile's `m * BN` outputs, all 6 warps. Every load is - # issued before any sum (indices clamped in bounds, so no branch separates - # them); the `KS` partials are summed in slice order (deterministic). - total = m * self.bn + # This CTA's share of the tile's `rows * BN` outputs (`rows` of this M tile), + # all 6 warps. Every load is issued before any sum (indices clamped in bounds, + # so no branch separates them); the `KS` partials are summed in slice order + # (deterministic). + row0 = mt * 128 + total = cutlass.min(m - row0, 128) * self.bn per = (total + self.ks - 1) // self.ks base = kr * per nthr = self.threads_per_cta @@ -479,10 +483,12 @@ class SharedStorage: z = acc * inv_hc sg = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.math.exp(-z)) if n < self.rank: - mLow[(row, n)] = (z * sg).to(self.io_dtype) + mLow[(row0 + row, n)] = (z * sg).to(self.io_dtype) if cutlass.const_expr(self.with_inject): if n >= self.rank and n < self.rank + self.hc: - mInj[(row, n - self.rank)] = (cutlass.Float32(2.0) * sg).to(self.io_dtype) + mInj[(row0 + row, n - self.rank)] = (cutlass.Float32(2.0) * sg).to( + self.io_dtype + ) # No CTA may exit (releasing its shared memory) while a peer still reads it. cute.arch.cluster_arrive_relaxed() @@ -538,10 +544,32 @@ def hc_decode_tc_first(normed, first, rank, hc_count, with_inject, bn, ks): return low, inject -#: Tile: `_BN` weight rows per CTA (7 tiles cover the 324 real rows of 336) and -#: a `_KS`-way K split, one cluster per column tile. -_BN = 48 -_KS = 16 +#: 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 + + +def _tile(rows: int) -> tuple[int, int]: + return next(t for hi, t in _TILES if rows <= hi) + + +def precompile(k: int, rank: int, hc_count: int, device) -> None: + """Compile every `(tile, with_inject)` variant `Qwen4ExpHcDecodeTcFirst` can launch + for a `[*, k]` stream, without launching: a variant first met by a serving step would + stall it for the compile. Placeholder operands (the compile reads only their + addresses and the `low` layout).""" + normed = torch.empty((MAX_ROWS, k), dtype=torch.bfloat16, device=device) + first = torch.empty((rank + hc_count, k), dtype=torch.bfloat16, device=device) + low = torch.empty((MAX_ROWS, rank), dtype=torch.bfloat16, device=device) + ptrs = (normed.data_ptr(), first.data_ptr(), low, low.data_ptr()) + for _, (bn, ks) in _TILES: + for with_inject in (True, False): + n_cols = rank + (hc_count if with_inject else 0) + _launcher(bn, ks, hc_count, rank, n_cols, k, with_inject, ptrs) class Qwen4ExpHcDecodeTcFirst(OpWrapper): @@ -565,9 +593,21 @@ class Qwen4ExpHcDecodeTcFirst(OpWrapper): # rank 320, hidden 2560, the padded [384, K] weight, bf16. CELLS: tuple[Cell, ...] = ( Cell( - why="checkpoint 4x2560, 324 cols of 384, inject, decode 128 rows (MAX_ROWS)", + why="checkpoint 4x2560, 324 cols of 384, inject, decode 128 rows (one M tile)", spec=dict(rows=128), ), + Cell( + why="checkpoint 4x2560, 324 cols of 384, inject, 256 rows (MAX_ROWS, two M tiles)", + spec=dict(rows=256), + ), + Cell( + why="checkpoint 4x2560, 324 cols of 384, inject, 129 rows (a 1-row second M tile)", + spec=dict(rows=129), + ), + 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), @@ -586,25 +626,29 @@ class Qwen4ExpHcDecodeTcFirst(OpWrapper): `_HC_TC_FIRST_MAX_ROWS` (= MAX_ROWS) rows, 97 per forward, between `qwen4_exp_hc_inject_norm` (or the PLE tail) and `qwen4_exp_hc_up_gate_mean(_fp4)`. Preconditions (bf16): `normed` - `[T, K]` contiguous, T <= 128; `first` the `.t()` view of a row-major + `[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 * KS. Only the R (+ C) real weight rows are read (TMA zero-fills + 64 * 16. Only the R (+ C) real weight rows are read (TMA zero-fills the tile past them), never the load-time padding. - Launch: grid (KS, cdiv(R + C, BN)) = (16, 7) CTAs of 6 warps, cluster - (KS, 1, 1), one CTA per SM (~227 KB shared memory: every K tile of the - slice in flight plus the padded fp32 partial). The partials are reduced - through distributed shared memory between two cluster barriers; no - global workspace. Deterministic (fixed K-slice order). Programmatic - dependent launch (TRTLLM_ENABLE_PDL=0 turns it off): weight stages - before `griddepcontrol.wait`, `normed` after it; dependents released - at kernel start. Graph-capturable (current stream, no host sync, no - persistent state). Compiles once per (BN, KS, C, R, N, K, inject). - Measured on GB300 (graph replay, 24 cold weights): 6.4 us at 128 rows - against 8.6 us for cuBLAS split-K + reduce + lowrank_act; tile sweep - BN 32 / 48 / 64 / 112 x KS 8 / 10 / 16, 48 x 16 fastest at every row - count. - Certified: bf16, 324 columns with the injection at 2 and 128 rows, 320 - columns without it at 128 rows; graph replay. + 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. + 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. 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. @@ -618,7 +662,9 @@ def __call__( hc_count: int, with_inject: bool, ): - return hc_decode_tc_first(normed, first, rank, hc_count, with_inject, _BN, _KS) + return hc_decode_tc_first( + normed, first, rank, hc_count, with_inject, *_tile(normed.shape[0]) + ) def reference( self, @@ -648,11 +694,11 @@ def is_valid( rows, k = normed.shape n_cols = rank + (hc_count if with_inject else 0) assert normed.dtype == first.dtype == torch.bfloat16, "bf16 operands" - assert rows <= MAX_ROWS, f"at most {MAX_ROWS} rows (one MMA tile)" + assert rows <= MAX_ROWS, f"at most {MAX_ROWS} rows (two MMA tiles)" assert normed.stride() == (k, 1), "normed contiguous [T, K]" assert first.shape[0] == k and first.shape[1] >= n_cols, "first [K, >= n_cols]" assert first.stride() == (1, k), "first must be the .t() view of a row-major [N, K] weight" - assert k % (64 * _KS) == 0, f"K must be a multiple of {64 * _KS}" + assert k % _K_ALIGN == 0, f"K must be a multiple of {_K_ALIGN}" for t in (normed, first): assert t.data_ptr() % 16 == 0, "16-byte aligned operands" 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 new file mode 100644 index 000000000000..9c4a5dbeadcc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_first_rescale_tc.py @@ -0,0 +1,969 @@ +# 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_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 new file mode 100644 index 000000000000..597216df40fc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_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/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 new file mode 100644 index 000000000000..69fcca405ae8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_rstd.py @@ -0,0 +1,1119 @@ +# 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/moe/fp4_block_scale_moe_renorm_permuted.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_permuted.py index f6efc5755b5a..90e2623cf65d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_permuted.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_permuted.py @@ -10,10 +10,9 @@ import torch -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from .fp4_block_scale_moe_renorm_routed import ACT_SWIGLU, ROUTING_RENORMALIZE, _ckpt +from .fp4_block_scale_moe_renorm_routed import _ckpt, run_renorm_moe from .fp4_block_scale_moe_renorm_routed import fp4_block_scale_moe_renorm_routed as _finalized @@ -58,6 +57,11 @@ class Fp4BlockScaleMoeRenormPermuted(OpWrapper): "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), @@ -83,7 +87,7 @@ class Fp4BlockScaleMoeRenormPermuted(OpWrapper): not a multiple of 256). Hard limits - T certified at 1, 64 and 128 (decode); the decode forward folds at + 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 """ @@ -104,34 +108,21 @@ def __call__( top_k: int, intermediate_size: int, ) -> tuple: - rows, topk_weights, topk_rows = torch.ops.trtllm.fp4_block_scale_moe_runner( + rows, topk_weights, topk_rows = run_renorm_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, False, - ACT_SWIGLU, ) return rows, topk_weights, topk_rows diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_routed.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_routed.py index d4a76c8155c4..20484966694c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_routed.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_routed.py @@ -7,6 +7,8 @@ from __future__ import annotations +from functools import cache + import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* @@ -14,6 +16,9 @@ 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 .renorm_moe_routing_op import renorm_moe_routing_op as _renorm @@ -29,6 +34,167 @@ #: 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: `_pinned_tactic` refuses a tactic the runner does not list as valid. +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]: + """`DECODE_TACTICS[bucket]`, once confirmed valid for the bucket's row count.""" + 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) + assert list(tactic) in [list(t) for t in valid], ( + f"pinned tactic {tactic} for {bucket} rows is not a valid config of this build's " + "trtllm-gen FP4 MoE runner; re-measure DECODE_TACTICS" + ) + 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.""" @@ -96,6 +262,12 @@ class Fp4BlockScaleMoeRenormRouted(OpWrapper): 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 @@ -120,34 +292,21 @@ def __call__( top_k: int, intermediate_size: int, ) -> torch.Tensor: - return torch.ops.trtllm.fp4_block_scale_moe_runner( + return run_renorm_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, True, - ACT_SWIGLU, )[0] def reference( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_moe_renorm_cute_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_moe_renorm_cute_dsl.py index 85cbcb72ae9f..0ebde14b0fac 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_moe_renorm_cute_dsl.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_moe_renorm_cute_dsl.py @@ -13,15 +13,83 @@ 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 .fp4_block_scale_moe_renorm_routed import CERTIFIED_PAIRS, _ckpt from .fp4_block_scale_moe_renorm_routed import fp4_block_scale_moe_renorm_routed as _routed -from .renorm_moe_routing_op import renorm_moe_routing_op #: Token tile of the grouped GEMMs (the kernels' M tile). TILE_TOKENS = 128 +#: 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) +_FinalizeKernel = _cdo.Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel + + +class _Fc2Kernel(_FinalizeKernel): + def __init__(self, *args, **kwargs): + super().__init__( + *args, mma_inst_tile_k=_FC2_MMA_INST_TILE_K, row_owned_epilogue=True, **kwargs + ) + + +class _Fc2HintKernel(_FinalizeKernel): + 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): @@ -77,12 +145,13 @@ class Fp4MoeRenormCuteDsl(OpWrapper): expert (folded into `output1_scale_scalar`); a checkpoint whose experts disagree on `down_proj.input_scale` cannot use this entry as is. - Kernels (all TRT-LLM's): `trtllm::moe_sort` (tile metadata), + Kernels (all TRT-LLM's): `trtllm::moe_topk_sort` (routing + tile metadata), `trtllm::moe_output_memset_inplace` (zeroes the rows the scatter-add - touches), `trtllm::cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell`, - `trtllm::cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell`; the two - GEMMs pick their tactics through the AutoTuner. Not CUDA-graph certified - (prefill steps run eagerly). + 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 @@ -104,67 +173,79 @@ def __call__( 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 - ids, weights = renorm_moe_routing_op(routing_logits, top_k, torch.bfloat16) + 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() - t2e, t2lim, e2p, p2e, _, live = torch.ops.trtllm.moe_sort( - token_selected_experts=ids, - token_final_scales=scales, - num_experts=num_experts, - top_k=top_k, - local_expert_offset=0, - local_num_experts=num_experts, - tile_tokens_dim=TILE_TOKENS, - ) out = torch.empty((tokens, hidden), dtype=torch.bfloat16, device=hidden_states.device) - torch.ops.trtllm.moe_output_memset_inplace( - input=out, - tile_idx_to_mn_limit=t2lim, - expanded_idx_to_permuted_idx=e2p, - permuted_idx_to_expanded_idx=p2e, - num_non_exiting_tiles=live, - tile_tokens_dim=TILE_TOKENS, - top_k=top_k, - ep_size=1, - enable_alltoall=False, - ) - h, h_sf = torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell( - input=hidden_states.view(_FP4), - weight=fc1_w.view(_FP4), - input_scale=hidden_states_scale.view(torch.uint8).view(tokens, hidden // 16), - weight_scale=fc1_s.view(torch.uint8), - alpha=fc1_alpha, - tile_idx_to_group_idx=t2e, - tile_idx_to_mn_limit=t2lim, - permuted_idx_to_expanded_idx=p2e, - num_non_exiting_tiles=live, - global_sf=fc2_input_scale, - num_experts=num_experts, - top_k=top_k, - num_local_experts=num_experts, - local_expert_offset=0, - tile_size=TILE_TOKENS, + + 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, ) - torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell( - input=h.view(_FP4), - weight=fc2_w.view(_FP4), - input_scale=h_sf.view(torch.uint8), - weight_scale=fc2_s.view(torch.uint8), - alpha=fc2_alpha, - output=out, - tile_idx_to_group_idx=t2e, - tile_idx_to_mn_limit=t2lim, - permuted_idx_to_expanded_idx=p2e, - num_non_exiting_tiles=live, - token_final_scales=scales, - num_experts=num_experts, - top_k=top_k, - num_local_experts=num_experts, - local_expert_offset=0, - tile_size=TILE_TOKENS, - output_dtype=torch.bfloat16, + 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 @@ -198,6 +279,7 @@ def is_valid( 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), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/nvfp4_moe_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/nvfp4_moe_gemv.py new file mode 100644 index 000000000000..183237318f1d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/nvfp4_moe_gemv.py @@ -0,0 +1,1787 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Small-row NVFP4 MoE layer entered with router logits, left unfinalized: Renormalize +routing + FC1 + SwiGLU + NVFP4 requantization + FC2 over the trtllm-gen expert layout, in +two PDL-chained CuTe DSL launches on the tcgen05 block-scaled tensor core +(`kind::mxf4nvf4.block_scale.block16`: e2m1 operands, ue4m3 16-element scales, fp32 in +TMEM), for 1 .. MAX_ROWS rows. Same outputs as `fp4_block_scale_moe_renorm_permuted`: the +permuted FC2 rows, the routing weights and the slot -> row map the next hyper-connection +kernel gathers. + +Geometry (asserted): 512 experts, top-10, hidden 2560, intermediate 640. + +Routing (both launches, every CTA, one warp per row): the runner's Renormalize policy -- +the 10 largest bf16 logits ranked by the packed key (order-preserving bf16 bits << 16 | +65535 - expert), ties to the lower expert, then an fp32 softmax over the 10 with the routing +kernel's arithmetic (libdevice `expf`, a 32-lane xor-butterfly sum, IEEE division), each +weight rounded once to bf16. The distinct experts of all rows, in ascending id order, are the +work list; a row's slot k is its k-th ranked expert. + +MMA: M 128 weight rows x N 8 rows x K 64 a step; A (weights) arrives by 3D TMA boxes in the +128- (FC1) / 64-byte (FC2) swizzled K-major layout, B (the rows) sits in 16-byte core +matrices, and both scale operands are 512-byte 128x4 atoms (the checkpoint's own for the +weights) moved into TMEM by `tcgen05.cp.32x128b.warpx4`. The TMA descriptors of an expert +stack are built once (`register_weights`) and read from global memory. + +Launch 1 (`_Fc1`): after `griddepcontrol.wait` it releases its dependents and routes; a +producer thread streams its 128-row FC1 tiles (5 K stages of 36 KB) into a shared-memory +ring, one thread issues the MMAs (the whole K of a tile in one TMEM accumulator), and four +epilogue warps read the accumulator: TMEM lane quarter q of tile t is the 32-row block +4 t + q, whose rows are the up and gate rows of 16 intermediate columns (row d's gate is +row d + 8, a shuffle away) -- one NVFP4 block. `up * output1_scale_scalar` times +`silu(gate * output1_scale_gate_scalar)` is requantized (scale `e4m3(amax / 6)`, codes +`e2m1(v / scale)`) into a per-slot record in global memory. + +Launch 2 (`_Fc2`): before its wait it routes again (from the logits, which launch 1 read +after its own wait, so they are complete), loads its 128-row FC2 tiles and their weight +scales into shared memory and the scales into TMEM -- while launch 1 still runs. After the +wait it copies each unit's rows of the slot records into its B operand, issues the MMAs, +scales by `output2_scale_scalar` and stores bf16 rows. Launch 1 trips no barrier launch 2 +waits on: the only cross-launch dependency is `griddepcontrol.wait`, so residency of either +grid is never required. + +The two grids do not share an SM (launch 2's CTAs start where no launch-1 CTA is resident), +so at one row, where the FC1 work is always 100 tiles, launch 1 takes 100 SMs (its whole +tile in flight) and launch 2 two CTAs on each of the other 52; at more rows each takes every +SM and launch 2 starts as launch 1 leaves. + +Outputs: `rows [10 T, 2560]` bf16 with slot (t, k) at row 10 t + k, `topk_weights [T, 10]` +bf16, `topk_rows [T, 10]` int32 = 10 t + k (a constant per T). +""" + +from __future__ import annotations + +import functools + +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 + +from ..gemm.fp8_gemv import ( + _asm, + _bar_sync, + _bulk_g2s, + _bulk_prefetch_l2, + _env_pdl, + _ir, + _lds4, + _lds_f32, + _mbar_arrive, + _mbar_expect_tx, + _mbar_init, + _mbar_init_fence, + _mbar_wait, + _sms, + _sts_f32, + _unpack, +) +from ..gemm.nvfp4_gemv import _lds32, _sts4 +from .fp4_block_scale_moe_renorm_permuted import finalize +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 + +NUM_EXPERTS = 512 +TOP_K = 10 +HIDDEN = 2560 +INTER = 640 +#: Most rows a launch takes (the n8 MMA tile). +MAX_ROWS = 8 +#: Compiled row buckets: a call with T rows runs the smallest bucket >= T. +ROW_BUCKETS = (1, 2, 4, 8) + +BLK = 32 # rows of one requantization block (FC1) / output block (FC2) of a tile +N1, K1 = 2 * INTER, HIDDEN # FC1: 1280 rows of 1280 code bytes, 160 scales +N2, K2 = HIDDEN, INTER # FC2: 2560 rows of 320 code bytes, 40 scales +NT1, NT2 = N1 // 128, N2 // 128 # 128-row tiles per expert: 10 and 20 +KQ1 = 5 # FC1 K stages a tile (512 columns each) +A1_BYTES = 128 * 256 # one FC1 stage: two 128-byte swizzled chunks of 128 rows +SLOT1 = A1_BYTES + 8 * 512 # + its 8 scale atoms +A2_BYTES = 128 * (K2 // 2) # one FC2 tile: five 64-byte swizzled chunks of 128 rows +SLOT2 = A2_BYTES + (K2 // 64) * 512 +NB_ROW = 8 * 16 # B operand: one 16-byte k-chunk of the 8 rows +B2_BYTES = 4352 # FC2 B (20 chunks) and its 10 scale atoms (+ the 384-byte read overhang) +TC_THREADS = 256 # producer, MMA issuer, 4 epilogue warps, 2 helpers +#: Per-slot record of the requantized intermediate: 320 code bytes, 40 e4m3 scales, pad. +REC = 384 +#: Most units one CTA takes in a call. +MAX_ITEMS = 64 +#: TMEM columns per CTA (both launches together use all 512 of an SM). FC1: two 8-column +#: accumulators, the weight scales of 2 stages (32 each), the activation scales (160); FC2: +#: two accumulators, the weight scales of 2 tiles (40 each), the B scales of 2 units (40 each). +TMEM_COLS = 256 +SFA1_COL = 16 # FC1 activation scales follow the slots' weight scales +SFA2_COL = 16 # FC2 B scales follow the slots' weight scales +#: Instruction descriptor: E2M1 A and B, UE4M3 16-element block scales, N 8, M 128. +IDESC = (1 << 7) | (1 << 10) | ((8 >> 3) << 17) | ((128 >> 4) << 24) + + +def SFB_BYTES(k: int) -> int: # noqa: N802 + """Bytes of the activation's scale atoms over K (128 a 64-column step + the overhang a + 512-byte atom copy reads past the last).""" + return (k // 64) * 128 + 384 + + +_SMEM_BUDGET = 227 << 10 +_COMPILE_OPTIONS = "--opt-level 2 --enable-tvm-ffi" + + +# ------------------------------------------------------------------------------------------ +# inline PTX helpers +# ------------------------------------------------------------------------------------------ + + +@dsl_user_op +def _ldg_cg4(gaddr, *, loc=None, ip=None): + """16 bytes from global through L2 only (data another grid wrote this step).""" + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_ir(gaddr, cutlass.Int64, loc, ip)], + "ld.global.cg.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_nc_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 _lds_u8(saddr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + T.i32(), + [_ir(saddr, cutlass.Int32, loc, ip)], + "ld.shared.u8 $0, [$1];", + "=r,r", + loc, + ip, + side=True, + ) + ) + + +@dsl_user_op +def _sts32(saddr, v, *, loc=None, ip=None): + _asm( + T.i32(), + [_ir(saddr, cutlass.Int32, loc, ip), _ir(v, cutlass.Int32, loc, ip)], + "{ st.shared.b32 [$1], $2; mov.u32 $0, 0; }", + "=r,r,r", + loc, + ip, + side=True, + ) + + +@dsl_user_op +def _sts_u8(saddr, v, *, loc=None, ip=None): + _asm( + T.i32(), + [_ir(saddr, cutlass.Int32, loc, ip), _ir(v, cutlass.Int32, loc, ip)], + "{ .reg .b16 h; cvt.u16.u32 h, $2; st.shared.u8 [$1], h; mov.u32 $0, 0; }", + "=r,r,r", + loc, + ip, + side=True, + ) + + +@dsl_user_op +def _atom_or_shared(saddr, v, *, loc=None, ip=None): + _asm( + T.i32(), + [_ir(saddr, cutlass.Int32, loc, ip), _ir(v, cutlass.Int32, loc, ip)], + "{ .reg .b32 t; atom.shared.or.b32 t, [$1], $2; mov.u32 $0, 0; }", + "=r,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 _stg_v2(gaddr, v0, v1, *, loc=None, ip=None): + _asm( + T.i32(), + [_ir(gaddr, cutlass.Int64, loc, ip)] + [_ir(v, cutlass.Int32, loc, ip) for v in (v0, v1)], + "{ st.global.v2.b32 [$1], {$2, $3}; mov.u32 $0, 0; }", + "=r,l,r,r", + loc, + ip, + side=True, + ) + + +@dsl_user_op +def _stg_u8(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.u8 [$1], h; mov.u32 $0, 0; }", + "=r,l,r", + loc, + ip, + side=True, + ) + + +@dsl_user_op +def _cp_async16(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_wait_all(*, loc=None, ip=None): + _asm( + T.i32(), + [], + "{ cp.async.commit_group; cp.async.wait_all; mov.u32 $0, 0; }", + "=r", + loc, + ip, + side=True, + ) + + +@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 _keys2(w, e0, *, loc=None, ip=None): + """A word of two bf16 logits (experts e0, e0 + 1) -> their ranking keys: the bits made + order-preserving as unsigned (positive: flip the sign bit; negative: flip all) in the + high half, 65535 - expert in the low half (ties rank the lower expert first).""" + res = _asm( + llvm.StructType.get_literal([T.i32()] * 2), + [_ir(w, cutlass.Int32, loc, ip), _ir(e0, cutlass.Int32, loc, ip)], + "{ .reg .b32 lo, hi, m, i; " + "and.b32 lo, $2, 65535; shr.b32 hi, $2, 16; " + "shr.b32 m, lo, 15; mul.lo.u32 m, m, 32767; or.b32 m, m, 32768; xor.b32 lo, lo, m; " + "shr.b32 m, hi, 15; mul.lo.u32 m, m, 32767; or.b32 m, m, 32768; xor.b32 hi, hi, m; " + "sub.u32 i, 65535, $3; shl.b32 lo, lo, 16; or.b32 $0, lo, i; " + "sub.u32 i, i, 1; shl.b32 hi, hi, 16; or.b32 $1, hi, i; }", + "=r,=r,r,r", + loc, + ip, + ) + return _unpack(res, 2, cutlass.Int32, loc, ip) + + +@dsl_user_op +def _umax(a, b, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + T.i32(), + [_ir(a, cutlass.Int32, loc, ip), _ir(b, cutlass.Int32, loc, ip)], + "max.u32 $0, $1, $2;", + "=r,r,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _redux_umax(a, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + T.i32(), + [_ir(a, cutlass.Int32, loc, ip)], + "redux.sync.max.u32 $0, $1, 0xffffffff;", + "=r,r", + loc, + ip, + side=True, + ) + ) + + +@dsl_user_op +def _clear_eq(k, win, *, loc=None, ip=None): + """`k`, or 0 where it equals `win`.""" + return cutlass.Int32( + _asm( + T.i32(), + [_ir(k, cutlass.Int32, loc, ip), _ir(win, cutlass.Int32, loc, ip)], + "{ .reg .pred p; setp.eq.b32 p, $1, $2; selp.b32 $0, 0, $1, p; }", + "=r,r,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _sel_eq(a, b, x, y, *, loc=None, ip=None): + """`x` where `a == b`, else `y`.""" + return cutlass.Int32( + _asm( + T.i32(), + [_ir(v, cutlass.Int32, loc, ip) for v in (a, b, x, y)], + "{ .reg .pred p; setp.eq.b32 p, $1, $2; selp.b32 $0, $3, $4, p; }", + "=r,r,r,r,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _key_value(key, *, loc=None, ip=None): + """The fp32 value of a ranking key's bf16 logit.""" + return cutlass.Float32( + _asm( + T.f32(), + [_ir(key, cutlass.Int32, loc, ip)], + "{ .reg .b32 o, m; shr.b32 o, $1, 16; shr.b32 m, o, 15; mul.lo.u32 m, m, 32767; " + "xor.b32 m, m, 65535; xor.b32 o, o, m; shl.b32 $0, o, 16; }", + "=f,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _shfl_idx_f32(v, src, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [_ir(v, cutlass.Float32, loc, ip), _ir(src, cutlass.Int32, loc, ip)], + "shfl.sync.idx.b32 $0, $1, $2, 31, 0xffffffff;", + "=f,f,r", + loc, + ip, + side=True, + ) + ) + + +@dsl_user_op +def _shfl_bfly_f32(v, m, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [_ir(v, cutlass.Float32, loc, ip), _ir(m, cutlass.Int32, loc, ip)], + "shfl.sync.bfly.b32 $0, $1, $2, 31, 0xffffffff;", + "=f,f,r", + loc, + ip, + side=True, + ) + ) + + +@dsl_user_op +def _expf(x, *, loc=None, ip=None): + """libdevice `expf` (the sequence nvcc emits without fast math).""" + return cutlass.Float32( + _asm( + T.f32(), + [_ir(x, cutlass.Float32, loc, ip)], + "{ .reg .f32 a, b, c, d, e, f; .reg .b32 s; " + "fma.rn.f32 a, $1, 0f3BBB989D, 0f3F000000; cvt.sat.f32.f32 a, a; " + "fma.rm.f32 b, a, 0f437C0000, 0f4B400001; add.f32 c, b, 0fCB40007F; neg.f32 c, c; " + "fma.rn.f32 d, $1, 0f3FB8AA3B, c; fma.rn.f32 d, $1, 0f32A57060, d; " + "mov.b32 s, b; shl.b32 s, s, 23; mov.b32 f, s; " + "ex2.approx.ftz.f32 e, d; mul.f32 $0, e, f; }", + "=f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _div_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [_ir(a, cutlass.Float32, loc, ip), _ir(b, cutlass.Float32, loc, ip)], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_bits_f32(x, *, loc=None, ip=None): + """`x` rounded to bf16: (bits in the low half, the rounded value as fp32).""" + res = _asm( + llvm.StructType.get_literal([T.i32(), T.f32()]), + [_ir(x, cutlass.Float32, loc, ip)], + "{ .reg .b16 h; cvt.rn.bf16.f32 h, $2; cvt.u32.u16 $0, h; cvt.f32.bf16 $1, h; }", + "=r,=f,f", + loc, + ip, + ) + return ( + cutlass.Int32(llvm.extractvalue(T.i32(), res, [0], loc=loc, ip=ip)), + cutlass.Float32(llvm.extractvalue(T.f32(), res, [1], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def _e4m3_sf(amax, *, loc=None, ip=None): + """The NVFP4 block scale of a block whose largest magnitude is `amax`: (e4m3 byte of + amax / 6, rounded to nearest and saturated, and its value as fp32).""" + res = _asm( + llvm.StructType.get_literal([T.i32(), T.f32()]), + [_ir(amax, cutlass.Float32, loc, ip)], + "{ .reg .f32 q; .reg .b16 h, lo, hi; .reg .b32 f; " + "div.rn.f32 q, $2, 0f40C00000; cvt.rn.satfinite.e4m3x2.f32 h, 0f00000000, q; " + "cvt.u32.u16 $0, h; cvt.rn.f16x2.e4m3x2 f, h; mov.b32 {lo, hi}, f; " + "cvt.f32.f16 $1, lo; }", + "=r,=f,f", + loc, + ip, + ) + return ( + cutlass.Int32(llvm.extractvalue(T.i32(), res, [0], loc=loc, ip=ip)), + cutlass.Float32(llvm.extractvalue(T.f32(), res, [1], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def _pos_or_one(x, *, loc=None, ip=None): + """`x` if positive, else 1 (a zero block scale divides nothing: its block is all zero + codes either way).""" + return cutlass.Float32( + _asm( + T.f32(), + [_ir(x, cutlass.Float32, loc, ip)], + "{ .reg .pred p; setp.gt.f32 p, $1, 0f00000000; selp.f32 $0, $1, 0f3F800000, p; }", + "=f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _fabs(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm(T.f32(), [_ir(x, cutlass.Float32, loc, ip)], "abs.f32 $0, $1;", "=f,f", loc, ip) + ) + + +@dsl_user_op +def _fmax(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [_ir(a, cutlass.Float32, loc, ip), _ir(b, cutlass.Float32, loc, ip)], + "max.f32 $0, $1, $2;", + "=f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _shfl_down_i32(v, d, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + T.i32(), + [_ir(v, cutlass.Int32, loc, ip), _ir(d, cutlass.Int32, loc, ip)], + "shfl.sync.down.b32 $0, $1, $2, 31, 0xffffffff;", + "=r,r,r", + loc, + ip, + side=True, + ) + ) + + +@dsl_user_op +def _e2m1_lo(x, *, loc=None, ip=None): + """`x` rounded (nearest even, saturating) to e2m1: its 4-bit code.""" + return cutlass.Int32( + _asm( + T.i32(), + [_ir(x, cutlass.Float32, loc, ip)], + "{ .reg .b8 b; .reg .b16 h; cvt.rn.satfinite.e2m1x2.f32 b, 0f00000000, $1; " + "cvt.u16.u8 h, b; cvt.u32.u16 $0, h; and.b32 $0, $0, 15; }", + "=r,f", + loc, + ip, + ) + ) + + +# ---- tcgen05 / TMA helpers ----------------------------------------------------------------- + + +def _desc_asm(reg: str, saddr: str, lbo: int, sbo: int, layout: int) -> str: + """PTX building the 64-bit shared-memory matrix descriptor `reg` of the shared address in + register `saddr` (start >> 4, LBO >> 4, SBO >> 4, version 1, layout type).""" + hi = (sbo >> 4) | (1 << 14) | (layout << 29) + return ( + f"shr.b32 {reg}lo, {saddr}, 4; and.b32 {reg}lo, {reg}lo, 16383; " + f"or.b32 {reg}lo, {reg}lo, {(lbo >> 4) << 16}; mov.b32 {reg}hi, {hi}; " + f"mov.b64 {reg}, {{{reg}lo, {reg}hi}}; " + ) + + +@dsl_user_op +def _tma3d(saddr, tmap, x, y, z, mbar, *, loc=None, ip=None): + _asm( + T.i32(), + [_ir(saddr, cutlass.Int32, loc, ip), _ir(tmap, cutlass.Int64, loc, ip)] + + [_ir(v, cutlass.Int32, loc, ip) for v in (x, y, z, mbar)], + "{ cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes " + "[$1], [$2, {$3, $4, $5}], [$6]; mov.u32 $0, 0; }", + "=r,r,l,r,r,r,r", + loc, + ip, + side=True, + ) + + +@dsl_user_op +def _prefetch_tmap(tmap, *, loc=None, ip=None): + _asm( + T.i32(), + [_ir(tmap, cutlass.Int64, loc, ip)], + "{ prefetch.tensormap [$1]; mov.u32 $0, 0; }", + "=r,l", + loc, + ip, + side=True, + ) + + +@dsl_user_op +def _tc_alloc(saddr, ncols=256, *, loc=None, ip=None): + _asm( + T.i32(), + [_ir(saddr, cutlass.Int32, loc, ip)], + f"{{ tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [$1], {ncols}; " + "tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned; mov.u32 $0, 0; }", + "=r,r", + loc, + ip, + side=True, + ) + + +@dsl_user_op +def _tc_dealloc(taddr, ncols=256, *, loc=None, ip=None): + _asm( + T.i32(), + [_ir(taddr, cutlass.Int32, loc, ip)], + f"{{ tcgen05.dealloc.cta_group::1.sync.aligned.b32 $1, {ncols}; mov.u32 $0, 0; }}", + "=r,r", + loc, + ip, + side=True, + ) + + +@dsl_user_op +def _tc_fence(after, *, loc=None, ip=None): + op = "after_thread_sync" if after else "before_thread_sync" + _asm(T.i32(), [], f"{{ tcgen05.fence::{op}; mov.u32 $0, 0; }}", "=r", loc, ip, side=True) + + +@dsl_user_op +def _fence_proxy_async(*, loc=None, ip=None): + _asm(T.i32(), [], "{ fence.proxy.async.shared::cta; mov.u32 $0, 0; }", "=r", loc, ip, side=True) + + +def _make_desc(lbo: int, sbo: int, layout: int): + @dsl_user_op + def _desc(saddr, *, loc=None, ip=None): + """The 64-bit shared-memory matrix descriptor of `saddr`.""" + return cutlass.Int64( + _asm( + T.i64(), + [_ir(saddr, cutlass.Int32, loc, ip)], + "{ .reg .b32 dlo, dhi; " + + _desc_asm("d", "$1", lbo, sbo, layout).replace("mov.b64 d,", "mov.b64 $0,") + + "}", + "=l,r", + loc, + ip, + ) + ) + + return _desc + + +_desc_sf = _make_desc(0, 128, 0) # a scale atom (tcgen05.cp) +_desc_b = _make_desc(8 * 16, 128, 0) # B: the 8 rows in 16-byte core matrices +_desc_a1 = _make_desc(16, 1024, 2) # FC1 A: 128-byte swizzle +_desc_a2 = _make_desc(16, 512, 4) # FC2 A: 64-byte swizzle + + +@dsl_user_op +def _tc_cp_d(taddr, desc, off, *, loc=None, ip=None): + """tcgen05.cp of the scale atom `off` bytes past descriptor `desc` (off % 16 == 0).""" + _asm( + T.i32(), + [_ir(taddr, cutlass.Int32, loc, ip), _ir(desc, cutlass.Int64, loc, ip)], + "{ .reg .b64 d; add.s64 d, $2, " + str(off >> 4) + "; " + "tcgen05.cp.cta_group::1.32x128b.warpx4 [$1], d; mov.u32 $0, 0; }", + "=r,r,l", + loc, + ip, + side=True, + ) + + +@dsl_user_op +def _mma_d(tacc, adesc, aoff, bdesc, boff, sfa, sfb, acc, *, loc=None, ip=None): + """tcgen05 block-scaled NVFP4 MMA (M 128, N 8, K 64) on descriptors `adesc` / `bdesc` + advanced by `aoff` / `boff` bytes.""" + _asm( + T.i32(), + [ + _ir(tacc, cutlass.Int32, loc, ip), + _ir(adesc, cutlass.Int64, loc, ip), + _ir(bdesc, cutlass.Int64, loc, ip), + _ir(sfa, cutlass.Int32, loc, ip), + _ir(sfb, cutlass.Int32, loc, ip), + _ir(acc, cutlass.Int32, loc, ip), + ], + "{ .reg .pred p; .reg .b64 a, b; " + f"add.s64 a, $2, {aoff >> 4}; add.s64 b, $3, {boff >> 4}; setp.ne.b32 p, $6, 0; " + "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 " + f"[$1], a, b, {IDESC}, [$4], [$5], p; mov.u32 $0, 0; }}", + "=r,r,l,l,r,r,r", + loc, + ip, + side=True, + ) + + +@dsl_user_op +def _cp_async4g(saddr, gaddr, *, loc=None, ip=None): + _asm( + T.i32(), + [_ir(saddr, cutlass.Int32, loc, ip), _ir(gaddr, cutlass.Int64, loc, ip)], + "{ cp.async.ca.shared.global [$1], [$2], 4; mov.u32 $0, 0; }", + "=r,r,l", + loc, + ip, + side=True, + ) + + +@dsl_user_op +def _tc_commit(mbar, *, loc=None, ip=None): + _asm( + T.i32(), + [_ir(mbar, cutlass.Int32, loc, ip)], + "{ tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [$1]; " + "mov.u32 $0, 0; }", + "=r,r", + loc, + ip, + side=True, + ) + + +@dsl_user_op +def _tc_ld8(taddr, *, loc=None, ip=None): + """Eight fp32 TMEM columns of this thread's lane (32x32b shape), waited.""" + res = _asm( + llvm.StructType.get_literal([T.f32()] * 8), + [_ir(taddr, cutlass.Int32, loc, ip)], + "{ tcgen05.ld.sync.aligned.32x32b.x8.b32 {$0, $1, $2, $3, $4, $5, $6, $7}, [$8]; " + "tcgen05.wait::ld.sync.aligned; }", + "=f,=f,=f,=f,=f,=f,=f,=f,r", + loc, + ip, + side=True, + ) + return _unpack(res, 8, cutlass.Float32, loc, ip) + + +@dsl_user_op +def _ldg_nc4p(gaddr, pred, *, loc=None, ip=None): + """16 bytes from global (L2-coherent) 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.cg.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_cg32p(gaddr, pred, *, loc=None, ip=None): + 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.cg.b32 $0, [$1]; }", + "=r,l,r", + loc, + ip, + side=True, + ) + ) + + +@dsl_user_op +def _sel_f(flag, a, b, *, loc=None, ip=None): + """`a` where `flag` != 0, else `b`.""" + return cutlass.Float32( + _asm( + T.f32(), + [ + _ir(flag, cutlass.Int32, loc, ip), + _ir(a, cutlass.Float32, loc, ip), + _ir(b, cutlass.Float32, loc, ip), + ], + "{ .reg .pred p; setp.ne.b32 p, $1, 0; selp.f32 $0, $2, $3, p; }", + "=f,r,f,f", + loc, + ip, + ) + ) + + +# ------------------------------------------------------------------------------------------ +# kernels +# ------------------------------------------------------------------------------------------ + + +class _Routed: + """Routing and the work list, shared by both launches. Shared-memory routing area + (`o_rt`): slot keys' experts `[8][10]` int32, weights `[8][10]` fp32, the selected-expert + bitmap (16 words), the work list `[80]` int32, per-entry row masks `[80]` and per-entry + row slots `[80][8]` bytes.""" + + RT_IDS, RT_WTS, RT_MAP, RT_LIST, RT_MASK, RT_SLOT = 0, 320, 640, 704, 1024, 1344 + RT_BYTES = 1344 + 640 + + @cute.jit + def _route(self, lg_base, n, lane, rt, topw_base, write_topw): + """Row `n`'s top-10 and softmax (one warp): slot k on lane k; results to shared + memory, the weights also to `topw_base` when `write_topw`.""" + row = lg_base + cutlass.Int64(n) * (NUM_EXPERTS * 2) + cutlass.Int64(lane * 32) + va = _ldg_cg4(row) + vb = _ldg_cg4(row + 16) + words = va + vb + keys = [] + for i in cutlass.range_constexpr(8): + k2 = _keys2(words[i], lane * 16 + 2 * i) + keys.append(k2[0]) + keys.append(k2[1]) + mine = cutlass.Int32(0) + for r in cutlass.range_constexpr(TOP_K): + m8 = [_umax(keys[2 * i], keys[2 * i + 1]) for i in range(8)] + m4 = [_umax(m8[2 * i], m8[2 * i + 1]) for i in range(4)] + m2 = [_umax(m4[0], m4[1]), _umax(m4[2], m4[3])] + win = _redux_umax(_umax(m2[0], m2[1])) + keys = [_clear_eq(k, win) for k in keys] + mine = _sel_eq(lane, cutlass.Int32(r), win, mine) + v = _key_value(mine) + vmax = _shfl_idx_f32(v, cutlass.Int32(0)) + e = cutlass.Float32(0.0) + if lane < TOP_K: + e = _expf(v - vmax) + s = e + for m in cutlass.range_constexpr(5): + s = s + _shfl_bfly_f32(s, cutlass.Int32(16 >> m)) + if lane < TOP_K: + wbits, wf = _bf16_bits_f32(_div_rn(e, s)) + expert = 65535 - (mine & 65535) + _sts32(rt + self.RT_IDS + (n * TOP_K + lane) * 4, expert) + _sts_f32(rt + self.RT_WTS + (n * TOP_K + lane) * 4, wf) + _atom_or_shared(rt + self.RT_MAP + (expert >> 5) * 4, cutlass.Int32(1) << (expert & 31)) + if write_topw: + _stg_b16(topw_base + cutlass.Int64(n * TOP_K + lane) * 2, wbits) + + @cute.jit + def _list(self, tidx, rows, rt): + """Thread i < 10 rows: slot (i // 10, i % 10) enters the work list at its expert's + rank among the selected experts (ascending id).""" + if tidx < rows * TOP_K: + n = tidx // TOP_K + k = tidx - n * TOP_K + e = cutlass.Int32(0) + e = _lds32(rt + self.RT_IDS + tidx * 4) + wd = e >> 5 + rank = _popc(_lds32(rt + self.RT_MAP + wd * 4) & ((cutlass.Int32(1) << (e & 31)) - 1)) + for w in cutlass.range_constexpr(16): + if w < wd: + rank = rank + _popc(_lds32(rt + self.RT_MAP + w * 4)) + _sts32(rt + self.RT_LIST + rank * 4, e) + _atom_or_shared(rt + self.RT_MASK + rank * 4, cutlass.Int32(1) << n) + _sts_u8(rt + self.RT_SLOT + rank * 8 + n, k) + + @cute.jit + def _ndist(self, rt): + nd = cutlass.Int32(0) + for w in cutlass.range_constexpr(4): + q = _lds4(rt + self.RT_MAP + w * 16) + nd = nd + _popc(q[0]) + _popc(q[1]) + _popc(q[2]) + _popc(q[3]) + return nd + + @cute.jit + def _setup_rt(self, tidx, rt): + """Zero the bitmap and the row masks.""" + if tidx < 16: + _sts32(rt + self.RT_MAP + tidx * 4, cutlass.Int32(0)) + if tidx < 80: + _sts32(rt + self.RT_MASK + tidx * 4, cutlass.Int32(0)) + + +class _Fc1(_Routed): + """Launch 1: routing, FC1 on the tensor core (tcgen05, accumulated in TMEM over all of + K), SwiGLU and the requantization into the slot records.""" + + def __init__(self, *, mt, slots, grid, pdl): + assert mt in ROW_BUCKETS + self.mt, self.slots, self.grid, self.pdl = mt, slots, grid, pdl + self.o_ring = 0 + self.o_b = slots * SLOT1 + self.o_sfb = self.o_b + (K1 // 32) * NB_ROW + self.o_rt = self.o_sfb + SFB_BYTES(K1) + self.o_scal = self.o_rt + -(-self.RT_BYTES // 128) * 128 + self.o_bar = self.o_scal + MAX_ITEMS * 8 + self.o_tm = self.o_bar + (2 * slots + 4) * 8 + self.smem = self.o_tm + 16 + 1024 + assert self.smem <= _SMEM_BUDGET, (self.smem, mt, slots) + # TMEM: two 8-column accumulators, 32 weight-scale columns a slot, the 160 + # activation-scale columns. + self.sfb_col = SFA1_COL + 32 * slots + self.tcols = 256 if self.sfb_col + (K1 // 64) * 4 <= 256 else 512 + + @cute.jit + def __call__( + self, + rec_t: cute.Tensor, + lg_ptr: cute.Pointer, + xq_ptr: cute.Pointer, + xsf_ptr: cute.Pointer, + tmap_ptr: cute.Pointer, + s_ptr: cute.Pointer, + o1_ptr: cute.Pointer, + o1g_ptr: cute.Pointer, + topw_ptr: cute.Pointer, + rows: cutlass.Int32, + stream, + ): + self.kernel( + rec_t.iterator.toint(), + lg_ptr.toint(), + xq_ptr.toint(), + xsf_ptr.toint(), + tmap_ptr.toint(), + s_ptr.toint(), + o1_ptr.toint(), + o1g_ptr.toint(), + topw_ptr.toint(), + rows, + ).launch( + grid=[self.grid, 1, 1], + block=[TC_THREADS, 1, 1], + smem=self.smem, + stream=stream, + use_pdl=self.pdl, + preferred_smem_carveout=100, + ) + + @cute.kernel + def kernel( + self, + rec_base: cutlass.Int64, + lg_base: cutlass.Int64, + xq_base: cutlass.Int64, + xsf_base: cutlass.Int64, + tmap: cutlass.Int64, + s_base: cutlass.Int64, + o1_base: cutlass.Int64, + o1g_base: cutlass.Int64, + topw_base: cutlass.Int64, + rows: 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() + mt: cutlass.Constexpr = self.mt + slots: cutlass.Constexpr = self.slots + grid: cutlass.Constexpr = self.grid + + smem = cutlass.utils.SmemAllocator() + base_t = smem.allocate_tensor( + cutlass.Int8, cute.make_layout((self.smem - 1024,), stride=(1,)), byte_alignment=1024 + ) + sbase = cutlass.Int32(base_t.iterator.toint()) + ring = sbase + self.o_ring + bq = sbase + self.o_b # B: k-chunk c (16 B) of row n at c * 128 + n * 16 + sfb = sbase + self.o_sfb # SFB: k-step j's atom at j * 128, row n's 4 scales at n * 16 + rt = sbase + self.o_rt + scal = sbase + self.o_scal # [item] (output1_scale_scalar, output1_scale_gate_scalar) + full = sbase + self.o_bar + empty = full + 8 * slots + accf = empty + 8 * slots + acce = accf + 16 + tmw = sbase + self.o_tm + + if warp == 0: + if lane == 0: + for sl in cutlass.range_constexpr(slots): + _mbar_init(full + 8 * sl, 1) + _mbar_init(empty + 8 * sl, 1) + for ab in cutlass.range_constexpr(2): + _mbar_init(accf + 8 * ab, 1) + _mbar_init(acce + 8 * ab, 4) + _mbar_init_fence() + if warp == 1: + _tc_alloc(tmw, self.tcols) + if tidx == 0: + _prefetch_tmap(tmap) + self._setup_rt(tidx, rt) + _tc_fence(False) + cute.arch.sync_threads() + _tc_fence(True) + tm = _lds32(tmw) + if const_expr(self.pdl): + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + + # The activation as the MMA's B operand (8 rows, rows past `rows` zero) and its scale + # atoms, loaded before the routing so the two latencies overlap. + n_bi: cutlass.Constexpr = -(-((K1 // 32) * 8) // TC_THREADS) + n_si: cutlass.Constexpr = -(-((K1 // 64) * 8) // TC_THREADS) + bv = [] + for it in cutlass.range_constexpr(n_bi): + q = tidx + it * TC_THREADS + n = q % 8 + c = q // 8 + ok = cutlass.Int32((q < (K1 // 32) * 8) & (n < rows)) + bv.append(_ldg_nc4p(xq_base + cutlass.Int64(n) * (K1 // 2) + c * 16, ok)) + sv = [] + for it in cutlass.range_constexpr(n_si): + q = tidx + it * TC_THREADS + n = q % 8 + j = q // 8 + ok = cutlass.Int32((q < (K1 // 64) * 8) & (n < rows)) + sv.append(_ldg_cg32p(xsf_base + cutlass.Int64(n) * (K1 // 16) + j * 4, ok)) + + if warp < rows: + self._route(lg_base, warp, lane, rt, topw_base, bidx == 0) + cute.arch.sync_threads() + self._list(tidx, rows, rt) + cute.arch.sync_threads() + n_units = self._ndist(rt) * NT1 + n_items = (n_units - bidx + grid - 1) // grid + + for it in cutlass.range_constexpr(n_bi): + q = tidx + it * TC_THREADS + if q < (K1 // 32) * 8: + v = bv[it] + _sts4(bq + q * 16, v[0], v[1], v[2], v[3]) + for it in cutlass.range_constexpr(n_si): + q = tidx + it * TC_THREADS + if q < (K1 // 64) * 8: + _sts32(sfb + q * 16, sv[it]) + if tidx < n_items: + e_i = _lds32(rt + self.RT_LIST + ((bidx + grid * tidx) // NT1) * 4) + _sts_f32(scal + tidx * 8, _ldg_nc_f32(o1_base + cutlass.Int64(e_i) * 4)) + _sts_f32(scal + tidx * 8 + 4, _ldg_nc_f32(o1g_base + cutlass.Int64(e_i) * 4)) + _fence_proxy_async() + _tc_fence(False) + cute.arch.sync_threads() + _tc_fence(True) + + if warp == 0: + # ---- producer: tile (expert rank u // 10, tile u % 10), five K stages ---- + if lane == 0: + for i in cutlass.range(n_items): + u = bidx + grid * i + rank = u // NT1 + tile = u - rank * NT1 + e = _lds32(rt + self.RT_LIST + rank * 4) + for kq in cutlass.range_constexpr(KQ1): + st = i * KQ1 + kq + sl = st % slots + if st >= slots: + _mbar_wait(empty + 8 * sl, ((st // slots) - 1) & 1) + slot = ring + sl * SLOT1 + _mbar_expect_tx(full + 8 * sl, SLOT1) + _tma3d(slot, tmap, 0, e * N1 + tile * 128, 2 * kq, full + 8 * sl) + _bulk_g2s( + slot + A1_BYTES, + s_base + + cutlass.Int64(e) * (N1 * (K1 // 16)) + + tile * (128 * (K1 // 16)) + + kq * (8 * 512), + 8 * 512, + full + 8 * sl, + ) + elif warp == 1: + # ---- MMA issuer ---- + if lane == 0: + dsfb = _desc_sf(sfb) + dbq = _desc_b(bq) + for c in cutlass.range_constexpr(K1 // 64): + _tc_cp_d(tm + self.sfb_col + 4 * c, dsfb, c * 128) + for i in cutlass.range(n_items): + ab = i % 2 + if i >= 2: + _mbar_wait(acce + 8 * ab, ((i // 2) - 1) & 1) + _tc_fence(True) + for kq in cutlass.range_constexpr(KQ1): + st = i * KQ1 + kq + sl = st % slots + _mbar_wait(full + 8 * sl, (st // slots) & 1) + _tc_fence(True) + slot = ring + sl * SLOT1 + tsfa = tm + SFA1_COL + sl * 32 + dsfa = _desc_sf(slot + A1_BYTES) + da = _desc_a1(slot) + for q in cutlass.range_constexpr(8): + _tc_cp_d(tsfa + 4 * q, dsfa, q * 512) + for q in cutlass.range_constexpr(8): + _mma_d( + tm + ab * 8, + da, + (q // 4) * 16384 + (q % 4) * 32, + dbq, + (2 * (8 * kq + q)) * NB_ROW, + tsfa + 4 * q, + tm + self.sfb_col + 4 * (8 * kq + q), + cutlass.Int32(0 if kq + q == 0 else 1), + ) + _tc_commit(empty + 8 * sl) + _tc_commit(accf + 8 * ab) + elif warp < 6: + # ---- epilogue: TMEM lane quarter = one 32-row block = one requant block ---- + quarter = warp % 4 + trow = tm + ((quarter * 32) << 16) + up_row = cutlass.Int32(((lane >> 3) & 1) == 0) + for i in cutlass.range(n_items): + ab = i % 2 + _mbar_wait(accf + 8 * ab, (i // 2) & 1) + _tc_fence(True) + r = _tc_ld8(trow + ab * 8) + _tc_fence(False) + cute.arch.sync_warp() + if lane == 0: + _mbar_arrive(acce + 8 * ab) + u = bidx + grid * i + rank = u // NT1 + b = (u - rank * NT1) * 4 + quarter + s_up = _lds_f32(scal + i * 8) + s_gt = _lds_f32(scal + i * 8 + 4) + tmask = _lds32(rt + self.RT_MASK + rank * 4) + for n in cutlass.range_constexpr(mt): + # Unit row d = lane: up rows have bit 3 clear, their gate row is d + 8. + gt = _shfl_bfly_f32(r[n], cutlass.Int32(8)) * s_gt + act = (r[n] * s_up) * _div_rn(gt, 1.0 + _expf(-gt)) + amax = _sel_f(up_row, _fabs(act), 0.0) + for sh in cutlass.range_constexpr(5): + amax = _fmax(amax, _shfl_bfly_f32(amax, cutlass.Int32(16 >> sh))) + sfb_, sff = _e4m3_sf(amax) + nib = _e2m1_lo(_div_rn(act, _pos_or_one(sff))) + # Intermediate column j = 2 (d % 8) + d // 16: byte d % 8, high nibble + # from rows 16 .. 23. + byte = nib | (_shfl_down_i32(nib, cutlass.Int32(16)) << 4) + w = ( + byte + | (_shfl_down_i32(byte, cutlass.Int32(1)) << 8) + | (_shfl_down_i32(byte, cutlass.Int32(2)) << 16) + | (_shfl_down_i32(byte, cutlass.Int32(3)) << 24) + ) + w1 = _shfl_down_i32(w, cutlass.Int32(4)) + if (lane == 0) & (n < rows) & (((tmask >> n) & 1) == 1): + slot_k = _lds_u8(rt + self.RT_SLOT + rank * 8 + n) + recp = rec_base + cutlass.Int64(n * TOP_K + slot_k) * REC + _stg_v2(recp + b * 8, w, w1) + _stg_u8(recp + 320 + b, sfb_) + _tc_fence(False) + cute.arch.sync_threads() + if warp == 1: + _tc_dealloc(tm, self.tcols) + + +class _Fc2(_Routed): + """Launch 2: routing, its FC2 tiles and their scales into shared memory / TMEM before the + wait; after it the B operand from the slot records, the MMAs and the bf16 rows.""" + + def __init__(self, *, mt, slots, grid, pdl): + assert mt in ROW_BUCKETS + self.mt, self.slots, self.grid, self.pdl = mt, slots, grid, pdl + # Batch mode (4+ slots: every unit of the CTA in shared memory before the wait, the CTA + # alone on its SM): one B buffer, accumulator and TMEM scale range per unit, all 512 + # TMEM columns; otherwise two of each, double-buffered. + self.batch = slots >= 4 + self.nb = slots if self.batch else 2 + self.tcols = 512 if self.batch else TMEM_COLS + self.o_ring = 0 + self.o_b = slots * SLOT2 + self.o_rt = self.o_b + self.nb * B2_BYTES + self.o_s2 = self.o_rt + -(-self.RT_BYTES // 128) * 128 + self.o_bar = self.o_s2 + MAX_ITEMS * 4 + self.o_tm = self.o_bar + (2 * slots + 4 + 4 * self.nb) * 8 + self.smem = self.o_tm + 16 + 1024 + self.acc_cols = 8 * self.nb + self.sfa_col = self.acc_cols + self.sfb_col = self.sfa_col + 40 * slots + assert self.sfb_col + 40 * self.nb <= self.tcols, slots + assert self.smem <= _SMEM_BUDGET, (self.smem, mt, slots) + + @cute.jit + def __call__( + self, + out_t: cute.Tensor, + lg_ptr: cute.Pointer, + rec_ptr: cute.Pointer, + tmap_ptr: cute.Pointer, + s_ptr: cute.Pointer, + o2_ptr: cute.Pointer, + w_ptr: cute.Pointer, + rows: cutlass.Int32, + stream, + ): + self.kernel( + out_t.iterator.toint(), + lg_ptr.toint(), + rec_ptr.toint(), + tmap_ptr.toint(), + s_ptr.toint(), + o2_ptr.toint(), + w_ptr.toint(), + rows, + ).launch( + grid=[self.grid, 1, 1], + block=[TC_THREADS, 1, 1], + smem=self.smem, + stream=stream, + use_pdl=self.pdl, + preferred_smem_carveout=100, + ) + + @cute.kernel + def kernel( + self, + out_base: cutlass.Int64, + lg_base: cutlass.Int64, + rec_base: cutlass.Int64, + tmap: cutlass.Int64, + s_base: cutlass.Int64, + o2_base: cutlass.Int64, + w_base: cutlass.Int64, + rows: 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() + mt: cutlass.Constexpr = self.mt + slots: cutlass.Constexpr = self.slots + grid: cutlass.Constexpr = self.grid + + smem = cutlass.utils.SmemAllocator() + base_t = smem.allocate_tensor( + cutlass.Int8, cute.make_layout((self.smem - 1024,), stride=(1,)), byte_alignment=1024 + ) + sbase = cutlass.Int32(base_t.iterator.toint()) + ring = sbase + self.o_ring + bb = sbase + self.o_b # [2] (B 2560 B, SFB atoms) + rt = sbase + self.o_rt + s2b = sbase + self.o_s2 # [item] output2_scale_scalar + nb: cutlass.Constexpr = self.nb + full = sbase + self.o_bar + empty = full + 8 * slots + accf = empty + 8 * slots + acce = accf + 8 * nb + bready = acce + 8 * nb + bfree = bready + 8 * nb + tmw = sbase + self.o_tm + + if warp == 0: + if lane == 0: + for sl in cutlass.range_constexpr(slots): + _mbar_init(full + 8 * sl, 1) + _mbar_init(empty + 8 * sl, 1) + for pb in cutlass.range_constexpr(nb): + _mbar_init(accf + 8 * pb, 1) + _mbar_init(acce + 8 * pb, 4) + _mbar_init(bready + 8 * pb, 1) + _mbar_init(bfree + 8 * pb, 1) + _mbar_init_fence() + if warp == 1: + _tc_alloc(tmw, self.tcols) + if tidx == 0: + _prefetch_tmap(tmap) + self._setup_rt(tidx, rt) + _tc_fence(False) + cute.arch.sync_threads() + _tc_fence(True) + tm = _lds32(tmw) + + # Routing before the wait: the logits are complete (launch 1 released this grid + # only after its own wait on the router GEMM), and are read through L2. + if warp < rows: + self._route(lg_base, warp, lane, rt, lg_base, False) + cute.arch.sync_threads() + self._list(tidx, rows, rt) + cute.arch.sync_threads() + n_units = self._ndist(rt) * NT2 + n_items = (n_units - bidx + grid - 1) // grid + if (warp >= 2) & (tidx - 64 < n_items): + e_i = _lds32(rt + self.RT_LIST + ((bidx + grid * (tidx - 64)) // NT2) * 4) + _sts_f32(s2b + (tidx - 64) * 4, _ldg_nc_f32(o2_base + cutlass.Int64(e_i) * 4)) + + if warp == 0: + # ---- producer: the first `slots` tiles now, L2 prefetch of the rest ---- + if lane == 0: + for i in cutlass.range(n_items): + u = bidx + grid * i + rank = u // NT2 + tile = u - rank * NT2 + e = _lds32(rt + self.RT_LIST + rank * 4) + sl = i % slots + sa = s_base + cutlass.Int64(e) * (N2 * (K2 // 16)) + tile * (128 * (K2 // 16)) + if i < slots: + slot = ring + sl * SLOT2 + _mbar_expect_tx(full + 8 * sl, SLOT2) + _tma3d(slot, tmap, 0, e * N2 + tile * 128, 0, full + 8 * sl) + _bulk_g2s(slot + A2_BYTES, sa, 10 * 512, full + 8 * sl) + else: + _bulk_prefetch_l2( + w_base + (cutlass.Int64(e) * N2 + tile * 128) * (K2 // 2), + 128 * (K2 // 2), + ) + _bulk_prefetch_l2(sa, 10 * 512) + for i in cutlass.range(slots, n_items): + u = bidx + grid * i + rank = u // NT2 + tile = u - rank * NT2 + e = _lds32(rt + self.RT_LIST + rank * 4) + sl = i % slots + _mbar_wait(empty + 8 * sl, ((i // slots) - 1) & 1) + slot = ring + sl * SLOT2 + _mbar_expect_tx(full + 8 * sl, SLOT2) + _tma3d(slot, tmap, 0, e * N2 + tile * 128, 0, full + 8 * sl) + _bulk_g2s( + slot + A2_BYTES, + s_base + cutlass.Int64(e) * (N2 * (K2 // 16)) + tile * (128 * (K2 // 16)), + 10 * 512, + full + 8 * sl, + ) + elif warp == 1: + # ---- MMA issuer: weight scales to TMEM as tiles land (before the wait too) ---- + if lane == 0: + # The weight scales of the units already in their slots, before the wait. + for i in cutlass.range(cutlass.min(n_items, slots)): + slot = ring + i * SLOT2 + _mbar_wait(full + 8 * i, 0) + _tc_fence(True) + dsfa = _desc_sf(slot + A2_BYTES) + for q in cutlass.range_constexpr(K2 // 64): + _tc_cp_d(tm + self.sfa_col + i * 40 + 4 * q, dsfa, q * 512) + for i in cutlass.range(n_items): + sl = i % slots + pb = i % nb + slot = ring + sl * SLOT2 + _mbar_wait(full + 8 * sl, (i // slots) & 1) + _tc_fence(True) + tsfa = tm + self.sfa_col + sl * 40 + if i >= slots: + dsfa = _desc_sf(slot + A2_BYTES) + for q in cutlass.range_constexpr(K2 // 64): + _tc_cp_d(tsfa + 4 * q, dsfa, q * 512) + if const_expr(self.batch): + _mbar_wait(bready, 0) + else: + _mbar_wait(bready + 8 * pb, (i // 2) & 1) + _tc_fence(True) + tsfb = tm + self.sfb_col + pb * 40 + bp = bb + pb * B2_BYTES + dsfb = _desc_sf(bp + (K2 // 32) * NB_ROW) + dbp = _desc_b(bp) + da = _desc_a2(slot) + for q in cutlass.range_constexpr(K2 // 64): + _tc_cp_d(tsfb + 4 * q, dsfb, q * 128) + if const_expr(not self.batch): + if i >= 2: + _mbar_wait(acce + 8 * pb, ((i // 2) - 1) & 1) + _tc_fence(True) + for q in cutlass.range_constexpr(K2 // 64): + _mma_d( + tm + pb * 8, + da, + (q // 2) * 8192 + (q % 2) * 32, + dbp, + (2 * q) * NB_ROW, + tsfa + 4 * q, + tsfb + 4 * q, + cutlass.Int32(0 if q == 0 else 1), + ) + if const_expr(not self.batch): + _tc_commit(empty + 8 * sl) + _tc_commit(bfree + 8 * pb) + _tc_commit(accf + 8 * pb) + elif warp >= 6: + # ---- B builders: B zeroed before the wait; after it each valid row's record + # pieces (20 x 16-byte code chunks, 10 scale words) by cp.async from global ---- + ptid = tidx - 192 + for it in cutlass.range_constexpr(-(-(nb * B2_BYTES // 16) // 64)): + q = ptid + it * 64 + if q < nb * B2_BYTES // 16: + z = cutlass.Int32(0) + _sts4(bb + q * 16, z, z, z, z) + _bar_sync(2, 64) + if const_expr(self.pdl): + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + n_pc: cutlass.Constexpr = mt * ((K2 // 32) + (K2 // 64)) + for i in cutlass.range(n_items): + pb = i % nb + u = bidx + grid * i + rank = u // NT2 + tmask = _lds32(rt + self.RT_MASK + rank * 4) + if const_expr(not self.batch): + if i >= 2: + _mbar_wait(bfree + 8 * pb, ((i // 2) - 1) & 1) + for it in cutlass.range_constexpr(-(-(B2_BYTES // 16) // 64)): + q = ptid + it * 64 + if q < B2_BYTES // 16: + z = cutlass.Int32(0) + _sts4(bb + pb * B2_BYTES + q * 16, z, z, z, z) + _bar_sync(2, 64) + bp = bb + pb * B2_BYTES + for it in cutlass.range_constexpr(-(-n_pc // 64)): + q = ptid + it * 64 + n = q // ((K2 // 32) + (K2 // 64)) + pc = q - n * ((K2 // 32) + (K2 // 64)) + if (q < n_pc) & (n < rows): + if ((tmask >> n) & 1) == 1: + sk = _lds_u8(rt + self.RT_SLOT + rank * 8 + n) + rec = rec_base + cutlass.Int64(n * TOP_K + sk) * REC + if pc < K2 // 32: + _cp_async16(bp + pc * NB_ROW + n * 16, rec + pc * 16) + else: + j = pc - K2 // 32 + _cp_async4g( + bp + (K2 // 32) * NB_ROW + j * 128 + n * 16, rec + 320 + j * 4 + ) + if const_expr(not self.batch): + _cp_async_wait_all() + _fence_proxy_async() + _bar_sync(2, 64) + if ptid == 0: + _mbar_arrive(bready + 8 * pb) + if const_expr(self.batch): + _cp_async_wait_all() + _fence_proxy_async() + _bar_sync(2, 64) + if ptid == 0: + _mbar_arrive(bready) + else: + # ---- epilogue: TMEM lane quarter = 32 output rows ---- + if const_expr(self.pdl): + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + quarter = warp % 4 + trow = tm + ((quarter * 32) << 16) + for i in cutlass.range(n_items): + pb = i % nb + _mbar_wait(accf + 8 * pb, (i // nb) & 1) + _tc_fence(True) + r = _tc_ld8(trow + pb * 8) + _tc_fence(False) + cute.arch.sync_warp() + if lane == 0: + _mbar_arrive(acce + 8 * pb) + u = bidx + grid * i + rank = u // NT2 + b = (u - rank * NT2) * 4 + quarter + # Unit row d = lane is output column 32 b + 4 (d % 8) + d // 8. + col = b * BLK + 4 * (lane % 8) + lane // 8 + s2 = _lds_f32(s2b + i * 4) + tmask = _lds32(rt + self.RT_MASK + rank * 4) + for n in cutlass.range_constexpr(mt): + if (n < rows) & (((tmask >> n) & 1) == 1): + sk = _lds_u8(rt + self.RT_SLOT + rank * 8 + n) + _stg_b16( + out_base + (cutlass.Int64(n * TOP_K + sk) * N2 + col) * 2, + _bf16_bits_f32(r[n] * s2)[0], + ) + _tc_fence(False) + cute.arch.sync_threads() + if warp == 1: + _tc_dealloc(tm, self.tcols) + + +# ------------------------------------------------------------------------------------------ +# host side +# ------------------------------------------------------------------------------------------ + + +def _bucket(rows: int) -> int: + for b in ROW_BUCKETS: + if rows <= b: + return b + raise AssertionError(f"{rows} rows > MAX_ROWS ({MAX_ROWS})") + + +#: Launch plan by row bucket: (launch 1 slots, launch 2 slots, launch 1 grid, launch 2 grid). +#: One FC1 stage (36 KB) or FC2 tile (45 KB) a slot. At one row the FC1 work is exactly 100 +#: tiles (10 experts x 10), so launch 1 takes 100 SMs and launch 2 runs two CTAs on each of +#: the other 52, where it routes and loads all of its FC2 tiles while launch 1 runs; a grid +#: of 0 means one CTA per SM. +_PLANS = {1: (5, 2, 100, 104), 2: (5, 2, 0, 0), 4: (5, 2, 0, 0), 8: (4, 2, 0, 0)} + + +def _dl(t: torch.Tensor): + return from_dlpack(t, assumed_align=16).mark_layout_dynamic(leading_dim=t.dim() - 1) + + +def _gptr(dtype, buf, align=16): + return make_ptr(dtype, buf.data_ptr(), cute.AddressSpace.gmem, assumed_align=align) + + +@functools.lru_cache(maxsize=None) +def _compiled(mt: int, slots1: int, slots2: int, grid1: int, grid2: int, pdl: bool): + """(launch 1, launch 2) for row bucket `mt`.""" + buf = torch.empty(64, dtype=torch.int32, device="cuda") + rec = torch.empty(mt * TOP_K, REC, dtype=torch.uint8, device="cuda") + out = torch.empty(mt * TOP_K, N2, dtype=torch.bfloat16, device="cuda") + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + f1 = cute.compile( + _Fc1(mt=mt, slots=slots1, grid=grid1, pdl=pdl), + _dl(rec), + _gptr(cutlass.BFloat16, buf), + _gptr(cutlass.Uint8, buf), + _gptr(cutlass.Uint8, buf), + _gptr(cutlass.Uint8, buf), + _gptr(cutlass.Uint8, buf), + _gptr(cutlass.Float32, buf), + _gptr(cutlass.Float32, buf), + _gptr(cutlass.BFloat16, buf), + cutlass.Int32(1), + stream, + options=_COMPILE_OPTIONS, + ) + f2 = cute.compile( + _Fc2(mt=mt, slots=slots2, grid=grid2, pdl=pdl), + _dl(out), + _gptr(cutlass.BFloat16, buf), + _gptr(cutlass.Uint8, buf), + _gptr(cutlass.Uint8, buf), + _gptr(cutlass.Uint8, buf), + _gptr(cutlass.Float32, buf), + _gptr(cutlass.Uint8, buf), + cutlass.Int32(1), + stream, + options=_COMPILE_OPTIONS, + ) + return f1, f2 + + +def _plan_key(mt: int, pdl: bool) -> tuple: + slots1, slots2, grid1, grid2 = _PLANS[mt] + sms = _sms() + grid1 = grid1 if grid1 > 0 else sms + grid2 = sms + grid2 if grid2 < 0 else (grid2 or sms) + return (mt, slots1, slots2, grid1, grid2, bool(pdl and _env_pdl())) + + +def precompile(max_rows: int = MAX_ROWS, pdl: bool = True) -> None: + """Compile every row bucket up to `max_rows` now: a first call mid-serving would stall on + the compile.""" + for mt in ROW_BUCKETS: + if mt <= _bucket(max_rows): + _compiled(*_plan_key(mt, pdl)) + + +#: TMA descriptors of registered expert stacks, 128 bytes each in one device buffer (the +#: kernels read them from global memory), keyed by (data_ptr, operand). +_TMAP_ROWS = 1024 +_tmaps: dict = {} +_tmap_buf: list = [] + + +def _encode(ptr: int, fc: int) -> bytes: + from cuda.bindings import driver as cu + + if fc == 1: # FC1: (128 B, rows, 10 chunks), box (128, 128, 2), 128-byte swizzle + ch, rows, kp, row_bytes, box_kp = 128, NUM_EXPERTS * N1, 10, K1 // 2, 2 + sw = cu.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B + else: # FC2: (64 B, rows, 5 chunks), box (64, 128, 5), 64-byte swizzle + ch, rows, kp, row_bytes, box_kp = 64, NUM_EXPERTS * N2, 5, K2 // 2, 5 + sw = cu.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_64B + err, tm = cu.cuTensorMapEncodeTiled( + cu.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT8, + 3, + ptr, + [cu.cuuint64_t(ch), cu.cuuint64_t(rows), cu.cuuint64_t(kp)], + [cu.cuuint64_t(row_bytes), cu.cuuint64_t(ch)], + [cu.cuuint32_t(ch), cu.cuuint32_t(128), cu.cuuint32_t(box_kp)], + [cu.cuuint32_t(1)] * 3, + cu.CUtensorMapInterleave.CU_TENSOR_MAP_INTERLEAVE_NONE, + sw, + cu.CUtensorMapL2promotion.CU_TENSOR_MAP_L2_PROMOTION_L2_256B, + cu.CUtensorMapFloatOOBfill.CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE, + ) + assert int(err) == 0, f"cuTensorMapEncodeTiled: {err}" + return b"".join(int(v).to_bytes(8, "little") for v in tm.opaque) + + +def _tmap(t: torch.Tensor, fc: int) -> int: + """Device address of `t`'s TMA descriptor (built and uploaded on first use, which must not + be inside a CUDA-graph capture: `register_weights` does it ahead).""" + key = (t.data_ptr(), fc) + row = _tmaps.get(key) + if row is None: + assert not torch.cuda.is_current_stream_capturing(), ( + "nvfp4_moe_gemv: expert stack not registered before CUDA-graph capture " + "(call register_weights at load)" + ) + if not _tmap_buf: + _tmap_buf.append(torch.zeros(_TMAP_ROWS, 128, dtype=torch.uint8, device=t.device)) + row = len(_tmaps) + assert row < _TMAP_ROWS, "TMA descriptor table full" + src = torch.frombuffer(bytearray(_encode(t.data_ptr(), fc)), dtype=torch.uint8) + _tmap_buf[0][row].copy_(src) + torch.cuda.synchronize(t.device) + _tmaps[key] = row + return _tmap_buf[0].data_ptr() + row * 128 + + +def register_weights(gemm1_weights: torch.Tensor, gemm2_weights: torch.Tensor) -> None: + """Build the TMA descriptors of one layer's expert stacks (trtllm-gen layout) now.""" + _tmap(gemm1_weights, 1) + _tmap(gemm2_weights, 2) + + +@functools.lru_cache(maxsize=None) +def _topk_rows(rows: int, device: torch.device) -> torch.Tensor: + return torch.arange(rows * TOP_K, dtype=torch.int32, device=device).view(rows, TOP_K) + + +def nvfp4_moe_gemv( + 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, + pdl: bool = True, +) -> tuple: + """`(rows [10 T, 2560] bf16, topk_weights [T, 10] bf16, topk_rows [T, 10] int32)`; see + the module docstring. Every operand dense, the expert operands in the trtllm-gen layout.""" + rows = routing_logits.shape[0] + dev = routing_logits.device + f1, f2 = _compiled(*_plan_key(_bucket(rows), pdl)) + t1 = _tmap(gemm1_weights, 1) + t2 = _tmap(gemm2_weights, 2) + rec = torch.empty(rows * TOP_K, REC, dtype=torch.uint8, device=dev) + out = torch.empty(rows * TOP_K, N2, dtype=torch.bfloat16, device=dev) + topw = torch.empty(rows, TOP_K, dtype=torch.bfloat16, device=dev) + lg = routing_logits.data_ptr() + f1( + rec, + lg, + hidden_states.data_ptr(), + hidden_states_scale.data_ptr(), + t1, + gemm1_weights_scale.data_ptr(), + output1_scale_scalar.data_ptr(), + output1_scale_gate_scalar.data_ptr(), + topw.data_ptr(), + rows, + ) + f2( + out, + lg, + rec.data_ptr(), + t2, + gemm2_weights_scale.data_ptr(), + output2_scale_scalar.data_ptr(), + gemm2_weights.data_ptr(), + rows, + ) + return out, topw, _topk_rows(rows, dev) + + +class Nvfp4MoeGemv(OpWrapper): + """`fp4_block_scale_moe_renorm_permuted` for 1 .. MAX_ROWS rows, in two PDL launches. + + Same arguments and the same three outputs: `rows` `[10 T, H]` bf16 (slot `(t, k)` at row + `10 t + k`), `topk_weights` `[T, 10]` bf16 in the runner's slot order, `topk_rows` + `[T, 10]` int32. `finalize(*outputs)` is the layer's combined output. + + Fusion boundary: Renormalize routing, FC1, SwiGLU, the intermediate's NVFP4 + requantization and FC2; the finalize is the consumer's (the hyper-connection fold). + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="STP decode, one row (con=1): the 1-row bucket", spec=_ckpt(1)), + Cell(why="2 rows: the 2-row bucket", spec=_ckpt(2)), + Cell(why="3 rows: a partial 4-row bucket", spec=_ckpt(3)), + Cell(why="MTP3 verify of one request (4 rows): the 4-row bucket", spec=_ckpt(4)), + Cell(why="8 rows: the largest bucket", spec=_ckpt(8)), + Cell(why="coarse 3-level logits: ties decide most slots", spec=_ckpt(4, dist="coarse")), + ) + + note = """ + Metadata consumed: none. Call site: `DecodeTarget._moe_block` under + `TRTLLM_MODELING_V2_MOE_GEMV=1` at the rows in `_MOE_GEMV_ROWS` (1 .. 4) when the finalize + folds into the next site. Preconditions: every operand contiguous; the expert stacks in + the trtllm-gen layout (FC1 rows interleaved and 32-row block shuffled, FC2 rows block + shuffled, both scale stacks 128x4 swizzled) and registered (`register_weights`) before a + CUDA-graph capture; `hidden_states_scale` the linear e4m3 layout. Two PDL launches; + CUDA-graph capturable; one compile per row bucket (`precompile` builds them all), the row + count a launch argument. + + Measured on GB300, us a layer, 48 distinct layers each behind its router projection in + one CUDA graph, rows correlated 0.85 (runner -> this): 14.42 -> 12.64 at 1 row, 20.59 -> + 17.99 at 2, 23.37 -> 22.72 at 4, 28.41 -> 31.59 at 8 (slower: not dispatched). + + Numerics: the routing is the runner's (bitwise ids and weights). The products of e2m1 + codes and e4m3 scales are exact and accumulate in fp32 on the tensor core; the SwiGLU is + `up * (gate / (1 + expf(-gate)))` in fp32 and the intermediate is requantized as + `quantize_nvfp4` defines it. Measured bitwise equal to the runner's slot rows on the + cells and on real checkpoint layers. + + Hard limits + T 1 .. MAX_ROWS (8); buckets 1, 2, 4, 8 + geometry 512 experts, top-10, hidden 2560, intermediate 640 + TMA `_TMAP_ROWS` registered expert stacks per process + """ + + 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: + return nvfp4_moe_gemv( + 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, + ) + + 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 _routed.reference(*args, **kwargs) + + 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: + 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; the kernels address it from data_ptr() as dense" + ) + assert (num_experts, top_k) in CERTIFIED_PAIRS and num_experts == NUM_EXPERTS + assert intermediate_size == INTER and hidden_states.shape[1] * 2 == HIDDEN + assert routing_logits.dtype == torch.bfloat16, "bf16 router logits" + assert 1 <= routing_logits.shape[0] <= MAX_ROWS, f"1 .. {MAX_ROWS} rows" + + +nvfp4_moe_gemv_op = Nvfp4MoeGemv() + +__all__ = ["finalize", "nvfp4_moe_gemv", "nvfp4_moe_gemv_op", "precompile", "MAX_ROWS"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py index fb79cf780dfd..85dd83a2f7c1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py @@ -9,6 +9,7 @@ import torch import triton +import triton.language as tl # Registration is lazy: the op is defined beside the mamba/GDN mixer, and # `import tensorrt_llm._torch.custom_ops` alone does not make @@ -24,6 +25,11 @@ #: programs. _PREFILL_MIN_ROWS = 32768 # capacity: fixed: a dispatch threshold, not a buffer size _PREFILL_ROWS = 8 +#: The amax copy's program: 64 rows on 8 warps. At 16384 tokens on GB300 108 vs 122 us +#: at 8 rows on 4 warps (bitwise-identical output and maximum; 16 rows on 4 warps 118, +#: 64 rows on 4 or 16 warps 145 / 126, 128 rows on 8 or 16 warps 371 / 180). +_AMAX_ROWS = 64 +_AMAX_WARPS = 8 def _prefill_multirow(x: torch.Tensor, z: torch.Tensor, gate_activation: str) -> bool: @@ -42,10 +48,64 @@ def _prefill_multirow(x: torch.Tensor, z: torch.Tensor, gate_activation: str) -> ) -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).""" +@triton.jit +def _rms_norm_gated_amax_kernel( + X, + Y, + W, + Z, + PARTS, + stride_x_row, + stride_y_row, + stride_z_tok, + M, + eps, + NUM_PARTS, + N: tl.constexpr, + ROWS: tl.constexpr, + HEADS_PER_TOK: tl.constexpr, + GATE_SIGMOID: tl.constexpr, +): + """`_rms_norm_gated_fwd_multirow_kernel`'s arithmetic, plus the maximum of the stored + `|y|` folded into `PARTS[program % NUM_PARTS]`.""" + rows = tl.program_id(0) * ROWS + tl.arange(0, ROWS) + row_mask = rows < M + cols = tl.arange(0, N) + mask2d = row_mask[:, None] + x_off = rows[:, None].to(tl.int64) * stride_x_row + cols[None, :] + x = tl.load(X + x_off, mask=mask2d, other=0.0).to(tl.float32) + var = tl.sum(x * x, axis=1) / N + rstd = 1.0 / tl.sqrt(var + eps) + w = tl.load(W + cols).to(tl.float32) + y = x * rstd[:, None] * w[None, :] + tok = rows // HEADS_PER_TOK + head = rows % HEADS_PER_TOK + z_off = tok[:, None].to(tl.int64) * stride_z_tok + head[:, None] * N + cols[None, :] + z = tl.load(Z + z_off, mask=mask2d, other=0.0).to(tl.float32) + gate = tl.sigmoid(z) + if not GATE_SIGMOID: + gate *= z + y *= gate + yq = y.to(Y.dtype.element_ty) + y_off = rows[:, None].to(tl.int64) * stride_y_row + cols[None, :] + tl.store(Y + y_off, yq, mask=mask2d) + peak = tl.max(tl.max(tl.where(mask2d, tl.abs(yq.to(tl.float32)), 0.0), axis=1), axis=0) + tl.atomic_max(PARTS + tl.program_id(0) % NUM_PARTS, peak) + + +def _prefill_launch(x, z, weight, eps, gate_activation, amax_parts=None): + """The multi-row kernel at `_PREFILL_ROWS` rows a program (no PDL: the grid covers every + SM); with `amax_parts` the target's copy that also folds the output's maximum, at + `_AMAX_ROWS` rows a program.""" rows, n = x.shape out = torch.empty_like(x) + if amax_parts is not None: + _rms_norm_gated_amax_kernel[(triton.cdiv(rows, _AMAX_ROWS),)]( + x, out, weight, z, amax_parts, x.stride(0), out.stride(0), z.stride(0), rows, eps, + amax_parts.shape[0], N=n, ROWS=_AMAX_ROWS, HEADS_PER_TOK=z.shape[1], + GATE_SIGMOID=gate_activation == "sigmoid", num_warps=_AMAX_WARPS, + ) # fmt: skip + return out _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", @@ -71,8 +131,12 @@ class RmsNormGatedTokenMajor(OpWrapper): 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. + With `amax_parts` (a zeroed fp32 buffer, `fp8_per_tensor_mm.amax_partials`) the + maximum of the stored `|out|` is also folded into it (atomic maxima, one slot per + program modulo its length), for the per-tensor FP8 projection that consumes `out`. + + Fusion boundary: norm, weight, gate and the optional quantize (or output maximum). + The caller owns both projections producing `x` and `z`, and every reshape. """ ARCHS = frozenset({Arch.SM_103}) @@ -90,11 +154,16 @@ class RmsNormGatedTokenMajor(OpWrapper): why="683 tokens: the first launch at 8 rows a program", spec=dict(tokens=683, z="in_proj", **_CKPT), ), + Cell( + why="prefill 16384 tokens folding the output maximum for the FP8 out projection", + spec=dict(tokens=16384, z="in_proj", amax=True, **_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`, + program; 8 from 32768 rows, without fp8 output; its amax copy 64 rows on 8 + warps) 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 @@ -136,12 +205,16 @@ def __call__( eps: float, fp8_scale: Optional[torch.Tensor] = None, gate_activation: str = "silu", + amax_parts: Optional[torch.Tensor] = None, ) -> 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( + return _prefill_launch(x, z, weight, eps, gate_activation, amax_parts) + out = torch.ops.trtllm.rms_norm_gated_token_major( x, z, weight, eps, fp8_scale, gate_activation ) + if amax_parts is not None: + amax_parts[:1].copy_(torch.reshape(torch.abs(out).amax().float(), [1])) + return out def reference( self, @@ -151,6 +224,7 @@ def reference( eps: float, fp8_scale: Optional[torch.Tensor] = None, gate_activation: str = "silu", + amax_parts: Optional[torch.Tensor] = None, ) -> torch.Tensor: """fp32 throughout; row `t * heads + h` pairs with `z[t, h]`. @@ -177,7 +251,11 @@ def is_valid( eps: float, fp8_scale: Optional[torch.Tensor] = None, gate_activation: str = "silu", + amax_parts: Optional[torch.Tensor] = None, ) -> None: + assert amax_parts is None or ( + fp8_scale is None and amax_parts.dtype == torch.float32 and amax_parts.dim() == 1 + ), "amax_parts: a 1-D fp32 buffer, bf16 output only" 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] " diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_gate_conv_inject_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_gate_conv_inject_norm.py index 6afb43f1c077..bb2814cac29b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_gate_conv_inject_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_gate_conv_inject_norm.py @@ -333,8 +333,12 @@ class PleGateConvInjectNorm(OpWrapper): launch: the conv over packed (row tile, channel block) programs, the per-sequence kernel then only publishing the state. - Returns `(hc_out, normed)`, fresh `[T, C*D]` in `hc`'s dtype. - `conv_states` is advanced in place. + 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 @@ -348,6 +352,7 @@ class PleGateConvInjectNorm(OpWrapper): # 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( @@ -371,11 +376,10 @@ class PleGateConvInjectNorm(OpWrapper): distinct live slots; `query_start_loc[-1] == T`. Operands 2-D with unit column stride; weights contiguous. Certified at the checkpoint's 4 x 2560 bf16 geometry with every slot live (the target never passes `pad_slot_id`), - at a 128 x 1-row decode, an 8192-row fresh prefill and a mixed 2-context + - 64-decode step, all on the 461-slot pool. Not certified: more than one - row per sequence on the per-sequence path (`1 < max_rows <= - PER_SEQ_MAX_ROWS`) or sequences shorter than `max_rows` -- a verify step - takes the unfused ops, so the target's per-sequence path runs one row. + 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 @@ -408,18 +412,27 @@ def __call__( 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] - hc_out = torch.empty((rows, width), dtype=hc.dtype, device=hc.device) - normed = torch.empty((rows, width), dtype=hc.dtype, device=hc.device) + 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 = 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)]( @@ -534,6 +547,8 @@ def reference( 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 @@ -610,6 +625,8 @@ def is_valid( 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}" @@ -652,6 +669,15 @@ def is_valid( 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 index dea3bb7c4136..f8c2236854cf 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_ngram_embed.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_ngram_embed.py @@ -328,6 +328,13 @@ class PleNgramEmbed(OpWrapper): 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", @@ -356,12 +363,10 @@ class PleNgramEmbed(OpWrapper): - token id * max multiplier < 2^63 (int64 wraps, as in `ngram_hash`). Certified at a 128 x 1-row decode on a host-resident table and on the - full checkpoint-height device table, and a mixed step of 2 contexts (one - fresh, one carrying) + 64 decode rows on the packed-row form. Not - certified: more than one row per sequence on the per-sequence path - (`1 < max_rows <= PER_SEQ_MAX_ROWS`) or zero-row sequences -- a verify - step takes the unfused ops, so the target's per-sequence path runs one - row. + 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. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/flashinfer_chunk_gated_delta_rule.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/flashinfer_chunk_gated_delta_rule.py index e6a503da7c0f..c0f7de9908f2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/flashinfer_chunk_gated_delta_rule.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/flashinfer_chunk_gated_delta_rule.py @@ -7,12 +7,16 @@ (tier-2 sourcing), the adapter in front of the installed FlashInfer's ``flashinfer.chunk_gated_delta_rule`` (CuTe-DSL tcgen05 kernel on sm_100/103). One call is: ``g -> exp(g)``, a Triton gather of the indexed pool slots into a -packed fp32 ``[N, HV, V, K]`` buffer, the FlashInfer kernel, and a Triton -scatter of the final states back into the same slots. Call-compatible with the -vendored Triton ``chunk_gated_delta_rule`` (``modules/fla/chunk``). +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. +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 @@ -20,11 +24,14 @@ 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. @@ -39,6 +46,93 @@ def _engine_slots(num_seqs: int) -> tuple[int, ...]: 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)`. @@ -51,20 +145,25 @@ class FlashinferChunkGatedDeltaRule(OpWrapper): computed 64 tokens at a time on tensor cores (WY form, bf16 operands) -- the same recurrence as the Triton op, a different sum. `g` is in log space - (the adapter exponentiates it); `scale=None` is `K ** -0.5`; + (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]]`. With + 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, recurrence, output, slot scatter. The conv, q/k/v/g/beta - preparation, slot allocation and the gated norm on `o` are the caller's. + 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). @@ -73,8 +172,8 @@ class FlashinferChunkGatedDeltaRule(OpWrapper): ARCHS = frozenset({Arch.SM_103}) # spec: geometry (H, HV, K, V); seqlens; `slots`-row pool with `indices` - # (or _engine_slots); output (None | "buffer"). Every cell is - # qwen3.8-flash-next-nvfp4's geometry, in place over an fp32 pool. + # (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)}", @@ -97,6 +196,29 @@ class FlashinferChunkGatedDeltaRule(OpWrapper): 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 = """ @@ -114,12 +236,28 @@ class FlashinferChunkGatedDeltaRule(OpWrapper): (`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. - - Memory per call: a packed fp32 copy of the named slots in and out (2 x - 3 MiB per sequence), `exp(g)`, and the output. q/k/v/beta are not copied + 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 @@ -130,8 +268,9 @@ class FlashinferChunkGatedDeltaRule(OpWrapper): Numerics (measured vs an fp64 recurrence, 1 x 8192 tokens, qwen3.8 geometry): o rel-L2 4.06e-3, max abs 1.26e-3; state rel-L2 3.31e-3 -- the Triton op's are - 4.00e-3 / 1.13e-3 / 3.26e-3. Not bitwise to it. No cold-call defect: - nothing in the call is autotuned. + 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 @@ -140,7 +279,7 @@ class FlashinferChunkGatedDeltaRule(OpWrapper): Not checked: indices distinct and in `[0, S)`; q/k unit-norm (by the caller or the flag). - Not certified: `state_workspace=`, bf16/fp16 pool, `scale` other than the + 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`. @@ -166,7 +305,28 @@ def __call__( 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, @@ -202,6 +362,8 @@ def reference( 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)`. @@ -217,10 +379,13 @@ def reference( 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: @@ -228,7 +393,8 @@ def reference( k_t = k_t / torch.sqrt(k_t.pow(2).sum(-1, keepdim=True) + 1e-6) q_t = q_t.repeat_interleave(repeats, dim=0) * scale k_t = k_t.repeat_interleave(repeats, dim=0) - state = state * g[0, t].to(dt).exp()[:, None, None] + 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 ] 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__tp1/ops/ssm/gdn_cached_replay_cute.py new file mode 100644 index 000000000000..f7d53da8fd26 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_cached_replay_cute.py @@ -0,0 +1,2715 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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). + +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 +m-tiles and each of its warps `tpc / warps` of them (strided by the warp count). + +Per m-tile a warp computes, on mma.sync m16n8k16 (bf16 in, fp32 out), with the 16 state rows as the +A operand and the T normalized keys and T scaled queries as the 8 (T <= 4) or 16 (T <= 8) B columns, + D[v, k_t] = h0[v, :] . k_t D[v, q_t] = h0[v, :] . q_t +The head dimension is permuted identically in A and B (lane quad c holds dims 32m+8c..32m+8c+7), so +16-byte loads give the fragments directly. A warp's first tile is loaded into registers, its other +tiles by 1-D bulk copies (TMA, mbarrier) into shared memory -- the CTA's whole state share is in +flight from its start -- with an L2 evict-first policy (the state is read once a step). The history +term is one more MMA: the history u rows (shared memory, ldmatrix.trans) times the [16 x 2T] +coefficient matrix (k_t . k_j or q_t . k_j) exp(G_t - G_j), rounded to bf16 as the Triton kernel +rounds it. The T x T forward substitution and the output combination run per state row on CUDA +cores. With `commit`, a row whose history wraps (`p + T > 16`) also writes the advanced checkpoint +`exp(g_start) h0 + sum_j exp(g_start - G_j) u_j k_j^T` over the tile it just used (a +[16 x 16] x [16 x 128] MMA, u split into a bf16 hi and lo part as the Triton commit does), so the +checkpoint is read once and written once and the engine's all-layer commit is not needed. + +Given the engine's replay work items, each CTA reads its row's (row, slot, pending count, half) in +one load; otherwise slot, then pending count / half. With PDL the state and history (not written by +the previous kernel) are read before `griddepcontrol.wait`; without it the step's own small inputs +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). +""" + +from __future__ import annotations + +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 as _ty +from cutlass.cutlass_dsl import dsl_user_op + +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 +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 +LOG2E = 1.4426950408889634 +LN2 = 0.6931471805599453 + +# --------------------------------------------------------------------------------------- +# 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( + _ty.i32() if ty is cutlass.Int32 else _ty.f32(), res, [i], loc=loc, ip=ip + ) + ) + for i in range(n) + ] + + +def _make_ldg4(qual): + @dsl_user_op + def _ldg4(gaddr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([_ty.i32()] * 4), + [_ir(gaddr, cutlass.Int64, loc, ip)], + f"ld.global{qual}.v4.b32 {{$0, $1, $2, $3}}, [$4];", + "=r,=r,=r,=r,l", + loc, + ip, + side=True, + ) + return _unpack(res, 4, cutlass.Int32, loc, ip) + + return _ldg4 + + +@dsl_user_op +def _policy_evict_first(*, loc=None, ip=None): + """An L2 cache policy: evict first (data read once per step).""" + return cutlass.Int64( + _asm(_ty.i64(), [], "createpolicy.fractional.L2::evict_first.b64 $0, 1.0;", "=l", loc, ip) + ) + + +@dsl_user_op +def _ldg4_hint(gaddr, pol, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([_ty.i32()] * 4), + [_ir(gaddr, cutlass.Int64, loc, ip), _ir(pol, cutlass.Int64, loc, ip)], + "ld.global.nc.L1::no_allocate.L2::cache_hint.v4.b32 {$0, $1, $2, $3}, [$4], $5;", + "=r,=r,=r,=r,l,l", + loc, + ip, + side=True, + ) + return _unpack(res, 4, cutlass.Int32, loc, ip) + + +_LDG4 = { + "plain": _make_ldg4(""), + "na": _make_ldg4(".L1::no_allocate"), + "nc": _make_ldg4(".nc.L1::no_allocate"), + "nc256": _make_ldg4(".nc.L1::no_allocate.L2::256B"), + "cs": _make_ldg4(".cs"), + "lu": _make_ldg4(".lu"), +} +_ldg4 = _LDG4["plain"] + + +@dsl_user_op +def _ldg4p(gaddr, pred, *, loc=None, ip=None): + """16 bytes, or zeros where `pred` is 0 (nothing read).""" + res = _asm( + llvm.StructType.get_literal([_ty.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.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 _stg4(gaddr, a, b, c, d, *, loc=None, ip=None): + _asm( + _ty.i32(), + [_ir(gaddr, cutlass.Int64, loc, ip)] + + [_ir(x, cutlass.Int32, loc, ip) 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=True, + ) + + +@dsl_user_op +def _stg_bf16(gaddr, x, *, loc=None, ip=None): + """One f32 rounded to bf16 and stored (2 bytes).""" + _asm( + _ty.i32(), + [_ir(gaddr, cutlass.Int64, loc, ip), _ir(x, 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 _stg_f32(gaddr, x, *, loc=None, ip=None): + _asm( + _ty.i32(), + [_ir(gaddr, cutlass.Int64, loc, ip), _ir(x, cutlass.Float32, loc, ip)], + "{ st.global.f32 [$1], $2; mov.u32 $0, 0; }", + "=r,l,f", + loc, + ip, + side=True, + ) + + +@dsl_user_op +def _ldg_i32(gaddr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _ty.i32(), + [_ir(gaddr, cutlass.Int64, loc, ip)], + "ld.global.b32 $0, [$1];", + "=r,l", + loc, + ip, + side=True, + ) + ) + + +@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).""" + return cutlass.Float32( + _asm( + _ty.f32(), + [_ir(gaddr, cutlass.Int64, loc, ip), _ir(pred, cutlass.Int32, loc, ip)], + "{ .reg .pred q; setp.ne.b32 q, $2, 0; mov.f32 $0, 0f00000000; @q ld.global.f32 $0, [$1]; }", + "=f,l,r", + loc, + ip, + side=True, + ) + ) + + +@dsl_user_op +def _ldg_b16p(gaddr, pred, *, loc=None, ip=None): + """16 bits zero-extended (0 where `pred` is 0); converted where used, so the load does not + stall the issuing thread.""" + return cutlass.Int32( + _asm( + _ty.i32(), + [_ir(gaddr, cutlass.Int64, loc, ip), _ir(pred, cutlass.Int32, loc, ip)], + "{ .reg .pred q; setp.ne.b32 q, $2, 0; mov.b32 $0, 0; @q ld.global.u16 $0, [$1]; }", + "=r,l,r", + loc, + ip, + side=True, + ) + ) + + +@dsl_user_op +def _b16_f32(x, *, loc=None, ip=None): + """bf16 bits in the low half -> f32.""" + return cutlass.Float32( + _asm( + _ty.f32(), + [_ir(x, cutlass.Int32, loc, ip)], + "{ .reg .b32 t; shl.b32 t, $1, 16; mov.b32 $0, t; }", + "=f,r", + loc, + ip, + side=True, + ) + ) + + +@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).""" + _asm( + _ty.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 _cp_commit(*, loc=None, ip=None): + _asm(_ty.i32(), [], "{ cp.async.commit_group; mov.u32 $0, 0; }", "=r", loc, ip, side=True) + + +@dsl_user_op +def _cp_wait_all(*, loc=None, ip=None): + _asm(_ty.i32(), [], "{ cp.async.wait_all; mov.u32 $0, 0; }", "=r", loc, ip, side=True) + + +@dsl_user_op +def _sts_i32(saddr, a, *, loc=None, ip=None): + _asm( + _ty.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 _sts_f32(saddr, a, *, loc=None, ip=None): + _asm( + _ty.i32(), + [_ir(saddr, cutlass.Int32, loc, ip), _ir(a, cutlass.Float32, loc, ip)], + "{ st.shared.f32 [$1], $2; mov.u32 $0, 0; }", + "=r,r,f", + loc, + ip, + side=True, + ) + + +@dsl_user_op +def _sts2_f32(saddr, a, b, *, loc=None, ip=None): + _asm( + _ty.i32(), + [ + _ir(saddr, cutlass.Int32, loc, ip), + _ir(a, cutlass.Float32, loc, ip), + _ir(b, cutlass.Float32, loc, ip), + ], + "{ st.shared.v2.f32 [$1], {$2, $3}; mov.u32 $0, 0; }", + "=r,r,f,f", + loc, + ip, + side=True, + ) + + +@dsl_user_op +def _sts_bf16(saddr, a, *, loc=None, ip=None): + _asm( + _ty.i32(), + [_ir(saddr, cutlass.Int32, loc, ip), _ir(a, cutlass.Float32, loc, ip)], + "{ .reg .b16 h; cvt.rn.bf16.f32 h, $2; st.shared.b16 [$1], h; mov.u32 $0, 0; }", + "=r,r,f", + loc, + ip, + side=True, + ) + + +@dsl_user_op +def _lds_i32(saddr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _ty.i32(), + [_ir(saddr, cutlass.Int32, loc, ip)], + "ld.shared.b32 $0, [$1];", + "=r,r", + loc, + ip, + side=True, + ) + ) + + +@dsl_user_op +def _lds_f32(saddr, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _ty.f32(), + [_ir(saddr, cutlass.Int32, loc, ip)], + "ld.shared.f32 $0, [$1];", + "=f,r", + loc, + ip, + side=True, + ) + ) + + +@dsl_user_op +def _lds_bf16(saddr, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _ty.f32(), + [_ir(saddr, cutlass.Int32, loc, ip)], + "{ .reg .b16 h; ld.shared.b16 h, [$1]; cvt.f32.bf16 $0, h; }", + "=f,r", + loc, + ip, + side=True, + ) + ) + + +@dsl_user_op +def _lds4_i32(saddr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([_ty.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 _lds4_f32(saddr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([_ty.f32()] * 4), + [_ir(saddr, cutlass.Int32, loc, ip)], + "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 _ldsm4t(saddr, *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([_ty.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 (bf16 A/B, f32 C/D); a: 4 regs, b: 2 regs, c: 4 f32.""" + res = _asm( + llvm.StructType.get_literal([_ty.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 _bf2_pack(lo, hi, *, loc=None, ip=None): + """Two f32 -> bf16x2 (lo in the low half).""" + return cutlass.Int32( + _asm( + _ty.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 _bf2_lo_v(w, *, loc=None, ip=None): + """`_bf2_lo` pinned in program order (for prefetched words: a pure conversion may be hoisted + up to the load and stall the issuing warp there).""" + return cutlass.Float32( + _asm( + _ty.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, + side=True, + ) + ) + + +@dsl_user_op +def _bf2_hi_v(w, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _ty.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, + side=True, + ) + ) + + +@dsl_user_op +def _bf2_lo(w, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _ty.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( + _ty.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, + ) + ) + + +@dsl_user_op +def _ex2(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _ty.f32(), + [_ir(x, cutlass.Float32, loc, ip)], + "ex2.approx.ftz.f32 $0, $1;", + "=f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _lg2(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _ty.f32(), + [_ir(x, cutlass.Float32, loc, ip)], + "lg2.approx.ftz.f32 $0, $1;", + "=f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _div_approx(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _ty.f32(), + [_ir(a, cutlass.Float32, loc, ip), _ir(b, cutlass.Float32, loc, ip)], + "div.approx.ftz.f32 $0, $1, $2;", + "=f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _div_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _ty.f32(), + [_ir(a, cutlass.Float32, loc, ip), _ir(b, cutlass.Float32, loc, ip)], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _sqrt_rn(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_ty.f32(), [_ir(x, cutlass.Float32, loc, ip)], "sqrt.rn.f32 $0, $1;", "=f,f", loc, ip) + ) + + +@dsl_user_op +def _mbar_init(saddr, count, *, loc=None, ip=None): + _asm( + _ty.i32(), + [_ir(saddr, 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 _fence_mbar_init(*, loc=None, ip=None): + _asm( + _ty.i32(), + [], + "{ fence.mbarrier_init.release.cluster; mov.u32 $0, 0; }", + "=r", + loc, + ip, + side=True, + ) + + +@dsl_user_op +def _bulk_g2s(saddr, gaddr, nbytes, mbar, *, loc=None, ip=None): + """mbarrier.arrive.expect_tx + one 1-D bulk copy global -> shared of `nbytes`.""" + _asm( + _ty.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), + ], + "{ mbarrier.arrive.expect_tx.shared::cta.b64 _, [$4], $3; " + "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 _bulk_g2s_hint(saddr, gaddr, nbytes, mbar, pol, *, loc=None, ip=None): + """`_bulk_g2s` with an L2 cache policy.""" + _asm( + _ty.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), + _ir(pol, cutlass.Int64, loc, ip), + ], + "{ 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], $5; " + "mov.u32 $0, 0; }", + "=r,r,l,r,r,l", + loc, + ip, + side=True, + ) + + +@dsl_user_op +def _mbar_wait(mbar, phase, *, loc=None, ip=None): + _asm( + _ty.i32(), + [_ir(mbar, cutlass.Int32, loc, ip), _ir(phase, cutlass.Int32, loc, ip)], + "{ .reg .pred P1; LAB_WAIT: mbarrier.try_wait.parity.shared::cta.b64 P1, [$1], $2; " + "@P1 bra DONE; bra LAB_WAIT; DONE: mov.u32 $0, 0; }", + "=r,r,r", + loc, + ip, + side=True, + ) + + +@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) + + +def _shfl_xor(x, m): + return cute.arch.shuffle_sync_bfly(x, offset=m) + + +# --------------------------------------------------------------------------------------- +# the kernel +# --------------------------------------------------------------------------------------- + + +class _Replay: + def __init__( + self, + *, + T, + tpc, + warps, + commit, + pdl, + ld="nc", + pf=1, + minb=1, + l2norm=True, + gating=True, + g_f32=False, + b_f32=False, + early_trigger=True, + 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, + ): + assert 1 <= T <= 8 + assert tpc in (24, 12, 8, 4, 2, 1) and tpc % warps == 0 + self.T = T + self.TP = 4 if T <= 4 else 8 + self.NT = self.TP // 4 + self.tpc, self.W = tpc, warps + self.tpw = tpc // warps + self.S = TILES // tpc + self.nhv = 3 if tpc == 24 else (2 if tpc == 12 else 1) + self.commit, self.pdl = commit, pdl + # pf = 0 (shared-memory tiles only): every state tile goes by bulk copy, issued as soon + # 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.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.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.evict_first = evict_first + nset = 2 if (persistent or 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 + self.o_mb = o + o += -(-(tset * 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 + self.o_x = o + o += warps * 16 * self.RP + self.o_kk = o + o += 16 * 8 * 4 + self.o_cf = o + o += self.nhv * 16 * 16 * 2 + self.o_gh = o + o += nset * self.nhv * 16 * 4 + self.o_dec = o + o += self.nhv * 16 * 4 + self.o_g = o + o += self.nhv * 8 * 4 + self.o_eg = o + o += self.nhv * 8 * 4 + self.o_beta = o + o += self.nhv * 8 * 4 + self.o_a = o + o += self.nhv * 64 * 4 + self.o_cq = o + o += self.nhv * 64 * 4 + self.o_egs = o + o += 16 + self.smem = -(-o // 128) * 128 + 128 + + @cute.jit + def __call__( + self, + sidx_t: cute.Tensor, + 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, + wi_ptr: cutlass.Int64, + n_rows: cutlass.Int32, + 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, + ) + + # -- straight-line helpers (traced inline) ------------------------------------------------ + + def _tile_addr(self, h_ptr, slot64, h_ss, h, sp, warp, i, g, c): + mt = sp * self.tpc + warp + self.W * i + hv = h * GRP + mt // 8 + v0 = (mt % 8) * 16 + base = h_ptr + (slot64 * h_ss + cutlass.Int64(hv * V + v0) * K) * 2 + return base + cutlass.Int64(g * 256 + c * 16) + + def _qk_raw(self, nt, q_ptr, k_ptr, q_sn, q_st, k_sn, k_st, n64, h, g, c, scale, live): + """Raw bf16 words of key / query column nt * 8 + g at this lane's head dims, the + column's kind and its scale.""" + T, TP = self.T, self.TP + zi = cutlass.Int32(0) + t = g % 4 if TP == 4 else g + valid = (t < T) & live + tc = cutlass.min(t, T - 1) + qa = q_ptr + (n64 * q_sn + cutlass.Int64(tc) * q_st + cutlass.Int64(h * K + c * 8)) * 2 + ka = k_ptr + (n64 * k_sn + cutlass.Int64(tc) * k_st + cutlass.Int64(h * K + c * 8)) * 2 + if TP == 4: + is_q = g >= 4 + src = cutlass.Int64(cutlass.select_(is_q, qa, ka)) + mult = cutlass.Float32(cutlass.select_(is_q, scale, cutlass.Float32(1.0))) + elif nt == 1: + is_q = g >= 0 + src = qa + mult = scale + else: + is_q = g < 0 + src = ka + mult = cutlass.Float32(1.0) + vpred = cutlass.Int32(cutlass.select_(valid, cutlass.Int32(1), zi)) + raw = [] + for jj in range(4): + raw.extend(_ldg4p(src + jj * 64, vpred)) + return raw, is_q, mult + + def _qk_frag(self, raw, is_q, mult): + """L2-normalized (queries also scaled) bf16 B fragment from `_qk_raw`.""" + ss = cutlass.Float32(0.0) + fl = [] + for e in range(16): + lo = _bf2_lo_v(raw[e]) + hi = _bf2_hi_v(raw[e]) + fl.append(lo) + fl.append(hi) + ss = ss + lo * lo + hi * hi + if self.l2norm: + ss = ss + _shfl_xor(ss, 1) + ss = ss + _shfl_xor(ss, 2) + inv = _div_rn(mult, _sqrt_rn(ss + 1e-6)) + return [_bf2_pack(fl[2 * e] * inv, fl[2 * e + 1] * inv) for e in range(16)] + return [ + cutlass.Int32( + cutlass.select_(is_q, _bf2_pack(fl[2 * e] * mult, fl[2 * e + 1] * mult), raw[e]) + ) + for e in range(16) + ] + + def _gate_raw(self, pred, hv, n64, g_ptr, g_st, b_ptr, b_st, alog_ptr, dtb_ptr): + """Raw g / beta of the T tokens for v head `hv` and its A_log / dt_bias (predicated).""" + T = self.T + zf = cutlass.Float32(0.0) + pi = cutlass.Int32(cutlass.select_(pred, cutlass.Int32(1), cutlass.Int32(0))) + alv, dtb = zf, zf + if self.gating: + alv = _ldg_f32p(alog_ptr + cutlass.Int64(hv) * 4, pi) + dtb = _ldg_f32p(dtb_ptr + cutlass.Int64(hv) * 4, pi) + graws, braws = [], [] + for t in range(T): + if self.g_f32: + graws.append( + _ldg_f32p(g_ptr + (n64 * T + t) * g_st * 4 + cutlass.Int64(hv) * 4, pi) + ) + else: + graws.append( + _ldg_b16p(g_ptr + (n64 * T + t) * g_st * 2 + cutlass.Int64(hv) * 2, pi) + ) + if self.b_f32: + braws.append( + _ldg_f32p(b_ptr + (n64 * T + t) * b_st * 4 + cutlass.Int64(hv) * 4, pi) + ) + else: + braws.append( + _ldg_b16p(b_ptr + (n64 * T + t) * b_st * 2 + cutlass.Int64(hv) * 2, pi) + ) + return graws, braws, alv, dtb + + def _smem_tile(self, s_st, s_mb, i, g, c, zi): + """This lane's A fragment words of the warp's staged tile i (after its bulk copy landed). + + Rows g and g + 8; chunk order by row parity keeps each quarter-warp's two rows on + disjoint banks. + """ + kslot = i - self.pf + _mbar_wait(s_mb + kslot * 8, zi) + tb = s_st + kslot * 4096 + g * 256 + c * 16 + godd = (g % 2) == 1 + raw = [] + for ii in range(4): + jsw = (g % 2) ^ ii + raw.append(_lds4_i32(tb + jsw * 64)) + raw.append(_lds4_i32(tb + 2048 + jsw * 64)) + cur = [] + for jj in range(4): + for rr in range(2): + for e in range(4): + cur.append( + cutlass.Int32( + cutlass.select_(godd, raw[2 * (jj ^ 1) + rr][e], raw[2 * jj + rr][e]) + ) + ) + return cur + + def _tile_dots(self, cur, bq, zf): + """D = h0 [16 x 128] times the T keys / T scaled queries. + + Two MMA chains over the head dimension (even / odd 16-wide k-steps), summed. + """ + NT = self.NT + acc0 = [[zf, zf, zf, zf] for _ in range(NT)] + acc1 = [[zf, zf, zf, zf] for _ in range(NT)] + for jj in range(4): + for s in range(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 range(NT): + bfr = (bq[nt][4 * jj + 2 * s], bq[nt][4 * jj + 2 * s + 1]) + if s == 0: + acc0[nt] = _mma(a, bfr, acc0[nt]) + else: + acc1[nt] = _mma(a, bfr, acc1[nt]) + return [[acc0[nt][e] + acc1[nt][e] for e in range(4)] for nt in range(NT)] + + def _g_f32(self, x): + return x if self.g_f32 else _b16_f32(x) + + def _b_f32(self, x): + return x if self.b_f32 else _b16_f32(x) + + def _load_tile(self, addr, pol=None): + regs = [] + for jj in range(4): + for off in (jj * 64, 2048 + jj * 64): + if pol is not None: + regs.extend(_ldg4_hint(addr + off, pol)) + else: + regs.extend(_LDG4[self.ld](addr + off)) + return regs + + @cute.kernel + def kernel( + 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, + wi_ptr: cutlass.Int64, + scale: cutlass.Float32, + ): + 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 + + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + warp = tidx // 32 + lane = tidx % 32 + g = lane // 4 + c = lane % 4 + sp = bidx % S + nh = bidx // S + h = nh % H + if const_expr(self.use_wi): + wi = _LDG4["plain"](wi_ptr + cutlass.Int64(nh // H) * 16) + n = wi[0] + else: + n = nh // H + n64 = cutlass.Int64(n) + 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_u = sbase + self.o_u + s_kh = sbase + self.o_kh + 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_gh = sbase + self.o_gh + 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 + + if const_expr(self.pdl and not self.state_before_wait and not self.qk_first): + cute.arch.griddepcontrol_wait() + + NSM: cutlass.Constexpr = self.nsm + s_st = sbase + self.o_st + warp * (NSM * 4096) + s_mb = sbase + self.o_mb + warp * (NSM * 8) + if const_expr(NSM > 0): + if lane == 0: + for k_ in cutlass.range_constexpr(NSM): + _mbar_init(s_mb + k_ * 8, cutlass.Int32(1)) + _fence_mbar_init() + + hv_first = (sp * TPC) // 8 + hv0 = h * GRP + hv_first + gi = tidx - (32 * (W - 1) if W > 1 else 0) + gpred = (gi >= 0) & (gi < NHV) + if const_expr(self.qk_first): + # the step's own inputs first: they are L2-resident and would otherwise queue behind + # the state stream + if const_expr(self.pdl): + cute.arch.griddepcontrol_wait() + if const_expr(self.early_trigger): + cute.arch.griddepcontrol_launch_dependents() + pre = [ + self._qk_raw( + nt, q_ptr, k_ptr, q_sn, q_st, k_sn, k_st, n64, h, g, c, scale, tidx >= 0 + ) + for nt in range(NT) + ] + gpre = self._gate_raw( + gpred, + hv0 + cutlass.min(cutlass.max(gi, zi), NHV - 1), + n64, + g_ptr, + g_st, + b_ptr, + b_st, + alog_ptr, + dtb_ptr, + ) + n_v0: cutlass.Constexpr = NHV * T * 16 + for it0 in cutlass.range_constexpr(-(-n_v0 // NTHR)): + ev0 = tidx + it0 * NTHR + if ev0 < n_v0: + hl0 = ev0 // (T * 16) + tv0 = (ev0 // 16) % T + qv0 = ev0 % 16 + sv0 = ( + v_ptr + + ( + n64 * v_sn + + cutlass.Int64(tv0) * v_st + + cutlass.Int64((hv0 + hl0) * V + qv0 * 8) + ) + * 2 + ) + _cp16z(s_v + (hl0 * T + tv0) * 256 + qv0 * 16, sv0, cutlass.Int32(16)) + _cp_commit() + + if const_expr(self.use_wi): + slot = wi[1] + else: + slot = _ldg_i32(sidx_ptr + n64 * 4) + if slot >= 0: + slot64 = cutlass.Int64(slot) + # ---- 1. the state tiles: bulk copies first, then the register tile ----------------- + pol = _policy_evict_first() if const_expr(self.evict_first) else None + if const_expr(NSM > 0): + if lane == 0: + for k_ in cutlass.range_constexpr(NSM): + if const_expr(self.evict_first): + _bulk_g2s_hint( + s_st + k_ * 4096, + self._tile_addr( + h_ptr, slot64, h_ss, h, sp, warp, self.pf + k_, 0, 0 + ), + cutlass.Int32(4096), + s_mb + k_ * 8, + pol, + ) + else: + _bulk_g2s( + s_st + k_ * 4096, + self._tile_addr( + h_ptr, slot64, h_ss, h, sp, warp, self.pf + k_, 0, 0 + ), + cutlass.Int32(4096), + s_mb + k_ * 8, + ) + tiles = [] + for i in cutlass.range_constexpr(self.pf): + tiles.append( + self._load_tile(self._tile_addr(h_ptr, slot64, h_ss, h, sp, warp, i, g, c), pol) + ) + + # ---- 2. history bookkeeping and the history rows into shared memory ------------ + if const_expr(self.use_wi): + p = wi[2] + rb = wi[3] + else: + p = _ldg_i32(pnat_ptr + slot64 * 4) + rb = _ldg_i32(buf_ptr + slot64 * 4) + wrap = (p + T) > HIST + wb = cutlass.Int32(cutlass.select_(wrap, 1 - rb, rb)) + woff = cutlass.Int32(cutlass.select_(wrap, zi, p)) + rslot = slot64 * 2 + cutlass.Int64(rb) + g_rd = og_ptr + rslot * (HV * HIST * 4) + wslot = slot64 * 2 + cutlass.Int64(wb) + u_rd = ou_ptr + rslot * (HIST * HV * V * 2) + k_rd = ok_ptr + rslot * (HIST * H * K * 2) + 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) + + n_u: cutlass.Constexpr = NHV * 256 + if const_expr(self.spec_hist): + # both halves, every row, before the pending count / half are known; rows j >= p + # are masked where they are used + for hb in cutlass.range_constexpr(2): + hslot = slot64 * 2 + hb + u_h = ou_ptr + hslot * (HIST * HV * V * 2) + k_h = ok_ptr + hslot * (HIST * H * K * 2) + g_h = og_ptr + hslot * (HV * HIST * 4) + for it in cutlass.range_constexpr(-(-n_u // NTHR)): + e = tidx + it * NTHR + if e < n_u: + hl = e // 256 + j = (e // 16) % 16 + qq = e % 16 + src = u_h + cutlass.Int64(((j * HV + hv0 + hl) * V + qq * 8) * 2) + _cp16z( + s_u + hb * (NHV * 16 * PU) + hl * (16 * PU) + j * PU + qq * 16, + src, + cutlass.Int32(16), + ) + for it in cutlass.range_constexpr(-(-256 // NTHR)): + e = tidx + it * NTHR + if e < 256: + j = e // 16 + qq = e % 16 + src = k_h + cutlass.Int64(((j * H + h) * K + qq * 8) * 2) + _cp16z(s_kh + hb * (16 * PK) + j * PK + qq * 16, src, cutlass.Int32(16)) + if tidx < NHV * 4: + hl = tidx // 4 + qq = tidx % 4 + src = g_h + cutlass.Int64(((hv0 + hl) * HIST + qq * 4) * 4) + _cp16z(s_gh + hb * (NHV * 64) + hl * 64 + qq * 16, src, cutlass.Int32(16)) + else: + for it in cutlass.range_constexpr(-(-n_u // NTHR)): + e = tidx + it * NTHR + if e < n_u: + hl = e // 256 + j = (e // 16) % 16 + qq = e % 16 + jc = cutlass.min( + j, cutlass.max(p - 1, zi) + ) # an in-bounds row for a zero fill + src = u_rd + cutlass.Int64(((jc * HV + hv0 + hl) * V + qq * 8) * 2) + _cp16z( + s_u + hl * (16 * PU) + j * PU + qq * 16, + src, + cutlass.Int32(cutlass.select_(j < p, cutlass.Int32(16), zi)), + ) + for it in cutlass.range_constexpr(-(-256 // NTHR)): + e = tidx + it * NTHR + if e < 256: + j = e // 16 + qq = e % 16 + jc = cutlass.min(j, cutlass.max(p - 1, zi)) + src = k_rd + cutlass.Int64(((jc * H + h) * K + qq * 8) * 2) + _cp16z( + s_kh + j * PK + qq * 16, + src, + cutlass.Int32(cutlass.select_(j < p, cutlass.Int32(16), zi)), + ) + if tidx < NHV * 4: + hl = tidx // 4 + qq = tidx % 4 + src = g_rd + cutlass.Int64(((hv0 + hl) * HIST + qq * 4) * 4) + _cp16z(s_gh + hl * 64 + qq * 16, src, cutlass.Int32(16)) + _cp_commit() + + if const_expr(self.spec_hist): + s_u = s_u + rb * (NHV * 16 * PU) + s_kh = s_kh + rb * (16 * PK) + s_gh = s_gh + rb * (NHV * 64) + + # ---- 3. the step's own inputs (written by the previous kernel) ------------------ + if const_expr(not self.qk_first): + if const_expr(self.pdl and self.state_before_wait): + cute.arch.griddepcontrol_wait() + if const_expr(self.pdl and self.early_trigger): + cute.arch.griddepcontrol_launch_dependents() + n_v: cutlass.Constexpr = NHV * T * 16 + for it in cutlass.range_constexpr(-(-n_v // NTHR)): + e = tidx + it * 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 + (hl * T + t) * 256 + qq * 16, src, cutlass.Int32(16)) + _cp_commit() + + # key / query B fragments: column nt * 8 + g of the m16n8 tiles; lane quad c holds + # head dims 32 jj + 8 c .. 32 jj + 8 c + 7. + bq = [] + for nt in cutlass.range_constexpr(NT): + if const_expr(self.qk_first): + rq = pre[nt] + else: + rq = self._qk_raw( + nt, q_ptr, k_ptr, q_sn, q_st, k_sn, k_st, n64, h, g, c, scale, tidx >= 0 + ) + bq.append(self._qk_frag(rq[0], rq[1], rq[2])) + + # the new keys into the history (one CTA of the head, warp 0) + if sp == 0: + if warp == 0: + # column g < T of tile 0 is key g (both column layouts) + 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], + ) + + # gating, cumulative decays, the G history (a thread per v head, off warp 0) + if gpred: + hl = gi + hv = hv0 + hl + gs = _ldg_f32p( + g_rd + cutlass.Int64(hv * HIST + p - 1) * 4, + cutlass.Int32(cutlass.select_(p > 0, cutlass.Int32(1), zi)), + ) + if const_expr(self.qk_first): + graws, braws, alv, dtb = gpre + else: + graws, braws, alv, dtb = self._gate_raw( + gpred, hv, n64, g_ptr, g_st, b_ptr, b_st, alog_ptr, dtb_ptr + ) + aexp = _exp(alv) + owns = (hv_first + hl) * 8 >= sp * TPC + gl = zf + for t in cutlass.range_constexpr(T): + graw = self._g_f32(graws[t]) + braw = self._b_f32(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)) + + # 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]) + + # the tiles' state dots need only the state and the new keys / queries: run them + # while the history is still in flight + early = [] + if const_expr(self.use_early): + for i in cutlass.range_constexpr(TPW): + if const_expr(NSM == 0): + if const_expr(i + self.pf < TPW): + tiles.append( + self._load_tile( + self._tile_addr( + h_ptr, slot64, h_ss, h, sp, warp, i + self.pf, g, c + ) + ) + ) + if const_expr(i < self.pf or NSM == 0): + early.append(self._tile_dots(tiles[i], bq, zf)) + else: + early.append( + self._tile_dots(self._smem_tile(s_st, s_mb, i, g, c, zi), bq, zf) + ) + + _cp_wait_all() + cute.arch.sync_threads() + + # ---- 4. history-key dots, commit operand, per-head scalars ---------------------- + 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] + ) + # history coefficients (k_t . k_j or q_t . k_j) exp(G_t - G_j), j < p, as bf16 + # B fragments [head][column][j] + 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 it in cutlass.range_constexpr(-(-256 // NTHR)): + e = tidx + it * 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): + wv = w4[s4] + if const_expr(self.spec_hist): + wv = cutlass.Int32(cutlass.select_(j < p, wv, zi)) + _sts_i32(s_kp + j * PU + (16 * jj + 4 * s4 + aa) * 4, wv) + + n_ts: cutlass.Constexpr = NHV * T * T + for it in cutlass.range_constexpr(-(-n_ts // NTHR)): + e = tidx + it * 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) + kk = _lds_f32(s_kk + (t * 8 + s) * 4) + qk = _lds_f32(s_kk + ((TP + t) * 8 + s) * 4) + av = cutlass.Float32(cutlass.select_(s < t, bt * kk * dec, zf)) + cv = cutlass.Float32(cutlass.select_(s <= t, qk * 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() + + # ---- 5. the warp's m-tiles ------------------------------------------------------ + jrow = (lane % 8) + 8 * (lane // 16) + v8 = 8 * ((lane // 8) % 2) + r = lane % 16 + hf = lane // 16 + godd = (g % 2) == 1 + mask01 = cutlass.Int32( + cutlass.select_(2 * c < p, cutlass.Int32(0xFFFF), zi) + ) | cutlass.Int32(cutlass.select_(2 * c + 1 < p, cutlass.Int32(-65536), zi)) + mask89 = cutlass.Int32( + cutlass.select_(2 * c + 8 < p, cutlass.Int32(0xFFFF), zi) + ) | cutlass.Int32(cutlass.select_(2 * c + 9 < p, cutlass.Int32(-65536), zi)) + for i in cutlass.range_constexpr(TPW): + if const_expr(NSM == 0 and not self.use_early): + if const_expr(i + self.pf < TPW): + tiles.append( + self._load_tile( + self._tile_addr(h_ptr, slot64, h_ss, h, sp, warp, i + self.pf, g, c) + ) + ) + if const_expr(i < self.pf or NSM == 0): + cur = tiles[i] + elif const_expr(self.use_early): + # a staged tile is read again from shared memory only for its commit + cur = None + else: + # rows g and g + 8 of the staged tile; chunk order by row parity keeps each + # quarter-warp's two rows on disjoint banks + kslot = i - self.pf + _mbar_wait(s_mb + kslot * 8, zi) + 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 + + if const_expr(self.use_early): + acc = early[i] + else: + acc = self._tile_dots(cur, bq, zf) + + au = _ldsm4t(s_u + hl * (16 * PU) + jrow * PU + (v0 + v8) * 2) + if const_expr(self.spec_hist): + au = [cutlass.Int32(au[q4] & (mask01 if q4 < 2 else mask89)) for q4 in range(4)] + 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() + + # per state row: the forward substitution and the outputs + 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_v + ((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() + + # the checkpoint commit of a wrapping row, over the tile just used + 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))) + if const_expr(cur is None): + curc = self._smem_tile(s_st, s_mb, i, g, c, zi) + else: + curc = cur + 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 = curc[8 * jj + s4] + wh = curc[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]) + + 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 +# --------------------------------------------------------------------------------------- + +_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.""" + if n <= 2: + return dict(tpc=2, warps=2, pf=0, minb=4, early_acc=True) + if n <= 8: + return dict(tpc=4, warps=4, pf=0, minb=4, early_acc=True) + if commit: + return dict(tpc=8, warps=4, pf=1, minb=4) + if n <= 16: + return dict(tpc=24, warps=8, pf=1, minb=2) + return dict(tpc=8, warps=4, pf=1, minb=5) + + +def _compile(key, kern): + fn = _COMPILED.get(key) + if fn is None: + dummy = torch.zeros(4, dtype=torch.int32, device="cuda") + i64 = cutlass.Int64(0) + fn = cute.compile( + kern, + from_dlpack(dummy, assumed_align=4).mark_layout_dynamic(leading_dim=0), + *([i64] * 24), + 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", + ) + _COMPILED[key] = fn + return fn + + +def cached_replay_update( + 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: Optional[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, + 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, + **ignored, +) -> torch.Tensor: + """Same contract as `fused_recurrent_gated_delta_rule_cached_replay_update` (bf16 state). + + Without work items a wrapping row commits its checkpoint in-kernel, as the Triton op does; + with work items + `use_all_layer_commit` it does not (the caller runs the all-layer commit). + `in_kernel_commit` overrides that choice (True: commit in-kernel even when the batch is + partitioned -- then the caller must NOT run the all-layer commit for this layer). + Given work items, the kernel reads each row's (row, slot, pending count, half) from them (one + load instead of two dependent ones), so they must list every one of the N rows, as the engine's + builder over all decode rows does (`config={"use_wi": False}` falls back to `state_indices`). + Work items with `in_kernel_commit=True` need no `use_all_layer_commit` (the caller runs none). + 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" + assert 1 <= T <= 8 + assert ssm_states.dtype == torch.bfloat16, "bf16 checkpoint state only" + assert ssm_states.stride()[1:] == (V * K, K, 1) + assert old_u.dtype == torch.bfloat16 and old_k.dtype == torch.bfloat16 + assert old_G.dtype == torch.float32 + for x in (old_u, old_k, old_G): + assert x.is_contiguous() + assert g.shape == (N, T, HV) and beta.shape == (N, T, HV) + assert g.stride(2) == 1 and beta.stride(2) == 1 + assert g.stride(0) == T * g.stride(1) and beta.stride(0) == T * beta.stride(1) + assert g.dtype in (torch.bfloat16, torch.float32) and beta.dtype in ( + torch.bfloat16, + torch.float32, + ) + if (replay_work_items is None) != (n_writes is None): + raise ValueError("replay_work_items and n_writes must either both be set or both be None") + if replay_work_items is not None and not use_all_layer_commit and not in_kernel_commit: + raise ValueError("Partitioned replay requires the all-layer commit") + if replay_work_items is None and use_all_layer_commit: + raise ValueError("use_all_layer_commit requires replay work items") + commit = replay_work_items is None if in_kernel_commit is None else bool(in_kernel_commit) + if scale is None: + scale = K**-0.5 + if output is None: + output = torch.empty((N, T, HV, V), dtype=q.dtype, device=q.device) + assert output.is_contiguous() and output.dtype == torch.bfloat16 + if N == 0: + return output + if packed_qkv is not None: + 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) + 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" + qp, kp, vp = q.data_ptr(), k.data_ptr(), v.data_ptr() + qs, ks, vs = q.stride()[:2], k.stride()[:2], v.stride()[:2] + gating = A_log is not None + if gating: + assert dt_bias is not None + A_log = A_log.float().contiguous() + dt_bias = dt_bias.float().contiguous() + 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)) + 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") + ) + 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, + commit=commit, + pdl=pdl, + l2norm=bool(use_qk_l2norm_in_kernel), + gating=gating, + g_f32=g.dtype == torch.float32, + b_f32=beta.dtype == torch.float32, + **cfg, + ) + key = tuple(sorted(kw.items())) + fn = _COMPILED.get(key) + if fn is None: + fn = _compile(key, _Replay(**kw)) + sidx = state_indices if state_indices.dtype == torch.int32 else state_indices.to(torch.int32) + fn( + sidx, + qp, + kp, + vp, + qs[0], + qs[1], + ks[0], + ks[1], + vs[0], + vs[1], + g.data_ptr(), + g.stride(1), + beta.data_ptr(), + beta.stride(1), + ap, + dp, + output.data_ptr(), + ssm_states.data_ptr(), + ssm_states.stride(0), + prev_num_accepted_tokens.data_ptr(), + cache_buf_idx.data_ptr(), + old_u.data_ptr(), + old_k.data_ptr(), + old_G.data_ptr(), + replay_work_items.data_ptr() if use_wi else 0, + N, + float(scale), + ) + return output + + +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). + + 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). + + 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 with the in-kernel + # commit at every step. + 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), + ), + Cell( + why="MTP3 verify at con32: N=32, T=4, in-kernel commit", + spec=dict(path="all_layer", n=32, t=4, steps=4, slots=72), + ), + Cell( + why="MTP3 verify at con256: N=256, T=4, in-kernel commit", + spec=dict(path="all_layer", n=256, t=4, steps=3, slots=300, layers=1), + ), + ) + + 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, `in_kernel_commit=True` + (with the work items the metadata builds, if given, and no all-layer commit); the target + switches the manager's all-layer commit off when it binds a bf16 pool. `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). + """ + + def __call__(self, *args, **kwargs) -> torch.Tensor: + return cached_replay_update(*args, **kwargs) + + def reference(self, *args, **kwargs): + raise NotImplementedError( + "the multi-step cell harness drives the Triton entry's reference (see the test)" + ) + + +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/ssm/gdn_decode_rows_gated_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py index 6ec2bd7904fb..7fb0f8144a6b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py @@ -1,20 +1,29 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """One Gated-DeltaNet decode step per row and the gated RMS norm on its output, streamed -over the fp32 state pool: the decode step at batch sizes past the cluster kernel's. +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, read -and written once -- so a CTA owns one whole state and streams it: warp w the value rows -[w * RPW, (w + 1) * RPW), lane l the key columns [4l, 4l + 4) of each, so a warp's -access is one contiguous 512-byte state row. The state is loaded before +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 @@ -35,7 +44,7 @@ from .gdn_decode_gated_norm_pdl import _recurrence_reference from .gdn_decode_gated_norm_pdl import gdn_decode_gated_norm_pdl as _one_launch -#: Key columns a lane (a 16-byte vector of a state row). +#: 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 @@ -101,6 +110,40 @@ def _st_state(addr, pred, a, b, c, d, *, loc=None, ip=None): ) +@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.""" @@ -162,6 +205,151 @@ def _round_bf16(x, *, loc=None, ip=None): ) +@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 @@ -187,12 +375,29 @@ def _rows_reduce(vals, lane): 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.""" + 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): + 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. @@ -202,6 +407,11 @@ def __init__(self, k_heads, v_heads, kdim, vdim, pdl, norm=True): 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__( @@ -243,9 +453,10 @@ def __call__( ).launch( grid=[self.hv, rows, 1], block=[32 * _WARPS, 1, 1], - smem=_WARPS * 4 + 64, + 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 @@ -272,9 +483,18 @@ def kernel( kh = hv // (self.hv // self.h) v0 = warp * self.rpw jl = lane // self.lpr # the value row whose sums this lane ends up holding - red = cutlass.utils.SmemAllocator().allocate_tensor( - cutlass.Float32, cute.make_layout((_WARPS,), stride=(1,)), byte_alignment=16 - ) + 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. @@ -284,15 +504,34 @@ def kernel( live = cutlass.Int32(1) # 64-bit: a coalesced cache page puts every GDN layer's state of a slot in one # page, so `slot * s_slot` passes int32. - base = ( - pool_base - + ( - cutlass.Int64(cutlass.max(slot, cutlass.Int32(0))) * s_slot - + cutlass.Int64(hv * s_head + v0 * s_row + lane * _PL) + 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 ) - * 4 - ) - h = [_ld_state(base + cutlass.Int64(j * s_row * 4), live) for j in range(self.rpw)] + 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) @@ -333,23 +572,40 @@ def kernel( 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 = h[j][e] * kn[e] + acc + 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): - h[j][e] = u * kn[e] + decay * h[j][e] - acc = h[j][e] * qn[e] + acc + 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) - _st_state(base + cutlass.Int64(j * s_row * 4), live, *h[j]) + 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 @@ -361,12 +617,20 @@ def kernel( y = o_b if cutlass.const_expr(self.norm): ss = cute.arch.warp_reduction_sum(o_b * o_b * cutlass.Float32(own)) - if lane == 0: - red[warp] = ss - cute.arch.sync_threads() - tot = red[0] - for i in cutlass.range_constexpr(1, _WARPS): - tot = tot + red[i] + 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) @@ -375,10 +639,84 @@ def kernel( _COMPILED: dict = {} -def launch(proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, pdl=True, out=None): +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.""" + (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 @@ -387,7 +725,11 @@ def launch(proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, p if rows == 0: return out pdl = bool(pdl and get_env_enable_pdl()) - key = (k_heads, v_heads, kdim, vdim, pdl, norm) + 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(), @@ -397,32 +739,13 @@ def launch(proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, p (norm_weight if norm else A_log).data_ptr(), out.data_ptr(), ) - fn = _COMPILED.get(key) + base_key = (k_heads, v_heads, kdim, vdim, pdl, norm, state_bf16) + fn = _COMPILED.get(base_key + (bulk,)) if fn is None: - gm = cute.AddressSpace.gmem - cargs = ( - from_dlpack(proj, assumed_align=8).mark_layout_dynamic(leading_dim=1), - make_ptr(cutlass.Float32, ptrs[1], gm, assumed_align=16), - make_ptr(cutlass.Int32, ptrs[2], gm, assumed_align=4), - make_ptr(cutlass.Float32, ptrs[3], gm, assumed_align=4), - make_ptr(cutlass.Float32, ptrs[4], gm, assumed_align=4), - make_ptr(cutlass.BFloat16, ptrs[5], gm, assumed_align=2), - make_ptr(cutlass.BFloat16, ptrs[6], gm, assumed_align=2), - ) - fn = cute.compile( - _GdnRowsKernel(k_heads, v_heads, kdim, vdim, pdl, norm), - *cargs, - cutlass.Int32(1), - cutlass.Int32(1), - cutlass.Int64(1), - cutlass.Int32(1), - cutlass.Int32(1), - cutlass.Int32(1), - cutlass.Float32(1e-6), - cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), - options="--opt-level 2 --enable-tvm-ffi", - ) - _COMPILED[key] = fn + 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, @@ -437,8 +760,8 @@ def launch(proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, p class GdnDecodeRowsGatedNorm(OpWrapper): - """A single-token Gated-DeltaNet step per row over the fp32 slot pool, then the - sigmoid-gated RMS norm of its output per v-head. With `proj` the decode input + """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):: @@ -466,14 +789,23 @@ class GdnDecodeRowsGatedNorm(OpWrapper): why="2 rows, coalesced pool (slot stride 36 layers, slot 100), row 1 a pad slot (-1)", spec=dict(rows=2, pool="coalesced", pad=True), ), + Cell(why="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 `[slots, HV, V, K]` with K innermost (any - outer strides, multiples of 4) and 16-byte aligned; `K == 128`, `V % 8 == 0` with + 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 @@ -483,21 +815,29 @@ class GdnDecodeRowsGatedNorm(OpWrapper): the entry. Launch: one CTA of 8 warps a (v-head, row), `HV * rows` CTAs, PDL (off with - TRTLLM_ENABLE_PDL=0). GB300, graph replay, 461-slot pool, us (FlashInfer's - recurrence + the gated norm entry -> this): 128 rows 132.2 -> 118.8 (6.78 TB/s - against an in-place `mul_` over the same bytes at 6.83), 64 rows 70.7 -> 62.7, 16 - rows 19.1 -> 14.9, 4 rows 8.5 -> 8.4. + 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. Inside torch's default band for bf16 / fp32 of the composition's - float64 reference, not bitwise FlashInfer's kernel. + 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). Under graph replay the op reads the - projection after `griddepcontrol.wait` (16 rows). Not certified: `pdl=False`. + 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__( @@ -529,7 +869,7 @@ def reference( pdl: bool = True, ) -> tuple[torch.Tensor, torch.Tensor]: """`(out, pool_after)`: `gdn_decode_gated_norm_pdl`'s reference, the same contract -- - the fp64 recurrence (o rounded to bf16, the pool to fp32), then the gated norm + 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 @@ -553,12 +893,10 @@ def is_valid( assert proj.dtype == torch.bfloat16 and proj.stride(1) == 1, "bf16 proj, unit column stride" assert proj.shape[1] == 2 * k_heads * kdim + 2 * v_heads * vdim + 2 * v_heads, "proj width" assert proj.data_ptr() % 8 == 0 and proj.stride(0) % 4 == 0, "8-byte aligned proj rows" - assert pool.dtype == torch.float32 and pool.stride(3) == 1 and kd == kdim, ( - "fp32 pool [S, HV, V, K]" - ) - assert pool.data_ptr() % 16 == 0 and all(s % 4 == 0 for s in pool.stride()[:3]), ( - "16-byte aligned pool rows" - ) + assert 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" @@ -595,6 +933,14 @@ class GdnDecodeRows(OpWrapper): 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 = """ @@ -634,7 +980,7 @@ def reference( out: torch.Tensor, pdl: bool = True, ) -> tuple[torch.Tensor, torch.Tensor]: - """`(o, pool_after)`: the fp64 recurrence, o rounded to bf16 (the pool to fp32).""" + """`(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( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py index a24cbbf3865a..99ea41dfd83d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py @@ -26,6 +26,7 @@ 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 @@ -45,6 +46,9 @@ #: 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 @@ -554,8 +558,10 @@ def _gdn_gates_kernel( 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.""" + """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) @@ -569,7 +575,10 @@ def _gdn_gates_kernel( gate_input <= softplus_threshold, tl.log(1.0 + tl.exp(gate_input)), gate_input ) g_off = t64 * NUM_V_HEADS + h - tl.store(g_ptr + g_off, -tl.exp(A_log) * softplus, mask=m) + 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) @@ -670,6 +679,7 @@ def _launch( 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 " @@ -696,7 +706,7 @@ def _launch( x, weight, conv_states, query_start_loc, cache_indices, has_initial_state, q, k, v, num_k_heads, num_v_heads, head_k_dim, pad_slot_id, l2_norm_eps, ) # fmt: skip - _gdn_gates_kernel[(triton.cdiv(num_tokens, 64),)]( + _gdn_gates_kernel[(triton.cdiv(num_tokens, _GATES_BLOCK_T),)]( a, a.stride(0), a.stride(1), @@ -711,7 +721,8 @@ def _launch( softplus_threshold, NUM_V_HEADS=num_v_heads, HP=triton.next_power_of_2(num_v_heads), - BLOCK_T=64, + 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))]( @@ -751,7 +762,7 @@ class GdnPrefillConvPost(OpWrapper): 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) + 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) @@ -796,6 +807,15 @@ class GdnPrefillConvPost(OpWrapper): 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( @@ -819,8 +839,10 @@ class GdnPrefillConvPost(OpWrapper): whole conv pool unsliced, `pad_slot_id` -1, `l2_norm_eps` 1e-6 and `softplus_threshold` 20. `x`, `a`, `b` are column (and row) views, row stride 16480, of the one input-projection output; the outputs feed - `flashinfer_chunk_gated_delta_rule` (the Triton chunk on the verify path) - with `use_qk_l2norm_in_kernel=False`. + `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 @@ -830,8 +852,9 @@ class GdnPrefillConvPost(OpWrapper): first token's sequence by a binary search of `query_start_loc` (sync-free, capture-safe), and a thread whose 16 tokens and windows lie inside one sequence takes a straight-line branch, the rest read the pool - token by token. The gates (Triton): one row-coalesced `[64, 48]` tile a - program. Then a per-sequence window publish (Triton), after both, so no + 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 @@ -852,8 +875,9 @@ class GdnPrefillConvPost(OpWrapper): Certified over the 461-slot pool, seen as one layer of a two-layer buffer: three contexts of 37, 1 and 128 tokens; a mixed step's 8065-token context - half, the first rows of an 8192-row projection; and 128 short contexts (1 to - 64 tokens, some shorter than the window). + 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 @@ -902,6 +926,7 @@ def __call__( 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, @@ -921,6 +946,7 @@ def __call__( pad_slot_id, l2_norm_eps, softplus_threshold, + linear_gate, ) def reference( @@ -942,12 +968,14 @@ def reference( 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; `beta` the fp32 sigmoid.""" + 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] @@ -983,6 +1011,8 @@ def l2norm(part): 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) @@ -1005,6 +1035,7 @@ def is_valid( 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. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py index 1df6b9943f24..f681feffe334 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py @@ -44,6 +44,7 @@ both ways: every key is consumed or in the non-load set. """ +import os from dataclasses import dataclass from types import SimpleNamespace from typing import Any, Callable @@ -141,6 +142,26 @@ def memory_kind(ptr: int) -> int: _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 @@ -151,14 +172,27 @@ class _RegisteredHostBuffer: GB300 the device's random row gathers from its pages run an order of magnitude slower than from registered pageable memory.""" - def __init__(self, nbytes: int) -> None: - self._storage = torch.empty(nbytes, dtype=torch.uint8) - self.pointer = self._storage.data_ptr() - CUASSERT( - cudart.cudaHostRegister( - self.pointer, nbytes, _CUDA_HOST_REGISTER_PORTABLE | _CUDA_HOST_REGISTER_MAPPED + def __init__(self, nbytes: int, device=None) -> None: + # The registration faults the pages in on this thread, so under the default local + # allocation policy they land on the NUMA node of the CPU it runs on. Pinned to the + # CPUs local to `device` meanwhile, the table sits next to the GPU that gathers from + # it: on GB300, random row gathers from the other socket's memory run ~1.7x slower, + # and with a process per GPU half of them would otherwise land there. + local = _local_cpus(device) if device is not None else None + previous = os.sched_getaffinity(0) if local else None + if local: + os.sched_setaffinity(0, local) + try: + self._storage = torch.empty(nbytes, dtype=torch.uint8) + self.pointer = self._storage.data_ptr() + CUASSERT( + cudart.cudaHostRegister( + self.pointer, nbytes, _CUDA_HOST_REGISTER_PORTABLE | _CUDA_HOST_REGISTER_MAPPED + ) ) - ) + finally: + if previous: + os.sched_setaffinity(0, previous) # Bound now: at interpreter shutdown the module may already be torn down. self._unregister = cudart.cudaHostUnregister @@ -184,7 +218,7 @@ def _host_resident_empty( ) assert device.type == "cuda", f"host offload needs a cuda device to map into; got {device}" - host = _RegisteredHostBuffer(rows * cols) + 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 @@ -394,6 +428,64 @@ def _mtp_fc2_s(t: torch.Tensor, core) -> torch.Tensor: 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 diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/kernels.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/kernels.py index 358b5f424b43..e049d7483503 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/kernels.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/kernels.py @@ -186,6 +186,115 @@ def triton_qsa_decode_token_mapping( ) +# Largest batch the generation token mapping holds in one tile of request lengths. +QSA_GENERATION_MAPPING_MAX_SEQS = 256 + + +@triton.jit +def _qsa_generation_token_mapping_kernel( + kv_lens, + seq_lens, + request_indices, + logical_positions, + sequence_lengths, + visible_blocks, + num_seqs, + num_rows, + COMPRESS_RATIO: tl.constexpr, + BLOCK_ROWS: tl.constexpr, + BLOCK_SEQS: tl.constexpr, +): + rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + valid = rows < num_rows + seqs = tl.arange(0, BLOCK_SEQS) + lengths = tl.load(seq_lens + seqs, mask=seqs < num_seqs, other=0).to(tl.int64) + # A row belongs to the first request whose inclusive end exceeds it; the + # lengths of the requests ending at or before it sum to its request's start. + ends = tl.cumsum(lengths, axis=0) + ended = (ends[None, :] <= rows[:, None].to(tl.int64)) & (seqs < num_seqs)[None, :] + request = tl.sum(ended.to(tl.int32), axis=1) + request_start = tl.sum(tl.where(ended, lengths[None, :], 0), axis=1) + kv_len = tl.load(kv_lens + request, mask=valid, other=0) + seq_len = tl.load(seq_lens + request, mask=valid, other=0) + logical = kv_len.to(tl.int64) - seq_len.to(tl.int64) + rows.to(tl.int64) - request_start + # Floor division, matching the integer floor division of the general mapping. + visible_num = logical + 1 + visible = visible_num // COMPRESS_RATIO + visible = tl.where( + (visible_num % COMPRESS_RATIO != 0) & (visible_num < 0), visible - 1, visible + ) + tl.store(request_indices + rows, request, mask=valid) + tl.store(logical_positions + rows, logical, mask=valid) + tl.store(sequence_lengths + rows, kv_len, mask=valid) + tl.store(visible_blocks + rows, visible.to(tl.int32), mask=valid) + + +def triton_qsa_generation_token_mapping( + *, + kv_lens: torch.Tensor, + seq_lens: torch.Tensor, + request_indices: torch.Tensor, + logical_positions: torch.Tensor, + sequence_lengths: torch.Tensor, + visible_blocks: torch.Tensor, + num_rows: int, + compress_ratio: int, +) -> None: + """Build request and position metadata for a batch without context requests. + + Each request owns ``seq_lens`` consecutive packed rows, read on the device, + so verify groups of any width are covered: the same mapping as a cumulative + sum of the lengths followed by a right-sided search of every row, in one + launch for up to ``QSA_GENERATION_MAPPING_MAX_SEQS`` requests. + """ + if kv_lens.ndim != 1 or seq_lens.ndim != 1: + raise ValueError("QSA generation lengths must be one-dimensional") + num_seqs = kv_lens.numel() + if seq_lens.numel() != num_seqs: + raise ValueError("QSA generation KV and sequence lengths must have matching shapes") + if num_seqs > QSA_GENERATION_MAPPING_MAX_SEQS: + raise ValueError( + f"QSA generation token mapping supports at most {QSA_GENERATION_MAPPING_MAX_SEQS} " + f"requests, got {num_seqs}" + ) + outputs = ( + request_indices, + logical_positions, + sequence_lengths, + visible_blocks, + ) + if any(output.numel() < num_rows for output in outputs): + raise ValueError("QSA generation token-mapping outputs are too small") + if kv_lens.dtype != torch.int32 or seq_lens.dtype != torch.int32: + raise ValueError("QSA generation lengths must use int32 storage") + expected_dtypes = (torch.int32, torch.int64, torch.int32, torch.int32) + if any(tensor.dtype != dtype for tensor, dtype in zip(outputs, expected_dtypes)): + raise ValueError("QSA generation token-mapping outputs have incompatible dtypes") + _require_same_cuda_device("QSA generation token mapping", kv_lens, seq_lens, *outputs) + if compress_ratio <= 0: + raise ValueError(f"QSA compression ratio must be positive, got {compress_ratio}") + if num_rows == 0 or num_seqs == 0: + return + + # One specialization covers every batch and verify width, so no graph shape + # compiles a variant of its own. + block_rows = 32 + _qsa_generation_token_mapping_kernel[(triton.cdiv(num_rows, block_rows),)]( + kv_lens, + seq_lens, + request_indices, + logical_positions, + sequence_lengths, + visible_blocks, + num_seqs, + num_rows, + COMPRESS_RATIO=compress_ratio, + BLOCK_ROWS=block_rows, + BLOCK_SEQS=QSA_GENERATION_MAPPING_MAX_SEQS, + num_warps=4, + ) + + @triton.jit def _qsa_paged_kv_store_kernel( k, diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py index f689b1af5320..76320dfc25b3 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py @@ -10,7 +10,11 @@ from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from .constants import QSA_KEY_ROLE_INDEX -from .kernels import triton_qsa_unscale_block_table +from .kernels import ( + QSA_GENERATION_MAPPING_MAX_SEQS, + triton_qsa_generation_token_mapping, + triton_qsa_unscale_block_table, +) from .params import QSASparseMetadataParams @@ -256,6 +260,18 @@ def _refresh_qsa_token_mapping(self) -> None: compress_ratio=self.sparse_metadata_params.compress_ratio, ) return + if self.num_contexts == 0 and num_seqs <= QSA_GENERATION_MAPPING_MAX_SEQS: + triton_qsa_generation_token_mapping( + kv_lens=self.kv_lens_cuda_runtime[:num_seqs], + seq_lens=seq_lens, + request_indices=self.qsa_req_idx_per_token, + logical_positions=self.qsa_logical_positions, + sequence_lengths=self.qsa_sequence_lengths, + visible_blocks=self.qsa_visible_blocks, + num_rows=num_tokens, + compress_ratio=self.sparse_metadata_params.compress_ratio, + ) + return torch.cumsum( seq_lens, dim=0, diff --git a/tensorrt_llm/_torch/modules/fla/cache_manager.py b/tensorrt_llm/_torch/modules/fla/cache_manager.py index e40fd30a500b..26bf0a440173 100644 --- a/tensorrt_llm/_torch/modules/fla/cache_manager.py +++ b/tensorrt_llm/_torch/modules/fla/cache_manager.py @@ -280,6 +280,7 @@ def __init__(self, tokens_per_gen_step: int, use_flashinfer_ring: bool = False) self._seed_request_counter = 0 self._seed_rank_offset = 0 self.use_flashinfer_ring = use_flashinfer_ring + self.all_layer_commit_enabled = True self.replay_step_width = tokens_per_gen_step self.replay_history_size = max(MIN_REPLAY_HISTORY_SIZE, tokens_per_gen_step) self.prev_num_accepted_tokens: torch.Tensor | None = None @@ -504,7 +505,11 @@ def reset_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: def update(self, batch: MambaAcceptanceBatch) -> None: num_decodes = batch.num_generations - if not self.use_flashinfer_ring and num_decodes >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE: + if ( + self.all_layer_commit_enabled + and not self.use_flashinfer_ring + and num_decodes >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE + ): self.commit_all_layers(batch.attention_metadata, num_decodes) metadata = self.get_replay_metadata() if metadata is None: diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index 6dc63c740d56..2a6de9882595 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -79,11 +79,14 @@ def _build_replay_work_items_triton(state_indices, prev_num_accepted_tokens, replay_step_width, replay_history_size): """Single-launch build of the write-first replay partition. - Interchangeable with :func:`_build_replay_work_items_torch`; the caller - picks between them. Kept to one CTA because the write-first offsets come - from an in-block ``tl.cumsum``. + Interchangeable with :func:`_build_replay_work_items_torch` for up to + ``_FUSED_GDN_REPLAY_WORK_ITEMS_MAX_BATCH_SIZE`` decodes; the caller picks + between them. Kept to one CTA because the write-first offsets come from an + in-block ``tl.cumsum``. The block always spans the largest batch, so a + single specialization serves every batch size. """ num_decodes = state_indices.shape[0] + assert num_decodes <= _FUSED_GDN_REPLAY_WORK_ITEMS_MAX_BATCH_SIZE _prepare_gdn_replay_work_items_kernel[(1, )]( state_indices, prev_num_accepted_tokens, @@ -98,7 +101,7 @@ def _build_replay_work_items_triton(state_indices, prev_num_accepted_tokens, cache_slot_field=REPLAY_WORK_CACHE_SLOT, pnat_field=REPLAY_WORK_PNAT, cache_buf_idx_field=REPLAY_WORK_CACHE_BUF_IDX, - BLOCK_SIZE=triton.next_power_of_2(num_decodes), + BLOCK_SIZE=_FUSED_GDN_REPLAY_WORK_ITEMS_MAX_BATCH_SIZE, num_warps=4, ) @@ -432,8 +435,9 @@ def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int, replay_step_width = replay_metadata.replay_step_width replay_history_size = replay_metadata.replay_history_size - if (use_gdn_all_layer_commit - and num_decodes <= _FUSED_GDN_REPLAY_WORK_ITEMS_MAX_BATCH_SIZE): + if (num_decodes <= _FUSED_GDN_REPLAY_WORK_ITEMS_MAX_BATCH_SIZE + and prev_num_accepted_tokens.is_contiguous() + and cache_buf_idx.is_contiguous()): build_work_items = _build_replay_work_items_triton else: build_work_items = _build_replay_work_items_torch diff --git a/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py b/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py index f304c6d5c16e..07db09129a98 100644 --- a/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py +++ b/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py @@ -756,7 +756,17 @@ def use_gdn_flashinfer_ring_replay(self) -> bool: @property def use_gdn_cached_replay_all_layer_commit(self) -> bool: state = self._speculative_state - return isinstance(state, GDNReplayState) and state.has_bound_states + return ( + isinstance(state, GDNReplayState) + and state.has_bound_states + and state.all_layer_commit_enabled + ) + + def disable_gdn_cached_replay_all_layer_commit(self) -> None: + """Use when the model's replay kernel commits checkpoints itself.""" + state = self._speculative_state + if isinstance(state, GDNReplayState): + state.all_layer_commit_enabled = False def get_qwen4_exp_ple_cache_params(config, *, total_layers: int, is_draft: bool): diff --git a/tensorrt_llm/_torch/moe/kernels/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py b/tensorrt_llm/_torch/moe/kernels/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py index 88b088e34d67..97972c82dccd 100644 --- a/tensorrt_llm/_torch/moe/kernels/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py +++ b/tensorrt_llm/_torch/moe/kernels/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py @@ -35,7 +35,9 @@ 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, @@ -47,6 +49,30 @@ 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. @@ -206,6 +232,9 @@ def __init__( 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. @@ -230,6 +259,18 @@ def __init__( 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 @@ -333,7 +374,7 @@ def _setup_attributes(self): # Compute mma/cluster/tile shapes mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) - mma_inst_tile_k = 4 + mma_inst_tile_k = self.mma_inst_tile_k self.mma_tiler = ( self.mma_tiler[0], self.mma_tiler[1], @@ -1429,36 +1470,70 @@ def kernel( ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status) # TMA load A/B - 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, - ) tBgB_slice = tBgB_0[(None, mma_tile_coord_mnl[1], None, expert_idx)] tBgSFB_slice = tBgSFB_0[(None, slice_n, None, expert_idx)] - 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, - ) + 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() @@ -1873,6 +1948,8 @@ def kernel( 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). @@ -1956,7 +2033,13 @@ def kernel( 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): + 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], @@ -1976,8 +2059,9 @@ def kernel( ) cute.arch.cp_async_bulk_commit_group() - cute.arch.cp_async_bulk_wait_group(0, read=True) - self.epilog_sync_barrier.arrive_and_wait() + 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) @@ -1997,6 +2081,8 @@ def kernel( ) 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 # diff --git a/tensorrt_llm/_torch/speculative/draft_loop_ops.py b/tensorrt_llm/_torch/speculative/draft_loop_ops.py new file mode 100644 index 000000000000..b6491caf8725 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/draft_loop_ops.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One-launch versions of the one-model draft loop's row bookkeeping, for a generation-only +step that a CUDA graph captures. + +* `draft_row_select`: after the first draft step over all `p` rows of every request, each + request continues from its last accepted row `g = batch_indices[i] * p + accepted[i] - 1 + + num_ctx_tokens` (a negative `g` wraps, as torch indexing does). One launch writes the row + ids, the selected hidden rows (the next step's input and the LM-head rows) and the next + step's positions `position_ids[..., g] + 1`; given the attention lengths, it also moves + them to one row per request (`seq_lens[i] = 1`, `kv_lens[i] -= draft_len - accepted[i]`). +* `next_new_tokens`: the next verify step's input, `[accepted_tokens[batch_indices[i], + accepted[i] - 1], next_draft_tokens[i, :]]` per request, in one launch. + +Both equal the composed torch ops they replace element for element (integer arithmetic and +copies only). +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +#: Most requests one launch covers. +MAX_REQUESTS = 256 +#: Hidden-state columns a program copies. +_BLOCK_H = 1024 + + +@triton.jit(do_not_specialize=["num_rows", "num_ctx_tokens"]) +def _row_select_kernel( + batch_indices, + accepted, + hidden, + hidden_stride, + positions, + pos_stride, + out_ids, + out_hidden, + out_pos, + kv_lens, + seq_lens, + num_rows, + num_ctx_tokens, + width, + B, + P: tl.constexpr, + DRAFT_LEN: tl.constexpr, + HAS_KV: tl.constexpr, + HAS_SEQ: tl.constexpr, + POS_ROWS: tl.constexpr, + POS_BLOCK: tl.constexpr, + BLOCK_H: tl.constexpr, +): + i = tl.program_id(0) + hb = tl.program_id(1) + acc = tl.load(accepted + i) + g = tl.load(batch_indices + i).to(tl.int64) * P + acc.to(tl.int64) - 1 + ids = g + num_ctx_tokens + g = tl.where(ids < 0, ids + num_rows, ids) + c = hb * BLOCK_H + tl.arange(0, BLOCK_H) + m = c < width + vals = tl.load(hidden + g * hidden_stride + c, mask=m) + tl.store(out_hidden + i.to(tl.int64) * width + c, vals, mask=m) + if hb == 0: + tl.store(out_ids + i, ids) + r = tl.arange(0, POS_BLOCK) + rm = r < POS_ROWS + pos = tl.load(positions + r * pos_stride + g, mask=rm) + tl.store(out_pos + r * B + i, pos + 1, mask=rm) + if HAS_KV: + kv = tl.load(kv_lens + i) + tl.store(kv_lens + i, kv - (DRAFT_LEN - acc)) + if HAS_SEQ: + tl.store(seq_lens + i, 1) + + +def draft_row_select( + batch_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + rows_per_request: int, + num_ctx_tokens: int = 0, + kv_lens: torch.Tensor | None = None, + seq_lens: torch.Tensor | None = None, + draft_len: int = 0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """`(gather_ids, hidden_states[gather_ids], position_ids[..., gather_ids] + 1)` with + `gather_ids = batch_indices * rows_per_request + num_accepted_tokens - 1 + num_ctx_tokens` + (int64), for `num_accepted_tokens.shape[0]` requests. `hidden_states` is `[rows, width]` + with unit column stride; `position_ids` has the rows along its last dim. In place, when + given: `kv_lens[:n] -= draft_len - num_accepted_tokens` and `seq_lens[:n] = 1`.""" + n = num_accepted_tokens.shape[0] + rows, width = hidden_states.shape + assert 0 < n <= MAX_REQUESTS and hidden_states.stride(1) == 1 + assert position_ids.shape[-1] == rows + pos2 = position_ids.reshape(-1, rows) + assert pos2.stride(1) == 1 + dev = hidden_states.device + out_ids = torch.empty(n, dtype=torch.int64, device=dev) + out_hidden = torch.empty((n, width), dtype=hidden_states.dtype, device=dev) + out_pos = torch.empty(position_ids.shape[:-1] + (n,), dtype=position_ids.dtype, device=dev) + _row_select_kernel[(n, triton.cdiv(width, _BLOCK_H))]( + batch_indices, + num_accepted_tokens, + hidden_states, + hidden_states.stride(0), + pos2, + pos2.stride(0), + out_ids, + out_hidden, + out_pos, + kv_lens if kv_lens is not None else out_ids, + seq_lens if seq_lens is not None else out_ids, + rows, + num_ctx_tokens, + width, + n, + P=rows_per_request, + DRAFT_LEN=draft_len, + HAS_KV=kv_lens is not None, + HAS_SEQ=seq_lens is not None, + POS_ROWS=pos2.shape[0], + POS_BLOCK=triton.next_power_of_2(pos2.shape[0]), + BLOCK_H=_BLOCK_H, + num_warps=4, + ) + return out_ids, out_hidden, out_pos + + +@triton.jit(do_not_specialize=["num_gens"]) +def _next_new_tokens_kernel( + accepted_tokens, + accepted_stride, + accepted_cols, + drafts, + draft_stride, + batch_indices, + num_accepted, + out, + num_gens, + D: tl.constexpr, + W_BLOCK: tl.constexpr, + BLOCK: tl.constexpr, +): + i = tl.arange(0, BLOCK) + rows = i < num_gens + b = tl.load(batch_indices + i, mask=rows, other=0).to(tl.int64) + col = tl.load(num_accepted + i, mask=rows, other=1).to(tl.int64) - 1 + col = tl.where(col < 0, col + accepted_cols, col) + golden = tl.load(accepted_tokens + b * accepted_stride + col, mask=rows) + j = tl.arange(0, W_BLOCK) + is_draft = (j >= 1) & (j < D + 1) + dr = tl.load( + drafts + i[:, None].to(tl.int64) * draft_stride + (j - 1)[None, :], + mask=rows[:, None] & is_draft[None, :], + ) + vals = tl.where((j == 0)[None, :], golden[:, None], dr) + tl.store( + out + i[:, None].to(tl.int64) * (D + 1) + j[None, :], + vals, + mask=rows[:, None] & (j < D + 1)[None, :], + ) + + +def next_new_tokens( + accepted_tokens: torch.Tensor, + next_draft_tokens: torch.Tensor, + batch_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, +) -> torch.Tensor: + """`torch.concat([accepted_tokens[batch_indices, num_accepted_tokens - 1, None], + next_draft_tokens], 1)` for `num_accepted_tokens.shape[0]` requests, in one launch.""" + n = num_accepted_tokens.shape[0] + d = next_draft_tokens.shape[1] + assert 0 < n <= MAX_REQUESTS + assert accepted_tokens.stride(1) == 1 and next_draft_tokens.stride(1) == 1 + assert accepted_tokens.dtype == next_draft_tokens.dtype + out = torch.empty((n, d + 1), dtype=accepted_tokens.dtype, device=accepted_tokens.device) + _next_new_tokens_kernel[(1,)]( + accepted_tokens, + accepted_tokens.stride(0), + accepted_tokens.shape[1], + next_draft_tokens, + next_draft_tokens.stride(0), + batch_indices, + num_accepted_tokens, + out, + n, + D=d, + W_BLOCK=triton.next_power_of_2(d + 1), + BLOCK=MAX_REQUESTS, + num_warps=4, + ) + return out diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index bac560ef8f36..bf62938990d7 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -21,6 +21,8 @@ from ..pyexecutor.llm_request import LlmRequest from ..pyexecutor.resource_manager import BaseResourceManager, SlotManager from ..pyexecutor.scheduler import ScheduledRequests +from .draft_loop_ops import MAX_REQUESTS as _ROW_SELECT_MAX_REQUESTS +from .draft_loop_ops import draft_row_select from .interface import (INVALID_PROMPT_LOOKAHEAD_TOKEN, SpecMetadata, SpecWorkerBase) from .mtp import _select_mtp_position_ids @@ -644,8 +646,12 @@ def _forward_impl(self, # prepare next new tokens to support overlap scheduler next_new_tokens = self._prepare_next_new_tokens( - accepted_tokens, next_draft_tokens, - spec_metadata.batch_indices_cuda, batch_size, num_accepted_tokens) + accepted_tokens, + next_draft_tokens, + spec_metadata.batch_indices_cuda, + batch_size, + num_accepted_tokens, + fused=spec_metadata.is_cuda_graph and num_contexts == 0) attn_metadata.use_spec_decoding = True @@ -772,9 +778,23 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, runtime_draft_len = spec_metadata.runtime_draft_len num_gens = batch_size - num_contexts next_draft_tokens = [] - last_tokens_idx = torch.cumsum( - attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1 + # The context rows' last tokens; a generation-only step has none. + last_tokens_idx = ( + torch.cumsum(attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - + 1 if num_contexts > 0 else None) position_ids = inputs["position_ids"] + # A captured generation-only MTP step selects each request's continuing + # row (ids, hidden row, next positions) in one launch. + fused_rows = ( + self.is_mtp_eagle and spec_metadata.is_cuda_graph + and num_contexts == 0 and 0 < num_gens <= _ROW_SELECT_MAX_REQUESTS + and self.guided_decoder is None + and not (self.model_config is not None + and self.model_config.mapping.enable_lm_head_tp_in_adp)) + fused_hidden = fused_positions = None + # Whether the row selection also moved the attention lengths to one row + # per request (the first step's metadata update below). + fused_lengths = False # Sparse backends that can reuse one indexer selection across the draft # loop expose this setter; each indexer still decides whether its own @@ -824,15 +844,34 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, # selection is the prefix slice [0, batch_size); each # ``[gather_ids]`` site below takes a slice as readily as an # index tensor. - if i == 0: + if (i == 0 and fused_rows and hidden_states.dim() == 2 + and hidden_states.stride(1) == 1 + and inputs["position_ids"].stride(-1) == 1): + # The draft step above was the last reader of the step's + # lengths, so the row selection also rewrites them. + (gather_ids, fused_hidden, + fused_positions) = draft_row_select( + spec_metadata.batch_indices_cuda[:num_gens], + num_accepted_tokens[num_contexts:], + hidden_states, + inputs["position_ids"], + runtime_draft_len + 1, + attn_metadata.num_ctx_tokens, + kv_lens=getattr(attn_metadata, 'kv_lens_cuda', None), + seq_lens=attn_metadata._seq_lens_cuda, + draft_len=runtime_draft_len) + fused_lengths = True + elif i == 0: start_ids_gen = ( spec_metadata.batch_indices_cuda[:num_gens] * (runtime_draft_len + 1)).long() gather_ids_gen = (start_ids_gen + num_accepted_tokens[num_contexts:] - 1 + attn_metadata.num_ctx_tokens) - gather_ids = torch.concat( - [last_tokens_idx[:num_contexts], gather_ids_gen], dim=0) + gather_ids = ( + gather_ids_gen if num_contexts == 0 else torch.concat( + [last_tokens_idx[:num_contexts], gather_ids_gen], + dim=0)) else: gather_ids = slice(0, batch_size) @@ -876,7 +915,9 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, if self.is_mtp_eagle: # One row per request: the LM-head input on this step and # the draft input on the next. - request_hidden_states = hidden_states[gather_ids] + request_hidden_states = (fused_hidden if i == 0 + and fused_hidden is not None else + hidden_states[gather_ids]) if use_lm_head_tp_in_adp: token_count = request_hidden_states.view( -1, request_hidden_states.shape[-1]).shape[0] @@ -974,8 +1015,11 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, hidden_states = request_hidden_states else: hidden_states = hidden_states_to_save[gather_ids] - position_ids = (_select_mtp_position_ids( - inputs["position_ids"], gather_ids) + 1) + if i == 0 and fused_positions is not None: + position_ids = fused_positions + else: + position_ids = (_select_mtp_position_ids( + inputs["position_ids"], gather_ids) + 1) inputs = { "input_ids": new_draft_token, "position_ids": position_ids, @@ -987,7 +1031,8 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, # Update attn_metadata for the next iteration. if i == 0: attn_metadata._seq_lens[:batch_size].fill_(1) - attn_metadata._seq_lens_cuda[:batch_size].fill_(1) + if not fused_lengths: + attn_metadata._seq_lens_cuda[:batch_size].fill_(1) attn_metadata.on_update() has_kv_cache = inputs[ "attn_metadata"].kv_cache_manager is not None @@ -997,7 +1042,8 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, num_contexts].fill_(1) if has_kv_cache: attn_metadata.num_contexts = 0 - if hasattr(attn_metadata, 'kv_lens_cuda'): + if hasattr(attn_metadata, + 'kv_lens_cuda') and not fused_lengths: attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= ( runtime_draft_len - num_accepted_tokens[num_contexts:]) @@ -1241,7 +1287,8 @@ def prepare_1st_drafter_inputs( num_contexts:, :spec_metadata.runtime_draft_len + 1].flatten() # get draft inputs - input_ids = torch.concat([input_ids_ctx, input_ids_gen], dim=0) + input_ids = (input_ids_gen if input_ids_ctx.numel() == 0 else + torch.concat([input_ids_ctx, input_ids_gen], dim=0)) return { "input_ids": input_ids, diff --git a/tensorrt_llm/_torch/speculative/fused_accept.py b/tensorrt_llm/_torch/speculative/fused_accept.py new file mode 100644 index 000000000000..0b5127fb6f42 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/fused_accept.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Strict (token-equality) acceptance of a generation-only verify step in one launch. + +`SpecWorkerBase._sample_and_accept_draft_tokens_base` builds the accepted tokens and counts +of a uniform verify step from the target's sampled tokens with a chain of small kernels: a +compare, a cumulative product, a row sum, the int32 copies and, under +`TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`, the counter bump, pool hash and threshold of +the synthetic acceptance. `strict_accept` writes the same two outputs from one program: + +* `accepted_tokens[i, :] = target_tokens[i * (D + 1) : (i + 1) * (D + 1)]` (int32); +* `num_accepted_tokens[i] = 1 + ` the number of leading draft tokens equal to the target's + token at the same position, or, with a forced count, `base` (mode 1) or + `min(base + (pool[(c * CS + i * SS) & (len(pool) - 1)] < frac), D + 1)` with the counter + `c` advanced by one first (mode 2), all in the arithmetic of the composed ops (int64 hash, + float32 compare). +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +#: Most generation requests one launch covers (one program). +MAX_REQUESTS = 256 + +FORCE_NONE = 0 +FORCE_BASE = 1 +FORCE_RANDOM = 2 + + +# Runtime scalars are not specialized: one compiled variant per force mode and draft length, +# whichever batch size a graph first records it at. +@triton.jit(do_not_specialize=["num_gens", "frac", "base", "pool_mask"]) +def _strict_accept_kernel( + target, + draft, + accepted, + num_accepted, + counter, + pool, + num_gens, + frac, + base, + pool_mask, + FORCE: tl.constexpr, + D: tl.constexpr, + D_BLOCK: tl.constexpr, + BLOCK: tl.constexpr, + COUNTER_STRIDE: tl.constexpr, + SLOT_STRIDE: tl.constexpr, +): + i = tl.arange(0, BLOCK) + j = tl.arange(0, D_BLOCK) + rows = i < num_gens + cols = j < D + 1 + tgt = tl.load(target + i[:, None] * (D + 1) + j[None, :], mask=rows[:, None] & cols[None, :]) + tl.store( + accepted + i[:, None] * (D + 1) + j[None, :], + tgt.to(tl.int32), + mask=rows[:, None] & cols[None, :], + ) + if FORCE == 0: + is_draft = j < D + dr = tl.load( + draft + i[:, None] * D + j[None, :], mask=rows[:, None] & is_draft[None, :], other=0 + ) + miss = (dr.to(tl.int64) != tgt.to(tl.int64)) & is_draft[None, :] + first_miss = tl.min(tl.where(miss, j[None, :], D), axis=1) + count = 1 + first_miss + elif FORCE == 1: + count = tl.zeros([BLOCK], dtype=tl.int32) + base + else: + c = tl.load(counter) + 1 + tl.store(counter, c) + idx = (c * COUNTER_STRIDE + i.to(tl.int64) * SLOT_STRIDE) & pool_mask + rand = tl.load(pool + idx, mask=rows, other=1.0) + extra = (rand < frac).to(tl.int32) + count = tl.minimum(base + extra, D + 1) + tl.store(num_accepted + i, count.to(tl.int32), mask=rows) + + +_WARMED: set = set() + + +def _launch( + target, + draft, + accepted, + num_accepted, + counter, + pool, + num_gens, + frac, + base, + pool_mask, + force_mode, + draft_len, + counter_stride, + slot_stride, +) -> None: + _strict_accept_kernel[(1,)]( + target, + draft, + accepted, + num_accepted, + counter, + pool, + num_gens, + frac, + base, + pool_mask, + FORCE=force_mode, + D=draft_len, + D_BLOCK=triton.next_power_of_2(draft_len + 1), + BLOCK=MAX_REQUESTS, + COUNTER_STRIDE=counter_stride, + SLOT_STRIDE=slot_stride, + num_warps=4, + ) + + +def _warm( + device, draft_len: int, target_dtype, draft_dtype, counter_stride: int, slot_stride: int +) -> None: + """Launch every force mode once on scratch buffers, so that a mode a CUDA graph first + records (the synthetic acceptance is off in the eager graph warmup) is compiled and + loaded already.""" + key = (device, draft_len, target_dtype, draft_dtype, counter_stride, slot_stride) + if key in _WARMED or torch.cuda.is_current_stream_capturing(): + return + target = torch.zeros(draft_len + 1, dtype=target_dtype, device=device) + draft = torch.zeros(max(draft_len, 1), dtype=draft_dtype, device=device) + accepted = torch.empty(draft_len + 1, dtype=torch.int32, device=device) + num_accepted = torch.empty(1, dtype=torch.int32, device=device) + counter = torch.zeros(1, dtype=torch.int64, device=device) + pool = torch.zeros(1, dtype=torch.float32, device=device) + for mode in (FORCE_NONE, FORCE_BASE, FORCE_RANDOM): + # strict_accept passes num_accepted as both pointers outside FORCE_RANDOM. + mode_counter, mode_pool = ( + (counter, pool) if mode == FORCE_RANDOM else (num_accepted, num_accepted) + ) + _launch( + target, + draft, + accepted, + num_accepted, + mode_counter, + mode_pool, + 1, + 0.5, + 1, + 0, + mode, + draft_len, + counter_stride, + slot_stride, + ) + _WARMED.add(key) + + +def strict_accept( + target_tokens: torch.Tensor, + draft_tokens: torch.Tensor, + num_gens: int, + draft_len: int, + force_mode: int = FORCE_NONE, + force_base: int = 0, + force_frac: float = 0.0, + counter: torch.Tensor | None = None, + pool: torch.Tensor | None = None, + counter_stride: int = 0, + slot_stride: int = 0, +) -> tuple[torch.Tensor, torch.Tensor]: + """`(accepted_tokens [num_gens, draft_len + 1] int32, num_accepted_tokens [num_gens] + int32)` of a generation-only uniform verify step (see the module docstring). + `target_tokens` holds `num_gens * (draft_len + 1)` sampled tokens, `draft_tokens` the + `num_gens * draft_len` proposals (any integer dtype, contiguous). Mode `FORCE_RANDOM` + advances `counter` (int64 `[1]`) in place and reads `pool` (float32, power-of-two + length). Graph-capturable; `num_gens` at most `MAX_REQUESTS`.""" + assert 0 < num_gens <= MAX_REQUESTS + assert target_tokens.is_contiguous() and draft_tokens.is_contiguous() + dev = target_tokens.device + accepted = torch.empty((num_gens, draft_len + 1), dtype=torch.int32, device=dev) + num_accepted = torch.empty(num_gens, dtype=torch.int32, device=dev) + if force_mode == FORCE_RANDOM: + assert counter is not None and pool is not None + pool_size = pool.numel() + assert pool_size & (pool_size - 1) == 0 + pool_mask = pool_size - 1 + else: + counter = pool = num_accepted + pool_mask = 0 + _warm(dev, draft_len, target_tokens.dtype, draft_tokens.dtype, counter_stride, slot_stride) + _launch( + target_tokens, + draft_tokens, + accepted, + num_accepted, + counter, + pool, + num_gens, + float(force_frac), + int(force_base), + pool_mask, + force_mode, + draft_len, + counter_stride, + slot_stride, + ) + return accepted, num_accepted diff --git a/tensorrt_llm/_torch/speculative/greedy_argmax.py b/tensorrt_llm/_torch/speculative/greedy_argmax.py new file mode 100644 index 000000000000..4d5b72fd4501 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/greedy_argmax.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Row-wise argmax of large `[rows, vocab]` logits (greedy sampling) in Triton. + +`torch.argmax(logits, dim=-1)` runs one generic reduce kernel that reads fp32 logits of a +248320-entry vocabulary well under the HBM rate (GB300, graph replay, HBM-cold: 76 us at +256 rows, 216 at 1024). Here each row is split into `splits` chunks; one program reduces a +chunk to a (value, index) partial, and a second launch combines a row's partials. With one +split a row the first launch writes the result itself. + +The result equals `torch.argmax` element for element: NaN ranks above every number (the +first NaN of a row wins), equal values (including -0.0 and +0.0) resolve to the lowest +index, and an all -inf row returns 0. bf16 / fp16 inputs are compared after an exact +widening to fp32. + +Inside a CUDA graph a few rows (`small_rows=True`, at most `SMALL_ROWS`) take the Triton +kernels too, at `_SMALL_SPLITS` chunks a row: torch's generic reduce replays as a memset and a +31-block reduce kernel there (~11 us a call at 1-4 rows of 248320), two short launches here. +Eager callers keep torch.argmax for such inputs (fewer host launches). +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +#: Smallest `rows * vocab` routed to the Triton kernels; smaller inputs keep torch.argmax. +MIN_NUMEL = 1 << 22 +#: Elements a program loads per iteration (4 warps: 16 a thread, four 16-byte fp32 loads). +_BLOCK = 2048 +_WARPS = 4 +_INT32_MAX = 2**31 - 1 +#: `(min_rows, splits)`, largest first: chunks a row by the row count, so that a launch has +#: ~1000-4000 programs. GB300 sweep over splits x block x warps at 64-1024 rows of 248320. +_SPLITS = ((1024, 1), (512, 2), (256, 8), (128, 16), (0, 32)) +#: The table for 16-bit logits (half the bytes a chunk): GB300 sweep over splits x block x +#: warps at 64-1024 rows of 248320 bf16, at this block and warp count. +_SPLITS_16BIT = ((512, 4), (128, 8), (64, 16), (0, 32)) +#: Most rows `small_rows` callers route to the Triton kernels below `MIN_NUMEL`. +SMALL_ROWS = 16 +#: Chunks a row for `small_rows` launches (GB300 graph replay, 248320 fp32 columns, 1-16 rows: +#: 2.4-5.0 us against 10.6-16.4 us for torch.argmax; 32 and 128 chunks within 0.5 us). +_SMALL_SPLITS = 64 + + +@triton.jit +def _combine(v1, i1, v2, i2): + """torch's `greater_or_nan` order: NaN first, then larger values, ties to the lower index.""" + n1 = v1 != v1 + n2 = v2 != v2 + lower = i1 < i2 + pick1 = tl.where(n1, (~n2) | lower, (~n2) & ((v1 > v2) | ((v1 == v2) & lower))) + return tl.where(pick1, v1, v2), tl.where(pick1, i1, i2) + + +@triton.jit +def _argmax_chunk_kernel( + x_ptr, + val_ptr, + idx_ptr, + out_ptr, + stride_row, + n_cols, + CHUNK: tl.constexpr, + BLOCK: tl.constexpr, + SPLITS: tl.constexpr, +): + row = tl.program_id(0) + split = tl.program_id(1) + base = x_ptr + row.to(tl.int64) * stride_row + start = split * CHUNK + offs = start + tl.arange(0, BLOCK) + mask = offs < n_cols + acc_v = tl.load(base + offs, mask=mask, other=float("-inf")).to(tl.float32) + acc_i = tl.where(mask, offs, 2147483647) + # Each lane walks its column positions in increasing index order, so a later element + # replaces the lane's running pick only when strictly larger, or NaN over a number. + for k in tl.range(BLOCK, CHUNK, BLOCK): + o = start + k + tl.arange(0, BLOCK) + m = o < n_cols + v = tl.load(base + o, mask=m, other=float("-inf")).to(tl.float32) + take = (acc_v == acc_v) & ((v != v) | (v > acc_v)) + acc_v = tl.where(take, v, acc_v) + acc_i = tl.where(take, o, acc_i) + best_v, best_i = tl.reduce((acc_v, acc_i), 0, _combine) + if SPLITS == 1: + tl.store(out_ptr + row, best_i.to(out_ptr.dtype.element_ty)) + else: + tl.store(val_ptr + row * SPLITS + split, best_v) + tl.store(idx_ptr + row * SPLITS + split, best_i) + + +@triton.jit +def _argmax_final_kernel(val_ptr, idx_ptr, out_ptr, SPLITS: tl.constexpr): + row = tl.program_id(0) + offs = row * SPLITS + tl.arange(0, SPLITS) + best_v, best_i = tl.reduce((tl.load(val_ptr + offs), tl.load(idx_ptr + offs)), 0, _combine) + tl.store(out_ptr + row, best_i.to(out_ptr.dtype.element_ty)) + + +def _table(dtype: torch.dtype) -> tuple: + return _SPLITS if dtype == torch.float32 else _SPLITS_16BIT + + +def _splits(rows: int, n_cols: int, dtype: torch.dtype = torch.float32) -> int: + if rows * n_cols < MIN_NUMEL: + return _SMALL_SPLITS + return next(splits for min_rows, splits in _table(dtype) if rows >= min_rows) + + +def _buffers(rows: int, splits: int, out: torch.Tensor): + """The partial (value, index) buffers of a launch; with one split the result itself.""" + if splits == 1: + return out, out + return ( + torch.empty(rows * splits, dtype=torch.float32, device=out.device), + torch.empty(rows * splits, dtype=torch.int32, device=out.device), + ) + + +_WARMED: set = set() + + +def _warm(logits: torch.Tensor, out_dtype: torch.dtype) -> None: + """Compile every split variant for this input's specialization (dtype, vocab, row stride, + output dtype) without launching, so that a later row count never compiles mid-serving.""" + key = (logits.dtype, logits.shape[1], logits.stride(0), out_dtype, logits.device) + if key in _WARMED: + return + n_cols = logits.shape[1] + out = torch.empty(1, dtype=out_dtype, device=logits.device) + for splits in sorted({s for _, s in _table(logits.dtype)} | {_SMALL_SPLITS}): + val, idx = _buffers(1, splits, out) + chunk = triton.cdiv(triton.cdiv(n_cols, splits), _BLOCK) * _BLOCK + _argmax_chunk_kernel.warmup( + logits, + val, + idx, + out, + logits.stride(0), + n_cols, + CHUNK=chunk, + BLOCK=_BLOCK, + SPLITS=splits, + num_warps=_WARPS, + grid=(1,), + ) + if splits > 1: + _argmax_final_kernel.warmup(val, idx, out, SPLITS=splits, num_warps=1, grid=(1,)) + _WARMED.add(key) + + +def eligible(logits: torch.Tensor, small_rows: bool = False) -> bool: + """Whether `greedy_argmax` takes the Triton kernels for `logits` (else torch.argmax): + from `MIN_NUMEL` elements on, or with `small_rows` at 1 to `SMALL_ROWS` rows.""" + return ( + logits.is_cuda + and logits.dim() == 2 + and logits.dtype in (torch.float32, torch.bfloat16, torch.float16) + and logits.stride(-1) == 1 + and ( + logits.numel() >= MIN_NUMEL + or (small_rows and 0 < logits.shape[0] <= SMALL_ROWS and logits.shape[1] > 0) + ) + and logits.shape[-1] < _INT32_MAX + ) + + +def greedy_argmax( + logits: torch.Tensor, out_dtype: torch.dtype = torch.int64, small_rows: bool = False +) -> torch.Tensor: + """`torch.argmax(logits, dim=-1).to(out_dtype)` for a 2-D `[rows, vocab]` tensor: the + Triton kernels when `eligible(logits, small_rows)`, else torch.argmax. Graph-capturable + (no host sync); callers pass `small_rows=True` for steps that a CUDA graph replays. The + first call for an input specialization compiles every split variant (`_warm`); the row + count then only selects among them.""" + if not eligible(logits, small_rows): + return torch.argmax(logits, dim=-1).to(out_dtype) + _warm(logits, out_dtype) + rows, n_cols = logits.shape + out = torch.empty(rows, dtype=out_dtype, device=logits.device) + splits = _splits(rows, n_cols, logits.dtype) + chunk = triton.cdiv(triton.cdiv(n_cols, splits), _BLOCK) * _BLOCK + val, idx = _buffers(rows, splits, out) + _argmax_chunk_kernel[(rows, splits)]( + logits, + val, + idx, + out, + logits.stride(0), + n_cols, + CHUNK=chunk, + BLOCK=_BLOCK, + SPLITS=splits, + num_warps=_WARPS, + ) + if splits > 1: + _argmax_final_kernel[(rows,)](val, idx, out, SPLITS=splits, num_warps=1) + return out diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 97b67a555523..299543914f31 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -14,6 +14,7 @@ # limitations under the License. import copy +import functools import os from abc import ABC, abstractmethod from contextlib import contextmanager @@ -33,6 +34,8 @@ TrtllmAttentionMetadata) from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE from ..pyexecutor.resource_manager import ResourceManagerType +from . import draft_loop_ops, fused_accept +from .greedy_argmax import greedy_argmax if TYPE_CHECKING: from ..pyexecutor.guided_decoder import CapturableGuidedDecoder @@ -456,6 +459,53 @@ def from_string(name: Optional[str]) -> "SpeculativeDecodingMode": _RNG_SLOT_SPAN = 32 +@functools.lru_cache(maxsize=4096) +def _normalize_one_model_sampling_params( + *, + temperature: Optional[float], + top_k: Optional[int], + top_p: Optional[float], + min_p: Optional[float], +) -> tuple[float, int, float, float, bool]: + """Convert request sampling params into normalized per-request scalars. + + Interpreted as TorchSampler does (resolve_row_params); greedy rows take + DISABLE_TEMP_VAL, and a disabled top_k takes DISABLE_TOPK_VAL rather than 0. + A pure function of its arguments, so batches re-scanned every step look the + distinct parameter sets up instead of resolving them again. + """ + from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import \ + GREEDY_TEMPERATURE_THRESHOLD + from tensorrt_llm._torch.pyexecutor.sampler.sampling_resolution import ( + SamplingMode, resolve_row_params) + + # Sentinel temperature for greedy / temperature-disabled rows. Must stay + # strictly below GREEDY_TEMPERATURE_THRESHOLD so the sampling kernels + # recognize these rows as greedy; small enough that even if a row were + # (incorrectly) sampled, softmax(logits / val) is effectively one-hot, + # and non-zero to avoid division by 0. + DISABLE_TEMP_VAL = GREEDY_TEMPERATURE_THRESHOLD / 10 + # Very large values disable topk. + DISABLE_TOPK_VAL = torch.iinfo(torch.int32).max + + row = resolve_row_params(temperature=temperature, + top_k=top_k, + top_p=top_p, + min_p=min_p) + is_greedy = row.mode is SamplingMode.GREEDY + # The sentinel is one-hot under softmax, so it must not reach a non-greedy + # row. + normalized_temperature = DISABLE_TEMP_VAL if is_greedy else row.temperature + normalized_top_k = row.top_k if row.top_k > 0 else DISABLE_TOPK_VAL + return ( + normalized_temperature, + normalized_top_k, + row.top_p, + row.min_p, + is_greedy, + ) + + @dataclass class SpecMetadata: """ @@ -880,21 +930,24 @@ def _populate_request_rng_state( dtype=torch.int64, device='cuda') - def _upload(dst: torch.Tensor, values: list[int]) -> None: - dst[:len(values)].copy_(torch.tensor(values, - dtype=torch.int64, - pin_memory=prefer_pinned()), - non_blocking=True) - # An all-greedy batch runs the argmax graph, which reads none of these # buffers. The offset windows above still advance, so skipping the # copies changes no later batch's streams. if self.is_all_greedy_sample: return - _upload(self.request_seeds, request_seeds) - _upload(self.request_offsets, request_offsets) - _upload(self.seeds, flat_seeds) - _upload(self.offsets, flat_offsets) + # One pinned staging tensor holds all four lists; each buffer is + # filled from its own slice. + dsts = (self.request_seeds, self.request_offsets, self.seeds, + self.offsets) + lists = (request_seeds, request_offsets, flat_seeds, flat_offsets) + staging = torch.tensor(sum(lists, []), + dtype=torch.int64, + pin_memory=prefer_pinned()) + start = 0 + for dst, values in zip(dsts, lists): + end = start + len(values) + dst[:len(values)].copy_(staging[start:end], non_blocking=True) + start = end def prepare_rejection_sampling_buffers(self): """ @@ -1207,53 +1260,6 @@ def _scan_one_model_sampling( """ from tensorrt_llm._torch.pyexecutor.llm_request import ( LlmRequestState, get_request_tokens_per_gen_step) - from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import \ - GREEDY_TEMPERATURE_THRESHOLD - from tensorrt_llm._torch.pyexecutor.sampler.sampling_resolution import ( - SamplingMode, resolve_row_params) - - # Sentinel temperature for greedy / temperature-disabled rows. Must stay - # strictly below GREEDY_TEMPERATURE_THRESHOLD so the sampling kernels - # recognize these rows as greedy; small enough that even if a row were - # (incorrectly) sampled, softmax(logits / val) is effectively one-hot, - # and non-zero to avoid division by 0. - DISABLE_TEMP_VAL = GREEDY_TEMPERATURE_THRESHOLD / 10 - # Very large values disable topk. - DISABLE_TOPK_VAL = torch.iinfo(torch.int32).max - - def _normalize_request_sampling_params( - *, - temperature: Optional[float], - top_k: Optional[int], - top_p: Optional[float], - min_p: Optional[float], - ) -> tuple[float, int, float, float, bool]: - """Convert request sampling params into normalized per-request scalars. - - Interpreted as TorchSampler does (resolve_row_params); greedy rows - take the sentinels above, and a disabled top_k takes - DISABLE_TOPK_VAL rather than 0. - """ - row = resolve_row_params(temperature=temperature, - top_k=top_k, - top_p=top_p, - min_p=min_p) - is_greedy = row.mode is SamplingMode.GREEDY - # The sentinel is one-hot under softmax, so it must not reach a - # non-greedy row. - normalized_temperature = (DISABLE_TEMP_VAL - if is_greedy else row.temperature) - normalized_top_k = row.top_k if row.top_k > 0 else DISABLE_TOPK_VAL - normalized_top_p = row.top_p - normalized_min_p = row.min_p - - return ( - normalized_temperature, - normalized_top_k, - normalized_top_p, - normalized_min_p, - is_greedy, - ) # Phase 1: collect per-request flags and normalized values. per_request_normalized: list[tuple[float, int, float, float, int]] = [] @@ -1286,7 +1292,7 @@ def _normalize_request_sampling_params( tp_val, mp_val, is_greedy, - ) = _normalize_request_sampling_params( + ) = _normalize_one_model_sampling_params( temperature=temp_val, top_k=tk_val, top_p=tp_val, @@ -1953,6 +1959,38 @@ def _ensure_force_accept_rng_state(self, device: torch.device) -> None: dtype=torch.int64, device=device) + def _force_accept_plan(self, runtime_draft_len: int, spec_metadata=None): + """``(mode, base_total, frac_part)`` of the synthetic acceptance this + step applies (see :meth:`_apply_force_accepted_tokens`): mode + ``FORCE_NONE`` keeps the verified counts (no override configured, or + the eager CUDA-graph warmup), ``FORCE_BASE`` sets every generation + request to ``base_total``, ``FORCE_RANDOM`` adds one token with + probability ``frac_part``. + """ + if self.force_num_accepted_tokens == 0.0: + return fused_accept.FORCE_NONE, 0, 0.0 + + if spec_metadata is not None: + is_warmup = (spec_metadata.is_cuda_graph + and not torch.cuda.is_current_stream_capturing()) + if is_warmup: + return fused_accept.FORCE_NONE, 0, 0.0 + + # Decompose into a deterministic integer part (always accepted) and a + # probabilistic fractional part. ``int(...)`` truncates toward zero, + # which matches floor for the supported non-negative range. + int_part = int(self.force_num_accepted_tokens) + frac_part = self.force_num_accepted_tokens - int_part + + # ``num_accepted_tokens`` counts the target token + accepted draft + # tokens, so the maximum reachable value is ``runtime_draft_len + 1``. + max_total = runtime_draft_len + 1 + base_total = min(int_part + 1, max_total) + + if frac_part > 0.0 and base_total < max_total: + return fused_accept.FORCE_RANDOM, base_total, frac_part + return fused_accept.FORCE_BASE, base_total, frac_part + def _apply_force_accepted_tokens(self, num_accepted_tokens, num_contexts, @@ -1993,27 +2031,13 @@ def _apply_force_accepted_tokens(self, Returns: Modified num_accepted_tokens tensor. """ - if self.force_num_accepted_tokens == 0.0: + force_mode, base_total, frac_part = self._force_accept_plan( + runtime_draft_len, spec_metadata) + if force_mode == fused_accept.FORCE_NONE: return num_accepted_tokens - - if spec_metadata is not None: - is_warmup = (spec_metadata.is_cuda_graph - and not torch.cuda.is_current_stream_capturing()) - if is_warmup: - return num_accepted_tokens - - # Decompose into a deterministic integer part (always accepted) and a - # probabilistic fractional part. ``int(...)`` truncates toward zero, - # which matches floor for the supported non-negative range. - int_part = int(self.force_num_accepted_tokens) - frac_part = self.force_num_accepted_tokens - int_part - - # ``num_accepted_tokens`` counts the target token + accepted draft - # tokens, so the maximum reachable value is ``runtime_draft_len + 1``. max_total = runtime_draft_len + 1 - base_total = min(int_part + 1, max_total) - if frac_part > 0.0 and base_total < max_total: + if force_mode == fused_accept.FORCE_RANDOM: self._ensure_force_accept_rng_state(num_accepted_tokens.device) # ``num_gens`` is fixed at CUDA-graph capture time (graphs are @@ -2083,6 +2107,12 @@ def _sample_and_accept_draft_tokens_base( if isinstance(logits, torch.Tensor) and logits.dim() == 1: logits = logits.unsqueeze(0) + num_gens = batch_size - num_contexts + if self._can_use_fused_strict_accept(logits, num_contexts, num_gens, + runtime_draft_len, spec_metadata): + return self._fused_strict_accept(logits, draft_tokens, num_gens, + runtime_draft_len, spec_metadata) + # Sample tokens using per-request sampling parameters target_tokens = self._sample_tokens_for_batch(logits, spec_metadata, num_contexts, batch_size) @@ -2148,6 +2178,44 @@ def _sample_and_accept_draft_tokens_base( spec_metadata.verify_lens.to(num_accepted_tokens.dtype)) return accepted_tokens, num_accepted_tokens + @staticmethod + def _can_use_fused_strict_accept(logits, num_contexts, num_gens, + runtime_draft_len, spec_metadata) -> bool: + """Whether a strict acceptance runs as ``fused_accept.strict_accept``: + a greedy, generation-only, uniform verify step of a CUDA-graph batch + (eager steps keep the composed ops and their fewer host launches).""" + return (spec_metadata is not None and spec_metadata.is_cuda_graph + and spec_metadata.is_all_greedy_sample + and not spec_metadata.is_ragged_verify and num_contexts == 0 + and 0 < num_gens <= fused_accept.MAX_REQUESTS + and isinstance(logits, torch.Tensor) and logits.dim() == 2 + and logits.shape[0] == num_gens * (runtime_draft_len + 1)) + + def _fused_strict_accept(self, logits, draft_tokens, num_gens, + runtime_draft_len, spec_metadata): + """The strict acceptance of `_sample_and_accept_draft_tokens_base` + for the steps `_can_use_fused_strict_accept` admits: the greedy target + tokens, then one launch for the accepted tokens and counts, including + the synthetic acceptance override.""" + target_tokens = self._sample_tokens_for_batch(logits, spec_metadata, 0, + num_gens) + force_mode, base_total, frac_part = self._force_accept_plan( + runtime_draft_len, spec_metadata) + if force_mode == fused_accept.FORCE_RANDOM: + self._ensure_force_accept_rng_state(logits.device) + return fused_accept.strict_accept( + target_tokens, + draft_tokens.contiguous(), + num_gens, + runtime_draft_len, + force_mode=force_mode, + force_base=base_total, + force_frac=frac_part, + counter=self._force_accept_rng_counter, + pool=self._force_accept_rng_pool, + counter_stride=_FORCE_ACCEPT_RNG_COUNTER_STRIDE, + slot_stride=_FORCE_ACCEPT_RNG_SLOT_STRIDE) + def _apply_occurrence_penalties( self, logits: torch.Tensor, draft_tokens: torch.Tensor, num_contexts: int, batch_size: int, @@ -2950,7 +3018,10 @@ def greedy_sample_draft_with_tp_gather(self, # requests with replicated full-vocab logits, so a per-rank argmax is # the correct proposal and a gather would desync the ranks (see # _draft_logits_are_sharded). Plain argmax; caller applies d2t. - return torch.argmax(logits, dim=-1).type(torch.int32) + return greedy_argmax(logits, + torch.int32, + small_rows=spec_metadata is not None + and spec_metadata.is_cuda_graph) def advanced_sample_draft_block(self, gen_logits: torch.Tensor, spec_metadata: "SpecMetadata", @@ -3167,9 +3238,13 @@ def _rollback_guided_decoder_after_verify(self, num_accepted_tokens): if self.guided_decoder is not None: self.guided_decoder.rollback_rejected_batch(num_accepted_tokens) - def _prepare_next_new_tokens(self, accepted_tokens, next_draft_tokens, - batch_indices_cuda, batch_size, - num_accepted_tokens): + def _prepare_next_new_tokens(self, + accepted_tokens, + next_draft_tokens, + batch_indices_cuda, + batch_size, + num_accepted_tokens, + fused: bool = False): """ Prepare next_new_tokens for overlap scheduler support. @@ -3179,10 +3254,21 @@ def _prepare_next_new_tokens(self, accepted_tokens, next_draft_tokens, batch_indices_cuda: Batch indices tensor batch_size: Number of requests num_accepted_tokens: [batch_size] - Number of accepted tokens per request + fused: One launch (``draft_loop_ops.next_new_tokens``) for a step + a CUDA graph captures; same values. Returns: next_new_tokens: [batch_size, runtime_draft_len + 1] - Input tokens for next iteration """ + if (fused and next_draft_tokens.dim() == 2 + and 0 < batch_size <= draft_loop_ops.MAX_REQUESTS + and num_accepted_tokens.shape[0] == batch_size + and accepted_tokens.dtype == next_draft_tokens.dtype + and accepted_tokens.stride(1) == 1 + and next_draft_tokens.stride(1) == 1): + return draft_loop_ops.next_new_tokens( + accepted_tokens, next_draft_tokens, + batch_indices_cuda[:batch_size], num_accepted_tokens) next_new_tokens = accepted_tokens[batch_indices_cuda[:batch_size], num_accepted_tokens - 1].unsqueeze(1) next_new_tokens = torch.concat([next_new_tokens, next_draft_tokens], @@ -3334,6 +3420,7 @@ def _sample_tokens_for_batch( seed=seed, offset=offset) else: - sampled_tokens = torch.argmax(logits, dim=-1) + sampled_tokens = greedy_argmax( + logits, small_rows=spec_metadata.is_cuda_graph) return sampled_tokens diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index f5183be9523d..c186e5df1858 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -342,6 +342,11 @@ l0_b200: - unittest/_torch/speculative/test_chain_speculative_sampling_op.py - unittest/_torch/speculative/test_dynamic_tree_rejection_op.py - unittest/_torch/speculative/test_dynamic_tree_sampling_dispatch.py + # Captured MTP decode-step bookkeeping kernels (Triton): the fused strict + # acceptance, the draft loop's row selection and the greedy argmax. + - unittest/_torch/speculative/test_fused_strict_accept.py + - unittest/_torch/speculative/test_draft_loop_ops.py + - unittest/_torch/speculative/test_greedy_argmax.py # One-model speculative decoding step-copy kernels and the engine's decode # step that runs them (SM 100; they skip on other GPUs). - unittest/_torch/kernels/test_spec_step_copies.py diff --git a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py index 864f8d66832a..aa241361d2ae 100644 --- a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py +++ b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py @@ -16,10 +16,12 @@ average_pool_qsa_keys, ) from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import ( + QSA_GENERATION_MAPPING_MAX_SEQS, _expand_launch, _qsa_index_scores_query_tile, triton_qsa_decode_pre_indexer, triton_qsa_decode_token_mapping, + triton_qsa_generation_token_mapping, triton_qsa_mtp_reuse_topk, triton_qsa_paged_index_scores, triton_qsa_paged_kv_store, @@ -534,6 +536,119 @@ def test_qsa_decode_token_mapping_matches_reference(rows: int) -> None: torch.testing.assert_close(visible_blocks, ((kv_lens - seq_lens + 1) // 4)) +def _general_token_mapping(kv_lens, seq_lens, num_tokens, compress_ratio): + """The cumulative-sum and search mapping used for batches with context rows.""" + num_seqs = seq_lens.numel() + cu_seq_lens = torch.zeros(num_seqs + 1, dtype=torch.int64, device=seq_lens.device) + torch.cumsum(seq_lens, dim=0, dtype=torch.int64, out=cu_seq_lens[1:]) + token_ids = torch.arange(num_tokens, dtype=torch.int64, device=seq_lens.device) + req_idx = torch.searchsorted(cu_seq_lens[1:], token_ids, right=True) + cached_lens = kv_lens.to(torch.int64) - seq_lens.to(torch.int64) + logical = cached_lens[req_idx] + token_ids - cu_seq_lens[:num_seqs][req_idx] + return ( + req_idx.to(torch.int32), + logical, + kv_lens[req_idx], + ((logical + 1) // compress_ratio).to(torch.int32), + ) + + +@pytest.mark.parametrize( + "seq_lens", + [ + [4], + [4] * 3, + [4] * 32, + [4] * QSA_GENERATION_MAPPING_MAX_SEQS, + [2] * 7, + [1, 4, 2, 0, 3, 4, 1], + [0, 5], + ], + ids=["one", "three", "batch32", "batch-max", "width2", "ragged", "leading-empty"], +) +@pytest.mark.parametrize("cached", [0, 3, 1029]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_qsa_generation_token_mapping_matches_general_mapping( + seq_lens: list[int], cached: int +) -> None: + seq_lens_cuda = torch.tensor(seq_lens, dtype=torch.int32, device="cuda") + kv_lens = seq_lens_cuda + cached + torch.arange(len(seq_lens), dtype=torch.int32, device="cuda") + num_tokens = sum(seq_lens) + outputs = ( + torch.full((num_tokens + 3,), -7, dtype=torch.int32, device="cuda"), + torch.full((num_tokens + 3,), -7, dtype=torch.int64, device="cuda"), + torch.full((num_tokens + 3,), -7, dtype=torch.int32, device="cuda"), + torch.full((num_tokens + 3,), -7, dtype=torch.int32, device="cuda"), + ) + + triton_qsa_generation_token_mapping( + kv_lens=kv_lens, + seq_lens=seq_lens_cuda, + request_indices=outputs[0], + logical_positions=outputs[1], + sequence_lengths=outputs[2], + visible_blocks=outputs[3], + num_rows=num_tokens, + compress_ratio=4, + ) + + expected = _general_token_mapping(kv_lens, seq_lens_cuda, num_tokens, 4) + for actual, reference in zip(outputs, expected): + torch.testing.assert_close(actual[:num_tokens], reference, rtol=0, atol=0) + # Rows past the batch are left untouched. + assert (actual[num_tokens:] == -7).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_qsa_generation_token_mapping_floors_negative_positions() -> None: + """Visible blocks floor like the general mapping even below position zero.""" + seq_lens = torch.tensor([4, 4], dtype=torch.int32, device="cuda") + kv_lens = torch.tensor([1, 6], dtype=torch.int32, device="cuda") + outputs = ( + torch.empty(8, dtype=torch.int32, device="cuda"), + torch.empty(8, dtype=torch.int64, device="cuda"), + torch.empty(8, dtype=torch.int32, device="cuda"), + torch.empty(8, dtype=torch.int32, device="cuda"), + ) + + triton_qsa_generation_token_mapping( + kv_lens=kv_lens, + seq_lens=seq_lens, + request_indices=outputs[0], + logical_positions=outputs[1], + sequence_lengths=outputs[2], + visible_blocks=outputs[3], + num_rows=8, + compress_ratio=4, + ) + + expected = _general_token_mapping(kv_lens, seq_lens, 8, 4) + assert expected[1].min() < -1 + for actual, reference in zip(outputs, expected): + torch.testing.assert_close(actual, reference, rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_qsa_generation_token_mapping_rejects_batches_past_one_tile() -> None: + num_seqs = QSA_GENERATION_MAPPING_MAX_SEQS + 1 + lengths = torch.ones(num_seqs, dtype=torch.int32, device="cuda") + outputs = [ + torch.empty(num_seqs, dtype=dtype, device="cuda") + for dtype in (torch.int32, torch.int64, torch.int32, torch.int32) + ] + with pytest.raises(ValueError, match="at most"): + triton_qsa_generation_token_mapping( + kv_lens=lengths, + seq_lens=lengths, + request_indices=outputs[0], + logical_positions=outputs[1], + sequence_lengths=outputs[2], + visible_blocks=outputs[3], + num_rows=num_seqs, + compress_ratio=4, + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") def test_qsa_spec_dec_remaps_rows_after_the_draft_loop_packs_one_row_per_request() -> None: """The first draft step repacks the batch, so the row mapping must follow it.""" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py index 1cfb3de32ee3..28cfce03ba4c 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py @@ -54,12 +54,21 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res out, gate_buffer, gate = _build(cell.spec, seed=_seed(cell.why)) gate_before = gate_buffer.clone() expected = op.reference(out, gate) + parts = None + if cell.spec.get("amax"): + registered = op(out.clone(), gate) + parts = torch.zeros(1024, dtype=torch.float32, device="cuda") - result = op(out, gate) + result = op(out, gate, parts) 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) + if parts is not None: + # The target's copy is the registered kernel's arithmetic, and the folded maximum + # is exactly the stored output's. + assert torch.equal(result, registered) + assert parts.max().item() == result.abs().max().float().item() def test_metadata_misuse_raises() -> None: diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_nvfp4_pages_to_bf16.py b/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 new file mode 100644 index 000000000000..99c939511024 --- /dev/null +++ b/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 @@ -0,0 +1,161 @@ +# 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.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, +) + +__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) + + +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_paged_index_scores_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_index_scores_verify.py new file mode 100644 index 000000000000..30ed1274f6b3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores_verify.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qsa_paged_index_scores_verify entry: every cell's visible columns bitwise +the row-wise `qsa_paged_index_scores`, on a verify step's layout (each request's rows +consecutive at consecutive positions) over the engine's table (request slots drawn from a +permutation, shuffled pages, -1 past each sequence) and, on the NVFP4 cells, the index +plane as a view of the trunk's padded page; columns past a row's bound stay untouched; a +graph replay with PDL reads the queries after its wait.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_index_scores_verify import ( # noqa: E501 + qsa_paged_index_scores_verify as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "qsa_paged_index_scores_verify requires a CUDA device" + +DEV = "cuda" +IH, ID, TPB, CR, WIDTH = 4, 128, 32, 4, 324 +#: The NVFP4 trunk's index page: 12 index buffers of TPB x ID, then the position page. +BUFFERS, PAD = 12, 192 + + +def _build(spec, seed): + g = torch.Generator(device=DEV).manual_seed(seed) + starts, per = spec["starts"], spec["per"] + rows = len(starts) * per + slots = len(starts) + 3 + slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(starts)] + need = [(s + per - 1) // TPB + 1 for s in starts] + pages = sum(need) + 5 + 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 slot, n in zip(slot_of.tolist(), need): + block_table[slot, :n] = perm[used : used + n] + used += n + keys = torch.randn(pages, TPB, 1, ID, generator=g, device=DEV).to(torch.bfloat16) + if spec.get("nvfp4"): + page = torch.zeros(pages, BUFFERS * TPB * ID + PAD, dtype=torch.bfloat16, device=DEV) + index_cache = page[:, 5 * TPB * ID : 6 * TPB * ID].view(pages, TPB, 1, ID) + index_cache.copy_(keys) + else: + index_cache = keys + scratch = torch.randn(rows, IH + 2, ID, generator=g, device=DEV).to(torch.bfloat16) + positions = torch.tensor( + [s + i for s in starts for i in range(per)], dtype=torch.int64, device=DEV + ) + requests = slot_of.to(DEV).repeat_interleave(per).to(torch.int32) + return dict( + q=scratch[:, :IH], index_cache=index_cache, block_table=block_table, + query_positions=positions, request_indices=requests, tokens_per_block=TPB, + compress_ratio=CR, rows_per_request=per, + ) # fmt: skip + + +def _visible_equal(got, want, positions) -> None: + for r, p in enumerate(positions.tolist()): + n = (p + 1) // CR + assert torch.equal(got[r, :n], want[r, :n]), f"row {r} (position {p})" + + +@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 = op.reference(**args) + got = op(**args) + torch.cuda.synchronize() + assert got.shape == want.shape and got.dtype == torch.float32 + _visible_equal(got, want, args["query_positions"]) + + +def test_reads_the_queries_after_the_wait_under_graph_replay() -> None: + args = _build(dict(starts=(8697, 5000, 1203), per=4), seed=5) + src = args["q"].clone() + work = torch.empty_like(src) + args["q"] = work + + def step(): + work.copy_(src * 1.0) + return op(**args, use_pdl=True) + + 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): + out = step() + for trial in range(3): + src.copy_(torch.randn(src.shape, device=DEV).to(src.dtype)) + graph.replay() + torch.cuda.synchronize() + want = op.reference(**{**args, "q": src}) + _visible_equal(out, want, args["query_positions"]) 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 new file mode 100644 index 000000000000..78dda077cb01 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_decode.py @@ -0,0 +1,233 @@ +# 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 + nvfp4_sparse_gqa_decode, +) +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)) + + +@pytest.mark.parametrize("rows", [72, 200]) +@pytest.mark.parametrize("gated", [False, True], ids=["ungated", "gated"]) +def test_lean_split_kernel_is_bitwise_the_standard_one(rows, gated) -> None: + """A one-split launch builds P V's B fragments from the raw V tile; its output equals the + standard kernel's (forced with `lean=False`) bit for bit.""" + args = _args( + dict(segments=tuple((100 + (499 * g) % 8000, 1) for g in range(rows)), gated=gated), + seed=rows + 5, + ) + call = dict(args) + gate = call.pop("output_gate", None) + positional = [call.pop(name) for name in ( + "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", + )] # fmt: skip + lean = nvfp4_sparse_gqa_decode(*positional, output_gate=gate, lean=True) + standard = nvfp4_sparse_gqa_decode(*positional, output_gate=gate, lean=False) + torch.cuda.synchronize() + assert torch.equal(lean, standard) 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 new file mode 100644 index 000000000000..851a8b21f281 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py @@ -0,0 +1,318 @@ +# 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), (3, None), (4, 1), (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, 128]) +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_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 new file mode 100644 index 000000000000..6239cc4f71af --- /dev/null +++ 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 @@ -0,0 +1,83 @@ +# 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] + rejected = (torch.rand(pages.shape[0], generator=g, device=DEV) < 0.5).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_sparse_gqa_rowtile.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py index fdbdc68c264d..1dc4209b0173 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py @@ -187,7 +187,12 @@ def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811 groups = math.ceil(rows / 10) assert launches[0][1] == (rows,) and launches[1][1] == (groups,) assert launches[2][1] == (min(rows, 256), spec["kv_heads"]) - assert launches[3][1] == (groups, spec["kv_heads"]) + # 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 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()) @@ -202,6 +207,23 @@ 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_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)) @@ -213,3 +235,53 @@ 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_plan_kernels_agree() -> None: + """The CuTe plan (int32 request indices) and the Triton plan (int64) give the same group + plan, slot lists, packed words (below each row-tile group's tile count) and lone rows, with + unmapped pages inside the tables and groups whose rows span requests.""" + spec = _with([(0, 300), (2500, 45), (9000, 64), (700, 3)], slots=8) + args = _build(spec, seed=9900) + table = args["block_table"].clone() + for request in (0, 2): + slot = int(args["request_indices"][sum(n for _, n in spec["segments"][:request])]) + table[slot, 1::3] = -1 + rows, tpb = args["q"].shape[0], spec["tokens_per_block"] + width = table.shape[1] + num_words = math.ceil(width * tpb / 64) + rows_per_tile = 128 // (spec["q_heads"] // spec["kv_heads"]) + groups = math.ceil(rows / rows_per_tile) + bits, row_info = _module._row_bits( + args["selected_tokens"], args["query_positions"], spec["compress_ratio"], tpb, width, + num_words, + ) # fmt: skip + + def plan(request_indices, alloc): + slots = torch.full((groups, num_words * 64), -7, dtype=torch.int32, device=DEV) + packed = alloc((rows, num_words), dtype=torch.int64, device=DEV) + group_plan = torch.full((groups, 2), -7, dtype=torch.int32, device=DEV) + lone = torch.full((rows,), -7, dtype=torch.int32, device=DEV) + count = torch.zeros((1,), dtype=torch.int32, device=DEV) + _module._rowtile_plan( + (groups,), request_indices, row_info, bits, table, slots, packed, group_plan, lone, + count, rows, rows_per_tile=rows_per_tile, tokens_per_block=tpb, + num_cache_pages=args["k_cache"].shape[0], + ) # fmt: skip + return slots, packed, group_plan, lone[: int(count)].sort().values + + assert _module._cute_plan_supported( + args["request_indices"].to(torch.int32), tpb, rows_per_tile, num_words + ) + assert not _module._cute_plan_supported( + args["request_indices"].to(torch.int64), tpb, rows_per_tile, num_words + ) + c_slots, c_packed, c_plan, c_lone = plan(args["request_indices"].to(torch.int32), torch.empty) + t_slots, t_packed, t_plan, t_lone = plan(args["request_indices"].to(torch.int64), torch.zeros) + assert torch.equal(c_plan, t_plan) and torch.equal(c_lone, t_lone) + assert int((c_plan[:, 0] >= 0).sum()) > 0 and c_lone.numel() > 0 + for g, (tiles, listed) in enumerate(c_plan.tolist()): + assert torch.equal(c_slots[g, :listed], t_slots[g, :listed]), g + if tiles >= 0: + members = slice(g * rows_per_tile, min(rows, (g + 1) * rows_per_tile)) + assert torch.equal(c_packed[members, :tiles], t_packed[members, :tiles]), g diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py index c8e63aaa8343..0013c41db1bb 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py @@ -44,6 +44,11 @@ def _build(spec, seed): 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]) 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 index 24b1b58b3521..73b77a7548e3 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue.py @@ -1,10 +1,16 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """GPU test for the qwen4_exp_attn_decode_prologue entry: every cell bitwise the -composition it replaces (`qwen4_exp_qk_norm_rope`, `qsa_paged_kv_store`, -`qsa_paged_index_store`, `qsa_prefill_compress`) on every buffer either writes -- the -projection, the index query, the K/V pool, the index plane, the position cache; a -stale-read check under graph replay with PDL.""" +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 @@ -23,6 +29,8 @@ 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, @@ -30,26 +38,51 @@ ) # 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) - pos = list(spec["pos"]) + layout = _rows(spec) + pos = [p for _, p in layout] rows = len(pos) - pages = SLOTS * WIDTH // 2 + 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) + block_table = torch.full((slots, WIDTH), -1, dtype=torch.int32, device=DEV) used = 0 - for r, p in enumerate(pos): + for req, p in last.items(): n = p // TPB + 1 - block_table[r + 1, :n] = perm[used : used + n] + block_table[req, :n] = perm[used : used + n] used += n - kv = torch.randn(pages, 2, HK, TPB, D, generator=g, device=DEV).to(torch.bfloat16) - index_cache = torch.randn(pages, TPB, 1, ID, generator=g, device=DEV).to(torch.bfloat16) - position_cache = torch.randint( - 0, TABLE, (pages, TPB, 3), generator=g, device=DEV, dtype=torch.int32 - ) - request = torch.arange(1, rows + 1, dtype=torch.int32, device=DEV) + 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 + 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) @@ -58,33 +91,80 @@ def _build(spec, seed): 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], - k_cache=kv[:, 0], v_cache=kv[:, 1], index_cache=index_cache, position_cache=position_cache, - request_indices=request, logical_positions=logical, block_table=block_table, **GEOM, + 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 -_WRITTEN = ("proj", "index_q", "k_cache", "v_cache", "index_cache", "position_cache") +def _bytes(storage, device) -> torch.Tensor: + return torch.empty(0, dtype=torch.uint8, device=device).set_(storage) def _clone(args): - out = dict(args) - kv = args["k_cache"]._base if args["k_cache"]._base is not None else None - if kv is not None: - kv2 = kv.clone() - out["k_cache"], out["v_cache"] = kv2[:, 0], kv2[:, 1] - iq = args["index_q"]._base if args["index_q"]._base is not None else args["index_q"] - iq2 = iq.clone() - out["index_q"] = iq2[:, : args["index_q"].shape[1]] - for name in ("proj", "index_cache", "position_cache"): - out[name] = args[name].clone() - return out + """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]) @@ -92,20 +172,30 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by certifying(op) args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) want = _clone(args) - op.reference(**want) + op.reference(**want[0]) got = _clone(args) - op(**got) + op(**got[0]) torch.cuda.synchronize() - for name in _WRITTEN: - assert torch.equal(got[name], want[name]), name - - -def test_reads_after_the_wait_under_graph_replay() -> None: + _assert_same(got, want) + + +@pytest.mark.parametrize( + "spec", + [ + dict(pos=(8699, 8700)), + dict(pos=(8699, 8700), nvfp4=True), + dict(pos=(8697, 8700), rows_per_request=4, nvfp4=True), + dict(pos=(8699, 8700), nvfp4=True, scores=True), + ], + ids=["bf16_pool", "nvfp4_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.""" - args = _build(dict(pos=(8699, 8700)), seed=23) + 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 = _clone(args) + work, work_copies = _clone(args) def step(): work["proj"].copy_(src * 1.0) @@ -124,8 +214,7 @@ def step(): before = _clone(work) graph.replay() torch.cuda.synchronize() - want = _clone(before) - want["proj"].copy_(src) - op.reference(**want) - for name in _WRITTEN: - assert torch.equal(work[name], want[name]), f"replay {trial}: {name}" + 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_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 new file mode 100644 index 000000000000..9d8456e59f62 --- /dev/null +++ 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 @@ -0,0 +1,101 @@ +# 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 new file mode 100644 index 000000000000..e6ca4c81177c --- /dev/null +++ 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 @@ -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.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), (3, 4), (32, 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/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 new file mode 100644 index 000000000000..2500b3fa4722 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_fp8_gemv.py @@ -0,0 +1,221 @@ +# 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 new file mode 100644 index 000000000000..27bbb8741835 --- /dev/null +++ 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 @@ -0,0 +1,102 @@ +# 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 + AMAX_PARTS, + E4M3_MAX, + amax_partials, + 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 + + +@pytest.mark.parametrize("m,k", [(16384, 6144), (3, 2560)]) +def test_producer_partials_give_the_same_bits(m, k) -> None: + """Partial maxima a producer folded (here: per row block, modulo the buffer) replace the + maximum pass: same scale, same codes, same output.""" + x, _, codes_t, scale = _operands(m, k, 2560, 11) + parts = amax_partials(x.device) + blocks = torch.split(x.abs().float().amax(1), 16) + for i, b in enumerate(blocks): + parts[i % AMAX_PARTS] = torch.maximum(parts[i % AMAX_PARTS], b.max()) + q0, s0 = fp8_quantize_per_tensor(x) + q1, s1 = fp8_quantize_per_tensor(x, parts) + assert torch.equal(s0, s1) and torch.equal(q0.view(torch.uint8), q1.view(torch.uint8)) + assert torch.equal( + fp8_per_tensor_mm(x, codes_t, scale, parts), fp8_per_tensor_mm(x, codes_t, scale) + ) + + +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_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 new file mode 100644 index 000000000000..959822b4b891 --- /dev/null +++ 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 @@ -0,0 +1,119 @@ +# 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 new file mode 100644 index 000000000000..7f5a273112d2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_nvfp4_gemv.py @@ -0,0 +1,326 @@ +# 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_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 new file mode 100644 index 000000000000..8758f80cac54 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_rows_tc.py @@ -0,0 +1,173 @@ +# 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, 3, 4, 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_decode_tc_first.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first.py index 4ec79d79cbe3..bf2d62d45c13 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first.py @@ -68,14 +68,15 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by _compare((low, inject), op.reference(*args)) -def test_guard_refuses_more_than_one_tile_of_rows() -> None: - normed, first, rank, c, wi = _build(dict(rows=129), seed=4) +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) -def test_cuda_graph_capture_and_replay() -> None: - normed, first, rank, c, wi = _build(dict(rows=8), seed=5) +@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): @@ -84,7 +85,7 @@ def test_cuda_graph_capture_and_replay() -> None: graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): g_low, g_inject = op(normed, first, rank, c, wi) - fresh = _build(dict(rows=8), seed=6)[0] + fresh = _build(dict(rows=rows), seed=6)[0] normed.copy_(fresh) graph.replay() torch.cuda.synchronize() 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 new file mode 100644 index 000000000000..1645a331fa44 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_first_rescale_tc.py @@ -0,0 +1,122 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the 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_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 new file mode 100644 index 000000000000..20558035d0b0 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_rstd.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp_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_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 new file mode 100644 index 000000000000..3f4c2f41f174 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_rstd.py @@ -0,0 +1,175 @@ +# 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/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py index c0ead50efd89..0d896c2bd078 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py @@ -8,7 +8,9 @@ fp32 reference. The NVFP4 operand construction (layout, activations, scalars, the band) is the pre-routed entry's test's, loaded by path so the two tests cannot drift apart on what a valid operand is. Beyond the cells: the refusal -of an uncertified (num_experts, top_k) pair. +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 importlib.util @@ -19,7 +21,13 @@ 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, @@ -125,3 +133,51 @@ def test_uncertified_pair_refused(): args, _, data, sf, _, logits, _, _, _, scalars = _operands(spec, 13) with validating(op), pytest.raises(AssertionError, match="certified Renormalize"): _routed(args, data, sf, logits, scalars, spec, top_k=9) + + +@pytest.mark.parametrize("bucket", sorted(DECODE_TACTICS)) +def test_decode_tactic_valid_and_bitwise(bucket): + """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, + ) + # Raises if the runner does not list the tactic as valid at this row count. + assert _pinned_tactic(bucket, num_experts, top_k, hidden, inter) == DECODE_TACTICS[bucket] + args, _, data, sf, _, logits, _, _, _, scalars = _operands(spec, 7000 + bucket) + 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_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 new file mode 100644 index 000000000000..6e47da43f2fe --- /dev/null +++ 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 @@ -0,0 +1,142 @@ +# 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_nvfp4_moe_gemv.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_nvfp4_moe_gemv.py new file mode 100644 index 000000000000..02a27053dafa --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_nvfp4_moe_gemv.py @@ -0,0 +1,117 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for ops/moe/nvfp4_moe_gemv. + +Per cell, on the runner test's exact operands (512 experts in the trtllm-gen layout, +power-of-two block scales): the routing weights are bitwise the trtllm-gen runner's in its +slot order, every slot's FC2 row is bitwise the runner's row for that slot (with power-of-two +scales every accumulation is exact, so accumulation order cannot show), and the finalized +rows sit inside the pre-routed entry's band against the fp32 reference. Beyond the cells: +CUDA-graph replay and the row-count refusal. The operand construction is the routed entry's +test's, loaded by path. +""" + +import importlib.util +import zlib +from pathlib import Path + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_block_scale_moe_renorm_permuted import ( # noqa: E501 + finalize, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_block_scale_moe_renorm_permuted import ( # noqa: E501 + fp4_block_scale_moe_renorm_permuted as runner, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.nvfp4_moe_gemv import ( # noqa: E501 + MAX_ROWS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.nvfp4_moe_gemv import ( # noqa: E501 + nvfp4_moe_gemv_op as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "nvfp4_moe_gemv requires a CUDA device" + + +def _routed_test(): + path = Path(__file__).resolve().parent / "test_modeling_v2_fp4_block_scale_moe_renorm_routed.py" + spec = importlib.util.spec_from_file_location("_fp4_moe_renorm_routed_test", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +R = _routed_test() +H = R.H + + +def _args(args, data, sf, logits, scalars, spec): + 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"], + spec["intermediate"], + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + spec = cell.spec + t, k, hidden = spec["tokens"], spec["top_k"], spec["hidden"] + seed = zlib.crc32(cell.why.encode()) % 2**31 + args, ref, data, sf, x, logits, ids, wts, g2, scalars = R._operands(spec, seed) + call = _args(args, data, sf, logits, scalars, spec) + + rows, topk_w, topk_rows = op(*call) + assert rows.shape == (t * k, hidden) and rows.dtype == torch.bfloat16 + assert topk_w.shape == (t, k) and topk_w.dtype == torch.bfloat16 + assert torch.equal(topk_rows.cpu(), torch.arange(t * k, dtype=torch.int32).view(t, k)) + + r_rows, r_w, r_idx = runner(*call) + assert torch.equal(topk_w, r_w), "routing weights differ from the runner's" + slot_rows = r_rows[r_idx.flatten().long()] + differing = (slot_rows.view(torch.int16) != rows.view(torch.int16)).any(dim=1).sum().item() + assert differing == 0, f"{differing} of {t * k} slot rows differ from the runner's" + + exp = op.reference(logits, x, k, ref["up"], ref["gate"], ref["down"], g2) + H._compare(finalize(rows, topk_w, topk_rows), exp.to(torch.bfloat16)) + + +def test_cuda_graph_capture_and_replay(): + spec = next(c.spec for c in op.CELLS if c.spec["tokens"] == 4) + args, _, data, sf, _, logits, _, _, _, scalars = R._operands(spec, 43) + call = list(_args(args, data, sf, logits, scalars, spec)) + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + for _ in range(3): + op(*call) + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + g_rows, g_w, g_idx = op(*call) + logits.copy_(torch.randn_like(logits.float()).to(logits.dtype)) + graph.replay() + torch.cuda.synchronize() + want = op(*call) + assert torch.equal(g_w, want[1]) + assert torch.equal(finalize(g_rows, g_w, g_idx), finalize(*want)) + + +def test_too_many_rows_refused(): + spec = dict(op.CELLS[0].spec, tokens=MAX_ROWS + 1) + args, _, data, sf, _, logits, _, _, _, scalars = R._operands(spec, 13) + with validating(op), pytest.raises(AssertionError, match="rows"): + op(*_args(args, data, sf, logits, scalars, spec)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py index 21569aeef4e0..69c68fd620c4 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py @@ -57,13 +57,22 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res 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()] + parts = None + if cell.spec.get("amax"): + plain = op(*args) + parts = torch.zeros(1024, dtype=torch.float32, device="cuda") - out = op(*args) + out = op(*args, amax_parts=parts) 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)) + if parts is not None: + # The target's copy keeps the shared kernel's arithmetic; the folded maximum is + # exactly the stored output's. + assert torch.equal(out, plain) + assert parts.max().item() == out.abs().max().float().item() def test_guard_refuses_a_weight_of_the_wrong_length() -> None: diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py index b2729a8e1ee6..de9c02457d60 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py @@ -149,6 +149,37 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by ) +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] diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py index 96f88e6dbdaa..c02361fd0d2e 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py @@ -63,7 +63,7 @@ def _build(spec, seed): 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 @@ -71,7 +71,14 @@ def _build(spec, seed): 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, @@ -106,3 +113,31 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res named = set(args["initial_state_indices"].tolist()) unnamed = [s for s in range(pool.shape[0]) if s not in named] assert torch.equal(pool[unnamed], pool_before[unnamed]), "unnamed slots written" + + +_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_gdn_block.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py index b551a1ec8097..6447b4cf57fd 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py @@ -4,12 +4,12 @@ Drives `DecodeTarget._gdn_block` and `PrefillTarget._gdn_block` at the checkpoint's geometry (hidden 2560, 16 k-heads, 48 v-heads, head dim 128, -fp32 recurrent pool of 461 slots) on a decode and a mixed (context + one-token +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` for decode, -`chunk_gated_delta_rule` over every row of a prefill step) around the same -projections, convolutions and epilogue. +(`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 @@ -45,6 +45,9 @@ 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" @@ -65,11 +68,12 @@ def _slots(n: int) -> list[int]: return ([SLOTS - 1, 0] + rest)[:n] -def _core(seed: int) -> types.SimpleNamespace: - """The attributes `_gdn_block` reads off the core, with layer 0's pools.""" +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 + ssm = (torch.randn(SLOTS, HV, V, K, generator=g, device=DEV) * 0.05).to(state) return types.SimpleNamespace( conv_dim=CONV, key_dim=KD, @@ -110,6 +114,14 @@ def lin(n_out, n_in): 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, ) @@ -142,9 +154,7 @@ def _decode_state(m: int): slot vector is the head of an engine-sized buffer.""" buf = torch.zeros(MAX_BATCH, dtype=torch.int32, device=DEV) buf[:m] = torch.tensor(_slots(m), dtype=torch.int32, device=DEV) - si_all = buf[:m] - qsl_long = torch.arange(m + 1, dtype=torch.int64, device=DEV) - return (si_all, False, m, 1, None, None), qsl_long + return (buf[:m], False, m, 1, None, None) def _prefill_state(ctx_lens: tuple, n_dec: int, has_init: tuple): @@ -210,8 +220,10 @@ def _post_conv(prefill, decode, a, b, L): ) -def _triton_decode_block(core, x, L, si_all, qsl_long): - """The decode block with the Triton single-token recurrence.""" +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( @@ -224,22 +236,22 @@ def _triton_decode_block(core, x, L, si_all, qsl_long): conv_state_indices=si_all, pad_slot_id=M._PAD_SLOT_ID, ) - qv, kv, vv = torch.split(mixed, [KD, KD, VD], 1) + 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=a, + a=torch.reshape(a.contiguous(), [m, 1, HV]), dt_bias=L.dt_bias, softplus_beta=M._SOFTPLUS_BETA, softplus_threshold=M._SOFTPLUS_THRESHOLD, - q=torch.reshape(qv, [1, m, H, K]), - k=torch.reshape(kv, [1, m, H, K]), - v=torch.reshape(vv, [1, m, HV, V]), - b=b, + 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=qsl_long, + cu_seqlens=None, output=None, ) return M.DecodeTarget._gdn_epilogue(_target(M.DecodeTarget, core), out, z, L, m) @@ -252,7 +264,7 @@ def _triton_prefill_block(core, x, L, st): T = x.shape[0] conv, ssm = core._bound.conv[0], core._bound.ssm[0] mixed, z, a, b = torch.split(M.cublas_mm(x, L.in_proj), [CONV, VD, HV, HV], 1) - M.reset_recurrent_state_rows(ssm, si_p, hi_p, conv) + reset_recurrent_state_rows(ssm, si_p, hi_p, conv) conv_args = dict( weight=L.conv_w, bias=None, @@ -337,43 +349,50 @@ def warm_triton_chunk(): torch.cuda.synchronize() -def test_decode_block() -> None: +#: 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), _layer(200 + m), _x(m, 300 + m) + core, L, x = _core(100 + m, STATES[state]), _layer(200 + m), _x(m, 300 + m) core_ref = _clone_core(core) - st, qsl_long = _decode_state(m) + 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], qsl_long) - _check(y, y_ref, core, core_ref, f"decode m={m}") + 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) + original = _core(100 + m, STATES[state]) assert torch.equal(core._bound.ssm[0][unnamed], original._bound.ssm[0][unnamed]), ( "unnamed slots written" ) -def test_prefill_block(warm_triton_chunk) -> None: # noqa: F811 — fixture by name +@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), _layer(500 + tokens), _x(tokens, 600 + tokens) + 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") + _check(y, y_ref, core, core_ref, f"prefill {ctx_lens}+{n_dec}x1, {state} pool") -def test_decode_block_replays_under_cuda_graph() -> None: +@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), _layer(701), _x(m, 702) - st, _ = _decode_state(m) + 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) 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 new file mode 100644 index 000000000000..764cc594b70d --- /dev/null +++ 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 @@ -0,0 +1,363 @@ +# 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: + assert run_cell(cell.why, cell.spec, "b") + + +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", [(1, 4), (4, 4), (5, 3), (32, 4), (8, 8)]) +@pytest.mark.parametrize("pdl", [False, 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_rows_gated_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py index adf92c5b6ba6..bce3e25f9062 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py @@ -3,11 +3,13 @@ """GPU test for the gdn_decode_rows_gated_norm entry. Cells against the composition's fp64 reference (the gated delta rule's recurrence, -then the gated norm entry's) in torch's default band for the bf16 output and the fp32 -pool, the pool compared whole so slots the call did not name are -gated too. Beyond the cells: a stale-read check under graph replay (a producer -rewrites the projection inside the graph right before the op), and the pool-dtype -guard. The `gdn_decode_rows` entry (the recurrence alone) has its own cells below. +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 @@ -15,6 +17,9 @@ 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, ) @@ -43,19 +48,20 @@ def _scattered(n: int, slots: int) -> tuple[int, ...]: 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) + 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 + 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: @@ -86,6 +92,39 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by 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 @@ -117,11 +156,72 @@ def step(): assert torch.equal(after, args[1]), f"replay {trial} pool" +_VARIANT_SPECS = ( + dict(rows=1, state="bf16"), + dict(rows=8, state="bf16", pad=True), + dict(rows=33, state="bf16"), + dict(rows=3), + dict(rows=33, pad=True), + dict(rows=2, pool="coalesced", state="bf16", pad=True), + dict(rows=2, pool="coalesced"), +) + + +@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: - """A bf16 pool (the target sends this entry an fp32 one only) is refused.""" + """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.bfloat16), *args[2:]) + 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_prefill_conv_post.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_prefill_conv_post.py index 8482175c4a3a..9cfca29e1712 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_prefill_conv_post.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_prefill_conv_post.py @@ -95,6 +95,7 @@ def randn(*shape): PAD, L2_NORM_EPS, SOFTPLUS_THRESHOLD, + spec.get("linear_gate", False), ] @@ -136,3 +137,20 @@ def test_guard_refuses_channel_strided_x() -> None: 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 new file mode 100644 index 000000000000..4b256dbf494b --- /dev/null +++ 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 @@ -0,0 +1,156 @@ +# 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, the work items the metadata builds over every decode row (`_replay_work`), PDL +after the fused in_proj GEMV -- against the call it replaces (contiguous `a`/`b`, no work +items, 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) and the history bitwise equal at every step. +""" + +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.mamba.mamba2_metadata import _build_replay_work_items_torch + +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) + mm = SimpleNamespace( + replay_work_items=torch.zeros(n, 4, dtype=torch.int32, device=DEV), + replay_n_writes=torch.zeros(1, dtype=torch.int32, device=DEV), + replay_num_decodes=n, + ) + wraps = 0 + for step in range(HIST // t + 3): + wraps += int((disp.pnat[slots.long()] + t > HIST).sum()) + # what the metadata's prepare builds before the forward + _build_replay_work_items_torch( + slots, disp.pnat, disp.buf, mm.replay_work_items, mm.replay_n_writes, t, HIST + ) + replay = M._replay_work(disp.core, mm, n) + assert replay[0] is not None and not replay[2], "the CuTe kernel takes the work items" + 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" + + +def test_no_work_items_without_the_metadata(): + """A metadata without the build (or built for another batch) hands no work items.""" + side = _Side( + torch.zeros(2, HV, V, K, dtype=BF, device=DEV), + (torch.zeros(1),) * 4, + torch.zeros(2, dtype=torch.int32, device=DEV), + torch.zeros(2, dtype=torch.int32, device=DEV), + 4, + (None, None), + ) + assert M._replay_work(side.core, SimpleNamespace(), 3) == (None, None, False) + mm = SimpleNamespace( + replay_work_items=torch.zeros(8, 4, dtype=torch.int32, device=DEV), + replay_n_writes=torch.zeros(1, dtype=torch.int32, device=DEV), + replay_num_decodes=5, + ) + assert M._replay_work(side.core, mm, 3) == (None, None, False) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_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 new file mode 100644 index 000000000000..1b71871ffcb6 --- /dev/null +++ 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 @@ -0,0 +1,100 @@ +# 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), (4, 4), (8, 2), (2, 8), (5, 3)]) +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/test_modeling_v2_gdn_replay_verify.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py index 0cf38eeb0c1d..b67a01545790 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py @@ -105,6 +105,8 @@ def __init__(self, replay, t, conv0, ssm0, weights, gating, max_batch): 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( diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_logits_widening.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_logits_widening.py index c4aea680b2e9..62dc201db2db 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_logits_widening.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_logits_widening.py @@ -48,3 +48,53 @@ def test_context_logits_and_no_metadata(head) -> None: 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_spec_rollback.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py index c56c098ff397..fc1099ea6537 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py @@ -12,7 +12,8 @@ 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. +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. """ import pytest @@ -175,9 +176,11 @@ 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): +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.""" + 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.""" rollback.discard() qsl = b.qsl rollback.begin( @@ -190,10 +193,11 @@ def _verify_step(b, rollback): logical_positions=b.logical[b.prefix :], block_table=b.block_table, tokens_per_block=TPB, + fused=fused, ) - rollback.stage_ple_pools(b.ngram, b.conv) - window = ngram_token_history(b.ids, b.ngram, b.slots, b.carry, qsl, EOS, NGRAM) - rollback.stage_ple_rows(window, b.x) + rollback.stage_ple_verify(b.ngram, b.conv, b.ids) + ngram_token_history(b.ids, b.ngram, b.slots, b.carry, 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 ) @@ -258,8 +262,8 @@ def _assert_state(b, expected, what): # (ctx_lens, gen_starts, p, gen_accepted). Positions straddle pages and land on -# compression anchors; one request accepts everything. `decode` and `mixed` are the accepted- -# prefix check; `decode_p8` drives the shared-page and capture tests. +# compression anchors; one request accepts everything. All three are the accepted-prefix check; +# `decode_p8` also drives the shared-page and capture tests. CASES = { "decode": ((), (5, 30, 61, 100), 4, (1, 2, 4, 3)), "mixed": ((7, 3), (12, 29, 95), 4, (2, 1, 3)), @@ -273,16 +277,31 @@ def _slots(num): return list(range(3, 3 + 2 * num, 2)) -@pytest.mark.parametrize("case", ["decode", "mixed"]) -def test_commit_matches_the_accepted_prefix(case): +#: 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"] + + +def _rollback(b, mode): + rollback = SpeculativeRollback() + if mode != "per_layer": + rollback.set_index_planes(b.index) + return rollback + + +@pytest.mark.parametrize("mode", MODES) +@pytest.mark.parametrize("case", ["decode", "mixed", "decode_p8"]) +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, SpeculativeRollback()) + _verify_step(b, _rollback(b, mode), mode == "fused") _assert_state(b, _reference(b), case) -def test_commit_over_the_strided_planes_of_a_shared_page(): +@pytest.mark.parametrize("mode", MODES) +def test_commit_over_the_strided_planes_of_a_shared_page(mode): """The handler stages and restores the index and position planes by advanced indexing, so planes that are strided views of one buffer (an NVFP4 KV cache's side planes) roll back like contiguous ones, and the bytes of the other layers' pages in the same buffer stay as @@ -293,7 +312,7 @@ def test_commit_over_the_strided_planes_of_a_shared_page(): 13, ctx_lens, gen_starts, p, _slots(len(ctx_lens) + len(gen_starts)), shared_page=True ) b.set_accepted(gen_accepted) - _verify_step(b, SpeculativeRollback()) + _verify_step(b, _rollback(b, mode), mode == "fused") _assert_state(b, _reference(b), case) @@ -328,21 +347,27 @@ def test_discard_drops_a_stale_staging(): _assert_state(b, (b.ngram0, b.conv0, b.index0, b.position0), "discarded") -def test_commit_is_sync_free(): +@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 = SpeculativeRollback() + 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() torch.cuda.synchronize() torch.cuda.set_sync_debug_mode("error") try: - _verify_step(b, rollback) + _verify_step(b, rollback, mode == "fused") finally: torch.cuda.set_sync_debug_mode(0) _assert_state(b, _reference(b), "sync-free") -def test_padding_duplicates_leave_the_unique_request_exact(): +@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. @@ -350,7 +375,7 @@ def test_padding_duplicates_leave_the_unique_request_exact(): slots = [3, 7, 7, 7] b = _Batch(17, (), gen_starts, p, slots, dup_of={2: 1, 3: 1}) b.set_accepted((2, 3, 3, 3)) - _verify_step(b, SpeculativeRollback()) + _verify_step(b, _rollback(b, mode), mode == "fused") # The unique request's slot is exactly its accepted prefix. pages = b.block_table[0].long() expected = _reference(_single(b, 0)) @@ -396,24 +421,26 @@ 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): +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.""" ctx_lens, gen_starts, p, _ = CASES[case] b = _Batch(18, ctx_lens, gen_starts, p, _slots(len(gen_starts))) - rollback = SpeculativeRollback() + rollback = _rollback(b, mode) + fused = mode == "fused" b.set_accepted((p,) * len(gen_starts)) side = torch.cuda.Stream() side.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(side): - _verify_step(b, rollback) + _verify_step(b, rollback, fused) torch.cuda.current_stream().wait_stream(side) b.reset() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - _verify_step(b, rollback) + _verify_step(b, rollback, fused) for accepted in _REPLAYS[case]: b.reset() diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py index 0e3fe28c3b64..f69e5f4c84e3 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py @@ -173,6 +173,45 @@ def test_bind_step_buffers_uses_manager_layer_views(replay: bool) -> None: assert view.shape == storage.shape[1:] +def test_bf16_replay_pool_switches_the_managers_all_layer_commit_off() -> None: + """A bf16 pool replays on the CuTe kernel, which commits every wrapping checkpoint itself. + + Binding it must leave the manager's all-layer commit off: kept on, the manager would commit + those checkpoints a second time after a step of 16 or more requests, and below that the + metadata would skip building the work items the kernel reads each row's slot from. + """ + state = GDNReplayState(tokens_per_gen_step=2) + state.intermediate_conv = _t(2, SLOTS, 2, G.conv_dim, G.conv_kernel - 1) + state.intermediate_indices = _t(SLOTS, dtype=torch.int32) + 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) + 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(torch.bfloat16), _ssm(torch.bfloat16)] + assert manager.use_gdn_cached_replay_all_layer_commit + + 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), + ) + replay = core._bind_step_buffers(manager, G, metadata, _conv(), _ssm(torch.bfloat16))["replay"] + + assert replay.cute and not replay.all_layer_commit + assert not manager.use_gdn_cached_replay_all_layer_commit + assert not state.all_layer_commit_enabled + + @pytest.mark.parametrize( "plane, match", [ diff --git a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py index adb9257fe190..7872609c44a4 100644 --- a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py +++ b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py @@ -64,8 +64,8 @@ def get_replay_state_update_metadata(self): @pytest.mark.parametrize("num_decodes", [0, 2]) def test_replay_mode_uses_one_metadata_query(monkeypatch, enabled, num_decodes): replay = SimpleNamespace( - prev_num_accepted_tokens=object(), - cache_buf_idx=object(), + prev_num_accepted_tokens=torch.zeros(10, dtype=torch.int32), + cache_buf_idx=torch.zeros(10, dtype=torch.int32), replay_step_width=3, replay_history_size=MIN_REPLAY_HISTORY_SIZE, ) @@ -83,7 +83,10 @@ def use_replay_state_update(self): metadata.replay_work_items = object() metadata.replay_n_writes = object() build = Mock() - monkeypatch.setattr(mamba2_metadata, "_build_replay_work_items_torch", build) + # Contiguous bookkeeping up to the fused maximum takes the one-launch builder, anything else + # the ATen one; either way the replay mode must come from the one metadata query. + for builder in ("_build_replay_work_items_torch", "_build_replay_work_items_triton"): + monkeypatch.setattr(mamba2_metadata, builder, build) Mamba2Metadata._prepare_replay_work_items(metadata, manager, 1 + num_decodes, 1) manager.get_replay_state_update_metadata.assert_called_once_with() assert metadata.replay_num_decodes == (num_decodes if enabled else 0) @@ -368,7 +371,7 @@ def get_replay_state_update_metadata(self): assert actual[0, REPLAY_WORK_PNAT] == 11 assert actual[0, REPLAY_WORK_CACHE_BUF_IDX] == 1 - @pytest.mark.parametrize("num_decodes", [16, 17, 40, 255, 256]) + @pytest.mark.parametrize("num_decodes", [1, 2, 3, 8, 16, 17, 40, 255, 256]) def test_replay_work_items_triton_matches_torch(self, num_decodes): """The two builders must agree everywhere the dispatch may pick either. @@ -460,6 +463,50 @@ def __getitem__(self, grid): assert bool(launches) is expect_fused + @pytest.mark.parametrize( + ("num_decodes", "expect_fused"), + [ + pytest.param(1, True, id="single"), + pytest.param(8, True, id="small"), + pytest.param(256, True, id="fused-max"), + pytest.param(257, False, id="above-fused-max"), + ], + ) + def test_prepare_replay_work_items_without_all_layer_commit_dispatch( + self, monkeypatch, num_decodes, expect_fused + ): + """Without the all-layer commit every batch up to the fused maximum builds in one launch, + with the same work items as the ATen path.""" + num_slots = num_decodes + 7 + prev_num_accepted_tokens = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 21 + cache_buf_idx = torch.arange(num_slots, dtype=torch.int32, device="cuda") % 2 + manager = _GdnReplayCacheManager(prev_num_accepted_tokens, cache_buf_idx) + manager.use_gdn_cached_replay_all_layer_commit = False + metadata = Mamba2Metadata(max_batch_size=num_decodes, chunk_size=8) + state_indices = torch.randperm(num_slots, device="cuda")[:num_decodes].to(torch.int32) + metadata.state_indices.copy_(state_indices) + + launches = [] + original_kernel = mamba2_metadata._prepare_gdn_replay_work_items_kernel + + class _CountingKernel: + def __getitem__(self, grid): + launches.append(grid) + return original_kernel[grid] + + monkeypatch.setattr( + mamba2_metadata, "_prepare_gdn_replay_work_items_kernel", _CountingKernel() + ) + + metadata._prepare_replay_work_items(manager, num_decodes, 0) + + assert bool(launches) is expect_fused + expected_items, expected_n_writes = _torch_reference_work_items( + state_indices, prev_num_accepted_tokens, cache_buf_idx + ) + torch.testing.assert_close(metadata.replay_work_items[:num_decodes], expected_items) + torch.testing.assert_close(metadata.replay_n_writes, expected_n_writes) + def test_prepare_gdn_replay_work_items_cuda_graph_replay(self): num_decodes = 40 num_slots = num_decodes + 7 diff --git a/tests/unittest/_torch/speculative/test_draft_loop_ops.py b/tests/unittest/_torch/speculative/test_draft_loop_ops.py new file mode 100644 index 000000000000..68ee7a2981a7 --- /dev/null +++ b/tests/unittest/_torch/speculative/test_draft_loop_ops.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""`draft_loop_ops` must equal the composed torch ops of the one-model draft loop element for +element, eagerly and replayed from a CUDA graph.""" + +import pytest +import torch + +from tensorrt_llm._torch.speculative.draft_loop_ops import draft_row_select, next_new_tokens +from tensorrt_llm._torch.speculative.mtp import _select_mtp_position_ids + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") + +DEV = "cuda" + + +def _row_select_reference(batch_indices, accepted, hidden, positions, p, num_ctx_tokens): + ids = (batch_indices * p).long() + accepted - 1 + num_ctx_tokens + return ids, hidden[ids], _select_mtp_position_ids(positions, ids) + 1 + + +def _inputs(n, p, width, pos_shape, g, accepted=None, stride_pad=0): + rows = n * p + hidden = torch.randn(rows, width + stride_pad, generator=g).to(DEV, torch.bfloat16) + hidden = hidden[:, :width] + if accepted is None: + accepted = torch.randint(1, p + 1, (n,), generator=g) + accepted = accepted.to(DEV, torch.int32) + batch_indices = torch.arange(n, dtype=torch.int32, device=DEV) + full = torch.randint(0, 1 << 20, pos_shape[:-1] + (rows + 7,), generator=g).to(DEV, torch.int32) + positions = full[..., :rows] + return batch_indices, accepted, hidden, positions + + +@pytest.mark.parametrize("pos_shape", [(3, 1, 0), (3, 0), (0,)]) +@pytest.mark.parametrize("n,p,width", [(1, 4, 10240), (7, 4, 1000), (64, 8, 2048), (256, 2, 64)]) +def test_row_select_matches_composed_ops(n, p, width, pos_shape): + g = torch.Generator().manual_seed(n * 31 + p) + batch_indices, accepted, hidden, positions = _inputs(n, p, width, pos_shape, g, stride_pad=5) + ref = _row_select_reference(batch_indices, accepted, hidden, positions, p, 0) + got = draft_row_select(batch_indices, accepted, hidden, positions, p, 0) + for a, b in zip(got, ref): + assert a.dtype == b.dtype and a.shape == b.shape + assert torch.equal(a, b) + + +def test_row_select_wraps_a_negative_row_like_torch(): + g = torch.Generator().manual_seed(3) + accepted = torch.tensor([0, 2, 4], dtype=torch.int32) + batch_indices, accepted, hidden, positions = _inputs(3, 4, 96, (3, 1, 0), g, accepted) + ref = _row_select_reference(batch_indices, accepted, hidden, positions, 4, 0) + got = draft_row_select(batch_indices, accepted, hidden, positions, 4, 0) + for a, b in zip(got, ref): + assert torch.equal(a, b) + + +def test_row_select_cuda_graph_follows_accepted_counts(): + g = torch.Generator().manual_seed(4) + batch_indices, accepted, hidden, positions = _inputs(16, 4, 512, (3, 1, 0), g) + draft_row_select(batch_indices, accepted, hidden, positions, 4, 0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = draft_row_select(batch_indices, accepted, hidden, positions, 4, 0) + for _ in range(4): + accepted.copy_(torch.randint(1, 5, (16,), generator=g).to(DEV, torch.int32)) + hidden.copy_(torch.randn(hidden.shape, generator=g).to(DEV, torch.bfloat16)) + graph.replay() + torch.cuda.synchronize() + ref = _row_select_reference(batch_indices, accepted, hidden, positions, 4, 0) + for a, b in zip(out, ref): + assert torch.equal(a, b) + + +def _next_new_reference(accepted_tokens, drafts, batch_indices, accepted): + golden = accepted_tokens[batch_indices, accepted - 1].unsqueeze(1) + return torch.concat([golden, drafts], dim=1) + + +@pytest.mark.parametrize("n,d", [(1, 3), (5, 1), (32, 3), (256, 7)]) +def test_next_new_tokens_matches_composed_ops(n, d): + g = torch.Generator().manual_seed(n + d) + accepted_tokens = torch.randint(0, 248320, (n, d + 1), generator=g).to(DEV, torch.int32) + drafts = torch.randint(0, 248320, (n, d), generator=g).to(DEV, torch.int32) + accepted = torch.randint(0, d + 2, (n,), generator=g).to(DEV, torch.int32) + batch_indices = torch.randperm(n, generator=g).to(DEV, torch.int32) + ref = _next_new_reference(accepted_tokens, drafts, batch_indices, accepted) + got = next_new_tokens(accepted_tokens, drafts, batch_indices, accepted) + assert got.dtype == ref.dtype and torch.equal(got, ref) + + +def test_next_new_tokens_cuda_graph(): + g = torch.Generator().manual_seed(9) + n, d = 8, 3 + accepted_tokens = torch.randint(0, 1000, (n, d + 1), generator=g).to(DEV, torch.int32) + drafts = torch.randint(0, 1000, (n, d), generator=g).to(DEV, torch.int32) + accepted = torch.ones(n, dtype=torch.int32, device=DEV) + batch_indices = torch.arange(n, dtype=torch.int32, device=DEV) + next_new_tokens(accepted_tokens, drafts, batch_indices, accepted) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = next_new_tokens(accepted_tokens, drafts, batch_indices, accepted) + for _ in range(3): + accepted.copy_(torch.randint(1, d + 2, (n,), generator=g).to(DEV, torch.int32)) + drafts.copy_(torch.randint(0, 1000, (n, d), generator=g).to(DEV, torch.int32)) + graph.replay() + torch.cuda.synchronize() + assert torch.equal( + out, _next_new_reference(accepted_tokens, drafts, batch_indices, accepted) + ) + + +@pytest.mark.parametrize("n,d", [(1, 3), (32, 3), (100, 7)]) +def test_row_select_moves_the_lengths_to_one_row(n, d): + """`kv_lens[:n] -= d - accepted` and `seq_lens[:n] = 1`, entries past `n` untouched.""" + g = torch.Generator().manual_seed(n * 3 + d) + p = d + 1 + batch_indices, accepted, hidden, positions = _inputs(n, p, 128, (3, 1, 0), g) + kv = torch.randint(100, 9000, (n + 5,), generator=g).to(DEV, torch.int32) + seq = torch.full((n + 5,), p, dtype=torch.int32, device=DEV) + kv_ref, seq_ref = kv.clone(), seq.clone() + kv_ref[:n] -= d - accepted + seq_ref[:n].fill_(1) + got = draft_row_select( + batch_indices, accepted, hidden, positions, p, 0, kv_lens=kv, seq_lens=seq, draft_len=d + ) + ref = _row_select_reference(batch_indices, accepted, hidden, positions, p, 0) + for a, b in zip(got, ref): + assert torch.equal(a, b) + assert torch.equal(kv, kv_ref) + assert torch.equal(seq, seq_ref) diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index 8cebeff8c38e..575c3368903a 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -106,7 +106,7 @@ def run_draft_forward(*args, **kwargs): worker, {"position_ids": torch.tensor([0], dtype=torch.int32)}, metadata, - SimpleNamespace(runtime_draft_len=1), + SimpleNamespace(runtime_draft_len=1, is_cuda_graph=False), draft_model=object(), draft_kv_cache_manager=draft_manager, num_contexts=0, diff --git a/tests/unittest/_torch/speculative/test_fused_strict_accept.py b/tests/unittest/_torch/speculative/test_fused_strict_accept.py new file mode 100644 index 000000000000..721a81059a13 --- /dev/null +++ b/tests/unittest/_torch/speculative/test_fused_strict_accept.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The strict acceptance of a CUDA-graph verify step (`fused_accept.strict_accept`) must +return what the composed ops of an eager step return: the same accepted tokens and counts, +including the synthetic acceptance override and its counter, replay after replay.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.speculative import fused_accept +from tensorrt_llm._torch.speculative.interface import SpecWorkerBase + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") + +VOCAB = 248320 + + +class _Worker(SpecWorkerBase): + @property + def max_draft_len(self) -> int: + return 8 + + def _forward_impl(self, *args, **kwargs): + raise NotImplementedError + + +def _worker(force: float) -> _Worker: + worker = _Worker() + worker.force_num_accepted_tokens = force + return worker + + +def _meta(is_cuda_graph: bool) -> SimpleNamespace: + return SimpleNamespace( + is_cuda_graph=is_cuda_graph, + is_all_greedy_sample=True, + is_ragged_verify=False, + enable_penalty=False, + ) + + +def _fill(logits, draft, gen, d, generator): + """Random logits with one planted maximum a row; draft tokens that repeat a random number + of the target's leading tokens, then differ (or repeat all of them).""" + rows = gen * (d + 1) + logits.copy_(torch.randn(rows, VOCAB, generator=generator).to(logits.device)) + target = torch.randint(0, VOCAB, (rows,), generator=generator) + logits[torch.arange(rows, device=logits.device), target.to(logits.device)] = 100.0 + target = target.reshape(gen, d + 1) + match = torch.randint(0, d + 1, (gen,), generator=generator) + proposal = torch.where( + torch.arange(d)[None, :] < match[:, None], target[:, :d], (target[:, :d] + 1) % VOCAB + ) + draft.copy_(proposal.to(draft.device, torch.int32)) + + +def _composed(worker, logits, draft, gen): + return worker._sample_and_accept_draft_tokens_base(logits, draft, 0, gen, _meta(False)) + + +def _fused(worker, logits, draft, gen): + meta = _meta(True) + assert worker._can_use_fused_strict_accept(logits, 0, gen, draft.shape[1], meta) + return worker._sample_and_accept_draft_tokens_base(logits, draft, 0, gen, meta) + + +@pytest.mark.parametrize("gen", [1, 4, 33]) +@pytest.mark.parametrize("d", [1, 3, 7]) +def test_verified_counts_match(gen, d): + g = torch.Generator().manual_seed(gen * 10 + d) + logits = torch.empty(gen * (d + 1), VOCAB, device="cuda") + draft = torch.empty(gen, d, dtype=torch.int32, device="cuda") + worker = _worker(0.0) + for _ in range(4): + _fill(logits, draft, gen, d, g) + ref = _composed(worker, logits, draft, gen) + got = _fused(worker, logits, draft, gen) + assert got[0].dtype == ref[0].dtype == torch.int32 + assert got[1].dtype == ref[1].dtype == torch.int32 + assert torch.equal(got[0], ref[0]) + assert torch.equal(got[1], ref[1]) + + +@pytest.mark.parametrize("force", [0.0, 2.3, 2.0, 1.7, 5.0, 3.5]) +@pytest.mark.parametrize("gen", [1, 4, 64]) +def test_cuda_graph_replays_match_eager_steps(force, gen): + """The fused step captured once and replayed against the composed eager step of a second + worker: both counters start at zero and advance once a step.""" + d = 3 + g = torch.Generator().manual_seed(int(force * 10) + gen) + logits = torch.empty(gen * (d + 1), VOCAB, device="cuda") + draft = torch.empty(gen, d, dtype=torch.int32, device="cuda") + eager, graphed = _worker(force), _worker(force) + graphed._ensure_force_accept_rng_state(logits.device) + _fill(logits, draft, gen, d, g) + # The graph warmup step: synthetic acceptance off, kernels compiled and loaded. + _fused(graphed, logits, draft, gen) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = _fused(graphed, logits, draft, gen) + for _ in range(12): + _fill(logits, draft, gen, d, g) + ref = _composed(eager, logits, draft, gen) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out[0], ref[0]) + assert torch.equal(out[1], ref[1]) + if eager._force_accept_rng_counter is not None: + assert torch.equal(graphed._force_accept_rng_counter, eager._force_accept_rng_counter) + + +def test_admission(): + meta = _meta(True) + logits = torch.empty(8, 16, device="cuda") + assert SpecWorkerBase._can_use_fused_strict_accept(logits, 0, 2, 3, meta) + # Context rows, eager steps, non-greedy and ragged batches keep the composed ops. + assert not SpecWorkerBase._can_use_fused_strict_accept(logits, 1, 2, 3, meta) + assert not SpecWorkerBase._can_use_fused_strict_accept(logits, 0, 2, 3, _meta(False)) + for key in ("is_all_greedy_sample",): + other = _meta(True) + setattr(other, key, False) + assert not SpecWorkerBase._can_use_fused_strict_accept(logits, 0, 2, 3, other) + other = _meta(True) + other.is_ragged_verify = True + assert not SpecWorkerBase._can_use_fused_strict_accept(logits, 0, 2, 3, other) + big = torch.empty((fused_accept.MAX_REQUESTS + 1) * 4, 16, device="cuda") + assert not SpecWorkerBase._can_use_fused_strict_accept( + big, 0, fused_accept.MAX_REQUESTS + 1, 3, meta + ) diff --git a/tests/unittest/_torch/speculative/test_greedy_argmax.py b/tests/unittest/_torch/speculative/test_greedy_argmax.py new file mode 100644 index 000000000000..18f2dd41fc7f --- /dev/null +++ b/tests/unittest/_torch/speculative/test_greedy_argmax.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""`greedy_argmax` must equal `torch.argmax(logits, dim=-1)` element for element.""" + +import pytest +import torch + +from tensorrt_llm._torch.speculative.greedy_argmax import ( + MIN_NUMEL, + SMALL_ROWS, + eligible, + greedy_argmax, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") + +VOCAB = 248320 + + +def _check(x: torch.Tensor, out_dtype: torch.dtype = torch.int64, small_rows: bool = False) -> None: + assert eligible(x, small_rows) + got = greedy_argmax(x, out_dtype, small_rows=small_rows) + ref = torch.argmax(x, dim=-1).to(out_dtype) + assert got.dtype == out_dtype + assert torch.equal(got, ref) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("rows", [17, 128, 256, 512, 1024]) +def test_random(rows, dtype): + torch.manual_seed(rows) + _check(torch.randn(rows, VOCAB, device="cuda").to(dtype)) + + +@pytest.mark.parametrize("cols", [VOCAB + 1, 4097, 151937, 1 << 20]) +def test_odd_vocab(cols): + rows = max(1, MIN_NUMEL // cols + 1) + torch.manual_seed(cols) + x = torch.randn(rows, cols, device="cuda") + _check(x) + _check(x, torch.int32) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("rows", [1, 2, 4, 7, 16]) +def test_small_rows(rows, dtype): + """Graph-replayed steps route 1 to `SMALL_ROWS` rows to the kernels below `MIN_NUMEL`.""" + torch.manual_seed(rows) + x = torch.randn(rows, VOCAB, device="cuda").to(dtype) + assert not eligible(x) + _check(x, small_rows=True) + _check(x, torch.int32, small_rows=True) + + +def test_small_rows_bound(): + x = torch.randn(SMALL_ROWS + 1, 1000, device="cuda") + assert not eligible(x, small_rows=True) + assert torch.equal(greedy_argmax(x, small_rows=True), torch.argmax(x, dim=-1)) + + +@pytest.mark.parametrize("cols", [1, 100, 2049, 151937, VOCAB + 1]) +def test_small_rows_odd_vocab(cols): + torch.manual_seed(cols) + x = torch.randn(4, cols, device="cuda") + _check(x, small_rows=True) + _check(x, torch.int32, small_rows=True) + + +@pytest.mark.parametrize("rows", [1, 4, 16, 64, 256, 1024]) +def test_ties_and_special_values(rows): + torch.manual_seed(rows) + # bf16-valued fp32 logits, as the widened head output: ties are frequent. + x = torch.randn(rows, VOCAB, device="cuda").bfloat16().float() + r = torch.arange(rows, device="cuda") + # A repeated row maximum at several positions, across chunk boundaries. + for c in (7, 4095, 4096, 31040, 31041, 124159, VOCAB - 1): + x[r % 4 == 0, c] = 100.0 + # +0.0 / -0.0 ties as the maximum. + x[r % 4 == 1] = -1.0 + x[r % 4 == 1, 9000] = -0.0 + x[r % 4 == 1, 100] = 0.0 + x[r % 4 == 1, 200000] = -0.0 + # All -inf; -inf with one finite value late in the row. + x[r % 8 == 2] = float("-inf") + x[r % 8 == 6] = float("-inf") + x[r % 8 == 6, VOCAB - 3] = -1e30 + # NaN: the first NaN wins over +inf and larger indices. + x[r % 4 == 3, 50] = float("inf") + x[r % 4 == 3, 70000] = float("nan") + x[r % 8 == 3, 3] = float("nan") + x[r % 4 == 3, VOCAB - 1] = float("nan") + small = rows <= SMALL_ROWS + _check(x, small_rows=small) + _check(x.bfloat16(), small_rows=small) + # One row entirely NaN. + x[0] = float("nan") + _check(x, small_rows=small) + + +def test_small_input_falls_back(): + x = torch.randn(2, 1000, device="cuda") + assert not eligible(x) + assert torch.equal(greedy_argmax(x), torch.argmax(x, dim=-1)) + + +def test_strided_rows(): + torch.manual_seed(1) + base = torch.randn(256, VOCAB + 64, device="cuda") + _check(base[:, :VOCAB]) + + +@pytest.mark.parametrize("rows", [1, 4, 256]) +def test_cuda_graph(rows): + torch.manual_seed(2) + x = torch.randn(rows, VOCAB, device="cuda") + out = greedy_argmax(x, small_rows=True) + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out = greedy_argmax(x, small_rows=True) + for _ in range(3): + x.copy_(torch.randn_like(x)) + g.replay() + torch.cuda.synchronize() + assert torch.equal(out, torch.argmax(x, dim=-1))