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 5, 11:00 UTC
[ROCm] Enable fused_silu_mul_block_quant on ROCm (#38817)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
Gregory Shtrasberg committed
Apr 8, 2026 at 11:23 UTC
56c976c1b5ccdf1297cb792911792565f00208b2
7 files changed
+28
-29
CMakeLists.txt
+2
-2
@@ -299,6 +299,7 @@ set(VLLM_EXT_SRC
299
"csrc/quantization/w8a8/int8/scaled_quant.cu"
300
"csrc/quantization/w8a8/fp8/common.cu"
301
"csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
302
+ "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
303
"csrc/quantization/gguf/gguf_kernel.cu"
304
"csrc/quantization/activation_kernels.cu"
305
"csrc/cuda_utils_kernels.cu"
@@ -340,8 +341,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
341
342
list(APPEND VLLM_EXT_SRC
343
"csrc/quantization/awq/gemm_kernels.cu"
343
- "csrc/cutlass_extensions/common.cpp"
344
- "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu")
344
+ "csrc/cutlass_extensions/common.cpp")
345
346
set_gencode_flags_for_srcs(
347
SRCS "${VLLM_EXT_SRC}"
csrc/ops.h
-2
@@ -143,13 +143,11 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
143
std::optional<torch::Tensor> residual,
144
int64_t group_size, bool is_scale_transposed);
145
146
-#ifndef USE_ROCM
146
void silu_and_mul_per_block_quant(torch::Tensor& out,
147
torch::Tensor const& input,
148
torch::Tensor& scales, int64_t group_size,
149
std::optional<torch::Tensor> scale_ub,
150
bool is_scale_transposed);
152
-#endif
151
152
void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
153
std::optional<torch::Tensor> key, int64_t head_size,
csrc/quantization/fused_kernels/quant_conversions.cuh
+1
-1
@@ -6,7 +6,7 @@
6
7
#include "libtorch_stable/quantization/vectorization.cuh"
8
// TODO(luka/varun):refactor common.cuh to use this file instead
9
-#include "quantization/w8a8/fp8/common.cuh"
9
+#include "../w8a8/fp8/common.cuh"
10
11
namespace vllm {
12
csrc/quantization/w8a8/fp8/common.cuh
+1
-1
@@ -1,7 +1,7 @@
1
#pragma once
2
3
#include "libtorch_stable/quantization/vectorization.cuh"
4
-#include "quantization/utils.cuh"
4
+#include "../../utils.cuh"
5
6
#include <cmath>
7
csrc/torch_bindings.cpp
+12
-11
@@ -110,6 +110,18 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
110
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
111
ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
112
113
+ // Fused SiLU+Mul + per-block quantization
114
+ ops.def(
115
+ "silu_and_mul_per_block_quant("
116
+ "Tensor! out, "
117
+ "Tensor input, "
118
+ "Tensor! scales, "
119
+ "int group_size, "
120
+ "Tensor? scale_ub=None, "
121
+ "bool is_scale_transposed=False) -> ()");
122
+ ops.impl("silu_and_mul_per_block_quant", torch::kCUDA,
123
+ &silu_and_mul_per_block_quant);
124
+
125
ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");
126
ops.impl("mul_and_silu", torch::kCUDA, &mul_and_silu);
127
@@ -233,17 +245,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
245
246
// Quantization ops
247
#ifndef USE_ROCM
236
- // Fused SiLU+Mul + per-block quantization
237
- ops.def(
238
- "silu_and_mul_per_block_quant("
239
- "Tensor! out, "
240
- "Tensor input, "
241
- "Tensor! scales, "
242
- "int group_size, "
243
- "Tensor? scale_ub=None, "
244
- "bool is_scale_transposed=False) -> ()");
245
- ops.impl("silu_and_mul_per_block_quant", torch::kCUDA,
246
- &silu_and_mul_per_block_quant);
248
// DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
249
ops.def(
250
"dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
tests/kernels/core/test_fused_silu_mul_block_quant.py
+10
-10
@@ -16,7 +16,7 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import (
16
from vllm.platforms import current_platform
17
18
DTYPES = [torch.float16, torch.bfloat16]
19
-QUANT_DTYPES = [torch.float8_e4m3fn, torch.int8]
19
+QUANT_DTYPES = [current_platform.fp8_dtype(), torch.int8]
20
VEC_HIDDEN_SIZES = [1024, 1025, 1027, 1029]
21
NUM_TOKENS_HIDDEN_SIZES = [
22
*[(1, i) for i in [64, *VEC_HIDDEN_SIZES, 2048, 5120]],
@@ -28,9 +28,7 @@ SCALE_UBS = [False]
28
GROUP_SIZES = [64, 128]
29
IS_SCALE_TRANSPOSED = [False, True]
30
SEEDS = [0]
31
-CUDA_DEVICES = [
32
- f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
33
-]
31
+CUDA_DEVICES = [i for i in range(1 if torch.accelerator.device_count() == 1 else 2)]
32
33
34
def ref_silu_and_mul_per_block_quant(
@@ -60,7 +58,7 @@ def ref_silu_and_mul_per_block_quant(
58
@pytest.mark.parametrize("group_size", GROUP_SIZES)
59
@pytest.mark.parametrize("is_scale_transposed", IS_SCALE_TRANSPOSED)
60
@pytest.mark.parametrize("seed", SEEDS)
63
-@pytest.mark.parametrize("device", CUDA_DEVICES)
61
+@pytest.mark.parametrize("device_idx", CUDA_DEVICES)
62
@torch.inference_mode()
63
def test_silu_and_mul_per_block_quant(
64
default_vllm_config,
@@ -72,9 +70,11 @@ def test_silu_and_mul_per_block_quant(
70
group_size: int,
71
is_scale_transposed: bool,
72
seed: int,
75
- device: str,
73
+ device_idx: str,
74
) -> None:
75
"""Test SiLU+Mul+Block Quantization kernel correctness."""
76
+ torch.accelerator.set_device_index(device_idx)
77
+ device = f"cuda:{device_idx}"
78
torch.random.manual_seed(seed)
79
torch.set_default_device(device)
80
@@ -147,7 +147,7 @@ def test_silu_block_quant_shapes(
147
out, scales = ops.silu_and_mul_per_block_quant(
148
x,
149
group_size=group_size,
150
- quant_dtype=torch.float8_e4m3fn,
150
+ quant_dtype=current_platform.fp8_dtype(),
151
is_scale_transposed=False,
152
)
153
assert out.shape == (num_tokens, hidden_size)
@@ -157,7 +157,7 @@ def test_silu_block_quant_shapes(
157
out, scales = ops.silu_and_mul_per_block_quant(
158
x,
159
group_size=group_size,
160
- quant_dtype=torch.float8_e4m3fn,
160
+ quant_dtype=current_platform.fp8_dtype(),
161
is_scale_transposed=True,
162
)
163
assert out.shape == (num_tokens, hidden_size)
@@ -177,12 +177,12 @@ def test_silu_block_quant_edge_cases(
177
out, scales = ops.silu_and_mul_per_block_quant(
178
x,
179
group_size=128,
180
- quant_dtype=torch.float8_e4m3fn,
180
+ quant_dtype=current_platform.fp8_dtype(),
181
is_scale_transposed=False,
182
)
183
184
assert out.shape == (batch_size, hidden_size)
185
- assert out.dtype == torch.float8_e4m3fn
185
+ assert out.dtype == current_platform.fp8_dtype()
186
assert scales.dtype == torch.float32
187
assert not torch.isnan(out.float()).any()
188
assert not torch.isnan(scales).any()
vllm/compilation/passes/fusion/act_quant_fusion.py
+2
-2
@@ -45,7 +45,7 @@ silu_and_mul_nvfp4_quant_supported = current_platform.is_cuda() and hasattr(
45
if silu_and_mul_nvfp4_quant_supported:
46
FUSED_OPS[kNvfp4Dynamic] = torch.ops._C.silu_and_mul_nvfp4_quant.default # noqa: E501
47
48
-if current_platform.is_cuda():
48
+if current_platform.is_cuda_alike():
49
FUSED_OPS[kFp8Dynamic128Sym] = torch.ops._C.silu_and_mul_per_block_quant.default
50
FUSED_OPS[kFp8Dynamic64Sym] = torch.ops._C.silu_and_mul_per_block_quant.default
51
@@ -301,7 +301,7 @@ class ActivationQuantFusionPass(VllmPatternMatcherPass):
301
pattern_silu_mul_nvfp4 = SiluMulNvfp4QuantPattern()
302
pattern_silu_mul_nvfp4.register(self.patterns)
303
304
- if current_platform.is_cuda():
304
+ if current_platform.is_cuda_alike():
305
for quant_key in [kFp8Dynamic128Sym, kFp8Dynamic64Sym]:
306
for is_scale_transposed in [False, True]:
307
for is_e8m0 in [True, False]: