@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 37ece593c1

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 3, 09:15 UTC

[Perf] Padded nvfp4 quant kernel to remove additional copy, 2.4%~5.7% e2e performance improvement (#42774)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

Wentao Ye committed May 18, 2026 at 19:38 UTC 37ece593c105b5bb818aa94885617b863d390d7f

5 files changed +119 -43

csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu

+30 -24

@@ -38,8 +38,8 @@ namespace vllm {

38 // Use UE4M3 by default.

39 template <class Type, bool UE8M0_SF = false>

40 __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))

41 - cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, int32_t num_padded_cols,

42 - Type const* __restrict__ in,

41 + cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, int32_t outputCols,

42 + int32_t num_padded_cols, Type const* __restrict__ in,

43 float const* __restrict__ SFScale,

44 uint32_t* __restrict__ out, uint32_t* __restrict__ SFout) {

45 using PackedVec = vllm::PackedVec<Type, CVT_FP4_PACK16>;

@@ -50,7 +50,7 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))

50 "Vec size is not matched.");

52 // Precompute SF layout parameter (constant for entire kernel).

53 - int32_t const numKTiles = (numCols + 63) / 64;

53 + int32_t const numKTiles = (outputCols + 63) / 64;

55 int sf_m = round_up<int>(numRows, 128);

56 int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x;

@@ -68,16 +68,17 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))

68 PackedVec in_vec;

69 int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;

71 - // If we are outside valid rows OR outside valid columns -> Use Zeros

72 - bool valid = (rowIdx < numRows) && (elem_idx < numCols);

71 + // If we are outside valid columns, feed zeros

72 + bool valid_input = (rowIdx < numRows) && (elem_idx < numCols);

73 + bool valid_output = (rowIdx < numRows) && (elem_idx < outputCols);

74 if constexpr (CVT_FP4_PACK16) {

75 ld256_cg_or_zero(reinterpret_cast<u32x8_t&>(in_vec),

76 &reinterpret_cast<const uint32_t*>(in)[inOffset * 8],

76 - valid);

77 + valid_input);

78 } else {

79 ld128_cg_or_zero(reinterpret_cast<uint4&>(in_vec),

80 &reinterpret_cast<const uint32_t*>(in)[inOffset * 4],

80 - valid);

81 + valid_input);

82 }

84 auto sf_out =

@@ -89,16 +90,16 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))

90 cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(

91 in_vec, global_scale, sf_out);

92 - // We do NOT write output for padding because the 'out' tensor is not

93 - // padded.

94 - if (valid) {

93 + if (valid_output) {

94 if constexpr (CVT_FP4_PACK16) {

96 - int64_t outOffset = rowIdx * (numCols / 8) + colIdx * 2;

95 + int64_t outOffset = rowIdx * (outputCols / 8) + colIdx * 2;

96 uint64_t packed64 =

97 (uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo);

98 reinterpret_cast<uint64_t*>(out)[outOffset >> 1] = packed64;

99 } else {

101 - out[inOffset] = out_val;

100 + int64_t outOffset =

101 + rowIdx * (outputCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;

102 + out[outOffset] = out_val;

103 }

104 }

105 }

@@ -109,7 +110,8 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))

110 template <class Type, bool UE8M0_SF = false>

111 __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))

112 cvt_fp16_to_fp4_sf_major(int32_t numRows, int32_t numCols,

112 - int32_t sf_n_unpadded, int32_t num_packed_cols,

113 + int32_t outputCols, int32_t sf_n_unpadded,

114 + int32_t num_packed_cols,

115 Type const* __restrict__ in,

116 float const* __restrict__ SFScale,

117 uint32_t* __restrict__ out,

@@ -136,7 +138,7 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))

138 PackedVec in_vec;

139 int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;

140

139 - // If we are outside valid rows OR outside valid columns -> Use Zeros

141 + // If we are outside valid columns, feed zeros

142 bool valid = (rowIdx < numRows) && (elem_idx < numCols);

143 if constexpr (CVT_FP4_PACK16) {

144 ld256_cg_or_zero(reinterpret_cast<u32x8_t&>(in_vec),

@@ -155,16 +157,16 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))

157 cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(

158 in_vec, global_scale, sf_out);

159

158 - // We do NOT write output for padding because the 'out' tensor is not

159 - // padded.

