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 15, 06:00 UTC
[AMD][Kernel][BugFix] fix test_rocm_compressed_tensors_w8a8 for rocm (#19509)
Signed-off-by: Randall Smith <Randall.Smith@amd.com>
rasmith committed
Jun 12, 2025 at 02:14 UTC
2e090bd5df974949651ad439517e0da4e981b508
2 files changed
+6
-9
vllm/_custom_ops.py
+2
-5
@@ -2,7 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import contextlib
5
-import importlib
5
from typing import TYPE_CHECKING, Optional, Union
6
7
import torch
@@ -706,10 +705,8 @@ def cutlass_scaled_mm(a: torch.Tensor,
705
706
cutlass_compatible_b = (b.shape[0] % 16 == 0 and b.shape[1] % 16 == 0)
707
if current_platform.is_rocm() or not cutlass_compatible_b:
709
- triton_scaled_mm_module = importlib.import_module(
710
- "vllm.model_executor.layers.quantization.compressed_tensors."
711
- "triton_scaled_mm")
712
- triton_scaled_mm = triton_scaled_mm_module.triton_scaled_mm
708
+ from vllm.model_executor.layers.quantization.compressed_tensors.triton_scaled_mm import ( # noqa
709
+ triton_scaled_mm)
710
return triton_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
711
712
out = torch.empty((m, n), dtype=out_dtype, device=a.device)
vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py
+4
-4
@@ -144,10 +144,10 @@ def triton_scaled_mm(input: torch.Tensor,
144
scale_b = scale_b.reshape(-1, 1) if scale_b.dim() <= 1 else scale_b
145
146
assert scale_a.dtype == scale_b.dtype and scale_a.is_floating_point()
147
- assert scale_a.shape == torch.Size([1, 1]) or scale_a.shape == torch.Size(
148
- [M, 1])
149
- assert scale_b.shape == torch.Size([1, 1]) or scale_b.shape == torch.Size(
150
- [N, 1])
147
+ assert scale_a.shape[1] == 1 and (scale_a.shape[0] == 1
148
+ or scale_a.shape[0] == M)
149
+ assert scale_b.shape[1] == 1 and (scale_b.shape[0] == 1
150
+ or scale_b.shape[0] == N)
151
assert out_dtype.is_floating_point
152
assert bias is None or bias.is_floating_point()
153
assert is_weak_contiguous(input)