# Backport of sgl-project/sglang#29858 (head 0314d9de3a), ported to the # current update_sliding_window_buffer(window_kv_indptr, translator, # req_pool_indices, window, seq_lens, bs, ...) signature on 06008c17. diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index c378c80ab0f8..c5839ba47e10 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -567,6 +567,25 @@ def _update_draft_extend_buffers( kv_indptr = self._fill_kv_indptr_and_indices( bs, kv_lens, req_pool_indices, self.cuda_graph_kv_indices ) + if self.sliding_window_size is not None and self.sliding_window_size > 0: + # Fill the swa-pool window buffers like target_verify does, so the + # SWA layers' draft-extend reads the smaller swa K/V buffer with + # window indices instead of falling back to the full-pool + # kv_indices (which overruns the swa buffer -> OOB read). + window_kv_indices = self.cuda_graph_window_kv_indices + window_kv_offsets = self.cuda_graph_window_kv_offsets + _, window_kv_indices, _, window_kv_offsets[:bs] = ( + update_sliding_window_buffer( + self.window_kv_indptr, + self.kv_index_translator, + req_pool_indices, + self.sliding_window_size, + kv_lens[:bs], + bs, + token_to_kv_pool=self.token_to_kv_pool, + window_kv_indices=window_kv_indices, + ) + ) return qo_indptr, kv_indptr, num_tokens_per_req def init_forward_metadata_out_graph( @@ -1153,10 +1172,12 @@ def _build_cuda_graph_forward_metadata( qo_indptr=self.qo_indptr[: bs + 1], custom_mask=None, mask_indptr=None, - window_kv_indptr=self.window_kv_indptr, - window_kv_indices=None, - window_num_kv_splits=None, - window_kv_offsets=None, + window_kv_indptr=self.window_kv_indptr[: bs + 1] if swa else None, + window_kv_indices=self.cuda_graph_window_kv_indices if swa else None, + window_num_kv_splits=( + self.cuda_graph_window_num_kv_splits if swa else None + ), + window_kv_offsets=self.cuda_graph_window_kv_offsets if swa else None, swa_out_cache_loc=swa_out_cache_loc, out_cache_loc_full_physical=out_cache_loc_full_physical, ) diff --git a/test/registered/unit/spec/test_triton_draft_extend_cudagraph_swa_window.py b/test/registered/unit/spec/test_triton_draft_extend_cudagraph_swa_window.py new file mode 100644 index 000000000000..747f490a21be --- /dev/null +++ b/test/registered/unit/spec/test_triton_draft_extend_cudagraph_swa_window.py @@ -0,0 +1,73 @@ +"""Regression test for the EAGLE draft-extend cuda-graph SWA window buffers. + +White-box unit test: it builds a bare ``TritonAttnBackend`` via ``__new__`` and +stubs only the cuda-graph buffers that ``_build_cuda_graph_forward_metadata`` +reads, so the draft-extend metadata can be exercised on CPU without a captured +graph or a real model. It is coupled to that method's internals on purpose, +which is the price of testing the metadata shape in isolation. + +The contract under test: on a hybrid sliding-window model the draft-extend +cuda-graph metadata must expose the swa-pool window buffers (a non-None +``window_kv_indices``) exactly like target_verify does. When it leaves them +None the SWA layers fall back to the full-pool kv_indices and read out of +bounds of the smaller swa K/V buffer. Non-SWA models must keep the window +fields None. +""" + +import unittest + +import torch + +from sglang.srt.layers.attention.triton_backend import TritonAttnBackend +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +BS = 3 + + +class TestTritonDraftExtendCudaGraphSWAWindow(CustomTestCase): + def _build_backend(self, sliding_window_size): + backend = TritonAttnBackend.__new__(TritonAttnBackend) + backend.sliding_window_size = sliding_window_size + backend.num_draft_tokens = 4 + backend.speculative_num_steps = 3 + backend.kv_indptr = torch.zeros(BS + 1, dtype=torch.int32) + backend.qo_indptr = torch.zeros(BS + 1, dtype=torch.int32) + backend.window_kv_indptr = torch.zeros(BS + 1, dtype=torch.int32) + backend.cuda_graph_kv_indices = torch.zeros(16, dtype=torch.int64) + backend.cuda_graph_window_kv_indices = torch.zeros(16, dtype=torch.int64) + backend.cuda_graph_window_num_kv_splits = torch.zeros(BS, dtype=torch.int32) + backend.cuda_graph_window_kv_offsets = torch.zeros(BS, dtype=torch.int32) + return backend + + def test_draft_extend_swa_exposes_window_buffers(self): + backend = self._build_backend(sliding_window_size=128) + md = backend._build_cuda_graph_forward_metadata( + BS, ForwardMode.DRAFT_EXTEND_V2, spec_info=None, swa_out_cache_loc=None + ) + # The draft step must be handed the swa-pool window buffers instead of + # None, so the SWA layers stay on the swa buffer rather than indexing it + # with full-pool locations. + self.assertIsNotNone(md.window_kv_indices) + self.assertIs(md.window_kv_indices, backend.cuda_graph_window_kv_indices) + self.assertIsNotNone(md.window_kv_indptr) + self.assertIs(md.window_kv_offsets, backend.cuda_graph_window_kv_offsets) + self.assertIsNotNone(md.window_num_kv_splits) + + def test_draft_extend_without_swa_keeps_window_none(self): + for size in (None, 0): + backend = self._build_backend(sliding_window_size=size) + md = backend._build_cuda_graph_forward_metadata( + BS, ForwardMode.DRAFT_EXTEND_V2, spec_info=None, swa_out_cache_loc=None + ) + self.assertIsNone(md.window_kv_indices, msg=f"size={size}") + self.assertIsNone(md.window_kv_indptr, msg=f"size={size}") + self.assertIsNone(md.window_kv_offsets, msg=f"size={size}") + self.assertIsNone(md.window_num_kv_splits, msg=f"size={size}") + + +if __name__ == "__main__": + unittest.main()