Accelerated API

def forward(
    self,
    x: Tensor,
    scale_or_granularity: Tensor | Granularity,
    out_dtype: torch.dtype = torch.float16,
) -> Tensor:
    """Apply the quantized linear transformation.

    Args:
        x: Activations shaped ``[..., in_features]``. When a scale tensor
            is supplied, these must already use the layer's quantized
            ``dtype``; when a granularity is supplied, these are
            full-precision activations quantized inside this call.
        scale_or_granularity: For prequantized ``x``, an FP32 tensorwise
            scale shaped ``[1, ..., 1]`` or slice scale shaped
            ``[..., 1]``. For full-precision ``x``, the granularity used
            to produce one of those scale shapes.
        out_dtype: Data type of the projected activations. Defaults to
            ``torch.float16``.

    Returns:
        Projected activations shaped ``[..., out_features]``.

    Raises:
        ValueError: ``x`` or its scale does not match the layer contract.
    """
    if x.ndim == 0 or x.shape[-1] != self.in_features:
        raise ValueError(
            f"expected input last dim {self.in_features}, got {tuple(x.shape)}"
        )

    if isinstance(scale_or_granularity, Tensor):
        # Prequantized path: preserve ``x [..., I]`` and its existing
        # tensorwise ``scale [1, ..., 1]`` or slice ``scale [..., 1]``.
        if x.dtype is not self.dtype:
            raise ValueError(
                f"expected quantized input dtype {self.dtype}, got {x.dtype}"
            )
        self._validate_scale(x, scale_or_granularity)
        quantized, scale = x, scale_or_granularity
    elif isinstance(scale_or_granularity, Granularity):
        # Dynamic path: quantization keeps ``x [..., I]`` unchanged in
        # shape and reduces either every dimension or only ``I`` for scale.
        if x.numel() == 0:
            return torch.empty(
                (*x.shape[:-1], self.out_features),
                device=x.device,
                dtype=out_dtype,
            )
        quantized, scale = quantize(
            x,
            self.dtype,
            scale_or_granularity,
            axis=-1,
        )
    else:
        raise ValueError(
            "scale_or_granularity must be a scale tensor or Granularity"
        )

    return self._forward_quantized(quantized, scale, out_dtype)

def __post_init__(self) -> None:
    """Validate and normalize attention configuration values."""
    context_dim = self.query_dim if self.context_dim is None else self.context_dim
    if self.query_dim <= 0:
        raise ValueError(f"query_dim must be positive; got {self.query_dim}")
    if context_dim <= 0:
        raise ValueError(f"context_dim must be positive; got {context_dim}")
    if self.n_heads <= 0:
        raise ValueError(f"n_heads must be positive; got {self.n_heads}")
    n_kv_heads = self.n_heads if self.n_kv_heads is None else self.n_kv_heads
    if n_kv_heads <= 0:
        raise ValueError(f"n_kv_heads must be positive; got {n_kv_heads}")
    if self.n_heads % n_kv_heads:
        raise ValueError(
            "n_heads must be divisible by n_kv_heads; "
            f"got {self.n_heads} and {n_kv_heads}"
        )
    if self.head_dim <= 0:
        raise ValueError(f"head_dim must be positive; got {self.head_dim}")
    if not isinstance(self.qk_norm_scope, QKNormScope):
        raise TypeError(
            f"qk_norm_scope must be a QKNormScope; got {self.qk_norm_scope!r}"
        )
    if not math.isfinite(self.qk_norm_eps) or self.qk_norm_eps <= 0:
        raise ValueError(
            f"qk_norm_eps must be finite and positive; got {self.qk_norm_eps}"
        )
    if self.rope_config is not None and not isinstance(
        self.rope_config, RoPEConfig
    ):
        raise TypeError(
            f"rope_config must be a RoPEConfig or None; got {self.rope_config!r}"
        )
    object.__setattr__(self, "context_dim", context_dim)
    object.__setattr__(self, "n_kv_heads", n_kv_heads)

def __post_init__(self) -> None:
    """Validate rotary policy enum values."""
    if not isinstance(self.style, RoPEStyle):
        raise TypeError(f"style must be a RoPEStyle; got {self.style!r}")
    if not isinstance(self.scope, RoPEScope):
        raise TypeError(f"scope must be a RoPEScope; got {self.scope!r}")

@abstractmethod
def compute_kv(
    self,
    context: Tensor,
    rope_freqs: Tensor | None = None,
) -> KVCacheT:
    """Project context and return a precomputed K/V cache.

    Use this stage to materialize static cross-attention context. The cache
    is ready for repeated :meth:`forward` calls; implementations decide its
    physical layout, precision, and ownership.

    Args:
        context: Key/value source, shape ``[..., S, context_dim]``.
        rope_freqs: Optional key positional data for the ``S`` context
            tokens. Applied only by before-cache RoPE; after-cache RoPE
            stores unrotated keys and consumes cache-relative data in
            :meth:`forward`. Ignored when ``rope_config`` is ``None``.

    Returns:
        Precomputed cache containing K/V for all ``S`` context tokens.
    """

