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] Reuse workspace for FP8+FP4 Marlin MoE (#20500)
Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Michael Goin committed
Sep 16, 2025 at 17:45 UTC
dbebb7f812123b4a0efe5b085582d3345fe7f740
3 files changed
+8
-4
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+4
-2
@@ -398,7 +398,8 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
398
quant_type_id=scalar_types.float4_e2m1f.id,
399
apply_router_weight_on_input=apply_router_weight_on_input,
400
global_num_experts=global_num_experts,
401
- expert_map=expert_map)
401
+ expert_map=expert_map,
402
+ workspace=layer.workspace)
403
404
# FlashInfer fused experts path
405
if self.fused_experts is not None:
@@ -940,7 +941,8 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
941
quant_type_id=scalar_types.float8_e4m3fn.id,
942
apply_router_weight_on_input=apply_router_weight_on_input,
943
global_num_experts=global_num_experts,
943
- expert_map=expert_map)
944
+ expert_map=expert_map,
945
+ workspace=layer.workspace)
946
947
assert self.fused_experts_func is not None
948
vllm/model_executor/layers/quantization/fp8.py
+2
-1
@@ -1103,7 +1103,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1103
quant_type_id=scalar_types.float8_e4m3fn.id,
1104
apply_router_weight_on_input=apply_router_weight_on_input,
1105
global_num_experts=global_num_experts,
1106
- expert_map=expert_map)
1106
+ expert_map=expert_map,
1107
+ workspace=layer.workspace)
1108
elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
1109
assert self.block_quant is None
1110
assert (not renormalize and custom_routing_function is not None)
vllm/model_executor/layers/quantization/modelopt.py
+2
-1
@@ -1474,7 +1474,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1474
quant_type_id=scalar_types.float4_e2m1f.id,
1475
apply_router_weight_on_input=apply_router_weight_on_input,
1476
global_num_experts=global_num_experts,
1477
- expert_map=expert_map)
1477
+ expert_map=expert_map,
1478
+ workspace=layer.workspace)
1479
1480
if self.fused_experts is not None:
1481
assert self.allow_flashinfer and \