chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Enable embedding lookup/lora_a logic for chunked backend (#17692)

Browse Source 
Co-authored-by: Bruce Wu <mogicianwu@fb.com>
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
Co-authored-by: Ethan (Yusheng) Su <yushengsu.thu@gmail.com>
This commit is contained in:
Bruce Wu
2026-03-16 11:37:58 -07:00
committed by GitHub
parent 061ec582bf
commit 70a6fb53af
15 changed files with 930 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
python/sglang/srt/layers/logits_processor.py
View File

@@ -688,6 +688,12 @@ class LogitsProcessor(nn.Module):
start_idx = i * chunk_size
end_idx = min((i + 1) * chunk_size, total_size)
# Notify lm_head LoRA about the current chunk so it can swap
# to the precomputed per-chunk batch_info. This is a no-op
# for non-LoRA lm_head modules.
if hasattr(lm_head, "set_lm_head_pass"):
lm_head.set_lm_head_pass(i)
# Get indices for this chunk
chunk_mask = (input_logprob_indices >= start_idx) & (
input_logprob_indices < end_idx
@@ -792,6 +798,13 @@ class LogitsProcessor(nn.Module):
]
input_token_logprobs.append(chunk_input_token_logprobs)
# Restore the full-pruned lm_head batch_info after chunk iteration.
if hasattr(lm_head, "reset_lm_head_pass"):
assert hasattr(
lm_head, "set_lm_head_pass"
), "lm_head must have set_lm_head_pass method and reset_lm_head_pass method at the same time"
lm_head.reset_lm_head_pass()
# Concatenate the results
input_token_logprobs = torch.cat(input_token_logprobs, dim=0)

4
python/sglang/srt/lora/backend/base_backend.py
View File

@@ -2,10 +2,11 @@ from typing import Tuple, Union
import torch
from sglang.srt.lora.backend.lmhead_mixing import LoRABackendLmHeadMixing
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
class BaseLoRABackend:
class BaseLoRABackend(LoRABackendLmHeadMixing):
"""Base class for different Lora backends.
Each backend has its own implementation of Lora kernels.
@@ -18,6 +19,7 @@ class BaseLoRABackend:
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()
def run_lora_a_embedding(
self,

140
python/sglang/srt/lora/backend/chunked_backend.py
View File

@@ -1,11 +1,20 @@
import dataclasses
from typing import List, Optional, Tuple
import torch
from sglang.srt.lora.backend.base_backend import BaseLoRABackend
from sglang.srt.lora.triton_ops import (
chunked_embedding_lora_a_forward,
chunked_sgmv_lora_expand_forward,
chunked_sgmv_lora_shrink_forward,
)
from sglang.srt.lora.utils import LoRABatchInfo, generate_sequence_lengths
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
@@ -33,13 +42,40 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
super().__init__(max_loras_per_batch, device)
self.max_chunk_size = server_args.max_lora_chunk_size
def run_lora_a_sgemm(
self, x: torch.Tensor, weights: torch.Tensor, *args, **kwargs
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,
*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=self.batch_info,
batch_info=batch_info,
num_slices=1,
)
@@ -49,16 +85,20 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
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=self.batch_info,
batch_info=batch_info,
slice_offsets=output_offset,
max_slice_size=max_slice_size,
base_output=base_output,
@@ -141,15 +181,18 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
Returns:
The determined chunk size
"""
if self.max_chunk_size <= MIN_CHUNK_SIZE:
return MIN_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:
@@ -253,6 +296,85 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
batch_info.permutation[: len(permutation)].copy_(permutation, non_blocking=True)
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):

64
python/sglang/srt/lora/backend/lmhead_mixing.py Normal file
View File

@@ -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

96
python/sglang/srt/lora/backend/triton_backend.py
View File

