From 8c096404f66d1318aa862deec724eeffc19fbb29 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 30 Jul 2026 05:02:28 +0000 Subject: [PATCH 1/2] Initial plan From 4116059c963d50d366d0ca5504c072c02208ed72 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 30 Jul 2026 05:09:38 +0000 Subject: [PATCH 2/2] fix: stamp weight dtype on GQA PackQKV Concat/Transpose intermediates The PackQKV rewrite built the packed QKV weight with `op.Concat` / `op.Transpose`, whose output values carry no declared type. When those intermediates were folded into initializers, there was no declared dtype to inherit, so fp16 weights could be widened to fp32 (mitigated downstream in #351 by const_value fallback + fold-pass stamping). Verified the issue still reproduces on main (packed Concat/Transpose outputs had `dtype is None`) and fixed it at the source: propagate the projection weight dtype (and bias dtype for the biased path) onto the new intermediates in both PackQKV rewrite sites. The downstream mitigation is kept as defense-in-depth. --- .../rewrite_rules/_group_query_attention.py | 28 ++++++++++++ .../_group_query_attention_test.py | 44 +++++++++++++++++++ 2 files changed, 72 insertions(+) diff --git a/src/mobius/rewrite_rules/_group_query_attention.py b/src/mobius/rewrite_rules/_group_query_attention.py index f90b772f..d109b6fe 100644 --- a/src/mobius/rewrite_rules/_group_query_attention.py +++ b/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) diff --git a/src/mobius/rewrite_rules/_group_query_attention_test.py b/src/mobius/rewrite_rules/_group_query_attention_test.py index 6d95812d..65c4bbf5 100644 --- a/src/mobius/rewrite_rules/_group_query_attention_test.py +++ b/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)