[LoRA] MoE LoRA Refactor

[jeejeelee]

Motivation

Currently, MoE LoRA is wired in by monkey-patching methods on the modular kernel at construction time. FusedMoEWithLoRA._inject_lora_into_fused_moe wraps FusedMoEKernel.apply, TritonExperts.activation, and TritonExperts.moe_sum with fwd_decorator / act_decorator / moe_sum_decorator, and smuggles tensors between them through moe_state_dict. This has several concrete problems:

1. Hacky and hard to maintain.

The LoRA contribution is hidden inside decorators on base-MoE methods. Someone reading the MoE experts code sees self.activation(...) and self.moe_sum(...) — they have no syntactic hint that these calls may actually run LoRA shrink/expand GEMMs, or that moe_state_dict carries cross-call state between them. Debugging requires holding the patching order in your head.

2. MoE changes don't see LoRA

Because the LoRA path lives outside the expert apply() functions, any refactor of TritonExperts / UnfusedOAITritonExperts / MarlinExperts has to re-discover the LoRA contract (what activation receives, what moe_sum receives, what shape intermediate_cache* has). Every MoE-side change risks breaking LoRA silently.

3. Hard to extend

Adding support for new features — EP, additional quantized backends, new expert impls — means replicating or working around the decorator chain, and the state-dict plumbing assumes one specific control flow. There is no extension point for an expert that wants to apply LoRA at a different point .

Change

Treat LoRA as a first-class concern in the modular MoE kernel rather than an external patch.

1. MoELoRAContext: a single explicit payload

A new MoELoRAContext dataclass (vllm/lora/lora_context.py) packages all of
the LoRA state a MoE forward pass needs:

• w13_lora_a_stacked / w13_lora_b_stacked / w2_lora_a_stacked /
  w2_lora_b_stacked
• adapter_enabled, max_loras
• routing/sharding info: top_k, w13_num_slices, fully_sharded, tp_rank,
  tp_size, local_num_experts
• the active punica_wrapper
• use_tuned_config (whether VLLM_TUNED_CONFIG_FOLDER is set)

FusedMoEWithLoRA.set_mapping builds this context once and stashes it on the base layer as FusedMoE._lora_context. MoERunnerBase forwards it into FusedMoEMethodBase.apply(..., lora_context=...), and FusedMoEModularMethod.apply / FusedMoEKernel.apply / FusedMoEKernelModularImpl.apply
thread it through to FusedMoEExpertsModular.apply(..., lora_context=...).

The context is the only LoRA surface area seen by the MoE code path —
there's no more hidden state passed between method wrappers.

2. LoRA compute inlined into the expert apply()

Expert implementations that support LoRA now call it directly inside their own apply() function, at the same logical point the decorators used to target:

• TritonExperts.apply (fused_moe.py): after the w13 GEMM and before
  activation, call self.apply_w13_lora(...) to add the LoRA delta to
  intermediate_cache1. After the w2 GEMM and before moe_sum, call
  self.apply_w2_lora(...) on intermediate_cache3, reusing the
  sorted_token_ids_lora tensors from the first call.
• UnfusedOAITritonExperts.apply (gpt_oss_triton_kernels_moe.py):
  same pattern, adjusted for the gather/scatter layout that its two
  matmul_ogs calls produce.
• MarlinExperts.apply (fused_marlin_moe.py): fused_marlin_moe
  consumes activation_func and moe_sum as callables, so the LoRA path
  wraps those two callables to inject apply_w13_lora / apply_w2_lora at
  the correct buffer state.

FusedMoEExperts.supports_lora() defaults to False. Each expert impl that has a validated LoRA path overrides it to True (TritonExperts, UnfusedOAITritonExperts, MarlinExperts). FusedMoEWithLoRA.__init__ asserts that the selected expert impl reports supports_lora(), and oracle/unquantized.py::select_unquantized_moe_backend
now filters the backend auto-selection by that flag so unsupported backends (FlashInfer / AITER) are transparently skipped when LoRA is enabled instead of silently producing wrong output or crashing later.

Because the LoRA shrink/expand is now visible in the expert source, anyone modifying TritonExperts.apply can see the LoRA call site and keep it correct; tests on the MoE path automatically cover the LoRA path as well.

3. LoRA computation stays in PunicaWrapper

