[LingBot World] Eleven bit-identical removals of repeated work in the realtime path (#7648)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

vllm_omni/experimental/ar_diffusion/kv_cache/config.py

@@ -1,10 +1,24 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Configuration for AR-Diffusion engine-level KV cache management."""
 
 from __future__ import annotations
 
+import os
 from dataclasses import dataclass
 
+KV_GATHER_ENV = "VLLM_OMNI_AR_DIFFUSION_KV_GATHER"
+
+
+def contiguous_kv_gather_enabled() -> bool:
+    """Whether the experimental contiguous-K/V gather attention path is on (``VLLM_OMNI_AR_DIFFUSION_KV_GATHER=1``).
+
+    The one place the switch is read: the KV manager consults it when it budgets and allocates the history
+    staging buffers, and the attention dispatch when it picks the path, so the two cannot disagree. It is the
+    only consumer of ``ARDiffusionKVConfig.reuse_history_staging``.
+    """
+    return os.environ.get(KV_GATHER_ENV, "0") == "1"
+
 
 @dataclass
 class ARDiffusionKVConfig:
@@ -30,6 +44,12 @@ class ARDiffusionKVConfig:
     # DiT graphs for every window-fill shape at load time via a synthetic rollout,
     # so the serving run is fast from the first chunk. No effect when eager.
     warmup_cudagraph: bool = True
+    # Keep one contiguous K/V staging buffer per layer and refresh only the tokens that changed.
+    # Within one AR block every forward attends the same history and differs only in the current
+    # chunk, so re-gathering the whole visible window per forward re-copies bytes that did not move.
+    # Off by default: it costs two max-sequence buffers per layer, per rank. See
+    # ARDiffusionPagedForwardContext.history_staging.
+    reuse_history_staging: bool = False
     # Also capture the post-window-boundary (reset-cycle) forward during warm-up.
     warmup_capture_reset: bool = False
 

vllm_omni/experimental/ar_diffusion/kv_cache/manager.py

@@ -15,7 +15,8 @@ from __future__ import annotations
 import inspect
 import os
 from collections.abc import Collection, Iterable, Sequence
-from typing import cast
+from dataclasses import dataclass
+from typing import Any, cast
 
 import torch
 from vllm.logger import init_logger
@@ -29,7 +30,11 @@ from vllm.v1.request import RequestStatus
 
 from vllm_omni.diffusion.diffusion_kv.layout import build_kv_cache_tensor
 from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec
-from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
+from vllm_omni.experimental.ar_diffusion.kv_cache.config import (
+    KV_GATHER_ENV,
+    ARDiffusionKVConfig,
+    contiguous_kv_gather_enabled,
+)
 from vllm_omni.experimental.ar_diffusion.kv_cache.paged import (
     ChunkWindowManager,
     ChunkWindowSpec,
@@ -158,6 +163,18 @@ def build_kv_manager(
     return KVCacheManager(config, **kwargs)
 
 
+@dataclass
+class HistoryStagingState:
+    """What the staged history window currently holds: whose session and which visible blocks.
+
+    Owned by the cache next to the buffers themselves; ``ARDiffusionPagedForwardContext`` reads it to decide
+    whether a forward may keep the staged history and writes back what it staged.
+    """
+
+    adapter: Any | None = None  # weakref to the session adapter the window was staged for
+    signature: tuple[Any, ...] | None = None
+
+
 class ARDiffusionKVCache:
     """Own the paged KV pool and KV-branch-local storage for one model.
 
@@ -180,6 +197,7 @@ class ARDiffusionKVCache:
         kv_branches: tuple[ARDiffusionKVBranchSpec, ...],
         session_capacity: int,
         cross_attention_lengths: dict[str, int] | None = None,
+        cross_attention_kv_heads: dict[str, int] | None = None,
         device: torch.device | None = None,
         frames_per_block: int = 1,
         max_scratch_tokens_per_branch: int = 0,
@@ -225,6 +243,18 @@ class ARDiffusionKVCache:
         self.head_size = head_size
         self.dtype = dtype
         self.cross_attention_lengths = dict(cross_attention_lengths or {})
+        # A cross-attention cache may hold more heads than the self-attention share: a model that keeps every
+        # local head on every rank, rather than sharding heads across the sequence-parallel group, stores the
+        # full local set here. Absent means the self-attention head count, which is what every caller did before.
+        self.cross_attention_kv_heads = {
+            name: int((cross_attention_kv_heads or {}).get(name, num_kv_heads)) for name in self.cross_attention_lengths
+        }
+        invalid_heads = {n: h for n, h in self.cross_attention_kv_heads.items() if h <= 0}
+        if invalid_heads:
+            raise ValueError(f"cross_attention_kv_heads must be positive, got {invalid_heads}")
+        unknown_heads = set(cross_attention_kv_heads or {}) - set(self.cross_attention_lengths)
+        if unknown_heads:
+            raise ValueError(f"cross_attention_kv_heads names unknown caches: {sorted(unknown_heads)}")
         invalid_cross = {name: length for name, length in self.cross_attention_lengths.items() if length <= 0}
         if invalid_cross:
             raise ValueError(f"cross_attention_lengths must be positive, got {invalid_cross}")
@@ -274,20 +304,46 @@ class ARDiffusionKVCache:
         self.scratch_reserved_bytes = self.scratch_num_blocks * page_size_bytes
         self.model_owned_state_bytes_per_session = model_owned_state_bytes_per_session
 
-        def _cross_pool_bytes(length: int) -> int:
-            return int(2 * len(self.kv_branches) * length * num_kv_heads * head_size * dtype.itemsize * num_layers)
+        def _cross_pool_bytes(length: int, heads: int) -> int:
+            return int(2 * len(self.kv_branches) * length * heads * head_size * dtype.itemsize * num_layers)
 
         self.cross_attention_bytes_per_session = sum(
-            _cross_pool_bytes(length) for length in self.cross_attention_lengths.values()
+            _cross_pool_bytes(length, self.cross_attention_kv_heads[name])
+            for name, length in self.cross_attention_lengths.items()
         )
 
         def _required_managed_blocks(capacity: int) -> int:
             resident_per_session = config.sink_chunks + config.window_chunks
             return self.num_local_kv_branches * (capacity * resident_per_session + self.frames_per_block) + 2
 
+        # reuse_history_staging keeps one contiguous K and V buffer per layer,
+        # sized to the padded visible window (sink + window plus the
+        # action-capacity block build_block_table always reserves), for the
+        # whole worker. It is allocated lazily on the first prepared forward,
+        # which is after admission, so it has to be reserved here or admission
+        # can succeed and that first forward run out of memory. Only the
+        # contiguous-gather attention path consumes it.
+        self.history_staging_reserved_bytes = 0
+        self.history_staging_tokens = 0
+        if config.reuse_history_staging:
+            if contiguous_kv_gather_enabled():
+                self.history_staging_tokens = (
+                    self.spec.sliding_window + config.sink_chunks * config.chunk_size + block_size
+                )
+                self.history_staging_reserved_bytes = int(
+                    2 * num_layers * self.history_staging_tokens * num_kv_heads * head_size * dtype.itemsize
+                )
+            else:
+                _log.warning(
+                    "reuse_history_staging is set but the contiguous K/V gather path (%s=1) is off: "
+                    "staging has no consumer and is neither budgeted nor allocated.",
+                    KV_GATHER_ENV,
+                )
+
         def _required_bytes(capacity: int) -> int:
             return (
                 self.scratch_reserved_bytes
+                + self.history_staging_reserved_bytes
                 + capacity * self.cross_attention_bytes_per_session
                 + capacity * self.model_owned_state_bytes_per_session
                 + _required_managed_blocks(capacity) * page_size_bytes
@@ -298,7 +354,7 @@ class ARDiffusionKVCache:
             raise ValueError(
                 "AR-Diffusion available device memory cannot fit one session: "
                 f"available={available_bytes} bytes, required={one_session_bytes} bytes "
-                "(managed self-attention + cross-attention + scratch + model-owned state)."
+                "(managed self-attention + cross-attention + scratch + model-owned state + K/V staging)."
             )
         self.memory_budget_bytes = max(self.configured_memory_budget_bytes, one_session_bytes)
         if self.memory_budget_bytes > self.configured_memory_budget_bytes:
@@ -325,6 +381,7 @@ class ARDiffusionKVCache:
         self_attn_budget_bytes = (
             self.memory_budget_bytes
             - self.scratch_reserved_bytes
+            - self.history_staging_reserved_bytes
             - self.cross_attention_reserved_bytes
             - self.model_owned_state_reserved_bytes
         )
@@ -374,6 +431,20 @@ class ARDiffusionKVCache:
                 dtype,
                 device,
             )
+        # reuse_history_staging: one contiguous (K, V) pair per layer, sized to the padded visible window this
+        # cache admits and budgeted above, owned here next to the pools it stages from. Allocated outside any
+        # capture and marked static, so compiled regions treat it as a stable input rather than a new tensor.
+        self.history_staging: list[tuple[torch.Tensor, torch.Tensor]] = []
+        self.history_staging_state = HistoryStagingState()
+        if self.history_staging_tokens and device is not None:
+            for _ in range(num_layers):
+                pair = (
+                    torch.empty((self.history_staging_tokens, num_kv_heads, head_size), device=device, dtype=dtype),
+                    torch.empty((self.history_staging_tokens, num_kv_heads, head_size), device=device, dtype=dtype),
+                )
+                for buffer in pair:
+                    torch._dynamo.mark_static_address(buffer)
+                self.history_staging.append(pair)
 
     # -- cross-attention pool access -------------------------------------------
     # Cross-attn KV is static once populated — write once (from text encoder),
@@ -455,7 +526,7 @@ class ARDiffusionKVCache:
         if not self._allocate_tensors:
             raise RuntimeError("AR-Diffusion cross-attention tensors require a configured pool device")
 
-        shape = (length, self.num_kv_heads, self.head_size)
+        shape = (length, self.cross_attention_kv_heads[cache_name], self.head_size)
         expected_input_shape = (1, *shape)
         k_pool = [torch.empty(shape, dtype=self.dtype, device=self.device) for _ in range(self.num_layers)]
         v_pool = [torch.empty(shape, dtype=self.dtype, device=self.device) for _ in range(self.num_layers)]

vllm_omni/experimental/ar_diffusion/kv_cache/paged.py

@@ -115,11 +115,27 @@ def allocate_kv_pool_with_views(
     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):
-        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))
+        # The flat slot views are what the paged-write custom op mutates inside the
+        # compiled block graphs, so they must be the *base* allocations: dynamo only
+        # honours mark_static_address on the tensor object it sees as a graph input,
+        # and a view of a static base is still classified as a mutated input, which
+        # makes inductor skip CUDA graphs for every block ("skipping cudagraphs due
+        # to mutated inputs"). The block-table layout is therefore the view here.
+        k = torch.empty(flat_shape, dtype=dtype, device=device)
+        v = torch.empty(flat_shape, dtype=dtype, device=device)
+        for pool in (k, v):
+            # Block 0 is the manager's null block: block tables are tail-padded with
+            # it and it is never written. The contiguous-K/V gather path may read
+            # (masked) rows from it, so it must hold finite values; every other
+            # block is fully written before it is read.
+            pool[:block_size].zero_()
+            # Session-lifetime storage: tell dynamo/inductor the address is static so
+            # CUDA-graph trees (mode="reduce-overhead") mutate it in place instead of
+            # copying ~300 MB per layer per replay (or skipping the graph).
+            torch._dynamo.mark_static_address(pool)
+        kv_pools.append([k.view(cache_shape), v.view(cache_shape)])
+        k_pools.append(k)
+        v_pools.append(v)
     return kv_pools, k_pools, v_pools
 
 

vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py

@@ -1,13 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Paged self-attention helpers for AR-Diffusion KV reuse."""
 
 from __future__ import annotations
 
+import weakref
 from dataclasses import dataclass, field
 from typing import Any, ClassVar, NamedTuple
 
 import torch
 
+from vllm_omni.experimental.ar_diffusion.kv_cache.config import contiguous_kv_gather_enabled
 from vllm_omni.experimental.ar_diffusion.kv_cache.paged import compute_slot_mapping
 
 _LAYER_IDX_TENSORS: dict[int, torch.Tensor] = {}
@@ -21,6 +24,19 @@ def _layer_idx_tensor(layer_idx: int) -> torch.Tensor:
     return t
 
 
+def _to_device_async(t: torch.Tensor, device: torch.device) -> torch.Tensor:
+    """Host->device copy that does not stall the CPU.
+
+    ``tensor.to(device)`` from pageable memory is synchronous: the CPU blocks until
+    every kernel already queued on the stream has finished. Pinned + non_blocking
+    keeps the CPU running ahead (the caching host allocator keeps the pinned
+    source alive until the copy's stream event completes).
+    """
+    if torch.device(device).type != "cuda" or t.device.type != "cpu":
+        return t.to(device=device)
+    return t.pin_memory().to(device=device, non_blocking=True)
+
+
 class ARDiffusionPagedLayerInputs(NamedTuple):
     """Compiled-region payload for one layer's paged self-attention.
 
@@ -48,6 +64,14 @@ class ARDiffusionPagedLayerInputs(NamedTuple):
     seq_lens: torch.Tensor
     max_query_len: int
     max_seq_len: int
+    # Contiguous K/V staging for this layer, when reuse_history_staging is on. ``reuse_history`` is a host
+    # bool decided once per forward (``_prepare_history_staging``): it selects one of two compiled variants.
+    stage_key: torch.Tensor | None = None
+    stage_value: torch.Tensor | None = None
+    reuse_history: bool = False
+    # Index of the current chunk's first block within the staged window: right after the visible history,
+    # which is not the end of the padded block table (that carries at least one action-capacity block).
+    stage_first_block: int = 0
 
 
 @dataclass
@@ -104,15 +128,15 @@ class ARDiffusionPagedForwardContext:
             start_block = start // self.block_size
             self.current_video_block_ids = [int(b) for b in table[start_block : start_block + n_blocks]]
             positions = torch.arange(start, start + self.seq_len, dtype=torch.long)
-            self.current_video_slot_mapping = compute_slot_mapping(table, positions, self.block_size).to(device=device)
+            self.current_video_slot_mapping = _to_device_async(
+                compute_slot_mapping(table, positions, self.block_size), device
+            )
         else:
             self.current_video_block_ids = self.kv_cache.scratch_block_ids(self.kv_branch, 0, n_blocks)
             positions = torch.arange(self.seq_len, dtype=torch.long)
-            self.current_video_slot_mapping = compute_slot_mapping(
-                self.current_video_block_ids,
-                positions,
-                self.block_size,
-            ).to(device=device)
+            self.current_video_slot_mapping = _to_device_async(
+                compute_slot_mapping(self.current_video_block_ids, positions, self.block_size), device
+            )
         self._allocated_video = True
 
     def ensure_action_slots(self, action_len: int, device: torch.device) -> None:
@@ -135,11 +159,9 @@ class ARDiffusionPagedForwardContext:
             action_blocks,
         )
         positions = torch.arange(action_len, dtype=torch.long)
