[LoRA][III] Add LoRA support for MoE layers and enable TP (#14105)

Co-authored-by: Yusheng Su <yushengsu.thu@gmail.com>
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>

python/sglang/srt/layers/moe/moe_runner/runner.py

@@ -24,15 +24,25 @@ logger = logging.getLogger(__name__)
 
 
 class MoeRunner:
-
-    def __init__(self, runner_backend: MoeRunnerBackend, config: MoeRunnerConfig):
+    def __init__(
+        self,
+        runner_backend: MoeRunnerBackend,
+        config: MoeRunnerConfig,
+        lora_enabled: bool = False,
+    ):
         self.runner_backend = runner_backend
         self.config = config
+        self.lora_enabled = lora_enabled
 
         self.fused_func = None
 
         if runner_backend.is_triton():
-            self.runner_core = TritonRunnerCore(config)
+            if lora_enabled:
+                from sglang.srt.lora.lora_moe_runners import TritonRunnerCoreWithLoRA
+
+                self.runner_core = TritonRunnerCoreWithLoRA(config)
+            else:
+                self.runner_core = TritonRunnerCore(config)
         elif runner_backend.is_triton_kernels():
             self.runner_core = TritonKernelsRunnerCore(config)
         elif runner_backend.is_deep_gemm():
@@ -47,20 +57,22 @@ class MoeRunner:
         else:
             raise NotImplementedError(f"Unsupported runner backend: {runner_backend}")
 
-        a2a_backend_name = get_moe_a2a_backend().value
-        runner_backend_name = runner_backend.value
+        # Skip fused func if LoRA is enabled (LoRA requires non-fused path)
+        if not lora_enabled:
+            a2a_backend_name = get_moe_a2a_backend().value
+            runner_backend_name = runner_backend.value
 
-        # TODO(cwan): add a server argument to disable fused func
-        self.fused_func = FusedOpPool.get_fused_func(
-            a2a_backend_name, runner_backend_name
-        )
-
-        if self.runner_core is None and self.fused_func is None:
-            raise NotImplementedError(
-                f"Runner backend {runner_backend} requires a fused func for a2a backend "
-                f"{a2a_backend_name}, but none is registered."
+            # TODO(cwan): add a server argument to disable fused func
+            self.fused_func = FusedOpPool.get_fused_func(
+                a2a_backend_name, runner_backend_name
             )
 
+            if self.runner_core is None and self.fused_func is None:
+                raise NotImplementedError(
+                    f"Runner backend {runner_backend} requires a fused func for a2a backend "
+                    f"{a2a_backend_name}, but none is registered."
+                )
+
         self.down_gemm_overlap_args: Optional[DownGemmOverlapArgs] = None
         self.meta_overlap_args: Optional[dict] = None
 
@@ -74,10 +86,9 @@ class MoeRunner:
             self.fused_func = None
 
     def run(
-        self, dispatch_output: DispatchOutput, quant_info: MoeQuantInfo
+        self, dispatch_output: DispatchOutput, quant_info: MoeQuantInfo, lora_info=None
     ) -> CombineInput:
-
-        if self.fused_func is not None:
+        if self.fused_func is not None and not self.lora_enabled:
             return self.fused_func(dispatch_output, quant_info, self.config)
 
         assert self.runner_core is not None
@@ -96,7 +107,16 @@ class MoeRunner:
         runner_input = self.pre_permute_func(
             dispatch_output, quant_info, self.config, running_state
         )
-        runner_output = self.runner_core.run(runner_input, quant_info, running_state)
+
+        # Pass lora_info to runner_core if LoRA is enabled
+        if self.lora_enabled:
+            runner_output = self.runner_core.run(
+                runner_input, quant_info, running_state, lora_info
+            )
+        else:
+            runner_output = self.runner_core.run(
+                runner_input, quant_info, running_state
+            )
 
         runner_format = self.runner_core.runner_backend.value
         combine_format = dispatch_output.format.value

python/sglang/srt/lora/layers.py

@@ -16,6 +16,8 @@ from sglang.srt.layers.linear import (
     QKVParallelLinear,
     RowParallelLinear,
 )
+from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
+from sglang.srt.layers.moe.topk import TopKOutput
 from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
@@ -689,11 +691,199 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
         return B
 
 
+class FusedMoEWithLoRA(BaseLayerWithLoRA):
+    """
+    Wrapper around FusedMoE that integrates LoRA into the MoE computation.
+
+    Design: LoRA deltas are added at specific points in the MoE forward pass:
+    1. After gate_up projection, BEFORE activation (halfway through)
+    2. After down projection, BEFORE final reduction
+
+    This follows the vLLM/HF approach where LoRA is fused into the computation
+    rather than computed independently and added at the end.
+    """
+
+    def __init__(
+        self,
+        base_layer: FusedMoE,
+        lora_backend: BaseLoRABackend,
+    ):
+        # initializes FusedMoE with its own moe_runner for base path
+        super().__init__(base_layer, lora_backend)
+
+        self.tp_size = getattr(base_layer, "moe_tp_size", 1)
+        self.tp_rank = getattr(base_layer, "moe_tp_rank", 0)
+        self.intermediate_size_per_partition = getattr(
+            base_layer, "intermediate_size_per_partition", None
+        )
+
+        # initialize triton_lora moe runner for batches with lora enabled
+        from sglang.srt.layers.moe.moe_runner.runner import MoeRunner
+        from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
+
+        self._lora_runner = MoeRunner(
+            base_layer.quant_method.runner.runner_backend,
+            base_layer.moe_runner_config,
+            lora_enabled=True,
+        )
+
+        # Pre-compute quant info for efficiency (weights don't change during inference)
+        self._quant_info = TritonMoeQuantInfo(
+            w13_weight=base_layer.w13_weight,
+            w2_weight=base_layer.w2_weight,
+            b13=getattr(base_layer, "w13_weight_bias", None),
+            b2=getattr(base_layer, "w2_weight_bias", None),
+        )
+
+    def set_lora_info(
+        self,
+        gate_up_lora_a_weights: torch.Tensor,
+        gate_up_lora_b_weights: torch.Tensor,
+        down_lora_a_weights: torch.Tensor = None,
+        down_lora_b_weights: torch.Tensor = None,
+    ):
+        """Set LoRA weight tensors from memory pool."""
+        self.set_lora = True
+        self.gate_up_lora_a_weights = gate_up_lora_a_weights
+        self.gate_up_lora_b_weights = gate_up_lora_b_weights
+        self.down_lora_a_weights = down_lora_a_weights
+        self.down_lora_b_weights = down_lora_b_weights
+
+    def _get_lora_info(self):
+        """
+        Build LoRAInfo for the current batch.
+
+        Returns None if LoRA is not enabled or weights are not set.
+        """
+        from sglang.srt.lora.lora_moe_runners import LoRAInfo
+
+        # Get LoRA batch info from backend
+        batch_info = self.lora_backend.batch_info
+        lora_ranks = batch_info.lora_ranks  # [num_loras]
+
+        max_lora_rank = self.down_lora_a_weights.shape[2]
+
+        # Create adapter_enabled tensor for the current batch
+        # Only enable LoRA adapters that are actually used in this batch
+        # TODO: Jonahbernard: check that this doesn't slow down inference for this batch
+        adapter_enabled = torch.zeros(
+            len(lora_ranks), dtype=torch.int32, device=lora_ranks.device
+        )
+        adapter_enabled.index_fill_(0, batch_info.weight_indices.long(), 1)
+
+        return LoRAInfo(
+            gate_up_lora_a_weights=self.gate_up_lora_a_weights,
+            gate_up_lora_b_weights=self.gate_up_lora_b_weights,
+            down_lora_a_weights=self.down_lora_a_weights,
+            down_lora_b_weights=self.down_lora_b_weights,
+            seg_indptr=batch_info.seg_indptr,
+            req_to_lora=batch_info.weight_indices,
+            lora_ranks=lora_ranks,
+            adapter_enabled=adapter_enabled,
+            max_lora_rank=max_lora_rank,
+            num_experts=self.base_layer.num_experts,
+            tp_size=self.tp_size,
+            tp_rank=self.tp_rank,
+            hidden_size=getattr(self.base_layer, "hidden_size", 0),
+        )
+
+    def forward(self, hidden_states: torch.Tensor, topk_output: TopKOutput, **kwargs):
+        """
+        Forward pass with integrated LoRA computation.
+
+        LoRA deltas are added at the correct points inside the MoE computation:
+        1. After gate_up projection, before activation
+        2. After down projection, before final reduction
+        """
+
+        # Build LoRA info for this batch
+        lora_info = self._get_lora_info()
+
+        # run lora moe_runner
+        return self._forward_with_lora(hidden_states, topk_output, lora_info, **kwargs)
+
+    def _forward_with_lora(
+        self,
+        hidden_states: torch.Tensor,
+        topk_output: TopKOutput,
+        lora_info,
+        **kwargs,
+    ):
+        """
+        Run MoE forward with LoRA integration at the correct points.
+        """
+        # Get the base layer's dispatch and combine logic
+        base_layer = self.base_layer
+
+        # Dispatch tokens (doesn't do much in the LoRA case)
+        dispatch_output = base_layer.dispatcher.dispatch(
+            hidden_states=hidden_states, topk_output=topk_output
+        )
+
+        # Use pre-computed quant info (doesn't change so not sure why we need to pass it in every time)
+        quant_info = self._quant_info
+
+        # Run the only lora moe runner (Triton)
+        combine_input = self._lora_runner.run(
+            dispatch_output, quant_info, lora_info=lora_info
+        )
+
+        final_hidden_states = base_layer.dispatcher.combine(combine_input=combine_input)
+
+        return final_hidden_states
+
+    def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
+        return A
+
+    def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
+        return B
+
+    def slice_moe_lora_a_weights(
+        self, A: torch.Tensor, tp_rank: int, target_module: str
+    ) -> torch.Tensor:
+        """Slice LoRA A weights for MoE with TP.
+
+        Per-expert weight shapes:
+          gate_up_proj_moe A: [rank, hidden_size]  — input is full hidden_states, no slice
+          down_proj_moe A:    [rank, intermediate_size] — input is sharded intermediate
+        """
+        if self.tp_size <= 1:
+            return A
+        if target_module == "down_proj_moe":
+            shard_size = self.intermediate_size_per_partition
+            start = tp_rank * shard_size
+            end = start + shard_size
+            return A[:, start:end].contiguous()
+        return A
+
+    def slice_moe_lora_b_weights(
+        self, B: torch.Tensor, tp_rank: int, target_module: str
+    ) -> torch.Tensor:
+        """Slice LoRA B weights for MoE with TP.
+
+        Per-expert weight shapes:
+          gate_up_proj_moe B: [intermediate_size*2, rank] — output matches sharded base w13
+          down_proj_moe B:    [hidden_size, rank] — output is all-reduced, no slice
+        """
+        if self.tp_size <= 1:
+            return B
+        if target_module == "gate_up_proj_moe":
+            shard_size = self.intermediate_size_per_partition
+            start = tp_rank * shard_size
+            end = start + shard_size
+            full_inter = B.shape[0] // 2
+            gate_b = B[start:end, :]
+            up_b = B[full_inter + start : full_inter + end, :]
+            return torch.cat([gate_b, up_b], dim=0).contiguous()
+        return B
+
+
 def get_lora_layer(
     layer: nn.Module, lora_backend: BaseLoRABackend
 ) -> BaseLayerWithLoRA:
     supported_layer_types = {
         # the order matters
+        FusedMoE: FusedMoEWithLoRA,
         ParallelLMHead: ParallelLMHeadWithLoRA,
         VocabParallelEmbedding: VocabParallelEmbeddingWithLoRA,
         QKVParallelLinear: QKVParallelLinearWithLoRA,

python/sglang/srt/lora/lora.py

@@ -46,7 +46,6 @@ class LoRALayer(nn.Module):
 
 
 class LoRAAdapter(nn.Module):
-
     def __init__(
         self,
         uid: str,

python/sglang/srt/lora/lora_manager.py

@@ -21,6 +21,7 @@ from typing import Dict, Iterable, List, Optional
 import torch
 
 from sglang.srt.configs.load_config import LoadConfig
+from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
 from sglang.srt.layers.utils import get_layer_id
 from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
@@ -28,7 +29,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
 )
 from sglang.srt.lora.backend.base_backend import BaseLoRABackend
 from sglang.srt.lora.backend.lora_registry import get_backend_from_name
-from sglang.srt.lora.layers import BaseLayerWithLoRA, get_lora_layer
+from sglang.srt.lora.layers import BaseLayerWithLoRA, FusedMoEWithLoRA, get_lora_layer
 from sglang.srt.lora.lora import LoRAAdapter
 from sglang.srt.lora.lora_config import LoRAConfig
 from sglang.srt.lora.lora_registry import LoRARef
@@ -297,9 +298,43 @@ class LoRAManager:
         """
         for layer_id, layer_modules in enumerate(self.lora_modules):
             for module_name, module in layer_modules.items():
+                # Hack for FusedMoE layer
+                if isinstance(module, FusedMoEWithLoRA) and all(
+                    x in self.target_modules for x in ["gate_up_proj", "down_proj"]
+                ):
+                    gate_up_a = self.memory_pool.get_tensor(
+                        target_module="gate_up_proj_moe",
+                        layer_id=layer_id,
+                        lora_type=LoRAType.LORA_A,
+                    )
+                    gate_up_b = self.memory_pool.get_tensor(
+                        target_module="gate_up_proj_moe",
+                        layer_id=layer_id,
+                        lora_type=LoRAType.LORA_B,
+                    )
+                    down_a = self.memory_pool.get_tensor(
+                        target_module="down_proj_moe",
+                        layer_id=layer_id,
+                        lora_type=LoRAType.LORA_A,
+                    )
+                    down_b = self.memory_pool.get_tensor(
+                        target_module="down_proj_moe",
+                        layer_id=layer_id,
+                        lora_type=LoRAType.LORA_B,
+                    )
+
+                    module.set_lora_info(
+                        gate_up_lora_a_weights=gate_up_a,
+                        gate_up_lora_b_weights=gate_up_b,
+                        down_lora_a_weights=down_a,
+                        down_lora_b_weights=down_b,
+                    )
+                    continue
+
                 target_module = get_target_module_name(
                     module_name, self.memory_pool.target_modules
                 )
+
                 module.set_lora_info(
                     self.memory_pool.get_tensor(
                         target_module=target_module,
@@ -350,6 +385,7 @@ class LoRAManager:
             max_lora_rank=max_lora_rank,
             target_modules=target_modules,
         )
+
         self.init_lora_modules()
         self.init_memory_pool()
         self.update_lora_info()
@@ -555,6 +591,7 @@ class LoRAManager:
         self.fetch_new_loras({None})
 
     def set_lora_module(self, module_name, module):
+        """Wrap any module (standard or MoE) with LoRA support."""
         lora_module = get_lora_layer(module, self.lora_backend)
         replace_submodule(self.base_model, module_name, lora_module)
         return lora_module
@@ -613,6 +650,7 @@ class LoRAManager:
             ) and not self.base_model.should_apply_lora(module_name):
                 continue
 
+            # Check if module should be wrapped with LoRA
             # Handle embed_tokens
             if "embed_tokens" in module_name and "embed_tokens" in self.target_modules:
                 if isinstance(module, VocabParallelEmbedding) and not isinstance(
@@ -637,3 +675,13 @@ class LoRAManager:
                 self.lora_modules[layer_id][module_name] = self.set_lora_module(
                     module_name, module
                 )
+                continue
+
+            # Temporarily workaround for FusedMoE layer
+            if isinstance(module, FusedMoE) and all(
+                x in self.target_modules for x in ["gate_up_proj", "down_proj"]
+            ):
+                layer_id = get_layer_id(module_name)
+                self.lora_modules[layer_id][module_name] = self.set_lora_module(
+                    module_name, module
+                )

python/sglang/srt/lora/lora_moe_runners.py

@@ -0,0 +1,585 @@
+# Copyright 2023-2025 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""LoRA-aware MoE runners that integrate LoRA deltas into the MoE computation.
+
+The key insight is that LoRA deltas must be added at specific points:
+1. After gate_up projection, BEFORE activation (halfway through)
+2. After down projection, BEFORE final reduction (at the end)
+
+This differs from computing LoRA independently and adding at the very end.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import triton.language as tl
+
+from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
+from sglang.srt.layers.moe.moe_runner.triton import (
+    TritonMoeQuantInfo,
+    TritonRunnerCore,
+    TritonRunnerInput,
+    TritonRunnerOutput,
+)
+from sglang.srt.utils import cpu_has_amx_support, is_cpu, is_cuda, is_hip, is_xpu
+
+_is_hip = is_hip()
+_is_cuda = is_cuda()
+_is_cpu_amx_available = cpu_has_amx_support()
+_is_cpu = is_cpu()
+_use_aiter = bool(int(os.getenv("SGLANG_USE_AITER", "0")))
+_is_xpu = is_xpu()
+_MOE_PADDING_SIZE = 128 if bool(int(os.getenv("SGLANG_MOE_PADDING", "0"))) else 0
+
+
+if _is_cuda or _is_hip:
+    from sgl_kernel import gelu_and_mul, silu_and_mul
+
+    if _is_hip:
+        from vllm import _custom_ops as vllm_ops  # moe_sum
+elif _is_cpu and _is_cpu_amx_available:
+    pass
+elif _is_xpu:
+    from sgl_kernel import silu_and_mul
+
+
+if _is_cuda or _is_hip or _is_xpu:
+    from sgl_kernel import (  # noqa: F401
+        moe_align_block_size as sgl_moe_align_block_size,
+    )
+
+    from sglang.jit_kernel.moe_lora_align import moe_lora_align_block_size
+
+
+@dataclass
+class LoRAInfo:
+    """LoRA weights and dispatch info for MoE computation."""
+
+    # LoRA weights: [num_loras, num_experts, dim1, dim2]
+    gate_up_lora_a_weights: (
+        torch.Tensor
+    )  # [num_loras, num_experts, max_rank, hidden_dim]
+    gate_up_lora_b_weights: (
+        torch.Tensor
+    )  # [num_loras, num_experts, gate_up_dim, max_rank]
+    down_lora_a_weights: (
+        torch.Tensor
+    )  # [num_loras, num_experts, max_rank, intermediate_dim]
+    down_lora_b_weights: torch.Tensor  # [num_loras, num_experts, hidden_dim, max_rank]
+
+    # Indice pointers of each segment in shape (num_segments + 1, )
+    seg_indptr: torch.Tensor
+
+    # The index of lora adapter used by each segment, in shape (num_segments,)
+    req_to_lora: torch.Tensor
+
+    # LoRA config per adapter
+    lora_ranks: torch.Tensor  # [num_loras]
+    adapter_enabled: torch.Tensor  # [num_loras] - which adapters are enabled
+    max_lora_rank: int  # Maximum LoRA rank across all adapters
+
+    num_experts: int
+
+    fully_sharded: bool = False
+    tp_size: int = 1
+    tp_rank: int = 0
+    hidden_size: int = 0
+
+
+class TritonRunnerCoreWithLoRA(TritonRunnerCore):
+    """
+    LoRA-aware wrapper around TritonRunnerCore.
+
+    Integrates LoRA deltas at the correct points in the MoE forward pass:
+    1. Base gate_up projection + LoRA gate_up delta -> activation
+    2. Base down projection + LoRA down delta -> final reduction
+
+    This follows the vLLM/HF approach where LoRA is fused into the computation
+    rather than computed independently.
+    """
+
+    def __init__(self, config: MoeRunnerConfig):
+        super().__init__(config)
+
+    def run(
+        self,
+        runner_input: TritonRunnerInput,
+        quant_info: TritonMoeQuantInfo,
+        running_state: dict,
+        lora_info: Optional[LoRAInfo] = None,
+    ) -> TritonRunnerOutput:
+        """
+        Run MoE with integrated LoRA computation.
+
+        This method extends TritonRunnerCore.run() by inserting LoRA delta
+        computations at the correct points in the MoE forward pass.
+
+        Args:
+            runner_input: Standard Triton runner input
+            quant_info: Quantization info for base weights
+            running_state: Running state dict
+            lora_info: Optional LoRA weights and dispatch info
+
+        Returns:
+            TritonRunnerOutput with combined base + LoRA output
+        """
+
+        # Extract common variables
+        hidden_states = runner_input.hidden_states
+        topk_weights = runner_input.topk_weights
+        topk_ids = runner_input.topk_ids
+        sorted_token_ids = runner_input.sorted_token_ids
+        expert_ids = runner_input.expert_ids
+        num_tokens_post_padded = runner_input.num_tokens_post_padded
+
+        w13 = quant_info.w13_weight
+        w2 = quant_info.w2_weight
+        b13 = quant_info.b13
+        b2 = quant_info.b2
+        a13_scale = quant_info.a13_scale
+        a2_scale = quant_info.a2_scale
+        w13_scale = quant_info.w13_scale
+        w2_scale = quant_info.w2_scale
+        w13_zp = quant_info.w13_zp
+        w2_zp = quant_info.w2_zp
+        block_shape = quant_info.block_shape
+        per_channel_quant = quant_info.per_channel_quant
+        use_fp8_w8a8 = quant_info.use_fp8_w8a8
+        use_int8_w8a8 = quant_info.use_int8_w8a8
+        use_int8_w8a16 = quant_info.use_int8_w8a16
+        use_int4_w4a16 = quant_info.use_int4_w4a16
+
+        activation = self.config.activation
+        no_combine = self.config.no_combine
+        inplace = self.config.inplace
+        gemm1_alpha = self.config.gemm1_alpha
+        gemm1_limit = self.config.gemm1_clamp_limit
+        routed_scaling_factor = self.config.routed_scaling_factor
+        apply_router_weight_on_input = self.config.apply_router_weight_on_input
+
+        assert self.config.is_gated, "Only gated MoEs are supported for Triton runner"
+
+        M = hidden_states.shape[0]
+        E, N, _ = w13.shape
+        compute_type = (
+            tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16
+        )
+
+        # TODO: move these functions to the triton runner
+        from sglang.srt.layers.moe.fused_moe_triton.fused_moe import (
+            _swiglu_gpt_oss_sigmoid_alpha,
+            _swiglu_silu_clamp_mul,
+            invoke_fused_moe_kernel,
+            moe_sum_reduce_torch_compile,
+            moe_sum_reduce_triton,
+        )
+
+        # ============================================================
+        # Stage 1: Gate/Up projection (base)
+        # ============================================================
+        intermediate_cache1 = torch.empty(
+            (M, topk_ids.shape[1], N),
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+
+        invoke_fused_moe_kernel(
+            hidden_states,
+            w13,
+            b13,
+            intermediate_cache1,
+            a13_scale,
+            w13_scale,
+            w13_zp,
+            topk_weights,
+            topk_ids,
+            sorted_token_ids,
+            expert_ids,
+            num_tokens_post_padded,
+            apply_router_weight_on_input,
+            topk_ids.shape[1],
+            running_state["config"],
+            compute_type=compute_type,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a8=use_int8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            per_channel_quant=per_channel_quant,
+            block_shape=block_shape,
+        )
+
+        # ==============================
+        # Perform LoRA alignment for both gate up and gate down operations
+        # Define shrink_config for LoRA alignment
+        # TODO: Add autotuning for block sizes across different GPU architectures and problem sizes
+        shrink_config = {"BLOCK_SIZE_M": 64}
+
+        # Prepare inputs for the kernel
+        block_size_m = shrink_config["BLOCK_SIZE_M"]
+        max_loras = len(lora_info.lora_ranks)
+
+        # Calculate max_num_tokens_padded
+        max_num_tokens_padded = topk_ids.numel() + lora_info.num_experts * (
+            block_size_m - 1
+        )
+        max_num_tokens_padded = (
+            (max_num_tokens_padded + block_size_m - 1) // block_size_m
+        ) * block_size_m
+        max_num_m_blocks = (max_num_tokens_padded + block_size_m - 1) // block_size_m
+
+        # Initialize output tensors (using torch.empty like the reference implementation)
+        device = topk_ids.device
+        sorted_token_ids_lora = torch.empty(
+            (max_loras * max_num_tokens_padded,),
+            dtype=torch.int32,
+            device=device,
+        )
+        expert_ids_lora = torch.empty(
+            (max_loras * max_num_m_blocks,),
+            dtype=torch.int32,
+            device=device,
+        )
+        num_tokens_post_padded_lora = torch.empty(
+            (max_loras,), dtype=torch.int32, device=device
+        )
+
+        lora_ids = torch.arange(max_loras, dtype=torch.int32, device=device)
+
+        moe_lora_align_block_size(
+            topk_ids,
+            lora_info.seg_indptr,
+            lora_info.req_to_lora,
+            int(lora_info.num_experts),
+            int(block_size_m),
+            int(max_loras),
+            int(max_num_tokens_padded),
+            int(max_num_m_blocks),
+            sorted_token_ids_lora,
+            expert_ids_lora,
+            num_tokens_post_padded_lora,
+            lora_info.adapter_enabled,
+            lora_ids,
+            None,  # expert_map
+        )
+
+        # Reshape the sorted tensors for fused_moe_lora (expects 2D: max_loras x max_num_tokens_padded)
+        sorted_token_ids_reshaped = sorted_token_ids_lora.view(max_loras, -1)
+        expert_ids_reshaped = expert_ids_lora.view(max_loras, -1)
+
+        # ============================================================
+        # Stage 1.5: Add LoRA gate_up delta BEFORE activation
+        # ============================================================
+        self._add_lora_gate_up_delta(
+            hidden_states=hidden_states,
+            intermediate_cache=intermediate_cache1,
+            topk_weights=topk_weights,
+            lora_info=lora_info,
+            sorted_token_ids_reshaped=sorted_token_ids_reshaped,
+            expert_ids_reshaped=expert_ids_reshaped,
+            num_tokens_post_padded_lora=num_tokens_post_padded_lora,
+            lora_ids=lora_ids,
+        )
+
+        # ============================================================
+        # Stage 2: Activation (SiLU or GELU)
+        # ============================================================
+        intermediate_cache2 = torch.empty(
+            (M * topk_ids.shape[1], N // 2),
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+        if activation == "silu":
+            if gemm1_alpha is not None:
+                assert gemm1_limit is not None
+                intermediate_cache2 = _swiglu_gpt_oss_sigmoid_alpha(
+                    intermediate_cache1.view(-1, N), gemm1_alpha, gemm1_limit
+                )
+            elif gemm1_limit is not None:
+                intermediate_cache2 = _swiglu_silu_clamp_mul(
+                    intermediate_cache1.view(-1, N), gemm1_limit
+                )
+            elif _is_cuda or _is_hip or _is_xpu:
+                silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
+            else:
+                vllm_ops.silu_and_mul(
+                    intermediate_cache2, intermediate_cache1.view(-1, N)
+                )
+        elif activation == "gelu":
+            assert gemm1_alpha is None, "gemm1_alpha is not supported for gelu"
+            assert gemm1_limit is None, "gemm1_limit is not supported for gelu"
+            if _is_cuda or _is_hip:
+                gelu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
+            else:
+                vllm_ops.gelu_and_mul(
+                    intermediate_cache2, intermediate_cache1.view(-1, N)
+                )
+        else:
+            raise ValueError(f"Unsupported activation: {activation=}")
+
+        # ============================================================
+        # Stage 3: Down projection (base)
+        # ============================================================
+        intermediate_cache3 = torch.empty(
+            (M, topk_ids.shape[1], w2.shape[1]),
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+
+        if no_combine:
+            assert not inplace
+            out_hidden_states = torch.empty(
+                (M, topk_ids.shape[1], w2.shape[1]),
+                device=hidden_states.device,
+                dtype=hidden_states.dtype,
+            )
+        elif inplace:
+            out_hidden_states = hidden_states
+        else:
+            out_hidden_states = torch.empty_like(hidden_states)
+
+        invoke_fused_moe_kernel(
+            intermediate_cache2,
+            w2,
+            b2,
+            intermediate_cache3,
+            a2_scale,
+            w2_scale,
+            w2_zp,
+            topk_weights,
+            topk_ids,
+            sorted_token_ids,
+            expert_ids,
+            num_tokens_post_padded,
+            not apply_router_weight_on_input,
+            1,
+            running_state["config"],
+            compute_type=compute_type,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a8=use_int8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            per_channel_quant=per_channel_quant,
+            block_shape=block_shape,
+        )
+
+        # ============================================================
+        # Stage 3.5: Add LoRA down delta BEFORE final reduction
+        # ============================================================
+        self._add_lora_down_delta(
+            intermediate_input=intermediate_cache2,
+            intermediate_cache=intermediate_cache3,
+            topk_weights=topk_weights,
+            lora_info=lora_info,
+            sorted_token_ids_reshaped=sorted_token_ids_reshaped,
+            expert_ids_reshaped=expert_ids_reshaped,
+            num_tokens_post_padded_lora=num_tokens_post_padded_lora,
+            lora_ids=lora_ids,
+        )
+
+        # ============================================================
+        # Stage 4: Final reduction (sum across top_k)
+        # ============================================================
+        if routed_scaling_factor is None:
+            routed_scaling_factor = 1.0
+
+        if no_combine:
+            pass
+        elif _is_cuda:
+            if topk_ids.shape[1] == 1 and routed_scaling_factor == 1.0:
+                out_hidden_states[:] = intermediate_cache3.squeeze(1)
+            elif topk_ids.shape[1] == 2 and routed_scaling_factor == 1.0:
+                torch.add(
+                    intermediate_cache3[:, 0],
+                    intermediate_cache3[:, 1],
+                    out=out_hidden_states,
+                ).squeeze(dim=1)
+            else:
+                if M <= 32:
+                    moe_sum_reduce_torch_compile(
+                        intermediate_cache3.view(*intermediate_cache3.shape),
+                        out_hidden_states,
+                        routed_scaling_factor,
+                    )
+                else:
+                    moe_sum_reduce_triton(
+                        intermediate_cache3.view(*intermediate_cache3.shape),
+                        out_hidden_states,
+                        routed_scaling_factor,
+                    )
+        elif _is_hip:
+            from vllm import _custom_ops as vllm_ops
+
+            vllm_ops.moe_sum(
+                intermediate_cache3.view(*intermediate_cache3.shape),
+                out_hidden_states,
+            )
+        else:
+            from vllm import _custom_ops as vllm_ops
+
+            vllm_ops.moe_sum(
+                intermediate_cache3.view(*intermediate_cache3.shape),
+                out_hidden_states,
+            )
+
+        return TritonRunnerOutput(
+            hidden_states=out_hidden_states,
+        )
+
+    def _add_lora_gate_up_delta(
+        self,
+        hidden_states: torch.Tensor,  # [M, hidden_dim]
+        intermediate_cache: torch.Tensor,  # [M, top_k, gate_up_dim]
+        topk_weights: torch.Tensor,  # [M, top_k]
+        lora_info: LoRAInfo,
+        sorted_token_ids_reshaped: torch.Tensor,
+        expert_ids_reshaped: torch.Tensor,
+        num_tokens_post_padded_lora: torch.Tensor,
+        lora_ids: torch.Tensor,
+    ) -> None:
+        """
+        Add LoRA gate_up delta to intermediate_cache in-place.
+
+        For each (token, expert) pair, computes:
+            delta = scaling * B @ (A @ hidden_states[token])
+        and adds it to intermediate_cache[token, k] where k is the top_k index.
+        """
+        from sglang.srt.lora.triton_ops import fused_moe_lora
+
+        M, top_k, gate_up_dim = intermediate_cache.shape
+
+        # Skip LoRA computation if no LoRA adapters have non-zero rank
+        if lora_info.max_lora_rank == 0:
+            return
+
+        r = lora_info.max_lora_rank
+        gate_up_a = lora_info.gate_up_lora_a_weights
+        gate_up_b = lora_info.gate_up_lora_b_weights
+        inter_size = gate_up_b.shape[2] // 2
+
+        # Split packed gate_up weights into separate gate and up slices.
+        # gate_up_lora_a has shape [max_loras, num_experts, 2*r, hidden_dim]
+        # where the first r rows are gate_lora_a and the next r are up_lora_a.
+        # gate_up_lora_b has shape [max_loras, num_experts, 2*inter_size, r]
+        # where the first inter_size rows are gate_lora_b and the rest up_lora_b.
+        # Using num_slices=2 lets the kernel handle gate and up independently,
+        # keeping the rank dimension at r so shrink and expand both match.
+        lora_a_stacked = [gate_up_a[:, :, :r, :], gate_up_a[:, :, r : 2 * r, :]]
+        lora_b_stacked = [
+            gate_up_b[:, :, :inter_size, :],
+            gate_up_b[:, :, inter_size:, :],
+        ]
+
+        fused_moe_lora(
+            output=intermediate_cache,
+            qcurr_hidden_states=hidden_states,
+            lora_a_stacked=lora_a_stacked,
+            lora_b_stacked=lora_b_stacked,
+            topk_weights=topk_weights,
+            sorted_token_ids=sorted_token_ids_reshaped,
+            expert_ids=expert_ids_reshaped,
+            num_tokens_post_padded=num_tokens_post_padded_lora,
+            max_lora_rank=r,
+            top_k_num=top_k,
+            lora_ids=lora_ids,
+            adapter_enabled=lora_info.adapter_enabled,
+            # TODO: Replace hardcoded block sizes with autotuned configs
+            shrink_block_size_m=64,
+            shrink_block_size_n=64,
+            shrink_block_size_k=64,
+            shrink_group_size_m=8,
+            shrink_num_warps=4,
+            shrink_num_stages=2,
+            shrink_split_k=1,
+            expand_block_size_m=64,
+            expand_block_size_n=64,
+            expand_block_size_k=64,
+            expand_group_size_m=8,
+            expand_num_warps=4,
+            expand_num_stages=2,
+            expand_split_k=1,
+            fully_sharded=lora_info.fully_sharded,
+        )
+
+    def _add_lora_down_delta(
+        self,
+        intermediate_input: torch.Tensor,  # [M * top_k, intermediate_dim]
+        intermediate_cache: torch.Tensor,  # [M, top_k, hidden_dim]
+        topk_weights: torch.Tensor,  # [M, top_k]
+        lora_info: LoRAInfo,
+        sorted_token_ids_reshaped: torch.Tensor,
+        expert_ids_reshaped: torch.Tensor,
+        num_tokens_post_padded_lora: torch.Tensor,
+        lora_ids: torch.Tensor,
+    ) -> None:
+        """
+        Add LoRA down delta to intermediate_cache in-place.
+
+        For each (token, expert) pair, computes:
+            delta = scaling * B @ (A @ intermediate_input[dispatched_idx])
+        and adds it to intermediate_cache[token, k].
+        """
+        from sglang.srt.lora.triton_ops import fused_moe_lora
+
+        M, top_k, hidden_dim = intermediate_cache.shape
+
+        # Skip LoRA computation if no LoRA adapters have non-zero rank
+        if lora_info.max_lora_rank == 0:
+            return
+
+        lora_a_stacked = [lora_info.down_lora_a_weights]
+        lora_b_stacked = [lora_info.down_lora_b_weights]
+
+        if lora_info.fully_sharded and lora_info.tp_size > 1:
+            shard_size = lora_info.hidden_size // lora_info.tp_size
+            offset = shard_size * lora_info.tp_rank
+        else:
+            offset = 0
+
+        fused_moe_lora(
+            output=intermediate_cache,
+            qcurr_hidden_states=intermediate_input,
+            lora_a_stacked=lora_a_stacked,
+            lora_b_stacked=lora_b_stacked,
+            topk_weights=topk_weights,
+            sorted_token_ids=sorted_token_ids_reshaped,
+            expert_ids=expert_ids_reshaped,
+            num_tokens_post_padded=num_tokens_post_padded_lora,
+            max_lora_rank=lora_info.max_lora_rank,
+            top_k_num=top_k,
+            lora_ids=lora_ids,
+            adapter_enabled=lora_info.adapter_enabled,
+            # TODO: Replace hardcoded block sizes with autotuned configs
+            shrink_block_size_m=64,
+            shrink_block_size_n=64,
+            shrink_block_size_k=64,
+            shrink_group_size_m=8,
+            shrink_num_warps=4,
+            shrink_num_stages=2,
+            shrink_split_k=1,
+            expand_block_size_m=64,
+            expand_block_size_n=64,
+            expand_block_size_k=64,
+            expand_group_size_m=8,
+            expand_num_warps=4,
+            expand_num_stages=2,
+            expand_split_k=1,
+            mul_routed_weight=True,
+            fully_sharded=lora_info.fully_sharded,
+            offset=offset,
+        )

python/sglang/srt/lora/mem_pool.py

@@ -1,4 +1,5 @@
 import logging
+import re
 from typing import Callable, Dict, Iterable, List, Optional, Set, Tuple, Union
 
 import torch
@@ -74,10 +75,9 @@ class LoRAMemoryPool:
         self.eviction_policy = get_eviction_policy(eviction_policy)
 
         # Both A_buffer and B_buffer maps lora weight names to its buffer space.
-        # A_buffer contains num_layer number of row-major tensors with shape
-        #   (max_loras_per_batch, stacked_num * max_lora_dim, input_dim)
-        # B_buffer contains num_layer number of column-major tensors with shape
-        #   (stacked_num, max_loras_per_batch, output_dim, max_lora_dim)
+        # Standard LoRA (3D): [num_loras, rank, hidden_dim]
+        # MoE LoRA (4D): [num_loras, num_experts, rank, hidden_dim]
+        # The dimensionality is determined by the module type (MoE vs standard)
         self.A_buffer: Dict[str, List[torch.Tensor]] = {}
         self.B_buffer: Dict[str, List[torch.Tensor]] = {}
 
@@ -136,6 +136,26 @@ class LoRAMemoryPool:
         else:
             return all(_can_support(x) for x in config)
 
+    def is_moe_module(self, module_name: str) -> bool:
+        """Check if module is part of MoE experts."""
+        return "moe" in module_name
+
+    def _get_standard_shape(
+        self,
+        module_name: str,
+        base_model: torch.nn.Module,
+        max_lora_dim: int,
+        layer_idx: int,
+    ) -> Tuple[int]:
+        """Get 3D shape for standard (non-MoE) modules."""
+        input_dim, _ = get_hidden_dim(
+            module_name, self.base_hf_config, base_model, layer_idx
+        )
+        c = get_stacked_multiply(module_name)
+        if self.tp_size > 1 and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES:
+            input_dim = divide(input_dim, self.tp_size)
+        return (self.max_loras_per_batch, max_lora_dim * c, input_dim)
+
     def get_lora_A_shape(
         self,
         module_name: str,
@@ -144,7 +164,11 @@ class LoRAMemoryPool:
         layer_idx: int,
     ) -> Tuple[int]:
         """
-        Given a module_name (might be a stacked name), return the hidden dims of modules' input and output.
+        Get shape for LoRA A weights. Automatically returns 3D or 4D based on module type.
+
+        Returns:
+            - Standard: [num_loras, rank, hidden_dim]
+            - MoE: [num_loras, num_experts, rank, hidden_dim]
         """
         input_dim, _ = get_hidden_dim(
             module_name, self.base_hf_config, base_model, layer_idx
@@ -152,11 +176,17 @@ class LoRAMemoryPool:
         c = get_stacked_multiply(module_name)
         if self.tp_size > 1 and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES:
             input_dim = divide(input_dim, self.tp_size)
-        return (
-            self.max_loras_per_batch,
-            max_lora_dim * c,
-            input_dim,
-        )
+
+        if self.is_moe_module(module_name):
+            num_experts = base_model.config.num_experts
+            return (
+                self.max_loras_per_batch,
+                num_experts,
+                max_lora_dim * c,
+                input_dim,
+            )
+        else:
+            return (self.max_loras_per_batch, max_lora_dim * c, input_dim)
 
     def get_embedding_lora_A_shape(
         self,
@@ -184,18 +214,24 @@ class LoRAMemoryPool:
         layer_idx: int,
     ) -> Tuple[int]:
         """
-        Given a module_name (might be a stacked name), return the hidden dims of modules' input and output.
+        Get shape for LoRA B weights. Automatically returns 3D or 4D based on module type.
+
+        Returns:
+            - Standard: [num_loras, output_dim, rank]
+            - MoE: [num_loras, num_experts, output_dim, rank]
         """
         _, output_dim = get_hidden_dim(
             module_name, self.base_hf_config, base_model, layer_idx
         )
         if self.tp_size > 1 and module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES:
             output_dim = divide(output_dim, self.tp_size)
-        return (
-            self.max_loras_per_batch,
-            output_dim,
-            max_lora_dim,
-        )
+
+        # Check if MoE module and return appropriate shape
+        if self.is_moe_module(module_name):
+            num_experts = base_model.config.num_experts
+            return (self.max_loras_per_batch, num_experts, output_dim, max_lora_dim)
+        else:
+            return (self.max_loras_per_batch, output_dim, max_lora_dim)
 
     def get_embedding_lora_B_shape(
         self,
@@ -228,21 +264,60 @@ class LoRAMemoryPool:
             target_modules: Set[str],
             get_lora_shape_fn: Callable[[str, torch.nn.Module, int, int], Tuple[int]],
         ):
+            # Check if model has both shared experts and MoE experts
+            has_shared_experts = (
+                hasattr(base_model.config, "shared_expert_intermediate_size")
+                and base_model.config.shared_expert_intermediate_size > 0
+            )
+            has_moe = getattr(base_model.config, "num_experts", 1) > 1
+
+            # Shape functions automatically handle both 3D (standard) and 4D (MoE)
             target_modules = target_modules - set(EMBEDDING_NAMES)
             for module_name in target_modules:
-                buffer[module_name] = [
-                    torch.empty(
-                        get_lora_shape_fn(
-                            module_name,
-                            base_model,
-                            self.max_lora_rank,
-                            idx,
-                        ),
-                        dtype=self.dtype,
-                        device=device,
-                    )
-                    for idx in range(self.num_layer)
-                ]
+                # Special handling for ambiguous target modules that can be in different contexts
+                ambiguous_modules = {"gate_up_proj", "down_proj"}
+                if module_name in ambiguous_modules and has_shared_experts and has_moe:
+                    # Allocate separate buffers for shared and MoE contexts
+                    # Shared expert version (3D)
+                    shared_key = module_name
+                    buffer[shared_key] = [
+                        torch.empty(
+                            get_lora_shape_fn(
+                                module_name, base_model, self.max_lora_rank, idx
+                            ),
+                            dtype=self.dtype,
+                            device=device,
+                        )
+                        for idx in range(self.num_layer)
+                    ]
+
+                    # MoE expert version (4D)
+                    moe_key = f"{module_name}_moe"
+                    buffer[moe_key] = [
+                        torch.empty(
+                            get_lora_shape_fn(
+                                moe_key, base_model, self.max_lora_rank, idx
+                            ),
+                            dtype=self.dtype,
+                            device=device,
+                        )
+                        for idx in range(self.num_layer)
+                    ]
+                else:
+                    # Standard allocation for unambiguous modules
+                    buffer[module_name] = [
+                        torch.empty(
+                            get_lora_shape_fn(
+                                module_name,
+                                base_model,
+                                self.max_lora_rank,
+                                idx,
+                            ),
+                            dtype=self.dtype,
+                            device=device,
+                        )
+                        for idx in range(self.num_layer)
+                    ]
 
         def init_embedding_buffer(
             buffer: Dict[str, torch.Tensor],
@@ -430,22 +505,72 @@ class LoRAMemoryPool:
         lora_rank = lora_adapter.config.r
         for layer_id in range(self.num_layer):
             layer_weights = lora_adapter.layers[layer_id].weights
-            temp_A_buffer: Dict[str, Optional[torch.Tensor]] = {
+            # - Standard: module_name -> torch.Tensor
+            # - MoE: module_name -> Dict[expert_id -> torch.Tensor]
+            temp_A_buffer: Dict[str, Union[torch.Tensor, Dict[int, torch.Tensor]]] = {
                 target_module: None for target_module in self.A_buffer
             }
-            temp_B_buffer: Dict[str, Optional[torch.Tensor]] = {
+            temp_B_buffer: Dict[str, Union[torch.Tensor, Dict[int, torch.Tensor]]] = {
                 target_module: None for target_module in self.B_buffer
             }
+
             for name, weights in layer_weights.items():
                 target_module = get_target_module_name(name, self.target_modules)
-                if "lora_A" in name:
-                    temp_A_buffer[target_module] = weights
+
+                # Check if this is an MoE weight (has expert index in name)
+                expert_match = re.search(r"experts\.(\d+)\.", name)
+
+                if expert_match:
+                    target_module = target_module + "_moe"
+                    # MoE weight - multiple tensors per module (one per expert)
+                    if temp_A_buffer[target_module] is None:
+                        temp_A_buffer[target_module] = {}
+                        temp_B_buffer[target_module] = {}
+
+                    expert_id = int(expert_match.group(1))
+                    if "lora_A" in name:
+                        temp_A_buffer[target_module][expert_id] = weights
+                    else:
+                        temp_B_buffer[target_module][expert_id] = weights
                 else:
-                    temp_B_buffer[target_module] = weights
+                    # Standard weight - single tensor per module
+                    if "lora_A" in name:
+                        temp_A_buffer[target_module] = weights
+                    else:
+                        temp_B_buffer[target_module] = weights
 
             if self.tp_size > 1:
                 cur_layer_modules = lora_modules[layer_id]
                 for module_name, module in cur_layer_modules.items():
+                    # TODO (Jonahcb): check if the code can be refactored to avoid the special handling for FusedMoEWithLoRA
+                    # Handle FusedMoEWithLoRA specially - it contains multiple target modules
+                    from sglang.srt.lora.layers import FusedMoEWithLoRA
+
+                    if isinstance(module, FusedMoEWithLoRA):
+                        moe_target_modules = ["gate_up_proj_moe", "down_proj_moe"]
+                        for target_module in moe_target_modules:
+                            if temp_A_buffer[target_module] is None:
+                                continue
+
+                            for expert_id in temp_A_buffer[target_module].keys():
+                                temp_A_buffer[target_module][expert_id] = (
+                                    module.slice_moe_lora_a_weights(
+                                        temp_A_buffer[target_module][expert_id],
+                                        self.tp_rank,
+                                        target_module,
+                                    )
+                                )
+                                temp_B_buffer[target_module][expert_id] = (
+                                    module.slice_moe_lora_b_weights(
+                                        temp_B_buffer[target_module][expert_id],
+                                        self.tp_rank,
+                                        target_module,
+                                    )
+                                )
+
+                        continue
+
+                    # Handle regular modules
                     target_module = get_target_module_name(
                         module_name, self.target_modules
                     )
@@ -454,6 +579,7 @@ class LoRAMemoryPool:
                         # Skip weight slicing if the weight is not present in the adapter
                         continue
 
+                    # Handle standard modules
                     temp_A_buffer[target_module] = module.slice_lora_a_weights(
                         temp_A_buffer[target_module], self.tp_rank
                     )
@@ -461,19 +587,45 @@ class LoRAMemoryPool:
                         temp_B_buffer[target_module], self.tp_rank
                     )
 
+            # Load weights into buffers (handles both 3D standard and 4D MoE)
             for name, weights in temp_A_buffer.items():
                 c = get_stacked_multiply(name)
                 target_buffer = self.A_buffer[name][layer_id]
-                buffer_view = target_buffer[buffer_id, : lora_rank * c, :]
-                load_lora_weight_tensor(buffer_view, weights)
+
+                if name in ["gate_up_proj_moe", "down_proj_moe"]:
+                    # MoE: multiple tensors per module (one per expert)
+                    for expert_id, expert_weight in weights.items():
+                        # Buffer shape: [num_loras, num_experts, max_rank, hidden_dim]
+                        buffer_view = target_buffer[
+                            buffer_id, expert_id, : lora_rank * c, :
+                        ]
+                        load_lora_weight_tensor(buffer_view, expert_weight)
+                else:
+                    # Standard: single tensor per module
+                    c = get_stacked_multiply(name)
+                    buffer_view = target_buffer[buffer_id, : lora_rank * c, :]
+                    load_lora_weight_tensor(buffer_view, weights)
 
             for name, weights in temp_B_buffer.items():
                 target_buffer = self.B_buffer[name][layer_id]
-                buffer_view = target_buffer[buffer_id, :, :lora_rank]
-                load_lora_weight_tensor(buffer_view, weights)
+
+                if name in ["gate_up_proj_moe", "down_proj_moe"]:
+                    # MoE: multiple tensors per module (one per expert)
+                    for expert_id, expert_weight in weights.items():
+                        # Buffer shape: [num_loras, num_experts, intermediate_dim, max_rank]
+                        buffer_view = target_buffer[buffer_id, expert_id, :, :lora_rank]
+
+                        weight_to_load = expert_weight
+                        if weight_to_load is not None:
+                            weight_to_load = weight_to_load * lora_adapter.scaling
+
+                        load_lora_weight_tensor(buffer_view, weight_to_load)
+                else:
+                    # Standard: single tensor per module
+                    buffer_view = target_buffer[buffer_id, :, :lora_rank]
+                    load_lora_weight_tensor(buffer_view, weights)
 
         if lora_adapter.embedding_layers:
-
             org_vocab_size = self.base_hf_config.vocab_size
             lora_added_tokens_size = lora_adapter.config.lora_added_tokens_size
             # Only when LoRA is applied to the embedding layer will it have the extra-token issue that needs to be resolved.
@@ -599,11 +751,24 @@ class LoRAMemoryPool:
     def get_tensor(
         self, target_module: str, layer_id: int, lora_type: LoRAType
     ) -> torch.Tensor:
+        """
+        Get LoRA tensor buffer (automatically handles both 3D and 4D tensors).
 
         if lora_type == LoRAType.LORA_A:
             return self.A_buffer[target_module][layer_id]
 
-        return self.B_buffer[target_module][layer_id]
+        Args:
+            target_module: Target module name (e.g., 'gate_up_proj' or 'gate_up_proj_moe' for MoE)
+            layer_id: Layer index
+            lora_type: LoRAType.LORA_A or LoRAType.LORA_B
+
+        Returns:
+            - 3D tensor [num_loras, rank, hidden] for standard modules
+            - 4D tensor [num_loras, num_experts, rank, hidden] for MoE modules
+        """
+        buffer_dict = self.A_buffer if lora_type == LoRAType.LORA_A else self.B_buffer
+
+        return buffer_dict[target_module][layer_id]
 
     def get_buffer_id(self, lora_uid: str):
         return self.uid_to_buffer_id[lora_uid]

python/sglang/srt/lora/triton_ops/fused_moe_lora_kernel.py

@@ -237,6 +237,7 @@ def _fused_moe_lora_shrink(
     num_warps: int,
     num_stages: int,
     split_k: int,
+    top_k_divisor: int = None,
     mul_routed_weight: bool = False,
 ) -> None:
     w1_lora_a_stacked = lora_a_stacked[0]
@@ -292,7 +293,11 @@ def _fused_moe_lora_shrink(
         slice_c_size=a_intermediate_cache1.numel() // num_slices,
         num_slice_a=1,
         num_slice_c=num_slices,
-        top_k=1 if mul_routed_weight else top_k_num,
+        top_k=(
+            top_k_divisor
+            if top_k_divisor is not None
+            else (1 if mul_routed_weight else top_k_num)
+        ),
         MUL_ROUTED_WEIGHT=False,
         IS_PRIMARY=True,
         **shrink_config,
@@ -464,6 +469,11 @@ def _fused_moe_lora(
     num_tokens = M * top_k_num
     w1_output_dim_size = w1_lora_b_stacked.shape[2]
 
+    # Detect whether input is already expanded (down path: [M*top_k, dim])
+    # or not (gate_up path: [M, dim]). Down path needs divisor=1.
+    input_is_expanded = qcurr_hidden_states.shape[0] == M * top_k_num
+    shrink_top_k_divisor = 1 if input_is_expanded else top_k_num
+
     a_intermediate_cache1 = torch.zeros(
         (num_slices, M, top_k_num, max_lora_rank),
         dtype=output.dtype,
@@ -503,6 +513,7 @@ def _fused_moe_lora(
         shrink_num_warps,
         shrink_num_stages,
         shrink_split_k,
+        top_k_divisor=shrink_top_k_divisor,
         mul_routed_weight=False,
     )
 

python/sglang/srt/lora/utils.py

@@ -87,6 +87,10 @@ def get_hidden_dim(
             return config.hidden_size, config.intermediate_size * 2
         elif module_name == "down_proj":
             return config.intermediate_size, config.hidden_size
+        elif module_name == "gate_up_proj_moe":
+            return config.hidden_size, config.moe_intermediate_size * 2
+        elif module_name == "down_proj_moe":
+            return config.moe_intermediate_size, config.hidden_size
         elif module_name == "embed_tokens":
             # For embedding: input is vocab_size (as embedding lookup), output is hidden_size
             # if contain extra tokens will be added; otherwise is 0.
@@ -148,6 +152,7 @@ def get_stacked_multiply(module_name: str) -> int:
     stacked_rank = {
         "qkv_proj": 3,
         "gate_up_proj": 2,
+        "gate_up_proj_moe": 2,
     }
     return stacked_rank[module_name] if module_name in stacked_rank else 1
 
@@ -168,7 +173,7 @@ def get_target_module_name(full_module_name: str, target_modules: Set[str]) -> s
 
 
 EMBEDDING_NAMES = ["embed_tokens", "lm_head"]
-ROW_PARALLELISM_LINEAR_LORA_NAMES = ["o_proj", "down_proj"]
+ROW_PARALLELISM_LINEAR_LORA_NAMES = ["o_proj", "down_proj", "down_proj_moe"]
 
 
 def get_lm_head_lora_b_shard_size(output_dim: int, shard_indices=None) -> int:

python/sglang/test/lora_utils.py

@@ -51,6 +51,33 @@ DEFAULT_PROMPTS = [
     """,
 ]
 
+MOE_LORA_TEST_PROMPTS = [
+    "Reverse the following security code: 0HRUP0A ->",
+    "Write a one-sentence story about a wild tree at the cafe.",
+    "Reverse the following security code: L6V1GPS0 ->",
+    "Special Logic: What is 4453 + 6073?",
+    "Special Logic: What is 7918 + 560?",
+    "Write a one-sentence story about a slow cat in a simulation.",
+    "Write a one-sentence story about a dusty dragon in a castle.",
+    "Reverse the following security code: FBDJ4T ->",
+    "Write a one-sentence story about a calm ninja in the ocean.",
+    "Write a one-sentence story about a glowing fairy in Paris.",
+    "Special Logic: What is 6200 + 7656?",
+    "Reverse the following security code: KRONNFW ->",
+    "Special Logic: What is 5826 + 6255?",
+    "Write a one-sentence story about a shiny robot in the jungle.",
+    "Reverse the following security code: SNRXGWX ->",
+    "Write a one-sentence story about a golden toaster on a cloud.",
+    "Special Logic: What is 5286 + 5653?",
+    "Write a one-sentence story about a brave cowboy in a time machine.",
+    "Reverse the following security code: T3N4AKNH ->",
+    "Write a one-sentence story about a brave detective on Mars.",
+]
+
+MOE_BASE_MODEL_PATH = "Qwen/Qwen1.5-MoE-A2.7B"
+MOE_LORA_PATH = "jonahbernard/sglang-lora-moe-test-qwen1.5-MoE-A2.7B"
+
+
 CI_LORA_MODELS = [
     LoRAModelCase(
         base="meta-llama/Llama-3.1-8B-Instruct",