vllm-omni/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py

# 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] = {}


def _layer_idx_tensor(layer_idx: int) -> torch.Tensor:
    t = _LAYER_IDX_TENSORS.get(layer_idx)
    if t is None:
        t = torch.tensor(layer_idx, dtype=torch.int64)
        _LAYER_IDX_TENSORS[layer_idx] = t
    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.

    A NamedTuple of plain tensors + ints so ``torch.compile`` treats every field
    as a pytree graph input (no object-attribute guards, no recompiles when only
    tensor *values* change). All layers of one KV branch forward share the same
    metadata tensor objects, built once by ``prepare()``.

    ``layer_idx`` is a 0-dim CPU tensor, NOT a python int: all 40 DiT blocks
    share one compiled code object, and an int here becomes a per-layer dynamo
    value guard (``layer_idx == k``) — 40 cache variants that blow the
    recompile limit. A tensor input guards on shape/dtype only, so one graph
    serves every layer.
    """

    layer_idx: torch.Tensor
    key_pool: torch.Tensor
    value_pool: torch.Tensor
    block_size: int
    seq_len: 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
    # 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
class ARDiffusionPagedForwardContext:
    """Mutable KV-branch state shared by all layer contexts in one forward."""

    kv_cache: Any
    adapter: Any
    kv_branch: str
    history_block_ids: list[int]
    seq_len: int
    commit_current: bool
    max_video_tokens: int
    current_video_block_ids: list[int] = field(default_factory=list)
    current_video_slot_mapping: torch.Tensor | None = None
    action_scratch_block_ids: list[int] = field(default_factory=list)
    action_slot_mapping: torch.Tensor | None = None
    query_len: int = 0
    kv_len: int = 0
    _allocated_video: bool = False
    _committed: bool = False
    _action_len: int = 0
    # Set once by prepare(); shared by all layers of the KV branch forward.
    block_table: torch.Tensor | None = None
    query_start_loc: torch.Tensor | None = None
    seq_lens: torch.Tensor | None = None
    max_query_len: int = 0
    max_seq_len: int = 0
    _prepared: bool = False

    @property
    def block_size(self) -> int:
        return int(self.kv_cache.block_size)

    @property
    def num_current_video_blocks(self) -> int:
        if self.seq_len % self.block_size != 0:
            raise AssertionError(
                "AR-Diffusion paged attention expects frame-aligned seq_len "
                f"(multiple of block_size={self.block_size}), got {self.seq_len}"
            )
        return self.seq_len // self.block_size

    def ensure_video_slots(self, device: torch.device) -> None:
        """Allocate/write targets for the current video tokens, once per KV branch."""
        if self._allocated_video:
            return

        n_blocks = self.num_current_video_blocks
        if self.commit_current:
            start = int(self.adapter.num_computed_tokens)
            self.kv_cache.allocate_token_slots(self.adapter, self.seq_len)
            table = self.kv_cache.block_table(self.adapter)
            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 = _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 = _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:
        """Reserve scratch slots for action/state K/V, if present."""
        if action_len <= 0:
            self.action_scratch_block_ids = []
            self.action_slot_mapping = torch.empty(0, dtype=torch.long, device=device)
            self._action_len = 0
            return

        self.ensure_video_slots(device)
        if self.action_slot_mapping is not None and self._action_len == action_len:
            return

        action_blocks = (action_len + self.block_size - 1) // self.block_size
        scratch_offset = 0 if self.commit_current else len(self.current_video_block_ids)
        self.action_scratch_block_ids = self.kv_cache.scratch_block_ids(
            self.kv_branch,
            scratch_offset,
            action_blocks,
        )
        positions = torch.arange(action_len, dtype=torch.long)
        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]:
        self.ensure_video_slots(device)
        if self.max_video_tokens % self.block_size != 0:
            raise AssertionError(
                "AR-Diffusion paged attention requires max_video_tokens to be block-aligned, "
                f"got max_video_tokens={self.max_video_tokens}, block_size={self.block_size}"
            )
        all_video_blocks = self.history_block_ids + self.current_video_block_ids
        max_video_blocks = self.max_video_tokens // self.block_size
        sink_blocks = int(self.kv_cache.spec.sink_chunks)
        if len(all_video_blocks) <= max_video_blocks:
            visible_video_blocks = all_video_blocks
        else:
            tail_blocks = max_video_blocks - sink_blocks
            visible_video_blocks = all_video_blocks[:sink_blocks]
            if tail_blocks:
                visible_video_blocks += all_video_blocks[-tail_blocks:]
        video_len = len(visible_video_blocks) * self.block_size
        return visible_video_blocks, video_len

    def build_block_table(
        self,
        *,
        action_len: int,
        query_len: int,
        device: torch.device,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
        """Build FlashAttention block-table metadata for one self-attn call.

        The block table is tail-padded to a fixed width and ``max_seq_len`` is a
        constant upper bound, so across window growth only tensor *values* change
        — tensor shapes and the int consts stay stable for ``torch.compile``
        (and, later, CUDA-graph capture). The kernel only dereferences the first
        ``ceil(seq_lens/block_size)`` entries, so padding is never read.
        """
        video_blocks, video_len = self.video_block_table(device)
        self.ensure_action_slots(action_len, device)
        action_blocks = self.action_scratch_block_ids if action_len > 0 else []
        block_ids = video_blocks + action_blocks
        if not block_ids:
            raise RuntimeError("AR-Diffusion paged attention needs at least current video KV blocks")

        # Fixed capacity: full visible video window + one action-capacity block.
        action_capacity_blocks = max(1, (action_len + self.block_size - 1) // self.block_size)
        width = max(self.max_video_tokens // self.block_size + action_capacity_blocks, len(block_ids))
        padded = block_ids + [0] * (width - len(block_ids))

        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)
        # 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:
        """Host-side, once-per-KV-branch setup (called OUTSIDE torch.compile).

        Allocates the current video/action slots (still lazy: only the KV branch a
        CFG-parallel rank actually runs reaches its ``_forward_blocks``), builds
        the padded block-table metadata ONCE for all layers, and publishes the
        pool registry for the fused custom op. The compiled per-layer code then
        only consumes prebuilt tensors via ``ARDiffusionPagedLayerInputs``.
        """
        if getattr(self, "_prepared", False):
            return
        self.ensure_video_slots(device)
        (
            self.block_table,
            self.query_start_loc,
            self.seq_lens,
            self.max_query_len,
            self.max_seq_len,
        ) = 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=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,
            action_slots=self.action_slot_mapping,
            block_table=self.block_table,
            query_start_loc=self.query_start_loc,
            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:
        self._committed = True


@dataclass
class ARDiffusionPagedLayerContext:
    """Layer-specific handle passed through a model's ``kv_cache`` slot."""

    is_ar_diffusion_paged_context: ClassVar[bool] = True
    layer_idx: int
    forward_ctx: ARDiffusionPagedForwardContext

    @property
    def kv_cache(self):
        return self.forward_ctx.kv_cache

    @property
    def adapter(self):
        return self.forward_ctx.adapter

    @property
    def kv_branch(self) -> str:
        return self.forward_ctx.kv_branch

    @property
    def history_block_ids(self) -> list[int]:
        return self.forward_ctx.history_block_ids

    @property
    def current_video_block_ids(self) -> list[int]:
        return self.forward_ctx.current_video_block_ids

    @property
    def current_video_slot_mapping(self) -> torch.Tensor | None:
        return self.forward_ctx.current_video_slot_mapping

    @property
    def action_scratch_block_ids(self) -> list[int]:
        return self.forward_ctx.action_scratch_block_ids

    @property
    def action_slot_mapping(self) -> torch.Tensor | None:
        return self.forward_ctx.action_slot_mapping

    @property
    def seq_len(self) -> int:
        return self.forward_ctx.seq_len

    @property
    def query_len(self) -> int:
        return self.forward_ctx.query_len

    @property
    def kv_len(self) -> int:
        return self.forward_ctx.kv_len

    @property
    def commit_current(self) -> bool:
        return self.forward_ctx.commit_current

    def to_layer_inputs(self) -> ARDiffusionPagedLayerInputs:
        """Compiled-region payload; requires ``forward_ctx.prepare()`` first."""
        return self.forward_ctx.layer_inputs(self.layer_idx)


