[Bugfix][Core] Preserve static KV staging inputs in AR-Diffusion (#7921)

Signed-off-by: MichaelCaoo <139663530+MichaelCao0@users.noreply.github.com>

tests/diffusion/ar_diffusion/test_history_staging_gpu.py

@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CUDA regressions for contiguous AR-Diffusion history staging."""
+
+import pytest
+import torch
+
+from tests.diffusion.ar_diffusion.test_paged_attention import (
+    BLOCK,
+    HEAD_DIM,
+    N_HEADS,
+    POS,
+    _commit_video_span,
+    _gpu_flash_attn_usable,
+    make_state,
+)
+from tests.helpers.mark import hardware_test
+from vllm_omni.experimental.ar_diffusion.kv_cache import paged_write_attn
+from vllm_omni.experimental.ar_diffusion.kv_cache.config import KV_GATHER_ENV
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+@pytest.mark.skipif(
+    torch.version.hip is not None or not _gpu_flash_attn_usable(), reason="usable CUDA FlashAttention is required"
+)
+@hardware_test(res={"cuda": ["L4", "H100"]}, num_cards=1)
+@pytest.mark.parametrize("window_chunks", [2, 4])
+@torch.inference_mode()
+def test_staged_attention_updates_only_the_active_window_gpu(monkeypatch, window_chunks):
+    """Full and reused gathers match fresh attention without touching spare capacity or stale history."""
+    monkeypatch.setenv(KV_GATHER_ENV, "1")
+    device, dtype = torch.device("cuda"), torch.bfloat16
+    kv, st = make_state(device=device, dtype=dtype, window_chunks=window_chunks, reuse_history_staging=True)
+    _commit_video_span(kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device)
+    stage_key, stage_value = kv.history_staging[0]
+    stage_key.fill_(-7)
+    stage_value.fill_(-7)
+    query = torch.randn(BLOCK, N_HEADS, HEAD_DIM, device=device, dtype=dtype)
+    for step in range(2):
+        ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
+        ctx.max_video_tokens = 2 * BLOCK
+        ctx.prepare(device, action_len=0, query_len=BLOCK)
+        inputs = ctx.layer_inputs(0)
+        assert inputs.reuse_history is (step > 0)
+        key = torch.randn_like(query)
+        value = torch.randn_like(query)
+        actual = paged_write_attn(inputs, query, key, value, None, None, HEAD_DIM**-0.5)
+        expected = paged_write_attn(
+            inputs._replace(stage_key=None, stage_value=None, reuse_history=False),
+            query,
+            key,
+            value,
+            None,
+            None,
+            HEAD_DIM**-0.5,
+        )
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        ids = inputs.block_table[0].long()
+        for staged, cache in ((stage_key, kv.key_cache(0)), (stage_value, kv.value_cache(0))):
+            torch.testing.assert_close(
+                staged[: inputs.max_seq_len], cache.index_select(0, ids).flatten(0, 1), rtol=0, atol=0
+            )
+            assert (staged[inputs.max_seq_len :] == -7).all()
+
+
+@pytest.mark.skipif(
+    torch.version.hip is not None or not _gpu_flash_attn_usable(), reason="usable CUDA FlashAttention is required"
+)
+@hardware_test(res={"cuda": ["L4", "H100"]}, num_cards=1)
+@pytest.mark.parametrize("window_chunks", [2, 4])
+@pytest.mark.parametrize("reuse_history", [False, True])
+@torch.inference_mode()
+def test_staged_attention_replays_inductor_cudagraph(monkeypatch, window_chunks, reuse_history):
+    """Inductor must capture the real mutable op and replay with fresh current K/V."""
+    monkeypatch.setenv(KV_GATHER_ENV, "1")
+    device, dtype = torch.device("cuda"), torch.bfloat16
+    kv, st = make_state(device=device, dtype=dtype, window_chunks=window_chunks, reuse_history_staging=True)
+    _commit_video_span(kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device)
+    stage_key, stage_value = kv.history_staging[0]
+    stage_key.fill_(-7)
+    stage_value.fill_(-7)
+    query = torch.randn(BLOCK, N_HEADS, HEAD_DIM, device=device, dtype=dtype)
+    key, value = torch.randn_like(query), torch.randn_like(query)
+
+    ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
+    ctx.max_video_tokens = 2 * BLOCK
+    ctx.prepare(device, action_len=0, query_len=BLOCK)
+    inputs = ctx.layer_inputs(0)
+    # Populate history before testing the reuse variant, as the first denoising probe does.
+    paged_write_attn(inputs, query, key, value, None, None, HEAD_DIM**-0.5)
+    if reuse_history:
+        ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
+        ctx.max_video_tokens = 2 * BLOCK
+        ctx.prepare(device, action_len=0, query_len=BLOCK)
+        inputs = ctx.layer_inputs(0)
+    assert inputs.reuse_history is reuse_history
+
+    torch._dynamo.reset()
+    try:
+        compiled = torch.compile(paged_write_attn, mode="reduce-overhead", fullgraph=True)
+        for _ in range(3):
+            torch.compiler.cudagraph_mark_step_begin()
+            compiled(inputs, query, key, value, None, None, HEAD_DIM**-0.5)
+        current_omni_platform.synchronize()
+
+        previous = None
+        for step in range(2):
+            key.normal_()
+            # Distinct value ranges ensure a stale replay cannot match the reference.
+            value.fill_(2.0 if step == 0 else -2.0)
+            torch.compiler.cudagraph_mark_step_begin()
+            with torch.profiler.profile(
+                activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA]
+            ) as prof:
+                actual = compiled(inputs, query, key, value, None, None, HEAD_DIM**-0.5)
+                current_omni_platform.synchronize()
+            assert any("cudaGraphLaunch" in event.key for event in prof.key_averages()), (
+                "Inductor did not replay a CUDA graph; mutable staging inputs may have lost their static annotation"
+            )
+            expected = paged_write_attn(
+                inputs._replace(stage_key=None, stage_value=None, reuse_history=False),
+                query,
+                key,
+                value,
+                None,
+                None,
+                HEAD_DIM**-0.5,
+            )
+            torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+            if previous is not None:
+                assert not torch.equal(actual, previous)
+            # Copy outside capture: graph-owned output storage is reused on the next step.
+            previous = actual.clone()
+            del actual
+            ids = inputs.block_table[0].long()
+            for staged, cache in ((stage_key, kv.key_cache(0)), (stage_value, kv.value_cache(0))):
+                torch.testing.assert_close(
+                    staged[: inputs.max_seq_len], cache.index_select(0, ids).flatten(0, 1), rtol=0, atol=0
+                )
+                assert (staged[inputs.max_seq_len :] == -7).all()
+    finally:
+        torch._dynamo.reset()

tests/diffusion/ar_diffusion/test_paged_attention.py

@@ -244,7 +244,8 @@ def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len,
 
 
 @pytest.mark.parametrize("history_chunks", [0, 1, 2, 3])
-def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypatch, history_chunks):
+@pytest.mark.parametrize("window_chunks", [2, 4])
+def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypatch, history_chunks, window_chunks):
     """A second probe of the same AR block restages its current K/V where the table actually holds it.
 
     The padded table always ends in at least one action-capacity block, and while the window is still
@@ -254,11 +255,14 @@ def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypa
     monkeypatch.setenv(KV_GATHER_ENV, "1")
     device = torch.device("cpu")
     dtype = torch.float32
-    kv, st = make_state(dtype=dtype, device=device, window_chunks=2, reuse_history_staging=True)
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=window_chunks, reuse_history_staging=True)
+    for buffer in kv.history_staging[0]:
+        buffer.fill_(-7)
     if history_chunks:
         _commit_video_span(kv, st, kv_branch=POS, n_chunks=history_chunks, dtype=dtype, device=device)
 
     ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
