diff --git a/src/lib/mlxcel-xla/src/emitter/builder.rs b/src/lib/mlxcel-xla/src/emitter/builder.rs index f9751155f..f5bcdb274 100644 --- a/src/lib/mlxcel-xla/src/emitter/builder.rs +++ b/src/lib/mlxcel-xla/src/emitter/builder.rs @@ -318,6 +318,25 @@ pub struct Builder { next: usize, precision: Precision, gemma3n_qmv: bool, + projection_layout: ProjectionLayout, +} + +/// Layout a graph declares for projection weights (issue #2859). +/// +/// The best layout depends on the graph shape, not on the model: IREE lowers a +/// single-row contraction against `[out, in]` into a separate full-weight +/// transpose plus a vecmat (slow), but the same contraction at the ragged +/// decode's row count into one transposed-B matmul (fast). The converse holds +/// for the pre-transposed layout. The single-sequence session therefore loads +/// [`ProjectionLayout::InOut`] and the continuous-batching engine keeps +/// [`ProjectionLayout::OutIn`]. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum ProjectionLayout { + /// The checkpoint's `[out, in]` layout, uploaded unchanged. + #[default] + OutIn, + /// `[in, out]`, transposed once by the loader at upload. + InOut, } fn dims_list(d: &[usize]) -> String { @@ -348,9 +367,25 @@ impl Builder { next: 0, precision: Precision::F32, gemma3n_qmv: false, + projection_layout: ProjectionLayout::OutIn, } } + /// Set the projection weight layout (builder-style). Default is + /// [`ProjectionLayout::OutIn`], the checkpoint layout. + #[must_use] + pub fn with_projection_layout(mut self, layout: ProjectionLayout) -> Self { + self.projection_layout = layout; + self + } + + /// The projection layout this graph declares. Read by the weight-argument + /// declaration and the layout-aware linear helpers. + #[must_use] + pub fn projection_layout(&self) -> ProjectionLayout { + self.projection_layout + } + /// Set the contraction precision (builder-style). Default is `F32`. #[must_use] pub fn with_precision(mut self, precision: Precision) -> Self { @@ -1837,6 +1872,47 @@ impl Builder { self.dot_general(x, w, &[], &[], &[1], &[1], vec![l, n]) } + /// y = x @ W for x:[K], W:[K,N] (weights stored [in,out], issue #2859). + /// + /// The checkpoint layout is `[out, in]`; the loader transposes it once at + /// upload. Contracting the shared axis instead of the weight's last axis + /// keeps IREE from materializing a `[out,in] -> [in,out]` transpose on every + /// decode step, which was 59% of the step's device time. + pub fn linear_t(&mut self, x: &Val, w: &Val) -> Val { + let k = x.ty.shape[0]; + assert_eq!(k, w.ty.shape[0], "linear_t contracts the weight's in axis"); + let n = w.ty.shape[1]; + self.dot_general(x, w, &[], &[], &[0], &[0], vec![n]) + } + + /// Sequence-batched form of [`Self::linear_t`]: y[L,N] = x[L,K] @ W[K,N]. + pub fn linear_seq_t(&mut self, x: &Val, w: &Val) -> Val { + let l = x.ty.shape[0]; + let k = x.ty.shape[1]; + assert_eq!( + k, w.ty.shape[0], + "linear_seq_t contracts the weight's in axis" + ); + let n = w.ty.shape[1]; + self.dot_general(x, w, &[], &[], &[1], &[0], vec![l, n]) + } + + /// Linear over a projection weight in the layout this graph declares. + pub fn linear_proj(&mut self, x: &Val, w: &Val) -> Val { + match self.projection_layout { + ProjectionLayout::OutIn => self.linear(x, w), + ProjectionLayout::InOut => self.linear_t(x, w), + } + } + + /// Sequence-batched [`Self::linear_proj`]. + pub fn linear_seq_proj(&mut self, x: &Val, w: &Val) -> Val { + match self.projection_layout { + ProjectionLayout::OutIn => self.linear_seq(x, w), + ProjectionLayout::InOut => self.linear_seq_t(x, w), + } + } + /// Linear projection over the final axis of an arbitrary-rank tensor. pub fn linear_last(&mut self, x: &Val, w: &Val) -> Val { let input = *x.ty.shape.last().expect("linear input rank"); @@ -1902,6 +1978,37 @@ mod tests { assert_eq!(default_precision("local-sync"), Precision::F32); } + /// The transposed-layout helpers contract the weight's in axis and keep the + /// output width from its last axis (issue #2859). + #[test] + fn transposed_linear_contracts_the_in_axis() { + let mut b = Builder::new(); + let x = Builder::arg(0, Ty::f32(vec![2048])); + let w = Builder::arg(1, Ty::f32(vec![2048, 8192])); + let y = b.linear_t(&x, &w); + assert_eq!(y.ty.shape, vec![8192]); + let body = b.body(); + assert!( + body.contains("dot_general %arg0, %arg1, contracting_dims = [0] x [0]"), + "{body}" + ); + assert!(body.contains("(tensor<2048xf32>, tensor<2048x8192xf32>) -> tensor<8192xf32>")); + } + + #[test] + fn transposed_linear_seq_contracts_the_in_axis() { + let mut b = Builder::new(); + let x = Builder::arg(0, Ty::f32(vec![64, 2048])); + let w = Builder::arg(1, Ty::f32(vec![2048, 8192])); + let y = b.linear_seq_t(&x, &w); + assert_eq!(y.ty.shape, vec![64, 8192]); + assert!( + b.body().contains("contracting_dims = [1] x [0]"), + "{}", + b.body() + ); + } + /// Issue #612: only bf16-on-Metal is rejected; every other pair lowers. The /// `metal-spirv` target has no bf16 compute, but f16 / f32 lower on Metal and /// bf16 lowers on CUDA, HIP, and the CPU targets. diff --git a/src/lib/mlxcel-xla/src/emitter/config.rs b/src/lib/mlxcel-xla/src/emitter/config.rs index 6b68cdcc3..d2b71c060 100644 --- a/src/lib/mlxcel-xla/src/emitter/config.rs +++ b/src/lib/mlxcel-xla/src/emitter/config.rs @@ -41,6 +41,8 @@ //! family is not validated, an unsupported activation, yarn RoPE, MiniCPM3's MLA) //! are rejected rather than mis-emitted. +use super::builder::ProjectionLayout; + /// How the RoPE inverse-frequency table is computed. Both kinds share the /// `outer(pos, inv_freq)` table build (see [`rope`](super::rope)); they differ /// only in `inv_freq`. @@ -457,6 +459,10 @@ pub struct Config { /// Per-request speech/vision adapters carried by the Phi4MM checkpoint. /// `None` preserves the existing dense graph signatures byte-for-byte. pub phi4_lora: Option, + /// Layout the graph declares for projection weights (issue #2859). Set by + /// the engine that emits the graph: the single-sequence session loads + /// `InOut`, the continuous-batching engine keeps the checkpoint's `OutIn`. + pub projection_layout: ProjectionLayout, } impl Config { @@ -517,6 +523,7 @@ impl Config { fused_qkv: false, fused_gate_up: false, phi4_lora: None, + projection_layout: ProjectionLayout::OutIn, } } @@ -1649,6 +1656,7 @@ impl Config { fused_qkv, fused_gate_up, phi4_lora, + projection_layout: ProjectionLayout::OutIn, }) } diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index aeed3ae1c..281ebea8a 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -173,8 +173,8 @@ pub(crate) use config::{MoeConfig, SharedExpertConfig}; // superseded uniform-B Stage-1 graph) is re-exported for the validation harness. #[allow(unused_imports)] pub(crate) use builder::{ - Precision, check_packed_supported, precision_env_override, quant_in_graph, resolve_precision, - resolve_precision_checked, target_flags_with_rocm_target, + Precision, ProjectionLayout, check_packed_supported, precision_env_override, quant_in_graph, + resolve_precision, resolve_precision_checked, target_flags_with_rocm_target, }; #[allow(unused_imports)] pub(crate) use model::{ diff --git a/src/lib/mlxcel-xla/src/emitter/model.rs b/src/lib/mlxcel-xla/src/emitter/model.rs index e2ea7ac60..991ca9273 100644 --- a/src/lib/mlxcel-xla/src/emitter/model.rs +++ b/src/lib/mlxcel-xla/src/emitter/model.rs @@ -33,7 +33,9 @@ //! `embed` matrix (see [`take_lm_head`]); a tied checkpoint emits no such arg and //! is byte-identical to before. -use super::builder::{Builder, Precision, Ty, Val, precision_from_env, quant_in_graph}; +use super::builder::{ + Builder, Precision, ProjectionLayout, Ty, Val, precision_from_env, quant_in_graph, +}; use super::config::{Config, DeepStackConfig, MropeLayout, NormStyle}; use super::moe::{self, MoeLayerW, MoeSharedW}; use super::rope; @@ -393,7 +395,15 @@ fn take_weight( Ty::new(vec![out, n_groups], "f16"), format!("{loc}.biases"), ); - b.dequant_affine(&packed, &scales, &biases, qc.bits, qc.group_size) + // The packed dequant always reconstructs the checkpoint's `[out, in]`. + // The `InOut` graphs consume the transposed form, which costs one + // in-graph transpose per weight (the packed opt-in already trades + // decode throughput for upload bytes, so it keeps working). + let w = b.dequant_affine(&packed, &scales, &biases, qc.bits, qc.group_size); + match b.projection_layout() { + ProjectionLayout::OutIn => w, + ProjectionLayout::InOut => b.transpose(&w, &[1, 0]), + } } // issue #572: on the f16 GPU path, declare the resident projection weight as // an f16 arg (uploaded f16-resident by the loader) instead of an f32 arg the @@ -407,7 +417,15 @@ fn take_weight( } else { "f32" }; - take_arg(decls, idx, Ty::new(vec![out, in_], elt), loc) + // 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 + // decode step. The `OutIn` graphs keep the checkpoint layout. + let shape = match b.projection_layout() { + ProjectionLayout::OutIn => vec![out, in_], + ProjectionLayout::InOut => vec![in_, out], + }; + take_arg(decls, idx, Ty::new(shape, elt), loc) } } } @@ -1532,8 +1550,19 @@ impl AttnLayout { } fn lora_branch(&self, b: &mut Builder, x: &Val, pair: &LoraPairW) -> Val { - let low_rank = self.linear(b, x, &pair.a); - let delta = self.linear(b, &low_rank, &pair.b); + // Phi4MM adapter tensors are declared by `modal` in the checkpoint's + // `[out, in]` layout (the loader does not transpose them), so this + // branch keeps the untransposed helpers (issue #2859). + let delta = match self { + AttnLayout::Single { .. } => { + let low_rank = b.linear(x, &pair.a); + b.linear(&low_rank, &pair.b) + } + _ => { + let low_rank = b.linear_seq(x, &pair.a); + b.linear_seq(&low_rank, &pair.b) + } + }; let scale = b.const_f32(pair.scale); let scale = b.broadcast(&scale, &[], delta.ty.shape.clone()); b.multiply(&delta, &scale) @@ -1578,12 +1607,12 @@ impl AttnLayout { } } - /// A rank-aware linear `x @ W^T`: `[K] -> [N]` for single decode, `[L, K] -> - /// [L, N]` for the seq kinds. Lets the shared MLP body emit at any rank. + /// A rank-aware linear over a transposed projection weight (issue #2859): + /// `[K] -> [N]` for single decode, `[L, K] -> [L, N]` for the seq kinds. fn linear(&self, b: &mut Builder, x: &Val, w: &Val) -> Val { match self { - AttnLayout::Single { .. } => b.linear(x, w), - _ => b.linear_seq(x, w), + AttnLayout::Single { .. } => b.linear_proj(x, w), + _ => b.linear_seq_proj(x, w), } } @@ -1618,7 +1647,7 @@ impl AttnLayout { .phi4_lora .as_ref() .map(|lora| self.modal_lora_delta(b, hn, &lora.qkv)); - let q = b.linear(hn, lw.wq.as_ref().expect("ordinary q projection")); + let q = b.linear_proj(hn, lw.wq.as_ref().expect("ordinary q projection")); let q = match &lora { Some(delta) => { let delta = b.slice(delta, &[(0, nq * d)]); @@ -1628,7 +1657,7 @@ impl AttnLayout { }; let q = add_proj_bias(b, q, &lw.bq); let q = b.reshape(&q, vec![nq, d]); - let kk = b.linear(hn, lw.wk.as_ref().expect("ordinary k projection")); + let kk = b.linear_proj(hn, lw.wk.as_ref().expect("ordinary k projection")); let kk = match &lora { Some(delta) => { let delta = b.slice(delta, &[(nq * d, nq * d + nkv * d)]); @@ -1638,7 +1667,7 @@ impl AttnLayout { }; let kk = add_proj_bias(b, kk, &lw.bk); let kk = b.reshape(&kk, vec![nkv, d]); - let vv = b.linear(hn, lw.wv.as_ref().expect("ordinary v projection")); + let vv = b.linear_proj(hn, lw.wv.as_ref().expect("ordinary v projection")); let vv = match &lora { Some(delta) => { let delta = b.slice(delta, &[(nq * d + nkv * d, nq * d + 2 * nkv * d)]); @@ -1670,7 +1699,7 @@ impl AttnLayout { .phi4_lora .as_ref() .map(|lora| self.modal_lora_delta(b, hn, &lora.qkv)); - let q = b.linear_seq(hn, lw.wq.as_ref().expect("ordinary q projection")); + let q = b.linear_seq_proj(hn, lw.wq.as_ref().expect("ordinary q projection")); let q = match &lora { Some(delta) => { let delta = b.slice(delta, &[(0, n), (0, nq * d)]); @@ -1680,7 +1709,7 @@ impl AttnLayout { }; let q = add_proj_bias_seq(b, q, &lw.bq, n, nq * d); let q = b.reshape(&q, vec![n, nq, d]); - let kk = b.linear_seq(hn, lw.wk.as_ref().expect("ordinary k projection")); + let kk = b.linear_seq_proj(hn, lw.wk.as_ref().expect("ordinary k projection")); let kk = match &lora { Some(delta) => { let delta = b.slice(delta, &[(0, n), (nq * d, nq * d + nkv * d)]); @@ -1690,7 +1719,7 @@ impl AttnLayout { }; let kk = add_proj_bias_seq(b, kk, &lw.bk, n, nkv * d); let kk = b.reshape(&kk, vec![n, nkv, d]); - let vv = b.linear_seq(hn, lw.wv.as_ref().expect("ordinary v projection")); + let vv = b.linear_seq_proj(hn, lw.wv.as_ref().expect("ordinary v projection")); let vv = match &lora { Some(delta) => { let delta = b.slice(delta, &[(0, n), (nq * d + nkv * d, nq * d + 2 * nkv * d)]); @@ -1965,9 +1994,9 @@ impl AttnLayout { fn o_proj(&self, b: &mut Builder, o: &Val, lw: &LayerW) -> Val { let out = match self { AttnLayout::Single { .. } => { - b.linear(o, lw.wo.as_ref().expect("ordinary output projection")) + b.linear_proj(o, lw.wo.as_ref().expect("ordinary output projection")) } - _ => b.linear_seq(o, lw.wo.as_ref().expect("ordinary output projection")), + _ => b.linear_seq_proj(o, lw.wo.as_ref().expect("ordinary output projection")), }; let out = match &lw.phi4_lora { Some(lora) => { @@ -2158,13 +2187,13 @@ fn emit_mla_attention( AttnLayout::Prefill { lp, .. } => Some(*lp), }; let q_latent = match rows { - None => b.linear(&hn, &mw.q_a), - Some(_) => b.linear_seq(&hn, &mw.q_a), + None => b.linear_proj(&hn, &mw.q_a), + Some(_) => b.linear_seq_proj(&hn, &mw.q_a), }; let q_latent = mla_rms_norm(b, &q_latent, &mw.q_a_norm, &k.eps, rows, m.q_lora_rank); let q = match rows { - None => b.linear(&q_latent, &mw.q_b), - Some(_) => b.linear_seq(&q_latent, &mw.q_b), + None => b.linear_proj(&q_latent, &mw.q_b), + Some(_) => b.linear_seq_proj(&q_latent, &mw.q_b), }; let q = match rows { None => b.reshape(&q, vec![c.n_q, m.q_head_dim()]), @@ -2184,8 +2213,8 @@ fn emit_mla_attention( ), }; let kv = match rows { - None => b.linear(&hn, &mw.kv_a), - Some(_) => b.linear_seq(&hn, &mw.kv_a), + None => b.linear_proj(&hn, &mw.kv_a), + Some(_) => b.linear_seq_proj(&hn, &mw.kv_a), }; let (latent, k_pe) = match rows { None => ( @@ -2450,8 +2479,8 @@ fn emit_mla_attention( } }; let output = match layout { - AttnLayout::Single { .. } => b.linear(&context, &mw.o), - _ => b.linear_seq(&context, &mw.o), + AttnLayout::Single { .. } => b.linear_proj(&context, &mw.o), + _ => b.linear_seq_proj(&context, &mw.o), }; let output = post_attn_norm(b, c, k, layout, output, lw); (hn, output) @@ -2604,7 +2633,9 @@ pub fn emit_decode(c: &Config, sample: bool) -> String { } pub fn emit_decode_with(c: &Config, sample: bool, precision: Precision) -> String { - let mut b = Builder::new().with_precision(precision); + let mut b = Builder::new() + .with_precision(precision) + .with_projection_layout(c.projection_layout); let (decls, a) = build_arg_schema(&mut b, c); let k = emit_consts(&mut b, c); @@ -2886,7 +2917,9 @@ pub fn emit_decode_batched_with( c.phi4_lora.is_none(), "Phi4MM requires per-row ragged decode; uniform batched decode cannot carry request-scoped adapter modes" ); - let mut b = Builder::new().with_precision(precision); + let mut b = Builder::new() + .with_precision(precision) + .with_projection_layout(c.projection_layout); let (decls, a) = build_batched_arg_schema(&mut b, c, bsz); let k = emit_consts(&mut b, c); @@ -2936,13 +2969,13 @@ pub fn emit_decode_batched_with( .as_ref() .expect("uniform-B batched decode is emitted only for pre-norm archs"); let hn = rms_norm_seq(&mut b, &x, in_ln, &k, bsz, h); // [B, H] - let q = b.linear_seq(&hn, lw.wq.as_ref().expect("ordinary q projection")); // [B, qd] + let q = b.linear_seq_proj(&hn, lw.wq.as_ref().expect("ordinary q projection")); // [B, qd] let q = add_proj_bias_seq(&mut b, q, &lw.bq, bsz, nq * d); let q = b.reshape(&q, vec![bsz, nq, d]); - let kk = b.linear_seq(&hn, lw.wk.as_ref().expect("ordinary k projection")); // [B, kv] + let kk = b.linear_seq_proj(&hn, lw.wk.as_ref().expect("ordinary k projection")); // [B, kv] let kk = add_proj_bias_seq(&mut b, kk, &lw.bk, bsz, nkv * d); let kk = b.reshape(&kk, vec![bsz, nkv, d]); - let vv = b.linear_seq(&hn, lw.wv.as_ref().expect("ordinary v projection")); // [B, kv] + let vv = b.linear_seq_proj(&hn, lw.wv.as_ref().expect("ordinary v projection")); // [B, kv] let vv = add_proj_bias_seq(&mut b, vv, &lw.bv, bsz, nkv * d); let vv = b.reshape(&vv, vec![bsz, nkv, d]); @@ -3015,7 +3048,7 @@ pub fn emit_decode_batched_with( ); let o = b.reshape(&o, vec![bsz, nq, d]); let o = b.reshape(&o, vec![bsz, nq * d]); - let attn_out = b.linear_seq(&o, lw.wo.as_ref().expect("ordinary output projection")); // [B, H] + let attn_out = b.linear_seq_proj(&o, lw.wo.as_ref().expect("ordinary output projection")); // [B, H] x = b.add(&x, &attn_out); // FFN: the MoE block (issue #500) on a MoE layer, else the dense SwiGLU MLP @@ -3046,8 +3079,8 @@ pub fn emit_decode_batched_with( bsz, h, ); - let gate = b.linear_seq(&hn2, lw.gate.as_ref().expect("batched decode: gate")); // [B, inter] - let up = b.linear_seq(&hn2, LayerW::dense(&lw.up)); // [B, inter] + let gate = b.linear_seq_proj(&hn2, lw.gate.as_ref().expect("batched decode: gate")); // [B, inter] + let up = b.linear_seq_proj(&hn2, LayerW::dense(&lw.up)); // [B, inter] let neg = b.negate(&gate); let ex = b.exponential(&neg); let one_b = b.broadcast(&k.one, &[], vec![bsz, c.inter]); @@ -3055,7 +3088,7 @@ pub fn emit_decode_batched_with( let sig = b.divide(&one_b, &denom); let silu = b.multiply(&gate, &sig); let act = b.multiply(&silu, &up); - let down = b.linear_seq(&act, LayerW::dense(&lw.down)); // [B, H] + let down = b.linear_seq_proj(&act, LayerW::dense(&lw.down)); // [B, H] x = b.add(&x, &down); } } @@ -3229,7 +3262,9 @@ pub fn emit_decode_ragged_with( sample: bool, precision: Precision, ) -> String { - let mut b = Builder::new().with_precision(precision); + let mut b = Builder::new() + .with_precision(precision) + .with_projection_layout(c.projection_layout); let (decls, a) = build_ragged_arg_schema(&mut b, c, bsz); let k = emit_consts(&mut b, c); // Constant row indices 0..bsz for the per-row KV-write dim-0 offsets. @@ -3855,7 +3890,9 @@ fn emit_prefill_module( "post-hook state capture is only valid for DeepStack prefill" ); let lp = c.context_capacity; - let mut b = Builder::new().with_precision(precision); + let mut b = Builder::new() + .with_precision(precision) + .with_projection_layout(c.projection_layout); let (decls, a) = build_prefill_arg_schema(&mut b, c, lp, input_kind); let k = emit_consts(&mut b, c); @@ -4089,7 +4126,9 @@ pub(crate) fn emit_moe_probe_with(c: &Config, n: usize, precision: Precision) -> let mw = take_moe_weights(&mut decls, &mut idx, c, 0); let hn = take_arg(&mut decls, &mut idx, Ty::f32(vec![n, h]), "hn".into()); - let mut b = Builder::new().with_precision(precision); + let mut b = Builder::new() + .with_precision(precision) + .with_projection_layout(c.projection_layout); let k = emit_consts(&mut b, c); let out = moe::moe_block(&mut b, c, m, &mw, &k, &hn, n); diff --git a/src/lib/mlxcel-xla/src/iree.rs b/src/lib/mlxcel-xla/src/iree.rs index 6825b90c7..52f547a63 100644 --- a/src/lib/mlxcel-xla/src/iree.rs +++ b/src/lib/mlxcel-xla/src/iree.rs @@ -59,12 +59,13 @@ use safetensors::{Dtype, SafeTensors}; use crate::emitter::{ Config, DeepStackConfig, Gemma3nConfig, Gemma3nWeightSpec, LlavaVisionConfig, Precision, - QuantConfig, VisionProjector, WeightScheme, check_packed_supported, emit_decode_ragged_with, - emit_decode_with, emit_gemma3n_decode_ragged_with_qmv, emit_gemma3n_decode_with_qmv, - emit_gemma3n_prefill_embeddings_ple_with_qmv, emit_gemma3n_prefill_with_qmv, - emit_prefill_embeddings_deepstack_with, emit_prefill_embeddings_with, emit_prefill_with, - gemma3n_qmv_artifact_identity, gemma3n_qmv_is_available, gemma3n_weight_specs, - precision_env_override, quant_in_graph, resolve_precision_checked, + ProjectionLayout, QuantConfig, VisionProjector, WeightScheme, check_packed_supported, + emit_decode_ragged_with, emit_decode_with, emit_gemma3n_decode_ragged_with_qmv, + emit_gemma3n_decode_with_qmv, emit_gemma3n_prefill_embeddings_ple_with_qmv, + emit_gemma3n_prefill_with_qmv, emit_prefill_embeddings_deepstack_with, + emit_prefill_embeddings_with, emit_prefill_with, gemma3n_qmv_artifact_identity, + gemma3n_qmv_is_available, gemma3n_weight_specs, precision_env_override, quant_in_graph, + resolve_precision_checked, }; #[cfg(feature = "diagnostics")] use crate::emitter::{ @@ -90,7 +91,7 @@ use crate::{DeepStackFeatures, DeepStackPreparedPrefill, Gemma3nDensePle, Gemma3 use crate::weights::{ QuantPart, WeightSpec, bf16_to_f32, dequantize_affine, dequantize_affine_bf16_sequential, dequantize_affine_f32, dequantize_affine_stacked, f16_to_f32, f32_le_to_f32, pack_f16, - slice_rows, weight_specs, + slice_rows, transpose_projection, weight_specs, }; /// Weight-buffer element dtype passed to the C shim (issue #516 per-weight ABI): @@ -570,6 +571,23 @@ impl RuntimeConfig { .map(Self::Dense) } + /// Select the projection weight layout this engine emits and uploads + /// (issue #2859). Gemma3n's own emitters keep the checkpoint layout. + fn with_projection_layout(mut self, layout: ProjectionLayout) -> Self { + if let Self::Dense(config) = &mut self { + config.projection_layout = layout; + } + self + } + + /// The projection weight layout the emitted graphs declare. + fn projection_layout(&self) -> ProjectionLayout { + match self { + Self::Dense(config) => config.projection_layout, + Self::Gemma3n(_) => ProjectionLayout::OutIn, + } + } + fn context_capacity(&self) -> usize { match self { Self::Dense(config) => config.context_capacity, @@ -1854,10 +1872,17 @@ mod checkpoint_name_tests { /// canonical tensor name also ensures fused tensors (`qkv_proj`, `gate_up_proj`, /// `att_proj`, `ff_proj`) are dequantized once per layer instead of once per /// slice. +/// +/// `projection_layout` is the layout the emitted graphs declare (issue #2859): +/// with [`ProjectionLayout::InOut`] each `Proj` / `Rows` weight is transposed +/// once here, from the checkpoint's `[out, in]` to `[in, out]`, so no decode +/// step pays for it; [`ProjectionLayout::OutIn`] uploads the checkpoint layout +/// unchanged. fn upload_weights_with_sink( model_dir: &Path, cfg: &RuntimeConfig, resident_f16: bool, + projection_layout: ProjectionLayout, mut sink: F, ) -> Result<(), String> where @@ -2130,23 +2155,42 @@ where sink(i, WeightBuf::F32(data), WDT_F32, rank, dims)?; } WeightSpec::Proj(_) => { - let rank = checked_ffi_int(shape.len(), &format!("weight {name} rank"))?; - for (k, &s) in shape.iter().enumerate() { - dims[k] = checked_ffi_i64(s, &format!("weight {name} dim {k}"))?; + if shape.len() != 2 { + return Err(format!( + "projection weight {name} is rank {} (expected 2)", + shape.len() + )); } - if resident_f16 { - let packed = pack_f16(data_ref); + let (out, in_) = (shape[0], shape[1]); + if matches!(projection_layout, ProjectionLayout::InOut) { + let transposed = transpose_projection(data_ref, out, in_); if remaining == 0 { drop(owned_widened.take()); } - sink(i, WeightBuf::F16(packed), WDT_F16, rank, dims)?; + dims[0] = checked_ffi_i64(in_, &format!("weight {name} in axis"))?; + dims[1] = checked_ffi_i64(out, &format!("weight {name} out axis"))?; + if resident_f16 { + sink(i, WeightBuf::F16(pack_f16(&transposed)), WDT_F16, 2, dims)?; + } else { + sink(i, WeightBuf::F32(transposed), WDT_F32, 2, dims)?; + } } else { - let data = if remaining == 0 { - owned_widened.take().unwrap().0 + dims[0] = checked_ffi_i64(out, &format!("weight {name} dim 0"))?; + dims[1] = checked_ffi_i64(in_, &format!("weight {name} dim 1"))?; + if resident_f16 { + let packed = pack_f16(data_ref); + if remaining == 0 { + drop(owned_widened.take()); + } + sink(i, WeightBuf::F16(packed), WDT_F16, 2, dims)?; } else { - data_ref.clone() - }; - sink(i, WeightBuf::F32(data), WDT_F32, rank, dims)?; + let data = if remaining == 0 { + owned_widened.take().unwrap().0 + } else { + data_ref.clone() + }; + sink(i, WeightBuf::F32(data), WDT_F32, 2, dims)?; + } } } WeightSpec::Rows { start, end, .. } => { @@ -2158,16 +2202,24 @@ where } let sliced = slice_rows(data_ref, shape[0], *start, *end) .map_err(|e| format!("row-slice {name}: {e}"))?; - dims[0] = - checked_ffi_i64(*end - *start, &format!("weight {name} sliced rows"))?; - dims[1] = checked_ffi_i64(shape[1], &format!("weight {name} dim 1"))?; + // Same layout rule as `Proj` (issue #2859). + let rows = *end - *start; + let in_ = shape[1]; + let (data, (r, c)) = match projection_layout { + ProjectionLayout::OutIn => (sliced, (rows, in_)), + ProjectionLayout::InOut => { + (transpose_projection(&sliced, rows, in_), (in_, rows)) + } + }; + dims[0] = checked_ffi_i64(r, &format!("weight {name} sliced dim 0"))?; + dims[1] = checked_ffi_i64(c, &format!("weight {name} sliced dim 1"))?; if remaining == 0 { drop(owned_widened.take()); } if resident_f16 { - sink(i, WeightBuf::F16(pack_f16(&sliced)), WDT_F16, 2, dims)?; + sink(i, WeightBuf::F16(pack_f16(&data)), WDT_F16, 2, dims)?; } else { - sink(i, WeightBuf::F32(sliced), WDT_F32, 2, dims)?; + sink(i, WeightBuf::F32(data), WDT_F32, 2, dims)?; } } WeightSpec::Elements { start, end, .. } => { @@ -2410,8 +2462,12 @@ fn create_ctx_with_diagnostics( (IREE runtime; see stderr for the annotated status)" )); } - if let Err(error) = - upload_weights_with_sink(model_dir, cfg, resident_f16, |i, buf, dtype, rank, dims| { + if let Err(error) = upload_weights_with_sink( + model_dir, + cfg, + resident_f16, + cfg.projection_layout(), + |i, buf, dtype, rank, dims| { let index = checked_ffi_int(i, "weight index")?; let rc = unsafe { xla_llama_upload_weight( @@ -2430,8 +2486,8 @@ fn create_ctx_with_diagnostics( )); } Ok(()) - }) - { + }, + ) { unsafe { xla_llama_free(ctx) }; return Err(error); } @@ -2461,7 +2517,12 @@ impl IreeLlama { /// weights resident, and readies the prefill / decode calls. pub fn load(model_dir: &Path, device: &str, context_capacity: usize) -> Result { let load_started = std::time::Instant::now(); - let cfg = RuntimeConfig::from_json(model_dir, context_capacity)?; + // The single-sequence graphs contract one row, where IREE materializes a + // separate weight transpose for the checkpoint layout; the pre-transposed + // layout removes it (issue #2859). The batching engine keeps `OutIn` + // because its ragged decode lowers that layout to one fast matmul. + let cfg = RuntimeConfig::from_json(model_dir, context_capacity)? + .with_projection_layout(ProjectionLayout::InOut); report_load_timing("language-config", load_started); runtime_ffi_dimensions(&cfg, cfg.weight_specs().len())?; let gemma3_vlm_contract = gemma3_vlm_compatibility(model_dir, &cfg)?; diff --git a/src/lib/mlxcel-xla/src/weights.rs b/src/lib/mlxcel-xla/src/weights.rs index ca371b8c8..6d3268aaf 100644 --- a/src/lib/mlxcel-xla/src/weights.rs +++ b/src/lib/mlxcel-xla/src/weights.rs @@ -460,6 +460,34 @@ pub(crate) fn slice_rows( Ok(buf[start * in_..end * in_].to_vec()) } +/// Transpose a row-major `[out, in]` f32 projection into the `[in, out]` layout +/// the graphs declare (issue #2859). +/// +/// The checkpoint stores projections `[out, in]`; the graph contracts the +/// weight's input axis so IREE never materializes a transpose per decode step. +/// Done once at upload, in 32x32 blocks so each block fits a cache line span. +pub(crate) fn transpose_projection(buf: &[f32], out: usize, in_: usize) -> Vec { + assert_eq!( + buf.len(), + out.saturating_mul(in_), + "projection buffer is [out, in]" + ); + const BLOCK: usize = 32; + let mut t = vec![0.0f32; buf.len()]; + for ob in (0..out).step_by(BLOCK) { + let o_end = (ob + BLOCK).min(out); + for ib in (0..in_).step_by(BLOCK) { + let i_end = (ib + BLOCK).min(in_); + for o in ob..o_end { + for i in ib..i_end { + t[i * out + o] = buf[o * in_ + i]; + } + } + } + } + t +} + /// bf16 little-endian bytes -> f32 (bf16 is the high 16 bits of f32). pub(crate) fn bf16_to_f32(bytes: &[u8]) -> Vec { bytes @@ -1392,6 +1420,26 @@ mod tests { assert!(slice_rows(&buf, 3, 0, 1).is_err()); } + /// The projection transpose moves `[out, in]` to `[in, out]` exactly. + #[test] + fn transpose_projection_moves_the_out_axis_last() { + // 2 rows (out) x 3 cols (in), row-major. + let buf = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; + assert_eq!( + transpose_projection(&buf, 2, 3), + vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0] + ); + // A larger-than-block shape keeps every element in its transposed slot. + let (out, in_) = (40, 70); + let buf: Vec = (0..out * in_).map(|x| x as f32).collect(); + let t = transpose_projection(&buf, out, in_); + for o in 0..out { + for i in 0..in_ { + assert_eq!(t[i * out + o], buf[o * in_ + i]); + } + } + } + /// f16 widening is exact against `f32 as` for representative values: zero, one, /// a fraction, a negative, the max normal, and a subnormal. #[test]