@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b952f4d3c3

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, 21:15 UTC

[v1] Add PrefixLM support to FlexAttention backend (#27938)

Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>

Isotr0py committed Dec 7, 2025 at 23:51 UTC b952f4d3c31068de3a98d680261be9a5caa04caa

16 files changed +173 -25

docs/models/supported_models.md

-20

@@ -740,23 +740,6 @@ Some models are supported only via the [Transformers modeling backend](#transfor

740 <sup>E</sup> Pre-computed embeddings can be inputted for this modality.

741 <sup>+</sup> Multiple items can be inputted per text prompt for this modality.

742

743 -!!! warning

744 - Both V0 and V1 support `Gemma3ForConditionalGeneration` for text-only inputs.

745 - However, there are differences in how they handle text + image inputs:

746 -

747 - V0 correctly implements the model's attention pattern:

748 - - Uses bidirectional attention between the image tokens corresponding to the same image

749 - - Uses causal attention for other tokens

750 - - Implemented via (naive) PyTorch SDPA with masking tensors

751 - - Note: May use significant memory for long prompts with image

752 -

753 - V1 currently uses a simplified attention pattern:

754 - - Uses causal attention for all tokens, including image tokens

755 - - Generates reasonable outputs but does not match the original model's attention for text + image inputs, especially when `{"do_pan_and_scan": true}`

756 - - Will be updated in the future to support the correct behavior

757 -

758 - This limitation exists because the model's mixed attention pattern (bidirectional for images, causal otherwise) is not yet supported by vLLM's attention backends.

759 -

743 !!! note

744 `Gemma3nForConditionalGeneration` is only supported on V1 due to shared KV caching and it depends on `timm>=1.0.17` to make use of its

745 MobileNet-v5 vision backbone.

@@ -776,9 +759,6 @@ Some models are supported only via the [Transformers modeling backend](#transfor

759 The official `openbmb/MiniCPM-V-2` doesn't work yet, so we need to use a fork (`HwwwH/MiniCPM-V-2`) for now.

760 For more details, please see: <https://github.com/vllm-project/vllm/pull/4087#issuecomment-2250397630>

761

779 -!!! warning

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

781 -

762 !!! note

763 For Qwen2.5-Omni and Qwen3-Omni, reading audio from video pre-processing (`--mm-processor-kwargs '{"use_audio_in_video": true}'`) is currently work in progress and not yet supported.

764

tests/models/multimodal/generation/test_common.py

-1

@@ -382,7 +382,6 @@ VLM_TEST_SETTINGS = {

382 auto_cls=AutoModelForImageTextToText,

383 vllm_runner_kwargs={"mm_processor_kwargs": {"do_pan_and_scan": True}},

384 patch_hf_runner=model_utils.gemma3_patch_hf_runner,

385 - num_logprobs=10,

385 ),

386 "glm4v": VLMTestInfo(

387 models=["zai-org/glm-4v-9b"],

tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/dummy_platform.py

@@ -30,5 +30,6 @@ class DummyPlatform(Platform):

30 use_mla,

31 has_sink,

32 use_sparse,

33 + use_mm_prefix,

34 ):

35 return "vllm_add_dummy_platform.dummy_attention_backend.DummyAttentionBackend" # noqa E501

vllm/attention/backends/abstract.py

@@ -166,6 +166,10 @@ class AttentionBackend(ABC):

166 def supports_sink(cls) -> bool:

167 return False

168

169 + @classmethod

170 + def supports_mm_prefix(cls) -> bool:

171 + return False

172 +

173 @classmethod

174 def is_sparse(cls) -> bool:

175 return False

@@ -207,6 +211,7 @@ class AttentionBackend(ABC):

211 use_mla: bool,

212 has_sink: bool,

213 use_sparse: bool,

214 + use_mm_prefix: bool,

215 device_capability: "DeviceCapability",

216 attn_type: str,

217 ) -> list[str]:

@@ -219,6 +224,10 @@ class AttentionBackend(ABC):

224 invalid_reasons.append("kv_cache_dtype not supported")

225 if not cls.supports_block_size(block_size):

226 invalid_reasons.append("block_size not supported")

227 + if use_mm_prefix and not cls.supports_mm_prefix():

228 + invalid_reasons.append(

229 + "partial multimodal token full attention not supported"

230 + )

231 if use_mla != cls.is_mla():

232 if use_mla:

233 invalid_reasons.append("MLA not supported")

vllm/attention/layer.py

@@ -230,6 +230,10 @@ class Attention(nn.Module, AttentionLayerBase):

230 self.sliding_window = sliding_window

231 self.has_sink = extra_impl_args.get("sinks") is not None

232

233 + # NOTE: model_config may be None during certain tests

234 + model_config = vllm_config.model_config

235 + self.use_mm_prefix = model_config is not None and model_config.is_mm_prefix_lm

236 +

237 # During model initialization, the default dtype is set as the model

238 # weight and activation dtype.

239 dtype = torch.get_default_dtype()

@@ -241,6 +245,7 @@ class Attention(nn.Module, AttentionLayerBase):

245 block_size,

246 use_mla=False,

247 has_sink=self.has_sink,

248 + use_mm_prefix=self.use_mm_prefix,

249 attn_type=attn_type,

250 )

251 else:

vllm/attention/selector.py

@@ -27,6 +27,7 @@ def get_attn_backend(

27 use_mla: bool = False,

28 has_sink: bool = False,

29 use_sparse: bool = False,

30 + use_mm_prefix: bool = False,

31 attn_type: str | None = None,

32 ) -> type[AttentionBackend]:

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

@@ -52,6 +53,7 @@ def get_attn_backend(

53 use_mla=use_mla,

54 has_sink=has_sink,

55 use_sparse=use_sparse,

56 + use_mm_prefix=use_mm_prefix,

57 attn_type=attn_type,

58 )

@@ -66,6 +68,7 @@ def _cached_get_attn_backend(

68 use_mla: bool = False,

69 has_sink: bool = False,

70 use_sparse: bool = False,

71 + use_mm_prefix: bool = False,

72 attn_type: str | None = None,

73 ) -> type[AttentionBackend]:

74 from vllm.platforms import current_platform

@@ -87,6 +90,7 @@ def _cached_get_attn_backend(

90 use_mla,

91 has_sink,

92 use_sparse,

93 + use_mm_prefix,

94 attn_type,

95 )

96 else:

@@ -99,6 +103,7 @@ def _cached_get_attn_backend(

103 use_mla,

104 has_sink,

105 use_sparse,

106 + use_mm_prefix,

107 attn_type,

108 )

109 if not attention_cls:

vllm/config/model.py

+14

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

4 import warnings

5 from collections.abc import Callable

6 from dataclasses import InitVar, field

7 +from functools import cached_property

8 from typing import TYPE_CHECKING, Any, Literal, cast, get_args

10 import torch

@@ -1217,6 +1218,19 @@ class ModelConfig:

1218 )

1219 return False

1220

1221 + @cached_property

1222 + def is_mm_prefix_lm(self) -> bool:

1223 + """Whether to use bidirectional attention for mm positions."""

1224 + MM_PREFIX_LM_MODELS = (

1225 + "gemma3",

1226 + # TODO(Isotr0py): Disable paligemma for now before

1227 + # we supports soft cap attention for FlexAttention

1228 + # "paligemma",

1229 + )

1230 + if not hasattr(self.hf_config, "model_type"):

1231 + return False

1232 + return self.hf_config.model_type in MM_PREFIX_LM_MODELS

1233 +

1234 def get_head_size(self) -> int:

1235 # TODO remove hard code

1236 if self.is_deepseek_mla:

vllm/multimodal/inputs.py

+25

@@ -175,6 +175,31 @@ class PlaceholderRange:

175

176 return int(self.is_embed.sum().item())

177

178 + def extract_embeds_range(self) -> list[tuple[int, int]]:

179 + """Extract the start and end indices of the embedded region in prompt.

180 +

181 + For example, given `PlaceholderRange(offset=2, length=5)` and

182 + `is_embed = [False, True, False, True, True]`, the output is

183 + `[(1 + offset, 1 + offset), (3 + offset, 4 + offset)]`.

184 +

185 + Returns:

186 + A tuple `(start, end)` representing the start and end

187 + indices (inclusive) of the embedded region.

188 + Returns full placeholder range if `is_embed` is `None`.

189 + """

