b56a4f2e6b
CP shared KV and HiCache now keep page-aligned physical ownership while preserving valid-token radix semantics. Repeated tiny EAGLE exact hits free duplicate tail pages instead of leaking one allocator page, owner-lane load-back uses page-vector admission/eviction, and single-DP idle schedulers avoid entering an unnecessary MLP-sync collective. The commit also records the current page-aligned cache contract and adds gated decode-side EAGLE accept diagnostics so future accept-length collapses can be tied to draft KV/state transfer evidence instead of more prefill cache speculation. Constraint: CP HiCache allocator ownership is page-granular while radix matching remains valid-token based. Constraint: New diagnostics must be gated and must not alter normal EAGLE, transfer, or cache behavior. Rejected: Padding short requests to cp_size or 2*cp_size pages | wastes KV capacity and still hides valid-tail lifecycle bugs. Rejected: Adding more unconditional collectives to prove CP consistency | hot-path collectives previously caused severe performance risk. Confidence: medium Scope-risk: broad Directive: Do not reintroduce silent fallback for CP shared KV/HiCache paths; warning-level fallback or fail-fast is intentional. Tested: git diff --check Tested: local py_compile for all modified Python files Tested: remote g0034 container py_compile for modified Python/test files Tested: remote g0034 container PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py test/registered/unit/managers/test_scheduler_dp_attn_mixin.py => 114 passed, 5 warnings, 2 subtests passed Not-tested: full ETE traffic rerun after this commit Not-tested: CUDA/TAI kernel benchmark coverage for all production shapes
38 lines
1.2 KiB
Python
38 lines
1.2 KiB
Python
from unittest import TestCase
|
|
from unittest.mock import patch
|
|
|
|
from sglang.srt.managers.scheduler_dp_attn_mixin import (
|
|
MLPSyncBatchInfo,
|
|
prepare_mlp_sync_batch_raw,
|
|
)
|
|
|
|
|
|
class _FakeTPGroup:
|
|
device_group = object()
|
|
cpu_group = object()
|
|
device = "cuda"
|
|
|
|
|
|
class TestSchedulerDPAttnMixin(TestCase):
|
|
def test_single_dp_idle_batch_does_not_enter_mlp_sync_collective(self):
|
|
def fail_all_gather(self, *args, **kwargs):
|
|
raise AssertionError("idle single-DP scheduler must not all-gather")
|
|
|
|
with patch.object(MLPSyncBatchInfo, "all_gather", fail_all_gather):
|
|
result = prepare_mlp_sync_batch_raw(
|
|
local_batch=None,
|
|
dp_size=1,
|
|
attn_tp_size=1,
|
|
attn_cp_size=8,
|
|
tp_group=_FakeTPGroup(),
|
|
get_idle_batch=lambda: (_ for _ in ()).throw(
|
|
AssertionError("idle single-DP scheduler must not build idle batch")
|
|
),
|
|
disable_cuda_graph=False,
|
|
require_mlp_tp_gather=True,
|
|
disable_overlap_schedule=True,
|
|
offload_tags=set(),
|
|
)
|
|
|
|
self.assertIsNone(result)
|