def is_ar_diffusion_paged_context(value: object) -> bool:
    return isinstance(value, ARDiffusionPagedLayerContext)


def _reference_paged_attention(
    query: torch.Tensor,
    key_cache: torch.Tensor,
    value_cache: torch.Tensor,
    block_table: torch.Tensor,
    query_start_loc: torch.Tensor,
    seq_lens: torch.Tensor,
    softmax_scale: float,
    *,
    causal: bool,
) -> torch.Tensor:
    if causal:
        raise NotImplementedError("AR-Diffusion paged self-attention uses causal=False")
    outs: list[torch.Tensor] = []
    block_size = key_cache.shape[1]
    for i in range(seq_lens.shape[0]):
        q_start = int(query_start_loc[i].item())
        q_end = int(query_start_loc[i + 1].item())
        kv_len = int(seq_lens[i].item())
        q = query[q_start:q_end]
        positions = torch.arange(kv_len, device=query.device)
        logical_blocks = torch.div(positions, block_size, rounding_mode="floor")
        offsets = positions % block_size
        physical_blocks = block_table[i, logical_blocks].long()
        k = key_cache[physical_blocks, offsets]
        v = value_cache[physical_blocks, offsets]
        scores = torch.einsum("qhd,khd->hqk", q.float(), k.float()) * float(softmax_scale)
        probs = torch.softmax(scores, dim=-1).to(v.dtype)
        outs.append(torch.einsum("hqk,khd->qhd", probs, v))
    return torch.cat(outs, dim=0)


_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.
    version = _FA_VERSION_BY_HEAD_SIZE.get(head_size)
    if version is None:
        try:
            from vllm.v1.attention.backends.fa_utils import get_flash_attn_version

            version = int(get_flash_attn_version(requires_alibi=False, head_size=head_size) or 2)
        except Exception:
            version = 2
        _FA_VERSION_BY_HEAD_SIZE[head_size] = version
    return version