190 + if self.is_embed is None:

191 + return [(self.offset, self.offset + self.length)]

192 +

193 + mask_i = self.is_embed.int()

194 + starts = torch.nonzero(

195 + torch.diff(mask_i, prepend=mask_i.new_zeros(1)) == 1

196 + ).flatten()

197 + ends = torch.nonzero(

198 + torch.diff(mask_i, append=mask_i.new_zeros(1)) == -1

199 + ).flatten()

200 + ranges = torch.stack((starts, ends), dim=1) + self.offset

201 + return [tuple(x) for x in ranges.tolist()]

202 +

203 def __eq__(self, other: object) -> bool:

204 if not isinstance(other, self.__class__):

205 return False

vllm/platforms/cpu.py

@@ -133,6 +133,7 @@ class CpuPlatform(Platform):

133 use_mla: bool,

134 has_sink: bool,

135 use_sparse: bool,

136 + use_mm_prefix: bool,

137 attn_type: str | None = None,

138 ) -> str:

139 if selected_backend and selected_backend != AttentionBackendEnum.CPU_ATTN:

vllm/platforms/cuda.py

+19

@@ -233,6 +233,20 @@ class CudaPlatformBase(Platform):

233 "Forcing kv cache block size to 64 for FlashMLASparse backend."

234 )

235

236 + scheduler_config = vllm_config.scheduler_config

237 + # Note: model_config may be None during testing

238 + if (

239 + model_config is not None

240 + and model_config.is_mm_prefix_lm

241 + and scheduler_config.is_multimodal_model

242 + and not scheduler_config.disable_chunked_mm_input

243 + ):

244 + logger.warning(

245 + "Forcing --disable_chunked_mm_input for models "

246 + "with multimodal-bidirectional attention."

247 + )

248 + scheduler_config.disable_chunked_mm_input = True

249 +

250 @classmethod

251 def get_current_memory_usage(

252 cls, device: torch.types.Device | None = None

@@ -268,6 +282,7 @@ class CudaPlatformBase(Platform):

282 use_mla,

283 has_sink,

284 use_sparse,

285 + use_mm_prefix,

286 device_capability,

287 attn_type,

288 ) -> tuple[

@@ -289,6 +304,7 @@ class CudaPlatformBase(Platform):

304 use_mla,

305 has_sink,

306 use_sparse,

307 + use_mm_prefix,

308 device_capability,

309 attn_type,

310 )

@@ -312,6 +328,7 @@ class CudaPlatformBase(Platform):

328 use_mla: bool,

329 has_sink: bool,

330 use_sparse: bool,

331 + use_mm_prefix: bool,

332 attn_type: str | None = None,

333 ) -> str:

334 if attn_type is None:

@@ -332,6 +349,7 @@ class CudaPlatformBase(Platform):

349 use_mla,

350 has_sink,

351 use_sparse,

352 + use_mm_prefix,

353 device_capability,

354 attn_type,

355 )

@@ -356,6 +374,7 @@ class CudaPlatformBase(Platform):

374 use_mla,

375 has_sink,

376 use_sparse,

377 + use_mm_prefix,

378 device_capability,

379 attn_type,

380 )

vllm/platforms/interface.py

@@ -239,6 +239,7 @@ class Platform:

239 use_mla: bool,

240 has_sink: bool,

241 use_sparse: bool,

242 + use_mm_prefix: bool,

243 attn_type: str | None = None,

244 ) -> str:

245 """Get the attention backend class of a device."""

vllm/platforms/rocm.py

@@ -216,6 +216,7 @@ class RocmPlatform(Platform):

216 use_mla,

217 has_sink,

218 use_sparse,

219 + use_mm_prefix,

220 attn_type: str | None = None,

221 ) -> str:

222 from vllm._aiter_ops import rocm_aiter_ops

vllm/platforms/tpu.py

+3 -2

@@ -62,8 +62,9 @@ class TpuPlatform(Platform):

62 kv_cache_dtype: str | None,

63 block_size: int,

64 use_mla: bool,

65 - has_sink,

66 - use_sparse,

65 + has_sink: bool,

66 + use_sparse: bool,

67 + use_mm_prefix: bool,

68 attn_type: str | None = None,

69 ) -> str:

