L3 3.0: host-slab accessor (layer_page_first <-> contiguous blob)
CpSharedL2SlabAccessor (HIGH-1): memcpy-gather a page's N_layers strided slab slices into a contiguous buffer (spill) + scatter a contiguous blob back into the strided slices (reload) — the cheap RAM bridge that lets host stay layer_page_first (inference perf) while disk is page-contiguous. Generic CpL3SlabLayout(n_layers, page_num, slice_bytes) with verified-from-source factories: for_mla (layer_num,page_num,page_size,kv_cache_dim,itemsize -> 2.74 MiB/page) and for_index (n_active_layers,indexer_page_num,indexer_page_stride -> 0.17/0.63 MiB/page). gather_into/ scatter_from take an offset so the spill writes straight after the 64B slot header (zero extra copy). Pure (mmap+ints), no torch. 8/8 tests. Canonical layer 0..N-1 order shared by spill+reload. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
115
python/sglang/srt/mem_cache/cp_l3_slab_accessor.py
Normal file
115
python/sglang/srt/mem_cache/cp_l3_slab_accessor.py
Normal file
@@ -0,0 +1,115 @@
|
||||
# Copyright 2023-2024 SGLang Team
|
||||
# Licensed under the Apache License, Version 2.0
|
||||
"""CP HiCache L3 — host-slab byte accessor (the layer_page_first <-> contiguous-blob bridge).
|
||||
|
||||
The CP shared-L2 host slab is ``layer_page_first`` (kept for inference perf): a page's bytes are N_layers
|
||||
*strided* slices, not contiguous. L3 wants one contiguous page blob on disk. This accessor does the
|
||||
**memcpy-gather** (spill: strided slab slices -> contiguous buffer) and **scatter** (reload: contiguous blob
|
||||
-> strided slab slices) — measured faster than vectored I/O under O_DIRECT, and the cheap RAM side of the
|
||||
spill/reload. The on-disk slice order (layer 0..N-1) is canonical: spill and reload use the same order.
|
||||
|
||||
Verified layouts (memory_pool_host.py, GLM + layer_page_first):
|
||||
* target_kv / draft_kv (MLA): slab tensor ``(layer_num, page_num, page_size, 1, kv_cache_dim)`` ->
|
||||
slice(layer,page) at ``layer*(page_num*slice) + page*slice``, slice = ``page_size*kv_cache_dim*itemsize``.
|
||||
* index_k (NSA): ``(n_active_layers, indexer_page_num, 1, indexer_page_stride)`` ->
|
||||
slice(slot,page) at ``slot*(indexer_page_num*stride) + page*stride``, slice = ``indexer_page_stride``.
|
||||
Both reduce to a generic ``(n_layers, layer_stride_bytes, slice_bytes)``. Pure (mmap + ints): no torch.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from sglang.srt.mem_cache.cp_l3_disk import _l3_failfast
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CpL3SlabLayout:
|
||||
n_layers: int # layer_num (target/draft) or index_active_layer_num (index_k)
|
||||
page_num: int # slab pages along the page axis (for bounds + layer stride)
|
||||
slice_bytes: int # bytes of one (layer, page) contiguous slice
|
||||
|
||||
def __post_init__(self):
|
||||
if self.n_layers <= 0 or self.page_num <= 0 or self.slice_bytes <= 0:
|
||||
_l3_failfast(
|
||||
f"CpL3SlabLayout requires positive dims; got n_layers={self.n_layers} "
|
||||
f"page_num={self.page_num} slice_bytes={self.slice_bytes}"
|
||||
)
|
||||
|
||||
@property
|
||||
def layer_stride_bytes(self) -> int:
|
||||
return self.page_num * self.slice_bytes
|
||||
|
||||
@property
|
||||
def page_blob_bytes(self) -> int:
|
||||
return self.n_layers * self.slice_bytes
|
||||
|
||||
@property
|
||||
def total_bytes(self) -> int:
|
||||
return self.n_layers * self.layer_stride_bytes
|
||||
|
||||
@classmethod
|
||||
def for_mla(cls, *, layer_num: int, page_num: int, page_size: int, kv_cache_dim: int, itemsize: int):
|
||||
return cls(n_layers=layer_num, page_num=page_num, slice_bytes=page_size * kv_cache_dim * itemsize)
|
||||
|
||||
@classmethod
|
||||
def for_index(cls, *, n_active_layers: int, indexer_page_num: int, indexer_page_stride_size: int):
|
||||
return cls(n_layers=n_active_layers, page_num=indexer_page_num, slice_bytes=indexer_page_stride_size)
|
||||
|
||||
|
||||
class CpSharedL2SlabAccessor:
|
||||
"""Gather/scatter one global page's strided layer slices to/from a contiguous buffer."""
|
||||
|
||||
def __init__(self, slab_mmap, layout: CpL3SlabLayout):
|
||||
self._mm = slab_mmap
|
||||
self._lo = layout
|
||||
if len(slab_mmap) < layout.total_bytes:
|
||||
_l3_failfast(
|
||||
f"slab mmap {len(slab_mmap)} bytes < layout total {layout.total_bytes} "
|
||||
f"(n_layers={layout.n_layers} page_num={layout.page_num} slice={layout.slice_bytes})"
|
||||
)
|
||||
|
||||
@property
|
||||
def page_blob_bytes(self) -> int:
|
||||
return self._lo.page_blob_bytes
|
||||
|
||||
def _slice_offset(self, layer: int, global_page: int) -> int:
|
||||
return layer * self._lo.layer_stride_bytes + global_page * self._lo.slice_bytes
|
||||
|
||||
def _check_page(self, global_page: int) -> None:
|
||||
if global_page < 0 or global_page >= self._lo.page_num:
|
||||
_l3_failfast(f"global_page {global_page} out of range [0,{self._lo.page_num})")
|
||||
|
||||
def gather_into(self, global_page: int, dst, dst_offset: int = 0) -> int:
|
||||
"""Copy this page's N_layers strided slices contiguously into ``dst[dst_offset:]``. Returns bytes
|
||||
written (page_blob_bytes). ``dst`` is typically the aligned disk-slot buffer at HEADER_BYTES."""
|
||||
self._check_page(global_page)
|
||||
sb = self._lo.slice_bytes
|
||||
blob = self._lo.page_blob_bytes
|
||||
if len(dst) < dst_offset + blob:
|
||||
_l3_failfast(f"dst too small: need {dst_offset + blob}, have {len(dst)}")
|
||||
dview = memoryview(dst)
|
||||
for layer in range(self._lo.n_layers):
|
||||
off = self._slice_offset(layer, global_page)
|
||||
dview[dst_offset + layer * sb : dst_offset + (layer + 1) * sb] = self._mm[off : off + sb]
|
||||
return blob
|
||||
|
||||
def gather(self, global_page: int) -> bytes:
|
||||
"""Convenience: return the contiguous page blob as bytes."""
|
||||
buf = bytearray(self._lo.page_blob_bytes)
|
||||
self.gather_into(global_page, buf, 0)
|
||||
return bytes(buf)
|
||||
|
||||
def scatter_from(self, global_page: int, src, src_offset: int = 0) -> int:
|
||||
"""Write a contiguous page blob (``src[src_offset:]``) back into this page's N_layers strided slices.
|
||||
Returns bytes written. Inverse of gather_into (same canonical layer order)."""
|
||||
self._check_page(global_page)
|
||||
sb = self._lo.slice_bytes
|
||||
blob = self._lo.page_blob_bytes
|
||||
if len(src) < src_offset + blob:
|
||||
_l3_failfast(f"src too small: need {src_offset + blob}, have {len(src)}")
|
||||
sview = memoryview(src)
|
||||
for layer in range(self._lo.n_layers):
|
||||
off = self._slice_offset(layer, global_page)
|
||||
self._mm[off : off + sb] = sview[src_offset + layer * sb : src_offset + (layer + 1) * sb]
|
||||
return blob
|
||||
Reference in New Issue
Block a user