[DSV4] Generalize attention metadata, sparse prefill, and KV pool over compress ratios (#39921)

test/registered/amd/test_dsv4_hip_bcg_metadata.py

@@ -44,7 +44,7 @@ class TestDSV4HipBreakableCudaGraphMetadata(unittest.TestCase):
             positions_casual=torch.tensor([base + 5], dtype=torch.int32),
             swa_page_indices=torch.tensor([[base + 6, base + 7]], dtype=torch.int32),
             swa_topk_lengths=torch.tensor([base + 8], dtype=torch.int32),
-            c4_sparse_topk=512,
+            index_topk=512,
             swa_out_cache_loc=torch.tensor([base + 9], dtype=torch.int32),
             unified=UnifiedKvMetadata(),
         )

test/registered/attention/unittests/dsv4/test_deepseek_v4.py

@@ -318,6 +318,7 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase):
             swa_page_size=128,
             seq_lens=torch.tensor([max_seq_len, max_seq_len], **int32),
             query_start_loc=torch.tensor([0, 1, 2], **int32),
+            query_pos=torch.tensor([max_seq_len - 1, max_seq_len - 1], **int32),
             swa_token_ids=torch.empty(0, **int32),
             swa_first_pos=torch.zeros(2, **int32),
             swa_gather_lens=torch.zeros(2, **int32),
@@ -340,7 +341,8 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase):
                 [[base + 11, base + 12], [base + 13, base + 14]], dtype=torch.int32
             ),
             swa_topk_lengths=torch.tensor([base + 15, base + 16], dtype=torch.int32),
-            c4_sparse_topk=128,
+            index_topk=128,
+            present_ratios=(4, 128),
         )
         metadata.c4_out_loc = torch.tensor([base + 17, base + 18], dtype=torch.int32)
         metadata.c128_out_loc = torch.tensor([base + 19, base + 20], dtype=torch.int32)
@@ -370,6 +372,107 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase):
         metadata.c128_flashmla_metadata = object()
         return metadata
 
