@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 06d490282f
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 11, 03:45 UTC

[NVFP4][Perf] Tune NVFP4 input quant kernel for small batch size (#30897)

Signed-off-by: mgoin <mgoin64@gmail.com>

Michael Goin committed Dec 21, 2025 at 12:41 UTC 06d490282f2bab6922137eb5230be9df5ebbe9c4
5 files changed +243 -97
benchmarks/kernels/bench_nvfp4_quant.py new
+177
@@ -0,0 +1,177 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import argparse
4 +import copy
5 +import itertools
6 +
7 +import torch
8 +from weight_shapes import WEIGHT_SHAPES
9 +
10 +from vllm import _custom_ops as ops
11 +from vllm.platforms import current_platform
12 +from vllm.scalar_type import scalar_types
13 +from vllm.triton_utils import triton
14 +from vllm.utils.flashinfer import flashinfer_fp4_quantize
15 +
16 +if not current_platform.has_device_capability(100):
17 + raise RuntimeError("NVFP4 requires compute capability of 10.0 (Blackwell)")
18 +
19 +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),
25 +}
26 +
27 +_enabled = [k for k, v in PROVIDER_CFGS.items() if v["enabled"]]
28 +
29 +
30 +def compute_global_scale(tensor: torch.Tensor) -> torch.Tensor:
31 + """Compute global scale for FP4 quantization."""
32 + amax = torch.abs(tensor).max().to(torch.float32)
33 + return FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / amax
34 +
35 +
36 +@triton.testing.perf_report(
37 + triton.testing.Benchmark(
38 + x_names=["batch_size"],
39 + x_vals=[1, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096],
40 + x_log=False,
41 + line_arg="provider",
42 + line_vals=_enabled,
43 + line_names=_enabled,
44 + ylabel="us (lower is better)",
45 + plot_name="NVFP4 Input Quantization Latency (us)",
46 + args={},
47 + )
48 +)
49 +def benchmark(batch_size, provider, N, K):
50 + M = batch_size
51 + device = "cuda"
52 + dtype = torch.bfloat16
53 +
54 + # Create input tensor
55 + a = torch.randn((M, K), device=device, dtype=dtype)
56 +
57 + # Compute global scale for activation
58 + a_global_scale = compute_global_scale(a)
59 +
60 + quantiles = [0.5, 0.2, 0.8]
61 +
62 + cfg = PROVIDER_CFGS[provider]
63 +
64 + if cfg["backend"] == "vllm":
65 + # 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 + elif cfg["backend"] == "flashinfer":
71 + # 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 + )
79 +
80 + # Convert ms to us for better readability at small batch sizes
81 + to_us = lambda t_ms: t_ms * 1000
82 + return to_us(ms), to_us(max_ms), to_us(min_ms)
83 +
84 +
85 +def prepare_shapes(args):
86 + out = []
87 + for model, tp_size in itertools.product(args.models, args.tp_sizes):
88 + for KN, tp_dim in copy.deepcopy(WEIGHT_SHAPES[model]):
89 + KN[tp_dim] //= tp_size
90 + KN.append(model)
91 + out.append(KN)
92 + return out
93 +
94 +
95 +def _test_accuracy_once(M: int, K: int, dtype: torch.dtype, device: str):
96 + """Test accuracy between vLLM and FlashInfer FP4 quantization."""
97 + # Create input tensor
98 + a = torch.randn((M, K), device=device, dtype=dtype)
99 +
100 + # Compute global scale
101 + a_global_scale = compute_global_scale(a)
102 +
103 + # vLLM quantization
104 + vllm_fp4, vllm_scale = ops.scaled_fp4_quant(a, a_global_scale)
105 +
106 + # FlashInfer quantization (with swizzled layout to match vLLM's output)
107 + flashinfer_fp4, flashinfer_scale = flashinfer_fp4_quantize(
108 + a, a_global_scale, is_sf_swizzled_layout=True
109 + )
110 + flashinfer_scale = flashinfer_scale.view(torch.float8_e4m3fn)
111 +
112 + # Compare outputs
113 + torch.testing.assert_close(
114 + vllm_fp4,
115 + flashinfer_fp4,
116 + )
117 + print(f"M={M}, K={K}, dtype={dtype}: PASSED")
118 +
119 +
120 +def test_accuracy():
121 + """Run accuracy tests across various shapes."""
122 + print("\n" + "=" * 60)
123 + print("Running accuracy tests: vLLM vs FlashInfer")
124 + print("=" * 60)
125 +
126 + device = "cuda"
127 + dtype = torch.bfloat16
128 +
129 + # Test various batch sizes and hidden dimensions
130 + Ms = [1, 1024]
131 + Ks = [4096]
132 +
133 + for M in Ms:
134 + for K in Ks:
135 + _test_accuracy_once(M, K, dtype, device)
136 +
137 + print("\nAll accuracy tests passed!")
138 +
139 +
140 +if __name__ == "__main__":
141 + parser = argparse.ArgumentParser(
142 + description="Benchmark NVFP4 quantization: vLLM vs FlashInfer"
143 + )
144 + parser.add_argument(
145 + "--models",
146 + nargs="+",
147 + type=str,
148 + default=["meta-llama/Llama-3.1-8B-Instruct"],
149 + choices=list(WEIGHT_SHAPES.keys()),
150 + )
151 + parser.add_argument("--tp-sizes", nargs="+", type=int, default=[1])
152 + parser.add_argument(
153 + "--save-path",
154 + type=str,
155 + default=None,
156 + help="Path to save benchmark results",
157 + )
158 + parser.add_argument(
159 + "--accuracy",
160 + action="store_true",
161 + help="Run accuracy tests",
162 + )
163 + args = parser.parse_args()
164 +
165 + if args.accuracy:
166 + test_accuracy()
167 +
168 + for K, N, model in prepare_shapes(args):
169 + print(f"\n{model}, N={N} K={K}")
170 + benchmark.run(
171 + print_data=True,
172 + save_path=args.save_path,
173 + N=N,
174 + K=K,
175 + )
176 +
177 + print("\nBenchmark finished!")
csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu
+4 -1
@@ -74,6 +74,9 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
74 static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
75 "Vec size is not matched.");
76
77 + // Precompute SF layout parameter (constant for entire kernel).
78 + int32_t const numKTiles = (numCols + 63) / 64;
79 +
80 // Get the global scaling factor, which will be applied to the SF.
81 // Note SFScale is the same as next GEMM's alpha, which is
82 // (448.f / (Alpha_A / 6.f)).
@@ -101,7 +104,7 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
104 auto sf_out =
105 cvt_quant_to_fp4_get_sf_out_offset<uint32_t,
106 CVT_FP4_NUM_THREADS_PER_SF>(
104 - rowIdx, colIdx, numCols, SFout);
107 + rowIdx, colIdx, numKTiles, SFout);
108
109 out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(out_silu_mul, SFScaleVal,
110 sf_out);
csrc/quantization/fp4/nvfp4_experts_quant.cu
+14 -17
@@ -25,6 +25,7 @@
25 #include <cuda_fp8.h>
26 #include "dispatch_utils.h"
27
28 +#include "cuda_utils.h"
29 #include "nvfp4_utils.cuh"
30 #include "launch_bounds_utils.h"
31
@@ -44,6 +45,9 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
45 static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
46 "Vec size is not matched.");
47
48 + // Precompute SF layout parameter (constant for entire kernel).
49 + int32_t const numKTiles = (numCols + 63) / 64;
50 +
51 int tid = blockIdx.x * blockDim.x + threadIdx.x;
52 int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
53
@@ -112,17 +116,13 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
116 // (448.f / (Alpha_A / 6.f)).
117 float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[expert_idx];
118
115 - int factor = CVT_FP4_SF_VEC_SIZE * 4;
116 - // The actual output_scales dim is computed from the padded numCols.
117 - int32_t numCols_padded = (numCols + factor - 1) / factor * factor;
118 - int numCols_SFout = numCols_padded / CVT_FP4_SF_VEC_SIZE / 4;
119 uint32_t* SFout_in_expert =
120 - SFout + output_scale_offset_by_experts[expert_idx] * numCols_SFout;
120 + SFout + output_scale_offset_by_experts[expert_idx] * numKTiles;
121
122 auto sf_out =
123 cvt_quant_to_fp4_get_sf_out_offset<uint32_t,
124 CVT_FP4_NUM_THREADS_PER_SF>(
125 - rowIdx_in_expert, colIdx, numCols, SFout_in_expert);
125 + rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
126
127 out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
128 }
@@ -140,6 +140,10 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
140 (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
141 static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
142 "Vec size is not matched.");
143 +
144 + // Precompute SF layout parameter (constant for entire kernel).
145 + int32_t const numKTiles = (numCols + 63) / 64;
146 +
147 extern __shared__ uint32_t shared_input_offsets[];
148
149 // Load input offsets into shared memory.
@@ -202,16 +206,13 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
206
207 float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[expert_idx];
208
205 - int factor = CVT_FP4_SF_VEC_SIZE * 4;
206 - int32_t numCols_padded = (numCols + factor - 1) / factor * factor;
207 - int numCols_SFout = numCols_padded / CVT_FP4_SF_VEC_SIZE / 4;
209 uint32_t* SFout_in_expert =
209 - SFout + output_scale_offset_by_experts[expert_idx] * numCols_SFout;
210 + SFout + output_scale_offset_by_experts[expert_idx] * numKTiles;
211
212 auto sf_out =
213 cvt_quant_to_fp4_get_sf_out_offset<uint32_t,
214 CVT_FP4_NUM_THREADS_PER_SF>(
214 - rowIdx_in_expert, colIdx, numCols, SFout_in_expert);
215 + rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
216
217 out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
218 }
@@ -222,12 +223,8 @@ void quant_impl(void* output, void* output_scale, void* input,
223 void* input_global_scale, void* input_offset_by_experts,
224 void* output_scale_offset_by_experts, int m_topk, int k,
225 int n_experts, cudaStream_t stream) {
225 - // TODO: this multiProcessorCount should be cached.
226 - int device;
227 - cudaGetDevice(&device);
228 - int multiProcessorCount;
229 - cudaDeviceGetAttribute(&multiProcessorCount, cudaDevAttrMultiProcessorCount,
230 - device);
226 + int multiProcessorCount =
227 + get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
228
229 // Grid, Block size.
230 // Each thread converts 8 values.
csrc/quantization/fp4/nvfp4_quant_kernels.cu
+20 -42
@@ -38,6 +38,12 @@ __host__ __device__ inline Int round_up(Int x, Int y) {
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 +
47 // Use UE4M3 by default.
48 template <class Type, bool UE8M0_SF = false>
49 __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
@@ -49,6 +55,9 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
55 static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
56 "Vec size is not matched.");
57
58 + // Precompute SF layout parameter (constant for entire kernel).
59 + int32_t const numKTiles = (numCols + 63) / 64;
60 +
61 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;
@@ -79,7 +88,7 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
88 auto sf_out =
89 cvt_quant_to_fp4_get_sf_out_offset<uint32_t,
90 CVT_FP4_NUM_THREADS_PER_SF>(
82 - rowIdx, colIdx, numCols, SFout);
91 + rowIdx, colIdx, numKTiles, SFout);
92
93 out_pos =
94 cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, global_scale, sf_out);
@@ -87,43 +96,6 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
96 }
97 }
98
90 -template <typename T>
91 -void invokeFP4Quantization(int m, int n, T const* input, float const* SFScale,
92 - int64_t* output, int32_t* SFOuput, bool useUE8M0,
93 - int multiProcessorCount, cudaStream_t stream) {
94 - // Grid, Block size.
95 - // Each thread converts 8 values.
96 - dim3 block(std::min(int(n / ELTS_PER_THREAD), 512));
97 - // Get number of blocks per SM
98 - int const numBlocksPerSM =
99 - vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
100 - dim3 grid(std::min(int(m), multiProcessorCount * numBlocksPerSM));
101 -
102 - // Launch the cvt kernel.
103 - if (useUE8M0) {
104 - cvt_fp16_to_fp4<T, true><<<grid, block, 0, stream>>>(
105 - m, n, input, SFScale, reinterpret_cast<uint32_t*>(output),
106 - reinterpret_cast<uint32_t*>(SFOuput));
107 - } else {
108 - cvt_fp16_to_fp4<T, false><<<grid, block, 0, stream>>>(
109 - m, n, input, SFScale, reinterpret_cast<uint32_t*>(output),
110 - reinterpret_cast<uint32_t*>(SFOuput));
111 - }
112 -}
113 -
114 -// Instantiate the function.
115 -template void invokeFP4Quantization(int m, int n, half const* input,
116 - float const* SFScale, int64_t* output,
117 - int32_t* SFOuput, bool useUE8M0,
118 - int multiProcessorCount,
119 - cudaStream_t stream);
120 -
121 -template void invokeFP4Quantization(int m, int n, __nv_bfloat16 const* input,
122 - float const* SFScale, int64_t* output,
123 - int32_t* SFOuput, bool useUE8M0,
124 - int multiProcessorCount,
125 - cudaStream_t stream);
126 -
99 } // namespace vllm
100
101 void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
@@ -147,13 +119,19 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
119 const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
120 auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
121
150 - // We don't support e8m0 scales at this moment.
151 - bool useUE8M0 = false;
122 + // Grid, Block size. Each thread converts 8 values.
123 + dim3 block(std::min(int(n / ELTS_PER_THREAD), 512));
124 + int const numBlocksPerSM =
125 + vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
126 + int effectiveRows = vllm::computeEffectiveRows(m);
127 + dim3 grid(std::min(effectiveRows, multiProcessorCount * numBlocksPerSM));
128
129 VLLM_DISPATCH_HALF_TYPES(input.scalar_type(), "nvfp4_quant_kernel", [&] {
130 using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
131 auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
156 - vllm::invokeFP4Quantization(m, n, input_ptr, input_sf_ptr, output_ptr,
157 - sf_out, useUE8M0, multiProcessorCount, stream);
132 + // NOTE: We don't support e8m0 scales at this moment.
133 + vllm::cvt_fp16_to_fp4<cuda_type, false><<<grid, block, 0, stream>>>(
134 + m, n, input_ptr, input_sf_ptr, reinterpret_cast<uint32_t*>(output_ptr),
135 + reinterpret_cast<uint32_t*>(sf_out));
136 });
137 }
csrc/quantization/fp4/nvfp4_utils.cuh
+28 -37
@@ -128,51 +128,42 @@ inline __device__ float reciprocal_approximate_ftz(float a) {
128 return b;
129 }
130
131 +// Compute SF output offset for swizzled tensor core layout.
132 +// SF layout: [numMTiles, numKTiles, 32, 4, 4]
133 +// Caller must precompute: numKTiles = (numCols + 63) / 64
134 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 +__device__ __forceinline__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset(
136 + int rowIdx, int colIdx, int32_t numKTiles, SFType* SFout) {
137 static_assert(CVT_FP4_NUM_THREADS_PER_SF == 1 ||
138 CVT_FP4_NUM_THREADS_PER_SF == 2);
139
140 // One pair of threads write one SF to global memory.
141 // TODO: stage through smem for packed STG.32
142 // 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;
143 + if (threadIdx.x % CVT_FP4_NUM_THREADS_PER_SF != 0) {
144 + return nullptr;
145 }
175 - return nullptr;
146 +
147 + // SF vector index (16 elements share one SF in the K dimension).
148 + int32_t kIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF;
149 + int32_t mIdx = rowIdx;
150 +
151 + // Decompose indices using bitwise ops (all divisors are powers of 2).
152 + // SF layout [numMTiles, numKTiles, 32 (mTile), 4 (mTile), 4(kTile)]
153 + int32_t mTileIdx = mIdx >> 7; // mIdx / 128
154 + int32_t outerMIdx = mIdx & 31; // mIdx % 32
155 + int32_t innerMIdx = (mIdx >> 5) & 3; // (mIdx / 32) % 4
156 + int32_t kTileIdx = kIdx >> 2; // kIdx / 4
157 + int32_t innerKIdx = kIdx & 3; // kIdx % 4
158 +
159 + // Compute global SF offset: mTileIdx * (numKTiles * 512) + kTileIdx * 512 +
160 + // outerMIdx * 16 + innerMIdx * 4 + innerKIdx
161 + // Use bitwise OR for non-overlapping lower bits.
162 + int64_t SFOffset = (static_cast<int64_t>(mTileIdx) * numKTiles + kTileIdx)
163 + << 9 |
164 + (outerMIdx << 4) | (innerMIdx << 2) | innerKIdx;
165 +
166 + return reinterpret_cast<uint8_t*>(SFout) + SFOffset;
167 }
168
169 // Quantizes the provided PackedVec into the uint32_t output