@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 66e674cdd5
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, 23:45 UTC

[Attention][UX][1/N] Add AttentionConfig and change attention env vars to CLI arguments (#26315)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: Matthew Bonanni <mbonanni001@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>

Matthew Bonanni committed Dec 5, 2025 at 12:48 UTC 66e674cdd549e89b341fb33405a2c02ec5f5ae8d
22 files changed +374 -332
tests/compile/test_fusion_attn.py
+2 -3
@@ -12,13 +12,13 @@ from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
12 from vllm.attention.backends.abstract import AttentionMetadata
13 from vllm.attention.backends.registry import AttentionBackendEnum
14 from vllm.attention.layer import Attention
15 -from vllm.attention.selector import global_force_attn_backend_context_manager
15 from vllm.compilation.fusion_attn import ATTN_OP, AttnFusionPass
16 from vllm.compilation.fx_utils import find_op_nodes
17 from vllm.compilation.matcher_utils import QUANT_OPS
18 from vllm.compilation.noop_elimination import NoOpEliminationPass
19 from vllm.compilation.post_cleanup import PostCleanupPass
20 from vllm.config import (
21 + AttentionConfig,
22 CacheConfig,
23 CompilationConfig,
24 CompilationMode,
@@ -335,6 +335,7 @@ def test_attention_quant_pattern(
335 custom_ops=custom_ops_list,
336 ),
337 cache_config=CacheConfig(cache_dtype="fp8"),
338 + attention_config=AttentionConfig(backend=backend),
339 )
340
341 # Create test inputs
@@ -352,7 +353,6 @@ def test_attention_quant_pattern(
353 with (
354 set_current_vllm_config(vllm_config_unfused),
355 set_forward_context(attn_metadata=None, vllm_config=vllm_config_unfused),
355 - global_force_attn_backend_context_manager(backend),
356 ):
357 model_unfused = model_class(
358 num_qo_heads=num_qo_heads,
@@ -378,7 +378,6 @@ def test_attention_quant_pattern(
378 with (
379 set_current_vllm_config(vllm_config),
380 set_forward_context(attn_metadata=None, vllm_config=vllm_config),
381 - global_force_attn_backend_context_manager(backend),
381 ):
382 model_fused = model_class(
383 num_qo_heads=num_qo_heads,
tests/v1/kv_connector/unit/test_nixl_connector.py
+30 -9
@@ -1151,13 +1151,29 @@ def test_register_kv_caches(dist_init, attn_backend, monkeypatch):
1151 }
1152
1153 # Store tensor info for validation
1154 - expected_tensor_size = shared_tensor[0].element_size() * shared_tensor[0].numel()
1155 - expected_base_addrs = [
1156 - shared_tensor[0].data_ptr(),
1157 - shared_tensor[1].data_ptr(),
1158 - unique_tensor[0].data_ptr(),
1159 - unique_tensor[1].data_ptr(),
1160 - ]
1154 + test_shape = backend_cls.get_kv_cache_shape(
1155 + num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
1156 + )
1157 + is_blocks_first = len(test_shape) == 5 and test_shape[0] == 1
1158 +
1159 + if is_blocks_first:
1160 + expected_tensor_size = shared_tensor.element_size() * shared_tensor.numel()
1161 + expected_base_addrs = [
1162 + shared_tensor.data_ptr(),
1163 + unique_tensor.data_ptr(),
1164 + ]
1165 + expected_num_entries = 2
1166 + else:
1167 + expected_tensor_size = (
1168 + shared_tensor[0].element_size() * shared_tensor[0].numel()
1169 + )
1170 + expected_base_addrs = [
1171 + shared_tensor[0].data_ptr(),
1172 + shared_tensor[1].data_ptr(),
1173 + unique_tensor[0].data_ptr(),
1174 + unique_tensor[1].data_ptr(),
1175 + ]
1176 + expected_num_entries = 4
1177
1178 with (
1179 patch(
@@ -1192,7 +1208,7 @@ def test_register_kv_caches(dist_init, attn_backend, monkeypatch):
1208 # Verify get_reg_descs was called with caches_data
1209 assert mock_wrapper_instance.get_reg_descs.called
1210 caches_data, _ = mock_wrapper_instance.get_reg_descs.call_args[0]
1195 - assert len(caches_data) == 4
1211 + assert len(caches_data) == expected_num_entries
1212
1213 for i, cache_entry in enumerate(caches_data):
1214 base_addr, size, _tp_rank, _ = cache_entry
@@ -1214,7 +1230,12 @@ def test_register_kv_caches(dist_init, attn_backend, monkeypatch):
1230 f"Expected {expected_blocks_count} blocks, got {len(blocks_data)}"
1231 )
1232
1217 - expected_block_len = expected_tensor_size // 2
1233 + num_blocks = 2
1234 + if is_blocks_first:
1235 + expected_block_len = expected_tensor_size // num_blocks // 2
1236 + else:
1237 + expected_block_len = expected_tensor_size // num_blocks
1238 +
1239 for i, block_entry in enumerate(blocks_data):
1240 block_start_addr, block_len, tp_rank = block_entry
1241 assert block_len == expected_block_len, (
tests/v1/worker/test_gpu_model_runner.py
+86 -86
@@ -6,8 +6,10 @@ import pytest
6 import torch
7
8 from vllm.attention.backends.abstract import MultipleOf
9 +from vllm.attention.backends.registry import AttentionBackendEnum
10 from vllm.attention.layer import Attention
11 from vllm.config import (
12 + AttentionConfig,
13 CacheConfig,
14 ModelConfig,
15 ParallelConfig,
@@ -765,7 +767,7 @@ def test_init_kv_cache_with_kv_sharing_valid():
767 current_platform.is_rocm(),
768 reason="Attention backend FLASHINFER is not supported on ROCm.",
769 )
768 -def test_hybrid_attention_mamba_tensor_shapes(monkeypatch):
770 +def test_hybrid_attention_mamba_tensor_shapes():
771 """
772 The GPU model runner creates different views into the
773 KVCacheTensors for the attention and mamba layers
@@ -806,11 +808,13 @@ def test_hybrid_attention_mamba_tensor_shapes(monkeypatch):
808 cache_dtype="auto",
809 )
810 parallel_config = ParallelConfig()
811 + attention_config = AttentionConfig(backend=AttentionBackendEnum.FLASHINFER)
812 vllm_config = VllmConfig(
813 model_config=model_config,
814 cache_config=cache_config,
815 scheduler_config=scheduler_config,
816 parallel_config=parallel_config,
817 + attention_config=attention_config,
818 )
819
820 layer_0 = "model.layers.0.self_attn.attn"
@@ -820,8 +824,7 @@ def test_hybrid_attention_mamba_tensor_shapes(monkeypatch):
824 layer_4 = "model.layers.4.mixer"
825 layer_5 = "model.layers.5.mixer"
826
823 - with set_current_vllm_config(vllm_config), monkeypatch.context() as m:
824 - m.setenv("VLLM_ATTENTION_BACKEND", "FLASHINFER")
827 + with set_current_vllm_config(vllm_config):
828 hf_config = vllm_config.model_config.hf_config
829 fwd_context = {}
830 for key in [layer_0, layer_1]:
@@ -851,10 +854,7 @@ def test_hybrid_attention_mamba_tensor_shapes(monkeypatch):
854 )
855 # suppress var not used error
856 assert fwd_context is not None
854 - vllm_ctx = vllm_config.compilation_config.static_forward_context
855 -
856 - with monkeypatch.context() as m:
857 - m.setenv("VLLM_ATTENTION_BACKEND", "FLASHINFER")
857 + vllm_ctx = vllm_config.compilation_config.static_forward_context
858
859 runner = GPUModelRunner(vllm_config, DEVICE)
860 kv_cache_spec = runner.get_kv_cache_spec()
@@ -865,94 +865,94 @@ def test_hybrid_attention_mamba_tensor_shapes(monkeypatch):
865 )[0]
866 runner.initialize_kv_cache(kv_cache_config)
867
868 - # random partition of blocks
869 - # blocks0 will be assigned to attention layers
870 - # blocks1 will be assigned to mamba layers
871 - num_blocks = kv_cache_config.num_blocks
872 - ind = np.arange(num_blocks)
873 - np.random.shuffle(ind)
874 - blocks0, blocks1 = ind[: (num_blocks // 2)], ind[(num_blocks // 2) :]
868 + # random partition of blocks
869 + # blocks0 will be assigned to attention layers
870 + # blocks1 will be assigned to mamba layers
871 + num_blocks = kv_cache_config.num_blocks
872 + ind = np.arange(num_blocks)
873 + np.random.shuffle(ind)
874 + blocks0, blocks1 = ind[: (num_blocks // 2)], ind[(num_blocks // 2) :]
875
876 - attn_shape = vllm_ctx[layer_0].kv_cache[0].shape
877 - conv_shape = vllm_ctx[layer_2].kv_cache[0][0].shape
878 - ssm_shape = vllm_ctx[layer_2].kv_cache[0][1].shape
876 + attn_shape = vllm_ctx[layer_0].kv_cache[0].shape
877 + conv_shape = vllm_ctx[layer_2].kv_cache[0][0].shape
878 + ssm_shape = vllm_ctx[layer_2].kv_cache[0][1].shape
879
880 - # assert we are using FlashInfer
881 - assert attn_shape[0] % num_blocks == 0
882 - block_split_ratio = attn_shape[0] // num_blocks
880 + # assert we are using FlashInfer
881 + assert attn_shape[0] % num_blocks == 0
882 + block_split_ratio = attn_shape[0] // num_blocks
883
884 - # use small blocks for testing to avoid memory issues
885 - test_block_size = min(2, len(blocks0), len(blocks1))
884 + # use small blocks for testing to avoid memory issues
885 + test_block_size = min(2, len(blocks0), len(blocks1))
886
887 - # use non-overlapping blocks to avoid data contamination
888 - # Split kernel blocks: first half for attention, second half for mamba
889 - mid_point = num_blocks // 2
887 + # use non-overlapping blocks to avoid data contamination
888 + # Split kernel blocks: first half for attention, second half for mamba
889 + mid_point = num_blocks // 2
890
891 - # attention uses kernel blocks from first half (mapped to logical blocks)
892 - kv_blocks_for_attention = np.array([0, 1])[:test_block_size]
891 + # attention uses kernel blocks from first half (mapped to logical blocks)
892 + kv_blocks_for_attention = np.array([0, 1])[:test_block_size]
893
894 - # mamba uses kernel blocks from second half
895 - kv_blocks_for_mamba = np.array([mid_point, mid_point + 1])[:test_block_size]
894 + # mamba uses kernel blocks from second half
895 + kv_blocks_for_mamba = np.array([mid_point, mid_point + 1])[:test_block_size]
896
897 - # create small constant tensors for testing with corrected shapes
898 - # attention: [block_size, ...] starting from dimension 2
899 - attn_constant_shape = attn_shape[2:]
900 - conv_constant_shape = conv_shape[1:]
901 - ssm_constant_shape = ssm_shape[1:]
897 + # create small constant tensors for testing with corrected shapes
898 + # attention: [block_size, ...] starting from dimension 2
899 + attn_constant_shape = attn_shape[2:]
900 + conv_constant_shape = conv_shape[1:]
901 + ssm_constant_shape = ssm_shape[1:]
902
903 - attn_blocks_constant = torch.full(
904 - (test_block_size, *attn_constant_shape), device=DEVICE, fill_value=3.33
905 - )
906 - conv_blocks_constant = torch.full(
907 - (test_block_size, *conv_constant_shape), device=DEVICE, fill_value=6.66
908 - )
909 - ssm_blocks_constant = torch.full(
910 - (test_block_size, *ssm_constant_shape), device=DEVICE, fill_value=9.99
911 - )
903 + attn_blocks_constant = torch.full(
904 + (test_block_size, *attn_constant_shape), device=DEVICE, fill_value=3.33
905 + )
906 + conv_blocks_constant = torch.full(
907 + (test_block_size, *conv_constant_shape), device=DEVICE, fill_value=6.66
908 + )
909 + ssm_blocks_constant = torch.full(
910 + (test_block_size, *ssm_constant_shape), device=DEVICE, fill_value=9.99
911 + )
912
913 - # Fill attention blocks with constants using kv block indices
914 - kernel_blocks_for_attention = kv_blocks_for_attention * block_split_ratio
915 -
916 - for layer in [layer_0, layer_1]:
917 - # attention: kv_cache[0][kernel_block_idx, kv_idx, ...]
918 - for i, kernel_block in enumerate(kernel_blocks_for_attention):
919 - vllm_ctx[layer].kv_cache[0][kernel_block, :] = attn_blocks_constant[i]
920 -
921 - # fill mamba blocks with constants using kernel block indices
922 - for layer in [layer_2, layer_3, layer_4, layer_5]:
923 - # mamba: kv_cache[0][component][kernel_block_idx, ...]
924 - for i, kv_block in enumerate(kv_blocks_for_mamba):
925 - vllm_ctx[layer].kv_cache[0][0][kv_block, :] = conv_blocks_constant[i]
926 - vllm_ctx[layer].kv_cache[0][1][kv_block, :] = ssm_blocks_constant[i]
927 -
928 - # verify attention and mamba contents are correct
929 - for layer in [layer_0, layer_1]:
930 - for i, kernel_block in enumerate(kernel_blocks_for_attention):
931 - actual_kv = vllm_ctx[layer].kv_cache[0][kernel_block, :]
932 - expected = attn_blocks_constant[i]
933 -
934 - # Check K and V separately
935 - assert torch.equal(actual_kv[0], expected)
936 - assert torch.equal(actual_kv[1], expected)
937 -
938 - for layer in [layer_2, layer_3, layer_4, layer_5]:
939 - for i, kv_block in enumerate(kv_blocks_for_mamba):
940 - actual_conv = vllm_ctx[layer].kv_cache[0][0][kv_block, :]
941 - actual_ssm = vllm_ctx[layer].kv_cache[0][1][kv_block, :]
942 - expected_conv = conv_blocks_constant[i]
943 - expected_ssm = ssm_blocks_constant[i]
944 -
945 - assert torch.equal(actual_conv, expected_conv)
946 - assert torch.equal(actual_ssm, expected_ssm)
947 -
948 - for layer in [layer_2, layer_3, layer_4, layer_5]:
949 - for i, kv_block in enumerate(kv_blocks_for_mamba):
950 - actual_conv = vllm_ctx[layer].kv_cache[0][0][kv_block, :]
951 - actual_ssm = vllm_ctx[layer].kv_cache[0][1][kv_block, :]
952 - expected_conv = conv_blocks_constant[i]
953 - expected_ssm = ssm_blocks_constant[i]
954 - assert torch.equal(actual_conv, expected_conv)
955 - assert torch.equal(actual_ssm, expected_ssm)
913 + # Fill attention blocks with constants using kv block indices
914 + kernel_blocks_for_attention = kv_blocks_for_attention * block_split_ratio
915 +
916 + for layer in [layer_0, layer_1]:
917 + # attention: kv_cache[0][kernel_block_idx, kv_idx, ...]
918 + for i, kernel_block in enumerate(kernel_blocks_for_attention):
919 + vllm_ctx[layer].kv_cache[0][kernel_block, :] = attn_blocks_constant[i]
920 +
921 + # fill mamba blocks with constants using kernel block indices
922 + for layer in [layer_2, layer_3, layer_4, layer_5]:
923 + # mamba: kv_cache[0][component][kernel_block_idx, ...]
924 + for i, kv_block in enumerate(kv_blocks_for_mamba):
925 + vllm_ctx[layer].kv_cache[0][0][kv_block, :] = conv_blocks_constant[i]
926 + vllm_ctx[layer].kv_cache[0][1][kv_block, :] = ssm_blocks_constant[i]
927 +
928 + # verify attention and mamba contents are correct
929 + for layer in [layer_0, layer_1]:
930 + for i, kernel_block in enumerate(kernel_blocks_for_attention):
931 + actual_kv = vllm_ctx[layer].kv_cache[0][kernel_block, :]
932 + expected = attn_blocks_constant[i]
933 +
934 + # Check K and V separately
935 + assert torch.equal(actual_kv[0], expected)
936 + assert torch.equal(actual_kv[1], expected)
937 +
938 + for layer in [layer_2, layer_3, layer_4, layer_5]:
939 + for i, kv_block in enumerate(kv_blocks_for_mamba):
940 + actual_conv = vllm_ctx[layer].kv_cache[0][0][kv_block, :]
941 + actual_ssm = vllm_ctx[layer].kv_cache[0][1][kv_block, :]
942 + expected_conv = conv_blocks_constant[i]
943 + expected_ssm = ssm_blocks_constant[i]
944 +
945 + assert torch.equal(actual_conv, expected_conv)
946 + assert torch.equal(actual_ssm, expected_ssm)
947 +
948 + for layer in [layer_2, layer_3, layer_4, layer_5]:
949 + for i, kv_block in enumerate(kv_blocks_for_mamba):
950 + actual_conv = vllm_ctx[layer].kv_cache[0][0][kv_block, :]
951 + actual_ssm = vllm_ctx[layer].kv_cache[0][1][kv_block, :]
952 + expected_conv = conv_blocks_constant[i]
953 + expected_ssm = ssm_blocks_constant[i]
954 + assert torch.equal(actual_conv, expected_conv)
955 + assert torch.equal(actual_ssm, expected_ssm)
956
957
958 def test_hybrid_block_table_initialization():
vllm/attention/backends/abstract.py
+10 -16
@@ -289,6 +289,16 @@ class AttentionImpl(ABC, Generic[T]):
289 # even if they can return lse (for efficiency reasons)
290 need_to_return_lse_for_decode: bool = False
291
292 + # Whether this attention implementation supports pre-quantized query input.
293 + # When True, the attention layer will quantize queries before passing them
294 + # to this backend, allowing torch.compile to fuse the quantization with
295 + # previous operations. This is typically supported when using FP8 KV cache
296 + # with compatible attention kernels (e.g., TRT-LLM).
297 + # Subclasses should set this in __init__.
298 + # TODO add support to more backends:
299 + # https://github.com/vllm-project/vllm/issues/25584
300 + supports_quant_query_input: bool = False
301 +
302 dcp_world_size: int
303 dcp_rank: int
304
@@ -368,22 +378,6 @@ class AttentionImpl(ABC, Generic[T]):
378 """
379 return False
380
371 - def supports_quant_query_input(self) -> bool:
372 - """
373 - Check if this attention implementation supports pre-quantized query input.
374 -
375 - When True, the attention layer will quantize queries before passing them
376 - to this backend, allowing torch.compile to fuse the quantization with
377 - previous operations. This is typically supported when using FP8 KV cache
378 - with compatible attention kernels (e.g., TRT-LLM).
379 - TODO add support to more backends:
380 - https://github.com/vllm-project/vllm/issues/25584
381 -
382 - Returns:
383 - bool: True if the implementation can accept pre-quantized queries.
384 - """
385 - return False
386 -
381 def process_weights_after_loading(self, act_dtype: torch.dtype):
382 pass
383
vllm/attention/layer.py
+2 -2
@@ -303,7 +303,7 @@ class Attention(nn.Module, AttentionLayerBase):
303 self.query_quant = None
304 if (
305 self.kv_cache_dtype.startswith("fp8")
306 - and self.impl.supports_quant_query_input()
306 + and self.impl.supports_quant_query_input
307 ):
308 self.query_quant = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
309
@@ -338,7 +338,7 @@ class Attention(nn.Module, AttentionLayerBase):
338 assert self.kv_cache_dtype in {"fp8", "fp8_e4m3"}
339
340 # check if query quantization is supported
341 - if self.impl.supports_quant_query_input():
341 + if self.impl.supports_quant_query_input:
342 query, _ = self.query_quant(query, self._q_scale)
343
344 if self.use_output:
vllm/attention/selector.py
+9 -128
@@ -2,19 +2,14 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import inspect
5 -import os
6 -from collections.abc import Generator
7 -from contextlib import contextmanager
5 from functools import cache
6 from typing import cast, get_args
7
8 import torch
9
13 -import vllm.envs as envs
10 from vllm.attention.backends.abstract import AttentionBackend
11 from vllm.attention.backends.registry import (
12 MAMBA_TYPE_TO_BACKEND_MAP,
17 - AttentionBackendEnum,
13 MambaAttentionBackendEnum,
14 )
15 from vllm.config.cache import CacheDType
@@ -24,60 +19,6 @@ from vllm.utils.import_utils import resolve_obj_by_qualname
19 logger = init_logger(__name__)
20
21
27 -def get_env_variable_attn_backend() -> AttentionBackendEnum | None:
28 - """
29 - Get the backend override specified by the vLLM attention
30 - backend environment variable, if one is specified.
31 -
32 - Returns:
33 -
34 - * AttentionBackendEnum value if an override is specified
35 - * None otherwise
36 - """
37 - backend_name = os.environ.get("VLLM_ATTENTION_BACKEND")
38 - if backend_name is None:
39 - return None
40 - if backend_name == "XFORMERS":
41 - raise ValueError(
42 - "Attention backend 'XFORMERS' has been removed (See PR #29262 for "
43 - "details). Please select a supported attention backend."
44 - )
45 - return AttentionBackendEnum[backend_name]
46 -
47 -
48 -# Global state allows a particular choice of backend
49 -# to be forced, overriding the logic which auto-selects
50 -# a backend based on system & workload configuration
51 -# (default behavior if this variable is None)
52 -#
53 -# THIS SELECTION TAKES PRECEDENCE OVER THE
54 -# VLLM_ATTENTION_BACKEND ENVIRONMENT VARIABLE
55 -forced_attn_backend: AttentionBackendEnum | None = None
56 -
57 -
58 -def global_force_attn_backend(attn_backend: AttentionBackendEnum | None) -> None:
59 - """
60 - Force all attention operations to use a specified backend.
61 -
62 - Passing `None` for the argument re-enables automatic
63 - backend selection.,
64 -
65 - Arguments:
66 -
67 - * attn_backend: backend selection (None to revert to auto)
68 - """
69 - global forced_attn_backend
70 - forced_attn_backend = attn_backend
71 -
72 -
73 -def get_global_forced_attn_backend() -> AttentionBackendEnum | None:
74 - """
75 - Get the currently-forced choice of attention backend,
76 - or None if auto-selection is currently enabled.
77 - """
78 - return forced_attn_backend
79 -
80 -
22 def get_attn_backend(
23 head_size: int,
24 dtype: torch.dtype,
@@ -97,7 +38,13 @@ def get_attn_backend(
38 f"Valid values are: {valid_cache_dtypes}"
39 )
40
41 + from vllm.config import get_current_vllm_config
42 +
43 + vllm_config = get_current_vllm_config()
44 + backend_enum = vllm_config.attention_config.backend
45 +
46 return _cached_get_attn_backend(
47 + backend=backend_enum,
48 head_size=head_size,
49 dtype=dtype,
50 kv_cache_dtype=cast(CacheDType | None, kv_cache_dtype),
@@ -111,6 +58,7 @@ def get_attn_backend(
58
59 @cache
60 def _cached_get_attn_backend(
61 + backend,
62 head_size: int,
63 dtype: torch.dtype,
64 kv_cache_dtype: CacheDType | None,
@@ -120,39 +68,6 @@ def _cached_get_attn_backend(
68 use_sparse: bool = False,
69 attn_type: str | None = None,
70 ) -> type[AttentionBackend]:
123 - # Check whether a particular choice of backend was
124 - # previously forced.
125 - #
126 - # THIS SELECTION OVERRIDES THE VLLM_ATTENTION_BACKEND
127 - # ENVIRONMENT VARIABLE.
128 - selected_backend = None
129 - backend_by_global_setting: AttentionBackendEnum | None = (
130 - get_global_forced_attn_backend()
131 - )
132 - if backend_by_global_setting is not None:
133 - selected_backend = backend_by_global_setting
134 - else:
135 - # Check the environment variable and override if specified
136 - backend_by_env_var: str | None = envs.VLLM_ATTENTION_BACKEND
137 - if backend_by_env_var is not None:
138 - if backend_by_env_var.endswith("_VLLM_V1"):
139 - logger.warning(
140 - "The suffix '_VLLM_V1' in the environment variable "
141 - "VLLM_ATTENTION_BACKEND is no longer necessary as "
142 - "V0 backends have been deprecated. "
143 - "Please remove this suffix from your "
144 - "environment variable setting.",
145 - )
146 - backend_by_env_var = backend_by_env_var.removesuffix("_VLLM_V1")
147 - try:
148 - selected_backend = AttentionBackendEnum[backend_by_env_var]
149 - except KeyError as e:
150 - raise ValueError(
151 - f"Invalid attention backend: '{backend_by_env_var}'. Valid "
152 - f"backends are: {list(AttentionBackendEnum.__members__.keys())}"
153 - ) from e
154 -
155 - # get device-specific attn_backend
71 from vllm.platforms import current_platform
72
73 sig = inspect.signature(current_platform.get_attn_backend_cls)
@@ -163,7 +78,7 @@ def _cached_get_attn_backend(
78 "remove it from your plugin code."
79 )
80 attention_cls = current_platform.get_attn_backend_cls(
166 - selected_backend,
81 + backend,
82 head_size,
83 dtype,
84 kv_cache_dtype,
@@ -176,7 +91,7 @@ def _cached_get_attn_backend(
91 )
92 else:
93 attention_cls = current_platform.get_attn_backend_cls(
179 - selected_backend,
94 + backend,
95 head_size,
96 dtype,
97 kv_cache_dtype,
@@ -232,37 +147,3 @@ def _cached_get_mamba_attn_backend(
147
148 mamba_attn_backend = selected_backend.get_class()
149 return mamba_attn_backend
235 -
236 -
237 -@contextmanager
238 -def global_force_attn_backend_context_manager(
239 - attn_backend: AttentionBackendEnum,
240 -) -> Generator[None, None, None]:
241 - """
242 - Globally force a vLLM attention backend override within a
243 - context manager, reverting the global attention backend
244 - override to its prior state upon exiting the context
245 - manager.
246 -
247 - Arguments:
248 -
249 - * attn_backend: attention backend to force
250 -
251 - Returns:
252 -
253 - * Generator
254 - """
255 -
256 - # Save the current state of the global backend override (if any)
257 - original_value = get_global_forced_attn_backend()
258 -
259 - # Globally force the new backend override
260 - global_force_attn_backend(attn_backend)
261 -
262 - # Yield control back to the enclosed code block
263 - try:
264 - yield
265 - finally:
266 - # Revert the original global backend override, if any
267 - global_force_attn_backend(original_value)
268 - _cached_get_attn_backend.cache_clear()
vllm/attention/utils/fa_utils.py
+6 -5
@@ -1,7 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from vllm import envs
4 from vllm.logger import init_logger
5 from vllm.platforms import current_platform
6
@@ -49,10 +48,12 @@ def get_flash_attn_version(requires_alibi: bool = False) -> int | None:
48 3 if (device_capability.major == 9 and is_fa_version_supported(3)) else 2
49 )
50
52 - # 2. override if passed by environment
53 - if envs.VLLM_FLASH_ATTN_VERSION is not None:
54 - assert envs.VLLM_FLASH_ATTN_VERSION in [2, 3]
55 - fa_version = envs.VLLM_FLASH_ATTN_VERSION
51 + # 2. override if passed by environment or config
52 + from vllm.config import get_current_vllm_config
53 +
54 + vllm_config = get_current_vllm_config()
55 + if vllm_config.attention_config.flash_attn_version is not None:
56 + fa_version = vllm_config.attention_config.flash_attn_version
57
58 # 3. fallback for unsupported combinations
59 if device_capability.major == 10 and fa_version == 3:
vllm/config/__init__.py
+3
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from vllm.config.attention import AttentionConfig
5 from vllm.config.cache import CacheConfig
6 from vllm.config.compilation import (
7 CompilationConfig,
@@ -46,6 +47,8 @@ from vllm.config.vllm import (
47 # __all__ should only contain classes and functions.
48 # Types and globals should be imported from their respective modules.
49 __all__ = [
50 + # From vllm.config.attention
51 + "AttentionConfig",
52 # From vllm.config.cache
53 "CacheConfig",
54 # From vllm.config.compilation
vllm/config/attention.py new
+114
@@ -0,0 +1,114 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import Any, Literal
5 +
6 +from pydantic import field_validator
7 +from pydantic.dataclasses import dataclass
8 +
9 +from vllm.attention.backends.registry import AttentionBackendEnum
10 +from vllm.config.utils import config
11 +from vllm.logger import init_logger
12 +
13 +logger = init_logger(__name__)
14 +
15 +
16 +@config
17 +@dataclass
18 +class AttentionConfig:
19 + """Configuration for attention mechanisms in vLLM."""
20 +
21 + backend: AttentionBackendEnum | None = None
22 + """Attention backend to use. If None, will be selected automatically."""
23 +
24 + flash_attn_version: Literal[2, 3] | None = None
25 + """Force vllm to use a specific flash-attention version (2 or 3).
26 + Only valid when using the flash-attention backend."""
27 +
28 + use_prefill_decode_attention: bool = False
29 + """Use separate prefill and decode kernels for attention instead of
30 + the unified triton kernel."""
31 +
32 + flash_attn_max_num_splits_for_cuda_graph: int = 32
33 + """Flash Attention max number splits for cuda graph decode."""
34 +
35 + use_cudnn_prefill: bool = False
36 + """Whether to use cudnn prefill."""
37 +
38 + use_trtllm_ragged_deepseek_prefill: bool = False
39 + """Whether to use TRTLLM ragged deepseek prefill."""
40 +
41 + use_trtllm_attention: bool | None = None
42 + """If set to True/False, use or don't use the TRTLLM attention backend
43 + in flashinfer. If None, auto-detect the attention backend in flashinfer."""
44 +
45 + disable_flashinfer_prefill: bool = False
46 + """Whether to disable flashinfer prefill."""
47 +
48 + disable_flashinfer_q_quantization: bool = False
49 + """If set, when using fp8 kv, do not quantize Q to fp8."""
50 +
51 + def compute_hash(self) -> str:
52 + """
53 + Provide a hash that uniquely identifies all the configs
54 + that affect the structure of the computation
55 + graph from input ids/embeddings to the final hidden states,
56 + excluding anything before input ids/embeddings and after
57 + the final hidden states.
58 + """
59 + from vllm.config.utils import get_hash_factors, hash_factors
60 +
61 + ignored_factors: list[str] = []
62 + factors = get_hash_factors(self, ignored_factors)
63 + return hash_factors(factors)
64 +
65 + @field_validator("backend", mode="before")
66 + @classmethod
67 + def validate_backend_before(cls, value: Any) -> Any:
68 + """Enable parsing of the `backend` enum type from string."""
69 + if isinstance(value, str):
70 + return AttentionBackendEnum[value.upper()]
71 + return value
72 +
73 + def _set_from_env_if_set(self, field_name: str, env_var_name: str) -> None:
74 + """Set field from env var if set, with deprecation warning."""
75 + from vllm import envs
76 +
77 + if envs.is_set(env_var_name):
78 + value = getattr(envs, env_var_name)
79 + if field_name == "backend":
80 + value = self.validate_backend_before(value)
81 + setattr(self, field_name, value)
82 + logger.warning_once(
83 + "Using %s environment variable is deprecated and will be removed in "
84 + "v0.14.0 or v1.0.0, whichever is soonest. Please use "
85 + "--attention-config.%s command line argument or "
86 + "AttentionConfig(%s=...) config field instead.",
87 + env_var_name,
88 + field_name,
89 + field_name,
90 + )
91 +
92 + def __post_init__(self) -> None:
93 + self._set_from_env_if_set("backend", "VLLM_ATTENTION_BACKEND")
94 + self._set_from_env_if_set("flash_attn_version", "VLLM_FLASH_ATTN_VERSION")
95 + self._set_from_env_if_set(
96 + "use_prefill_decode_attention", "VLLM_V1_USE_PREFILL_DECODE_ATTENTION"
97 + )
98 + self._set_from_env_if_set(
99 + "flash_attn_max_num_splits_for_cuda_graph",
100 + "VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH",
101 + )
102 + self._set_from_env_if_set("use_cudnn_prefill", "VLLM_USE_CUDNN_PREFILL")
103 + self._set_from_env_if_set(
104 + "use_trtllm_ragged_deepseek_prefill",
105 + "VLLM_USE_TRTLLM_RAGGED_DEEPSEEK_PREFILL",
106 + )
107 + self._set_from_env_if_set("use_trtllm_attention", "VLLM_USE_TRTLLM_ATTENTION")
108 + self._set_from_env_if_set(
109 + "disable_flashinfer_prefill", "VLLM_DISABLE_FLASHINFER_PREFILL"
110 + )
111 + self._set_from_env_if_set(
112 + "disable_flashinfer_q_quantization",
113 + "VLLM_FLASHINFER_DISABLE_Q_QUANTIZATION",
114 + )
vllm/config/model.py
-13
@@ -4,7 +4,6 @@
4 import warnings
5 from collections.abc import Callable
6 from dataclasses import InitVar, field
7 -from importlib.util import find_spec
7 from typing import TYPE_CHECKING, Any, Literal, cast, get_args
8
9 import torch
@@ -467,18 +466,6 @@ class ModelConfig:
466
467 self.maybe_pull_model_tokenizer_for_runai(self.model, self.tokenizer)
468
470 - if (
471 - (backend := envs.VLLM_ATTENTION_BACKEND)
472 - and backend == "FLASHINFER"
473 - and find_spec("flashinfer") is None
474 - ):
475 - raise ValueError(
476 - "VLLM_ATTENTION_BACKEND is set to FLASHINFER, but flashinfer "
477 - "module was not found. See "
478 - "https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile " # noqa: E501
479 - "for instructions on how to install it."
480 - )
481 -
469 from vllm.platforms import current_platform
470
471 if self.override_attention_dtype is not None and not current_platform.is_rocm():
vllm/config/vllm.py
+7
@@ -27,6 +27,7 @@ from vllm.transformers_utils.runai_utils import is_runai_obj_uri
27 from vllm.utils import random_uuid
28 from vllm.utils.hashing import safe_hash
29
30 +from .attention import AttentionConfig
31 from .cache import CacheConfig
32 from .compilation import CompilationConfig, CompilationMode, CUDAGraphMode
33 from .device import DeviceConfig
@@ -192,6 +193,8 @@ class VllmConfig:
193 """Device configuration."""
194 load_config: LoadConfig = Field(default_factory=LoadConfig)
195 """Load configuration."""
196 + attention_config: AttentionConfig = Field(default_factory=AttentionConfig)
197 + """Attention configuration."""
198 lora_config: LoRAConfig | None = None
199 """LoRA configuration."""
200 speculative_config: SpeculativeConfig | None = None
@@ -279,6 +282,10 @@ class VllmConfig:
282 vllm_factors.append(self.load_config.compute_hash())
283 else:
284 vllm_factors.append("None")
285 + if self.attention_config:
286 + vllm_factors.append(self.attention_config.compute_hash())
287 + else:
288 + vllm_factors.append("None")
289 if self.lora_config:
290 vllm_factors.append(self.lora_config.compute_hash())
291 else:
vllm/engine/arg_utils.py
+30 -1
@@ -34,6 +34,7 @@ from typing_extensions import TypeIs
34 import vllm.envs as envs
35 from vllm.attention.backends.registry import AttentionBackendEnum
36 from vllm.config import (
37 + AttentionConfig,
38 CacheConfig,
39 CompilationConfig,
40 ConfigType,
@@ -527,6 +528,7 @@ class EngineArgs:
528
529 pooler_config: PoolerConfig | None = ModelConfig.pooler_config
530 compilation_config: CompilationConfig = get_field(VllmConfig, "compilation_config")
531 + attention_config: AttentionConfig = get_field(VllmConfig, "attention_config")
532 worker_cls: str = ParallelConfig.worker_cls
533 worker_extension_cls: str = ParallelConfig.worker_extension_cls
534
@@ -542,6 +544,7 @@ class EngineArgs:
544 )
545 model_impl: str = ModelConfig.model_impl
546 override_attention_dtype: str = ModelConfig.override_attention_dtype
547 + attention_backend: AttentionBackendEnum | None = AttentionConfig.backend
548
549 calculate_kv_scales: bool = CacheConfig.calculate_kv_scales
550 mamba_cache_dtype: MambaDType = CacheConfig.mamba_cache_dtype
@@ -580,6 +583,8 @@ class EngineArgs:
583 # CompilationConfig object
584 if isinstance(self.compilation_config, dict):
585 self.compilation_config = CompilationConfig(**self.compilation_config)
586 + if isinstance(self.attention_config, dict):
587 + self.attention_config = AttentionConfig(**self.attention_config)
588 if isinstance(self.eplb_config, dict):
589 self.eplb_config = EPLBConfig(**self.eplb_config)
590 # Setup plugins
@@ -717,6 +722,16 @@ class EngineArgs:
722 "--pt-load-map-location", **load_kwargs["pt_load_map_location"]
723 )
724
725 + # Attention arguments
726 + attention_kwargs = get_kwargs(AttentionConfig)
727 + attention_group = parser.add_argument_group(
728 + title="AttentionConfig",
729 + description=AttentionConfig.__doc__,
730 + )
731 + attention_group.add_argument(
732 + "--attention-backend", **attention_kwargs["backend"]
733 + )
734 +
735 # Structured outputs arguments
736 structured_outputs_kwargs = get_kwargs(StructuredOutputsConfig)
737 structured_outputs_group = parser.add_argument_group(
@@ -1140,6 +1155,9 @@ class EngineArgs:
1155 vllm_group.add_argument(
1156 "--compilation-config", "-cc", **vllm_kwargs["compilation_config"]
1157 )
1158 + vllm_group.add_argument(
1159 + "--attention-config", "-ac", **vllm_kwargs["attention_config"]
1160 + )
1161 vllm_group.add_argument(
1162 "--additional-config", **vllm_kwargs["additional_config"]
1163 )
@@ -1693,6 +1711,16 @@ class EngineArgs:
1711 if model_config.quantization == "bitsandbytes":
1712 self.quantization = self.load_format = "bitsandbytes"
1713
1714 + # Attention config overrides
1715 + attention_config = copy.deepcopy(self.attention_config)
1716 + if self.attention_backend is not None:
1717 + if attention_config.backend is not None:
1718 + raise ValueError(
1719 + "attention_backend and attention_config.backend "
1720 + "are mutually exclusive"
1721 + )
1722 + attention_config.backend = self.attention_backend
1723 +
1724 load_config = self.create_load_config()
1725
1726 # Pass reasoning_parser into StructuredOutputsConfig
@@ -1750,9 +1778,10 @@ class EngineArgs:
1778 parallel_config=parallel_config,
1779 scheduler_config=scheduler_config,
1780 device_config=device_config,
1781 + load_config=load_config,
1782 + attention_config=attention_config,
1783 lora_config=lora_config,
1784 speculative_config=speculative_config,
1755 - load_config=load_config,
1785 structured_outputs_config=self.structured_outputs_config,
1786 observability_config=observability_config,
1787 compilation_config=compilation_config,
vllm/model_executor/models/config.py
+7 -4
@@ -4,7 +4,7 @@ from copy import deepcopy
4 from math import lcm
5 from typing import TYPE_CHECKING
6
7 -import vllm.envs as envs
7 +from vllm.attention.backends.registry import AttentionBackendEnum
8 from vllm.logger import init_logger
9 from vllm.model_executor.models import ModelRegistry
10 from vllm.platforms import current_platform
@@ -331,6 +331,7 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
331 # Enable FULL_AND_PIECEWISE by default
332 MambaModelConfig.verify_and_update_config(vllm_config)
333
334 + attention_config = vllm_config.attention_config
335 cache_config = vllm_config.cache_config
336 model_config = vllm_config.model_config
337 parallel_config = vllm_config.parallel_config
@@ -347,7 +348,9 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
348 # * CUTLASS_MLA backend: kernel_block_size 128 alignment
349 # * Other MLA backends: kernel_block_size 64 alignment
350 if model_config.use_mla:
350 - use_cutlass_mla = envs.VLLM_ATTENTION_BACKEND == "CUTLASS_MLA"
351 + use_cutlass_mla = (
352 + attention_config.backend == AttentionBackendEnum.CUTLASS_MLA
353 + )
354 kernel_block_alignment_size = 128 if use_cutlass_mla else 64
355 attn_page_size_1_token = MLAAttentionSpec(
356 block_size=1,
@@ -361,8 +364,8 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
364 current_platform.is_device_capability(100)
365 and model_config.get_head_size() == 256
366 and (
364 - envs.VLLM_ATTENTION_BACKEND is None
365 - or envs.VLLM_ATTENTION_BACKEND == "FLASHINFER"
367 + attention_config.backend is None
368 + or attention_config.backend == AttentionBackendEnum.FLASHINFER
369 )
370 ):
371 # https://github.com/flashinfer-ai/flashinfer/issues/1993 reports that`
vllm/model_executor/models/vision.py
+2 -5
@@ -11,7 +11,7 @@ import torch
11 from transformers import PretrainedConfig
12
13 from vllm.attention.backends.registry import AttentionBackendEnum
14 -from vllm.config import VllmConfig
14 +from vllm.config import VllmConfig, get_current_vllm_config
15 from vllm.distributed import (
16 get_tensor_model_parallel_rank,
17 get_tensor_model_parallel_world_size,
@@ -91,10 +91,7 @@ def get_vit_attn_backend(
91 if attn_backend_override is not None:
92 return attn_backend_override
93
94 - # Lazy import to avoid circular dependency
95 - from vllm.attention.selector import get_env_variable_attn_backend
96 -
97 - selected_backend: AttentionBackendEnum | None = get_env_variable_attn_backend()
94 + selected_backend = get_current_vllm_config().attention_config.backend
95 if selected_backend is not None:
96 return selected_backend
97
vllm/platforms/cuda.py
+13 -10
@@ -14,7 +14,6 @@ from typing_extensions import ParamSpec
14
15 # import custom ops, trigger op registration
16 import vllm._C # noqa
17 -import vllm.envs as envs
17 from vllm.attention.backends.abstract import AttentionType
18 from vllm.attention.backends.registry import AttentionBackendEnum
19 from vllm.logger import init_logger
@@ -149,6 +148,8 @@ class CudaPlatformBase(Platform):
148
149 @classmethod
150 def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
151 + from vllm.attention.backends.registry import AttentionBackendEnum
152 +
153 parallel_config = vllm_config.parallel_config
154 model_config = vllm_config.model_config
155
@@ -171,7 +172,7 @@ class CudaPlatformBase(Platform):
172 and cache_config.block_size is not None
173 ):
174 use_sparse = hasattr(vllm_config.model_config.hf_config, "index_topk")
174 - # If `VLLM_ATTENTION_BACKEND` is not set and we are using MLA,
175 + # If `--attention-config.backend` is not set and we are using MLA,
176 # then we default to FlashMLA backend for non-blackwell GPUs,
177 # else we default to CutlassMLA. For each case, we force the
178 # required block_size.
@@ -179,23 +180,25 @@ class CudaPlatformBase(Platform):
180 use_cutlass_mla = False
181 use_flashinfer_mla = False
182
182 - if envs.VLLM_ATTENTION_BACKEND is None:
183 + if vllm_config.attention_config.backend is None:
184 # Default case
185 if cls.is_device_capability(100):
186 # Blackwell => Force CutlassMLA.
187 use_cutlass_mla = True
187 - # TODO: This does not work, because the
188 - # global_force_attn_backend_context_manager is not set.
189 - # See vllm/attention/selector.py:_cached_get_attn_backend
190 - envs.VLLM_ATTENTION_BACKEND = "CUTLASS_MLA"
188 + # Set the backend in AttentionConfig so it's used during
189 + # backend selection
190 + vllm_config.attention_config.backend = (
191 + AttentionBackendEnum.CUTLASS_MLA
192 + )
193 else:
194 # Not Blackwell
195 use_flashmla = True
196 else:
197 # Forced case
196 - use_flashmla = envs.VLLM_ATTENTION_BACKEND == "FLASHMLA"
197 - use_cutlass_mla = envs.VLLM_ATTENTION_BACKEND == "CUTLASS_MLA"
198 - use_flashinfer_mla = envs.VLLM_ATTENTION_BACKEND == "FLASHINFER_MLA"
198 + backend = vllm_config.attention_config.backend
199 + use_flashmla = backend == AttentionBackendEnum.FLASHMLA
200 + use_cutlass_mla = backend == AttentionBackendEnum.CUTLASS_MLA
201 + use_flashinfer_mla = backend == AttentionBackendEnum.FLASHINFER_MLA
202
203 from vllm.attention.ops.flashmla import is_flashmla_dense_supported
204
vllm/utils/flashinfer.py
+14 -23
@@ -267,21 +267,16 @@ def supports_trtllm_attention() -> bool:
267 return current_platform.is_device_capability(100) and has_nvidia_artifactory()
268
269
270 -@functools.cache
271 -def _force_use_trtllm_attention(env_value: bool | None) -> bool | None:
272 - """Cache the env value for VLLM_USE_TRTLLM_ATTENTION"""
273 - if env_value is not None:
274 - logger.info_once("VLLM_USE_TRTLLM_ATTENTION is set to %s", env_value)
275 - return env_value
276 -
277 -
270 def force_use_trtllm_attention() -> bool | None:
271 """
280 - Return `None` if VLLM_USE_TRTLLM_ATTENTION is not set,
272 + Return `None` if --attention-config.use_trtllm_attention is not set,
273 return `True` if TRTLLM attention is forced to be used,
274 return `False` if TRTLLM attention is forced to be not used.
275 """
284 - return _force_use_trtllm_attention(envs.VLLM_USE_TRTLLM_ATTENTION)
276 + from vllm.config import get_current_vllm_config
277 +
278 + vllm_config = get_current_vllm_config()
279 + return vllm_config.attention_config.use_trtllm_attention
280
281
282 def can_use_trtllm_attention(num_qo_heads: int, num_kv_heads: int) -> bool:
@@ -307,7 +302,7 @@ def use_trtllm_attention(
302 """Return `True` if TRTLLM attention is used."""
303 force_use_trtllm = force_use_trtllm_attention()
304
310 - # Environment variable is set to 0 - respect it
305 + # CLI argument is set to 0 - respect it
306 if force_use_trtllm is not None and not force_use_trtllm:
307 return False
308
@@ -324,7 +319,7 @@ def use_trtllm_attention(
319 if force_use_trtllm:
320 logger.warning_once(
321 "TRTLLM attention is not supported on this platform, "
327 - "but VLLM_USE_TRTLLM_ATTENTION is set to 1"
322 + "but --attention-config.use_trtllm_attention is set to 1"
323 )
324 return False
325
@@ -333,7 +328,8 @@ def use_trtllm_attention(
328 if force_use_trtllm:
329 logger.warning_once(
330 "TRTLLM attention is not supported for this combination of "
336 - "query and key heads, but VLLM_USE_TRTLLM_ATTENTION is set to 1"
331 + "query and key heads, but --attention-config.use_trtllm_attention is "
332 + "set to 1"
333 )
334 return False
335
@@ -354,7 +350,7 @@ def use_trtllm_attention(
350 return True
351
352 if force_use_trtllm is None:
357 - # Environment variable not set - use auto-detection
353 + # CLI argument not set - use auto-detection
354 if is_prefill:
355 # Prefill auto-detection
356 use_trtllm = kv_cache_dtype == "auto"
@@ -367,8 +363,10 @@ def use_trtllm_attention(
363 logger.warning_once("Using TRTLLM decode attention (auto-detected).")
364 return use_trtllm
365
370 - # Environment variable is set to 1 - respect it
371 - logger.info_once("Using TRTLLM attention (VLLM_USE_TRTLLM_ATTENTION is set to 1)")
366 + # CLI argument is set to 1 - respect it
367 + logger.info_once(
368 + "Using TRTLLM attention (--attention-config.use_trtllm_attention is set to 1)"
369 + )
370 return True
371
372
@@ -500,12 +498,6 @@ def flashinfer_scaled_fp8_mm(
498 return output
499
500
503 -@functools.cache
504 -def flashinfer_disable_q_quantization() -> bool:
505 - """Cache result which only depends on the environment"""
506 - return envs.VLLM_FLASHINFER_DISABLE_Q_QUANTIZATION
507 -
508 -
501 __all__ = [
502 "has_flashinfer",
503 "flashinfer_trtllm_fp8_block_scale_moe",
@@ -526,7 +518,6 @@ __all__ = [
518 "supports_trtllm_attention",
519 "can_use_trtllm_attention",
520 "use_trtllm_attention",
529 - "flashinfer_disable_q_quantization",
521 "flashinfer_scaled_fp4_mm",
522 "flashinfer_scaled_fp8_mm",
523 ]
vllm/v1/attention/backends/flash_attn.py
+5 -4
@@ -8,7 +8,6 @@ from typing import ClassVar
8 import numpy as np
9 import torch
10
11 -from vllm import envs
11 from vllm.attention.backends.abstract import (
12 AttentionBackend,
13 AttentionImpl,
@@ -264,6 +263,7 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
263 self.parallel_config = vllm_config.parallel_config
264 self.cache_config = vllm_config.cache_config
265 self.compilation_config = vllm_config.compilation_config
266 + self.attention_config = vllm_config.attention_config
267
268 self.num_heads_q = self.model_config.get_num_attention_heads(
269 self.parallel_config
@@ -304,7 +304,9 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
304 # When using cuda graph, we need to set the upper bound of the
305 # number of splits so that large enough intermediate buffers are
306 # pre-allocated during capture.
307 - self.max_num_splits = envs.VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH
307 + self.max_num_splits = (
308 + self.attention_config.flash_attn_max_num_splits_for_cuda_graph
309 + )
310
311 # Sliding window size to be used with the AOT scheduler will be
312 # populated on first build() call.
@@ -554,8 +556,7 @@ class FlashAttentionImpl(AttentionImpl):
556 "heads in the layer"
557 )
558
557 - def supports_quant_query_input(self) -> bool:
558 - return True
559 + self.supports_quant_query_input = True
560
561 def forward(
562 self,
vllm/v1/attention/backends/flashinfer.py
+14 -12
@@ -26,7 +26,7 @@ from vllm.attention.backends.abstract import (
26 )
27 from vllm.attention.ops.common import cp_lse_ag_out_rs
28 from vllm.attention.ops.merge_attn_states import merge_attn_states
29 -from vllm.config import CUDAGraphMode, VllmConfig
29 +from vllm.config import CUDAGraphMode, VllmConfig, get_current_vllm_config
30 from vllm.config.cache import CacheDType
31 from vllm.distributed.parallel_state import get_dcp_group
32 from vllm.logger import init_logger
@@ -43,7 +43,6 @@ from vllm.platforms.interface import DeviceCapability
43 from vllm.triton_utils import tl, triton
44 from vllm.utils.flashinfer import (
45 can_use_trtllm_attention,
46 - flashinfer_disable_q_quantization,
46 use_trtllm_attention,
47 )
48 from vllm.utils.math_utils import cdiv
@@ -362,7 +361,8 @@ class FlashInferBackend(AttentionBackend):
361 supports_trtllm_attention,
362 )
363
365 - # Respect explicit disable flag (e.g., VLLM_USE_TRTLLM_ATTENTION=0)
364 + # Respect explicit disable flag (e.g.,
365 + # --attention-config.use_trtllm_attention=0)
366 if force_use_trtllm_attention() is False:
367 return False
368
@@ -500,11 +500,14 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
500 self.kv_cache_dtype = self.kv_cache_spec.dtype
501
502 # Use model dtype as q dtype when TRTLLM attn is not supported, or
503 - # VLLM_FLASHINFER_DISABLE_Q_QUANTIZATION is set to 1. Otherwise, try to
504 - # use fp8 q if kv cache is fp8, and will fall back to model dtype
503 + # --attention-config.disable_flashinfer_q_quantization is set to 1. Otherwise,
504 + # try to use fp8 q if kv cache is fp8, and will fall back to model dtype
505 # if TRTLLM attention kernel is not used when building attn metadata
506 can_use_trtllm = can_use_trtllm_attention(self.num_qo_heads, self.num_kv_heads)
507 - if can_use_trtllm and not flashinfer_disable_q_quantization():
507 + if (
508 + can_use_trtllm
509 + and not vllm_config.attention_config.disable_flashinfer_q_quantization
510 + ):
511 self.q_data_type = self.kv_cache_dtype
512 else:
513 self.q_data_type = self.model_config.dtype
@@ -1035,6 +1038,11 @@ class FlashInferImpl(AttentionImpl):
1038 self.sinks = sinks
1039
1040 self.support_trtllm_attn = can_use_trtllm_attention(num_heads, num_kv_heads)
1041 + vllm_config = get_current_vllm_config()
1042 + self.supports_quant_query_input = (
1043 + self.support_trtllm_attn
1044 + and not vllm_config.attention_config.disable_flashinfer_q_quantization
1045 + )
1046 self.bmm1_scale: float | None = None
1047 self.bmm2_scale: float | None = None
1048 self.o_sf_scale: float | None = None
@@ -1046,12 +1054,6 @@ class FlashInferImpl(AttentionImpl):
1054 and quant_key in (kFp8StaticTensorSym, kNvfp4Quant)
1055 )
1056
1049 - def supports_quant_query_input(self) -> bool:
1050 - if flashinfer_disable_q_quantization():
1051 - return False
1052 -
1053 - return self.support_trtllm_attn
1054 -
1057 # FlashInfer requires attention sinks to be float32
1058 def process_weights_after_loading(self, act_dtype: torch.dtype):
1059 if self.sinks is not None and self.sinks.dtype != torch.float32:
vllm/v1/attention/backends/mla/common.py
+14 -5
@@ -438,19 +438,25 @@ A = TypeVar("A")
438 def use_flashinfer_prefill() -> bool:
439 # For blackwell default to flashinfer prefill if it's available since
440 # it is faster than FA2.
441 + from vllm.config import get_current_vllm_config
442 +
443 + vllm_config = get_current_vllm_config()
444 return (
442 - not envs.VLLM_DISABLE_FLASHINFER_PREFILL
445 + not vllm_config.attention_config.disable_flashinfer_prefill
446 and flashinfer_available
444 - and not envs.VLLM_USE_CUDNN_PREFILL
445 - and not envs.VLLM_USE_TRTLLM_RAGGED_DEEPSEEK_PREFILL
447 + and not vllm_config.attention_config.use_cudnn_prefill
448 + and not vllm_config.attention_config.use_trtllm_ragged_deepseek_prefill
449 and current_platform.is_device_capability(100)
450 )
451
452
453 def use_cudnn_prefill() -> bool:
454 + from vllm.config import get_current_vllm_config
455 +
456 + vllm_config = get_current_vllm_config()
457 return (
458 flashinfer_available
453 - and envs.VLLM_USE_CUDNN_PREFILL
459 + and vllm_config.attention_config.use_cudnn_prefill
460 and current_platform.is_device_capability(100)
461 and has_nvidia_artifactory()
462 )
@@ -458,9 +464,12 @@ def use_cudnn_prefill() -> bool:
464
465 def use_trtllm_ragged_deepseek_prefill() -> bool:
466 """Check if TRT-LLM ragged DeepSeek prefill should be used."""
467 + from vllm.config import get_current_vllm_config
468 +
469 + vllm_config = get_current_vllm_config()
470 return (
471 flashinfer_available
463 - and envs.VLLM_USE_TRTLLM_RAGGED_DEEPSEEK_PREFILL
472 + and vllm_config.attention_config.use_trtllm_ragged_deepseek_prefill
473 and current_platform.is_device_capability(100)
474 )
475
vllm/v1/attention/backends/mla/flashattn_mla.py
+3 -2
@@ -6,7 +6,6 @@ from typing import ClassVar
6
7 import torch
8
9 -from vllm import envs
9 from vllm.attention.backends.abstract import (
10 AttentionLayer,
11 AttentionType,
@@ -131,7 +130,9 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
130 # When using cuda graph, we need to set the upper bound of the
131 # number of splits so that large enough intermediate buffers are
132 # pre-allocated during capture.
134 - self.max_num_splits = envs.VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH
133 + self.max_num_splits = (
134 + vllm_config.attention_config.flash_attn_max_num_splits_for_cuda_graph
135 + )
136
137 if vllm_is_batch_invariant():
138 self.max_num_splits = 1
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 VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "
168 + "Set --attention-config.backend=FLEX_ATTENTION to use "
169 "FlexAttention backend which supports all head sizes."
170 )
171
vllm/v1/attention/backends/triton_attn.py
+2 -3
@@ -210,9 +210,6 @@ class TritonAttentionImpl(AttentionImpl):
210 def fused_output_quant_supported(self, quant_key: QuantKey):
211 return quant_key == kFp8StaticTensorSym
212
213 - def supports_quant_query_input(self) -> bool:
214 - return current_platform.is_cuda()
215 -
213 def __init__(
214 self,
215 num_heads: int,
@@ -262,6 +259,8 @@ class TritonAttentionImpl(AttentionImpl):
259 f"num_heads: {num_heads}."
260 )
261
262 + self.supports_quant_query_input = current_platform.is_cuda()
263 +
264 def forward(
265 self,
266 layer: torch.nn.Module,