Emit bias-aware external-KV (TensorScatter + nonpad_kv_seqlen) attention for float-bias decoders

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()

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,

src/mobius/components/_attention.py

@@ -62,8 +62,8 @@ class StaticCacheState(NamedTuple):
     ``nonpad_kv_seqlen`` to indicate valid token counts.
 
     Fields:
-        key_cache: Pre-allocated key cache [B, max_seq, kv_hidden] 3D.
-        value_cache: Pre-allocated value cache [B, max_seq, kv_hidden] 3D.
+        key_cache: Pre-allocated key cache [B, max_seq_len, kv_hidden] 3D.
+        value_cache: Pre-allocated value cache [B, max_seq_len, kv_hidden] 3D.
         write_indices: Position to write new tokens [B] int64.
         nonpad_kv_seqlen: Valid KV length per batch entry [B] int64.
     """
@@ -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:
@@ -114,10 +119,12 @@ def _apply_attention(
             unmasking encoded in the bias.
 
     Note:
-        Both paths default to ``is_causal=1`` on the Attention op, which
-        enables built-in causal masking. This means ``attn_mask`` should
-        encode only padding information (as a bool mask), not causality,
-        unless ``is_causal=0`` is passed explicitly.
+        This applies to the DYNAMIC cache path only. There, the Attention op
+        defaults to ``is_causal=1`` for built-in causal masking, so
+        ``attn_mask`` should encode only padding information (as a bool mask),
+        not causality, unless ``is_causal=0`` is passed explicitly. In STATIC
+        cache mode the incoming ``is_causal`` argument is ignored — causality
+        is derived from ``attn_mask`` presence (see above).
 
     Note:
         ``nonpad_kv_seqlen`` (input #6) is only valid in static cache mode
@@ -149,36 +156,53 @@ 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 CUDA Flash path, drives the fully-masked-row
+        # zero guard (LaunchZeroFullyMaskedRows).  In bias mode the additive
+        # bias already encodes the same ``slot < nonpad`` validity.  The
+        # cross-repo invariant is ``nonpad == write_indices + valid_token_count``
+        # (the count of UNPADDED query tokens), which equals
+        # ``write_indices + S_q`` only when the chunk is unpadded — S_q is the
+        # PADDED chunk width.  When the chunk is unpadded, every query row keeps
+        # its own diagonal slot valid, so a fully-masked (all-``dtype.min``) row
+        # never arises.  With intra-prompt padding plus a sliding window,
+        # however, a pad-token query row CAN fall outside every valid slot and
+        # become fully masked.  In that case the CPU MEA path this bias mode uses
+        # does NOT apply the Flash zero-guard: it returns a finite mean-of-V row
+        # (not NaN, not exactly 0).  This finite-row behavior was empirically
+        # verified on ORT 1.27 CPU MEA; it is an observed ORT-version behavior,
+        # not a permanent op-spec invariant — see test_fully_masked_row_stays_finite.
+        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,
             q_num_heads=num_attention_heads,
             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

src/mobius/components/_common.py

@@ -264,6 +264,146 @@ 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: 1-D ``[1]`` int64 tensor holding the query-chunk length ``S_q``
+            (e.g. ``op.Shape(input_ids, start=1, end=2)``); squeezed to a scalar
+            internally.
+        nonpad_kv_seqlen: ``[B]`` int64 count of valid cache slots *after* the
+            current chunk is scattered. The cross-repo invariant is
+            ``nonpad_kv_seqlen == write_indices + valid_token_count``, where
+            ``valid_token_count`` is the number of *unpadded* query tokens in the
+            chunk. This equals ``write_indices + S_q`` only when the chunk is
+            **unpadded** (``S_q`` is the padded chunk width). With intra-prompt
+            padding plus a sliding window, pad-token query rows can fall outside
+            every valid slot and become fully masked — see the fully-masked-row
+            behavior note in ``_apply_attention``.
+        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``.