[Feature] Enable TRITON_ATTN for Batch Invariance

[frankwang28]

Purpose

This PR adds TRITON_ATTN support for batch invariance.

Related / parent issue: #27433

Test Plan

Run tests with and without the or is_batch_invariant check in the triton_unified_attention's unified_attention method.

Test Result

Tests are run on a B200 (do not have access to a Hopper GPU to validate there 🙁)

Without:

CUDA_VISIBLE_DEVICES=2 VLLM_TEST_SEED=12345 pytest tests/v1/determinism/test_batch_invariance.py::test_logprobs_bitwise_batch_invariance_bs1_vs_bsN[TRITON_ATTN] -s

...

FAILED tests/v1/determinism/test_batch_invariance.py::test_logprobs_bitwise_batch_invariance_bs1_vs_bsN[TRITON_ATTN] - Failed: Batch invariance violated in 128/128 prompts. See output above for details.
===================================================================== 1 failed, 3 warnings in 80.20s (0:01:20) =====================================================================

CUDA_VISIBLE_DEVICES=2 VLLM_TEST_SEED=12345 VLLM_TEST_MODEL=openai/gpt-oss-120b pytest tests/v1/determinism/test_batch_invariance.py::test_logprobs_bitwise_batch_invariance_bs1_vs_bsN[TRITON_ATTN] -s

...

FAILED tests/v1/determinism/test_batch_invariance.py::test_logprobs_bitwise_batch_invariance_bs1_vs_bsN[TRITON_ATTN] - Failed: Batch invariance violated in 128/128 prompts. See output above for details.
==================================================================== 1 failed, 3 warnings in 268.72s (0:04:28) =====================================================================

With:

CUDA_VISIBLE_DEVICES=2 VLLM_TEST_SEED=12345 pytest tests/v1/determinism/test_batch_invariance.py::test_logprobs_bitwise_batch_invariance_bs1_vs_bsN[TRITON_ATTN] -s

...

========================================================================== 1 passed, 3 warnings in 33.91s ==========================================================================

CUDA_VISIBLE_DEVICES=2 VLLM_TEST_SEED=12345 VLLM_TEST_MODEL=openai/gpt-oss-120b pytest tests/v1/determinism/test_batch_invariance.py::test_logprobs_bitwise_batch_invariance_bs1_vs_bsN[TRITON_ATTN] -s

...

===================================================================== 1 passed, 3 warnings in 65.82s (0:01:05) =====================================================================

Doing some more testing using my own test suite, Triton seems to also be decode invariant (prefilling part of a decoded sequence and decoding the rest of a sequence seems logprob identical).

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request effectively enables batch invariance for the TRITON_ATTN backend. The changes are well-structured and logical. By adding TRITON_ATTN to the list of decode-invariant backends and forcing the use of the deterministic 2D Triton kernel when batch invariance is enabled, the PR successfully addresses the non-determinism issue. The accompanying test updates ensure that this new capability is properly verified. The code is clean and the changes are correct. Excellent work!

[yewentao256]

LGTM, thanks for the work!
Since you have tested gpt oss, could you also added it in the doc? https://docs.vllm.ai/en/latest/features/batch_invariance/#tested-models

[mergify]

Documentation preview: https://vllm--33688.org.readthedocs.build/en/33688/