[nvbugs/6276981][fix] DSA chunked-prefill indexer: always q-split when eligible

When the indexer chunked-prefill is gated by q_split_eligible (TP > 1, no
attention DP) but apply_q_split is False (chunk smaller than
q_split_threshold), every TP rank computes the full chunk's topk indices
independently via fp8_mqa_logits / fp8_fp4_mqa_logits. Those DeepGEMM
kernels are not bit-exact across launches, so per-rank topk indices
diverge for the same tokens. The downstream MLA attention then attends
to different KV positions on different ranks, corrupting KV-cache
writes. Short generations (MMLU's 2-token answers) hide it; long ones
(GSM8K's 256 tokens) compound it into garbage and 0% accuracy.

Force the q-split + allgather path whenever eligible: small chunks pay
a microscopic allgather instead of redundant per-rank logits compute,
and the per-token canonical owner from the slice + allgather erases
any rank-local nondeterminism before downstream layers read the
indices. q_split_threshold < 0 still fully disables eligibility.

Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>

Author: tensorrt-cicd

tensorrt_llm/_torch/attention_backend/sparse/dsa.py

@@ -2016,8 +2016,19 @@ def sparse_attn_indexer(
                         num_k_tokens, gather_head_dim)
 
                     chunk_num_token = chunk.token_end - chunk.token_start
-                    apply_q_split = q_split_eligible and chunk_num_token >= q_split_threshold
-                    if apply_q_split:
+                    # Always q-split when eligible: redundant per-rank
+                    # `fp8_mqa_logits` / `fp8_fp4_mqa_logits` are not bit-exact
+                    # across launches, so different TP ranks produce different
+                    # topk indices for the same tokens. The downstream MLA
+                    # attention then attends to divergent KV positions on
+                    # different ranks, corrupting KV-cache writes -- invisible
+                    # for short generations (e.g. MMLU's 2-token answers) but
+                    # accumulates into garbage over long ones (GSM8K's 256
+                    # tokens) -> 0% accuracy. The q-split + allgather path
+                    # canonicalizes ownership per-token, so any rank-local
+                    # nondeterminism is erased before downstream layers read
+                    # the indices. q_split_threshold < 0 still fully disables.
+                    if q_split_eligible:
                         chunk_q_start = chunk_num_token * tp_rank // tp_size
                         chunk_q_end = chunk_num_token * (tp_rank + 1) // tp_size
                     else:
@@ -2065,7 +2076,7 @@ def sparse_attn_indexer(
                             global_q_start:global_q_end, :topk_indices.
                             shape[-1]] = topk_indices.to(dtype=torch.int32)
 
-                    if apply_q_split:
+                    if q_split_eligible:
                         q_sizes = [(r + 1) * chunk_num_token // tp_size -
                                    r * chunk_num_token // tp_size
                                    for r in range(tp_size)]