+    ctx.max_video_tokens = 2 * BLOCK
     ctx.ensure_video_slots(device)
     key_cache, value_cache = kv.key_cache(0), kv.value_cache(0)
 
@@ -275,6 +279,8 @@ def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypa
         n_blocks = max_seq_len // BLOCK
         block_ids = block_table[0, :n_blocks].to(torch.long)
         stage_k, stage_v = ctx.history_staging(0)
+        # The custom op narrows the manager-owned buffers before calling this helper.
+        stage_k, stage_v = stage_k[:max_seq_len], stage_v[:max_seq_len]
         paged_attention_module._stage_window(
             stage_k,
             stage_v,
@@ -287,6 +293,8 @@ def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypa
         )
         full_k = key_cache.index_select(0, block_ids).reshape(n_blocks * BLOCK, N_HEADS, HEAD_DIM)
         full_v = value_cache.index_select(0, block_ids).reshape(n_blocks * BLOCK, N_HEADS, HEAD_DIM)
+        for buffer in kv.history_staging[0]:
+            assert (buffer[max_seq_len:] == -7).all()
         return stage_k, stage_v, full_k, full_v, n_blocks
 
     stage_k, stage_v, full_k, full_v, n_blocks = probe(1)
@@ -323,6 +331,24 @@ def test_staging_is_only_allocated_for_the_gather_path(monkeypatch, gather_enabl
     assert (stage_k is not None) is gather_enabled
 
 
+@pytest.mark.parametrize("window_chunks", [2, 4])
+def test_layer_inputs_preserve_static_staging_tensors(monkeypatch, window_chunks):
+    """Compiled inputs must retain the manager's static-address annotations, including with spare capacity."""
+    monkeypatch.setenv(KV_GATHER_ENV, "1")
+    device = torch.device("cpu")
+    kv, st = make_state(device=device, window_chunks=window_chunks, reuse_history_staging=True)
+    stage_key, stage_value = kv.history_staging[0]
+    for _ in range(2):
+        ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
+        ctx.max_video_tokens = 2 * BLOCK
+        ctx.prepare(device, action_len=0, query_len=BLOCK)
+        inputs = ctx.layer_inputs(0)
+        assert inputs.stage_key is stage_key
+        assert inputs.stage_value is stage_value
+        assert inputs.max_seq_len == 3 * BLOCK
+        assert stage_key.shape[0] == (window_chunks + 1) * BLOCK
+
+
 @pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
 @pytest.mark.parametrize("history_chunks", [1, 3])
 @pytest.mark.parametrize("action_len", [0, 3])
@@ -363,8 +389,8 @@ def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len,
         query[0],
         current_k[0],
         current_v[0],
-        action_k[0] if action_len else None,
-        action_v[0] if action_len else None,
+        action_k[0] if action_k is not None else None,
+        action_v[0] if action_v is not None else None,
         HEAD_DIM**-0.5,
     ).unsqueeze(0)
 
