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, 00:00 UTC
[FIX] FP4 quantization kernel padding initialization bug (#31097)
Signed-off-by: <> Co-authored-by: root <root@gpu-193.slurm-workers-slurm.slurm.svc.cluster.local> Co-authored-by: root <root@gpu-951.slurm-workers-slurm.slurm.svc.cluster.local>
danielafrimi committed
Dec 23, 2025 at 18:45 UTC
b94f80ffb8f279b97931bdc2d39bd9328503d7d1
1 file changed
+28
-18
csrc/quantization/fp4/nvfp4_quant_kernels.cu
+28
-18
@@ -35,7 +35,7 @@ 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;
38
+ return ((x + y - 1) / y) * y;
39
}
40
41
// Compute effective rows for grid configuration with swizzled SF layouts.
@@ -61,37 +61,47 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
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;
64
- for (int row = numRows + blockIdx.x; row < sf_m; row += gridDim.x) {
65
- // Each thread writes 4 uint32_t elements.
66
- for (int col = sf_n_unpadded + threadIdx.x * 4; col < sf_n_int;
67
- col += blockDim.x * 4) {
68
- SFout[row * sf_n_int + col] = 0x00;
69
- }
70
- }
64
+ int num_padded_cols = sf_n_int * 4 * CVT_FP4_SF_VEC_SIZE;
65
66
// Get the global scaling factor, which will be applied to the SF.
67
// Note SFScale is the same as next GEMM's alpha, which is
68
// (448.f / (Alpha_A / 6.f)).
69
float const global_scale = SFScale == nullptr ? 1.0f : SFScale[0];
70
77
- // Input tensor row/col loops.
78
- for (int rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) {
79
- for (int colIdx = threadIdx.x; colIdx < numCols / CVT_FP4_ELTS_PER_THREAD;
71
+ // Iterate over all rows and cols including padded ones -
72
+ // ensures we visit every single scale factor address to initialize it.
73
+ 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
+
79
+ PackedVec in_vec;
80
int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
82
- PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
83
- // Get the output tensor offset.
84
- // Same as inOffset because 8 elements are packed into one uint32_t.
85
- int64_t outOffset = inOffset;
86
- auto& out_pos = out[outOffset];
81
+
82
+ // 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
+
86
+ } else {
87
+ // Valid Region: Load actual data
88
+ in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
89
+ }
90
91
auto sf_out =
92
cvt_quant_to_fp4_get_sf_out_offset<uint32_t,
93
CVT_FP4_NUM_THREADS_PER_SF>(
94
rowIdx, colIdx, numKTiles, SFout);
95
93
- out_pos =
96
+ auto out_val =
97
cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, global_scale, sf_out);
98
+
99
+ // We do NOT write output for padding because the 'out' tensor is not
100
+ // 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;
104
+ }
105
}
106
}
107
}
@@ -134,4 +144,4 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
144
m, n, input_ptr, input_sf_ptr, reinterpret_cast<uint32_t*>(output_ptr),
145
reinterpret_cast<uint32_t*>(sf_out));
146
});
137
-}
147
+}
\ No newline at end of file