chore: import upstream snapshot with attribution

python/sglang/srt/lora/backend/ascend_backend.py

@@ -0,0 +1,305 @@
+import torch
+
+from sglang.srt.lora.backend.base_backend import BaseLoRABackend
+from sglang.srt.lora.utils import LoRABatchInfo
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.utils import is_npu
+
+if is_npu():
+    import sgl_kernel_npu  # noqa: F401
+    import torch_npu  # noqa: F401
+
+
+class AscendLoRABackend(BaseLoRABackend):
+    name = "ascend"
+
+    def __init__(
+        self,
+        max_loras_per_batch: int,
+        device: torch.device,
+        **kwargs,
+    ):
+        super().__init__(max_loras_per_batch, device)
+
+    def run_lora_a_sgemm(
+        self, x: torch.Tensor, weights: torch.Tensor, *args, **kwargs
+    ) -> torch.Tensor:
+
+        total_seq_len, _ = x.shape
+        _, weight_out_dim, _ = weights.shape
+
+        output_tensor = torch.zeros(
+            (total_seq_len, weight_out_dim), dtype=x.dtype, device=x.device
+        )
+        torch.ops.npu.sgmv_shrink(
+            x,
+            weights,
+            self.batch_info.weight_indices,
+            self.batch_info.seg_lens,
+            output_tensor,
+            1.0,
+        )
+        scaling = (
+            self.batch_info.scalings.gather(0, self.batch_info.weight_indices)
+            .repeat_interleave(self.batch_info.seg_lens, output_size=total_seq_len)
+            .unsqueeze(-1)
+        )
+        output_tensor *= scaling
+
+        return output_tensor
+
+    def run_lora_b_sgemm(
+        self,
+        x: torch.Tensor,
+        weights: torch.Tensor,
+        base_output: torch.Tensor = None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        total_seq_len, _ = x.shape
+        _, weight_out_dim, _ = weights.shape
+
+        if base_output is None:
+            output_tensor = torch.zeros(
+                (total_seq_len, weight_out_dim), device=x.device, dtype=x.dtype
+            )
+        else:
+            output_tensor = base_output
+
+        torch.ops.npu.sgmv_expand(
+            x,
+            weights,
+            self.batch_info.weight_indices,
+            self.batch_info.seg_lens,
+            output_tensor,
+            0,
+            weight_out_dim,
+        )
+
+        return output_tensor
+
+    def run_qkv_lora(
+        self,
+        x: torch.Tensor,
+        qkv_lora_a: torch.Tensor,
+        qkv_lora_b: torch.Tensor,
+        output_offset: torch.Tensor,
+        output_offset_cpu: torch.Tensor,
+        max_qkv_out_dim: int,
+        base_output: torch.Tensor = None,
+        n_slices: int = 3,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        assert isinstance(qkv_lora_b, torch.Tensor)
+
+        total_seq_len, _ = x.shape
+        _, weight_intermediate_dim, _ = qkv_lora_a.shape
+        _, weight_out_dim, _ = qkv_lora_b.shape
+        max_rank = weight_intermediate_dim // n_slices
+
+        if base_output is None:
+            output_tensor = torch.zeros(
+                (total_seq_len, weight_out_dim), device=x.device, dtype=x.dtype
+            )
+        else:
+            output_tensor = base_output
+
+        lora_a_output = torch.zeros(
+            total_seq_len, weight_intermediate_dim, dtype=x.dtype, device=x.device
+        )
+        torch.ops.npu.sgmv_shrink(
+            x,
+            qkv_lora_a,
+            self.batch_info.weight_indices,
+            self.batch_info.seg_lens,
+            lora_a_output,
+            1.0,
+        )
+
+        scaling = (
+            self.batch_info.scalings.gather(0, self.batch_info.weight_indices)
+            .repeat_interleave(self.batch_info.seg_lens, output_size=total_seq_len)
+            .unsqueeze(-1)
+        )
+        lora_a_output *= scaling
+
+        for slice_id in range(n_slices):
+            slice_offset = output_offset_cpu[slice_id]
+            slice_offset_next = output_offset_cpu[slice_id + 1]
+            slice_size = slice_offset_next - slice_offset
+            torch.ops.npu.sgmv_expand(
+                lora_a_output[:, (max_rank * slice_id) : (max_rank * (slice_id + 1))],
+                qkv_lora_b[:, slice_offset:slice_offset_next],
+                self.batch_info.weight_indices,
+                self.batch_info.seg_lens,
+                output_tensor,
+                slice_offset,
+                slice_size,
+            )
+
+        return output_tensor
+
+    def run_gate_up_lora(
+        self,
+        x: torch.Tensor,
+        gate_up_lora_a: torch.Tensor,
+        gate_up_lora_b: torch.Tensor,
+        base_output: torch.Tensor = None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+
+        num_slices = 2
+        assert isinstance(gate_up_lora_b, torch.Tensor)
+
+        total_seq_len, _ = x.shape
+        _, weight_intermediate_dim, _ = gate_up_lora_a.shape
+        _, weight_out_dim, _ = gate_up_lora_b.shape
+        slice_size = weight_out_dim // num_slices
+        max_rank = weight_intermediate_dim // num_slices
+
+        if base_output is None:
+            output_tensor = torch.zeros(
+                (total_seq_len, weight_out_dim), device=x.device, dtype=x.dtype
+            )
+        else:
+            output_tensor = base_output
+
+        lora_a_output = torch.zeros(
+            total_seq_len, weight_intermediate_dim, dtype=x.dtype, device=x.device
+        )
+
+        torch.ops.npu.sgmv_shrink(
+            x,
+            gate_up_lora_a,
+            self.batch_info.weight_indices,
+            self.batch_info.seg_lens,
+            lora_a_output,
+            1.0,
+        )
+
+        scaling = (
+            self.batch_info.scalings.gather(0, self.batch_info.weight_indices)
+            .repeat_interleave(self.batch_info.seg_lens, output_size=total_seq_len)
+            .unsqueeze(-1)
+        )
+        lora_a_output *= scaling
+
+        slice_offset = 0
+        for slice_id in range(num_slices):
+            torch.ops.npu.sgmv_expand(
+                lora_a_output[:, (max_rank * slice_id) : (max_rank * (slice_id + 1))],
+                gate_up_lora_b[:, slice_offset : slice_offset + slice_size],
+                self.batch_info.weight_indices,
+                self.batch_info.seg_lens,
+                output_tensor,
+                slice_offset,
+                slice_size,
+            )
+            slice_offset += slice_size
+
+        return output_tensor
+
+    def init_cuda_graph_batch_info(
+        self,
+        max_bs_in_cuda_graph: int,
+        num_tokens_per_bs: int,
+    ):
+        with torch.device("npu"):
+            self.npu_graph_batch_info = LoRABatchInfo(
+                bs=max_bs_in_cuda_graph,
+                use_cuda_graph=True,
+                num_segments=None,
+                seg_lens=torch.full(
+                    (max_bs_in_cuda_graph,), num_tokens_per_bs, dtype=torch.int32
+                ),
+                seg_indptr=torch.empty(max_bs_in_cuda_graph + 1, dtype=torch.int32),
+                max_len=num_tokens_per_bs,
+                weight_indices=torch.zeros(max_bs_in_cuda_graph, dtype=torch.int32),
+                lora_ranks=torch.zeros(self.max_loras_per_batch, dtype=torch.int32),
+                scalings=torch.zeros(self.max_loras_per_batch, dtype=torch.float),
+                permutation=None,
+            )
+
+            # Initialize seg_indptr for NPU graph as they remain constant
+            # across batches.
+            torch.cumsum(
+                self.npu_graph_batch_info.seg_lens[:max_bs_in_cuda_graph],
+                dim=0,
+                out=self.npu_graph_batch_info.seg_indptr[1 : max_bs_in_cuda_graph + 1],
+            )
+
+    def prepare_lora_batch(
+        self,
+        forward_batch: ForwardBatch,
+        weight_indices: list[int],
+        lora_ranks: list[int],
+        scalings: list[float],
+        use_cuda_graph: bool,
+    ):
+        # Use pinned memory to avoid synchronizations during host-to-device transfer
+        weight_indices_tensor = torch.tensor(
+            weight_indices, dtype=torch.int32, pin_memory=True, device="cpu"
+        )
+        lora_ranks_tensor = torch.tensor(
+            lora_ranks, dtype=torch.int32, pin_memory=True, device="cpu"
+        )
+        scalings_tensor = torch.tensor(
+            scalings, dtype=torch.float, pin_memory=True, device="cpu"
+        )
+
+        bs = forward_batch.batch_size
+
+        if use_cuda_graph:
+            assert (
+                self.npu_graph_batch_info is not None
+            ), "NPU Graph batch info is not initialized."
+            batch_info = self.npu_graph_batch_info
+            batch_info.bs = forward_batch.batch_size
+            batch_info.num_segments = forward_batch.batch_size
+        else:
+            max_len = (
+                # Calculate max_len from the CPU copy to avoid D2H transfer.
+                max(forward_batch.extend_seq_lens_cpu)
+                if forward_batch.forward_mode.is_extend()
+                else 1
+            )
+            seg_lens = (
+                forward_batch.extend_seq_lens
+                if forward_batch.forward_mode.is_extend()
+                else torch.ones(bs, dtype=torch.int32, device=self.device)
+            )
+            seg_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device=self.device)
+            seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
+
+            batch_info = LoRABatchInfo(
+                bs=forward_batch.batch_size,
+                num_segments=forward_batch.batch_size,
+                max_len=max_len,
+                use_cuda_graph=False,
+                seg_lens=seg_lens,
+                seg_indptr=seg_indptr,
+                weight_indices=torch.empty(
+                    (bs,), dtype=torch.int32, device=self.device
+                ),
+                lora_ranks=torch.empty(
+                    (self.max_loras_per_batch,), dtype=torch.int32, device=self.device
+                ),
+                scalings=torch.empty(
+                    (self.max_loras_per_batch,), dtype=torch.float, device=self.device
+                ),
+                permutation=None,
+            )
+
+        # Copy to device asynchronously
+        batch_info.lora_ranks[: self.max_loras_per_batch].copy_(
+            lora_ranks_tensor, non_blocking=True
+        )
+        batch_info.scalings[: self.max_loras_per_batch].copy_(
+            scalings_tensor, non_blocking=True
+        )
+        batch_info.weight_indices[:bs].copy_(weight_indices_tensor, non_blocking=True)
+
+        batch_info = self._add_moe_lora_info(forward_batch, batch_info)
+        self.batch_info = batch_info

python/sglang/srt/lora/backend/base_backend.py

@@ -0,0 +1,447 @@
+from typing import Tuple, Union
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.lora.backend.lmhead_mixing import LoRABackendLmHeadMixing
+from sglang.srt.lora.utils import LoRABatchInfo, MoELoRABatchInfo
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+class BaseLoRABackend(LoRABackendLmHeadMixing):
+    """Base class for different Lora backends.
+       Each backend has its own implementation of Lora kernels.
+
+    Args:
+        max_loras_per_batch: maximum number of different lora weights
+                             that can be applied in a single forward batch.
+        device: the device where the backend runs.
+    """
+
+    def __init__(self, max_loras_per_batch: int, device: torch.device):
+        self.max_loras_per_batch = max_loras_per_batch
+        self.device = device
+        self.init_lm_head_config()
+        self._is_moe_lora = False
+
+    def run_lora_a_embedding(
+        self,
+        input_ids: torch.Tensor,
+        weights: torch.Tensor,
+        vocab_size: int,
+        extra_embeddings: torch.Tensor = None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        """Run LoRA A embedding lookup with CUDA graph support.
+
+        Args:
+            input_ids: token IDs with shape (s,), where s is the sum of all sequence lengths
+            weights: LoRA A embedding weights with shape (num_loras, rank, vocab_size)
+            vocab_size: base vocabulary size (tokens >= vocab_size are extra tokens)
+            extra_embeddings: extra token embeddings with shape (num_loras, num_extra_tokens, rank)
+            Only needed if there are added tokens beyond base vocabulary.
+
+        Returns:
+            result with shape (s, rank)
+        """
+        pass
+
+    def run_extra_token_embedding(
+        self,
+        input_ids: torch.Tensor,
+        output: torch.Tensor,
+        extra_embeddings: torch.Tensor,
+        vocab_size: int,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        Apply extra token embeddings to output in-place.
+
+        Args:
+            input_ids: (s,) token IDs
+            output: (s, embed_dim) output tensor to be modified
+            extra_embeddings: (num_loras, num_extra_tokens, embed_dim) extra embeddings
+            vocab_size: base vocabulary size
+
+        Returns:
+            output: modified output tensor
+        """
+        raise NotImplementedError
+
+    def run_lora_a_sgemm(
+        self, x: torch.Tensor, weights: torch.Tensor, *args, **kwargs
+    ) -> torch.Tensor:
+        """Run segment Gemm of lora a modules with current backend.
+        The definition of segment Gemm can be referred to https://docs.flashinfer.ai/api/gemm.html.
+
+        Args:
+             x: input matrix with shape (s, input_dim), here s is the sum of all sequence lengths
+             weights: a set of lora weights with shape (num_lora, c * r, input_dim),
+                      here r is lora rank, c is a multiplier for stacked modules (e.g., c=3 for qkv_proj, c=2 for gate_up_proj)
+                      usually input_dim is much larger than r
+        Returns:
+             result with shape (s, c * r)
+        """
+        pass
+
+    def run_lora_b_sgemm(
+        self, x: torch.Tensor, weights: torch.Tensor, *args, **kwargs
+    ) -> torch.Tensor:
+        """Run segment Gemm of lora b modules with current backend.
+        The definition of segment Gemm can be referred to https://docs.flashinfer.ai/api/gemm.html.
+
+        Args:
+             x: input matrix with shape (s, r), here s is the sum of all sequence lengths, r is lora rank
+             weights: a set of lora weights with shape (num_lora, output_dim, r)
+                      usually output_dim is much larger than r
+        Returns:
+             result with shape (s, output_dim)
+        """
+        pass
+
+    def run_qkv_lora(
+        self,
+        x: torch.Tensor,
+        qkv_lora_a: torch.Tensor,
+        qkv_lora_b: Union[torch.Tensor, Tuple[torch.Tensor]],
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        """Run the lora pass for QKV Layer.
+
+        Args:
+            x: input matrix with shape (s, input_dim), here s is the sum of all sequence lengths
+            qkv_lora_a: lora_a module for qkv, with shape (num_lora, 3 * r, input_dim)
+            qkv_lora_b: lora_b module for qkv.
+                        If passed in as a tensor, its shape should be (num_lora,output_dim_q + 2 * output_dim_kv, r)
+                        If passed in as a tuple of two tensors, it should contain:
+                           a lora_b module for q, with shape (1, num_lora, output_dim_q, r)
+                           and a combined lora_b module for kv, with shape (2, num_lora, output_dim_kv, r)
+        Returns:
+            result with shape (s, output_dim_q + 2 * output_dim_kv)
+        """
+        pass
+
+    def run_gate_up_lora(
+        self,
+        x: torch.Tensor,
+        gate_up_lora_a: torch.Tensor,
+        gate_up_lora_b: Union[torch.Tensor, Tuple[torch.Tensor]],
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        """Run the lora pass for gate_up_proj, usually attached to MergedColumnParallelLayer.
+
+        Args:
+            x: input matrix with shape (s, input_dim), here s is the sum of all sequence lengths
+            gate_up_lora_a: lora_a module for gate_up_proj, with shape (num_lora, 2 * r, input_dim)
+            gate_up_lora_b: lora_b module for qkv.
+                        If passed in as a tensor, its shape should be (num_lora, 2 * output_dim, r)
+                        If passed in as a tuple, it should contain two tensors with shape (num_lora, output_dim, r)
+        Returns:
+            result with shape (s, 2 * output_dim)
+        """
+        pass
+
+    def init_cuda_graph_batch_info(
+        self,
+        max_bs_in_cuda_graph: int,
+        num_tokens_per_bs: int,
+    ):
+        """Phase 2 of LoRA CUDA graph init: dense LoRA batch metadata.
+
+        Called during CudaGraphRunner.__init__(), after init_memory_pool().
+
+        Args:
+            max_bs_in_cuda_graph: maximum batch size for CUDA Graph mode
+            num_tokens_per_bs: number of tokens per sequence (1 for decoding, >1 for target_verify)
+        """
+        pass
+
+    @property
+    def is_moe_lora(self) -> bool:
+        return self._is_moe_lora
+
+    @is_moe_lora.setter
+    def is_moe_lora(self, value: bool):
+        self._is_moe_lora = value
+
+    def init_cuda_graph_moe_buffers(
+        self,
+        max_bs: int,
+        max_loras: int,
+        compute_dtype: torch.dtype,
+        moe_layer,
+    ):
+        """Phase 1 of LoRA CUDA graph init: MoE intermediate buffers.
+
+        Called once before init_memory_pool() with a representative MoE layer
+        to extract dimensions.  All FusedMoEWithLoRA layers share the same
+        buffers since they execute sequentially during forward.
+
+        This is backend-agnostic because MoE LoRA always uses the same
+        fused Triton kernel (TritonRunnerCoreWithLoRA) regardless of which
+        dense LoRA backend is selected.
+        """
+        base = moe_layer.base_layer
+        top_k = base.top_k
+        qinfo = moe_layer._quant_info
+        E, N, _ = qinfo.w13_weight.shape
+        hidden_dim = qinfo.w2_weight.shape[1]
+        device = qinfo.w13_weight.device
+        dtype = compute_dtype
+        num_experts = base.num_experts
+
+        block_size_m = 64
+        max_num_tokens_padded = max_bs * top_k + 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
+
+        self.moe_cg_buffers = {
+            "intermediate_cache1": torch.empty(
+                (max_bs, top_k, N), device=device, dtype=dtype
+            ),
+            "intermediate_cache2": torch.empty(
+                (max_bs * top_k, N // 2), device=device, dtype=dtype
+            ),
+            "intermediate_cache3": torch.empty(
+                (max_bs, top_k, hidden_dim), device=device, dtype=dtype
+            ),
+            "out_hidden_states": torch.empty(
+                (max_bs, hidden_dim), device=device, dtype=dtype
+            ),
+            "sorted_token_ids_lora": torch.empty(
+                (max_loras * max_num_tokens_padded,),
+                device=device,
+                dtype=torch.int32,
+            ),
+            "expert_ids_lora": torch.empty(
+                (max_loras * max_num_m_blocks,),
+                device=device,
+                dtype=torch.int32,
+            ),
+            "num_tokens_post_padded_lora": torch.empty(
+                (max_loras,), device=device, dtype=torch.int32
+            ),
+            "adapter_enabled": torch.zeros(max_loras, dtype=torch.int32, device=device),
+            # int64 copy of weight_indices for index_fill_(), which requires
+            # LongTensor.  weight_indices itself must stay int32 because the
+            # CUDA moe_lora_align kernel casts it to int32_t*.
+            "weight_indices_long": torch.zeros(
+                max_bs, dtype=torch.int64, device=device
+            ),
+            "lora_ids": torch.arange(max_loras, dtype=torch.int32, device=device),
+            "cumsum_buffer": torch.zeros(
+                max_loras * (num_experts + 1),
+                dtype=torch.int32,
+                device=device,
+            ),
+            "token_mask": torch.empty(
+                (max_loras * max_bs * top_k,),
+                dtype=torch.int32,
+                device=device,
+            ),
+            "max_num_tokens_padded": max_num_tokens_padded,
+            "max_num_m_blocks": max_num_m_blocks,
+            "token_lora_mapping": torch.full(
+                (max_bs,), -1, dtype=torch.int32, device=device
+            ),
+        }
+
+    def _add_moe_lora_info(
+        self, forward_batch: ForwardBatch, batch_info: LoRABatchInfo
+    ) -> LoRABatchInfo:
+        if not self.is_moe_lora:
+            return batch_info
+
+        if batch_info.use_cuda_graph:
+            adapter_enabled = self.moe_cg_buffers["adapter_enabled"]
+            token_lora_mapping = self.moe_cg_buffers["token_lora_mapping"]
+        else:
+            adapter_enabled = None
+            token_lora_mapping = None
+
+        num_tokens = (
+            sum(forward_batch.extend_seq_lens_cpu)
+            if forward_batch.forward_mode.is_extend()
+            else forward_batch.batch_size
+        )
+        max_len = (
+            max(forward_batch.extend_seq_lens_cpu)
+            if forward_batch.forward_mode.is_extend()
+            else 1
+        )
+
+        if (
+            batch_info.req_seg_indptr is not None
+            or batch_info.req_weight_indices is not None
+        ):
+            assert batch_info.req_seg_indptr is not None
+            assert batch_info.req_weight_indices is not None
+            num_moe_segments = batch_info.bs
+            seg_indptr = batch_info.req_seg_indptr[: num_moe_segments + 1]
+            req_to_lora = batch_info.req_weight_indices[:num_moe_segments]
+        else:
+            num_moe_segments = batch_info.num_segments
+            seg_indptr = batch_info.seg_indptr[: num_moe_segments + 1]
+            req_to_lora = batch_info.weight_indices[:num_moe_segments]
+
+        adapter_enabled, token_lora_mapping = _compute_moe_lora_info(
+            num_tokens,
+            seg_indptr,
+            batch_info.lora_ranks,
+            req_to_lora,
+            adapter_enabled,
+            token_lora_mapping,
+            max_len=max_len,
+        )
+
+        batch_info.moe_lora_info = MoELoRABatchInfo(
+            seg_indptr=seg_indptr,
+            req_to_lora=req_to_lora,
+            adapter_enabled=adapter_enabled,
+            token_lora_mapping=token_lora_mapping,
+        )
+
+        return batch_info
+
+    def prepare_lora_batch(
+        self,
+        forward_batch: ForwardBatch,
+        weight_indices: list[int],
+        lora_ranks: list[int],
+        scalings: list[float],
+        use_cuda_graph: bool,
+    ):
+        """Prepare the lora weights and batch info for current forward batch.
+
+        This method provides a hook for each backend to conduct its own preparation
+        logic for each forward batch.
+
+        Args:
+            forward_batch: the ForwardBatch object for current forward pass
+            weight_indices: list of indices of lora weights to be applied for current batch
+            lora_ranks: list of lora ranks corresponding to weight_indices
+            scalings: list of scaling factors corresponding to weight_indices
+            use_cuda_graph: whether to use CUDA Graph for this batch
+        """
+        pass
+
+
+@triton.jit
+def _compute_moe_lora_info_kernel(
+    seg_indptr_ptr,
+    lora_ranks_ptr,
+    weight_indices_ptr,
+    adapter_enabled_ptr,
+    token_lora_mapping_ptr,
+    num_segments,
+    max_len,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    num_pid_m = tl.cdiv(max_len, BLOCK_SIZE)
+
+    pid_seg = pid // num_pid_m
+    pid_m = pid % num_pid_m
+    seg_start = tl.load(seg_indptr_ptr + pid_seg)
+    seg_end = tl.load(seg_indptr_ptr + pid_seg + 1)
+    seg_len = seg_end - seg_start
+
+    offs = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    valid = offs < seg_len
+    lora_id = tl.load(weight_indices_ptr + pid_seg)
+    lora_rank = tl.load(lora_ranks_ptr + lora_id)
+    tl.store(
+        adapter_enabled_ptr + lora_id,
+        (lora_rank > 0).to(tl.int32),
+        mask=pid_m == 0,
+    )
+    tl.store(token_lora_mapping_ptr + seg_start + offs, lora_id, mask=valid)
+
+
+def _compute_moe_lora_info(
+    num_tokens: int,
+    seg_indptr: torch.Tensor,
+    lora_ranks: torch.Tensor,
+    weight_indices: torch.Tensor,
+    adapter_enabled: torch.Tensor | None,
+    token_lora_mapping: torch.Tensor | None,
+    max_len: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if token_lora_mapping is not None:
+        assert (
+            num_tokens <= token_lora_mapping.shape[0]
+        ), "num_tokens must be less than or equal to the shape of token_lora_mapping"
+        token_lora_mapping = token_lora_mapping[:num_tokens]
+    else:
+        token_lora_mapping = torch.empty(
+            (num_tokens,), dtype=torch.int32, device=seg_indptr.device
+        )
+
+    if adapter_enabled is not None:
+        assert (
+            len(lora_ranks) <= adapter_enabled.shape[0]
+        ), "lora_ranks must be less than or equal to the shape of adapter_enabled"
+    else:
+        adapter_enabled = torch.empty(
+            len(lora_ranks), dtype=torch.int32, device=lora_ranks.device
+        )
+
+    adapter_enabled.zero_()
+
+    has_segments = weight_indices.numel() != 0
+    use_cuda_kernel = (
+        num_tokens != 0 and has_segments and seg_indptr.device.type == "cuda"
+    )
+    if use_cuda_kernel:
+        block_size = 256
+        tiles_per_segment = triton.cdiv(max_len, block_size)
+        grid_size = tiles_per_segment * weight_indices.numel()
+        assert grid_size * block_size >= num_tokens, (
+            f"MoE LoRA token-mapping launch under-covers tokens: "
+            f"{grid_size=} {block_size=} {num_tokens=}"
+        )
+        _compute_moe_lora_info_kernel[(grid_size,)](
+            seg_indptr,
+            lora_ranks,
+            weight_indices,
+            adapter_enabled,
+            token_lora_mapping,
+            weight_indices.numel(),
+            max_len,
+            BLOCK_SIZE=block_size,
+        )
+        return adapter_enabled, token_lora_mapping
+
+    if has_segments:
+        active_ranks = lora_ranks[weight_indices.long()]
+        adapter_enabled.scatter_(
+            0, weight_indices.long(), (active_ranks > 0).to(torch.int32)
+        )
+    if num_tokens == 0:
+        return adapter_enabled, token_lora_mapping
+    if not has_segments:
+        token_lora_mapping.fill_(-1)
+        return adapter_enabled, token_lora_mapping
+
+    token_positions = torch.arange(
+        num_tokens, device=seg_indptr.device, dtype=torch.int32
+    )
+    # There is a torch.compile bug so we can't use seg_indptr[1:] here.
+    # Instead we pass seg_indptr and then subtract 1 from the result.
+    # This works because seg_indptr[0] == 0.
+    req_indices = (
+        torch.searchsorted(seg_indptr.to(torch.int32), token_positions, right=True) - 1
+    )
+
+    token_lora_mapping = torch.index_select(
+        weight_indices.to(torch.int32), 0, req_indices, out=token_lora_mapping
+    )
+
+    return adapter_enabled, token_lora_mapping

python/sglang/srt/lora/backend/chunked_backend.py

@@ -0,0 +1,525 @@
+import dataclasses
+from typing import List, Optional, Tuple
+
+import torch
+
+from sglang.kernels.ops.gemm.chunked_embedding_lora_a import (
+    chunked_embedding_lora_a_forward,
+)
+from sglang.kernels.ops.gemm.chunked_sgmv_expand import chunked_sgmv_lora_expand_forward
+from sglang.kernels.ops.gemm.chunked_sgmv_shrink import chunked_sgmv_lora_shrink_forward
+from sglang.srt.lora.backend.base_backend import BaseLoRABackend
+from sglang.srt.lora.utils import (
+    LoRABatchInfo,
+    generate_sequence_lengths,
+    get_lm_head_pruned_lens,
+    merge_and_chunk_segments,
+)
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.server_args import ServerArgs
+
+MIN_CHUNK_SIZE = 16
+
+
+class ChunkedSgmvLoRABackend(BaseLoRABackend):
+    """
+    Chunked LoRA backend using segmented matrix-vector multiplication.
+
+    This backend is largely based on the SGMV (Segmented Gather Matrix-Vector multiplication) algorithm
+    introduced in the Punica paper (https://arxiv.org/pdf/2310.18547). One main variation made here is to
+    segment the input sequences into fixed-size chunks, which reduces excessive kernel launches especially
+    when the LoRA distribution is skewed.
+    """
+
+    name = "csgmv"
+
+    def __init__(
+        self,
+        max_loras_per_batch: int,
+        device: torch.device,
+        server_args: ServerArgs,
+    ):
+        super().__init__(max_loras_per_batch, device)
+        self.max_chunk_size = server_args.max_lora_chunk_size
+
+    def run_lora_a_embedding(
+        self,
+        input_ids: torch.Tensor,
+        weights: torch.Tensor,
+        vocab_size: int,
+        extra_embeddings: torch.Tensor = None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        assert (
+            extra_embeddings is None
+        ), "Extra embeddings for lora a is not supported yet in chunked backend"
+        return chunked_embedding_lora_a_forward(
+            input_ids=input_ids,
+            weights=weights,
+            batch_info=self.batch_info,
+            vocab_size=vocab_size,
+        )
+
+    def run_lora_a_sgemm(
+        self,
+        x: torch.Tensor,
+        weights: torch.Tensor,
+        pruned_batch_info: LoRABatchInfo = None,
+        stack_num: int = 1,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        batch_info = (
+            pruned_batch_info if pruned_batch_info is not None else self.batch_info
+        )
+        return chunked_sgmv_lora_shrink_forward(
+            x=x,
+            weights=weights,
+            batch_info=batch_info,
+            num_slices=stack_num,
+        )
+
+    def run_lora_b_sgemm(
+        self,
+        x: torch.Tensor,
+        weights: torch.Tensor,
+        output_offset: torch.Tensor,
+        base_output: torch.Tensor = None,
+        pruned_batch_info: LoRABatchInfo = None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        # For simple lora B, we use slice offsets [0, output_dim]
+        output_dim = weights.shape[-2]
+        max_slice_size = output_dim
+        batch_info = (
+            pruned_batch_info if pruned_batch_info is not None else self.batch_info
+        )
+        return chunked_sgmv_lora_expand_forward(
+            x=x,
+            weights=weights,
+            batch_info=batch_info,
+            slice_offsets=output_offset,
+            max_slice_size=max_slice_size,
+            base_output=base_output,
+        )
+
+    def run_qkv_lora(
+        self,
+        x: torch.Tensor,
+        qkv_lora_a: torch.Tensor,
+        qkv_lora_b: torch.Tensor,
+        output_offset: torch.Tensor,
+        max_qkv_out_dim: int,
+        base_output: torch.Tensor = None,
+        n_slices: int = 3,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+
+        # x: (s, input_dim)
+        # qkv_lora_a: (num_lora, n_slices * r, input_dim)
+        # qkv_lora_b: (num_lora, total_output_dim, r)
+        assert isinstance(qkv_lora_b, torch.Tensor)
+
+        lora_a_output = chunked_sgmv_lora_shrink_forward(
+            x=x,
+            weights=qkv_lora_a,
+            batch_info=self.batch_info,
+            num_slices=n_slices,
+        )
+        lora_output = chunked_sgmv_lora_expand_forward(
+            x=lora_a_output,
+            weights=qkv_lora_b,
+            batch_info=self.batch_info,
+            slice_offsets=output_offset,
+            max_slice_size=max_qkv_out_dim,
+            base_output=base_output,
+        )
+        return lora_output
+
+    def run_gate_up_lora(
+        self,
+        x: torch.Tensor,
+        gate_up_lora_a: torch.Tensor,
+        gate_up_lora_b: torch.Tensor,
+        output_offset: torch.Tensor,
+        base_output: torch.Tensor = None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+
+        # x: (s, input_dim)
+        # gate_up_lora_a: (num_lora, 2 * r, input_dim)
+        # gate_up_lora_b: (num_lora, 2 * output_dim, r)
+        assert isinstance(gate_up_lora_b, torch.Tensor)
+        output_dim = gate_up_lora_b.shape[-2] // 2
+
+        # lora_a_output: (s, 2 * r)
+        lora_a_output = chunked_sgmv_lora_shrink_forward(
+            x=x,
+            weights=gate_up_lora_a,
+            batch_info=self.batch_info,
+            num_slices=2,
+        )
+        lora_output = chunked_sgmv_lora_expand_forward(
+            x=lora_a_output,
+            weights=gate_up_lora_b,
+            batch_info=self.batch_info,
+            slice_offsets=output_offset,
+            max_slice_size=output_dim,
+            base_output=base_output,
+        )
+        return lora_output
+
+    def _determine_chunk_size(self, forward_batch: ForwardBatch) -> int:
+        """
+        Heuristically determine the chunk size based on token token number in a batch.
+
+        Args:
+            forward_batch (ForwardBatch): The batch information containing sequence lengths.
+
+        Returns:
+            The determined chunk size
+        """
+        num_tokens = (
+            forward_batch.extend_num_tokens
+            if forward_batch.forward_mode.is_extend()
+            else forward_batch.batch_size
+        )
+        return self._determine_chunk_size_for_tokens(num_tokens)
+
+    def _determine_chunk_size_for_tokens(self, num_tokens: int) -> int:
+        """Determine chunk size given a token count directly."""
+        if self.max_chunk_size <= MIN_CHUNK_SIZE:
+            return MIN_CHUNK_SIZE
+
+        if num_tokens >= 256:
+            chunk_size = 128
+        elif num_tokens >= 64:
+            chunk_size = 32
+        else:  # num_tokens < 64
+            chunk_size = 16
+        return min(self.max_chunk_size, chunk_size)
+
+    @staticmethod
+    def _build_req_seg_indptr(forward_batch: ForwardBatch) -> torch.Tensor:
+        """Build per-request cumulative token boundaries on CPU (pinned)."""
+        bs = forward_batch.batch_size
+        if forward_batch.forward_mode.is_decode():
+            indptr = torch.arange(bs + 1, dtype=torch.int32, pin_memory=True)
+        else:
+            seg_lens = generate_sequence_lengths(forward_batch, device="cpu")
+            indptr = torch.zeros(bs + 1, dtype=torch.int32, pin_memory=True)
+            torch.cumsum(seg_lens, dim=0, out=indptr[1:])
+        return indptr
+
+    def init_cuda_graph_batch_info(
+        self,
+        max_bs_in_cuda_graph: int,
+        num_tokens_per_bs: int,
+    ):
+        max_num_segments = (
+            (num_tokens_per_bs + MIN_CHUNK_SIZE - 1) // MIN_CHUNK_SIZE
+        ) * max_bs_in_cuda_graph
+        max_num_tokens = max_bs_in_cuda_graph * num_tokens_per_bs
+        with torch.device("cuda"):
+            self.cuda_graph_batch_info = LoRABatchInfo(
+                bs=max_bs_in_cuda_graph,
+                use_cuda_graph=True,
+                seg_lens=torch.zeros(max_num_segments, dtype=torch.int32),
+                seg_indptr=torch.zeros(max_num_segments + 1, dtype=torch.int32),
+                weight_indices=torch.zeros(max_num_segments, dtype=torch.int32),
+                permutation=torch.zeros(max_num_tokens, dtype=torch.int32),
+                lora_ranks=torch.zeros(self.max_loras_per_batch, dtype=torch.int32),
+                scalings=torch.zeros(self.max_loras_per_batch, dtype=torch.float),
+                num_segments=None,  # Set per batch
+                max_len=None,  # Not used in CSGMV backend
+                req_seg_indptr=torch.zeros(max_bs_in_cuda_graph + 1, dtype=torch.int32),
+                req_weight_indices=torch.zeros(max_bs_in_cuda_graph, dtype=torch.int32),
+            )
+
+    def prepare_lora_batch(
+        self,
+        forward_batch: ForwardBatch,
+        weight_indices: list[int],
+        lora_ranks: list[int],
+        scalings: list[float],
+        use_cuda_graph: bool,
+    ):
+        chunk_size = self._determine_chunk_size(forward_batch)
+
+        permutation, weight_indices_reordered = ChunkedSgmvLoRABackend._get_permutation(
+            seq_weight_indices=weight_indices,
+            forward_batch=forward_batch,
+        )
+
+        seg_weight_indices, seg_indptr = self._get_segments_info(
+            weights_reordered=weight_indices_reordered,
+            chunk_size=chunk_size,
+        )
+        num_segments = len(seg_weight_indices)
+
+        lora_ranks_tensor = torch.tensor(
+            lora_ranks, dtype=torch.int32, pin_memory=True, device="cpu"
+        )
+        scalings_tensor = torch.tensor(
+            scalings, dtype=torch.float, pin_memory=True, device="cpu"
+        )
+
+        bs = forward_batch.batch_size
+        req_wi_tensor = torch.tensor(
+            weight_indices, dtype=torch.int32, pin_memory=True, device="cpu"
+        )
+        req_seg_indptr_cpu = self._build_req_seg_indptr(forward_batch)
+        max_num_segments = 0
+        has_unused_cuda_graph_segments = False
+
+        if not use_cuda_graph:
+            batch_info = LoRABatchInfo(
+                bs=bs,
+                num_segments=num_segments,
+                max_len=chunk_size,
+                use_cuda_graph=False,
+                seg_indptr=torch.empty(
+                    (num_segments + 1,), dtype=torch.int32, device=self.device
+                ),
+                weight_indices=torch.empty(
+                    (num_segments,), dtype=torch.int32, device=self.device
+                ),
+                lora_ranks=torch.empty(
+                    (self.max_loras_per_batch,), dtype=torch.int32, device=self.device
+                ),
+                scalings=torch.empty(
+                    (self.max_loras_per_batch,), dtype=torch.float, device=self.device
+                ),
+                permutation=torch.empty(
+                    (len(permutation),), dtype=torch.int32, device=self.device
+                ),
+                seg_lens=None,
+                req_seg_indptr=torch.empty(
+                    (bs + 1,), dtype=torch.int32, device=self.device
+                ),
+                req_weight_indices=torch.empty(
+                    (bs,), dtype=torch.int32, device=self.device
+                ),
+            )
+        else:
+            batch_info = self.cuda_graph_batch_info
+            batch_info.bs = bs
+            batch_info.num_segments = num_segments
+            batch_info.max_len = chunk_size
+            max_num_segments = batch_info.weight_indices.shape[0]
+            has_unused_cuda_graph_segments = num_segments < max_num_segments
+
+        # Copy to device asynchronously
+        batch_info.lora_ranks[: self.max_loras_per_batch].copy_(
+            lora_ranks_tensor, non_blocking=True
+        )
+        batch_info.scalings[: self.max_loras_per_batch].copy_(
+            scalings_tensor, non_blocking=True
+        )
+        batch_info.weight_indices[:num_segments].copy_(
+            seg_weight_indices, non_blocking=True
+        )
+        if has_unused_cuda_graph_segments:
+            batch_info.weight_indices[num_segments:max_num_segments].zero_()
+        batch_info.seg_indptr[: num_segments + 1].copy_(seg_indptr, non_blocking=True)
+        if has_unused_cuda_graph_segments:
+            batch_info.seg_indptr[num_segments + 1 : max_num_segments + 1].fill_(
+                int(seg_indptr[-1])
+            )
+        batch_info.permutation[: len(permutation)].copy_(permutation, non_blocking=True)
+        batch_info.req_seg_indptr[: bs + 1].copy_(req_seg_indptr_cpu, non_blocking=True)
+        batch_info.req_weight_indices[:bs].copy_(req_wi_tensor, non_blocking=True)
+
+        batch_info = self._add_moe_lora_info(forward_batch, batch_info)
+
+        self.batch_info = batch_info
+        self.lm_head_batch_info, self.lm_head_pass_batch_infos = (
+            self._prepare_lm_head_batch_info(forward_batch, weight_indices, batch_info)
+        )
+
+    def _prepare_lm_head_batch_info(
+        self,
+        forward_batch: ForwardBatch,
+        weight_indices: list[int],
+        batch_info: LoRABatchInfo,
+    ) -> Tuple[Optional[LoRABatchInfo], Optional[List[LoRABatchInfo]]]:
+
+        # Precompute lm_head_batch_info for pruned lm_head LoRA
+        pruned_lens = get_lm_head_pruned_lens(forward_batch)
+        lm_head_batch_info = None
+        lm_head_pass_batch_infos = None
+
+        if pruned_lens is not None:
+            pruned_total = sum(pruned_lens)
+            chunk_size = self._determine_chunk_size_for_tokens(pruned_total)
+            lm_head_segments = merge_and_chunk_segments(
+                weight_indices, pruned_lens, chunk_size=chunk_size
+            )
+            lm_head_batch_info = self._build_lm_head_batch_info(
+                lm_head_segments, batch_info, chunk_size, pruned_total
+            )
+
+            # Precompute per-pass batch_infos for logprobs chunking
+            pass_segments = self._get_lm_head_pass_segments(weight_indices, pruned_lens)
+            if pass_segments is not None:
+                lm_head_pass_batch_infos = []
+                for seg_wi, seg_lens_list in pass_segments:
+                    pass_total = sum(seg_lens_list)
+                    pass_chunk_size = self._determine_chunk_size_for_tokens(pass_total)
+                    chunked_segments = merge_and_chunk_segments(
+                        seg_wi, seg_lens_list, chunk_size=pass_chunk_size
+                    )
+                    lm_head_pass_batch_infos.append(
+                        self._build_lm_head_batch_info(
+                            chunked_segments,
+                            batch_info,
+                            pass_chunk_size,
+                            pass_total,
+                        )
+                    )
+
+        return lm_head_batch_info, lm_head_pass_batch_infos
+
+    def _build_lm_head_batch_info(
+        self,
+        lm_head_segments: Tuple[List[int], List[int]],
+        batch_info: LoRABatchInfo,
+        chunk_size: int,
+        expected_tokens: int,
+    ) -> LoRABatchInfo:
+        seg_weight_indices_cpu, seg_lens_cpu = lm_head_segments
+        pruned_total = sum(seg_lens_cpu)
+        num_segments = len(seg_weight_indices_cpu)
+
+        weight_indices = torch.tensor(
+            seg_weight_indices_cpu, dtype=torch.int32, device=self.device
+        )
+        seg_lens = torch.tensor(seg_lens_cpu, dtype=torch.int32, device=self.device)
+        seg_indptr = torch.zeros(
+            (num_segments + 1,), dtype=torch.int32, device=self.device
+        )
+        seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
+
+        # Identity permutation (lm_head tokens are in original order)
+        permutation = torch.arange(pruned_total, dtype=torch.int32, device=self.device)
+
+        return dataclasses.replace(
+            batch_info,
+            num_segments=num_segments,
+            max_len=chunk_size,
+            seg_indptr=seg_indptr,
+            weight_indices=weight_indices,
+            permutation=permutation,
+            expected_tokens=expected_tokens,
+        )
+
+    @staticmethod
+    def _get_permutation(seq_weight_indices, forward_batch: ForwardBatch):
+        """
+        Computes permutation indices for reordering tokens by their LoRA adapter assignments.
+
+        This function implements the "gather" step in Chunked Segmented Gather Matrix Vector
+        multiplication by creating a permutation that groups tokens by their LoRA adapter.
+        Tokens using the same LoRA adapter are placed together to enable efficient batched
+        computation.
+
+        Example:
+            seq_weight_indices = [0, 1, 0]  # 3 sequences using adapters [0, 1, 0]
+            extend_seq_lens = [2, 1, 3]     # sequence lengths [2, 1, 3 tokens]
+
+            # Creates row_weight_indices: [0, 0, 1, 0, 0, 0] (6 tokens total)
+            # Returns permutation: [0, 1, 3, 4, 5, 2] (groups adapter 0 tokens together)
+            # weights_reordered: [0, 0, 0, 0, 0, 1] (sorted by adapter)
+
+        Args:
+            seq_weight_indices: List of LoRA adapter indices for each sequence
+            forward_batch (ForwardBatch): Batch information containing sequence lengths
+
+        Returns:
+            tuple: (permutation, weights_reordered) where:
+                - permutation: Token reordering indices to group by adapter
+                - weights_reordered: Sorted adapter indices for each token
+        """
+        with torch.device("cpu"):
+            seq_weight_indices = torch.tensor(seq_weight_indices, dtype=torch.int32)
+            seg_lens_cpu = generate_sequence_lengths(forward_batch)
+
+            row_weight_indices = torch.repeat_interleave(
+                seq_weight_indices, seg_lens_cpu
+            )
+            permutation = torch.empty(
+                (len(row_weight_indices),), dtype=torch.long, pin_memory=True
+            )
+            torch.argsort(row_weight_indices, stable=True, out=permutation)
+            weights_reordered = row_weight_indices[permutation]
+
+            return permutation, weights_reordered
+
+    def _get_segments_info(self, weights_reordered: torch.Tensor, chunk_size: int):
+        """
+        Computes segment information for chunked SGMV operations.
+
+        This function takes the reordered weight indices and creates segments of fixed size
+        (self.segment_size) for efficient kernel execution. Each segment contains tokens
+        that use the same LoRA adapter, enabling vectorized computation.
+
+        The segmentation is necessary because:
+        1. GPU kernels work efficiently on fixed-size blocks
+        2. Large groups of tokens using the same adapter are split into manageable chunks
+        3. Each segment can be processed independently in parallel
+
+        Example:
+            weights_reordered = [0, 0, 0, 0, 0, 1]  # 5 tokens with adapter 0, 1 with adapter 1
+            segment_size = 3
+
+            # Creates segments:
+            # Segment 0: tokens 0-2 (adapter 0), length=3
+            # Segment 1: tokens 3-4 (adapter 0), length=2
+            # Segment 2: token 5 (adapter 1), length=1
+
+            # Returns:
+            # weight_indices_list: [0, 0, 1] (adapter for each segment)
+            # seg_indptr: [0, 3, 5, 6] (cumulative segment boundaries)
+
+        Args:
+            weights_reordered (torch.Tensor): Sorted adapter indices for each token
+            chunk_size (int): Fixed size for each segment
+
+        Returns:
+            tuple: (weight_indices_list, seg_indptr) where:
+                - weight_indices_list: LoRA adapter index for each segment
+                - seg_indptr: Cumulative segment boundaries (CSR-style indptr)
+        """
+        with torch.device("cpu"):
+            unique_weights, counts = torch.unique_consecutive(
+                weights_reordered, return_counts=True
+            )
+
+            weight_indices_list = []
+            seg_lens_list = []
+
+            for weight_idx, group_len in zip(unique_weights, counts):
+                group_len = group_len.item()
+                num_segs = (group_len + chunk_size - 1) // chunk_size
+
+                weight_indices_list.extend([weight_idx.item()] * num_segs)
+                seg_lens_list.extend([chunk_size] * (num_segs - 1))
+                seg_lens_list.append(group_len - (num_segs - 1) * chunk_size)
+
+            seg_lens = torch.tensor(seg_lens_list, dtype=torch.int32)
+
+            weight_indices_list = torch.tensor(
+                weight_indices_list, dtype=torch.int32, pin_memory=True
+            )
+
+            seg_indptr = torch.empty(
+                (len(seg_lens) + 1,), dtype=torch.int32, pin_memory=True
+            )
+            seg_indptr[0] = 0
+            seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
+
+            return weight_indices_list, seg_indptr

python/sglang/srt/lora/backend/lmhead_mixing.py

@@ -0,0 +1,64 @@
+from typing import List, Optional, Tuple
+
+from sglang.srt.environ import envs
+from sglang.srt.lora.utils import LoRABatchInfo, build_lm_head_pass_segments
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+class LoRABackendLmHeadMixing:
+    def init_lm_head_config(self):
+        self.lm_head_batch_info = None
+        # Precomputed per-pass lm_head batch_infos.  When the logits processor
+        # calls lm_head in multiple passes (chunked logprobs), each pass gets
+        # its own batch_info from this list.
+        self.lm_head_pass_batch_infos = None
+        # Current pass index.  When set, apply_lora uses
+        # lm_head_pass_batch_infos[idx] instead of lm_head_batch_info.
+        self._lm_head_pass_idx = None
+
+    def _get_lm_head_pass_segments(
+        self,
+        weight_indices: list[int],
+        pruned_lens: List[int],
+    ) -> Optional[List[Tuple[List[int], List[int]]]]:
+        """Compute per-pass segment info for lm_head LoRA logprobs chunking.
+
+        When LogitsProcessor splits pruned states into fixed-size passes,
+        each pass needs its own segmentation so that lm_head LoRA operates
+        on the correct adapter assignments.  This method returns the generic
+        per-pass (seg_weight_indices, seg_lens) tuples; each backend is
+        responsible for converting them into backend-specific LoRABatchInfo.
+
+        Returns None if logprobs chunking is disabled or the pruned token
+        count does not exceed the logprobs chunk size.
+        """
+        logprobs_chunk_size = envs.SGLANG_LOGITS_PROCESSER_CHUNK_SIZE.get()
+        enable_logprobs_chunk = envs.SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK.get()
+        pruned_total = sum(pruned_lens)
+
+        if not enable_logprobs_chunk or pruned_total <= logprobs_chunk_size:
+            return None
+
+        return build_lm_head_pass_segments(
+            weight_indices, pruned_lens, logprobs_chunk_size
+        )
+
+    def _prepare_lm_head_batch_info(
+        self,
+        forward_batch: ForwardBatch,
+        weight_indices: list[int],
+        batch_info: LoRABatchInfo,
+    ) -> Tuple[Optional[LoRABatchInfo], Optional[List[LoRABatchInfo]]]:
+        """Prepare the lm_head batch info for the current forward batch."""
+        """It returns a tuple of (lm_head_batch_info, lm_head_pass_batch_infos)."""
+        pass
+
+    def _build_lm_head_batch_info(
+        self,
+        lm_head_segments: Tuple[List[int], List[int]],
+        batch_info: LoRABatchInfo,
+        chunk_size: int,
+        expected_tokens: int,
+    ) -> LoRABatchInfo:
+        """Build a LoRABatchInfo for pruned lm_head input."""
+        pass

python/sglang/srt/lora/backend/lora_registry.py

@@ -0,0 +1,61 @@
+import logging
+from typing import Type
+
+from sglang.srt.lora.backend.base_backend import BaseLoRABackend
+
+logger = logging.getLogger(__name__)
+
+LORA_SUPPORTED_BACKENDS = {}
+
+
+def register_lora_backend(name):
+    def decorator(fn):
+        LORA_SUPPORTED_BACKENDS[name] = fn
+        return fn
+
+    return decorator
+
+
+@register_lora_backend("triton")
+def create_triton_backend():
+    from sglang.srt.lora.backend.triton_backend import TritonLoRABackend
+
+    return TritonLoRABackend
+
+
+@register_lora_backend("csgmv")
+def create_triton_csgmv_backend():
+    from sglang.srt.lora.backend.chunked_backend import ChunkedSgmvLoRABackend
+
+    return ChunkedSgmvLoRABackend
+
+
+@register_lora_backend("ascend")
+def create_ascend_backend():
+    from sglang.srt.lora.backend.ascend_backend import AscendLoRABackend
+
+    return AscendLoRABackend
+
+
+@register_lora_backend("torch_native")
+def create_torch_native_backend():
+    from sglang.srt.lora.backend.torch_backend import TorchNativeLoRABackend
+
+    return TorchNativeLoRABackend
+
+
+@register_lora_backend("flashinfer")
+def create_flashinfer_backend():
+    raise ValueError(
+        "FlashInfer LoRA backend has been deprecated, please use `triton` instead."
+    )
+
+
+def get_backend_from_name(name: str) -> Type[BaseLoRABackend]:
+    """
+    Get corresponding backend class from backend's name
+    """
+    if name not in LORA_SUPPORTED_BACKENDS:
+        raise ValueError(f"Invalid backend: {name}")
+    lora_backend = LORA_SUPPORTED_BACKENDS[name]()
+    return lora_backend

python/sglang/srt/lora/backend/torch_backend.py

@@ -0,0 +1,302 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+from sglang.srt.lora.backend.base_backend import BaseLoRABackend
+from sglang.srt.lora.torch_ops import (
+    sgemm_lora_a_embedding_fwd,
+    sgemm_lora_a_fwd,
+    sgemm_lora_b_fwd,
+)
+from sglang.srt.lora.utils import LoRABatchInfo, generate_sequence_lengths
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+@dataclass
+class TorchNativeLoRABatchInfo(LoRABatchInfo):
+    # ranks of each lora adapter, in shape (lora_num,) placed on cpu device
+    lora_ranks_cpu: Optional[torch.Tensor] = None
+
+    # Indice pointers of each segment in shape (num_segments + 1, ) placed on cpu device
+    seg_indptr_cpu: Optional[torch.Tensor] = None
+
+    # Lengths of each segments in shape (num_segments,) placed on cpu device
+    seg_lens_cpu: Optional[torch.Tensor] = None
+
+    # The index of lora adapter used by each segment, in shape (num_segments,) placed on cpu device
+    weight_indices_cpu: Optional[torch.Tensor] = None
+
+    # Scaling factors for each lora adapter, in shape (lora_num,) placed on cpu device
+    scalings_cpu: Optional[torch.Tensor] = None
+
+
+class TorchNativeLoRABackend(BaseLoRABackend):
+    name = "torch_native"
+
+    def __init__(
+        self,
+        max_loras_per_batch: int,
+        device: torch.device,
+        **kwargs,
+    ):
+        super().__init__(max_loras_per_batch, device)
+
+    def run_lora_a_embedding(
+        self,
+        input_ids: torch.Tensor,
+        weights: torch.Tensor,
+        vocab_size: int,
+        extra_embeddings: torch.Tensor = None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        assert (
+            extra_embeddings is None
+        ), "Extra embeddings for lora a is not supported yet in chunked backend"
+        output_tensor = sgemm_lora_a_embedding_fwd(
+            inputs=input_ids,
+            weights=weights,
+            batch_info=self.batch_info,
+            vocab_size=vocab_size,
+        )
+
+        return output_tensor
+
+    def run_lora_a_sgemm(
+        self,
+        x: torch.Tensor,
+        weights: torch.Tensor,
+        stack_num: int = 1,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        output_tensor = sgemm_lora_a_fwd(
+            inputs=x,
+            weights=weights,
+            batch_info=self.batch_info,
+            num_slices=stack_num,
+        )
+
+        return output_tensor
+
+    def run_lora_b_sgemm(
+        self,
+        x: torch.Tensor,
+        weights: torch.Tensor,
+        output_offset_cpu: torch.Tensor,
+        base_output: torch.Tensor = None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        _, weight_out_dim, _ = weights.shape
+
+        output_tensor = sgemm_lora_b_fwd(
+            inputs=x,
+            weights=weights,
+            batch_info=self.batch_info,
+            slice_offsets=output_offset_cpu,
+            base_output=base_output,
+        )
+
+        return output_tensor
+
+    def run_qkv_lora(
+        self,
+        x: torch.Tensor,
+        qkv_lora_a: torch.Tensor,
+        qkv_lora_b: torch.Tensor,
+        output_offset: torch.Tensor,
+        output_offset_cpu: torch.Tensor,
+        max_qkv_out_dim: int,
+        base_output: torch.Tensor = None,
+        n_slices: int = 3,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        lora_a_output = sgemm_lora_a_fwd(
+            inputs=x,
+            weights=qkv_lora_a,
+            batch_info=self.batch_info,
+            num_slices=n_slices,
+        )
+
+        output_tensor = sgemm_lora_b_fwd(
+            inputs=lora_a_output,
+            weights=qkv_lora_b,
+            batch_info=self.batch_info,
+            slice_offsets=output_offset_cpu,
+            base_output=base_output,
+        )
+
+        return output_tensor
+
+    def run_gate_up_lora(
+        self,
+        x: torch.Tensor,
+        gate_up_lora_a: torch.Tensor,
+        gate_up_lora_b: torch.Tensor,
+        output_offset_cpu: torch.Tensor,
+        base_output: torch.Tensor = None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        num_slices = len(output_offset_cpu) - 1
+        _, weight_out_dim, _ = gate_up_lora_b.shape
+
+        lora_a_output = sgemm_lora_a_fwd(
+            inputs=x,
+            weights=gate_up_lora_a,
+            batch_info=self.batch_info,
+            num_slices=num_slices,
+        )
+
+        output_tensor = sgemm_lora_b_fwd(
+            inputs=lora_a_output,
+            weights=gate_up_lora_b,
+            batch_info=self.batch_info,
+            slice_offsets=output_offset_cpu,
+            base_output=base_output,
+        )
+
+        return output_tensor
+
+    def init_cuda_graph_batch_info(
+        self,
+        max_bs_in_cuda_graph: int,
+        num_tokens_per_bs: int,
+    ):
+        with torch.device("cuda"):
+            self.cuda_graph_batch_info = TorchNativeLoRABatchInfo(
+                use_cuda_graph=True,
+                bs=max_bs_in_cuda_graph,
+                num_segments=self.max_loras_per_batch,
+                seg_lens=torch.full(
+                    (max_bs_in_cuda_graph,), num_tokens_per_bs, dtype=torch.int32
+                ),
+                seg_indptr=torch.zeros(max_bs_in_cuda_graph + 1, dtype=torch.int32),
+                weight_indices=torch.zeros(max_bs_in_cuda_graph, dtype=torch.int32),
+                lora_ranks=torch.zeros(self.max_loras_per_batch, dtype=torch.int32),
+                scalings=torch.zeros(self.max_loras_per_batch, dtype=torch.float),
+                permutation=None,
+                max_len=num_tokens_per_bs,
+            )
+
+            # Initialize seg_indptr for CUDA graph as they remain constant
+            # across batches.
+            torch.cumsum(
+                self.cuda_graph_batch_info.seg_lens[:max_bs_in_cuda_graph],
+                dim=0,
+                out=self.cuda_graph_batch_info.seg_indptr[1 : max_bs_in_cuda_graph + 1],
+            )
+
+    def prepare_lora_batch(
+        self,
+        forward_batch: ForwardBatch,
+        weight_indices: list[int],
+        lora_ranks: list[int],
+        scalings: list[float],
+        use_cuda_graph: bool,
+    ):
+        # Do not use merge optimization for graph mode
+        # Use pinned memory to avoid synchronizations during host-to-device transfer
+        original_seq_lens_cpu = generate_sequence_lengths(forward_batch, device="cpu")
+        if not use_cuda_graph:
+            original_weight_indices_tensor = torch.tensor(
+                weight_indices, dtype=torch.int32, device="cpu"
+            )
+
+            unique_weight_indices_tensor, inverse_weight_indices_tensor = (
+                torch.unique_consecutive(
+                    original_weight_indices_tensor, return_inverse=True
+                )
+            )
+
+            seg_lens_cpu = (
+                torch.zeros_like(
+                    unique_weight_indices_tensor, dtype=torch.int32, device="cpu"
+                )
+                .scatter_add_(
+                    0,
+                    inverse_weight_indices_tensor,
+                    original_seq_lens_cpu,
+                )
+                .pin_memory()
+            )
+
+            weight_indices_tensor = unique_weight_indices_tensor.pin_memory()
+        else:
+            weight_indices_tensor = torch.repeat_interleave(
+                torch.tensor(weight_indices, dtype=torch.int32, device="cpu"),
+                original_seq_lens_cpu,
+            ).pin_memory()
+            seg_lens_cpu = torch.ones_like(weight_indices_tensor).pin_memory()
+
+        seg_indptr_cpu = torch.zeros(
+            (len(seg_lens_cpu) + 1,), dtype=torch.int32, pin_memory=True
+        )
+        seg_indptr_cpu[1:] = torch.cumsum(seg_lens_cpu, dim=0)
+        lora_ranks_tensor = torch.tensor(
+            lora_ranks, dtype=torch.int32, pin_memory=True, device="cpu"
+        )
+        scalings_tensor = torch.tensor(
+            scalings, dtype=torch.float, pin_memory=True, device="cpu"
+        )
+
+        bs = forward_batch.batch_size
+        num_segments = len(weight_indices_tensor)
+
+        if use_cuda_graph:
+            assert (
+                self.cuda_graph_batch_info is not None
+            ), "CUDA Graph batch info is not initialized."
+            batch_info = self.cuda_graph_batch_info
+            batch_info.bs = forward_batch.batch_size
+            batch_info.num_segments = num_segments
+        else:
+            max_len = max(seg_lens_cpu)
+
+            batch_info = TorchNativeLoRABatchInfo(
+                bs=forward_batch.batch_size,
+                num_segments=num_segments,
+                max_len=max_len,
+                use_cuda_graph=False,
+                seg_lens=torch.empty((bs,), dtype=torch.int32, device=self.device),
+                seg_indptr=torch.empty(
+                    (bs + 1,), dtype=torch.int32, device=self.device
+                ),
+                weight_indices=torch.empty(
+                    (bs,), dtype=torch.int32, device=self.device
+                ),
+                lora_ranks=torch.empty(
+                    (self.max_loras_per_batch,), dtype=torch.int32, device=self.device
+                ),
+                scalings=torch.empty(
+                    (self.max_loras_per_batch,), dtype=torch.float, device=self.device
+                ),
+                permutation=None,
+            )
+
+        # Copy to device asynchronously
+        batch_info.lora_ranks[: self.max_loras_per_batch].copy_(
+            lora_ranks_tensor, non_blocking=True
+        )
+        batch_info.scalings[: self.max_loras_per_batch].copy_(
+            scalings_tensor, non_blocking=True
+        )
+        batch_info.weight_indices[:num_segments].copy_(
+            weight_indices_tensor, non_blocking=True
+        )
+        batch_info.seg_indptr[: len(seg_indptr_cpu)].copy_(
+            seg_indptr_cpu, non_blocking=True
+        )
+        batch_info.seg_lens[: len(seg_lens_cpu)].copy_(seg_lens_cpu, non_blocking=True)
+
+        batch_info.lora_ranks_cpu = lora_ranks_tensor
+        batch_info.seg_indptr_cpu = seg_indptr_cpu
+        batch_info.seg_lens_cpu = seg_lens_cpu
+        batch_info.weight_indices_cpu = weight_indices_tensor
+        batch_info.scalings_cpu = scalings_tensor
+
+        batch_info = self._add_moe_lora_info(forward_batch, batch_info)
+        self.batch_info = batch_info

python/sglang/srt/lora/backend/triton_backend.py

@@ -0,0 +1,376 @@
+import dataclasses
+from typing import List, Optional, Tuple
+
+import torch
+
+from sglang.kernels.ops.gemm.embedding_lora_a import embedding_lora_a_fwd
+from sglang.kernels.ops.gemm.gate_up_lora_b import gate_up_lora_b_fwd
+from sglang.kernels.ops.gemm.qkv_lora_b import qkv_lora_b_fwd
+from sglang.kernels.ops.gemm.sgemm_lora_a import sgemm_lora_a_fwd
+from sglang.kernels.ops.gemm.sgemm_lora_b import sgemm_lora_b_fwd
+from sglang.srt.lora.backend.base_backend import BaseLoRABackend
+from sglang.srt.lora.utils import (
+    LoRABatchInfo,
+    get_lm_head_pruned_lens,
+    merge_and_chunk_segments,
+)
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+class TritonLoRABackend(BaseLoRABackend):
+    name = "triton"
+
+    def __init__(
+        self,
+        max_loras_per_batch: int,
+        device: torch.device,
+        **kwargs,
+    ):
+        super().__init__(max_loras_per_batch, device)
+
+    def run_lora_a_embedding(
+        self,
+        input_ids: torch.Tensor,
+        weights: torch.Tensor,
+        vocab_size: int,
+        extra_embeddings: torch.Tensor = None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        """Run LoRA A embedding lookup using Triton kernel."""
+        return embedding_lora_a_fwd(
+            input_ids=input_ids,
+            weights=weights,
+            batch_info=self.batch_info,
+            vocab_size=vocab_size,
+            extra_embeddings=extra_embeddings,
+        )
+
+    def _sgemm_info(self, pruned_batch_info=None):
+        """Return the sgemm batch_info (merged segments when available)."""
+        if pruned_batch_info is not None:
+            return pruned_batch_info
+        return getattr(self, "sgemm_batch_info", None) or self.batch_info
+
+    def run_lora_a_sgemm(
+        self,
+        x: torch.Tensor,
+        weights: torch.Tensor,
+        pruned_batch_info: LoRABatchInfo = None,
+        stack_num: int = 1,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        return sgemm_lora_a_fwd(
+            x, weights, self._sgemm_info(pruned_batch_info), stack_num=stack_num
+        )
+
+    def run_lora_b_sgemm(
+        self,
+        x: torch.Tensor,
+        weights: torch.Tensor,
+        base_output: torch.Tensor = None,
+        pruned_batch_info: LoRABatchInfo = None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        return sgemm_lora_b_fwd(
+            x, weights, self._sgemm_info(pruned_batch_info), base_output
+        )
+
+    def run_qkv_lora(
+        self,
+        x: torch.Tensor,
+        qkv_lora_a: torch.Tensor,
+        qkv_lora_b: torch.Tensor,
+        output_offset: torch.Tensor,
+        max_qkv_out_dim: int,
+        base_output: torch.Tensor = None,
+        n_slices: int = 3,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+
+        # x: (s, input_dim)
+        # qkv_lora_a: (num_lora, n_slices * r, input_dim)
+        # qkv_lora_b: (num_lora, total_output_dim, r)
+        assert isinstance(qkv_lora_b, torch.Tensor)
+
+        sgemm_info = self._sgemm_info()
+        lora_a_output = sgemm_lora_a_fwd(x, qkv_lora_a, sgemm_info, stack_num=n_slices)
+        lora_output = qkv_lora_b_fwd(
+            lora_a_output,
+            qkv_lora_b,
+            sgemm_info,
+            output_offset,
+            max_qkv_out_dim,
+            base_output,
+            n_slices=n_slices,
+        )
+        return lora_output
+
+    def run_gate_up_lora(
+        self,
+        x: torch.Tensor,
+        gate_up_lora_a: torch.Tensor,
+        gate_up_lora_b: torch.Tensor,
+        base_output: torch.Tensor = None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+
+        # x: (s, input_dim)
+        # gate_up_lora_a: (num_lora, 2 * r, input_dim)
+        # gate_up_lora_b: (num_lora, 2 * output_dim, r)
+        assert isinstance(gate_up_lora_b, torch.Tensor)
+        output_dim = gate_up_lora_b.shape[-2] // 2
+
+        sgemm_info = self._sgemm_info()
+        # lora_a_output: (s, 2 * r)
+        lora_a_output = sgemm_lora_a_fwd(x, gate_up_lora_a, sgemm_info, stack_num=2)
+        lora_output = gate_up_lora_b_fwd(
+            lora_a_output,
+            gate_up_lora_b,
+            sgemm_info,
+            output_dim,
+            base_output,
+        )
+        return lora_output
+
+    def init_cuda_graph_batch_info(
+        self,
+        max_bs_in_cuda_graph: int,
+        num_tokens_per_bs: int,
+    ):
+        max_tokens = max_bs_in_cuda_graph * num_tokens_per_bs
+        mlpb = self.max_loras_per_batch
+        with torch.device("cuda"):
+            self.cuda_graph_batch_info = LoRABatchInfo(
+                bs=max_bs_in_cuda_graph,
+                use_cuda_graph=True,
+                num_segments=None,
+                seg_lens=torch.full(
+                    (max_bs_in_cuda_graph,), num_tokens_per_bs, dtype=torch.int32
+                ),
+                seg_indptr=torch.zeros(max_bs_in_cuda_graph + 1, dtype=torch.int32),
+                max_len=num_tokens_per_bs,
+                weight_indices=torch.zeros(max_bs_in_cuda_graph, dtype=torch.int32),
+                lora_ranks=torch.zeros(mlpb, dtype=torch.int32),
+                scalings=torch.zeros(mlpb, dtype=torch.float),
+                permutation=None,
+            )
+
+            torch.cumsum(
+                self.cuda_graph_batch_info.seg_lens[:max_bs_in_cuda_graph],
+                dim=0,
+                out=self.cuda_graph_batch_info.seg_indptr[1 : max_bs_in_cuda_graph + 1],
+            )
+
+            # Sgemm batch_info with segments merged by adapter.
+            # Updated each batch by compute_sgemm_routing().
+            self.cuda_graph_sgemm_batch_info = LoRABatchInfo(
+                bs=mlpb,
+                use_cuda_graph=True,
+                num_segments=mlpb,
+                seg_lens=torch.zeros(mlpb, dtype=torch.int32),
+                seg_indptr=torch.zeros(mlpb + 1, dtype=torch.int32),
+                max_len=max_tokens,
+                weight_indices=torch.arange(mlpb, dtype=torch.int32),
+                lora_ranks=torch.zeros(mlpb, dtype=torch.int32),
+                scalings=torch.zeros(mlpb, dtype=torch.float),
+                permutation=torch.zeros(max_tokens, dtype=torch.int32),
+            )
+
+    def compute_sgemm_routing(self, use_cuda_graph: bool):
+        """Sort tokens by adapter and build merged segments for sgemm LoRA."""
+        bi = self.batch_info
+        bs = bi.bs
+        mlpb = self.max_loras_per_batch
+        wi = bi.weight_indices[:bs]
+
+        perm = torch.argsort(wi, stable=True).to(torch.int32)
+        sorted_wi = wi[perm]
+        adapter_ids = torch.arange(mlpb, device=wi.device, dtype=torch.int32)
+        seg_starts = torch.searchsorted(sorted_wi, adapter_ids)
+        seg_ends = torch.searchsorted(sorted_wi, adapter_ids, right=True)
+        seg_lens = seg_ends - seg_starts
+
+        if use_cuda_graph:
+            sgemm = getattr(self, "cuda_graph_sgemm_batch_info", None)
+            if sgemm is None:
+                return
+            sgemm.permutation[:bs] = perm
+            sgemm.seg_lens[:] = seg_lens
+            sgemm.seg_indptr[0:1].zero_()
+            torch.cumsum(sgemm.seg_lens, dim=0, out=sgemm.seg_indptr[1:])
+            sgemm.max_len = bs
+            sgemm.lora_ranks[:mlpb] = bi.lora_ranks[:mlpb]
+            sgemm.scalings[:mlpb] = bi.scalings[:mlpb]
+        else:
+            seg_indptr = torch.zeros(mlpb + 1, dtype=torch.int32, device=wi.device)
+            seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
+            sgemm = LoRABatchInfo(
+                bs=mlpb,
+                use_cuda_graph=False,
+                num_segments=mlpb,
+                seg_lens=seg_lens,
+                seg_indptr=seg_indptr,
+                max_len=bs,
+                weight_indices=adapter_ids,
+                lora_ranks=bi.lora_ranks[:mlpb].clone(),
+                scalings=bi.scalings[:mlpb].clone(),
+                permutation=perm,
+            )
+
+        self.sgemm_batch_info = sgemm
+
+    def prepare_lora_batch(
+        self,
+        forward_batch: ForwardBatch,
+        weight_indices: list[int],
+        lora_ranks: list[int],
+        scalings: list[float],
+        use_cuda_graph: bool,
+    ):
+        # Use pinned memory to avoid synchronizations during host-to-device transfer
+        weight_indices_tensor = torch.tensor(
+            weight_indices, dtype=torch.int32, pin_memory=True, device="cpu"
+        )
+        lora_ranks_tensor = torch.tensor(
+            lora_ranks, dtype=torch.int32, pin_memory=True, device="cpu"
+        )
+        scalings_tensor = torch.tensor(
+            scalings, dtype=torch.float, pin_memory=True, device="cpu"
+        )
+
+        bs = forward_batch.batch_size
+
+        if use_cuda_graph:
+            assert (
+                self.cuda_graph_batch_info is not None
+            ), "CUDA Graph batch info is not initialized."
+            batch_info = self.cuda_graph_batch_info
+            batch_info.bs = forward_batch.batch_size
+            batch_info.num_segments = forward_batch.batch_size
+        else:
+            max_len = (
+                # Calculate max_len from the CPU copy to avoid D2H transfer.
+                max(forward_batch.extend_seq_lens_cpu)
+                if forward_batch.forward_mode.is_extend()
+                else 1
+            )
+            seg_lens = (
+                forward_batch.extend_seq_lens
+                if forward_batch.forward_mode.is_extend()
+                else torch.ones(bs, dtype=torch.int32, device=self.device)
+            )
+            seg_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device=self.device)
+            seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
+
+            batch_info = LoRABatchInfo(
+                bs=forward_batch.batch_size,
+                num_segments=forward_batch.batch_size,
+                max_len=max_len,
+                use_cuda_graph=False,
+                seg_lens=seg_lens,
+                seg_indptr=seg_indptr,
+                weight_indices=torch.empty(
+                    (bs,), dtype=torch.int32, device=self.device
+                ),
+                lora_ranks=torch.empty(
+                    (self.max_loras_per_batch,), dtype=torch.int64, device=self.device
+                ),
+                scalings=torch.empty(
+                    (self.max_loras_per_batch,), dtype=torch.float, device=self.device
+                ),
+                permutation=None,
+            )
+
+        # Copy to device asynchronously
+        batch_info.lora_ranks[: self.max_loras_per_batch].copy_(
+            lora_ranks_tensor, non_blocking=True
+        )
+        batch_info.scalings[: self.max_loras_per_batch].copy_(
+            scalings_tensor, non_blocking=True
+        )
+        batch_info.weight_indices[:bs].copy_(weight_indices_tensor, non_blocking=True)
+
+        batch_info = self._add_moe_lora_info(forward_batch, batch_info)
+        self.batch_info = batch_info
+
+        # Biggest win is in decode.
+        is_decode = not forward_batch.forward_mode.is_extend()
+        if is_decode:
+            self.compute_sgemm_routing(use_cuda_graph)
+        else:
+            self.sgemm_batch_info = None
+
+        self.lm_head_batch_info, self.lm_head_pass_batch_infos = (
+            self._prepare_lm_head_batch_info(forward_batch, weight_indices, batch_info)
+        )
+
+    def _prepare_lm_head_batch_info(
+        self,
+        forward_batch: ForwardBatch,
+        weight_indices: list[int],
+        batch_info: LoRABatchInfo,
+    ) -> Tuple[Optional[LoRABatchInfo], Optional[List[LoRABatchInfo]]]:
+
+        # Precompute lm_head_batch_info for pruned lm_head LoRA
+        pruned_lens = get_lm_head_pruned_lens(forward_batch)
+        lm_head_batch_info = None
+        lm_head_pass_batch_infos = None
+
+        if pruned_lens is not None:
+            pruned_total = sum(pruned_lens)
+            lm_head_segments = merge_and_chunk_segments(
+                weight_indices, pruned_lens, chunk_size=pruned_total
+            )
+            lm_head_batch_info = self._build_lm_head_batch_info(
+                lm_head_segments, batch_info, pruned_total
+            )
+
+            # Precompute per-pass batch_infos for logprobs chunking
+            pass_segments = self._get_lm_head_pass_segments(weight_indices, pruned_lens)
+            if pass_segments is not None:
+                lm_head_pass_batch_infos = []
+                for seg_wi, seg_lens_list in pass_segments:
+                    pass_total = sum(seg_lens_list)
+                    merged_segments = merge_and_chunk_segments(
+                        seg_wi, seg_lens_list, chunk_size=pass_total
+                    )
+                    self.lm_head_pass_batch_infos.append(
+                        self._build_lm_head_batch_info(
+                            merged_segments, batch_info, pass_total
+                        )
+                    )
+
+        return lm_head_batch_info, lm_head_pass_batch_infos
+
+    def _build_lm_head_batch_info(
+        self,
+        lm_head_segments: Tuple[List[int], List[int]],
+        batch_info: LoRABatchInfo,
+        expected_tokens: int,
+    ) -> LoRABatchInfo:
+        seg_weight_indices_cpu, seg_lens_cpu = lm_head_segments
+        num_segments = len(seg_weight_indices_cpu)
+
+        seg_lens = torch.tensor(seg_lens_cpu, dtype=torch.int32, device=self.device)
+        seg_indptr = torch.zeros(
+            (num_segments + 1,), dtype=torch.int32, device=self.device
+        )
+        seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
+
+        return dataclasses.replace(
+            batch_info,
+            bs=num_segments,
+            num_segments=num_segments,
+            max_len=max(seg_lens_cpu),
+            seg_lens=seg_lens,
+            seg_indptr=seg_indptr,
+            weight_indices=torch.tensor(
+                seg_weight_indices_cpu, dtype=torch.int32, device=self.device
+            ),
+            expected_tokens=expected_tokens,
+        )

python/sglang/srt/lora/deepseek_mla_correction.py

@@ -0,0 +1,117 @@
+"""LoRA correction for absorbed-MLA ``kv_b_proj``.
+
+The absorbed-MLA path in ``DeepseekV2AttentionMLA`` bypasses
+``kv_b_proj.forward()`` and folds the K/V contribution into two BMMs against
+the pre-computed ``w_kc`` / ``w_vc`` weights, so a standard
+``ColumnParallelLinearWithLoRA`` wrapper would never see the activations and
+the LoRA delta would silently be dropped. These helpers inject the missing
+delta on top of the absorbed intermediates via the SGMM-style Triton kernels
+in ``triton_ops/kv_b_lora_absorbed.py``.
+
+Used from ``deepseek_common/attention_forward_methods/forward_mla.py``. Call
+sites should gate the call with :func:`is_kv_b_lora_active` so non-LoRA
+forwards take a single ``getattr`` and skip the helper entirely.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional, Tuple
+
+import torch
+
+from sglang.kernels.ops.gemm.kv_b_lora_absorbed import (
+    step_a_q_fwd,
+    step_a_v_fwd,
+    step_b_q_fwd,
+    step_b_v_fwd,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.lora.utils import LoRABatchInfo
+    from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
+
+
+def is_kv_b_lora_active(attn_module: DeepseekV2AttentionMLA) -> bool:
+    """Cheap precondition check used at call sites in the attention forward
+    to skip the entire LoRA-correction path when no ``kv_b_proj`` adapter is
+    wrapped on this module (the common case)."""
+    return getattr(attn_module.kv_b_proj, "set_lora", False)
+
+
+def _get_state(
+    attn_module: DeepseekV2AttentionMLA,
+) -> Optional[Tuple[torch.Tensor, torch.Tensor, LoRABatchInfo]]:
+    if not is_kv_b_lora_active(attn_module):
+        return None
+    if not hasattr(attn_module.kv_b_proj, "A_buffer"):
+        return None
+    lora_backend = attn_module.kv_b_proj.lora_backend
+    if not hasattr(lora_backend, "batch_info"):
+        return None
+    batch_info = lora_backend.batch_info
+    if batch_info is None:
+        return None
+
+    # Triton backend exposes _sgemm_info() to group decode-shape repeats of
+    # the same adapter; csgmv-style backends just expose batch_info directly.
+    sgemm_info = getattr(lora_backend, "_sgemm_info", None)
+    if callable(sgemm_info):
+        batch_info = sgemm_info()
+    return attn_module.kv_b_proj.A_buffer, attn_module.kv_b_proj.B_buffer, batch_info
+
+
+def apply_q_correction(
+    attn_module: DeepseekV2AttentionMLA,
+    q_nope: torch.Tensor,
+    q_nope_out: torch.Tensor,
+) -> torch.Tensor:
+    """LoRA correction for the absorbed ``q_nope @ w_kc`` path.
+
+    Computes ``q_nope_out += q_nope @ B_kc @ A * scaling`` per token, per
+    active LoRA slot via two SGMM-style Triton kernels. Factored along the
+    LoRA-A/B boundary so we never materialise ``B @ A`` (~268M FMAs per layer
+    per slot in the naive implementation)::
+
+      step A_q : ``(S,H,qk_nope) @ B_kc[slot, h] (qk_nope, rank) -> (S,H,rank)``
+      step B_q : ``(S,H,rank)    @ A[slot] (rank, kv_lora_rank)  -> += q_nope_out``
+    """
+    state = _get_state(attn_module)
+    if state is None:
+        return q_nope_out
+    A_buf, B_buf, batch_info = state
+
+    full_K_per_head = attn_module.qk_nope_head_dim + attn_module.v_head_dim
+    q_lora_a = step_a_q_fwd(q_nope, B_buf, batch_info, full_K_per_head)
+    return step_b_q_fwd(q_lora_a, A_buf, batch_info, q_nope_out)
+
+
+def apply_v_correction(
+    attn_module: DeepseekV2AttentionMLA,
+    attn_output: torch.Tensor,
+    attn_bmm_flat: torch.Tensor,
+) -> torch.Tensor:
+    """LoRA correction for the absorbed ``attn_output @ w_vc`` path.
+
+    Computes ``attn_bmm_flat += attn_output @ A.T @ B_vc.T * scaling`` per
+    token, per active LoRA slot. ``attn_bmm_flat`` is the flat
+    ``(S, H*v_head_dim)`` view of the absorbed BMM result; we pass strides
+    matching the implicit ``(S, H, v_head_dim)`` layout to step B_v.
+    """
+    state = _get_state(attn_module)
+    if state is None:
+        return attn_bmm_flat
+    A_buf, B_buf, batch_info = state
+
+    attn_lora_a = step_a_v_fwd(attn_output, A_buf, batch_info)
+    base_view = attn_bmm_flat.view(
+        -1, attn_module.num_local_heads, attn_module.v_head_dim
+    )
+    step_b_v_fwd(
+        attn_lora_a,
+        B_buf,
+        batch_info,
+        base_view,
+        attn_module.qk_nope_head_dim,
+        attn_module.v_head_dim,
+    )
+    return attn_bmm_flat

python/sglang/srt/lora/eviction_policy.py

@@ -0,0 +1,139 @@
+# Copyright 2023-2024 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.
+# ==============================================================================
+
+"""
+Eviction policies for LoRA adapter memory management.
+"""
+
+import logging
+import time
+from abc import ABC, abstractmethod
+from collections import OrderedDict
+from typing import Optional, Set
+
+logger = logging.getLogger(__name__)
+
+
+class EvictionPolicy(ABC):
+    """Abstract base class for LoRA adapter eviction policies."""
+
+    @abstractmethod
+    def mark_used(self, uid: Optional[str]) -> None:
+        """Marks an adapter as used."""
+        pass
+
+    @abstractmethod
+    def select_victim(self, candidates: Set[Optional[str]]) -> Optional[str]:
+        """Selects an adapter to evict from candidates."""
+        pass
+
+    @abstractmethod
+    def remove(self, uid: Optional[str]) -> None:
+        """Removes an adapter from the policy's tracking."""
+        pass
+
+
+class LRUEvictionPolicy(EvictionPolicy):
+    """LRU eviction policy - evicts the least recently used adapter."""
+
+    def __init__(self):
+        self.access_order = OrderedDict()  # key=uid, value=last_access_time
+        self.total_accesses = 0
+        self.eviction_count = 0
+
+    def mark_used(self, uid: Optional[str]) -> None:
+        if uid is not None:
+            current_time = time.monotonic()
+            # Remove and re-add to move to end (most recent)
+            self.access_order.pop(uid, None)
+            self.access_order[uid] = current_time
+            self.total_accesses += 1
+            logger.debug(f"LoRA {uid} marked as used at {current_time}")
+
+    def select_victim(self, candidates: Set[Optional[str]]) -> Optional[str]:
+        """Select the least recently used adapter from candidates."""
+        # Iterate through access_order (oldest first) to find LRU victim
+        for uid in list(self.access_order.keys()):
+            if uid in candidates:
+                logger.debug(f"Selected LoRA {uid} for eviction (LRU)")
+                self.eviction_count += 1
+                return uid
+
+        # If no tracked UID found in candidates, check if None is available
+        # This happens when the batch consists entirely of LoRA requests
+        # and None (base model) is the only eviction candidate
+        if None in candidates:
+            logger.debug("Selected None (base model) for eviction")
+            self.eviction_count += 1
+            return None
+
+        # Should never reach here if candidates is non-empty
+        assert False, f"Failed to select LRU victim from candidates: {candidates}"
+
+    def remove(self, uid: Optional[str]) -> None:
+        if uid is not None:
+            self.access_order.pop(uid, None)
+            logger.debug(f"Removed LoRA {uid} from LRU tracking")
+
+
+class FIFOEvictionPolicy(EvictionPolicy):
+    """FIFO eviction policy - for backward compatibility."""
+
+    def __init__(self):
+        self.insertion_order = (
+            OrderedDict()
+        )  # key=uid, OrderedDict maintains insertion order
+        self.eviction_count = 0
+
+    def mark_used(self, uid: Optional[str]) -> None:
+        """For FIFO, we only track insertion order (not access time)."""
+        if uid is not None and uid not in self.insertion_order:
+            self.insertion_order[uid] = (
+                True  # Value unused, OrderedDict tracks insertion order
+            )
+
+    def select_victim(self, candidates: Set[Optional[str]]) -> Optional[str]:
+        """Select the first inserted adapter from candidates."""
+        # Iterate through insertion_order (oldest first) to find FIFO victim
+        for uid in list(self.insertion_order.keys()):
+            if uid in candidates:
+                logger.debug(f"Selected LoRA {uid} for eviction (FIFO)")
+                self.eviction_count += 1
+                return uid
+
+        # If no tracked UID found in candidates, check if None is available
+        # This happens when the batch consists entirely of LoRA requests
+        # and None (base model) is the only eviction candidate
+        if None in candidates:
+            logger.debug("Selected None (base model) for eviction")
+            self.eviction_count += 1
+            return None
+
+        # Should never reach here if candidates is non-empty
+        assert False, f"Failed to select FIFO victim from candidates: {candidates}"
+
+    def remove(self, uid: Optional[str]) -> None:
+        if uid is not None:
+            self.insertion_order.pop(uid, None)
+
+
+def get_eviction_policy(policy_name: str) -> EvictionPolicy:
+    """Factory function to create eviction policy instances."""
+    policies = {
+        "fifo": FIFOEvictionPolicy,
+        "lru": LRUEvictionPolicy,
+    }
+    if policy_name not in policies:
+        raise ValueError(f"Unknown eviction policy: {policy_name}")
+    return policies[policy_name]()

python/sglang/srt/lora/lora.py

@@ -0,0 +1,436 @@
+# Copyright 2023-2024 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.
+# ==============================================================================
+
+# Integrates "S-LoRA: Serving Thousands of Concurrent LoRA Adapters"
+# and "Punica: Multi-Tenant LoRA Serving"
+
+# LoRA layers class inheritance adapted from:
+# https://github.com/vllm-project/vllm/blob/4abf6336ec65c270343eb895e7b18786e9274176/vllm/lora/layers.py
+
+import logging
+import re
+from typing import Dict, List, Optional
+
+import torch
+from torch import nn
+
+from sglang.srt.configs.load_config import LoadConfig
+from sglang.srt.layers.utils import get_layer_id
+from sglang.srt.lora.backend.base_backend import BaseLoRABackend
+from sglang.srt.lora.lora_config import LoRAConfig
+from sglang.srt.model_loader.loader import DefaultModelLoader
+from sglang.srt.utils.hf_transformers_utils import AutoConfig
+
+# Matches both per-expert keys ("...experts.0.<module>...") and shared-outer
+# keys ("...experts.<module>..."), while excluding "shared_experts." (where the
+# preceding char is "_", not ".").
+_ROUTED_EXPERT_PATTERN = re.compile(r"\.experts\.")
+
+logger = logging.getLogger(__name__)
+
+
+class LoRALayer(nn.Module):
+    def __init__(self, config: LoRAConfig, base_hf_config: AutoConfig):
+        super().__init__()
+        self.config: LoRAConfig = config
+        self.base_hf_config: AutoConfig = base_hf_config
+
+        # lora weights in cpu. The weights are loaded from checkpoint.
+        self.weights: Dict[str, torch.Tensor] = {}
+        self.pinned_weights: Dict[str, torch.Tensor] = {}
+
+
+class LoRAAdapter(nn.Module):
+    def __init__(
+        self,
+        uid: str,
+        config: LoRAConfig,
+        base_hf_config: AutoConfig,
+        load_config: LoadConfig,
+        lora_backend: BaseLoRABackend,
+        base_model: Optional[torch.nn.Module] = None,
+    ):
+        super().__init__()
+        self.uid: str = uid
+        self.config: LoRAConfig = config
+        assert self.config.hf_config["peft_type"].lower() == "lora"
+        self.base_hf_config: AutoConfig = base_hf_config
+        self.load_config: LoadConfig = load_config
+        self.lora_backend: BaseLoRABackend = lora_backend
+        self.scaling: float = self.config.lora_alpha / self.config.r
+
+        # Bypass nn.Module.__setattr__ so the base model is held as a plain
+        # reference rather than auto-registered as a submodule (which would
+        # leak its parameters into our state_dict / parameters() / .to()).
+        object.__setattr__(self, "base_model", base_model)
+        object.__setattr__(
+            self,
+            "_moe_is_gated_by_layer",
+            self._build_moe_gated_map(base_model) if base_model is not None else {},
+        )
+
+        self.layers: List[LoRALayer] = nn.ModuleList(
+            [
+                LoRALayer(config, base_hf_config)
+                for _ in range(base_hf_config.num_hidden_layers)
+            ]
+        )
+
+        self.embedding_layers: Dict[str, torch.Tensor] = {}
+        self.pinned_embedding_layers: Dict[str, torch.Tensor] = {}
+        self.added_tokens_embeddings: Dict[str, torch.Tensor] = {}
+        self.pinned_added_tokens_embeddings: Dict[str, torch.Tensor] = {}
+
+    @staticmethod
+    def _build_moe_gated_map(base_model: torch.nn.Module) -> Dict[int, bool]:
+        """Map layer_id -> moe_runner_config.is_gated for FusedMoE base layers.
+
+        Only used by normalize_gate_up_proj to decide whether per-expert
+        gate_proj weights should be zero-padded and stacked (gated → c=2 buffer)
+        or just renamed (non-gated → c=1 buffer via model's get_stacked_multiply
+        override on gate_up_proj_moe).
+
+        Adapters can be loaded both before `init_lora_modules` (initial
+        --lora-paths) and after (dynamic API loads), so the FusedMoE may
+        appear either directly or under a `BaseLayerWithLoRA.base_layer`.
+        """
+        from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
+
+        gated_map: Dict[int, bool] = {}
+        for name, module in base_model.named_modules():
+            inner = (
+                module
+                if isinstance(module, FusedMoE)
+                else getattr(module, "base_layer", None)
+            )
+            if not isinstance(inner, FusedMoE):
+                continue
+            layer_id = get_layer_id(name)
+            if layer_id is not None:
+                gated_map[layer_id] = bool(inner.moe_runner_config.is_gated)
+        return gated_map
+
+    def _is_non_gated_moe_weight(self, weight_name: str) -> bool:
+        """True iff this adapter weight targets a non-gated MoE expert.
+
+        Such weights flow into the `gate_up_proj_moe` buffer, which the model
+        overrides to stacked_multiply=1 — so the weight must be stored without
+        being stacked with a synthetic up_proj zero-pad.
+
+        Matches both adapter key conventions:
+        - per-expert: ``...experts.0.<module>...`` (one tensor per expert)
+        - shared-outer: ``...experts.<module>...`` (3D tensor with the expert
+          dim baked into the shape)
+        """
+        if not _ROUTED_EXPERT_PATTERN.search(weight_name):
+            return False
+        layer_id = get_layer_id(weight_name)
+        if layer_id is None:
+            return False
+        return self._moe_is_gated_by_layer.get(layer_id) is False
+
+    def initialize_weights(self):
+        model_path = self.config.path
+        loader = DefaultModelLoader(self.load_config)
+        revision = getattr(self.config.hf_config, "revision", None)
+
+        # Get normalized target modules for filtering
+        for name, loaded_weight in loader._get_weights_iterator(
+            DefaultModelLoader.Source(
+                model_path, revision=revision, fall_back_to_pt=True
+            )
+        ):
+            self._process_weight(name, loaded_weight)
+
+        self._normalize_weights()
+
+    def initialize_weights_from_tensors(self, tensors: Dict[str, torch.Tensor]):
+        for name, tensor in tensors.items():
+            self._process_weight(name, tensor)
+
+        self._normalize_weights()
+
+    def _process_weight(self, name: str, loaded_weight: torch.Tensor):
+        from sglang.srt.lora.utils import get_normalized_target_modules
+
+        normalized_target_modules = get_normalized_target_modules(
+            self.config.target_modules
+        )
+
+        # Remap PEFT "unembed_tokens" key to "lm_head" so the weight is
+        # recognized and loaded into the correct buffer.
+        if "unembed_tokens" in name:
+            name = name.replace("unembed_tokens", "lm_head")
+
+        layer_id = get_layer_id(name)
+        if layer_id is not None:
+            self.layers[layer_id].weights[name] = loaded_weight.cpu()
+        elif "embed_tokens" in name or "lm_head" in name:
+            # Check if this module is declared in target_modules before loading.
+            # When normalized_target_modules is {"all"} (e.g. target_modules was
+            # "all-linear"), we allow loading since the server-level
+            # --lora-target-modules will govern which modules are active.
+            module_name = "embed_tokens" if "embed_tokens" in name else "lm_head"
+            if (
+                "all" in normalized_target_modules
+                or module_name in normalized_target_modules
+            ):
+                self.embedding_layers[name] = loaded_weight.cpu()
+            else:
+                logger.debug(
+                    f"Skipping {name} as '{module_name}' is not in adapter's target_modules: {self.config.target_modules}"
+                )
+        elif "input_embeddings" in name or "output_embeddings" in name:
+            # added/extra token emb
+            self.added_tokens_embeddings[name] = loaded_weight.cpu()
+            assert loaded_weight.shape[0] == self.config.lora_added_tokens_size, (
+                f"LoRA adapter {self.uid} has lora_added_tokens_size {self.config.lora_added_tokens_size} specified in the config, "
+                f"but the loaded weight '{name}' has shape {loaded_weight.shape[0]} in first dimension"
+            )
+
+    def _normalize_weights(self):
+        for layer in self.layers:
+            weight_names = list(layer.weights.keys())
+            self.normalize_qkv_proj(weight_names, layer.weights)
+            self._rename_expert_w_to_proj(layer.weights)
+            # Stack gate_proj + x_proj → in_proj for Mamba layers (before gate_up normalization)
+            self._normalize_in_proj(layer.weights)
+            # Stack in_proj_q + in_proj_k + in_proj_v + in_proj_z → in_proj_qkvz for GDN layers
+            self._normalize_in_proj_qkvz(layer.weights)
+            weight_names = list(layer.weights.keys())
+            self.normalize_gate_up_proj(weight_names, layer.weights)
+            weight_names = list(layer.weights.keys())
+            self.normalize_fused_qkv_a_proj(weight_names, layer.weights)
+
+    def normalize_qkv_proj(
+        self, weight_names: List[str], weights: Dict[str, torch.Tensor]
+    ):
+        # Collect target q/k/v modules. This process is necessary since there might be no lora attached to k_proj
+        target_module = set()
+        for weight_name in weight_names:
+            if "k_proj" in weight_name:
+                target_module.add("k_proj")
+            if "q_proj" in weight_name:
+                target_module.add("q_proj")
+            if "v_proj" in weight_name:
+                target_module.add("v_proj")
+            if "qkv_proj" in weight_name:
+                target_module.add("qkv_proj")
+        if len(target_module) == 0:
+            return
+
+        for weight_name in weight_names:
+            # We assume every lora adaptor should contain lora modules for q_proj
+            if "q_proj" in weight_name:
+                q_name = weight_name
+                k_name = weight_name.replace("q_proj", "k_proj")
+                v_name = weight_name.replace("q_proj", "v_proj")
+                qkv_name = weight_name.replace("q_proj", "qkv_proj")
+
+                # If k_proj doesn't have lora, initialize it to zero
+                k_proj_weight = (
+                    weights[k_name]
+                    if "k_proj" in target_module
+                    else torch.zeros_like(weights[v_name])
+                )
+                weights[qkv_name] = torch.cat(
+                    (
+                        weights[q_name],
+                        k_proj_weight,
+                        weights[v_name],
+                    ),
+                    0,
+                )
+                weights.pop(q_name)
+                if "k_proj" in target_module:
+                    weights.pop(k_name)
+                weights.pop(v_name)
+            elif "qkv_proj" in weight_name:
+                # If qkv_proj is already stacked, we normalize it following the SGL convention.
+                qkv_name = weight_name
+                q_name = weight_name.replace("qkv_proj", "q_proj")
+                k_name = weight_name.replace("qkv_proj", "k_proj")
+                v_name = weight_name.replace("qkv_proj", "v_proj")
+                if "lora_A" in weight_name:
+                    weights[qkv_name] = weights[qkv_name].repeat(3, 1)
+                # else: no-op as LoRA B weight is already stacked.
+
+    def _rename_expert_w_to_proj(self, weights: Dict[str, torch.Tensor]):
+        """Rename w1 -> gate_proj, w3 -> up_proj, w2 -> down_proj so that
+        normalize_gate_up_proj can stack them into gate_up_proj."""
+        renames = {}
+        for name in list(weights.keys()):
+            new_name = name
+            if ".w1." in name:
+                new_name = name.replace(".w1.", ".gate_proj.")
+            elif ".w3." in name:
+                new_name = name.replace(".w3.", ".up_proj.")
+            elif ".w2." in name:
+                new_name = name.replace(".w2.", ".down_proj.")
+            if new_name != name:
+                renames[name] = new_name
+        for old_name, new_name in renames.items():
+            weights[new_name] = weights.pop(old_name)
+
+    def _normalize_in_proj(self, weights: Dict[str, torch.Tensor]):
+        """Stack gate_proj + x_proj → in_proj for Mamba layers.
+
+        Detects Mamba layers by the presence of both gate_proj and x_proj.
+        Must run BEFORE normalize_gate_up_proj to prevent gate_proj from
+        being consumed by the gate+up stacking.
+        """
+        # Find gate_proj weights that have a matching x_proj (Mamba pattern)
+        for weight_name in list(weights.keys()):
+            if "gate_proj" not in weight_name:
+                continue
+            x_name = weight_name.replace("gate_proj", "x_proj")
+            if x_name not in weights:
+                continue
+            # This is a Mamba layer: stack gate_proj + x_proj → in_proj
+            in_proj_name = weight_name.replace("gate_proj", "in_proj")
+            cat_dim = weights[weight_name].dim() - 2
+            weights[in_proj_name] = torch.cat(
+                (weights[weight_name], weights[x_name]), cat_dim
+            )
+            weights.pop(weight_name)
+            weights.pop(x_name)
+
+    def _normalize_in_proj_qkvz(self, weights: Dict[str, torch.Tensor]):
+        """Normalize in_proj_qkvz weights for GDN (GatedDeltaNet) layers like
+        Qwen3.5.
+
+        Two adapter formats are handled:
+
+        1. Split: ``in_proj_q + in_proj_k + in_proj_v + in_proj_z`` are present
+           as separate weights → concatenate them into ``in_proj_qkvz``.
+
+        2. Already-merged: the adapter has a single ``in_proj_qkvz`` weight
+           (PEFT trained against SGLang's fused Linear). The stacked buffer
+           expects four per-slice ``A`` blocks, so repeat ``lora_A`` 4× along
+           the rank dim. ``lora_B`` is already full-output-dim and matches
+           the buffer directly.
+        """
+        for weight_name in list(weights.keys()):
+            if "in_proj_q." in weight_name:
+                k_name = weight_name.replace("in_proj_q", "in_proj_k")
+                v_name = weight_name.replace("in_proj_q", "in_proj_v")
+                z_name = weight_name.replace("in_proj_q", "in_proj_z")
+                if (
+                    k_name not in weights
+                    or v_name not in weights
+                    or z_name not in weights
+                ):
+                    continue
+                qkvz_name = weight_name.replace("in_proj_q", "in_proj_qkvz")
+                cat_dim = weights[weight_name].dim() - 2
+                weights[qkvz_name] = torch.cat(
+                    (
+                        weights[weight_name],
+                        weights[k_name],
+                        weights[v_name],
+                        weights[z_name],
+                    ),
+                    cat_dim,
+                )
+                weights.pop(weight_name)
+                weights.pop(k_name)
+                weights.pop(v_name)
+                weights.pop(z_name)
+            elif "in_proj_qkvz" in weight_name and "lora_A" in weight_name:
+                # Already-merged adapter: replicate the shared A across the 4
+                # stacked slots the buffer expects (q, k, v, z).
+                ndim = weights[weight_name].dim()
+                repeat_dims = [1] * ndim
+                repeat_dims[ndim - 2] = 4
+                weights[weight_name] = weights[weight_name].repeat(*repeat_dims)
+            # else (in_proj_qkvz lora_B, or unrelated): no-op.
+
+    def normalize_gate_up_proj(
+        self, weight_names: List[str], weights: Dict[str, torch.Tensor]
+    ):
+        for weight_name in weight_names:
+            if "gate_proj" in weight_name:
+                up_name = weight_name.replace("gate_proj", "up_proj")
+                gate_up_name = weight_name.replace("gate_proj", "gate_up_proj")
+                # PEFT can ship up_proj in two forms when there's no real
+                # up_proj content: the key may be absent, or present as a
+                # numel-zero placeholder. Treat both as "no up_proj".
+                if up_name not in weights or weights[up_name].numel() == 0:
+                    if self._is_non_gated_moe_weight(weight_name):
+                        # Non-gated MoE expert: the gate_up_proj_moe buffer
+                        # uses stacked_multiply=1 (per model override), so just
+                        # rename without stacking.
+                        weights[gate_up_name] = weights.pop(weight_name)
+                        if up_name in weights:
+                            weights.pop(up_name)
+                        continue
+                    # Gated path: buffer expects stacked [2r, hidden] (c=2);
+                    # synthesize a properly-shaped zero up_proj.
+                    weights[up_name] = torch.zeros_like(weights[weight_name])
+                cat_dim = weights[weight_name].dim() - 2
+                weights[gate_up_name] = torch.cat(
+                    (weights[weight_name], weights[up_name]), cat_dim
+                )
+                weights.pop(weight_name)
+                weights.pop(up_name)
+            elif "gate_up_proj" in weight_name:
+                # If gate_up_proj is already stacked, we normalize it following the SGL convention
+                gate_up_name = weight_name
+                if "lora_A" in weight_name:
+                    ndim = weights[gate_up_name].dim()
+                    repeat_dims = [1] * ndim
+                    repeat_dims[ndim - 2] = 2
+                    weights[gate_up_name] = weights[gate_up_name].repeat(*repeat_dims)
+                # else: no-op as LoRA B weight is already stacked.
+        # Orphan up_proj weights (no matching gate_proj) are kept as-is.
+        # Models with non-gated MLP/shared-experts declare up_proj in
+        # supported_lora_modules so they get their own buffer and wrapping.
+
+    def normalize_fused_qkv_a_proj(
+        self, weight_names: List[str], weights: Dict[str, torch.Tensor]
+    ):
+        """Fuse separate q_a_proj and kv_a_proj_with_mqa LoRA weights into
+        a single fused_qkv_a_proj_with_mqa entry (concat along dim 0 for
+        both A and B), matching the DeepSeek MLA fused projection layout."""
+        for weight_name in weight_names:
+            if "q_a_proj" not in weight_name:
+                continue
+            if "fused_qkv_a_proj_with_mqa" in weight_name:
+                continue
+
+            q_a_name = weight_name
+            kv_a_name = weight_name.replace("q_a_proj", "kv_a_proj_with_mqa")
+            fused_name = weight_name.replace("q_a_proj", "fused_qkv_a_proj_with_mqa")
+
+            kv_a_weight = (
+                weights[kv_a_name]
+                if kv_a_name in weights
+                else torch.zeros_like(weights[q_a_name])
+            )
+
+            weights[fused_name] = torch.cat((weights[q_a_name], kv_a_weight), dim=0)
+            weights.pop(q_a_name)
+            if kv_a_name in weights:
+                weights.pop(kv_a_name)
+
+    def pin_weights_in_cpu(self):
+        for layer in self.layers:
+            for name, weight in layer.weights.items():
+                layer.weights[name] = weight.pin_memory()
+
+        for name, weight in self.embedding_layers.items():
+            self.embedding_layers[name] = weight.pin_memory()
+
+        for name, weight in self.added_tokens_embeddings.items():
+            self.added_tokens_embeddings[name] = weight.pin_memory()

python/sglang/srt/lora/lora_config.py

@@ -0,0 +1,108 @@
+# Copyright 2023-2024 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.
+# ==============================================================================
+
+import json
+import logging
+import os
+from typing import Dict, Optional
+
+from huggingface_hub import snapshot_download
+
+logger = logging.getLogger(__name__)
+
+
+class LoRAConfig:
+    def __init__(
+        self,
+        path: Optional[str] = None,
+        config_dict: Optional[Dict] = None,
+        added_tokens_config: Optional[Dict] = None,
+        base_vocab_size: Optional[int] = None,
+    ) -> None:
+        self.path = path
+
+        if config_dict is not None:
+            self.hf_config = config_dict
+            self.added_tokens_config = added_tokens_config
+        else:
+            self.hf_config = self.get_lora_config()
+            self.added_tokens_config = self.get_added_tokens_config()
+
+        self.target_modules = self.hf_config["target_modules"]
+        self.r = self.hf_config["r"]
+        self.lora_alpha = self.hf_config["lora_alpha"]
+        self.use_dora = self.hf_config.get("use_dora", False)
+
+        # Filter fake added tokens: tokens with ID < base_vocab_size are already
+        # part of the base vocabulary and should not be treated as added tokens.
+        # This commonly happens when added_tokens.json is copied from the base
+        # model's tokenizer.
+        if self.added_tokens_config and base_vocab_size is not None:
+            self.added_tokens_config = {
+                token: token_id
+                for token, token_id in self.added_tokens_config.items()
+                if token_id >= base_vocab_size
+            }
+
+        self.lora_added_tokens_size = (
+            len(self.added_tokens_config) if self.added_tokens_config is not None else 0
+        )
+
+    @classmethod
+    def from_dict(
+        cls,
+        config_dict: Dict,
+        added_tokens_config: Optional[Dict] = None,
+        base_vocab_size: Optional[int] = None,
+    ) -> "LoRAConfig":
+        return cls(
+            config_dict=config_dict,
+            added_tokens_config=added_tokens_config,
+            base_vocab_size=base_vocab_size,
+        )
+
+    def get_lora_config(self, dummy=False):
+        if dummy:
+            raise NotImplementedError()
+        else:
+            if not os.path.isdir(self.path):
+                weights_dir = snapshot_download(self.path, allow_patterns=["*.json"])
+            else:
+                weights_dir = self.path
+            config_name = "adapter_config.json"
+            with open(os.path.join(weights_dir, config_name), "r") as f:
+                return json.load(f)
+
+    def get_added_tokens_config(self):
+        """Load added tokens from the LoRA adapter if the file exists."""
+        # Determine the weights directory
+        if not os.path.isdir(self.path):
+            weights_dir = snapshot_download(self.path, allow_patterns=["*.json"])
+        else:
+            weights_dir = self.path
+
+        # Construct the path to added_tokens.json
+        added_tokens_path = os.path.join(weights_dir, "added_tokens.json")
+
+        # Return None if the file doesn't exist (optional for standard LoRA adapters)
+        if not os.path.exists(added_tokens_path):
+            return None
+
+        # Load and return the added tokens
+        try:
+            with open(added_tokens_path, "r") as f:
+                return json.load(f)
+        except json.JSONDecodeError as e:
+            logger.warning(f"Failed to parse added_tokens.json: {e}")
+            return None

python/sglang/srt/lora/lora_drainer.py

@@ -0,0 +1,191 @@
+# Copyright 2023-2024 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.
+# ==============================================================================
+
+import logging
+import time
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+from sglang.srt.managers.schedule_batch import Req
+
+logger = logging.getLogger(__name__)
+
+DRAIN_SCHEDULE_TOLERANCE = 1.2
+
+
+@dataclass
+class AdapterStats:
+    num_waiting_reqs: int = 0
+    max_wait_time_secs: float = 0.0
+    max_remaining_tokens: int = 0
+    is_draining_for: Optional[str] = None
+
+    def _reset_stats(self):
+        self.num_waiting_reqs = 0
+        self.max_wait_time_secs = 0.0
+        self.max_remaining_tokens = 0
+
+    def is_starving(self, drain_wait_threshold: float):
+        return (
+            self.max_wait_time_secs > drain_wait_threshold and self.num_waiting_reqs > 0
+        )
+
+
+class LoRADrainer:
+    """
+    Drainer for LoRA requests that manages draining. It tracks:
+    - Number of waiting requests per adapter
+    - Maximum wait time for requests needing each adapter
+    - Maximum number of tokens needed for running requests for each adapter
+    """
+
+    def __init__(self, max_loras_per_batch: int, max_wait_time_secs: float = 0.0):
+        self.max_loras_per_batch = max_loras_per_batch
+        self.max_wait_time_secs = max_wait_time_secs
+        self.adapter_to_stats: Dict[Optional[str], AdapterStats] = defaultdict(
+            AdapterStats
+        )
+
+    def update_draining_state(
+        self,
+        waiting_queue: List[Req],
+        running_reqs: List[Req],
+    ) -> None:
+        """
+        Update LoRA drainer state based on current waiting queue and running requests.
+
+        This method updates adapter statistics, identifies starving adapters that need
+        to be scheduled, and marks adapters for draining to make room for starving ones.
+        """
+        self._update_adapter_stats(waiting_queue, running_reqs)
+        self._update_draining_loras(running_reqs)
+        self._update_fully_drained_loras(running_reqs)
+
+    def _update_adapter_stats(
+        self,
+        waiting_queue: List[Req],
+        running_reqs: List[Req],
+    ) -> None:
+        for stats in self.adapter_to_stats.values():
+            stats._reset_stats()
+
+        for req in waiting_queue:
+            stats = self.adapter_to_stats[req.lora_id]
+
+            stats.num_waiting_reqs += 1
+            stats.max_wait_time_secs = max(
+                stats.max_wait_time_secs,
+                time.monotonic() - req.time_stats.wait_queue_entry_time,
+            )
+
+        for req in running_reqs:
+            stats = self.adapter_to_stats[req.lora_id]
+
+            stats.max_remaining_tokens = max(
+                stats.max_remaining_tokens,
+                req.sampling_params.max_new_tokens - len(req.output_ids),
+            )
+
+    def _update_draining_loras(self, running_reqs: List[Req]) -> None:
+        """
+        Select LoRA adapters to drain based on starvation detection.
+
+        This method identifies adapters in the waiting queue that are "starving"
+        (waiting too long) and marks currently running adapters as "draining"
+        to make room for the starving adapters. Draining adapters will not
+        accept new requests, allowing them to complete and free up slots.
+        """
+        running_adapter_ids = {req.lora_id for req in running_reqs}
+        if len(running_adapter_ids) < self.max_loras_per_batch:
+            return None
+
+        starving_adapters = set()
+        draining_for_adapters = set()
+        for adapter_id, stats in self.adapter_to_stats.items():
+            if stats.is_starving(self.max_wait_time_secs):
+                starving_adapters.add(adapter_id)
+
+            draining_for_adapter = stats.is_draining_for
+            if draining_for_adapter is not None:
+                draining_for_adapters.add(draining_for_adapter)
+
+        new_starving_adapters = starving_adapters - draining_for_adapters
+        if not new_starving_adapters:
+            return None
+
+        sorted_new_starving_adapters = sorted(
+            new_starving_adapters,
+            key=lambda adapter: self.adapter_to_stats[adapter].max_wait_time_secs,
+            reverse=True,
+        )
+
+        eligible_to_drain_adapters = {
+            adapter
+            for adapter in running_adapter_ids
+            if self.adapter_to_stats[adapter].is_draining_for is None
+        }
+
+        for starving_adapter in sorted_new_starving_adapters:
+            if not eligible_to_drain_adapters:
+                break
+
+            min_eligible_adapter = min(
+                eligible_to_drain_adapters,
+                key=lambda adapter_id: self.adapter_to_stats[
+                    adapter_id
+                ].max_remaining_tokens,
+            )
+
+            self.adapter_to_stats[min_eligible_adapter].is_draining_for = (
+                starving_adapter
+            )
+            logger.debug(
+                f"LoRA adapter {min_eligible_adapter} is draining for {starving_adapter}"
+            )
+
+            eligible_to_drain_adapters.remove(min_eligible_adapter)
+
+    def _update_fully_drained_loras(self, running_reqs: List[Req]) -> None:
+        """
+        Clear draining state for adapters that have fully drained.
+
+        An adapter is considered fully drained when it was marked as draining
+        but no longer has any running requests.
+        """
+        running_adapter_ids = {req.lora_id for req in running_reqs}
+        for adapter_id, stats in self.adapter_to_stats.items():
+            if stats.is_draining_for is None:
+                continue
+
+            if adapter_id not in running_adapter_ids:
+                logger.debug(f"LoRA adapter {adapter_id} finished draining")
+                stats.is_draining_for = None
+
+    def can_schedule(self, req: Req) -> bool:
+        """
+        Check if a request can be scheduled based on draining state.
+
+        If the adapter for this request is currently draining, only allow
+        scheduling if the request's max_new_tokens is within tolerance of
+        the max remaining tokens for the draining adapter.
+        """
+        stats = self.adapter_to_stats[req.lora_id]
+        if not stats.is_draining_for:
+            return True
+
+        return (
+            req.sampling_params.max_new_tokens
+            <= stats.max_remaining_tokens * DRAIN_SCHEDULE_TOLERANCE
+        )

python/sglang/srt/lora/lora_manager.py

@@ -0,0 +1,858 @@
+# Copyright 2023-2024 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.
+# ==============================================================================
+
+# Integrates "S-LoRA: Serving Thousands of Concurrent LoRA Adapters"
+# and "Punica: Multi-Tenant LoRA Serving"
+
+import logging
+from typing import Dict, Iterable, List, Optional
+
+import torch
+
+from sglang.srt.configs.load_config import LoadConfig
+from sglang.srt.environ import envs
+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,
+    VocabParallelEmbedding,
+)
+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, 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
+from sglang.srt.lora.mem_pool import LoRAMemoryPool
+from sglang.srt.lora.utils import (
+    DSA_INDEXER_LORA_NAMES,
+    EMBEDDING_NAMES,
+    LoRAType,
+    auto_detect_lora_target_modules,
+    get_normalized_target_modules,
+    get_target_module_name,
+)
+from sglang.srt.managers.io_struct import LoRAUpdateOutput
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils import replace_submodule
+from sglang.srt.utils.hf_transformers_utils import AutoConfig
+
+_SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get()
+
+logger = logging.getLogger(__name__)
+
+
+class LoRAManager:
+    def __init__(
+        self,
+        base_model: torch.nn.Module,
+        base_hf_config: AutoConfig,
+        max_loras_per_batch: int,
+        load_config: LoadConfig,
+        dtype: torch.dtype,
+        server_args: ServerArgs,
+        lora_backend: str = "triton",
+        tp_size: int = 1,
+        tp_rank: int = 0,
+        max_lora_rank: Optional[int] = None,
+        target_modules: Optional[Iterable[str]] = None,
+        lora_paths: Optional[List[LoRARef]] = None,
+    ):
+        self.base_model: torch.nn.Module = base_model
+        if hasattr(base_hf_config, "get_text_config"):
+            self.base_hf_config: AutoConfig = base_hf_config.get_text_config()
+        else:
+            self.base_hf_config: AutoConfig = base_hf_config
+        self.max_loras_per_batch: int = max_loras_per_batch
+        self.load_config: LoadConfig = load_config
+        self.dtype: torch.dtype = dtype
+        self.device: torch.device = next(self.base_model.parameters()).device
+        self.tp_size: int = tp_size
+        self.tp_rank: int = tp_rank
+        self.lora_added_tokens_size: Optional[int] = None
+        self.enable_lora_overlap_loading: Optional[bool] = (
+            server_args.enable_lora_overlap_loading
+        )
+
+        self.eviction_policy = server_args.lora_eviction_policy
+        self._experts_shared_outer_override: Optional[bool] = (
+            server_args.experts_shared_outer_loras
+        )
+        self.lora_use_virtual_experts: bool = server_args.lora_use_virtual_experts
+        self.lora_strict_loading: bool = getattr(
+            server_args, "lora_strict_loading", False
+        )
+
+        # LoRA backend for running sgemm kernels
+        logger.info(f"Using {lora_backend} as backend of LoRA kernels.")
+        backend_type = get_backend_from_name(lora_backend)
+        self.lora_backend: BaseLoRABackend = backend_type(
+            max_loras_per_batch=max_loras_per_batch,
+            device=self.device,
+            server_args=server_args,
+        )
+
+        # Initialize mutable internal state of the LoRAManager.
+        self.init_state(
+            max_lora_rank=max_lora_rank,
+            target_modules=target_modules,
+            lora_paths=lora_paths,
+        )
+
+    def init_cuda_graph_batch_info(
+        self, max_bs_in_cuda_graph: int, num_tokens_per_bs: int
+    ):
+        """Phase 2 of LoRA CUDA graph init: dense LoRA batch metadata.
+
+        Called during CudaGraphRunner.__init__(), after init_memory_pool().
+        Phase 1 (MoE buffers) is handled earlier via init_cuda_graph_moe_buffers().
+        """
+        self.max_bs_in_cuda_graph = max_bs_in_cuda_graph
+        self.lora_backend.init_cuda_graph_batch_info(
+            max_bs_in_cuda_graph=max_bs_in_cuda_graph,
+            num_tokens_per_bs=num_tokens_per_bs,
+        )
+
+        # ===== TO BE REFACTORED ====
+        # Pre-create the experimental LoRA two-stream side stream now (gated) so the
+        # torch.cuda.Stream() call never lands inside a cuda-graph capture region.
+        if _SGLANG_EXPERIMENTAL_LORA_OPTI:
+            from sglang.srt.lora.trtllm_lora_temp import (
+                init_lora_two_stream_resources,
+            )
+
+            init_lora_two_stream_resources(self.device)
+        # ===== END TO BE REFACTORED ====
+
+    def init_cuda_graph_moe_buffers(
+        self, max_bs: int, max_loras: int, compute_dtype, moe_layer
+    ):
+        """Phase 1 of LoRA CUDA graph init: MoE intermediate buffers.
+
+        Called before init_memory_pool() so memory profiling accounts for them.
+        Phase 2 (dense batch metadata) is handled later via init_cuda_graph_batch_info().
+        """
+        self.lora_backend.init_cuda_graph_moe_buffers(
+            max_bs=max_bs,
+            max_loras=max_loras,
+            compute_dtype=compute_dtype,
+            moe_layer=moe_layer,
+        )
+
+    def create_lora_update_result(
+        self, success: bool, error_message: str = ""
+    ) -> LoRAUpdateOutput:
+        return LoRAUpdateOutput(
+            success=success,
+            error_message=error_message,
+            loaded_adapters={
+                lora_ref.lora_name: lora_ref.lora_path
+                for lora_ref in self.lora_refs.values()
+            },
+        )
+
+    def load_lora_adapter(self, lora_ref: LoRARef) -> LoRAUpdateOutput:
+        """
+        Load a single LoRA adapter from the specified path.
+
+        Args:
+            lora_ref (LoRARef): The LoRARef object containing the LoRA name, path, and ID.
+        """
+        assert (
+            lora_ref.lora_name is not None and lora_ref.lora_path is not None
+        ), "LoRARef must have both lora_name and lora_path set for loading."
+        assert (
+            lora_ref.lora_id not in self.loras
+        ), f"LoRA adapter with ID {lora_ref.lora_id} is already loaded. This should have been verified before request is sent to the backend."
+
+        try:
+            # load configs
+            new_adapter = LoRAConfig(
+                lora_ref.lora_path,
+                base_vocab_size=self.base_hf_config.vocab_size,
+            )
+            self.validate_new_adapter(new_adapter, lora_ref)
+            self.configs[lora_ref.lora_id] = new_adapter
+
+            # load weights
+            self.load_lora_weights(lora_ref)
+
+            # keep metadata for displayed messages
+            self.lora_refs[lora_ref.lora_id] = lora_ref
+            self.num_pinned_loras += int(lora_ref.pinned)
+        except Exception as e:
+            return self.create_lora_update_result(
+                success=False,
+                error_message=str(e),
+            )
+
+        return self.create_lora_update_result(success=True)
+
+    def validate_new_adapter(self, lora_config: LoRAConfig, lora_ref: LoRARef):
+        """
+        Validate if an adapter can be loaded into the current LoRA memory pool and generate error if it is incompatible.
+        """
+        if lora_config.lora_added_tokens_size > 0:
+            raise ValueError(
+                f"Failed to load {lora_ref.lora_name} because LoRA serving currently doesn't support adapters that add tokens to the vocabulary"
+            )
+
+        if lora_config.use_dora:
+            raise ValueError(
+                f"Failed to load {lora_ref.lora_name} because LoRA serving currently doesn't support DoRA adapters"
+            )
+
+        # Check if this LoRA adapter is already loaded
+        for existing_lora_ref in self.lora_refs.values():
+            if lora_ref.lora_name == existing_lora_ref.lora_name:
+                raise ValueError(
+                    f"Failed to load LoRA adapter {lora_ref.lora_name} because it is already loaded"
+                )
+
+            if lora_ref.lora_path == existing_lora_ref.lora_path:
+                logger.warning(
+                    f"{lora_ref.lora_path} is already loaded with name: {existing_lora_ref.lora_name}, "
+                    f"but another copy is being loaded with name: {lora_ref.lora_name}"
+                )
+
+        # Check if the LoRA adapter shape is compatible with the current LoRA memory pool configuration.
+        memory_pool = getattr(self, "memory_pool", None)
+        incompatible = memory_pool and not memory_pool.can_support(lora_config)
+        if incompatible:
+            raise ValueError(
+                f"LoRA adapter {lora_ref.lora_name} with rank {lora_config.r} is incompatible with the current "
+                "LoRA memory pool configuration. Please ensure that the LoRA adapter's rank is within the configured "
+                "`--max-lora-rank` and that the target modules are included in `--lora-target-modules`."
+            )
+
+        # Ensure pinned LoRA adapters does not exceed maximal limit or cause starvation.
+        if lora_ref.pinned and self.num_pinned_loras >= self.max_loras_per_batch - 1:
+            raise ValueError(
+                f"Failed to load LoRA adapter {lora_ref.lora_name} as a pinned adapter. It is not allowed to pin all slots "
+                "in the LoRA memory pool to avoid starvation for unpinned adapters and base models. Please increase your "
+                "`--max-loras-per-batch` or load it as unpinned LoRA adapters."
+            )
+
+    def unload_lora_adapter(self, lora_ref: LoRARef) -> LoRAUpdateOutput:
+        """
+        Unload LoRA adapters by their names. This will remove the adapters from the memory pool and
+        delete the corresponding LoRA modules.
+        """
+
+        adapter = self.configs.get(lora_ref.lora_id)
+        lora_ref = self.lora_refs.get(lora_ref.lora_id)
+        assert (
+            adapter is not None and lora_ref is not None
+        ), f"LoRA adapter with ID {lora_ref.lora_id} is not loaded. This should have been verified before request is sent to the backend."
+
+        try:
+            del self.configs[lora_ref.lora_id]
+            del self.loras[lora_ref.lora_id]
+            del self.lora_refs[lora_ref.lora_id]
+            self.num_pinned_loras -= int(lora_ref.pinned)
+        except Exception as e:
+            return self.create_lora_update_result(
+                success=False,
+                error_message=str(e),
+            )
+
+        return self.create_lora_update_result(success=True)
+
+    def validate_lora_batch(self, lora_ids: set[Optional[str]]) -> bool:
+        """
+        Validate if the LoRA IDs in the batch can be loaded into the current LoRA memory pool.
+        """
+        if len(lora_ids) > self.max_loras_per_batch:
+            return False
+
+        # skip pinned LoRA check if no pinned LoRA adapters are loaded.
+        if self.num_pinned_loras == 0:
+            return True
+
+        # counting the number of pinned LoRA adapters in the batch.
+        pinned_loras_in_batch = 0
+        for lora_id in lora_ids:
+            if lora_id is not None:
+                lora_ref = self.lora_refs.get(lora_id)
+                assert (
+                    lora_ref is not None
+                ), f"LoRA ID {lora_id} not found in lora_refs."
+                pinned_loras_in_batch += int(lora_ref.pinned)
+
+        assert pinned_loras_in_batch <= self.num_pinned_loras, (
+            f"Number of pinned LoRA adapters in the batch ({pinned_loras_in_batch}) exceeds the total number of pinned adapters "
+            f"({self.num_pinned_loras}). This indicates a bug in the LoRA loading logic."
+        )
+
+        required_slots = len(lora_ids) - pinned_loras_in_batch
+        mem_pool_vacancy = self.memory_pool.max_loras_per_batch - self.num_pinned_loras
+
+        return required_slots <= mem_pool_vacancy
+
+    def fetch_new_loras(
+        self, new_loras: set[Optional[str]], running_loras: set[Optional[str]] = set()
+    ):
+        # Load active loras into lora memory pool
+        cur_uids = new_loras | running_loras
+
+        assert len(cur_uids) <= self.max_loras_per_batch
+        self.memory_pool.prepare_lora_batch(
+            cur_uids=cur_uids,
+            lora_adapters=self.loras,
+            lora_modules=self.lora_modules,
+            lora_refs=self.lora_refs.copy(),  # copy snapshot of current lora_refs to avoid mutation during the batch preparation.
+            lora_embed_tokens_module=self.embed_tokens_module,  # merge into embedding or lora module
+            lora_lm_head_module=self.lm_head_module,  # merge into embedding or lora module
+        )
+
+    def prepare_lora_batch(self, forward_batch: ForwardBatch):
+        # set up batch info shared by all lora modules
+        bs = forward_batch.batch_size
+
+        use_cuda_graph = (
+            hasattr(self, "max_bs_in_cuda_graph")
+            and bs <= self.max_bs_in_cuda_graph
+            and forward_batch.forward_mode.is_cuda_graph()
+        )
+
+        weight_indices = [0] * len(forward_batch.lora_ids)
+        lora_ranks = [0] * self.max_loras_per_batch
+        scalings = [0] * self.max_loras_per_batch
+        for i, uid in enumerate(forward_batch.lora_ids):
+            if uid not in self.memory_pool.uid_to_buffer_id:
+                continue
+            weight_indices[i] = self.memory_pool.get_buffer_id(uid)
+            if uid is not None:
+                lora = self.loras[uid]
+                lora_ranks[weight_indices[i]] = lora.config.r
+                scalings[weight_indices[i]] = lora.scaling
+        # Do in-place updates when CUDA graph is enabled and the batch forward mode
+        # could use CUDA graph.
+        self.lora_backend.prepare_lora_batch(
+            forward_batch=forward_batch,
+            weight_indices=weight_indices,
+            lora_ranks=lora_ranks,
+            scalings=scalings,
+            use_cuda_graph=use_cuda_graph,
+        )
+        self.lora_backend.batch_info.has_active_lora = any(
+            lora_ranks[wi] > 0 for wi in weight_indices
+        )
+
+    def update_lora_info(self):
+        """
+        Update all LoRA modules to associate them with the latest memory buffer.
+        """
+        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_key = (
+                        "gate_up_proj_moe"
+                        if "gate_up_proj_moe" in self.memory_pool.A_buffer
+                        else "gate_up_proj"
+                    )
+                    down_key = (
+                        "down_proj_moe"
+                        if "down_proj_moe" in self.memory_pool.A_buffer
+                        else "down_proj"
+                    )
+                    gate_up_a = self.memory_pool.get_tensor(
+                        target_module=gate_up_key,
+                        layer_id=layer_id,
+                        lora_type=LoRAType.LORA_A,
+                    )
+                    gate_up_b = self.memory_pool.get_tensor(
+                        target_module=gate_up_key,
+                        layer_id=layer_id,
+                        lora_type=LoRAType.LORA_B,
+                    )
+                    down_a = self.memory_pool.get_tensor(
+                        target_module=down_key,
+                        layer_id=layer_id,
+                        lora_type=LoRAType.LORA_A,
+                    )
+                    down_b = self.memory_pool.get_tensor(
+                        target_module=down_key,
+                        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,
+                        layer_id=layer_id,
+                        lora_type=LoRAType.LORA_A,
+                    ),
+                    self.memory_pool.get_tensor(
+                        target_module=target_module,
+                        layer_id=layer_id,
+                        lora_type=LoRAType.LORA_B,
+                    ),
+                )
+
+        # Update embedding layer if present - gotta merge (refer to PR codebase)
+        if self.embed_tokens_module is not None:
+            self.embed_tokens_module.set_lora_info(
+                self.memory_pool.get_embedding_tensor("added_tokens", LoRAType.LORA_A),
+                self.memory_pool.get_embedding_tensor("embed_tokens", LoRAType.LORA_A),
+                self.memory_pool.get_embedding_tensor("embed_tokens", LoRAType.LORA_B),
+            )
+
+        # Update lm_head layer if present
+        if self.lm_head_module is not None:
+            self.lm_head_module.set_lora_info(
+                self.memory_pool.get_embedding_tensor("lm_head", LoRAType.LORA_A),
+                self.memory_pool.get_embedding_tensor("lm_head", LoRAType.LORA_B),
+            )
+
+    def init_state(
+        self,
+        max_lora_rank: Optional[int] = None,
+        target_modules: Optional[Iterable[str]] = None,
+        lora_paths: Optional[List[LoRARef]] = None,
+    ):
+        """
+        Initialize the internal (mutable) state of the LoRAManager.
+
+        When `lora_paths` is provided and not empty, it might be used for inferring LoRA shape info such as
+        the target modules and max_lora_rank.
+        """
+
+        assert lora_paths or (
+            max_lora_rank is not None and target_modules is not None
+        ), "When no initial --lora-paths is provided, you need to specify both --max-lora-rank and --lora-target-modules for LoRA initialization."
+
+        self.init_lora_adapters(lora_paths)
+        self.init_lora_shapes(
+            max_lora_rank=max_lora_rank,
+            target_modules=target_modules,
+        )
+
+        if self._experts_shared_outer_override is not None:
+            self.experts_shared_outer_loras = self._experts_shared_outer_override
+        else:
+            self.experts_shared_outer_loras = self._detect_shared_outer_loras()
+        if self.experts_shared_outer_loras:
+            logger.info(
+                "Shared outer LoRA mode enabled: gate_up lora_A and "
+                "down lora_B will be shared across experts (expert_dim=1)."
+            )
+
+        self.init_lora_modules()
+        self.init_memory_pool()
+        self.update_lora_info()
+
+    def init_lora_adapters(self, lora_paths: Optional[List[LoRARef]] = None):
+        # Configs of all active LoRA adapters, indexed by LoRA ID.
+        self.configs: Dict[str, LoRAConfig] = {}
+
+        # LoRA adapter weights cached in CPU memory, indexed by LoRA ID.
+        self.loras: Dict[str, LoRAAdapter] = {}
+
+        # Mapping from LoRA ID to LoRARef object.
+        self.lora_refs: Dict[str, LoRARef] = {}
+
+        # Count of pinned LoRA adapters.
+        self.num_pinned_loras: int = 0
+
+        if lora_paths:
+            for lora_ref in lora_paths:
+                result = self.load_lora_adapter(lora_ref)
+                if not result.success:
+                    raise RuntimeError(
+                        f"Failed to load LoRA adapter {lora_ref.lora_name}: {result.error_message}"
+                    )
+
+    def _detect_shared_outer_loras(self) -> bool:
+        """Auto-detect shared outer LoRA format from loaded adapter weights.
+
+        MoE adapters with shared outer experts store 3D tensors where
+        dim[0]=1 indicates weights shared across all experts, while
+        dim[0]=num_experts indicates per-expert weights.
+        Returns True if gate_up lora_A has expert_dim=1 (shared).
+
+        All loaded adapters that expose a 3D gate_up lora_A must agree;
+        mixed formats raise RuntimeError.
+        """
+        shared_outer: Optional[bool] = None
+        for adapter_id, adapter in self.loras.items():
+            found = False
+            for layer in adapter.layers:
+                for name, weight in layer.weights.items():
+                    if (
+                        "gate_up_proj" in name
+                        and "lora_A" in name
+                        and weight.dim() == 3
+                    ):
+                        is_shared = weight.shape[0] == 1
+                        if shared_outer is None:
+                            shared_outer = is_shared
+                        elif shared_outer != is_shared:
+                            raise RuntimeError(
+                                "Mixed shared-outer LoRA formats detected across "
+                                f"loaded adapters (conflict in adapter '{adapter_id}'). "
+                                "All MoE adapters must either all use shared outer "
+                                "experts (expert_dim=1) or all use per-expert weights."
+                            )
+                        found = True
+                        break
+                if found:
+                    break
+        return bool(shared_outer) if shared_outer is not None else False
+
+    def init_lora_shapes(
+        self,
+        max_lora_rank: Optional[int] = None,
+        target_modules: Optional[Iterable[str]] = None,
+    ):
+        """Infer LoRA target modules and max_lora_rank from loaded adapters if not provided."""
+
+        if target_modules and target_modules == {"all"}:
+            self.target_modules = auto_detect_lora_target_modules(self.base_model)
+            self.target_modules.update(EMBEDDING_NAMES)
+            logger.info(
+                "CLI --lora-target-modules='all' resolved to %s "
+                "by inspecting the base model.",
+                sorted(self.target_modules),
+            )
+            target_modules = self.target_modules
+        elif target_modules:
+            self.target_modules = get_normalized_target_modules(target_modules)
+        else:
+            self.target_modules = set()
+
+        for lora_id, config in self.configs.items():
+            # Handle PEFT shorthand strings like "all-linear" or "all".
+            if isinstance(config.target_modules, str):
+                if config.target_modules in ("all-linear", "all"):
+                    if target_modules is not None:
+                        # CLI --lora-target-modules already provided; skip
+                        # per-adapter inference for this adapter.
+                        continue
+                    else:
+                        # Resolve by scanning the base model for all
+                        # LoRA-compatible linear modules.
+                        adapter_target_modules = auto_detect_lora_target_modules(
+                            self.base_model
+                        )
+                        logger.info(
+                            "LoRA adapter '%s' uses target_modules='%s'. "
+                            "Resolved to %s by inspecting the base model.",
+                            self.lora_refs[lora_id].lora_name,
+                            config.target_modules,
+                            sorted(adapter_target_modules),
+                        )
+                        self.target_modules.update(adapter_target_modules)
+                        continue
+                else:
+                    raise ValueError(
+                        f"SGLang does not recognize target_modules="
+                        f"'{config.target_modules}'. Please use a list of module "
+                        "name suffixes in the adapter's PEFT config, or explicitly "
+                        "specify --lora-target-modules during server startup."
+                    )
+
+            if not isinstance(config.target_modules, list):
+                raise ValueError(
+                    f"SGLang currently only supports inferring LoRA target modules when a list of "
+                    "suffixes is provided in `target_modules` field of PEFT config. Please explicitly "
+                    "specify `--lora-target-modules` during server startup. You can specify `all` to "
+                    "enable all support modules types. "
+                )
+
+            adapter_target_modules = get_normalized_target_modules(
+                config.target_modules
+            )
+
+            if target_modules is not None:
+                # When `--lora-target-modules` is provided, validate adapter target modules is a subset of the specified target modules.
+                if not adapter_target_modules.issubset(self.target_modules):
+                    unsupported_modules = adapter_target_modules - self.target_modules
+                    lora_name = self.lora_refs[lora_id].lora_name
+                    raise ValueError(
+                        f"LoRA adapter '{lora_name}' contains target modules {sorted(unsupported_modules)} "
+                        f"that are not included in the specified --lora-target-modules {sorted(self.target_modules)}. "
+                        f"Please update --lora-target-modules to include all required modules: "
+                        f"{sorted(self.target_modules | adapter_target_modules)}, or use 'all' to enable all supported modules."
+                    )
+            else:
+                # Otherwise, infer target_modules from adapter configs.
+                self.target_modules.update(adapter_target_modules)
+
+        # Fusion folds wk + weights_proj into wk_weights_proj, so the modules
+        # LoRA wraps are absent and an indexer-targeted adapter is silently dropped.
+        indexer_targets = self.target_modules & DSA_INDEXER_LORA_NAMES
+        if indexer_targets:
+            from sglang.srt.layers.attention.dsa.dsa_indexer import (
+                _use_dsa_indexer_fusion,
+            )
+
+            if _use_dsa_indexer_fusion:
+                raise ValueError(
+                    f"LoRA targets the DSA indexer ({sorted(indexer_targets)}), which is "
+                    "incompatible with DSA indexer Q/K fusion. Set "
+                    "SGLANG_DISABLE_DSA_INDEXER_FUSION=1 to disable fusion and use indexer LoRA."
+                )
+
+        if max_lora_rank is not None:
+            self.max_lora_rank = max_lora_rank
+        else:
+            self.max_lora_rank = max(
+                [x.r for x in self.configs.values()],
+                default=0,
+            )
+
+        # Auto-infer self.lora_added_vocab_size from loaded LoRA configs
+        # This happens automatically without requiring user input
+        # if self.lora_added_vocab_size is None:
+        if self.lora_added_tokens_size is None:
+            inferred_extra_vocab_size = next(
+                (
+                    x.lora_added_tokens_size
+                    for x in self.configs.values()
+                    if x.lora_added_tokens_size > 0
+                ),
+                0,
+            )
+            if inferred_extra_vocab_size > 0:
+                logger.info(
+                    f"self.lora_added_tokens_size={inferred_extra_vocab_size} from LoRA adapters."
+                )
+            self.lora_added_tokens_size = inferred_extra_vocab_size
+
+    def load_lora_weights(self, lora_ref: LoRARef):
+        """
+        Load the weights of a LoRA adapter to CPU memory and conducts post-loading validation.
+        """
+        lora_adapter = LoRAAdapter(
+            lora_ref.lora_id,
+            self.configs[lora_ref.lora_id],
+            self.base_hf_config,
+            self.load_config,
+            self.lora_backend,
+            base_model=self.base_model,
+        )
+        lora_adapter.initialize_weights()
+
+        self.loras[lora_ref.lora_id] = lora_adapter
+
+    def load_lora_weights_from_tensors(
+        self, lora_ref: LoRARef, tensors: Dict[str, torch.Tensor]
+    ):
+        """
+        Load the weights of a LoRA adapter from tensors to CPU memory.
+        """
+        lora_adapter = LoRAAdapter(
+            lora_ref.lora_id,
+            self.configs[lora_ref.lora_id],
+            self.base_hf_config,
+            self.load_config,
+            self.lora_backend,
+            base_model=self.base_model,
+        )
+        lora_adapter.initialize_weights_from_tensors(tensors)
+        self.loras[lora_ref.lora_id] = lora_adapter
+
+    def load_lora_adapter_from_tensors(
+        self,
+        lora_ref: LoRARef,
+        tensors: Dict[str, torch.Tensor],
+        config_dict: Dict,
+        added_tokens_config: Optional[Dict] = None,
+    ) -> LoRAUpdateOutput:
+        """
+        Load a single LoRA adapter from tensors and config dict.
+        """
+        assert (
+            lora_ref.lora_name is not None and lora_ref.lora_path is not None
+        ), "LoRARef must have both lora_name and lora_path set for loading."
+        assert (
+            lora_ref.lora_id not in self.loras
+        ), f"LoRA adapter with ID {lora_ref.lora_id} is already loaded. This should have been verified before request is sent to the backend."
+
+        try:
+            new_adapter = LoRAConfig.from_dict(
+                config_dict,
+                added_tokens_config,
+                base_vocab_size=self.base_hf_config.vocab_size,
+            )
+            self.validate_new_adapter(new_adapter, lora_ref)
+            self.configs[lora_ref.lora_id] = new_adapter
+
+            self.load_lora_weights_from_tensors(lora_ref, tensors)
+
+            self.lora_refs[lora_ref.lora_id] = lora_ref
+            self.num_pinned_loras += int(lora_ref.pinned)
+        except Exception as e:
+            return self.create_lora_update_result(
+                success=False,
+                error_message=str(e),
+            )
+
+        return self.create_lora_update_result(success=True)
+
+    def init_memory_pool(self):
+        """(Re)initialize the LoRA memory pool based on the current configurations."""
+        self.memory_pool = LoRAMemoryPool(
+            base_hf_config=self.base_hf_config,
+            max_loras_per_batch=self.max_loras_per_batch,
+            dtype=self.dtype,
+            tp_size=self.tp_size,
+            tp_rank=self.tp_rank,
+            max_lora_rank=self.max_lora_rank,
+            target_modules=self.target_modules,
+            base_model=self.base_model,
+            eviction_policy=self.eviction_policy,
+            lora_added_tokens_size=self.lora_added_tokens_size,
+            experts_shared_outer_loras=self.experts_shared_outer_loras,
+            strict_loading=self.lora_strict_loading,
+            enable_lora_overlap_loading=self.enable_lora_overlap_loading,
+        )
+
+        # Initializing memory pool with base model
+        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
+
+    def init_lora_modules(self):
+        # Look-up table that essentially maps (layer_index, module_name) to the corresponding LoRA module.
+        self.lora_modules: List[Dict[str, BaseLayerWithLoRA]] = [
+            {} for _ in range(self.base_hf_config.num_hidden_layers)
+        ]
+
+        self.embed_tokens_module: Optional[BaseLayerWithLoRA] = None
+        self.lm_head_module: Optional[BaseLayerWithLoRA] = None
+
+        # When tie_word_embeddings=True, lm_head is the same Python object as
+        # embed_tokens. PyTorch's named_modules() deduplicates by object identity,
+        # so lm_head will not appear as a separate entry in the scan below,
+        # preventing LoRA from wrapping it. To fix this, we create a new
+        # ParallelLMHead that shares the same base weight tensor (no extra GPU
+        # memory) so that named_modules() yields it as an independent module.
+        if "lm_head" in self.target_modules:
+            lm_head = getattr(self.base_model, "lm_head", None)
+            embed_tokens = None
+            for name, mod in self.base_model.named_modules():
+                if name.endswith("embed_tokens"):
+                    embed_tokens = mod
+                    break
+            if (
+                lm_head is not None
+                and embed_tokens is not None
+                and lm_head is embed_tokens
+            ):
+                logger.info(
+                    "lm_head is tied with embed_tokens. Creating a separate "
+                    "ParallelLMHead that shares the base weight for LoRA support."
+                )
+                untied_lm_head = ParallelLMHead(
+                    num_embeddings=embed_tokens.org_vocab_size,
+                    embedding_dim=embed_tokens.embedding_dim,
+                    params_dtype=embed_tokens.weight.dtype,
+                    org_num_embeddings=embed_tokens.org_vocab_size,
+                )
+                # Share the base weight tensor — no additional GPU memory.
+                untied_lm_head.weight = embed_tokens.weight
+                # Replace the model attribute so named_modules() sees it
+                # independently.
+                self.base_model.lm_head = untied_lm_head
+
+        for module_name, module in self.base_model.named_modules():
+            # Handle embed_tokens and lm_head before the should_apply_lora gate,
+            # since VL models' should_apply_lora patterns only match language
+            # model layers and would incorrectly skip these.
+            # Handle embed_tokens
+            if "embed_tokens" in module_name and "embed_tokens" in self.target_modules:
+                if isinstance(module, VocabParallelEmbedding) and not isinstance(
+                    module, BaseLayerWithLoRA
+                ):
+                    lora_module = self.set_lora_module(module_name, module)
+                    self.embed_tokens_module = lora_module
+                    continue
+            # Handle lm_head
+            if "lm_head" in module_name and "lm_head" in self.target_modules:
+                if isinstance(module, ParallelLMHead) and not isinstance(
+                    module, BaseLayerWithLoRA
+                ):
+                    lora_module = self.set_lora_module(module_name, module)
+                    self.lm_head_module = lora_module
+                    continue
+
+            # Handle DeepSeek MLA fused projection: set the boundary
+            # between q_a and kv_a output partitions so the LoRA layer
+            # can apply separate B projections for each.
+            if (
+                "fused_qkv_a_proj_with_mqa" in self.target_modules
+                and module_name.endswith("fused_qkv_a_proj_with_mqa")
+            ):
+                from sglang.srt.lora.layers import ReplicatedLinearWithLoRA
+
+                layer_id = get_layer_id(module_name)
+                if layer_id is None:
+                    continue
+                lora_module = self.set_lora_module(module_name, module)
+                if isinstance(lora_module, ReplicatedLinearWithLoRA):
+                    q_lora_rank = getattr(self.base_hf_config, "q_lora_rank", None) or 0
+                    lora_module.first_output_dim = q_lora_rank
+                self.lora_modules[layer_id][module_name] = lora_module
+                continue
+
+            # The module should be converted if it is included in target_names
+            parts = module_name.split(".")
+            if (
+                parts[-1] in self.target_modules
+                or ".".join(parts[-2:]) in self.target_modules
+            ):
+                layer_id = get_layer_id(module_name)
+                if layer_id is None:
+                    continue
+                self.lora_modules[layer_id][module_name] = self.set_lora_module(
+                    module_name, module
+                )
+                continue
+
+            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)
+                if layer_id is None:
+                    # FusedMoE submodules outside the decoder layer hierarchy
+                    # (e.g. nested helpers under non-".layers." prefixes) have
+                    # no resolvable layer id; skip them so we don't index
+                    # `self.lora_modules` with `None`.
+                    continue
+                lora_module = self.set_lora_module(module_name, module)
+                lora_module.experts_shared_outer_loras = self.experts_shared_outer_loras
+                lora_module.lora_use_virtual_experts = self.lora_use_virtual_experts
+                self.lora_modules[layer_id][module_name] = lora_module

python/sglang/srt/lora/lora_moe_runner_marlin.py

@@ -0,0 +1,200 @@
+"""Marlin MoE runner core with hook support for LoRA injection.
+
+Uses Marlin int4/int8 kernels for the base MoE projections.
+LoRA deltas are injected via hooks.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
+from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo
+from sglang.srt.utils import is_cuda
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.token_dispatcher import (
+        StandardCombineInput,
+        StandardDispatchOutput,
+    )
+
+_is_cuda = is_cuda()
+
+if _is_cuda:
+    from sglang.jit_kernel.moe_wna16_marlin import moe_wna16_marlin_gemm
+    from sglang.kernels.ops.activation import silu_and_mul
+    from sglang.srt.layers.moe.fused_moe_triton.fused_marlin_moe import (
+        get_scalar_type,
+    )
+    from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import (
+        moe_align_block_size,
+    )
+    from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_kernels import (
+        moe_sum_reduce_triton,
+    )
+    from sglang.srt.layers.quantization.marlin_utils import marlin_make_workspace
+
+
+class MarlinLoraRunnerCore:
+    """
+    MoE runner using Marlin kernels for base projections, with hooks for LoRA.
+
+    Pipeline:
+      1. moe_wna16_marlin_gemm (gate_up)
+      1.5. hooks.after_gate_up
+      2. silu_and_mul
+      3. moe_wna16_marlin_gemm (down)
+      3.5. hooks.after_down
+      4. moe_sum_reduce
+    """
+
+    def __init__(self, config: MoeRunnerConfig):
+        self.config = config
+
+    def run_from_dispatch(
+        self,
+        dispatch_output: StandardDispatchOutput,
+        quant_info: MarlinMoeQuantInfo,
+        runner_config: MoeRunnerConfig,
+        hooks=None,
+    ) -> StandardCombineInput:
+        from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
+
+        assert hooks is not None, "hooks must be provided for MarlinLoraRunnerCore"
+
+        hidden_states = dispatch_output.hidden_states
+        topk_output = dispatch_output.topk_output
+        topk_weights = topk_output.topk_weights
+        topk_ids = topk_output.topk_ids
+
+        assert runner_config.activation == "silu", "Only SiLU activation is supported."
+        assert (
+            torch.cuda.get_device_capability(hidden_states.device)[0] >= 9
+        ), "MarlinLoraRunnerCore requires CUDA compute capability >= 9"
+        inplace = runner_config.inplace
+        routed_scaling_factor = runner_config.routed_scaling_factor
+
+        M, K = hidden_states.shape
+        E = quant_info.w13_qweight.shape[0]
+        N = quant_info.w2_qweight.shape[1] * 16
+        topk = topk_ids.shape[1]
+        num_bits = quant_info.weight_bits
+
+        for block_size_m in [8, 16, 32, 48, 64]:
+            if M * topk / E / block_size_m < 0.9:
+                break
+
+        sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
+            topk_ids, block_size_m, E
+        )
+
+        from sglang.srt.runtime_context import get_resources
+
+        buffers = get_resources().buffers
+        workspace = buffers.get("marlin_lora_workspace")
+        if workspace is None or workspace.device != hidden_states.device:
+            workspace = marlin_make_workspace(hidden_states.device, max_blocks_per_sm=4)
+            buffers["marlin_lora_workspace"] = workspace
+
+        scalar_type1 = get_scalar_type(num_bits, quant_info.w13_qzeros is not None)
+        scalar_type2 = get_scalar_type(num_bits, quant_info.w2_qzeros is not None)
+
+        # Stage 1: Gate/Up (Marlin)
+        intermediate_cache1 = torch.empty(
+            (M * topk, 2 * N), device=hidden_states.device, dtype=hidden_states.dtype
+        )
+        intermediate_cache1 = moe_wna16_marlin_gemm(
+            hidden_states,
+            intermediate_cache1,
+            quant_info.w13_qweight,
+            None,
+            quant_info.w13_scales,
+            None,
+            quant_info.w13_qzeros,
+            quant_info.w13_g_idx,
+            quant_info.w13_g_idx_sort_indices,
+            workspace,
+            sorted_token_ids,
+            expert_ids,
+            num_tokens_post_padded,
+            topk_weights,
+            moe_block_size=block_size_m,
+            top_k=topk,
+            mul_topk_weights=False,
+            is_ep=quant_info.expert_map is not None,
+            b_q_type=scalar_type1,
+            size_m=M,
+            size_n=2 * N,
+            size_k=K,
+            is_k_full=quant_info.is_k_full,
+            use_atomic_add=True,
+            use_fp32_reduce=True,
+            is_zp_float=False,
+        )
+
+        # Hook: after gate_up
+        if hooks.after_gate_up:
+            intermediate_cache1_3d = intermediate_cache1.view(M, topk, 2 * N)
+            hooks.after_gate_up(
+                hidden_states, intermediate_cache1_3d, topk_weights, topk_ids
+            )
+
+        # Stage 2: Activation
+        intermediate_cache2 = torch.empty(
+            (M * topk, N), device=hidden_states.device, dtype=hidden_states.dtype
+        )
+        silu_and_mul(intermediate_cache1.view(-1, 2 * N), intermediate_cache2)
+
+        # Stage 3: Down (Marlin)
+        intermediate_cache3 = torch.empty(
+            (M * topk, K), device=hidden_states.device, dtype=hidden_states.dtype
+        )
+        if quant_info.expert_map is not None:
+            intermediate_cache3.zero_()
+
+        intermediate_cache3 = moe_wna16_marlin_gemm(
+            intermediate_cache2,
+            intermediate_cache3,
+            quant_info.w2_qweight,
+            None,
+            quant_info.w2_scales,
+            None,
+            quant_info.w2_qzeros,
+            quant_info.w2_g_idx,
+            quant_info.w2_g_idx_sort_indices,
+            workspace,
+            sorted_token_ids,
+            expert_ids,
+            num_tokens_post_padded,
+            topk_weights,
+            moe_block_size=block_size_m,
+            top_k=1,
+            mul_topk_weights=True,
+            is_ep=quant_info.expert_map is not None,
+            b_q_type=scalar_type2,
+            size_m=M * topk,
+            size_n=K,
+            size_k=N,
+            is_k_full=quant_info.is_k_full,
+            use_atomic_add=True,
+            use_fp32_reduce=True,
+            is_zp_float=False,
+        )
+        intermediate_cache3 = intermediate_cache3.view(M, topk, K)
+
+        # Hook: after down
+        if hooks.after_down:
+            hooks.after_down(
+                intermediate_cache2, intermediate_cache3, topk_weights, topk_ids
+            )
+
+        # Stage 4: Reduction
+        output = hidden_states if inplace else torch.empty_like(hidden_states)
+        if routed_scaling_factor is None:
+            routed_scaling_factor = 1.0
+        # NOTE: fusion opportunity here
+        moe_sum_reduce_triton(intermediate_cache3, output, routed_scaling_factor)
+
+        return StandardCombineInput(hidden_states=output)

python/sglang/srt/lora/lora_moe_runners.py

@@ -0,0 +1,556 @@
+# 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 hooks for MoE runners.
+
+LoRA deltas are injected at two points in the MoE pipeline:
+1. After gate_up projection, BEFORE activation
+2. After down projection, BEFORE final reduction
+
+This module provides hook closures that any MoE backend can call at those points,
+without needing a per-backend LoRA runner subclass.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable
+
+import torch
+
+from sglang.srt.model_executor.runner import get_is_capture_mode
+from sglang.srt.utils import is_cuda, is_hip, is_xpu, next_power_of_2
+
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+_is_xpu = is_xpu()
+
+if _is_cuda or _is_hip or _is_xpu:
+    from sglang.jit_kernel.moe_lora_align import moe_lora_align_block_size
+
+
+def _get_moe_lora_block_config(max_lora_rank: int) -> dict:
+    """Compute rank-aware block sizes for MoE LoRA kernels.
+
+    Shrink: output dim is the rank -> cap BLOCK_SIZE_N to avoid waste.
+    Expand: input dim is the rank -> cap BLOCK_SIZE_K similarly.
+    """
+    if max_lora_rank <= 0:
+        rank_pow2 = 64
+    else:
+        rank_pow2 = next_power_of_2(max_lora_rank)
+
+    shrink_n = min(64, rank_pow2)
+    expand_k = max(16, min(64, rank_pow2))
+
+    return {
+        "shrink_block_size_n": shrink_n,
+        "expand_block_size_k": expand_k,
+    }
+
+
+_SPARSITY_FACTOR = 8
+
+
+def _naive_moe_lora_align_block_size(
+    topk_ids: torch.Tensor,
+    seg_indptr: torch.Tensor,
+    req_to_lora: torch.Tensor,
+    num_experts: int,
+    block_size_m: int,
+    max_loras: int,
+    max_num_tokens_padded: int,
+    max_num_m_blocks: int,
+    adapter_enabled: torch.Tensor,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Construct LoRA token-expert alignment on CPU for small batches.
+
+    When the number of tokens is very small, the overhead of launching the
+    CUDA-based moe_lora_align_block_size kernel exceeds the actual
+    computation. This function builds the same data structures using simple
+    Python loops on CPU and transfers the result to GPU in one shot.
+    """
+    M, top_k = topk_ids.shape
+    num_valid_tokens = M * top_k
+
+    sorted_token_ids = torch.full(
+        (max_loras * max_num_tokens_padded,),
+        num_valid_tokens,
+        dtype=torch.int32,
+    )
+    expert_ids_out = torch.full((max_loras * max_num_m_blocks,), -1, dtype=torch.int32)
+    num_tokens_post_padded = torch.zeros(max_loras, dtype=torch.int32)
+
+    seg_indptr_list = seg_indptr.cpu().tolist()
+    req_to_lora_list = req_to_lora.cpu().tolist()
+    topk_ids_list = topk_ids.cpu().tolist()
+    adapter_enabled_list = adapter_enabled.cpu().tolist()
+
+    for lora_id in range(max_loras):
+        if not adapter_enabled_list[lora_id]:
+            continue
+
+        pairs: list[tuple[int, int]] = []
+        for seg_idx in range(len(seg_indptr_list) - 1):
+            if req_to_lora_list[seg_idx] != lora_id:
+                continue
+            start = seg_indptr_list[seg_idx]
+            end = seg_indptr_list[seg_idx + 1]
+            for m in range(start, end):
+                for k in range(top_k):
+                    pairs.append((topk_ids_list[m][k], m * top_k + k))
+
+        if not pairs:
+            continue
+
+        pairs.sort()
+
+        base_t = lora_id * max_num_tokens_padded
+        base_e = lora_id * max_num_m_blocks
+        pos = 0
+        block_idx = 0
+        i = 0
+        while i < len(pairs):
+            cur_expert = pairs[i][0]
+            group_start = pos
+            while i < len(pairs) and pairs[i][0] == cur_expert:
+                sorted_token_ids[base_t + pos] = pairs[i][1]
+                pos += 1
+                i += 1
+            group_len = pos - group_start
+            padded_len = ((group_len + block_size_m - 1) // block_size_m) * block_size_m
+            num_blocks = padded_len // block_size_m
+            for b in range(num_blocks):
+                expert_ids_out[base_e + block_idx + b] = cur_expert
+            block_idx += num_blocks
+            pos = group_start + padded_len
+
+        num_tokens_post_padded[lora_id] = pos
+
+    return (
+        sorted_token_ids.to(device),
+        expert_ids_out.to(device),
+        num_tokens_post_padded.to(device),
+    )
+
+
+@dataclass
+class LoRAInfo:
+    """LoRA weights and dispatch info for MoE computation."""
+
+    # LoRA weights: [num_loras, num_experts_or_1, dim1, dim2]
+    # When experts_shared_outer_loras=True:
+    #   gate_up_lora_a: [num_loras, 1, max_rank, hidden_dim] (shared)
+    #   down_lora_b: [num_loras, 1, hidden_dim, max_rank] (shared)
+    gate_up_lora_a_weights: (
+        torch.Tensor
+    )  # [num_loras, num_experts_or_1, 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_or_1, 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
+    token_lora_mapping: torch.Tensor  # [num_tokens] - adapter used by each token
+    max_lora_rank: int  # Maximum LoRA rank across all adapters
+
+    num_experts: int
+    has_active_lora: bool = True
+    experts_shared_outer_loras: bool = False
+    cg_buffers: dict | None = None
+
+    fully_sharded: bool = False
+    tp_size: int = 1
+    tp_rank: int = 0
+    hidden_size: int = 0
+    lora_use_virtual_experts: bool = False
+
+
+@dataclass
+class LoRAHooks:
+    """Hook callbacks for injecting LoRA deltas into the MoE pipeline."""
+
+    after_gate_up: (
+        Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], None] | None
+    ) = None
+    after_down: (
+        Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], None] | None
+    ) = None
+
+
+def _compute_lora_alignment(
+    topk_ids: torch.Tensor,
+    lora_info: LoRAInfo,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Compute LoRA alignment tensors for the non-virtual-expert (classic) path.
+
+    Returns: (sorted_token_ids_reshaped, expert_ids_reshaped, num_tokens_post_padded_lora, lora_ids)
+    """
+    cg = lora_info.cg_buffers if get_is_capture_mode() else None
+    shrink_config = {"BLOCK_SIZE_M": 64}
+    M = topk_ids.shape[0]
+    block_size_m = shrink_config["BLOCK_SIZE_M"]
+    max_loras = len(lora_info.lora_ranks)
+
+    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
+
+    device = topk_ids.device
+
+    use_naive = (
+        cg is None
+        and M * topk_ids.shape[1] * _SPARSITY_FACTOR
+        <= lora_info.num_experts * max_loras
+    )
+
+    if use_naive:
+        sorted_token_ids_lora, expert_ids_lora, num_tokens_post_padded_lora = (
+            _naive_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),
+                lora_info.adapter_enabled,
+                device,
+            )
+        )
+        lora_ids = torch.arange(max_loras, dtype=torch.int32, device=device)
+    else:
+        if cg is not None:
+            sorted_token_ids_lora = cg["sorted_token_ids_lora"][
+                : max_loras * max_num_tokens_padded
+            ]
+            expert_ids_lora = cg["expert_ids_lora"][: max_loras * max_num_m_blocks]
+            num_tokens_post_padded_lora = cg["num_tokens_post_padded_lora"][:max_loras]
+        else:
+            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
+            )
+
+        if cg is not None and "lora_ids" in cg:
+            lora_ids = cg["lora_ids"][:max_loras]
+        else:
+            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,
+            cumsum_buffer=cg.get("cumsum_buffer") if cg is not None else None,
+            token_mask=cg.get("token_mask") if cg is not None else None,
+        )
+
+    return (
+        sorted_token_ids_lora.view(max_loras, -1),
+        expert_ids_lora.view(max_loras, -1),
+        num_tokens_post_padded_lora,
+        lora_ids,
+    )
+
+
+def _add_lora_gate_up_delta(
+    hidden_states: torch.Tensor,
+    intermediate_cache: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    lora_info: LoRAInfo,
+    token_lora_mapping: torch.Tensor | None,
+    sorted_token_ids_reshaped: torch.Tensor | None,
+    expert_ids_reshaped: torch.Tensor | None,
+    num_tokens_post_padded_lora: torch.Tensor | None,
+    lora_ids: torch.Tensor | None,
+    routing_cache: dict | None = None,
+) -> None:
+    """Add LoRA gate_up delta to intermediate_cache in-place."""
+    from sglang.kernels.ops.moe.fused_moe_lora_kernel import fused_moe_lora
+    from sglang.kernels.ops.moe.virtual_experts import merged_experts_fused_moe_lora_add
+
+    if lora_info is None or lora_info.max_lora_rank == 0:
+        return
+    if not get_is_capture_mode() and not lora_info.has_active_lora:
+        return
+
+    M, top_k, gate_up_dim = intermediate_cache.shape
+    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
+
+    if lora_info.experts_shared_outer_loras and not lora_info.lora_use_virtual_experts:
+        gate_up_a = gate_up_a.expand(-1, lora_info.num_experts, -1, -1)
+
+    # Detect gated vs non-gated from A buffer rank dimension.
+    # Gated: A has 2*r rows (gate + up). Non-gated: A has 1*r rows (w1 only).
+    is_gated = gate_up_a.shape[2] > r
+    if is_gated:
+        inter_size = gate_up_b.shape[2] // 2
+        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:, :],
+        ]
+    else:
+        lora_a_stacked = [gate_up_a]
+        lora_b_stacked = [gate_up_b]
+
+    if lora_info.lora_use_virtual_experts:
+        merged_experts_fused_moe_lora_add(
+            output=intermediate_cache,
+            hidden_states=hidden_states,
+            lora_a=gate_up_a,
+            lora_b=gate_up_b,
+            topk_ids=topk_ids,
+            topk_weights=topk_weights,
+            token_lora_mapping=token_lora_mapping,
+            mul_routed_weight=False,
+            experts_shared_outer_loras_a=lora_info.experts_shared_outer_loras,
+            experts_shared_outer_loras_b=False,
+            routing_cache=routing_cache,
+        )
+    else:
+        blk = _get_moe_lora_block_config(r)
+        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,
+            shrink_block_size_m=64,
+            shrink_block_size_n=blk["shrink_block_size_n"],
+            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=blk["expand_block_size_k"],
+            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(
+    intermediate_input: torch.Tensor,
+    intermediate_cache: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    lora_info: LoRAInfo,
+    token_lora_mapping: torch.Tensor | None,
+    sorted_token_ids_reshaped: torch.Tensor | None,
+    expert_ids_reshaped: torch.Tensor | None,
+    num_tokens_post_padded_lora: torch.Tensor | None,
+    lora_ids: torch.Tensor | None,
+    routing_cache: dict | None = None,
+) -> None:
+    """Add LoRA down delta to intermediate_cache in-place."""
+    from sglang.kernels.ops.moe.fused_moe_lora_kernel import fused_moe_lora
+    from sglang.kernels.ops.moe.virtual_experts import merged_experts_fused_moe_lora_add
+
+    if lora_info.max_lora_rank == 0:
+        return
+
+    M, top_k, hidden_dim = intermediate_cache.shape
+
+    down_lora_a = lora_info.down_lora_a_weights
+    down_lora_b = lora_info.down_lora_b_weights
+    if lora_info.experts_shared_outer_loras and not lora_info.lora_use_virtual_experts:
+        down_lora_b = down_lora_b.expand(-1, lora_info.num_experts, -1, -1)
+
+    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
+
+    if lora_info.lora_use_virtual_experts:
+        merged_experts_fused_moe_lora_add(
+            output=intermediate_cache,
+            hidden_states=intermediate_input,
+            lora_a=down_lora_a,
+            lora_b=down_lora_b,
+            topk_ids=topk_ids,
+            topk_weights=topk_weights,
+            token_lora_mapping=token_lora_mapping,
+            mul_routed_weight=True,
+            experts_shared_outer_loras_a=False,
+            experts_shared_outer_loras_b=lora_info.experts_shared_outer_loras,
+            routing_cache=routing_cache,
+        )
+    else:
+        blk = _get_moe_lora_block_config(lora_info.max_lora_rank)
+        fused_moe_lora(
+            output=intermediate_cache,
+            qcurr_hidden_states=intermediate_input,
+            lora_a_stacked=[down_lora_a],
+            lora_b_stacked=[down_lora_b],
+            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,
+            shrink_block_size_m=64,
+            shrink_block_size_n=blk["shrink_block_size_n"],
+            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=blk["expand_block_size_k"],
+            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,
+        )
+
+
+def build_lora_hooks(
+    hidden_states: torch.Tensor,
+    lora_info: LoRAInfo,
+    topk_ids: torch.Tensor,
+) -> LoRAHooks:
+    """Build LoRA hook closures for injection into any MoE runner.
+
+    Computes token_lora_mapping and alignment tensors once, then returns
+    closures that capture them for the two injection points.
+    """
+    if lora_info is None or lora_info.max_lora_rank == 0:
+        return LoRAHooks()
+    # Skip alignment/mapping work entirely when the batch has no active adapter.
+    # During CUDA graph capture we still need to record the kernels into the
+    # graph (adapter_enabled is all-zero, kernels early-exit on GPU).
+    if not get_is_capture_mode() and not lora_info.has_active_lora:
+        return LoRAHooks()
+
+    # Compute alignment / mapping (once, shared by both hooks)
+    token_lora_mapping: torch.Tensor | None = None
+    sorted_token_ids_reshaped: torch.Tensor | None = None
+    expert_ids_reshaped: torch.Tensor | None = None
+    num_tokens_post_padded_lora: torch.Tensor | None = None
+    lora_ids: torch.Tensor | None = None
+
+    if lora_info.lora_use_virtual_experts:
+        token_lora_mapping = lora_info.token_lora_mapping
+    else:
+        (
+            sorted_token_ids_reshaped,
+            expert_ids_reshaped,
+            num_tokens_post_padded_lora,
+            lora_ids,
+        ) = _compute_lora_alignment(topk_ids, lora_info)
+
+    # Shared routing cache: gate_up and down reuse routing for same (num_experts, shared_outer, block_size)
+    routing_cache: dict = {}
+
+    def after_gate_up(
+        hidden_states: torch.Tensor,
+        intermediate_cache1: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+    ) -> None:
+        _add_lora_gate_up_delta(
+            hidden_states,
+            intermediate_cache1,
+            topk_weights,
+            topk_ids,
+            lora_info,
+            token_lora_mapping,
+            sorted_token_ids_reshaped,
+            expert_ids_reshaped,
+            num_tokens_post_padded_lora,
+            lora_ids,
+            routing_cache=routing_cache,
+        )
+
+    def after_down(
+        intermediate_input: torch.Tensor,
+        intermediate_cache3: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+    ) -> None:
+        _add_lora_down_delta(
+            intermediate_input,
+            intermediate_cache3,
+            topk_weights,
+            topk_ids,
+            lora_info,
+            token_lora_mapping,
+            sorted_token_ids_reshaped,
+            expert_ids_reshaped,
+            num_tokens_post_padded_lora,
+            lora_ids,
+            routing_cache=routing_cache,
+        )
+
+    return LoRAHooks(after_gate_up=after_gate_up, after_down=after_down)

python/sglang/srt/lora/lora_overlap_loader.py

@@ -0,0 +1,93 @@
+import logging
+from enum import Enum, auto
+from typing import Dict, Optional
+
+import torch
+from torch.cuda import Event as CudaEvent
+from torch.cuda import Stream as CudaStream
+from torch.cuda import StreamContext as CudaStreamContext
+
+from sglang.srt.lora.lora_manager import LoRAManager
+
+logger = logging.getLogger(__name__)
+
+
+class LoRAOverlapLoadStatus(Enum):
+    LOADED = auto()
+    LOADING = auto()
+    NOT_LOADED = auto()
+
+
+class LoRAOverlapLoader:
+    def __init__(self, lora_manager):
+        self.lora_manager: LoRAManager = lora_manager
+        self.device_module = torch.get_device_module(self.lora_manager.device)
+        self.load_stream: CudaStream = self.device_module.Stream()
+        self.load_stream_context: CudaStreamContext = self.device_module.stream(
+            self.load_stream
+        )
+        self.lora_to_overlap_load_event: Dict[Optional[str], CudaEvent] = {}
+
+    def try_overlap_load_lora(
+        self, lora_id: Optional[str], running_loras: set[Optional[str]]
+    ) -> bool:
+        """
+        Check a LoRA adapter's asynchronous load status, and try to load it if there's capacity
+        in the memory pool. Returns whether or not the adapter has been loaded.
+        """
+        # Drain completed async loads before status/capacity checks so finished
+        # adapters no longer count as in-flight.
+        self._drain_completed_overlap_loads()
+
+        lora_pipeline_load_status = self._check_overlap_load_status(lora_id)
+        if lora_pipeline_load_status == LoRAOverlapLoadStatus.LOADING:
+            return False
+        elif lora_pipeline_load_status == LoRAOverlapLoadStatus.NOT_LOADED:
+            res = self._try_start_overlap_load(lora_id, running_loras)
+            if res:
+                logger.debug(f"Loading LoRA adapter {lora_id} asynchronously")
+
+            return False
+        else:
+            assert lora_pipeline_load_status == LoRAOverlapLoadStatus.LOADED
+            return True
+
+    def _check_overlap_load_status(
+        self, lora_id: Optional[str]
+    ) -> LoRAOverlapLoadStatus:
+        if lora_id in self.lora_to_overlap_load_event:
+            return LoRAOverlapLoadStatus.LOADING
+
+        # After completed events have been drained, a memory-pool entry with no
+        # pending event is safe to use on the current stream.
+        if lora_id in self.lora_manager.memory_pool.uid_to_buffer_id:
+            return LoRAOverlapLoadStatus.LOADED
+
+        return LoRAOverlapLoadStatus.NOT_LOADED
+
+    def _drain_completed_overlap_loads(self) -> None:
+        completed_loads = [
+            (lora_id, event)
+            for lora_id, event in self.lora_to_overlap_load_event.items()
+            if event.query()
+        ]
+        for lora_id, event in completed_loads:
+            torch.cuda.current_stream().wait_event(event)
+            del self.lora_to_overlap_load_event[lora_id]
+
+    def _try_start_overlap_load(
+        self, lora_id: Optional[str], running_loras: set[Optional[str]]
+    ) -> bool:
+        loras_to_be_loaded = running_loras | self.lora_to_overlap_load_event.keys()
+
+        new_lora_set = {lora_id} | loras_to_be_loaded
+        if not self.lora_manager.validate_lora_batch(new_lora_set):
+            return False
+
+        with self.load_stream_context:
+            self.lora_manager.fetch_new_loras({lora_id}, loras_to_be_loaded)
+            event = self.device_module.Event()
+            event.record(self.load_stream)
+
+        self.lora_to_overlap_load_event[lora_id] = event
+        return True

python/sglang/srt/lora/lora_registry.py

@@ -0,0 +1,263 @@
+# Copyright 2023-2024 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.
+# ==============================================================================
+
+
+import asyncio
+from collections import OrderedDict
+from typing import Dict, List, Optional, Union
+from uuid import NAMESPACE_URL, uuid4, uuid5
+
+import msgspec
+from msgspec.structs import fields
+
+from sglang.srt.utils import ConcurrentCounter
+from sglang.srt.utils.aio_rwlock import RWLock
+
+
+class LoRARef(msgspec.Struct, frozen=True, array_like=True):
+    """
+    Reference record for a LoRA model.
+
+    This object guarantees a unique ``lora_id`` and may include ``lora_name``, ``lora_path``, and ``pinned``.
+    The ID eliminates conflicts from reused LoRA names or paths and can be used to generate deterministic cache
+    keys (e.g., radix cache).
+    """
+
+    lora_id: str = msgspec.field(default_factory=lambda: uuid4().hex)
+    lora_name: Optional[str] = None
+    lora_path: Optional[str] = None
+    pinned: Optional[bool] = None
+
+    def __post_init__(self):
+        if self.lora_id is None:
+            raise ValueError("lora_id cannot be None")
+
+    @staticmethod
+    def deterministic_id(lora_name: str, lora_path: str) -> str:
+        """Stable ``lora_id`` for ``--lora-paths`` adapters.
+
+        Each node in a multi-node launch parses ``--lora-paths`` independently;
+        ``uuid4`` would mint a different id per node for the same adapter,
+        breaking cross-node lookups when the master broadcasts a request id.
+        """
+        return uuid5(NAMESPACE_URL, f"{lora_name}\0{lora_path}").hex
+
+    def __str__(self) -> str:
+        parts = [
+            f"{f.name}={value}"
+            for f in fields(self)
+            if (value := getattr(self, f.name)) is not None
+        ]
+        return f"{self.__class__.__name__}({', '.join(parts)})"
+
+
+class LoRARegistry:
+    """
+    The central registry to keep track of available LoRA adapters and ongoing LoRA requests.
+
+    The `LoRARegistry` resides in the tokenizer manager process and acts as the single source of truth for all
+    available LoRA adapters. It supports concurrent inference and dynamic adapter updates through a two-phase
+    update / eventual consistency model between the tokenizer manager process and the scheduler processes.
+    """
+
+    def __init__(self, lora_paths: Optional[List[LoRARef]] = None):
+        assert lora_paths is None or all(
+            isinstance(lora, LoRARef) for lora in lora_paths
+        ), (
+            "server_args.lora_paths should have been normalized to LoRARef objects during server initialization. "
+            "Please file an issue if you see this error."
+        )
+
+        # A read-write lock to ensure adapters loading / unloading operations are exclusive.
+        # Please note that the counter increment/decrement operations are not synchronized through this
+        # lock, as they are designed to be non-blocking and can be performed concurrently.
+        self._registry_lock = RWLock()
+        # An ordered dictionary to hold LoRARef objects, mapping from LoRA name to LoRARef.
+        # The LoRARefs are stored in LRU order, such that LoRA adapters that have been
+        # most recently used are stored at the end. Note that lookups count for accesses.
+        # Ties are broken arbitrarily.
+        self._registry: OrderedDict[str, LoRARef] = OrderedDict()
+        # Counters for ongoing requests, mapping from LoRA ID to ConcurrentCounter.
+        self._counters: Dict[str, ConcurrentCounter] = {}
+
+        # Initialize the registry with provided LoRA paths, if present.
+        if lora_paths:
+            for lora_ref in lora_paths:
+                self._register_adapter(lora_ref)
+
+    async def register(self, lora_ref: LoRARef):
+        """
+        Register a new LoRARef object in the registry.
+
+        Args:
+            lora_ref (LoRARef): The LoRARef object to register.
+        """
+        async with self._registry_lock.writer_lock:
+            self._register_adapter(lora_ref)
+
+    async def unregister(self, lora_name: str) -> str:
+        """
+        Unregister a LoRARef object from the registry and returns the removed LoRA ID.
+
+        Args:
+            lora_name (str): The name of the LoRA model to unregister.
+        """
+        async with self._registry_lock.writer_lock:
+            lora_ref = self._registry.get(lora_name, None)
+            if lora_ref is None:
+                raise ValueError(
+                    f"LoRA with name {lora_name} does not exist. Loaded LoRAs: {self._registry.keys()}"
+                )
+            del self._registry[lora_name]
+
+        return lora_ref.lora_id
+
+    async def acquire(self, lora_name: Union[str, List[str]]) -> Union[str, List[str]]:
+        """
+        Queries registry for LoRA IDs based on LoRA names and start tracking the usage of the corresponding LoRA adapters
+        by incrementing its counter.
+        """
+
+        def _lookup(name: str) -> str:
+            if name is None:
+                return None
+
+            lora_ref = self._registry.get(name, None)
+            if lora_ref is None:
+                raise ValueError(
+                    f"The following requested LoRA adapters are not loaded: {name}\n"
+                    f"Loaded adapters: {self._registry.keys()}."
+                )
+            self._registry.move_to_end(name)
+            return lora_ref.lora_id
+
+        if isinstance(lora_name, str):
+            async with self._registry_lock.writer_lock:
+                lora_id = _lookup(lora_name)
+
+            await self._counters[lora_id].increment(notify_all=False)
+            return lora_id
+        elif isinstance(lora_name, list):
+            async with self._registry_lock.writer_lock:
+                lora_ids = [_lookup(name) for name in lora_name]
+
+            # Increment the counters only after all IDs are looked up.
+            await asyncio.gather(
+                *[
+                    self._counters[id].increment(notify_all=False)
+                    for id in lora_ids
+                    if id is not None
+                ]
+            )
+            return lora_ids
+        else:
+            raise TypeError("lora_name must be either a string or a list of strings.")
+
+    async def release(self, lora_id: Union[str, List[str]]):
+        """
+        Decrements the usage counter for a LoRA adapter, indicating that it is no longer in use.
+        """
+
+        async with self._registry_lock.reader_lock:
+            if isinstance(lora_id, str):
+                await self._counters[lora_id].decrement()
+            elif isinstance(lora_id, list):
+                await asyncio.gather(
+                    *[
+                        self._counters[id].decrement()
+                        for id in lora_id
+                        if id is not None
+                    ]
+                )
+            else:
+                raise TypeError("lora_id must be either a string or a list of strings.")
+
+    async def wait_for_unload(self, lora_id: str):
+        """
+        Waits until the usage counter for a LoRA adapter reaches zero, indicating that it is no longer in use.
+        This is useful for ensuring that a LoRA adapter can be safely unloaded.
+
+        This method itself is not synchronized, which is safe because it should only be called during LoRA unloading,
+        which itself is guaranteed to be sequential.
+        """
+        assert (
+            lora_id not in self._registry
+        ), "wait_for_unload should only be called after the LoRA adapter has been unregistered. "
+        assert (
+            lora_id in self._counters
+        ), "The LoRA ID should still have a counter if it has been registered before."
+
+        # Wait until no requests are using this LoRA adapter.
+        await self._counters[lora_id].wait_for_zero()
+        del self._counters[lora_id]
+
+    async def get_unregistered_loras(self, lora_name: set[str]):
+        """
+        Returns all LoRA adapters in lora_name that are not found in self._registry.
+        """
+        async with self._registry_lock.writer_lock:
+            unregistered_loras = []
+
+            for name in lora_name:
+                if name in self._registry:
+                    # This counts as a lookup, so we want to update the cache
+                    self._registry.move_to_end(name)
+                else:
+                    unregistered_loras.append(name)
+
+            return unregistered_loras
+
+    async def lru_lora_name(self, exclude_pinned=False):
+        """
+        Returns the least recently used LoRA adapter.
+        If exclude_pinned is True, then return the LRU LoRA adapter that isn't pinned.
+        """
+        async with self._registry_lock.reader_lock:
+            if not exclude_pinned:
+                return next(iter(self._registry), None)
+
+            for lora_name, lora_ref in self._registry.items():
+                if not lora_ref.pinned:
+                    return lora_name
+            else:
+                return None
+
+    def _register_adapter(self, lora_ref: LoRARef):
+        """
+        Internal helper method to register a LoRA adapter.
+        """
+
+        if lora_ref.lora_name in self._registry:
+            raise ValueError(
+                f"LoRA with name {lora_ref.lora_name} already exists. Loaded LoRAs: {self._registry.keys()}"
+            )
+        self._registry[lora_ref.lora_name] = lora_ref
+        self._counters[lora_ref.lora_id] = ConcurrentCounter()
+        return lora_ref
+
+    @property
+    def num_registered_loras(self) -> int:
+        """
+        Returns the total number of LoRA adapters currently registered.
+        """
+        return len(self._registry)
+
+    def get_all_adapters(self) -> Dict[str, LoRARef]:
+        """
+        Returns a dictionary of all registered LoRA adapters.
+
+        Returns:
+            Dict[str, LoRARef]: A dictionary mapping LoRA names to LoRARef objects.
+        """
+        return dict(self._registry)

python/sglang/srt/lora/torch_ops/__init__.py

@@ -0,0 +1,109 @@
+from typing import Optional
+
+import torch
+
+from sglang.srt.lora.utils import LoRABatchInfo
+
+from .graph_lora_ops import (
+    sgemm_lora_a_embedding_graph_fwd,
+    sgemm_lora_a_graph_fwd,
+    sgemm_lora_b_graph_fwd,
+)
+from .lora_ops import sgemm_lora_a_embedding_fwd as sgemm_lora_a_embedding_control_fwd
+from .lora_ops import sgemm_lora_a_fwd as sgemm_lora_a_control_fwd
+from .lora_ops import sgemm_lora_b_fwd as sgemm_lora_b_control_fwd
+
+
+def sgemm_lora_a_embedding_fwd(
+    inputs: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    vocab_size: int,
+) -> torch.Tensor:
+    output: torch.Tensor
+    if batch_info.use_cuda_graph:
+        output = sgemm_lora_a_embedding_graph_fwd(
+            inputs,
+            weights,
+            batch_info.weight_indices,
+            batch_info.seg_lens,
+            batch_info.scalings,
+            vocab_size,
+        )
+    else:
+        output = sgemm_lora_a_embedding_control_fwd(
+            inputs,
+            weights,
+            batch_info.weight_indices_cpu,
+            batch_info.seg_lens_cpu,
+            batch_info.lora_ranks_cpu,
+            batch_info.scalings_cpu,
+            vocab_size,
+        )
+    return output
+
+
+def sgemm_lora_a_fwd(
+    inputs: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    num_slices: int = 1,
+) -> torch.Tensor:
+    output: torch.Tensor
+    if batch_info.use_cuda_graph:
+        output = sgemm_lora_a_graph_fwd(
+            inputs,
+            weights,
+            batch_info.weight_indices,
+            batch_info.seg_lens,
+            batch_info.scalings,
+            num_slices,
+        )
+    else:
+        output = sgemm_lora_a_control_fwd(
+            inputs,
+            weights,
+            batch_info.weight_indices_cpu,
+            batch_info.seg_lens_cpu,
+            batch_info.lora_ranks_cpu,
+            batch_info.scalings_cpu,
+            num_slices,
+        )
+    return output
+
+
+def sgemm_lora_b_fwd(
+    inputs: torch.Tensor,
+    weights: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    slice_offsets: torch.Tensor,
+    base_output: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    output: torch.Tensor
+    if batch_info.use_cuda_graph:
+        output = sgemm_lora_b_graph_fwd(
+            inputs,
+            weights,
+            batch_info.weight_indices,
+            batch_info.seg_lens,
+            slice_offsets,
+            base_output,
+        )
+    else:
+        output = sgemm_lora_b_control_fwd(
+            inputs,
+            weights,
+            batch_info.weight_indices_cpu,
+            batch_info.seg_lens_cpu,
+            batch_info.lora_ranks_cpu,
+            slice_offsets,
+            base_output,
+        )
+    return output
+
+
+__all__ = [
+    "sgemm_lora_a_embedding_fwd",
+    "sgemm_lora_a_fwd",
+    "sgemm_lora_b_fwd",
+]

python/sglang/srt/lora/torch_ops/graph_lora_ops.py

@@ -0,0 +1,120 @@
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+
+
+def sgemm_lora_a_embedding_graph_fwd(
+    inputs: torch.Tensor,
+    weights: torch.Tensor,
+    weight_indices: torch.Tensor,
+    seg_len_tensor: torch.Tensor,
+    scaling_tensor: torch.Tensor,
+    vocab_size: int,
+) -> torch.Tensor:
+    total_seq_len = inputs.shape[0]
+    if weights.numel() == 0:
+        return torch.zeros(total_seq_len, 0, dtype=weights.dtype, device=weights.device)
+
+    num_loras, max_rank, _ = weights.shape
+
+    output = torch.zeros(
+        total_seq_len, max_rank, dtype=weights.dtype, device=weights.device
+    )
+
+    for lora_idx in range(num_loras):
+
+        batch_token_mask = weight_indices[:total_seq_len] == lora_idx
+
+        x_seq = torch.where(batch_token_mask, inputs, 0)
+        w_seq = weights[lora_idx]
+
+        output.add_(
+            scaling_tensor[lora_idx]
+            * torch.where(
+                batch_token_mask.unsqueeze(1), F.embedding(x_seq, w_seq.t()), 0
+            )
+        )
+
+    return output
+
+
+def sgemm_lora_a_graph_fwd(
+    inputs: torch.Tensor,
+    weights: torch.Tensor,
+    weight_indices: torch.Tensor,
+    seg_len_tensor: torch.Tensor,
+    scaling_tensor: torch.Tensor,
+    num_slices: int = 1,
+) -> torch.Tensor:
+    total_seq_len, input_dim = inputs.shape
+    if weights.numel() == 0:
+        return torch.zeros(total_seq_len, 0, dtype=inputs.dtype, device=inputs.device)
+
+    num_loras, weight_out_dim, _ = weights.shape
+    max_rank = weight_out_dim // num_slices
+
+    output = torch.zeros(
+        total_seq_len, num_slices * max_rank, dtype=inputs.dtype, device=inputs.device
+    )
+
+    for lora_idx in range(num_loras):
+
+        batch_token_mask = (weight_indices[:total_seq_len] == lora_idx).unsqueeze(1)
+
+        x_seq = torch.where(batch_token_mask, inputs, 0)
+        w_seq = weights[lora_idx]
+
+        output.add_(scaling_tensor[lora_idx] * torch.mm(x_seq, w_seq.t()))
+
+    return output
+
+
+def sgemm_lora_b_graph_fwd(
+    inputs: torch.Tensor,
+    weights: torch.Tensor,
+    weight_indices: torch.Tensor,
+    seg_len_tensor: torch.Tensor,
+    slice_offsets: torch.Tensor,
+    base_output: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    total_seq_len, input_dim = inputs.shape
+    num_loras, weight_out_dim, _ = weights.shape
+    total_output_dim = slice_offsets[-1].item() if len(slice_offsets) > 0 else 0
+
+    if weights.numel() == 0:
+        return torch.zeros(
+            total_seq_len, total_output_dim, dtype=inputs.dtype, device=inputs.device
+        )
+
+    num_slices = len(slice_offsets) - 1
+    max_rank = input_dim // num_slices
+
+    if base_output is not None:
+        output = base_output
+    else:
+        output = torch.zeros(
+            total_seq_len, total_output_dim, dtype=inputs.dtype, device=inputs.device
+        )
+
+    for lora_idx in range(num_loras):
+
+        batch_token_mask = (weight_indices[:total_seq_len] == lora_idx).unsqueeze(1)
+        inputs_masked = torch.where(batch_token_mask, inputs, 0)
+
+        for slice_idx in range(num_slices):
+            slice_start_input = slice_idx * max_rank
+            slice_end_input = (slice_idx + 1) * max_rank
+
+            slice_start_output = slice_offsets[slice_idx]
+            slice_end_output = slice_offsets[slice_idx + 1]
+
+            x_slice = inputs_masked[..., slice_start_input:slice_end_input]
+            w_slice = weights[
+                lora_idx, slice_start_output:slice_end_output
+            ]  # (slice_dim, max_rank)
+            output[..., slice_start_output:slice_end_output].add_(
+                torch.mm(x_slice, w_slice.t())
+            )
+
+    return output

python/sglang/srt/lora/torch_ops/lora_ops.py

@@ -0,0 +1,146 @@
+from typing import Optional
+
+import torch
+
+
+def sgemm_lora_a_embedding_fwd(
+    inputs: torch.Tensor,
+    weights: torch.Tensor,
+    weight_indices: torch.Tensor,
+    seg_len_tensor: torch.Tensor,
+    lora_ranks: torch.Tensor,
+    scaling_tensor: torch.Tensor,
+    vocab_size: int,
+) -> torch.Tensor:
+    total_seq_len = inputs.shape[0]
+    if weights.numel() == 0:
+        return torch.zeros(total_seq_len, 0, dtype=weights.dtype, device=weights.device)
+
+    num_loras, max_rank, _ = weights.shape
+
+    output = torch.zeros(
+        total_seq_len, max_rank, dtype=weights.dtype, device=weights.device
+    )
+
+    token_offset = 0
+    for lora_idx, seq_len in zip(weight_indices, seg_len_tensor):
+        if seq_len == 0:
+            continue
+
+        rank = lora_ranks[lora_idx]
+        if rank > 0:
+
+            x_seq = inputs[token_offset : token_offset + seq_len]
+            w_seq = weights[lora_idx, :rank]
+
+            result = torch.nn.functional.embedding(x_seq, w_seq.T)
+            output[token_offset : token_offset + seq_len, :rank] = (
+                scaling_tensor[lora_idx].item() * result
+            )
+
+        token_offset += seq_len
+
+    return output
+
+
+def sgemm_lora_a_fwd(
+    inputs: torch.Tensor,
+    weights: torch.Tensor,
+    weight_indices: torch.Tensor,
+    seg_len_tensor: torch.Tensor,
+    lora_ranks: torch.Tensor,
+    scaling_tensor: torch.Tensor,
+    num_slices: int = 1,
+) -> torch.Tensor:
+    total_seq_len, input_dim = inputs.shape
+    if weights.numel() == 0:
+        return torch.zeros(total_seq_len, 0, dtype=inputs.dtype, device=inputs.device)
+
+    num_loras, weight_out_dim, _ = weights.shape
+    max_rank = weight_out_dim // num_slices
+
+    output = torch.zeros(
+        total_seq_len, num_slices * max_rank, dtype=inputs.dtype, device=inputs.device
+    )
+
+    token_offset = 0
+    for lora_idx, seq_len in zip(weight_indices, seg_len_tensor):
+        if seq_len == 0:
+            continue
+
+        rank = lora_ranks[lora_idx]
+        if rank > 0:
+
+            x_seq = inputs[token_offset : token_offset + seq_len]
+            w_seq = weights[lora_idx, : num_slices * rank]
+
+            output[token_offset : token_offset + seq_len, : num_slices * rank].addmm_(
+                x_seq,
+                w_seq.T,
+                beta=0,
+                alpha=scaling_tensor[lora_idx].item(),
+            )
+
+        token_offset += seq_len
+
+    return output
+
+
+def sgemm_lora_b_fwd(
+    inputs: torch.Tensor,
+    weights: torch.Tensor,
+    weight_indices: torch.Tensor,
+    seg_len_tensor: torch.Tensor,
+    lora_ranks: torch.Tensor,
+    slice_offsets: torch.Tensor,
+    base_output: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    total_seq_len, _ = inputs.shape
+    num_loras, weight_out_dim, _ = weights.shape
+    total_output_dim = slice_offsets[-1].item() if len(slice_offsets) > 0 else 0
+
+    if weights.numel() == 0:
+        return torch.zeros(
+            total_seq_len, total_output_dim, dtype=inputs.dtype, device=inputs.device
+        )
+
+    num_slices = len(slice_offsets) - 1
+
+    if base_output is not None:
+        output = base_output
+    else:
+        output = torch.zeros(
+            total_seq_len, total_output_dim, dtype=inputs.dtype, device=inputs.device
+        )
+
+    token_offset = 0
+    for lora_idx, seq_len in zip(weight_indices, seg_len_tensor):
+        if seq_len == 0:
+            continue
+
+        rank = lora_ranks[lora_idx]
+        if rank > 0:
+
+            for slice_idx in range(num_slices):
+                slice_start_input = slice_idx * rank
+                slice_end_input = (slice_idx + 1) * rank
+
+                slice_start_output = slice_offsets[slice_idx]
+                slice_end_output = slice_offsets[slice_idx + 1]
+
+                x_slice = inputs[
+                    token_offset : token_offset + seq_len,
+                    slice_start_input:slice_end_input,
+                ]  # (seq_len, rank)
+                w_slice = weights[
+                    lora_idx, slice_start_output:slice_end_output, :rank
+                ]  # (slice_dim, rank)
+
+                output[
+                    token_offset : token_offset + seq_len,
+                    slice_start_output:slice_end_output,
+                ].addmm_(x_slice, w_slice.T)
+
+        token_offset += seq_len
+
+    return output

python/sglang/srt/lora/trtllm_lora_temp/__init__.py

@@ -0,0 +1,224 @@
+"""Two-stream LoRA overlap (O1 + O7 + O8 + O9) — installed as a monkey-patch.
+
+Activates when env ``SGLANG_LORA_TWO_STREAM=1``. Triggered exactly once via
+:func:`install_two_stream_overrides` (called at end of ``sglang/srt/lora/layers.py``).
+
+When enabled, these call sites are redirected to side-stream-overlapped versions
+defined entirely in this package:
+
+  * ``QKVParallelLinearWithLoRA.forward``  → :mod:`.attention.qkv_proj_lora_forward`
+  * ``RowParallelLinearWithLoRA.forward``  → :mod:`.attention.row_parallel_lora_forward`
+  * ``MergedColumnParallelLinearWithLoRA.forward`` → :mod:`.merged_column.merged_column_lora_forward`
+  * ``fused_experts_none_to_experimental_sgl_trtllm_fp8_lora`` →
+    :mod:`.moe_overlap.fused_experts_none_to_experimental_sgl_trtllm_fp8_lora_two_stream`
+
+When disabled (env unset), ``install_two_stream_overrides`` is a no-op and all
+the original functions / methods in ``sglang/srt/lora/layers.py`` and
+``sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py`` run unchanged.
+
+Per-batch gating still happens inside the patched callables — they fall back
+to the saved-original implementation for non-decode batches (token count above
+``SGLANG_TWO_STREAM_MAX_TOKENS`` default 256), so prefill stays on the serial
+path even with the patch installed.
+"""
+
+from typing import Callable, Optional
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
+
+
+def is_two_stream_active(x: torch.Tensor) -> bool:
+    """Per-batch gate (two-stream now always-on). True iff batch is decode-shaped (<= SGLANG_TWO_STREAM_MAX_TOKENS)."""
+    return x.shape[0] <= lora_envs.SGLANG_TWO_STREAM_MAX_TOKENS.get()
+
+
+def get_lora_side_stream() -> torch.cuda.Stream:
+    """Lazily allocate a single shared LoRA side stream.
+
+    Within one decode layer the three sites (qkv → attn → o_proj → moe_gate_up)
+    run sequentially, so one stream suffices and avoids extra graph-capture
+    nodes from per-site streams.
+    """
+    from sglang.srt.runtime_context import get_stream
+
+    return get_stream("lora_side")
+
+
+def init_lora_two_stream_resources(device: Optional[torch.device] = None) -> None:
+    """Eagerly create the side stream before cuda-graph capture begins.
+
+    ``torch.cuda.Stream()`` is a driver call that must not run inside a
+    cuda-graph capture region. Since :func:`get_lora_side_stream` is otherwise
+    lazy, the first eligible decode forward would create it — which can fall
+    inside capture if warmup didn't happen to exercise a two-stream batch.
+    Calling this from a pre-capture hook pins creation to init/warmup on the
+    correct device.
+    """
+    if device is not None:
+        with torch.cuda.device(device):
+            get_lora_side_stream()
+    else:
+        get_lora_side_stream()
+
+
+def lora_overlap_alloc_stream() -> Optional[torch.cuda.Stream]:
+    """Stream to allocate side-stream LoRA-shrink OUTPUT buffers on, or None for default behavior.
+
+    A buffer allocated *inside* ``with torch.cuda.stream(side)`` is tagged to the side stream, so the
+    caching allocator may free/reuse it on the side stream's schedule — before the MAIN stream (the real
+    consumer, via the LoRA-B expand) is done. Under cuda-graph replay that's a premature-reuse WAR ->
+    qwen3.5 mamba decode garbage. With ``SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC`` this returns the MAIN
+    stream so the op allocates the output on the consumer stream (like the MoE O1 ``gate_up_delta``),
+    making a single shared side stream graph-safe. Call on the MAIN stream BEFORE forking to the side.
+    """
+    if lora_envs.SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC.get():
+        return torch.cuda.current_stream()
+    return None
+
+
+# References to the original implementations, captured at install time so the
+# patched callables can defer to them for non-decode batches.
+_ORIGINAL_QKV_FORWARD: Optional[Callable] = None
+_ORIGINAL_ROW_FORWARD: Optional[Callable] = None
+_ORIGINAL_MERGED_FORWARD: Optional[Callable] = None
+_ORIGINAL_COLUMN_FORWARD: Optional[Callable] = None
+_ORIGINAL_REPLICATED_FORWARD: Optional[Callable] = None
+_ORIGINAL_MOE_LORA_FUNC: Optional[Callable] = None
+_ORIGINAL_FP4_MOE_LORA_FUNC: Optional[Callable] = None
+_ORIGINAL_BF16_MOE_LORA_FUNC: Optional[Callable] = None
+_INSTALLED: bool = False
+
+
+def get_original_qkv_forward() -> Callable:
+    return _ORIGINAL_QKV_FORWARD
+
+
+def get_original_row_forward() -> Callable:
+    return _ORIGINAL_ROW_FORWARD
+
+
+def get_original_merged_column_forward() -> Callable:
+    return _ORIGINAL_MERGED_FORWARD
+
+
+def get_original_column_forward() -> Callable:
+    return _ORIGINAL_COLUMN_FORWARD
+
+
+def get_original_replicated_forward() -> Callable:
+    return _ORIGINAL_REPLICATED_FORWARD
+
+
+def get_original_moe_lora_func() -> Callable:
+    return _ORIGINAL_MOE_LORA_FUNC
+
+
+def get_original_fp4_moe_lora_func() -> Callable:
+    return _ORIGINAL_FP4_MOE_LORA_FUNC
+
+
+def get_original_bf16_moe_lora_func() -> Callable:
+    return _ORIGINAL_BF16_MOE_LORA_FUNC
+
+
+def install_two_stream_overrides() -> None:
+    """Install the side-stream overlapped overrides if ``SGLANG_LORA_TWO_STREAM=1``.
+
+    Idempotent: subsequent calls are a no-op. Patches:
+
+      1. ``QKVParallelLinearWithLoRA.forward`` (O7 — qkv LoRA shrink overlap)
+      2. ``RowParallelLinearWithLoRA.forward`` (O8 — o_proj LoRA shrink overlap)
+      3. ``MergedColumnParallelLinearWithLoRA.forward`` (O9 — merged-column LoRA
+         shrink overlap: dense gate_up + mamba in_proj_qkvz)
+      4. ``lora_dispatch.fused_experts_none_to_experimental_sgl_trtllm_fp8_lora``
+         (O1 — MoE gate_up LoRA overlap), plus its fp4 (O1-fp4) and bf16
+         (O1-bf16) siblings
+
+    The saved originals are exposed via :func:`get_original_qkv_forward`,
+    :func:`get_original_row_forward`, :func:`get_original_moe_lora_func` so the
+    new versions can fall back when their per-batch gate says single-stream.
+    """
+    global _INSTALLED, _ORIGINAL_QKV_FORWARD, _ORIGINAL_ROW_FORWARD, _ORIGINAL_MERGED_FORWARD, _ORIGINAL_COLUMN_FORWARD, _ORIGINAL_REPLICATED_FORWARD, _ORIGINAL_MOE_LORA_FUNC, _ORIGINAL_FP4_MOE_LORA_FUNC, _ORIGINAL_BF16_MOE_LORA_FUNC
+
+    if _INSTALLED:
+        return
+
+    from sglang.srt.lora.layers import (
+        ColumnParallelLinearWithLoRA,
+        MergedColumnParallelLinearWithLoRA,
+        QKVParallelLinearWithLoRA,
+        ReplicatedLinearWithLoRA,
+        RowParallelLinearWithLoRA,
+    )
+    from sglang.srt.lora.trtllm_lora_temp.attention import (
+        column_parallel_lora_forward,
+        qkv_proj_lora_forward,
+        replicated_lora_forward,
+        row_parallel_lora_forward,
+    )
+    from sglang.srt.lora.trtllm_lora_temp.merged_column import (
+        merged_column_lora_forward,
+    )
+
+    # Capture all originals before patching: QKV / MergedColumn subclass
+    # ColumnParallel, so the plain-Column O10 patch must not clobber the
+    # subclasses' own (3-/2-slice) forwards captured here as their fallbacks.
+    _ORIGINAL_QKV_FORWARD = QKVParallelLinearWithLoRA.forward
+    _ORIGINAL_ROW_FORWARD = RowParallelLinearWithLoRA.forward
+    _ORIGINAL_MERGED_FORWARD = MergedColumnParallelLinearWithLoRA.forward
+    _ORIGINAL_COLUMN_FORWARD = ColumnParallelLinearWithLoRA.forward
+    _ORIGINAL_REPLICATED_FORWARD = ReplicatedLinearWithLoRA.forward
+    QKVParallelLinearWithLoRA.forward = qkv_proj_lora_forward
+    RowParallelLinearWithLoRA.forward = row_parallel_lora_forward
+    MergedColumnParallelLinearWithLoRA.forward = merged_column_lora_forward
+    # O10 (MLA q_b_proj / kv_b_proj) + O11 (MLA fused_qkv_a_proj_with_mqa).
+    ColumnParallelLinearWithLoRA.forward = column_parallel_lora_forward
+    ReplicatedLinearWithLoRA.forward = replicated_lora_forward
+
+    import sglang.srt.lora.trtllm_lora_temp.lora_dispatch as ft
+    from sglang.srt.lora.trtllm_lora_temp.moe_overlap import (
+        fused_experts_none_to_experimental_sgl_trtllm_bf16_lora_two_stream,
+        fused_experts_none_to_experimental_sgl_trtllm_fp4_lora_two_stream,
+        fused_experts_none_to_experimental_sgl_trtllm_fp8_lora_two_stream,
+    )
+
+    # O1 (FP8 Qwen) + O1-fp4 (NVFP4 Kimi) + O1-bf16 (unquantized Qwen): MoE gate_up
+    # LoRA overlap. Each patched fn falls back to its saved single-stream original
+    # for non-decode batches.
+    _ORIGINAL_MOE_LORA_FUNC = ft.fused_experts_none_to_experimental_sgl_trtllm_fp8_lora
+    _ORIGINAL_FP4_MOE_LORA_FUNC = (
+        ft.fused_experts_none_to_experimental_sgl_trtllm_fp4_lora
+    )
+    _ORIGINAL_BF16_MOE_LORA_FUNC = (
+        ft.fused_experts_none_to_experimental_sgl_trtllm_bf16_lora
+    )
+    ft.fused_experts_none_to_experimental_sgl_trtllm_fp8_lora = (
+        fused_experts_none_to_experimental_sgl_trtllm_fp8_lora_two_stream
+    )
+    ft.fused_experts_none_to_experimental_sgl_trtllm_fp4_lora = (
+        fused_experts_none_to_experimental_sgl_trtllm_fp4_lora_two_stream
+    )
+    ft.fused_experts_none_to_experimental_sgl_trtllm_bf16_lora = (
+        fused_experts_none_to_experimental_sgl_trtllm_bf16_lora_two_stream
+    )
+
+    _INSTALLED = True
+
+
+__all__ = [
+    "is_two_stream_active",
+    "get_lora_side_stream",
+    "init_lora_two_stream_resources",
+    "get_original_qkv_forward",
+    "get_original_row_forward",
+    "get_original_merged_column_forward",
+    "get_original_column_forward",
+    "get_original_replicated_forward",
+    "get_original_moe_lora_func",
+    "get_original_fp4_moe_lora_func",
+    "get_original_bf16_moe_lora_func",
+    "install_two_stream_overrides",
+]

python/sglang/srt/lora/trtllm_lora_temp/attention.py

@@ -0,0 +1,253 @@
+"""Two-stream attention LoRA forward implementations (O7 + O8).
+
+These are monkey-patched onto :class:`QKVParallelLinearWithLoRA` and
+:class:`RowParallelLinearWithLoRA` by
+:func:`sglang.srt.lora.trtllm_lora_temp.install_two_stream_overrides` when
+``SGLANG_LORA_TWO_STREAM=1``. The saved-original forward methods are
+preserved and called for batches where two-stream isn't active.
+"""
+
+import torch
+
+from sglang.srt.distributed import (
+    split_tensor_along_last_dim,
+    tensor_model_parallel_all_gather,
+    tensor_model_parallel_all_reduce,
+)
+from sglang.srt.layers.moe.utils import should_skip_mlp_all_reduce
+from sglang.srt.lora.trtllm_lora_temp import (
+    get_lora_side_stream,
+    get_original_column_forward,
+    get_original_qkv_forward,
+    get_original_replicated_forward,
+    get_original_row_forward,
+    is_two_stream_active,
+    lora_overlap_alloc_stream,
+)
+from sglang.srt.runtime_context import get_parallel
+
+
+def qkv_proj_lora_forward(self, input_: torch.Tensor):
+    """O7 — side-stream LoRA-A shrink ‖ base qkv_proj GEMM.
+
+    The shrink reads ``input_`` and the LoRA-A weights — same input as the
+    base GEMM, no write conflict. The expand needs the shrink intermediate
+    AND base_output, so it runs after the rejoin on the main stream.
+    """
+    if not self.set_lora or not is_two_stream_active(input_):
+        return get_original_qkv_forward()(self, input_)
+
+    from sglang.kernels.ops.gemm.trtllm_lora_temp.qkv_lora_b import qkv_lora_b_fwd
+    from sglang.kernels.ops.gemm.trtllm_lora_temp.sgemm_lora_a import sgemm_lora_a_fwd
+
+    bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None
+    side_stream = get_lora_side_stream()
+    # sgemm_info is host-side (LoRABatchInfo); compute once, share both calls.
+    sgemm_info = self.lora_backend._sgemm_info()
+
+    _alloc = lora_overlap_alloc_stream()  # capture MAIN stream here (before the fork)
+    side_stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(side_stream):
+        shrink_intermediate = sgemm_lora_a_fwd(
+            input_, self.A_buffer_qkv, sgemm_info, stack_num=3, out_alloc_stream=_alloc
+        )
+
+    # Base qkv_proj GEMM on main, concurrent with the side-stream shrink.
+    output_parallel = self.base_layer.quant_method.apply(self.base_layer, input_, bias)
+
+    # Rejoin: expand reads both side-produced shrink_intermediate and base_output.
+    torch.cuda.current_stream().wait_stream(side_stream)
+    output_parallel = qkv_lora_b_fwd(
+        shrink_intermediate,
+        self.B_buffer_qkv,
+        sgemm_info,
+        self.output_offset,
+        self.max_qkv_out_dim,
+        output_parallel,
+        n_slices=3,
+    )
+
+    if self.base_layer.gather_output:
+        output = tensor_model_parallel_all_gather(output_parallel)
+    else:
+        output = output_parallel
+    output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
+    return output, output_bias
+
+
+def row_parallel_lora_forward(
+    self, input_: torch.Tensor, skip_all_reduce: bool = False, forward_batch=None
+):
+    """O8 — side-stream LoRA-A shrink ‖ base row-parallel (o_proj) GEMM.
+
+    Mirrors O7 but the row-parallel context adds: input split per TP rank
+    (when not already parallel), bias on rank 0 only, optional cross-rank
+    all-reduce on both base output and lora_a intermediate when reducing.
+
+    Falls back to the saved-original :meth:`forward` for non-decode batches
+    or when LoRA isn't set on this layer.
+    """
+    # We need ``input_parallel`` to gate the per-batch decode check (its
+    # token-count drives the threshold, not the unsplit ``input_``).
+    if self.base_layer.input_is_parallel:
+        input_parallel = input_
+    else:
+        tp_rank = get_parallel().tp_rank
+        splitted_input = split_tensor_along_last_dim(
+            input_, num_partitions=self.base_layer.tp_size
+        )
+        input_parallel = splitted_input[tp_rank].contiguous()
+
+    if not self.set_lora or not is_two_stream_active(input_parallel):
+        return get_original_row_forward()(self, input_, skip_all_reduce, forward_batch)
+
+    bias_ = (
+        None
+        if (self.base_layer.tp_rank > 0 or self.base_layer.skip_bias_add)
+        else self.base_layer.bias
+    )
+
+    side_stream = get_lora_side_stream()
+    _alloc = lora_overlap_alloc_stream()  # capture MAIN stream here (before the fork)
+    side_stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(side_stream):
+        lora_a_output = self.lora_backend.run_lora_a_sgemm(
+            input_parallel, self.A_buffer, out_alloc_stream=_alloc
+        )
+
+    # Base row-parallel GEMM on main, concurrent with the side-stream shrink.
+    output_parallel = self.base_layer.quant_method.apply(
+        self.base_layer, input_parallel, bias=bias_
+    )
+
+    torch.cuda.current_stream().wait_stream(side_stream)
+
+    should_reduce = (
+        self.base_layer.reduce_results
+        and self.base_layer.tp_size > 1
+        and not skip_all_reduce
+        and not should_skip_mlp_all_reduce()
+    )
+
+    if should_reduce:
+        output_ = tensor_model_parallel_all_reduce(output_parallel)
+        lora_a_output = tensor_model_parallel_all_reduce(lora_a_output)
+        output_ = self.lora_backend.run_lora_b_sgemm(
+            x=lora_a_output,
+            weights=self.B_buffer,
+            output_offset=self.output_offset,
+            output_offset_cpu=self.output_offset_cpu,
+            base_output=output_,
+        )
+    else:
+        # Two-stream already produced lora_a_output on the side stream; finish
+        # the LoRA with just the expand atomic-add against output_parallel.
+        output_parallel = self.lora_backend.run_lora_b_sgemm(
+            x=lora_a_output,
+            weights=self.B_buffer,
+            output_offset=self.output_offset,
+            output_offset_cpu=self.output_offset_cpu,
+            base_output=output_parallel,
+        )
+        output_ = output_parallel
+
+    output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
+    return output_, output_bias
+
+
+def column_parallel_lora_forward(self, input_: torch.Tensor):
+    """O10 — side-stream LoRA-A shrink ‖ base ColumnParallel GEMM.
+
+    Covers DeepSeek/Kimi MLA's ``q_b_proj`` / ``kv_b_proj`` (plain
+    :class:`ColumnParallelLinearWithLoRA` — the merged/QKV subclasses keep their
+    own O9/O7 overrides). The shrink reads ``input_`` (same as the base GEMM, no
+    write conflict); the expand needs the shrink intermediate AND base_output,
+    so it runs after the rejoin. Byte-identical to the saved-original forward
+    for non-decode batches or when LoRA isn't set on this layer.
+    """
+    if not self.set_lora or not is_two_stream_active(input_):
+        return get_original_column_forward()(self, input_)
+
+    bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None
+    side_stream = get_lora_side_stream()
+
+    _alloc = lora_overlap_alloc_stream()  # capture MAIN stream here (before the fork)
+    side_stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(side_stream):
+        lora_a_output = self.lora_backend.run_lora_a_sgemm(
+            input_, self.A_buffer, out_alloc_stream=_alloc
+        )
+
+    # Base ColumnParallel GEMM on main, concurrent with the side-stream shrink.
+    output_parallel = self.base_layer.quant_method.apply(self.base_layer, input_, bias)
+
+    # Rejoin: expand reads both the side-produced shrink and base_output.
+    torch.cuda.current_stream().wait_stream(side_stream)
+    output_parallel = self.lora_backend.run_lora_b_sgemm(
+        x=lora_a_output,
+        weights=self.B_buffer,
+        output_offset=self.output_offset,
+        output_offset_cpu=self.output_offset_cpu,
+        base_output=output_parallel,
+    )
+
+    if self.base_layer.gather_output:
+        output = tensor_model_parallel_all_gather(output_parallel)
+    else:
+        output = output_parallel
+    output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
+    return output, output_bias
+
+
+def replicated_lora_forward(self, x: torch.Tensor):
+    """O11 — side-stream LoRA-A shrink ‖ base ReplicatedLinear GEMM.
+
+    Covers DeepSeek/Kimi MLA's ``fused_qkv_a_proj_with_mqa``
+    (:class:`ReplicatedLinearWithLoRA`, no TP sharding). Handles both the
+    single-projection (``first_output_dim == 0``) and the fused q_a+kv_a
+    (``> 0``, ``n_slices=2``) cases, splitting the same A-shrink / B-expand the
+    backend's ``run_qkv_lora`` composes internally — A on the side stream, B on
+    the main after the rejoin. Falls back to the saved-original otherwise.
+    """
+    if not self.set_lora or not is_two_stream_active(x):
+        return get_original_replicated_forward()(self, x)
+
+    bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None
+    side_stream = get_lora_side_stream()
+    first_dim = self.first_output_dim
+
+    _alloc = lora_overlap_alloc_stream()  # capture MAIN stream here (before the fork)
+    side_stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(side_stream):
+        lora_a_output = self.lora_backend.run_lora_a_sgemm(
+            x,
+            self.A_buffer,
+            stack_num=(2 if first_dim > 0 else 1),
+            out_alloc_stream=_alloc,
+        )
+
+    # Base ReplicatedLinear GEMM on main, concurrent with the side-stream shrink.
+    output = self.base_layer.quant_method.apply(self.base_layer, x, bias)
+
+    torch.cuda.current_stream().wait_stream(side_stream)
+    if first_dim == 0:
+        output = self.lora_backend.run_lora_b_sgemm(
+            x=lora_a_output,
+            weights=self.B_buffer,
+            output_offset=self._output_offset,
+            base_output=output,
+        )
+    else:
+        from sglang.kernels.ops.gemm.trtllm_lora_temp.qkv_lora_b import qkv_lora_b_fwd
+
+        output = qkv_lora_b_fwd(
+            lora_a_output,
+            self.B_buffer,
+            self.lora_backend._sgemm_info(),
+            self._output_offset,
+            self._max_out_dim,
+            output,
+            n_slices=2,
+        )
+    output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
+    return output, output_bias

python/sglang/srt/lora/trtllm_lora_temp/deepseek_mla_correction.py

@@ -0,0 +1,238 @@
+"""LoRA correction for absorbed-MLA ``kv_b_proj``.
+
+The absorbed-MLA path in ``DeepseekV2AttentionMLA`` bypasses
+``kv_b_proj.forward()`` and folds the K/V contribution into two BMMs against
+the pre-computed ``w_kc`` / ``w_vc`` weights, so a standard
+``ColumnParallelLinearWithLoRA`` wrapper would never see the activations and
+the LoRA delta would silently be dropped. These helpers inject the missing
+delta on top of the absorbed intermediates via the SGMM-style Triton kernels
+in ``triton_ops/kv_b_lora_absorbed.py``.
+
+Used from ``deepseek_common/attention_forward_methods/forward_mla.py``. Call
+sites should gate the call with :func:`is_kv_b_lora_active` so non-LoRA
+forwards take a single ``getattr`` and skip the helper entirely.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional, Tuple
+
+import torch
+
+# The four step kernels live in triton_ops; importing it pulls the LoRA kernel
+# modules (and specialized_expand) into the process. They are only ever reached
+# after _get_state returns a non-None state (a kv_b LoRA adapter is wrapped), so
+# defer the import to that success path: a no-LoRA forward never imports it here.
+step_a_q_fwd = step_a_v_fwd = step_b_q_fwd = step_b_v_fwd = None
+
+
+def _ensure_step_kernels() -> None:
+    global step_a_q_fwd, step_a_v_fwd, step_b_q_fwd, step_b_v_fwd
+    if step_a_q_fwd is None:
+        from sglang.kernels.ops.gemm.trtllm_lora_temp.kv_b_lora_absorbed import (
+            step_a_q_fwd as _aq,
+        )
+        from sglang.kernels.ops.gemm.trtllm_lora_temp.kv_b_lora_absorbed import (
+            step_a_v_fwd as _av,
+        )
+        from sglang.kernels.ops.gemm.trtllm_lora_temp.kv_b_lora_absorbed import (
+            step_b_q_fwd as _bq,
+        )
+        from sglang.kernels.ops.gemm.trtllm_lora_temp.kv_b_lora_absorbed import (
+            step_b_v_fwd as _bv,
+        )
+
+        step_a_q_fwd, step_a_v_fwd, step_b_q_fwd, step_b_v_fwd = _aq, _av, _bq, _bv
+
+
+if TYPE_CHECKING:
+    from sglang.srt.lora.utils import LoRABatchInfo
+    from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
+
+
+def is_kv_b_lora_active(attn_module: DeepseekV2AttentionMLA) -> bool:
+    """Cheap precondition check used at call sites in the attention forward
+    to skip the entire LoRA-correction path when no ``kv_b_proj`` adapter is
+    wrapped on this module (the common case)."""
+    return getattr(attn_module.kv_b_proj, "set_lora", False)
+
+
+def _get_state(
+    attn_module: DeepseekV2AttentionMLA,
+) -> Optional[Tuple[torch.Tensor, torch.Tensor, LoRABatchInfo]]:
+    if not is_kv_b_lora_active(attn_module):
+        return None
+    if not hasattr(attn_module.kv_b_proj, "A_buffer"):
+        return None
+    lora_backend = attn_module.kv_b_proj.lora_backend
+    if not hasattr(lora_backend, "batch_info"):
+        return None
+    batch_info = lora_backend.batch_info
+    if batch_info is None:
+        return None
+
+    # Triton backend exposes _sgemm_info() to group decode-shape repeats of
+    # the same adapter; csgmv-style backends just expose batch_info directly.
+    sgemm_info = getattr(lora_backend, "_sgemm_info", None)
+    if callable(sgemm_info):
+        batch_info = sgemm_info()
+    # Non-None state ⇒ a kv_b adapter is active here; load the step kernels now
+    # (cached after the first active forward). No-LoRA forwards return above and
+    # never import triton_ops.
+    _ensure_step_kernels()
+    return attn_module.kv_b_proj.A_buffer, attn_module.kv_b_proj.B_buffer, batch_info
+
+
+def apply_q_correction(
+    attn_module: DeepseekV2AttentionMLA,
+    q_nope: torch.Tensor,
+    q_nope_out: torch.Tensor,
+) -> torch.Tensor:
+    """LoRA correction for the absorbed ``q_nope @ w_kc`` path.
+
+    Computes ``q_nope_out += q_nope @ B_kc @ A * scaling`` per token, per
+    active LoRA slot via two SGMM-style Triton kernels. Factored along the
+    LoRA-A/B boundary so we never materialise ``B @ A`` (~268M FMAs per layer
+    per slot in the naive implementation)::
+
+      step A_q : ``(S,H,qk_nope) @ B_kc[slot, h] (qk_nope, rank) -> (S,H,rank)``
+      step B_q : ``(S,H,rank)    @ A[slot] (rank, kv_lora_rank)  -> += q_nope_out``
+    """
+    state = _get_state(attn_module)
+    if state is None:
+        return q_nope_out
+    A_buf, B_buf, batch_info = state
+
+    full_K_per_head = attn_module.qk_nope_head_dim + attn_module.v_head_dim
+    q_lora_a = step_a_q_fwd(q_nope, B_buf, batch_info, full_K_per_head)
+    return step_b_q_fwd(q_lora_a, A_buf, batch_info, q_nope_out)
+
+
+def apply_v_correction(
+    attn_module: DeepseekV2AttentionMLA,
+    attn_output: torch.Tensor,
+    attn_bmm_flat: torch.Tensor,
+) -> torch.Tensor:
+    """LoRA correction for the absorbed ``attn_output @ w_vc`` path.
+
+    Computes ``attn_bmm_flat += attn_output @ A.T @ B_vc.T * scaling`` per
+    token, per active LoRA slot. ``attn_bmm_flat`` is the flat
+    ``(S, H*v_head_dim)`` view of the absorbed BMM result; we pass strides
+    matching the implicit ``(S, H, v_head_dim)`` layout to step B_v.
+    """
+    state = _get_state(attn_module)
+    if state is None:
+        return attn_bmm_flat
+    A_buf, B_buf, batch_info = state
+
+    attn_lora_a = step_a_v_fwd(attn_output, A_buf, batch_info)
+    base_view = attn_bmm_flat.view(
+        -1, attn_module.num_local_heads, attn_module.v_head_dim
+    )
+    step_b_v_fwd(
+        attn_lora_a,
+        B_buf,
+        batch_info,
+        base_view,
+        attn_module.qk_nope_head_dim,
+        attn_module.v_head_dim,
+    )
+    return attn_bmm_flat
+
+
+# ---------------------------------------------------------------------------
+# Two-stream overlap (O12) for the absorbed kv_b correction.
+#
+# Each correction factors into an input-only A-step (reads q_nope / attn_output,
+# independent of the absorbed bmm output) and a B-step that adds into that bmm
+# output. ``*_prepare`` forks the A-step onto the shared LoRA side stream so it
+# overlaps the main-stream ``q_nope @ w_kc`` / ``attn_output @ w_vc`` bmm;
+# ``*_apply`` rejoins and runs the B-step.
+#
+# Gated by ``SGLANG_LORA_TWO_STREAM`` (decode batches only) via
+# ``is_two_stream_active``. When inactive, ``*_prepare`` returns None and
+# ``*_apply`` falls back to the serial ``apply_*_correction`` (or a no-op when no
+# kv_b adapter is wrapped), so the deepseek call sites stay byte-identical with
+# two-stream off. Same fork/join (``wait_stream``) idiom as the O7/O8 attention
+# overrides — cuda-graph-capture safe.
+# ---------------------------------------------------------------------------
+
+
+def _kv_b_two_stream_state(attn_module, x):
+    from sglang.srt.lora.trtllm_lora_temp import (
+        get_lora_side_stream,
+        is_two_stream_active,
+    )
+
+    if not is_two_stream_active(x):
+        return None
+    state = _get_state(attn_module)
+    if state is None:
+        return None
+    A_buf, B_buf, batch_info = state
+    return A_buf, B_buf, batch_info, get_lora_side_stream()
+
+
+def kv_b_lora_q_prepare(attn_module, q_nope):
+    """Fork the q-correction A-step onto the side stream (``step_a_q`` reads only
+    ``q_nope``) so it overlaps the main-stream ``q_nope @ w_kc`` bmm. Returns a
+    handle for :func:`kv_b_lora_q_apply`, or None when two-stream is inactive."""
+    st = _kv_b_two_stream_state(attn_module, q_nope)
+    if st is None:
+        return None
+    A_buf, B_buf, batch_info, side_stream = st
+    full_K_per_head = attn_module.qk_nope_head_dim + attn_module.v_head_dim
+    side_stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(side_stream):
+        q_lora_a = step_a_q_fwd(q_nope, B_buf, batch_info, full_K_per_head)
+    return q_lora_a, A_buf, batch_info, side_stream
+
+
+def kv_b_lora_q_apply(attn_module, q_nope, q_nope_out, handle):
+    """Finish the q-correction: two-stream (rejoin + B-step) when ``handle`` is
+    set, else the serial correction, else a no-op. Single call replacing the
+    ``if is_kv_b_lora_active: apply_q_correction`` at the call site."""
+    if handle is not None:
+        q_lora_a, A_buf, batch_info, side_stream = handle
+        torch.cuda.current_stream().wait_stream(side_stream)
+        return step_b_q_fwd(q_lora_a, A_buf, batch_info, q_nope_out)
+    if is_kv_b_lora_active(attn_module):
+        return apply_q_correction(attn_module, q_nope, q_nope_out)
+    return q_nope_out
+
+
+def kv_b_lora_v_prepare(attn_module, attn_output):
+    """Fork the v-correction A-step onto the side stream (``step_a_v`` reads only
+    ``attn_output``) so it overlaps the main-stream ``attn_output @ w_vc`` bmm.
+    Returns a handle for :func:`kv_b_lora_v_apply`, or None when inactive."""
+    st = _kv_b_two_stream_state(attn_module, attn_output)
+    if st is None:
+        return None
+    A_buf, B_buf, batch_info, side_stream = st
+    side_stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(side_stream):
+        attn_lora_a = step_a_v_fwd(attn_output, A_buf, batch_info)
+    return attn_lora_a, B_buf, batch_info, side_stream
+
+
+def kv_b_lora_v_apply(attn_module, attn_output, attn_bmm_flat, handle):
+    """Finish the v-correction: two-stream (rejoin + B-step) when ``handle`` is
+    set, else the serial correction, else a no-op."""
+    if handle is not None:
+        attn_lora_a, B_buf, batch_info, side_stream = handle
+        torch.cuda.current_stream().wait_stream(side_stream)
+        base_view = attn_bmm_flat.view(
+            -1, attn_module.num_local_heads, attn_module.v_head_dim
+        )
+        step_b_v_fwd(
+            attn_lora_a,
+            B_buf,
+            batch_info,
+            base_view,
+            attn_module.qk_nope_head_dim,
+            attn_module.v_head_dim,
+        )
+        return attn_bmm_flat
+    if is_kv_b_lora_active(attn_module):
+        return apply_v_correction(attn_module, attn_output, attn_bmm_flat)
+    return attn_bmm_flat

python/sglang/srt/lora/trtllm_lora_temp/environ.py

@@ -0,0 +1,135 @@
+"""Local env registry for the experimental TRT-LLM LoRA fast path.
+
+Every flag here is gated by the single global master switch
+``SGLANG_EXPERIMENTAL_LORA_OPTI`` (defined in ``sglang.srt.environ``). When the master
+switch is OFF (the default), every flag reads ``False`` (its default is
+suppressed), so the no-LoRA path, other MoE backends, and the default
+(non-experimental) LoRA path are byte-identical to upstream.
+
+Keeping these flags out of the global ``Envs`` class is deliberate: the only
+sglang-global addition for this feature is ``SGLANG_EXPERIMENTAL_LORA_OPTI``; all the
+fine-grained opt switches live here, next to the code that consumes them.
+
+Default policy (applies only when ``SGLANG_EXPERIMENTAL_LORA_OPTI=1``):
+  * **common** flags — used by BOTH the qwen3.5 (FP8) and kimi (NVFP4) configs —
+    default ``True`` so they need not be repeated on every launch command.
+  * **non-shared** flags default ``False`` and must be set explicitly in the
+    launch environment for the model that needs them.
+
+C++-getenv flags (read via ``getenv`` in the JIT launcher, not in Python) are
+listed at the bottom for documentation only; set them in the launch env.
+"""
+
+import os
+
+from sglang.srt.environ import envs
+
+
+def experimental_lora_enabled() -> bool:
+    """Master gate. All flags below are forced off unless this is set."""
+    return envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get()
+
+
+_TRUE = {"1", "true", "yes", "on"}
+
+
+class _GatedBool:
+    def __init__(self, name: str, default: bool):
+        self._name = name
+        self._default = default
+
+    def get(self) -> bool:
+        if not experimental_lora_enabled():
+            return False
+        raw = os.environ.get(self._name)
+        if raw is None:
+            return self._default
+        return raw.strip().lower() in _TRUE
+
+
+class _GatedInt:
+    def __init__(self, name: str, default: int):
+        self._name = name
+        self._default = default
+
+    def get(self) -> int:
+        # Only consulted on the experimental path; return the default otherwise.
+        if not experimental_lora_enabled():
+            return self._default
+        raw = os.environ.get(self._name)
+        return int(raw) if raw is not None else self._default
+
+
+class _LoraEnvs:
+    # ---- common (qwen3.5 ∩ kimi): default True when experimental is on ----
+    SGLANG_ENABLE_LORA_SHRINK_SPLIT_K = _GatedBool(
+        "SGLANG_ENABLE_LORA_SHRINK_SPLIT_K", True
+    )
+    SGLANG_OPT_LORA_FUSED_MERGED_ALIGN = _GatedBool(
+        "SGLANG_OPT_LORA_FUSED_MERGED_ALIGN", True
+    )
+    SGLANG_OPT_LORA_FUSED_TOPK_PACK = _GatedBool(
+        "SGLANG_OPT_LORA_FUSED_TOPK_PACK", True
+    )
+    SGLANG_OPT_LORA_QKV_B_STORE = _GatedBool("SGLANG_OPT_LORA_QKV_B_STORE", True)
+    # F1-①: prefill routing reuse — unify the A (shrink) stage's routing BLOCK_SIZE_M with
+    # the B stage's at prefill (>=512 tokens) so the per-layer routing_cache key matches
+    # across stages and the Triton align/sort runs once per layer-forward instead of once
+    # per stage (4x at prefill). Dtype-agnostic (the chain is shared by fp8/nvfp4/bf16).
+    # Decode (<512) keeps the opt1 fused merged-align path and its tuned shrink block.
+    SGLANG_OPT_LORA_PREFILL_ROUTING_REUSE = _GatedBool(
+        "SGLANG_OPT_LORA_PREFILL_ROUTING_REUSE", True
+    )
+
+    # ---- correctness fixes: on by default when experimental ----
+    # gate_up gated-split fix (up_A shrink for the up half); set =0 only to A/B bisect.
+    SGLANG_ENABLE_LORA_MOE_GATEUP_GATED_SPLIT = _GatedBool(
+        "SGLANG_ENABLE_LORA_MOE_GATEUP_GATED_SPLIT", True
+    )
+    # feed bf16 router logits straight to the JIT kimi gate (bitwise-identical).
+    SGLANG_OPT_KIMI_GATE_BF16_INPUT = _GatedBool(
+        "SGLANG_OPT_KIMI_GATE_BF16_INPUT", True
+    )
+
+    # ---- non-shared: default False, set explicitly in the launch env ----
+    # kimi (NVFP4):
+    SGLANG_OPT_USE_JIT_KERNEL_KIMI_GATE = _GatedBool(
+        "SGLANG_OPT_USE_JIT_KERNEL_KIMI_GATE", False
+    )
+    SGLANG_OPT_USE_JIT_KERNEL_MOE_ALIGN = _GatedBool(
+        "SGLANG_OPT_USE_JIT_KERNEL_MOE_ALIGN", False
+    )
+    # qwen3.5 (FP8):
+    SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC = _GatedBool(
+        "SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC", False
+    )
+    # (SGLANG_OPT_LORA_DOWN_FINALIZE_OVERLAP removed: net-neutral + base/decode-corruption hazard; serial down-LoRA only.)
+    SGLANG_OPT_LORA_SHARED_ADD_OVERLAP = _GatedBool(
+        "SGLANG_OPT_LORA_SHARED_ADD_OVERLAP", False
+    )
+    SGLANG_OPT_LORA_CUBLAS = _GatedBool("SGLANG_OPT_LORA_CUBLAS", False)
+    SGLANG_OPT_LORA_CUBLAS_A = _GatedBool("SGLANG_OPT_LORA_CUBLAS_A", False)
+    SGLANG_OPT_LORA_CUBLAS_B = _GatedBool("SGLANG_OPT_LORA_CUBLAS_B", False)
+    SGLANG_OPT_LORA_CUBLAS_GATE_UP = _GatedBool("SGLANG_OPT_LORA_CUBLAS_GATE_UP", False)
+    SGLANG_OPT_LORA_CUBLAS_QKV = _GatedBool("SGLANG_OPT_LORA_CUBLAS_QKV", False)
+    SGLANG_OPT_LORA_CUBLAS_KV_B = _GatedBool("SGLANG_OPT_LORA_CUBLAS_KV_B", False)
+    # diagnostics / tuning:
+    SGLANG_OPT_LORA_SHRINK_TUNE = _GatedBool("SGLANG_OPT_LORA_SHRINK_TUNE", False)
+
+    # kimi NVFP4 permute+quant fuse — read in jit_kernel/trtllm_lora_temp/core.py (Python) to pass
+    # a bool to the kernel, AND C++-side via getenv in the launcher. Default off (kimi-only).
+    SGLANG_OPT_FUSED_PERMUTE_QUANT = _GatedBool("SGLANG_OPT_FUSED_PERMUTE_QUANT", False)
+
+    # ---- integer knob ----
+    # decode two-stream token ceiling (consulted only on the experimental path).
+    SGLANG_TWO_STREAM_MAX_TOKENS = _GatedInt("SGLANG_TWO_STREAM_MAX_TOKENS", 256)
+
+
+lora_envs = _LoraEnvs()
+
+# ---------------------------------------------------------------------------
+# C++-getenv-only flags (read via getenv in jit_kernel .cu launchers, NOT in Python).
+# Set them in the launch env on the model that needs them; default off:
+#   SGLANG_OPT_FUSED_MOE_ACTIVATION_QUANT_FUSE (kimi NVFP4 act+down-quant fuse)
+#   SGLANG_OPT_FUSED_MOE_ACTIVATION_VEC      (kimi NVFP4 vectorized activation)
+# ---------------------------------------------------------------------------

python/sglang/srt/lora/trtllm_lora_temp/experimental_sgl_trtllm_moe.py

@@ -0,0 +1,203 @@
+from typing import Optional
+
+import torch
+
+from sglang.srt.utils.custom_op import register_custom_op
+
+
+def _fake_fp8_block_scale_moe(
+    routing_logits: torch.Tensor,
+    routing_bias: Optional[torch.Tensor],
+    hidden_states: torch.Tensor,
+    hidden_states_scale: torch.Tensor,
+    gemm1_weights: torch.Tensor,
+    gemm1_weights_scale: torch.Tensor,
+    gemm2_weights: torch.Tensor,
+    gemm2_weights_scale: torch.Tensor,
+    num_experts: int,
+    top_k: int,
+    n_group: Optional[int],
+    topk_group: Optional[int],
+    intermediate_size: int,
+    local_expert_offset: int,
+    local_num_experts: int,
+    routed_scaling_factor: Optional[float],
+    routing_method_type: int = 0,
+    use_shuffled_weight: bool = False,
+    weight_layout: int = 0,
+    enable_pdl: Optional[bool] = None,
+    tune_max_num_tokens: int = 8192,
+    fp8_quantization_type: Optional[int] = None,
+    activation_type: Optional[int] = None,
+) -> torch.Tensor:
+    return torch.empty(
+        hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device
+    )
+
+
+@register_custom_op(fake_impl=_fake_fp8_block_scale_moe)
+def sgl_trtllm_fp8_block_scale_moe_wrapper(
+    routing_logits: torch.Tensor,
+    routing_bias: Optional[torch.Tensor],
+    hidden_states: torch.Tensor,
+    hidden_states_scale: torch.Tensor,
+    gemm1_weights: torch.Tensor,
+    gemm1_weights_scale: torch.Tensor,
+    gemm2_weights: torch.Tensor,
+    gemm2_weights_scale: torch.Tensor,
+    num_experts: int,
+    top_k: int,
+    n_group: Optional[int],
+    topk_group: Optional[int],
+    intermediate_size: int,
+    local_expert_offset: int,
+    local_num_experts: int,
+    routed_scaling_factor: Optional[float],
+    routing_method_type: int = 0,
+    use_shuffled_weight: bool = False,
+    weight_layout: int = 0,
+    enable_pdl: Optional[bool] = None,
+    tune_max_num_tokens: int = 8192,
+    fp8_quantization_type: Optional[int] = None,
+    activation_type: Optional[int] = None,
+) -> torch.Tensor:
+    try:
+        from flashinfer.fused_moe import Fp8QuantizationType
+        from flashinfer.fused_moe.core import ActivationType
+    except ImportError as e:
+        raise ImportError(
+            "experimental_sgl_trtllm requires flashinfer-python to provide "
+            "TRTLLM enums and cubin-loader utilities."
+        ) from e
+
+    from sglang.jit_kernel.trtllm_lora_temp import trtllm_fp8_block_scale_moe
+
+    kwargs = {
+        "routing_logits": routing_logits,
+        "routing_bias": routing_bias,
+        "hidden_states": hidden_states,
+        "hidden_states_scale": hidden_states_scale,
+        "gemm1_weights": gemm1_weights,
+        "gemm1_weights_scale": gemm1_weights_scale,
+        "gemm2_weights": gemm2_weights,
+        "gemm2_weights_scale": gemm2_weights_scale,
+        "num_experts": num_experts,
+        "top_k": top_k,
+        "n_group": n_group,
+        "topk_group": topk_group,
+        "intermediate_size": intermediate_size,
+        "local_expert_offset": local_expert_offset,
+        "local_num_experts": local_num_experts,
+        "routed_scaling_factor": routed_scaling_factor,
+        "routing_method_type": routing_method_type,
+        "use_shuffled_weight": use_shuffled_weight,
+        "weight_layout": weight_layout,
+        "enable_pdl": enable_pdl,
+        "tune_max_num_tokens": tune_max_num_tokens,
+    }
+    if fp8_quantization_type is not None:
+        kwargs["fp8_quantization_type"] = Fp8QuantizationType(fp8_quantization_type)
+    if activation_type is not None:
+        kwargs["activation_type"] = ActivationType(activation_type)
+
+    return trtllm_fp8_block_scale_moe(**kwargs)
+
+
+def _fake_fp8_block_scale_routed_moe(
+    topk_ids: torch.Tensor,
+    routing_bias: Optional[torch.Tensor],
+    hidden_states: torch.Tensor,
+    hidden_states_scale: torch.Tensor,
+    gemm1_weights: torch.Tensor,
+    gemm1_weights_scale: torch.Tensor,
+    gemm2_weights: torch.Tensor,
+    gemm2_weights_scale: torch.Tensor,
+    num_experts: int,
+    top_k: int,
+    n_group: Optional[int],
+    topk_group: Optional[int],
+    intermediate_size: int,
+    local_expert_offset: int,
+    local_num_experts: int,
+    routed_scaling_factor: Optional[float],
+    routing_method_type: int = 0,
+    use_shuffled_weight: bool = False,
+    weight_layout: int = 0,
+    enable_pdl: Optional[bool] = None,
+    tune_max_num_tokens: int = 8192,
+    fp8_quantization_type: Optional[int] = None,
+    activation_type: Optional[int] = None,
+) -> torch.Tensor:
+    return torch.empty(
+        hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device
+    )
+
+
+@register_custom_op(fake_impl=_fake_fp8_block_scale_routed_moe)
+def sgl_trtllm_fp8_block_scale_routed_moe_wrapper(
+    topk_ids: torch.Tensor,
+    routing_bias: Optional[torch.Tensor],
+    hidden_states: torch.Tensor,
+    hidden_states_scale: torch.Tensor,
+    gemm1_weights: torch.Tensor,
+    gemm1_weights_scale: torch.Tensor,
+    gemm2_weights: torch.Tensor,
+    gemm2_weights_scale: torch.Tensor,
+    num_experts: int,
+    top_k: int,
+    n_group: Optional[int],
+    topk_group: Optional[int],
+    intermediate_size: int,
+    local_expert_offset: int,
+    local_num_experts: int,
+    routed_scaling_factor: Optional[float],
+    routing_method_type: int = 0,
+    use_shuffled_weight: bool = False,
+    weight_layout: int = 0,
+    enable_pdl: Optional[bool] = None,
+    tune_max_num_tokens: int = 8192,
+    fp8_quantization_type: Optional[int] = None,
+    activation_type: Optional[int] = None,
+) -> torch.Tensor:
+    try:
+        from flashinfer.fused_moe import Fp8QuantizationType
+        from flashinfer.fused_moe.core import ActivationType
+    except ImportError as e:
+        raise ImportError(
+            "experimental_sgl_trtllm requires flashinfer-python to provide "
+            "TRTLLM enums and cubin-loader utilities."
+        ) from e
+
+    from sglang.jit_kernel.trtllm_lora_temp import (
+        trtllm_fp8_block_scale_routed_moe,
+    )
+
+    kwargs = {
+        "topk_ids": topk_ids,
+        "routing_bias": routing_bias,
+        "hidden_states": hidden_states,
+        "hidden_states_scale": hidden_states_scale,
+        "gemm1_weights": gemm1_weights,
+        "gemm1_weights_scale": gemm1_weights_scale,
+        "gemm2_weights": gemm2_weights,
+        "gemm2_weights_scale": gemm2_weights_scale,
+        "num_experts": num_experts,
+        "top_k": top_k,
+        "n_group": n_group,
+        "topk_group": topk_group,
+        "intermediate_size": intermediate_size,
+        "local_expert_offset": local_expert_offset,
+        "local_num_experts": local_num_experts,
+        "routed_scaling_factor": routed_scaling_factor,
+        "routing_method_type": routing_method_type,
+        "use_shuffled_weight": use_shuffled_weight,
+        "weight_layout": weight_layout,
+        "enable_pdl": enable_pdl,
+        "tune_max_num_tokens": tune_max_num_tokens,
+    }
+    if fp8_quantization_type is not None:
+        kwargs["fp8_quantization_type"] = Fp8QuantizationType(fp8_quantization_type)
+    if activation_type is not None:
+        kwargs["activation_type"] = ActivationType(activation_type)
+
+    return trtllm_fp8_block_scale_routed_moe(**kwargs)

python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py

@@ -0,0 +1,614 @@
+"""experimental_sgl_trtllm MoE LoRA dispatch (original single-stream).
+
+This is the LoRA-enabled fused-experts path added by the trtllm-lora work — it
+was originally a function in ``layers/moe/moe_runner/flashinfer_trtllm.py`` and
+is now hosted here so that file holds only a re-export. The function name
+remains ``fused_experts_none_to_experimental_sgl_trtllm_fp8_lora`` for
+import-site stability.
+
+When ``SGLANG_LORA_TWO_STREAM=1`` is set, this is the function the
+``install_two_stream_overrides()`` monkey-patch swaps for the side-stream
+version in :mod:`sglang.srt.lora.trtllm_lora_temp.moe_overlap`. Otherwise it runs as
+the active path.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.srt.distributed import get_tp_group
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    use_symmetric_memory,
+)
+from sglang.srt.layers.dp_attention import is_allocation_symmetric
+from sglang.srt.layers.quantization.fp8_kernel import per_token_group_quant_fp8
+from sglang.srt.utils.common import next_power_of_2
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
+    from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
+        FlashInferTrtllmBf16MoeQuantInfo,
+        FlashInferTrtllmFp4MoeQuantInfo,
+        FlashInferTrtllmFp8MoeQuantInfo,
+    )
+    from sglang.srt.layers.moe.token_dispatcher import (
+        StandardCombineInput,
+        StandardDispatchOutput,
+    )
+
+
+def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora(
+    dispatch_output: StandardDispatchOutput,
+    quant_info: FlashInferTrtllmFp8MoeQuantInfo,
+    runner_config: MoeRunnerConfig,
+    lora_info,
+) -> StandardCombineInput:
+    from flashinfer.fused_moe import Fp8QuantizationType
+
+    from sglang.jit_kernel.trtllm_lora_temp import (
+        trtllm_fp8_block_scale_moe_lora_finalize,
+        trtllm_fp8_block_scale_routed_moe_lora,
+    )
+    from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk
+    from sglang.kernels.ops.moe.trtllm_lora_temp.virtual_experts import (
+        merged_experts_fused_moe_lora_add,
+    )
+    from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
+    from sglang.srt.layers.moe.topk import TopKOutputChecker
+    from sglang.srt.layers.moe.utils import RoutingMethodType
+    from sglang.srt.lora.lora_moe_runners import build_lora_hooks
+    from sglang.srt.lora.trtllm_lora_temp.sgl_fp8_moe import (
+        fused_experts_fp8_sgl,
+    )
+    from sglang.srt.lora.trtllm_lora_temp.shared_add_overlap import (
+        maybe_overlap_staged_shared_add,
+    )
+    from sglang.srt.model_executor.runner_utils.capture_mode import get_is_capture_mode
+
+    assert runner_config.activation == "silu" and runner_config.is_gated, (
+        "experimental_sgl_trtllm LoRA currently supports the gated SwiGLU FP8 "
+        "Qwen path only."
+    )
+    assert quant_info.block_quant and not quant_info.use_mxfp8, (
+        "experimental_sgl_trtllm LoRA currently supports DeepSeekFp8 block-quant "
+        "checkpoints only."
+    )
+    assert quant_info.weight_block_k is not None
+    assert quant_info.w13_weight_scale_inv is not None
+    assert quant_info.w2_weight_scale_inv is not None
+
+    hidden_states = dispatch_output.hidden_states
+    topk_output = dispatch_output.topk_output
+    assert TopKOutputChecker.format_is_standard(topk_output)
+    assert runner_config.top_k is not None
+
+    if not get_is_capture_mode() and not lora_info.has_active_lora:
+        return fused_experts_fp8_sgl(
+            dispatch_output,
+            quant_info,
+            runner_config,
+            use_routed_topk=True,
+        )
+
+    topk_ids = topk_output.topk_ids
+    topk_weights = topk_output.topk_weights
+    use_virtual_lora_store = bool(
+        lora_info.lora_use_virtual_experts and lora_info.max_lora_rank > 0
+    )
+    if use_virtual_lora_store:
+        hooks = None
+        token_lora_mapping = lora_info.token_lora_mapping
+        fused_lora_routing_cache: dict = {}
+    else:
+        hooks = build_lora_hooks(hidden_states, lora_info, topk_ids)
+        token_lora_mapping = None
+        fused_lora_routing_cache = {}
+
+    # Fuse the per-token scale transpose into the quant kernel (column-major scales) so the
+    # `.t()` is a free view -> drops the standalone ~2us transpose+copy. Byte/shape-identical.
+    a_q, a_sf = per_token_group_quant_fp8(
+        hidden_states, quant_info.weight_block_k, column_major_scales=True
+    )
+    a_sf_t = a_sf.t()
+
+    # EP-aware LoRA: under MoE EP each rank computes the delta only for the experts it
+    # owns (passed via local_expert_offset/local_num_experts below). gate_up_delta stays
+    # new_empty even though non-owned [token, k] slots are then left unwritten -- the
+    # trtllm MoE is itself EP-aware, so those slots never feed the all-reduced output.
+    gate_up_delta_shape = (
+        hidden_states.shape[0],
+        runner_config.top_k,
+        quant_info.w13_weight.shape[1],
+    )
+    gate_up_delta = (
+        hidden_states.new_empty(gate_up_delta_shape)
+        if use_virtual_lora_store
+        else hidden_states.new_zeros(gate_up_delta_shape)
+    )
+    if use_virtual_lora_store:
+        merged_experts_fused_moe_lora_add(
+            output=gate_up_delta,
+            hidden_states=hidden_states,
+            lora_a=lora_info.gate_up_lora_a_weights,
+            lora_b=lora_info.gate_up_lora_b_weights,
+            topk_ids=topk_ids,
+            topk_weights=topk_weights,
+            token_lora_mapping=token_lora_mapping,
+            mul_routed_weight=False,
+            experts_shared_outer_loras_a=lora_info.experts_shared_outer_loras,
+            experts_shared_outer_loras_b=False,
+            routing_cache=fused_lora_routing_cache,
+            fuse_add_to_output=False,
+            use_direct_expand_add=lora_info.max_lora_rank <= 64,
+            local_expert_offset=quant_info.local_expert_offset,
+            local_num_experts=quant_info.local_num_experts,
+        )
+    elif hooks.after_gate_up is not None:
+        hooks.after_gate_up(hidden_states, gate_up_delta, topk_weights, topk_ids)
+
+    activation_lora_input = torch.empty(
+        (hidden_states.shape[0], runner_config.top_k, quant_info.intermediate_size),
+        dtype=hidden_states.dtype,
+        device=hidden_states.device,
+    )
+
+    # SGLANG_OPT_LORA_FUSED_TOPK_PACK: the routed pack may already have been produced
+    # fused inside the gating kernel (StandardTopKOutput.packed_topk_ids) — including
+    # the padded-region id=-1 mask. Fall back to the separate pack otherwise.
+    packed_topk_ids = getattr(topk_output, "packed_topk_ids", None)
+    if packed_topk_ids is None:
+        packed_topk_ids = fused_pack_topk(
+            topk_ids=topk_ids,
+            topk_weights=topk_weights,
+        )
+
+    direct_down_output = None
+    if use_virtual_lora_store:
+        with use_symmetric_memory(
+            get_tp_group(), disabled=not is_allocation_symmetric()
+        ):
+            direct_down_output = torch.empty(
+                hidden_states.shape[0],
+                hidden_states.shape[1],
+                dtype=hidden_states.dtype,
+                device=hidden_states.device,
+            )
+
+    moe_result = trtllm_fp8_block_scale_routed_moe_lora(
+        topk_ids=packed_topk_ids,
+        routing_bias=None,
+        hidden_states=a_q,
+        hidden_states_scale=a_sf_t,
+        gemm1_weights=quant_info.w13_weight,
+        gemm1_weights_scale=quant_info.w13_weight_scale_inv,
+        gemm2_weights=quant_info.w2_weight,
+        gemm2_weights_scale=quant_info.w2_weight_scale_inv,
+        gate_up_lora_delta=gate_up_delta,
+        activation_lora_input=activation_lora_input,
+        num_experts=quant_info.global_num_experts,
+        top_k=runner_config.top_k,
+        n_group=None,
+        topk_group=None,
+        intermediate_size=quant_info.intermediate_size,
+        local_expert_offset=quant_info.local_expert_offset,
+        local_num_experts=quant_info.local_num_experts,
+        routed_scaling_factor=(
+            runner_config.routed_scaling_factor
+            if runner_config.routed_scaling_factor is not None
+            else 1.0
+        ),
+        routing_method_type=(
+            RoutingMethodType.TopK
+            if quant_info.routing_method_type == RoutingMethodType.DeepSeekV3
+            else quant_info.routing_method_type
+        ),
+        use_shuffled_weight=False,
+        do_finalize=use_virtual_lora_store,
+        output=(
+            direct_down_output
+            if direct_down_output is not None
+            else torch.empty_like(hidden_states)
+        ),
+        tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
+        fp8_quantization_type=Fp8QuantizationType.DeepSeekFp8,
+        activation_type=quant_info.activation_type,
+    )
+    if use_virtual_lora_store:
+        output = moe_result
+        # Shared-add overlap: the trtllm op above already finalized `output`, so the
+        # staged shared-expert add (if any) can run on the main stream concurrent with
+        # the down-LoRA shrink below; the expand waits on it via expand_wait_event.
+        shared_add_done = maybe_overlap_staged_shared_add(output)
+        merged_experts_fused_moe_lora_add(
+            output=output,
+            hidden_states=activation_lora_input.view(-1, quant_info.intermediate_size),
+            lora_a=lora_info.down_lora_a_weights,
+            lora_b=lora_info.down_lora_b_weights,
+            topk_ids=topk_ids,
+            topk_weights=topk_weights,
+            token_lora_mapping=token_lora_mapping,
+            mul_routed_weight=True,
+            experts_shared_outer_loras_a=False,
+            experts_shared_outer_loras_b=lora_info.experts_shared_outer_loras,
+            routing_cache=fused_lora_routing_cache,
+            fuse_add_to_output=False,
+            fuse_sum_all_reduce=True,
+            use_direct_expand_add=lora_info.max_lora_rank <= 64,
+            local_expert_offset=quant_info.local_expert_offset,
+            local_num_experts=quant_info.local_num_experts,
+            expand_wait_event=shared_add_done,
+        )
+        return StandardCombineInput(hidden_states=output)
+
+    gemm2_output, expert_weights, expanded_idx_to_permuted_idx = moe_result
+
+    down_delta_shape = (
+        hidden_states.shape[0],
+        runner_config.top_k,
+        hidden_states.shape[1],
+    )
+    down_delta = (
+        hidden_states.new_empty(down_delta_shape)
+        if use_virtual_lora_store
+        else hidden_states.new_zeros(down_delta_shape)
+    )
+    if use_virtual_lora_store:
+        merged_experts_fused_moe_lora_add(
+            output=down_delta,
+            hidden_states=activation_lora_input.view(-1, quant_info.intermediate_size),
+            lora_a=lora_info.down_lora_a_weights,
+            lora_b=lora_info.down_lora_b_weights,
+            topk_ids=topk_ids,
+            topk_weights=topk_weights,
+            token_lora_mapping=token_lora_mapping,
+            mul_routed_weight=True,
+            experts_shared_outer_loras_a=False,
+            experts_shared_outer_loras_b=lora_info.experts_shared_outer_loras,
+            routing_cache=fused_lora_routing_cache,
+            fuse_add_to_output=False,
+        )
+    elif hooks.after_down is not None:
+        hooks.after_down(
+            activation_lora_input.view(-1, quant_info.intermediate_size),
+            down_delta,
+            topk_weights,
+            topk_ids,
+        )
+
+    with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()):
+        output = torch.empty(
+            hidden_states.shape[0],
+            hidden_states.shape[1],
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
+    output = trtllm_fp8_block_scale_moe_lora_finalize(
+        gemm2_output=gemm2_output,
+        expert_weights=expert_weights,
+        expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx,
+        down_lora_delta=down_delta,
+        output=output,
+        routed_scaling_factor=(
+            runner_config.routed_scaling_factor
+            if runner_config.routed_scaling_factor is not None
+            else 1.0
+        ),
+    )
+
+    return StandardCombineInput(hidden_states=output)
+
+
+def fused_experts_none_to_experimental_sgl_trtllm_bf16_lora(
+    dispatch_output: StandardDispatchOutput,
+    quant_info: FlashInferTrtllmBf16MoeQuantInfo,
+    runner_config: MoeRunnerConfig,
+    lora_info,
+) -> StandardCombineInput:
+    """BF16 sibling of ``fused_experts_none_to_experimental_sgl_trtllm_fp8_lora``.
+
+    Decomposed (unfused-activation) MoE-LoRA, bf16 end-to-end (no quantization):
+    routing -> gather -> gate_up grouped GEMM (raw 2*inter, bf16) -> activation that
+    adds ``gate_up_lora_delta`` pre-SwiGLU and captures ``activation_lora_input`` ->
+    down grouped GEMM -> finalize, then the virtual-experts down-LoRA is merged into
+    the output. Single-stream version (no two-stream overlap yet — phase 2).
+    """
+    from sglang.jit_kernel.trtllm_lora_temp import trtllm_bf16_routed_moe_lora
+    from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk
+    from sglang.kernels.ops.moe.trtllm_lora_temp.virtual_experts import (
+        merged_experts_fused_moe_lora_add,
+    )
+    from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
+        fused_experts_none_to_flashinfer_trtllm_bf16,
+        get_activation_type,
+    )
+    from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
+    from sglang.srt.layers.moe.topk import TopKOutputChecker
+    from sglang.srt.layers.moe.utils import RoutingMethodType
+    from sglang.srt.model_executor.runner_utils.capture_mode import get_is_capture_mode
+
+    assert (
+        runner_config.activation == "silu" and runner_config.is_gated
+    ), "experimental_sgl_trtllm BF16 LoRA currently supports the gated SwiGLU path only."
+
+    hidden_states = dispatch_output.hidden_states
+    topk_output = dispatch_output.topk_output
+    assert TopKOutputChecker.format_is_standard(topk_output)
+    assert runner_config.top_k is not None
+
+    # No active LoRA in a non-capture decode -> plain (fast) bf16 path.
+    if not get_is_capture_mode() and not lora_info.has_active_lora:
+        return fused_experts_none_to_flashinfer_trtllm_bf16(
+            dispatch_output, quant_info, runner_config, use_routed_topk=True
+        )
+
+    topk_ids = topk_output.topk_ids
+    topk_weights = topk_output.topk_weights
+    use_virtual_lora_store = bool(
+        lora_info.lora_use_virtual_experts and lora_info.max_lora_rank > 0
+    )
+    assert use_virtual_lora_store, "BF16 trtllm LoRA requires virtual-experts."
+    token_lora_mapping = lora_info.token_lora_mapping
+    fused_lora_routing_cache: dict = {}
+
+    inter = runner_config.intermediate_size_per_partition
+
+    # Gated gate_up LoRA delta (same shape/semantics as the fp8/fp4 paths). EP args scope
+    # the delta to this rank's experts, matching the EP-aware trtllm MoE base.
+    gate_up_delta = hidden_states.new_empty(
+        (hidden_states.shape[0], runner_config.top_k, 2 * inter)
+    )
+    merged_experts_fused_moe_lora_add(
+        output=gate_up_delta,
+        hidden_states=hidden_states,
+        lora_a=lora_info.gate_up_lora_a_weights,
+        lora_b=lora_info.gate_up_lora_b_weights,
+        topk_ids=topk_ids,
+        topk_weights=topk_weights,
+        token_lora_mapping=token_lora_mapping,
+        mul_routed_weight=False,
+        experts_shared_outer_loras_a=lora_info.experts_shared_outer_loras,
+        experts_shared_outer_loras_b=False,
+        routing_cache=fused_lora_routing_cache,
+        fuse_add_to_output=False,
+        use_direct_expand_add=lora_info.max_lora_rank <= 64,
+        local_expert_offset=quant_info.local_expert_offset,
+        local_num_experts=runner_config.num_local_experts,
+    )
+
+    activation_lora_input = torch.empty(
+        (hidden_states.shape[0], runner_config.top_k, inter),
+        dtype=hidden_states.dtype,
+        device=hidden_states.device,
+    )
+
+    packed_topk_ids = getattr(topk_output, "packed_topk_ids", None)
+    if packed_topk_ids is None:
+        packed_topk_ids = fused_pack_topk(
+            topk_ids=topk_ids,
+            topk_weights=topk_weights,
+        )
+
+    routing_method_type = runner_config.routing_method_type
+    if routing_method_type is None:
+        routing_method_type = RoutingMethodType.Default
+    elif routing_method_type == RoutingMethodType.DeepSeekV3:
+        routing_method_type = RoutingMethodType.TopK
+
+    with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()):
+        direct_down_output = torch.empty(
+            hidden_states.shape[0],
+            hidden_states.shape[1],
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
+
+    output = trtllm_bf16_routed_moe_lora(
+        topk_ids=packed_topk_ids,
+        routing_bias=None,
+        hidden_states=hidden_states,
+        gemm1_weights=quant_info.gemm1_weights,
+        gemm2_weights=quant_info.gemm2_weights,
+        gate_up_lora_delta=gate_up_delta,
+        activation_lora_input=activation_lora_input,
+        num_experts=quant_info.global_num_experts,
+        top_k=runner_config.top_k,
+        intermediate_size=inter,
+        local_expert_offset=quant_info.local_expert_offset,
+        local_num_experts=runner_config.num_local_experts,
+        routed_scaling_factor=(
+            runner_config.routed_scaling_factor
+            if runner_config.routed_scaling_factor is not None
+            else 1.0
+        ),
+        routing_method_type=routing_method_type,
+        do_finalize=True,
+        output=direct_down_output,
+        activation_type=get_activation_type(
+            runner_config.activation, is_gated=runner_config.is_gated
+        ),
+    )
+
+    merged_experts_fused_moe_lora_add(
+        output=output,
+        hidden_states=activation_lora_input.view(-1, inter),
+        lora_a=lora_info.down_lora_a_weights,
+        lora_b=lora_info.down_lora_b_weights,
+        topk_ids=topk_ids,
+        topk_weights=topk_weights,
+        token_lora_mapping=token_lora_mapping,
+        mul_routed_weight=True,
+        experts_shared_outer_loras_a=False,
+        experts_shared_outer_loras_b=lora_info.experts_shared_outer_loras,
+        routing_cache=fused_lora_routing_cache,
+        fuse_add_to_output=False,
+        fuse_sum_all_reduce=True,
+        use_direct_expand_add=lora_info.max_lora_rank <= 64,
+        local_expert_offset=quant_info.local_expert_offset,
+        local_num_experts=runner_config.num_local_experts,
+    )
+    return StandardCombineInput(hidden_states=output)
+
+
+def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora(
+    dispatch_output: StandardDispatchOutput,
+    quant_info: FlashInferTrtllmFp4MoeQuantInfo,
+    runner_config: MoeRunnerConfig,
+    lora_info,
+) -> StandardCombineInput:
+    """NVFP4 sibling of ``fused_experts_none_to_experimental_sgl_trtllm_fp8_lora``.
+
+    Decomposed (unfused-activation) MoE-LoRA: routing -> gather -> gate_up grouped
+    GEMM (raw 2*inter) -> activation that adds ``gate_up_lora_delta`` pre-SwiGLU and
+    captures ``activation_lora_input`` -> NvFP4 quant -> down grouped GEMM -> finalize,
+    then the virtual-experts down-LoRA is merged into the output. Single-stream
+    version; ``moe_overlap.py`` provides the two-stream variant.
+    """
+    from sglang.jit_kernel.trtllm_lora_temp import (
+        trtllm_fp4_block_scale_routed_moe_lora,
+    )
+    from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk
+    from sglang.kernels.ops.moe.trtllm_lora_temp.virtual_experts import (
+        merged_experts_fused_moe_lora_add,
+    )
+    from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
+        fused_experts_none_to_flashinfer_trtllm_fp4,
+    )
+    from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
+    from sglang.srt.layers.moe.topk import TopKOutputChecker
+    from sglang.srt.model_executor.runner_utils.capture_mode import get_is_capture_mode
+
+    assert (
+        runner_config.activation == "silu" and runner_config.is_gated
+    ), "experimental_sgl_trtllm NVFP4 LoRA currently supports the gated SwiGLU path only."
+
+    hidden_states = dispatch_output.hidden_states
+    topk_output = dispatch_output.topk_output
+    assert TopKOutputChecker.format_is_standard(topk_output)
+    assert runner_config.top_k is not None
+
+    # No active LoRA in a non-capture decode -> plain (fast) FP4 path.
+    if not get_is_capture_mode() and not lora_info.has_active_lora:
+        return fused_experts_none_to_flashinfer_trtllm_fp4(
+            dispatch_output, quant_info, runner_config, use_routed_topk=True
+        )
+
+    topk_ids = topk_output.topk_ids
+    topk_weights = topk_output.topk_weights
+    use_virtual_lora_store = bool(
+        lora_info.lora_use_virtual_experts and lora_info.max_lora_rank > 0
+    )
+    assert use_virtual_lora_store, "NVFP4 trtllm LoRA requires virtual-experts."
+    token_lora_mapping = lora_info.token_lora_mapping
+    fused_lora_routing_cache: dict = {}
+
+    inter = quant_info.intermediate_size_per_partition
+
+    # Path 3: feed bf16 hidden DIRECTLY to the op (no python pre-quant). The op permutes the
+    # bf16 hidden (moe::dev::permute, token->expert order) then NvFP4-quantizes ONCE with the
+    # 1/(448*6) global + per-token scale — eliminating the dequant->permute->requant round-trip
+    # (and its magnitude bug) that pre-quantized fp4 input forced. Requires the layer to be in
+    # per-token-activation mode (SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION=1) so that
+    # g1_scale_c == g1_alphas and g2_alphas == w2_weight_scale_2, making the decomposed
+    # gate_up(g1_alphas)/SwiGLU/down(g2_alphas) scale composition match the plain fused path.
+
+    gate_up_delta_shape = (
+        hidden_states.shape[0],
+        runner_config.top_k,
+        quant_info.w13_weight.shape[1],
+    )
+    # Gated gate_up LoRA delta: a single merged_experts call on the full stacked [gate_A; up_A]
+    # lora_a (rank 2r) and [gate_B; up_B] lora_b (rank r), via the rank-specialized direct expand
+    # (use_direct_expand_add, rank <= 64). EP args scope the delta to this rank's experts, matching
+    # the EP-aware trtllm MoE base.
+    gate_up_delta = hidden_states.new_empty(gate_up_delta_shape)
+    merged_experts_fused_moe_lora_add(
+        output=gate_up_delta,
+        hidden_states=hidden_states,
+        lora_a=lora_info.gate_up_lora_a_weights,
+        lora_b=lora_info.gate_up_lora_b_weights,
+        topk_ids=topk_ids,
+        topk_weights=topk_weights,
+        token_lora_mapping=token_lora_mapping,
+        mul_routed_weight=False,
+        experts_shared_outer_loras_a=lora_info.experts_shared_outer_loras,
+        experts_shared_outer_loras_b=False,
+        routing_cache=fused_lora_routing_cache,
+        fuse_add_to_output=False,
+        use_direct_expand_add=lora_info.max_lora_rank <= 64,
+        local_expert_offset=quant_info.local_expert_offset,
+        local_num_experts=quant_info.local_num_experts,
+    )
+
+    activation_lora_input = torch.empty(
+        (hidden_states.shape[0], runner_config.top_k, inter),
+        dtype=hidden_states.dtype,
+        device=hidden_states.device,
+    )
+
+    packed_topk_ids = fused_pack_topk(
+        topk_ids=topk_ids,
+        topk_weights=topk_weights,
+    )
+
+    with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()):
+        direct_down_output = torch.empty(
+            hidden_states.shape[0],
+            hidden_states.shape[1],
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
+
+    output = trtllm_fp4_block_scale_routed_moe_lora(
+        topk_ids=packed_topk_ids,
+        routing_bias=None,
+        hidden_states=hidden_states,
+        hidden_states_scale=None,
+        gemm1_weights=quant_info.w13_weight,
+        gemm1_weights_scale=quant_info.w13_weight_scale.view(torch.float8_e4m3fn),
+        gemm2_weights=quant_info.w2_weight,
+        gemm2_weights_scale=quant_info.w2_weight_scale.view(torch.float8_e4m3fn),
+        output1_scales_scalar=quant_info.g1_scale_c,
+        output1_scales_gate_scalar=quant_info.g1_alphas,
+        output2_scales_scalar=quant_info.g2_alphas,
+        gate_up_lora_delta=gate_up_delta,
+        activation_lora_input=activation_lora_input,
+        num_experts=quant_info.global_num_experts,
+        top_k=runner_config.top_k,
+        intermediate_size=inter,
+        local_expert_offset=quant_info.local_expert_offset,
+        local_num_experts=quant_info.local_num_experts,
+        routed_scaling_factor=(
+            runner_config.routed_scaling_factor
+            if runner_config.routed_scaling_factor is not None
+            else 1.0
+        ),
+        routing_method_type=quant_info.routing_method_type,
+        do_finalize=True,
+        output=direct_down_output,
+    )
+
+    merged_experts_fused_moe_lora_add(
+        output=output,
+        hidden_states=activation_lora_input.view(-1, inter),
+        lora_a=lora_info.down_lora_a_weights,
+        lora_b=lora_info.down_lora_b_weights,
+        topk_ids=topk_ids,
+        topk_weights=topk_weights,
+        token_lora_mapping=token_lora_mapping,
+        mul_routed_weight=True,
+        experts_shared_outer_loras_a=False,
+        experts_shared_outer_loras_b=lora_info.experts_shared_outer_loras,
+        routing_cache=fused_lora_routing_cache,
+        fuse_add_to_output=False,
+        fuse_sum_all_reduce=True,
+        use_direct_expand_add=lora_info.max_lora_rank <= 64,
+        # EP-aware: scope the down delta to this rank's experts, matching the gate_up
+        # call above and the FP8 down call. Harmless at EP=1 (local==global, Kimi today);
+        # required for correctness if MoE-EP is turned on later (otherwise non-owned
+        # experts' deltas get over-counted by the fuse_sum_all_reduce).
+        local_expert_offset=quant_info.local_expert_offset,
+        local_num_experts=quant_info.local_num_experts,
+    )
+    return StandardCombineInput(hidden_states=output)

python/sglang/srt/lora/trtllm_lora_temp/lora_layer.py

@@ -0,0 +1,188 @@
+"""experimental_sgl_trtllm-specific bits of ``FusedMoEWithLoRA``.
+
+This file holds the two trtllm-specific code blocks that used to be inlined
+inside the ``FusedMoEWithLoRA`` class in ``lora/layers.py``:
+
+  - :func:`init_experimental_sgl_trtllm_lora` — builds the FP8 block-scale
+    ``FlashInferTrtllmFp8MoeQuantInfo`` and stores it on the layer instance.
+    Called from ``FusedMoEWithLoRA.__init__`` when the runner backend is the
+    experimental_sgl_trtllm MoE.
+  - :func:`dispatch_experimental_sgl_trtllm_lora` — dispatches the LoRA fused
+    experts call. Called from ``FusedMoEWithLoRA.run`` for the same backend.
+
+Keeping them here means ``lora/layers.py`` only has tiny ``if backend == ...:
+init(self, base_layer)`` / ``dispatch(...)`` injection points for the new
+trtllm path instead of ~70 lines of inlined logic.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
+
+
+def init_experimental_sgl_trtllm_lora(layer, base_layer) -> None:
+    """Build and store the trtllm FP8 LoRA quant info on the layer.
+
+    Sets ``layer._lora_runner = None`` (trtllm path doesn't use ``MoeRunner``)
+    and ``layer._quant_info`` to a fully-populated
+    ``FlashInferTrtllmFp8MoeQuantInfo`` (or ``FlashInferTrtllmFp4MoeQuantInfo`` for
+    NVFP4 / modelopt checkpoints like Kimi-K2.5-NVFP4).
+    """
+    from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
+        FlashInferTrtllmFp8MoeQuantInfo,
+        get_activation_type,
+    )
+    from sglang.srt.layers.moe.utils import RoutingMethodType
+
+    # ---- NVFP4 (modelopt) path ----
+    # The fp4 weight loader sets ``g1_scale_c`` on the FusedMoE layer (see
+    # ModelOptNvFp4FusedMoEMethod.apply). Mirror the non-LoRA construction in
+    # modelopt_quant.py (~L2099) so the fp4 LoRA dispatch gets the same payload.
+    # NOTE(w13 layout): the decomposed op runs the gate_up projection as a *non-gated*
+    # Gemm2-style GEMM and lets the activation kernel do the SwiGLU split (silu(first)*second),
+    # so w13 must be [Gate, Up] with shuffle_matrix_a but WITHOUT reorder_rows_for_gated_act_gemm.
+    # The TRTLLM fp4 path uses load_up_proj_weight_first=False (=> [Gate, Up]); verify the
+    # processed w13 layout against this at e2e (acc gate) and re-prep if mismatched.
+    if hasattr(base_layer, "g1_scale_c"):
+        from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
+            FlashInferTrtllmFp4MoeQuantInfo,
+        )
+
+        layer._lora_runner = None
+        layer._quant_info = FlashInferTrtllmFp4MoeQuantInfo(
+            w13_weight=base_layer.w13_weight.data,
+            w2_weight=base_layer.w2_weight.data,
+            w13_weight_scale=base_layer.w13_weight_scale.data,
+            w2_weight_scale=base_layer.w2_weight_scale.data,
+            g1_scale_c=base_layer.g1_scale_c.data,
+            g1_alphas=base_layer.g1_alphas.data,
+            g2_alphas=base_layer.g2_alphas.data,
+            w13_input_scale_quant=base_layer.w13_input_scale_quant,
+            global_num_experts=int(base_layer.num_experts),
+            local_expert_offset=int(base_layer.moe_ep_rank)
+            * int(base_layer.num_local_experts),
+            local_num_experts=int(base_layer.num_local_experts),
+            intermediate_size_per_partition=int(
+                base_layer.intermediate_size_per_partition
+            ),
+            routing_method_type=int(
+                getattr(base_layer, "routing_method_type", None)
+                or RoutingMethodType.DeepSeekV3
+            ),
+        )
+        return
+
+    quant_method = base_layer.quant_method
+    quant_config = getattr(quant_method, "quant_config", None)
+    weight_block_size = getattr(quant_config, "weight_block_size", None)
+    if weight_block_size is None:
+        weight_block_size = getattr(quant_method, "weight_block_size", None)
+    use_mxfp8 = bool(getattr(quant_config, "use_mxfp8", False))
+
+    # ---- BF16 (unquantized) path ----
+    # No quant_config / block scales => the checkpoint is bf16. The bf16 LoRA dispatch
+    # runs the decomposed trtllm pipeline (sgl_trtllm_bf16_routed_moe_lora): permute ->
+    # raw gate_up GEMM -> LoRA-aware activation -> down GEMM, all bf16 — using the SAME
+    # prepared w13/w2 tensors (shuffled + BlockMajorK) the plain trtllm_bf16 path consumes.
+    # intermediate_size / local_num_experts / routing_method_type come from
+    # base_layer.moe_runner_config at dispatch time (the bf16 quant-info is minimal).
+    if quant_config is None and not getattr(quant_method, "block_quant", False):
+        from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
+            FlashInferTrtllmBf16MoeQuantInfo,
+        )
+
+        layer._lora_runner = None
+        layer._quant_info = FlashInferTrtllmBf16MoeQuantInfo(
+            gemm1_weights=base_layer.w13_weight.data,
+            gemm2_weights=base_layer.w2_weight.data,
+            global_num_experts=int(base_layer.num_experts),
+            local_expert_offset=int(base_layer.moe_ep_rank)
+            * int(base_layer.num_local_experts),
+        )
+        # Expose w13_weight/w2_weight on the bf16 quant-info so the backend-agnostic
+        # cuda-graph MoE buffer init (BaseLoRABackend.init_cuda_graph_moe_buffers) reads
+        # expert dims uniformly with the FP8/FP4 quant-infos (which name them w13/w2).
+        # The bf16 BlockMajorK weights are 4-D [E, N, K/128, 128]; collapsing the inner
+        # dims to a 3-D [E, N, K] view (free for contiguous weights) makes the upstream
+        # `E, N, _ = w13_weight.shape` dim-extraction work without touching base_backend.
+        _g1 = layer._quant_info.gemm1_weights
+        _g2 = layer._quant_info.gemm2_weights
+        layer._quant_info.w13_weight = _g1.reshape(_g1.shape[0], _g1.shape[1], -1)
+        layer._quant_info.w2_weight = _g2.reshape(_g2.shape[0], _g2.shape[1], -1)
+        return
+
+    assert getattr(
+        quant_method, "block_quant", False
+    ), "experimental_sgl_trtllm LoRA currently requires FP8 block quant."
+    assert (
+        not use_mxfp8
+    ), "experimental_sgl_trtllm LoRA currently targets the non-MX FP8 Qwen path."
+    assert (
+        weight_block_size is not None
+    ), "experimental_sgl_trtllm LoRA needs the FP8 weight block size."
+    w13_weight_scale = getattr(base_layer, "w13_weight_scale_inv", None)
+    if w13_weight_scale is None:
+        w13_weight_scale = getattr(base_layer, "w13_weight_scale", None)
+    w2_weight_scale = getattr(base_layer, "w2_weight_scale_inv", None)
+    if w2_weight_scale is None:
+        w2_weight_scale = getattr(base_layer, "w2_weight_scale", None)
+    assert w13_weight_scale is not None and w2_weight_scale is not None
+
+    layer._lora_runner = None
+    layer._quant_info = FlashInferTrtllmFp8MoeQuantInfo(
+        w13_weight=base_layer.w13_weight,
+        w2_weight=base_layer.w2_weight,
+        global_num_experts=int(base_layer.num_experts),
+        local_expert_offset=int(base_layer.moe_ep_rank)
+        * int(base_layer.num_local_experts),
+        local_num_experts=int(base_layer.num_local_experts),
+        intermediate_size=base_layer.w2_weight.shape[2],
+        routing_method_type=int(
+            getattr(base_layer, "routing_method_type", None)
+            or RoutingMethodType.DeepSeekV3
+        ),
+        block_quant=True,
+        use_mxfp8=False,
+        weight_block_k=weight_block_size[1],
+        w13_weight_scale_inv=w13_weight_scale,
+        w2_weight_scale_inv=w2_weight_scale,
+        activation_type=get_activation_type(
+            base_layer.moe_runner_config.activation,
+            is_gated=base_layer.moe_runner_config.is_gated,
+        ),
+    )
+
+
+def dispatch_experimental_sgl_trtllm_lora(
+    dispatch_output, quant_info, base_layer, lora_info
+) -> StandardCombineInput:
+    """Call the trtllm fused-experts LoRA function for a single layer.
+
+    Looked up at call time so the install-time monkey-patch in
+    :mod:`sglang.srt.lora.trtllm_lora_temp` (the two-stream override) takes effect.
+    """
+    import sglang.srt.lora.trtllm_lora_temp.lora_dispatch as ft
+    from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
+        FlashInferTrtllmBf16MoeQuantInfo,
+        FlashInferTrtllmFp4MoeQuantInfo,
+    )
+
+    # Resolve the fused-experts fn on the module at CALL TIME so the install-time
+    # two-stream monkey-patch (sglang.srt.lora.trtllm_lora_temp) takes effect. Route by
+    # quant dtype: NVFP4 -> fp4 LoRA op, BF16 (unquantized) -> bf16 LoRA op, else FP8.
+    if isinstance(quant_info, FlashInferTrtllmFp4MoeQuantInfo):
+        fused_fn = ft.fused_experts_none_to_experimental_sgl_trtllm_fp4_lora
+    elif isinstance(quant_info, FlashInferTrtllmBf16MoeQuantInfo):
+        fused_fn = ft.fused_experts_none_to_experimental_sgl_trtllm_bf16_lora
+    else:
+        fused_fn = ft.fused_experts_none_to_experimental_sgl_trtllm_fp8_lora
+
+    return fused_fn(
+        dispatch_output,
+        quant_info,
+        base_layer.moe_runner_config,
+        lora_info,
+    )

python/sglang/srt/lora/trtllm_lora_temp/merged_column.py

@@ -0,0 +1,89 @@
+"""Two-stream MergedColumnParallelLinear LoRA forward (O9).
+
+Monkey-patched onto :class:`MergedColumnParallelLinearWithLoRA` by
+:func:`sglang.srt.lora.trtllm_lora_temp.install_two_stream_overrides` when
+``SGLANG_LORA_TWO_STREAM=1``. Covers the merged-column LoRA modules not handled
+by O7 (QKV) / O8 (o_proj) / O1 (MoE experts):
+
+  * Qwen3.5 mamba ``in_proj_qkvz`` (a MergedColumnParallelLinear, every mamba layer)
+  * dense ``gate_up_proj`` MLP layers (e.g. Qwen3-VL non-expert MLP)
+
+Same shape as O7: the LoRA-A shrink reads ``input_`` (same input as the base
+GEMM, no write conflict) and runs on the side stream concurrent with the base
+merged-column GEMM on the main stream; the LoRA-B expand needs both the shrink
+output and base_output, so it runs after the rejoin on the main stream. The
+expand mirrors ``MergedColumnParallelLinearWithLoRA.apply_lora`` — gate_up vs
+general n-slice.
+"""
+
+import torch
+
+from sglang.srt.distributed import tensor_model_parallel_all_gather
+from sglang.srt.lora.trtllm_lora_temp import (
+    get_lora_side_stream,
+    get_original_merged_column_forward,
+    is_two_stream_active,
+    lora_overlap_alloc_stream,
+)
+
+
+def merged_column_lora_forward(self, input_: torch.Tensor):
+    """O9 — side-stream LoRA-A shrink ‖ base merged-column GEMM."""
+    if not self.set_lora or not is_two_stream_active(input_):
+        return get_original_merged_column_forward()(self, input_)
+
+    from sglang.kernels.ops.gemm.trtllm_lora_temp.gate_up_lora_b import (
+        gate_up_lora_b_fwd,
+    )
+    from sglang.kernels.ops.gemm.trtllm_lora_temp.qkv_lora_b import qkv_lora_b_fwd
+    from sglang.kernels.ops.gemm.trtllm_lora_temp.sgemm_lora_a import sgemm_lora_a_fwd
+
+    bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None
+    side_stream = get_lora_side_stream()
+    # sgemm_info is host-side (LoRABatchInfo); compute once, share both calls.
+    sgemm_info = self.lora_backend._sgemm_info()
+    lora_n_slices = self._get_lora_n_slices()
+    use_gate_up = lora_n_slices == 2 and self.use_gate_up_lora
+
+    # Shrink on side stream, concurrent with the base merged-column GEMM on main.
+    _alloc = lora_overlap_alloc_stream()  # capture MAIN stream here (before the fork)
+    side_stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(side_stream):
+        shrink_intermediate = sgemm_lora_a_fwd(
+            input_,
+            self.A_buffer,
+            sgemm_info,
+            stack_num=lora_n_slices,
+            out_alloc_stream=_alloc,
+        )
+
+    output_parallel = self.base_layer.quant_method.apply(self.base_layer, input_, bias)
+
+    # Rejoin: expand reads both side-produced shrink_intermediate and base_output.
+    torch.cuda.current_stream().wait_stream(side_stream)
+    if use_gate_up:
+        output_dim = self.B_buffer.shape[-2] // 2
+        output_parallel = gate_up_lora_b_fwd(
+            shrink_intermediate,
+            self.B_buffer,
+            sgemm_info,
+            output_dim,
+            output_parallel,
+        )
+    else:
+        output_parallel = qkv_lora_b_fwd(
+            shrink_intermediate,
+            self.B_buffer,
+            sgemm_info,
+            self.output_offset,
+            self.max_out_dim,
+            output_parallel,
+            n_slices=lora_n_slices,
+        )
+
+    if self.base_layer.gather_output:
+        output = tensor_model_parallel_all_gather(output_parallel)
+    else:
+        output = output_parallel
+    output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
+    return output, output_bias

python/sglang/srt/lora/trtllm_lora_temp/moe_overlap.py

@@ -0,0 +1,786 @@
+"""Two-stream MoE LoRA dispatch (O1).
+
+Monkey-patches ``fused_experts_none_to_experimental_sgl_trtllm_fp8_lora`` in
+``layers/moe/moe_runner/flashinfer_trtllm.py`` (when
+``SGLANG_LORA_TWO_STREAM=1``) so the gate_up LoRA shrink+expand runs on a
+side stream concurrent with the main-stream FP8 quant.
+
+Batches that don't qualify for two-stream (prefill / non-virtual-lora /
+batch without active LoRA) fall through to the saved-original function so
+their behavior is byte-identical to the unpatched code path.
+"""
+
+import torch
+
+from sglang.srt.lora.trtllm_lora_temp import (
+    get_lora_side_stream,
+    get_original_bf16_moe_lora_func,
+    get_original_fp4_moe_lora_func,
+    get_original_moe_lora_func,
+    is_two_stream_active,
+)
+
+# GEMM1-LoRA overlap: keep LoRA-ready events recorded during cuda-graph capture alive so the
+# captured cross-stream wait (resolved inside the trtllm op before activation) isn't torn down
+# before graph instantiation. Only appended while capturing; eager runs rely on CUDA's
+# deferred cudaEventDestroy.
+_LORA_OVERLAP_EVENTS: list = []
+
+
+def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora_two_stream(
+    dispatch_output,
+    quant_info,
+    runner_config,
+    lora_info,
+):
+    """Drop-in replacement for the like-named function in flashinfer_trtllm.py.
+
+    Two-stream fast path: only fires when the batch is decode-shaped AND uses
+    virtual-experts LoRA. Everything else delegates to the original function.
+    """
+    hidden_states = dispatch_output.hidden_states
+
+    use_virtual_lora_store = bool(
+        lora_info.lora_use_virtual_experts and lora_info.max_lora_rank > 0
+    )
+    # Two-stream requires virtual-experts LoRA AND a decode-shaped batch.
+    # Fall back to the original implementation for anything else (prefill,
+    # non-virtual LoRA, non-LoRA capture, etc.).
+    if not (use_virtual_lora_store and is_two_stream_active(hidden_states)):
+        return get_original_moe_lora_func()(
+            dispatch_output, quant_info, runner_config, lora_info
+        )
+
+    # ---- two-stream fast path ----
+    from flashinfer.fused_moe import Fp8QuantizationType
+
+    from sglang.jit_kernel.trtllm_lora_temp import (
+        trtllm_fp8_block_scale_routed_moe_lora,
+    )
+    from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk
+    from sglang.kernels.ops.moe.trtllm_lora_temp.virtual_experts import (
+        merged_experts_fused_moe_lora_add,
+    )
+    from sglang.srt.distributed import get_tp_group
+    from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+        use_symmetric_memory,
+    )
+    from sglang.srt.layers.dp_attention import is_allocation_symmetric
+    from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
+    from sglang.srt.layers.moe.topk import TopKOutputChecker
+    from sglang.srt.layers.moe.utils import RoutingMethodType
+    from sglang.srt.layers.quantization.fp8_kernel import per_token_group_quant_fp8
+    from sglang.srt.lora.trtllm_lora_temp.shared_add_overlap import (
+        maybe_overlap_staged_shared_add,
+    )
+    from sglang.srt.utils.common import next_power_of_2
+
+    assert runner_config.activation == "silu" and runner_config.is_gated, (
+        "experimental_sgl_trtllm LoRA currently supports the gated SwiGLU FP8 "
+        "Qwen path only."
+    )
+    assert quant_info.block_quant and not quant_info.use_mxfp8, (
+        "experimental_sgl_trtllm LoRA currently supports DeepSeekFp8 block-quant "
+        "checkpoints only."
+    )
+    assert quant_info.weight_block_k is not None
+    assert quant_info.w13_weight_scale_inv is not None
+    assert quant_info.w2_weight_scale_inv is not None
+
+    topk_output = dispatch_output.topk_output
+    assert TopKOutputChecker.format_is_standard(topk_output)
+    assert runner_config.top_k is not None
+
+    topk_ids = topk_output.topk_ids
+    topk_weights = topk_output.topk_weights
+    token_lora_mapping = lora_info.token_lora_mapping
+    fused_lora_routing_cache: dict = {}
+
+    side_stream = get_lora_side_stream()
+
+    # EP-aware LoRA: under MoE EP each rank computes the delta only for its owned experts
+    # (passed via local_expert_offset/local_num_experts below). gate_up_delta stays
+    # new_empty even though non-owned [token, k] slots are then left unwritten -- the
+    # trtllm MoE is itself EP-aware, so those slots never feed the all-reduced output.
+    gate_up_delta_shape = (
+        hidden_states.shape[0],
+        runner_config.top_k,
+        quant_info.w13_weight.shape[1],
+    )
+    gate_up_delta = hidden_states.new_empty(gate_up_delta_shape)
+
+    def _run_gate_up_lora():
+        merged_experts_fused_moe_lora_add(
+            output=gate_up_delta,
+            hidden_states=hidden_states,
+            lora_a=lora_info.gate_up_lora_a_weights,
+            lora_b=lora_info.gate_up_lora_b_weights,
+            topk_ids=topk_ids,
+            topk_weights=topk_weights,
+            token_lora_mapping=token_lora_mapping,
+            mul_routed_weight=False,
+            experts_shared_outer_loras_a=lora_info.experts_shared_outer_loras,
+            experts_shared_outer_loras_b=False,
+            routing_cache=fused_lora_routing_cache,
+            fuse_add_to_output=False,
+            use_direct_expand_add=lora_info.max_lora_rank <= 64,
+            local_expert_offset=quant_info.local_expert_offset,
+            local_num_experts=quant_info.local_num_experts,
+            intermediate_buffer=gate_up_lora_intermediate,
+        )
+
+    # GEMM1-LoRA overlap: fire the gate_up LoRA on the side stream + record an event; the
+    # trtllm op waits on it right before activation (the only consumer of gate_up_delta), so
+    # permute+GEMM1 overlap the side-stream LoRA shrink/expand instead of joining before the
+    # whole op.
+    lora_event = torch.cuda.Event()
+
+    # Hoist every side-chain allocation onto the MAIN stream (cuda-graph
+    # allocator safety -- see the "routing" stage in virtual_experts.py):
+    # pre-warm the routing cache and pre-allocate the shrink intermediate here,
+    # so the side-stream block below launches kernels only.
+    merged_experts_fused_moe_lora_add(
+        output=gate_up_delta,
+        hidden_states=hidden_states,
+        lora_a=lora_info.gate_up_lora_a_weights,
+        lora_b=lora_info.gate_up_lora_b_weights,
+        topk_ids=topk_ids,
+        topk_weights=topk_weights,
+        token_lora_mapping=token_lora_mapping,
+        mul_routed_weight=False,
+        experts_shared_outer_loras_a=lora_info.experts_shared_outer_loras,
+        experts_shared_outer_loras_b=False,
+        routing_cache=fused_lora_routing_cache,
+        stage="routing",
+        local_expert_offset=quant_info.local_expert_offset,
+        local_num_experts=quant_info.local_num_experts,
+    )
+    gate_up_lora_intermediate = hidden_states.new_empty(
+        (
+            hidden_states.shape[0],
+            topk_ids.shape[1],
+            lora_info.gate_up_lora_a_weights.shape[2],
+        )
+    )
+
+    # O1 fork — gate_up shrink/expand on side stream concurrent with the main-stream
+    # per-token-group FP8 quant + the trtllm op's permute+GEMM1 below.
+    side_stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(side_stream):
+        _run_gate_up_lora()
+        lora_event.record()
+
+    # Fuse the per-token scale transpose into the quant kernel: column-major scales make
+    # the `.t()` a free view, dropping the standalone ~2us transpose+copy. The trtllm MoE
+    # kernel wants the [K, M]-contiguous scale, which `.t()` of the column-major buffer is
+    # exactly -- byte/shape-identical to the old `a_sf.t().contiguous()`.
+    a_q, a_sf = per_token_group_quant_fp8(
+        hidden_states, quant_info.weight_block_k, column_major_scales=True
+    )
+    a_sf_t = a_sf.t()
+
+    activation_lora_input = torch.empty(
+        (hidden_states.shape[0], runner_config.top_k, quant_info.intermediate_size),
+        dtype=hidden_states.dtype,
+        device=hidden_states.device,
+    )
+
+    # SGLANG_OPT_LORA_FUSED_TOPK_PACK: the routed pack may already have been produced
+    # fused inside the gating kernel (StandardTopKOutput.packed_topk_ids) — including
+    # the padded-region id=-1 mask. Fall back to the separate pack otherwise.
+    packed_topk_ids = getattr(topk_output, "packed_topk_ids", None)
+    if packed_topk_ids is None:
+        packed_topk_ids = fused_pack_topk(
+            topk_ids=topk_ids,
+            topk_weights=topk_weights,
+        )
+
+    with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()):
+        direct_down_output = torch.empty(
+            hidden_states.shape[0],
+            hidden_states.shape[1],
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
+
+    # No pre-op join: the trtllm op waits on lora_event right before its activation kernel,
+    # so permute+GEMM1 run concurrent with the side-stream LoRA. Keep the event alive through
+    # cuda-graph capture so the captured cross-stream wait isn't torn down before instantiation.
+    if torch.cuda.is_current_stream_capturing():
+        _LORA_OVERLAP_EVENTS.append(lora_event)
+    lora_ready_handle = lora_event.cuda_event
+
+    # Down-LoRA/finalize overlap (env-gated): the op records gemm2_done_event right after the
+    # base down GEMM (before finalize); the side stream waits on it and runs ONLY the down-proj
+    # LoRA shrink (gemm A) + routing prep concurrent with finalizeKernel. The expand-add
+    # (gemm B) atomic-adds into `output` -- which finalize WRITES concurrently -- so it stays
+    # on the main stream after the op (post-finalize), exactly like the serial path.
+    # DISABLED: the down/finalize overlap is bench-verified net-neutral-to-negative AND
+    # corrupts the base/decode path — the captured gemm2_done cross-stream event under
+    # cuda-graph replay perturbs no-active-LoRA (base) requests (qwen base gsm8k 0.81 -> 0.56
+    # with it on; bisect-confirmed). The serial down-LoRA path below is used unconditionally.
+    down_overlap = False
+    gemm2_done_handle = 0
+    if down_overlap:
+        gemm2_done_event = torch.cuda.Event()
+        # Materialize the underlying cudaEvent (torch creates it lazily on first record) so
+        # .cuda_event is a real handle; the op re-records it after GEMM2.
+        gemm2_done_event.record()
+        gemm2_done_handle = gemm2_done_event.cuda_event
+
+    moe_result = trtllm_fp8_block_scale_routed_moe_lora(
+        topk_ids=packed_topk_ids,
+        routing_bias=None,
+        hidden_states=a_q,
+        hidden_states_scale=a_sf_t,
+        gemm1_weights=quant_info.w13_weight,
+        gemm1_weights_scale=quant_info.w13_weight_scale_inv,
+        gemm2_weights=quant_info.w2_weight,
+        gemm2_weights_scale=quant_info.w2_weight_scale_inv,
+        gate_up_lora_delta=gate_up_delta,
+        activation_lora_input=activation_lora_input,
+        lora_ready_event=lora_ready_handle,
+        num_experts=quant_info.global_num_experts,
+        top_k=runner_config.top_k,
+        n_group=None,
+        topk_group=None,
+        intermediate_size=quant_info.intermediate_size,
+        local_expert_offset=quant_info.local_expert_offset,
+        local_num_experts=quant_info.local_num_experts,
+        routed_scaling_factor=(
+            runner_config.routed_scaling_factor
+            if runner_config.routed_scaling_factor is not None
+            else 1.0
+        ),
+        routing_method_type=(
+            RoutingMethodType.TopK
+            if quant_info.routing_method_type == RoutingMethodType.DeepSeekV3
+            else quant_info.routing_method_type
+        ),
+        use_shuffled_weight=False,
+        do_finalize=True,
+        output=direct_down_output,
+        tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
+        fp8_quantization_type=Fp8QuantizationType.DeepSeekFp8,
+        activation_type=quant_info.activation_type,
+        gemm2_done_event=gemm2_done_handle,
+    )
+
+    output = moe_result
+
+    def _run_down_lora(
+        out, stage="all", intermediate_buffer=None, expand_wait_event=None
+    ):
+        return merged_experts_fused_moe_lora_add(
+            output=out,
+            hidden_states=activation_lora_input.view(-1, quant_info.intermediate_size),
+            lora_a=lora_info.down_lora_a_weights,
+            lora_b=lora_info.down_lora_b_weights,
+            topk_ids=topk_ids,
+            topk_weights=topk_weights,
+            token_lora_mapping=token_lora_mapping,
+            mul_routed_weight=True,
+            experts_shared_outer_loras_a=False,
+            experts_shared_outer_loras_b=lora_info.experts_shared_outer_loras,
+            routing_cache=fused_lora_routing_cache,
+            fuse_add_to_output=False,
+            fuse_sum_all_reduce=True,
+            use_direct_expand_add=lora_info.max_lora_rank <= 64,
+            local_expert_offset=quant_info.local_expert_offset,
+            local_num_experts=quant_info.local_num_experts,
+            stage=stage,
+            intermediate_buffer=intermediate_buffer,
+            expand_wait_event=expand_wait_event,
+        )
+
+    # Shared-add overlap: the trtllm op above already finalized `output`, so the
+    # staged shared-expert add (if any) can run on the producer (main) stream
+    # concurrent with the down-LoRA shrink below; the expand waits on it via
+    # expand_wait_event before atomic-adding into the same buffer.
+    shared_add_done = maybe_overlap_staged_shared_add(output)
+
+    if down_overlap:
+        # Fork at "base down GEMM done": ONLY the shrink (gemm A) + routing prep run on the
+        # side stream, concurrent with the main-stream finalizeKernel. The expand-add (gemm B)
+        # joins back on the MAIN stream after the op -- i.e. strictly after finalize wrote
+        # `output` -- and atomic-adds into it exactly like the serial path: same kernels,
+        # same buffers, identical numerics; the shrink just starts earlier.
+        # The shrink intermediate is allocated HERE (main = consumer stream of the expand),
+        # per the SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC lesson on side-stream allocations.
+        down_intermediate = hidden_states.new_empty(
+            (
+                hidden_states.shape[0],
+                topk_ids.shape[1],
+                lora_info.down_lora_a_weights.shape[2],
+            )
+        )
+        side_stream.wait_event(gemm2_done_event)
+        shrink_done_event = torch.cuda.Event()
+        with torch.cuda.stream(side_stream):
+            _run_down_lora(
+                output, stage="shrink", intermediate_buffer=down_intermediate
+            )
+            shrink_done_event.record()
+        if torch.cuda.is_current_stream_capturing():
+            _LORA_OVERLAP_EVENTS.append(gemm2_done_event)
+            _LORA_OVERLAP_EVENTS.append(shrink_done_event)
+        torch.cuda.current_stream().wait_event(shrink_done_event)
+        _run_down_lora(
+            output,
+            stage="expand",
+            intermediate_buffer=down_intermediate,
+            expand_wait_event=shared_add_done,
+        )
+    else:
+        _run_down_lora(output, expand_wait_event=shared_add_done)
+    return StandardCombineInput(hidden_states=output)
+
+
+def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora_two_stream(
+    dispatch_output,
+    quant_info,
+    runner_config,
+    lora_info,
+):
+    """Two-stream NVFP4 sibling of the FP8 two-stream MoE LoRA dispatch.
+
+    Fires only for virtual-experts LoRA + decode-shaped batches; everything else
+    delegates to the saved-original single-stream FP4 dispatch (byte-identical).
+    """
+    hidden_states = dispatch_output.hidden_states
+
+    use_virtual_lora_store = bool(
+        lora_info.lora_use_virtual_experts and lora_info.max_lora_rank > 0
+    )
+    if not (use_virtual_lora_store and is_two_stream_active(hidden_states)):
+        return get_original_fp4_moe_lora_func()(
+            dispatch_output, quant_info, runner_config, lora_info
+        )
+
+    # ---- two-stream fast path ----
+    from sglang.jit_kernel.trtllm_lora_temp import (
+        trtllm_fp4_block_scale_routed_moe_lora,
+    )
+    from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk
+    from sglang.kernels.ops.moe.trtllm_lora_temp.virtual_experts import (
+        merged_experts_fused_moe_lora_add,
+    )
+    from sglang.srt.distributed import get_tp_group
+    from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+        use_symmetric_memory,
+    )
+    from sglang.srt.layers.dp_attention import is_allocation_symmetric
+    from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
+    from sglang.srt.layers.moe.topk import TopKOutputChecker
+
+    assert (
+        runner_config.activation == "silu" and runner_config.is_gated
+    ), "experimental_sgl_trtllm NVFP4 LoRA currently supports the gated SwiGLU path only."
+    topk_output = dispatch_output.topk_output
+    assert TopKOutputChecker.format_is_standard(topk_output)
+    assert runner_config.top_k is not None
+
+    topk_ids = topk_output.topk_ids
+    topk_weights = topk_output.topk_weights
+    token_lora_mapping = lora_info.token_lora_mapping
+    fused_lora_routing_cache: dict = {}
+
+    # Down-proj LoRA runs serially on the main stream (after the trtllm op) by default. The old
+    # side-stream down-overlap was removed: bench-verified net-neutral-to-negative (the extra
+    # side-stream all-reduce cancels any overlap gain), AND its act_ready_event cross-stream sync
+    # corrupted decode state under sustained heavy LoRA load (cuda-graph replay -> persistent
+    # garbage). SGLANG_OPT_LORA_DOWN_FINALIZE_OVERLAP=1 re-introduces a more conservative variant:
+    # fork at gemm2_done (not act_ready), and only the SHRINK (gemm A) overlaps the finalize
+    # kernel -- the expand-add (gemm B) stays on the MAIN stream post-finalize (it writes the
+    # same `output` finalize writes), so its kernels/numerics match the serial path exactly.
+    inter = quant_info.intermediate_size_per_partition
+    side_stream = get_lora_side_stream()
+
+    gate_up_delta = hidden_states.new_empty(
+        (hidden_states.shape[0], runner_config.top_k, quant_info.w13_weight.shape[1])
+    )
+
+    def _run_gate_up_lora():
+        merged_experts_fused_moe_lora_add(
+            output=gate_up_delta,
+            hidden_states=hidden_states,
+            lora_a=lora_info.gate_up_lora_a_weights,
+            lora_b=lora_info.gate_up_lora_b_weights,
+            topk_ids=topk_ids,
+            topk_weights=topk_weights,
+            token_lora_mapping=token_lora_mapping,
+            mul_routed_weight=False,
+            experts_shared_outer_loras_a=lora_info.experts_shared_outer_loras,
+            experts_shared_outer_loras_b=False,
+            routing_cache=fused_lora_routing_cache,
+            fuse_add_to_output=False,
+            use_direct_expand_add=lora_info.max_lora_rank <= 64,
+            local_expert_offset=quant_info.local_expert_offset,
+            local_num_experts=quant_info.local_num_experts,
+            intermediate_buffer=gate_up_lora_intermediate,
+        )
+
+    # Hoist every side-chain allocation onto the MAIN stream (cuda-graph
+    # allocator safety -- see the "routing" stage in virtual_experts.py).
+    merged_experts_fused_moe_lora_add(
+        output=gate_up_delta,
+        hidden_states=hidden_states,
+        lora_a=lora_info.gate_up_lora_a_weights,
+        lora_b=lora_info.gate_up_lora_b_weights,
+        topk_ids=topk_ids,
+        topk_weights=topk_weights,
+        token_lora_mapping=token_lora_mapping,
+        mul_routed_weight=False,
+        experts_shared_outer_loras_a=lora_info.experts_shared_outer_loras,
+        experts_shared_outer_loras_b=False,
+        routing_cache=fused_lora_routing_cache,
+        stage="routing",
+        local_expert_offset=quant_info.local_expert_offset,
+        local_num_experts=quant_info.local_num_experts,
+    )
+    gate_up_lora_intermediate = hidden_states.new_empty(
+        (
+            hidden_states.shape[0],
+            topk_ids.shape[1],
+            lora_info.gate_up_lora_a_weights.shape[2],
+        )
+    )
+
+    # O1-fp4 fork: gate_up shrink/expand on the side stream, concurrent with the
+    # FP4 op's permute + gate_up GEMM1 below. The op waits on lora_event right
+    # before its activation kernel (the only consumer of gate_up_delta).
+    lora_event = torch.cuda.Event()
+    side_stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(side_stream):
+        _run_gate_up_lora()
+        lora_event.record()
+
+    activation_lora_input = torch.empty(
+        (hidden_states.shape[0], runner_config.top_k, inter),
+        dtype=hidden_states.dtype,
+        device=hidden_states.device,
+    )
+    packed_topk_ids = fused_pack_topk(
+        topk_ids=topk_ids,
+        topk_weights=topk_weights,
+    )
+    with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()):
+        direct_down_output = torch.empty(
+            hidden_states.shape[0],
+            hidden_states.shape[1],
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
+
+    # Keep the event alive through cuda-graph capture so the captured wait inside
+    # the FP4 op isn't torn down before instantiation (eager relies on deferred destroy).
+    if torch.cuda.is_current_stream_capturing():
+        _LORA_OVERLAP_EVENTS.append(lora_event)
+    lora_ready_handle = lora_event.cuda_event
+
+    # Down-LoRA/finalize overlap (env-gated, see the comment above): record gemm2_done inside
+    # the op; the side stream runs only the down-LoRA shrink + routing prep concurrent with
+    # finalize; the expand-add joins back on the main stream after the op.
+    # DISABLED: the down/finalize overlap is bench-verified net-neutral-to-negative AND
+    # corrupts the base/decode path — the captured gemm2_done cross-stream event under
+    # cuda-graph replay perturbs no-active-LoRA (base) requests (qwen base gsm8k 0.81 -> 0.56
+    # with it on; bisect-confirmed). The serial down-LoRA path below is used unconditionally.
+    down_overlap = False
+    gemm2_done_handle = 0
+    if down_overlap:
+        gemm2_done_event = torch.cuda.Event()
+        # Materialize the underlying cudaEvent (torch creates it lazily on first record) so
+        # .cuda_event is a real handle; the op re-records it after the down GEMM.
+        gemm2_done_event.record()
+        gemm2_done_handle = gemm2_done_event.cuda_event
+
+    output = trtllm_fp4_block_scale_routed_moe_lora(
+        topk_ids=packed_topk_ids,
+        routing_bias=None,
+        hidden_states=hidden_states,
+        hidden_states_scale=None,
+        gemm1_weights=quant_info.w13_weight,
+        gemm1_weights_scale=quant_info.w13_weight_scale.view(torch.float8_e4m3fn),
+        gemm2_weights=quant_info.w2_weight,
+        gemm2_weights_scale=quant_info.w2_weight_scale.view(torch.float8_e4m3fn),
+        output1_scales_scalar=quant_info.g1_scale_c,
+        output1_scales_gate_scalar=quant_info.g1_alphas,
+        output2_scales_scalar=quant_info.g2_alphas,
+        gate_up_lora_delta=gate_up_delta,
+        activation_lora_input=activation_lora_input,
+        lora_ready_event=lora_ready_handle,
+        num_experts=quant_info.global_num_experts,
+        top_k=runner_config.top_k,
+        intermediate_size=inter,
+        local_expert_offset=quant_info.local_expert_offset,
+        local_num_experts=quant_info.local_num_experts,
+        routed_scaling_factor=(
+            runner_config.routed_scaling_factor
+            if runner_config.routed_scaling_factor is not None
+            else 1.0
+        ),
+        routing_method_type=quant_info.routing_method_type,
+        do_finalize=True,
+        output=direct_down_output,
+        gemm2_done_event=gemm2_done_handle,
+    )
+
+    def _run_down_lora(out, stage="all", intermediate_buffer=None):
+        return merged_experts_fused_moe_lora_add(
+            output=out,
+            hidden_states=activation_lora_input.view(-1, inter),
+            lora_a=lora_info.down_lora_a_weights,
+            lora_b=lora_info.down_lora_b_weights,
+            topk_ids=topk_ids,
+            topk_weights=topk_weights,
+            token_lora_mapping=token_lora_mapping,
+            mul_routed_weight=True,
+            experts_shared_outer_loras_a=False,
+            experts_shared_outer_loras_b=lora_info.experts_shared_outer_loras,
+            routing_cache=fused_lora_routing_cache,
+            fuse_add_to_output=False,
+            fuse_sum_all_reduce=True,
+            use_direct_expand_add=lora_info.max_lora_rank <= 64,
+            local_expert_offset=quant_info.local_expert_offset,
+            local_num_experts=quant_info.local_num_experts,
+            stage=stage,
+            intermediate_buffer=intermediate_buffer,
+        )
+
+    if down_overlap:
+        # Fork at "base down GEMM done": shrink (gemm A) + routing prep on the side stream
+        # concurrent with the main-stream finalize; the expand-add (gemm B) joins back on the
+        # MAIN stream after the op (strictly post-finalize) -- serial-path kernels/numerics.
+        # Intermediate allocated on main (= consumer stream of the expand), per the
+        # SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC lesson.
+        down_intermediate = hidden_states.new_empty(
+            (
+                hidden_states.shape[0],
+                topk_ids.shape[1],
+                lora_info.down_lora_a_weights.shape[2],
+            )
+        )
+        side_stream.wait_event(gemm2_done_event)
+        shrink_done_event = torch.cuda.Event()
+        with torch.cuda.stream(side_stream):
+            _run_down_lora(
+                output, stage="shrink", intermediate_buffer=down_intermediate
+            )
+            shrink_done_event.record()
+        if torch.cuda.is_current_stream_capturing():
+            _LORA_OVERLAP_EVENTS.append(gemm2_done_event)
+            _LORA_OVERLAP_EVENTS.append(shrink_done_event)
+        torch.cuda.current_stream().wait_event(shrink_done_event)
+        _run_down_lora(output, stage="expand", intermediate_buffer=down_intermediate)
+    else:
+        _run_down_lora(output)
+    return StandardCombineInput(hidden_states=output)
+
+
+def fused_experts_none_to_experimental_sgl_trtllm_bf16_lora_two_stream(
+    dispatch_output,
+    quant_info,
+    runner_config,
+    lora_info,
+):
+    """Two-stream BF16 sibling of the FP8/FP4 two-stream MoE LoRA dispatches.
+
+    O1-bf16 fork: the gate_up LoRA shrink/expand runs on the side stream
+    concurrent with the bf16 op's routing + permute + gate_up GEMM; the op
+    waits on ``lora_ready_event`` right before its activation kernel (the only
+    consumer of ``gate_up_delta``). Fires only for virtual-experts LoRA +
+    decode-shaped batches; everything else delegates to the saved-original
+    single-stream bf16 dispatch (byte-identical). Down-LoRA stays serial on
+    the main stream — the down/finalize overlap was bench-verified
+    net-neutral-to-negative on the FP8/FP4 paths and corrupted the base
+    decode path under cuda-graph replay (see the comment in the FP4 variant).
+    """
+    hidden_states = dispatch_output.hidden_states
+
+    use_virtual_lora_store = bool(
+        lora_info.lora_use_virtual_experts and lora_info.max_lora_rank > 0
+    )
+    if not (use_virtual_lora_store and is_two_stream_active(hidden_states)):
+        return get_original_bf16_moe_lora_func()(
+            dispatch_output, quant_info, runner_config, lora_info
+        )
+
+    # ---- two-stream fast path ----
+    from sglang.jit_kernel.trtllm_lora_temp import trtllm_bf16_routed_moe_lora
+    from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk
+    from sglang.kernels.ops.moe.trtllm_lora_temp.virtual_experts import (
+        merged_experts_fused_moe_lora_add,
+    )
+    from sglang.srt.distributed import get_tp_group
+    from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+        use_symmetric_memory,
+    )
+    from sglang.srt.layers.dp_attention import is_allocation_symmetric
+    from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
+        get_activation_type,
+    )
+    from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
+    from sglang.srt.layers.moe.topk import TopKOutputChecker
+    from sglang.srt.layers.moe.utils import RoutingMethodType
+
+    assert (
+        runner_config.activation == "silu" and runner_config.is_gated
+    ), "experimental_sgl_trtllm BF16 LoRA currently supports the gated SwiGLU path only."
+    topk_output = dispatch_output.topk_output
+    assert TopKOutputChecker.format_is_standard(topk_output)
+    assert runner_config.top_k is not None
+
+    topk_ids = topk_output.topk_ids
+    topk_weights = topk_output.topk_weights
+    token_lora_mapping = lora_info.token_lora_mapping
+    fused_lora_routing_cache: dict = {}
+
+    inter = runner_config.intermediate_size_per_partition
+    side_stream = get_lora_side_stream()
+
+    gate_up_delta = hidden_states.new_empty(
+        (hidden_states.shape[0], runner_config.top_k, 2 * inter)
+    )
+
+    def _run_gate_up_lora():
+        merged_experts_fused_moe_lora_add(
+            output=gate_up_delta,
+            hidden_states=hidden_states,
+            lora_a=lora_info.gate_up_lora_a_weights,
+            lora_b=lora_info.gate_up_lora_b_weights,
+            topk_ids=topk_ids,
+            topk_weights=topk_weights,
+            token_lora_mapping=token_lora_mapping,
+            mul_routed_weight=False,
+            experts_shared_outer_loras_a=lora_info.experts_shared_outer_loras,
+            experts_shared_outer_loras_b=False,
+            routing_cache=fused_lora_routing_cache,
+            fuse_add_to_output=False,
+            use_direct_expand_add=lora_info.max_lora_rank <= 64,
+            local_expert_offset=quant_info.local_expert_offset,
+            local_num_experts=runner_config.num_local_experts,
+            intermediate_buffer=gate_up_lora_intermediate,
+        )
+
+    # O1-bf16 fork: gate_up shrink/expand on the side stream, concurrent with the
+    # bf16 op's routing + permute + gate_up GEMM below. The op waits on lora_event
+    # right before its activation kernel (the only consumer of gate_up_delta).
+    lora_event = torch.cuda.Event()
+
+    # Hoist every side-chain allocation onto the MAIN stream (cuda-graph allocator
+    # safety -- see the "routing" stage in virtual_experts.py): pre-warm the routing
+    # cache and pre-allocate the shrink intermediate here, so the side-stream block
+    # below launches kernels only. Without this, the routing tensors + shrink
+    # intermediate get allocated inside the side-stream context during cuda-graph
+    # capture, where cross-stream tracking is off -> pool blocks reused with no graph
+    # edge -> '!!!!' decode corruption at max-loras>=2 (matches the fp8/fp4 fix).
+    merged_experts_fused_moe_lora_add(
+        output=gate_up_delta,
+        hidden_states=hidden_states,
+        lora_a=lora_info.gate_up_lora_a_weights,
+        lora_b=lora_info.gate_up_lora_b_weights,
+        topk_ids=topk_ids,
+        topk_weights=topk_weights,
+        token_lora_mapping=token_lora_mapping,
+        mul_routed_weight=False,
+        experts_shared_outer_loras_a=lora_info.experts_shared_outer_loras,
+        experts_shared_outer_loras_b=False,
+        routing_cache=fused_lora_routing_cache,
+        stage="routing",
+        local_expert_offset=quant_info.local_expert_offset,
+        local_num_experts=runner_config.num_local_experts,
+    )
+    gate_up_lora_intermediate = hidden_states.new_empty(
+        (
+            hidden_states.shape[0],
+            topk_ids.shape[1],
+            lora_info.gate_up_lora_a_weights.shape[2],
+        )
+    )
+    side_stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(side_stream):
+        _run_gate_up_lora()
+        lora_event.record()
+
+    activation_lora_input = torch.empty(
+        (hidden_states.shape[0], runner_config.top_k, inter),
+        dtype=hidden_states.dtype,
+        device=hidden_states.device,
+    )
+    packed_topk_ids = getattr(topk_output, "packed_topk_ids", None)
+    if packed_topk_ids is None:
+        packed_topk_ids = fused_pack_topk(
+            topk_ids=topk_ids,
+            topk_weights=topk_weights,
+        )
+
+    routing_method_type = runner_config.routing_method_type
+    if routing_method_type is None:
+        routing_method_type = RoutingMethodType.Default
+    elif routing_method_type == RoutingMethodType.DeepSeekV3:
+        routing_method_type = RoutingMethodType.TopK
+
+    with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()):
+        direct_down_output = torch.empty(
+            hidden_states.shape[0],
+            hidden_states.shape[1],
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
+
+    # Keep the event alive through cuda-graph capture so the captured wait inside
+    # the bf16 op isn't torn down before instantiation (eager relies on deferred destroy).
+    if torch.cuda.is_current_stream_capturing():
+        _LORA_OVERLAP_EVENTS.append(lora_event)
+    lora_ready_handle = lora_event.cuda_event
+
+    output = trtllm_bf16_routed_moe_lora(
+        topk_ids=packed_topk_ids,
+        routing_bias=None,
+        hidden_states=hidden_states,
+        gemm1_weights=quant_info.gemm1_weights,
+        gemm2_weights=quant_info.gemm2_weights,
+        gate_up_lora_delta=gate_up_delta,
+        activation_lora_input=activation_lora_input,
+        num_experts=quant_info.global_num_experts,
+        top_k=runner_config.top_k,
+        intermediate_size=inter,
+        local_expert_offset=quant_info.local_expert_offset,
+        local_num_experts=runner_config.num_local_experts,
+        routed_scaling_factor=(
+            runner_config.routed_scaling_factor
+            if runner_config.routed_scaling_factor is not None
+            else 1.0
+        ),
+        routing_method_type=routing_method_type,
+        do_finalize=True,
+        output=direct_down_output,
+        activation_type=get_activation_type(
+            runner_config.activation, is_gated=runner_config.is_gated
+        ),
+        lora_ready_event=lora_ready_handle,
+        # Down-LoRA/finalize overlap intentionally NOT wired (gemm2_done_event=0):
+        # bench-verified net-neutral-to-negative on FP8/FP4 and corrupts the base
+        # decode path under cuda-graph replay. Serial down-LoRA below.
+        gemm2_done_event=0,
+    )
+
+    merged_experts_fused_moe_lora_add(
+        output=output,
+        hidden_states=activation_lora_input.view(-1, inter),
+        lora_a=lora_info.down_lora_a_weights,
+        lora_b=lora_info.down_lora_b_weights,
+        topk_ids=topk_ids,
+        topk_weights=topk_weights,
+        token_lora_mapping=token_lora_mapping,
+        mul_routed_weight=True,
+        experts_shared_outer_loras_a=False,
+        experts_shared_outer_loras_b=lora_info.experts_shared_outer_loras,
+        routing_cache=fused_lora_routing_cache,
+        fuse_add_to_output=False,
+        fuse_sum_all_reduce=True,
+        use_direct_expand_add=lora_info.max_lora_rank <= 64,
+        local_expert_offset=quant_info.local_expert_offset,
+        local_num_experts=runner_config.num_local_experts,
+    )
+    return StandardCombineInput(hidden_states=output)

python/sglang/srt/lora/trtllm_lora_temp/sgl_backend.py

@@ -0,0 +1,33 @@
+"""Registers the ``experimental_sgl_trtllm`` MoE fused-func.
+
+``MoeRunner.__init__`` requires a registered fused-func at CONSTRUCTION time even
+for the LoRA case, because LoRA is attached *after* the MoE layer is built (so
+``lora_enabled`` is False inside ``MoeRunner.__init__``). At run time the runner
+skips this for the LoRA path; the no-LoRA path delegates entirely to the upstream
+flashinfer_trtllm dispatch (all quant types), so no-LoRA is identical to the stock backend.
+
+Registration fires at model-build time via a one-line import of this module in
+``moe_runner/flashinfer_trtllm.py`` (the module already imported there for the
+trtllm weight-prep). Keeping the dispatch body here keeps that file otherwise
+pristine; the sgl FP8 LoRA dispatch lives in ``sgl_fp8_moe.py`` (used only by the LoRA path).
+"""
+
+from sglang.srt.layers.moe.moe_runner.base import register_fused_func
+
+
+@register_fused_func("none", "experimental_sgl_trtllm")
+def fused_experts_none_to_experimental_sgl_trtllm(
+    dispatch_output, quant_info, runner_config
+):
+    # No-LoRA on the experimental_sgl_trtllm backend == upstream flashinfer_trtllm for EVERY
+    # quant type (FP8 / NVFP4 / bf16). When LoRA is disabled the runner calls this fused-func,
+    # so delegating entirely to upstream keeps the no-LoRA path byte-identical to the stock
+    # backend. The new sgl kernels (sgl_fp8_moe, trtllm_*_routed_moe_lora) run ONLY on the LoRA
+    # dispatch (lora_dispatch.py), never here.
+    from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
+        fused_experts_none_to_flashinfer_trtllm,
+    )
+
+    return fused_experts_none_to_flashinfer_trtllm(
+        dispatch_output, quant_info, runner_config
+    )

python/sglang/srt/lora/trtllm_lora_temp/sgl_fp8_moe.py

@@ -0,0 +1,248 @@
+"""Copy of upstream flashinfer-trtllm FP8 MoE dispatch, wired to experimental_sgl_trtllm_moe
+block-scale wrappers (LoRA-capable) so moe_runner/flashinfer_trtllm.py stays pristine. Body is
+verbatim from upstream; helper imports are call-time (cycle-safe); two FP8 wrappers shadowed.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
+    from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
+        FlashInferTrtllmFp8MoeQuantInfo,
+    )
+    from sglang.srt.layers.moe.token_dispatcher.standard import (
+        StandardCombineInput,
+        StandardDispatchOutput,
+    )
+
+
+def fused_experts_fp8_sgl(
+    dispatch_output: StandardDispatchOutput,
+    quant_info: FlashInferTrtllmFp8MoeQuantInfo,
+    runner_config: MoeRunnerConfig,
+    use_routed_topk: bool = False,
+) -> StandardCombineInput:
+    # Lazy (call-time) imports so this module never triggers the flashinfer_trtllm
+    # <-> quantization import cycle at load time.
+    from flashinfer.fused_moe import Fp8QuantizationType
+
+    from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk
+    from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
+        get_tp_group,
+        is_allocation_symmetric,
+        next_power_of_2,
+        per_token_group_quant_fp8,
+        scaled_fp8_quant,
+        trtllm_fp8_per_tensor_scale_moe_wrapper,
+        use_symmetric_memory,
+    )
+    from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
+    from sglang.srt.layers.moe.topk import TopKOutputChecker
+    from sglang.srt.layers.moe.utils import RoutingMethodType
+    from sglang.srt.lora.trtllm_lora_temp.experimental_sgl_trtllm_moe import (
+        sgl_trtllm_fp8_block_scale_moe_wrapper as trtllm_fp8_block_scale_moe_wrapper,
+    )
+    from sglang.srt.lora.trtllm_lora_temp.experimental_sgl_trtllm_moe import (
+        sgl_trtllm_fp8_block_scale_routed_moe_wrapper as trtllm_fp8_block_scale_routed_moe_wrapper,
+    )
+
+    _SUPPORTED_FP8_ACTIVATIONS = {"silu", "relu2"}
+    assert runner_config.activation in _SUPPORTED_FP8_ACTIVATIONS, (
+        f"Only {_SUPPORTED_FP8_ACTIVATIONS} are supported for FP8 MoE, "
+        f"got '{runner_config.activation}'."
+    )
+    assert not runner_config.no_combine, "no_combine is not supported for flashinfer."
+
+    hidden_states = dispatch_output.hidden_states
+    topk_output = dispatch_output.topk_output
+    if TopKOutputChecker.format_is_bypassed(topk_output):
+        router_logits = topk_output.router_logits
+        topk_config = topk_output.topk_config
+        correction_bias = (
+            None
+            if topk_config.correction_bias is None
+            else topk_config.correction_bias.to(hidden_states.dtype)
+        )
+    else:
+        router_logits = None
+        topk_config = None
+        correction_bias = None
+
+    routing_method_type = quant_info.routing_method_type
+    fp8_quantization_type = (
+        Fp8QuantizationType.MxFp8
+        if quant_info.use_mxfp8
+        else Fp8QuantizationType.DeepSeekFp8
+    )
+    use_shuffled_weight = quant_info.use_mxfp8
+
+    if quant_info.block_quant:
+        assert quant_info.weight_block_k is not None
+        assert quant_info.w13_weight_scale_inv is not None
+        assert quant_info.w2_weight_scale_inv is not None
+
+        if quant_info.use_mxfp8:
+            assert quant_info.weight_block_k == 32
+            from flashinfer import mxfp8_quantize
+
+            a_q, a_sf = mxfp8_quantize(hidden_states, False)
+            # FlashInfer TRT-LLM MxFP8 expects token-major activation scales:
+            # [num_tokens, hidden_size // 32] (no transpose).
+            a_sf_t = a_sf.view(torch.uint8).reshape(hidden_states.shape[0], -1)
+        else:
+            a_q, a_sf = per_token_group_quant_fp8(
+                hidden_states, quant_info.weight_block_k
+            )
+            a_sf_t = a_sf.t().contiguous()
+
+        # Allocate output inside symmetric memory context
+        with use_symmetric_memory(
+            get_tp_group(), disabled=not is_allocation_symmetric()
+        ):
+            symm_output = torch.empty(
+                hidden_states.shape[0],
+                hidden_states.shape[1],
+                dtype=hidden_states.dtype,
+                device=hidden_states.device,
+            )
+
+        # Move kernel call outside context manager to avoid graph breaks
+        # during torch.compile for piecewise cuda graph.
+        # Use custom op wrapper for torch.compile compatibility.
+        if use_routed_topk:
+            assert (
+                runner_config.top_k is not None
+            ), "runner_config.top_k is required for flashinfer_trtllm_routed."
+            assert TopKOutputChecker.format_is_standard(topk_output)
+            packed_topk_ids = fused_pack_topk(
+                topk_ids=topk_output.topk_ids,
+                topk_weights=topk_output.topk_weights,
+            )
+
+            output = trtllm_fp8_block_scale_routed_moe_wrapper(
+                topk_ids=packed_topk_ids,
+                routing_bias=None,
+                hidden_states=a_q,
+                hidden_states_scale=a_sf_t,
+                gemm1_weights=quant_info.w13_weight,
+                gemm1_weights_scale=quant_info.w13_weight_scale_inv,
+                gemm2_weights=quant_info.w2_weight,
+                gemm2_weights_scale=quant_info.w2_weight_scale_inv,
+                num_experts=quant_info.global_num_experts,
+                top_k=runner_config.top_k,
+                n_group=None,
+                topk_group=None,
+                intermediate_size=quant_info.intermediate_size,
+                local_expert_offset=quant_info.local_expert_offset,
+                local_num_experts=quant_info.local_num_experts,
+                routed_scaling_factor=(
+                    runner_config.routed_scaling_factor
+                    if runner_config.routed_scaling_factor is not None
+                    else 1.0
+                ),
+                routing_method_type=(
+                    RoutingMethodType.TopK
+                    if routing_method_type == RoutingMethodType.DeepSeekV3
+                    else routing_method_type
+                ),
+                use_shuffled_weight=use_shuffled_weight,
+                tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
+                fp8_quantization_type=int(fp8_quantization_type),
+                activation_type=quant_info.activation_type,
+            )
+        else:
+            assert TopKOutputChecker.format_is_bypassed(topk_output)
+
+            output = trtllm_fp8_block_scale_moe_wrapper(
+                routing_logits=router_logits,
+                routing_bias=correction_bias,
+                hidden_states=a_q,
+                hidden_states_scale=a_sf_t,
+                gemm1_weights=quant_info.w13_weight,
+                gemm1_weights_scale=quant_info.w13_weight_scale_inv,
+                gemm2_weights=quant_info.w2_weight,
+                gemm2_weights_scale=quant_info.w2_weight_scale_inv,
+                num_experts=quant_info.global_num_experts,
+                top_k=topk_config.top_k,
+                n_group=topk_config.num_expert_group,
+                topk_group=topk_config.topk_group,
+                intermediate_size=quant_info.intermediate_size,
+                local_expert_offset=quant_info.local_expert_offset,
+                local_num_experts=quant_info.local_num_experts,
+                routed_scaling_factor=(
+                    runner_config.routed_scaling_factor
+                    if runner_config.routed_scaling_factor is not None
+                    else 1.0
+                ),
+                routing_method_type=routing_method_type,
+                use_shuffled_weight=use_shuffled_weight,
+                tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
+                fp8_quantization_type=int(fp8_quantization_type),
+                activation_type=quant_info.activation_type,
+            )
+        # TODO: Once https://github.com/flashinfer-ai/flashinfer/issues/2703 is fixed, pass output to moe kernel and remove this copy.
+        symm_output.copy_(output)
+        output = symm_output
+    else:
+        assert TopKOutputChecker.format_is_bypassed(topk_output)
+        assert quant_info.w13_input_scale is not None
+        assert quant_info.output1_scales_scalar is not None
+        assert quant_info.output1_scales_gate_scalar is not None
+        assert quant_info.output2_scales_scalar is not None
+
+        a_q, _ = scaled_fp8_quant(hidden_states, quant_info.w13_input_scale)
+        routing_bias_cast = (
+            None if correction_bias is None else correction_bias.to(torch.bfloat16)
+        )
+
+        # Allocate output inside symmetric memory context
+        with use_symmetric_memory(
+            get_tp_group(), disabled=not is_allocation_symmetric()
+        ):
+            symm_output = torch.empty(
+                hidden_states.shape[0],
+                hidden_states.shape[1],
+                dtype=torch.bfloat16,
+                device=hidden_states.device,
+            )
+
+        # Move kernel call outside context manager to avoid graph breaks
+        # during torch.compile for piecewise cuda graph.
+        # Use custom op wrapper for torch.compile compatibility.
+
+        router_logits = router_logits.to(torch.bfloat16)
+
+        output = trtllm_fp8_per_tensor_scale_moe_wrapper(
+            routing_logits=router_logits,
+            routing_bias=routing_bias_cast,
+            hidden_states=a_q,
+            gemm1_weights=quant_info.w13_weight,
+            output1_scales_scalar=quant_info.output1_scales_scalar,
+            output1_scales_gate_scalar=quant_info.output1_scales_gate_scalar,
+            gemm2_weights=quant_info.w2_weight,
+            output2_scales_scalar=quant_info.output2_scales_scalar,
+            num_experts=quant_info.global_num_experts,
+            top_k=topk_config.top_k,
+            n_group=topk_config.num_expert_group,
+            topk_group=topk_config.topk_group,
+            intermediate_size=int(quant_info.w2_weight.shape[2]),
+            local_expert_offset=quant_info.local_expert_offset,
+            local_num_experts=quant_info.local_num_experts,
+            routed_scaling_factor=(
+                runner_config.routed_scaling_factor
+                if runner_config.routed_scaling_factor is not None
+                else 1.0
+            ),
+            use_routing_scales_on_input=False,
+            routing_method_type=routing_method_type,
+            tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
+            activation_type=quant_info.activation_type,
+        )
+        symm_output.copy_(output)
+        output = symm_output
+
+    return StandardCombineInput(hidden_states=output)

python/sglang/srt/lora/trtllm_lora_temp/shared_add_overlap.py

@@ -0,0 +1,123 @@
+"""Overlap the MoE shared-expert add with the down-LoRA shrink (gemm A).
+
+In the Qwen dual-stream decode path (``Qwen2MoeSparseMoeBlock.forward_normal_dual_stream``)
+the shared expert runs on the main stream while the routed experts + trtllm LoRA
+MoE run on the alt stream; ``final_hidden_states += shared_output`` then runs on
+the main stream only after the WHOLE alt-stream chain (base MoE -> finalize ->
+down-LoRA shrink -> down-LoRA expand) joins — putting a ~2us elementwise add at
+the tail of every MoE layer's critical path.
+
+The add's real dependency is only the base-MoE finalize (the trtllm op with
+``do_finalize=True`` writing the output buffer): the down-LoRA shrink writes a
+separate intermediate, and the down-LoRA expand atomic-adds into the same output
+buffer, so addition order is commutative — the only hard constraint is that the
+non-atomic shared add and the expand must not run CONCURRENTLY on the buffer.
+(The cross-rank ``tensor_model_parallel_all_reduce`` happens after all of this
+at the model layer, unchanged.)
+
+Protocol (gated by ``SGLANG_OPT_LORA_SHARED_ADD_OVERLAP``):
+
+1. The model layer stages ``(shared_output, producer_stream)`` via
+   :func:`stage_shared_expert_add` after computing the shared expert and before
+   forking the routed experts to the alt stream.
+2. The LoRA dispatch, right after the trtllm op returns (finalize done), calls
+   :func:`maybe_overlap_staged_shared_add(output)`: it records ``base_ready`` on
+   the current (alt) stream, enqueues ``wait(base_ready); output += shared;
+   record(add_done)`` on the producer (main) stream — main-stream program order
+   already guarantees ``shared_output`` is ready there — and returns ``add_done``.
+3. The down-LoRA ``merged_experts_fused_moe_lora_add`` waits on ``add_done``
+   right before launching the expand kernel, so the shrink (+ stage-B routing)
+   overlaps the add and the expand never races it.
+4. After the dual-stream join the model layer calls
+   :func:`unstage_shared_expert_add`; if the dispatch never consumed the staging
+   (prefill / non-virtual-store / fallback paths) it gets the tensor back and
+   performs the original add itself — byte-identical fallback behavior.
+
+The state is a single slot: MoE layers run sequentially within one scheduler
+process, and the stage/consume pair lives within a single layer forward.
+"""
+
+from typing import Optional, Tuple
+
+import torch
+
+from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
+
+_PENDING: Optional[Tuple[torch.Tensor, torch.cuda.Stream]] = None
+
+# Keep events recorded during cuda-graph capture alive so the captured
+# cross-stream waits aren't torn down before graph instantiation (same pattern
+# as moe_overlap._LORA_OVERLAP_EVENTS). Eager runs rely on deferred destroy.
+_SHARED_ADD_EVENTS: list = []
+
+
+def shared_add_overlap_enabled() -> bool:
+    return lora_envs.SGLANG_OPT_LORA_SHARED_ADD_OVERLAP.get()
+
+
+def stage_shared_expert_add(
+    shared_output: torch.Tensor, producer_stream: torch.cuda.Stream
+) -> None:
+    """Stage the shared-expert output for the LoRA dispatch to add.
+
+    ``producer_stream`` is the stream ``shared_output`` was computed on (the
+    main stream); the overlapped add is enqueued there so its data dependency
+    on the shared expert is carried by stream program order.
+    """
+    global _PENDING
+    _PENDING = (shared_output, producer_stream)
+
+
+def unstage_shared_expert_add() -> Optional[torch.Tensor]:
+    """Reclaim a staged-but-unconsumed shared add (fallback paths).
+
+    Returns the staged tensor if the dispatch did NOT consume it (the model
+    layer must then do the add itself), or None if it was consumed (the add is
+    already enqueued on the producer stream).
+    """
+    global _PENDING
+    if _PENDING is None:
+        return None
+    shared_output, _ = _PENDING
+    _PENDING = None
+    return shared_output
+
+
+def maybe_overlap_staged_shared_add(output: torch.Tensor) -> Optional[torch.cuda.Event]:
+    """Enqueue the staged shared-expert add overlapped with the down-LoRA shrink.
+
+    Call from the LoRA dispatch right after the base-MoE finalize has been
+    enqueued on the current stream with ``output`` fully written. Returns the
+    ``add_done`` event the down-LoRA expand must wait on before atomic-adding
+    into ``output``, or None when nothing was staged.
+    """
+    global _PENDING
+    if _PENDING is None:
+        return None
+    shared_output, producer_stream = _PENDING
+    current_stream = torch.cuda.current_stream()
+    if producer_stream == current_stream:
+        # Single-stream caller: nothing to overlap. Leave the staging in place
+        # so the model layer reclaims it and does the add as before.
+        return None
+    if torch.cuda.is_current_stream_capturing():
+        # The cross-stream producer-stream add_ (ordered via base_ready/add_done
+        # events) is NOT cuda-graph-capture-safe: it corrupts `output` on replay.
+        # Fall back to the serial caller-side add -- leave the staging so the model
+        # layer reclaims it via unstage_shared_expert_add and adds shared_output
+        # after current_stream.wait_stream(alt_stream).
+        return None
+    _PENDING = None
+
+    base_ready = torch.cuda.Event()
+    base_ready.record(current_stream)
+    add_done = torch.cuda.Event()
+    with torch.cuda.stream(producer_stream):
+        producer_stream.wait_event(base_ready)
+        output.add_(shared_output)
+        add_done.record(producer_stream)
+
+    if torch.cuda.is_current_stream_capturing():
+        _SHARED_ADD_EVENTS.append(base_ready)
+        _SHARED_ADD_EVENTS.append(add_done)
+    return add_done

python/sglang/srt/lora/trtllm_lora_temp/specialized_expand.py

@@ -0,0 +1,234 @@
+"""Rank-specialized LoRA-B expand for virtual-expert LoRA.
+
+The kernel here was originally a chunk in ``lora/triton_ops/virtual_experts.py``.
+It is rank-specialized: the ``R`` dimension (LoRA rank) is a triton
+``constexpr``, so each rank value used at runtime gets its own JIT-compiled
+specialization (R=16, R=32, R=64 are all supported up to the ``R <= 64`` assert,
+with no perf interaction between them — each gets its own kernel).
+
+Called from :mod:`sglang.kernels.ops.moe.virtual_experts` when
+``use_direct_expand_add=True`` (the trtllm-lora path uses this when
+``max_lora_rank <= 64``); the generic ``invoke_fused_moe_kernel`` is used
+when that flag is False (incl. ranks above 64).
+"""
+
+from typing import Any
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
+
+
+@triton.jit
+def _moe_lora_expand_add_kernel(
+    # Pointers
+    a_ptr,  # [num_tokens * top_k, rank]
+    b_ptr,  # [num_virtual_experts, N, rank]
+    c_ptr,  # [num_tokens, N]
+    topk_weights_ptr,
+    sorted_token_ids_ptr,
+    expert_ids_ptr,
+    num_tokens_post_padded_ptr,
+    # Dimensions
+    N,
+    R: tl.constexpr,
+    num_valid_tokens,
+    # Strides
+    stride_am,
+    stride_ar,
+    stride_be,
+    stride_bn,
+    stride_br,
+    stride_cm,
+    stride_cn,
+    # Constexprs
+    router_topk: tl.constexpr,
+    MUL_ROUTED_WEIGHT: tl.constexpr,
+    FUSE_SUM_ALL_REDUCE: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_R: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    GATED_A_HALF: tl.constexpr,
+):
+    """Rank-specialized LoRA-B expand for virtual-expert LoRA.
+
+    ``GATED_A_HALF`` > 0 enables the gate/up split for a gated (SwiGLU) gate_up
+    LoRA: the intermediate (A) has ``2*R`` columns (gate-shrink ``[0:R]`` then
+    up-shrink ``[R:2R]``) and the output has ``2*GATED_A_HALF`` columns (gate
+    then up). Output tiles in the up half (column >= ``GATED_A_HALF``) read the
+    up-shrink columns ``[R:2R]`` instead of ``[0:R]``. ``GATED_A_HALF`` must be
+    a multiple of ``BLOCK_SIZE_N`` so no tile straddles the gate/up boundary.
+    ``GATED_A_HALF == 0`` is the non-gated path (read ``[0:R]`` for all tiles).
+    """
+    pid = tl.program_id(0)
+    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
+    num_pid_m = tl.cdiv(num_tokens_post_padded, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+
+    if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
+        return
+
+    offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
+    offs_token = tl.load(sorted_token_ids_ptr + offs_token_id).to(tl.int64)
+    token_mask = offs_token < num_valid_tokens
+
+    off_expert = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
+    if off_expert == -1:
+        if not FUSE_SUM_ALL_REDUCE:
+            offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)
+            c_ptrs = (
+                c_ptr + offs_token[:, None] * stride_cm + offs_n[None, :] * stride_cn
+            )
+            c_mask = token_mask[:, None] & (offs_n[None, :] < N)
+            zeros = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=c_ptr.dtype.element_ty)
+            tl.store(c_ptrs, zeros, mask=c_mask)
+        return
+
+    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)
+    offs_r = tl.arange(0, BLOCK_SIZE_R).to(tl.int64)
+    rank_mask = offs_r < R
+
+    # Gated gate_up split: the up-half output tiles read up-shrink (A columns [R:2R]); gate-half
+    # tiles read gate-shrink (A columns [0:R]). GATED_A_HALF == 0 -> always read [0:R] (non-gated).
+    a_col = offs_r
+    if GATED_A_HALF > 0:
+        a_col = offs_r + tl.where(pid_n * BLOCK_SIZE_N >= GATED_A_HALF, R, 0)
+
+    a = tl.load(
+        a_ptr + offs_token[:, None] * stride_am + a_col[None, :] * stride_ar,
+        mask=token_mask[:, None] & rank_mask[None, :],
+        other=0.0,
+    )
+    b = tl.load(
+        b_ptr
+        + off_expert * stride_be
+        + offs_n[None, :] * stride_bn
+        + offs_r[:, None] * stride_br,
+        mask=(offs_n[None, :] < N) & rank_mask[:, None],
+        other=0.0,
+    )
+
+    accumulator = tl.dot(a, b, out_dtype=tl.float32)
+    if MUL_ROUTED_WEIGHT:
+        moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0.0)
+        accumulator *= moe_weight[:, None]
+
+    if FUSE_SUM_ALL_REDUCE:
+        offs_token_out = offs_token // router_topk
+    else:
+        offs_token_out = offs_token
+    c_ptrs = c_ptr + offs_token_out[:, None] * stride_cm + offs_n[None, :] * stride_cn
+    c_mask = token_mask[:, None] & (offs_n[None, :] < N)
+    if FUSE_SUM_ALL_REDUCE:
+        tl.atomic_add(c_ptrs, accumulator.to(c_ptr.dtype.element_ty), mask=c_mask)
+    else:
+        tl.store(c_ptrs, accumulator.to(c_ptr.dtype.element_ty), mask=c_mask)
+
+
+def _invoke_moe_lora_expand_add(
+    intermediate: torch.Tensor,
+    weight: torch.Tensor,
+    output: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_padded: torch.Tensor,
+    config: "dict[str, Any]",
+    mul_routed_weight: bool,
+    fuse_sum_all_reduce: bool,
+    force_block_size_n: "int | None" = None,
+) -> None:
+    """Launch the rank-specialized LoRA-B expand kernel.
+
+    ``R`` (= ``weight.shape[2]``) up to 64 is supported. ``BLOCK_SIZE_R`` is
+    set to ``next_power_of_2(R)`` so each rank value pairs with the smallest
+    tile that covers it (R=16 → BSR=16, R=32 → BSR=32, R=64 → BSR=64).
+    Triton compiles a separate specialization per (R, BLOCK_SIZE_R) combo
+    so different ranks don't interfere with each other's perf.
+    """
+    N = weight.shape[1]
+    R = weight.shape[2]
+    assert R <= 64, f"direct LoRA expand/add expects rank <= 64, got {R}"
+
+    block_size_m = config["BLOCK_SIZE_M"]
+    # BLOCK_SIZE_N defaults to 128 when N % 128 == 0 (a good N-divisible tile that also keeps
+    # the gated gate_up split aligned). ``force_block_size_n`` lets a tuner/bench override it;
+    # down-proj has no gate/up boundary, so any divisor of N is valid there.
+    if force_block_size_n is not None:
+        block_size_n = force_block_size_n
+    else:
+        block_size_n = 128 if N % 128 == 0 else config["BLOCK_SIZE_N"]
+    group_size_m = config.get("GROUP_SIZE_M", 1)
+    block_size_r = triton.next_power_of_2(R)
+
+    # gate_up LoRA: the shrink stacks gate_A and up_A, so the intermediate has 2*R columns
+    # ([0:R] = gate-shrink x@gate_A^T, [R:2R] = up-shrink x@up_A^T). The up output half
+    # (column >= N/2) must contract the up-shrink [R:2R], not gate_A's [0:R]. Reading [0:R]
+    # for both halves (the previous hardcode) computed the up delta from gate_A and dropped
+    # up_A -- wrong whenever gate_A != up_A (the normal independently-trained gate/up case;
+    # verified >100% rel error vs a PEFT reference on the real Qwen3.5 adapter). The earlier
+    # "vs cutlass" justification for reading [0:R] was unreliable (the cutlass reference shared
+    # the same bug). Detect the gated layout from the intermediate width and split in-kernel.
+    inter_width = intermediate.shape[1]
+    assert inter_width in (R, 2 * R), (
+        f"LoRA expand intermediate width must be R ({R}, non-gated) or 2*R "
+        f"({2 * R}, gated gate_up), got {inter_width}"
+    )
+    # Lazy import to avoid the trtllm_moe <-> triton_ops package import cycle at load time.
+
+    gated = inter_width == 2 * R
+    use_gated_split = (
+        gated and lora_envs.SGLANG_ENABLE_LORA_MOE_GATEUP_GATED_SPLIT.get()
+    )
+    gated_a_half = (N // 2) if use_gated_split else 0
+    if use_gated_split:
+        assert N % 2 == 0 and (N // 2) % block_size_n == 0, (
+            f"gated gate_up split needs N/2 ({N // 2}) divisible by BLOCK_SIZE_N "
+            f"({block_size_n})"
+        )
+
+    grid = (
+        triton.cdiv(sorted_token_ids.shape[0], block_size_m)
+        * triton.cdiv(N, block_size_n),
+    )
+
+    _moe_lora_expand_add_kernel[grid](
+        intermediate,
+        weight,
+        output,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        N,
+        R,
+        topk_ids.numel(),
+        intermediate.stride(0),
+        intermediate.stride(1),
+        weight.stride(0),
+        weight.stride(1),
+        weight.stride(2),
+        output.stride(-2),
+        output.stride(-1),
+        router_topk=topk_ids.shape[1],
+        MUL_ROUTED_WEIGHT=mul_routed_weight,
+        FUSE_SUM_ALL_REDUCE=fuse_sum_all_reduce,
+        BLOCK_SIZE_M=block_size_m,
+        BLOCK_SIZE_N=block_size_n,
+        BLOCK_SIZE_R=block_size_r,
+        GROUP_SIZE_M=group_size_m,
+        GATED_A_HALF=gated_a_half,
+        num_warps=config.get("num_warps", 4),
+        num_stages=1,
+    )

python/sglang/srt/lora/utils.py

@@ -0,0 +1,584 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import Iterable, List, Optional, Set, Tuple, Union
+
+import torch
+
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.utils.hf_transformers_utils import AutoConfig
+
+
+@dataclass
+class MoELoRABatchInfo:
+    # Per-request segment indptrs used by MoE LoRA routing, shape (bs + 1,).
+    seg_indptr: torch.Tensor
+
+    # Per-request adapter index used by MoE LoRA routing, shape (bs,).
+    req_to_lora: torch.Tensor
+
+    # A mask indicating if lora adapter is enabled. Shape (num_loras,)
+    adapter_enabled: torch.Tensor
+
+    # A mapping of which lora adapter is used for each token. Shape (num_tokens,)
+    # If a token has no lora adapter, the value is -1.
+    token_lora_mapping: torch.Tensor
+
+
+@dataclass
+class LoRABatchInfo:
+    # The forward mode is using CUDA Graph.
+    use_cuda_graph: bool
+
+    # Batch size
+    bs: int
+
+    # Number of segments. For triton backend, it is equal to batch size.
+    num_segments: int
+
+    # 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,)
+    weight_indices: torch.Tensor
+
+    # ranks of each lora adapter, in shape (lora_num,)
+    lora_ranks: torch.Tensor
+
+    # scaling of each lora adapter, in shape (lora_num,)
+    scalings: torch.Tensor
+
+    # Maximum segment length of current batch
+    max_len: Optional[int]
+
+    # Lengths of each segments in shape (num_segments,)
+    seg_lens: Optional[torch.Tensor]
+
+    # The logical (re)ordering of input rows (tokens), in shape (num_tokens,)
+    permutation: Optional[torch.Tensor]
+
+    # Total number of tokens this batch info expects (host-side int).
+    # Used by lm_head LoRA to validate input shape without GPU sync.
+    expected_tokens: Optional[int] = None
+
+    # CPU-side flag: True when at least one request uses a LoRA adapter.
+    # Computed from Python lists in prepare_lora_batch to avoid GPU sync.
+    has_active_lora: bool = False
+
+    # Per-request segment indptrs, shape (bs + 1,). Required by MoE virtual
+    # experts which map tokens to requests regardless of the dense-LoRA
+    # backend's internal segmentation.  For the triton backend these are
+    # identical to seg_indptr/weight_indices; for csgmv they differ because
+    # its segments are chunked across adapters.
+    req_seg_indptr: Optional[torch.Tensor] = None
+
+    # Per-request adapter index, shape (bs,).
+    req_weight_indices: Optional[torch.Tensor] = None
+
+    # MoE LoRA batch info
+    moe_lora_info: Optional[MoELoRABatchInfo] = None
+
+
+class LoRAType(Enum):
+    LORA_A = 0
+    LORA_B = 1
+
+
+def copy_weight_into_buffer(
+    buffer_view: torch.Tensor,
+    weight: torch.Tensor,
+) -> None:
+    """
+    Copy a LoRA weight tensor into a destination buffer.
+
+    When a pinned CPU source has a dtype mismatch with a device destination,
+    cast on the destination device instead of doing the conversion on CPU.
+    """
+    if weight.dtype == buffer_view.dtype:
+        buffer_view.copy_(weight, non_blocking=True)
+        return
+
+    if weight.device.type == "cpu" and buffer_view.device.type != "cpu":
+        weight = weight.to(device=buffer_view.device, non_blocking=True)
+
+    buffer_view.copy_(weight.to(dtype=buffer_view.dtype), non_blocking=True)
+
+
+def get_hidden_dim(
+    module_name: str,
+    config: AutoConfig,
+    base_model: torch.nn.Module,
+    layer_idx: int,
+    lora_added_vocab_size: int = 0,
+) -> Tuple[int]:
+    """
+    Given a module_name (might be a stacked name), return the hidden dims of modules' input and output.
+    """
+
+    if hasattr(base_model, "get_hidden_dim"):
+        return base_model.get_hidden_dim(module_name, layer_idx)
+    else:
+        """
+        WARNING: get_hidden_dim() is not defined,
+        which is used to get the hidden dim for different lora modules
+        Use the default one, but please check if it is correct for your model.
+        Please implement the function in the model class if it is not.
+        You can reference this function in llama.py.
+        """
+        head_dim = getattr(
+            config, "head_dim", config.hidden_size // config.num_attention_heads
+        )
+        if module_name == "qkv_proj":
+            return config.hidden_size, head_dim * (
+                config.num_attention_heads + config.num_key_value_heads * 2
+            )
+        elif module_name == "o_proj":
+            o_head_dim = getattr(config, "v_head_dim", None) or head_dim
+            return (
+                o_head_dim * config.num_attention_heads,
+                config.hidden_size,
+            )
+        elif module_name == "gate_up_proj":
+            inter = config.intermediate_size
+            first_k = getattr(config, "first_k_dense_replace", None)
+            moe_freq = getattr(config, "moe_layer_freq", 1)
+            if (
+                first_k is not None
+                and layer_idx >= first_k
+                and layer_idx % moe_freq == 0
+            ):
+                moe_inter = getattr(config, "moe_intermediate_size", None)
+                n_shared = getattr(config, "n_shared_experts", None)
+                if moe_inter is not None and n_shared is not None:
+                    inter = moe_inter * n_shared
+            return config.hidden_size, inter * 2
+        elif module_name == "down_proj":
+            inter = config.intermediate_size
+            first_k = getattr(config, "first_k_dense_replace", None)
+            moe_freq = getattr(config, "moe_layer_freq", 1)
+            if (
+                first_k is not None
+                and layer_idx >= first_k
+                and layer_idx % moe_freq == 0
+            ):
+                moe_inter = getattr(config, "moe_intermediate_size", None)
+                n_shared = getattr(config, "n_shared_experts", None)
+                if moe_inter is not None and n_shared is not None:
+                    inter = moe_inter * n_shared
+            return inter, config.hidden_size
+        elif module_name == "fused_qkv_a_proj_with_mqa":
+            q_lora_rank = getattr(config, "q_lora_rank", None) or 0
+            kv_lora_rank = config.kv_lora_rank
+            qk_rope_head_dim = config.qk_rope_head_dim
+            return (
+                config.hidden_size,
+                q_lora_rank + kv_lora_rank + qk_rope_head_dim,
+            )
+        elif module_name == "q_b_proj":
+            return (
+                config.q_lora_rank,
+                config.num_attention_heads
+                * (config.qk_nope_head_dim + config.qk_rope_head_dim),
+            )
+        elif module_name == "kv_b_proj":
+            return (
+                config.kv_lora_rank,
+                config.num_attention_heads
+                * (config.qk_nope_head_dim + config.v_head_dim),
+            )
+        elif module_name in DSA_INDEXER_LORA_NAMES:
+            from sglang.srt.configs.model_config import (
+                get_dsa_index_head_dim,
+                get_dsa_index_n_heads,
+            )
+
+            if module_name == "indexer.wq_b":
+                return (
+                    config.q_lora_rank,
+                    get_dsa_index_n_heads(config) * get_dsa_index_head_dim(config),
+                )
+            elif module_name == "indexer.wk":
+                return config.hidden_size, get_dsa_index_head_dim(config)
+            else:  # indexer.weights_proj
+                return config.hidden_size, get_dsa_index_n_heads(config)
+        elif module_name == "gate_up_proj_moe":
+            moe_inter = (
+                getattr(config, "moe_intermediate_size", None)
+                or config.intermediate_size
+            )
+            return config.hidden_size, moe_inter * 2
+        elif module_name == "down_proj_moe":
+            moe_inter = (
+                getattr(config, "moe_intermediate_size", None)
+                or config.intermediate_size
+            )
+            return moe_inter, 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.
+            return config.vocab_size + lora_added_vocab_size, config.hidden_size
+        elif module_name == "lm_head":
+            # For lm_head: input is hidden_size, output is vocab_size
+            # if contain extra tokens will be added; otherwise is 0.
+            return config.hidden_size, config.vocab_size + lora_added_vocab_size
+        else:
+            raise NotImplementedError(
+                "get_hidden_dim not implemented for " + module_name
+            )
+
+
+def get_normalized_target_modules(
+    target_modules: Union[str, Iterable[str]],
+) -> set[str]:
+    """
+    Mapping a list of target module name to names of the normalized LoRA weights.
+    Handles both base module names (e.g., "gate_proj") and prefixed module names (e.g., "feed_forward.gate_proj").
+
+    Also handles PEFT shorthand strings like "all-linear" or "all" by returning
+    {"all"} as a sentinel value.  Callers that need a concrete module set
+    should use :func:`auto_detect_lora_target_modules` to resolve the shorthand
+    against the loaded base model.
+    """
+    # Handle PEFT shorthand strings — return {"all"} as sentinel.
+    # Callers can resolve to concrete names via auto_detect_lora_target_modules().
+    if isinstance(target_modules, str):
+        if target_modules not in ["all", "all-linear"]:
+            raise ValueError(
+                "Only 'all' or 'all-linear' can be used as the string for target module"
+            )
+        return {"all"}
+
+    params_mapping = {
+        "q_proj": "qkv_proj",
+        "k_proj": "qkv_proj",
+        "v_proj": "qkv_proj",
+        "gate_proj": "gate_up_proj",
+        "up_proj": "gate_up_proj",
+        "out_proj": "out_proj",
+        "embed_tokens": "embed_tokens",
+        "vocab_emb": "embed_tokens",
+        "embeddings": "embed_tokens",
+        "word_embeddings": "embed_tokens",
+        "lm_head": "lm_head",
+        "output": "lm_head",
+        "unembed_tokens": "lm_head",
+        "q_a_proj": "fused_qkv_a_proj_with_mqa",
+        "kv_a_proj_with_mqa": "fused_qkv_a_proj_with_mqa",
+        # DSA indexer projections are qualified with their parent module name
+        # because the bare leaf names collide with unrelated modules in other
+        # models (e.g. DeepSeek-V4 attention `wq_b`, Pixtral vision `wk`).
+        "wq_b": "indexer.wq_b",
+        "wk": "indexer.wk",
+        "weights_proj": "indexer.weights_proj",
+    }
+
+    result = set()
+    for name in target_modules:
+        base_name = name.split(".")[-1]
+        normalized_name = params_mapping.get(base_name, base_name)
+        result.add(normalized_name)
+    return result
+
+
+def get_stacked_multiply(
+    module_name: str, base_model: Optional[torch.nn.Module] = None
+) -> int:
+    """
+    Mapping a lora module name to its magnification at output dimension.
+    Models can override via a get_stacked_multiply(module_name) method.
+    """
+    if base_model is not None and hasattr(base_model, "get_stacked_multiply"):
+        return base_model.get_stacked_multiply(module_name)
+    stacked_rank = {
+        "qkv_proj": 3,
+        "in_proj_qkvz": 4,  # GDN packed input projection
+        "gate_up_proj": 2,
+        "gate_up_proj_moe": 2,
+        "in_proj": 2,
+        "fused_qkv_a_proj_with_mqa": 2,
+    }
+    return stacked_rank[module_name] if module_name in stacked_rank else 1
+
+
+def get_target_module_name(full_module_name: str, target_modules: Set[str]) -> str:
+    """
+    Get the target module name in target_modules that can match full_module_name.
+
+    If there is a target module name in target_modules that can match full_module_name, return this name
+    Else raise ValueError.
+
+    When multiple target modules match (e.g. both "up_proj" and "gate_up_proj"
+    are substrings), the longest match wins to avoid ambiguity.
+    """
+    best = None
+    for target_module in target_modules:
+        if target_module in full_module_name:
+            if best is None or len(target_module) > len(best):
+                best = target_module
+    if best is not None:
+        return best
+    raise ValueError(
+        f"Cannot find target module name for {full_module_name} in {target_modules}"
+    )
+
+
+EMBEDDING_NAMES = ["embed_tokens", "lm_head"]
+ROW_PARALLELISM_LINEAR_LORA_NAMES = ["o_proj", "out_proj", "down_proj", "down_proj_moe"]
+DSA_INDEXER_LORA_NAMES = frozenset(
+    {"indexer.wq_b", "indexer.wk", "indexer.weights_proj"}
+)
+REPLICATED_LINEAR_LORA_NAMES = [
+    "fused_qkv_a_proj_with_mqa",
+    "fc1_latent_proj",
+    "fc2_latent_proj",
+    *DSA_INDEXER_LORA_NAMES,
+]
+
+# Normalized module names that the LoRA system fully supports
+# (i.e. get_hidden_dim, init_buffers, and init_lora_modules can handle them).
+_KNOWN_LORA_TARGET_MODULES = frozenset(
+    {
+        "qkv_proj",
+        "o_proj",
+        "out_proj",
+        "in_proj",
+        "in_proj_qkvz",
+        "up_proj",
+        "gate_up_proj",
+        "down_proj",
+        "fc1_latent_proj",
+        "fc2_latent_proj",
+        "embed_tokens",
+        "lm_head",
+        "fused_qkv_a_proj_with_mqa",
+        "q_b_proj",
+        "kv_b_proj",
+    }
+    | DSA_INDEXER_LORA_NAMES
+)
+
+
+def auto_detect_lora_target_modules(model: "torch.nn.Module") -> set:
+    """Discover LoRA-compatible modules by inspecting the base model.
+
+    Walks the model graph and returns the set of *normalized* target-module
+    names that (a) actually exist in the model and (b) the LoRA memory pool
+    can handle.  This is used to resolve PEFT shorthands like ``"all-linear"``
+    without requiring the user to enumerate modules on the CLI.
+    """
+    from sglang.srt.layers.linear import LinearBase
+    from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
+    from sglang.srt.layers.vocab_parallel_embedding import (
+        ParallelLMHead,
+        VocabParallelEmbedding,
+    )
+
+    raw_names: set = set()
+    dsa_indexer_leaf_names = {
+        target_name.split(".")[-1] for target_name in DSA_INDEXER_LORA_NAMES
+    }
+    for name, module in model.named_modules():
+        if isinstance(module, FusedMoE):
+            raw_names.add("gate_up_proj")
+            raw_names.add("down_proj")
+        elif isinstance(module, ParallelLMHead):
+            raw_names.add("lm_head")
+        elif isinstance(module, VocabParallelEmbedding):
+            raw_names.add("embed_tokens")
+        elif isinstance(module, LinearBase):
+            parts = name.split(".")
+            leaf_name = parts[-1]
+            parent_qualified_name = ".".join(parts[-2:])
+            if parent_qualified_name in DSA_INDEXER_LORA_NAMES:
+                raw_names.add(parent_qualified_name)
+            elif leaf_name in dsa_indexer_leaf_names:
+                # Bare DSA indexer leaf names are ambiguous across model
+                # families. Only auto-detect them when the actual module path
+                # proves they are under an `indexer` parent.
+                continue
+            else:
+                raw_names.add(leaf_name)
+
+    normalized = get_normalized_target_modules(raw_names)
+    result = normalized & _KNOWN_LORA_TARGET_MODULES
+
+    # Allow models to declare additional LoRA-compatible modules that
+    # cannot be auto-discovered or need to bypass normalization
+    # (e.g. Mamba in_proj, non-gated up_proj).
+    if hasattr(model, "supported_lora_modules"):
+        result.update(set(model.supported_lora_modules) & _KNOWN_LORA_TARGET_MODULES)
+
+    return result
+
+
+def get_lm_head_lora_b_shard_size(output_dim: int, shard_indices=None) -> int:
+    """Get the LoRA B output dimension for lm_head, accounting for TP.
+
+    lm_head is column-parallel, so its LoRA B must be sharded along the
+    vocab dimension to match the base output.  When shard_indices is
+    provided, the returned size reflects the base model's actual per-rank
+    vocab partition.
+
+    Args:
+        output_dim: Full (unsharded) output dimension (vocab_size).
+        shard_indices: VocabParallelEmbeddingShardIndices from the base
+            ParallelLMHead layer.  When provided, returns the per-rank
+            org vocab size from the base model's actual sharding.
+    """
+    if shard_indices is not None:
+        return shard_indices.num_org_elements
+    return output_dim
+
+
+def generate_sequence_lengths(
+    forward_batch: ForwardBatch, device: Optional[torch.device] = None
+) -> torch.Tensor:
+
+    device = torch.get_default_device() if device is None else device
+    with torch.device(device):
+        if forward_batch.forward_mode.is_decode():
+            seg_lens = torch.ones(forward_batch.batch_size, dtype=torch.int32)
+        elif forward_batch.forward_mode.is_target_verify():
+            seg_lens = torch.full(
+                size=(forward_batch.batch_size,),
+                fill_value=forward_batch.spec_info.draft_token_num,
+                dtype=torch.int32,
+            )
+        elif forward_batch.forward_mode.is_extend():
+            seg_lens = (
+                forward_batch.extend_seq_lens
+                if forward_batch.extend_seq_lens.device == device
+                else torch.tensor(
+                    forward_batch.extend_seq_lens_cpu,
+                    dtype=torch.int32,
+                )
+            )
+        else:
+            raise ValueError(f"Unsupported forward mode: {forward_batch.forward_mode}")
+    return seg_lens
+
+
+def get_lm_head_pruned_lens(
+    forward_batch: ForwardBatch,
+) -> Optional[List[int]]:
+    """
+    Compute per-sequence pruned lengths for lm_head LoRA.
+
+    Returns a list of pruned lengths (one per sequence) if pruning applies,
+    or None if lm_head pruning is not applicable for this batch.
+
+    Pruning rules:
+    - Extend without logprobs: 1 token per sequence
+    - Extend with logprobs: max(extend_len - logprob_start_len, 1) per sequence
+    - Decode / target_verify / draft_extend_v2: no pruning
+
+    IMPORTANT: This must stay in sync with LogitsProcessor._get_pruned_states()
+    in sglang/srt/layers/logits_processor.py, which determines how many tokens
+    per sequence are passed to lm_head. If the pruning conditions or lengths
+    there change, this function must be updated to match, otherwise the
+    lm_head LoRA will operate on incorrectly shaped inputs.
+    """
+    lm_head_pruning = (
+        forward_batch.forward_mode.is_extend()
+        and not forward_batch.forward_mode.is_target_verify()
+        and not forward_batch.forward_mode.is_draft_extend_v2()
+    )
+
+    if not lm_head_pruning:
+        return None
+
+    if forward_batch.return_logprob:
+        pruned_lens = []
+        for ext_len, start_len in zip(
+            forward_batch.extend_seq_lens_cpu,
+            forward_batch.extend_logprob_start_lens_cpu,
+        ):
+            pruned_lens.append(1 if ext_len == start_len else ext_len - start_len)
+    else:
+        pruned_lens = [1] * forward_batch.batch_size
+
+    return pruned_lens
+
+
+def merge_and_chunk_segments(
+    weight_indices: list[int],
+    pruned_lens: List[int],
+    chunk_size: int,
+) -> Tuple[List[int], List[int]]:
+    """
+    Merge consecutive same-adapter sequences and chunk at chunk_size boundaries.
+
+    Merges consecutive sequences that use the same adapter into single
+    segments, splitting any segment that exceeds chunk_size.
+
+    Args:
+        weight_indices: Per-sequence adapter indices.
+        pruned_lens: Per-sequence pruned token counts.
+        chunk_size: Maximum segment length before splitting.
+
+    Returns:
+        (seg_weight_indices, seg_lens): Merged and chunked segments.
+    """
+    seg_weight_indices: List[int] = []
+    seg_lens: List[int] = []
+    for wi, pl in zip(weight_indices, pruned_lens):
+        if seg_weight_indices and seg_weight_indices[-1] == wi:
+            seg_lens[-1] += pl
+        else:
+            seg_weight_indices.append(wi)
+            seg_lens.append(pl)
+        # Split the last segment if it exceeds chunk_size
+        while seg_lens[-1] > chunk_size:
+            remainder = seg_lens[-1] - chunk_size
+            seg_lens[-1] = chunk_size
+            seg_weight_indices.append(wi)
+            seg_lens.append(remainder)
+
+    return seg_weight_indices, seg_lens
+
+
+def build_lm_head_pass_segments(
+    weight_indices: List[int],
+    pruned_lens: List[int],
+    logprobs_chunk_size: int,
+) -> List[Tuple[List[int], List[int]]]:
+    """
+    Precompute per-pass segment info for lm_head LoRA logprobs processing.
+
+    When LogitsProcessor uses chunked logprobs processing
+    (process_input_logprobs_by_chunk), pruned hidden states are split into
+    fixed-size passes.  Each pass needs its own segmentation
+    (weight_indices, seg_lens) so that lm_head LoRA operates on the
+    correct adapter assignments per pass.
+
+    Args:
+        weight_indices: Per-sequence adapter indices.
+        pruned_lens: Per-sequence pruned token counts.
+        logprobs_chunk_size: Fixed pass size used by LogitsProcessor.
+
+    Returns:
+        List of (seg_weight_indices, seg_lens) tuples, one per pass.
+    """
+    # Expand to per-token weight index
+    token_wi: List[int] = []
+    for wi, pl in zip(weight_indices, pruned_lens):
+        token_wi.extend([wi] * pl)
+    total = len(token_wi)
+    num_passes = (total + logprobs_chunk_size - 1) // logprobs_chunk_size
+
+    result: List[Tuple[List[int], List[int]]] = []
+    for i in range(num_passes):
+        start = i * logprobs_chunk_size
+        end = min((i + 1) * logprobs_chunk_size, total)
+
+        # Run-length encode the pass's adapter indices
+        seg_wi: List[int] = []
+        seg_lens: List[int] = []
+        for t in range(start, end):
+            if seg_wi and seg_wi[-1] == token_wi[t]:
+                seg_lens[-1] += 1
+            else:
+                seg_wi.append(token_wi[t])
+                seg_lens.append(1)
+        result.append((seg_wi, seg_lens))
+
+    return result