diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 204b81d95a58..f5962e693752 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,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 diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index dfd056b822ab..42683d244dbb 100644 --- a/tests/integration/test_lists/waives.txt +++ b/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)