193 lines
8.2 KiB
Python
193 lines
8.2 KiB
Python
|
# SPDX-License-Identifier: Apache-2.0
|
||
|
|
"""CUDA-only regression test for the AR-Diffusion paged KV write op.
|
||
|
|
|
||
|
|
Split out of ``test_kv_cache.py``: that module is swept into CI's
|
||
|
|
``core_model and cpu`` job via a module-level ``pytest.mark.cpu``, and
|
||
|
|
``hardware_test(..., num_cards=1)`` adds no ``skipif`` guard for a missing
|
||
|
|
GPU, so this test was being collected (and would fail to import a CUDA
|
||
|
|
device) on CPU-only runners. Keeping it in its own, non-``cpu`` module means
|
||
|
|
it is only ever collected by a job that explicitly selects ``cuda`` for this
|
||
|
|
path -- none does yet, so for now this is intentionally not wired into any
|
||
|
|
CI job.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from tests.helpers.mark import hardware_test
|
||
|
|
from vllm_omni.experimental.ar_diffusion.kv_cache import allocate_kv_pool_with_views
|
||
|
|
|
||
|
|
BLOCK = 16
|
||
|
|
|
||
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
|
||
|
|
|
||
|
|
|
||
|
|
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||
|
|
def test_real_paged_write_op_does_not_clone_the_pool_on_cuda():
|
||
|
|
"""Same assertion as the CPU test, against the real op's schema and CUDA codegen.
|
||
|
|
|
||
|
|
The CPU test compiles a stand-in with the same ``mutates_args`` contract,
|
||
|
|
because the real op needs a device and an attention kernel. That leaves
|
||
|
|
one thing unverified: whether Triton codegen reinplaces the same way,
|
||
|
|
which is the path the reported regression was observed on.
|
||
|
|
|
||
|
|
This does not call the real op directly: ``ar_diffusion_paged_write_attn``
|
||
|
|
runs FlashAttention's paged ``flash_attn_varlen_func`` internally, and
|
||
|
|
compiling that under ``fullgraph=True`` can fail for reasons that have
|
||
|
|
nothing to do with cloning (a missing kernel for this toy shape, a graph
|
||
|
|
break, or an autotune illegal-memory-access on the control pool). Instead
|
||
|
|
this registers a second stand-in that copies the real op's full argument
|
||
|
|
list and ``mutates_args`` declaration -- so reinplace sees the identical
|
||
|
|
contract the real op presents -- but skips straight to the pool writes,
|
||
|
|
the same way the real op's own implementation does before it calls
|
||
|
|
attention. That isolates "does Inductor clone this schema on CUDA" from
|
||
|
|
"does FlashAttention compile", which is a separate, already-covered
|
||
|
|
concern.
|
||
|
|
|
||
|
|
Note the pools passed here are the flat views, since that is what the op
|
||
|
|
takes -- it indexes them by slot and unflattens internally.
|
||
|
|
"""
|
||
|
|
import re
|
||
|
|
|
||
|
|
from torch._inductor import config as inductor_config
|
||
|
|
from torch._inductor.utils import run_and_get_code
|
||
|
|
|
||
|
|
device = torch.device("cuda")
|
||
|
|
num_blocks, heads, dim, num_tokens = 8, 4, 64, 8
|
||
|
|
pool_numel = num_blocks * BLOCK * heads * dim
|
||
|
|
# The size a genuine clone would produce: K and V pool-sized allocations,
|
||
|
|
# coalesced into one buffer the way `auto_functionalized_v2` was cloning
|
||
|
|
# them before this PR's fix. Matching this exact size (rather than "any
|
||
|
|
# allocation at least as large as one pool") avoids false positives from
|
||
|
|
# unrelated workspace/scratch buffers Inductor may emit in that range.
|
||
|
|
clone_numel = 2 * pool_numel
|
||
|
|
|
||
|
|
if "write" not in dir(getattr(torch.ops, "test_kv_pool_clone_schema", object())):
|
||
|
|
|
||
|
|
@torch.library.custom_op("test_kv_pool_clone_schema::write", mutates_args=("key_pool", "value_pool"))
|
||
|
|
def _write(
|
||
|
|
query: torch.Tensor,
|
||
|
|
k_curr: torch.Tensor,
|
||
|
|
v_curr: torch.Tensor,
|
||
|
|
k_act: torch.Tensor | None,
|
||
|
|
v_act: torch.Tensor | None,
|
||
|
|
key_pool: torch.Tensor,
|
||
|
|
value_pool: torch.Tensor,
|
||
|
|
block_size: int,
|
||
|
|
video_slots: torch.Tensor,
|
||
|
|
action_slots: torch.Tensor,
|
||
|
|
block_table: torch.Tensor,
|
||
|
|
query_start_loc: torch.Tensor,
|
||
|
|
seq_lens: torch.Tensor,
|
||
|
|
max_query_len: int,
|
||
|
|
max_seq_len: int,
|
||
|
|
softmax_scale: float,
|
||
|
|
) -> torch.Tensor:
|
||
|
|
# Same pool-write side effect as `_paged_write_attn_impl`, minus
|
||
|
|
# the `ar_diffusion_paged_attention` call -- that is the part
|
||
|
|
# this test is not trying to exercise.
|
||
|
|
key_pool[video_slots] = k_curr.to(key_pool.dtype)
|
||
|
|
value_pool[video_slots] = v_curr.to(value_pool.dtype)
|
||
|
|
if k_act is not None and v_act is not None and k_act.shape[0] > 0:
|
||
|
|
key_pool[action_slots] = k_act.to(key_pool.dtype)
|
||
|
|
value_pool[action_slots] = v_act.to(value_pool.dtype)
|
||
|
|
return query * 2
|
||
|
|
|
||
|
|
@_write.register_fake
|
||
|
|
def _(
|
||
|
|
query,
|
||
|
|
k_curr,
|
||
|
|
v_curr,
|
||
|
|
k_act,
|
||
|
|
v_act,
|
||
|
|
key_pool,
|
||
|
|
value_pool,
|
||
|
|
block_size,
|
||
|
|
video_slots,
|
||
|
|
action_slots,
|
||
|
|
block_table,
|
||
|
|
query_start_loc,
|
||
|
|
seq_lens,
|
||
|
|
max_query_len,
|
||
|
|
max_seq_len,
|
||
|
|
softmax_scale,
|
||
|
|
):
|
||
|
|
return torch.empty_like(query)
|
||
|
|
|
||
|
|
class _Holder(torch.nn.Module):
|
||
|
|
def __init__(self, key_flat, value_flat):
|
||
|
|
super().__init__()
|
||
|
|
self.register_buffer("key_flat", key_flat)
|
||
|
|
self.register_buffer("value_flat", value_flat)
|
||
|
|
|
||
|
|
def forward(self, query, k_curr, v_curr, video_slots, action_slots, block_table, query_start_loc, seq_lens):
|
||
|
|
return torch.ops.test_kv_pool_clone_schema.write(
|
||
|
|
query,
|
||
|
|
k_curr,
|
||
|
|
v_curr,
|
||
|
|
None,
|
||
|
|
None,
|
||
|
|
self.key_flat,
|
||
|
|
self.value_flat,
|
||
|
|
BLOCK,
|
||
|
|
video_slots,
|
||
|
|
action_slots,
|
||
|
|
block_table,
|
||
|
|
query_start_loc,
|
||
|
|
seq_lens,
|
||
|
|
num_tokens,
|
||
|
|
num_tokens,
|
||
|
|
dim**-0.5,
|
||
|
|
)
|
||
|
|
|
||
|
|
def compiled_pool_allocations(key_flat, value_flat) -> list[int]:
|
||
|
|
args = (
|
||
|
|
torch.randn(num_tokens, heads, dim, dtype=torch.bfloat16, device=device),
|
||
|
|
torch.randn(num_tokens, heads, dim, dtype=torch.bfloat16, device=device),
|
||
|
|
torch.randn(num_tokens, heads, dim, dtype=torch.bfloat16, device=device),
|
||
|
|
torch.arange(num_tokens, device=device),
|
||
|
|
torch.empty(0, dtype=torch.long, device=device),
|
||
|
|
torch.arange(num_blocks, device=device, dtype=torch.int32).unsqueeze(0),
|
||
|
|
torch.tensor([0, num_tokens], device=device, dtype=torch.int32),
|
||
|
|
torch.tensor([num_tokens], device=device, dtype=torch.int32),
|
||
|
|
)
|
||
|
|
torch._dynamo.reset()
|
||
|
|
# See the CPU test: inductor's cache outlives torch._dynamo.reset().
|
||
|
|
with inductor_config.patch(force_disable_caches=True):
|
||
|
|
_, codes = run_and_get_code(
|
||
|
|
torch.compile(_Holder(key_flat, value_flat), backend="inductor", fullgraph=True), *args
|
||
|
|
)
|
||
|
|
sizes = []
|
||
|
|
for shape in re.findall(r"empty_strided_cuda\(\(([\d, ]*?)\)", "\n".join(codes)):
|
||
|
|
dims = [int(part) for part in shape.replace(" ", "").strip(",").split(",") if part]
|
||
|
|
numel = 1
|
||
|
|
for dim_size in dims:
|
||
|
|
numel *= dim_size
|
||
|
|
if dims and numel >= pool_numel:
|
||
|
|
sizes.append(numel)
|
||
|
|
return sizes
|
||
|
|
|
||
|
|
# Positive control: one allocation, K and V as its two halves. This must
|
||
|
|
# produce an allocation of exactly `clone_numel` -- that is the clone
|
||
|
|
# signature the assertion below checks for, not "no allocation this big".
|
||
|
|
shared = torch.empty(2, num_blocks * BLOCK, heads, dim, dtype=torch.bfloat16, device=device)
|
||
|
|
control = compiled_pool_allocations(shared[0], shared[1])
|
||
|
|
assert clone_numel in control, (
|
||
|
|
f"the shared layout produced no {clone_numel}-element allocation ({control}), so this test can no "
|
||
|
|
"longer tell the two layouts apart -- inductor's codegen or the detector changed"
|
||
|
|
)
|
||
|
|
|
||
|
|
_, k_pools, v_pools = allocate_kv_pool_with_views(
|
||
|
|
num_blocks=num_blocks,
|
||
|
|
block_size=BLOCK,
|
||
|
|
num_layers=1,
|
||
|
|
num_kv_heads=heads,
|
||
|
|
head_dim=dim,
|
||
|
|
dtype=torch.bfloat16,
|
||
|
|
device=device,
|
||
|
|
)
|
||
|
|
separate = compiled_pool_allocations(k_pools[0], v_pools[0])
|
||
|
|
assert clone_numel not in separate, (
|
||
|
|
f"the compiled graph allocates a {clone_numel}-element buffer, so the pool is being cloned "
|
||
|
|
f"rather than written in place: {separate} (control saw {control})"
|
||
|
|
)
|