[kv-shard 2/4] Sharded pools (#37615)

Co-authored-by: Zhangheng <hzh0425@apache.org>

test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py

@@ -25,6 +25,7 @@ import torch
 from sglang.srt.configs.model_config import AttentionArch
 from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend
 from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator
+from sglang.srt.mem_cache.page_interleave_pool import PageInterleaveKVPoolMixin
 from sglang.srt.runtime_context import get_context
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
@@ -32,7 +33,13 @@ from sglang.test.test_utils import CustomTestCase
 register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
 
 
-def _make_prefill_aware_swa_runner(*, pool_size: int, max_context_len: int = 64):
+class _FakeShardPool(PageInterleaveKVPoolMixin):
+    pass
+
+
+def _make_prefill_aware_swa_runner(
+    *, pool_size: int, max_context_len: int = 64, token_to_kv_pool=None
+):
     """A minimal fake ModelRunner that reaches FlashAttentionBackend.__init__'s
     is_prefill_aware_swa branch (mirrors how models like
     python/sglang/srt/models/unlimited_ocr.py opt in)."""
@@ -63,7 +70,7 @@ def _make_prefill_aware_swa_runner(*, pool_size: int, max_context_len: int = 64)
         enable_prefill_cp=False,
         enable_dp_attention=False,
     )
-    token_to_kv_pool = object()
+    token_to_kv_pool = token_to_kv_pool if token_to_kv_pool is not None else object()
     token_to_kv_pool_allocator = object()
     return SimpleNamespace(
         sliding_window_size=None,
@@ -94,6 +101,16 @@ def _make_prefill_aware_swa_runner(*, pool_size: int, max_context_len: int = 64)
 
 @unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
 class TestPrefillAwareSwaPrefillLensBound(CustomTestCase):
+    def test_sharded_pool_requests_cpu_sequence_lengths(self):
+        runner = _make_prefill_aware_swa_runner(
+            pool_size=8, token_to_kv_pool=_FakeShardPool()
+        )
+
+        with get_context().override_server_args():
+            backend = FlashAttentionBackend(runner)
+
+        self.assertTrue(backend.needs_cpu_seq_lens)
+
     def test_buffer_covers_full_req_pool_idx_range(self):
         pool_size = 8
         runner = _make_prefill_aware_swa_runner(pool_size=pool_size)

test/registered/unit/layers/attention/test_kv_shard_hooks.py

@@ -0,0 +1,112 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import sys
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from sglang.srt.layers.attention.kv_shard_hooks import (
+    get_kv_shard_pool,
+    prepare_kv_shard_forward,
+)
+from sglang.srt.mem_cache.page_interleave_pool import PageInterleaveKVPoolMixin
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+class _RecordingPool(PageInterleaveKVPoolMixin):
+    def __init__(self):
+        self.begin_args = None
+        self.begin_calls = 0
+        self.end_calls = 0
+
+    def begin_shard_extend(self, *args):
+        self.begin_args = args
+        self.begin_calls += 1
+
+    def end_shard_extend(self):
+        self.end_calls += 1
+
+
+def _batch(mode: ForwardMode):
+    return SimpleNamespace(
+        forward_mode=mode,
+        req_pool_indices=torch.tensor([3], dtype=torch.int64),
+        extend_prefix_lens_cpu=[64],
+        seq_lens_cpu=torch.tensor([128], dtype=torch.int64),
+    )
+
+
+def test_detects_only_page_interleaved_pools():
+    pool = _RecordingPool()
+
+    assert get_kv_shard_pool(pool) is pool
+    assert get_kv_shard_pool(object()) is None
+
+
+@pytest.mark.parametrize(
+    "mode, active",
+    [
+        (ForwardMode.EXTEND, True),
+        (ForwardMode.MIXED, True),
+        (ForwardMode.SPLIT_PREFILL, True),
+        (ForwardMode.DECODE, False),
+        (ForwardMode.IDLE, False),
+        (ForwardMode.TARGET_VERIFY, False),
+        (ForwardMode.DRAFT_EXTEND_V2, False),
+    ],
+)
+def test_prepare_updates_the_pool_lifecycle(mode, active):
+    pool = _RecordingPool()
+    req_to_token = torch.arange(8)
+    batch = _batch(mode)
+
+    assert prepare_kv_shard_forward(pool, req_to_token, batch) is active
+    assert pool.begin_calls == int(active)
+    assert pool.end_calls == int(not active)
+    if active:
+        assert all(
+            actual is expected
+            for actual, expected in zip(
+                pool.begin_args,
+                (
+                    req_to_token,
+                    batch.req_pool_indices,
+                    batch.extend_prefix_lens_cpu,
+                    batch.seq_lens_cpu,
+                ),
+            )
+        )
+    else:
+        assert pool.begin_args is None
+
+
+@pytest.mark.parametrize(
+    "missing_field",
+    ["req_pool_indices", "extend_prefix_lens_cpu", "seq_lens_cpu"],
+)
+def test_extend_requires_host_metadata(missing_field):
+    batch = _batch(ForwardMode.EXTEND)
+    setattr(batch, missing_field, None)
+
+    with pytest.raises(RuntimeError, match="requires request indices and CPU"):
+        prepare_kv_shard_forward(_RecordingPool(), torch.empty(0), batch)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))