MoE LoRA still respects the PunicaWrapper logic , the actual shrink/expand compute is not moved. Two new methods on PunicaWrapperBase — add_lora_w13 and add_lora_w2 — encapsulate config lookup (tuned vs. heuristic), moe_lora_align_block_size, and the add_lora_fused_moe call. PunicaWrapperGPU provides the concrete implementation. FusedMoEExpertsModular has thin helpers apply_w13_lora / apply_w2_lora that just forward the context fields to these methods.

Test Plan

All the LoRA and MoE tests on CI should pass correctly

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review to trigger a review and subscribe this PR to future pushes, or @claude review once for a one-time review.

_{Tip: disable this comment in your organization's Code Review settings.}

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @jeejeelee.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[gemini-code-assist]

Code Review

This pull request refactors the LoRA implementation for Fused MoE layers by replacing the legacy decorator-based monkey-patching with an explicit MoELoRAContext propagation through the modular kernel path. It updates various quantization methods and MoE backends to handle this context. Feedback highlights potential AttributeError issues: one regarding the access of fused_experts on FusedMoEKernel, and another concerning the missing block_shape property in the FusedMoEExperts base class and its subclasses like MarlinExperts, which would cause runtime crashes when LoRA is enabled.

[gemini-code-assist]

The check isinstance(moe_kernel.fused_experts, FusedMoEExpertsModular) might fail because FusedMoEKernel typically wraps the implementation in an impl attribute. You should likely check moe_kernel.impl.fused_experts instead, or ensure that FusedMoEKernel exposes fused_experts as a property.

Suggested change

	isinstance(moe_kernel.fused_experts, FusedMoEExpertsModular)
	isinstance(moe_kernel.impl.fused_experts, FusedMoEExpertsModular)

[bnellnm]

You could use the is_monolithic property instead of isinstance.

You might also need to add a supports_lora method to FusedMoEKernel

[gemini-code-assist]

The apply_w13_lora method (and apply_w2_lora below) calls self.block_shape, but block_shape is not defined as an abstract property in the FusedMoEExperts base class. This will cause an AttributeError at runtime for any expert implementation that does not explicitly define it (e.g., MarlinExperts or UnfusedOAITritonExperts). Please add block_shape as an abstract property to FusedMoEExperts and implement it in all subclasses that support LoRA.

[gemini-code-assist]

MarlinExperts claims to support LoRA but does not implement the block_shape property required by apply_w13_lora and apply_w2_lora in the base class. This will lead to a crash when LoRA is enabled with Marlin quantization.

[bnellnm]

Do we know if this case is ever hit now? Most methods have been switched over to the new MK initialization pattern (_setup_kernel)

[bnellnm]

Does this happen at runtime or is this part of the LoRA setup?

[jeejeelee]

This is LoRA setup, not runtime. MoELoRAContext captures references to it so the experts kernel sees fresh values without rebinding

[bnellnm]

It's probably too much for this PR but we could consider having separate subclasses for experts that support LoRA (so that the LoRA code could be completely isolated) and the setup in FusedMoEWithLoRA could construct the proper LoRA MK instead of rewriting or hijacking the existing MK.

[bnellnm]

I think at least one of the gemini comments (the one about isinstance) are relevant and should be fixed. Otherwise, I think it looks pretty good.

[jeejeelee]

It's probably too much for this PR but we could consider having separate subclasses for experts that support LoRA (so that the LoRA code could be completely isolated) and the setup in FusedMoEWithLoRA could construct the proper LoRA MK instead of rewriting or hijacking the existing MK.

Good point, I'll look into it further. Thanks!

[bnellnm]

LGTM. Thanks for the good work! Btw, Do you know if the select_gemm_impl codepath ever gets triggered? It should largely be defunct now and will be removed once everything is migrated over to _setup_kernel.

[jeejeelee]

LGTM. Thanks for the good work! Btw, Do you know if the select_gemm_impl codepath ever gets triggered? It should largely be defunct now and will be removed once everything is migrated over to _setup_kernel.

Thank you for the impressive reivew firstly.

Yes, it looks like still gets triggered — specifically by the unmigrated quant methods (AWQ-Marlin, compressed_tensors_moe).

[bnellnm]

FYI, my guess is that #40794 is causing the LoRA failure. There was a similar issue when the truncate came before the reduce in a prior PR that was fixed by moving the trunc afterwards. I'm not sure what the best solution is here. Calling .contiguous() on the result of the truncation should "fix" the problem but feels like a bandaid.