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, 14:45 UTC
[FIX] Always support TP > 4 for FP4 Gemm (#31099)
[FIX] Always support TP > 4 for FP4 Gemm (#31099) Signed-off-by: dafrimi <dafrimi@nvidia.com> Co-authored-by: root <root@gpu-51.slurm-workers-slurm.slurm.svc.cluster.local>
danielafrimi committed
Jan 26, 2026 at 20:04 UTC
67fe677c53e7a5f5d27a4f07edd66f44bea0495c
3 files changed
+114
-6
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
+26
-4
@@ -16,6 +16,9 @@ from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import
16
)
17
from vllm.model_executor.layers.quantization.utils.quant_utils import (
18
cutlass_fp4_supported,
19
+ pad_nvfp4_activation_for_cutlass,
20
+ pad_nvfp4_weight_for_cutlass,
21
+ slice_nvfp4_output,
22
swizzle_blockscale,
23
)
24
from vllm.model_executor.parameter import (
@@ -159,9 +162,20 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
162
if self.backend == "fbgemm":
163
swizzled_weight_scale = swizzled_weight_scale.view(-1).view(torch.uint8)
164
layer.weight_scale = Parameter(swizzled_weight_scale, requires_grad=False)
162
- layer.weight_packed = Parameter(
163
- layer.weight_packed.data, requires_grad=False
164
- )
165
+
166
+ # Pad weights for CUTLASS/FlashInfer kernel alignment (K and N
167
+ # divisible by 32). fbgemm has its own layout requirements.
168
+ if self.backend in ("cutlass", "flashinfer-cutlass"):
169
+ weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
170
+ layer.weight_packed.data
171
+ )
172
+ layer.weights_padding_cols = weights_padding_cols
173
+ layer.weight_packed = Parameter(weight, requires_grad=False)
174
+ else:
175
+ layer.weights_padding_cols = 0
176
+ layer.weight_packed = Parameter(
177
+ layer.weight_packed.data, requires_grad=False
178
+ )
179
180
layer.alpha = Parameter(
181
1 / (layer.input_global_scale * layer.weight_global_scale),
@@ -187,7 +201,8 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
201
return out
202
203
output_dtype = x.dtype
190
- output_shape = [*x.shape[:-1], layer.weight_packed.shape[0]]
204
+ output_size = layer.output_size_per_partition
205
+ output_shape = [*x.shape[:-1], output_size]
206
207
# quantize BF16 or FP16 to (FP4 and interleaved block scale)
208
x_fp4, x_blockscale = scaled_fp4_quant(
@@ -197,6 +212,10 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
212
backend=self.backend,
213
)
214
215
+ # Pad activations to match weight K-dimension padding
216
+ weights_padding_cols = getattr(layer, "weights_padding_cols", 0)
217
+ x_fp4 = pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_cols)
218
+
219
mm_args = (
220
x_fp4,
221
layer.weight_packed,
@@ -221,6 +240,9 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
240
assert self.backend == "cutlass"
241
out = cutlass_scaled_fp4_mm(*mm_args)
242
243
+ # Slice output to remove N-dimension padding
244
+ out = slice_nvfp4_output(out, output_size)
245
+
246
if bias is not None:
247
out = out + bias
248
return out.view(*output_shape)
vllm/model_executor/layers/quantization/modelopt.py
+21
-2
@@ -84,6 +84,9 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
84
kFp8StaticTokenSym,
85
kNvfp4Dynamic,
86
kNvfp4Static,
87
+ pad_nvfp4_activation_for_cutlass,
88
+ pad_nvfp4_weight_for_cutlass,
89
+ slice_nvfp4_output,
90
swizzle_blockscale,
91
)
92
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
@@ -1280,9 +1283,16 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
1283
layer.weight_scale = Parameter(weight_scale, requires_grad=False)
1284
layer.weight = Parameter(weight, requires_grad=False)
1285
else:
1286
+ # Swizzle block scales and pad the packed NVFP4 weights for kernel
1287
+ # alignment (CUTLASS/FlashInfer require K and N divisible by 32).
1288
swizzled_weight_scale = swizzle_blockscale(layer.weight_scale)
1289
layer.weight_scale = Parameter(swizzled_weight_scale, requires_grad=False)
1285
- layer.weight = Parameter(layer.weight.data, requires_grad=False)
1290
+
1291
+ weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
1292
+ layer.weight.data
1293
+ )
1294
+ layer.weights_padding_cols = weights_padding_cols
1295
+ layer.weight = Parameter(weight, requires_grad=False)
1296
1297
def apply(
1298
self,
@@ -1304,7 +1314,6 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
1314
)
1315
1316
output_dtype = x.dtype
1307
- output_shape = [x.shape[0], layer.weight.shape[0]]
1317
1318
# quantize BF16 or FP16 to (FP4 and interleaved block scale)
1319
x_fp4, x_blockscale = scaled_fp4_quant(
@@ -1319,6 +1328,12 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
1328
assert layer.weight_scale.dtype == torch.float8_e4m3fn
1329
assert layer.alpha.dtype == torch.float32
1330
1331
+ # Pad activations to match weight K-dimension padding
1332
+ weights_padding_cols = getattr(layer, "weights_padding_cols", 0)
1333
+ output_size = layer.output_size_per_partition
1334
+ output_shape = [x.shape[0], output_size]
1335
+ x_fp4 = pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_cols)
1336
+
1337
mm_args = (
1338
x_fp4,
1339
layer.weight,
@@ -1327,6 +1342,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
1342
layer.alpha,
1343
output_dtype,
1344
)
1345
+
1346
if self.backend.startswith("flashinfer-"):
1347
backend_name = self.backend[len("flashinfer-") :]
1348
out = flashinfer_scaled_fp4_mm(*mm_args, backend=backend_name)
@@ -1334,6 +1350,9 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
1350
assert self.backend == "cutlass"
1351
out = cutlass_scaled_fp4_mm(*mm_args)
1352
1353
+ # Slice output to remove N-dimension padding
1354
+ out = slice_nvfp4_output(out, output_size)
1355
+
1356
if bias is not None:
1357
out = out + bias
1358
return out.view(*output_shape)
vllm/model_executor/layers/quantization/utils/quant_utils.py
+67
@@ -868,3 +868,70 @@ def convert_packed_uint4b8_to_signed_int4_inplace(t: torch.Tensor) -> torch.Tens
868
t |= ((nib - 8) & 0xF) << shift
869
870
return t
871
+
872
+
873
+def round_up(x: int, m: int) -> int:
874
+ """Round up x to the nearest multiple of m."""
875
+ return (x + m - 1) // m * m
876
+
877
+
878
+def pad_nvfp4_weight_for_cutlass(
879
+ weight: torch.Tensor,
880
+ alignment: int = 32,
881
+) -> tuple[torch.Tensor, int]:
882
+ """
883
+ Pad packed NVFP4 weights so that both N (rows) and K (columns) satisfy
884
+ the alignment constraints required by CUTLASS / FlashInfer FP4 kernels.
885
+
886
+ CUTLASS FP4 kernel requires both K and N matrix dimensions to be divisible
887
+ by 32 for aligned memory access and efficient tensor core operations.
888
+ """
889
+ weight_current_rows = weight.shape[0]
890
+
891
+ # Pad N dimension (rows) if not aligned
892
+ if weight_current_rows % alignment != 0:
893
+ total_rows = round_up(weight_current_rows, alignment)
894
+ pad_rows = total_rows - weight_current_rows
895
+ weight = torch.nn.functional.pad(weight, (0, 0, 0, pad_rows)).contiguous()
896
+
897
+ # Check K dimension alignment
898
+ # 2 FP4 items are packed per byte in the input dimension
899
+ weight_current_col_bytes = weight.shape[1]
900
+ weight_current_col_elements = weight_current_col_bytes * 2
901
+
902
+ weights_padding_bytes = 0
903
+ if weight_current_col_elements % alignment != 0:
904
+ total_cols = round_up(weight_current_col_elements, alignment)
905
+ pad_cols = total_cols - weight_current_col_elements
906
+ # Convert from FP4 element count to bytes (2 FP4 values per byte)
907
+ # pad_cols is always even since alignment=32 and current elements are even
908
+ pad_bytes = pad_cols // 2
909
+ weight = torch.nn.functional.pad(weight, (0, pad_bytes, 0, 0)).contiguous()
910
+ weights_padding_bytes = pad_bytes
911
+
912
+ return weight, weights_padding_bytes
913
+
914
+
915
+def pad_nvfp4_activation_for_cutlass(
916
+ x_fp4: torch.Tensor,
917
+ weights_padding_bytes: int,
918
+) -> torch.Tensor:
919
+ """
920
+ Pad packed FP4 activations to match the K-dimension padding applied to weights.
921
+ The padding is in bytes (tensor dimension), not FP4 elements.
922
+ """
923
+ if weights_padding_bytes > 0:
924
+ return torch.nn.functional.pad(x_fp4, (0, weights_padding_bytes)).contiguous()
925
+ return x_fp4
926
+
927
+
928
+def slice_nvfp4_output(
929
+ out: torch.Tensor,
930
+ output_size: int,
931
+) -> torch.Tensor:
932
+ """
933
+ Slice the output tensor to remove padding in N dimension if weight was padded.
934
+ """
935
+ if out.shape[-1] != output_size:
936
+ return out[..., :output_size].contiguous()
937
+ return out