[AMD][DSV4] Fix unified-KV pool sizing and SWA ring accounting (#30315)

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

@@ -7,7 +7,11 @@ import pytest
 import torch
 import triton
 
-from sglang.kernels.ops.attention.dsv4 import compress_forward
+from sglang.kernels.ops.attention.dsv4 import (
+    CompressorDecodePlan,
+    CompressorPrefillPlan,
+    compress_forward,
+)
 from sglang.srt.utils import get_device
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.kernels.deepseek_v4.common import (
@@ -122,6 +126,91 @@ def _make_inputs(
 # -----------------------------------------------------------------------------
 
 
+@pytest.mark.parametrize("ring_size", [8, 16])
+@pytest.mark.parametrize(
+    ("gpu_inputs", "use_cuda_graph"),
+    [(False, False), (True, False), (True, True)],
+)
+def test_unified_request_ring_plans_ignore_full_to_state(
+    ring_size: int, gpu_inputs: bool, use_cuda_graph: bool
+) -> None:
+    """C4 plans must address state by request slot, not the full-cache map."""
+    device = torch.device(get_device())
+    req_pool_indices = torch.tensor([2, 5], dtype=torch.int64, device=device)
+    req_to_token = torch.zeros((6, 16), dtype=torch.int32, device=device)
+    full_to_state = torch.zeros(1, dtype=torch.int64, device=device)
+    seq_lens = torch.tensor([8, 12], dtype=torch.int64)
+    extend_lens = torch.tensor([4, 4], dtype=torch.int64)
+    if gpu_inputs:
+        seq_lens = seq_lens.to(device)
+        extend_lens = extend_lens.to(device)
+
+    prefill = CompressorPrefillPlan.generate(
+        compress_ratio=RATIO,
+        req_pool_indices=req_pool_indices,
+        seq_lens=seq_lens,
+        extend_lens=extend_lens,
+        req_to_token=req_to_token,
+        full_to_state=full_to_state,
+        swa_page_size=256,
+        ring_size=ring_size,
+        num_q_tokens=8,
+        use_cuda_graph=use_cuda_graph,
+        use_req_ring=True,
+    )
+    plan_c = prefill.plan_c.view(torch.int32).reshape(-1, 4).cpu()
+    plan_w = prefill.plan_w.view(torch.int32).reshape(-1, 2).cpu()
+
+    valid_c = plan_c[plan_c[:, 2] >= 0]
+    got_reads = {
+        int(row[1].item()) & 0xFFFF: (int(row[2].item()), int(row[3].item()))
+        for row in valid_c
+    }
+    expected_reads = {
+        3: (
+            (2 * ring_size + 3 % ring_size) // RATIO,
+            (2 * ring_size + 7 % ring_size) // RATIO,
+        ),
+        7: (
+            (5 * ring_size + 7 % ring_size) // RATIO,
+            (5 * ring_size + 11 % ring_size) // RATIO,
+        ),
+    }
+    assert got_reads == expected_reads
+
+    valid_w = plan_w[plan_w[:, 1] >= 0]
+    got_writes = {int(row[0].item()): int(row[1].item()) for row in valid_w}
+    expected_writes = {
+        **{j: 2 * ring_size + (4 + j) % ring_size for j in range(4)},
+        **{4 + j: 5 * ring_size + (8 + j) % ring_size for j in range(4)},
+    }
+    assert got_writes == expected_writes
+    assert {got_writes[j] for j in range(4)}.isdisjoint(
+        {got_writes[j] for j in range(4, 8)}
+    )
+
+    decode = CompressorDecodePlan.generate(
+        compress_ratio=RATIO,
+        req_pool_indices=req_pool_indices,
+        req_to_token=req_to_token,
+        full_to_state=full_to_state,
+        seq_lens=torch.tensor([8, 12], dtype=torch.int64, device=device),
+        swa_page_size=256,
+        ring_size=ring_size,
+        use_req_ring=True,
+    )
+    got_decode = decode.plan_d.view(torch.int32).reshape(-1, 4).cpu()
+    expected_decode = torch.tensor(
+        [
+            [8, 2 * ring_size + 7 % ring_size, *expected_reads[3]],
+            [12, 5 * ring_size + 11 % ring_size, *expected_reads[7]],
+        ],
+        dtype=torch.int32,
+    )
+    assert torch.equal(got_decode, expected_decode)
+    assert got_decode[0, 1] != got_decode[1, 1]
+
+
 @pytest.mark.parametrize("mode", ["legacy", "paged"])
 @pytest.mark.parametrize("seq_len", [4, 8, 32, 256, 1024])
 def test_prefill_no_context(mode: str, seq_len: int) -> None:

test/registered/unit/managers/test_prefill_adder.py

@@ -9,10 +9,7 @@ from sglang.srt.managers.schedule_policy import (
     PrefillAdder,
     estimate_prefill_extend_tile_metrics,
 )
-from sglang.srt.mem_cache.base_prefix_cache import (
-    DecLockRefResult,
-    IncLockRefResult,
-)
+from sglang.srt.mem_cache.base_prefix_cache import DecLockRefResult, IncLockRefResult
 from sglang.srt.runtime_context import get_context
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.srt.utils.common import Range
@@ -71,6 +68,13 @@ class TestPrefillAdder(CustomTestCase):
         allocator.swa_available_size.return_value = swa_available_size
         allocator.available_size.return_value = available_size
         allocator.size_swa = size_swa
+        # get_kvcache().[_unified_kv] gates the unified-KV SWA-ring accounting
+        # path in schedule_policy.add_chunked_req / rem_swa_tokens. A bare
+        # MagicMock auto-creates any attribute access as a truthy Mock, so
+        # without this the getattr(..., "_unified_kv", False) default never
+        # triggers and these tests silently exercise the unified-KV branch
+        # instead of the standard hybrid-SWA one they intend to cover.
+        allocator.get_kvcache.return_value._unified_kv = False
         return allocator
 
     def create_running_batch(self, reqs=None) -> MagicMock:

test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py

@@ -23,6 +23,9 @@ class _FakeAllocator:
         self.alloc_calls = []
         self.extend_calls = []
 
+    def get_kvcache(self):
+        return None
+
     def available_size(self):
         return 1 << 30
 

test/registered/unit/mem_cache/test_dsv4_c4_state_lifecycle.py

@@ -0,0 +1,119 @@
+"""CPU/mock tests for unified DSV4 C4 request-state lifecycle."""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import torch
+
+from sglang.srt.mem_cache.allocation import alloc_req_slots
+from sglang.srt.mem_cache.deepseek_v4_compress_state import KVAndScore
+from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
+from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+from sglang.srt.model_executor.pool_configurator import DSV4PoolConfigurator
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+def _request(req_pool_idx=None, *, reused=False):
+    return SimpleNamespace(
+        kv=SimpleNamespace(
+            req_pool_idx=req_pool_idx,
+            kv_committed_len=1 if reused else 0,
+            kv_allocated_len=1 if reused else 0,
+            holds_kv=reused,
+        ),
+        inflight_middle_chunks=1 if reused else 0,
+    )
+
+
+def _c4_pool(rows: int, width: int, ring_size: int):
+    return SimpleNamespace(
+        ratio=4,
+        ring_size=ring_size,
+        kv_score_buffer=KVAndScore(torch.full((rows, width), 7.0)),
+    )
+
+
+class TestUnifiedC4StateLifecycle(unittest.TestCase):
+    def test_pool_size_is_exact_request_ring_product(self):
+        configurator = object.__new__(DSV4PoolConfigurator)
+        configurator.disaggregation_mode = "decode"
+        configurator.disaggregation_decode_extra_slots = 3
+        configurator.c4_ring_size = 16
+
+        self.assertEqual(configurator._unified_c4_state_pool_size(10), 14 * 16)
+
+    def test_clear_resets_only_selected_request_rings(self):
+        ring_size = 8
+        logical_rows = 4 * ring_size
+        physical_rows = logical_rows + ring_size + 4
+        attn = _c4_pool(physical_rows, width=12, ring_size=ring_size)
+        indexer = _c4_pool(physical_rows, width=8, ring_size=ring_size)
+        c128 = SimpleNamespace(
+            ratio=128,
+            ring_size=128,
+            kv_score_buffer=KVAndScore(torch.full((physical_rows, 8), 9.0)),
+        )
+
+        token_pool = object.__new__(DeepSeekV4TokenToKVPool)
+        token_pool._unified_kv = True
+        token_pool.compress_state_pools = [attn, c128]
+        token_pool.indexer_compress_state_pools = [indexer, None]
+        token_pool.get_ring_size = MagicMock(return_value=ring_size)
+
+        token_pool.clear_c4_req_states([1, 3])
+
+        selected = torch.tensor(list(range(8, 16)) + list(range(24, 32)))
+        untouched = torch.tensor(list(range(0, 8)) + list(range(16, 24)))
+        for pool in (attn, indexer):
+            state = pool.kv_score_buffer.kv_score
+            half = state.shape[-1] // 2
+            self.assertTrue(
+                torch.equal(
+                    state[selected, :half], torch.zeros_like(state[selected, :half])
+                )
+            )
+            self.assertTrue(torch.isneginf(state[selected, half:]).all())
+            self.assertTrue((state[untouched] == 7).all())
+            self.assertTrue((state[logical_rows:] == 7).all())
+        self.assertTrue((c128.kv_score_buffer.kv_score == 9).all())
+
+    def test_alloc_clears_new_slots_but_not_reused_slots(self):
+        req_pool = ReqToTokenPool(3, 16, "cpu", enable_memory_saver=False)
+        token_pool = MagicMock()
+        reused = _request()
+
+        # First admission: a brand-new slot, so its C4 ring must be cleared.
+        (reused_idx,) = alloc_req_slots(
+            req_pool, [reused], None, token_to_kv_pool=token_pool
+        )
+        token_pool.clear_c4_req_states.assert_called_once_with([reused_idx])
+
+        # Chunked continuation reuses the same slot -- clearing it here would
+        # wipe the state captured by the previous chunk.
+        token_pool.clear_c4_req_states.reset_mock()
+        reused.kv.req_pool_idx = reused_idx
+        reused.kv.kv_committed_len = 1
+        reused.kv.kv_allocated_len = 1
+        reused.kv.holds_kv = True
+        reused.inflight_middle_chunks = 1
+        self.assertEqual(
+            alloc_req_slots(req_pool, [reused], None, token_to_kv_pool=token_pool),
+            [reused_idx],
+        )
+        token_pool.clear_c4_req_states.assert_not_called()
+
+        # Mixed batch: only the newly allocated slot is cleared.
+        fresh = _request()
+        indices = alloc_req_slots(
+            req_pool, [reused, fresh], None, token_to_kv_pool=token_pool
+        )
+        self.assertEqual(indices[0], reused_idx)
+        self.assertNotEqual(indices[1], reused_idx)
+        token_pool.clear_c4_req_states.assert_called_once_with([indices[1]])
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_hisparse_allocator.py

@@ -131,6 +131,7 @@ class TestDeepSeekV4HiSparseAllocator(CustomTestCase):
         queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
         queue.req_to_token_pool = req_to_token_pool
         queue.token_to_kv_pool_allocator = allocator
+        queue.token_to_kv_pool = None
         queue.tree_cache = SimpleNamespace(
             evictable_size=MagicMock(return_value=0),
             protected_size=MagicMock(return_value=0),

test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py

@@ -36,6 +36,11 @@ def _make_self(*, page_size: int, full_available: int, swa_available: int):
 
     return SimpleNamespace(
         page_size=page_size,
+        # alloc_extend branches on self._unified to skip the vestigial paged SWA
+        # allocator on the unified-KV path. This stub exercises the standard
+        # hybrid-SWA path, so pin it False rather than letting the attribute go
+        # missing (SimpleNamespace raises instead of defaulting).
+        _unified=False,
         full_attn_allocator=SimpleNamespace(
             available_size=lambda: full_available,
             alloc_extend=MagicMock(return_value=full_indices),