[Perf][Kernel] Initialize NVFP4 padding in quant kernel (#53568)
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Co-authored-by: Codex <codex@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Roberto L. Castro committed
Sep 1, 2026 at 18:49 UTC
514c7314a06be9e43d00d3f126a5aaa54211274f
3 files changed
+91
-27
csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu
+48
-19
@@ -36,7 +36,7 @@
36
namespace vllm {
37
38
// Use UE4M3 by default.
39
-template <class Type, bool UE8M0_SF = false>
39
+template <class Type, bool UE8M0_SF = false, bool PADDED_OUTPUT = 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 outputCols,
42
int32_t num_padded_cols, Type const* __restrict__ in,
@@ -55,7 +55,7 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
55
#endif
56
57
// Precompute SF layout parameter (constant for entire kernel).
58
- int32_t const numKTiles = (outputCols + 63) / 64;
58
+ int32_t const numKTiles = ((PADDED_OUTPUT ? outputCols : numCols) + 63) / 64;
59
60
int sf_m = round_up<int>(numRows, 128);
61
int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x;
@@ -74,16 +74,15 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
74
int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
75
76
// If we are outside valid columns, feed zeros
77
- bool valid_input = (rowIdx < numRows) && (elem_idx < numCols);
78
- bool valid_output = (rowIdx < numRows) && (elem_idx < outputCols);
77
+ bool valid = (rowIdx < numRows) && (elem_idx < numCols);
78
if constexpr (CVT_FP4_PACK16) {
79
ld256_cg_or_zero(reinterpret_cast<u32x8_t&>(in_vec),
80
&reinterpret_cast<const uint32_t*>(in)[inOffset * 8],
82
- valid_input);
81
+ valid);
82
} else {
83
ld128_cg_or_zero(reinterpret_cast<uint4&>(in_vec),
84
&reinterpret_cast<const uint32_t*>(in)[inOffset * 4],
86
- valid_input);
85
+ valid);
86
}
87
88
auto sf_out =
@@ -91,20 +90,42 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
90
CVT_FP4_NUM_THREADS_PER_SF>(
91
rowIdx, colIdx, numKTiles, SFout);
92
94
- auto out_val =
95
- cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
96
- in_vec, global_scale, sf_out);
93
+ fp4_packed_t out_val{};
94
+ if (valid) {
95
+ out_val =
96
+ cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
97
+ in_vec, global_scale, sf_out);
98
+ } else if (sf_out != nullptr) {
99
+ // The swizzled scale tensor is allocated with torch.empty. Explicitly
100
+ // initialize padded rows and columns in this kernel so downstream
101
+ // NVFP4 GEMMs never observe stale scale factors.
102
+ *sf_out = 0;
103
+ }
104
98
- if (valid_output) {
105
+ if constexpr (PADDED_OUTPUT) {
106
+ bool valid_output = (rowIdx < numRows) && (elem_idx < outputCols);
107
+ if (valid_output) {
108
+ if constexpr (CVT_FP4_PACK16) {
109
+ int64_t outOffset = rowIdx * (outputCols / 8) + colIdx * 2;
110
+ uint64_t packed64 =
111
+ (uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo);
112
+ reinterpret_cast<uint64_t*>(out)[outOffset >> 1] = packed64;
113
+ } else {
114
+ int64_t outOffset =
115
+ rowIdx * (outputCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
116
+ reinterpret_cast<fp4_packed_t*>(out)[outOffset] = out_val;
117
+ }
118
+ }
119
+ } else if (valid) {
120
+ // #32520 fast path: input and output widths are identical, so avoid
121
+ // the padded-output predicate and use the original address arithmetic.
122
if constexpr (CVT_FP4_PACK16) {
100
- int64_t outOffset = rowIdx * (outputCols / 8) + colIdx * 2;
123
+ int64_t outOffset = rowIdx * (numCols / 8) + colIdx * 2;
124
uint64_t packed64 =
125
(uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo);
126
reinterpret_cast<uint64_t*>(out)[outOffset >> 1] = packed64;
127
} else {
105
- int64_t outOffset =
106
- rowIdx * (outputCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
107
- out[outOffset] = out_val;
128
+ reinterpret_cast<fp4_packed_t*>(out)[inOffset] = out_val;
129
}
130
}
131
}
@@ -244,11 +265,19 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
265
attrs[0].val.programmaticStreamSerializationAllowed = 1;
266
config.numAttrs = 1;
267
config.attrs = attrs;
247
- cudaLaunchKernelEx(&config, vllm::cvt_fp16_to_fp4<cuda_type, false>,
248
- m, n, output_n, num_padded_cols, input_ptr,
249
- input_sf_ptr,
250
- reinterpret_cast<uint32_t*>(output_ptr),
251
- reinterpret_cast<uint32_t*>(sf_out));
268
+ if (output_n == n) {
269
+ cudaLaunchKernelEx(
270
+ &config, vllm::cvt_fp16_to_fp4<cuda_type, false, false>, m, n,
271
+ output_n, num_padded_cols, input_ptr, input_sf_ptr,
272
+ reinterpret_cast<uint32_t*>(output_ptr),
273
+ reinterpret_cast<uint32_t*>(sf_out));
274
+ } else {
275
+ cudaLaunchKernelEx(
276
+ &config, vllm::cvt_fp16_to_fp4<cuda_type, false, true>, m, n,
277
+ output_n, num_padded_cols, input_ptr, input_sf_ptr,
278
+ reinterpret_cast<uint32_t*>(output_ptr),
279
+ reinterpret_cast<uint32_t*>(sf_out));
280
+ }
281
});
282
} else {
283
int num_packed_cols = output_n / CVT_FP4_ELTS_PER_THREAD;
tests/kernels/quantization/test_nvfp4_quant.py
+40
-1
@@ -119,7 +119,7 @@ def ref_nvfp4_quant(x, global_scale):
119
return cast_to_fp4(clipped_x), scale.squeeze(-1)
120
121
122
-def recover_swizzled_scales(scale, m, n):
122
+def recover_swizzled_scales(scale, m, n, include_padding=False):
123
round_up = lambda x, y: (x + y - 1) // y * y
124
rounded_m = round_up(m, 128)
125
scale_n = n // BLOCK_SIZE
@@ -128,6 +128,8 @@ def recover_swizzled_scales(scale, m, n):
128
tmp = torch.reshape(scale, (1, rounded_m // 128, rounded_n // 4, 32, 4, 4))
129
tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5))
130
result = torch.reshape(tmp, (rounded_m, rounded_n)).to(torch.float32)
131
+ if include_padding:
132
+ return result
133
return result[:m, :scale_n]
134
135
@@ -264,6 +266,43 @@ def test_quantize_to_fp4_with_padded_output(
266
assert torch.count_nonzero(scale_ans[:, n // BLOCK_SIZE :]) == 0
267
268
269
+@pytest.mark.parametrize(
270
+ ("shape", "padded_n"),
271
+ [((90, 80), None), ((90, 48), 64)],
272
+)
273
+@torch.inference_mode()
274
+def test_swizzled_scale_padding_initialized_in_kernel(
275
+ shape: tuple[int, int], padded_n: int | None
276
+) -> None:
277
+ """The out variant must overwrite dirty scale padding without a memset."""
278
+ torch.set_default_device("cuda:0")
279
+ m, n = shape
280
+ physical_n = padded_n if padded_n is not None else n
281
+ rounded_m = round_up(m, 128)
282
+ rounded_scale_n = round_up(physical_n // BLOCK_SIZE, 4)
283
+
284
+ x = torch.randn((m, n), dtype=torch.float16)
285
+ global_scale = torch.tensor(1.0, dtype=torch.float32)
286
+ out = torch.full((m, physical_n // 2), 0xFF, dtype=torch.uint8)
287
+ out_scale = torch.full(
288
+ (rounded_m, rounded_scale_n // 4), 0x7F7F7F7F, dtype=torch.int32
289
+ )
290
+
291
+ torch.ops._C.scaled_fp4_quant.out(
292
+ x,
293
+ global_scale,
294
+ True,
295
+ output=out,
296
+ output_scale=out_scale,
297
+ )
298
+
299
+ full_scale = recover_swizzled_scales(
300
+ out_scale.view(torch.float8_e4m3fn), m, physical_n, include_padding=True
301
+ )
302
+ assert torch.count_nonzero(full_scale[m:, :]) == 0
303
+ assert torch.count_nonzero(full_scale[:m, n // BLOCK_SIZE :]) == 0
304
+
305
+
306
@pytest.mark.parametrize("pad_shape", PAD_SHAPES)
307
@torch.inference_mode()
308
def test_quantize_to_fp4_padded(pad_shape: tuple[int, int]) -> None:
vllm/_custom_ops.py
+3
-7
@@ -57,13 +57,9 @@ def create_fp4_scale_tensor(
57
rounded_m = round_up(m, 128)
58
scale_n = n // block_size
59
rounded_n = round_up(scale_n, 4)
60
- # Must be zero-initialized: the swizzled scale buffer is padded to
61
- # (round_up(m, 128), round_up(scale_n, 4) // 4) but the NVFP4 quant
62
- # kernel does not write every padded element that the downstream
63
- # NVFP4 GEMM reads. torch.empty leaves those padded scale factors
64
- # uninitialized, which corrupts dequantization and causes a severe
65
- # Blackwell NVFP4 decode throughput/output-length regression.
66
- return torch.zeros(
60
+ # The NVFP4 quant kernel explicitly zeroes every padded scale entry,
61
+ # so no separate zero-initialization kernel is required here.
62
+ return torch.empty(
63
(rounded_m, rounded_n // 4), device=device, dtype=torch.int32
64
)
65
else: