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

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 10, 12:45 UTC

[MM Encoder]: Make MMEncoderAttention's `scale` takes effect properly (#31950)

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

Isotr0py committed Jan 8, 2026 at 18:33 UTC 2972a05473d64a21c7390a5676bbb07bd8d2c8da

11 files changed +32 -8

vllm/attention/layers/mm_encoder_attention.py

@@ -133,6 +133,7 @@ class MMEncoderAttention(CustomOp):

133 q=query,

134 k=key,

135 v=value,

136 + scale=self.scale,

137 cu_seqlens=cu_seqlens,

138 )

139 if is_reshaped:

@@ -167,6 +168,7 @@ class MMEncoderAttention(CustomOp):

168 q=query,

169 k=key,

170 v=value,

171 + scale=self.scale,

172 cu_seqlens=cu_seqlens,

173 max_seqlen=max_seqlen,

174 batch_size=bsz,

vllm/attention/ops/vit_attn_wrappers.py

+21 -8

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

27 batch_size: int,

28 is_rocm_aiter: bool,

29 fa_version: int | None,

30 + scale: float | None = None,

31 cu_seqlens: torch.Tensor | None = None,

32 max_seqlen: torch.Tensor | None = None,

33 ) -> torch.Tensor:

@@ -57,6 +58,7 @@ def flash_attn_maxseqlen_wrapper(

58 max_seqlen_k=max_seqlen,

59 dropout_p=0.0,

60 causal=False,

61 + softmax_scale=scale,

62 **kwargs,

63 )

64 context_layer = einops.rearrange(output, "(b s) h d -> b s h d", b=batch_size)

@@ -67,11 +69,12 @@ def flash_attn_maxseqlen_wrapper_fake(

69 q: torch.Tensor,

70 k: torch.Tensor,

71 v: torch.Tensor,

70 - cu_seqlens: torch.Tensor,

71 - max_seqlen: torch.Tensor,

72 batch_size: int,

73 is_rocm_aiter: bool,

74 fa_version: int | None,

75 + scale: float | None,

76 + cu_seqlens: torch.Tensor | None,

77 + max_seqlen: torch.Tensor | None,

78 ) -> torch.Tensor:

79 return torch.empty_like(q)

@@ -90,6 +93,7 @@ def vit_flash_attn_wrapper(

93 batch_size: int,

94 is_rocm_aiter: bool,

95 fa_version: int | None,

96 + scale: float | None = None,

97 cu_seqlens: torch.Tensor | None = None,

98 max_seqlen: torch.Tensor | None = None,

99 ) -> torch.Tensor:

@@ -100,18 +104,24 @@ def vit_flash_attn_wrapper(

104 batch_size,

105 is_rocm_aiter,

106 fa_version,

107 + scale,

108 cu_seqlens,

109 max_seqlen,

110 )

111

112

108 -def apply_sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:

113 +def apply_sdpa(

114 + q: torch.Tensor,

115 + k: torch.Tensor,

116 + v: torch.Tensor,

117 + scale: float | None = None,

118 +) -> torch.Tensor:

119 """

120 Input shape:

121 (batch_size x seq_len x num_heads x head_size)

122 """

123 q, k, v = (einops.rearrange(x, "b s h d -> b h s d") for x in [q, k, v])

114 - output = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0)

124 + output = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, scale=scale)

125 output = einops.rearrange(output, "b h s d -> b s h d ")

126 return output

127

@@ -122,6 +132,7 @@ def torch_sdpa_wrapper(

132 q: torch.Tensor,

133 k: torch.Tensor,

134 v: torch.Tensor,

135 + scale: float | None = None,

136 cu_seqlens: torch.Tensor | None = None,

137 ) -> torch.Tensor:

138 # Never remove the contiguous logic for ROCm

