9fb52e523a · tduncan/Nvidia-Optimization-Stack-VLLM

.buildkite/scripts/hardware_ci/run-amd-test.sh

+1 -2

@@ -107,10 +107,9 @@ fi

107

108 if [[ $commands == *" kernels/attention"* ]]; then

109 commands="${commands} \

110 - --ignore=kernels/attention/stest_attention_selector.py \

110 + --ignore=kernels/attention/test_attention_selector.py \

111 --ignore=kernels/attention/test_blocksparse_attention.py \

112 --ignore=kernels/attention/test_encoder_decoder_attn.py \

113 - --ignore=kernels/attention/test_attention_selector.py \

113 --ignore=kernels/attention/test_flash_attn.py \

114 --ignore=kernels/attention/test_flashinfer.py \

115 --ignore=kernels/attention/test_prefix_prefill.py \

docs/models/supported_models.md

+2 -8

@@ -626,9 +626,6 @@ Specified using `--task generate`.

626 !!! note

627 Only `InternVLChatModel` with Qwen2.5 text backbone (`OpenGVLab/InternVL3-2B`, `OpenGVLab/InternVL2.5-1B` etc) has video inputs support currently.

628

629 -!!! note

630 - `h2oai/h2ovl-mississippi-2b` will be available in V1 once we support head size 80.

631 -

629 !!! note

630 To use `TIGER-Lab/Mantis-8B-siglip-llama3`, you have to pass `--hf_overrides '{"architectures": ["MantisForConditionalGeneration"]}'` when running vLLM.

631

@@ -671,11 +668,8 @@ Specified using `--task generate`.

668 Our PaliGemma implementations have the same problem as Gemma 3 (see above) for both V0 and V1.

669

670 !!! note

674 - To use Qwen2.5-Omni, you have to install Hugging Face Transformers library from source via

675 - `pip install git+https://github.com/huggingface/transformers.git`.

676 -

677 - Read audio from video pre-processing is currently supported on V0 (but not V1), because overlapping modalities is not yet supported in V1.

678 - `--mm-processor-kwargs '{"use_audio_in_video": true}'`.

671 + For Qwen2.5-Omni, reading audio from video pre-processing (`--mm-processor-kwargs '{"use_audio_in_video": true}'`)

672 + is currently supported on V0 (but not V1), because overlapping modalities is not yet supported in V1.

673

674 #### Transcription

675

examples/offline_inference/vision_language.py

+2 -2

@@ -98,7 +98,7 @@ def run_blip2(questions: list[str], modality: str) -> ModelRequestData:

98 # See https://huggingface.co/Salesforce/blip2-opt-2.7b/discussions/15#64ff02f3f8cf9e4f5b038262 #noqa

99 prompts = [f"Question: {question} Answer:" for question in questions]

100 engine_args = EngineArgs(

101 - model="Salesforce/blip2-opt-6.7b",

101 + model="Salesforce/blip2-opt-2.7b",

102 limit_mm_per_prompt={modality: 1},

103 )

104

@@ -971,7 +971,7 @@ def run_pixtral_hf(questions: list[str], modality: str) -> ModelRequestData:

971 )

972

973

974 -# Qwen

974 +# Qwen-VL

975 def run_qwen_vl(questions: list[str], modality: str) -> ModelRequestData:

976 assert modality == "image"

977

tests/kernels/attention/test_attention_selector.py

+12 -1

