From 1e1b17facc9b405d70c41018e3f15ca38ea8102e Mon Sep 17 00:00:00 2001 From: titaiwang Date: Fri, 19 Jun 2026 23:44:53 +0000 Subject: [PATCH 1/4] feat: bias-aware external-KV static-cache attention (slice A, #366) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire mobius's maskless external-KV static-cache graph together with its float-bias builder so float-bias decoders (sliding window, Gemma4 block overlay) can use the opset-24 external KV cache (TensorScatter + nonpad_kv_seqlen) on the standard ONNX Attention op with is_causal=0, instead of giving up the decode fast-path. Part of #349. Gated behind MOBIUS_STATIC_CACHE_BIAS (default off): no shipped model's emission changes unless the flag is set AND the model declares a bias need (sliding_window). ORT-testable on CPU MEA today (no genai, no Flash, no onnxruntime#28958). Deliverables: - (A) components/_attention.py: static branch derives (mask_arg, is_causal) from attn_mask presence — bias present => is_causal=0 (STRICTLY paired, guarding the double-causal bug); maskless default => is_causal=1. nonpad_kv_seqlen stays as Attention input #6 in both modes. - (B) components/_common.py: new create_static_cache_attention_bias building a (B,1,S_q,max_seq) additive bias (causal + sliding window + Gemma4 block overlay + padding) keyed on absolute query positions write_indices+arange, dense KV slot ids, and KV validity slot < nonpad_kv_seqlen. Exported. - (C) models/base.py: TextModel threads the bias in the static branch when the flag is on, the model has a sliding window, and the cache is a StaticCacheState; behavior is exactly the maskless path otherwise. - (D) _flags.py: MOBIUS_STATIC_CACHE_BIAS -> flags.static_cache_bias. - (E) tests/static_cache_bias_parity_test.py: ORT CPU MEA parity vs an independent numpy dense reference (prefill, decode, sliding-window sweep, nonpad padding clamp) plus graph-wiring tests for (C)+(D). 9 tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang --- src/mobius/_flags.py | 28 ++ src/mobius/components/__init__.py | 2 + src/mobius/components/_attention.py | 51 ++- src/mobius/components/_common.py | 132 ++++++ src/mobius/models/base.py | 54 ++- tests/static_cache_bias_parity_test.py | 569 +++++++++++++++++++++++++ 6 files changed, 815 insertions(+), 21 deletions(-) create mode 100644 tests/static_cache_bias_parity_test.py diff --git a/src/mobius/_flags.py b/src/mobius/_flags.py index d81077fc..6090cb2e 100644 --- a/src/mobius/_flags.py +++ b/src/mobius/_flags.py @@ -90,6 +90,13 @@ class _Flags: - Use native ``MatMulNBits bits=2`` for Tencent SEQ Q1_0 (smaller, semantically faithful, but ~20x slower on CPU EP pending an MLAS fast path). + * - ``static_cache_bias`` + - ``MOBIUS_STATIC_CACHE_BIAS`` + - ``False`` + - Emit a float additive attention bias (causal + sliding window + + block overlay + padding) on the external-KV static-cache + ``Attention`` path (``is_causal=0``) for float-bias decoders, + instead of the maskless ``is_causal=1`` default. """ suppress_dedup_warning: bool = dataclasses.field( @@ -150,6 +157,27 @@ class _Flags: opt in to the smaller native form once kernel performance lands. """ + static_cache_bias: bool = dataclasses.field( + default_factory=lambda: _env_bool("MOBIUS_STATIC_CACHE_BIAS", False) + ) + """Emit a float additive attention bias on the external-KV static-cache + ``Attention`` path instead of the maskless ``is_causal=1`` default. + + When ``True`` (and the model declares a bias need, e.g. a sliding window + or a block-overlay hook), :class:`~mobius.models.base.TextModel` builds a + ``(B, 1, S_q, max_seq)`` additive bias via + :func:`~mobius.components.create_static_cache_attention_bias` (causal + + sliding window + block overlay + padding, keyed on absolute query + positions with KV validity ``slot < nonpad_kv_seqlen``) and threads it + into the static-cache ``Attention`` op with ``is_causal=0``. This lets a + single standard-``Attention`` graph carry an arbitrary additive bias that + ``com.microsoft.GroupQueryAttention`` cannot express, while still using the + opset-24 external KV cache (``TensorScatter`` + ``nonpad_kv_seqlen``). + + Default ``False``: the maskless ``is_causal=1`` static-cache emission is + unchanged, so no shipped model's graph changes unless this flag is set. + """ + # Global singleton — import and use this directly. flags = _Flags() diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index 87f0fc56..630e447c 100644 --- a/src/mobius/components/__init__.py +++ b/src/mobius/components/__init__.py @@ -74,6 +74,7 @@ "create_decoder_layer", "create_padding_mask", "create_sliding_window_mask", + "create_static_cache_attention_bias", "get_activation", "initialize_rope", "make_quantized_linear_factory", @@ -118,6 +119,7 @@ create_attention_bias, create_padding_mask, create_sliding_window_mask, + create_static_cache_attention_bias, ) from mobius.components._conv import ( BatchNorm2d, diff --git a/src/mobius/components/_attention.py b/src/mobius/components/_attention.py index 2ff62d77..dce9db1e 100644 --- a/src/mobius/components/_attention.py +++ b/src/mobius/components/_attention.py @@ -101,7 +101,12 @@ def _apply_attention( Static cache mode (``static_cache is not None``): Scatters new key/value into the static cache via TensorScatter, then attends over the full cache using ``nonpad_kv_seqlen``. - Also uses ``is_causal=1``. + Uses ``is_causal=1`` when ``attn_mask`` is ``None`` (maskless + default), or ``is_causal=0`` when ``attn_mask`` is a float additive + bias (the bias then carries the full causal + sliding + block + + padding mask). The incoming ``is_causal`` argument is ignored in + this mode; causality is derived from ``attn_mask`` presence so the + two can never disagree. Returns ``(attn_output, updated_key_cache, updated_value_cache)``. Args: @@ -149,28 +154,36 @@ def _apply_attention( axis=1, ) # [B, max_seq, kv_hidden] - # Attend over the full cache. We pass None for attn_mask and use - # is_causal=1 instead — the Attention op handles causal + padding - # masking internally via is_causal + nonpad_kv_seqlen. Using - # create_attention_bias() here would produce incorrect causality - # during prefill because it cannot represent the relationship - # between query positions and the full cache length. + # External-cache masking. Two modes, selected by whether the caller + # supplied a float additive bias: # - # NOTE: The ONNX Attention spec supports attn_mask alongside - # nonpad_kv_seqlen for custom masking (e.g., user-defined masks - # beyond causal + padding). Currently we rely on is_causal=1 + - # nonpad_kv_seqlen for standard LLM causal + padding masking. - # TODO(titaiwang): Support user-provided attn_mask in external - # cache mode for advanced use cases (e.g., prefix masking, - # document boundaries in batched inference). - # TODO(titaiwang): Support sliding window (circular cache mode) - # with static cache for long-context models that use local - # attention windows. + # * attn_mask is None (default, maskless): pass None and use + # is_causal=1 — the Attention op derives causal + padding masking + # internally from is_causal + nonpad_kv_seqlen. This is the + # Flash-eligible form (onnx#8068 / onnxruntime#28958). + # * attn_mask is not None (float-bias decoders): pass the bias and + # STRICTLY pair it with is_causal=0. The bias already bakes in the + # FULL mask (causal + sliding + Gemma4 block overlay + padding), so + # leaving is_causal=1 would double-apply causality and cancel any + # bidirectional unmasking encoded in the bias. This routes ORT to + # the MEA external-cache path (Flash is precluded by any bias). + # + # nonpad_kv_seqlen stays as input #6 in BOTH modes: it bounds the valid + # KV prefix and, on the Flash path, drives the fully-masked-row zero + # guard. In bias mode the additive bias already encodes the same + # ``slot < nonpad`` validity, and with the contract-consistent feed + # (nonpad == write_indices + S_q) every query row keeps its own diagonal + # slot valid — so no row degenerates to an all-``dtype.min`` (NaN) + # softmax in normal operation. + if attn_mask is not None: + mask_arg, causal = attn_mask, 0 + else: + mask_arg, causal = None, 1 attn_output, _, _ = op.Attention( query, updated_k, updated_v, - None, # no attn_mask — is_causal handles masking + mask_arg, None, # no past_key (full cache is already provided) None, # no past_value static_cache.nonpad_kv_seqlen, @@ -178,7 +191,7 @@ def _apply_attention( kv_num_heads=num_key_value_heads, scale=scale, softcap=softcap, - is_causal=is_causal, + is_causal=causal, _outputs=3, ) return attn_output, updated_k, updated_v diff --git a/src/mobius/components/_common.py b/src/mobius/components/_common.py index f89cb189..d7c02d72 100644 --- a/src/mobius/components/_common.py +++ b/src/mobius/components/_common.py @@ -264,6 +264,138 @@ def create_attention_bias( return op.Unsqueeze(attention_bias, [1]) +def create_static_cache_attention_bias( + op: OpBuilder, + *, + write_indices: ir.Value, + seq_len: ir.Value, + nonpad_kv_seqlen: ir.Value, + max_seq_len: int, + sliding_window: int | None = None, + block_sequence_ids: ir.Value | None = None, + dtype: ir.DataType = ir.DataType.FLOAT, +) -> ir.Value: + """Build a float additive attention bias for the external-KV static cache. + + Sibling to :func:`create_attention_bias`. That function masks against a + *dense* ``[batch, total_length]`` ``attention_mask`` whose width grows with + the sequence; this one masks against the *pre-allocated* static cache whose + KV axis has a fixed width ``max_seq_len`` and whose valid prefix length is + given per batch by ``nonpad_kv_seqlen``. It reuses the same masking rules + (causal, sliding window, Gemma4 bidirectional block overlay, padding), only + re-keyed onto the static-cache geometry: + + * **KV axis width** is ``max_seq_len`` (every cache slot), not the running + ``total_length``. + * **KV positions** are dense slot ids ``0 .. max_seq_len - 1`` (the absolute + cache position each slot holds), so the causal/sliding comparisons use + absolute positions directly instead of cumsum indices. + * **Q positions** are absolute: ``write_indices[b] + arange(S_q)`` — the + cache slots the current chunk is being scattered into. + * **KV validity** is ``slot < nonpad_kv_seqlen[b]`` (replaces the cumsum + padding term): a slot is real iff it lies inside the valid prefix. + + The result is meant to be passed as the ``attn_mask`` of the static-cache + ``Attention`` op with ``is_causal=0`` (see + :func:`~mobius.components._attention._apply_attention`). + + Args: + op: The OpBuilder. + write_indices: ``[B]`` int64 start slot per batch (where the current + chunk is scattered — the same tensor fed to ``TensorScatter``). + seq_len: Scalar int64 query-chunk length ``S_q``. + nonpad_kv_seqlen: ``[B]`` int64 count of valid cache slots *after* the + current chunk is scattered (``write_indices + S_q`` under the + bottom-right contract). + max_seq_len: Static width of the pre-allocated cache KV axis. + sliding_window: Optional local-attention window; when set, a query at + absolute position ``q`` attends slot ``k`` only if ``q - k < w``. + block_sequence_ids: Optional ``[B, S_q]`` int64 block id per *current* + query position (``>= 0`` for vision tokens sharing a block, ``-1`` + for text). When set, a bidirectional overlay is OR-ed onto the + causal/sliding mask: two positions in the same block may attend to + each other regardless of causal order — mirroring HuggingFace + ``blockwise_overlay`` for Gemma4. The per-slot KV block ids are + built by scattering ``block_sequence_ids`` into a ``-1``-filled + ``[B, max_seq_len]`` buffer at ``write_indices`` (the same scatter + geometry as K/V), so within a single forward only the current + chunk's block ids participate; cross-step persistence of past block + ids is a caller concern (out of scope for this primitive). + dtype: Output dtype; masked entries use ``dtype.min``. + + Returns: + Additive bias of shape ``(B, 1, S_q, max_seq_len)``: ``0`` where the + query may attend the slot, ``dtype.min`` where masked. + """ + zero_scalar = op.Constant(value_int=0) + one_scalar = op.Constant(value_int=1) + + # Q absolute positions: write_indices[b] + arange(S_q) -> (B, S_q, 1). + seq_scalar = op.Squeeze(seq_len, op.Constant(value_ints=[0])) + q_offsets = op.Range(zero_scalar, seq_scalar, one_scalar) # (S_q,) + q_abs_2d = op.Add( + op.Unsqueeze(write_indices, [1]), # (B, 1) + op.Unsqueeze(q_offsets, [0]), # (1, S_q) + ) # (B, S_q) + q_abs = op.Unsqueeze(q_abs_2d, [2]) # (B, S_q, 1) + + # KV positions: dense cache slot ids 0 .. max_seq_len - 1 -> (max_seq_len,). + kv_slots = op.Range( + zero_scalar, op.Constant(value_int=max_seq_len), one_scalar + ) # (max_seq_len,) + + # Causal: a query at absolute position q may attend slot k iff q >= k. + # Broadcasting (B, S_q, 1) >= (max_seq_len,) -> (B, S_q, max_seq_len). + full_mask = op.GreaterOrEqual(q_abs, kv_slots) + + if sliding_window is not None: + # Local window: keep slots within `sliding_window` of the query. + dist = op.Sub(q_abs, kv_slots) # (B, S_q, max_seq_len) + within_window = op.Less(dist, op.Constant(value_int=sliding_window)) + full_mask = op.And(full_mask, within_window) + + if block_sequence_ids is not None: + # Bidirectional vision-block overlay (OR-ed BEFORE the padding AND, + # matching create_attention_bias / HF blockwise_overlay ordering). + # + # q_group: block id per current query position -> (B, S_q, 1). + q_group = op.Unsqueeze(block_sequence_ids, [2]) # (B, S_q, 1) + # kv_group: block id per cache slot -> (B, 1, max_seq_len). Built by + # scattering the current chunk's block ids into a -1 buffer at + # write_indices, exactly as K/V are scattered into the cache. + batch_dim = op.Shape(write_indices, start=0, end=1) # (1,) == [B] + cache_shape = op.Concat( + batch_dim, op.Constant(value_ints=[max_seq_len]), axis=0 + ) # (2,) == [B, max_seq_len] + neg1_buffer = op.ConstantOfShape( + cache_shape, value=ir.tensor(np.array([-1], dtype=np.int64)) + ) # (B, max_seq_len) of -1 + kv_group_2d = op.TensorScatter( + neg1_buffer, block_sequence_ids, write_indices, axis=1 + ) # (B, max_seq_len) + kv_group = op.Unsqueeze(kv_group_2d, [1]) # (B, 1, max_seq_len) + same_block = op.And( + op.Equal(q_group, kv_group), + op.GreaterOrEqual(q_group, op.Constant(value_int=0)), + ) + full_mask = op.Or(full_mask, same_block) + + # KV validity (padding): a slot is real iff it lies inside the valid + # prefix, slot < nonpad_kv_seqlen[b]. Broadcasting (max_seq_len,) < + # (B, 1) -> (B, max_seq_len) -> (B, 1, max_seq_len). + valid_kv = op.Less(kv_slots, op.Unsqueeze(nonpad_kv_seqlen, [1])) + valid_kv = op.Unsqueeze(valid_kv, [1]) # (B, 1, max_seq_len) + full_mask = op.And(full_mask, valid_kv) + + # Convert to float bias: 0 where attended, dtype.min where masked. + mask_value = float(dtype.min) + attention_bias = op.Where(full_mask, 0.0, mask_value) # (B, S_q, max_seq_len) + attention_bias = op.Cast(attention_bias, to=dtype) + + # Unsqueeze to (B, 1, S_q, max_seq_len). + return op.Unsqueeze(attention_bias, [1]) + + def build_packed_token_offset(op: OpBuilder, cu_seqlens) -> ir.Value: """Build ``token_offset`` for ``com.microsoft::PackedMultiHeadAttention``. diff --git a/src/mobius/models/base.py b/src/mobius/models/base.py index a32701f0..8dfe9509 100644 --- a/src/mobius/models/base.py +++ b/src/mobius/models/base.py @@ -19,6 +19,7 @@ from mobius._build_context import ep_capabilities, get_build_dtype from mobius._configs import ArchitectureConfig, CausalLMConfig +from mobius._flags import flags from mobius._weight_utils import ( preprocess_awq_weights, preprocess_gptq_weights, @@ -32,10 +33,11 @@ Linear, RMSNorm, create_padding_mask, + create_static_cache_attention_bias, initialize_rope, make_quantized_linear_factory, ) -from mobius.components._attention import GQAContext +from mobius.components._attention import GQAContext, StaticCacheState from mobius.components._rotary_embedding import BaseRope, _MRopeBase @@ -73,6 +75,54 @@ def __init__(self, config: ArchitectureConfig, mlp_class: type | None = None): self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = initialize_rope(config) + # Sliding-window models declare a local-attention span; it drives the + # optional static-cache float bias (flags.static_cache_bias). Standard + # full-attention models leave this None, so the bias path is a no-op + # for them even when the flag is set. + self._sliding_window: int | None = getattr(config, "sliding_window", None) + + def _maybe_static_cache_bias( + self, + op: OpBuilder, + input_ids: ir.Value | None, + past_key_values: list | None, + ) -> ir.Value | None: + """Optionally build the static-cache float additive attention bias. + + Returns ``None`` (maskless ``is_causal=1`` default) unless ALL hold: + * ``flags.static_cache_bias`` is set, AND + * the model declares a bias need (``self._sliding_window`` is set), AND + * the cache is the opset-24 external cache (``StaticCacheState``). + + When emitted, the bias is a ``(B, 1, S_q, max_seq)`` additive mask keyed + on absolute query positions with KV validity ``slot < nonpad_kv_seqlen``; + ``_apply_attention`` then pairs it with ``is_causal=0``. The + ``write_indices`` / ``nonpad_kv_seqlen`` graph inputs are shared across + all layers, so the first layer's cache state carries them. + """ + if not flags.static_cache_bias or self._sliding_window is None: + return None + if not past_key_values: + return None + first = past_key_values[0] + if not isinstance(first, StaticCacheState): + return None + if input_ids is None: + return None + + # Static cache KV axis width is a concrete int: [B, max_seq, kv_hidden]. + max_seq_len = int(first.key_cache.shape[1]) + seq_len = op.Shape(input_ids, start=1, end=2) # (1,) int64 == [S_q] + return create_static_cache_attention_bias( + op, + write_indices=first.write_indices, + seq_len=seq_len, + nonpad_kv_seqlen=first.nonpad_kv_seqlen, + max_seq_len=max_seq_len, + sliding_window=self._sliding_window, + dtype=self._dtype, + ) + def forward( self, op: OpBuilder, @@ -166,7 +216,7 @@ def forward( attention_mask=attention_mask, ) else: - attention_bias = None + attention_bias = self._maybe_static_cache_bias(op, input_ids, past_key_values) present_key_values = [] past_kvs = past_key_values or [None] * len(self.layers) diff --git a/tests/static_cache_bias_parity_test.py b/tests/static_cache_bias_parity_test.py new file mode 100644 index 00000000..92f119e5 --- /dev/null +++ b/tests/static_cache_bias_parity_test.py @@ -0,0 +1,569 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Static-cache *bias* numerical-parity tests (mobius #366, Part of #349). + +These tests guard the bias-aware external-KV static-cache attention combination +that ``flags.static_cache_bias`` enables: ONNX ``Attention`` with ``is_causal=0`` ++ a float additive ``attn_mask`` (Attention input #3) + ``nonpad_kv_seqlen`` +(input #6), fed by two ``TensorScatter`` writes into a pre-allocated KV cache. +The additive bias is built in-graph by +:func:`mobius.components.create_static_cache_attention_bias` and carries the +FULL mask geometry (causal + sliding window + Gemma4 block overlay + padding), +keyed on absolute query positions with KV validity ``slot < nonpad_kv_seqlen``. + +Unlike the maskless ``is_causal=1`` static-cache path (``static_cache_parity_test``) +— which is Flash-eligible but rejected by CPU/pre-#28958 kernels for +``S_q != total_kv`` with no ``past_key`` — the **bias** path routes ORT to the +MEA external-cache combination, which **runs on the CPU EP today** (no genai, no +Flash, no onnxruntime#28958). So these tests run unconditionally on CPU. + +The reference is an independent NumPy *dense* attention that gathers the valid +cache slots and applies the SAME causal + sliding + block-overlay + padding mask +the in-graph bias encodes. Because the bias geometry (absolute query positions +vs. dense cache slot ids, GQA head sharing, the bottom-right contract) is the +only thing under test, an independent re-derivation of the mask in NumPy is the +authoritative guard (mirrors ``create_attention_bias``'s parity strategy). + +Two cases: + +* prefill chunk — ``write_indices=0``, ``S_q=N``, ``nonpad=N`` (the whole chunk + is the valid region). +* decode step — ``write_indices=N``, ``S_q=1``, ``nonpad=N+1`` against a + pre-populated cache (the single query attends the ``N`` past slots + itself). + +Plus a padding-clamp guard (risk ii): an under-clamped ``nonpad < S_q`` feed +masks the cache tail; because this bias keys query positions on +``write_indices + arange(S_q)`` (not the bottom-right contract), every row still +attends its diagonal slot, so the output stays finite (no all-``dtype.min`` NaN) +and matches the dense reference. + +Run:: + + pytest tests/static_cache_bias_parity_test.py -v +""" + +from __future__ import annotations + +import numpy as np +import onnx_ir as ir +import pytest +from onnxscript import GraphBuilder + +from mobius._constants import OPSET_VERSION +from mobius._flags import override_flags +from mobius._testing.ort_inference import OnnxModelSession +from mobius.components import create_static_cache_attention_bias + +# --------------------------------------------------------------------------- +# Standalone bias-attention graph builder +# --------------------------------------------------------------------------- + + +def _build_static_cache_bias_graph( + *, + batch: int, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + max_seq_len: int, + query_len: int, + sliding_window: int | None, + use_block_overlay: bool, + dtype: ir.DataType = ir.DataType.FLOAT, +) -> ir.Model: + """Build a standalone graph mirroring the bias static-cache attention branch. + + Replicates ``_apply_attention``'s static *bias* path: scatter ``key`` / + ``value`` into the pre-allocated caches via ``TensorScatter``, build the + ``(B, 1, S_q, max_seq)`` additive bias via + :func:`create_static_cache_attention_bias`, then run ``Attention`` with that + bias as input #3, ``is_causal=0``, and ``nonpad_kv_seqlen`` as input #6. + + Inputs: ``query`` ``[B, S_q, num_q_heads*head_dim]``; ``key`` / ``value`` + ``[B, S_q, num_kv_heads*head_dim]`` (GQA); ``key_cache`` / + ``value_cache`` ``[B, max_seq_len, num_kv_heads*head_dim]``; + ``write_indices`` and ``nonpad_kv_seqlen`` ``[B]`` int64; optionally + ``block_sequence_ids`` ``[B, S_q]`` int64. + Outputs: ``attn_output`` ``[B, S_q, num_q_heads*head_dim]``, the updated + caches, and ``bias`` (for debugging / reference cross-check). + """ + q_hidden = num_q_heads * head_dim + kv_hidden = num_kv_heads * head_dim + + def _value(name: str, dims: list[int], dt: ir.DataType) -> ir.Value: + return ir.Value(name=name, shape=ir.Shape(dims), type=ir.TensorType(dt)) + + query = _value("query", [batch, query_len, q_hidden], dtype) + key = _value("key", [batch, query_len, kv_hidden], dtype) + value = _value("value", [batch, query_len, kv_hidden], dtype) + key_cache = _value("key_cache", [batch, max_seq_len, kv_hidden], dtype) + value_cache = _value("value_cache", [batch, max_seq_len, kv_hidden], dtype) + write_indices = _value("write_indices", [batch], ir.DataType.INT64) + nonpad_kv_seqlen = _value("nonpad_kv_seqlen", [batch], ir.DataType.INT64) + + inputs = [query, key, value, key_cache, value_cache, write_indices, nonpad_kv_seqlen] + block_sequence_ids = None + if use_block_overlay: + block_sequence_ids = _value( + "block_sequence_ids", [batch, query_len], ir.DataType.INT64 + ) + inputs.append(block_sequence_ids) + + graph = ir.Graph( + inputs=inputs, + outputs=[], + nodes=[], + name="static_cache_bias_probe", + opset_imports={"": OPSET_VERSION}, + ) + op = GraphBuilder(graph).op + + # Scatter new K/V into the pre-allocated cache: cache[b, write[b] + t] = upd[b, t]. + updated_k = op.TensorScatter(key_cache, key, write_indices, axis=1) + updated_v = op.TensorScatter(value_cache, value, write_indices, axis=1) + + # Build the additive bias in-graph (the unit under test). + seq_len = op.Constant(value_ints=[query_len]) # (1,) int64 == [S_q] + bias = create_static_cache_attention_bias( + op, + write_indices=write_indices, + seq_len=seq_len, + nonpad_kv_seqlen=nonpad_kv_seqlen, + max_seq_len=max_seq_len, + sliding_window=sliding_window, + block_sequence_ids=block_sequence_ids, + dtype=dtype, + ) + + scale = 1.0 / np.sqrt(head_dim) + # is_causal=0 STRICTLY paired with the bias (the bias already encodes + # causality; is_causal=1 would double-apply it). + attn_output, _, _ = op.Attention( + query, + updated_k, + updated_v, + bias, # attn_mask input #3 — the additive float bias + None, # no past_key (full cache is provided) + None, # no past_value + nonpad_kv_seqlen, # input #6 — drives the fully-masked-row zero guard + q_num_heads=num_q_heads, + kv_num_heads=num_kv_heads, + scale=float(scale), + is_causal=0, + _outputs=3, + ) + + attn_output.name = "attn_output" + updated_k.name = "updated_key_cache" + updated_v.name = "updated_value_cache" + bias.name = "bias" + graph.outputs.extend([attn_output, updated_k, updated_v, bias]) + + return ir.Model(graph, ir_version=10) + + +# --------------------------------------------------------------------------- +# Independent NumPy dense-attention reference +# --------------------------------------------------------------------------- + + +def _dense_reference( + query: np.ndarray, + full_key_cache: np.ndarray, + full_value_cache: np.ndarray, + *, + write_indices: np.ndarray, + nonpad_kv_seqlen: np.ndarray, + max_seq_len: int, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + sliding_window: int | None, + block_sequence_ids: np.ndarray | None, +) -> np.ndarray: + """Dense reference applying the SAME mask the in-graph bias encodes. + + ``full_*_cache`` are the *post-scatter* caches (current chunk already + written). For each (batch, query, head) we build the boolean mask in the + exact rule order ``create_static_cache_attention_bias`` uses — causal, + AND sliding window, OR block overlay, AND padding validity — gather the + valid cache slots (with GQA head sharing) and softmax over them. A query + row with no valid slot yields exactly ``0`` (the kernel's zero guard). + """ + batch, query_len, q_hidden = query.shape + scale = 1.0 / np.sqrt(head_dim) + group = num_q_heads // num_kv_heads + out = np.zeros((batch, query_len, q_hidden), dtype=np.float64) + + kv_slots = np.arange(max_seq_len) + for b in range(batch): + wi = int(write_indices[b]) + npad = int(nonpad_kv_seqlen[b]) + q_abs = wi + np.arange(query_len) # absolute query positions + + # Per-slot KV block ids: -1 buffer with the current chunk scattered in + # (mirrors the in-graph TensorScatter into a -1 buffer). + kv_group = np.full((max_seq_len,), -1, dtype=np.int64) + if block_sequence_ids is not None: + kv_group[wi : wi + query_len] = block_sequence_ids[b] + + # Cache reshaped to (max_seq, num_kv_heads, head_dim) for GQA gather. + k_heads = full_key_cache[b].reshape(max_seq_len, num_kv_heads, head_dim) + v_heads = full_value_cache[b].reshape(max_seq_len, num_kv_heads, head_dim) + q_heads = query[b].reshape(query_len, num_q_heads, head_dim) + + for t in range(query_len): + mask = q_abs[t] >= kv_slots # causal + if sliding_window is not None: + mask &= (q_abs[t] - kv_slots) < sliding_window # local window + if block_sequence_ids is not None and block_sequence_ids[b, t] >= 0: + mask |= kv_group == block_sequence_ids[b, t] # bidirectional block + mask &= kv_slots < npad # padding validity + idx = np.nonzero(mask)[0] + if idx.size == 0: + continue # structurally empty -> exactly 0 (zero guard) + for h in range(num_q_heads): + kv_h = h // group # GQA: query head -> shared kv head + kk = k_heads[idx, kv_h] # (n_valid, head_dim) + vv = v_heads[idx, kv_h] + scores = (q_heads[t, h] @ kk.T) * scale + scores = scores - scores.max() + weights = np.exp(scores) + weights = weights / weights.sum() + out[b, t, h * head_dim : (h + 1) * head_dim] = weights @ vv + return out + + +# --------------------------------------------------------------------------- +# Shared tiny bias-decoder geometry (hidden=64, GQA kv