minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

feat(unified-memory): read unified pool from attention backends fa3/flashinfer/trtllm_mha/flashmla (#34613)

Browse Source 
Co-authored-by: Caihua Li <caihua.li@bytedance.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>
This commit is contained in:
caihuali95
2026-08-30 23:58:24 -07:00
committed by GitHub
co-authored by Caihua Li Claude Fable 5 Cheng Wan
parent 29578d5578
commit 8bb776dc48
31 changed files with 1182 additions and 757 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_full_loc_fast_path.py
+38 -38
View File
@@ -59,9 +59,11 @@ class _RecordingPool:
class TestUnifiedSWARouting(unittest.TestCase):
"""`UnifiedSWAKVPool.set_kv_buffer` routing: full layers write the full-physical
`full_loc`; SWA layers write the swa-physical `swa_loc`. Both come from the
write metadata; the pool never translates."""
"""`UnifiedSWAKVPool.set_kv_buffer` routing: full layers write `full_loc`
when present (triton's capture-stable buffer), else the rebound generic
`loc` -- the same id space once the loc is rebound; SWA layers write the swa-physical
`swa_loc`, which has no fallback (a different id space). The pool never
translates."""
def _make_bare_pool(self):
from sglang.srt.mem_cache.unified_memory_pool import UnifiedSWAKVPool
@@ -96,21 +98,29 @@ class TestUnifiedSWARouting(unittest.TestCase):
self.assertIsNot(forwarded, virtual_loc)
self.assertNotIn("already_physical", kwargs)
def test_full_layer_requires_full_loc(self):
def test_full_layer_falls_back_to_generic_loc(self):
"""Bug regression: fa3 x unified-SWA crashed at gpt-oss
cuda-graph capture because every backend except triton bundles the
2-arg KVWriteLoc(loc, swa) and the full-layer door demanded an explicit
full_loc. Once the loc is rebound the generic `loc` IS the full-side kernel-facing id
(rebind_write_loc runs at ForwardBatch construction),
so the door must fall back to it -- the pool still never translates."""
pool = self._make_bare_pool()
virtual_loc = torch.tensor([10, 11, 12], dtype=torch.int64)
rebound_loc = torch.tensor([10, 11, 12], dtype=torch.int64)
swa_phys = torch.tensor([1, 2, 0], dtype=torch.int64)
layer = types.SimpleNamespace(layer_id=0)
# No full_loc precomputed -> fail loud (the unified memory pool must precompute
# out_cache_loc_full_physical) rather than write a virtual loc as physical.
with self.assertRaises(AssertionError):
pool.set_kv_buffer(
layer,
_loc_info(virtual_loc, swa_phys),
torch.zeros(3, 4, 8),
torch.zeros(3, 4, 8),
)
pool.set_kv_buffer(
layer,
_loc_info(rebound_loc, swa_phys),
torch.zeros(3, 4, 8),
torch.zeros(3, 4, 8),
)
self.assertEqual(len(pool.full_kv_pool.calls), 1)
forwarded, kwargs = pool.full_kv_pool.calls[0]
self.assertIs(forwarded, rebound_loc)
self.assertNotIn("already_physical", kwargs)
def test_swa_layer_writes_swa_loc(self):
pool = self._make_bare_pool()
@@ -264,19 +274,17 @@ class TestHybridLinearMLARouting(unittest.TestCase):
- `set_kv_buffer` (MLA branch) mirrors the MHA branch — write the
pre-translated `KVWriteLoc.full_loc` when present (unified pool, where it
carries the DENSE loc), else the raw `loc` (static pool, already physical).
- `set_mla_kv_buffer` forwards `loc` untouched (kernel-facing since the
ForwardBatch rebind); `get_mla_kv_buffer` applies `_full_translate`
exactly once (its indices are req_to_token-produced, virtual under the
unified pool)."""
- `set_mla_kv_buffer` / `get_mla_kv_buffer` forward `loc` untouched:
writes are kernel-facing since the ForwardBatch rebind, and read
indices are translated at their production sites."""
def _make_bare_pool(self, translate=None):
def _make_bare_pool(self):
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
pool = object.__new__(HybridLinearKVPool)
pool.full_kv_pool = _RecordingMLAPool()
pool.use_mla = True
pool.full_attention_layer_id_mapping = {0: 0}
pool._full_translate = translate if translate is not None else (lambda ids: ids)
return pool
def test_mla_writes_full_loc_from_write_loc(self):
@@ -328,28 +336,20 @@ class TestHybridLinearMLARouting(unittest.TestCase):
self.assertEqual(len(pool.full_kv_pool.mla_set_calls), 1)
self.assertIs(pool.full_kv_pool.mla_set_calls[0], loc)
def test_get_mla_kv_buffer_translates_exactly_once(self):
"""READ door: `loc` is produced from req_to_token (VIRTUAL under the
unified pool), so the get side still translates here — exactly once.
The WRITE door (case above) never translates: the split is the write
flip's contract."""
calls = []
def translate(ids):
calls.append(ids)
return ids + 100
pool = self._make_bare_pool(translate=translate)
virtual_loc = torch.tensor([4, 5], dtype=torch.int64)
def test_get_mla_kv_buffer_door_never_translates(self):
"""Kernel-facing contract, read side: `loc` is a read-index tensor
already translated at its production site
(fetch_mha_one_shot_kv_indices / prepare_chunked_kv_indices); the
door forwards it UNTOUCHED -- a re-added door translate would
double-translate every unified MLA prefix read."""
pool = self._make_bare_pool()
loc = torch.tensor([104, 105], dtype=torch.int64)
layer = types.SimpleNamespace(layer_id=0)
pool.get_mla_kv_buffer(layer, virtual_loc)
pool.get_mla_kv_buffer(layer, loc)
self.assertEqual(len(calls), 1)
self.assertEqual(len(pool.full_kv_pool.mla_get_calls), 1)
self.assertTrue(
torch.all(pool.full_kv_pool.mla_get_calls[0] == virtual_loc + 100)
)
self.assertIs(pool.full_kv_pool.mla_get_calls[0], loc)
if __name__ == "__main__":
test/registered/unit/mem_cache/test_multi_ended_allocator.py
+76
View File
@@ -31,11 +31,13 @@ import torch
from sglang.srt.mem_cache.multi_ended_allocator import (
MultiEndedAllocator,
UnifiedMambaTokenToKVPoolAllocator,
UnifiedSWATokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.unified_memory_pool import (
MambaSubPoolSpec,
MHASubPoolSpec,
MLASubPoolSpec,
UnifiedKVPool,
)
@@ -2746,5 +2748,79 @@ class TestSWACompositeDenseSurface(unittest.TestCase):
self.assertTrue(bool((got[in_tomb] == 0).all().item()))
class TestPs64MLACompositeFeasibility(unittest.TestCase):
"""The Kimi/flashmla shape: MLA + mamba composite at page_size=64 (the
flashmla arg snap). Large pages stress every sizing derivation at once —
the 64-token sink-page floor, the ps*entry_bytes per-layer-view tail pad, and
the page-granular alloc — so this pins that the factory-shaped
construction stays FEASIBLE and the dense surface stays on-formula when
the page size jumps from the usual 1..4 to 64."""
PS = 64
LAYERS = 3
def _build(self):
full = MLASubPoolSpec(
name="full",
layer_num=self.LAYERS,
kv_lora_rank=64,
qk_rope_head_dim=16,
store_dtype=torch.float16,
grow_direction="down",
)
mamba = MambaSubPoolSpec(
name="mamba",
layer_num=2,
conv_state_shapes=((8, 16),),
conv_dtype=torch.bfloat16,
temporal_state_shape=(4, 8, 8),
temporal_dtype=torch.float32,
grow_direction="up",
)
n_full = 8 * self.PS # 8 pages incl. the sink page
total = n_full * full.entry_bytes() + 16 * mamba.entry_bytes()
pool = UnifiedKVPool(
total_bytes=total,
sub_pool_specs=[full, mamba],
device=_DEV,
enable_memory_saver=False,
page_size=self.PS,
)
full_kv = _FakeKVCache(pool.max_slots("full"))
full_kv.attach_allocator = lambda allocator: None
mamba_kv = _FakeKVCache(pool.max_slots("mamba"))
mamba_kv.attach_allocator = lambda allocator: None
mamba_kv._copy_from_physical = lambda src, dst: None
class _FakeHybridLinearKVPool:
full_kv_pool = full_kv
mamba_pool = mamba_kv
return UnifiedMambaTokenToKVPoolAllocator(
unified_buffer=pool,
kvcache=_FakeHybridLinearKVPool(),
device=_DEV,
page_size=self.PS,
need_sort=False,
forward_stream=None,
)
def test_construction_alloc_and_dense_formula(self):
a = self._build()
# MLA: one latent row per layer, so the spec reports LAYERS blocks.
self.assertEqual(a.kernel_page_multiplier, self.LAYERS)
v = a.alloc(2 * self.PS)
self.assertIsNotNone(v, "2-page alloc infeasible at ps=64")
# Page-aligned virtual run (page-granular allocator invariant).
self.assertEqual(int(v[0].item()) % self.PS, 0)
# Dense translate follows the affine formula at ps=64, and every id
# fits int32 (the canonical narrows on store).
v2p = a.full_v2p_page_table
want = v2p[v // self.PS] * (self.PS * self.LAYERS) + v % self.PS
got = a.translate_kv_loc_for_kernel(v)
self.assertTrue(torch.equal(got, want), "kernel-facing formula broke at ps=64")
self.assertTrue(bool((got < 2**31).all().item()))
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_unified_mla_block_table.py
+106 -227
View File
@@ -20,20 +20,26 @@ 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_for_kernel` over the
token-level kv_indices;
- `normal_decode_set_metadata` in-kernel, for fa3's captured-decode page table.
Since the read-path translator, ONE builder computes that formula for every
family — `build_kv_read_table` (the canonical) — and the backends only
differ in how they consume it:
- trtllm_mla / cutedsl_mla / tokenspeed_mla / flashmla: rows filled straight
into their padded block tables (`KVIndexTranslator.build_into`, prefix-only so
the backends' own -1 / stale tail sentinels survive);
- the flashinfer updaters: token ids reconstructed from the canonical by
`create_flashinfer_kv_indices_triton[ENTRY_PAGE_SIZE=ps]`;
- fa3's captured decode: `normal_decode_set_metadata` copies the canonical
rows' live prefixes (src_is_read_table=True).
Covered here:
- kernel identity: `v2p_ptr=None, PAGE_MULT=1` is byte-identical to main;
- 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 kernel-facing translate the flashinfer updaters apply agrees with the
page-level block table the trtllm path builds;
- the static `create_flashmla_kv_indices_triton` (no id-space knowledge left)
still matches the plain token//ps reference;
- the canonical route against the python dense reference, for several page
sizes, ragged sequence lengths and a non-identity v2p permutation;
- lanes past a row's live prefix keep the backend's -1 sentinel (prefix-only
discipline — the trtllm/flashmla tail contract);
- the token-level kernel-facing translate the flashinfer updaters used to apply
agrees with the canonical page table (page-affinity of the id space);
- 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.
@@ -57,6 +63,28 @@ _LAYERS = 24 # K3 MLA full-attention layer count
def _fill_block_table(
req_to_token, req_pool_indices, seq_lens, page_size, *, v2p, mult
):
"""The unified route: canonical builder into a -1-filled block table
(exactly what KVIndexTranslator.build_into does for trtllm_mla/flashmla)."""
from sglang.kernels.ops.kvcache.kv_read_table import build_kv_read_table
bs = req_pool_indices.shape[0]
max_blocks = (int(seq_lens.max().item()) + page_size - 1) // page_size
out = torch.full((bs, max_blocks), -1, dtype=torch.int32, device=_DEV)
build_kv_read_table(
req_to_token=req_to_token,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens.to(torch.int64),
v2p=v2p,
multiplier=mult,
page_size=page_size,
max_pages=max_blocks,
out=out,
)
return out
def _fill_block_table_static(req_to_token, req_pool_indices, seq_lens, page_size):
"""The static-pool route: the stripped flashmla kernel, token//ps verbatim."""
from sglang.kernels.ops.kvcache.kv_indices import (
create_flashmla_kv_indices_triton,
get_num_kv_index_blocks_flashmla,
@@ -75,8 +103,6 @@ def _fill_block_table(
req_to_token.stride(0),
max_blocks,
PAGED_SIZE=page_size,
v2p_ptr=v2p,
PAGE_MULT=mult,
)
return out
@@ -122,11 +148,12 @@ class TestDenseBlockTable(unittest.TestCase):
v2p[0] = 0 # page 0 is the reserved sink
return req_to_token, req_pool_indices, seq_lens, v2p
def test_identity_when_hooks_absent(self):
"""v2p_ptr=None / PAGE_MULT=1 must reproduce the pre-change behaviour."""
def test_static_kernel_matches_reference(self):
"""The stripped (id-space-free) flashmla kernel is byte-identical to the
plain token//ps reference -- guards the v2p-arg removal itself."""
for page_size in (1, 32, 64):
rt, rpi, sl, _ = self._make_batch(page_size)
got = _fill_block_table(rt, rpi, sl, page_size, v2p=None, mult=1)
got = _fill_block_table_static(rt, rpi, sl, page_size)
want = _reference(rt, rpi, sl, page_size, v2p=None, mult=1)
self.assertTrue(
torch.equal(got.long(), want), f"page_size={page_size}: {got} != {want}"
@@ -166,8 +193,9 @@ class TestDenseBlockTable(unittest.TestCase):
)
def test_padded_lanes_stay_untouched(self):
"""Lanes past a request's page count keep the -1 fill: the masked v2p load
must not write a translated value (nor read out of bounds)."""
"""Lanes past a request's page count keep the -1 fill: the prefix-only
canonical build must never write a backend's tail sentinel (the
trtllm/flashmla block-table contract)."""
page_size = 64
rt, rpi, sl, v2p = self._make_batch(page_size)
got = _fill_block_table(rt, rpi, sl, page_size, v2p=v2p, mult=_LAYERS)
@@ -207,43 +235,77 @@ class TestDenseBlockTable(unittest.TestCase):
@unittest.skipUnless(_HAS_CUDA, "requires CUDA")
class TestFa3MetadataDenseBlockTable(unittest.TestCase):
"""fa3 folds the unified remap into `normal_decode_set_metadata`, the fused
gather that writes its captured-decode page table, so the kernel itself has
to get the mapping right. Two kernels back it: a page_size == 1 / no-SWA fast
path (what Kimi-Linear takes, since fa3 imposes no page-size constraint) and
a general one.
"""fa3's captured-decode page table is written by `normal_decode_set_metadata`
fed with the translator's read table kernel page table
(src_is_read_table=True): the fused kernel copies the canonical
rows' live prefixes into the capture-stable buffer. Pinned END-TO-END:
build_kv_read_table -> wrapper -> page_table must equal the python
reference of the kernel-facing formula, on both the page_size == 1 / no-SWA fast
path (what Kimi-Linear takes) and the general kernel. The static call
(no source flag) stays byte-identical to the pre-translator kernel.
"""
def _run(self, page_size, *, v2p, mult, bs=5, max_ctx=2048):
from sglang.kernels.ops.attention.metadata import normal_decode_set_metadata
from sglang.kernels.ops.kvcache.kv_read_table import (
build_kv_read_table,
)
maker = TestDenseBlockTable._make_batch
rt, rpi, sl, v2p_full = maker(self, page_size, bs=bs, max_ctx=max_ctx)
v2p_arg = v2p_full if v2p else None
max_pages = (max_ctx + page_size - 1) // page_size
page_table = torch.zeros((bs, max_pages), dtype=torch.int32, device=_DEV)
cache_seqlens = torch.zeros((bs,), dtype=torch.int32, device=_DEV)
cu_seqlens_k = torch.zeros((bs + 1,), dtype=torch.int32, device=_DEV)
strided = torch.arange(0, max_ctx, page_size, device=_DEV)
max_seq_pages = (int(sl.max().item()) + page_size - 1) // page_size
normal_decode_set_metadata(
cache_seqlens,
cu_seqlens_k,
page_table,
rt,
rpi,
strided,
max_seq_pages,
sl.to(torch.int64),
0,
page_size,
v2p_page_table=v2p_arg,
kernel_page_multiplier=mult,
)
if v2p:
# The translator's canonical, then the wrapper copies its rows.
canonical = torch.zeros((bs, max_pages), dtype=torch.int32, device=_DEV)
build_kv_read_table(
req_to_token=rt,
req_pool_indices=rpi,
seq_lens=sl.to(torch.int64),
v2p=v2p_full,
multiplier=mult,
page_size=page_size,
max_pages=max_pages,
out=canonical,
)
rows = torch.arange(bs, dtype=torch.int64, device=_DEV)
normal_decode_set_metadata(
cache_seqlens,
cu_seqlens_k,
page_table,
canonical,
rows,
max_seq_pages,
sl.to(torch.int64),
0,
page_size,
None,
None,
src_is_read_table=True,
)
else:
normal_decode_set_metadata(
cache_seqlens,
cu_seqlens_k,
page_table,
rt,
rpi,
max_seq_pages,
sl.to(torch.int64),
0,
page_size,
None,
None,
)
torch.cuda.synchronize()
want = _reference(rt, rpi, sl, page_size, v2p=v2p_arg, mult=mult)
want = _reference(
rt, rpi, sl, page_size, v2p=(v2p_full if v2p else None), mult=mult
)
return page_table, want, sl
def _assert_live_prefix(self, got, want, sl, page_size):
@@ -287,192 +349,9 @@ class TestFa3MetadataDenseBlockTable(unittest.TestCase):
"test batch degenerated: v2p is the identity on the pages used",
)
def test_agrees_with_flashmla_block_table(self):
"""fa3 and trtllm_mla build the same table two different ways; a
disagreement means one family is addressing the wrong pages."""
for page_size in (1, 64):
got, _, sl = self._run(page_size, v2p=True, mult=_LAYERS)
rt, rpi, sl2, v2p = TestDenseBlockTable._make_batch(self, page_size)
other = _fill_block_table(
rt, rpi, sl2, page_size, v2p=v2p, mult=_LAYERS
).long()
for r in range(got.shape[0]):
n_pages = (int(sl[r].item()) + page_size - 1) // page_size
self.assertTrue(
torch.equal(got[r, :n_pages].long(), other[r, :n_pages]),
f"fa3 and flashmla block tables disagree (row {r}, ps={page_size})",
)
class TestUnifiedMLAHookDetection(unittest.TestCase):
"""`unified_mla_hooks` decides whether the paged MLA backends translate at
all. Getting the predicate wrong is silent: the block table and KV write loc
stay in virtual id space and address the wrong pages once virtual and
physical diverge (e.g. after compaction)."""
@staticmethod
def _probe(**attrs):
from sglang.srt.layers.attention.unified_mem_hooks import (
unified_mla_hooks,
)
class _Alloc:
pass
alloc = _Alloc()
for k, v in attrs.items():
setattr(alloc, k, v)
return unified_mla_hooks(alloc)
def test_static_pool_disables_every_hook(self):
"""No v2p table -> statically-partitioned pool; req_to_token is already
physical, so all hooks must stay off (byte-identical to pre-change)."""
hooks = self._probe()
self.assertFalse(hooks.enabled)
self.assertIsNone(hooks.v2p_page_table)
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_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_for_kernel)
self.assertEqual(hooks.kernel_page_multiplier, _LAYERS)
def test_single_full_attention_layer_pool_is_still_unified(self):
"""REGRESSION: `kernel_page_multiplier == 1` does NOT mean static.
A hybrid MLA config with exactly one full-attention layer (e.g. a
pipeline-parallel rank owning a single MLA layer) has multiplier 1, yet
its locs are still virtual. Detecting on `multiplier > 1` would disable
the v2p gather here and corrupt reads/writes after compaction.
"""
table = torch.arange(8)
hooks = self._probe(
full_v2p_page_table=table,
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_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)
@unittest.skipUnless(_HAS_CUDA, "requires CUDA")
class TestInPlaceKvIndicesTranslate(unittest.TestCase):
"""The flashinfer decode updater must translate kv_indices IN PLACE.
Under cuda-graph replay the `kv_indices` it is handed IS the capture-stable
buffer the captured wrapper reads (`fast_decode_kwargs["kv_indices"]`), and
`fast_mla_decode_plan` ignores its `kv_indices` argument -- so rebinding the
local name to a fresh translated tensor leaves the graph reading VIRTUAL ids.
These pin the write-back contract that fix relies on.
"""
def _allocator(self, page_size=1, n_full_tokens=4096):
from sglang.srt.mem_cache.multi_ended_allocator import MultiEndedAllocator
from sglang.srt.mem_cache.unified_memory_pool import (
MambaSubPoolSpec,
MLASubPoolSpec,
UnifiedKVPool,
)
full = MLASubPoolSpec(
name="full",
layer_num=_LAYERS,
kv_lora_rank=512,
qk_rope_head_dim=64,
store_dtype=torch.bfloat16,
grow_direction="down",
)
mamba = MambaSubPoolSpec(
name="mamba",
layer_num=2,
conv_state_shapes=((8, 16),),
conv_dtype=torch.bfloat16,
temporal_state_shape=(4, 8, 8),
temporal_dtype=torch.float32,
grow_direction="up",
)
pool = UnifiedKVPool(
total_bytes=full.entry_bytes() * n_full_tokens + mamba.entry_bytes() * 16,
sub_pool_specs=[full, mamba],
device=_DEV,
enable_memory_saver=False,
page_size=page_size,
)
class _Stub:
def move_kv_cache(self, dst, src):
pass
full_alloc = MultiEndedAllocator(
kvcache=_Stub(),
unified_buffer=pool,
sub_pool_name="full",
device=_DEV,
is_id_owner=True,
page_size=page_size,
kernel_page_multiplier=_LAYERS,
)
mamba_alloc = MultiEndedAllocator(
kvcache=_Stub(),
unified_buffer=pool,
sub_pool_name="mamba",
device=_DEV,
is_id_owner=True,
)
full_alloc.bind_peer(mamba_alloc)
mamba_alloc.bind_peer(full_alloc)
return full_alloc
def test_int32_buffer_prefix_translated_tail_untouched(self):
"""Mirrors the updater: an int32 capture-stable buffer holding VIRTUAL
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)
self.assertIsNotNone(virt)
n = virt.numel()
# Capture-stable int32 buffer: [:n] freshly filled with virtual ids by
# create_flashinfer_kv_indices_triton, tail = stale junk from a bigger replay.
buf = torch.full((n * 3,), 2**30, dtype=torch.int32, device=_DEV)
buf[:n] = virt.to(torch.int32)
tail_before = buf[n:].clone()
valid = buf[:n]
valid.copy_(alloc.translate_kv_loc_for_kernel(valid))
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 kernel-facing ids in the stable buffer",
)
self.assertTrue(
torch.equal(buf[n:], tail_before),
"stale tail was modified -- it can hold ids outside the v2p table",
)
def test_dense_ids_differ_from_virtual(self):
"""Guard the guard: if dense == virtual the in-place test proves nothing."""
alloc = self._allocator()
virt = alloc.alloc(64)
self.assertIsNotNone(virt)
self.assertFalse(
torch.equal(alloc.translate_kv_loc_for_kernel(virt), virt),
"kernel-facing ids coincide with virtual ids; pick a different allocation",
)
# (The old fa3<->flashmla agreement case is gone: both families now
# consume the SAME canonical builder, so cross-family agreement holds by
# construction and the per-family cases above cover the two consumers.)
if __name__ == "__main__":