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 14, 05:30 UTC
Fix nvfp4 swizzling (#23140)
Signed-off-by: yiliu30 <yi4.liu@intel.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Yi Liu committed
Aug 22, 2025 at 00:54 UTC
0278f1ac3ac90c245a9a000cb11eb3146103cfd1
2 files changed
+5
-26
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
+3
-24
@@ -12,6 +12,8 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
12
CompressedTensorsScheme)
13
from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( # noqa: E501
14
run_nvfp4_emulations)
15
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
16
+ swizzle_blockscale)
17
from vllm.model_executor.parameter import (GroupQuantScaleParameter,
18
ModelWeightParameter,
19
PerTensorScaleParameter)
@@ -83,29 +85,6 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
85
weight_loader=weight_loader)
86
layer.register_parameter("input_global_scale", input_global_scale)
87
86
- def swizzle_blockscale(self, scale: torch.tensor):
87
- assert (scale.dtype == torch.float8_e4m3fn)
88
- # Pad and blockwise interleave weight_scale
89
- scale_ndim = scale.ndim
90
- if scale.ndim == 2:
91
- scale = scale.unsqueeze(0)
92
- assert scale.ndim == 3
93
- B, M, K = scale.shape
94
- round_up_multiple = lambda x, m: (x + m - 1) // m * m
95
- M_padded = round_up_multiple(M, 128)
96
- K_padded = round_up_multiple(K, 4)
97
- padded_scale = torch.zeros((B, M_padded, K_padded), dtype=scale.dtype)
98
- padded_scale[:B, :M, :K] = scale
99
- batches, rows, cols = padded_scale.shape
100
- assert rows % 128 == 0
101
- assert cols % 4 == 0
102
- padded_scale = padded_scale.reshape(batches, rows // 128, 4, 32,
103
- cols // 4, 4)
104
- swizzled_scale = padded_scale.permute((0, 1, 4, 3, 2, 5))
105
- swizzled_scale = swizzled_scale.contiguous().cuda()
106
- return (swizzled_scale.reshape(M, K)
107
- if scale_ndim == 2 else swizzled_scale.reshape(B, M, K))
108
-
88
def process_weights_after_loading(self, layer) -> None:
89
90
global_input_scale = layer.input_global_scale.max().to(torch.float32)
@@ -137,7 +116,7 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
116
requires_grad=False)
117
layer.weight_packed = Parameter(weight, requires_grad=False)
118
else:
140
- swizzled_weight_scale = self.swizzle_blockscale(layer.weight_scale)
119
+ swizzled_weight_scale = swizzle_blockscale(layer.weight_scale)
120
layer.weight_scale_swizzled = Parameter(swizzled_weight_scale,
121
requires_grad=False)
122
layer.weight_packed = Parameter(layer.weight_packed.data,
vllm/model_executor/layers/quantization/utils/quant_utils.py
+2
-2
@@ -552,8 +552,8 @@ def swizzle_blockscale(scale: torch.Tensor) -> torch.Tensor:
552
swizzled = padded.permute(0, 1, 4, 3, 2, 5).contiguous().cuda()
553
554
if scale_ndim == 2:
555
- return swizzled.reshape(M, K)
556
- return swizzled.reshape(B, M, K)
555
+ return swizzled.reshape(M_padded, K_padded)
556
+ return swizzled.reshape(B, M_padded, K_padded)
557
558
559
def cutlass_fp4_supported() -> bool: