chore: import upstream snapshot with attribution

python/sglang/srt/mem_cache/layout/page_major.py

@@ -0,0 +1,229 @@
+"""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 strided 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_page_major_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`` in the page-major layer-major layout.
+
+    Each returned view is 4-D ``(num_pages, page_size, head_num, head_dim*)``
+    with constant strides:
+
+        stride[0] = page_bytes / itemsize     # next page
+        stride[1] = k_row_bytes / itemsize    # next slot within layer L's K block
+        stride[2] = head_dim                  # next head
+        stride[3] = 1                         # next element
+
+    V is analogous with ``v_row_bytes`` / ``v_head_dim``. A token id ``t`` reads
+    page ``t // page_size``, slot ``t % page_size``.
+    """
+    itemsize = store_dtype.itemsize
+    k_row_bytes = head_num * head_dim * itemsize
+    v_row_bytes = head_num * v_head_dim * itemsize
+    entry_bytes = layer_num * (k_row_bytes + v_row_bytes)
+    page_bytes = page_size * entry_bytes
+    assert anchor_bytes % itemsize == 0
+    assert k_row_bytes % itemsize == 0
+    assert v_row_bytes % itemsize == 0
+    assert page_bytes % itemsize == 0
+
+    as_dtype_view = raw.view(store_dtype)
+    stride_page = page_bytes // itemsize
+    stride_tok_k = k_row_bytes // itemsize
+    stride_tok_v = v_row_bytes // itemsize
+
+    k_shape = (num_pages, page_size, head_num, head_dim)
+    v_shape = (num_pages, page_size, head_num, v_head_dim)
+    k_stride = (stride_page, stride_tok_k, head_dim, 1)
+    v_stride = (stride_page, stride_tok_v, v_head_dim, 1)
+
+    k_buffer: List[torch.Tensor] = []
+    v_buffer: List[torch.Tensor] = []
+    for layer in range(layer_num):
+        # Layer L's K block starts at L * page_size * (k_row + v_row); V follows.
+        k_base_bytes = anchor_bytes + layer * page_size * (k_row_bytes + v_row_bytes)
+        v_base_bytes = k_base_bytes + page_size * k_row_bytes
+        assert k_base_bytes % itemsize == 0
+        assert v_base_bytes % itemsize == 0
+        k_buffer.append(
+            torch.as_strided(
+                as_dtype_view,
+                size=k_shape,
+                stride=k_stride,
+                storage_offset=k_base_bytes // itemsize,
+            )
+        )
+        v_buffer.append(
+            torch.as_strided(
+                as_dtype_view,
+                size=v_shape,
+                stride=v_stride,
+                storage_offset=v_base_bytes // itemsize,
+            )
+        )
+    return k_buffer, v_buffer
+
+
+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