From ecd0b60aad2f4e28dd00ababfc1402690d88cbed Mon Sep 17 00:00:00 2001 From: Jee Jee Li Date: Sat, 9 May 2026 15:31:23 +0800 Subject: [PATCH] [LoRA] Initial EP support for LoRA (#40867) Signed-off-by: Jee Jee Li Signed-off-by: Jee Jee Li --- tests/lora/test_gptoss_tp.py | 1 + tests/lora/test_qwen3moe_tp.py | 14 ++-- vllm/lora/layers/fused_moe.py | 54 ++++++++++---- vllm/lora/model_manager.py | 72 ++++++++++++++++--- vllm/lora/punica_wrapper/punica_base.py | 5 ++ vllm/lora/punica_wrapper/punica_gpu.py | 52 ++++++++++---- vllm/lora/punica_wrapper/punica_xpu.py | 1 + .../layers/fused_moe/lora_context.py | 7 ++ .../layers/fused_moe/lora_experts_mixin.py | 1 + .../layers/fused_moe/oracle/int8.py | 3 - .../layers/fused_moe/oracle/mxfp4.py | 22 ------ .../layers/fused_moe/oracle/mxfp8.py | 2 - .../fused_moe/prepare_finalize/naive_dp_ep.py | 48 +++++++++++-- 13 files changed, 209 insertions(+), 73 deletions(-) diff --git a/tests/lora/test_gptoss_tp.py b/tests/lora/test_gptoss_tp.py index 68dd87233ac..64866073465 100644 --- a/tests/lora/test_gptoss_tp.py +++ b/tests/lora/test_gptoss_tp.py @@ -129,6 +129,7 @@ def test_gpt_oss_lora_tp2( tensor_parallel_size=2, gpu_memory_utilization=0.8, fully_sharded_loras=fully_sharded_loras, + enable_expert_parallel=not fully_sharded_loras, compilation_config=vllm.config.CompilationConfig( # Avoid OOM cudagraph_specialize_lora=False, ), diff --git a/tests/lora/test_qwen3moe_tp.py b/tests/lora/test_qwen3moe_tp.py index fcac4275cc4..9af142f6f38 100644 --- a/tests/lora/test_qwen3moe_tp.py +++ b/tests/lora/test_qwen3moe_tp.py @@ -5,6 +5,8 @@ # NOTE To avoid overloading the CI pipeline, this test script will not # be triggered on CI and is primarily intended for local testing and verification. +import pytest + import vllm from vllm.lora.request import LoRARequest @@ -82,15 +84,15 @@ def test_qwen3moe_lora(qwen3moe_lora_files): @multi_gpu_test(num_gpus=2) -def test_qwen3moe_lora_tp2(qwen3moe_lora_files): +@pytest.mark.parametrize("ep", [False, True]) +def test_qwen3moe_lora_tp2(ep, qwen3moe_lora_files): llm = vllm.LLM( MODEL_PATH, max_model_len=1024, enable_lora=True, max_loras=4, - enforce_eager=True, trust_remote_code=True, - enable_chunked_prefill=True, + enable_expert_parallel=ep, tensor_parallel_size=2, ) @@ -99,15 +101,15 @@ def test_qwen3moe_lora_tp2(qwen3moe_lora_files): @multi_gpu_test(num_gpus=4) -def test_qwen3moe_lora_tp4(qwen3moe_lora_files): +@pytest.mark.parametrize("ep", [False, True]) +def test_qwen3moe_lora_tp4(ep, qwen3moe_lora_files): llm = vllm.LLM( MODEL_PATH, max_model_len=1024, enable_lora=True, max_loras=4, - enforce_eager=True, trust_remote_code=True, - enable_chunked_prefill=True, + enable_expert_parallel=ep, tensor_parallel_size=4, ) diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 2f9a4701b0d..8cb32f07965 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -7,10 +7,6 @@ from transformers import PretrainedConfig from vllm import envs from vllm.config.lora import LoRAConfig -from vllm.distributed.parallel_state import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) from vllm.distributed.utils import divide from vllm.lora.layers.base import BaseLayerWithLoRA from vllm.model_executor.layers.fused_moe import FusedMoE @@ -30,15 +26,12 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): def __init__(self, base_layer: FusedMoE) -> None: super().__init__() self.base_layer = base_layer - - assert not self.base_layer.use_ep, ( - "EP support for Fused MoE LoRA is not implemented yet." - ) - assert not self.base_layer.quant_method.is_monolithic, ( - "Monolithic kernels are not supported for Fused MoE LoRA." - ) - self.tp_size = get_tensor_model_parallel_world_size() - self.tp_rank = get_tensor_model_parallel_rank() + self._ep_check() + # Use the MoE-aware TP rank/size: when EP is active, FusedMoE collapses + # moe_parallel_config.tp_size to 1 (experts are sharded across the + # TP group instead). + self.tp_size = self.base_layer.tp_size + self.tp_rank = self.base_layer.tp_rank self.device = _get_lora_device(base_layer) # For non-gated MoE (is_act_and_mul=False), only 1 slice is needed # since there's only up_proj (w1), not gate_proj + up_proj (w1 + w3) @@ -65,7 +58,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): "For quantized MoE, mix LoRAExpertsMixin into the experts class " "and consume self._lora_context in apply()." ) - self._fused_experts = moe_kernel.fused_experts + self._moe_kernel = moe_kernel self.base_layer._replace_quant_method( FusedMoEModularMethod(self.base_layer.quant_method, moe_kernel) ) @@ -150,6 +143,26 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): ), ) + def _ep_check(self): + if self.base_layer.use_ep: + moe_config = self.base_layer.moe_config + all2all_backend = moe_config.moe_parallel_config.all2all_backend + assert all2all_backend == "allgather_reducescatter", ( + "Fused MoE LoRA with EP currently only supports " + f"all2all_backend='allgather_reducescatter', got '{all2all_backend}'." + ) + assert not moe_config.moe_parallel_config.is_sequence_parallel + + def _verify_ep_fs(self, lora_config: LoRAConfig): + # EP and fully_sharded LoRA both partition along the same TP group — + # EP on the expert dim, fully_sharded on the LoRA rank dim — with + # mutually contradictory assumptions about which rank holds which + # expert's rank-shard. + assert not (self.base_layer.use_ep and lora_config.fully_sharded_loras), ( + "Fused MoE LoRA does not support enable_expert_parallel=True " + "together with fully_sharded_loras=True. Disable one of them." + ) + def create_lora_weights( self, max_loras: int, @@ -157,6 +170,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): model_config: PretrainedConfig | None = None, ) -> None: """Initializes lora matrices.""" + + self._verify_ep_fs(lora_config) self.max_loras = lora_config.max_loras self.fully_sharded = lora_config.fully_sharded_loras @@ -282,6 +297,10 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): w1_lora_a, w2_lora_a, w3_lora_a = lora_a w1_lora_b, w2_lora_b, w3_lora_b = lora_b + + # EP slicing is done once at add time in + # LoRAModelManager._slice_moe_lora_ep, so by here the cached + # tensors already match the local-expert dim of the stacked buffers. assert ( num_experts == w1_lora_a.shape[0] @@ -326,7 +345,11 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): def set_mapping(self, punica_wrapper): super().set_mapping(punica_wrapper) - self._fused_experts.set_lora_context(self._build_lora_context()) + lora_context = self._build_lora_context() + self._moe_kernel.fused_experts.set_lora_context(lora_context) + prepare_finalize = self._moe_kernel.prepare_finalize + if hasattr(prepare_finalize, "set_lora_context"): + prepare_finalize.set_lora_context(lora_context) def forward(self, *args, **kwargs): return self.base_layer.forward(*args, **kwargs) @@ -400,6 +423,7 @@ class FusedMoE3DWithLoRA(FusedMoEWithLoRA): """Initializes lora matrices.""" assert isinstance(model_config, PretrainedConfig) + self._verify_ep_fs(lora_config) self._base_model = model_config.architectures[0] self.max_loras = lora_config.max_loras self.fully_sharded = lora_config.fully_sharded_loras diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index 52ff8ebc91f..015a783342b 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -14,6 +14,7 @@ from vllm.logger import init_logger from vllm.lora.layers import ( BaseLayerWithLoRA, FusedMoE3DWithLoRA, + FusedMoEWithLoRA, LoRAMapping, LoRAMappingType, ) @@ -562,6 +563,10 @@ class LoRAModelManager: else: parts = module_name.split(".") replacements = self.packed_modules_mapping[parts[-1]] + if module.__class__.__name__ == "FusedMoEWithLoRA": + replacements = replacements[ + : len(module.lora_a_stacked) // self.lora_slots + ] subloras: list[LoRALayerWeights | None] = [] for i, r in enumerate(replacements): lora = LoRALayerWeights.create_dummy_lora_weights( @@ -716,6 +721,8 @@ class LoRAModelManager: for module_name, module in self.modules.items(): if isinstance(module, FusedMoE3DWithLoRA): self._stack_moe_lora_weights(lora_model, module, module_name) + elif isinstance(module, FusedMoEWithLoRA): + self._slice_moe_lora_ep(lora_model, module, module_name) first_lora: LoRALayerWeights = next(iter(lora_model.loras.values())) assert first_lora.lora_a is not None @@ -762,23 +769,33 @@ class LoRAModelManager: assert gate_up_proj_lora is not None assert down_proj_lora is not None if self._is_3d_moe_model: - num_experts = module.w13_lora_a_stacked[0].shape[1] + local_num_experts = module.w13_lora_a_stacked[0].shape[1] + # The checkpoint holds weights for all global experts, but + # each EP rank owns only local_num_experts. Reshape against + # the adapter's actual expert count, then slice this rank's + # owned expert range before it gets copied into the local + # stacked buffer. For non-EP (local == global) this is a + # no-op slice. + global_num_experts = module.base_layer.global_num_experts + ep_rank = module.base_layer.ep_rank + expert_start = ep_rank * local_num_experts + expert_end = expert_start + local_num_experts # (num_experts,rank,input_size) gate_up_proj_lora.lora_a = gate_up_proj_lora.lora_a.reshape( - num_experts, -1, gate_up_proj_lora.lora_a.shape[-1] - ) + global_num_experts, -1, gate_up_proj_lora.lora_a.shape[-1] + )[expert_start:expert_end].contiguous() down_proj_lora.lora_a = down_proj_lora.lora_a.reshape( - num_experts, -1, down_proj_lora.lora_a.shape[-1] - ) + global_num_experts, -1, down_proj_lora.lora_a.shape[-1] + )[expert_start:expert_end].contiguous() # (output_size,rank,num_experts) gate_up_proj_lora.lora_b = gate_up_proj_lora.lora_b.reshape( - gate_up_proj_lora.lora_b.shape[0], -1, num_experts - ) + gate_up_proj_lora.lora_b.shape[0], -1, global_num_experts + )[..., expert_start:expert_end] down_proj_lora.lora_b = down_proj_lora.lora_b.reshape( - down_proj_lora.lora_b.shape[0], -1, num_experts - ) + down_proj_lora.lora_b.shape[0], -1, global_num_experts + )[..., expert_start:expert_end] # (num_experts,output_size,rank) gate_up_proj_lora.lora_b = gate_up_proj_lora.lora_b.permute( @@ -828,6 +845,43 @@ class LoRAModelManager: module_lora.lora_a = lora_a module_lora.lora_b = lora_b + def _slice_moe_lora_ep( + self, + lora_model: LoRAModel, + module: FusedMoEWithLoRA, + module_name: str, + ) -> None: + """Slice the cached LoRA tensors down to this rank's local experts. + + The 2D MoE checkpoint enters as a list of per-(w1/w2/w3) tensors of + shape (num_experts, rank, in) / (num_experts, out, rank). When EP + is active each rank only owns local_num_experts; without this slice + the CPU LoRAModel keeps the full global weight and set_lora has to + re-slice on every activation. + """ + if not module.base_layer.use_ep: + return + module_lora = self._get_lora_layer_weights(lora_model, module_name) + if module_lora is None or not isinstance(module_lora.lora_a, list): + return + + local_num_experts = module.base_layer.local_num_experts + global_num_experts = module.base_layer.global_num_experts + ep_rank = module.base_layer.ep_rank + expert_start = ep_rank * local_num_experts + expert_end = expert_start + local_num_experts + + new_lora_a: list[torch.Tensor | None] = [] + new_lora_b: list[torch.Tensor | None] = [] + for a, b in zip(module_lora.lora_a, module_lora.lora_b): + if a is not None and b is not None and a.shape[0] == global_num_experts: + a = a[expert_start:expert_end].contiguous() + b = b[expert_start:expert_end].contiguous() + new_lora_a.append(a) + new_lora_b.append(b) + module_lora.lora_a = new_lora_a + module_lora.lora_b = new_lora_b + def _get_lora_layer_weights( self, lora_model: LoRAModel, module_name: str ) -> LoRALayerWeights | None: diff --git a/vllm/lora/punica_wrapper/punica_base.py b/vllm/lora/punica_wrapper/punica_base.py index 4ab66dccdc2..0448a6d00cd 100644 --- a/vllm/lora/punica_wrapper/punica_base.py +++ b/vllm/lora/punica_wrapper/punica_base.py @@ -514,6 +514,7 @@ class PunicaWrapperBase(PunicaWrapperABC): num_slices: int, fully_sharded: bool, use_tuned_config: bool, + token_lora_mapping: torch.Tensor | None = None, ) -> tuple[ torch.Tensor | None, torch.Tensor | None, @@ -522,6 +523,10 @@ class PunicaWrapperBase(PunicaWrapperABC): ]: """Apply w13 LoRA to y (intermediate_cache1) in-place before activation. + When `token_lora_mapping` is provided it overrides the punica_wrapper's + global mapping — used by EP+LoRA to pass the per-rank-local mapping + after all-to-all dispatch. + Returns (sorted_token_ids_lora, expert_ids_lora, num_tokens_post_padded_lora, token_lora_mapping) for reuse by add_lora_w2. diff --git a/vllm/lora/punica_wrapper/punica_gpu.py b/vllm/lora/punica_wrapper/punica_gpu.py index 44d1dbd5072..bf951e07494 100644 --- a/vllm/lora/punica_wrapper/punica_gpu.py +++ b/vllm/lora/punica_wrapper/punica_gpu.py @@ -335,25 +335,49 @@ class PunicaWrapperGPU(PunicaWrapperBase): expert_map: torch.Tensor | None = None, pad_sorted_ids: bool = False, naive_block_assignment: bool = False, + token_lora_mapping: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """ Aligns tokens and experts into block-sized chunks for LoRA-based mixture-of-experts (MoE) execution. + + When `token_lora_mapping` is provided, it overrides the global mapping + read from `self.token_mapping_meta`. This is how EP+LoRA injects the + per-rank-local token→LoRA map after all-to-all dispatch. """ - (token_lora_mapping, _, _, _, lora_ids, _, _) = ( - self.token_mapping_meta.meta_args( - num_tokens, self.lora_config.specialize_active_lora - ) + ( + token_lora_mapping_meta, + _, + _, + _, + lora_ids, + _, + _, + ) = self.token_mapping_meta.meta_args( + num_tokens, self.lora_config.specialize_active_lora + ) + if token_lora_mapping is None: + token_lora_mapping = token_lora_mapping_meta + # Under EP the caller passes local_num_experts but topk_ids carries + # GLOBAL expert indices. The CUDA kernel uses num_experts to size + # its bucketing table; with EP we must size by global_num_experts + # so global topk_ids don't overflow. expert_map inside the kernel + # then translates global→local so the output expert_ids are local + # (mirrors the non-LoRA moe_align_block_size behavior). + kernel_num_experts = ( + expert_map.numel() if expert_map is not None else num_experts ) if naive_block_assignment: expert_ids = topk_ids.reshape(-1) sorted_ids = None num_tokens_post_pad = None else: - max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1) + max_num_tokens_padded = topk_ids.numel() + kernel_num_experts * ( + block_size - 1 + ) if pad_sorted_ids: max_num_tokens_padded = round_up(max_num_tokens_padded, block_size) - if topk_ids.numel() < num_experts: + if topk_ids.numel() < kernel_num_experts: max_num_tokens_padded = topk_ids.numel() * block_size sorted_ids = torch.empty( (max_loras * max_num_tokens_padded,), @@ -361,9 +385,12 @@ class PunicaWrapperGPU(PunicaWrapperBase): device=topk_ids.device, ) max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size) - # Expert ids must be set default to -1 to prevent a blank block - expert_ids = torch.empty( + # Expert ids are initialized to -1 so unused (lora, expert) + # slots don't drive the LoRA Triton kernel into the wrong bucket. + # The kernel overwrites only active slots. + expert_ids = torch.full( (max_loras * max_num_m_blocks,), + -1, dtype=torch.int32, device=topk_ids.device, ) @@ -374,7 +401,7 @@ class PunicaWrapperGPU(PunicaWrapperBase): ops.moe_lora_align_block_size( topk_ids, token_lora_mapping, - num_experts, + kernel_num_experts, block_size, max_loras, max_num_tokens_padded, @@ -384,11 +411,10 @@ class PunicaWrapperGPU(PunicaWrapperBase): num_tokens_post_pad, adapter_enabled, lora_ids, + expert_map, ) - if expert_map is not None: - expert_ids = expert_map[expert_ids] - return None, sorted_ids, expert_ids, num_tokens_post_pad + return token_lora_mapping, sorted_ids, expert_ids, num_tokens_post_pad def add_lora_fused_moe( self, @@ -480,6 +506,7 @@ class PunicaWrapperGPU(PunicaWrapperBase): num_slices: int, fully_sharded: bool, use_tuned_config: bool, + token_lora_mapping: torch.Tensor | None = None, ) -> tuple[ torch.Tensor | None, torch.Tensor | None, @@ -558,6 +585,7 @@ class PunicaWrapperGPU(PunicaWrapperBase): adapter_enabled, expert_map, naive_block_assignment=naive_block_assignment, + token_lora_mapping=token_lora_mapping, ) _sorted = sorted_token_ids_lora diff --git a/vllm/lora/punica_wrapper/punica_xpu.py b/vllm/lora/punica_wrapper/punica_xpu.py index 20dde67b068..58316cb7597 100755 --- a/vllm/lora/punica_wrapper/punica_xpu.py +++ b/vllm/lora/punica_wrapper/punica_xpu.py @@ -461,6 +461,7 @@ class PunicaWrapperXPU(PunicaWrapperBase): num_slices: int, fully_sharded: bool, use_tuned_config: bool, + token_lora_mapping: torch.Tensor | None = None, ) -> tuple[ torch.Tensor | None, torch.Tensor | None, diff --git a/vllm/model_executor/layers/fused_moe/lora_context.py b/vllm/model_executor/layers/fused_moe/lora_context.py index 92500a7bb47..ab1f0bfc147 100644 --- a/vllm/model_executor/layers/fused_moe/lora_context.py +++ b/vllm/model_executor/layers/fused_moe/lora_context.py @@ -42,3 +42,10 @@ class MoELoRAContext: # 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 + + # Per-rank token→LoRA mapping after EP dispatch. Set by + # FusedMoEPrepareAndFinalizeModular.prepare() when EP+LoRA is active, read + # by LoRAExpertsMixin helpers in place of punica_wrapper's global mapping. + # None means no dispatch happened (non-EP path), in which case callers + # fall back to punica_wrapper.token_mapping_meta. + local_token_lora_mapping: torch.Tensor | None = None diff --git a/vllm/model_executor/layers/fused_moe/lora_experts_mixin.py b/vllm/model_executor/layers/fused_moe/lora_experts_mixin.py index c609c5cf56b..10707b91b70 100644 --- a/vllm/model_executor/layers/fused_moe/lora_experts_mixin.py +++ b/vllm/model_executor/layers/fused_moe/lora_experts_mixin.py @@ -70,6 +70,7 @@ class LoRAExpertsMixin: lora_context.w13_num_slices, lora_context.fully_sharded, lora_context.use_tuned_config, + token_lora_mapping=lora_context.local_token_lora_mapping, ) def apply_w2_lora( diff --git a/vllm/model_executor/layers/fused_moe/oracle/int8.py b/vllm/model_executor/layers/fused_moe/oracle/int8.py index cdb1be108b5..ebdd20d54dc 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int8.py @@ -79,9 +79,6 @@ def select_int8_moe_backend( Note: Shape-specific fallbacks may still occur at runtime. """ - if config.is_lora_enabled: - return Int8MoeBackend.TRITON, backend_to_kernel_cls(Int8MoeBackend.TRITON)[0] - AVAILABLE_BACKENDS = _get_priority_backends(config) activation_format = ( diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 1cc94a347c4..f2742fd18b7 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -301,28 +301,6 @@ def select_mxfp4_moe_backend( """ # If activation_key is explicitly provided (e.g., W4A8), use it requested_activation_key = activation_key - device_capability = current_platform.get_device_capability() - triton_kernels_supported = ( - has_triton_kernels() - and device_capability is not None - and (9, 0) <= device_capability < (11, 0) - ) - - # LoRA: separate experts backend path - if config.is_lora_enabled: - if not current_platform.is_cuda(): - # ROCm: Triton mxfp4 LoRA hits GPU memory faults due to - # triton_kernels.tensor.Tensor / HIP read-only page issues - # during weight swizzle and LoRA forward. Needs work from - # the triton_kernels/aiter side. - raise NotImplementedError("Mxfp4 LoRA is currently only supported on CUDA.") - if envs.VLLM_MXFP4_USE_MARLIN is False and triton_kernels_supported: - logger.info_once("Using Triton backend for mxfp4 lora") - return Mxfp4MoeBackend.TRITON_UNFUSED, backend_to_kernel_cls( - Mxfp4MoeBackend.TRITON_UNFUSED - )[0] - logger.info_once("Using Marlin backend for mxfp4 lora") - return Mxfp4MoeBackend.MARLIN, backend_to_kernel_cls(Mxfp4MoeBackend.MARLIN)[0] activation_format = ( mk.FusedMoEActivationFormat.BatchedExperts diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py index c67def149b9..8133902d519 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py @@ -61,8 +61,6 @@ def select_mxfp8_moe_backend( Returns: A tuple of (fp8_backend, experts_cls). """ - if config.is_lora_enabled: - raise NotImplementedError("LoRA is not supported for MXFP8 MoE.") runner_backend = config.moe_backend if runner_backend != "auto": diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py index 54d77101a3f..b8633726c72 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py @@ -84,6 +84,14 @@ class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular super().__init__() self.is_sequence_parallel = is_sequence_parallel self._num_dispatchers = num_dispatchers + # Set by FusedMoEWithLoRA.set_mapping() when LoRA is active. When + # present, prepare() dispatches the per-token LoRA mapping alongside + # hidden_states and writes the gathered result back to the context so + # experts can use the per-rank-local mapping. + self._lora_context = None + + def set_lora_context(self, ctx) -> None: + self._lora_context = ctx @property def activation_format(self) -> mk.FusedMoEActivationFormat: @@ -124,22 +132,54 @@ class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular a1q, scales = _quantize_and_setup_dispatch(a1, quant_config, defer_input_quant) + # When LoRA is active, dispatch the per-token LoRA id along with + # hidden_states so every rank receives the correct mapping for the + # tokens it ends up processing. The punica_wrapper stores indices as + # int64 but the moe_lora_align_block_size kernel expects int32, so + # pull the pre-cast view from token_mapping_meta. + lora_ctx = self._lora_context + local_token_lora_mapping = None + if lora_ctx is not None: + local_token_lora_mapping = ( + lora_ctx.punica_wrapper.token_mapping_meta.token_lora_mapping[ + : a1.shape[0] + ] + ) + + extra_tensors: list[torch.Tensor] | None = None + if scales is not None: + extra_tensors = list(scales) + if local_token_lora_mapping is not None: + if extra_tensors is None: + extra_tensors = [] + extra_tensors.append(local_token_lora_mapping) + res = get_ep_group().dispatch( a1q, topk_weights, topk_ids, is_sequence_parallel=self.is_sequence_parallel, - extra_tensors=scales, + extra_tensors=extra_tensors, ) - if scales is None: + if extra_tensors is None: assert len(res) == 3 a1q, topk_weights, topk_ids = res a1q_scale = None else: assert len(res) == 4 - a1q, topk_weights, topk_ids, scales = res - a1q_scale = _unwrap_scale_and_prepare_for_moe(scales, quant_config) + a1q, topk_weights, topk_ids, gathered_extras = res + gathered_extras = list(gathered_extras) + if local_token_lora_mapping is not None: + dispatched_lora_mapping = gathered_extras.pop() + assert lora_ctx is not None + lora_ctx.local_token_lora_mapping = dispatched_lora_mapping + if scales is not None: + a1q_scale = _unwrap_scale_and_prepare_for_moe( + gathered_extras, quant_config + ) + else: + a1q_scale = None return a1q, a1q_scale, None, topk_ids, topk_weights