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 TYPE_CHECKING, 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, visible_window_blocks

if TYPE_CHECKING:
    from vllm.v1.attention.backend import MultipleOf

_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
    # Tokens already committed when this forward began. Recorded rather than
    # derived from the block count: the last block of the history is only
    # partly written when a chunk is not a whole number of blocks, so
    # blocks * block_size overstates it and would make attention read past
    # what was written.
    _history_tokens: int = 0
    # Scratch blocks the current video tokens occupy. Not the same as
    # len(current_video_block_ids): when the history ends mid-block, the first
    # entry there is the history's own tail block, which is managed, not
    # scratch. Action K/V is placed after this, so it must not count it.
    _scratch_blocks_used: int = 0
    # Slot offset of this chunk's first token inside its first block. The
    # chunk's blocks then hold positions ``history - offset + j * block_size``,
    # which is how the read path places them among the history's blocks.
    _current_offset: 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 chunk_size(self) -> int:
        """Tokens in one chunk -- one latent frame. The eviction unit."""
        return int(self.kv_cache.spec.chunk_size)

    @property
    def start_offset(self) -> int:
        """Where in its first block this forward's tokens begin.

        Non-zero exactly when the committed history does not end on a block
        boundary, which is the normal case once a frame is not a whole number
        of blocks: 1560 tokens per frame against 16-token blocks leaves the
        history 8 slots into its last block on every odd-numbered chunk.
        """
        return int(self.adapter.num_computed_tokens) % self.block_size

    @property
    def sink_tokens(self) -> int:
        """Leading tokens the window keeps however far it has slid."""
        return int(self.kv_cache.spec.sink_chunks) * self.chunk_size

    @property
    def window_tokens(self) -> int:
        """Most recent tokens the window keeps, this forward's own included."""
        return self.max_video_tokens - self.sink_tokens

    @property
    def max_video_blocks(self) -> int:
        """Most blocks the visible window can span, which fixes the table width.

        The window is two token ranges, the sink and the most recent tokens, and
        each is converted to blocks on its own because each can straddle a block
        edge on its own. Rounding the two ranges' *sum* up once was a block short
        whenever both straddled: at 832x480 a 9-frame sink and a 9-frame window
        are each 8 tokens past a 16-token edge but 28080 tokens together, a whole
        number of blocks, so it was short on every tick once the window had
        filled.

        The sink ends at a fixed position, so its blocks are just its tokens
        rounded up. The recent range starts at a whole number of chunks, so its
        start sits at most ``block_size - gcd(chunk_size, block_size)`` tokens
        past a block edge, and it can need that much more room. A chunk that is
        a whole number of blocks never straddles, so every geometry that paged
        one frame per block keeps exactly the width it had.

        The width is a constant of the configuration, which is what keeps the
        block table a fixed shape as the window fills.
        """
        return visible_window_blocks(
            chunk_size=self.chunk_size,
            block_size=self.block_size,
            sink_tokens=self.sink_tokens,
            window_tokens=self.window_tokens,
        )

    @property
    def num_current_video_blocks(self) -> int:
        """Blocks this forward's tokens occupy, counted from where they start.

        A chunk need not be a whole number of blocks, so the count depends on
        the offset it begins at, not only on its length: thirty tokens
        starting at position thirty span three sixteen-token blocks, not two.
        Rounding up leaves the tail of the last block unwritten, which is only
        safe because nothing reads past ``kv_len`` -- see
        :meth:`video_block_table`.

        Both paths count the same way. The committing path writes straight
        after the history; the scratch path is made to line up with it by
        :meth:`ensure_video_slots`.
        """
        return -(-(self.start_offset + 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
        self._history_tokens = int(self.adapter.num_computed_tokens)
        if self.commit_current:
            start = int(self.adapter.num_computed_tokens)
            self._current_offset = start % self.block_size
            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:
            # The scratch region cannot simply start at slot zero. The kernel
            # reads the history and this chunk as one contiguous run, so if the
            # history stopped mid-block, starting here at zero would leave the
            # rest of that block unwritten and the kernel would read those dead
            # slots as if they were tokens, shifting the whole sequence.
            #
            # Instead the chunk begins where the history left off, spilling its
            # first few tokens into the free tail of the history's own last
            # block. That block is already this session's -- blocks are
            # allocated whole -- and those slots hold nothing yet; the
            # committing forward overwrites the same slots with the clean K/V
            # later, so the committed state is unchanged either way.
            # A history whose tail block is no longer resident has nothing to
            # spill into; the window then starts at this chunk and zero is right.
            offset = self.start_offset if self.history_block_ids else 0
            self._current_offset = offset
            self._scratch_blocks_used = n_blocks - (1 if offset else 0)
            scratch_ids = self.kv_cache.scratch_block_ids(self.kv_branch, 0, self._scratch_blocks_used)
            tail_block = [self.history_block_ids[-1]] if offset else []
            self.current_video_block_ids = tail_block + scratch_ids
            positions = torch.arange(offset, offset + 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
        # Count scratch blocks, not entries: the video list may lead with the
        # history's tail block, which lives in the managed pool.
        scratch_offset = 0 if self.commit_current else self._scratch_blocks_used
        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]:
        """Blocks the attention reads, and how many of their slots it reads.

        The kernel reads the listed blocks as one contiguous run and stops after
        the returned length, so two things must hold: every block holding a
        token the window keeps is listed, in token order, and only the last
        listed block may be partly written.

        Blocks are chosen by the token positions they hold rather than counted
        off the ends of the resident list. The window is the sink ``[0, S)``
        plus the most recent tokens, and each of those can straddle a block
        edge on its own. Taking a fixed number of blocks off the tail dropped
        the block holding the window's first tokens exactly when both
        straddled, and a length capped by a block count rather than counted
        from what was written then ran past the last written slot.

        Whole blocks are read, so each boundary block contributes up to
        ``block_size - 1`` tokens the window has already let go of. They are
        real K/V at their own positions; a slightly wider window is a far
        smaller error than a truncated one.
        """
        self.ensure_video_slots(device)
        block_size = self.block_size
        history = self._history_tokens
        end = history + self.seq_len
        sink_end = self.sink_tokens
        recent_start = end - self.window_tokens

        # Position of each block's first slot. History comes from the block
        # table, where an evicted block is a null placeholder, so a block's index
        # is its position -- in storage positions, which is what
        # num_computed_tokens counts once compaction has dropped an evicted gap.
        # Compaction shifts storage by whole chunks and whole blocks, so the sink,
        # the window and every block offset sit exactly where they would in model
        # positions. Reading the table after ensure_video_slots means a block
        # this chunk's own allocation evicted is never listed. The current
        # chunk's blocks follow from where its first token was placed; its first
        # block is the history's own tail block whenever the history stopped
        # part way through one, and keeps the history's position.
        first_position: dict[int, int] = {}
        #
        # Only two index ranges can hold a token the window keeps: the sink's
        # blocks, and the blocks from the recent window's start onward. Evicted
        # positions stay in the table as null entries until compaction drops
        # them a whole number of chunks at a time, so the table can still be
        # much longer than the window; reading just these two ranges, without
        # copying the table, keeps this at the window's size.
        history_blocks = -(-history // block_size)
        sink_blocks = min(-(-sink_end // block_size), history_blocks)
        recent_first_block = max(max(recent_start, 0) // block_size, sink_blocks)
        indices = [*range(sink_blocks), *range(recent_first_block, history_blocks)]
        for index, block in zip(indices, self.kv_cache.block_ids_at(self.adapter, indices)):
            if block != self.kv_cache.null_block_id:
                first_position.setdefault(block, index * block_size)
        current_start = history - self._current_offset
        for index, block in enumerate(self.current_video_block_ids):
            first_position.setdefault(int(block), current_start + index * block_size)

        visible_video_blocks: list[int] = []
        video_len = 0
        for block, first in sorted(first_position.items(), key=lambda item: item[1]):
            if first >= sink_end and first + block_size <= recent_start:
                continue  # wholly between the sink and the recent window
            if video_len % block_size:
                raise RuntimeError(
                    "AR-Diffusion paged attention would read past a partly written block: "
                    f"block {visible_video_blocks[-1]} is followed by block {block}."
                )
            visible_video_blocks.append(block)
            video_len += min(block_size, end - first)
        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)
        # Action K/V is listed after the video blocks and the kernel reads the
        # table as one run, so the video run has to end on a block edge. If it
        # stops part way through its last block, the run reads that block's
        # unwritten slots as the first action tokens and never reaches the last
        # ones -- a wrong attention with every shape intact. A frame that is a
        # whole number of blocks always ends on an edge, so this refuses only
        # geometries that could not page at all before the paging unit was
        # separated from the frame.
        if action_len > 0 and video_len % self.block_size:
            raise ValueError(
                "AR-Diffusion paged attention cannot place action tokens after a partly written video block: "
                f"the video run ends {video_len % self.block_size} slots into block {video_blocks[-1]}, so the "
                "kernel would read its unwritten slots as action tokens. Action tokens need tokens_per_frame "
                f"to be a whole number of blocks (block_size={self.block_size})."
            )
        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)
        # Both of these count in blocks. Deriving them from the token count by
        # flooring would understate the capacity whenever the window is not a
        # whole number of blocks: the width would fall back on len(block_ids)
        # and start tracking how full the window is -- the exact shape churn
        # the fixed width exists to prevent -- and max_seq_len could come out
        # below the kv_len actually being passed, by up to block_size - 1.
        capacity_blocks = self.max_video_blocks + action_capacity_blocks
        width = max(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(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
        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
        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
        # A chunk that is not a whole number of blocks shares its first block with the history's tail (see
        # ensure_video_slots), and the attention op restages the current chunk in whole blocks of its own,
        # so such a forward restages the whole window instead.
        block_aligned = self.start_offset == 0 and self.seq_len % self.block_size == 0
        reuse = same_session and state.signature == signature and block_aligned
        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]:
        """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
        return self.kv_cache.history_staging[layer_idx]

    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:]
    # 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:]
    )


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 supported_kernel_block_sizes() -> list[int | MultipleOf]:
    """Block sizes the kernel this module dispatches to will accept.

    Same shape as vLLM's ``AttentionBackend.get_supported_kernel_block_sizes``
    -- a plain int is that exact size, ``MultipleOf(b)`` is any positive
    multiple of ``b`` -- but answered here rather than read off a backend,
    because AR-Diffusion does not go through backend selection. It calls
    ``flash_attn_varlen_func`` itself, choosing between vLLM's CUDA build and
    ROCm's AITER a few lines below, and only this module knows which.

    It lives next to that choice so there is one place to change. The
    constraint is a property of the kernel, not of the card, and the kernels
    reachable from here agree on 16 today: vLLM's CUDA FlashAttention, ROCm
    AITER, and upstream ``flash_attn`` all advertise ``MultipleOf(16)``. Other
    backends in the same tree do not -- ``hpc_attn`` accepts only 64, and
    FlashInfer advertises pages of 128 or more solely on Blackwell -- so a
    caller must treat this as data to be queried, never as the number 16.

    vLLM's FlashAttention backend does advertise a single 128-token page when FA4
    runs its dedicated head-size-256 kernel on SM100/SM110. That kernel is not
    reachable from here: ``get_flash_attn_version`` selects it only when the
    caller passes ``supports_fa4_hd256=True``, and :func:`_resolve_fa_version`
    does not, so a head size of 256 falls back to FA2, which pages at any
    multiple of 16. Passing that flag later has to change this answer as well.
    """
    from vllm.v1.attention.backend import MultipleOf

    return [MultipleOf(16)]


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.
        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")
        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 = rocm_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.
            # 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.
                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 as paged_flash_attn_varlen_func

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

        out = torch.empty_like(query_flat)
        paged_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,
    )