[MoE] Fix output_shape calculation in Attention layer to handle 3D query inputs

[AndreasKaratzas]

Fixes a regression introduced in #28775 where the output_shape calculation in Attention.forward() assumes 2D query input, causing failures when models pass 3D query tensors.

Problem

The new code added in #28775:

if output_shape is None:
    output_shape = torch.Size((*query.shape[:-1], self.num_heads * self.head_size_v))

This breaks when query is 3D [num_tokens, num_heads, head_dim]:

• query.shape[:-1] = (num_tokens, num_heads)
• output_shape = (num_tokens, num_heads, hidden_size) ← incorrect

This causes DeepseekV2Attention (used by DeepSeek-V2/V3 with MLA disabled and MTP layers) to produce incorrect output shapes, leading to downstream failures in FP8 quantization kernels.

Fix

Use query.shape[0] to always get num_tokens, making the calculation robust to both 2D and 3D inputs:

if output_shape is None:
    num_tokens = query.shape[0]
    output_shape = torch.Size((num_tokens, self.num_heads * self.head_size_v))

Testing

On ROCm (8 x mi325 machine):

• pytest -v -s tests/v1/e2e/test_spec_decode.py::test_mtp_correctness[deepseek]
• pytest -v -s tests/v1/e2e/test_spec_decode.py::test_eagle_correctness[TRITON_ATTN-deepseek_eagle]

Fixes issue observed on ROCm but the bug exists on all platforms.

[AndreasKaratzas]

cc @DarkLight1337 @yt0428

[gemini-code-assist]

Code Review

The pull request modifies the vllm/attention/layer.py file to adjust the output_shape calculation in the forward method, explicitly extracting num_tokens from the query shape to handle both 2D and 3D query inputs. Additionally, in vllm/model_executor/layers/quantization/fp8.py, the DeepGEMM backend activation logic was updated to first respect the user's VLLM_USE_DEEP_GEMM setting and then disable DeepGEMM if the platform does not support it, logging an informational message. A review comment highlighted that the DeepGEMM logging logic could be misleading; if a user explicitly disables DeepGEMM, the current message incorrectly attributes the disabling to platform incompatibility. The suggested fix proposes to only log the platform incompatibility message if DeepGEMM was initially requested by the user but is not supported by the platform.

[gemini-code-assist]

The current logic for checking DeepGEMM support can produce a misleading log message. If a user explicitly disables DeepGEMM by setting VLLM_USE_DEEP_GEMM=0, is_deep_gemm_supported() will return False, causing the message "DeepGEMM is disabled because the platform does not support it" to be logged. This is inaccurate because the user disabled it, not the platform.

The check should only log a message if the user intended to use DeepGEMM, but it's not supported by the platform. I've suggested a change to correct this logic and make the log message more precise.

Suggested change

	if not is_deep_gemm_supported():
	use_deep_gemm = False
	logger.info_once(
	"DeepGEMM is disabled because the platform does not support it.",
	scope="local",
	)
	if use_deep_gemm and not is_deep_gemm_supported():
	use_deep_gemm = False
	logger.info_once(
	"DeepGEMM was requested but is disabled because the platform does not support it.",
	scope="local",
	)

[AndreasKaratzas]

This is effectively the next check that is performed. And the message in the next if statement is the same with the proposed one. So I think this modification is unnecessary.

[LucasWilkinson]

These changes are unrelated to the intent of the PR; why did you add this?

[AndreasKaratzas]

I can move it to a different PR if that's what you are asking. On ROCm right now the message logged during a run is that DeepGemm is requested but not found, which is not that accurate because DeepGemm is not a ROCm supported feature. So I put together this short block that renders a more precise check first.

[LucasWilkinson]

got it, thanks 👍

[LucasWilkinson]

got it, thanks 👍

[LucasWilkinson]

LGTM, thanks for the contribution!