70 if use_sparse:

vllm/platforms/xpu.py

+2 -1

@@ -48,7 +48,8 @@ class XPUPlatform(Platform):

48 block_size: int,

49 use_mla: bool,

50 has_sink: bool,

51 - use_sparse,

51 + use_sparse: bool,

52 + use_mm_prefix: bool,

53 attn_type: str | None = None,

54 ) -> str:

55 from vllm.v1.attention.backends.utils import set_kv_cache_layout

vllm/v1/attention/backends/flex_attention.py

+62

@@ -17,6 +17,7 @@ from torch.nn.attention.flex_attention import (

17 and_masks,

18 create_block_mask,

19 flex_attention,

20 + or_masks,

21 )

23 from vllm.attention.backends.abstract import (

@@ -42,6 +43,7 @@ from vllm.v1.kv_cache_interface import AttentionSpec

44 logger = init_logger(__name__)

46 +torch._dynamo.config.recompile_limit = 16

47 create_block_mask_compiled = torch.compile(

48 create_block_mask, fullgraph=True, mode="reduce-overhead"

49 )

@@ -91,6 +93,11 @@ class FlexAttentionBackend(AttentionBackend):

93 """FlexAttention supports both decoder and encoder-only attention."""

94 return attn_type in (AttentionType.DECODER, AttentionType.ENCODER_ONLY)

96 + @classmethod

97 + def supports_mm_prefix(cls) -> bool:

98 + """FlexAttention supports full attention for image tokens."""

99 + return True

100 +

101 @staticmethod

102 def get_impl_cls() -> type["FlexAttentionImpl"]:

103 return FlexAttentionImpl

@@ -316,6 +323,7 @@ class FlexAttentionMetadata:

323 kv_block_size: int = 16

324 transformed_score_mod: _score_mod_signature | None = None

325 sliding_window: int | None = None

326 + mm_prefix_range: dict[int, list[tuple[int, int]]] | None = None

327

328 @cached_property

329 def logical_block_ids(self):

@@ -443,6 +451,45 @@ class FlexAttentionMetadata:

451

452 return final_mask_mod if self.causal else sliding_window_mask_mod

453

454 + def get_prefix_lm_mask_mod(self) -> _mask_mod_signature:

455 + """Creates the prefix LM mask_mod function for FlexAttention."""

456 +

457 + assert self.doc_ids is not None

458 + request_lookup = self.doc_ids

459 +

460 + def prefix_lm_mask_mod(

461 + b: torch.Tensor,

462 + h: torch.Tensor,

463 + cu_q_idx: torch.Tensor,

464 + q_idx: torch.Tensor,

465 + kv_idx: torch.Tensor,

466 + ):

467 + mask = torch.zeros_like(q_idx, dtype=torch.bool)

468 + for req, doc_range_lst in (self.mm_prefix_range or {}).items():

469 + req_mask = request_lookup[cu_q_idx] == req

470 + for start, end in doc_range_lst:

471 + doc_mask_q = (q_idx >= start) & (q_idx <= end)

472 + doc_mask_kv = (kv_idx >= start) & (kv_idx <= end)

473 + mask = mask | (req_mask & doc_mask_q & doc_mask_kv)

474 + return mask

475 +

476 + def final_mask_mod(

477 + b: torch.Tensor,

478 + h: torch.Tensor,

479 + q_idx: torch.Tensor,

480 + physical_kv_idx: torch.Tensor,

481 + ) -> torch.Tensor:

482 + (is_valid, logical_q_idx, logical_kv_idx) = (

483 + self._convert_physical_to_logical(self.doc_ids, q_idx, physical_kv_idx)

484 + )

485 + return torch.where(

486 + is_valid,

487 + prefix_lm_mask_mod(b, h, q_idx, logical_q_idx, logical_kv_idx),

488 + False,

489 + )

490 +

491 + return final_mask_mod

492 +

493 def get_mask_mod(self):

494 # Stage-1: initialize the base mask_mod

495 # (causal mask for decoder or bidirectional mask for encoder)

@@ -456,6 +503,10 @@ class FlexAttentionMetadata:

503 # Add sliding window mask for sliding window attention

504 sliding_window_mask_mod = self.get_sliding_window_mask_mod()

505 mask_mod = and_masks(mask_mod, sliding_window_mask_mod)

506 + if self.mm_prefix_range:

507 + # Add prefix LM mask for vision-language prefix LM attention

508 + prefix_lm_mask_mod = self.get_prefix_lm_mask_mod()

509 + mask_mod = or_masks(mask_mod, prefix_lm_mask_mod)

510 return mask_mod

511

512 def get_transformed_score_mod(self) -> _score_mod_signature | None:

@@ -709,6 +760,7 @@ class FlexAttentionImpl(AttentionImpl):

760 sliding_window: int | None

761 alibi_slopes: torch.Tensor | None

762 logits_soft_cap: float | None

763 + mm_prefix_range: dict[int, list[tuple[int, int]]] | None = None

764

765 def __init__(

766 self,

@@ -810,11 +862,21 @@ class FlexAttentionImpl(AttentionImpl):

862

863 num_actual_tokens = attn_metadata.num_actual_tokens

864

865 + needs_rebuild_block_mask = False

866 if attn_metadata.sliding_window != self.sliding_window:

867 attn_metadata.sliding_window = self.sliding_window

868 if attn_metadata.direct_build:

869 # update mask mod in attention metadata

870 attn_metadata.mask_mod = attn_metadata.get_mask_mod()

871 + needs_rebuild_block_mask = True

872 +

873 + if self.mm_prefix_range != getattr(attn_metadata, "mm_prefix_range", None):

874 + self.mm_prefix_range = attn_metadata.mm_prefix_range

875 + attn_metadata.mask_mod = attn_metadata.get_mask_mod()

876 + needs_rebuild_block_mask = True

877 +

878 + if needs_rebuild_block_mask:

879 + if attn_metadata.direct_build and attn_metadata.causal:

880 attn_metadata.block_mask = attn_metadata._build_block_mask_direct()

881 else:

882 attn_metadata.block_mask = attn_metadata.build_block_mask()

vllm/v1/worker/gpu_model_runner.py

+25 -1

@@ -48,7 +48,10 @@ from vllm.distributed.parallel_state import (

48 is_global_first_rank,

49 prepare_communication_buffer_for_model,

50 )

51 -from vllm.forward_context import BatchDescriptor, set_forward_context

51 +from vllm.forward_context import (

52 + BatchDescriptor,

53 + set_forward_context,

54 +)

55 from vllm.logger import init_logger

56 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

57 from vllm.model_executor.layers.rotary_embedding import (

@@ -329,6 +332,7 @@ class GPUModelRunner(

332 self.use_alibi = model_config.uses_alibi

333

334 self.cascade_attn_enabled = not self.model_config.disable_cascade_attn

335 + self.is_mm_prefix_lm = self.model_config.is_mm_prefix_lm

336

337 # Multi-modal data support

338 self.mm_registry = MULTIMODAL_REGISTRY

@@ -1700,6 +1704,26 @@ class GPUModelRunner(

1704 for layer_name in attn_group.layer_names:

1705 attn_metadata[layer_name] = attn_metadata_i

1706

1707 + if self.is_mm_prefix_lm:

1708 + req_doc_ranges = {}

1709 + for req_id in self.input_batch.req_ids:

1710 + image_doc_ranges = []

1711 + req_state = self.requests[req_id]

1712 + for mm_feature in req_state.mm_features:

1713 + pos_info = mm_feature.mm_position

1714 + img_doc_range = pos_info.extract_embeds_range()

1715 + image_doc_ranges.extend(img_doc_range)

1716 + req_idx = self.input_batch.req_id_to_index[req_id]

1717 + req_doc_ranges[req_idx] = image_doc_ranges

1718 +

1719 + if isinstance(attn_metadata, list):

1720 + for ub_metadata in attn_metadata:

1721 + for _metadata in ub_metadata.values():

1722 + _metadata.mm_prefix_range = req_doc_ranges # type: ignore[attr-defined]

1723 + else:

1724 + for _metadata in attn_metadata.values():

1725 + _metadata.mm_prefix_range = req_doc_ranges # type: ignore[attr-defined]

1726 +

1727 if spec_decode_common_attn_metadata is not None and (

1728 num_reqs != num_reqs_padded or num_tokens != num_tokens_padded

1729 ):