@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ecd0b60aad

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

[LoRA] Initial EP support for LoRA (#40867)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com> Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

Jee Jee Li committed May 9, 2026 at 15:31 UTC ecd0b60aad2f4e28dd00ababfc1402690d88cbed

13 files changed +209 -73

tests/lora/test_gptoss_tp.py

@@ -129,6 +129,7 @@ def test_gpt_oss_lora_tp2(

129 tensor_parallel_size=2,

130 gpu_memory_utilization=0.8,

131 fully_sharded_loras=fully_sharded_loras,

132 + enable_expert_parallel=not fully_sharded_loras,

133 compilation_config=vllm.config.CompilationConfig( # Avoid OOM

134 cudagraph_specialize_lora=False,

135 ),

tests/lora/test_qwen3moe_tp.py

+8 -6

@@ -5,6 +5,8 @@

5 # NOTE To avoid overloading the CI pipeline, this test script will not

6 # be triggered on CI and is primarily intended for local testing and verification.

8 +import pytest

9 +

10 import vllm

11 from vllm.lora.request import LoRARequest

@@ -82,15 +84,15 @@ def test_qwen3moe_lora(qwen3moe_lora_files):

86 @multi_gpu_test(num_gpus=2)

85 -def test_qwen3moe_lora_tp2(qwen3moe_lora_files):

87 +@pytest.mark.parametrize("ep", [False, True])

88 +def test_qwen3moe_lora_tp2(ep, qwen3moe_lora_files):

89 llm = vllm.LLM(

90 MODEL_PATH,

91 max_model_len=1024,

92 enable_lora=True,

93 max_loras=4,

91 - enforce_eager=True,

94 trust_remote_code=True,

93 - enable_chunked_prefill=True,

95 + enable_expert_parallel=ep,

96 tensor_parallel_size=2,

97 )

@@ -99,15 +101,15 @@ def test_qwen3moe_lora_tp2(qwen3moe_lora_files):

101

102

103 @multi_gpu_test(num_gpus=4)

102 -def test_qwen3moe_lora_tp4(qwen3moe_lora_files):

104 +@pytest.mark.parametrize("ep", [False, True])

105 +def test_qwen3moe_lora_tp4(ep, qwen3moe_lora_files):

106 llm = vllm.LLM(

107 MODEL_PATH,

108 max_model_len=1024,

109 enable_lora=True,

110 max_loras=4,

108 - enforce_eager=True,

111 trust_remote_code=True,

110 - enable_chunked_prefill=True,

112 + enable_expert_parallel=ep,

113 tensor_parallel_size=4,

114 )

115

vllm/lora/layers/fused_moe.py

+39 -15

@@ -7,10 +7,6 @@ from transformers import PretrainedConfig

8 from vllm import envs

9 from vllm.config.lora import LoRAConfig

10 -from vllm.distributed.parallel_state import (

11 - get_tensor_model_parallel_rank,

12 - get_tensor_model_parallel_world_size,

13 -)

10 from vllm.distributed.utils import divide

11 from vllm.lora.layers.base import BaseLayerWithLoRA

12 from vllm.model_executor.layers.fused_moe import FusedMoE

@@ -30,15 +26,12 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

26 def __init__(self, base_layer: FusedMoE) -> None:

27 super().__init__()

28 self.base_layer = base_layer

33 -

34 - assert not self.base_layer.use_ep, (

35 - "EP support for Fused MoE LoRA is not implemented yet."

36 - )

37 - assert not self.base_layer.quant_method.is_monolithic, (

38 - "Monolithic kernels are not supported for Fused MoE LoRA."

39 - )

40 - self.tp_size = get_tensor_model_parallel_world_size()

41 - self.tp_rank = get_tensor_model_parallel_rank()

29 + self._ep_check()

30 + # Use the MoE-aware TP rank/size: when EP is active, FusedMoE collapses

31 + # moe_parallel_config.tp_size to 1 (experts are sharded across the

32 + # TP group instead).

33 + self.tp_size = self.base_layer.tp_size

34 + self.tp_rank = self.base_layer.tp_rank

35 self.device = _get_lora_device(base_layer)

36 # For non-gated MoE (is_act_and_mul=False), only 1 slice is needed

37 # since there's only up_proj (w1), not gate_proj + up_proj (w1 + w3)

@@ -65,7 +58,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

58 "For quantized MoE, mix LoRAExpertsMixin into the experts class "

59 "and consume self._lora_context in apply()."

60 )

68 - self._fused_experts = moe_kernel.fused_experts

61 + self._moe_kernel = moe_kernel

62 self.base_layer._replace_quant_method(

63 FusedMoEModularMethod(self.base_layer.quant_method, moe_kernel)

64 )

@@ -150,6 +143,26 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

143 ),

