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
[Perf] Remove duplicated NVFP4 blockscales to save memory (#23379)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Aug 26, 2025 at 07:16 UTC
d52358c1e07768266e3db92e847cd28af87ca4b9
3 files changed
+30
-35
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+10
-10
@@ -246,13 +246,13 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
246
return
247
248
# swizzle weight scales
249
- layer.w13_blockscale_swizzled = torch.nn.Parameter(swizzle_blockscale(
249
+ layer.w13_weight_scale = torch.nn.Parameter(swizzle_blockscale(
250
layer.w13_weight_scale),
251
- requires_grad=False)
251
+ requires_grad=False)
252
253
- layer.w2_blockscale_swizzled = torch.nn.Parameter(swizzle_blockscale(
253
+ layer.w2_weight_scale = torch.nn.Parameter(swizzle_blockscale(
254
layer.w2_weight_scale),
255
- requires_grad=False)
255
+ requires_grad=False)
256
257
# w13
258
w13_input_global_scale = layer.w13_input_global_scale.max(
@@ -383,8 +383,8 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
383
activation=activation,
384
global_num_experts=global_num_experts,
385
expert_map=expert_map,
386
- w1_scale=layer.w13_blockscale_swizzled,
387
- w2_scale=layer.w2_blockscale_swizzled,
386
+ w1_scale=layer.w13_weight_scale,
387
+ w2_scale=layer.w2_weight_scale,
388
apply_router_weight_on_input=apply_router_weight_on_input,
389
)
390
@@ -406,8 +406,8 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
406
activation=activation,
407
global_num_experts=global_num_experts,
408
expert_map=expert_map,
409
- w1_scale=layer.w13_blockscale_swizzled,
410
- w2_scale=layer.w2_blockscale_swizzled,
409
+ w1_scale=layer.w13_weight_scale,
410
+ w2_scale=layer.w2_weight_scale,
411
g1_alphas=layer.g1_alphas,
412
g2_alphas=layer.g2_alphas,
413
a1_gscale=layer.w13_input_scale_quant,
@@ -427,8 +427,8 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
427
a=x,
428
w1_fp4=layer.w13_weight,
429
w2_fp4=layer.w2_weight,
430
- w1_blockscale=layer.w13_blockscale_swizzled,
431
- w2_blockscale=layer.w2_blockscale_swizzled,
430
+ w1_blockscale=layer.w13_weight_scale,
431
+ w2_blockscale=layer.w2_weight_scale,
432
g1_alphas=layer.g1_alphas,
433
g2_alphas=layer.g2_alphas,
434
a1_gscale=layer.w13_input_scale_quant,
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
+5
-6
@@ -112,13 +112,12 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
112
torch.uint8), epilogue_tile_m).reshape(
113
weight_scale.shape).view(torch.float8_e4m3fn))
114
115
- layer.weight_scale_swizzled = Parameter(weight_scale,
116
- requires_grad=False)
115
+ layer.weight_scale = Parameter(weight_scale, requires_grad=False)
116
layer.weight_packed = Parameter(weight, requires_grad=False)
117
else:
118
swizzled_weight_scale = swizzle_blockscale(layer.weight_scale)
120
- layer.weight_scale_swizzled = Parameter(swizzled_weight_scale,
121
- requires_grad=False)
119
+ layer.weight_scale = Parameter(swizzled_weight_scale,
120
+ requires_grad=False)
121
layer.weight_packed = Parameter(layer.weight_packed.data,
122
requires_grad=False)
123
@@ -136,7 +135,7 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
135
x=x,
136
input_global_scale=layer.input_global_scale,
137
weight=layer.weight_packed,
139
- weight_scale_swizzled=layer.weight_scale_swizzled,
138
+ weight_scale_swizzled=layer.weight_scale,
139
weight_global_scale=layer.weight_global_scale)
140
if bias is not None:
141
out = out + bias
@@ -149,7 +148,7 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
148
x_fp4, x_blockscale = scaled_fp4_quant(x, layer.input_global_scale)
149
150
mm_args = (x_fp4, layer.weight_packed, x_blockscale,
152
- layer.weight_scale_swizzled, layer.alpha, output_dtype)
151
+ layer.weight_scale, layer.alpha, output_dtype)
152
if self.backend == "flashinfer-trtllm":
153
out = flashinfer_scaled_fp4_mm(*mm_args, backend="trtllm")
154
elif self.backend == "flashinfer-cutlass":
vllm/model_executor/layers/quantization/modelopt.py
+15
-19
@@ -907,20 +907,18 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
907
torch.uint8), epilogue_tile_m).reshape(
908
weight_scale.shape).view(torch.float8_e4m3fn))
909
910
- layer.weight_scale_swizzled = Parameter(weight_scale,
911
- requires_grad=False)
910
+ layer.weight_scale = Parameter(weight_scale, requires_grad=False)
911
layer.weight = Parameter(weight, requires_grad=False)
912
else:
913
swizzled_weight_scale = swizzle_blockscale(layer.weight_scale)
915
- layer.weight_scale_swizzled = Parameter(swizzled_weight_scale,
916
- requires_grad=False)
914
+ layer.weight_scale = Parameter(swizzled_weight_scale,
915
+ requires_grad=False)
916
layer.weight = Parameter(layer.weight.data, requires_grad=False)
917
918
if self.backend == "marlin":
919
prepare_fp4_layer_for_marlin(layer)
920
del layer.alpha
921
del layer.input_scale
923
- del layer.weight_scale_swizzled
922
923
def apply(
924
self,
@@ -951,14 +949,14 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
949
assert (x_fp4.dtype == torch.uint8)
950
assert (layer.weight.dtype == torch.uint8)
951
assert (x_blockscale.dtype == torch.float8_e4m3fn)
954
- assert (layer.weight_scale_swizzled.dtype == torch.float8_e4m3fn)
952
+ assert (layer.weight_scale.dtype == torch.float8_e4m3fn)
953
assert (layer.alpha.dtype == torch.float32)
954
955
mm_args = (
956
x_fp4,
957
layer.weight,
958
x_blockscale,
961
- layer.weight_scale_swizzled,
959
+ layer.weight_scale,
960
layer.alpha,
961
output_dtype,
962
)
@@ -1320,16 +1318,16 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1318
"Weight Blockscale must be represented as FP8-E4M3")
1319
w13_blockscale_swizzled = swizzle_blockscale(
1320
layer.w13_weight_scale)
1323
- layer.w13_blockscale_swizzled = Parameter(w13_blockscale_swizzled,
1324
- requires_grad=False)
1321
+ layer.w13_weight_scale = Parameter(w13_blockscale_swizzled,
1322
+ requires_grad=False)
1323
1324
assert (layer.w2_weight_scale.shape[2] % 16 == 0), (
1325
"Expected weight_scale.dim(1) to be divisible by 16")
1326
assert (layer.w2_weight_scale.dtype == torch.float8_e4m3fn), (
1327
"Weight Blockscale must be represented as FP8-E4M3")
1328
w2_blockscale_swizzled = swizzle_blockscale(layer.w2_weight_scale)
1331
- layer.w2_blockscale_swizzled = Parameter(w2_blockscale_swizzled,
1332
- requires_grad=False)
1329
+ layer.w2_weight_scale = Parameter(w2_blockscale_swizzled,
1330
+ requires_grad=False)
1331
layer.w2_weight = Parameter(layer.w2_weight.data,
1332
requires_grad=False)
1333
@@ -1339,8 +1337,6 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1337
del layer.g2_alphas
1338
del layer.w13_input_scale_quant
1339
del layer.w2_input_scale_quant
1342
- del layer.w13_blockscale_swizzled
1343
- del layer.w2_blockscale_swizzled
1340
1341
def apply(
1342
self,
@@ -1474,8 +1470,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1470
activation=activation,
1471
global_num_experts=global_num_experts,
1472
expert_map=expert_map,
1477
- w1_scale=layer.w13_blockscale_swizzled,
1478
- w2_scale=layer.w2_blockscale_swizzled,
1473
+ w1_scale=layer.w13_weight_scale,
1474
+ w2_scale=layer.w2_weight_scale,
1475
apply_router_weight_on_input=apply_router_weight_on_input,
1476
)
1477
elif (self.allow_flashinfer
@@ -1489,8 +1485,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1485
w2=layer.w2_weight,
1486
topk_weights=topk_weights,
1487
topk_ids=topk_ids,
1492
- w1_scale=layer.w13_blockscale_swizzled,
1493
- w2_scale=layer.w2_blockscale_swizzled,
1488
+ w1_scale=layer.w13_weight_scale,
1489
+ w2_scale=layer.w2_weight_scale,
1490
g1_alphas=layer.g1_alphas,
1491
g2_alphas=layer.g2_alphas,
1492
a1_gscale=layer.w13_input_scale_quant,
@@ -1510,8 +1506,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1506
a=x,
1507
w1_fp4=layer.w13_weight,
1508
w2_fp4=layer.w2_weight,
1513
- w1_blockscale=layer.w13_blockscale_swizzled,
1514
- w2_blockscale=layer.w2_blockscale_swizzled,
1509
+ w1_blockscale=layer.w13_weight_scale,
1510
+ w2_blockscale=layer.w2_weight_scale,
1511
g1_alphas=layer.g1_alphas,
1512
g2_alphas=layer.g2_alphas,
1513
a1_gscale=layer.w13_input_scale_quant,