PagedAttention#

  • Domain: ai.rt

  • Since version: 1

Appends immutable KV pages and computes grouped-query attention for finite Q/K/V tensors of shape [batch,heads,L,D]. Q may have a multiple of the K/V head count; K and V have equal head counts and all inputs have equal batch and new sequence dimensions. The past/present values are structured caches whose blocks form a dynamic sequence of pages with INT64 scalar start/length and logical key/value tensors [batch,kv_heads,capacity,head_size]. Runtime pages may store dense FLOAT/FLOAT16/BFLOAT16 or inline affine integer data. Shape inference produces Y with Q’s batch, head and sequence dimensions, V’s head size and the input floating-point type. Storage policy belongs to the registered kernel and is not part of the operator attributes.

Inputs

  • Q (T): Queries [batch,q_heads,L,key_head_size].

  • K (T): New keys [batch,kv_heads,L,key_head_size].

  • V (T): New values [batch,kv_heads,L,value_head_size].

  • past (C): Paged cache structured value.

Outputs

  • Y (T): Attention result [batch,q_heads,L,value_head_size].

  • present (C): Cache after appending the new pages.

Attributes

  • block_size (int): Positive maximum token capacity for each new page.

  • is_causal (int): Whether to apply causal masking; only 0 or 1 is valid.

  • left_window_size (int): Past-token window; -1 is unbounded and non-negative values limit it.

  • max_tokens (int): Positive maximum retained token count, including past and new tokens.

Type Constraints

  • T: Constrain Q/K/V and Y to one floating-point type. Allowed types: tensor(bfloat16), tensor(float), tensor(float16).

  • C: Named paged-cache structure. Allowed types: struct.