perf(xla): remove the per-step LM-head transpose and f32-to-f16 conversion in the single-sequence decode graph #2863

Problem

After #2859 (PR #2861, merged bb1072d4) moved the dense projection weights to an [in, out] layout for the single-sequence graph, a decode step on gfx1151 is 47.0 ms of device time over 318 dispatches for llama-3.2-1b-instruct-4bit (hip, release build). The tied LM head still uses the checkpoint layout and costs 16.0 ms of that step:

The emitted head is stablehlo.dot_general %x, %embed, contracting_dims = [0] x [1] against the f32 params['embed'] tensor [vocab, hidden] (see src/lib/mlxcel-xla/assets/llama-3.2-1b/decode.mlir, the dot_general before the final reduce). On this IREE pin (3.12.0rc20260721) that form lowers to a materialized full-weight transpose fused with the f32 -> f16 convert, then a vecmat, so every decode step streams and rewrites 525 MB of head weights. MLX ROCm does the same step in 4.1 ms with no transposes.

Controlled experiment from #2859: compiling [0] x [1] against [out, in] produces two dispatches (transpose + vecmat), while [0] x [0] against a pre-transposed [in, out] operand produces one.

Proposed change

1. For ProjectionLayout::InOut graphs only (the single-sequence IreeLlama path), upload the head operand once as [hidden, vocab] and emit the head matmul with contracting dims [0] x [0] (reuse Builder::linear_t / linear_seq_t).
2. Untied checkpoints (lm_head.weight present): replace the single Whole head argument with the transposed projection spec, so no extra tensor is added.
3. Tied checkpoints (llama-3.2-1b): the embed tensor stays [vocab, hidden] for the gather, and one additional transposed head argument is uploaded. Declare it f16 when the precision is f16 and the device supports f16-resident weights, matching the projection arguments; this also halves the head's resident bytes for tied models that currently hold it as f32.
4. Keep ProjectionLayout::OutIn (the ragged batch engine and every committed golden) unchanged, and keep Gemma3n's own emitters unchanged.

The extra tied-head copy costs hidden * vocab * 2 bytes (525 MB for llama-3.2-1b). The issue must state that cost and, if it turns out too large for a checkpoint, gate the extra copy so the graph falls back to the current head form instead of failing the load.

Acceptance criteria

• The single-sequence decode trace for llama-3.2-1b on gfx1151 contains no elementwise_transpose dispatch for the head, and the per-step device time for the head drops from 16.0 ms to the vecmat plus reduction.
• Greedy output stays token-identical to the current binary on llama-3.2-1b (tied) and at least one untied checkpoint, and the xla_batch_bench REFERENCE-EXACT gate still passes.
• Committed goldens are unchanged (the change is InOut-only); cargo test -p mlxcel-xla --lib and clippy at default and iree features stay green.
• Measured before/after numbers for B=1 n=64 and n=128 at capacity 256 in the PR body.

Out of scope

The remaining ~31 ms of the step is 100 projection vecmats at about 0.31 ms each, while MLX's qmv_wide path measures about 0.04 ms per call; that is an IREE codegen/dispatch-efficiency gap, not a graph-layout problem, and needs its own investigation with the profiler. The argmax reduction (3.86 ms for 128256 f32 logits) may also deserve its own look once the head operand is f16.

Refs #2859