minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[LoRA] Fix chunked SGMV (csgmv) CUDA graph segment replay (#28371)

Browse Source
This commit is contained in:
Yanbin Jiang
2026-06-16 19:19:07 -07:00
committed by GitHub
parent 71b090a8e7
commit 093908d4c0
7 changed files with 392 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/lora/backend/chunked_backend.py
+10
View File
@@ -273,6 +273,8 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
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(
@@ -308,6 +310,8 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
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_(
@@ -319,7 +323,13 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
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)
python/sglang/srt/lora/triton_ops/chunked_embedding_lora_a.py
+12 -5
View File
@@ -5,7 +5,7 @@ import triton.language as tl
from sglang.srt.lora.utils import LoRABatchInfo
@triton.jit
@triton.jit(do_not_specialize=["num_segments"])
def _chunked_embedding_lora_a_kernel(
# Pointers to tensors
input_ids,
@@ -39,6 +39,10 @@ def _chunked_embedding_lora_a_kernel(
# If chunk id is larger than actual number of chunks, skip
if chunk_idx >= num_segments:
return
chunk_start = tl.load(seg_indptr + chunk_idx)
chunk_end = tl.load(seg_indptr + chunk_idx + 1)
if chunk_start == chunk_end:
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)
@@ -46,8 +50,6 @@ def _chunked_embedding_lora_a_kernel(
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
@@ -108,8 +110,13 @@ def chunked_embedding_lora_a_forward(
# Block size for rank dimension
BLOCK_RANK = 128
num_segments = batch_info.num_segments
segment_grid = (
batch_info.weight_indices.shape[0]
if batch_info.use_cuda_graph
else 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,)
grid = (segment_grid,)
output = torch.zeros((S, rank), device=input_ids.device, dtype=weights.dtype)
_chunked_embedding_lora_a_kernel[grid](
@@ -127,7 +134,7 @@ def chunked_embedding_lora_a_forward(
batch_info.seg_indptr,
batch_info.weight_indices,
batch_info.lora_ranks,
batch_info.num_segments,
segment_grid,
batch_info.permutation,
BLOCK_RANK,
)
python/sglang/srt/lora/triton_ops/chunked_sgmv_expand.py
+13 -6
View File
@@ -67,6 +67,11 @@ def _chunked_lora_expand_kernel(
if pid_s >= num_segs:
return
seg_start = tl.load(seg_indptr + pid_s)
seg_end = tl.load(seg_indptr + pid_s + 1)
if seg_start == seg_end:
return
# Current block computes sequence with batch_id,
# which starts from row seg_start of x with length seg_len.
# qkv_id decides which of q,k,v to compute (0: q, 1: k, 2: v)
@@ -77,9 +82,6 @@ def _chunked_lora_expand_kernel(
if cur_rank == 0:
return
seg_start = tl.load(seg_indptr + pid_s)
seg_end = tl.load(seg_indptr + pid_s + 1)
slice_id = tl.program_id(axis=1)
slice_start = tl.load(slice_offsets + slice_id)
slice_end = tl.load(slice_offsets + slice_id + 1)
@@ -91,7 +93,7 @@ def _chunked_lora_expand_kernel(
# Map logical sequence index to physical index
s_offset_logical = tl.arange(0, BLOCK_M) + seg_start
s_offset_physical = tl.load(
permutation + s_offset_logical, mask=s_offset_logical < seg_end
permutation + s_offset_logical, mask=s_offset_logical < seg_end, other=0
)
# Create pointers for the first block of x and weights[batch_id][n_start: n_end][:]
@@ -184,11 +186,16 @@ def chunked_sgmv_lora_expand_forward(
BLOCK_N = config["BLOCK_N"]
num_segments = batch_info.num_segments
segment_grid = (
batch_info.weight_indices.shape[0]
if batch_info.use_cuda_graph
else num_segments
)
grid = (
triton.cdiv(max_slice_size, BLOCK_N),
num_slices, # number of slices in the input/output
batch_info.bs if batch_info.use_cuda_graph else num_segments,
segment_grid,
)
if base_output is None:
@@ -215,7 +222,7 @@ def chunked_sgmv_lora_expand_forward(
weight_indices=batch_info.weight_indices,
lora_ranks=batch_info.lora_ranks,
permutation=batch_info.permutation,
num_segs=num_segments,
num_segs=segment_grid,
scalings=batch_info.scalings,
slice_offsets=slice_offsets,
# constants
python/sglang/srt/lora/triton_ops/chunked_sgmv_shrink.py
+13 -6
View File
@@ -61,6 +61,11 @@ def _chunked_lora_shrink_kernel(
pid_n = tl.program_id(0)
seg_start = tl.load(seg_indptr + pid_s)
seg_end = tl.load(seg_indptr + pid_s + 1)
if seg_start == seg_end:
return
# Current block computes sequence with batch_id,
# which starts from row seg_start of x with length seg_len
w_index = tl.load(weight_indices + pid_s)
@@ -70,16 +75,13 @@ def _chunked_lora_shrink_kernel(
if rank == 0:
return
seg_start = tl.load(seg_indptr + pid_s)
seg_end = tl.load(seg_indptr + pid_s + 1)
# Adjust N dim according to the specific LoRA adapter
cur_n = tl.minimum(N, rank * NUM_SLICES)
# Map logical sequence index to physical index
s_offset_logical = tl.arange(0, BLOCK_M) + seg_start
s_offset_physical = tl.load(
permutation + s_offset_logical, mask=s_offset_logical < seg_end
permutation + s_offset_logical, mask=s_offset_logical < seg_end, other=0
)
n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
@@ -154,9 +156,14 @@ def chunked_sgmv_lora_shrink_forward(
assert x.shape[-1] == K
num_segments = batch_info.num_segments
segment_grid = (
batch_info.weight_indices.shape[0]
if batch_info.use_cuda_graph
else num_segments
)
grid = (
triton.cdiv(N, BLOCK_N),
batch_info.bs if batch_info.use_cuda_graph else num_segments,
segment_grid,
)
# Optional launch params from tuned config
@@ -175,7 +182,7 @@ def chunked_sgmv_lora_shrink_forward(
weight_indices=batch_info.weight_indices,
lora_ranks=batch_info.lora_ranks,
permutation=batch_info.permutation,
num_segs=num_segments,
num_segs=segment_grid,
# constants
N=N,
K=K,
python/sglang/srt/lora/triton_ops/kv_b_lora_absorbed.py
+9 -5
View File
@@ -90,7 +90,11 @@ def _max_segment_len(batch_info: LoRABatchInfo) -> int:
def _segment_grid_size(batch_info: LoRABatchInfo, num_segments: int) -> int:
return batch_info.bs if batch_info.use_cuda_graph else num_segments
return (
batch_info.weight_indices.shape[0]
if batch_info.use_cuda_graph
else num_segments
)
# ---------------------------------------------------------------------------
@@ -273,7 +277,7 @@ def step_a_q_fwd(
batch_info.weight_indices,
batch_info.lora_ranks,
batch_info.permutation,
num_segments,
segment_grid,
FULL_K=full_K_per_head,
SORTED_BY_ADAPTER=sorted_by_adapter,
BLOCK_S=_BLOCK_S,
@@ -462,7 +466,7 @@ def step_b_q_fwd(
batch_info.lora_ranks,
batch_info.permutation,
batch_info.scalings,
num_segments,
segment_grid,
SORTED_BY_ADAPTER=sorted_by_adapter,
BLOCK_S=_BLOCK_S,
BLOCK_N=_STEP_B_BLOCK_N,
@@ -641,7 +645,7 @@ def step_a_v_fwd(
batch_info.weight_indices,
batch_info.lora_ranks,
batch_info.permutation,
num_segments,
segment_grid,
SORTED_BY_ADAPTER=sorted_by_adapter,
BLOCK_S=_BLOCK_S,
BLOCK_N=_STEP_A_BLOCK_N,
@@ -838,7 +842,7 @@ def step_b_v_fwd(
batch_info.lora_ranks,
batch_info.permutation,
batch_info.scalings,
num_segments,
segment_grid,
FULL_K=full_K_per_head,
QK_NOPE_OFFSET=qk_nope_head_dim,
SORTED_BY_ADAPTER=sorted_by_adapter,
python/sglang/srt/lora/trtllm_lora_temp/triton_ops/kv_b_lora_absorbed.py
+9 -5
View File
@@ -90,7 +90,11 @@ def _max_segment_len(batch_info: LoRABatchInfo) -> int:
def _segment_grid_size(batch_info: LoRABatchInfo, num_segments: int) -> int:
return batch_info.bs if batch_info.use_cuda_graph else num_segments
return (
batch_info.weight_indices.shape[0]
if batch_info.use_cuda_graph
else num_segments
)
# ---------------------------------------------------------------------------
@@ -297,7 +301,7 @@ def step_a_q_fwd(
batch_info.weight_indices,
batch_info.lora_ranks,
batch_info.permutation,
num_segments,
segment_grid,
FULL_K=full_K_per_head,
SORTED_BY_ADAPTER=sorted_by_adapter,
K_DIV=(qk_nope_dim % _STEP_A_Q_BLOCK_K == 0),
@@ -521,7 +525,7 @@ def step_b_q_fwd(
batch_info.lora_ranks,
batch_info.permutation,
batch_info.scalings,
num_segments,
segment_grid,
SORTED_BY_ADAPTER=sorted_by_adapter,
N_DIV=(kv_lora_rank % _STEP_B_Q_BLOCK_N == 0),
BLOCK_S=_BLOCK_S,
@@ -726,7 +730,7 @@ def step_a_v_fwd(
batch_info.weight_indices,
batch_info.lora_ranks,
batch_info.permutation,
num_segments,
segment_grid,
SORTED_BY_ADAPTER=sorted_by_adapter,
K_DIV=(kv_lora_rank % _STEP_A_V_BLOCK_K == 0),
BLOCK_S=_BLOCK_S,
@@ -939,7 +943,7 @@ def step_b_v_fwd(
batch_info.lora_ranks,
batch_info.permutation,
batch_info.scalings,
num_segments,
segment_grid,
FULL_K=full_K_per_head,
QK_NOPE_OFFSET=qk_nope_head_dim,
SORTED_BY_ADAPTER=sorted_by_adapter,
test/registered/lora/test_chunked_sgmv_backend.py
+326
View File
@@ -11,6 +11,10 @@ from sglang.srt.lora.triton_ops import (
chunked_embedding_lora_a_forward,
chunked_sgmv_lora_expand_forward,
chunked_sgmv_lora_shrink_forward,
step_a_q_fwd,
step_a_v_fwd,
step_b_q_fwd,
step_b_v_fwd,
)
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
@@ -825,6 +829,328 @@ class TestChunkedSGMV(unittest.TestCase):
f"decode skewed batch_size={batch_size}",
)
def _make_cuda_graph_batch_info(self, bs: int, num_loras: int) -> LoRABatchInfo:
return LoRABatchInfo(
use_cuda_graph=True,
bs=bs,
num_segments=None,
max_len=CHUNK_SIZE,
seg_lens=None,
seg_indptr=torch.zeros(bs + 1, dtype=torch.int32, device=self.device),
weight_indices=torch.zeros(bs, dtype=torch.int32, device=self.device),
lora_ranks=torch.zeros(num_loras, dtype=torch.int32, device=self.device),
scalings=torch.ones(num_loras, dtype=torch.float, device=self.device),
permutation=torch.arange(bs, dtype=torch.int32, device=self.device),
)
def _set_cuda_graph_segment_state(
self,
batch_info: LoRABatchInfo,
lora_ranks: List[int],
weight_indices: List[int],
seg_indptr: List[int],
):
num_segments = len(weight_indices)
total_tokens = seg_indptr[-1]
device = batch_info.weight_indices.device
batch_info.lora_ranks.zero_()
batch_info.lora_ranks[: len(lora_ranks)].copy_(
torch.tensor(lora_ranks, dtype=torch.int32, device=device)
)
batch_info.weight_indices.zero_()
batch_info.weight_indices[:num_segments].copy_(
torch.tensor(weight_indices, dtype=torch.int32, device=device)
)
batch_info.seg_indptr.fill_(total_tokens)
batch_info.seg_indptr[: num_segments + 1].copy_(
torch.tensor(seg_indptr, dtype=torch.int32, device=device)
)
batch_info.num_segments = num_segments
def _set_cuda_graph_capture_state(self, batch_info: LoRABatchInfo, bs: int):
self._set_cuda_graph_segment_state(
batch_info=batch_info,
lora_ranks=[0] * batch_info.lora_ranks.shape[0],
weight_indices=[0],
seg_indptr=[0, bs],
)
def _set_cuda_graph_replay_state(
self, batch_info: LoRABatchInfo, max_rank: int, bs: int
):
self._set_cuda_graph_segment_state(
batch_info=batch_info,
lora_ranks=[max_rank] * batch_info.lora_ranks.shape[0],
weight_indices=[1, 2, 3, 4],
seg_indptr=[0, 2, 4, 6, bs],
)
@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required")
def test_cuda_graph_shrink_replay_with_more_segments_than_capture(self):
"""CUDA graph replay must honor updated shrink segment metadata."""
reset_kernel_cache()
bs = 8
num_loras = 5
max_rank = 8
input_dim = 64
num_slices = 1
x = torch.randn(bs, input_dim, dtype=self.dtype, device=self.device)
weights = torch.randn(
num_loras, max_rank, input_dim, dtype=self.dtype, device=self.device
)
batch_info = self._make_cuda_graph_batch_info(bs, num_loras)
self._set_cuda_graph_replay_state(batch_info, max_rank, bs)
expected = chunked_sgmv_lora_shrink_forward(
x, weights, batch_info, num_slices=num_slices
).clone()
torch.cuda.synchronize()
self._set_cuda_graph_capture_state(batch_info, bs)
warmup_stream = torch.cuda.Stream()
warmup_stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(warmup_stream):
for _ in range(3):
chunked_sgmv_lora_shrink_forward(
x, weights, batch_info, num_slices=num_slices
)
torch.cuda.current_stream().wait_stream(warmup_stream)
torch.cuda.synchronize()
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
captured_output = chunked_sgmv_lora_shrink_forward(
x, weights, batch_info, num_slices=num_slices
)
self._set_cuda_graph_replay_state(batch_info, max_rank, bs)
captured_output.zero_()
graph.replay()
torch.cuda.synchronize()
torch.testing.assert_close(
captured_output[:, :max_rank],
expected[:, :max_rank],
rtol=self.RTOL,
atol=self.ATOL,
)
@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required")
def test_cuda_graph_expand_replay_with_more_segments_than_capture(self):
"""CUDA graph replay must honor updated expand segment metadata."""
reset_kernel_cache()
bs = 8
num_loras = 5
max_rank = 8
output_dim = 32
slice_offsets = torch.tensor(
[0, output_dim], dtype=torch.int32, device=self.device
)
x = torch.randn(bs, max_rank, dtype=self.dtype, device=self.device)
weights = torch.randn(
num_loras, output_dim, max_rank, dtype=self.dtype, device=self.device
)
base_output = torch.randn(bs, output_dim, dtype=self.dtype, device=self.device)
graph_base_output = base_output.clone()
batch_info = self._make_cuda_graph_batch_info(bs, num_loras)
self._set_cuda_graph_replay_state(batch_info, max_rank, bs)
expected = chunked_sgmv_lora_expand_forward(
x,
weights,
batch_info,
slice_offsets,
output_dim,
base_output=base_output.clone(),
).clone()
torch.cuda.synchronize()
self._set_cuda_graph_capture_state(batch_info, bs)
warmup_stream = torch.cuda.Stream()
warmup_stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(warmup_stream):
for _ in range(3):
graph_base_output.copy_(base_output)
chunked_sgmv_lora_expand_forward(
x,
weights,
batch_info,
slice_offsets,
output_dim,
base_output=graph_base_output,
)
torch.cuda.current_stream().wait_stream(warmup_stream)
torch.cuda.synchronize()
graph_base_output.copy_(base_output)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
captured_output = chunked_sgmv_lora_expand_forward(
x,
weights,
batch_info,
slice_offsets,
output_dim,
base_output=graph_base_output,
)
self._set_cuda_graph_replay_state(batch_info, max_rank, bs)
graph_base_output.copy_(base_output)
graph.replay()
torch.cuda.synchronize()
torch.testing.assert_close(
captured_output,
expected,
rtol=self.RTOL,
atol=self.ATOL,
)
@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required")
def test_prepare_lora_batch_cuda_graph_zero_length_tail(self):
"""prepare_lora_batch must neutralize stale CUDA graph tail segments."""
class MockForwardBatch:
def __init__(self, batch_size):
self.batch_size = batch_size
self.forward_mode = ForwardMode.DECODE
mock_server_args = type(
"ServerArgs", (object,), {"max_lora_chunk_size": CHUNK_SIZE}
)
backend = ChunkedSgmvLoRABackend(
max_loras_per_batch=5, device=self.device, server_args=mock_server_args
)
backend.init_cuda_graph_batch_info(max_bs_in_cuda_graph=8, num_tokens_per_bs=1)
lora_ranks = [8] * 5
scalings = [1.0] * 5
backend.prepare_lora_batch(
forward_batch=MockForwardBatch(8),
weight_indices=[0, 1, 2, 3, 4, 0, 1, 2],
lora_ranks=lora_ranks,
scalings=scalings,
use_cuda_graph=True,
)
backend.prepare_lora_batch(
forward_batch=MockForwardBatch(2),
weight_indices=[0, 0],
lora_ranks=lora_ranks,
scalings=scalings,
use_cuda_graph=True,
)
torch.cuda.synchronize()
batch_info = backend.batch_info
self.assertEqual(batch_info.num_segments, 1)
torch.testing.assert_close(
batch_info.weight_indices.cpu(),
torch.tensor([0] * 8, dtype=torch.int32),
)
torch.testing.assert_close(
batch_info.seg_indptr.cpu(),
torch.tensor([0, 2, 2, 2, 2, 2, 2, 2, 2], dtype=torch.int32),
)
@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required")
def test_kv_b_cuda_graph_replay_with_more_segments_than_capture(self):
"""Absorbed MLA kv_b LoRA kernels must replay dynamic segment metadata."""
bs = 8
num_loras = 5
max_rank = 8
num_heads = 2
qk_nope_head_dim = 16
v_head_dim = 16
kv_lora_rank = 32
full_k_per_head = qk_nope_head_dim + v_head_dim
q_nope = torch.randn(
bs, num_heads, qk_nope_head_dim, dtype=self.dtype, device=self.device
)
attn_output = torch.randn(
bs, num_heads, kv_lora_rank, dtype=self.dtype, device=self.device
)
a_buf = torch.randn(
num_loras, max_rank, kv_lora_rank, dtype=self.dtype, device=self.device
)
b_buf = torch.randn(
num_loras,
num_heads * full_k_per_head,
max_rank,
dtype=self.dtype,
device=self.device,
)
base_q = torch.randn(
bs, num_heads, kv_lora_rank, dtype=self.dtype, device=self.device
)
base_v = torch.randn(
bs, num_heads, v_head_dim, dtype=self.dtype, device=self.device
)
graph_base_q = base_q.clone()
graph_base_v = base_v.clone()
batch_info = self._make_cuda_graph_batch_info(bs, num_loras)
def run_kv_b(base_q_out, base_v_out):
q_lora_a = step_a_q_fwd(q_nope, b_buf, batch_info, full_k_per_head)
q_out = step_b_q_fwd(q_lora_a, a_buf, batch_info, base_q_out)
v_lora_a = step_a_v_fwd(attn_output, a_buf, batch_info)
v_out = step_b_v_fwd(
v_lora_a,
b_buf,
batch_info,
base_v_out,
qk_nope_head_dim,
v_head_dim,
)
return q_out, v_out
self._set_cuda_graph_replay_state(batch_info, max_rank, bs)
expected_q, expected_v = run_kv_b(base_q.clone(), base_v.clone())
expected_q = expected_q.clone()
expected_v = expected_v.clone()
torch.cuda.synchronize()
self._set_cuda_graph_capture_state(batch_info, bs)
warmup_stream = torch.cuda.Stream()
warmup_stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(warmup_stream):
for _ in range(3):
graph_base_q.copy_(base_q)
graph_base_v.copy_(base_v)
run_kv_b(graph_base_q, graph_base_v)
torch.cuda.current_stream().wait_stream(warmup_stream)
torch.cuda.synchronize()
graph_base_q.copy_(base_q)
graph_base_v.copy_(base_v)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
captured_q, captured_v = run_kv_b(graph_base_q, graph_base_v)
self._set_cuda_graph_replay_state(batch_info, max_rank, bs)
graph_base_q.copy_(base_q)
graph_base_v.copy_(base_v)
graph.replay()
torch.cuda.synchronize()
torch.testing.assert_close(
captured_q,
expected_q,
rtol=self.RTOL,
atol=self.ATOL,
)
torch.testing.assert_close(
captured_v,
expected_v,
rtol=self.RTOL,
atol=self.ATOL,
)
class TestLmHeadPruningConsistency(unittest.TestCase):
"""Verify get_lm_head_pruned_lens (LoRA) stays consistent with