@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8cd174fa35
Fork 0

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 4, 15:00 UTC

[LoRA] MoE LoRA Refactor (#40338)

Jee Jee Li committed Apr 26, 2026 at 09:55 UTC 8cd174fa358326d5cc4195446be2ebcd65c481ce
14 files changed +736 -328
vllm/lora/layers/fused_moe.py
+36 -297
@@ -1,6 +1,5 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import functools
3
4 import torch
5 import torch.nn as nn
@@ -14,31 +13,17 @@ from vllm.distributed.parallel_state import (
13 )
14 from vllm.distributed.utils import divide
15 from vllm.lora.layers.base import BaseLayerWithLoRA
17 -from vllm.lora.ops.triton_ops.utils import get_lora_op_configs
16 from vllm.model_executor.layers.fused_moe import FusedMoE
19 -from vllm.model_executor.layers.fused_moe.config import (
20 - _get_config_dtype_str,
21 -)
22 -from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import (
23 - UnfusedOAITritonExperts,
24 -)
25 -from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
26 - MarlinExperts,
27 -)
28 -from vllm.model_executor.layers.fused_moe.fused_moe import (
29 - TritonExperts,
30 -)
17 from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import (
18 FusedMoEModularMethod,
19 )
34 -from vllm.model_executor.layers.fused_moe.modular_kernel import (
35 - FusedMoEKernel,
36 -)
20 +from vllm.model_executor.layers.fused_moe.lora_context import MoELoRAContext
21 +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
22 from vllm.model_executor.layers.fused_moe.prepare_finalize import (
23 MoEPrepareAndFinalizeNoDPEPModular,
24 )
25
41 -from .utils import _get_lora_device, try_get_optimal_moe_lora_config
26 +from .utils import _get_lora_device
27
28
29 class FusedMoEWithLoRA(BaseLayerWithLoRA):
@@ -58,299 +43,49 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
43 # For non-gated MoE (is_act_and_mul=False), only 1 slice is needed
44 # since there's only up_proj (w1), not gate_proj + up_proj (w1 + w3)
45 self._w13_slices = 2 if base_layer.moe_config.is_act_and_mul else 1
61 - self._inject_lora_into_fused_moe()
62 -
63 - def _normalize_keys(self, config: dict[str, int | None]) -> dict[str, int | None]:
64 - normalized_config = {}
65 - for key, value in config.items():
66 - if key.islower():
67 - if key.startswith("block_"):
68 - normalized_key = "BLOCK_SIZE_" + key.split("_")[-1].upper()
69 - else:
70 - normalized_key = key.upper()
71 - else:
72 - normalized_key = key
73 - normalized_config[normalized_key] = value
74 - return normalized_config
75 -
76 - def _get_lora_moe_configs(
77 - self,
78 - op_prefix: str,
79 - num_loras: int,
80 - rank: int,
81 - num_slices: int,
82 - M: int,
83 - layer: FusedMoE,
84 - top_k: int,
85 - config_dtype: str,
86 - ):
87 - if envs.VLLM_TUNED_CONFIG_FOLDER:
88 - hidden_size = layer.hidden_size
89 - intermediate_size = (
90 - self.w2_lora_a_stacked[0].shape[-1]
91 - if op_prefix == "w2"
92 - else self.w13_lora_b_stacked[0].shape[-2]
93 - )
94 - shrink_config = get_lora_op_configs(
95 - op_type=f"fused_moe_lora_{op_prefix}_shrink",
96 - max_loras=num_loras,
97 - batch=M,
98 - hidden_size=hidden_size,
99 - rank=rank,
100 - num_slices=num_slices,
101 - moe_intermediate_size=intermediate_size,
102 - )
103 - expand_config = get_lora_op_configs(
104 - op_type=f"fused_moe_lora_{op_prefix}_expand",
105 - max_loras=num_loras,
106 - batch=M,
107 - hidden_size=hidden_size, # lora_a_stacked.shape[-1],
108 - rank=rank,
109 - num_slices=num_slices,
110 - moe_intermediate_size=intermediate_size, # lora_b_stacked.shape[-2],
111 - )
112 - else: # fall back to the default config
113 - get_config_func = functools.partial(
114 - try_get_optimal_moe_lora_config,
115 - w1_shape=layer.w13_weight.shape,
116 - w2_shape=layer.w2_weight.shape,
117 - rank=rank,
118 - top_k=top_k,
119 - dtype=config_dtype,
120 - M=M,
121 - block_shape=layer.quant_method.moe_quant_config.block_shape,
122 - )
123 - shrink_config = get_config_func(
124 - op_type=f"fused_moe_lora_{op_prefix}_shrink"
125 - )
126 - expand_config = get_config_func(
127 - op_type=f"fused_moe_lora_{op_prefix}_expand"
128 - )
129 - shrink_config = self._normalize_keys(shrink_config)
130 - expand_config = self._normalize_keys(expand_config)
131 - return shrink_config, expand_config
132 -
133 - def _inject_lora_into_fused_moe(self):
134 - moe_state_dict = {}
135 - top_k = self.base_layer.top_k
46
47 self.base_layer.ensure_moe_quant_config_init()
138 - quant_config = self.base_layer.quant_method.moe_quant_config
139 -
48 if getattr(self.base_layer.quant_method, "supports_internal_mk", False):
141 - # Use the existing modular kernel from the quant method
142 - m_fused_moe_fn = self.base_layer.quant_method.moe_kernel
49 + moe_kernel = self.base_layer.quant_method.moe_kernel
50 # Don't let the kernel own shared experts so the runner can
51 # overlap them with routed experts via a separate CUDA stream.
145 - m_fused_moe_fn.shared_experts = None
52 + moe_kernel.shared_experts = None
53 else:
147 - # Create a new modular kernel via select_gemm_impl.
148 - # Don't pass shared_experts to the kernel so the runner can
149 - # overlap them with routed experts via a separate CUDA stream.
54 prepare_finalize = MoEPrepareAndFinalizeNoDPEPModular()
151 - m_fused_moe_fn = FusedMoEKernel(
55 + moe_kernel = FusedMoEKernel(
56 prepare_finalize,
57 self.base_layer.quant_method.select_gemm_impl(
58 prepare_finalize, self.base_layer
59 ),
60 )
157 -
158 - if quant_config.use_mxfp4_w4a16:
159 - assert isinstance(
160 - m_fused_moe_fn.impl.fused_experts,
161 - (MarlinExperts, UnfusedOAITritonExperts),
162 - )
163 - else:
164 - assert isinstance(m_fused_moe_fn.impl.fused_experts, TritonExperts)
165 -
166 - def fwd_decorator(layer, func):
167 - def wrapper(*args, **kwargs):
168 - moe_state_dict["hidden_states"] = kwargs["hidden_states"]
169 - moe_state_dict["topk_ids"] = kwargs["topk_ids"]
170 - moe_state_dict["topk_weights"] = kwargs["topk_weights"]
171 - moe_state_dict["expert_map"] = kwargs["expert_map"]
172 - moe_state_dict["apply_router_weight_on_input"] = kwargs[
173 - "apply_router_weight_on_input"
174 - ]
175 - result = func(*args, **kwargs)
176 - return result
177 -
178 - return wrapper
179 -
180 - def act_decorator(layer, func):
181 - def wrapper(*args, **kwargs):
182 - _, output, input = args
183 -
184 - hidden_states = moe_state_dict["hidden_states"]
185 - topk_weights = moe_state_dict["topk_weights"]
186 - curr_topk_ids = moe_state_dict["topk_ids"]
187 -
188 - expert_map = moe_state_dict["expert_map"]
189 -
190 - config_dtype = _get_config_dtype_str(
191 - dtype=hidden_states.dtype,
192 - use_fp8_w8a8=False,
193 - use_int8_w8a16=False,
194 - use_int4_w4a16=False,
195 - )
196 - num_tokens = hidden_states.size(0)
197 - M = num_tokens
198 - max_lora_rank = self.w13_lora_a_stacked[0].shape[-2]
199 - shrink_config, expand_config = self._get_lora_moe_configs(
200 - op_prefix="w13",
201 - num_loras=self.max_loras,
202 - rank=max_lora_rank,
203 - num_slices=self._w13_slices,
204 - M=M,
205 - layer=layer,
206 - top_k=top_k,
207 - config_dtype=config_dtype,
208 - )
209 -
210 - # SPARSITY_FACTOR is a heuristic margin ensuring tokens * top_k
211 - # activates only a small fraction of total experts * loras.
212 - SPARSITY_FACTOR = 8
213 - naive_block_assignment = (
214 - expert_map is None
215 - and num_tokens * top_k * SPARSITY_FACTOR
216 - <= self.base_layer.local_num_experts * self.max_loras
217 - )
218 -
219 - # get the block size of m from customized config or default config
220 - (
221 - token_lora_mapping,
222 - sorted_token_ids_lora,
223 - expert_ids_lora,
224 - num_tokens_post_padded_lora,
225 - ) = self.punica_wrapper.moe_lora_align_block_size(
226 - curr_topk_ids,
227 - num_tokens,
228 - shrink_config["BLOCK_SIZE_M"],
229 - self.base_layer.local_num_experts,
230 - self.max_loras,
231 - self.adapter_enabled,
232 - expert_map,
233 - naive_block_assignment=naive_block_assignment,
234 - )
235 -
236 - moe_state_dict["sorted_token_ids_lora"] = sorted_token_ids_lora
237 - moe_state_dict["expert_ids_lora"] = expert_ids_lora
238 - moe_state_dict["num_tokens_post_padded_lora"] = (
239 - num_tokens_post_padded_lora
240 - )
241 - moe_state_dict["token_lora_mapping"] = token_lora_mapping
242 -
243 - if sorted_token_ids_lora is not None:
244 - expert_ids_lora = expert_ids_lora.view(self.max_loras, -1)
245 - sorted_token_ids_lora = sorted_token_ids_lora.view(
246 - self.max_loras, -1
247 - )
248 - #
249 -
250 - self.punica_wrapper.add_lora_fused_moe(
251 - input.view(-1, top_k, input.shape[-1]),
252 - hidden_states,
253 - self.w13_lora_a_stacked,
254 - self.w13_lora_b_stacked,
255 - topk_weights,
256 - sorted_token_ids_lora,
257 - expert_ids_lora,
258 - num_tokens_post_padded_lora,
259 - max_lora_rank,
260 - top_k,
261 - shrink_config, ## pass the shrink config
262 - expand_config, ## pass the expand config
263 - self.adapter_enabled,
264 - fully_sharded=self.fully_sharded,
265 - token_lora_mapping=token_lora_mapping,
266 - )
267 -
268 - result = func(*args, **kwargs)
269 -
270 - moe_state_dict["intermediate_cache2"] = output
271 - return result
272 -
273 - return wrapper
274 -
275 - def moe_sum_decorator(layer, func):
276 - def wrapper(*args, **kwargs):
277 - hidden_states = moe_state_dict["hidden_states"]
278 - topk_weights = moe_state_dict["topk_weights"]
279 -
280 - config_dtype = _get_config_dtype_str(
281 - dtype=hidden_states.dtype,
282 - use_fp8_w8a8=False,
283 - use_int8_w8a16=False,
284 - use_int4_w4a16=False,
285 - )
286 - num_tokens = hidden_states.size(0)
287 - M = num_tokens
288 - max_lora_rank = self.w2_lora_a_stacked[0].shape[-2]
289 - shrink_config, expand_config = self._get_lora_moe_configs(
290 - op_prefix="w2",
291 - num_loras=self.max_loras,
292 - rank=max_lora_rank,
293 - num_slices=1,
294 - M=M,
295 - layer=layer,
296 - top_k=top_k,
297 - config_dtype=config_dtype,
298 - )
299 -
300 - sorted_token_ids_lora = moe_state_dict["sorted_token_ids_lora"]
301 - expert_ids_lora = moe_state_dict["expert_ids_lora"]
302 - num_tokens_post_padded_lora = moe_state_dict[
303 - "num_tokens_post_padded_lora"
304 - ]
305 - token_lora_mapping = moe_state_dict.get("token_lora_mapping")
306 -
307 - if sorted_token_ids_lora is not None:
308 - expert_ids_lora = expert_ids_lora.view(self.max_loras, -1)
309 - sorted_token_ids_lora = sorted_token_ids_lora.view(
310 - self.max_loras, -1
311 - )
312 - intermediate_cache2 = moe_state_dict["intermediate_cache2"]
313 - intermediate_cache3 = args[0]
314 -
315 - shard_size_w2 = divide(self.base_layer.hidden_size, self.tp_size)
316 -
317 - self.punica_wrapper.add_lora_fused_moe(
318 - intermediate_cache3,
319 - intermediate_cache2,
320 - self.w2_lora_a_stacked,
321 - self.w2_lora_b_stacked,
322 - topk_weights,
323 - sorted_token_ids_lora,
324 - expert_ids_lora,
325 - num_tokens_post_padded_lora,
326 - max_lora_rank,
327 - top_k,
328 - shrink_config, ## pass the shrink config
329 - expand_config, ## pass the expand config
330 - self.adapter_enabled,
331 - True,
332 - fully_sharded=self.fully_sharded,
333 - offset=shard_size_w2 * self.tp_rank if self.fully_sharded else 0,
334 - token_lora_mapping=token_lora_mapping,
335 - )
336 -
337 - result = func(*args, **kwargs)
338 - return result
339 -
340 - return wrapper
341 -
342 - fused_experts = m_fused_moe_fn.impl.fused_experts
343 -
344 - m_fused_moe_fn.apply = fwd_decorator(self.base_layer, m_fused_moe_fn.apply)
345 - fused_experts.activation = act_decorator(
346 - self.base_layer, fused_experts.activation
347 - )
348 - fused_experts.moe_sum = moe_sum_decorator(
349 - self.base_layer, fused_experts.moe_sum
61 + assert moe_kernel.supports_lora(), (
62 + f"{type(moe_kernel.fused_experts).__name__} does not support LoRA. "
63 + "For unquantized MoE, set moe_backend='triton' or moe_backend='auto' "
64 + "(auto selects Triton automatically when LoRA is enabled). "
65 + "For quantized MoE, mix LoRAExpertsMixin into the experts class "
66 + "and consume self._lora_context in apply()."
67 )
351 - # TODO(bnell): find a less intrusive way to handle this.
68 + self._fused_experts = moe_kernel.fused_experts
69 self.base_layer._replace_quant_method(
353 - FusedMoEModularMethod(self.base_layer.quant_method, m_fused_moe_fn)
70 + FusedMoEModularMethod(self.base_layer.quant_method, moe_kernel)
71 + )
72 +
73 + def _build_lora_context(self):
74 + return MoELoRAContext(
75 + w13_lora_a_stacked=self.w13_lora_a_stacked,
76 + w13_lora_b_stacked=self.w13_lora_b_stacked,
77 + w2_lora_a_stacked=self.w2_lora_a_stacked,
78 + w2_lora_b_stacked=self.w2_lora_b_stacked,
79 + adapter_enabled=self.adapter_enabled,
80 + max_loras=self.max_loras,
81 + top_k=self.base_layer.top_k,
82 + w13_num_slices=self._w13_slices,
83 + fully_sharded=self.fully_sharded,
84 + tp_rank=self.tp_rank,
85 + tp_size=self.tp_size,
86 + local_num_experts=self.base_layer.local_num_experts,
87 + punica_wrapper=self.punica_wrapper,
88 + use_tuned_config=bool(envs.VLLM_TUNED_CONFIG_FOLDER),
89 )
90
91 def _create_lora_a_weights(
@@ -589,6 +324,10 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
324 index, :, : sliced_w2_lora_b.shape[1], : sliced_w2_lora_b.shape[2]
325 ].copy_(sliced_w2_lora_b, non_blocking=True)
326
327 + def set_mapping(self, punica_wrapper):
328 + super().set_mapping(punica_wrapper)
329 + self._fused_experts.set_lora_context(self._build_lora_context())
330 +
331 def forward(self, *args, **kwargs):
332 return self.base_layer.forward(*args, **kwargs)
333
vllm/lora/layers/utils.py
+5 -4
@@ -90,11 +90,12 @@ def try_get_optimal_moe_lora_config(
90 top_k: int,
91 dtype: str | None,
92 M: int,
93 - block_shape: list[int] | None = None,
93 ) -> dict[str, int | None]:
95 - config = try_get_optimal_moe_config(
96 - w1_shape, w2_shape, top_k, dtype, M, block_shape
97 - ).copy()
94 + # LoRA shrink/expand operates on bf16/fp16 adapters regardless of the
95 + # base MoE weight's block-wise quantization, so block_shape is omitted
96 + # from the config lookup — the non-quantized branch in get_default_config
97 + # ignores it anyway.
98 + config = try_get_optimal_moe_config(w1_shape, w2_shape, top_k, dtype, M).copy()
99 if op_type in [
100 "fused_moe_lora_w13_shrink",
101 "fused_moe_lora_w2_shrink",
vllm/lora/ops/triton_ops/utils.py
+17
@@ -321,3 +321,20 @@ def supports_pdl(device: torch.device | None = None) -> bool:
321 def supports_tma(device: torch.device | None = None) -> bool:
322 # TMA requires compute capability SM90 or above
323 return current_platform.is_cuda() and current_platform.has_device_capability(90)
324 +
325 +
326 +def _normalize_lora_config_keys(
327 + config: dict[str, int | None],
328 +) -> dict[str, int | None]:
329 + """Normalize Triton config dict keys to uppercase BLOCK_SIZE_* format."""
330 + out: dict[str, int | None] = {}
331 + for key, val in config.items():
332 + if key.islower():
333 + if key.startswith("block_"):
334 + nk = "BLOCK_SIZE_" + key.split("_")[-1].upper()
335 + else:
336 + nk = key.upper()
337 + else:
338 + nk = key
339 + out[nk] = val
340 + return out
vllm/lora/punica_wrapper/punica_base.py
+62
@@ -493,3 +493,65 @@ class PunicaWrapperBase(PunicaWrapperABC):
493 """
494 # TODO: implement it based on torch ops
495 raise NotImplementedError
496 +
497 + def add_lora_w13(
498 + self,
499 + y: torch.Tensor,
500 + x: torch.Tensor,
501 + lora_a_stacked: tuple[torch.Tensor, ...],
502 + lora_b_stacked: tuple[torch.Tensor, ...],
503 + topk_ids: torch.Tensor,
504 + topk_weights: torch.Tensor,
505 + expert_map: torch.Tensor | None,
506 + w1: torch.Tensor,
507 + w2: torch.Tensor,
508 + num_tokens: int,
509 + top_k_num: int,
510 + max_loras: int,
511 + adapter_enabled: torch.Tensor,
512 + local_num_experts: int,
513 + top_k: int,
514 + num_slices: int,
515 + fully_sharded: bool,
516 + use_tuned_config: bool,
517 + ) -> tuple[
518 + torch.Tensor | None,
519 + torch.Tensor | None,
520 + torch.Tensor | None,
521 + torch.Tensor | None,
522 + ]:
523 + """Apply w13 LoRA to y (intermediate_cache1) in-place before activation.
524 +
525 + Returns (sorted_token_ids_lora, expert_ids_lora,
526 + num_tokens_post_padded_lora, token_lora_mapping)
527 + for reuse by add_lora_w2.
528 + """
529 + raise NotImplementedError
530 +
531 + def add_lora_w2(
532 + self,
533 + y: torch.Tensor,
534 + x: torch.Tensor,
535 + lora_a_stacked: tuple[torch.Tensor, ...],
536 + lora_b_stacked: tuple[torch.Tensor, ...],
537 + topk_weights: torch.Tensor,
538 + sorted_token_ids_lora: torch.Tensor | None,
539 + expert_ids_lora: torch.Tensor | None,
540 + num_tokens_post_padded_lora: torch.Tensor | None,
541 + token_lora_mapping: torch.Tensor | None,
542 + num_tokens: int,
543 + w1: torch.Tensor,
544 + w2: torch.Tensor,
545 + top_k_num: int,
546 + max_loras: int,
547 + adapter_enabled: torch.Tensor,
548 + top_k: int,
549 + fully_sharded: bool,
550 + tp_rank: int,
551 + use_tuned_config: bool,
552 + ) -> None:
553 + """Apply w2 LoRA to y (intermediate_cache3) in-place before moe_sum.
554 +
555 + Reuses routing tensors returned by add_lora_w13.
556 + """
557 + raise NotImplementedError
vllm/lora/punica_wrapper/punica_gpu.py
+236
@@ -459,3 +459,239 @@ class PunicaWrapperGPU(PunicaWrapperBase):
459 fully_sharded,
460 offset,
461 )
462 +
463 + def add_lora_w13(
464 + self,
465 + y: torch.Tensor,
466 + x: torch.Tensor,
467 + lora_a_stacked: tuple[torch.Tensor, ...],
468 + lora_b_stacked: tuple[torch.Tensor, ...],
469 + topk_ids: torch.Tensor,
470 + topk_weights: torch.Tensor,
471 + expert_map: torch.Tensor | None,
472 + w1: torch.Tensor,
473 + w2: torch.Tensor,
474 + num_tokens: int,
475 + top_k_num: int,
476 + max_loras: int,
477 + adapter_enabled: torch.Tensor,
478 + local_num_experts: int,
479 + top_k: int,
480 + num_slices: int,
481 + fully_sharded: bool,
482 + use_tuned_config: bool,
483 + ) -> tuple[
484 + torch.Tensor | None,
485 + torch.Tensor | None,
486 + torch.Tensor | None,
487 + torch.Tensor | None,
488 + ]:
489 + import functools
490 +
491 + from vllm.lora.layers.utils import try_get_optimal_moe_lora_config
492 + from vllm.lora.ops.triton_ops.utils import (
493 + _normalize_lora_config_keys,
494 + get_lora_op_configs,
495 + )
496 + from vllm.model_executor.layers.fused_moe.config import _get_config_dtype_str
497 +
498 + config_dtype = _get_config_dtype_str(
499 + dtype=x.dtype,
500 + use_fp8_w8a8=False,
501 + use_int8_w8a16=False,
502 + use_int4_w4a16=False,
503 + )
504 + max_lora_rank = lora_a_stacked[0].shape[-2]
505 +
506 + if use_tuned_config:
507 + shrink_config = get_lora_op_configs(
508 + op_type="fused_moe_lora_w13_shrink",
509 + max_loras=max_loras,
510 + batch=num_tokens,
511 + hidden_size=x.shape[-1],
512 + rank=max_lora_rank,
513 + num_slices=num_slices,
514 + moe_intermediate_size=lora_b_stacked[0].shape[-2],
515 + )
516 + expand_config = get_lora_op_configs(
517 + op_type="fused_moe_lora_w13_expand",
518 + max_loras=max_loras,
519 + batch=num_tokens,
520 + hidden_size=x.shape[-1],
521 + rank=max_lora_rank,
522 + num_slices=num_slices,
523 + moe_intermediate_size=lora_b_stacked[0].shape[-2],
524 + )
525 + else:
526 + get_config = functools.partial(
527 + try_get_optimal_moe_lora_config,
528 + w1_shape=w1.shape,
529 + w2_shape=w2.shape,
530 + rank=max_lora_rank,
531 + top_k=top_k,
532 + dtype=config_dtype,
533 + M=num_tokens,
534 + )
535 + shrink_config = get_config(op_type="fused_moe_lora_w13_shrink")
536 + expand_config = get_config(op_type="fused_moe_lora_w13_expand")
537 +
538 + shrink_config = _normalize_lora_config_keys(shrink_config)
539 + expand_config = _normalize_lora_config_keys(expand_config)
540 +
541 + SPARSITY_FACTOR = 8
542 + naive_block_assignment = (
543 + expert_map is None
544 + and num_tokens * top_k * SPARSITY_FACTOR <= local_num_experts * max_loras
545 + )
546 +
547 + (
548 + token_lora_mapping,
549 + sorted_token_ids_lora,
550 + expert_ids_lora,
551 + num_tokens_post_padded_lora,
552 + ) = self.moe_lora_align_block_size(
553 + topk_ids,
554 + num_tokens,
555 + int(shrink_config.get("BLOCK_SIZE_M") or 64),
556 + local_num_experts,
557 + max_loras,
558 + adapter_enabled,
559 + expert_map,
560 + naive_block_assignment=naive_block_assignment,
561 + )
562 +
563 + _sorted = sorted_token_ids_lora
564 + _eids = expert_ids_lora
565 + if _sorted is not None:
566 + _eids = _eids.view(max_loras, -1)
567 + _sorted = _sorted.view(max_loras, -1)
568 +
569 + self.add_lora_fused_moe(
570 + y.view(-1, top_k_num, y.shape[-1]),
571 + x,
572 + lora_a_stacked,
573 + lora_b_stacked,
574 + topk_weights,
575 + _sorted,
576 + _eids,
577 + num_tokens_post_padded_lora,
578 + max_lora_rank,
579 + top_k,
580 + shrink_config,
581 + expand_config,
582 + adapter_enabled,
583 + fully_sharded=fully_sharded,
584 + token_lora_mapping=token_lora_mapping,
585 + )
586 +
587 + return (
588 + sorted_token_ids_lora,
589 + expert_ids_lora,
590 + num_tokens_post_padded_lora,
591 + token_lora_mapping,
592 + )
593 +
594 + def add_lora_w2(
595 + self,
596 + y: torch.Tensor,
597 + x: torch.Tensor,
598 + lora_a_stacked: tuple[torch.Tensor, ...],
599 + lora_b_stacked: tuple[torch.Tensor, ...],
600 + topk_weights: torch.Tensor,
601 + sorted_token_ids_lora: torch.Tensor | None,
602 + expert_ids_lora: torch.Tensor | None,
603 + num_tokens_post_padded_lora: torch.Tensor | None,
604 + token_lora_mapping: torch.Tensor | None,
605 + num_tokens: int,
606 + w1: torch.Tensor,
607 + w2: torch.Tensor,
608 + top_k_num: int,
609 + max_loras: int,
610 + adapter_enabled: torch.Tensor,
611 + top_k: int,
612 + fully_sharded: bool,
613 + tp_rank: int,
614 + use_tuned_config: bool,
615 + ) -> None:
616 + import functools
617 +
618 + from vllm.lora.layers.utils import try_get_optimal_moe_lora_config
619 + from vllm.lora.ops.triton_ops.utils import (
620 + _normalize_lora_config_keys,
621 + get_lora_op_configs,
622 + )
623 + from vllm.model_executor.layers.fused_moe.config import _get_config_dtype_str
624 +
625 + config_dtype = _get_config_dtype_str(
626 + dtype=x.dtype,
627 + use_fp8_w8a8=False,
628 + use_int8_w8a16=False,
629 + use_int4_w4a16=False,
630 + )
631 + max_lora_rank = lora_a_stacked[0].shape[-2]
632 +
633 + if use_tuned_config:
634 + shrink_config = get_lora_op_configs(
635 + op_type="fused_moe_lora_w2_shrink",
636 + max_loras=max_loras,
637 + batch=num_tokens,
638 + hidden_size=y.shape[-1],
639 + rank=max_lora_rank,
640 + num_slices=1,
641 + moe_intermediate_size=lora_a_stacked[0].shape[-1],
642 + )
643 + expand_config = get_lora_op_configs(
644 + op_type="fused_moe_lora_w2_expand",
645 + max_loras=max_loras,
646 + batch=num_tokens,
647 + hidden_size=y.shape[-1],
648 + rank=max_lora_rank,
649 + num_slices=1,
650 + moe_intermediate_size=lora_a_stacked[0].shape[-1],
651 + )
652 + else:
653 + get_config = functools.partial(
654 + try_get_optimal_moe_lora_config,
655 + w1_shape=w1.shape,
656 + w2_shape=w2.shape,
657 + rank=max_lora_rank,
658 + top_k=top_k,
659 + dtype=config_dtype,
660 + M=num_tokens,
661 + )
662 + shrink_config = get_config(op_type="fused_moe_lora_w2_shrink")
663 + expand_config = get_config(op_type="fused_moe_lora_w2_expand")
664 +
665 + shrink_config = _normalize_lora_config_keys(shrink_config)
666 + expand_config = _normalize_lora_config_keys(expand_config)
667 +
668 + _sorted = sorted_token_ids_lora
669 + _eids = expert_ids_lora
670 + if _sorted is not None:
671 + assert _eids is not None
672 + _eids = _eids.view(max_loras, -1)
673 + _sorted = _sorted.view(max_loras, -1)
674 +
675 + # w2_lora_b shape[-2] is hidden_size // tp_size when fully_sharded
676 + shard_size = lora_b_stacked[0].shape[-2]
677 + offset = shard_size * tp_rank if fully_sharded else 0
678 +
679 + self.add_lora_fused_moe(
680 + y,
681 + x,
682 + lora_a_stacked,
683 + lora_b_stacked,
684 + topk_weights,
685 + _sorted,
686 + _eids,
687 + num_tokens_post_padded_lora,
688 + max_lora_rank,
689 + top_k,
690 + shrink_config,
691 + expand_config,
692 + adapter_enabled,
693 + True, # mul_routed_weight
694 + fully_sharded=fully_sharded,
695 + offset=offset,
696 + token_lora_mapping=token_lora_mapping,
697 + )
vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py
+52 -3
@@ -1,7 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -
4 import torch
5
6 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -16,6 +15,7 @@ from vllm.model_executor.layers.fused_moe.config import (
15 FusedMoEQuantConfig,
16 RoutingMethodType,
17 )
18 +from vllm.model_executor.layers.fused_moe.lora_experts_mixin import LoRAExpertsMixin
19 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
20 TopKWeightAndReduceNoOP,
21 )
@@ -654,7 +654,7 @@ class OAITritonExperts(BaseOAITritonExperts):
654 )
655
656
657 -class UnfusedOAITritonExperts(BaseOAITritonExperts):
657 +class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts):
658 """
659 A Triton based MoE expert class that operates on expert standard
660 format and explicitly keeps the activation and reduction (moe_sum) steps
@@ -721,6 +721,7 @@ class UnfusedOAITritonExperts(BaseOAITritonExperts):
721 if quant_config is None:
722 quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
723
724 + global_topk_ids = topk_ids
725 if expert_map is not None:
726 topk_ids = expert_map[topk_ids]
727
@@ -775,10 +776,40 @@ class UnfusedOAITritonExperts(BaseOAITritonExperts):
776 y=intermediate_cache1,
777 )
778
779 + # w13 LoRA: gather the activation input from expert-sorted
780 + # intermediate_cache1, then add the LoRA delta in-place on that copy
781 + # before passing it to activation — exactly mirroring the old
782 + # decorator approach which modified the gathered tensor in-place.
783 + act_input = intermediate_cache1.view(-1, N)[gather_indx.dst_indx]
784 +
785 + sorted_token_ids_lora = None
786 + expert_ids_lora = None
787 + num_tokens_post_padded_lora = None
788 + token_lora_mapping = None
789 + lora_context = self._lora_context
790 + if lora_context is not None:
791 + (
792 + sorted_token_ids_lora,
793 + expert_ids_lora,
794 + num_tokens_post_padded_lora,
795 + token_lora_mapping,
796 + ) = self.apply_w13_lora(
797 + lora_context,
798 + y=act_input,
799 + x=hidden_states,
800 + topk_ids=global_topk_ids,
801 + topk_weights=topk_weights,
802 + expert_map=expert_map,
803 + w1=w1,
804 + w2=w2,
805 + num_tokens=M,
806 + top_k_num=topk,
807 + )
808 +
809 self.activation(
810 activation,
811 intermediate_cache2,
781 - intermediate_cache1.view(-1, N)[gather_indx.dst_indx],
812 + act_input,
813 )
814
815 # matmul_ogs grouped reduction fuse sum across multiple experts:
@@ -797,6 +828,24 @@ class UnfusedOAITritonExperts(BaseOAITritonExperts):
828 y=intermediate_cache3,
829 )
830
831 + # w2 LoRA: after matmul_ogs with scatter_indx, intermediate_cache3 is
832 + # in token-topk order, matching the (M, topk, K) layout add_lora_w2 expects.
833 + if lora_context is not None:
834 + self.apply_w2_lora(
835 + lora_context,
836 + y=intermediate_cache3.view(-1, topk, K),
837 + x=intermediate_cache2,
838 + topk_weights=topk_weights,
839 + sorted_token_ids_lora=sorted_token_ids_lora,
840 + expert_ids_lora=expert_ids_lora,
841 + num_tokens_post_padded_lora=num_tokens_post_padded_lora,
842 + token_lora_mapping=token_lora_mapping,
843 + num_tokens=M,
844 + w1=w1,
845 + w2=w2,
846 + top_k_num=topk,
847 + )
848 +
849 self.moe_sum(intermediate_cache3.view(-1, topk, K), output)
850
851
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+106 -6
@@ -17,6 +17,7 @@ from vllm.model_executor.layers.fused_moe.config import (
17 FusedMoEParallelConfig,
18 FusedMoEQuantConfig,
19 )
20 +from vllm.model_executor.layers.fused_moe.lora_experts_mixin import LoRAExpertsMixin
21 from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
22 batched_moe_align_block_size,
23 moe_align_block_size,
@@ -655,7 +656,7 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
656 return E, M, N, K, topk
657
658
658 -class MarlinExperts(MarlinExpertsBase):
659 +class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
660 """Marlin-based fused MoE expert implementation."""
661
662 def supports_expert_map(self) -> bool:
@@ -720,7 +721,108 @@ class MarlinExperts(MarlinExpertsBase):
721 ):
722 assert self.w1_scale is not None
723 assert self.w2_scale is not None
723 - fused_marlin_moe(
724 +
725 + ctx = self._lora_context
726 + if ctx is None:
727 + fused_marlin_moe(
728 + hidden_states=hidden_states,
729 + w1=w1,
730 + w2=w2,
731 + bias1=self.w1_bias,
732 + bias2=self.w2_bias,
733 + w1_scale=self.w1_scale,
734 + w2_scale=self.w2_scale,
735 + topk_weights=topk_weights,
736 + topk_ids=topk_ids,
737 + global_scale1=self.g1_alphas,
738 + global_scale2=self.g2_alphas,
739 + quant_type_id=self.quant_type_id,
740 + apply_router_weight_on_input=apply_router_weight_on_input,
741 + global_num_experts=global_num_experts,
742 + activation=activation,
743 + activation_func=self.activation,
744 + moe_sum=self.moe_sum,
745 + expert_map=expert_map,
746 + output=output,
747 + # Workspaces are swapped in workspace_shapes() to account for proper
748 + # output buffer allocation. Please refer to workspace_shapes().
749 + intermediate_cache13=workspace2,
750 + intermediate_cache2=workspace13,
751 + g_idx1=self.w13_g_idx,
752 + g_idx2=self.w2_g_idx,
753 + sort_indices1=self.w13_g_idx_sort_indices,
754 + sort_indices2=self.w2_g_idx_sort_indices,
755 + is_k_full=self.is_k_full,
756 + input_dtype=self.input_dtype,
757 + )
758 + return
759 +
760 + # LoRA path: wrap activation_func and moe_sum to inject LoRA at the
761 + # two natural injection points.
762 + #
763 + # Marlin uses moe_align_block_size (same as TritonExperts) so
764 + # intermediate_cache1 is indexed by flat (token, expert) pair index,
765 + # which is compatible with add_lora_fused_moe's scatter mechanism.
766 +
767 + M = hidden_states.size(0)
768 + top_k_num = topk_ids.size(1)
769 + lora_state: dict = {}
770 +
771 + def activation_with_lora(
772 + act_enum: MoEActivation,
773 + act_output: torch.Tensor,
774 + act_input: torch.Tensor,
775 + ) -> None:
776 + # act_input = intermediate_cache1 (M*topk, 2N for gated)
777 + # act_output = intermediate_cache2 (M*topk, N)
778 +
779 + (
780 + sorted_token_ids_lora,
781 + expert_ids_lora,
782 + num_tokens_post_padded_lora,
783 + token_lora_mapping,
784 + ) = self.apply_w13_lora(
785 + ctx,
786 + y=act_input,
787 + x=hidden_states,
788 + topk_ids=topk_ids,
789 + topk_weights=topk_weights,
790 + expert_map=expert_map,
791 + w1=w1,
792 + w2=w2,
793 + num_tokens=M,
794 + top_k_num=top_k_num,
795 + )
796 + lora_state.update(
797 + {
798 + "sorted": sorted_token_ids_lora,
799 + "eids": expert_ids_lora,
800 + "npad": num_tokens_post_padded_lora,
801 + "tlm": token_lora_mapping,
802 + }
803 + )
804 + self.activation(act_enum, act_output, act_input)
805 + lora_state["cache2"] = act_output
806 +
807 + def moe_sum_with_lora(moe_out: torch.Tensor, out: torch.Tensor) -> None:
808 + # moe_out shape: (M, topk, K)
809 + self.apply_w2_lora(
810 + ctx,
811 + y=moe_out,
812 + x=lora_state["cache2"],
813 + topk_weights=topk_weights,
814 + sorted_token_ids_lora=lora_state["sorted"],
815 + expert_ids_lora=lora_state["eids"],
816 + num_tokens_post_padded_lora=lora_state["npad"],
817 + token_lora_mapping=lora_state["tlm"],
818 + num_tokens=M,
819 + w1=w1,
820 + w2=w2,
821 + top_k_num=top_k_num,
822 + )
823 + self.moe_sum(moe_out, out)
824 +
825 + return fused_marlin_moe(
826 hidden_states=hidden_states,
827 w1=w1,
828 w2=w2,
@@ -736,12 +838,10 @@ class MarlinExperts(MarlinExpertsBase):
838 apply_router_weight_on_input=apply_router_weight_on_input,
839 global_num_experts=global_num_experts,
840 activation=activation,
739 - activation_func=self.activation,
740 - moe_sum=self.moe_sum,
841 + activation_func=activation_with_lora,
842 + moe_sum=moe_sum_with_lora,
843 expert_map=expert_map,
844 output=output,
743 - # Workspaces are swapped in workspace_shapes() to account for proper
744 - # output buffer allocation. Please refer to workspace_shapes().
845 intermediate_cache13=workspace2,
846 intermediate_cache2=workspace13,
847 g_idx1=self.w13_g_idx,
vllm/model_executor/layers/fused_moe/fused_moe.py
+48 -1
@@ -25,6 +25,7 @@ from vllm.model_executor.layers.fused_moe.config import (
25 FusedMoEQuantConfig,
26 _get_config_dtype_str,
27 )
28 +from vllm.model_executor.layers.fused_moe.lora_experts_mixin import LoRAExpertsMixin
29 from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
30 moe_align_block_size,
31 )
@@ -1886,7 +1887,7 @@ def fused_experts_impl(
1887 return out_hidden_states
1888
1889
1889 -class TritonExperts(mk.FusedMoEExpertsModular):
1890 +class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
1891 """Triton-based fused MoE expert implementation."""
1892
1893 def __init__(
@@ -2094,6 +2095,33 @@ class TritonExperts(mk.FusedMoEExpertsModular):
2095 B_bias=self.w1_bias,
2096 )
2097
2098 + # LoRA w13: applied to intermediate_cache1 before activation, using
2099 + # hidden_states as the lora_a input. moe_lora_align_block_size is
2100 + # called once here and results reused for the w2 LoRA below.
2101 + sorted_token_ids_lora = None
2102 + expert_ids_lora = None
2103 + num_tokens_post_padded_lora = None
2104 + token_lora_mapping = None
2105 + lora_context = self._lora_context
2106 + if lora_context is not None:
2107 + (
2108 + sorted_token_ids_lora,
2109 + expert_ids_lora,
2110 + num_tokens_post_padded_lora,
2111 + token_lora_mapping,
2112 + ) = self.apply_w13_lora(
2113 + lora_context,
2114 + y=intermediate_cache1,
2115 + x=hidden_states,
2116 + topk_ids=topk_ids,
2117 + topk_weights=topk_weights,
2118 + expert_map=expert_map,
2119 + w1=w1,
2120 + w2=w2,
2121 + num_tokens=num_tokens,
2122 + top_k_num=top_k_num,
2123 + )
2124 +
2125 self.activation(
2126 activation, intermediate_cache2, intermediate_cache1.view(-1, N)
2127 )
@@ -2132,6 +2160,25 @@ class TritonExperts(mk.FusedMoEExpertsModular):
2160 B_bias=self.w2_bias,
2161 )
2162
2163 + # LoRA w2: applied to intermediate_cache3 before moe_sum, using the
2164 + # unquantized intermediate_cache2 as the lora_a input. Reuses the
2165 + # sorted_token_ids_lora computed above.
2166 + if lora_context is not None:
2167 + self.apply_w2_lora(
2168 + lora_context,
2169 + y=intermediate_cache3,
2170 + x=intermediate_cache2,
2171 + topk_weights=topk_weights,
2172 + sorted_token_ids_lora=sorted_token_ids_lora,
2173 + expert_ids_lora=expert_ids_lora,
2174 + num_tokens_post_padded_lora=num_tokens_post_padded_lora,
2175 + token_lora_mapping=token_lora_mapping,
2176 + num_tokens=num_tokens,
2177 + w1=w1,
2178 + w2=w2,
2179 + top_k_num=top_k_num,
2180 + )
2181 +
2182 # separate function is required for MoE + LoRA
2183 self.moe_sum(intermediate_cache3, output)
2184
vllm/model_executor/layers/fused_moe/lora_context.py new
+44
@@ -0,0 +1,44 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from dataclasses import dataclass
5 +
6 +import torch
7 +
8 +from vllm.lora.punica_wrapper.punica_base import PunicaWrapperBase
9 +
10 +
11 +@dataclass
12 +class MoELoRAContext:
13 + """
14 + Carries all LoRA state for one MoE forward pass.
15 +
16 + Built by FusedMoEWithLoRA.forward() and propagated explicitly through the
17 + modular kernel path (FusedMoEKernel -> FusedMoEExpertsModular.apply) so
18 + that TritonExperts.apply() can compute the LoRA contribution inline,
19 + replacing the decorator-based monkey-patch approach.
20 + """
21 +
22 + # LoRA weight tensors (same shapes as FusedMoEWithLoRA attributes)
23 + w13_lora_a_stacked: tuple[torch.Tensor, ...]
24 + w13_lora_b_stacked: tuple[torch.Tensor, ...]
25 + w2_lora_a_stacked: tuple[torch.Tensor, ...]
26 + w2_lora_b_stacked: tuple[torch.Tensor, ...]
27 +
28 + # (max_loras + 1,) int32; slot 0 is the "no-adapter" sentinel
29 + adapter_enabled: torch.Tensor
30 +
31 + # Metadata
32 + max_loras: int
33 + top_k: int
34 + w13_num_slices: int # 2 = gated (gate + up), 1 = non-gated or 3D-fused
35 + fully_sharded: bool
36 + tp_rank: int
37 + tp_size: int
38 + local_num_experts: int
39 +
40 + punica_wrapper: PunicaWrapperBase
41 +
42 + # Whether VLLM_TUNED_CONFIG_FOLDER is set; selects get_lora_op_configs vs
43 + # try_get_optimal_moe_lora_config for Triton kernel tile configs.
44 + use_tuned_config: bool
vllm/model_executor/layers/fused_moe/lora_experts_mixin.py new
+111
@@ -0,0 +1,111 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from vllm.model_executor.layers.fused_moe.lora_context import MoELoRAContext
7 +
8 +
9 +class LoRAExpertsMixin:
10 + """
11 + Mixin for FusedMoEExpertsModular subclasses that natively handle
12 + MoELoRAContext inside their apply() implementation.
13 +
14 + Mixing this class in:
15 + - Flips supports_lora() to True so _can_fused_experts_support lets
16 + LoRA through the gate check.
17 + - Stashes a MoELoRAContext on the experts instance via
18 + set_lora_context(), which apply() consumes from self._lora_context.
19 + - Provides apply_w13_lora / apply_w2_lora helpers that dispatch to
20 + the PunicaWrapper kernels.
21 +
22 + The helper methods are pure functions of their inputs; all required
23 + state is on lora_context or passed as arguments.
24 + """
25 +
26 + _lora_context: MoELoRAContext | None = None
27 +
28 + def set_lora_context(self, ctx: MoELoRAContext) -> None:
29 + self._lora_context = ctx
30 +
31 + @staticmethod
32 + def supports_lora() -> bool:
33 + return True
34 +
35 + def apply_w13_lora(
36 + self,
37 + lora_context: MoELoRAContext,
38 + *,
39 + y: torch.Tensor,
40 + x: torch.Tensor,
41 + topk_ids: torch.Tensor,
42 + topk_weights: torch.Tensor,
43 + expert_map: torch.Tensor | None,
44 + w1: torch.Tensor,
45 + w2: torch.Tensor,
46 + num_tokens: int,
47 + top_k_num: int,
48 + ) -> tuple[
49 + torch.Tensor | None,
50 + torch.Tensor | None,
51 + torch.Tensor | None,
52 + torch.Tensor | None,
53 + ]:
54 + return lora_context.punica_wrapper.add_lora_w13(
55 + y,
56 + x,
57 + lora_context.w13_lora_a_stacked,
58 + lora_context.w13_lora_b_stacked,
59 + topk_ids,
60 + topk_weights,
61 + expert_map,
62 + w1,
63 + w2,
64 + num_tokens,
65 + top_k_num,
66 + lora_context.max_loras,
67 + lora_context.adapter_enabled,
68 + lora_context.local_num_experts,
69 + lora_context.top_k,
70 + lora_context.w13_num_slices,
71 + lora_context.fully_sharded,
72 + lora_context.use_tuned_config,
73 + )
74 +
75 + def apply_w2_lora(
76 + self,
77 + lora_context: MoELoRAContext,
78 + *,
79 + y: torch.Tensor,
80 + x: torch.Tensor,
81 + topk_weights: torch.Tensor,
82 + sorted_token_ids_lora: torch.Tensor | None,
83 + expert_ids_lora: torch.Tensor | None,
84 + num_tokens_post_padded_lora: torch.Tensor | None,
85 + token_lora_mapping: torch.Tensor | None,
86 + num_tokens: int,
87 + w1: torch.Tensor,
88 + w2: torch.Tensor,
89 + top_k_num: int,
90 + ) -> None:
91 + lora_context.punica_wrapper.add_lora_w2(
92 + y,
93 + x,
94 + lora_context.w2_lora_a_stacked,
95 + lora_context.w2_lora_b_stacked,
96 + topk_weights,
97 + sorted_token_ids_lora,
98 + expert_ids_lora,
99 + num_tokens_post_padded_lora,
100 + token_lora_mapping,
101 + num_tokens,
102 + w1,
103 + w2,
104 + top_k_num,
105 + lora_context.max_loras,
106 + lora_context.adapter_enabled,
107 + lora_context.top_k,
108 + lora_context.fully_sharded,
109 + lora_context.tp_rank,
110 + lora_context.use_tuned_config,
111 + )
vllm/model_executor/layers/fused_moe/modular_kernel.py
+14
@@ -570,6 +570,8 @@ class FusedMoEExperts(ABC):
570 return False, _make_reason(f"{activation_format.value} activation format")
571 elif envs.VLLM_BATCH_INVARIANT and not cls._supports_batch_invariance():
572 return False, _make_reason("batch invariance")
573 + elif moe_config.is_lora_enabled and not cls.supports_lora():
574 + return False, _make_reason("LoRA")
575 return True, None
576
577 @staticmethod
@@ -734,6 +736,15 @@ class FusedMoEExperts(ABC):
736 def g2_alphas(self) -> torch.Tensor | None:
737 return self.quant_config.g2_alphas
738
739 + @staticmethod
740 + def supports_lora() -> bool:
741 + """Return True if this expert impl natively handles LoRA.
742 +
743 + LoRA-aware experts should mix in LoRAExpertsMixin, which flips this
744 + to True and provides the per-forward LoRA state plumbing.
745 + """
746 + return False
747 +
748 @abstractmethod
749 def supports_expert_map(self) -> bool:
750 """
@@ -1527,6 +1538,9 @@ class FusedMoEKernel:
1538 def fused_experts(self) -> FusedMoEExperts:
1539 return self.impl.fused_experts
1540
1541 + def supports_lora(self) -> bool:
1542 + return self.fused_experts.supports_lora()
1543 +
1544 def _post_init_setup(self):
1545 """
1546 Resolve any leftover setup dependencies between self.prepare_finalize
vllm/model_executor/layers/fused_moe/oracle/fp8.py
-3
@@ -220,9 +220,6 @@ def select_fp8_moe_backend(
220 Note: Shape-specific fallbacks may still occur at runtime.
221 """
222
223 - if config.is_lora_enabled:
224 - return Fp8MoeBackend.TRITON, backend_to_kernel_cls(Fp8MoeBackend.TRITON)[0]
225 -
223 # NOTE: the kernels are selected in the following order.
224 AVAILABLE_BACKENDS = _get_priority_backends(config, weight_key, activation_key)
225
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
-13
@@ -214,19 +214,6 @@ def select_unquantized_moe_backend(
214 return backend, k_cls
215 raise ValueError(_make_log_unsupported(backend, reason))
216
217 - # LoRA needs Triton's unfused activation/reduction hooks. Selecting the
218 - # backend here ensures weights stay in a LoRA-compatible layout instead of
219 - # being permuted for a backend like FlashInfer or AITER during load.
220 - if moe_config.is_lora_enabled:
221 - backend = UnquantizedMoeBackend.TRITON
222 - if activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
223 - backend = UnquantizedMoeBackend.BATCHED_TRITON
224 - return _return_or_raise(
225 - backend,
226 - moe_config,
227 - activation_format,
228 - )
229 -
217 runner_backend = moe_config.moe_backend
218 if runner_backend != "auto":
219 requested_backend = map_unquantized_backend(runner_backend)
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+5 -1
@@ -297,7 +297,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
297 shared_experts_input: torch.Tensor | None,
298 ) -> torch.Tensor:
299 return self.forward_native(
300 - layer, x, topk_weights, topk_ids, shared_experts_input
300 + layer,
301 + x,
302 + topk_weights,
303 + topk_ids,
304 + shared_experts_input,
305 )
306
307 def apply_monolithic(