@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 16a45b3a28
Fork 0

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 14, 05:30 UTC

[NVIDIA] Support SiluMul + NVFP4 quant fusion (#23671)

Signed-off-by: jindih <jindih@nvidia.com> Signed-off-by: elvischenv <219235043+elvischenv@users.noreply.github.com> Co-authored-by: jindih <jindih@nvidia.com> Co-authored-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Luka Govedic <lgovedic@redhat.com>

elvischenv committed Aug 29, 2025 at 03:36 UTC 16a45b3a281805ea4d4ff3908cef512fdf6d9f84
11 files changed +746 -64
.buildkite/test-pipeline.yaml
+2
@@ -668,6 +668,7 @@ steps:
668 # Quantization
669 - pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8'
670 - pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py
671 + - pytest -v -s tests/kernels/quantization/test_silu_nvfp4_quant_fusion.py
672 - pytest -v -s tests/kernels/quantization/test_nvfp4_scaled_mm.py
673 - pytest -v -s tests/kernels/quantization/test_flashinfer_scaled_mm.py
674 - pytest -v -s tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py
@@ -677,6 +678,7 @@ steps:
678 - pytest -v -s tests/compile/test_fusion_all_reduce.py
679 - pytest -v -s tests/compile/test_fusion_attn.py::test_attention_quant_pattern
680 - pytest -v -s tests/kernels/moe/test_flashinfer.py
681 + - pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
682
683 ##### 1 GPU test #####
684 ##### multi gpus test #####
CMakeLists.txt
+2
@@ -541,6 +541,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
541 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_ARCHS)
542 set(SRCS
543 "csrc/quantization/fp4/nvfp4_quant_kernels.cu"
544 + "csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
545 "csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu")
546 set_gencode_flags_for_srcs(
547 SRCS "${SRCS}"
@@ -559,6 +560,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
560 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_ARCHS)
561 set(SRCS
562 "csrc/quantization/fp4/nvfp4_quant_kernels.cu"
563 + "csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
564 "csrc/quantization/fp4/nvfp4_experts_quant.cu"
565 "csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
566 "csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
csrc/dispatch_utils.h
+16
@@ -19,6 +19,13 @@
19 #define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
20 AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
21
22 +#define VLLM_DISPATCH_CASE_HALF_TYPES(...) \
23 + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
24 + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
25 +
26 +#define VLLM_DISPATCH_HALF_TYPES(TYPE, NAME, ...) \
27 + AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_HALF_TYPES(__VA_ARGS__))
28 +
29 // ROCm devices might use either fn or fnuz, so set up dispatch table for both.
30 // A host-based check at runtime will create a preferred FP8 type for ROCm
31 // such that the correct kernel is dispatched.
@@ -45,6 +52,15 @@
52 #define VLLM_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \
53 AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__))
54
55 +#define AT_DISPATCH_BYTE_CASE(enum_type, ...) \
56 + AT_PRIVATE_CASE_TYPE_USING_HINT(enum_type, byte_t, __VA_ARGS__)
57 +
58 +#define VLLM_DISPATCH_CASE_BYTE_TYPES(...) \
59 + AT_DISPATCH_BYTE_CASE(at::ScalarType::Byte, __VA_ARGS__)
60 +
61 +#define VLLM_DISPATCH_BYTE_TYPES(TYPE, NAME, ...) \
62 + AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_BYTE_TYPES(__VA_ARGS__))
63 +
64 #define VLLM_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \
65 AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__))
66
csrc/ops.h
+8
@@ -130,6 +130,14 @@ void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
130 void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
131 torch::Tensor& scale);
132
133 +#ifndef USE_ROCM
134 +
135 +void silu_and_mul_nvfp4_quant(torch::Tensor& out,
136 + torch::Tensor& output_block_scale,
137 + torch::Tensor& input,
138 + torch::Tensor& input_global_scale);
139 +#endif
140 +
141 void mul_and_silu(torch::Tensor& out, torch::Tensor& input);
142
143 void gelu_and_mul(torch::Tensor& out, torch::Tensor& input);
csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu new
+368
@@ -0,0 +1,368 @@
1 +/*
2 + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
3 + *
4 + * Licensed under the Apache License, Version 2.0 (the "License");
5 + * you may not use this file except in compliance with the License.
6 + * You may obtain a copy of the License at
7 + *
8 + * http://www.apache.org/licenses/LICENSE-2.0
9 + *
10 + * Unless required by applicable law or agreed to in writing, software
11 + * distributed under the License is distributed on an "AS IS" BASIS,
12 + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 + * See the License for the specific language governing permissions and
14 + * limitations under the License.
15 + */
16 +
17 +#include <torch/all.h>
18 +
19 +#include <cuda_runtime_api.h>
20 +#include <cuda_runtime.h>
21 +
22 +#include <ATen/cuda/CUDAContext.h>
23 +#include <c10/cuda/CUDAGuard.h>
24 +
25 +#include <cuda_fp8.h>
26 +#include "dispatch_utils.h"
27 +
28 +#include "cuda_utils.h"
29 +
30 +namespace vllm {
31 +
32 +// Get type2 from type or vice versa (applied to half and bfloat16)
33 +template <typename T>
34 +struct TypeConverter {
35 + using Type = half2;
36 +}; // keep for generality
37 +
38 +template <>
39 +struct TypeConverter<half2> {
40 + using Type = c10::Half;
41 +};
42 +
43 +template <>
44 +struct TypeConverter<c10::Half> {
45 + using Type = half2;
46 +};
47 +
48 +template <>
49 +struct TypeConverter<__nv_bfloat162> {
50 + using Type = c10::BFloat16;
51 +};
52 +
53 +template <>
54 +struct TypeConverter<c10::BFloat16> {
55 + using Type = __nv_bfloat162;
56 +};
57 +
58 +#define ELTS_PER_THREAD 8
59 +
60 +constexpr int CVT_FP4_ELTS_PER_THREAD = 8;
61 +constexpr int CVT_FP4_SF_VEC_SIZE = 16;
62 +
63 +// Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t).
64 +inline __device__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) {
65 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
66 + uint32_t val;
67 + asm volatile(
68 + "{\n"
69 + ".reg .b8 byte0;\n"
70 + ".reg .b8 byte1;\n"
71 + ".reg .b8 byte2;\n"
72 + ".reg .b8 byte3;\n"
73 + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n"
74 + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n"
75 + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
76 + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
77 + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
78 + "}"
79 + : "=r"(val)
80 + : "f"(array[0]), "f"(array[1]), "f"(array[2]), "f"(array[3]),
81 + "f"(array[4]), "f"(array[5]), "f"(array[6]), "f"(array[7]));
82 + return val;
83 +#else
84 + return 0;
85 +#endif
86 +}
87 +
88 +// Convert 4 float2 values into 8 e2m1 values (represented as one uint32_t).
89 +inline __device__ uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) {
90 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
91 + uint32_t val;
92 + asm volatile(
93 + "{\n"
94 + ".reg .b8 byte0;\n"
95 + ".reg .b8 byte1;\n"
96 + ".reg .b8 byte2;\n"
97 + ".reg .b8 byte3;\n"
98 + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n"
99 + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n"
100 + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
101 + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
102 + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
103 + "}"
104 + : "=r"(val)
105 + : "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y),
106 + "f"(array[2].x), "f"(array[2].y), "f"(array[3].x), "f"(array[3].y));
107 + return val;
108 +#else
109 + return 0;
110 +#endif
111 +}
112 +
113 +// Fast reciprocal.
114 +inline __device__ float reciprocal_approximate_ftz(float a) {
115 + float b;
116 + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a));
117 + return b;
118 +}
119 +
120 +template <class SFType, int CVT_FP4_NUM_THREADS_PER_SF>
121 +__device__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset(int rowIdx, int colIdx,
122 + int numCols,
123 + SFType* SFout) {
124 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
125 + static_assert(CVT_FP4_NUM_THREADS_PER_SF == 1 ||
126 + CVT_FP4_NUM_THREADS_PER_SF == 2);
127 +
128 + // One pair of threads write one SF to global memory.
129 + // TODO: stage through smem for packed STG.32
130 + // is it better than STG.8 from 4 threads ?
131 + if (threadIdx.x % CVT_FP4_NUM_THREADS_PER_SF == 0) {
132 + // SF vector index (16 elements share one SF in the K dimension).
133 + int32_t kIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF;
134 + int32_t mIdx = rowIdx;
135 +
136 + // SF layout [numMTiles, numKTiles, 32 (mTile), 4 (mTile), 4(kTile)]
137 + // --> index [mTileIdx, kTileIdx, outerMIdx, innerMIdx, innerKIdx]
138 +
139 + int32_t mTileIdx = mIdx / (32 * 4);
140 + // SF vector size 16.
141 + int factor = CVT_FP4_SF_VEC_SIZE * 4;
142 + int32_t numKTiles = (numCols + factor - 1) / factor;
143 + int64_t mTileStride = numKTiles * 32 * 4 * 4;
144 +
145 + int32_t kTileIdx = (kIdx / 4);
146 + int64_t kTileStride = 32 * 4 * 4;
147 +
148 + // M tile layout [32, 4] is column-major.
149 + int32_t outerMIdx = (mIdx % 32);
150 + int64_t outerMStride = 4 * 4;
151 +
152 + int32_t innerMIdx = (mIdx % (32 * 4)) / 32;
153 + int64_t innerMStride = 4;
154 +
155 + int32_t innerKIdx = (kIdx % 4);
156 + int64_t innerKStride = 1;
157 +
158 + // Compute the global offset.
159 + int64_t SFOffset = mTileIdx * mTileStride + kTileIdx * kTileStride +
160 + outerMIdx * outerMStride + innerMIdx * innerMStride +
161 + innerKIdx * innerKStride;
162 +
163 + return reinterpret_cast<uint8_t*>(SFout) + SFOffset;
164 + }
165 +#endif
166 + return nullptr;
167 +}
168 +
169 +// Define a 16 bytes packed data type.
170 +template <class Type>
171 +struct PackedVec {
172 + typename TypeConverter<Type>::Type elts[4];
173 +};
174 +
175 +template <>
176 +struct PackedVec<__nv_fp8_e4m3> {
177 + __nv_fp8x2_e4m3 elts[8];
178 +};
179 +
180 +template <class Type>
181 +__inline__ __device__ PackedVec<Type> compute_silu(PackedVec<Type>& vec,
182 + PackedVec<Type>& vec2) {
183 + PackedVec<Type> result;
184 +#pragma unroll
185 + for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; ++i) {
186 + if constexpr (std::is_same_v<Type, c10::Half>) {
187 + half2 val(0.5f, 0.5f);
188 + half2 t0 = __hmul2(vec.elts[i], val);
189 + half2 t1 = __hfma2(h2tanh(t0), val, val);
190 + half2 t2 = __hmul2(vec.elts[i], t1);
191 + result.elts[i] = __hmul2(t2, vec2.elts[i]);
192 + } else {
193 + __nv_bfloat162 val(0.5f, 0.5f);
194 + __nv_bfloat162 t0 = __hmul2(vec.elts[i], val);
195 + __nv_bfloat162 t1 = __hfma2(h2tanh(t0), val, val);
196 + __nv_bfloat162 t2 = __hmul2(vec.elts[i], t1);
197 + result.elts[i] = __hmul2(t2, vec2.elts[i]);
198 + }
199 + }
200 + return result;
201 +}
202 +
203 +// Quantizes the provided PackedVec into the uint32_t output
204 +template <class Type, bool UE8M0_SF = false>
205 +__device__ uint32_t silu_and_cvt_warp_fp16_to_fp4(PackedVec<Type>& vec,
206 + PackedVec<Type>& vec2,
207 + float SFScaleVal,
208 + uint8_t* SFout) {
209 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
210 + PackedVec<Type> out_silu = compute_silu(vec, vec2);
211 + // Get absolute maximum values among the local 8 values.
212 + auto localMax = __habs2(out_silu.elts[0]);
213 +
214 + // Local maximum value.
215 + #pragma unroll
216 + for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
217 + localMax = __hmax2(localMax, __habs2(out_silu.elts[i]));
218 + }
219 +
220 + // Get the absolute maximum among all 16 values (two threads).
221 + localMax = __hmax2(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax);
222 + // Get the final absolute maximum values.
223 + float vecMax = float(__hmax(localMax.x, localMax.y));
224 +
225 + // Get the SF (max value of the vector / max value of e2m1).
226 + // maximum value of e2m1 = 6.0.
227 + // TODO: use half as compute data type.
228 + float SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f));
229 + // 8 bits representation of the SF.
230 + uint8_t fp8SFVal;
231 + // Write the SF to global memory (STG.8).
232 + if constexpr (UE8M0_SF) {
233 + // Extract the 8 exponent bits from float32.
234 + // float 32bits = 1 sign bit + 8 exponent bits + 23 mantissa bits.
235 + uint32_t tmp = reinterpret_cast<uint32_t&>(SFValue) >> 23;
236 + fp8SFVal = tmp & 0xff;
237 + // Convert back to fp32.
238 + reinterpret_cast<uint32_t&>(SFValue) = tmp << 23;
239 + } else {
240 + // Here SFValue is always positive, so E4M3 is the same as UE4M3.
241 + __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue);
242 + reinterpret_cast<__nv_fp8_e4m3&>(fp8SFVal) = tmp;
243 + // Convert back to fp32.
244 + SFValue = float(tmp);
245 + }
246 + // Get the output scale.
247 + // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) *
248 + // reciprocal(SFScaleVal))
249 + float outputScale =
250 + SFValue != 0 ? reciprocal_approximate_ftz(
251 + SFValue * reciprocal_approximate_ftz(SFScaleVal))
252 + : 0.0f;
253 +
254 + if (SFout) {
255 + // Write the SF to global memory (STG.8).
256 + *SFout = fp8SFVal;
257 + }
258 +
259 + // Convert the input to float.
260 + float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
261 +
262 + #pragma unroll
263 + for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
264 + if constexpr (std::is_same_v<Type, c10::Half>) {
265 + fp2Vals[i] = __half22float2(out_silu.elts[i]);
266 + } else {
267 + fp2Vals[i] = __bfloat1622float2(out_silu.elts[i]);
268 + }
269 + fp2Vals[i].x *= outputScale;
270 + fp2Vals[i].y *= outputScale;
271 + }
272 +
273 + // Convert to e2m1 values.
274 + uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals);
275 +
276 + // Write the e2m1 values to global memory.
277 + return e2m1Vec;
278 +#else
279 + return 0;
280 +#endif
281 +}
282 +
283 +// Use UE4M3 by default.
284 +template <class Type, bool UE8M0_SF = false>
285 +__global__ void
286 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
287 +__launch_bounds__(1024, 4) silu_and_cvt_fp16_to_fp4(
288 +#else
289 +silu_and_cvt_fp16_to_fp4(
290 +#endif
291 + int32_t numRows, int32_t numCols, Type const* in, float const* SFScale,
292 + uint32_t* out, uint32_t* SFout) {
293 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
294 + using PackedVec = PackedVec<Type>;
295 + static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
296 + (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
297 + static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
298 + "Vec size is not matched.");
299 +
300 + // Get the global scaling factor, which will be applied to the SF.
301 + // Note SFScale is the same as next GEMM's alpha, which is
302 + // (448.f / (Alpha_A / 6.f)).
303 + float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[0];
304 +
305 + // Input tensor row/col loops.
306 + for (int rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) {
307 + for (int colIdx = threadIdx.x; colIdx < numCols / CVT_FP4_ELTS_PER_THREAD;
308 + colIdx += blockDim.x) {
309 + int64_t inOffset =
310 + rowIdx * (numCols * 2 / CVT_FP4_ELTS_PER_THREAD) + colIdx;
311 + int64_t inOffset2 = rowIdx * (numCols * 2 / CVT_FP4_ELTS_PER_THREAD) +
312 + numCols / CVT_FP4_ELTS_PER_THREAD + colIdx;
313 + PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
314 + PackedVec in_vec2 = reinterpret_cast<PackedVec const*>(in)[inOffset2];
315 +
316 + // Get the output tensor offset.
317 + // Same as inOffset because 8 elements are packed into one uint32_t.
318 + int64_t outOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
319 + ;
320 + auto& out_pos = out[outOffset];
321 +
322 + auto sf_out =
323 + cvt_quant_to_fp4_get_sf_out_offset<uint32_t,
324 + CVT_FP4_NUM_THREADS_PER_SF>(
325 + rowIdx, colIdx, numCols, SFout);
326 +
327 + out_pos = silu_and_cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(
328 + in_vec, in_vec2, SFScaleVal, sf_out);
329 + }
330 + }
331 +#endif
332 +}
333 +
334 +} // namespace vllm
335 +
336 +void silu_and_mul_nvfp4_quant(torch::Tensor& output, // [..., d]
337 + torch::Tensor& output_sf,
338 + torch::Tensor& input, // [..., 2 * d]
339 + torch::Tensor& input_sf) {
340 + TORCH_CHECK(input.dtype() == torch::kFloat16 ||
341 + input.dtype() == torch::kBFloat16);
342 + int32_t m = input.size(0);
343 + int32_t n = input.size(1) / 2;
344 + TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");
345 + int multiProcessorCount =
346 + get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
347 + auto input_sf_ptr = static_cast<float const*>(input_sf.data_ptr());
348 + auto sf_out = static_cast<int32_t*>(output_sf.data_ptr());
349 + auto output_ptr = static_cast<int64_t*>(output.data_ptr());
350 + const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
351 + auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
352 + dim3 block(std::min(int(n / ELTS_PER_THREAD), 1024));
353 + int const numBlocksPerSM = 2048 / block.x;
354 + dim3 grid(std::min(int(m), multiProcessorCount * numBlocksPerSM));
355 + VLLM_DISPATCH_HALF_TYPES(
356 + input.scalar_type(), "act_and_mul_quant_kernel", [&] {
357 + auto input_ptr = reinterpret_cast<scalar_t const*>(input.data_ptr());
358 + VLLM_DISPATCH_BYTE_TYPES(
359 + output.scalar_type(), "fused_act_and_mul_quant_kernel_nvfp4_type",
360 + [&] {
361 + vllm::silu_and_cvt_fp16_to_fp4<scalar_t>
362 + <<<grid, block, 0, stream>>>(
363 + m, n, input_ptr, input_sf_ptr,
364 + reinterpret_cast<uint32_t*>(output_ptr),
365 + reinterpret_cast<uint32_t*>(sf_out));
366 + });
367 + });
368 +}
csrc/torch_bindings.cpp
+7
@@ -115,6 +115,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
115 "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
116 ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
117
118 +#ifndef USE_ROCM
119 + ops.def(
120 + "silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, "
121 + "Tensor input, Tensor input_global_scale) -> ()");
122 + ops.impl("silu_and_mul_nvfp4_quant", torch::kCUDA, &silu_and_mul_nvfp4_quant);
123 +#endif
124 +
125 ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");
126 ops.impl("mul_and_silu", torch::kCUDA, &mul_and_silu);
127
tests/compile/test_silu_mul_quant_fusion.py
+70 -27
@@ -4,32 +4,41 @@ import pytest
4 import torch
5
6 import vllm.envs as envs
7 -from vllm.compilation.activation_quant_fusion import ActivationQuantFusionPass
8 -from vllm.compilation.fx_utils import find_auto_fn, find_auto_fn_maybe
7 +from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
8 +# yapf conflicts with isort for this block
9 +# yapf: disable
10 +from vllm.compilation.activation_quant_fusion import (
11 + FUSED_OPS, SILU_MUL_OP, ActivationQuantFusionPass)
12 +# yapf: enable
13 +from vllm.compilation.fusion import QUANT_OPS
14 from vllm.compilation.noop_elimination import NoOpEliminationPass
15 from vllm.config import CompilationConfig, PassConfig, VllmConfig
16 from vllm.model_executor.layers.activation import SiluAndMul
17 from vllm.model_executor.layers.quantization.utils.quant_utils import (
13 - GroupShape)
18 + GroupShape, kFp8StaticTensorSym, kNvfp4Quant)
19 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
20 Fp8LinearOp)
21 from vllm.platforms import current_platform
22
23 from .backend import TestBackend
24
25 +FP8_DTYPE = current_platform.fp8_dtype()
26 +FP4_DTYPE = torch.uint8
27
21 -class TestModel(torch.nn.Module):
28
23 - def __init__(self, hidden_size: int, force_fp8_e4m3fnuz: bool, *args,
24 - **kwargs):
25 - super().__init__(*args, **kwargs)
29 +def is_nvfp4_supported():
30 + return current_platform.has_device_capability(100)
31 +
32 +
33 +class TestSiluMulFp8QuantModel(torch.nn.Module):
34 +
35 + def __init__(self, hidden_size: int, force_fp8_e4m3fnuz: bool, **kwargs):
36 + super().__init__()
37 self.silu_and_mul = SiluAndMul()
38 self.wscale = torch.rand(1, dtype=torch.float32)
39 self.scale = torch.rand(1, dtype=torch.float32)
40
30 - self.w = (torch.rand(
31 - hidden_size,
32 - hidden_size).to(dtype=current_platform.fp8_dtype()).t())
41 + self.w = torch.rand(hidden_size, hidden_size).to(dtype=FP8_DTYPE).t()
42
43 self.fp8_linear = Fp8LinearOp(
44 force_fp8_e4m3fnuz=force_fp8_e4m3fnuz,
@@ -45,14 +54,56 @@ class TestModel(torch.nn.Module):
54 input_scale=self.wscale)
55 return x2
56
57 + def ops_in_model_before(self):
58 + return [SILU_MUL_OP, QUANT_OPS[kFp8StaticTensorSym]]
59 +
60 + def ops_in_model_after(self):
61 + return [FUSED_OPS[kFp8StaticTensorSym]]
62 +
63 +
64 +class TestSiluMulNvfp4QuantModel(torch.nn.Module):
65 +
66 + def __init__(self, hidden_size: int, **kwargs):
67 + super().__init__()
68 + self.silu_and_mul = SiluAndMul()
69 + self.w = torch.randint(256, (hidden_size, hidden_size // 2),
70 + dtype=FP4_DTYPE)
71 + self.wscale = torch.randn(hidden_size,
72 + hidden_size // 16).to(dtype=FP8_DTYPE)
73 + self.wscale2 = torch.rand(1, dtype=torch.float32)
74 + self.scale = torch.rand(1, dtype=torch.float32)
75
49 -@pytest.mark.parametrize("num_tokens", [256])
50 -@pytest.mark.parametrize("hidden_size", [64])
76 + def forward(self, x):
77 + y = self.silu_and_mul(x)
78 + y_quant, y_block_scale = scaled_fp4_quant(y, 1 / self.scale)
79 + out = cutlass_scaled_fp4_mm(a=y_quant,
80 + b=self.w,
81 + block_scale_a=y_block_scale,
82 + block_scale_b=self.wscale,
83 + alpha=self.scale * self.wscale2,
84 + out_dtype=y.dtype)
85 + return out
86 +
87 + def ops_in_model_before(self):
88 + return [SILU_MUL_OP, QUANT_OPS[kNvfp4Quant]]
89 +
90 + def ops_in_model_after(self):
91 + return [FUSED_OPS[kNvfp4Quant]]
92 +
93 +
94 +@pytest.mark.parametrize("num_tokens", [64])
95 +@pytest.mark.parametrize("hidden_size", [128])
96 +@pytest.mark.parametrize(
97 + "model_class", [TestSiluMulFp8QuantModel, TestSiluMulNvfp4QuantModel]
98 + if is_nvfp4_supported() else [TestSiluMulFp8QuantModel])
99 @pytest.mark.parametrize("force_fp8_e4m3fnuz", [True, False])
100 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"],
101 reason="Only test on CUDA and ROCm")
54 -def test_fusion_silu_and_mul_quant(num_tokens, hidden_size,
102 +def test_fusion_silu_and_mul_quant(num_tokens, hidden_size, model_class,
103 force_fp8_e4m3fnuz):
104 + if model_class == TestSiluMulNvfp4QuantModel and force_fp8_e4m3fnuz:
105 + pytest.skip("Duplicate tests for NVFP4")
106 +
107 torch.set_default_device("cuda")
108 torch.set_default_dtype(torch.float16)
109
@@ -63,7 +114,8 @@ def test_fusion_silu_and_mul_quant(num_tokens, hidden_size,
114 fusion_pass = ActivationQuantFusionPass(config)
115
116 backend = TestBackend(NoOpEliminationPass(config), fusion_pass)
66 - model = TestModel(hidden_size, force_fp8_e4m3fnuz)
117 + model = model_class(hidden_size=hidden_size,
118 + force_fp8_e4m3fnuz=force_fp8_e4m3fnuz)
119
120 # First dimension dynamic
121 x = torch.rand(num_tokens, hidden_size * 2)
@@ -80,17 +132,8 @@ def test_fusion_silu_and_mul_quant(num_tokens, hidden_size,
132 atol=1e-3,
133 rtol=1e-3)
134
83 - # Check substitution worked
84 - pre_nodes = backend.graph_pre_pass.nodes
85 - post_nodes = backend.graph_post_pass.nodes
86 -
87 - silu_and_mul_quant = torch.ops._C.silu_and_mul_quant.default
88 - fp8_quant = torch.ops._C.static_scaled_fp8_quant.default
89 -
90 - # In pre-nodes, fp8 quant should be present and fused kernels should not
91 - assert find_auto_fn_maybe(pre_nodes, silu_and_mul_quant) is None
92 - find_auto_fn(pre_nodes, fp8_quant)
135 + # In pre-nodes, quant op should be present and fused kernels should not
136 + backend.check_before_ops(model.ops_in_model_before())
137
94 - # In post-nodes, fused kernels should be present and fp8 quant should not
95 - find_auto_fn(post_nodes, silu_and_mul_quant)
96 - assert find_auto_fn_maybe(post_nodes, fp8_quant) is None
138 + # In post-nodes, fused kernels should be present and quant op should not
139 + backend.check_after_ops(model.ops_in_model_after())
tests/kernels/quantization/test_silu_nvfp4_quant_fusion.py new
+126
@@ -0,0 +1,126 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import pytest
4 +import torch
5 +
6 +from tests.kernels.utils import opcheck
7 +from vllm.model_executor.layers.activation import SiluAndMul
8 +from vllm.platforms import current_platform
9 +from vllm.scalar_type import scalar_types
10 +
11 +if not current_platform.has_device_capability(100):
12 + pytest.skip(reason="Nvfp4 Requires compute capability of 10 or above.",
13 + allow_module_level=True)
14 +
15 +DTYPES = [torch.float16, torch.bfloat16]
16 +SHAPES = [(128, 64), (128, 128), (256, 64), (256, 128)]
17 +SEEDS = [42]
18 +CUDA_DEVICES = ['cuda:0']
19 +
20 +FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
21 +FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
22 +
23 +BLOCK_SIZE = 16
24 +
25 +
26 +def ref_impl(silu_and_mul: SiluAndMul, x: torch.Tensor,
27 + global_scale: torch.Tensor,
28 + ref_output_scale: torch.Tensor) -> torch.Tensor:
29 + silu_and_mul_out = silu_and_mul.forward_native(x)
30 + assert not current_platform.is_rocm()
31 + assert silu_and_mul_out.ndim >= 1, (
32 + f'input.ndim needs to be >= 1, but got {silu_and_mul_out.ndim}.')
33 + other_dims = 1 if silu_and_mul_out.ndim == 1 else -1
34 + silu_and_mul_out = silu_and_mul_out.reshape(other_dims,
35 + silu_and_mul_out.shape[-1])
36 + m, n = silu_and_mul_out.shape
37 + device = silu_and_mul_out.device
38 +
39 + # Two fp4 values will be packed into an uint8.
40 + out = torch.empty((m, n // 2), device=device, dtype=torch.uint8)
41 +
42 + output_scale = ref_output_scale
43 +
44 + torch.ops._C.scaled_fp4_quant(out, silu_and_mul_out, output_scale,
45 + global_scale)
46 +
47 + return out, output_scale
48 +
49 +
50 +def ops_impl(x: torch.Tensor, global_scale: torch.Tensor,
51 + ref_output_scale: torch.Tensor) -> torch.Tensor:
52 + out_shape = (x.shape[0], x.shape[1] // 4)
53 + output_scale = ref_output_scale
54 + out = torch.empty(out_shape, dtype=torch.uint8, device=x.device)
55 + torch.ops._C.silu_and_mul_nvfp4_quant(out, output_scale, x, global_scale)
56 + return out, output_scale
57 +
58 +
59 +@pytest.mark.parametrize("dtype", DTYPES)
60 +@pytest.mark.parametrize("shape", SHAPES)
61 +@pytest.mark.parametrize("seed", SEEDS)
62 +@pytest.mark.parametrize("device", CUDA_DEVICES)
63 +@torch.inference_mode()
64 +def test_quantize_to_fp4(
65 + dtype: torch.dtype,
66 + shape: tuple[int, int],
67 + seed: int,
68 + device: str,
69 +) -> None:
70 + current_platform.seed_everything(seed)
71 + torch.set_default_device(device)
72 +
73 + m, n = shape
74 +
75 + x = torch.randn((m, n), dtype=dtype)
76 + tensor_amax = torch.abs(x).max().to(torch.float32)
77 + global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
78 +
79 + block_size = 16
80 +
81 + assert n % block_size == 0, (
82 + f'last dim has to be multiple of 16, but got {n}.')
83 + assert x.dtype in (torch.float16, torch.bfloat16), (
84 + f'input.dtype needs to be fp16 or bf16 but got {x.dtype}.')
85 +
86 + round_up = lambda x, y: (x + y - 1) // y * y
87 + rounded_m = round_up(x.shape[0], 128)
88 + scale_n = x.shape[1] // (2 * block_size)
89 + rounded_n = round_up(scale_n, 4)
90 + output_scale = torch.empty((rounded_m, rounded_n // 4),
91 + device=x.device,
92 + dtype=torch.int32)
93 +
94 + layer = SiluAndMul()
95 +
96 + ref_out, ref_out_scale = ref_impl(layer, x, global_scale, output_scale)
97 +
98 + fusion_out, fusion_out_scale = ops_impl(x, global_scale, output_scale)
99 +
100 + assert ref_out.dtype == torch.uint8
101 + assert fusion_out.dtype == torch.uint8
102 + assert ref_out.shape == fusion_out.shape
103 +
104 + assert ref_out_scale.dtype == torch.int32
105 + assert fusion_out_scale.dtype == torch.int32
106 + assert ref_out_scale.shape == fusion_out_scale.shape
107 +
108 + # Allow up to 2% of mismatched values since BF16 has accuracy issues.
109 + mis_threshold = 0.02
110 + atol = 0.4
111 + rtol = 0.4
112 + ref_logits = ref_out[-1]
113 + fusion_logits = fusion_out[-1]
114 +
115 + mis_count = torch.sum(
116 + torch.abs(fusion_logits - ref_logits) > (atol +
117 + rtol * torch.abs(ref_logits)))
118 + mis_ratio = mis_count / fusion_logits.numel()
119 +
120 + assert mis_ratio < mis_threshold, \
121 + f"Mismatch ratio {mis_ratio} exceeds threshold {mis_threshold}"
122 +
123 + torch.testing.assert_close(ref_out_scale, fusion_out_scale)
124 +
125 + opcheck(torch.ops._C.silu_and_mul_nvfp4_quant,
126 + (fusion_out, fusion_out_scale, x, global_scale))
vllm/compilation/activation_quant_fusion.py
+135 -35
@@ -1,55 +1,154 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from abc import ABC, abstractmethod
5 +
6 import torch
7 from torch._higher_order_ops.auto_functionalize import auto_functionalized
8 from torch._inductor.pattern_matcher import (PatternMatcherPass, fwd_only,
9 register_replacement)
10 +from torch._ops import OpOverload
11
12 from vllm.config import VllmConfig
13 from vllm.logger import init_logger
14 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
15 + QuantKey, kFp8StaticTensorSym, kNvfp4Quant, kStaticTensorScale)
16 from vllm.platforms import current_platform
17
18 +from .fusion import QUANT_OPS, empty_bf16, empty_fp32, empty_i32
19 from .inductor_pass import enable_fake_mode
20 from .vllm_inductor_pass import VllmInductorPass
21
22 logger = init_logger(__name__)
23
24 +FP8_DTYPE = current_platform.fp8_dtype()
25 +FP4_DTYPE = torch.uint8
26 +
27 +SILU_MUL_OP = torch.ops._C.silu_and_mul.default
28
19 -def silu_mul_pattern_static(result: torch.Tensor,
20 - result_silu_mul: torch.Tensor, input: torch.Tensor,
21 - scale: torch.Tensor):
22 - at1 = auto_functionalized(torch.ops._C.silu_and_mul.default,
23 - result=result_silu_mul,
24 - input=input)
25 - at2 = auto_functionalized(torch.ops._C.static_scaled_fp8_quant.default,
26 - result=result,
27 - input=at1[1],
28 - scale=scale)
29 - return at2[1]
29 +FUSED_OPS: dict[QuantKey, OpOverload] = {
30 + kFp8StaticTensorSym: torch.ops._C.silu_and_mul_quant.default, # noqa: E501
31 +}
32 +if current_platform.is_cuda() and hasattr(torch.ops._C,
33 + "silu_and_mul_nvfp4_quant"):
34 + FUSED_OPS[
35 + kNvfp4Quant] = torch.ops._C.silu_and_mul_nvfp4_quant.default # noqa: E501
36
37
32 -def silu_mul_replacement_static(result: torch.Tensor,
33 - result_silu_mul: torch.Tensor,
34 - input: torch.Tensor, scale: torch.Tensor):
35 - at = auto_functionalized(torch.ops._C.silu_and_mul_quant.default,
36 - result=result,
37 - input=input,
38 - scale=scale)
39 - return at[1]
38 +class ActivationQuantPattern(ABC):
39 + """
40 + The base class for Activation+Quant fusions.
41 + Should not be used directly.
42 + """
43
44 + def __init__(
45 + self,
46 + quant_key: QuantKey,
47 + ):
48 + self.quant_key = quant_key
49 + self.quant_dtype = quant_key.dtype
50
42 -def empty_bf16(*args, **kwargs):
43 - return torch.empty(*args, **kwargs, dtype=torch.bfloat16, device="cuda")
51 + assert self.quant_key in QUANT_OPS, \
52 + f"unsupported quantization scheme {self.quant_key}"
53 + self.QUANT_OP = QUANT_OPS[self.quant_key]
54
55 + assert self.quant_key in FUSED_OPS, \
56 + f"unsupported fusion scheme {self.quant_key}"
57 + self.FUSED_OP = FUSED_OPS[self.quant_key]
58
46 -def empty_fp8(*args, **kwargs):
47 - fp8 = current_platform.fp8_dtype()
48 - return torch.empty(*args, **kwargs, dtype=fp8, device="cuda")
59 + def empty_quant(self, *args, **kwargs):
60 + kwargs = {'dtype': self.quant_dtype, 'device': "cuda", **kwargs}
61 + return torch.empty(*args, **kwargs)
62
63 + @abstractmethod
64 + def register(self, pm_pass: PatternMatcherPass):
65 + raise NotImplementedError
66
51 -def empty_fp32(*args, **kwargs):
52 - return torch.empty(*args, **kwargs, dtype=torch.float32, device="cuda")
67 +
68 +class SiluMulFp8StaticQuantPattern(ActivationQuantPattern):
69 + """
70 + Fusion for SiluMul+Fp8StaticQuant Pattern
71 + """
72 +
73 + def __init__(self, symmetric: bool = True):
74 + quant_key = QuantKey(dtype=FP8_DTYPE,
75 + scale=kStaticTensorScale,
76 + symmetric=symmetric)
77 + super().__init__(quant_key)
78 +
79 + def register(self, pm_pass: PatternMatcherPass):
80 +
81 + def pattern(result: torch.Tensor, result_silu_mul: torch.Tensor,
82 + input: torch.Tensor, scale: torch.Tensor):
83 + at1 = auto_functionalized(SILU_MUL_OP,
84 + result=result_silu_mul,
85 + input=input)
86 + at2 = auto_functionalized(self.QUANT_OP,
87 + result=result,
88 + input=at1[1],
89 + scale=scale)
90 + return at2[1]
91 +
92 + def replacement(result: torch.Tensor, result_silu_mul: torch.Tensor,
93 + input: torch.Tensor, scale: torch.Tensor):
94 + at = auto_functionalized(self.FUSED_OP,
95 + result=result,
96 + input=input,
97 + scale=scale)
98 + return at[1]
99 +
100 + inputs = [
101 + self.empty_quant(5, 4), # result
102 + empty_bf16(5, 4), # result_silu_mul
103 + empty_bf16(5, 4), # input
104 + empty_fp32(1, 1) # scale
105 + ]
106 +
107 + register_replacement(pattern, replacement, inputs, fwd_only, pm_pass)
108 +
109 +
110 +class SiluMulNvfp4QuantPattern(ActivationQuantPattern):
111 + """
112 + Fusion for SiluMul+Nvfp4Quant Pattern
113 + """
114 +
115 + def __init__(self):
116 + super().__init__(kNvfp4Quant)
117 +
118 + def register(self, pm_pass: PatternMatcherPass):
119 +
120 + def pattern(result: torch.Tensor, output_scale: torch.Tensor,
121 + result_silu_mul: torch.Tensor, input: torch.Tensor,
122 + scale: torch.Tensor):
123 + at1 = auto_functionalized(SILU_MUL_OP,
124 + result=result_silu_mul,
125 + input=input)
126 + at2 = auto_functionalized(self.QUANT_OP,
127 + output=result,
128 + input=at1[1],
129 + output_scale=output_scale,
130 + input_scale=scale)
131 + return at2[1], at2[2]
132 +
133 + def replacement(result: torch.Tensor, output_scale: torch.Tensor,
134 + result_silu_mul: torch.Tensor, input: torch.Tensor,
135 + scale: torch.Tensor):
136 + at = auto_functionalized(self.FUSED_OP,
137 + result=result,
138 + result_block_scale=output_scale,
139 + input=input,
140 + input_global_scale=scale)
141 + return at[1], at[2]
142 +
143 + inputs = [
144 + self.empty_quant(5, 32), # result
145 + empty_i32(128, 4), # output_scale
146 + empty_bf16(5, 64), # result_silu_mul
147 + empty_bf16(5, 64), # input
148 + empty_fp32(1, 1) # scale
149 + ]
150 +
151 + register_replacement(pattern, replacement, inputs, fwd_only, pm_pass)
152
153
154 class ActivationQuantFusionPass(VllmInductorPass):
@@ -69,15 +168,11 @@ class ActivationQuantFusionPass(VllmInductorPass):
168 self.patterns: PatternMatcherPass = PatternMatcherPass(
169 pass_name="activation_quant_fusion_pass")
170
72 - inputs = [
73 - empty_fp8(5, 4), # Quant output
74 - empty_bf16(5, 4), # Silu_and_mul output
75 - empty_bf16(5, 4), # Input
76 - empty_fp32(1, 1) # Scale
77 - ]
78 - register_replacement(silu_mul_pattern_static,
79 - silu_mul_replacement_static, inputs, fwd_only,
80 - self.patterns)
171 + pattern_silu_mul_fp8 = SiluMulFp8StaticQuantPattern()
172 + pattern_silu_mul_fp8.register(self.patterns)
173 +
174 + pattern_silu_mul_nvfp4 = SiluMulNvfp4QuantPattern()
175 + pattern_silu_mul_nvfp4.register(self.patterns)
176
177 def __call__(self, graph: torch.fx.Graph):
178 self.begin()
@@ -89,3 +184,8 @@ class ActivationQuantFusionPass(VllmInductorPass):
184
185 self.dump_graph(graph, "after_act_quant_fusion")
186 self.end_and_log()
187 +
188 + def uuid(self):
189 + return VllmInductorPass.hash_source(self, ActivationQuantPattern,
190 + SiluMulFp8StaticQuantPattern,
191 + SiluMulNvfp4QuantPattern)
vllm/compilation/fix_functionalization.py
+7
@@ -97,6 +97,13 @@ class FixFunctionalizationPass(VllmInductorPass):
97 node,
98 mutated_args,
99 args=('result', 'input', 'scale'))
100 + elif at_target == torch.ops._C.silu_and_mul_nvfp4_quant.default:
101 + mutated_args = {1: 'result', 2: 'result_block_scale'}
102 + self.defunctionalize(graph,
103 + node,
104 + mutated_args,
105 + args=('result', 'result_block_scale',
106 + 'input', 'input_global_scale'))
107 else:
108 continue # skip the count
109
vllm/model_executor/layers/quantization/modelopt.py
+5 -2
@@ -885,6 +885,10 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
885 layer.alpha = Parameter(layer.input_scale * layer.weight_scale_2,
886 requires_grad=False)
887
888 + # Calculate `1 / input_scale` so that we don't need to do so at runtime
889 + layer.input_scale_inv = Parameter(
890 + (1 / layer.input_scale).to(torch.float32), requires_grad=False)
891 +
892 # Swizzle the weight blockscale.
893 # contracting dimension is input dimension
894 # block_size = 16;
@@ -941,8 +945,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
945 output_shape = [x.shape[0], layer.weight.shape[0]]
946
947 # quantize BF16 or FP16 to (FP4 and interleaved block scale)
944 - s_quant = 1 / layer.input_scale
945 - x_fp4, x_blockscale = scaled_fp4_quant(x, s_quant)
948 + x_fp4, x_blockscale = scaled_fp4_quant(x, layer.input_scale_inv)
949
950 # validate dtypes of quantized input, input block scale,
951 # weight and weight_blockscale