recipes/docker-images/mimo26-spark/patches/sglang/29858.diff

# 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()