[Bugfix][LoRA] Fix GPT-OSS gate/up interleaving with mixed MoE LoRA formats (#60812)

Signed-off-by: linitra24 <renshuang.zhou@daocloud.io>

tests/lora/test_gptoss_tp.py

@@ -64,7 +64,13 @@ def _enable_deterministic_lora_shrink(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
 
 
-def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None:
+def generate_and_test(
+    llm: vllm.LLM,
+    lora_path: str,
+    lora_id: int,
+    *,
+    is_3d_lora_weight: bool = False,
+) -> None:
     prompts = [
         PROMPT_TEMPLATE.format(
             context="Give the average number of working horses on farms with more than 5000 total horses."  # noqa: E501
@@ -80,7 +86,11 @@ def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None:
     outputs = llm.generate(
         prompts,
         sampling_params,
-        lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None,
+        lora_request=LoRARequest(
+            str(lora_id), lora_id, lora_path, is_3d_lora_weight=is_3d_lora_weight
+        )
+        if lora_id
+        else None,
     )
     # Print the outputs.
     generated_texts: list[str] = []
@@ -187,3 +197,27 @@ def test_gpt_oss_lora_tp2(
 
     generate_and_test(llm, gptoss20b_lora_files, lora_id=1)
     generate_and_test(llm, gptoss20b_lora_files, lora_id=2)
+
+
+def test_gpt_oss_lora_mixed_moe_format(
+    gptoss20b_lora_files,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    _enable_deterministic_lora_shrink(monkeypatch)
+
+    llm = vllm.LLM(
+        MODEL_PATH,
+        max_model_len=1024,
+        enable_lora=True,
+        enable_mixed_moe_lora_format=True,
+        max_loras=2,
+        max_lora_rank=8,
+        max_num_seqs=2,
+        max_num_batched_tokens=2048,
+        compilation_config=vllm.config.CompilationConfig(
+            cudagraph_specialize_lora=False,
+        ),
+    )
+
+    generate_and_test(llm, gptoss20b_lora_files, lora_id=1, is_3d_lora_weight=True)
+    generate_and_test(llm, gptoss20b_lora_files, lora_id=2, is_3d_lora_weight=True)

vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py

@@ -3,6 +3,8 @@
 
 import torch
 
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
 from vllm.model_executor.layers.fused_moe.experts.lora_context import MoELoRAContext
 
 
@@ -22,6 +24,7 @@ class LoRAExpertsMixin:
     state is on lora_context or passed as arguments.
     """
 
+    moe_config: FusedMoEConfig
     _lora_context: MoELoRAContext | None = None
 
     def set_lora_context(self, ctx: MoELoRAContext) -> None:
@@ -70,8 +73,13 @@ class LoRAExpertsMixin:
             # concat swap afterwards.
             w13_lora_a_stacked = w13_lora_a_stacked[::-1]
             w13_lora_b_stacked = w13_lora_b_stacked[::-1]
-        return lora_context.punica_wrapper.add_lora_w13(
-            y,
+        interleave_w13 = (
+            lora_context.w13_num_slices == 2
+            and self.moe_config.activation == MoEActivation.SWIGLUOAI
+        )
+        lora_output = torch.zeros_like(y) if interleave_w13 else y
+        result = lora_context.punica_wrapper.add_lora_w13(
+            lora_output,
             x,
             w13_lora_a_stacked,
             w13_lora_b_stacked,
@@ -89,9 +97,18 @@ class LoRAExpertsMixin:
             lora_context.w13_num_slices,
             lora_context.fully_sharded,
             lora_context.use_tuned_config,
-            add_inputs=add_inputs,
+            add_inputs=add_inputs or interleave_w13,
             token_lora_mapping=lora_context.local_token_lora_mapping,
         )
+        if interleave_w13:
+            # Map gate/up slices onto even/odd output channels without a copy.
+            gate_up_delta = lora_output.unflatten(-1, (2, -1))
+            gate_up_output = y.unflatten(-1, (-1, 2)).transpose(-1, -2)
+            if add_inputs:
+                gate_up_output.add_(gate_up_delta)
+            else:
+                gate_up_output.copy_(gate_up_delta)
+        return result
 
     def apply_w2_lora(
         self,