From adee02915937dd266a549e409e358470369ce607 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Mon, 9 Nov 2026 19:13:22 -0800 Subject: [PATCH 01/19] [None][fix] DSpark FP8 producers: round UE8M0 scales up explicitly Use cvt.rp.satfinite.ue8m0x2.f32 for the scale used by rcp_approx and the stored packed scale, instead of the generic Float8E8M0FNU conversion. Signed-off-by: Zhenhuan Chen --- .../kernels/blackwell/dspark_rmsnorm_rope.py | 38 ++++++++++++++++--- .../dspark_rmsnorm_rope_custom_op.py | 8 ++-- 2 files changed, 37 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark_rmsnorm_rope.py b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark_rmsnorm_rope.py index 74a872a60439..fe74fb02a69d 100644 --- a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark_rmsnorm_rope.py +++ b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark_rmsnorm_rope.py @@ -7,6 +7,8 @@ import cutlass import cutlass.cute as cute +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op try: from cuda.bindings import driver as cuda @@ -14,6 +16,30 @@ from cuda import cuda +@dsl_user_op +def _cvt_rp_ue8m0(value, *, loc=None, ip=None) -> cutlass.Uint8: + """Round a positive FP32 value up to its raw UE8M0 exponent byte.""" + packed = llvm.inline_asm( + T.i16(), + [cutlass.Float32(value).ir_value(loc=loc, ip=ip)], + "cvt.rp.satfinite.ue8m0x2.f32 $0, 0f00000000, $1;", + "=h,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Uint8(cutlass.Int16(packed) & cutlass.Int16(0xFF)) + + +@dsl_user_op +def _ue8m0_to_f32(exponent, *, loc=None, ip=None) -> cutlass.Float32: + """Widen a nonzero UE8M0 exponent byte to its FP32 power of two.""" + bits = cutlass.Uint32(exponent) << cutlass.Uint32(23) + return cutlass.Float32(llvm.bitcast(T.f32(), bits.ir_value(loc=loc, ip=ip))) + + class DSparkRMSNormRoPEKernel: """Apply optional RMSNorm and last-dimension RoPE with one warp per row.""" @@ -409,9 +435,10 @@ def kernel( ) block_absmax = cute.arch.warp_reduction_max(local_absmax) + # The 1e-10 floor keeps the exponent byte nonzero (normal FP32 scale). block_absmax = cute.arch.fmax(block_absmax, cutlass.Float32(1e-10)) - scale = (block_absmax * cutlass.Float32(1.0 / 448.0)).to(cutlass.Float8E8M0FNU) - quant_scale = cute.arch.rcp_approx(cutlass.Float32(scale)) + scale = _cvt_rp_ue8m0(block_absmax * cutlass.Float32(1.0 / 448.0)) + quant_scale = cute.arch.rcp_approx(_ue8m0_to_f32(scale)) for item in cutlass.range_constexpr(self.values_per_thread): dim = block_base + tidx * self.values_per_thread + item @@ -476,7 +503,7 @@ def kernel( lane = tidx % cute.arch.WARP_SIZE row = row_block * self.num_warps + warp total_rows = first.shape[0] + second.shape[0] - scale = cutlass.Float8E8M0FNU(0.0) + scale = cutlass.Uint8(0) if cutlass.dynamic_expr(row < total_rows): values = cute.make_rmem_tensor((self.values_per_thread,), cutlass.BFloat16) @@ -505,9 +532,10 @@ def kernel( ) block_absmax = cute.arch.warp_reduction_max(local_absmax) + # The 1e-10 floor keeps the exponent byte nonzero (normal FP32 scale). block_absmax = cute.arch.fmax(block_absmax, cutlass.Float32(1e-10)) - scale = (block_absmax * cutlass.Float32(1.0 / 448.0)).to(cutlass.Float8E8M0FNU) - quant_scale = cute.arch.rcp_approx(cutlass.Float32(scale)) + scale = _cvt_rp_ue8m0(block_absmax * cutlass.Float32(1.0 / 448.0)) + quant_scale = cute.arch.rcp_approx(_ue8m0_to_f32(scale)) for item in cutlass.range_constexpr(self.values_per_thread): dim = block_base + lane * self.values_per_thread + item output[row, dim] = (cutlass.Float32(values[item]) * quant_scale).to( diff --git a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py index d244ba0e719e..2b384dc798a1 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py @@ -690,7 +690,7 @@ def _compile_dspark_rmsnorm_fp8_quant(hidden_dim: int, eps: float): cutlass.Float8E4M3FN, (rows, hidden_dim), stride_order=(1, 0) ) scale_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Float8E8M0FNU, (scale_elems,), stride_order=(0,) + cutlass.Uint8, (scale_elems,), stride_order=(0,) ) stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) kernel = DSparkRMSNormFP8QuantKernel(hidden_dim, eps) @@ -730,7 +730,7 @@ def _compile_dspark_concat_fp8_quant(hidden_dim: int): cutlass.Float8E4M3FN, (output_rows, hidden_dim), stride_order=(1, 0) ) scale_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Float8E8M0FNU, (scale_elems,), stride_order=(0,) + cutlass.Uint8, (scale_elems,), stride_order=(0,) ) stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) kernel = DSparkConcatFP8QuantKernel(hidden_dim) @@ -765,7 +765,7 @@ def cute_dsl_dspark_rmsnorm_fp8_quant( x_flat, weight, output, - output_scale.view(torch.float8_e8m0fnu), + output_scale, ) return output, output_scale @@ -805,7 +805,7 @@ def cute_dsl_dspark_concat_fp8_quant( first, second, output, - output_scale.view(torch.float8_e8m0fnu), + output_scale, ) return output, output_scale From a4321c14f311995a486a3529a7b7003ec4f487f1 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Mon, 9 Nov 2026 19:15:19 -0800 Subject: [PATCH 02/19] [None][fix] DSpark: one BF16 attention binder; V4.1 rejects fp8 projections Bind BF16 projections with the checkpoint FP8 block size (32 for V4.1) in one helper, and fail on TRTLLM_DSPARK_LINEAR_BACKEND=fp8 for drafters without FP8. Signed-off-by: Zhenhuan Chen --- docs/source/features/speculative-decoding.md | 4 +- tensorrt_llm/_torch/models/modeling_dspark.py | 76 ++++++------------- .../_torch/models/modeling_dspark_v41.py | 2 + .../hw_agnostic/test_dspark_attention.py | 50 ++++++++++++ .../hw_agnostic/test_dspark_v41.py | 10 +++ 5 files changed, 89 insertions(+), 53 deletions(-) diff --git a/docs/source/features/speculative-decoding.md b/docs/source/features/speculative-decoding.md index 94213b82efb5..5cf562337e9f 100644 --- a/docs/source/features/speculative-decoding.md +++ b/docs/source/features/speculative-decoding.md @@ -360,7 +360,9 @@ use FP8 on Rubin SM107 with TP=1 or attention DP, and BF16 otherwise. Set bf16 to use dequantized BF16 draft weights. An explicit fp8 selection requires TP=1 or attention DP; unsupported tensor-parallel mappings are rejected during initialization. Non-Rubin FP8 is opt-in and requires workload -accuracy and acceptance-length validation. +accuracy and acceptance-length validation. The variable applies only to +DeepSeek-V4 drafters: DeepSeek-V4.1 DSpark always uses BF16 projections and +rejects fp8. ### Packed FP4 MLA wave planning diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 701f3bfc6f76..3aa926a8ca48 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -1248,6 +1248,8 @@ class DSv4DSparkDraftModel(nn.Module): block_cls = DSv4DSparkBlock uses_lagged_pre = False checkpoint_fp8_block_size = 128 + # Whether post_load_weights(weights) can bind the loaded FP8 attention Linears. + supports_fp8_attn_projections = True confidence_uses_normalized_hidden = True def __init__( @@ -1364,7 +1366,9 @@ def __init__( self._freqs_cap = _runtime_position_cap(model_config, config, self.block_size + 2) self._freqs_table_cache: Dict = {} self._context_wkv_fp8 = None - self._use_bf16_linear = _use_dspark_bf16_linear(model_config.mapping) + self._use_bf16_linear = _use_dspark_bf16_linear( + model_config.mapping, fp8_supported=self.supports_fp8_attn_projections + ) def post_load_weights(self, weights: Optional[Dict] = None) -> None: """Transform draft linears and bind the selected attention backend.""" @@ -1401,51 +1405,9 @@ def _block_dequant( scale = scale.repeat_interleave(block, 1)[:, :input_features] return (weight * scale).to(torch.bfloat16) - def _cache_attn_weights(self, src: Dict) -> None: - """Populate each stage's ``_dspark_attn`` from a dict of raw ``mtp.{s}.attn.*`` - tensors (source-agnostic core shared by the two public entry points). - - The captured-context attention runs the validated ``_dspark_attention_forward_legacy`` - free function on reference-layout bf16 weights dequantized here. Sourcing the - separate ``wq_a``/``wkv`` (plain 128×128 block scale) sidesteps the TRT-LLM - ``MLA`` module's fused + interleaved fp8 storage (``kv_a_proj_with_mqa`` fuses - ``q_a``+``kv`` and stores the scale interleaved). This mirrors the - golden-validated dequant exactly. - """ - self._use_bf16_linear = True - self._context_wkv_fp8 = None - for s, stage in enumerate(self.mtp_layers): - pref = f"mtp.{s}.attn." - dev = stage.input_layernorm.weight.device - - def deq(name: str, fp8: bool) -> torch.Tensor: - w = src[f"{pref}{name}.weight"].to(dev) - if fp8: - return self._block_dequant( - w, - src[f"{pref}{name}.scale"].to(dev), - block=self.checkpoint_fp8_block_size, - ) - return w.to(torch.bfloat16) - - stage._dspark_attn = dict( - wq_a=deq("wq_a", True), - q_norm_w=src[f"{pref}q_norm.weight"].to(dev).to(torch.bfloat16), - wq_b=deq("wq_b", True), - wkv=deq("wkv", True), - kv_norm_w=src[f"{pref}kv_norm.weight"].to(dev).to(torch.bfloat16), - # wo_a IS fp8+scale in the checkpoint (verified); always dequant. - wo_a=deq("wo_a", True), - wo_b=deq("wo_b", True), - attn_sink=src[f"{pref}attn_sink"].to(dev).float(), - ) - def cache_attn_weights_from_checkpoint(self, ckpt_dir: str, weight_map: Dict[str, str]) -> None: """Populate ``_dspark_attn`` by reading the ``mtp.{s}.attn.*`` tensors from the - checkpoint shards on disk, then dequantizing via :meth:`_cache_attn_weights`. - - TODO(step 3): source these from the loaded ``MLA`` modules instead, once the - fused/interleaved fp8 scale layout is decoded, to drop the checkpoint I/O. + checkpoint shards on disk, then dequantizing via :meth:`_bind_bf16_attn_projections`. """ from safetensors import safe_open @@ -1459,14 +1421,14 @@ def cache_attn_weights_from_checkpoint(self, ckpt_dir: str, weight_map: Dict[str with safe_open(os.path.join(ckpt_dir, shard), framework="pt", device="cpu") as f: for k in ks: raw[k] = f.get_tensor(k) - self._cache_attn_weights(raw) + self._bind_bf16_attn_projections(raw) def cache_attn_weights_from_state_dict(self, weights: Dict) -> None: """Populate ``_dspark_attn`` from an already-loaded in-memory ``weights`` dict - (no extra disk I/O); used on the one-engine load path - (``DSv4DSparkForCausalLM.load_weights``). Delegates to :meth:`_cache_attn_weights`. + (no extra disk I/O). DeepSeek-V4.1 DSpark loads its BF16 attention projections + this way; delegates to :meth:`_bind_bf16_attn_projections`. """ - self._cache_attn_weights(weights) + self._bind_bf16_attn_projections(weights) def _dspark_freqs_table(self, device: torch.device) -> torch.Tensor: """Return the plain-RoPE table cached for ``(device, cap)``. @@ -2156,7 +2118,12 @@ def forward_head( ) def _bind_bf16_attn_projections(self, weights: Dict) -> None: - """Bind all five captured-context attention projections as BF16 tensors.""" + """Bind all five captured-context attention projections as BF16 tensors. + + Dequantizes the raw ``mtp.{s}.attn.*`` checkpoint tensors with the + checkpoint's FP8 block size and disables every FP8 projection path. + """ + self._use_bf16_linear = True for stage_id, stage in enumerate(self.mtp_layers): prefix = f"mtp.{stage_id}.attn." device = stage.input_layernorm.weight.device @@ -2164,7 +2131,7 @@ def _bind_bf16_attn_projections(self, weights: Dict) -> None: def dequant(name: str) -> torch.Tensor: weight = weights[f"{prefix}{name}.weight"].to(device) scale = weights[f"{prefix}{name}.scale"].to(device) - return self._block_dequant(weight, scale) + return self._block_dequant(weight, scale, block=self.checkpoint_fp8_block_size) stage._dspark_context_wkv_slice = None stage._dspark_attn = dict( @@ -4496,15 +4463,20 @@ def _fused_fp8_output_projection_from_quantized( return projected.view(batch, block, projected.shape[-1]) -def _use_dspark_bf16_linear(mapping) -> bool: +def _use_dspark_bf16_linear(mapping, fp8_supported: bool = True) -> bool: """Keep BF16 outside Rubin unless FP8 is explicitly requested.""" requested = os.environ.get(_DSPARK_LINEAR_BACKEND_ENV) sharded = mapping.tp_size > 1 and not mapping.enable_attention_dp if requested is None: - return sharded or get_sm_version() != 107 + return not fp8_supported or sharded or get_sm_version() != 107 backend = requested.lower() if backend not in ("fp8", "bf16"): raise ValueError(f"{_DSPARK_LINEAR_BACKEND_ENV} must be 'fp8' or 'bf16', got {backend!r}") + if backend == "fp8" and not fp8_supported: + raise ValueError( + f"{_DSPARK_LINEAR_BACKEND_ENV}=fp8 is not supported by this DSpark drafter; " + "it binds BF16 attention projections" + ) if backend == "fp8" and sharded: raise ValueError(f"{_DSPARK_LINEAR_BACKEND_ENV}=fp8 requires TP=1 or attention DP") return backend == "bf16" diff --git a/tensorrt_llm/_torch/models/modeling_dspark_v41.py b/tensorrt_llm/_torch/models/modeling_dspark_v41.py index 4b030bb6e991..5ab86b3aa403 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark_v41.py +++ b/tensorrt_llm/_torch/models/modeling_dspark_v41.py @@ -144,6 +144,8 @@ class DSv41DSparkDraftModel(DSv4DSparkDraftModel): block_cls = DSv41DSparkBlock uses_lagged_pre = True checkpoint_fp8_block_size = 32 + # V4.1 loads attention through cache_attn_weights_from_state_dict (BF16). + supports_fp8_attn_projections = False # Reference DSparkBlock.forward_head passes collapsed x to confidence, # while only the shared vocabulary projection consumes norm(x). confidence_uses_normalized_hidden = False diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py index f5247270b44e..71a84df8303e 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py @@ -561,6 +561,22 @@ def test_use_dspark_bf16_linear_rejects_sharded_fp8(monkeypatch): modeling_dspark._use_dspark_bf16_linear(mapping) +@pytest.mark.parametrize("value,expected", [(None, True), ("bf16", True), ("fp8", None)]) +def test_use_dspark_bf16_linear_without_fp8_support(monkeypatch, value, expected): + monkeypatch.setattr(modeling_dspark, "get_sm_version", lambda: 107) + if value is None: + monkeypatch.delenv("TRTLLM_DSPARK_LINEAR_BACKEND", raising=False) + else: + monkeypatch.setenv("TRTLLM_DSPARK_LINEAR_BACKEND", value) + for tp_size, enable_attention_dp in ((1, False), (2, True), (2, False)): + mapping = types.SimpleNamespace(tp_size=tp_size, enable_attention_dp=enable_attention_dp) + if expected is None: + with pytest.raises(ValueError, match="not supported by this DSpark drafter"): + modeling_dspark._use_dspark_bf16_linear(mapping, fp8_supported=False) + else: + assert modeling_dspark._use_dspark_bf16_linear(mapping, fp8_supported=False) + + def test_use_dspark_bf16_linear_rejects_invalid_backend(monkeypatch): monkeypatch.setenv("TRTLLM_DSPARK_LINEAR_BACKEND", "auto") mapping = types.SimpleNamespace(tp_size=1, enable_attention_dp=False) @@ -1048,6 +1064,40 @@ def test_bind_bf16_attn_projections_disables_every_fp8_projection(): assert model._context_wkv_fp8 is None +@pytest.mark.parametrize("block", [128, 32]) +def test_bind_bf16_attn_projections_uses_checkpoint_block_size(block: int) -> None: + model = DSv4DSparkDraftModel.__new__(DSv4DSparkDraftModel) + torch.nn.Module.__init__(model) + model.checkpoint_fp8_block_size = block + model._use_bf16_linear = False + stage = types.SimpleNamespace(input_layernorm=types.SimpleNamespace(weight=torch.ones(1))) + model.mtp_layers = [stage] + prefix = "mtp.0.attn." + generator = torch.Generator().manual_seed(block) + weights = { + f"{prefix}q_norm.weight": torch.ones(2), + f"{prefix}kv_norm.weight": torch.ones(3), + f"{prefix}attn_sink": torch.zeros(6), + } + for name in ("wq_a", "wq_b", "wkv", "wo_a", "wo_b"): + weights[f"{prefix}{name}.weight"] = torch.randn(128, 128, generator=generator).to( + torch.float8_e4m3fn + ) + weights[f"{prefix}{name}.scale"] = torch.rand( + 128 // block, 128 // block, generator=generator + ) + + model._bind_bf16_attn_projections(weights) + + assert model._use_bf16_linear is True + assert stage._dspark_context_wkv_slice is None + for name in ("wq_a", "wq_b", "wkv", "wo_a", "wo_b"): + weight = weights[f"{prefix}{name}.weight"].float() + scale = weights[f"{prefix}{name}.scale"] + expected = weight * scale.repeat_interleave(block, 0).repeat_interleave(block, 1) + torch.testing.assert_close(stage._dspark_attn[name], expected.to(torch.bfloat16)) + + def test_context_wkv_fp8_quantizes_once_per_window_write(monkeypatch): num_stages, hidden, head_dim, window = 3, 4, 2, 8 model = DSv4DSparkDraftModel.__new__(DSv4DSparkDraftModel) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py index 1a8b77066a7b..f6c0fdb602b3 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py @@ -253,6 +253,16 @@ def test_v41_block32_dequant_uses_each_tile_scale(): ) +def test_v41_draft_rejects_fp8_attention_projections(monkeypatch): + monkeypatch.setenv("TRTLLM_DSPARK_LINEAR_BACKEND", "fp8") + assert DSv41DSparkDraftModel.supports_fp8_attn_projections is False + mapping = SimpleNamespace(tp_size=1, enable_attention_dp=False) + with pytest.raises(ValueError, match="not supported by this DSpark drafter"): + modeling_dspark._use_dspark_bf16_linear( + mapping, fp8_supported=DSv41DSparkDraftModel.supports_fp8_attn_projections + ) + + def test_v41_draft_uses_smaller_expert_bank_without_changing_target(monkeypatch): captured = {} From 941e201be59b37fa19a63ddc92424cd065aa2b6f Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Mon, 9 Nov 2026 19:16:24 -0800 Subject: [PATCH 03/19] [None][chore] DSpark: check fused WKV width, document ring mapping, merge imports Reject a q_a+wkv projection whose width is not q_lora_rank + head_dim, since every context-WKV consumer slices the trailing head_dim columns. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/models/modeling_dspark.py | 33 ++++++++++++------- .../hw_agnostic/test_dspark_attention.py | 5 +++ 2 files changed, 26 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 3aa926a8ca48..f5e60a0f61bf 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -137,25 +137,19 @@ class or the edge would become a cycle. _DSPARK_LINEAR_BACKEND_ENV = "TRTLLM_DSPARK_LINEAR_BACKEND" -if IS_CUTLASS_DSL_AVAILABLE: - from ..custom_ops import dspark_attention_custom_op # noqa: F401 - from ..custom_ops.dspark_rmsnorm_rope_custom_op import ( - cute_dsl_dspark_concat_fp8_quant, - cute_dsl_dspark_rmsnorm_fp8_quant, - cute_dsl_dspark_rmsnorm_rope, - cute_dsl_dspark_rmsnorm_rope_draft_block, - is_fused_dspark_rmsnorm_rope_supported, - ) - - if IS_CUTLASS_DSL_AVAILABLE: from ..custom_ops.dspark_attention_custom_op import ( fused_dsv4_dspark_attention, is_fused_dsv4_dspark_attention_supported, ) from ..custom_ops.dspark_rmsnorm_rope_custom_op import ( + cute_dsl_dspark_concat_fp8_quant, + cute_dsl_dspark_rmsnorm_fp8_quant, + cute_dsl_dspark_rmsnorm_rope, cute_dsl_dspark_rmsnorm_rope_cache_write, + cute_dsl_dspark_rmsnorm_rope_draft_block, is_fused_dspark_attention_preparation_supported, + is_fused_dspark_rmsnorm_rope_supported, ) # ---------------------------------------------------------------------------- @@ -607,7 +601,15 @@ def _dspark_ring_attention_metadata( valid_len: torch.Tensor, window_size: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Map a rolling window to two aliases of the same physical page.""" + """Map a rolling window to two aliases of the same physical page. + + The paged kernel reads rows ``[end - 128, end)`` of a ``2 * window_size`` + ring. Both page-table columns alias the request's one ``window_size``-row page + (the page size must equal the kernel's 128-row window), and ``end`` is shifted + by ``window_size`` so every window row lands in one of the two aliases and the + kernel's ``min(128, position)`` bound never truncates ``valid_len``. Capacity + ``2 * window_size`` keeps the capacity mask from dropping rows. + """ slots_i32 = slots.to(torch.int32) return ( torch.stack((slots_i32, slots_i32), dim=1), @@ -2177,6 +2179,13 @@ def _bind_attn_projections(self) -> None: ) qkv_in_features = attn.kv_a_proj_with_mqa.in_features q_lora_rank = attn.q_a_layernorm.weight.shape[0] + # Every WKV consumer slices the trailing head_dim columns after q_lora_rank. + qkv_out_features = attn.kv_a_proj_with_mqa.out_features + if qkv_out_features != q_lora_rank + head_dim: + raise ValueError( + "DSpark fused q_a+wkv projection must have q_lora_rank + head_dim = " + f"{q_lora_rank} + {head_dim} outputs, got {qkv_out_features}" + ) use_fused_qkv_input_projection = _can_use_fused_fp8_input_projection( attn.kv_a_proj_with_mqa, qkv_in_features ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py index 71a84df8303e..7afb43e2def0 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py @@ -838,6 +838,7 @@ def test_bind_attn_projections_reuses_loaded_mla_modules(monkeypatch): qkv_proj = Mock() qkv_proj.in_features = 4 + qkv_proj.out_features = 14 q_b_proj = Mock() q_b_proj.in_features = 6 o_b_proj = Mock() @@ -934,6 +935,10 @@ def test_bind_attn_projections_reuses_loaded_mla_modules(monkeypatch): assert model._context_wkv_fp8 is context_wkv assert stage._dspark_context_wkv_slice is None + qkv_proj.out_features = 15 + with pytest.raises(ValueError, match="q_lora_rank \\+ head_dim"): + DSv4DSparkDraftModel._bind_attn_projections(model) + @pytest.mark.cpu_only @pytest.mark.parametrize( From fbf3b553e8411448ebd442156e7578184875cd13 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Mon, 9 Nov 2026 19:17:08 -0800 Subject: [PATCH 04/19] [None][chore] DSpark: restore the dspark_attention_forward name and export It is the live non-batched path and the reference for the batched path. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/models/modeling_dspark.py | 9 +++++---- .../speculative/hw_agnostic/test_dspark_attention.py | 6 +++--- .../speculative/hw_agnostic/test_dspark_cuda_graph.py | 4 ++-- .../_torch/speculative/hw_agnostic/test_dspark_v41.py | 4 ++-- 4 files changed, 12 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index f5e60a0f61bf..1605ad24291d 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -479,7 +479,7 @@ def _rmsnorm_rope_batched( return t -def _dspark_attention_forward_legacy( +def dspark_attention_forward( x: torch.Tensor, main_x: torch.Tensor, start_pos: int, @@ -662,7 +662,7 @@ def dspark_attention_forward_batched( This path belongs to the embedded DeepSeek-V4-Pro draft. Standalone DSpark drafters (including Kimi-K3-DSpark) use the DFlash paged-KV decode instead. - Numerically identical, per request, to :func:`_dspark_attention_forward_legacy`, but + Numerically identical, per request, to :func:`dspark_attention_forward`, but free of host syncs and data-dependent shapes so it can be captured into a CUDA graph (the one-engine drafter runs inside the target's graph). The differences from the scalar path are purely mechanical: @@ -1600,7 +1600,7 @@ def write_context_windows( """Write captured-context ``main_kv`` into the rolling per-stage KV windows. Replicates exactly the per-position context write that - :func:`_dspark_attention_forward_legacy` performs each generation step + :func:`dspark_attention_forward` performs each generation step (``main_kv = RoPE_pos(rmsnorm(wkv @ main_x))`` written at ``window[pos % window_size]``), but for an arbitrary set of ``(captured-hidden, position)`` pairs. Used to (a) seed a request's @@ -1801,7 +1801,7 @@ def _forward_stage( **self._attn_params, ) else: - attn = _dspark_attention_forward_legacy( + attn = dspark_attention_forward( layer_input, main_x, start_pos, @@ -4246,6 +4246,7 @@ def _build_dspark_draft(model_config, draft_config, lm_head, model): "precompute_dspark_freqs_cis", "apply_dspark_rotary", "apply_dspark_rotary_batched", + "dspark_attention_forward", "dspark_attention_forward_batched", # Block draft I/O. "build_draft_input_ids", diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py index 7afb43e2def0..ae9033b50061 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py @@ -30,10 +30,10 @@ import tensorrt_llm._torch.models.modeling_dspark as modeling_dspark from tensorrt_llm._torch.models.modeling_dspark import ( DSv4DSparkDraftModel, - _dspark_attention_forward_legacy, _output_projection, _qkv_projections, apply_dspark_rotary, + dspark_attention_forward, dspark_sparse_attn, get_dspark_topk_idxs, precompute_dspark_freqs_cis, @@ -487,7 +487,7 @@ def _make_attn_inputs(seed=0): def _run(g): - return _dspark_attention_forward_legacy( + return dspark_attention_forward( g["x"], g["main_x"], g["start_pos"], @@ -1250,7 +1250,7 @@ def test_rolling_attention_uses_projection_modules(monkeypatch, batched): args = (torch.full((g["b"],), g["start_pos"]),) tail = (torch.arange(g["b"]),) else: - fn = _dspark_attention_forward_legacy + fn = dspark_attention_forward args = (g["start_pos"],) tail = () expected = fn(g["x"], g["main_x"], *args, expected_cache, *tail, **weights) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py index ccb5d852a1f9..95fa963affed 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py @@ -28,9 +28,9 @@ import torch from tensorrt_llm._torch.models.modeling_dspark import ( - _dspark_attention_forward_legacy, apply_dspark_rotary, apply_dspark_rotary_batched, + dspark_attention_forward, dspark_attention_forward_batched, get_dspark_topk_idxs, get_dspark_topk_idxs_batched, @@ -112,7 +112,7 @@ def _scalar_reference(g): """Run the validated scalar attention once per request and stack the outputs.""" outs = [] for i, sp in enumerate(g["start_positions"]): - out_i = _dspark_attention_forward_legacy( + out_i = dspark_attention_forward( g["x"][i : i + 1], g["main_x"][i : i + 1], int(sp), diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py index f6c0fdb602b3..db9df17e20e7 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py @@ -370,7 +370,7 @@ def test_v41_stage_collapses_with_previous_not_current_mix(monkeypatch): model_config=SimpleNamespace(mapping=SimpleNamespace(enable_attention_dp=False, tp_size=1)), ) monkeypatch.setattr( - modeling_dspark, "_dspark_attention_forward_legacy", Mock(return_value=attention_output) + modeling_dspark, "dspark_attention_forward", Mock(return_value=attention_output) ) result, next_pre = modeling_dspark.DSv4DSparkDraftModel._forward_stage( model, @@ -431,7 +431,7 @@ def capture_attention(q, *args): x, main_x, torch.tensor([1]), window, torch.tensor([0]), **kwargs ) else: - modeling_dspark._dspark_attention_forward_legacy(x, main_x, 1, window, **kwargs) + modeling_dspark.dspark_attention_forward(x, main_x, 1, window, **kwargs) q = F.linear(x, q_a) q = q * torch.rsqrt(q.square().mean(-1, keepdim=True) + 1e-20) * norm_weight q = F.linear(q, q_b).unflatten(-1, (heads, head_dim)) From 6fa912123255df5d0fc6b52c428bdd9bdc80ac39 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Mon, 9 Nov 2026 19:19:02 -0800 Subject: [PATCH 05/19] [None][fix] DSpark attention op: reject zero page stride, report rejected values Split merged input checks so each error names the offending value, drop unused _DSPARK_* constants, and rename identifiers that no longer match their contents. Signed-off-by: Zhenhuan Chen --- .../kernels/blackwell/dspark/attention.py | 18 ++-- .../blackwell/dspark/attention_kernel.py | 20 ++-- .../custom_ops/dspark_attention_custom_op.py | 100 +++++++++++------- .../dspark_rmsnorm_rope_custom_op.py | 35 +++--- .../test_dspark_cute_dsl_attention.py | 15 +-- 5 files changed, 111 insertions(+), 77 deletions(-) diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention.py b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention.py index 6b29cb9c2871..7834b24fb06e 100644 --- a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention.py +++ b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention.py @@ -205,18 +205,18 @@ class DSparkAttention(DSparkAttentionKernel): pv_tiler_mn = (128, 256) qk_tiler_k = 128 page_size_draft = 8 - page_size_win = 128 + window_len = 128 @staticmethod def _compute_grid( - o: cute.Tensor, + q_latent: cute.Tensor, cluster_shape_mnk: cute.Shape, max_active_clusters: int, ) -> Tuple[DSparkPersistentTileSchedulerParams, cute.Shape]: """Build scheduler parameters and cap the persistent DSpark grid.""" tile_sched_params = DSparkPersistentTileSchedulerParams( - cute.size(o.shape[3]), - cute.size(o.shape[2]), + cute.size(q_latent.shape[3]), + cute.size(q_latent.shape[2]), cluster_shape_mnk, ) grid = DSparkPersistentTileScheduler.get_grid_shape(tile_sched_params, max_active_clusters) @@ -237,7 +237,7 @@ def __init__( mma_pv_tiler_mn: Tuple[int, int], max_active_clusters: int, page_size_draft: int, - page_size_win: int, + window_len: int, skip_correction_threshold: float, *, arch_str: str, @@ -256,7 +256,7 @@ def __init__( "mma_qk_tiler_mn": (mma_qk_tiler_mn, self.qk_tiler_mn), "mma_pv_tiler_mn": (mma_pv_tiler_mn, self.pv_tiler_mn), "page_size_draft": (page_size_draft, self.page_size_draft), - "page_size_win": (page_size_win, self.page_size_win), + "window_len": (window_len, self.window_len), "mma_qk_tiler_k": (mma_qk_tiler_k, self.qk_tiler_k), } mismatches = [ @@ -269,7 +269,9 @@ def __init__( if seq_len_q not in (5, 6): raise ValueError(f"DSpark block size must be 5 or 6, got {seq_len_q}") if history_page_size not in (16, 32, 64, 128, 256): - raise ValueError("DSpark history page size must be 16, 32, 64, 128, or 256") + raise ValueError( + f"DSpark history page size must be 16, 32, 64, 128, or 256, got {history_page_size}" + ) if inverse_rope_dim not in (0, 64): raise ValueError(f"DSpark inverse_rope_dim must be 0 or 64, got {inverse_rope_dim}") if output_groups is not None and ( @@ -285,7 +287,7 @@ def __init__( mma_pv_tiler_mn, max_active_clusters, page_size_draft, - page_size_win, + window_len, skip_correction_threshold, arch_str=arch_str, seq_len_q=seq_len_q, diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention_kernel.py b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention_kernel.py index 876f298662fc..a4625958e270 100644 --- a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention_kernel.py +++ b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention_kernel.py @@ -58,12 +58,8 @@ epilogue into a single launch. KV arrives as two paged streams (the 128-row context window and the draft block's own rows); see ``attention.py`` for the supported DSpark contract and -``tests/unittest/_torch/speculative/test_dspark_paged_attention.py`` for -validation coverage. - -The split-KV/LSE/variable-sequence machinery of the generic attention kernel and its -standalone harness have been removed, and the inverse-RoPE epilogue fusion was -added here. +``tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py`` for +validation coverage. This specialization has no split-KV/LSE reduction. """ LOG2_E = 1.4426950408889634074 @@ -150,7 +146,7 @@ def __init__( mma_pv_tiler_mn: Tuple[int, int], max_active_clusters: int, page_size_draft: int, - page_size_win: int, + window_len: int, skip_correction_threshold: float, *, arch_str: str, @@ -181,7 +177,7 @@ def __init__( self.max_active_clusters = max_active_clusters self.skip_correction_threshold = skip_correction_threshold self.page_size_draft = page_size_draft - self.page_size_win = page_size_win + self.window_len = window_len # Physical page sizes define the page-table/cache ABI. TMA page spans # normally match them, but a specialization may use a wider logical # span so one OOB-zero-filling TMA replaces many tiny page copies. @@ -1180,9 +1176,9 @@ def is_score_valid( total_seq_len, ) if cutlass.const_expr(self.window_valid_len_from_tensor): - age = ( - window_end_pos - 1 - local_col + cutlass.Int32(self.page_size_win) - ) % cutlass.Int32(self.page_size_win) + age = (window_end_pos - 1 - local_col + cutlass.Int32(self.window_len)) % cutlass.Int32( + self.window_len + ) valid = valid and (k_index != 0 or cute.elem_less(age, window_valid_len)) return valid @@ -2745,7 +2741,7 @@ def _create_tile_scheduler(tile_sched_params, blk_coord: cute.Coord, grid_shape: @staticmethod def _compute_grid( - o: cute.Tensor, + q_latent: cute.Tensor, cluster_shape_mnk: Tuple[int, int, int], max_active_clusters: int, ): diff --git a/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py index b3030ecfc748..6102e62e59b4 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py @@ -26,13 +26,6 @@ cute_dsl_dspark_rmsnorm_rope_draft_block, ) -_DSPARK_NUM_HEADS = 128 -_DSPARK_HEAD_DIM = 512 -_DSPARK_WINDOW_SIZE = 128 -_DSPARK_PAGE_SIZE = 8 -_DSPARK_BLOCK_SIZES = (5, 6) -_DSPARK_ROPE_DIM = 64 - _DSPARK_CONTEXT_WKV_MIN_M_BUCKET = 128 @@ -51,12 +44,13 @@ def _gen_dspark_context_wkv_m_profiles(max_m: int) -> Tuple[int, ...]: return tuple(buckets) -_DSV4_DSPARK_NUM_HEADS = (64, 128) +_DSV4_DSPARK_SUPPORTED_NUM_HEADS = (64, 128) _DSV4_DSPARK_HEAD_DIM = 512 _DSV4_DSPARK_WINDOW_SIZE = 128 _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE = 8 _DSV4_DSPARK_BLOCK_SIZES = (5, 6) _DSV4_DSPARK_ROPE_DIM = 64 +_DSV4_DSPARK_PAGE_SIZES = (16, 32, 64, 128, 256) _dspark_attention_kernel_cache: dict[ @@ -183,41 +177,71 @@ def _run_paged_dspark_attention( raise RuntimeError("DSpark fused FP8 output requires SM107") batch, block, heads, dim = q.shape page_size = kv_pages.shape[1] - if ( - block not in _DSV4_DSPARK_BLOCK_SIZES - or heads not in _DSV4_DSPARK_NUM_HEADS - or dim != _DSV4_DSPARK_HEAD_DIM - or page_size not in (16, 32, 64, 128, 256) - ): + if block not in _DSV4_DSPARK_BLOCK_SIZES: raise ValueError( - "DSpark requires block 5/6, 64/128 heads, head_dim 512 and 16/32/64/128/256-token pages" + f"DSpark block size must be one of {_DSV4_DSPARK_BLOCK_SIZES}, got {block}" ) - if q.dtype != torch.bfloat16 or draft_block.dtype != q.dtype or kv_pages.dtype != q.dtype: - raise ValueError("DSpark Q, draft KV and pages must use BF16") - if draft_block.shape != (batch, 8, dim) or kv_pages.shape[2] != dim: - raise ValueError("DSpark draft KV must have shape [batch, 8, 512]") - if ( - block_tables.dtype != torch.int32 - or block_tables.shape[0] != batch - or block_tables.stride(1) != 1 + if heads not in _DSV4_DSPARK_SUPPORTED_NUM_HEADS: + raise ValueError(f"DSpark requires {_DSV4_DSPARK_SUPPORTED_NUM_HEADS} heads, got {heads}") + if dim != _DSV4_DSPARK_HEAD_DIM: + raise ValueError(f"DSpark head_dim must be {_DSV4_DSPARK_HEAD_DIM}, got {dim}") + if page_size not in _DSV4_DSPARK_PAGE_SIZES: + raise ValueError( + f"DSpark page size must be one of {_DSV4_DSPARK_PAGE_SIZES}, got {page_size}" + ) + for name, value in (("Q", q), ("draft KV", draft_block), ("pages", kv_pages)): + if value.dtype != torch.bfloat16: + raise ValueError(f"DSpark {name} must use BF16, got {value.dtype}") + expected_draft_shape = (batch, _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE, dim) + if draft_block.shape != expected_draft_shape: + raise ValueError( + f"DSpark draft KV must have shape {expected_draft_shape}, got {tuple(draft_block.shape)}" + ) + if kv_pages.ndim != 3 or kv_pages.shape[2] != dim: + raise ValueError( + f"DSpark pages must have shape [pages, page_size, {dim}], got {tuple(kv_pages.shape)}" + ) + if block_tables.dtype != torch.int32 or block_tables.stride(1) != 1: + raise ValueError( + "DSpark page tables must have contiguous INT32 rows, got " + f"{block_tables.dtype} with strides {block_tables.stride()}" + ) + if block_tables.shape[0] != batch: + raise ValueError(f"DSpark page tables must have {batch} rows, got {block_tables.shape[0]}") + for name, value in ( + ("positions", positions), + ("valid lengths", valid_lengths), + ("capacities", capacities), ): - raise ValueError("DSpark page tables must have contiguous INT32 rows") - for value in (positions, valid_lengths, capacities): if value.dtype != torch.int64 or value.shape != (batch,) or not value.is_contiguous(): raise ValueError( - "DSpark positions, lengths and capacities must be contiguous INT64 batch vectors" + f"DSpark {name} must be a contiguous INT64 [{batch}] vector, got " + f"{value.dtype} {tuple(value.shape)} with strides {value.stride()}" ) if attn_sink.dtype != torch.float32 or attn_sink.shape != (heads,): - raise ValueError("DSpark attention sink must be FP32 with one entry per head") - if inverse_rope_freqs.dtype != torch.float32 or inverse_rope_freqs.shape != ( - batch, - block, - 32, - 2, + raise ValueError( + f"DSpark attention sink must be FP32 [{heads}], got " + f"{attn_sink.dtype} {tuple(attn_sink.shape)}" + ) + expected_freqs_shape = (batch, block, _DSV4_DSPARK_ROPE_DIM // 2, 2) + if ( + inverse_rope_freqs.dtype != torch.float32 + or inverse_rope_freqs.shape != expected_freqs_shape ): - raise ValueError("DSpark inverse RoPE frequencies must be FP32 [batch, block, 32, 2]") - if kv_pages.stride(1) != dim or kv_pages.stride(2) != 1 or kv_pages.stride(0) % 8: - raise ValueError("DSpark pages require contiguous tokens and an aligned page stride") + raise ValueError( + f"DSpark inverse RoPE frequencies must be FP32 {list(expected_freqs_shape)}, got " + f"{inverse_rope_freqs.dtype} {tuple(inverse_rope_freqs.shape)}" + ) + if ( + kv_pages.stride(0) <= 0 + or kv_pages.stride(0) % 8 + or kv_pages.stride(1) != dim + or kv_pages.stride(2) != 1 + ): + raise ValueError( + "DSpark pages require contiguous tokens and a positive 16-byte-aligned page " + f"stride, got strides {kv_pages.stride()}" + ) for value in (q, draft_block, attn_sink, inverse_rope_freqs): if not value.is_contiguous(): raise ValueError("DSpark query, draft, sink and frequencies must be contiguous") @@ -459,7 +483,7 @@ def is_dsv4_dspark_attention_config_supported( """Return whether static model geometry matches the DSV4 specialization.""" return ( block_size in _DSV4_DSPARK_BLOCK_SIZES - and num_heads in _DSV4_DSPARK_NUM_HEADS + and num_heads in _DSV4_DSPARK_SUPPORTED_NUM_HEADS and head_dim == _DSV4_DSPARK_HEAD_DIM and window_size == _DSV4_DSPARK_WINDOW_SIZE ) @@ -501,9 +525,9 @@ def is_fused_dsv4_dspark_attention_supported( ) if q.shape[1] not in _DSV4_DSPARK_BLOCK_SIZES: return _log_unsupported(f"draft block size must be 5 or 6, got {q.shape[1]}", "block_size") - if q.shape[2] not in _DSV4_DSPARK_NUM_HEADS or q.shape[3] != _DSV4_DSPARK_HEAD_DIM: + if q.shape[2] not in _DSV4_DSPARK_SUPPORTED_NUM_HEADS or q.shape[3] != _DSV4_DSPARK_HEAD_DIM: return _log_unsupported( - "q must have 128 heads and head_dim 512; " + f"q must have {_DSV4_DSPARK_SUPPORTED_NUM_HEADS} heads and head_dim 512; " f"got {q.shape[2]} heads and head_dim {q.shape[3]}", "q_shape", ) diff --git a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py index 2b384dc798a1..8389f0a427e6 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py @@ -629,22 +629,31 @@ def _( def _validate_dspark_preparation( x: torch.Tensor, weight: torch.Tensor, freqs: torch.Tensor ) -> None: - if ( - x.ndim != 3 - or x.shape[-1] != _DSV4_DSPARK_HEAD_DIM - or x.dtype != torch.bfloat16 - or weight.dtype != x.dtype - or weight.shape != (_DSV4_DSPARK_HEAD_DIM,) - or freqs.dtype != torch.float32 - or freqs.shape != (x.shape[0] * x.shape[1], _DSV4_DSPARK_ROPE_DIM // 2, 2) - or not all( - t.is_cuda and t.device == x.device and t.is_contiguous() for t in (x, weight, freqs) + if x.ndim != 3 or x.shape[-1] != _DSV4_DSPARK_HEAD_DIM or x.dtype != torch.bfloat16: + raise ValueError( + f"DSpark preparation input must be BF16 [batch, tokens, {_DSV4_DSPARK_HEAD_DIM}], " + f"got {x.dtype} {tuple(x.shape)}" ) - ): + if weight.dtype != x.dtype or weight.shape != (_DSV4_DSPARK_HEAD_DIM,): + raise ValueError( + f"DSpark preparation weight must be BF16 [{_DSV4_DSPARK_HEAD_DIM}], " + f"got {weight.dtype} {tuple(weight.shape)}" + ) + expected_freqs_shape = (x.shape[0] * x.shape[1], _DSV4_DSPARK_ROPE_DIM // 2, 2) + if freqs.dtype != torch.float32 or freqs.shape != expected_freqs_shape: raise ValueError( - "DSpark preparation requires contiguous BF16 [batch, tokens, 512] inputs, " - "BF16 [512] weights and FP32 [batch * tokens, 32, 2] frequencies" + f"DSpark preparation frequencies must be FP32 {list(expected_freqs_shape)}, " + f"got {freqs.dtype} {tuple(freqs.shape)}" ) + for name, tensor in (("input", x), ("weight", weight), ("frequencies", freqs)): + if not tensor.is_cuda or tensor.device != x.device: + raise ValueError( + f"DSpark preparation {name} must be on {x.device}, got {tensor.device}" + ) + if not tensor.is_contiguous(): + raise ValueError( + f"DSpark preparation {name} must be contiguous, got strides {tensor.stride()}" + ) def precompile_dspark_fp8_input_fusions( diff --git a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py index 4add57165f85..5e01b463cea6 100644 --- a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py +++ b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py @@ -183,16 +183,17 @@ def test_fused_dsv4_dspark_attention_support_gate_logs_rejection(monkeypatch, in "invalid_case,reason", [ ("q_dtype", "must use BF16"), - ("block_size", "requires block 5/6"), - ("heads", "64/128 heads"), - ("page_size", "16/32/64/128/256-token pages"), + ("block_size", "block size must be one of"), + ("heads", r"requires \(64, 128\) heads"), + ("page_size", "page size must be one of"), + ("page_stride", "positive 16-byte-aligned page stride"), ("draft_shape", "draft KV must have shape"), ("draft_layout", "must be contiguous"), ("table_dtype", "contiguous INT32 rows"), ("table_layout", "contiguous INT32 rows"), - ("positions", "contiguous INT64 batch vectors"), - ("valid_lengths", "contiguous INT64 batch vectors"), - ("capacities", "contiguous INT64 batch vectors"), + ("positions", "positions must be a contiguous INT64"), + ("valid_lengths", "valid lengths must be a contiguous INT64"), + ("capacities", "capacities must be a contiguous INT64"), ("sink_dtype", "sink must be FP32"), ("freqs_shape", "inverse RoPE frequencies must be FP32"), ], @@ -209,6 +210,8 @@ def test_fused_dsv4_dspark_attention_rejects_invalid_inputs_before_launch( inputs["q"] = inputs["q"][:, :, :24].contiguous() elif invalid_case == "page_size": inputs["kv_pages"] = inputs["kv_pages"][:, :24] + elif invalid_case == "page_stride": + inputs["kv_pages"] = inputs["kv_pages"][:1].expand_as(inputs["kv_pages"]) elif invalid_case == "draft_shape": inputs["draft_block"] = inputs["draft_block"][:, :7].contiguous() elif invalid_case == "draft_layout": From fce1445b80af6c7f6858412e2bf872acbac4e0d0 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Mon, 9 Nov 2026 19:19:52 -0800 Subject: [PATCH 06/19] [None][chore] DSpark draft-block prep: drop the per-device compile cache key TensorRT-LLM runs one GPU per process and the sibling DSpark kernels are not keyed by device either. Signed-off-by: Zhenhuan Chen --- .../_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py | 6 ++---- .../_torch/speculative/test_dspark_cute_dsl_attention.py | 3 +-- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py index 8389f0a427e6..3cf4f7aa0586 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py @@ -355,8 +355,7 @@ def _compile_dspark_rmsnorm_rope_cache_write(eps: float): @functools.cache -def _compile_dspark_rmsnorm_rope_draft_block(block_size: int, eps: float, device: int): - # The device argument scopes this compilation cache to its CUDA device. +def _compile_dspark_rmsnorm_rope_draft_block(block_size: int, eps: float): if torch.cuda.is_current_stream_capturing(): raise RuntimeError( "DSpark RMSNorm/RoPE draft-block must be warmed up before CUDA graph capture" @@ -603,8 +602,7 @@ def cute_dsl_dspark_rmsnorm_rope_draft_block( block_size = x.shape[1] if block_size not in (5, 6): raise ValueError("DSpark draft block size must be 5 or 6") - with torch.cuda.device(x.device): - compiled = _compile_dspark_rmsnorm_rope_draft_block(block_size, eps, x.device.index) + compiled = _compile_dspark_rmsnorm_rope_draft_block(block_size, eps) output = x.new_empty((x.shape[0], _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE, _DSV4_DSPARK_HEAD_DIM)) compiled( x.view(-1, _DSV4_DSPARK_HEAD_DIM), diff --git a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py index 5e01b463cea6..dd4ed2384ab7 100644 --- a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py +++ b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py @@ -400,7 +400,6 @@ def unexpected_compile(*args, **kwargs): def test_preparation_cache_misses_reject_cuda_graph_capture(monkeypatch): preparation_op._compile_dspark_rmsnorm_rope_cache_write.cache_clear() preparation_op._compile_dspark_rmsnorm_rope_draft_block.cache_clear() - device = torch.cuda.current_device() monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) def unexpected_compile(*args, **kwargs): @@ -410,7 +409,7 @@ def unexpected_compile(*args, **kwargs): with pytest.raises(RuntimeError, match="cache-write must be warmed up"): preparation_op._compile_dspark_rmsnorm_rope_cache_write(1e-6) with pytest.raises(RuntimeError, match="draft-block must be warmed up"): - preparation_op._compile_dspark_rmsnorm_rope_draft_block(5, 1e-6, device) + preparation_op._compile_dspark_rmsnorm_rope_draft_block(5, 1e-6) def test_preparation_self_jit_covers_dynamic_batches(): From 1266c260f6515f2d88d4af5c9bfe091d1f9a6370 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Mon, 9 Nov 2026 19:21:53 -0800 Subject: [PATCH 07/19] [None][chore] DSpark: remove the unwired fused_dsv4_dspark_attention_fp8 op No model path calls the FP8-output attention variant; drop the op, its output_groups kernel epilogue and its test until a caller consumes it. Signed-off-by: Zhenhuan Chen --- .../kernels/blackwell/dspark/attention.py | 17 +- .../blackwell/dspark/attention_kernel.py | 216 +----------------- .../custom_ops/dspark_attention_custom_op.py | 108 +-------- .../test_dspark_cute_dsl_attention.py | 68 +----- 4 files changed, 21 insertions(+), 388 deletions(-) diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention.py b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention.py index 7834b24fb06e..f30e68044396 100644 --- a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention.py +++ b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention.py @@ -3,7 +3,7 @@ """Blackwell DSpark attention specialization and persistent scheduler.""" -from typing import Optional, Tuple, Type +from typing import Tuple, Type import cutlass import cutlass.cute as cute @@ -244,7 +244,6 @@ def __init__( seq_len_q: int = block_size, mma_qk_tiler_k: int = qk_tiler_k, inverse_rope_dim: int = 0, - output_groups: int | None = None, history_page_size: int = 32, num_heads: int = num_heads, ): @@ -274,12 +273,6 @@ def __init__( ) if inverse_rope_dim not in (0, 64): raise ValueError(f"DSpark inverse_rope_dim must be 0 or 64, got {inverse_rope_dim}") - if output_groups is not None and ( - output_groups <= 0 or self.num_heads % output_groups != 0 - ): - raise ValueError( - f"DSpark output groups must divide {self.num_heads}, got {output_groups}" - ) super().__init__( acc_dtype, @@ -299,7 +292,6 @@ def __init__( self.tma_page_size_draft = self.qk_tiler_mn[1] self.fixed_cache_seq_len = self.window_size + seq_len_q self.inverse_rope_dim = inverse_rope_dim - self.output_groups = output_groups self.history_page_size = history_page_size # Subdivide the 64-row loads at physical-page and cyclic-window boundaries. self.history_tile_sizes = tuple(sorted({1, 4, 16, min(history_page_size, 64)})) @@ -322,7 +314,6 @@ def wrapper( attn_sink_ptr: cute.Pointer, inverse_rope_freqs_ptr: cute.Pointer, o_tensor: cute.Tensor, - o_scale_tensor: Optional[cute.Tensor], softmax_scale: cutlass.Float32, stream: cuda.CUstream, ): @@ -364,10 +355,7 @@ def wrapper( page_table_ptr, layout=cute.make_layout((table_width, batch_size), stride=(1, table_stride)), ) - if cutlass.const_expr(self.output_groups is None): - output = cute.make_tensor(o_tensor.iterator, layout=q_layout) - else: - output = o_tensor + output = cute.make_tensor(o_tensor.iterator, layout=q_layout) cache_seqs = cute.make_tensor( cache_seqs_ptr, layout=cute.make_layout((batch_size,), stride=(1,)), @@ -397,7 +385,6 @@ def wrapper( draft, page_table, output, - o_scale_tensor, cache_seqs, valid_len, capacities, diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention_kernel.py b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention_kernel.py index a4625958e270..a69e5551b12b 100644 --- a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention_kernel.py +++ b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention_kernel.py @@ -37,10 +37,9 @@ import cutlass.pipeline as pipeline import cutlass.utils as utils import cutlass.utils.blackwell_helpers as sm100_utils -from cutlass._mlir.dialects import llvm from cutlass.cute.arch import Arch from cutlass.cute.nvgpu import OperandMajorMode, tcgen05 -from cutlass.cutlass_dsl import BaseDSL, Int32, T, dsl_user_op +from cutlass.cutlass_dsl import BaseDSL from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait try: @@ -65,75 +64,6 @@ LOG2_E = 1.4426950408889634074 -@dsl_user_op -def _store_global_b64(gmem_ptr, value_lo: Int32, value_hi: Int32, *, loc=None, ip=None) -> None: - gmem_addr = gmem_ptr.toint(loc=loc, ip=ip) - llvm.inline_asm( - None, - [ - gmem_addr.ir_value(loc=loc, ip=ip), - Int32(value_lo).ir_value(loc=loc, ip=ip), - Int32(value_hi).ir_value(loc=loc, ip=ip), - ], - "{\n .reg .b64 packed;\n mov.b64 packed, {$1, $2};\n st.global.u64 [$0], packed;\n}", - "l,r,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -@dsl_user_op -def _load_shared_b32(smem_ptr, *, loc=None, ip=None) -> Int32: - smem_addr = smem_ptr.toint(loc=loc, ip=ip) - # The FP8 epilogue reuses this scratch space for each output column tile. - # Keep each read ordered with its stores and prevent reuse across tiles. - return Int32( - llvm.inline_asm( - T.i32(), - [smem_addr.ir_value(loc=loc, ip=ip)], - "ld.shared.b32 $0, [$1];", - "=r,r,~{memory}", - has_side_effects=True, - loc=loc, - ip=ip, - ) - ) - - -@dsl_user_op -def _store_shared_b128( - smem_ptr, - value0: Int32, - value1: Int32, - value2: Int32, - value3: Int32, - *, - loc=None, - ip=None, -) -> None: - smem_addr = smem_ptr.toint(loc=loc, ip=ip) - llvm.inline_asm( - None, - [ - smem_addr.ir_value(loc=loc, ip=ip), - Int32(value0).ir_value(loc=loc, ip=ip), - Int32(value1).ir_value(loc=loc, ip=ip), - Int32(value2).ir_value(loc=loc, ip=ip), - Int32(value3).ir_value(loc=loc, ip=ip), - ], - "st.shared.v4.b32 [$0], {$1, $2, $3, $4};", - "r,r,r,r,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - class DSparkAttentionKernel: # The DSpark accumulator/correction layouts occupy a fixed 512-column # TMEM footprint on every supported architecture. @@ -288,7 +218,6 @@ def __call__( c_latent_draft: cute.Tensor, page_table_win: cute.Tensor, o: cute.Tensor, - o_scale: Optional[cute.Tensor], cache_seqs: Optional[cute.Tensor], window_valid_lens: Optional[cute.Tensor], capacities: cute.Tensor, @@ -327,18 +256,8 @@ def __call__( raise ValueError("c_latent_draft must have leading dimension 1") if cutlass.const_expr(c_latent_win.stride[1] != 1): raise ValueError("c_latent_win must have leading dimension 1") - if cutlass.const_expr(self.output_groups is None): - if cutlass.const_expr(o.element_type is not cutlass.BFloat16 or o.stride[1] != 1): - raise ValueError("o must be leading-dimension-1 BF16") - else: - if cutlass.const_expr( - o.element_type is not cutlass.Float8E4M3FN - or o_scale is None - or o_scale.element_type is not cutlass.Float32 - or o.stride[2] != 1 - or o_scale.stride[2] != 1 - ): - raise ValueError("quantized o must be contiguous FP8 with contiguous FP32 scales") + if cutlass.const_expr(o.element_type is not cutlass.BFloat16 or o.stride[1] != 1): + raise ValueError("o must be leading-dimension-1 BF16") c_latent_draft_transpose_layout = cute.select(c_latent_draft.layout, mode=[1, 0, 2]) c_latent_draft_transpose = cute.make_tensor( @@ -555,17 +474,6 @@ class AttentionKernelSharedStorage: epilogue_smem_exchange: cute.struct.MemRange[ self.acc_dtype, self.num_compute_warps * self.threads_per_warp ] - fp8_epilogue_scratch: cute.struct.Align[ - cute.struct.MemRange[ - cutlass.Int32, - ( - self.num_compute_warps * self.threads_per_warp * 32 - if self.output_groups is not None - else 1 - ), - ], - 128, - ] smem_q_latent: cute.struct.Align[ cute.struct.MemRange[self.q_dtype, cute.cosize(q_latent_smem_layout_staged)], 1024, @@ -604,7 +512,6 @@ class AttentionKernelSharedStorage: tma_tensor_c_latent_transpose_draft, page_table_win, o, - o_scale, cache_seqs, window_valid_lens, capacities, @@ -681,7 +588,6 @@ def kernel( mCLT_draft: cute.Tensor, mPT_win: cute.Tensor, mO: Optional[cute.Tensor], - mOScale: Optional[cute.Tensor], cache_seqs: cute.Tensor, window_valid_lens: cute.Tensor, capacities: cute.Tensor, @@ -805,15 +711,6 @@ def kernel( epilogue_smem_exchange = storage.epilogue_smem_exchange.get_tensor( cute.make_layout(self.num_compute_warps * self.threads_per_warp) ) - fp8_epilogue_scratch = None - if cutlass.const_expr(self.output_groups is not None): - fp8_epilogue_scratch = storage.fp8_epilogue_scratch.get_tensor( - cute.make_layout( - (self.num_compute_warps, self.threads_per_warp, 32), - stride=(self.threads_per_warp * 32, 32, 1), - ) - ) - # # Cluster wait before tensor memory alloc # @@ -1044,7 +941,7 @@ def kernel( compute_common_params = SimpleNamespace( blk_coord=blk_coord, smem_exchange=softmax_smem_exchange, - mO=mQL if cutlass.const_expr(self.output_groups is not None) else mO, + mO=mO, K=self.get_cache_seq_len(cache_seqs, blk_coord[2]), window_valid_len=self.get_window_valid_len( window_valid_lens, capacities, cache_seqs, blk_coord[2] @@ -1099,7 +996,7 @@ def kernel( compute_common_params = SimpleNamespace( blk_coord=blk_coord, smem_exchange=epilogue_smem_exchange, - mO=mQL if cutlass.const_expr(self.output_groups is not None) else mO, + mO=mO, K=self.get_cache_seq_len(cache_seqs, blk_coord[2]), L=mCL_draft.shape[1], H=mQL.shape[0], @@ -1113,9 +1010,6 @@ def kernel( output_scale=output_scale, softmax_scale_log2=softmax_scale_log2, mFreqs=inverse_rope_freqs, - mOQuant=mO, - mOScale=mOScale, - sOQuant=fp8_epilogue_scratch, ) p_cor_consumer_state, mma_o_consumer_state = self.correction( compute_common_params, @@ -2497,102 +2391,12 @@ def epilogue( tR2G_rO_src[i] = (re * cos + im * sin).to(self.o_dtype) tR2G_rO_src[i + 1] = (im * cos - re * sin).to(self.o_dtype) - if cutlass.const_expr(self.output_groups is not None and cute.size(tR2G_rO_src) != 128): - raise ValueError("DSpark FP8 output requires one 128-value block per thread") - if cute.elem_less(tTR_cO[0][0], common_params.H): - if cutlass.const_expr(self.output_groups is None): - cute.autovec_copy( - tR2G_rO_src, - tR2G_rO_dst, - l1c_evict_priority=cute.nvgpu.CacheEvictionPriority.NO_ALLOCATE, - ) - else: - # Each correction thread owns one complete 1x128 quant - # block. Use eight independent chains so the amax is not - # serialized behind 128 dependent fmax instructions. - partial_absmax = cute.make_rmem_tensor((8,), cutlass.Float32) - for i in cutlass.range_constexpr(8): - partial_absmax[i] = cutlass.Float32(0.0) - for i in cutlass.range_constexpr(0, cute.size(tR2G_rO_src), 8): - for lane in cutlass.range_constexpr(8): - value = cutlass.Float32(tR2G_rO_src[i + lane]) - partial_absmax[lane] = cute.arch.fmax( - partial_absmax[lane], - cute.arch.fmax(value, -value), - ) - for i in cutlass.range_constexpr(4): - partial_absmax[i] = cute.arch.fmax(partial_absmax[i], partial_absmax[i + 4]) - for i in cutlass.range_constexpr(2): - partial_absmax[i] = cute.arch.fmax(partial_absmax[i], partial_absmax[i + 2]) - block_absmax = cute.arch.fmax(partial_absmax[0], partial_absmax[1]) - block_absmax = cute.arch.fmax(block_absmax, cutlass.Float32(1e-10)) - quant_scale = cutlass.Float32(448.0) / block_absmax - - head = tTR_cO[0][0] - group_size = self.num_heads // self.output_groups - group = head // group_size - row = common_params.blk_coord[2] * self.seq_len_q + common_params.blk_coord[1] - scale_block = ((head % group_size) * self.latent_dim + tTR_cO[0][1]) // 128 - epilogue_params.mOScale[group, scale_block, row] = ( - cutlass.Float32(1.0) / quant_scale - ) - quantized = cute.make_rmem_tensor( - (cute.size(tR2G_rO_src),), cutlass.Float8E4M3FN - ) - for i in cutlass.range( - cute.size(tR2G_rO_src), vectorize=True, unroll_full=True - ): - quantized[i] = (cutlass.Float32(tR2G_rO_src[i]) * quant_scale).to( - cutlass.Float8E4M3FN - ) - packed_output = cute.recast_tensor(quantized, cutlass.Int32) - lane_idx = tidx % self.threads_per_warp - epilogue_warp_idx = tidx // self.threads_per_warp - scratch_row = cute.domain_offset( - (epilogue_warp_idx, lane_idx, 0), epilogue_params.sOQuant - ).iterator - for i in cutlass.range_constexpr(0, 32, 4): - _store_shared_b128( - scratch_row + i, - packed_output[i], - packed_output[i + 1], - packed_output[i + 2], - packed_output[i + 3], - ) - - # Four adjacent lanes cooperatively drain one scratch row, - # producing fully coalesced 32-byte global sectors. - cute.arch.sync_warp() - row_group_idx = lane_idx // 4 - row_lane_idx = lane_idx % 4 - warp_head_base = head - lane_idx - for row_batch_idx in cutlass.range_constexpr(4): - output_row = row_batch_idx * 8 + row_group_idx - output_head = warp_head_base + output_row - output_group = output_head // group_size - output_group_k = (output_head % group_size) * self.latent_dim + tTR_cO[0][1] - source_row = cute.domain_offset( - (epilogue_warp_idx, output_row, 0), - epilogue_params.sOQuant, - ).iterator - for segment_idx in cutlass.range_constexpr(4): - word_offset = segment_idx * 8 + row_lane_idx * 2 - value_lo = _load_shared_b32(source_row + word_offset) - value_hi = _load_shared_b32(source_row + word_offset + 1) - _store_global_b64( - cute.domain_offset( - ( - output_group, - row, - output_group_k + segment_idx * 32 + row_lane_idx * 8, - ), - epilogue_params.mOQuant, - ).iterator, - value_lo, - value_hi, - ) - cute.arch.sync_warp() + cute.autovec_copy( + tR2G_rO_src, + tR2G_rO_dst, + l1c_evict_priority=cute.nvgpu.CacheEvictionPriority.NO_ALLOCATE, + ) cute.arch.fence_view_async_tmem_load() common_params.mma_o_pipeline.consumer_release(mma_o_consumer_state) diff --git a/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py index 6102e62e59b4..a2c4a1fd27fd 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py @@ -53,9 +53,7 @@ def _gen_dspark_context_wkv_m_profiles(max_m: int) -> Tuple[int, ...]: _DSV4_DSPARK_PAGE_SIZES = (16, 32, 64, 128, 256) -_dspark_attention_kernel_cache: dict[ - tuple[int, int, int, str, int | None], Callable[..., None] -] = {} +_dspark_attention_kernel_cache: dict[tuple[int, int, int, str], Callable[..., None]] = {} def _make_compile_gmem_pointer(dtype: Type[Numeric], assumed_align: int) -> Pointer: @@ -72,7 +70,6 @@ def _compile_dspark_attention( block_size: int, arch_str: str, page_size: int, - output_groups: int | None, num_heads: int, ) -> Callable[..., None]: """Compile the pointer host wrapper without runtime tensor specimens.""" @@ -84,22 +81,6 @@ def _compile_dspark_attention( stride_order=(3, 2, 1, 0), assumed_align=16, ) - output_scale_fake = None - if output_groups is not None: - group_dim = num_heads * head_dim // output_groups - rows = batch * block_size - output_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Float8E4M3FN, - (output_groups, rows, group_dim), - stride_order=(2, 1, 0), - assumed_align=16, - ) - output_scale_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Float32, - (output_groups, group_dim // 128, cute.sym_int(divisibility=4)), - stride_order=(2, 1, 0), - assumed_align=16, - ) stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) # Real tensor specimens used to make PyTorch's primary CUDA context current @@ -121,7 +102,6 @@ def _compile_dspark_attention( inverse_rope_dim=_DSV4_DSPARK_ROPE_DIM, arch_str=arch_str, history_page_size=page_size, - output_groups=output_groups, num_heads=num_heads, ) # Scalars and typed, aligned compile-only pointers define the runtime ABI. @@ -144,7 +124,6 @@ def _compile_dspark_attention( _make_compile_gmem_pointer(cutlass.Float32, 4), _make_compile_gmem_pointer(cutlass.Float32, 4), output_fake, - output_scale_fake, cutlass.Float32(1.0), stream_fake, options="--opt-level 2 --enable-tvm-ffi", @@ -167,14 +146,11 @@ def _run_paged_dspark_attention( attn_sink: torch.Tensor, inverse_rope_freqs: torch.Tensor, softmax_scale: float, - output_groups: int | None = None, -) -> tuple[torch.Tensor, torch.Tensor | None]: +) -> torch.Tensor: """Attend to a 128-token context window and fuse the inverse-RoPE epilogue.""" arch_str = _get_dspark_arch_str() if arch_str is None: raise RuntimeError("Embedded DSpark attention requires SM100, SM103, or SM107") - if output_groups is not None and arch_str != "sm_107": - raise RuntimeError("DSpark fused FP8 output requires SM107") batch, block, heads, dim = q.shape page_size = kv_pages.shape[1] if block not in _DSV4_DSPARK_BLOCK_SIZES: @@ -260,26 +236,14 @@ def _run_paged_dspark_attention( ) ): raise ValueError("All DSpark attention inputs must be on the same CUDA device") - cache_key = (block, heads, page_size, arch_str, output_groups) + cache_key = (block, heads, page_size, arch_str) compiled = _dspark_attention_kernel_cache.get(cache_key) if compiled is None: if torch.cuda.is_current_stream_capturing(): raise RuntimeError("DSpark paged attention must run eagerly before graph capture") - compiled = _compile_dspark_attention(block, arch_str, page_size, output_groups, heads) + compiled = _compile_dspark_attention(block, arch_str, page_size, heads) _dspark_attention_kernel_cache[cache_key] = compiled output = torch.empty_like(q) - output_scale = None - if output_groups is not None: - rows = batch * block - group_dim = heads * dim // output_groups - output = torch.empty( - (output_groups, rows, group_dim), dtype=torch.float8_e4m3fn, device=q.device - ) - output_scale = torch.empty( - (output_groups, group_dim // 128, (rows + 3) // 4 * 4), - dtype=torch.float32, - device=q.device, - ) compiled( batch, kv_pages.shape[0], @@ -296,10 +260,9 @@ def _run_paged_dspark_attention( attn_sink.data_ptr(), inverse_rope_freqs.data_ptr(), output, - output_scale, softmax_scale, ) - return output, output_scale + return output @torch.library.custom_op( @@ -328,7 +291,7 @@ def fused_dsv4_dspark_attention( attn_sink, inverse_rope_freqs, softmax_scale, - )[0] + ) @torch.library.register_fake("trtllm::fused_dsv4_dspark_attention") @@ -347,65 +310,6 @@ def _( return torch.empty_like(q) -@torch.library.custom_op( - "trtllm::fused_dsv4_dspark_attention_fp8", mutates_args=(), device_types="cuda" -) -def fused_dsv4_dspark_attention_fp8( - q: torch.Tensor, - draft_block: torch.Tensor, - kv_pages: torch.Tensor, - block_tables: torch.Tensor, - positions: torch.Tensor, - valid_lengths: torch.Tensor, - capacities: torch.Tensor, - attn_sink: torch.Tensor, - inverse_rope_freqs: torch.Tensor, - softmax_scale: float, - output_groups: int, -) -> tuple[torch.Tensor, torch.Tensor]: - output, scales = _run_paged_dspark_attention( - q, - draft_block, - kv_pages, - block_tables, - positions, - valid_lengths, - capacities, - attn_sink, - inverse_rope_freqs, - softmax_scale, - output_groups, - ) - assert scales is not None - return output, scales - - -@torch.library.register_fake("trtllm::fused_dsv4_dspark_attention_fp8") -def _( - q: torch.Tensor, - draft_block: torch.Tensor, - kv_pages: torch.Tensor, - block_tables: torch.Tensor, - positions: torch.Tensor, - valid_lengths: torch.Tensor, - capacities: torch.Tensor, - attn_sink: torch.Tensor, - inverse_rope_freqs: torch.Tensor, - softmax_scale: float, - output_groups: int, -) -> tuple[torch.Tensor, torch.Tensor]: - rows = q.shape[0] * q.shape[1] - group_dim = q.shape[2] * q.shape[3] // output_groups - return ( - torch.empty((output_groups, rows, group_dim), dtype=torch.float8_e4m3fn, device=q.device), - torch.empty( - (output_groups, group_dim // 128, (rows + 3) // 4 * 4), - dtype=torch.float32, - device=q.device, - ), - ) - - if IS_CUTLASS_DSL_RUBIN_AVAILABLE: from .cute_dsl_custom_ops import CuteDSLMXFP8RubinLinear, _get_mxfp8_gemm_alpha diff --git a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py index dd4ed2384ab7..9d9316e18f38 100644 --- a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py +++ b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py @@ -231,7 +231,7 @@ def unexpected_launch(*args): pytest.fail("invalid inputs reached compilation or a kernel launch") monkeypatch.setattr(attention_op, "_get_dspark_arch_str", lambda: "sm_100") - key = (6, 128, 32, attention_op._get_dspark_arch_str(), None) + key = (6, 128, 32, attention_op._get_dspark_arch_str()) monkeypatch.setattr(attention_op, "_dspark_attention_kernel_cache", {key: unexpected_launch}) monkeypatch.setattr(attention_op, "_compile_dspark_attention", unexpected_launch) with pytest.raises(ValueError, match=reason): @@ -363,11 +363,11 @@ def unexpected_compile_op(*args, **kwargs): patch.setattr(torch, "empty", unexpected_compile_op) patch.setattr(torch, "empty_like", unexpected_compile_op) patch.setattr(attention_op.cute.runtime, "from_dlpack", unexpected_compile_op) - compiled = attention_op._compile_dspark_attention(5, arch_str, 32, None, 128) + compiled = attention_op._compile_dspark_attention(5, arch_str, 32, 128) monkeypatch.setattr( attention_op, "_dspark_attention_kernel_cache", - {(5, 128, 32, arch_str, None): compiled}, + {(5, 128, 32, arch_str): compiled}, ) inputs = _make_inputs(41, block=5, positions=[257, 9]) @@ -533,68 +533,6 @@ def test_batched_attention_respects_partial_window_valid_len(): torch.testing.assert_close(actual[row], expected, rtol=5e-2, atol=3e-2) -@pytest.mark.skipif( - get_sm_version() != 107, - reason="The fused DSpark FP8 output projection is a Rubin specialization", -) -def test_fused_dsv4_dspark_attention_fp8_matches_separate_quant_and_replays_graph(): - from tensorrt_llm._torch.custom_ops.dspark_attention_custom_op import ( - fused_dsv4_dspark_attention, - fused_dsv4_dspark_attention_fp8, - ) - - groups = 16 - inputs = _make_inputs(47, page_size=256, positions=[257, 390], valid_lengths=[3, 5]) - q = inputs["q"] - block = q.shape[1] - - def separate_quantize(): - bf16 = fused_dsv4_dspark_attention(**inputs) - rows = q.shape[0] * block - grouped = bf16.reshape(rows, groups, -1).permute(1, 0, 2).contiguous() - quant_blocks = grouped.view(groups, rows, -1, 128) - amax = quant_blocks.abs().amax(dim=-1).float().clamp_min(1e-10) - quant_scale = 448.0 / amax - quantized = (quant_blocks.float() * quant_scale.unsqueeze(-1)).to(torch.float8_e4m3fn) - scale_rows = (rows + 3) // 4 * 4 - scales = torch.zeros( - groups, - quant_blocks.shape[2], - scale_rows, - device=q.device, - dtype=torch.float32, - ) - scales[..., :rows] = (1.0 / quant_scale).transpose(1, 2) - return quantized.view_as(grouped), scales - - def assert_quantized_close(actual_fp8, actual_scale, expected_fp8, expected_scale): - rows = q.shape[0] * block - torch.testing.assert_close( - actual_scale[..., :rows], expected_scale[..., :rows], rtol=1e-6, atol=0 - ) - actual = actual_fp8.float() - expected = expected_fp8.float() - mismatch_rate = torch.count_nonzero(actual != expected) / actual.numel() - assert mismatch_rate <= 1e-3 - # CuTe and Torch can choose neighboring E4M3 values at a rounding tie. - torch.testing.assert_close(actual, expected, rtol=0.125, atol=0.015625) - - actual_fp8, actual_scale = fused_dsv4_dspark_attention_fp8(**inputs, output_groups=groups) - expected_fp8, expected_scale = separate_quantize() - assert_quantized_close(actual_fp8, actual_scale, expected_fp8, expected_scale) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured_fp8, captured_scale = fused_dsv4_dspark_attention_fp8( - **inputs, output_groups=groups - ) - q.copy_(torch.randn_like(q)) - inputs["valid_lengths"].copy_(torch.tensor([7, 2], device=q.device)) - graph.replay() - expected_fp8, expected_scale = separate_quantize() - assert_quantized_close(captured_fp8, captured_scale, expected_fp8, expected_scale) - - @pytest.mark.skipif( get_sm_version() != 107 or not IS_CUTLASS_DSL_RUBIN_AVAILABLE, reason="The fused DSpark FP8 input producers require SM107 Rubin CuTe DSL", From a6aeac25c85bfd0040434da720db7c2829c93207 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Mon, 9 Nov 2026 19:23:15 -0800 Subject: [PATCH 08/19] [None][chore] DFlash TRTLLM: drop context-kernel leftovers after the folded decode Remove the always-False has_context_attention check and size the multi-CTA counters by the largest head count either launch shape uses. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/models/modeling_dflash.py | 2 +- tensorrt_llm/_torch/speculative/dflash.py | 1 - tensorrt_llm/_torch/speculative/dflash_attention.py | 5 ----- .../hw_agnostic/test_dspark_noncausal_window.py | 7 ++----- 4 files changed, 3 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index 1d94771554b1..a9352c4edc6a 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -1724,7 +1724,7 @@ def _prepare_dflash_trtllm_gen_buffers( sm_count = self._dflash_trtllm_gen_sm_count counter_bytes = trtllm_gen_ops.get_multi_ctas_kv_counter_size( - num_heads * block_size, max_batch_size, sm_count + max(heads for _, heads in shapes), max_batch_size, sm_count ) counters = self._dflash_trtllm_gen_counters counters_need_allocation = ( diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index d7427811ea9d..8868a6690642 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -1571,7 +1571,6 @@ def _lazy_init_ctx_buffers( num_kv_heads=nkv, head_dim=hd, tokens_per_block=page_size, - has_context_attention=False, folded_block_size=( self._compute_block_size if DFLASH_FOLDED_DECODE_LOWERING in lowerings diff --git a/tensorrt_llm/_torch/speculative/dflash_attention.py b/tensorrt_llm/_torch/speculative/dflash_attention.py index 7513e46543e3..8dbf8cc28344 100644 --- a/tensorrt_llm/_torch/speculative/dflash_attention.py +++ b/tensorrt_llm/_torch/speculative/dflash_attention.py @@ -190,7 +190,6 @@ def validate_dflash_trtllm_gen_runtime( num_kv_heads: int, head_dim: int, tokens_per_block: int, - has_context_attention: bool, folded_block_size: Optional[int] = None, ) -> None: """Fail before cache allocation when DFlash's kernel shape is unsupported.""" @@ -220,7 +219,3 @@ def validate_dflash_trtllm_gen_runtime( "DFlash TRTLLM attention does not support context-cache page size " f"{tokens_per_block}; supported sizes are {supported}." ) - if has_context_attention and head_dim in FlashInferTrtllmGenFmha.UNSUPPORTED_HEAD_SIZES_CONTEXT: - raise RuntimeError( - f"DFlash TRTLLM context attention does not support head dimension {head_dim}." - ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py index d799c27c3af3..97e363df3fe1 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py @@ -245,10 +245,8 @@ def test_decode_layers_are_not_held_to_context_kernel_head_sizes(head_dim: int) dtype=torch.bfloat16, num_heads=16, num_kv_heads=2, head_dim=head_dim, tokens_per_block=32 ) with patch("tensorrt_llm._torch.speculative.dflash_attention.get_dflash_trtllm_gen_ops"): - validate_dflash_trtllm_gen_runtime(**shape, has_context_attention=False) - validate_dflash_trtllm_gen_runtime(**shape, has_context_attention=False) - with pytest.raises(RuntimeError, match="head dimension"): - validate_dflash_trtllm_gen_runtime(**shape, has_context_attention=True) + validate_dflash_trtllm_gen_runtime(**shape) + validate_dflash_trtllm_gen_runtime(**shape, folded_block_size=4) # --------------------------------------------------------------------------- @@ -635,7 +633,6 @@ def test_folded_runtime_accepts_large_blocks_by_chunking( num_kv_heads=kv_heads, head_dim=128, tokens_per_block=32, - has_context_attention=False, folded_block_size=block, ) chunk = get_dflash_folded_decode_chunk_size(block, heads, kv_heads) From 54ccfb2fac9cca5c0365b5d133534daae112abe1 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Mon, 9 Nov 2026 19:23:57 -0800 Subject: [PATCH 09/19] [None][chore] Revert no-op Helix V-up flag parametrization and autotuner rewrite TRTLLM_HELIX_V_UP_BEFORE_A2A is not read anywhere, so the parametrization only reran the same cases; the autotuner tactics rewrite changed nothing. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/autotuner.py | 3 +-- .../attention/multi_gpu/test_mla_helix.py | 20 ++----------------- 2 files changed, 3 insertions(+), 20 deletions(-) diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index da35c1f3ad68..9ef87cc93d04 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -2239,11 +2239,10 @@ def _maybe_parallelize_tactics(self, # Each rank of the tuning group only tunes its share of tactics; # _merge_cache_data then keeps the fastest across the group. group_size, group_index = self._tuning_group() - tactics = [ + return [ tactic for idx, tactic in enumerate(all_valid_tactics) if (offset + idx) % group_size == group_index ] - return tactics else: return all_valid_tactics diff --git a/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py b/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py index 088ce147e2cf..f8304e4ffc8b 100644 --- a/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py +++ b/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py @@ -643,18 +643,10 @@ def _v_up_builder(kv_b_weight, kv_b_scale, *, num_heads, cp_size, cp_rank, nope, return module -@pytest.mark.parametrize("obsolete_flag", [None, "0", "1"]) @pytest.mark.parametrize("cp_rank", [0, 1]) -def test_helix_v_up_weight_from_fp8_kv_b( - cp_rank: int, obsolete_flag: str | None, monkeypatch: pytest.MonkeyPatch -) -> None: +def test_helix_v_up_weight_from_fp8_kv_b(cp_rank: int): """MLA dequantizes every tp-local head's V weight from the loaded FP8 kv_b_proj; this rank's dequantized V chunk becomes a view of it.""" - # K3 leaves loaded FP8 block scales for the MLA post-load walk to bind. - if obsolete_flag is None: - monkeypatch.delenv("TRTLLM_HELIX_V_UP_BEFORE_A2A", raising=False) - else: - monkeypatch.setenv("TRTLLM_HELIX_V_UP_BEFORE_A2A", obsolete_flag) h, n, v, k, cp = 4, 128, 128, 256, 2 generator = torch.Generator().manual_seed(cp_rank) weight = (torch.randn(h * (n + v), k, generator=generator) * 64).to(torch.float8_e4m3fn) @@ -679,15 +671,7 @@ def test_helix_v_up_weight_from_fp8_kv_b( ) -@pytest.mark.parametrize("obsolete_flag", [None, "0", "1"]) -def test_helix_v_up_weight_views_bf16_kv_b( - obsolete_flag: str | None, monkeypatch: pytest.MonkeyPatch -) -> None: - # The common MLA binding owns the BF16 K3 layout as well. - if obsolete_flag is None: - monkeypatch.delenv("TRTLLM_HELIX_V_UP_BEFORE_A2A", raising=False) - else: - monkeypatch.setenv("TRTLLM_HELIX_V_UP_BEFORE_A2A", obsolete_flag) +def test_helix_v_up_weight_views_bf16_kv_b(): h, n, v, k = 4, 64, 128, 256 weight = torch.randn(h * (n + v), k, dtype=torch.bfloat16) module = _v_up_builder(weight, None, num_heads=h, cp_size=2, cp_rank=0, nope=n, v=v, k=k) From af72cfac555130555ac598f09ed8d695c6e7fa00 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Mon, 9 Nov 2026 19:25:23 -0800 Subject: [PATCH 10/19] [None][chore] CuteDSL FC12 MoE: drop the unreachable input-ready event join FC12 has no quantize_input_async, so the scheduler never hands it an x_ready_event; remove the plumbing and the event cases of its test. Signed-off-by: Zhenhuan Chen --- .../_torch/moe/fused_moe/cutedsl/fused_fc12.py | 15 --------------- tests/unittest/_torch/moe/test_moe_backend.py | 15 +++------------ 2 files changed, 3 insertions(+), 27 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/cutedsl/fused_fc12.py b/tensorrt_llm/_torch/moe/fused_moe/cutedsl/fused_fc12.py index cf1d8d682752..afc97716f4e3 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/cutedsl/fused_fc12.py +++ b/tensorrt_llm/_torch/moe/fused_moe/cutedsl/fused_fc12.py @@ -311,13 +311,6 @@ def _build_local_weight_view(self) -> NvFp4WeightView: quantize_input = quantize_input_nvfp4 prequantized_nvfp4_input_scale = prequantized_nvfp4_input_scale - def _join_pending_x_ready(self) -> None: - """Wait for the input-ready event handed to ``run_moe_nvfp4`` (once per forward).""" - event = getattr(self, "_pending_x_ready_event", None) - if event is not None: - event.wait() - self._pending_x_ready_event = None - def run_moe(self, ctx: MoERunContext, *, workspace: dict | None = None) -> torch.Tensor: del workspace # the fused kernel allocates its own intermediates plan = require_comm_plan(self, ctx) @@ -329,7 +322,6 @@ def run_moe(self, ctx: MoERunContext, *, workspace: dict | None = None) -> torch moe_output=plan.moe_output, enable_alltoall=plan.enable_alltoall, weight_view=self._build_local_weight_view(), - x_ready_event=ctx.x_ready_event, sort_tables=ctx.sort_tables, ) @@ -344,7 +336,6 @@ def run_moe_nvfp4( moe_output: torch.Tensor | None = None, enable_alltoall: bool = False, weight_view: NvFp4WeightView | None = None, - x_ready_event: torch.cuda.Event | None = None, sort_tables: MoESortTables | None = None, ) -> torch.Tensor: assert self.has_nvfp4 @@ -359,13 +350,9 @@ def run_moe_nvfp4( assert moe_output.size() == (token_final_scales.size(0), self.hidden_size) assert moe_output.dtype == output_dtype - # Join before the fused kernel reads x, or before returning an empty batch. - self._pending_x_ready_event = x_ready_event - # Empty micro-batches: skip autotuning (synthetic grouped-GEMM inputs # require at least one output row). if token_selected_experts.size(0) == 0: - self._join_pending_x_ready() return moe_output effective_top_k = token_selected_experts.size(-1) @@ -470,8 +457,6 @@ def run_moe_nvfp4_impl( num_non_exiting_tiles, ) = sort_outputs - self._join_pending_x_ready() - # One fused op: gather + FC1 GEMM + gated act (SwiGLU or SiTU) + requant # + FC2 GEMM + finalize (scatter-add into moe_output = a2a combine # workspace). diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index b4560f665ec5..533e7ff3a625 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -6440,19 +6440,17 @@ def _runner(identity): @pytest.mark.cpu_only @pytest.mark.parametrize("num_tokens", [0, 4]) -@pytest.mark.parametrize("with_event", [False, True]) @pytest.mark.parametrize("provided_output", [False, True]) @pytest.mark.parametrize("sort_tile_size", [None, 128, 256]) @pytest.mark.parametrize("tuning_dispatch", [False, True]) -def test_cutedsl_fc12_input_ready_event( +def test_cutedsl_fc12_run_moe_sort_tables( monkeypatch: pytest.MonkeyPatch, num_tokens: int, - with_event: bool, provided_output: bool, sort_tile_size: int | None, tuning_dispatch: bool, ) -> None: - """Each forward joins before input consumption with cached or fresh sort tables.""" + """Each forward reuses matching sort tables and skips empty batches.""" backend = CuteDslFc12FusedMoE.__new__(CuteDslFc12FusedMoE) torch.nn.Module.__init__(backend) backend._weights_created = True @@ -6499,9 +6497,6 @@ def choose_one(name, runners, config, inputs): for _ in range(2): calls.clear() sort.reset_mock() - event = Mock() if with_event else None - if event is not None: - event.wait.side_effect = lambda: calls.append("wait") result = backend.run_moe( MoERunContext( x=torch.zeros((num_tokens, backend.hidden_size // 2), dtype=torch.uint8), @@ -6514,17 +6509,13 @@ def choose_one(name, runners, config, inputs): input_sf_swizzled=False, payload_in_workspace=False, ), - x_ready_event=event, sort_tables=sort_tables, ) ) kernel_calls = 1 + int(tuning_dispatch) if num_tokens else 0 - assert calls == (["wait"] if with_event else []) + ["kernel"] * kernel_calls + assert calls == ["kernel"] * kernel_calls expected_sorts = int(tuning_dispatch) + int(sort_tile_size != 128) if num_tokens else 0 assert sort.call_count == expected_sorts - if event is not None: - event.wait.assert_called_once_with() - assert getattr(backend, "_pending_x_ready_event", None) is None assert result.shape == output.shape assert result.dtype == torch.bfloat16 if provided_output: From b84ed541578d9e00f127e2a8d4006ff576ba1e32 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Mon, 9 Nov 2026 19:26:32 -0800 Subject: [PATCH 11/19] [None][fix] Rejection sampling: reject padding slots that alias the scratch row Bound padding slot ids by dummy_slot_row instead of the draft_probs row count, which includes the shared scratch row. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/speculative/interface.py | 8 +++----- .../test_rejection_buffers_guard.py | 19 +++++++++++++++++++ 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index d68ebfff9b8d..4f8e18a2b609 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -1213,16 +1213,14 @@ def write_padding_onehot_draft_probs(self, padding_slot_ids, draft_len): return onehot_vocab = (self.draft_probs_last_dim if self.draft_probs_last_dim > 0 else self.draft_probs_vocab_size) - bad = [ - s for s in padding_slot_ids - if s is None or s < 0 or s >= self.draft_probs.shape[0] - ] + # The last draft_probs row is the shared scratch row (dummy_slot_row). + bad = [s for s in padding_slot_ids if s < 0 or s >= self.dummy_slot_row] if bad: # Writing a different row would leave the real slot's distribution # stale and break lossless acceptance, so fail loudly instead. raise RuntimeError( f"rejection sampling: padding slot ids {bad[:8]} outside " - f"draft_probs rows [0, {self.draft_probs.shape[0]}) " + f"slot rows [0, {self.dummy_slot_row}) " f"(num_seq_slots={self.num_seq_slots}, " f"max_num_requests={self.max_num_requests}, " f"is_cuda_graph={self.is_cuda_graph})") diff --git a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py index 619181632773..e6a56984c131 100644 --- a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py +++ b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py @@ -51,6 +51,25 @@ def _alloc_meta(**over): return types.SimpleNamespace(**base) +def test_padding_onehot_rejects_the_scratch_row(): + m = _alloc_meta(draft_probs_last_dim=0, dummy_slot_row=0, is_cuda_graph=False) + SpecMetadata.prepare_rejection_sampling_buffers(m) + assert m.dummy_slot_row == R + m.draft_probs.fill_(0.5) + + SpecMetadata.write_padding_onehot_draft_probs(m, [0, R - 1], 2) + for slot in (0, R - 1): + assert torch.all(m.draft_probs[slot, :2, 0] == 1.0) + assert torch.all(m.draft_probs[slot, :2, 1:] == 0.0) + assert torch.all(m.draft_probs[slot, 2:] == 0.5) + assert torch.all(m.draft_probs[1 : R - 1] == 0.5) + + for bad_slot in (R, -1): + with pytest.raises(RuntimeError, match="padding slot ids"): + SpecMetadata.write_padding_onehot_draft_probs(m, [bad_slot], 2) + assert torch.all(m.draft_probs[R] == 0.5) + + def test_prepare_buffers_allocates_when_enabled(): # Shared draft/target vocab: draft_probs + batch_slot_ids are allocated; # full_draft_probs is skipped (only needed when the vocabs differ). From 3f1e1357c49fa23dfa9c84a5155bd6edd9153b4b Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Mon, 9 Nov 2026 19:27:34 -0800 Subject: [PATCH 12/19] [None][chore] FP4 MLA packed plan test: pin the default rule plan as well Parametrize the balance test over wave_plan so a change to the production default (wave_plan=False) fails a test. Signed-off-by: Zhenhuan Chen --- .../_torch/attention/test_fp4_mla_packed.py | 40 +++++++++++++------ 1 file changed, 28 insertions(+), 12 deletions(-) diff --git a/tests/unittest/_torch/attention/test_fp4_mla_packed.py b/tests/unittest/_torch/attention/test_fp4_mla_packed.py index 7f2c7a13e296..87b1203acf2f 100644 --- a/tests/unittest/_torch/attention/test_fp4_mla_packed.py +++ b/tests/unittest/_torch/attention/test_fp4_mla_packed.py @@ -66,36 +66,51 @@ def _plan_mod(): @pytest.mark.parametrize( - "tiles,capacity,wave_tasks,expected_tiles_per_task,expected_begins", + "tiles,capacity,wave_tasks,wave_plan,expected_tiles_per_task,expected_begins", [ # Even batch under one wave (helix16 B4, 6 m-tiles -> 17 wave tasks): one wave of 32-tile tasks. - ([128, 128, 128, 128], 24, 17, 32, [0, 4, 8, 12, 16]), + ([128, 128, 128, 128], 24, 17, True, 32, [0, 4, 8, 12, 16]), + ([128, 128, 128, 128], 24, 17, False, 32, [0, 4, 8, 12, 16]), # One long request among short ones: 16-tile tasks; the one task past the wave is a short tail. - ([1, 79, 3, 176], 24, 17, 16, [0, 1, 6, 7, 18]), + ([1, 79, 3, 176], 24, 17, True, 16, [0, 1, 6, 7, 18]), + # The rule takes most of the slots (two-wave target, fit by the fallback share). + ([1, 79, 3, 176], 24, 17, False, 13, [0, 1, 8, 9, 23]), # Full machine, longest 1.5x above the wave share: 100-tile tasks, the tails in a second wave. - ([128] * 15 + [200], 48, 17, 100, list(range(0, 34, 2))), + ([128] * 15 + [200], 48, 17, True, 100, list(range(0, 34, 2))), + # The rule's two-wave target also halves the short requests. + ([128] * 15 + [200], 48, 17, False, 67, list(range(0, 32, 2)) + [33]), # Full machine, longest within 1.5x of the wave share: no split at all. - ([128] * 15 + [130], 48, 17, 130, list(range(17))), + ([128] * 15 + [130], 48, 17, True, 130, list(range(17))), + ([128] * 15 + [130], 48, 17, False, 130, list(range(17))), # Strong skew in a full machine: eleven 94-tile tasks for the long request, short tails behind. - ([128] * 15 + [1024], 48, 17, 94, list(range(0, 32, 2)) + [41]), + ([128] * 15 + [1024], 48, 17, True, 94, list(range(0, 32, 2)) + [41]), + # The rule gives the long request twelve tasks. + ([128] * 15 + [1024], 48, 17, False, 86, list(range(0, 32, 2)) + [42]), # Strong skew in a small batch: 94-tile tasks, 11 for the long request, fill one wave. - ([128, 128, 128, 1024], 24, 17, 94, [0, 2, 4, 6, 17]), + ([128, 128, 128, 1024], 24, 17, True, 94, [0, 2, 4, 6, 17]), + # The rule's two-wave target, fit by the fallback share. + ([128, 128, 128, 1024], 24, 17, False, 71, [0, 2, 4, 6, 21]), # A single request fills one wave with 16 tasks; a short one is not worth splitting. - ([128], 16, 17, 8, [0, 16]), - ([8], 16, 17, 8, [0, 1]), + ([128], 16, 17, True, 8, [0, 16]), + ([8], 16, 17, True, 8, [0, 1]), + # The rule fills one wave too, or honours the minimum task length. + ([128], 16, 17, False, 8, [0, 16]), + ([8], 16, 17, False, 2, [0, 4]), # No spare slots: one task per request, however long. - ([5, 900], 2, 17, 900, list(range(3))), + ([5, 900], 2, 17, True, 900, list(range(3))), + ([5, 900], 2, 17, False, 900, list(range(3))), ], ) def test_plan_packed_tasks_balances_by_kv_length( tiles: list[int], capacity: int, wave_tasks: int, + wave_plan: bool, expected_tiles_per_task: int, expected_begins: list[int], ) -> None: tiles_per_task, begins = _plan_mod().plan_packed_tasks( - tiles, capacity, wave_tasks, wave_plan=True + tiles, capacity, wave_tasks, wave_plan=wave_plan ) assert tiles_per_task == expected_tiles_per_task assert begins == expected_begins @@ -152,7 +167,8 @@ def test_plan_task_order_runs_full_tasks_then_longest_tails() -> None: def test_plan_without_wave_plan_keeps_the_rule_and_request_order() -> None: reduce_mod = _plan_mod() plan = reduce_mod.plan_packed_tasks - # TRTLLM_FP4_MLA_PACKED_WAVE_PLAN=0: the rule's length, as before the wave simulation. + # The default plan (TRTLLM_FP4_MLA_PACKED_WAVE_PLAN=0): rule-based task length, request + # order, no wave simulation. assert plan([1, 79, 3, 176], 24, 17, wave_plan=False) == (13, [0, 1, 8, 9, 23]) assert plan([128] * 15 + [1024], 48, 17, wave_plan=False) == ( 86, From ef653dde7166b6c78028990fc8abc72e5cbbf7d3 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Mon, 9 Nov 2026 19:29:40 -0800 Subject: [PATCH 13/19] [None][chore] FP4 MLA mufu16: delete the dead score-slot row-sum protocol Remove the uncalled _runtime_t336_producer_tile with its rowsum helper, barrier ids, barrier inits and budget, and name the row-sum part buffer after its use. Signed-off-by: Zhenhuan Chen --- .../fp4_mla/fp4_mla_cutedsl_mufu16.py | 419 +----------------- 1 file changed, 10 insertions(+), 409 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py index a70c73d37fad..35d0ebf54cb1 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py @@ -612,7 +612,6 @@ class KvCache3DLayout: + NUM_QK_Q_PIPELINE_STAGE * 2 + SMEM_P4_QK_BARRIER_STAGES * 2 + SMEM_P4_V_PIPELINE_STAGES * 2 - + SMEM_P4_TMEM_SCORE_PIPELINE_STAGES * 2 + SMEM_P4_P4_CONTROL_MBAR_STAGES * 2 + SMEM_P4_PAIR_OVERLAP_MBAR_COUNT + SMEM_P4_QK_SCORE_RING_SLOTS @@ -923,8 +922,6 @@ def _qk_compact_ring_stage_byte_offset( # TMEM allocation is published to the producers and the PV issue warp. TMEM_BAR_THREADS = 32 * (2 + SMEM_P4_TOTAL_PRODUCER_WARPS) O_STORE_BAR_ID = 2 -SMEM_P4_P4_SLOT0_BAR_ID = 6 -SMEM_P4_P4_SLOT1_BAR_ID = 7 SMEM_P4_QK1_SPLIT_ARM_BAR_ID = 8 SMEM_P4_QK1_SPLIT_ARM_BAR_THREADS = 32 * 2 SMEM_P4_QSF_OWNER_BAR_ID = 9 @@ -3569,28 +3566,6 @@ def _arrive_smem_p4_mapped_leader_mbar(leader_mbar) -> None: _mbarrier_arrive_release_cta_shared_cluster(leader_mbar) -@cute.jit -def _store_smem_p4_n256_rowsum_part( - sRowSumParts, - warp_idx: ctm.Int32, - tidx: ctm.Int32, - producer_warp_base: ctm.Int32, - bank_idx: ctm.Int32, - score_half_idx: ctm.Int32, - value: ctm.Float32, -) -> None: - local_warp_idx = warp_idx - producer_warp_base - lane_idx = tidx & ctm.Int32(31) - row_band = local_warp_idx & ctm.Int32(1) - col_band = local_warp_idx >> ctm.Int32(1) - local_row = row_band * ctm.Int32(32) + lane_idx - slot_idx = bank_idx * ctm.Int32(SMEM_P4_SCORE_SLOTS_PER_KV_TILE) + score_half_idx - row_state_offset = slot_idx * ctm.Int32( - SMEM_P4_SCORE_SLOT_ROW_STATE_STRIDE - ) + col_band * ctm.Int32(SMEM_P4_BMM1_M) - sRowSumParts[row_state_offset + local_row] = value - - @cute.jit def _smem_p4_qk_score_group_col( score_slot_idx: ctm.Constexpr, col_band: ctm.Int32, local_group_idx: ctm.Constexpr @@ -5189,369 +5164,6 @@ def _runtime_t336_prepare_p_source( ) -@cute.jit -def _runtime_t336_producer_tile( - sAtomicRunningRowMax, - sP, - sPSF, - sScoreSlotRowMax, - sRowScaleRing, - qk_full_mbar, - pair_overlap_sync_mbars, - leader_pair_overlap_sync_mbars, - leader_pair_overlap_sync_mbars_dsmem, - p4_smem_slot0_producer, - p4_smem_slot0_mbar_ptr, - tmem_score_slot0_consumer, - tmem_score_slot0_mbar_ptr, - pv_full_mbar, - corr_o_acc_tmem, - warp_idx: ctm.Int32, - tidx: ctm.Int32, - p4_tmem_base_col_id: ctm.Int32, - p4_tmem_base_row_id: ctm.Int32, - valid_k_for_l: ctm.Int32, - softmax_scale_log2: ctm.Float32, - pv_psf_rescale: ctm.Float32, - running_row_max: ctm.Float32, - running_row_anchor: ctm.Float32, - carried_dq1024: ctm.Int32, - carried_old_scale: ctm.Float32, - stream_li_idx: ctm.Int32, - has_previous_tile: ctm.Constexpr, - tail_needs_mask: ctm.Boolean, - is_final_tile: ctm.Constexpr, - use_mixed_imlp: ctm.Constexpr, -) -> tuple: - bank_idx = stream_li_idx % ctm.Int32(SMEM_P4_TMEM_BANKS) - qk_phase = ( - stream_li_idx - // ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) - % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) - ) - rowsum_full_mbar = pair_overlap_sync_mbars.subview(SMEM_P4_PAIR_P4_READY_MBAR_OFFSET + bank_idx) - pair_rowmax_read_done_mbar = pair_overlap_sync_mbars.subview( - SMEM_P4_PAIR_ROWMAX_READ_DONE_MBAR_OFFSET + bank_idx - ) - direct_p_ready_mbar = pair_overlap_sync_mbars.subview( - SMEM_P4_PAIR_DIRECT_P_READY_MBAR_OFFSET + bank_idx - ) - leader_pair_p_source_ready_mbar = leader_pair_overlap_sync_mbars.subview( - SMEM_P4_PAIR_P_SOURCE_READY_MBAR_OFFSET + bank_idx - ) - leader_pair_p_source_ready_dsmem_mbar = leader_pair_overlap_sync_mbars_dsmem.subview( - SMEM_P4_PAIR_P_SOURCE_READY_MBAR_OFFSET + bank_idx - ) - qk_score_chunk_col_offset = ctm.Int32(SMEM_P4_QK_ACC_COL_OFFSET) + bank_idx * ctm.Int32( - SMEM_P4_QK_PIPELINE_SLOT_STRIDE - ) - p_stage_base = bank_idx * ctm.Int32(SMEM_P4_SCORE_SLOTS_PER_KV_TILE) - score_half_idx = warp_idx >> ctm.Int32(2) - producer_warp_base = score_half_idx << ctm.Int32(2) - producer_local_warp_idx = warp_idx & ctm.Int32(3) - producer_row_band = producer_local_warp_idx & ctm.Int32(1) - producer_col_band = producer_local_warp_idx >> ctm.Int32(1) - is_compute_warp = score_half_idx == ctm.Int32(0) - score_tmem_addresses = _prepare_p4_n256_score_half_tmem_addresses( - warp_idx, p4_tmem_base_col_id, p4_tmem_base_row_id, qk_score_chunk_col_offset - ) - next_running_row_max = running_row_max - next_running_row_anchor = running_row_anchor - next_carried_dq1024 = carried_dq1024 - next_carried_old_scale = carried_old_scale - if is_compute_warp: - tmem_score_slot0_consumer.wait() - if cutlass.const_expr(is_final_tile): - score_values, group_max_values, group_valid_values = ( - _runtime_load_final_p4_n256_score_half_from_tmem( - sAtomicRunningRowMax, - qk_full_mbar.subview(bank_idx), - score_tmem_addresses, - producer_col_band, - warp_idx, - tidx, - kv_tile_idx=stream_li_idx, - valid_k=valid_k_for_l, - tail_needs_mask=tail_needs_mask, - qk_phase=qk_phase, - producer_warp_base=producer_warp_base, - score_half_idx=score_half_idx, - ) - ) - else: - score_values, group_max_values, group_valid_values = _load_p4_n256_score_half_from_tmem( - sAtomicRunningRowMax, - qk_full_mbar.subview(bank_idx), - score_tmem_addresses, - producer_col_band, - warp_idx, - tidx, - kv_tile_idx=stream_li_idx, - valid_k=valid_k_for_l, - mask_valid_k=False, - qk_phase=qk_phase, - wait_for_qk_full=True, - producer_warp_base=producer_warp_base, - score_half_idx=score_half_idx, - ) - packed_p4_words = ctm.Vector.from_elements( - tuple((ctm.Int32(0) for _ in range(8))), dtype=ctm.Int32 - ) - group_p_sums = ctm.Vector.from_elements( - tuple((ctm.Float32(0.0) for _ in range(4))), dtype=ctm.Float32 - ) - if cutlass.const_expr(is_final_tile): - # ``tail_needs_mask`` is CTA-uniform. Keep the NaN-safe validity path - # for a genuinely partial final tile, but do not execute its per-SF16 - # predicates for a full final tile. Both cases remain in one runtime-KV - # binary; no request-dependent valid_k value enters the JIT key. - if tail_needs_mask: - packed_p4_words, group_p_sums = _prepack_p4_n256_score_half_before_rowmax( - score_values, - group_max_values, - group_valid_values, - warp_idx, - softmax_scale_log2, - kv_tile_idx=stream_li_idx, - valid_k=valid_k_for_l, - producer_warp_base=producer_warp_base, - score_half_idx=score_half_idx, - mask_valid_k=True, - use_mixed_imlp=use_mixed_imlp, - ) - else: - packed_p4_words, group_p_sums = _prepack_p4_n256_score_half_before_rowmax( - score_values, - group_max_values, - group_valid_values, - warp_idx, - softmax_scale_log2, - kv_tile_idx=stream_li_idx, - valid_k=valid_k_for_l, - producer_warp_base=producer_warp_base, - score_half_idx=score_half_idx, - mask_valid_k=False, - use_mixed_imlp=use_mixed_imlp, - ) - else: - packed_p4_words, group_p_sums = _prepack_p4_n256_score_half_before_rowmax( - score_values, - group_max_values, - group_valid_values, - warp_idx, - softmax_scale_log2, - kv_tile_idx=stream_li_idx, - valid_k=valid_k_for_l, - producer_warp_base=producer_warp_base, - score_half_idx=score_half_idx, - mask_valid_k=False, - use_mixed_imlp=use_mixed_imlp, - ) - if is_compute_warp: - p4_smem_slot0_producer.acquire() - prims.barrier( - barrier_id=SMEM_P4_DUAL_HALF_BAR_ID, - number_of_threads=SMEM_P4_DUAL_HALF_BAR_THREADS, - ) - new_row_max = running_row_max - new_row_valid = ctm.Boolean(True) - if cutlass.const_expr(is_final_tile): - if tail_needs_mask: - new_row_max, new_row_valid = _load_p4_n256_atomic_running_rowmax_with_validity( - sAtomicRunningRowMax, - warp_idx, - tidx, - producer_warp_base=producer_warp_base, - ) - else: - new_row_max = _load_p4_n256_atomic_running_rowmax( - sAtomicRunningRowMax, - warp_idx, - tidx, - producer_warp_base=producer_warp_base, - ) - else: - new_row_max = _load_p4_n256_atomic_running_rowmax( - sAtomicRunningRowMax, warp_idx, tidx, producer_warp_base=producer_warp_base - ) - if is_compute_warp: - if prims.elect_sync(): - prims.mbarrier_arrive(pair_rowmax_read_done_mbar) - lane_rebase = ctm.Boolean(False) - anchor_delta_log2 = ctm.Float32(0.0) - if cutlass.const_expr(is_final_tile): - if tail_needs_mask: - next_running_row_anchor, lane_rebase, anchor_delta_log2 = ( - _select_smem_p4_lazy_anchor_with_validity( - new_row_max, - new_row_valid, - running_row_anchor, - softmax_scale_log2, - ) - ) - else: - next_running_row_anchor, lane_rebase, anchor_delta_log2 = ( - _select_smem_p4_lazy_anchor_runtime( - new_row_max, - running_row_anchor, - softmax_scale_log2, - has_previous_tile, - ) - ) - _store_prepacked_p4_n256_score_half_to_direct_a( - sP, - packed_p4_words, - warp_idx, - tidx, - producer_warp_base=producer_warp_base, - p_stage_base=p_stage_base, - score_half_idx=score_half_idx, - ) - if prims.elect_sync(): - prims.mbarrier_arrive(direct_p_ready_mbar) - if cutlass.const_expr(not is_final_tile): - next_running_row_anchor, lane_rebase, anchor_delta_log2 = ( - _select_smem_p4_lazy_anchor_runtime( - new_row_max, running_row_anchor, softmax_scale_log2, has_previous_tile - ) - ) - current_old_scale = ctm.Float32(1.0) - warp_rebase = cute.arch.vote_any_sync(lane_rebase) - next_running_row_max = new_row_max - if cutlass.const_expr(is_final_tile): - owned_psf_word, prequant_row_sum = _runtime_finalize_final_p4_n256_score_half_psf_source( - sPSF, - group_max_values, - group_valid_values, - group_p_sums, - next_running_row_anchor, - warp_idx, - tidx, - softmax_scale_log2, - pv_psf_rescale, - tail_needs_mask=tail_needs_mask, - producer_warp_base=producer_warp_base, - score_half_idx=score_half_idx, - bank_idx=bank_idx, - use_mixed_imlp=use_mixed_imlp, - ) - else: - owned_psf_word, prequant_row_sum = _finalize_p4_n256_score_half_psf_source( - sPSF, - group_max_values, - group_p_sums, - next_running_row_anchor, - warp_idx, - tidx, - softmax_scale_log2, - pv_psf_rescale, - kv_tile_idx=stream_li_idx, - valid_k=valid_k_for_l, - mask_valid_k=False, - producer_warp_base=producer_warp_base, - score_half_idx=score_half_idx, - bank_idx=bank_idx, - use_mixed_imlp=use_mixed_imlp, - ) - if has_previous_tile: - common_source_ready_pred = prims.elect_sync() & ~warp_rebase - if common_source_ready_pred: - prims.mbarrier_arrive( - leader_pair_p_source_ready_dsmem_mbar, - count=1, - scope=prims.MemScope.CLUSTER, - relaxed=True, - ) - if warp_rebase: - pv_done_li_idx = stream_li_idx - ctm.Int32(1) - pv_done_slot = pv_done_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) - pv_done_phase = ( - pv_done_li_idx - // ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) - % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) - ) - while not prims.mbarrier_try_wait_parity( - pv_full_mbar.subview(pv_done_slot), pv_done_phase, time_limit=10000000 - ): - pass - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - current_old_scale = _compute_smem_p4_lazy_rebase_scale(lane_rebase, anchor_delta_log2) - _rescale_smem_p4_resident_o_in_tmem(corr_o_acc_tmem, current_old_scale) - if prims.elect_sync(): - _arrive_smem_p4_mapped_leader_mbar(leader_pair_p_source_ready_mbar) - elif prims.elect_sync(): - _arrive_smem_p4_mapped_leader_mbar(leader_pair_p_source_ready_mbar) - if is_compute_warp: - if cutlass.const_expr(not is_final_tile): - while not prims.mbarrier_try_wait_parity( - direct_p_ready_mbar, qk_phase, time_limit=10000000 - ): - pass - cute.arch.mbarrier_arrive( - p4_smem_slot0_mbar_ptr + p4_smem_slot0_producer.current_handle().index, - cutlass.Int32(0), - ) - p4_smem_slot0_producer.advance() - cute.arch.mbarrier_arrive( - tmem_score_slot0_mbar_ptr - + SMEM_P4_TMEM_SCORE_PIPELINE_STAGES - + tmem_score_slot0_consumer.current_handle().index, - cutlass.Int32(0), - ) - tmem_score_slot0_consumer.advance() - if is_compute_warp: - prims.barrier( - barrier_id=SMEM_P4_P4_SLOT0_BAR_ID, - number_of_threads=SMEM_P4_SCORE_SLOT_PRODUCER_THREADS, - ) - else: - prims.barrier( - barrier_id=SMEM_P4_P4_SLOT1_BAR_ID, - number_of_threads=SMEM_P4_CORRECTION_THREADS, - ) - row_sum = prequant_row_sum - rowsum_empty_mbar = pair_overlap_sync_mbars.subview( - SMEM_P4_PAIR_HALF1_ROWSUM_READY_MBAR_OFFSET + bank_idx - ) - if stream_li_idx >= ctm.Int32(SMEM_P4_TMEM_BANKS): - previous_rowsum_phase = ( - (stream_li_idx - ctm.Int32(SMEM_P4_TMEM_BANKS)) - // ctm.Int32(SMEM_P4_TMEM_BANKS) - % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) - ) - while not prims.mbarrier_try_wait_parity( - rowsum_empty_mbar, previous_rowsum_phase, time_limit=10000000 - ): - pass - _store_smem_p4_n256_rowsum_part( - sScoreSlotRowMax, - warp_idx, - tidx, - producer_warp_base=producer_warp_base, - bank_idx=bank_idx, - score_half_idx=score_half_idx, - value=row_sum, - ) - if is_compute_warp & (producer_col_band == ctm.Int32(0)): - row_state_local_row = producer_row_band * ctm.Int32(32) + (tidx & ctm.Int32(31)) - sRowScaleRing[bank_idx * ctm.Int32(SMEM_P4_BMM1_M) + row_state_local_row] = ( - current_old_scale - ) - if prims.elect_sync(): - prims.mbarrier_arrive(rowsum_full_mbar) - if not is_compute_warp: - while not prims.mbarrier_try_wait_parity( - pair_rowmax_read_done_mbar, qk_phase, time_limit=10000000 - ): - pass - return ( - next_running_row_max, - next_running_row_anchor, - next_carried_dq1024, - next_carried_old_scale, - ) - - @cute.jit def _issue_pipelined_score_loads( qk_full_mbar: ctm.Array, @@ -5709,9 +5321,10 @@ def _finalize_pipelined_tile( ) -> tuple: """Finish tile ``stream_li_idx`` from its prepacked P words. - Same steps as ``_runtime_t336_prepare_p_source`` from the dual-half barrier on followed by the - tail of ``_runtime_t336_producer_tile_v23``: row anchor, P and PSF stores, lazy O rescale and - the P_SOURCE_READY arrive. + Follows ``_runtime_t336_prepare_p_source`` from the dual-half barrier on and the tail of + ``_runtime_t336_producer_tile_v23`` (row anchor, P and PSF stores, lazy O rescale and the + P_SOURCE_READY arrive), except that it reads the parity-indexed ``sPipelinedRowMax``, never + masks K, and does not wait on ``pair_rowmax_read_done``. """ bank_idx = stream_li_idx % ctm.Int32(SMEM_P4_TMEM_BANKS) p_stage_base = bank_idx * ctm.Int32(SMEM_P4_SCORE_SLOTS_PER_KV_TILE) @@ -6481,7 +6094,7 @@ def _run_mla_decode_body( space=ctm.AddressSpace.smem, alignment=128, ) - sScoreSlotRowMax = ctm.Array( + sRowSumParts = ctm.Array( ctm.Float32, SMEM_P4_STAGE_ROW_STATE_SMEM_FLOATS, space=ctm.AddressSpace.smem, @@ -6580,18 +6193,6 @@ def _run_mla_decode_body( for pv_full_mbar_idx in ctm.range_constexpr(SMEM_P4_QK_PIPELINE_SLOTS): prims.mbarrier_init(pv_full_mbar.subview(pv_full_mbar_idx), 1) for pair_slot_idx in ctm.range_constexpr(SMEM_P4_QK_PIPELINE_SLOTS): - prims.mbarrier_init( - pair_overlap_sync_mbars.subview( - SMEM_P4_PAIR_P4_READY_MBAR_OFFSET + pair_slot_idx - ), - SMEM_P4_TOTAL_PRODUCER_WARPS, - ) - prims.mbarrier_init( - pair_overlap_sync_mbars.subview( - SMEM_P4_PAIR_HALF1_ROWSUM_READY_MBAR_OFFSET + pair_slot_idx - ), - 1, - ) prims.mbarrier_init( pair_overlap_sync_mbars.subview( SMEM_P4_PAIR_FINAL_PV_ISSUED_MBAR_OFFSET + pair_slot_idx @@ -7376,7 +6977,7 @@ def _run_mla_decode_body( part_local_warp_idx = warp_idx & ctm.Int32(3) part_idx = (warp_idx >> ctm.Int32(2)) * ctm.Int32(2) + (part_local_warp_idx >> ctm.Int32(1)) part_row = (part_local_warp_idx & ctm.Int32(1)) * ctm.Int32(32) + (tidx & ctm.Int32(31)) - sScoreSlotRowMax[part_idx * ctm.Int32(SMEM_P4_BMM1_M) + part_row] = running_part_sum + sRowSumParts[part_idx * ctm.Int32(SMEM_P4_BMM1_M) + part_row] = running_part_sum prims.barrier( barrier_id=SMEM_P4_DUAL_HALF_BAR_ID, number_of_threads=SMEM_P4_DUAL_HALF_BAR_THREADS, @@ -7384,11 +6985,11 @@ def _run_mla_decode_body( if is_correction_warp: final_row_state_local_row = corr_row_band * ctm.Int32(32) + (tidx & ctm.Int32(31)) final_anchor_row_sum = ( - sScoreSlotRowMax[final_row_state_local_row] - + sScoreSlotRowMax[ctm.Int32(2 * SMEM_P4_BMM1_M) + final_row_state_local_row] + sRowSumParts[final_row_state_local_row] + + sRowSumParts[ctm.Int32(2 * SMEM_P4_BMM1_M) + final_row_state_local_row] ) + ( - sScoreSlotRowMax[ctm.Int32(SMEM_P4_BMM1_M) + final_row_state_local_row] - + sScoreSlotRowMax[ctm.Int32(3 * SMEM_P4_BMM1_M) + final_row_state_local_row] + sRowSumParts[ctm.Int32(SMEM_P4_BMM1_M) + final_row_state_local_row] + + sRowSumParts[ctm.Int32(3 * SMEM_P4_BMM1_M) + final_row_state_local_row] ) # Packed rows may be empty while their tile is not (Helix tokens # whose KV lives on another rank); judge emptiness per row. From 793558aca25f02f8ad6dbf16c56c8ad3ffbd1f9e Mon Sep 17 00:00:00 2001 From: Xin Guan <294044352+xguannv@users.noreply.github.com> Date: Tue, 10 Nov 2026 14:32:06 +0000 Subject: [PATCH 14/19] [None][chore] DSpark attention op: one source for the shape contract, clearer rejections - Take the head dim, RoPE dim, window, draft-block storage and block sizes from dspark_rmsnorm_rope_custom_op instead of a second copy; only the attention-only head counts and page sizes stay here. - Fold the single-caller _run_paged_dspark_attention into fused_dsv4_dspark_attention. - Report "DSpark num_heads must be one of (64, 128), got N". - Check the pages' rank and head dim before reading page_size from them. - The preparation validator reports a non-CUDA input as such instead of "must be on cpu, got cpu". Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com> --- .../custom_ops/dspark_attention_custom_op.py | 60 ++++++------------- .../dspark_rmsnorm_rope_custom_op.py | 6 +- .../test_dspark_cute_dsl_attention.py | 2 +- 3 files changed, 25 insertions(+), 43 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py index a2c4a1fd27fd..7bbb15b52a9f 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py @@ -21,6 +21,11 @@ from ..kernels.blackwell.utils import make_ptr from ..utils import next_positive_power_of_2 from .dspark_rmsnorm_rope_custom_op import ( + _DSV4_DSPARK_BLOCK_SIZES, + _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE, + _DSV4_DSPARK_HEAD_DIM, + _DSV4_DSPARK_ROPE_DIM, + _DSV4_DSPARK_WINDOW_SIZE, _get_dspark_arch_str, cute_dsl_dspark_rmsnorm_rope_cache_write, cute_dsl_dspark_rmsnorm_rope_draft_block, @@ -44,12 +49,9 @@ def _gen_dspark_context_wkv_m_profiles(max_m: int) -> Tuple[int, ...]: return tuple(buckets) +# Attention-only limits; the shared DSpark shape contract (head and RoPE dims, window, +# draft block sizes and storage) comes from dspark_rmsnorm_rope_custom_op. _DSV4_DSPARK_SUPPORTED_NUM_HEADS = (64, 128) -_DSV4_DSPARK_HEAD_DIM = 512 -_DSV4_DSPARK_WINDOW_SIZE = 128 -_DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE = 8 -_DSV4_DSPARK_BLOCK_SIZES = (5, 6) -_DSV4_DSPARK_ROPE_DIM = 64 _DSV4_DSPARK_PAGE_SIZES = (16, 32, 64, 128, 256) @@ -135,7 +137,10 @@ def _compile_dspark_attention( return compiled -def _run_paged_dspark_attention( +@torch.library.custom_op( + "trtllm::fused_dsv4_dspark_attention", mutates_args=(), device_types="cuda" +) +def fused_dsv4_dspark_attention( q: torch.Tensor, draft_block: torch.Tensor, kv_pages: torch.Tensor, @@ -152,15 +157,21 @@ def _run_paged_dspark_attention( if arch_str is None: raise RuntimeError("Embedded DSpark attention requires SM100, SM103, or SM107") batch, block, heads, dim = q.shape - page_size = kv_pages.shape[1] if block not in _DSV4_DSPARK_BLOCK_SIZES: raise ValueError( f"DSpark block size must be one of {_DSV4_DSPARK_BLOCK_SIZES}, got {block}" ) if heads not in _DSV4_DSPARK_SUPPORTED_NUM_HEADS: - raise ValueError(f"DSpark requires {_DSV4_DSPARK_SUPPORTED_NUM_HEADS} heads, got {heads}") + raise ValueError( + f"DSpark num_heads must be one of {_DSV4_DSPARK_SUPPORTED_NUM_HEADS}, got {heads}" + ) if dim != _DSV4_DSPARK_HEAD_DIM: raise ValueError(f"DSpark head_dim must be {_DSV4_DSPARK_HEAD_DIM}, got {dim}") + if kv_pages.ndim != 3 or kv_pages.shape[2] != dim: + raise ValueError( + f"DSpark pages must have shape [pages, page_size, {dim}], got {tuple(kv_pages.shape)}" + ) + page_size = kv_pages.shape[1] if page_size not in _DSV4_DSPARK_PAGE_SIZES: raise ValueError( f"DSpark page size must be one of {_DSV4_DSPARK_PAGE_SIZES}, got {page_size}" @@ -173,10 +184,6 @@ def _run_paged_dspark_attention( raise ValueError( f"DSpark draft KV must have shape {expected_draft_shape}, got {tuple(draft_block.shape)}" ) - if kv_pages.ndim != 3 or kv_pages.shape[2] != dim: - raise ValueError( - f"DSpark pages must have shape [pages, page_size, {dim}], got {tuple(kv_pages.shape)}" - ) if block_tables.dtype != torch.int32 or block_tables.stride(1) != 1: raise ValueError( "DSpark page tables must have contiguous INT32 rows, got " @@ -265,35 +272,6 @@ def _run_paged_dspark_attention( return output -@torch.library.custom_op( - "trtllm::fused_dsv4_dspark_attention", mutates_args=(), device_types="cuda" -) -def fused_dsv4_dspark_attention( - q: torch.Tensor, - draft_block: torch.Tensor, - kv_pages: torch.Tensor, - block_tables: torch.Tensor, - positions: torch.Tensor, - valid_lengths: torch.Tensor, - capacities: torch.Tensor, - attn_sink: torch.Tensor, - inverse_rope_freqs: torch.Tensor, - softmax_scale: float, -) -> torch.Tensor: - return _run_paged_dspark_attention( - q, - draft_block, - kv_pages, - block_tables, - positions, - valid_lengths, - capacities, - attn_sink, - inverse_rope_freqs, - softmax_scale, - ) - - @torch.library.register_fake("trtllm::fused_dsv4_dspark_attention") def _( q: torch.Tensor, diff --git a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py index 3cf4f7aa0586..a055cb9ef1ff 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py @@ -644,7 +644,11 @@ def _validate_dspark_preparation( f"got {freqs.dtype} {tuple(freqs.shape)}" ) for name, tensor in (("input", x), ("weight", weight), ("frequencies", freqs)): - if not tensor.is_cuda or tensor.device != x.device: + if not tensor.is_cuda: + raise ValueError( + f"DSpark preparation {name} must be a CUDA tensor, got {tensor.device}" + ) + if tensor.device != x.device: raise ValueError( f"DSpark preparation {name} must be on {x.device}, got {tensor.device}" ) diff --git a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py index 9d9316e18f38..ff102085f943 100644 --- a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py +++ b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py @@ -184,7 +184,7 @@ def test_fused_dsv4_dspark_attention_support_gate_logs_rejection(monkeypatch, in [ ("q_dtype", "must use BF16"), ("block_size", "block size must be one of"), - ("heads", r"requires \(64, 128\) heads"), + ("heads", r"num_heads must be one of \(64, 128\), got 24"), ("page_size", "page size must be one of"), ("page_stride", "positive 16-byte-aligned page stride"), ("draft_shape", "draft KV must have shape"), From 1efaab71bb7fc13d298cf668554627cfce09beea Mon Sep 17 00:00:00 2001 From: Xin Guan <294044352+xguannv@users.noreply.github.com> Date: Tue, 10 Nov 2026 14:32:36 +0000 Subject: [PATCH 15/19] [None][chore] DSpark FP8 producers: reuse the dense GEMM UE8M0 helpers; test 448 * 2^k boundaries - The RMSNorm and concat FP8 producers use _round_up_ue8m0 and _u32_as_f32 from kernels/blackwell/dense_gemm_persistent.py instead of local copies of the same cvt.rp.satfinite.ue8m0x2.f32 conversion and exponent widening. Both copies keep the low byte of the packed result, so the scale bytes and FP8 values are unchanged; the quantization multiplier is still rcp_approx of the widened scale. - Rename the raw exponent byte from scale to scale_byte. - The producer test adds blocks whose amax is exactly 448 * 2^k (k = 1, 8, 40) and one of 1e3, for both producers, and checks the bytes (127 + k, and 129 for 1e3) as well as the bitwise match with fp8_quantize_1x128_packed_ue8m0. Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com> --- .../kernels/blackwell/dspark_rmsnorm_rope.py | 44 +++++------------ .../test_dspark_cute_dsl_attention.py | 47 +++++++++++++++++++ 2 files changed, 59 insertions(+), 32 deletions(-) diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark_rmsnorm_rope.py b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark_rmsnorm_rope.py index fe74fb02a69d..5a924affedf1 100644 --- a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark_rmsnorm_rope.py +++ b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark_rmsnorm_rope.py @@ -7,37 +7,13 @@ import cutlass import cutlass.cute as cute -from cutlass._mlir.dialects import llvm -from cutlass.cutlass_dsl import T, dsl_user_op try: from cuda.bindings import driver as cuda except ImportError: from cuda import cuda - -@dsl_user_op -def _cvt_rp_ue8m0(value, *, loc=None, ip=None) -> cutlass.Uint8: - """Round a positive FP32 value up to its raw UE8M0 exponent byte.""" - packed = llvm.inline_asm( - T.i16(), - [cutlass.Float32(value).ir_value(loc=loc, ip=ip)], - "cvt.rp.satfinite.ue8m0x2.f32 $0, 0f00000000, $1;", - "=h,f", - has_side_effects=False, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - return cutlass.Uint8(cutlass.Int16(packed) & cutlass.Int16(0xFF)) - - -@dsl_user_op -def _ue8m0_to_f32(exponent, *, loc=None, ip=None) -> cutlass.Float32: - """Widen a nonzero UE8M0 exponent byte to its FP32 power of two.""" - bits = cutlass.Uint32(exponent) << cutlass.Uint32(23) - return cutlass.Float32(llvm.bitcast(T.f32(), bits.ir_value(loc=loc, ip=ip))) +from tensorrt_llm._torch.kernels.blackwell.dense_gemm_persistent import _round_up_ue8m0, _u32_as_f32 class DSparkRMSNormRoPEKernel: @@ -437,8 +413,10 @@ def kernel( block_absmax = cute.arch.warp_reduction_max(local_absmax) # The 1e-10 floor keeps the exponent byte nonzero (normal FP32 scale). block_absmax = cute.arch.fmax(block_absmax, cutlass.Float32(1e-10)) - scale = _cvt_rp_ue8m0(block_absmax * cutlass.Float32(1.0 / 448.0)) - quant_scale = cute.arch.rcp_approx(_ue8m0_to_f32(scale)) + scale_byte = _round_up_ue8m0(block_absmax * cutlass.Float32(1.0 / 448.0)) + quant_scale = cute.arch.rcp_approx( + _u32_as_f32(cutlass.Uint32(scale_byte) << cutlass.Uint32(23)) + ) for item in cutlass.range_constexpr(self.values_per_thread): dim = block_base + tidx * self.values_per_thread + item @@ -454,7 +432,7 @@ def kernel( + ((row % 128) // 32) * 4 ) for replica in cutlass.range_constexpr(4): - output_scale[scale_offset + replica] = scale + output_scale[scale_offset + replica] = scale_byte class DSparkConcatFP8QuantKernel: @@ -503,7 +481,7 @@ def kernel( lane = tidx % cute.arch.WARP_SIZE row = row_block * self.num_warps + warp total_rows = first.shape[0] + second.shape[0] - scale = cutlass.Uint8(0) + scale_byte = cutlass.Uint8(0) if cutlass.dynamic_expr(row < total_rows): values = cute.make_rmem_tensor((self.values_per_thread,), cutlass.BFloat16) @@ -534,8 +512,10 @@ def kernel( block_absmax = cute.arch.warp_reduction_max(local_absmax) # The 1e-10 floor keeps the exponent byte nonzero (normal FP32 scale). block_absmax = cute.arch.fmax(block_absmax, cutlass.Float32(1e-10)) - scale = _cvt_rp_ue8m0(block_absmax * cutlass.Float32(1.0 / 448.0)) - quant_scale = cute.arch.rcp_approx(_ue8m0_to_f32(scale)) + scale_byte = _round_up_ue8m0(block_absmax * cutlass.Float32(1.0 / 448.0)) + quant_scale = cute.arch.rcp_approx( + _u32_as_f32(cutlass.Uint32(scale_byte) << cutlass.Uint32(23)) + ) for item in cutlass.range_constexpr(self.values_per_thread): dim = block_base + lane * self.values_per_thread + item output[row, dim] = (cutlass.Float32(values[item]) * quant_scale).to( @@ -550,4 +530,4 @@ def kernel( + ((row % 128) // 32) * 4 ) for replica in cutlass.range_constexpr(4): - output_scale[scale_offset + replica] = scale + output_scale[scale_offset + replica] = scale_byte diff --git a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py index ff102085f943..ef52b38b127b 100644 --- a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py +++ b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py @@ -627,6 +627,53 @@ def active_scales(scale: torch.Tensor, rows: int) -> torch.Tensor: atol=0, ) + # A block amax of exactly 448 * 2**k keeps the scale 2**k (byte 127 + k); 1e3 rounds up to + # 4 (byte 129). One block per amax, the rest of each block at most half of it. + boundary_amax = (448.0 * 2.0, 448.0 * 2.0**8, 448.0 * 2.0**40, 1.0e3) + boundary_bytes = torch.tensor([128, 135, 167, 129], dtype=torch.uint8) + boundary_blocks = torch.empty(len(boundary_amax), 128) + for index, amax in enumerate(boundary_amax): + boundary_blocks[index] = (torch.rand(128) - 0.5) * amax + boundary_blocks[index, 7 * index + 3] = amax if index % 2 else -amax + boundary_blocks = boundary_blocks.to(device="cuda", dtype=torch.bfloat16) + assert torch.equal( + boundary_blocks.float().abs().amax(dim=1), + torch.tensor(boundary_amax, device="cuda").to(torch.bfloat16).float(), + ) + rows = len(boundary_amax) + + # Concat producer: row r carries the boundary block in its quant block r. + boundary_input = torch.randn(rows, hidden_dim, device="cuda", dtype=torch.bfloat16) + for index in range(rows): + boundary_input[index, index * 128 : (index + 1) * 128] = boundary_blocks[index] + expected_fp8, expected_scale = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(boundary_input) + actual_fp8, actual_scale = cute_dsl_dspark_concat_fp8_quant( + boundary_input[:1], boundary_input[1:] + ) + torch.testing.assert_close(actual_fp8.float(), expected_fp8.float(), rtol=0, atol=0) + actual_active = active_scales(actual_scale, rows) + torch.testing.assert_close(actual_active, active_scales(expected_scale, rows), rtol=0, atol=0) + quant_blocks = hidden_dim // 128 + for index in range(rows): + offset = (index * quant_blocks + index) * 4 + assert torch.all(actual_active[offset : offset + 4].cpu() == boundary_bytes[index]) + + # RMSNorm producer: +-1 inputs have an RMS of 1, so the weight sets each block amax. + signs = torch.randint(0, 2, (rows, hidden_dim), device="cuda", dtype=torch.int32) + norm_input = (signs * 2 - 1).to(torch.bfloat16) + boundary_weight = torch.randn(hidden_dim, device="cuda", dtype=torch.bfloat16) + boundary_weight[: rows * 128] = boundary_blocks.abs().flatten() + expected_norm = _rmsnorm(norm_input, boundary_weight, eps) + expected_fp8, expected_scale = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(expected_norm) + actual_fp8, actual_scale = cute_dsl_dspark_rmsnorm_fp8_quant(norm_input, boundary_weight, eps) + torch.testing.assert_close(actual_fp8.float(), expected_fp8.float(), rtol=0, atol=0) + actual_active = active_scales(actual_scale, rows) + torch.testing.assert_close(actual_active, active_scales(expected_scale, rows), rtol=0, atol=0) + for row in range(rows): + for index in range(rows): + offset = (row * quant_blocks + index) * 4 + assert torch.all(actual_active[offset : offset + 4].cpu() == boundary_bytes[index]) + # Runtime row counts and row strides reuse the two initialization artifacts. assert _compile_dspark_concat_fp8_quant.cache_info().misses == 1 assert _compile_dspark_rmsnorm_fp8_quant.cache_info().misses == 1 From 1b38a68fbcad50c51d345439e5492ee2cba03217 Mon Sep 17 00:00:00 2001 From: Xin Guan <294044352+xguannv@users.noreply.github.com> Date: Tue, 10 Nov 2026 14:33:07 +0000 Subject: [PATCH 16/19] [None][fix] DSpark: post_load_weights(weights) is the only attention-weight entry point - DeepSeek-V4.1 binds its attention projections through dspark_model.post_load_weights(weights), the same call V4 makes; V4.1 always selects BF16, so this binds the same BF16 projections as before. - Delete cache_attn_weights_from_checkpoint (no caller) and cache_attn_weights_from_state_dict. _bind_bf16_attn_projections no longer sets _use_bf16_linear, so nothing can switch an explicit FP8 selection to BF16 after load. - The forward() errors for unbound attention weights point at load_weights. - The V4.1 fp8 rejection names V4.1 and says to unset TRTLLM_DSPARK_LINEAR_BACKEND or set it to bf16; _use_dspark_bf16_linear documents fp8_supported. - dspark_attention_forward documents every argument with shapes and dtypes. - Tests: the embedded loading test checks that V4 and V4.1 both call post_load_weights(weights) once; the BF16 binder test is parametrized over the 128 and 32 checkpoint block sizes (replacing the separate block-size test); the no-FP8-support selector test is parametrized per mapping and split into accept / reject cases; the new and changed tests are annotated. Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_dspark.py | 103 +++++++++++------- .../_torch/models/modeling_dspark_v41.py | 5 +- .../hw_agnostic/test_dspark_attention.py | 91 ++++++---------- ...test_dspark_embedded_confidence_loading.py | 20 ++-- .../hw_agnostic/test_dspark_v41.py | 2 +- 5 files changed, 117 insertions(+), 104 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 1605ad24291d..787307232902 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -522,12 +522,58 @@ def dspark_attention_forward( block's own positions, attention-sink softmax, inverse-RoPE on the output, and a grouped low-rank O projection (``wo_a`` einsum + ``wo_b``). - Projections accept tensor weights or loaded projection modules; ``wo_a`` is the raw grouped weight matrix - ``[n_groups * o_lora_rank, n_heads * head_dim // n_groups]``. ``kv_cache`` is - ``[b, window_size, head_dim]`` and is updated functionally (cloned). + Projections accept tensor weights or loaded projection modules. Activations + and weights are BF16 unless stated otherwise. + + Args: + x: ``[b, block, dim]`` block layer input. + main_x: ``[b, 1, dim]`` projected captured context. + start_pos: absolute decode position (> 0); the context K/V row is written + at ``start_pos % window_size``. + kv_cache: ``[b, window_size, head_dim]`` rolling captured-context window. + Cloned before the write unless ``persist`` is set. + wq_a: ``[q_lora_rank, dim]`` Q down-projection; may be None when + ``qkv_proj`` is given. + q_norm_w: ``[q_lora_rank]`` RMSNorm weight of the Q latent. + wq_b: ``[n_heads * head_dim, q_lora_rank]`` Q up-projection weight, or a + callable that applies it. + wkv: ``[head_dim, dim]`` MQA K/V projection shared by all heads; may be + None when ``qkv_proj`` is given. + kv_norm_w: ``[head_dim]`` RMSNorm weight of K/V. + wo_a: ``[n_groups * o_lora_rank, n_heads * head_dim // n_groups]`` raw + grouped O down-projection; may be None when ``o_a_proj`` is given. + wo_b: ``[dim, n_groups * o_lora_rank]`` O up-projection weight, or a + callable that applies it. + attn_sink: ``[n_heads]`` FP32 per-head sink logits. + n_heads: number of query heads. + head_dim: query/key/value head dimension (the MQA latent dimension). + rope_head_dim: number of trailing head dimensions that RoPE rotates. + n_groups: number of O-projection groups. + o_lora_rank: per-group rank of the O down-projection. + window_size: number of rows in the rolling context window. + eps: RMSNorm epsilon. + softmax_scale: attention logit scale. + freqs_cis: ``[maxlen, rope_head_dim // 2]`` complex plain-RoPE table; + must satisfy ``maxlen > start_pos + block``. + persist: write the context K/V row into ``kv_cache`` in place, for a + worker-owned window that carries across decode steps. + q_b_norm_enabled: apply the per-head RMS normalization after ``wq_b``. + qkv_proj: fused ``wq_a`` + ``wkv`` projection whose output is + ``[..., q_lora_rank + head_dim]``; replaces ``wq_a`` and ``wkv``. + o_a_proj: ``[n_groups, o_lora_rank, n_heads * head_dim // n_groups]`` + grouped O down-projection, BF16 or FP8 E4M3; replaces ``wo_a``. + o_a_scale: FP32 block scales of an FP8 ``o_a_proj``. + o_a_dequant: optional dequantized ``o_a_proj`` passed to the FP8 grouped + BMM. + use_fused_output_projection: run the fused FP8 output projection with an + FP8 ``o_a_proj``; None decides from the weight shapes. + use_fused_qkv_input_projection: quantize the ``qkv_proj`` input to FP8 in + one fused producer. + use_fused_q_input_projection: produce the FP8 ``wq_b`` input with the + fused RMSNorm + FP8 quantization kernel. Returns: - ``[b, block_size, dim]`` attention output (residual stream contribution). + ``[b, block, dim]`` attention output (residual stream contribution). """ assert start_pos > 0, "DSpark draft attention runs at generation (start_pos > 0)" b, block, _ = x.shape @@ -1407,31 +1453,6 @@ def _block_dequant( scale = scale.repeat_interleave(block, 1)[:, :input_features] return (weight * scale).to(torch.bfloat16) - def cache_attn_weights_from_checkpoint(self, ckpt_dir: str, weight_map: Dict[str, str]) -> None: - """Populate ``_dspark_attn`` by reading the ``mtp.{s}.attn.*`` tensors from the - checkpoint shards on disk, then dequantizing via :meth:`_bind_bf16_attn_projections`. - """ - from safetensors import safe_open - - prefixes = tuple(f"mtp.{s}.attn." for s in range(len(self.mtp_layers))) - shards: Dict[str, list] = {} - for k in weight_map: - if k.startswith(prefixes): - shards.setdefault(weight_map[k], []).append(k) - raw: Dict[str, torch.Tensor] = {} - for shard, ks in shards.items(): - with safe_open(os.path.join(ckpt_dir, shard), framework="pt", device="cpu") as f: - for k in ks: - raw[k] = f.get_tensor(k) - self._bind_bf16_attn_projections(raw) - - def cache_attn_weights_from_state_dict(self, weights: Dict) -> None: - """Populate ``_dspark_attn`` from an already-loaded in-memory ``weights`` dict - (no extra disk I/O). DeepSeek-V4.1 DSpark loads its BF16 attention projections - this way; delegates to :meth:`_bind_bf16_attn_projections`. - """ - self._bind_bf16_attn_projections(weights) - def _dspark_freqs_table(self, device: torch.device) -> torch.Tensor: """Return the plain-RoPE table cached for ``(device, cap)``. @@ -1909,8 +1930,8 @@ def forward( assert start_pos > 0, "DSpark draft runs at generation (start_pos > 0)" if getattr(self.mtp_layers[0], "_dspark_attn", None) is None: raise RuntimeError( - "DSpark attention weights not cached; call " - "cache_attn_weights_from_checkpoint(ckpt_dir, weight_map) after loading." + "DSpark attention weights are not bound; load the draft checkpoint with " + "load_weights(weights) before running the drafter." ) x, main_x, draft_ids = self.forward_embed(main_hidden, bonus_token_ids) main_x = main_x.unsqueeze(1) # [T, 1, hidden] for the MQA K/V projection @@ -1983,8 +2004,8 @@ def forward_batched( """ if getattr(self.mtp_layers[0], "_dspark_attn", None) is None: raise RuntimeError( - "DSpark attention weights not cached; call " - "cache_attn_weights_from_checkpoint(ckpt_dir, weight_map) after loading." + "DSpark attention weights are not bound; load the draft checkpoint with " + "load_weights(weights) before running the drafter." ) x, main_x, draft_ids = self.forward_embed(main_hidden, bonus_token_ids) main_x = main_x.unsqueeze(1) # [G, 1, hidden] for the MQA K/V projection @@ -2124,8 +2145,9 @@ def _bind_bf16_attn_projections(self, weights: Dict) -> None: Dequantizes the raw ``mtp.{s}.attn.*`` checkpoint tensors with the checkpoint's FP8 block size and disables every FP8 projection path. + Only :meth:`post_load_weights` calls this, and only when the BF16 + backend is selected. """ - self._use_bf16_linear = True for stage_id, stage in enumerate(self.mtp_layers): prefix = f"mtp.{stage_id}.attn." device = stage.input_layernorm.weight.device @@ -4474,7 +4496,13 @@ def _fused_fp8_output_projection_from_quantized( def _use_dspark_bf16_linear(mapping, fp8_supported: bool = True) -> bool: - """Keep BF16 outside Rubin unless FP8 is explicitly requested.""" + """Return whether the DSpark draft attention projections run in BF16. + + Without ``TRTLLM_DSPARK_LINEAR_BACKEND``, FP8 is used only on SM107 with + unsharded attention (TP=1 or attention DP). ``fp8_supported=False`` marks a + drafter that binds BF16 projections only (DeepSeek-V4.1); it always gets + BF16, and an explicit ``fp8`` request raises. + """ requested = os.environ.get(_DSPARK_LINEAR_BACKEND_ENV) sharded = mapping.tp_size > 1 and not mapping.enable_attention_dp if requested is None: @@ -4484,8 +4512,9 @@ def _use_dspark_bf16_linear(mapping, fp8_supported: bool = True) -> bool: raise ValueError(f"{_DSPARK_LINEAR_BACKEND_ENV} must be 'fp8' or 'bf16', got {backend!r}") if backend == "fp8" and not fp8_supported: raise ValueError( - f"{_DSPARK_LINEAR_BACKEND_ENV}=fp8 is not supported by this DSpark drafter; " - "it binds BF16 attention projections" + f"{_DSPARK_LINEAR_BACKEND_ENV}=fp8 is not supported by this DSpark drafter " + "(DeepSeek-V4.1 binds BF16 attention projections); unset " + f"{_DSPARK_LINEAR_BACKEND_ENV} or set it to bf16" ) if backend == "fp8" and sharded: raise ValueError(f"{_DSPARK_LINEAR_BACKEND_ENV}=fp8 requires TP=1 or attention DP") diff --git a/tensorrt_llm/_torch/models/modeling_dspark_v41.py b/tensorrt_llm/_torch/models/modeling_dspark_v41.py index 5ab86b3aa403..fec12082a08a 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark_v41.py +++ b/tensorrt_llm/_torch/models/modeling_dspark_v41.py @@ -144,7 +144,7 @@ class DSv41DSparkDraftModel(DSv4DSparkDraftModel): block_cls = DSv41DSparkBlock uses_lagged_pre = True checkpoint_fp8_block_size = 32 - # V4.1 loads attention through cache_attn_weights_from_state_dict (BF16). + # V4.1 binds BF16 attention projections; the FP8 backend is rejected. supports_fp8_attn_projections = False # Reference DSparkBlock.forward_head passes collapsed x to confidence, # while only the shared vocabulary projection consumes norm(x). @@ -305,8 +305,7 @@ def load_weights(self, weights: dict[str, torch.Tensor], weight_mapper=None, **k "DeepSeek-V4.1 DSpark weight load is incomplete: " f"unconsumed keys {unexpected[:8]}; unfed parameters {unfed[:8]}" ) - self.dspark_model.post_load_weights() - self.dspark_model.cache_attn_weights_from_state_dict(weights) + self.dspark_model.post_load_weights(weights) logger.info(f"[DSpark V4.1] loaded {self.num_stages} lagged-pre draft stages") diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py index ae9033b50061..0b42b23e6dbf 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py @@ -561,20 +561,32 @@ def test_use_dspark_bf16_linear_rejects_sharded_fp8(monkeypatch): modeling_dspark._use_dspark_bf16_linear(mapping) -@pytest.mark.parametrize("value,expected", [(None, True), ("bf16", True), ("fp8", None)]) -def test_use_dspark_bf16_linear_without_fp8_support(monkeypatch, value, expected): +_DSPARK_MAPPINGS = [(1, False), (2, True), (2, False)] + + +@pytest.mark.parametrize("tp_size,enable_attention_dp", _DSPARK_MAPPINGS) +@pytest.mark.parametrize("value", [None, "bf16"]) +def test_use_dspark_bf16_linear_without_fp8_support( + monkeypatch: pytest.MonkeyPatch, value: str | None, tp_size: int, enable_attention_dp: bool +) -> None: monkeypatch.setattr(modeling_dspark, "get_sm_version", lambda: 107) if value is None: monkeypatch.delenv("TRTLLM_DSPARK_LINEAR_BACKEND", raising=False) else: monkeypatch.setenv("TRTLLM_DSPARK_LINEAR_BACKEND", value) - for tp_size, enable_attention_dp in ((1, False), (2, True), (2, False)): - mapping = types.SimpleNamespace(tp_size=tp_size, enable_attention_dp=enable_attention_dp) - if expected is None: - with pytest.raises(ValueError, match="not supported by this DSpark drafter"): - modeling_dspark._use_dspark_bf16_linear(mapping, fp8_supported=False) - else: - assert modeling_dspark._use_dspark_bf16_linear(mapping, fp8_supported=False) + mapping = types.SimpleNamespace(tp_size=tp_size, enable_attention_dp=enable_attention_dp) + assert modeling_dspark._use_dspark_bf16_linear(mapping, fp8_supported=False) is True + + +@pytest.mark.parametrize("tp_size,enable_attention_dp", _DSPARK_MAPPINGS) +def test_use_dspark_bf16_linear_without_fp8_support_rejects_fp8( + monkeypatch: pytest.MonkeyPatch, tp_size: int, enable_attention_dp: bool +) -> None: + monkeypatch.setattr(modeling_dspark, "get_sm_version", lambda: 107) + monkeypatch.setenv("TRTLLM_DSPARK_LINEAR_BACKEND", "fp8") + mapping = types.SimpleNamespace(tp_size=tp_size, enable_attention_dp=enable_attention_dp) + with pytest.raises(ValueError, match="not supported by this DSpark drafter"): + modeling_dspark._use_dspark_bf16_linear(mapping, fp8_supported=False) def test_use_dspark_bf16_linear_rejects_invalid_backend(monkeypatch): @@ -1027,54 +1039,12 @@ def transform() -> None: assert events == ["transform", "bind"] -def test_bind_bf16_attn_projections_disables_every_fp8_projection(): - model = DSv4DSparkDraftModel.__new__(DSv4DSparkDraftModel) - torch.nn.Module.__init__(model) - model._context_wkv_fp8 = object() - stage = types.SimpleNamespace(input_layernorm=types.SimpleNamespace(weight=torch.ones(1))) - model.mtp_layers = [stage] - prefix = "mtp.0.attn." - shapes = { - "wq_a": (2, 4), - "wq_b": (6, 2), - "wkv": (3, 4), - "wo_a": (8, 3), - "wo_b": (4, 8), - } - weights = { - f"{prefix}q_norm.weight": torch.ones(2), - f"{prefix}kv_norm.weight": torch.ones(3), - f"{prefix}attn_sink": torch.zeros(6), - } - for index, (name, shape) in enumerate(shapes.items(), start=1): - weights[f"{prefix}{name}.weight"] = torch.full(shape, index, dtype=torch.float8_e4m3fn) - weights[f"{prefix}{name}.scale"] = torch.full((1, 1), index + 0.5) - - model._bind_bf16_attn_projections(weights) - - bound = stage._dspark_attn - for name in ("wq_a", "wq_b", "wkv", "wo_a", "wo_b"): - expected = model._block_dequant( - weights[f"{prefix}{name}.weight"], weights[f"{prefix}{name}.scale"] - ) - assert bound[name].dtype == torch.bfloat16 - torch.testing.assert_close(bound[name], expected) - assert bound["qkv_proj"] is None - assert bound["o_a_proj"] is None - assert bound["o_a_scale"] is None - assert bound["o_a_dequant"] is None - assert bound["use_fused_output_projection"] is False - assert bound["use_fused_qkv_input_projection"] is False - assert bound["use_fused_q_input_projection"] is False - assert model._context_wkv_fp8 is None - - @pytest.mark.parametrize("block", [128, 32]) -def test_bind_bf16_attn_projections_uses_checkpoint_block_size(block: int) -> None: +def test_bind_bf16_attn_projections_disables_every_fp8_projection(block: int) -> None: model = DSv4DSparkDraftModel.__new__(DSv4DSparkDraftModel) torch.nn.Module.__init__(model) model.checkpoint_fp8_block_size = block - model._use_bf16_linear = False + model._context_wkv_fp8 = object() stage = types.SimpleNamespace(input_layernorm=types.SimpleNamespace(weight=torch.ones(1))) model.mtp_layers = [stage] prefix = "mtp.0.attn." @@ -1094,13 +1064,22 @@ def test_bind_bf16_attn_projections_uses_checkpoint_block_size(block: int) -> No model._bind_bf16_attn_projections(weights) - assert model._use_bf16_linear is True - assert stage._dspark_context_wkv_slice is None + bound = stage._dspark_attn for name in ("wq_a", "wq_b", "wkv", "wo_a", "wo_b"): weight = weights[f"{prefix}{name}.weight"].float() scale = weights[f"{prefix}{name}.scale"] expected = weight * scale.repeat_interleave(block, 0).repeat_interleave(block, 1) - torch.testing.assert_close(stage._dspark_attn[name], expected.to(torch.bfloat16)) + assert bound[name].dtype == torch.bfloat16 + torch.testing.assert_close(bound[name], expected.to(torch.bfloat16)) + assert bound["qkv_proj"] is None + assert bound["o_a_proj"] is None + assert bound["o_a_scale"] is None + assert bound["o_a_dequant"] is None + assert bound["use_fused_output_projection"] is False + assert bound["use_fused_qkv_input_projection"] is False + assert bound["use_fused_q_input_projection"] is False + assert stage._dspark_context_wkv_slice is None + assert model._context_wkv_fp8 is None def test_context_wkv_fp8_quantizes_once_per_window_write(monkeypatch): diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_embedded_confidence_loading.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_embedded_confidence_loading.py index 89b9199050df..626825adb52a 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_embedded_confidence_loading.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_embedded_confidence_loading.py @@ -15,7 +15,9 @@ pytestmark = pytest.mark.cpu_only -def _wrapper(monkeypatch, v41, use_bf16_linear: bool = True): +def _wrapper( + monkeypatch: pytest.MonkeyPatch, v41: bool, use_bf16_linear: bool = True +) -> tuple[nn.Module, dict[str, torch.Tensor]]: module = modeling_dspark_v41 if v41 else modeling_dspark cls = module.DSv41DSparkForCausalLM if v41 else module.DSv4DSparkForCausalLM wrapper = object.__new__(cls) @@ -46,7 +48,6 @@ def _wrapper(monkeypatch, v41, use_bf16_linear: bool = True): ) ) model.post_load_weights = Mock() - model.cache_attn_weights_from_state_dict = Mock() wrapper.dspark_model = model wrapper.config = model.config wrapper.num_stages = 1 @@ -62,8 +63,8 @@ def _wrapper(monkeypatch, v41, use_bf16_linear: bool = True): @pytest.mark.parametrize("use_bf16_linear", [False, True]) @pytest.mark.parametrize("bias", [False, True]) def test_embedded_confidence_load_validates_and_copies_real_parameters( - monkeypatch, v41, bias, use_bf16_linear -): + monkeypatch: pytest.MonkeyPatch, v41: bool, bias: bool, use_bf16_linear: bool +) -> None: wrapper, weights = _wrapper(monkeypatch, v41, use_bf16_linear=use_bf16_linear) head = wrapper.dspark_model.mtp_layers[-1].confidence_head prefix = "mtp_layers.0.confidence_head." @@ -78,12 +79,15 @@ def test_embedded_confidence_load_validates_and_copies_real_parameters( assert (head.proj.bias is not None) is bias if bias: assert torch.equal(head.proj.bias, weights[prefix + "proj.bias"]) - wrapper.dspark_model.post_load_weights.assert_called_once() + # V4 and V4.1 both bind attention from the raw checkpoint in one call. + wrapper.dspark_model.post_load_weights.assert_called_once_with(weights) @pytest.mark.parametrize("v41", [False, True]) @pytest.mark.parametrize("fault", ["unknown", "missing_weight", "broadcast_weight"]) -def test_embedded_confidence_rejects_incomplete_or_unknown_tensors(monkeypatch, v41, fault): +def test_embedded_confidence_rejects_incomplete_or_unknown_tensors( + monkeypatch: pytest.MonkeyPatch, v41: bool, fault: str +) -> None: wrapper, weights = _wrapper(monkeypatch, v41) prefix = "mtp_layers.0.confidence_head." if fault == "unknown": @@ -99,7 +103,9 @@ def test_embedded_confidence_rejects_incomplete_or_unknown_tensors(monkeypatch, @pytest.mark.parametrize("v41", [False, True]) -def test_entirely_absent_confidence_head_remains_supported(monkeypatch, v41): +def test_entirely_absent_confidence_head_remains_supported( + monkeypatch: pytest.MonkeyPatch, v41: bool +) -> None: wrapper, weights = _wrapper(monkeypatch, v41) weights.pop("mtp_layers.0.confidence_head.proj.weight") wrapper.load_weights(weights) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py index db9df17e20e7..ff7c3bd154b1 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py @@ -253,7 +253,7 @@ def test_v41_block32_dequant_uses_each_tile_scale(): ) -def test_v41_draft_rejects_fp8_attention_projections(monkeypatch): +def test_v41_draft_rejects_fp8_attention_projections(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("TRTLLM_DSPARK_LINEAR_BACKEND", "fp8") assert DSv41DSparkDraftModel.supports_fp8_attn_projections is False mapping = SimpleNamespace(tp_size=1, enable_attention_dp=False) From a020441070e115035adc69a759718b035a828ed5 Mon Sep 17 00:00:00 2001 From: Xin Guan <294044352+xguannv@users.noreply.github.com> Date: Tue, 10 Nov 2026 14:33:35 +0000 Subject: [PATCH 17/19] [None][chore] Annotate the rejection-guard padding test and the Helix V-up tests Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com> --- tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py | 4 ++-- .../_torch/speculative/test_rejection_buffers_guard.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py b/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py index f8304e4ffc8b..8e1629f3d897 100644 --- a/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py +++ b/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py @@ -644,7 +644,7 @@ def _v_up_builder(kv_b_weight, kv_b_scale, *, num_heads, cp_size, cp_rank, nope, @pytest.mark.parametrize("cp_rank", [0, 1]) -def test_helix_v_up_weight_from_fp8_kv_b(cp_rank: int): +def test_helix_v_up_weight_from_fp8_kv_b(cp_rank: int) -> None: """MLA dequantizes every tp-local head's V weight from the loaded FP8 kv_b_proj; this rank's dequantized V chunk becomes a view of it.""" h, n, v, k, cp = 4, 128, 128, 256, 2 @@ -671,7 +671,7 @@ def test_helix_v_up_weight_from_fp8_kv_b(cp_rank: int): ) -def test_helix_v_up_weight_views_bf16_kv_b(): +def test_helix_v_up_weight_views_bf16_kv_b() -> None: h, n, v, k = 4, 64, 128, 256 weight = torch.randn(h * (n + v), k, dtype=torch.bfloat16) module = _v_up_builder(weight, None, num_heads=h, cp_size=2, cp_rank=0, nope=n, v=v, k=k) diff --git a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py index e6a56984c131..6481ca226d92 100644 --- a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py +++ b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py @@ -34,7 +34,7 @@ R, K, V = 8, 4, 32 # max_num_requests, max_draft_len, vocab_size -def _alloc_meta(**over): +def _alloc_meta(**over: object) -> types.SimpleNamespace: base = dict( use_rejection_sampling=True, draft_probs=None, @@ -51,7 +51,7 @@ def _alloc_meta(**over): return types.SimpleNamespace(**base) -def test_padding_onehot_rejects_the_scratch_row(): +def test_padding_onehot_rejects_the_scratch_row() -> None: m = _alloc_meta(draft_probs_last_dim=0, dummy_slot_row=0, is_cuda_graph=False) SpecMetadata.prepare_rejection_sampling_buffers(m) assert m.dummy_slot_row == R From 6e14cb3ed99b31435ab101a4d5d7370ff59efada Mon Sep 17 00:00:00 2001 From: Xin Guan <294044352+xguannv@users.noreply.github.com> Date: Tue, 10 Nov 2026 14:33:59 +0000 Subject: [PATCH 18/19] [None][chore] FP4 MLA packed plan test: drop the duplicated default-plan asserts The wave_plan=False task lengths are already rows of test_plan_packed_tasks_balances_by_kv_length; the default-plan test keeps only its request-order check. Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com> --- .../_torch/attention/test_fp4_mla_packed.py | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) diff --git a/tests/unittest/_torch/attention/test_fp4_mla_packed.py b/tests/unittest/_torch/attention/test_fp4_mla_packed.py index 87b1203acf2f..f2ce91688b4c 100644 --- a/tests/unittest/_torch/attention/test_fp4_mla_packed.py +++ b/tests/unittest/_torch/attention/test_fp4_mla_packed.py @@ -164,17 +164,10 @@ def test_plan_task_order_runs_full_tasks_then_longest_tails() -> None: assert order([5, 9, 5, 2], 10) == [(1, 0), (0, 0), (2, 0), (3, 0)] -def test_plan_without_wave_plan_keeps_the_rule_and_request_order() -> None: - reduce_mod = _plan_mod() - plan = reduce_mod.plan_packed_tasks - # The default plan (TRTLLM_FP4_MLA_PACKED_WAVE_PLAN=0): rule-based task length, request - # order, no wave simulation. - assert plan([1, 79, 3, 176], 24, 17, wave_plan=False) == (13, [0, 1, 8, 9, 23]) - assert plan([128] * 15 + [1024], 48, 17, wave_plan=False) == ( - 86, - list(range(0, 32, 2)) + [42], - ) - order = reduce_mod.plan_task_order([24, 30, 3, 14], 10, wave_plan=False) +def test_task_order_without_wave_plan_keeps_request_order() -> None: + # The default plan (TRTLLM_FP4_MLA_PACKED_WAVE_PLAN=0) keeps request order; its rule-based + # task lengths are the wave_plan=False rows of test_plan_packed_tasks_balances_by_kv_length. + order = _plan_mod().plan_task_order([24, 30, 3, 14], 10, wave_plan=False) assert order == [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2), (2, 0), (3, 0), (3, 1)] From 7385f850410c286e8e963fcb84c1ddb47c39c05f Mon Sep 17 00:00:00 2001 From: Xin Guan <294044352+xguannv@users.noreply.github.com> Date: Tue, 10 Nov 2026 14:34:33 +0000 Subject: [PATCH 19/19] [None][chore] FP4 MLA mufu16: delete code left unreferenced by the dead row-sum protocol Deletes, in fp4_mla_cutedsl_mufu16.py only: - _runtime_load_final_p4_n256_score_half_from_tmem and _runtime_finalize_final_p4_n256_score_half_psf_source, whose only caller was the removed _runtime_t336_producer_tile; - SMEM_P4_CORRECTION_THREADS; - the static SMEM accounting that nothing reads: SMEM_P4_STATIC_SMEM_BYTES (set through globals().update) and everything that only fed it, namely static_smem_bytes, SMEM_P4_CONTROL_AND_ALIGNMENT_BYTES, SMEM_P4_CONTROL_BYTES_RAW, SMEM_P4_CONTROL_MBAR_COUNT, SMEM_P4_PAGE_ID_PLAN_BYTES, SMEM_P4_RUNTIME_SCALE_PAIR_BYTES and the q/qsf/k/ksf compact byte locals in _refresh_sf_vec_config. Reference check (AST over tensorrt_llm/ and tests/: Name loads, attribute accesses, keyword names, import aliases and string constants): before deleting, none of these names had a read in this module outside the deleted code, and no other file reads them from this module. The only other occurrences are the separate definitions and uses inside fp4_mla_cutedsl_mufu16_fused_v_transpose.py, which does not import from this module and is unchanged. After deleting, the operands of the deleted expressions were checked again and all still have readers. The diff only deletes lines (221), and no remaining line reads a deleted name, so no remaining constant, SMEM layout field, barrier offset or barrier count changes value. Not changed here because they still feed the layout: the SMEM_P4_PAIR_P4_READY_MBAR_OFFSET and SMEM_P4_PAIR_HALF1_ROWSUM_READY_MBAR_OFFSET names (later barrier offsets are computed from them), and the sRowSumParts size (SMEM_P4_STAGE_ROW_STATE_SMEM_FLOATS, a shared-memory layout field). Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com> --- .../fp4_mla/fp4_mla_cutedsl_mufu16.py | 221 ------------------ 1 file changed, 221 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py index 35d0ebf54cb1..646e46750b20 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py @@ -261,11 +261,9 @@ def _initial_float_from_argv(option: str, default: float) -> float: (INT32_MAX - (KV_TILE - 1)) // SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT ) * SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT SMEM_P4_PAGE_ID_PLAN_INTS = SMEM_P4_PAGE_PLAN_PROFILE_KV // TRTLLM_PAGE_SIZE -SMEM_P4_PAGE_ID_PLAN_BYTES = SMEM_P4_PAGE_ID_PLAN_INTS * ctm.Int32.bytes # softmax scale, PV rescale, producer mask bound (Int32 bits), pad. SMEM_P4_RUNTIME_SCALE_PAIR_FLOATS = 4 SMEM_P4_RUNTIME_MASK_BOUND_WORD = 2 -SMEM_P4_RUNTIME_SCALE_PAIR_BYTES = SMEM_P4_RUNTIME_SCALE_PAIR_FLOATS * ctm.Float32.bytes SMEM_P4_LAZY_ANCHOR_REBASE_LOG2 = _initial_float_from_argv("--lazy-anchor-threshold-log2", 1.0) if not 0.0 <= SMEM_P4_LAZY_ANCHOR_REBASE_LOG2 <= 8.0: raise ValueError( @@ -487,7 +485,6 @@ class KvCache3DLayout: SMEM_P4_TOTAL_PRODUCER_WARPS = SMEM_P4_COMPUTE_WARPS + SMEM_P4_CORRECTION_WARPS SMEM_P4_STORE_WARPS = SMEM_P4_CORRECTION_WARPS SMEM_P4_COMPUTE_THREADS = 32 * SMEM_P4_COMPUTE_WARPS -SMEM_P4_CORRECTION_THREADS = 32 * SMEM_P4_CORRECTION_WARPS SMEM_P4_SCORE_SLOT_PRODUCER_THREADS = SMEM_P4_COMPUTE_THREADS THREADS_PER_CTA = 32 * 12 SMEM_P4_TMEM_ALLOC_EXCLUSIVE = True @@ -605,21 +602,6 @@ class KvCache3DLayout: SMEM_P4_V_DATA_STAGE_BYTES = B_STAGE_BYTES * SMEM_P4_N_OUT_TILES SMEM_P4_V_SFB_TMA_STAGE_BYTES = SFA_TMA_STAGE_BYTES * SMEM_P4_PV_SFB_N_CHUNKS * SMEM_P4_N_OUT_TILES SMEM_P4_V_TILE_TMA_BYTES = SMEM_P4_V_DATA_STAGE_BYTES + SMEM_P4_V_SFB_TMA_STAGE_BYTES -SMEM_P4_CONTROL_MBAR_COUNT = ( - SMEM_P4_QK_COMPLETION_MBARS - + SMEM_P4_QK_PIPELINE_SLOTS - + 1 - + NUM_QK_Q_PIPELINE_STAGE * 2 - + SMEM_P4_QK_BARRIER_STAGES * 2 - + SMEM_P4_V_PIPELINE_STAGES * 2 - + SMEM_P4_P4_CONTROL_MBAR_STAGES * 2 - + SMEM_P4_PAIR_OVERLAP_MBAR_COUNT - + SMEM_P4_QK_SCORE_RING_SLOTS - + 2 - + SMEM_P4_QK_BULK_READY_MBARS -) -SMEM_P4_CONTROL_BYTES_RAW = SMEM_P4_CONTROL_MBAR_COUNT * 8 + 8 -SMEM_P4_CONTROL_AND_ALIGNMENT_BYTES = (SMEM_P4_CONTROL_BYTES_RAW + 127) // 128 * 128 def _refresh_sf_vec_config(sf_vec_size: int) -> None: @@ -749,32 +731,6 @@ def _refresh_sf_vec_config(sf_vec_size: int) -> None: stage_rowsum_tmem_cols = SMEM_P4_TMEM_BANKS * SMEM_P4_SCORE_SLOTS_PER_KV_TILE rowmeta_tmem_cols = stage_rowmax_tmem_cols + stage_rowsum_tmem_cols + running_rowmeta_tmem_cols rowmeta_tmem_col_offset = SMEM_P4_SCALE_FEED_TMEM_COLS[1] - q_compact_tile_bytes = SMEM_P4_QK_FULL_DATA_STAGES * QK_A_STAGE_BYTES + QK_TAIL_A_STAGE_BYTES - qsf_compact_tile_bytes = ( - SMEM_P4_QK_FULL_DATA_STAGES * qk_sfa_stage_bytes + 512 * qk_tail_data_stage_sf_chunks - ) - k_compact_slot_bytes = SMEM_P4_QK_FULL_DATA_STAGES * QK_B_STAGE_BYTES + QK_TAIL_B_STAGE_BYTES - ksf_compact_slot_bytes = ( - SMEM_P4_QK_FULL_DATA_STAGES * qk_sfb_stage_bytes - + 512 * qk_k_tail_data_stage_sf_chunks * qk_sfb_n_chunks - ) - static_smem_bytes = ( - SMEM_P4_CONTROL_AND_ALIGNMENT_BYTES - + q_compact_tile_bytes - + SMEM_P4_QK_KTILE_RING_STAGES * k_compact_slot_bytes - + qsf_compact_tile_bytes - + SMEM_P4_QK_KTILE_RING_STAGES * ksf_compact_slot_bytes - + SMEM_P4_P4_SMEM_WORDS * 4 - + SMEM_P4_V_PIPELINE_STAGES * B_STAGE_BYTES * SMEM_P4_N_OUT_TILES - + SMEM_P4_P4_OWNER_SMEM_SF_MAILBOX_WORDS * 4 - + SMEM_P4_V_PIPELINE_STAGES * sfa_tma_stage_bytes * pv_sfb_n_chunks * SMEM_P4_N_OUT_TILES - + SMEM_P4_STAGE_ROW_STATE_SMEM_FLOATS * 4 - + SMEM_P4_ATOMIC_RUNNING_ROWMAX_SMEM_ROWS * 4 - # Row maxima of the pipelined leading tiles, by tile parity. - + 2 * SMEM_P4_ATOMIC_RUNNING_ROWMAX_SMEM_ROWS * 4 - + SMEM_P4_PAGE_ID_PLAN_BYTES - + SMEM_P4_RUNTIME_SCALE_PAIR_BYTES - ) globals().update( SF_VEC_SIZE=sf_vec_size, MMA_SCALE_CHUNK_K=mma_scale_chunk_k, @@ -899,7 +855,6 @@ def _refresh_sf_vec_config(sf_vec_size: int) -> None: SMEM_P4_QK_FULL_DATA_STAGES * qk_sfb_stage_bytes + 512 * qk_k_tail_data_stage_sf_chunks * qk_sfb_n_chunks ), - SMEM_P4_STATIC_SMEM_BYTES=static_smem_bytes, SMEM_P4_V_DATA_STAGE_BYTES=B_STAGE_BYTES * SMEM_P4_N_OUT_TILES, SMEM_P4_V_SFB_TMA_STAGE_BYTES=sfa_tma_stage_bytes * pv_sfb_n_chunks * SMEM_P4_N_OUT_TILES, SMEM_P4_V_TILE_TMA_BYTES=B_STAGE_BYTES * SMEM_P4_N_OUT_TILES @@ -3758,72 +3713,6 @@ def _load_p4_n256_score_half_from_tmem( ) -@cute.jit -def _runtime_load_final_p4_n256_score_half_from_tmem( - sAtomicRunningRowMax, - qk_full_mbar, - score_tmem_addresses, - producer_col_band: ctm.Int32, - warp_idx: ctm.Int32, - tidx: ctm.Int32, - kv_tile_idx: ctm.Int32, - valid_k: ctm.Int32, - tail_needs_mask: ctm.Boolean, - qk_phase: ctm.Int32, - producer_warp_base: ctm.Int32, - score_half_idx: ctm.Int32, -) -> tuple: - score_half_n: ctm.Constexpr = SMEM_P4_BMM1_N // SMEM_P4_N256_SCORE_HALVES - score_half_groups_per_col_band: ctm.Constexpr = ( - score_half_n // SF_VEC_SIZE // SMEM_P4_TMEM_WARP_N - ) - score_value_count: ctm.Constexpr = score_half_groups_per_col_band * SF_VEC_SIZE - score_values = ctm.Vector.from_elements( - tuple((ctm.Float32(0.0) for _ in range(score_value_count))), dtype=ctm.Float32 - ) - group_max_values = ctm.Vector.from_elements( - tuple((ctm.Float32(-ctm.Float32.inf) for _ in range(score_half_groups_per_col_band))), - dtype=ctm.Float32, - ) - group_valid_values = ctm.Vector.from_elements( - tuple((ctm.Boolean(False) for _ in range(score_half_groups_per_col_band))), - dtype=ctm.Boolean, - ) - if tail_needs_mask: - score_values, group_max_values, group_valid_values = _load_p4_n256_score_half_from_tmem( - sAtomicRunningRowMax, - qk_full_mbar, - score_tmem_addresses, - producer_col_band, - warp_idx, - tidx, - kv_tile_idx=kv_tile_idx, - valid_k=valid_k, - mask_valid_k=True, - qk_phase=qk_phase, - wait_for_qk_full=True, - producer_warp_base=producer_warp_base, - score_half_idx=score_half_idx, - ) - else: - score_values, group_max_values, group_valid_values = _load_p4_n256_score_half_from_tmem( - sAtomicRunningRowMax, - qk_full_mbar, - score_tmem_addresses, - producer_col_band, - warp_idx, - tidx, - kv_tile_idx=kv_tile_idx, - valid_k=valid_k, - mask_valid_k=False, - qk_phase=qk_phase, - wait_for_qk_full=True, - producer_warp_base=producer_warp_base, - score_half_idx=score_half_idx, - ) - return (score_values, group_max_values, group_valid_values) - - @cute.jit def _load_p4_n256_atomic_running_rowmax( sAtomicRunningRowMax, @@ -4173,116 +4062,6 @@ def _finalize_p4_n256_score_half_psf_source( return (denominator_sf_word, prequant_row_sum) -@cute.jit -def _runtime_finalize_final_p4_n256_score_half_psf_source( - sPSF, - group_max_values, - group_valid_values, - group_p_sums, - new_row_max: ctm.Float32, - warp_idx: ctm.Int32, - tidx: ctm.Int32, - softmax_scale_log2: ctm.Float32, - pv_psf_rescale: ctm.Float32, - tail_needs_mask: ctm.Boolean, - producer_warp_base: ctm.Int32, - score_half_idx: ctm.Int32, - bank_idx: ctm.Int32, - use_mixed_imlp: ctm.Constexpr, -) -> tuple: - local_warp_idx = warp_idx - producer_warp_base - lane_idx = tidx & ctm.Int32(31) - score_half_n: ctm.Constexpr = SMEM_P4_BMM1_N // SMEM_P4_N256_SCORE_HALVES - score_half_groups_per_col_band: ctm.Constexpr = ( - score_half_n // SF_VEC_SIZE // SMEM_P4_TMEM_WARP_N - ) - sf_bias = -new_row_max * softmax_scale_log2 - ctm.Float32(LOG2_6) - if cutlass.const_expr(score_half_groups_per_col_band != 4): - raise ValueError("packed P_SFA conversion requires four groups per col band") - denominator_sf_word = ctm.Int32(0) - pv_sf_word = ctm.Int32(0) - prequant_row_sum = ctm.Float32(0.0) - if tail_needs_mask: - sf_exact_0 = ctm.Float32(0.0) - sf_exact_1 = ctm.Float32(0.0) - sf_exact_2 = ctm.Float32(0.0) - sf_exact_3 = ctm.Float32(0.0) - for local_group_idx in ctm.range_constexpr(score_half_groups_per_col_band): - group_valid = group_valid_values[local_group_idx] - sf_exact = ctm.Float32(0.0) - if group_valid: - sf_exact = cute.exp2( - group_max_values[local_group_idx] * softmax_scale_log2 + sf_bias, - fastmath=True, - ) - prequant_row_sum += sf_exact * group_p_sums[local_group_idx] - if cutlass.const_expr(local_group_idx == 0): - sf_exact_0 = sf_exact - elif cutlass.const_expr(local_group_idx == 1): - sf_exact_1 = sf_exact - elif cutlass.const_expr(local_group_idx == 2): - sf_exact_2 = sf_exact - else: - sf_exact_3 = sf_exact - denominator_sf_word = _pack_e4m3x4_natural(sf_exact_0, sf_exact_1, sf_exact_2, sf_exact_3) - pv_sf_word = _pack_e4m3x4_natural( - sf_exact_0 * pv_psf_rescale, - sf_exact_1 * pv_psf_rescale, - sf_exact_2 * pv_psf_rescale, - sf_exact_3 * pv_psf_rescale, - ) - elif cutlass.const_expr(use_mixed_imlp): - sf_h2_01 = _softmax_exp2_pair_packed_f16x2( - group_max_values[0], group_max_values[1], softmax_scale_log2, sf_bias - ) - sf_h2_23 = _softmax_exp2_pair_packed_f16x2( - group_max_values[2], group_max_values[3], softmax_scale_log2, sf_bias - ) - denominator_sf_word = _pack_e4m3x4_natural_from_f16x2(sf_h2_01, sf_h2_23) - pv_sf_word = _scale_pack_e4m3x4_natural_from_f16x2(sf_h2_01, sf_h2_23, pv_psf_rescale) - prequant_row_sum = _p4_n256_prequant_rowsum_from_f16x2_psf(sf_h2_01, sf_h2_23, group_p_sums) - else: - sf_args_01 = nvvm_fma_packed_f32x2( - (group_max_values[0], group_max_values[1]), - (softmax_scale_log2, softmax_scale_log2), - (sf_bias, sf_bias), - ) - sf_args_23 = nvvm_fma_packed_f32x2( - (group_max_values[2], group_max_values[3]), - (softmax_scale_log2, softmax_scale_log2), - (sf_bias, sf_bias), - ) - sf_exact_0 = cute.exp2(sf_args_01[0], fastmath=True) - sf_exact_1 = cute.exp2(sf_args_01[1], fastmath=True) - sf_exact_2 = cute.exp2(sf_args_23[0], fastmath=True) - sf_exact_3 = cute.exp2(sf_args_23[1], fastmath=True) - denominator_sf_word = _pack_e4m3x4_natural(sf_exact_0, sf_exact_1, sf_exact_2, sf_exact_3) - pv_sf_word = _pack_e4m3x4_natural( - sf_exact_0 * pv_psf_rescale, - sf_exact_1 * pv_psf_rescale, - sf_exact_2 * pv_psf_rescale, - sf_exact_3 * pv_psf_rescale, - ) - row_sum_parts01 = nvvm_mul_packed_f32x2( - (sf_exact_0, sf_exact_1), (group_p_sums[0], group_p_sums[1]) - ) - row_sum_parts23 = nvvm_mul_packed_f32x2( - (sf_exact_2, sf_exact_3), (group_p_sums[2], group_p_sums[3]) - ) - prequant_row_sum = _sum_p4_n256_psf_row_parts( - row_sum_parts01[0], - row_sum_parts01[1], - row_sum_parts23[0], - row_sum_parts23[1], - ) - psf_source_word_offset = _smem_p4_n256_psf_source_word_offset( - bank_idx, score_half_idx, local_warp_idx, lane_idx - ) - sPSF[psf_source_word_offset] = pv_sf_word - prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) - return (denominator_sf_word, prequant_row_sum) - - @cute.jit def _reload_p4_n256_owned_psf_word_after_ready( sPSF,