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 7, 06:30 UTC
[Misc] Add `--attention-backend auto` option (#35738)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Mar 4, 2026 at 16:12 UTC
18e01a0a10e37ed7a705b46373b9b004f03b9e6b
3 files changed
+54
-9
tests/kernels/attention/test_attention_selector.py
+42
@@ -293,6 +293,48 @@ def test_invalid_backend():
293
AttentionConfig(backend=AttentionBackendEnum["INVALID"])
294
295
296
+@pytest.mark.parametrize("auto_value", ["auto", "AUTO", "Auto"])
297
+def test_auto_backend_string(auto_value: str):
298
+ """Test that 'auto' string value triggers automatic backend selection."""
299
+ # Using "auto" should result in backend=None (automatic selection)
300
+ attention_config = AttentionConfig(backend=auto_value)
301
+ assert attention_config.backend is None
302
+
303
+
304
+def test_auto_backend_selection_behavior():
305
+ """Test that 'auto' backend behaves same as None (automatic selection)."""
306
+ # Create config with explicit "auto"
307
+ auto_config = AttentionConfig(backend="auto")
308
+
309
+ # Create config with None (default)
310
+ none_config = AttentionConfig(backend=None)
311
+
312
+ # Both should have backend=None
313
+ assert auto_config.backend is None
314
+ assert none_config.backend is None
315
+
316
+ # Both configs should result in the same automatic backend selection
317
+ vllm_config_auto = VllmConfig(attention_config=auto_config)
318
+ vllm_config_none = VllmConfig(attention_config=none_config)
319
+
320
+ with (
321
+ set_current_vllm_config(vllm_config_auto),
322
+ patch("vllm.platforms.current_platform", CpuPlatform()),
323
+ ):
324
+ backend_auto = get_attn_backend(16, torch.float16, None, 16)
325
+
326
+ _cached_get_attn_backend.cache_clear()
327
+
328
+ with (
329
+ set_current_vllm_config(vllm_config_none),
330
+ patch("vllm.platforms.current_platform", CpuPlatform()),
331
+ ):
332
+ backend_none = get_attn_backend(16, torch.float16, None, 16)
333
+
334
+ # Both should select the same backend
335
+ assert backend_auto.get_name() == backend_none.get_name()
336
+
337
+
338
@pytest.mark.parametrize(
339
"backend_name,flash_attn_version,should_succeed",
340
[
vllm/config/attention.py
+8
-2
@@ -14,7 +14,7 @@ class AttentionConfig:
14
"""Configuration for attention mechanisms in vLLM."""
15
16
backend: AttentionBackendEnum | None = None
17
- """Attention backend to use. If None, will be selected automatically."""
17
+ """Attention backend to use. Use "auto" or None for automatic selection."""
18
19
flash_attn_version: Literal[2, 3, 4] | None = None
20
"""Force vllm to use a specific flash-attention version (2, 3, or 4).
@@ -63,7 +63,13 @@ class AttentionConfig:
63
@field_validator("backend", mode="before")
64
@classmethod
65
def validate_backend_before(cls, value: Any) -> Any:
66
- """Enable parsing of the `backend` enum type from string."""
66
+ """Enable parsing of the `backend` enum type from string.
67
+
68
+ The special value "auto" is treated as None, which triggers
69
+ automatic backend selection.
70
+ """
71
if isinstance(value, str):
72
+ if value.lower() == "auto":
73
+ return None
74
return AttentionBackendEnum[value.upper()]
75
return value
vllm/engine/arg_utils.py
+4
-7
@@ -1816,13 +1816,10 @@ class EngineArgs:
1816
"attention_backend and attention_config.backend "
1817
"are mutually exclusive"
1818
)
1819
- # Convert string to enum if needed (CLI parsing returns a string)
1820
- if isinstance(self.attention_backend, str):
1821
- attention_config.backend = AttentionBackendEnum[
1822
- self.attention_backend.upper()
1823
- ]
1824
- else:
1825
- attention_config.backend = self.attention_backend
1819
+ # Reuse the validator to handle "auto" and string-to-enum conversion
1820
+ attention_config.backend = AttentionConfig.validate_backend_before(
1821
+ self.attention_backend
1822
+ )
1823
1824
# Kernel config overrides
1825
kernel_config = copy.deepcopy(self.kernel_config)