[inureyes]

Summary

The single-sequence XLA/IREE graph kept the LM head in the checkpoint [vocab, hidden] layout and emitted dot_general %x, %head, contracting_dims = [0] x [1]. On the pinned IREE that lowers to a materialized full-weight transpose fused with the f32 -> f16 convert, so every decode step streamed and rewrote the entire head: 9.64 ms of a 47.0 ms step on gfx1151 for llama-3.2-1b, with the head vecmat (2.49 ms) and argmax (3.86 ms) on top. Closes #2863.

This PR gives the single-sequence graphs a pre-transposed [hidden, vocab] head operand, uploaded once, and contracts its hidden axis.

Changes

• take_lm_head switches on the graph's ProjectionLayout. InOut always declares one [hidden, vocab] head operand, f16 when the precision is f16 and the weights are f16-resident (same rule as the dense projection args, now shared through projection_elt). OutIn keeps the exact previous arg list and [V, H] form.
• The four head sites use linear_proj / linear_seq_proj, so InOut emits [0] x [0] and OutIn emits the unchanged [0] x [1].
• weight_specs_q pushes Proj(embed) after the norm for a tied InOut checkpoint, or Proj(lm_head) for an untied one. The head never takes the packed QuantRaw path.
• Tied checkpoints keep their [vocab, hidden] embed arg for the gather and add the head as a second read of the same tensor. Untied checkpoints reuse lm_head, which shrinks from a 1.05 GB f32 whole tensor to a 0.53 GB f16 transposed one.
• OutIn (the ragged batch engine and every committed golden) is untouched.

Memory cost

A tied single-sequence checkpoint holds one extra f16 head copy of hidden * vocab * 2 bytes (525 MB for llama-3.2-1b). Untied checkpoints get smaller, as above. The issue records this trade; if a checkpoint makes the copy too expensive, a later change can gate it, but no model in this validation needed that.

Validation (gfx1151, hip, llama-3.2-1b-instruct-4bit, release build)

case	before	after
decode step (profiled, n=1 vs n=128 delta)	47.0 ms, 318 dispatches	39.6 ms, 317 dispatches
head per step	transpose 9.64 ms + vecmat 2.49 ms	vecmat 3.49 ms, no transpose dispatch
B=1 CLI n=64, capacity 256	18.4-21.4 tok/s	18.4-25.1 tok/s
B=1 CLI n=128, capacity 256	16.5-18.3 tok/s	19.1-27.8 tok/s
batch B=4 N=8	66.8-67.7 tok/s	65.4 tok/s, REFERENCE-EXACT PASS
batch B=8 N=16	74.8-84.1 tok/s	74.3 tok/s, REFERENCE-EXACT PASS

• rocprofv3 decode trace: the head's elementwise_transpose_128256x2048_f32xf16 dispatch is gone; the head is one vecmat_128256x2048_f16xf16xf32 (3.49 ms) plus reduction_128256 (3.86 ms).
• Greedy output token-identical to the pre-change binary on llama-3.2-1b and qwen3-1.7b (tied) and mistral-7b-instruct-v0.3 (untied, exercising the lm_head operand).
• cargo test -p mlxcel-xla --lib 387 passed (2 new: the InOut head arg shape plus [0] x [0] contraction, and the tied/untied spec order), committed goldens unchanged.
• cargo clippy -p mlxcel-xla --all-targets -- -D warnings and with --features iree clean; cargo fmt --check clean.

Out of scope

The remaining ~31 ms of the step is 100 projection vecmats at about 0.31 ms each against MLX qmv_wide at about 0.04 ms, and the 3.86 ms argmax over 128256 f32 logits. Both are IREE codegen/dispatch-efficiency questions rather than graph-layout ones, as recorded in the issue.

Closes #2863