From 681ad0c371575664833aa49f8a9e26021ad64cdf Mon Sep 17 00:00:00 2001 From: Caihua Li Date: Sat, 12 Sep 2026 01:03:12 -0700 Subject: [PATCH 01/15] refactor(unified-memory): add the token-major entry layout types Introduce the description of a token-major slot entry, alongside the existing page-major builders and used by nobody yet: DensePart one row family inside the entry (K, V, or the MLA latent): layer_num rows of row_shape, layer l at offset_bytes + l * layer_stride_bytes. DenseEntryLayout the entry's byte map -- its parts plus entry_bytes, which is the slot stride of every view built over it. validate() rejects overlapping parts, an unaligned offset/stride/row, and any part reaching outside the entry. build_dense_views per-layer views of one part: a strided (num_pages * page_size, *row_shape) tensor indexed by the PHYSICAL token id, with the whole entry as the slot stride. Two alignments come with it: ENTRY_ALIGN_BYTES pads a slot so entries stay vector-store aligned, ROW_ALIGN_BYTES is what validate() holds each part to. Parts may differ in row width -- only their offsets differ, never the stride -- which is what lets one entry hold K and V rows of different widths. The page-major builders cannot: they lay out one uniform row-block array per page and address it by a re-numbered kernel-facing id. build_page_major_mamba_views already had a nested contiguous_strides identical to the one build_dense_views needs, so it is hoisted to module scope as _contiguous_strides rather than copied. --- .../sglang/srt/mem_cache/layout/page_major.py | 133 ++++++++++++++++-- 1 file changed, 123 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/mem_cache/layout/page_major.py b/python/sglang/srt/mem_cache/layout/page_major.py index ff81583181ab..129b5c2fa4ec 100644 --- a/python/sglang/srt/mem_cache/layout/page_major.py +++ b/python/sglang/srt/mem_cache/layout/page_major.py @@ -17,8 +17,12 @@ from typing import List, Sequence, Tuple +import msgspec import torch +ENTRY_ALIGN_BYTES = 32 +ROW_ALIGN_BYTES = 16 + def _prod(shape: Sequence[int]) -> int: out = 1 @@ -27,6 +31,123 @@ def _prod(shape: Sequence[int]) -> int: return out +def _contiguous_strides(shape: Sequence[int]) -> Tuple[int, ...]: + strides = [] + acc = 1 + for s in reversed(shape): + strides.append(acc) + acc *= int(s) + return tuple(reversed(strides)) + + +def align_entry_bytes(num_bytes: int) -> int: + """Round a slot's byte sum up to the entry alignment.""" + return -(-num_bytes // ENTRY_ALIGN_BYTES) * ENTRY_ALIGN_BYTES + + +class DensePart(msgspec.Struct, frozen=True, kw_only=True): + """One row family inside the entry: ``layer_num`` rows of ``row_shape``, + layer ``l`` at ``offset_bytes + l * layer_stride_bytes``.""" + + name: str + offset_bytes: int + layer_stride_bytes: int + layer_num: int + row_shape: Tuple[int, ...] + dtype: torch.dtype + + def row_bytes(self) -> int: + return _prod(self.row_shape) * self.dtype.itemsize + + def layer_offset_bytes(self, layer: int) -> int: + return self.offset_bytes + layer * self.layer_stride_bytes + + +class DenseEntryLayout(msgspec.Struct, frozen=True, kw_only=True): + """Byte layout of one slot's entry; ``entry_bytes`` is the slot stride of + every view built over it.""" + + entry_bytes: int + parts: Tuple[DensePart, ...] + + def __post_init__(self): + self.validate() + + def part(self, name: str) -> DensePart: + for p in self.parts: + if p.name == name: + return p + raise KeyError(f"no part {name!r} in {[p.name for p in self.parts]}") + + def validate(self) -> None: + assert self.entry_bytes % ENTRY_ALIGN_BYTES == 0, ( + f"entry_bytes={self.entry_bytes} is not a multiple of {ENTRY_ALIGN_BYTES}" + ) + spans = [] + for p in self.parts: + row = p.row_bytes() + assert ( + p.offset_bytes % ROW_ALIGN_BYTES == 0 + and p.layer_stride_bytes % ROW_ALIGN_BYTES == 0 + and row % ROW_ALIGN_BYTES == 0 + ), ( + f"part {p.name!r}: offset {p.offset_bytes}, layer stride " + f"{p.layer_stride_bytes} and row {row} B must all be multiples of " + f"{ROW_ALIGN_BYTES} (vector stores)" + ) + for l in range(p.layer_num): + lo = p.layer_offset_bytes(l) + assert 0 <= lo and lo + row <= self.entry_bytes, ( + f"part {p.name!r} layer {l} spans [{lo}, {lo + row}) outside " + f"the {self.entry_bytes}-byte entry" + ) + spans.append((lo, lo + row, p.name, l)) + spans.sort() + for a, b in zip(spans, spans[1:]): + assert a[1] <= b[0], ( + f"parts overlap: {a[2]}[{a[3]}] [{a[0]}, {a[1]}) and " + f"{b[2]}[{b[3]}] [{b[0]}, {b[1]})" + ) + + +def build_dense_views( + raw: torch.Tensor, + *, + layout: DenseEntryLayout, + part: DensePart, + page_size: int, + num_pages: int, + anchor_bytes: int = 0, +) -> List[torch.Tensor]: + """Per-layer views of one part: ``(num_pages * page_size, *row_shape)`` with + slot stride ``layout.entry_bytes``, indexed by the physical token id + ``page * page_size + slot`` (``paged_view`` regroups them by page). + """ + itemsize = part.dtype.itemsize + assert layout.entry_bytes % itemsize == 0 and anchor_bytes % itemsize == 0 + n_rows = num_pages * page_size + end = anchor_bytes + n_rows * layout.entry_bytes + assert end <= raw.numel() * raw.itemsize, ( + f"build_dense_views: {n_rows} slots of {layout.entry_bytes} B end at byte " + f"{end} but the raw buffer holds only {raw.numel() * raw.itemsize} bytes" + ) + as_dtype_view = raw.view(part.dtype) + stride = (layout.entry_bytes // itemsize, *_contiguous_strides(part.row_shape)) + views: List[torch.Tensor] = [] + for layer in range(part.layer_num): + base_bytes = anchor_bytes + part.layer_offset_bytes(layer) + assert base_bytes % itemsize == 0 + views.append( + torch.as_strided( + as_dtype_view, + size=(n_rows, *part.row_shape), + stride=stride, + storage_offset=base_bytes // itemsize, + ) + ) + return views + + def mha_entry_bytes( *, layer_num: int, head_num: int, head_dim: int, v_head_dim: int, itemsize: int ) -> int: @@ -214,14 +335,6 @@ def build_page_major_mamba_views( temporal_dtype=temporal_dtype, ) - def contiguous_strides(shape: Sequence[int]) -> Tuple[int, ...]: - strides = [] - acc = 1 - for s in reversed(shape): - strides.append(acc) - acc *= int(s) - return tuple(reversed(strides)) - conv_itemsize = conv_dtype.itemsize assert entry_bytes % conv_itemsize == 0, ( f"misaligned mamba spec: per-slot entry_bytes={entry_bytes} is not a " @@ -243,7 +356,7 @@ def contiguous_strides(shape: Sequence[int]) -> Tuple[int, ...]: stride = ( inner_shape_bytes // conv_itemsize, conv_slot_stride_elems, - ) + contiguous_strides(shape) + ) + _contiguous_strides(shape) conv_views.append( torch.as_strided( as_conv_dtype, @@ -283,7 +396,7 @@ def contiguous_strides(shape: Sequence[int]) -> Tuple[int, ...]: stride = ( inner_shape_bytes // itemsize, entry_bytes // itemsize, - ) + contiguous_strides(temporal_state_shape) + ) + _contiguous_strides(temporal_state_shape) temporal_view = torch.as_strided( as_temporal_dtype, size=(layer_num, max_slots) + tuple(temporal_state_shape), From 53a3b24dceb3a3b69548cde5846e9aa1fcbf0b48 Mon Sep 17 00:00:00 2001 From: Caihua Li Date: Sat, 12 Sep 2026 01:03:53 -0700 Subject: [PATCH 02/15] feat(unified-memory): store the unified pool token-major Switch the unified pool's byte law from the page-major envelope to the token-major entry, so that a kernel-facing id IS the physical token id. Before, one page held a uniform array of row-blocks, one per layer per side: page bytes = [L0_K * ps | L0_V * ps | L1_K * ps | L1_V * ps | ...] Every per-layer view was contiguous, and the block origin was folded into its storage offset, so layer l had to be addressed by a re-numbered id `(t // ps) * (ps * blocks_per_page) + t % ps`. Views ran past the last page envelope, which is what view_tail_pad_bytes sized. Now one slot's entry holds every part of that token at a fixed offset: page bytes = [ entry(slot 0) | entry(slot 1) | ... ] entry bytes = [ K_0 | V_0 | K_1 | V_1 | ... ] so the per-layer view is strided with the entry as its slot stride and is indexed by `page * ps + slot` directly. Three consequences: * No re-numbering. kernel_page_multiplier is pinned to 1 (asserted, not read from the spec), and translate_kv_loc_for_kernel now computes what translate_kv_loc does. The plumbing is removed in a later commit. * No tail pad. The views end exactly at the last slot, so the raw buffer is allocated at total_bytes. * Parts may differ in row width. Only the startup screen still refuses asymmetric K/V, because the backends' write and read paths are not audited for it -- not because the views cannot hold it. entry_bytes() is rounded up to ENTRY_ALIGN_BYTES here rather than in the layout types: it feeds page_bytes(), buffer sizing, the sub-pool anchors, move_kv_cache envelopes and PD transfer, so the alignment has to flip together with them. The view tests are rewritten against the new byte law `byte(t, l, K) = t * entry + l * (k + v)`, each cross-checked against an independent (page, slot) description at page sizes 1 and 4. test_hicache.py's unified L2 case built its pool with 8 B rows (fp16, one head of width 4), below the 16 B row alignment every token-major part needs; it now uses 16 B rows, the smallest the layout admits. The entry, and so each page the test copies, stays 32 B. fa4, trtllm_mla, cutedsl_mla and tokenspeed_mla are admitted under --enable-unified-memory and supported by this change, but not yet validated. From this commit on their kernels read per-layer views whose slot stride is a whole entry, and every validation run was on H20 (SM90), where those kernels do not run. The same holds for trtllm_mha's SM100 kernels (trtllm-gen context and decode) and for decode context parallelism (--dcp-size > 1). Reviewers with SM100-class hardware: please run them before merging. --- python/sglang/srt/arg_groups/kv_cache_hook.py | 20 +- .../mem_cache/allocator/unified_sub_pool.py | 26 +- .../sglang/srt/mem_cache/layout/page_major.py | 40 +- .../srt/mem_cache/unified_memory_pool.py | 220 +++++---- .../kernels/ops/kvcache/test_hicache.py | 8 +- .../unit/managers/test_prefill_adder.py | 10 +- .../test_scheduler_init_req_max_new_tokens.py | 10 +- .../unit/mem_cache/test_full_loc_fast_path.py | 2 +- .../mem_cache/test_kv_index_translator.py | 158 +++--- .../unit/mem_cache/test_layout_compat.py | 66 ++- .../mem_cache/test_multi_ended_allocator.py | 50 +- .../test_pd_envelope_transfer_layout.py | 157 +++--- .../mem_cache/test_prefill_memory_budget.py | 30 +- .../mem_cache/test_unified_handout_zeroing.py | 16 +- .../unit/mem_cache/test_unified_mha_views.py | 462 +++++++++++------- .../mem_cache/test_unified_mla_block_table.py | 36 +- .../mem_cache/test_unified_mla_gpu_parity.py | 2 +- .../unit/mem_cache/test_unified_mla_views.py | 280 ++++++----- .../mem_cache/test_unified_npool_sweep.py | 12 +- .../test_page_major_backend_allowlist.py | 50 +- 20 files changed, 987 insertions(+), 668 deletions(-) diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index 196d8fdfd94d..09d9966245d4 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -656,7 +656,7 @@ def _validate_unified_memory_dcp(server_args: Any) -> None: def handle_page_major_kv_layout(server_args: Any): - # The unified pool stores state in the page-major envelope-strided layout, so + # The unified pool stores state in the page-major envelope layout, so # enabling it implies --enable-page-major-kv-layout — routing it through the # single page-major path + stride-aware Triton asserts (set before the guard). @@ -684,17 +684,18 @@ def handle_page_major_kv_layout(server_args: Any): assert unified_memory_supported_for_model( model_config, use_mla_backend=use_mla_backend(server_args) ), ( - "--enable-unified-memory requires uniform K/V rows " - "(head_dim == v_head_dim); this model has " + "--enable-unified-memory does not yet admit asymmetric K/V rows " + "(head_dim != v_head_dim); this model has " f"head_dim={model_config.head_dim}, " f"v_head_dim={model_config.v_head_dim}, " f"swa_head_dim={model_config.swa_head_dim}, " - f"swa_v_head_dim={model_config.swa_v_head_dim}. The unified " - "pool's per-layer views require a uniform row width; run " - "this model without --enable-unified-memory." + f"swa_v_head_dim={model_config.swa_v_head_dim}. The token-major " + "views can hold them, but the backends' write and read paths are " + "not audited for it; run this model without --enable-unified-memory." ) # Allow-list. Every backend below reads through the translator, so what - # gates one is only whether its kernels can address the per-layer views: + # gates one is only whether its kernels address the per-layer views by + # their strides (the slot stride is the whole entry, not one row): # * MLA models: the full paged MLA family, incl. flashmla (ps=64 # snap). # * MHA/SWA models: fa3 / fa4 / flashinfer / trtllm_mha alongside @@ -727,9 +728,8 @@ def handle_page_major_kv_layout(server_args: Any): "--enable-page-major-kv-layout: the resolved attention backends " f"{sorted(backends)} are not in the allowed set " f"{sorted(allowed_full)} for this configuration (unified memory " - "allows the per-layer-view families; plain page-major keeps the " - "envelope-strided views only Triton reads). Pass a compatible " - "--attention-backend." + "allows the stride-aware per-layer-view families). Pass a " + "compatible --attention-backend." ) # The Mamba/KDA state is stored in envelope-strided views; only # stride-audited kernels may read it (Stage 4 audit, per slot): diff --git a/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py b/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py index 2594572d72d6..d37413b2836b 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py +++ b/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py @@ -316,13 +316,14 @@ def __init__( self.entry_bytes = spec.entry_bytes() self.min_slot_index = unified_buffer.min_slot_index(sub_pool_name) self.is_id_owner = is_id_owner - # Kernel-facing page-stride scale, from the spec that owns the layout; - # `kernel_page_multiplier=` overrides it only for tests. - self.kernel_page_multiplier = ( - spec.blocks_per_page() - if kernel_page_multiplier is None - else kernel_page_multiplier + # Kernel-facing ids are the physical token ids: the token-major views + # step slots by the whole entry, so there is no per-page block scale. + # The kwarg is accepted (and must be 1) until the plumbing is removed. + assert kernel_page_multiplier in (None, 1), ( + f"kernel_page_multiplier must be 1 (token-major views); got " + f"{kernel_page_multiplier}" ) + self.kernel_page_multiplier = 1 # Zero page envelopes on hand-out -- see _maybe_zero_pages. self._zero_pages_on_alloc = isinstance(kvcache, UnifiedMLATokenToKVPool) # Overlap mode: `free` drops a wait_stream(forward_stream) barrier so its @@ -1060,14 +1061,15 @@ def translate_kv_loc_for_kernel( *, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: - """Virtual token ids -> kernel-facing ids: + """Virtual token ids -> kernel-facing ids, which under the token-major + views are the physical token ids: - kernel_id(t) = (t // ps) * (ps * kernel_page_multiplier) + t % ps + kernel_id(t) = v2p[t // ps] * ps + t % ps - Internal machinery (compaction, in-flight write sets) MUST keep using - `translate_kv_loc`: kernel-facing ids are for kernels only. Tombstones (-1) - clamp to kernel-facing id 0, the page-0 sink. int64 out; a consumer whose - kernel ABI wants int32 narrows where it fills that buffer. + (`kernel_page_multiplier` is pinned to 1, so this equals + `translate_kv_loc`.) Tombstones (-1) clamp to id 0, the page-0 sink. + int64 out; a consumer whose kernel ABI wants int32 narrows where it + fills that buffer. """ with record_function("MultiEndedAlloc.translate_kv_loc_for_kernel"): return self._translate_loc_fused(virt_tokens, dcp_size=1, out=out) diff --git a/python/sglang/srt/mem_cache/layout/page_major.py b/python/sglang/srt/mem_cache/layout/page_major.py index 129b5c2fa4ec..066eb93f4ed4 100644 --- a/python/sglang/srt/mem_cache/layout/page_major.py +++ b/python/sglang/srt/mem_cache/layout/page_major.py @@ -1,18 +1,24 @@ -"""Page-granularity envelope (page-major, layer-major within a page) cache views. +"""Page-granularity envelope (page-major, token-major within a page) cache views. A pool of this layout keeps all layers of all slots in one contiguous byte -buffer. The buffer is split into pages of ``page_size`` slots; within a page, -each layer's K and V (or each Mamba conv/temporal tensor) are grouped together: - - page bytes = [L0_K * ps | L0_V * ps | L1_K * ps | L1_V * ps | ...] - -Across pages the layout is envelope-major (one ``page_bytes`` block per page). -At ``page_size == 1`` a page is a single slot, so the within-page block is the -per-slot ``[L0_K | L0_V | L1_K | L1_V | ...]`` envelope (token-granularity). - -These builders produce per-layer views into a raw ``uint8`` buffer; they hold -no allocator/ownership state. ``anchor_bytes`` is the byte offset of the -pool's region inside the raw buffer (0 for a standalone pool). +buffer, split into pages of ``page_size`` slots. Within a page the slots follow +one another, and one slot's ENTRY holds every part of that token -- K and V of +every layer for MHA, the latent row of every layer for MLA -- at a fixed byte +offset: + + page bytes = [ entry(slot 0) | entry(slot 1) | ... | entry(slot ps-1) ] + entry bytes = [ K_0 | V_0 | K_1 | V_1 | ... ] (MHA) + [ lat_0 | lat_1 | ... ] (MLA) + +Every per-layer view is therefore a flat ``(num_pages * page_size, *row_shape)`` +tensor with slot stride ``entry_bytes`` and storage offset +``anchor + part offset``, indexed by the PHYSICAL token id +``page * page_size + slot``. Parts may differ in row width (K vs V); only their +offsets differ, never the stride. + +These builders produce views into a raw ``uint8`` buffer; they hold no +allocator/ownership state. ``anchor_bytes`` is the byte offset of the pool's +region inside the raw buffer (0 for a standalone pool). """ from typing import List, Sequence, Tuple @@ -151,10 +157,10 @@ def build_dense_views( def mha_entry_bytes( *, layer_num: int, head_num: int, head_dim: int, v_head_dim: int, itemsize: int ) -> int: - """Bytes occupied by one slot across all layers (K and V).""" + """Bytes occupied by one slot across all layers (K and V), aligned.""" k_row_bytes = head_num * head_dim * itemsize v_row_bytes = head_num * v_head_dim * itemsize - return layer_num * (k_row_bytes + v_row_bytes) + return align_entry_bytes(layer_num * (k_row_bytes + v_row_bytes)) def build_mha_views( @@ -226,8 +232,8 @@ def build_mha_views( def mla_entry_bytes(*, layer_num: int, kv_cache_dim: int, itemsize: int) -> int: - """Bytes occupied by one MLA slot across all layers (single latent row, no V).""" - return layer_num * kv_cache_dim * itemsize + """Bytes occupied by one MLA slot across all layers (latent rows), aligned.""" + return align_entry_bytes(layer_num * kv_cache_dim * itemsize) def build_mla_views( diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index f243a5a51bfc..021930941bb4 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -37,8 +37,10 @@ from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.environ import envs from sglang.srt.mem_cache.layout.page_major import ( - build_mha_views, - build_mla_views, + DenseEntryLayout, + DensePart, + align_entry_bytes, + build_dense_views, build_page_major_mamba_views, ) from sglang.srt.mem_cache.memory_pool import ( @@ -159,14 +161,40 @@ def v_row_bytes(self) -> int: return self.head_num * self.v_head_dim * self.store_dtype.itemsize def entry_bytes(self) -> int: - return self.layer_num * (self.k_row_bytes() + self.v_row_bytes()) + return align_entry_bytes( + self.layer_num * (self.k_row_bytes() + self.v_row_bytes()) + ) - # Page-major byte math: within a page block K/V group per layer - # [L0_K*ps | L0_V*ps | L1_K*ps | ...]; at ps==1 this collapses to the per-slot envelope. + # Token-major entry: [K_0 | V_0 | K_1 | V_1 | ...] per slot; a page is + # page_size such entries back to back. def page_bytes(self, page_size: int) -> int: return page_size * self.entry_bytes() + def layout(self) -> DenseEntryLayout: + layer_stride = self.k_row_bytes() + self.v_row_bytes() + return DenseEntryLayout( + entry_bytes=self.entry_bytes(), + parts=( + DensePart( + name="k", + offset_bytes=0, + layer_stride_bytes=layer_stride, + layer_num=self.layer_num, + row_shape=(self.head_num, self.head_dim), + dtype=self.store_dtype, + ), + DensePart( + name="v", + offset_bytes=self.k_row_bytes(), + layer_stride_bytes=layer_stride, + layer_num=self.layer_num, + row_shape=(self.head_num, self.v_head_dim), + dtype=self.store_dtype, + ), + ), + ) + def view_tail_pad_bytes(self, page_size: int) -> int: return page_size * self.entry_bytes() @@ -205,8 +233,26 @@ def __post_init__(self): def kv_cache_dim(self) -> int: return self.kv_lora_rank + self.qk_rope_head_dim + def row_bytes(self) -> int: + return self.kv_cache_dim * self.store_dtype.itemsize + def entry_bytes(self) -> int: - return self.layer_num * self.kv_cache_dim * self.store_dtype.itemsize + return align_entry_bytes(self.layer_num * self.row_bytes()) + + def layout(self) -> DenseEntryLayout: + return DenseEntryLayout( + entry_bytes=self.entry_bytes(), + parts=( + DensePart( + name="kv", + offset_bytes=0, + layer_stride_bytes=self.row_bytes(), + layer_num=self.layer_num, + row_shape=(1, self.kv_cache_dim), + dtype=self.store_dtype, + ), + ), + ) def view_tail_pad_bytes(self, page_size: int) -> int: return page_size * self.entry_bytes() @@ -296,10 +342,10 @@ def _reserved_floor_bytes(sub_pool_specs: List[SubPoolSpec], page_size: int) -> class UnifiedKVPool: """One physical `uint8` byte buffer shared by N sub-pools, each exposing - per-layer views over its own byte range (contiguous per layer for KV, - strided for the Mamba state). Two END pools (one grow-up, one grow-down) - own the buffer's ends; optional "float" MIDDLE pools live between their - frontiers. Allocators keep byte ranges disjoint; no usage tracking here. + per-layer strided views over its own byte range. Two END pools (one + grow-up, one grow-down) own the buffer's ends; optional "float" MIDDLE + pools live between their frontiers. Allocators keep byte ranges disjoint; + no usage tracking here. """ def __init__( @@ -347,13 +393,8 @@ def __init__( self.memory_saver_adapter = TorchMemorySaverAdapter.create( enable=enable_memory_saver ) - self.view_tail_pad_bytes = max( - spec.view_tail_pad_bytes(page_size) for spec in sub_pool_specs - ) with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): - self._raw = torch.empty( - total_bytes + self.view_tail_pad_bytes, dtype=torch.uint8, device=device - ) + self._raw = torch.empty(total_bytes, dtype=torch.uint8, device=device) if envs.SGLANG_DEBUG_POISON_POOL.get(): # Debug: bf16-NaN-fill so NaN-unsafe reads of never-written bytes # fail deterministically. @@ -368,7 +409,7 @@ def __init__( self._max_slots: Dict[str, int] = {} self._anchor_bytes: Dict[str, int] = {} self._min_slot_index: Dict[str, int] = {} - # MHA: (k_buffer, v_buffer); MLA: [per-layer per-layer views]; + # MHA: (k_buffer, v_buffer); MLA: [per-layer views]; # Mamba: (conv_state_list, temporal_state) self._mha_views: Dict[str, Tuple[List[torch.Tensor], List[torch.Tensor]]] = {} self._mla_views: Dict[str, List[torch.Tensor]] = {} @@ -491,21 +532,20 @@ def _build_mha_views( page_size: int, ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]: num_pages = max_slots // page_size - _assert_kernel_id_bound( - sub_pool_name=spec.name, - n_rows=num_pages * spec.blocks_per_page() * page_size, - ) - return build_mha_views( - self._raw, - layer_num=spec.layer_num, - head_num=spec.head_num, - head_dim=spec.head_dim, - v_head_dim=spec.v_head_dim, - store_dtype=spec.store_dtype, - page_size=page_size, - num_pages=num_pages, - anchor_bytes=anchor_bytes, + _assert_kernel_id_bound(sub_pool_name=spec.name, n_rows=num_pages * page_size) + layout = spec.layout() + k_views, v_views = ( + build_dense_views( + self._raw, + layout=layout, + part=layout.part(name), + page_size=page_size, + num_pages=num_pages, + anchor_bytes=anchor_bytes, + ) + for name in ("k", "v") ) + return k_views, v_views def _build_mla_views( self, @@ -515,15 +555,12 @@ def _build_mla_views( page_size: int, ) -> List[torch.Tensor]: num_pages = max_slots // page_size - _assert_kernel_id_bound( - sub_pool_name=spec.name, - n_rows=num_pages * spec.blocks_per_page() * page_size, - ) - return build_mla_views( + _assert_kernel_id_bound(sub_pool_name=spec.name, n_rows=num_pages * page_size) + layout = spec.layout() + return build_dense_views( self._raw, - layer_num=spec.layer_num, - kv_cache_dim=spec.kv_cache_dim, - store_dtype=spec.store_dtype, + layout=layout, + part=layout.part("kv"), page_size=page_size, num_pages=num_pages, anchor_bytes=anchor_bytes, @@ -545,16 +582,13 @@ def _build_mamba_views( class UnifiedMHATokenToKVPool(MHATokenToKVPool): - """MHA KV pool whose per-layer `k_buffer`/`v_buffer` are `build_mha_views` - views into a `UnifiedKVPool` (requires uniform K/V rows). - - Views are contiguous `(n_rows, head_num, head_dim)`; locs are - - kernel_id(t) = (t // ps) * (ps * 2 * layer_num) + t % ps + """MHA KV pool whose per-layer `k_buffer`/`v_buffer` are token-major views + into a `UnifiedKVPool` (see `build_dense_views`). - which is layer- and K/V-independent, each view's storage_offset folding in - its block origin (layer l's K at block 2l, V at 2l+1). `move_kv_cache` is - the exception: compaction passes REAL physical token ids. + Each view is a strided `(num_pages * ps, head_num, head_dim)` tensor whose + slot stride is the whole entry, indexed by the PHYSICAL token id; K and V + of every layer sit at fixed offsets inside that entry, so K and V rows may + differ in width. `move_kv_cache` relocates whole page envelopes. """ def __init__( @@ -577,7 +611,7 @@ def __init__( self._v_views = v_views self._num_pages = max_slots // page_size self._page_bytes = page_size * spec.entry_bytes() - view_rows = self._num_pages * spec.blocks_per_page() * page_size + view_rows = self._num_pages * page_size super().__init__( size=view_rows - page_size, @@ -595,7 +629,6 @@ def __init__( enable_kv_cache_copy=False, kv_cache_layout="page_major", ) - self.kernel_page_blocks = spec.blocks_per_page() def _create_buffers(self): self.k_buffer = self._k_views @@ -611,9 +644,8 @@ def get_kv_size_bytes(self): return 0, 0 # UnifiedKVPool logs the total; per-sub-pool would double-count def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): - """Relocate slots by whole page envelope. - `tgt_loc`/`src_loc` are REAL physical token ids, not kernel-facing ids. - """ + """Relocate slots by whole page envelope (`tgt_loc`/`src_loc` are + physical token ids).""" if tgt_loc.numel() == 0: return # The envelope view below starts at byte 0, so this sub-pool must be @@ -654,19 +686,15 @@ def grow_direction(self) -> str: """Physical growth direction used by this sub-pool's L1 allocator.""" return self._unified_buffer.spec(self._sub_pool_name).grow_direction - def _physical_to_kernel_indices(self, indices: torch.Tensor) -> torch.Tensor: - return (indices // self.page_size) * ( - self.page_size * self.kernel_page_blocks - ) + indices % self.page_size - def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): - """Translate physical host-pool ids for the page-major parent path.""" - return super().get_cpu_copy(self._physical_to_kernel_indices(indices)) + """A physical token id IS the kernel-facing id, so the host copy needs + no rewrite; the override drops the arguments the parent ignores.""" + return super().get_cpu_copy(indices) def load_cpu_copy( self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None ): - super().load_cpu_copy(kv_cache_cpu, self._physical_to_kernel_indices(indices)) + super().load_cpu_copy(kv_cache_cpu, indices) def set_kv_buffer_prefix_valid(self, *args, **kwargs): raise NotImplementedError( @@ -676,21 +704,15 @@ def set_kv_buffer_prefix_valid(self, *args, **kwargs): class UnifiedMLATokenToKVPool(MLATokenToKVPool): - """MLA KV pool whose per-layer `kv_buffer` entries are kernel-facing views into a - `UnifiedKVPool` (see `build_mla_views`). - - Loc-space contract: every loc this pool receives through the KVCache API - (`set_kv_buffer` / `set_mla_kv_buffer` / `get_mla_kv_buffer`, and the - kv_indices consumed by attention kernels reading `get_key_buffer` / - `get_value_buffer`) is a kernel-facing id — the `translate_kv_loc_for_kernel` output - - kernel_id(t) = (t // ps) * (ps * layer_num) + t % ps - - which is layer-independent (the layer offset is folded into each view's - storage_offset), so the stock `MLATokenToKVPool` read/write methods work on - the views unmodified. The ONE exception is `move_kv_cache`: the allocator's - compaction calls it with REAL physical token ids, and it is overridden to - relocate whole page envelopes on the raw buffer. + """MLA KV pool whose per-layer `kv_buffer` entries are token-major views + into a `UnifiedKVPool` (see `build_dense_views`). + + Every loc this pool receives through the KVCache API is a PHYSICAL token + id (`translate_kv_loc` of the virtual id). Each view is a strided + `(num_pages * ps, 1, kv_cache_dim)` tensor whose slot stride is the whole + entry, the layer offset folded into its storage_offset, so the stock + `MLATokenToKVPool` read/write methods work on the views unmodified. + `move_kv_cache` relocates whole page envelopes on the raw buffer. """ def __init__( @@ -714,10 +736,10 @@ def __init__( max_slots = unified_buffer.max_slots(sub_pool_name) self._num_pages = max_slots // page_size self._page_bytes = page_size * spec.entry_bytes() - self._view_rows = self._num_pages * spec.blocks_per_page() * page_size + self._view_rows = self._num_pages * page_size super().__init__( - # OOB checks bound locs by `size + page_size`; kernel-facing ids run to + # OOB checks bound locs by `size + page_size`; physical ids run to # `_view_rows` (page 0 is the reserved padding sink). size=self._view_rows - page_size, page_size=page_size, @@ -728,7 +750,6 @@ def __init__( device=unified_buffer.device, enable_memory_saver=False, # buffer owned by UnifiedKVPool ) - self.kernel_page_blocks = spec.blocks_per_page() def _create_buffers(self): self.kv_buffer = self._kv_views @@ -749,36 +770,30 @@ def get_contiguous_buf_infos(self): ``raw_ptr + physical_page_id * page_envelope_bytes``. The transfer item is the whole page envelope (all layers of one page) - rather than a per-layer region, because the per-layer per-layer views - overlap and index in kernel-facing ids. Both sides must therefore build the - pool with identical specs. + rather than a per-layer region: the layers interleave inside each + slot's entry. Both sides must therefore build the pool with identical + specs. """ # The address formula omits the anchor; a nonzero one would mis-address. assert self._unified_buffer.anchor_bytes(self._sub_pool_name) == 0 raw = self._unified_buffer._raw return [raw.data_ptr()], [raw.numel()], [self._page_bytes] - def _physical_to_kernel_indices(self, indices: torch.Tensor) -> torch.Tensor: - """Physical TOKEN ids -> the kernel-facing ids this class's `kv_buffer` - views are indexed by; the formula is the one in the class docstring.""" - return (indices // self.page_size) * ( - self.page_size * self.kernel_page_blocks - ) + indices % self.page_size - def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): - """Translate physical host-pool ids for the page-major parent path.""" - return super().get_cpu_copy(self._physical_to_kernel_indices(indices)) + """A physical token id IS the kernel-facing id, so the host copy needs + no rewrite; the override drops the arguments the parent ignores.""" + return super().get_cpu_copy(indices) def load_cpu_copy( self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None ): - super().load_cpu_copy(kv_cache_cpu, self._physical_to_kernel_indices(indices)) + super().load_cpu_copy(kv_cache_cpu, indices) def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): """Relocate whole page envelopes. - `tgt_loc`/`src_loc` are REAL physical token ids (NOT kernel-facing ids): both - compaction paths expand page ids into page-major-ordered token runs + `tgt_loc`/`src_loc` are physical token ids: both compaction paths + expand page ids into page-major-ordered token runs (`pages[:, None] * ps + offsets`), relied on here to recover the page lists. One contiguous envelope copy replaces the per-layer strided moves. """ @@ -1458,15 +1473,13 @@ def init_unified_mamba_pools( if use_mla_backend: logger.info( "[unified-memory-pool] full_layers=%d, mamba_layers=%d, kv_lora_rank=%d, " - "qk_rope_head_dim=%d, page_size=%d (per-layer views, kernel_page_multiplier=%d, " - "view_tail_pad=%d B)", + "qk_rope_head_dim=%d, page_size=%d (token-major views, entry_bytes=%d B)", len(full_attention_layer_ids), len(mamba_layer_ids), kv_lora_rank, qk_rope_head_dim, page_size, - len(full_attention_layer_ids), - shared_pool.view_tail_pad_bytes, + full_spec.entry_bytes(), ) else: logger.info( @@ -1478,8 +1491,7 @@ def init_unified_mamba_pools( head_dim, page_size, is_draft_worker, - "per-layer views, kernel_page_multiplier=%d, view_tail_pad=%d B" - % (full_spec.blocks_per_page(), shared_pool.view_tail_pad_bytes), + "token-major views, entry_bytes=%d B" % full_spec.entry_bytes(), ) logger.info( "[unified-memory-pool] total_bytes=%d, max_total_num_tokens=%d, max_mamba_cache_size=%d, " @@ -1903,12 +1915,8 @@ def init_unified_swa_pools( logger.info("[unified-memory-pool] UNIFIED MEMORY POOL ENABLED -- path=SWA hybrid") logger.info( "[unified-memory-pool] %s", - "per-layer views, kernel_page_multiplier full=%d swa=%d, view_tail_pad=%d B" - % ( - full_spec.blocks_per_page(), - swa_spec.blocks_per_page(), - shared_pool.view_tail_pad_bytes, - ), + "token-major views, entry_bytes full=%d swa=%d B" + % (full_spec.entry_bytes(), swa_spec.entry_bytes()), ) logger.info( "[unified-memory-pool] full_layers=%d, swa_layers=%d, head_num=%d, head_dim=%d, " diff --git a/test/registered/kernels/ops/kvcache/test_hicache.py b/test/registered/kernels/ops/kvcache/test_hicache.py index 83c61eee76c0..4fe68280ac95 100644 --- a/test/registered/kernels/ops/kvcache/test_hicache.py +++ b/test/registered/kernels/ops/kvcache/test_hicache.py @@ -54,11 +54,11 @@ def test_unified_l2_waits_for_index_producer(direction, backend): device=DEVICE, kv_cache_dtype=torch.float16, head_num=1, - head_dim=4, - v_head_dim=4, + head_dim=8, + v_head_dim=8, swa_head_num=1, - swa_head_dim=4, - swa_v_head_dim=4, + swa_head_dim=8, + swa_v_head_dim=8, page_size=1, start_layer=0, end_layer=2, diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index 5d3877f18a53..0e475d5ea63e 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -154,17 +154,17 @@ def create_shared_adder(self, *, num_mixed_decode_tokens=0): device="cpu", kv_cache_dtype=torch.float16, head_num=1, - head_dim=4, - v_head_dim=4, + head_dim=8, + v_head_dim=8, swa_head_num=1, - swa_head_dim=4, - swa_v_head_dim=4, + swa_head_dim=8, + swa_v_head_dim=8, page_size=4, start_layer=0, end_layer=2, swa_attention_layer_ids=[1], full_attention_layer_ids=[0], - total_bytes=1024, + total_bytes=2048, enable_memory_saver=False, need_sort=False, lazy_compaction=True, diff --git a/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py b/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py index d2d034dba8a4..c9cbaa81f878 100644 --- a/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py +++ b/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py @@ -232,17 +232,17 @@ def test_unified_budget_rounds_prompt_and_decode_together(self): device="cpu", kv_cache_dtype=torch.float16, head_num=1, - head_dim=4, - v_head_dim=4, + head_dim=8, + v_head_dim=8, swa_head_num=1, - swa_head_dim=4, - swa_v_head_dim=4, + swa_head_dim=8, + swa_v_head_dim=8, page_size=4, start_layer=0, end_layer=2, swa_attention_layer_ids=[1], full_attention_layer_ids=[0], - total_bytes=384, + total_bytes=768, enable_memory_saver=False, need_sort=False, lazy_compaction=True, diff --git a/test/registered/unit/mem_cache/test_full_loc_fast_path.py b/test/registered/unit/mem_cache/test_full_loc_fast_path.py index 8c301633e5d9..f72ebe024e5c 100644 --- a/test/registered/unit/mem_cache/test_full_loc_fast_path.py +++ b/test/registered/unit/mem_cache/test_full_loc_fast_path.py @@ -170,7 +170,7 @@ def _make_bare_pool(self, page_size, v2p, multiplier=1): return pool def test_tombstoned_id_lands_on_sink(self): - for ps, mult in ((1, 1), (4, 1), (4, 6)): + for ps, mult in ((1, 1), (4, 1)): v2p = torch.tensor([0, -1, 2], dtype=torch.int64) pool = self._make_bare_pool(ps, v2p, multiplier=mult) # Virtual ids covering the tombstoned page (index 1) and a live one. diff --git a/test/registered/unit/mem_cache/test_kv_index_translator.py b/test/registered/unit/mem_cache/test_kv_index_translator.py index 45768e04a2c0..de767f56f694 100644 --- a/test/registered/unit/mem_cache/test_kv_index_translator.py +++ b/test/registered/unit/mem_cache/test_kv_index_translator.py @@ -13,15 +13,34 @@ # ============================================================================== """KVIndexTranslator -- the read-path id translator. -CPU-only: these exercise the builder's pure-torch reference path, not the -Triton kernel. +Covers, CPU-only (the builder's pure-torch reference path; GPU parity of the +Triton kernel is a later CUDA CI pin): + - strict passthrough: a non-unified source returns the SAME req_to_token / + req_pool_indices objects -- zero tensor ops, no copies (the property that + makes backend re-pointing byte-identical for every non-unified server); + - static SWA pools keep their legacy full->swa mapping on the view; + - the read table matches the hand formula + entry[b, c] = clamp(v2p[req_to_token[req[b], c*ps] // ps] * mult, 0) + over the REAL SWA composite's tables (full AND swa, ps in {1, 4}, + multiplier in {1, 2L}), with the swa table built from VIRTUAL ids; + - sink routing: dead lanes (seq_len 0), -1 req_to_token entries, and + tombstoned v2p pages all read entry 0; + - the capture contract: buffers are zero-filled and idempotent; a refresh + updates ONLY the live prefix (stale tails and rows beyond bs keep prior + contents); the returned table is the WHOLE buffer (pointer-stable); + - the eager-view memo: a single source-resident slot keyed by batch + identity (same batch shares one build; the next batch replaces it; a + dead batch never matches); + - the two-phase write contract: the rebind touches only the full side, and + the sliding-window write loc derives POINTWISE from the kernel-facing values + (pads, slices, and fresh copies included), for both pool families. + + python -m pytest test/registered/unit/mem_cache/test_kv_index_translator.py -v """ -from sglang.srt.runtime_context import publish, reset_context -from sglang.srt.server_args import ServerArgs from sglang.test.ci.ci_register import register_cpu_ci -register_cpu_ci(est_time=10, suite="base-a-test-cpu") +register_cpu_ci(est_time=8, suite="base-a-test-cpu") import unittest from types import SimpleNamespace @@ -35,6 +54,8 @@ from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator, KVReadTables from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.mem_cache.unified_memory_pool import MHASubPoolSpec, UnifiedKVPool +from sglang.srt.runtime_context import publish, reset_context +from sglang.srt.server_args import ServerArgs from sglang.srt.state_capturer.base import BaseTopkCapturer from sglang.test.test_utils import CustomTestCase @@ -43,7 +64,7 @@ _SWA_L = 3 -def _build_composite(ps, collapse=False, n_full_pages=16, n_swa_pages=8): +def _build_composite(ps, n_full_pages=16, n_swa_pages=8): full_spec = MHASubPoolSpec( name="full", layer_num=_FULL_L, @@ -80,11 +101,6 @@ def _build_composite(ps, collapse=False, n_full_pages=16, n_swa_pages=8): need_sort=False, forward_stream=None, ) - if collapse: - # The multiplier-1 arm, where kernel-facing ids ARE the physical ones. - # No unified sub-pool reports 1 today, so pin the regime here. - allocator.full_attn_allocator.kernel_page_multiplier = 1 - allocator.swa_attn_allocator.kernel_page_multiplier = 1 # The fake IS the runner's token_to_kv_pool, and the real UnifiedSWAKVPool # carries the pool-level full->swa translate, so the fake must too. kvcache.translate_loc_from_full_to_swa = allocator.translate_loc_from_full_to_swa @@ -126,8 +142,9 @@ def setUp(self): publish(ServerArgs(model_path="dummy"), role="tokenizer") def test_non_unified_returns_same_objects(self): - """Strict passthrough: no copy, no branch. Any tensor op on the - non-unified path breaks byte-identity for every static-pool server.""" + """The strict-passthrough property: no copy, no branch, the exact + tensors backends read today. A regression here (any tensor op on the + non-unified path) breaks byte-identity for every static-pool server.""" req_to_token = torch.arange(64, dtype=torch.int64).reshape(4, 16) src = KVIndexTranslator( req_to_token=req_to_token, @@ -155,7 +172,7 @@ def test_non_unified_returns_same_objects(self): def _alloc_and_fill(allocator, ps, lens): """Allocate per-request virtual runs and write them into a fake - req_to_token.""" + req_to_token; returns (req_to_token, req_pool_indices, seq_lens).""" width = 16 * ps req_to_token = torch.full((len(lens), width), -1, dtype=torch.int64) for r, n in enumerate(lens): @@ -177,14 +194,14 @@ def setUp(self): self.addCleanup(reset_context) publish(ServerArgs(model_path="dummy"), role="tokenizer") - def test_read_table_matches_reference_across_multipliers(self): - """Both read tables must equal the independent per-element derivation, - across page sizes and both multiplier regimes (MLA=1, MHA=2L); the swa - table agreeing over VIRTUAL ids proves it is never chained through - full-physical.""" + def test_read_table_matches_reference(self): + """The load-bearing formula pin: full AND swa read tables equal + the independent per-element derivation across page sizes (kernel-facing + ids are the physical ones: multiplier 1). The swa table agreeing with + a formula over VIRTUAL ids is also the never-chained-through- + full-physical proof.""" for ps in (1, 4): - for collapse in (True, False): - allocator = _build_composite(ps, collapse=collapse) + for allocator in (_build_composite(ps),): full_mult = allocator.kernel_page_multiplier swa_mult = allocator.swa_kernel_page_multiplier req_to_token, rows, seq_lens = _alloc_and_fill( @@ -229,8 +246,10 @@ def test_read_table_matches_reference_across_multipliers(self): ) def test_packed_stream_equals_the_rectangle_it_replaces(self): - """The two builders must agree element for element: - packed[indptr[b] + p] == ids[b, p // ps] * ps + p % ps.""" + """The packed builder and the rectangle must agree element for element: + packed[indptr[b] + p] == ids[b, p // ps] * ps + p % ps. Consumers that + plan over the stream and consumers that read the page table have to see + the same KV, so a change to either builder alone turns this red.""" for ps in (1, 4): allocator = _build_composite(ps) req_to_token, rows, seq_lens = _alloc_and_fill( @@ -261,9 +280,9 @@ def test_packed_stream_equals_the_rectangle_it_replaces(self): ) def test_sink_routing(self): - """Dead lanes, -1 slots inside the live prefix, and tombstoned v2p - pages must ALL read entry 0; one wild entry is a captured-graph OOB - read at replay.""" + """Dead lanes (seq_len 0), -1 slots inside the live prefix, and + tombstoned v2p pages must ALL read entry 0 -- one wild entry is a + captured-graph OOB read at replay.""" ps = 4 allocator = _build_composite(ps) req_to_token, rows, seq_lens = _alloc_and_fill( @@ -298,9 +317,12 @@ def setUp(self): publish(ServerArgs(model_path="dummy"), role="tokenizer") def test_prefix_filled_tail_sentinel_preserved_width_capped(self): - """The -1 tail sentinel belongs to the backend, and a table padded - WIDER than the req_to_token page span (trtllm's LCM alignment) must be - capped rather than trip the builder's width assert.""" + """Three contracts in one batch: entries equal the read-table formula, + lanes past each row's live pages keep the backend's -1 sentinel + (prefix-only -- a tail write scatters the trtllm sentinel contract), + and a table padded WIDER than the req_to_token page span (trtllm's + LCM alignment) is capped instead of tripping the builder's width + assert.""" ps = 4 allocator = _build_composite(ps) full_mult = allocator.kernel_page_multiplier @@ -334,8 +356,8 @@ def test_prefix_filled_tail_sentinel_preserved_width_capped(self): ) def test_passthrough_source_refuses(self): - """Callers dispatch on `reads_are_translated`; a passthrough source has - no v2p to build from and must fail loud, not fill garbage.""" + """Callers dispatch on `enabled`; a passthrough source has no v2p to + build from and must fail loud, not fill garbage.""" src = KVIndexTranslator( req_to_token=torch.zeros((2, 4), dtype=torch.int64), token_to_kv_pool_allocator=SimpleNamespace(), @@ -355,11 +377,13 @@ class TestPoolOwnership(unittest.TestCase): """A runner only gets the kernel-facing id space when the pool IT reads and writes is the one the allocator's ids address. - Guarded shape: a runner handed a SHARED allocator while owning a SEPARATE - KV buffer sized to the allocator's SLOT count. Probing the allocator alone - reports "unified" for that runner, so its indices would be mapped into the - composite's kernel-facing space (ids up to num_pages * multiplier) and used - to address a buffer with only num_slots rows. + Guarded shape: a runner handed a SHARED allocator (one slot index space, + one req_to_token) while owning a SEPARATE KV buffer sized to the + allocator's SLOT count. Probing the allocator alone reports "unified" for + that runner, so its indices would be mapped into the composite's + kernel-facing space (kernel-facing ids up to num_pages * multiplier) and then used + to address a buffer with only num_slots rows -- out of bounds on both the + read gather and the KV store. """ def setUp(self): @@ -369,10 +393,11 @@ def setUp(self): publish(ServerArgs(model_path="dummy"), role="tokenizer") def test_real_factory_bundle_satisfies_the_ownership_identity(self): - """The guard rests on `allocator.get_kvcache() is token_to_kv_pool`, so - a factory returning a pool the allocator does not hold would silently - disable the unified path for EVERY model. Pinned against the real - factory, not this file's own construction.""" + """The guard rests on `allocator.get_kvcache() is token_to_kv_pool` + holding for a REAL target bundle. If a factory ever returned a pool + the allocator does not hold, the guard would silently disable the + unified path for EVERY model -- so pin it against the real factory + rather than against this file's own construction.""" from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools bundle = init_unified_swa_pools( @@ -422,9 +447,9 @@ def test_runner_with_its_own_pool_is_disabled(self): self.assertFalse(src.is_translating) def test_disabled_source_is_strict_passthrough(self): - """Such a runner must see RAW virtual ids: they index its own pool - directly, and a translate here is the out-of-bounds bug the ownership - identity exists to prevent.""" + """Consequence of the guard: such a runner must see RAW virtual ids on + the read table -- they index its own pool directly. A translate here is + the out-of-bounds bug the ownership identity exists to prevent.""" alloc = _build_composite(ps=1) req_to_token = torch.arange(16, dtype=torch.int32, device=_DEV).view(2, 8) src = KVIndexTranslator( @@ -479,17 +504,16 @@ def test_caller_owned_table_is_returned_whole_and_filled_prefix_only(self): into=tables, ) self.assertIs(view.ids, cap, "the caller's table comes back WHOLE") - want = allocator.full_v2p_page_table[req_to_token[1, ::ps][:2] // ps] * ( - 2 * _FULL_L - ) + want = allocator.full_v2p_page_table[req_to_token[1, ::ps][:2] // ps] self.assertTrue(torch.equal(cap[0, :2], want.to(torch.int32))) self.assertTrue(bool((cap[0, 2:] == 7).all()), "stale tail was cleared") self.assertTrue(bool((cap[1:] == 7).all()), "rows beyond bs were touched") def test_row_ids_not_reallocated_across_builds(self): - """`row_ids` is one arange sized from the request pool and sliced per - build; a per-build `torch.arange` would be correct but costs an - allocation and a launch on every replay prep.""" + """`row_ids` is a constant arange sized once from the request pool, so + builds at different batch sizes hand back slices of ONE buffer. A + per-build `torch.arange` would be correct but would spend an allocation + and a launch on every replay prep.""" allocator = _build_composite(1) req_to_token, rows, seq_lens = _alloc_and_fill(allocator, 1, lens=[4, 2, 3]) src = _make_source(allocator, req_to_token, 1) @@ -534,9 +558,10 @@ def __init__( class TestViewMemo(unittest.TestCase): - """The eager view is memoized ON THE SOURCE in a single slot keyed by batch - identity, so per-batch state stays out of the ForwardBatch while one - metadata build's many consumers still share one table build.""" + """The eager view is memoized ON THE SOURCE in a single slot keyed by + batch identity -- per-batch state stays out of the ForwardBatch (it does + not scale with the number of id spaces), and one metadata build's many + consumers still share one table build.""" def setUp(self): # The code under test reads its config from the bags. @@ -653,7 +678,7 @@ def test_capture_capacity_covers_dcp_widened_mamba_ids(self): name="full", layer_num=1, head_num=1, - head_dim=4, + head_dim=8, store_dtype=torch.float16, grow_direction="up", ), @@ -752,8 +777,10 @@ def test_topk_capture_round_trips_request_token_ids(self): ) def test_swa_write_loc_round_trips_from_full_side(self): - """Derived property: `field(full(t)) == swa(t)` for any virtual run t, - across page sizes and multipliers.""" + """The derived property behind phase 2: for any virtual run t, + deriving from the kernel-facing full-side values must equal the direct + virtual->swa translate -- `field(full(t)) == swa(t)` across page sizes + and multipliers.""" for ps in (1, 4, 64): src, _, rows, seq_lens, _, want_full, want_swa = self._built( ps=ps, n=3 * ps @@ -762,15 +789,30 @@ def test_swa_write_loc_round_trips_from_full_side(self): self.assertTrue(torch.equal(got, want_swa)) def test_pad_lanes_derive_to_sink(self): - """The DP pad appends zeros, and kernel-facing 0 is the reserved - padding slot in every id space, so pad lanes derive to swa slot 0 with - no `num_live` bookkeeping.""" + """The DP pad appends zeros; kernel-facing 0 is the reserved padding slot in + every id space, so pad lanes must derive to swa slot 0 with no + `num_live` bookkeeping.""" src, _, rows, seq_lens, _, want_full, want_swa = self._built(n=3) padded = torch.cat([want_full, want_full.new_zeros(2)]) got = self._field(src, rows, seq_lens, padded) self.assertTrue(torch.equal(got[:3], want_swa)) self.assertTrue(bool((got[3:] == 0).all()), "pad lanes must land on slot 0") + def test_slice_and_copy_derive_pointwise_without_handover(self): + """REGRESSION (design): the retired identity-resolver refused any + tensor it had not been handed -- a TBO child's re-padded slice or a + registry's fresh copy raised. Value-based derivation must accept + both, pointwise, with no adopt/handover call.""" + src, _, rows, seq_lens, _, want_full, want_swa = self._built(n=4) + padded = torch.cat([want_full, want_full.new_zeros(2)]) + # TBO-child shape: a slice crossing the pad boundary. + got = self._field(src, rows, seq_lens, padded[2:6]) + self.assertTrue(torch.equal(got[:2], want_swa[2:4])) + self.assertTrue(bool((got[2:] == 0).all())) + # Registry shape: a fresh equal-value copy. + got2 = self._field(src, rows, seq_lens, want_full.clone()) + self.assertTrue(torch.equal(got2, want_swa)) + def test_tombstoned_swa_page_clamps_to_sink(self): src, allocator, rows, seq_lens, virt, want_full, _ = self._built(ps=1, n=2) allocator.swa_v2p_page_table[int(virt[0])] = -1 diff --git a/test/registered/unit/mem_cache/test_layout_compat.py b/test/registered/unit/mem_cache/test_layout_compat.py index 19a5e836da36..25c5d33b22f8 100644 --- a/test/registered/unit/mem_cache/test_layout_compat.py +++ b/test/registered/unit/mem_cache/test_layout_compat.py @@ -11,29 +11,81 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""``move_kv_cache_native`` -- the stock per-layer 3-D move static-pool -compaction rides on -- must stay byte-exact. CPU-only. +"""Unit tests for the page-major envelope byte layout. -The page-major envelope layout and the per-layer views over it are covered by -``test_unified_mha_views.py``, which pins the view addressing against the -envelope formula byte for byte. +The subject here is the ENVELOPE -- the byte layout the unified pool stores its +KV in -- pinned through ``MHASubPoolSpec``'s offset math. The 3-D per-layer views +the pool exposes over the same bytes are covered by +``test_unified_mha_views.py``, which also pins the view addressing +against the envelope formula byte for byte. + +Verifies that: +1. ``MHASubPoolSpec.layout()`` -- the descriptor every per-layer view is + built through -- places each layer's K/V at the token-major entry offsets, + independent of the page size. +2. ``move_kv_cache_native`` (the stock per-layer 3-D move) stays byte-exact. + +CPU-only -- no GPU / Triton needed. + + python -m pytest test/registered/unit/mem_cache/test_layout_compat.py -v """ from sglang.test.ci.ci_register import register_cpu_ci -register_cpu_ci(est_time=9, suite="base-a-test-cpu") +register_cpu_ci(est_time=6, suite="base-a-test-cpu") import unittest import torch from sglang.srt.mem_cache.memory_pool import move_kv_cache_native +from sglang.srt.mem_cache.unified_memory_pool import MHASubPoolSpec + +_DEV = "cpu" + + +def _make_mha_spec(name, grow, layer_num=2, head_num=2, head_dim=4): + return MHASubPoolSpec( + name=name, + layer_num=layer_num, + head_num=head_num, + head_dim=head_dim, + store_dtype=torch.float16, + grow_direction=grow, + ) + + +class TestMHASpecLayerOffsets(unittest.TestCase): + """Verify the layout descriptor's part offsets against the entry math.""" + + def test_offsets_match_token_major_entry(self): + spec = _make_mha_spec("full", "up", layer_num=3, head_num=2, head_dim=4) + # Layer L's K sits at L * (k_row + v_row) inside the slot's entry and + # its V right after that K. The page size never enters: a page is + # page_size entries back to back. + k_row = spec.k_row_bytes() + v_row = spec.v_row_bytes() + layout = spec.layout() + self.assertEqual(layout.entry_bytes, spec.entry_bytes()) + for L in range(spec.layer_num): + self.assertEqual( + layout.part("k").layer_offset_bytes(L), L * (k_row + v_row) + ) + self.assertEqual( + layout.part("v").layer_offset_bytes(L), L * (k_row + v_row) + k_row + ) + + def test_page_bytes(self): + spec = _make_mha_spec("full", "up", layer_num=3, head_num=2, head_dim=4) + # page_bytes = page_size * entry_bytes (preserved invariant) + for ps in [1, 8, 64, 256]: + self.assertEqual(spec.page_bytes(ps), ps * spec.entry_bytes()) class TestMoveKVCacheNative(unittest.TestCase): def test_move_kv_cache_3d_path_unchanged(self): """The stock per-layer 3-D move must relocate exactly the named token - rows, byte-identically; compaction on static pools rides on it.""" + rows, byte-identically -- compaction on static pools rides on it.""" k = [torch.zeros((32, 2, 4), dtype=torch.float16) for _ in range(2)] v = [torch.zeros((32, 2, 4), dtype=torch.float16) for _ in range(2)] for L in range(2): diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index 2f82ba87d206..44d2a123ab76 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -1873,16 +1873,13 @@ def test_paged_pool_translate_helper_returns_physical_tokens(self): "v2p_page[virt_pages] * page_size + offsets.", ) - # The composite emits KERNEL-FACING ids, not the physical token ids this - # helper returns; they coincide only at multiplier 1, which nothing uses. - swa_mult = allocator.swa_kernel_page_multiplier - self.assertEqual(swa_mult, 2 * swa_spec.layer_num) + # The composite emits the swa sub-pool's PHYSICAL token ids (the + # kernel-facing id space of the token-major views). composite_out = allocator.translate_loc_from_full_to_swa(v_tokens) - expected_kernel = swa_phys_pages_direct * (PS * swa_mult) + offsets_in self.assertTrue( - bool((composite_out.long() == expected_kernel.long()).all().item()), + bool((composite_out.long() == expected.long()).all().item()), "REGRESSION: translate_loc_from_full_to_swa must emit the swa " - "sub-pool's kernel-facing ids (phys_page * ps * blocks_per_page + offset).", + "sub-pool's physical token ids (phys_page * ps + offset).", ) @@ -2426,7 +2423,7 @@ def test_grow_down_fast_path(self): class TestSWACompositeKernelIdSurface(unittest.TestCase): """The SWA composite's kernel-facing id surface. Attention backends probe for `translate_kv_loc_for_kernel` / `full_v2p_page_table`, and every id must - follow `kernel_id(t) = v2p[t // ps] * (ps * mult) + t % ps`.""" + follow `kernel_id(t) = v2p[t // ps] * ps + t % ps`.""" def setUp(self): # The code under test reads its config from the bags. @@ -2477,16 +2474,16 @@ def _build(self): ) def test_full_kernel_translate_matches_formula(self): - """Both sides scale by their OWN sub-pool's block count, and the full - kernel id follows v2p[t // ps] * (ps * mult) + t % ps.""" - mult = 2 * self.FULL_L + """Kernel-facing ids ARE the physical token ids under the token-major + views, so both sides pin at multiplier 1 and the id follows + v2p[t // ps] * ps + t % ps.""" a = self._build() - self.assertEqual(a.kernel_page_multiplier, 2 * self.FULL_L) - self.assertEqual(a.swa_kernel_page_multiplier, 2 * self.SWA_L) + self.assertEqual(a.kernel_page_multiplier, 1) + self.assertEqual(a.swa_kernel_page_multiplier, 1) v = a.alloc(3 * self.PS) self.assertIsNotNone(v) v2p = a.full_attn_allocator.virtual_to_physical - expected = v2p[v // self.PS] * (self.PS * mult) + v % self.PS + expected = v2p[v // self.PS] * self.PS + v % self.PS self.assertTrue(torch.equal(a.translate_kv_loc_for_kernel(v), expected)) # The PHYSICAL translate must stay unscaled -- compaction and the byte # machinery depend on it staying in physical space. @@ -2499,12 +2496,11 @@ def test_kernel_translate_accepts_an_int32_page_table(self): for ps in (1, 4): with self.subTest(page_size=ps): self.PS = ps - mult = 2 * self.FULL_L a = self._build() v = a.alloc(4 * ps) self.assertIsNotNone(v) v2p = a.full_attn_allocator.virtual_to_physical - expected = v2p[v // ps] * (ps * mult) + v % ps + expected = v2p[v // ps] * ps + v % ps page_table = v.to(torch.int32).view(2, -1) got = a.translate_kv_loc_for_kernel(page_table) self.assertEqual(got.shape, page_table.shape) @@ -2514,17 +2510,15 @@ def test_kernel_translate_accepts_an_int32_page_table(self): a.translate_kv_loc_for_kernel(page_table, out=dst) self.assertTrue(torch.equal(dst.reshape(-1), expected)) - def test_swa_translate_scales_page_stride(self): - mult = 2 * self.SWA_L + def test_swa_translate_matches_formula(self): a = self._build() v = a.alloc(3 * self.PS) self.assertIsNotNone(v) v2p_swa = a.swa_attn_allocator.virtual_to_physical - expected = v2p_swa[v // self.PS] * (self.PS * mult) + v % self.PS + expected = v2p_swa[v // self.PS] * self.PS + v % self.PS self.assertTrue(torch.equal(a.translate_loc_from_full_to_swa(v), expected)) def test_swa_transfer_page_is_physical_not_kernel_scaled(self): - mult = 2 * self.SWA_L a = self._build() v = a.alloc(3 * self.PS) self.assertIsNotNone(v) @@ -2540,17 +2534,17 @@ def test_swa_transfer_page_is_physical_not_kernel_scaled(self): physical_pages.tolist(), ) + # The read path is physical too: the token-major entry leaves no + # per-page block scale for a kernel id to carry. kernel_tokens = a.translate_loc_from_full_to_swa(v) self.assertEqual( kv_to_page_indices(kernel_tokens, self.PS).tolist(), - (physical_pages * mult).tolist(), + physical_pages.tolist(), ) def test_swa_kernel_tombstone_still_lands_on_sink(self): - """The scaled stride must not break the tombstone clamp: a tombstoned - page's ids (v2p == -1 -> -stride + offset, negative for every in-page - offset) still land on the sink, never negative.""" - mult = 2 * self.SWA_L + """A tombstoned page's ids (v2p == -1 -> -ps + offset, negative for + every in-page offset) still land on the sink, never negative.""" a = self._build() v = a.alloc(2 * self.PS) self.assertIsNotNone(v) @@ -2624,8 +2618,8 @@ class _FakeHybridLinearKVPool: def test_construction_alloc_and_kernel_formula(self): a = self._build() - # MLA: one latent row per layer, so the spec reports LAYERS blocks. - self.assertEqual(a.kernel_page_multiplier, self.LAYERS) + # Token-major views: the kernel id is the physical token id. + self.assertEqual(a.kernel_page_multiplier, 1) v = a.alloc(2 * self.PS) self.assertIsNotNone(v, "2-page alloc infeasible at ps=64") # Page-aligned virtual run (page-granular allocator invariant). @@ -2633,7 +2627,7 @@ def test_construction_alloc_and_kernel_formula(self): # The kernel translate follows the affine formula at ps=64, and every id # fits int32 (the canonical narrows on store). v2p = a.full_v2p_page_table - want = v2p[v // self.PS] * (self.PS * self.LAYERS) + v % self.PS + want = v2p[v // self.PS] * self.PS + v % self.PS got = a.translate_kv_loc_for_kernel(v) self.assertTrue(torch.equal(got, want), "kernel-facing formula broke at ps=64") self.assertTrue(bool((got < 2**31).all().item())) diff --git a/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py b/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py index 8ecfb9676711..58d3d08c9a87 100644 --- a/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py +++ b/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py @@ -18,8 +18,9 @@ import torch from sglang.srt.mem_cache.layout.page_major import ( - build_mha_views, - build_mla_views, + DenseEntryLayout, + DensePart, + build_dense_views, build_page_major_mamba_views, mamba_entry_bytes, mha_entry_bytes, @@ -33,29 +34,34 @@ class TestMLAEnvelopeTransferAddressing(CustomTestCase): def test_page_envelope_matches_per_layer_views(self): - """Every (page, layer, slot) row written through the MLA views - must land at raw_ptr + page * page_envelope_bytes + layer-block offset, - i.e. inside the page's transfer envelope.""" + """Every (page, layer, slot) row written through the MLA views must + land at raw_ptr + page * page_envelope_bytes + slot * entry_bytes + + layer * row_bytes, i.e. inside the page's transfer envelope.""" layer_num, page_size, kv_dim, num_pages = 3, 4, 8, 6 store_dtype = torch.bfloat16 row_bytes = kv_dim * store_dtype.itemsize - page_bytes = page_size * layer_num * row_bytes - self.assertEqual( - page_bytes, - page_size - * mla_entry_bytes( - layer_num=layer_num, - kv_cache_dim=kv_dim, - itemsize=store_dtype.itemsize, + entry_bytes = mla_entry_bytes( + layer_num=layer_num, kv_cache_dim=kv_dim, itemsize=store_dtype.itemsize + ) + page_bytes = page_size * entry_bytes + layout = DenseEntryLayout( + entry_bytes=entry_bytes, + parts=( + DensePart( + name="kv", + offset_bytes=0, + layer_stride_bytes=row_bytes, + layer_num=layer_num, + row_shape=(1, kv_dim), + dtype=store_dtype, + ), ), ) - # +1 page envelope of tail pad, as UnifiedKVPool allocates for MLA. - raw = torch.zeros((num_pages + 1) * page_bytes, dtype=torch.uint8) - views = build_mla_views( + raw = torch.zeros(num_pages * page_bytes, dtype=torch.uint8) + views = build_dense_views( raw, - layer_num=layer_num, - kv_cache_dim=kv_dim, - store_dtype=store_dtype, + layout=layout, + part=layout.part("kv"), page_size=page_size, num_pages=num_pages, anchor_bytes=0, @@ -64,14 +70,10 @@ def test_page_envelope_matches_per_layer_views(self): for page in range(num_pages): for layer in range(layer_num): for off in range(page_size): - kernel_id = page * layer_num * page_size + off + token = page * page_size + off val = torch.randn(kv_dim, dtype=store_dtype) - views[layer][kernel_id, 0] = val - start = ( - page * page_bytes - + layer * page_size * row_bytes - + off * row_bytes - ) + views[layer][token, 0] = val + start = page * page_bytes + off * entry_bytes + layer * row_bytes got = raw[start : start + row_bytes].view(store_dtype) self.assertTrue(torch.equal(got, val), (page, layer, off)) @@ -79,8 +81,9 @@ def test_page_envelope_matches_per_layer_views(self): class TestMHAEnvelopeTransferAddressing(CustomTestCase): """The MHA counterpart of the MLA case above. - An MHA page envelope holds ``2 * layer_num`` row-blocks (layer l's K at - block 2l, its V at 2l+1). PD ships that whole envelope as one item, so a + An MHA page envelope holds ``page_size`` slot entries, each carrying every + layer's K and V at a fixed offset. PD ships that whole envelope as one + item, so a row written through ANY per-layer view must land inside its own page's ``page_envelope_bytes`` block -- otherwise the transfer would carry a page's K but another page's V and every kernel would still read fine @@ -101,26 +104,45 @@ def test_page_envelope_matches_per_layer_views(self): row_bytes = head_num * head_dim * store_dtype.itemsize self.assertEqual(page_bytes, page_size * 2 * layer_num * row_bytes) - # One page envelope of tail pad, as UnifiedKVPool allocates for MHA. - raw = torch.zeros((num_pages + 1) * page_bytes, dtype=torch.uint8) - k_views, v_views = build_mha_views( - raw, - layer_num=layer_num, - head_num=head_num, - head_dim=head_dim, - v_head_dim=head_dim, - store_dtype=store_dtype, - page_size=page_size, - num_pages=num_pages, - anchor_bytes=0, + layout = DenseEntryLayout( + entry_bytes=entry_bytes, + parts=( + DensePart( + name="k", + offset_bytes=0, + layer_stride_bytes=2 * row_bytes, + layer_num=layer_num, + row_shape=(head_num, head_dim), + dtype=store_dtype, + ), + DensePart( + name="v", + offset_bytes=row_bytes, + layer_stride_bytes=2 * row_bytes, + layer_num=layer_num, + row_shape=(head_num, head_dim), + dtype=store_dtype, + ), + ), + ) + raw = torch.zeros(num_pages * page_bytes, dtype=torch.uint8) + k_views, v_views = ( + build_dense_views( + raw, + layout=layout, + part=layout.part(name), + page_size=page_size, + num_pages=num_pages, + anchor_bytes=0, + ) + for name in ("k", "v") ) - blocks = 2 * layer_num for page in range(num_pages): for layer in range(layer_num): for is_v, views in ((0, k_views), (1, v_views)): for pos in range(page_size): - row = page * blocks * page_size + pos + row = page * page_size + pos views[layer][row].fill_(1) (nz,) = torch.nonzero(raw, as_tuple=True) lo, hi = int(nz.min()), int(nz.max()) @@ -142,7 +164,7 @@ def test_envelope_move_is_a_whole_page_copy(self): """Relocating a page envelope must move every layer's K and V with it; this is what `UnifiedMHATokenToKVPool.move_kv_cache` relies on and what makes a physical page id a valid PD transfer index after compaction.""" - layer_num, page_size, head_num, head_dim, num_pages = 2, 2, 1, 4, 4 + layer_num, page_size, head_num, head_dim, num_pages = 2, 2, 1, 8, 4 store_dtype = torch.bfloat16 entry_bytes = mha_entry_bytes( layer_num=layer_num, @@ -152,23 +174,44 @@ def test_envelope_move_is_a_whole_page_copy(self): itemsize=store_dtype.itemsize, ) page_bytes = page_size * entry_bytes - raw = torch.zeros((num_pages + 1) * page_bytes, dtype=torch.uint8) - k_views, v_views = build_mha_views( - raw, - layer_num=layer_num, - head_num=head_num, - head_dim=head_dim, - v_head_dim=head_dim, - store_dtype=store_dtype, - page_size=page_size, - num_pages=num_pages, - anchor_bytes=0, + row_bytes = head_num * head_dim * store_dtype.itemsize + layout = DenseEntryLayout( + entry_bytes=entry_bytes, + parts=( + DensePart( + name="k", + offset_bytes=0, + layer_stride_bytes=2 * row_bytes, + layer_num=layer_num, + row_shape=(head_num, head_dim), + dtype=store_dtype, + ), + DensePart( + name="v", + offset_bytes=row_bytes, + layer_stride_bytes=2 * row_bytes, + layer_num=layer_num, + row_shape=(head_num, head_dim), + dtype=store_dtype, + ), + ), + ) + raw = torch.zeros(num_pages * page_bytes, dtype=torch.uint8) + k_views, v_views = ( + build_dense_views( + raw, + layout=layout, + part=layout.part(name), + page_size=page_size, + num_pages=num_pages, + anchor_bytes=0, + ) + for name in ("k", "v") ) - blocks = 2 * layer_num # Distinct content in source page 1, every layer, K and V. for layer in range(layer_num): for pos in range(page_size): - row = 1 * blocks * page_size + pos + row = 1 * page_size + pos k_views[layer][row].fill_(layer + 1) v_views[layer][row].fill_(-(layer + 1)) @@ -177,7 +220,7 @@ def test_envelope_move_is_a_whole_page_copy(self): for layer in range(layer_num): for pos in range(page_size): - row = 3 * blocks * page_size + pos + row = 3 * page_size + pos self.assertTrue( torch.all(k_views[layer][row] == layer + 1), f"K layer {layer} did not ride the envelope move", diff --git a/test/registered/unit/mem_cache/test_prefill_memory_budget.py b/test/registered/unit/mem_cache/test_prefill_memory_budget.py index 2a73dfffc9e3..b1dbad1abf54 100644 --- a/test/registered/unit/mem_cache/test_prefill_memory_budget.py +++ b/test/registered/unit/mem_cache/test_prefill_memory_budget.py @@ -30,25 +30,34 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu") -def _shared_allocator(*, page_size=4, total_bytes=1024): +def _shared_allocator(*, page_size=4, total_bytes=2048, pool_tokens=None): + # `pool_tokens` sizes both sub-pools directly, for a test whose arithmetic + # needs an exact token count rather than whatever a byte budget buys. + sized = ( + dict( + full_max_total_num_tokens=pool_tokens, swa_max_total_num_tokens=pool_tokens + ) + if pool_tokens is not None + else dict(total_bytes=total_bytes) + ) return init_unified_swa_pools( device="cpu", kv_cache_dtype=torch.float16, head_num=1, - head_dim=4, - v_head_dim=4, + head_dim=8, + v_head_dim=8, swa_head_num=1, - swa_head_dim=4, - swa_v_head_dim=4, + swa_head_dim=8, + swa_v_head_dim=8, page_size=page_size, start_layer=0, end_layer=2, swa_attention_layer_ids=[1], full_attention_layer_ids=[0], - total_bytes=total_bytes, enable_memory_saver=False, need_sort=False, lazy_compaction=True, + **sized, ).token_to_kv_pool_allocator @@ -148,8 +157,11 @@ def test_common_eviction_dispatches_joint_reclaim(self): class TestSharedPrefillAdmission(unittest.TestCase): def _new_admission(self, page_size, pool_pages, *, ignore_eos=False): allocator = _shared_allocator( - page_size=page_size, total_bytes=pool_pages * page_size * 16 + page_size=page_size, pool_tokens=pool_pages * page_size ) + # The admission arithmetic below is tuned to this pool; a resize would + # measure something else, so pin the geometry rather than assume it. + self.assertEqual(allocator.size_full, pool_pages * page_size) req = Req( rid="unaligned-prompt", origin_input_text=None, @@ -191,7 +203,7 @@ def _new_admission(self, page_size, pool_pages, *, ignore_eos=False): def test_unaligned_final_chunk_makes_progress(self): for page_size in (4, 64): with self.subTest(page_size=page_size): - allocator, req, adder = self._new_admission(page_size, pool_pages=7) + allocator, req, adder = self._new_admission(page_size, pool_pages=6) req.prefix_indices = allocator.alloc(page_size) self.assertIsNotNone(req.prefix_indices) self.assertTrue(allocator.can_reserve(page_size + 2, page_size + 2)) @@ -204,7 +216,7 @@ def test_unaligned_ignore_eos_enters_empty_pool(self): for page_size in (4, 64): with self.subTest(page_size=page_size): allocator, req, adder = self._new_admission( - page_size, pool_pages=6, ignore_eos=True + page_size, pool_pages=5, ignore_eos=True ) self.assertEqual(len(req.prefix_indices), 0) self.assertTrue(allocator.can_reserve(2 * page_size + 2, 2 * page_size)) diff --git a/test/registered/unit/mem_cache/test_unified_handout_zeroing.py b/test/registered/unit/mem_cache/test_unified_handout_zeroing.py index 3bedc70b7c32..3691a577c30d 100644 --- a/test/registered/unit/mem_cache/test_unified_handout_zeroing.py +++ b/test/registered/unit/mem_cache/test_unified_handout_zeroing.py @@ -20,7 +20,7 @@ BF16_NAN = 0x7FC1 # LE bf16 NaN bit pattern, as SGLANG_DEBUG_POISON_POOL fills -def _build(device, page_size=1, kernel_page_multiplier=None): +def _build(device, page_size=1): """A tiny MLA+mamba unified pool + full-side allocator. Mirrors init_unified_mamba_pools' construction just enough for the @@ -64,9 +64,6 @@ def _build(device, page_size=1, kernel_page_multiplier=None): device=device, is_id_owner=True, page_size=page_size, - kernel_page_multiplier=( - layer_num if kernel_page_multiplier is None else kernel_page_multiplier - ), ) return buf, kvcache, allocator @@ -116,12 +113,11 @@ def test_fresh_and_recycled_pages_zeroed(self): pages2 = self._phys_pages(allocator, out2) self.assertTrue((env[pages2] == 0).all().item()) - def test_zeroing_enabled_for_single_layer_multiplier(self): - # A shard owning exactly ONE full-attention MLA layer has - # kernel_page_multiplier == 1 but its pool is still - # UnifiedMLATokenToKVPool with the same NaN-unsafe partial-page - # reads — zeroing must key on the pool type, not on multiplier > 1. - buf, kvcache, allocator = _build("cuda", kernel_page_multiplier=1) + def test_zeroing_keys_on_pool_type(self): + # UnifiedMLATokenToKVPool has NaN-unsafe partial-page reads whatever + # its geometry; zeroing must key on the pool type, not on any + # id-space scale. + buf, kvcache, allocator = _build("cuda") self.assertTrue(allocator._zero_pages_on_alloc) self._poison(buf) out = allocator.alloc(8) diff --git a/test/registered/unit/mem_cache/test_unified_mha_views.py b/test/registered/unit/mem_cache/test_unified_mha_views.py index cd0d66d0c228..9d0fa45ed7cc 100644 --- a/test/registered/unit/mem_cache/test_unified_mha_views.py +++ b/test/registered/unit/mem_cache/test_unified_mha_views.py @@ -11,19 +11,31 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""MHA K/V views for the unified memory pool (uniform-row hybrid models), CPU-only. +"""MHA K/V views for the unified memory pool (token-major dense views). + +Covers, CPU-only (pure torch -- no GPU / Triton kernels): + - `MHASubPoolSpec.layout()`: K and V of every layer sit at fixed offsets + inside one slot entry, K and V rows may differ in width, alignment is + enforced at construction; + - `build_dense_views` addressing: view_l[t] for PHYSICAL token t must land + exactly at the byte the envelope formula assigns to (page, slot, layer, + K|V) -- cross-checked against an independent 4-D (page, slot) description; + - K and V of one token share ONE id (the per-layer origin shift does the + disambiguation), with no aliasing across the 2L views; + - a too-short buffer and misaligned rows fail loud at construction. Addressing law under test: - kernel_id(t) = (t // ps) * (ps * 2L) + t % ps - K of layer l at block 2l, V at block 2l+1; blocks are ps rows of - head_num*head_dim elements, at offsets identical to - MHASubPoolSpec.layer_k/v_offset_in_page when rows are uniform. + byte(t, l, K) = t * entry_bytes + l * (k_row + v_row) + byte(t, l, V) = byte(t, l, K) + k_row + t = page * page_size + slot (the physical token id IS the kernel id) + + python -m pytest test/registered/unit/mem_cache/test_unified_mha_views.py -v """ from sglang.test.ci.ci_register import register_cpu_ci -register_cpu_ci(est_time=11, suite="base-a-test-cpu") +register_cpu_ci(est_time=8, suite="base-a-test-cpu") import unittest from types import SimpleNamespace @@ -31,7 +43,12 @@ import torch from sglang.srt.environ import envs -from sglang.srt.mem_cache.layout.page_major import build_mha_views +from sglang.srt.mem_cache.layout.page_major import ( + ENTRY_ALIGN_BYTES, + build_dense_views, + mha_entry_bytes, + paged_view, +) from sglang.srt.mem_cache.unified_memory_pool import ( MHASubPoolSpec, UnifiedKVPool, @@ -40,18 +57,24 @@ _DEV = "cpu" # `set_kv_buffer` dispatches on the PLATFORM (memory_pool._is_cuda, resolved at -# import), not on the tensors it is handed, so cases driving it build there. +# import), not on the tensors it is handed, so cases driving it must build on +# the platform's device. The rest of this file is byte arithmetic, so CPU. _STORE_DEV = "cuda" if torch.cuda.is_available() else "cpu" -# Geometry kept tiny so every byte offset is hand-checkable. +# Small-but-nontrivial MHA geometry: L=2 layers, H=2 heads, D=4, so every byte +# offset is hand-checkable. One entry = 2 layers x (K row + V row) = 64 B. _L = 2 _H = 2 _D = 4 -_ROW = _H * _D # row elements _DTYPE = torch.bfloat16 _ITEM = _DTYPE.itemsize -_BLOCKS = 2 * _L _HAS_FP8 = hasattr(torch, "float8_e4m3fn") +_K_ROW = _H * _D * _ITEM +_V_ROW = _H * _D * _ITEM +_ENTRY = mha_entry_bytes( + layer_num=_L, head_num=_H, head_dim=_D, v_head_dim=_D, itemsize=_ITEM +) +assert _ENTRY == _L * (_K_ROW + _V_ROW) # no alignment pad at this geometry def _mha_spec(head_dim=_D, v_head_dim=None, layer_num=_L, grow="down"): @@ -66,150 +89,244 @@ def _mha_spec(head_dim=_D, v_head_dim=None, layer_num=_L, grow="down"): ) -def _kernel_id(t, ps): - return (t // ps) * (ps * _BLOCKS) + t % ps +def _make_raw(ps, num_pages, entry=_ENTRY, short=0): + n = num_pages * ps * entry - short + return torch.zeros(n, dtype=torch.uint8, device=_DEV) -def _make_raw(ps, num_pages, pad_pages=1): - page_bytes = ps * _BLOCKS * _ROW * _ITEM - raw = torch.zeros( - (num_pages + pad_pages) * page_bytes, dtype=torch.uint8, device=_DEV +def _build_views(raw, ps, num_pages, head_dim=_D, v_head_dim=None, layer_num=_L): + layout = _mha_spec(head_dim, v_head_dim, layer_num).layout() + kw = dict(layout=layout, page_size=ps, num_pages=num_pages) + return ( + build_dense_views(raw, part=layout.part("k"), **kw), + build_dense_views(raw, part=layout.part("v"), **kw), ) - return raw -def _build_views(raw, ps, num_pages, head_dim=_D, v_head_dim=_D, layer_num=_L): - return build_mha_views( - raw, - layer_num=layer_num, +def _reference_paged_views( + raw, *, page_size, num_pages, anchor_bytes=0, head_dim=_D, v_head_dim=_D +): + """Independent 4-D description of the token-major envelope: per-layer + ``(num_pages, page_size, head_num, head_dim)`` views addressed by + ``(page, slot)``, so the flat builder's addressing can be cross-checked + against a second, independently derived description of the same bytes.""" + k_row = _H * head_dim * _ITEM + v_row = _H * v_head_dim * _ITEM + entry = mha_entry_bytes( + layer_num=_L, head_num=_H, head_dim=head_dim, v_head_dim=v_head_dim, - store_dtype=_DTYPE, - page_size=ps, - num_pages=num_pages, + itemsize=_ITEM, ) - - -def _reference_strided_views(raw, *, page_size, num_pages, anchor_bytes=0): - """Independent 4-D strided description of the same page-major envelope, - addressed by ``(page, slot)`` -- the oracle for the view builder.""" - k_row_bytes = _ROW * _ITEM - v_row_bytes = _ROW * _ITEM - page_bytes = page_size * _L * (k_row_bytes + v_row_bytes) as_dtype_view = raw.view(_DTYPE) - k_stride = (page_bytes // _ITEM, k_row_bytes // _ITEM, _D, 1) - v_stride = (page_bytes // _ITEM, v_row_bytes // _ITEM, _D, 1) - shape = (num_pages, page_size, _H, _D) k_views, v_views = [], [] for layer in range(_L): - k_base = anchor_bytes + layer * page_size * (k_row_bytes + v_row_bytes) - v_base = k_base + page_size * k_row_bytes - k_views.append( - torch.as_strided( - as_dtype_view, - size=shape, - stride=k_stride, - storage_offset=k_base // _ITEM, - ) - ) - v_views.append( - torch.as_strided( - as_dtype_view, - size=shape, - stride=v_stride, - storage_offset=v_base // _ITEM, + k_base = anchor_bytes + layer * (k_row + v_row) + v_base = k_base + k_row + parts = ((k_base, head_dim, k_views), (v_base, v_head_dim, v_views)) + for base, dim, out in parts: + out.append( + torch.as_strided( + as_dtype_view, + size=(num_pages, page_size, _H, dim), + stride=(page_size * entry // _ITEM, entry // _ITEM, dim, 1), + storage_offset=base // _ITEM, + ) ) - ) return k_views, v_views class TestMHASpecSurface(unittest.TestCase): - def test_asymmetric_rows_refused_by_the_view_builder(self): - """ServerArgs screens asymmetric-KV models out of - --enable-unified-memory; this guards a caller reaching the builder.""" + def test_layout_parts_match_the_entry_formula(self): + """Every per-layer view is built through the layout descriptor, so its + part offsets are pinned against the closed-form entry math. The page + size never enters (a page is page_size entries back to back).""" spec = _mha_spec() - raw = torch.zeros(1 << 16, dtype=torch.uint8) - with self.assertRaises(AssertionError): - build_mha_views( - raw, - layer_num=spec.layer_num, - head_num=spec.head_num, - head_dim=6, - v_head_dim=4, - store_dtype=spec.store_dtype, - page_size=1, - num_pages=4, + layout = spec.layout() + self.assertEqual(layout.entry_bytes, spec.entry_bytes()) + for l in range(_L): + self.assertEqual( + layout.part("k").layer_offset_bytes(l), l * (_K_ROW + _V_ROW) + ) + self.assertEqual( + layout.part("v").layer_offset_bytes(l), l * (_K_ROW + _V_ROW) + _K_ROW ) + def test_entry_bytes_matches_layout_helper_and_is_aligned(self): + spec = _mha_spec() + self.assertEqual( + spec.entry_bytes(), + mha_entry_bytes( + layer_num=_L, head_num=_H, head_dim=_D, v_head_dim=_D, itemsize=_ITEM + ), + ) + # 48 B of rows round up to one 64 B entry; the parts still fit inside. + padded = MHASubPoolSpec( + name="full", + layer_num=1, + head_num=1, + head_dim=16, + v_head_dim=8, + store_dtype=_DTYPE, + grow_direction="down", + ) + self.assertEqual(padded.entry_bytes(), ENTRY_ALIGN_BYTES * 2) + self.assertEqual(padded.entry_bytes() % ENTRY_ALIGN_BYTES, 0) + padded.layout() # validates + + def test_asymmetric_rows_are_admitted(self): + """K and V are two parts of one entry at their own offsets; they no + longer have to be the same kind of row (the MiMoV2 shape, scaled).""" + spec = _mha_spec(head_dim=8, v_head_dim=4) + layout = spec.layout() + self.assertEqual(layout.part("k").row_bytes(), _H * 8 * _ITEM) + self.assertEqual(layout.part("v").row_bytes(), _H * 4 * _ITEM) + self.assertEqual(layout.part("v").offset_bytes, _H * 8 * _ITEM) + + def test_misaligned_rows_fail_loud(self): + """A row that is not a multiple of 16 B would break the vector stores + every write kernel issues, so the layout refuses it at construction.""" + with self.assertRaises(AssertionError): + _mha_spec(head_dim=6, v_head_dim=6).layout() + class TestMHAViews(unittest.TestCase): - def test_view_shapes_are_stock_mha(self): + def test_view_shapes_and_strides(self): ps, num_pages = 4, 6 k_views, v_views = _build_views(_make_raw(ps, num_pages), ps, num_pages) - n_rows = num_pages * _BLOCKS * ps + n_rows = num_pages * ps self.assertEqual(len(k_views), _L) self.assertEqual(len(v_views), _L) for v in (*k_views, *v_views): - # The stock MHATokenToKVPool per-layer signature: 3-D, packed rows. + # The stock MHATokenToKVPool per-layer signature: 3-D, one row per + # physical token, the slot stride being the whole entry. self.assertEqual(tuple(v.shape), (n_rows, _H, _D)) - self.assertEqual(v.stride(), (_ROW, _D, 1)) + self.assertEqual(v.stride(), (_ENTRY // _ITEM, _D, 1)) + for l in range(_L): + self.assertEqual( + (v_views[l].storage_offset() - k_views[l].storage_offset()) * _ITEM, + _K_ROW, + ) - def test_addressing_matches_strided_reference(self): - """Cross-readback both ways: a write through the strided oracle at - (page, slot) must read back through the view at kernel_id(t).""" + def test_addressing_matches_paged_reference(self): + """Cross-readback: bytes written through the reference (page, slot) + views must be read back through the flat views at t = page*ps+slot, + for both K and V of every layer -- and vice versa.""" for ps in (1, 4): num_pages = 5 raw = _make_raw(ps, num_pages) - sk, sv = _reference_strided_views(raw, page_size=ps, num_pages=num_pages) + sk, sv = _reference_paged_views(raw, page_size=ps, num_pages=num_pages) dk, dv = _build_views(raw, ps, num_pages) probes = [(0, 0, 0), (1, 1, ps - 1), (4, 0, ps // 2), (3, 1, 0)] - # strided-write -> view-read for p, l, s in probes: t = p * ps + s - d = _kernel_id(t, ps) sk[l][p, s] = float(p * 100 + l * 10 + s + 1) sv[l][p, s] = float(p * 100 + l * 10 + s + 2) self.assertTrue( - torch.all(dk[l][d] == float(p * 100 + l * 10 + s + 1)), + torch.all(dk[l][t] == float(p * 100 + l * 10 + s + 1)), f"K (p={p}, l={l}, s={s}, ps={ps}) view readback off-formula", ) self.assertTrue( - torch.all(dv[l][d] == float(p * 100 + l * 10 + s + 2)), + torch.all(dv[l][t] == float(p * 100 + l * 10 + s + 2)), f"V (p={p}, l={l}, s={s}, ps={ps}) view readback off-formula", ) - # view-write -> strided-read for p, l, s in probes: t = p * ps + s - d = _kernel_id(t, ps) - dk[l][d] = float(p * 100 + l * 10 + s + 3) - dv[l][d] = float(p * 100 + l * 10 + s + 4) - self.assertTrue( - torch.all(sk[l][p, s] == float(p * 100 + l * 10 + s + 3)) - ) - self.assertTrue( - torch.all(sv[l][p, s] == float(p * 100 + l * 10 + s + 4)) - ) - - def test_k_and_v_share_one_kernel_id_without_aliasing(self): - """One kernel-facing id, 2L distinct cells (K and V of every layer): writes - through all 2L views at the SAME id must not clobber each other.""" + dk[l][t] = float(p * 100 + l * 10 + s + 3) + dv[l][t] = float(p * 100 + l * 10 + s + 4) + want_k = float(p * 100 + l * 10 + s + 3) + want_v = float(p * 100 + l * 10 + s + 4) + self.assertTrue(torch.all(sk[l][p, s] == want_k)) + self.assertTrue(torch.all(sv[l][p, s] == want_v)) + + def test_byte_addresses_match_envelope_formula(self): + """The per-layer view's byte address for token ``t``, layer ``l`` must + equal the hand-computed envelope formula. Independent of any view + builder -- this is the raw layout contract every envelope consumer + (moves, sizing, transfer math) relies on.""" + for ps in (1, 4): + num_pages = 5 + dk, dv = _build_views(_make_raw(ps, num_pages), ps, num_pages) + for t in (0, 1, ps, 3 * ps + (ps - 1), 4 * ps): + for l in range(_L): + expected_k = t * _ENTRY + l * (_K_ROW + _V_ROW) + expected_v = expected_k + _K_ROW + got_k = (dk[l].storage_offset() + t * dk[l].stride(0)) * _ITEM + got_v = (dv[l].storage_offset() + t * dv[l].stride(0)) * _ITEM + self.assertEqual(got_k, expected_k, f"K t={t} l={l} ps={ps}") + self.assertEqual(got_v, expected_v, f"V t={t} l={l} ps={ps}") + + def test_k_and_v_share_one_id_without_aliasing(self): + """One id, 2L distinct cells (K and V of every layer): writes through + all 2L views at the SAME id must not clobber each other.""" ps, num_pages = 4, 4 dk, dv = _build_views(_make_raw(ps, num_pages), ps, num_pages) t = 2 * ps + 1 # page 2, slot 1 - d = _kernel_id(t, ps) for l in range(_L): - dk[l][d] = float(2 * l + 1) - dv[l][d] = float(2 * l + 2) + dk[l][t] = float(2 * l + 1) + dv[l][t] = float(2 * l + 2) for l in range(_L): - self.assertTrue(torch.all(dk[l][d] == float(2 * l + 1))) - self.assertTrue(torch.all(dv[l][d] == float(2 * l + 2))) + self.assertTrue(torch.all(dk[l][t] == float(2 * l + 1))) + self.assertTrue(torch.all(dv[l][t] == float(2 * l + 2))) + + def test_asymmetric_views_address_their_own_rows(self): + head_dim, v_head_dim = 8, 4 + k_row, v_row = _H * head_dim * _ITEM, _H * v_head_dim * _ITEM + entry = mha_entry_bytes( + layer_num=_L, + head_num=_H, + head_dim=head_dim, + v_head_dim=v_head_dim, + itemsize=_ITEM, + ) + ps, num_pages = 4, 3 + raw = _make_raw(ps, num_pages, entry=entry) + dk, dv = _build_views(raw, ps, num_pages, head_dim, v_head_dim) + self.assertEqual(tuple(dk[0].shape[1:]), (_H, head_dim)) + self.assertEqual(tuple(dv[0].shape[1:]), (_H, v_head_dim)) + t = 2 * ps + 3 + for l in range(_L): + dk[l][t] = float(l + 1) + dv[l][t] = float(l + 11) + flat = raw.view(_DTYPE) + for l in range(_L): + k0 = (t * entry + l * (k_row + v_row)) // _ITEM + v0 = k0 + k_row // _ITEM + self.assertTrue(torch.all(flat[k0 : k0 + k_row // _ITEM] == float(l + 1))) + self.assertTrue(torch.all(flat[v0 : v0 + v_row // _ITEM] == float(l + 11))) + + def test_views_fill_the_buffer_exactly(self): + """No tail pad: the last view's last byte is the buffer's last byte.""" + ps, num_pages = 2, 4 + raw = _make_raw(ps, num_pages) + _, dv = _build_views(raw, ps, num_pages) + last = dv[_L - 1] + end = (last.storage_offset() + (last.shape[0] - 1) * last.stride(0)) * _ITEM + self.assertEqual(end + _V_ROW, raw.numel()) - def test_missing_tail_pad_fails_loud(self): + def test_short_buffer_fails_loud(self): ps, num_pages = 2, 4 - raw = _make_raw(ps, num_pages, pad_pages=0) with self.assertRaises(AssertionError): - _build_views(raw, ps, num_pages) + _build_views(_make_raw(ps, num_pages, short=1), ps, num_pages) + + def test_paged_view_regroups_by_page(self): + """BUG REGRESSION at page_size 1: a `view`-built split gives the size-1 + slot dim the ROW stride, not the entry stride, so the paged view + misreported the slot stride of every MHA/SWA layer at ps=1.""" + num_pages = 3 + for ps in (1, 4): + dk, _ = _build_views(_make_raw(ps, num_pages), ps, num_pages) + paged = paged_view(dk[1], ps) + self.assertEqual( + tuple(paged.stride()), + (ps * _ENTRY // _ITEM, _ENTRY // _ITEM, _D, 1), + ps, + ) + for t in range(num_pages * ps): + self.assertEqual( + paged[t // ps, t % ps].data_ptr(), dk[1][t].data_ptr(), (ps, t) + ) # ---- pool level ---- @@ -244,16 +361,27 @@ def _make_pool(ps=1, full_spec=None, device=_DEV): class TestUnifiedKVPoolViews(unittest.TestCase): - def test_every_mha_sub_pool_is_per_layer_contiguous(self): + def test_every_mha_sub_pool_is_a_slot_strided_view(self): """The unified pool has ONE MHA layout: both sub-pools come back as - stock 3-D per-layer views, whatever their page size.""" + 3-D per-layer views whose slot stride is their own entry.""" for ps in (1, 4): pool = _make_pool(ps=ps) - for name in ("full", "swa"): + for name, spec in (("full", _mha_spec()), ("swa", _swa_spec())): k, v = pool.mha_views_for(name) - self.assertEqual(k[0].dim(), 3, f"{name} K at ps={ps}") - self.assertEqual(v[0].dim(), 3, f"{name} V at ps={ps}") - self.assertTrue(k[0].is_contiguous()) + for t in (k[0], v[0], k[-1], v[-1]): + self.assertEqual(t.dim(), 3, f"{name} at ps={ps}") + self.assertEqual(t.stride(0) * _ITEM, spec.entry_bytes()) + + def test_raw_is_exactly_the_budget(self): + """The views end at the last slot, so the pool allocates exactly the + byte budget: no tail pad, nothing to under-allocate.""" + for ps in (1, 4): + kv = _make_pool(ps) + self.assertEqual( + kv._raw.numel(), + _mha_spec().entry_bytes() * _N_FULL + + _swa_spec().entry_bytes() * _N_SWA, + ) def _layer(l): @@ -271,21 +399,23 @@ def _make_pool_and_kv(ps, device=_DEV): class TestUnifiedMHATokenToKVPool(unittest.TestCase): - def test_size_is_view_row_bound(self): + def test_size_is_slot_bound(self): """`size` drives BOTH the python OOB check and the store kernel's - device-side size_limit; it must be the view row bound, not slot count.""" + device-side size_limit; `size + page_size` must be the view row count.""" for ps in (1, 4): unified_kv, pool_under_test = _make_pool_and_kv(ps) - n_rows = (unified_kv.max_slots("full") // ps) * _BLOCKS * ps + n_rows = (unified_kv.max_slots("full") // ps) * ps self.assertEqual(pool_under_test.size, n_rows - ps) + self.assertEqual(pool_under_test.k_buffer[0].shape[0], n_rows) def test_stock_write_lands_on_envelope_truth(self): - """Byte-identity: the inherited `set_kv_buffer` at kernel-facing locs must - produce the same bytes as writes through the strided oracle.""" + """Byte-identity: the pool's stock inherited `set_kv_buffer` at physical + token ids must produce exactly the bytes that direct writes through + the reference (page, slot) views produce at the same cells -- pins the + whole write path (loc -> strided view -> raw bytes) end to end.""" for ps in (1, 4): kv, pool = _make_pool_and_kv(ps, device=_STORE_DEV) - # An independent strided view of the SAME sub-pool region. - sk, sv = _reference_strided_views( + sk, sv = _reference_paged_views( kv._raw, page_size=ps, num_pages=kv.max_slots("full") // ps, @@ -296,20 +426,10 @@ def test_stock_write_lands_on_envelope_truth(self): toks = torch.tensor( [p * ps + s for (p, s) in probes], device=_STORE_DEV ) - kernel_locs = (toks // ps) * (ps * _BLOCKS) + toks % ps - k = torch.full( - (len(probes), _H, _D), - float(l + 1), - dtype=_DTYPE, - device=_STORE_DEV, - ) - v = torch.full( - (len(probes), _H, _D), - float(l + 101), - dtype=_DTYPE, - device=_STORE_DEV, - ) - pool.set_kv_buffer(_layer(l), kernel_locs, k, v) + shape = (len(probes), _H, _D) + k = torch.full(shape, float(l + 1), dtype=_DTYPE, device=_STORE_DEV) + v = torch.full(shape, float(l + 101), dtype=_DTYPE, device=_STORE_DEV) + pool.set_kv_buffer(_layer(l), toks, k, v) for p, s in probes: self.assertTrue( torch.all(sk[l][p, s] == float(l + 1)), @@ -321,13 +441,13 @@ def test_stock_write_lands_on_envelope_truth(self): ) def test_move_kv_cache_relocates_whole_envelopes(self): - """Compaction hands PHYSICAL token runs, not kernel-facing ids; the - override must relocate exactly the page envelopes those runs name.""" + """Compaction hands PHYSICAL token runs. The override must relocate + exactly the page envelopes those runs name.""" ps = 4 kv, pool = _make_pool_and_kv(ps) - live = kv._raw.numel() - kv.view_tail_pad_bytes + live = kv._raw.numel() seed = (torch.arange(live, dtype=torch.float32) % 251).to(torch.uint8) - kv._raw[:live] = seed + kv._raw[:] = seed page_bytes = ps * _mha_spec().entry_bytes() src_pages, tgt_pages = torch.tensor([5, 6]), torch.tensor([2, 3]) @@ -341,7 +461,7 @@ def test_move_kv_cache_relocates_whole_envelopes(self): sp * page_bytes : (sp + 1) * page_bytes ] self.assertTrue( - torch.equal(kv._raw[:live], want), + torch.equal(kv._raw, want), "envelope move did not relocate exactly the named pages", ) @@ -370,8 +490,10 @@ def test_pd_registration_is_one_whole_envelope(self): self.assertLessEqual(pool._num_pages * item_lens[0], lens[0]) def test_hnd_env_cannot_hijack_layout(self): - """SGLANG_USE_HND_KVCACHE must not flip this pool's layout: HND indexes - 4-D while the per-layer views are 3-D, so the pinned label has to win.""" + """SGLANG_USE_HND_KVCACHE=1 used to flip the inherited env-driven + layout selector, putting the pool in a mode whose code paths do not + match its buffers (HND indexes 4-D; the per-layer views are 3-D). The + pinned label must win.""" with envs.SGLANG_USE_HND_KVCACHE.override(True): _, pool = _make_pool_and_kv(1) self.assertFalse(pool.use_hnd) @@ -379,6 +501,10 @@ def test_hnd_env_cannot_hijack_layout(self): class TestFactoryViews(unittest.TestCase): + """The real SWA factory builds the sub-pools and the composite allocator. + End-to-end over that factory, kernel-facing ids are the physical ones and + the rebind must emit BOTH write locs.""" + def setUp(self): # `KVIndexTranslator.__init__` asks the parallel context for # `attn_dcp_size`, which is a quotient of the configured leaves and is @@ -391,15 +517,9 @@ def setUp(self): self.addCleanup(reset_context) publish(ServerArgs(model_path="dummy"), role="test") - """Over the real SWA factory: matching kernel-facing multipliers in the - composite allocator, and a rebind that emits both write locs.""" - - # _swa_factory geometry: L_full = L_swa = 2, uniform 8/8 dims, ps = 1. - FULL_MULT = 4 # 2 * L_full - SWA_MULT = 4 # 2 * L_swa - def _bundle(self): - # Kept tiny so the per-layer views build on CPU. + # Self-contained tiny SWA-factory bundle (L_full = L_swa = 2, uniform + # 8/8 dims, ps = 1) -- small enough that per-layer views build on CPU. from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools return init_unified_swa_pools( @@ -422,16 +542,14 @@ def _bundle(self): need_sort=False, ) - def test_factory_wires_matching_multipliers(self): + def test_factory_pins_unit_multipliers(self): b = self._bundle() - pool = b.unified_memory_pool alloc = b.token_to_kv_pool_allocator - self.assertEqual(alloc.kernel_page_multiplier, self.FULL_MULT) - self.assertEqual(alloc.swa_kernel_page_multiplier, self.SWA_MULT) + self.assertEqual(alloc.kernel_page_multiplier, 1) + self.assertEqual(alloc.swa_kernel_page_multiplier, 1) # Sub-pools expose stock 3-D per-layer views. self.assertEqual(b.token_to_kv_pool.full_kv_pool.k_buffer[0].dim(), 3) self.assertEqual(b.token_to_kv_pool.swa_kv_pool.k_buffer[0].dim(), 3) - self.assertGreater(pool.view_tail_pad_bytes, 0) @unittest.skipUnless(_HAS_FP8, "requires torch.float8_e4m3fn") def test_swa_factory_preserves_fp8_logical_dtype(self): @@ -463,38 +581,48 @@ def test_swa_factory_preserves_fp8_logical_dtype(self): self.assertEqual(pool.store_dtype, torch.uint8) self.assertEqual(pool.k_buffer[0].dtype, torch.uint8) - def test_rebind_emits_kernel_facing_full_and_build_derives_swa(self): - """rebind_write_loc rebinds out_cache_loc to FULL-kernel-facing ids, and - the SWA write loc is derived pointwise from those kernel-facing values.""" + def test_rebind_emits_physical_full_and_build_derives_swa(self): + """End-to-end over the real factory: rebind_write_loc rebinds + out_cache_loc to FULL-side physical ids (phase 1), and the per-batch + build derives the SWA write loc pointwise from those values (phase 2) + -- both checked against the v2p tables over the VIRTUAL ids.""" from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator + from sglang.srt.runtime_context import get_parallel b = self._bundle() alloc = b.token_to_kv_pool_allocator v = alloc.alloc(4) self.assertIsNotNone(v) - expected_full = alloc.full_v2p_page_table[v] * self.FULL_MULT # ps=1 - expected_swa = alloc.swa_v2p_page_table[v] * self.SWA_MULT + expected_full = alloc.full_v2p_page_table[v] # ps=1 + expected_swa = alloc.swa_v2p_page_table[v] class _FB: pass fb = _FB() fb.out_cache_loc = v.clone() - source = KVIndexTranslator( - req_to_token=torch.zeros((2, 8), dtype=torch.int64), - token_to_kv_pool_allocator=alloc, - token_to_kv_pool=b.token_to_kv_pool, - page_size=1, - device="cpu", - ) - self.assertTrue(source.is_translating) - source.rebind_write_loc(fb) - self.assertTrue(torch.equal(fb.out_cache_loc, expected_full)) - self.assertTrue( - torch.equal( - source.sliding_window_write_loc_for(fb.out_cache_loc), expected_swa + # The translating paths read the DCP topology twice: construction reads + # `attn_dcp_size` to decide whether the read translate defers, and the + # write translate reads `attn_dcp_rank`, which derives from + # `dcp_enabled`. Neither has a pre-publish default and nothing is + # published in a unit test, so state both and hold the scope across + # every call that translates. + with get_parallel().override(attn_dcp_size=1, dcp_enabled=False): + source = KVIndexTranslator( + req_to_token=torch.zeros((2, 8), dtype=torch.int64), + token_to_kv_pool_allocator=alloc, + token_to_kv_pool=b.token_to_kv_pool, + page_size=1, + device="cpu", + ) + self.assertTrue(source.is_translating) + source.rebind_write_loc(fb) + self.assertTrue(torch.equal(fb.out_cache_loc, expected_full)) + self.assertTrue( + torch.equal( + source.sliding_window_write_loc_for(fb.out_cache_loc), expected_swa + ) ) - ) if __name__ == "__main__": diff --git a/test/registered/unit/mem_cache/test_unified_mla_block_table.py b/test/registered/unit/mem_cache/test_unified_mla_block_table.py index a8204c7e31d5..e4835a79024c 100644 --- a/test/registered/unit/mem_cache/test_unified_mla_block_table.py +++ b/test/registered/unit/mem_cache/test_unified_mla_block_table.py @@ -15,10 +15,11 @@ memory pool (Kimi-Linear). `req_to_token` holds VIRTUAL token ids, while the per-layer MLA views are -contiguous (`build_mla_views`). The paged MLA backends therefore need their -page-level block table filled with kernel-facing page ids: +token-major (`build_dense_views`), indexed by the physical token id. The paged +MLA backends therefore need their page-level block table filled with physical +page ids: - kernel_page(virtual_page) = v2p[virtual_page] * layer_num + kernel_page(virtual_page) = v2p[virtual_page] Since the read-path translator, ONE builder computes that formula for every family — `build_index_table` (the canonical) — and the backends only @@ -57,7 +58,7 @@ _HAS_CUDA = torch.cuda.is_available() _DEV = "cuda" -_LAYERS = 24 # K3 MLA full-attention layer count +_MULT = 1 # kernel-facing ids are the physical ones def _fill_block_table( @@ -162,22 +163,19 @@ def test_static_kernel_matches_reference(self): def test_block_table_matches_reference(self): for page_size in (1, 32, 64): rt, rpi, sl, v2p = self._make_batch(page_size) - got = _fill_block_table(rt, rpi, sl, page_size, v2p=v2p, mult=_LAYERS) - want = _reference(rt, rpi, sl, page_size, v2p=v2p, mult=_LAYERS) + got = _fill_block_table(rt, rpi, sl, page_size, v2p=v2p, mult=_MULT) + want = _reference(rt, rpi, sl, page_size, v2p=v2p, mult=_MULT) self.assertTrue( torch.equal(got.long(), want), f"page_size={page_size}:\ngot ={got}\nwant={want}", ) def test_single_full_attention_layer_still_maps_v2p(self): - """A config with exactly ONE full-attention layer (e.g. a PP rank owning a - single MLA layer) has `kernel_page_multiplier == 1`, but its req_to_token - still holds VIRTUAL ids. The kernel-facing id collapses onto the physical id, so - the v2p gather alone IS the whole translation -- it must not be skipped. - - Regression guard for detecting the unified pool via `multiplier > 1`: - that predicate treats this config as a static pool and leaves the block - table in virtual id space. + """The kernel-facing id IS the physical id, so the v2p gather alone is + the whole translation -- it must not be skipped. Regression guard for + detecting the unified pool via `multiplier > 1`: that predicate would + treat every unified pool as static and leave the block table in + virtual id space. """ for page_size in (1, 64): rt, rpi, sl, v2p = self._make_batch(page_size) @@ -198,7 +196,7 @@ def test_padded_lanes_stay_untouched(self): trtllm/flashmla block-table contract).""" page_size = 64 rt, rpi, sl, v2p = self._make_batch(page_size) - got = _fill_block_table(rt, rpi, sl, page_size, v2p=v2p, mult=_LAYERS) + got = _fill_block_table(rt, rpi, sl, page_size, v2p=v2p, mult=_MULT) for r in range(got.shape[0]): n_pages = (int(sl[r].item()) + page_size - 1) // page_size self.assertTrue( @@ -213,14 +211,14 @@ def test_agrees_with_token_level_translate(self): page_size = 64 rt, rpi, sl, v2p = self._make_batch(page_size) block_table = _fill_block_table( - rt, rpi, sl, page_size, v2p=v2p, mult=_LAYERS + rt, rpi, sl, page_size, v2p=v2p, mult=_MULT ).long() for r in range(rt.shape[0]): n = int(sl[r].item()) virt_tokens = rt[r, :n].long() # translate_kv_loc_for_kernel's formula, applied to token ids. kernel_tokens = ( - v2p[virt_tokens // page_size] * (page_size * _LAYERS) + v2p[virt_tokens // page_size] * (page_size * _MULT) + virt_tokens % page_size ) # The block-table entry scaled by page_size must be the kernel-facing id of @@ -326,11 +324,11 @@ def test_identity_when_hooks_absent(self): self._assert_live_prefix(got, want, sl, page_size) def test_translated_mapping_ps1_fast_path(self): - got, want, sl = self._run(1, v2p=True, mult=_LAYERS) + got, want, sl = self._run(1, v2p=True, mult=_MULT) self._assert_live_prefix(got, want, sl, 1) def test_translated_mapping_general_path(self): - got, want, sl = self._run(64, v2p=True, mult=_LAYERS) + got, want, sl = self._run(64, v2p=True, mult=_MULT) self._assert_live_prefix(got, want, sl, 64) def test_single_full_attention_layer(self): diff --git a/test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py b/test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py index 9cb7996903fb..63efebdeb8c8 100644 --- a/test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py +++ b/test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py @@ -43,7 +43,7 @@ def _kernel_id(t: torch.Tensor, ps: int) -> torch.Tensor: - return (t // ps) * (ps * _L) + t % ps + return t # token-major views: the kernel id is the physical token id def _make_pools(ps: int, n_tokens: int = 4096): diff --git a/test/registered/unit/mem_cache/test_unified_mla_views.py b/test/registered/unit/mem_cache/test_unified_mla_views.py index 0c0347327df6..0bb957585056 100644 --- a/test/registered/unit/mem_cache/test_unified_mla_views.py +++ b/test/registered/unit/mem_cache/test_unified_mla_views.py @@ -11,19 +11,30 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""MLA views for the unified memory pool (MLA-hybrid-Mamba, Kimi K3), CPU-only. - -Addressing law under test: the (page, layer, slot) cell sits at envelope byte -offset `p*(L*ps*D) + l*(ps*D) + s*D`, reached through the kernel-facing id -`(t // ps) * (ps * L) + t % ps`. - -GPU parity of the read/write kernels (set_mla_kv_buffer TMA path etc.) lives in -`test_unified_mla_gpu_parity.py`. +"""MLA views for the unified memory pool (MLA-hybrid-Mamba, Kimi K3). + +Covers, CPU-only (pure torch -- no GPU / Triton kernels): + - `MLASubPoolSpec` byte math (aligned entry, one latent row per layer); + - `build_dense_views` addressing: view_l[t] for PHYSICAL token t must land + at the envelope byte `t * entry_bytes + l * row_bytes`, the per-layer + views must not alias at equal ids, and a short buffer must fail loud; + - `UnifiedKVPool` MLA plumbing: the allocation is exactly the budget and + the reserved sink floor covers the whole page-0 envelope; + - `UnifiedMLATokenToKVPool`: buffer wiring, V-as-prefix-slice, and the + page-envelope `move_kv_cache` (physical token ids, page-major runs); + - `MultiEndedAllocator.translate_kv_loc_for_kernel`: identical to the + physical translate, tombstone clamp to the sink, `out=` contract, the + pinned multiplier, and correctness across eager compaction. + +GPU parity of the actual read/write kernels (set_mla_kv_buffer TMA path etc.) +lives in the server-level tests, not here. + + python -m pytest test/registered/unit/mem_cache/test_unified_mla_views.py -v """ from sglang.test.ci.ci_register import register_cpu_ci -register_cpu_ci(est_time=10, suite="base-a-test-cpu") +register_cpu_ci(est_time=8, suite="base-a-test-cpu") import unittest from unittest import mock @@ -34,7 +45,12 @@ UnifiedMambaTokenToKVPoolAllocator, ) from sglang.srt.mem_cache.allocator.unified_sub_pool import MultiEndedAllocator -from sglang.srt.mem_cache.layout.page_major import build_mla_views +from sglang.srt.mem_cache.layout.page_major import ( + ENTRY_ALIGN_BYTES, + build_dense_views, + mla_entry_bytes, + paged_row_view, +) from sglang.srt.mem_cache.unified_memory_pool import ( MambaSubPoolSpec, MLASubPoolSpec, @@ -45,14 +61,18 @@ _DEV = "cpu" -# Geometry kept tiny so every byte offset is hand-checkable; real K3 is -# L=24, D=576 (=512+64). +# Small-but-nontrivial MLA geometry: L=3 layers, D=8 (=6+2), so every byte +# offset is hand-checkable. Real K3 is L=24, D=576 (=512+64). 3 rows x 16 B +# = 48 B of payload round up to one 64 B entry. _L = 3 _LORA = 6 _ROPE = 2 _D = _LORA + _ROPE _DTYPE = torch.bfloat16 _ITEM = _DTYPE.itemsize +_ROW = _D * _ITEM +_ENTRY = mla_entry_bytes(layer_num=_L, kv_cache_dim=_D, itemsize=_ITEM) +_E_ELEMS = _ENTRY // _ITEM def _mla_spec(grow="down", layer_num=_L): @@ -92,11 +112,27 @@ def _make_unified(page_size=1, n_full_tokens=64, n_mamba_slots=8): return pool, full, mamba -def _kernel_id(t, ps, layer_num): - return (t // ps) * (ps * layer_num) + t % ps +def _build_views(raw, ps, num_pages): + layout = _mla_spec().layout() + return build_dense_views( + raw, layout=layout, part=layout.part("kv"), page_size=ps, num_pages=num_pages + ) class TestMLASubPoolSpec(unittest.TestCase): + def test_entry_bytes_and_dim(self): + spec = _mla_spec() + self.assertEqual(spec.kv_cache_dim, _D) + self.assertEqual(spec.row_bytes(), _ROW) + self.assertEqual(spec.entry_bytes(), _ENTRY) + self.assertGreaterEqual(spec.entry_bytes(), _L * _ROW) + self.assertEqual(spec.entry_bytes() % ENTRY_ALIGN_BYTES, 0) + self.assertEqual(spec.get_dtype(), _DTYPE) + layout = spec.layout() + self.assertEqual(layout.entry_bytes, _ENTRY) + for l in range(_L): + self.assertEqual(layout.part("kv").layer_offset_bytes(l), l * _ROW) + def test_rejects_nonpositive_dims(self): with self.assertRaises(AssertionError): MLASubPoolSpec( @@ -110,79 +146,72 @@ def test_rejects_nonpositive_dims(self): class TestMLAViews(unittest.TestCase): - def _make_raw(self, ps, num_pages, pad_pages=1): - page_bytes = ps * _L * _D * _ITEM - raw = torch.zeros( - (num_pages + pad_pages) * page_bytes, dtype=torch.uint8, device=_DEV - ) - return raw, page_bytes + def _make_raw(self, ps, num_pages, short=0): + n = num_pages * ps * _ENTRY - short + return torch.zeros(n, dtype=torch.uint8, device=_DEV) def test_view_addressing_matches_envelope_formula(self): for ps in (1, 4): num_pages = 6 - raw, _ = self._make_raw(ps, num_pages) - views = build_mla_views( - raw, - layer_num=_L, - kv_cache_dim=_D, - store_dtype=_DTYPE, - page_size=ps, - num_pages=num_pages, - ) + raw = self._make_raw(ps, num_pages) + views = _build_views(raw, ps, num_pages) self.assertEqual(len(views), _L) - n_rows = num_pages * _L * ps + n_rows = num_pages * ps for v in views: self.assertEqual(tuple(v.shape), (n_rows, 1, _D)) - # contiguous in the (row, dim) sense: .view(-1, ps, D) legality - self.assertEqual(v.stride(0), _D) - self.assertEqual(v.stride(2), 1) + self.assertEqual(v.stride(), (_E_ELEMS, _D, 1)) flat = raw.view(_DTYPE) for p, l, s in [(0, 0, 0), (1, 2, ps - 1), (4, 1, ps // 2), (5, 2, 0)]: t = p * ps + s marker = float(p * 100 + l * 10 + s + 1) - views[l][_kernel_id(t, ps, _L)] = marker - # envelope formula, in elements - elem = p * (_L * ps * _D) + l * (ps * _D) + s * _D + views[l][t] = marker + elem = t * _E_ELEMS + l * _D # envelope formula, in elements self.assertTrue( torch.all(flat[elem : elem + _D] == marker), f"(p={p}, l={l}, s={s}, ps={ps}) landed off-formula", ) + def test_paged_row_view_keeps_the_slot_stride(self): + """BUG REGRESSION at page_size 1: the paged MLA backends hand the + kernels `paged_row_view(kv)`, whose dim 1 must carry the entry stride + at every page size; a `view`-built split gave the size-1 slot dim the + row stride instead.""" + num_pages = 3 + for ps in (1, 4): + views = _build_views(self._make_raw(ps, num_pages), ps, num_pages) + paged = paged_row_view(views[1], ps) + self.assertEqual(tuple(paged.shape), (num_pages, ps, _D), ps) + self.assertEqual(tuple(paged.stride()), (ps * _E_ELEMS, _E_ELEMS, 1), ps) + for t in range(num_pages * ps): + self.assertEqual( + paged[t // ps, t % ps].data_ptr(), + views[1][t].data_ptr(), + (ps, t), + ) + def test_views_do_not_alias_across_layers(self): - ps = 4 - num_pages = 4 - raw, _ = self._make_raw(ps, num_pages) - views = build_mla_views( - raw, - layer_num=_L, - kv_cache_dim=_D, - store_dtype=_DTYPE, - page_size=ps, - num_pages=num_pages, - ) + ps, num_pages = 4, 4 + views = _build_views(self._make_raw(ps, num_pages), ps, num_pages) t = 2 * ps + 1 # page 2, slot 1 - d = _kernel_id(t, ps, _L) for l in range(_L): - views[l][d] = float(l + 1) + views[l][t] = float(l + 1) for l in range(_L): - self.assertTrue(torch.all(views[l][d] == float(l + 1))) + self.assertTrue(torch.all(views[l][t] == float(l + 1))) - def test_missing_tail_pad_fails_loud(self): - ps = 2 - num_pages = 4 - raw, _ = self._make_raw(ps, num_pages, pad_pages=0) + def test_short_buffer_fails_loud(self): + ps, num_pages = 2, 4 with self.assertRaises(AssertionError): - build_mla_views( - raw, - layer_num=_L, - kv_cache_dim=_D, - store_dtype=_DTYPE, - page_size=ps, - num_pages=num_pages, - ) + _build_views(self._make_raw(ps, num_pages, short=1), ps, num_pages) class TestUnifiedKVPoolMLA(unittest.TestCase): + def test_raw_is_exactly_the_budget(self): + pool, full, mamba = _make_unified(page_size=4) + total = full.entry_bytes() * 64 + mamba.entry_bytes() * 8 + self.assertEqual(pool.max_slots("full"), total // full.entry_bytes()) + self.assertEqual(pool.max_slots("mamba"), total // mamba.entry_bytes()) + self.assertEqual(pool._raw.numel(), total) + def test_reserved_floor_covers_page0_envelope(self): ps = 4 pool, full, mamba = _make_unified(page_size=ps) @@ -199,6 +228,7 @@ def test_mla_views_accessor(self): views = pool.mla_views_for("full") self.assertEqual(len(views), _L) self.assertIs(pool.mla_spec("full"), full) + self.assertEqual(views[0].stride(0) * _ITEM, full.entry_bytes()) class TestUnifiedMLATokenToKVPool(unittest.TestCase): @@ -216,6 +246,7 @@ def test_buffers_and_prefix_value_slice(self): pool, kv_pool = self._make(ps=1) self.assertEqual(len(kv_pool.kv_buffer), _L) self.assertEqual(kv_pool.get_kv_size_bytes(), 0) + self.assertEqual(kv_pool.size, pool.max_slots("full") - 1) k = kv_pool.get_key_buffer(1) v = kv_pool.get_value_buffer(1) self.assertEqual(k.shape[-1], _D) @@ -226,26 +257,15 @@ def test_buffers_and_prefix_value_slice(self): def test_cpu_copy_round_trips_through_physical_ids(self): """REGRESSION: the host copy for decode retraction is addressed by - PHYSICAL token ids, but this pool's `kv_buffer` views are indexed by - kernel-facing ids. Without the rewrite the parent read a different row - and the restore silently returned other tokens' KV.""" + PHYSICAL token ids. Under the token-major entry those ARE the ids the + `kv_buffer` views take, so the round trip must need no rewrite; one + would read a different row and restore other tokens' KV.""" for ps in (1, 4): with self.subTest(page_size=ps): pool, kv_pool = self._make(ps=ps) phys = torch.tensor([0, 1, ps, ps + 1], dtype=torch.int64) - self.assertTrue( - torch.equal( - kv_pool._physical_to_kernel_indices(phys), - torch.tensor( - [_kernel_id(int(t), ps, _L) for t in phys], - dtype=torch.int64, - ), - ) - ) for layer in range(_L): - kv_pool.get_key_buffer(layer)[ - kv_pool._physical_to_kernel_indices(phys) - ] = float(layer + 1) + kv_pool.get_key_buffer(layer)[phys] = float(layer + 1) with ( get_parallel().override(dcp_enabled=False, attn_dcp_rank=0), @@ -258,33 +278,22 @@ def test_cpu_copy_round_trips_through_physical_ids(self): kv_pool.load_cpu_copy(saved, phys) for layer in range(_L): - restored = kv_pool.get_key_buffer(layer)[ - kv_pool._physical_to_kernel_indices(phys) - ] + restored = kv_pool.get_key_buffer(layer)[phys] self.assertTrue( torch.all(restored == float(layer + 1)), f"layer {layer} did not round-trip at page_size {ps}", ) def test_move_kv_cache_moves_page_envelopes(self): - """Whole page envelopes relocate, in raw bytes and (at ps=4) as read - back through the per-layer views at the destination kernel ids.""" for ps in (1, 4): pool, kv_pool = self._make(ps=ps) num_pages = pool.max_slots("full") // ps - page_bytes = ps * _L * _D * _ITEM + page_bytes = ps * pool.mla_spec("full").entry_bytes() env = pool._raw[: num_pages * page_bytes].view(num_pages, page_bytes) - src_pages = torch.tensor([num_pages - 2, num_pages - 4, num_pages - 3]) - dst_pages = torch.tensor([2, 3, 5]) + src_pages = torch.tensor([num_pages - 2, num_pages - 4]) + dst_pages = torch.tensor([2, 3]) env[src_pages[0]] = 7 env[src_pages[1]] = 9 - if ps == 4: - # write through the views at src, expect it at dst after the move - for l in range(_L): - for s in range(ps): - kv_pool.kv_buffer[l][ - _kernel_id(int(src_pages[2]) * ps + s, ps, _L) - ] = float(l * ps + s + 1) # page-major token runs, exactly how compaction expands pages offsets = torch.arange(ps, dtype=torch.int64) src_t = (src_pages[:, None] * ps + offsets).reshape(-1) @@ -292,15 +301,27 @@ def test_move_kv_cache_moves_page_envelopes(self): kv_pool.move_kv_cache(dst_t, src_t) self.assertTrue(torch.all(env[dst_pages[0]] == 7), f"ps={ps}") self.assertTrue(torch.all(env[dst_pages[1]] == 9), f"ps={ps}") - if ps == 4: - for l in range(_L): - for s in range(ps): - got = kv_pool.kv_buffer[l][ - _kernel_id(int(dst_pages[2]) * ps + s, ps, _L) - ] - self.assertTrue( - torch.all(got == float(l * ps + s + 1)), f"(l={l}, s={s})" - ) + + def test_move_then_readback(self): + ps = 4 + pool, kv_pool = self._make(ps=ps) + num_pages = pool.max_slots("full") // ps + src_page, dst_page = num_pages - 3, 5 + # write through the views at src, expect it at dst after the move + for l in range(_L): + for s in range(ps): + kv_pool.kv_buffer[l][src_page * ps + s] = float(l * ps + s + 1) + offsets = torch.arange(ps, dtype=torch.int64) + kv_pool.move_kv_cache( + (torch.tensor([dst_page])[:, None] * ps + offsets).reshape(-1), + (torch.tensor([src_page])[:, None] * ps + offsets).reshape(-1), + ) + for l in range(_L): + for s in range(ps): + got = kv_pool.kv_buffer[l][dst_page * ps + s] + self.assertTrue( + torch.all(got == float(l * ps + s + 1)), f"(l={l}, s={s})" + ) class _FakeKVCache: @@ -312,7 +333,7 @@ def move_kv_cache(self, dst_loc: torch.Tensor, src_loc: torch.Tensor): class TestTranslateKvLocForKernel(unittest.TestCase): - def _build(self, ps=1, n_full_tokens=64, multiplier=_L): + def _build(self, ps=1, n_full_tokens=64, multiplier=None): pool, full, mamba = _make_unified(page_size=ps, n_full_tokens=n_full_tokens) full_alloc = MultiEndedAllocator( kvcache=_FakeKVCache(pool.max_slots("full")), @@ -334,33 +355,20 @@ def _build(self, ps=1, n_full_tokens=64, multiplier=_L): mamba_alloc.bind_peer(full_alloc) return full_alloc - def test_kernel_id_matches_formula(self): - """kernel id = (phys // ps) * (ps * multiplier) + phys % ps, across - page sizes, the multiplier-1 physical fallback, and eager compaction.""" - for ps, multiplier in ((1, _L), (4, _L), (1, 1)): - with self.subTest(page_size=ps, multiplier=multiplier): - alloc = self._build(ps=ps, multiplier=multiplier) - a = alloc.alloc(4 * ps) - b = alloc.alloc(4 * ps) - c = alloc.alloc(4 * ps) - self.assertIsNotNone(c) - - def check(virt): - phys = alloc.translate_kv_loc(virt) - expected = (phys // ps) * (ps * multiplier) + phys % ps - self.assertTrue( - torch.all(alloc.translate_kv_loc_for_kernel(virt) == expected) - ) - - for virt in (a, b, c): - check(virt) - alloc.free(b) # eager compaction relocates survivors - for virt in (a, c): - check(virt) + def test_kernel_id_is_the_physical_id(self): + for ps in (1, 4): + alloc = self._build(ps=ps) + v = alloc.alloc(3 * ps) + self.assertIsNotNone(v) + phys = alloc.translate_kv_loc(v) + kernel = alloc.translate_kv_loc_for_kernel(v) + self.assertTrue(torch.equal(kernel, phys), f"ps={ps}") + v2p = alloc.virtual_to_physical + self.assertTrue(torch.equal(kernel, v2p[v // ps] * ps + v % ps)) def test_tombstone_clamps_to_sink(self): alloc = self._build(ps=1) - # never-allocated virtual ids -> v2p == -1 -> kernel-facing id 0 + # never-allocated virtual ids -> v2p == -1 -> id 0 virt = torch.tensor([alloc.min_slot_index + 1], dtype=torch.int64) kernel = alloc.translate_kv_loc_for_kernel(virt) self.assertTrue(torch.all(kernel == 0)) @@ -380,6 +388,26 @@ def test_out_matches_and_aliases(self): alloc.translate_kv_loc_for_kernel(x, out=x) self.assertTrue(torch.all(x == no_out)) + def test_multiplier_is_pinned_to_one(self): + self.assertEqual(self._build(ps=1).kernel_page_multiplier, 1) + self.assertEqual(self._build(ps=1, multiplier=1).kernel_page_multiplier, 1) + with self.assertRaises(AssertionError): + self._build(ps=1, multiplier=_L) + + def test_kernel_id_follows_compaction(self): + alloc = self._build(ps=1) + a = alloc.alloc(4) + b = alloc.alloc(4) + c = alloc.alloc(4) + self.assertIsNotNone(c) + alloc.free(b) # eager compaction relocates survivors + for run in (a, c): + self.assertTrue( + torch.equal( + alloc.translate_kv_loc_for_kernel(run), alloc.translate_kv_loc(run) + ) + ) + class _RecordingHybridPool: """Stands in for `UnifiedHybridLinearKVPool`, recording the ids it is handed.""" diff --git a/test/registered/unit/mem_cache/test_unified_npool_sweep.py b/test/registered/unit/mem_cache/test_unified_npool_sweep.py index 6af10bdd4453..d0133183cba9 100644 --- a/test/registered/unit/mem_cache/test_unified_npool_sweep.py +++ b/test/registered/unit/mem_cache/test_unified_npool_sweep.py @@ -221,16 +221,14 @@ def test_float_mha_views_shape_and_roundtrip(self): self.assertEqual(len(k_views), spec.layer_num) self.assertEqual(len(v_views), spec.layer_num) num_pages = pool.max_slots("swa") // page_size - blocks = 2 * spec.layer_num # K at block 2l, V at 2l+1 - n_rows = num_pages * blocks * page_size + n_rows = num_pages * page_size for k in (*k_views, *v_views): - # Stock 3-D per-layer MHA signature; the row index is the - # kernel-facing id, each view's storage_offset folding in its block - # origin (see `build_mha_views`). + # 3-D per-layer views indexed by the physical token id; the slot + # stride is the whole entry (see `build_dense_views`). self.assertEqual(tuple(k.shape), (n_rows, spec.head_num, spec.head_dim)) + self.assertEqual(k.stride(0) * k.element_size(), spec.entry_bytes()) # Round-trip: a float view is a real strided window into _raw. - slot = pool.min_slot_index("swa") - row = (slot // page_size) * (page_size * blocks) + slot % page_size + row = pool.min_slot_index("swa") pattern = ( torch.arange(spec.head_num * spec.head_dim, dtype=torch.float32) .reshape(spec.head_num, spec.head_dim) diff --git a/test/registered/unit/server_args/test_page_major_backend_allowlist.py b/test/registered/unit/server_args/test_page_major_backend_allowlist.py index 5cb295b14489..628b62276a0e 100644 --- a/test/registered/unit/server_args/test_page_major_backend_allowlist.py +++ b/test/registered/unit/server_args/test_page_major_backend_allowlist.py @@ -13,19 +13,30 @@ # ============================================================================== """`--enable-page-major-kv-layout` full-attention backend allowlist. -`handle_page_major_kv_layout` gates two ways, because the per-layer views the -unified pool exposes are all the allowlisted backends can read: an MLA arm -(`build_mla_views`), an MHA/SWA arm (`build_mha_views`), and no page-major arm -at all without the unified pool. `fa3` is the resolved default on pre-Blackwell -hosts, so its absence from an arm makes `--enable-unified-memory` fail to boot -under its own default configuration. - -The same handler screens the pool itself: the MHA/SWA per-layer views need -uniform K/V rows, so an asymmetric-K/V model (MiMoV2: head_dim 192 != -v_head_dim 128) is rejected on EVERY backend, Triton included. MLA models are -exempt -- their sub-pool keeps one latent row per layer, and real MLA configs -(Kimi-Linear: head_dim 72, v_head_dim 128) report asymmetric dims while running -the unified pool today. +Two-way gate (see `_handle_page_major_kv_layout`), because the unified pool +exposes per-layer views and nothing else: + * unified-memory MLA models allow the whole wired + paged MLA family -- `fa3`, `flashinfer`'s MLA backend, `trtllm_mla` with + its `cutedsl_mla` / `tokenspeed_mla` subclasses, and `flashmla` (ps=64 + snap); + * unified-memory MHA/SWA models allow `fa3` / + `fa4` / `flashinfer` / `trtllm_mha` alongside Triton; + * plain `--enable-page-major-kv-layout` without the unified pool keeps the + envelope-strided 4-D views only the stride-aware Triton kernels read. + +The same handler also screens the pool itself: asymmetric K/V rows are not +admitted yet, so an asymmetric-K/V model (MiMoV2: head_dim 192 != v_head_dim +128) cannot run `--enable-unified-memory` at all and is rejected on EVERY +backend, Triton included. MLA models are exempt -- their sub-pool keeps +one latent row per layer, and several MLA configs (Kimi-Linear: head_dim 72, +v_head_dim 128) report asymmetric dims while running the unified pool today. + +Pinned here so no arm silently widens to an unwired backend (`cutlass_mla`, +`aiter`) and no arm silently narrows: `fa3` is the resolved default on +pre-Blackwell hosts, so its absence from an arm makes `--enable-unified-memory` +fail to boot under its own default configuration. + + python -m pytest test/registered/unit/server_args/test_page_major_backend_allowlist.py -v """ import unittest @@ -96,7 +107,7 @@ class TestPageMajorBackendAllowlist(unittest.TestCase): "tokenspeed_mla", "flashmla", ) - # Wired for the per-layer MHA/SWA views (uniform-row models). + # Wired for the per-layer MHA/SWA views. PER_LAYER_VIEW_MHA_BACKENDS = ("fa3", "fa4", "flashinfer", "trtllm_mha") # MLA-family kernels that must never leak into the MHA arm. MLA_ONLY_BACKENDS = ("trtllm_mla", "cutedsl_mla", "tokenspeed_mla", "flashmla") @@ -108,8 +119,8 @@ def test_triton_allowed_on_every_arm(self): the MLA nor the MHA arm can narrow away.""" for use_mla in (True, False): self.assertTrue(_accepts("triton", use_mla=use_mla)) - # The uniform-row screen is a property of the model, not of the - # backend, so it rejects even Triton. + # Asymmetric K/V rows are not admitted yet; that is a property of the + # model, not of the backend, so the screen rejects even Triton. self.assertFalse(_accepts("triton", use_mla=False, has_asymmetric_kv=True)) def test_per_layer_view_mla_backends_allowed_under_unified_mla(self): @@ -119,11 +130,11 @@ def test_per_layer_view_mla_backends_allowed_under_unified_mla(self): f"{backend} should be allowed with the unified-memory MLA pool", ) - def test_per_layer_view_mha_backends_allowed_for_uniform_row_models(self): + def test_per_layer_view_mha_backends_allowed_under_unified_mha(self): for backend in self.PER_LAYER_VIEW_MHA_BACKENDS: self.assertTrue( _accepts(backend, use_mla=False), - f"{backend} should be allowed for a uniform-row MHA model", + f"{backend} should be allowed with the unified-memory MHA pool", ) def test_mla_only_backends_rejected_for_mha(self): @@ -134,7 +145,8 @@ def test_mla_only_backends_rejected_for_mha(self): ) def test_asymmetric_kv_mha_model_cannot_use_unified_memory(self): - """The rejection is the POOL's, not a backend's, so it must fire on + """head_dim != v_head_dim (MiMoV2): not admitted under the unified pool + yet. The rejection is the POOL's, not a backend's, so it must fire on every backend -- Triton included.""" for backend in ("triton",) + self.PER_LAYER_VIEW_MHA_BACKENDS: self.assertFalse( From ef45ec87a6b590477571e34d3b0a84577b606a5b Mon Sep 17 00:00:00 2001 From: Caihua Li Date: Sat, 12 Sep 2026 01:04:30 -0700 Subject: [PATCH 03/15] refactor(unified-memory): delete the page-major envelope machinery Nothing reads any of it once the unified pool stores token-major; each symbol below has zero remaining references: build_mha_views / build_mla_views superseded by build_dense_views. SubPoolSpec.view_tail_pad_bytes the views no longer overrun the last page, so there is no pad to size. SubPoolSpec.blocks_per_page kernel-facing ids are physical ids, so there is no per-page block scale. KVCache.kernel_page_blocks same, and it was the probe's only input. maybe_detect_kernel_facing_loc the probe told a kernel-facing id from a virtual one by its in-page offset. With one id space left there is nothing to separate, and it was vacuous at blocks_per_page 1 anyway. The OOB probe at the same write doors is unchanged. The unified page-envelope host pool had one caller of `kernel_page_blocks`: it divided the ids it is handed by the kernel-view page stride to recover physical slots. Those ids are physical here, so the conversion goes, and its unit test reads pages straight off them. --- .../sglang/srt/mem_cache/layout/page_major.py | 131 ------------------ python/sglang/srt/mem_cache/memory_pool.py | 18 +-- .../sglang/srt/mem_cache/pool_host/unified.py | 7 - .../srt/mem_cache/unified_memory_pool.py | 28 ---- python/sglang/srt/utils/async_probe.py | 21 --- .../unit/mem_cache/test_full_loc_fast_path.py | 1 - .../mem_cache/test_swa_cpu_copy_filter.py | 4 +- .../test_unified_out_cache_loc_rebind.py | 20 --- 8 files changed, 2 insertions(+), 228 deletions(-) diff --git a/python/sglang/srt/mem_cache/layout/page_major.py b/python/sglang/srt/mem_cache/layout/page_major.py index 066eb93f4ed4..2141f5e37304 100644 --- a/python/sglang/srt/mem_cache/layout/page_major.py +++ b/python/sglang/srt/mem_cache/layout/page_major.py @@ -163,142 +163,11 @@ def mha_entry_bytes( return align_entry_bytes(layer_num * (k_row_bytes + v_row_bytes)) -def build_mha_views( - raw: torch.Tensor, - *, - layer_num: int, - head_num: int, - head_dim: int, - v_head_dim: int, - store_dtype: torch.dtype, - page_size: int, - num_pages: int, - anchor_bytes: int = 0, -) -> Tuple[List[torch.Tensor], List[torch.Tensor]]: - """Per-layer K/V views over ``raw`` for uniform-row MHA. - - The page envelope ``[L0_K*ps | L0_V*ps | L1_K*ps | ...]`` is a uniform - array of ``2*layer_num`` row-blocks when K and V rows are equally wide, so - it is a valid paged pool under - - kernel_id(t) = (t // ps) * (ps * 2 * layer_num) + t % ps - - with layer ``l``'s K at block ``2l`` and its V at block ``2l+1``. Each view - is a contiguous ``(num_pages * 2 * layer_num * ps, head_num, head_dim)``. - - Views overlap by ``ps`` rows per block, safe because an id always resolves - inside its own block; the last view runs ``(2*layer_num - 1) * ps`` rows - past the envelope, so ``raw`` needs ``UnifiedKVPool.view_tail_pad_bytes``. - """ - assert head_dim == v_head_dim, ( - f"build_mha_views requires uniform rows (head_dim == v_head_dim); " - f"got head_dim={head_dim}, v_head_dim={v_head_dim}. Asymmetric-KV " - "models cannot use the unified pool (screened out at startup)." - ) - itemsize = store_dtype.itemsize - row_elems = head_num * head_dim - row_bytes = row_elems * itemsize - blocks = 2 * layer_num - page_bytes = page_size * blocks * row_bytes - n_rows = num_pages * blocks * page_size - assert anchor_bytes % itemsize == 0 - last_view_end = ( - anchor_bytes + (blocks - 1) * page_size * row_bytes + n_rows * row_bytes - ) - assert last_view_end <= raw.numel() * raw.itemsize, ( - f"build_mha_views: block {blocks - 1}'s view ends at byte " - f"{last_view_end} but the raw buffer holds only " - f"{raw.numel() * raw.itemsize} bytes; allocate the tail pad " - f"(one page envelope = {page_bytes} B) via view_tail_pad_bytes" - ) - - as_dtype_view = raw.view(store_dtype) - k_buffer: List[torch.Tensor] = [] - v_buffer: List[torch.Tensor] = [] - for layer in range(layer_num): - k_base_bytes = anchor_bytes + (2 * layer) * page_size * row_bytes - v_base_bytes = k_base_bytes + page_size * row_bytes - for base_bytes, out in ((k_base_bytes, k_buffer), (v_base_bytes, v_buffer)): - assert base_bytes % itemsize == 0 - out.append( - torch.as_strided( - as_dtype_view, - size=(n_rows, head_num, head_dim), - stride=(row_elems, head_dim, 1), - storage_offset=base_bytes // itemsize, - ) - ) - return k_buffer, v_buffer - - def mla_entry_bytes(*, layer_num: int, kv_cache_dim: int, itemsize: int) -> int: """Bytes occupied by one MLA slot across all layers (latent rows), aligned.""" return align_entry_bytes(layer_num * kv_cache_dim * itemsize) -def build_mla_views( - raw: torch.Tensor, - *, - layer_num: int, - kv_cache_dim: int, - store_dtype: torch.dtype, - page_size: int, - num_pages: int, - anchor_bytes: int = 0, -) -> List[torch.Tensor]: - """Per-layer views over ``raw`` for MLA in the page-major layout. - - The page envelope is ``[L0_latent * ps | L1_latent * ps | ...]``. Because all - MLA layers share one uniform row size (``kv_cache_dim``), the envelope is - itself a valid paged pool under a re-numbered index space: folding the - layer offset ``l * ps * kv_cache_dim`` into each view's storage_offset makes - every per-layer view a plain CONTIGUOUS ``(num_pages * layer_num * ps, 1, - kv_cache_dim)`` tensor, addressed by the layer-independent kernel-facing id - - kernel_id(t) = (t // ps) * (ps * layer_num) + t % ps (t = physical token) - - so one shared block table (entry = page * layer_num) serves every layer, and - kernels that require ``.view(-1, page_size, kv_cache_dim)`` (trtllm/cutlass/ - flashmla) work on the views natively. - - The views overlap each other (view ``l+1`` is view ``l`` shifted by ``ps`` - rows); that is safe because layer ``l`` is only ever indexed at kernel-facing ids, - which always resolve to layer-``l`` bytes relative to view ``l``'s origin. - Layer ``layer_num-1``'s view extends ``(layer_num-1) * ps`` rows past the - last page envelope, so ``raw`` must carry at least one extra page envelope - of tail padding (``UnifiedKVPool``'s ``view_tail_pad_bytes``). - """ - itemsize = store_dtype.itemsize - row_bytes = kv_cache_dim * itemsize - page_bytes = page_size * layer_num * row_bytes - n_rows = num_pages * layer_num * page_size - assert anchor_bytes % itemsize == 0 - last_view_end = ( - anchor_bytes + (layer_num - 1) * page_size * row_bytes + (n_rows * row_bytes) - ) - assert last_view_end <= raw.numel() * raw.itemsize, ( - f"build_mla_views: layer {layer_num - 1}'s view ends at byte " - f"{last_view_end} but the raw buffer holds only " - f"{raw.numel() * raw.itemsize} bytes; allocate the tail pad " - f"(one page envelope = {page_bytes} B) via view_tail_pad_bytes" - ) - - as_dtype_view = raw.view(store_dtype) - views: List[torch.Tensor] = [] - for layer in range(layer_num): - base_bytes = anchor_bytes + layer * page_size * row_bytes - assert base_bytes % itemsize == 0 - views.append( - torch.as_strided( - as_dtype_view, - size=(n_rows, 1, kv_cache_dim), - stride=(kv_cache_dim, kv_cache_dim, 1), - storage_offset=base_bytes // itemsize, - ) - ) - return views - - def mamba_entry_bytes( *, layer_num: int, diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 746439bb344c..93b5eeabf9b3 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -86,7 +86,6 @@ next_power_of_2, ) from sglang.srt.utils.async_probe import ( - maybe_detect_kernel_facing_loc, maybe_detect_oob, ) from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter @@ -1845,10 +1844,6 @@ def __init__( ): self.size = size self.page_size = page_size - # Row-blocks one page holds in this pool's kernel-facing id space; >1 - # only for the unified pool's per-layer views, and then a write loc must - # have been translated into that space first. - self.kernel_page_blocks = 1 self.dtype = dtype self.device = device if dtype in (torch.float8_e5m2, torch.float8_e4m3fn, torch.float8_e4m3fnuz): @@ -2615,9 +2610,6 @@ def set_kv_buffer( # Catch stale slot ids here instead of as illegal-addr / silent KV # corruption in the store_kvcache write (gated on SGLANG_ENABLE_ASYNC_ASSERT). maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MHA)") - maybe_detect_kernel_facing_loc( - loc, self.page_size, self.kernel_page_blocks, "set_kv_buffer (MHA)" - ) layer_id = ( layer_id_override if layer_id_override is not None else layer.layer_id ) @@ -4530,9 +4522,7 @@ def get_kv_buffer(self, layer_id: int): # Has the WRITE loc arriving here already had the DCP owner rule resolved? # False: this pool takes a WIDENED loc. The unified pool resolves it in - # `KVIndexTranslator.rebind_write_loc` and flips this. Not derivable from - # `kernel_page_blocks`: that is `layer_num`, so a rank owning one - # full-attention layer is translated with blocks_per_page 1. + # `KVIndexTranslator.rebind_write_loc` and flips this. write_loc_is_dcp_resolved = False @property @@ -4564,9 +4554,6 @@ def set_kv_buffer( ): loc, _, _ = unwrap_write_loc(loc_info) maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MLA)") - maybe_detect_kernel_facing_loc( - loc, self.page_size, self.kernel_page_blocks, "set_kv_buffer (MLA)" - ) layer_id = ( layer_id_override if layer_id_override is not None else layer.layer_id ) @@ -4651,9 +4638,6 @@ def set_mla_kv_buffer( (self.size + self.page_size) * self._write_loc_dcp_span, "set_mla_kv_buffer (MLA)", ) - maybe_detect_kernel_facing_loc( - loc, self.page_size, self.kernel_page_blocks, "set_mla_kv_buffer (MLA)" - ) layer_id = ( layer_id_override if layer_id_override is not None else layer.layer_id ) diff --git a/python/sglang/srt/mem_cache/pool_host/unified.py b/python/sglang/srt/mem_cache/pool_host/unified.py index f3aa017274c1..b1841645705c 100644 --- a/python/sglang/srt/mem_cache/pool_host/unified.py +++ b/python/sglang/srt/mem_cache/pool_host/unified.py @@ -788,13 +788,6 @@ def prepare_transfer_indices( host_indices = self._shared_backing.translate_indices( self.pool_label, host_indices ) - # The L2 engine resolves device IDs using the per-layer HiCache contract. - # This pool copies whole envelopes, so discard the kernel-view page stride. - kernel_page_size = self.page_size * self.device_pool.kernel_page_blocks - device_indices = ( - device_indices // kernel_page_size * self.page_size - + device_indices % self.page_size - ) if io_backend == "kernel": if not host_indices.is_cuda: host_indices = host_indices.to( diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index 021930941bb4..741914360be3 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -112,19 +112,6 @@ def get_dtype(self) -> torch.dtype: """Storage dtype (informational). Multi-dtype subclasses return the dominant buffer's.""" raise NotImplementedError - def view_tail_pad_bytes(self, page_size: int) -> int: - """Bytes this sub-pool's views reach PAST its last page envelope.""" - return 0 - - def blocks_per_page(self) -> int: - """Row-blocks one page holds in this sub-pool's kernel-facing id space. - - The page envelope is a uniform array of equally wide row-blocks, so a - kernel-facing id is the physical page scaled by this count (see - `MultiEndedAllocator.translate_kv_loc_for_kernel`). 1 means the kernel-facing ids are the physical ones. - """ - return 1 - @dataclass(frozen=True, kw_only=True) class MHASubPoolSpec(SubPoolSpec): @@ -195,13 +182,6 @@ def layout(self) -> DenseEntryLayout: ), ) - def view_tail_pad_bytes(self, page_size: int) -> int: - return page_size * self.entry_bytes() - - def blocks_per_page(self) -> int: - """Row-blocks per page in the kernel-facing id space (one K + one V per layer).""" - return 2 * self.layer_num - def get_dtype(self) -> torch.dtype: return self.store_dtype @@ -254,14 +234,6 @@ def layout(self) -> DenseEntryLayout: ), ) - def view_tail_pad_bytes(self, page_size: int) -> int: - return page_size * self.entry_bytes() - - def blocks_per_page(self) -> int: - """One latent row per layer, so L blocks per page (MHA has 2L: a K - block and a V block per layer).""" - return self.layer_num - def get_dtype(self) -> torch.dtype: return self.store_dtype diff --git a/python/sglang/srt/utils/async_probe.py b/python/sglang/srt/utils/async_probe.py index be514dc4e5fc..abe25cd9a0e5 100644 --- a/python/sglang/srt/utils/async_probe.py +++ b/python/sglang/srt/utils/async_probe.py @@ -141,27 +141,6 @@ def maybe_detect_oob(indices: Optional[torch.Tensor], low: int, high: int, msg: ) -def maybe_detect_kernel_facing_loc( - indices: Optional[torch.Tensor], page_size: int, blocks_per_page: int, msg: str -): - """Async check that a write loc is in the pool's KERNEL-FACING id space. - - A kernel-facing id is `phys_page * (page_size * blocks_per_page) + offset` - with `offset < page_size`, so its remainder modulo the page stride is - below page_size; a VIRTUAL id satisfies that only in the first block. - Vacuous at blocks_per_page 1. Virtual ids are in range for the OOB probe, - so this is the only check that separates them. - """ - if blocks_per_page <= 1 or not envs.SGLANG_ENABLE_ASYNC_ASSERT.get(): - return - if indices is None or indices.numel() == 0: - return - torch._assert_async( - (indices % (page_size * blocks_per_page) < page_size).all(), - f"write loc outside the kernel-facing id space (virtual ids?): {msg}", - ) - - def maybe_detect_page_aligned( indices: Optional[torch.Tensor], page_size: int, msg: str ): diff --git a/test/registered/unit/mem_cache/test_full_loc_fast_path.py b/test/registered/unit/mem_cache/test_full_loc_fast_path.py index f72ebe024e5c..d78877a55056 100644 --- a/test/registered/unit/mem_cache/test_full_loc_fast_path.py +++ b/test/registered/unit/mem_cache/test_full_loc_fast_path.py @@ -297,7 +297,6 @@ def _bare_mla_pool(self): pool = object.__new__(MLATokenToKVPool) pool.size = 64 pool.page_size = 1 - pool.kernel_page_blocks = 1 pool.start_layer = 0 pool.dtype = torch.float16 pool.store_dtype = torch.float16 diff --git a/test/registered/unit/mem_cache/test_swa_cpu_copy_filter.py b/test/registered/unit/mem_cache/test_swa_cpu_copy_filter.py index 3927adad88d7..1fb754090914 100644 --- a/test/registered/unit/mem_cache/test_swa_cpu_copy_filter.py +++ b/test/registered/unit/mem_cache/test_swa_cpu_copy_filter.py @@ -87,9 +87,7 @@ def test_offload_uses_physical_swa_pages(self): expected_pages = allocator.swa_v2p_page_table[ indices[tombstone_tokens::4] // 4 ] - pages = transfers[0].device_indices[::4] // ( - 4 * manager.kv_cache.swa_kv_pool.kernel_page_blocks - ) + pages = transfers[0].device_indices[::4] // 4 buffer = manager.kv_cache.swa_kv_pool.get_page_envelope_buffer() self.assertTrue( torch.equal(buffer[pages].clone(), buffer[expected_pages]) diff --git a/test/registered/unit/model_executor/test_unified_out_cache_loc_rebind.py b/test/registered/unit/model_executor/test_unified_out_cache_loc_rebind.py index 4cc9b006e515..17c3f1649f29 100644 --- a/test/registered/unit/model_executor/test_unified_out_cache_loc_rebind.py +++ b/test/registered/unit/model_executor/test_unified_out_cache_loc_rebind.py @@ -151,26 +151,6 @@ def test_pad_lanes_derive_to_sink_and_slices_stay_pointwise(self): sub = src._swa_write_loc_unified(fb.out_cache_loc[1:5]) self.assertTrue(torch.equal(sub, loc[1:5])) - def test_the_probe_separates_kernel_facing_from_virtual_ids(self): - """A skipped rebind is the failure mode this contract has no other - guard against: virtual ids stay inside the OOB probe's bounds (they are - `blocks_per_page` times SMALLER than a kernel-facing id), so the store lands on - the wrong slots and only the output is wrong. The kernel-facing probe - is what separates them -- the in-page offset of a kernel-facing id is always - below page_size, and a virtual id's is not unless it happens to fall in - the first block.""" - for page_size, blocks in ((1, 8), (4, 6)): - with self.subTest(page_size=page_size, blocks=blocks): - stride = page_size * blocks - virt = torch.arange(1, 2 * stride, dtype=torch.int64) - kernel = (virt // page_size) * stride + virt % page_size - in_space = kernel % stride < page_size - self.assertTrue(bool(in_space.all()), "kernel-facing ids must pass") - # Virtual ids pass only in the first block; that is why the - # probe needs a batch, not one id, to be conclusive. - caught = ~(virt % stride < page_size) - self.assertTrue(bool(caught.any()), "virtual ids must be caught") - def test_empty_loc_rebinds_to_empty(self): src = _armed_source( torch.arange(8, dtype=torch.int64), torch.arange(8, dtype=torch.int64) From e7ae772889a21e2501fe578197a54954b9ca0ffe Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Mon, 28 Sep 2026 06:31:04 -0700 Subject: [PATCH 04/15] fix(unified-memory): refuse trtllm_mha decode cuda graphs under the unified pool Under the unified pool, trtllm_mha refills its cuda-graph page table out of graph, before each replay, bounded by `cache_seqlens_int32`. That buffer is written only by the fused metadata kernel recorded inside the graph, so at refill time it still holds the previous replay's lengths: once a sequence grows past them, the kernel reads page-table columns that were never translated. Decoding then produces garbage without any error (GSM8K drops to ~0 with graphs on and is at baseline with them off). Until the refill derives its lengths from the batch, refuse the combination at startup and point at `--disable-cuda-graph` or another decode attention backend. trtllm_mha as the prefill backend only, or with decode capture disabled, still boots. --- python/sglang/srt/arg_groups/kv_cache_hook.py | 15 +++++++++ .../test_unified_prefill_cuda_graph_gate.py | 32 +++++++++++++++++-- 2 files changed, 45 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index 09d9966245d4..dbb775f7f9e6 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -608,6 +608,21 @@ def handle_unified_memory_pool(server_args: Any) -> None: sorted(full_cg_backends), sorted(backends), ) + # trtllm_mha refills its cuda-graph page table out of graph, before the + # replay, from `cache_seqlens_int32` -- which only the recorded in-graph + # metadata kernel writes. The refill therefore translates the previous + # replay's lengths, and the kernel reads untranslated columns once a + # sequence grows past them. Refuse the combination until that refill + # derives its lengths from the batch. + _, decode_backend = attention_backends_of(resolved_view(server_args)) + if decode_backend == "trtllm_mha": + assert _cg_cfg is None or _cg_cfg.decode.backend == Backend.DISABLED, ( + "--enable-unified-memory does not yet support decode cuda graphs " + "with the trtllm_mha attention backend: its replay refills the page " + "table from the previous replay's sequence lengths. Pass " + "--disable-cuda-graph, or pick another decode attention backend " + "(fa3 / fa4 / flashinfer / triton)." + ) def _validate_unified_memory_dcp(server_args: Any) -> None: diff --git a/test/registered/unit/server_args/test_unified_prefill_cuda_graph_gate.py b/test/registered/unit/server_args/test_unified_prefill_cuda_graph_gate.py index ac79697f078a..df2e375270fe 100644 --- a/test/registered/unit/server_args/test_unified_prefill_cuda_graph_gate.py +++ b/test/registered/unit/server_args/test_unified_prefill_cuda_graph_gate.py @@ -30,6 +30,10 @@ would fail on a flag the user never typed), and decode capture is never touched. +Decode capture has one refusal of its own: trtllm_mha refills its graph page +table before the replay from lengths only the recorded kernel writes, so a +replay translates the previous step's lengths. That combination does not boot. + python -m pytest test/registered/unit/server_args/test_unified_prefill_cuda_graph_gate.py -v """ @@ -47,12 +51,12 @@ register_cpu_ci(est_time=8, suite="base-a-test-cpu") -def _run_handler(*, prefill_backend, attention_backends): +def _run_handler(*, prefill_backend, attention_backends, decode_backend=Backend.FULL): """Run just `handle_unified_memory_pool` over a minimal stand-in.""" sa = ServerArgs(model_path="dummy") cg = SimpleNamespace( prefill=SimpleNamespace(backend=prefill_backend), - decode=SimpleNamespace(backend=Backend.FULL), + decode=SimpleNamespace(backend=decode_backend), ) for name, value in { "enable_unified_memory": True, @@ -109,5 +113,29 @@ def test_already_disabled_is_a_no_op(self): self.assertEqual(cg.decode.backend, Backend.FULL) +class TestUnifiedTrtllmMhaDecodeGraphGate(unittest.TestCase): + def test_trtllm_mha_decode_capture_is_refused(self): + with self.assertRaisesRegex(AssertionError, "trtllm_mha"): + _run_handler( + prefill_backend=Backend.BREAKABLE, + attention_backends=("trtllm_mha", "trtllm_mha"), + ) + + def test_trtllm_mha_without_decode_capture_boots(self): + cg = _run_handler( + prefill_backend=Backend.BREAKABLE, + attention_backends=("trtllm_mha", "trtllm_mha"), + decode_backend=Backend.DISABLED, + ) + self.assertEqual(cg.decode.backend, Backend.DISABLED) + + def test_trtllm_mha_prefill_with_another_decode_backend_boots(self): + cg = _run_handler( + prefill_backend=Backend.BREAKABLE, + attention_backends=("trtllm_mha", "fa3"), + ) + self.assertEqual(cg.decode.backend, Backend.FULL) + + if __name__ == "__main__": unittest.main() From 95530483cacb4f3368e73a8d97c97bf513c3c820 Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Mon, 28 Sep 2026 06:40:20 -0700 Subject: [PATCH 05/15] test(unified-memory): keep main's test edits in the token-major rewrites The token-major test rewrites had been made on an older copy of these files and undid later upstream edits along the way. Restore them, keeping only the changes the layout itself requires: - `est_time` goes back to main's values in the four retimed files; - the tests pruned as redundant stay pruned: `TestMHASpecLayerOffsets` and `test_page_bytes` (both restate `test_layout_parts_match_the_entry_formula` and a one-line formula), the write-loc `..._without_handover` case, and the MLA `test_entry_bytes_and_dim` and `test_kernel_id_follows_compaction` (the view-addressing and `test_kernel_id_is_the_physical_id` cases pin the same things). The move readback is folded back into `test_move_kv_cache_moves_page_envelopes`, as on main; - docstrings keep main's wording where the layout does not change their meaning (`reads_are_translated`, not `enabled`; no plain page-major arm and no `cutlass_mla` in the allowlist docstring), and drop the retired 2L multiplier regime and the one-element loop it left behind. Also delete `mha_entry_bytes` / `mla_entry_bytes`: nothing outside the tests called them, and a test comparing the spec against a second copy of the same formula checks nothing. The tests now derive the entry size in closed form from the row sizes and `ENTRY_ALIGN_BYTES`. --- .../sglang/srt/mem_cache/layout/page_major.py | 14 -- .../mem_cache/test_kv_index_translator.py | 225 +++++++----------- .../unit/mem_cache/test_layout_compat.py | 66 +---- .../test_pd_envelope_transfer_layout.py | 35 ++- .../unit/mem_cache/test_unified_mha_views.py | 63 ++--- .../unit/mem_cache/test_unified_mla_views.py | 106 +++------ .../test_page_major_backend_allowlist.py | 44 ++-- 7 files changed, 173 insertions(+), 380 deletions(-) diff --git a/python/sglang/srt/mem_cache/layout/page_major.py b/python/sglang/srt/mem_cache/layout/page_major.py index 2141f5e37304..b7b0b629d2c4 100644 --- a/python/sglang/srt/mem_cache/layout/page_major.py +++ b/python/sglang/srt/mem_cache/layout/page_major.py @@ -154,20 +154,6 @@ def build_dense_views( return views -def mha_entry_bytes( - *, layer_num: int, head_num: int, head_dim: int, v_head_dim: int, itemsize: int -) -> int: - """Bytes occupied by one slot across all layers (K and V), aligned.""" - k_row_bytes = head_num * head_dim * itemsize - v_row_bytes = head_num * v_head_dim * itemsize - return align_entry_bytes(layer_num * (k_row_bytes + v_row_bytes)) - - -def mla_entry_bytes(*, layer_num: int, kv_cache_dim: int, itemsize: int) -> int: - """Bytes occupied by one MLA slot across all layers (latent rows), aligned.""" - return align_entry_bytes(layer_num * kv_cache_dim * itemsize) - - def mamba_entry_bytes( *, layer_num: int, diff --git a/test/registered/unit/mem_cache/test_kv_index_translator.py b/test/registered/unit/mem_cache/test_kv_index_translator.py index de767f56f694..bae4e724c9e9 100644 --- a/test/registered/unit/mem_cache/test_kv_index_translator.py +++ b/test/registered/unit/mem_cache/test_kv_index_translator.py @@ -13,34 +13,15 @@ # ============================================================================== """KVIndexTranslator -- the read-path id translator. -Covers, CPU-only (the builder's pure-torch reference path; GPU parity of the -Triton kernel is a later CUDA CI pin): - - strict passthrough: a non-unified source returns the SAME req_to_token / - req_pool_indices objects -- zero tensor ops, no copies (the property that - makes backend re-pointing byte-identical for every non-unified server); - - static SWA pools keep their legacy full->swa mapping on the view; - - the read table matches the hand formula - entry[b, c] = clamp(v2p[req_to_token[req[b], c*ps] // ps] * mult, 0) - over the REAL SWA composite's tables (full AND swa, ps in {1, 4}, - multiplier in {1, 2L}), with the swa table built from VIRTUAL ids; - - sink routing: dead lanes (seq_len 0), -1 req_to_token entries, and - tombstoned v2p pages all read entry 0; - - the capture contract: buffers are zero-filled and idempotent; a refresh - updates ONLY the live prefix (stale tails and rows beyond bs keep prior - contents); the returned table is the WHOLE buffer (pointer-stable); - - the eager-view memo: a single source-resident slot keyed by batch - identity (same batch shares one build; the next batch replaces it; a - dead batch never matches); - - the two-phase write contract: the rebind touches only the full side, and - the sliding-window write loc derives POINTWISE from the kernel-facing values - (pads, slices, and fresh copies included), for both pool families. - - python -m pytest test/registered/unit/mem_cache/test_kv_index_translator.py -v +CPU-only: these exercise the builder's pure-torch reference path, not the +Triton kernel. """ +from sglang.srt.runtime_context import publish, reset_context +from sglang.srt.server_args import ServerArgs from sglang.test.ci.ci_register import register_cpu_ci -register_cpu_ci(est_time=8, suite="base-a-test-cpu") +register_cpu_ci(est_time=10, suite="base-a-test-cpu") import unittest from types import SimpleNamespace @@ -54,8 +35,6 @@ from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator, KVReadTables from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.mem_cache.unified_memory_pool import MHASubPoolSpec, UnifiedKVPool -from sglang.srt.runtime_context import publish, reset_context -from sglang.srt.server_args import ServerArgs from sglang.srt.state_capturer.base import BaseTopkCapturer from sglang.test.test_utils import CustomTestCase @@ -142,9 +121,8 @@ def setUp(self): publish(ServerArgs(model_path="dummy"), role="tokenizer") def test_non_unified_returns_same_objects(self): - """The strict-passthrough property: no copy, no branch, the exact - tensors backends read today. A regression here (any tensor op on the - non-unified path) breaks byte-identity for every static-pool server.""" + """Strict passthrough: no copy, no branch. Any tensor op on the + non-unified path breaks byte-identity for every static-pool server.""" req_to_token = torch.arange(64, dtype=torch.int64).reshape(4, 16) src = KVIndexTranslator( req_to_token=req_to_token, @@ -172,7 +150,7 @@ def test_non_unified_returns_same_objects(self): def _alloc_and_fill(allocator, ps, lens): """Allocate per-request virtual runs and write them into a fake - req_to_token; returns (req_to_token, req_pool_indices, seq_lens).""" + req_to_token.""" width = 16 * ps req_to_token = torch.full((len(lens), width), -1, dtype=torch.int64) for r, n in enumerate(lens): @@ -195,61 +173,57 @@ def setUp(self): publish(ServerArgs(model_path="dummy"), role="tokenizer") def test_read_table_matches_reference(self): - """The load-bearing formula pin: full AND swa read tables equal - the independent per-element derivation across page sizes (kernel-facing - ids are the physical ones: multiplier 1). The swa table agreeing with - a formula over VIRTUAL ids is also the never-chained-through- - full-physical proof.""" + """Both read tables must equal the independent per-element derivation + across page sizes; the swa table agreeing over VIRTUAL ids proves it is + never chained through full-physical.""" for ps in (1, 4): - for allocator in (_build_composite(ps),): - full_mult = allocator.kernel_page_multiplier - swa_mult = allocator.swa_kernel_page_multiplier - req_to_token, rows, seq_lens = _alloc_and_fill( - allocator, ps, lens=[5 * ps, 2 * ps, 3 * ps - 1] - ) - src = _make_source(allocator, req_to_token, ps) - self.assertTrue(src.is_translating) - width = 6 - view = src.build_index_table( - req_pool_indices=rows, seq_lens=seq_lens, max_pages=width - ) - self.assertTrue(view.is_translated) - self.assertEqual(view.entry_page_size, ps) - self.assertTrue( - torch.equal(view.row_ids, torch.arange(3, dtype=torch.int64)) - ) - want_full = _reference_table( - req_to_token, - rows, - seq_lens, - allocator.full_v2p_page_table, - full_mult, - ps, - width, - ) - want_swa = _reference_table( - req_to_token, - rows, - seq_lens, - allocator.swa_v2p_page_table, - swa_mult, - ps, - width, - ) - self.assertTrue( - torch.equal(view.ids, want_full), - f"full read table off-formula (ps={ps}, mult={full_mult})", - ) - self.assertTrue( - torch.equal(view.sliding_window_ids, want_swa), - f"swa read table off-formula (ps={ps}, mult={swa_mult})", - ) + allocator = _build_composite(ps) + full_mult = allocator.kernel_page_multiplier + swa_mult = allocator.swa_kernel_page_multiplier + req_to_token, rows, seq_lens = _alloc_and_fill( + allocator, ps, lens=[5 * ps, 2 * ps, 3 * ps - 1] + ) + src = _make_source(allocator, req_to_token, ps) + self.assertTrue(src.is_translating) + width = 6 + view = src.build_index_table( + req_pool_indices=rows, seq_lens=seq_lens, max_pages=width + ) + self.assertTrue(view.is_translated) + self.assertEqual(view.entry_page_size, ps) + self.assertTrue( + torch.equal(view.row_ids, torch.arange(3, dtype=torch.int64)) + ) + want_full = _reference_table( + req_to_token, + rows, + seq_lens, + allocator.full_v2p_page_table, + full_mult, + ps, + width, + ) + want_swa = _reference_table( + req_to_token, + rows, + seq_lens, + allocator.swa_v2p_page_table, + swa_mult, + ps, + width, + ) + self.assertTrue( + torch.equal(view.ids, want_full), + f"full read table off-formula (ps={ps}, mult={full_mult})", + ) + self.assertTrue( + torch.equal(view.sliding_window_ids, want_swa), + f"swa read table off-formula (ps={ps}, mult={swa_mult})", + ) def test_packed_stream_equals_the_rectangle_it_replaces(self): - """The packed builder and the rectangle must agree element for element: - packed[indptr[b] + p] == ids[b, p // ps] * ps + p % ps. Consumers that - plan over the stream and consumers that read the page table have to see - the same KV, so a change to either builder alone turns this red.""" + """The two builders must agree element for element: + packed[indptr[b] + p] == ids[b, p // ps] * ps + p % ps.""" for ps in (1, 4): allocator = _build_composite(ps) req_to_token, rows, seq_lens = _alloc_and_fill( @@ -280,9 +254,9 @@ def test_packed_stream_equals_the_rectangle_it_replaces(self): ) def test_sink_routing(self): - """Dead lanes (seq_len 0), -1 slots inside the live prefix, and - tombstoned v2p pages must ALL read entry 0 -- one wild entry is a - captured-graph OOB read at replay.""" + """Dead lanes, -1 slots inside the live prefix, and tombstoned v2p + pages must ALL read entry 0; one wild entry is a captured-graph OOB + read at replay.""" ps = 4 allocator = _build_composite(ps) req_to_token, rows, seq_lens = _alloc_and_fill( @@ -317,12 +291,9 @@ def setUp(self): publish(ServerArgs(model_path="dummy"), role="tokenizer") def test_prefix_filled_tail_sentinel_preserved_width_capped(self): - """Three contracts in one batch: entries equal the read-table formula, - lanes past each row's live pages keep the backend's -1 sentinel - (prefix-only -- a tail write scatters the trtllm sentinel contract), - and a table padded WIDER than the req_to_token page span (trtllm's - LCM alignment) is capped instead of tripping the builder's width - assert.""" + """The -1 tail sentinel belongs to the backend, and a table padded + WIDER than the req_to_token page span (trtllm's LCM alignment) must be + capped rather than trip the builder's width assert.""" ps = 4 allocator = _build_composite(ps) full_mult = allocator.kernel_page_multiplier @@ -356,8 +327,8 @@ def test_prefix_filled_tail_sentinel_preserved_width_capped(self): ) def test_passthrough_source_refuses(self): - """Callers dispatch on `enabled`; a passthrough source has no v2p to - build from and must fail loud, not fill garbage.""" + """Callers dispatch on `reads_are_translated`; a passthrough source has + no v2p to build from and must fail loud, not fill garbage.""" src = KVIndexTranslator( req_to_token=torch.zeros((2, 4), dtype=torch.int64), token_to_kv_pool_allocator=SimpleNamespace(), @@ -377,13 +348,11 @@ class TestPoolOwnership(unittest.TestCase): """A runner only gets the kernel-facing id space when the pool IT reads and writes is the one the allocator's ids address. - Guarded shape: a runner handed a SHARED allocator (one slot index space, - one req_to_token) while owning a SEPARATE KV buffer sized to the - allocator's SLOT count. Probing the allocator alone reports "unified" for - that runner, so its indices would be mapped into the composite's - kernel-facing space (kernel-facing ids up to num_pages * multiplier) and then used - to address a buffer with only num_slots rows -- out of bounds on both the - read gather and the KV store. + Guarded shape: a runner handed a SHARED allocator while owning a SEPARATE + KV buffer sized to the allocator's SLOT count. Probing the allocator alone + reports "unified" for that runner, so its indices would be mapped into the + composite's kernel-facing space (ids up to num_pages * multiplier) and used + to address a buffer with only num_slots rows. """ def setUp(self): @@ -393,11 +362,10 @@ def setUp(self): publish(ServerArgs(model_path="dummy"), role="tokenizer") def test_real_factory_bundle_satisfies_the_ownership_identity(self): - """The guard rests on `allocator.get_kvcache() is token_to_kv_pool` - holding for a REAL target bundle. If a factory ever returned a pool - the allocator does not hold, the guard would silently disable the - unified path for EVERY model -- so pin it against the real factory - rather than against this file's own construction.""" + """The guard rests on `allocator.get_kvcache() is token_to_kv_pool`, so + a factory returning a pool the allocator does not hold would silently + disable the unified path for EVERY model. Pinned against the real + factory, not this file's own construction.""" from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools bundle = init_unified_swa_pools( @@ -447,9 +415,9 @@ def test_runner_with_its_own_pool_is_disabled(self): self.assertFalse(src.is_translating) def test_disabled_source_is_strict_passthrough(self): - """Consequence of the guard: such a runner must see RAW virtual ids on - the read table -- they index its own pool directly. A translate here is - the out-of-bounds bug the ownership identity exists to prevent.""" + """Such a runner must see RAW virtual ids: they index its own pool + directly, and a translate here is the out-of-bounds bug the ownership + identity exists to prevent.""" alloc = _build_composite(ps=1) req_to_token = torch.arange(16, dtype=torch.int32, device=_DEV).view(2, 8) src = KVIndexTranslator( @@ -510,10 +478,9 @@ def test_caller_owned_table_is_returned_whole_and_filled_prefix_only(self): self.assertTrue(bool((cap[1:] == 7).all()), "rows beyond bs were touched") def test_row_ids_not_reallocated_across_builds(self): - """`row_ids` is a constant arange sized once from the request pool, so - builds at different batch sizes hand back slices of ONE buffer. A - per-build `torch.arange` would be correct but would spend an allocation - and a launch on every replay prep.""" + """`row_ids` is one arange sized from the request pool and sliced per + build; a per-build `torch.arange` would be correct but costs an + allocation and a launch on every replay prep.""" allocator = _build_composite(1) req_to_token, rows, seq_lens = _alloc_and_fill(allocator, 1, lens=[4, 2, 3]) src = _make_source(allocator, req_to_token, 1) @@ -558,10 +525,9 @@ def __init__( class TestViewMemo(unittest.TestCase): - """The eager view is memoized ON THE SOURCE in a single slot keyed by - batch identity -- per-batch state stays out of the ForwardBatch (it does - not scale with the number of id spaces), and one metadata build's many - consumers still share one table build.""" + """The eager view is memoized ON THE SOURCE in a single slot keyed by batch + identity, so per-batch state stays out of the ForwardBatch while one + metadata build's many consumers still share one table build.""" def setUp(self): # The code under test reads its config from the bags. @@ -777,10 +743,8 @@ def test_topk_capture_round_trips_request_token_ids(self): ) def test_swa_write_loc_round_trips_from_full_side(self): - """The derived property behind phase 2: for any virtual run t, - deriving from the kernel-facing full-side values must equal the direct - virtual->swa translate -- `field(full(t)) == swa(t)` across page sizes - and multipliers.""" + """Derived property: `field(full(t)) == swa(t)` for any virtual run t, + across page sizes.""" for ps in (1, 4, 64): src, _, rows, seq_lens, _, want_full, want_swa = self._built( ps=ps, n=3 * ps @@ -789,30 +753,15 @@ def test_swa_write_loc_round_trips_from_full_side(self): self.assertTrue(torch.equal(got, want_swa)) def test_pad_lanes_derive_to_sink(self): - """The DP pad appends zeros; kernel-facing 0 is the reserved padding slot in - every id space, so pad lanes must derive to swa slot 0 with no - `num_live` bookkeeping.""" + """The DP pad appends zeros, and kernel-facing 0 is the reserved + padding slot in every id space, so pad lanes derive to swa slot 0 with + no `num_live` bookkeeping.""" src, _, rows, seq_lens, _, want_full, want_swa = self._built(n=3) padded = torch.cat([want_full, want_full.new_zeros(2)]) got = self._field(src, rows, seq_lens, padded) self.assertTrue(torch.equal(got[:3], want_swa)) self.assertTrue(bool((got[3:] == 0).all()), "pad lanes must land on slot 0") - def test_slice_and_copy_derive_pointwise_without_handover(self): - """REGRESSION (design): the retired identity-resolver refused any - tensor it had not been handed -- a TBO child's re-padded slice or a - registry's fresh copy raised. Value-based derivation must accept - both, pointwise, with no adopt/handover call.""" - src, _, rows, seq_lens, _, want_full, want_swa = self._built(n=4) - padded = torch.cat([want_full, want_full.new_zeros(2)]) - # TBO-child shape: a slice crossing the pad boundary. - got = self._field(src, rows, seq_lens, padded[2:6]) - self.assertTrue(torch.equal(got[:2], want_swa[2:4])) - self.assertTrue(bool((got[2:] == 0).all())) - # Registry shape: a fresh equal-value copy. - got2 = self._field(src, rows, seq_lens, want_full.clone()) - self.assertTrue(torch.equal(got2, want_swa)) - def test_tombstoned_swa_page_clamps_to_sink(self): src, allocator, rows, seq_lens, virt, want_full, _ = self._built(ps=1, n=2) allocator.swa_v2p_page_table[int(virt[0])] = -1 diff --git a/test/registered/unit/mem_cache/test_layout_compat.py b/test/registered/unit/mem_cache/test_layout_compat.py index 25c5d33b22f8..19a5e836da36 100644 --- a/test/registered/unit/mem_cache/test_layout_compat.py +++ b/test/registered/unit/mem_cache/test_layout_compat.py @@ -11,81 +11,29 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Unit tests for the page-major envelope byte layout. +"""``move_kv_cache_native`` -- the stock per-layer 3-D move static-pool +compaction rides on -- must stay byte-exact. CPU-only. -The subject here is the ENVELOPE -- the byte layout the unified pool stores its -KV in -- pinned through ``MHASubPoolSpec``'s offset math. The 3-D per-layer views -the pool exposes over the same bytes are covered by -``test_unified_mha_views.py``, which also pins the view addressing -against the envelope formula byte for byte. - -Verifies that: -1. ``MHASubPoolSpec.layout()`` -- the descriptor every per-layer view is - built through -- places each layer's K/V at the token-major entry offsets, - independent of the page size. -2. ``move_kv_cache_native`` (the stock per-layer 3-D move) stays byte-exact. - -CPU-only -- no GPU / Triton needed. - - python -m pytest test/registered/unit/mem_cache/test_layout_compat.py -v +The page-major envelope layout and the per-layer views over it are covered by +``test_unified_mha_views.py``, which pins the view addressing against the +envelope formula byte for byte. """ from sglang.test.ci.ci_register import register_cpu_ci -register_cpu_ci(est_time=6, suite="base-a-test-cpu") +register_cpu_ci(est_time=9, suite="base-a-test-cpu") import unittest import torch from sglang.srt.mem_cache.memory_pool import move_kv_cache_native -from sglang.srt.mem_cache.unified_memory_pool import MHASubPoolSpec - -_DEV = "cpu" - - -def _make_mha_spec(name, grow, layer_num=2, head_num=2, head_dim=4): - return MHASubPoolSpec( - name=name, - layer_num=layer_num, - head_num=head_num, - head_dim=head_dim, - store_dtype=torch.float16, - grow_direction=grow, - ) - - -class TestMHASpecLayerOffsets(unittest.TestCase): - """Verify the layout descriptor's part offsets against the entry math.""" - - def test_offsets_match_token_major_entry(self): - spec = _make_mha_spec("full", "up", layer_num=3, head_num=2, head_dim=4) - # Layer L's K sits at L * (k_row + v_row) inside the slot's entry and - # its V right after that K. The page size never enters: a page is - # page_size entries back to back. - k_row = spec.k_row_bytes() - v_row = spec.v_row_bytes() - layout = spec.layout() - self.assertEqual(layout.entry_bytes, spec.entry_bytes()) - for L in range(spec.layer_num): - self.assertEqual( - layout.part("k").layer_offset_bytes(L), L * (k_row + v_row) - ) - self.assertEqual( - layout.part("v").layer_offset_bytes(L), L * (k_row + v_row) + k_row - ) - - def test_page_bytes(self): - spec = _make_mha_spec("full", "up", layer_num=3, head_num=2, head_dim=4) - # page_bytes = page_size * entry_bytes (preserved invariant) - for ps in [1, 8, 64, 256]: - self.assertEqual(spec.page_bytes(ps), ps * spec.entry_bytes()) class TestMoveKVCacheNative(unittest.TestCase): def test_move_kv_cache_3d_path_unchanged(self): """The stock per-layer 3-D move must relocate exactly the named token - rows, byte-identically -- compaction on static pools rides on it.""" + rows, byte-identically; compaction on static pools rides on it.""" k = [torch.zeros((32, 2, 4), dtype=torch.float16) for _ in range(2)] v = [torch.zeros((32, 2, 4), dtype=torch.float16) for _ in range(2)] for L in range(2): diff --git a/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py b/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py index 58d3d08c9a87..a6afaea907b4 100644 --- a/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py +++ b/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py @@ -18,13 +18,12 @@ import torch from sglang.srt.mem_cache.layout.page_major import ( + ENTRY_ALIGN_BYTES, DenseEntryLayout, DensePart, build_dense_views, build_page_major_mamba_views, mamba_entry_bytes, - mha_entry_bytes, - mla_entry_bytes, ) from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -32,6 +31,14 @@ register_cpu_ci(est_time=11, suite="base-a-test-cpu") +def _entry_bytes(row_bytes_per_layer, layer_num): + """One slot's rows across every layer, rounded up to the entry alignment -- + derived here, not read back from the layout under test.""" + return ( + -(-(row_bytes_per_layer * layer_num) // ENTRY_ALIGN_BYTES) * ENTRY_ALIGN_BYTES + ) + + class TestMLAEnvelopeTransferAddressing(CustomTestCase): def test_page_envelope_matches_per_layer_views(self): """Every (page, layer, slot) row written through the MLA views must @@ -40,9 +47,7 @@ def test_page_envelope_matches_per_layer_views(self): layer_num, page_size, kv_dim, num_pages = 3, 4, 8, 6 store_dtype = torch.bfloat16 row_bytes = kv_dim * store_dtype.itemsize - entry_bytes = mla_entry_bytes( - layer_num=layer_num, kv_cache_dim=kv_dim, itemsize=store_dtype.itemsize - ) + entry_bytes = _entry_bytes(row_bytes, layer_num) page_bytes = page_size * entry_bytes layout = DenseEntryLayout( entry_bytes=entry_bytes, @@ -93,15 +98,9 @@ class TestMHAEnvelopeTransferAddressing(CustomTestCase): def test_page_envelope_matches_per_layer_views(self): layer_num, page_size, head_num, head_dim, num_pages = 3, 4, 2, 8, 6 store_dtype = torch.bfloat16 - entry_bytes = mha_entry_bytes( - layer_num=layer_num, - head_num=head_num, - head_dim=head_dim, - v_head_dim=head_dim, - itemsize=store_dtype.itemsize, - ) - page_bytes = page_size * entry_bytes row_bytes = head_num * head_dim * store_dtype.itemsize + entry_bytes = _entry_bytes(2 * row_bytes, layer_num) + page_bytes = page_size * entry_bytes self.assertEqual(page_bytes, page_size * 2 * layer_num * row_bytes) layout = DenseEntryLayout( @@ -166,15 +165,9 @@ def test_envelope_move_is_a_whole_page_copy(self): makes a physical page id a valid PD transfer index after compaction.""" layer_num, page_size, head_num, head_dim, num_pages = 2, 2, 1, 8, 4 store_dtype = torch.bfloat16 - entry_bytes = mha_entry_bytes( - layer_num=layer_num, - head_num=head_num, - head_dim=head_dim, - v_head_dim=head_dim, - itemsize=store_dtype.itemsize, - ) - page_bytes = page_size * entry_bytes row_bytes = head_num * head_dim * store_dtype.itemsize + entry_bytes = _entry_bytes(2 * row_bytes, layer_num) + page_bytes = page_size * entry_bytes layout = DenseEntryLayout( entry_bytes=entry_bytes, parts=( diff --git a/test/registered/unit/mem_cache/test_unified_mha_views.py b/test/registered/unit/mem_cache/test_unified_mha_views.py index 9d0fa45ed7cc..d8ef9721c97c 100644 --- a/test/registered/unit/mem_cache/test_unified_mha_views.py +++ b/test/registered/unit/mem_cache/test_unified_mha_views.py @@ -11,31 +11,18 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""MHA K/V views for the unified memory pool (token-major dense views). - -Covers, CPU-only (pure torch -- no GPU / Triton kernels): - - `MHASubPoolSpec.layout()`: K and V of every layer sit at fixed offsets - inside one slot entry, K and V rows may differ in width, alignment is - enforced at construction; - - `build_dense_views` addressing: view_l[t] for PHYSICAL token t must land - exactly at the byte the envelope formula assigns to (page, slot, layer, - K|V) -- cross-checked against an independent 4-D (page, slot) description; - - K and V of one token share ONE id (the per-layer origin shift does the - disambiguation), with no aliasing across the 2L views; - - a too-short buffer and misaligned rows fail loud at construction. +"""MHA K/V views for the unified memory pool (token-major dense views), CPU-only. Addressing law under test: byte(t, l, K) = t * entry_bytes + l * (k_row + v_row) byte(t, l, V) = byte(t, l, K) + k_row t = page * page_size + slot (the physical token id IS the kernel id) - - python -m pytest test/registered/unit/mem_cache/test_unified_mha_views.py -v """ from sglang.test.ci.ci_register import register_cpu_ci -register_cpu_ci(est_time=8, suite="base-a-test-cpu") +register_cpu_ci(est_time=11, suite="base-a-test-cpu") import unittest from types import SimpleNamespace @@ -46,7 +33,6 @@ from sglang.srt.mem_cache.layout.page_major import ( ENTRY_ALIGN_BYTES, build_dense_views, - mha_entry_bytes, paged_view, ) from sglang.srt.mem_cache.unified_memory_pool import ( @@ -57,12 +43,10 @@ _DEV = "cpu" # `set_kv_buffer` dispatches on the PLATFORM (memory_pool._is_cuda, resolved at -# import), not on the tensors it is handed, so cases driving it must build on -# the platform's device. The rest of this file is byte arithmetic, so CPU. +# import), not on the tensors it is handed, so cases driving it build there. _STORE_DEV = "cuda" if torch.cuda.is_available() else "cpu" -# Small-but-nontrivial MHA geometry: L=2 layers, H=2 heads, D=4, so every byte -# offset is hand-checkable. One entry = 2 layers x (K row + V row) = 64 B. +# Geometry kept tiny so every byte offset is hand-checkable. _L = 2 _H = 2 _D = 4 @@ -71,9 +55,16 @@ _HAS_FP8 = hasattr(torch, "float8_e4m3fn") _K_ROW = _H * _D * _ITEM _V_ROW = _H * _D * _ITEM -_ENTRY = mha_entry_bytes( - layer_num=_L, head_num=_H, head_dim=_D, v_head_dim=_D, itemsize=_ITEM -) + + +def _entry_bytes(head_dim=_D, v_head_dim=_D): + """One slot's K and V rows across every layer, rounded up to the entry + alignment -- derived here, not read back from the spec under test.""" + rows = _L * _H * (head_dim + v_head_dim) * _ITEM + return -(-rows // ENTRY_ALIGN_BYTES) * ENTRY_ALIGN_BYTES + + +_ENTRY = _entry_bytes() assert _ENTRY == _L * (_K_ROW + _V_ROW) # no alignment pad at this geometry @@ -112,13 +103,7 @@ def _reference_paged_views( against a second, independently derived description of the same bytes.""" k_row = _H * head_dim * _ITEM v_row = _H * v_head_dim * _ITEM - entry = mha_entry_bytes( - layer_num=_L, - head_num=_H, - head_dim=head_dim, - v_head_dim=v_head_dim, - itemsize=_ITEM, - ) + entry = _entry_bytes(head_dim, v_head_dim) as_dtype_view = raw.view(_DTYPE) k_views, v_views = [], [] for layer in range(_L): @@ -153,14 +138,8 @@ def test_layout_parts_match_the_entry_formula(self): layout.part("v").layer_offset_bytes(l), l * (_K_ROW + _V_ROW) + _K_ROW ) - def test_entry_bytes_matches_layout_helper_and_is_aligned(self): - spec = _mha_spec() - self.assertEqual( - spec.entry_bytes(), - mha_entry_bytes( - layer_num=_L, head_num=_H, head_dim=_D, v_head_dim=_D, itemsize=_ITEM - ), - ) + def test_entry_bytes_matches_the_closed_form_and_is_aligned(self): + self.assertEqual(_mha_spec().entry_bytes(), _ENTRY) # 48 B of rows round up to one 64 B entry; the parts still fit inside. padded = MHASubPoolSpec( name="full", @@ -273,13 +252,7 @@ def test_k_and_v_share_one_id_without_aliasing(self): def test_asymmetric_views_address_their_own_rows(self): head_dim, v_head_dim = 8, 4 k_row, v_row = _H * head_dim * _ITEM, _H * v_head_dim * _ITEM - entry = mha_entry_bytes( - layer_num=_L, - head_num=_H, - head_dim=head_dim, - v_head_dim=v_head_dim, - itemsize=_ITEM, - ) + entry = _entry_bytes(head_dim, v_head_dim) ps, num_pages = 4, 3 raw = _make_raw(ps, num_pages, entry=entry) dk, dv = _build_views(raw, ps, num_pages, head_dim, v_head_dim) diff --git a/test/registered/unit/mem_cache/test_unified_mla_views.py b/test/registered/unit/mem_cache/test_unified_mla_views.py index 0bb957585056..f87fd5c6025e 100644 --- a/test/registered/unit/mem_cache/test_unified_mla_views.py +++ b/test/registered/unit/mem_cache/test_unified_mla_views.py @@ -11,30 +11,19 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""MLA views for the unified memory pool (MLA-hybrid-Mamba, Kimi K3). - -Covers, CPU-only (pure torch -- no GPU / Triton kernels): - - `MLASubPoolSpec` byte math (aligned entry, one latent row per layer); - - `build_dense_views` addressing: view_l[t] for PHYSICAL token t must land - at the envelope byte `t * entry_bytes + l * row_bytes`, the per-layer - views must not alias at equal ids, and a short buffer must fail loud; - - `UnifiedKVPool` MLA plumbing: the allocation is exactly the budget and - the reserved sink floor covers the whole page-0 envelope; - - `UnifiedMLATokenToKVPool`: buffer wiring, V-as-prefix-slice, and the - page-envelope `move_kv_cache` (physical token ids, page-major runs); - - `MultiEndedAllocator.translate_kv_loc_for_kernel`: identical to the - physical translate, tombstone clamp to the sink, `out=` contract, the - pinned multiplier, and correctness across eager compaction. - -GPU parity of the actual read/write kernels (set_mla_kv_buffer TMA path etc.) -lives in the server-level tests, not here. - - python -m pytest test/registered/unit/mem_cache/test_unified_mla_views.py -v +"""MLA views for the unified memory pool (MLA-hybrid-Mamba, Kimi K3), CPU-only. + +Addressing law under test: the (layer, token) cell sits at envelope byte +offset `t * entry_bytes + l * row_bytes`, and the kernel-facing id of token `t` +is its physical token id. + +GPU parity of the read/write kernels (set_mla_kv_buffer TMA path etc.) lives in +`test_unified_mla_gpu_parity.py`. """ from sglang.test.ci.ci_register import register_cpu_ci -register_cpu_ci(est_time=8, suite="base-a-test-cpu") +register_cpu_ci(est_time=10, suite="base-a-test-cpu") import unittest from unittest import mock @@ -48,7 +37,6 @@ from sglang.srt.mem_cache.layout.page_major import ( ENTRY_ALIGN_BYTES, build_dense_views, - mla_entry_bytes, paged_row_view, ) from sglang.srt.mem_cache.unified_memory_pool import ( @@ -61,9 +49,8 @@ _DEV = "cpu" -# Small-but-nontrivial MLA geometry: L=3 layers, D=8 (=6+2), so every byte -# offset is hand-checkable. Real K3 is L=24, D=576 (=512+64). 3 rows x 16 B -# = 48 B of payload round up to one 64 B entry. +# Geometry kept tiny so every byte offset is hand-checkable; real K3 is +# L=24, D=576 (=512+64). 3 rows x 16 B round up to one 64 B entry. _L = 3 _LORA = 6 _ROPE = 2 @@ -71,7 +58,8 @@ _DTYPE = torch.bfloat16 _ITEM = _DTYPE.itemsize _ROW = _D * _ITEM -_ENTRY = mla_entry_bytes(layer_num=_L, kv_cache_dim=_D, itemsize=_ITEM) +# One entry = every layer's latent row, rounded up to the entry alignment. +_ENTRY = -(-(_L * _ROW) // ENTRY_ALIGN_BYTES) * ENTRY_ALIGN_BYTES _E_ELEMS = _ENTRY // _ITEM @@ -120,19 +108,6 @@ def _build_views(raw, ps, num_pages): class TestMLASubPoolSpec(unittest.TestCase): - def test_entry_bytes_and_dim(self): - spec = _mla_spec() - self.assertEqual(spec.kv_cache_dim, _D) - self.assertEqual(spec.row_bytes(), _ROW) - self.assertEqual(spec.entry_bytes(), _ENTRY) - self.assertGreaterEqual(spec.entry_bytes(), _L * _ROW) - self.assertEqual(spec.entry_bytes() % ENTRY_ALIGN_BYTES, 0) - self.assertEqual(spec.get_dtype(), _DTYPE) - layout = spec.layout() - self.assertEqual(layout.entry_bytes, _ENTRY) - for l in range(_L): - self.assertEqual(layout.part("kv").layer_offset_bytes(l), l * _ROW) - def test_rejects_nonpositive_dims(self): with self.assertRaises(AssertionError): MLASubPoolSpec( @@ -285,15 +260,24 @@ def test_cpu_copy_round_trips_through_physical_ids(self): ) def test_move_kv_cache_moves_page_envelopes(self): + """Whole page envelopes relocate, in raw bytes and (at ps=4) as read + back through the per-layer views at the destination ids.""" for ps in (1, 4): pool, kv_pool = self._make(ps=ps) num_pages = pool.max_slots("full") // ps page_bytes = ps * pool.mla_spec("full").entry_bytes() env = pool._raw[: num_pages * page_bytes].view(num_pages, page_bytes) - src_pages = torch.tensor([num_pages - 2, num_pages - 4]) - dst_pages = torch.tensor([2, 3]) + src_pages = torch.tensor([num_pages - 2, num_pages - 4, num_pages - 3]) + dst_pages = torch.tensor([2, 3, 5]) env[src_pages[0]] = 7 env[src_pages[1]] = 9 + if ps == 4: + # write through the views at src, expect it at dst after the move + for l in range(_L): + for s in range(ps): + kv_pool.kv_buffer[l][int(src_pages[2]) * ps + s] = float( + l * ps + s + 1 + ) # page-major token runs, exactly how compaction expands pages offsets = torch.arange(ps, dtype=torch.int64) src_t = (src_pages[:, None] * ps + offsets).reshape(-1) @@ -301,27 +285,13 @@ def test_move_kv_cache_moves_page_envelopes(self): kv_pool.move_kv_cache(dst_t, src_t) self.assertTrue(torch.all(env[dst_pages[0]] == 7), f"ps={ps}") self.assertTrue(torch.all(env[dst_pages[1]] == 9), f"ps={ps}") - - def test_move_then_readback(self): - ps = 4 - pool, kv_pool = self._make(ps=ps) - num_pages = pool.max_slots("full") // ps - src_page, dst_page = num_pages - 3, 5 - # write through the views at src, expect it at dst after the move - for l in range(_L): - for s in range(ps): - kv_pool.kv_buffer[l][src_page * ps + s] = float(l * ps + s + 1) - offsets = torch.arange(ps, dtype=torch.int64) - kv_pool.move_kv_cache( - (torch.tensor([dst_page])[:, None] * ps + offsets).reshape(-1), - (torch.tensor([src_page])[:, None] * ps + offsets).reshape(-1), - ) - for l in range(_L): - for s in range(ps): - got = kv_pool.kv_buffer[l][dst_page * ps + s] - self.assertTrue( - torch.all(got == float(l * ps + s + 1)), f"(l={l}, s={s})" - ) + if ps == 4: + for l in range(_L): + for s in range(ps): + got = kv_pool.kv_buffer[l][int(dst_pages[2]) * ps + s] + self.assertTrue( + torch.all(got == float(l * ps + s + 1)), f"(l={l}, s={s})" + ) class _FakeKVCache: @@ -394,20 +364,6 @@ def test_multiplier_is_pinned_to_one(self): with self.assertRaises(AssertionError): self._build(ps=1, multiplier=_L) - def test_kernel_id_follows_compaction(self): - alloc = self._build(ps=1) - a = alloc.alloc(4) - b = alloc.alloc(4) - c = alloc.alloc(4) - self.assertIsNotNone(c) - alloc.free(b) # eager compaction relocates survivors - for run in (a, c): - self.assertTrue( - torch.equal( - alloc.translate_kv_loc_for_kernel(run), alloc.translate_kv_loc(run) - ) - ) - class _RecordingHybridPool: """Stands in for `UnifiedHybridLinearKVPool`, recording the ids it is handed.""" diff --git a/test/registered/unit/server_args/test_page_major_backend_allowlist.py b/test/registered/unit/server_args/test_page_major_backend_allowlist.py index 628b62276a0e..6d71695b71a4 100644 --- a/test/registered/unit/server_args/test_page_major_backend_allowlist.py +++ b/test/registered/unit/server_args/test_page_major_backend_allowlist.py @@ -13,30 +13,19 @@ # ============================================================================== """`--enable-page-major-kv-layout` full-attention backend allowlist. -Two-way gate (see `_handle_page_major_kv_layout`), because the unified pool -exposes per-layer views and nothing else: - * unified-memory MLA models allow the whole wired - paged MLA family -- `fa3`, `flashinfer`'s MLA backend, `trtllm_mla` with - its `cutedsl_mla` / `tokenspeed_mla` subclasses, and `flashmla` (ps=64 - snap); - * unified-memory MHA/SWA models allow `fa3` / - `fa4` / `flashinfer` / `trtllm_mha` alongside Triton; - * plain `--enable-page-major-kv-layout` without the unified pool keeps the - envelope-strided 4-D views only the stride-aware Triton kernels read. - -The same handler also screens the pool itself: asymmetric K/V rows are not -admitted yet, so an asymmetric-K/V model (MiMoV2: head_dim 192 != v_head_dim -128) cannot run `--enable-unified-memory` at all and is rejected on EVERY -backend, Triton included. MLA models are exempt -- their sub-pool keeps -one latent row per layer, and several MLA configs (Kimi-Linear: head_dim 72, -v_head_dim 128) report asymmetric dims while running the unified pool today. - -Pinned here so no arm silently widens to an unwired backend (`cutlass_mla`, -`aiter`) and no arm silently narrows: `fa3` is the resolved default on -pre-Blackwell hosts, so its absence from an arm makes `--enable-unified-memory` -fail to boot under its own default configuration. - - python -m pytest test/registered/unit/server_args/test_page_major_backend_allowlist.py -v +`handle_page_major_kv_layout` gates two ways, because the per-layer views the +unified pool exposes are all the allowlisted backends can read: an MLA arm, an +MHA/SWA arm, and no page-major arm at all without the unified pool. `fa3` is +the resolved default on pre-Blackwell hosts, so its absence from an arm makes +`--enable-unified-memory` fail to boot under its own default configuration. + +The same handler screens the pool itself: asymmetric K/V rows are not admitted +yet (the token-major views can hold them, but the backends' read and write +paths are not audited for them), so an asymmetric-K/V model (MiMoV2: head_dim +192 != v_head_dim 128) is rejected on EVERY backend, Triton included. MLA +models are exempt -- their sub-pool keeps one latent row per layer, and real +MLA configs (Kimi-Linear: head_dim 72, v_head_dim 128) report asymmetric dims +while running the unified pool today. """ import unittest @@ -119,8 +108,8 @@ def test_triton_allowed_on_every_arm(self): the MLA nor the MHA arm can narrow away.""" for use_mla in (True, False): self.assertTrue(_accepts("triton", use_mla=use_mla)) - # Asymmetric K/V rows are not admitted yet; that is a property of the - # model, not of the backend, so the screen rejects even Triton. + # The asymmetric-K/V screen is a property of the model, not of the + # backend, so it rejects even Triton. self.assertFalse(_accepts("triton", use_mla=False, has_asymmetric_kv=True)) def test_per_layer_view_mla_backends_allowed_under_unified_mla(self): @@ -145,8 +134,7 @@ def test_mla_only_backends_rejected_for_mha(self): ) def test_asymmetric_kv_mha_model_cannot_use_unified_memory(self): - """head_dim != v_head_dim (MiMoV2): not admitted under the unified pool - yet. The rejection is the POOL's, not a backend's, so it must fire on + """The rejection is the POOL's, not a backend's, so it must fire on every backend -- Triton included.""" for backend in ("triton",) + self.PER_LAYER_VIEW_MHA_BACKENDS: self.assertFalse( From 7b8ac5676087417a0d46740e2c2ee4f7b71bfc51 Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Mon, 28 Sep 2026 06:42:37 -0700 Subject: [PATCH 06/15] refactor(unified-memory): drop the pass-through CPU-copy overrides With the kernel-to-physical id rewrite gone, `get_cpu_copy` and `load_cpu_copy` on `UnifiedMHATokenToKVPool` and `UnifiedMLATokenToKVPool` only forwarded to their parents. The parents already accept `mamba_indices` / `req_pool_index` and ignore them, so delete all four and let the inherited methods serve the unified pools. --- .../srt/mem_cache/unified_memory_pool.py | 20 ------------------- 1 file changed, 20 deletions(-) diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index 741914360be3..07f37921a712 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -658,16 +658,6 @@ def grow_direction(self) -> str: """Physical growth direction used by this sub-pool's L1 allocator.""" return self._unified_buffer.spec(self._sub_pool_name).grow_direction - def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): - """A physical token id IS the kernel-facing id, so the host copy needs - no rewrite; the override drops the arguments the parent ignores.""" - return super().get_cpu_copy(indices) - - def load_cpu_copy( - self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None - ): - super().load_cpu_copy(kv_cache_cpu, indices) - def set_kv_buffer_prefix_valid(self, *args, **kwargs): raise NotImplementedError( "prefix-valid commit is unsupported under the unified layout " @@ -751,16 +741,6 @@ def get_contiguous_buf_infos(self): raw = self._unified_buffer._raw return [raw.data_ptr()], [raw.numel()], [self._page_bytes] - def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): - """A physical token id IS the kernel-facing id, so the host copy needs - no rewrite; the override drops the arguments the parent ignores.""" - return super().get_cpu_copy(indices) - - def load_cpu_copy( - self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None - ): - super().load_cpu_copy(kv_cache_cpu, indices) - def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): """Relocate whole page envelopes. From d2612c31bf2750b9c19433e8421712b44d3618f4 Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Mon, 28 Sep 2026 06:42:51 -0700 Subject: [PATCH 07/15] docs(unified-memory): describe the token-major entry where the text still said layer-major The token-major layout made several descriptions false: - the `--enable-page-major-kv-layout` help and the server-arguments page said a page is "layer-major within a page"; a page is now page_size token entries, each holding every layer's K/V; - `_reserved_floor_bytes` and the matching `UnifiedKVPool` comment justified reserving all of page 0 by saying a slot-0 write touches layer blocks spread across the page. It touches only slot 0's entry; the whole page is reserved because page-aware sub-pools allocate whole pages; - `handle_page_major_kv_layout` described the unified pool as the page-major envelope layout and kept an allowlist arm for plain page-major, which the assert above it already rejects. Drop the dead arm and its comment. --- .../advanced_features/server_arguments.mdx | 2 +- python/sglang/srt/arg_groups/fields/memory.py | 8 ++++---- python/sglang/srt/arg_groups/kv_cache_hook.py | 12 ++++-------- .../srt/mem_cache/unified_memory_pool.py | 18 +++++++++--------- 4 files changed, 18 insertions(+), 22 deletions(-) diff --git a/docs/docs/advanced_features/server_arguments.mdx b/docs/docs/advanced_features/server_arguments.mdx index 8f4a52b5e008..b0892a68b5bc 100644 --- a/docs/docs/advanced_features/server_arguments.mdx +++ b/docs/docs/advanced_features/server_arguments.mdx @@ -534,7 +534,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s `--enable-page-major-kv-layout` - Enable the page-major KV layout: lay out the Mamba state and full/SWA KV caches in a page-granularity envelope (page is the outermost axis, layer-major within a page) instead of the default per-layer (layer-major) layout. Requires the Triton attention / linear-attn / Mamba backends (`--attention-backend triton`, and for hybrid models `--linear-attn-backend triton --mamba-backend triton`). + Enable the page-major KV layout: lay out the Mamba state and full/SWA KV caches in a page-granularity envelope (page is the outermost axis; within a page, each token's entry holds every layer's K/V) instead of the default per-layer (layer-major) layout. Requires the Triton attention / linear-attn / Mamba backends (`--attention-backend triton`, and for hybrid models `--linear-attn-backend triton --mamba-backend triton`). `False` bool flag (set to enable) diff --git a/python/sglang/srt/arg_groups/fields/memory.py b/python/sglang/srt/arg_groups/fields/memory.py index 992caf722244..a76165cd5afe 100644 --- a/python/sglang/srt/arg_groups/fields/memory.py +++ b/python/sglang/srt/arg_groups/fields/memory.py @@ -71,10 +71,10 @@ class Memory(msgspec.Struct): enable_page_major_kv_layout: A[ bool, "Enable the page-major KV layout: lay out the Mamba state and full/SWA " - "KV caches in a page-granularity envelope (page is the outermost axis, " - "layer-major within a page) instead of the default per-layer " - "(layer-major) layout. Requires the Triton attention / linear-attn / " - "Mamba backends.", + "KV caches in a page-granularity envelope (page is the outermost axis; " + "within a page, each token's entry holds every layer's K/V) instead of " + "the default per-layer (layer-major) layout. Requires the Triton " + "attention / linear-attn / Mamba backends.", ] = False enable_unified_memory: A[ bool, diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index dbb775f7f9e6..b7ebc1271681 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -671,9 +671,8 @@ def _validate_unified_memory_dcp(server_args: Any) -> None: def handle_page_major_kv_layout(server_args: Any): - # The unified pool stores state in the page-major envelope layout, so - # enabling it implies --enable-page-major-kv-layout — routing it through the - # single page-major path + stride-aware Triton asserts (set before the guard). + # --enable-unified-memory implies --enable-page-major-kv-layout, so the + # unified pool goes through this one gate (declared before the guard). cfg = resolving_view(server_args) if cfg.enable_unified_memory: @@ -715,9 +714,8 @@ def handle_page_major_kv_layout(server_args: Any): # snap). # * MHA/SWA models: fa3 / fa4 / flashinfer / trtllm_mha alongside # Triton. fa4 is the fa3 class. - # * Without the unified pool, plain page-major stays Triton-only. # Names are the RESOLVED ids from attention_backends_of. - if cfg.enable_unified_memory and use_mla_backend(server_args): + if use_mla_backend(server_args): allowed_full = { "triton", "fa3", @@ -727,7 +725,7 @@ def handle_page_major_kv_layout(server_args: Any): "tokenspeed_mla", "flashmla", } - elif cfg.enable_unified_memory: + else: allowed_full = { "triton", "fa3", @@ -735,8 +733,6 @@ def handle_page_major_kv_layout(server_args: Any): "flashinfer", "trtllm_mha", } - else: - allowed_full = {"triton"} backends = set(attention_backends_of(resolved_view(server_args))) backends.discard(None) assert backends <= allowed_full, ( diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index 07f37921a712..a1e2c8a3b279 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -292,12 +292,12 @@ def _reserved_floor_bytes(sub_pool_specs: List[SubPoolSpec], page_size: int) -> """Bytes at the bottom of the buffer reserved as the slot-0 padding sink. Slot-0 dummy writes for every sub-pool land here; each sub-pool's first - allocatable slot is chosen so real data starts past it. For a PAGE-AWARE - sub-pool the slot-0 write touches layer blocks spread across the whole - page-0 envelope (page_size * entry_bytes), not just one slot envelope -- - but a mamba sub-pool is page_size=1, so its entry is charged ONCE. Charging - a mamba entry per page would reserve page_size * ~100 MB of buffer that the - sink never touches. + allocatable slot is chosen so real data starts past it. A slot-0 write + touches only slot 0's entry, but a PAGE-AWARE sub-pool allocates whole + pages, so its first allocatable page starts past all of page 0 + (page_size * entry_bytes) -- while a mamba sub-pool is page_size=1, so its + entry is charged ONCE. Charging a mamba entry per page would reserve + page_size * ~100 MB of buffer that the sink never touches. Single source of truth: `UnifiedKVPool` reserves exactly this, and the factories' bs=1 feasibility floors charge exactly this. @@ -389,9 +389,9 @@ def __init__( # Slot-0 dummy writes for both pools land in the reserved low-byte sink; # each pool's first allocatable slot is chosen so real data starts past it. - # For a page-aware sub-pool the slot-0 write touches layer blocks spread - # across the WHOLE page-0 envelope (up to page_size * entry_bytes), not - # just one slot envelope — reserve the max of both. + # A page-aware sub-pool allocates whole pages, so the sink is all of page + # 0 (up to page_size * entry_bytes), not one slot entry -- reserve the + # max of both. reserved_floor = _reserved_floor_bytes(self.sub_pool_specs, page_size) for spec in self.sub_pool_specs: From a2a96ce55c63f4e764c060b7f501e59adaa1f91c Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Mon, 28 Sep 2026 06:53:30 -0700 Subject: [PATCH 08/15] test(unified-memory): import the paged-view helpers from their module The token-major view tests use `paged_view` / `paged_row_view`, which now live in `mem_cache/layout/paged_view.py`. Import lines only. --- test/registered/unit/mem_cache/test_unified_mha_views.py | 7 ++----- test/registered/unit/mem_cache/test_unified_mla_views.py | 7 ++----- 2 files changed, 4 insertions(+), 10 deletions(-) diff --git a/test/registered/unit/mem_cache/test_unified_mha_views.py b/test/registered/unit/mem_cache/test_unified_mha_views.py index d8ef9721c97c..77efaed0c9e2 100644 --- a/test/registered/unit/mem_cache/test_unified_mha_views.py +++ b/test/registered/unit/mem_cache/test_unified_mha_views.py @@ -30,11 +30,8 @@ import torch from sglang.srt.environ import envs -from sglang.srt.mem_cache.layout.page_major import ( - ENTRY_ALIGN_BYTES, - build_dense_views, - paged_view, -) +from sglang.srt.mem_cache.layout.page_major import ENTRY_ALIGN_BYTES, build_dense_views +from sglang.srt.mem_cache.layout.paged_view import paged_view from sglang.srt.mem_cache.unified_memory_pool import ( MHASubPoolSpec, UnifiedKVPool, diff --git a/test/registered/unit/mem_cache/test_unified_mla_views.py b/test/registered/unit/mem_cache/test_unified_mla_views.py index f87fd5c6025e..3e765b64239d 100644 --- a/test/registered/unit/mem_cache/test_unified_mla_views.py +++ b/test/registered/unit/mem_cache/test_unified_mla_views.py @@ -34,11 +34,8 @@ UnifiedMambaTokenToKVPoolAllocator, ) from sglang.srt.mem_cache.allocator.unified_sub_pool import MultiEndedAllocator -from sglang.srt.mem_cache.layout.page_major import ( - ENTRY_ALIGN_BYTES, - build_dense_views, - paged_row_view, -) +from sglang.srt.mem_cache.layout.page_major import ENTRY_ALIGN_BYTES, build_dense_views +from sglang.srt.mem_cache.layout.paged_view import paged_row_view from sglang.srt.mem_cache.unified_memory_pool import ( MambaSubPoolSpec, MLASubPoolSpec, From 695859593bdbf80360e1d321d90751bb5c5ec28e Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Mon, 28 Sep 2026 06:17:41 -0700 Subject: [PATCH 09/15] refactor(unified-memory): tidy the token-major layout code - MHASubPoolSpec: the entry diagram documents `layout()`, not a gap between two methods. - _assert_kernel_id_bound: ids no longer scale with the layer count, so stop suggesting it as a remedy. - build_dense_views: give the itemsize assert a message; drop the per-layer offset assert that validate()'s row alignment already implies. - page_major: say that the Mamba views use the same per-slot entry. - MultiEndedAllocator: drop the note about future plumbing removal. --- .../sglang/srt/mem_cache/allocator/unified_sub_pool.py | 1 - python/sglang/srt/mem_cache/layout/page_major.py | 10 +++++++--- python/sglang/srt/mem_cache/unified_memory_pool.py | 7 +++---- 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py b/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py index d37413b2836b..459513c999df 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py +++ b/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py @@ -318,7 +318,6 @@ def __init__( self.is_id_owner = is_id_owner # Kernel-facing ids are the physical token ids: the token-major views # step slots by the whole entry, so there is no per-page block scale. - # The kwarg is accepted (and must be 1) until the plumbing is removed. assert kernel_page_multiplier in (None, 1), ( f"kernel_page_multiplier must be 1 (token-major views); got " f"{kernel_page_multiplier}" diff --git a/python/sglang/srt/mem_cache/layout/page_major.py b/python/sglang/srt/mem_cache/layout/page_major.py index b7b0b629d2c4..ae91e64b1bb1 100644 --- a/python/sglang/srt/mem_cache/layout/page_major.py +++ b/python/sglang/srt/mem_cache/layout/page_major.py @@ -14,7 +14,8 @@ tensor with slot stride ``entry_bytes`` and storage offset ``anchor + part offset``, indexed by the PHYSICAL token id ``page * page_size + slot``. Parts may differ in row width (K vs V); only their -offsets differ, never the stride. +offsets differ, never the stride. The Mamba state +(``build_page_major_mamba_views``) uses the same per-slot entry at page size 1. These builders produce views into a raw ``uint8`` buffer; they hold no allocator/ownership state. ``anchor_bytes`` is the byte offset of the pool's @@ -130,7 +131,10 @@ def build_dense_views( ``page * page_size + slot`` (``paged_view`` regroups them by page). """ itemsize = part.dtype.itemsize - assert layout.entry_bytes % itemsize == 0 and anchor_bytes % itemsize == 0 + assert layout.entry_bytes % itemsize == 0 and anchor_bytes % itemsize == 0, ( + f"build_dense_views: entry {layout.entry_bytes} B and anchor " + f"{anchor_bytes} B must be multiples of the {part.dtype} itemsize" + ) n_rows = num_pages * page_size end = anchor_bytes + n_rows * layout.entry_bytes assert end <= raw.numel() * raw.itemsize, ( @@ -141,8 +145,8 @@ def build_dense_views( stride = (layout.entry_bytes // itemsize, *_contiguous_strides(part.row_shape)) views: List[torch.Tensor] = [] for layer in range(part.layer_num): + # Part offsets and layer strides are ROW_ALIGN_BYTES-aligned (validate()). base_bytes = anchor_bytes + part.layer_offset_bytes(layer) - assert base_bytes % itemsize == 0 views.append( torch.as_strided( as_dtype_view, diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index a1e2c8a3b279..3ea7d1907901 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -152,13 +152,12 @@ def entry_bytes(self) -> int: self.layer_num * (self.k_row_bytes() + self.v_row_bytes()) ) - # Token-major entry: [K_0 | V_0 | K_1 | V_1 | ...] per slot; a page is - # page_size such entries back to back. - def page_bytes(self, page_size: int) -> int: return page_size * self.entry_bytes() def layout(self) -> DenseEntryLayout: + """Token-major entry ``[K_0 | V_0 | K_1 | V_1 | ...]`` per slot; a page + is ``page_size`` such entries back to back.""" layer_stride = self.k_row_bytes() + self.v_row_bytes() return DenseEntryLayout( entry_bytes=self.entry_bytes(), @@ -284,7 +283,7 @@ def _assert_kernel_id_bound(*, sub_pool_name: str, n_rows: int) -> None: assert n_rows < 2**31, ( f"sub-pool {sub_pool_name!r}: kernel-facing id space has {n_rows} rows, " f"exceeding the int32 bound (2^31) that read-index buffers assume. " - "Reduce max_total_num_tokens or the layer count." + "Reduce max_total_num_tokens." ) From 11be13056c2e4cd5f63aaaa7c0745e28a30a7442 Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Mon, 28 Sep 2026 06:55:12 -0700 Subject: [PATCH 10/15] refactor(unified-memory): take a part name and a slot count in build_dense_views `build_dense_views` took the part itself next to the layout it must come from, so the two could disagree, and `page_size` / `num_pages` only to multiply them: a token-major view is `num_slots` entries long whatever the page size. Take `part_name` (looked up in `layout`) and `num_slots`. --- .../sglang/srt/mem_cache/layout/page_major.py | 19 +++++++++---------- .../srt/mem_cache/unified_memory_pool.py | 18 ++++++++---------- .../test_pd_envelope_transfer_layout.py | 15 ++++++--------- .../unit/mem_cache/test_unified_mha_views.py | 6 +++--- .../unit/mem_cache/test_unified_mla_views.py | 2 +- 5 files changed, 27 insertions(+), 33 deletions(-) diff --git a/python/sglang/srt/mem_cache/layout/page_major.py b/python/sglang/srt/mem_cache/layout/page_major.py index ae91e64b1bb1..787606482961 100644 --- a/python/sglang/srt/mem_cache/layout/page_major.py +++ b/python/sglang/srt/mem_cache/layout/page_major.py @@ -121,24 +121,23 @@ def build_dense_views( raw: torch.Tensor, *, layout: DenseEntryLayout, - part: DensePart, - page_size: int, - num_pages: int, + part_name: str, + num_slots: int, anchor_bytes: int = 0, ) -> List[torch.Tensor]: - """Per-layer views of one part: ``(num_pages * page_size, *row_shape)`` with - slot stride ``layout.entry_bytes``, indexed by the physical token id - ``page * page_size + slot`` (``paged_view`` regroups them by page). + """Per-layer views of one part: ``(num_slots, *row_shape)`` with slot + stride ``layout.entry_bytes``, indexed by the physical token id + (``paged_view`` regroups them by page). """ + part = layout.part(part_name) itemsize = part.dtype.itemsize assert layout.entry_bytes % itemsize == 0 and anchor_bytes % itemsize == 0, ( f"build_dense_views: entry {layout.entry_bytes} B and anchor " f"{anchor_bytes} B must be multiples of the {part.dtype} itemsize" ) - n_rows = num_pages * page_size - end = anchor_bytes + n_rows * layout.entry_bytes + end = anchor_bytes + num_slots * layout.entry_bytes assert end <= raw.numel() * raw.itemsize, ( - f"build_dense_views: {n_rows} slots of {layout.entry_bytes} B end at byte " + f"build_dense_views: {num_slots} slots of {layout.entry_bytes} B end at byte " f"{end} but the raw buffer holds only {raw.numel() * raw.itemsize} bytes" ) as_dtype_view = raw.view(part.dtype) @@ -150,7 +149,7 @@ def build_dense_views( views.append( torch.as_strided( as_dtype_view, - size=(n_rows, *part.row_shape), + size=(num_slots, *part.row_shape), stride=stride, storage_offset=base_bytes // itemsize, ) diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index 3ea7d1907901..88f1bb88799f 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -502,16 +502,15 @@ def _build_mha_views( max_slots: int, page_size: int, ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]: - num_pages = max_slots // page_size - _assert_kernel_id_bound(sub_pool_name=spec.name, n_rows=num_pages * page_size) + num_slots = max_slots // page_size * page_size + _assert_kernel_id_bound(sub_pool_name=spec.name, n_rows=num_slots) layout = spec.layout() k_views, v_views = ( build_dense_views( self._raw, layout=layout, - part=layout.part(name), - page_size=page_size, - num_pages=num_pages, + part_name=name, + num_slots=num_slots, anchor_bytes=anchor_bytes, ) for name in ("k", "v") @@ -525,15 +524,14 @@ def _build_mla_views( max_slots: int, page_size: int, ) -> List[torch.Tensor]: - num_pages = max_slots // page_size - _assert_kernel_id_bound(sub_pool_name=spec.name, n_rows=num_pages * page_size) + num_slots = max_slots // page_size * page_size + _assert_kernel_id_bound(sub_pool_name=spec.name, n_rows=num_slots) layout = spec.layout() return build_dense_views( self._raw, layout=layout, - part=layout.part("kv"), - page_size=page_size, - num_pages=num_pages, + part_name="kv", + num_slots=num_slots, anchor_bytes=anchor_bytes, ) diff --git a/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py b/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py index a6afaea907b4..b94f6e14d4fa 100644 --- a/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py +++ b/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py @@ -66,9 +66,8 @@ def test_page_envelope_matches_per_layer_views(self): views = build_dense_views( raw, layout=layout, - part=layout.part("kv"), - page_size=page_size, - num_pages=num_pages, + part_name="kv", + num_slots=num_pages * page_size, anchor_bytes=0, ) torch.manual_seed(0) @@ -129,9 +128,8 @@ def test_page_envelope_matches_per_layer_views(self): build_dense_views( raw, layout=layout, - part=layout.part(name), - page_size=page_size, - num_pages=num_pages, + part_name=name, + num_slots=num_pages * page_size, anchor_bytes=0, ) for name in ("k", "v") @@ -194,9 +192,8 @@ def test_envelope_move_is_a_whole_page_copy(self): build_dense_views( raw, layout=layout, - part=layout.part(name), - page_size=page_size, - num_pages=num_pages, + part_name=name, + num_slots=num_pages * page_size, anchor_bytes=0, ) for name in ("k", "v") diff --git a/test/registered/unit/mem_cache/test_unified_mha_views.py b/test/registered/unit/mem_cache/test_unified_mha_views.py index 77efaed0c9e2..7d2deb4c8986 100644 --- a/test/registered/unit/mem_cache/test_unified_mha_views.py +++ b/test/registered/unit/mem_cache/test_unified_mha_views.py @@ -84,10 +84,10 @@ def _make_raw(ps, num_pages, entry=_ENTRY, short=0): def _build_views(raw, ps, num_pages, head_dim=_D, v_head_dim=None, layer_num=_L): layout = _mha_spec(head_dim, v_head_dim, layer_num).layout() - kw = dict(layout=layout, page_size=ps, num_pages=num_pages) + kw = dict(layout=layout, num_slots=num_pages * ps) return ( - build_dense_views(raw, part=layout.part("k"), **kw), - build_dense_views(raw, part=layout.part("v"), **kw), + build_dense_views(raw, part_name="k", **kw), + build_dense_views(raw, part_name="v", **kw), ) diff --git a/test/registered/unit/mem_cache/test_unified_mla_views.py b/test/registered/unit/mem_cache/test_unified_mla_views.py index 3e765b64239d..9d93e61be718 100644 --- a/test/registered/unit/mem_cache/test_unified_mla_views.py +++ b/test/registered/unit/mem_cache/test_unified_mla_views.py @@ -100,7 +100,7 @@ def _make_unified(page_size=1, n_full_tokens=64, n_mamba_slots=8): def _build_views(raw, ps, num_pages): layout = _mla_spec().layout() return build_dense_views( - raw, layout=layout, part=layout.part("kv"), page_size=ps, num_pages=num_pages + raw, layout=layout, part_name="kv", num_slots=num_pages * ps ) From 80f1eb3c9e306712589ba4b72b04f2db6fdc2b34 Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Mon, 28 Sep 2026 12:22:04 -0700 Subject: [PATCH 11/15] refactor(unified-memory): name the entry-layout module token_major After the switch to per-slot entries, `mem_cache/layout/page_major.py` holds the token-major entry types (`DensePart`, `DenseEntryLayout`, `build_dense_views`) and the Mamba per-slot views; nothing in it lays pages out layer-major any more. Rename it to `token_major.py`, `build_page_major_mamba_views` to `build_mamba_entry_views` (next to `mamba_entry_bytes`), and the Mamba view test after the module. Pure rename: no body changes. The `--enable-page-major-kv-layout` flag keeps its name. --- .../mem_cache/layout/{page_major.py => token_major.py} | 4 ++-- python/sglang/srt/mem_cache/memory_pool.py | 8 ++++---- python/sglang/srt/mem_cache/unified_memory_pool.py | 6 +++--- .../unit/layers/test_mamba_state_scatter_triton.py | 6 +++--- .../mem_cache/test_flashkda_strided_state_access.py | 6 +++--- .../mem_cache/test_kda_fused_decode_strided_state.py | 6 +++--- .../unit/mem_cache/test_pd_envelope_transfer_layout.py | 6 +++--- ...page_major_layout.py => test_token_major_layout.py} | 10 +++++----- .../unit/mem_cache/test_unified_mamba_views.py | 8 +++----- .../unit/mem_cache/test_unified_mha_views.py | 2 +- .../unit/mem_cache/test_unified_mla_views.py | 2 +- 11 files changed, 31 insertions(+), 33 deletions(-) rename python/sglang/srt/mem_cache/layout/{page_major.py => token_major.py} (98%) rename test/registered/unit/mem_cache/{test_page_major_layout.py => test_token_major_layout.py} (84%) diff --git a/python/sglang/srt/mem_cache/layout/page_major.py b/python/sglang/srt/mem_cache/layout/token_major.py similarity index 98% rename from python/sglang/srt/mem_cache/layout/page_major.py rename to python/sglang/srt/mem_cache/layout/token_major.py index 787606482961..e5e9431f19b4 100644 --- a/python/sglang/srt/mem_cache/layout/page_major.py +++ b/python/sglang/srt/mem_cache/layout/token_major.py @@ -15,7 +15,7 @@ ``anchor + part offset``, indexed by the PHYSICAL token id ``page * page_size + slot``. Parts may differ in row width (K vs V); only their offsets differ, never the stride. The Mamba state -(``build_page_major_mamba_views``) uses the same per-slot entry at page size 1. +(``build_mamba_entry_views``) uses the same per-slot entry at page size 1. These builders produce views into a raw ``uint8`` buffer; they hold no allocator/ownership state. ``anchor_bytes`` is the byte offset of the pool's @@ -173,7 +173,7 @@ def mamba_entry_bytes( return total -def build_page_major_mamba_views( +def build_mamba_entry_views( raw: torch.Tensor, *, layer_num: int, diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 93b5eeabf9b3..f0f5a37e8d6f 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -60,8 +60,8 @@ from sglang.srt.mem_cache.allocator.mamba import MambaSlotAllocator from sglang.srt.mem_cache.index_key_cache import IndexKeyCache from sglang.srt.mem_cache.kv_vmm_backing import KvVmmBufferOwner -from sglang.srt.mem_cache.layout.page_major import ( - build_page_major_mamba_views, +from sglang.srt.mem_cache.layout.token_major import ( + build_mamba_entry_views, mamba_entry_bytes, ) from sglang.srt.mem_cache.utils import ( @@ -587,7 +587,7 @@ def __init__( # Page-granularity envelope layout (page_size==1 for state): all # mamba layers/slots share one contiguous byte buffer; conv and # temporal are strided views into it (see mem_cache/layout/ - # page_major.py). Only the standard CUDA Triton path is supported. + # token_major.py). Only the standard CUDA Triton path is supported. assert not _is_npu and not (_is_cpu and _cpu_has_amx_support), ( "envelope_layout mamba is only supported on the CUDA path" ) @@ -602,7 +602,7 @@ def __init__( self._raw = torch.zeros( max_slots * entry_bytes, dtype=torch.uint8, device=device ) - conv_state, temporal_state = build_page_major_mamba_views( + conv_state, temporal_state = build_mamba_entry_views( self._raw, layer_num=num_mamba_layers, conv_state_shapes=conv_state_shape, diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index 88f1bb88799f..7f0d2f8d64ff 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -36,12 +36,12 @@ from sglang.kernels.ops.kvcache.zero_pages import zero_pages from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.environ import envs -from sglang.srt.mem_cache.layout.page_major import ( +from sglang.srt.mem_cache.layout.token_major import ( DenseEntryLayout, DensePart, align_entry_bytes, build_dense_views, - build_page_major_mamba_views, + build_mamba_entry_views, ) from sglang.srt.mem_cache.memory_pool import ( HybridLinearKVPool, @@ -538,7 +538,7 @@ def _build_mla_views( def _build_mamba_views( self, spec: MambaSubPoolSpec, anchor_bytes: int, max_slots: int ) -> Tuple[List[torch.Tensor], torch.Tensor]: - return build_page_major_mamba_views( + return build_mamba_entry_views( self._raw, layer_num=spec.layer_num, conv_state_shapes=spec.conv_state_shapes, diff --git a/test/registered/unit/layers/test_mamba_state_scatter_triton.py b/test/registered/unit/layers/test_mamba_state_scatter_triton.py index 145f20fe7c23..a76d39cee3d6 100644 --- a/test/registered/unit/layers/test_mamba_state_scatter_triton.py +++ b/test/registered/unit/layers/test_mamba_state_scatter_triton.py @@ -31,8 +31,8 @@ fused_mamba_state_scatter_with_mask = None _FUSED_IMPORT_ERROR = e -from sglang.srt.mem_cache.layout.page_major import ( - build_page_major_mamba_views, +from sglang.srt.mem_cache.layout.token_major import ( + build_mamba_entry_views, mamba_entry_bytes, ) from sglang.srt.utils import get_device @@ -269,7 +269,7 @@ def _make_envelope_views(device="cpu"): temporal_dtype=temporal_dtype, ) raw = torch.zeros(slots * entry, dtype=torch.uint8, device=device) - conv_views, temporal = build_page_major_mamba_views( + conv_views, temporal = build_mamba_entry_views( raw, layer_num=layers, conv_state_shapes=conv_shapes, diff --git a/test/registered/unit/mem_cache/test_flashkda_strided_state_access.py b/test/registered/unit/mem_cache/test_flashkda_strided_state_access.py index 8d576b4f53c2..d821e99610fc 100644 --- a/test/registered/unit/mem_cache/test_flashkda_strided_state_access.py +++ b/test/registered/unit/mem_cache/test_flashkda_strided_state_access.py @@ -34,8 +34,8 @@ import torch from sglang.srt.layers.attention.linear.kernels.kda_flashkda import FlashKDAKernel -from sglang.srt.mem_cache.layout.page_major import ( - build_page_major_mamba_views, +from sglang.srt.mem_cache.layout.token_major import ( + build_mamba_entry_views, mamba_entry_bytes, ) @@ -68,7 +68,7 @@ def _make_strided_temporal_views(): temporal_dtype=_TEMPORAL_DTYPE, ) raw = torch.zeros(_SLOTS * entry, dtype=torch.uint8, device=_DEV) - conv_views, temporal = build_page_major_mamba_views( + conv_views, temporal = build_mamba_entry_views( raw, layer_num=_LAYERS, conv_state_shapes=_CONV_SHAPES, diff --git a/test/registered/unit/mem_cache/test_kda_fused_decode_strided_state.py b/test/registered/unit/mem_cache/test_kda_fused_decode_strided_state.py index 86a48f46976a..c2a866443f11 100644 --- a/test/registered/unit/mem_cache/test_kda_fused_decode_strided_state.py +++ b/test/registered/unit/mem_cache/test_kda_fused_decode_strided_state.py @@ -43,8 +43,8 @@ _CONV_STATE_W, covered, ) -from sglang.srt.mem_cache.layout.page_major import ( - build_page_major_mamba_views, +from sglang.srt.mem_cache.layout.token_major import ( + build_mamba_entry_views, mamba_entry_bytes, ) @@ -85,7 +85,7 @@ def _make_strided_temporal_view(heads: int): temporal_dtype=_TEMPORAL_DTYPE, ) raw = torch.zeros(_SLOTS * entry, dtype=torch.uint8, device=_DEV) - _conv_views, temporal = build_page_major_mamba_views( + _conv_views, temporal = build_mamba_entry_views( raw, layer_num=_LAYERS, conv_state_shapes=_CONV_SHAPES, diff --git a/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py b/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py index b94f6e14d4fa..4341594bf69b 100644 --- a/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py +++ b/test/registered/unit/mem_cache/test_pd_envelope_transfer_layout.py @@ -17,12 +17,12 @@ import torch -from sglang.srt.mem_cache.layout.page_major import ( +from sglang.srt.mem_cache.layout.token_major import ( ENTRY_ALIGN_BYTES, DenseEntryLayout, DensePart, build_dense_views, - build_page_major_mamba_views, + build_mamba_entry_views, mamba_entry_bytes, ) from sglang.test.ci.ci_register import register_cpu_ci @@ -239,7 +239,7 @@ def test_slot_envelope_is_self_contained(self): temporal_dtype=temporal_dtype, ) raw = torch.zeros(max_slots * entry, dtype=torch.uint8) - conv_views, temporal_view = build_page_major_mamba_views( + conv_views, temporal_view = build_mamba_entry_views( raw, layer_num=layer_num, conv_state_shapes=conv_shapes, diff --git a/test/registered/unit/mem_cache/test_page_major_layout.py b/test/registered/unit/mem_cache/test_token_major_layout.py similarity index 84% rename from test/registered/unit/mem_cache/test_page_major_layout.py rename to test/registered/unit/mem_cache/test_token_major_layout.py index 090c3592b5e4..21bfd9281aac 100644 --- a/test/registered/unit/mem_cache/test_page_major_layout.py +++ b/test/registered/unit/mem_cache/test_token_major_layout.py @@ -1,13 +1,13 @@ """CPU correctness tests for the page-major envelope Mamba state views. -Covers the standalone ``build_page_major_mamba_views`` builder (no allocator / +Covers the standalone ``build_mamba_entry_views`` builder (no allocator / shared pool): conv / temporal state views with correct shapes and no aliasing across layers / slots. The unified pool stores its Mamba/KDA state through these views. Runs on CPU — pure-torch advanced indexing, no Triton. - python -m pytest test/registered/unit/mem_cache/test_page_major_layout.py -v + python -m pytest test/registered/unit/mem_cache/test_token_major_layout.py -v """ from sglang.test.ci.ci_register import register_cpu_ci @@ -18,8 +18,8 @@ import torch -from sglang.srt.mem_cache.layout.page_major import ( - build_page_major_mamba_views, +from sglang.srt.mem_cache.layout.token_major import ( + build_mamba_entry_views, mamba_entry_bytes, ) @@ -40,7 +40,7 @@ def test_conv_temporal_shapes_no_alias(self): temporal_dtype=temp_dt, ) raw = torch.zeros(slots * entry, dtype=torch.uint8, device=_DEV) - conv_views, temporal = build_page_major_mamba_views( + conv_views, temporal = build_mamba_entry_views( raw, layer_num=layers, conv_state_shapes=conv_shapes, diff --git a/test/registered/unit/mem_cache/test_unified_mamba_views.py b/test/registered/unit/mem_cache/test_unified_mamba_views.py index b9f1eef0a31f..6268d693ed80 100644 --- a/test/registered/unit/mem_cache/test_unified_mamba_views.py +++ b/test/registered/unit/mem_cache/test_unified_mamba_views.py @@ -249,8 +249,8 @@ class TestKDAFlashInferEnvelopeStateContract(unittest.TestCase): def _build_kda_views(heads_per_rank: int = 12): """Real K3 KDA envelope views on CPU for one attn-TP shard; 2 slots suffice, the per-slot geometry is slot-count independent.""" - from sglang.srt.mem_cache.layout.page_major import ( - build_page_major_mamba_views, + from sglang.srt.mem_cache.layout.token_major import ( + build_mamba_entry_views, mamba_entry_bytes, ) @@ -258,9 +258,7 @@ def _build_kda_views(heads_per_rank: int = 12): entry_bytes = mamba_entry_bytes(**geom) max_slots = 2 raw = torch.empty(max_slots * entry_bytes, dtype=torch.uint8, device="cpu") - _, temporal_view = build_page_major_mamba_views( - raw, max_slots=max_slots, **geom - ) + _, temporal_view = build_mamba_entry_views(raw, max_slots=max_slots, **geom) return geom, entry_bytes, temporal_view def test_k3_tp8_envelope_view_matches_recurrent_kda_contract(self): diff --git a/test/registered/unit/mem_cache/test_unified_mha_views.py b/test/registered/unit/mem_cache/test_unified_mha_views.py index 7d2deb4c8986..3ee7ac0bd1c5 100644 --- a/test/registered/unit/mem_cache/test_unified_mha_views.py +++ b/test/registered/unit/mem_cache/test_unified_mha_views.py @@ -30,8 +30,8 @@ import torch from sglang.srt.environ import envs -from sglang.srt.mem_cache.layout.page_major import ENTRY_ALIGN_BYTES, build_dense_views from sglang.srt.mem_cache.layout.paged_view import paged_view +from sglang.srt.mem_cache.layout.token_major import ENTRY_ALIGN_BYTES, build_dense_views from sglang.srt.mem_cache.unified_memory_pool import ( MHASubPoolSpec, UnifiedKVPool, diff --git a/test/registered/unit/mem_cache/test_unified_mla_views.py b/test/registered/unit/mem_cache/test_unified_mla_views.py index 9d93e61be718..fe031fabf24d 100644 --- a/test/registered/unit/mem_cache/test_unified_mla_views.py +++ b/test/registered/unit/mem_cache/test_unified_mla_views.py @@ -34,8 +34,8 @@ UnifiedMambaTokenToKVPoolAllocator, ) from sglang.srt.mem_cache.allocator.unified_sub_pool import MultiEndedAllocator -from sglang.srt.mem_cache.layout.page_major import ENTRY_ALIGN_BYTES, build_dense_views from sglang.srt.mem_cache.layout.paged_view import paged_row_view +from sglang.srt.mem_cache.layout.token_major import ENTRY_ALIGN_BYTES, build_dense_views from sglang.srt.mem_cache.unified_memory_pool import ( MambaSubPoolSpec, MLASubPoolSpec, From e8ef731c762184f31468605830f5eec937661da5 Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Mon, 28 Sep 2026 12:27:09 -0700 Subject: [PATCH 12/15] fix(unified-memory): tag HiCache storage keys with the unified page format The unified host pool persists whole page envelopes. Their size is the same under the old layer-major layout and the token-major one, and storage keys record only the model and ranks, so a page written by an older build (or by a static pool) would be read back as a token-major page without any error. Give host pools a `storage_format_tag` (None by default) and set it on `UnifiedPageEnvelopeHostPool`. The controller appends it to the model name every storage backend namespaces its keys with, so pages of another format miss instead of loading. The page hash is untouched: it has to stay equal to the router's. --- .../sglang/srt/managers/cache_controller.py | 12 +++ python/sglang/srt/mem_cache/pool_host/base.py | 3 + .../sglang/srt/mem_cache/pool_host/unified.py | 2 + .../test_hicache_storage_format_tag.py | 76 +++++++++++++++++++ 4 files changed, 93 insertions(+) create mode 100644 test/registered/unit/mem_cache/test_hicache_storage_format_tag.py diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 115720b06fe0..8ea82e881185 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -50,6 +50,17 @@ device_module = get_device_module() +def storage_model_name( + model_name: Optional[str], host_pool: HostKVCache +) -> Optional[str]: + """The model name storage backends key pages on, tagged with the host + pool's page format when it has one of its own.""" + tag = host_pool.storage_format_tag + if tag is None: + return model_name + return f"{model_name}-{tag}" if model_name else tag + + class LayerLoadingEvent: def __init__(self, num_layers: int): self._num_layers = num_layers @@ -730,6 +741,7 @@ def _generate_storage_config( ) attn_cp_rank, attn_cp_size = self.get_attn_cp_rank_and_size() + model_name = storage_model_name(model_name, self.storage_host_pool) return HiCacheStorageConfig( tp_rank=self.tp_rank, diff --git a/python/sglang/srt/mem_cache/pool_host/base.py b/python/sglang/srt/mem_cache/pool_host/base.py index 826c754bba30..bb8b90665904 100644 --- a/python/sglang/srt/mem_cache/pool_host/base.py +++ b/python/sglang/srt/mem_cache/pool_host/base.py @@ -157,6 +157,9 @@ class HostKVCache(abc.ABC): dcp_size = 1 dcp_rank = 0 shared_allocation_domain = None + # Names this pool's page byte format in storage keys when it has one of its + # own, so pages persisted in another format miss instead of loading. + storage_format_tag: Optional[str] = None def __init__( self, diff --git a/python/sglang/srt/mem_cache/pool_host/unified.py b/python/sglang/srt/mem_cache/pool_host/unified.py index b1841645705c..48ca720a20ff 100644 --- a/python/sglang/srt/mem_cache/pool_host/unified.py +++ b/python/sglang/srt/mem_cache/pool_host/unified.py @@ -610,6 +610,8 @@ def translate_index(self, name: str, index: int) -> int: class UnifiedPageEnvelopeHostPool(HostKVCache): """Host mirror that transfers complete unified-memory page envelopes.""" + storage_format_tag = "unified-token-major" + def __init__( self, device_pool, diff --git a/test/registered/unit/mem_cache/test_hicache_storage_format_tag.py b/test/registered/unit/mem_cache/test_hicache_storage_format_tag.py new file mode 100644 index 000000000000..bb9d2ad6384c --- /dev/null +++ b/test/registered/unit/mem_cache/test_hicache_storage_format_tag.py @@ -0,0 +1,76 @@ +"""A host pool with its own page byte format namespaces its storage keys. + +The unified pool's host mirror stores whole page envelopes, whose bytes are +token-major entries. A page persisted under another format has the same size, +so only the key can keep it from being read back as a unified page: the +storage model name carries the pool's format tag, and an untagged page misses. + +CPU-only. + + python -m pytest test/registered/unit/mem_cache/test_hicache_storage_format_tag.py -v +""" + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=6, suite="base-a-test-cpu") + +import shutil +import tempfile +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.managers.cache_controller import storage_model_name +from sglang.srt.mem_cache.hicache_storage import HiCacheFile, HiCacheStorageConfig +from sglang.srt.mem_cache.pool_host.base import HostKVCache +from sglang.srt.mem_cache.pool_host.unified import UnifiedPageEnvelopeHostPool +from sglang.test.test_utils import CustomTestCase + + +def _config(model_name): + return HiCacheStorageConfig( + tp_rank=0, + tp_size=1, + pp_rank=0, + pp_size=1, + attn_cp_rank=0, + attn_cp_size=1, + is_mla_model=False, + enable_storage_metrics=False, + is_page_first_layout=True, + model_name=model_name, + ) + + +class TestStorageModelName(CustomTestCase): + def test_default_pools_keep_the_model_name(self): + self.assertIsNone(HostKVCache.storage_format_tag) + pool = SimpleNamespace(storage_format_tag=None) + self.assertEqual(storage_model_name("org/model", pool), "org/model") + self.assertIsNone(storage_model_name(None, pool)) + + def test_unified_pool_tags_the_model_name(self): + tag = UnifiedPageEnvelopeHostPool.storage_format_tag + self.assertIsNotNone(tag) + self.assertEqual( + storage_model_name("org/model", UnifiedPageEnvelopeHostPool), + f"org/model-{tag}", + ) + self.assertEqual(storage_model_name(None, UnifiedPageEnvelopeHostPool), tag) + + def test_untagged_page_misses_under_the_unified_name(self): + root = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, root, ignore_errors=True) + untagged = HiCacheFile(_config("org/model"), file_path=root) + tagged = HiCacheFile( + _config(storage_model_name("org/model", UnifiedPageEnvelopeHostPool)), + file_path=root, + ) + self.assertTrue(untagged.set("page-0", torch.zeros(64, dtype=torch.uint8))) + self.assertTrue(untagged.exists("page-0")) + self.assertFalse(tagged.exists("page-0")) + + +if __name__ == "__main__": + unittest.main() From 2c4e2d29351028f7f6e91f71dede489b439e7de1 Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Mon, 28 Sep 2026 12:27:46 -0700 Subject: [PATCH 13/15] docs(unified-memory): state the backend requirement the startup checks enforce The `--enable-page-major-kv-layout` and `--enable-unified-memory` help and the server-arguments page still said the Triton attention, linear-attention and Mamba backends are required, and the page also said unified memory supports neither PD disaggregation nor speculative decoding. Say what holds: the backends must be stride-aware, and the allowed set and the supported PD / speculative configurations are checked at startup. The page-major flag is only accepted together with unified memory. --- docs/docs/advanced_features/server_arguments.mdx | 4 ++-- python/sglang/srt/arg_groups/fields/memory.py | 15 +++++++-------- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/docs/docs/advanced_features/server_arguments.mdx b/docs/docs/advanced_features/server_arguments.mdx index b0892a68b5bc..d8453e95a92e 100644 --- a/docs/docs/advanced_features/server_arguments.mdx +++ b/docs/docs/advanced_features/server_arguments.mdx @@ -534,13 +534,13 @@ Please consult the documentation below and [server_args.py](https://github.com/s `--enable-page-major-kv-layout` - Enable the page-major KV layout: lay out the Mamba state and full/SWA KV caches in a page-granularity envelope (page is the outermost axis; within a page, each token's entry holds every layer's K/V) instead of the default per-layer (layer-major) layout. Requires the Triton attention / linear-attn / Mamba backends (`--attention-backend triton`, and for hybrid models `--linear-attn-backend triton --mamba-backend triton`). + Store the Mamba state and full/SWA KV caches as one entry per token slot, holding every layer's state, in a single page-granular buffer instead of one buffer per layer. Implied by `--enable-unified-memory`, and currently supported only with it. `False` bool flag (set to enable) `--enable-unified-memory` - For hybrid Mamba/GDN and hybrid SWA models, replace the statically-partitioned pools (full-attention KV + SWA/Mamba conv/SSM state) with a single byte buffer split dynamically between the sub-pools, so KV-vs-state capacity flexes with the workload instead of being fixed at startup. Implies `--enable-page-major-kv-layout`. Requires the Triton attention / linear-attn / Mamba backends; monolithic (decode) cuda-graph capture only; not yet compatible with PD disaggregation or speculative decoding. + For hybrid Mamba/GDN and hybrid SWA models, replace the statically-partitioned pools (full-attention KV + SWA/Mamba conv/SSM state) with a single byte buffer split dynamically between the sub-pools, so KV-vs-state capacity flexes with the workload instead of being fixed at startup. Implies `--enable-page-major-kv-layout`. Requires stride-aware attention, linear-attention and Mamba backends; the allowed backends and the supported PD-disaggregation and speculative-decoding configurations are validated at startup. `False` bool flag (set to enable) diff --git a/python/sglang/srt/arg_groups/fields/memory.py b/python/sglang/srt/arg_groups/fields/memory.py index a76165cd5afe..79d230d6bec4 100644 --- a/python/sglang/srt/arg_groups/fields/memory.py +++ b/python/sglang/srt/arg_groups/fields/memory.py @@ -70,19 +70,18 @@ class Memory(msgspec.Struct): ] = False enable_page_major_kv_layout: A[ bool, - "Enable the page-major KV layout: lay out the Mamba state and full/SWA " - "KV caches in a page-granularity envelope (page is the outermost axis; " - "within a page, each token's entry holds every layer's K/V) instead of " - "the default per-layer (layer-major) layout. Requires the Triton " - "attention / linear-attn / Mamba backends.", + "Store the Mamba state and full/SWA KV caches as one entry per token " + "slot, holding every layer's state, in a single page-granular buffer " + "instead of one buffer per layer. Implied by --enable-unified-memory, " + "and currently supported only with it.", ] = False enable_unified_memory: A[ bool, "Replace the statically-partitioned hybrid-model pools (full-attn KV + " "SWA/Mamba state) with one byte buffer split dynamically between " - "sub-pools. Requires the Triton attention / linear-attn / Mamba " - "backends. Supported PD-disaggregation and speculative-decoding " - "configurations are validated at startup.", + "sub-pools. Requires stride-aware attention, linear-attention and Mamba " + "backends; the allowed backends and the supported PD-disaggregation and " + "speculative-decoding configurations are validated at startup.", ] = False enable_session_radix_cache: A[ bool, From d531e36c7fe9e0cc808e2c397de2a486c078bf58 Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Mon, 28 Sep 2026 12:29:11 -0700 Subject: [PATCH 14/15] fix(unified-memory): refuse misaligned per-rank rows with a clear error The token-major entry lays every part out in 16-byte units, so a model whose per-rank K, V or latent row is not a multiple of 16 bytes cannot use the unified pool. That surfaced as a bare assertion inside `DenseEntryLayout.validate()`. Check it where the sub-pool builds its layout and raise a `ValueError` naming the sub-pool, the row and the way out; the assertion stays as the layout's own invariant. --- .../sglang/srt/mem_cache/unified_memory_pool.py | 16 ++++++++++++++++ .../unit/mem_cache/test_unified_mha_views.py | 8 ++++---- .../unit/mem_cache/test_unified_mla_views.py | 12 ++++++++++++ 3 files changed, 32 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index 7f0d2f8d64ff..7bcc63ecf994 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -37,6 +37,7 @@ from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.environ import envs from sglang.srt.mem_cache.layout.token_major import ( + ROW_ALIGN_BYTES, DenseEntryLayout, DensePart, align_entry_bytes, @@ -66,6 +67,19 @@ def _prod(iterable) -> int: return out +def _check_row_alignment(sub_pool: str, **row_bytes: int) -> None: + """Entry parts are laid out in ROW_ALIGN_BYTES units; a model whose + per-rank rows are not cannot use the unified pool.""" + for part, nbytes in row_bytes.items(): + if nbytes % ROW_ALIGN_BYTES: + raise ValueError( + f"--enable-unified-memory needs every per-rank KV row to be a " + f"multiple of {ROW_ALIGN_BYTES} bytes; sub-pool {sub_pool!r} has a " + f"{nbytes}-byte {part} row. Run this model without " + "--enable-unified-memory." + ) + + def _store_dtype_for(kv_cache_dtype: torch.dtype) -> torch.dtype: if kv_cache_dtype in ( torch.float8_e5m2, @@ -158,6 +172,7 @@ def page_bytes(self, page_size: int) -> int: def layout(self) -> DenseEntryLayout: """Token-major entry ``[K_0 | V_0 | K_1 | V_1 | ...]`` per slot; a page is ``page_size`` such entries back to back.""" + _check_row_alignment(self.name, K=self.k_row_bytes(), V=self.v_row_bytes()) layer_stride = self.k_row_bytes() + self.v_row_bytes() return DenseEntryLayout( entry_bytes=self.entry_bytes(), @@ -219,6 +234,7 @@ def entry_bytes(self) -> int: return align_entry_bytes(self.layer_num * self.row_bytes()) def layout(self) -> DenseEntryLayout: + _check_row_alignment(self.name, latent=self.row_bytes()) return DenseEntryLayout( entry_bytes=self.entry_bytes(), parts=( diff --git a/test/registered/unit/mem_cache/test_unified_mha_views.py b/test/registered/unit/mem_cache/test_unified_mha_views.py index 3ee7ac0bd1c5..93ad8a276665 100644 --- a/test/registered/unit/mem_cache/test_unified_mha_views.py +++ b/test/registered/unit/mem_cache/test_unified_mha_views.py @@ -160,10 +160,10 @@ def test_asymmetric_rows_are_admitted(self): self.assertEqual(layout.part("v").row_bytes(), _H * 4 * _ITEM) self.assertEqual(layout.part("v").offset_bytes, _H * 8 * _ITEM) - def test_misaligned_rows_fail_loud(self): - """A row that is not a multiple of 16 B would break the vector stores - every write kernel issues, so the layout refuses it at construction.""" - with self.assertRaises(AssertionError): + def test_misaligned_rows_are_refused(self): + """Entry parts are laid out in 16-byte units (the write kernels' + vector stores), so a model whose per-rank row is not cannot run.""" + with self.assertRaisesRegex(ValueError, "24-byte K row"): _mha_spec(head_dim=6, v_head_dim=6).layout() diff --git a/test/registered/unit/mem_cache/test_unified_mla_views.py b/test/registered/unit/mem_cache/test_unified_mla_views.py index fe031fabf24d..80fe619c7d95 100644 --- a/test/registered/unit/mem_cache/test_unified_mla_views.py +++ b/test/registered/unit/mem_cache/test_unified_mla_views.py @@ -116,6 +116,18 @@ def test_rejects_nonpositive_dims(self): grow_direction="down", ) + def test_misaligned_latent_row_is_refused(self): + spec = MLASubPoolSpec( + name="full", + layer_num=_L, + kv_lora_rank=5, + qk_rope_head_dim=_ROPE, + store_dtype=_DTYPE, + grow_direction="down", + ) + with self.assertRaisesRegex(ValueError, "14-byte latent row"): + spec.layout() + class TestMLAViews(unittest.TestCase): def _make_raw(self, ps, num_pages, short=0): From 94238102c2ad8016d40d3694b0fb4df99172a0e2 Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Mon, 28 Sep 2026 12:30:38 -0700 Subject: [PATCH 15/15] style(unified-memory): keep the addressing rule in docstrings, drop the retelling The token-major module docstring now states the entry and the byte formula once; the unified pool classes point at it instead of restating the layout, and the trtllm_mha gate comment says what is refused and why in three lines. --- python/sglang/srt/arg_groups/kv_cache_hook.py | 9 +++---- .../srt/mem_cache/layout/token_major.py | 25 ++++++------------- .../srt/mem_cache/unified_memory_pool.py | 21 ++++++---------- 3 files changed, 18 insertions(+), 37 deletions(-) diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index b7ebc1271681..a82fbd186648 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -608,12 +608,9 @@ def handle_unified_memory_pool(server_args: Any) -> None: sorted(full_cg_backends), sorted(backends), ) - # trtllm_mha refills its cuda-graph page table out of graph, before the - # replay, from `cache_seqlens_int32` -- which only the recorded in-graph - # metadata kernel writes. The refill therefore translates the previous - # replay's lengths, and the kernel reads untranslated columns once a - # sequence grows past them. Refuse the combination until that refill - # derives its lengths from the batch. + # trtllm_mha refills its graph page table before replay from + # `cache_seqlens_int32`, which only the in-graph metadata kernel writes, so + # it uses the previous replay's lengths. Refuse until it uses the batch's. _, decode_backend = attention_backends_of(resolved_view(server_args)) if decode_backend == "trtllm_mha": assert _cg_cfg is None or _cg_cfg.decode.backend == Backend.DISABLED, ( diff --git a/python/sglang/srt/mem_cache/layout/token_major.py b/python/sglang/srt/mem_cache/layout/token_major.py index e5e9431f19b4..98322d7f8431 100644 --- a/python/sglang/srt/mem_cache/layout/token_major.py +++ b/python/sglang/srt/mem_cache/layout/token_major.py @@ -1,25 +1,16 @@ -"""Page-granularity envelope (page-major, token-major within a page) cache views. +"""Token-major cache views over one raw byte buffer. -A pool of this layout keeps all layers of all slots in one contiguous byte -buffer, split into pages of ``page_size`` slots. Within a page the slots follow -one another, and one slot's ENTRY holds every part of that token -- K and V of -every layer for MHA, the latent row of every layer for MLA -- at a fixed byte -offset: +Each token slot owns one ENTRY holding every layer's state at fixed offsets, +and a page is ``page_size`` entries back to back: - page bytes = [ entry(slot 0) | entry(slot 1) | ... | entry(slot ps-1) ] entry bytes = [ K_0 | V_0 | K_1 | V_1 | ... ] (MHA) [ lat_0 | lat_1 | ... ] (MLA) + byte(t, part, l) = anchor + t * entry_bytes + part offset + l * layer stride -Every per-layer view is therefore a flat ``(num_pages * page_size, *row_shape)`` -tensor with slot stride ``entry_bytes`` and storage offset -``anchor + part offset``, indexed by the PHYSICAL token id -``page * page_size + slot``. Parts may differ in row width (K vs V); only their -offsets differ, never the stride. The Mamba state -(``build_mamba_entry_views``) uses the same per-slot entry at page size 1. - -These builders produce views into a raw ``uint8`` buffer; they hold no -allocator/ownership state. ``anchor_bytes`` is the byte offset of the pool's -region inside the raw buffer (0 for a standalone pool). +so every per-layer view is ``(num_slots, *row_shape)`` with slot stride +``entry_bytes``, indexed by the physical token id. The Mamba state +(``build_mamba_entry_views``) uses the same per-slot entry. ``anchor_bytes`` +is the pool region's offset inside the raw buffer. """ from typing import List, Sequence, Tuple diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index 7bcc63ecf994..2dcaedabda5a 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -568,12 +568,9 @@ def _build_mamba_views( class UnifiedMHATokenToKVPool(MHATokenToKVPool): """MHA KV pool whose per-layer `k_buffer`/`v_buffer` are token-major views - into a `UnifiedKVPool` (see `build_dense_views`). - - Each view is a strided `(num_pages * ps, head_num, head_dim)` tensor whose - slot stride is the whole entry, indexed by the PHYSICAL token id; K and V - of every layer sit at fixed offsets inside that entry, so K and V rows may - differ in width. `move_kv_cache` relocates whole page envelopes. + into a `UnifiedKVPool`: `(num_pages * ps, head_num, head_dim)` with the whole + entry as slot stride, indexed by the physical token id. `move_kv_cache` + relocates whole page envelopes. """ def __init__( @@ -680,14 +677,10 @@ def set_kv_buffer_prefix_valid(self, *args, **kwargs): class UnifiedMLATokenToKVPool(MLATokenToKVPool): """MLA KV pool whose per-layer `kv_buffer` entries are token-major views - into a `UnifiedKVPool` (see `build_dense_views`). - - Every loc this pool receives through the KVCache API is a PHYSICAL token - id (`translate_kv_loc` of the virtual id). Each view is a strided - `(num_pages * ps, 1, kv_cache_dim)` tensor whose slot stride is the whole - entry, the layer offset folded into its storage_offset, so the stock - `MLATokenToKVPool` read/write methods work on the views unmodified. - `move_kv_cache` relocates whole page envelopes on the raw buffer. + into a `UnifiedKVPool`: `(num_pages * ps, 1, kv_cache_dim)` with the whole + entry as slot stride, indexed by the physical token id (`translate_kv_loc` + of the virtual id), so the stock `MLATokenToKVPool` read/write methods work + on them unmodified. `move_kv_cache` relocates whole page envelopes. """ def __init__(