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 4, 22:30 UTC
[Bugfix] Fix scaled_mm output narrowing for 3D input tensors (#38093)
Signed-off-by: nemanjaudovic <nudovic@amd.com>
nemanjaudovic committed
Apr 20, 2026 at 10:58 UTC
58631d7c3f9984f979e3cd5abf20a61590b36b1f
1 file changed
+15
-4
vllm/model_executor/kernels/linear/scaled_mm/pytorch.py
+15
-4
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import math
5
6
import torch
7
@@ -13,6 +14,13 @@ from .ScaledMMLinearKernel import (
14
)
15
16
17
+def _get_num_tokens(output_shape: list) -> int:
18
+ # torch._scaled_mm works with 2D tensors, so input tensors are
19
+ # flattened if they are 3D. If output_shape is 3D, num_tokens is
20
+ # the product of all dims except the last (hidden dim).
21
+ return math.prod(output_shape[:-1])
22
+
23
+
24
class TorchFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
25
"""
26
Base class for FP8 linear kernels using Torch.
@@ -78,7 +86,8 @@ class PerTensorTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
86
if type(output) is tuple and len(output) == 2:
87
output = output[0]
88
81
- return torch.narrow(output, 0, 0, output_shape[0]).view(*output_shape)
89
+ num_tokens = _get_num_tokens(output_shape)
90
+ return torch.narrow(output, 0, 0, num_tokens).view(*output_shape)
91
92
93
class RowWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
@@ -145,7 +154,8 @@ class RowWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
154
bias=bias,
155
)
156
148
- return torch.narrow(output, 0, 0, output_shape[0]).view(*output_shape)
157
+ num_tokens = _get_num_tokens(output_shape)
158
+ return torch.narrow(output, 0, 0, num_tokens).view(*output_shape)
159
160
161
class ChannelWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
@@ -206,8 +216,9 @@ class ChannelWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
216
output = output[0]
217
218
# Unpad (undo num_token_padding)
209
- output = torch.narrow(output, 0, 0, output_shape[0])
210
- x_scale = torch.narrow(As, 0, 0, output_shape[0])
219
+ num_tokens = _get_num_tokens(output_shape)
220
+ output = torch.narrow(output, 0, 0, num_tokens)
221
+ x_scale = torch.narrow(As, 0, 0, num_tokens)
222
223
# DQ
224
# C = sw * sx * (X * W) + bias