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 9, 07:30 UTC
Disable Cascade Attention for Batch Invariance (#32561)
Signed-off-by: frankwang28 <frank.wbb@hotmail.com> Signed-off-by: Frank Wang <41319051+frankwang28@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Frank Wang committed
Jan 30, 2026 at 07:00 UTC
8f5d51203b56e8f80cca01dc775e6bf3804ddd55
6 files changed
+60
-9
tests/v1/determinism/test_batch_invariance.py
+11
-3
@@ -188,7 +188,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
188
llm = LLM(
189
model=model_name,
190
tensor_parallel_size=tp_size,
191
- max_num_seqs=32,
191
+ max_num_seqs=128,
192
max_model_len=8192,
193
dtype="bfloat16", # not everything is supported
194
gpu_memory_utilization=0.9,
@@ -197,12 +197,20 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
197
)
198
199
# Use more realistic prompts for better token generation
200
- prompts = [_random_prompt(10, 50) for i in range(32)]
200
+ prompts = [_random_prompt(10, 50) for _ in range(32)]
201
+
202
+ # TODO: Update prompts to have ragged lengths in order to test chunked prefill
203
+ # The above tests are not currently long enough to exercise chunking.
204
+ # prompts = (
205
+ # [_random_prompt(10, 50) for _ in range(28)]
206
+ # + [_random_prompt(256, 512) for _ in range(50)]
207
+ # + [_random_prompt(2048, 4096) for _ in range(50)]
208
+ # )
209
210
sp = SamplingParams(
211
temperature=0.6,
212
top_p=1.0,
205
- max_tokens=8,
213
+ max_tokens=16,
214
seed=1234,
215
logprobs=5,
216
)
tests/v1/determinism/utils.py
+6
-4
@@ -7,7 +7,6 @@ import pytest
7
import torch
8
9
from vllm.platforms import current_platform
10
-from vllm.utils.flashinfer import has_flashinfer
10
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
11
12
skip_unsupported = pytest.mark.skipif(
@@ -22,8 +21,10 @@ BACKENDS: list[str] = [
21
"TRITON_MLA",
22
]
23
25
-if has_flashinfer():
26
- BACKENDS.append("FLASHINFER")
24
+# FlashInfer temporarily disabled due to invariant CTA sizes.
25
+# See FlashInfer issue #2424
26
+# if has_flashinfer():
27
+# BACKENDS.append("FLASHINFER")
28
29
if flash_attn_supports_mla():
30
BACKENDS.append("FLASH_ATTN_MLA")
@@ -78,9 +79,10 @@ def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str:
79
# For longer prompts, repeat context
80
padding_text = (
81
" This is an interesting topic that deserves more explanation. "
82
+ # TODO: Update to * (target_words // 10) to better align with word ratio
83
* (target_words // 50)
84
)
83
- base_prompt = base_prompt + padding_text
85
+ base_prompt = padding_text + base_prompt
86
87
return base_prompt
88
vllm/config/vllm.py
+12
@@ -959,6 +959,18 @@ class VllmConfig:
959
"when cudagraph_mode piecewise cudagraphs is used, "
960
f"cudagraph_mode={self.compilation_config.cudagraph_mode}"
961
)
962
+ from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant
963
+
964
+ if (
965
+ self.model_config
966
+ and vllm_is_batch_invariant()
967
+ and not self.model_config.disable_cascade_attn
968
+ ):
969
+ self.model_config.disable_cascade_attn = True
970
+ logger.warning_once(
971
+ "Disabling cascade attention when VLLM_BATCH_INVARIANT is enabled.",
972
+ scope="local",
973
+ )
974
975
if self.parallel_config.use_ubatching:
976
a2a_backend = self.parallel_config.all2all_backend
vllm/model_executor/layers/batch_invariant.py
+3
-1
@@ -1005,7 +1005,9 @@ def override_envs_for_invariance(
1005
):
1006
supported_backends = [
1007
AttentionBackendEnum.FLASH_ATTN, # best supported backend
1008
- AttentionBackendEnum.FLASHINFER,
1008
+ # FlashInfer temporarily disabled due to invariant CTA sizes.
1009
+ # See FlashInfer issue #2424
1010
+ # AttentionBackendEnum.FLASHINFER,
1011
AttentionBackendEnum.FLASH_ATTN_MLA,
1012
AttentionBackendEnum.TRITON_MLA,
1013
# Not yet supported MLA backends
vllm/model_executor/layers/linear.py
+14
-1
@@ -18,11 +18,18 @@ from vllm.distributed import (
18
)
19
from vllm.logger import init_logger
20
from vllm.model_executor.custom_op import PluggableLayer
21
+from vllm.model_executor.layers.batch_invariant import (
22
+ linear_batch_invariant,
23
+ vllm_is_batch_invariant,
24
+)
25
from vllm.model_executor.layers.quantization.base_config import (
26
QuantizationConfig,
27
QuantizeMethodBase,
28
)
25
-from vllm.model_executor.layers.utils import dispatch_unquantized_gemm
29
+from vllm.model_executor.layers.utils import (
30
+ dispatch_unquantized_gemm,
31
+ is_layer_moe_router_gate,
32
+)
33
from vllm.model_executor.parameter import (
34
BasevLLMParameter,
35
BlockQuantScaleParameter,
@@ -236,6 +243,12 @@ class UnquantizedLinearMethod(LinearMethodBase):
243
x: torch.Tensor,
244
bias: torch.Tensor | None = None,
245
) -> torch.Tensor:
246
+ if (
247
+ vllm_is_batch_invariant()
248
+ and current_platform.is_cuda_alike()
249
+ and is_layer_moe_router_gate(getattr(layer, "prefix", ""))
250
+ ):
251
+ return linear_batch_invariant(x, layer.weight, bias)
252
return dispatch_unquantized_gemm()(layer, x, layer.weight, bias)
253
254
vllm/model_executor/layers/utils.py
+14
@@ -16,6 +16,20 @@ from vllm.utils.torch_utils import direct_register_custom_op
16
17
logger = init_logger(__name__)
18
19
+MOE_LAYER_ROUTER_GATE_SUFFIXES = {
20
+ "gate",
21
+ "router",
22
+ "router_gate",
23
+ "shared_expert_gate",
24
+ "expert_gate",
25
+}
26
+
27
+
28
+def is_layer_moe_router_gate(prefix: str) -> bool:
29
+ if not prefix:
30
+ return False
31
+ return prefix.rsplit(".", 1)[-1] in MOE_LAYER_ROUTER_GATE_SUFFIXES
32
+
33
34
def shuffle_weight(w: torch.Tensor) -> torch.Tensor:
35
# Shuffle weight along the last dimension so that