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, 19:00 UTC
[MoE] Fix output_shape calculation in Attention layer to handle 3D query inputs (#31596)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jan 2, 2026 at 09:46 UTC
6ef770df7c3f0d135c2f3a594c461949113aae91
2 files changed
+17
-2
vllm/attention/layer.py
+4
-1
@@ -357,8 +357,11 @@ class Attention(nn.Module, AttentionLayerBase):
357
358
if self.use_output:
359
if output_shape is None:
360
+ # Handle both 2D [num_tokens, hidden] and
361
+ # 3D [num_tokens, heads, head_dim] query
362
+ num_tokens = query.shape[0]
363
output_shape = torch.Size(
361
- (*query.shape[:-1], self.num_heads * self.head_size_v)
364
+ (num_tokens, self.num_heads * self.head_size_v)
365
)
366
output_shape = output_shape if output_shape is not None else query.shape
367
output = torch.empty(output_shape, dtype=output_dtype, device=query.device)
vllm/model_executor/layers/quantization/fp8.py
+13
-1
@@ -180,7 +180,19 @@ def get_fp8_moe_backend(
180
scope="local",
181
)
182
183
- if envs.VLLM_USE_DEEP_GEMM and moe_use_deep_gemm and block_quant:
183
+ # Determine if we should use DeepGEMM (top-level enable switch)
184
+ # - If explicitly set by user, respect their choice
185
+ # - If not platform supports DeepGEMM, disable it
186
+ # This helps avoid warning messages on unsupported platforms.
187
+ use_deep_gemm = envs.VLLM_USE_DEEP_GEMM
188
+ if not is_deep_gemm_supported():
189
+ use_deep_gemm = False
190
+ logger.info_once(
191
+ "DeepGEMM is disabled because the platform does not support it.",
192
+ scope="local",
193
+ )
194
+
195
+ if use_deep_gemm and moe_use_deep_gemm and block_quant:
196
if not has_deep_gemm():
197
logger.warning_once(
198
"DeepGEMM backend requested but not available.", scope="local"