@@ -1,3 +1,6 @@
import dataclasses
from typing import List, Optional, Tuple
import torch
from sglang.srt.lora.backend.base_backend import BaseLoRABackend
@@ -8,7 +11,11 @@ from sglang.srt.lora.triton_ops import (
sgemm_lora_a_fwd,
sgemm_lora_b_fwd,
)
from sglang.srt.lora.utils import LoRABatchInfo
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
@@ -42,19 +49,31 @@ class TritonLoRABackend(BaseLoRABackend):
)
def run_lora_a_sgemm(
self, x: torch.Tensor, weights: torch.Tensor, *args, **kwargs
self,
x: torch.Tensor,
weights: torch.Tensor,
pruned_batch_info: LoRABatchInfo = None,
*args,
**kwargs,
) -> torch.Tensor:
return sgemm_lora_a_fwd(x, weights, self.batch_info)
batch_info = (
pruned_batch_info if pruned_batch_info is not None else self.batch_info
)
return sgemm_lora_a_fwd(x, weights, batch_info)
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.batch_info, base_output)
batch_info = (
pruned_batch_info if pruned_batch_info is not None else self.batch_info
)
return sgemm_lora_b_fwd(x, weights, batch_info, base_output)
def run_qkv_lora(
self,
@@ -214,3 +233,72 @@ class TritonLoRABackend(BaseLoRABackend):
batch_info.weight_indices[:bs].copy_(weight_indices_tensor, non_blocking=True)
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)
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,
)

64
python/sglang/srt/lora/layers.py
View File

@@ -240,8 +240,46 @@ class ParallelLMHeadWithLoRA(BaseLayerWithLoRA):
self.lm_head_A_buffer = lm_head_A_buffer # (num_loras, rank, hidden_dim)
self.lm_head_B_buffer = lm_head_B_buffer # (num_loras, vocab_size, rank)
def _get_lm_head_batch_info(self, num_tokens: int):
"""Resolve and validate the active lm_head batch_info.
When the logits processor calls lm_head in multiple passes
(chunked logprobs), _lm_head_pass_idx selects a precomputed
per-pass batch_info. Otherwise the full-pruned batch_info is used.
Returns None when no lm_head pruning applies (decode, no LoRA, etc.).
"""
pass_idx = self.lora_backend._lm_head_pass_idx
if (
pass_idx is not None
and self.lora_backend.lm_head_pass_batch_infos is not None
):
batch_info = self.lora_backend.lm_head_pass_batch_infos[pass_idx]
else:
batch_info = self.lora_backend.lm_head_batch_info
if batch_info is not None:
if batch_info.use_cuda_graph:
raise RuntimeError(
"lm_head LoRA with pruned batch info is not supported "
"under CUDA graph. lm_head pruning should only occur "
"during extend, which does not use CUDA graph."
)
if num_tokens != batch_info.expected_tokens:
raise RuntimeError(
f"lm_head LoRA input token count mismatch: got "
f"{num_tokens} tokens but lm_head_batch_info expects "
f"{batch_info.expected_tokens}. This likely means "
f"a pruning step in LogitsProcessor._get_pruned_states is "
f"not reflected in get_lm_head_pruned_lens()."
)
return batch_info
def apply_lora(
self, base_output: torch.Tensor, hidden_states: torch.Tensor
self,
base_output: torch.Tensor,
hidden_states: torch.Tensor,
) -> torch.Tensor:
"""
Apply LoRA to LM head layer.
@@ -250,9 +288,13 @@ class ParallelLMHeadWithLoRA(BaseLayerWithLoRA):
= hidden @ W^T + hidden @ A^T @ B^T
= base_output + (hidden @ A^T) @ B^T
"""
lm_head_batch_info = self._get_lm_head_batch_info(hidden_states.shape[0])
# Apply lora_A^T: hidden_states @ A^T
lora_a_output = self.lora_backend.run_lora_a_sgemm(
hidden_states, self.lm_head_A_buffer
hidden_states,
self.lm_head_A_buffer,
pruned_batch_info=lm_head_batch_info,
)
# Apply lora_B^T: lora_a_output @ B^T
@@ -261,6 +303,7 @@ class ParallelLMHeadWithLoRA(BaseLayerWithLoRA):
weights=self.lm_head_B_buffer,
output_offset=self.output_offset,
base_output=base_output,
pruned_batch_info=lm_head_batch_info,
)
return lora_output
@@ -277,6 +320,23 @@ class ParallelLMHeadWithLoRA(BaseLayerWithLoRA):
return base_output
# ------------------------------------------------------------------
# Multi-pass lm_head support (chunked logprobs)
# ------------------------------------------------------------------
def set_lm_head_pass(self, pass_idx: int):
"""Set the active lm_head pass index before a logprobs chunk.
Called by LogitsProcessor.process_input_logprobs_by_chunk() before
each chunk's _get_logits call. _get_lm_head_batch_info() will
resolve to lm_head_pass_batch_infos[pass_idx].
"""
self.lora_backend._lm_head_pass_idx = pass_idx
def reset_lm_head_pass(self):
"""Reset the lm_head pass index after all passes are done."""
self.lora_backend._lm_head_pass_idx = None
def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
# For TP=1, no slicing needed
# For TP>1, need to modify code in: sglang/python/sglang/srt/lora/mem_pool.py