+    def test_present_ratios_gate_per_ratio_buffers(self):
+        from sglang.srt.layers.attention import deepseek_v4_backend as be
+
+        with mock.patch.object(be, "_create_flashmla_metadata", side_effect=object):
+            c4_only = self._make_core_metadata(0)
+            c4_only.present_ratios = (4,)
+            c4_only.index_topk = 512
+            c4_only.c128_page_indices = None
+            c4_only.c128_topk_lengths_clamp1 = None
+            c4_only.init_flashmla_related(is_prefill=True)
+            self.assertTrue(c4_only.has_c4)
+            self.assertFalse(c4_only.has_c128)
+            self.assertEqual(c4_only.sparse_page_indices(4).shape[0], 2)
+            self.assertIsNotNone(c4_only.sparse_raw_indices(4))
+            self.assertIsNone(c4_only.sparse_page_indices(128))
+            self.assertIsNotNone(c4_only.c4_flashmla_metadata)
+            self.assertIsNone(c4_only.c128_flashmla_metadata)
+
+            c128_only = self._make_core_metadata(0)
+            c128_only.present_ratios = (128,)
+            c128_only.index_topk = 512
+            c128_only.c4_topk_lengths_clamp1 = None
+            c128_only.init_flashmla_related(is_prefill=True)
+            self.assertFalse(c128_only.has_c4)
+            self.assertIsNone(c128_only.sparse_page_indices(4))
+            self.assertIsNone(c128_only.sparse_raw_indices(4))
+            self.assertIs(
+                c128_only.sparse_page_indices(128), c128_only.c128_page_indices
+            )
+            self.assertIsNone(c128_only.c4_flashmla_metadata)
+            self.assertIsNotNone(c128_only.c128_flashmla_metadata)
+
+        # Replay metadata must describe the same set of ratios as its source.
+        src = self._make_core_metadata(100)
+        src.present_ratios = (4,)
+        with self.assertRaises(AssertionError):
+            self._make_core_metadata(0).copy_(src)
+        with self.assertRaises(AssertionError):
+            self._make_core_metadata(0).refresh_for_breakable_cuda_graph_replay_(src)
+
+    def test_sparse_topk_accessors_route_by_ratio(self):
+        metadata = self._make_core_metadata(0)
+        page_indices = torch.full((2, 4), 3, dtype=torch.int32)
+        lengths = torch.tensor([1, 2], dtype=torch.int32)
+        raw = torch.full((2, 4), 5, dtype=torch.int32)
+
+        metadata.set_sparse_topk(
+            4, page_indices=page_indices, topk_lengths=lengths, raw_indices=raw
+        )
+        self.assertIs(metadata.sparse_page_indices(4), page_indices)
+        self.assertIs(metadata.sparse_topk_lengths(4), lengths)
+        self.assertIs(metadata.sparse_raw_indices(4), raw)
+
+        metadata.set_sparse_topk(128, page_indices=page_indices, topk_lengths=lengths)
+        self.assertIs(metadata.c128_page_indices, page_indices)
+        self.assertIs(metadata.sparse_topk_lengths(128), lengths)
+        with self.assertRaises(AssertionError):
+            metadata.set_sparse_topk(
+                128, page_indices=page_indices, topk_lengths=lengths, raw_indices=raw
+            )
+        with self.assertRaises(ValueError):
+            metadata.sparse_raw_indices(128)
+        for bad_ratio in (0, 7):
+            with self.assertRaises(ValueError):
+                metadata.sparse_page_indices(bad_ratio)
+            with self.assertRaises(ValueError):
+                metadata.set_sparse_topk(
+                    bad_ratio, page_indices=page_indices, topk_lengths=lengths
+                )
+
+    def test_cp_reindex_slices_present_fields_and_skips_absent_ones(self):
+        from sglang.srt.layers.attention import deepseek_v4_backend as be
+
+        metadata = self._make_core_metadata(0)
+        metadata.present_ratios = (4,)
+        metadata.c128_page_indices = None
+        metadata.c128_topk_lengths_clamp1 = None
+        parallel = SimpleNamespace(attn_cp_rank=1, attn_cp_size=2)
+        with mock.patch.object(be, "get_parallel", return_value=parallel):
+            metadata.apply_cp_reindex()
+
+        # Rank 1 of 2 keeps row 1 of every per-token field.
+        self.assertEqual(metadata.seq_lens_casual.tolist(), [8])
+        self.assertEqual(metadata.positions_casual.tolist(), [10])
+        self.assertEqual(metadata.page_table.tolist(), [[3, 4]])
+        self.assertEqual(metadata.swa_page_indices.tolist(), [[13, 14]])
+        self.assertEqual(metadata.swa_topk_lengths.tolist(), [16])
+        self.assertEqual(metadata.c4_topk_lengths_raw.tolist(), [22])
+        self.assertEqual(metadata.c4_topk_lengths_clamp1.tolist(), [24])
+        self.assertIsNone(metadata.c128_page_indices)
+        self.assertIsNone(metadata.c128_topk_lengths_clamp1)
+        # Cache-write locations stay in global logical order.
+        self.assertEqual(metadata.raw_out_loc.tolist(), [5, 6])
+        self.assertEqual(metadata.c4_out_loc.tolist(), [17, 18])
+
+        missing = self._make_core_metadata(0)
+        missing.swa_topk_lengths = None
+        with mock.patch.object(be, "get_parallel", return_value=parallel):
+            with self.assertRaises(AssertionError):
+                missing.apply_cp_reindex()
+
     def test_bcg_is_explicit_and_dsv4_backend_opt_in_only(self):
         from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
         from sglang.srt.layers.attention.deepseek_v4_backend import (
@@ -452,7 +555,7 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase):
                 metadata = backend.forward_metadata
                 if builds:
                     backend._build_sparse_prefill_chunk_cache.assert_called_once_with(
-                        batch, num_qo_tokens=num_qo_tokens
+                        batch, metadata.core_attn_metadata, num_qo_tokens=num_qo_tokens
                     )
                     self.assertIs(metadata.sparse_prefill_cache, cache)
                 else:
