diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 281ebea8a..d402ef6cb 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/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) // ====================================================================== diff --git a/src/lib/mlxcel-xla/src/emitter/model.rs b/src/lib/mlxcel-xla/src/emitter/model.rs index 991ca9273..f1fb8b4f8 100644 --- a/src/lib/mlxcel-xla/src/emitter/model.rs +++ b/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, idx: &mut usize, c: &Config) -> Option { - 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, idx: &mut usize, c: &Config, b: &Builder) -> Option { + 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) -> &'a Val { lm_head.as_ref().unwrap_or(embed) } @@ -797,7 +833,7 @@ fn build_arg_schema(b: &mut Builder, c: &Config) -> (Vec, 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 Vec { 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();