@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / fcb9df99bd
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 9, 16:00 UTC

[Perf][Kernel] Optimize FP4 quantization kernels (SM100F) (#32520)

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>

Roberto L. Castro committed Jan 25, 2026 at 02:45 UTC fcb9df99bd7d0e532bcf2891db2a85bd927605fe
18 files changed +509 -152
benchmarks/kernels/bench_nvfp4_quant.py
+55 -22
@@ -20,8 +20,12 @@ FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
20 FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
21
22 PROVIDER_CFGS = {
23 - "vllm": dict(backend="vllm", enabled=True),
24 - "flashinfer": dict(backend="flashinfer", enabled=True),
23 + "vllm": dict(backend="vllm", is_sf_swizzled_layout=False, enabled=True),
24 + "vllm-swizzle": dict(backend="vllm", is_sf_swizzled_layout=True, enabled=True),
25 + "flashinfer": dict(backend="flashinfer", is_sf_swizzled_layout=False, enabled=True),
26 + "flashinfer-swizzle": dict(
27 + backend="flashinfer", is_sf_swizzled_layout=True, enabled=True
28 + ),
29 }
30
31 _enabled = [k for k, v in PROVIDER_CFGS.items() if v["enabled"]]
@@ -36,7 +40,7 @@ def compute_global_scale(tensor: torch.Tensor) -> torch.Tensor:
40 @triton.testing.perf_report(
41 triton.testing.Benchmark(
42 x_names=["batch_size"],
39 - x_vals=[1, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096],
43 + x_vals=[1, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192],
44 x_log=False,
45 line_arg="provider",
46 line_vals=_enabled,
@@ -63,19 +67,36 @@ def benchmark(batch_size, provider, N, K):
67
68 if cfg["backend"] == "vllm":
69 # vLLM's FP4 quantization
66 - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
67 - lambda: ops.scaled_fp4_quant(a, a_global_scale),
68 - quantiles=quantiles,
69 - )
70 + if cfg["is_sf_swizzled_layout"]:
71 + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
72 + lambda: ops.scaled_fp4_quant(
73 + a, a_global_scale, is_sf_swizzled_layout=True
74 + ),
75 + quantiles=quantiles,
76 + )
77 + else:
78 + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
79 + lambda: ops.scaled_fp4_quant(
80 + a, a_global_scale, is_sf_swizzled_layout=False
81 + ),
82 + quantiles=quantiles,
83 + )
84 elif cfg["backend"] == "flashinfer":
85 # FlashInfer's FP4 quantization
72 - # Use is_sf_swizzled_layout=True to match vLLM's output format
73 - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
74 - lambda: flashinfer_fp4_quantize(
75 - a, a_global_scale, is_sf_swizzled_layout=True
76 - ),
77 - quantiles=quantiles,
78 - )
86 + if cfg["is_sf_swizzled_layout"]:
87 + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
88 + lambda: flashinfer_fp4_quantize(
89 + a, a_global_scale, is_sf_swizzled_layout=True
90 + ),
91 + quantiles=quantiles,
92 + )
93 + else:
94 + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
95 + lambda: flashinfer_fp4_quantize(
96 + a, a_global_scale, is_sf_swizzled_layout=False
97 + ),
98 + quantiles=quantiles,
99 + )
100
101 # Convert ms to us for better readability at small batch sizes
102 to_us = lambda t_ms: t_ms * 1000
@@ -92,7 +113,9 @@ def prepare_shapes(args):
113 return out
114
115
95 -def _test_accuracy_once(M: int, K: int, dtype: torch.dtype, device: str):
116 +def _test_accuracy_once(
117 + M: int, K: int, dtype: torch.dtype, device: str, is_sf_swizzled_layout: bool
118 +):
119 """Test accuracy between vLLM and FlashInfer FP4 quantization."""
120 # Create input tensor
121 a = torch.randn((M, K), device=device, dtype=dtype)
@@ -101,11 +124,13 @@ def _test_accuracy_once(M: int, K: int, dtype: torch.dtype, device: str):
124 a_global_scale = compute_global_scale(a)
125
126 # vLLM quantization
104 - vllm_fp4, vllm_scale = ops.scaled_fp4_quant(a, a_global_scale)
127 + vllm_fp4, vllm_scale = ops.scaled_fp4_quant(
128 + a, a_global_scale, is_sf_swizzled_layout=is_sf_swizzled_layout
129 + )
130
131 # FlashInfer quantization (with swizzled layout to match vLLM's output)
132 flashinfer_fp4, flashinfer_scale = flashinfer_fp4_quantize(
108 - a, a_global_scale, is_sf_swizzled_layout=True
133 + a, a_global_scale, is_sf_swizzled_layout=is_sf_swizzled_layout
134 )
135 flashinfer_scale = flashinfer_scale.view(torch.float8_e4m3fn)
136
@@ -114,7 +139,14 @@ def _test_accuracy_once(M: int, K: int, dtype: torch.dtype, device: str):
139 vllm_fp4,
140 flashinfer_fp4,
141 )
117 - print(f"M={M}, K={K}, dtype={dtype}: PASSED")
142 + # Compare scales
143 + torch.testing.assert_close(
144 + vllm_scale,
145 + flashinfer_scale,
146 + )
147 + print(
148 + f"M={M}, K={K}, dtype={dtype}, is_sf_swizzled_layout={is_sf_swizzled_layout}: PASSED" # noqa: E501
149 + )
150
151
152 def test_accuracy():
@@ -130,9 +162,10 @@ def test_accuracy():
162 Ms = [1, 1024]
163 Ks = [4096]
164
133 - for M in Ms:
134 - for K in Ks:
135 - _test_accuracy_once(M, K, dtype, device)
165 + for is_sf_swizzled_layout in [True, False]:
166 + for M in Ms:
167 + for K in Ks:
168 + _test_accuracy_once(M, K, dtype, device, is_sf_swizzled_layout)
169
170 print("\nAll accuracy tests passed!")
171
@@ -145,7 +178,7 @@ if __name__ == "__main__":
178 "--models",
179 nargs="+",
180 type=str,
148 - default=["meta-llama/Llama-3.1-8B-Instruct"],
181 + default=["meta-llama/Llama-3.3-70B-Instruct"],
182 choices=list(WEIGHT_SHAPES.keys()),
183 )
184 parser.add_argument("--tp-sizes", nargs="+", type=int, default=[1])
csrc/ops.h
+2 -1
@@ -293,7 +293,8 @@ std::vector<torch::Tensor> cutlass_sparse_compress(torch::Tensor const& a);
293
294 void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
295 torch::Tensor& output_scale,
296 - torch::Tensor const& input_scale);
296 + torch::Tensor const& input_scale,
297 + bool is_sf_swizzled_layout);
298
299 void scaled_fp4_experts_quant(
300 torch::Tensor& output, torch::Tensor& output_scale,
csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu
+59 -20
@@ -27,17 +27,24 @@
27
28 #include "cuda_utils.h"
29 #include "launch_bounds_utils.h"
30 +
31 +// Define before including nvfp4_utils.cuh so the header
32 +// can use this macro during compilation.
33 +#define NVFP4_ENABLE_ELTS16 1
34 #include "nvfp4_utils.cuh"
35
36 namespace vllm {
37
38 // Use UE4M3 by default.
39 template <class Type, bool UE8M0_SF = false>
36 -__global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
37 - silu_mul_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
38 - float const* SFScale, uint32_t* out,
39 - uint32_t* SFout) {
40 - using PackedVec = PackedVec<Type>;
40 +__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
41 + silu_mul_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols,
42 + int32_t num_padded_cols,
43 + Type const* __restrict__ in,
44 + float const* __restrict__ SFScale,
45 + uint32_t* __restrict__ out,
46 + uint32_t* __restrict__ SFout) {
47 + using PackedVec = vllm::PackedVec<Type>;
48 static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
49 (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
50 static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
@@ -49,34 +56,60 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
56 // Get the global scaling factor, which will be applied to the SF.
57 // Note SFScale is the same as next GEMM's alpha, which is
58 // (448.f / (Alpha_A / 6.f)).
52 - float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[0];
59 + float const SFScaleVal = (SFScale == nullptr) ? 1.0f : SFScale[0];
60 +
61 + int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x;
62 + int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD;
63
64 // Input tensor row/col loops.
65 for (int rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) {
56 - for (int colIdx = threadIdx.x; colIdx < numCols / CVT_FP4_ELTS_PER_THREAD;
57 - colIdx += blockDim.x) {
66 + if (colIdx < num_padded_cols) {
67 + PackedVec in_vec;
68 + PackedVec in_vec2;
69 int64_t inOffset =
70 rowIdx * (numCols * 2 / CVT_FP4_ELTS_PER_THREAD) + colIdx;
71 int64_t inOffset2 = rowIdx * (numCols * 2 / CVT_FP4_ELTS_PER_THREAD) +
72 numCols / CVT_FP4_ELTS_PER_THREAD + colIdx;
62 - PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
63 - PackedVec in_vec2 = reinterpret_cast<PackedVec const*>(in)[inOffset2];
73
65 - // Get the output tensor offset.
66 - // Same as inOffset because 8 elements are packed into one uint32_t.
67 - int64_t outOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
68 - auto& out_pos = out[outOffset];
74 + bool valid = (rowIdx < numRows) && (elem_idx < numCols);
75 + if constexpr (CVT_FP4_PACK16) {
76 + ld256_or_zero_cg_u32<Type>(
77 + in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 8],
78 + valid);
79 + ld256_or_zero_cg_u32<Type>(
80 + in_vec2, &reinterpret_cast<const uint32_t*>(in)[inOffset2 * 8],
81 + valid);
82 + } else {
83 + ld128_or_zero_cg_u32<Type>(
84 + in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 4],
85 + valid);
86 + ld128_or_zero_cg_u32<Type>(
87 + in_vec2, &reinterpret_cast<const uint32_t*>(in)[inOffset2 * 4],
88 + valid);
89 + }
90
91 // Compute silu and mul
71 - PackedVec out_silu_mul = compute_silu_mul(in_vec, in_vec2);
92 + PackedVec out_silu_mul = compute_silu_mul<Type>(in_vec, in_vec2);
93
94 auto sf_out =
95 cvt_quant_to_fp4_get_sf_out_offset<uint32_t,
96 CVT_FP4_NUM_THREADS_PER_SF>(
97 rowIdx, colIdx, numKTiles, SFout);
98
78 - out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(out_silu_mul, SFScaleVal,
79 - sf_out);
99 + auto out_val =
100 + cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
101 + out_silu_mul, SFScaleVal, sf_out);
102 +
103 + if (valid) {
104 + if constexpr (CVT_FP4_PACK16) {
105 + int64_t outOffset = rowIdx * (numCols / 8) + colIdx * 2;
106 + uint64_t packed64 =
107 + (uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo);
108 + reinterpret_cast<uint64_t*>(out)[outOffset >> 1] = packed64;
109 + } else {
110 + out[inOffset] = out_val;
111 + }
112 + }
113 }
114 }
115 }
@@ -103,17 +136,23 @@ void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output, // [..., d]
136 auto output_ptr = static_cast<int64_t*>(output.data_ptr());
137 const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
138 auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
106 - dim3 block(std::min(int(n / ELTS_PER_THREAD), 1024));
139 + dim3 block(std::min(int(n / ELTS_PER_THREAD), 512));
140 int const numBlocksPerSM =
141 vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
109 - dim3 grid(std::min(int(m), multiProcessorCount * numBlocksPerSM));
142 +
143 + int sf_n_unpadded = int(n / CVT_FP4_SF_VEC_SIZE);
144 +
145 + int grid_y = vllm::div_round_up(sf_n_unpadded, static_cast<int>(block.x));
146 + int grid_x = std::min(
147 + int(m), std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
148 + dim3 grid(grid_x, grid_y);
149
150 VLLM_DISPATCH_HALF_TYPES(
151 input.scalar_type(), "silu_and_mul_nvfp4_quant_kernel", [&] {
152 using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
153 auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
154 vllm::silu_mul_cvt_fp16_to_fp4<cuda_type><<<grid, block, 0, stream>>>(
116 - m, n, input_ptr, input_sf_ptr,
155 + m, n, sf_n_unpadded, input_ptr, input_sf_ptr,
156 reinterpret_cast<uint32_t*>(output_ptr),
157 reinterpret_cast<uint32_t*>(sf_out));
158 });
csrc/quantization/fp4/nvfp4_experts_quant.cu
+4 -4
@@ -140,8 +140,8 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
140 CVT_FP4_NUM_THREADS_PER_SF>(
141 rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
142
143 - out_pos =
144 - cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(quant_input, SFScaleVal, sf_out);
143 + out_pos = cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
144 + quant_input, SFScaleVal, sf_out);
145 }
146 }
147
@@ -246,8 +246,8 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
246 CVT_FP4_NUM_THREADS_PER_SF>(
247 rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
248
249 - out_pos =
250 - cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(quant_input, SFScaleVal, sf_out);
249 + out_pos = cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
250 + quant_input, SFScaleVal, sf_out);
251 }
252 }
253
csrc/quantization/fp4/nvfp4_quant_entry.cu
+6 -3
@@ -21,7 +21,8 @@
21 void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
22 torch::Tensor const& input,
23 torch::Tensor const& output_sf,
24 - torch::Tensor const& input_sf);
24 + torch::Tensor const& input_sf,
25 + bool is_sf_swizzled_layout);
26 #endif
27
28 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
@@ -51,10 +52,12 @@ void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
52 #endif
53
54 void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
54 - torch::Tensor& output_sf, torch::Tensor const& input_sf) {
55 + torch::Tensor& output_sf, torch::Tensor const& input_sf,
56 + bool is_sf_swizzled_layout) {
57 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
58 (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
57 - return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf);
59 + return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf,
60 + is_sf_swizzled_layout);
61 #endif
62 TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 quantization kernel");
63 }
csrc/quantization/fp4/nvfp4_quant_kernels.cu
+141 -47
@@ -27,29 +27,23 @@
27
28 #include "cuda_utils.h"
29 #include "launch_bounds_utils.h"
30 +
31 +// Define before including nvfp4_utils.cuh so the header
32 +// can use this macro during compilation.
33 +#define NVFP4_ENABLE_ELTS16 1
34 #include "nvfp4_utils.cuh"
35
36 namespace vllm {
37
34 -template <typename Int>
35 -__host__ __device__ inline Int round_up(Int x, Int y) {
36 - static_assert(std::is_integral_v<Int>,
37 - "round_up argument must be integral type");
38 - return ((x + y - 1) / y) * y;
39 -}
40 -
41 -// Compute effective rows for grid configuration with swizzled SF layouts.
42 -inline int computeEffectiveRows(int m) {
43 - constexpr int ROW_TILE = 128;
44 - return round_up(m, ROW_TILE);
45 -}
46 -
38 // Use UE4M3 by default.
39 template <class Type, bool UE8M0_SF = false>
40 __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
50 - cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
51 - float const* SFScale, uint32_t* out, uint32_t* SFout) {
52 - using PackedVec = PackedVec<Type>;
41 + cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, int32_t num_padded_cols,
42 + Type const* __restrict__ in,
43 + float const* __restrict__ SFScale,
44 + uint32_t* __restrict__ out, uint32_t* __restrict__ SFout) {
45 + using PackedVec = vllm::PackedVec<Type>;
46 +
47 static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
48 (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
49 static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
@@ -59,33 +53,31 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
53 int32_t const numKTiles = (numCols + 63) / 64;
54
55 int sf_m = round_up<int>(numRows, 128);
62 - int sf_n_unpadded = numCols / CVT_FP4_SF_VEC_SIZE;
63 - int sf_n_int = round_up<int>(sf_n_unpadded, 4) / 4;
64 - int num_padded_cols = sf_n_int * 4 * CVT_FP4_SF_VEC_SIZE;
56 + int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x;
57 + int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD;
58
59 // Get the global scaling factor, which will be applied to the SF.
60 // Note SFScale is the same as next GEMM's alpha, which is
61 // (448.f / (Alpha_A / 6.f)).
69 - float const global_scale = SFScale == nullptr ? 1.0f : SFScale[0];
62 + float const global_scale = (SFScale == nullptr) ? 1.0f : SFScale[0];
63
64 // Iterate over all rows and cols including padded ones -
65 // ensures we visit every single scale factor address to initialize it.
66 for (int rowIdx = blockIdx.x; rowIdx < sf_m; rowIdx += gridDim.x) {
74 - for (int colIdx = threadIdx.x;
75 - colIdx < num_padded_cols / CVT_FP4_ELTS_PER_THREAD;
76 - colIdx += blockDim.x) {
77 - int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD;
78 -
67 + if (colIdx < num_padded_cols) {
68 PackedVec in_vec;
69 int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
70
71 // If we are outside valid rows OR outside valid columns -> Use Zeros
83 - if (rowIdx >= numRows || elem_idx >= numCols) {
84 - memset(&in_vec, 0, sizeof(PackedVec));
85 -
72 + bool valid = (rowIdx < numRows) && (elem_idx < numCols);
73 + if constexpr (CVT_FP4_PACK16) {
74 + ld256_or_zero_cg_u32<Type>(
75 + in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 8],
76 + valid);
77 } else {
87 - // Valid Region: Load actual data
88 - in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
78 + ld128_or_zero_cg_u32<Type>(
79 + in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 4],
80 + valid);
81 }
82
83 auto sf_out =
@@ -94,13 +86,85 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
86 rowIdx, colIdx, numKTiles, SFout);
87
88 auto out_val =
97 - cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, global_scale, sf_out);
89 + cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
90 + in_vec, global_scale, sf_out);
91 +
92 + // We do NOT write output for padding because the 'out' tensor is not
93 + // padded.
94 + if (valid) {
95 + if constexpr (CVT_FP4_PACK16) {
96 + int64_t outOffset = rowIdx * (numCols / 8) + colIdx * 2;
97 + uint64_t packed64 =
98 + (uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo);
99 + reinterpret_cast<uint64_t*>(out)[outOffset >> 1] = packed64;
100 + } else {
101 + out[inOffset] = out_val;
102 + }
103 + }
104 + }
105 + }
106 +}
107 +
108 +// Use UE4M3 by default.
109 +template <class Type, bool UE8M0_SF = false>
110 +__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
111 + cvt_fp16_to_fp4_sf_major(int32_t numRows, int32_t numCols,
112 + int32_t sf_n_unpadded, Type const* __restrict__ in,
113 + float const* __restrict__ SFScale,
114 + uint32_t* __restrict__ out,
115 + uint32_t* __restrict__ SFout) {
116 + using PackedVec = PackedVec<Type>;
117 +
118 + static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
119 + (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
120 + static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
121 + "Vec size is not matched.");
122 +
123 + int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x;
124 + int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD;
125 +
126 + // Get the global scaling factor, which will be applied to the SF.
127 + // Note SFScale is the same as next GEMM's alpha, which is
128 + // (448.f / (Alpha_A / 6.f)).
129 + float const global_scale = (SFScale == nullptr) ? 1.0f : SFScale[0];
130 +
131 + // Iterate over all rows and cols including padded ones -
132 + // ensures we visit every single scale factor address to initialize it.
133 + for (int rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) {
134 + if (colIdx < sf_n_unpadded) {
135 + PackedVec in_vec;
136 + int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
137 +
138 + // If we are outside valid rows OR outside valid columns -> Use Zeros
139 + bool valid = (rowIdx < numRows) && (elem_idx < numCols);
140 + if constexpr (CVT_FP4_PACK16) {
141 + ld256_or_zero_cg_u32<Type>(
142 + in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 8],
143 + valid);
144 + } else {
145 + ld128_or_zero_cg_u32<Type>(
146 + in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 4],
147 + valid);
148 + }
149 +
150 + auto sf_out =
151 + sf_out_rowmajor_u8<uint32_t>(rowIdx, colIdx, sf_n_unpadded, SFout);
152 +
153 + auto out_val =
154 + cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
155 + in_vec, global_scale, sf_out);
156
157 // We do NOT write output for padding because the 'out' tensor is not
158 // padded.
101 - if (rowIdx < numRows && elem_idx < numCols) {
102 - // Same as inOffset because 8 elements are packed into one uint32_t.
103 - out[inOffset] = out_val;
159 + if (valid) {
160 + if constexpr (CVT_FP4_PACK16) {
161 + int64_t outOffset = rowIdx * (numCols / 8) + colIdx * 2;
162 + uint64_t packed64 =
163 + (uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo);
164 + reinterpret_cast<uint64_t*>(out)[outOffset >> 1] = packed64;
165 + } else {
166 + out[inOffset] = out_val;
167 + }
168 }
169 }
170 }
@@ -111,7 +175,8 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
175 void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
176 torch::Tensor const& input,
177 torch::Tensor const& output_sf,
114 - torch::Tensor const& input_sf) {
178 + torch::Tensor const& input_sf,
179 + bool is_sf_swizzled_layout) {
180 int32_t m = input.size(0);
181 int32_t n = input.size(1);
182
@@ -129,19 +194,48 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
194 const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
195 auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
196
197 + int sf_n_unpadded = int(n / CVT_FP4_SF_VEC_SIZE);
198 +
199 // Grid, Block size. Each thread converts 8 values.
200 dim3 block(std::min(int(n / ELTS_PER_THREAD), 512));
201 int const numBlocksPerSM =
202 vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
136 - int effectiveRows = vllm::computeEffectiveRows(m);
137 - dim3 grid(std::min(effectiveRows, multiProcessorCount * numBlocksPerSM));
138 -
139 - VLLM_DISPATCH_HALF_TYPES(input.scalar_type(), "nvfp4_quant_kernel", [&] {
140 - using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
141 - auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
142 - // NOTE: We don't support e8m0 scales at this moment.
143 - vllm::cvt_fp16_to_fp4<cuda_type, false><<<grid, block, 0, stream>>>(
144 - m, n, input_ptr, input_sf_ptr, reinterpret_cast<uint32_t*>(output_ptr),
145 - reinterpret_cast<uint32_t*>(sf_out));
146 - });
147 -}
\ No newline at end of file
203 +
204 + if (is_sf_swizzled_layout) {
205 + int sf_n_int = int(vllm::round_up(sf_n_unpadded, 4) / 4);
206 + int32_t num_padded_cols =
207 + sf_n_int * 4 * CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
208 +
209 + int grid_y = vllm::div_round_up(num_padded_cols, static_cast<int>(block.x));
210 + int grid_x =
211 + std::min(vllm::computeEffectiveRows(m),
212 + std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
213 + dim3 grid(grid_x, grid_y);
214 +
215 + VLLM_DISPATCH_HALF_TYPES(input.scalar_type(), "nvfp4_quant_kernel", [&] {
216 + using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
217 + auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
218 + // NOTE: We don't support e8m0 scales at this moment.
219 + vllm::cvt_fp16_to_fp4<cuda_type, false><<<grid, block, 0, stream>>>(
220 + m, n, num_padded_cols, input_ptr, input_sf_ptr,
221 + reinterpret_cast<uint32_t*>(output_ptr),
222 + reinterpret_cast<uint32_t*>(sf_out));
223 + });
224 + } else {
225 + int grid_y = vllm::div_round_up(sf_n_unpadded, static_cast<int>(block.x));
226 + int grid_x = std::min(
227 + m, std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
228 + dim3 grid(grid_x, grid_y);
229 +
230 + VLLM_DISPATCH_HALF_TYPES(input.scalar_type(), "nvfp4_quant_kernel", [&] {
231 + using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
232 + auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
233 + // NOTE: We don't support e8m0 scales at this moment.
234 + vllm::cvt_fp16_to_fp4_sf_major<cuda_type, false>
235 + <<<grid, block, 0, stream>>>(m, n, sf_n_unpadded, input_ptr,
236 + input_sf_ptr,
237 + reinterpret_cast<uint32_t*>(output_ptr),
238 + reinterpret_cast<uint32_t*>(sf_out));
239 + });
240 + }
241 +}
csrc/quantization/fp4/nvfp4_utils.cuh
+171 -25
@@ -19,9 +19,17 @@
19 #include <cuda_runtime.h>
20 #include <cuda_fp8.h>
21
22 -#define ELTS_PER_THREAD 8
23 -
22 +#if (defined(NVFP4_ENABLE_ELTS16) && (CUDART_VERSION >= 12090) && \
23 + defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100)
24 + #define ELTS_PER_THREAD 16
25 +constexpr int CVT_FP4_ELTS_PER_THREAD = 16;
26 +constexpr bool CVT_FP4_PACK16 = true;
27 +#else
28 + #define ELTS_PER_THREAD 8
29 constexpr int CVT_FP4_ELTS_PER_THREAD = 8;
30 +constexpr bool CVT_FP4_PACK16 = false;
31 +#endif
32 +
33 constexpr int CVT_FP4_SF_VEC_SIZE = 16;
34
35 namespace vllm {
@@ -68,19 +76,46 @@ struct TypeConverter<__nv_bfloat16> {
76 using Type = __nv_bfloat162;
77 };
78
79 +#if (defined(NVFP4_ENABLE_ELTS16) && (CUDART_VERSION >= 12090) && \
80 + defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100)
81 +// Define a 32 bytes packed data type.
82 +template <class Type>
83 +struct alignas(32) PackedVec {
84 + typename TypeConverter<Type>::Type elts[8];
85 +};
86 +#else
87 // Define a 16 bytes packed data type.
88 template <class Type>
73 -struct PackedVec {
89 +struct alignas(16) PackedVec {
90 typename TypeConverter<Type>::Type elts[4];
91 };
92 +#endif
93
94 template <>
95 struct PackedVec<__nv_fp8_e4m3> {
96 __nv_fp8x2_e4m3 elts[8];
97 };
98
99 +template <typename Int>
100 +__host__ __device__ inline Int round_up(Int x, Int y) {
101 + static_assert(std::is_integral_v<Int>,
102 + "round_up argument must be integral type");
103 + return ((x + y - 1) / y) * y;
104 +}
105 +
106 +template <typename Int>
107 +__host__ __device__ __forceinline__ Int div_round_up(Int x, Int y) {
108 + return (x + y - 1) / y;
109 +}
110 +
111 +// Compute effective rows for grid configuration with swizzled SF layouts.
112 +inline int computeEffectiveRows(int m) {
113 + constexpr int ROW_TILE = 128;
114 + return round_up(m, ROW_TILE);
115 +}
116 +
117 // 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]) {
118 +inline __device__ uint32_t fp32_vec8_to_e2m1(float (&array)[8]) {
119 uint32_t val;
120 asm volatile(
121 "{\n"
@@ -101,7 +136,7 @@ inline __device__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) {
136 }
137
138 // 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]) {
139 +__device__ __forceinline__ uint32_t fp32_vec8_to_e2m1(float2 (&array)[4]) {
140 uint32_t val;
141 asm volatile(
142 "{\n"
@@ -114,20 +149,115 @@ inline __device__ uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) {
149 "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
150 "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
151 "mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
117 - "}"
152 + "}\n"
153 : "=r"(val)
154 : "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y),
155 "f"(array[2].x), "f"(array[2].y), "f"(array[3].x), "f"(array[3].y));
156 return val;
157 }
158
159 +struct u32x2 {
160 + uint32_t lo, hi;
161 +};
162 +
163 +using fp4_packed_t = std::conditional_t<CVT_FP4_PACK16, u32x2, uint32_t>;
164 +
165 +__device__ __forceinline__ u32x2 fp32_vec16_to_e2m1(float2 (&array)[8]) {
166 + u32x2 out;
167 + asm volatile(
168 + "{\n"
169 + ".reg .b8 b0;\n"
170 + ".reg .b8 b1;\n"
171 + ".reg .b8 b2;\n"
172 + ".reg .b8 b3;\n"
173 + ".reg .b8 b4;\n"
174 + ".reg .b8 b5;\n"
175 + ".reg .b8 b6;\n"
176 + ".reg .b8 b7;\n"
177 + "cvt.rn.satfinite.e2m1x2.f32 b0, %3, %2;\n"
178 + "cvt.rn.satfinite.e2m1x2.f32 b1, %5, %4;\n"
179 + "cvt.rn.satfinite.e2m1x2.f32 b2, %7, %6;\n"
180 + "cvt.rn.satfinite.e2m1x2.f32 b3, %9, %8;\n"
181 + "cvt.rn.satfinite.e2m1x2.f32 b4, %11, %10;\n"
182 + "cvt.rn.satfinite.e2m1x2.f32 b5, %13, %12;\n"
183 + "cvt.rn.satfinite.e2m1x2.f32 b6, %15, %14;\n"
184 + "cvt.rn.satfinite.e2m1x2.f32 b7, %17, %16;\n"
185 + "mov.b32 %0, {b0, b1, b2, b3};\n"
186 + "mov.b32 %1, {b4, b5, b6, b7};\n"
187 + "}\n"
188 + : "=r"(out.lo), "=r"(out.hi)
189 + : "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y),
190 + "f"(array[2].x), "f"(array[2].y), "f"(array[3].x), "f"(array[3].y),
191 + "f"(array[4].x), "f"(array[4].y), "f"(array[5].x), "f"(array[5].y),
192 + "f"(array[6].x), "f"(array[6].y), "f"(array[7].x), "f"(array[7].y));
193 + return out;
194 +}
195 +
196 +__device__ __forceinline__ uint32_t pack_fp4(float2 (&v)[4]) {
197 + return fp32_vec8_to_e2m1(v);
198 +}
199 +
200 +__device__ __forceinline__ u32x2 pack_fp4(float2 (&v)[8]) {
201 + return fp32_vec16_to_e2m1(v);
202 +}
203 +
204 // Fast reciprocal.
125 -inline __device__ float reciprocal_approximate_ftz(float a) {
205 +__device__ __forceinline__ float reciprocal_approximate_ftz(float a) {
206 float b;
127 - asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a));
207 + asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(b) : "f"(a));
208 return b;
209 }
210
211 +template <class Type>
212 +__device__ __forceinline__ void ld128_or_zero_cg_u32(PackedVec<Type>& out,
213 + const void* ptr,
214 + bool pred) {
215 + uint32_t r0, r1, r2, r3;
216 +
217 + asm volatile(
218 + "{\n"
219 + " .reg .pred pr;\n"
220 + " setp.ne.u32 pr, %4, 0;\n"
221 + " mov.u32 %0, 0;\n"
222 + " mov.u32 %1, 0;\n"
223 + " mov.u32 %2, 0;\n"
224 + " mov.u32 %3, 0;\n"
225 + " @pr ld.global.cg.v4.u32 {%0,%1,%2,%3}, [%5];\n"
226 + "}\n"
227 + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
228 + : "r"((int)pred), "l"(ptr));
229 +
230 + *reinterpret_cast<uint4*>(&out) = uint4{r0, r1, r2, r3};
231 +}
232 +
233 +template <class Type>
234 +__device__ __forceinline__ void ld256_or_zero_cg_u32(PackedVec<Type>& out,
235 + const void* ptr,
236 + bool pred) {
237 + uint32_t r0, r1, r2, r3, r4, r5, r6, r7;
238 +
239 + asm volatile(
240 + "{\n"
241 + " .reg .pred pr;\n"
242 + " setp.ne.u32 pr, %8, 0;\n"
243 + " mov.u32 %0, 0;\n"
244 + " mov.u32 %1, 0;\n"
245 + " mov.u32 %2, 0;\n"
246 + " mov.u32 %3, 0;\n"
247 + " mov.u32 %4, 0;\n"
248 + " mov.u32 %5, 0;\n"
249 + " mov.u32 %6, 0;\n"
250 + " mov.u32 %7, 0;\n"
251 + " @pr ld.global.cg.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%9];\n"
252 + "}\n"
253 + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3), "=r"(r4), "=r"(r5), "=r"(r6),
254 + "=r"(r7)
255 + : "r"((int)pred), "l"(ptr));
256 +
257 + reinterpret_cast<uint4*>(&out)[0] = uint4{r0, r1, r2, r3};
258 + reinterpret_cast<uint4*>(&out)[1] = uint4{r4, r5, r6, r7};
259 +}
260 +
261 // Compute SF output offset for swizzled tensor core layout.
262 // SF layout: [numMTiles, numKTiles, 32, 4, 4]
263 // Caller must precompute: numKTiles = (numCols + 63) / 64
@@ -166,21 +296,41 @@ __device__ __forceinline__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset(
296 return reinterpret_cast<uint8_t*>(SFout) + SFOffset;
297 }
298
299 +template <class SFType>
300 +__device__ __forceinline__ uint8_t* sf_out_rowmajor_u8(int row, int pack,
301 + int packs_per_row_sf,
302 + SFType* SFout) {
303 + constexpr int PACK = CVT_FP4_ELTS_PER_THREAD;
304 + constexpr int THREADS_PER_SF =
305 + CVT_FP4_SF_VEC_SIZE / PACK; // 1 if PACK=16, 2 else PACK=8
306 +
307 + if (threadIdx.x % THREADS_PER_SF != 0) return nullptr;
308 +
309 + int sf_col =
310 + pack / THREADS_PER_SF; // PACK=16 => sf_col=pack; PACK=8 => sf_col=pack/2
311 + int64_t off = (int64_t)row * packs_per_row_sf + sf_col;
312 +
313 + return (uint8_t*)SFout + off;
314 +}
315 +
316 // Quantizes the provided PackedVec into the uint32_t output
170 -template <class Type, bool UE8M0_SF = false>
171 -__device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal,
172 - uint8_t* SFout) {
317 +template <class Type, int CVT_FP4_NUM_THREADS_PER_SF, bool UE8M0_SF = false>
318 +__device__ __forceinline__ fp4_packed_t
319 +cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal, uint8_t* SFout) {
320 // Get absolute maximum values among the local 8 values.
321 auto localMax = __habs2(vec.elts[0]);
322
176 -// Local maximum value.
323 + // Local maximum value.
324 #pragma unroll
325 for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
326 localMax = __hmax2(localMax, __habs2(vec.elts[i]));
327 }
328
329 // Get the absolute maximum among all 16 values (two threads).
183 - localMax = __hmax2(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax);
330 +
331 + if constexpr (CVT_FP4_NUM_THREADS_PER_SF == 2) {
332 + localMax = __hmax2(__shfl_xor_sync(0xffffffffu, localMax, 1), localMax);
333 + }
334 // Get the final absolute maximum values.
335 float vecMax = float(__hmax(localMax.x, localMax.y));
336
@@ -205,18 +355,17 @@ __device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal,
355 // Convert back to fp32.
356 SFValue = float(tmp);
357 }
358 +
359 + // Write the SF to global memory (STG.8).
360 + if (SFout) *SFout = fp8SFVal;
361 +
362 // Get the output scale.
363 // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) *
364 // reciprocal(SFScaleVal))
365 float outputScale =
212 - SFValue != 0 ? reciprocal_approximate_ftz(
213 - SFValue * reciprocal_approximate_ftz(SFScaleVal))
214 - : 0.0f;
215 -
216 - if (SFout) {
217 - // Write the SF to global memory (STG.8).
218 - *SFout = fp8SFVal;
219 - }
366 + SFValue != 0.0f ? reciprocal_approximate_ftz(
367 + SFValue * reciprocal_approximate_ftz(SFScaleVal))
368 + : 0.0f;
369
370 // Convert the input to float.
371 float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
@@ -233,10 +382,7 @@ __device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal,
382 }
383
384 // Convert to e2m1 values.
236 - uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals);
237 -
238 - // Write the e2m1 values to global memory.
239 - return e2m1Vec;
385 + return pack_fp4(fp2Vals);
386 }
387
388 // silu in float32
csrc/torch_bindings.cpp
+2 -1
@@ -546,7 +546,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
546 // Compute NVFP4 block quantized tensor.
547 ops.def(
548 "scaled_fp4_quant(Tensor! output, Tensor input,"
549 - " Tensor! output_scale, Tensor input_scale) -> ()");
549 + " Tensor! output_scale, Tensor input_scale, bool "
550 + "is_sf_swizzled_layout) -> ()");
551 ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant);
552
553 // Compute NVFP4 experts quantization.
tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py
+6 -2
@@ -107,10 +107,14 @@ def test_flashinfer_nvfp4_gemm(
107 # from checkpoints are in linear scales.
108 # So instead of needing to swizzle for cutlass as in modelopt.py,
109 # we need to unswizzle for trtllm here.
110 - a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a_dtype, a_global_scale, backend)
110 + a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(
111 + a_dtype, a_global_scale, is_sf_swizzled_layout=True, backend=backend
112 + )
113 is_sf_128x4_layout = not (backend == "trtllm" and m <= 32)
114
113 - b_fp4, b_scale_interleaved = ops.scaled_fp4_quant(b_dtype, b_global_scale)
115 + b_fp4, b_scale_interleaved = ops.scaled_fp4_quant(
116 + b_dtype, b_global_scale, is_sf_swizzled_layout=True
117 + )
118
119 # get_ref_results unswizzles the scales internally.
120 expected_out = get_ref_results(
tests/kernels/quantization/test_nvfp4_quant.py
+28
@@ -27,6 +27,12 @@ PAD_SHAPES = [
27 (150, 128),
28 (150, 48),
29 (90, 80),
30 + (128, 512),
31 + (128, 1024),
32 + (128, 2048),
33 + (64, 7168),
34 + (64, 7152),
35 + (32, 14336),
36 ]
37 SEEDS = [42]
38 CUDA_DEVICES = ["cuda:0"]
@@ -173,3 +179,25 @@ def test_quantize_to_fp4_padded(pad_shape: tuple[int, int]) -> None:
179 out_ans = cast_from_fp4(out, m, n)
180 torch.testing.assert_close(out_ans, out_ref)
181 torch.testing.assert_close(scale_ans, scale_ref)
182 +
183 +
184 +@pytest.mark.parametrize("pad_shape", PAD_SHAPES)
185 +@torch.inference_mode()
186 +def test_quantize_to_fp4_padded_no_sf_swizzled(pad_shape: tuple[int, int]) -> None:
187 + dtype = torch.float16
188 + set_random_seed(42)
189 + torch.set_default_device("cuda:0")
190 +
191 + m, n = pad_shape
192 +
193 + x = torch.randn((m, n), dtype=dtype)
194 +
195 + tensor_amax = torch.abs(x).max().to(torch.float32)
196 + global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
197 + out_ref, scale_ref = ref_nvfp4_quant(x, global_scale)
198 +
199 + out, out_scale = ops.scaled_fp4_quant(x, global_scale, is_sf_swizzled_layout=False)
200 + scale_ans = out_scale.to(torch.float32)
201 + out_ans = cast_from_fp4(out, m, n)
202 + torch.testing.assert_close(out_ans, out_ref)
203 + torch.testing.assert_close(scale_ans, scale_ref)
vllm/_custom_ops.py
+18 -13
@@ -1534,6 +1534,7 @@ def permute_cols(a: torch.Tensor, perm: torch.Tensor) -> torch.Tensor:
1534 def scaled_fp4_quant(
1535 input: torch.Tensor,
1536 input_global_scale: torch.Tensor,
1537 + is_sf_swizzled_layout: bool = True,
1538 backend: str = "none",
1539 ) -> tuple[torch.Tensor, torch.Tensor]:
1540 """
@@ -1577,22 +1578,26 @@ def scaled_fp4_quant(
1578 else:
1579 # Two fp4 values will be packed into an uint8.
1580 output = torch.empty((m, n // 2), device=device, dtype=torch.uint8)
1581 + if is_sf_swizzled_layout:
1582 + # We use the rounded values to store the swizzled values. Due to the
1583 + # requirement of the Tensor Core, the minimum tile is 128x4 for the scales.
1584 + # So, we first pad the scales to multiples of 128 and 4. Then, the scales
1585 + # (in float8_e4m3fn) are packed into an int32 for every 4 values. More:
1586 + # https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-b-layout-4x
1587 + round_up = lambda x, y: (x + y - 1) // y * y
1588 + rounded_m = round_up(m, 128)
1589 + scale_n = n // block_size
1590 + rounded_n = round_up(scale_n, 4)
1591 + output_scale = torch.empty(
1592 + (rounded_m, rounded_n // 4), device=device, dtype=torch.int32
1593 + )
1594 + else:
1595 + output_scale = torch.empty((m, n // 16), device=device, dtype=torch.uint8)
1596
1581 - # We use the rounded values to store the swizzled values. Due to the
1582 - # requirement of the Tensor Core, the minimum tile is 128x4 for the scales.
1583 - # So, we first pad the scales to multiples of 128 and 4. Then, the scales
1584 - # (in float8_e4m3fn) are packed into an int32 for every 4 values. More:
1585 - # https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-b-layout-4x
1586 - round_up = lambda x, y: (x + y - 1) // y * y
1587 - rounded_m = round_up(m, 128)
1588 - scale_n = n // block_size
1589 - rounded_n = round_up(scale_n, 4)
1590 - output_scale = torch.empty(
1591 - (rounded_m, rounded_n // 4), device=device, dtype=torch.int32
1597 + torch.ops._C.scaled_fp4_quant(
1598 + output, input, output_scale, input_global_scale, is_sf_swizzled_layout
1599 )
1600
1594 - torch.ops._C.scaled_fp4_quant(output, input, output_scale, input_global_scale)
1595 -
1601 output_scale = output_scale.view(torch.float8_e4m3fn)
1602 return output, output_scale
1603
vllm/compilation/activation_quant_fusion.py
+1
@@ -152,6 +152,7 @@ class SiluMulNvfp4QuantPattern(ActivationQuantPattern):
152 input=result_silu_mul,
153 output_scale=output_scale,
154 input_scale=scale,
155 + is_sf_swizzled_layout=True,
156 )
157 return at[1], at[2]
158
vllm/compilation/collective_fusion.py
+2
@@ -946,6 +946,7 @@ class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):
946 input=rms,
947 output_scale=output_scale,
948 input_scale=input_global_scale,
949 + is_sf_swizzled_layout=True,
950 )
951
952 # quant_out, allreduce_output, output_scale
@@ -1043,6 +1044,7 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
1044 input=rms,
1045 output_scale=output_scale,
1046 input_scale=input_global_scale,
1047 + is_sf_swizzled_layout=True,
1048 )
1049
1050 # quant_out, allreduce_output, output_scale
vllm/compilation/fusion_attn.py
+1
@@ -248,6 +248,7 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
248 input=attn_out_view,
249 output_scale=output_scale,
250 input_scale=input_scale,
251 + is_sf_swizzled_layout=True,
252 )
253 output_scale_view = torch.ops.aten.view.dtype(at2[2], FP8_DTYPE)
254 return at2[1], output_scale_view
vllm/model_executor/layers/fused_moe/utils.py
+1 -4
@@ -24,7 +24,6 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
24 mxfp8_e4m3_quantize,
25 )
26 from vllm.triton_utils import tl, triton
27 -from vllm.utils.flashinfer import flashinfer_fp4_quantize
27 from vllm.utils.math_utils import cdiv
28 from vllm.utils.torch_utils import is_torch_equal_or_newer
29
@@ -117,9 +116,7 @@ def _nvfp4_quantize(
116 A_scale: torch.Tensor | None,
117 is_sf_swizzled_layout: bool,
118 ) -> tuple[torch.Tensor, torch.Tensor]:
120 - return flashinfer_fp4_quantize(
121 - A, A_scale, is_sf_swizzled_layout=is_sf_swizzled_layout
122 - )
119 + return ops.scaled_fp4_quant(A, A_scale, is_sf_swizzled_layout=is_sf_swizzled_layout)
120
121
122 def _fp8_quantize(
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
+4 -1
@@ -191,7 +191,10 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
191
192 # quantize BF16 or FP16 to (FP4 and interleaved block scale)
193 x_fp4, x_blockscale = scaled_fp4_quant(
194 - x, layer.input_global_scale, self.backend
194 + x,
195 + layer.input_global_scale,
196 + is_sf_swizzled_layout=True,
197 + backend=self.backend,
198 )
199
200 mm_args = (
vllm/model_executor/layers/quantization/modelopt.py
+3 -1
@@ -1307,7 +1307,9 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
1307 output_shape = [x.shape[0], layer.weight.shape[0]]
1308
1309 # quantize BF16 or FP16 to (FP4 and interleaved block scale)
1310 - x_fp4, x_blockscale = scaled_fp4_quant(x, layer.input_scale_inv, self.backend)
1310 + x_fp4, x_blockscale = scaled_fp4_quant(
1311 + x, layer.input_scale_inv, is_sf_swizzled_layout=True, backend=self.backend
1312 + )
1313
1314 # validate dtypes of quantized input, input block scale,
1315 # weight and weight_blockscale
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+5 -8
@@ -8,6 +8,7 @@ import torch
8
9 import vllm.envs as envs
10 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11 +from vllm import _custom_ops as ops
12 from vllm.logger import init_logger
13 from vllm.model_executor.layers.fused_moe.config import (
14 FusedMoEConfig,
@@ -341,10 +342,8 @@ def flashinfer_trtllm_fp4_moe(
342 hidden_states_fp4, hidden_states_scale_linear_fp4 = x
343 else:
344 # hidden_states is the already quantized
344 - (hidden_states_fp4, hidden_states_scale_linear_fp4) = flashinfer.fp4_quantize(
345 - x,
346 - layer.a1_gscale,
347 - is_sf_swizzled_layout=False,
345 + (hidden_states_fp4, hidden_states_scale_linear_fp4) = ops.scaled_fp4_quant(
346 + x, layer.a1_gscale, is_sf_swizzled_layout=False
347 )
348
349 # Determine routing method type
@@ -443,10 +442,8 @@ def flashinfer_trtllm_fp4_routed_moe(
442 hidden_states_fp4, hidden_states_scale_linear_fp4 = x
443 else:
444 # Quantize input to FP4
446 - (hidden_states_fp4, hidden_states_scale_linear_fp4) = flashinfer.fp4_quantize(
447 - x,
448 - layer.a1_gscale,
449 - is_sf_swizzled_layout=False,
445 + (hidden_states_fp4, hidden_states_scale_linear_fp4) = ops.scaled_fp4_quant(
446 + x, layer.a1_gscale, is_sf_swizzled_layout=False
447 )
448
449 # Call TRT-LLM FP4 block-scale MoE kernel