Stamp weight dtype on GQA PackQKV Concat/Transpose intermediates

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)