597 lines
25 KiB
Python
597 lines
25 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Tests for AR-Diffusion paged self-attention contexts."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import subprocess
|
|
from importlib.util import find_spec
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec
|
|
from vllm_omni.experimental.ar_diffusion.kv_cache import (
|
|
ARDiffusionKVCache,
|
|
ARDiffusionKVConfig,
|
|
ARDiffusionPagedLayerContext,
|
|
ARDiffusionPagedLayerInputs,
|
|
ar_diffusion_paged_attention,
|
|
paged_write_attn,
|
|
)
|
|
from vllm_omni.experimental.ar_diffusion.kv_cache import paged_attention as paged_attention_module
|
|
from vllm_omni.experimental.ar_diffusion.kv_cache.config import KV_GATHER_ENV
|
|
from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
|
|
|
|
|
BLOCK = 16
|
|
N_HEADS = 4
|
|
HEAD_DIM = 64
|
|
POS = "positive"
|
|
NEG = "negative"
|
|
|
|
|
|
def make_state(
|
|
*, num_layers=1, window_chunks=2, dtype=torch.float32, device=torch.device("cpu"), reuse_history_staging=False
|
|
):
|
|
cfg = ARDiffusionKVConfig(
|
|
enable=True, chunk_size=BLOCK, window_chunks=window_chunks, reuse_history_staging=reuse_history_staging
|
|
)
|
|
kv = ARDiffusionKVCache(
|
|
cfg,
|
|
num_layers=num_layers,
|
|
num_kv_heads=N_HEADS,
|
|
head_size=HEAD_DIM,
|
|
dtype=dtype,
|
|
block_size=BLOCK,
|
|
max_model_len=4096,
|
|
available_bytes=1 << 26,
|
|
kv_branches=(ARDiffusionKVBranchSpec(POS, 0), ARDiffusionKVBranchSpec(NEG, 1)),
|
|
session_capacity=2,
|
|
frames_per_block=2,
|
|
max_scratch_tokens_per_branch=BLOCK,
|
|
device=device,
|
|
)
|
|
pos = kv.begin_request("r-pos")
|
|
neg = kv.begin_request("r-neg")
|
|
return kv, ARDiffusionKVState(kv, "s1", {POS: pos, NEG: neg}, num_layers=num_layers)
|
|
|
|
|
|
def _dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
|
|
scores = torch.einsum("bqhd,bkhd->bhqk", query.float(), key.float()) * (HEAD_DIM**-0.5)
|
|
probs = torch.softmax(scores, dim=-1).to(value.dtype)
|
|
return torch.einsum("bhqk,bkhd->bqhd", probs, value)
|
|
|
|
|
|
def _gpu_flash_attn_usable() -> bool:
|
|
if not torch.cuda.is_available():
|
|
return False
|
|
if torch.version.hip is not None:
|
|
for module in ("aiter", "flash_attn"):
|
|
try:
|
|
imported = __import__(module, fromlist=["flash_attn_varlen_func"])
|
|
if getattr(imported, "flash_attn_varlen_func", None) is not None:
|
|
return True
|
|
except ImportError:
|
|
pass
|
|
return False
|
|
try:
|
|
spec = find_spec("vllm.vllm_flash_attn")
|
|
if spec is None or spec.origin is None:
|
|
return True
|
|
fa2_so = Path(spec.origin).parent / "_vllm_fa2_C.abi3.so"
|
|
linked = subprocess.check_output(["ldd", str(fa2_so)], text=True, timeout=5)
|
|
except Exception:
|
|
return True
|
|
if "libcudart.so.13" not in linked:
|
|
return True
|
|
try:
|
|
out = subprocess.check_output(
|
|
["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"],
|
|
text=True,
|
|
timeout=5,
|
|
)
|
|
driver_major = int(out.splitlines()[0].split(".")[0])
|
|
except Exception:
|
|
return True
|
|
return driver_major >= 580
|
|
|
|
|
|
def _commit_video_span(
|
|
kv: ARDiffusionKVCache,
|
|
st: ARDiffusionKVState,
|
|
*,
|
|
kv_branch: str,
|
|
n_chunks: int,
|
|
dtype: torch.dtype,
|
|
device: torch.device,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
ctx = st.get_kv_caches(kv_branch, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
|
|
ctx.ensure_video_slots(device)
|
|
k = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
|
v = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
|
kv._k_pools[0][ctx.current_video_slot_mapping] = k[0]
|
|
kv._v_pools[0][ctx.current_video_slot_mapping] = v[0]
|
|
st.commit_paged_context(kv_branch)
|
|
return k, v
|
|
|
|
|
|
def test_paged_context_allocates_lazily_and_commits_after_forward():
|
|
_, st = make_state()
|
|
|
|
contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=True)
|
|
ctx = contexts[0].forward_ctx
|
|
assert isinstance(contexts[0], ARDiffusionPagedLayerContext)
|
|
assert st.adapter(POS).completed_chunks == 0
|
|
assert ctx.current_video_slot_mapping is None
|
|
|
|
ctx.ensure_video_slots(torch.device("cpu"))
|
|
assert st.adapter(POS).completed_chunks == 0
|
|
assert len(ctx.current_video_block_ids) == 1
|
|
|
|
st.commit_paged_context(POS)
|
|
assert st.adapter(POS).completed_chunks == 1
|
|
assert st._committed[POS] == BLOCK
|
|
|
|
|
|
def test_scratch_video_and_action_blocks_do_not_commit():
|
|
kv, st = make_state()
|
|
|
|
ctx = st.get_kv_caches(POS, seq_len=2 * BLOCK, commit_current=False)[0].forward_ctx
|
|
ctx.ensure_video_slots(torch.device("cpu"))
|
|
ctx.ensure_action_slots(3, torch.device("cpu"))
|
|
|
|
assert ctx.current_video_block_ids == kv.scratch_block_ids(POS, 0, 2)
|
|
assert ctx.action_scratch_block_ids == kv.scratch_block_ids(POS, 2, 1)
|
|
st.commit_paged_context(POS)
|
|
assert st.adapter(POS).completed_chunks == 0
|
|
assert st._committed[POS] == 0
|
|
|
|
|
|
def test_pipeline_kv_get_paged_path_has_no_gather_backend():
|
|
kv, st = make_state()
|
|
assert not hasattr(kv, "gather_window_all_layers")
|
|
|
|
from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
|
|
|
|
pipeline = DreamZeroPipeline.__new__(DreamZeroPipeline)
|
|
pipeline._ar_diffusion_kv_state = st
|
|
contexts = pipeline._kv_get(MagicMock(), False, seq_len=BLOCK, update_kv_cache=False)
|
|
|
|
assert len(contexts) == 1
|
|
assert isinstance(contexts[0], ARDiffusionPagedLayerContext)
|
|
|
|
|
|
@pytest.mark.parametrize("history_chunks", [0, 1, 3])
|
|
@pytest.mark.parametrize("action_len", [0, 3])
|
|
@pytest.mark.parametrize("commit_current", [False, True])
|
|
def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len, commit_current):
|
|
torch.manual_seed(0)
|
|
device = torch.device("cpu")
|
|
dtype = torch.float32
|
|
kv, st = make_state(dtype=dtype, device=device, window_chunks=2)
|
|
|
|
history_k_parts: list[torch.Tensor] = []
|
|
history_v_parts: list[torch.Tensor] = []
|
|
if history_chunks:
|
|
k, v = _commit_video_span(
|
|
kv,
|
|
st,
|
|
kv_branch=POS,
|
|
n_chunks=history_chunks,
|
|
dtype=dtype,
|
|
device=device,
|
|
)
|
|
history_k_parts.append(k)
|
|
history_v_parts.append(v)
|
|
|
|
ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=commit_current)[0].forward_ctx
|
|
ctx.ensure_video_slots(device)
|
|
current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
|
current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
|
kv._k_pools[0][ctx.current_video_slot_mapping] = current_k[0]
|
|
kv._v_pools[0][ctx.current_video_slot_mapping] = current_v[0]
|
|
|
|
action_k = action_v = None
|
|
if action_len:
|
|
ctx.ensure_action_slots(action_len, device)
|
|
action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
|
action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
|
kv._k_pools[0][ctx.action_slot_mapping] = action_k[0]
|
|
kv._v_pools[0][ctx.action_slot_mapping] = action_v[0]
|
|
|
|
query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
|
block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
|
|
action_len=action_len,
|
|
query_len=query.shape[1],
|
|
device=device,
|
|
)
|
|
paged = ar_diffusion_paged_attention(
|
|
query,
|
|
kv.key_cache(0),
|
|
kv.value_cache(0),
|
|
block_table=block_table,
|
|
query_start_loc=query_start_loc,
|
|
seq_lens=seq_lens,
|
|
max_query_len=max_query_len,
|
|
max_seq_len=max_seq_len,
|
|
softmax_scale=HEAD_DIM**-0.5,
|
|
causal=False,
|
|
)
|
|
|
|
if history_k_parts:
|
|
history_k = torch.cat(history_k_parts, dim=1)
|
|
history_v = torch.cat(history_v_parts, dim=1)
|
|
else:
|
|
history_k = torch.empty(1, 0, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
|
history_v = torch.empty(1, 0, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
|
new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :]
|
|
new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :]
|
|
if action_len:
|
|
new_k = torch.cat([new_k, action_k], dim=1)
|
|
new_v = torch.cat([new_v, action_v], dim=1)
|
|
ref = _dense_attention(query, new_k, new_v)
|
|
|
|
torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5)
|
|
|
|
before = st.adapter(POS).completed_chunks
|
|
st.commit_paged_context(POS)
|
|
assert st.adapter(POS).completed_chunks == before + (1 if commit_current else 0)
|
|
|
|
|
|
@pytest.mark.parametrize("history_chunks", [0, 1, 2, 3])
|
|
@pytest.mark.parametrize("window_chunks", [2, 4])
|
|
def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypatch, history_chunks, window_chunks):
|
|
"""A second probe of the same AR block restages its current K/V where the table actually holds it.
|
|
|
|
The padded table always ends in at least one action-capacity block, and while the window is still
|
|
growing in unused window capacity too, so "the last current_blocks entries" is padding. The refresh
|
|
has to start right after the visible history: empty (window growing), partial, full, and after a slide.
|
|
"""
|
|
monkeypatch.setenv(KV_GATHER_ENV, "1")
|
|
device = torch.device("cpu")
|
|
dtype = torch.float32
|
|
kv, st = make_state(dtype=dtype, device=device, window_chunks=window_chunks, reuse_history_staging=True)
|
|
for buffer in kv.history_staging[0]:
|
|
buffer.fill_(-7)
|
|
if history_chunks:
|
|
_commit_video_span(kv, st, kv_branch=POS, n_chunks=history_chunks, dtype=dtype, device=device)
|
|
|
|
ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
|
|
ctx.max_video_tokens = 2 * BLOCK
|
|
ctx.ensure_video_slots(device)
|
|
key_cache, value_cache = kv.key_cache(0), kv.value_cache(0)
|
|
|
|
def probe(step: int):
|
|
# Each probe of the block writes new current K/V, then stages.
|
|
kv._k_pools[0][ctx.current_video_slot_mapping] = torch.full((BLOCK, N_HEADS, HEAD_DIM), float(step))
|
|
kv._v_pools[0][ctx.current_video_slot_mapping] = torch.full((BLOCK, N_HEADS, HEAD_DIM), -float(step))
|
|
# Store the metadata on the context the way prepare() does.
|
|
(ctx.block_table, ctx.query_start_loc, ctx.seq_lens, ctx.max_query_len, ctx.max_seq_len) = (
|
|
ctx.build_block_table(action_len=0, query_len=BLOCK, device=device)
|
|
)
|
|
block_table, max_seq_len = ctx.block_table, ctx.max_seq_len
|
|
ctx._prepare_history_staging(0)
|
|
n_blocks = max_seq_len // BLOCK
|
|
block_ids = block_table[0, :n_blocks].to(torch.long)
|
|
stage_k, stage_v = ctx.history_staging(0)
|
|
# The custom op narrows the manager-owned buffers before calling this helper.
|
|
stage_k, stage_v = stage_k[:max_seq_len], stage_v[:max_seq_len]
|
|
paged_attention_module._stage_window(
|
|
stage_k,
|
|
stage_v,
|
|
key_cache,
|
|
value_cache,
|
|
block_ids,
|
|
n_blocks,
|
|
BLOCK,
|
|
first_block=ctx.stage_first_block if ctx.reuse_history else 0,
|
|
)
|
|
full_k = key_cache.index_select(0, block_ids).reshape(n_blocks * BLOCK, N_HEADS, HEAD_DIM)
|
|
full_v = value_cache.index_select(0, block_ids).reshape(n_blocks * BLOCK, N_HEADS, HEAD_DIM)
|
|
for buffer in kv.history_staging[0]:
|
|
assert (buffer[max_seq_len:] == -7).all()
|
|
return stage_k, stage_v, full_k, full_v, n_blocks
|
|
|
|
stage_k, stage_v, full_k, full_v, n_blocks = probe(1)
|
|
assert ctx.reuse_history is False
|
|
assert torch.equal(stage_k, full_k) and torch.equal(stage_v, full_v)
|
|
|
|
stage_k, stage_v, full_k, full_v, n_blocks = probe(2)
|
|
assert ctx.reuse_history is True
|
|
visible_history_blocks = min(history_chunks, 2 - 1) # window of 2 blocks minus the current one
|
|
assert ctx.stage_first_block == visible_history_blocks
|
|
# The padded table is wider than the live window, so the end-of-table guess is not the offset.
|
|
assert n_blocks - 1 != ctx.stage_first_block
|
|
assert torch.equal(stage_k, full_k) and torch.equal(stage_v, full_v)
|
|
|
|
|
|
@pytest.mark.parametrize("gather_enabled", [False, True])
|
|
def test_staging_is_only_allocated_for_the_gather_path(monkeypatch, gather_enabled):
|
|
if gather_enabled:
|
|
monkeypatch.setenv(KV_GATHER_ENV, "1")
|
|
else:
|
|
monkeypatch.delenv(KV_GATHER_ENV, raising=False)
|
|
device = torch.device("cpu")
|
|
kv, st = make_state(device=device, reuse_history_staging=True)
|
|
# The manager owns the pairs and allocates them only when their consumer is on.
|
|
assert bool(kv.history_staging) is gather_enabled
|
|
ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
|
|
ctx.ensure_video_slots(device)
|
|
(ctx.block_table, ctx.query_start_loc, ctx.seq_lens, ctx.max_query_len, ctx.max_seq_len) = ctx.build_block_table(
|
|
action_len=0, query_len=BLOCK, device=device
|
|
)
|
|
ctx._prepare_history_staging(0)
|
|
assert ctx.staging_enabled is gather_enabled
|
|
stage_k, _ = ctx.history_staging(0)
|
|
assert (stage_k is not None) is gather_enabled
|
|
|
|
|
|
@pytest.mark.parametrize("window_chunks", [2, 4])
|
|
def test_layer_inputs_preserve_static_staging_tensors(monkeypatch, window_chunks):
|
|
"""Compiled inputs must retain the manager's static-address annotations, including with spare capacity."""
|
|
monkeypatch.setenv(KV_GATHER_ENV, "1")
|
|
device = torch.device("cpu")
|
|
kv, st = make_state(device=device, window_chunks=window_chunks, reuse_history_staging=True)
|
|
stage_key, stage_value = kv.history_staging[0]
|
|
for _ in range(2):
|
|
ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
|
|
ctx.max_video_tokens = 2 * BLOCK
|
|
ctx.prepare(device, action_len=0, query_len=BLOCK)
|
|
inputs = ctx.layer_inputs(0)
|
|
assert inputs.stage_key is stage_key
|
|
assert inputs.stage_value is stage_value
|
|
assert inputs.max_seq_len == 3 * BLOCK
|
|
assert stage_key.shape[0] == (window_chunks + 1) * BLOCK
|
|
|
|
|
|
@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
|
|
@pytest.mark.parametrize("history_chunks", [1, 3])
|
|
@pytest.mark.parametrize("action_len", [0, 3])
|
|
@pytest.mark.parametrize("commit_current", [False, True])
|
|
def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len, commit_current):
|
|
torch.manual_seed(0)
|
|
device = torch.device("cuda")
|
|
dtype = torch.float16
|
|
kv, st = make_state(dtype=dtype, device=device, window_chunks=2)
|
|
|
|
history_k, history_v = _commit_video_span(
|
|
kv,
|
|
st,
|
|
kv_branch=POS,
|
|
n_chunks=history_chunks,
|
|
dtype=dtype,
|
|
device=device,
|
|
)
|
|
|
|
layer_ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=commit_current)[0]
|
|
ctx = layer_ctx.forward_ctx
|
|
current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
|
current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
|
action_k = action_v = None
|
|
if action_len:
|
|
action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
|
action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
|
|
|
query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
|
|
|
# The production path: once-per-forward host prep, then the fused
|
|
# write+attend custom op consuming the NamedTuple payload.
|
|
ctx.prepare(device=device, action_len=action_len, query_len=query.shape[1])
|
|
inputs = layer_ctx.to_layer_inputs()
|
|
assert isinstance(inputs, ARDiffusionPagedLayerInputs)
|
|
paged = paged_write_attn(
|
|
inputs,
|
|
query[0],
|
|
current_k[0],
|
|
current_v[0],
|
|
action_k[0] if action_k is not None else None,
|
|
action_v[0] if action_v is not None else None,
|
|
HEAD_DIM**-0.5,
|
|
).unsqueeze(0)
|
|
|
|
# Direct python-fn call on the same (already written) pools must be
|
|
# bit-exact: identical kernel, identical inputs.
|
|
direct = ar_diffusion_paged_attention(
|
|
query,
|
|
kv.key_cache(0),
|
|
kv.value_cache(0),
|
|
block_table=ctx.block_table,
|
|
query_start_loc=ctx.query_start_loc,
|
|
seq_lens=ctx.seq_lens,
|
|
max_query_len=ctx.max_query_len,
|
|
max_seq_len=ctx.max_seq_len,
|
|
softmax_scale=HEAD_DIM**-0.5,
|
|
causal=False,
|
|
)
|
|
assert torch.equal(paged, direct)
|
|
|
|
new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :]
|
|
new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :]
|
|
if action_len:
|
|
new_k = torch.cat([new_k, action_k], dim=1)
|
|
new_v = torch.cat([new_v, action_v], dim=1)
|
|
ref = _dense_attention(query, new_k, new_v)
|
|
|
|
torch.testing.assert_close(paged, ref, rtol=2e-2, atol=2e-2)
|
|
|
|
|
|
def test_block_table_padded_to_fixed_width():
|
|
"""Shapes must be constant across window growth: only values change."""
|
|
device = torch.device("cpu")
|
|
kv, st = make_state(window_chunks=2)
|
|
|
|
ctx1 = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=True)[0].forward_ctx
|
|
ctx1.prepare(device=device, action_len=0, query_len=BLOCK)
|
|
st.commit_paged_context(POS)
|
|
|
|
ctx2 = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=True)[0].forward_ctx
|
|
ctx2.prepare(device=device, action_len=0, query_len=BLOCK)
|
|
st.commit_paged_context(POS)
|
|
|
|
# 1-block vs 2-block visible history: same table width, same max_seq_len.
|
|
assert ctx1.block_table.shape == ctx2.block_table.shape
|
|
assert ctx1.max_seq_len == ctx2.max_seq_len
|
|
expected_width = kv.spec.sliding_window // kv.block_size + 1
|
|
assert ctx1.block_table.shape == (1, expected_width)
|
|
# Real lengths live in seq_lens, not the padded table.
|
|
assert int(ctx1.seq_lens[0]) == BLOCK
|
|
assert int(ctx2.seq_lens[0]) == 2 * BLOCK
|
|
|
|
|
|
def test_prepare_is_idempotent_and_layers_share_metadata():
|
|
device = torch.device("cpu")
|
|
kv, st = make_state(num_layers=2)
|
|
contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)
|
|
fctx = contexts[0].forward_ctx
|
|
fctx.prepare(device=device, action_len=0, query_len=BLOCK)
|
|
table = fctx.block_table
|
|
fctx.prepare(device=device, action_len=0, query_len=BLOCK)
|
|
assert fctx.block_table is table # memoized, not rebuilt
|
|
|
|
i0, i1 = contexts[0].to_layer_inputs(), contexts[1].to_layer_inputs()
|
|
# 0-dim tensors (NOT python ints) so dynamo doesn't install per-layer
|
|
# value guards on the shared block code object.
|
|
assert isinstance(i0.layer_idx, torch.Tensor) and int(i0.layer_idx) == 0
|
|
assert isinstance(i1.layer_idx, torch.Tensor) and int(i1.layer_idx) == 1
|
|
# All layers share the same metadata tensor objects.
|
|
assert i0.block_table is i1.block_table
|
|
assert i0.seq_lens is i1.seq_lens
|
|
assert i0.video_slots is i1.video_slots
|
|
assert i0.key_pool is kv._k_pools[0]
|
|
assert i0.value_pool is kv._v_pools[0]
|
|
assert i1.key_pool is kv._k_pools[1]
|
|
assert i1.value_pool is kv._v_pools[1]
|
|
|
|
|
|
def test_layer_inputs_before_prepare_raises():
|
|
_, st = make_state()
|
|
layer_ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0]
|
|
with pytest.raises(RuntimeError, match="before prepare"):
|
|
layer_ctx.to_layer_inputs()
|
|
|
|
|
|
def test_custom_op_registration_idempotent():
|
|
import importlib
|
|
import sys
|
|
|
|
assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn")
|
|
mod = "vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention"
|
|
saved = sys.modules.pop(mod)
|
|
try:
|
|
importlib.import_module(mod) # re-registration must not raise
|
|
finally:
|
|
sys.modules[mod] = saved
|
|
assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn")
|
|
|
|
|
|
def test_custom_op_mutable_arguments_cannot_be_elided_as_defaults():
|
|
# Older PyTorch ADInplaceOrView handlers index positional mutable inputs
|
|
# directly. Default-valued trailing inputs can disappear before that handler.
|
|
schema = torch.ops.vllm_omni.ar_diffusion_paged_write_attn.default._schema
|
|
mutable = [arg for arg in schema.arguments if arg.alias_info is not None and arg.alias_info.is_write]
|
|
assert {arg.name for arg in mutable} == {"key_pool", "value_pool", "stage_key", "stage_value"}
|
|
assert all(not arg.has_default_value() for arg in mutable)
|
|
|
|
|
|
@pytest.mark.parametrize("reuse_history_staging", [False, True])
|
|
def test_custom_op_compiles_fullgraph_without_recompile_on_value_change(monkeypatch, reuse_history_staging):
|
|
"""The op must trace as one opaque node: fullgraph OK, and changed tensor
|
|
VALUES (new slots / block ids) must not trigger recompilation."""
|
|
import torch._dynamo
|
|
|
|
device = torch.device("cpu")
|
|
monkeypatch.setenv(KV_GATHER_ENV, "1")
|
|
kv, st = make_state(num_layers=2, window_chunks=2, reuse_history_staging=reuse_history_staging)
|
|
if reuse_history_staging:
|
|
# Hold the host-side staging offset fixed while table/slot values change.
|
|
_commit_video_span(kv, st, kv_branch=POS, n_chunks=2, dtype=torch.float32, device=device)
|
|
|
|
def run_one_forward(commit):
|
|
contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=commit)
|
|
fctx = contexts[0].forward_ctx
|
|
fctx.prepare(device=device, action_len=0, query_len=BLOCK)
|
|
q = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
|
|
k = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
|
|
v = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
|
|
# Both layers use one compiled function; the tensor-valued layer index
|
|
# must not specialize it.
|
|
for layer_ctx in contexts:
|
|
out = compiled(layer_ctx.to_layer_inputs(), q, k, v)
|
|
st.commit_paged_context(POS)
|
|
return out
|
|
|
|
torch._dynamo.reset()
|
|
try:
|
|
from torch._dynamo.testing import CompileCounter
|
|
|
|
counter = CompileCounter()
|
|
|
|
def fn(inputs, q, k, v):
|
|
return paged_write_attn(inputs, q, k, v, None, None, HEAD_DIM**-0.5) * 1.0
|
|
|
|
compiled = torch.compile(fn, backend=counter, fullgraph=True)
|
|
|
|
run_one_forward(commit=True) # history grows between calls ->
|
|
run_one_forward(commit=True) # block-table VALUES change, shapes don't
|
|
run_one_forward(commit=False)
|
|
|
|
assert counter.frame_count == 1, f"recompiled: frame_count={counter.frame_count}"
|
|
finally:
|
|
# Leave a clean dynamo state for later suites in the same pytest process
|
|
# (e.g. model_executor transformers models).
|
|
torch._dynamo.reset()
|
|
|
|
|
|
@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
|
|
@pytest.mark.parametrize("history_chunks", [0, 1, 3])
|
|
@pytest.mark.parametrize("action_len", [0, 3])
|
|
def test_contiguous_kv_gather_path_matches_paged_path_gpu(monkeypatch, history_chunks, action_len):
|
|
"""VLLM_OMNI_AR_DIFFUSION_KV_GATHER=1 gathers the visible blocks and runs varlen FA3 without a block table.
|
|
|
|
Covers an empty, partial and full window plus a partially filled action block:
|
|
the gather reads the tail-padding null block, so its (zeroed) rows must be masked.
|
|
"""
|
|
|
|
torch.manual_seed(0)
|
|
device = torch.device("cuda")
|
|
dtype = torch.float16
|
|
kv, st = make_state(dtype=dtype, device=device, window_chunks=2)
|
|
if history_chunks:
|
|
_commit_video_span(kv, st, kv_branch=POS, n_chunks=history_chunks, dtype=dtype, device=device)
|
|
|
|
layer_ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0]
|
|
ctx = layer_ctx.forward_ctx
|
|
current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
|
current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
|
action_k = action_v = None
|
|
if action_len:
|
|
action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
|
action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
|
query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
|
|
ctx.prepare(device=device, action_len=action_len, query_len=query.shape[1])
|
|
inputs = layer_ctx.to_layer_inputs()
|
|
|
|
def run() -> torch.Tensor:
|
|
return paged_write_attn(
|
|
inputs,
|
|
query[0],
|
|
current_k[0],
|
|
current_v[0],
|
|
action_k[0] if action_k is not None else None,
|
|
action_v[0] if action_v is not None else None,
|
|
HEAD_DIM**-0.5,
|
|
).unsqueeze(0)
|
|
|
|
monkeypatch.delenv(KV_GATHER_ENV, raising=False)
|
|
paged = run()
|
|
monkeypatch.setenv(KV_GATHER_ENV, "1")
|
|
gathered = run()
|
|
assert torch.isfinite(gathered).all()
|
|
# Same kernel family on the same K/V: only accumulation order differs.
|
|
torch.testing.assert_close(gathered, paged, rtol=2e-3, atol=2e-3)
|