@@ -556,8 +582,8 @@ def test_contiguous_kv_gather_path_matches_paged_path_gpu(monkeypatch, history_c
             query[0],
             current_k[0],
             current_v[0],
-            action_k[0] if action_len else None,
-            action_v[0] if action_len else None,
+            action_k[0] if action_k is not None else None,
+            action_v[0] if action_v is not None else None,
             HEAD_DIM**-0.5,
         ).unsqueeze(0)
 

vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py

@@ -264,6 +264,7 @@ class ARDiffusionPagedForwardContext:
         # gather path -- their only consumer -- is switched on.
         if not self.kv_cache.history_staging:
             return
+        assert self.block_table is not None
         if action_len or self.block_table.shape[0] != 1:
             # Action tokens live in scratch blocks outside the video window; staging them is not modelled.
             return
@@ -291,12 +292,14 @@ class ARDiffusionPagedForwardContext:
         self.stage_first_block = self.kv_len // self.block_size - self.num_current_video_blocks
 
     def history_staging(self, layer_idx: int) -> tuple[torch.Tensor | None, torch.Tensor | None]:
-        """This layer's staging pair, narrowed to this forward's staged window (a prefix of the manager's buffer)."""
+        """Return the manager-owned tensors marked static for CUDA Graph input mutation.
+
+        Fresh prefix views lose the bases' static-address annotation. Narrow the
+        window inside the attention custom op instead of at the compiled boundary.
+        """
         if not self.staging_enabled:
             return None, None
