From 9b070f6aa2d2565b21fa44c70c67d1eb05a9555f Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Thu, 5 Nov 2026 18:48:02 +0900 Subject: [PATCH] perf(xla): load single-sequence graphs with transposed weights (#2859) The single-sequence decode graph contracted each projection's last axis against the checkpoint's `[out, in]` layout, and IREE lowered that to a full-weight transpose plus a vecmat on every decode step. Profiling on gfx1151 showed 112 weight transposes per step at 96.5 ms/step (57 ms of transposes, plus a 9.65 ms f32 LM-head transpose), against 4.1 ms/step for the MLX ROCm path. The graph now declares projection weights `[in, out]` and contracts the shared axis, and the loader transposes each `Proj` / `Rows` weight once at upload (`transpose_projection`, 32x32 blocked). Graph emission picks the layout through `Config::projection_layout` and the `linear_proj` / `linear_seq_proj` helpers, so a graph's weight declarations and its contractions cannot disagree. The layout is per engine, not global, because the best lowering depends on the graph shape. The ragged batch decode already lowered the old `[out, in]` contraction to one transposed-B `linalg.matmul` (fast at its row count), while the canonical matmul over `[in, out]` was about 8x slower per call; the single-row case was the opposite. `IreeLlama` therefore emits `InOut` and `IreeRaggedLlama` keeps the checkpoint's `OutIn`, and every committed golden stays byte-identical on the `OutIn` path. Phi4MM LoRA adapter tensors and the LM head stay untransposed because the loader uploads them `Whole`. Validation on gfx1151 (hip, llama-3.2-1b-instruct-4bit, release build): B=1 CLI at capacity 256, n=64 10.5-11.4 -> 18.6-21.4 tok/s and n=128 9.0-10.0 -> 16.5-18.3 tok/s; batch B=4 N=8 63.9-69.2 -> 66.8-67.7 tok/s and B=8 N=16 80.1-80.5 -> 74.8-84.1 tok/s (unchanged path, within run-to-run spread); REFERENCE-EXACT PASS on both batch sizes. Greedy output is token-identical to the previous binary on llama-3.2-1b, qwen3-1.7b, qwen2.5-0.5b and gemma-3-4b. mlxcel-xla lib tests 385 pass with unchanged goldens; clippy clean at default and `iree` features. Refs #2859 --- src/lib/mlxcel-xla/src/emitter/builder.rs | 107 ++++++++++++++++++++ src/lib/mlxcel-xla/src/emitter/config.rs | 8 ++ src/lib/mlxcel-xla/src/emitter/mod.rs | 4 +- src/lib/mlxcel-xla/src/emitter/model.rs | 113 ++++++++++++++------- src/lib/mlxcel-xla/src/iree.rs | 117 ++++++++++++++++------ src/lib/mlxcel-xla/src/weights.rs | 48 +++++++++ 6 files changed, 330 insertions(+), 67 deletions(-) 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]