@abstractmethod
def forward(
    self,
    x: Tensor,
    kv_cache: KVCacheT,
    rope_freqs: Tensor | None = None,
    *,
    attn_mask: AttentionMask | None = None,
) -> Tensor:
    """Apply the configured attention type to ``x`` and ``kv_cache``.

    Implementations own the complete operation so fused backends need not
    expose independently callable cache-update or cache-query stages.

    Args:
        x: Query tokens, shape ``[..., L, query_dim]``.
        kv_cache: Streaming cache for self-attention or precomputed static
            cache for cross-attention. A streaming cache must already be in
            its current-chunk update phase.
        rope_freqs: Optional positional data. Before-cache RoPE expects the
            current ``L`` query/key positions. After-cache RoPE expects
            cache-relative positions for all visible keys and selects the
            current query positions from the cache write interval. Ignored
            when ``rope_config`` is ``None``.
        attn_mask: Optional dense or block-sparse visibility mask. Backend
            implementations define which representation they support.

    Returns:
        Attention result with shape ``[..., L, query_dim]``.
    """

def allocate_kv_cache(
    self,
    batch_size: int,
    chunk_size: int,
    window_size: int,
    sink_size: int,
    device: torch.device | str,
    dtype: torch.dtype,
) -> BlockKVCache:
    """Allocate a native-precision rolling K/V cache.

    The returned cache is empty. For every chunk, call
    :meth:`BlockKVCache.before_update`, invoke self-attention, and call
    :meth:`BlockKVCache.after_update` with the same chunk index.

    Args:
        batch_size: Product ``B`` of the input's leading dimensions.
        chunk_size: Exact number of current tokens ``L`` per update.
        window_size: Number of rolling context tokens retained after the sink.
        sink_size: Number of initial context tokens that are never evicted.
        device: Device on which to allocate K/V storage.
        dtype: Data type used by K/V storage.

    Returns:
        Block cache with K/V storage shaped
        ``[B, sink_size + window_size, H, D]``.
    """
    # BlockKVCache rolls sequence dimension 1 while preserving independent
    # batch and head axes for SDPA.
    kv_heads = self.attention_config.n_kv_heads
    assert kv_heads is not None
    cache_shape = (
        batch_size,
        sink_size + window_size,
        kv_heads,
        self.attention_config.head_dim,
    )
    return BlockKVCache(
        k_shape=cache_shape,
        v_shape=cache_shape,
        seq_dim=1,
        chunk_size=chunk_size,
        window_size=window_size,
        sink_size=sink_size,
        device=device,
        dtype=dtype,
    )

def compute_kv(
    self,
    context: Tensor,
    rope_freqs: Tensor | None = None,
) -> BlockKVCache:
    """Project static context into a finalized, reusable K/V cache.

    Args:
        context: Static key/value source, shape
            ``[..., S, context_dim]``. Leading dimensions flatten into ``B``.
        rope_freqs: Optional key rotations with shape ``[S, 1, 1, D]``.
            Applied only by before-cache RoPE; after-cache RoPE stores
            unrotated keys. Ignored when ``rope_config`` is ``None``.

    Returns:
        Static cache with K/V shape ``[B, S, H, D]``, ready for repeated
        :meth:`forward` calls without lifecycle bookkeeping.

    Raises:
        ValueError: Context width or RoPE geometry is incompatible with the
            configured attention geometry.
    """
    key, value = self._project_kv(context)
    if (
        self.attention_config.rope_config is not None
        and self.attention_config.rope_config.scope is RoPEScope.BEFORE_KV_CACHE
        and rope_freqs is not None
    ):
        key = self._apply_rope(key, rope_freqs)
    # ``from_tensor`` completes its one write internally, so static context
    # can be queried immediately and never enters the rolling lifecycle.
    return BlockKVCache.from_tensor(key, value, seq_dim=1)

