Remove full-cache scans from CP owner-lane allocation

The CP shared-KV allocator was still doing total-cache-sized CPU work in the scheduler hot path. That cannot be hidden by GPU overlap, so owner-lane allocation now maintains per-owner free/release buckets and consumes request-sized prefixes instead of rebuilding masks over the full free-page tensor on each request.\n\nThe benchmark was extended to isolate L1 stats, selection, and allocation costs, and the CPU layout tests now install a complete sgl_kernel stub before importing SGLang helpers so remote unit collection does not abort in native extension loading.\n\nConstraint: Allocator CPU work blocks scheduler progress and cannot overlap with GPU forward execution.\nConstraint: CPU unit tests must not load native sgl_kernel on remote images where the loader can SIGABRT.\nRejected: Keep contiguous-run search over full free_pages | still scales with cache capacity and measured multi-ms overhead.\nRejected: Treat remote collection abort as an environment-only issue | it prevented allocator regression coverage and was fixable with a test-local stub.\nConfidence: high\nScope-risk: moderate\nDirective: CP owner-lane allocation is bucket-based; do not reintroduce full free_pages scans on the hot path without benchmark evidence.\nTested: Local py_compile for touched files\nTested: Local benchmark unit test, 6 passed\nTested: Remote benchmark unit test, 6 passed\nTested: Remote test_alloc_pages_with_owners.py, 10 passed\nTested: Remote test_cp_shared_kv_layout.py, 27 passed\nTested: Remote production allocator microbench shows select/alloc p50 reduced from ms-scale to sub-ms scale\nNot-tested: Full ETE traffic run after allocator bucket change
2026-06-02 08:41:00 +08:00
parent ce3a20d11b
commit 7c8fa2f71c
5 changed files with 906 additions and 115 deletions
@@ -1,20 +1,77 @@
+import sys
+import types
 import unittest
 from unittest.mock import patch

 import numpy as np
 import torch

+# This CPU unit file must not import the native sgl_kernel package during test
+# collection.  Some remote images abort inside sgl_kernel's architecture-specific
+# loader instead of raising ImportError/RuntimeError, so a try/except fallback is
+# not sufficient here.  Install a minimal stub before importing sglang helpers.
+if "sgl_kernel" not in sys.modules:
+    sgl_kernel_stub = types.ModuleType("sgl_kernel")
+    sgl_kernel_stub.__file__ = "sgl_kernel_stub.py"
+    sgl_kernel_stub.__path__ = []
+
+    def _sgl_kernel_getattr(name):
+        if name.startswith("__"):
+            raise AttributeError(name)
+        fn = lambda *args, **kwargs: None
+        setattr(sgl_kernel_stub, name, fn)
+        return fn
+
+    sgl_kernel_stub.__getattr__ = _sgl_kernel_getattr
+    sys.modules["sgl_kernel"] = sgl_kernel_stub
+
+if "sgl_kernel.kvcacheio" not in sys.modules:
+    kvcacheio_stub = types.ModuleType("sgl_kernel.kvcacheio")
+    kvcacheio_stub.__file__ = "sgl_kernel_kvcacheio_stub.py"
+
+    def _kvcacheio_getattr(name):
+        if name.startswith("__"):
+            raise AttributeError(name)
+        fn = lambda *args, **kwargs: None
+        setattr(kvcacheio_stub, name, fn)
+        return fn
+
+    kvcacheio_stub.__getattr__ = _kvcacheio_getattr
+    sys.modules["sgl_kernel.kvcacheio"] = kvcacheio_stub
+
+if "sgl_kernel.quantization" not in sys.modules:
+    quantization_stub = types.ModuleType("sgl_kernel.quantization")
+    quantization_stub.__file__ = "sgl_kernel_quantization_stub.py"
+
+    def _quantization_getattr(name):
+        if name.startswith("__"):
+            raise AttributeError(name)
+        fn = lambda *args, **kwargs: None
+        setattr(quantization_stub, name, fn)
+        return fn
+
+    quantization_stub.__getattr__ = _quantization_getattr
+    sys.modules["sgl_kernel.quantization"] = quantization_stub
+
 _sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT")
-try:
-    _sgl_kernel_lib.define(
+for _schema in (
+    "sgl_per_token_group_quant_8bit(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
+    "sgl_per_token_group_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
+    "sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()",
+    "fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor",
+    "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor",
+    (
        "moe_fused_gate(Tensor input_tensor, Tensor? bias, int num_expert_group, "
        "int topk_group, int topk, int num_fused_shared_experts, "
        "float routed_scaling_factor, bool apply_routed_scaling_factor_on_output) "
        "-> (Tensor, Tensor)"
-    )
-except RuntimeError as exc:
-    if "already" not in str(exc).lower() and "duplicate" not in str(exc).lower():
-        raise
+    ),
+):
+    try:
+        _sgl_kernel_lib.define(_schema)
+    except RuntimeError as exc:
+        if "already" not in str(exc).lower() and "duplicate" not in str(exc).lower():
+            raise

 from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
 from sglang.test.ci.ci_register import register_cpu_ci
@@ -392,7 +449,7 @@ class TestCPSharedPagedAllocator(CustomTestCase):
        self.assertEqual(allocator.free_pages.tolist(), [2, 3, 4])
        self.assertEqual(allocator.release_pages.tolist(), [])

-    def test_contiguous_owner_lane_selection_prefers_later_physical_run(self):
+    def test_owner_lane_selection_uses_sorted_contiguous_bucket_prefix(self):
        from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator

        page_size = 64
@@ -418,11 +475,13 @@ class TestCPSharedPagedAllocator(CustomTestCase):

        self.assertIsNotNone(locs)
        logical_pages = locs.view(-1, page_size)[:, 0] // page_size
-        self.assertEqual(logical_pages.tolist(), [9, 13, 17])
+        self.assertEqual(logical_pages.tolist(), [1, 5, 9])
        self.assertEqual(
            ((logical_pages - 1) % cp_size).tolist(),
            [0, 0, 0],
        )
+        physical_pages = torch.div(logical_pages - 1, cp_size, rounding_mode="floor")
+        self.assertTrue(torch.all(physical_pages[1:] - physical_pages[:-1] == 1))
        for selected_page in logical_pages.tolist():
            self.assertNotIn(selected_page, allocator.free_pages.tolist())