def _rocm_flash_attn_varlen_func():
    """Resolve a ROCm varlen kernel, preferring AITER when available.

    The caller gathers paged KV into packed tensors before invoking this
    function, avoiding AITER releases whose ``block_table`` kernel is broken.
    """
    try:
        from aiter import flash_attn_varlen_func

        return flash_attn_varlen_func
    except ImportError:
        from flash_attn import flash_attn_varlen_func

        return flash_attn_varlen_func


def ar_diffusion_paged_attention(
    query: torch.Tensor,
    key_cache: torch.Tensor,
    value_cache: 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,
    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.

    ``query`` may be ``(B, L, H, D)`` or already flattened as ``(T, H, D)``.
    ``key_cache`` / ``value_cache`` are ``(num_blocks, block_size, H, D)``.
    """
    batched = query.dim() == 4
    if batched:
        batch, q_len = query.shape[:2]
        query_flat = query.reshape(batch * q_len, *query.shape[2:])
    else:
        query_flat = query

    if not query_flat.is_cuda:
        out = _reference_paged_attention(
            query_flat,
            key_cache,
            value_cache,
            block_table,
            query_start_loc,
            seq_lens,
            softmax_scale,
            causal=causal,
        )
    elif torch.version.hip is not None:
        # vllm.vllm_flash_attn contains CUDA-only extensions. ROCm's AITER and
        # upstream flash-attn expose the standard cu_seqlens_k API instead of
        # 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()
        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")
        offsets = positions % key_cache.shape[1]
        physical_blocks = block_table[:, logical_blocks].long()
        gathered_k = key_cache[physical_blocks, offsets]
        gathered_v = value_cache[physical_blocks, offsets]
        valid = positions.unsqueeze(0) < seq_lens.unsqueeze(1)
        packed_k = gathered_k[valid]
        packed_v = gathered_v[valid]
        out = flash_attn_varlen_func(
            q=query_flat,
            k=packed_k,
            v=packed_v,
            cu_seqlens_q=query_start_loc,
            cu_seqlens_k=cu_seqlens_k,
            max_seqlen_q=int(max_query_len),
            max_seqlen_k=int(max_seq_len),
            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

        fa_version = _resolve_fa_version(query_flat.shape[-1])

        out = torch.empty_like(query_flat)
        flash_attn_varlen_func(
            q=query_flat,
            k=key_cache,
            v=value_cache,
            out=out,
            cu_seqlens_q=query_start_loc,
            max_seqlen_q=int(max_query_len),
            seqused_k=seq_lens,
            max_seqlen_k=int(max_seq_len),
            softmax_scale=float(softmax_scale),
            causal=causal,
            block_table=block_table,
            fa_version=fa_version,
        )

    if batched:
        return out.reshape(query.shape)
    return out


# ── Fused write+attend custom op (torch.compile-safe) ──────────────────────
#
# One opaque op per layer keeps the compiled DiT block fullgraph: dynamo treats
# it as a single graph node (no eager island, no graph breaks), and the K/V slot
# writes happen inside the op so write→read ordering with the block-table kernel
# is internal. The flat pools are explicit mutable inputs: Inductor/CUDA Graph
# must track their storage lifetime instead of observing an undeclared mutation
# through a process-global registry.
def _paged_write_attn_impl(
    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,
    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)
    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)
    key_cache = key_pool.unflatten(0, (-1, block_size))
    value_cache = value_pool.unflatten(0, (-1, block_size))
    return ar_diffusion_paged_attention(
        query,
        key_cache,
        value_cache,
        block_table=block_table,
        query_start_loc=query_start_loc,
        seq_lens=seq_lens,
        max_query_len=max_query_len,
        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", "stage_key", "stage_value"),
    )
    def _paged_write_attn_op(
        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,
        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,
            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,
            stage_key,
            stage_value,
            reuse_history,
            stage_first_block,
        )

    @_paged_write_attn_op.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,
        stage_key=None,
        stage_value=None,
        reuse_history=False,
        stage_first_block=0,
    ):
        return torch.empty_like(query)


def paged_write_attn(
    inputs: ARDiffusionPagedLayerInputs, query, k_curr, v_curr, k_act, v_act, softmax_scale: float
) -> torch.Tensor:
    """Model-facing entry: routes through the custom op (traceable in fullgraph)."""
    return torch.ops.vllm_omni.ar_diffusion_paged_write_attn(
        query,
        k_curr,
        v_curr,
        k_act,
        v_act,
        inputs.key_pool,
        inputs.value_pool,
        inputs.block_size,
        inputs.video_slots,
        inputs.action_slots,
        inputs.block_table,
        inputs.query_start_loc,
        inputs.seq_lens,
        inputs.max_query_len,
        inputs.max_seq_len,
        softmax_scale,
        inputs.stage_key,
        inputs.stage_value,
        inputs.reuse_history,
        inputs.stage_first_block,
    )