A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 22:45 UTC
[Feature] Enable `TRITON_ATTN` for Batch Invariance (#33688)
Signed-off-by: frankwang28 <frank.wbb@hotmail.com>
Frank Wang committed
Feb 3, 2026 at 21:27 UTC
45f8fd6f979de059e2d8a03ea2ddc40c06394931
4 files changed
+13
-4
docs/features/batch_invariance.md
+1
@@ -108,6 +108,7 @@ Batch invariance has been tested and verified on the following models:
108
- **Qwen3 (MoE)**: `Qwen/Qwen3-30B-A3B`, `Qwen/Qwen3-Next-80B-A3B-Instruct`
109
- **Qwen2.5**: `Qwen/Qwen2.5-0.5B-Instruct`, `Qwen/Qwen2.5-1.5B-Instruct`, `Qwen/Qwen2.5-3B-Instruct`, `Qwen/Qwen2.5-7B-Instruct`, `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-32B-Instruct`
110
- **Llama 3**: `meta-llama/Llama-3.1-8B-Instruct`, `meta-llama/Llama-3.2-1B-Instruct`
111
+- **GPT-OSS**: `openai/gpt-oss-20b`, `openai/gpt-oss-120b`
112
113
Other models may also work, but these have been explicitly validated. If you encounter issues with a specific model, please report them on the [GitHub issue tracker](https://github.com/vllm-project/vllm/issues/new/choose).
114
tests/v1/determinism/utils.py
+1
@@ -18,6 +18,7 @@ skip_unsupported = pytest.mark.skipif(
18
19
BACKENDS: list[str] = [
20
"FLASH_ATTN",
21
+ "TRITON_ATTN",
22
"TRITON_MLA",
23
]
24
vllm/model_executor/layers/batch_invariant.py
+6
-3
@@ -1003,8 +1003,11 @@ def vllm_is_batch_invariant() -> bool:
1003
def override_envs_for_invariance(
1004
attention_backend: AttentionBackendEnum | None,
1005
):
1006
- supported_backends = [
1006
+ decode_invariant_backends = [
1007
AttentionBackendEnum.FLASH_ATTN, # best supported backend
1008
+ AttentionBackendEnum.TRITON_ATTN,
1009
+ ]
1010
+ supported_backends = decode_invariant_backends + [
1011
# FlashInfer temporarily disabled due to invariant CTA sizes.
1012
# See FlashInfer issue #2424
1013
# AttentionBackendEnum.FLASHINFER,
@@ -1025,9 +1028,9 @@ def override_envs_for_invariance(
1028
"one of the supported backends before enabling batch_invariant."
1029
)
1030
raise RuntimeError(error)
1028
- if attention_backend != supported_backends[0]:
1031
+ if attention_backend not in decode_invariant_backends:
1032
warning = (
1030
- "You are using a decode-invariant form of batch invariance. "
1033
+ "You are using a non-decode-invariant form of batch invariance. "
1034
"This will not be invariant between prefill and decode."
1035
)
1036
logger.warning_once(warning, scope="local")
vllm/v1/attention/ops/triton_unified_attention.py
+5
-1
@@ -10,10 +10,12 @@
10
import torch
11
12
from vllm.logger import init_logger
13
+from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant
14
from vllm.platforms import current_platform
15
from vllm.triton_utils import tl, triton
16
17
logger = init_logger(__name__)
18
+is_batch_invariant = vllm_is_batch_invariant()
19
float8_info = torch.finfo(current_platform.fp8_dtype())
20
21
@@ -972,7 +974,8 @@ def unified_attention(
974
# Launch the 2D kernel if
975
# 1. No intermediate tiled softmax buffers for the 3D kernel have been allocated, or
976
# 2. The batch includes at least one prefill request, or
975
- # 3. The number of sequences exceeds the configured threshold
977
+ # 3. The number of sequences exceeds the configured threshold, or
978
+ # 4. Batch invariance is enabled
979
if (
980
seq_threshold_3D is None
981
or num_par_softmax_segments is None
@@ -981,6 +984,7 @@ def unified_attention(
984
or softmax_segm_expsum is None
985
or max_seqlen_q > 1
986
or num_seqs > seq_threshold_3D
987
+ or is_batch_invariant
988
):
989
kernel_unified_attention_2d[
990
(