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,
+        )