Skip to content

Revert "[Bugfix] Fix scaled_mm output narrowing for 3D input tensors" (#38093) #40423

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
vllm-agent wants to merge 1 commit into from
Closed

Revert "[Bugfix] Fix scaled_mm output narrowing for 3D input tensors" (#38093) #40423

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 4 additions & 15 deletions vllm/model_executor/kernels/linear/scaled_mm/pytorch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import math

import torch

Expand All @@ -14,13 +13,6 @@
)


def _get_num_tokens(output_shape: list) -> int:
# torch._scaled_mm works with 2D tensors, so input tensors are
# flattened if they are 3D. If output_shape is 3D, num_tokens is
# the product of all dims except the last (hidden dim).
return math.prod(output_shape[:-1])


class TorchFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
"""
Base class for FP8 linear kernels using Torch.
Expand Down Expand Up @@ -86,8 +78,7 @@ def apply_scaled_mm(
if type(output) is tuple and len(output) == 2:
output = output[0]

num_tokens = _get_num_tokens(output_shape)
return torch.narrow(output, 0, 0, num_tokens).view(*output_shape)
return torch.narrow(output, 0, 0, output_shape[0]).view(*output_shape)

@gemini-code-assist gemini-code-assist Bot Apr 21, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Reverting this change re-introduces a bug for 3D input tensors. When the input x is 3D (e.g., [batch, seq, hidden]), output_shape is [batch, seq, out_features]. Using output_shape[0] as the narrowing size only selects the first batch tokens, whereas the flattened GEMM output contains batch * seq tokens. This will cause a RuntimeError in the subsequent .view(*output_shape) call due to a size mismatch (unless seq == 1). Since the base class FP8ScaledMMLinearKernel.apply_weights explicitly supports 3D tensors by flattening them, this implementation should correctly unpad all tokens using the product of all dimensions except the last one (e.g., math.prod(output_shape[:-1])).



class RowWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
Expand Down Expand Up @@ -154,8 +145,7 @@ def apply_scaled_mm(
bias=bias,
)

num_tokens = _get_num_tokens(output_shape)
return torch.narrow(output, 0, 0, num_tokens).view(*output_shape)
return torch.narrow(output, 0, 0, output_shape[0]).view(*output_shape)

@gemini-code-assist gemini-code-assist Bot Apr 21, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Similar to the issue in PerTensorTorchFP8ScaledMMLinearKernel, using output_shape[0] here is incorrect for 3D tensors and will lead to a shape mismatch error during .view(*output_shape) if the sequence length is greater than 1.



class ChannelWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
Expand Down Expand Up @@ -216,9 +206,8 @@ def apply_scaled_mm(
output = output[0]

# Unpad (undo num_token_padding)
num_tokens = _get_num_tokens(output_shape)
output = torch.narrow(output, 0, 0, num_tokens)
x_scale = torch.narrow(As, 0, 0, num_tokens)
output = torch.narrow(output, 0, 0, output_shape[0])
x_scale = torch.narrow(As, 0, 0, output_shape[0])
Comment on lines +209 to +210

@gemini-code-assist gemini-code-assist Bot Apr 21, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

For 3D tensors, narrowing output and x_scale to output_shape[0] (the batch size) instead of the total number of tokens (batch * seq) is incorrect. This will result in a size mismatch during the dequantization step or the final .view(*output_shape) call.


# DQ
# C = sw * sx * (X * W) + bias
Expand Down