160 - if (valid) {

160 + if (rowIdx < numRows) {

161 if constexpr (CVT_FP4_PACK16) {

162 - int64_t outOffset = rowIdx * (numCols / 8) + colIdx * 2;

162 + int64_t outOffset = rowIdx * (outputCols / 8) + colIdx * 2;

163 uint64_t packed64 =

164 (uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo);

165 reinterpret_cast<uint64_t*>(out)[outOffset >> 1] = packed64;

166 } else {

167 - out[inOffset] = out_val;

167 + int64_t outOffset =

168 + rowIdx * (outputCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;

169 + out[outOffset] = out_val;

170 }

171 }

172 }

@@ -180,8 +182,11 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,

182 bool is_sf_swizzled_layout) {

183 int32_t m = input.size(0);

184 int32_t n = input.size(1);

185 + int32_t output_n = output.size(1) * 2;

186

187 STD_TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");

188 + STD_TORCH_CHECK(output_n % 16 == 0,

189 + "The output tensor width must be a multiple of 16.");

190 STD_TORCH_CHECK(

191 input.scalar_type() == torch::headeronly::ScalarType::Half ||

192 input.scalar_type() == torch::headeronly::ScalarType::BFloat16,

@@ -197,7 +202,7 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,

202 input.get_device_index());

203 auto stream = get_current_cuda_stream(input.get_device_index());

204

200 - int sf_n_unpadded = int(n / CVT_FP4_SF_VEC_SIZE);

205 + int output_sf_n_unpadded = int(output_n / CVT_FP4_SF_VEC_SIZE);

206

207 // Grid, Block size. Each thread converts 8 values.

208 dim3 block(std::min(int(n / ELTS_PER_THREAD), 512));

@@ -205,7 +210,7 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,

210 vllm_runtime_blocks_per_sm(static_cast<int>(block.x));

211

212 if (is_sf_swizzled_layout) {

208 - int sf_n_int = int(vllm::round_up(sf_n_unpadded, 4) / 4);

213 + int sf_n_int = int(vllm::round_up(output_sf_n_unpadded, 4) / 4);

214 int32_t num_padded_cols =

215 sf_n_int * 4 * CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;

216

@@ -220,12 +225,12 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,

225 using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;

226 auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());

227 vllm::cvt_fp16_to_fp4<cuda_type, false><<<grid, block, 0, stream>>>(

223 - m, n, num_padded_cols, input_ptr, input_sf_ptr,

228 + m, n, output_n, num_padded_cols, input_ptr, input_sf_ptr,

229 reinterpret_cast<uint32_t*>(output_ptr),

230 reinterpret_cast<uint32_t*>(sf_out));

231 });

232 } else {

228 - int num_packed_cols = n / CVT_FP4_ELTS_PER_THREAD;

233 + int num_packed_cols = output_n / CVT_FP4_ELTS_PER_THREAD;

234 int grid_y = vllm::div_round_up(num_packed_cols, static_cast<int>(block.x));

235 int grid_x = std::min(

236 m, std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));

@@ -237,7 +242,8 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,

242 auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());

243 vllm::cvt_fp16_to_fp4_sf_major<cuda_type, false>

244 <<<grid, block, 0, stream>>>(

240 - m, n, sf_n_unpadded, num_packed_cols, input_ptr, input_sf_ptr,

245 + m, n, output_n, output_sf_n_unpadded, num_packed_cols,

246 + input_ptr, input_sf_ptr,

247 reinterpret_cast<uint32_t*>(output_ptr),

248 reinterpret_cast<uint32_t*>(sf_out));

249 });

tests/kernels/quantization/test_nvfp4_quant.py

+59

@@ -34,6 +34,7 @@ PAD_SHAPES = [

34 (64, 7152),

35 (32, 14336),

36 ]

37 +PADDED_OUTPUT_SHAPES = [(128, 48), (128, 80), (150, 48), (150, 80), (64, 7152)]

38 SEEDS = [42]

39 CUDA_DEVICES = ["cuda:0"]

@@ -130,6 +131,10 @@ def recover_swizzled_scales(scale, m, n):

131 return result[:m, :scale_n]

132

133

134 +def round_up(x: int, y: int) -> int:

135 + return (x + y - 1) // y * y

136 +

137 +

138 @pytest.mark.parametrize("dtype", DTYPES)

139 @pytest.mark.parametrize("shape", SHAPES)

140 @pytest.mark.parametrize("seed", SEEDS)

@@ -205,6 +210,60 @@ def test_python_util_matches_cpp_allocation(

210 )

211

212

213 +@pytest.mark.parametrize("shape", PADDED_OUTPUT_SHAPES)

214 +@pytest.mark.parametrize("is_sf_swizzled_layout", [True, False])

215 +@torch.inference_mode()

216 +def test_quantize_to_fp4_with_padded_output(

217 + shape: tuple[int, int],

218 + is_sf_swizzled_layout: bool,

219 +) -> None:

220 + from vllm._custom_ops import create_fp4_output_tensors

221 +

222 + dtype = torch.float16

223 + set_random_seed(42)