-        self.action_slot_mapping = compute_slot_mapping(
-            self.action_scratch_block_ids,
-            positions,
-            self.block_size,
-        ).to(device=device)
+        self.action_slot_mapping = _to_device_async(
+            compute_slot_mapping(self.action_scratch_block_ids, positions, self.block_size), device
+        )
         self._action_len = action_len
 
     def video_block_table(self, device: torch.device) -> tuple[list[int], int]:
@@ -192,9 +214,12 @@ class ARDiffusionPagedForwardContext:
         self.query_len = int(query_len)
         self.kv_len = int(video_len + action_len)
         max_seq_len = int(self.max_video_tokens + action_capacity_blocks * self.block_size)
-        block_table = torch.tensor([padded], dtype=torch.int32, device=device)
-        query_start_loc = torch.tensor([0, self.query_len], dtype=torch.int32, device=device)
-        seq_lens = torch.tensor([self.kv_len], dtype=torch.int32, device=device)
+        # Built on the host once per AR block; the copies go through the same
+        # pinned, non-blocking path as the slot mappings above so the CPU does
+        # not wait for the stream to drain three times before the first layer.
+        block_table = _to_device_async(torch.tensor([padded], dtype=torch.int32), device)
+        query_start_loc = _to_device_async(torch.tensor([0, self.query_len], dtype=torch.int32), device)
+        seq_lens = _to_device_async(torch.tensor([self.kv_len], dtype=torch.int32), device)
         return block_table, query_start_loc, seq_lens, self.query_len, max_seq_len
 
     def prepare(self, device: torch.device, action_len: int, query_len: int) -> None:
