CP shared-KV: global-position-rotated zigzag owner-lane balancing (default on)

In NSA in-seq-split prefill-CP shared-KV the owner builder restarted every chunk/request at owner 0, starving physical lane 0 -> load-back livelock (g0057: ~64k cp_load_back_owner_lane_capacity_failed + ~4k owner_lane_exhausted). Rotate the owner/compute assignment by phase=((extend_prefix_len//page_size)+req_index) %cp_size so chunks/requests/draft pages spread across lanes, keeping owner==compute (the local_loc_owner_mismatch invariant) by rotating the zigzag compute split by the same phase. Cache-safe: reload replays recorded page_owners, never recomputes phase. - owner builder + helpers (cp_owner_lane_phase two-term prefix+req_index, cp_in_seq_reverse_map, cp_rotated_per_rank_actual_token) in cp_shared_kv_compute_owner.py - zigzag compute split, gather-back reverse map, bs>1 torch rerange, last-token owner, request_phases on NSAContextParallelMetadata, scalar-path + last-token fail-loud guards in nsa/utils.py - mqa-logits prefill buffer estimator made rotation-aware (per-request phase) so it no longer under-sizes the admission budget - flag SGLANG_CP_SHARED_KV_OWNER_ROTATION (default True; =0 is the legacy escape hatch -- phase forced to 0 = byte-for-byte legacy, needs no new kernel) - startup require_tai_kernel_version("0.0.2") gate (utils/common.py, model_runner) - unit tests (11): tiny-chunk pathology [40,0x7]->[5x8], synchronized-batch spread [0x8]->[0..7], gather-back reassembly, per-rank-token Requires tai-kernel >= 0.0.2 (phase-aware in_seq_all_gather_rerange). Reviewed by 3 adversarial agents (owner==compute consistency, gather-back equivalence, integration safety): SHIP-WITH-NITS, all findings fixed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-28 14:09:58 +00:00
parent ce0e486327
commit 014948c5d7
7 changed files with 587 additions and 66 deletions
@@ -0,0 +1,234 @@
+"""Unit tests for the global-position-rotated zigzag owner assignment.
+
+Pure-Python (no torch): drives `build_in_seq_page_compute_owners` directly.
+
+Root cause it guards (see docs_internal/cp_shared_kv_owner_lane_rotation_design.md):
+the legacy owner builder restarts every chunk/request at owner 0, so under chunked
+prefill / bs>1 / EAGLE every first-segment page piles onto physical lane 0 and
+starves it. The rotation makes owner a function of the page's global sequence
+position (phase = extend_prefix_len // page_size % cp_size), balancing all lanes.
+"""
+
+import importlib.util
+import os
+import unittest
+
+_MODULE_PATH = os.path.join(
+    os.path.dirname(__file__),
+    "../../../../python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py",
+)
+
+
+def _load_owner_module():
+    spec = importlib.util.spec_from_file_location(
+        "cp_shared_kv_compute_owner_under_test", os.path.abspath(_MODULE_PATH)
+    )
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+cow = _load_owner_module()
+PAGE = 64
+CP = 8
+
+
+def _set_rotation(enabled: bool):
+    # Patch the flag accessor directly so the test stays torch/env-free.
+    cow.cp_owner_lane_rotation_enabled = lambda: enabled
+
+
+def _seq_owner_histogram(total_tokens, chunk_tokens, *, rotate, page=PAGE, cp=CP):
+    _set_rotation(rotate)
+    owners = []
+    prefix = 0
+    while prefix < total_tokens:
+        ext = min(chunk_tokens, total_tokens - prefix)
+        # A new owned page is minted only at a page boundary (page-floored prefix).
+        if prefix % page == 0:
+            o = cow.build_in_seq_page_compute_owners(
+                extend_len=ext, extend_prefix_len=prefix, page_size=page, cp_size=cp
+            )
+            if o:
+                owners += o
+        prefix += ext
+    hist = [0] * cp
+    for o in owners:
+        hist[o] += 1
+    return hist
+
+
+class TestCpOwnerLaneRotation(unittest.TestCase):
+    def test_phase_equals_prefix_pages_mod_cp(self):
+        _set_rotation(True)
+        for p_pages in range(0, 3 * CP):
+            prefix = p_pages * PAGE
+            owners = cow.build_in_seq_page_compute_owners(
+                extend_len=PAGE, extend_prefix_len=prefix, page_size=PAGE, cp_size=CP
+            )
+            self.assertEqual(owners, [p_pages % CP], msg=f"prefix_pages={p_pages}")
+
+    def test_disabled_is_legacy_identity(self):
+        # phase forced to 0 -> every chunk starts at owner 0 (the legacy behavior).
+        _set_rotation(False)
+        owners = cow.build_in_seq_page_compute_owners(
+            extend_len=PAGE, extend_prefix_len=10 * PAGE, page_size=PAGE, cp_size=CP
+        )
+        self.assertEqual(owners, [0])
+
+    def test_tiny_chunks_balance_under_rotation(self):
+        # The production pathology: many 1-page continuation chunks.
+        off = _seq_owner_histogram(PAGE * 40, PAGE, rotate=False)
+        on = _seq_owner_histogram(PAGE * 40, PAGE, rotate=True)
+        self.assertEqual(off, [40, 0, 0, 0, 0, 0, 0, 0])  # legacy starves lane 0..
+        self.assertEqual(max(on) - min(on), 0)  # ..rotation spreads perfectly
+        self.assertEqual(on, [5] * CP)
+
+    def test_balanced_for_any_chunking(self):
+        # For any chunk size the rotated histogram is (a) never worse than legacy
+        # and (b) bounded by the unavoidable intra-chunk zigzag imbalance
+        # ceil(chunk_pages / cp_size) -- the legacy code instead grows the lane-0
+        # deficit without bound as chunks accumulate.
+        total = PAGE * 30
+        for chunk_pages in (1, 2, 3, 5, 7, 13):
+            on = _seq_owner_histogram(total, PAGE * chunk_pages, rotate=True)
+            off = _seq_owner_histogram(total, PAGE * chunk_pages, rotate=False)
+            bound = -(-chunk_pages // CP)  # ceil(chunk_pages / cp_size)
+            self.assertLessEqual(
+                max(on) - min(on), bound, msg=f"chunk_pages={chunk_pages} hist={on}"
+            )
+            self.assertLessEqual(
+                max(on) - min(on),
+                max(off) - min(off),
+                msg=f"rotation worse than legacy: chunk_pages={chunk_pages} on={on} off={off}",
+            )
+
+    def test_single_big_forward_unchanged(self):
+        # A single large forward already spreads owners; rotation must not regress it.
+        off = _seq_owner_histogram(PAGE * 40, PAGE * 40, rotate=False)
+        on = _seq_owner_histogram(PAGE * 40, PAGE * 40, rotate=True)
+        self.assertEqual(max(off) - min(off), 0)
+        self.assertEqual(max(on) - min(on), 0)
+
+    def test_batch_builder_matches_per_request(self):
+        _set_rotation(True)
+        extend_lens = [PAGE, 2 * PAGE, PAGE]
+        prefix_lens = [3 * PAGE, 5 * PAGE, 0]
+        batch = cow.build_batch_in_seq_page_compute_owners(
+            extend_lens=extend_lens,
+            extend_prefix_lens=prefix_lens,
+            page_size=PAGE,
+            cp_size=CP,
+        )
+        expected = []
+        for req_index, (e, p) in enumerate(zip(extend_lens, prefix_lens)):
+            expected += cow.build_in_seq_page_compute_owners(
+                extend_len=e,
+                extend_prefix_len=p,
+                page_size=PAGE,
+                cp_size=CP,
+                req_index=req_index,
+            )
+        self.assertEqual(batch, expected)
+        # req0 prefix3+idx0=phase3 -> [3]; req1 prefix5+idx1=phase6, 2 pages -> [6,7];
+        # req2 prefix0+idx2=phase2 -> [2].
+        self.assertEqual(batch, [3, 6, 7, 2])
+
+    def test_synchronized_short_requests_spread(self):
+        # The case the prefix term alone cannot fix: a batch of brand-new (prefix=0)
+        # 1-page requests. Legacy/prefix-only maps them all to owner 0; the req_index
+        # term spreads them round-robin across lanes.
+        _set_rotation(True)
+        on = cow.build_batch_in_seq_page_compute_owners(
+            extend_lens=[PAGE] * CP,
+            extend_prefix_lens=[0] * CP,
+            page_size=PAGE,
+            cp_size=CP,
+        )
+        self.assertEqual(on, list(range(CP)))  # owners 0,1,...,cp-1
+        _set_rotation(False)
+        off = cow.build_batch_in_seq_page_compute_owners(
+            extend_lens=[PAGE] * CP,
+            extend_prefix_lens=[0] * CP,
+            page_size=PAGE,
+            cp_size=CP,
+        )
+        self.assertEqual(off, [0] * CP)  # the pathology
+
+
+class TestCpInSeqReverseMap(unittest.TestCase):
+    """The gather-back reverse map used by _build_batch_metadata_from_plan."""
+
+    def _reassemble(self, split_list, cp, phase):
+        # Build the rank-major all-gather buffer (rank j holds group g_j=(j-phase)%cp),
+        # apply the real helper's (reverse_split_len, cp_reverse_index), and return the
+        # reconstructed flat segment-marker sequence.
+        from itertools import accumulate  # noqa: F401
+
+        N2 = cp * 2
+        rsl, cri = cow.cp_in_seq_reverse_map(split_list, cp, phase)
+        trimmed = []
+        for j in range(cp):
+            g = (j - phase) % cp
+            for seg in (g, N2 - 1 - g):
+                trimmed += [seg] * split_list[seg]
+        pieces, off = [], 0
+        for length in rsl:
+            pieces.append(trimmed[off : off + length])
+            off += length
+        out = []
+        for i in cri:
+            out += pieces[i]
+        return out
+
+    def test_reconstructs_sequence_order_all_phases(self):
+        import random
+
+        for cp in (2, 4, CP):
+            N2 = cp * 2
+            for trial in range(60):
+                rng = random.Random(trial * 13 + cp)
+                split_list = [rng.randint(0, 5) for _ in range(N2)]
+                expected = []
+                for s in range(N2):
+                    expected += [s] * split_list[s]
+                for phase in range(cp):
+                    self.assertEqual(
+                        self._reassemble(split_list, cp, phase),
+                        expected,
+                        msg=f"cp={cp} phase={phase} split={split_list}",
+                    )
+
+    def test_identity_at_phase_zero(self):
+        split_list = [1, 2, 1, 3, 2, 1, 2, 1]  # cp=4
+        rsl, cri = cow.cp_in_seq_reverse_map(split_list, 4, 0)
+        self.assertEqual(cri, [0, 2, 4, 6, 7, 5, 3, 1])  # legacy reverse index
+        self.assertEqual(
+            rsl,
+            [split_list[0], split_list[7], split_list[1], split_list[6],
+             split_list[2], split_list[5], split_list[3], split_list[4]],
+        )
+
+
+class TestCpRotatedPerRankActualToken(unittest.TestCase):
+    def test_phase_shifts_which_rank_owns_tokens(self):
+        # req0 (zigzag-balanced -> 9 per rank) + req1 whose only tokens are in segment 0
+        # (10 tokens, owned by group 0 = the rank at (0+phase)). The phase moves those 10.
+        sls = [[1, 2, 3, 4, 5, 6, 7, 8], [10, 0, 0, 0, 0, 0, 0, 0]]
+        self.assertEqual(
+            cow.cp_rotated_per_rank_actual_token(sls, [0, 1], 4), [9, 19, 9, 9]
+        )
+        self.assertEqual(
+            cow.cp_rotated_per_rank_actual_token(sls, [0, 0], 4), [19, 9, 9, 9]
+        )
+
+    def test_identity_at_phase_zero(self):
+        N2 = CP * 2
+        sls = [[i + 1 for i in range(N2)], [1] * N2]
+        got = cow.cp_rotated_per_rank_actual_token(sls, [0] * len(sls), CP)
+        expected = [sum(sl[r] + sl[N2 - 1 - r] for sl in sls) for r in range(CP)]
+        self.assertEqual(got, expected)
+
+
+if __name__ == "__main__":
+    unittest.main()