iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[LoRA] MoE LoRA Refactor (#40338)

Browse Source
This commit is contained in:
Jee Jee Li
2026-04-26 01:55:19 +00:00
committed by GitHub
parent c798593f0d
commit 8cd174fa35
14 changed files with 736 additions and 328 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py
+52 -3
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -16,6 +15,7 @@ from vllm.model_executor.layers.fused_moe.config import (
FusedMoEQuantConfig,
RoutingMethodType,
)
from vllm.model_executor.layers.fused_moe.lora_experts_mixin import LoRAExpertsMixin
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
TopKWeightAndReduceNoOP,
)
@@ -654,7 +654,7 @@ class OAITritonExperts(BaseOAITritonExperts):
)
class UnfusedOAITritonExperts(BaseOAITritonExperts):
class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts):
"""
A Triton based MoE expert class that operates on expert standard
format and explicitly keeps the activation and reduction (moe_sum) steps
@@ -721,6 +721,7 @@ class UnfusedOAITritonExperts(BaseOAITritonExperts):
if quant_config is None:
quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
global_topk_ids = topk_ids
if expert_map is not None:
topk_ids = expert_map[topk_ids]
@@ -775,10 +776,40 @@ class UnfusedOAITritonExperts(BaseOAITritonExperts):
y=intermediate_cache1,
)
# w13 LoRA: gather the activation input from expert-sorted
# intermediate_cache1, then add the LoRA delta in-place on that copy
# before passing it to activation — exactly mirroring the old
# decorator approach which modified the gathered tensor in-place.
act_input = intermediate_cache1.view(-1, N)[gather_indx.dst_indx]
sorted_token_ids_lora = None
expert_ids_lora = None
num_tokens_post_padded_lora = None
token_lora_mapping = None
lora_context = self._lora_context
if lora_context is not None:
(
sorted_token_ids_lora,
expert_ids_lora,
num_tokens_post_padded_lora,
token_lora_mapping,
) = self.apply_w13_lora(
lora_context,
y=act_input,
x=hidden_states,
topk_ids=global_topk_ids,
topk_weights=topk_weights,
expert_map=expert_map,
w1=w1,
w2=w2,
num_tokens=M,
top_k_num=topk,
)
self.activation(
activation,
intermediate_cache2,
intermediate_cache1.view(-1, N)[gather_indx.dst_indx],
act_input,
)
# matmul_ogs grouped reduction fuse sum across multiple experts:
@@ -797,6 +828,24 @@ class UnfusedOAITritonExperts(BaseOAITritonExperts):
y=intermediate_cache3,
)
# w2 LoRA: after matmul_ogs with scatter_indx, intermediate_cache3 is
# in token-topk order, matching the (M, topk, K) layout add_lora_w2 expects.
if lora_context is not None:
self.apply_w2_lora(
lora_context,
y=intermediate_cache3.view(-1, topk, K),
x=intermediate_cache2,
topk_weights=topk_weights,
sorted_token_ids_lora=sorted_token_ids_lora,
expert_ids_lora=expert_ids_lora,
num_tokens_post_padded_lora=num_tokens_post_padded_lora,
token_lora_mapping=token_lora_mapping,
num_tokens=M,
w1=w1,
w2=w2,
top_k_num=topk,
)
self.moe_sum(intermediate_cache3.view(-1, topk, K), output)
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+106 -6
View File
@@ -17,6 +17,7 @@ from vllm.model_executor.layers.fused_moe.config import (
FusedMoEParallelConfig,
FusedMoEQuantConfig,
)
from vllm.model_executor.layers.fused_moe.lora_experts_mixin import LoRAExpertsMixin
from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
batched_moe_align_block_size,
moe_align_block_size,
@@ -655,7 +656,7 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
return E, M, N, K, topk
class MarlinExperts(MarlinExpertsBase):
class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
"""Marlin-based fused MoE expert implementation."""
def supports_expert_map(self) -> bool:
@@ -720,7 +721,108 @@ class MarlinExperts(MarlinExpertsBase):
):
assert self.w1_scale is not None
assert self.w2_scale is not None
fused_marlin_moe(
ctx = self._lora_context
if ctx is None:
fused_marlin_moe(
hidden_states=hidden_states,
w1=w1,
w2=w2,
bias1=self.w1_bias,
bias2=self.w2_bias,
w1_scale=self.w1_scale,
w2_scale=self.w2_scale,
topk_weights=topk_weights,
topk_ids=topk_ids,
global_scale1=self.g1_alphas,
global_scale2=self.g2_alphas,
quant_type_id=self.quant_type_id,
apply_router_weight_on_input=apply_router_weight_on_input,
global_num_experts=global_num_experts,
activation=activation,
activation_func=self.activation,
moe_sum=self.moe_sum,
expert_map=expert_map,
output=output,
# Workspaces are swapped in workspace_shapes() to account for proper
# output buffer allocation. Please refer to workspace_shapes().
intermediate_cache13=workspace2,
intermediate_cache2=workspace13,
g_idx1=self.w13_g_idx,
g_idx2=self.w2_g_idx,
sort_indices1=self.w13_g_idx_sort_indices,
sort_indices2=self.w2_g_idx_sort_indices,
is_k_full=self.is_k_full,
input_dtype=self.input_dtype,
)
return
# LoRA path: wrap activation_func and moe_sum to inject LoRA at the
# two natural injection points.
#
# Marlin uses moe_align_block_size (same as TritonExperts) so
# intermediate_cache1 is indexed by flat (token, expert) pair index,
# which is compatible with add_lora_fused_moe's scatter mechanism.
M = hidden_states.size(0)
top_k_num = topk_ids.size(1)
lora_state: dict = {}
def activation_with_lora(
act_enum: MoEActivation,
act_output: torch.Tensor,
act_input: torch.Tensor,
) -> None:
# act_input = intermediate_cache1 (M*topk, 2N for gated)
# act_output = intermediate_cache2 (M*topk, N)
(
sorted_token_ids_lora,
expert_ids_lora,
num_tokens_post_padded_lora,
token_lora_mapping,
) = self.apply_w13_lora(
ctx,
y=act_input,
x=hidden_states,
topk_ids=topk_ids,
topk_weights=topk_weights,
expert_map=expert_map,
w1=w1,
w2=w2,
num_tokens=M,
top_k_num=top_k_num,
)
lora_state.update(
{
"sorted": sorted_token_ids_lora,
"eids": expert_ids_lora,
"npad": num_tokens_post_padded_lora,
"tlm": token_lora_mapping,
}
)
self.activation(act_enum, act_output, act_input)
lora_state["cache2"] = act_output
def moe_sum_with_lora(moe_out: torch.Tensor, out: torch.Tensor) -> None:
# moe_out shape: (M, topk, K)
self.apply_w2_lora(
ctx,
y=moe_out,
x=lora_state["cache2"],
topk_weights=topk_weights,
sorted_token_ids_lora=lora_state["sorted"],
expert_ids_lora=lora_state["eids"],
num_tokens_post_padded_lora=lora_state["npad"],
token_lora_mapping=lora_state["tlm"],
num_tokens=M,
w1=w1,
w2=w2,
top_k_num=top_k_num,
)
self.moe_sum(moe_out, out)
return fused_marlin_moe(
hidden_states=hidden_states,
w1=w1,
w2=w2,
@@ -736,12 +838,10 @@ class MarlinExperts(MarlinExpertsBase):
apply_router_weight_on_input=apply_router_weight_on_input,
global_num_experts=global_num_experts,
activation=activation,
activation_func=self.activation,
moe_sum=self.moe_sum,
activation_func=activation_with_lora,
moe_sum=moe_sum_with_lora,
expert_map=expert_map,
output=output,
# Workspaces are swapped in workspace_shapes() to account for proper
# output buffer allocation. Please refer to workspace_shapes().
intermediate_cache13=workspace2,
intermediate_cache2=workspace13,
g_idx1=self.w13_g_idx,
vllm/model_executor/layers/fused_moe/fused_moe.py
+48 -1
View File
@@ -25,6 +25,7 @@ from vllm.model_executor.layers.fused_moe.config import (
FusedMoEQuantConfig,
_get_config_dtype_str,
)
from vllm.model_executor.layers.fused_moe.lora_experts_mixin import LoRAExpertsMixin
from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
moe_align_block_size,
)
@@ -1886,7 +1887,7 @@ def fused_experts_impl(
return out_hidden_states
class TritonExperts(mk.FusedMoEExpertsModular):
class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
"""Triton-based fused MoE expert implementation."""
def __init__(
@@ -2094,6 +2095,33 @@ class TritonExperts(mk.FusedMoEExpertsModular):
B_bias=self.w1_bias,
)
# LoRA w13: applied to intermediate_cache1 before activation, using
# hidden_states as the lora_a input. moe_lora_align_block_size is
# called once here and results reused for the w2 LoRA below.
sorted_token_ids_lora = None
expert_ids_lora = None
num_tokens_post_padded_lora = None
token_lora_mapping = None
lora_context = self._lora_context
if lora_context is not None:
(
sorted_token_ids_lora,
expert_ids_lora,
num_tokens_post_padded_lora,
token_lora_mapping,
) = self.apply_w13_lora(
lora_context,
y=intermediate_cache1,
x=hidden_states,
topk_ids=topk_ids,
topk_weights=topk_weights,
expert_map=expert_map,
w1=w1,
w2=w2,
num_tokens=num_tokens,
top_k_num=top_k_num,
)
self.activation(
activation, intermediate_cache2, intermediate_cache1.view(-1, N)
)
@@ -2132,6 +2160,25 @@ class TritonExperts(mk.FusedMoEExpertsModular):
B_bias=self.w2_bias,
)
# LoRA w2: applied to intermediate_cache3 before moe_sum, using the
# unquantized intermediate_cache2 as the lora_a input. Reuses the
# sorted_token_ids_lora computed above.
if lora_context is not None:
self.apply_w2_lora(
lora_context,
y=intermediate_cache3,
x=intermediate_cache2,
topk_weights=topk_weights,
sorted_token_ids_lora=sorted_token_ids_lora,
expert_ids_lora=expert_ids_lora,
num_tokens_post_padded_lora=num_tokens_post_padded_lora,
token_lora_mapping=token_lora_mapping,
num_tokens=num_tokens,
w1=w1,
w2=w2,
top_k_num=top_k_num,
)
# separate function is required for MoE + LoRA
self.moe_sum(intermediate_cache3, output)
vllm/model_executor/layers/fused_moe/lora_context.py
+44
View File
@@ -0,0 +1,44 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass
import torch
from vllm.lora.punica_wrapper.punica_base import PunicaWrapperBase
@dataclass
class MoELoRAContext:
"""
Carries all LoRA state for one MoE forward pass.
Built by FusedMoEWithLoRA.forward() and propagated explicitly through the
modular kernel path (FusedMoEKernel -> FusedMoEExpertsModular.apply) so
that TritonExperts.apply() can compute the LoRA contribution inline,
replacing the decorator-based monkey-patch approach.
"""
# LoRA weight tensors (same shapes as FusedMoEWithLoRA attributes)
w13_lora_a_stacked: tuple[torch.Tensor, ...]
w13_lora_b_stacked: tuple[torch.Tensor, ...]
w2_lora_a_stacked: tuple[torch.Tensor, ...]
w2_lora_b_stacked: tuple[torch.Tensor, ...]
# (max_loras + 1,) int32; slot 0 is the "no-adapter" sentinel
adapter_enabled: torch.Tensor
# Metadata
max_loras: int
top_k: int
w13_num_slices: int # 2 = gated (gate + up), 1 = non-gated or 3D-fused
fully_sharded: bool
tp_rank: int
tp_size: int
local_num_experts: int
punica_wrapper: PunicaWrapperBase
# Whether VLLM_TUNED_CONFIG_FOLDER is set; selects get_lora_op_configs vs
# try_get_optimal_moe_lora_config for Triton kernel tile configs.
use_tuned_config: bool
vllm/model_executor/layers/fused_moe/lora_experts_mixin.py
+111
View File
@@ -0,0 +1,111 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm.model_executor.layers.fused_moe.lora_context import MoELoRAContext
class LoRAExpertsMixin:
"""
Mixin for FusedMoEExpertsModular subclasses that natively handle
MoELoRAContext inside their apply() implementation.
Mixing this class in:
- Flips supports_lora() to True so _can_fused_experts_support lets
LoRA through the gate check.
- Stashes a MoELoRAContext on the experts instance via
set_lora_context(), which apply() consumes from self._lora_context.
- Provides apply_w13_lora / apply_w2_lora helpers that dispatch to
the PunicaWrapper kernels.
The helper methods are pure functions of their inputs; all required
state is on lora_context or passed as arguments.
"""
_lora_context: MoELoRAContext | None = None
def set_lora_context(self, ctx: MoELoRAContext) -> None:
self._lora_context = ctx
@staticmethod
def supports_lora() -> bool:
return True
def apply_w13_lora(
self,
lora_context: MoELoRAContext,
*,
y: torch.Tensor,
x: torch.Tensor,
topk_ids: torch.Tensor,
topk_weights: torch.Tensor,
expert_map: torch.Tensor | None,
w1: torch.Tensor,
w2: torch.Tensor,
num_tokens: int,
top_k_num: int,
) -> tuple[
torch.Tensor | None,
torch.Tensor | None,
torch.Tensor | None,
torch.Tensor | None,
]:
return lora_context.punica_wrapper.add_lora_w13(
y,
x,
lora_context.w13_lora_a_stacked,
lora_context.w13_lora_b_stacked,
topk_ids,
topk_weights,
expert_map,
w1,
w2,
num_tokens,
top_k_num,
lora_context.max_loras,
lora_context.adapter_enabled,
lora_context.local_num_experts,
lora_context.top_k,
lora_context.w13_num_slices,
lora_context.fully_sharded,
lora_context.use_tuned_config,
)
def apply_w2_lora(
self,
lora_context: MoELoRAContext,
*,
y: torch.Tensor,
x: torch.Tensor,
topk_weights: torch.Tensor,
sorted_token_ids_lora: torch.Tensor | None,
expert_ids_lora: torch.Tensor | None,
num_tokens_post_padded_lora: torch.Tensor | None,
token_lora_mapping: torch.Tensor | None,
num_tokens: int,
w1: torch.Tensor,
w2: torch.Tensor,
top_k_num: int,
) -> None:
lora_context.punica_wrapper.add_lora_w2(
y,
x,
lora_context.w2_lora_a_stacked,
lora_context.w2_lora_b_stacked,
topk_weights,
sorted_token_ids_lora,
expert_ids_lora,
num_tokens_post_padded_lora,
token_lora_mapping,
num_tokens,
w1,
w2,
top_k_num,
lora_context.max_loras,
lora_context.adapter_enabled,
lora_context.top_k,
lora_context.fully_sharded,
lora_context.tp_rank,
lora_context.use_tuned_config,
)
vllm/model_executor/layers/fused_moe/modular_kernel.py
+14
View File
@@ -570,6 +570,8 @@ class FusedMoEExperts(ABC):
return False, _make_reason(f"{activation_format.value} activation format")
elif envs.VLLM_BATCH_INVARIANT and not cls._supports_batch_invariance():
return False, _make_reason("batch invariance")
elif moe_config.is_lora_enabled and not cls.supports_lora():
return False, _make_reason("LoRA")
return True, None
@staticmethod
@@ -734,6 +736,15 @@ class FusedMoEExperts(ABC):
def g2_alphas(self) -> torch.Tensor | None:
return self.quant_config.g2_alphas
@staticmethod
def supports_lora() -> bool:
"""Return True if this expert impl natively handles LoRA.
LoRA-aware experts should mix in LoRAExpertsMixin, which flips this
to True and provides the per-forward LoRA state plumbing.
"""
return False
@abstractmethod
def supports_expert_map(self) -> bool:
"""
@@ -1527,6 +1538,9 @@ class FusedMoEKernel:
def fused_experts(self) -> FusedMoEExperts:
return self.impl.fused_experts
def supports_lora(self) -> bool:
return self.fused_experts.supports_lora()
def _post_init_setup(self):
"""
Resolve any leftover setup dependencies between self.prepare_finalize
vllm/model_executor/layers/fused_moe/oracle/fp8.py
-3
View File
@@ -220,9 +220,6 @@ def select_fp8_moe_backend(
Note: Shape-specific fallbacks may still occur at runtime.
"""
if config.is_lora_enabled:
return Fp8MoeBackend.TRITON, backend_to_kernel_cls(Fp8MoeBackend.TRITON)[0]
# NOTE: the kernels are selected in the following order.
AVAILABLE_BACKENDS = _get_priority_backends(config, weight_key, activation_key)
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
-13
View File
@@ -214,19 +214,6 @@ def select_unquantized_moe_backend(
return backend, k_cls
raise ValueError(_make_log_unsupported(backend, reason))
# LoRA needs Triton's unfused activation/reduction hooks. Selecting the
# backend here ensures weights stay in a LoRA-compatible layout instead of
# being permuted for a backend like FlashInfer or AITER during load.
if moe_config.is_lora_enabled:
backend = UnquantizedMoeBackend.TRITON
if activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
backend = UnquantizedMoeBackend.BATCHED_TRITON
return _return_or_raise(
backend,
moe_config,
activation_format,
)
runner_backend = moe_config.moe_backend
if runner_backend != "auto":
requested_backend = map_unquantized_backend(runner_backend)
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+5 -1
View File
@@ -297,7 +297,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
shared_experts_input: torch.Tensor | None,
) -> torch.Tensor:
return self.forward_native(
layer, x, topk_weights, topk_ids, shared_experts_input
layer,
x,
topk_weights,
topk_ids,
shared_experts_input,
)
def apply_monolithic(