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 8, 06:15 UTC
[Bugfix] Add method to swap quant_method on FusedMoE to fix LoRA issues (#34453)
Signed-off-by: Bill Nell <bnell@redhat.com>
bnellnm committed
Feb 15, 2026 at 23:10 UTC
5bff999d12dd061c102381b0c9c5d364c5953ea2
2 files changed
+21
-12
vllm/lora/layers/fused_moe.py
+3
-2
@@ -338,8 +338,9 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
338
fused_experts.moe_sum = moe_sum_decorator(
339
self.base_layer, fused_experts.moe_sum
340
)
341
- self.base_layer.quant_method = FusedMoEModularMethod(
342
- self.base_layer.quant_method, m_fused_moe_fn
341
+ # TODO(bnell): find a less intrusive way to handle this.
342
+ self.base_layer._replace_quant_method(
343
+ FusedMoEModularMethod(self.base_layer.quant_method, m_fused_moe_fn)
344
)
345
346
def _create_lora_a_weights(
vllm/model_executor/layers/fused_moe/layer.py
+18
-10
@@ -655,6 +655,16 @@ class FusedMoE(CustomOp):
655
enable_dbo=self.vllm_config.parallel_config.enable_dbo,
656
)
657
658
+ # TODO(bnell): This method is provided as a hook so vllm/lora/layers/fused_moe.py
659
+ # can safely swap out the quant_method. We should figure out a less
660
+ # intrusive way to do this.
661
+ def _replace_quant_method(self, mk: FusedMoEMethodBase):
662
+ self.quant_method = mk
663
+ # We need to force reconstruction of runner because we're swapping out
664
+ # the quant_method with a FusedMoEModularMethod. This logic can go
665
+ # away once the FusedMoEModularMethod is eliminated.
666
+ self.runner = self._init_runner()
667
+
668
# Note: maybe_init_modular_kernel should only be called by
669
# prepare_communication_buffer_for_model.
670
# This is called after all weight loading and post-processing, so it
@@ -676,17 +686,15 @@ class FusedMoE(CustomOp):
686
logger.debug(
687
"%s for %s(%s)", prepare_finalize.__class__.__name__, self, id(self)
688
)
679
- self.quant_method = FusedMoEModularMethod.make(
680
- self,
681
- self.quant_method,
682
- prepare_finalize,
683
- self.shared_experts,
684
- inplace=not self.moe_config.disable_inplace,
689
+ self._replace_quant_method(
690
+ FusedMoEModularMethod.make(
691
+ self,
692
+ self.quant_method,
693
+ prepare_finalize,
694
+ self.shared_experts,
695
+ inplace=not self.moe_config.disable_inplace,
696
+ )
697
)
686
- # We need to force reconstruction of runner because we're swapping out
687
- # the quant_method with a FusedMoEModularMethod. This logic can go
688
- # away once the FusedMoEModularMethod is eliminated.
689
- self.runner = self._init_runner()
698
699
@property
700
def shared_experts(self) -> torch.nn.Module | None: