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

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 9, 00:00 UTC

[MM] Pass `prefix` parameter to MMEncoderAttention (#33674)

Signed-off-by: shen-shanshan <467638484@qq.com>

Shanshan Shen committed Feb 3, 2026 at 22:47 UTC 5c4f2dd6ef2009d81f4a765b5c2a7278fc389ef3

15 files changed +58 -11

vllm/model_executor/models/aimv2.py

+4 -1

@@ -127,7 +127,10 @@ class AIMv2Attention(nn.Module):

127 self.num_heads_per_partition = divide(self.num_heads, self.tp_size)

128

129 self.attn = MMEncoderAttention(

130 - self.num_heads_per_partition, self.head_dim, self.scale

130 + self.num_heads_per_partition,

131 + self.head_dim,

132 + self.scale,

133 + prefix=prefix,

134 )

135

136 def forward(self, x: torch.Tensor) -> torch.Tensor:

vllm/model_executor/models/blip.py

+4 -1

@@ -123,7 +123,10 @@ class BlipAttention(nn.Module):

123 self.num_heads_per_partition = divide(self.num_heads, self.tp_size)

124

125 self.attn = MMEncoderAttention(

126 - self.num_heads_per_partition, self.head_dim, self.scale

126 + self.num_heads_per_partition,

127 + self.head_dim,

128 + self.scale,

129 + prefix=prefix,

130 )

131

132 def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):

vllm/model_executor/models/glm4_1v.py

@@ -296,6 +296,7 @@ class Glm4vVisionAttention(nn.Module):

296 num_heads=self.num_attention_heads_per_partition,

297 head_size=self.hidden_size_per_attention_head,

298 scale=self.hidden_size_per_attention_head**-0.5,

299 + prefix=prefix,

300 )

301

302 self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True)

vllm/model_executor/models/glm4v.py

+4 -1

@@ -136,7 +136,10 @@ class EVA2CLIPAttention(nn.Module):

136 )

137

138 self.attn = MMEncoderAttention(

139 - self.num_heads_per_rank, self.head_dim, self.scale

139 + self.num_heads_per_rank,

140 + self.head_dim,

141 + self.scale,

142 + prefix=prefix,

143 )

144 self.output_dropout = torch.nn.Dropout(config.dropout_prob)

145

vllm/model_executor/models/idefics2_vision_model.py

+4 -1

@@ -163,7 +163,10 @@ class Idefics2VisionAttention(nn.Module):

163 )

164 # Use unified MMEncoderAttention with Flash Attention support

165 self.attn = MMEncoderAttention(

166 - self.num_heads_per_partition, self.head_dim, self.scale

166 + self.num_heads_per_partition,

167 + self.head_dim,

168 + self.scale,

169 + prefix=prefix,

170 )

171

172 def forward(

vllm/model_executor/models/intern_vit.py

+4 -1

@@ -212,7 +212,10 @@ class InternParallelAttention(nn.Module):

212 )

213

214 self.attn = MMEncoderAttention(

215 - self.num_heads_per_partition, self.head_dim, self.scale

215 + self.num_heads_per_partition,

216 + self.head_dim,

217 + self.scale,

218 + prefix=prefix,

219 )

220

221 def _apply_qk_norm(self, q: torch.Tensor, k: torch.Tensor):

vllm/model_executor/models/interns1_vit.py

+12 -2

@@ -170,6 +170,7 @@ class InternSdpaAttention(nn.Module):

170 config: PretrainedConfig,

171 *,

172 num_dummy_heads: int = 0,

173 + prefix: str = "",

174 ) -> None:

175 super().__init__()

176

@@ -215,7 +216,12 @@ class InternSdpaAttention(nn.Module):

216 self.projection_layer = nn.Linear(self.dummy_dim, self.embed_dim)

217

218 # Use unified MMEncoderAttention with automatic backend selection

218 - self.attn = MMEncoderAttention(self.num_heads, self.head_dim, self.scale)

219 + self.attn = MMEncoderAttention(

220 + self.num_heads,

221 + self.head_dim,

222 + self.scale,

223 + prefix=prefix,

224 + )

225

226 def forward(self, x: torch.Tensor) -> torch.Tensor:

