[https://nvbugs/6276981][fix] Force the q-split + allgather code path whenever q_split_eligible=True (drop…

[tensorrt-cicd]

Summary

• Root cause: When q_split_eligible=True but apply_q_split=False (chunk smaller than threshold), each TP rank ran fp8_mqa_logits independently on the full chunk; the DeepGEMM kernel is not bit-exact across launches, so per-rank topk indices diverged and downstream MLA attention attended to different KV positions on different ranks, corrupting KV-cache writes.
• Fix: Force the q-split + allgather code path whenever q_split_eligible=True (drop the chunk_num_token >= q_split_threshold gate). The per-token canonical owner from slice + allgather erases per-rank nondeterminism.
• Automated fix generated by repair-bot

Test plan

• Verify fix on the same GPU type as the original failure
• Check for regressions in related tests

Links

• Bug: https://nvbugs/6276981

Dev Engineer Review

• Updated Indexer.sparse_attn_indexer so eligible prefill chunks always use q-splitting and allgather.
• Removed the chunk_num_token >= q_split_threshold condition.
• Preserved q_split_threshold < 0 as the disable condition.
• This gives all tensor-parallel ranks a canonical topk_indices_buffer and avoids rank-dependent DeepGEMM results.
• No public or exported declarations changed.
• No configuration or test-list changes were found.

QA Engineer Review

No test changes.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: e8cb3431-f135-4648-84ac-b267f9ccd724

📥 Commits

Reviewing files that changed from the base of the PR and between 8e588da and b3ee62f.

📒 Files selected for processing (1)

• tensorrt_llm/_torch/attention_backend/sparse/dsa.py

🚧 Files skipped from review as they are similar to previous changes (1)

• tensorrt_llm/_torch/attention_backend/sparse/dsa.py

---

Walkthrough

The sparse attention indexer now applies q-splitting to all eligible prefill chunks. It also gathers per-rank top-k indices for the same eligibility condition. The negative threshold remains a disabling condition.

Changes

Q-split logic update

Layer / File(s)	Summary
Unconditional q-split and top-k allgather tensorrt_llm/_torch/attention_backend/sparse/dsa.py	Eligible prefill chunks use q_split_eligible for per-rank query processing and topk_indices_buffer allgather. The chunk-size threshold no longer gates these operations.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: perkzzheng

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 inconclusive)

Check name	Status	Explanation	Resolution
Description check	❓ Inconclusive	The description explains the root cause, solution, and test plan. However, the PR checklist is incomplete—checkboxes are not marked and several required items lack confirmation.	Complete the PR checklist by checking items that apply and confirming alignment with CODING GUIDELINES, test coverage, and documentation requirements.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly references the NVBugs ID, specifies the fix type, and summarizes the main change: forcing q-split + allgather when q_split_eligible=True and dropping the threshold gate.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[trtllm-agent]

NVBug 6276981 is closed as Bug - Fix unknown. The linked bug appears resolved elsewhere or for a reason that does not prove this PR is redundant. This PR should be judged on its own merits; repair-bot is not auto-closing it.

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.