@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / adc3ddb430
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

[Bugfix][Misc] Fix silu_and_mul_nvfp4_quant issue and extract common utils for nvfp4 kernel source files (#23727)

Signed-off-by: elvischenv <219235043+elvischenv@users.noreply.github.com> Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>

elvischenv committed Sep 5, 2025 at 05:25 UTC adc3ddb4309d4843a5425eed702bb3e1d942fc2d
11 files changed +382 -718
.buildkite/test-pipeline.yaml
+2 -2
@@ -666,7 +666,7 @@ steps:
666 # Quantization
667 - pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8'
668 - pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py
669 - # - pytest -v -s tests/kernels/quantization/test_silu_nvfp4_quant_fusion.py
669 + - pytest -v -s tests/kernels/quantization/test_silu_nvfp4_quant_fusion.py
670 - pytest -v -s tests/kernels/quantization/test_nvfp4_scaled_mm.py
671 - pytest -v -s tests/kernels/quantization/test_flashinfer_scaled_mm.py
672 - pytest -v -s tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py
@@ -676,7 +676,7 @@ steps:
676 - pytest -v -s tests/compile/test_fusion_all_reduce.py
677 - pytest -v -s tests/compile/test_fusion_attn.py::test_attention_quant_pattern
678 - pytest -v -s tests/kernels/moe/test_flashinfer.py
679 - # - pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
679 + - pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
680
681 ##### 1 GPU test #####
682 ##### multi gpus test #####
csrc/dispatch_utils.h
-9
@@ -52,15 +52,6 @@
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 -
55 #define VLLM_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \
56 AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__))
57
csrc/ops.h
+1 -2
@@ -130,8 +130,7 @@ 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 -#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
134 - (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
133 +#ifndef USE_ROCM
134 void silu_and_mul_nvfp4_quant(torch::Tensor& out,
135 torch::Tensor& output_block_scale,
136 torch::Tensor& input,
csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu
+28 -184
@@ -26,164 +26,17 @@
26 #include "dispatch_utils.h"
27
28 #include "cuda_utils.h"
29 +#include "nvfp4_utils.cuh"
30
31 namespace vllm {
32
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 -
33 template <class Type>
34 __inline__ __device__ PackedVec<Type> compute_silu(PackedVec<Type>& vec,
35 PackedVec<Type>& vec2) {
36 PackedVec<Type> result;
37 #pragma unroll
38 for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; ++i) {
186 - if constexpr (std::is_same_v<Type, c10::Half>) {
39 + if constexpr (std::is_same_v<Type, half>) {
40 half2 val(0.5f, 0.5f);
41 half2 t0 = __hmul2(vec.elts[i], val);
42 half2 t1 = __hfma2(h2tanh(t0), val, val);
@@ -206,13 +59,12 @@ __device__ uint32_t silu_and_cvt_warp_fp16_to_fp4(PackedVec<Type>& vec,
59 PackedVec<Type>& vec2,
60 float SFScaleVal,
61 uint8_t* SFout) {
209 -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
62 PackedVec<Type> out_silu = compute_silu(vec, vec2);
63 // Get absolute maximum values among the local 8 values.
64 auto localMax = __habs2(out_silu.elts[0]);
65
214 - // Local maximum value.
215 - #pragma unroll
66 +// Local maximum value.
67 +#pragma unroll
68 for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
69 localMax = __hmax2(localMax, __habs2(out_silu.elts[i]));
70 }
@@ -259,9 +111,9 @@ __device__ uint32_t silu_and_cvt_warp_fp16_to_fp4(PackedVec<Type>& vec,
111 // Convert the input to float.
112 float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
113
262 - #pragma unroll
114 +#pragma unroll
115 for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
264 - if constexpr (std::is_same_v<Type, c10::Half>) {
116 + if constexpr (std::is_same_v<Type, half>) {
117 fp2Vals[i] = __half22float2(out_silu.elts[i]);
118 } else {
119 fp2Vals[i] = __bfloat1622float2(out_silu.elts[i]);
@@ -275,22 +127,14 @@ __device__ uint32_t silu_and_cvt_warp_fp16_to_fp4(PackedVec<Type>& vec,
127
128 // Write the e2m1 values to global memory.
129 return e2m1Vec;
278 -#else
279 - return 0;
280 -#endif
130 }
131
132 // Use UE4M3 by default.
133 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)
134 +__global__ void __launch_bounds__(1024, 4)
135 + silu_and_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
136 + float const* SFScale, uint32_t* out,
137 + uint32_t* SFout) {
138 using PackedVec = PackedVec<Type>;
139 static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
140 (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
@@ -328,22 +172,25 @@ silu_and_cvt_fp16_to_fp4(
172 in_vec, in_vec2, SFScaleVal, sf_out);
173 }
174 }
331 -#endif
175 }
176
177 } // namespace vllm
178
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);
179 +void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output, // [..., d]
180 + torch::Tensor& output_sf,
181 + torch::Tensor& input, // [..., 2 * d]
182 + torch::Tensor& input_sf) {
183 int32_t m = input.size(0);
184 int32_t n = input.size(1) / 2;
185 +
186 TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");
187 + TORCH_CHECK(input.scalar_type() == at::ScalarType::Half ||
188 + input.scalar_type() == at::ScalarType::BFloat16,
189 + "Unsupported input data type for quantize_to_fp4.");
190 +
191 int multiProcessorCount =
192 get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
193 +
194 auto input_sf_ptr = static_cast<float const*>(input_sf.data_ptr());
195 auto sf_out = static_cast<int32_t*>(output_sf.data_ptr());
196 auto output_ptr = static_cast<int64_t*>(output.data_ptr());
@@ -352,17 +199,14 @@ void silu_and_mul_nvfp4_quant(torch::Tensor& output, // [..., d]
199 dim3 block(std::min(int(n / ELTS_PER_THREAD), 1024));
200 int const numBlocksPerSM = 2048 / block.x;
201 dim3 grid(std::min(int(m), multiProcessorCount * numBlocksPerSM));
202 +
203 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 - });
204 + input.scalar_type(), "silu_and_mul_nvfp4_quant_kernel", [&] {
205 + using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
206 + auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
207 + vllm::silu_and_cvt_fp16_to_fp4<cuda_type><<<grid, block, 0, stream>>>(
208 + m, n, input_ptr, input_sf_ptr,
209 + reinterpret_cast<uint32_t*>(output_ptr),
210 + reinterpret_cast<uint32_t*>(sf_out));
211 });
212 }
csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu
+16
@@ -1,3 +1,19 @@
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 #include <cutlass/arch/arch.h>
19
csrc/quantization/fp4/nvfp4_experts_quant.cu
+47 -263
@@ -1,247 +1,42 @@
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
6 -#include <cuda_runtime.h>
25 #include <cuda_fp8.h>
26 +#include "dispatch_utils.h"
27
9 -template <typename T>
10 -struct TypeConverter {
11 - using Type = half2;
12 -}; // keep for generality
13 -
14 -template <>
15 -struct TypeConverter<half2> {
16 - using Type = half;
17 -};
18 -
19 -template <>
20 -struct TypeConverter<half> {
21 - using Type = half2;
22 -};
23 -
24 -template <>
25 -struct TypeConverter<__nv_bfloat162> {
26 - using Type = __nv_bfloat16;
27 -};
28 -
29 -template <>
30 -struct TypeConverter<__nv_bfloat16> {
31 - using Type = __nv_bfloat162;
32 -};
33 -
34 -#define ELTS_PER_THREAD 8
35 -
36 -constexpr int CVT_FP4_ELTS_PER_THREAD = 8;
37 -constexpr int CVT_FP4_SF_VEC_SIZE = 16;
38 -
39 -// Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t).
40 -inline __device__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) {
41 -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
42 - uint32_t val;
43 - asm volatile(
44 - "{\n"
45 - ".reg .b8 byte0;\n"
46 - ".reg .b8 byte1;\n"
47 - ".reg .b8 byte2;\n"
48 - ".reg .b8 byte3;\n"
49 - "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n"
50 - "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n"
51 - "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
52 - "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
53 - "mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
54 - "}"
55 - : "=r"(val)
56 - : "f"(array[0]), "f"(array[1]), "f"(array[2]), "f"(array[3]),
57 - "f"(array[4]), "f"(array[5]), "f"(array[6]), "f"(array[7]));
58 - return val;
59 -#else
60 - return 0;
61 -#endif
62 -}
63 -
64 -// Convert 4 float2 values into 8 e2m1 values (represented as one uint32_t).
65 -inline __device__ uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) {
66 -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
67 - uint32_t val;
68 - asm volatile(
69 - "{\n"
70 - ".reg .b8 byte0;\n"
71 - ".reg .b8 byte1;\n"
72 - ".reg .b8 byte2;\n"
73 - ".reg .b8 byte3;\n"
74 - "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n"
75 - "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n"
76 - "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
77 - "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
78 - "mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
79 - "}"
80 - : "=r"(val)
81 - : "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y),
82 - "f"(array[2].x), "f"(array[2].y), "f"(array[3].x), "f"(array[3].y));
83 - return val;
84 -#else
85 - return 0;
86 -#endif
87 -}
88 -
89 -// Fast reciprocal.
90 -inline __device__ float reciprocal_approximate_ftz(float a) {
91 - float b;
92 - asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a));
93 - return b;
94 -}
28 +#include "nvfp4_utils.cuh"
29
96 -template <class SFType, int CVT_FP4_NUM_THREADS_PER_SF>
97 -__device__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset(int rowIdx, int colIdx,
98 - int numCols,
99 - SFType* SFout) {
100 -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
101 - static_assert(CVT_FP4_NUM_THREADS_PER_SF == 1 ||
102 - CVT_FP4_NUM_THREADS_PER_SF == 2);
103 -
104 - // One pair of threads write one SF to global memory.
105 - // TODO: stage through smem for packed STG.32
106 - // is it better than STG.8 from 4 threads ?
107 - if (threadIdx.x % CVT_FP4_NUM_THREADS_PER_SF == 0) {
108 - // SF vector index (16 elements share one SF in the K dimension).
109 - int32_t kIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF;
110 - int32_t mIdx = rowIdx;
111 -
112 - // SF layout [numMTiles, numKTiles, 32 (mTile), 4 (mTile), 4(kTile)]
113 - // --> index [mTileIdx, kTileIdx, outerMIdx, innerMIdx, innerKIdx]
114 -
115 - int32_t mTileIdx = mIdx / (32 * 4);
116 - // SF vector size 16.
117 - int factor = CVT_FP4_SF_VEC_SIZE * 4;
118 - int32_t numKTiles = (numCols + factor - 1) / factor;
119 - int64_t mTileStride = numKTiles * 32 * 4 * 4;
120 -
121 - int32_t kTileIdx = (kIdx / 4);
122 - int64_t kTileStride = 32 * 4 * 4;
123 -
124 - // M tile layout [32, 4] is column-major.
125 - int32_t outerMIdx = (mIdx % 32);
126 - int64_t outerMStride = 4 * 4;
127 -
128 - int32_t innerMIdx = (mIdx % (32 * 4)) / 32;
129 - int64_t innerMStride = 4;
130 -
131 - int32_t innerKIdx = (kIdx % 4);
132 - int64_t innerKStride = 1;
133 -
134 - // Compute the global offset.
135 - int64_t SFOffset = mTileIdx * mTileStride + kTileIdx * kTileStride +
136 - outerMIdx * outerMStride + innerMIdx * innerMStride +
137 - innerKIdx * innerKStride;
138 -
139 - return reinterpret_cast<uint8_t*>(SFout) + SFOffset;
140 - }
141 -#endif
142 - return nullptr;
143 -}
144 -
145 -// Define a 16 bytes packed data type.
146 -template <class Type>
147 -struct PackedVec {
148 - typename TypeConverter<Type>::Type elts[4];
149 -};
150 -
151 -template <>
152 -struct PackedVec<__nv_fp8_e4m3> {
153 - __nv_fp8x2_e4m3 elts[8];
154 -};
155 -
156 -// Quantizes the provided PackedVec into the uint32_t output
157 -template <class Type, bool UE8M0_SF = false>
158 -__device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal,
159 - uint8_t* SFout) {
160 -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
161 - // Get absolute maximum values among the local 8 values.
162 - auto localMax = __habs2(vec.elts[0]);
163 -
164 - // Local maximum value.
165 - #pragma unroll
166 - for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
167 - localMax = __hmax2(localMax, __habs2(vec.elts[i]));
168 - }
169 -
170 - // Get the absolute maximum among all 16 values (two threads).
171 - localMax = __hmax2(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax);
172 - // Get the final absolute maximum values.
173 - float vecMax = float(__hmax(localMax.x, localMax.y));
174 -
175 - // Get the SF (max value of the vector / max value of e2m1).
176 - // maximum value of e2m1 = 6.0.
177 - // TODO: use half as compute data type.
178 - float SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f));
179 - // 8 bits representation of the SF.
180 - uint8_t fp8SFVal;
181 - // Write the SF to global memory (STG.8).
182 - if constexpr (UE8M0_SF) {
183 - // Extract the 8 exponent bits from float32.
184 - // float 32bits = 1 sign bit + 8 exponent bits + 23 mantissa bits.
185 - uint32_t tmp = reinterpret_cast<uint32_t&>(SFValue) >> 23;
186 - fp8SFVal = tmp & 0xff;
187 - // Convert back to fp32.
188 - reinterpret_cast<uint32_t&>(SFValue) = tmp << 23;
189 - } else {
190 - // Here SFValue is always positive, so E4M3 is the same as UE4M3.
191 - __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue);
192 - reinterpret_cast<__nv_fp8_e4m3&>(fp8SFVal) = tmp;
193 - // Convert back to fp32.
194 - SFValue = float(tmp);
195 - }
196 - // Get the output scale.
197 - // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) *
198 - // reciprocal(SFScaleVal))
199 - float outputScale =
200 - SFValue != 0 ? reciprocal_approximate_ftz(
201 - SFValue * reciprocal_approximate_ftz(SFScaleVal))
202 - : 0.0f;
203 -
204 - if (SFout) {
205 - // Write the SF to global memory (STG.8).
206 - *SFout = fp8SFVal;
207 - }
208 -
209 - // Convert the input to float.
210 - float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
211 -
212 - #pragma unroll
213 - for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
214 - if constexpr (std::is_same_v<Type, half>) {
215 - fp2Vals[i] = __half22float2(vec.elts[i]);
216 - } else {
217 - fp2Vals[i] = __bfloat1622float2(vec.elts[i]);
218 - }
219 - fp2Vals[i].x *= outputScale;
220 - fp2Vals[i].y *= outputScale;
221 - }
222 -
223 - // Convert to e2m1 values.
224 - uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals);
225 -
226 - // Write the e2m1 values to global memory.
227 - return e2m1Vec;
228 -#else
229 - return 0;
230 -#endif
231 -}
30 +namespace vllm {
31
32 // Use UE4M3 by default.
33 template <class Type, bool UE8M0_SF = false, bool SMALL_NUM_EXPERTS = false>
235 -__global__ void
236 -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
237 -__launch_bounds__(512, 4) cvt_fp16_to_fp4(
238 -#else
239 -cvt_fp16_to_fp4(
240 -#endif
241 - int32_t numRows, int32_t numCols, Type const* in, float const* SFScale,
242 - uint32_t* out, uint32_t* SFout, uint32_t* input_offset_by_experts,
243 - uint32_t* output_scale_offset_by_experts, int n_experts, bool low_latency) {
244 -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
34 +__global__ void __launch_bounds__(512, 4)
35 + cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
36 + float const* SFScale, uint32_t* out, uint32_t* SFout,
37 + uint32_t* input_offset_by_experts,
38 + uint32_t* output_scale_offset_by_experts, int n_experts,
39 + bool low_latency) {
40 using PackedVec = PackedVec<Type>;
41 static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
42 (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
@@ -299,8 +94,8 @@ cvt_fp16_to_fp4(
94 &input_offset_by_experts[chunk_start + 12]));
95 local_offsets[16] = __ldca(&input_offset_by_experts[chunk_start + 16]);
96
302 - // Check against the 16 loaded offsets
303 - #pragma unroll
97 +// Check against the 16 loaded offsets
98 +#pragma unroll
99 for (int i = 0; i < 16; i++) {
100 if (rowIdx >= local_offsets[i] && rowIdx < local_offsets[i + 1]) {
101 rowIdx_in_expert = rowIdx - local_offsets[i];
@@ -330,21 +125,15 @@ cvt_fp16_to_fp4(
125
126 out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
127 }
333 -#endif
128 }
129
130 // Kernel for LARGE_M_TOPK = true (large m_topk optimized version)
131 template <class Type, bool UE8M0_SF = false, bool SMALL_NUM_EXPERTS = false>
338 -__global__ void
339 -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
340 -__launch_bounds__(1024, 4) cvt_fp16_to_fp4(
341 -#else
342 -cvt_fp16_to_fp4(
343 -#endif
344 - int32_t numRows, int32_t numCols, Type const* in, float const* SFScale,
345 - uint32_t* out, uint32_t* SFout, uint32_t* input_offset_by_experts,
346 - uint32_t* output_scale_offset_by_experts, int n_experts) {
347 -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
132 +__global__ void __launch_bounds__(1024, 4)
133 + cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
134 + float const* SFScale, uint32_t* out, uint32_t* SFout,
135 + uint32_t* input_offset_by_experts,
136 + uint32_t* output_scale_offset_by_experts, int n_experts) {
137 using PackedVec = PackedVec<Type>;
138 static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
139 (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
@@ -425,7 +214,6 @@ cvt_fp16_to_fp4(
214
215 out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
216 }
428 -#endif
217 }
218
219 template <typename T>
@@ -501,6 +289,8 @@ void quant_impl(void* output, void* output_scale, void* input,
289 }
290 }
291
292 +} // namespace vllm
293 +
294 /*Quantization entry for fp4 experts quantization*/
295 #define CHECK_TH_CUDA(x, m) TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
296 #define CHECK_CONTIGUOUS(x, m) \
@@ -560,23 +350,17 @@ void scaled_fp4_experts_quant_sm100a(
350 // 4 means 4 fp8 values are packed into one int32
351 TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
352
563 - auto in_dtype = input.dtype();
353 const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
354 const cudaStream_t stream =
355 at::cuda::getCurrentCUDAStream(input.get_device());
567 - if (in_dtype == at::ScalarType::Half) {
568 - quant_impl<half>(output.data_ptr(), output_scale.data_ptr(),
569 - input.data_ptr(), input_global_scale.data_ptr(),
570 - input_offset_by_experts.data_ptr(),
571 - output_scale_offset_by_experts.data_ptr(), m_topk, k,
572 - n_experts, stream);
573 - } else if (in_dtype == at::ScalarType::BFloat16) {
574 - quant_impl<__nv_bfloat16>(output.data_ptr(), output_scale.data_ptr(),
575 - input.data_ptr(), input_global_scale.data_ptr(),
576 - input_offset_by_experts.data_ptr(),
577 - output_scale_offset_by_experts.data_ptr(), m_topk,
578 - k, n_experts, stream);
579 - } else {
580 - TORCH_CHECK(false, "Expected input data type to be half or bfloat16");
581 - }
356 +
357 + VLLM_DISPATCH_HALF_TYPES(
358 + input.scalar_type(), "nvfp4_experts_quant_kernel", [&] {
359 + using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
360 + vllm::quant_impl<cuda_type>(
361 + output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
362 + input_global_scale.data_ptr(), input_offset_by_experts.data_ptr(),
363 + output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
364 + stream);
365 + });
366 }
csrc/quantization/fp4/nvfp4_quant_entry.cu
+18
@@ -32,6 +32,14 @@ void scaled_fp4_experts_quant_sm100a(
32 torch::Tensor const& output_scale_offset_by_experts);
33 #endif
34
35 +#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
36 + (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
37 +void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output,
38 + torch::Tensor& output_sf,
39 + torch::Tensor& input,
40 + torch::Tensor& input_sf);
41 +#endif
42 +
43 void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
44 torch::Tensor& output_sf, torch::Tensor const& input_sf) {
45 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
@@ -54,3 +62,13 @@ void scaled_fp4_experts_quant(
62 TORCH_CHECK_NOT_IMPLEMENTED(false,
63 "No compiled nvfp4 experts quantization kernel");
64 }
65 +
66 +void silu_and_mul_nvfp4_quant(torch::Tensor& output, torch::Tensor& output_sf,
67 + torch::Tensor& input, torch::Tensor& input_sf) {
68 +#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
69 + (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
70 + return silu_and_mul_nvfp4_quant_sm1xxa(output, output_sf, input, input_sf);
71 +#endif
72 + TORCH_CHECK_NOT_IMPLEMENTED(
73 + false, "No compiled silu_and_mul nvfp4 quantization kernel");
74 +}
csrc/quantization/fp4/nvfp4_quant_kernels.cu
+17 -254
@@ -23,245 +23,18 @@
23 #include <c10/cuda/CUDAGuard.h>
24
25 #include <cuda_fp8.h>
26 +#include "dispatch_utils.h"
27
28 #include "cuda_utils.h"
29 +#include "nvfp4_utils.cuh"
30
29 -// Get type2 from type or vice versa (applied to half and bfloat16)
30 -template <typename T>
31 -struct TypeConverter {
32 - using Type = half2;
33 -}; // keep for generality
34 -
35 -template <>
36 -struct TypeConverter<half2> {
37 - using Type = half;
38 -};
39 -
40 -template <>
41 -struct TypeConverter<half> {
42 - using Type = half2;
43 -};
44 -
45 -template <>
46 -struct TypeConverter<__nv_bfloat162> {
47 - using Type = __nv_bfloat16;
48 -};
49 -
50 -template <>
51 -struct TypeConverter<__nv_bfloat16> {
52 - using Type = __nv_bfloat162;
53 -};
54 -
55 -#define ELTS_PER_THREAD 8
56 -
57 -constexpr int CVT_FP4_ELTS_PER_THREAD = 8;
58 -constexpr int CVT_FP4_SF_VEC_SIZE = 16;
59 -
60 -// Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t).
61 -inline __device__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) {
62 -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
63 - uint32_t val;
64 - asm volatile(
65 - "{\n"
66 - ".reg .b8 byte0;\n"
67 - ".reg .b8 byte1;\n"
68 - ".reg .b8 byte2;\n"
69 - ".reg .b8 byte3;\n"
70 - "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n"
71 - "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n"
72 - "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
73 - "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
74 - "mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
75 - "}"
76 - : "=r"(val)
77 - : "f"(array[0]), "f"(array[1]), "f"(array[2]), "f"(array[3]),
78 - "f"(array[4]), "f"(array[5]), "f"(array[6]), "f"(array[7]));
79 - return val;
80 -#else
81 - return 0;
82 -#endif
83 -}
84 -
85 -// Convert 4 float2 values into 8 e2m1 values (represented as one uint32_t).
86 -inline __device__ uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) {
87 -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
88 - uint32_t val;
89 - asm volatile(
90 - "{\n"
91 - ".reg .b8 byte0;\n"
92 - ".reg .b8 byte1;\n"
93 - ".reg .b8 byte2;\n"
94 - ".reg .b8 byte3;\n"
95 - "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n"
96 - "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n"
97 - "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
98 - "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
99 - "mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
100 - "}"
101 - : "=r"(val)
102 - : "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y),
103 - "f"(array[2].x), "f"(array[2].y), "f"(array[3].x), "f"(array[3].y));
104 - return val;
105 -#else
106 - return 0;
107 -#endif
108 -}
109 -
110 -// Fast reciprocal.
111 -inline __device__ float reciprocal_approximate_ftz(float a) {
112 - float b;
113 - asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a));
114 - return b;
115 -}
116 -
117 -template <class SFType, int CVT_FP4_NUM_THREADS_PER_SF>
118 -__device__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset(int rowIdx, int colIdx,
119 - int numCols,
120 - SFType* SFout) {
121 -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
122 - static_assert(CVT_FP4_NUM_THREADS_PER_SF == 1 ||
123 - CVT_FP4_NUM_THREADS_PER_SF == 2);
124 -
125 - // One pair of threads write one SF to global memory.
126 - // TODO: stage through smem for packed STG.32
127 - // is it better than STG.8 from 4 threads ?
128 - if (threadIdx.x % CVT_FP4_NUM_THREADS_PER_SF == 0) {
129 - // SF vector index (16 elements share one SF in the K dimension).
130 - int32_t kIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF;
131 - int32_t mIdx = rowIdx;
132 -
133 - // SF layout [numMTiles, numKTiles, 32 (mTile), 4 (mTile), 4(kTile)]
134 - // --> index [mTileIdx, kTileIdx, outerMIdx, innerMIdx, innerKIdx]
135 -
136 - int32_t mTileIdx = mIdx / (32 * 4);
137 - // SF vector size 16.
138 - int factor = CVT_FP4_SF_VEC_SIZE * 4;
139 - int32_t numKTiles = (numCols + factor - 1) / factor;
140 - int64_t mTileStride = numKTiles * 32 * 4 * 4;
141 -
142 - int32_t kTileIdx = (kIdx / 4);
143 - int64_t kTileStride = 32 * 4 * 4;
144 -
145 - // M tile layout [32, 4] is column-major.
146 - int32_t outerMIdx = (mIdx % 32);
147 - int64_t outerMStride = 4 * 4;
148 -
149 - int32_t innerMIdx = (mIdx % (32 * 4)) / 32;
150 - int64_t innerMStride = 4;
151 -
152 - int32_t innerKIdx = (kIdx % 4);
153 - int64_t innerKStride = 1;
154 -
155 - // Compute the global offset.
156 - int64_t SFOffset = mTileIdx * mTileStride + kTileIdx * kTileStride +
157 - outerMIdx * outerMStride + innerMIdx * innerMStride +
158 - innerKIdx * innerKStride;
159 -
160 - return reinterpret_cast<uint8_t*>(SFout) + SFOffset;
161 - }
162 -#endif
163 - return nullptr;
164 -}
165 -
166 -// Define a 16 bytes packed data type.
167 -template <class Type>
168 -struct PackedVec {
169 - typename TypeConverter<Type>::Type elts[4];
170 -};
171 -
172 -template <>
173 -struct PackedVec<__nv_fp8_e4m3> {
174 - __nv_fp8x2_e4m3 elts[8];
175 -};
176 -
177 -// Quantizes the provided PackedVec into the uint32_t output
178 -template <class Type, bool UE8M0_SF = false>
179 -__device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal,
180 - uint8_t* SFout) {
181 -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
182 - // Get absolute maximum values among the local 8 values.
183 - auto localMax = __habs2(vec.elts[0]);
184 -
185 - // Local maximum value.
186 - #pragma unroll
187 - for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
188 - localMax = __hmax2(localMax, __habs2(vec.elts[i]));
189 - }
190 -
191 - // Get the absolute maximum among all 16 values (two threads).
192 - localMax = __hmax2(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax);
193 - // Get the final absolute maximum values.
194 - float vecMax = float(__hmax(localMax.x, localMax.y));
195 -
196 - // Get the SF (max value of the vector / max value of e2m1).
197 - // maximum value of e2m1 = 6.0.
198 - // TODO: use half as compute data type.
199 - float SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f));
200 - // 8 bits representation of the SF.
201 - uint8_t fp8SFVal;
202 - // Write the SF to global memory (STG.8).
203 - if constexpr (UE8M0_SF) {
204 - // Extract the 8 exponent bits from float32.
205 - // float 32bits = 1 sign bit + 8 exponent bits + 23 mantissa bits.
206 - uint32_t tmp = reinterpret_cast<uint32_t&>(SFValue) >> 23;
207 - fp8SFVal = tmp & 0xff;
208 - // Convert back to fp32.
209 - reinterpret_cast<uint32_t&>(SFValue) = tmp << 23;
210 - } else {
211 - // Here SFValue is always positive, so E4M3 is the same as UE4M3.
212 - __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue);
213 - reinterpret_cast<__nv_fp8_e4m3&>(fp8SFVal) = tmp;
214 - // Convert back to fp32.
215 - SFValue = float(tmp);
216 - }
217 - // Get the output scale.
218 - // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) *
219 - // reciprocal(SFScaleVal))
220 - float outputScale =
221 - SFValue != 0 ? reciprocal_approximate_ftz(
222 - SFValue * reciprocal_approximate_ftz(SFScaleVal))
223 - : 0.0f;
224 -
225 - if (SFout) {
226 - // Write the SF to global memory (STG.8).
227 - *SFout = fp8SFVal;
228 - }
229 -
230 - // Convert the input to float.
231 - float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
232 -
233 - #pragma unroll
234 - for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
235 - if constexpr (std::is_same_v<Type, half>) {
236 - fp2Vals[i] = __half22float2(vec.elts[i]);
237 - } else {
238 - fp2Vals[i] = __bfloat1622float2(vec.elts[i]);
239 - }
240 - fp2Vals[i].x *= outputScale;
241 - fp2Vals[i].y *= outputScale;
242 - }
243 -
244 - // Convert to e2m1 values.
245 - uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals);
246 -
247 - // Write the e2m1 values to global memory.
248 - return e2m1Vec;
249 -#else
250 - return 0;
251 -#endif
252 -}
31 +namespace vllm {
32
33 // Use UE4M3 by default.
34 template <class Type, bool UE8M0_SF = false>
256 -__global__ void
257 -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
258 -__launch_bounds__(512, 4) cvt_fp16_to_fp4(
259 -#else
260 -cvt_fp16_to_fp4(
261 -#endif
262 - int32_t numRows, int32_t numCols, Type const* in, float const* SFScale,
263 - uint32_t* out, uint32_t* SFout) {
264 -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
35 +__global__ void __launch_bounds__(512, 4)
36 + cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
37 + float const* SFScale, uint32_t* out, uint32_t* SFout) {
38 using PackedVec = PackedVec<Type>;
39 static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
40 (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
@@ -293,7 +66,6 @@ cvt_fp16_to_fp4(
66 cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
67 }
68 }
296 -#endif
69 }
70
71 template <typename T>
@@ -332,6 +104,8 @@ template void invokeFP4Quantization(int m, int n, __nv_bfloat16 const* input,
104 int multiProcessorCount,
105 cudaStream_t stream);
106
107 +} // namespace vllm
108 +
109 void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
110 torch::Tensor const& input,
111 torch::Tensor const& output_sf,
@@ -340,6 +114,9 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
114 int32_t n = input.size(1);
115
116 TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");
117 + TORCH_CHECK(input.scalar_type() == at::ScalarType::Half ||
118 + input.scalar_type() == at::ScalarType::BFloat16,
119 + "Unsupported input data type for quantize_to_fp4.");
120
121 int multiProcessorCount =
122 get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
@@ -353,24 +130,10 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
130 // We don't support e8m0 scales at this moment.
131 bool useUE8M0 = false;
132
356 - switch (input.scalar_type()) {
357 - case torch::kHalf: {
358 - auto input_ptr = reinterpret_cast<half const*>(input.data_ptr());
359 - invokeFP4Quantization(m, n, input_ptr, input_sf_ptr, output_ptr, sf_out,
360 - useUE8M0, multiProcessorCount, stream);
361 - break;
362 - }
363 - case torch::kBFloat16: {
364 - auto input_ptr = reinterpret_cast<__nv_bfloat16 const*>(input.data_ptr());
365 - invokeFP4Quantization(m, n, input_ptr, input_sf_ptr, output_ptr, sf_out,
366 - useUE8M0, multiProcessorCount, stream);
367 - break;
368 - }
369 - default: {
370 - std::cerr << "Observing: " << input.scalar_type()
371 - << " for the input datatype which is invalid";
372 - throw std::runtime_error(
373 - "Unsupported input data type for quantize_to_fp4.");
374 - }
375 - }
133 + VLLM_DISPATCH_HALF_TYPES(input.scalar_type(), "nvfp4_quant_kernel", [&] {
134 + using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
135 + auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
136 + vllm::invokeFP4Quantization(m, n, input_ptr, input_sf_ptr, output_ptr,
137 + sf_out, useUE8M0, multiProcessorCount, stream);
138 + });
139 }
csrc/quantization/fp4/nvfp4_utils.cuh new
+251
@@ -0,0 +1,251 @@
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 +#pragma once
18 +
19 +#include <cuda_runtime.h>
20 +#include <cuda_fp8.h>
21 +
22 +#define ELTS_PER_THREAD 8
23 +
24 +constexpr int CVT_FP4_ELTS_PER_THREAD = 8;
25 +constexpr int CVT_FP4_SF_VEC_SIZE = 16;
26 +
27 +namespace vllm {
28 +
29 +// Convert PyTorch cpp type to CUDA type
30 +template <typename T>
31 +struct CUDATypeConverter {
32 + using Type = T;
33 +};
34 +
35 +template <>
36 +struct CUDATypeConverter<at::Half> {
37 + using Type = half;
38 +};
39 +
40 +template <>
41 +struct CUDATypeConverter<at::BFloat16> {
42 + using Type = __nv_bfloat16;
43 +};
44 +
45 +// Get type2 from type or vice versa (applied to half and bfloat16)
46 +template <typename T>
47 +struct TypeConverter {
48 + using Type = half2;
49 +}; // keep for generality
50 +
51 +template <>
52 +struct TypeConverter<half2> {
53 + using Type = half;
54 +};
55 +
56 +template <>
57 +struct TypeConverter<half> {
58 + using Type = half2;
59 +};
60 +
61 +template <>
62 +struct TypeConverter<__nv_bfloat162> {
63 + using Type = __nv_bfloat16;
64 +};
65 +
66 +template <>
67 +struct TypeConverter<__nv_bfloat16> {
68 + using Type = __nv_bfloat162;
69 +};
70 +
71 +// Define a 16 bytes packed data type.
72 +template <class Type>
73 +struct PackedVec {
74 + typename TypeConverter<Type>::Type elts[4];
75 +};
76 +
77 +template <>
78 +struct PackedVec<__nv_fp8_e4m3> {
79 + __nv_fp8x2_e4m3 elts[8];
80 +};
81 +
82 +// Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t).
83 +inline __device__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) {
84 + uint32_t val;
85 + asm volatile(
86 + "{\n"
87 + ".reg .b8 byte0;\n"
88 + ".reg .b8 byte1;\n"
89 + ".reg .b8 byte2;\n"
90 + ".reg .b8 byte3;\n"
91 + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n"
92 + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n"
93 + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
94 + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
95 + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
96 + "}"
97 + : "=r"(val)
98 + : "f"(array[0]), "f"(array[1]), "f"(array[2]), "f"(array[3]),
99 + "f"(array[4]), "f"(array[5]), "f"(array[6]), "f"(array[7]));
100 + return val;
101 +}
102 +
103 +// Convert 4 float2 values into 8 e2m1 values (represented as one uint32_t).
104 +inline __device__ uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) {
105 + uint32_t val;
106 + asm volatile(
107 + "{\n"
108 + ".reg .b8 byte0;\n"
109 + ".reg .b8 byte1;\n"
110 + ".reg .b8 byte2;\n"
111 + ".reg .b8 byte3;\n"
112 + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n"
113 + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n"
114 + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
115 + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
116 + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
117 + "}"
118 + : "=r"(val)
119 + : "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y),
120 + "f"(array[2].x), "f"(array[2].y), "f"(array[3].x), "f"(array[3].y));
121 + return val;
122 +}
123 +
124 +// Fast reciprocal.
125 +inline __device__ float reciprocal_approximate_ftz(float a) {
126 + float b;
127 + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a));
128 + return b;
129 +}
130 +
131 +template <class SFType, int CVT_FP4_NUM_THREADS_PER_SF>
132 +__device__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset(int rowIdx, int colIdx,
133 + int numCols,
134 + SFType* SFout) {
135 + static_assert(CVT_FP4_NUM_THREADS_PER_SF == 1 ||
136 + CVT_FP4_NUM_THREADS_PER_SF == 2);
137 +
138 + // One pair of threads write one SF to global memory.
139 + // TODO: stage through smem for packed STG.32
140 + // is it better than STG.8 from 4 threads ?
141 + if (threadIdx.x % CVT_FP4_NUM_THREADS_PER_SF == 0) {
142 + // SF vector index (16 elements share one SF in the K dimension).
143 + int32_t kIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF;
144 + int32_t mIdx = rowIdx;
145 +
146 + // SF layout [numMTiles, numKTiles, 32 (mTile), 4 (mTile), 4(kTile)]
147 + // --> index [mTileIdx, kTileIdx, outerMIdx, innerMIdx, innerKIdx]
148 +
149 + int32_t mTileIdx = mIdx / (32 * 4);
150 + // SF vector size 16.
151 + int factor = CVT_FP4_SF_VEC_SIZE * 4;
152 + int32_t numKTiles = (numCols + factor - 1) / factor;
153 + int64_t mTileStride = numKTiles * 32 * 4 * 4;
154 +
155 + int32_t kTileIdx = (kIdx / 4);
156 + int64_t kTileStride = 32 * 4 * 4;
157 +
158 + // M tile layout [32, 4] is column-major.
159 + int32_t outerMIdx = (mIdx % 32);
160 + int64_t outerMStride = 4 * 4;
161 +
162 + int32_t innerMIdx = (mIdx % (32 * 4)) / 32;
163 + int64_t innerMStride = 4;
164 +
165 + int32_t innerKIdx = (kIdx % 4);
166 + int64_t innerKStride = 1;
167 +
168 + // Compute the global offset.
169 + int64_t SFOffset = mTileIdx * mTileStride + kTileIdx * kTileStride +
170 + outerMIdx * outerMStride + innerMIdx * innerMStride +
171 + innerKIdx * innerKStride;
172 +
173 + return reinterpret_cast<uint8_t*>(SFout) + SFOffset;
174 + }
175 + return nullptr;
176 +}
177 +
178 +// Quantizes the provided PackedVec into the uint32_t output
179 +template <class Type, bool UE8M0_SF = false>
180 +__device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal,
181 + uint8_t* SFout) {
182 + // Get absolute maximum values among the local 8 values.
183 + auto localMax = __habs2(vec.elts[0]);
184 +
185 +// Local maximum value.
186 +#pragma unroll
187 + for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
188 + localMax = __hmax2(localMax, __habs2(vec.elts[i]));
189 + }
190 +
191 + // Get the absolute maximum among all 16 values (two threads).
192 + localMax = __hmax2(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax);
193 + // Get the final absolute maximum values.
194 + float vecMax = float(__hmax(localMax.x, localMax.y));
195 +
196 + // Get the SF (max value of the vector / max value of e2m1).
197 + // maximum value of e2m1 = 6.0.
198 + // TODO: use half as compute data type.
199 + float SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f));
200 + // 8 bits representation of the SF.
201 + uint8_t fp8SFVal;
202 + // Write the SF to global memory (STG.8).
203 + if constexpr (UE8M0_SF) {
204 + // Extract the 8 exponent bits from float32.
205 + // float 32bits = 1 sign bit + 8 exponent bits + 23 mantissa bits.
206 + uint32_t tmp = reinterpret_cast<uint32_t&>(SFValue) >> 23;
207 + fp8SFVal = tmp & 0xff;
208 + // Convert back to fp32.
209 + reinterpret_cast<uint32_t&>(SFValue) = tmp << 23;
210 + } else {
211 + // Here SFValue is always positive, so E4M3 is the same as UE4M3.
212 + __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue);
213 + reinterpret_cast<__nv_fp8_e4m3&>(fp8SFVal) = tmp;
214 + // Convert back to fp32.
215 + SFValue = float(tmp);
216 + }
217 + // Get the output scale.
218 + // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) *
219 + // reciprocal(SFScaleVal))
220 + float outputScale =
221 + SFValue != 0 ? reciprocal_approximate_ftz(
222 + SFValue * reciprocal_approximate_ftz(SFScaleVal))
223 + : 0.0f;
224 +
225 + if (SFout) {
226 + // Write the SF to global memory (STG.8).
227 + *SFout = fp8SFVal;
228 + }
229 +
230 + // Convert the input to float.
231 + float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
232 +
233 +#pragma unroll
234 + for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
235 + if constexpr (std::is_same_v<Type, half>) {
236 + fp2Vals[i] = __half22float2(vec.elts[i]);
237 + } else {
238 + fp2Vals[i] = __bfloat1622float2(vec.elts[i]);
239 + }
240 + fp2Vals[i].x *= outputScale;
241 + fp2Vals[i].y *= outputScale;
242 + }
243 +
244 + // Convert to e2m1 values.
245 + uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals);
246 +
247 + // Write the e2m1 values to global memory.
248 + return e2m1Vec;
249 +}
250 +
251 +} // namespace vllm
csrc/torch_bindings.cpp
+1 -2
@@ -115,8 +115,7 @@ 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 -#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
119 - (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
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) -> ()");
tests/kernels/quantization/test_silu_nvfp4_quant_fusion.py
+1 -2
@@ -8,8 +8,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 - and hasattr(torch.ops._C, "silu_and_mul_nvfp4_quant")):
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