Reduce CP shared KV materialize and direct-write overhead

Shared KV now relies on page-aligned CP metadata and compute-owner page allocation so persistent MLA KV and NSA index shards can be written by the rank that computed them. The compatibility read path keeps the dense full-view contract for existing topk and attention kernels, but removes duplicated prev/next index materialize, adds optional tai materialize integration, and tightens tests/docs around the fallback boundaries. Constraint: Decode remains non-CP while prefill CP owns the shared-KV changes Constraint: Existing attention/topk kernels still expect dense full-view KV/index inputs Rejected: Change attention kernels to read owner-sharded KV directly | larger semantic change reserved for later phases Rejected: Merge index K/scale storage with MLA KV storage | would couple topk and attention cache lifecycles before materialize overhead is isolated Confidence: medium Scope-risk: broad Directive: Do not remove fallback logging or debug-gated assertions without reproducing long-context chunked/radix-hit paths Tested: git diff --check --cached Not-tested: Local pytest/runtime server verification not run in this commit step per current workflow constraints
2026-05-02 07:07:28 +08:00
parent 2317952a01
commit 5769b63082
17 changed files with 2524 additions and 96 deletions
@@ -1,11 +1,18 @@
 import unittest
+from types import SimpleNamespace
+from unittest.mock import patch

 from sglang.srt.layers.attention.nsa.utils import (
+    NSAContextParallelMetadata,
    _get_in_seq_last_token_owner_and_offset,
    build_page_aligned_in_seq_split_list,
    build_token_balanced_in_seq_split_list,
+    can_cp_split,
+    cp_split_and_rebuild_1d,
+    get_cp_shared_kv_local_out_cache_loc,
    split_in_seq_cp_local_pair,
 )
+from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
 from sglang.test.ci.ci_register import register_cpu_ci

 register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
@@ -93,6 +100,85 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
        self.assertFalse(info.page_aligned)
        self.assertEqual(split_list, build_token_balanced_in_seq_split_list(512, 8))

