Cheng Wan
b99175dc7d
Last of four; stacked on #38048. The record is the operator's input; the bags are what is in effect. A reader that takes the record and reads a field off it gets the input, which is the wrong one of the two whenever resolution decided something -- and the mistake is silent, because for most fields and most launches the two agree. Several of these files already read both ways, sometimes in the same expression: ```python get_tokenizer( get_serving().tokenizer_path, tokenizer_mode=server_args.tokenizer_mode, # the input, not the decision ... ) ``` Sixty-odd files convert. Record field reads in runtime code go from 199 to 11. Nine parameters that the conversion emptied are dropped along with the argument at every call site -- the dead-parameter ratchet is what names them. ### "Runs after its process publishes" is a per-entry-point claim Most converted reads sit in the serving and model-executor layers, which only exist after publication, or in the two subprocess entry points, which publish first thing. Three places are not like that, and they keep reading the record they were handed: - **`HttpServerEngineAdapter`** launches the server as a *child*. The parent resolves the record and never publishes, so the adapter's own reads -- the launch banner, the API key in its readiness loop, the TP width in `update_weights_from_tensor` -- are of `self.server_args`. A bag read here fails closed in a bare process, or answers for an unrelated engine in one that happens to have published. - **`serve_grpc`** reads its sidecar port before the integrated servicer builds the `Engine` that publishes. The comment above that line already said so and already bound `cfg = resolving_view(server_args)` for it; the sidecar port and the port it derives from read `cfg`. - **`initialize_dp_attention`** runs from callers whose publish is not guaranteed, so its one predicate stays on the resolution view. `ROLE_NAMESPACE_SETS["dp_controller"]` gains `observability` and `serving`, because the controller's metrics gate, tracing setup and worker-port broadcast now read those namespaces. Under `SGLANG_ROLE_NAMESPACES=enforce` that set is what the process may read, so a conversion that reaches a new namespace has to widen it in the same change. ## Three things worth a reviewer's attention **Eleven reads were `getattr(record, "field", default)`.** An AST scan for attribute access does not see those, so the census that said "43 readers" was counting the shape it could match rather than the thing it was after. `incremental_streaming_output` was read that way twice, and the transcription tests were the only reason it surfaced. **Not every record read is a bag read waiting to happen.** A multimodal processor's `base_gpu_id` is the instance's, not the process's: two engines in one process keep different ones, and `test_publishing_another_config_does_not_move_the_device` exists to say so. It stays on the record while `rl_on_policy_target` beside it moves. `RequestMetricsExporter` is the same shape -- it is handed the directory it writes to, and a test builds several with different ones. `configure_logger` is a third: 17 call sites, one of which passes an `argparse.Namespace`, so it is not a global-context reader at all. Those eleven remaining reads are the ones with a reason. **The fixtures move with the code.** Tests that hung config off a mock manager now publish a record, which is what the serving layer reads; where a test states a value it says so with `override_server_args` instead of assigning through the mock. `test_hisparse_unit` is the last of them: it stubbed a `server_args` onto a fake scheduler to say the decode radix cache was off, and the value it was standing in for is the published default, so the stub goes and the class publishes. ## Two things CI caught that a local sweep could not **`unittest.TestCase.enterContext` is Python 3.11+.** The converted fixtures used it at 18 sites; `requires-python` is `>=3.10` and CI runs 3.10, so every one of them raised `AttributeError` there while passing on a newer local interpreter. They call `enter_override(self, ...)` now -- a four-line helper in `sglang/test/test_utils.py` over the override's own `install()` / `restore()`. **A batched sweep cannot see a missing publish.** Three fixtures needed a published config and did not have one; each *passed* inside a shard where some other file had published, and failed when run alone. The affected cases are `test_serving_completions` (which set `incremental_streaming_output` on the mock manager's record, where nothing reads it now), `test_qwen3_vl_feature_materialization` (same shape for `mm_enable_dp_encoder`), and the two Qwen Rust tests -- whose fixture already carried the comment `# Non-auto: get_resolved_model_impl would choke on a SimpleNamespace` next to the `model_impl` it sets, which is exactly what happened once `get_mm_processor_cls` started reading that value from the bag. Its `publish` mirrors `model_impl` now, like the four fields it already mirrored. ## Verification A full registered-unit sweep (648 files) against this stack's merge-base: 19 failures on both sides, the same 19, none of them config. That sweep is what caught 23 failures the file-scoped runs missed -- and, later, that the narrower 139-file list did not even contain the files this change reaches. It is also what caught the `test_hisparse_unit` fixture above: the file passes inside a shard where something else published, and fails when it is run on its own, which is why every failing file is re-run alone before it is counted.
272 lines
11 KiB
Python
272 lines
11 KiB
Python
import unittest
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from sglang.srt.managers.schedule_batch import ReqKvInfo
|
|
from sglang.srt.mem_cache.allocator.hisparse import (
|
|
DeepSeekV4HiSparseTokenToKVPoolAllocator,
|
|
)
|
|
from sglang.srt.runtime_context import get_context, publish, reset_context
|
|
from sglang.srt.server_args import ServerArgs
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
from sglang.test.test_utils import CustomTestCase
|
|
|
|
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
|
|
|
|
|
|
class TestDeepSeekV4HiSparseAllocator(CustomTestCase):
|
|
def setUp(self):
|
|
# The code under test reads its config from the bags.
|
|
reset_context()
|
|
self.addCleanup(reset_context)
|
|
publish(ServerArgs(model_path="dummy"), role="tokenizer")
|
|
|
|
def test_forwards_swa_tail_allocation_to_logical_allocator(self):
|
|
allocator = object.__new__(DeepSeekV4HiSparseTokenToKVPoolAllocator)
|
|
logical_allocator = MagicMock(spec=["alloc_extend_swa_tail"])
|
|
allocator.logical_attn_allocator = logical_allocator
|
|
|
|
expected = torch.tensor([8, 9, 10], dtype=torch.int64)
|
|
logical_allocator.alloc_extend_swa_tail.return_value = expected
|
|
|
|
prefix_lens = torch.tensor([0], dtype=torch.int64)
|
|
prefix_lens_cpu = torch.tensor([0], dtype=torch.int64)
|
|
seq_lens = torch.tensor([512], dtype=torch.int64)
|
|
seq_lens_cpu = torch.tensor([512], dtype=torch.int64)
|
|
last_loc = torch.tensor([-1], dtype=torch.int64)
|
|
|
|
result = allocator.alloc_extend_swa_tail(
|
|
prefix_lens=prefix_lens,
|
|
prefix_lens_cpu=prefix_lens_cpu,
|
|
seq_lens=seq_lens,
|
|
seq_lens_cpu=seq_lens_cpu,
|
|
last_loc=last_loc,
|
|
extend_num_tokens=512,
|
|
swa_tail_len=128,
|
|
)
|
|
|
|
self.assertIs(result, expected)
|
|
logical_allocator.alloc_extend_swa_tail.assert_called_once()
|
|
_, kwargs = logical_allocator.alloc_extend_swa_tail.call_args
|
|
self.assertIs(kwargs["prefix_lens"], prefix_lens)
|
|
self.assertIs(kwargs["prefix_lens_cpu"], prefix_lens_cpu)
|
|
self.assertIs(kwargs["seq_lens"], seq_lens)
|
|
self.assertIs(kwargs["seq_lens_cpu"], seq_lens_cpu)
|
|
self.assertIs(kwargs["last_loc"], last_loc)
|
|
self.assertEqual(kwargs["extend_num_tokens"], 512)
|
|
self.assertEqual(kwargs["swa_tail_len"], 128)
|
|
|
|
def test_hisparse_budget_uses_full_logical_capacity_for_swa_tail(self):
|
|
from sglang.srt.disaggregation.decode import DecodePreallocQueue
|
|
|
|
queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
|
|
logical_allocator = SimpleNamespace(
|
|
available_size=MagicMock(return_value=32),
|
|
full_available_size=MagicMock(return_value=512),
|
|
)
|
|
queue.token_to_kv_pool_allocator = SimpleNamespace(
|
|
logical_attn_allocator=logical_allocator
|
|
)
|
|
queue.scheduler = SimpleNamespace(enable_hisparse=True, last_batch=None)
|
|
queue.retracted_queue = []
|
|
queue.num_reserved_decode_tokens = 0
|
|
queue._uses_swa_tail_prealloc = MagicMock(return_value=True)
|
|
queue._need_space_for_single_req = MagicMock(return_value=0)
|
|
queue._active_reserved_tokens = MagicMock(return_value=0)
|
|
|
|
budget = queue._allocatable_token_budgets()
|
|
|
|
self.assertEqual(budget, 512)
|
|
logical_allocator.full_available_size.assert_called_once_with()
|
|
logical_allocator.available_size.assert_not_called()
|
|
|
|
def test_hisparse_prealloc_uses_swa_tail_for_direct_host_path(self):
|
|
from sglang.srt.disaggregation.decode import DecodePreallocQueue
|
|
|
|
fill_len = 512
|
|
sliding_window_size = 200
|
|
# _swa_tail_len floors the window start to a page boundary:
|
|
# floor_align(512 - 200, 256) = 256, so the tail is one whole page.
|
|
swa_tail_len = 256
|
|
kv_loc = torch.arange(512, 512 + fill_len, dtype=torch.int64)
|
|
host_indices = torch.arange(1000, 1128, dtype=torch.int64)
|
|
|
|
req = SimpleNamespace(
|
|
rid="req-0",
|
|
origin_input_ids=list(range(fill_len)),
|
|
output_ids=[],
|
|
kv=ReqKvInfo(),
|
|
)
|
|
|
|
def set_extend_range(start, end):
|
|
req.extend_range = SimpleNamespace(start=start, end=end, length=end - start)
|
|
|
|
req.set_extend_range = set_extend_range
|
|
|
|
class ReqToTokenPool:
|
|
def __init__(self):
|
|
self.writes = []
|
|
|
|
def alloc(self, reqs):
|
|
for item in reqs:
|
|
item.kv.req_pool_idx = 0
|
|
return torch.tensor([0], dtype=torch.int64)
|
|
|
|
def write(self, indices, values):
|
|
self.writes.append((indices, values))
|
|
|
|
req_to_token_pool = ReqToTokenPool()
|
|
allocator = SimpleNamespace(
|
|
device=torch.device("cpu"),
|
|
page_size=256,
|
|
available_size=MagicMock(return_value=fill_len),
|
|
swa_available_size=MagicMock(return_value=swa_tail_len),
|
|
alloc_extend_swa_tail=MagicMock(return_value=kv_loc),
|
|
alloc_logical_only=MagicMock(return_value=kv_loc),
|
|
)
|
|
regular_host_alloc = MagicMock(return_value=host_indices)
|
|
coordinator = SimpleNamespace(
|
|
mem_pool_host=SimpleNamespace(alloc_paged_token_slots=regular_host_alloc),
|
|
req_to_host_pool=object(),
|
|
req_to_host_pool_allocated_len=object(),
|
|
host_token_len=MagicMock(side_effect=lambda token_len: token_len // 4),
|
|
)
|
|
queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
|
|
queue.req_to_token_pool = req_to_token_pool
|
|
queue.token_to_kv_pool_allocator = allocator
|
|
queue.tree_cache = SimpleNamespace(
|
|
evictable_size=MagicMock(return_value=0),
|
|
protected_size=MagicMock(return_value=0),
|
|
)
|
|
queue.scheduler = SimpleNamespace(
|
|
enable_hisparse=True,
|
|
hisparse_coordinator=coordinator,
|
|
server_args=SimpleNamespace(disaggregation_decode_enable_radix_cache=False),
|
|
sliding_window_size=sliding_window_size,
|
|
)
|
|
queue._uses_swa_tail_prealloc = MagicMock(return_value=True)
|
|
|
|
result = queue._pre_alloc(req)
|
|
|
|
self.assertTrue(torch.equal(result, host_indices))
|
|
allocator.alloc_extend_swa_tail.assert_called_once()
|
|
allocator.alloc_logical_only.assert_not_called()
|
|
_, kwargs = allocator.alloc_extend_swa_tail.call_args
|
|
self.assertEqual(kwargs["extend_num_tokens"], fill_len)
|
|
self.assertEqual(kwargs["swa_tail_len"], swa_tail_len)
|
|
self.assertEqual(req.kv.swa_evicted_seqlen, fill_len - swa_tail_len)
|
|
self.assertEqual(req.kv.kv_allocated_len, fill_len)
|
|
self.assertEqual(req.kv.kv_committed_len, fill_len)
|
|
self.assertEqual(req.extend_range.length, fill_len)
|
|
self.assertEqual(len(req_to_token_pool.writes), 1)
|
|
coordinator.host_token_len.assert_called_once_with(fill_len)
|
|
regular_host_alloc.assert_called_once_with(
|
|
coordinator.req_to_host_pool,
|
|
coordinator.req_to_host_pool_allocated_len,
|
|
req.kv.req_pool_idx,
|
|
0,
|
|
len(host_indices),
|
|
)
|
|
self.assertTrue(torch.equal(req_to_token_pool.writes[0][1], kv_loc))
|
|
|
|
# C4 indexer/C128 use the logical allocator's full-page IDs. They do not
|
|
# use either the independently allocated host pages or C4 sparse slots.
|
|
np.testing.assert_array_equal(
|
|
np.unique(kv_loc.numpy() // allocator.page_size),
|
|
np.array([2, 3]),
|
|
)
|
|
|
|
def test_mooncake_uses_separate_host_and_device_page_indices(self):
|
|
from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager
|
|
|
|
manager = object.__new__(MooncakeKVManager)
|
|
manager.is_mla_backend = True
|
|
manager.is_hybrid_mla_backend = False
|
|
manager.enable_custom_mem_pool = False
|
|
manager._transfer_data = MagicMock(return_value=0)
|
|
|
|
with ThreadPoolExecutor(max_workers=1) as executor:
|
|
ret = manager._send_kvcache_generic(
|
|
mooncake_session_id="session",
|
|
src_data_ptrs=[1000, 2000, 3000],
|
|
dst_data_ptrs=[10000, 20000, 30000],
|
|
item_lens=[100, 100, 100],
|
|
prefill_data_indices=np.array([1, 2], dtype=np.int32),
|
|
dst_data_indices=np.array([7, 8], dtype=np.int32),
|
|
executor=executor,
|
|
dst_device_data_indices=np.array([21, 22], dtype=np.int32),
|
|
dst_device_data_ptrs={20000, 30000},
|
|
)
|
|
|
|
self.assertEqual(ret, 0)
|
|
manager._transfer_data.assert_called_once_with(
|
|
"session",
|
|
[
|
|
(1100, 10700, 200),
|
|
(2100, 22100, 200),
|
|
(3100, 32100, 200),
|
|
],
|
|
)
|
|
|
|
def test_mooncake_derives_device_buffers_from_local_pp_layout(self):
|
|
from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager
|
|
|
|
manager = object.__new__(MooncakeKVManager)
|
|
manager.kv_args = SimpleNamespace(
|
|
kv_data_ptrs=[1000, 2000, 3000],
|
|
kv_item_lens=[100, 100, 100],
|
|
kv_layer_ids=[],
|
|
mla_compression_ratios=[4, 128, 4, 128],
|
|
prefill_start_layer=0,
|
|
prefill_end_layer=2,
|
|
)
|
|
manager._send_kvcache_generic = MagicMock(return_value=0)
|
|
executor = MagicMock()
|
|
|
|
# send_kvcache reads the memory bag (the unified-memory envelope-layout
|
|
# check), so the context has to be published. This is the non-unified
|
|
# path -- pin that explicitly rather than leaning on the default.
|
|
with get_context().override_server_args(enable_unified_memory=False):
|
|
manager.send_kvcache(
|
|
"session",
|
|
np.array([1], dtype=np.int32),
|
|
[10000, 20000, 30000],
|
|
np.array([7], dtype=np.int32),
|
|
executor,
|
|
dst_device_kv_indices=np.array([21], dtype=np.int32),
|
|
)
|
|
|
|
kwargs = manager._send_kvcache_generic.call_args.kwargs
|
|
self.assertEqual(kwargs["dst_device_data_ptrs"], {20000, 30000})
|
|
|
|
def test_mooncake_transfer_metadata_carries_device_page_indices(self):
|
|
from sglang.srt.disaggregation.mooncake.conn import TransferInfo
|
|
|
|
host_pages = np.array([7, 8], dtype=np.int32)
|
|
device_pages = np.array([21, 22], dtype=np.int32)
|
|
info = TransferInfo.from_zmq(
|
|
[
|
|
b"9",
|
|
b"127.0.0.1",
|
|
b"12345",
|
|
b"session",
|
|
host_pages.tobytes(),
|
|
b"0",
|
|
b"",
|
|
b"1",
|
|
b"0",
|
|
device_pages.tobytes(),
|
|
]
|
|
)
|
|
|
|
np.testing.assert_array_equal(info.dst_kv_indices, host_pages)
|
|
np.testing.assert_array_equal(info.dst_device_kv_indices, device_pages)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|