555e7225bc · tduncan/Nvidia-Optimization-Stack-VLLM

tests/v1/attention/test_attention_backends.py

+11 -8

@@ -198,7 +198,8 @@ class MockAttentionLayer:

198

199

200 def run_attention_backend(backend: _Backend, kv_cache_spec: FullAttentionSpec,

201 - vllm_config, device: torch.device,

201 + layer_names: list[str], vllm_config,

202 + device: torch.device,

203 common_attn_metadata: CommonAttentionMetadata,

204 query: torch.Tensor, key: torch.Tensor,

205 value: torch.Tensor,

@@ -211,31 +212,33 @@ def run_attention_backend(backend: _Backend, kv_cache_spec: FullAttentionSpec,

212 if backend == _Backend.FLASHINFER_VLLM_V1:

213 import unittest.mock

214

214 - from vllm.v1.attention.backends.flashinfer import PerLayerParameters

215 + from vllm.v1.attention.backends.utils import PerLayerParameters

216

216 - def mock_get_per_layer_parameters(vllm_config, impl_cls):

217 + def mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls):

218 # Return mock parameters for a single layer

219 head_size = vllm_config.model_config.get_head_size()

220 return {

220 - "mock_layer":

221 + layer_name:

222 PerLayerParameters(

223 window_left=-1, # No sliding window

224 logits_soft_cap=0.0, # No soft cap

225 sm_scale=1.0 / (head_size**0.5) # Standard scale

226 )

227 + for layer_name in layer_names

228 }

229

230 with unittest.mock.patch(

231 'vllm.v1.attention.backends.flashinfer.get_per_layer_parameters',

232 mock_get_per_layer_parameters):

231 - builder = builder_cls(kv_cache_spec, vllm_config, device)

233 + builder = builder_cls(kv_cache_spec, layer_names, vllm_config,

234 + device)

235 attn_metadata = builder.build(

236 common_prefix_len=0,

237 common_attn_metadata=common_attn_metadata,

238 )

239 else:

240 # Build metadata

238 - builder = builder_cls(kv_cache_spec, vllm_config, device)

241 + builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device)

242 attn_metadata = builder.build(

243 common_prefix_len=0,

244 common_attn_metadata=common_attn_metadata,

@@ -427,8 +430,8 @@ def test_backend_correctness(batch_spec_name: str, model: str):

430 set_kv_cache_layout("HND")

431

432 backend_output = run_attention_backend(backend_name, kv_cache_spec,

430 - vllm_config, device,

431 - common_attn_metadata,

433 + ["placeholder"], vllm_config,

434 + device, common_attn_metadata,

435 query_vllm, key_vllm,

436 value_vllm,

437 kv_cache_for_backend)

tests/v1/spec_decode/test_eagle.py

+1

@@ -305,6 +305,7 @@ def test_propose(num_speculative_tokens):

305 _Backend.FLASH_ATTN_VLLM_V1)

306 attn_metadata_builder = attn_metadata_builder_cls(

307 kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),

308 + layer_names=proposer.attn_layer_names,

309 vllm_config=proposer.vllm_config,

310 device=device,

311 )

tests/v1/worker/test_gpu_model_runner.py

+2 -1

@@ -745,7 +745,8 @@ def test_hybrid_attention_mamba_tensor_shapes(monkeypatch):

745 layer_4 = "model.layers.4.mixer"

746 layer_5 = "model.layers.5.mixer"

747

748 - with set_current_vllm_config(vllm_config):

748 + with set_current_vllm_config(vllm_config), monkeypatch.context() as m:

749 + m.setenv("VLLM_ATTENTION_BACKEND", "FLASHINFER")

750 hf_config = vllm_config.model_config.hf_config

751 fwd_context = {}

752 for key in [layer_0, layer_1]:

vllm/config.py

+24 -8

@@ -740,8 +740,8 @@ class ModelConfig:

740 isinstance(sliding_window, list))

741

742 if not self.disable_sliding_window and has_interleaved_attention:

743 - if (backend :=

744 - envs.VLLM_ATTENTION_BACKEND) in ("XFORMERS", "FLASHINFER"):

743 + if not envs.VLLM_USE_V1 and (backend := envs.VLLM_ATTENTION_BACKEND

744 + ) in ("XFORMERS", "FLASHINFER"):

745 sliding_window_len_min = get_min_sliding_window(

746 self.hf_text_config.sliding_window)

747