-        key, value = self.kv_cache.history_staging[layer_idx]
-        tokens = int(self.max_seq_len)
-        return key[:tokens], value[:tokens]
+        return self.kv_cache.history_staging[layer_idx]
 
     def layer_inputs(self, layer_idx: int) -> ARDiffusionPagedLayerInputs:
         if not getattr(self, "_prepared", False):
@@ -435,12 +438,11 @@ def _stage_window(stage_key, stage_value, key_cache, value_cache, block_ids, n_b
     earlier forward of the same AR block and cannot have changed since.
     """
     ids = block_ids[first_block:]
-    blocks = n_blocks - first_block
-    stage_key.view(n_blocks, block_size, *key_cache.shape[2:])[first_block:].copy_(
-        key_cache.index_select(0, ids).view(blocks, block_size, *key_cache.shape[2:])
-    )
-    stage_value.view(n_blocks, block_size, *value_cache.shape[2:])[first_block:].copy_(
-        value_cache.index_select(0, ids).view(blocks, block_size, *value_cache.shape[2:])
+    # Write directly into the caller-owned window, avoiding an intermediate
+    # gathered tensor and its copy on every layer and denoising step.
+    torch.index_select(key_cache, 0, ids, out=stage_key.view(n_blocks, block_size, *key_cache.shape[2:])[first_block:])
+    torch.index_select(
+        value_cache, 0, ids, out=stage_value.view(n_blocks, block_size, *value_cache.shape[2:])[first_block:]
     )
 
 
@@ -522,7 +524,7 @@ def ar_diffusion_paged_attention(
         # vLLM's seqused_k/fa_version API. The ROCm flash-attn paged kernel also
         # requires 128-token blocks, while AR-Diffusion uses frame-aligned
         # 16-token blocks, so gather the visible blocks on-device first.
-        flash_attn_varlen_func = _rocm_flash_attn_varlen_func()
+        rocm_flash_attn_varlen_func = _rocm_flash_attn_varlen_func()
         cu_seqlens_k = torch.cat([seq_lens.new_zeros(1), torch.cumsum(seq_lens, dim=0, dtype=torch.int32)])
         positions = torch.arange(int(max_seq_len), device=query_flat.device)
         logical_blocks = torch.div(positions, key_cache.shape[1], rounding_mode="floor")
@@ -533,7 +535,7 @@ def ar_diffusion_paged_attention(
         valid = positions.unsqueeze(0) < seq_lens.unsqueeze(1)
         packed_k = gathered_k[valid]
         packed_v = gathered_v[valid]
-        out = flash_attn_varlen_func(
+        out = rocm_flash_attn_varlen_func(
             q=query_flat,
             k=packed_k,
             v=packed_v,
@@ -579,8 +581,10 @@ def ar_diffusion_paged_attention(
             # byte-for-byte what a full gather would produce. When the history did move -- new block,
             # slid window, different session -- reuse_history is 0 and the whole window is re-gathered,
             # which is the same work the unstaged path always does.
-            k_flat = stage_key
-            v_flat = stage_value
+            # Keep the marked base tensors as graph inputs, but only stage the
+            # active prefix when the manager allocated a larger capacity.
+            k_flat = stage_key[:max_seq_len]
+            v_flat = stage_value[:max_seq_len]
             if reuse_history:
                 # The caller's metadata must describe a block-aligned current chunk inside the staged
                 # window; anything else is a wrong offset, not a reason to quietly restage everything.
@@ -621,12 +625,12 @@ def ar_diffusion_paged_attention(
             fa_version=fa_version,
         )
     else:
-        from vllm.vllm_flash_attn import flash_attn_varlen_func
+        from vllm.vllm_flash_attn import flash_attn_varlen_func as paged_flash_attn_varlen_func
 
         fa_version = _resolve_fa_version(query_flat.shape[-1])
 
         out = torch.empty_like(query_flat)
-        flash_attn_varlen_func(
+        paged_flash_attn_varlen_func(
             q=query_flat,
             k=key_cache,
             v=value_cache,