[Bugfix] Fix scaled_mm output narrowing for 3D input tensors

vllm/model_executor/kernels/linear/scaled_mm/pytorch.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import math
 
 import torch
 
@@ -13,6 +14,13 @@
 )
 
 
+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.
@@ -78,7 +86,8 @@ def apply_scaled_mm(
         if type(output) is tuple and len(output) == 2:
             output = output[0]
 
-        return torch.narrow(output, 0, 0, output_shape[0]).view(*output_shape)
+        num_tokens = _get_num_tokens(output_shape)
+        return torch.narrow(output, 0, 0, num_tokens).view(*output_shape)
 
 
 class RowWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
@@ -145,7 +154,8 @@ def apply_scaled_mm(
             bias=bias,
         )
 
-        return torch.narrow(output, 0, 0, output_shape[0]).view(*output_shape)
+        num_tokens = _get_num_tokens(output_shape)
+        return torch.narrow(output, 0, 0, num_tokens).view(*output_shape)
 
 
 class ChannelWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
@@ -206,8 +216,9 @@ def apply_scaled_mm(
             output = output[0]
 
         # Unpad (undo num_token_padding)
-        output = torch.narrow(output, 0, 0, output_shape[0])
-        x_scale = torch.narrow(As, 0, 0, output_shape[0])
+        num_tokens = _get_num_tokens(output_shape)
+        output = torch.narrow(output, 0, 0, num_tokens)
+        x_scale = torch.narrow(As, 0, 0, num_tokens)
 
         # DQ
         # C = sw * sx * (X * W) + bias