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, 08:45 UTC
[TRTLLM] Remove the MoE GEMM weight name change (#30713)
Signed-off-by: Ming Yang <minos.future@gmail.com>
Ming Yang committed
Dec 16, 2025 at 08:01 UTC
ce12b407f2ecec0a72f426c55e72a8af806b3f5c
3 files changed
+16
-40
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+4
-12
@@ -469,16 +469,14 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
469
)
470
logger.debug_once("Finished shuffling weights for TRT-LLM MOE")
471
472
- layer.gemm1_weights_fp4_shuffled = Parameter(
472
+ layer.w13_weight = Parameter(
473
gemm1_weights_fp4_shuffled, requires_grad=False
474
)
475
- layer.gemm2_weights_fp4_shuffled = Parameter(
476
- gemm2_weights_fp4_shuffled, requires_grad=False
477
- )
478
- layer.gemm1_scales_fp4_shuffled = Parameter(
475
+ layer.w2_weight = Parameter(gemm2_weights_fp4_shuffled, requires_grad=False)
476
+ layer.w13_weight_scale = Parameter(
477
gemm1_scales_fp4_shuffled, requires_grad=False
478
)
481
- layer.gemm2_scales_fp4_shuffled = Parameter(
479
+ layer.w2_weight_scale = Parameter(
480
gemm2_scales_fp4_shuffled, requires_grad=False
481
)
482
@@ -487,12 +485,6 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
485
(layer.w2_input_scale_quant * layer.g1_alphas).to(torch.float32),
486
requires_grad=False,
487
)
490
-
491
- # Clean up weights that won't be used by TRT-LLM
492
- del layer.w2_weight
493
- del layer.w2_weight_scale
494
- del layer.w13_weight
495
- del layer.w13_weight_scale
488
else:
489
# swizzle weight scales
490
layer.w13_weight_scale = torch.nn.Parameter(
vllm/model_executor/layers/quantization/modelopt.py
+4
-12
@@ -1458,16 +1458,14 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1458
)
1459
logger.debug_once("Finished shuffling weights for TRT-LLM MOE")
1460
1461
- layer.gemm1_weights_fp4_shuffled = Parameter(
1461
+ layer.w13_weight = Parameter(
1462
gemm1_weights_fp4_shuffled, requires_grad=False
1463
)
1464
- layer.gemm2_weights_fp4_shuffled = Parameter(
1465
- gemm2_weights_fp4_shuffled, requires_grad=False
1466
- )
1467
- layer.gemm1_scales_fp4_shuffled = Parameter(
1464
+ layer.w2_weight = Parameter(gemm2_weights_fp4_shuffled, requires_grad=False)
1465
+ layer.w13_weight_scale = Parameter(
1466
gemm1_scales_fp4_shuffled, requires_grad=False
1467
)
1470
- layer.gemm2_scales_fp4_shuffled = Parameter(
1468
+ layer.w2_weight_scale = Parameter(
1469
gemm2_scales_fp4_shuffled, requires_grad=False
1470
)
1471
@@ -1476,12 +1474,6 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1474
(layer.w2_input_scale_quant * layer.g1_alphas).to(torch.float32),
1475
requires_grad=False,
1476
)
1479
-
1480
- # Clean up weights that won't be used by TRT-LLM
1481
- del layer.w2_weight
1482
- del layer.w2_weight_scale
1483
- del layer.w13_weight
1484
- del layer.w13_weight_scale
1477
elif self.use_marlin:
1478
# Marlin processing
1479
prepare_moe_fp4_layer_for_marlin(layer)
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+8
-16
@@ -301,18 +301,14 @@ def flashinfer_trtllm_fp4_moe(
301
hidden_states_scale=hidden_states_scale_linear_fp4.view(
302
torch.float8_e4m3fn
303
).flatten(),
304
- gemm1_weights=layer.gemm1_weights_fp4_shuffled.data,
305
- gemm1_weights_scale=layer.gemm1_scales_fp4_shuffled.data.view(
306
- torch.float8_e4m3fn
307
- ),
304
+ gemm1_weights=layer.w13_weight.data,
305
+ gemm1_weights_scale=layer.w13_weight_scale.data.view(torch.float8_e4m3fn),
306
gemm1_bias=None,
307
gemm1_alpha=None,
308
gemm1_beta=None,
309
gemm1_clamp_limit=None,
312
- gemm2_weights=layer.gemm2_weights_fp4_shuffled.data,
313
- gemm2_weights_scale=layer.gemm2_scales_fp4_shuffled.data.view(
314
- torch.float8_e4m3fn
315
- ),
310
+ gemm2_weights=layer.w2_weight.data,
311
+ gemm2_weights_scale=layer.w2_weight_scale.data.view(torch.float8_e4m3fn),
312
gemm2_bias=None,
313
output1_scale_scalar=layer.g1_scale_c.data,
314
output1_scale_gate_scalar=layer.g1_alphas.data,
@@ -380,18 +376,14 @@ def flashinfer_trtllm_fp4_routed_moe(
376
hidden_states_scale=hidden_states_scale_linear_fp4.view(
377
torch.float8_e4m3fn
378
).flatten(),
383
- gemm1_weights=layer.gemm1_weights_fp4_shuffled.data,
384
- gemm1_weights_scale=layer.gemm1_scales_fp4_shuffled.data.view(
385
- torch.float8_e4m3fn
386
- ),
379
+ gemm1_weights=layer.w13_weight.data,
380
+ gemm1_weights_scale=layer.w13_weight_scale.data.view(torch.float8_e4m3fn),
381
gemm1_bias=None,
382
gemm1_alpha=None,
383
gemm1_beta=None,
384
gemm1_clamp_limit=None,
391
- gemm2_weights=layer.gemm2_weights_fp4_shuffled.data,
392
- gemm2_weights_scale=layer.gemm2_scales_fp4_shuffled.data.view(
393
- torch.float8_e4m3fn
394
- ),
385
+ gemm2_weights=layer.w2_weight.data,
386
+ gemm2_weights_scale=layer.w2_weight_scale.data.view(torch.float8_e4m3fn),
387
gemm2_bias=None,
388
output1_scale_scalar=layer.g1_scale_c.data,
389
output1_scale_gate_scalar=layer.g1_alphas.data,