@@ -218,15 +243,71 @@ class ARDiffusionPagedForwardContext:
         ) = self.build_block_table(action_len=action_len, query_len=query_len, device=device)
         if self.action_slot_mapping is None:
             self.action_slot_mapping = torch.empty(0, dtype=torch.long, device=device)
+        self._prepare_history_staging(action_len)
         self._prepared = True
 
+    def _prepare_history_staging(self, action_len: int) -> None:
+        """Decide whether this forward can keep the history already staged, and from where.
+
+        The staged window is only reusable when nothing about it moved: the same session adapter, the same
+        visible history blocks, the same sequence bounds. Any difference -- a new AR block, a slid window, a
+        different request -- and the buffer is rebuilt from the pools.
+
+        The decision is a host bool, deliberately. It has to be, for the gather to actually be skipped: a
+        device flag could only select between two results already computed. It is the same value for all
+        layers of a forward, so it costs two compiled variants in total rather than one per layer.
+        """
+        self.staging_enabled = False
+        self.reuse_history = False
+        self.stage_first_block = 0
+        # The manager allocates the staging pairs only when reuse_history_staging is on and the contiguous
+        # gather path -- their only consumer -- is switched on.
+        if not self.kv_cache.history_staging:
+            return
+        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
+        capacity = int(self.kv_cache.history_staging[0][0].shape[0])
+        if int(self.max_seq_len) > capacity:
+            raise RuntimeError(
+                f"history staging buffers hold {capacity} tokens but this forward stages {int(self.max_seq_len)}"
+            )
+        state = self.kv_cache.history_staging_state
+        signature = (
+            int(self.adapter.num_computed_tokens),
+            tuple(self.history_block_ids),
+            int(self.max_seq_len),
+            int(self.seq_len),
+        )
+        same_session = state.adapter is not None and state.adapter() is self.adapter
+        reuse = same_session and state.signature == signature
+        state.adapter = weakref.ref(self.adapter)
+        state.signature = signature
+        self.staging_enabled = True
+        self.reuse_history = bool(reuse)
+        # ``kv_len`` is the visible video window here (no action tokens on this path): the current chunk's
+        # blocks are its last ``num_current_video_blocks`` entries, and the table's trailing padding blocks
+        # come after them.
+        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)."""
+        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]
+
     def layer_inputs(self, layer_idx: int) -> ARDiffusionPagedLayerInputs:
         if not getattr(self, "_prepared", False):
             raise RuntimeError("ARDiffusionPagedForwardContext.layer_inputs() before prepare()")
+        key_pool = self.kv_cache._k_pools[layer_idx]
+        value_pool = self.kv_cache._v_pools[layer_idx]
+        stage_key, stage_value = self.history_staging(layer_idx)
         return ARDiffusionPagedLayerInputs(
             layer_idx=_layer_idx_tensor(layer_idx),
-            key_pool=self.kv_cache._k_pools[layer_idx],
-            value_pool=self.kv_cache._v_pools[layer_idx],
+            key_pool=key_pool,
+            value_pool=value_pool,
             block_size=int(self.kv_cache.block_size),
             seq_len=int(self.seq_len),
             video_slots=self.current_video_slot_mapping,
@@ -236,6 +317,10 @@ class ARDiffusionPagedForwardContext:
             seq_lens=self.seq_lens,
             max_query_len=int(self.max_query_len),
             max_seq_len=int(self.max_seq_len),
+            stage_key=stage_key,
+            stage_value=stage_value,
+            reuse_history=self.reuse_history,
+            stage_first_block=int(self.stage_first_block),
         )
 
     def mark_committed(self) -> None:
@@ -342,6 +427,23 @@ def _reference_paged_attention(
 _FA_VERSION_BY_HEAD_SIZE: dict[int, int] = {}
 
 
+def _stage_window(stage_key, stage_value, key_cache, value_cache, block_ids, n_blocks, block_size, *, first_block):
+    """Gather visible blocks from ``first_block`` onward into the staging buffers.
+
+    ``first_block`` is 0 for a full restage and the index of the current chunk's first block when the
+    history is already staged. Skipping the leading blocks is the whole point: they were gathered by an
+    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:])
+    )
+
+
 def _resolve_fa_version(head_size: int) -> int:
     # get_flash_attn_version -> current_platform.get_device_capability() is not
     # dynamo-traceable, and the answer is fixed per head size for the process.