144 )

145

146 + def _ep_check(self):

147 + if self.base_layer.use_ep:

148 + moe_config = self.base_layer.moe_config

149 + all2all_backend = moe_config.moe_parallel_config.all2all_backend

150 + assert all2all_backend == "allgather_reducescatter", (

151 + "Fused MoE LoRA with EP currently only supports "

152 + f"all2all_backend='allgather_reducescatter', got '{all2all_backend}'."

153 + )

154 + assert not moe_config.moe_parallel_config.is_sequence_parallel

155 +

156 + def _verify_ep_fs(self, lora_config: LoRAConfig):

157 + # EP and fully_sharded LoRA both partition along the same TP group —

158 + # EP on the expert dim, fully_sharded on the LoRA rank dim — with

159 + # mutually contradictory assumptions about which rank holds which

160 + # expert's rank-shard.

161 + assert not (self.base_layer.use_ep and lora_config.fully_sharded_loras), (

162 + "Fused MoE LoRA does not support enable_expert_parallel=True "

163 + "together with fully_sharded_loras=True. Disable one of them."

164 + )

165 +

166 def create_lora_weights(

167 self,

168 max_loras: int,

@@ -157,6 +170,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

170 model_config: PretrainedConfig | None = None,

171 ) -> None:

172 """Initializes lora matrices."""

173 +

174 + self._verify_ep_fs(lora_config)

175 self.max_loras = lora_config.max_loras

176 self.fully_sharded = lora_config.fully_sharded_loras

177

@@ -282,6 +297,10 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

297

298 w1_lora_a, w2_lora_a, w3_lora_a = lora_a

299 w1_lora_b, w2_lora_b, w3_lora_b = lora_b

300 +

301 + # EP slicing is done once at add time in

302 + # LoRAModelManager._slice_moe_lora_ep, so by here the cached

303 + # tensors already match the local-expert dim of the stacked buffers.

304 assert (

305 num_experts

306 == w1_lora_a.shape[0]

@@ -326,7 +345,11 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

345

346 def set_mapping(self, punica_wrapper):

347 super().set_mapping(punica_wrapper)

329 - self._fused_experts.set_lora_context(self._build_lora_context())

348 + lora_context = self._build_lora_context()

349 + self._moe_kernel.fused_experts.set_lora_context(lora_context)

350 + prepare_finalize = self._moe_kernel.prepare_finalize

351 + if hasattr(prepare_finalize, "set_lora_context"):

352 + prepare_finalize.set_lora_context(lora_context)

353

354 def forward(self, *args, **kwargs):

355 return self.base_layer.forward(*args, **kwargs)

@@ -400,6 +423,7 @@ class FusedMoE3DWithLoRA(FusedMoEWithLoRA):

423 """Initializes lora matrices."""

424

425 assert isinstance(model_config, PretrainedConfig)

426 + self._verify_ep_fs(lora_config)

427 self._base_model = model_config.architectures[0]

428 self.max_loras = lora_config.max_loras

429 self.fully_sharded = lora_config.fully_sharded_loras

vllm/lora/model_manager.py

+63 -9

@@ -14,6 +14,7 @@ from vllm.logger import init_logger

14 from vllm.lora.layers import (

15 BaseLayerWithLoRA,

16 FusedMoE3DWithLoRA,

17 + FusedMoEWithLoRA,

18 LoRAMapping,

19 LoRAMappingType,

20 )

@@ -562,6 +563,10 @@ class LoRAModelManager:

563 else:

564 parts = module_name.split(".")

565 replacements = self.packed_modules_mapping[parts[-1]]

566 + if module.__class__.__name__ == "FusedMoEWithLoRA":

567 + replacements = replacements[

568 + : len(module.lora_a_stacked) // self.lora_slots

569 + ]

570 subloras: list[LoRALayerWeights | None] = []

571 for i, r in enumerate(replacements):

572 lora = LoRALayerWeights.create_dummy_lora_weights(

@@ -716,6 +721,8 @@ class LoRAModelManager:

721 for module_name, module in self.modules.items():

722 if isinstance(module, FusedMoE3DWithLoRA):

723 self._stack_moe_lora_weights(lora_model, module, module_name)

724 + elif isinstance(module, FusedMoEWithLoRA):

725 + self._slice_moe_lora_ep(lora_model, module, module_name)

726

727 first_lora: LoRALayerWeights = next(iter(lora_model.loras.values()))

728 assert first_lora.lora_a is not None

@@ -762,23 +769,33 @@ class LoRAModelManager:

769 assert gate_up_proj_lora is not None

770 assert down_proj_lora is not None

771 if self._is_3d_moe_model:

765 - num_experts = module.w13_lora_a_stacked[0].shape[1]

772 + local_num_experts = module.w13_lora_a_stacked[0].shape[1]

773 + # The checkpoint holds weights for all global experts, but

774 + # each EP rank owns only local_num_experts. Reshape against

775 + # the adapter's actual expert count, then slice this rank's

776 + # owned expert range before it gets copied into the local

777 + # stacked buffer. For non-EP (local == global) this is a

778 + # no-op slice.

779 + global_num_experts = module.base_layer.global_num_experts

780 + ep_rank = module.base_layer.ep_rank

781 + expert_start = ep_rank * local_num_experts

782 + expert_end = expert_start + local_num_experts

783

784 # (num_experts,rank,input_size)

785 gate_up_proj_lora.lora_a = gate_up_proj_lora.lora_a.reshape(

769 - num_experts, -1, gate_up_proj_lora.lora_a.shape[-1]

770 - )

786 + global_num_experts, -1, gate_up_proj_lora.lora_a.shape[-1]

787 + )[expert_start:expert_end].contiguous()

788 down_proj_lora.lora_a = down_proj_lora.lora_a.reshape(

772 - num_experts, -1, down_proj_lora.lora_a.shape[-1]

773 - )

789 + global_num_experts, -1, down_proj_lora.lora_a.shape[-1]

790 + )[expert_start:expert_end].contiguous()

