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, 18:45 UTC
[Feature] Batch-Invariant Support for FA2 and LoRA (#30018)
Signed-off-by: quanliu <18646313696@163.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
quanliu committed
Dec 9, 2025 at 23:01 UTC
5dcd593baf9ceaad68a0e5e16f20f6ccccd10797
3 files changed
+23
-3
tests/v1/determinism/test_batch_invariance.py
+10
@@ -10,6 +10,7 @@ from utils import (
10
BACKENDS,
11
_extract_step_logprobs,
12
_random_prompt,
13
+ is_device_capability_below_90,
14
resolve_model_name,
15
skip_unsupported,
16
)
@@ -17,6 +18,8 @@ from utils import (
18
import vllm.model_executor.layers.batch_invariant as batch_invariant
19
from vllm import LLM, SamplingParams
20
21
+IS_DEVICE_CAPABILITY_BELOW_90 = is_device_capability_below_90()
22
+
23
24
@skip_unsupported
25
@pytest.mark.timeout(1000)
@@ -190,6 +193,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
193
max_model_len=8192,
194
dtype="bfloat16", # not everything is supported
195
gpu_memory_utilization=0.9,
196
+ enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
197
)
198
199
# Use more realistic prompts for better token generation
@@ -393,6 +397,8 @@ def test_simple_generation(backend, monkeypatch: pytest.MonkeyPatch):
397
gpu_memory_utilization=0.9,
398
max_model_len=2048,
399
dtype="bfloat16",
400
+ enable_prefix_caching=False,
401
+ enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
402
)
403
404
prompt = "the capital of france is"
@@ -459,6 +465,7 @@ def test_logprobs_without_batch_invariance_should_fail(
465
max_num_seqs=32,
466
max_model_len=8192,
467
dtype="bfloat16",
468
+ enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
469
)
470
471
# build ragged prompts to change shapes significantly across BS=1 vs BS=N
@@ -682,6 +689,7 @@ def test_decode_logprobs_match_prefill_logprobs(
689
max_num_seqs=32,
690
max_model_len=8192,
691
dtype="bfloat16",
692
+ enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
693
)
694
695
# Use a few test prompts
@@ -925,6 +933,8 @@ def LLM_with_max_seqs(
933
max_model_len=max_model_len,
934
dtype="bfloat16",
935
tensor_parallel_size=int(os.getenv("VLLM_TP_SIZE", "1")),
936
+ enable_prefix_caching=False,
937
+ enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
938
# Enable for MOE models
939
# enable_expert_parallel=True,
940
)
tests/v1/determinism/utils.py
+8
-2
@@ -11,8 +11,10 @@ from vllm.platforms import current_platform
11
from vllm.utils.flashinfer import has_flashinfer
12
13
skip_unsupported = pytest.mark.skipif(
14
- not (current_platform.is_cuda() and current_platform.has_device_capability(90)),
15
- reason="Requires CUDA and >= Hopper (SM90)",
14
+ not (current_platform.is_cuda() and current_platform.has_device_capability(80)),
15
+ # Supports testing on Ampere and Ada Lovelace devices.
16
+ # Note: For devices with SM < 90, batch invariance does not support CUDA Graphs.
17
+ reason="Requires CUDA and >= Ampere (SM80)",
18
)
19
20
BACKENDS: list[str] = [
@@ -97,3 +99,7 @@ def _extract_step_logprobs(request_output):
99
return t, inner.token_ids
100
101
return None, None
102
+
103
+
104
+def is_device_capability_below_90() -> bool:
105
+ return not current_platform.has_device_capability(90)
vllm/model_executor/layers/batch_invariant.py
+5
-1
@@ -935,7 +935,11 @@ def enable_batch_invariant_mode():
935
936
# Batch invariant matmuls are no longer needed after cublas overrides
937
if not is_torch_equal_or_newer("2.10.0.dev"):
938
- if current_platform.is_device_capability(100):
938
+ if (
939
+ current_platform.is_device_capability(100)
940
+ or current_platform.is_device_capability(80)
941
+ or current_platform.is_device_capability(89)
942
+ ):
943
# For PyTorch 2.9, B200 uses GEMV for bs=1
944
# Requires https://github.com/pytorch/pytorch/pull/166735
945
_batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, "CUDA")