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)