def forward(
    self,
    x: Tensor,
    kv_cache: BlockKVCache,
    rope_freqs: Tensor | None = None,
    *,
    attn_mask: AttentionMask | None = None,
) -> Tensor:
    """Apply self- or cross-attention using the configured cache lifecycle.

    Args:
        x: Query tokens, shape ``[..., L, query_dim]``.
        kv_cache: Prepared rolling cache for self-attention or precomputed
            static cache for cross-attention.
        rope_freqs: Optional positional data. Before-cache RoPE expects the
            current chunk. After-cache RoPE expects positions covering the
            query and visible cache. Ignored when ``rope_config`` is
            ``None``.
        attn_mask: Optional dense boolean visibility mask shaped ``[L, S]``.

    Returns:
        Output-projected tokens with the same shape as ``x``.
    """
    query_rope_freqs, key_rope_freqs = self._slice_rope_freqs(
        rope_freqs, kv_cache, x.shape[-2]
    )
    if self.attention_type is AttentionType.SELF_ATTENTION:
        kv_cache = self._update_kv(x, kv_cache, key_rope_freqs)
    return self._query_kv(
        x,
        kv_cache,
        query_rope_freqs,
        key_rope_freqs,
        attn_mask=attn_mask,
    )

def __post_init__(self) -> None:
    """Validate the projection quantization dtypes."""
    if self.projection is not None and self.projection not in DTYPE_MAX:
        raise ValueError(
            f"unsupported projection quantization dtype: {self.projection}"
        )
    if (
        self.output_projection is not None
        and self.output_projection not in DTYPE_MAX
    ):
        raise ValueError(
            "unsupported output projection quantization dtype: "
            f"{self.output_projection}"
        )
    if not isinstance(self.quantized_sdpa, bool):
        raise TypeError(
            f"quantized_sdpa must be a bool; got {self.quantized_sdpa!r}"
        )

def __post_init__(self) -> None:
    """Validate optimized implementation policy values."""
    if not isinstance(self.qkv_fusion_option, QKVFusionOption):
        raise TypeError(
            "qkv_fusion_option must be a QKVFusionOption; "
            f"got {self.qkv_fusion_option!r}"
        )
    if not isinstance(self.quantization, QuantizationOption):
        raise TypeError(
            f"quantization must be a QuantizationOption; got {self.quantization!r}"
        )
    if not isinstance(self.flex_attention, FlexAttentionOptions):
        raise TypeError(
            "flex_attention must be a FlexAttentionOptions; "
            f"got {self.flex_attention!r}"
        )
    if not isinstance(self.sdpa_backend, SDPABackend):
        raise TypeError(
            f"sdpa_backend must be an SDPABackend; got {self.sdpa_backend!r}"
        )
    if not isinstance(self.use_tma, bool):
        raise TypeError(f"use_tma must be a bool; got {self.use_tma!r}")

def allocate_kv_cache(
    self,
    batch_size: int,
    chunk_size: int,
    window_size: int,
    sink_size: int,
    device: torch.device | str,
    dtype: torch.dtype,
) -> BlockKVCache:
    """Allocate a rolling cache for native or quantized SDPA.

    Args:
        batch_size: Flattened batch size ``B``.
        chunk_size: Number of current tokens ``L`` written per update.
        window_size: Number of rolling context tokens retained after the sink.
        sink_size: Number of initial context tokens that are never evicted.
        device: Device on which to allocate K/V storage.
        dtype: FP16 or BF16 activation dtype. Quantized SDPA stores the cache
            in FP8 e4m3 instead.

    Returns:
        Block cache with K/V storage shaped
        ``[B, sink_size + window_size, H, D]``.

    """
    # Keep ``[B, S, H, D]`` as the public cache shape: ``BlockKVCache``
    # rolls and slices axis 1, and both attention backends accept that logical
    # order.
    kv_heads = self.attention_config.n_kv_heads
    assert kv_heads is not None
    cache_shape = (
        batch_size,
        sink_size + window_size,
        kv_heads,
        self.attention_config.head_dim,
    )
    self._validate_cuda_device(device)
    return BlockKVCache(
        k_shape=cache_shape,
        v_shape=cache_shape,
        seq_dim=1,
        chunk_size=chunk_size,
        window_size=window_size,
        sink_size=sink_size,
        device=device,
        dtype=(
            torch.float8_e4m3fn
            if self.optimized_impl_config.quantization.quantized_sdpa
            else dtype
        ),
    )

@torch.no_grad()
def compute_kv(
    self,
    context: Tensor,
    rope_freqs: Tensor | None = None,
) -> BlockKVCache:
    """Project complete context into a reusable static K/V cache.

    Args:
        context: Context tokens shaped ``[..., S, C]``.
        rope_freqs: Optional key rotations shaped ``[S, 1, 1, D]``.
            Applied only by before-cache RoPE; after-cache RoPE stores
            unrotated keys. Ignored when ``rope_config`` is ``None``.

    Returns:
        Filled cache with logical K/V shape ``[B, S, H, D]``. Its sequence
        length and window both equal S, so subsequent forward calls read all
        context.
    """
    key, value = self._project_kv(context)
    if (
        self.attention_config.rope_config is not None
        and self.attention_config.rope_config.scope is RoPEScope.BEFORE_KV_CACHE
        and rope_freqs is not None
    ):
        # Position affects key directions used in Q·K; values are never rotated.
        key = self._apply_rope(key, rope_freqs)

    if self.optimized_impl_config.quantization.quantized_sdpa:
        key = key.to(torch.float8_e4m3fn)
        value = value.to(torch.float8_e4m3fn)

    return BlockKVCache.from_tensor(key, value, seq_dim=1)

