From c2fea3fa2f40077350de67e552084e6fa7c80976 Mon Sep 17 00:00:00 2001 From: James Liu Date: Tue, 22 Sep 2026 20:28:34 +0000 Subject: [PATCH 1/5] [DCP] Use logical token capacity for PD admission and load reporting --- python/sglang/srt/disaggregation/prefill.py | 2 +- python/sglang/srt/managers/scheduler.py | 5 +- python/sglang/srt/managers/tp_worker.py | 12 +- .../sglang/srt/model_executor/model_runner.py | 15 ++ .../managers/test_dcp_logical_capacity.py | 240 ++++++++++++++++++ 5 files changed, 261 insertions(+), 13 deletions(-) create mode 100644 test/registered/unit/managers/test_dcp_logical_capacity.py diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index a2623adcff56..754ed3a1fa03 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -192,7 +192,7 @@ def __init__( self.scheduler = scheduler self.scheduler_stage_metrics = scheduler_stage_metrics self.max_total_num_tokens = ( - self.scheduler.tp_worker.model_runner.effective_max_total_num_tokens + self.scheduler.tp_worker.model_runner.effective_logical_max_total_num_tokens ) self.transfer_backend = transfer_backend if envs.SGLANG_DISAGG_STAGING_BUFFER.get(): diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index ccdd14f191f0..1454532ffacc 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2360,8 +2360,7 @@ def init_pool_stats_observer(self) -> None: enable_hisparse=self.enable_hisparse, full_tokens_per_layer=self.full_tokens_per_layer, swa_tokens_per_layer=self.swa_tokens_per_layer, - max_total_num_tokens=self.max_total_num_tokens - * get_parallel().attn_dcp_size, + max_total_num_tokens=self.max_total_num_tokens, get_last_batch=lambda: self.last_batch, get_running_batch=lambda: self.running_batch, ) @@ -2546,7 +2545,7 @@ def init_req_max_new_tokens(self, req): max_new_tokens = self.token_to_kv_pool_allocator.max_new_tokens_for_memory( input_len, max_new_tokens, - token_capacity=self.max_total_num_tokens * get_parallel().attn_dcp_size, + token_capacity=self.max_total_num_tokens, sliding_window_size=self.sliding_window_size, chunk_size=self.chunked_prefill_size, ) diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index 9fa3a10993ec..a90470d2ea9d 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -449,9 +449,7 @@ def alloc_memory_pool( assert self.model_runner.max_running_requests > 0, "max_running_request is zero" max_req_len = min( self.model_config.context_len - 1, - self.model_runner.effective_max_total_num_tokens - * get_parallel().attn_dcp_size - - 1, + self.model_runner.effective_logical_max_total_num_tokens - 1, ) assert max_req_len > 0, "Memory pool size is too small" @@ -575,17 +573,13 @@ def register_hisparse_coordinator(self, coordinator): def get_worker_info(self): max_req_len = min( self.model_config.context_len - 1, - self.model_runner.effective_max_total_num_tokens - * get_parallel().attn_dcp_size - - 1, + self.model_runner.effective_logical_max_total_num_tokens - 1, ) max_req_input_len = max_req_len - 5 if self.dllm_algorithm is not None: max_req_input_len -= self.dllm_algorithm.block_size return ( - self.model_runner.req_to_token_pool.schedulable_token_capacity( - self.model_runner.max_total_num_tokens - ), + self.model_runner.logical_max_total_num_tokens, get_schedule().max_prefill_tokens, self.model_runner.max_running_requests, get_schedule().max_queued_requests, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index bd0da74046d0..fa474093fee2 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1409,6 +1409,21 @@ def unload_lora_adapter(self, lora_ref: LoRARef): """Unload a lora adapter that was previously loaded during initialization or dynamic loading.""" return self.lora_manager.unload_lora_adapter(lora_ref) + @property + def logical_max_total_num_tokens(self): + """Request-token capacity; the DCP allocator already accounts for sharding.""" + capacity = self.max_total_num_tokens + if get_parallel().attn_dcp_size > 1 and not self.is_hybrid_swa: + capacity = self.token_to_kv_pool_allocator.size + return self.req_to_token_pool.schedulable_token_capacity(capacity) + + @property + def effective_logical_max_total_num_tokens(self): + """Logical request limit, preserving hybrid SWA's separate pool bounds.""" + if get_parallel().attn_dcp_size > 1 and not self.is_hybrid_swa: + return self.logical_max_total_num_tokens + return self.effective_max_total_num_tokens + @property def effective_max_total_num_tokens(self): """Return the max token pool size considering hybrid swa settings.""" diff --git a/test/registered/unit/managers/test_dcp_logical_capacity.py b/test/registered/unit/managers/test_dcp_logical_capacity.py new file mode 100644 index 000000000000..5218ddb39eee --- /dev/null +++ b/test/registered/unit/managers/test_dcp_logical_capacity.py @@ -0,0 +1,240 @@ +"""Logical DCP capacities must agree across validation, admission and telemetry.""" + +import unittest +from types import SimpleNamespace as NS +from unittest.mock import Mock, patch + +import torch + +from sglang.srt.disaggregation.decode import DecodePreallocQueue +from sglang.srt.disaggregation.prefill import PrefillBootstrapQueue +from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.managers.tp_worker import TpModelWorker +from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator +from sglang.srt.mem_cache.memory_pool import ReqToTokenPool +from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.observability.metrics_collector import SchedulerStats +from sglang.srt.runtime_context import get_parallel +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, published_topology + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +PHYSICAL = 690240 +CONTEXT = 1048576 + + +def make_worker(dcp_size, *, allocator_size=None): + logical = PHYSICAL * dcp_size if allocator_size is None else allocator_size + kv = NS(size=PHYSICAL, mem_usage=8.89) + allocator = PagedTokenToKVPoolAllocator( + logical, 64 * dcp_size, torch.uint8, "cpu", kv, False + ) + runner = ModelRunner.__new__(ModelRunner) + runner.server_args = NS(dcp_size=dcp_size) + runner.kv_cache_configurator = NS( + hybrid_swa_token_capacity=lambda **kw: kw["full_capacity"] or kw["swa_capacity"] + ) + runner.is_hybrid_swa = False + runner.max_total_num_tokens = PHYSICAL + runner.max_running_requests = 64 + runner.token_to_kv_pool_allocator = allocator + runner.token_to_kv_pool = kv + runner.req_to_token_pool = ReqToTokenPool.__new__(ReqToTokenPool) + runner.req_to_token_pool.size = 96 + runner.req_to_token_pool.max_context_len = CONTEXT + runner.req_to_token_pool._aux_cache = None + runner.forward_stream = None + runner.weight_load_mem_usage = 0 + return NS( + model_runner=runner, + model_config=NS(context_len=CONTEXT), + server_args=NS(max_prefill_tokens=16384, max_queued_requests=None), + random_seed=0, + device="cpu", + graph_memory_usage={}, + ) + + +def make_scheduler(worker): + info = TpModelWorker.get_worker_info(worker) + runner = worker.model_runner + scheduler = NS( + tp_worker=worker, + token_to_kv_pool_allocator=runner.token_to_kv_pool_allocator, + req_to_token_pool=runner.req_to_token_pool, + tree_cache=NS(evictable_size=lambda: 0), + session_controller=None, + hisparse_coordinator=None, + is_hybrid_swa=False, + is_hybrid_ssm=False, + enable_hisparse=False, + sliding_window_size=None, + chunked_prefill_size=16384, + full_tokens_per_layer=None, + swa_tokens_per_layer=None, + max_total_num_tokens=info[0], + max_running_requests=info[2], + max_req_len=info[4], + max_req_input_len=info[5], + max_new_tokens_limit=None, + page_size=64 * runner.server_args.dcp_size, + server_args=NS(dcp_size=runner.server_args.dcp_size, enable_lora=False), + running_batch=NS(reqs=[]), + last_batch=None, + waiting_queue=[], + chunked_req=None, + disaggregation_mode=DisaggregationMode.DECODE, + disagg_decode_prealloc_queue=NS(queue=[], retracted_queue=[]), + disagg_decode_transfer_queue=NS(queue=[]), + spec_algorithm=NS(is_none=lambda: True), + metrics_reporter=NS(stats=SchedulerStats()), + ) + Scheduler.init_pool_stats_observer(scheduler) + Scheduler.init_load_inquirer(scheduler) + return scheduler + + +class TestDcpLogicalCapacity(CustomTestCase): + def make_worker(self, dcp_size, **kwargs): + self.enterContext(published_topology(tp_size=dcp_size, dcp_size=dcp_size)) + return make_worker(dcp_size, **kwargs) + + def setUp(self): + for config in ( + patch( + "sglang.srt.managers.tp_worker.get_schedule", + return_value=NS(max_prefill_tokens=16384, max_queued_requests=None), + ), + patch( + "sglang.srt.managers.scheduler_components.load_inquirer.get_lora", + return_value=NS(enable_lora=False), + ), + ): + config.start() + self.addCleanup(config.stop) + + def test_auxiliary_dense_capacity_applies_after_dcp_translation(self): + worker = self.make_worker(8) + worker.model_config.context_len = 4_000_000 + worker.model_runner.req_to_token_pool._aux_cache = NS(dense_capacity=2_000_000) + info = TpModelWorker.get_worker_info(worker) + self.assertEqual(info[0], 2_000_000) + self.assertEqual(info[4], 1_999_999) + + def test_worker_capacity_uses_effective_dcp(self): + for dcp_size in (1, 2, 8): + with self.subTest(dcp_size=dcp_size): + worker = self.make_worker(dcp_size) + # The runner's configured DCP is not the published attention group. + worker.model_runner.server_args.dcp_size = 8 + info = TpModelWorker.get_worker_info(worker) + capacity = PHYSICAL * dcp_size + self.assertEqual(info[0], capacity) + self.assertEqual(info[4], min(CONTEXT, capacity) - 1) + self.assertEqual(info[5], info[4] - 5) + + def test_allocator_is_authoritative_not_another_dcp_multiplier(self): + worker = self.make_worker(8, allocator_size=PHYSICAL * 4) + self.assertEqual(TpModelWorker.get_worker_info(worker)[0], PHYSICAL * 4) + # The size is already logical even when a draft runner's local buffer + # configuration also covers the full logical token domain. + worker.model_runner.max_total_num_tokens = PHYSICAL * 8 + self.assertEqual(TpModelWorker.get_worker_info(worker)[0], PHYSICAL * 4) + + def test_non_dcp_and_hybrid_swa_bounds(self): + runner = self.make_worker(1, allocator_size=PHYSICAL * 2).model_runner + self.assertEqual(runner.logical_max_total_num_tokens, PHYSICAL) + runner.is_hybrid_swa = True + runner.swa_max_total_num_tokens = PHYSICAL // 4 + for dcp_size, full_capacity, expected in ( + (1, PHYSICAL // 2, PHYSICAL // 2), + (1, 0, PHYSICAL // 4), + (8, PHYSICAL // 2, PHYSICAL // 2), + ): + with self.subTest(dcp_size=dcp_size, full_capacity=full_capacity): + self.enterContext(get_parallel().override(attn_dcp_size=dcp_size)) + runner.full_max_total_num_tokens = full_capacity + self.assertEqual( + runner.effective_logical_max_total_num_tokens, expected + ) + self.assertEqual(runner.logical_max_total_num_tokens, PHYSICAL) + + def test_output_budget_does_not_multiply_logical_capacity_again(self): + worker = self.make_worker(8) + worker.model_config.context_len = PHYSICAL * 16 + scheduler = make_scheduler(worker) + req = NS( + rid="near-capacity", + origin_input_ids=range(scheduler.max_total_num_tokens - 1024), + sampling_params=NS(max_new_tokens=600, min_new_tokens=0), + ) + Scheduler.init_req_max_new_tokens(scheduler, req) + self.assertEqual(req.sampling_params.max_new_tokens, 511) + + def test_one_million_token_context_is_not_clipped_to_per_rank_rows(self): + scheduler = make_scheduler(self.make_worker(8)) + self.assertEqual(scheduler.max_req_input_len, CONTEXT - 6) + req = NS( + rid="dcp-long", + origin_input_ids=range(1_000_000), + output_ids=[], + sampling_params=NS(max_new_tokens=64, min_new_tokens=0), + ) + Scheduler.init_req_max_new_tokens(scheduler, req) + self.assertEqual(req.sampling_params.max_new_tokens, 64) + self.assertFalse( + PrefillBootstrapQueue._check_if_req_exceed_kv_capacity( + NS( + max_total_num_tokens=scheduler.tp_worker.model_runner.effective_logical_max_total_num_tokens + ), + req, + ) + ) + queue = DecodePreallocQueue.__new__(DecodePreallocQueue) + queue.max_total_num_tokens = scheduler.max_total_num_tokens + queue.token_to_kv_pool_allocator = scheduler.token_to_kv_pool_allocator + queue.token_to_kv_pool = scheduler.tp_worker.model_runner.token_to_kv_pool + queue.num_reserved_decode_tokens = 0 + queue.scheduler = scheduler + scheduler.output_streamer = Mock() + self.assertFalse( + DecodePreallocQueue._check_if_req_exceed_kv_capacity(queue, req) + ) + # Context safety still applies independently of the larger pool. + req.origin_input_ids = range(CONTEXT - 32) + Scheduler.init_req_max_new_tokens(scheduler, req) + self.assertEqual(req.sampling_params.max_new_tokens, 30) + + req.origin_input_ids = range(scheduler.max_total_num_tokens + 1) + req.return_logprob = False + with patch("sglang.srt.disaggregation.decode.prepare_abort") as abort: + self.assertTrue( + DecodePreallocQueue._check_if_req_exceed_kv_capacity(queue, req) + ) + abort.assert_called_once() + queue.scheduler.output_streamer.stream_output.assert_called_once() + + def test_load_usage_uses_logical_capacity_exactly_once(self): + for dcp_size in (1, 8): + with self.subTest(dcp_size=dcp_size): + scheduler = make_scheduler(self.make_worker(dcp_size)) + allocator = scheduler.token_to_kv_pool_allocator + # Keep half the real allocator's pages available; no CUDA calls. + allocator.free_pages = allocator.free_pages[ + : len(allocator.free_pages) // 2 + ] + loads = scheduler.load_inquirer.get_loads() + self.assertEqual(loads.max_total_num_tokens, allocator.size) + self.assertEqual(loads.memory.token_capacity, allocator.size) + used = allocator.size - (allocator.num_pages // 2) * allocator.page_size + self.assertEqual(loads.num_used_tokens, used) + self.assertEqual(loads.token_usage, 0.5) + self.assertEqual( + scheduler.pool_stats_observer.max_total_num_tokens, allocator.size + ) + + +if __name__ == "__main__": + unittest.main() From fae067007e1ce68d173c49fb0c68cde0311d2d2c Mon Sep 17 00:00:00 2001 From: James Liu Date: Tue, 22 Sep 2026 20:44:20 +0000 Subject: [PATCH 2/5] Derive DCP logical capacity from configured rows in the KV configurator Under DCP with the unified Mamba pool, the allocator size also counts Mamba state bytes as tokens. Compute logical capacity as the configured per-rank rows times attn_dcp_size in KVCacheConfigurator (draft sizes are already widened; SWA allocators never widen), divide sglang:num_pages by the widened allocator page size, and use the 3.10-safe enter_scope helper in the test. Co-Authored-By: Claude Opus 5.5 --- python/sglang/srt/managers/scheduler.py | 3 +- .../srt/mem_cache/kv_cache_configurator.py | 8 ++ .../sglang/srt/model_executor/model_runner.py | 17 +-- .../managers/test_dcp_logical_capacity.py | 128 +++++++++++++----- 4 files changed, 115 insertions(+), 41 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 1454532ffacc..b5b23cade936 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1255,7 +1255,8 @@ def emit_metrics_constants(self) -> None: # TODO: max_running_requests_under_SLO has no setter — dead chain. max_running_requests_under_SLO=None, page_size=self.page_size, - num_pages=self.max_total_num_tokens // self.page_size, + num_pages=self.max_total_num_tokens + // self.token_to_kv_pool_allocator.page_size, context_len=self.model_config.context_len, startup_available_gpu_memory_gb=self.startup_available_gpu_memory_gb, ) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 2157a39fccdb..3ab13adb1adc 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -322,6 +322,14 @@ def hybrid_swa_token_capacity( ) return full_capacity or swa_capacity + def logical_token_capacity(self, *, max_total_num_tokens: int) -> int: + """Request tokens a pool of `max_total_num_tokens` per-rank rows holds.""" + # The paged target allocator widens each row into attn_dcp_size ids; draft + # sizes are already widened (loc_space_scale) and SWA allocators never widen. + if self.is_hybrid_swa or self.is_draft_worker: + return max_total_num_tokens + return max_total_num_tokens * get_parallel().attn_dcp_size + def _build_fp4_quant_method(self, *, num_layers: int): if not is_float4_e2m1fn_x2(self.kv_cache_dtype): return None diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index fa474093fee2..54acaa1c2131 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1411,18 +1411,19 @@ def unload_lora_adapter(self, lora_ref: LoRARef): @property def logical_max_total_num_tokens(self): - """Request-token capacity; the DCP allocator already accounts for sharding.""" - capacity = self.max_total_num_tokens - if get_parallel().attn_dcp_size > 1 and not self.is_hybrid_swa: - capacity = self.token_to_kv_pool_allocator.size - return self.req_to_token_pool.schedulable_token_capacity(capacity) + """Request-token capacity in logical tokens, not per-rank DCP rows.""" + return self.req_to_token_pool.schedulable_token_capacity( + self.kv_cache_configurator.logical_token_capacity( + max_total_num_tokens=self.max_total_num_tokens + ) + ) @property def effective_logical_max_total_num_tokens(self): """Logical request limit, preserving hybrid SWA's separate pool bounds.""" - if get_parallel().attn_dcp_size > 1 and not self.is_hybrid_swa: - return self.logical_max_total_num_tokens - return self.effective_max_total_num_tokens + if self.is_hybrid_swa: + return self.effective_max_total_num_tokens + return self.logical_max_total_num_tokens @property def effective_max_total_num_tokens(self): diff --git a/test/registered/unit/managers/test_dcp_logical_capacity.py b/test/registered/unit/managers/test_dcp_logical_capacity.py index 5218ddb39eee..a0a129e391ce 100644 --- a/test/registered/unit/managers/test_dcp_logical_capacity.py +++ b/test/registered/unit/managers/test_dcp_logical_capacity.py @@ -12,12 +12,21 @@ from sglang.srt.managers.scheduler import Scheduler from sglang.srt.managers.tp_worker import TpModelWorker from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator +from sglang.srt.mem_cache.allocator.unified_mamba import ( + UnifiedMambaTokenToKVPoolAllocator, +) +from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator from sglang.srt.mem_cache.memory_pool import ReqToTokenPool +from sglang.srt.mem_cache.unified_memory_pool import ( + MambaSubPoolSpec, + MLASubPoolSpec, + UnifiedKVPool, +) from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.observability.metrics_collector import SchedulerStats from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase, published_topology +from sglang.test.test_utils import CustomTestCase, enter_scope, published_topology register_cpu_ci(est_time=5, suite="base-a-test-cpu") @@ -25,17 +34,56 @@ CONTEXT = 1048576 -def make_worker(dcp_size, *, allocator_size=None): - logical = PHYSICAL * dcp_size if allocator_size is None else allocator_size +def make_configurator(*, is_hybrid_swa=False, is_draft_worker=False): + configurator = KVCacheConfigurator.__new__(KVCacheConfigurator) + configurator.is_hybrid_swa = is_hybrid_swa + configurator.is_draft_worker = is_draft_worker + return configurator + + +def make_unified_mamba_allocator(*, n_full_tokens, n_mamba_slots): + full = MLASubPoolSpec( + name="full", + layer_num=3, + kv_lora_rank=6, + qk_rope_head_dim=2, + store_dtype=torch.bfloat16, + grow_direction="down", + ) + mamba = MambaSubPoolSpec( + name="mamba", + layer_num=2, + conv_state_shapes=((4, 3),), + conv_dtype=torch.float32, + temporal_state_shape=(2, 2, 2), + temporal_dtype=torch.float32, + grow_direction="up", + ) + pool = UnifiedKVPool( + total_bytes=full.entry_bytes() * n_full_tokens + + mamba.entry_bytes() * n_mamba_slots, + sub_pool_specs=[full, mamba], + device="cpu", + enable_memory_saver=False, + page_size=1, + ) + kvcache = NS( + full_kv_pool=NS(buf=torch.empty(pool.max_slots("full"))), + mamba_pool=NS(buf=torch.empty(pool.max_slots("mamba"))), + ) + return UnifiedMambaTokenToKVPoolAllocator( + unified_buffer=pool, kvcache=kvcache, device="cpu", page_size=1 + ) + + +def make_worker(dcp_size): kv = NS(size=PHYSICAL, mem_usage=8.89) allocator = PagedTokenToKVPoolAllocator( - logical, 64 * dcp_size, torch.uint8, "cpu", kv, False + PHYSICAL * dcp_size, 64 * dcp_size, torch.uint8, "cpu", kv, False ) runner = ModelRunner.__new__(ModelRunner) runner.server_args = NS(dcp_size=dcp_size) - runner.kv_cache_configurator = NS( - hybrid_swa_token_capacity=lambda **kw: kw["full_capacity"] or kw["swa_capacity"] - ) + runner.kv_cache_configurator = make_configurator() runner.is_hybrid_swa = False runner.max_total_num_tokens = PHYSICAL runner.max_running_requests = 64 @@ -79,7 +127,7 @@ def make_scheduler(worker): max_req_len=info[4], max_req_input_len=info[5], max_new_tokens_limit=None, - page_size=64 * runner.server_args.dcp_size, + page_size=64, server_args=NS(dcp_size=runner.server_args.dcp_size, enable_lora=False), running_batch=NS(reqs=[]), last_batch=None, @@ -90,6 +138,10 @@ def make_scheduler(worker): disagg_decode_transfer_queue=NS(queue=[]), spec_algorithm=NS(is_none=lambda: True), metrics_reporter=NS(stats=SchedulerStats()), + metrics_collector=Mock(), + draft_worker=None, + startup_available_gpu_memory_gb=None, + model_config=worker.model_config, ) Scheduler.init_pool_stats_observer(scheduler) Scheduler.init_load_inquirer(scheduler) @@ -97,9 +149,14 @@ def make_scheduler(worker): class TestDcpLogicalCapacity(CustomTestCase): - def make_worker(self, dcp_size, **kwargs): - self.enterContext(published_topology(tp_size=dcp_size, dcp_size=dcp_size)) - return make_worker(dcp_size, **kwargs) + def make_worker(self, dcp_size, **server_args_fields): + enter_scope( + self, + published_topology( + tp_size=dcp_size, dcp_size=dcp_size, **server_args_fields + ), + ) + return make_worker(dcp_size) def setUp(self): for config in ( @@ -115,14 +172,6 @@ def setUp(self): config.start() self.addCleanup(config.stop) - def test_auxiliary_dense_capacity_applies_after_dcp_translation(self): - worker = self.make_worker(8) - worker.model_config.context_len = 4_000_000 - worker.model_runner.req_to_token_pool._aux_cache = NS(dense_capacity=2_000_000) - info = TpModelWorker.get_worker_info(worker) - self.assertEqual(info[0], 2_000_000) - self.assertEqual(info[4], 1_999_999) - def test_worker_capacity_uses_effective_dcp(self): for dcp_size in (1, 2, 8): with self.subTest(dcp_size=dcp_size): @@ -135,18 +184,28 @@ def test_worker_capacity_uses_effective_dcp(self): self.assertEqual(info[4], min(CONTEXT, capacity) - 1) self.assertEqual(info[5], info[4] - 5) - def test_allocator_is_authoritative_not_another_dcp_multiplier(self): - worker = self.make_worker(8, allocator_size=PHYSICAL * 4) - self.assertEqual(TpModelWorker.get_worker_info(worker)[0], PHYSICAL * 4) - # The size is already logical even when a draft runner's local buffer - # configuration also covers the full logical token domain. - worker.model_runner.max_total_num_tokens = PHYSICAL * 8 - self.assertEqual(TpModelWorker.get_worker_info(worker)[0], PHYSICAL * 4) - - def test_non_dcp_and_hybrid_swa_bounds(self): - runner = self.make_worker(1, allocator_size=PHYSICAL * 2).model_runner - self.assertEqual(runner.logical_max_total_num_tokens, PHYSICAL) + def test_capacity_is_rows_times_dcp_not_allocator_size(self): + """The unified Mamba allocator's size also counts Mamba state bytes as + tokens, so capacity must come from the configured rows, not the allocator.""" + runner = self.make_worker(8).model_runner + runner.max_total_num_tokens = 64 + runner.token_to_kv_pool_allocator = make_unified_mamba_allocator( + n_full_tokens=64, n_mamba_slots=8 + ) + self.assertGreater(runner.token_to_kv_pool_allocator.size, 64 * 8) + self.assertEqual(runner.logical_max_total_num_tokens, 64 * 8) + # Draft sizes are already widened by loc_space_scale; SWA never widens. + for configurator in ( + make_configurator(is_draft_worker=True), + make_configurator(is_hybrid_swa=True), + ): + runner.kv_cache_configurator = configurator + self.assertEqual(runner.logical_max_total_num_tokens, 64) + + def test_hybrid_swa_bounds_do_not_widen_under_dcp(self): + runner = self.make_worker(1).model_runner runner.is_hybrid_swa = True + runner.kv_cache_configurator = make_configurator(is_hybrid_swa=True) runner.swa_max_total_num_tokens = PHYSICAL // 4 for dcp_size, full_capacity, expected in ( (1, PHYSICAL // 2, PHYSICAL // 2), @@ -154,7 +213,7 @@ def test_non_dcp_and_hybrid_swa_bounds(self): (8, PHYSICAL // 2, PHYSICAL // 2), ): with self.subTest(dcp_size=dcp_size, full_capacity=full_capacity): - self.enterContext(get_parallel().override(attn_dcp_size=dcp_size)) + enter_scope(self, get_parallel().override(attn_dcp_size=dcp_size)) runner.full_max_total_num_tokens = full_capacity self.assertEqual( runner.effective_logical_max_total_num_tokens, expected @@ -219,7 +278,9 @@ def test_one_million_token_context_is_not_clipped_to_per_rank_rows(self): def test_load_usage_uses_logical_capacity_exactly_once(self): for dcp_size in (1, 8): with self.subTest(dcp_size=dcp_size): - scheduler = make_scheduler(self.make_worker(dcp_size)) + scheduler = make_scheduler( + self.make_worker(dcp_size, enable_metrics=True) + ) allocator = scheduler.token_to_kv_pool_allocator # Keep half the real allocator's pages available; no CUDA calls. allocator.free_pages = allocator.free_pages[ @@ -234,6 +295,9 @@ def test_load_usage_uses_logical_capacity_exactly_once(self): self.assertEqual( scheduler.pool_stats_observer.max_total_num_tokens, allocator.size ) + Scheduler.emit_metrics_constants(scheduler) + constants = scheduler.metrics_collector.emit_constants.call_args + self.assertEqual(constants.kwargs["num_pages"], PHYSICAL // 64) if __name__ == "__main__": From b042f16977c1e54c6907596d9ae88ef50800d432 Mon Sep 17 00:00:00 2001 From: James Liu Date: Wed, 23 Sep 2026 21:53:05 +0000 Subject: [PATCH 3/5] Give the DCP capacity test's fake worker a dllm_algorithm get_worker_info now reads self.dllm_algorithm (#34061). Set it to None, as _init_dllm_algorithm does for non-dLLM models. Co-Authored-By: Claude Opus 5.5 --- test/registered/unit/managers/test_dcp_logical_capacity.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test/registered/unit/managers/test_dcp_logical_capacity.py b/test/registered/unit/managers/test_dcp_logical_capacity.py index a0a129e391ce..2bc73484ebb6 100644 --- a/test/registered/unit/managers/test_dcp_logical_capacity.py +++ b/test/registered/unit/managers/test_dcp_logical_capacity.py @@ -102,6 +102,7 @@ def make_worker(dcp_size): random_seed=0, device="cpu", graph_memory_usage={}, + dllm_algorithm=None, ) From 1fadd81ae5707b2d2327c8ebc887f547973119be Mon Sep 17 00:00:00 2001 From: James Liu Date: Fri, 25 Sep 2026 07:25:00 +0000 Subject: [PATCH 4/5] Reuse loc_space_scale for DCP logical capacity; trim unit test Co-Authored-By: Claude Opus 5.5 --- .../srt/mem_cache/kv_cache_configurator.py | 11 +- .../managers/test_dcp_logical_capacity.py | 123 ++++++------------ 2 files changed, 47 insertions(+), 87 deletions(-) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 7366f073b972..13210daa44d3 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -318,11 +318,14 @@ def hybrid_swa_token_capacity( def logical_token_capacity(self, *, max_total_num_tokens: int) -> int: """Request tokens a pool of `max_total_num_tokens` per-rank rows holds.""" - # The paged target allocator widens each row into attn_dcp_size ids; draft - # sizes are already widened (loc_space_scale) and SWA allocators never widen. - if self.is_hybrid_swa or self.is_draft_worker: + # SWA allocators never widen under DCP. + if self.is_hybrid_swa: return max_total_num_tokens - return max_total_num_tokens * get_parallel().attn_dcp_size + # Target rows widen into attn_dcp_size ids; draft sizes already carry + # loc_space_scale. + return ( + max_total_num_tokens * get_parallel().attn_dcp_size // self.loc_space_scale + ) def _build_fp4_quant_method(self, *, num_layers: int): if not is_float4_e2m1fn_x2(self.kv_cache_dtype): diff --git a/test/registered/unit/managers/test_dcp_logical_capacity.py b/test/registered/unit/managers/test_dcp_logical_capacity.py index 2bc73484ebb6..871ff9c0544d 100644 --- a/test/registered/unit/managers/test_dcp_logical_capacity.py +++ b/test/registered/unit/managers/test_dcp_logical_capacity.py @@ -24,7 +24,6 @@ ) from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.observability.metrics_collector import SchedulerStats -from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase, enter_scope, published_topology @@ -131,9 +130,7 @@ def make_scheduler(worker): page_size=64, server_args=NS(dcp_size=runner.server_args.dcp_size, enable_lora=False), running_batch=NS(reqs=[]), - last_batch=None, waiting_queue=[], - chunked_req=None, disaggregation_mode=DisaggregationMode.DECODE, disagg_decode_prealloc_queue=NS(queue=[], retracted_queue=[]), disagg_decode_transfer_queue=NS(queue=[]), @@ -173,21 +170,15 @@ def setUp(self): config.start() self.addCleanup(config.stop) - def test_worker_capacity_uses_effective_dcp(self): + def test_logical_capacity(self): for dcp_size in (1, 2, 8): with self.subTest(dcp_size=dcp_size): - worker = self.make_worker(dcp_size) - # The runner's configured DCP is not the published attention group. - worker.model_runner.server_args.dcp_size = 8 - info = TpModelWorker.get_worker_info(worker) - capacity = PHYSICAL * dcp_size - self.assertEqual(info[0], capacity) - self.assertEqual(info[4], min(CONTEXT, capacity) - 1) - self.assertEqual(info[5], info[4] - 5) + info = TpModelWorker.get_worker_info(self.make_worker(dcp_size)) + self.assertEqual(info[0], PHYSICAL * dcp_size) + self.assertEqual(info[4], min(CONTEXT, PHYSICAL * dcp_size) - 1) - def test_capacity_is_rows_times_dcp_not_allocator_size(self): - """The unified Mamba allocator's size also counts Mamba state bytes as - tokens, so capacity must come from the configured rows, not the allocator.""" + # Unified Mamba's allocator.size also counts Mamba state bytes, so + # capacity must come from the configured rows. runner = self.make_worker(8).model_runner runner.max_total_num_tokens = 64 runner.token_to_kv_pool_allocator = make_unified_mamba_allocator( @@ -195,61 +186,32 @@ def test_capacity_is_rows_times_dcp_not_allocator_size(self): ) self.assertGreater(runner.token_to_kv_pool_allocator.size, 64 * 8) self.assertEqual(runner.logical_max_total_num_tokens, 64 * 8) - # Draft sizes are already widened by loc_space_scale; SWA never widens. - for configurator in ( - make_configurator(is_draft_worker=True), - make_configurator(is_hybrid_swa=True), - ): - runner.kv_cache_configurator = configurator - self.assertEqual(runner.logical_max_total_num_tokens, 64) - def test_hybrid_swa_bounds_do_not_widen_under_dcp(self): - runner = self.make_worker(1).model_runner + # Draft sizes already carry loc_space_scale; SWA never widens. + runner.kv_cache_configurator = make_configurator(is_draft_worker=True) + self.assertEqual(runner.logical_max_total_num_tokens, 64) runner.is_hybrid_swa = True runner.kv_cache_configurator = make_configurator(is_hybrid_swa=True) - runner.swa_max_total_num_tokens = PHYSICAL // 4 - for dcp_size, full_capacity, expected in ( - (1, PHYSICAL // 2, PHYSICAL // 2), - (1, 0, PHYSICAL // 4), - (8, PHYSICAL // 2, PHYSICAL // 2), - ): - with self.subTest(dcp_size=dcp_size, full_capacity=full_capacity): - enter_scope(self, get_parallel().override(attn_dcp_size=dcp_size)) - runner.full_max_total_num_tokens = full_capacity - self.assertEqual( - runner.effective_logical_max_total_num_tokens, expected - ) - self.assertEqual(runner.logical_max_total_num_tokens, PHYSICAL) - - def test_output_budget_does_not_multiply_logical_capacity_again(self): - worker = self.make_worker(8) - worker.model_config.context_len = PHYSICAL * 16 - scheduler = make_scheduler(worker) - req = NS( - rid="near-capacity", - origin_input_ids=range(scheduler.max_total_num_tokens - 1024), - sampling_params=NS(max_new_tokens=600, min_new_tokens=0), - ) - Scheduler.init_req_max_new_tokens(scheduler, req) - self.assertEqual(req.sampling_params.max_new_tokens, 511) + runner.full_max_total_num_tokens = 32 + runner.swa_max_total_num_tokens = 16 + self.assertEqual(runner.logical_max_total_num_tokens, 64) + self.assertEqual(runner.effective_logical_max_total_num_tokens, 32) - def test_one_million_token_context_is_not_clipped_to_per_rank_rows(self): + def test_one_million_token_request_is_admitted_under_dcp8(self): scheduler = make_scheduler(self.make_worker(8)) self.assertEqual(scheduler.max_req_input_len, CONTEXT - 6) req = NS( rid="dcp-long", origin_input_ids=range(1_000_000), output_ids=[], + return_logprob=False, sampling_params=NS(max_new_tokens=64, min_new_tokens=0), ) Scheduler.init_req_max_new_tokens(scheduler, req) self.assertEqual(req.sampling_params.max_new_tokens, 64) self.assertFalse( PrefillBootstrapQueue._check_if_req_exceed_kv_capacity( - NS( - max_total_num_tokens=scheduler.tp_worker.model_runner.effective_logical_max_total_num_tokens - ), - req, + NS(max_total_num_tokens=scheduler.max_total_num_tokens), req ) ) queue = DecodePreallocQueue.__new__(DecodePreallocQueue) @@ -262,43 +224,38 @@ def test_one_million_token_context_is_not_clipped_to_per_rank_rows(self): self.assertFalse( DecodePreallocQueue._check_if_req_exceed_kv_capacity(queue, req) ) - # Context safety still applies independently of the larger pool. - req.origin_input_ids = range(CONTEXT - 32) - Scheduler.init_req_max_new_tokens(scheduler, req) - self.assertEqual(req.sampling_params.max_new_tokens, 30) req.origin_input_ids = range(scheduler.max_total_num_tokens + 1) - req.return_logprob = False with patch("sglang.srt.disaggregation.decode.prepare_abort") as abort: self.assertTrue( DecodePreallocQueue._check_if_req_exceed_kv_capacity(queue, req) ) abort.assert_called_once() - queue.scheduler.output_streamer.stream_output.assert_called_once() - def test_load_usage_uses_logical_capacity_exactly_once(self): - for dcp_size in (1, 8): - with self.subTest(dcp_size=dcp_size): - scheduler = make_scheduler( - self.make_worker(dcp_size, enable_metrics=True) - ) - allocator = scheduler.token_to_kv_pool_allocator - # Keep half the real allocator's pages available; no CUDA calls. - allocator.free_pages = allocator.free_pages[ - : len(allocator.free_pages) // 2 - ] - loads = scheduler.load_inquirer.get_loads() - self.assertEqual(loads.max_total_num_tokens, allocator.size) - self.assertEqual(loads.memory.token_capacity, allocator.size) - used = allocator.size - (allocator.num_pages // 2) * allocator.page_size - self.assertEqual(loads.num_used_tokens, used) - self.assertEqual(loads.token_usage, 0.5) - self.assertEqual( - scheduler.pool_stats_observer.max_total_num_tokens, allocator.size - ) - Scheduler.emit_metrics_constants(scheduler) - constants = scheduler.metrics_collector.emit_constants.call_args - self.assertEqual(constants.kwargs["num_pages"], PHYSICAL // 64) + # Output budget counts the logical capacity once. + worker = self.make_worker(8) + worker.model_config.context_len = PHYSICAL * 16 + scheduler = make_scheduler(worker) + req.origin_input_ids = range(scheduler.max_total_num_tokens - 1024) + req.sampling_params.max_new_tokens = 600 + Scheduler.init_req_max_new_tokens(scheduler, req) + self.assertEqual(req.sampling_params.max_new_tokens, 511) + + def test_reporting_counts_dcp_once(self): + scheduler = make_scheduler(self.make_worker(8, enable_metrics=True)) + allocator = scheduler.token_to_kv_pool_allocator + # Keep half the real allocator's pages available; no CUDA calls. + allocator.free_pages = allocator.free_pages[: len(allocator.free_pages) // 2] + loads = scheduler.load_inquirer.get_loads() + self.assertEqual(loads.max_total_num_tokens, allocator.size) + self.assertEqual(loads.memory.token_capacity, allocator.size) + self.assertEqual(loads.token_usage, 0.5) + self.assertEqual( + scheduler.pool_stats_observer.max_total_num_tokens, allocator.size + ) + Scheduler.emit_metrics_constants(scheduler) + constants = scheduler.metrics_collector.emit_constants.call_args + self.assertEqual(constants.kwargs["num_pages"], PHYSICAL // 64) if __name__ == "__main__": From 9e3f8efdd980652f4a1d22aaadd083766f9bae2d Mon Sep 17 00:00:00 2001 From: James Liu Date: Fri, 25 Sep 2026 07:31:14 +0000 Subject: [PATCH 5/5] Trim DCP capacity test to the capacity rules Co-Authored-By: Claude Opus 5.5 --- .../managers/test_dcp_logical_capacity.py | 129 +----------------- 1 file changed, 3 insertions(+), 126 deletions(-) diff --git a/test/registered/unit/managers/test_dcp_logical_capacity.py b/test/registered/unit/managers/test_dcp_logical_capacity.py index 871ff9c0544d..3d7a92322dae 100644 --- a/test/registered/unit/managers/test_dcp_logical_capacity.py +++ b/test/registered/unit/managers/test_dcp_logical_capacity.py @@ -2,16 +2,11 @@ import unittest from types import SimpleNamespace as NS -from unittest.mock import Mock, patch +from unittest.mock import patch import torch -from sglang.srt.disaggregation.decode import DecodePreallocQueue -from sglang.srt.disaggregation.prefill import PrefillBootstrapQueue -from sglang.srt.disaggregation.utils import DisaggregationMode -from sglang.srt.managers.scheduler import Scheduler from sglang.srt.managers.tp_worker import TpModelWorker -from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.unified_mamba import ( UnifiedMambaTokenToKVPoolAllocator, ) @@ -23,7 +18,6 @@ UnifiedKVPool, ) from sglang.srt.model_executor.model_runner import ModelRunner -from sglang.srt.observability.metrics_collector import SchedulerStats from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase, enter_scope, published_topology @@ -77,83 +71,30 @@ def make_unified_mamba_allocator(*, n_full_tokens, n_mamba_slots): def make_worker(dcp_size): kv = NS(size=PHYSICAL, mem_usage=8.89) - allocator = PagedTokenToKVPoolAllocator( - PHYSICAL * dcp_size, 64 * dcp_size, torch.uint8, "cpu", kv, False - ) runner = ModelRunner.__new__(ModelRunner) runner.server_args = NS(dcp_size=dcp_size) runner.kv_cache_configurator = make_configurator() runner.is_hybrid_swa = False runner.max_total_num_tokens = PHYSICAL runner.max_running_requests = 64 - runner.token_to_kv_pool_allocator = allocator runner.token_to_kv_pool = kv runner.req_to_token_pool = ReqToTokenPool.__new__(ReqToTokenPool) runner.req_to_token_pool.size = 96 runner.req_to_token_pool.max_context_len = CONTEXT runner.req_to_token_pool._aux_cache = None runner.forward_stream = None - runner.weight_load_mem_usage = 0 return NS( model_runner=runner, model_config=NS(context_len=CONTEXT), - server_args=NS(max_prefill_tokens=16384, max_queued_requests=None), random_seed=0, device="cpu", - graph_memory_usage={}, dllm_algorithm=None, ) -def make_scheduler(worker): - info = TpModelWorker.get_worker_info(worker) - runner = worker.model_runner - scheduler = NS( - tp_worker=worker, - token_to_kv_pool_allocator=runner.token_to_kv_pool_allocator, - req_to_token_pool=runner.req_to_token_pool, - tree_cache=NS(evictable_size=lambda: 0), - session_controller=None, - hisparse_coordinator=None, - is_hybrid_swa=False, - is_hybrid_ssm=False, - enable_hisparse=False, - sliding_window_size=None, - chunked_prefill_size=16384, - full_tokens_per_layer=None, - swa_tokens_per_layer=None, - max_total_num_tokens=info[0], - max_running_requests=info[2], - max_req_len=info[4], - max_req_input_len=info[5], - max_new_tokens_limit=None, - page_size=64, - server_args=NS(dcp_size=runner.server_args.dcp_size, enable_lora=False), - running_batch=NS(reqs=[]), - waiting_queue=[], - disaggregation_mode=DisaggregationMode.DECODE, - disagg_decode_prealloc_queue=NS(queue=[], retracted_queue=[]), - disagg_decode_transfer_queue=NS(queue=[]), - spec_algorithm=NS(is_none=lambda: True), - metrics_reporter=NS(stats=SchedulerStats()), - metrics_collector=Mock(), - draft_worker=None, - startup_available_gpu_memory_gb=None, - model_config=worker.model_config, - ) - Scheduler.init_pool_stats_observer(scheduler) - Scheduler.init_load_inquirer(scheduler) - return scheduler - - class TestDcpLogicalCapacity(CustomTestCase): - def make_worker(self, dcp_size, **server_args_fields): - enter_scope( - self, - published_topology( - tp_size=dcp_size, dcp_size=dcp_size, **server_args_fields - ), - ) + def make_worker(self, dcp_size): + enter_scope(self, published_topology(tp_size=dcp_size, dcp_size=dcp_size)) return make_worker(dcp_size) def setUp(self): @@ -162,10 +103,6 @@ def setUp(self): "sglang.srt.managers.tp_worker.get_schedule", return_value=NS(max_prefill_tokens=16384, max_queued_requests=None), ), - patch( - "sglang.srt.managers.scheduler_components.load_inquirer.get_lora", - return_value=NS(enable_lora=False), - ), ): config.start() self.addCleanup(config.stop) @@ -197,66 +134,6 @@ def test_logical_capacity(self): self.assertEqual(runner.logical_max_total_num_tokens, 64) self.assertEqual(runner.effective_logical_max_total_num_tokens, 32) - def test_one_million_token_request_is_admitted_under_dcp8(self): - scheduler = make_scheduler(self.make_worker(8)) - self.assertEqual(scheduler.max_req_input_len, CONTEXT - 6) - req = NS( - rid="dcp-long", - origin_input_ids=range(1_000_000), - output_ids=[], - return_logprob=False, - sampling_params=NS(max_new_tokens=64, min_new_tokens=0), - ) - Scheduler.init_req_max_new_tokens(scheduler, req) - self.assertEqual(req.sampling_params.max_new_tokens, 64) - self.assertFalse( - PrefillBootstrapQueue._check_if_req_exceed_kv_capacity( - NS(max_total_num_tokens=scheduler.max_total_num_tokens), req - ) - ) - queue = DecodePreallocQueue.__new__(DecodePreallocQueue) - queue.max_total_num_tokens = scheduler.max_total_num_tokens - queue.token_to_kv_pool_allocator = scheduler.token_to_kv_pool_allocator - queue.token_to_kv_pool = scheduler.tp_worker.model_runner.token_to_kv_pool - queue.num_reserved_decode_tokens = 0 - queue.scheduler = scheduler - scheduler.output_streamer = Mock() - self.assertFalse( - DecodePreallocQueue._check_if_req_exceed_kv_capacity(queue, req) - ) - - req.origin_input_ids = range(scheduler.max_total_num_tokens + 1) - with patch("sglang.srt.disaggregation.decode.prepare_abort") as abort: - self.assertTrue( - DecodePreallocQueue._check_if_req_exceed_kv_capacity(queue, req) - ) - abort.assert_called_once() - - # Output budget counts the logical capacity once. - worker = self.make_worker(8) - worker.model_config.context_len = PHYSICAL * 16 - scheduler = make_scheduler(worker) - req.origin_input_ids = range(scheduler.max_total_num_tokens - 1024) - req.sampling_params.max_new_tokens = 600 - Scheduler.init_req_max_new_tokens(scheduler, req) - self.assertEqual(req.sampling_params.max_new_tokens, 511) - - def test_reporting_counts_dcp_once(self): - scheduler = make_scheduler(self.make_worker(8, enable_metrics=True)) - allocator = scheduler.token_to_kv_pool_allocator - # Keep half the real allocator's pages available; no CUDA calls. - allocator.free_pages = allocator.free_pages[: len(allocator.free_pages) // 2] - loads = scheduler.load_inquirer.get_loads() - self.assertEqual(loads.max_total_num_tokens, allocator.size) - self.assertEqual(loads.memory.token_capacity, allocator.size) - self.assertEqual(loads.token_usage, 0.5) - self.assertEqual( - scheduler.pool_stats_observer.max_total_num_tokens, allocator.size - ) - Scheduler.emit_metrics_constants(scheduler) - constants = scheduler.metrics_collector.emit_constants.call_args - self.assertEqual(constants.kwargs["num_pages"], PHYSICAL // 64) - if __name__ == "__main__": unittest.main()