diff --git a/python/sglang/srt/lora/backend/chunked_backend.py b/python/sglang/srt/lora/backend/chunked_backend.py index 91b160c47..7a47d10cd 100644 --- a/python/sglang/srt/lora/backend/chunked_backend.py +++ b/python/sglang/srt/lora/backend/chunked_backend.py @@ -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) diff --git a/python/sglang/srt/lora/triton_ops/chunked_embedding_lora_a.py b/python/sglang/srt/lora/triton_ops/chunked_embedding_lora_a.py index bf3cdbe30..d8a6406fe 100644 --- a/python/sglang/srt/lora/triton_ops/chunked_embedding_lora_a.py +++ b/python/sglang/srt/lora/triton_ops/chunked_embedding_lora_a.py @@ -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, ) diff --git a/python/sglang/srt/lora/triton_ops/chunked_sgmv_expand.py b/python/sglang/srt/lora/triton_ops/chunked_sgmv_expand.py index 50830be5a..e1968ff99 100644 --- a/python/sglang/srt/lora/triton_ops/chunked_sgmv_expand.py +++ b/python/sglang/srt/lora/triton_ops/chunked_sgmv_expand.py @@ -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 diff --git a/python/sglang/srt/lora/triton_ops/chunked_sgmv_shrink.py b/python/sglang/srt/lora/triton_ops/chunked_sgmv_shrink.py index e625588aa..8788a65de 100644 --- a/python/sglang/srt/lora/triton_ops/chunked_sgmv_shrink.py +++ b/python/sglang/srt/lora/triton_ops/chunked_sgmv_shrink.py @@ -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, diff --git a/python/sglang/srt/lora/triton_ops/kv_b_lora_absorbed.py b/python/sglang/srt/lora/triton_ops/kv_b_lora_absorbed.py index c1942c85e..e4ed6a35b 100644 --- a/python/sglang/srt/lora/triton_ops/kv_b_lora_absorbed.py +++ b/python/sglang/srt/lora/triton_ops/kv_b_lora_absorbed.py @@ -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, diff --git a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/kv_b_lora_absorbed.py b/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/kv_b_lora_absorbed.py index e6af8ba6f..1f6adca7c 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/kv_b_lora_absorbed.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/triton_ops/kv_b_lora_absorbed.py @@ -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, diff --git a/test/registered/lora/test_chunked_sgmv_backend.py b/test/registered/lora/test_chunked_sgmv_backend.py index 253889090..e075043a7 100644 --- a/test/registered/lora/test_chunked_sgmv_backend.py +++ b/test/registered/lora/test_chunked_sgmv_backend.py @@ -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