[LoRA] Fix chunked SGMV (csgmv) CUDA graph segment replay (#28371)
This commit is contained in:
Yanbin Jiang
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user