feat(unified-memory): dense KV views for uniform-row MHA/SWA models (#34602)
Co-authored-by: Caihua Li <caihua.li@bytedance.com> Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>
This commit is contained in:
caihuali95
co-authored by
Caihua Li
Cheng Wan
Caihua Li
Cheng Wan
parent
007ef5e23a
commit
4bea51d885
@@ -1,10 +1,10 @@
|
||||
"""
|
||||
End-to-end accuracy test for the page-major KV layout on a hybrid-SWA MoE model.
|
||||
End-to-end accuracy test for the unified memory pool on a hybrid-SWA MoE model.
|
||||
|
||||
Launches gpt-oss-20b with ``--enable-page-major-kv-layout`` on the Triton
|
||||
attention backend and checks that GSM8K accuracy holds. This exercises the
|
||||
SWA + full-attention KV pools under the page-granularity envelope layout
|
||||
(SWAKVPool routes both sub-pools through PageMajorMHATokenToKVPool).
|
||||
Launches gpt-oss-20b with ``--enable-unified-memory`` on the Triton attention
|
||||
backend and checks that GSM8K accuracy holds. This exercises the SWA +
|
||||
full-attention KV sub-pools stored as per-layer views in the unified
|
||||
page-major envelope.
|
||||
|
||||
Registered to the label-gated ``run-ci-extra`` suite (opt-in, not per-commit).
|
||||
|
||||
@@ -22,9 +22,17 @@ from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST_MXFP4_WITH_MOE
|
||||
|
||||
register_cuda_ci(est_time=420, stage="extra-a", runner_config="1-gpu-large")
|
||||
|
||||
_UNIFIED_COMMON_ARGS = [
|
||||
"--enable-unified-memory",
|
||||
"--mem-fraction-static",
|
||||
"0.70",
|
||||
"--cuda-graph-backend-prefill=disabled",
|
||||
]
|
||||
|
||||
class TestPageMajorGptOss(DefaultServerBase):
|
||||
"""Page-major KV layout on gpt-oss-20b (hybrid-SWA MoE), Triton backend."""
|
||||
|
||||
class TestUnifiedGptOssTriton(DefaultServerBase):
|
||||
"""Unified pool on gpt-oss-20b (hybrid-SWA MoE), Triton pinned: dense
|
||||
MHA/SWA views through the reference backend."""
|
||||
|
||||
model = DEFAULT_MODEL_NAME_FOR_TEST_MXFP4_WITH_MOE
|
||||
|
||||
@@ -33,16 +41,7 @@ class TestPageMajorGptOss(DefaultServerBase):
|
||||
num_shots = 5
|
||||
parallel = 32
|
||||
|
||||
other_args = [
|
||||
"--enable-page-major-kv-layout",
|
||||
# The envelope's strided 4-D K/V views are only read by the Triton
|
||||
# attention kernels (the layout's validator enforces this).
|
||||
"--attention-backend",
|
||||
"triton",
|
||||
"--mem-fraction-static",
|
||||
"0.70",
|
||||
"--cuda-graph-backend-prefill=disabled",
|
||||
]
|
||||
other_args = _UNIFIED_COMMON_ARGS + ["--attention-backend", "triton"]
|
||||
|
||||
def test_gsm8k(self):
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_few_shot_gsm8k
|
||||
|
||||
@@ -1,12 +1,12 @@
|
||||
"""
|
||||
End-to-end accuracy test for the page-major KV layout on a GDN-hybrid model.
|
||||
End-to-end accuracy test for the unified memory pool on a GDN-hybrid model.
|
||||
|
||||
Launches Qwen3.5-4B (a gated-delta-net / linear-attention hybrid) with
|
||||
``--enable-page-major-kv-layout`` on the Triton attention + linear-attn + Mamba
|
||||
backends and checks that GSM8K accuracy holds. This exercises the page-major
|
||||
path most prone to subtle bugs: the Mamba conv/SSM state stored as a strided
|
||||
envelope view, plus the full-attention KV pool, both read/written by the GDN
|
||||
prefill and decode kernels.
|
||||
``--enable-unified-memory`` on the Triton attention + linear-attn + Mamba
|
||||
backends and checks that GSM8K accuracy holds. This exercises the unified
|
||||
envelope's most bug-prone path: the Mamba conv/SSM state stored as a strided
|
||||
envelope view, plus the full-attention KV stored as per-layer views,
|
||||
both read/written by the GDN prefill and decode kernels.
|
||||
|
||||
Registered to the label-gated ``run-ci-extra`` suite (opt-in, not per-commit).
|
||||
|
||||
@@ -24,35 +24,34 @@ from sglang.test.test_utils import DEFAULT_HYBRID_GDN_SMALL_MODEL_NAME_FOR_TEST
|
||||
|
||||
register_cuda_ci(est_time=300, stage="extra-a", runner_config="1-gpu-large")
|
||||
|
||||
_UNIFIED_COMMON_ARGS = [
|
||||
"--trust-remote-code",
|
||||
"--mem-fraction-static",
|
||||
"0.85",
|
||||
"--enable-unified-memory",
|
||||
"--linear-attn-backend",
|
||||
"triton",
|
||||
"--mamba-backend",
|
||||
"triton",
|
||||
]
|
||||
|
||||
class TestPageMajorQwenHybrid(DefaultServerBase):
|
||||
"""Page-major KV layout on Qwen3.5-4B (GDN-hybrid), Triton backends."""
|
||||
|
||||
class TestUnifiedQwenHybridTriton(DefaultServerBase):
|
||||
"""Unified pool on Qwen3.5-4B (GDN-hybrid), Triton pinned: dense
|
||||
full-attention views + strided conv/SSM state through the reference
|
||||
backends."""
|
||||
|
||||
model = DEFAULT_HYBRID_GDN_SMALL_MODEL_NAME_FOR_TEST
|
||||
|
||||
# Measured in this harness: baseline (no page-major) and page-major both
|
||||
# ~0.86; the 0.80 threshold leaves margin for run-to-run noise while still
|
||||
# catching the prefill-state corruption that page-major hit before the
|
||||
# gather/scatter fix in gdn_backend.forward_extend (which dropped it to ~0.61).
|
||||
# Measured ~0.86 in this harness on both the static pools and the envelope
|
||||
# layout; 0.80 leaves noise margin and still catches a corrupted prefill
|
||||
# state, which reads ~0.61.
|
||||
gsm8k_threshold = 0.80
|
||||
num_gsm8k_questions = 200
|
||||
num_shots = 5
|
||||
parallel = 32
|
||||
|
||||
other_args = [
|
||||
"--trust-remote-code",
|
||||
"--mem-fraction-static",
|
||||
"0.85",
|
||||
"--enable-page-major-kv-layout",
|
||||
# Only the Triton attention / linear-attn / Mamba kernels read the
|
||||
# strided envelope K/V and conv/SSM state (enforced by the validator).
|
||||
"--attention-backend",
|
||||
"triton",
|
||||
"--linear-attn-backend",
|
||||
"triton",
|
||||
"--mamba-backend",
|
||||
"triton",
|
||||
]
|
||||
other_args = _UNIFIED_COMMON_ARGS + ["--attention-backend", "triton"]
|
||||
|
||||
def test_gsm8k(self):
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_few_shot_gsm8k
|
||||
|
||||
@@ -23,7 +23,7 @@ always PHYSICAL. Two routing contracts are pinned here:
|
||||
`UnifiedSWAKVPool` asserts it's present (the unified memory pool always precomputes
|
||||
it); `HybridLinearKVPool` falls back to `loc` for a static (non-shared) pool,
|
||||
where `loc` is itself already physical.
|
||||
2. SWA. The swa-physical loc rides the backend `swa_out_cache_loc` rail
|
||||
2. SWA. The swa-physical loc rides the backend `swa_out_cache_loc` slot
|
||||
(`KVWriteLoc.swa_loc`) and is written directly.
|
||||
|
||||
Pure dispatch tests: the inner sub-pools are recording stubs, so no GPU / real
|
||||
@@ -127,7 +127,7 @@ class TestUnifiedSWARouting(unittest.TestCase):
|
||||
|
||||
self.assertEqual(len(pool.swa_kv_pool.calls), 1)
|
||||
forwarded, kwargs = pool.swa_kv_pool.calls[0]
|
||||
# SWA write rides the backend rail: forward the swa-physical loc directly.
|
||||
# SWA write rides the backend slot: forward the swa-physical loc directly.
|
||||
self.assertIs(forwarded, swa_phys)
|
||||
self.assertNotIn("already_physical", kwargs)
|
||||
# Full pool untouched for an SWA layer.
|
||||
@@ -138,7 +138,7 @@ class TestUnifiedSWARouting(unittest.TestCase):
|
||||
virtual_loc = torch.tensor([10, 11, 12], dtype=torch.int64)
|
||||
|
||||
layer = types.SimpleNamespace(layer_id=1) # SWA layer
|
||||
# No swa_loc bundled -> the rail contract is violated; must assert
|
||||
# No swa_loc bundled -> the write-loc contract is violated; must assert
|
||||
# rather than silently writing wrong (un-translated) locations.
|
||||
with self.assertRaises(AssertionError):
|
||||
pool.set_kv_buffer(
|
||||
@@ -149,6 +149,46 @@ class TestUnifiedSWARouting(unittest.TestCase):
|
||||
)
|
||||
|
||||
|
||||
class TestUnifiedSWATombstoneClamp(unittest.TestCase):
|
||||
"""`UnifiedSWAKVPool.translate_loc_from_full_to_swa` must clamp tombstoned
|
||||
ids to the reserved padding sink (0).
|
||||
|
||||
A token whose swa page was freed carries -1 in `virtual_to_physical`. Before
|
||||
the clamp, that produced a negative id, which a captured graph stores at a
|
||||
negative offset from the buffer base. The composite allocator's method of
|
||||
the same name already clamped; this path did not.
|
||||
"""
|
||||
|
||||
def _make_bare_pool(self, page_size, v2p, multiplier=1):
|
||||
from sglang.srt.mem_cache.multi_ended_allocator import MultiEndedAllocator
|
||||
from sglang.srt.mem_cache.unified_memory_pool import UnifiedSWAKVPool
|
||||
|
||||
# A real sub-allocator (not a stand-in): the translation reads its v2p
|
||||
# table, and the pool reaches it through the allocator's own method.
|
||||
swa_allocator = object.__new__(MultiEndedAllocator)
|
||||
swa_allocator.page_size = page_size
|
||||
swa_allocator.virtual_to_physical = v2p
|
||||
swa_allocator.kernel_page_multiplier = multiplier
|
||||
pool = object.__new__(UnifiedSWAKVPool)
|
||||
pool._swa_allocator = swa_allocator
|
||||
return pool
|
||||
|
||||
def test_tombstoned_id_lands_on_sink(self):
|
||||
for ps, mult in ((1, 1), (4, 1), (4, 6)):
|
||||
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.
|
||||
kv_indices = torch.tensor([0, ps, 2 * ps], dtype=torch.int64)
|
||||
out = pool.translate_loc_from_full_to_swa(kv_indices)
|
||||
self.assertEqual(out.dtype, torch.int64)
|
||||
self.assertTrue(
|
||||
bool((out >= 0).all().item()),
|
||||
f"tombstoned swa id stayed negative at page_size={ps}, "
|
||||
f"multiplier={mult}: {out}",
|
||||
)
|
||||
self.assertEqual(int(out[1].item()), 0)
|
||||
|
||||
|
||||
class TestHybridLinearFullLocRouting(unittest.TestCase):
|
||||
"""`HybridLinearKVPool.set_kv_buffer` (non-MLA) writes the full-physical
|
||||
`full_loc` from the write metadata when present (unified memory pool), else the
|
||||
|
||||
@@ -11,18 +11,18 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
"""Unit tests for the page-major layer-major byte layout.
|
||||
"""Unit tests for the page-major envelope byte layout.
|
||||
|
||||
The subject here is the ENVELOPE — the byte layout the unified pool stores its
|
||||
KV in — pinned through ``MHASubPoolSpec``'s offset math. The dense 3-D views
|
||||
the pool exposes over the same bytes are covered by
|
||||
``test_unified_mha_views.py``, which also pins the view addressing
|
||||
against the envelope formula byte for byte.
|
||||
|
||||
Verifies that:
|
||||
1. The new 4-D ``_build_mha_views`` output exposes correct byte addresses
|
||||
for each (layer, page, tok_in_page, head, dim) — under both the
|
||||
degenerate ``page_size=1`` case (byte-identical to the old per-token
|
||||
envelope) and the new ``page_size>1`` layer-major case.
|
||||
2. ``MHASubPoolSpec.layer_k_offset_in_page`` /
|
||||
``layer_v_offset_in_page`` math matches the layout intent.
|
||||
3. ``set_kv_buffer`` round-trips correctly for both page sizes.
|
||||
4. Compaction (``move_kv_cache_native``) moves the right bytes for both
|
||||
page sizes via the 4-D advanced indexing path.
|
||||
1. ``MHASubPoolSpec.layer_k_offset_in_page`` / ``layer_v_offset_in_page`` math
|
||||
matches the layout intent at ``page_size == 1`` and ``> 1``.
|
||||
2. ``move_kv_cache_native`` (the stock per-layer 3-D move) stays byte-exact.
|
||||
|
||||
CPU-only — no GPU / Triton needed.
|
||||
|
||||
@@ -38,11 +38,7 @@ import unittest
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.memory_pool import move_kv_cache_native
|
||||
from sglang.srt.mem_cache.unified_memory_pool import (
|
||||
MambaSubPoolSpec,
|
||||
MHASubPoolSpec,
|
||||
UnifiedKVPool,
|
||||
)
|
||||
from sglang.srt.mem_cache.unified_memory_pool import MHASubPoolSpec
|
||||
|
||||
_DEV = "cpu"
|
||||
|
||||
@@ -58,18 +54,6 @@ def _make_mha_spec(name, grow, layer_num=2, head_num=2, head_dim=4):
|
||||
)
|
||||
|
||||
|
||||
def _make_mamba_spec(name, grow, layer_num=2):
|
||||
return MambaSubPoolSpec(
|
||||
name=name,
|
||||
layer_num=layer_num,
|
||||
conv_state_shapes=((4, 3),),
|
||||
conv_dtype=torch.float32,
|
||||
temporal_state_shape=(2, 2, 2),
|
||||
temporal_dtype=torch.float32,
|
||||
grow_direction=grow,
|
||||
)
|
||||
|
||||
|
||||
class TestMHASpecLayerOffsets(unittest.TestCase):
|
||||
"""Verify ``layer_k_offset_in_page`` / ``layer_v_offset_in_page`` math."""
|
||||
|
||||
@@ -113,175 +97,10 @@ class TestMHASpecLayerOffsets(unittest.TestCase):
|
||||
self.assertEqual(spec.page_bytes(ps), ps * spec.entry_bytes())
|
||||
|
||||
|
||||
class TestBuildMHAViews(unittest.TestCase):
|
||||
"""Verify the 4-D view shape + strides at both page sizes."""
|
||||
|
||||
def _build(self, page_size, layer_num=3, head_num=2, head_dim=4, n_full_slots=64):
|
||||
full = _make_mha_spec(
|
||||
"full", "up", layer_num=layer_num, head_num=head_num, head_dim=head_dim
|
||||
)
|
||||
swa = _make_mha_spec(
|
||||
"swa", "down", layer_num=2, head_num=head_num, head_dim=head_dim
|
||||
)
|
||||
# Pad to ensure max_slots % page_size == 0 in both sub-pools.
|
||||
# entry_bytes is fixed per spec; size accordingly.
|
||||
total = full.entry_bytes() * n_full_slots + swa.entry_bytes() * n_full_slots
|
||||
pool = UnifiedKVPool(
|
||||
total_bytes=total,
|
||||
sub_pool_specs=[full, swa],
|
||||
device=_DEV,
|
||||
enable_memory_saver=False,
|
||||
page_size=page_size,
|
||||
)
|
||||
return pool, full
|
||||
|
||||
def test_view_shape_is_4d(self):
|
||||
for ps in [1, 8]:
|
||||
pool, spec = self._build(page_size=ps)
|
||||
k_views, v_views = pool.mha_views_for("full")
|
||||
self.assertEqual(len(k_views), spec.layer_num)
|
||||
max_slots = pool.max_slots("full")
|
||||
for L in range(spec.layer_num):
|
||||
self.assertEqual(k_views[L].ndim, 4)
|
||||
self.assertEqual(
|
||||
tuple(k_views[L].shape),
|
||||
(max_slots // ps, ps, spec.head_num, spec.head_dim),
|
||||
)
|
||||
self.assertEqual(
|
||||
tuple(v_views[L].shape),
|
||||
(max_slots // ps, ps, spec.head_num, spec.v_head_dim),
|
||||
)
|
||||
|
||||
def test_strides_at_page_size_1_match_envelope(self):
|
||||
"""At ps=1, the 4-D view's stride[0] equals what today's 3-D view's
|
||||
stride[0] would have been (= entry_bytes / itemsize)."""
|
||||
pool, spec = self._build(page_size=1, layer_num=4, head_num=3, head_dim=8)
|
||||
k_views, _ = pool.mha_views_for("full")
|
||||
itemsize = spec.store_dtype.itemsize
|
||||
for L in range(spec.layer_num):
|
||||
# stride[0] = page_bytes/itemsize = entry_bytes/itemsize at ps=1
|
||||
self.assertEqual(k_views[L].stride(0), spec.entry_bytes() // itemsize)
|
||||
# stride[1] = k_row/itemsize (within-page token stride)
|
||||
self.assertEqual(k_views[L].stride(1), spec.k_row_bytes() // itemsize)
|
||||
# stride[2] = head_dim (head stride)
|
||||
self.assertEqual(k_views[L].stride(2), spec.head_dim)
|
||||
# stride[3] = 1 (innermost)
|
||||
self.assertEqual(k_views[L].stride(3), 1)
|
||||
|
||||
def test_strides_at_page_size_gt_1(self):
|
||||
pool, spec = self._build(page_size=8, layer_num=4, head_num=3, head_dim=8)
|
||||
k_views, _ = pool.mha_views_for("full")
|
||||
itemsize = spec.store_dtype.itemsize
|
||||
for L in range(spec.layer_num):
|
||||
# page_bytes = 8 * 4 * (k_row + v_row); stride[0] = that / itemsize
|
||||
self.assertEqual(k_views[L].stride(0), spec.page_bytes(8) // itemsize)
|
||||
# token stride within layer L's K block = k_row/itemsize
|
||||
self.assertEqual(k_views[L].stride(1), spec.k_row_bytes() // itemsize)
|
||||
self.assertEqual(k_views[L].stride(2), spec.head_dim)
|
||||
self.assertEqual(k_views[L].stride(3), 1)
|
||||
|
||||
def test_distinct_layers_dont_alias_at_page_size_gt_1(self):
|
||||
"""Writes to layer 0 must not affect layer 1's K/V values (under
|
||||
layer-major within-page layout)."""
|
||||
pool, spec = self._build(page_size=8, layer_num=3, head_num=2, head_dim=4)
|
||||
k_views, v_views = pool.mha_views_for("full")
|
||||
# Set page 0, token 3, layer 0 K to a distinct pattern.
|
||||
target_val = 0.5
|
||||
k_views[0][0, 3] = target_val
|
||||
# Layer 1 K at the same (page, tok) should remain at default (0.0).
|
||||
self.assertFalse(torch.all(k_views[1][0, 3] == target_val))
|
||||
self.assertTrue(torch.all(k_views[1][0, 3] == 0.0))
|
||||
# And layer 0 V at the same (page, tok) should remain at default.
|
||||
self.assertFalse(torch.all(v_views[0][0, 3] == target_val))
|
||||
self.assertTrue(torch.all(v_views[0][0, 3] == 0.0))
|
||||
|
||||
def test_distinct_pages_dont_alias_at_page_size_gt_1(self):
|
||||
"""Writes to one page must not affect another page."""
|
||||
pool, spec = self._build(page_size=8, layer_num=3, head_num=2, head_dim=4)
|
||||
k_views, _ = pool.mha_views_for("full")
|
||||
# Set page 0, token 3, layer 0 K to a distinct pattern.
|
||||
k_views[0][0, 3] = 1.25
|
||||
# Page 1, token 3, layer 0 K should remain at default.
|
||||
self.assertTrue(torch.all(k_views[0][1, 3] == 0.0))
|
||||
|
||||
|
||||
class TestMoveKVCacheNative4D(unittest.TestCase):
|
||||
"""Verify ``move_kv_cache_native`` handles 4-D buffers at both
|
||||
page_size=1 (degenerate envelope) and page_size>1 (layer-major)."""
|
||||
|
||||
def _build_buffer(
|
||||
self, page_size, layer_num=2, head_num=2, head_dim=4, n_full_slots=64
|
||||
):
|
||||
full = _make_mha_spec(
|
||||
"full", "up", layer_num=layer_num, head_num=head_num, head_dim=head_dim
|
||||
)
|
||||
swa = _make_mha_spec(
|
||||
"swa", "down", layer_num=2, head_num=head_num, head_dim=head_dim
|
||||
)
|
||||
total = full.entry_bytes() * n_full_slots + swa.entry_bytes() * n_full_slots
|
||||
pool = UnifiedKVPool(
|
||||
total_bytes=total,
|
||||
sub_pool_specs=[full, swa],
|
||||
device=_DEV,
|
||||
enable_memory_saver=False,
|
||||
page_size=page_size,
|
||||
)
|
||||
return pool
|
||||
|
||||
def test_move_kv_cache_page_size_1(self):
|
||||
pool = self._build_buffer(page_size=1, layer_num=2, head_num=2, head_dim=4)
|
||||
k_views, v_views = pool.mha_views_for("full")
|
||||
# Write distinct markers at source slots 5, 6.
|
||||
for L in range(2):
|
||||
k_views[L][5, 0] = float(L + 1)
|
||||
v_views[L][5, 0] = -float(L + 1)
|
||||
k_views[L][6, 0] = float(L + 10)
|
||||
v_views[L][6, 0] = -float(L + 10)
|
||||
# Move 5 -> 8 and 6 -> 9.
|
||||
move_kv_cache_native(
|
||||
k_views,
|
||||
v_views,
|
||||
tgt_loc=torch.tensor([8, 9], dtype=torch.int64),
|
||||
src_loc=torch.tensor([5, 6], dtype=torch.int64),
|
||||
page_size=1,
|
||||
)
|
||||
for L in range(2):
|
||||
self.assertTrue(torch.all(k_views[L][8, 0] == float(L + 1)))
|
||||
self.assertTrue(torch.all(v_views[L][8, 0] == -float(L + 1)))
|
||||
self.assertTrue(torch.all(k_views[L][9, 0] == float(L + 10)))
|
||||
self.assertTrue(torch.all(v_views[L][9, 0] == -float(L + 10)))
|
||||
|
||||
def test_move_kv_cache_page_size_gt_1(self):
|
||||
ps = 8
|
||||
pool = self._build_buffer(page_size=ps, layer_num=2, head_num=2, head_dim=4)
|
||||
k_views, v_views = pool.mha_views_for("full")
|
||||
# Write markers at token ids 5 and 14 (different pages).
|
||||
for L in range(2):
|
||||
# token 5 = (page 0, tok 5)
|
||||
k_views[L][0, 5] = float(L + 1)
|
||||
v_views[L][0, 5] = -float(L + 1)
|
||||
# token 14 = (page 1, tok 6)
|
||||
k_views[L][1, 6] = float(L + 10)
|
||||
v_views[L][1, 6] = -float(L + 10)
|
||||
# Move token 5 -> token 23 (page 2, tok 7) and 14 -> 31 (page 3, tok 7).
|
||||
move_kv_cache_native(
|
||||
k_views,
|
||||
v_views,
|
||||
tgt_loc=torch.tensor([23, 31], dtype=torch.int64),
|
||||
src_loc=torch.tensor([5, 14], dtype=torch.int64),
|
||||
page_size=ps,
|
||||
)
|
||||
for L in range(2):
|
||||
# 23 = page 2, tok 7
|
||||
self.assertTrue(torch.all(k_views[L][2, 7] == float(L + 1)))
|
||||
self.assertTrue(torch.all(v_views[L][2, 7] == -float(L + 1)))
|
||||
# 31 = page 3, tok 7
|
||||
self.assertTrue(torch.all(k_views[L][3, 7] == float(L + 10)))
|
||||
self.assertTrue(torch.all(v_views[L][3, 7] == -float(L + 10)))
|
||||
|
||||
def test_move_kv_cache_3d_legacy_path_unchanged(self):
|
||||
"""move_kv_cache_native(3-D, page_size=1) must take the legacy
|
||||
else-branch and be byte-identical to today."""
|
||||
class TestMoveKVCacheNative(unittest.TestCase):
|
||||
def test_move_kv_cache_3d_path_unchanged(self):
|
||||
"""The stock per-layer 3-D move must relocate exactly the named token
|
||||
rows, byte-identically — compaction on static pools rides on it."""
|
||||
k = [torch.zeros((32, 2, 4), dtype=torch.float16) for _ in range(2)]
|
||||
v = [torch.zeros((32, 2, 4), dtype=torch.float16) for _ in range(2)]
|
||||
for L in range(2):
|
||||
@@ -292,64 +111,11 @@ class TestMoveKVCacheNative4D(unittest.TestCase):
|
||||
v,
|
||||
tgt_loc=torch.tensor([7], dtype=torch.int64),
|
||||
src_loc=torch.tensor([5], dtype=torch.int64),
|
||||
page_size=1,
|
||||
)
|
||||
for L in range(2):
|
||||
self.assertTrue(torch.all(k[L][7] == float(L + 1)))
|
||||
self.assertTrue(torch.all(v[L][7] == -float(L + 1)))
|
||||
|
||||
|
||||
class TestByteIdentityAtPageSize1(unittest.TestCase):
|
||||
"""Verify that at page_size=1 the new 4-D view describes the SAME
|
||||
physical bytes as the old 3-D view would have. The view
|
||||
semantics differ (4-D vs 3-D shape) but the underlying byte layout is
|
||||
identical — confirmed by manually computing expected byte offsets and
|
||||
matching them against the 4-D view's strides + storage_offset.
|
||||
"""
|
||||
|
||||
def test_byte_addresses_match_envelope(self):
|
||||
spec = _make_mha_spec("full", "up", layer_num=4, head_num=2, head_dim=4)
|
||||
ps = 1
|
||||
# Build pool.
|
||||
total = spec.entry_bytes() * 64 + spec.entry_bytes() * 32
|
||||
pool = UnifiedKVPool(
|
||||
total_bytes=total,
|
||||
sub_pool_specs=[
|
||||
spec,
|
||||
_make_mha_spec("swa", "down", layer_num=2),
|
||||
],
|
||||
device=_DEV,
|
||||
enable_memory_saver=False,
|
||||
page_size=ps,
|
||||
)
|
||||
k_views, v_views = pool.mha_views_for("full")
|
||||
# For each (layer, slot), compute the expected byte address under
|
||||
# the envelope layout and verify the 4-D view's data_ptr +
|
||||
# advanced indexing agrees.
|
||||
max_slots = pool.max_slots("full")
|
||||
itemsize = spec.store_dtype.itemsize
|
||||
base_addr = pool._raw.data_ptr()
|
||||
for L in range(spec.layer_num):
|
||||
for s in range(0, max_slots, max(1, max_slots // 4)):
|
||||
# Envelope: bytes for slot s, layer L's K start at:
|
||||
# s * entry_bytes + L * (k_row + v_row)
|
||||
expected_k_byte_offset = s * spec.entry_bytes() + L * (
|
||||
spec.k_row_bytes() + spec.v_row_bytes()
|
||||
)
|
||||
# 4-D view: k_views[L][page=s, tok=0, head=0, dim=0]
|
||||
# storage_offset of the element [s, 0, 0, 0]:
|
||||
view_offset_elems = (
|
||||
k_views[L].storage_offset()
|
||||
+ s * k_views[L].stride(0)
|
||||
+ 0 * k_views[L].stride(1)
|
||||
+ 0 * k_views[L].stride(2)
|
||||
+ 0 * k_views[L].stride(3)
|
||||
)
|
||||
view_byte_offset = view_offset_elems * itemsize
|
||||
# 4-D view sits over `_raw.view(spec.store_dtype)`, which
|
||||
# has data_ptr == _raw.data_ptr() (same backing storage).
|
||||
self.assertEqual(view_byte_offset, expected_k_byte_offset)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -438,8 +438,7 @@ class TestMultiEndedAllocator(unittest.TestCase):
|
||||
|
||||
def test_translate_kv_loc_dtype_assertion(self):
|
||||
"""REGRESSION: wrong-dtype `out=` (int32 instead of int64) raises
|
||||
AssertionError. Guards against the copy/paste hazard where someone
|
||||
might allocate the full-physical buffer with the SWA int32 pattern."""
|
||||
AssertionError -- `out=` must match the v2p dtype the gather writes."""
|
||||
_, full_alloc, _, full_kv, _ = self._build_pair()
|
||||
v = self._alloc(full_alloc, full_kv, 5)
|
||||
wrong_dtype = torch.empty(v.shape, dtype=torch.int32, device=_DEV)
|
||||
@@ -911,28 +910,28 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase):
|
||||
def test_swa_translate_loc_from_full_to_swa_with_out_writes_inplace(self):
|
||||
"""REGRESSION: `translate_loc_from_full_to_swa(v, out=buf)`
|
||||
must modify `buf` in place AND preserve `buf.data_ptr()`. `out=`
|
||||
buffer MUST be int32 (matches SWA Triton kernel contract)."""
|
||||
buffer is int64 — every id the allocator emits is."""
|
||||
_, allocator, _ = self._build()
|
||||
v = allocator.alloc(4)
|
||||
self.assertIsNotNone(v)
|
||||
buf = torch.empty(v.shape, dtype=torch.int32, device=_DEV)
|
||||
buf = torch.empty(v.shape, dtype=torch.int64, device=_DEV)
|
||||
ptr_before = buf.data_ptr()
|
||||
ret = allocator.translate_loc_from_full_to_swa(v, out=buf)
|
||||
self.assertIs(ret, buf)
|
||||
self.assertEqual(buf.data_ptr(), ptr_before)
|
||||
# Byte-identical to the no-out form:
|
||||
no_out = allocator.translate_loc_from_full_to_swa(v)
|
||||
self.assertEqual(no_out.dtype, torch.int32)
|
||||
self.assertEqual(no_out.dtype, torch.int64)
|
||||
self.assertTrue(bool((buf == no_out).all().item()))
|
||||
|
||||
def test_swa_translate_loc_from_full_to_swa_dtype_assertion(self):
|
||||
"""REGRESSION: wrong-dtype `out=` (int64 instead of int32)
|
||||
raises AssertionError. Guards against accidentally reusing the int64
|
||||
full-physical buffer pattern for the SWA precompute."""
|
||||
"""REGRESSION: wrong-dtype `out=` (int32 instead of int64) raises
|
||||
AssertionError. Guards against reintroducing a narrowed SWA write loc: the
|
||||
allocator emits int64 and consumers narrow at their own buffer."""
|
||||
_, allocator, _ = self._build()
|
||||
v = allocator.alloc(4)
|
||||
self.assertIsNotNone(v)
|
||||
wrong_dtype = torch.empty(v.shape, dtype=torch.int64, device=_DEV)
|
||||
wrong_dtype = torch.empty(v.shape, dtype=torch.int32, device=_DEV)
|
||||
with self.assertRaises(AssertionError):
|
||||
allocator.translate_loc_from_full_to_swa(v, out=wrong_dtype)
|
||||
|
||||
@@ -947,21 +946,48 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase):
|
||||
# Inject a tombstone on the swa side at one of the live virtual ids.
|
||||
v_tomb = int(v[1].item())
|
||||
allocator.swa_attn_allocator.virtual_to_physical[v_tomb] = -1
|
||||
# No-out form: result must be int32 AND every entry >= 0.
|
||||
# No-out form: result must be int64 AND every entry >= 0.
|
||||
out = allocator.translate_loc_from_full_to_swa(v)
|
||||
self.assertEqual(out.dtype, torch.int32)
|
||||
self.assertEqual(out.dtype, torch.int64)
|
||||
self.assertTrue(
|
||||
bool((out >= 0).all().item()),
|
||||
"translate_loc_from_full_to_swa must clamp tombstoned to >=0",
|
||||
)
|
||||
self.assertEqual(int(out[1].item()), 0)
|
||||
# out= form (int32 buffer) must also clamp.
|
||||
buf = torch.empty(v.shape, dtype=torch.int32, device=_DEV)
|
||||
# out= form must also clamp.
|
||||
buf = torch.empty(v.shape, dtype=torch.int64, device=_DEV)
|
||||
ret = allocator.translate_loc_from_full_to_swa(v, out=buf)
|
||||
self.assertIs(ret, buf)
|
||||
self.assertTrue(bool((buf >= 0).all().item()))
|
||||
self.assertEqual(int(buf[1].item()), 0)
|
||||
|
||||
def test_swa_slot_zero_sink_invariant_survives_churn(self):
|
||||
"""PINNED INVARIANT (swa side of the physical-loc contract): BOTH maps
|
||||
send virtual 0 to physical 0 — `translate_kv_loc(zeros) == zeros` AND
|
||||
`translate_loc_from_full_to_swa(zeros) == zeros` — after init and
|
||||
after alloc/free/free_swa churn. Cuda-graph capture replaced the
|
||||
capture-time translate with zero-fill/copy of the zero-filled static
|
||||
buffers; that is only equivalent while slot 0 stays the sink in both
|
||||
sub-pools."""
|
||||
_, allocator, kvcache = self._build()
|
||||
zeros64 = torch.zeros(4, dtype=torch.int64)
|
||||
|
||||
def check():
|
||||
self.assertTrue(torch.equal(allocator.translate_kv_loc(zeros64), zeros64))
|
||||
self.assertTrue(
|
||||
torch.equal(allocator.translate_loc_from_full_to_swa(zeros64), zeros64)
|
||||
)
|
||||
|
||||
check()
|
||||
a = self._alloc(allocator, kvcache, 5)
|
||||
b = self._alloc(allocator, kvcache, 5)
|
||||
allocator.free_swa(a) # tombstone swa side only
|
||||
self._free(allocator, kvcache, b) # full free (compaction on both)
|
||||
self._free(allocator, kvcache, a)
|
||||
c = self._alloc(allocator, kvcache, 3)
|
||||
self._free(allocator, kvcache, c)
|
||||
check()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# page_size > 1 — paged unit tests
|
||||
@@ -1912,13 +1938,17 @@ class TestPagedMultiEndedAllocator(unittest.TestCase):
|
||||
"v2p_page[virt_pages] * page_size + offsets.",
|
||||
)
|
||||
|
||||
# And the composite allocator's translate method must produce the
|
||||
# same token-granular result (same page math).
|
||||
# The composite emits KERNEL-FACING ids, not the physical token ids
|
||||
# this helper returns; they coincide only at multiplier 1, which no
|
||||
# sub-pool uses.
|
||||
swa_mult = allocator.swa_kernel_page_multiplier
|
||||
self.assertEqual(swa_mult, 2 * swa_spec.layer_num)
|
||||
composite_out = allocator.translate_loc_from_full_to_swa(v_tokens)
|
||||
expected_dense = swa_phys_pages_direct * (PS * swa_mult) + offsets_in
|
||||
self.assertTrue(
|
||||
bool((swa_phys.long() == composite_out.long()).all().item()),
|
||||
"REGRESSION: the UnifiedSWAKVPool helper and the composite "
|
||||
"allocator's translate_loc_from_full_to_swa must agree.",
|
||||
bool((composite_out.long() == expected_dense.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).",
|
||||
)
|
||||
|
||||
|
||||
@@ -2565,5 +2595,131 @@ class TestO3FusedAllocBind(unittest.TestCase):
|
||||
self.assertEqual(int(sa.physical_to_virtual[p].item()), v)
|
||||
|
||||
|
||||
class TestSWACompositeDenseSurface(unittest.TestCase):
|
||||
"""The SWA composite's dense (kernel-facing) id surface.
|
||||
|
||||
Presence of `translate_kv_loc_for_kernel` / `full_v2p_page_table` is what flips
|
||||
the attention backends' kernel-facing-first probes, and the `page_stride` scale in
|
||||
`translate_loc_from_full_to_swa` is what carries the swa kernel-facing space.
|
||||
Everything must collapse
|
||||
byte-identically at multiplier 1 — the strided arm every existing SWA model
|
||||
runs — and follow `kernel_id(t) = v2p[t//ps]*(ps*mult) + t%ps` otherwise.
|
||||
"""
|
||||
|
||||
PS = 4
|
||||
FULL_L = 4
|
||||
SWA_L = 2
|
||||
|
||||
def _build(self):
|
||||
full_spec = MHASubPoolSpec(
|
||||
name="full",
|
||||
layer_num=self.FULL_L,
|
||||
head_num=2,
|
||||
head_dim=4,
|
||||
store_dtype=torch.float16,
|
||||
grow_direction="up",
|
||||
)
|
||||
swa_spec = MHASubPoolSpec(
|
||||
name="swa",
|
||||
layer_num=self.SWA_L,
|
||||
head_num=2,
|
||||
head_dim=4,
|
||||
store_dtype=torch.float16,
|
||||
grow_direction="down",
|
||||
)
|
||||
n_full, n_swa = 64, 32 # tokens = 16 / 8 pages at PS=4
|
||||
total = n_full * full_spec.entry_bytes() + n_swa * swa_spec.entry_bytes()
|
||||
pool = UnifiedKVPool(
|
||||
total_bytes=total,
|
||||
sub_pool_specs=[full_spec, swa_spec],
|
||||
device=_DEV,
|
||||
enable_memory_saver=False,
|
||||
page_size=self.PS,
|
||||
)
|
||||
kvcache = _FakeUnifiedSWAKVPool(pool)
|
||||
return UnifiedSWATokenToKVPoolAllocator(
|
||||
unified_buffer=pool,
|
||||
kvcache=kvcache,
|
||||
device=_DEV,
|
||||
full_max_total_num_tokens=n_full,
|
||||
swa_max_total_num_tokens=n_swa,
|
||||
page_size=self.PS,
|
||||
need_sort=False,
|
||||
forward_stream=None,
|
||||
)
|
||||
|
||||
def test_multipliers_come_from_the_specs(self):
|
||||
"""Both sides scale by their OWN sub-pool's block count, and the
|
||||
composite exposes the raw v2p tables unwrapped. Nothing injects the
|
||||
scale: a spec whose views are dense cannot be paired with a
|
||||
physical-id multiplier, which is the state that writes physical ids
|
||||
into view rows."""
|
||||
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.assertIs(a.full_v2p_page_table, a.full_attn_allocator.virtual_to_physical)
|
||||
self.assertIs(a.swa_v2p_page_table, a.swa_attn_allocator.virtual_to_physical)
|
||||
|
||||
def test_full_dense_translate_matches_formula(self):
|
||||
mult = 2 * self.FULL_L
|
||||
a = self._build()
|
||||
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
|
||||
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.
|
||||
phys = v2p[v // self.PS] * self.PS + v % self.PS
|
||||
self.assertTrue(torch.equal(a.translate_kv_loc(v), phys))
|
||||
|
||||
def test_dense_translate_accepts_an_int32_page_table(self):
|
||||
"""REGRESSION: fa3 translates its own page table, which is int32 and
|
||||
2-D. A gather that requires an int64 index (`torch.take`) crashes the
|
||||
scheduler there while every int64 caller stays green. Both page sizes:
|
||||
at ps == 1 the index IS the caller's tensor, at ps > 1 it is derived."""
|
||||
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
|
||||
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)
|
||||
self.assertTrue(torch.equal(got.reshape(-1), expected))
|
||||
# `out=` takes the same int32 index; the buffer stays int64.
|
||||
dst = torch.empty(page_table.shape, dtype=torch.int64, device=_DEV)
|
||||
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
|
||||
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
|
||||
self.assertTrue(torch.equal(a.translate_loc_from_full_to_swa(v), expected))
|
||||
|
||||
def test_swa_dense_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 = self._build()
|
||||
v = a.alloc(2 * self.PS)
|
||||
self.assertIsNotNone(v)
|
||||
tomb_page = int(v[0].item()) // self.PS
|
||||
a.swa_attn_allocator.virtual_to_physical[tomb_page] = -1
|
||||
got = a.translate_loc_from_full_to_swa(v)
|
||||
self.assertTrue(bool((got >= 0).all().item()))
|
||||
in_tomb = v // self.PS == tomb_page
|
||||
self.assertTrue(bool((got[in_tomb] == 0).all().item()))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -1,12 +1,9 @@
|
||||
"""CPU correctness tests for the page-major layer-major envelope layout.
|
||||
"""CPU correctness tests for the page-major envelope Mamba state views.
|
||||
|
||||
Covers the standalone view builders (no allocator / shared pool):
|
||||
|
||||
- ``build_page_major_mha_views``: 4-D K/V views with correct addressing at
|
||||
page_size 1 (token-granularity envelope) and > 1 (layer-major within a page),
|
||||
and no aliasing across layers / slots.
|
||||
- ``build_page_major_mamba_views``: conv / temporal state views.
|
||||
- ``move_kv_cache_native`` 4-D branch: relocating token rows preserves data.
|
||||
Covers the standalone ``build_page_major_mamba_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.
|
||||
|
||||
@@ -23,107 +20,10 @@ import torch
|
||||
|
||||
from sglang.srt.mem_cache.layout.page_major import (
|
||||
build_page_major_mamba_views,
|
||||
build_page_major_mha_views,
|
||||
mamba_entry_bytes,
|
||||
mha_entry_bytes,
|
||||
)
|
||||
from sglang.srt.mem_cache.memory_pool import move_kv_cache_native
|
||||
|
||||
_DEV = "cpu"
|
||||
_DT = torch.float32
|
||||
|
||||
|
||||
def _make_mha_views(layer_num, head_num, head_dim, v_head_dim, page_size, num_pages):
|
||||
entry = mha_entry_bytes(
|
||||
layer_num=layer_num,
|
||||
head_num=head_num,
|
||||
head_dim=head_dim,
|
||||
v_head_dim=v_head_dim,
|
||||
itemsize=_DT.itemsize,
|
||||
)
|
||||
raw = torch.zeros(num_pages * page_size * entry, dtype=torch.uint8, device=_DEV)
|
||||
k, v = build_page_major_mha_views(
|
||||
raw,
|
||||
layer_num=layer_num,
|
||||
head_num=head_num,
|
||||
head_dim=head_dim,
|
||||
v_head_dim=v_head_dim,
|
||||
store_dtype=_DT,
|
||||
page_size=page_size,
|
||||
num_pages=num_pages,
|
||||
)
|
||||
return raw, k, v
|
||||
|
||||
|
||||
class TestPageMajorMHAViews(unittest.TestCase):
|
||||
def test_view_shapes(self):
|
||||
_, k, v = _make_mha_views(3, 2, 4, 4, page_size=2, num_pages=4)
|
||||
self.assertEqual(len(k), 3)
|
||||
for t in k:
|
||||
self.assertEqual(tuple(t.shape), (4, 2, 2, 4))
|
||||
for t in v:
|
||||
self.assertEqual(tuple(t.shape), (4, 2, 2, 4))
|
||||
|
||||
def test_no_aliasing_ps1(self):
|
||||
# Every (layer, slot) cell must be independently addressable.
|
||||
layer_num, slots = 3, 5
|
||||
_, k, v = _make_mha_views(layer_num, 2, 4, 4, page_size=1, num_pages=slots)
|
||||
for L in range(layer_num):
|
||||
for s in range(slots):
|
||||
k[L][s, 0] = float(100 + L * 10 + s)
|
||||
v[L][s, 0] = float(200 + L * 10 + s)
|
||||
for L in range(layer_num):
|
||||
for s in range(slots):
|
||||
self.assertTrue(torch.all(k[L][s, 0] == float(100 + L * 10 + s)))
|
||||
self.assertTrue(torch.all(v[L][s, 0] == float(200 + L * 10 + s)))
|
||||
|
||||
def test_page_slot_addressing_ps_gt1(self):
|
||||
# token id t -> page t // ps, slot t % ps; no aliasing across tokens.
|
||||
ps, pages = 2, 4
|
||||
total = ps * pages
|
||||
_, k, _ = _make_mha_views(2, 1, 2, 2, page_size=ps, num_pages=pages)
|
||||
for L in range(2):
|
||||
for t in range(total):
|
||||
k[L][t // ps, t % ps, 0] = float(1000 + L * 100 + t)
|
||||
for L in range(2):
|
||||
for t in range(total):
|
||||
self.assertEqual(
|
||||
float(k[L][t // ps, t % ps, 0, 0].item()), 1000 + L * 100 + t
|
||||
)
|
||||
|
||||
def test_asymmetric_v_head_dim(self):
|
||||
_, k, v = _make_mha_views(2, 2, 6, 4, page_size=1, num_pages=3)
|
||||
self.assertEqual(tuple(k[0].shape), (3, 1, 2, 6))
|
||||
self.assertEqual(tuple(v[0].shape), (3, 1, 2, 4))
|
||||
|
||||
|
||||
class TestPageMajorMove(unittest.TestCase):
|
||||
def test_move_ps1(self):
|
||||
slots = 6
|
||||
_, k, v = _make_mha_views(2, 1, 4, 4, page_size=1, num_pages=slots)
|
||||
for L in range(2):
|
||||
for s in range(slots):
|
||||
k[L][s, 0] = float(s + 1)
|
||||
v[L][s, 0] = float(-(s + 1))
|
||||
tgt = torch.tensor([0, 1], dtype=torch.int64)
|
||||
src = torch.tensor([4, 5], dtype=torch.int64)
|
||||
move_kv_cache_native(k, v, tgt, src, page_size=1)
|
||||
for L in range(2):
|
||||
self.assertTrue(torch.all(k[L][0, 0] == 5.0))
|
||||
self.assertTrue(torch.all(k[L][1, 0] == 6.0))
|
||||
self.assertTrue(torch.all(v[L][0, 0] == -5.0))
|
||||
|
||||
def test_move_ps_gt1(self):
|
||||
ps, pages = 2, 4
|
||||
total = ps * pages
|
||||
_, k, v = _make_mha_views(1, 1, 2, 2, page_size=ps, num_pages=pages)
|
||||
for t in range(total):
|
||||
k[0][t // ps, t % ps, 0] = float(t + 1)
|
||||
tgt = torch.tensor([0, 3], dtype=torch.int64) # page0 slot0, page1 slot1
|
||||
src = torch.tensor([6, 7], dtype=torch.int64) # page3 slot0, page3 slot1
|
||||
move_kv_cache_native(k, v, tgt, src, page_size=ps)
|
||||
self.assertEqual(float(k[0][0, 0, 0, 0].item()), 7.0)
|
||||
self.assertEqual(float(k[0][1, 1, 0, 0].item()), 8.0)
|
||||
|
||||
|
||||
class TestMambaEnvelopeViews(unittest.TestCase):
|
||||
|
||||
@@ -18,7 +18,7 @@ import unittest
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.layout.page_major import (
|
||||
build_dense_mla_views,
|
||||
build_mla_views,
|
||||
build_page_major_mamba_views,
|
||||
mamba_entry_bytes,
|
||||
mla_entry_bytes,
|
||||
@@ -31,7 +31,7 @@ register_cpu_ci(est_time=60, suite="base-a-test-cpu")
|
||||
|
||||
class TestMLAEnvelopeTransferAddressing(CustomTestCase):
|
||||
def test_page_envelope_matches_dense_views(self):
|
||||
"""Every (page, layer, slot) row written through the dense MLA views
|
||||
"""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."""
|
||||
layer_num, page_size, kv_dim, num_pages = 3, 4, 8, 6
|
||||
@@ -49,7 +49,7 @@ class TestMLAEnvelopeTransferAddressing(CustomTestCase):
|
||||
)
|
||||
# +1 page envelope of tail pad, as UnifiedKVPool allocates for MLA.
|
||||
raw = torch.zeros((num_pages + 1) * page_bytes, dtype=torch.uint8)
|
||||
views = build_dense_mla_views(
|
||||
views = build_mla_views(
|
||||
raw,
|
||||
layer_num=layer_num,
|
||||
kv_cache_dim=kv_dim,
|
||||
|
||||
@@ -1,412 +0,0 @@
|
||||
"""Parity tests for the `store_cache_4d` Triton kernel.
|
||||
|
||||
The kernel writes K/V into the 4-D page-major envelope view. These tests prove
|
||||
it produces byte-identical output to the legacy advanced-indexing path on
|
||||
representative fixtures:
|
||||
|
||||
- ``page_size = 1`` (envelope-degenerate, the critical compatibility case)
|
||||
- ``page_size > 1`` (layer-major within page)
|
||||
- both int32 and int64 ``loc`` dtypes
|
||||
- bf16 and fp8_e5m2 view dtypes
|
||||
- asymmetric ``head_dim != v_head_dim``
|
||||
- empty ``loc`` (no-op)
|
||||
|
||||
Skipped on CPU — Triton requires a GPU.
|
||||
|
||||
python -m pytest test/registered/unit/mem_cache/test_store_cache_4d.py -v
|
||||
"""
|
||||
|
||||
import importlib.util
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
_HAS_CUDA = torch.cuda.is_available()
|
||||
# The set_kv_buffer integration test needs UnifiedMHATokenToKVPool, which only
|
||||
# exists once the shared-KV-pool feature lands; skip it where absent.
|
||||
_HAS_SHARED_POOL = (
|
||||
importlib.util.find_spec("sglang.srt.mem_cache.unified_memory_pool") is not None
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
def _legacy_advanced_indexing_write(
|
||||
k_view: torch.Tensor,
|
||||
v_view: torch.Tensor,
|
||||
cache_k: torch.Tensor,
|
||||
cache_v: torch.Tensor,
|
||||
loc: torch.Tensor,
|
||||
page_size: int,
|
||||
) -> None:
|
||||
"""Reference implementation: the legacy bypass-super() advanced-indexing
|
||||
path that the Triton kernel replaces. Used as the byte-identity oracle
|
||||
for the parity tests below.
|
||||
"""
|
||||
if page_size == 1:
|
||||
k_view[loc, 0] = cache_k
|
||||
v_view[loc, 0] = cache_v
|
||||
else:
|
||||
page_id = loc // page_size
|
||||
tok_in_p = loc % page_size
|
||||
k_view[page_id, tok_in_p] = cache_k
|
||||
v_view[page_id, tok_in_p] = cache_v
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_CUDA, "Triton kernels require CUDA")
|
||||
class TestStoreCache4D(unittest.TestCase):
|
||||
"""Byte-identity parity vs the legacy advanced-indexing write path."""
|
||||
|
||||
def _make_view_and_cache(
|
||||
self,
|
||||
num_pages: int,
|
||||
page_size: int,
|
||||
head_num: int,
|
||||
head_dim: int,
|
||||
v_head_dim: int,
|
||||
N: int,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
loc_dtype: torch.dtype = torch.int64,
|
||||
seed: int = 0xC0FFEE,
|
||||
):
|
||||
torch.manual_seed(seed)
|
||||
# The unified memory pool's views are 4-D `(num_pages, page_size, head_num,
|
||||
# head_dim)` with the trailing two dims contiguous. We allocate two
|
||||
# independent contiguous buffers (one for the kernel-under-test,
|
||||
# one as the legacy-path target) so we can compare them.
|
||||
k_view = torch.zeros(
|
||||
(num_pages, page_size, head_num, head_dim),
|
||||
dtype=dtype,
|
||||
device="cuda",
|
||||
)
|
||||
v_view = torch.zeros(
|
||||
(num_pages, page_size, head_num, v_head_dim),
|
||||
dtype=dtype,
|
||||
device="cuda",
|
||||
)
|
||||
cache_k = torch.randn(
|
||||
(N, head_num, head_dim), dtype=torch.float32, device="cuda"
|
||||
).to(dtype)
|
||||
cache_v = torch.randn(
|
||||
(N, head_num, v_head_dim), dtype=torch.float32, device="cuda"
|
||||
).to(dtype)
|
||||
# Valid loc values in [0, num_pages * page_size); generate without
|
||||
# duplicates so the comparison is unambiguous (advanced-indexing
|
||||
# with duplicates is order-undefined for both paths).
|
||||
total_slots = num_pages * page_size
|
||||
assert N <= total_slots
|
||||
loc = torch.randperm(total_slots, device="cuda")[:N].to(loc_dtype)
|
||||
return k_view, v_view, cache_k, cache_v, loc
|
||||
|
||||
def _check_parity(
|
||||
self,
|
||||
num_pages: int,
|
||||
page_size: int,
|
||||
head_num: int,
|
||||
head_dim: int,
|
||||
v_head_dim: int,
|
||||
N: int,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
loc_dtype: torch.dtype = torch.int64,
|
||||
):
|
||||
from sglang.kernels.ops.kvcache.cache_move import store_cache_4d
|
||||
|
||||
# Two independent target buffers — one for the kernel, one for the
|
||||
# legacy reference path.
|
||||
k_kernel, v_kernel, cache_k, cache_v, loc = self._make_view_and_cache(
|
||||
num_pages,
|
||||
page_size,
|
||||
head_num,
|
||||
head_dim,
|
||||
v_head_dim,
|
||||
N,
|
||||
dtype=dtype,
|
||||
loc_dtype=loc_dtype,
|
||||
)
|
||||
k_legacy = k_kernel.clone()
|
||||
v_legacy = v_kernel.clone()
|
||||
|
||||
# Kernel-under-test
|
||||
store_cache_4d(k_kernel, v_kernel, cache_k, cache_v, loc, page_size)
|
||||
# Legacy reference
|
||||
_legacy_advanced_indexing_write(
|
||||
k_legacy, v_legacy, cache_k, cache_v, loc, page_size
|
||||
)
|
||||
|
||||
# Byte-identical comparison — the kernel must reproduce the
|
||||
# advanced-indexing path bit-for-bit, NOT just numerically close.
|
||||
# For fp8 dtypes, torch.equal works on the integer bit pattern.
|
||||
self.assertTrue(
|
||||
torch.equal(k_kernel, k_legacy),
|
||||
f"K view mismatch: ps={page_size}, dtype={dtype}, "
|
||||
f"loc_dtype={loc_dtype}, N={N}",
|
||||
)
|
||||
self.assertTrue(
|
||||
torch.equal(v_kernel, v_legacy),
|
||||
f"V view mismatch: ps={page_size}, dtype={dtype}, "
|
||||
f"loc_dtype={loc_dtype}, N={N}",
|
||||
)
|
||||
|
||||
# ---- Test 1: ps=1 envelope-degenerate (the critical compat case) ----
|
||||
|
||||
def test_store_cache_4d_ps1_byte_identical(self):
|
||||
"""At page_size=1 the kernel constexpr-folds to the slot-major
|
||||
envelope view. Output must be byte-identical to advanced indexing.
|
||||
This protects against byte-layout regression."""
|
||||
self._check_parity(
|
||||
num_pages=64,
|
||||
page_size=1,
|
||||
head_num=4,
|
||||
head_dim=128,
|
||||
v_head_dim=128,
|
||||
N=16,
|
||||
)
|
||||
|
||||
# ---- Test 2: ps>1 layer-major within page ----
|
||||
|
||||
def test_store_cache_4d_ps_gt1_byte_identical(self):
|
||||
"""At page_size > 1 the kernel splits loc into (page_id, tok_in_p)
|
||||
and writes via the 4-D stride. Output must match the equivalent
|
||||
advanced-indexing write."""
|
||||
self._check_parity(
|
||||
num_pages=8,
|
||||
page_size=64,
|
||||
head_num=4,
|
||||
head_dim=128,
|
||||
v_head_dim=128,
|
||||
N=128,
|
||||
)
|
||||
|
||||
# ---- Test 3: int32 loc dtype ----
|
||||
|
||||
def test_store_cache_4d_int32_loc(self):
|
||||
"""The SWA-side path passes int32 loc (matches the SWA Triton
|
||||
kernel contract). PyTorch advanced indexing tolerates either
|
||||
int32 or int64; the kernel must too."""
|
||||
self._check_parity(
|
||||
num_pages=32,
|
||||
page_size=1,
|
||||
head_num=4,
|
||||
head_dim=64,
|
||||
v_head_dim=64,
|
||||
N=10,
|
||||
loc_dtype=torch.int32,
|
||||
)
|
||||
|
||||
# ---- Test 4: int64 loc dtype (already exercised, explicit) ----
|
||||
|
||||
# ---- Test 5: bf16 dtype (the production case) ----
|
||||
|
||||
# ---- Test 6: fp8_e5m2 dtype ----
|
||||
|
||||
def test_store_cache_4d_dtype_fp8_e5m2(self):
|
||||
"""fp8_e5m2 is used for KV-cache quantization. Caller is responsible
|
||||
for the cast; the kernel sees same-dtype source and destination."""
|
||||
self._check_parity(
|
||||
num_pages=16,
|
||||
page_size=64,
|
||||
head_num=4,
|
||||
head_dim=128,
|
||||
v_head_dim=128,
|
||||
N=64,
|
||||
dtype=torch.float8_e5m2,
|
||||
)
|
||||
|
||||
# ---- Test 7: empty loc (no-op) ----
|
||||
|
||||
def test_store_cache_4d_empty_loc(self):
|
||||
"""N=0 must be a no-op: no kernel launch, no exception, no buffer
|
||||
mutation."""
|
||||
from sglang.kernels.ops.kvcache.cache_move import store_cache_4d
|
||||
|
||||
k_view = torch.zeros((8, 4, 4, 64), dtype=torch.bfloat16, device="cuda")
|
||||
v_view = torch.zeros((8, 4, 4, 64), dtype=torch.bfloat16, device="cuda")
|
||||
k_before = k_view.clone()
|
||||
v_before = v_view.clone()
|
||||
cache_k = torch.empty((0, 4, 64), dtype=torch.bfloat16, device="cuda")
|
||||
cache_v = torch.empty((0, 4, 64), dtype=torch.bfloat16, device="cuda")
|
||||
loc = torch.empty((0,), dtype=torch.int64, device="cuda")
|
||||
|
||||
store_cache_4d(k_view, v_view, cache_k, cache_v, loc, page_size=4)
|
||||
|
||||
# Buffers must be unchanged.
|
||||
self.assertTrue(torch.equal(k_view, k_before))
|
||||
self.assertTrue(torch.equal(v_view, v_before))
|
||||
|
||||
# ---- Test 8: head_dim != v_head_dim (asymmetric, e.g. MLA-style) ----
|
||||
|
||||
def test_store_cache_4d_v_head_dim_differs(self):
|
||||
"""When v_head_dim != head_dim, the kernel's K and V branches use
|
||||
different per-token strides. Exercises the stride_k_tok ≠
|
||||
stride_v_tok branch."""
|
||||
self._check_parity(
|
||||
num_pages=8,
|
||||
page_size=16,
|
||||
head_num=2,
|
||||
head_dim=128,
|
||||
v_head_dim=64,
|
||||
N=16,
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_CUDA, "Triton kernels require CUDA")
|
||||
class TestStoreCache4DAssertions(unittest.TestCase):
|
||||
"""The wrapper's contract assertions must fire on bad inputs."""
|
||||
|
||||
def test_rejects_non_contiguous_view_trailing_dim(self):
|
||||
"""Wrapper requires `stride[-1] == 1` and `stride[-2] == head_dim`
|
||||
(the trailing two dims must be contiguous). A permutation that
|
||||
breaks this should trigger AssertionError."""
|
||||
from sglang.kernels.ops.kvcache.cache_move import store_cache_4d
|
||||
|
||||
# Build a 4-D view, then permute the last two dims → trailing
|
||||
# contiguity violated.
|
||||
k_view = torch.zeros(
|
||||
(4, 4, 4, 64), dtype=torch.bfloat16, device="cuda"
|
||||
).permute(
|
||||
0, 1, 3, 2
|
||||
) # now shape (4, 4, 64, 4); strides broken
|
||||
v_view = torch.zeros((4, 4, 4, 64), dtype=torch.bfloat16, device="cuda")
|
||||
cache_k = torch.zeros((2, 4, 64), dtype=torch.bfloat16, device="cuda")
|
||||
cache_v = torch.zeros((2, 4, 64), dtype=torch.bfloat16, device="cuda")
|
||||
loc = torch.arange(2, dtype=torch.int64, device="cuda")
|
||||
with self.assertRaises(AssertionError):
|
||||
store_cache_4d(k_view, v_view, cache_k, cache_v, loc, page_size=4)
|
||||
|
||||
def test_rejects_dtype_mismatch(self):
|
||||
"""All four tensors must share a dtype; the caller is responsible
|
||||
for any cast before the call."""
|
||||
from sglang.kernels.ops.kvcache.cache_move import store_cache_4d
|
||||
|
||||
k_view = torch.zeros((4, 4, 4, 64), dtype=torch.bfloat16, device="cuda")
|
||||
v_view = torch.zeros((4, 4, 4, 64), dtype=torch.bfloat16, device="cuda")
|
||||
cache_k = torch.zeros((2, 4, 64), dtype=torch.float16, device="cuda")
|
||||
cache_v = torch.zeros((2, 4, 64), dtype=torch.bfloat16, device="cuda")
|
||||
loc = torch.arange(2, dtype=torch.int64, device="cuda")
|
||||
with self.assertRaises(AssertionError):
|
||||
store_cache_4d(k_view, v_view, cache_k, cache_v, loc, page_size=4)
|
||||
|
||||
|
||||
@unittest.skipUnless(
|
||||
_HAS_CUDA and _HAS_SHARED_POOL,
|
||||
"Triton kernels require CUDA; UnifiedMHATokenToKVPool required",
|
||||
)
|
||||
class TestStoreCache4DThroughSetKVBuffer(unittest.TestCase):
|
||||
"""Integration parity test — exercises the kernel through the FULL
|
||||
``UnifiedMHATokenToKVPool.set_kv_buffer`` path (the direct PHYSICAL write +
|
||||
the dtype cast; the pool no longer translates). Confirms it produces
|
||||
bit-identical output to a PyTorch advanced-indexing reference write.
|
||||
"""
|
||||
|
||||
def _build_pool(self, page_size: int):
|
||||
"""Build a small UnifiedMHATokenToKVPool. The pool writes PHYSICAL locs
|
||||
directly (no allocator / v2p translate), so `set_kv_buffer` receives the
|
||||
already-physical write location."""
|
||||
import torch as _t
|
||||
|
||||
from sglang.srt.mem_cache.unified_memory_pool import (
|
||||
MHASubPoolSpec,
|
||||
UnifiedKVPool,
|
||||
UnifiedMHATokenToKVPool,
|
||||
)
|
||||
|
||||
spec = MHASubPoolSpec(
|
||||
name="full",
|
||||
layer_num=2,
|
||||
head_num=4,
|
||||
head_dim=64,
|
||||
store_dtype=_t.bfloat16,
|
||||
grow_direction="up",
|
||||
)
|
||||
total = spec.entry_bytes() * 64
|
||||
# Use a peer to satisfy the two-sub-pool contract.
|
||||
peer = MHASubPoolSpec(
|
||||
name="swa",
|
||||
layer_num=1,
|
||||
head_num=4,
|
||||
head_dim=64,
|
||||
store_dtype=_t.bfloat16,
|
||||
grow_direction="down",
|
||||
)
|
||||
pool = UnifiedKVPool(
|
||||
total_bytes=total + peer.entry_bytes() * 16,
|
||||
sub_pool_specs=[spec, peer],
|
||||
device="cuda",
|
||||
enable_memory_saver=False,
|
||||
page_size=page_size,
|
||||
)
|
||||
kv_pool = UnifiedMHATokenToKVPool(
|
||||
unified_buffer=pool,
|
||||
sub_pool_name="full",
|
||||
page_size=page_size,
|
||||
start_layer=0,
|
||||
end_layer=2,
|
||||
enable_alt_stream=False,
|
||||
)
|
||||
|
||||
return kv_pool
|
||||
|
||||
def _run_set_kv_buffer_and_compare(self, page_size: int):
|
||||
import torch as _t
|
||||
|
||||
kv_pool = self._build_pool(page_size)
|
||||
|
||||
# A fake `layer` object with the minimum interface
|
||||
# `set_kv_buffer` reads: `.layer_id`.
|
||||
class _FakeLayer:
|
||||
layer_id = 0
|
||||
|
||||
layer = _FakeLayer()
|
||||
head_num, head_dim = 4, 64
|
||||
N = 16
|
||||
# Generate valid loc in range [0, num_pages * page_size).
|
||||
num_pages = kv_pool.k_buffer[0].shape[0]
|
||||
total = num_pages * page_size
|
||||
assert N <= total
|
||||
loc = _t.randperm(total, device="cuda")[:N].to(_t.int64)
|
||||
cache_k = _t.randn((N, head_num, head_dim), dtype=_t.bfloat16, device="cuda")
|
||||
cache_v = _t.randn((N, head_num, head_dim), dtype=_t.bfloat16, device="cuda")
|
||||
|
||||
# Production path: the Triton `store_cache_4d` kernel via set_kv_buffer.
|
||||
kv_pool.set_kv_buffer(layer, loc, cache_k.clone(), cache_v.clone())
|
||||
k_kernel = kv_pool.k_buffer[0].clone()
|
||||
v_kernel = kv_pool.v_buffer[0].clone()
|
||||
|
||||
# Reference: PyTorch advanced-indexing into a fresh view at the same
|
||||
# (physical) loc, with no dtype cast (store_dtype == dtype) — the exact
|
||||
# write the kernel performs.
|
||||
kv_pool.k_buffer[0].zero_()
|
||||
kv_pool.v_buffer[0].zero_()
|
||||
k_view = kv_pool.k_buffer[0]
|
||||
v_view = kv_pool.v_buffer[0]
|
||||
if page_size == 1:
|
||||
k_view[loc, 0] = cache_k
|
||||
v_view[loc, 0] = cache_v
|
||||
else:
|
||||
page_id = loc // page_size
|
||||
tok_in_p = loc % page_size
|
||||
k_view[page_id, tok_in_p] = cache_k
|
||||
v_view[page_id, tok_in_p] = cache_v
|
||||
k_ref = kv_pool.k_buffer[0].clone()
|
||||
v_ref = kv_pool.v_buffer[0].clone()
|
||||
|
||||
self.assertTrue(
|
||||
_t.equal(k_kernel, k_ref),
|
||||
f"K view mismatch through set_kv_buffer at ps={page_size}",
|
||||
)
|
||||
self.assertTrue(
|
||||
_t.equal(v_kernel, v_ref),
|
||||
f"V view mismatch through set_kv_buffer at ps={page_size}",
|
||||
)
|
||||
|
||||
def test_integration_ps1(self):
|
||||
self._run_set_kv_buffer_and_compare(page_size=1)
|
||||
|
||||
def test_integration_ps64(self):
|
||||
self._run_set_kv_buffer_and_compare(page_size=64)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,190 +0,0 @@
|
||||
"""Triton-kernel parity test for the page-aware decode / extend kernels.
|
||||
|
||||
Verifies that the modified decode / extend Triton kernels produce
|
||||
bit-identical output when called against:
|
||||
|
||||
(a) the legacy 3-D ``[N, head, dim]`` KV view (PAGE_SIZE=1 default),
|
||||
(b) the new 4-D ``[num_pages, page_size, head, dim]`` view with
|
||||
``page_size=1`` (degenerate envelope — same physical bytes as (a)),
|
||||
(c) the new 4-D view with ``page_size>1`` (layer-major), using the
|
||||
same logical KV data but routed via page-aware address math.
|
||||
|
||||
Output for (a) vs (b) must be bit-identical at PAGE_SIZE=1 (the kernel
|
||||
specializes to the legacy branch). Output for (c) must match a hand-
|
||||
computed reference SDPA result (same logical attention; different byte
|
||||
layout).
|
||||
|
||||
Skipped on CPU — Triton requires a GPU.
|
||||
|
||||
python -m pytest test/registered/unit/mem_cache/test_triton_kernel_layout.py -v
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
|
||||
_HAS_CUDA = torch.cuda.is_available()
|
||||
|
||||
register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small")
|
||||
register_amd_ci(est_time=5, stage="stage-b", runner_config="1-gpu-small-amd")
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_CUDA, "Triton kernels require CUDA")
|
||||
class TestTritonKernelLayoutParity(unittest.TestCase):
|
||||
"""Decode + extend kernel parity across (3-D, 4-D ps=1, 4-D ps>1)."""
|
||||
|
||||
def _setup_decode_inputs(
|
||||
self, bs=2, head_num=2, head_dim=8, num_slots=64, dtype=torch.float16
|
||||
):
|
||||
torch.manual_seed(0xC0FFEE)
|
||||
# Logical KV: shape [num_slots, head_num, head_dim]
|
||||
logical_kv_k = torch.randn(
|
||||
num_slots, head_num, head_dim, dtype=dtype, device="cuda"
|
||||
)
|
||||
logical_kv_v = torch.randn(
|
||||
num_slots, head_num, head_dim, dtype=dtype, device="cuda"
|
||||
)
|
||||
q = torch.randn(bs, head_num, head_dim, dtype=dtype, device="cuda")
|
||||
# All requests use the first `seq_len` slots.
|
||||
seq_len = 16
|
||||
kv_indices_per_req = torch.arange(seq_len, dtype=torch.int64, device="cuda")
|
||||
kv_indices = kv_indices_per_req.repeat(bs) # [bs * seq_len]
|
||||
kv_indptr = torch.tensor(
|
||||
[i * seq_len for i in range(bs + 1)], dtype=torch.int32, device="cuda"
|
||||
)
|
||||
return q, logical_kv_k, logical_kv_v, kv_indptr, kv_indices, seq_len
|
||||
|
||||
def _run_decode(self, q, k_buf, v_buf, kv_indptr, kv_indices, page_size):
|
||||
from sglang.kernels.ops.attention.decode_attention import (
|
||||
decode_attention_fwd,
|
||||
)
|
||||
|
||||
bs, head_num, head_dim = q.shape
|
||||
max_kv_splits = 4
|
||||
attn_logits = torch.empty(
|
||||
(bs, head_num, max_kv_splits, head_dim),
|
||||
dtype=torch.float32,
|
||||
device="cuda",
|
||||
)
|
||||
attn_lse = torch.empty(
|
||||
(bs, head_num, max_kv_splits),
|
||||
dtype=torch.float32,
|
||||
device="cuda",
|
||||
)
|
||||
o = torch.empty_like(q)
|
||||
num_kv_splits = torch.full(
|
||||
(bs,), max_kv_splits, dtype=torch.int32, device="cuda"
|
||||
)
|
||||
decode_attention_fwd(
|
||||
q,
|
||||
k_buf,
|
||||
v_buf,
|
||||
o,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
attn_logits,
|
||||
attn_lse,
|
||||
num_kv_splits,
|
||||
max_kv_splits,
|
||||
sm_scale=1.0 / (head_dim**0.5),
|
||||
k_scale=1.0,
|
||||
v_scale=1.0,
|
||||
logit_cap=0.0,
|
||||
page_size=page_size,
|
||||
)
|
||||
return o
|
||||
|
||||
def test_decode_3d_vs_4d_ps1_byte_identical(self):
|
||||
"""(a) vs (b): same physical bytes, different view shape.
|
||||
Triton specializes PAGE_SIZE=1 to the legacy branch; output must
|
||||
be bit-identical (modulo non-deterministic FP add ordering, which
|
||||
we sidestep here since the kernels use deterministic reductions
|
||||
for fixed input + grid)."""
|
||||
q, k, v, kv_indptr, kv_indices, seq_len = self._setup_decode_inputs()
|
||||
# (a) legacy 3-D view
|
||||
o_3d = self._run_decode(q, k, v, kv_indptr, kv_indices, page_size=1)
|
||||
# (b) 4-D view: reshape SAME physical bytes to (num_pages=N, 1, head, dim)
|
||||
num_slots = k.shape[0]
|
||||
k_4d = k.view(num_slots, 1, *k.shape[1:])
|
||||
v_4d = v.view(num_slots, 1, *v.shape[1:])
|
||||
o_4d_ps1 = self._run_decode(q, k_4d, v_4d, kv_indptr, kv_indices, page_size=1)
|
||||
# bit-identical (same byte layout, same PAGE_SIZE specialization)
|
||||
self.assertTrue(torch.equal(o_3d, o_4d_ps1))
|
||||
|
||||
def test_extend_3d_vs_4d_ps1_byte_identical(self):
|
||||
"""Same parity check for extend kernel."""
|
||||
from sglang.kernels.ops.attention.extend_attention import (
|
||||
extend_attention_fwd,
|
||||
)
|
||||
|
||||
torch.manual_seed(0xDEADBEEF)
|
||||
# head_dim must be >= 16: the extend kernel's QK^T tl.dot requires the
|
||||
# contraction dim K (= head_dim) >= 16 on modern GPU archs (Hopper+).
|
||||
head_num, head_dim = 2, 32
|
||||
num_slots = 32
|
||||
dtype = torch.float16
|
||||
bs = 2
|
||||
prefix_len = 8
|
||||
extend_len = 4
|
||||
|
||||
k_buffer = torch.randn(
|
||||
num_slots, head_num, head_dim, dtype=dtype, device="cuda"
|
||||
)
|
||||
v_buffer = torch.randn(
|
||||
num_slots, head_num, head_dim, dtype=dtype, device="cuda"
|
||||
)
|
||||
q_extend = torch.randn(
|
||||
bs * extend_len, head_num, head_dim, dtype=dtype, device="cuda"
|
||||
)
|
||||
k_extend = torch.randn(
|
||||
bs * extend_len, head_num, head_dim, dtype=dtype, device="cuda"
|
||||
)
|
||||
v_extend = torch.randn(
|
||||
bs * extend_len, head_num, head_dim, dtype=dtype, device="cuda"
|
||||
)
|
||||
o = torch.empty_like(q_extend)
|
||||
|
||||
qo_indptr = torch.tensor(
|
||||
[i * extend_len for i in range(bs + 1)], dtype=torch.int32, device="cuda"
|
||||
)
|
||||
kv_indptr = torch.tensor(
|
||||
[i * prefix_len for i in range(bs + 1)], dtype=torch.int32, device="cuda"
|
||||
)
|
||||
kv_indices = torch.arange(prefix_len, dtype=torch.int64, device="cuda").repeat(
|
||||
bs
|
||||
)
|
||||
|
||||
def run(k_buf, v_buf, page_size):
|
||||
o_out = torch.empty_like(q_extend)
|
||||
extend_attention_fwd(
|
||||
q_extend,
|
||||
k_extend,
|
||||
v_extend,
|
||||
o_out,
|
||||
k_buf,
|
||||
v_buf,
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
custom_mask=None,
|
||||
is_causal=True,
|
||||
mask_indptr=None,
|
||||
max_len_extend=extend_len,
|
||||
k_scale=1.0,
|
||||
v_scale=1.0,
|
||||
sm_scale=1.0 / (head_dim**0.5),
|
||||
page_size=page_size,
|
||||
)
|
||||
return o_out
|
||||
|
||||
o_3d = run(k_buffer, v_buffer, page_size=1)
|
||||
k_4d = k_buffer.view(num_slots, 1, *k_buffer.shape[1:])
|
||||
v_4d = v_buffer.view(num_slots, 1, *v_buffer.shape[1:])
|
||||
o_4d_ps1 = run(k_4d, v_4d, page_size=1)
|
||||
self.assertTrue(torch.equal(o_3d, o_4d_ps1))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -47,7 +47,6 @@ def _build(device, page_size=1, kernel_page_multiplier=None):
|
||||
device=device,
|
||||
enable_memory_saver=False,
|
||||
page_size=page_size,
|
||||
view_tail_pad_bytes=page_size * full_spec.entry_bytes(),
|
||||
)
|
||||
kvcache = UnifiedMLATokenToKVPool(
|
||||
unified_buffer=buf,
|
||||
|
||||
@@ -0,0 +1,531 @@
|
||||
# Copyright 2023-2026 SGLang Team
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
"""Dense MHA K/V views for the unified memory pool (uniform-row hybrid models).
|
||||
|
||||
Covers, CPU-only (pure torch — no GPU / Triton kernels):
|
||||
- `build_mha_views` refuses an asymmetric-KV spec: its addressing
|
||||
assumes one uniform row width, so it is the boundary that checks;
|
||||
- `build_mha_views` addressing: view_l[kernel_id(t)] must land exactly at
|
||||
the page-major envelope byte offset the STRIDED builder assigns to the same
|
||||
(page, slot, layer, K|V) cell — the two builders are views over one truth;
|
||||
- K and V of one token share ONE kernel-facing id (per-layer origin shift does the
|
||||
disambiguation), with no aliasing across the 2*L overlapping views;
|
||||
- the missing-tail-pad and asymmetric-dims cases fail loud at construction.
|
||||
|
||||
Addressing law under test (the derived property everything else builds on):
|
||||
|
||||
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 — offsets identical to
|
||||
MHASubPoolSpec.layer_k/v_offset_in_page when rows are uniform.
|
||||
|
||||
python -m pytest test/registered/unit/mem_cache/test_unified_mha_views.py -v
|
||||
"""
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=8, suite="base-a-test-cpu")
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.mem_cache.layout.page_major import (
|
||||
build_mha_views,
|
||||
mha_entry_bytes,
|
||||
)
|
||||
from sglang.srt.mem_cache.unified_memory_pool import (
|
||||
MHASubPoolSpec,
|
||||
UnifiedKVPool,
|
||||
UnifiedMHATokenToKVPool,
|
||||
)
|
||||
|
||||
_DEV = "cpu"
|
||||
# `set_kv_buffer` dispatches on the PLATFORM (memory_pool._is_cuda, resolved at
|
||||
# import), not on the tensors it is handed, so cases driving it must build on
|
||||
# the platform's device. The rest of this file is byte arithmetic, so CPU.
|
||||
_STORE_DEV = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
# Small-but-nontrivial MHA geometry: L=2 layers, H=2 heads, D=4, so every byte
|
||||
# offset is hand-checkable. blocks = 2L = 4 per page.
|
||||
_L = 2
|
||||
_H = 2
|
||||
_D = 4
|
||||
_ROW = _H * _D # row elements
|
||||
_DTYPE = torch.bfloat16
|
||||
_ITEM = _DTYPE.itemsize
|
||||
_BLOCKS = 2 * _L
|
||||
|
||||
|
||||
def _mha_spec(head_dim=_D, v_head_dim=None, layer_num=_L, grow="down"):
|
||||
return MHASubPoolSpec(
|
||||
name="full",
|
||||
layer_num=layer_num,
|
||||
head_num=_H,
|
||||
head_dim=head_dim,
|
||||
v_head_dim=v_head_dim,
|
||||
store_dtype=_DTYPE,
|
||||
grow_direction=grow,
|
||||
)
|
||||
|
||||
|
||||
def _kernel_id(t, ps):
|
||||
return (t // ps) * (ps * _BLOCKS) + t % ps
|
||||
|
||||
|
||||
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
|
||||
)
|
||||
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 page-major envelope.
|
||||
|
||||
This is the retired production strided builder, kept here as the oracle:
|
||||
per-layer ``(num_pages, page_size, head_num, head_dim)`` views addressed by
|
||||
``(page, slot)``, so the builder's addressing can be cross-checked
|
||||
against a second, independently-derived description of the same bytes.
|
||||
"""
|
||||
k_row_bytes = _ROW * _ITEM
|
||||
v_row_bytes = _ROW * _ITEM
|
||||
page_bytes = page_size * _L * (k_row_bytes + v_row_bytes)
|
||||
as_dtype_view = raw.view(_DTYPE)
|
||||
k_stride = (page_bytes // _ITEM, k_row_bytes // _ITEM, _D, 1)
|
||||
v_stride = (page_bytes // _ITEM, v_row_bytes // _ITEM, _D, 1)
|
||||
shape = (num_pages, page_size, _H, _D)
|
||||
k_views, v_views = [], []
|
||||
for layer in range(_L):
|
||||
k_base = anchor_bytes + layer * page_size * (k_row_bytes + v_row_bytes)
|
||||
v_base = k_base + page_size * k_row_bytes
|
||||
k_views.append(
|
||||
torch.as_strided(
|
||||
as_dtype_view,
|
||||
size=shape,
|
||||
stride=k_stride,
|
||||
storage_offset=k_base // _ITEM,
|
||||
)
|
||||
)
|
||||
v_views.append(
|
||||
torch.as_strided(
|
||||
as_dtype_view,
|
||||
size=shape,
|
||||
stride=v_stride,
|
||||
storage_offset=v_base // _ITEM,
|
||||
)
|
||||
)
|
||||
return k_views, v_views
|
||||
|
||||
|
||||
class TestMHADenseSpecSurface(unittest.TestCase):
|
||||
def test_asymmetric_rows_refused_by_the_view_builder(self):
|
||||
"""The row-block array exists only for uniform rows, so the builder
|
||||
whose addressing depends on it is the one that refuses (the MiMoV2
|
||||
shape, scaled down). ServerArgs screens such models out of
|
||||
--enable-unified-memory long before we get here; this is the check for
|
||||
a caller that reaches the builder directly."""
|
||||
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,
|
||||
)
|
||||
|
||||
def test_spec_offsets_equal_block_origins(self):
|
||||
"""The spec's byte math and the view builder's origins are two
|
||||
independent derivations of the envelope; under uniform rows they must
|
||||
agree: layer_k_offset(l) == (2l)*ps*row, layer_v_offset(l) == (2l+1)*ps*row."""
|
||||
spec = _mha_spec()
|
||||
for ps in (1, 4):
|
||||
row = spec.k_row_bytes()
|
||||
for l in range(_L):
|
||||
self.assertEqual(spec.layer_k_offset_in_page(l, ps), (2 * l) * ps * row)
|
||||
self.assertEqual(
|
||||
spec.layer_v_offset_in_page(l, ps), (2 * l + 1) * ps * row
|
||||
)
|
||||
|
||||
def test_entry_bytes_matches_layout_helper(self):
|
||||
spec = _mha_spec()
|
||||
self.assertEqual(
|
||||
spec.entry_bytes(),
|
||||
mha_entry_bytes(
|
||||
layer_num=_L, head_num=_H, head_dim=_D, v_head_dim=_D, itemsize=_ITEM
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class TestDenseMHAViews(unittest.TestCase):
|
||||
def test_view_shapes_are_stock_mha(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
|
||||
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.
|
||||
self.assertEqual(tuple(v.shape), (n_rows, _H, _D))
|
||||
self.assertEqual(v.stride(), (_ROW, _D, 1))
|
||||
|
||||
def test_addressing_matches_strided_reference(self):
|
||||
"""Cross-readback: bytes written through the reference STRIDED views at
|
||||
(page, slot) must be read back through the views at kernel_id(t),
|
||||
for both K and V of every layer — and vice versa. This pins that the
|
||||
view builder and the independent strided description agree on the
|
||||
same physical envelope."""
|
||||
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)
|
||||
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)),
|
||||
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)),
|
||||
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_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: page origin + layer-block origin +
|
||||
slot offset. Independent of any view builder — this is the raw layout
|
||||
contract every envelope consumer (moves, sizing, transfer math) relies
|
||||
on."""
|
||||
k_row = _ROW * _ITEM
|
||||
v_row = _ROW * _ITEM
|
||||
for ps in (1, 4):
|
||||
num_pages = 5
|
||||
page_bytes = ps * _L * (k_row + v_row)
|
||||
dk, dv = _build_views(_make_raw(ps, num_pages), ps, num_pages)
|
||||
for t in (0, 1, ps, 3 * ps + (ps - 1), 4 * ps):
|
||||
d = _kernel_id(t, ps)
|
||||
for L in range(_L):
|
||||
expected_k = (
|
||||
(t // ps) * page_bytes
|
||||
+ L * ps * (k_row + v_row)
|
||||
+ (t % ps) * k_row
|
||||
)
|
||||
expected_v = (
|
||||
(t // ps) * page_bytes
|
||||
+ L * ps * (k_row + v_row)
|
||||
+ ps * k_row
|
||||
+ (t % ps) * v_row
|
||||
)
|
||||
got_k = (dk[L].storage_offset() + d * dk[L].stride(0)) * _ITEM
|
||||
got_v = (dv[L].storage_offset() + d * 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_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."""
|
||||
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)
|
||||
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)))
|
||||
|
||||
def test_missing_tail_pad_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)
|
||||
|
||||
def test_asymmetric_dims_rejected(self):
|
||||
ps, num_pages = 2, 4
|
||||
raw = _make_raw(ps, num_pages)
|
||||
with self.assertRaises(AssertionError):
|
||||
_build_views(raw, ps, num_pages, head_dim=6, v_head_dim=4)
|
||||
|
||||
|
||||
# ---- pool level ----
|
||||
|
||||
_N_FULL = 32 # full-attn token slots per pool in the fixtures below
|
||||
_N_SWA = 16
|
||||
|
||||
|
||||
def _swa_spec(grow="up", head_dim=_D, v_head_dim=None):
|
||||
return MHASubPoolSpec(
|
||||
name="swa",
|
||||
layer_num=_L,
|
||||
head_num=_H,
|
||||
head_dim=head_dim,
|
||||
v_head_dim=v_head_dim,
|
||||
store_dtype=_DTYPE,
|
||||
grow_direction=grow,
|
||||
)
|
||||
|
||||
|
||||
def _make_pool(ps=1, full_spec=None, device=_DEV):
|
||||
full = full_spec if full_spec is not None else _mha_spec()
|
||||
swa = _swa_spec()
|
||||
total = full.entry_bytes() * _N_FULL + swa.entry_bytes() * _N_SWA
|
||||
return UnifiedKVPool(
|
||||
total_bytes=total,
|
||||
sub_pool_specs=[full, swa],
|
||||
device=device,
|
||||
enable_memory_saver=False,
|
||||
page_size=ps,
|
||||
)
|
||||
|
||||
|
||||
class TestUnifiedKVPoolDenseViews(unittest.TestCase):
|
||||
def test_every_mha_sub_pool_is_per_layer_contiguous(self):
|
||||
"""The unified pool has ONE MHA layout: both sub-pools come back as
|
||||
stock 3-D per-layer views, whatever their page size."""
|
||||
for ps in (1, 4):
|
||||
pool = _make_pool(ps=ps)
|
||||
for name in ("full", "swa"):
|
||||
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())
|
||||
|
||||
def test_tail_pad_is_derived_from_the_specs(self):
|
||||
"""The per-layer views hang past the last page envelope, so the pool
|
||||
over-allocates one envelope of the widest sub-pool. Derived here, not
|
||||
passed in, so no construction site can under-allocate it."""
|
||||
for ps in (1, 4):
|
||||
kv = _make_pool(ps)
|
||||
full, swa = _mha_spec(), _swa_spec()
|
||||
self.assertEqual(
|
||||
kv.view_tail_pad_bytes,
|
||||
ps * max(full.entry_bytes(), swa.entry_bytes()),
|
||||
f"tail pad at ps={ps}",
|
||||
)
|
||||
self.assertEqual(
|
||||
kv._raw.numel(),
|
||||
full.entry_bytes() * _N_FULL
|
||||
+ swa.entry_bytes() * _N_SWA
|
||||
+ kv.view_tail_pad_bytes,
|
||||
"the pad extends the allocation only",
|
||||
)
|
||||
|
||||
|
||||
def _layer(l):
|
||||
return SimpleNamespace(layer_id=l)
|
||||
|
||||
|
||||
def _make_pool_and_kv(ps, device=_DEV):
|
||||
kv = _make_pool(ps=ps, device=device)
|
||||
return kv, UnifiedMHATokenToKVPool(
|
||||
unified_buffer=kv,
|
||||
sub_pool_name="full",
|
||||
page_size=ps,
|
||||
enable_alt_stream=False,
|
||||
)
|
||||
|
||||
|
||||
class TestUnifiedMHATokenToKVPool(unittest.TestCase):
|
||||
def test_size_is_view_row_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."""
|
||||
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
|
||||
self.assertEqual(pool_under_test.size, n_rows - ps)
|
||||
|
||||
def test_stock_write_lands_on_envelope_truth(self):
|
||||
"""Byte-identity: the pool's stock inherited `set_kv_buffer` at kernel-facing
|
||||
locs must produce exactly the bytes that direct writes through STRIDED
|
||||
views over the same envelope produce at the same (page, slot, layer)
|
||||
cells. The strided views are built here purely as the independent
|
||||
description of the envelope — pins the whole write path (loc -> 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(
|
||||
kv._raw,
|
||||
page_size=ps,
|
||||
num_pages=kv.max_slots("full") // ps,
|
||||
anchor_bytes=kv.anchor_bytes("full"),
|
||||
)
|
||||
probes = [(1, 0), (2, ps - 1), (5, ps // 2)]
|
||||
for l in range(_L):
|
||||
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)
|
||||
for p, s in probes:
|
||||
self.assertTrue(
|
||||
torch.all(sk[l][p, s] == float(l + 1)),
|
||||
f"K (l={l}, p={p}, s={s}, ps={ps}) not at the envelope cell",
|
||||
)
|
||||
self.assertTrue(
|
||||
torch.all(sv[l][p, s] == float(l + 101)),
|
||||
f"V (l={l}, p={p}, s={s}, ps={ps}) not at the envelope cell",
|
||||
)
|
||||
|
||||
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 — red if it is
|
||||
lost, since the inherited per-layer move would apply physical ids to
|
||||
the row space."""
|
||||
ps = 4
|
||||
kv, pool = _make_pool_and_kv(ps)
|
||||
live = kv._raw.numel() - kv.view_tail_pad_bytes
|
||||
seed = (torch.arange(live, dtype=torch.float32) % 251).to(torch.uint8)
|
||||
kv._raw[:live] = seed
|
||||
page_bytes = ps * _mha_spec().entry_bytes()
|
||||
|
||||
src_pages, tgt_pages = torch.tensor([5, 6]), torch.tensor([2, 3])
|
||||
offs = torch.arange(ps)
|
||||
run = lambda p: (p[:, None] * ps + offs).reshape(-1)
|
||||
pool.move_kv_cache(run(tgt_pages), run(src_pages))
|
||||
|
||||
want = seed.clone()
|
||||
for sp, tp in zip(src_pages.tolist(), tgt_pages.tolist()):
|
||||
want[tp * page_bytes : (tp + 1) * page_bytes] = seed[
|
||||
sp * page_bytes : (sp + 1) * page_bytes
|
||||
]
|
||||
self.assertTrue(
|
||||
torch.equal(kv._raw[:live], want),
|
||||
"envelope move did not relocate exactly the named pages",
|
||||
)
|
||||
|
||||
def test_transfer_entry_points_fail_loud(self):
|
||||
"""PD / CPU-copy entry points assume per-layer buffers indexed by TOKEN
|
||||
id; against the row space they would silently mis-index (or hit a
|
||||
missing-attr AttributeError). Every one of them must raise."""
|
||||
_, pool = _make_pool_and_kv(1)
|
||||
with self.assertRaises(NotImplementedError):
|
||||
pool.get_contiguous_buf_infos()
|
||||
with self.assertRaises(NotImplementedError):
|
||||
pool.get_cpu_copy(torch.tensor([1]))
|
||||
with self.assertRaises(NotImplementedError):
|
||||
pool.load_cpu_copy(None, torch.tensor([1]))
|
||||
with self.assertRaises(NotImplementedError):
|
||||
pool.set_kv_buffer_prefix_valid()
|
||||
|
||||
def test_hnd_env_cannot_hijack_layout(self):
|
||||
"""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)
|
||||
self.assertEqual(pool.kv_cache_layout, "page_major")
|
||||
|
||||
|
||||
class TestFactoryDenseViews(unittest.TestCase):
|
||||
"""The real SWA factory builds both sub-pools and wires the matching
|
||||
kernel-facing multipliers into the composite allocator."""
|
||||
|
||||
# _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):
|
||||
# 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(
|
||||
device="cpu",
|
||||
kv_cache_dtype=torch.float16,
|
||||
head_num=2,
|
||||
head_dim=8,
|
||||
v_head_dim=8,
|
||||
swa_head_num=2,
|
||||
swa_head_dim=8,
|
||||
swa_v_head_dim=8,
|
||||
page_size=1,
|
||||
start_layer=0,
|
||||
end_layer=4,
|
||||
swa_attention_layer_ids=[1, 3],
|
||||
full_attention_layer_ids=[0, 2],
|
||||
full_max_total_num_tokens=64,
|
||||
swa_max_total_num_tokens=32,
|
||||
enable_memory_saver=False,
|
||||
need_sort=False,
|
||||
)
|
||||
|
||||
def test_factory_wires_matching_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)
|
||||
# Sub-pools are the dense class exposing 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)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+21
-22
@@ -15,15 +15,15 @@
|
||||
memory pool (Kimi-Linear).
|
||||
|
||||
`req_to_token` holds VIRTUAL token ids, while the per-layer MLA views are dense
|
||||
(`build_dense_mla_views`). The paged MLA backends therefore need their page-level
|
||||
block table filled with DENSE page ids:
|
||||
(`build_mla_views`). The paged MLA backends therefore need their page-level
|
||||
block table filled with kernel-facing page ids:
|
||||
|
||||
dense_page(virtual_page) = v2p[virtual_page] * layer_num
|
||||
|
||||
Three backend families reach that same formula by different routes:
|
||||
- `create_flashmla_kv_indices_triton` in-kernel via `v2p_ptr` / `PAGE_MULT`
|
||||
(trtllm_mla / cutedsl_mla / tokenspeed_mla);
|
||||
- the flashinfer_mla updaters, post-gathering `translate_kv_loc_dense` over the
|
||||
- the flashinfer_mla updaters, post-gathering `translate_kv_loc_for_kernel` over the
|
||||
token-level kv_indices;
|
||||
- `normal_decode_set_metadata` in-kernel, for fa3's captured-decode page table.
|
||||
|
||||
@@ -32,13 +32,13 @@ Covered here:
|
||||
- kernel dense mapping against the python reference, for several page sizes,
|
||||
ragged sequence lengths and a non-identity v2p permutation;
|
||||
- padded block-table lanes never index the v2p table out of bounds;
|
||||
- the token-level dense translate the flashinfer updaters apply agrees with the
|
||||
- the token-level kernel-facing translate the flashinfer updaters apply agrees with the
|
||||
page-level block table the trtllm path builds;
|
||||
- fa3's fused metadata kernels agree with the same reference, on both the
|
||||
page_size == 1 fast path (which is what Kimi-Linear takes: fa3 imposes no
|
||||
page-size constraint) and the general path.
|
||||
|
||||
python -m pytest test/registered/unit/mem_cache/test_unified_mla_dense_block_table.py -v
|
||||
python -m pytest test/registered/unit/mem_cache/test_unified_mla_block_table.py -v
|
||||
"""
|
||||
|
||||
import unittest
|
||||
@@ -145,7 +145,7 @@ class TestDenseBlockTable(unittest.TestCase):
|
||||
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 dense id collapses onto the physical id, so
|
||||
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`:
|
||||
@@ -180,7 +180,7 @@ class TestDenseBlockTable(unittest.TestCase):
|
||||
|
||||
def test_agrees_with_token_level_dense_translate(self):
|
||||
"""The flashinfer updaters translate TOKEN ids with
|
||||
`translate_kv_loc_dense`; the trtllm path builds PAGE ids in-kernel. Both
|
||||
`translate_kv_loc_for_kernel`; the trtllm path builds PAGE ids in-kernel. Both
|
||||
must address the same dense page block."""
|
||||
page_size = 64
|
||||
rt, rpi, sl, v2p = self._make_batch(page_size)
|
||||
@@ -190,12 +190,12 @@ class TestDenseBlockTable(unittest.TestCase):
|
||||
for r in range(rt.shape[0]):
|
||||
n = int(sl[r].item())
|
||||
virt_tokens = rt[r, :n].long()
|
||||
# translate_kv_loc_dense's formula, applied to token ids.
|
||||
# translate_kv_loc_for_kernel's formula, applied to token ids.
|
||||
dense_tokens = (
|
||||
v2p[virt_tokens // page_size] * (page_size * _LAYERS)
|
||||
+ virt_tokens % page_size
|
||||
)
|
||||
# The block-table entry scaled by page_size must be the dense id of
|
||||
# The block-table entry scaled by page_size must be the kernel-facing id of
|
||||
# each page's first token.
|
||||
first_of_page = dense_tokens[::page_size]
|
||||
n_pages = (n + page_size - 1) // page_size
|
||||
@@ -330,19 +330,19 @@ class TestUnifiedMLAHookDetection(unittest.TestCase):
|
||||
hooks = self._probe()
|
||||
self.assertFalse(hooks.enabled)
|
||||
self.assertIsNone(hooks.v2p_page_table)
|
||||
self.assertIsNone(hooks.translate_kv_loc_dense)
|
||||
self.assertIsNone(hooks.translate_kv_loc_for_kernel)
|
||||
self.assertEqual(hooks.kernel_page_multiplier, 1)
|
||||
|
||||
def test_multi_layer_unified_pool(self):
|
||||
table = torch.arange(8)
|
||||
hooks = self._probe(
|
||||
full_v2p_page_table=table,
|
||||
translate_kv_loc_dense=lambda x, **kw: x,
|
||||
translate_kv_loc_for_kernel=lambda x, **kw: x,
|
||||
kernel_page_multiplier=_LAYERS,
|
||||
)
|
||||
self.assertTrue(hooks.enabled)
|
||||
self.assertIs(hooks.v2p_page_table, table)
|
||||
self.assertIsNotNone(hooks.translate_kv_loc_dense)
|
||||
self.assertIsNotNone(hooks.translate_kv_loc_for_kernel)
|
||||
self.assertEqual(hooks.kernel_page_multiplier, _LAYERS)
|
||||
|
||||
def test_single_full_attention_layer_pool_is_still_unified(self):
|
||||
@@ -356,13 +356,13 @@ class TestUnifiedMLAHookDetection(unittest.TestCase):
|
||||
table = torch.arange(8)
|
||||
hooks = self._probe(
|
||||
full_v2p_page_table=table,
|
||||
translate_kv_loc_dense=lambda x, **kw: x,
|
||||
translate_kv_loc_for_kernel=lambda x, **kw: x,
|
||||
kernel_page_multiplier=1,
|
||||
)
|
||||
self.assertTrue(hooks.enabled, "single-layer unified pool read as static")
|
||||
self.assertIs(hooks.v2p_page_table, table)
|
||||
self.assertIsNotNone(hooks.translate_kv_loc_dense)
|
||||
# Multiplier stays 1: dense id == physical id, so the v2p gather alone is
|
||||
self.assertIsNotNone(hooks.translate_kv_loc_for_kernel)
|
||||
# Multiplier stays 1: kernel-facing id == physical id, so the v2p gather alone is
|
||||
# the whole translation and PAGE_MULT must not scale it.
|
||||
self.assertEqual(hooks.kernel_page_multiplier, 1)
|
||||
|
||||
@@ -409,7 +409,6 @@ class TestInPlaceKvIndicesTranslate(unittest.TestCase):
|
||||
device=_DEV,
|
||||
enable_memory_saver=False,
|
||||
page_size=page_size,
|
||||
view_tail_pad_bytes=page_size * full.entry_bytes(),
|
||||
)
|
||||
|
||||
class _Stub:
|
||||
@@ -438,7 +437,7 @@ class TestInPlaceKvIndicesTranslate(unittest.TestCase):
|
||||
|
||||
def test_int32_buffer_prefix_translated_tail_untouched(self):
|
||||
"""Mirrors the updater: an int32 capture-stable buffer holding VIRTUAL
|
||||
ids in [:n] gets the dense ids written back in place, narrowed to int32,
|
||||
ids in [:n] gets the kernel-facing ids written back in place, narrowed to int32,
|
||||
with the stale tail left alone (it must never index the v2p table)."""
|
||||
alloc = self._allocator()
|
||||
virt = alloc.alloc(64)
|
||||
@@ -452,13 +451,13 @@ class TestInPlaceKvIndicesTranslate(unittest.TestCase):
|
||||
tail_before = buf[n:].clone()
|
||||
|
||||
valid = buf[:n]
|
||||
valid.copy_(alloc.translate_kv_loc_dense(valid))
|
||||
valid.copy_(alloc.translate_kv_loc_for_kernel(valid))
|
||||
|
||||
expected = alloc.translate_kv_loc_dense(virt)
|
||||
expected = alloc.translate_kv_loc_for_kernel(virt)
|
||||
self.assertEqual(buf.dtype, torch.int32)
|
||||
self.assertTrue(
|
||||
torch.equal(buf[:n].long(), expected),
|
||||
"in-place translate did not land dense ids in the stable buffer",
|
||||
"in-place translate did not land kernel-facing ids in the stable buffer",
|
||||
)
|
||||
self.assertTrue(
|
||||
torch.equal(buf[n:], tail_before),
|
||||
@@ -471,8 +470,8 @@ class TestInPlaceKvIndicesTranslate(unittest.TestCase):
|
||||
virt = alloc.alloc(64)
|
||||
self.assertIsNotNone(virt)
|
||||
self.assertFalse(
|
||||
torch.equal(alloc.translate_kv_loc_dense(virt), virt),
|
||||
"dense ids coincide with virtual ids; pick a different allocation",
|
||||
torch.equal(alloc.translate_kv_loc_for_kernel(virt), virt),
|
||||
"kernel-facing ids coincide with virtual ids; pick a different allocation",
|
||||
)
|
||||
|
||||
|
||||
@@ -11,16 +11,16 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
"""GPU parity of the dense-view `UnifiedMLATokenToKVPool` against the stock
|
||||
"""GPU parity of the per-layer-view `UnifiedMLATokenToKVPool` against the stock
|
||||
`MLATokenToKVPool` on real K3 MLA geometry (L=24, D=512+64).
|
||||
|
||||
The unified pool receives DENSE locs (dense(t) = (t//ps)*(ps*L) + t%ps); the
|
||||
The unified pool receives kernel-facing locs (kernel_id(t) = (t//ps)*(ps*L) + t%ps); the
|
||||
reference pool receives the raw token ids. Every (layer, token) cell must hold
|
||||
identical bytes afterwards. Covers:
|
||||
|
||||
- `set_mla_kv_buffer` under BOTH kernel paths — the Triton fallback
|
||||
(n_loc < 768) and the TMA JIT fast path (n_loc >= 768, which flattens the
|
||||
buffer via `.view(shape[0], -1)`, only legal because dense views are
|
||||
buffer via `.view(shape[0], -1)`, only legal because per-layer views are
|
||||
contiguous);
|
||||
- `set_kv_buffer` (combined pre-concatenated write, the Triton-backend path);
|
||||
- `get_mla_kv_buffer` roundtrip;
|
||||
@@ -48,7 +48,7 @@ _D = _LORA + _ROPE
|
||||
_DTYPE = torch.bfloat16
|
||||
|
||||
|
||||
def _dense(t: torch.Tensor, ps: int) -> torch.Tensor:
|
||||
def _kernel_id(t: torch.Tensor, ps: int) -> torch.Tensor:
|
||||
return (t // ps) * (ps * _L) + t % ps
|
||||
|
||||
|
||||
@@ -85,7 +85,6 @@ def _make_pools(ps: int, n_tokens: int = 4096):
|
||||
device=_DEV,
|
||||
enable_memory_saver=False,
|
||||
page_size=ps,
|
||||
view_tail_pad_bytes=ps * full.entry_bytes(),
|
||||
)
|
||||
unified = UnifiedMLATokenToKVPool(
|
||||
unified_buffer=pool,
|
||||
@@ -118,7 +117,7 @@ def _rand_locs(max_tokens: int, ps: int, n: int) -> torch.Tensor:
|
||||
class TestUnifiedMLAPoolGPUParity(unittest.TestCase):
|
||||
def _assert_parity(self, unified, ref, locs, ps, layers=range(_L)):
|
||||
for l in layers:
|
||||
got = unified.get_key_buffer(l)[_dense(locs, ps)]
|
||||
got = unified.get_key_buffer(l)[_kernel_id(locs, ps)]
|
||||
want = ref.get_key_buffer(l)[locs]
|
||||
torch.testing.assert_close(got, want, rtol=0, atol=0)
|
||||
|
||||
@@ -130,7 +129,7 @@ class TestUnifiedMLAPoolGPUParity(unittest.TestCase):
|
||||
layer = types.SimpleNamespace(layer_id=l)
|
||||
nope = torch.randn(n_loc, 1, _LORA, dtype=_DTYPE, device=_DEV)
|
||||
rope = torch.randn(n_loc, 1, _ROPE, dtype=_DTYPE, device=_DEV)
|
||||
unified.set_mla_kv_buffer(layer, _dense(locs, ps), nope, rope)
|
||||
unified.set_mla_kv_buffer(layer, _kernel_id(locs, ps), nope, rope)
|
||||
ref.set_mla_kv_buffer(layer, locs, nope, rope)
|
||||
torch.cuda.synchronize()
|
||||
self._assert_parity(unified, ref, locs, ps)
|
||||
@@ -156,7 +155,7 @@ class TestUnifiedMLAPoolGPUParity(unittest.TestCase):
|
||||
for l in (0, _L // 2, _L - 1):
|
||||
layer = types.SimpleNamespace(layer_id=l)
|
||||
k = torch.randn(n_loc, 1, _D, dtype=_DTYPE, device=_DEV)
|
||||
unified.set_kv_buffer(layer, _dense(locs, ps), k, None)
|
||||
unified.set_kv_buffer(layer, _kernel_id(locs, ps), k, None)
|
||||
ref.set_kv_buffer(layer, locs, k, None)
|
||||
torch.cuda.synchronize()
|
||||
self._assert_parity(unified, ref, locs, ps, layers=(0, _L // 2, _L - 1))
|
||||
@@ -170,8 +169,8 @@ class TestUnifiedMLAPoolGPUParity(unittest.TestCase):
|
||||
layer = types.SimpleNamespace(layer_id=3)
|
||||
nope = torch.randn(n_loc, 1, _LORA, dtype=_DTYPE, device=_DEV)
|
||||
rope = torch.randn(n_loc, 1, _ROPE, dtype=_DTYPE, device=_DEV)
|
||||
unified.set_mla_kv_buffer(layer, _dense(locs, ps), nope, rope)
|
||||
got_nope, got_rope = unified.get_mla_kv_buffer(layer, _dense(locs, ps))
|
||||
unified.set_mla_kv_buffer(layer, _kernel_id(locs, ps), nope, rope)
|
||||
got_nope, got_rope = unified.get_mla_kv_buffer(layer, _kernel_id(locs, ps))
|
||||
torch.cuda.synchronize()
|
||||
torch.testing.assert_close(got_nope, nope, rtol=0, atol=0)
|
||||
torch.testing.assert_close(got_rope, rope, rtol=0, atol=0)
|
||||
@@ -188,14 +187,15 @@ class TestUnifiedMLAPoolGPUParity(unittest.TestCase):
|
||||
for l in range(_L):
|
||||
layer = types.SimpleNamespace(layer_id=l)
|
||||
k = torch.randn(n_loc, 1, _D, dtype=_DTYPE, device=_DEV)
|
||||
unified.set_kv_buffer(layer, _dense(src_t, ps), k, None)
|
||||
unified.set_kv_buffer(layer, _kernel_id(src_t, ps), k, None)
|
||||
before = [
|
||||
unified.get_key_buffer(l)[_dense(src_t, ps)].clone() for l in range(_L)
|
||||
unified.get_key_buffer(l)[_kernel_id(src_t, ps)].clone()
|
||||
for l in range(_L)
|
||||
]
|
||||
unified.move_kv_cache(dst_t, src_t)
|
||||
torch.cuda.synchronize()
|
||||
for l in range(_L):
|
||||
got = unified.get_key_buffer(l)[_dense(dst_t, ps)]
|
||||
got = unified.get_key_buffer(l)[_kernel_id(dst_t, ps)]
|
||||
torch.testing.assert_close(got, before[l], rtol=0, atol=0)
|
||||
|
||||
|
||||
|
||||
@@ -15,15 +15,15 @@
|
||||
|
||||
Covers, CPU-only (pure torch — no GPU / Triton kernels):
|
||||
- `MLASubPoolSpec` byte math;
|
||||
- `build_dense_mla_views` addressing: view_l[dense(t)] must land exactly at
|
||||
- `build_mla_views` addressing: view_l[kernel_id(t)] must land exactly at
|
||||
the page-major envelope byte offset `p*(L*ps*D) + l*(ps*D) + s*D`, the
|
||||
overlapping per-layer views must not alias at equal dense ids, and the
|
||||
overlapping per-layer views must not alias at equal kernel-facing ids, and the
|
||||
missing-tail-pad case must fail loud;
|
||||
- `UnifiedKVPool` MLA plumbing: `view_tail_pad_bytes` extends the allocation
|
||||
only, and the reserved sink floor covers the whole page-0 envelope;
|
||||
- `UnifiedMLATokenToKVPool`: buffer wiring, V-as-prefix-slice, and the
|
||||
page-envelope `move_kv_cache` (REAL physical token ids, page-major runs);
|
||||
- `MultiEndedAllocator.translate_kv_loc_dense`: dense = v2p-page * (ps*L) +
|
||||
- `MultiEndedAllocator.translate_kv_loc_for_kernel`: dense = v2p-page * (ps*L) +
|
||||
offset, tombstone clamp to the sink, `out=` contract, multiplier-1
|
||||
fallback, and correctness across eager compaction.
|
||||
|
||||
@@ -42,7 +42,7 @@ import unittest
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.layout.page_major import (
|
||||
build_dense_mla_views,
|
||||
build_mla_views,
|
||||
mla_entry_bytes,
|
||||
)
|
||||
from sglang.srt.mem_cache.multi_ended_allocator import MultiEndedAllocator
|
||||
@@ -98,12 +98,11 @@ def _make_unified(page_size=1, n_full_tokens=64, n_mamba_slots=8):
|
||||
device=_DEV,
|
||||
enable_memory_saver=False,
|
||||
page_size=page_size,
|
||||
view_tail_pad_bytes=page_size * full.entry_bytes(),
|
||||
)
|
||||
return pool, full, mamba
|
||||
|
||||
|
||||
def _dense(t, ps, layer_num):
|
||||
def _kernel_id(t, ps, layer_num):
|
||||
return (t // ps) * (ps * layer_num) + t % ps
|
||||
|
||||
|
||||
@@ -142,7 +141,7 @@ class TestDenseMLAViews(unittest.TestCase):
|
||||
for ps in (1, 4):
|
||||
num_pages = 6
|
||||
raw, _ = self._make_raw(ps, num_pages)
|
||||
views = build_dense_mla_views(
|
||||
views = build_mla_views(
|
||||
raw,
|
||||
layer_num=_L,
|
||||
kv_cache_dim=_D,
|
||||
@@ -151,9 +150,9 @@ class TestDenseMLAViews(unittest.TestCase):
|
||||
num_pages=num_pages,
|
||||
)
|
||||
self.assertEqual(len(views), _L)
|
||||
n_dense = num_pages * _L * ps
|
||||
n_rows = num_pages * _L * ps
|
||||
for v in views:
|
||||
self.assertEqual(tuple(v.shape), (n_dense, 1, _D))
|
||||
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)
|
||||
@@ -161,7 +160,7 @@ class TestDenseMLAViews(unittest.TestCase):
|
||||
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][_dense(t, ps, _L)] = marker
|
||||
views[l][_kernel_id(t, ps, _L)] = marker
|
||||
# envelope formula, in elements
|
||||
elem = p * (_L * ps * _D) + l * (ps * _D) + s * _D
|
||||
self.assertTrue(
|
||||
@@ -173,7 +172,7 @@ class TestDenseMLAViews(unittest.TestCase):
|
||||
ps = 4
|
||||
num_pages = 4
|
||||
raw, _ = self._make_raw(ps, num_pages)
|
||||
views = build_dense_mla_views(
|
||||
views = build_mla_views(
|
||||
raw,
|
||||
layer_num=_L,
|
||||
kv_cache_dim=_D,
|
||||
@@ -182,7 +181,7 @@ class TestDenseMLAViews(unittest.TestCase):
|
||||
num_pages=num_pages,
|
||||
)
|
||||
t = 2 * ps + 1 # page 2, slot 1
|
||||
d = _dense(t, ps, _L)
|
||||
d = _kernel_id(t, ps, _L)
|
||||
for l in range(_L):
|
||||
views[l][d] = float(l + 1)
|
||||
for l in range(_L):
|
||||
@@ -193,7 +192,7 @@ class TestDenseMLAViews(unittest.TestCase):
|
||||
num_pages = 4
|
||||
raw, _ = self._make_raw(ps, num_pages, pad_pages=0)
|
||||
with self.assertRaises(AssertionError):
|
||||
build_dense_mla_views(
|
||||
build_mla_views(
|
||||
raw,
|
||||
layer_num=_L,
|
||||
kv_cache_dim=_D,
|
||||
@@ -279,7 +278,7 @@ class TestUnifiedMLATokenToKVPool(unittest.TestCase):
|
||||
# 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][_dense(src_page * ps + s, ps, _L)] = float(
|
||||
kv_pool.kv_buffer[l][_kernel_id(src_page * ps + s, ps, _L)] = float(
|
||||
l * ps + s + 1
|
||||
)
|
||||
offsets = torch.arange(ps, dtype=torch.int64)
|
||||
@@ -289,7 +288,7 @@ class TestUnifiedMLATokenToKVPool(unittest.TestCase):
|
||||
)
|
||||
for l in range(_L):
|
||||
for s in range(ps):
|
||||
got = kv_pool.kv_buffer[l][_dense(dst_page * ps + s, ps, _L)]
|
||||
got = kv_pool.kv_buffer[l][_kernel_id(dst_page * ps + s, ps, _L)]
|
||||
self.assertTrue(
|
||||
torch.all(got == float(l * ps + s + 1)), f"(l={l}, s={s})"
|
||||
)
|
||||
@@ -331,7 +330,7 @@ class TestTranslateKvLocDense(unittest.TestCase):
|
||||
v = alloc.alloc(8)
|
||||
self.assertIsNotNone(v)
|
||||
phys = alloc.translate_kv_loc(v)
|
||||
dense = alloc.translate_kv_loc_dense(v)
|
||||
dense = alloc.translate_kv_loc_for_kernel(v)
|
||||
self.assertTrue(torch.all(dense == phys * _L))
|
||||
|
||||
def test_dense_matches_formula_paged(self):
|
||||
@@ -340,15 +339,15 @@ class TestTranslateKvLocDense(unittest.TestCase):
|
||||
v = alloc.alloc(3 * ps)
|
||||
self.assertIsNotNone(v)
|
||||
phys = alloc.translate_kv_loc(v)
|
||||
dense = alloc.translate_kv_loc_dense(v)
|
||||
dense = alloc.translate_kv_loc_for_kernel(v)
|
||||
expected = (phys // ps) * (ps * _L) + phys % ps
|
||||
self.assertTrue(torch.all(dense == expected))
|
||||
|
||||
def test_tombstone_clamps_to_sink(self):
|
||||
alloc = self._build(ps=1)
|
||||
# never-allocated virtual ids -> v2p == -1 -> dense id 0
|
||||
# never-allocated virtual ids -> v2p == -1 -> kernel-facing id 0
|
||||
virt = torch.tensor([alloc.min_slot_index + 1], dtype=torch.int64)
|
||||
dense = alloc.translate_kv_loc_dense(virt)
|
||||
dense = alloc.translate_kv_loc_for_kernel(virt)
|
||||
self.assertTrue(torch.all(dense == 0))
|
||||
|
||||
def test_out_matches_and_aliases(self):
|
||||
@@ -356,14 +355,14 @@ class TestTranslateKvLocDense(unittest.TestCase):
|
||||
alloc = self._build(ps=ps)
|
||||
v = alloc.alloc(2 * ps)
|
||||
self.assertIsNotNone(v)
|
||||
no_out = alloc.translate_kv_loc_dense(v)
|
||||
no_out = alloc.translate_kv_loc_for_kernel(v)
|
||||
out = torch.empty_like(v)
|
||||
ret = alloc.translate_kv_loc_dense(v, out=out)
|
||||
ret = alloc.translate_kv_loc_for_kernel(v, out=out)
|
||||
self.assertIs(ret, out)
|
||||
self.assertTrue(torch.all(out == no_out))
|
||||
# canonical in-place aliasing: translate(x, out=x)
|
||||
x = v.clone()
|
||||
alloc.translate_kv_loc_dense(x, out=x)
|
||||
alloc.translate_kv_loc_for_kernel(x, out=x)
|
||||
self.assertTrue(torch.all(x == no_out))
|
||||
|
||||
def test_multiplier_one_falls_back_to_physical(self):
|
||||
@@ -371,7 +370,7 @@ class TestTranslateKvLocDense(unittest.TestCase):
|
||||
v = alloc.alloc(4)
|
||||
self.assertIsNotNone(v)
|
||||
self.assertTrue(
|
||||
torch.all(alloc.translate_kv_loc_dense(v) == alloc.translate_kv_loc(v))
|
||||
torch.all(alloc.translate_kv_loc_for_kernel(v) == alloc.translate_kv_loc(v))
|
||||
)
|
||||
|
||||
def test_dense_follows_compaction(self):
|
||||
@@ -383,8 +382,8 @@ class TestTranslateKvLocDense(unittest.TestCase):
|
||||
alloc.free(b) # eager compaction relocates survivors
|
||||
phys_a = alloc.translate_kv_loc(a)
|
||||
phys_c = alloc.translate_kv_loc(c)
|
||||
self.assertTrue(torch.all(alloc.translate_kv_loc_dense(a) == phys_a * _L))
|
||||
self.assertTrue(torch.all(alloc.translate_kv_loc_dense(c) == phys_c * _L))
|
||||
self.assertTrue(torch.all(alloc.translate_kv_loc_for_kernel(a) == phys_a * _L))
|
||||
self.assertTrue(torch.all(alloc.translate_kv_loc_for_kernel(c) == phys_c * _L))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -15,9 +15,9 @@
|
||||
|
||||
The page-major envelope K/V views are strided, which only the Triton attention
|
||||
kernels read. The one exception is the unified-memory MLA pool: it exposes each
|
||||
layer as a DENSE contiguous view (`build_dense_mla_views`), so the paged MLA
|
||||
layer as a contiguous view (`build_mla_views`), so the paged MLA
|
||||
backends can read it directly once their kv_indices / block tables are remapped
|
||||
to dense ids -- `fa3`, `flashinfer`'s MLA backend, and `trtllm_mla` with its
|
||||
to kernel-facing ids -- `fa3`, `flashinfer`'s MLA backend, and `trtllm_mla` with its
|
||||
`cutedsl_mla` / `tokenspeed_mla` subclasses.
|
||||
|
||||
Pinned here so the exception cannot silently widen to a backend that has no
|
||||
@@ -47,6 +47,7 @@ def _accepts(
|
||||
unified: bool = True,
|
||||
linear_decode: str | None = None,
|
||||
linear_prefill: str | None = None,
|
||||
has_asymmetric_kv: bool = False,
|
||||
) -> bool:
|
||||
"""Run just `_handle_page_major_kv_layout` against a minimal stand-in.
|
||||
|
||||
@@ -72,7 +73,12 @@ def _accepts(
|
||||
sa,
|
||||
"_model_config",
|
||||
SimpleNamespace(
|
||||
attention_arch=AttentionArch.MLA if use_mla else AttentionArch.MHA
|
||||
attention_arch=AttentionArch.MLA if use_mla else AttentionArch.MHA,
|
||||
has_asymmetric_kv=has_asymmetric_kv,
|
||||
head_dim=192 if has_asymmetric_kv else 128,
|
||||
v_head_dim=128,
|
||||
swa_head_dim=128,
|
||||
swa_v_head_dim=128,
|
||||
),
|
||||
)
|
||||
try:
|
||||
@@ -106,7 +112,7 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
|
||||
)
|
||||
|
||||
def test_dense_mla_backends_rejected_for_mha(self):
|
||||
"""The dense-view exception is MLA-only -- MHA sub-pools stay strided."""
|
||||
"""The per-layer-view exception is MLA-only -- MHA sub-pools stay strided."""
|
||||
for backend in self.DENSE_MLA_BACKENDS:
|
||||
self.assertFalse(
|
||||
_accepts(backend, use_mla=False),
|
||||
@@ -122,6 +128,40 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
|
||||
f"{backend} must stay rejected without --enable-unified-memory",
|
||||
)
|
||||
|
||||
def test_plain_page_major_arm_is_gated_at_boot(self):
|
||||
"""The strided views were removed: --enable-page-major-kv-layout
|
||||
without --enable-unified-memory must be rejected up front for EVERY
|
||||
backend, Triton included, until the per-layer-view reimplementation."""
|
||||
for backend in ("triton",) + self.DENSE_MLA_BACKENDS:
|
||||
for use_mla in (True, False):
|
||||
self.assertFalse(
|
||||
_accepts(backend, use_mla=use_mla, unified=False),
|
||||
f"{backend} must be rejected on the static page-major arm",
|
||||
)
|
||||
|
||||
def test_asymmetric_kv_mha_model_cannot_use_unified_memory(self):
|
||||
"""head_dim != v_head_dim (MiMoV2): no uniform rows, so no per-layer views
|
||||
and no unified pool. The rejection is the POOL's, not a backend's, so
|
||||
it must fire on every backend -- Triton included."""
|
||||
for backend in ("triton",) + self.DENSE_MLA_BACKENDS:
|
||||
self.assertFalse(
|
||||
_accepts(backend, use_mla=False, has_asymmetric_kv=True),
|
||||
f"--enable-unified-memory + {backend} must be rejected for an "
|
||||
"asymmetric-K/V model",
|
||||
)
|
||||
|
||||
def test_asymmetric_dims_do_not_screen_out_mla(self):
|
||||
"""MLA stores one latent row per layer, so its K/V head dims never have
|
||||
to agree -- and real MLA configs report them as unequal (Kimi-Linear:
|
||||
head_dim 72, v_head_dim 128). Screening on `has_asymmetric_kv` alone
|
||||
would lock every one of them out of the unified pool."""
|
||||
for backend in ("triton",) + self.DENSE_MLA_BACKENDS:
|
||||
self.assertTrue(
|
||||
_accepts(backend, use_mla=True, has_asymmetric_kv=True),
|
||||
f"{backend} must stay allowed for an MLA model with asymmetric "
|
||||
"K/V head dims",
|
||||
)
|
||||
|
||||
def test_unwired_backends_always_rejected(self):
|
||||
for backend in self.UNWIRED_BACKENDS:
|
||||
for use_mla in (True, False):
|
||||
@@ -131,15 +171,10 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
|
||||
)
|
||||
|
||||
def test_helion_linear_attention_is_kda_only(self):
|
||||
for unified in (True, False):
|
||||
for phase in ("decode", "prefill"):
|
||||
kwargs = {f"linear_{phase}": "helion"}
|
||||
self.assertTrue(
|
||||
_accepts("triton", use_mla=True, unified=unified, **kwargs)
|
||||
)
|
||||
self.assertFalse(
|
||||
_accepts("triton", use_mla=False, unified=unified, **kwargs)
|
||||
)
|
||||
for phase in ("decode", "prefill"):
|
||||
kwargs = {f"linear_{phase}": "helion"}
|
||||
self.assertTrue(_accepts("triton", use_mla=True, **kwargs))
|
||||
self.assertFalse(_accepts("triton", use_mla=False, **kwargs))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user