Address titaiwangms review: conditional nonpad invariant + inputs_embeds bias

PR #367 review fixes (no bias-math behavior change):

MAJOR (doc accuracy + test): the cross-repo invariant is
nonpad == write_indices + valid_token_count (UNPADDED query-token count),
which equals write + S_q only when the chunk is unpadded (S_q is the
PADDED chunk width). With intra-prompt padding + a sliding window,
pad-token query rows can become fully masked.
- create_static_cache_attention_bias docstring + the _apply_attention
  comment now state the CONDITIONAL invariant.
- The 'CPU MEA returns a finite mean-of-V row, not NaN' claim is anchored
  to the empirically verified ORT 1.27 behavior (S_q=8, nonpad=3,
  window=4: rows 6-7 fully masked, stayed finite) rather than asserted as
  a permanent op-spec invariant.
- Add test_fully_masked_row_stays_finite, which proves (via an
  independent numpy mask) that rows 6-7 are fully masked and asserts the
  CPU MEA output stays finite — the prior parity tests assert idx.size>0
  and never reach this boundary.

MINOR: _maybe_static_cache_bias no longer returns None when input_ids is
None. It derives S_q from an always-present tensor (hidden_states), so
inputs_embeds-driven forwards still receive the bias.

NITS:
- StaticCacheState docstring uses max_seq_len (was max_seq); the is_causal
  Note in _apply_attention is scoped to the DYNAMIC path (static mode
  ignores the incoming is_causal).
- Guard the max_seq_len cast in _maybe_static_cache_bias: a symbolic
  key_cache KV dim now raises a descriptive TypeError naming the symbolic
  dim instead of an opaque int() failure.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: titaiwang <titaiwang@microsoft.com>

Author: titaiwangms

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.
     """
@@ -119,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
@@ -171,14 +173,19 @@ def _apply_attention(
         # 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, and with the
-        # contract-consistent feed (nonpad == write_indices + S_q) every query
-        # row keeps its own diagonal slot valid — so a fully-masked
-        # (all-``dtype.min``) row never arises in normal operation.  Even if one
-        # were forced (an out-of-contract feed), 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 and not exactly 0 (the zero guard is
-        # CUDA-Flash-specific).
+        # 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:

src/mobius/components/_common.py

@@ -307,8 +307,14 @@ def create_static_cache_attention_bias(
             (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 (``write_indices + S_q`` under the
-            bottom-right contract).
+            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``.

src/mobius/models/base.py

@@ -84,7 +84,7 @@ def __init__(self, config: ArchitectureConfig, mlp_class: type | None = None):
     def _maybe_static_cache_bias(
         self,
         op: OpBuilder,
-        input_ids: ir.Value | None,
+        seq_len_source: ir.Value,
         past_key_values: list | None,
     ) -> ir.Value | None:
         """Optionally build the static-cache float additive attention bias.
