perf(xla): pre-transpose the LM head for single-sequence decode

src/lib/mlxcel-xla/src/emitter/mod.rs

@@ -1456,6 +1456,46 @@ mod tests {
         assert_eq!(f32.matches("xf16> loc(").count(), 0);
     }
 
+    /// issue #2863: with the `InOut` layout the single-sequence graphs take the
+    /// LM head as one transposed `[H, V]` operand and contract the hidden axis,
+    /// so IREE never materializes the per-step weight transpose and convert. The
+    /// `OutIn` graph keeps the checkpoint `[V, H]` head and its `[0] x [1]` form
+    /// (and therefore its byte-exact goldens).
+    #[test]
+    fn inout_graphs_take_the_lm_head_transposed() {
+        use super::builder::Precision;
+        let mut c = Config::llama_3_2_1b();
+        c.projection_layout = ProjectionLayout::InOut;
+        let mlir = emit_decode_with(&c, false, Precision::F16);
+        assert!(
+            mlir.contains("tensor<2048x128256xf16> loc(\"params['embed']\")"),
+            "the InOut head is one transposed f16 arg"
+        );
+        let head = mlir
+            .lines()
+            .find(|line| line.contains("stablehlo.dot_general") && line.contains("128256"))
+            .expect("head dot");
+        assert!(
+            head.contains("contracting_dims = [0] x [0]"),
+            "head dot contracts the hidden axis: {head}"
+        );
+        assert!(
+            head.contains("tensor<2048xf16>, tensor<2048x128256xf16>"),
+            "head operands are the hidden vector and the transposed head: {head}"
+        );
+
+        let out = emit_decode_with(&Config::llama_3_2_1b(), false, Precision::F16);
+        assert!(!out.contains("tensor<2048x128256xf16>"));
+        let head = out
+            .lines()
+            .find(|line| line.contains("stablehlo.dot_general") && line.contains("128256"))
+            .expect("head dot");
+        assert!(
+            head.contains("contracting_dims = [0] x [1]"),
+            "the OutIn head keeps the checkpoint layout: {head}"
+        );
+    }
+
     // ======================================================================
     // dense arch pack (issue #498)
     // ======================================================================

src/lib/mlxcel-xla/src/emitter/model.rs

