From 52f2fe456a97a894f95bd20de71de502dbee5c93 Mon Sep 17 00:00:00 2001
From: Kevin Flansburg <kflansburg@cloudflare.com>
Date: Wed, 3 Jun 2026 00:33:41 -0600
Subject: [PATCH] fix(disagg): correct DSA/SWA state-page transfer mismatch in
 PD disaggregation (#27004)

---
 .../sglang/srt/disaggregation/common/utils.py | 12 ++++-
 python/sglang/srt/disaggregation/prefill.py   |  8 +++-
 .../test_disaggregation_wire.py               | 45 +++++++++++++++++++
 3 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/python/sglang/srt/disaggregation/common/utils.py b/python/sglang/srt/disaggregation/common/utils.py
index 1084b7536..fadce5383 100644
--- a/python/sglang/srt/disaggregation/common/utils.py
+++ b/python/sglang/srt/disaggregation/common/utils.py
@@ -98,9 +98,19 @@ def group_concurrent_contiguous(
     src_indices: npt.NDArray[np.int32], dst_indices: npt.NDArray[np.int32]
 ) -> Tuple[List[npt.NDArray[np.int32]], List[npt.NDArray[np.int32]]]:
     """Vectorised NumPy implementation."""
-    if src_indices.size == 0:
+    # src/dst indices are transferred pairwise, so an empty side means there is
+    # nothing to transfer. Guarding both sides (not just src) avoids a cryptic
+    # NumPy broadcast error from np.diff() below when only one side is empty, e.g.
+    # a non-empty prefill DSA/SWA state list paired with an empty decode registration.
+    if src_indices.size == 0 or dst_indices.size == 0:
         return [], []
 
+    if src_indices.size != dst_indices.size:
+        raise ValueError(
+            "group_concurrent_contiguous requires equal-length src/dst index arrays, "
+            f"got {src_indices.size} and {dst_indices.size}"
+        )
+
     brk = np.where((np.diff(src_indices) != 1) | (np.diff(dst_indices) != 1))[0] + 1
     src_groups = np.split(src_indices, brk)
     dst_groups = np.split(dst_indices, brk)
diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py
index 96dbe7d48..483195c39 100644
--- a/python/sglang/srt/disaggregation/prefill.py
+++ b/python/sglang/srt/disaggregation/prefill.py
@@ -947,7 +947,13 @@ class SchedulerDisaggregationPrefillMixin:
         if last_chunk:
             self.disagg_metadata_buffers.set_buf(req)
 
-            seq_len = len(req.fill_ids)
+            # fill_ids includes the token sampled during prefill, but decode
+            # registers state pages over origin_input_ids (DecodePreallocQueue)
+            # and the main pool send is clamped to end_idx above. Matching that
+            # length here avoids emitting an extra state page when the sampled
+            # token crosses a page boundary, which mismatched src/dst lengths in
+            # group_concurrent_contiguous.
+            seq_len = min(len(req.fill_ids), len(req.origin_input_ids))
 
             def _mamba_payload():
                 return [
diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py
index 334b3121d..76ee5874b 100644
--- a/test/registered/unit/disaggregation/test_disaggregation_wire.py
+++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py
@@ -3,6 +3,7 @@ import unittest
 import numpy as np
 
 from sglang.srt.disaggregation.common.utils import (
+    group_concurrent_contiguous,
     pack_int_lists,
     pack_list_of_buffers,
     unpack_int_lists,
@@ -45,5 +46,49 @@ class TestDisaggregationWire(unittest.TestCase):
         self.assertEqual(unpack_list_of_buffers(pack_list_of_buffers(bufs)), bufs)
 
 
+class TestGroupConcurrentContiguous(unittest.TestCase):
+    @staticmethod
+    def _arr(values):
+        return np.array(values, dtype=np.int32)
+
+    def test_single_contiguous_group(self):
+        src = self._arr([10, 11, 12])
+        dst = self._arr([5, 6, 7])
+        self.assertEqual(
+            group_concurrent_contiguous(src, dst),
+            ([[10, 11, 12]], [[5, 6, 7]]),
+        )
+
+    def test_splits_on_discontiguous_indices(self):
+        src = self._arr([10, 11, 20])
+        dst = self._arr([5, 6, 7])
+        self.assertEqual(
+            group_concurrent_contiguous(src, dst),
+            ([[10, 11], [20]], [[5, 6], [7]]),
+        )
+
+    def test_both_empty(self):
+        self.assertEqual(
+            group_concurrent_contiguous(self._arr([]), self._arr([])), ([], [])
+        )
+
+    def test_empty_src_nonempty_dst(self):
+        self.assertEqual(
+            group_concurrent_contiguous(self._arr([]), self._arr([1, 2])), ([], [])
+        )
+
+    def test_nonempty_src_empty_dst(self):
+        # Regression: a non-empty source paired with an empty destination must not
+        # raise a NumPy broadcast error (observed transferring DSA sparse-attention
+        # state on a disaggregated GLM deployment when decode registered zero dst indices).
+        self.assertEqual(
+            group_concurrent_contiguous(self._arr([1, 2]), self._arr([])), ([], [])
+        )
+
+    def test_mismatched_nonempty_lengths_raise(self):
+        with self.assertRaises(ValueError):
+            group_concurrent_contiguous(self._arr([1, 2, 3]), self._arr([1, 2]))
+
+
 if __name__ == "__main__":
     unittest.main()