GQA PackQKV rewrite leaves Concat/Transpose intermediates without a declared dtype

[titaiwangms]

Summary

The GQA PackQKV rewrite produces Concat/Transpose intermediate ir.Values
that carry no declared .dtype. This is the most likely source of the fp16
dtype-drop that PR #351 mitigated downstream — the declared type is not carried
on these intermediates (code-consistent and mechanistically the cause; not
separately reproduced in isolation here). The mitigation shipped in #351 is robust and
proven by tests, so this issue is a robustness / clarity follow-up to fix the
problem at its source — not an active correctness bug.

Root cause

In src/mobius/rewrite_rules/_group_query_attention.py, both PackQKV replacement
paths build the packed weight as graph ops:

# ~lines 344-346 (and again ~lines 467-469)
packed_w  = op.Concat(q_w, k_w, v_w, axis=0)
packed_wt = op.Transpose(packed_w, perm=[1, 0])

The intermediate ir.Values returned by op.Concat(...) / op.Transpose(...)
have no declared .dtype. For an fp16 model, the packed-QKV MatMul(hidden, Transpose(Concat(W_q, W_k, W_v))) therefore feeds the fold passes a Concat
output whose declared type is unset, and the fold passes would otherwise default
the folded initializer to FLOAT — silently widening fp16 weights to fp32.

Note: the cast site hypothesised during review — _cast_module_dtype in
src/mobius/_builder.py (line 101, param.type = ir.TensorType(dtype)) —
already stamps the declared dtype on module parameters (pre-existing code).
It is not the drop site. The unset declared type originates at the GQA rewrite
intermediates above, which is why the source fix belongs there.

Current mitigation (already shipped in #351)

PR #351 added a two-layer defense that keeps the fold path correct without
touching the rewrite:

1. const_value fallback — initializer_dtype() in
   src/mobius/_passes/_dtype_utils.py resolves the dtype from const_value
   when an initializer's declared type is missing.
2. Fold-pass dtype stamping — FoldConcatInitializersPass and
   FoldTransposeInitializersPass stamp the resolved dtype onto the new
   packed/transposed initializer (type=ir.TensorType(...) at
   _fold_concat.py:150 and _fold_transpose.py:128) instead of inheriting the
   unset type or defaulting to FLOAT.

This is guarded end-to-end by src/mobius/_passes/_fold_dtype_e2e_test.py, which
drives the real export pipeline and asserts the packed-QKV weights stay fp16
through fold + serialization round-trip. So the fold path is robust today.

Proposed fix

Fix the dtype at the source so the type stays consistent from where the packed
weight is created, rather than relying solely on downstream recovery:

• Stamp the declared dtype on the packed_w / packed_wt intermediates at both
  GQA rewrite sites (_group_query_attention.py ~344-346 and ~467-469) — e.g.
  set their TensorType from the incoming q_w/k_w/v_w dtype; or
• Reorder shape/type inference to run before the fold passes so the Concat
  output is type-stamped before folding.

The fold-pass stamping + const_value fallback should be kept as
defense-in-depth even after the source fix.

Why this was deferred from #351: it is a multi-call-site change in a file
outside that bug-fix PR's scope, and validating a GQA rewrite change requires
the onnxruntime_easy test dependency (the
src/mobius/rewrite_rules/_group_query_attention_test.py suite cannot collect
without it), so it was kept out of the focused bug-fix PR.

References

• PR Fix fp16/GQA static-cache export correctness (split from #340) #351 — the bug-fix split that added the two-layer mitigation.
• Raised by @justinchuby's review comment on _fold_transpose.py
  ("which process drops the dtype? can we fix it too?").