@@ -646,12 +749,10 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase):
         for max_seq_len in (3, 4, 255, 256, 259, 260):
             with self.subTest(max_seq_len=max_seq_len):
                 cache = self._make_sparse_prefill_cache(max_seq_len)
-                cache.ensure_c4(page_table, c4_page_size=64)
+                gather = cache.ensure_compressed(4, page_table, c_page_size=64)
                 expected_extent = max(max_seq_len // 4, 1)
-                self.assertEqual(cache.c4_flat_token_ids.numel(), 2 * expected_extent)
-                self.assertEqual(
-                    cache.c4_compressed_base.tolist(), [0, expected_extent]
-                )
+                self.assertEqual(gather.flat_token_ids.numel(), 2 * expected_extent)
+                self.assertEqual(gather.compressed_base.tolist(), [0, expected_extent])
 
     def test_sparse_prefill_c128_uses_live_extent(self):
         from sglang.srt.layers.attention.dsv4 import sparse_prefill_utils
@@ -670,9 +771,9 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase):
                     "combine_topk_swa_indices",
                     return_value=combined,
                 ) as combine:
-                    cache.ensure_c128(page_indices)
+                    gather = cache.ensure_c128(page_indices)
 
-                self.assertEqual(cache.c128_flat_token_ids.numel(), 2 * expected_extent)
+                self.assertEqual(gather.flat_token_ids.numel(), 2 * expected_extent)
                 self.assertEqual(combine.call_args.kwargs["topk"], expected_extent)
                 self.assertEqual(
                     combine.call_args.kwargs["topk_indices"].shape,

test/registered/kernel/attention/test_combine_topk_swa_indices.py

@@ -0,0 +1,177 @@
+"""Unit tests for ``combine_topk_swa_indices`` (DSV4 sparse prefill).
+
+Checks the Triton kernel against a per-row torch reference:
+
+1. ``test_trailing_extend_matches_reference``: the V4 layout, where the query
+   rows are the trailing extend tokens of each request and every top-k entry
+   inside the scanned prefix is valid.
+2. ``test_negative_one_holes_are_kept``: ``-1`` entries inside the top-k prefix
+   stay ``-1`` instead of being shifted by ``compressed_base``.
+3. ``test_non_trailing_query_positions``: absolute ``query_pos`` that are not
+   the trailing extend tokens, with a cross-chunk ``query_start_loc`` offset.
+4. ``test_swa_only_layer``: ``topk == 0`` writes only the window.
+"""
+
+import pytest
+import torch
+
+from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import (
+    combine_topk_swa_indices,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="combine_topk_swa_indices requires CUDA"
+)
+
+DEVICE = "cuda"
+WINDOW = 128
+# flash_mla_sparse_fwd reads rows in 128-wide tiles; the combined width is
+# padded to that multiple.
+TOPK_ALIGNMENT = 128
+
+
+def _i32(values):
+    return torch.tensor(values, dtype=torch.int32, device=DEVICE)
+
+
+def _reference(
+    topk_indices,
+    query_start_loc,
+    query_pos,
+    seq_lens,
+    gather_lens,
+    compressed_base,
+    swa_base,
+    window_size,
+    compress_ratio,
+    topk,
+):
+    num_tokens = topk_indices.shape[0]
+    width = -(-(topk + window_size) // TOPK_ALIGNMENT) * TOPK_ALIGNMENT
+    out = torch.full((num_tokens, width), -1, dtype=torch.int32, device=DEVICE)
+    lens = torch.zeros(num_tokens, dtype=torch.int32, device=DEVICE)
+    qsl = query_start_loc.tolist()
+    base = qsl[0]
+    for r in range(seq_lens.shape[0]):
+        gather_start = int(seq_lens[r]) - int(gather_lens[r])
+        for token_idx in range(qsl[r] - base, qsl[r + 1] - base):
+            pos = int(query_pos[token_idx])
+            topk_len = min((pos + 1) // compress_ratio, topk)
+            swa_len = min(pos + 1, window_size)
+            vals = topk_indices[token_idx, :topk_len]
+            out[token_idx, :topk_len] = torch.where(
+                vals >= 0, vals + compressed_base[r], torch.full_like(vals, -1)
+            )
+            window = torch.arange(swa_len, dtype=torch.int32, device=DEVICE)
+            out[token_idx, topk_len : topk_len + swa_len] = (
+                swa_base[r] + window + pos - swa_len + 1 - gather_start
+            )
+            lens[token_idx] = topk_len + swa_len
+    return out, lens
+
+
+def _check(**kwargs):
+    got_idx, got_lens = combine_topk_swa_indices(**kwargs)
+    ref_idx, ref_lens = _reference(**kwargs)
+    assert torch.equal(got_lens, ref_lens), (got_lens.tolist(), ref_lens.tolist())
+    assert torch.equal(got_idx, ref_idx)
+
+
+def _trailing_case(seq_lens, extend_lens, topk, compress_ratio, seed=0):
+    """Rows are the trailing ``extend_lens[r]`` tokens of request ``r``."""
+    gen = torch.Generator(device="cpu").manual_seed(seed)
+    query_pos = []
+    for seq_len, extend_len in zip(seq_lens, extend_lens):
+        query_pos.extend(range(seq_len - extend_len, seq_len))
+    num_tokens = len(query_pos)
+    topk_indices = torch.randint(
+        0, 1 << 20, (num_tokens, topk), generator=gen, dtype=torch.int32
+    ).to(DEVICE)
+    starts = [0]
+    for extend_len in extend_lens:
+        starts.append(starts[-1] + extend_len)
+    gather_lens = [min(s, e + WINDOW - 1) for s, e in zip(seq_lens, extend_lens)]
+    return dict(
+        topk_indices=topk_indices,
+        query_start_loc=_i32(starts),
+        query_pos=_i32(query_pos),
+        seq_lens=_i32(seq_lens),
+        gather_lens=_i32(gather_lens),
+        compressed_base=_i32([1000 * r for r in range(len(seq_lens))]),
+        swa_base=_i32([5000 + 300 * r for r in range(len(seq_lens))]),
+        window_size=WINDOW,
+        compress_ratio=compress_ratio,
+        topk=topk,
+    )
+
+
+@pytest.mark.parametrize("compress_ratio", [4, 128])
+@pytest.mark.parametrize(
+    "seq_lens, extend_lens",
+    [([96, 144], [3, 2]), ([7, 300, 1000], [7, 130, 5]), ([1], [1])],
+)
+def test_trailing_extend_matches_reference(seq_lens, extend_lens, compress_ratio):
+    _check(
+        **_trailing_case(seq_lens, extend_lens, topk=64, compress_ratio=compress_ratio)
+    )
+
+
+def test_negative_one_holes_are_kept():
+    case = _trailing_case([512, 640], [4, 4], topk=64, compress_ratio=4)
+    topk_indices = case["topk_indices"]
+    # Holes inside the scanned prefix (every row scans the full top-k here).
+    topk_indices[:, 0] = -1
+    topk_indices[:, 5] = -1
+    topk_indices[3, 10:20] = -1
+    _check(**case)
+    got_idx, got_lens = combine_topk_swa_indices(**case)
+    assert (got_idx[:, 0] == -1).all()
+    assert (got_idx[:, 5] == -1).all()
+    assert (got_idx[3, 10:20] == -1).all()
+    # The scanned prefix still counts the holes.
+    assert int(got_lens[0]) == 64 + WINDOW
+
+
+def test_non_trailing_query_positions():
+    # Two requests; each rank of a two-way interleave holds every other row of
+    # the extend, so the query positions are not the trailing tokens. The
+    # query_start_loc carries a cross-chunk offset that the kernel rebases.
+    seq_lens = [96, 144]
+    extend_lens = [6, 4]
+    query_pos = [90, 92, 94, 140, 142]
+    starts = [10, 13, 15]
+    num_tokens = len(query_pos)
+    topk = 32
+    topk_indices = torch.arange(
+        num_tokens * topk, dtype=torch.int32, device=DEVICE
+    ).view(num_tokens, topk)
+    gather_lens = [min(s, e + WINDOW - 1) for s, e in zip(seq_lens, extend_lens)]
+    _check(
+        topk_indices=topk_indices,
+        query_start_loc=_i32(starts),
+        query_pos=_i32(query_pos),
+        seq_lens=_i32(seq_lens),
+        gather_lens=_i32(gather_lens),
+        compressed_base=_i32([0, 24]),
+        swa_base=_i32([48, 181]),
+        window_size=WINDOW,
+        compress_ratio=4,
+        topk=topk,
+    )
+
+
+def test_swa_only_layer():
+    case = _trailing_case([200, 50], [2, 2], topk=0, compress_ratio=4)
+    case["topk_indices"] = torch.zeros((4, 1), dtype=torch.int32, device=DEVICE)
+    _check(**case)
+    _, got_lens = combine_topk_swa_indices(**case)
+    assert got_lens.tolist() == [WINDOW, WINDOW, 49, 50]
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__]))