@@ -5065,13 +5065,29 @@ def assert_hashable(text):

5065 T = TypeVar("T")

5066

5067

5068 -def get_layers_from_vllm_config(vllm_config: VllmConfig,

5069 - layer_type: type[T]) -> dict[str, T]:

5068 +def get_layers_from_vllm_config(

5069 + vllm_config: VllmConfig,

5070 + layer_type: type[T],

5071 + layer_names: Optional[list[str]] = None) -> dict[str, T]:

5072 + """

5073 + Get layers from the vLLM config.

5074 +

5075 + Args:

5076 + vllm_config: The vLLM config.

5077 + layer_type: The type of the layer to get.

5078 + layer_names: The names of the layers to get. If None, return all layers.

5079 + """

5080 +

5081 + if layer_names is None:

5082 + layer_names = list(

5083 + vllm_config.compilation_config.static_forward_context.keys())

5084 +

5085 + forward_context = vllm_config.compilation_config.static_forward_context

5086 +

5087 return {

5071 - layer_name: layer

5072 - for layer_name, layer in

5073 - vllm_config.compilation_config.static_forward_context.items()

5074 - if isinstance(layer, layer_type)

5088 + layer_name: forward_context[layer_name]

5089 + for layer_name in layer_names

5090 + if isinstance(forward_context[layer_name], layer_type)

5091 }

5092

5093

vllm/v1/attention/backends/cpu_attn.py

+2 -2

@@ -315,8 +315,8 @@ class TorchSDPAMetadata(AttentionMetadata):

315

316 class TorchSDPAMetadataBuilderV1(AttentionMetadataBuilder[TorchSDPAMetadata]):

317

318 - def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,

319 - device: torch.device) -> None:

318 + def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],

319 + vllm_config: VllmConfig, device: torch.device) -> None:

320 self.kv_cache_spec = kv_cache_spec

321 self.vllm_config = vllm_config

322 self.scheduler_config = vllm_config.scheduler_config

vllm/v1/attention/backends/flash_attn.py

+2 -2

@@ -148,8 +148,8 @@ class FlashAttentionMetadataBuilder(

148 AttentionMetadataBuilder[FlashAttentionMetadata]):

149 full_cudagraph_supported: ClassVar[bool] = get_flash_attn_version() == 3

150

151 - def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,

152 - device: torch.device):

151 + def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],

152 + vllm_config: VllmConfig, device: torch.device):

153 self.vllm_config = vllm_config

154 self.model_config = vllm_config.model_config

155 self.parallel_config = vllm_config.parallel_config

vllm/v1/attention/backends/flashinfer.py

+7 -11

@@ -21,10 +21,9 @@ from vllm.platforms import current_platform

21 from vllm.utils import cdiv

22 from vllm.v1.attention.backends.flash_attn import use_cascade_attention

23 from vllm.v1.attention.backends.utils import (

24 - AttentionMetadataBuilder, CommonAttentionMetadata, PerLayerParameters,

25 - get_kv_cache_layout, get_per_layer_parameters,

26 - infer_global_hyperparameters, reorder_batch_to_split_decodes_and_prefills,

27 - split_decodes_and_prefills)

24 + AttentionMetadataBuilder, CommonAttentionMetadata, get_kv_cache_layout,

25 + get_per_layer_parameters, infer_global_hyperparameters,

26 + reorder_batch_to_split_decodes_and_prefills, split_decodes_and_prefills)

27 from vllm.v1.kv_cache_interface import AttentionSpec

28

29 if TYPE_CHECKING:

@@ -219,8 +218,8 @@ class FlashInferMetadata:

218

219 class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

220

222 - def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,

223 - device: torch.device):

221 + def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],

222 + vllm_config: VllmConfig, device: torch.device):

223 self.device = device

224 self._workspace_buffer = None

225 self._prefill_wrapper = None # Wrapper for prefill/append

@@ -228,7 +227,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

227 self._cascade_wrapper = None # Wrapper for cascade attention

228

229 # Global hyperparameters shared by all attention layers

231 - self.global_hyperparameters: Optional[PerLayerParameters] = None

230 + self.global_hyperparameters = infer_global_hyperparameters(

231 + get_per_layer_parameters(vllm_config, layer_names, FlashInferImpl))

232

233 self.vllm_config = vllm_config

234 self.cache_config = vllm_config.cache_config

@@ -283,10 +283,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

283

284 def _plan(self, num_prefills: int, num_decodes: int,

285 attn_metadata: FlashInferMetadata):

