[unified-memory] PD disaggregation for every unified pool shape (#37506)
This commit is contained in:
Cheng Wan
@@ -9,6 +9,9 @@ from sglang.test.server_fixtures.disaggregation_fixture import (
|
||||
register_cuda_ci(est_time=236, stage="base-b", runner_config="2-gpu-large")
|
||||
|
||||
KIMI_LINEAR_MODEL = "yujiepan/kimi-linear-tiny-random"
|
||||
# Smallest in-tree GDN hybrid: MHA full attention + gated-delta-net linear
|
||||
# layers, i.e. the unified pool's MHA sub-pool rather than the MLA one.
|
||||
QWEN_GDN_MODEL = "Qwen/Qwen3.5-0.8B"
|
||||
SERVER_ENV = {"SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_DEEPGEMM": "0"}
|
||||
|
||||
# --attention-backend and --enable-deterministic-inference are deliberately
|
||||
@@ -63,5 +66,21 @@ class TestUnifiedMemoryDisaggregationChunkedPrefill(TestUnifiedMemoryDisaggregat
|
||||
extra_decode_args = _chunked_args
|
||||
|
||||
|
||||
class TestUnifiedMemoryDisaggregationMHA(TestUnifiedMemoryDisaggregation):
|
||||
"""The MHA full-attention sub-pool over the wire.
|
||||
|
||||
Kimi-Linear above exercises the MLA sub-pool, whose whole-envelope
|
||||
registration has always been the one PD supports. An MHA envelope is a
|
||||
different shape -- `2 * layer_num` row-blocks per page instead of
|
||||
`layer_num` -- and it reaches a different branch of
|
||||
`_send_kvcache_generic`: without `force_flat` the MHA branch halves the
|
||||
single registered region into K and V, computes `num_kv_layers = 0` and
|
||||
transfers NOTHING, which shows up as garbage decode rather than an error.
|
||||
Logprob parity against a non-PD unified reference is what catches that.
|
||||
"""
|
||||
|
||||
model = QWEN_GDN_MODEL
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,67 @@
|
||||
"""PD disaggregation for a hybrid-SWA model on the unified memory pool.
|
||||
|
||||
A hybrid-SWA model ships TWO attention components: the full-attention KV on the
|
||||
ordinary `kv_data_ptrs` channel and the sliding-window KV as `StateType.SWA`.
|
||||
Under `--enable-unified-memory` both are whole page envelopes into the SAME raw
|
||||
buffer, distinguished only by their per-page stride, and each is addressed by
|
||||
its OWN sub-pool's physical page id -- the full and SWA sides run independent
|
||||
compactions, so one virtual token names two unrelated physical pages.
|
||||
|
||||
That makes three ways to be silently wrong rather than loud:
|
||||
* shipping virtual ids (the base `translate_kv_indices_for_transfer` is the
|
||||
identity, and virtual ids address real bytes);
|
||||
* shipping the SWA side's KERNEL-FACING ids, which the read path uses, in
|
||||
place of its physical ones;
|
||||
* letting compaction relocate a page mid-transfer, which the SWA allocator
|
||||
had no `set_disagg_move_gate` to prevent.
|
||||
|
||||
Logprob parity against a non-PD unified reference catches all three; GSM8K on
|
||||
gpt-oss is too noisy to (single-server unified and static both score 0.570 at
|
||||
200 questions, and PD runs of each span 0.540-0.610).
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kits.pd_parity_kit import PDLogprobParityMixin
|
||||
from sglang.test.server_fixtures.disaggregation_fixture import (
|
||||
PDDisaggregationServerBase,
|
||||
)
|
||||
from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST_MXFP4_WITH_MOE
|
||||
|
||||
register_cuda_ci(est_time=1200, stage="extra-a", runner_config="2-gpu-large")
|
||||
|
||||
UNIFIED_SWA_ARGS = [
|
||||
"--skip-tokenizer-init",
|
||||
"--random-seed",
|
||||
"1",
|
||||
"--enable-unified-memory",
|
||||
# gpt-oss uses attention sinks, which flashinfer does not support; triton
|
||||
# reads both sub-pools' per-layer views.
|
||||
"--attention-backend",
|
||||
"triton",
|
||||
"--mem-fraction-static",
|
||||
"0.7",
|
||||
"--cuda-graph-backend-decode",
|
||||
"disabled",
|
||||
"--cuda-graph-backend-prefill",
|
||||
"disabled",
|
||||
]
|
||||
|
||||
|
||||
class TestUnifiedMemoryDisaggregationSWA(
|
||||
PDLogprobParityMixin, PDDisaggregationServerBase
|
||||
):
|
||||
"""1 prefill + 1 decode, both unified, vs a non-PD unified reference."""
|
||||
|
||||
model = DEFAULT_MODEL_NAME_FOR_TEST_MXFP4_WITH_MOE
|
||||
prefill_tp_size = 1
|
||||
decode_tp_size = 1
|
||||
decode_base_gpu_id = 1
|
||||
baseline_args = UNIFIED_SWA_ARGS
|
||||
extra_prefill_args = UNIFIED_SWA_ARGS
|
||||
extra_decode_args = UNIFIED_SWA_ARGS
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,86 @@
|
||||
"""PD disaggregation for a TRI-pool model on the unified memory pool.
|
||||
|
||||
Inkling is the only in-tree architecture that is both mambaish and hybrid-SWA,
|
||||
so one unified buffer carries three components with three independent
|
||||
compactions -- ``[conv state (up END) | swa (FLOAT) | full (down END)]`` -- and
|
||||
PD must ship all three per request: full KV on the ``kv_data_ptrs`` channel,
|
||||
sliding-window KV as ``StateType.SWA``, ShortConv state as ``StateType.MAMBA``
|
||||
(via the ``req_to_token_pool`` fallback, since the KV pool here is a
|
||||
``UnifiedSWAKVPool`` rather than a ``HybridLinearKVPool``).
|
||||
|
||||
Two failures this pins that the 2-pool cases cannot:
|
||||
|
||||
* the FLOAT sub-pool moves for reasons neither END does, so a move gate that
|
||||
reaches only full and swa still lets a conv slot relocate under an
|
||||
in-flight state transfer;
|
||||
* ``page_size > 1`` turns on the decode node's SWA-tail prealloc, whose
|
||||
static body allocates the swa side independently -- an assertion failure
|
||||
against this composite's single virtual id space, and, once that is
|
||||
handled, the first path that can bind the WRONG swa pages.
|
||||
|
||||
Logprob parity against a non-PD unified reference is the check: the tiny
|
||||
``test`` revision is undertrained, so answer quality carries no signal, but a
|
||||
dropped or misaddressed component moves logprobs immediately.
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kits.pd_parity_kit import PDLogprobParityMixin
|
||||
from sglang.test.server_fixtures.disaggregation_fixture import (
|
||||
PDDisaggregationServerBase,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=900, stage="extra-a", runner_config="2-gpu-large")
|
||||
|
||||
_MODEL_PATH = os.environ.get("INKLING_TEST_MODEL_PATH", "thinkingmachines/Inkling")
|
||||
_MODEL_REVISION = os.environ.get("INKLING_TEST_MODEL_REVISION", "test")
|
||||
|
||||
# The unified radix tree is what merges the three components into one tree.
|
||||
SERVER_ENV = {"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}
|
||||
|
||||
UNIFIED_TRI_ARGS = [
|
||||
"--skip-tokenizer-init",
|
||||
"--random-seed",
|
||||
"1",
|
||||
"--enable-unified-memory",
|
||||
# Unified requires the Triton strided page-major read/write paths.
|
||||
"--attention-backend",
|
||||
"triton",
|
||||
"--page-size",
|
||||
"128",
|
||||
"--mamba-radix-cache-strategy",
|
||||
"extra_buffer",
|
||||
"--swa-full-tokens-ratio",
|
||||
"0.1",
|
||||
"--mamba-full-memory-ratio",
|
||||
"0.1",
|
||||
"--mem-fraction-static",
|
||||
"0.5",
|
||||
# Inkling defaults to a FULL prefill graph, which unified rejects at boot.
|
||||
"--cuda-graph-backend-prefill",
|
||||
"disabled",
|
||||
"--revision",
|
||||
_MODEL_REVISION,
|
||||
]
|
||||
|
||||
|
||||
class TestUnifiedMemoryDisaggregationTriPool(
|
||||
PDLogprobParityMixin, PDDisaggregationServerBase
|
||||
):
|
||||
"""1 prefill + 1 decode, both unified, vs a non-PD unified reference."""
|
||||
|
||||
model = _MODEL_PATH
|
||||
extra_prefill_env = SERVER_ENV
|
||||
extra_decode_env = SERVER_ENV
|
||||
prefill_tp_size = 1
|
||||
decode_tp_size = 1
|
||||
decode_base_gpu_id = 1
|
||||
baseline_args = UNIFIED_TRI_ARGS
|
||||
extra_prefill_args = UNIFIED_TRI_ARGS
|
||||
extra_decode_args = UNIFIED_TRI_ARGS
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -186,5 +186,210 @@ class TestMoveGateRejectsNonPdNode(CustomTestCase):
|
||||
unified_memory_disagg_move_gate(scheduler)
|
||||
|
||||
|
||||
class TestUnifiedAllocatorsPublishTheTransferContract(CustomTestCase):
|
||||
"""Every unified composite allocator must OVERRIDE the two PD hooks.
|
||||
|
||||
`BaseTokenToKVPoolAllocator.translate_kv_indices_for_transfer` is the
|
||||
IDENTITY, and `set_disagg_move_gate` exists only where a composite defines
|
||||
it. Inheriting either is silent, not loud: identity puts VIRTUAL ids on the
|
||||
wire (they address real bytes, so the peer gets plausible garbage), and a
|
||||
missing gate lets lazy compaction relocate pages under in-flight RDMA.
|
||||
An AST-level check because instantiating these composites needs a GPU.
|
||||
"""
|
||||
|
||||
# Composites that own the full-side virtual ids and so must define the
|
||||
# transfer translate themselves.
|
||||
_COMPOSITES = (
|
||||
"UnifiedMambaTokenToKVPoolAllocator",
|
||||
"UnifiedSWATokenToKVPoolAllocator",
|
||||
)
|
||||
# Every composite must define the gate setter, including the tri-pool,
|
||||
# which inherits the SWA translates (same full side) but has a THIRD
|
||||
# member the 2-pool setter does not reach.
|
||||
_GATE_COMPOSITES = _COMPOSITES + ("UnifiedMambaSWATokenToKVPoolAllocator",)
|
||||
|
||||
@staticmethod
|
||||
def _own_methods(cls_name: str) -> Set[str]:
|
||||
"""Names this class defines ITSELF, inheritance excluded.
|
||||
|
||||
Resolved off the class object rather than by parsing a named module:
|
||||
these composites have already been moved once (out of
|
||||
`multi_ended_allocator` into `allocator/unified_*`), and a hardcoded
|
||||
module path turns that kind of move into a test failure that says
|
||||
nothing about the contract. `__dict__` needs no GPU -- it is the class
|
||||
body, not an instance.
|
||||
"""
|
||||
from sglang.srt.mem_cache.allocator import (
|
||||
unified_hybrid_swa,
|
||||
unified_mamba,
|
||||
)
|
||||
|
||||
for mod in (unified_mamba, unified_hybrid_swa):
|
||||
cls = getattr(mod, cls_name, None)
|
||||
if cls is not None:
|
||||
return set(vars(cls))
|
||||
raise AssertionError(f"class {cls_name} not found in the unified allocators")
|
||||
|
||||
def test_transfer_translate_is_not_inherited_identity(self):
|
||||
for name in self._COMPOSITES:
|
||||
with self.subTest(composite=name):
|
||||
self.assertIn(
|
||||
"translate_kv_indices_for_transfer",
|
||||
self._own_methods(name),
|
||||
f"{name} inherits the identity transfer translate; PD would "
|
||||
"ship VIRTUAL ids and corrupt KV without any error",
|
||||
)
|
||||
|
||||
# Every sub-allocator attribute a composite can hold. The stub carries all
|
||||
# of them regardless of composite, so the assertion is on what installation
|
||||
# REACHES rather than on what the stub was given.
|
||||
_MEMBER_ATTRS = ("full_attn_allocator", "swa_attn_allocator", "mamba_allocator")
|
||||
|
||||
# The members each composite's gate must reach. The tri-pool row is the one
|
||||
# that matters: it inherits the setter, so an enumeration written inside
|
||||
# that setter would silently leave the third member ungated.
|
||||
_EXPECTED_COVERAGE = {
|
||||
"UnifiedMambaTokenToKVPoolAllocator": {
|
||||
"full_attn_allocator",
|
||||
"mamba_allocator",
|
||||
},
|
||||
"UnifiedSWATokenToKVPoolAllocator": {
|
||||
"full_attn_allocator",
|
||||
"swa_attn_allocator",
|
||||
},
|
||||
"UnifiedMambaSWATokenToKVPoolAllocator": {
|
||||
"full_attn_allocator",
|
||||
"swa_attn_allocator",
|
||||
"mamba_allocator",
|
||||
},
|
||||
}
|
||||
|
||||
def _members_reached(self, cls_name: str, slot: str) -> Set[str]:
|
||||
"""Install one gate on a stub composite and report which members got it.
|
||||
|
||||
`object.__new__` skips `__init__` (which needs a GPU); the setter reads
|
||||
only `lazy_compaction` and the member attributes.
|
||||
"""
|
||||
from sglang.srt.mem_cache.allocator import unified_hybrid_swa, unified_mamba
|
||||
|
||||
cls = getattr(unified_mamba, cls_name, None) or getattr(
|
||||
unified_hybrid_swa, cls_name
|
||||
)
|
||||
alloc = object.__new__(cls)
|
||||
alloc.lazy_compaction = True
|
||||
for attr in self._MEMBER_ATTRS:
|
||||
member = type("_Member", (), {})()
|
||||
member.disagg_move_gate = None
|
||||
member.host_transfer_move_gate = None
|
||||
setattr(alloc, attr, member)
|
||||
|
||||
def gate() -> bool:
|
||||
return True
|
||||
|
||||
alloc.set_disagg_move_gate(gate)
|
||||
return {
|
||||
attr
|
||||
for attr in self._MEMBER_ATTRS
|
||||
if getattr(getattr(alloc, attr), slot) is gate
|
||||
}
|
||||
|
||||
def test_the_gate_reaches_every_member(self):
|
||||
"""A gate that reaches only some members is not a weaker gate, it is no
|
||||
gate: the ungated end relocates its own pages under the very transfer
|
||||
the gate was installed for.
|
||||
"""
|
||||
for name, expected in self._EXPECTED_COVERAGE.items():
|
||||
with self.subTest(composite=name):
|
||||
self.assertEqual(
|
||||
self._members_reached(name, "disagg_move_gate"),
|
||||
expected,
|
||||
f"{name}.disagg_move_gate does not cover every member",
|
||||
)
|
||||
|
||||
def test_gate_setters_do_not_enumerate_members_themselves(self):
|
||||
"""The structural half of the rule above: a setter that names its
|
||||
members is one a new member silently escapes. Installation must go
|
||||
through the shared helper, which drives off `_move_gate_targets`.
|
||||
"""
|
||||
import inspect
|
||||
|
||||
from sglang.srt.mem_cache.allocator import unified_hybrid_swa, unified_mamba
|
||||
|
||||
for name in self._EXPECTED_COVERAGE:
|
||||
cls = getattr(unified_mamba, name, None) or getattr(
|
||||
unified_hybrid_swa, name
|
||||
)
|
||||
if "set_disagg_move_gate" not in vars(cls):
|
||||
continue # inherited, and the inherited one is checked above
|
||||
with self.subTest(composite=name):
|
||||
body = inspect.getsource(cls.set_disagg_move_gate)
|
||||
self.assertIn("install_move_gate", body)
|
||||
self.assertNotIn("_move_gate = ", body)
|
||||
|
||||
def test_swa_composite_translates_the_swa_side_separately(self):
|
||||
"""The SWA sub-pool runs its OWN compaction, so a full-side physical id
|
||||
does not name the SWA page holding the same virtual token. The read-path
|
||||
`translate_loc_from_full_to_swa` cannot stand in either: it returns
|
||||
kernel-facing ids, and the transfer addresses raw page envelopes."""
|
||||
self.assertIn(
|
||||
"translate_swa_indices_for_transfer",
|
||||
self._own_methods("UnifiedSWATokenToKVPoolAllocator"),
|
||||
)
|
||||
|
||||
|
||||
class TestEverySwaAllocatorAnswersTheTransferTranslate(CustomTestCase):
|
||||
"""Any allocator with a full->SWA read translate needs the transfer sibling.
|
||||
|
||||
`_swa_payload` on both PD sides calls
|
||||
`translate_swa_indices_for_transfer` on whatever allocator the scheduler
|
||||
holds. Most get it by inheriting `SWATokenToKVPoolAllocator`, but a
|
||||
composite that merely DELEGATES the read translate (the DSV4 HiSparse
|
||||
allocator derives from `BaseTokenToKVPoolAllocator`) inherits neither the
|
||||
default nor an override, and PD aborts with an AttributeError the moment a
|
||||
sliding-window payload is built.
|
||||
|
||||
Derived from the live class tree rather than a hand-kept list: a list would
|
||||
pass forever the day someone adds the next delegating composite.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _allocator_classes():
|
||||
import importlib
|
||||
import inspect
|
||||
import pkgutil
|
||||
|
||||
import sglang.srt.mem_cache.allocator as pkg
|
||||
from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator
|
||||
|
||||
found = {}
|
||||
for mod_info in pkgutil.iter_modules(pkg.__path__):
|
||||
try:
|
||||
mod = importlib.import_module(
|
||||
f"sglang.srt.mem_cache.allocator.{mod_info.name}"
|
||||
)
|
||||
except Exception:
|
||||
continue # optional backends need hardware this runner may lack
|
||||
for _, cls in inspect.getmembers(mod, inspect.isclass):
|
||||
if issubclass(cls, BaseTokenToKVPoolAllocator):
|
||||
found[cls.__name__] = cls
|
||||
return found
|
||||
|
||||
def test_read_translate_implies_transfer_translate(self):
|
||||
classes = self._allocator_classes()
|
||||
# Guard the guard: an import failure that empties this set would make
|
||||
# the assertion below vacuous.
|
||||
self.assertIn("SWATokenToKVPoolAllocator", classes)
|
||||
for name, cls in sorted(classes.items()):
|
||||
if not hasattr(cls, "translate_loc_from_full_to_swa"):
|
||||
continue
|
||||
with self.subTest(allocator=name):
|
||||
self.assertTrue(
|
||||
hasattr(cls, "translate_swa_indices_for_transfer"),
|
||||
f"{name} translates full->SWA for reads but cannot answer "
|
||||
"translate_swa_indices_for_transfer; PD's _swa_payload "
|
||||
"calls it on whatever allocator the scheduler holds",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -18,9 +18,11 @@ import unittest
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.layout.page_major import (
|
||||
build_mha_views,
|
||||
build_mla_views,
|
||||
build_page_major_mamba_views,
|
||||
mamba_entry_bytes,
|
||||
mha_entry_bytes,
|
||||
mla_entry_bytes,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
@@ -74,6 +76,118 @@ class TestMLAEnvelopeTransferAddressing(CustomTestCase):
|
||||
self.assertTrue(torch.equal(got, val), (page, layer, off))
|
||||
|
||||
|
||||
class TestMHAEnvelopeTransferAddressing(CustomTestCase):
|
||||
"""The MHA counterpart of the MLA case above.
|
||||
|
||||
An MHA page envelope holds ``2 * layer_num`` row-blocks (layer l's K at
|
||||
block 2l, its V at 2l+1). PD ships that whole envelope as one item, so a
|
||||
row written through ANY per-layer view must land inside its own page's
|
||||
``page_envelope_bytes`` block -- otherwise the transfer would carry a
|
||||
page's K but another page's V and every kernel would still read fine
|
||||
locally.
|
||||
"""
|
||||
|
||||
def test_page_envelope_matches_per_layer_views(self):
|
||||
layer_num, page_size, head_num, head_dim, num_pages = 3, 4, 2, 8, 6
|
||||
store_dtype = torch.bfloat16
|
||||
entry_bytes = mha_entry_bytes(
|
||||
layer_num=layer_num,
|
||||
head_num=head_num,
|
||||
head_dim=head_dim,
|
||||
v_head_dim=head_dim,
|
||||
itemsize=store_dtype.itemsize,
|
||||
)
|
||||
page_bytes = page_size * entry_bytes
|
||||
row_bytes = head_num * head_dim * store_dtype.itemsize
|
||||
self.assertEqual(page_bytes, page_size * 2 * layer_num * row_bytes)
|
||||
|
||||
# One page envelope of tail pad, as UnifiedKVPool allocates for MHA.
|
||||
raw = torch.zeros((num_pages + 1) * page_bytes, dtype=torch.uint8)
|
||||
k_views, v_views = build_mha_views(
|
||||
raw,
|
||||
layer_num=layer_num,
|
||||
head_num=head_num,
|
||||
head_dim=head_dim,
|
||||
v_head_dim=head_dim,
|
||||
store_dtype=store_dtype,
|
||||
page_size=page_size,
|
||||
num_pages=num_pages,
|
||||
anchor_bytes=0,
|
||||
)
|
||||
|
||||
blocks = 2 * layer_num
|
||||
for page in range(num_pages):
|
||||
for layer in range(layer_num):
|
||||
for is_v, views in ((0, k_views), (1, v_views)):
|
||||
for pos in range(page_size):
|
||||
row = page * blocks * page_size + pos
|
||||
views[layer][row].fill_(1)
|
||||
(nz,) = torch.nonzero(raw, as_tuple=True)
|
||||
lo, hi = int(nz.min()), int(nz.max())
|
||||
self.assertGreaterEqual(
|
||||
lo,
|
||||
page * page_bytes,
|
||||
f"page={page} layer={layer} v={is_v} pos={pos} "
|
||||
"wrote below its page envelope",
|
||||
)
|
||||
self.assertLess(
|
||||
hi,
|
||||
(page + 1) * page_bytes,
|
||||
f"page={page} layer={layer} v={is_v} pos={pos} "
|
||||
"wrote past its page envelope",
|
||||
)
|
||||
views[layer][row].zero_()
|
||||
|
||||
def test_envelope_move_is_a_whole_page_copy(self):
|
||||
"""Relocating a page envelope must move every layer's K and V with it;
|
||||
this is what `UnifiedMHATokenToKVPool.move_kv_cache` relies on and what
|
||||
makes a physical page id a valid PD transfer index after compaction."""
|
||||
layer_num, page_size, head_num, head_dim, num_pages = 2, 2, 1, 4, 4
|
||||
store_dtype = torch.bfloat16
|
||||
entry_bytes = mha_entry_bytes(
|
||||
layer_num=layer_num,
|
||||
head_num=head_num,
|
||||
head_dim=head_dim,
|
||||
v_head_dim=head_dim,
|
||||
itemsize=store_dtype.itemsize,
|
||||
)
|
||||
page_bytes = page_size * entry_bytes
|
||||
raw = torch.zeros((num_pages + 1) * page_bytes, dtype=torch.uint8)
|
||||
k_views, v_views = build_mha_views(
|
||||
raw,
|
||||
layer_num=layer_num,
|
||||
head_num=head_num,
|
||||
head_dim=head_dim,
|
||||
v_head_dim=head_dim,
|
||||
store_dtype=store_dtype,
|
||||
page_size=page_size,
|
||||
num_pages=num_pages,
|
||||
anchor_bytes=0,
|
||||
)
|
||||
blocks = 2 * layer_num
|
||||
# Distinct content in source page 1, every layer, K and V.
|
||||
for layer in range(layer_num):
|
||||
for pos in range(page_size):
|
||||
row = 1 * blocks * page_size + pos
|
||||
k_views[layer][row].fill_(layer + 1)
|
||||
v_views[layer][row].fill_(-(layer + 1))
|
||||
|
||||
env = raw[: num_pages * page_bytes].view(num_pages, page_bytes)
|
||||
env[3] = env[1]
|
||||
|
||||
for layer in range(layer_num):
|
||||
for pos in range(page_size):
|
||||
row = 3 * blocks * page_size + pos
|
||||
self.assertTrue(
|
||||
torch.all(k_views[layer][row] == layer + 1),
|
||||
f"K layer {layer} did not ride the envelope move",
|
||||
)
|
||||
self.assertTrue(
|
||||
torch.all(v_views[layer][row] == -(layer + 1)),
|
||||
f"V layer {layer} did not ride the envelope move",
|
||||
)
|
||||
|
||||
|
||||
class TestMambaEnvelopeTransferAddressing(CustomTestCase):
|
||||
def test_slot_envelope_is_self_contained(self):
|
||||
"""A slot's conv+temporal state for all layers must live exactly in
|
||||
|
||||
@@ -345,11 +345,12 @@ class TestUnifiedMHATokenToKVPool(unittest.TestCase):
|
||||
)
|
||||
|
||||
def test_transfer_entry_points_fail_loud(self):
|
||||
"""PD / CPU-copy entry points assume per-layer buffers indexed by TOKEN
|
||||
id and would silently mis-index the row space, so each must raise."""
|
||||
"""The entry points that assume per-layer buffers indexed by TOKEN id
|
||||
would silently mis-index against the row space (or hit a missing-attr
|
||||
AttributeError), so each must raise. `get_contiguous_buf_infos` is NOT
|
||||
among them: PD addresses this pool as whole page envelopes, pinned by
|
||||
`test_pd_registration_is_one_whole_envelope` below."""
|
||||
_, 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):
|
||||
@@ -357,6 +358,24 @@ class TestUnifiedMHATokenToKVPool(unittest.TestCase):
|
||||
with self.assertRaises(NotImplementedError):
|
||||
pool.set_kv_buffer_prefix_valid()
|
||||
|
||||
def test_pd_registration_is_one_whole_envelope(self):
|
||||
"""PD registers ONE region -- the whole raw buffer -- with the page
|
||||
envelope as the item, so the transfer engine addresses it as
|
||||
`raw_ptr + physical_page * page_envelope_bytes`. Per-layer regions
|
||||
would be wrong here: the per-layer views overlap inside the envelope
|
||||
and index in kernel-facing ids, not token ids."""
|
||||
kv, pool = _make_pool_and_kv(1)
|
||||
ptrs, lens, item_lens = pool.get_contiguous_buf_infos()
|
||||
self.assertEqual(len(ptrs), 1)
|
||||
self.assertEqual(len(lens), 1)
|
||||
self.assertEqual(len(item_lens), 1)
|
||||
self.assertEqual(ptrs[0], kv._raw.data_ptr())
|
||||
self.assertEqual(lens[0], kv._raw.numel())
|
||||
self.assertEqual(item_lens[0], pool._page_bytes)
|
||||
# The whole addressable page range must fit the registered region, or
|
||||
# the last page's write would run off the end of the RDMA mapping.
|
||||
self.assertLessEqual(pool._num_pages * item_lens[0], lens[0])
|
||||
|
||||
def test_hnd_env_cannot_hijack_layout(self):
|
||||
"""SGLANG_USE_HND_KVCACHE must not flip this pool's layout: HND indexes
|
||||
4-D while the per-layer views are 3-D, so the pinned label has to win."""
|
||||
|
||||
Reference in New Issue
Block a user