227 """x shape: (B, N, C)"""

@@ -313,7 +319,11 @@ class InternS1VisionLayer(nn.Module):

319 num_dummy_heads: int,

320 prefix: str = "",

321 ):

316 - return InternSdpaAttention(config, num_dummy_heads=num_dummy_heads)

322 + return InternSdpaAttention(

323 + config,

324 + num_dummy_heads=num_dummy_heads,

325 + prefix=prefix,

326 + )

327

328 def forward(

329 self,

vllm/model_executor/models/mllama4.py

+4 -1

@@ -254,7 +254,10 @@ class Llama4VisionAttention(nn.Module):

254 self.scaling = self.head_dim**-0.5

255

256 self.attn = MMEncoderAttention(

257 - self.num_local_heads, self.head_dim, self.scaling

257 + self.num_local_heads,

258 + self.head_dim,

259 + self.scaling,

260 + prefix=prefix,

261 )

262

263 if use_data_parallel:

vllm/model_executor/models/molmo.py

+5 -1

@@ -231,7 +231,11 @@ class MultiHeadDotProductAttention(nn.Module):

231

232 self.scale = self.head_dim**-0.5

233 self.attn = MMEncoderAttention(

234 - self.num_heads, self.head_dim, self.scale, num_kv_heads=self.num_kv_heads

234 + self.num_heads,

235 + self.head_dim,

236 + self.scale,

237 + num_kv_heads=self.num_kv_heads,

238 + prefix=prefix,

239 )

240

241 def forward(

vllm/model_executor/models/molmo2.py

@@ -611,6 +611,7 @@ class ImagePoolingAttention(nn.Module):

611 self.head_dim,

612 self.scale,

613 num_kv_heads=self.num_kv_heads,

614 + prefix=prefix,

615 )

616

617 def forward_sdpa(

vllm/model_executor/models/qwen2_5_vl.py

@@ -345,6 +345,7 @@ class Qwen2_5_VisionAttention(nn.Module):

345 num_heads=self.num_attention_heads_per_partition,

346 head_size=self.hidden_size_per_attention_head,

347 scale=self.hidden_size_per_attention_head**-0.5,

348 + prefix=prefix,

349 )

350

351 self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True)

vllm/model_executor/models/qwen2_vl.py

@@ -319,6 +319,7 @@ class Qwen2VisionAttention(nn.Module):

319 num_heads=self.num_attention_heads_per_partition,

320 head_size=self.hidden_size_per_attention_head,

321 scale=self.hidden_size_per_attention_head**-0.5,

322 + prefix=prefix,

323 )

324

325 self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True)

vllm/model_executor/models/qwen3_omni_moe_thinker.py

@@ -194,6 +194,7 @@ class Qwen3OmniMoeAudioAttention(nn.Module):

194 num_heads=self.num_local_heads,

195 head_size=self.head_dim,

196 scale=self.scaling,

197 + prefix=prefix,

198 )

199

200 def forward(

vllm/model_executor/models/step3_vl.py

+6 -1

@@ -759,7 +759,12 @@ class Step3VisionAttention(nn.Module):

759 )

760

761 # Use unified MMEncoderAttention with automatic backend selection

762 - self.attn = MMEncoderAttention(self.num_heads, self.head_dim, self.scale)

762 + self.attn = MMEncoderAttention(

763 + self.num_heads,

764 + self.head_dim,

765 + self.scale,

766 + prefix=prefix,

767 + )

768

769 def forward(

770 self,

vllm/model_executor/models/step_vl.py

+6 -1

@@ -220,7 +220,12 @@ class PerceptionEncoderVisionAttention(nn.Module):

220 prefix=f"{prefix}.out_proj",

221 disable_tp=use_data_parallel,

222 )

223 - self.attn = MMEncoderAttention(self.num_heads, self.head_dim, self.scale)

223 + self.attn = MMEncoderAttention(

224 + self.num_heads,

225 + self.head_dim,

226 + self.scale,

227 + prefix=prefix,

228 + )

229 self.rope = PerceptionEncoderRope2D(

230 dim=self.head_dim,

231 max_grid_height=max_grid_height,