16
python/sglang/srt/lora/mem_pool.py
View File

@@ -533,6 +533,22 @@ class LoRAMemoryPool:
:lora_rank,
]
load_lora_weight_tensor(buffer_view, lora_b_weights)
else:
# Zero out embedding/lm_head buffers for adapters without embedding LoRA
# to avoid using garbage values from uninitialized memory
for k in self.embedding_A_buffer.keys():
self.embedding_A_buffer[k][buffer_id].zero_()
for k in self.embedding_B_buffer.keys():
self.embedding_B_buffer[k][buffer_id].zero_()
for k in self.lm_head_A_buffer.keys():
self.lm_head_A_buffer[k][buffer_id].zero_()
for k in self.lm_head_B_buffer.keys():
self.lm_head_B_buffer[k][buffer_id].zero_()
if (
self.lora_added_tokens_size > 0
and "input_embeddings" in self.new_embeddings_buffer
):
self.new_embeddings_buffer["input_embeddings"][buffer_id].zero_()
def get_embedding_tensor(
self, target_module: str, lora_type: LoRAType

2
python/sglang/srt/lora/triton_ops/__init__.py
View File

@@ -1,3 +1,4 @@
from .chunked_embedding_lora_a import chunked_embedding_lora_a_forward
from .chunked_sgmv_expand import chunked_sgmv_lora_expand_forward
from .chunked_sgmv_shrink import chunked_sgmv_lora_shrink_forward
from .embedding_lora_a import embedding_lora_a_fwd
@@ -13,5 +14,6 @@ __all__ = [
"sgemm_lora_b_fwd",
"chunked_sgmv_lora_shrink_forward",
"chunked_sgmv_lora_expand_forward",
"chunked_embedding_lora_a_forward",
"embedding_lora_a_fwd",
]

135
python/sglang/srt/lora/triton_ops/chunked_embedding_lora_a.py Normal file
View File

@@ -0,0 +1,135 @@
import torch
import triton
import triton.language as tl
from sglang.srt.lora.utils import LoRABatchInfo
@triton.jit
def _chunked_embedding_lora_a_kernel(
# Pointers to tensors
input_ids,
weights,
output,
# Dimensions
vocab_size,
rank,
num_loras,
# Strides
w_stride_0, # stride for lora index
w_stride_1, # stride for rank
w_stride_2, # stride for vocab
output_stride_0,
output_stride_1,
# Chunk info
seg_indptr,
weight_indices,
lora_ranks,
num_segments,
permutation,
# Meta-parameters
BLOCK_RANK: tl.constexpr,
):
"""
Embedding lookup for LoRA A weights without support for extra tokens.
Each program handles one chunk of tokens across rank dimension
"""
chunk_idx = tl.program_id(axis=0)
# If chunk id is larger than actual number of chunks, skip
if chunk_idx >= num_segments:
return
# Load LoRA adapter index for this segment, then look up the rank
lora_index = tl.load(weight_indices + chunk_idx)
rank_val = tl.load(lora_ranks + lora_index)
# If rank is 0, skip
if rank_val == 0:
return
# for each token in chunk, load embedding across rank dimension
chunk_start = tl.load(seg_indptr + chunk_idx)
chunk_end = tl.load(seg_indptr + chunk_idx + 1)
for c in range(chunk_start, chunk_end):
s_index = tl.load(permutation + c)
# Load the token ID
token_id = tl.load(input_ids + s_index)
# Process in chunks of BLOCK_RANK dimensions
num_blocks = tl.cdiv(rank_val, BLOCK_RANK)
for block_id in range(num_blocks):
rank_offset = tl.arange(0, BLOCK_RANK) + block_id * BLOCK_RANK
rank_mask = rank_offset < rank_val
# Use regular LoRA A weights
# weights shape: (num_loras, rank, vocab_size)
# We need to load weights[lora_index, rank_offset, token_id]
weight_ptr = (
weights
+ lora_index * w_stride_0
+ rank_offset * w_stride_1
+ token_id * w_stride_2
)
emb_values = tl.load(weight_ptr, mask=rank_mask, other=0.0)
# Write to output
output_ptr = (
output + s_index * output_stride_0 + rank_offset * output_stride_1
)
tl.store(output_ptr, emb_values, mask=rank_mask)
def chunked_embedding_lora_a_forward(
input_ids: torch.Tensor,
weights: torch.Tensor,
batch_info: LoRABatchInfo,
vocab_size: int,
) -> torch.Tensor:
"""
Chunked Forward pass for LoRA A embedding lookup; each program handles one chunk of embedding lookup work
belonging to the same adapter
Args:
input_ids: (s,) token IDs
weights: (num_loras, rank, vocab_size) LoRA A embedding weights
batch_info: LoRABatchInfo containing batch information
vocab_size: base vocabulary size
Returns:
output: (s, rank) embedded features
"""
assert input_ids.is_contiguous()
assert weights.is_contiguous()
assert len(input_ids.shape) == 1
assert len(weights.shape) == 3
S = input_ids.shape[0]
num_loras = weights.shape[0]
rank = weights.shape[1]
# Block size for rank dimension
BLOCK_RANK = 128
num_segments = batch_info.num_segments
# 1D Grid: one program per chunk of embedding lookup work
grid = (batch_info.bs if batch_info.use_cuda_graph else num_segments,)
output = torch.zeros((S, rank), device=input_ids.device, dtype=weights.dtype)
_chunked_embedding_lora_a_kernel[grid](
input_ids,
weights,
output,
vocab_size,
rank,
num_loras,
weights.stride(0),
weights.stride(1),
weights.stride(2),
output.stride(0),
output.stride(1),
batch_info.seg_indptr,
batch_info.weight_indices,
batch_info.lora_ranks,
batch_info.num_segments,
batch_info.permutation,
BLOCK_RANK,
)
return output

4
python/sglang/srt/lora/triton_ops/chunked_sgmv_expand.py
View File

@@ -8,7 +8,9 @@ from sglang.srt.lora.utils import LoRABatchInfo
from sglang.srt.utils import cached_triton_kernel
@cached_triton_kernel(lambda _, kwargs: (kwargs["NUM_SLICES"], kwargs["BLOCK_M"]))
@cached_triton_kernel(
lambda _, kwargs: (kwargs["NUM_SLICES"], kwargs["BLOCK_M"], kwargs["OUTPUT_DIM"])
)
@triton.jit(do_not_specialize=["num_segs"])
def _chunked_lora_expand_kernel(
# Pointers to matrices

133
python/sglang/srt/lora/utils.py
View File

@@ -1,6 +1,6 @@
from dataclasses import dataclass
from enum import Enum
from typing import Iterable, Optional, Set, Tuple, Union
from typing import Iterable, List, Optional, Set, Tuple, Union
import torch
@@ -40,6 +40,10 @@ class LoRABatchInfo:
# 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
class LoRAType(Enum):
LORA_A = 0
@@ -193,3 +197,130 @@ def generate_sequence_lengths(
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

11
python/sglang/srt/server_args.py
View File

@@ -5931,17 +5931,6 @@ class ServerArgs:
), "If 'all' is specified in --lora-target-modules, it should be the only module specified."
self.lora_target_modules = set(SUPPORTED_LORA_TARGET_MODULES)
# When using the chunked SGMV backend, skip embedding / lm_head layers for now,
# since it does not support these yet (TODO: implement embedding / lm_head support)
if self.lora_backend == "csgmv":
logger.warning(
"LoRA backend 'csgmv' does not yet support embedding or lm_head layers; "
"dropping 'embed_tokens' and 'lm_head' from --lora-target-modules=all. "
"To apply LoRA to these, use --lora-backend triton."
)
self.lora_target_modules.discard("embed_tokens")
self.lora_target_modules.discard("lm_head")
# Ensure sufficient information is provided for LoRA initialization.
assert self.lora_paths or (
self.max_lora_rank and self.lora_target_modules

61
python/sglang/test/lora_utils.py
View File

@@ -201,6 +201,67 @@ def reference_sgmv_shrink(
return output
def reference_embedding_lora_a_shrink(
input_ids: torch.Tensor,
weights: torch.Tensor,
weight_indices: torch.Tensor,
seq_lengths: torch.Tensor,
lora_ranks: torch.Tensor,
vocab_size: int,
) -> torch.Tensor:
"""
Simple sequence-level reference implementation of embedding LoRA A shrink operation.
Args:
input_ids: (total_seq_len,) - Token IDs
weights: (num_loras, max_rank, vocab_size) - LoRA A embedding weights
weight_indices: LoRA idx for each sequence
seq_lengths: Length of each sequence
lora_ranks: LoRA rank for each LoRA adapters
vocab_size: Base vocabulary size
Returns:
output: (total_seq_len, max_rank) - Embedded features
"""
if weights.numel() == 0:
total_tokens = input_ids.shape[0]
return torch.zeros(total_tokens, 0, dtype=weights.dtype, device=weights.device)
total_tokens = input_ids.shape[0]
_, max_rank, _ = weights.shape
output = torch.zeros(
total_tokens, max_rank, dtype=weights.dtype, device=weights.device
)
token_offset = 0
for lora_idx, seq_len, rank in zip(
weight_indices,
seq_lengths,
lora_ranks[weight_indices],
):
if seq_len == 0:
continue
if rank > 0:
# Get token IDs for this sequence
seq_input_ids = input_ids[token_offset : token_offset + seq_len]
# Clamp token IDs to vocab size
clamped_ids = torch.clamp(seq_input_ids, max=vocab_size - 1)
# Lookup embeddings: weights[lora_idx, :rank, token_ids] -> (seq_len, rank)
# weights shape: (num_loras, max_rank, vocab_size)
lora_weights = weights[lora_idx, :rank, :] # (rank, vocab_size)
embeddings = lora_weights[:, clamped_ids].t() # (seq_len, rank)
output[token_offset : token_offset + seq_len, :rank] = embeddings
token_offset += seq_len
return output
def reference_sgmv_expand(
x: torch.Tensor,
weights: torch.Tensor,

180
test/manual/lora/test_chunked_sgmv_backend.py → test/registered/lora/test_chunked_sgmv_backend.py
View File

@@ -5,18 +5,28 @@ from typing import List, Optional, Tuple
import torch
from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessor
from sglang.srt.lora.backend.chunked_backend import ChunkedSgmvLoRABackend
from sglang.srt.lora.triton_ops import (
chunked_embedding_lora_a_forward,
chunked_sgmv_lora_expand_forward,
chunked_sgmv_lora_shrink_forward,
)
from sglang.srt.lora.triton_ops.chunked_sgmv_expand import _chunked_lora_expand_kernel
from sglang.srt.lora.triton_ops.chunked_sgmv_shrink import _chunked_lora_shrink_kernel
from sglang.srt.lora.utils import LoRABatchInfo
from sglang.test.lora_utils import reference_sgmv_expand, reference_sgmv_shrink
from sglang.srt.lora.utils import LoRABatchInfo, get_lm_head_pruned_lens
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.lora_utils import (
reference_embedding_lora_a_shrink,
reference_sgmv_expand,
reference_sgmv_shrink,
)
CHUNK_SIZE = 16
register_cuda_ci(est_time=60, suite="nightly-1-gpu", nightly=True)
def reset_kernel_cache():
_chunked_lora_shrink_kernel._clear_cache()
@@ -100,6 +110,7 @@ class TestChunkedSGMV(unittest.TestCase):
self.dtype = torch.float16
self.input_dim = 2560 # Hidden dimension
self.max_seq_len = 1024
self.vocab_size = 32000 # Vocabulary size for embedding tests
# LoRA configurations: name -> (rank, output_q, output_k, output_v)
self.lora_configs = {
@@ -285,6 +296,42 @@ class TestChunkedSGMV(unittest.TestCase):
return stacked
def create_embedding_lora_a_weights(self, lora_ranks: torch.Tensor) -> torch.Tensor:
"""Create LoRA A weights for embedding lookup.
Args:
lora_ranks: Tensor of ranks for each LoRA adapter
Returns:
Tensor of shape (num_loras, max_rank, vocab_size)
"""
lora_ranks_cpu = lora_ranks.cpu().numpy()
num_loras = len(lora_ranks_cpu)
max_rank = int(lora_ranks_cpu.max()) if num_loras > 0 else 0
if max_rank == 0:
return torch.empty(
num_loras, 0, self.vocab_size, dtype=self.dtype, device=self.device
)
weights = torch.zeros(
num_loras, max_rank, self.vocab_size, dtype=self.dtype, device=self.device
)
for i, rank in enumerate(lora_ranks_cpu):
if rank > 0:
weights[i, :rank, :] = torch.randn(
rank, self.vocab_size, dtype=self.dtype, device=self.device
)
return weights
def create_test_input_ids(self, total_tokens: int) -> torch.Tensor:
"""Create random token IDs for embedding test."""
return torch.randint(
0, self.vocab_size, (total_tokens,), dtype=torch.int64, device=self.device
)
def create_test_batch(
self,
batch_composition: BatchComposition,
@@ -461,6 +508,36 @@ class TestChunkedSGMV(unittest.TestCase):
chunked_shrink, reference_shrink, rtol=self.RTOL, atol=self.ATOL
)
# Test chunked embedding LoRA A forward
# Create embedding-specific LoRA A weights with shape (num_loras, rank, vocab_size)
embedding_lora_a = self.create_embedding_lora_a_weights(
batch_info.lora_ranks
)
# Create input_ids (token indices) instead of hidden states
total_tokens = x.shape[0]
input_ids = self.create_test_input_ids(total_tokens)
chunked_shrink_embeddings = chunked_embedding_lora_a_forward(
input_ids, embedding_lora_a, batch_info, self.vocab_size
)
reference_shrink_embeddings = reference_embedding_lora_a_shrink(
input_ids,
embedding_lora_a,
lora_assignments_tensor,
seq_lengths_tensor,
lora_ranks_tensor,
self.vocab_size,
)
torch.testing.assert_close(
chunked_shrink_embeddings,
reference_shrink_embeddings,
rtol=self.RTOL,
atol=self.ATOL,
msg=f"Shrink test embedding loRA A operation failed for batch_size={batch_size}",
)
def test_expand_basic(self):
"""Test basic expand operation against PyTorch reference"""
for batch_size in [1, 2, 16, 64]:
@@ -644,5 +721,104 @@ class TestChunkedSGMV(unittest.TestCase):
)
class TestLmHeadPruningConsistency(unittest.TestCase):
"""Verify get_lm_head_pruned_lens (LoRA) stays consistent with
LogitsProcessor._get_pruned_states (logits_processor).
If this test fails, it likely means one side was changed without
updating the other. See cross-references in both functions.
"""
def _make_mock_forward_batch(
self,
forward_mode,
extend_seq_lens_cpu,
return_logprob=False,
logprob_start_lens_cpu=None,
):
class MockForwardBatch:
pass
batch = MockForwardBatch()
batch.forward_mode = forward_mode
batch.batch_size = len(extend_seq_lens_cpu)
batch.return_logprob = return_logprob
batch.extend_seq_lens_cpu = extend_seq_lens_cpu
batch.extend_logprob_start_lens_cpu = logprob_start_lens_cpu
return batch
def _count_pruned_states_tokens(
self,
forward_mode,
extend_seq_lens_cpu,
return_logprob=False,
logprob_start_lens_cpu=None,
):
"""Call _get_pruned_states and return the number of output tokens."""
total_tokens = sum(extend_seq_lens_cpu)
hidden_states = torch.zeros(total_tokens, 4)
logits_meta = LogitsMetadata(
forward_mode=forward_mode,
extend_return_logprob=return_logprob,
extend_seq_lens=torch.tensor(extend_seq_lens_cpu, dtype=torch.int64),
extend_seq_lens_cpu=extend_seq_lens_cpu,
extend_logprob_start_lens_cpu=logprob_start_lens_cpu,
)
# _get_pruned_states does not use self, so pass None
result = LogitsProcessor._get_pruned_states(
None, hidden_states, None, None, logits_meta
)
pruned_states = result[0]
return pruned_states.shape[0]
def _assert_consistency(
self,
forward_mode,
extend_seq_lens_cpu,
return_logprob=False,
logprob_start_lens_cpu=None,
):
mock_batch = self._make_mock_forward_batch(
forward_mode,
extend_seq_lens_cpu,
return_logprob,
logprob_start_lens_cpu,
)
pruned_lens = get_lm_head_pruned_lens(mock_batch)
actual_count = self._count_pruned_states_tokens(
forward_mode,
extend_seq_lens_cpu,
return_logprob,
logprob_start_lens_cpu,
)
if pruned_lens is None:
expected_count = sum(extend_seq_lens_cpu)
else:
expected_count = sum(pruned_lens)
self.assertEqual(
expected_count,
actual_count,
f"get_lm_head_pruned_lens expects {expected_count} tokens, "
f"but _get_pruned_states produces {actual_count}. "
f"These functions must stay in sync — see their cross-reference comments.",
)
def test_extend_no_logprob(self):
self._assert_consistency(ForwardMode.EXTEND, [4, 5, 6])
def test_extend_with_logprob(self):
self._assert_consistency(
ForwardMode.EXTEND,
[4, 5, 6],
return_logprob=True,
logprob_start_lens_cpu=[0, 5, 3],
)
if __name__ == "__main__":
unittest.main()

53
test/registered/lora/test_lora_hf_sgl_logprob_diff.py
View File

@@ -28,6 +28,7 @@ Usage:
"""
import multiprocessing as mp
import os
import unittest
from typing import Any, Dict, List, Optional, Tuple
@@ -47,10 +48,13 @@ register_amd_ci(
suite="stage-b-test-small-1-gpu-amd",
)
# Test configuration constants
LORA_BACKEND = "triton"
BASE_MODEL = "meta-llama/Llama-2-7b-hf"
LORA_PATHS = ["yushengsu/sglang_lora_logprob_diff_without_tuning"]
LORA_BACKEND = "csgmv"
DISABLE_CUDA_GRAPH = False
LORA_TARGET_MODULES = None
LOGPROB_THRESHOLD = 1e-01
MAX_NEW_TOKENS = 32
# Default test prompts
DEFAULT_TEST_PROMPTS = [
@@ -442,7 +446,7 @@ class TestLoRAHFSGLLogprobDifference(CustomTestCase):
model_path: str,
lora_paths: List[str],
prompts: List[str],
max_new_tokens: int = 32,
max_new_tokens: int = MAX_NEW_TOKENS,
torch_dtype: torch.dtype = torch.float16,
lora_backend: str = LORA_BACKEND,
port: int = DEFAULT_PORT_FOR_SRT_TEST_RUNNER,
@@ -506,32 +510,51 @@ class TestLoRAHFSGLLogprobDifference(CustomTestCase):
"""
Basic test comparing HF and SGLang LoRA logprobs with small model.
"""
model_path = "meta-llama/Llama-2-7b-hf"
lora_paths = ["yushengsu/sglang_lora_logprob_diff_without_tuning"]
prompts = DEFAULT_TEST_PROMPTS[:2] # Use fewer prompts for faster testing
self._run_comparison_test(
model_path=model_path,
lora_paths=lora_paths,
model_path=BASE_MODEL,
lora_paths=LORA_PATHS,
prompts=prompts,
max_new_tokens=32,
)
def test_lora_logprob_comparison_full(self):
"""
Full test comparing HF and SGLang LoRA logprobs with all prompts.
"""
model_path = "meta-llama/Llama-2-7b-hf"
lora_paths = ["yushengsu/sglang_lora_logprob_diff_without_tuning"]
prompts = DEFAULT_TEST_PROMPTS
self._run_comparison_test(
model_path=model_path,
lora_paths=lora_paths,
prompts=prompts,
max_new_tokens=32,
model_path=BASE_MODEL,
lora_paths=LORA_PATHS,
prompts=DEFAULT_TEST_PROMPTS,
)
def test_lora_logprob_comparison_chunked(self):
"""
Test with logprobs chunking enabled and a small chunk size so that
even short prompts trigger the multi-pass lm_head LoRA path.
"""
saved = {}
env_overrides = {
"SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK": "true",
"SGLANG_LOGITS_PROCESSER_CHUNK_SIZE": "4",
}
for key, val in env_overrides.items():
saved[key] = os.environ.get(key)
os.environ[key] = val
try:
self._run_comparison_test(
model_path=BASE_MODEL,
lora_paths=LORA_PATHS,
prompts=DEFAULT_TEST_PROMPTS,
)
finally:
for key, orig in saved.items():
if orig is None:
os.environ.pop(key, None)
else:
os.environ[key] = orig
if __name__ == "__main__":
try: