diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml index 1ea89ae97..d7e14ed64 100644 --- a/.buildkite/cuda/test-ready.yml +++ b/.buildkite/cuda/test-ready.yml @@ -76,7 +76,7 @@ steps: - label: "Diffusion ยท Other Test" commands: - - timeout 15m pytest -sv tests/diffusion/*.py -m 'core_model and cuda' --run-level "core_model" + - timeout 15m pytest -sv tests/diffusion/*.py tests/diffusion/ar_diffusion -m 'core_model and cuda' --run-level "core_model" mirror_hardwares: l4_1 - group: ":card_index_dividers: Tiny Model Tests [Single GPU]" diff --git a/tests/diffusion/ar_diffusion/test_kv_cache.py b/tests/diffusion/ar_diffusion/test_kv_cache.py index 086666063..4d590e716 100644 --- a/tests/diffusion/ar_diffusion/test_kv_cache.py +++ b/tests/diffusion/ar_diffusion/test_kv_cache.py @@ -25,12 +25,9 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import ( ) from vllm_omni.experimental.ar_diffusion.kv_cache.paged import chunk_window_skipped_tokens -pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] - - BLOCK = 16 -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False): @@ -155,14 +152,155 @@ def test_paged_pool_layout_exposes_flat_slot_views(): dtype=torch.float32, device=torch.device("cpu"), ) - assert kv_pools[0].shape == (2, 4, BLOCK, 4, 64) + key_cache, value_cache = kv_pools[0] + assert key_cache.shape == value_cache.shape == (4, BLOCK, 4, 64) assert k_pools[0].shape == (4 * BLOCK, 4, 64) assert v_pools[0].shape == (4 * BLOCK, 4, 64) + # The flat views must alias their block-shaped cache, so a slot write is + # visible through the layout the attention kernel reads. k_pools[0][BLOCK + 3].fill_(7) v_pools[0][2 * BLOCK + 5].fill_(11) - assert torch.equal(kv_pools[0][0, 1, 3], k_pools[0][BLOCK + 3]) - assert torch.equal(kv_pools[0][1, 2, 5], v_pools[0][2 * BLOCK + 5]) + assert torch.equal(key_cache[1, 3], k_pools[0][BLOCK + 3]) + assert torch.equal(value_cache[2, 5], v_pools[0][2 * BLOCK + 5]) + + +def test_key_and_value_caches_do_not_share_storage(): + """K and V must be separate allocations, not halves of one tensor. + + The paged-write custom op declares both as mutated. When they alias one + storage, inductor's reinplace pass can re-inplace only the first of the + two, and auto_functionalized_v2's clone of the entire second pool survives + into the compiled graph -- one full pool copied per compiled region per + denoising step. + """ + kv_pools, k_pools, v_pools = allocate_kv_pool_with_views( + num_blocks=4, + block_size=BLOCK, + num_layers=2, + num_kv_heads=4, + head_dim=64, + dtype=torch.float32, + device=torch.device("cpu"), + ) + seen = set() + for layer, (key_cache, value_cache) in enumerate(kv_pools): + k_storage = key_cache.untyped_storage().data_ptr() + v_storage = value_cache.untyped_storage().data_ptr() + assert k_storage != v_storage, f"layer {layer}: K and V share one allocation" + assert k_pools[layer].untyped_storage().data_ptr() == k_storage + assert v_pools[layer].untyped_storage().data_ptr() == v_storage + seen.update((k_storage, v_storage)) + # Every layer's K and V are distinct allocations as well. + assert len(seen) == 2 * len(kv_pools) + + # Negative control: writing V must not disturb K. + k_pools[0].fill_(0) + v_pools[0].fill_(5) + assert torch.equal(k_pools[0], torch.zeros_like(k_pools[0])) + + +def _pool_sized_allocations(code: str, pool_numel: int) -> list[int]: + """Sizes of generated allocations at least as large as one pool.""" + import re + + sizes = [] + for shape in re.findall(r"empty_strided_\w+\(\(([\d, ]*?)\)", code): + 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 + + +def _generated_code_for_pools(key_cache, value_cache): + """Compile a paged-write-shaped op over these pools, return the generated code.""" + from torch._inductor import config as inductor_config + from torch._inductor.utils import run_and_get_code + + class _Holder(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("key_cache", key_cache) + self.register_buffer("value_cache", value_cache) + + def forward(self, value): + return torch.ops.test_kv_pool_clone.write(self.key_cache, self.value_cache, value) + + torch._dynamo.reset() + # Inductor caches compiled code across ``torch._dynamo.reset()``, so the + # second call in a process would otherwise return the first one's code and + # both branches of this test would describe the same graph. + with inductor_config.patch(force_disable_caches=True): + _, codes = run_and_get_code(torch.compile(_Holder(), backend="inductor", fullgraph=True), torch.ones(1)) + return "\n".join(codes) + + +def test_compiled_paged_write_does_not_clone_the_pool(): + """The compiled graph must write the pools in place, not through a clone. + + ``test_key_and_value_caches_do_not_share_storage`` pins the mechanism this + change introduces -- two allocations rather than one. This pins the effect + that mechanism buys, which is the reason the change exists: with K and V as + halves of one tensor, inductor's reinplace pass re-inplaces only the first, + and ``auto_functionalized_v2``'s clone of the pool survives into the + generated code as a pool-sized allocation plus a copy, once per compiled + region per denoising step. + + The shared layout is compiled here as a positive control. Without it, a + codegen change that stops emitting ``empty_strided_*`` would leave the + detector matching nothing and this test passing for the wrong reason. + """ + if "write" not in dir(getattr(torch.ops, "test_kv_pool_clone", object())): + + @torch.library.custom_op("test_kv_pool_clone::write", mutates_args=("key_pool", "value_pool")) + def _write(key_pool: torch.Tensor, value_pool: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + key_pool[0, 0, 0, 0] = value[0] + value_pool[0, 0, 0, 0] = value[0] + return value * 2 + + @_write.register_fake + def _(key_pool, value_pool, value): + return torch.empty_like(value) + + # The stand-in models the real op's contract; that declaration is what + # drives the reinplace pass, so a drift here would silently test nothing. + real_schema = torch.ops.vllm_omni.ar_diffusion_paged_write_attn.default._schema + real_mutated = [arg.name for arg in real_schema.arguments if arg.alias_info and arg.alias_info.is_write] + assert real_mutated == ["key_pool", "value_pool"], ( + f"the real paged-write op now mutates {real_mutated}; this test still models two pools" + ) + + num_blocks, heads, dim = 8, 4, 16 + cache_shape = (num_blocks, BLOCK, heads, dim) + pool_numel = num_blocks * BLOCK * heads * dim + + # Positive control: one allocation, K and V aliasing it. + shared = torch.empty(2, *cache_shape) + shared_allocations = _pool_sized_allocations(_generated_code_for_pools(shared[0], shared[1]), pool_numel) + assert shared_allocations, ( + "the shared layout produced no pool-sized allocation, so this test can no " + "longer tell the two layouts apart -- inductor's codegen or the detector changed" + ) + + # The allocator under test. + kv_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.float32, + device=torch.device("cpu"), + ) + key_cache, value_cache = kv_pools[0] + separate_allocations = _pool_sized_allocations(_generated_code_for_pools(key_cache, value_cache), pool_numel) + assert not separate_allocations, ( + "the compiled graph allocates a pool-sized buffer, so the pool is being cloned " + f"rather than written in place: {separate_allocations} (control saw {shared_allocations})" + ) def test_build_manager_allocate_free_roundtrip(): diff --git a/tests/diffusion/ar_diffusion/test_kv_cache_cuda.py b/tests/diffusion/ar_diffusion/test_kv_cache_cuda.py new file mode 100644 index 000000000..7789c479f --- /dev/null +++ b/tests/diffusion/ar_diffusion/test_kv_cache_cuda.py @@ -0,0 +1,193 @@ +# 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})" + ) diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py index da61afe4d..f18a00ada 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py @@ -350,7 +350,9 @@ class ARDiffusionKVCache: self.null_block_id = self.manager.block_pool.null_block.block_id # Allocate the per-layer paged K/V pools on the given device. - self._kv_pools: list[torch.Tensor] = [] + # Per layer, ``[k_cache, v_cache]`` -- separate allocations; see + # allocate_kv_pool_with_views for why they are not one tensor. + self._kv_pools: list[list[torch.Tensor]] = [] self._k_pools: list[torch.Tensor] = [] self._v_pools: list[torch.Tensor] = [] if device is not None: diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py b/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py index 6b7d7ffe5..ddc1937a2 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py @@ -91,22 +91,35 @@ def allocate_kv_pool_with_views( head_dim: int, dtype: torch.dtype, device: torch.device, -) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]: +) -> tuple[list[list[torch.Tensor]], list[torch.Tensor], list[torch.Tensor]]: """Allocate vLLM-style paged KV pools plus flat slot-write views. - The owning tensor follows FlashAttention's block-table cache layout: - ``(2, num_blocks, block_size, num_kv_heads, head_dim)``, where dim-0 is - ``[K, V]``. The flat K/V views keep ``slot = block_id * block_size + - offset`` writes simple without changing the kernel-facing cache layout. + Each cache keeps FlashAttention's block-table layout, + ``(num_blocks, block_size, num_kv_heads, head_dim)``. The flat K/V views + keep ``slot = block_id * block_size + offset`` writes simple without + changing the kernel-facing layout, and ``kv_pools[layer][0]`` / + ``[1]`` still address K and V. + + K and V are allocated separately rather than as the two halves of one + ``(2, ...)`` tensor. Sharing one allocation makes them views of the same + storage, and the paged-write custom op declares both as mutated. Inductor's + reinplace pass can then re-inplace only the first of the two -- the second + hits the "mutated arg aliases a graph input" case in + ``_inductor/fx_passes/reinplace.py`` -- so ``auto_functionalized_v2``'s + clone of the whole V pool survives into the compiled graph. That clone is + one full pool per compiled region per denoising step. """ - kv_pools: list[torch.Tensor] = [] + kv_pools: list[list[torch.Tensor]] = [] k_pools: list[torch.Tensor] = [] v_pools: list[torch.Tensor] = [] + cache_shape = (num_blocks, block_size, num_kv_heads, head_dim) + flat_shape = (num_blocks * block_size, num_kv_heads, head_dim) for _ in range(num_layers): - kv = torch.empty(2, num_blocks, block_size, num_kv_heads, head_dim, dtype=dtype, device=device) - kv_pools.append(kv) - k_pools.append(kv[0].reshape(num_blocks * block_size, num_kv_heads, head_dim)) - v_pools.append(kv[1].reshape(num_blocks * block_size, num_kv_heads, head_dim)) + k = torch.empty(cache_shape, dtype=dtype, device=device) + v = torch.empty(cache_shape, dtype=dtype, device=device) + kv_pools.append([k, v]) + k_pools.append(k.reshape(flat_shape)) + v_pools.append(v.reshape(flat_shape)) return kv_pools, k_pools, v_pools