wehub-resource-sync
77 lines
2.7 KiB
Python
77 lines
2.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Slot-mapping / block-table tests for the AR-Diffusion engine (Phase 1, Step 3)."""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.experimental.ar_diffusion.kv_cache import (
|
|
ARDiffusionRequestAdapter,
|
|
ChunkWindowSpec,
|
|
build_kv_manager,
|
|
chunk_slot_mapping,
|
|
compute_slot_mapping,
|
|
resident_block_ids,
|
|
)
|
|
|
|
BLOCK = 16
|
|
|
|
|
|
def make_spec(*, window_chunks=2):
|
|
return ChunkWindowSpec(
|
|
block_size=BLOCK,
|
|
num_kv_heads=4,
|
|
head_size=64,
|
|
dtype=torch.float16,
|
|
sliding_window=window_chunks * BLOCK,
|
|
chunk_size=BLOCK,
|
|
window_chunks=window_chunks,
|
|
)
|
|
|
|
|
|
def test_slot_mapping_matches_block_offsets():
|
|
# block table [5, 2, 9], block_size 4. slot = block*4 + (pos % 4).
|
|
slots = compute_slot_mapping([5, 2, 9], [0, 1, 4, 5, 8, 9], block_size=4)
|
|
assert slots.tolist() == [20, 21, 8, 9, 36, 37]
|
|
|
|
|
|
def test_slot_mapping_rejects_bad_block_size():
|
|
with pytest.raises(ValueError):
|
|
compute_slot_mapping([1], [0], block_size=0)
|
|
|
|
|
|
def test_chunk_slot_mapping_targets_current_chunk():
|
|
# After 2 committed chunks (num_computed=32), the in-flight chunk maps to the
|
|
# block covering positions 32..47 — i.e. block table index 2.
|
|
block_ids = [7, 4, 3] # block index 2 -> physical block 3
|
|
slots = chunk_slot_mapping(block_ids, num_computed_tokens=32, chunk_size=4, block_size=16)
|
|
# positions 32..35 -> block_index 2 -> physical 3 -> slots 48..51
|
|
assert slots.tolist() == [48, 49, 50, 51]
|
|
|
|
|
|
def test_resident_block_ids_excludes_null():
|
|
assert resident_block_ids([0, 1, 0, 2, 3], null_block_id=0) == [1, 2, 3]
|
|
|
|
|
|
def test_blocktable_build_single_request():
|
|
"""Integration: a real KVCacheManager block table -> a valid slot mapping.
|
|
|
|
Every slot for the in-flight chunk must land inside a real (non-null) block
|
|
that the manager actually allocated for the request.
|
|
"""
|
|
spec = make_spec()
|
|
mgr = build_kv_manager(spec, ["l0"], num_blocks=16, max_model_len=1024)
|
|
null_id = mgr.block_pool.null_block.block_id
|
|
adapter = ARDiffusionRequestAdapter("req", chunk_size=BLOCK)
|
|
|
|
mgr.allocate_slots(adapter, num_new_tokens=BLOCK)
|
|
block_ids = mgr.get_block_ids(adapter.request_id)[0]
|
|
real = set(resident_block_ids(block_ids, null_id))
|
|
assert real, "expected at least one real block after allocation"
|
|
|
|
slots = chunk_slot_mapping(block_ids, adapter.num_computed_tokens, spec.chunk_size, spec.block_size)
|
|
# Each slot resolves to a real allocated block, none to null_block.
|
|
blocks_used = {int(s) // spec.block_size for s in slots}
|
|
assert blocks_used <= real
|
|
assert null_id not in blocks_used
|
|
assert len(slots) == spec.chunk_size
|