286 - if self.global_hyperparameters is None:

287 - self.global_hyperparameters = infer_global_hyperparameters(

288 - get_per_layer_parameters(self.vllm_config, FlashInferImpl))

289 -

286 if attn_metadata.use_cascade:

287 attn_metadata.cascade_wrapper = self._get_cascade_wrapper()

288 attn_metadata.cascade_wrapper.plan(

vllm/v1/attention/backends/flex_attention.py

+2 -2

@@ -258,8 +258,8 @@ class FlexAttentionMetadata:

258 class FlexAttentionMetadataBuilder(

259 AttentionMetadataBuilder[FlexAttentionMetadata]):

260

261 - def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,

262 - device: torch.device):

261 + def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],

262 + vllm_config: VllmConfig, device: torch.device):

263 self.model_config = vllm_config.model_config

264 self.parallel_config = vllm_config.parallel_config

265 self.cache_config = vllm_config.cache_config

vllm/v1/attention/backends/mamba_attn.py

+2 -2

@@ -87,8 +87,8 @@ class Mamba2AttentionMetadata:

87 class Mamba2AttentionMetadataBuilder(

88 AttentionMetadataBuilder[Mamba2AttentionMetadata]):

89

90 - def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,

91 - device: torch.device):

90 + def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],

91 + vllm_config: VllmConfig, device: torch.device):

92 assert isinstance(kv_cache_spec, MambaSpec)

93 self.kv_cache_spec = kv_cache_spec

94 self.chunk_size = vllm_config.model_config.get_mamba_chunk_size()

vllm/v1/attention/backends/mla/common.py

+3 -1

@@ -406,6 +406,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

406

407 def __init__(self,

408 kv_cache_spec: AttentionSpec,

409 + layer_names: list[str],

410 vllm_config: VllmConfig,

411 device: torch.device,

412 metadata_cls: Optional[type[M]] = None):

@@ -471,7 +472,8 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

472 BatchPrefillWithRaggedKVCacheWrapper] = []

473

474 self._global_hyperparameters = infer_global_hyperparameters(

474 - get_per_layer_parameters(vllm_config, MLACommonImpl))

475 + get_per_layer_parameters(vllm_config, layer_names,

476 + MLACommonImpl))

477

478 if self._use_cudnn_prefill:

479 self.cudnn_workspace = torch.empty(

vllm/v1/attention/backends/mla/flashmla.py

+4 -3

@@ -56,9 +56,10 @@ class FlashMLAMetadata(MLACommonMetadata[FlashMLADecodeMetadata]):

56 class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):

57 full_cudagraph_supported: ClassVar[bool] = True # Decode-only

58

59 - def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,

60 - device: torch.device):

61 - super().__init__(kv_cache_spec, vllm_config, device, FlashMLAMetadata)

59 + def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],

60 + vllm_config: VllmConfig, device: torch.device):

61 + super().__init__(kv_cache_spec, layer_names, vllm_config, device,

62 + FlashMLAMetadata)

63

64 self.compilation_config = vllm_config.compilation_config

65 self.num_q_heads = vllm_config.model_config.get_num_attention_heads(

vllm/v1/attention/backends/mla/rocm_aiter_mla.py

+4 -3

@@ -66,9 +66,10 @@ class AiterMLAMetadata(MLACommonMetadata[AiterMLADecodeMetadata]):

66 class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

67 full_cudagraph_supported: ClassVar[bool] = True # decode only

68

69 - def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,

70 - device: torch.device):

71 - super().__init__(kv_cache_spec, vllm_config, device, AiterMLAMetadata)

69 + def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],

70 + vllm_config: VllmConfig, device: torch.device):

71 + super().__init__(kv_cache_spec, layer_names, vllm_config, device,

72 + AiterMLAMetadata)

73 assert self.kv_cache_spec.block_size == 1, "AITER MLA" \

74 "only supports block size 1."

75

vllm/v1/attention/backends/rocm_aiter_fa.py

+2 -2

@@ -231,8 +231,8 @@ class AiterFlashAttentionMetadataBuilder(

231 AttentionMetadataBuilder[AiterFlashAttentionMetadata]):

232 full_cudagraph_supported: ClassVar[bool] = True

233

234 - def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,

235 - device: torch.device):

234 + def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],

235 + vllm_config: VllmConfig, device: torch.device):

236 self.vllm_config = vllm_config

237 self.model_config = vllm_config.model_config

