[NVBUG-6241842][fix] DSA DSL atom-split: guard against MTP draft next…

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,16 +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}`. 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:
+                    if dsl_atom_split:
                         factor = metadata.dsl_expand_factor
                         eff_next_n = metadata.dsl_atom
                         exp_B = num_generations * factor
@@ -2180,16 +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.
+                    # 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:
+                    if dsl_atom_split:
                         factor = metadata.dsl_expand_factor
                         atom = metadata.dsl_atom
                         exp_B = num_generations * factor

tests/integration/test_lists/waives.txt

@@ -34,8 +34,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_chunked_prefill[latency] SKIP (https://nvbugs/6276981)
 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)