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 3, 17:45 UTC
[CI] Fix pre-commit issue (#42563)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
May 13, 2026 at 15:37 UTC
3f611f6106f11cde5ff48b0f8a6dd317241a16f4
1 file changed
+10
-9
vllm/model_executor/layers/quantization/quark/quark_moe.py
+10
-9
@@ -1644,7 +1644,7 @@ class QuarkNvfp4MoEMethod(QuarkMoEMethod):
1644
layer.register_parameter("w2_input_scale_2", w2_input_scale_2)
1645
set_weight_attrs(w2_input_scale_2, extra_weight_attrs)
1646
1647
- def process_weights_after_loading(self, layer: FusedMoE) -> None:
1647
+ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
1648
"""
1649
Convert NVFP4 MoE weights into kernel format and setup the kernel.
1650
"""
@@ -1698,12 +1698,11 @@ class QuarkNvfp4MoEMethod(QuarkMoEMethod):
1698
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
1699
if self.moe_quant_config:
1700
assert self.experts_cls is not None
1701
- self.moe_mk = make_nvfp4_moe_kernel(
1701
+ self.moe_kernel = make_nvfp4_moe_kernel(
1702
moe_quant_config=self.moe_quant_config,
1703
moe_config=self.moe,
1704
experts_cls=self.experts_cls,
1705
- shared_experts=layer.shared_experts,
1706
- routing_tables=layer._maybe_init_expert_routing_tables(),
1705
+ routing_tables=layer._expert_routing_tables(),
1706
)
1707
1708
def get_fused_moe_quant_config(
@@ -1721,14 +1720,15 @@ class QuarkNvfp4MoEMethod(QuarkMoEMethod):
1720
1721
def apply(
1722
self,
1724
- layer: FusedMoE,
1723
+ layer: RoutedExperts,
1724
x: torch.Tensor,
1725
topk_weights: torch.Tensor,
1726
topk_ids: torch.Tensor,
1728
- shared_experts_input: Any | None,
1729
- ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1730
- assert self.moe_mk is not None
1731
- return self.moe_mk.apply(
1727
+ shared_experts: SharedExperts | None,
1728
+ shared_experts_input: torch.Tensor | None,
1729
+ ) -> torch.Tensor:
1730
+ assert self.moe_kernel is not None
1731
+ return self.moe_kernel.apply(
1732
x,
1733
layer.w13_weight,
1734
layer.w2_weight,
@@ -1738,5 +1738,6 @@ class QuarkNvfp4MoEMethod(QuarkMoEMethod):
1738
global_num_experts=layer.global_num_experts,
1739
expert_map=layer.expert_map,
1740
apply_router_weight_on_input=layer.apply_router_weight_on_input,
1741
+ shared_experts=shared_experts,
1742
shared_experts_input=shared_experts_input,
1743
)