238 self.parallel_config = vllm_config.parallel_config

vllm/v1/attention/backends/triton_attn.py

+2 -2

@@ -59,8 +59,8 @@ class TritonAttentionMetadataBuilder(

59 AttentionMetadataBuilder[TritonAttentionMetadata]):

60 full_cudagraph_supported: ClassVar[bool] = True

61

62 - def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,

63 - device: torch.device):

62 + def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],

63 + vllm_config: VllmConfig, device: torch.device):

64 self.device = device

65 self.block_size = kv_cache_spec.block_size

66 self.kv_cache_spec = kv_cache_spec

vllm/v1/attention/backends/utils.py

+9 -5

@@ -70,8 +70,8 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):

70 full_cudagraph_supported: ClassVar[bool] = False

71

72 @abstractmethod

73 - def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,

74 - device: torch.device):

73 + def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],

74 + vllm_config: VllmConfig, device: torch.device):

75 self.kv_cache_spec = kv_cache_spec

76

77 @abstractmethod

@@ -164,14 +164,14 @@ class PerLayerParameters:

164

165

166 def get_per_layer_parameters(

167 - vllm_config: VllmConfig,

167 + vllm_config: VllmConfig, layer_names: list[str],

168 cls_: type['AttentionImpl']) -> dict[str, PerLayerParameters]:

169 """

170 - Scan all attention layers and determine some hyperparameters

170 + Scan layers in `layer_names` and determine some hyperparameters

171 to use during `plan`.

172 """

173

174 - layers = get_layers_from_vllm_config(vllm_config, Attention)

174 + layers = get_layers_from_vllm_config(vllm_config, Attention, layer_names)

175 per_layer_params: dict[str, PerLayerParameters] = {}

176

177 for key, layer in layers.items():

@@ -208,6 +208,10 @@ def infer_global_hyperparameters(

208 param_sets = list(per_layer_params.values())

209 global_params = param_sets[0]

210 for params in param_sets:

211 + if params.window_left != global_params.window_left:

212 + raise ValueError(

213 + "Window left is not the same for all layers. One potential fix "

214 + "is to set disable_sliding_window=True")

215 assert params == global_params, (

216 "FlashInfer backend currently only supports models in which all "

217 "layers share the same values for the following hyperparameters: "

vllm/v1/worker/gpu_model_runner.py

+8 -5

@@ -2521,7 +2521,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2521 elapsed_time, cuda_graph_size / (1 << 30))

2522

2523 def _initialize_single_attn_backend(

2524 - self, kv_cache_spec: KVCacheSpec

2524 + self, kv_cache_spec: KVCacheSpec, layer_names: list[str]

2525 ) -> tuple[AttentionBackend, AttentionMetadataBuilder]:

2526 if isinstance(kv_cache_spec, AttentionSpec):

2527 attn_backend_i = get_attn_backend(

@@ -2551,6 +2551,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2551

2552 attn_metadata_builder_i = attn_backend_i.get_builder_cls()(

2553 kv_cache_spec,

2554 + layer_names,

2555 self.vllm_config,

2556 self.device,

2557 )

@@ -2574,8 +2575,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2575 kv_cache_config.kv_cache_groups):

2576 kv_cache_spec = kv_cache_group_spec.kv_cache_spec

2577

2577 - attn_backend_i, attn_metadata_builder_i = \

2578 - self._initialize_single_attn_backend(kv_cache_spec)

2578 + attn_backend_i, attn_metadata_builder_i = (

2579 + self._initialize_single_attn_backend(

2580 + kv_cache_spec, kv_cache_group_spec.layer_names))

2581 self.attn_backends.append(attn_backend_i)

2582 self.attn_metadata_builders.append(attn_metadata_builder_i)

2583

@@ -2606,8 +2608,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2608 assert len(attn_specs) == len(attn_layers), \

2609 "All or none of the layers are expected to be encoder-only"

2610

2609 - attn_backend, attn_metadata_builder = \

2610 - self._initialize_single_attn_backend(attn_specs[0])

2611 + attn_backend, attn_metadata_builder = (

2612 + self._initialize_single_attn_backend(attn_specs[0],

2613 + attn_layers.keys()))

2614 self.attn_backends.append(attn_backend)

2615 self.attn_metadata_builders.append(attn_metadata_builder)

2616 self.is_encoder_only_model = True

[v1][attention] Support Hybrid Allocator + FlashInfer (#21412)