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