791

792 # (output_size,rank,num_experts)

793 gate_up_proj_lora.lora_b = gate_up_proj_lora.lora_b.reshape(

777 - gate_up_proj_lora.lora_b.shape[0], -1, num_experts

778 - )

794 + gate_up_proj_lora.lora_b.shape[0], -1, global_num_experts

795 + )[..., expert_start:expert_end]

796 down_proj_lora.lora_b = down_proj_lora.lora_b.reshape(

780 - down_proj_lora.lora_b.shape[0], -1, num_experts

781 - )

797 + down_proj_lora.lora_b.shape[0], -1, global_num_experts

798 + )[..., expert_start:expert_end]

799

800 # (num_experts,output_size,rank)

801 gate_up_proj_lora.lora_b = gate_up_proj_lora.lora_b.permute(

@@ -828,6 +845,43 @@ class LoRAModelManager:

845 module_lora.lora_a = lora_a

846 module_lora.lora_b = lora_b

847

848 + def _slice_moe_lora_ep(

849 + self,

850 + lora_model: LoRAModel,

851 + module: FusedMoEWithLoRA,

852 + module_name: str,

853 + ) -> None:

854 + """Slice the cached LoRA tensors down to this rank's local experts.

855 +

856 + The 2D MoE checkpoint enters as a list of per-(w1/w2/w3) tensors of

857 + shape (num_experts, rank, in) / (num_experts, out, rank). When EP

858 + is active each rank only owns local_num_experts; without this slice

859 + the CPU LoRAModel keeps the full global weight and set_lora has to

860 + re-slice on every activation.

861 + """

862 + if not module.base_layer.use_ep:

863 + return

864 + module_lora = self._get_lora_layer_weights(lora_model, module_name)

865 + if module_lora is None or not isinstance(module_lora.lora_a, list):

866 + return

867 +

868 + local_num_experts = module.base_layer.local_num_experts

869 + global_num_experts = module.base_layer.global_num_experts

870 + ep_rank = module.base_layer.ep_rank

871 + expert_start = ep_rank * local_num_experts

872 + expert_end = expert_start + local_num_experts

873 +

874 + new_lora_a: list[torch.Tensor | None] = []

875 + new_lora_b: list[torch.Tensor | None] = []

876 + for a, b in zip(module_lora.lora_a, module_lora.lora_b):

877 + if a is not None and b is not None and a.shape[0] == global_num_experts:

878 + a = a[expert_start:expert_end].contiguous()

879 + b = b[expert_start:expert_end].contiguous()

880 + new_lora_a.append(a)

881 + new_lora_b.append(b)

882 + module_lora.lora_a = new_lora_a

883 + module_lora.lora_b = new_lora_b

884 +

885 def _get_lora_layer_weights(

886 self, lora_model: LoRAModel, module_name: str

887 ) -> LoRALayerWeights | None:

vllm/lora/punica_wrapper/punica_base.py

@@ -514,6 +514,7 @@ class PunicaWrapperBase(PunicaWrapperABC):

514 num_slices: int,

515 fully_sharded: bool,

516 use_tuned_config: bool,

517 + token_lora_mapping: torch.Tensor | None = None,

518 ) -> tuple[

519 torch.Tensor | None,

520 torch.Tensor | None,

@@ -522,6 +523,10 @@ class PunicaWrapperBase(PunicaWrapperABC):

523 ]:

