@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 28b484e5d7

[Model] Pixtral: use packed multimodal encoder attention (#52185)

Signed-off-by: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Oliver Holworthy committed Aug 26, 2026 at 15:14 UTC 28b484e5d741519feb72e5c0bf6ebc3b08176292

2 files changed +157 -71

tests/models/multimodal/generation/test_pixtral.py

+34

@@ -5,6 +5,7 @@ from dataclasses import asdict

5 from typing import TYPE_CHECKING, Any

7 import pytest

8 +import torch

9 from mistral_common.multimodal import download_image

10 from mistral_common.protocol.instruct.chunk import ImageURLChunk

11 from mistral_common.protocol.instruct.request import ChatCompletionRequest

@@ -15,7 +16,9 @@ from transformers import AutoProcessor

16 from vllm import SamplingParams, TextPrompt, TokensPrompt

17 from vllm.inputs import MultiModalDataBuiltins

18 from vllm.logprobs import Logprob, SampleLogprobs

19 +from vllm.model_executor.models.pixtral import _make_packed_sequence_metadata

20 from vllm.platforms import current_platform

21 +from vllm.v1.attention.backends.registry import AttentionBackendEnum

23 from ....utils import VLLM_PATH, large_gpu_test

24 from ...utils import check_logprobs_close

@@ -123,6 +126,37 @@ FIXTURE_LOGPROBS_CHAT = {

126 OutputsLogprobs = list[tuple[list[int], str, SampleLogprobs | None]]

127

128

129 +@pytest.mark.parametrize(

130 + "backend",

131 + [

132 + AttentionBackendEnum.FLASH_ATTN,

133 + AttentionBackendEnum.FLASHINFER,

134 + AttentionBackendEnum.TORCH_SDPA,

135 + ],

136 +)

137 +def test_packed_sequence_metadata(backend: AttentionBackendEnum) -> None:

138 + cu_seqlens, max_seqlen, sequence_lengths = _make_packed_sequence_metadata(

139 + [4, 6],

140 + backend,

141 + hidden_size=64,

142 + tp_size=1,

143 + device=torch.device("cpu"),

144 + )

145 +

146 + assert cu_seqlens.dtype == torch.int32

147 + assert max_seqlen.dtype == torch.int32

148 + if backend == AttentionBackendEnum.FLASHINFER:

149 + assert max_seqlen.item() >= 6

150 + assert sequence_lengths is not None

151 + assert sequence_lengths.dtype == torch.int32

152 + assert len(cu_seqlens) % 2 == 0

153 + else:

154 + expected_max_seqlen = 6 if backend == AttentionBackendEnum.FLASH_ATTN else 0

155 + assert max_seqlen.item() == expected_max_seqlen

156 + assert cu_seqlens.tolist() == [0, 4, 10]

157 + assert sequence_lengths is None

158 +

159 +

160 # For the test author to store golden output in JSON

161 def _dump_outputs_w_logprobs(

162 outputs: OutputsLogprobs,

vllm/model_executor/models/pixtral.py

+123 -71

@@ -6,6 +6,7 @@ from collections.abc import Iterable, Mapping, Sequence

6 from dataclasses import dataclass, fields

7 from typing import Annotated, Literal

9 +import numpy as np

10 import torch

11 import torch.nn as nn

12 from mistral_common.protocol.instruct.chunk import ImageChunk, TextChunk

@@ -26,6 +27,7 @@ from vllm.config.multimodal import BaseDummyOptions

27 from vllm.distributed import divide, get_tensor_model_parallel_world_size

28 from vllm.inputs import MultiModalDataDict

29 from vllm.model_executor.layers.activation import SiluAndMul, get_act_and_mul_fn

30 +from vllm.model_executor.layers.attention import MMEncoderAttention

31 from vllm.model_executor.layers.conv import Conv2dLayer

32 from vllm.model_executor.layers.layernorm import RMSNorm

33 from vllm.model_executor.layers.linear import (

@@ -59,7 +61,6 @@ from vllm.multimodal.processing.processor import (

61 PromptUpdate,

62 PromptUpdateDetails,

63 )

62 -from vllm.platforms import current_platform

64 from vllm.sequence import IntermediateTensors

65 from vllm.tokenizers import cached_tokenizer_from_config

66 from vllm.tokenizers.mistral import MistralTokenizer

@@ -69,6 +70,7 @@ from vllm.transformers_utils.processors.pixtral import (

70 )

71 from vllm.utils.collection_utils import is_list_of

72 from vllm.utils.tensor_schema import TensorSchema, TensorShape

73 +from vllm.v1.attention.backends.registry import AttentionBackendEnum

75 from .interfaces import (

76 MultiModalEmbeddings,

@@ -87,19 +89,35 @@ from .vision import (

89 resolve_visual_encoder_outputs,

90 )

90 -try:

91 - # Note: vLLM does not install xformers by default.

92 - from xformers import ops as xops

92 +PATCH_MERGE = "patch_merge"

94 - if current_platform.is_cuda() and current_platform.has_device_capability(100):

95 - # Xformers FA is not compatible with B200

96 - USE_XFORMERS_OPS = False

97 - else:

98 - USE_XFORMERS_OPS = True

99 -except ImportError:

100 - USE_XFORMERS_OPS = False

102 -PATCH_MERGE = "patch_merge"

95 +def _make_packed_sequence_metadata(

96 + sequence_lengths: list[int],

97 + attn_backend: AttentionBackendEnum,

98 + hidden_size: int,

99 + tp_size: int,

100 + device: torch.device,

101 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:

102 + lengths = np.array(sequence_lengths, dtype=np.int32)

103 + cu_seqlens = np.concatenate(

104 + [np.zeros(1, dtype=np.int32), lengths.cumsum(dtype=np.int32)]

105 + )

106 + sequence_lengths_tensor = MMEncoderAttention.maybe_compute_seq_lens(

107 + attn_backend, cu_seqlens, device

108 + )

109 + max_seqlen = torch.tensor(

110 + MMEncoderAttention.compute_max_seqlen(attn_backend, cu_seqlens),

111 + dtype=torch.int32,

112 + )

113 + cu_seqlens_tensor = MMEncoderAttention.maybe_recompute_cu_seqlens(

114 + attn_backend,

115 + cu_seqlens,

116 + hidden_size,

117 + tp_size,

118 + device,

119 + )

120 + return cu_seqlens_tensor, max_seqlen, sequence_lengths_tensor

121

122

123 def _is_layer_none_or_staged(layer: nn.Module) -> bool:

@@ -766,12 +784,19 @@ class Attention(nn.Module):

784

785 tp_size = 1 if disable_tp else get_tensor_model_parallel_world_size()

786 self.n_heads = divide(args.num_attention_heads, tp_size)

787 + self.attn = MMEncoderAttention(

788 + num_heads=self.n_heads,

789 + head_size=self.head_dim,

790 + prefix=f"{prefix}.attn",

791 + )

792

793 def forward(

794 self,

795 x: torch.Tensor,

773 - mask: torch.Tensor,

796 freqs_cis: torch.Tensor,

797 + cu_seqlens: torch.Tensor,

798 + max_seqlen: torch.Tensor,

799 + sequence_lengths: torch.Tensor | None,

800 ) -> torch.Tensor:

801 batch, patches, _ = x.shape

802

@@ -782,15 +807,14 @@ class Attention(nn.Module):

807 v = v.reshape(batch, patches, self.n_heads, self.head_dim)

808

809 q, k = apply_rotary_emb_vit(q, k, freqs_cis=freqs_cis)

785 -

786 - if USE_XFORMERS_OPS:

787 - out = xops.memory_efficient_attention(q, k, v, attn_bias=mask)

788 - else:

789 - q = q.transpose(1, 2)

790 - k = k.transpose(1, 2)

791 - v = v.transpose(1, 2)

792 - out = nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask)

793 - out = out.transpose(1, 2)

810 + out = self.attn(

811 + q,

812 + k,

813 + v,

814 + cu_seqlens=cu_seqlens,

815 + max_seqlen=max_seqlen,

816 + sequence_lengths=sequence_lengths,

817 + )

818

819 out = out.reshape(batch, patches, self.n_heads * self.head_dim)

820 out, _ = self.o_proj(out)

@@ -825,11 +849,17 @@ class TransformerBlock(nn.Module):

849 def forward(

850 self,

851 x: torch.Tensor,

828 - mask: torch.Tensor,

852 freqs_cis: torch.Tensor,

853 + cu_seqlens: torch.Tensor,

854 + max_seqlen: torch.Tensor,

855 + sequence_lengths: torch.Tensor | None,

856 ) -> torch.Tensor:

857 r = self.attention.forward(

832 - self.attention_norm(x), mask=mask, freqs_cis=freqs_cis

858 + self.attention_norm(x),

859 + freqs_cis=freqs_cis,

860 + cu_seqlens=cu_seqlens,

861 + max_seqlen=max_seqlen,

862 + sequence_lengths=sequence_lengths,

863 )

864 h = x + r

865 r = self.feed_forward.forward(self.ffn_norm(h))

@@ -860,11 +890,19 @@ class Transformer(nn.Module):

890 def forward(

891 self,

892 x: torch.Tensor,

863 - mask: torch.Tensor,

893 freqs_cis: torch.Tensor | None,

894 + cu_seqlens: torch.Tensor,

895 + max_seqlen: torch.Tensor,

896 + sequence_lengths: torch.Tensor | None,

897 ) -> torch.Tensor:

898 for layer in self.layers:

867 - x = layer(x, mask=mask, freqs_cis=freqs_cis)

899 + x = layer(

900 + x,

901 + freqs_cis=freqs_cis,

902 + cu_seqlens=cu_seqlens,

903 + max_seqlen=max_seqlen,

904 + sequence_lengths=sequence_lengths,

905 + )

906 return x

907

908

@@ -971,20 +1009,21 @@ class VisionTransformer(nn.Module):

1009 positions = position_meshgrid(patch_embeds_list).to(self.device)

1010 freqs_cis = self.freqs_cis[positions[:, 0], positions[:, 1]]

1011

974 - # pass through Transformer with a block diagonal mask delimiting images

975 - if USE_XFORMERS_OPS:

976 - mask = xops.fmha.attn_bias.BlockDiagonalMask.from_seqlens(

977 - [p.shape[-2] * p.shape[-1] for p in patch_embeds_list],

978 - )

979 - else:

980 - from transformers.models.pixtral.modeling_pixtral import (

981 - generate_block_attention_mask,

982 - )

983 -

984 - mask = generate_block_attention_mask(

985 - [p.shape[-2] * p.shape[-1] for p in patch_embeds_list], patch_embeds

986 - )

987 - out = self.transformer(patch_embeds, mask=mask, freqs_cis=freqs_cis)

1012 + attention = self.transformer.layers[0].attention.attn

1013 + cu_seqlens, max_seqlen, sequence_lengths = _make_packed_sequence_metadata(

1014 + embed_sizes,

1015 + attention.attn_backend,

1016 + self.args.hidden_size,

1017 + 1 if is_vit_use_data_parallel() else get_tensor_model_parallel_world_size(),

1018 + patch_embeds.device,

1019 + )

1020 + out = self.transformer(

1021 + patch_embeds,

1022 + freqs_cis=freqs_cis,

1023 + cu_seqlens=cu_seqlens,

1024 + max_seqlen=max_seqlen,

1025 + sequence_lengths=sequence_lengths,

1026 + )

1027

1028 # squeeze dim 0 and split into separate tensors for each image

1029 return torch.split(out.squeeze(0), embed_sizes)

@@ -1244,36 +1283,39 @@ class PixtralHFAttention(nn.Module):

1283 1 if use_data_parallel else get_tensor_model_parallel_world_size()

1284 )

1285 self.n_heads = divide(config.num_attention_heads, self.tp_size)

1286 + self.attn = MMEncoderAttention(

1287 + num_heads=self.n_heads,

1288 + head_size=self.head_dim,

1289 + prefix=f"{prefix}.attn",

1290 + )

1291

1292 def forward(

1293 self,

1294 hidden_states: torch.Tensor,

1251 - attention_mask: torch.Tensor,

1295 position_embeddings: torch.Tensor,

1296 + cu_seqlens: torch.Tensor,

1297 + max_seqlen: torch.Tensor,

1298 + sequence_lengths: torch.Tensor | None,

1299 ) -> tuple[torch.Tensor, torch.Tensor | None]:

1300 batch, patches, _ = hidden_states.size()

1301

1302 qkv_states, _ = self.qkv_proj(hidden_states)

1303 q, k, v = qkv_states.chunk(3, dim=-1)

1304

1259 - # Transpose q and k to apply HF's Rotary Position Embedding

1305 + # Transpose q and k to apply HF's Rotary Position Embedding.

1306 q = q.view(batch, patches, self.n_heads, self.head_dim).transpose(1, 2)

1307 k = k.view(batch, patches, self.n_heads, self.head_dim).transpose(1, 2)

1308 v = v.view(batch, patches, self.n_heads, self.head_dim)

1309 cos, sin = position_embeddings

1310 q, k = apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=0)

1265 -

1266 - if USE_XFORMERS_OPS:

1267 - # Transpose q and k back for attention

1268 - q = q.transpose(1, 2).contiguous()

1269 - k = k.transpose(1, 2).contiguous()

1270 - out = xops.memory_efficient_attention(q, k, v, attn_bias=attention_mask)

1271 - else:

1272 - v = v.transpose(1, 2)

1273 - out = nn.functional.scaled_dot_product_attention(

1274 - q, k, v, attn_mask=attention_mask

1275 - )

1276 - out = out.transpose(1, 2)

1311 + out = self.attn(

1312 + q.transpose(1, 2).contiguous(),

1313 + k.transpose(1, 2).contiguous(),

1314 + v,

1315 + cu_seqlens=cu_seqlens,

1316 + max_seqlen=max_seqlen,

1317 + sequence_lengths=sequence_lengths,

1318 + )

1319

1320 out = out.reshape(batch, patches, self.n_heads * self.head_dim)

1321 attn_output, _ = self.o_proj(out)

@@ -1307,13 +1349,17 @@ class PixtralHFTransformerBlock(nn.Module):

1349 def forward(

1350 self,

1351 hidden_states: torch.Tensor,

1310 - attention_mask: torch.Tensor,

1352 position_embeddings: torch.Tensor,

1353 + cu_seqlens: torch.Tensor,

1354 + max_seqlen: torch.Tensor,

1355 + sequence_lengths: torch.Tensor | None,

1356 ) -> torch.Tensor:

1357 r, _ = self.attention.forward(

1358 self.attention_norm(hidden_states),

1315 - attention_mask=attention_mask,

1359 position_embeddings=position_embeddings,

1360 + cu_seqlens=cu_seqlens,

1361 + max_seqlen=max_seqlen,

1362 + sequence_lengths=sequence_lengths,

1363 )

1364 h = hidden_states + r

1365 r = self.feed_forward.forward(self.ffn_norm(h))

@@ -1351,14 +1397,22 @@ class PixtralHFTransformer(nn.Module):

1397 def forward(

1398 self,

1399 x: torch.Tensor,

1354 - attention_mask: torch.Tensor,

1400 position_embeddings: torch.Tensor,

1401 + cu_seqlens: torch.Tensor,

1402 + max_seqlen: torch.Tensor,

1403 + sequence_lengths: torch.Tensor | None,

1404 return_all_hidden_states: bool,

1405 ) -> torch.Tensor:

1406 hidden_states_pool = [x]

1407

1408 for layer in self.layers:

1361 - x = layer(x, attention_mask, position_embeddings)

1409 + x = layer(

1410 + x,

1411 + position_embeddings,

1412 + cu_seqlens,

1413 + max_seqlen,

1414 + sequence_lengths,

1415 + )

1416 if return_all_hidden_states:

1417 hidden_states_pool.append(x)

1418 # If we have multiple feature sample layers, we return all hidden

@@ -1453,23 +1507,21 @@ class PixtralHFVisionModel(nn.Module):

1507 ).to(self.device)

1508 position_embedding = self.patch_positional_embedding(patch_embeds, position_ids)

1509

1456 - if USE_XFORMERS_OPS:

1457 - attention_mask = xops.fmha.attn_bias.BlockDiagonalMask.from_seqlens(

1458 - [p.shape[-2] * p.shape[-1] for p in patch_embeds_list],

1459 - )

1460 - else:

1461 - from transformers.models.pixtral.modeling_pixtral import (

1462 - generate_block_attention_mask,

1463 - )

1464 -

1465 - attention_mask = generate_block_attention_mask(

1466 - [p.shape[-2] * p.shape[-1] for p in patch_embeds_list], patch_embeds

1467 - )

1510 + attention = self.transformer.layers[0].attention.attn

1511 + cu_seqlens, max_seqlen, sequence_lengths = _make_packed_sequence_metadata(

1512 + embed_sizes,

1513 + attention.attn_backend,

1514 + self.config.hidden_size,

1515 + self.transformer.layers[0].attention.tp_size,

1516 + patch_embeds.device,

1517 + )

1518

1519 out = self.transformer(

1520 patch_embeds,

1471 - attention_mask,

1521 position_embedding,

1522 + cu_seqlens,

1523 + max_seqlen,

1524 + sequence_lengths,

1525 return_all_hidden_states=select_layers is not None,

1526 )

1527