+    def test_page_aligned_split_allows_radix_hit_suffix_with_one_page_per_rank(self):
+        split_list, info = build_page_aligned_in_seq_split_list(
+            total_len=512,
+            extend_len=512,
+            extend_prefix_len=54464,
+            page_size=64,
+            cp_size=8,
+        )
+
+        self.assertTrue(info.page_aligned)
+        self.assertEqual(sum(split_list), 512)
+        self.assertEqual(split_list[:8], [64] * 8)
+        self.assertEqual(split_list[8:], [0] * 8)
+        self.assert_page_aligned_boundaries(
+            split_list, extend_prefix_len=54464, extend_len=512, page_size=64
+        )
+
+    def test_page_aligned_split_falls_back_when_radix_hit_suffix_has_zero_rank(self):
+        split_list, info = build_page_aligned_in_seq_split_list(
+            total_len=256,
+            extend_len=256,
+            extend_prefix_len=54464,
+            page_size=64,
+            cp_size=8,
+        )
+
+        self.assertFalse(info.page_aligned)
+        self.assertEqual(split_list, build_token_balanced_in_seq_split_list(256, 8))
+
+    def test_can_cp_split_uses_replicated_compute_for_short_radix_hit_suffix(self):
+        class Mode:
+            def is_context_parallel_extend(self):
+                return True
+
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            extend_seq_lens_cpu=[256],
+            extend_prefix_lens_cpu=[54464],
+            token_to_kv_pool=SimpleNamespace(page_size=64),
+            forward_mode=Mode(),
+        )
+
+        with (
+            patch(
+                "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
+                return_value=False,
+            ),
+            patch(
+                "sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp",
+                return_value=True,
+            ),
+        ):
+            self.assertFalse(can_cp_split(256, 8, True, forward_batch))
+
+    def test_can_cp_split_keeps_cp_for_radix_hit_suffix_with_one_page_per_rank(self):
+        class Mode:
+            def is_context_parallel_extend(self):
+                return True
+
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            extend_seq_lens_cpu=[512],
+            extend_prefix_lens_cpu=[54464],
+            token_to_kv_pool=SimpleNamespace(page_size=64),
+            forward_mode=Mode(),
+        )
+
+        with (
+            patch(
+                "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
+                return_value=False,
+            ),
+            patch(
+                "sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp",
+                return_value=True,
+            ),
+        ):
+            self.assertTrue(can_cp_split(512, 8, True, forward_batch))
+
    def test_page_aligned_split_adds_padding_tokens_to_last_segment(self):
        split_list, info = build_page_aligned_in_seq_split_list(
            total_len=1040,
@@ -152,6 +238,309 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
        with self.assertRaisesRegex(RuntimeError, "local in-seq CP length mismatch"):
            split_in_seq_cp_local_pair(torch.arange(9), 5, 5, name="q_fp8")

+    def test_cp_split_and_rebuild_1d_matches_in_seq_zigzag_order(self):
+        import torch
+        from types import SimpleNamespace
+
+        forward_batch = SimpleNamespace(
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                split_list=[2, 2, 2, 2, 2, 2, 2, 2],
+                zigzag_index=[1, 6],
+            )
+        )
+
+        local_locs = cp_split_and_rebuild_1d(forward_batch, torch.arange(16))
+
+        self.assertEqual(local_locs.tolist(), [2, 3, 12, 13])
+
+    def test_local_out_cache_loc_requires_compute_owner_pages(self):
+        import torch
+        from types import SimpleNamespace
+
+        page_size = 4
+        # Segment order for cp_size=4, cp_rank=1 is segment 1 then 6.
+        # The logical page ids below deliberately encode the same owners through
+        # (logical_page - 1) % cp_size:
+        #   segment 1 -> page 2 owner 1
+        #   segment 6 -> page 6 owner 1
+        segment_pages = [1, 2, 3, 4, 8, 7, 6, 5]
+        out_cache_loc = torch.cat(
+            [
+                torch.arange(page * page_size, (page + 1) * page_size)
+                for page in segment_pages
+            ]
+        )
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            cp_shared_kv_layout=CpSharedKVLayout(
+                page_size=page_size,
+                cp_size=4,
+                cp_rank=1,
+            ),
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                split_list=[page_size] * 8,
+                zigzag_index=[1, 6],
+                page_aligned=True,
+                page_size=page_size,
+                extend_prefix_len=0,
+            ),
+            out_cache_loc=out_cache_loc,
+        )
+
+        local_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch)
+
+        self.assertIsNotNone(local_locs)
+        self.assertEqual(
+            local_locs.tolist(),
+            list(range(2 * page_size, 3 * page_size))
+            + list(range(6 * page_size, 7 * page_size)),
+        )
+
+    def test_local_out_cache_loc_falls_back_when_owner_mismatch(self):
+        import torch
+        from types import SimpleNamespace
+
+        page_size = 4
+        out_cache_loc = torch.arange(page_size * 8, page_size * 16)
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            cp_shared_kv_layout=CpSharedKVLayout(
+                page_size=page_size,
+                cp_size=4,
+                cp_rank=1,
+            ),
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                split_list=[page_size] * 8,
+                zigzag_index=[1, 6],
+                page_aligned=True,
+                page_size=page_size,
+                extend_prefix_len=0,
+            ),
+            out_cache_loc=out_cache_loc,
+        )
+
+        self.assertIsNone(get_cp_shared_kv_local_out_cache_loc(forward_batch))
+
+    def test_local_out_cache_loc_logs_every_fallback_event(self):
+        import torch
+        from types import SimpleNamespace
+
+        from sglang.srt.layers.attention.nsa import utils as nsa_utils
+
+        page_size = 4
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            cp_shared_kv_layout=CpSharedKVLayout(
+                page_size=page_size,
+                cp_size=4,
+                cp_rank=1,
+            ),
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                split_list=[page_size] * 8,
+                zigzag_index=[1, 6],
+                page_aligned=False,
+                page_size=page_size,
+                extend_prefix_len=0,
+            ),
+            out_cache_loc=torch.arange(page_size * 8, page_size * 16),
+        )
+
+        with self.assertLogs(
+            "sglang.srt.layers.attention.nsa.utils", level="INFO"
+        ) as cm:
+            self.assertIsNone(get_cp_shared_kv_local_out_cache_loc(forward_batch))
+            self.assertIsNone(get_cp_shared_kv_local_out_cache_loc(forward_batch))
+
+        self.assertEqual(len(cm.output), 2)
+        self.assertIn("CP shared KV direct-write fallback", cm.output[0])
+        self.assertIn("metadata is not page-aligned", cm.output[0])
+        self.assertIn("CP shared KV direct-write fallback", cm.output[1])
+        self.assertIn("metadata is not page-aligned", cm.output[1])
+
+    def test_indexer_direct_write_does_not_log_missing_metadata_for_non_cp_batch(self):
+        import torch
+        from types import SimpleNamespace
+
+        from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
+
+        indexer = object.__new__(Indexer)
+        indexer.nsa_enable_prefill_cp = True
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            nsa_cp_metadata=None,
+        )
+
+        with self.assertNoLogs(
+            "sglang.srt.layers.attention.nsa.utils", level="INFO"
+        ):
+            stored = Indexer._store_cp_shared_local_index_k_cache(
+                indexer,
+                forward_batch,
+                layer_id=0,
+                local_key=torch.empty(0),
+                act_quant=None,
+            )
+
+        self.assertFalse(stored)
+
+    def test_indexer_in_seq_cp_pair_materializes_index_once_for_prev_next(self):
+        import torch
+
+        from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
+
+        indexer = object.__new__(Indexer)
+        logical_pages = torch.tensor([[1, 2, 3, 4]], dtype=torch.int32)
+        materialized_index = torch.tensor([11], dtype=torch.int32)
+        dense_pages = torch.tensor([[1, 2, 3, 4]], dtype=torch.int32)
+        materialize_calls = []
+        topk_calls = []
+
+        class Metadata:
+            def get_page_table_64(self):
+                return logical_pages
+
+        def fake_materialize(forward_batch, layer_id, logical_page_table):
+            materialize_calls.append((layer_id, logical_page_table))
+            return materialized_index, dense_pages
+
+        def fake_get_topk(
+            forward_batch,
+            layer_id,
+            q_fp8,
+            weights,
+            metadata,
+            kv_len,
+            actual_seq_q,
+            cp_index=None,
+            current_index_kv=None,
+            shared_index_buffer=None,
+            shared_block_tables=None,
+        ):
+            topk_calls.append(
+                {
+                    "kv_len": kv_len,
+                    "actual_seq_q": actual_seq_q,
+                    "shared_index_buffer": shared_index_buffer,
+                    "shared_block_tables": shared_block_tables,
+                    "current_index_kv": current_index_kv,
+                }
+            )
+            return torch.full((actual_seq_q, 2), len(topk_calls), dtype=torch.int32)
+
+        indexer._maybe_materialize_shared_index_buffer = fake_materialize
+        indexer._get_topk_ragged_with_cp = fake_get_topk
+
+        forward_batch = type(
+            "ForwardBatchStub",
+            (),
+            {
+                "nsa_cp_metadata": NSAContextParallelMetadata(
+                    kv_len_prev=5,
+                    kv_len_next=9,
+                    actual_seq_q_prev=3,
+                    actual_seq_q_next=2,
+                )
+            },
+        )()
+        q_fp8 = torch.arange(5 * 4, dtype=torch.float32).view(5, 4)
+        weights = torch.arange(5 * 2, dtype=torch.float32).view(5, 2)
+
+        result = Indexer._get_topk_in_seq_cp_pair(
+            indexer,
+            forward_batch,
+            layer_id=7,
+            q_fp8=q_fp8,
+            weights=weights,
+            metadata=Metadata(),
+            current_index_kv=None,
+        )
+
+        self.assertEqual(len(materialize_calls), 1)
+        self.assertIs(materialize_calls[0][1], logical_pages)
+        self.assertEqual(len(topk_calls), 2)
+        self.assertIs(topk_calls[0]["shared_index_buffer"], materialized_index)
+        self.assertIs(topk_calls[1]["shared_index_buffer"], materialized_index)
+        self.assertIs(topk_calls[0]["shared_block_tables"], dense_pages)
+        self.assertIs(topk_calls[1]["shared_block_tables"], dense_pages)
+        self.assertIsNone(topk_calls[0]["current_index_kv"])
+        self.assertEqual(topk_calls[0]["kv_len"], 5)
+        self.assertEqual(topk_calls[1]["kv_len"], 9)
+        self.assertEqual(result.tolist(), [[1, 1], [1, 1], [1, 1], [2, 2], [2, 2]])
+
+    def test_indexer_in_seq_cp_pair_skips_materialize_when_current_index_reused(self):
+        import torch
+
+        from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
+
+        indexer = object.__new__(Indexer)
+        current_index_kv = (torch.tensor([1]), torch.tensor([2]))
+        materialize_calls = []
+        topk_calls = []
+
+        class Metadata:
+            def get_page_table_64(self):
+                raise AssertionError("current index reuse should not read page table")
+
+        def fake_materialize(forward_batch, layer_id, logical_page_table):
+            materialize_calls.append((layer_id, logical_page_table))
+            raise AssertionError("current index reuse should not materialize")
+
+        def fake_get_topk(
+            forward_batch,
+            layer_id,
+            q_fp8,
+            weights,
+            metadata,
+            kv_len,
+            actual_seq_q,
+            cp_index=None,
+            current_index_kv=None,
+            shared_index_buffer=None,
+            shared_block_tables=None,
+        ):
+            topk_calls.append(
+                {
+                    "current_index_kv": current_index_kv,
+                    "shared_index_buffer": shared_index_buffer,
+                    "shared_block_tables": shared_block_tables,
+                }
+            )
+            return torch.full((actual_seq_q, 2), len(topk_calls), dtype=torch.int32)
+
+        indexer._maybe_materialize_shared_index_buffer = fake_materialize
+        indexer._get_topk_ragged_with_cp = fake_get_topk
+
+        forward_batch = type(
+            "ForwardBatchStub",
+            (),
+            {
+                "nsa_cp_metadata": NSAContextParallelMetadata(
+                    kv_len_prev=5,
+                    kv_len_next=9,
+                    actual_seq_q_prev=3,
+                    actual_seq_q_next=2,
+                )
+            },
+        )()
+
+        result = Indexer._get_topk_in_seq_cp_pair(
+            indexer,
+            forward_batch,
+            layer_id=7,
+            q_fp8=torch.empty(5, 4),
+            weights=torch.empty(5, 2),
+            metadata=Metadata(),
+            current_index_kv=current_index_kv,
+        )
+
+        self.assertEqual(materialize_calls, [])
+        self.assertEqual(len(topk_calls), 2)
+        self.assertIs(topk_calls[0]["current_index_kv"], current_index_kv)
+        self.assertIs(topk_calls[1]["current_index_kv"], current_index_kv)
+        self.assertIsNone(topk_calls[0]["shared_index_buffer"])
+        self.assertIsNone(topk_calls[1]["shared_block_tables"])
+        self.assertEqual(result.tolist(), [[1, 1], [1, 1], [1, 1], [2, 2], [2, 2]])
+

 if __name__ == "__main__":
    unittest.main()