524 """Apply w13 LoRA to y (intermediate_cache1) in-place before activation.

525

526 + When `token_lora_mapping` is provided it overrides the punica_wrapper's

527 + global mapping — used by EP+LoRA to pass the per-rank-local mapping

528 + after all-to-all dispatch.

529 +

530 Returns (sorted_token_ids_lora, expert_ids_lora,

531 num_tokens_post_padded_lora, token_lora_mapping)

532 for reuse by add_lora_w2.

vllm/lora/punica_wrapper/punica_gpu.py

+40 -12

@@ -335,25 +335,49 @@ class PunicaWrapperGPU(PunicaWrapperBase):

335 expert_map: torch.Tensor | None = None,

336 pad_sorted_ids: bool = False,

337 naive_block_assignment: bool = False,

338 + token_lora_mapping: torch.Tensor | None = None,

339 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

340 """

341 Aligns tokens and experts into block-sized chunks for LoRA-based

342 mixture-of-experts (MoE) execution.

343 +

344 + When `token_lora_mapping` is provided, it overrides the global mapping

345 + read from `self.token_mapping_meta`. This is how EP+LoRA injects the

346 + per-rank-local token→LoRA map after all-to-all dispatch.

347 """

343 - (token_lora_mapping, _, _, _, lora_ids, _, _) = (

344 - self.token_mapping_meta.meta_args(

345 - num_tokens, self.lora_config.specialize_active_lora

346 - )

348 + (

349 + token_lora_mapping_meta,

350 + _,

351 + _,

352 + _,

353 + lora_ids,

354 + _,

355 + _,

356 + ) = self.token_mapping_meta.meta_args(

357 + num_tokens, self.lora_config.specialize_active_lora

358 + )

359 + if token_lora_mapping is None:

360 + token_lora_mapping = token_lora_mapping_meta

361 + # Under EP the caller passes local_num_experts but topk_ids carries

362 + # GLOBAL expert indices. The CUDA kernel uses num_experts to size

363 + # its bucketing table; with EP we must size by global_num_experts

364 + # so global topk_ids don't overflow. expert_map inside the kernel

365 + # then translates global→local so the output expert_ids are local

366 + # (mirrors the non-LoRA moe_align_block_size behavior).

367 + kernel_num_experts = (

368 + expert_map.numel() if expert_map is not None else num_experts

369 )

370 if naive_block_assignment:

371 expert_ids = topk_ids.reshape(-1)

372 sorted_ids = None

373 num_tokens_post_pad = None

374 else:

353 - max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)

375 + max_num_tokens_padded = topk_ids.numel() + kernel_num_experts * (

376 + block_size - 1

377 + )

378 if pad_sorted_ids:

379 max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)

356 - if topk_ids.numel() < num_experts:

380 + if topk_ids.numel() < kernel_num_experts:

381 max_num_tokens_padded = topk_ids.numel() * block_size

382 sorted_ids = torch.empty(

383 (max_loras * max_num_tokens_padded,),

@@ -361,9 +385,12 @@ class PunicaWrapperGPU(PunicaWrapperBase):

385 device=topk_ids.device,

386 )

387 max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)

364 - # Expert ids must be set default to -1 to prevent a blank block

365 - expert_ids = torch.empty(

388 + # Expert ids are initialized to -1 so unused (lora, expert)

389 + # slots don't drive the LoRA Triton kernel into the wrong bucket.

390 + # The kernel overwrites only active slots.

391 + expert_ids = torch.full(

392 (max_loras * max_num_m_blocks,),

393 + -1,

394 dtype=torch.int32,

395 device=topk_ids.device,

396 )

@@ -374,7 +401,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):

401 ops.moe_lora_align_block_size(

402 topk_ids,

403 token_lora_mapping,

377 - num_experts,

404 + kernel_num_experts,

405 block_size,

406 max_loras,

407 max_num_tokens_padded,

@@ -384,11 +411,10 @@ class PunicaWrapperGPU(PunicaWrapperBase):

411 num_tokens_post_pad,

412 adapter_enabled,

413 lora_ids,

414 + expert_map,

415 )

388 - if expert_map is not None:

389 - expert_ids = expert_map[expert_ids]

416

391 - return None, sorted_ids, expert_ids, num_tokens_post_pad

417 + return token_lora_mapping, sorted_ids, expert_ids, num_tokens_post_pad

418

419 def add_lora_fused_moe(

420 self,

@@ -480,6 +506,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):

506 num_slices: int,

507 fully_sharded: bool,

508 use_tuned_config: bool,

509 + token_lora_mapping: torch.Tensor | None = None,

510 ) -> tuple[

511 torch.Tensor | None,

512 torch.Tensor | None,

@@ -558,6 +585,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):

585 adapter_enabled,

586 expert_map,

587 naive_block_assignment=naive_block_assignment,

588 + token_lora_mapping=token_lora_mapping,

589 )

590

591 _sorted = sorted_token_ids_lora

vllm/lora/punica_wrapper/punica_xpu.py

@@ -461,6 +461,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):

461 num_slices: int,

462 fully_sharded: bool,

463 use_tuned_config: bool,

464 + token_lora_mapping: torch.Tensor | None = None,

465 ) -> tuple[

466 torch.Tensor | None,

467 torch.Tensor | None,

vllm/model_executor/layers/fused_moe/lora_context.py

@@ -42,3 +42,10 @@ class MoELoRAContext:

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

45 +

46 + # Per-rank token→LoRA mapping after EP dispatch. Set by

47 + # FusedMoEPrepareAndFinalizeModular.prepare() when EP+LoRA is active, read

48 + # by LoRAExpertsMixin helpers in place of punica_wrapper's global mapping.

49 + # None means no dispatch happened (non-EP path), in which case callers

50 + # fall back to punica_wrapper.token_mapping_meta.

51 + local_token_lora_mapping: torch.Tensor | None = None

vllm/model_executor/layers/fused_moe/lora_experts_mixin.py

@@ -70,6 +70,7 @@ class LoRAExpertsMixin:

70 lora_context.w13_num_slices,

71 lora_context.fully_sharded,

72 lora_context.use_tuned_config,

73 + token_lora_mapping=lora_context.local_token_lora_mapping,

74 )

76 def apply_w2_lora(

vllm/model_executor/layers/fused_moe/oracle/int8.py

-3

@@ -79,9 +79,6 @@ def select_int8_moe_backend(

79 Note: Shape-specific fallbacks may still occur at runtime.

80 """

