[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

Summary by CodeRabbit

Release Notes

• Performance

  • Refined sparse attention path selection logic during prefill operations with distributed tensor configurations. Chunk processing now consistently applies distributed synchronization when specific eligibility criteria are met.

• Tests

  • Re-enabled a previously skipped test case covering multi-GPU inference with chunked prefill, strengthening quality assurance for distributed configurations.

[coderabbitai]

Caution

Review failed

An error occurred during the review process. Please try again later.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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