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