From 6053105587ea41a0a6c21a94b53c6b2daa9f26e7 Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Wed, 3 Jun 2026 04:57:42 +0000 Subject: [PATCH 1/3] [NVBUG-6241845][fix] DSA DSL atom-split: guard against MTP draft next_n=1 `_pick_dsl_expand` caches `(factor, atom)` on the metadata once per `prepare()` with the invariant `factor * atom == 1 + max_draft_tokens` (the target-pass next_n). MTPEagle reuses the same `attn_metadata` for its multi-step draft loop: after iteration `i=0` it mutates `seq_lens` to 1 so subsequent iterations run 1-token incremental decode with `next_n=1`. The DSL FP4/FP8 atom-split reshape `(num_gen, next_n, H, D) -> (num_gen*factor, atom, H, D)` is only valid when the caller actually supplies `next_n == factor * atom` tokens per request. For MTP=3 + FP4 + atom-split (`factor=2, atom=2`), the `i>=1` draft iteration's `q` only has `num_gen * 1` tokens, so the reshape crashed with `shape '[48, 2, 64, 64]' is invalid for input of size 98304` (4x mismatch). Add an explicit guard `next_n == factor * atom` to both FP4 and FP8 atom-split entries. Target / draft `i=0` satisfy the invariant and continue to benefit from atom-split; MTPEagle `i>=1` falls back to the kernel-native `next_n=1` path (atom-split has no benefit for 1 token per request anyway). Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../_torch/attention_backend/sparse/dsa.py | 30 ++++++++++++++----- 1 file changed, 22 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 48b2b2b3d8de..4d0155a99578 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -2155,12 +2155,20 @@ def sparse_attn_indexer( # DSL FP4 kernel natively supports next_n ∈ {1, 2, 3}. # The wave-aware picker in `_pick_dsl_expand` is run # once per metadata prepare and the result cached on - # `metadata.dsl_{expand_factor, atom}`. Trigger expand - # whenever the picker decided to split (factor > 1), - # regardless of next_n — this lets next_n ∈ {2, 3} also - # benefit from atom-split when low-batch leaves SMs idle, - # in addition to the mandatory next_n=4 case. - if metadata.dsl_expand_factor > 1: + # `metadata.dsl_{expand_factor, atom}` with the invariant + # `factor * atom == 1 + max_draft_tokens` (the target/ + # verify-time next_n). MTPEagle reuses the same metadata + # for the multi-step draft loop: after iteration i=0 it + # mutates `seq_lens` to 1 so subsequent iterations run + # 1-token incremental decode (next_n=1). The atom-split + # reshape `(num_gen, next_n, ...) -> (num_gen*factor, + # atom, ...)` is only valid when the caller actually + # supplies `next_n == factor * atom` tokens per request; + # gate on this so the i≥1 draft calls fall back to the + # kernel-native next_n=1 path (atom-split has no benefit + # for 1 token/request anyway — there's nothing to split). + if metadata.dsl_expand_factor > 1 and \ + next_n == metadata.dsl_expand_factor * metadata.dsl_atom: factor = metadata.dsl_expand_factor eff_next_n = metadata.dsl_atom exp_B = num_generations * factor @@ -2184,12 +2192,18 @@ def sparse_attn_indexer( # split (factor > 1) — typically benefits small-batch / # low-ntask configs by raising SM utilization at the cost # of factor× KV HBM re-reads. Picker decision was cached - # on metadata.{dsl_expand_factor, dsl_atom} during prepare. + # on metadata.{dsl_expand_factor, dsl_atom} during prepare + # with invariant `factor * atom == 1 + max_draft_tokens`. + # Same guard as the FP4 path above: skip the reshape when + # the caller's `next_n` doesn't match the picker's + # assumption (MTPEagle i≥1 draft iterations supply + # `next_n=1` after the post-i=0 metadata mutation). dsl_q = q_decode fp8_ctx_lens = dsl_context_lens fp8_block_table = block_table fp8_schedule_meta = metadata.scheduler_metadata_buffer - if metadata.dsl_expand_factor > 1: + if metadata.dsl_expand_factor > 1 and \ + next_n == metadata.dsl_expand_factor * metadata.dsl_atom: factor = metadata.dsl_expand_factor atom = metadata.dsl_atom exp_B = num_generations * factor From 7862937ff30fca8fd54903a42f8f0db2bcb485db Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Sun, 7 Jun 2026 08:33:16 +0000 Subject: [PATCH 2/3] refactor: extract dsl_atom_split variable for FP4/FP8 atom-split guard Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../_torch/attention_backend/sparse/dsa.py | 48 +++++++------------ 1 file changed, 18 insertions(+), 30 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 4d0155a99578..7c0b60782b2a 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -2138,6 +2138,19 @@ def sparse_attn_indexer( # avoids materializing a 2D contiguous tensor per call. dsl_context_lens = metadata.kv_lens_cuda_runtime[ num_contexts:num_contexts + num_generations] + # Wave-aware atom-split: the picker in `_pick_dsl_expand` caches + # (factor, atom) on metadata with invariant + # `factor * atom == 1 + max_draft_tokens` (the target/verify-time + # next_n). MTPEagle reuses the same metadata for its multi-step + # draft loop; after i=0 it mutates seq_lens to 1, so i≥1 + # iterations run with next_n=1. The reshape + # `(num_gen, next_n, ...) -> (num_gen*factor, atom, ...)` is only + # valid when the caller actually supplies next_n == factor * atom + # tokens; gate here so i≥1 draft calls fall back to the + # kernel-native next_n=1 path. + dsl_atom_split = (metadata.dsl_expand_factor > 1 + and next_n == metadata.dsl_expand_factor * + metadata.dsl_atom) if self.use_fp4: # FP4 DSL signature splits DG's (q, sf_q) tuple into two # separate args and requires q.dtype == uint8 (q_decode @@ -2151,24 +2164,7 @@ def sparse_attn_indexer( dsl_q = q_decode.view(torch.uint8) dsl_block_table = block_table dsl_schedule_meta = metadata.scheduler_metadata_buffer - - # DSL FP4 kernel natively supports next_n ∈ {1, 2, 3}. - # The wave-aware picker in `_pick_dsl_expand` is run - # once per metadata prepare and the result cached on - # `metadata.dsl_{expand_factor, atom}` with the invariant - # `factor * atom == 1 + max_draft_tokens` (the target/ - # verify-time next_n). MTPEagle reuses the same metadata - # for the multi-step draft loop: after iteration i=0 it - # mutates `seq_lens` to 1 so subsequent iterations run - # 1-token incremental decode (next_n=1). The atom-split - # reshape `(num_gen, next_n, ...) -> (num_gen*factor, - # atom, ...)` is only valid when the caller actually - # supplies `next_n == factor * atom` tokens per request; - # gate on this so the i≥1 draft calls fall back to the - # kernel-native next_n=1 path (atom-split has no benefit - # for 1 token/request anyway — there's nothing to split). - if metadata.dsl_expand_factor > 1 and \ - next_n == metadata.dsl_expand_factor * metadata.dsl_atom: + if dsl_atom_split: factor = metadata.dsl_expand_factor eff_next_n = metadata.dsl_atom exp_B = num_generations * factor @@ -2188,22 +2184,14 @@ def sparse_attn_indexer( max_seq_len) else: # FP8 DSL kernel natively supports next_n ∈ {1, 2, 3, 4}. - # Apply wave-aware atom-split when the picker decided to - # split (factor > 1) — typically benefits small-batch / - # low-ntask configs by raising SM utilization at the cost - # of factor× KV HBM re-reads. Picker decision was cached - # on metadata.{dsl_expand_factor, dsl_atom} during prepare - # with invariant `factor * atom == 1 + max_draft_tokens`. - # Same guard as the FP4 path above: skip the reshape when - # the caller's `next_n` doesn't match the picker's - # assumption (MTPEagle i≥1 draft iterations supply - # `next_n=1` after the post-i=0 metadata mutation). + # Atom-split benefits small-batch / low-ntask configs by + # raising SM utilization at the cost of factor× KV HBM + # re-reads; guard logic shared via dsl_atom_split above. dsl_q = q_decode fp8_ctx_lens = dsl_context_lens fp8_block_table = block_table fp8_schedule_meta = metadata.scheduler_metadata_buffer - if metadata.dsl_expand_factor > 1 and \ - next_n == metadata.dsl_expand_factor * metadata.dsl_atom: + if dsl_atom_split: factor = metadata.dsl_expand_factor atom = metadata.dsl_atom exp_B = num_generations * factor From 1434aef7b3ba79a0190c98d51d443024d542a98b Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Mon, 8 Jun 2026 00:52:03 +0000 Subject: [PATCH 3/3] [NVBUG-6241845][NVBUG-6241842][fix] Remove waives for fp4_indexer_dsl_mtp2/mtp3 tests Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 3866d63e2df4..145fc2221061 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -32,8 +32,6 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_fp8_blockscale[baseline] accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_fp8_blockscale[latency_default] SKIP (https://nvbugs/6185196) accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[baseline] SKIP (https://nvbugs/6185196) accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[baseline_mtp1] SKIP (https://nvbugs/6185196) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[fp4_indexer_dsl_mtp2] SKIP (https://nvbugs/6241842) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[fp4_indexer_dsl_mtp3] SKIP (https://nvbugs/6241845) accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_piecewise_cuda_graph[baseline] SKIP (https://nvbugs/6185196) accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_piecewise_cuda_graph[mtp3_fp8kv_chunked] SKIP (https://nvbugs/5989920) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6084720)