@@ -94,11 +94,18 @@ def _maybe_static_cache_bias(
           * 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.
+        When emitted, the bias is a ``(B, 1, S_q, max_seq_len)`` 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.
+
+        Args:
+            seq_len_source: An always-present ``[B, S_q, ...]`` tensor (e.g.
+                ``hidden_states``) whose dim 1 is the query length ``S_q``. Using
+                this instead of ``input_ids`` keeps the bias enabled for
+                ``inputs_embeds``-driven forwards (where ``input_ids`` is None).
         """
         if not flags.static_cache_bias or self._sliding_window is None:
             return None
@@ -107,12 +114,20 @@ def _maybe_static_cache_bias(
         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]
+        # Static cache KV axis width is a concrete int: [B, max_seq_len, kv_hidden].
+        # Guard against a symbolic dim, which would otherwise raise an opaque
+        # TypeError downstream. Static-cache always allocates a fixed width today.
+        max_seq_len = first.key_cache.shape[1]
+        if not isinstance(max_seq_len, int):
+            raise TypeError(
+                "static-cache bias requires a concrete key_cache KV dimension "
+                f"(axis 1), but got symbolic dim {max_seq_len!r}. The static "
+                "cache must be allocated with a fixed max_seq_len."
+            )
+        # S_q lives at dim 1 of both input_ids ([B, S_q]) and hidden_states
+        # ([B, S_q, hidden]), so the bias works for either forward entry point.
+        seq_len = op.Shape(seq_len_source, start=1, end=2)  # (1,) int64 == [S_q]
         return create_static_cache_attention_bias(
             op,
             write_indices=first.write_indices,
@@ -216,7 +231,9 @@ def forward(
                     attention_mask=attention_mask,
                 )
             else:
-                attention_bias = self._maybe_static_cache_bias(op, input_ids, past_key_values)
+                attention_bias = self._maybe_static_cache_bias(
+                    op, hidden_states, past_key_values
+                )
 
         present_key_values = []
         past_kvs = past_key_values or [None] * len(self.layers)

tests/static_cache_bias_parity_test.py

@@ -519,6 +519,115 @@ def test_nonpad_padding_clamp_matches_dense_reference():
     np.testing.assert_allclose(onnx_attn, ref, atol=1e-4, rtol=1e-4)
 
 
+def _fully_masked_query_rows(
+    *,
+    query_len: int,
+    write_idx: int,
+    nonpad: int,
+    sliding_window: int,
+    max_seq_len: int,
+) -> list[int]:
+    """Return the query-row indices that are fully masked (zero valid slots).
+
+    Mirrors ``_dense_reference``'s rule order (causal AND sliding window AND
+    padding validity, no block overlay) so a test can prove it genuinely
+    exercises the all-masked-row boundary instead of silently skipping it.
+    """
+    kv_slots = np.arange(max_seq_len)
+    empty_rows: list[int] = []
+    for t in range(query_len):
+        q_abs = write_idx + t
+        mask = q_abs >= kv_slots  # causal
+        mask &= (q_abs - kv_slots) < sliding_window  # local window
+        mask &= kv_slots < nonpad  # padding validity
+        if not mask.any():
+            empty_rows.append(t)
+    return empty_rows
+
+
+def test_fully_masked_row_stays_finite():
+    """Intra-prompt padding + sliding window can fully mask a pad-token row.
+
+    Counters the ``nonpad == write + S_q`` framing: ``S_q`` is the PADDED chunk
+    width, while ``nonpad == write + valid_token_count`` uses the UNPADDED count.
+    When the chunk carries trailing pad tokens (``nonpad < write + S_q``) AND a
+    sliding window is active, the high pad-token query rows can fall outside
+    every valid slot and become fully masked (all-``dtype.min`` bias row).
+
+    Unlike the other parity tests (which assert ``idx.size > 0`` via
+    ``_dense_reference`` and never reach this boundary), this test deliberately
+    drives the empty-row case and asserts the CPU MEA external-cache path keeps
+    the output FINITE (a finite mean-of-V row, not NaN, not exactly 0).
+
+    Scenario reproduces the reviewer's ORT 1.27 finding (``S_q=8, nonpad=3,
+    window=4`` → rows 6-7 have zero valid slots and stay finite). This is an
+    observed ORT-version behavior (verified on 1.27), not a permanent op-spec
+    guarantee — see the note in ``_apply_attention``.
+    """
+    query_len, write_idx, nonpad = 8, 0, 3
+
+    # Prove the config genuinely produces fully-masked rows (else the test would
+    # silently pass without exercising the boundary it claims to cover).
+    empty_rows = _fully_masked_query_rows(
+        query_len=query_len,
+        write_idx=write_idx,
+        nonpad=nonpad,
+        sliding_window=_SLIDING_WINDOW,
+        max_seq_len=_MAX_SEQ_LEN,
+    )
+    assert empty_rows == [6, 7], (
+        f"expected rows 6-7 fully masked for the reviewer's scenario, got {empty_rows}"
+    )
+
+    rng = np.random.default_rng(11)
+    q_hidden = _NUM_Q_HEADS * _HEAD_DIM
+    kv_hidden = _NUM_KV_HEADS * _HEAD_DIM
+    query = rng.standard_normal((_BATCH, query_len, q_hidden)).astype(np.float32)
+    key = rng.standard_normal((_BATCH, query_len, kv_hidden)).astype(np.float32)
+    value = rng.standard_normal((_BATCH, query_len, kv_hidden)).astype(np.float32)
+    key_cache = np.zeros((_BATCH, _MAX_SEQ_LEN, kv_hidden), dtype=np.float32)
+    value_cache = np.zeros((_BATCH, _MAX_SEQ_LEN, kv_hidden), dtype=np.float32)
+    write_indices = np.full((_BATCH,), write_idx, dtype=np.int64)
+    nonpad_kv_seqlen = np.full((_BATCH,), nonpad, dtype=np.int64)
+
+    model = _build_static_cache_bias_graph(
+        batch=_BATCH,
+        num_q_heads=_NUM_Q_HEADS,
+        num_kv_heads=_NUM_KV_HEADS,
+        head_dim=_HEAD_DIM,
+        max_seq_len=_MAX_SEQ_LEN,
+        query_len=query_len,
+        sliding_window=_SLIDING_WINDOW,
+        use_block_overlay=False,
+    )
+    session = OnnxModelSession(model)  # CPU EP (MEA external-cache path)
+    try:
+        out = session.run(
+            {
+                "query": query,
+                "key": key,
+                "value": value,
+                "key_cache": key_cache,
+                "value_cache": value_cache,
+                "write_indices": write_indices,
+                "nonpad_kv_seqlen": nonpad_kv_seqlen,
+            }
+        )
+    finally:
+        session.close()
+
+    attn_output = out["attn_output"]
+    # The core empirically-verified claim: even fully-masked rows stay finite
+    # (no NaN/Inf) on the CPU MEA path — the Flash zero-guard is not applied.
+    assert np.isfinite(attn_output).all(), (
+        "fully-masked query row produced NaN/Inf on the CPU MEA external-cache path"
+    )
+    for t in empty_rows:
+        assert np.isfinite(attn_output[:, t, :]).all(), (
+            f"fully-masked query row {t} is not finite"
+        )
+
+
 @pytest.mark.parametrize("sliding_window", [None, 2, 4])
 def test_prefill_sliding_window_variants(sliding_window):
     """Prefill parity across no-window and tight/loose sliding windows.