@@ -385,6 +487,11 @@ def ar_diffusion_paged_attention(
     max_seq_len: int,
     softmax_scale: float,
     causal: bool = False,
+    stage_key: torch.Tensor | None = None,
+    stage_value: torch.Tensor | None = None,
+    reuse_history: bool = False,
+    current_tokens: int = 0,
+    stage_first_block: int = 0,
 ) -> torch.Tensor:
     """Run non-causal paged attention over a vLLM block table.
 
@@ -437,6 +544,82 @@ def ar_diffusion_paged_attention(
             softmax_scale=float(softmax_scale),
             causal=causal,
         )
+    elif contiguous_kv_gather_enabled() and block_table.shape[0] == 1:
+        # Experimental (VLLM_OMNI_AR_DIFFUSION_KV_GATHER=1): FA3's paged-KV path with the
+        # frame-sized block (1560 tokens, not a multiple of the FA3 K/V tile)
+        # runs ~15-17% slower than the same kernel on contiguous K/V. Gather the
+        # visible blocks into a contiguous buffer once per layer and attend
+        # without a block table; ``seqused_k`` masks the tail-padding block.
+        from vllm.vllm_flash_attn import flash_attn_varlen_func
+
+        fa_version = _resolve_fa_version(query_flat.shape[-1])
+        block_size = key_cache.shape[1]
+        n_blocks = int(max_seq_len) // block_size
+        if n_blocks * block_size != int(max_seq_len):
+            raise ValueError("the contiguous K/V gather path requires max_seq_len to be block-aligned")
+        block_ids = block_table[0, :n_blocks].to(torch.long)
+        if stage_key is None or stage_value is None:
+            # Fresh allocations on purpose: a module-level cached buffer that is first
+            # allocated inside a CUDA-graph-trees warm-up run lives in the graph pool
+            # untracked ("tensor(s) in the cudagraph pool not tracked as outputs").
+            k_buf = key_cache.index_select(0, block_ids)
+            v_buf = value_cache.index_select(0, block_ids)
+            k_flat = k_buf.view(n_blocks * block_size, *key_cache.shape[2:])
+            v_flat = v_buf.view(n_blocks * block_size, *value_cache.shape[2:])
+        else:
+            # Caller-owned staging, allocated outside any capture and marked static, so the buffer this
+            # writes into is the same address every forward and the graph pool never owns it.
+            #
+            # Only the current chunk moved. The block table is the visible history followed by the
+            # current chunk's blocks and then padding (at least the action-capacity block, plus unused
+            # window capacity while the window is still growing), so the rows this forward changed start
+            # at ``stage_first_block`` -- the caller's count of visible history blocks -- and NOT at
+            # ``n_blocks - current_blocks``, which would refresh padding and leave the current K/V stale.
+            # Everything before them was staged by an earlier forward of the same AR block and is still
+            # 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
+            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.
+                current_blocks = current_tokens // block_size
+                assert current_tokens > 0 and current_blocks * block_size == current_tokens, (
+                    f"staged reuse needs a block-aligned current chunk, got {current_tokens} tokens"
+                )
+                assert 0 <= stage_first_block and stage_first_block + current_blocks <= n_blocks, (
+                    f"stage_first_block={stage_first_block} + {current_blocks} blocks "
+                    f"exceeds the {n_blocks}-block window"
+                )
+            _stage_window(
+                k_flat,
+                v_flat,
+                key_cache,
+                value_cache,
+                block_ids,
+                n_blocks,
+                block_size,
+                first_block=stage_first_block if reuse_history else 0,
+            )
+        out = torch.empty_like(query_flat)
+        # Varlen K/V: cu_seqlens_k = [0, kv_len] built on device from seq_lens
+        # (no host sync). Rows past kv_len (the tail-padding block) are never
+        # read, exactly like the ROCm gather path above.
+        cu_seqlens_k = torch.cat((seq_lens.new_zeros(1), seq_lens.to(torch.int32)))
+        flash_attn_varlen_func(
+            q=query_flat,
+            k=k_flat,
+            v=v_flat,
+            out=out,
+            cu_seqlens_q=query_start_loc,
+            max_seqlen_q=int(max_query_len),
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_k=int(max_seq_len),
+            softmax_scale=float(softmax_scale),
+            causal=causal,
+            fa_version=fa_version,
+        )
     else:
         from vllm.vllm_flash_attn import flash_attn_varlen_func
 
@@ -488,6 +671,10 @@ def _paged_write_attn_impl(
     max_query_len: int,
     max_seq_len: int,
     softmax_scale: float,
+    stage_key: torch.Tensor | None = None,
+    stage_value: torch.Tensor | None = None,
+    reuse_history: bool = False,
+    stage_first_block: int = 0,
 ) -> torch.Tensor:
     key_pool[video_slots] = k_curr.to(key_pool.dtype)
     value_pool[video_slots] = v_curr.to(value_pool.dtype)
@@ -507,16 +694,23 @@ def _paged_write_attn_impl(
         max_seq_len=max_seq_len,
         softmax_scale=softmax_scale,
         causal=False,
+        stage_key=stage_key,
+        stage_value=stage_value,
+        reuse_history=reuse_history,
+        current_tokens=int(k_curr.shape[0]),
+        stage_first_block=stage_first_block,
     )
 
 
 # hasattr guard keeps registration idempotent across test re-imports that pop
 # the module from sys.modules (same as sage_attn3.py).
 if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"):
-
+    # Keep staging arguments required even when their value is None. The dispatcher
+    # strips trailing defaults, and older PyTorch mutation handlers index the
+    # positional arguments without restoring them before bumping version counters.
     @torch.library.custom_op(
         "vllm_omni::ar_diffusion_paged_write_attn",
-        mutates_args=("key_pool", "value_pool"),
+        mutates_args=("key_pool", "value_pool", "stage_key", "stage_value"),
     )
     def _paged_write_attn_op(
         query: torch.Tensor,
@@ -535,6 +729,10 @@ if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"):
         max_query_len: int,
         max_seq_len: int,
         softmax_scale: float,
+        stage_key: torch.Tensor | None,
+        stage_value: torch.Tensor | None,
+        reuse_history: bool,
+        stage_first_block: int,
     ) -> torch.Tensor:
         return _paged_write_attn_impl(
             query,
@@ -553,6 +751,10 @@ if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"):
             max_query_len,
             max_seq_len,
             softmax_scale,
+            stage_key,
+            stage_value,
+            reuse_history,
+            stage_first_block,
         )
 
     @_paged_write_attn_op.register_fake
@@ -573,6 +775,10 @@ if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"):
         max_query_len,
         max_seq_len,
         softmax_scale,
+        stage_key=None,
+        stage_value=None,
+        reuse_history=False,
+        stage_first_block=0,
     ):
         return torch.empty_like(query)
 
@@ -598,4 +804,8 @@ def paged_write_attn(
         inputs.max_query_len,
         inputs.max_seq_len,
         softmax_scale,
+        inputs.stage_key,
+        inputs.stage_value,
+        inputs.reuse_history,
+        inputs.stage_first_block,
     )