82 - if config.is_lora_enabled:

83 - return Int8MoeBackend.TRITON, backend_to_kernel_cls(Int8MoeBackend.TRITON)[0]

84 -

82 AVAILABLE_BACKENDS = _get_priority_backends(config)

84 activation_format = (

vllm/model_executor/layers/fused_moe/oracle/mxfp4.py

-22

@@ -301,28 +301,6 @@ def select_mxfp4_moe_backend(

301 """

302 # If activation_key is explicitly provided (e.g., W4A8), use it

303 requested_activation_key = activation_key

304 - device_capability = current_platform.get_device_capability()

305 - triton_kernels_supported = (

306 - has_triton_kernels()

307 - and device_capability is not None

308 - and (9, 0) <= device_capability < (11, 0)

309 - )

310 -

311 - # LoRA: separate experts backend path

312 - if config.is_lora_enabled:

313 - if not current_platform.is_cuda():

314 - # ROCm: Triton mxfp4 LoRA hits GPU memory faults due to

315 - # triton_kernels.tensor.Tensor / HIP read-only page issues

316 - # during weight swizzle and LoRA forward. Needs work from

317 - # the triton_kernels/aiter side.

318 - raise NotImplementedError("Mxfp4 LoRA is currently only supported on CUDA.")

319 - if envs.VLLM_MXFP4_USE_MARLIN is False and triton_kernels_supported:

320 - logger.info_once("Using Triton backend for mxfp4 lora")

321 - return Mxfp4MoeBackend.TRITON_UNFUSED, backend_to_kernel_cls(

322 - Mxfp4MoeBackend.TRITON_UNFUSED

323 - )[0]

324 - logger.info_once("Using Marlin backend for mxfp4 lora")

325 - return Mxfp4MoeBackend.MARLIN, backend_to_kernel_cls(Mxfp4MoeBackend.MARLIN)[0]

304

305 activation_format = (

306 mk.FusedMoEActivationFormat.BatchedExperts

vllm/model_executor/layers/fused_moe/oracle/mxfp8.py

-2

@@ -61,8 +61,6 @@ def select_mxfp8_moe_backend(

61 Returns:

62 A tuple of (fp8_backend, experts_cls).

63 """

64 - if config.is_lora_enabled:

65 - raise NotImplementedError("LoRA is not supported for MXFP8 MoE.")

65 runner_backend = config.moe_backend

66 if runner_backend != "auto":

vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py

+44 -4

@@ -84,6 +84,14 @@ class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular

84 super().__init__()

85 self.is_sequence_parallel = is_sequence_parallel

86 self._num_dispatchers = num_dispatchers

87 + # Set by FusedMoEWithLoRA.set_mapping() when LoRA is active. When

88 + # present, prepare() dispatches the per-token LoRA mapping alongside

89 + # hidden_states and writes the gathered result back to the context so

90 + # experts can use the per-rank-local mapping.

91 + self._lora_context = None

92 +

93 + def set_lora_context(self, ctx) -> None:

94 + self._lora_context = ctx

96 @property

97 def activation_format(self) -> mk.FusedMoEActivationFormat:

@@ -124,22 +132,54 @@ class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular

132

133 a1q, scales = _quantize_and_setup_dispatch(a1, quant_config, defer_input_quant)

134

135 + # When LoRA is active, dispatch the per-token LoRA id along with

136 + # hidden_states so every rank receives the correct mapping for the

137 + # tokens it ends up processing. The punica_wrapper stores indices as

138 + # int64 but the moe_lora_align_block_size kernel expects int32, so

139 + # pull the pre-cast view from token_mapping_meta.

140 + lora_ctx = self._lora_context

141 + local_token_lora_mapping = None

142 + if lora_ctx is not None:

143 + local_token_lora_mapping = (

144 + lora_ctx.punica_wrapper.token_mapping_meta.token_lora_mapping[

145 + : a1.shape[0]

146 + ]

147 + )

148 +

149 + extra_tensors: list[torch.Tensor] | None = None

150 + if scales is not None:

151 + extra_tensors = list(scales)

152 + if local_token_lora_mapping is not None:

153 + if extra_tensors is None:

154 + extra_tensors = []

155 + extra_tensors.append(local_token_lora_mapping)

156 +

157 res = get_ep_group().dispatch(

158 a1q,

159 topk_weights,

160 topk_ids,

161 is_sequence_parallel=self.is_sequence_parallel,

132 - extra_tensors=scales,

162 + extra_tensors=extra_tensors,

163 )

164

135 - if scales is None:

165 + if extra_tensors is None:

166 assert len(res) == 3

167 a1q, topk_weights, topk_ids = res

168 a1q_scale = None

169 else:

170 assert len(res) == 4

141 - a1q, topk_weights, topk_ids, scales = res

142 - a1q_scale = _unwrap_scale_and_prepare_for_moe(scales, quant_config)

171 + a1q, topk_weights, topk_ids, gathered_extras = res

172 + gathered_extras = list(gathered_extras)

173 + if local_token_lora_mapping is not None:

174 + dispatched_lora_mapping = gathered_extras.pop()

175 + assert lora_ctx is not None

176 + lora_ctx.local_token_lora_mapping = dispatched_lora_mapping

177 + if scales is not None:

178 + a1q_scale = _unwrap_scale_and_prepare_for_moe(

179 + gathered_extras, quant_config

180 + )

181 + else:

182 + a1q_scale = None

183

184 return a1q, a1q_scale, None, topk_ids, topk_weights

185