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 13, 10:45 UTC
Use FLASHINFER MLA backend when testing fp8_kv_scale_compile (#28491)
Signed-off-by: adabeyta <aabeyta@redhat.com>
Adrian Abeyta committed
Nov 11, 2025 at 18:34 UTC
d23539549a6db54ab152ce4e566c31f6891ddab5
1 file changed
+16
-4
tests/compile/test_full_graph.py
+16
-4
@@ -10,6 +10,7 @@ import torch
10
11
from tests.quantization.utils import is_quant_method_supported
12
from vllm import LLM, SamplingParams
13
+from vllm.attention.backends.registry import AttentionBackendEnum
14
from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode, PassConfig
15
from vllm.platforms import current_platform
16
from vllm.utils.torch_utils import is_torch_equal_or_newer
@@ -184,13 +185,24 @@ def test_custom_compile_config(
185
[CompilationMode.NONE, CompilationMode.VLLM_COMPILE],
186
)
187
@pytest.mark.parametrize(
187
- "model",
188
+ "model, backend",
189
[
189
- "Qwen/Qwen2-0.5B", # Standard attention model
190
- "deepseek-ai/DeepSeek-V2-Lite", # MLA (Multi-head Latent Attention) model
190
+ ("Qwen/Qwen2-0.5B", None), # Standard attention model
191
+ (
192
+ "deepseek-ai/DeepSeek-V2-Lite",
193
+ AttentionBackendEnum.FLASHINFER_MLA,
194
+ ), # MLA (Multi-head Latent Attention) model
195
],
196
)
193
-def test_fp8_kv_scale_compile(compilation_mode: int, model: str):
197
+def test_fp8_kv_scale_compile(
198
+ monkeypatch: pytest.MonkeyPatch,
199
+ compilation_mode: int,
200
+ model: str,
201
+ backend: AttentionBackendEnum | None,
202
+):
203
+ if backend:
204
+ monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend.name)
205
+
206
model_kwargs = {
207
"quantization": "fp8",
208
"kv_cache_dtype": "fp8_e4m3",