@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7eb6cb6c18
Fork 0

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, 07:30 UTC

[Attention] Update tests to remove deprecated env vars (#30563)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

Matthew Bonanni committed Dec 17, 2025 at 12:49 UTC 7eb6cb6c18a948fb49824154cb3ece1e32d12cf8
34 files changed +582 -449
.buildkite/scripts/hardware_ci/run-xpu-test.sh
+1 -1
@@ -39,7 +39,7 @@ docker run \
39 python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend ray
40 python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp
41 python3 examples/offline_inference/basic/generate.py --model Intel/Qwen2.5-0.5B-W4A16-G128-AutoRound-LLMC-TEST-ONLY --enforce-eager
42 - VLLM_ATTENTION_BACKEND=TRITON_ATTN python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
42 + python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN
43 cd tests
44 pytest -v -s v1/core
45 pytest -v -s v1/engine
tests/basic_correctness/test_basic_correctness.py
+41 -48
@@ -67,7 +67,6 @@ def _fix_prompt_embed_outputs(
67 @pytest.mark.parametrize("model_executor", ["uni", "mp"])
68 @pytest.mark.parametrize("enable_prompt_embeds", [True, False])
69 def test_models(
70 - monkeypatch: pytest.MonkeyPatch,
70 hf_runner,
71 model: str,
72 backend: str,
@@ -77,48 +76,46 @@ def test_models(
76 model_executor: str,
77 enable_prompt_embeds: bool,
78 ) -> None:
80 - with monkeypatch.context() as m:
81 - m.setenv("VLLM_ATTENTION_BACKEND", backend)
82 -
83 - # 5042 tokens for gemma2
84 - # gemma2 has alternating sliding window size of 4096
85 - # we need a prompt with more than 4096 tokens to test the sliding window
86 - prompt = (
87 - "The following numbers of the sequence "
88 - + ", ".join(str(i) for i in range(1024))
89 - + " are:"
90 - )
91 - example_prompts = [prompt]
92 -
93 - with hf_runner(model) as hf_model:
94 - hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
95 - if enable_prompt_embeds:
96 - with torch.no_grad():
97 - prompt_embeds = hf_model.get_prompt_embeddings(example_prompts)
98 -
99 - with VllmRunner(
100 - model,
101 - max_model_len=8192,
102 - enforce_eager=enforce_eager,
103 - enable_prompt_embeds=enable_prompt_embeds,
104 - gpu_memory_utilization=0.7,
105 - async_scheduling=async_scheduling,
106 - distributed_executor_backend=model_executor,
107 - ) as vllm_model:
108 - if enable_prompt_embeds:
109 - vllm_outputs = vllm_model.generate_greedy(prompt_embeds, max_tokens)
110 - vllm_outputs = _fix_prompt_embed_outputs(
111 - vllm_outputs, hf_model, example_prompts
112 - )
113 - else:
114 - vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
79 + # 5042 tokens for gemma2
80 + # gemma2 has alternating sliding window size of 4096
81 + # we need a prompt with more than 4096 tokens to test the sliding window
82 + prompt = (
83 + "The following numbers of the sequence "
84 + + ", ".join(str(i) for i in range(1024))
85 + + " are:"
86 + )
87 + example_prompts = [prompt]
88 +
89 + with hf_runner(model) as hf_model:
90 + hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
91 + if enable_prompt_embeds:
92 + with torch.no_grad():
93 + prompt_embeds = hf_model.get_prompt_embeddings(example_prompts)
94 +
95 + with VllmRunner(
96 + model,
97 + max_model_len=8192,
98 + enforce_eager=enforce_eager,
99 + enable_prompt_embeds=enable_prompt_embeds,
100 + gpu_memory_utilization=0.7,
101 + async_scheduling=async_scheduling,
102 + distributed_executor_backend=model_executor,
103 + attention_config={"backend": backend},
104 + ) as vllm_model:
105 + if enable_prompt_embeds:
106 + vllm_outputs = vllm_model.generate_greedy(prompt_embeds, max_tokens)
107 + vllm_outputs = _fix_prompt_embed_outputs(
108 + vllm_outputs, hf_model, example_prompts
109 + )
110 + else:
111 + vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
112
116 - check_outputs_equal(
117 - outputs_0_lst=hf_outputs,
118 - outputs_1_lst=vllm_outputs,
119 - name_0="hf",
120 - name_1="vllm",
121 - )
113 + check_outputs_equal(
114 + outputs_0_lst=hf_outputs,
115 + outputs_1_lst=vllm_outputs,
116 + name_0="hf",
117 + name_1="vllm",
118 + )
119
120
121 @multi_gpu_test(num_gpus=2)
@@ -161,12 +158,6 @@ def test_models_distributed(
158 ): # noqa
159 pytest.skip("enable_prompt_embeds does not work with ray compiled dag.")
160
164 - if attention_backend:
165 - monkeypatch_context.setenv(
166 - "VLLM_ATTENTION_BACKEND",
167 - attention_backend,
168 - )
169 -
161 for k, v in extra_env.items():
162 monkeypatch_context.setenv(k, v)
163
@@ -178,6 +169,7 @@ def test_models_distributed(
169 # if we run HF first, the cuda initialization will be done and it
170 # will hurt multiprocessing backend with fork method
171 # (the default method).
172 + attention_config = {"backend": attention_backend} if attention_backend else None
173 with vllm_runner(
174 model,
175 dtype=dtype,
@@ -185,6 +177,7 @@ def test_models_distributed(
177 distributed_executor_backend=distributed_executor_backend,
178 enable_prompt_embeds=enable_prompt_embeds,
179 gpu_memory_utilization=0.7,
180 + attention_config=attention_config,
181 ) as vllm_model:
182 if enable_prompt_embeds:
183 with hf_runner(model, dtype=dtype) as hf_model:
tests/compile/distributed/test_fusions_e2e.py
+6 -3
@@ -208,7 +208,8 @@ def test_attn_quant(
208 # To capture subprocess logs, we need to know whether spawn or fork is used.
209 # Force spawn as it is more general.
210 monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
211 - monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend.name)
211 +
212 + model_kwargs["attention_config"] = {"backend": backend.name}
213
214 compilation_config = CompilationConfig(
215 # Testing properties
@@ -297,7 +298,8 @@ def test_tp2_attn_quant_allreduce_rmsnorm(
298 # To capture subprocess logs, we need to know whether spawn or fork is used.
299 # Force spawn as it is more general.
300 monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
300 - monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend.name)
301 +
302 + model_kwargs["attention_config"] = {"backend": backend.name}
303
304 compilation_config = CompilationConfig(
305 # Testing properties
@@ -409,7 +411,8 @@ def test_tp2_attn_quant_async_tp(
411 # To capture subprocess logs, we need to know whether spawn or fork is used.
412 # Force spawn as it is more general.
413 monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
412 - monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend.name)
414 +
415 + model_kwargs["attention_config"] = {"backend": backend.name}
416
417 compilation_config = CompilationConfig(
418 # Testing properties
tests/compile/fullgraph/test_basic_correctness.py
+40 -42
@@ -89,7 +89,6 @@ class TestSetting:
89 ],
90 )
91 def test_compile_correctness(
92 - monkeypatch: pytest.MonkeyPatch,
92 test_setting: TestSetting,
93 ):
94 # this test is run under multiple suits, with different GPUs.
@@ -107,49 +106,48 @@ def test_compile_correctness(
106 f"{cuda_device_count_stateless()}"
107 )
108
110 - with monkeypatch.context() as m:
111 - m.setenv("VLLM_ATTENTION_BACKEND", attn_backend)
112 - final_args = [
113 - *model_args,
114 - "-pp",
115 - str(pp_size),
116 - "-tp",
117 - str(tp_size),
118 - "-cc.cudagraph_mode=none",
119 - ]
109 + final_args = [
110 + *model_args,
111 + "-pp",
112 + str(pp_size),
113 + "-tp",
114 + str(tp_size),
115 + "-cc.cudagraph_mode=none",
116 + f"--attention-backend={attn_backend}",
117 + ]
118
121 - all_args: list[list[str]] = []
122 - all_envs: list[dict[str, str] | None] = []
119 + all_args: list[list[str]] = []
120 + all_envs: list[dict[str, str] | None] = []
121
124 - for comp_mode in [
125 - CompilationMode.STOCK_TORCH_COMPILE,
126 - CompilationMode.DYNAMO_TRACE_ONCE,
127 - CompilationMode.VLLM_COMPILE,
128 - ]:
129 - for mode in [CompilationMode.NONE, comp_mode]:
130 - all_args.append(
131 - final_args + [f"-cc.mode={mode.name}", "-cc.backend=inductor"]
132 - )
133 -
134 - # inductor will change the output, so we only compare if the output
135 - # is close, not exactly the same.
136 - compare_all_settings(
137 - model,
138 - all_args,
139 - all_envs,
140 - method=method if method != "generate" else "generate_close",
122 + for comp_mode in [
123 + CompilationMode.STOCK_TORCH_COMPILE,
124 + CompilationMode.DYNAMO_TRACE_ONCE,
125 + CompilationMode.VLLM_COMPILE,
126 + ]:
127 + for mode in [CompilationMode.NONE, comp_mode]:
128 + all_args.append(
129 + final_args + [f"-cc.mode={mode.name}", "-cc.backend=inductor"]
130 )
142 - all_envs.clear()
143 - all_args.clear()
131
145 - for mode in [
146 - CompilationMode.NONE,
147 - CompilationMode.STOCK_TORCH_COMPILE,
148 - CompilationMode.DYNAMO_TRACE_ONCE,
149 - CompilationMode.VLLM_COMPILE,
150 - ]:
151 - all_args.append(final_args + [f"-cc.mode={mode.name}", "-cc.backend=eager"])
152 - all_envs.append({})
153 - all_envs.append({})
132 + # inductor will change the output, so we only compare if the output
133 + # is close, not exactly the same.
134 + compare_all_settings(
135 + model,
136 + all_args,
137 + all_envs,
138 + method=method if method != "generate" else "generate_close",
139 + )
140 + all_envs.clear()
141 + all_args.clear()
142 +
143 + for mode in [
144 + CompilationMode.NONE,
145 + CompilationMode.STOCK_TORCH_COMPILE,
146 + CompilationMode.DYNAMO_TRACE_ONCE,
147 + CompilationMode.VLLM_COMPILE,
148 + ]:
149 + all_args.append(final_args + [f"-cc.mode={mode.name}", "-cc.backend=eager"])
150 + all_envs.append({})
151 + all_envs.append({})
152
155 - compare_all_settings(model, all_args * 3, all_envs, method=method)
153 + compare_all_settings(model, all_args * 3, all_envs, method=method)
tests/compile/fullgraph/test_full_cudagraph.py
+3 -10
@@ -74,7 +74,6 @@ def llm_pair(request):
74 # Force native sampler to avoid potential nondeterminism in FlashInfer
75 # when per-request generators are not used in V1.
76 "VLLM_USE_FLASHINFER_SAMPLER": "0",
77 - **backend_config.env_vars,
77 }
78 with temporary_environ(env_vars):
79 full = LLM(
@@ -170,16 +169,10 @@ class TestFullCUDAGraph:
169
170 @pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
171 def test_full_cudagraph_with_invalid_backend():
173 - with (
174 - temporary_environ(
175 - {
176 - "VLLM_ATTENTION_BACKEND": "FLEX_ATTENTION",
177 - # Flex_Attention is not supported with full cuda graph
178 - }
179 - ),
180 - pytest.raises(RuntimeError),
181 - ):
172 + # Flex_Attention is not supported with full cuda graph
173 + with pytest.raises(RuntimeError):
174 LLM(
175 model="Qwen/Qwen2-1.5B-Instruct",
176 compilation_config=CompilationConfig(cudagraph_mode="FULL"),
177 + attention_config={"backend": "FLEX_ATTENTION"},
178 )
tests/compile/fullgraph/test_full_graph.py
+3 -4
@@ -197,20 +197,19 @@ def test_custom_compile_config(
197 ],
198 )
199 def test_fp8_kv_scale_compile(
200 - monkeypatch: pytest.MonkeyPatch,
200 compilation_mode: int,
201 model: str,
202 backend: AttentionBackendEnum | None,
203 ):
205 - if backend:
206 - monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend.name)
207 -
204 model_kwargs = {
205 "quantization": "fp8",
206 "kv_cache_dtype": "fp8_e4m3",
207 "calculate_kv_scales": True,
208 "max_model_len": 512,
209 }
210 + if backend:
211 + model_kwargs["attention_config"] = {"backend": backend.name}
212 +
213 run_model(compilation_mode, model, **model_kwargs)
214
215
tests/distributed/test_context_parallel.py
+1 -3
@@ -219,14 +219,12 @@ def _test_cp_gsm8k(
219 ]
220 )
221
222 - server_env = {}
222 if attn_backend:
224 - server_env["VLLM_ATTENTION_BACKEND"] = attn_backend
223 + server_args.append(f"--attention-backend={attn_backend}")
224
225 with RemoteOpenAIServer(
226 model_id,
227 server_args,
229 - env_dict=server_env,
228 max_wait_seconds=720,
229 ) as remote_server:
230 host = f"http://{remote_server.host}"
tests/distributed/test_pp_cudagraph.py
+12 -14
@@ -20,23 +20,21 @@ from ..utils import compare_two_settings, create_new_process_for_each_test
20 )
21 @create_new_process_for_each_test()
22 def test_pp_cudagraph(
23 - monkeypatch: pytest.MonkeyPatch,
23 PP_SIZE: int,
24 MODEL_NAME: str,
25 ATTN_BACKEND: LiteralString,
26 ):
28 - with monkeypatch.context() as m:
29 - cudagraph_args = [
30 - # use half precision for speed and memory savings in CI environment
31 - "--dtype",
32 - "float16",
33 - "--pipeline-parallel-size",
34 - str(PP_SIZE),
35 - "--distributed-executor-backend",
36 - "mp",
37 - ]
38 - m.setenv("VLLM_ATTENTION_BACKEND", ATTN_BACKEND)
27 + cudagraph_args = [
28 + # use half precision for speed and memory savings in CI environment
29 + "--dtype",
30 + "float16",
31 + "--pipeline-parallel-size",
32 + str(PP_SIZE),
33 + "--distributed-executor-backend",
34 + "mp",
35 + f"--attention-backend={ATTN_BACKEND}",
36 + ]
37
40 - eager_args = cudagraph_args + ["--enforce-eager"]
38 + eager_args = cudagraph_args + ["--enforce-eager"]
39
42 - compare_two_settings(MODEL_NAME, eager_args, cudagraph_args)
40 + compare_two_settings(MODEL_NAME, eager_args, cudagraph_args)
tests/engine/test_arg_utils.py
+134 -1
@@ -9,7 +9,7 @@ from typing import Annotated, Literal
9
10 import pytest
11
12 -from vllm.config import CompilationConfig, config
12 +from vllm.config import AttentionConfig, CompilationConfig, config
13 from vllm.engine.arg_utils import (
14 EngineArgs,
15 contains_type,
@@ -298,6 +298,139 @@ def test_compilation_config():
298 )
299
300
301 +def test_attention_config():
302 + from vllm.attention.backends.registry import AttentionBackendEnum
303 +
304 + parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
305 +
306 + # default value
307 + args = parser.parse_args([])
308 + assert args is not None
309 + engine_args = EngineArgs.from_cli_args(args)
310 + assert engine_args.attention_config == AttentionConfig()
311 +
312 + # set backend via dot notation
313 + args = parser.parse_args(["--attention-config.backend", "FLASH_ATTN"])
314 + assert args is not None
315 + engine_args = EngineArgs.from_cli_args(args)
316 + assert engine_args.attention_config.backend is not None
317 + assert engine_args.attention_config.backend.name == "FLASH_ATTN"
318 +
319 + # set backend via --attention-backend shorthand
320 + args = parser.parse_args(["--attention-backend", "FLASHINFER"])
321 + assert args is not None
322 + engine_args = EngineArgs.from_cli_args(args)
323 + assert engine_args.attention_backend is not None
324 + assert engine_args.attention_backend == "FLASHINFER"
325 +
326 + # set all fields via dot notation
327 + args = parser.parse_args(
328 + [
329 + "--attention-config.backend",
330 + "FLASH_ATTN",
331 + "--attention-config.flash_attn_version",
332 + "3",
333 + "--attention-config.use_prefill_decode_attention",
334 + "true",
335 + "--attention-config.flash_attn_max_num_splits_for_cuda_graph",
336 + "16",
337 + "--attention-config.use_cudnn_prefill",
338 + "true",
339 + "--attention-config.use_trtllm_ragged_deepseek_prefill",
340 + "true",
341 + "--attention-config.use_trtllm_attention",
342 + "true",
343 + "--attention-config.disable_flashinfer_prefill",
344 + "true",
345 + "--attention-config.disable_flashinfer_q_quantization",
346 + "true",
347 + ]
348 + )
349 + assert args is not None
350 + engine_args = EngineArgs.from_cli_args(args)
351 + assert engine_args.attention_config.backend is not None
352 + assert engine_args.attention_config.backend.name == "FLASH_ATTN"
353 + assert engine_args.attention_config.flash_attn_version == 3
354 + assert engine_args.attention_config.use_prefill_decode_attention is True
355 + assert engine_args.attention_config.flash_attn_max_num_splits_for_cuda_graph == 16
356 + assert engine_args.attention_config.use_cudnn_prefill is True
357 + assert engine_args.attention_config.use_trtllm_ragged_deepseek_prefill is True
358 + assert engine_args.attention_config.use_trtllm_attention is True
359 + assert engine_args.attention_config.disable_flashinfer_prefill is True
360 + assert engine_args.attention_config.disable_flashinfer_q_quantization is True
361 +
362 + # set to string form of a dict with all fields
363 + args = parser.parse_args(
364 + [
365 + "--attention-config="
366 + '{"backend": "FLASHINFER", "flash_attn_version": 2, '
367 + '"use_prefill_decode_attention": false, '
368 + '"flash_attn_max_num_splits_for_cuda_graph": 8, '
369 + '"use_cudnn_prefill": false, '
370 + '"use_trtllm_ragged_deepseek_prefill": false, '
371 + '"use_trtllm_attention": false, '
372 + '"disable_flashinfer_prefill": false, '
373 + '"disable_flashinfer_q_quantization": false}',
374 + ]
375 + )
376 + assert args is not None
377 + engine_args = EngineArgs.from_cli_args(args)
378 + assert engine_args.attention_config.backend is not None
379 + assert engine_args.attention_config.backend.name == "FLASHINFER"
380 + assert engine_args.attention_config.flash_attn_version == 2
381 + assert engine_args.attention_config.use_prefill_decode_attention is False
382 + assert engine_args.attention_config.flash_attn_max_num_splits_for_cuda_graph == 8
383 + assert engine_args.attention_config.use_cudnn_prefill is False
384 + assert engine_args.attention_config.use_trtllm_ragged_deepseek_prefill is False
385 + assert engine_args.attention_config.use_trtllm_attention is False
386 + assert engine_args.attention_config.disable_flashinfer_prefill is False
387 + assert engine_args.attention_config.disable_flashinfer_q_quantization is False
388 +
389 + # test --attention-backend flows into VllmConfig.attention_config
390 + args = parser.parse_args(
391 + [
392 + "--model",
393 + "facebook/opt-125m",
394 + "--attention-backend",
395 + "FLASH_ATTN",
396 + ]
397 + )
398 + assert args is not None
399 + engine_args = EngineArgs.from_cli_args(args)
400 + vllm_config = engine_args.create_engine_config()
401 + assert vllm_config.attention_config.backend == AttentionBackendEnum.FLASH_ATTN
402 +
403 + # test --attention-config.backend flows into VllmConfig.attention_config
404 + args = parser.parse_args(
405 + [
406 + "--model",
407 + "facebook/opt-125m",
408 + "--attention-config.backend",
409 + "FLASHINFER",
410 + ]
411 + )
412 + assert args is not None
413 + engine_args = EngineArgs.from_cli_args(args)
414 + vllm_config = engine_args.create_engine_config()
415 + assert vllm_config.attention_config.backend == AttentionBackendEnum.FLASHINFER
416 +
417 + # test --attention-backend and --attention-config.backend are mutually exclusive
418 + args = parser.parse_args(
419 + [
420 + "--model",
421 + "facebook/opt-125m",
422 + "--attention-backend",
423 + "FLASH_ATTN",
424 + "--attention-config.backend",
425 + "FLASHINFER",
426 + ]
427 + )
428 + assert args is not None
429 + engine_args = EngineArgs.from_cli_args(args)
430 + with pytest.raises(ValueError, match="mutually exclusive"):
431 + engine_args.create_engine_config()
432 +
433 +
434 def test_prefix_cache_default():
435 parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
436 args = parser.parse_args([])
tests/entrypoints/openai/test_serving_chat.py
+4 -9
@@ -76,15 +76,10 @@ def default_server_args(with_tool_parser: bool):
76
77
78 @pytest.fixture(scope="module")
79 -def gptoss_server(
80 - monkeypatch_module: pytest.MonkeyPatch, default_server_args: list[str]
81 -):
82 - with monkeypatch_module.context() as m:
83 - m.setenv("VLLM_ATTENTION_BACKEND", "TRITON_ATTN")
84 - with RemoteOpenAIServer(
85 - GPT_OSS_MODEL_NAME, default_server_args
86 - ) as remote_server:
87 - yield remote_server
79 +def gptoss_server(default_server_args: list[str]):
80 + server_args = default_server_args + ["--attention-backend=TRITON_ATTN"]
81 + with RemoteOpenAIServer(GPT_OSS_MODEL_NAME, server_args) as remote_server:
82 + yield remote_server
83
84
85 @pytest_asyncio.fixture
tests/kernels/attention/test_attention_selector.py
+27 -25
@@ -6,7 +6,9 @@ from unittest.mock import patch
6 import pytest
7 import torch
8
9 +from vllm.attention.backends.registry import AttentionBackendEnum
10 from vllm.attention.selector import _cached_get_attn_backend, get_attn_backend
11 +from vllm.config import AttentionConfig, VllmConfig, set_current_vllm_config
12 from vllm.platforms import current_platform
13 from vllm.platforms.cpu import CpuPlatform
14 from vllm.platforms.cuda import CudaPlatform
@@ -73,18 +75,18 @@ def generate_params():
75
76
77 @pytest.mark.parametrize("device, name, use_mla, block_size", generate_params())
76 -def test_env(
78 +def test_backend_selection(
79 device: str,
80 name: str,
81 use_mla: bool,
82 block_size: int,
81 - monkeypatch: pytest.MonkeyPatch,
83 ):
84 """Test attention backend selection with valid device-backend pairs."""
84 - with monkeypatch.context() as m:
85 - m.setenv("VLLM_ATTENTION_BACKEND", name)
86 - m.setenv("VLLM_MLA_DISABLE", "1" if use_mla else "0")
85 + # Create AttentionConfig with the specified backend
86 + attention_config = AttentionConfig(backend=AttentionBackendEnum[name])
87 + vllm_config = VllmConfig(attention_config=attention_config)
88
89 + with set_current_vllm_config(vllm_config):
90 if device == "cpu":
91 with patch("vllm.platforms.current_platform", CpuPlatform()):
92 backend = get_attn_backend(16, torch.float16, None, block_size)
@@ -217,27 +219,32 @@ def test_env(
219 @pytest.mark.parametrize("device", ["cpu", "cuda"])
220 def test_fp32_fallback(device: str):
221 """Test attention backend selection with fp32."""
220 - if device == "cpu":
221 - with patch("vllm.platforms.current_platform", CpuPlatform()):
222 - backend = get_attn_backend(16, torch.float32, None, 16)
223 - assert backend.get_name() == "CPU_ATTN"
222 + # Use default config (no backend specified)
223 + vllm_config = VllmConfig()
224
225 - elif device == "cuda":
226 - with patch("vllm.platforms.current_platform", CudaPlatform()):
227 - backend = get_attn_backend(16, torch.float32, None, 16)
228 - assert backend.get_name() == "FLEX_ATTENTION"
225 + with set_current_vllm_config(vllm_config):
226 + if device == "cpu":
227 + with patch("vllm.platforms.current_platform", CpuPlatform()):
228 + backend = get_attn_backend(16, torch.float32, None, 16)
229 + assert backend.get_name() == "CPU_ATTN"
230 +
231 + elif device == "cuda":
232 + with patch("vllm.platforms.current_platform", CudaPlatform()):
233 + backend = get_attn_backend(16, torch.float32, None, 16)
234 + assert backend.get_name() == "FLEX_ATTENTION"
235
236
237 def test_flash_attn(monkeypatch: pytest.MonkeyPatch):
238 """Test FlashAttn validation."""
239 pytest.skip(
240 "Skipping as current backend selector does not "
235 - "handle fallbacks when a backend is set via env var."
241 + "handle fallbacks when a backend is explicitly set."
242 )
243
238 - with monkeypatch.context() as m:
239 - m.setenv("VLLM_ATTENTION_BACKEND", "FLASH_ATTN")
244 + attention_config = AttentionConfig(backend=AttentionBackendEnum.FLASH_ATTN)
245 + vllm_config = VllmConfig(attention_config=attention_config)
246
247 + with set_current_vllm_config(vllm_config):
248 # Unsupported CUDA arch
249 monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _=None: (7, 5))
250 backend = get_attn_backend(16, torch.float16, None, 16)
@@ -277,15 +284,10 @@ def test_flash_attn(monkeypatch: pytest.MonkeyPatch):
284 assert backend.get_name() != "FLASH_ATTN"
285
286
280 -def test_invalid_env(monkeypatch: pytest.MonkeyPatch):
287 +def test_invalid_backend():
288 """Test that invalid attention backend names raise ValueError."""
289 with (
283 - monkeypatch.context() as m,
284 - patch("vllm.platforms.current_platform", CudaPlatform()),
290 + pytest.raises(ValueError),
291 ):
286 - m.setenv("VLLM_ATTENTION_BACKEND", "INVALID")
287 -
288 - # Should raise ValueError for invalid backend
289 - with pytest.raises(ValueError) as exc_info:
290 - get_attn_backend(32, torch.float16, None, 16)
291 - assert "Invalid value 'INVALID'" in str(exc_info.value)
292 + # Invalid backend name should raise ValueError when creating enum
293 + AttentionConfig(backend=AttentionBackendEnum["INVALID"])
tests/kernels/attention/test_rocm_attention_selector.py
+39 -21
@@ -4,7 +4,9 @@
4 import pytest
5 import torch
6
7 +from vllm.attention.backends.registry import AttentionBackendEnum
8 from vllm.attention.selector import _cached_get_attn_backend, get_attn_backend
9 +from vllm.config import AttentionConfig, VllmConfig, set_current_vllm_config
10 from vllm.platforms.rocm import RocmPlatform
11
12
@@ -16,40 +18,56 @@ def clear_cache():
18
19 @pytest.mark.skip(reason="Skipped for now. Should be revisited.")
20 def test_selector(monkeypatch: pytest.MonkeyPatch):
19 - with monkeypatch.context() as m:
20 - m.setenv("VLLM_ATTENTION_BACKEND", "ROCM_ATTN")
21 + # Set the current platform to ROCm using monkeypatch
22 + monkeypatch.setattr("vllm.attention.selector.current_platform", RocmPlatform())
23
22 - # Set the current platform to ROCm using monkeypatch
23 - monkeypatch.setattr("vllm.attention.selector.current_platform", RocmPlatform())
24 + # Test standard ROCm attention
25 + attention_config = AttentionConfig(backend=AttentionBackendEnum.ROCM_ATTN)
26 + vllm_config = VllmConfig(attention_config=attention_config)
27
25 - # Test standard ROCm attention
28 + with set_current_vllm_config(vllm_config):
29 backend = get_attn_backend(16, torch.float16, torch.float16, 16, False)
30 assert backend.get_name() == "ROCM_FLASH" or backend.get_name() == "TRITON_ATTN"
31
29 - # MLA test for deepseek related
32 + # MLA test for deepseek related
33 + # Change the attention backend to triton MLA
34 + attention_config = AttentionConfig(backend=AttentionBackendEnum.TRITON_MLA)
35 + vllm_config = VllmConfig(attention_config=attention_config)
36
31 - # change the attention backend to triton MLA
32 - m.setenv("VLLM_ATTENTION_BACKEND", "TRITON_MLA")
37 + with set_current_vllm_config(vllm_config):
38 backend = get_attn_backend(576, torch.bfloat16, "auto", 16, False, use_mla=True)
39 assert backend.get_name() == "TRITON_MLA"
40
36 - # If attention backend is None
37 - # If use_mla is true
38 - # The selected backend is triton MLA
39 - m.setenv("VLLM_ATTENTION_BACKEND", "")
41 + # If attention backend is None
42 + # If use_mla is true
43 + # The selected backend is triton MLA
44 + attention_config = AttentionConfig(backend=None)
45 + vllm_config = VllmConfig(attention_config=attention_config)
46 +
47 + with set_current_vllm_config(vllm_config):
48 backend = get_attn_backend(576, torch.bfloat16, "auto", 16, False, use_mla=True)
49 assert backend.get_name() == "TRITON_MLA"
50
43 - # change the attention backend to AITER MLA
44 - m.setenv("VLLM_ATTENTION_BACKEND", "ROCM_AITER_MLA")
51 + # Change the attention backend to AITER MLA
52 + attention_config = AttentionConfig(backend=AttentionBackendEnum.ROCM_AITER_MLA)
53 + vllm_config = VllmConfig(attention_config=attention_config)
54 +
55 + with set_current_vllm_config(vllm_config):
56 backend = get_attn_backend(576, torch.bfloat16, "auto", 1, False, use_mla=True)
57 assert backend.get_name() == "ROCM_AITER_MLA"
58
48 - # If attention backend is None
49 - # If use_mla is true
50 - # If VLLM_ROCM_USE_AITER is enabled
51 - # The selected backend is ROCM_AITER_MLA
52 - m.setenv("VLLM_ATTENTION_BACKEND", "")
59 + # If attention backend is None
60 + # If use_mla is true
61 + # If VLLM_ROCM_USE_AITER is enabled
62 + # The selected backend is ROCM_AITER_MLA
63 + with monkeypatch.context() as m:
64 m.setenv("VLLM_ROCM_USE_AITER", "1")
54 - backend = get_attn_backend(576, torch.bfloat16, "auto", 1, False, use_mla=True)
55 - assert backend.get_name() == "ROCM_AITER_MLA"
65 +
66 + attention_config = AttentionConfig(backend=None)
67 + vllm_config = VllmConfig(attention_config=attention_config)
68 +
69 + with set_current_vllm_config(vllm_config):
70 + backend = get_attn_backend(
71 + 576, torch.bfloat16, "auto", 1, False, use_mla=True
72 + )
73 + assert backend.get_name() == "ROCM_AITER_MLA"
tests/kernels/test_flex_attention.py
+44 -51
@@ -37,7 +37,7 @@ def set_seed(seed):
37 not torch.cuda.is_available() or TORCH_VERSION < MINIMUM_TORCH_VERSION,
38 reason="CUDA not available or PyTorch version < 2.7",
39 )
40 -def test_flex_attention_vs_default_backend(vllm_runner, monkeypatch):
40 +def test_flex_attention_vs_default_backend(vllm_runner):
41 """Test that FlexAttention produces the same outputs as the default backend.
42
43 This test compares the outputs from the FlexAttention backend with
@@ -54,35 +54,32 @@ def test_flex_attention_vs_default_backend(vllm_runner, monkeypatch):
54 ]
55
56 # Run with flex attention
57 - with monkeypatch.context() as m:
58 - m.setenv("VLLM_ATTENTION_BACKEND", "FLEX_ATTENTION")
59 -
60 - set_seed(seed)
61 - with vllm_runner(
62 - model_name,
63 - runner="generate",
64 - tensor_parallel_size=1,
65 - num_gpu_blocks_override=128,
66 - enforce_eager=True,
67 - ) as llm_flex:
68 - output_flex = llm_flex.generate_greedy_logprobs(
69 - prompts, max_tokens, num_logprobs
70 - )
57 + set_seed(seed)
58 + with vllm_runner(
59 + model_name,
60 + runner="generate",
61 + tensor_parallel_size=1,
62 + num_gpu_blocks_override=128,
63 + enforce_eager=True,
64 + attention_config={"backend": "FLEX_ATTENTION"},
65 + ) as llm_flex:
66 + output_flex = llm_flex.generate_greedy_logprobs(
67 + prompts, max_tokens, num_logprobs
68 + )
69
70 # Run with default backend
73 - with monkeypatch.context() as m:
74 - set_seed(seed)
75 - with vllm_runner(
76 - model_name,
77 - runner="generate",
78 - tensor_parallel_size=1,
79 - num_gpu_blocks_override=128,
80 - enforce_eager=True,
81 - gpu_memory_utilization=0.85,
82 - ) as llm_default:
83 - output_default = llm_default.generate_greedy_logprobs(
84 - prompts, max_tokens, num_logprobs
85 - )
71 + set_seed(seed)
72 + with vllm_runner(
73 + model_name,
74 + runner="generate",
75 + tensor_parallel_size=1,
76 + num_gpu_blocks_override=128,
77 + enforce_eager=True,
78 + gpu_memory_utilization=0.85,
79 + ) as llm_default:
80 + output_default = llm_default.generate_greedy_logprobs(
81 + prompts, max_tokens, num_logprobs
82 + )
83
84 check_logprobs_close(
85 outputs_0_lst=output_flex,
@@ -96,7 +93,7 @@ def test_flex_attention_vs_default_backend(vllm_runner, monkeypatch):
93 not torch.cuda.is_available() or TORCH_VERSION < MINIMUM_TORCH_VERSION,
94 reason="CUDA not available or PyTorch version < 2.7",
95 )
99 -def test_encoder_flex_attention_vs_default_backend(vllm_runner, monkeypatch):
96 +def test_encoder_flex_attention_vs_default_backend(vllm_runner):
97 """Test that FlexAttention produces the same outputs as the default backend.
98
99 This test compares the outputs from the FlexAttention backend with
@@ -110,30 +107,26 @@ def test_encoder_flex_attention_vs_default_backend(vllm_runner, monkeypatch):
107 ]
108
109 # Run with flex attention
113 - with monkeypatch.context() as m:
114 - m.setenv("VLLM_ATTENTION_BACKEND", "FLEX_ATTENTION")
115 - with vllm_runner(
116 - model_name,
117 - runner="pooling",
118 - dtype=torch.bfloat16,
119 - tensor_parallel_size=1,
120 - max_model_len=100,
121 - enforce_eager=True,
122 - ) as llm_flex:
123 - flex_outputs = llm_flex.embed(prompts)
110 + with vllm_runner(
111 + model_name,
112 + runner="pooling",
113 + dtype=torch.bfloat16,
114 + tensor_parallel_size=1,
115 + max_model_len=100,
116 + enforce_eager=True,
117 + attention_config={"backend": "FLEX_ATTENTION"},
118 + ) as llm_flex:
119 + flex_outputs = llm_flex.embed(prompts)
120
121 # Run with default backend
126 - with (
127 - monkeypatch.context() as m,
128 - vllm_runner(
129 - model_name,
130 - runner="pooling",
131 - dtype=torch.bfloat16,
132 - tensor_parallel_size=1,
133 - max_model_len=100,
134 - enforce_eager=True,
135 - ) as llm_default,
136 - ):
122 + with vllm_runner(
123 + model_name,
124 + runner="pooling",
125 + dtype=torch.bfloat16,
126 + tensor_parallel_size=1,
127 + max_model_len=100,
128 + enforce_eager=True,
129 + ) as llm_default:
130 default_outputs = llm_default.embed(prompts)
131
132 check_embeddings_close(
tests/models/multimodal/generation/test_granite_speech.py
+9 -3
@@ -35,10 +35,12 @@ audio_lora_path = MODEL_NAME
35 models = [MODEL_NAME]
36
37
38 -@pytest.fixture(autouse=True)
39 -def set_attention_backend_for_rocm(monkeypatch):
38 +@pytest.fixture
39 +def granite_speech_attention_config():
40 + """Return attention config for Granite Speech tests on ROCm."""
41 if current_platform.is_rocm():
41 - monkeypatch.setenv("VLLM_ATTENTION_BACKEND", "TRITON_ATTN")
42 + return {"backend": "TRITON_ATTN"}
43 + return None
44
45
46 def run_test(
@@ -53,6 +55,7 @@ def run_test(
55 num_logprobs: int,
56 tensor_parallel_size: int,
57 distributed_executor_backend: str | None = None,
58 + attention_config: dict | None = None,
59 ):
60 """Inference result should be the same between hf and vllm.
61
@@ -80,6 +83,7 @@ def run_test(
83 enable_lora=True,
84 max_lora_rank=64,
85 enforce_eager=True,
86 + attention_config=attention_config,
87 ) as vllm_model:
88 lora_request = LoRARequest("audio", 1, audio_lora_path)
89 vllm_outputs_per_case = [
@@ -131,6 +135,7 @@ def test_models(
135 vllm_runner,
136 model: str,
137 audio_assets: AudioTestAssets,
138 + granite_speech_attention_config,
139 dtype: str,
140 max_model_len: int,
141 max_tokens: int,
@@ -157,4 +162,5 @@ def test_models(
162 max_tokens=max_tokens,
163 num_logprobs=num_logprobs,
164 tensor_parallel_size=1,
165 + attention_config=granite_speech_attention_config,
166 )
tests/models/multimodal/pooling/conftest.py
+9 -15
@@ -2,23 +2,17 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Pytest configuration for vLLM pooling tests."""
4
5 -import os
6 -import warnings
5 +import pytest
6
7 from vllm.platforms import current_platform
8
9
11 -def pytest_collection_modifyitems(config, items):
12 - """Set FLEX_ATTENTION backend for SigLIP tests on ROCm."""
13 - if not current_platform.is_rocm():
14 - return
10 +@pytest.fixture
11 +def siglip_attention_config():
12 + """Return attention config for SigLIP tests on ROCm.
13
16 - siglip_tests = [item for item in items if "test_siglip" in item.nodeid]
17 -
18 - if siglip_tests:
19 - os.environ["VLLM_ATTENTION_BACKEND"] = "FLEX_ATTENTION"
20 - warnings.warn(
21 - "ROCm: Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION for SigLIP tests",
22 - UserWarning,
23 - stacklevel=1,
24 - )
14 + On ROCm, SigLIP tests require FLEX_ATTENTION backend.
15 + """
16 + if current_platform.is_rocm():
17 + return {"backend": "FLEX_ATTENTION"}
18 + return None
tests/models/multimodal/pooling/test_siglip.py
+8
@@ -38,6 +38,7 @@ def _run_test(
38 *,
39 dtype: str,
40 tokenization_kwargs: dict[str, Any] | None = None,
41 + attention_config: dict[str, Any] | None = None,
42 ) -> None:
43 if tokenization_kwargs is None:
44 tokenization_kwargs = {}
@@ -49,6 +50,7 @@ def _run_test(
50 enforce_eager=True,
51 max_model_len=64,
52 gpu_memory_utilization=0.7,
53 + attention_config=attention_config,
54 ) as vllm_model:
55 vllm_outputs = vllm_model.embed(
56 input_texts, images=input_images, tokenization_kwargs=tokenization_kwargs
@@ -90,6 +92,7 @@ def test_models_text(
92 hf_runner,
93 vllm_runner,
94 image_assets,
95 + siglip_attention_config,
96 model: str,
97 dtype: str,
98 ) -> None:
@@ -108,6 +111,7 @@ def test_models_text(
111 "padding": "max_length",
112 "max_length": 64,
113 }, # siglip2 was trained with this padding setting.
114 + attention_config=siglip_attention_config,
115 )
116
117
@@ -117,6 +121,7 @@ def test_models_image(
121 hf_runner,
122 vllm_runner,
123 image_assets,
124 + siglip_attention_config,
125 model: str,
126 dtype: str,
127 ) -> None:
@@ -133,6 +138,7 @@ def test_models_image(
138 input_images,
139 model,
140 dtype=dtype,
141 + attention_config=siglip_attention_config,
142 )
143
144
@@ -141,6 +147,7 @@ def test_models_image(
147 def test_models_text_image_no_crash(
148 vllm_runner,
149 image_assets,
150 + siglip_attention_config,
151 model: str,
152 dtype: str,
153 ) -> None:
@@ -154,6 +161,7 @@ def test_models_text_image_no_crash(
161 enforce_eager=True,
162 max_model_len=64,
163 gpu_memory_utilization=0.7,
164 + attention_config=siglip_attention_config,
165 ) as vllm_model:
166 with pytest.raises(ValueError, match="not both"):
167 vllm_model.embed(texts, images=images)
tests/models/quantization/test_fp8.py
+2 -1
@@ -75,7 +75,6 @@ def test_models(
75
76 with monkeypatch.context() as m:
77 m.setenv("TOKENIZERS_PARALLELISM", "true")
78 - m.setenv("VLLM_ATTENTION_BACKEND", backend)
78
79 MAX_MODEL_LEN = 1024
80 NUM_LOG_PROBS = 8
@@ -86,6 +85,7 @@ def test_models(
85 tensor_parallel_size=tensor_parallel_size,
86 enforce_eager=enforce_eager,
87 kv_cache_dtype="auto",
88 + attention_config={"backend": backend},
89 ) as vllm_model:
90 baseline_outputs = vllm_model.generate_greedy_logprobs(
91 example_prompts, max_tokens, NUM_LOG_PROBS
@@ -97,6 +97,7 @@ def test_models(
97 tensor_parallel_size=tensor_parallel_size,
98 enforce_eager=enforce_eager,
99 kv_cache_dtype=kv_cache_dtype,
100 + attention_config={"backend": backend},
101 ) as vllm_model:
102 test_outputs = vllm_model.generate_greedy_logprobs(
103 example_prompts, max_tokens, NUM_LOG_PROBS
tests/models/test_initialization.py
+7 -5
@@ -108,11 +108,12 @@ def can_initialize(
108 patch.object(V1EngineCore, "_initialize_kv_caches", _initialize_kv_caches_v1),
109 monkeypatch.context() as m,
110 ):
111 - if model_arch == "GptOssForCausalLM":
112 - # FIXME: A hack to bypass FA3 assertion because our CI's L4 GPU
113 - # has cc==8.9 which hasn't supported FA3 yet. Remove this hack when
114 - # L4 supports FA3.
115 - m.setenv("VLLM_ATTENTION_BACKEND", "TRITON_ATTN")
111 + # FIXME: A hack to bypass FA3 assertion because our CI's L4 GPU
112 + # has cc==8.9 which hasn't supported FA3 yet. Remove this hack when
113 + # L4 supports FA3.
114 + attention_config = (
115 + {"backend": "TRITON_ATTN"} if model_arch == "GptOssForCausalLM" else None
116 + )
117 if model_arch == "WhisperForConditionalGeneration":
118 m.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
119
@@ -143,6 +144,7 @@ def can_initialize(
144 else "vllm",
145 hf_overrides=hf_overrides_fn,
146 max_num_seqs=model_info.max_num_seqs,
147 + attention_config=attention_config,
148 )
149
150
tests/v1/attention/test_rocm_attention_backends_selection.py
+3 -9
@@ -94,26 +94,20 @@ def mock_on_gfx9():
94 None,
95 AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path(),
96 ),
97 - # Test Case 9: VLLM_V1_USE_PREFILL_DECODE_ATTENTION=1
98 - (
99 - {"VLLM_V1_USE_PREFILL_DECODE_ATTENTION": "1"},
100 - None,
101 - AttentionBackendEnum.ROCM_ATTN.get_path(),
102 - ),
103 - # Test Case 10: VLLM_ROCM_USE_AITER=1 + explicit TRITON_ATTN
97 + # Test Case 9: VLLM_ROCM_USE_AITER=1 + explicit TRITON_ATTN
98 (
99 {"VLLM_ROCM_USE_AITER": "1"},
100 "TRITON_ATTN",
101 AttentionBackendEnum.TRITON_ATTN.get_path(),
102 ),
109 - # Test Case 11: VLLM_ROCM_USE_AITER=1 + VLLM_ROCM_USE_AITER_MHA=0
103 + # Test Case 10: VLLM_ROCM_USE_AITER=1 + VLLM_ROCM_USE_AITER_MHA=0
104 # (explicitly disabled)
105 (
106 {"VLLM_ROCM_USE_AITER": "1", "VLLM_ROCM_USE_AITER_MHA": "0"},
107 None,
108 AttentionBackendEnum.TRITON_ATTN.get_path(),
109 ),
116 - # Test Case 12: VLLM_ROCM_USE_AITER=1 + explicit ROCM_ATTN
110 + # Test Case 11: VLLM_ROCM_USE_AITER=1 + explicit ROCM_ATTN
111 (
112 {"VLLM_ROCM_USE_AITER": "1"},
113 "ROCM_ATTN",
tests/v1/attention/utils.py
+22 -25
@@ -249,8 +249,8 @@ def create_dummy_kv_cache(
249 @dataclass
250 class BackendConfig:
251 name: str
252 - env_vars: dict
253 - comp_config: dict # compilation config
252 + attention_config: dict
253 + comp_config: dict
254 specific_gpu_arch: tuple | None = None
255
256
@@ -259,10 +259,10 @@ full_cg_backend_configs = {
259 # FA3 on Hopper
260 "FA3": BackendConfig(
261 name="FA3",
262 - env_vars={
263 - "VLLM_ATTENTION_BACKEND": "FLASH_ATTN",
264 - "VLLM_FLASH_ATTN_VERSION": "3",
265 - "VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH": "16",
262 + attention_config={
263 + "backend": "FLASH_ATTN",
264 + "flash_attn_version": 3,
265 + "flash_attn_max_num_splits_for_cuda_graph": 16,
266 },
267 comp_config={
268 "cudagraph_mode": "FULL",
@@ -272,9 +272,7 @@ full_cg_backend_configs = {
272 # FlashMLA on Hopper
273 "FlashMLA": BackendConfig(
274 name="FlashMLA",
275 - env_vars={
276 - "VLLM_ATTENTION_BACKEND": "FLASHMLA",
277 - },
275 + attention_config={"backend": "FLASHMLA"},
276 comp_config={
277 "cudagraph_mode": "FULL_AND_PIECEWISE",
278 },
@@ -283,9 +281,7 @@ full_cg_backend_configs = {
281 # Cutlass MLA on Blackwell
282 "CutlassMLA": BackendConfig(
283 name="CutlassMLA",
286 - env_vars={
287 - "VLLM_ATTENTION_BACKEND": "CUTLASS_MLA",
288 - },
284 + attention_config={"backend": "CUTLASS_MLA"},
285 comp_config={
286 "cudagraph_mode": "FULL_AND_PIECEWISE",
287 },
@@ -294,9 +290,7 @@ full_cg_backend_configs = {
290 # FlashInfer MLA on Blackwell
291 "FlashInferMLA": BackendConfig(
292 name="FlashInferMLA",
297 - env_vars={
298 - "VLLM_ATTENTION_BACKEND": "FLASHINFER_MLA",
299 - },
293 + attention_config={"backend": "FLASHINFER_MLA"},
294 comp_config={
295 "cudagraph_mode": "FULL_AND_PIECEWISE",
296 },
@@ -305,9 +299,9 @@ full_cg_backend_configs = {
299 # FlashAttention MLA on Hopper
300 "FlashAttentionMLA": BackendConfig(
301 name="FlashAttentionMLA",
308 - env_vars={
309 - "VLLM_ATTENTION_BACKEND": "FLASH_ATTN_MLA",
310 - "VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH": "16",
302 + attention_config={
303 + "backend": "FLASH_ATTN_MLA",
304 + "flash_attn_max_num_splits_for_cuda_graph": 16,
305 },
306 comp_config={
307 "cudagraph_mode": "FULL_DECODE_ONLY",
@@ -317,10 +311,10 @@ full_cg_backend_configs = {
311 # FA2
312 "FA2": BackendConfig(
313 name="FA2",
320 - env_vars={
321 - "VLLM_ATTENTION_BACKEND": "FLASH_ATTN",
322 - "VLLM_FLASH_ATTN_VERSION": "2",
323 - "VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH": "16",
314 + attention_config={
315 + "backend": "FLASH_ATTN",
316 + "flash_attn_version": 2,
317 + "flash_attn_max_num_splits_for_cuda_graph": 16,
318 },
319 comp_config={
320 "cudagraph_mode": "FULL_AND_PIECEWISE",
@@ -329,7 +323,7 @@ full_cg_backend_configs = {
323 # Triton Attention
324 "TritonAttn": BackendConfig(
325 name="TritonAttn",
332 - env_vars={"VLLM_ATTENTION_BACKEND": "TRITON_ATTN"},
326 + attention_config={"backend": "TRITON_ATTN"},
327 comp_config={
328 "cudagraph_mode": "FULL_AND_PIECEWISE",
329 },
@@ -337,14 +331,17 @@ full_cg_backend_configs = {
331 # FlashInfer
332 "FlashInfer": BackendConfig(
333 name="FlashInfer",
340 - env_vars={"VLLM_ATTENTION_BACKEND": "FLASHINFER"},
334 + attention_config={"backend": "FLASHINFER"},
335 comp_config={
336 "cudagraph_mode": "FULL_AND_PIECEWISE",
337 },
338 ),
339 "RocmAttn": BackendConfig(
340 name="RocmAttn",
347 - env_vars={"VLLM_V1_USE_PREFILL_DECODE_ATTENTION": "1"},
341 + attention_config={
342 + "backend": "ROCM_ATTN",
343 + "use_prefill_decode_attention": True,
344 + },
345 comp_config={
346 "cudagraph_mode": "FULL",
347 },
tests/v1/cudagraph/test_cudagraph_mode.py
+7 -26
@@ -1,7 +1,5 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import contextlib
4 -import os
3 import weakref
4 from contextlib import ExitStack
5
@@ -13,26 +11,6 @@ from vllm import LLM
11 from vllm.config import CompilationConfig, CompilationMode
12 from vllm.platforms import current_platform
13
16 -
17 -@contextlib.contextmanager
18 -def temporary_environ(env_vars):
19 - """
20 - Temporarily set environment variables and restore them afterward.
21 - We have to do this vs monkeypatch because monkeypatch doesn't work
22 - with "module" scoped fixtures.
23 - """
24 - original_env = {k: os.environ.get(k) for k in env_vars}
25 - try:
26 - os.environ.update(env_vars)
27 - yield
28 - finally:
29 - for k, v in original_env.items():
30 - if v is None:
31 - os.environ.pop(k, None)
32 - else:
33 - os.environ[k] = v
34 -
35 -
14 # test attention backend and cudagraph_mode combo
15 # (backend_name, cudagraph_mode, supported)
16 if current_platform.is_rocm():
@@ -68,9 +46,9 @@ def test_backend_and_cudagraph_mode_combo(backend_name, cudagraph_mode, supporte
46 ):
47 pytest.skip("Only Hopper GPUs support FA3 and FlashMLA")
48
71 - env_vars = backend_configs[backend_name].env_vars
49 + attention_config = backend_config.attention_config
50
73 - with temporary_environ(env_vars), ExitStack() as stack:
51 + with ExitStack() as stack:
52 if not supported:
53 stack.enter_context(pytest.raises(Exception))
54
@@ -80,6 +58,7 @@ def test_backend_and_cudagraph_mode_combo(backend_name, cudagraph_mode, supporte
58 trust_remote_code=True,
59 gpu_memory_utilization=0.45,
60 max_model_len=1024,
61 + attention_config=attention_config,
62 compilation_config=CompilationConfig(
63 mode=CompilationMode.VLLM_COMPILE, cudagraph_mode=cudagraph_mode
64 ),
@@ -122,9 +101,10 @@ combo_cases_2 = [
101 def test_cudagraph_compilation_combo(
102 backend_name, cudagraph_mode, compilation_mode, supported
103 ):
125 - env_vars = backend_configs[backend_name].env_vars
104 + backend_config = backend_configs[backend_name]
105 + attention_config = backend_config.attention_config
106
127 - with temporary_environ(env_vars), ExitStack() as stack:
107 + with ExitStack() as stack:
108 if not supported:
109 stack.enter_context(pytest.raises(Exception))
110
@@ -134,6 +114,7 @@ def test_cudagraph_compilation_combo(
114 trust_remote_code=True,
115 gpu_memory_utilization=0.45,
116 max_model_len=1024,
117 + attention_config=attention_config,
118 compilation_config=CompilationConfig(
119 mode=compilation_mode, cudagraph_mode=cudagraph_mode
120 ),
tests/v1/determinism/test_batch_invariance.py
+13 -12
@@ -28,7 +28,7 @@ IS_DEVICE_CAPABILITY_BELOW_90 = is_device_capability_below_90()
28 BACKENDS,
29 )
30 def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
31 - backend, monkeypatch: pytest.MonkeyPatch
31 + backend,
32 ):
33 """
34 Ensures that the same request (the 'needle' prompt) yields identical output
@@ -54,7 +54,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
54 seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
55 random.seed(seed)
56
57 - monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend)
57 + attention_config = {"backend": backend}
58 # Allow overrides from environment (useful for CI tuning)
59 # "facebook/opt-125m" is too small, doesn't reliably test determinism
60 model = resolve_model_name(backend)
@@ -92,6 +92,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
92 max_num_seqs=max_batch_size,
93 gpu_memory_utilization=gpu_mem_util,
94 max_model_len=max_model_len,
95 + attention_config=attention_config,
96 )
97
98 # Baseline generation for the needle prompt alone.
@@ -106,6 +107,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
107 max_num_seqs=max_batch_size,
108 gpu_memory_utilization=gpu_mem_util,
109 max_model_len=max_model_len,
110 + attention_config=attention_config,
111 )
112
113 mismatches = 0
@@ -163,10 +165,8 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
165 BACKENDS,
166 )
167 def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
166 - backend, monkeypatch: pytest.MonkeyPatch
168 + backend,
169 ):
168 - monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend)
169 -
170 seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
171 random.seed(seed)
172 model_name = resolve_model_name(backend)
@@ -193,6 +193,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
193 dtype="bfloat16", # not everything is supported
194 gpu_memory_utilization=0.9,
195 enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
196 + attention_config={"backend": backend},
197 )
198
199 # Use more realistic prompts for better token generation
@@ -381,12 +382,11 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
382 "backend",
383 BACKENDS,
384 )
384 -def test_simple_generation(backend, monkeypatch: pytest.MonkeyPatch):
385 +def test_simple_generation(backend):
386 """
387 Simple test that runs the model with a basic prompt and prints the output.
388 Useful for quick smoke testing and debugging.
389 """
389 - monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend)
390 model = resolve_model_name(backend)
391
392 llm = LLM(
@@ -398,6 +398,7 @@ def test_simple_generation(backend, monkeypatch: pytest.MonkeyPatch):
398 dtype="bfloat16",
399 enable_prefix_caching=False,
400 enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
401 + attention_config={"backend": backend},
402 )
403
404 prompt = "the capital of france is"
@@ -444,8 +445,6 @@ def test_logprobs_without_batch_invariance_should_fail(
445 The test will PASS if we detect differences (proving batch invariance matters).
446 The test will FAIL if everything matches (suggesting batch invariance isn't needed).
447 """
447 - monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend)
448 -
448 # CRITICAL: Disable batch invariance for this test
449 monkeypatch.setenv("VLLM_BATCH_INVARIANT", "0")
450 monkeypatch.setattr(batch_invariant, "VLLM_BATCH_INVARIANT", False)
@@ -465,6 +464,7 @@ def test_logprobs_without_batch_invariance_should_fail(
464 max_model_len=8192,
465 dtype="bfloat16",
466 enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
467 + attention_config={"backend": backend},
468 )
469
470 # build ragged prompts to change shapes significantly across BS=1 vs BS=N
@@ -649,7 +649,7 @@ def test_logprobs_without_batch_invariance_should_fail(
649 @skip_unsupported
650 @pytest.mark.parametrize("backend", ["FLASH_ATTN"])
651 def test_decode_logprobs_match_prefill_logprobs(
652 - backend, monkeypatch: pytest.MonkeyPatch
652 + backend,
653 ):
654 """
655 Test that verifies decode logprobs match prefill logprobs.
@@ -664,8 +664,6 @@ def test_decode_logprobs_match_prefill_logprobs(
664 This ensures that the logprobs from decode are consistent with what
665 we would get if we ran prefill on each prefix.
666 """
667 - monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend)
668 -
667 seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
668 random.seed(seed)
669 model_name = resolve_model_name(backend)
@@ -689,6 +687,7 @@ def test_decode_logprobs_match_prefill_logprobs(
687 max_model_len=8192,
688 dtype="bfloat16",
689 enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
690 + attention_config={"backend": backend},
691 )
692
693 # Use a few test prompts
@@ -920,6 +919,7 @@ def LLM_with_max_seqs(
919 max_num_seqs: int,
920 gpu_memory_utilization: float,
921 max_model_len: int,
922 + attention_config: dict | None = None,
923 ) -> LLM:
924 """
925 Helper to construct an LLM with a specific max_num_seqs (batch-size limit)
@@ -934,6 +934,7 @@ def LLM_with_max_seqs(
934 tensor_parallel_size=int(os.getenv("VLLM_TP_SIZE", "1")),
935 enable_prefix_caching=False,
936 enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
937 + attention_config=attention_config,
938 # Enable for MOE models
939 # enable_expert_parallel=True,
940 )
tests/v1/determinism/test_online_batch_invariance.py
+2 -3
@@ -136,11 +136,9 @@ def _compare_bs1_vs_bsn_single_process(
136 @skip_unsupported
137 @pytest.mark.parametrize("backend", BACKENDS)
138 def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
139 - backend: str, monkeypatch: pytest.MonkeyPatch
139 + backend: str,
140 ) -> None:
141 random.seed(int(os.getenv("VLLM_TEST_SEED", "12345")))
142 - # Override backend for this test (and the RemoteOpenAIServer child process).
143 - monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend)
142 model_name = resolve_model_name(backend)
143 prompts_all = [_random_prompt(10, 50) for _ in range(32)]
144
@@ -156,6 +154,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
154 server_args: list[str] = [
155 "--max-model-len=8192",
156 "--max-num-seqs=32",
157 + f"--attention-backend={backend}",
158 ]
159 if tp_size:
160 server_args += ["-tp", tp_size]
tests/v1/e2e/test_async_scheduling.py
+13 -9
@@ -142,16 +142,17 @@ def run_tests(
142 """Test consistency of combos of async scheduling, preemption,
143 uni/multiproc executor with spec decoding."""
144
145 - with monkeypatch.context() as m:
146 - # avoid precision errors
147 - if current_platform.is_rocm():
148 - if is_testing_with_spec_decoding:
149 - # Use TRITON_ATTN for spec decoding test for consistency
150 - m.setenv("VLLM_ATTENTION_BACKEND", "TRITON_ATTN")
151 - else:
152 - m.setenv("VLLM_ATTENTION_BACKEND", "ROCM_AITER_FA")
145 + # Determine attention config based on platform
146 + if current_platform.is_rocm():
147 + if is_testing_with_spec_decoding:
148 + # Use TRITON_ATTN for spec decoding test for consistency
149 + attention_config = {"backend": "TRITON_ATTN"}
150 else:
154 - m.setenv("VLLM_ATTENTION_BACKEND", "FLEX_ATTENTION")
151 + attention_config = {"backend": "ROCM_AITER_FA"}
152 + else:
153 + attention_config = {"backend": "FLEX_ATTENTION"}
154 +
155 + with monkeypatch.context() as m:
156 # lock matmul precision to full FP32 (IEEE)
157 m.setenv("VLLM_FLOAT32_MATMUL_PRECISION", "ieee")
158 # m.setenv("VLLM_BATCH_INVARIANT", "1")
@@ -174,6 +175,7 @@ def run_tests(
175 spec_config,
176 test_prefill_chunking=test_prefill_chunking,
177 is_testing_with_spec_decoding=is_testing_with_spec_decoding,
178 + attention_config=attention_config,
179 )
180 outputs.append(test_results)
181
@@ -262,6 +264,7 @@ def run_test(
264 spec_config: dict[str, Any] | None,
265 test_prefill_chunking: bool,
266 is_testing_with_spec_decoding: bool = False,
267 + attention_config: dict[str, Any] | None = None,
268 ):
269 spec_decoding = spec_config is not None
270 cache_arg: dict[str, Any] = (
@@ -301,6 +304,7 @@ def run_test(
304 dtype=dtype,
305 speculative_config=spec_config,
306 disable_log_stats=False,
307 + attention_config=attention_config,
308 **cache_arg,
309 ) as vllm_model:
310 results = []
tests/v1/e2e/test_cascade_attention.py
+16 -17
@@ -10,7 +10,7 @@ from ...utils import create_new_process_for_each_test
10
11 @create_new_process_for_each_test()
12 @pytest.mark.parametrize("attn_backend", ["FLASH_ATTN", "FLASHINFER"])
13 -def test_cascade_attention(example_system_message, monkeypatch, attn_backend):
13 +def test_cascade_attention(example_system_message, attn_backend):
14 prompt = "\n<User>: Implement fibonacci sequence in Python.\n<Claude>:"
15
16 if attn_backend == "FLASHINFER":
@@ -19,19 +19,18 @@ def test_cascade_attention(example_system_message, monkeypatch, attn_backend):
19 "needs investigation. See issue #25679."
20 )
21
22 - with monkeypatch.context() as m:
23 - m.setenv("VLLM_ATTENTION_BACKEND", attn_backend)
24 -
25 - llm = LLM(model="Qwen/Qwen2-1.5B-Instruct")
26 - sampling_params = SamplingParams(temperature=0.0, max_tokens=100)
27 -
28 - # No cascade attention.
29 - single_prompt = [example_system_message + prompt]
30 - responses = llm.generate(single_prompt, sampling_params)
31 - ref_output = responses[0].outputs[0].text
32 -
33 - # (Probably) Use cascade attention.
34 - prompts = [example_system_message + prompt] * 64
35 - responses = llm.generate(prompts, sampling_params)
36 - for response in responses:
37 - assert response.outputs[0].text == ref_output
22 + llm = LLM(
23 + model="Qwen/Qwen2-1.5B-Instruct", attention_config={"backend": attn_backend}
24 + )
25 + sampling_params = SamplingParams(temperature=0.0, max_tokens=100)
26 +
27 + # No cascade attention.
28 + single_prompt = [example_system_message + prompt]
29 + responses = llm.generate(single_prompt, sampling_params)
30 + ref_output = responses[0].outputs[0].text
31 +
32 + # (Probably) Use cascade attention.
33 + prompts = [example_system_message + prompt] * 64
34 + responses = llm.generate(prompts, sampling_params)
35 + for response in responses:
36 + assert response.outputs[0].text == ref_output
tests/v1/e2e/test_spec_decode.py
+24 -19
@@ -438,25 +438,26 @@ def test_eagle_correctness(
438 should be the same when using eagle speculative decoding.
439 model_setup: (method, model_name, eagle_model_name, tp_size)
440 """
441 - with monkeypatch.context() as m:
442 - if "Llama-4-Scout" in model_setup[1] and attn_backend == "FLASH_ATTN":
443 - # Scout requires default backend selection
444 - # because vision encoder has head_dim 88 being incompatible
445 - # with FLASH_ATTN and needs to fall back to Flex Attn
446 -
447 - # pass if not ROCm
448 - if current_platform.is_rocm():
449 - # TODO: Enable Flex Attn for spec_decode on ROCm
450 - pytest.skip("Flex Attn for spec_decode not supported on ROCm currently")
451 - else:
452 - m.setenv("VLLM_MLA_DISABLE", "1")
453 - m.setenv("VLLM_ATTENTION_BACKEND", attn_backend)
441 + # Determine attention config
442 + # Scout requires default backend selection because vision encoder has
443 + # head_dim 88 being incompatible with FLASH_ATTN and needs to fall back
444 + # to Flex Attn
445 + if "Llama-4-Scout" in model_setup[1] and attn_backend == "FLASH_ATTN":
446 + if current_platform.is_rocm():
447 + # TODO: Enable Flex Attn for spec_decode on ROCm
448 + pytest.skip("Flex Attn for spec_decode not supported on ROCm currently")
449 + attention_config = None # Let it fall back to default
450 + else:
451 + attention_config = {"backend": attn_backend}
452 +
453 + if attn_backend == "TRITON_ATTN" and not current_platform.is_rocm():
454 + pytest.skip(
455 + "TRITON_ATTN does not support "
456 + "multi-token eagle spec decode on current platform"
457 + )
458
455 - if attn_backend == "TRITON_ATTN" and not current_platform.is_rocm():
456 - pytest.skip(
457 - "TRITON_ATTN does not support "
458 - "multi-token eagle spec decode on current platform"
459 - )
459 + with monkeypatch.context() as m:
460 + m.setenv("VLLM_MLA_DISABLE", "1")
461
462 if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm():
463 if "deepseek" in model_setup[1].lower():
@@ -471,7 +472,10 @@ def test_eagle_correctness(
472 max_num_batched_tokens = 128 if enable_chunked_prefill else max_model_len
473
474 ref_llm = LLM(
474 - model=model_name, max_model_len=max_model_len, tensor_parallel_size=tp_size
475 + model=model_name,
476 + max_model_len=max_model_len,
477 + tensor_parallel_size=tp_size,
478 + attention_config=attention_config,
479 )
480 ref_outputs = ref_llm.chat(test_prompts, sampling_config)
481 del ref_llm
@@ -492,6 +496,7 @@ def test_eagle_correctness(
496 max_num_batched_tokens=max_num_batched_tokens,
497 enable_chunked_prefill=enable_chunked_prefill,
498 model_impl=model_impl,
499 + attention_config=attention_config,
500 )
501 spec_outputs = spec_llm.chat(test_prompts, sampling_config)
502 matches = 0
tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh
+20 -2
@@ -3,21 +3,29 @@ set -xe
3
4 # Parse command line arguments
5 KV_BUFFER_DEVICE="cuda" # Default to cuda
6 +ATTENTION_BACKEND="" # Default to empty (use vllm default)
7 while [[ $# -gt 0 ]]; do
8 case $1 in
9 --kv_buffer_device)
10 KV_BUFFER_DEVICE="$2"
11 shift 2
12 ;;
13 + --attention-backend)
14 + ATTENTION_BACKEND="$2"
15 + shift 2
16 + ;;
17 *)
18 echo "Unknown option $1"
14 - echo "Usage: $0 [--kv_buffer_device <cuda|cpu>]"
19 + echo "Usage: $0 [--kv_buffer_device <cuda|cpu>] [--attention-backend <backend>]"
20 exit 1
21 ;;
22 esac
23 done
24
25 echo "Running accuracy tests with kv_buffer_device=$KV_BUFFER_DEVICE"
26 +if [[ -n "$ATTENTION_BACKEND" ]]; then
27 + echo "Using attention backend: $ATTENTION_BACKEND"
28 +fi
29
30 DECODER_KV_LAYOUT=${DECODER_KV_LAYOUT:-"HND"} # Default to HND, optional NHD
31 if [[ "$DECODER_KV_LAYOUT" == "NHD" ]]; then
@@ -148,6 +156,11 @@ run_tests_for_model() {
156 --tensor-parallel-size $PREFILLER_TP_SIZE \
157 --kv-transfer-config '$KV_CONFIG'"
158
159 + # Add attention backend config if specified
160 + if [[ -n "$ATTENTION_BACKEND" ]]; then
161 + BASE_CMD="${BASE_CMD} --attention-backend=$ATTENTION_BACKEND"
162 + fi
163 +
164 if [ -n "$model_args" ]; then
165 FULL_CMD="$BASE_CMD $model_args"
166 else
@@ -188,7 +201,12 @@ run_tests_for_model() {
201 --block-size ${DECODE_BLOCK_SIZE} \
202 --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
203 --kv-transfer-config '$KV_CONFIG'"
191 -
204 +
205 + # Add attention backend config if specified
206 + if [[ -n "$ATTENTION_BACKEND" ]]; then
207 + BASE_CMD="${BASE_CMD} --attention-backend=$ATTENTION_BACKEND"
208 + fi
209 +
210 # DP-EP attention mode
211 if [[ -z "$DP_EP" ]]; then
212 BASE_CMD="${BASE_CMD} --tensor-parallel-size $DECODER_TP_SIZE"
tests/v1/kv_connector/nixl_integration/tp_config_sweep_accuracy_test.sh
+6 -6
@@ -15,14 +15,14 @@ configs=(
15
16 run_tests() {
17 local label=$1
18 - local extra_env=$2
18 + local extra_args=$2
19
20 echo "=== Running tests (${label}) ==="
21 for cfg in "${configs[@]}"; do
22 - echo "-> Running with ${cfg} ${extra_env:+and ${extra_env}}"
22 + echo "-> Running with ${cfg} ${extra_args:+and ${extra_args}}"
23 # Use 'env' to safely set variables without eval
24 - if ! env ${extra_env} ${cfg} bash "${SCRIPT}"; then
25 - echo "❌ Test failed for config: ${cfg} ${extra_env:+(${extra_env})}"
24 + if ! env ${cfg} bash "${SCRIPT}" ${extra_args}; then
25 + echo "❌ Test failed for config: ${cfg} ${extra_args:+(${extra_args})}"
26 exit 1
27 fi
28 done
@@ -34,8 +34,8 @@ run_tests "default backend" ""
34
35 # Check if FLASHINFER is set (non-empty)
36 if [[ -n "${FLASHINFER:-}" ]]; then
37 - echo "FLASHINFER is set, rerunning with VLLM_ATTENTION_BACKEND=FLASHINFER"
38 - run_tests "FLASHINFER backend" "VLLM_ATTENTION_BACKEND=FLASHINFER"
37 + echo "FLASHINFER is set, rerunning with --attention-backend FLASHINFER"
38 + run_tests "FLASHINFER backend" "--attention-backend FLASHINFER"
39 else
40 echo "FLASHINFER not set, skipping FLASHINFER runs."
41 fi
tests/v1/kv_connector/unit/test_nixl_connector.py
+2 -4
@@ -1132,7 +1132,7 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):
1132 "TRITON_ATTN",
1133 ],
1134 )
1135 -def test_register_kv_caches(dist_init, attn_backend, monkeypatch):
1135 +def test_register_kv_caches(dist_init, attn_backend):
1136 """
1137 Test that register_kv_caches() properly calls nixl_wrapper methods with
1138 correct data.
@@ -1144,9 +1144,7 @@ def test_register_kv_caches(dist_init, attn_backend, monkeypatch):
1144 block layout info
1145 """
1146
1147 - monkeypatch.setenv("VLLM_ATTENTION_BACKEND", attn_backend)
1148 -
1149 - vllm_config = create_vllm_config()
1147 + vllm_config = create_vllm_config(attention_backend=attn_backend)
1148
1149 # Import the appropriate backend based on the parameter
1150 if attn_backend == "FLASH_ATTN":
tests/v1/kv_connector/unit/utils.py
+4
@@ -11,6 +11,7 @@ import torch
11
12 from vllm import SamplingParams
13 from vllm.config import (
14 + AttentionConfig,
15 CacheConfig,
16 DeviceConfig,
17 KVTransferConfig,
@@ -94,6 +95,7 @@ def create_vllm_config(
95 dtype: str = "float16",
96 cache_dtype: str = "auto",
97 hf_overrides: dict[str, Any] | None = None,
98 + attention_backend: str | None = None,
99 ) -> VllmConfig:
100 """Initialize VllmConfig For Testing."""
101 model_config = ModelConfig(
@@ -124,12 +126,14 @@ def create_vllm_config(
126 enable_permute_local_kv=enable_permute_local_kv,
127 kv_connector_extra_config=kv_connector_extra_config or {},
128 )
129 + attention_config = AttentionConfig(backend=attention_backend)
130 return VllmConfig(
131 scheduler_config=scheduler_config,
132 model_config=model_config,
133 cache_config=cache_config,
134 kv_transfer_config=kv_transfer_config,
135 device_config=DeviceConfig("cpu"),
136 + attention_config=attention_config,
137 )
138
139
tests/v1/kv_offload/test_cpu_offloading.py
+7 -8
@@ -13,7 +13,6 @@ from vllm import LLM, SamplingParams, TokensPrompt
13 from vllm.config import KVEventsConfig, KVTransferConfig
14 from vllm.distributed.kv_events import BlockStored, KVEventBatch
15 from vllm.platforms import current_platform
16 -from vllm.utils.system_utils import set_env_var
16
17 CPU_BLOCK_SIZES = [48]
18 ATTN_BACKENDS = ["FLASH_ATTN"]
@@ -180,13 +179,13 @@ def test_cpu_offloading(cpu_block_size: int, attn_backend: str) -> None:
179 topic="test",
180 )
181
183 - with set_env_var("VLLM_ATTENTION_BACKEND", attn_backend):
184 - llm = LLM(
185 - model="meta-llama/Llama-3.2-1B-Instruct",
186 - gpu_memory_utilization=0.5,
187 - kv_events_config=kv_events_config,
188 - kv_transfer_config=kv_transfer_config,
189 - )
182 + llm = LLM(
183 + model="meta-llama/Llama-3.2-1B-Instruct",
184 + gpu_memory_utilization=0.5,
185 + kv_events_config=kv_events_config,
186 + kv_transfer_config=kv_transfer_config,
187 + attention_config={"backend": attn_backend},
188 + )
189
190 events_endpoint = events_endpoint.replace("*", "127.0.0.1")
191 subscriber = MockSubscriber(events_endpoint, topic=kv_events_config.topic)
tests/v1/spec_decode/test_eagle.py
+12 -7
@@ -15,6 +15,7 @@ from tests.v1.attention.utils import (
15 )
16 from vllm.attention.backends.registry import AttentionBackendEnum
17 from vllm.config import (
18 + AttentionConfig,
19 CacheConfig,
20 DeviceConfig,
21 ModelConfig,
@@ -38,6 +39,7 @@ eagle3_dir = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B"
39 def _create_proposer(
40 method: str,
41 num_speculative_tokens: int,
42 + attention_backend: str | None = None,
43 speculative_token_tree: list[tuple[int, ...]] | None = None,
44 ) -> EagleProposer:
45 model_config = ModelConfig(model=model_dir, runner="generate", max_model_len=100)
@@ -70,6 +72,7 @@ def _create_proposer(
72 max_model_len=model_config.max_model_len,
73 is_encoder_decoder=model_config.is_encoder_decoder,
74 ),
75 + attention_config=AttentionConfig(backend=attention_backend),
76 )
77
78 return EagleProposer(vllm_config=vllm_config, device=current_platform.device_type)
@@ -331,8 +334,6 @@ def test_load_model(
334 use_distinct_lm_head,
335 monkeypatch,
336 ):
334 - monkeypatch.setenv("VLLM_ATTENTION_BACKEND", attn_backend)
335 -
337 if attn_backend == "TRITON_ATTN" and not current_platform.is_rocm():
338 pytest.skip(
339 "TRITON_ATTN does not support "
@@ -394,7 +395,9 @@ def test_load_model(
395 assert not isinstance(target_model, SupportsMultiModal)
396
397 # Create proposer using the helper function
397 - proposer = _create_proposer(method, num_speculative_tokens=8)
398 + proposer = _create_proposer(
399 + method, num_speculative_tokens=8, attention_backend=attn_backend
400 + )
401
402 # Call the method under test
403 proposer.load_model(target_model)
@@ -420,8 +423,6 @@ def test_load_model(
423 @pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform())
424 @pytest.mark.parametrize("num_speculative_tokens", [1, 3, 8])
425 def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):
423 - monkeypatch.setenv("VLLM_ATTENTION_BACKEND", attn_backend)
424 -
426 if attn_backend == "TRITON_ATTN" and not current_platform.is_rocm():
427 pytest.skip(
428 "TRITON_ATTN does not support "
@@ -449,7 +450,9 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):
450 seq_lens = [seq_len_1, seq_len_2]
451
452 # Create proposer first so we can use its actual hidden_size
452 - proposer = _create_proposer("eagle", num_speculative_tokens)
453 + proposer = _create_proposer(
454 + "eagle", num_speculative_tokens, attention_backend=attn_backend
455 + )
456 # Get the hidden_size from the proposer to ensure consistency
457 hidden_size = proposer.hidden_size
458
@@ -622,7 +625,9 @@ def test_propose_tree(spec_token_tree):
625
626 # Create proposer first so we can use its actual hidden_size.
627 proposer = _create_proposer(
625 - "eagle", num_speculative_tokens, speculative_token_tree=spec_token_tree
628 + "eagle",
629 + num_speculative_tokens,
630 + speculative_token_tree=spec_token_tree,
631 )
632 # Get the hidden_size from the proposer to ensure consistency.
633 hidden_size = proposer.hidden_size
tests/v1/spec_decode/test_max_len.py
+40 -45
@@ -38,53 +38,48 @@ def test_ngram_max_len(num_speculative_tokens: int):
38 def test_eagle_max_len(
39 monkeypatch: pytest.MonkeyPatch, num_speculative_tokens: int, attn_backend: str
40 ):
41 - with monkeypatch.context() as m:
42 - m.setenv("VLLM_ATTENTION_BACKEND", attn_backend)
43 -
44 - if attn_backend == "TRITON_ATTN" and not current_platform.is_rocm():
45 - pytest.skip(
46 - "TRITON_ATTN does not support "
47 - "multi-token eagle spec decode on current platform"
48 - )
41 + if attn_backend == "TRITON_ATTN" and not current_platform.is_rocm():
42 + pytest.skip(
43 + "TRITON_ATTN does not support "
44 + "multi-token eagle spec decode on current platform"
45 + )
46
50 - if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm():
51 - m.setenv("VLLM_ROCM_USE_AITER", "1")
47 + if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm():
48 + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
49
53 - llm = LLM(
54 - model="meta-llama/Meta-Llama-3-8B-Instruct",
55 - enforce_eager=True, # For faster initialization.
56 - speculative_config={
57 - "method": "eagle",
58 - "model": "yuhuili/EAGLE-LLaMA3-Instruct-8B",
59 - "num_speculative_tokens": num_speculative_tokens,
60 - "max_model_len": 80,
61 - },
62 - max_model_len=200,
50 + llm = LLM(
51 + model="meta-llama/Meta-Llama-3-8B-Instruct",
52 + enforce_eager=True, # For faster initialization.
53 + speculative_config={
54 + "method": "eagle",
55 + "model": "yuhuili/EAGLE-LLaMA3-Instruct-8B",
56 + "num_speculative_tokens": num_speculative_tokens,
57 + "max_model_len": 80,
58 + },
59 + max_model_len=200,
60 + attention_config={"backend": attn_backend},
61 + )
62 + sampling_params = SamplingParams(max_tokens=200, ignore_eos=True)
63 + outputs = llm.generate(_PROMPTS, sampling_params)
64 + for o in outputs:
65 + assert o.outputs[0].finish_reason == "length", (
66 + "This test is only meaningful if the output is truncated due to max length"
67 )
64 - sampling_params = SamplingParams(max_tokens=200, ignore_eos=True)
65 - outputs = llm.generate(_PROMPTS, sampling_params)
66 - for o in outputs:
67 - assert o.outputs[0].finish_reason == "length", (
68 - "This test is only meaningful if the output "
69 - "is truncated due to max length"
70 - )
68
72 - sampling_params = SamplingParams(
73 - max_tokens=200,
74 - structured_outputs=StructuredOutputsParams(
75 - regex="^" + "a b c d e " * 15 + "$"
76 - ),
69 + sampling_params = SamplingParams(
70 + max_tokens=200,
71 + structured_outputs=StructuredOutputsParams(regex="^" + "a b c d e " * 15 + "$"),
72 + )
73 + output = llm.generate(_PROMPTS, sampling_params)
74 + for o in output:
75 + assert o.prompt_token_ids is not None
76 + assert (
77 + len(o.prompt_token_ids)
78 + < 80
79 + < len(o.prompt_token_ids) + len(o.outputs[0].token_ids)
80 + <= 200
81 + ), (
82 + "This test is only meaningful if the output "
83 + "is longer than the eagle max length"
84 )
78 - output = llm.generate(_PROMPTS, sampling_params)
79 - for o in output:
80 - assert o.prompt_token_ids is not None
81 - assert (
82 - len(o.prompt_token_ids)
83 - < 80
84 - < len(o.prompt_token_ids) + len(o.outputs[0].token_ids)
85 - <= 200
86 - ), (
87 - "This test is only meaningful if the output "
88 - "is longer than the eagle max length"
89 - )
90 - assert o.outputs[0].text == "a b c d e " * 15
85 + assert o.outputs[0].text == "a b c d e " * 15
vllm/v1/attention/backends/rocm_attn.py
+1 -1
@@ -165,7 +165,7 @@ class RocmAttentionBackend(AttentionBackend):
165 raise ValueError(
166 f"Head size {head_size} is not supported by {attn_type}. "
167 f"Supported head sizes are: {cls.get_supported_head_sizes()}. "
168 - "Set --attention-config.backend=FLEX_ATTENTION to use "
168 + "Set --attention-backend=FLEX_ATTENTION to use "
169 "FlexAttention backend which supports all head sizes."
170 )
171