@@ -132,7 +143,7 @@ def torch_sdpa_wrapper(

143 v = v.contiguous()

144

145 if cu_seqlens is None:

135 - return apply_sdpa(q, k, v)

146 + return apply_sdpa(q, k, v, scale=scale)

147

148 outputs = []

149

@@ -141,7 +152,7 @@ def torch_sdpa_wrapper(

152 k_chunks = torch.split(k, lens, dim=1)

153 v_chunks = torch.split(v, lens, dim=1)

154 for q_i, k_i, v_i in zip(q_chunks, k_chunks, v_chunks):

144 - output_i = apply_sdpa(q_i, k_i, v_i)

155 + output_i = apply_sdpa(q_i, k_i, v_i, scale=scale)

156 outputs.append(output_i)

157 context_layer = torch.cat(outputs, dim=1)

158 return context_layer

@@ -151,7 +162,8 @@ def torch_sdpa_wrapper_fake(

162 q: torch.Tensor,

163 k: torch.Tensor,

164 v: torch.Tensor,

154 - cu_seqlens: torch.Tensor,

165 + scale: float | None,

166 + cu_seqlens: torch.Tensor | None,

167 ) -> torch.Tensor:

168 return torch.empty_like(q)

169

@@ -167,6 +179,7 @@ def vit_torch_sdpa_wrapper(

179 q: torch.Tensor,

180 k: torch.Tensor,

181 v: torch.Tensor,

182 + scale: float | None = None,

183 cu_seqlens: torch.Tensor | None = None,

184 ) -> torch.Tensor:

172 - return torch.ops.vllm.torch_sdpa_wrapper(q, k, v, cu_seqlens)

185 + return torch.ops.vllm.torch_sdpa_wrapper(q, k, v, scale, cu_seqlens)

vllm/model_executor/models/dots_ocr.py

@@ -271,6 +271,7 @@ class DotsVisionAttention(nn.Module):

271 self.attn = MMEncoderAttention(

272 num_heads=self.num_attention_heads_per_partition,

273 head_size=self.hidden_size_per_attention_head,

274 + scale=self.hidden_size_per_attention_head**-0.5,

275 multimodal_config=multimodal_config,

276 prefix=f"{prefix}.attn",

277 )

vllm/model_executor/models/ernie45_vl.py

@@ -152,6 +152,7 @@ class Ernie4_5_VisionAttention(nn.Module):

152 self.attn = MMEncoderAttention(

153 num_heads=self.num_attention_heads_per_partition,

154 head_size=self.hidden_size_per_attention_head,

155 + scale=self.hidden_size_per_attention_head**-0.5,

156 multimodal_config=multimodal_config,

157 prefix=f"{prefix}.attn",

158 )

vllm/model_executor/models/glm4_1v.py

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

304 self.attn = MMEncoderAttention(

305 num_heads=self.num_attention_heads_per_partition,

306 head_size=self.hidden_size_per_attention_head,

307 + scale=self.hidden_size_per_attention_head**-0.5,

308 multimodal_config=multimodal_config,

309 )

310

vllm/model_executor/models/glmasr.py

@@ -188,6 +188,7 @@ class GlmAsrEncoderAttention(nn.Module):

188 self.attn = MMEncoderAttention(

189 num_heads=self.num_heads_per_rank,

190 head_size=self.head_dim,

191 + scale=self.head_dim**-0.5,

192 num_kv_heads=self.num_kv_heads_per_rank,

193 prefix=f"{prefix}.attn",

194 )

vllm/model_executor/models/isaac.py

@@ -984,6 +984,7 @@ class Siglip2VisionAttention(nn.Module):

984 self.attn = MMEncoderAttention(

985 num_heads=self.num_attention_heads_per_partition,

986 head_size=self.hidden_size_per_attention_head,

987 + scale=self.hidden_size_per_attention_head**-0.5,

988 prefix=f"{prefix}.attn",

989 multimodal_config=multimodal_config,

990 )

vllm/model_executor/models/moonvit.py

@@ -390,6 +390,7 @@ class MoonVitEncoderLayer(nn.Module):

390 self.attn = MMEncoderAttention(

391 num_heads=self.num_attention_heads_per_partition,

392 head_size=self.hidden_size_per_attention_head,

393 + scale=self.hidden_size_per_attention_head**-0.5,

394 multimodal_config=multimodal_config,

395 prefix=f"{prefix}.attn",

396 )

vllm/model_executor/models/paddleocr_vl.py

@@ -564,6 +564,7 @@ class SiglipAttention(nn.Module):

564 self.attn = MMEncoderAttention(

565 num_heads=self.num_attention_heads_per_partition,

566 head_size=self.hidden_size_per_attention_head,

567 + scale=self.hidden_size_per_attention_head**-0.5,

568 multimodal_config=multimodal_config,

569 prefix=f"{prefix}.attn",

570 )

vllm/model_executor/models/qwen2_5_vl.py

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

352 self.attn = MMEncoderAttention(

353 num_heads=self.num_attention_heads_per_partition,

354 head_size=self.hidden_size_per_attention_head,

355 + scale=self.hidden_size_per_attention_head**-0.5,

356 multimodal_config=multimodal_config,

357 )

358

vllm/model_executor/models/qwen2_vl.py

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

327 self.attn = MMEncoderAttention(

328 num_heads=self.num_attention_heads_per_partition,

329 head_size=self.hidden_size_per_attention_head,

330 + scale=self.hidden_size_per_attention_head**-0.5,

331 multimodal_config=multimodal_config,

332 )

333