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 11, 20:00 UTC
[Feature] Batch invariant: Enable `TRITON_MLA` without prefix-caching (#29125)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Dec 8, 2025 at 19:31 UTC
d9417096d1347ead26b7c0cc1afb22fc0028e6e8
5 files changed
+43
-7
tests/v1/determinism/test_batch_invariance.py
+1
-5
@@ -185,7 +185,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
185
llm = LLM(
186
model=model_name,
187
tensor_parallel_size=tp_size,
188
- enable_prefix_caching=False,
188
+ # enable_prefix_caching=False,
189
max_num_seqs=32,
190
max_model_len=8192,
191
dtype="bfloat16", # not everything is supported
@@ -393,7 +393,6 @@ def test_simple_generation(backend, monkeypatch: pytest.MonkeyPatch):
393
gpu_memory_utilization=0.9,
394
max_model_len=2048,
395
dtype="bfloat16",
396
- enable_prefix_caching=False,
396
)
397
398
prompt = "the capital of france is"
@@ -457,7 +456,6 @@ def test_logprobs_without_batch_invariance_should_fail(
456
llm = LLM(
457
model=model_name,
458
tensor_parallel_size=tp_size,
460
- enable_prefix_caching=False,
459
max_num_seqs=32,
460
max_model_len=8192,
461
dtype="bfloat16",
@@ -681,7 +679,6 @@ def test_decode_logprobs_match_prefill_logprobs(
679
llm = LLM(
680
model=model_name,
681
tensor_parallel_size=tp_size,
684
- enable_prefix_caching=False,
682
max_num_seqs=32,
683
max_model_len=8192,
684
dtype="bfloat16",
@@ -928,7 +925,6 @@ def LLM_with_max_seqs(
925
max_model_len=max_model_len,
926
dtype="bfloat16",
927
tensor_parallel_size=int(os.getenv("VLLM_TP_SIZE", "1")),
931
- enable_prefix_caching=False,
928
# Enable for MOE models
929
# enable_expert_parallel=True,
930
)
tests/v1/determinism/test_online_batch_invariance.py
+4
-1
@@ -153,7 +153,10 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
153
}
154
155
tp_size = os.getenv("VLLM_TP_SIZE", "1")
156
- server_args: list[str] = []
156
+ server_args: list[str] = [
157
+ "--max-model-len=8192",
158
+ "--max-num-seqs=32",
159
+ ]
160
if tp_size:
161
server_args += ["-tp", tp_size]
162
tests/v1/determinism/utils.py
+1
@@ -17,6 +17,7 @@ skip_unsupported = pytest.mark.skipif(
17
18
BACKENDS: list[str] = [
19
"FLASH_ATTN",
20
+ "TRITON_MLA",
21
]
22
23
if has_flashinfer():
vllm/attention/layer.py
+36
@@ -25,6 +25,7 @@ from vllm.config.vllm import VllmConfig
25
from vllm.forward_context import ForwardContext, get_forward_context
26
from vllm.logger import init_logger
27
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
28
+from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant
29
from vllm.model_executor.layers.linear import (
30
ColumnParallelLinear,
31
UnquantizedLinearMethod,
@@ -251,6 +252,24 @@ class Attention(nn.Module, AttentionLayerBase):
252
else:
253
self.attn_backend = attn_backend
254
255
+ # prefix caching + batch invariance is currently not supported for
256
+ # FLASHINFER and TRITON_MLA.
257
+ if (
258
+ cache_config is not None
259
+ and cache_config.enable_prefix_caching
260
+ and vllm_is_batch_invariant()
261
+ and (
262
+ self.attn_backend.get_name() == "FLASHINFER"
263
+ or self.attn_backend.get_name() == "TRITON_MLA"
264
+ )
265
+ ):
266
+ logger.warning_once(
267
+ "Disabling prefix caching for FLASHINFER/TRITON_MLA "
268
+ "with batch invariance, as it is not yet supported.",
269
+ scope="local",
270
+ )
271
+ cache_config.enable_prefix_caching = False
272
+
273
impl_cls = self.attn_backend.get_impl_cls()
274
self.impl = impl_cls(
275
num_heads,
@@ -628,6 +647,23 @@ class MLAAttention(nn.Module, AttentionLayerBase):
647
use_mla=True,
648
use_sparse=use_sparse,
649
)
650
+
651
+ if (
652
+ cache_config is not None
653
+ and cache_config.enable_prefix_caching
654
+ and vllm_is_batch_invariant()
655
+ and (
656
+ self.attn_backend.get_name() == "TRITON_MLA"
657
+ or self.attn_backend.get_name() == "FLASHINFER"
658
+ )
659
+ ):
660
+ logger.warning_once(
661
+ "Disabling prefix caching for TRITON_MLA / FLASHINFER "
662
+ "with batch invariance, as it is not yet supported.",
663
+ scope="local",
664
+ )
665
+ cache_config.enable_prefix_caching = False
666
+
667
impl_cls = cast(type[MLAAttentionImpl], self.attn_backend.get_impl_cls())
668
self.impl = impl_cls(
669
num_heads=self.num_heads,
vllm/model_executor/layers/batch_invariant.py
+1
-1
@@ -1006,11 +1006,11 @@ def override_envs_for_invariance():
1006
"FLASH_ATTN", # best supported backend
1007
"FLASHINFER",
1008
"FLASH_ATTN_MLA",
1009
+ "TRITON_MLA",
1010
# Not yet supported MLA backends
1011
# "FLASHMLA",
1012
# "FLEX_ATTENTION", # IMA issue even if we disable batch invariance
1013
# "FLASHINFER_MLA", https://github.com/vllm-project/vllm/pull/28967
1013
- # "TRITON_MLA",
1014
]
1015
if curr_attn_backend not in supported_backends:
1016
error = (