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, 16:00 UTC
Feature/silu block quant fusion v1 (#32996)
Signed-off-by: Monishver Chandrasekaran <monishverchandrasekaran@gmail.com>
Monishver committed
Apr 1, 2026 at 11:50 UTC
c09ad767cda74fc7fb587afaf3a92d714e04e1d5
11 files changed
+830
-9
CMakeLists.txt
+2
-1
@@ -340,7 +340,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
340
341
list(APPEND VLLM_EXT_SRC
342
"csrc/quantization/awq/gemm_kernels.cu"
343
- "csrc/cutlass_extensions/common.cpp")
343
+ "csrc/cutlass_extensions/common.cpp"
344
+ "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu")
345
346
set_gencode_flags_for_srcs(
347
SRCS "${VLLM_EXT_SRC}"
benchmarks/fused_kernels/silu_mul_block_quant_benchmark.py
new
+211
@@ -0,0 +1,211 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from collections.abc import Callable, Iterable
5
+from dataclasses import dataclass
6
+from itertools import product
7
+
8
+import torch
9
+import torch.nn.functional as F
10
+import torch.utils.benchmark as TBenchmark
11
+from torch.utils.benchmark import Measurement as TMeasurement
12
+from tqdm import tqdm
13
+
14
+import vllm._custom_ops as ops
15
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
16
+ per_token_group_quant_fp8,
17
+)
18
+
19
+
20
+@dataclass
21
+class bench_params_t:
22
+ num_tokens: int
23
+ hidden_size: int
24
+ dtype: torch.dtype
25
+ group_size: int # Changed from list[int] to int
26
+
27
+ def description(self):
28
+ return (
29
+ f"N {self.num_tokens} "
30
+ f"x D {self.hidden_size} "
31
+ f"x DT {self.dtype} "
32
+ f"x GS {self.group_size}"
33
+ )
34
+
35
+
36
+def get_bench_params() -> list[bench_params_t]:
37
+ """Test configurations covering common model sizes."""
38
+ NUM_TOKENS = [16, 128, 512, 2048]
39
+ HIDDEN_SIZES = [1024, 2048, 4096, 5120, 14336] # Common FFN sizes
40
+ DTYPES = [torch.float16, torch.bfloat16]
41
+ GROUP_SIZES = [64, 128] # Changed from [[1, 64], [1, 128]]
42
+
43
+ combinations = product(NUM_TOKENS, HIDDEN_SIZES, DTYPES, GROUP_SIZES)
44
+ bench_params = list(
45
+ map(lambda x: bench_params_t(x[0], x[1], x[2], x[3]), combinations)
46
+ )
47
+ return bench_params
48
+
49
+
50
+# Reference implementations
51
+def unfused_fp8_impl(
52
+ x: torch.Tensor,
53
+ quant_dtype: torch.dtype,
54
+ group_size: int, # Changed from list[int]
55
+):
56
+ """Unfused: SiLU+Mul then per-tensor quantize."""
57
+ hidden = x.shape[-1] // 2
58
+ gate, up = x.split(hidden, dim=-1)
59
+
60
+ # SiLU(gate) * up
61
+ silu_out = F.silu(gate) * up
62
+
63
+ # Per-tensor quantize (no group_size used here)
64
+ silu_out, _ = ops.scaled_fp8_quant(silu_out)
65
+
66
+
67
+def unfused_groupwise_fp8_impl(
68
+ x: torch.Tensor,
69
+ quant_dtype: torch.dtype,
70
+ group_size: int, # Changed from list[int]
71
+):
72
+ """Unfused: SiLU+Mul then group-wise quantize."""
73
+ hidden = x.shape[-1] // 2
74
+ gate, up = x.split(hidden, dim=-1)
75
+
76
+ # SiLU(gate) * up
77
+ silu_out = F.silu(gate) * up
78
+
79
+ # Group quantize - use group_size directly
80
+ silu_out, _ = per_token_group_quant_fp8(
81
+ silu_out, group_size=group_size, use_ue8m0=False
82
+ )
83
+
84
+
85
+def fused_impl(
86
+ x: torch.Tensor,
87
+ quant_dtype: torch.dtype,
88
+ group_size: int,
89
+):
90
+ """Fused: SiLU+Mul+Block Quantization in single kernel."""
91
+ out, _ = ops.silu_and_mul_per_block_quant(
92
+ x,
93
+ group_size=group_size,
94
+ quant_dtype=quant_dtype,
95
+ is_scale_transposed=False,
96
+ )
97
+
98
+
99
+# Bench functions
100
+def bench_fn(
101
+ x: torch.Tensor,
102
+ quant_dtype: torch.dtype,
103
+ group_size: int,
104
+ label: str,
105
+ sub_label: str,
106
+ fn: Callable,
107
+ description: str,
108
+) -> TMeasurement:
109
+ min_run_time = 1
110
+
111
+ globals = {
112
+ "x": x,
113
+ "quant_dtype": quant_dtype,
114
+ "group_size": group_size,
115
+ "fn": fn,
116
+ }
117
+ return TBenchmark.Timer(
118
+ stmt="fn(x, quant_dtype, group_size)",
119
+ globals=globals,
120
+ label=label,
121
+ sub_label=sub_label,
122
+ description=description,
123
+ ).blocked_autorange(min_run_time=min_run_time)
124
+
125
+
126
+def bench(params: bench_params_t, label: str, sub_label: str) -> Iterable[TMeasurement]:
127
+ """Run benchmarks for all implementations."""
128
+ # Make inputs: [num_tokens, hidden_size * 2] for [gate || up]
129
+ scale = 1 / params.hidden_size
130
+ x = (
131
+ torch.randn(
132
+ params.num_tokens,
133
+ params.hidden_size * 2,
134
+ dtype=params.dtype,
135
+ device="cuda",
136
+ )
137
+ * scale
138
+ )
139
+
140
+ timers = []
141
+
142
+ # Unfused per-tensor FP8
143
+ timers.append(
144
+ bench_fn(
145
+ x,
146
+ torch.float8_e4m3fn,
147
+ params.group_size,
148
+ label,
149
+ sub_label,
150
+ unfused_fp8_impl,
151
+ "unfused_fp8_impl",
152
+ )
153
+ )
154
+
155
+ # Unfused group-wise FP8
156
+ timers.append(
157
+ bench_fn(
158
+ x,
159
+ torch.float8_e4m3fn,
160
+ params.group_size,
161
+ label,
162
+ sub_label,
163
+ unfused_groupwise_fp8_impl,
164
+ "unfused_groupwise_fp8_impl",
165
+ )
166
+ )
167
+
168
+ # Fused group-wise FP8
169
+ timers.append(
170
+ bench_fn(
171
+ x,
172
+ torch.float8_e4m3fn,
173
+ params.group_size,
174
+ label,
175
+ sub_label,
176
+ fused_impl,
177
+ "fused_groupwise_fp8_impl",
178
+ )
179
+ )
180
+
181
+ return timers
182
+
183
+
184
+def print_timers(timers: Iterable[TMeasurement]):
185
+ compare = TBenchmark.Compare(timers)
186
+ compare.print()
187
+
188
+
189
+def main():
190
+ torch.set_default_device("cuda")
191
+ bench_params = get_bench_params()
192
+
193
+ print(f"Running {len(bench_params)} benchmark configurations...")
194
+ print(
195
+ f"This will take approximately {len(bench_params) * 3} seconds (1s per variant)"
196
+ )
197
+ print()
198
+
199
+ timers = []
200
+ for bp in tqdm(bench_params):
201
+ result_timers = bench(bp, "silu-mul-block-quant", bp.description())
202
+ timers.extend(result_timers)
203
+
204
+ print("\n" + "=" * 80)
205
+ print("FINAL COMPARISON - ALL RESULTS")
206
+ print("=" * 80)
207
+ print_timers(timers)
208
+
209
+
210
+if __name__ == "__main__":
211
+ main()
csrc/ops.h
+6
@@ -142,6 +142,12 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
142
std::optional<torch::Tensor> residual,
143
int64_t group_size, bool is_scale_transposed);
144
145
+void silu_and_mul_per_block_quant(torch::Tensor& out,
146
+ torch::Tensor const& input,
147
+ torch::Tensor& scales, int64_t group_size,
148
+ std::optional<torch::Tensor> scale_ub,
149
+ bool is_scale_transposed);
150
+
151
void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
152
std::optional<torch::Tensor> key, int64_t head_size,
153
torch::Tensor& cos_sin_cache, bool is_neox);
csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu
new
+169
@@ -0,0 +1,169 @@
1
+// SPDX-License-Identifier: Apache-2.0
2
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+#include <ATen/cuda/CUDAContext.h>
5
+#include <c10/cuda/CUDAGuard.h>
6
+
7
+#include "../../dispatch_utils.h"
8
+#include "quant_conversions.cuh"
9
+#include "../w8a8/fp8/common.cuh"
10
+
11
+namespace vllm {
12
+
13
+// Logic: one thread block per (token, group) pair
14
+
15
+template <typename scalar_t, typename scalar_out_t, bool is_scale_transposed,
16
+ int32_t group_size>
17
+__global__ void silu_and_mul_per_block_quant_kernel(
18
+ scalar_out_t* __restrict__ out, // Output: [num_tokens, hidden_size] in
19
+ // FP8/INT8
20
+ float* __restrict__ scales, // Output: [num_tokens, hidden_size /
21
+ // group_size] or [hidden_size / group_size,
22
+ // num_tokens]
23
+ scalar_t const* __restrict__ input, // Input: [num_tokens, hidden_size * 2]
24
+ float const* scale_ub, // Optional scale upper bound
25
+ int32_t const hidden_size // Output hidden size (input is 2x this)
26
+) {
27
+ static_assert((group_size & (group_size - 1)) == 0,
28
+ "group_size must be a power of 2 for correct reduction");
29
+
30
+ // Grid: (num_tokens, num_groups)
31
+ int const token_idx = blockIdx.x;
32
+ int const group_idx = blockIdx.y;
33
+ int const tid = threadIdx.x; // tid in [0, group_size)
34
+ int const num_tokens = gridDim.x;
35
+
36
+ // Input layout: [gate || up] concatenated along last dimension
37
+ int const input_stride = hidden_size * 2;
38
+ int const group_start = group_idx * group_size;
39
+
40
+ // Pointers to this token's data
41
+ scalar_t const* token_input_gate =
42
+ input + token_idx * input_stride + group_start;
43
+ scalar_t const* token_input_up = token_input_gate + hidden_size;
44
+ scalar_out_t* token_output = out + token_idx * hidden_size + group_start;
45
+
46
+ // Scale pointer for this group
47
+ int const num_groups = gridDim.y;
48
+ float* group_scale_ptr = is_scale_transposed
49
+ ? scales + group_idx * num_tokens + token_idx
50
+ : scales + token_idx * num_groups + group_idx;
51
+
52
+ // Shared memory for reduction (compile-time sized)
53
+ __shared__ float shared_max[group_size];
54
+
55
+ // Step 1: Each thread loads one element, computes SiLU, stores in register
56
+ float gate = static_cast<float>(token_input_gate[tid]);
57
+ float up = static_cast<float>(token_input_up[tid]);
58
+
59
+ // Compute SiLU(gate) * up
60
+ float sigmoid_gate = 1.0f / (1.0f + expf(-gate));
61
+ float silu_gate = gate * sigmoid_gate;
62
+ float result = silu_gate * up; // Keep in register
63
+
64
+ // Step 2: Reduce to find group max
65
+ shared_max[tid] = fabsf(result);
66
+ __syncthreads();
67
+
68
+// Power-of-2 reduction (group_size guaranteed to be power of 2)
69
+#pragma unroll
70
+ for (int stride = group_size / 2; stride > 0; stride >>= 1) {
71
+ if (tid < stride) {
72
+ shared_max[tid] = fmaxf(shared_max[tid], shared_max[tid + stride]);
73
+ }
74
+ __syncthreads();
75
+ }
76
+
77
+ // Step 3: Compute scale (thread 0), broadcast via shared memory
78
+ if (tid == 0) {
79
+ float group_max = shared_max[0];
80
+
81
+ float const quant_range = quant_type_max_v<scalar_out_t>;
82
+ float group_scale = group_max / quant_range;
83
+
84
+ // Apply scale upper bound if provided
85
+ if (scale_ub != nullptr) {
86
+ group_scale = fminf(group_scale, *scale_ub);
87
+ }
88
+
89
+ // Use minimum safe scaling factor
90
+ group_scale = fmaxf(group_scale, min_scaling_factor<scalar_out_t>::val());
91
+
92
+ // Store scale to global memory
93
+ *group_scale_ptr = group_scale;
94
+
95
+ // Reuse shared_max[0] to broadcast scale
96
+ shared_max[0] = group_scale;
97
+ }
98
+ __syncthreads();
99
+
100
+ float group_scale = shared_max[0];
101
+
102
+ // Step 4: Quantize and write output
103
+ token_output[tid] =
104
+ vllm::ScaledQuant<scalar_out_t, false>::quant_fn(result, group_scale);
105
+}
106
+
107
+} // namespace vllm
108
+
109
+void silu_and_mul_per_block_quant(torch::Tensor& out,
110
+ torch::Tensor const& input,
111
+ torch::Tensor& scales, int64_t group_size,
112
+ std::optional<torch::Tensor> scale_ub,
113
+ bool is_scale_transposed) {
114
+ static c10::ScalarType kFp8Type = is_fp8_ocp()
115
+ ? c10::ScalarType::Float8_e4m3fn
116
+ : c10::ScalarType::Float8_e4m3fnuz;
117
+
118
+ TORCH_CHECK(out.dtype() == kFp8Type || out.dtype() == torch::kInt8);
119
+ TORCH_CHECK(out.is_contiguous() && input.is_contiguous());
120
+ TORCH_CHECK(
121
+ input.dtype() == torch::kFloat16 || input.dtype() == torch::kBFloat16,
122
+ "Input must be FP16 or BF16");
123
+ TORCH_CHECK(scales.dtype() == torch::kFloat32, "Scales must be FP32");
124
+ TORCH_CHECK(group_size == 128 || group_size == 64,
125
+ "Unsupported group size: ", group_size);
126
+
127
+ if (scale_ub.has_value()) {
128
+ TORCH_CHECK(out.dtype() == kFp8Type);
129
+ }
130
+
131
+ int32_t hidden_size = out.size(-1);
132
+ auto num_tokens = input.size(0);
133
+ int32_t num_groups = hidden_size / group_size;
134
+
135
+ TORCH_CHECK(input.size(-1) == hidden_size * 2,
136
+ "input last dim must be 2x output hidden_size");
137
+ TORCH_CHECK(hidden_size % group_size == 0,
138
+ "hidden_size must be divisible by group_size");
139
+
140
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
141
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
142
+
143
+ dim3 grid(num_tokens, num_groups);
144
+ dim3 block(group_size);
145
+
146
+ VLLM_DISPATCH_FLOATING_TYPES(
147
+ input.scalar_type(), "silu_and_mul_per_block_quant", [&] {
148
+ using scalar_in_t = scalar_t;
149
+
150
+ VLLM_DISPATCH_QUANT_TYPES(
151
+ out.scalar_type(), "silu_and_mul_per_block_quant", [&] {
152
+ using scalar_out_t = scalar_t;
153
+
154
+ VLLM_DISPATCH_GROUP_SIZE(group_size, gs, [&] {
155
+ VLLM_DISPATCH_BOOL(is_scale_transposed, transpose_scale, [&] {
156
+ vllm::silu_and_mul_per_block_quant_kernel<
157
+ scalar_in_t, scalar_out_t, transpose_scale, gs>
158
+ <<<grid, block, 0, stream>>>(
159
+ out.data_ptr<scalar_out_t>(),
160
+ scales.data_ptr<float>(),
161
+ input.data_ptr<scalar_in_t>(),
162
+ scale_ub.has_value() ? scale_ub->data_ptr<float>()
163
+ : nullptr,
164
+ hidden_size);
165
+ });
166
+ });
167
+ });
168
+ });
169
+}
\ No newline at end of file
csrc/torch_bindings.cpp
+12
-1
@@ -2,7 +2,6 @@
2
#include "cuda_utils.h"
3
#include "ops.h"
4
#include "core/registration.h"
5
-
5
#include <torch/library.h>
6
#include <torch/version.h>
7
@@ -110,6 +109,18 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
109
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
110
ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
111
112
+ // Fused SiLU+Mul + per-block quantization
113
+ ops.def(
114
+ "silu_and_mul_per_block_quant("
115
+ "Tensor! out, "
116
+ "Tensor input, "
117
+ "Tensor! scales, "
118
+ "int group_size, "
119
+ "Tensor? scale_ub=None, "
120
+ "bool is_scale_transposed=False) -> ()");
121
+ ops.impl("silu_and_mul_per_block_quant", torch::kCUDA,
122
+ &silu_and_mul_per_block_quant);
123
+
124
ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");
125
ops.impl("mul_and_silu", torch::kCUDA, &mul_and_silu);
126
docs/design/fusions.md
+3
-1
@@ -45,7 +45,7 @@ The table below lists the quantization schemes supported by each fusion on each
45
| `enable_sp` | FP16/BF16, FP8 static† | FP16/BF16, FP8 static | FP16/BF16† | FP16/BF16† | — |
46
| `fuse_gemm_comms` | FP16/BF16, FP8 static† | FP16/BF16, FP8 static | FP16/BF16† | FP16/BF16† | — |
47
| `fuse_norm_quant` | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | — | FP8 static, FP8 per-token, FP8 per-group |
48
-| `fuse_act_quant` | FP8 static, NVFP4 | FP8 static | FP8 static | — | FP8 per-group |
48
+| `fuse_act_quant` | FP8 static, NVFP4 | FP8 static, FP8 per-group (128/64) | FP8 static, FP8 per-group (128/64) | — | FP8 per-group |
49
| `fuse_act_padding` | — | — | — | — | FP16/BF16 |
50
51
\* `fuse_attn_quant` support depends on the attention backend in use; not all backends support
@@ -305,6 +305,7 @@ Note that AITER fusions are in a separate pass in `vllm.compilation.passes.fusio
305
Supported quantization scheme/hardware combinations:
306
307
- FP8 static per-tensor: CUDA & HIP kernel
308
+- FP8 dynamic per-group (128/64): CUDA kernel (sm89+, not active when DeepGemm is used on sm100+)
309
- NVFP4 dynamic: CUDA sm100+ only with FlashInfer
310
- FP8 per-token-group (128): ROCm AITER only
311
@@ -313,6 +314,7 @@ Supported quantization scheme/hardware combinations:
314
- Pass: [`vllm/compilation/passes/fusion/act_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/act_quant_fusion.py)
315
- ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py)
316
- CUDA/HIP kernels: [`csrc/quantization/`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/)
317
+- Fused SiLU+Mul+BlockQuant kernel: [`csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu)
318
319
### RMSNorm + Padding (`fuse_act_padding`)
320
tests/compile/fusions_e2e/models.py
+2
-3
@@ -150,9 +150,8 @@ deepseek_v3_fp8 = ModelFusionInfo(
150
# - post_attn_layernorm + MLP
151
# 2 per MoE layer (remaining) due to MoE wrapping
152
rms_quant_fusion=n_layers * 2 + min(3, n_layers), # add for 3 dense layers
153
- # TODO silu+block quant
154
- # act_quant_fusion=min(3, n_layers), # dense layers only
155
- act_quant_fusion=0,
153
+ # silu+block quant
154
+ act_quant_fusion=min(3, n_layers), # dense layers only
155
# MLA attn + quant not supported yet:
156
# https://github.com/vllm-project/vllm/issues/35792
157
attn_quant_fusion=0,
tests/compile/passes/test_silu_mul_quant_fusion.py
+71
-3
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import itertools
4
+from functools import partial
5
6
import pytest
7
import torch
@@ -34,13 +35,16 @@ from vllm.model_executor.kernels.linear import (
35
ROCmFP8ScaledMMLinearKernel,
36
)
37
from vllm.model_executor.layers.activation import SiluAndMul
38
+from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
39
from vllm.model_executor.layers.quantization.utils.fp8_utils import W8A8BlockFp8LinearOp
40
from vllm.model_executor.layers.quantization.utils.quant_utils import (
41
GroupShape,
42
+ kFp8Dynamic128Sym,
43
kFp8StaticTensorSym,
44
kNvfp4Dynamic,
45
)
46
from vllm.platforms import current_platform
47
+from vllm.utils.deep_gemm import is_deep_gemm_supported
48
49
FP8_DTYPE = current_platform.fp8_dtype()
50
FP4_DTYPE = torch.uint8
@@ -165,6 +169,48 @@ class TestSiluMulGroupFp8QuantModel(torch.nn.Module):
169
return [torch.ops.vllm.rocm_aiter_act_mul_and_fp8_group_quant]
170
171
172
+class TestSiluMulBlockQuantModel(torch.nn.Module):
173
+ quant_key = kFp8Dynamic128Sym
174
+
175
+ def __init__(self, hidden_size: int, is_scale_transposed: bool = False, **kwargs):
176
+ super().__init__()
177
+ self.silu_and_mul = SiluAndMul()
178
+ self.is_scale_transposed = is_scale_transposed
179
+ self.quant_fp8 = QuantFP8(
180
+ static=False,
181
+ group_shape=GroupShape(1, 128),
182
+ column_major_scales=is_scale_transposed,
183
+ compile_native=False,
184
+ )
185
+
186
+ self.enable_silu_mul_custom_op = self.silu_and_mul.enabled()
187
+ self.enable_quant_fp8_custom_op = self.quant_fp8.enabled()
188
+
189
+ def forward(self, x):
190
+ y = self.silu_and_mul(x)
191
+ out, scale = self.quant_fp8(y)
192
+ group_size = self.quant_key.scale.group_shape[1]
193
+ scale_expanded = scale.repeat_interleave(group_size, dim=1)
194
+ dequant = out.to(dtype=torch.float32) * scale_expanded
195
+ return (dequant,)
196
+
197
+ def ops_in_model_before(self):
198
+ ops = []
199
+ if self.enable_silu_mul_custom_op:
200
+ ops.append(SILU_MUL_OP)
201
+ # When silu custom op is disabled, aten.mul.Tensor also appears
202
+ # in dequant code, so we skip checking it to avoid false positives.
203
+ ops.append(
204
+ QUANT_OPS[self.quant_key]
205
+ if self.enable_quant_fp8_custom_op
206
+ else torch.ops.aten.reciprocal.default
207
+ )
208
+ return ops
209
+
210
+ def ops_in_model_after(self):
211
+ return [FUSED_OPS[self.quant_key]]
212
+
213
+
214
ROCM_KERNELS = [ROCmFP8ScaledMMLinearKernel, PerTensorTorchFP8ScaledMMLinearKernel]
215
CUDA_KERNELS = [
216
FlashInferFP8ScaledMMLinearKernel,
@@ -200,6 +246,19 @@ TEST_KERNELS = ROCM_KERNELS if current_platform.is_rocm() else CUDA_KERNELS
246
not current_platform.is_rocm(), reason="ROCm only"
247
),
248
),
249
+ # Block quant fusion for per-group FP8 (CUDA only).
250
+ *[
251
+ pytest.param(
252
+ partial(TestSiluMulBlockQuantModel, is_scale_transposed=transposed),
253
+ True,
254
+ None,
255
+ marks=pytest.mark.skipif(
256
+ not current_platform.is_cuda(), reason="CUDA only"
257
+ ),
258
+ id=f"TestSiluMulBlockQuant-transposed={transposed}",
259
+ )
260
+ for transposed in [False, True]
261
+ ],
262
],
263
)
264
@pytest.mark.skipif(
@@ -213,6 +272,7 @@ def test_fusion_silu_and_mul_quant(
272
TestSiluMulFp8QuantModel
273
| TestSiluMulNvfp4QuantModel
274
| TestSiluMulGroupFp8QuantModel
275
+ | TestSiluMulBlockQuantModel
276
],
277
enable_silu_mul_custom_op: bool,
278
enable_quant_fp8_custom_op: bool,
@@ -223,6 +283,12 @@ def test_fusion_silu_and_mul_quant(
283
pytest.skip("NVFP4 is not supported on this GPU.")
284
if model_class is TestSiluMulGroupFp8QuantModel and not IS_AITER_FOUND:
285
pytest.skip("AITER is not supported on this GPU.")
286
+ if (
287
+ isinstance(model_class, partial)
288
+ and model_class.func is TestSiluMulBlockQuantModel
289
+ and is_deep_gemm_supported()
290
+ ):
291
+ pytest.skip("SiluMul+BlockQuant fusion not applicable with DeepGemm")
292
293
torch.set_default_device("cuda")
294
torch.set_default_dtype(dtype)
@@ -269,11 +335,13 @@ def test_fusion_silu_and_mul_quant(
335
result2 = model2(x)
336
337
# Check that it gives the same answer
272
- if model_class == TestSiluMulFp8QuantModel:
338
+ if isinstance(model, TestSiluMulFp8QuantModel):
339
atol, rtol = 1e-3, 1e-3
274
- elif model_class == TestSiluMulNvfp4QuantModel:
340
+ elif isinstance(model, TestSiluMulNvfp4QuantModel):
341
atol, rtol = 1e-1, 1e-1
276
- elif model_class == TestSiluMulGroupFp8QuantModel:
342
+ elif isinstance(
343
+ model, (TestSiluMulGroupFp8QuantModel, TestSiluMulBlockQuantModel)
344
+ ):
345
atol, rtol = 5e-2, 5e-2
346
347
torch.testing.assert_close(
tests/kernels/core/test_fused_silu_mul_block_quant.py
new
+189
@@ -0,0 +1,189 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+import torch.nn.functional as F
7
+
8
+import vllm._custom_ops as ops
9
+from tests.kernels.utils import opcheck
10
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
11
+ per_token_group_quant_fp8,
12
+)
13
+from vllm.model_executor.layers.quantization.utils.int8_utils import (
14
+ per_token_group_quant_int8,
15
+)
16
+from vllm.platforms import current_platform
17
+
18
+DTYPES = [torch.float16, torch.bfloat16]
19
+QUANT_DTYPES = [torch.float8_e4m3fn, 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]],
23
+ *[(16, i) for i in [64, *VEC_HIDDEN_SIZES, 5120]],
24
+ *[(128, i) for i in [64, *VEC_HIDDEN_SIZES]],
25
+ *[(512, i) for i in [64, 5120]],
26
+]
27
+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
+]
34
+
35
+
36
+def ref_silu_and_mul_per_block_quant(
37
+ x: torch.Tensor,
38
+ quant_dtype: torch.dtype,
39
+ group_size: int,
40
+) -> tuple[torch.Tensor, torch.Tensor]:
41
+ """Reference implementation: unfused SiLU+Mul then group quantization."""
42
+ hidden = x.shape[-1] // 2
43
+ gate, up = x.split(hidden, dim=-1)
44
+ silu_out = F.silu(gate) * up
45
+
46
+ if quant_dtype == current_platform.fp8_dtype():
47
+ return per_token_group_quant_fp8(
48
+ silu_out, group_size=group_size, use_ue8m0=False
49
+ )
50
+ elif quant_dtype == torch.int8:
51
+ return per_token_group_quant_int8(silu_out, group_size=group_size)
52
+ else:
53
+ raise ValueError(f"Unsupported quant_dtype: {quant_dtype}")
54
+
55
+
56
+@pytest.mark.parametrize("num_tokens, hidden_size", NUM_TOKENS_HIDDEN_SIZES)
57
+@pytest.mark.parametrize("has_scale_ub", SCALE_UBS)
58
+@pytest.mark.parametrize("dtype", DTYPES)
59
+@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
60
+@pytest.mark.parametrize("group_size", GROUP_SIZES)
61
+@pytest.mark.parametrize("is_scale_transposed", IS_SCALE_TRANSPOSED)
62
+@pytest.mark.parametrize("seed", SEEDS)
63
+@pytest.mark.parametrize("device", CUDA_DEVICES)
64
+@torch.inference_mode()
65
+def test_silu_and_mul_per_block_quant(
66
+ default_vllm_config,
67
+ num_tokens: int,
68
+ hidden_size: int,
69
+ has_scale_ub: bool,
70
+ dtype: torch.dtype,
71
+ quant_dtype: torch.dtype,
72
+ group_size: int,
73
+ is_scale_transposed: bool,
74
+ seed: int,
75
+ device: str,
76
+) -> None:
77
+ """Test SiLU+Mul+Block Quantization kernel correctness."""
78
+ torch.random.manual_seed(seed)
79
+ torch.set_default_device(device)
80
+
81
+ if hidden_size % group_size != 0:
82
+ return
83
+
84
+ if has_scale_ub:
85
+ pytest.skip("Scale upper bound not yet supported")
86
+
87
+ scale = 1 / hidden_size
88
+ x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device) * scale
89
+
90
+ # Reference implementation
91
+ ref_out, ref_scales = ref_silu_and_mul_per_block_quant(x, quant_dtype, group_size)
92
+
93
+ # Fused kernel implementation
94
+ ops_out, ops_scales = ops.silu_and_mul_per_block_quant(
95
+ x, group_size, quant_dtype, None, is_scale_transposed
96
+ )
97
+
98
+ # Check for NaN/Inf
99
+ assert not torch.isnan(ops_out.float()).any(), "Kernel output contains NaN"
100
+ assert not torch.isinf(ops_out.float()).any(), "Kernel output contains Inf"
101
+ assert not torch.isnan(ops_scales).any(), "Kernel scales contain NaN"
102
+ assert not torch.isinf(ops_scales).any(), "Kernel scales contain Inf"
103
+
104
+ # Check dtypes
105
+ assert ref_out.dtype == quant_dtype
106
+ assert ops_out.dtype == quant_dtype
107
+
108
+ # Check scales match
109
+ torch.testing.assert_close(ref_scales, ops_scales, rtol=1e-5, atol=1e-5)
110
+
111
+ # Check output correctness via dequantized values
112
+ ref_scales_expanded = ref_scales.repeat_interleave(group_size, dim=1)
113
+ ops_scales_expanded = ops_scales.repeat_interleave(group_size, dim=1)
114
+ ref_deq = ref_out.to(dtype=torch.float32) * ref_scales_expanded
115
+ ops_deq = ops_out.to(dtype=torch.float32) * ops_scales_expanded
116
+ torch.testing.assert_close(ref_deq, ops_deq, atol=5e-2, rtol=5e-2)
117
+
118
+ # opcheck
119
+ output = torch.empty(num_tokens, hidden_size, device=device, dtype=quant_dtype)
120
+ num_groups = hidden_size // group_size
121
+ if is_scale_transposed:
122
+ scales = torch.empty(num_groups, num_tokens, device=device, dtype=torch.float32)
123
+ else:
124
+ scales = torch.empty(num_tokens, num_groups, device=device, dtype=torch.float32)
125
+ opcheck(
126
+ torch.ops._C.silu_and_mul_per_block_quant,
127
+ (output, x, scales, group_size, None, is_scale_transposed),
128
+ )
129
+
130
+
131
+@pytest.mark.parametrize("dtype", [torch.float16])
132
+@pytest.mark.parametrize("hidden_size", [4096])
133
+@pytest.mark.parametrize("num_tokens", [128])
134
+@pytest.mark.parametrize("group_size", [128])
135
+def test_silu_block_quant_shapes(
136
+ default_vllm_config,
137
+ dtype: torch.dtype,
138
+ hidden_size: int,
139
+ num_tokens: int,
140
+ group_size: int,
141
+):
142
+ """Test that output shapes are correct."""
143
+ torch.set_default_device("cuda")
144
+ x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device="cuda")
145
+
146
+ # Row-major scales
147
+ out, scales = ops.silu_and_mul_per_block_quant(
148
+ x,
149
+ group_size=group_size,
150
+ quant_dtype=torch.float8_e4m3fn,
151
+ is_scale_transposed=False,
152
+ )
153
+ assert out.shape == (num_tokens, hidden_size)
154
+ assert scales.shape == (num_tokens, hidden_size // group_size)
155
+
156
+ # Column-major scales (logical shape same after .t() in _custom_ops)
157
+ out, scales = ops.silu_and_mul_per_block_quant(
158
+ x,
159
+ group_size=group_size,
160
+ quant_dtype=torch.float8_e4m3fn,
161
+ is_scale_transposed=True,
162
+ )
163
+ assert out.shape == (num_tokens, hidden_size)
164
+ assert scales.shape == (num_tokens, hidden_size // group_size)
165
+
166
+
167
+@pytest.mark.parametrize("dtype", [torch.float16])
168
+@pytest.mark.parametrize("batch_size", [1, 16, 256])
169
+@pytest.mark.parametrize("hidden_size", [1024, 5120, 14336])
170
+def test_silu_block_quant_edge_cases(
171
+ default_vllm_config, dtype: torch.dtype, batch_size: int, hidden_size: int
172
+):
173
+ """Test edge cases: single token, large batch, large hidden size."""
174
+ torch.set_default_device("cuda")
175
+ x = torch.randn(batch_size, hidden_size * 2, dtype=dtype, device="cuda")
176
+
177
+ out, scales = ops.silu_and_mul_per_block_quant(
178
+ x,
179
+ group_size=128,
180
+ quant_dtype=torch.float8_e4m3fn,
181
+ is_scale_transposed=False,
182
+ )
183
+
184
+ assert out.shape == (batch_size, hidden_size)
185
+ assert out.dtype == torch.float8_e4m3fn
186
+ assert scales.dtype == torch.float32
187
+ assert not torch.isnan(out.float()).any()
188
+ assert not torch.isnan(scales).any()
189
+ assert not torch.isinf(scales).any()
vllm/_custom_ops.py
+50
@@ -579,6 +579,56 @@ def rms_norm_per_block_quant(
579
return output, scales
580
581
582
+# fused silu_and_mul + block quant
583
+def silu_and_mul_per_block_quant(
584
+ input: torch.Tensor,
585
+ group_size: int, # Changed from list[int]
586
+ quant_dtype: torch.dtype,
587
+ scale_ub: torch.Tensor | None = None,
588
+ is_scale_transposed: bool = False,
589
+) -> tuple[torch.Tensor, torch.Tensor]:
590
+ assert input.ndim == 2, f"input must be 2D [batch, hidden*2], got {input.shape}"
591
+ assert input.shape[-1] % 2 == 0, (
592
+ f"input last dim must be even (gate||up layout), got {input.shape[-1]}"
593
+ )
594
+
595
+ # Output is half the width of input (after silu_and_mul)
596
+ num_tokens = input.shape[0]
597
+ hidden_size = input.shape[-1] // 2 # Divide by 2 because input is [gate || up]
598
+
599
+ # Allocate output tensor (FP8 or INT8)
600
+ output = torch.empty(
601
+ (num_tokens, hidden_size), device=input.device, dtype=quant_dtype
602
+ )
603
+
604
+ # Allocate scales tensor
605
+ num_groups = hidden_size // group_size # Directly use group_size
606
+ if is_scale_transposed:
607
+ scales = torch.empty(
608
+ (num_groups, num_tokens),
609
+ device=input.device,
610
+ dtype=torch.float32,
611
+ ).t()
612
+ else:
613
+ scales = torch.empty(
614
+ (num_tokens, num_groups),
615
+ device=input.device,
616
+ dtype=torch.float32,
617
+ )
618
+
619
+ # Call the C++ kernel
620
+ torch.ops._C.silu_and_mul_per_block_quant(
621
+ output,
622
+ input,
623
+ scales,
624
+ group_size, # Pass directly as int
625
+ scale_ub,
626
+ is_scale_transposed,
627
+ )
628
+
629
+ return output, scales
630
+
631
+
632
# quantization ops
633
# awq
634
def awq_dequantize(
vllm/compilation/passes/fusion/act_quant_fusion.py
+115
@@ -17,6 +17,8 @@ from vllm.config import VllmConfig
17
from vllm.logger import init_logger
18
from vllm.model_executor.layers.quantization.utils.quant_utils import (
19
QuantKey,
20
+ kFp8Dynamic64Sym,
21
+ kFp8Dynamic128Sym,
22
kFp8StaticTensorSym,
23
kNvfp4Dynamic,
24
)
@@ -43,6 +45,10 @@ 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():
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
+
52
53
class ActivationQuantPattern(ABC):
54
"""
@@ -174,6 +180,102 @@ class SiluMulNvfp4QuantPattern(ActivationQuantPattern):
180
register_replacement(pattern, replacement, self.get_inputs(), fwd_only, pm_pass)
181
182
183
+class SiluMulBlockQuantPattern(ActivationQuantPattern):
184
+ """
185
+ Fusion for SiluMul+BlockQuant (FP8 dynamic per-group) Pattern.
186
+ Supports group_size 128 and 64 via QuantKey.
187
+ Parameterized on is_scale_transposed for different scale layouts.
188
+ """
189
+
190
+ def __init__(
191
+ self,
192
+ quant_key: QuantKey,
193
+ is_scale_transposed: bool = False,
194
+ is_e8m0: bool = False,
195
+ is_tma_aligned: bool = False,
196
+ ) -> None:
197
+ super().__init__(quant_key)
198
+ self.quant_matcher = MatcherQuantFP8(
199
+ quant_key,
200
+ has_col_major_scales=is_scale_transposed,
201
+ is_e8m0=is_e8m0,
202
+ is_tma_aligned=is_tma_aligned,
203
+ )
204
+ self.group_size = quant_key.scale.group_shape[1]
205
+ self.is_scale_transposed = is_scale_transposed
206
+ self.is_e8m0 = is_e8m0
207
+ self.is_tma_aligned = is_tma_aligned
208
+
209
+ def get_inputs(self) -> list[torch.Tensor]:
210
+ scale = self.quant_matcher.empty_f32(1, 1)
211
+ return self.silu_and_mul_matcher.inputs() + [scale]
212
+
213
+ def register(self, pm_pass: PatternMatcherPass) -> None:
214
+ is_scale_transposed = self.is_scale_transposed
215
+
216
+ def pattern(
217
+ input: torch.Tensor,
218
+ scale: torch.Tensor,
219
+ ) -> tuple[torch.Tensor, torch.Tensor]:
220
+ silu_out = self.silu_and_mul_matcher(input)
221
+ result = torch.empty(
222
+ silu_out.shape,
223
+ device=silu_out.device,
224
+ dtype=self.quant_dtype,
225
+ )
226
+ assert scale is not None
227
+ finfo = torch.finfo(self.quant_dtype)
228
+ _, result, scale = auto_functionalized(
229
+ self.quant_matcher.QUANT_OP,
230
+ input=silu_out,
231
+ output_q=result,
232
+ output_s=scale,
233
+ group_size=self.group_size,
234
+ eps=1e-10,
235
+ fp8_min=finfo.min,
236
+ fp8_max=finfo.max,
237
+ scale_ue8m0=self.is_e8m0,
238
+ dummy_is_scale_transposed=is_scale_transposed,
239
+ dummy_is_tma_aligned=self.is_tma_aligned,
240
+ )
241
+ return result, scale
242
+
243
+ def replacement(
244
+ input: torch.Tensor,
245
+ scale: torch.Tensor,
246
+ ) -> tuple[torch.Tensor, torch.Tensor]:
247
+ d = input.shape[-1] // 2
248
+ output_shape = input.shape[:-1] + (d,)
249
+ result = torch.empty(
250
+ output_shape, device=input.device, dtype=self.quant_dtype
251
+ )
252
+ if is_scale_transposed:
253
+ scale = torch.empty(
254
+ (d // self.group_size, input.shape[0]),
255
+ device=input.device,
256
+ dtype=torch.float32,
257
+ ).permute(-1, -2)
258
+ else:
259
+ scale = torch.empty(
260
+ (input.shape[0], d // self.group_size),
261
+ device=input.device,
262
+ dtype=torch.float32,
263
+ )
264
+ at = auto_functionalized(
265
+ self.FUSED_OP,
266
+ out=result,
267
+ input=input,
268
+ scales=scale,
269
+ group_size=self.group_size,
270
+ scale_ub=None,
271
+ is_scale_transposed=is_scale_transposed,
272
+ )
273
+ return at[1], at[2]
274
+
275
+ inps = self.get_inputs()
276
+ register_replacement(pattern, replacement, inps, fwd_only, pm_pass)
277
+
278
+
279
class ActivationQuantFusionPass(VllmPatternMatcherPass):
280
"""
281
This pass fuses a pre-defined set of custom ops into fused ops.
@@ -199,6 +301,18 @@ class ActivationQuantFusionPass(VllmPatternMatcherPass):
301
pattern_silu_mul_nvfp4 = SiluMulNvfp4QuantPattern()
302
pattern_silu_mul_nvfp4.register(self.patterns)
303
304
+ if current_platform.is_cuda():
305
+ for quant_key in [kFp8Dynamic128Sym, kFp8Dynamic64Sym]:
306
+ for is_scale_transposed in [False, True]:
307
+ for is_e8m0 in [True, False]:
308
+ for is_tma_aligned in [False, True]:
309
+ SiluMulBlockQuantPattern(
310
+ quant_key,
311
+ is_scale_transposed=is_scale_transposed,
312
+ is_e8m0=is_e8m0,
313
+ is_tma_aligned=is_tma_aligned,
314
+ ).register(self.patterns)
315
+
316
self.dump_patterns(config, self.patterns)
317
318
@VllmInductorPass.time_and_log
@@ -212,4 +326,5 @@ class ActivationQuantFusionPass(VllmPatternMatcherPass):
326
ActivationQuantPattern,
327
SiluMulFp8StaticQuantPattern,
328
SiluMulNvfp4QuantPattern,
329
+ SiluMulBlockQuantPattern,
330
)