[unified-memory] Token-major dense views for the unified memory pool (3/7)

docs/docs/advanced_features/server_arguments.mdx

@@ -534,13 +534,13 @@ Please consult the documentation below and [server_args.py](https://github.com/s
     </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-page-major-kv-layout`</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>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`).</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>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.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`False`</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
     </tr>
         <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-unified-memory`</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>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.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>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.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`False`</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
     </tr>

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,

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,25 +722,22 @@ def handle_page_major_kv_layout(server_args: Any):
             "tokenspeed_mla",
             "flashmla",
         }
-    elif cfg.enable_unified_memory:
+    else:
         allowed_full = {
             "triton",
             "fa3",
             "fa4",
             "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):

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,

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)