224 + torch.set_default_device("cuda:0")

225 +

226 + m, n = shape

227 + padded_n = round_up(n, 32)

228 + assert padded_n > n

229 +

230 + x = torch.randn((m, n), dtype=dtype)

231 + tensor_amax = torch.abs(x).max().to(torch.float32)

232 + global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax

233 + out_ref, scale_ref = ref_nvfp4_quant(x, global_scale)

234 +

235 + out, out_scale = ops.scaled_fp4_quant(

236 + x,

237 + global_scale,

238 + is_sf_swizzled_layout=is_sf_swizzled_layout,

239 + padded_n=padded_n,

240 + )

241 + py_out, py_scale = create_fp4_output_tensors(

242 + m,

243 + n,

244 + torch.device("cuda:0"),

245 + is_sf_swizzled_layout,

246 + padded_n=padded_n,

247 + )

248 +

249 + assert out.shape == (m, padded_n // 2)

250 + assert out.shape == py_out.shape

251 + assert out_scale.shape == py_scale.view(torch.float8_e4m3fn).shape

252 +

253 + out_ans = cast_from_fp4(out[:, : n // 2], m, n)

254 + torch.testing.assert_close(out_ans, out_ref)

255 + assert torch.count_nonzero(out[:, n // 2 :]) == 0

256 +

257 + if is_sf_swizzled_layout:

258 + scale_ans = recover_swizzled_scales(out_scale, m, padded_n)

259 + torch.testing.assert_close(scale_ans[:, : n // BLOCK_SIZE], scale_ref)

260 + assert torch.count_nonzero(scale_ans[:, n // BLOCK_SIZE :]) == 0

261 + else:

262 + scale_ans = out_scale.to(torch.float32)

263 + torch.testing.assert_close(scale_ans[:, : n // BLOCK_SIZE], scale_ref)

264 + assert torch.count_nonzero(scale_ans[:, n // BLOCK_SIZE :]) == 0

265 +

266 +

267 @pytest.mark.parametrize("pad_shape", PAD_SHAPES)

268 @torch.inference_mode()

269 def test_quantize_to_fp4_padded(pad_shape: tuple[int, int]) -> None:

vllm/_custom_ops.py

+24 -5

@@ -69,15 +69,20 @@ def create_fp4_output_tensors(

69 n: int,

70 device: torch.device,

71 is_sf_swizzled_layout: bool,

72 + padded_n: int | None = None,

73 ) -> tuple[torch.Tensor, torch.Tensor]:

74 """

75 Allocate both output tensors for scaled_fp4_quant:

76 (quantized_output, output_scale).

77 - Must match the C++ scaled_fp4_quant_func allocation exactly.

78 + Must match the C++ scaled_fp4_quant_func allocation exactly when

79 + ``padded_n`` is ``None``. When ``padded_n`` is provided, allocate a larger

80 + packed-FP4 output/scale buffer so the quantization kernel can write

81 + CUTLASS-compatible K padding directly

82 """

79 - output = torch.empty((m, n // 2), device=device, dtype=torch.uint8)

80 - output_scale = create_fp4_scale_tensor(m, n, device, is_sf_swizzled_layout)

83 + physical_n = padded_n if padded_n is not None else n

84 + output = torch.empty((m, physical_n // 2), device=device, dtype=torch.uint8)

85 + output_scale = create_fp4_scale_tensor(m, physical_n, device, is_sf_swizzled_layout)

86 return output, output_scale

@@ -1597,6 +1602,7 @@ def scaled_fp4_quant(

1602 input_global_scale: torch.Tensor,

1603 is_sf_swizzled_layout: bool = True,

1604 backend: str = "none",

1605 + padded_n: int | None = None,

1606 ) -> tuple[torch.Tensor, torch.Tensor]:

1607 """

1608 Quantize input tensor to FP4 and return quantized tensor and scale.

@@ -1611,6 +1617,8 @@ def scaled_fp4_quant(

1617 input: The input tensor to be quantized to FP4

1618 input_global_scale: A scalar scaling factor for the entire tensor.

1619 use_8x4_sf_layout: Whether to use the 8x4 or 128x4 layout for the scaling

1620 + padded_n: Optional padded K dimension. When provided, the quantized

1621 + output and scale tensors are allocated for ``padded_n``

1622

1623 Returns:

1624 tuple[torch.Tensor, torch.Tensor]: The output tensor in FP4 but every

@@ -1628,9 +1636,16 @@ def scaled_fp4_quant(

1636 assert input.dtype in (torch.float16, torch.bfloat16), (

1637 f"input.dtype needs to be fp16 or bf16 but got {input.dtype}."

1638 )

1639 + if padded_n is not None:

1640 + assert padded_n >= n, f"padded_n must be >= n, got padded_n={padded_n}, n={n}."

1641 + assert padded_n % block_size == 0, (

1642 + f"padded_n has to be a multiple of {block_size}, but got {padded_n}."

1643 + )

1644

1645 use_8x4_sf_layout = True if "trtllm" in backend and m <= 32 else False # noqa: SIM210

1633 -

1646 + if use_8x4_sf_layout and padded_n is not None and padded_n != n:

1647 + # TODO: support this case

1648 + raise ValueError("padded_n is not supported with TRTLLM 8x4 scale layout.")

1649 if use_8x4_sf_layout:

1650 output, output_scale = flashinfer_quant_nvfp4_8x4_sf_layout(

1651 input, input_global_scale

@@ -1638,7 +1653,11 @@ def scaled_fp4_quant(

1653 else:

1654 # Pre-allocate and call .out variant (same behavior as old in-place API)

1655 output, output_scale = create_fp4_output_tensors(

1641 - m, n, input.device, is_sf_swizzled_layout

1656 + m,

1657 + n,

1658 + input.device,

1659 + is_sf_swizzled_layout,

1660 + padded_n=padded_n,

1661 )

1662 torch.ops._C.scaled_fp4_quant.out(

1663 input,

vllm/model_executor/kernels/linear/nvfp4/cutlass.py

+2 -5

@@ -9,7 +9,6 @@ from vllm._custom_ops import (

9 )

10 from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (

11 cutlass_fp4_supported,

12 - pad_nvfp4_activation_for_cutlass,

12 pad_nvfp4_weight_for_cutlass,

13 slice_nvfp4_output,

14 swizzle_blockscale,

@@ -52,16 +51,14 @@ class CutlassNvFp4LinearKernel(NvFp4LinearKernel):

51 output_size = layer.output_size_per_partition

52 output_dtype = x.dtype

53 output_shape = [*x.shape[:-1], output_size]

54 + weights_padding_bytes = getattr(layer, "weights_padding_cols", 0)

56 x_fp4, x_blockscale = scaled_fp4_quant(

57 x,

58 layer.input_global_scale_inv,

59 is_sf_swizzled_layout=True,

60 backend="cutlass",

61 - )

62 -

63 - x_fp4 = pad_nvfp4_activation_for_cutlass(

64 - x_fp4, getattr(layer, "weights_padding_cols", 0)

61 + padded_n=x.shape[-1] + weights_padding_bytes * 2,

62 )

64 out = cutlass_scaled_fp4_mm(

vllm/model_executor/kernels/linear/nvfp4/flashinfer.py

+4 -9

@@ -5,7 +5,6 @@ import torch

6 from vllm._custom_ops import scaled_fp4_quant

7 from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (

8 - pad_nvfp4_activation_for_cutlass,

8 pad_nvfp4_weight_for_cutlass,

9 slice_nvfp4_output,

10 swizzle_blockscale,

@@ -58,16 +57,14 @@ class FlashInferCutlassNvFp4LinearKernel(NvFp4LinearKernel):

57 output_size = layer.output_size_per_partition

58 output_dtype = x.dtype

59 output_shape = [*x.shape[:-1], output_size]

60 + weights_padding_bytes = getattr(layer, "weights_padding_cols", 0)

62 x_fp4, x_blockscale = scaled_fp4_quant(

63 x,

64 layer.input_global_scale_inv,

65 is_sf_swizzled_layout=True,

66 backend="flashinfer-cutlass",

67 - )

68 -

69 - x_fp4 = pad_nvfp4_activation_for_cutlass(

70 - x_fp4, getattr(layer, "weights_padding_cols", 0)

67 + padded_n=x.shape[-1] + weights_padding_bytes * 2,

68 )

70 out = flashinfer_scaled_fp4_mm(

@@ -189,16 +186,14 @@ class FlashInferCudnnNvFp4LinearKernel(NvFp4LinearKernel):

186 output_size = layer.output_size_per_partition

187 output_dtype = x.dtype

188 output_shape = [*x.shape[:-1], output_size]

189 + weights_padding_bytes = getattr(layer, "weights_padding_cols", 0)

190

191 x_fp4, x_blockscale = scaled_fp4_quant(

192 x,

193 layer.input_global_scale_inv,

194 is_sf_swizzled_layout=True,

195 backend="flashinfer-cudnn",

198 - )

199 -

200 - x_fp4 = pad_nvfp4_activation_for_cutlass(

201 - x_fp4, getattr(layer, "weights_padding_cols", 0)

196 + padded_n=x.shape[-1] + weights_padding_bytes * 2,

197 )

198

199 out = flashinfer_scaled_fp4_mm(