[AR-Diffusion] Allocate K and V separately so the compiled path stops cloning the pool (#6463)

Signed-off-by: linzhenpl07 <linzhenpl07@gmail.com>
Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

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():

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})"
+    )