@torch.no_grad()
def forward(
    self,
    x: Tensor,
    kv_cache: BlockKVCache,
    rope_freqs: Tensor | None = None,
    *,
    attn_mask: AttentionMask | None = None,
) -> Tensor:
    """Apply self- or cross-attention using the configured cache lifecycle.

    Self-attention updates the prepared rolling cache and computes Q in one
    backend-owned branch. Cross-attention computes Q while leaving its
    precomputed static cache unchanged.

    Args:
        x: Query tokens shaped ``[..., L, Q]``.
        kv_cache: Prepared rolling cache for self-attention or precomputed
            static cache for cross-attention.
        rope_freqs: Optional positional data. Before-cache RoPE expects the
            current chunk. After-cache RoPE expects positions covering the
            query and visible cache. Ignored when ``rope_config`` is
            ``None``.
        attn_mask: Optional dense or block-sparse visibility mask.

    Returns:
        Output-projected tokens with the same shape and dtype as ``x``.
    """
    query_rope_freqs, key_rope_freqs = self._slice_rope_freqs(
        rope_freqs, kv_cache, x.shape[-2]
    )
    if self.attention_type is AttentionType.SELF_ATTENTION:
        query = self._update_kv_and_compute_query(x, kv_cache, query_rope_freqs)
    else:
        query = self._compute_query(x, query_rope_freqs)
        self._validate_cache(kv_cache, x)

    # ``cached_k/v`` expose only the valid prefix while a rolling cache fills,
    # and the complete fixed-size buffer after it reaches steady state.
    key = kv_cache.cached_k()
    if (
        self.attention_config.rope_config is not None
        and self.attention_config.rope_config.scope is RoPEScope.AFTER_KV_CACHE
        and key_rope_freqs is not None
    ):
        # The shared RoPE kernel is in-place; keep cache storage unrotated so
        # rolling positions can be applied again on the next attention call.
        key = self._apply_rope(key.to(x.dtype, copy=True), key_rope_freqs)
    value = kv_cache.cached_v()
    if self.optimized_impl_config.quantization.quantized_sdpa:
        query = query.to(torch.float8_e4m3fn)
        key = key.to(torch.float8_e4m3fn)
        value = value.to(torch.float8_e4m3fn)
    output = self._attention(
        query,
        key,
        value,
        attn_mask=attn_mask,
        output_dtype=x.dtype,
    )
    sequence_length = x.shape[-2]
    output = output.reshape(-1, sequence_length, self.attention_config.inner_dim)
    output = self._project_output(output)
    return output.reshape(
        x.shape[:-2] + (sequence_length, self.attention_config.query_dim)
    )

def __post_init__(self) -> None:
    """Reject invalid values before they reach a compiled kernel."""
    for name, value in (
        ("block_size", self.block_size),
        ("mask_block_m", self.mask_block_m),
        ("mask_block_n", self.mask_block_n),
        ("block_m", self.block_m),
        ("block_n", self.block_n),
    ):
        if value is not None and (value < 16 or value & (value - 1)):
            raise ValueError(
                f"FlexAttention {name} must be a power of two >= 16; got {value}"
            )
    if self.compile_dynamic is not None and not isinstance(
        self.compile_dynamic, bool
    ):
        raise TypeError(
            f"compile_dynamic must be bool or None; got {self.compile_dynamic!r}"
        )
    if self.num_warps is not None and self.num_warps not in (1, 2, 4, 8):
        raise ValueError(
            f"num_warps must be one of 1, 2, 4, or 8; got {self.num_warps}"
        )
    if self.num_stages is not None and self.num_stages <= 0:
        raise ValueError(f"num_stages must be positive; got {self.num_stages}")
    for name, value in (
        ("prescale_qk", self.prescale_qk),
        ("use_tma", self.use_tma),
    ):
        if value is not None and not isinstance(value, bool):
            raise TypeError(f"{name} must be bool or None; got {value!r}")
    if self.backend is not None and self.backend not in ("TRITON", "FLASH"):
        raise ValueError(
            f"backend must be TRITON, FLASH, or None; got {self.backend!r}"
        )
    if not isinstance(self.rows_guaranteed_safe, bool):
        raise TypeError(
            f"rows_guaranteed_safe must be bool; got {self.rows_guaranteed_safe!r}"
        )