@@ -412,11 +412,7 @@ fn take_weight(
         // token-exact while halving the weight's per-step DRAM read. f32 / bf16
         // precision and the fused layouts keep the f32 arg (byte-identical goldens).
         _ => {
-            let elt = if b.precision() == Precision::F16 && c.supports_f16_resident() {
-                "f16"
-            } else {
-                "f32"
-            };
+            let elt = projection_elt(b, c);
             // issue #2859: in the `InOut` layout the loader transposes the
             // checkpoint's `[out, in]` projection once at upload, so the dense
             // forward can contract the shared axis instead of transposing every
@@ -430,29 +426,69 @@ fn take_weight(
     }
 }
 
-/// Take the untied LM head weight `params['lm_head']` (`[V, H]`), or `None` for a
-/// tied checkpoint (which reuses `embed` for the final projection). Called right
-/// after `final_norm` and before the layers, so the weight arg order is embed,
-/// final_norm, [lm_head when untied], layers..., matching `weight_names` in
-/// `iree.rs`. For a tied model nothing is emitted, so the graph stays byte-
-/// identical (the guard that keeps every tied checkpoint unchanged).
-fn take_lm_head(decls: &mut Vec<ArgDecl>, idx: &mut usize, c: &Config) -> Option<Val> {
-    if c.tie_word_embeddings {
-        None
+/// Element type of a projection weight argument: f16 on the f16 GPU path, whose
+/// loader uploads the weight f16-resident (issue #572), else f32. Shared by the
+/// dense projection args and the single-sequence head operand (issue #2863) so
+/// the two cannot drift from the loader's `resident_f16` decision.
+fn projection_elt(b: &Builder, c: &Config) -> &'static str {
+    if b.precision() == Precision::F16 && c.supports_f16_resident() {
+        "f16"
     } else {
-        Some(take_arg(
-            decls,
-            idx,
-            Ty::f32(vec![c.vocab, c.hidden]),
-            "params['lm_head']".into(),
-        ))
+        "f32"
+    }
+}
+
+/// Take the final logits operand for the layout this graph declares. Called
+/// right after `final_norm` and before the layers, so the weight arg order is
+/// embed, final_norm, [lm_head or head operand], layers..., matching
+/// `weight_names` in `iree.rs`.
+///
+/// In [`ProjectionLayout::OutIn`] (the ragged batch engine) this is the untied
+/// `params['lm_head']` (`[V, H]`), or `None` for a tied checkpoint that reuses
+/// `embed`; that path stays byte-identical. In [`ProjectionLayout::InOut`] (the
+/// single-sequence graphs, issue #2863) the head always gets its own `[H, V]`
+/// operand that the loader uploads transposed once, so the per-step
+/// `dot_general` contracts the hidden axis and IREE stops materializing a
+/// full-weight transpose fused with the f32 -> f16 convert on every decode step.
+/// A tied checkpoint keeps its `[V, H]` embed arg for the gather and adds this
+/// second arg, which the loader reads from the same tensor.
+fn take_lm_head(decls: &mut Vec<ArgDecl>, idx: &mut usize, c: &Config, b: &Builder) -> Option<Val> {
+    match b.projection_layout() {
+        ProjectionLayout::InOut => {
+            let elt = projection_elt(b, c);
+            let loc = if c.tie_word_embeddings {
+                "params['embed']"
+            } else {
+                "params['lm_head']"
+            };
+            Some(take_arg(
+                decls,
+                idx,
+                Ty::new(vec![c.hidden, c.vocab], elt),
+                loc.into(),
+            ))
+        }
+        ProjectionLayout::OutIn => {
+            if c.tie_word_embeddings {
+                None
+            } else {
+                Some(take_arg(
+                    decls,
+                    idx,
+                    Ty::f32(vec![c.vocab, c.hidden]),
+                    "params['lm_head']".into(),
+                ))
+            }
+        }
     }
 }
 
 /// The weight the final logits projection multiplies by: the dedicated `lm_head`
 /// for an untied checkpoint, else the tied token-embedding matrix. Both are
 /// `[V, H]` (`linear` computes `x @ W^T`), so the tail is identical apart from
-/// which buffer it reads.
+/// which buffer it reads. In [`ProjectionLayout::InOut`] the taken head is
+/// always `Some` (the transposed `[H, V]` operand), so the `embed` fallback only
+/// serves the batch engine.
 fn head_weight<'a>(embed: &'a Val, lm_head: &'a Option<Val>) -> &'a Val {
     lm_head.as_ref().unwrap_or(embed)
 }
@@ -797,7 +833,7 @@ fn build_arg_schema(b: &mut Builder, c: &Config) -> (Vec<ArgDecl>, Args) {
         "params['final_norm']".into(),
     );
     let final_norm_bias = take_final_norm_bias(&mut decls, &mut idx, c);
-    let lm_head = take_lm_head(&mut decls, &mut idx, c);
+    let lm_head = take_lm_head(&mut decls, &mut idx, c, b);
 
     let mut layers = Vec::with_capacity(c.n_layers);
     for li in 0..c.n_layers {
@@ -2727,7 +2763,7 @@ pub fn emit_decode_with(c: &Config, sample: bool, precision: Precision) -> Strin
     // Granite / MiniCPM divide), then optional on-device argmax ---
     let final_w = norm_w(&mut b, &a.final_norm, c, &k, h);
     let xf = normalize(&mut b, c, &k, &x, &final_w, a.final_norm_bias.as_ref(), h);
-    let logits = b.linear(&xf, head_weight(&a.embed, &a.lm_head)); // [V]
+    let logits = b.linear_proj(&xf, head_weight(&a.embed, &a.lm_head)); // [V]
     let logits = apply_logit_scale(&mut b, c, logits);
     let (out_val, out_ty) = if sample {
         let tok = b.argmax(&logits);
@@ -2802,7 +2838,7 @@ fn build_batched_arg_schema(
         "params['final_norm']".into(),
     );
     let final_norm_bias = take_final_norm_bias(&mut decls, &mut idx, c);
-    let lm_head = take_lm_head(&mut decls, &mut idx, c);
+    let lm_head = take_lm_head(&mut decls, &mut idx, c, b);
 
     let mut layers = Vec::with_capacity(c.n_layers);
     for li in 0..c.n_layers {
@@ -3096,7 +3132,7 @@ pub fn emit_decode_batched_with(
     // --- tail: final norm + LM head (tied embed or untied lm_head) -> [B, V],
     // optional per-row argmax ---
     let xf = rms_norm_seq(&mut b, &x, &a.final_norm, &k, bsz, h); // [B, H]
-    let logits = b.linear_seq(&xf, head_weight(&a.embed, &a.lm_head)); // [B, V]
+    let logits = b.linear_seq_proj(&xf, head_weight(&a.embed, &a.lm_head)); // [B, V]
     let (out_val, out_ty) = if sample {
         let tok = b.argmax_batched(&logits);
         (tok.name, Ty::new(vec![bsz], "i32").render())
@@ -3166,7 +3202,7 @@ fn build_ragged_arg_schema(b: &mut Builder, c: &Config, bsz: usize) -> (Vec<ArgD
         "params['final_norm']".into(),
     );
     let final_norm_bias = take_final_norm_bias(&mut decls, &mut idx, c);
-    let lm_head = take_lm_head(&mut decls, &mut idx, c);
+    let lm_head = take_lm_head(&mut decls, &mut idx, c, b);
 
     let mut layers = Vec::with_capacity(c.n_layers);
     for li in 0..c.n_layers {
@@ -3355,7 +3391,7 @@ pub fn emit_decode_ragged_with(
         bsz,
         h,
     );
-    let logits = b.linear_seq(&xf, head_weight(&a.embed, &a.lm_head)); // [B, V]
+    let logits = b.linear_seq_proj(&xf, head_weight(&a.embed, &a.lm_head)); // [B, V]
     // Final logit scaling (Gemma soft-cap / Cohere multiply / Granite / MiniCPM
     // divide), per row; argmax-invariant but kept for exactness.
     let logits = apply_logit_scale(&mut b, c, logits);
@@ -3461,7 +3497,7 @@ fn build_prefill_arg_schema(
         "params['final_norm']".into(),
     );
     let final_norm_bias = take_final_norm_bias(&mut decls, &mut idx, c);
-    let lm_head = take_lm_head(&mut decls, &mut idx, c);
+    let lm_head = take_lm_head(&mut decls, &mut idx, c, b);
 
     let mut layers = Vec::with_capacity(c.n_layers);
     for li in 0..c.n_layers {
@@ -4054,7 +4090,7 @@ fn emit_prefill_module(
     let last_idx = b.subtract(&a.real_len, &one_i); // real_len - 1
     let last_row = b.dynamic_slice(&xf, &[&last_idx, &k.c0], vec![1, h]); // [1, H]
     let last = b.reshape(&last_row, vec![h]); // [H]
-    let logits = b.linear(&last, head_weight(&a.embed, &a.lm_head)); // [V]
+    let logits = b.linear_proj(&last, head_weight(&a.embed, &a.lm_head)); // [V]
     let logits = apply_logit_scale(&mut b, c, logits);
     let (out_val, out_ty) = if sample {
         let tok = b.argmax(&logits);

src/lib/mlxcel-xla/src/weights.rs

@@ -23,7 +23,7 @@
 //! feature-gated `iree.rs`) so it is unit-tested without the IREE runtime and stays
 //! in lock-step with the emitter's arg order (`emitter::model::take_layer_weights`).
 
-use crate::emitter::{Config, quant_in_graph};
+use crate::emitter::{Config, ProjectionLayout, quant_in_graph};
 
 /// One checkpoint weight the loader reads, in the emitter's arg order. Most are a
 /// whole safetensors tensor; a Phi3 checkpoint fuses q/k/v into one `qkv_proj` and
@@ -197,8 +197,25 @@ fn weight_specs_q(cfg: &Config, quant: bool) -> Vec<WeightSpec> {
     if cfg.norm_bias {
         out.push(WeightSpec::Whole("model.norm.bias".to_string()));
     }
-    if !cfg.tie_word_embeddings {
-        out.push(WeightSpec::Whole(s.lm_head.to_string()));
+    // The final logits operand, in the layout the graph declares (issue #2863).
+    // `OutIn` keeps the checkpoint's `[V, H]` head (the untied `lm_head`, or the
+    // tied `embed` for which no extra spec is pushed). `InOut` always uploads a
+    // transposed `[H, V]` operand: the untied `lm_head`, or a second read of the
+    // tied `embed` tensor after the gather's own `Whole` spec.
+    match cfg.projection_layout {
+        ProjectionLayout::OutIn => {
+            if !cfg.tie_word_embeddings {
+                out.push(WeightSpec::Whole(s.lm_head.to_string()));
+            }
+        }
+        ProjectionLayout::InOut => {
+            let name = if cfg.tie_word_embeddings {
+                s.embed.to_string()
+            } else {
+                s.lm_head.to_string()
+            };
+            out.push(WeightSpec::Proj(name));
+        }
     }
     for i in 0..cfg.n_layers {
         let p = format!("{}{i}.", s.layer_stem);
@@ -1065,6 +1082,59 @@ mod tests {
         }
     }
 
+    #[test]
+    fn inout_head_spec_is_a_transposed_projection() {
+        // Tied (llama-3.2-1b): the embed tensor is read twice, once as the
+        // gather's whole tensor and once as the transposed head operand.
+        let tied = Config::llama_3_2_1b();
+        let out_in = weight_specs_q(&tied, false);
+        assert_eq!(
+            out_in[0],
+            WeightSpec::Whole("model.embed_tokens.weight".to_string())
+        );
+        assert_eq!(
+            out_in[1],
+            WeightSpec::Whole("model.norm.weight".to_string())
+        );
+        assert!(
+            !out_in.iter().any(
+                |spec| matches!(spec, WeightSpec::Proj(name) if name == "model.embed_tokens.weight")
+            ),
+            "the tied OutIn head adds no projection spec"
+        );
+
+        let mut inout = tied.clone();
+        inout.projection_layout = ProjectionLayout::InOut;
+        let specs = weight_specs_q(&inout, false);
+        assert_eq!(
+            specs[0],
+            WeightSpec::Whole("model.embed_tokens.weight".to_string())
+        );
+        assert_eq!(
+            specs[2],
+            WeightSpec::Proj("model.embed_tokens.weight".to_string()),
+            "the InOut head follows the norm as a transposed projection"
+        );
+
+        // Untied: the existing lm_head arg becomes the transposed projection
+        // instead of a fourth whole tensor.
+        let mut untied = tied;
+        untied.tie_word_embeddings = false;
+        untied.projection_layout = ProjectionLayout::InOut;
+        let specs = weight_specs_q(&untied, true);
+        assert_eq!(
+            specs[2],
+            WeightSpec::Proj("lm_head.weight".to_string()),
+            "the untied InOut head is the lm_head as a projection"
+        );
+        assert!(
+            !specs
+                .iter()
+                .any(|spec| matches!(spec, WeightSpec::Whole(name) if name == "lm_head.weight")),
+            "the packed path never declares a whole lm_head arg"
+        );
+    }
+
     #[test]
     fn phi4mm_weight_schema_uses_base_layers_and_complete_adapter_pairs() {
         let config = phi4_config();