Optimize C128 state pool allocation using request state pool (#28612)

test/registered/jit/benchmark/bench_online_c128_mtp.py

@@ -43,7 +43,6 @@ class BenchmarkCase:
     seq_lens: torch.Tensor
     req_pool_indices: torch.Tensor
     req_to_token: torch.Tensor
-    full_to_swa: torch.Tensor
     ape: torch.Tensor
     state: torch.Tensor
     layer_bs: int
@@ -86,10 +85,6 @@ def make_case(batch_size: int, num_verify_tokens: int) -> BenchmarkCase:
     req_to_token = make_req_to_token(batch_size, max_seq_len, num_chunks)
 
     num_full_locs = batch_size * num_chunks
-    full_to_swa = (
-        torch.arange(num_full_locs, dtype=torch.int64, device=DEFAULT_DEVICE)
-        * SWA_PAGE_SIZE
-    )
 
     state_slot_stride = num_full_locs
     state = torch.empty(
@@ -112,7 +107,6 @@ def make_case(batch_size: int, num_verify_tokens: int) -> BenchmarkCase:
         seq_lens=seq_lens,
         req_pool_indices=req_pool_indices,
         req_to_token=req_to_token,
-        full_to_swa=full_to_swa,
         ape=ape,
         state=state,
         layer_bs=batch_size,
@@ -127,11 +121,9 @@ def call_write_prefix(module, case: BenchmarkCase) -> None:
         case.seq_lens,
         case.req_pool_indices,
         case.req_to_token,
-        case.full_to_swa,
         case.ape,
         case.state,
         case.layer_bs,
-        SWA_PAGE_SIZE,
         case.num_verify_tokens,
         case.state_slot_stride,
     )
@@ -153,8 +145,10 @@ def call_write_prefix(module, case: BenchmarkCase) -> None:
 def benchmark(
     batch_size: int, num_verify_tokens: int, launch_mode: str
 ) -> tuple[float, float, float]:
-    module = _jit_online_c128_mtp_module(HEAD_DIM)
     case = make_case(batch_size, num_verify_tokens)
+    module = _jit_online_c128_mtp_module(
+        HEAD_DIM, case.seq_lens.dtype, case.req_pool_indices.dtype
+    )
     fn = lambda: call_write_prefix(module, case)
 
     if launch_mode == "cuda_graph":

test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py

@@ -126,6 +126,7 @@ class TestUnifiedDeepSeekV4FlashHiCacheL3(AccuracyTwoPassMixin, CustomTestCase):
 
     l3_prefetch_page_size = 256
     l3_prefetch_prompt_pages = 4
+    max_running_requests = 4
 
     @classmethod
     def setUpClass(cls):
@@ -164,6 +165,8 @@ class TestUnifiedDeepSeekV4FlashHiCacheL3(AccuracyTwoPassMixin, CustomTestCase):
                 "file",
                 "--swa-full-tokens-ratio",
                 "0.25",
+                "--max-running-requests",
+                str(cls.max_running_requests),
             ],
             env={
                 "SGLANG_DSV4_FP4_EXPERTS": "0",

test/registered/unit/disaggregation/test_disaggregation_wire.py

@@ -9,6 +9,7 @@ from sglang.srt.disaggregation.common.utils import (
     unpack_int_lists,
     unpack_list_of_buffers,
 )
+from sglang.srt.disaggregation.utils import get_dsv4_c128_state_indices
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=2, suite="base-a-test-cpu")
@@ -90,5 +91,31 @@ class TestGroupConcurrentContiguous(unittest.TestCase):
             group_concurrent_contiguous(self._arr([1, 2, 3]), self._arr([1, 2]))
 
 
+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),
+            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),
+            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),
+            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),
+            np.array([15], dtype=np.int32),
+        )
+
+
 if __name__ == "__main__":
     unittest.main()