@@ -172,7 +172,7 @@ def test_env(

172 expected = "FLASHINFER_VLLM_V1" if use_v1 else name

173 assert backend.get_name() == expected

174 else:

175 - backend = get_attn_backend(16,

175 + backend = get_attn_backend(32,

176 torch.float16,

177 torch.float16,

178 block_size,

@@ -181,6 +181,17 @@ def test_env(

181 expected = "FLASH_ATTN_VLLM_V1" if use_v1 else name

182 assert backend.get_name() == expected

183

184 + if use_v1:

185 + backend = get_attn_backend(16,

186 + torch.float16,

187 + torch.float16,

188 + block_size,

189 + False,

190 + use_mla=use_mla)

191 + assert backend.get_name() == "FLEX_ATTENTION", (

192 + "Should fallback to FlexAttention if head size is "

193 + "not supported by FlashAttention")

194 +

195

196 @pytest.mark.parametrize("device", ["cpu", "cuda"])

197 @pytest.mark.parametrize("use_v1", [True, False])

tests/models/multimodal/generation/test_common.py

+2 -7

@@ -33,9 +33,6 @@ if current_platform.is_rocm():

33 os.environ["VLLM_USE_TRITON_FLASH_ATTN"] = "0"

34

35 REQUIRES_V0_MODELS = [

36 - # V1 Test: no way to fall back for head_dim = 80

37 - # https://github.com/vllm-project/vllm/issues/14524

38 - "qwen_vl",

36 # V1 Test: not enough KV cache space in C1.

37 "fuyu",

38 ]

@@ -221,8 +218,7 @@ VLM_TEST_SETTINGS = {

218 marks=[large_gpu_mark(min_gb=32)],

219 ),

220 "blip2": VLMTestInfo(

224 - # TODO: Change back to 2.7b once head_dim = 80 is supported

225 - models=["Salesforce/blip2-opt-6.7b"],

221 + models=["Salesforce/blip2-opt-2.7b"],

222 test_type=VLMTestType.IMAGE,

223 prompt_formatter=lambda img_prompt: f"Question: {img_prompt} Answer:",

224 img_idx_to_prompt=lambda idx: "",

@@ -340,8 +336,7 @@ VLM_TEST_SETTINGS = {

336 "h2ovl": VLMTestInfo(

337 models = [

338 "h2oai/h2ovl-mississippi-800m",

343 - # TODO: Re-enable once head_dim = 80 is supported

344 - # "h2oai/h2ovl-mississippi-2b",

339 + "h2oai/h2ovl-mississippi-2b",

340 ],

341 test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),

342 prompt_formatter=lambda img_prompt: f"<|prompt|>{img_prompt}<|end|><|answer|>", # noqa: E501

tests/models/quantization/test_gguf.py

+1 -1

@@ -83,7 +83,7 @@ MODELS = [

83 QWEN2_CONFIG,

84 PHI3_CONFIG,

85 GPT2_CONFIG,

86 - # STABLELM_CONFIG, # enable this when v1 support head_size=80

86 + STABLELM_CONFIG,

87 DOLPHIN_CONFIG,

88 # STARCODER_CONFIG, # broken

89 ]

tests/models/registry.py

+6 -9

@@ -240,8 +240,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

240 "OrionForCausalLM": _HfExamplesInfo("OrionStarAI/Orion-14B-Chat",

241 trust_remote_code=True),

242 "PersimmonForCausalLM": _HfExamplesInfo("adept/persimmon-8b-chat"),

243 - "PhiForCausalLM": _HfExamplesInfo("microsoft/phi-2", v0_only=True),

243 + "PhiForCausalLM": _HfExamplesInfo("microsoft/phi-2"),

244 "Phi3ForCausalLM": _HfExamplesInfo("microsoft/Phi-3-mini-4k-instruct"),

245 + # Blocksparse attention not supported in V1 yet

246 "Phi3SmallForCausalLM": _HfExamplesInfo("microsoft/Phi-3-small-8k-instruct",

247 trust_remote_code=True,

248 v0_only=True),

@@ -258,10 +259,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

259 "Qwen3MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen3-30B-A3B"),

260 "Qwen3ForSequenceClassification": _HfExamplesInfo("tomaarsen/Qwen3-Reranker-0.6B-seq-cls"), # noqa: E501

261 "RWForCausalLM": _HfExamplesInfo("tiiuae/falcon-40b"),

261 - "StableLMEpochForCausalLM": _HfExamplesInfo("stabilityai/stablelm-zephyr-3b", # noqa: E501

262 - v0_only=True),

263 - "StableLmForCausalLM": _HfExamplesInfo("stabilityai/stablelm-3b-4e1t",

264 - v0_only=True),

262 + "StableLMEpochForCausalLM": _HfExamplesInfo("stabilityai/stablelm-zephyr-3b"), # noqa: E501

263 + "StableLmForCausalLM": _HfExamplesInfo("stabilityai/stablelm-3b-4e1t"),

264 "Starcoder2ForCausalLM": _HfExamplesInfo("bigcode/starcoder2-3b"),

265 "SolarForCausalLM": _HfExamplesInfo("upstage/solar-pro-preview-instruct"),

266 "TeleChat2ForCausalLM": _HfExamplesInfo("Tele-AI/TeleChat2-3B",

@@ -330,8 +329,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {

329 "AriaForConditionalGeneration": _HfExamplesInfo("rhymes-ai/Aria"),

330 "AyaVisionForConditionalGeneration": _HfExamplesInfo("CohereForAI/aya-vision-8b"), # noqa: E501

331 "Blip2ForConditionalGeneration": _HfExamplesInfo("Salesforce/blip2-opt-2.7b", # noqa: E501

333 - extras={"6b": "Salesforce/blip2-opt-6.7b"}, # noqa: E501

334 - v0_only=True),

332 + extras={"6b": "Salesforce/blip2-opt-6.7b"}), # noqa: E501

333 "ChameleonForConditionalGeneration": _HfExamplesInfo("facebook/chameleon-7b"), # noqa: E501

334 "DeepseekVLV2ForCausalLM": _HfExamplesInfo("deepseek-ai/deepseek-vl2-tiny", # noqa: E501

335 extras={"fork": "Isotr0py/deepseek-vl2-tiny"}, # noqa: E501

@@ -359,8 +357,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {

357 trust_remote_code=True),

358 "KimiVLForConditionalGeneration": _HfExamplesInfo("moonshotai/Kimi-VL-A3B-Instruct", # noqa: E501

359 extras={"thinking": "moonshotai/Kimi-VL-A3B-Thinking"}, # noqa: E501

362 - trust_remote_code=True,

363 - v0_only=True),

360 + trust_remote_code=True),

361 "Llama4ForConditionalGeneration": _HfExamplesInfo("meta-llama/Llama-4-Scout-17B-16E-Instruct", # noqa: E501

362 max_model_len=10240),

363 "LlavaForConditionalGeneration": _HfExamplesInfo("llava-hf/llava-1.5-7b-hf",

tests/models/test_initialization.py

+2 -1

@@ -22,7 +22,8 @@ def test_can_initialize(model_arch: str, monkeypatch: pytest.MonkeyPatch):

22 model_info.check_transformers_version(on_fail="skip")

23

24 # FIXME: Possible memory leak in the previous tests?

25 - if model_arch == "GraniteSpeechForConditionalGeneration":

25 + if model_arch in ("GraniteSpeechForConditionalGeneration",

26 + "KimiVLForConditionalGeneration"):

27 pytest.skip("Avoid OOM")

28

29 # Avoid OOM and reduce initialization time by only using 1 layer

vllm/attention/layer.py

+2 -1

@@ -310,7 +310,8 @@ class MultiHeadAttention(nn.Module):

310 # currently, only torch_sdpa is supported on rocm

311 self.attn_backend = _Backend.TORCH_SDPA

312 else:

313 - if backend in {_Backend.FLASH_ATTN, _Backend.FLASH_ATTN_VLLM_V1}:

313 + if backend in (_Backend.FLASH_ATTN, _Backend.FLASH_ATTN_VLLM_V1,

314 + _Backend.FLEX_ATTENTION):

315 backend = _Backend.XFORMERS

316

317 self.attn_backend = backend if backend in {

vllm/attention/selector.py

+30 -3

@@ -4,7 +4,7 @@

4 import os

5 from contextlib import contextmanager

6 from functools import cache

7 -from typing import Generator, Optional, Type

7 +from typing import Generator, Optional, Union

8

9 import torch

10

@@ -79,6 +79,33 @@ def get_global_forced_attn_backend() -> Optional[_Backend]:

79 return forced_attn_backend

80

81

82 +def supports_head_size(

83 + attn_backend: Union[str, type[AttentionBackend]],

84 + head_size: int,

85 +) -> bool:

86 + if isinstance(attn_backend, str):

87 + try:

88 + attn_backend = resolve_obj_by_qualname(attn_backend)

89 + except ImportError:

90 + return False

91 +

92 + assert isinstance(attn_backend, type)

93 +

94 + # TODO: Update the interface once V0 is removed

95 + if get_supported_head_sizes := getattr(attn_backend,

96 + "get_supported_head_sizes", None):

97 + return head_size in get_supported_head_sizes()

98 + if validate_head_size := getattr(attn_backend, "validate_head_size", None):

99 + try:

100 + validate_head_size(head_size)

101 + return True

102 + except Exception:

103 + return False

104 +

105 + raise NotImplementedError(f"{attn_backend.__name__} does not support "

106 + "head size validation")

107 +

108 +

109 def get_attn_backend(

110 head_size: int,

111 dtype: torch.dtype,

@@ -87,7 +114,7 @@ def get_attn_backend(

114 is_attention_free: bool,

115 is_blocksparse: bool = False,

116 use_mla: bool = False,

90 -) -> Type[AttentionBackend]:

117 +) -> type[AttentionBackend]:

118 """Selects which attention backend to use and lazily imports it."""

119 # Accessing envs.* behind an @lru_cache decorator can cause the wrong

120 # value to be returned from the cache if the value changes between calls.

@@ -115,7 +142,7 @@ def _cached_get_attn_backend(

142 is_blocksparse: bool = False,

143 use_v1: bool = False,

144 use_mla: bool = False,

118 -) -> Type[AttentionBackend]:

145 +) -> type[AttentionBackend]:

146 if is_blocksparse:

147 logger.info("Using BlocksparseFlashAttention backend.")

148 from vllm.attention.backends.blocksparse_attn import (

vllm/config.py

+1 -1

@@ -2319,7 +2319,7 @@ class SchedulerConfig:

2319

2320 if self.max_num_batched_tokens > self.max_num_seqs * self.max_model_len:

2321 logger.warning(

2322 - "max_num_batched_tokens (%d) exceeds max_num_seqs"

2322 + "max_num_batched_tokens (%d) exceeds max_num_seqs "

2323 "* max_model_len (%d). This may lead to unexpected behavior.",

2324 self.max_num_batched_tokens,

2325 self.max_num_seqs * self.max_model_len)

vllm/platforms/cuda.py

+28 -16

@@ -234,35 +234,44 @@ class CudaPlatformBase(Platform):

234 return ("vllm.attention.backends."

235 "flashmla.FlashMLABackend")

236 if use_v1:

237 + FLASHINFER_V1 = "vllm.v1.attention.backends.flashinfer.FlashInferBackend" # noqa: E501

238 + FLEX_ATTENTION_V1 = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend" # noqa: E501

239 + TRITON_ATTN_VLLM_V1 = "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend" # noqa: E501

240 + FLASH_ATTN_V1 = "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend" # noqa: E501

241 +

242 if selected_backend == _Backend.FLASHINFER:

243 logger.info_once("Using FlashInfer backend on V1 engine.")

239 - return "vllm.v1.attention.backends.flashinfer.FlashInferBackend"

244 + return FLASHINFER_V1

245 elif selected_backend == _Backend.FLEX_ATTENTION:

241 - logger.info("Using FlexAttenion backend on V1 engine.")

242 - return "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend" # noqa: E501

246 + logger.info_once("Using FlexAttention backend on V1 engine.")

247 + return FLEX_ATTENTION_V1

248 elif selected_backend == _Backend.TRITON_ATTN_VLLM_V1:

249 logger.info_once("Using Triton backend on V1 engine.")

245 - return ("vllm.v1.attention.backends."

246 - "triton_attn.TritonAttentionBackend")

250 + return TRITON_ATTN_VLLM_V1

251 elif selected_backend == _Backend.FLASH_ATTN:

252 logger.info_once("Using Flash Attention backend on V1 engine.")

249 - return ("vllm.v1.attention.backends."

250 - "flash_attn.FlashAttentionBackend")

253 + return FLASH_ATTN_V1

254 +

255 + from vllm.attention.selector import supports_head_size

256

257 # Default backends for V1 engine

253 - # Prefer FlashInfer for Blackwell GPUs if installed

258 + # FP32 is only supported by FlexAttention

259 if dtype not in (torch.float16, torch.bfloat16):

260 logger.info_once(

256 - f"Using FlexAttenion backend for {dtype} on V1 engine.")

257 - return "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend" # noqa: E501

258 - if cls.is_device_capability(100):

261 + "Using FlexAttention backend for %s on V1 engine.",

262 + dtype,

263 + )

264 + return FLEX_ATTENTION_V1

265 +

266 + # Prefer FlashInfer for Blackwell GPUs if installed

267 + if cls.is_device_capability(100) and \

268 + supports_head_size(FLASHINFER_V1, head_size):

269 try:

270 import flashinfer # noqa: F401

271 logger.info_once(

272 "Using FlashInfer backend on V1 engine by default for "

273 "Blackwell (SM 10.0) GPUs.")

264 - return ("vllm.v1.attention.backends."

265 - "flashinfer.FlashInferBackend")

274 + return FLASHINFER_V1

275 except ImportError:

276 logger.info_once(

277 "FlashInfer failed to import for V1 engine on "

@@ -270,10 +279,13 @@ class CudaPlatformBase(Platform):

279 "install FlashInfer for better performance.")

280 pass

281 # FlashAttention is the default for SM 8.0+ GPUs

273 - if cls.has_device_capability(80):

282 + if cls.has_device_capability(80) and \

283 + supports_head_size(FLASH_ATTN_V1, head_size):

284 logger.info_once("Using Flash Attention backend on V1 engine.")

275 - return ("vllm.v1.attention.backends."

276 - "flash_attn.FlashAttentionBackend")

285 + return FLASH_ATTN_V1

286 +

287 + logger.info_once("Using FlexAttention backend on V1 engine.")

288 + return FLEX_ATTENTION_V1

289

290 # Backends for V0 engine

291 if selected_backend == _Backend.FLASHINFER:

vllm/v1/attention/backends/cpu_attn.py

+18 -2

@@ -3,7 +3,8 @@

3 import numpy as np

4 import torch

5

6 -from vllm.attention.backends.abstract import AttentionMetadata

6 +from vllm.attention.backends.abstract import (AttentionBackend,

7 + AttentionMetadata)

8 from vllm.attention.backends.torch_sdpa import (TorchSDPABackendImpl,

9 TorchSDPAMetadata)

10 from vllm.attention.backends.utils import CommonAttentionState

@@ -17,9 +18,24 @@ from vllm.v1.worker.cpu_model_runner import CPUModelRunner

18 from vllm.v1.worker.gpu_input_batch import InputBatch

19

20

20 -class TorchSDPABackend:

21 +class TorchSDPABackend(AttentionBackend):

22 accept_output_buffer: bool = False

23

24 + @classmethod

25 + def get_supported_head_sizes(cls) -> list[int]:

26 + return PagedAttention.get_supported_head_sizes()

27 +

28 + @classmethod

29 + def validate_head_size(cls, head_size: int) -> None:

30 + supported_head_sizes = cls.get_supported_head_sizes()

31 + if head_size not in supported_head_sizes:

32 + attn_type = cls.__name__.removesuffix("Backend")

33 + raise ValueError(

34 + f"Head size {head_size} is not supported by {attn_type}. "

35 + f"Supported head sizes are: {supported_head_sizes}. "

36 + "Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "

37 + "FlexAttention backend which supports all head sizes.")

38 +

39 @staticmethod

40 def get_name() -> str:

41 return "TORCH_SDPA_VLLM_V1"

vllm/v1/attention/backends/flash_attn.py

+14 -8

@@ -44,10 +44,21 @@ class FlashAttentionBackend(AttentionBackend):

44

45 accept_output_buffer: bool = True

46

47 - @staticmethod

48 - def get_supported_head_sizes() -> list[int]:

47 + @classmethod

48 + def get_supported_head_sizes(cls) -> list[int]:

49 return [32, 64, 96, 128, 160, 192, 224, 256]

50

51 + @classmethod

52 + def validate_head_size(cls, head_size: int) -> None:

53 + supported_head_sizes = cls.get_supported_head_sizes()

54 + if head_size not in supported_head_sizes:

55 + attn_type = cls.__name__.removesuffix("Backend")

56 + raise ValueError(

57 + f"Head size {head_size} is not supported by {attn_type}. "

58 + f"Supported head sizes are: {supported_head_sizes}. "

59 + "Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "

60 + "FlexAttention backend which supports all head sizes.")

61 +

62 @staticmethod

63 def get_name() -> str:

64 return "FLASH_ATTN_VLLM_V1"

@@ -416,12 +427,7 @@ class FlashAttentionImpl(AttentionImpl):

427

428 self.num_queries_per_kv = self.num_heads // self.num_kv_heads

429

419 - support_head_sizes = FlashAttentionBackend.get_supported_head_sizes()

420 - if head_size not in support_head_sizes:

421 - raise ValueError(

422 - f"Head size {head_size} is not supported by FlashAttention. "

423 - f"Supported head sizes are: {support_head_sizes}. "

424 - "Set VLLM_USE_V1=0 to use another attention backend.")

430 + FlashAttentionBackend.validate_head_size(head_size)

431

432 if attn_type != AttentionType.DECODER:

433 raise NotImplementedError("Encoder self-attention and "

vllm/v1/attention/backends/flashinfer.py

+16 -10

@@ -38,10 +38,22 @@ class FlashInferBackend(AttentionBackend):

38

39 accept_output_buffer: bool = True

40

41 - @staticmethod

42 - def get_supported_head_sizes() -> list[int]:

41 + @classmethod

42 + def get_supported_head_sizes(cls) -> list[int]:

43 + # https://github.com/flashinfer-ai/flashinfer/blob/3d55c71a62052c590c130897d3a3db49b14fcc34/include/flashinfer/utils.cuh#L157

44 return [64, 128, 256]

45

46 + @classmethod

47 + def validate_head_size(cls, head_size: int) -> None:

48 + supported_head_sizes = cls.get_supported_head_sizes()

49 + if head_size not in supported_head_sizes:

50 + attn_type = cls.__name__.removesuffix("Backend")

51 + raise ValueError(

52 + f"Head size {head_size} is not supported by {attn_type}. "

53 + f"Supported head sizes are: {supported_head_sizes}. "

54 + "Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "

55 + "FlexAttention backend which supports all head sizes.")

56 +

57 @staticmethod

58 def get_name() -> str:

59 return "FLASHINFER_VLLM_V1"

@@ -207,14 +219,8 @@ class FlashInferMetadata:

219 return self.qo_indptr

220

221 def __post_init__(self):

210 - # Refer to

211 - # https://github.com/flashinfer-ai/flashinfer/blob/3d55c71a62052c590c130897d3a3db49b14fcc34/include/flashinfer/utils.cuh#L157

212 - supported_head_sizes = FlashInferBackend.get_supported_head_sizes()

213 - if self.head_dim is not None and self.head_dim \

214 - not in supported_head_sizes:

215 - raise ValueError(

216 - f"Only {supported_head_sizes} are supported for head_dim,",

217 - f" received {self.head_dim}.")

222 + if self.head_dim is not None:

223 + FlashInferBackend.validate_head_size(self.head_dim)

224

225

226 class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

vllm/v1/attention/backends/flex_attention.py

+20 -19

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Attention layer with FlashAttention."""

4 -

4 +from collections import defaultdict

5 from dataclasses import dataclass

6 from typing import TYPE_CHECKING, Any, Optional

7

@@ -21,9 +21,6 @@ from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,

21 from vllm.v1.kv_cache_interface import AttentionSpec

22 from vllm.v1.worker.block_table import BlockTable

23

24 -if current_platform.is_cuda():

25 - pass

26 -

24 logger = init_logger(__name__)

25

26 if TYPE_CHECKING:

@@ -45,9 +42,9 @@ def _offsets_to_doc_ids_tensor(offsets: torch.Tensor) -> torch.Tensor:

42 class FlexAttentionBackend(AttentionBackend):

43 accept_output_buffer: bool = True

44

48 - @staticmethod

49 - def get_supported_head_sizes() -> list[int]:

50 - return [16, 32, 64, 96, 128, 160, 192, 224, 256]

45 + @classmethod

46 + def validate_head_size(cls, head_size: int) -> None:

47 + return # FlexAttention supports any head size

48

49 @staticmethod

50 def get_name() -> str:

@@ -384,12 +381,8 @@ class FlexAttentionImpl(AttentionImpl):

381 raise NotImplementedError(

382 "FlexAttention does not support kv sharing yet.")

383

387 - support_head_sizes = FlexAttentionBackend.get_supported_head_sizes()

388 - if head_size not in support_head_sizes:

389 - raise ValueError(

390 - f"Head size {head_size} is not supported by FlashAttention. "

391 - f"Supported head sizes are: {support_head_sizes}. "

392 - "Set VLLM_USE_V1=0 to use another attention backend.")

384 + FlexAttentionBackend.validate_head_size(head_size)

385 +

386 if is_quantized_kv_cache(self.kv_cache_dtype):

387 raise NotImplementedError(

388 "FlexAttention does not support quantized kv-cache. Yet")

@@ -464,12 +457,20 @@ class FlexAttentionImpl(AttentionImpl):

457 # Doesn't work for now -> constraint violation

458 # torch._dynamo.try_mark_dynamic(query, 2)

459

467 - # default M=64, N=64 may run out of shared memory on

468 - # some GPUs with fp32, so we use smaller M and N.

469 - extra_kernel_options = {

470 - "BLOCK_M": 32,

471 - "BLOCK_N": 32

472 - } if query.dtype == torch.float32 else {}

460 + # default M=64, N=64 may run out of shared memory on some GPUs

461 + # TODO: Explicit configs for each GPU?

462 + # Not sure how to calculate the shared memory requirement

463 + extra_kernel_options = defaultdict[str, int](lambda: 64)

464 + if query.dtype == torch.float32:

465 + extra_kernel_options["BLOCK_M"] //= 2

466 + extra_kernel_options["BLOCK_N"] //= 2

467 + if current_platform.is_cuda():

468 + device_props = torch.cuda.get_device_properties()

469 + max_shared_memory = device_props.shared_memory_per_block_optin

470 + if max_shared_memory < 144 * 1024:

471 + extra_kernel_options["BLOCK_M"] //= 2

472 + extra_kernel_options["BLOCK_N"] //= 2

473 +

474 out = flex_attention_compiled(

475 query,

476 key_cache,

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

+15 -8

@@ -254,10 +254,21 @@ class MLACommonBackend(AttentionBackend):

254 ) -> tuple[int, ...]:

255 return (num_blocks, block_size, head_size)

256

257 - @staticmethod

258 - def get_supported_head_sizes() -> list[int]:

257 + @classmethod

258 + def get_supported_head_sizes(cls) -> list[int]:

259 return [576]

260

261 + @classmethod

262 + def validate_head_size(cls, head_size: int) -> None:

263 + supported_head_sizes = cls.get_supported_head_sizes()

264 + if head_size not in supported_head_sizes:

265 + attn_type = cls.__name__.removesuffix("Backend")

266 + raise ValueError(

267 + f"Head size {head_size} is not supported by {attn_type}. "

268 + f"Supported head sizes are: {supported_head_sizes}. "

269 + "Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "

270 + "FlexAttention backend which supports all head sizes.")

271 +

272

273 @dataclass

274 class MLACommonPrefillMetadata:

@@ -320,12 +331,8 @@ class MLACommonMetadata(Generic[D]):

331 prefill: Optional[MLACommonPrefillMetadata] = None

332

333 def __post_init__(self):

323 - supported_head_sizes = MLACommonBackend.get_supported_head_sizes()

324 - if self.head_dim is not None and self.head_dim \

325 - not in supported_head_sizes:

326 - raise ValueError(

327 - f"Only {supported_head_sizes} are supported for head_dim,",

328 - f"received {self.head_dim}.")

334 + if self.head_dim is not None:

335 + MLACommonBackend.validate_head_size(self.head_dim)

336

337

338 M = TypeVar("M", bound=MLACommonMetadata)

vllm/v1/attention/backends/rocm_aiter_fa.py

+14 -10

@@ -314,10 +314,21 @@ class AiterFlashAttentionBackend(AttentionBackend):

314

315 accept_output_buffer: bool = True

316

317 - @staticmethod

318 - def get_supported_head_sizes() -> list[int]:

317 + @classmethod

318 + def get_supported_head_sizes(cls) -> list[int]:

319 return [32, 64, 96, 128, 160, 192, 224, 256]

320

321 + @classmethod

322 + def validate_head_size(cls, head_size: int) -> None:

323 + supported_head_sizes = cls.get_supported_head_sizes()

324 + if head_size not in supported_head_sizes:

325 + attn_type = cls.__name__.removesuffix("Backend")

326 + raise ValueError(

327 + f"Head size {head_size} is not supported by {attn_type}. "

328 + f"Supported head sizes are: {supported_head_sizes}. "

329 + "Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "

330 + "FlexAttention backend which supports all head sizes.")

331 +

332 @staticmethod

333 def get_name() -> str:

334 return "FLASH_ATTN_VLLM_V1"

@@ -428,14 +439,7 @@ class AiterFlashAttentionImpl(AttentionImpl):

439 assert self.num_heads % self.num_kv_heads == 0

440 self.num_queries_per_kv = self.num_heads // self.num_kv_heads

441

431 - support_head_sizes = \

432 - AiterFlashAttentionBackend.get_supported_head_sizes()

433 - if head_size not in support_head_sizes:

434 - raise ValueError(

435 - f"Head size {head_size} is not supported by "

436 - "AiterFlashAttention. "

437 - f"Supported head sizes are: {support_head_sizes}. "

438 - "Set VLLM_USE_V1=0 to use another attention backend.")

442 + AiterFlashAttentionBackend.validate_head_size(head_size)

443

444 if attn_type != AttentionType.DECODER:

445 raise NotImplementedError("Encoder self-attention and "

vllm/v1/attention/backends/triton_attn.py

+14 -7

@@ -190,10 +190,21 @@ class TritonAttentionBackend(AttentionBackend):

190

191 accept_output_buffer: bool = True

192

193 - @staticmethod

194 - def get_supported_head_sizes() -> list[int]:

193 + @classmethod

194 + def get_supported_head_sizes(cls) -> list[int]:

195 return [32, 64, 96, 128, 160, 192, 224, 256]

196

197 + @classmethod

198 + def validate_head_size(cls, head_size: int) -> None:

199 + supported_head_sizes = cls.get_supported_head_sizes()

200 + if head_size not in supported_head_sizes:

201 + attn_type = cls.__name__.removesuffix("Backend")

202 + raise ValueError(

203 + f"Head size {head_size} is not supported by {attn_type}. "

204 + f"Supported head sizes are: {supported_head_sizes}. "

205 + "Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "

206 + "FlexAttention backend which supports all head sizes.")

207 +

208 @staticmethod

209 def get_name() -> str:

210 return "TRITON_ATTN_VLLM_V1"

@@ -268,11 +279,7 @@ class TritonAttentionImpl(AttentionImpl):

279

280 self.num_queries_per_kv = self.num_heads // self.num_kv_heads

281

271 - support_head_sizes = TritonAttentionBackend.get_supported_head_sizes()

272 - if head_size not in support_head_sizes:

273 - raise ValueError(

274 - f"Head size {head_size} is not supported by TritonAttention. "

275 - f"Supported head sizes are: {support_head_sizes}.")

282 + TritonAttentionBackend.validate_head_size(head_size)

283

284 if attn_type != AttentionType.DECODER:

285 raise NotImplementedError("Encoder self-attention and "

vllm/v1/spec_decode/eagle.py

+2 -2

@@ -12,8 +12,8 @@ from vllm.logger import init_logger

12 from vllm.model_executor.model_loader import get_model

13 from vllm.model_executor.models import supports_multimodal

14 from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM

15 -from vllm.v1.attention.backends.flash_attn import (CommonAttentionMetadata,

16 - FlashAttentionMetadata)

15 +from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata

16 +from vllm.v1.attention.backends.utils import CommonAttentionMetadata

17 from vllm.v1.kv_cache_interface import KVCacheConfig

18 from vllm.v1.sample.metadata import SamplingMetadata

19 from vllm.v1.spec_decode.utils import prepare_eagle_input_kernel

[V1] Support any head size for FlexAttention backend (#20467)