diff --git a/docs/docs/advanced_features/server_arguments.mdx b/docs/docs/advanced_features/server_arguments.mdx
index 8f4a52b5e008..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, 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`). |
+ 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 992caf722244..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, "
- "layer-major within a page) 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,
diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py
index 196d8fdfd94d..a82fbd186648 100644
--- a/python/sglang/srt/arg_groups/kv_cache_hook.py
+++ b/python/sglang/srt/arg_groups/kv_cache_hook.py
@@ -608,6 +608,18 @@ def handle_unified_memory_pool(server_args: Any) -> None:
sorted(full_cg_backends),
sorted(backends),
)
+ # 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, (
+ "--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:
@@ -656,9 +668,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-strided 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:
@@ -684,24 +695,24 @@ 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
# 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",
@@ -711,7 +722,7 @@ def handle_page_major_kv_layout(server_args: Any):
"tokenspeed_mla",
"flashmla",
}
- elif cfg.enable_unified_memory:
+ else:
allowed_full = {
"triton",
"fa3",
@@ -719,17 +730,14 @@ 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, (
"--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/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/allocator/unified_sub_pool.py b/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py
index 2594572d72d6..459513c999df 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,13 @@ 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.
+ 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 +1060,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
deleted file mode 100644
index ff81583181ab..000000000000
--- a/python/sglang/srt/mem_cache/layout/page_major.py
+++ /dev/null
@@ -1,293 +0,0 @@
-"""Page-granularity envelope (page-major, layer-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).
-"""
-
-from typing import List, Sequence, Tuple
-
-import torch
-
-
-def _prod(shape: Sequence[int]) -> int:
- out = 1
- for s in shape:
- out *= int(s)
- return out
-
-
-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)."""
- 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)
-
-
-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 (single latent row, no V)."""
- return 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,
- conv_state_shapes: Sequence[Sequence[int]],
- conv_dtype: torch.dtype,
- temporal_state_shape: Sequence[int],
- temporal_dtype: torch.dtype,
-) -> int:
- """Bytes occupied by one Mamba slot across all layers (conv + temporal)."""
- total = 0
- for shape in conv_state_shapes:
- total += layer_num * _prod(shape) * conv_dtype.itemsize
- total += layer_num * _prod(temporal_state_shape) * temporal_dtype.itemsize
- return total
-
-
-def build_page_major_mamba_views(
- raw: torch.Tensor,
- *,
- layer_num: int,
- conv_state_shapes: Sequence[Sequence[int]],
- conv_dtype: torch.dtype,
- temporal_state_shape: Sequence[int],
- temporal_dtype: torch.dtype,
- max_slots: int,
- anchor_bytes: int = 0,
-) -> Tuple[List[torch.Tensor], torch.Tensor]:
- """Per-slot envelope views over ``raw`` for Mamba state.
-
- Layout per slot: ``[conv[0] rows × layers][conv[1] rows × layers]...
- [temporal rows × layers]``. Each returned view has shape
- ``(num_layers, max_slots, *inner_shape)`` matching ``MambaPool.State.conv[i]``
- / ``.temporal``. Mamba state is always token-granular (page_size == 1).
- """
- entry_bytes = mamba_entry_bytes(
- layer_num=layer_num,
- conv_state_shapes=conv_state_shapes,
- conv_dtype=conv_dtype,
- temporal_state_shape=temporal_state_shape,
- 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 "
- f"multiple of the conv-state itemsize {conv_itemsize} B"
- )
- assert anchor_bytes % conv_itemsize == 0, (
- f"misaligned mamba spec: anchor_bytes={anchor_bytes} is not a multiple "
- f"of the conv-state itemsize {conv_itemsize} B"
- )
- as_conv_dtype = raw.view(conv_dtype)
- conv_slot_stride_elems = entry_bytes // conv_itemsize
-
- offset_bytes_within_entry = 0
- conv_views: List[torch.Tensor] = []
- for shape in conv_state_shapes:
- inner_shape_bytes = _prod(shape) * conv_itemsize
- assert inner_shape_bytes % conv_itemsize == 0
- offset_elems = (anchor_bytes + offset_bytes_within_entry) // conv_itemsize
- stride = (
- inner_shape_bytes // conv_itemsize,
- conv_slot_stride_elems,
- ) + contiguous_strides(shape)
- conv_views.append(
- torch.as_strided(
- as_conv_dtype,
- size=(layer_num, max_slots) + tuple(shape),
- stride=stride,
- storage_offset=offset_elems,
- )
- )
- offset_bytes_within_entry += layer_num * inner_shape_bytes
-
- # The temporal view's storage_offset is computed in temporal-dtype elements
- # by integer-dividing a byte offset by itemsize, so every term of that byte
- # offset (entry stride, anchor, the conv region) must be a whole multiple of
- # itemsize or the offset truncates and mis-places the view.
- itemsize = temporal_dtype.itemsize
- assert entry_bytes % itemsize == 0, (
- f"misaligned mamba spec: per-slot entry_bytes={entry_bytes} is not a "
- f"multiple of the temporal-state itemsize {itemsize} B; the temporal "
- f"view's storage_offset would truncate and mis-place the state"
- )
- assert anchor_bytes % itemsize == 0, (
- f"misaligned mamba spec: anchor_bytes={anchor_bytes} is not a multiple "
- f"of the temporal-state itemsize {itemsize} B"
- )
- inner_shape_bytes = _prod(temporal_state_shape) * itemsize
- assert inner_shape_bytes % itemsize == 0, (
- f"misaligned mamba spec: temporal inner_shape_bytes={inner_shape_bytes} "
- f"is not a multiple of the temporal-state itemsize {itemsize} B"
- )
- assert (anchor_bytes + offset_bytes_within_entry) % itemsize == 0, (
- f"misaligned mamba spec: temporal region byte offset "
- f"{anchor_bytes + offset_bytes_within_entry} is not a multiple of the "
- f"temporal-state itemsize {itemsize} B"
- )
- offset_elems = (anchor_bytes + offset_bytes_within_entry) // itemsize
- as_temporal_dtype = raw.view(temporal_dtype)
- stride = (
- inner_shape_bytes // itemsize,
- entry_bytes // itemsize,
- ) + contiguous_strides(temporal_state_shape)
- temporal_view = torch.as_strided(
- as_temporal_dtype,
- size=(layer_num, max_slots) + tuple(temporal_state_shape),
- stride=stride,
- storage_offset=offset_elems,
- )
- return conv_views, temporal_view
diff --git a/python/sglang/srt/mem_cache/layout/token_major.py b/python/sglang/srt/mem_cache/layout/token_major.py
new file mode 100644
index 000000000000..98322d7f8431
--- /dev/null
+++ b/python/sglang/srt/mem_cache/layout/token_major.py
@@ -0,0 +1,261 @@
+"""Token-major cache views over one raw byte buffer.
+
+Each token slot owns one ENTRY holding every layer's state at fixed offsets,
+and a page is ``page_size`` entries back to back:
+
+ 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
+
+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
+
+import msgspec
+import torch
+
+ENTRY_ALIGN_BYTES = 32
+ROW_ALIGN_BYTES = 16
+
+
+def _prod(shape: Sequence[int]) -> int:
+ out = 1
+ for s in shape:
+ out *= int(s)
+ 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_name: str,
+ num_slots: int,
+ anchor_bytes: int = 0,
+) -> List[torch.Tensor]:
+ """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"
+ )
+ end = anchor_bytes + num_slots * layout.entry_bytes
+ assert end <= raw.numel() * raw.itemsize, (
+ 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)
+ 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)
+ views.append(
+ torch.as_strided(
+ as_dtype_view,
+ size=(num_slots, *part.row_shape),
+ stride=stride,
+ storage_offset=base_bytes // itemsize,
+ )
+ )
+ return views
+
+
+def mamba_entry_bytes(
+ *,
+ layer_num: int,
+ conv_state_shapes: Sequence[Sequence[int]],
+ conv_dtype: torch.dtype,
+ temporal_state_shape: Sequence[int],
+ temporal_dtype: torch.dtype,
+) -> int:
+ """Bytes occupied by one Mamba slot across all layers (conv + temporal)."""
+ total = 0
+ for shape in conv_state_shapes:
+ total += layer_num * _prod(shape) * conv_dtype.itemsize
+ total += layer_num * _prod(temporal_state_shape) * temporal_dtype.itemsize
+ return total
+
+
+def build_mamba_entry_views(
+ raw: torch.Tensor,
+ *,
+ layer_num: int,
+ conv_state_shapes: Sequence[Sequence[int]],
+ conv_dtype: torch.dtype,
+ temporal_state_shape: Sequence[int],
+ temporal_dtype: torch.dtype,
+ max_slots: int,
+ anchor_bytes: int = 0,
+) -> Tuple[List[torch.Tensor], torch.Tensor]:
+ """Per-slot envelope views over ``raw`` for Mamba state.
+
+ Layout per slot: ``[conv[0] rows × layers][conv[1] rows × layers]...
+ [temporal rows × layers]``. Each returned view has shape
+ ``(num_layers, max_slots, *inner_shape)`` matching ``MambaPool.State.conv[i]``
+ / ``.temporal``. Mamba state is always token-granular (page_size == 1).
+ """
+ entry_bytes = mamba_entry_bytes(
+ layer_num=layer_num,
+ conv_state_shapes=conv_state_shapes,
+ conv_dtype=conv_dtype,
+ temporal_state_shape=temporal_state_shape,
+ temporal_dtype=temporal_dtype,
+ )
+
+ conv_itemsize = conv_dtype.itemsize
+ assert entry_bytes % conv_itemsize == 0, (
+ f"misaligned mamba spec: per-slot entry_bytes={entry_bytes} is not a "
+ f"multiple of the conv-state itemsize {conv_itemsize} B"
+ )
+ assert anchor_bytes % conv_itemsize == 0, (
+ f"misaligned mamba spec: anchor_bytes={anchor_bytes} is not a multiple "
+ f"of the conv-state itemsize {conv_itemsize} B"
+ )
+ as_conv_dtype = raw.view(conv_dtype)
+ conv_slot_stride_elems = entry_bytes // conv_itemsize
+
+ offset_bytes_within_entry = 0
+ conv_views: List[torch.Tensor] = []
+ for shape in conv_state_shapes:
+ inner_shape_bytes = _prod(shape) * conv_itemsize
+ assert inner_shape_bytes % conv_itemsize == 0
+ offset_elems = (anchor_bytes + offset_bytes_within_entry) // conv_itemsize
+ stride = (
+ inner_shape_bytes // conv_itemsize,
+ conv_slot_stride_elems,
+ ) + _contiguous_strides(shape)
+ conv_views.append(
+ torch.as_strided(
+ as_conv_dtype,
+ size=(layer_num, max_slots) + tuple(shape),
+ stride=stride,
+ storage_offset=offset_elems,
+ )
+ )
+ offset_bytes_within_entry += layer_num * inner_shape_bytes
+
+ # The temporal view's storage_offset is computed in temporal-dtype elements
+ # by integer-dividing a byte offset by itemsize, so every term of that byte
+ # offset (entry stride, anchor, the conv region) must be a whole multiple of
+ # itemsize or the offset truncates and mis-places the view.
+ itemsize = temporal_dtype.itemsize
+ assert entry_bytes % itemsize == 0, (
+ f"misaligned mamba spec: per-slot entry_bytes={entry_bytes} is not a "
+ f"multiple of the temporal-state itemsize {itemsize} B; the temporal "
+ f"view's storage_offset would truncate and mis-place the state"
+ )
+ assert anchor_bytes % itemsize == 0, (
+ f"misaligned mamba spec: anchor_bytes={anchor_bytes} is not a multiple "
+ f"of the temporal-state itemsize {itemsize} B"
+ )
+ inner_shape_bytes = _prod(temporal_state_shape) * itemsize
+ assert inner_shape_bytes % itemsize == 0, (
+ f"misaligned mamba spec: temporal inner_shape_bytes={inner_shape_bytes} "
+ f"is not a multiple of the temporal-state itemsize {itemsize} B"
+ )
+ assert (anchor_bytes + offset_bytes_within_entry) % itemsize == 0, (
+ f"misaligned mamba spec: temporal region byte offset "
+ f"{anchor_bytes + offset_bytes_within_entry} is not a multiple of the "
+ f"temporal-state itemsize {itemsize} B"
+ )
+ offset_elems = (anchor_bytes + offset_bytes_within_entry) // itemsize
+ as_temporal_dtype = raw.view(temporal_dtype)
+ stride = (
+ inner_shape_bytes // itemsize,
+ entry_bytes // itemsize,
+ ) + _contiguous_strides(temporal_state_shape)
+ temporal_view = torch.as_strided(
+ as_temporal_dtype,
+ size=(layer_num, max_slots) + tuple(temporal_state_shape),
+ stride=stride,
+ storage_offset=offset_elems,
+ )
+ return conv_views, temporal_view
diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py
index 746439bb344c..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 (
@@ -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
@@ -588,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"
)
@@ -603,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,
@@ -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/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 f3aa017274c1..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,
@@ -788,13 +790,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 f243a5a51bfc..2dcaedabda5a 100644
--- a/python/sglang/srt/mem_cache/unified_memory_pool.py
+++ b/python/sglang/srt/mem_cache/unified_memory_pool.py
@@ -36,10 +36,13 @@
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 (
- build_mha_views,
- build_mla_views,
- build_page_major_mamba_views,
+from sglang.srt.mem_cache.layout.token_major import (
+ ROW_ALIGN_BYTES,
+ DenseEntryLayout,
+ DensePart,
+ align_entry_bytes,
+ build_dense_views,
+ build_mamba_entry_views,
)
from sglang.srt.mem_cache.memory_pool import (
HybridLinearKVPool,
@@ -64,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,
@@ -110,19 +126,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):
@@ -159,20 +162,39 @@ 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())
-
- # 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.
+ return align_entry_bytes(
+ self.layer_num * (self.k_row_bytes() + self.v_row_bytes())
+ )
def page_bytes(self, page_size: int) -> int:
return page_size * self.entry_bytes()
- 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 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(),
+ 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 get_dtype(self) -> torch.dtype:
return self.store_dtype
@@ -205,16 +227,27 @@ def __post_init__(self):
def kv_cache_dim(self) -> int:
return self.kv_lora_rank + self.qk_rope_head_dim
- def entry_bytes(self) -> int:
- return self.layer_num * self.kv_cache_dim * self.store_dtype.itemsize
-
- def view_tail_pad_bytes(self, page_size: int) -> int:
- return page_size * self.entry_bytes()
+ def row_bytes(self) -> int:
+ return self.kv_cache_dim * self.store_dtype.itemsize
- 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 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=(
+ 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 get_dtype(self) -> torch.dtype:
return self.store_dtype
@@ -266,7 +299,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."
)
@@ -274,12 +307,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.
@@ -296,10 +329,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 +380,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 +396,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]] = {}
@@ -376,9 +404,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:
@@ -490,22 +518,20 @@ 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 * 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,
+ 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_name=name,
+ num_slots=num_slots,
+ anchor_bytes=anchor_bytes,
+ )
+ for name in ("k", "v")
)
+ return k_views, v_views
def _build_mla_views(
self,
@@ -514,25 +540,21 @@ 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 * spec.blocks_per_page() * page_size,
- )
- return build_mla_views(
+ 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,
- layer_num=spec.layer_num,
- kv_cache_dim=spec.kv_cache_dim,
- store_dtype=spec.store_dtype,
- page_size=page_size,
- num_pages=num_pages,
+ layout=layout,
+ part_name="kv",
+ num_slots=num_slots,
anchor_bytes=anchor_bytes,
)
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,
@@ -545,16 +567,10 @@ 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
-
- 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.
+ """MHA KV pool whose per-layer `k_buffer`/`v_buffer` are token-major views
+ 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__(
@@ -577,7 +593,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 +611,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 +626,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,20 +668,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 _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))
-
- 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))
-
def set_kv_buffer_prefix_valid(self, *args, **kwargs):
raise NotImplementedError(
"prefix-valid commit is unsupported under the unified layout "
@@ -676,21 +676,11 @@ 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`: `(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__(
@@ -714,10 +704,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 +718,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 +738,20 @@ 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))
-
- 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))
-
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 +1431,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 +1449,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 +1873,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/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/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/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/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_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_full_loc_fast_path.py b/test/registered/unit/mem_cache/test_full_loc_fast_path.py
index 8c301633e5d9..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
@@ -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.
@@ -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_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()
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_kv_index_translator.py b/test/registered/unit/mem_cache/test_kv_index_translator.py
index 45768e04a2c0..bae4e724c9e9 100644
--- a/test/registered/unit/mem_cache/test_kv_index_translator.py
+++ b/test/registered/unit/mem_cache/test_kv_index_translator.py
@@ -43,7 +43,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 +80,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
@@ -177,56 +172,54 @@ 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):
+ """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 collapse in (True, False):
- allocator = _build_composite(ps, collapse=collapse)
- 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 two builders must agree element for element:
@@ -479,9 +472,7 @@ 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")
@@ -653,7 +644,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",
),
@@ -753,7 +744,7 @@ 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."""
+ across page sizes."""
for ps in (1, 4, 64):
src, _, rows, seq_lens, _, want_full, want_swa = self._built(
ps=ps, n=3 * ps
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..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,13 +17,13 @@
import torch
-from sglang.srt.mem_cache.layout.page_major import (
- build_mha_views,
- build_mla_views,
- build_page_major_mamba_views,
+from sglang.srt.mem_cache.layout.token_major import (
+ ENTRY_ALIGN_BYTES,
+ DenseEntryLayout,
+ DensePart,
+ build_dense_views,
+ build_mamba_entry_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
@@ -31,47 +31,53 @@
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 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 = _entry_bytes(row_bytes, layer_num)
+ 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,
- page_size=page_size,
- num_pages=num_pages,
+ layout=layout,
+ part_name="kv",
+ num_slots=num_pages * page_size,
anchor_bytes=0,
)
torch.manual_seed(0)
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 +85,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
@@ -90,37 +97,49 @@ 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)
- # 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_name=name,
+ num_slots=num_pages * page_size,
+ 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,33 +161,47 @@ 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,
- head_num=head_num,
- head_dim=head_dim,
- v_head_dim=head_dim,
- itemsize=store_dtype.itemsize,
- )
+ row_bytes = head_num * head_dim * store_dtype.itemsize
+ entry_bytes = _entry_bytes(2 * row_bytes, layer_num)
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,
+ 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_name=name,
+ num_slots=num_pages * page_size,
+ 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 +210,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",
@@ -206,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_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_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/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_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_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 cd0d66d0c228..93ad8a276665 100644
--- a/test/registered/unit/mem_cache/test_unified_mha_views.py
+++ b/test/registered/unit/mem_cache/test_unified_mha_views.py
@@ -11,14 +11,13 @@
# 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), CPU-only.
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)
"""
from sglang.test.ci.ci_register import register_cpu_ci
@@ -31,7 +30,8 @@
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.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,
@@ -47,11 +47,22 @@
_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
+
+
+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
def _mha_spec(head_dim=_D, v_head_dim=None, layer_num=_L, grow="down"):
@@ -66,150 +77,226 @@ 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, num_slots=num_pages * ps)
+ return (
+ build_dense_views(raw, part_name="k", **kw),
+ build_dense_views(raw, part_name="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,
- head_num=_H,
- head_dim=head_dim,
- v_head_dim=v_head_dim,
- store_dtype=_DTYPE,
- page_size=ps,
- num_pages=num_pages,
- )
-
-
-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)
+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 = _entry_bytes(head_dim, v_head_dim)
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_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",
+ 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_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()
+
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 = _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)
+ 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 +331,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 +369,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 +396,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 +411,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 +431,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 +460,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 +471,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 +487,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 +512,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 +551,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..80fe619c7d95 100644
--- a/test/registered/unit/mem_cache/test_unified_mla_views.py
+++ b/test/registered/unit/mem_cache/test_unified_mla_views.py
@@ -13,9 +13,9 @@
# ==============================================================================
"""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`.
+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`.
@@ -34,7 +34,8 @@
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.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,
@@ -46,13 +47,17 @@
_DEV = "cpu"
# Geometry kept tiny so every byte offset is hand-checkable; real K3 is
-# L=24, D=576 (=512+64).
+# L=24, D=576 (=512+64). 3 rows x 16 B 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
+# 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
def _mla_spec(grow="down", layer_num=_L):
@@ -92,8 +97,11 @@ 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_name="kv", num_slots=num_pages * ps
+ )
class TestMLASubPoolSpec(unittest.TestCase):
@@ -108,81 +116,86 @@ 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, 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 +212,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 +230,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 +241,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,9 +262,7 @@ 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}",
@@ -268,11 +270,11 @@ 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 kernel ids."""
+ 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 * _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])
@@ -282,9 +284,9 @@ def test_move_kv_cache_moves_page_envelopes(self):
# 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)
+ 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)
@@ -295,9 +297,7 @@ def test_move_kv_cache_moves_page_envelopes(self):
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)
- ]
+ 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})"
)
@@ -312,7 +312,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 +334,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 +367,12 @@ 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)
+
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/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)
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..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
@@ -14,18 +14,18 @@
"""`--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.
+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
@@ -96,7 +96,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,7 +108,7 @@ 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
+ # 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))
@@ -119,11 +119,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):
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()