Stamp weight dtype on GQA PackQKV Concat/Transpose intermediates (#435)

The `PackQKVForGQA` / `PackQKVWithBiasForGQA` rewrites build the packed
weight with `op.Concat` / `op.Transpose`, whose output values carry no
declared type. When those intermediates are later folded into
initializers there is no declared dtype to inherit, so an fp16 model
could end up with fp32 packed weights — the drop that #351 mitigated
downstream.

The issue asked whether this still reproduces. It does. On `main`, a
CPU-EP llama build gives:

```
Concat    node_Concat_0    [FLOAT, FLOAT, FLOAT] -> None
Transpose node_Transpose_1 [None]                -> None
```

### Changes

- **`_group_query_attention.py`** — new `_propagate_dtype(source,
*targets)` helper resolves the source parameter's dtype through the
existing `initializer_dtype()` (declared type, `const_value` fallback,
fail-closed on mismatch) and stamps it on the rewrite-created values.
- **`PackQKVForGQA.rewrite`** — stamps `packed_w` / `packed_wt` from
`q_w`.
- **`PackQKVWithBiasForGQA.rewrite`** — same, plus `packed_bias` from
`bias_q`. The bias `Concat` folds into an initializer too and had the
identical untyped-output problem, so it is covered here rather than left
for a follow-up.
- **Tests** — parameterized fp32/fp16 regression tests for both rewrite
paths, asserting the `Concat`/`Transpose`/bias-`Concat` outputs declare
the model dtype. They fail on unmodified source (`assert None ==
FLOAT16`).

```python
packed_w = op.Concat(q_w, k_w, v_w, axis=0)
packed_wt = op.Transpose(packed_w, perm=[1, 0])
# Concat/Transpose outputs have no declared type; carry the weight dtype
# forward so the folded packed initializer keeps the model dtype.
_propagate_dtype(q_w, packed_w, packed_wt)
```

The #351 fold-pass stamping and `const_value` fallback are left
untouched as defense-in-depth.

Note: the `onnxruntime_easy` collection blocker cited as the reason this
was deferred from #351 no longer applies — it ships in the `testing`
extra, so `src/mobius/rewrite_rules/_group_query_attention_test.py` runs
directly.

<!-- START COPILOT CODING AGENT SUFFIX -->

- Fixes #355

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: Justin Chu <justinchuby@users.noreply.github.com>

Author: Copilot

src/mobius/rewrite_rules/_group_query_attention.py

@@ -43,12 +43,33 @@
 
 from __future__ import annotations
 
+import onnx_ir as ir
 from onnxscript.rewriter._basics import MatchFailureError, MatchResult
 from onnxscript.rewriter._rewrite_rule import (
     RewriteRuleClassBase,
     RewriteRuleSet,
 )
 
+from mobius._passes._dtype_utils import initializer_dtype
+
+
+def _propagate_dtype(source: ir.Value, *targets: ir.Value) -> None:
+    """Stamp ``source``'s dtype onto rewrite-created intermediate values.
+
+    Values produced by the replacement builder (``op.Concat(...)``,
+    ``op.Transpose(...)``) carry no declared type. When those intermediates are
+    later folded into initializers, a missing declared type forces the fold
+    passes to recover the dtype from ``const_value`` — or, absent that, to
+    default to ``FLOAT``, silently widening fp16 weights. Stamping the dtype of
+    the parameter the value is derived from keeps the type consistent from the
+    point the packed weight is created.
+    """
+    dtype = initializer_dtype(source)
+    if dtype is None:
+        return
+    for target in targets:
+        target.dtype = dtype
+
 
 def _has_unequal_kv_head_dimensions(k, v, past_key, past_value) -> bool:
     """Return whether static K/V shapes prove incompatible GQA head dimensions."""
@@ -364,6 +385,9 @@ def rewrite(
         packed_w = op.Concat(q_w, k_w, v_w, axis=0)
         # Transpose packed weight: (q_out+k_out+v_out, hidden) → (hidden, q_out+k_out+v_out)
         packed_wt = op.Transpose(packed_w, perm=[1, 0])
+        # Concat/Transpose outputs have no declared type; carry the weight dtype
+        # forward so the folded packed initializer keeps the model dtype.
+        _propagate_dtype(q_w, packed_w, packed_wt)
         packed_qkv = op.MatMul(hidden, packed_wt)
 
         # Recover remaining GQA inputs and attributes from the matched node
@@ -487,10 +511,14 @@ def rewrite(
         packed_w = op.Concat(q_w, k_w, v_w, axis=0)
         # Transpose packed weight: (q_out+k_out+v_out, hidden) → (hidden, q_out+k_out+v_out)
         packed_wt = op.Transpose(packed_w, perm=[1, 0])
+        # Concat/Transpose outputs have no declared type; carry the weight dtype
+        # forward so the folded packed initializer keeps the model dtype.
+        _propagate_dtype(q_w, packed_w, packed_wt)
         packed_mm = op.MatMul(hidden, packed_wt)
 
         # Concat biases: (q_out + k_out + v_out,)
         packed_bias = op.Concat(bias_q, bias_k, bias_v, axis=0)
+        _propagate_dtype(bias_q, packed_bias)
         # GQA has no bias input — the Add stays in the graph.
         packed_qkv = op.Add(packed_mm, packed_bias)
 

src/mobius/rewrite_rules/_group_query_attention_test.py

@@ -701,6 +701,50 @@ def test_packed_qkv_with_bias_uses_concat_nodes(self):
                 "Transpose input should be Concat of W_q, W_k, W_v"
             )
 
+    @pytest.mark.parametrize("dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16])
+    def test_packed_weight_intermediates_declare_weight_dtype(self, dtype):
+        """Concat/Transpose intermediates carry the projection weight dtype.
+
+        The replacement builder leaves new values untyped. Without an explicit
+        stamp, folding ``Transpose(Concat(W_q, W_k, W_v))`` into an initializer
+        has no declared type to inherit and can widen fp16 weights to fp32.
+        """
+        config = dataclasses.replace(_LLAMA_CONFIG, dtype=dtype)
+        m = build_from_module(registry.get("llama")(config), config)["model"]
+
+        rewrite(m, pattern_rewrite_rules=group_query_attention_rules())
+        rewrite(m, pattern_rewrite_rules=pack_qkv_for_gqa_rules())
+
+        gqa_nodes = [n for n in m.graph if n.op_type == "GroupQueryAttention"]
+        assert len(gqa_nodes) == config.num_hidden_layers
+
+        for gqa in gqa_nodes:
+            transpose = gqa.inputs[0].producer().inputs[1].producer()
+            assert transpose.op_type == "Transpose"
+            concat = transpose.inputs[0].producer()
+            assert concat.op_type == "Concat"
+            assert concat.outputs[0].dtype == dtype
+            assert transpose.outputs[0].dtype == dtype
+
+    @pytest.mark.parametrize("dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16])
+    def test_packed_bias_intermediate_declares_bias_dtype(self, dtype):
+        """The packed-bias Concat intermediate carries the bias dtype."""
+        config = dataclasses.replace(_QWEN2_BIAS_CONFIG, dtype=dtype)
+        m = build_from_module(registry.get("qwen2")(config), config)["model"]
+
+        rewrite(m, pattern_rewrite_rules=group_query_attention_rules())
+        rewrite(m, pattern_rewrite_rules=pack_qkv_for_gqa_rules())
+
+        gqa_nodes = [n for n in m.graph if n.op_type == "GroupQueryAttention"]
+        assert len(gqa_nodes) == config.num_hidden_layers
+
+        for gqa in gqa_nodes:
+            add = gqa.inputs[0].producer()
+            assert add.op_type == "Add"
+            bias_concat = add.inputs[1].producer()
+            assert bias_concat.op_type == "Concat"
+            assert bias_concat.outputs[0].dtype == dtype
+
     def test_packed_qkv_with_bias_runs_with_ort(self):
         """Biased packed-QKV GQA model runs correctly with ORT."""
         model = registry.get("qwen2")(_QWEN2_BIAS_CONFIG)