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, 12:15 UTC
[Quantizaton] [AMD] Add support for running DeepSeek int8 w8a8 MoE on ROCm (#17558)
Signed-off-by: Randall Smith <Randall.Smith@amd.com>
rasmith committed
May 2, 2025 at 23:41 UTC
e3d0a1d190678b49d541fea2bd3db8d3ce9f0430
2 files changed
+29
-3
vllm/_custom_ops.py
+2
-2
@@ -559,7 +559,6 @@ def cutlass_scaled_mm(a: torch.Tensor,
559
scale_a.shape * [1, 128] == a.shape
560
scale_b.shape * [128, 128] == b.shape
561
"""
562
- assert (b.shape[0] % 16 == 0 and b.shape[1] % 16 == 0)
562
assert (out_dtype is torch.bfloat16 or out_dtype is torch.float16)
563
assert bias is None or bias.shape[0] == b.shape[
564
1] and bias.dtype == out_dtype
@@ -567,7 +566,8 @@ def cutlass_scaled_mm(a: torch.Tensor,
566
m = a.shape[0]
567
n = b.shape[1]
568
570
- if current_platform.is_rocm():
569
+ cutlass_compatible_b = (b.shape[0] % 16 == 0 and b.shape[1] % 16 == 0)
570
+ if current_platform.is_rocm() or not cutlass_compatible_b:
571
triton_scaled_mm_module = importlib.import_module(
572
"vllm.model_executor.layers.quantization.compressed_tensors."
573
"triton_scaled_mm")
vllm/model_executor/layers/quantization/utils/int8_utils.py
+27
-1
@@ -85,6 +85,32 @@ def block_dequant(
85
return x_dq_block
86
87
88
+if current_platform.is_rocm():
89
+ from triton.language import core
90
+
91
+ # NOTE: This can be removed when hip.libdevice.round() is available.
92
+ @core.extern
93
+ def round_f32(arg0, _builder=None):
94
+ return core.extern_elementwise("",
95
+ "", [arg0], {
96
+ (core.dtype("fp32"), ):
97
+ ("llvm.round", core.dtype("fp32")),
98
+ (core.dtype("fp64"), ):
99
+ ("llvm.round", core.dtype("fp64")),
100
+ },
101
+ is_pure=True,
102
+ _builder=_builder)
103
+
104
+ @triton.jit
105
+ def round_int8(x):
106
+ return round_f32(x).to(tl.int8)
107
+else:
108
+
109
+ @triton.jit
110
+ def round_int8(x):
111
+ return tl.extra.cuda.libdevice.round(x).to(tl.int8)
112
+
113
+
114
@triton.jit
115
def _per_token_quant_int8(
116
x_ptr,
@@ -106,7 +132,7 @@ def _per_token_quant_int8(
132
absmax = tl.maximum(tl.max(tl.abs(x)), 1e-10)
133
scale_x = absmax / 127
134
x_q = x * (127 / absmax)
109
- x_q = tl.extra.cuda.libdevice.round(x_q).to(tl.int8)
135
+ x_q = round_int8(x_q)
136
137
tl.store(xq_ptr + row_id * stride_xq + cols, x_q, mask=mask)
138
tl.store(scale_ptr + row_id, scale_x)