fp16/GQA export fixes + capture-safe static-cache attention (supersedes #328)

CHANGELOG.md

@@ -7,6 +7,66 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Static-cache Attention Causal-Mask Fix
+
+#### Fixed
+
+- Static-cache exports (`--static-cache`, the ONNX `Attention` + `TensorScatter`
+  in-place KV-cache path) now drive the `Attention` op with `is_causal=0` plus an
+  explicit offset-aware causal mask instead of `is_causal=1`. The opset-24
+  `Attention` CUDA kernel rejects `is_causal=1` together with `nonpad_kv_seqlen`
+  when the query length differs from the total KV length and there is no
+  `past_key` (the `causal_cross_no_past` guard in `attention.cc`). Because the
+  static cache is pre-allocated to `max_seq_len`, `S_q != total_kv` in **both**
+  prefill and decode, so that guard fired and ORT raised `NOT_IMPLEMENTED` at
+  session init. A new `create_static_cache_causal_mask` helper builds a mask that
+  keeps key slot `j` for a query at absolute position `write_indices[b] + t` iff
+  `j <= write_indices[b] + t`, serving prefill (triangular) and decode (prefix)
+  with a single rule and subsuming the padding bound. The formulation is
+  branchless (no `If` phase-split), so the exported graph stays compatible with
+  CUDA Graph capture (see the export-gotchas skill, "Graph-capture
+  compatibility"); decode and prefill both run on the Memory-Efficient Attention
+  kernel.
+
+---
+
+### fp16 GQA Export Fix
+
+#### Fixed
+
+- Native fp16 GroupQueryAttention exports (e.g. `microsoft/Phi-3.5-mini-instruct`
+  with `--dtype f16 --execution-provider cuda`) no longer emit fp32 packed-QKV /
+  transposed weights. Previously the fold passes (`FoldConcatInitializersPass`,
+  `FoldTransposedInitializerPass`) defaulted a folded initializer's dtype to
+  `FLOAT` when the source `Value`'s declared type had been dropped during fp16
+  casting, producing a model onnxruntime rejected at load with a
+  `MatMul` type-parameter error (`tensor(float16)` vs `tensor(float)`) on both
+  CPU and CUDA EPs. A new `mobius._passes._dtype_utils.initializer_dtype()`
+  helper now resolves the effective dtype from `const_value` when the type
+  annotation is missing, so fp16 GQA models load directly with no manual
+  post-cast.
+
+---
+
+### GQA Present KV-Cache Shape Fix
+
+#### Fixed
+
+- GroupQueryAttention exports now declare correct `present.{i}.key` /
+  `present.{i}.value` graph-output shapes and dtype. The GQA contrib op's shape
+  inference mis-derived the present KV `head_dim` (e.g. 32 instead of 96 on
+  `microsoft/Phi-3.5-mini-instruct`), so the present KV-cache outputs declared a
+  `head_dim` inconsistent with the (correct) `past_key_values` inputs. ORT logged
+  `Error merging shape info ... lenient merge` (64 warnings on Phi-3.5) and any
+  consumer that chains `present` → `past` and trusts declared shapes (e.g.
+  `onnxruntime-genai`) saw mismatched past-vs-present KV cache types. This is a
+  metadata / declared-shape correction only — runtime numerics are unchanged
+  (weights byte-identical, next-token parity 20/20). `_register_kv_cache_outputs`
+  now stamps the present KV outputs symmetric to the past inputs. Affects
+  GQA-fusion packed-QKV exports (Phi-3.5, Llama-3.2, Qwen2, Mistral, Phi-3-GQA).
+
+---
+
 ### WebGPU Shape Op Support
 
 #### Changed

src/mobius/_passes/_dtype_utils.py

@@ -0,0 +1,55 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Shared dtype helpers for graph passes that materialize new initializers.
+
+Passes such as :class:`~mobius._passes.FoldConcatInitializersPass` and
+:class:`~mobius._passes.FoldTransposedInitializerPass` pre-compute new
+initializers from existing ones. They must stamp the *correct* dtype on the
+result, otherwise an fp16 model can silently end up with fp32 weights that
+onnxruntime rejects at load time (a MatMul binding fp16 and fp32 to the same
+type parameter ``T``).
+"""
+
+from __future__ import annotations
+
+import logging
+
+import onnx_ir as ir
+
+logger = logging.getLogger(__name__)
+
+
+def initializer_dtype(value: ir.Value) -> ir.DataType | None:
+    """Return the effective dtype of an initializer ``value``.
+
+    Uses the value's declared ``type`` dtype, but falls back to the dtype of its
+    ``const_value`` when the type annotation is missing. When both are present
+    but disagree, the ``const_value`` dtype wins (it is the data actually
+    serialized) and a warning is logged, since a healthy initializer should
+    never have a declared type that contradicts its data.
+
+    Graph building can drop the declared ``type`` on an initializer while its
+    actual tensor data (``const_value``) still carries the correct dtype. In
+    that situation, defaulting to ``ir.DataType.FLOAT`` would emit fp32 weights
+    into an otherwise fp16 model. Reading the dtype from ``const_value`` keeps
+    folded initializers consistent with the weights they are derived from.
+
+    Returns ``None`` only when neither the declared type nor ``const_value`` is
+    available; callers decide on a final fallback.
+    """
+    declared = value.dtype
+    const_dtype = value.const_value.dtype if value.const_value is not None else None
+
+    if declared is not None and const_dtype is not None and declared != const_dtype:
+        logger.warning(
+            "Initializer %r declares dtype %s but its data is %s; using the data "
+            "dtype. This indicates stale type metadata.",
+            value.name,
+            declared,
+            const_dtype,
+        )
+        return const_dtype
+    if declared is not None:
+        return declared
+    return const_dtype

src/mobius/_passes/_dtype_utils_test.py

@@ -0,0 +1,49 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Tests for the shared initializer dtype helper."""
+
+from __future__ import annotations
+
+import logging
+
+import numpy as np
+import onnx_ir as ir
+
+from mobius._passes._dtype_utils import initializer_dtype
+
+
+def _value(dtype: ir.DataType | None, const: np.ndarray | None) -> ir.Value:
+    v = ir.Value(name="w")
+    if dtype is not None:
+        v.dtype = dtype
+    if const is not None:
+        v.const_value = ir.tensor(const)
+    return v
+
+
+class TestInitializerDtype:
+    def test_uses_declared_dtype_when_present(self):
+        v = _value(ir.DataType.FLOAT16, np.ones((2,), np.float16))
+        assert initializer_dtype(v) == ir.DataType.FLOAT16
+
+    def test_falls_back_to_const_value_when_declared_missing(self):
+        """The core fix: a dropped declared type must not hide the real dtype."""
+        v = _value(None, np.ones((2,), np.float16))
+        assert v.dtype is None
+        assert initializer_dtype(v) == ir.DataType.FLOAT16
+
+    def test_const_value_wins_on_disagreement(self, caplog):
+        """Stale declared metadata must not override the serialized data dtype."""
+        v = _value(ir.DataType.FLOAT, np.ones((2,), np.float16))
+        with caplog.at_level(logging.WARNING, logger="mobius._passes._dtype_utils"):
+            assert initializer_dtype(v) == ir.DataType.FLOAT16
+        assert any(
+            record.levelno == logging.WARNING
+            and "stale type metadata" in record.getMessage().lower()
+            for record in caplog.records
+        ), "expected a warning when declared dtype disagrees with const_value"
+
+    def test_returns_none_when_nothing_available(self):
+        v = _value(None, None)
+        assert initializer_dtype(v) is None

src/mobius/_passes/_fold_concat.py

@@ -27,6 +27,8 @@
 import numpy as np
 import onnx_ir as ir
 
+from mobius._passes._dtype_utils import initializer_dtype
+
 logger = logging.getLogger(__name__)
 
 
@@ -63,6 +65,19 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
             if not all(v is not None and v.is_initializer() for v in inputs):
                 continue
 
+            # All inputs must have loaded weights. Folding before weights are
+            # loaded would force the dtype to default to FLOAT and bake a wrong
+            # type into the packed initializer. Mirror FoldTransposedInitializer's
+            # guard and skip with a warning so the Concat survives for a later run.
+            if any(v.const_value is None for v in inputs):  # type: ignore[union-attr]
+                logger.warning(
+                    "FoldConcatInitializersPass: skipping Concat %r — an input "
+                    "initializer has no const_value (pass ran before weights "
+                    "were loaded, or a weight is missing).",
+                    node.name,
+                )
+                continue
+
             axis_attr = node.attributes.get("axis")
             axis: int = axis_attr.value if axis_attr is not None else 0
 
@@ -79,8 +94,10 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
 
             # Require uniform dtype — mixed-dtype concat is unusual and likely
             # a modelling error; skip and warn rather than silently produce
-            # a wrong result.
-            dtypes = [v.dtype for v in inputs]  # type: ignore[union-attr]
+            # a wrong result. Resolve each dtype from the declared type or, when
+            # that is missing, from the loaded ``const_value`` so an fp16 weight
+            # whose type annotation was dropped is not mistaken for fp32.
+            dtypes = [initializer_dtype(v) for v in inputs]  # type: ignore[union-attr]
             if len(set(dtypes)) > 1:
                 logger.warning(
                     "FoldConcatInitializersPass: skipping Concat with mixed dtypes %s"
@@ -89,6 +106,10 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
                 )
                 continue
 
+            # Authoritative dtype for the packed initializer. Falls back to FLOAT
+            # only when neither the declared type nor const_value is available.
+            packed_dtype = dtypes[0] or ir.DataType.FLOAT
+
             # Build a name for the packed initializer from the input names and axis.
             # The name encodes both the ordered input names and the axis, so two Concat
             # nodes with the same inputs in the same order along the same axis will
@@ -104,7 +125,12 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
             if packed_name in model.graph.initializers:
                 existing_val = model.graph.initializers[packed_name]
                 out_val.replace_all_uses_with(existing_val, replace_graph_outputs=True)
-                model.graph.remove(node)
+                # safe=True detaches the node from its inputs before removal so
+                # the folded q/k/v initializers' use lists are cleared. Otherwise
+                # `inp.uses()` still points at the removed Concat and the dead
+                # pre-pack weights survive RemoveUnusedNodesPass (~1.8 GB of
+                # orphaned weights serialized into the fp16 GQA model).
+                model.graph.remove(node, safe=True)
                 folded_count += 1
                 modified = True
 
@@ -117,7 +143,12 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
                 )
                 continue
 
-            new_val = ir.Value(name=packed_name, shape=out_shape, type=out_val.type)
+            # Stamp the resolved dtype on the new initializer's type so the
+            # declared type, the LazyTensor, and the materialized data all agree
+            # (out_val.type may be missing after stage-2 rewrites).
+            new_val = ir.Value(
+                name=packed_name, shape=out_shape, type=ir.TensorType(packed_dtype)
+            )
 
             captured_inputs = list(inputs)  # capture for closure
 
@@ -133,13 +164,17 @@ def _make_packed(
                     arrays.append(v.const_value.numpy())
                 return ir.tensor(np.concatenate(arrays, axis=ax))
 
-            dtype = inputs[0].dtype or ir.DataType.FLOAT  # type: ignore[union-attr]
-            new_val.const_value = ir.LazyTensor(_make_packed, dtype=dtype, shape=out_shape)
+            new_val.const_value = ir.LazyTensor(
+                _make_packed, dtype=packed_dtype, shape=out_shape
+            )
 
             model.graph.initializers[new_val.name] = new_val
 
             out_val.replace_all_uses_with(new_val, replace_graph_outputs=True)
-            model.graph.remove(node)
+            # safe=True detaches the node from its inputs before removal so the
+            # folded q/k/v initializers' use lists are cleared; otherwise the
+            # dead pre-pack weights survive RemoveUnusedNodesPass.
+            model.graph.remove(node, safe=True)
             folded_count += 1
             modified = True