test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py

@@ -17,6 +17,7 @@ import pytest
 import torch
 
 from sglang.kernels.ops.attention.dsv4.index_buf_accessor import SetKAndS
+from sglang.kernels.ops.attention.dsv4.kv_layout import KVLayout
 from sglang.kernels.ops.attention.dsv4.quant_k_cache import (
     quant_to_nope_fp8_rope_bf16_pack_triton,
 )
@@ -25,6 +26,7 @@ from sglang.kernels.ops.attention.sparse_mla_q8kv8_prefill_sm90 import (
 )
 from sglang.srt.layers.attention.deepseek_v4_backend import DeepseekV4AttnBackend
 from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import (
+    CompressedGather,
     SparsePrefillChunkCache,
     SparsePrefillWorkspace,
     use_dsv4_q8kv8_sparse_prefill,
@@ -90,6 +92,13 @@ class _TokenToKVPool:
         _ = layer_id
         return self._extra_key_buffer
 
+    def get_swa_key_layout(self) -> KVLayout:
+        return KVLayout.V4
+
+    def get_extra_key_layout(self, layer_id: int) -> KVLayout:
+        _ = layer_id
+        return KVLayout.V4
+
 
 def _sm90_available() -> bool:
     return torch.cuda.is_available() and is_sm90_supported()
@@ -219,7 +228,12 @@ def _make_sparse_prefill_case(
         * 0.05
     ).to(torch.bfloat16)
     attn_sink = torch.zeros(local_heads, dtype=torch.float32, device=device)
-    core_attn_metadata = SimpleNamespace()
+    # position + 1 of the five query rows: seq_lens [96, 144], extend [3, 2]
+    core_attn_metadata = SimpleNamespace(
+        seq_lens_casual=torch.tensor(
+            [94, 95, 96, 143, 144], dtype=torch.int32, device=device
+        )
+    )
     return backend, forward_batch, token_to_kv_pool, q, attn_sink, core_attn_metadata
 
 
@@ -236,10 +250,17 @@ def _populate_compress_metadata(
         core_attn_metadata.c4_sparse_raw_indices = torch.zeros(
             (16, 1), dtype=torch.int32, device=device
         )
+        # The sparse prefill path selects the ratio's raw top-k through this accessor.
+        core_attn_metadata.sparse_raw_indices = lambda ratio: (
+            core_attn_metadata.c4_sparse_raw_indices if ratio == 4 else None
+        )
     elif compress_ratio == 128:
         core_attn_metadata.c128_page_indices = torch.zeros(
             (16, 1), dtype=torch.int32, device=device
         )
+        core_attn_metadata.sparse_page_indices = lambda ratio: (
+            core_attn_metadata.c128_page_indices if ratio == 128 else None
+        )
 
 
 @contextmanager
@@ -248,9 +269,9 @@ def _patched_compressed_sparse_cache_paths(compress_ratio: int):
         yield
         return
 
-    old_ensure_c4 = SparsePrefillChunkCache.ensure_c4
+    old_ensure_compressed = SparsePrefillChunkCache.ensure_compressed
     old_ensure_c128 = SparsePrefillChunkCache.ensure_c128
-    old_combine_c4_layer = SparsePrefillChunkCache.combine_c4_layer
+    old_combine_compressed = SparsePrefillChunkCache.combine_compressed
 
     def _with_compressed_prefix(cache: SparsePrefillChunkCache, n_compressed: int):
         shifted_swa = torch.where(
@@ -272,33 +293,48 @@ def _patched_compressed_sparse_cache_paths(compress_ratio: int):
     def fake_ensure_c128(self, c128_page_indices):
         _ = c128_page_indices
         n_compressed = 8
-        self.c128_flat_token_ids = torch.arange(
-            n_compressed, dtype=torch.int64, device=self.swa_token_ids.device
-        )
-        self.c128_combined_indices, self.c128_combined_lens = _with_compressed_prefix(
-            self, n_compressed
+        device = self.swa_token_ids.device
+        combined_indices, combined_lens = _with_compressed_prefix(self, n_compressed)
+        gather = CompressedGather(
+            flat_token_ids=torch.arange(n_compressed, dtype=torch.int64, device=device),
+            compressed_base=torch.zeros(
+                self.num_reqs, dtype=torch.int32, device=device
+            ),
+            swa_base=torch.zeros(self.num_reqs, dtype=torch.int32, device=device),
+            combined_indices=combined_indices,
+            combined_lens=combined_lens,
         )
+        self.compressed[128] = gather
+        return gather
 
-    def fake_ensure_c4(self, page_table, extra_page_size):
+    def fake_ensure_compressed(self, compress_ratio, page_table, extra_page_size):
         _ = page_table, extra_page_size
         n_compressed = 8
-        self.c4_flat_token_ids = torch.arange(
-            n_compressed, dtype=torch.int64, device=self.swa_token_ids.device
+        device = self.swa_token_ids.device
+        gather = CompressedGather(
+            flat_token_ids=torch.arange(n_compressed, dtype=torch.int64, device=device),
+            compressed_base=torch.zeros(
+                self.num_reqs, dtype=torch.int32, device=device
+            ),
+            swa_base=torch.zeros(self.num_reqs, dtype=torch.int32, device=device),
         )
+        self.compressed[compress_ratio] = gather
+        return gather
 
-    def fake_combine_c4_layer(self, c4_sparse_raw_indices):
-        _ = c4_sparse_raw_indices
-        return _with_compressed_prefix(self, self.c4_flat_token_ids.shape[0])
+    def fake_combine_compressed(self, compress_ratio, sparse_raw_indices):
+        _ = sparse_raw_indices
+        n_compressed = self.compressed[compress_ratio].flat_token_ids.shape[0]
+        return _with_compressed_prefix(self, n_compressed)
 
     SparsePrefillChunkCache.ensure_c128 = fake_ensure_c128
-    SparsePrefillChunkCache.ensure_c4 = fake_ensure_c4
-    SparsePrefillChunkCache.combine_c4_layer = fake_combine_c4_layer
+    SparsePrefillChunkCache.ensure_compressed = fake_ensure_compressed
+    SparsePrefillChunkCache.combine_compressed = fake_combine_compressed
     try:
         yield
     finally:
-        SparsePrefillChunkCache.ensure_c4 = old_ensure_c4
+        SparsePrefillChunkCache.ensure_compressed = old_ensure_compressed
         SparsePrefillChunkCache.ensure_c128 = old_ensure_c128
-        SparsePrefillChunkCache.combine_c4_layer = old_combine_c4_layer
+        SparsePrefillChunkCache.combine_compressed = old_combine_compressed
 
 
 def _make_q8kv8_kernel_args(

test/registered/unit/disaggregation/test_disaggregation_wire.py

@@ -35,8 +35,6 @@ from sglang.srt.disaggregation.utils import (
     MetadataBuffers,
     build_transfer_entry_pairs,
     compute_mamba_state_slice_byte_blocks,
-    get_dsv4_c4_state_indices,
-    get_dsv4_c128_state_indices,
     get_qsa_pending_state_indices,
     setup_state_kv_args,
     should_send_replicated_state,
@@ -45,6 +43,11 @@ from sglang.srt.environ import envs
 from sglang.srt.layers.attention.dsa.utils import should_use_dsa_fused_topk
 from sglang.srt.managers.overlap_utils import FutureMap, RelayPayload
 from sglang.srt.managers.schedule_batch import ReqKvInfo
+from sglang.srt.mem_cache.deepseek_v4_compress_state import (
+    CompressStatePool,
+    c4_state_transfer_indices,
+    request_scoped_state_transfer_indices,
+)
 from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
 from sglang.srt.mem_cache.qsa_kv_pool import (
     QSA_ROPE_STATE_LAYER_ID,
@@ -775,8 +778,8 @@ class TestEagleDsaSeedTransfer(CustomTestCase):
 class TestDSV4C4StateIndices(unittest.TestCase):
     def test_non_mtp_to_mtp_maps_the_same_logical_positions(self):
         # seq_len=13 keeps logical positions [8, 13) for the overlap C4 state.
-        src = get_dsv4_c4_state_indices(2, 13, ring_size=8)
-        dst = get_dsv4_c4_state_indices(2, 13, ring_size=16)
+        src = c4_state_transfer_indices(2, 13, ring_size=8)
+        dst = c4_state_transfer_indices(2, 13, ring_size=16)
 
         np.testing.assert_array_equal(src, np.array([16, 17, 18, 19, 20]))
         np.testing.assert_array_equal(dst, np.array([40, 41, 42, 43, 44]))
@@ -784,59 +787,126 @@ class TestDSV4C4StateIndices(unittest.TestCase):
 
     def test_ring_wrap_preserves_position_order(self):
         np.testing.assert_array_equal(
-            get_dsv4_c4_state_indices(0, 10, ring_size=8),
+            c4_state_transfer_indices(0, 10, ring_size=8),
             np.array([4, 5, 6, 7, 0, 1], dtype=np.int32),
         )
 
     def test_short_and_empty_sequences(self):
         np.testing.assert_array_equal(
-            get_dsv4_c4_state_indices(3, 3, ring_size=8),
+            c4_state_transfer_indices(3, 3, ring_size=8),
             np.array([24, 25, 26], dtype=np.int32),
         )
         np.testing.assert_array_equal(
-            get_dsv4_c4_state_indices(3, 0, ring_size=8),
+            c4_state_transfer_indices(3, 0, ring_size=8),
             np.empty((0,), dtype=np.int32),
         )
 
     def test_invalid_ring_size_is_rejected(self):
         with self.assertRaises(ValueError):
-            get_dsv4_c4_state_indices(0, 8, ring_size=4)
+            c4_state_transfer_indices(0, 8, ring_size=4)
         with self.assertRaises(ValueError):
-            get_dsv4_c4_state_indices(0, 8, ring_size=10)
+            c4_state_transfer_indices(0, 8, ring_size=10)
 
 
 class TestDSV4C128StateIndices(unittest.TestCase):
     def test_online_aligned_boundary_has_no_partial_state(self):
         np.testing.assert_array_equal(
-            get_dsv4_c128_state_indices(7, 256, online=True, ring_size=1),
+            request_scoped_state_transfer_indices(
+                7, 256, ratio=128, online=True, ring_size=1
+            ),
             np.empty((0,), dtype=np.int32),
         )
 
     def test_online_partial_boundary_uses_request_slot(self):
         np.testing.assert_array_equal(
-            get_dsv4_c128_state_indices(7, 257, online=True, ring_size=1),
+            request_scoped_state_transfer_indices(
+                7, 257, ratio=128, online=True, ring_size=1
+            ),
             np.array([7], dtype=np.int32),
         )
 
     def test_offline_aligned_boundary_has_no_partial_state(self):
         np.testing.assert_array_equal(
-            get_dsv4_c128_state_indices(7, 256, online=False, ring_size=128),
+            request_scoped_state_transfer_indices(
+                7, 256, ratio=128, online=False, ring_size=128
+            ),
             np.empty((0,), dtype=np.int32),
         )
 
     def test_offline_partial_boundary_uses_request_local_page(self):
         np.testing.assert_array_equal(
-            get_dsv4_c128_state_indices(7, 129, online=False, ring_size=256),
+            request_scoped_state_transfer_indices(
+                7, 129, ratio=128, online=False, ring_size=256
+            ),
             np.array([15], dtype=np.int32),
         )
 
 
+def _make_state_pool(*, ratio, request_scoped, online=False, ring_size=256):
+    pool = object.__new__(CompressStatePool)
+    pool.ratio = ratio
+    pool.request_scoped = request_scoped
+    pool.online = online
+    pool.ring_size = ring_size
+    return pool
+
+
+class TestDSV4RequestStateTransfer(unittest.TestCase):
+    def _kv(self, *pools):
+        kv = object.__new__(DeepSeekV4TokenToKVPool)
+        kv.compress_state_pools = [None, *pools]
+        return kv
+
+    def test_pool_delegates_to_its_request_scoped_state_pools(self):
+        # One state pool per compressed layer; all c128 layers share the ring layout.
+        kv = self._kv(
+            _make_state_pool(ratio=4, request_scoped=False),
+            *[
+                _make_state_pool(ratio=128, request_scoped=True, ring_size=256)
+                for _ in range(20)
+            ],
+        )
+        np.testing.assert_array_equal(
+            kv.request_state_transfer_indices(7, 129),
+            request_scoped_state_transfer_indices(
+                7, 129, ratio=128, online=False, ring_size=256
+            ),
+        )
+        np.testing.assert_array_equal(
+            kv.request_state_transfer_indices(7, 256), np.empty((0,), dtype=np.int32)
+        )
+
+    def test_online_pool_ships_the_request_row(self):
+        kv = self._kv(
+            _make_state_pool(ratio=128, request_scoped=True, online=True, ring_size=1)
+        )
+        np.testing.assert_array_equal(
+            kv.request_state_transfer_indices(7, 257), np.array([7], dtype=np.int32)
+        )
+
+    def test_requires_request_scoped_pools_with_one_ring_layout(self):
+        with self.assertRaises(AssertionError):
+            self._kv(
+                _make_state_pool(ratio=4, request_scoped=False)
+            ).request_state_transfer_indices(0, 5)
+        with self.assertRaises(AssertionError):
+            self._kv(
+                _make_state_pool(ratio=128, request_scoped=True, ring_size=128),
+                _make_state_pool(ratio=128, request_scoped=True, ring_size=256),
+            ).request_state_transfer_indices(0, 5)
+
+    def test_page_scoped_pool_has_no_transfer_indices(self):
+        with self.assertRaises(AssertionError):
+            _make_state_pool(ratio=4, request_scoped=False).transfer_indices(0, 5)
+
+
 def _buf_infos(*ptrs):
     return list(ptrs), [ptr + 100 for ptr in ptrs], [ptr + 200 for ptr in ptrs]
 
 
 def _make_dsv4_target(*, unified, mapping=None):
     pool = object.__new__(DeepSeekV4TokenToKVPool)
+    pool.compression_ratios = [0, 4, 128]
     pool._unified_kv = unified
     pool.page_size = 256
     pool.sliding_window = 128